diff --git a/.github/CODEOWNERS b/.github/CODEOWNERS index 6e45ec0faef..8543394b789 100644 --- a/.github/CODEOWNERS +++ b/.github/CODEOWNERS @@ -36,23 +36,86 @@ # Backend (default owner, more specific rules below will override) /api/ @QuantumGhost -# Backend - Tests (default owners, domain-specific rules below will override) -/api/tests/ @laipz8200 @QuantumGhost -/api/tests/*/core/rag/ @JohnJyong -/api/tests/*/services/rag_pipeline/ @JohnJyong -/api/tests/*/services/dataset_service/ @JohnJyong -/api/tests/*/controllers/console/datasets/ @JohnJyong -/api/tests/*/controllers/service_api/dataset/ @JohnJyong -/api/tests/*/vdb/ @JohnJyong -/api/tests/*/core/plugin/ @WH-2099 -/api/tests/*/services/plugin/ @WH-2099 -/api/tests/*/controllers/inner_api/plugin/ @WH-2099 -/api/tests/*/controllers/trigger/ @CourTeous33 -/api/tests/*/core/trigger/ @CourTeous33 -/api/tests/*/trigger/ @CourTeous33 -/api/tests/*/controllers/console/billing/ @zyssyz123 -/api/tests/*/services/enterprise/ @GareArc -/api/tests/*/migrations/ @snakevash @laipz8200 @MRZHUH +# Backend - Tests (inherit the production owners) +/api/tests/**/core/rag/ @JohnJyong +/api/tests/**/services/knowledge/ @JohnJyong +/api/tests/**/services/rag_pipeline/ @JohnJyong +/api/tests/**/services/test_hit_testing_service.py @JohnJyong +/api/tests/**/services/entities/knowledge_entities/ @JohnJyong +/api/tests/**/services/entities/external_knowledge_entities/ @JohnJyong +/api/tests/**/controllers/console/datasets/ @JohnJyong +/api/tests/**/controllers/service_api/dataset/ @JohnJyong +/api/tests/**/models/test_dataset.py @JohnJyong +/api/tests/**/tasks/rag_pipeline/ @JohnJyong +/api/tests/**/tasks/test_add_document_to_index_task.py @JohnJyong +/api/tests/**/tasks/test_batch_clean_document_task.py @JohnJyong +/api/tests/**/tasks/test_clean_document_task.py @JohnJyong +/api/tests/**/tasks/test_clean_notion_document_task.py @JohnJyong +/api/tests/**/tasks/test_document_indexing_task.py @JohnJyong +/api/tests/**/tasks/test_document_indexing_sync_task.py @JohnJyong +/api/tests/**/tasks/test_document_indexing_update_task.py @JohnJyong +/api/tests/**/tasks/test_duplicate_document_indexing_task.py @JohnJyong +/api/tests/**/tasks/test_recover_document_indexing_task.py @JohnJyong +/api/tests/**/tasks/test_remove_document_from_index_task.py @JohnJyong +/api/tests/**/tasks/test_retry_document_indexing_task.py @JohnJyong +/api/tests/**/tasks/test_sync_website_document_indexing_task.py @JohnJyong +/api/tests/**/tasks/test_batch_create_segment_to_index_task.py @JohnJyong +/api/tests/**/tasks/test_create_segment_to_index_task.py @JohnJyong +/api/tests/**/tasks/test_delete_segment_from_index_task.py @JohnJyong +/api/tests/**/tasks/test_disable_segment_from_index_task.py @JohnJyong +/api/tests/**/tasks/test_disable_segments_from_index_task.py @JohnJyong +/api/tests/**/tasks/test_enable_segment_to_index_task.py @JohnJyong +/api/tests/**/tasks/test_enable_segments_to_index_task.py @JohnJyong +/api/tests/**/tasks/test_clean_dataset_task.py @JohnJyong +/api/tests/**/tasks/test_deal_dataset_index_update_task.py @JohnJyong +/api/tests/**/tasks/test_deal_dataset_vector_index_task.py @JohnJyong +/api/tests/**/core/plugin/ @WH-2099 +/api/tests/**/services/plugin/ @WH-2099 +/api/tests/**/controllers/console/workspace/test_plugin.py @WH-2099 +/api/tests/**/controllers/inner_api/plugin/ @WH-2099 +/api/tests/**/tasks/test_process_tenant_plugin_autoupgrade_check_task.py @WH-2099 +/api/tests/**/controllers/trigger/ @CourTeous33 +/api/tests/**/controllers/console/app/test_workflow_trigger.py @CourTeous33 +/api/tests/**/controllers/console/workspace/test_trigger_providers.py @CourTeous33 +/api/tests/**/core/trigger/ @CourTeous33 +/api/tests/**/core/app/layers/test_trigger_post_layer.py @CourTeous33 +/api/tests/**/services/trigger/ @CourTeous33 +/api/tests/**/models/test_trigger.py @CourTeous33 +/api/tests/**/repositories/test_workflow_trigger_log_repository.py @CourTeous33 +/api/tests/**/repositories/test_sqlalchemy_workflow_trigger_log_repository.py @CourTeous33 +/api/tests/**/libs/test_schedule_utils.py @CourTeous33 +/api/tests/**/services/workflow/test_scheduler.py @CourTeous33 +/api/tests/**/schedule/test_trigger_provider_refresh_task.py @CourTeous33 +/api/tests/**/schedule/test_workflow_schedule_task.py @CourTeous33 +/api/tests/**/tasks/test_trigger_processing_tasks.py @CourTeous33 +/api/tests/**/tasks/test_trigger_subscription_refresh_tasks.py @CourTeous33 +/api/tests/**/tasks/test_workflow_schedule_tasks.py @CourTeous33 +/api/tests/**/tasks/workflow_cfs_scheduler/ @CourTeous33 +/api/tests/**/events/event_handlers/test_sync_plugin_trigger_when_app_created.py @CourTeous33 +/api/tests/**/events/event_handlers/test_update_app_triggers_when_app_published_workflow_updated.py @CourTeous33 +/api/tests/**/events/event_handlers/test_sync_workflow_schedule_when_app_published.py @CourTeous33 +/api/tests/**/events/event_handlers/test_sync_webhook_when_app_created.py @CourTeous33 +/api/tests/**/services/test_async_workflow_service.py @Mairuis +/api/tests/**/tasks/test_async_workflow_tasks.py @Mairuis +/api/tests/**/services/test_billing_service.py @zyssyz123 +/api/tests/**/controllers/console/billing/ @zyssyz123 +/api/tests/**/configs/enterprise/ @GareArc +/api/tests/**/services/enterprise/ @GareArc +/api/tests/**/services/test_feature_service.py @GareArc +/api/tests/**/controllers/console/test_feature.py @GareArc +/api/tests/**/controllers/web/test_feature.py @GareArc +/api/tests/**/migrations/ @snakevash @laipz8200 @MRZHUH +/api/tests/**/configs/middleware/vdb/* @JohnJyong + +# E2E scenarios retain names spanning the corresponding production modules. +/api/tests/e2e_tests/*/services/test_dataset* @JohnJyong +/api/tests/e2e_tests/*/services/test_document* @JohnJyong +/api/tests/e2e_tests/*/services/test_metadata* @JohnJyong +/api/tests/e2e_tests/*/services/test_schedule_service.py @CourTeous33 +/api/tests/e2e_tests/*/services/test_webhook* @CourTeous33 +/api/tests/e2e_tests/*/services/test_trigger* @CourTeous33 +/api/tests/e2e_tests/*/trigger/ @CourTeous33 +/api/tests/e2e_tests/*/commands/test_legacy_model_type_migration.py @QuantumGhost # Backend - RAG (Retrieval Augmented Generation) /api/core/rag/ @JohnJyong diff --git a/.github/workflows/api-tests.yml b/.github/workflows/api-tests.yml index f6ecfdc3e0b..0af85371ac6 100644 --- a/.github/workflows/api-tests.yml +++ b/.github/workflows/api-tests.yml @@ -45,7 +45,9 @@ jobs: --write-plan unit-test-plan.json \ --split-threshold 60 \ --ignore api/tests/unit_tests/controllers \ + --ignore api/tests/e2e_tests/contracts/controllers \ api/tests/unit_tests \ + api/tests/e2e_tests/contracts \ api/providers/vdb/*/tests/unit_tests \ api/providers/trace/*/tests/unit_tests - name: Upload shared plan @@ -110,7 +112,9 @@ jobs: --shard-index "${{ matrix.shardIndex }}" \ --shard-total "${{ matrix.shardTotal }}" \ --ignore api/tests/unit_tests/controllers \ + --ignore api/tests/e2e_tests/contracts/controllers \ api/tests/unit_tests \ + api/tests/e2e_tests/contracts \ api/providers/vdb/*/tests/unit_tests \ api/providers/trace/*/tests/unit_tests > "$RUNNER_TEMP/unit-test-files.txt" uv run --project api pytest \ @@ -193,7 +197,8 @@ jobs: -o 'addopts=-p no:opik -p no:langsmith_plugin --cov=./api --cov-branch --import-mode=importlib' \ --cov-report= \ --junitxml=test-results/controllers.xml \ - api/tests/unit_tests/controllers + api/tests/unit_tests/controllers \ + api/tests/e2e_tests/contracts/controllers - name: Upload controller test reports if: ${{ !cancelled() }} @@ -265,9 +270,9 @@ jobs: -o 'addopts=-p no:opik -p no:langsmith_plugin --cov=./api --cov-branch --import-mode=importlib' \ --cov-report= \ --timeout "${PYTEST_TIMEOUT:-180}" \ - api/tests/integration_tests/workflow \ - api/tests/integration_tests/tools \ - api/tests/test_containers_integration_tests + api/tests/e2e_tests/middleware/workflow \ + api/tests/e2e_tests/middleware/tools \ + api/tests/e2e_tests/containers - name: Generate coverage report run: uv run --project api coverage xml -o coverage.xml diff --git a/.github/workflows/knowledge-fs-ci.yml b/.github/workflows/knowledge-fs-ci.yml index d584281479b..65644e260f6 100644 --- a/.github/workflows/knowledge-fs-ci.yml +++ b/.github/workflows/knowledge-fs-ci.yml @@ -41,7 +41,7 @@ jobs: filters: | knowledge-fs: - 'api/knowledge-fs-contract/**' - - 'api/tests/knowledge_fs_vector_integration/**' + - 'api/tests/e2e_tests/knowledge_fs_vector/**' - 'api/services/tidb_binding_service.py' - 'api/tests/unit_tests/services/test_tidb_binding_service.py' - 'api/providers/vdb/vdb-tidb-on-qdrant/**' @@ -82,18 +82,20 @@ jobs: - 'api/services/knowledge_fs/**' - 'api/services/knowledge_fs_capability.py' - 'api/tests/unit_tests/dev/test_generate_knowledge_fs_contract.py' - - 'api/tests/unit_tests/controllers/console/workspace/test_rbac.py' + - 'api/tests/unit_tests/extensions/test_ext_celery.py' + - 'api/tests/unit_tests/migrations/versions/test_2026_10_21_1500-b9d3f7a5c102_add_upgrade_reconciliation_lease.py' + - 'api/tests/unit_tests/models/test_agent_config_entities.py' + - 'api/tests/unit_tests/controllers/console/workspace/rbac/test_roles.py' - 'api/tests/unit_tests/core/agent/test_base_agent_runner.py' - 'api/tests/unit_tests/core/app/apps/agent_app/test_runtime_request_builder.py' - 'api/tests/unit_tests/core/workflow/nodes/agent_v2/test_runtime_request_builder.py' - - 'api/tests/unit_tests/core/workflow/nodes/tool/test_tool_node_runtime.py' - 'api/tests/unit_tests/core/workflow/test_node_runtime.py' - 'api/tests/unit_tests/services/enterprise/test_rbac_service.py' - - 'api/tests/unit_tests/repositories/workspace/test_workspace_members.py' + - 'api/tests/unit_tests/repositories/workspace/test_workspace_repository.py' - 'api/tests/unit_tests/services/test_agent_tool_inner_service.py' - 'api/tests/unit_tests/services/test_vector_space_admission_service.py' - 'api/tests/unit_tests/services/test_knowledge_fs_capability.py' - - 'api/tests/unit_tests/services/test_knowledge_fs_product_operations.py' + - 'api/tests/unit_tests/services/knowledge_fs/test_product_operations.py' - 'api/pyproject.toml' - 'api/uv.lock' - 'dify-agent/src/dify_agent/layers/dify_core_tools/client.py' @@ -251,19 +253,21 @@ jobs: if [[ "$path" == *knowledge_fs* ]]; then add_test_target "$path" fi - done < <(find api/tests/unit_tests -type f -name '*.py' -print0) + done < <(find api/tests/unit_tests api/tests/e2e_tests/contracts api/tests/helpers/unit -type f -name '*.py' -print0) test_touchpoints=( + api/tests/unit_tests/extensions/test_ext_celery.py + api/tests/unit_tests/migrations/versions/test_2026_10_21_1500-b9d3f7a5c102_add_upgrade_reconciliation_lease.py + api/tests/unit_tests/models/test_agent_config_entities.py api/tests/unit_tests/controllers/console/datasets/test_datasets.py api/tests/unit_tests/controllers/console/test_wraps.py - api/tests/unit_tests/controllers/console/workspace/test_rbac.py + api/tests/unit_tests/controllers/console/workspace/rbac/test_roles.py api/tests/unit_tests/core/agent/test_base_agent_runner.py api/tests/unit_tests/core/app/apps/agent_app/test_runtime_request_builder.py api/tests/unit_tests/core/workflow/nodes/agent_v2/test_runtime_request_builder.py - api/tests/unit_tests/core/workflow/nodes/tool/test_tool_node_runtime.py api/tests/unit_tests/core/workflow/test_node_runtime.py api/tests/unit_tests/services/enterprise/test_rbac_service.py - api/tests/unit_tests/repositories/workspace/test_workspace_members.py + api/tests/unit_tests/repositories/workspace/test_workspace_repository.py api/tests/unit_tests/services/test_agent_tool_inner_service.py api/tests/unit_tests/services/test_vector_space_admission_service.py api/tests/unit_tests/services/test_tidb_binding_service.py @@ -520,4 +524,4 @@ jobs: WEAVIATE_GRPC_ENDPOINT: grpc://localhost:50051 run: >- uv run --project api --no-sync pytest - api/tests/knowledge_fs_vector_integration --no-cov -q + api/tests/e2e_tests/knowledge_fs_vector --no-cov -q diff --git a/.github/workflows/main-ci.yml b/.github/workflows/main-ci.yml index 92973d5057d..779debb08b6 100644 --- a/.github/workflows/main-ci.yml +++ b/.github/workflows/main-ci.yml @@ -126,7 +126,7 @@ jobs: - '.github/actions/wait-e2e-web-build/**' vdb: - 'api/core/rag/datasource/**' - - 'api/tests/integration_tests/vdb/**' + - 'api/tests/e2e_tests/middleware/vdb/**' - 'api/conftest.py' - 'api/tests/pytest_dify.py' - 'api/providers/vdb/*/tests/**' diff --git a/.github/workflows/post-merge.yml b/.github/workflows/post-merge.yml index adb0aa6d581..8b9b96c1c68 100644 --- a/.github/workflows/post-merge.yml +++ b/.github/workflows/post-merge.yml @@ -50,7 +50,7 @@ jobs: - 'dify-agent-runtime/**' - 'api/pyproject.toml' - 'api/uv.lock' - - 'api/tests/integration_tests/.env.example' + - 'api/tests/e2e_tests/middleware/.env.example' - 'api/clients/agent_backend/**' - 'api/core/app/apps/agent_app/**' - 'api/core/model_manager.py' diff --git a/.github/workflows/vdb-tests-full.yml b/.github/workflows/vdb-tests-full.yml index bbe5b5a1814..0ea7ba06918 100644 --- a/.github/workflows/vdb-tests-full.yml +++ b/.github/workflows/vdb-tests-full.yml @@ -57,7 +57,7 @@ jobs: # tiflash # - name: Check VDB Ready (TiDB) - # run: uv run --project api python api/providers/vdb/tidb-vector/tests/integration_tests/check_tiflash_ready.py + # run: uv run --project api python api/providers/vdb/vdb-tidb-vector/tests/e2e_tests/check_tiflash_ready.py - name: Test Vector Stores run: | @@ -65,4 +65,4 @@ jobs: --start-vdb \ --vdb-services "weaviate,qdrant,couchbase-server,etcd,minio,milvus-standalone,pgvecto-rs,pgvector,chroma,elasticsearch,oceanbase" \ --timeout "${PYTEST_TIMEOUT:-180}" \ - api/providers/vdb/*/tests/integration_tests + api/providers/vdb/*/tests/e2e_tests diff --git a/.github/workflows/vdb-tests.yml b/.github/workflows/vdb-tests.yml index dabd413a34d..515c201cf5f 100644 --- a/.github/workflows/vdb-tests.yml +++ b/.github/workflows/vdb-tests.yml @@ -54,14 +54,14 @@ jobs: # tiflash # - name: Check VDB Ready (TiDB) - # run: uv run --project api python api/providers/vdb/tidb-vector/tests/integration_tests/check_tiflash_ready.py + # run: uv run --project api python api/providers/vdb/vdb-tidb-vector/tests/e2e_tests/check_tiflash_ready.py - name: Test Vector Stores run: | uv run --project api pytest \ --start-vdb \ --timeout "${PYTEST_TIMEOUT:-180}" \ - api/providers/vdb/vdb-chroma/tests/integration_tests \ - api/providers/vdb/vdb-pgvector/tests/integration_tests \ - api/providers/vdb/vdb-qdrant/tests/integration_tests \ - api/providers/vdb/vdb-weaviate/tests/integration_tests + api/providers/vdb/vdb-chroma/tests/e2e_tests \ + api/providers/vdb/vdb-pgvector/tests/e2e_tests \ + api/providers/vdb/vdb-qdrant/tests/e2e_tests \ + api/providers/vdb/vdb-weaviate/tests/e2e_tests diff --git a/Makefile b/Makefile index d63f1ab7ec8..fde14308248 100644 --- a/Makefile +++ b/Makefile @@ -111,11 +111,12 @@ test: echo "Running backend unit tests"; \ uv run --project api --dev pytest -p no:benchmark --timeout "$${PYTEST_TIMEOUT:-20}" -n auto \ api/tests/unit_tests \ + api/tests/e2e_tests/contracts \ api/providers/vdb/*/tests/unit_tests \ api/providers/trace/*/tests/unit_tests \ - --ignore=api/tests/unit_tests/controllers; \ + --ignore=api/tests/unit_tests/controllers --ignore=api/tests/e2e_tests/contracts/controllers; \ uv run --project api --dev pytest -p no:benchmark --timeout "$${PYTEST_TIMEOUT:-20}" -n auto --cov-append \ - api/tests/unit_tests/controllers; \ + api/tests/unit_tests/controllers api/tests/e2e_tests/contracts/controllers; \ fi @echo "✅ Unit tests complete" @@ -129,26 +130,27 @@ test-all: echo "Running backend unit tests"; \ uv run --project api --dev pytest -p no:benchmark --timeout "$${PYTEST_TIMEOUT:-20}" -n auto \ api/tests/unit_tests \ + api/tests/e2e_tests/contracts \ api/providers/vdb/*/tests/unit_tests \ api/providers/trace/*/tests/unit_tests \ - --ignore=api/tests/unit_tests/controllers; \ + --ignore=api/tests/unit_tests/controllers --ignore=api/tests/e2e_tests/contracts/controllers; \ uv run --project api --dev pytest -p no:benchmark --timeout "$${PYTEST_TIMEOUT:-20}" -n auto --cov-append \ - api/tests/unit_tests/controllers; \ + api/tests/unit_tests/controllers api/tests/e2e_tests/contracts/controllers; \ echo "Running backend integration tests"; \ uv run --project api --dev pytest -p no:benchmark --start-middleware -n auto \ --timeout "$${PYTEST_TIMEOUT:-180}" \ --cov-append \ - api/tests/integration_tests/workflow \ - api/tests/integration_tests/tools \ - api/tests/test_containers_integration_tests; \ + api/tests/e2e_tests/middleware/workflow \ + api/tests/e2e_tests/middleware/tools \ + api/tests/e2e_tests/containers; \ echo "Running VDB smoke tests"; \ uv run --project api --dev pytest --start-vdb \ --timeout "$${PYTEST_TIMEOUT:-180}" \ --cov-append \ - api/providers/vdb/vdb-chroma/tests/integration_tests \ - api/providers/vdb/vdb-pgvector/tests/integration_tests \ - api/providers/vdb/vdb-qdrant/tests/integration_tests \ - api/providers/vdb/vdb-weaviate/tests/integration_tests; \ + api/providers/vdb/vdb-chroma/tests/e2e_tests \ + api/providers/vdb/vdb-pgvector/tests/e2e_tests \ + api/providers/vdb/vdb-qdrant/tests/e2e_tests \ + api/providers/vdb/vdb-weaviate/tests/e2e_tests; \ fi @echo "✅ Tests complete" diff --git a/api/AGENTS.md b/api/AGENTS.md index c86f39de366..b7e56d5f90b 100644 --- a/api/AGENTS.md +++ b/api/AGENTS.md @@ -13,6 +13,8 @@ Run backend checks from the repository root: Run direct Python commands through `uv run --project api`. Docker-backed integration suites are normally CI-owned. Do not start long-running services as part of routine agent work. +Follow the backend developer guide's [test layout and naming requirements](README.md#testing) when adding or moving unit tests. + ## Architecture And Boundaries - Treat `api/core/` as a migration-only area: do not add files or move implementations into it. Move responsibilities out; limit changes to existing callers to the adaptations needed for that migration. diff --git a/api/README.md b/api/README.md index a3f04df82dc..6c46b0167df 100644 --- a/api/README.md +++ b/api/README.md @@ -60,6 +60,10 @@ Leave `SECRET_KEY` empty to let Dify generate a persistent key in the storage di ## Testing +Backend unit-test directories must exactly mirror the production code's relative directory structure, and filenames must match the production filenames with a `test_` prefix. For example, `api/libs/schedule_utils.py` maps to `api/tests/unit_tests/libs/test_schedule_utils.py`. + +Apply the same rule to provider-local suites: `/src//module.py` maps to `/tests/unit_tests/src//test_module.py`. + Run from the repository root: ```bash diff --git a/api/controllers/API_SCHEMA_GUIDE.md b/api/controllers/API_SCHEMA_GUIDE.md index 81f556cf63a..36666a18d47 100644 --- a/api/controllers/API_SCHEMA_GUIDE.md +++ b/api/controllers/API_SCHEMA_GUIDE.md @@ -365,7 +365,7 @@ through `application_services()` nor tell which branch of a shared translator an had to be corrected by hand anyway and it was dropped. One gate checks served responses. Every 4xx/5xx JSON response an `ExternalApi` serves in a controller unit test is -checked against that API's generated document (`tests/unit_tests/controllers/error_contract.py`). The body must be +checked against that API's generated document (`tests/helpers/unit/controllers/error_contract.py`). The body must be the standard error body, unless the status declares its own model, and its `code` must be one the operation documents for that status or one of the surface's generic codes (`ExternalApi.implicit_error_codes`). A new business error without a declaration fails there on every path a test exercises, so give a new error path a test that serves it. @@ -410,7 +410,7 @@ For schema and documentation changes, run focused tests and generate Swagger JSO ```bash uv run --project api pytest api/tests/unit_tests/controllers/common/test_schema.py -uv run --project api pytest api/tests/unit_tests/commands/test_generate_swagger_specs.py api/tests/unit_tests/controllers/test_swagger.py +uv run --project api pytest api/tests/unit_tests/dev/test_generate_swagger_specs.py api/tests/unit_tests/controllers/service_api/test___init__.py uv run --project api python api/dev/generate_swagger_specs.py --output-dir /tmp/dify-openapi-check ``` diff --git a/api/providers/trace/trace-aliyun/tests/unit_tests/aliyun_trace/data_exporter/test_traceclient.py b/api/providers/trace/trace-aliyun/tests/unit_tests/src/dify_trace_aliyun/data_exporter/test_traceclient.py similarity index 100% rename from api/providers/trace/trace-aliyun/tests/unit_tests/aliyun_trace/data_exporter/test_traceclient.py rename to api/providers/trace/trace-aliyun/tests/unit_tests/src/dify_trace_aliyun/data_exporter/test_traceclient.py diff --git a/api/providers/trace/trace-aliyun/tests/unit_tests/aliyun_trace/entities/test_aliyun_trace_entity.py b/api/providers/trace/trace-aliyun/tests/unit_tests/src/dify_trace_aliyun/entities/test_aliyun_trace_entity.py similarity index 100% rename from api/providers/trace/trace-aliyun/tests/unit_tests/aliyun_trace/entities/test_aliyun_trace_entity.py rename to api/providers/trace/trace-aliyun/tests/unit_tests/src/dify_trace_aliyun/entities/test_aliyun_trace_entity.py diff --git a/api/providers/trace/trace-aliyun/tests/unit_tests/aliyun_trace/entities/test_semconv.py b/api/providers/trace/trace-aliyun/tests/unit_tests/src/dify_trace_aliyun/entities/test_semconv.py similarity index 100% rename from api/providers/trace/trace-aliyun/tests/unit_tests/aliyun_trace/entities/test_semconv.py rename to api/providers/trace/trace-aliyun/tests/unit_tests/src/dify_trace_aliyun/entities/test_semconv.py diff --git a/api/providers/trace/trace-aliyun/tests/unit_tests/aliyun_trace/test_aliyun_trace.py b/api/providers/trace/trace-aliyun/tests/unit_tests/src/dify_trace_aliyun/test_aliyun_trace.py similarity index 99% rename from api/providers/trace/trace-aliyun/tests/unit_tests/aliyun_trace/test_aliyun_trace.py rename to api/providers/trace/trace-aliyun/tests/unit_tests/src/dify_trace_aliyun/test_aliyun_trace.py index 25ffc570c9c..0225221faf5 100644 --- a/api/providers/trace/trace-aliyun/tests/unit_tests/aliyun_trace/test_aliyun_trace.py +++ b/api/providers/trace/trace-aliyun/tests/unit_tests/src/dify_trace_aliyun/test_aliyun_trace.py @@ -60,7 +60,7 @@ from core.ops.entities.trace_entity import ( WorkflowTraceInfo, ) from graphon.enums import BuiltinNodeTypes, WorkflowNodeExecutionMetadataKey, WorkflowNodeExecutionStatus -from tests.unit_tests.core.ops.trace_fixtures import workflow_node_execution +from tests.helpers.unit.core.ops.trace_fixtures import workflow_node_execution class RecordingTraceClient: diff --git a/api/providers/trace/trace-aliyun/tests/unit_tests/test_config_entity.py b/api/providers/trace/trace-aliyun/tests/unit_tests/src/dify_trace_aliyun/test_config.py similarity index 100% rename from api/providers/trace/trace-aliyun/tests/unit_tests/test_config_entity.py rename to api/providers/trace/trace-aliyun/tests/unit_tests/src/dify_trace_aliyun/test_config.py diff --git a/api/providers/trace/trace-aliyun/tests/unit_tests/aliyun_trace/test_aliyun_trace_utils.py b/api/providers/trace/trace-aliyun/tests/unit_tests/src/dify_trace_aliyun/test_utils.py similarity index 99% rename from api/providers/trace/trace-aliyun/tests/unit_tests/aliyun_trace/test_aliyun_trace_utils.py rename to api/providers/trace/trace-aliyun/tests/unit_tests/src/dify_trace_aliyun/test_utils.py index 7571b168076..f38034100ec 100644 --- a/api/providers/trace/trace-aliyun/tests/unit_tests/aliyun_trace/test_aliyun_trace_utils.py +++ b/api/providers/trace/trace-aliyun/tests/unit_tests/src/dify_trace_aliyun/test_utils.py @@ -44,7 +44,7 @@ from core.rag.models.document import Document from graphon.enums import WorkflowNodeExecutionStatus from models import EndUser from models.enums import EndUserType -from tests.unit_tests.core.ops.trace_fixtures import workflow_node_execution +from tests.helpers.unit.core.ops.trace_fixtures import workflow_node_execution def test_get_user_id_from_message_data_no_end_user(monkeypatch: pytest.MonkeyPatch): diff --git a/api/providers/trace/trace-arize-phoenix/tests/unit_tests/arize_phoenix_trace/test_arize_phoenix_trace.py b/api/providers/trace/trace-arize-phoenix/tests/unit_tests/src/dify_trace_arize_phoenix/test_arize_phoenix_trace.py similarity index 100% rename from api/providers/trace/trace-arize-phoenix/tests/unit_tests/arize_phoenix_trace/test_arize_phoenix_trace.py rename to api/providers/trace/trace-arize-phoenix/tests/unit_tests/src/dify_trace_arize_phoenix/test_arize_phoenix_trace.py diff --git a/api/providers/trace/trace-arize-phoenix/tests/unit_tests/test_config_entity.py b/api/providers/trace/trace-arize-phoenix/tests/unit_tests/src/dify_trace_arize_phoenix/test_config.py similarity index 100% rename from api/providers/trace/trace-arize-phoenix/tests/unit_tests/test_config_entity.py rename to api/providers/trace/trace-arize-phoenix/tests/unit_tests/src/dify_trace_arize_phoenix/test_config.py diff --git a/api/providers/trace/trace-langfuse/tests/unit_tests/langfuse_trace/test_langfuse_trace.py b/api/providers/trace/trace-langfuse/tests/unit_tests/langfuse_trace/test_langfuse_trace.py deleted file mode 100644 index 951a1559cc3..00000000000 --- a/api/providers/trace/trace-langfuse/tests/unit_tests/langfuse_trace/test_langfuse_trace.py +++ /dev/null @@ -1,777 +0,0 @@ -"""Unit tests for Langfuse trace translation with real SQLite-backed lookups.""" - -import collections -import logging -from datetime import UTC, datetime, timedelta -from types import SimpleNamespace -from typing import override -from unittest.mock import MagicMock - -import pytest -from dify_trace_langfuse.config import LangfuseConfig -from dify_trace_langfuse.entities.langfuse_trace_entity import ( - LangfuseGeneration, - LangfuseSpan, - LangfuseTrace, - LevelEnum, - UnitEnum, -) -from dify_trace_langfuse.langfuse_trace import LangFuseDataTrace -from sqlalchemy.orm import Session - -from core.ops.entities.trace_entity import ( - DatasetRetrievalTraceInfo, - GenerateNameTraceInfo, - MessageTraceInfo, - ModerationTraceInfo, - SuggestedQuestionTraceInfo, - ToolTraceInfo, - TraceTaskName, - WorkflowTraceInfo, -) -from graphon.enums import BuiltinNodeTypes -from models import EndUser -from models.enums import EndUserType, MessageStatus -from tests.unit_tests.core.ops.trace_fixtures import message_trace_info, tool_trace_info, workflow_trace_info - - -def _dt() -> datetime: - return datetime(2024, 1, 1, 0, 0, 0, tzinfo=UTC) - - -@pytest.fixture -def langfuse_config(): - return LangfuseConfig(public_key="pk-123", secret_key="sk-123", host="https://cloud.langfuse.com") - - -@pytest.fixture -def trace_instance(langfuse_config, monkeypatch: pytest.MonkeyPatch): - # Mock Langfuse client to avoid network calls - mock_client = MagicMock() - monkeypatch.setattr("dify_trace_langfuse.langfuse_trace.LangfuseAPI", lambda **kwargs: mock_client) - - instance = LangFuseDataTrace(langfuse_config) - return instance - - -def test_init(langfuse_config, monkeypatch: pytest.MonkeyPatch): - mock_langfuse = MagicMock() - monkeypatch.setattr("dify_trace_langfuse.langfuse_trace.LangfuseAPI", mock_langfuse) - monkeypatch.setenv("FILES_URL", "http://test.url") - - instance = LangFuseDataTrace(langfuse_config) - - mock_langfuse.assert_called_once() - kwargs = mock_langfuse.call_args.kwargs - assert kwargs["username"] == langfuse_config.public_key - assert kwargs["password"] == langfuse_config.secret_key - assert kwargs["base_url"] == langfuse_config.host - assert instance.file_base_url == "http://test.url" - - -def test_same_public_key_uses_independent_api_clients(monkeypatch: pytest.MonkeyPatch): - clients = [MagicMock(), MagicMock()] - create_client = MagicMock(side_effect=clients) - monkeypatch.setattr("dify_trace_langfuse.langfuse_trace.LangfuseAPI", create_client) - - first = LangFuseDataTrace( - LangfuseConfig(public_key="shared", secret_key="secret-a", host="https://tenant-a.example") - ) - second = LangFuseDataTrace( - LangfuseConfig(public_key="shared", secret_key="secret-b", host="https://tenant-b.example") - ) - - assert first.langfuse_client is clients[0] - assert second.langfuse_client is clients[1] - assert [ - (call.kwargs["base_url"], call.kwargs["username"], call.kwargs["password"]) - for call in create_client.call_args_list - ] == [ - ("https://tenant-a.example", "shared", "secret-a"), - ("https://tenant-b.example", "shared", "secret-b"), - ] - - -def test_close_is_idempotent(langfuse_config, monkeypatch: pytest.MonkeyPatch): - http_client = MagicMock() - monkeypatch.setattr("dify_trace_langfuse.langfuse_trace.httpx.Client", lambda **kwargs: http_client) - monkeypatch.setattr("dify_trace_langfuse.langfuse_trace.LangfuseAPI", MagicMock()) - - instance = LangFuseDataTrace(langfuse_config) - instance.close() - instance.close() - - http_client.close.assert_called_once() - - -def test_trace_dispatch(trace_instance, monkeypatch: pytest.MonkeyPatch): - methods = [ - "workflow_trace", - "message_trace", - "moderation_trace", - "suggested_question_trace", - "dataset_retrieval_trace", - "tool_trace", - "generate_name_trace", - ] - mocks = {method: MagicMock() for method in methods} - for method, m in mocks.items(): - monkeypatch.setattr(trace_instance, method, m) - - # WorkflowTraceInfo - info = workflow_trace_info() - trace_instance.trace(info) - mocks["workflow_trace"].assert_called_once_with(info) - - # MessageTraceInfo - info = message_trace_info() - trace_instance.trace(info) - mocks["message_trace"].assert_called_once_with(info) - - # ModerationTraceInfo - info = ModerationTraceInfo(metadata={}, flagged=False, action="allow", preset_response="", query="hello") - trace_instance.trace(info) - mocks["moderation_trace"].assert_called_once_with(info) - - # SuggestedQuestionTraceInfo - info = SuggestedQuestionTraceInfo(metadata={}, total_tokens=0, suggested_question=[], level="info") - trace_instance.trace(info) - mocks["suggested_question_trace"].assert_called_once_with(info) - - # DatasetRetrievalTraceInfo - info = DatasetRetrievalTraceInfo(metadata={}) - trace_instance.trace(info) - mocks["dataset_retrieval_trace"].assert_called_once_with(info) - - # ToolTraceInfo - info = tool_trace_info() - trace_instance.trace(info) - mocks["tool_trace"].assert_called_once_with(info) - - # GenerateNameTraceInfo - info = GenerateNameTraceInfo(tenant_id="tenant-1", metadata={}) - trace_instance.trace(info) - mocks["generate_name_trace"].assert_called_once_with(info) - - -@pytest.mark.parametrize("sqlite3_session", [()], indirect=True) -def test_workflow_trace_with_message_id( - trace_instance, monkeypatch: pytest.MonkeyPatch, sqlite3_session: Session -) -> None: - # Setup trace info - trace_info = WorkflowTraceInfo( - workflow_id="wf-1", - tenant_id="tenant-1", - workflow_run_id="run-1", - workflow_run_elapsed_time=1.0, - workflow_run_status="succeeded", - workflow_run_inputs={"input": "hi"}, - workflow_run_outputs={"output": "hello"}, - workflow_run_version="1.0", - message_id="msg-1", - conversation_id="conv-1", - total_tokens=100, - file_list=[], - query="hi", - start_time=_dt(), - end_time=_dt() + timedelta(seconds=1), - trace_id="trace-1", - metadata={"app_id": "app-1", "user_id": "user-1"}, - workflow_app_log_id="log-1", - error="", - ) - - monkeypatch.setattr( - "dify_trace_langfuse.langfuse_trace.db", - SimpleNamespace(engine=sqlite3_session.get_bind(), session=sqlite3_session), - ) - - # Mock node executions - node_llm = MagicMock() - node_llm.id = "node-llm" - node_llm.title = "LLM Node" - node_llm.node_type = BuiltinNodeTypes.LLM - node_llm.status = "succeeded" - node_llm.process_data = { - "model_mode": "chat", - "model_name": "gpt-4", - "model_provider": "openai", - "usage": {"prompt_tokens": 10, "completion_tokens": 20}, - } - node_llm.inputs = {"prompts": "p"} - node_llm.outputs = {"text": "t"} - node_llm.created_at = _dt() - node_llm.elapsed_time = 0.5 - node_llm.metadata = {"foo": "bar"} - - node_other = MagicMock() - node_other.id = "node-other" - node_other.title = "Other Node" - node_other.node_type = BuiltinNodeTypes.CODE - node_other.status = "failed" - node_other.process_data = None - node_other.inputs = {"code": "print"} - node_other.outputs = {"result": "ok"} - node_other.created_at = None # Trigger datetime.now() branch - node_other.elapsed_time = 0.2 - node_other.metadata = None - - repo = MagicMock() - repo.get_by_workflow_execution.return_value = [node_llm, node_other] - - mock_factory = MagicMock() - mock_factory.create_workflow_node_execution_repository.return_value = repo - monkeypatch.setattr("dify_trace_langfuse.langfuse_trace.DifyCoreRepositoryFactory", mock_factory) - - monkeypatch.setattr(trace_instance, "get_service_account_with_tenant", lambda app_id: MagicMock()) - - # Track calls to add_trace, add_span, add_generation - trace_instance.add_trace = MagicMock() - trace_instance.add_span = MagicMock() - trace_instance.add_generation = MagicMock() - - trace_instance.workflow_trace(trace_info) - - # Verify add_trace (Workflow Level) - trace_instance.add_trace.assert_called_once() - trace_data = trace_instance.add_trace.call_args[1]["langfuse_trace_data"] - assert trace_data.id == "trace-1" - assert trace_data.name == TraceTaskName.MESSAGE_TRACE - assert "message" in trace_data.tags - assert "workflow" in trace_data.tags - - # Verify add_span (Workflow Run Span) - assert trace_instance.add_span.call_count >= 1 - # First span should be workflow run span because message_id is present - workflow_span = trace_instance.add_span.call_args_list[0][1]["langfuse_span_data"] - assert workflow_span.id == "run-1" - assert workflow_span.name == TraceTaskName.WORKFLOW_TRACE - - # Verify Generation for LLM node - trace_instance.add_generation.assert_called_once() - gen_data = trace_instance.add_generation.call_args[1]["langfuse_generation_data"] - assert gen_data.id == "node-llm" - assert gen_data.usage.input == 10 - assert gen_data.usage.output == 20 - - # Verify normal span for Other node - # Second add_span call - other_span = trace_instance.add_span.call_args_list[1][1]["langfuse_span_data"] - assert other_span.id == "node-other" - assert other_span.level == LevelEnum.ERROR - - -@pytest.mark.parametrize("sqlite3_session", [()], indirect=True) -def test_workflow_trace_no_message_id( - trace_instance, monkeypatch: pytest.MonkeyPatch, sqlite3_session: Session -) -> None: - trace_info = WorkflowTraceInfo( - workflow_id="wf-1", - tenant_id="tenant-1", - workflow_run_id="run-1", - workflow_run_elapsed_time=1.0, - workflow_run_status="succeeded", - workflow_run_inputs={}, - workflow_run_outputs={}, - workflow_run_version="1.0", - total_tokens=0, - file_list=[], - query="", - message_id=None, - conversation_id="conv-1", - start_time=_dt(), - end_time=_dt(), - trace_id=None, # Should fallback to workflow_run_id - metadata={"app_id": "app-1"}, - workflow_app_log_id="log-1", - error="", - ) - - monkeypatch.setattr( - "dify_trace_langfuse.langfuse_trace.db", - SimpleNamespace(engine=sqlite3_session.get_bind(), session=sqlite3_session), - ) - repo = MagicMock() - repo.get_by_workflow_execution.return_value = [] - mock_factory = MagicMock() - mock_factory.create_workflow_node_execution_repository.return_value = repo - monkeypatch.setattr("dify_trace_langfuse.langfuse_trace.DifyCoreRepositoryFactory", mock_factory) - monkeypatch.setattr(trace_instance, "get_service_account_with_tenant", lambda app_id: MagicMock()) - - trace_instance.add_trace = MagicMock() - trace_instance.workflow_trace(trace_info) - - trace_instance.add_trace.assert_called_once() - trace_data = trace_instance.add_trace.call_args[1]["langfuse_trace_data"] - assert trace_data.id == "run-1" - assert trace_data.name == TraceTaskName.WORKFLOW_TRACE - - -@pytest.mark.parametrize("sqlite3_session", [()], indirect=True) -def test_workflow_trace_missing_app_id( - trace_instance, monkeypatch: pytest.MonkeyPatch, sqlite3_session: Session -) -> None: - trace_info = WorkflowTraceInfo( - workflow_id="wf-1", - tenant_id="tenant-1", - workflow_run_id="run-1", - workflow_run_elapsed_time=1.0, - workflow_run_status="succeeded", - workflow_run_inputs={}, - workflow_run_outputs={}, - workflow_run_version="1.0", - total_tokens=0, - file_list=[], - query="", - message_id=None, - conversation_id="conv-1", - start_time=_dt(), - end_time=_dt(), - metadata={}, # Missing app_id - workflow_app_log_id="log-1", - error="", - ) - monkeypatch.setattr( - "dify_trace_langfuse.langfuse_trace.db", - SimpleNamespace(engine=sqlite3_session.get_bind(), session=sqlite3_session), - ) - - with pytest.raises(ValueError, match="No app_id found in trace_info metadata"): - trace_instance.workflow_trace(trace_info) - - -def test_message_trace_basic(trace_instance, monkeypatch: pytest.MonkeyPatch): - message_data = MagicMock() - message_data.id = "msg-1" - message_data.from_account_id = "acc-1" - message_data.from_end_user_id = None - message_data.provider_response_latency = 0.5 - message_data.conversation_id = "conv-1" - message_data.total_price = 0.01 - message_data.model_id = "gpt-4" - message_data.answer = "hello" - message_data.status = MessageStatus.NORMAL - message_data.error = None - - trace_info = MessageTraceInfo( - message_id="msg-1", - message_data=message_data, - inputs={"query": "hi"}, - outputs={"answer": "hello"}, - message_tokens=10, - answer_tokens=20, - total_tokens=30, - start_time=_dt(), - end_time=_dt() + timedelta(seconds=1), - trace_id="trace-1", - metadata={"foo": "bar"}, - conversation_mode="chat", - conversation_model="gpt-4", - file_list=[], - error=None, - ) - - trace_instance.add_trace = MagicMock() - trace_instance.add_generation = MagicMock() - - trace_instance.message_trace(trace_info) - - trace_instance.add_trace.assert_called_once() - trace_instance.add_generation.assert_called_once() - - gen_data = trace_instance.add_generation.call_args[0][0] - assert gen_data.name == "llm" - assert gen_data.usage.total == 30 - - -@pytest.mark.parametrize("sqlite3_session", [(EndUser,)], indirect=True) -def test_message_trace_with_end_user(trace_instance, monkeypatch: pytest.MonkeyPatch, sqlite3_session: Session) -> None: - message_data = MagicMock() - message_data.id = "msg-1" - message_data.from_account_id = "acc-1" - message_data.from_end_user_id = "end-user-1" - message_data.conversation_id = "conv-1" - message_data.status = MessageStatus.NORMAL - message_data.model_id = "gpt-4" - message_data.error = "" - message_data.answer = "hello" - message_data.total_price = 0.0 - message_data.provider_response_latency = 0.1 - - trace_info = MessageTraceInfo( - message_id="msg-1", - message_data=message_data, - inputs={}, - outputs={}, - message_tokens=0, - answer_tokens=0, - total_tokens=0, - start_time=_dt(), - end_time=_dt(), - metadata={}, - conversation_mode="chat", - conversation_model="gpt-4", - file_list=[], - error=None, - ) - - end_user = EndUser( - id="end-user-1", - tenant_id="tenant-1", - app_id="app-1", - type=EndUserType.BROWSER, - session_id="session-id-123", - ) - engine = sqlite3_session.get_bind() - with Session(engine) as write_session: - write_session.add(end_user) - write_session.commit() - - with Session(engine) as read_session: - monkeypatch.setattr( - "dify_trace_langfuse.langfuse_trace.db", - SimpleNamespace(engine=engine, session=read_session), - ) - - trace_instance.add_trace = MagicMock() - trace_instance.add_generation = MagicMock() - - trace_instance.message_trace(trace_info) - - trace_data = trace_instance.add_trace.call_args[1]["langfuse_trace_data"] - assert trace_data.user_id == "session-id-123" - assert trace_data.metadata["user_id"] == "session-id-123" - - -def test_message_trace_none_data(trace_instance): - trace_info = SimpleNamespace(message_data=None, file_list=[], metadata={}) - trace_instance.add_trace = MagicMock() - trace_instance.message_trace(trace_info) - trace_instance.add_trace.assert_not_called() - - -def test_moderation_trace(trace_instance): - message_data = MagicMock() - message_data.created_at = _dt() - - trace_info = ModerationTraceInfo( - message_id="msg-1", - message_data=message_data, - inputs={"q": "hi"}, - action="stop", - flagged=True, - preset_response="blocked", - start_time=None, - end_time=None, - metadata={"foo": "bar"}, - trace_id="trace-1", - query="hi", - ) - - trace_instance.add_span = MagicMock() - trace_instance.moderation_trace(trace_info) - - trace_instance.add_span.assert_called_once() - span_data = trace_instance.add_span.call_args[1]["langfuse_span_data"] - assert span_data.name == TraceTaskName.MODERATION_TRACE - assert span_data.output["flagged"] is True - - -def test_suggested_question_trace(trace_instance): - message_data = MagicMock() - message_data.status = MessageStatus.NORMAL - message_data.error = None - - trace_info = SuggestedQuestionTraceInfo( - message_id="msg-1", - message_data=message_data, - inputs="hi", - suggested_question=["q1"], - total_tokens=10, - level="info", - start_time=_dt(), - end_time=_dt(), - metadata={}, - trace_id="trace-1", - ) - - trace_instance.add_generation = MagicMock() - trace_instance.suggested_question_trace(trace_info) - - trace_instance.add_generation.assert_called_once() - gen_data = trace_instance.add_generation.call_args[1]["langfuse_generation_data"] - assert gen_data.name == TraceTaskName.SUGGESTED_QUESTION_TRACE - assert gen_data.usage.unit == UnitEnum.CHARACTERS - - -def test_dataset_retrieval_trace(trace_instance): - message_data = MagicMock() - message_data.created_at = _dt() - message_data.updated_at = _dt() - - trace_info = DatasetRetrievalTraceInfo( - message_id="msg-1", - message_data=message_data, - inputs="query", - documents=[{"id": "doc1"}], - start_time=None, - end_time=None, - metadata={}, - trace_id="trace-1", - ) - - trace_instance.add_span = MagicMock() - trace_instance.dataset_retrieval_trace(trace_info) - - trace_instance.add_span.assert_called_once() - span_data = trace_instance.add_span.call_args[1]["langfuse_span_data"] - assert span_data.name == TraceTaskName.DATASET_RETRIEVAL_TRACE - assert span_data.output["documents"] == [{"id": "doc1"}] - - -def test_tool_trace(trace_instance): - trace_info = ToolTraceInfo( - message_id="msg-1", - message_data=MagicMock(), - inputs={}, - outputs={}, - tool_name="my_tool", - tool_inputs={"a": 1}, - tool_outputs="result_string", - time_cost=0.1, - start_time=_dt(), - end_time=_dt(), - metadata={}, - trace_id="trace-1", - tool_config={}, - tool_parameters={}, - error="some error", - ) - - trace_instance.add_span = MagicMock() - trace_instance.tool_trace(trace_info) - - trace_instance.add_span.assert_called_once() - span_data = trace_instance.add_span.call_args[1]["langfuse_span_data"] - assert span_data.name == "my_tool" - assert span_data.level == LevelEnum.ERROR - - -def test_generate_name_trace(trace_instance): - trace_info = GenerateNameTraceInfo( - inputs={"q": "hi"}, - outputs={"name": "new"}, - tenant_id="tenant-1", - conversation_id="conv-1", - start_time=_dt(), - end_time=_dt(), - metadata={"m": 1}, - ) - - trace_instance.add_trace = MagicMock() - trace_instance.add_span = MagicMock() - - trace_instance.generate_name_trace(trace_info) - - trace_instance.add_trace.assert_called_once() - trace_instance.add_span.assert_called_once() - - trace_data = trace_instance.add_trace.call_args[1]["langfuse_trace_data"] - assert trace_data.name == TraceTaskName.GENERATE_NAME_TRACE - assert trace_data.user_id == "tenant-1" - - span_data = trace_instance.add_span.call_args[1]["langfuse_span_data"] - assert span_data.trace_id == "conv-1" - - -def test_add_trace_success(trace_instance): - data = LangfuseTrace(id="t1", name="trace") - trace_instance.add_trace(data) - trace_instance.langfuse_client.ingestion.batch.assert_called_once() - - -def test_add_trace_error(trace_instance): - trace_instance.langfuse_client.ingestion.batch.side_effect = Exception("error") - data = LangfuseTrace(id="t1", name="trace") - with pytest.raises(ValueError, match="LangFuse Failed to create trace: error"): - trace_instance.add_trace(data) - - -def test_add_span_success(trace_instance): - data = LangfuseSpan(id="s1", name="span", trace_id="t1") - trace_instance.add_span(data) - trace_instance.langfuse_client.ingestion.batch.assert_called_once() - - -def test_add_span_error(trace_instance): - trace_instance.langfuse_client.ingestion.batch.side_effect = Exception("error") - data = LangfuseSpan(id="s1", name="span", trace_id="t1") - with pytest.raises(ValueError, match="LangFuse Failed to create span: error"): - trace_instance.add_span(data) - - -def test_update_span(trace_instance): - span = MagicMock() - data = LangfuseSpan(id="s1", name="span", trace_id="t1") - trace_instance.update_span(span, data) - span.end.assert_called_once() - - -def test_add_generation_success(trace_instance): - data = LangfuseGeneration(id="g1", name="gen", trace_id="t1") - trace_instance.add_generation(data) - trace_instance.langfuse_client.ingestion.batch.assert_called_once() - - -def test_add_generation_error(trace_instance): - trace_instance.langfuse_client.ingestion.batch.side_effect = Exception("error") - data = LangfuseGeneration(id="g1", name="gen", trace_id="t1") - with pytest.raises(ValueError, match="LangFuse Failed to create generation: error"): - trace_instance.add_generation(data) - - -def test_update_generation(trace_instance): - gen = MagicMock() - data = LangfuseGeneration(id="g1", name="gen", trace_id="t1") - trace_instance.update_generation(gen, data) - gen.end.assert_called_once() - - -def test_api_check_success(trace_instance): - trace_instance.langfuse_client.projects.get.return_value = MagicMock(data=[MagicMock()]) - assert trace_instance.api_check() is True - - -def test_api_check_rejects_empty_project_list(trace_instance): - trace_instance.langfuse_client.projects.get.return_value = MagicMock(data=[]) - with pytest.raises(ValueError, match="no project found for the provided credentials"): - trace_instance.api_check() - - -def test_api_check_error(trace_instance): - trace_instance.langfuse_client.projects.get.side_effect = Exception("fail") - with pytest.raises(ValueError, match="LangFuse API check failed: fail"): - trace_instance.api_check() - - -def test_get_project_key_success(trace_instance): - mock_data = MagicMock() - mock_data.id = "proj-1" - trace_instance.langfuse_client.projects.get.return_value = MagicMock(data=[mock_data]) - assert trace_instance.get_project_key() == "proj-1" - - -def test_get_project_key_error(trace_instance): - trace_instance.langfuse_client.projects.get.side_effect = Exception("fail") - with pytest.raises(ValueError, match="LangFuse get project key failed: fail"): - trace_instance.get_project_key() - - -def test_moderation_trace_none(trace_instance): - trace_info = ModerationTraceInfo( - message_id="m", - message_data=None, - inputs={}, - action="s", - flagged=False, - preset_response="", - query="", - metadata={}, - ) - trace_instance.add_span = MagicMock() - trace_instance.moderation_trace(trace_info) - trace_instance.add_span.assert_not_called() - - -def test_suggested_question_trace_none(trace_instance): - trace_info = SuggestedQuestionTraceInfo( - message_id="m", message_data=None, inputs={}, suggested_question=[], total_tokens=0, level="i", metadata={} - ) - trace_instance.add_generation = MagicMock() - trace_instance.suggested_question_trace(trace_info) - trace_instance.add_generation.assert_not_called() - - -def test_dataset_retrieval_trace_none(trace_instance): - trace_info = DatasetRetrievalTraceInfo(message_id="m", message_data=None, inputs={}, documents=[], metadata={}) - trace_instance.add_span = MagicMock() - trace_instance.dataset_retrieval_trace(trace_info) - trace_instance.add_span.assert_not_called() - - -def test_langfuse_trace_entity_with_list_dict_input(): - # To cover lines 29-31 in langfuse_trace_entity.py - # We need to mock replace_text_with_content or just check if it works - # Actually replace_text_with_content is imported from core.ops.utils - data = LangfuseTrace(id="t1", name="n", input=[{"text": "hello"}]) - assert isinstance(data.input, list) - assert data.input[0]["content"] == "hello" - - -@pytest.mark.parametrize("sqlite3_session", [()], indirect=True) -def test_workflow_trace_handles_usage_extraction_error( - trace_instance, - monkeypatch: pytest.MonkeyPatch, - caplog: pytest.LogCaptureFixture, - sqlite3_session: Session, -): - # Setup trace info to trigger LLM node usage extraction - trace_info = WorkflowTraceInfo( - workflow_id="wf-1", - tenant_id="t", - workflow_run_id="r", - workflow_run_elapsed_time=1.0, - workflow_run_status="s", - workflow_run_inputs={}, - workflow_run_outputs={}, - workflow_run_version="1", - total_tokens=0, - file_list=[], - query="", - message_id=None, - conversation_id="c", - start_time=_dt(), - end_time=_dt(), - metadata={"app_id": "app-1"}, - workflow_app_log_id="l", - error="", - ) - - node = MagicMock() - node.id = "n1" - node.title = "LLM Node" - node.node_type = BuiltinNodeTypes.LLM - node.status = "succeeded" - - class BadDict(collections.UserDict): - @override - def get(self, key, default=None): - if key == "usage": - raise Exception("Usage extraction failed") - return super().get(key, default) - - node.process_data = BadDict({"model_mode": "chat", "model_name": "gpt-4", "usage": True, "prompts": ["p"]}) - node.created_at = _dt() - node.elapsed_time = 0.1 - node.metadata = {} - node.outputs = {} - - repo = MagicMock() - repo.get_by_workflow_execution.return_value = [node] - mock_factory = MagicMock() - mock_factory.create_workflow_node_execution_repository.return_value = repo - monkeypatch.setattr("dify_trace_langfuse.langfuse_trace.DifyCoreRepositoryFactory", mock_factory) - monkeypatch.setattr( - "dify_trace_langfuse.langfuse_trace.db", - SimpleNamespace(engine=sqlite3_session.get_bind(), session=sqlite3_session), - ) - monkeypatch.setattr(trace_instance, "get_service_account_with_tenant", lambda app_id: MagicMock()) - - trace_instance.add_trace = MagicMock() - trace_instance.add_generation = MagicMock() - - with caplog.at_level(logging.ERROR): - trace_instance.workflow_trace(trace_info) - - assert "Failed to extract usage" in caplog.text - trace_instance.add_generation.assert_called_once() diff --git a/api/providers/trace/trace-langfuse/tests/unit_tests/test_config_entity.py b/api/providers/trace/trace-langfuse/tests/unit_tests/src/dify_trace_langfuse/test_config.py similarity index 100% rename from api/providers/trace/trace-langfuse/tests/unit_tests/test_config_entity.py rename to api/providers/trace/trace-langfuse/tests/unit_tests/src/dify_trace_langfuse/test_config.py diff --git a/api/providers/trace/trace-langfuse/tests/unit_tests/src/dify_trace_langfuse/test_langfuse_trace.py b/api/providers/trace/trace-langfuse/tests/unit_tests/src/dify_trace_langfuse/test_langfuse_trace.py new file mode 100644 index 00000000000..bf13dad521e --- /dev/null +++ b/api/providers/trace/trace-langfuse/tests/unit_tests/src/dify_trace_langfuse/test_langfuse_trace.py @@ -0,0 +1,911 @@ +import collections +import logging +from datetime import UTC, datetime, timedelta +from types import SimpleNamespace +from typing import cast, override +from unittest.mock import MagicMock, patch + +import pytest +from dify_trace_langfuse.config import LangfuseConfig +from dify_trace_langfuse.entities.langfuse_trace_entity import ( + LangfuseGeneration, + LangfuseSpan, + LangfuseTrace, + LevelEnum, + UnitEnum, +) +from dify_trace_langfuse.langfuse_trace import LangFuseDataTrace +from sqlalchemy.orm import Session + +from core.ops.entities.trace_entity import ( + DatasetRetrievalTraceInfo, + GenerateNameTraceInfo, + MessageTraceInfo, + ModerationTraceInfo, + SuggestedQuestionTraceInfo, + ToolTraceInfo, + TraceTaskName, + WorkflowTraceInfo, +) +from graphon.enums import BuiltinNodeTypes +from models import EndUser +from models.enums import EndUserType, MessageStatus +from tests.helpers.unit.core.ops.trace_fixtures import message_trace_info, tool_trace_info, workflow_trace_info + +"""Unit tests for Langfuse trace translation with real SQLite-backed lookups.""" + + +def _dt() -> datetime: + return datetime(2024, 1, 1, 0, 0, 0, tzinfo=UTC) + + +class TestLangfuseTraceCases: + @staticmethod + @pytest.fixture + def langfuse_config(): + return LangfuseConfig(public_key="pk-123", secret_key="sk-123", host="https://cloud.langfuse.com") + + @staticmethod + @pytest.fixture + def trace_instance(langfuse_config, monkeypatch: pytest.MonkeyPatch): + # Mock Langfuse client to avoid network calls + mock_client = MagicMock() + monkeypatch.setattr("dify_trace_langfuse.langfuse_trace.LangfuseAPI", lambda **kwargs: mock_client) + + instance = LangFuseDataTrace(langfuse_config) + return instance + + @staticmethod + def test_init(langfuse_config, monkeypatch: pytest.MonkeyPatch): + mock_langfuse = MagicMock() + monkeypatch.setattr("dify_trace_langfuse.langfuse_trace.LangfuseAPI", mock_langfuse) + monkeypatch.setenv("FILES_URL", "http://test.url") + + instance = LangFuseDataTrace(langfuse_config) + + mock_langfuse.assert_called_once() + kwargs = mock_langfuse.call_args.kwargs + assert kwargs["username"] == langfuse_config.public_key + assert kwargs["password"] == langfuse_config.secret_key + assert kwargs["base_url"] == langfuse_config.host + assert instance.file_base_url == "http://test.url" + + @staticmethod + def test_same_public_key_uses_independent_api_clients(monkeypatch: pytest.MonkeyPatch): + clients = [MagicMock(), MagicMock()] + create_client = MagicMock(side_effect=clients) + monkeypatch.setattr("dify_trace_langfuse.langfuse_trace.LangfuseAPI", create_client) + + first = LangFuseDataTrace( + LangfuseConfig(public_key="shared", secret_key="secret-a", host="https://tenant-a.example") + ) + second = LangFuseDataTrace( + LangfuseConfig(public_key="shared", secret_key="secret-b", host="https://tenant-b.example") + ) + + assert first.langfuse_client is clients[0] + assert second.langfuse_client is clients[1] + assert [ + (call.kwargs["base_url"], call.kwargs["username"], call.kwargs["password"]) + for call in create_client.call_args_list + ] == [ + ("https://tenant-a.example", "shared", "secret-a"), + ("https://tenant-b.example", "shared", "secret-b"), + ] + + @staticmethod + def test_close_is_idempotent(langfuse_config, monkeypatch: pytest.MonkeyPatch): + http_client = MagicMock() + monkeypatch.setattr("dify_trace_langfuse.langfuse_trace.httpx.Client", lambda **kwargs: http_client) + monkeypatch.setattr("dify_trace_langfuse.langfuse_trace.LangfuseAPI", MagicMock()) + + instance = LangFuseDataTrace(langfuse_config) + instance.close() + instance.close() + + http_client.close.assert_called_once() + + @staticmethod + def test_trace_dispatch(trace_instance, monkeypatch: pytest.MonkeyPatch): + methods = [ + "workflow_trace", + "message_trace", + "moderation_trace", + "suggested_question_trace", + "dataset_retrieval_trace", + "tool_trace", + "generate_name_trace", + ] + mocks = {method: MagicMock() for method in methods} + for method, m in mocks.items(): + monkeypatch.setattr(trace_instance, method, m) + + # WorkflowTraceInfo + info = workflow_trace_info() + trace_instance.trace(info) + mocks["workflow_trace"].assert_called_once_with(info) + + # MessageTraceInfo + info = message_trace_info() + trace_instance.trace(info) + mocks["message_trace"].assert_called_once_with(info) + + # ModerationTraceInfo + info = ModerationTraceInfo(metadata={}, flagged=False, action="allow", preset_response="", query="hello") + trace_instance.trace(info) + mocks["moderation_trace"].assert_called_once_with(info) + + # SuggestedQuestionTraceInfo + info = SuggestedQuestionTraceInfo(metadata={}, total_tokens=0, suggested_question=[], level="info") + trace_instance.trace(info) + mocks["suggested_question_trace"].assert_called_once_with(info) + + # DatasetRetrievalTraceInfo + info = DatasetRetrievalTraceInfo(metadata={}) + trace_instance.trace(info) + mocks["dataset_retrieval_trace"].assert_called_once_with(info) + + # ToolTraceInfo + info = tool_trace_info() + trace_instance.trace(info) + mocks["tool_trace"].assert_called_once_with(info) + + # GenerateNameTraceInfo + info = GenerateNameTraceInfo(tenant_id="tenant-1", metadata={}) + trace_instance.trace(info) + mocks["generate_name_trace"].assert_called_once_with(info) + + @staticmethod + @pytest.mark.parametrize("sqlite3_session", [()], indirect=True) + def test_workflow_trace_with_message_id( + trace_instance, monkeypatch: pytest.MonkeyPatch, sqlite3_session: Session + ) -> None: + # Setup trace info + trace_info = WorkflowTraceInfo( + workflow_id="wf-1", + tenant_id="tenant-1", + workflow_run_id="run-1", + workflow_run_elapsed_time=1.0, + workflow_run_status="succeeded", + workflow_run_inputs={"input": "hi"}, + workflow_run_outputs={"output": "hello"}, + workflow_run_version="1.0", + message_id="msg-1", + conversation_id="conv-1", + total_tokens=100, + file_list=[], + query="hi", + start_time=_dt(), + end_time=_dt() + timedelta(seconds=1), + trace_id="trace-1", + metadata={"app_id": "app-1", "user_id": "user-1"}, + workflow_app_log_id="log-1", + error="", + ) + + monkeypatch.setattr( + "dify_trace_langfuse.langfuse_trace.db", + SimpleNamespace(engine=sqlite3_session.get_bind(), session=sqlite3_session), + ) + + # Mock node executions + node_llm = MagicMock() + node_llm.id = "node-llm" + node_llm.title = "LLM Node" + node_llm.node_type = BuiltinNodeTypes.LLM + node_llm.status = "succeeded" + node_llm.process_data = { + "model_mode": "chat", + "model_name": "gpt-4", + "model_provider": "openai", + "usage": {"prompt_tokens": 10, "completion_tokens": 20}, + } + node_llm.inputs = {"prompts": "p"} + node_llm.outputs = {"text": "t"} + node_llm.created_at = _dt() + node_llm.elapsed_time = 0.5 + node_llm.metadata = {"foo": "bar"} + + node_other = MagicMock() + node_other.id = "node-other" + node_other.title = "Other Node" + node_other.node_type = BuiltinNodeTypes.CODE + node_other.status = "failed" + node_other.process_data = None + node_other.inputs = {"code": "print"} + node_other.outputs = {"result": "ok"} + node_other.created_at = None # Trigger datetime.now() branch + node_other.elapsed_time = 0.2 + node_other.metadata = None + + repo = MagicMock() + repo.get_by_workflow_execution.return_value = [node_llm, node_other] + + mock_factory = MagicMock() + mock_factory.create_workflow_node_execution_repository.return_value = repo + monkeypatch.setattr("dify_trace_langfuse.langfuse_trace.DifyCoreRepositoryFactory", mock_factory) + + monkeypatch.setattr(trace_instance, "get_service_account_with_tenant", lambda app_id: MagicMock()) + + # Track calls to add_trace, add_span, add_generation + trace_instance.add_trace = MagicMock() + trace_instance.add_span = MagicMock() + trace_instance.add_generation = MagicMock() + + trace_instance.workflow_trace(trace_info) + + # Verify add_trace (Workflow Level) + trace_instance.add_trace.assert_called_once() + trace_data = trace_instance.add_trace.call_args[1]["langfuse_trace_data"] + assert trace_data.id == "trace-1" + assert trace_data.name == TraceTaskName.MESSAGE_TRACE + assert "message" in trace_data.tags + assert "workflow" in trace_data.tags + + # Verify add_span (Workflow Run Span) + assert trace_instance.add_span.call_count >= 1 + # First span should be workflow run span because message_id is present + workflow_span = trace_instance.add_span.call_args_list[0][1]["langfuse_span_data"] + assert workflow_span.id == "run-1" + assert workflow_span.name == TraceTaskName.WORKFLOW_TRACE + + # Verify Generation for LLM node + trace_instance.add_generation.assert_called_once() + gen_data = trace_instance.add_generation.call_args[1]["langfuse_generation_data"] + assert gen_data.id == "node-llm" + assert gen_data.usage.input == 10 + assert gen_data.usage.output == 20 + + # Verify normal span for Other node + # Second add_span call + other_span = trace_instance.add_span.call_args_list[1][1]["langfuse_span_data"] + assert other_span.id == "node-other" + assert other_span.level == LevelEnum.ERROR + + @staticmethod + @pytest.mark.parametrize("sqlite3_session", [()], indirect=True) + def test_workflow_trace_no_message_id( + trace_instance, monkeypatch: pytest.MonkeyPatch, sqlite3_session: Session + ) -> None: + trace_info = WorkflowTraceInfo( + workflow_id="wf-1", + tenant_id="tenant-1", + workflow_run_id="run-1", + workflow_run_elapsed_time=1.0, + workflow_run_status="succeeded", + workflow_run_inputs={}, + workflow_run_outputs={}, + workflow_run_version="1.0", + total_tokens=0, + file_list=[], + query="", + message_id=None, + conversation_id="conv-1", + start_time=_dt(), + end_time=_dt(), + trace_id=None, # Should fallback to workflow_run_id + metadata={"app_id": "app-1"}, + workflow_app_log_id="log-1", + error="", + ) + + monkeypatch.setattr( + "dify_trace_langfuse.langfuse_trace.db", + SimpleNamespace(engine=sqlite3_session.get_bind(), session=sqlite3_session), + ) + repo = MagicMock() + repo.get_by_workflow_execution.return_value = [] + mock_factory = MagicMock() + mock_factory.create_workflow_node_execution_repository.return_value = repo + monkeypatch.setattr("dify_trace_langfuse.langfuse_trace.DifyCoreRepositoryFactory", mock_factory) + monkeypatch.setattr(trace_instance, "get_service_account_with_tenant", lambda app_id: MagicMock()) + + trace_instance.add_trace = MagicMock() + trace_instance.workflow_trace(trace_info) + + trace_instance.add_trace.assert_called_once() + trace_data = trace_instance.add_trace.call_args[1]["langfuse_trace_data"] + assert trace_data.id == "run-1" + assert trace_data.name == TraceTaskName.WORKFLOW_TRACE + + @staticmethod + @pytest.mark.parametrize("sqlite3_session", [()], indirect=True) + def test_workflow_trace_missing_app_id( + trace_instance, monkeypatch: pytest.MonkeyPatch, sqlite3_session: Session + ) -> None: + trace_info = WorkflowTraceInfo( + workflow_id="wf-1", + tenant_id="tenant-1", + workflow_run_id="run-1", + workflow_run_elapsed_time=1.0, + workflow_run_status="succeeded", + workflow_run_inputs={}, + workflow_run_outputs={}, + workflow_run_version="1.0", + total_tokens=0, + file_list=[], + query="", + message_id=None, + conversation_id="conv-1", + start_time=_dt(), + end_time=_dt(), + metadata={}, # Missing app_id + workflow_app_log_id="log-1", + error="", + ) + monkeypatch.setattr( + "dify_trace_langfuse.langfuse_trace.db", + SimpleNamespace(engine=sqlite3_session.get_bind(), session=sqlite3_session), + ) + + with pytest.raises(ValueError, match="No app_id found in trace_info metadata"): + trace_instance.workflow_trace(trace_info) + + @staticmethod + def test_message_trace_basic(trace_instance, monkeypatch: pytest.MonkeyPatch): + message_data = MagicMock() + message_data.id = "msg-1" + message_data.from_account_id = "acc-1" + message_data.from_end_user_id = None + message_data.provider_response_latency = 0.5 + message_data.conversation_id = "conv-1" + message_data.total_price = 0.01 + message_data.model_id = "gpt-4" + message_data.answer = "hello" + message_data.status = MessageStatus.NORMAL + message_data.error = None + + trace_info = MessageTraceInfo( + message_id="msg-1", + message_data=message_data, + inputs={"query": "hi"}, + outputs={"answer": "hello"}, + message_tokens=10, + answer_tokens=20, + total_tokens=30, + start_time=_dt(), + end_time=_dt() + timedelta(seconds=1), + trace_id="trace-1", + metadata={"foo": "bar"}, + conversation_mode="chat", + conversation_model="gpt-4", + file_list=[], + error=None, + ) + + trace_instance.add_trace = MagicMock() + trace_instance.add_generation = MagicMock() + + trace_instance.message_trace(trace_info) + + trace_instance.add_trace.assert_called_once() + trace_instance.add_generation.assert_called_once() + + gen_data = trace_instance.add_generation.call_args[0][0] + assert gen_data.name == "llm" + assert gen_data.usage.total == 30 + + @staticmethod + @pytest.mark.parametrize("sqlite3_session", [(EndUser,)], indirect=True) + def test_message_trace_with_end_user( + trace_instance, monkeypatch: pytest.MonkeyPatch, sqlite3_session: Session + ) -> None: + message_data = MagicMock() + message_data.id = "msg-1" + message_data.from_account_id = "acc-1" + message_data.from_end_user_id = "end-user-1" + message_data.conversation_id = "conv-1" + message_data.status = MessageStatus.NORMAL + message_data.model_id = "gpt-4" + message_data.error = "" + message_data.answer = "hello" + message_data.total_price = 0.0 + message_data.provider_response_latency = 0.1 + + trace_info = MessageTraceInfo( + message_id="msg-1", + message_data=message_data, + inputs={}, + outputs={}, + message_tokens=0, + answer_tokens=0, + total_tokens=0, + start_time=_dt(), + end_time=_dt(), + metadata={}, + conversation_mode="chat", + conversation_model="gpt-4", + file_list=[], + error=None, + ) + + end_user = EndUser( + id="end-user-1", + tenant_id="tenant-1", + app_id="app-1", + type=EndUserType.BROWSER, + session_id="session-id-123", + ) + engine = sqlite3_session.get_bind() + with Session(engine) as write_session: + write_session.add(end_user) + write_session.commit() + + with Session(engine) as read_session: + monkeypatch.setattr( + "dify_trace_langfuse.langfuse_trace.db", + SimpleNamespace(engine=engine, session=read_session), + ) + + trace_instance.add_trace = MagicMock() + trace_instance.add_generation = MagicMock() + + trace_instance.message_trace(trace_info) + + trace_data = trace_instance.add_trace.call_args[1]["langfuse_trace_data"] + assert trace_data.user_id == "session-id-123" + assert trace_data.metadata["user_id"] == "session-id-123" + + @staticmethod + def test_message_trace_none_data(trace_instance): + trace_info = SimpleNamespace(message_data=None, file_list=[], metadata={}) + trace_instance.add_trace = MagicMock() + trace_instance.message_trace(trace_info) + trace_instance.add_trace.assert_not_called() + + @staticmethod + def test_moderation_trace(trace_instance): + message_data = MagicMock() + message_data.created_at = _dt() + + trace_info = ModerationTraceInfo( + message_id="msg-1", + message_data=message_data, + inputs={"q": "hi"}, + action="stop", + flagged=True, + preset_response="blocked", + start_time=None, + end_time=None, + metadata={"foo": "bar"}, + trace_id="trace-1", + query="hi", + ) + + trace_instance.add_span = MagicMock() + trace_instance.moderation_trace(trace_info) + + trace_instance.add_span.assert_called_once() + span_data = trace_instance.add_span.call_args[1]["langfuse_span_data"] + assert span_data.name == TraceTaskName.MODERATION_TRACE + assert span_data.output["flagged"] is True + + @staticmethod + def test_suggested_question_trace(trace_instance): + message_data = MagicMock() + message_data.status = MessageStatus.NORMAL + message_data.error = None + + trace_info = SuggestedQuestionTraceInfo( + message_id="msg-1", + message_data=message_data, + inputs="hi", + suggested_question=["q1"], + total_tokens=10, + level="info", + start_time=_dt(), + end_time=_dt(), + metadata={}, + trace_id="trace-1", + ) + + trace_instance.add_generation = MagicMock() + trace_instance.suggested_question_trace(trace_info) + + trace_instance.add_generation.assert_called_once() + gen_data = trace_instance.add_generation.call_args[1]["langfuse_generation_data"] + assert gen_data.name == TraceTaskName.SUGGESTED_QUESTION_TRACE + assert gen_data.usage.unit == UnitEnum.CHARACTERS + + @staticmethod + def test_dataset_retrieval_trace(trace_instance): + message_data = MagicMock() + message_data.created_at = _dt() + message_data.updated_at = _dt() + + trace_info = DatasetRetrievalTraceInfo( + message_id="msg-1", + message_data=message_data, + inputs="query", + documents=[{"id": "doc1"}], + start_time=None, + end_time=None, + metadata={}, + trace_id="trace-1", + ) + + trace_instance.add_span = MagicMock() + trace_instance.dataset_retrieval_trace(trace_info) + + trace_instance.add_span.assert_called_once() + span_data = trace_instance.add_span.call_args[1]["langfuse_span_data"] + assert span_data.name == TraceTaskName.DATASET_RETRIEVAL_TRACE + assert span_data.output["documents"] == [{"id": "doc1"}] + + @staticmethod + def test_tool_trace(trace_instance): + trace_info = ToolTraceInfo( + message_id="msg-1", + message_data=MagicMock(), + inputs={}, + outputs={}, + tool_name="my_tool", + tool_inputs={"a": 1}, + tool_outputs="result_string", + time_cost=0.1, + start_time=_dt(), + end_time=_dt(), + metadata={}, + trace_id="trace-1", + tool_config={}, + tool_parameters={}, + error="some error", + ) + + trace_instance.add_span = MagicMock() + trace_instance.tool_trace(trace_info) + + trace_instance.add_span.assert_called_once() + span_data = trace_instance.add_span.call_args[1]["langfuse_span_data"] + assert span_data.name == "my_tool" + assert span_data.level == LevelEnum.ERROR + + @staticmethod + def test_generate_name_trace(trace_instance): + trace_info = GenerateNameTraceInfo( + inputs={"q": "hi"}, + outputs={"name": "new"}, + tenant_id="tenant-1", + conversation_id="conv-1", + start_time=_dt(), + end_time=_dt(), + metadata={"m": 1}, + ) + + trace_instance.add_trace = MagicMock() + trace_instance.add_span = MagicMock() + + trace_instance.generate_name_trace(trace_info) + + trace_instance.add_trace.assert_called_once() + trace_instance.add_span.assert_called_once() + + trace_data = trace_instance.add_trace.call_args[1]["langfuse_trace_data"] + assert trace_data.name == TraceTaskName.GENERATE_NAME_TRACE + assert trace_data.user_id == "tenant-1" + + span_data = trace_instance.add_span.call_args[1]["langfuse_span_data"] + assert span_data.trace_id == "conv-1" + + @staticmethod + def test_add_trace_success(trace_instance): + data = LangfuseTrace(id="t1", name="trace") + trace_instance.add_trace(data) + trace_instance.langfuse_client.ingestion.batch.assert_called_once() + + @staticmethod + def test_add_trace_error(trace_instance): + trace_instance.langfuse_client.ingestion.batch.side_effect = Exception("error") + data = LangfuseTrace(id="t1", name="trace") + with pytest.raises(ValueError, match="LangFuse Failed to create trace: error"): + trace_instance.add_trace(data) + + @staticmethod + def test_add_span_success(trace_instance): + data = LangfuseSpan(id="s1", name="span", trace_id="t1") + trace_instance.add_span(data) + trace_instance.langfuse_client.ingestion.batch.assert_called_once() + + @staticmethod + def test_add_span_error(trace_instance): + trace_instance.langfuse_client.ingestion.batch.side_effect = Exception("error") + data = LangfuseSpan(id="s1", name="span", trace_id="t1") + with pytest.raises(ValueError, match="LangFuse Failed to create span: error"): + trace_instance.add_span(data) + + @staticmethod + def test_update_span(trace_instance): + span = MagicMock() + data = LangfuseSpan(id="s1", name="span", trace_id="t1") + trace_instance.update_span(span, data) + span.end.assert_called_once() + + @staticmethod + def test_add_generation_success(trace_instance): + data = LangfuseGeneration(id="g1", name="gen", trace_id="t1") + trace_instance.add_generation(data) + trace_instance.langfuse_client.ingestion.batch.assert_called_once() + + @staticmethod + def test_add_generation_error(trace_instance): + trace_instance.langfuse_client.ingestion.batch.side_effect = Exception("error") + data = LangfuseGeneration(id="g1", name="gen", trace_id="t1") + with pytest.raises(ValueError, match="LangFuse Failed to create generation: error"): + trace_instance.add_generation(data) + + @staticmethod + def test_update_generation(trace_instance): + gen = MagicMock() + data = LangfuseGeneration(id="g1", name="gen", trace_id="t1") + trace_instance.update_generation(gen, data) + gen.end.assert_called_once() + + @staticmethod + def test_api_check_success(trace_instance): + trace_instance.langfuse_client.projects.get.return_value = MagicMock(data=[MagicMock()]) + assert trace_instance.api_check() is True + + @staticmethod + def test_api_check_rejects_empty_project_list(trace_instance): + trace_instance.langfuse_client.projects.get.return_value = MagicMock(data=[]) + with pytest.raises(ValueError, match="no project found for the provided credentials"): + trace_instance.api_check() + + @staticmethod + def test_api_check_error(trace_instance): + trace_instance.langfuse_client.projects.get.side_effect = Exception("fail") + with pytest.raises(ValueError, match="LangFuse API check failed: fail"): + trace_instance.api_check() + + @staticmethod + def test_get_project_key_success(trace_instance): + mock_data = MagicMock() + mock_data.id = "proj-1" + trace_instance.langfuse_client.projects.get.return_value = MagicMock(data=[mock_data]) + assert trace_instance.get_project_key() == "proj-1" + + @staticmethod + def test_get_project_key_error(trace_instance): + trace_instance.langfuse_client.projects.get.side_effect = Exception("fail") + with pytest.raises(ValueError, match="LangFuse get project key failed: fail"): + trace_instance.get_project_key() + + @staticmethod + def test_moderation_trace_none(trace_instance): + trace_info = ModerationTraceInfo( + message_id="m", + message_data=None, + inputs={}, + action="s", + flagged=False, + preset_response="", + query="", + metadata={}, + ) + trace_instance.add_span = MagicMock() + trace_instance.moderation_trace(trace_info) + trace_instance.add_span.assert_not_called() + + @staticmethod + def test_suggested_question_trace_none(trace_instance): + trace_info = SuggestedQuestionTraceInfo( + message_id="m", message_data=None, inputs={}, suggested_question=[], total_tokens=0, level="i", metadata={} + ) + trace_instance.add_generation = MagicMock() + trace_instance.suggested_question_trace(trace_info) + trace_instance.add_generation.assert_not_called() + + @staticmethod + def test_dataset_retrieval_trace_none(trace_instance): + trace_info = DatasetRetrievalTraceInfo(message_id="m", message_data=None, inputs={}, documents=[], metadata={}) + trace_instance.add_span = MagicMock() + trace_instance.dataset_retrieval_trace(trace_info) + trace_instance.add_span.assert_not_called() + + @staticmethod + def test_langfuse_trace_entity_with_list_dict_input(): + # To cover lines 29-31 in langfuse_trace_entity.py + # We need to mock replace_text_with_content or just check if it works + # Actually replace_text_with_content is imported from core.ops.utils + data = LangfuseTrace(id="t1", name="n", input=[{"text": "hello"}]) + assert isinstance(data.input, list) + assert data.input[0]["content"] == "hello" + + @staticmethod + @pytest.mark.parametrize("sqlite3_session", [()], indirect=True) + def test_workflow_trace_handles_usage_extraction_error( + trace_instance, + monkeypatch: pytest.MonkeyPatch, + caplog: pytest.LogCaptureFixture, + sqlite3_session: Session, + ): + # Setup trace info to trigger LLM node usage extraction + trace_info = WorkflowTraceInfo( + workflow_id="wf-1", + tenant_id="t", + workflow_run_id="r", + workflow_run_elapsed_time=1.0, + workflow_run_status="s", + workflow_run_inputs={}, + workflow_run_outputs={}, + workflow_run_version="1", + total_tokens=0, + file_list=[], + query="", + message_id=None, + conversation_id="c", + start_time=_dt(), + end_time=_dt(), + metadata={"app_id": "app-1"}, + workflow_app_log_id="l", + error="", + ) + + node = MagicMock() + node.id = "n1" + node.title = "LLM Node" + node.node_type = BuiltinNodeTypes.LLM + node.status = "succeeded" + + class BadDict(collections.UserDict): + @override + def get(self, key, default=None): + if key == "usage": + raise Exception("Usage extraction failed") + return super().get(key, default) + + node.process_data = BadDict({"model_mode": "chat", "model_name": "gpt-4", "usage": True, "prompts": ["p"]}) + node.created_at = _dt() + node.elapsed_time = 0.1 + node.metadata = {} + node.outputs = {} + + repo = MagicMock() + repo.get_by_workflow_execution.return_value = [node] + mock_factory = MagicMock() + mock_factory.create_workflow_node_execution_repository.return_value = repo + monkeypatch.setattr("dify_trace_langfuse.langfuse_trace.DifyCoreRepositoryFactory", mock_factory) + monkeypatch.setattr( + "dify_trace_langfuse.langfuse_trace.db", + SimpleNamespace(engine=sqlite3_session.get_bind(), session=sqlite3_session), + ) + monkeypatch.setattr(trace_instance, "get_service_account_with_tenant", lambda app_id: MagicMock()) + + trace_instance.add_trace = MagicMock() + trace_instance.add_generation = MagicMock() + + with caplog.at_level(logging.ERROR): + trace_instance.workflow_trace(trace_info) + + assert "Failed to extract usage" in caplog.text + trace_instance.add_generation.assert_called_once() + + +"""Tests for Langfuse TTFT reporting support.""" + + +def _create_trace_instance() -> LangFuseDataTrace: + with patch("dify_trace_langfuse.langfuse_trace.LangfuseAPI", autospec=True): + return LangFuseDataTrace( + LangfuseConfig( + public_key="public-key", + secret_key="secret-key", + host="https://cloud.langfuse.com", + ) + ) + + +class TestLangfuseTraceCasesUnitTests: + class TestLangFuseDataTraceCompletionStartTime: + def test_message_trace_reports_completion_start_time(self): + trace = _create_trace_instance() + start_time = datetime(2026, 3, 11, 13, 0, 0) + trace_info = MessageTraceInfo( + trace_id="trace-123", + message_id="message-123", + message_data=SimpleNamespace( + id="message-123", + from_account_id="account-1", + from_end_user_id=None, + conversation_id="conversation-1", + model_id="gpt-4o-mini", + answer="hi there", + status="normal", + error="", + total_price=0.12, + provider_response_latency=3.5, + ), + conversation_model="chat", + message_tokens=10, + answer_tokens=20, + total_tokens=30, + error="", + inputs="hello", + outputs="hi there", + file_list=[], + start_time=start_time, + end_time=start_time + timedelta(seconds=3.5), + metadata={}, + message_file_data=None, + conversation_mode="chat", + gen_ai_server_time_to_first_token=1.2, + llm_streaming_time_to_generate=2.3, + is_streaming_request=True, + ) + + with patch.object(trace, "add_trace"), patch.object(trace, "add_generation") as add_generation: + trace.message_trace(trace_info) + + generation = add_generation.call_args.args[0] + assert generation.completion_start_time == start_time + timedelta(seconds=1.2) + + def test_workflow_trace_reports_completion_start_time_from_llm_usage(self): + trace = _create_trace_instance() + start_time = datetime(2026, 3, 11, 13, 0, 0) + node_execution = SimpleNamespace( + id="node-exec-1", + title="Chat LLM", + node_type=BuiltinNodeTypes.LLM, + status="succeeded", + process_data={ + "model_mode": "chat", + "model_name": "gpt-4o-mini", + "usage": { + "prompt_tokens": 10, + "completion_tokens": 20, + "time_to_first_token": 1.2, + }, + }, + inputs={"question": "hello"}, + outputs={"text": "hi there"}, + created_at=start_time, + elapsed_time=3.5, + metadata={}, + ) + trace_info = WorkflowTraceInfo( + trace_id="trace-123", + workflow_data={}, + conversation_id=None, + workflow_app_log_id=None, + workflow_id="workflow-1", + tenant_id="tenant-1", + workflow_run_id="workflow-run-1", + workflow_run_elapsed_time=3.5, + workflow_run_status="succeeded", + workflow_run_inputs={"question": "hello"}, + workflow_run_outputs={"answer": "hi there"}, + workflow_run_version="1", + error="", + total_tokens=30, + file_list=[], + query="hello", + metadata={"app_id": "app-1", "user_id": "user-1"}, + start_time=start_time, + end_time=start_time + timedelta(seconds=3.5), + ) + repository = MagicMock() + repository.get_by_workflow_execution.return_value = [node_execution] + + with ( + patch.object(trace, "add_trace"), + patch.object(trace, "add_span"), + patch.object(trace, "add_generation") as add_generation, + patch.object(trace, "get_service_account_with_tenant", return_value=MagicMock()), + patch("dify_trace_langfuse.langfuse_trace.db", MagicMock()), + patch( + "dify_trace_langfuse.langfuse_trace.DifyCoreRepositoryFactory.create_workflow_node_execution_repository", + return_value=repository, + ), + ): + trace.workflow_trace(trace_info) + + generation = add_generation.call_args.kwargs["langfuse_generation_data"] + assert generation.completion_start_time == start_time + timedelta(seconds=1.2) + + def test_ignores_invalid_ttft_values(self): + trace = _create_trace_instance() + start_time = datetime(2026, 3, 11, 13, 0, 0) + + assert trace._get_completion_start_time(start_time, None) is None + assert trace._get_completion_start_time(start_time, -1) is None + assert trace._get_completion_start_time(start_time, cast(float | int | None, "invalid")) is None diff --git a/api/providers/trace/trace-langfuse/tests/unit_tests/test_langfuse_trace.py b/api/providers/trace/trace-langfuse/tests/unit_tests/test_langfuse_trace.py deleted file mode 100644 index 05dd5edd545..00000000000 --- a/api/providers/trace/trace-langfuse/tests/unit_tests/test_langfuse_trace.py +++ /dev/null @@ -1,138 +0,0 @@ -"""Tests for Langfuse TTFT reporting support.""" - -from datetime import datetime, timedelta -from types import SimpleNamespace -from typing import cast -from unittest.mock import MagicMock, patch - -from dify_trace_langfuse.config import LangfuseConfig -from dify_trace_langfuse.langfuse_trace import LangFuseDataTrace - -from core.ops.entities.trace_entity import MessageTraceInfo, WorkflowTraceInfo -from graphon.enums import BuiltinNodeTypes - - -def _create_trace_instance() -> LangFuseDataTrace: - with patch("dify_trace_langfuse.langfuse_trace.LangfuseAPI", autospec=True): - return LangFuseDataTrace( - LangfuseConfig( - public_key="public-key", - secret_key="secret-key", - host="https://cloud.langfuse.com", - ) - ) - - -class TestLangFuseDataTraceCompletionStartTime: - def test_message_trace_reports_completion_start_time(self): - trace = _create_trace_instance() - start_time = datetime(2026, 3, 11, 13, 0, 0) - trace_info = MessageTraceInfo( - trace_id="trace-123", - message_id="message-123", - message_data=SimpleNamespace( - id="message-123", - from_account_id="account-1", - from_end_user_id=None, - conversation_id="conversation-1", - model_id="gpt-4o-mini", - answer="hi there", - status="normal", - error="", - total_price=0.12, - provider_response_latency=3.5, - ), - conversation_model="chat", - message_tokens=10, - answer_tokens=20, - total_tokens=30, - error="", - inputs="hello", - outputs="hi there", - file_list=[], - start_time=start_time, - end_time=start_time + timedelta(seconds=3.5), - metadata={}, - message_file_data=None, - conversation_mode="chat", - gen_ai_server_time_to_first_token=1.2, - llm_streaming_time_to_generate=2.3, - is_streaming_request=True, - ) - - with patch.object(trace, "add_trace"), patch.object(trace, "add_generation") as add_generation: - trace.message_trace(trace_info) - - generation = add_generation.call_args.args[0] - assert generation.completion_start_time == start_time + timedelta(seconds=1.2) - - def test_workflow_trace_reports_completion_start_time_from_llm_usage(self): - trace = _create_trace_instance() - start_time = datetime(2026, 3, 11, 13, 0, 0) - node_execution = SimpleNamespace( - id="node-exec-1", - title="Chat LLM", - node_type=BuiltinNodeTypes.LLM, - status="succeeded", - process_data={ - "model_mode": "chat", - "model_name": "gpt-4o-mini", - "usage": { - "prompt_tokens": 10, - "completion_tokens": 20, - "time_to_first_token": 1.2, - }, - }, - inputs={"question": "hello"}, - outputs={"text": "hi there"}, - created_at=start_time, - elapsed_time=3.5, - metadata={}, - ) - trace_info = WorkflowTraceInfo( - trace_id="trace-123", - workflow_data={}, - conversation_id=None, - workflow_app_log_id=None, - workflow_id="workflow-1", - tenant_id="tenant-1", - workflow_run_id="workflow-run-1", - workflow_run_elapsed_time=3.5, - workflow_run_status="succeeded", - workflow_run_inputs={"question": "hello"}, - workflow_run_outputs={"answer": "hi there"}, - workflow_run_version="1", - error="", - total_tokens=30, - file_list=[], - query="hello", - metadata={"app_id": "app-1", "user_id": "user-1"}, - start_time=start_time, - end_time=start_time + timedelta(seconds=3.5), - ) - repository = MagicMock() - repository.get_by_workflow_execution.return_value = [node_execution] - - with ( - patch.object(trace, "add_trace"), - patch.object(trace, "add_span"), - patch.object(trace, "add_generation") as add_generation, - patch.object(trace, "get_service_account_with_tenant", return_value=MagicMock()), - patch("dify_trace_langfuse.langfuse_trace.db", MagicMock()), - patch( - "dify_trace_langfuse.langfuse_trace.DifyCoreRepositoryFactory.create_workflow_node_execution_repository", - return_value=repository, - ), - ): - trace.workflow_trace(trace_info) - - generation = add_generation.call_args.kwargs["langfuse_generation_data"] - assert generation.completion_start_time == start_time + timedelta(seconds=1.2) - - def test_ignores_invalid_ttft_values(self): - trace = _create_trace_instance() - start_time = datetime(2026, 3, 11, 13, 0, 0) - - assert trace._get_completion_start_time(start_time, None) is None - assert trace._get_completion_start_time(start_time, -1) is None - assert trace._get_completion_start_time(start_time, cast(float | int | None, "invalid")) is None diff --git a/api/providers/trace/trace-langsmith/tests/unit_tests/test_config_entity.py b/api/providers/trace/trace-langsmith/tests/unit_tests/src/dify_trace_langsmith/test_config.py similarity index 100% rename from api/providers/trace/trace-langsmith/tests/unit_tests/test_config_entity.py rename to api/providers/trace/trace-langsmith/tests/unit_tests/src/dify_trace_langsmith/test_config.py diff --git a/api/providers/trace/trace-langsmith/tests/unit_tests/langsmith_trace/test_langsmith_trace.py b/api/providers/trace/trace-langsmith/tests/unit_tests/src/dify_trace_langsmith/test_langsmith_trace.py similarity index 99% rename from api/providers/trace/trace-langsmith/tests/unit_tests/langsmith_trace/test_langsmith_trace.py rename to api/providers/trace/trace-langsmith/tests/unit_tests/src/dify_trace_langsmith/test_langsmith_trace.py index 0c0c8b5423a..264ab4d2862 100644 --- a/api/providers/trace/trace-langsmith/tests/unit_tests/langsmith_trace/test_langsmith_trace.py +++ b/api/providers/trace/trace-langsmith/tests/unit_tests/src/dify_trace_langsmith/test_langsmith_trace.py @@ -29,7 +29,7 @@ from core.ops.entities.trace_entity import ( from graphon.enums import BuiltinNodeTypes, WorkflowNodeExecutionMetadataKey from models import EndUser from models.enums import EndUserType -from tests.unit_tests.core.ops.trace_fixtures import message_trace_info, tool_trace_info, workflow_trace_info +from tests.helpers.unit.core.ops.trace_fixtures import message_trace_info, tool_trace_info, workflow_trace_info def _dt() -> datetime: diff --git a/api/providers/trace/trace-mlflow/tests/unit_tests/mlflow_trace/test_mlflow_trace.py b/api/providers/trace/trace-mlflow/tests/unit_tests/src/dify_trace_mlflow/test_mlflow_trace.py similarity index 100% rename from api/providers/trace/trace-mlflow/tests/unit_tests/mlflow_trace/test_mlflow_trace.py rename to api/providers/trace/trace-mlflow/tests/unit_tests/src/dify_trace_mlflow/test_mlflow_trace.py diff --git a/api/providers/trace/trace-opik/tests/unit_tests/opik_trace/test_opik_trace.py b/api/providers/trace/trace-opik/tests/unit_tests/opik_trace/test_opik_trace.py deleted file mode 100644 index 5e8d40f03f7..00000000000 --- a/api/providers/trace/trace-opik/tests/unit_tests/opik_trace/test_opik_trace.py +++ /dev/null @@ -1,717 +0,0 @@ -"""Unit tests for Opik trace translation with SQLite-backed lookups.""" - -import collections -import logging -from datetime import UTC, datetime, timedelta -from types import SimpleNamespace -from typing import override -from unittest.mock import MagicMock - -import pytest -from dify_trace_opik.config import OpikConfig -from dify_trace_opik.opik_trace import OpikDataTrace, prepare_opik_uuid, wrap_dict, wrap_metadata -from sqlalchemy.orm import Session - -from core.ops.entities.trace_entity import ( - DatasetRetrievalTraceInfo, - GenerateNameTraceInfo, - MessageTraceInfo, - ModerationTraceInfo, - SuggestedQuestionTraceInfo, - ToolTraceInfo, - TraceTaskName, - WorkflowTraceInfo, -) -from graphon.enums import BuiltinNodeTypes, WorkflowNodeExecutionMetadataKey -from models import EndUser -from models.enums import EndUserType, MessageStatus -from tests.unit_tests.core.ops.trace_fixtures import message_trace_info, tool_trace_info, workflow_trace_info - - -def _dt() -> datetime: - return datetime(2024, 1, 1, 0, 0, 0, tzinfo=UTC) - - -@pytest.fixture -def opik_config(): - return OpikConfig( - project="test-project", workspace="test-workspace", url="https://cloud.opik.com/api/", api_key="api-key-123" - ) - - -@pytest.fixture -def trace_instance(opik_config, monkeypatch: pytest.MonkeyPatch): - mock_client = MagicMock() - monkeypatch.setattr("dify_trace_opik.opik_trace.Opik", lambda **kwargs: mock_client) - - instance = OpikDataTrace(opik_config) - return instance - - -def test_wrap_dict(): - assert wrap_dict("input", {"a": 1}) == {"a": 1} - assert wrap_dict("input", "hello") == {"input": "hello"} - - -def test_wrap_metadata(): - assert wrap_metadata({"a": 1}, b=2) == {"a": 1, "b": 2, "created_from": "dify"} - - -def test_prepare_opik_uuid(): - # Test with valid datetime and uuid string - dt = datetime(2024, 1, 1) - uuid_str = "b3e8e918-472e-4b69-8051-12502c34fc07" - result = prepare_opik_uuid(dt, uuid_str) - assert result is not None - # We won't test the exact uuid7 value but just that it returns a string id - - # Test with None dt and uuid_str - result = prepare_opik_uuid(None, None) - assert result is not None - - -def test_init(opik_config, monkeypatch: pytest.MonkeyPatch): - mock_opik = MagicMock() - monkeypatch.setattr("dify_trace_opik.opik_trace.Opik", mock_opik) - monkeypatch.setenv("FILES_URL", "http://test.url") - - instance = OpikDataTrace(opik_config) - - mock_opik.assert_called_once_with( - project_name=opik_config.project, - workspace=opik_config.workspace, - host=opik_config.url, - api_key=opik_config.api_key, - ) - assert instance.file_base_url == "http://test.url" - assert instance.project == opik_config.project - - -def test_trace_dispatch(trace_instance, monkeypatch: pytest.MonkeyPatch): - methods = [ - "workflow_trace", - "message_trace", - "moderation_trace", - "suggested_question_trace", - "dataset_retrieval_trace", - "tool_trace", - "generate_name_trace", - ] - mocks = {method: MagicMock() for method in methods} - for method, m in mocks.items(): - monkeypatch.setattr(trace_instance, method, m) - - # WorkflowTraceInfo - info = workflow_trace_info() - trace_instance.trace(info) - mocks["workflow_trace"].assert_called_once_with(info) - - # MessageTraceInfo - info = message_trace_info() - trace_instance.trace(info) - mocks["message_trace"].assert_called_once_with(info) - - # ModerationTraceInfo - info = ModerationTraceInfo(metadata={}, flagged=False, action="allow", preset_response="", query="hello") - trace_instance.trace(info) - mocks["moderation_trace"].assert_called_once_with(info) - - # SuggestedQuestionTraceInfo - info = SuggestedQuestionTraceInfo(metadata={}, total_tokens=0, suggested_question=[], level="info") - trace_instance.trace(info) - mocks["suggested_question_trace"].assert_called_once_with(info) - - # DatasetRetrievalTraceInfo - info = DatasetRetrievalTraceInfo(metadata={}) - trace_instance.trace(info) - mocks["dataset_retrieval_trace"].assert_called_once_with(info) - - # ToolTraceInfo - info = tool_trace_info() - trace_instance.trace(info) - mocks["tool_trace"].assert_called_once_with(info) - - # GenerateNameTraceInfo - info = GenerateNameTraceInfo(tenant_id="tenant-1", metadata={}) - trace_instance.trace(info) - mocks["generate_name_trace"].assert_called_once_with(info) - - -@pytest.mark.parametrize("sqlite3_session", [()], indirect=True) -def test_workflow_trace_with_message_id( - trace_instance, monkeypatch: pytest.MonkeyPatch, sqlite3_session: Session -) -> None: - # Define constants for better readability - workflow_id = "fb05c7cd-6cec-4add-8a84-df03a408b4ce" - workflow_run_id = "33c67568-7a8a-450e-8916-a5f135baeaef" - message_id = "04ec3956-85f3-488a-8539-1017251dc8c6" - conversation_id = "d3d01066-23ae-4830-9ce4-eb5640b42a7e" - trace_id = "bf26d929-6f15-4c2f-9abc-761c217056f3" - workflow_app_log_id = "ca0e018e-edd4-43fb-a05a-ea001ca8ef4b" - llm_node_id = "80d7dfa8-08f4-4ab7-aa37-0ca7d27207e3" - code_node_id = "b9cd9a7b-c534-4aa9-b5da-efd454140900" - - trace_info = WorkflowTraceInfo( - workflow_id=workflow_id, - tenant_id="tenant-1", - workflow_run_id=workflow_run_id, - workflow_run_elapsed_time=1.0, - workflow_run_status="succeeded", - workflow_run_inputs={"input": "hi"}, - workflow_run_outputs={"output": "hello"}, - workflow_run_version="1.0", - message_id=message_id, - conversation_id=conversation_id, - total_tokens=100, - file_list=[], - query="hi", - start_time=_dt(), - end_time=_dt() + timedelta(seconds=1), - trace_id=trace_id, - metadata={"app_id": "app-1", "user_id": "user-1"}, - workflow_app_log_id=workflow_app_log_id, - error="", - ) - - monkeypatch.setattr( - "dify_trace_opik.opik_trace.db", - SimpleNamespace(engine=sqlite3_session.get_bind(), session=sqlite3_session), - ) - - node_llm = MagicMock() - node_llm.id = llm_node_id - node_llm.title = "LLM Node" - node_llm.node_type = BuiltinNodeTypes.LLM - node_llm.status = "succeeded" - node_llm.process_data = { - "model_mode": "chat", - "model_name": "gpt-4", - "model_provider": "openai", - "usage": {"prompt_tokens": 10, "completion_tokens": 20, "total_tokens": 30}, - } - node_llm.inputs = {"prompts": "p"} - node_llm.outputs = {"text": "t"} - node_llm.created_at = _dt() - node_llm.elapsed_time = 0.5 - node_llm.metadata = {"foo": "bar"} - - node_other = MagicMock() - node_other.id = code_node_id - node_other.title = "Other Node" - node_other.node_type = BuiltinNodeTypes.CODE - node_other.status = "failed" - node_other.process_data = None - node_other.inputs = {"code": "print"} - node_other.outputs = {"result": "ok"} - node_other.created_at = None - node_other.elapsed_time = 0.2 - node_other.metadata = {WorkflowNodeExecutionMetadataKey.TOTAL_TOKENS.value: 10} - - repo = MagicMock() - repo.get_by_workflow_execution.return_value = [node_llm, node_other] - - mock_factory = MagicMock() - mock_factory.create_workflow_node_execution_repository.return_value = repo - monkeypatch.setattr("dify_trace_opik.opik_trace.DifyCoreRepositoryFactory", mock_factory) - - monkeypatch.setattr(trace_instance, "get_service_account_with_tenant", lambda app_id: MagicMock()) - - trace_instance.add_trace = MagicMock() - trace_instance.add_span = MagicMock() - - trace_instance.workflow_trace(trace_info) - - trace_instance.add_trace.assert_called_once() - trace_data = trace_instance.add_trace.call_args[1].get("opik_trace_data", trace_instance.add_trace.call_args[0][0]) - assert trace_data["name"] == TraceTaskName.MESSAGE_TRACE - assert "message" in trace_data["tags"] - assert "workflow" in trace_data["tags"] - - assert trace_instance.add_span.call_count >= 1 - - -@pytest.mark.parametrize("sqlite3_session", [()], indirect=True) -def test_workflow_trace_no_message_id( - trace_instance, monkeypatch: pytest.MonkeyPatch, sqlite3_session: Session -) -> None: - # Define constants for better readability - workflow_id = "f0708b36-b1d7-42b3-a876-1d01b7d8f1a3" - workflow_run_id = "d42ec285-c2fd-4248-8866-5c9386b101ac" - conversation_id = "88a17f2e-9436-4472-bab9-4b1601d5af3c" - workflow_app_log_id = "41780d0d-ffba-4220-bc0c-401e4c89cdfb" - - trace_info = WorkflowTraceInfo( - workflow_id=workflow_id, - tenant_id="tenant-1", - workflow_run_id=workflow_run_id, - workflow_run_elapsed_time=1.0, - workflow_run_status="succeeded", - workflow_run_inputs={}, - workflow_run_outputs={}, - workflow_run_version="1.0", - total_tokens=0, - file_list=[], - query="", - message_id=None, - conversation_id=conversation_id, - start_time=_dt(), - end_time=_dt(), - trace_id=None, - metadata={"app_id": "app-1"}, - workflow_app_log_id=workflow_app_log_id, - error="", - ) - - monkeypatch.setattr( - "dify_trace_opik.opik_trace.db", - SimpleNamespace(engine=sqlite3_session.get_bind(), session=sqlite3_session), - ) - repo = MagicMock() - repo.get_by_workflow_execution.return_value = [] - mock_factory = MagicMock() - mock_factory.create_workflow_node_execution_repository.return_value = repo - monkeypatch.setattr("dify_trace_opik.opik_trace.DifyCoreRepositoryFactory", mock_factory) - monkeypatch.setattr(trace_instance, "get_service_account_with_tenant", lambda app_id: MagicMock()) - - trace_instance.add_trace = MagicMock() - trace_instance.workflow_trace(trace_info) - - trace_instance.add_trace.assert_called_once() - - -@pytest.mark.parametrize("sqlite3_session", [()], indirect=True) -def test_workflow_trace_missing_app_id( - trace_instance, monkeypatch: pytest.MonkeyPatch, sqlite3_session: Session -) -> None: - trace_info = WorkflowTraceInfo( - workflow_id="5745f1b8-f8e6-4859-8110-996acb6c8d6a", - tenant_id="tenant-1", - workflow_run_id="46f53304-1659-464b-bee5-116585f0bec8", - workflow_run_elapsed_time=1.0, - workflow_run_status="succeeded", - workflow_run_inputs={}, - workflow_run_outputs={}, - workflow_run_version="1.0", - total_tokens=0, - file_list=[], - query="", - message_id=None, - conversation_id="83f86b89-caef-4de8-a0f9-f164eddae1ea", - start_time=_dt(), - end_time=_dt(), - metadata={}, - workflow_app_log_id="339760b2-4b94-4532-8c81-133a97e4680e", - error="", - ) - monkeypatch.setattr( - "dify_trace_opik.opik_trace.db", - SimpleNamespace(engine=sqlite3_session.get_bind(), session=sqlite3_session), - ) - - with pytest.raises(ValueError, match="No app_id found in trace_info metadata"): - trace_instance.workflow_trace(trace_info) - - -def test_message_trace_basic(trace_instance, monkeypatch: pytest.MonkeyPatch): - # Define constants for better readability - message_data_id = "e3a26712-8cac-4a25-94a4-a3bff21ee3ab" - conversation_id = "9d3f3751-7521-4c19-9307-20e3cf6789a3" - message_trace_id = "710ace2f-bca8-41be-858c-54da42742a77" - opit_trace_id = "f7dfd978-0d10-4549-8abf-00f2cbc49d2c" - - message_data = MagicMock() - message_data.id = message_data_id - message_data.from_account_id = "acc-1" - message_data.from_end_user_id = None - message_data.provider_response_latency = 0.5 - message_data.conversation_id = conversation_id - message_data.total_price = 0.01 - message_data.model_id = "gpt-4" - message_data.answer = "hello" - message_data.status = MessageStatus.NORMAL - message_data.error = None - - trace_info = MessageTraceInfo( - message_id=message_trace_id, - message_data=message_data, - inputs={"query": "hi"}, - outputs={"answer": "hello"}, - message_tokens=10, - answer_tokens=20, - total_tokens=30, - start_time=_dt(), - end_time=_dt() + timedelta(seconds=1), - trace_id=opit_trace_id, - metadata={"foo": "bar"}, - conversation_mode="chat", - conversation_model="gpt-4", - file_list=[], - error=None, - message_file_data=MagicMock(url="test.png"), - ) - - trace_instance.add_trace = MagicMock(return_value=MagicMock(id="trace_id_1")) - trace_instance.add_span = MagicMock() - - trace_instance.message_trace(trace_info) - - trace_instance.add_trace.assert_called_once() - trace_instance.add_span.assert_called_once() - - -@pytest.mark.parametrize("sqlite3_session", [(EndUser,)], indirect=True) -def test_message_trace_with_end_user(trace_instance, monkeypatch: pytest.MonkeyPatch, sqlite3_session: Session) -> None: - write_session = sqlite3_session - message_data = MagicMock() - message_data.id = "85411059-79fb-4deb-a76c-c2e215f1b97e" - message_data.from_account_id = "acc-1" - message_data.from_end_user_id = "end-user-1" - message_data.conversation_id = "7d9f96d8-3be2-4e93-9c0e-922ff98dccc6" - message_data.status = MessageStatus.NORMAL - message_data.model_id = "gpt-4" - message_data.error = "" - message_data.answer = "hello" - message_data.total_price = 0.0 - message_data.provider_response_latency = 0.1 - - trace_info = MessageTraceInfo( - message_id="6bff35c7-33b7-4acb-ba21-44569a0327d0", - message_data=message_data, - inputs={}, - outputs={}, - message_tokens=0, - answer_tokens=0, - total_tokens=0, - start_time=_dt(), - end_time=_dt(), - metadata={}, - conversation_mode="chat", - conversation_model="gpt-4", - file_list=["url1"], - error=None, - ) - - end_user = EndUser( - id="end-user-1", - tenant_id="tenant-1", - app_id="app-1", - type=EndUserType.BROWSER, - session_id="session-id-123", - ) - write_session.add(end_user) - write_session.commit() - - with Session(write_session.get_bind()) as read_session: - monkeypatch.setattr( - "dify_trace_opik.opik_trace.db", - SimpleNamespace(engine=read_session.get_bind(), session=read_session), - ) - - trace_instance.add_trace = MagicMock(return_value=MagicMock(id="trace_id_2")) - trace_instance.add_span = MagicMock() - - trace_instance.message_trace(trace_info) - - trace_data = trace_instance.add_trace.call_args[0][0] - assert trace_data["metadata"]["user_id"] == "acc-1" - assert trace_data["metadata"]["end_user_id"] == "session-id-123" - - -def test_message_trace_none_data(trace_instance): - trace_info = SimpleNamespace(message_data=None, file_list=[], message_file_data=None, metadata={}) - trace_instance.add_trace = MagicMock() - trace_instance.message_trace(trace_info) - trace_instance.add_trace.assert_not_called() - - -def test_moderation_trace(trace_instance): - message_data = MagicMock() - message_data.created_at = _dt() - message_data.updated_at = _dt() - - trace_info = ModerationTraceInfo( - message_id="489d0dfd-065c-4106-8f9c-daded296c92d", - message_data=message_data, - inputs={"q": "hi"}, - action="stop", - flagged=True, - preset_response="blocked", - start_time=None, - end_time=None, - metadata={"foo": "bar"}, - trace_id="6f16cf18-9f4b-4955-8b6b-43cfa10978fc", - query="hi", - ) - - trace_instance.add_span = MagicMock() - trace_instance.moderation_trace(trace_info) - - trace_instance.add_span.assert_called_once() - span_data = trace_instance.add_span.call_args[0][0] - assert span_data["name"] == TraceTaskName.MODERATION_TRACE - assert span_data["output"]["flagged"] is True - - -def test_moderation_trace_none(trace_instance): - trace_info = ModerationTraceInfo( - message_id="cd732e4e-37f1-4c7e-8c64-820308bedcbf", - message_data=None, - inputs={}, - action="s", - flagged=False, - preset_response="", - query="", - metadata={}, - ) - trace_instance.add_span = MagicMock() - trace_instance.moderation_trace(trace_info) - trace_instance.add_span.assert_not_called() - - -def test_suggested_question_trace(trace_instance): - message_data = MagicMock() - message_data.created_at = _dt() - message_data.updated_at = _dt() - - trace_info = SuggestedQuestionTraceInfo( - message_id="7de55bda-a91d-477e-98ab-85c53c438469", - message_data=message_data, - inputs="hi", - suggested_question=["q1"], - total_tokens=10, - level="info", - start_time=_dt(), - end_time=_dt(), - metadata={}, - trace_id="a6687292-68c7-42ba-ae51-285579944d7b", - ) - - trace_instance.add_span = MagicMock() - trace_instance.suggested_question_trace(trace_info) - - trace_instance.add_span.assert_called_once() - span_data = trace_instance.add_span.call_args[0][0] - assert span_data["name"] == TraceTaskName.SUGGESTED_QUESTION_TRACE - - -def test_suggested_question_trace_none(trace_instance): - trace_info = SuggestedQuestionTraceInfo( - message_id="23696fc5-7e7f-46ec-bce8-1adc3c7f297d", - message_data=None, - inputs={}, - suggested_question=[], - total_tokens=0, - level="i", - metadata={}, - ) - trace_instance.add_span = MagicMock() - trace_instance.suggested_question_trace(trace_info) - trace_instance.add_span.assert_not_called() - - -def test_dataset_retrieval_trace(trace_instance): - message_data = MagicMock() - message_data.created_at = _dt() - message_data.updated_at = _dt() - - trace_info = DatasetRetrievalTraceInfo( - message_id="3e1a819f-c391-4950-adfd-96f82e5419a1", - message_data=message_data, - inputs="query", - documents=[{"id": "doc1"}], - start_time=None, - end_time=None, - metadata={}, - trace_id="41361000-e9be-4d11-b5e4-ab27ce0817d6", - ) - - trace_instance.add_span = MagicMock() - trace_instance.dataset_retrieval_trace(trace_info) - - trace_instance.add_span.assert_called_once() - span_data = trace_instance.add_span.call_args[0][0] - assert span_data["name"] == TraceTaskName.DATASET_RETRIEVAL_TRACE - - -def test_dataset_retrieval_trace_none(trace_instance): - trace_info = DatasetRetrievalTraceInfo( - message_id="35d6d44c-bccb-4e6e-8bd8-859257723ea8", message_data=None, inputs={}, documents=[], metadata={} - ) - trace_instance.add_span = MagicMock() - trace_instance.dataset_retrieval_trace(trace_info) - trace_instance.add_span.assert_not_called() - - -def test_tool_trace(trace_instance): - trace_info = ToolTraceInfo( - message_id="99db92c4-2254-496a-b5cc-18153315ce35", - message_data=MagicMock(), - inputs={}, - outputs={}, - tool_name="my_tool", - tool_inputs={"a": 1}, - tool_outputs="result_string", - time_cost=0.1, - start_time=_dt(), - end_time=_dt(), - metadata={}, - trace_id="a15a5fcb-7ffd-4458-8330-208f4cb1f796", - tool_config={}, - tool_parameters={}, - error="some error", - ) - - trace_instance.add_span = MagicMock() - trace_instance.tool_trace(trace_info) - - trace_instance.add_span.assert_called_once() - span_data = trace_instance.add_span.call_args[0][0] - assert span_data["name"] == "my_tool" - - -def test_generate_name_trace(trace_instance): - trace_info = GenerateNameTraceInfo( - inputs={"q": "hi"}, - outputs={"name": "new"}, - tenant_id="tenant-1", - conversation_id="271fe28f-6b86-416b-8d6b-bbbbfa9db791", - start_time=_dt(), - end_time=_dt(), - metadata={"921f010e-6878-4831-ae6b-271bf68c56fb": 1}, - ) - - trace_instance.add_trace = MagicMock(return_value=MagicMock(id="trace_id_3")) - trace_instance.add_span = MagicMock() - - trace_instance.generate_name_trace(trace_info) - - trace_instance.add_trace.assert_called_once() - trace_instance.add_span.assert_called_once() - - trace_data = trace_instance.add_trace.call_args[0][0] - assert trace_data["name"] == TraceTaskName.GENERATE_NAME_TRACE - - span_data = trace_instance.add_span.call_args[0][0] - assert span_data["trace_id"] == "trace_id_3" - - -def test_add_trace_success(trace_instance): - trace_data = {"id": "t1", "name": "trace"} - trace_instance.opik_client.trace.return_value = MagicMock(id="t1") - trace = trace_instance.add_trace(trace_data) - trace_instance.opik_client.trace.assert_called_once() - assert trace.id == "t1" - - -def test_add_trace_error(trace_instance): - trace_instance.opik_client.trace.side_effect = Exception("error") - trace_data = {"id": "t1", "name": "trace"} - with pytest.raises(ValueError, match="Opik Failed to create trace: error"): - trace_instance.add_trace(trace_data) - - -def test_add_span_success(trace_instance): - span_data = {"id": "s1", "name": "span", "trace_id": "t1"} - trace_instance.add_span(span_data) - trace_instance.opik_client.span.assert_called_once() - - -def test_add_span_error(trace_instance): - trace_instance.opik_client.span.side_effect = Exception("error") - span_data = {"id": "s1", "name": "span", "trace_id": "t1"} - with pytest.raises(ValueError, match="Opik Failed to create span: error"): - trace_instance.add_span(span_data) - - -def test_api_check_success(trace_instance): - trace_instance.opik_client.auth_check.return_value = True - assert trace_instance.api_check() is True - - -def test_api_check_error(trace_instance): - trace_instance.opik_client.auth_check.side_effect = Exception("fail") - with pytest.raises(ValueError, match="Opik API check failed: fail"): - trace_instance.api_check() - - -def test_get_project_url_success(trace_instance): - trace_instance.opik_client.get_project_url.return_value = "http://project.url" - assert trace_instance.get_project_url() == "http://project.url" - trace_instance.opik_client.get_project_url.assert_called_once_with(project_name=trace_instance.project) - - -def test_get_project_url_error(trace_instance): - trace_instance.opik_client.get_project_url.side_effect = Exception("fail") - with pytest.raises(ValueError, match="Opik get run url failed: fail"): - trace_instance.get_project_url() - - -@pytest.mark.parametrize("sqlite3_session", [()], indirect=True) -def test_workflow_trace_usage_extraction_error_fixed( - trace_instance, - monkeypatch: pytest.MonkeyPatch, - caplog: pytest.LogCaptureFixture, - sqlite3_session: Session, -) -> None: - trace_info = WorkflowTraceInfo( - workflow_id="86a52565-4a6b-4a1b-9bfd-98e4595e70de", - tenant_id="66e8e918-472e-4b69-8051-12502c34fc07", - workflow_run_id="8403965c-3344-4d22-a8fe-d8d55cee64d9", - workflow_run_elapsed_time=1.0, - workflow_run_status="s", - workflow_run_inputs={}, - workflow_run_outputs={}, - workflow_run_version="1", - total_tokens=0, - file_list=[], - query="", - message_id=None, - conversation_id="7a02cb9d-6949-4c59-a89d-f25bbc881e0e", - start_time=_dt(), - end_time=_dt(), - metadata={"app_id": "77e8e918-472e-4b69-8051-12502c34fc07"}, - workflow_app_log_id="82268424-e193-476c-a6db-f473388ee5fe", - error="", - ) - - node = MagicMock() - node.id = "88e8e918-472e-4b69-8051-12502c34fc07" - node.title = "LLM Node" - node.node_type = BuiltinNodeTypes.LLM - node.status = "succeeded" - - class BadDict(collections.UserDict): - @override - def get(self, key, default=None): - if key == "usage": - raise Exception("Usage extraction failed") - return super().get(key, default) - - node.process_data = BadDict({"model_mode": "chat", "model_name": "gpt-4", "usage": True, "prompts": ["p"]}) - node.created_at = _dt() - node.elapsed_time = 0.1 - node.metadata = {} - node.outputs = {} - - repo = MagicMock() - repo.get_by_workflow_execution.return_value = [node] - mock_factory = MagicMock() - mock_factory.create_workflow_node_execution_repository.return_value = repo - monkeypatch.setattr("dify_trace_opik.opik_trace.DifyCoreRepositoryFactory", mock_factory) - monkeypatch.setattr( - "dify_trace_opik.opik_trace.db", - SimpleNamespace(engine=sqlite3_session.get_bind(), session=sqlite3_session), - ) - monkeypatch.setattr(trace_instance, "get_service_account_with_tenant", lambda app_id: MagicMock()) - - trace_instance.add_trace = MagicMock() - trace_instance.add_span = MagicMock() - - with caplog.at_level(logging.ERROR): - trace_instance.workflow_trace(trace_info) - - assert "Failed to extract usage" in caplog.text - assert trace_instance.add_span.call_count >= 1 - # Verify that at least one of the spans is for the LLM Node - span_names = [call.args[0]["name"] for call in trace_instance.add_span.call_args_list] - assert "LLM Node" in span_names diff --git a/api/providers/trace/trace-opik/tests/unit_tests/test_config_entity.py b/api/providers/trace/trace-opik/tests/unit_tests/src/dify_trace_opik/test_config.py similarity index 100% rename from api/providers/trace/trace-opik/tests/unit_tests/test_config_entity.py rename to api/providers/trace/trace-opik/tests/unit_tests/src/dify_trace_opik/test_config.py diff --git a/api/providers/trace/trace-opik/tests/unit_tests/src/dify_trace_opik/test_opik_trace.py b/api/providers/trace/trace-opik/tests/unit_tests/src/dify_trace_opik/test_opik_trace.py new file mode 100644 index 00000000000..c228c00ebef --- /dev/null +++ b/api/providers/trace/trace-opik/tests/unit_tests/src/dify_trace_opik/test_opik_trace.py @@ -0,0 +1,1052 @@ +from __future__ import annotations + +import collections +import logging +import uuid +from datetime import UTC, datetime, timedelta +from types import SimpleNamespace +from typing import cast, override +from unittest.mock import MagicMock, patch + +import pytest +from dify_trace_opik.config import OpikConfig +from dify_trace_opik.opik_trace import OpikDataTrace, _seed_to_uuid4, prepare_opik_uuid, wrap_dict, wrap_metadata +from sqlalchemy.orm import Session + +from core.ops.entities.trace_entity import ( + DatasetRetrievalTraceInfo, + GenerateNameTraceInfo, + MessageTraceInfo, + ModerationTraceInfo, + SuggestedQuestionTraceInfo, + ToolTraceInfo, + TraceTaskName, + WorkflowTraceInfo, +) +from graphon.enums import BuiltinNodeTypes, WorkflowNodeExecutionMetadataKey +from models import EndUser +from models.enums import EndUserType, MessageStatus +from tests.helpers.unit.core.ops.trace_fixtures import message_trace_info, tool_trace_info, workflow_trace_info + +"""Unit tests for Opik trace translation with SQLite-backed lookups.""" + + +def _dt() -> datetime: + return datetime(2024, 1, 1, 0, 0, 0, tzinfo=UTC) + + +class TestOpikTraceCases: + @staticmethod + @pytest.fixture + def opik_config(): + return OpikConfig( + project="test-project", workspace="test-workspace", url="https://cloud.opik.com/api/", api_key="api-key-123" + ) + + @staticmethod + @pytest.fixture + def trace_instance(opik_config, monkeypatch: pytest.MonkeyPatch): + mock_client = MagicMock() + monkeypatch.setattr("dify_trace_opik.opik_trace.Opik", lambda **kwargs: mock_client) + + instance = OpikDataTrace(opik_config) + return instance + + @staticmethod + def test_wrap_dict(): + assert wrap_dict("input", {"a": 1}) == {"a": 1} + assert wrap_dict("input", "hello") == {"input": "hello"} + + @staticmethod + def test_wrap_metadata(): + assert wrap_metadata({"a": 1}, b=2) == {"a": 1, "b": 2, "created_from": "dify"} + + @staticmethod + def test_prepare_opik_uuid(): + # Test with valid datetime and uuid string + dt = datetime(2024, 1, 1) + uuid_str = "b3e8e918-472e-4b69-8051-12502c34fc07" + result = prepare_opik_uuid(dt, uuid_str) + assert result is not None + # We won't test the exact uuid7 value but just that it returns a string id + + # Test with None dt and uuid_str + result = prepare_opik_uuid(None, None) + assert result is not None + + @staticmethod + def test_init(opik_config, monkeypatch: pytest.MonkeyPatch): + mock_opik = MagicMock() + monkeypatch.setattr("dify_trace_opik.opik_trace.Opik", mock_opik) + monkeypatch.setenv("FILES_URL", "http://test.url") + + instance = OpikDataTrace(opik_config) + + mock_opik.assert_called_once_with( + project_name=opik_config.project, + workspace=opik_config.workspace, + host=opik_config.url, + api_key=opik_config.api_key, + ) + assert instance.file_base_url == "http://test.url" + assert instance.project == opik_config.project + + @staticmethod + def test_trace_dispatch(trace_instance, monkeypatch: pytest.MonkeyPatch): + methods = [ + "workflow_trace", + "message_trace", + "moderation_trace", + "suggested_question_trace", + "dataset_retrieval_trace", + "tool_trace", + "generate_name_trace", + ] + mocks = {method: MagicMock() for method in methods} + for method, m in mocks.items(): + monkeypatch.setattr(trace_instance, method, m) + + # WorkflowTraceInfo + info = workflow_trace_info() + trace_instance.trace(info) + mocks["workflow_trace"].assert_called_once_with(info) + + # MessageTraceInfo + info = message_trace_info() + trace_instance.trace(info) + mocks["message_trace"].assert_called_once_with(info) + + # ModerationTraceInfo + info = ModerationTraceInfo(metadata={}, flagged=False, action="allow", preset_response="", query="hello") + trace_instance.trace(info) + mocks["moderation_trace"].assert_called_once_with(info) + + # SuggestedQuestionTraceInfo + info = SuggestedQuestionTraceInfo(metadata={}, total_tokens=0, suggested_question=[], level="info") + trace_instance.trace(info) + mocks["suggested_question_trace"].assert_called_once_with(info) + + # DatasetRetrievalTraceInfo + info = DatasetRetrievalTraceInfo(metadata={}) + trace_instance.trace(info) + mocks["dataset_retrieval_trace"].assert_called_once_with(info) + + # ToolTraceInfo + info = tool_trace_info() + trace_instance.trace(info) + mocks["tool_trace"].assert_called_once_with(info) + + # GenerateNameTraceInfo + info = GenerateNameTraceInfo(tenant_id="tenant-1", metadata={}) + trace_instance.trace(info) + mocks["generate_name_trace"].assert_called_once_with(info) + + @staticmethod + @pytest.mark.parametrize("sqlite3_session", [()], indirect=True) + def test_workflow_trace_with_message_id( + trace_instance, monkeypatch: pytest.MonkeyPatch, sqlite3_session: Session + ) -> None: + # Define constants for better readability + workflow_id = "fb05c7cd-6cec-4add-8a84-df03a408b4ce" + workflow_run_id = "33c67568-7a8a-450e-8916-a5f135baeaef" + message_id = "04ec3956-85f3-488a-8539-1017251dc8c6" + conversation_id = "d3d01066-23ae-4830-9ce4-eb5640b42a7e" + trace_id = "bf26d929-6f15-4c2f-9abc-761c217056f3" + workflow_app_log_id = "ca0e018e-edd4-43fb-a05a-ea001ca8ef4b" + llm_node_id = "80d7dfa8-08f4-4ab7-aa37-0ca7d27207e3" + code_node_id = "b9cd9a7b-c534-4aa9-b5da-efd454140900" + + trace_info = WorkflowTraceInfo( + workflow_id=workflow_id, + tenant_id="tenant-1", + workflow_run_id=workflow_run_id, + workflow_run_elapsed_time=1.0, + workflow_run_status="succeeded", + workflow_run_inputs={"input": "hi"}, + workflow_run_outputs={"output": "hello"}, + workflow_run_version="1.0", + message_id=message_id, + conversation_id=conversation_id, + total_tokens=100, + file_list=[], + query="hi", + start_time=_dt(), + end_time=_dt() + timedelta(seconds=1), + trace_id=trace_id, + metadata={"app_id": "app-1", "user_id": "user-1"}, + workflow_app_log_id=workflow_app_log_id, + error="", + ) + + monkeypatch.setattr( + "dify_trace_opik.opik_trace.db", + SimpleNamespace(engine=sqlite3_session.get_bind(), session=sqlite3_session), + ) + + node_llm = MagicMock() + node_llm.id = llm_node_id + node_llm.title = "LLM Node" + node_llm.node_type = BuiltinNodeTypes.LLM + node_llm.status = "succeeded" + node_llm.process_data = { + "model_mode": "chat", + "model_name": "gpt-4", + "model_provider": "openai", + "usage": {"prompt_tokens": 10, "completion_tokens": 20, "total_tokens": 30}, + } + node_llm.inputs = {"prompts": "p"} + node_llm.outputs = {"text": "t"} + node_llm.created_at = _dt() + node_llm.elapsed_time = 0.5 + node_llm.metadata = {"foo": "bar"} + + node_other = MagicMock() + node_other.id = code_node_id + node_other.title = "Other Node" + node_other.node_type = BuiltinNodeTypes.CODE + node_other.status = "failed" + node_other.process_data = None + node_other.inputs = {"code": "print"} + node_other.outputs = {"result": "ok"} + node_other.created_at = None + node_other.elapsed_time = 0.2 + node_other.metadata = {WorkflowNodeExecutionMetadataKey.TOTAL_TOKENS.value: 10} + + repo = MagicMock() + repo.get_by_workflow_execution.return_value = [node_llm, node_other] + + mock_factory = MagicMock() + mock_factory.create_workflow_node_execution_repository.return_value = repo + monkeypatch.setattr("dify_trace_opik.opik_trace.DifyCoreRepositoryFactory", mock_factory) + + monkeypatch.setattr(trace_instance, "get_service_account_with_tenant", lambda app_id: MagicMock()) + + trace_instance.add_trace = MagicMock() + trace_instance.add_span = MagicMock() + + trace_instance.workflow_trace(trace_info) + + trace_instance.add_trace.assert_called_once() + trace_data = trace_instance.add_trace.call_args[1].get( + "opik_trace_data", trace_instance.add_trace.call_args[0][0] + ) + assert trace_data["name"] == TraceTaskName.MESSAGE_TRACE + assert "message" in trace_data["tags"] + assert "workflow" in trace_data["tags"] + + assert trace_instance.add_span.call_count >= 1 + + @staticmethod + @pytest.mark.parametrize("sqlite3_session", [()], indirect=True) + def test_workflow_trace_no_message_id( + trace_instance, monkeypatch: pytest.MonkeyPatch, sqlite3_session: Session + ) -> None: + # Define constants for better readability + workflow_id = "f0708b36-b1d7-42b3-a876-1d01b7d8f1a3" + workflow_run_id = "d42ec285-c2fd-4248-8866-5c9386b101ac" + conversation_id = "88a17f2e-9436-4472-bab9-4b1601d5af3c" + workflow_app_log_id = "41780d0d-ffba-4220-bc0c-401e4c89cdfb" + + trace_info = WorkflowTraceInfo( + workflow_id=workflow_id, + tenant_id="tenant-1", + workflow_run_id=workflow_run_id, + workflow_run_elapsed_time=1.0, + workflow_run_status="succeeded", + workflow_run_inputs={}, + workflow_run_outputs={}, + workflow_run_version="1.0", + total_tokens=0, + file_list=[], + query="", + message_id=None, + conversation_id=conversation_id, + start_time=_dt(), + end_time=_dt(), + trace_id=None, + metadata={"app_id": "app-1"}, + workflow_app_log_id=workflow_app_log_id, + error="", + ) + + monkeypatch.setattr( + "dify_trace_opik.opik_trace.db", + SimpleNamespace(engine=sqlite3_session.get_bind(), session=sqlite3_session), + ) + repo = MagicMock() + repo.get_by_workflow_execution.return_value = [] + mock_factory = MagicMock() + mock_factory.create_workflow_node_execution_repository.return_value = repo + monkeypatch.setattr("dify_trace_opik.opik_trace.DifyCoreRepositoryFactory", mock_factory) + monkeypatch.setattr(trace_instance, "get_service_account_with_tenant", lambda app_id: MagicMock()) + + trace_instance.add_trace = MagicMock() + trace_instance.workflow_trace(trace_info) + + trace_instance.add_trace.assert_called_once() + + @staticmethod + @pytest.mark.parametrize("sqlite3_session", [()], indirect=True) + def test_workflow_trace_missing_app_id( + trace_instance, monkeypatch: pytest.MonkeyPatch, sqlite3_session: Session + ) -> None: + trace_info = WorkflowTraceInfo( + workflow_id="5745f1b8-f8e6-4859-8110-996acb6c8d6a", + tenant_id="tenant-1", + workflow_run_id="46f53304-1659-464b-bee5-116585f0bec8", + workflow_run_elapsed_time=1.0, + workflow_run_status="succeeded", + workflow_run_inputs={}, + workflow_run_outputs={}, + workflow_run_version="1.0", + total_tokens=0, + file_list=[], + query="", + message_id=None, + conversation_id="83f86b89-caef-4de8-a0f9-f164eddae1ea", + start_time=_dt(), + end_time=_dt(), + metadata={}, + workflow_app_log_id="339760b2-4b94-4532-8c81-133a97e4680e", + error="", + ) + monkeypatch.setattr( + "dify_trace_opik.opik_trace.db", + SimpleNamespace(engine=sqlite3_session.get_bind(), session=sqlite3_session), + ) + + with pytest.raises(ValueError, match="No app_id found in trace_info metadata"): + trace_instance.workflow_trace(trace_info) + + @staticmethod + def test_message_trace_basic(trace_instance, monkeypatch: pytest.MonkeyPatch): + # Define constants for better readability + message_data_id = "e3a26712-8cac-4a25-94a4-a3bff21ee3ab" + conversation_id = "9d3f3751-7521-4c19-9307-20e3cf6789a3" + message_trace_id = "710ace2f-bca8-41be-858c-54da42742a77" + opit_trace_id = "f7dfd978-0d10-4549-8abf-00f2cbc49d2c" + + message_data = MagicMock() + message_data.id = message_data_id + message_data.from_account_id = "acc-1" + message_data.from_end_user_id = None + message_data.provider_response_latency = 0.5 + message_data.conversation_id = conversation_id + message_data.total_price = 0.01 + message_data.model_id = "gpt-4" + message_data.answer = "hello" + message_data.status = MessageStatus.NORMAL + message_data.error = None + + trace_info = MessageTraceInfo( + message_id=message_trace_id, + message_data=message_data, + inputs={"query": "hi"}, + outputs={"answer": "hello"}, + message_tokens=10, + answer_tokens=20, + total_tokens=30, + start_time=_dt(), + end_time=_dt() + timedelta(seconds=1), + trace_id=opit_trace_id, + metadata={"foo": "bar"}, + conversation_mode="chat", + conversation_model="gpt-4", + file_list=[], + error=None, + message_file_data=MagicMock(url="test.png"), + ) + + trace_instance.add_trace = MagicMock(return_value=MagicMock(id="trace_id_1")) + trace_instance.add_span = MagicMock() + + trace_instance.message_trace(trace_info) + + trace_instance.add_trace.assert_called_once() + trace_instance.add_span.assert_called_once() + + @staticmethod + @pytest.mark.parametrize("sqlite3_session", [(EndUser,)], indirect=True) + def test_message_trace_with_end_user( + trace_instance, monkeypatch: pytest.MonkeyPatch, sqlite3_session: Session + ) -> None: + write_session = sqlite3_session + message_data = MagicMock() + message_data.id = "85411059-79fb-4deb-a76c-c2e215f1b97e" + message_data.from_account_id = "acc-1" + message_data.from_end_user_id = "end-user-1" + message_data.conversation_id = "7d9f96d8-3be2-4e93-9c0e-922ff98dccc6" + message_data.status = MessageStatus.NORMAL + message_data.model_id = "gpt-4" + message_data.error = "" + message_data.answer = "hello" + message_data.total_price = 0.0 + message_data.provider_response_latency = 0.1 + + trace_info = MessageTraceInfo( + message_id="6bff35c7-33b7-4acb-ba21-44569a0327d0", + message_data=message_data, + inputs={}, + outputs={}, + message_tokens=0, + answer_tokens=0, + total_tokens=0, + start_time=_dt(), + end_time=_dt(), + metadata={}, + conversation_mode="chat", + conversation_model="gpt-4", + file_list=["url1"], + error=None, + ) + + end_user = EndUser( + id="end-user-1", + tenant_id="tenant-1", + app_id="app-1", + type=EndUserType.BROWSER, + session_id="session-id-123", + ) + write_session.add(end_user) + write_session.commit() + + with Session(write_session.get_bind()) as read_session: + monkeypatch.setattr( + "dify_trace_opik.opik_trace.db", + SimpleNamespace(engine=read_session.get_bind(), session=read_session), + ) + + trace_instance.add_trace = MagicMock(return_value=MagicMock(id="trace_id_2")) + trace_instance.add_span = MagicMock() + + trace_instance.message_trace(trace_info) + + trace_data = trace_instance.add_trace.call_args[0][0] + assert trace_data["metadata"]["user_id"] == "acc-1" + assert trace_data["metadata"]["end_user_id"] == "session-id-123" + + @staticmethod + def test_message_trace_none_data(trace_instance): + trace_info = SimpleNamespace(message_data=None, file_list=[], message_file_data=None, metadata={}) + trace_instance.add_trace = MagicMock() + trace_instance.message_trace(trace_info) + trace_instance.add_trace.assert_not_called() + + @staticmethod + def test_moderation_trace(trace_instance): + message_data = MagicMock() + message_data.created_at = _dt() + message_data.updated_at = _dt() + + trace_info = ModerationTraceInfo( + message_id="489d0dfd-065c-4106-8f9c-daded296c92d", + message_data=message_data, + inputs={"q": "hi"}, + action="stop", + flagged=True, + preset_response="blocked", + start_time=None, + end_time=None, + metadata={"foo": "bar"}, + trace_id="6f16cf18-9f4b-4955-8b6b-43cfa10978fc", + query="hi", + ) + + trace_instance.add_span = MagicMock() + trace_instance.moderation_trace(trace_info) + + trace_instance.add_span.assert_called_once() + span_data = trace_instance.add_span.call_args[0][0] + assert span_data["name"] == TraceTaskName.MODERATION_TRACE + assert span_data["output"]["flagged"] is True + + @staticmethod + def test_moderation_trace_none(trace_instance): + trace_info = ModerationTraceInfo( + message_id="cd732e4e-37f1-4c7e-8c64-820308bedcbf", + message_data=None, + inputs={}, + action="s", + flagged=False, + preset_response="", + query="", + metadata={}, + ) + trace_instance.add_span = MagicMock() + trace_instance.moderation_trace(trace_info) + trace_instance.add_span.assert_not_called() + + @staticmethod + def test_suggested_question_trace(trace_instance): + message_data = MagicMock() + message_data.created_at = _dt() + message_data.updated_at = _dt() + + trace_info = SuggestedQuestionTraceInfo( + message_id="7de55bda-a91d-477e-98ab-85c53c438469", + message_data=message_data, + inputs="hi", + suggested_question=["q1"], + total_tokens=10, + level="info", + start_time=_dt(), + end_time=_dt(), + metadata={}, + trace_id="a6687292-68c7-42ba-ae51-285579944d7b", + ) + + trace_instance.add_span = MagicMock() + trace_instance.suggested_question_trace(trace_info) + + trace_instance.add_span.assert_called_once() + span_data = trace_instance.add_span.call_args[0][0] + assert span_data["name"] == TraceTaskName.SUGGESTED_QUESTION_TRACE + + @staticmethod + def test_suggested_question_trace_none(trace_instance): + trace_info = SuggestedQuestionTraceInfo( + message_id="23696fc5-7e7f-46ec-bce8-1adc3c7f297d", + message_data=None, + inputs={}, + suggested_question=[], + total_tokens=0, + level="i", + metadata={}, + ) + trace_instance.add_span = MagicMock() + trace_instance.suggested_question_trace(trace_info) + trace_instance.add_span.assert_not_called() + + @staticmethod + def test_dataset_retrieval_trace(trace_instance): + message_data = MagicMock() + message_data.created_at = _dt() + message_data.updated_at = _dt() + + trace_info = DatasetRetrievalTraceInfo( + message_id="3e1a819f-c391-4950-adfd-96f82e5419a1", + message_data=message_data, + inputs="query", + documents=[{"id": "doc1"}], + start_time=None, + end_time=None, + metadata={}, + trace_id="41361000-e9be-4d11-b5e4-ab27ce0817d6", + ) + + trace_instance.add_span = MagicMock() + trace_instance.dataset_retrieval_trace(trace_info) + + trace_instance.add_span.assert_called_once() + span_data = trace_instance.add_span.call_args[0][0] + assert span_data["name"] == TraceTaskName.DATASET_RETRIEVAL_TRACE + + @staticmethod + def test_dataset_retrieval_trace_none(trace_instance): + trace_info = DatasetRetrievalTraceInfo( + message_id="35d6d44c-bccb-4e6e-8bd8-859257723ea8", message_data=None, inputs={}, documents=[], metadata={} + ) + trace_instance.add_span = MagicMock() + trace_instance.dataset_retrieval_trace(trace_info) + trace_instance.add_span.assert_not_called() + + @staticmethod + def test_tool_trace(trace_instance): + trace_info = ToolTraceInfo( + message_id="99db92c4-2254-496a-b5cc-18153315ce35", + message_data=MagicMock(), + inputs={}, + outputs={}, + tool_name="my_tool", + tool_inputs={"a": 1}, + tool_outputs="result_string", + time_cost=0.1, + start_time=_dt(), + end_time=_dt(), + metadata={}, + trace_id="a15a5fcb-7ffd-4458-8330-208f4cb1f796", + tool_config={}, + tool_parameters={}, + error="some error", + ) + + trace_instance.add_span = MagicMock() + trace_instance.tool_trace(trace_info) + + trace_instance.add_span.assert_called_once() + span_data = trace_instance.add_span.call_args[0][0] + assert span_data["name"] == "my_tool" + + @staticmethod + def test_generate_name_trace(trace_instance): + trace_info = GenerateNameTraceInfo( + inputs={"q": "hi"}, + outputs={"name": "new"}, + tenant_id="tenant-1", + conversation_id="271fe28f-6b86-416b-8d6b-bbbbfa9db791", + start_time=_dt(), + end_time=_dt(), + metadata={"921f010e-6878-4831-ae6b-271bf68c56fb": 1}, + ) + + trace_instance.add_trace = MagicMock(return_value=MagicMock(id="trace_id_3")) + trace_instance.add_span = MagicMock() + + trace_instance.generate_name_trace(trace_info) + + trace_instance.add_trace.assert_called_once() + trace_instance.add_span.assert_called_once() + + trace_data = trace_instance.add_trace.call_args[0][0] + assert trace_data["name"] == TraceTaskName.GENERATE_NAME_TRACE + + span_data = trace_instance.add_span.call_args[0][0] + assert span_data["trace_id"] == "trace_id_3" + + @staticmethod + def test_add_trace_success(trace_instance): + trace_data = {"id": "t1", "name": "trace"} + trace_instance.opik_client.trace.return_value = MagicMock(id="t1") + trace = trace_instance.add_trace(trace_data) + trace_instance.opik_client.trace.assert_called_once() + assert trace.id == "t1" + + @staticmethod + def test_add_trace_error(trace_instance): + trace_instance.opik_client.trace.side_effect = Exception("error") + trace_data = {"id": "t1", "name": "trace"} + with pytest.raises(ValueError, match="Opik Failed to create trace: error"): + trace_instance.add_trace(trace_data) + + @staticmethod + def test_add_span_success(trace_instance): + span_data = {"id": "s1", "name": "span", "trace_id": "t1"} + trace_instance.add_span(span_data) + trace_instance.opik_client.span.assert_called_once() + + @staticmethod + def test_add_span_error(trace_instance): + trace_instance.opik_client.span.side_effect = Exception("error") + span_data = {"id": "s1", "name": "span", "trace_id": "t1"} + with pytest.raises(ValueError, match="Opik Failed to create span: error"): + trace_instance.add_span(span_data) + + @staticmethod + def test_api_check_success(trace_instance): + trace_instance.opik_client.auth_check.return_value = True + assert trace_instance.api_check() is True + + @staticmethod + def test_api_check_error(trace_instance): + trace_instance.opik_client.auth_check.side_effect = Exception("fail") + with pytest.raises(ValueError, match="Opik API check failed: fail"): + trace_instance.api_check() + + @staticmethod + def test_get_project_url_success(trace_instance): + trace_instance.opik_client.get_project_url.return_value = "http://project.url" + assert trace_instance.get_project_url() == "http://project.url" + trace_instance.opik_client.get_project_url.assert_called_once_with(project_name=trace_instance.project) + + @staticmethod + def test_get_project_url_error(trace_instance): + trace_instance.opik_client.get_project_url.side_effect = Exception("fail") + with pytest.raises(ValueError, match="Opik get run url failed: fail"): + trace_instance.get_project_url() + + @staticmethod + @pytest.mark.parametrize("sqlite3_session", [()], indirect=True) + def test_workflow_trace_usage_extraction_error_fixed( + trace_instance, + monkeypatch: pytest.MonkeyPatch, + caplog: pytest.LogCaptureFixture, + sqlite3_session: Session, + ) -> None: + trace_info = WorkflowTraceInfo( + workflow_id="86a52565-4a6b-4a1b-9bfd-98e4595e70de", + tenant_id="66e8e918-472e-4b69-8051-12502c34fc07", + workflow_run_id="8403965c-3344-4d22-a8fe-d8d55cee64d9", + workflow_run_elapsed_time=1.0, + workflow_run_status="s", + workflow_run_inputs={}, + workflow_run_outputs={}, + workflow_run_version="1", + total_tokens=0, + file_list=[], + query="", + message_id=None, + conversation_id="7a02cb9d-6949-4c59-a89d-f25bbc881e0e", + start_time=_dt(), + end_time=_dt(), + metadata={"app_id": "77e8e918-472e-4b69-8051-12502c34fc07"}, + workflow_app_log_id="82268424-e193-476c-a6db-f473388ee5fe", + error="", + ) + + node = MagicMock() + node.id = "88e8e918-472e-4b69-8051-12502c34fc07" + node.title = "LLM Node" + node.node_type = BuiltinNodeTypes.LLM + node.status = "succeeded" + + class BadDict(collections.UserDict): + @override + def get(self, key, default=None): + if key == "usage": + raise Exception("Usage extraction failed") + return super().get(key, default) + + node.process_data = BadDict({"model_mode": "chat", "model_name": "gpt-4", "usage": True, "prompts": ["p"]}) + node.created_at = _dt() + node.elapsed_time = 0.1 + node.metadata = {} + node.outputs = {} + + repo = MagicMock() + repo.get_by_workflow_execution.return_value = [node] + mock_factory = MagicMock() + mock_factory.create_workflow_node_execution_repository.return_value = repo + monkeypatch.setattr("dify_trace_opik.opik_trace.DifyCoreRepositoryFactory", mock_factory) + monkeypatch.setattr( + "dify_trace_opik.opik_trace.db", + SimpleNamespace(engine=sqlite3_session.get_bind(), session=sqlite3_session), + ) + monkeypatch.setattr(trace_instance, "get_service_account_with_tenant", lambda app_id: MagicMock()) + + trace_instance.add_trace = MagicMock() + trace_instance.add_span = MagicMock() + + with caplog.at_level(logging.ERROR): + trace_instance.workflow_trace(trace_info) + + assert "Failed to extract usage" in caplog.text + assert trace_instance.add_span.call_count >= 1 + # Verify that at least one of the spans is for the LLM Node + span_names = [call.args[0]["name"] for call in trace_instance.add_span.call_args_list] + assert "LLM Node" in span_names + + +"""Tests for OpikDataTrace workflow_trace changes. + +Covers: +- _seed_to_uuid4 helper: produces valid UUID4 strings deterministically +- prepare_opik_uuid helper: basic contract +- workflow_trace without message_id now creates a root span parented to None +- workflow_trace without message_id: node spans parent to root_span_id (not workflow_app_log_id) +- workflow_trace with message_id still creates root span keyed on workflow_run_id (unchanged path) +""" + + +# A stable UUID4 used as the workflow_run_id throughout all tests. +_WORKFLOW_RUN_ID = "a3f1b2c4-d5e6-4f78-9a0b-c1d2e3f4a5b6" + + +# --------------------------------------------------------------------------- +# Helpers +# --------------------------------------------------------------------------- + + +def _make_workflow_trace_info( + *, + message_id: str | None = None, + workflow_app_log_id: str | None = None, + workflow_run_id: str = _WORKFLOW_RUN_ID, +) -> WorkflowTraceInfo: + """Return a minimal WorkflowTraceInfo suitable for unit testing.""" + return WorkflowTraceInfo( + message_id=message_id, + workflow_id="wf-id", + tenant_id="tenant-id", + workflow_run_id=workflow_run_id, + workflow_app_log_id=workflow_app_log_id, + workflow_run_elapsed_time=1.5, + workflow_run_status="succeeded", + workflow_run_inputs={"query": "hello"}, + workflow_run_outputs={"result": "world"}, + workflow_run_version="1", + total_tokens=42, + file_list=[], + query="hello", + start_time=datetime(2025, 1, 1, 12, 0, 0), + end_time=datetime(2025, 1, 1, 12, 0, 1), + metadata={"app_id": "app-abc"}, + conversation_id=None, + ) + + +def _make_opik_trace_instance() -> OpikDataTrace: + """Construct an OpikDataTrace with the Opik SDK client mocked out.""" + with patch("dify_trace_opik.opik_trace.Opik"): + from dify_trace_opik.config import OpikConfig + + config = OpikConfig(api_key="key", project="test-project", url="https://www.comet.com/opik/api/") + instance = OpikDataTrace(config) + + instance.add_trace = MagicMock(return_value=MagicMock(id="mock-trace-id")) + instance.add_span = MagicMock() + instance.get_service_account_with_tenant = MagicMock(return_value=MagicMock()) + return instance + + +def _add_trace_mock(instance: OpikDataTrace) -> MagicMock: + return cast(MagicMock, instance.add_trace) + + +def _add_span_mock(instance: OpikDataTrace) -> MagicMock: + return cast(MagicMock, instance.add_span) + + +class TestOpikTraceCasesUnitTests: + # --------------------------------------------------------------------------- + # _seed_to_uuid4 + # --------------------------------------------------------------------------- + + class TestSeedToUuid4: + def test_returns_valid_uuid4_string(self): + result = _seed_to_uuid4("some-arbitrary-seed") + parsed = uuid.UUID(result) + assert parsed.version == 4 + + def test_is_deterministic(self): + assert _seed_to_uuid4("seed-abc") == _seed_to_uuid4("seed-abc") + + def test_different_seeds_give_different_results(self): + assert _seed_to_uuid4("seed-1") != _seed_to_uuid4("seed-2") + + def test_workflow_run_id_with_root_suffix_is_valid_uuid4(self): + """The primary use-case: deriving a root-span UUID from workflow_run_id + '-root'.""" + seed = _WORKFLOW_RUN_ID + "-root" + result = _seed_to_uuid4(seed) + parsed = uuid.UUID(result) + assert parsed.version == 4 + + def test_seed_and_seed_root_produce_different_uuids(self): + """Root span UUID must differ from the base workflow UUID to avoid ID collisions.""" + base = _seed_to_uuid4(_WORKFLOW_RUN_ID) + with_root = _seed_to_uuid4(_WORKFLOW_RUN_ID + "-root") + assert base != with_root + + # --------------------------------------------------------------------------- + # prepare_opik_uuid + # --------------------------------------------------------------------------- + + class TestPrepareOpikUuid: + def test_is_deterministic(self): + dt = datetime(2025, 6, 15, 10, 30, 0) + uid = str(uuid.uuid4()) + assert prepare_opik_uuid(dt, uid) == prepare_opik_uuid(dt, uid) + + def test_different_uuids_give_different_results(self): + dt = datetime(2025, 6, 15, 10, 30, 0) + assert prepare_opik_uuid(dt, str(uuid.uuid4())) != prepare_opik_uuid(dt, str(uuid.uuid4())) + + def test_none_datetime_does_not_raise(self): + assert prepare_opik_uuid(None, str(uuid.uuid4())) is not None + + def test_none_uuid_does_not_raise(self): + assert prepare_opik_uuid(datetime(2025, 1, 1), None) is not None + + # --------------------------------------------------------------------------- + # workflow_trace — no message_id (new code path) + # --------------------------------------------------------------------------- + + class TestWorkflowTraceWithoutMessageId: + def _run(self, trace_info: WorkflowTraceInfo, node_executions: list | None = None): + instance = _make_opik_trace_instance() + fake_repo = MagicMock() + fake_repo.get_by_workflow_execution.return_value = node_executions or [] + + with ( + patch("dify_trace_opik.opik_trace.db") as mock_db, + patch("dify_trace_opik.opik_trace.sessionmaker"), + patch( + "dify_trace_opik.opik_trace.DifyCoreRepositoryFactory.create_workflow_node_execution_repository", + return_value=fake_repo, + ), + ): + mock_db.engine = MagicMock() + instance.workflow_trace(trace_info) + + return instance + + def _expected_root_span_id(self, trace_info: WorkflowTraceInfo): + return prepare_opik_uuid( + trace_info.start_time, + _seed_to_uuid4(trace_info.workflow_run_id + "-root"), + ) + + def test_root_span_is_created(self): + trace_info = _make_workflow_trace_info(message_id=None) + instance = self._run(trace_info) + assert _add_span_mock(instance).called + + def test_root_span_id_matches_expected(self): + trace_info = _make_workflow_trace_info(message_id=None) + instance = self._run(trace_info) + + expected = self._expected_root_span_id(trace_info) + root_span_kwargs = _add_span_mock(instance).call_args_list[0][0][0] + assert root_span_kwargs["id"] == expected + + def test_root_span_has_no_parent(self): + trace_info = _make_workflow_trace_info(message_id=None) + instance = self._run(trace_info) + + root_span_kwargs = _add_span_mock(instance).call_args_list[0][0][0] + assert root_span_kwargs["parent_span_id"] is None + + def test_trace_name_is_workflow_trace(self): + """Without message_id, the Opik trace itself should be named WORKFLOW_TRACE.""" + trace_info = _make_workflow_trace_info(message_id=None) + instance = self._run(trace_info) + + trace_kwargs = _add_trace_mock(instance).call_args_list[0][0][0] + assert trace_kwargs["name"] == TraceTaskName.WORKFLOW_TRACE + + def test_root_span_name_is_workflow_trace(self): + trace_info = _make_workflow_trace_info(message_id=None) + instance = self._run(trace_info) + + root_span_kwargs = _add_span_mock(instance).call_args_list[0][0][0] + assert root_span_kwargs["name"] == TraceTaskName.WORKFLOW_TRACE + + def test_root_span_has_workflow_tag(self): + trace_info = _make_workflow_trace_info(message_id=None) + instance = self._run(trace_info) + + root_span_kwargs = _add_span_mock(instance).call_args_list[0][0][0] + assert "workflow" in root_span_kwargs["tags"] + + def test_node_execution_spans_are_parented_to_root(self): + """Node spans must use root_span_id as parent, not any other ID.""" + trace_info = _make_workflow_trace_info(message_id=None) + expected_root_span_id = self._expected_root_span_id(trace_info) + + node_exec = MagicMock() + node_exec.id = str(uuid.uuid4()) + node_exec.title = "LLM Node" + node_exec.node_type = "llm" + node_exec.status = "succeeded" + node_exec.process_data = {} + node_exec.inputs = {"prompt": "hi"} + node_exec.outputs = {"text": "hello"} + node_exec.created_at = datetime(2025, 1, 1, 12, 0, 0) + node_exec.elapsed_time = 0.5 + node_exec.metadata = {} + + instance = self._run(trace_info, node_executions=[node_exec]) + + # call_args_list[0] = root span, [1] = node execution span + add_span = _add_span_mock(instance) + assert add_span.call_count == 2 + node_span_kwargs = add_span.call_args_list[1][0][0] + assert node_span_kwargs["parent_span_id"] == expected_root_span_id + + def test_node_span_not_parented_to_workflow_app_log_id(self): + """Old behaviour derived parent from workflow_app_log_id; that must no longer apply.""" + trace_info = _make_workflow_trace_info( + message_id=None, + workflow_app_log_id=str(uuid.uuid4()), + ) + + node_exec = MagicMock() + node_exec.id = str(uuid.uuid4()) + node_exec.title = "Tool Node" + node_exec.node_type = "tool" + node_exec.status = "succeeded" + node_exec.process_data = {} + node_exec.inputs = {} + node_exec.outputs = {} + node_exec.created_at = datetime(2025, 1, 1, 12, 0, 0) + node_exec.elapsed_time = 0.2 + node_exec.metadata = {} + + instance = self._run(trace_info, node_executions=[node_exec]) + + old_parent_id = prepare_opik_uuid(trace_info.start_time, trace_info.workflow_app_log_id) + node_span_kwargs = _add_span_mock(instance).call_args_list[1][0][0] + assert node_span_kwargs["parent_span_id"] != old_parent_id + + def test_root_span_id_differs_from_trace_id(self): + """The root span must have a different ID from the Opik trace to maintain correct hierarchy.""" + trace_info = _make_workflow_trace_info(message_id=None) + dify_trace_id = trace_info.trace_id or trace_info.workflow_run_id + opik_trace_id = prepare_opik_uuid(trace_info.start_time, dify_trace_id) + root_span_id = self._expected_root_span_id(trace_info) + assert root_span_id != opik_trace_id + + # --------------------------------------------------------------------------- + # workflow_trace — with message_id (unchanged path, guard against regression) + # --------------------------------------------------------------------------- + + class TestWorkflowTraceWithMessageId: + _MESSAGE_ID = str(uuid.uuid4()) + + def _run(self, trace_info: WorkflowTraceInfo, node_executions: list | None = None): + instance = _make_opik_trace_instance() + fake_repo = MagicMock() + fake_repo.get_by_workflow_execution.return_value = node_executions or [] + + with ( + patch("dify_trace_opik.opik_trace.db") as mock_db, + patch("dify_trace_opik.opik_trace.sessionmaker"), + patch( + "dify_trace_opik.opik_trace.DifyCoreRepositoryFactory.create_workflow_node_execution_repository", + return_value=fake_repo, + ), + ): + mock_db.engine = MagicMock() + instance.workflow_trace(trace_info) + + return instance + + def test_trace_name_is_message_trace(self): + """With message_id, the Opik trace should be named MESSAGE_TRACE.""" + trace_info = _make_workflow_trace_info(message_id=self._MESSAGE_ID) + instance = self._run(trace_info) + + trace_kwargs = _add_trace_mock(instance).call_args_list[0][0][0] + assert trace_kwargs["name"] == TraceTaskName.MESSAGE_TRACE + + def test_root_span_uses_workflow_run_id_directly(self): + """When message_id is set, root_span_id = prepare_opik_uuid(start_time, workflow_run_id).""" + trace_info = _make_workflow_trace_info(message_id=self._MESSAGE_ID) + instance = self._run(trace_info) + + expected_root_span_id = prepare_opik_uuid(trace_info.start_time, trace_info.workflow_run_id) + root_span_kwargs = _add_span_mock(instance).call_args_list[0][0][0] + assert root_span_kwargs["id"] == expected_root_span_id + + def test_root_span_id_differs_from_no_message_id_case(self): + """The two branches must produce different root span IDs for the same workflow_run_id.""" + id_with_message = prepare_opik_uuid( + datetime(2025, 1, 1, 12, 0, 0), + _WORKFLOW_RUN_ID, + ) + id_without_message = prepare_opik_uuid( + datetime(2025, 1, 1, 12, 0, 0), + _seed_to_uuid4(_WORKFLOW_RUN_ID + "-root"), + ) + assert id_with_message != id_without_message + + def test_node_spans_parented_to_workflow_run_root_span(self): + """Node spans must still parent to root_span_id derived from workflow_run_id.""" + trace_info = _make_workflow_trace_info(message_id=self._MESSAGE_ID) + expected_root_span_id = prepare_opik_uuid(trace_info.start_time, trace_info.workflow_run_id) + + node_exec = MagicMock() + node_exec.id = str(uuid.uuid4()) + node_exec.title = "LLM" + node_exec.node_type = "llm" + node_exec.status = "succeeded" + node_exec.process_data = {} + node_exec.inputs = {} + node_exec.outputs = {} + node_exec.created_at = datetime(2025, 1, 1, 12, 0, 0) + node_exec.elapsed_time = 0.3 + node_exec.metadata = {} + + instance = self._run(trace_info, node_executions=[node_exec]) + + node_span_kwargs = _add_span_mock(instance).call_args_list[1][0][0] + assert node_span_kwargs["parent_span_id"] == expected_root_span_id diff --git a/api/providers/trace/trace-opik/tests/unit_tests/test_opik_trace.py b/api/providers/trace/trace-opik/tests/unit_tests/test_opik_trace.py deleted file mode 100644 index 2e0796c2912..00000000000 --- a/api/providers/trace/trace-opik/tests/unit_tests/test_opik_trace.py +++ /dev/null @@ -1,340 +0,0 @@ -"""Tests for OpikDataTrace workflow_trace changes. - -Covers: -- _seed_to_uuid4 helper: produces valid UUID4 strings deterministically -- prepare_opik_uuid helper: basic contract -- workflow_trace without message_id now creates a root span parented to None -- workflow_trace without message_id: node spans parent to root_span_id (not workflow_app_log_id) -- workflow_trace with message_id still creates root span keyed on workflow_run_id (unchanged path) -""" - -from __future__ import annotations - -import uuid -from datetime import datetime -from typing import cast -from unittest.mock import MagicMock, patch - -from dify_trace_opik.opik_trace import OpikDataTrace, _seed_to_uuid4, prepare_opik_uuid - -from core.ops.entities.trace_entity import TraceTaskName, WorkflowTraceInfo - -# A stable UUID4 used as the workflow_run_id throughout all tests. -_WORKFLOW_RUN_ID = "a3f1b2c4-d5e6-4f78-9a0b-c1d2e3f4a5b6" - - -# --------------------------------------------------------------------------- -# Helpers -# --------------------------------------------------------------------------- - - -def _make_workflow_trace_info( - *, - message_id: str | None = None, - workflow_app_log_id: str | None = None, - workflow_run_id: str = _WORKFLOW_RUN_ID, -) -> WorkflowTraceInfo: - """Return a minimal WorkflowTraceInfo suitable for unit testing.""" - return WorkflowTraceInfo( - message_id=message_id, - workflow_id="wf-id", - tenant_id="tenant-id", - workflow_run_id=workflow_run_id, - workflow_app_log_id=workflow_app_log_id, - workflow_run_elapsed_time=1.5, - workflow_run_status="succeeded", - workflow_run_inputs={"query": "hello"}, - workflow_run_outputs={"result": "world"}, - workflow_run_version="1", - total_tokens=42, - file_list=[], - query="hello", - start_time=datetime(2025, 1, 1, 12, 0, 0), - end_time=datetime(2025, 1, 1, 12, 0, 1), - metadata={"app_id": "app-abc"}, - conversation_id=None, - ) - - -def _make_opik_trace_instance() -> OpikDataTrace: - """Construct an OpikDataTrace with the Opik SDK client mocked out.""" - with patch("dify_trace_opik.opik_trace.Opik"): - from dify_trace_opik.config import OpikConfig - - config = OpikConfig(api_key="key", project="test-project", url="https://www.comet.com/opik/api/") - instance = OpikDataTrace(config) - - instance.add_trace = MagicMock(return_value=MagicMock(id="mock-trace-id")) - instance.add_span = MagicMock() - instance.get_service_account_with_tenant = MagicMock(return_value=MagicMock()) - return instance - - -def _add_trace_mock(instance: OpikDataTrace) -> MagicMock: - return cast(MagicMock, instance.add_trace) - - -def _add_span_mock(instance: OpikDataTrace) -> MagicMock: - return cast(MagicMock, instance.add_span) - - -# --------------------------------------------------------------------------- -# _seed_to_uuid4 -# --------------------------------------------------------------------------- - - -class TestSeedToUuid4: - def test_returns_valid_uuid4_string(self): - result = _seed_to_uuid4("some-arbitrary-seed") - parsed = uuid.UUID(result) - assert parsed.version == 4 - - def test_is_deterministic(self): - assert _seed_to_uuid4("seed-abc") == _seed_to_uuid4("seed-abc") - - def test_different_seeds_give_different_results(self): - assert _seed_to_uuid4("seed-1") != _seed_to_uuid4("seed-2") - - def test_workflow_run_id_with_root_suffix_is_valid_uuid4(self): - """The primary use-case: deriving a root-span UUID from workflow_run_id + '-root'.""" - seed = _WORKFLOW_RUN_ID + "-root" - result = _seed_to_uuid4(seed) - parsed = uuid.UUID(result) - assert parsed.version == 4 - - def test_seed_and_seed_root_produce_different_uuids(self): - """Root span UUID must differ from the base workflow UUID to avoid ID collisions.""" - base = _seed_to_uuid4(_WORKFLOW_RUN_ID) - with_root = _seed_to_uuid4(_WORKFLOW_RUN_ID + "-root") - assert base != with_root - - -# --------------------------------------------------------------------------- -# prepare_opik_uuid -# --------------------------------------------------------------------------- - - -class TestPrepareOpikUuid: - def test_is_deterministic(self): - dt = datetime(2025, 6, 15, 10, 30, 0) - uid = str(uuid.uuid4()) - assert prepare_opik_uuid(dt, uid) == prepare_opik_uuid(dt, uid) - - def test_different_uuids_give_different_results(self): - dt = datetime(2025, 6, 15, 10, 30, 0) - assert prepare_opik_uuid(dt, str(uuid.uuid4())) != prepare_opik_uuid(dt, str(uuid.uuid4())) - - def test_none_datetime_does_not_raise(self): - assert prepare_opik_uuid(None, str(uuid.uuid4())) is not None - - def test_none_uuid_does_not_raise(self): - assert prepare_opik_uuid(datetime(2025, 1, 1), None) is not None - - -# --------------------------------------------------------------------------- -# workflow_trace — no message_id (new code path) -# --------------------------------------------------------------------------- - - -class TestWorkflowTraceWithoutMessageId: - def _run(self, trace_info: WorkflowTraceInfo, node_executions: list | None = None): - instance = _make_opik_trace_instance() - fake_repo = MagicMock() - fake_repo.get_by_workflow_execution.return_value = node_executions or [] - - with ( - patch("dify_trace_opik.opik_trace.db") as mock_db, - patch("dify_trace_opik.opik_trace.sessionmaker"), - patch( - "dify_trace_opik.opik_trace.DifyCoreRepositoryFactory.create_workflow_node_execution_repository", - return_value=fake_repo, - ), - ): - mock_db.engine = MagicMock() - instance.workflow_trace(trace_info) - - return instance - - def _expected_root_span_id(self, trace_info: WorkflowTraceInfo): - return prepare_opik_uuid( - trace_info.start_time, - _seed_to_uuid4(trace_info.workflow_run_id + "-root"), - ) - - def test_root_span_is_created(self): - trace_info = _make_workflow_trace_info(message_id=None) - instance = self._run(trace_info) - assert _add_span_mock(instance).called - - def test_root_span_id_matches_expected(self): - trace_info = _make_workflow_trace_info(message_id=None) - instance = self._run(trace_info) - - expected = self._expected_root_span_id(trace_info) - root_span_kwargs = _add_span_mock(instance).call_args_list[0][0][0] - assert root_span_kwargs["id"] == expected - - def test_root_span_has_no_parent(self): - trace_info = _make_workflow_trace_info(message_id=None) - instance = self._run(trace_info) - - root_span_kwargs = _add_span_mock(instance).call_args_list[0][0][0] - assert root_span_kwargs["parent_span_id"] is None - - def test_trace_name_is_workflow_trace(self): - """Without message_id, the Opik trace itself should be named WORKFLOW_TRACE.""" - trace_info = _make_workflow_trace_info(message_id=None) - instance = self._run(trace_info) - - trace_kwargs = _add_trace_mock(instance).call_args_list[0][0][0] - assert trace_kwargs["name"] == TraceTaskName.WORKFLOW_TRACE - - def test_root_span_name_is_workflow_trace(self): - trace_info = _make_workflow_trace_info(message_id=None) - instance = self._run(trace_info) - - root_span_kwargs = _add_span_mock(instance).call_args_list[0][0][0] - assert root_span_kwargs["name"] == TraceTaskName.WORKFLOW_TRACE - - def test_root_span_has_workflow_tag(self): - trace_info = _make_workflow_trace_info(message_id=None) - instance = self._run(trace_info) - - root_span_kwargs = _add_span_mock(instance).call_args_list[0][0][0] - assert "workflow" in root_span_kwargs["tags"] - - def test_node_execution_spans_are_parented_to_root(self): - """Node spans must use root_span_id as parent, not any other ID.""" - trace_info = _make_workflow_trace_info(message_id=None) - expected_root_span_id = self._expected_root_span_id(trace_info) - - node_exec = MagicMock() - node_exec.id = str(uuid.uuid4()) - node_exec.title = "LLM Node" - node_exec.node_type = "llm" - node_exec.status = "succeeded" - node_exec.process_data = {} - node_exec.inputs = {"prompt": "hi"} - node_exec.outputs = {"text": "hello"} - node_exec.created_at = datetime(2025, 1, 1, 12, 0, 0) - node_exec.elapsed_time = 0.5 - node_exec.metadata = {} - - instance = self._run(trace_info, node_executions=[node_exec]) - - # call_args_list[0] = root span, [1] = node execution span - add_span = _add_span_mock(instance) - assert add_span.call_count == 2 - node_span_kwargs = add_span.call_args_list[1][0][0] - assert node_span_kwargs["parent_span_id"] == expected_root_span_id - - def test_node_span_not_parented_to_workflow_app_log_id(self): - """Old behaviour derived parent from workflow_app_log_id; that must no longer apply.""" - trace_info = _make_workflow_trace_info( - message_id=None, - workflow_app_log_id=str(uuid.uuid4()), - ) - - node_exec = MagicMock() - node_exec.id = str(uuid.uuid4()) - node_exec.title = "Tool Node" - node_exec.node_type = "tool" - node_exec.status = "succeeded" - node_exec.process_data = {} - node_exec.inputs = {} - node_exec.outputs = {} - node_exec.created_at = datetime(2025, 1, 1, 12, 0, 0) - node_exec.elapsed_time = 0.2 - node_exec.metadata = {} - - instance = self._run(trace_info, node_executions=[node_exec]) - - old_parent_id = prepare_opik_uuid(trace_info.start_time, trace_info.workflow_app_log_id) - node_span_kwargs = _add_span_mock(instance).call_args_list[1][0][0] - assert node_span_kwargs["parent_span_id"] != old_parent_id - - def test_root_span_id_differs_from_trace_id(self): - """The root span must have a different ID from the Opik trace to maintain correct hierarchy.""" - trace_info = _make_workflow_trace_info(message_id=None) - dify_trace_id = trace_info.trace_id or trace_info.workflow_run_id - opik_trace_id = prepare_opik_uuid(trace_info.start_time, dify_trace_id) - root_span_id = self._expected_root_span_id(trace_info) - assert root_span_id != opik_trace_id - - -# --------------------------------------------------------------------------- -# workflow_trace — with message_id (unchanged path, guard against regression) -# --------------------------------------------------------------------------- - - -class TestWorkflowTraceWithMessageId: - _MESSAGE_ID = str(uuid.uuid4()) - - def _run(self, trace_info: WorkflowTraceInfo, node_executions: list | None = None): - instance = _make_opik_trace_instance() - fake_repo = MagicMock() - fake_repo.get_by_workflow_execution.return_value = node_executions or [] - - with ( - patch("dify_trace_opik.opik_trace.db") as mock_db, - patch("dify_trace_opik.opik_trace.sessionmaker"), - patch( - "dify_trace_opik.opik_trace.DifyCoreRepositoryFactory.create_workflow_node_execution_repository", - return_value=fake_repo, - ), - ): - mock_db.engine = MagicMock() - instance.workflow_trace(trace_info) - - return instance - - def test_trace_name_is_message_trace(self): - """With message_id, the Opik trace should be named MESSAGE_TRACE.""" - trace_info = _make_workflow_trace_info(message_id=self._MESSAGE_ID) - instance = self._run(trace_info) - - trace_kwargs = _add_trace_mock(instance).call_args_list[0][0][0] - assert trace_kwargs["name"] == TraceTaskName.MESSAGE_TRACE - - def test_root_span_uses_workflow_run_id_directly(self): - """When message_id is set, root_span_id = prepare_opik_uuid(start_time, workflow_run_id).""" - trace_info = _make_workflow_trace_info(message_id=self._MESSAGE_ID) - instance = self._run(trace_info) - - expected_root_span_id = prepare_opik_uuid(trace_info.start_time, trace_info.workflow_run_id) - root_span_kwargs = _add_span_mock(instance).call_args_list[0][0][0] - assert root_span_kwargs["id"] == expected_root_span_id - - def test_root_span_id_differs_from_no_message_id_case(self): - """The two branches must produce different root span IDs for the same workflow_run_id.""" - id_with_message = prepare_opik_uuid( - datetime(2025, 1, 1, 12, 0, 0), - _WORKFLOW_RUN_ID, - ) - id_without_message = prepare_opik_uuid( - datetime(2025, 1, 1, 12, 0, 0), - _seed_to_uuid4(_WORKFLOW_RUN_ID + "-root"), - ) - assert id_with_message != id_without_message - - def test_node_spans_parented_to_workflow_run_root_span(self): - """Node spans must still parent to root_span_id derived from workflow_run_id.""" - trace_info = _make_workflow_trace_info(message_id=self._MESSAGE_ID) - expected_root_span_id = prepare_opik_uuid(trace_info.start_time, trace_info.workflow_run_id) - - node_exec = MagicMock() - node_exec.id = str(uuid.uuid4()) - node_exec.title = "LLM" - node_exec.node_type = "llm" - node_exec.status = "succeeded" - node_exec.process_data = {} - node_exec.inputs = {} - node_exec.outputs = {} - node_exec.created_at = datetime(2025, 1, 1, 12, 0, 0) - node_exec.elapsed_time = 0.3 - node_exec.metadata = {} - - instance = self._run(trace_info, node_executions=[node_exec]) - - node_span_kwargs = _add_span_mock(instance).call_args_list[1][0][0] - assert node_span_kwargs["parent_span_id"] == expected_root_span_id diff --git a/api/providers/trace/trace-tencent/tests/unit_tests/tencent_trace/test_client.py b/api/providers/trace/trace-tencent/tests/unit_tests/src/dify_trace_tencent/test_client.py similarity index 100% rename from api/providers/trace/trace-tencent/tests/unit_tests/tencent_trace/test_client.py rename to api/providers/trace/trace-tencent/tests/unit_tests/src/dify_trace_tencent/test_client.py diff --git a/api/providers/trace/trace-tencent/tests/unit_tests/tencent_trace/test_span_builder.py b/api/providers/trace/trace-tencent/tests/unit_tests/src/dify_trace_tencent/test_span_builder.py similarity index 99% rename from api/providers/trace/trace-tencent/tests/unit_tests/tencent_trace/test_span_builder.py rename to api/providers/trace/trace-tencent/tests/unit_tests/src/dify_trace_tencent/test_span_builder.py index 905b0a11ffe..e1be98258c3 100644 --- a/api/providers/trace/trace-tencent/tests/unit_tests/tencent_trace/test_span_builder.py +++ b/api/providers/trace/trace-tencent/tests/unit_tests/src/dify_trace_tencent/test_span_builder.py @@ -23,7 +23,7 @@ from core.ops.entities.trace_entity import ( ) from core.rag.models.document import Document from graphon.enums import WorkflowNodeExecutionMetadataKey, WorkflowNodeExecutionStatus -from tests.unit_tests.core.ops.trace_fixtures import ( +from tests.helpers.unit.core.ops.trace_fixtures import ( message_trace_info, tool_trace_info, workflow_node_execution, diff --git a/api/providers/trace/trace-tencent/tests/unit_tests/tencent_trace/test_tencent_trace.py b/api/providers/trace/trace-tencent/tests/unit_tests/src/dify_trace_tencent/test_tencent_trace.py similarity index 99% rename from api/providers/trace/trace-tencent/tests/unit_tests/tencent_trace/test_tencent_trace.py rename to api/providers/trace/trace-tencent/tests/unit_tests/src/dify_trace_tencent/test_tencent_trace.py index 6a5b09b7bd6..5f53ecba4c8 100644 --- a/api/providers/trace/trace-tencent/tests/unit_tests/tencent_trace/test_tencent_trace.py +++ b/api/providers/trace/trace-tencent/tests/unit_tests/src/dify_trace_tencent/test_tencent_trace.py @@ -24,7 +24,7 @@ from graphon.enums import BuiltinNodeTypes from models import Account, App, Tenant, TenantAccountJoin from models.account import TenantAccountRole from models.model import AppMode, IconType -from tests.unit_tests.core.ops.trace_fixtures import ( +from tests.helpers.unit.core.ops.trace_fixtures import ( message_trace_info, tool_trace_info, workflow_node_execution, diff --git a/api/providers/trace/trace-tencent/tests/unit_tests/tencent_trace/test_tencent_trace_utils.py b/api/providers/trace/trace-tencent/tests/unit_tests/src/dify_trace_tencent/test_utils.py similarity index 100% rename from api/providers/trace/trace-tencent/tests/unit_tests/tencent_trace/test_tencent_trace_utils.py rename to api/providers/trace/trace-tencent/tests/unit_tests/src/dify_trace_tencent/test_utils.py diff --git a/api/providers/trace/trace-weave/tests/unit_tests/test_config_entity.py b/api/providers/trace/trace-weave/tests/unit_tests/src/dify_trace_weave/test_config.py similarity index 100% rename from api/providers/trace/trace-weave/tests/unit_tests/test_config_entity.py rename to api/providers/trace/trace-weave/tests/unit_tests/src/dify_trace_weave/test_config.py diff --git a/api/providers/trace/trace-weave/tests/unit_tests/weave_trace/test_weave_trace.py b/api/providers/trace/trace-weave/tests/unit_tests/src/dify_trace_weave/test_weave_trace.py similarity index 100% rename from api/providers/trace/trace-weave/tests/unit_tests/weave_trace/test_weave_trace.py rename to api/providers/trace/trace-weave/tests/unit_tests/src/dify_trace_weave/test_weave_trace.py diff --git a/api/providers/vdb/vdb-alibabacloud-mysql/tests/unit_tests/test_alibabacloud_mysql_vector.py b/api/providers/vdb/vdb-alibabacloud-mysql/tests/unit_tests/src/dify_vdb_alibabacloud_mysql/test_alibabacloud_mysql_vector.py similarity index 89% rename from api/providers/vdb/vdb-alibabacloud-mysql/tests/unit_tests/test_alibabacloud_mysql_vector.py rename to api/providers/vdb/vdb-alibabacloud-mysql/tests/unit_tests/src/dify_vdb_alibabacloud_mysql/test_alibabacloud_mysql_vector.py index e699fb27135..f0184c549f9 100644 --- a/api/providers/vdb/vdb-alibabacloud-mysql/tests/unit_tests/test_alibabacloud_mysql_vector.py +++ b/api/providers/vdb/vdb-alibabacloud-mysql/tests/unit_tests/src/dify_vdb_alibabacloud_mysql/test_alibabacloud_mysql_vector.py @@ -3,13 +3,78 @@ import unittest from typing import override from unittest.mock import MagicMock, patch +import dify_vdb_alibabacloud_mysql.alibabacloud_mysql_vector as alibaba_module import pytest from dify_vdb_alibabacloud_mysql.alibabacloud_mysql_vector import ( AlibabaCloudMySQLVector, AlibabaCloudMySQLVectorConfig, + AlibabaCloudMySQLVectorFactory, ) from core.rag.models.document import Document +from models.dataset import Dataset + + +def test_validate_distance_function_accepts_supported_values(): + factory = AlibabaCloudMySQLVectorFactory() + + assert factory._validate_distance_function("cosine") == "cosine" + assert factory._validate_distance_function("euclidean") == "euclidean" + + +def test_validate_distance_function_rejects_unsupported_values(): + factory = AlibabaCloudMySQLVectorFactory() + + with pytest.raises(ValueError, match="Invalid distance function"): + factory._validate_distance_function("dot_product") + + +def test_factory_init_vector_uses_existing_index_struct_class_prefix(monkeypatch: pytest.MonkeyPatch): + factory = AlibabaCloudMySQLVectorFactory() + dataset = Dataset( + id="dataset-1", index_struct=json.dumps({"vector_store": {"class_prefix": "existing_collection"}}) + ) + + monkeypatch.setattr(alibaba_module.dify_config, "ALIBABACLOUD_MYSQL_HOST", "host") + monkeypatch.setattr(alibaba_module.dify_config, "ALIBABACLOUD_MYSQL_PORT", 3306) + monkeypatch.setattr(alibaba_module.dify_config, "ALIBABACLOUD_MYSQL_USER", "user") + monkeypatch.setattr(alibaba_module.dify_config, "ALIBABACLOUD_MYSQL_PASSWORD", "password") + monkeypatch.setattr(alibaba_module.dify_config, "ALIBABACLOUD_MYSQL_DATABASE", "db") + monkeypatch.setattr(alibaba_module.dify_config, "ALIBABACLOUD_MYSQL_MAX_CONNECTION", 5) + monkeypatch.setattr(alibaba_module.dify_config, "ALIBABACLOUD_MYSQL_CHARSET", "utf8mb4") + monkeypatch.setattr(alibaba_module.dify_config, "ALIBABACLOUD_MYSQL_DISTANCE_FUNCTION", "cosine") + monkeypatch.setattr(alibaba_module.dify_config, "ALIBABACLOUD_MYSQL_HNSW_M", 6) + + with patch.object(alibaba_module, "AlibabaCloudMySQLVector", return_value="vector") as vector_cls: + result = factory.init_vector(dataset, attributes=[], embeddings=MagicMock(), session=None) + + assert result == "vector" + assert vector_cls.call_args.kwargs["collection_name"] == "existing_collection" + + +def test_factory_init_vector_generates_collection_name_when_index_struct_is_missing(monkeypatch: pytest.MonkeyPatch): + factory = AlibabaCloudMySQLVectorFactory() + dataset = Dataset(id="dataset-2") + + monkeypatch.setattr(alibaba_module.Dataset, "gen_collection_name_by_id", lambda dataset_id: f"COL_{dataset_id}") + monkeypatch.setattr(alibaba_module.dify_config, "ALIBABACLOUD_MYSQL_HOST", "host") + monkeypatch.setattr(alibaba_module.dify_config, "ALIBABACLOUD_MYSQL_PORT", 3306) + monkeypatch.setattr(alibaba_module.dify_config, "ALIBABACLOUD_MYSQL_USER", "user") + monkeypatch.setattr(alibaba_module.dify_config, "ALIBABACLOUD_MYSQL_PASSWORD", "password") + monkeypatch.setattr(alibaba_module.dify_config, "ALIBABACLOUD_MYSQL_DATABASE", "db") + monkeypatch.setattr(alibaba_module.dify_config, "ALIBABACLOUD_MYSQL_MAX_CONNECTION", 5) + monkeypatch.setattr(alibaba_module.dify_config, "ALIBABACLOUD_MYSQL_CHARSET", "utf8mb4") + monkeypatch.setattr(alibaba_module.dify_config, "ALIBABACLOUD_MYSQL_DISTANCE_FUNCTION", "euclidean") + monkeypatch.setattr(alibaba_module.dify_config, "ALIBABACLOUD_MYSQL_HNSW_M", 12) + + with patch.object(alibaba_module, "AlibabaCloudMySQLVector", return_value="vector") as vector_cls: + result = factory.init_vector(dataset, attributes=[], embeddings=MagicMock(), session=None) + + assert result == "vector" + vector_cls.assert_called_once() + assert vector_cls.call_args.kwargs["collection_name"] == "COL_dataset-2" + assert dataset.index_struct is not None + try: from mysql.connector import Error as MySQLError diff --git a/api/providers/vdb/vdb-alibabacloud-mysql/tests/unit_tests/test_alibabacloud_mysql_factory.py b/api/providers/vdb/vdb-alibabacloud-mysql/tests/unit_tests/test_alibabacloud_mysql_factory.py deleted file mode 100644 index 51902fd7445..00000000000 --- a/api/providers/vdb/vdb-alibabacloud-mysql/tests/unit_tests/test_alibabacloud_mysql_factory.py +++ /dev/null @@ -1,69 +0,0 @@ -import json -from unittest.mock import MagicMock, patch - -import dify_vdb_alibabacloud_mysql.alibabacloud_mysql_vector as alibaba_module -import pytest -from dify_vdb_alibabacloud_mysql.alibabacloud_mysql_vector import AlibabaCloudMySQLVectorFactory - -from models.dataset import Dataset - - -def test_validate_distance_function_accepts_supported_values(): - factory = AlibabaCloudMySQLVectorFactory() - - assert factory._validate_distance_function("cosine") == "cosine" - assert factory._validate_distance_function("euclidean") == "euclidean" - - -def test_validate_distance_function_rejects_unsupported_values(): - factory = AlibabaCloudMySQLVectorFactory() - - with pytest.raises(ValueError, match="Invalid distance function"): - factory._validate_distance_function("dot_product") - - -def test_factory_init_vector_uses_existing_index_struct_class_prefix(monkeypatch: pytest.MonkeyPatch): - factory = AlibabaCloudMySQLVectorFactory() - dataset = Dataset( - id="dataset-1", index_struct=json.dumps({"vector_store": {"class_prefix": "existing_collection"}}) - ) - - monkeypatch.setattr(alibaba_module.dify_config, "ALIBABACLOUD_MYSQL_HOST", "host") - monkeypatch.setattr(alibaba_module.dify_config, "ALIBABACLOUD_MYSQL_PORT", 3306) - monkeypatch.setattr(alibaba_module.dify_config, "ALIBABACLOUD_MYSQL_USER", "user") - monkeypatch.setattr(alibaba_module.dify_config, "ALIBABACLOUD_MYSQL_PASSWORD", "password") - monkeypatch.setattr(alibaba_module.dify_config, "ALIBABACLOUD_MYSQL_DATABASE", "db") - monkeypatch.setattr(alibaba_module.dify_config, "ALIBABACLOUD_MYSQL_MAX_CONNECTION", 5) - monkeypatch.setattr(alibaba_module.dify_config, "ALIBABACLOUD_MYSQL_CHARSET", "utf8mb4") - monkeypatch.setattr(alibaba_module.dify_config, "ALIBABACLOUD_MYSQL_DISTANCE_FUNCTION", "cosine") - monkeypatch.setattr(alibaba_module.dify_config, "ALIBABACLOUD_MYSQL_HNSW_M", 6) - - with patch.object(alibaba_module, "AlibabaCloudMySQLVector", return_value="vector") as vector_cls: - result = factory.init_vector(dataset, attributes=[], embeddings=MagicMock(), session=None) - - assert result == "vector" - assert vector_cls.call_args.kwargs["collection_name"] == "existing_collection" - - -def test_factory_init_vector_generates_collection_name_when_index_struct_is_missing(monkeypatch: pytest.MonkeyPatch): - factory = AlibabaCloudMySQLVectorFactory() - dataset = Dataset(id="dataset-2") - - monkeypatch.setattr(alibaba_module.Dataset, "gen_collection_name_by_id", lambda dataset_id: f"COL_{dataset_id}") - monkeypatch.setattr(alibaba_module.dify_config, "ALIBABACLOUD_MYSQL_HOST", "host") - monkeypatch.setattr(alibaba_module.dify_config, "ALIBABACLOUD_MYSQL_PORT", 3306) - monkeypatch.setattr(alibaba_module.dify_config, "ALIBABACLOUD_MYSQL_USER", "user") - monkeypatch.setattr(alibaba_module.dify_config, "ALIBABACLOUD_MYSQL_PASSWORD", "password") - monkeypatch.setattr(alibaba_module.dify_config, "ALIBABACLOUD_MYSQL_DATABASE", "db") - monkeypatch.setattr(alibaba_module.dify_config, "ALIBABACLOUD_MYSQL_MAX_CONNECTION", 5) - monkeypatch.setattr(alibaba_module.dify_config, "ALIBABACLOUD_MYSQL_CHARSET", "utf8mb4") - monkeypatch.setattr(alibaba_module.dify_config, "ALIBABACLOUD_MYSQL_DISTANCE_FUNCTION", "euclidean") - monkeypatch.setattr(alibaba_module.dify_config, "ALIBABACLOUD_MYSQL_HNSW_M", 12) - - with patch.object(alibaba_module, "AlibabaCloudMySQLVector", return_value="vector") as vector_cls: - result = factory.init_vector(dataset, attributes=[], embeddings=MagicMock(), session=None) - - assert result == "vector" - vector_cls.assert_called_once() - assert vector_cls.call_args.kwargs["collection_name"] == "COL_dataset-2" - assert dataset.index_struct is not None diff --git a/api/providers/vdb/vdb-analyticdb/tests/integration_tests/test_analyticdb.py b/api/providers/vdb/vdb-analyticdb/tests/e2e_tests/test_analyticdb.py similarity index 100% rename from api/providers/vdb/vdb-analyticdb/tests/integration_tests/test_analyticdb.py rename to api/providers/vdb/vdb-analyticdb/tests/e2e_tests/test_analyticdb.py diff --git a/api/providers/vdb/vdb-analyticdb/tests/unit_tests/test_analyticdb_vector.py b/api/providers/vdb/vdb-analyticdb/tests/unit_tests/src/dify_vdb_analyticdb/test_analyticdb_vector.py similarity index 100% rename from api/providers/vdb/vdb-analyticdb/tests/unit_tests/test_analyticdb_vector.py rename to api/providers/vdb/vdb-analyticdb/tests/unit_tests/src/dify_vdb_analyticdb/test_analyticdb_vector.py diff --git a/api/providers/vdb/vdb-analyticdb/tests/unit_tests/test_analyticdb_vector_openapi.py b/api/providers/vdb/vdb-analyticdb/tests/unit_tests/src/dify_vdb_analyticdb/test_analyticdb_vector_openapi.py similarity index 100% rename from api/providers/vdb/vdb-analyticdb/tests/unit_tests/test_analyticdb_vector_openapi.py rename to api/providers/vdb/vdb-analyticdb/tests/unit_tests/src/dify_vdb_analyticdb/test_analyticdb_vector_openapi.py diff --git a/api/providers/vdb/vdb-analyticdb/tests/unit_tests/test_analyticdb_vector_sql.py b/api/providers/vdb/vdb-analyticdb/tests/unit_tests/src/dify_vdb_analyticdb/test_analyticdb_vector_sql.py similarity index 100% rename from api/providers/vdb/vdb-analyticdb/tests/unit_tests/test_analyticdb_vector_sql.py rename to api/providers/vdb/vdb-analyticdb/tests/unit_tests/src/dify_vdb_analyticdb/test_analyticdb_vector_sql.py diff --git a/api/providers/vdb/vdb-baidu/tests/integration_tests/conftest.py b/api/providers/vdb/vdb-baidu/tests/e2e_tests/conftest.py similarity index 100% rename from api/providers/vdb/vdb-baidu/tests/integration_tests/conftest.py rename to api/providers/vdb/vdb-baidu/tests/e2e_tests/conftest.py diff --git a/api/providers/vdb/vdb-baidu/tests/integration_tests/test_baidu.py b/api/providers/vdb/vdb-baidu/tests/e2e_tests/test_baidu.py similarity index 100% rename from api/providers/vdb/vdb-baidu/tests/integration_tests/test_baidu.py rename to api/providers/vdb/vdb-baidu/tests/e2e_tests/test_baidu.py diff --git a/api/providers/vdb/vdb-baidu/tests/unit_tests/test_baidu_vector.py b/api/providers/vdb/vdb-baidu/tests/unit_tests/src/dify_vdb_baidu/test_baidu_vector.py similarity index 100% rename from api/providers/vdb/vdb-baidu/tests/unit_tests/test_baidu_vector.py rename to api/providers/vdb/vdb-baidu/tests/unit_tests/src/dify_vdb_baidu/test_baidu_vector.py diff --git a/api/providers/vdb/vdb-chroma/tests/integration_tests/test_chroma.py b/api/providers/vdb/vdb-chroma/tests/e2e_tests/test_chroma.py similarity index 100% rename from api/providers/vdb/vdb-chroma/tests/integration_tests/test_chroma.py rename to api/providers/vdb/vdb-chroma/tests/e2e_tests/test_chroma.py diff --git a/api/providers/vdb/vdb-chroma/tests/unit_tests/test_chroma_vector.py b/api/providers/vdb/vdb-chroma/tests/unit_tests/src/dify_vdb_chroma/test_chroma_vector.py similarity index 100% rename from api/providers/vdb/vdb-chroma/tests/unit_tests/test_chroma_vector.py rename to api/providers/vdb/vdb-chroma/tests/unit_tests/src/dify_vdb_chroma/test_chroma_vector.py diff --git a/api/providers/vdb/vdb-clickzetta/tests/README.md b/api/providers/vdb/vdb-clickzetta/tests/README.md index 61b7b6c5fb5..52f1c05748a 100644 --- a/api/providers/vdb/vdb-clickzetta/tests/README.md +++ b/api/providers/vdb/vdb-clickzetta/tests/README.md @@ -17,7 +17,7 @@ export CLICKZETTA_SCHEMA=dify The CI invocation from the repository root is: ```bash -uv run --project api pytest api/providers/vdb/vdb-clickzetta/tests/integration_tests/ +uv run --project api pytest api/providers/vdb/vdb-clickzetta/tests/e2e_tests/ ``` For local verification, run the provider unit tests: diff --git a/api/providers/vdb/vdb-clickzetta/tests/integration_tests/test_clickzetta.py b/api/providers/vdb/vdb-clickzetta/tests/e2e_tests/test_clickzetta.py similarity index 100% rename from api/providers/vdb/vdb-clickzetta/tests/integration_tests/test_clickzetta.py rename to api/providers/vdb/vdb-clickzetta/tests/e2e_tests/test_clickzetta.py diff --git a/api/providers/vdb/vdb-clickzetta/tests/integration_tests/test_docker_integration.py b/api/providers/vdb/vdb-clickzetta/tests/e2e_tests/test_docker_integration.py similarity index 100% rename from api/providers/vdb/vdb-clickzetta/tests/integration_tests/test_docker_integration.py rename to api/providers/vdb/vdb-clickzetta/tests/e2e_tests/test_docker_integration.py diff --git a/api/providers/vdb/vdb-clickzetta/tests/unit_tests/test_clickzetta_vector.py b/api/providers/vdb/vdb-clickzetta/tests/unit_tests/src/dify_vdb_clickzetta/test_clickzetta_vector.py similarity index 100% rename from api/providers/vdb/vdb-clickzetta/tests/unit_tests/test_clickzetta_vector.py rename to api/providers/vdb/vdb-clickzetta/tests/unit_tests/src/dify_vdb_clickzetta/test_clickzetta_vector.py diff --git a/api/providers/vdb/vdb-couchbase/tests/integration_tests/test_couchbase.py b/api/providers/vdb/vdb-couchbase/tests/e2e_tests/test_couchbase.py similarity index 100% rename from api/providers/vdb/vdb-couchbase/tests/integration_tests/test_couchbase.py rename to api/providers/vdb/vdb-couchbase/tests/e2e_tests/test_couchbase.py diff --git a/api/providers/vdb/vdb-couchbase/tests/unit_tests/test_couchbase_vector.py b/api/providers/vdb/vdb-couchbase/tests/unit_tests/src/dify_vdb_couchbase/test_couchbase_vector.py similarity index 100% rename from api/providers/vdb/vdb-couchbase/tests/unit_tests/test_couchbase_vector.py rename to api/providers/vdb/vdb-couchbase/tests/unit_tests/src/dify_vdb_couchbase/test_couchbase_vector.py diff --git a/api/providers/vdb/vdb-elasticsearch/tests/integration_tests/test_elasticsearch.py b/api/providers/vdb/vdb-elasticsearch/tests/e2e_tests/test_elasticsearch.py similarity index 100% rename from api/providers/vdb/vdb-elasticsearch/tests/integration_tests/test_elasticsearch.py rename to api/providers/vdb/vdb-elasticsearch/tests/e2e_tests/test_elasticsearch.py diff --git a/api/providers/vdb/vdb-elasticsearch/tests/unit_tests/test_elasticsearch_ja_vector.py b/api/providers/vdb/vdb-elasticsearch/tests/unit_tests/src/dify_vdb_elasticsearch/test_elasticsearch_ja_vector.py similarity index 100% rename from api/providers/vdb/vdb-elasticsearch/tests/unit_tests/test_elasticsearch_ja_vector.py rename to api/providers/vdb/vdb-elasticsearch/tests/unit_tests/src/dify_vdb_elasticsearch/test_elasticsearch_ja_vector.py diff --git a/api/providers/vdb/vdb-elasticsearch/tests/unit_tests/test_elasticsearch_vector.py b/api/providers/vdb/vdb-elasticsearch/tests/unit_tests/src/dify_vdb_elasticsearch/test_elasticsearch_vector.py similarity index 100% rename from api/providers/vdb/vdb-elasticsearch/tests/unit_tests/test_elasticsearch_vector.py rename to api/providers/vdb/vdb-elasticsearch/tests/unit_tests/src/dify_vdb_elasticsearch/test_elasticsearch_vector.py diff --git a/api/providers/vdb/vdb-hologres/tests/integration_tests/conftest.py b/api/providers/vdb/vdb-hologres/tests/e2e_tests/conftest.py similarity index 100% rename from api/providers/vdb/vdb-hologres/tests/integration_tests/conftest.py rename to api/providers/vdb/vdb-hologres/tests/e2e_tests/conftest.py diff --git a/api/providers/vdb/vdb-hologres/tests/integration_tests/test_hologres.py b/api/providers/vdb/vdb-hologres/tests/e2e_tests/test_hologres.py similarity index 100% rename from api/providers/vdb/vdb-hologres/tests/integration_tests/test_hologres.py rename to api/providers/vdb/vdb-hologres/tests/e2e_tests/test_hologres.py diff --git a/api/providers/vdb/vdb-hologres/tests/unit_tests/test_hologres_vector.py b/api/providers/vdb/vdb-hologres/tests/unit_tests/src/dify_vdb_hologres/test_hologres_vector.py similarity index 100% rename from api/providers/vdb/vdb-hologres/tests/unit_tests/test_hologres_vector.py rename to api/providers/vdb/vdb-hologres/tests/unit_tests/src/dify_vdb_hologres/test_hologres_vector.py diff --git a/api/providers/vdb/vdb-huawei-cloud/tests/integration_tests/conftest.py b/api/providers/vdb/vdb-huawei-cloud/tests/e2e_tests/conftest.py similarity index 100% rename from api/providers/vdb/vdb-huawei-cloud/tests/integration_tests/conftest.py rename to api/providers/vdb/vdb-huawei-cloud/tests/e2e_tests/conftest.py diff --git a/api/providers/vdb/vdb-huawei-cloud/tests/integration_tests/test_huawei_cloud.py b/api/providers/vdb/vdb-huawei-cloud/tests/e2e_tests/test_huawei_cloud.py similarity index 100% rename from api/providers/vdb/vdb-huawei-cloud/tests/integration_tests/test_huawei_cloud.py rename to api/providers/vdb/vdb-huawei-cloud/tests/e2e_tests/test_huawei_cloud.py diff --git a/api/providers/vdb/vdb-huawei-cloud/tests/unit_tests/test_huawei_cloud_vector.py b/api/providers/vdb/vdb-huawei-cloud/tests/unit_tests/src/dify_vdb_huawei_cloud/test_huawei_cloud_vector.py similarity index 100% rename from api/providers/vdb/vdb-huawei-cloud/tests/unit_tests/test_huawei_cloud_vector.py rename to api/providers/vdb/vdb-huawei-cloud/tests/unit_tests/src/dify_vdb_huawei_cloud/test_huawei_cloud_vector.py diff --git a/api/providers/vdb/vdb-iris/tests/integration_tests/test_iris.py b/api/providers/vdb/vdb-iris/tests/e2e_tests/test_iris.py similarity index 100% rename from api/providers/vdb/vdb-iris/tests/integration_tests/test_iris.py rename to api/providers/vdb/vdb-iris/tests/e2e_tests/test_iris.py diff --git a/api/providers/vdb/vdb-iris/tests/unit_tests/test_iris_vector.py b/api/providers/vdb/vdb-iris/tests/unit_tests/src/dify_vdb_iris/test_iris_vector.py similarity index 100% rename from api/providers/vdb/vdb-iris/tests/unit_tests/test_iris_vector.py rename to api/providers/vdb/vdb-iris/tests/unit_tests/src/dify_vdb_iris/test_iris_vector.py diff --git a/api/providers/vdb/vdb-lindorm/tests/integration_tests/test_lindorm.py b/api/providers/vdb/vdb-lindorm/tests/e2e_tests/test_lindorm.py similarity index 100% rename from api/providers/vdb/vdb-lindorm/tests/integration_tests/test_lindorm.py rename to api/providers/vdb/vdb-lindorm/tests/e2e_tests/test_lindorm.py diff --git a/api/providers/vdb/vdb-lindorm/tests/unit_tests/test_lindorm_vector.py b/api/providers/vdb/vdb-lindorm/tests/unit_tests/src/dify_vdb_lindorm/test_lindorm_vector.py similarity index 100% rename from api/providers/vdb/vdb-lindorm/tests/unit_tests/test_lindorm_vector.py rename to api/providers/vdb/vdb-lindorm/tests/unit_tests/src/dify_vdb_lindorm/test_lindorm_vector.py diff --git a/api/providers/vdb/vdb-matrixone/tests/integration_tests/test_matrixone.py b/api/providers/vdb/vdb-matrixone/tests/e2e_tests/test_matrixone.py similarity index 100% rename from api/providers/vdb/vdb-matrixone/tests/integration_tests/test_matrixone.py rename to api/providers/vdb/vdb-matrixone/tests/e2e_tests/test_matrixone.py diff --git a/api/providers/vdb/vdb-matrixone/tests/unit_tests/test_matrixone_vector.py b/api/providers/vdb/vdb-matrixone/tests/unit_tests/src/dify_vdb_matrixone/test_matrixone_vector.py similarity index 100% rename from api/providers/vdb/vdb-matrixone/tests/unit_tests/test_matrixone_vector.py rename to api/providers/vdb/vdb-matrixone/tests/unit_tests/src/dify_vdb_matrixone/test_matrixone_vector.py diff --git a/api/providers/vdb/vdb-milvus/tests/integration_tests/test_milvus.py b/api/providers/vdb/vdb-milvus/tests/e2e_tests/test_milvus.py similarity index 100% rename from api/providers/vdb/vdb-milvus/tests/integration_tests/test_milvus.py rename to api/providers/vdb/vdb-milvus/tests/e2e_tests/test_milvus.py diff --git a/api/providers/vdb/vdb-milvus/tests/unit_tests/test_milvus.py b/api/providers/vdb/vdb-milvus/tests/unit_tests/src/dify_vdb_milvus/test_milvus_vector.py similarity index 100% rename from api/providers/vdb/vdb-milvus/tests/unit_tests/test_milvus.py rename to api/providers/vdb/vdb-milvus/tests/unit_tests/src/dify_vdb_milvus/test_milvus_vector.py diff --git a/api/providers/vdb/vdb-myscale/tests/integration_tests/test_myscale.py b/api/providers/vdb/vdb-myscale/tests/e2e_tests/test_myscale.py similarity index 100% rename from api/providers/vdb/vdb-myscale/tests/integration_tests/test_myscale.py rename to api/providers/vdb/vdb-myscale/tests/e2e_tests/test_myscale.py diff --git a/api/providers/vdb/vdb-myscale/tests/unit_tests/test_myscale_vector.py b/api/providers/vdb/vdb-myscale/tests/unit_tests/src/dify_vdb_myscale/test_myscale_vector.py similarity index 100% rename from api/providers/vdb/vdb-myscale/tests/unit_tests/test_myscale_vector.py rename to api/providers/vdb/vdb-myscale/tests/unit_tests/src/dify_vdb_myscale/test_myscale_vector.py diff --git a/api/providers/vdb/vdb-oceanbase/tests/integration_tests/bench_oceanbase.py b/api/providers/vdb/vdb-oceanbase/tests/e2e_tests/bench_oceanbase.py similarity index 100% rename from api/providers/vdb/vdb-oceanbase/tests/integration_tests/bench_oceanbase.py rename to api/providers/vdb/vdb-oceanbase/tests/e2e_tests/bench_oceanbase.py diff --git a/api/providers/vdb/vdb-oceanbase/tests/integration_tests/test_oceanbase.py b/api/providers/vdb/vdb-oceanbase/tests/e2e_tests/test_oceanbase.py similarity index 100% rename from api/providers/vdb/vdb-oceanbase/tests/integration_tests/test_oceanbase.py rename to api/providers/vdb/vdb-oceanbase/tests/e2e_tests/test_oceanbase.py diff --git a/api/providers/vdb/vdb-oceanbase/tests/unit_tests/test_oceanbase_vector.py b/api/providers/vdb/vdb-oceanbase/tests/unit_tests/src/dify_vdb_oceanbase/test_oceanbase_vector.py similarity index 100% rename from api/providers/vdb/vdb-oceanbase/tests/unit_tests/test_oceanbase_vector.py rename to api/providers/vdb/vdb-oceanbase/tests/unit_tests/src/dify_vdb_oceanbase/test_oceanbase_vector.py diff --git a/api/providers/vdb/vdb-opengauss/tests/integration_tests/test_opengauss.py b/api/providers/vdb/vdb-opengauss/tests/e2e_tests/test_opengauss.py similarity index 100% rename from api/providers/vdb/vdb-opengauss/tests/integration_tests/test_opengauss.py rename to api/providers/vdb/vdb-opengauss/tests/e2e_tests/test_opengauss.py diff --git a/api/providers/vdb/vdb-opengauss/tests/unit_tests/test_opengauss.py b/api/providers/vdb/vdb-opengauss/tests/unit_tests/src/dify_vdb_opengauss/test_opengauss.py similarity index 100% rename from api/providers/vdb/vdb-opengauss/tests/unit_tests/test_opengauss.py rename to api/providers/vdb/vdb-opengauss/tests/unit_tests/src/dify_vdb_opengauss/test_opengauss.py diff --git a/api/providers/vdb/vdb-opensearch/tests/unit_tests/src/dify_vdb_opensearch/test_opensearch_vector.py b/api/providers/vdb/vdb-opensearch/tests/unit_tests/src/dify_vdb_opensearch/test_opensearch_vector.py new file mode 100644 index 00000000000..264c2d4803f --- /dev/null +++ b/api/providers/vdb/vdb-opensearch/tests/unit_tests/src/dify_vdb_opensearch/test_opensearch_vector.py @@ -0,0 +1,690 @@ +import importlib +import json +import sys +import types +from types import SimpleNamespace +from unittest.mock import MagicMock, patch + +import pytest +from pydantic import ValidationError + +from core.rag.datasource.vdb.field import Field +from core.rag.models.document import Document +from extensions import ext_redis +from models.dataset import Dataset + + +def _opensearch_build_fake_opensearch_modules(): + """Build fake opensearchpy modules to avoid the ``from events import Events`` + namespace collision (opensearch-py #756).""" + opensearchpy = types.ModuleType("opensearchpy") + opensearchpy_helpers = types.ModuleType("opensearchpy.helpers") + + class BulkIndexError(Exception): + def __init__(self, errors): + super().__init__("bulk error") + self.errors = errors + + class Urllib3AWSV4SignerAuth: + def __init__(self, credentials, region, service): + self.credentials = credentials + self.region = region + self.service = service + + class Urllib3HttpConnection: + pass + + class _IndicesClient: + def __init__(self): + self.exists = MagicMock(return_value=False) + self.create = MagicMock() + self.delete = MagicMock() + + class OpenSearch: + def __init__(self, **kwargs): + self.kwargs = kwargs + self.indices = _IndicesClient() + self.search = MagicMock(return_value={"hits": {"hits": []}}) + self.get = MagicMock() + + helpers = SimpleNamespace(bulk=MagicMock()) + + opensearchpy.OpenSearch = OpenSearch + opensearchpy.Urllib3AWSV4SignerAuth = Urllib3AWSV4SignerAuth + opensearchpy.Urllib3HttpConnection = Urllib3HttpConnection + opensearchpy.helpers = helpers + opensearchpy_helpers.BulkIndexError = BulkIndexError + + return { + "opensearchpy": opensearchpy, + "opensearchpy.helpers": opensearchpy_helpers, + } + + +def _opensearch_config(module, **overrides): + values = { + "host": "localhost", + "port": 9200, + "secure": False, + "user": "admin", + "password": "password", + } + values.update(overrides) + return module.OpenSearchConfig.model_validate(values) + + +def get_example_text() -> str: + return "This is a sample text for testing purposes." + + +class TestOpensearchCases: + @staticmethod + @pytest.fixture + def opensearch_module(monkeypatch: pytest.MonkeyPatch): + for name, module in _opensearch_build_fake_opensearch_modules().items(): + monkeypatch.setitem(sys.modules, name, module) + + import dify_vdb_opensearch.opensearch_vector as module + + return importlib.reload(module) + + class TestOpenSearchConfig: + def test_to_opensearch_params(self, opensearch_module): + config = _opensearch_config(opensearch_module, secure=True) + params = config.to_opensearch_params() + + assert params["hosts"] == [{"host": "localhost", "port": 9200}] + assert params["use_ssl"] is True + assert params["verify_certs"] is True + assert params["connection_class"].__name__ == "Urllib3HttpConnection" + assert params["http_auth"] == ("admin", "password") + + def test_to_opensearch_params_with_aws_managed_iam(self, opensearch_module, monkeypatch: pytest.MonkeyPatch): + class _Session: + def get_credentials(self): + return "creds" + + boto3 = types.ModuleType("boto3") + boto3.Session = _Session + monkeypatch.setitem(sys.modules, "boto3", boto3) + + config = _opensearch_config( + opensearch_module, + secure=True, + auth_method="aws_managed_iam", + aws_region="ap-southeast-2", + aws_service="aoss", + host="aoss-endpoint.ap-southeast-2.aoss.amazonaws.com", + port=9201, + ) + params = config.to_opensearch_params() + + assert params["hosts"] == [{"host": "aoss-endpoint.ap-southeast-2.aoss.amazonaws.com", "port": 9201}] + assert params["use_ssl"] is True + assert params["verify_certs"] is True + assert params["connection_class"].__name__ == "Urllib3HttpConnection" + assert params["http_auth"].credentials == "creds" + assert params["http_auth"].region == "ap-southeast-2" + assert params["http_auth"].service == "aoss" + + class TestOpenSearchVector: + COLLECTION_NAME = "test_collection" + EXAMPLE_DOC_ID = "example_doc_id" + + def _make_vector(self, module): + vector = module.OpenSearchVector(self.COLLECTION_NAME, _opensearch_config(module)) + vector._client = MagicMock() + return vector + + @pytest.mark.parametrize( + ("search_response", "expected_length", "expected_doc_id"), + [ + ( + { + "hits": { + "total": {"value": 1}, + "hits": [ + { + "_source": { + "page_content": get_example_text(), + "metadata": {"document_id": "example_doc_id"}, + } + } + ], + } + }, + 1, + "example_doc_id", + ), + ({"hits": {"total": {"value": 0}, "hits": []}}, 0, None), + ], + ) + def test_search_by_full_text(self, opensearch_module, search_response, expected_length, expected_doc_id): + vector = self._make_vector(opensearch_module) + vector._client.search.return_value = search_response + + hits = vector.search_by_full_text(query=get_example_text()) + assert len(hits) == expected_length + if expected_length > 0: + assert hits[0].metadata["document_id"] == expected_doc_id + + def test_search_by_vector(self, opensearch_module): + vector = self._make_vector(opensearch_module) + query_vector = [0.1] * 128 + mock_response = { + "hits": { + "total": {"value": 1}, + "hits": [ + { + "_source": { + Field.CONTENT_KEY: get_example_text(), + Field.METADATA_KEY: {"document_id": self.EXAMPLE_DOC_ID}, + }, + "_score": 1.0, + } + ], + } + } + vector._client.search.return_value = mock_response + + hits = vector.search_by_vector(query_vector=query_vector) + + assert len(hits) > 0 + assert hits[0].metadata["document_id"] == self.EXAMPLE_DOC_ID + + def test_get_ids_by_metadata_field(self, opensearch_module): + vector = self._make_vector(opensearch_module) + mock_response = {"hits": {"total": {"value": 1}, "hits": [{"_id": "mock_id"}]}} + vector._client.search.return_value = mock_response + + doc = Document(page_content="Test content", metadata={"document_id": self.EXAMPLE_DOC_ID}) + embedding = [0.1] * 128 + + opensearch_module.helpers.bulk.reset_mock() + vector.add_texts([doc], [embedding]) + + ids = vector.get_ids_by_metadata_field(key="document_id", value=self.EXAMPLE_DOC_ID) + assert len(ids) == 1 + assert ids[0] == "mock_id" + + def test_add_texts(self, opensearch_module): + vector = self._make_vector(opensearch_module) + vector._client.index.return_value = {"result": "created"} + + doc = Document(page_content="Test content", metadata={"document_id": self.EXAMPLE_DOC_ID}) + embedding = [0.1] * 128 + + opensearch_module.helpers.bulk.reset_mock() + vector.add_texts([doc], [embedding]) + + mock_response = {"hits": {"total": {"value": 1}, "hits": [{"_id": "mock_id"}]}} + vector._client.search.return_value = mock_response + + ids = vector.get_ids_by_metadata_field(key="document_id", value=self.EXAMPLE_DOC_ID) + assert len(ids) == 1 + assert ids[0] == "mock_id" + + def test_delete_nonexistent_index(self, opensearch_module): + """ignore_unavailable=True handles non-existent indices gracefully.""" + vector = self._make_vector(opensearch_module) + vector.delete() + + vector._client.indices.delete.assert_called_once_with( + index=self.COLLECTION_NAME.lower(), ignore_unavailable=True + ) + + def test_delete_existing_index(self, opensearch_module): + vector = self._make_vector(opensearch_module) + vector.delete() + + vector._client.indices.delete.assert_called_once_with( + index=self.COLLECTION_NAME.lower(), ignore_unavailable=True + ) + + @staticmethod + @pytest.fixture(scope="module") + def setup_mock_redis(): + ext_redis.redis_client.get = MagicMock(return_value=None) + ext_redis.redis_client.set = MagicMock(return_value=None) + + mock_redis_lock = MagicMock() + mock_redis_lock.__enter__ = MagicMock() + mock_redis_lock.__exit__ = MagicMock() + ext_redis.redis_client.lock = MagicMock(return_value=mock_redis_lock) + + @pytest.mark.usefixtures("setup_mock_redis") + class TestOpenSearchVectorWithRedis: + COLLECTION_NAME = "test_collection" + EXAMPLE_DOC_ID = "example_doc_id" + + def _make_vector(self, module): + vector = module.OpenSearchVector(self.COLLECTION_NAME, _opensearch_config(module)) + vector._client = MagicMock() + return vector + + def test_search_by_full_text(self, opensearch_module): + vector = self._make_vector(opensearch_module) + search_response = { + "hits": { + "total": {"value": 1}, + "hits": [ + {"_source": {"page_content": get_example_text(), "metadata": {"document_id": "example_doc_id"}}} + ], + } + } + vector._client.search.return_value = search_response + + hits = vector.search_by_full_text(query=get_example_text()) + assert len(hits) == 1 + assert hits[0].metadata["document_id"] == "example_doc_id" + + def test_get_ids_by_metadata_field(self, opensearch_module): + vector = self._make_vector(opensearch_module) + mock_response = {"hits": {"total": {"value": 1}, "hits": [{"_id": "mock_id"}]}} + vector._client.search.return_value = mock_response + + doc = Document(page_content="Test content", metadata={"document_id": self.EXAMPLE_DOC_ID}) + embedding = [0.1] * 128 + + opensearch_module.helpers.bulk.reset_mock() + vector.add_texts([doc], [embedding]) + + ids = vector.get_ids_by_metadata_field(key="document_id", value=self.EXAMPLE_DOC_ID) + assert len(ids) == 1 + assert ids[0] == "mock_id" + + def test_add_texts(self, opensearch_module): + vector = self._make_vector(opensearch_module) + vector._client.index.return_value = {"result": "created"} + + doc = Document(page_content="Test content", metadata={"document_id": self.EXAMPLE_DOC_ID}) + embedding = [0.1] * 128 + + opensearch_module.helpers.bulk.reset_mock() + vector.add_texts([doc], [embedding]) + + mock_response = {"hits": {"total": {"value": 1}, "hits": [{"_id": "mock_id"}]}} + vector._client.search.return_value = mock_response + + ids = vector.get_ids_by_metadata_field(key="document_id", value=self.EXAMPLE_DOC_ID) + assert len(ids) == 1 + assert ids[0] == "mock_id" + + def test_search_by_vector(self, opensearch_module): + vector = self._make_vector(opensearch_module) + query_vector = [0.1] * 128 + mock_response = { + "hits": { + "total": {"value": 1}, + "hits": [ + { + "_source": { + Field.CONTENT_KEY: get_example_text(), + Field.METADATA_KEY: {"document_id": self.EXAMPLE_DOC_ID}, + }, + "_score": 1.0, + } + ], + } + } + vector._client.search.return_value = mock_response + + hits = vector.search_by_vector(query_vector=query_vector) + assert len(hits) > 0 + assert hits[0].metadata["document_id"] == self.EXAMPLE_DOC_ID + + +# TODO(wylswz): There's a known issue with namespace collision +# https://github.com/langgenius/dify/issues/34732 +def _opensearch_vector_build_fake_opensearch_modules(): + opensearchpy = types.ModuleType("opensearchpy") + opensearchpy_helpers = types.ModuleType("opensearchpy.helpers") + + class BulkIndexError(Exception): + def __init__(self, errors): + super().__init__("bulk error") + self.errors = errors + + class Urllib3AWSV4SignerAuth: + def __init__(self, credentials, region, service): + self.credentials = credentials + self.region = region + self.service = service + + class Urllib3HttpConnection: + pass + + class _IndicesClient: + def __init__(self): + self.exists = MagicMock(return_value=False) + self.create = MagicMock() + self.delete = MagicMock() + + class OpenSearch: + def __init__(self, **kwargs): + self.kwargs = kwargs + self.indices = _IndicesClient() + self.search = MagicMock(return_value={"hits": {"hits": []}}) + self.get = MagicMock() + + helpers = SimpleNamespace(bulk=MagicMock()) + + opensearchpy.OpenSearch = OpenSearch + opensearchpy.Urllib3AWSV4SignerAuth = Urllib3AWSV4SignerAuth + opensearchpy.Urllib3HttpConnection = Urllib3HttpConnection + opensearchpy.helpers = helpers + opensearchpy_helpers.BulkIndexError = BulkIndexError + + return { + "opensearchpy": opensearchpy, + "opensearchpy.helpers": opensearchpy_helpers, + } + + +def _opensearch_vector_config(module, **overrides): + values = { + "host": "localhost", + "port": 9200, + "secure": True, + "verify_certs": True, + "auth_method": "basic", + "user": "admin", + "password": "secret", + } + values.update(overrides) + return module.OpenSearchConfig.model_validate(values) + + +class TestOpensearchVectorCases: + @staticmethod + @pytest.fixture + def opensearch_module(monkeypatch: pytest.MonkeyPatch): + for name, module in _opensearch_vector_build_fake_opensearch_modules().items(): + monkeypatch.setitem(sys.modules, name, module) + + import dify_vdb_opensearch.opensearch_vector as module + + return importlib.reload(module) + + @staticmethod + @pytest.mark.parametrize( + ("field", "value", "message"), + [ + ("host", "", "config OPENSEARCH_HOST is required"), + ("port", 0, "config OPENSEARCH_PORT is required"), + ], + ) + def test_config_validation_required_fields(opensearch_module, field, value, message): + values = _opensearch_vector_config(opensearch_module).model_dump() + values[field] = value + + with pytest.raises(ValidationError, match=message): + opensearch_module.OpenSearchConfig.model_validate(values) + + @staticmethod + def test_config_validation_for_aws_auth_and_https_fields(opensearch_module): + values = { + "host": "localhost", + "port": 9200, + "secure": True, + "verify_certs": True, + "auth_method": "aws_managed_iam", + "user": "admin", + "password": "secret", + } + with pytest.raises(ValidationError, match="OPENSEARCH_AWS_REGION"): + opensearch_module.OpenSearchConfig.model_validate(values) + + values = _opensearch_vector_config(opensearch_module).model_dump() + values["OPENSEARCH_SECURE"] = False + values["OPENSEARCH_VERIFY_CERTS"] = True + with pytest.raises(ValidationError, match="verify_certs=True requires secure"): + opensearch_module.OpenSearchConfig.model_validate(values) + + @staticmethod + def test_create_aws_managed_iam_auth(opensearch_module, monkeypatch: pytest.MonkeyPatch): + class _Session: + def get_credentials(self): + return "creds" + + boto3 = types.ModuleType("boto3") + boto3.Session = _Session + monkeypatch.setitem(sys.modules, "boto3", boto3) + + config = _opensearch_vector_config( + opensearch_module, + auth_method="aws_managed_iam", + aws_region="us-east-1", + aws_service="es", + ) + auth = config.create_aws_managed_iam_auth() + + assert auth.credentials == "creds" + assert auth.region == "us-east-1" + assert auth.service == "es" + + @staticmethod + def test_to_opensearch_params_supports_basic_and_aws(opensearch_module): + basic_params = _opensearch_vector_config(opensearch_module).to_opensearch_params() + assert basic_params["http_auth"] == ("admin", "secret") + + aws_config = _opensearch_vector_config( + opensearch_module, + auth_method="aws_managed_iam", + aws_region="us-west-2", + aws_service="es", + ) + with patch.object(opensearch_module.OpenSearchConfig, "create_aws_managed_iam_auth", return_value="iam-auth"): + aws_params = aws_config.to_opensearch_params() + + assert aws_params["http_auth"] == "iam-auth" + + @staticmethod + def test_init_and_create_delegate_calls(opensearch_module): + vector = opensearch_module.OpenSearchVector("Collection_1", _opensearch_vector_config(opensearch_module)) + vector.create_collection = MagicMock() + vector.add_texts = MagicMock() + docs = [Document(page_content="hello", metadata={"doc_id": "seg-1"})] + + vector.create(docs, [[0.1, 0.2]]) + + assert vector.get_type() == "opensearch" + vector.create_collection.assert_called_once_with([[0.1, 0.2]], [{"doc_id": "seg-1"}]) + vector.add_texts.assert_called_once_with(docs, [[0.1, 0.2]]) + + @staticmethod + def test_add_texts_supports_regular_and_aoss_clients(opensearch_module, monkeypatch: pytest.MonkeyPatch): + vector = opensearch_module.OpenSearchVector( + "Collection_1", _opensearch_vector_config(opensearch_module, aws_service="es") + ) + docs = [ + Document(page_content="a", metadata={"doc_id": "1"}), + Document(page_content="b", metadata={"doc_id": "2"}), + ] + + monkeypatch.setattr(opensearch_module, "uuid4", lambda: SimpleNamespace(hex="generated-id")) + opensearch_module.helpers.bulk.reset_mock() + vector.add_texts(docs, [[0.1], [0.2]]) + actions = opensearch_module.helpers.bulk.call_args.kwargs["actions"] + assert len(actions) == 2 + assert all("_id" in action for action in actions) + + vector._client_config.aws_service = "aoss" + opensearch_module.helpers.bulk.reset_mock() + vector.add_texts(docs, [[0.3], [0.4]]) + aoss_actions = opensearch_module.helpers.bulk.call_args.kwargs["actions"] + assert all("_id" not in action for action in aoss_actions) + + @staticmethod + def test_metadata_lookup_and_delete_by_metadata_field(opensearch_module): + vector = opensearch_module.OpenSearchVector("collection_1", _opensearch_vector_config(opensearch_module)) + vector._client.search.return_value = {"hits": {"hits": [{"_id": "id-1"}, {"_id": "id-2"}]}} + + assert vector.get_ids_by_metadata_field("document_id", "doc-1") == ["id-1", "id-2"] + + vector._client.search.return_value = {"hits": {"hits": []}} + assert vector.get_ids_by_metadata_field("document_id", "doc-1") is None + + vector.get_ids_by_metadata_field = MagicMock(return_value=["id-1"]) + vector.delete_by_ids = MagicMock() + vector.delete_by_metadata_field("document_id", "doc-1") + vector.delete_by_ids.assert_called_once_with(["id-1"]) + + @staticmethod + def test_delete_by_ids_branches_and_bulk_error_handling(opensearch_module): + vector = opensearch_module.OpenSearchVector("collection_1", _opensearch_vector_config(opensearch_module)) + opensearch_module.helpers.bulk.reset_mock() + vector._client.indices.exists.return_value = False + vector.delete_by_ids(["doc-1"]) + opensearch_module.helpers.bulk.assert_not_called() + + vector._client.indices.exists.return_value = True + vector.get_ids_by_metadata_field = MagicMock(side_effect=[["es-1"], None]) + vector.delete_by_ids(["doc-1", "doc-2"]) + opensearch_module.helpers.bulk.assert_called_once() + + opensearch_module.helpers.bulk.reset_mock() + vector.get_ids_by_metadata_field = MagicMock(return_value=["es-404"]) + opensearch_module.helpers.bulk.side_effect = opensearch_module.BulkIndexError( + [{"delete": {"status": 404, "_id": "es-404"}}] + ) + vector.delete_by_ids(["doc-404"]) + assert opensearch_module.helpers.bulk.call_count == 1 + + opensearch_module.helpers.bulk.side_effect = None + + @staticmethod + def test_delete_and_text_exists(opensearch_module): + vector = opensearch_module.OpenSearchVector("collection_1", _opensearch_vector_config(opensearch_module)) + vector.delete() + vector._client.indices.delete.assert_called_once_with(index="collection_1", ignore_unavailable=True) + + vector._client.get.return_value = {"_id": "id-1"} + assert vector.text_exists("id-1") is True + vector._client.get.side_effect = RuntimeError("not found") + assert vector.text_exists("id-1") is False + + @staticmethod + def test_search_by_vector_validates_and_builds_documents(opensearch_module): + vector = opensearch_module.OpenSearchVector("collection_1", _opensearch_vector_config(opensearch_module)) + + with pytest.raises(ValueError, match="query_vector should be a list"): + vector.search_by_vector("not-a-list") + + with pytest.raises(ValueError, match="should be floats"): + vector.search_by_vector([0.1, 1]) + + vector._client.search.return_value = { + "hits": { + "hits": [ + { + "_source": { + opensearch_module.Field.CONTENT_KEY: "doc-1", + opensearch_module.Field.METADATA_KEY: None, + }, + "_score": 0.9, + }, + { + "_source": { + opensearch_module.Field.CONTENT_KEY: "doc-2", + opensearch_module.Field.METADATA_KEY: {"doc_id": "2"}, + }, + "_score": 0.1, + }, + ] + } + } + docs = vector.search_by_vector([0.1, 0.2], top_k=2, score_threshold=0.5) + assert len(docs) == 1 + assert docs[0].page_content == "doc-1" + assert docs[0].metadata["score"] == pytest.approx(0.9) + + vector.search_by_vector([0.1, 0.2], top_k=3, document_ids_filter=["doc-a", "doc-b"]) + query = vector._client.search.call_args.kwargs["body"] + assert "script_score" in query["query"] + + @staticmethod + def test_search_by_vector_reraises_client_error(opensearch_module): + vector = opensearch_module.OpenSearchVector("collection_1", _opensearch_vector_config(opensearch_module)) + vector._client.search.side_effect = RuntimeError("boom") + + with pytest.raises(RuntimeError, match="boom"): + vector.search_by_vector([0.1, 0.2]) + + @staticmethod + def test_search_by_full_text_and_filters(opensearch_module): + vector = opensearch_module.OpenSearchVector("collection_1", _opensearch_vector_config(opensearch_module)) + vector._client.search.return_value = { + "hits": { + "hits": [ + { + "_source": { + opensearch_module.Field.METADATA_KEY: {"doc_id": "1"}, + opensearch_module.Field.VECTOR: [0.1], + opensearch_module.Field.CONTENT_KEY: "matched text", + } + }, + ] + } + } + + docs = vector.search_by_full_text("hello", document_ids_filter=["d-1"]) + + assert len(docs) == 1 + assert docs[0].page_content == "matched text" + query = vector._client.search.call_args.kwargs["body"] + assert query["query"]["bool"]["filter"] == [{"terms": {"metadata.document_id": ["d-1"]}}] + + @staticmethod + def test_create_collection_cache_and_create_path(opensearch_module, monkeypatch: pytest.MonkeyPatch): + lock = MagicMock() + lock.__enter__.return_value = None + lock.__exit__.return_value = None + monkeypatch.setattr(opensearch_module.redis_client, "lock", MagicMock(return_value=lock)) + monkeypatch.setattr(opensearch_module.redis_client, "set", MagicMock()) + + vector = opensearch_module.OpenSearchVector("Collection_1", _opensearch_vector_config(opensearch_module)) + + monkeypatch.setattr(opensearch_module.redis_client, "get", MagicMock(return_value=1)) + vector._client.indices.create.reset_mock() + vector.create_collection([[0.1, 0.2]]) + vector._client.indices.create.assert_not_called() + + monkeypatch.setattr(opensearch_module.redis_client, "get", MagicMock(return_value=None)) + vector._client.indices.exists.return_value = False + vector.create_collection([[0.1, 0.2]]) + vector._client.indices.create.assert_called_once() + index_body = vector._client.indices.create.call_args.kwargs["body"] + assert index_body["mappings"]["properties"]["vector"]["dimension"] == 2 + opensearch_module.redis_client.set.assert_called() + + @staticmethod + def test_opensearch_factory_initializes_expected_collection_name( + opensearch_module, monkeypatch: pytest.MonkeyPatch + ): + factory = opensearch_module.OpenSearchVectorFactory() + dataset_with_index = Dataset( + id="dataset-1", index_struct=json.dumps({"vector_store": {"class_prefix": "EXISTING_COLLECTION"}}) + ) + dataset_without_index = Dataset(id="dataset-2") + + monkeypatch.setattr(opensearch_module.Dataset, "gen_collection_name_by_id", lambda _id: "AUTO_COLLECTION") + monkeypatch.setattr(opensearch_module.dify_config, "OPENSEARCH_HOST", "localhost") + monkeypatch.setattr(opensearch_module.dify_config, "OPENSEARCH_PORT", 9200) + monkeypatch.setattr(opensearch_module.dify_config, "OPENSEARCH_SECURE", True) + monkeypatch.setattr(opensearch_module.dify_config, "OPENSEARCH_VERIFY_CERTS", True) + monkeypatch.setattr(opensearch_module.dify_config, "OPENSEARCH_AUTH_METHOD", "basic") + monkeypatch.setattr(opensearch_module.dify_config, "OPENSEARCH_USER", "admin") + monkeypatch.setattr(opensearch_module.dify_config, "OPENSEARCH_PASSWORD", "secret") + monkeypatch.setattr(opensearch_module.dify_config, "OPENSEARCH_AWS_REGION", None) + monkeypatch.setattr(opensearch_module.dify_config, "OPENSEARCH_AWS_SERVICE", None) + + with patch.object(opensearch_module, "OpenSearchVector", return_value="vector") as vector_cls: + result_1 = factory.init_vector(dataset_with_index, attributes=[], embeddings=MagicMock(), session=None) + result_2 = factory.init_vector(dataset_without_index, attributes=[], embeddings=MagicMock(), session=None) + + assert result_1 == "vector" + assert result_2 == "vector" + assert vector_cls.call_args_list[0].kwargs["collection_name"] == "existing_collection" + assert vector_cls.call_args_list[1].kwargs["collection_name"] == "auto_collection" + assert dataset_without_index.index_struct is not None diff --git a/api/providers/vdb/vdb-opensearch/tests/unit_tests/test_opensearch.py b/api/providers/vdb/vdb-opensearch/tests/unit_tests/test_opensearch.py deleted file mode 100644 index b2b004a4de9..00000000000 --- a/api/providers/vdb/vdb-opensearch/tests/unit_tests/test_opensearch.py +++ /dev/null @@ -1,332 +0,0 @@ -import importlib -import sys -import types -from types import SimpleNamespace -from unittest.mock import MagicMock - -import pytest - -from core.rag.datasource.vdb.field import Field -from core.rag.models.document import Document -from extensions import ext_redis - - -def _build_fake_opensearch_modules(): - """Build fake opensearchpy modules to avoid the ``from events import Events`` - namespace collision (opensearch-py #756).""" - opensearchpy = types.ModuleType("opensearchpy") - opensearchpy_helpers = types.ModuleType("opensearchpy.helpers") - - class BulkIndexError(Exception): - def __init__(self, errors): - super().__init__("bulk error") - self.errors = errors - - class Urllib3AWSV4SignerAuth: - def __init__(self, credentials, region, service): - self.credentials = credentials - self.region = region - self.service = service - - class Urllib3HttpConnection: - pass - - class _IndicesClient: - def __init__(self): - self.exists = MagicMock(return_value=False) - self.create = MagicMock() - self.delete = MagicMock() - - class OpenSearch: - def __init__(self, **kwargs): - self.kwargs = kwargs - self.indices = _IndicesClient() - self.search = MagicMock(return_value={"hits": {"hits": []}}) - self.get = MagicMock() - - helpers = SimpleNamespace(bulk=MagicMock()) - - opensearchpy.OpenSearch = OpenSearch - opensearchpy.Urllib3AWSV4SignerAuth = Urllib3AWSV4SignerAuth - opensearchpy.Urllib3HttpConnection = Urllib3HttpConnection - opensearchpy.helpers = helpers - opensearchpy_helpers.BulkIndexError = BulkIndexError - - return { - "opensearchpy": opensearchpy, - "opensearchpy.helpers": opensearchpy_helpers, - } - - -@pytest.fixture -def opensearch_module(monkeypatch: pytest.MonkeyPatch): - for name, module in _build_fake_opensearch_modules().items(): - monkeypatch.setitem(sys.modules, name, module) - - import dify_vdb_opensearch.opensearch_vector as module - - return importlib.reload(module) - - -def _config(module, **overrides): - values = { - "host": "localhost", - "port": 9200, - "secure": False, - "user": "admin", - "password": "password", - } - values.update(overrides) - return module.OpenSearchConfig.model_validate(values) - - -def get_example_text() -> str: - return "This is a sample text for testing purposes." - - -class TestOpenSearchConfig: - def test_to_opensearch_params(self, opensearch_module): - config = _config(opensearch_module, secure=True) - params = config.to_opensearch_params() - - assert params["hosts"] == [{"host": "localhost", "port": 9200}] - assert params["use_ssl"] is True - assert params["verify_certs"] is True - assert params["connection_class"].__name__ == "Urllib3HttpConnection" - assert params["http_auth"] == ("admin", "password") - - def test_to_opensearch_params_with_aws_managed_iam(self, opensearch_module, monkeypatch: pytest.MonkeyPatch): - class _Session: - def get_credentials(self): - return "creds" - - boto3 = types.ModuleType("boto3") - boto3.Session = _Session - monkeypatch.setitem(sys.modules, "boto3", boto3) - - config = _config( - opensearch_module, - secure=True, - auth_method="aws_managed_iam", - aws_region="ap-southeast-2", - aws_service="aoss", - host="aoss-endpoint.ap-southeast-2.aoss.amazonaws.com", - port=9201, - ) - params = config.to_opensearch_params() - - assert params["hosts"] == [{"host": "aoss-endpoint.ap-southeast-2.aoss.amazonaws.com", "port": 9201}] - assert params["use_ssl"] is True - assert params["verify_certs"] is True - assert params["connection_class"].__name__ == "Urllib3HttpConnection" - assert params["http_auth"].credentials == "creds" - assert params["http_auth"].region == "ap-southeast-2" - assert params["http_auth"].service == "aoss" - - -class TestOpenSearchVector: - COLLECTION_NAME = "test_collection" - EXAMPLE_DOC_ID = "example_doc_id" - - def _make_vector(self, module): - vector = module.OpenSearchVector(self.COLLECTION_NAME, _config(module)) - vector._client = MagicMock() - return vector - - @pytest.mark.parametrize( - ("search_response", "expected_length", "expected_doc_id"), - [ - ( - { - "hits": { - "total": {"value": 1}, - "hits": [ - { - "_source": { - "page_content": get_example_text(), - "metadata": {"document_id": "example_doc_id"}, - } - } - ], - } - }, - 1, - "example_doc_id", - ), - ({"hits": {"total": {"value": 0}, "hits": []}}, 0, None), - ], - ) - def test_search_by_full_text(self, opensearch_module, search_response, expected_length, expected_doc_id): - vector = self._make_vector(opensearch_module) - vector._client.search.return_value = search_response - - hits = vector.search_by_full_text(query=get_example_text()) - assert len(hits) == expected_length - if expected_length > 0: - assert hits[0].metadata["document_id"] == expected_doc_id - - def test_search_by_vector(self, opensearch_module): - vector = self._make_vector(opensearch_module) - query_vector = [0.1] * 128 - mock_response = { - "hits": { - "total": {"value": 1}, - "hits": [ - { - "_source": { - Field.CONTENT_KEY: get_example_text(), - Field.METADATA_KEY: {"document_id": self.EXAMPLE_DOC_ID}, - }, - "_score": 1.0, - } - ], - } - } - vector._client.search.return_value = mock_response - - hits = vector.search_by_vector(query_vector=query_vector) - - assert len(hits) > 0 - assert hits[0].metadata["document_id"] == self.EXAMPLE_DOC_ID - - def test_get_ids_by_metadata_field(self, opensearch_module): - vector = self._make_vector(opensearch_module) - mock_response = {"hits": {"total": {"value": 1}, "hits": [{"_id": "mock_id"}]}} - vector._client.search.return_value = mock_response - - doc = Document(page_content="Test content", metadata={"document_id": self.EXAMPLE_DOC_ID}) - embedding = [0.1] * 128 - - opensearch_module.helpers.bulk.reset_mock() - vector.add_texts([doc], [embedding]) - - ids = vector.get_ids_by_metadata_field(key="document_id", value=self.EXAMPLE_DOC_ID) - assert len(ids) == 1 - assert ids[0] == "mock_id" - - def test_add_texts(self, opensearch_module): - vector = self._make_vector(opensearch_module) - vector._client.index.return_value = {"result": "created"} - - doc = Document(page_content="Test content", metadata={"document_id": self.EXAMPLE_DOC_ID}) - embedding = [0.1] * 128 - - opensearch_module.helpers.bulk.reset_mock() - vector.add_texts([doc], [embedding]) - - mock_response = {"hits": {"total": {"value": 1}, "hits": [{"_id": "mock_id"}]}} - vector._client.search.return_value = mock_response - - ids = vector.get_ids_by_metadata_field(key="document_id", value=self.EXAMPLE_DOC_ID) - assert len(ids) == 1 - assert ids[0] == "mock_id" - - def test_delete_nonexistent_index(self, opensearch_module): - """ignore_unavailable=True handles non-existent indices gracefully.""" - vector = self._make_vector(opensearch_module) - vector.delete() - - vector._client.indices.delete.assert_called_once_with( - index=self.COLLECTION_NAME.lower(), ignore_unavailable=True - ) - - def test_delete_existing_index(self, opensearch_module): - vector = self._make_vector(opensearch_module) - vector.delete() - - vector._client.indices.delete.assert_called_once_with( - index=self.COLLECTION_NAME.lower(), ignore_unavailable=True - ) - - -@pytest.fixture(scope="module") -def setup_mock_redis(): - ext_redis.redis_client.get = MagicMock(return_value=None) - ext_redis.redis_client.set = MagicMock(return_value=None) - - mock_redis_lock = MagicMock() - mock_redis_lock.__enter__ = MagicMock() - mock_redis_lock.__exit__ = MagicMock() - ext_redis.redis_client.lock = MagicMock(return_value=mock_redis_lock) - - -@pytest.mark.usefixtures("setup_mock_redis") -class TestOpenSearchVectorWithRedis: - COLLECTION_NAME = "test_collection" - EXAMPLE_DOC_ID = "example_doc_id" - - def _make_vector(self, module): - vector = module.OpenSearchVector(self.COLLECTION_NAME, _config(module)) - vector._client = MagicMock() - return vector - - def test_search_by_full_text(self, opensearch_module): - vector = self._make_vector(opensearch_module) - search_response = { - "hits": { - "total": {"value": 1}, - "hits": [ - {"_source": {"page_content": get_example_text(), "metadata": {"document_id": "example_doc_id"}}} - ], - } - } - vector._client.search.return_value = search_response - - hits = vector.search_by_full_text(query=get_example_text()) - assert len(hits) == 1 - assert hits[0].metadata["document_id"] == "example_doc_id" - - def test_get_ids_by_metadata_field(self, opensearch_module): - vector = self._make_vector(opensearch_module) - mock_response = {"hits": {"total": {"value": 1}, "hits": [{"_id": "mock_id"}]}} - vector._client.search.return_value = mock_response - - doc = Document(page_content="Test content", metadata={"document_id": self.EXAMPLE_DOC_ID}) - embedding = [0.1] * 128 - - opensearch_module.helpers.bulk.reset_mock() - vector.add_texts([doc], [embedding]) - - ids = vector.get_ids_by_metadata_field(key="document_id", value=self.EXAMPLE_DOC_ID) - assert len(ids) == 1 - assert ids[0] == "mock_id" - - def test_add_texts(self, opensearch_module): - vector = self._make_vector(opensearch_module) - vector._client.index.return_value = {"result": "created"} - - doc = Document(page_content="Test content", metadata={"document_id": self.EXAMPLE_DOC_ID}) - embedding = [0.1] * 128 - - opensearch_module.helpers.bulk.reset_mock() - vector.add_texts([doc], [embedding]) - - mock_response = {"hits": {"total": {"value": 1}, "hits": [{"_id": "mock_id"}]}} - vector._client.search.return_value = mock_response - - ids = vector.get_ids_by_metadata_field(key="document_id", value=self.EXAMPLE_DOC_ID) - assert len(ids) == 1 - assert ids[0] == "mock_id" - - def test_search_by_vector(self, opensearch_module): - vector = self._make_vector(opensearch_module) - query_vector = [0.1] * 128 - mock_response = { - "hits": { - "total": {"value": 1}, - "hits": [ - { - "_source": { - Field.CONTENT_KEY: get_example_text(), - Field.METADATA_KEY: {"document_id": self.EXAMPLE_DOC_ID}, - }, - "_score": 1.0, - } - ], - } - } - vector._client.search.return_value = mock_response - - hits = vector.search_by_vector(query_vector=query_vector) - assert len(hits) > 0 - assert hits[0].metadata["document_id"] == self.EXAMPLE_DOC_ID diff --git a/api/providers/vdb/vdb-opensearch/tests/unit_tests/test_opensearch_vector.py b/api/providers/vdb/vdb-opensearch/tests/unit_tests/test_opensearch_vector.py deleted file mode 100644 index f7ebb39c6ac..00000000000 --- a/api/providers/vdb/vdb-opensearch/tests/unit_tests/test_opensearch_vector.py +++ /dev/null @@ -1,362 +0,0 @@ -import importlib -import json -import sys -import types -from types import SimpleNamespace -from unittest.mock import MagicMock, patch - -import pytest -from pydantic import ValidationError - -from core.rag.models.document import Document -from models.dataset import Dataset - - -# TODO(wylswz): There's a known issue with namespace collision -# https://github.com/langgenius/dify/issues/34732 -def _build_fake_opensearch_modules(): - opensearchpy = types.ModuleType("opensearchpy") - opensearchpy_helpers = types.ModuleType("opensearchpy.helpers") - - class BulkIndexError(Exception): - def __init__(self, errors): - super().__init__("bulk error") - self.errors = errors - - class Urllib3AWSV4SignerAuth: - def __init__(self, credentials, region, service): - self.credentials = credentials - self.region = region - self.service = service - - class Urllib3HttpConnection: - pass - - class _IndicesClient: - def __init__(self): - self.exists = MagicMock(return_value=False) - self.create = MagicMock() - self.delete = MagicMock() - - class OpenSearch: - def __init__(self, **kwargs): - self.kwargs = kwargs - self.indices = _IndicesClient() - self.search = MagicMock(return_value={"hits": {"hits": []}}) - self.get = MagicMock() - - helpers = SimpleNamespace(bulk=MagicMock()) - - opensearchpy.OpenSearch = OpenSearch - opensearchpy.Urllib3AWSV4SignerAuth = Urllib3AWSV4SignerAuth - opensearchpy.Urllib3HttpConnection = Urllib3HttpConnection - opensearchpy.helpers = helpers - opensearchpy_helpers.BulkIndexError = BulkIndexError - - return { - "opensearchpy": opensearchpy, - "opensearchpy.helpers": opensearchpy_helpers, - } - - -@pytest.fixture -def opensearch_module(monkeypatch: pytest.MonkeyPatch): - for name, module in _build_fake_opensearch_modules().items(): - monkeypatch.setitem(sys.modules, name, module) - - import dify_vdb_opensearch.opensearch_vector as module - - return importlib.reload(module) - - -def _config(module, **overrides): - values = { - "host": "localhost", - "port": 9200, - "secure": True, - "verify_certs": True, - "auth_method": "basic", - "user": "admin", - "password": "secret", - } - values.update(overrides) - return module.OpenSearchConfig.model_validate(values) - - -@pytest.mark.parametrize( - ("field", "value", "message"), - [ - ("host", "", "config OPENSEARCH_HOST is required"), - ("port", 0, "config OPENSEARCH_PORT is required"), - ], -) -def test_config_validation_required_fields(opensearch_module, field, value, message): - values = _config(opensearch_module).model_dump() - values[field] = value - - with pytest.raises(ValidationError, match=message): - opensearch_module.OpenSearchConfig.model_validate(values) - - -def test_config_validation_for_aws_auth_and_https_fields(opensearch_module): - values = { - "host": "localhost", - "port": 9200, - "secure": True, - "verify_certs": True, - "auth_method": "aws_managed_iam", - "user": "admin", - "password": "secret", - } - with pytest.raises(ValidationError, match="OPENSEARCH_AWS_REGION"): - opensearch_module.OpenSearchConfig.model_validate(values) - - values = _config(opensearch_module).model_dump() - values["OPENSEARCH_SECURE"] = False - values["OPENSEARCH_VERIFY_CERTS"] = True - with pytest.raises(ValidationError, match="verify_certs=True requires secure"): - opensearch_module.OpenSearchConfig.model_validate(values) - - -def test_create_aws_managed_iam_auth(opensearch_module, monkeypatch: pytest.MonkeyPatch): - class _Session: - def get_credentials(self): - return "creds" - - boto3 = types.ModuleType("boto3") - boto3.Session = _Session - monkeypatch.setitem(sys.modules, "boto3", boto3) - - config = _config( - opensearch_module, - auth_method="aws_managed_iam", - aws_region="us-east-1", - aws_service="es", - ) - auth = config.create_aws_managed_iam_auth() - - assert auth.credentials == "creds" - assert auth.region == "us-east-1" - assert auth.service == "es" - - -def test_to_opensearch_params_supports_basic_and_aws(opensearch_module): - basic_params = _config(opensearch_module).to_opensearch_params() - assert basic_params["http_auth"] == ("admin", "secret") - - aws_config = _config( - opensearch_module, - auth_method="aws_managed_iam", - aws_region="us-west-2", - aws_service="es", - ) - with patch.object(opensearch_module.OpenSearchConfig, "create_aws_managed_iam_auth", return_value="iam-auth"): - aws_params = aws_config.to_opensearch_params() - - assert aws_params["http_auth"] == "iam-auth" - - -def test_init_and_create_delegate_calls(opensearch_module): - vector = opensearch_module.OpenSearchVector("Collection_1", _config(opensearch_module)) - vector.create_collection = MagicMock() - vector.add_texts = MagicMock() - docs = [Document(page_content="hello", metadata={"doc_id": "seg-1"})] - - vector.create(docs, [[0.1, 0.2]]) - - assert vector.get_type() == "opensearch" - vector.create_collection.assert_called_once_with([[0.1, 0.2]], [{"doc_id": "seg-1"}]) - vector.add_texts.assert_called_once_with(docs, [[0.1, 0.2]]) - - -def test_add_texts_supports_regular_and_aoss_clients(opensearch_module, monkeypatch: pytest.MonkeyPatch): - vector = opensearch_module.OpenSearchVector("Collection_1", _config(opensearch_module, aws_service="es")) - docs = [ - Document(page_content="a", metadata={"doc_id": "1"}), - Document(page_content="b", metadata={"doc_id": "2"}), - ] - - monkeypatch.setattr(opensearch_module, "uuid4", lambda: SimpleNamespace(hex="generated-id")) - opensearch_module.helpers.bulk.reset_mock() - vector.add_texts(docs, [[0.1], [0.2]]) - actions = opensearch_module.helpers.bulk.call_args.kwargs["actions"] - assert len(actions) == 2 - assert all("_id" in action for action in actions) - - vector._client_config.aws_service = "aoss" - opensearch_module.helpers.bulk.reset_mock() - vector.add_texts(docs, [[0.3], [0.4]]) - aoss_actions = opensearch_module.helpers.bulk.call_args.kwargs["actions"] - assert all("_id" not in action for action in aoss_actions) - - -def test_metadata_lookup_and_delete_by_metadata_field(opensearch_module): - vector = opensearch_module.OpenSearchVector("collection_1", _config(opensearch_module)) - vector._client.search.return_value = {"hits": {"hits": [{"_id": "id-1"}, {"_id": "id-2"}]}} - - assert vector.get_ids_by_metadata_field("document_id", "doc-1") == ["id-1", "id-2"] - - vector._client.search.return_value = {"hits": {"hits": []}} - assert vector.get_ids_by_metadata_field("document_id", "doc-1") is None - - vector.get_ids_by_metadata_field = MagicMock(return_value=["id-1"]) - vector.delete_by_ids = MagicMock() - vector.delete_by_metadata_field("document_id", "doc-1") - vector.delete_by_ids.assert_called_once_with(["id-1"]) - - -def test_delete_by_ids_branches_and_bulk_error_handling(opensearch_module): - vector = opensearch_module.OpenSearchVector("collection_1", _config(opensearch_module)) - opensearch_module.helpers.bulk.reset_mock() - vector._client.indices.exists.return_value = False - vector.delete_by_ids(["doc-1"]) - opensearch_module.helpers.bulk.assert_not_called() - - vector._client.indices.exists.return_value = True - vector.get_ids_by_metadata_field = MagicMock(side_effect=[["es-1"], None]) - vector.delete_by_ids(["doc-1", "doc-2"]) - opensearch_module.helpers.bulk.assert_called_once() - - opensearch_module.helpers.bulk.reset_mock() - vector.get_ids_by_metadata_field = MagicMock(return_value=["es-404"]) - opensearch_module.helpers.bulk.side_effect = opensearch_module.BulkIndexError( - [{"delete": {"status": 404, "_id": "es-404"}}] - ) - vector.delete_by_ids(["doc-404"]) - assert opensearch_module.helpers.bulk.call_count == 1 - - opensearch_module.helpers.bulk.side_effect = None - - -def test_delete_and_text_exists(opensearch_module): - vector = opensearch_module.OpenSearchVector("collection_1", _config(opensearch_module)) - vector.delete() - vector._client.indices.delete.assert_called_once_with(index="collection_1", ignore_unavailable=True) - - vector._client.get.return_value = {"_id": "id-1"} - assert vector.text_exists("id-1") is True - vector._client.get.side_effect = RuntimeError("not found") - assert vector.text_exists("id-1") is False - - -def test_search_by_vector_validates_and_builds_documents(opensearch_module): - vector = opensearch_module.OpenSearchVector("collection_1", _config(opensearch_module)) - - with pytest.raises(ValueError, match="query_vector should be a list"): - vector.search_by_vector("not-a-list") - - with pytest.raises(ValueError, match="should be floats"): - vector.search_by_vector([0.1, 1]) - - vector._client.search.return_value = { - "hits": { - "hits": [ - { - "_source": { - opensearch_module.Field.CONTENT_KEY: "doc-1", - opensearch_module.Field.METADATA_KEY: None, - }, - "_score": 0.9, - }, - { - "_source": { - opensearch_module.Field.CONTENT_KEY: "doc-2", - opensearch_module.Field.METADATA_KEY: {"doc_id": "2"}, - }, - "_score": 0.1, - }, - ] - } - } - docs = vector.search_by_vector([0.1, 0.2], top_k=2, score_threshold=0.5) - assert len(docs) == 1 - assert docs[0].page_content == "doc-1" - assert docs[0].metadata["score"] == pytest.approx(0.9) - - vector.search_by_vector([0.1, 0.2], top_k=3, document_ids_filter=["doc-a", "doc-b"]) - query = vector._client.search.call_args.kwargs["body"] - assert "script_score" in query["query"] - - -def test_search_by_vector_reraises_client_error(opensearch_module): - vector = opensearch_module.OpenSearchVector("collection_1", _config(opensearch_module)) - vector._client.search.side_effect = RuntimeError("boom") - - with pytest.raises(RuntimeError, match="boom"): - vector.search_by_vector([0.1, 0.2]) - - -def test_search_by_full_text_and_filters(opensearch_module): - vector = opensearch_module.OpenSearchVector("collection_1", _config(opensearch_module)) - vector._client.search.return_value = { - "hits": { - "hits": [ - { - "_source": { - opensearch_module.Field.METADATA_KEY: {"doc_id": "1"}, - opensearch_module.Field.VECTOR: [0.1], - opensearch_module.Field.CONTENT_KEY: "matched text", - } - }, - ] - } - } - - docs = vector.search_by_full_text("hello", document_ids_filter=["d-1"]) - - assert len(docs) == 1 - assert docs[0].page_content == "matched text" - query = vector._client.search.call_args.kwargs["body"] - assert query["query"]["bool"]["filter"] == [{"terms": {"metadata.document_id": ["d-1"]}}] - - -def test_create_collection_cache_and_create_path(opensearch_module, monkeypatch: pytest.MonkeyPatch): - lock = MagicMock() - lock.__enter__.return_value = None - lock.__exit__.return_value = None - monkeypatch.setattr(opensearch_module.redis_client, "lock", MagicMock(return_value=lock)) - monkeypatch.setattr(opensearch_module.redis_client, "set", MagicMock()) - - vector = opensearch_module.OpenSearchVector("Collection_1", _config(opensearch_module)) - - monkeypatch.setattr(opensearch_module.redis_client, "get", MagicMock(return_value=1)) - vector._client.indices.create.reset_mock() - vector.create_collection([[0.1, 0.2]]) - vector._client.indices.create.assert_not_called() - - monkeypatch.setattr(opensearch_module.redis_client, "get", MagicMock(return_value=None)) - vector._client.indices.exists.return_value = False - vector.create_collection([[0.1, 0.2]]) - vector._client.indices.create.assert_called_once() - index_body = vector._client.indices.create.call_args.kwargs["body"] - assert index_body["mappings"]["properties"]["vector"]["dimension"] == 2 - opensearch_module.redis_client.set.assert_called() - - -def test_opensearch_factory_initializes_expected_collection_name(opensearch_module, monkeypatch: pytest.MonkeyPatch): - factory = opensearch_module.OpenSearchVectorFactory() - dataset_with_index = Dataset( - id="dataset-1", index_struct=json.dumps({"vector_store": {"class_prefix": "EXISTING_COLLECTION"}}) - ) - dataset_without_index = Dataset(id="dataset-2") - - monkeypatch.setattr(opensearch_module.Dataset, "gen_collection_name_by_id", lambda _id: "AUTO_COLLECTION") - monkeypatch.setattr(opensearch_module.dify_config, "OPENSEARCH_HOST", "localhost") - monkeypatch.setattr(opensearch_module.dify_config, "OPENSEARCH_PORT", 9200) - monkeypatch.setattr(opensearch_module.dify_config, "OPENSEARCH_SECURE", True) - monkeypatch.setattr(opensearch_module.dify_config, "OPENSEARCH_VERIFY_CERTS", True) - monkeypatch.setattr(opensearch_module.dify_config, "OPENSEARCH_AUTH_METHOD", "basic") - monkeypatch.setattr(opensearch_module.dify_config, "OPENSEARCH_USER", "admin") - monkeypatch.setattr(opensearch_module.dify_config, "OPENSEARCH_PASSWORD", "secret") - monkeypatch.setattr(opensearch_module.dify_config, "OPENSEARCH_AWS_REGION", None) - monkeypatch.setattr(opensearch_module.dify_config, "OPENSEARCH_AWS_SERVICE", None) - - with patch.object(opensearch_module, "OpenSearchVector", return_value="vector") as vector_cls: - result_1 = factory.init_vector(dataset_with_index, attributes=[], embeddings=MagicMock(), session=None) - result_2 = factory.init_vector(dataset_without_index, attributes=[], embeddings=MagicMock(), session=None) - - assert result_1 == "vector" - assert result_2 == "vector" - assert vector_cls.call_args_list[0].kwargs["collection_name"] == "existing_collection" - assert vector_cls.call_args_list[1].kwargs["collection_name"] == "auto_collection" - assert dataset_without_index.index_struct is not None diff --git a/api/providers/vdb/vdb-oracle/tests/integration_tests/test_oraclevector.py b/api/providers/vdb/vdb-oracle/tests/e2e_tests/test_oraclevector.py similarity index 100% rename from api/providers/vdb/vdb-oracle/tests/integration_tests/test_oraclevector.py rename to api/providers/vdb/vdb-oracle/tests/e2e_tests/test_oraclevector.py diff --git a/api/providers/vdb/vdb-oracle/tests/unit_tests/test_oraclevector.py b/api/providers/vdb/vdb-oracle/tests/unit_tests/src/dify_vdb_oracle/test_oraclevector.py similarity index 100% rename from api/providers/vdb/vdb-oracle/tests/unit_tests/test_oraclevector.py rename to api/providers/vdb/vdb-oracle/tests/unit_tests/src/dify_vdb_oracle/test_oraclevector.py diff --git a/api/providers/vdb/vdb-pgvecto-rs/tests/integration_tests/test_pgvecto_rs.py b/api/providers/vdb/vdb-pgvecto-rs/tests/e2e_tests/test_pgvecto_rs.py similarity index 100% rename from api/providers/vdb/vdb-pgvecto-rs/tests/integration_tests/test_pgvecto_rs.py rename to api/providers/vdb/vdb-pgvecto-rs/tests/e2e_tests/test_pgvecto_rs.py diff --git a/api/providers/vdb/vdb-pgvecto-rs/tests/unit_tests/test_pgvecto_rs.py b/api/providers/vdb/vdb-pgvecto-rs/tests/unit_tests/src/dify_vdb_pgvecto_rs/test_pgvecto_rs.py similarity index 100% rename from api/providers/vdb/vdb-pgvecto-rs/tests/unit_tests/test_pgvecto_rs.py rename to api/providers/vdb/vdb-pgvecto-rs/tests/unit_tests/src/dify_vdb_pgvecto_rs/test_pgvecto_rs.py diff --git a/api/providers/vdb/vdb-pgvector/tests/integration_tests/test_pgvector.py b/api/providers/vdb/vdb-pgvector/tests/e2e_tests/test_pgvector.py similarity index 100% rename from api/providers/vdb/vdb-pgvector/tests/integration_tests/test_pgvector.py rename to api/providers/vdb/vdb-pgvector/tests/e2e_tests/test_pgvector.py diff --git a/api/providers/vdb/vdb-pgvector/tests/unit_tests/test_pgvector.py b/api/providers/vdb/vdb-pgvector/tests/unit_tests/src/dify_vdb_pgvector/test_pgvector.py similarity index 100% rename from api/providers/vdb/vdb-pgvector/tests/unit_tests/test_pgvector.py rename to api/providers/vdb/vdb-pgvector/tests/unit_tests/src/dify_vdb_pgvector/test_pgvector.py diff --git a/api/providers/vdb/vdb-qdrant/tests/integration_tests/test_qdrant.py b/api/providers/vdb/vdb-qdrant/tests/e2e_tests/test_qdrant.py similarity index 100% rename from api/providers/vdb/vdb-qdrant/tests/integration_tests/test_qdrant.py rename to api/providers/vdb/vdb-qdrant/tests/e2e_tests/test_qdrant.py diff --git a/api/providers/vdb/vdb-qdrant/tests/unit_tests/test_qdrant_vector.py b/api/providers/vdb/vdb-qdrant/tests/unit_tests/src/dify_vdb_qdrant/test_qdrant_vector.py similarity index 100% rename from api/providers/vdb/vdb-qdrant/tests/unit_tests/test_qdrant_vector.py rename to api/providers/vdb/vdb-qdrant/tests/unit_tests/src/dify_vdb_qdrant/test_qdrant_vector.py diff --git a/api/providers/vdb/vdb-relyt/tests/unit_tests/test_relyt_vector.py b/api/providers/vdb/vdb-relyt/tests/unit_tests/src/dify_vdb_relyt/test_relyt_vector.py similarity index 100% rename from api/providers/vdb/vdb-relyt/tests/unit_tests/test_relyt_vector.py rename to api/providers/vdb/vdb-relyt/tests/unit_tests/src/dify_vdb_relyt/test_relyt_vector.py diff --git a/api/providers/vdb/vdb-tablestore/tests/integration_tests/test_tablestore.py b/api/providers/vdb/vdb-tablestore/tests/e2e_tests/test_tablestore.py similarity index 100% rename from api/providers/vdb/vdb-tablestore/tests/integration_tests/test_tablestore.py rename to api/providers/vdb/vdb-tablestore/tests/e2e_tests/test_tablestore.py diff --git a/api/providers/vdb/vdb-tablestore/tests/unit_tests/test_tablestore_vector.py b/api/providers/vdb/vdb-tablestore/tests/unit_tests/src/dify_vdb_tablestore/test_tablestore_vector.py similarity index 100% rename from api/providers/vdb/vdb-tablestore/tests/unit_tests/test_tablestore_vector.py rename to api/providers/vdb/vdb-tablestore/tests/unit_tests/src/dify_vdb_tablestore/test_tablestore_vector.py diff --git a/api/providers/vdb/vdb-tencent/tests/integration_tests/conftest.py b/api/providers/vdb/vdb-tencent/tests/e2e_tests/conftest.py similarity index 100% rename from api/providers/vdb/vdb-tencent/tests/integration_tests/conftest.py rename to api/providers/vdb/vdb-tencent/tests/e2e_tests/conftest.py diff --git a/api/providers/vdb/vdb-tencent/tests/integration_tests/test_tencent.py b/api/providers/vdb/vdb-tencent/tests/e2e_tests/test_tencent.py similarity index 100% rename from api/providers/vdb/vdb-tencent/tests/integration_tests/test_tencent.py rename to api/providers/vdb/vdb-tencent/tests/e2e_tests/test_tencent.py diff --git a/api/providers/vdb/vdb-tencent/tests/unit_tests/test_tencent_vector.py b/api/providers/vdb/vdb-tencent/tests/unit_tests/src/dify_vdb_tencent/test_tencent_vector.py similarity index 100% rename from api/providers/vdb/vdb-tencent/tests/unit_tests/test_tencent_vector.py rename to api/providers/vdb/vdb-tencent/tests/unit_tests/src/dify_vdb_tencent/test_tencent_vector.py diff --git a/api/providers/vdb/vdb-tidb-on-qdrant/tests/unit_tests/test_tidb_on_qdrant_vector.py b/api/providers/vdb/vdb-tidb-on-qdrant/tests/unit_tests/src/dify_vdb_tidb_on_qdrant/test_tidb_on_qdrant_vector.py similarity index 100% rename from api/providers/vdb/vdb-tidb-on-qdrant/tests/unit_tests/test_tidb_on_qdrant_vector.py rename to api/providers/vdb/vdb-tidb-on-qdrant/tests/unit_tests/src/dify_vdb_tidb_on_qdrant/test_tidb_on_qdrant_vector.py diff --git a/api/providers/vdb/vdb-tidb-on-qdrant/tests/unit_tests/test_tidb_service.py b/api/providers/vdb/vdb-tidb-on-qdrant/tests/unit_tests/src/dify_vdb_tidb_on_qdrant/test_tidb_service.py similarity index 100% rename from api/providers/vdb/vdb-tidb-on-qdrant/tests/unit_tests/test_tidb_service.py rename to api/providers/vdb/vdb-tidb-on-qdrant/tests/unit_tests/src/dify_vdb_tidb_on_qdrant/test_tidb_service.py diff --git a/api/providers/vdb/vdb-tidb-vector/tests/integration_tests/check_tiflash_ready.py b/api/providers/vdb/vdb-tidb-vector/tests/e2e_tests/check_tiflash_ready.py similarity index 100% rename from api/providers/vdb/vdb-tidb-vector/tests/integration_tests/check_tiflash_ready.py rename to api/providers/vdb/vdb-tidb-vector/tests/e2e_tests/check_tiflash_ready.py diff --git a/api/providers/vdb/vdb-tidb-vector/tests/integration_tests/test_tidb_vector.py b/api/providers/vdb/vdb-tidb-vector/tests/e2e_tests/test_tidb_vector.py similarity index 100% rename from api/providers/vdb/vdb-tidb-vector/tests/integration_tests/test_tidb_vector.py rename to api/providers/vdb/vdb-tidb-vector/tests/e2e_tests/test_tidb_vector.py diff --git a/api/providers/vdb/vdb-tidb-vector/tests/unit_tests/test_tidb_vector.py b/api/providers/vdb/vdb-tidb-vector/tests/unit_tests/src/dify_vdb_tidb_vector/test_tidb_vector.py similarity index 100% rename from api/providers/vdb/vdb-tidb-vector/tests/unit_tests/test_tidb_vector.py rename to api/providers/vdb/vdb-tidb-vector/tests/unit_tests/src/dify_vdb_tidb_vector/test_tidb_vector.py diff --git a/api/providers/vdb/vdb-upstash/tests/integration_tests/conftest.py b/api/providers/vdb/vdb-upstash/tests/e2e_tests/conftest.py similarity index 100% rename from api/providers/vdb/vdb-upstash/tests/integration_tests/conftest.py rename to api/providers/vdb/vdb-upstash/tests/e2e_tests/conftest.py diff --git a/api/providers/vdb/vdb-upstash/tests/integration_tests/test_upstash_vector.py b/api/providers/vdb/vdb-upstash/tests/e2e_tests/test_upstash_vector.py similarity index 100% rename from api/providers/vdb/vdb-upstash/tests/integration_tests/test_upstash_vector.py rename to api/providers/vdb/vdb-upstash/tests/e2e_tests/test_upstash_vector.py diff --git a/api/providers/vdb/vdb-upstash/tests/unit_tests/test_upstash_vector.py b/api/providers/vdb/vdb-upstash/tests/unit_tests/src/dify_vdb_upstash/test_upstash_vector.py similarity index 100% rename from api/providers/vdb/vdb-upstash/tests/unit_tests/test_upstash_vector.py rename to api/providers/vdb/vdb-upstash/tests/unit_tests/src/dify_vdb_upstash/test_upstash_vector.py diff --git a/api/providers/vdb/vdb-vastbase/tests/integration_tests/test_vastbase_vector.py b/api/providers/vdb/vdb-vastbase/tests/e2e_tests/test_vastbase_vector.py similarity index 100% rename from api/providers/vdb/vdb-vastbase/tests/integration_tests/test_vastbase_vector.py rename to api/providers/vdb/vdb-vastbase/tests/e2e_tests/test_vastbase_vector.py diff --git a/api/providers/vdb/vdb-vastbase/tests/unit_tests/test_vastbase_vector.py b/api/providers/vdb/vdb-vastbase/tests/unit_tests/src/dify_vdb_vastbase/test_vastbase_vector.py similarity index 100% rename from api/providers/vdb/vdb-vastbase/tests/unit_tests/test_vastbase_vector.py rename to api/providers/vdb/vdb-vastbase/tests/unit_tests/src/dify_vdb_vastbase/test_vastbase_vector.py diff --git a/api/providers/vdb/vdb-vikingdb/tests/integration_tests/conftest.py b/api/providers/vdb/vdb-vikingdb/tests/e2e_tests/conftest.py similarity index 100% rename from api/providers/vdb/vdb-vikingdb/tests/integration_tests/conftest.py rename to api/providers/vdb/vdb-vikingdb/tests/e2e_tests/conftest.py diff --git a/api/providers/vdb/vdb-vikingdb/tests/integration_tests/test_vikingdb.py b/api/providers/vdb/vdb-vikingdb/tests/e2e_tests/test_vikingdb.py similarity index 100% rename from api/providers/vdb/vdb-vikingdb/tests/integration_tests/test_vikingdb.py rename to api/providers/vdb/vdb-vikingdb/tests/e2e_tests/test_vikingdb.py diff --git a/api/providers/vdb/vdb-vikingdb/tests/unit_tests/test_vikingdb_vector.py b/api/providers/vdb/vdb-vikingdb/tests/unit_tests/src/dify_vdb_vikingdb/test_vikingdb_vector.py similarity index 100% rename from api/providers/vdb/vdb-vikingdb/tests/unit_tests/test_vikingdb_vector.py rename to api/providers/vdb/vdb-vikingdb/tests/unit_tests/src/dify_vdb_vikingdb/test_vikingdb_vector.py diff --git a/api/providers/vdb/vdb-weaviate/tests/integration_tests/test_weaviate.py b/api/providers/vdb/vdb-weaviate/tests/e2e_tests/test_weaviate.py similarity index 100% rename from api/providers/vdb/vdb-weaviate/tests/integration_tests/test_weaviate.py rename to api/providers/vdb/vdb-weaviate/tests/e2e_tests/test_weaviate.py diff --git a/api/providers/vdb/vdb-weaviate/tests/unit_tests/test_weaviate_vector.py b/api/providers/vdb/vdb-weaviate/tests/unit_tests/src/dify_vdb_weaviate/test_weaviate_vector.py similarity index 74% rename from api/providers/vdb/vdb-weaviate/tests/unit_tests/test_weaviate_vector.py rename to api/providers/vdb/vdb-weaviate/tests/unit_tests/src/dify_vdb_weaviate/test_weaviate_vector.py index 6ed323c43db..43e34730970 100644 --- a/api/providers/vdb/vdb-weaviate/tests/unit_tests/test_weaviate_vector.py +++ b/api/providers/vdb/vdb-weaviate/tests/unit_tests/src/dify_vdb_weaviate/test_weaviate_vector.py @@ -1,12 +1,3 @@ -"""Unit tests for Weaviate vector database implementation. - -Focuses on verifying that doc_type is properly handled in: -- Collection schema creation (_create_collection) -- Property migration (_ensure_properties) -- Vector search result metadata (search_by_vector) -- Full-text search result metadata (search_by_full_text) -""" - import datetime import json import unittest @@ -16,12 +7,348 @@ from unittest.mock import MagicMock, patch import pytest from dify_vdb_weaviate import weaviate_vector as weaviate_vector_module -from dify_vdb_weaviate.weaviate_vector import WeaviateConfig, WeaviateVector +from dify_vdb_weaviate.weaviate_vector import ( + _INTEGRATION_NAME, + WeaviateConfig, + WeaviateVector, + _integration_version, + _register_integration, +) from sqlalchemy.orm import Session from core.rag.models.document import Document from models.dataset import Dataset +"""Regression tests for issue #41714. + +The pre-fix code generated Weaviate object ids via +``uuid5(URL_NAMESPACE, page_content)`` (UUID v5 from content). The +cleanup path (``batch_clean_document_task`` -> ``index_processor.clean`` +-> ``vector.delete_by_ids``) passes the segment's ``index_node_id`` +(random UUID v4 from the database), so the two never agreed and +``delete_by_id`` silently no-op'd on every cleanup pass. + +The fix: + +1. ``_get_uuids`` now uses the segment's ``doc_id`` (== ``index_node_id``) + from each document's metadata. New objects line up with the + cleanup path one-for-one. +2. Documents without a ``doc_id`` keep a stable slot via a freshly + generated ``uuid4`` so the parallel ``objs`` list in ``add_texts`` + never goes out of alignment with the input documents. +3. ``delete_by_ids`` keeps the best-effort direct-UUID delete (the + fresh-write path), then falls through to a ``doc_id``-metadata + ``contains_any`` filter so the legacy UUID5 objects are reaped + too -- backwards-compatible with the rows already in production. +""" + + +class _FakeUnexpectedStatusCodeError(Exception): + """Stand-in for ``weaviate.exceptions.UnexpectedStatusCodeError``. + + The real class has a read-only ``status_code`` property; the + production code only reads ``.status_code`` off the exception, so + a plain attribute is enough. We patch the module-level reference + in the production module so the ``except`` clause catches ours. + """ + + def __init__(self, status_code: int) -> None: + super().__init__(f"status={status_code}") + self.status_code = status_code + + +def _make_vector() -> tuple[WeaviateVector, MagicMock]: + """Build a ``WeaviateVector`` without invoking ``__init__`` so we + never reach ``_init_client`` (which would try to connect). + + Only ``_client`` and ``_collection_name`` are touched by the + production code paths exercised here (``_get_uuids`` and + ``delete_by_ids``), so the rest of the instance dict is left bare. + + Returns the ``WeaviateVector`` together with its mocked client so + tests can reach ``client.collections.use(...).data`` (the same + path the production code walks) without pyrefly tripping over + chained ``.return_value`` access through the dynamic ``_client`` + instance attribute. + """ + # pyrefly cannot resolve chained ``.return_value`` access through + # dynamically-attached attributes (it infers ``v._client`` as the + # missing-on-class attribute and falls through to the stdlib + # ``collections`` module), so keep the mocked client in a local + # variable whose type pyrefly can see. + client: MagicMock = MagicMock() + client.collections.exists.return_value = True + client.collections.use.return_value = client.data + + v = WeaviateVector.__new__(WeaviateVector) + v._collection_name = "Test_Collection" + v._client = client + return v, client + + +def _collection(client: MagicMock) -> MagicMock: + """Resolve ``client.collections.use().data`` as a typed MagicMock. + + Centralised so the same ``client.data`` instance is reused -- the + production code reaches the Weaviate collection's ``data`` object + via ``v._client.collections.use(name).data``; we mirror that path + so ``delete_by_id`` / ``delete_many`` side effects wired on the + returned mock are observed by the production calls. + """ + col: MagicMock = client.collections.use.return_value + return col.data + + +class TestGetUuidsUsesDocId: + """_get_uuids must return one id per input document, positionally aligned.""" + + def test_doc_id_passes_through(self) -> None: + v, _client = _make_vector() + docs = [Document(page_content="x", metadata={"doc_id": "doc-aaa-1"})] + + uuids = v._get_uuids(docs) + + assert uuids == ["doc-aaa-1"] + + def test_missing_doc_id_falls_back_to_fresh_uuid4(self) -> None: + """A document without ``doc_id`` still needs a slot in the + returned list so the parallel ``objs`` list in ``add_texts`` + never goes out of alignment with the input documents. + """ + v, _client = _make_vector() + docs = [ + Document(page_content="x", metadata={"doc_id": "doc-aaa-1"}), + Document(page_content="y", metadata={}), + Document(page_content="z", metadata={"doc_id": "doc-ccc-3"}), + ] + + uuids = v._get_uuids(docs) + + assert uuids[0] == "doc-aaa-1" + assert uuids[2] == "doc-ccc-3" + # The middle slot is filled, not dropped. + assert len(uuids) == 3 + assert uuids[1] != "" + # And the fallback is a valid UUID4. + import uuid as _uuid + + parsed = _uuid.UUID(uuids[1]) + assert parsed.version == 4 + + def test_all_doc_ids_present(self) -> None: + v, _client = _make_vector() + docs = [ + Document(page_content="x", metadata={"doc_id": "a"}), + Document(page_content="y", metadata={"doc_id": "b"}), + Document(page_content="z", metadata={"doc_id": "c"}), + ] + + assert v._get_uuids(docs) == ["a", "b", "c"] + + +class TestDeleteByIdsBackwardCompatible: + """#41714 backwards-compatibility: legacy UUID5 objects whose Weaviate + UUID no longer matches ``index_node_id`` are reaped via the + ``doc_id``-metadata filter.""" + + def test_legacy_object_with_matching_doc_id_is_reaped(self) -> None: + """A legacy object whose ``doc_id`` matches the supplied + ``index_node_id`` must be reaped even though its Weaviate UUID + no longer matches. + """ + v, client = _make_vector() + col = _collection(client) + col.delete_by_id.side_effect = _FakeUnexpectedStatusCodeError(404) + + with patch.object(weaviate_vector_module, "UnexpectedStatusCodeError", _FakeUnexpectedStatusCodeError): + v.delete_by_ids(["doc-aaa-1"]) + + # Best-effort direct delete ran for each id. + assert col.delete_by_id.call_count == 1 + # And the metadata catch-up filter ran with the same ids. + col.delete_many.assert_called_once() + kwargs = col.delete_many.call_args.kwargs + where = kwargs["where"] + # Filter.by_property("doc_id").contains_any(["doc-aaa-1"]) + assert "doc_id" in str(where) + assert "doc-aaa-1" in str(where) + + def test_fresh_object_is_reaped_by_direct_delete(self) -> None: + """A fresh object whose Weaviate UUID equals the supplied + ``index_node_id`` is reaped by the direct delete, and the + metadata-filter pass is a redundant no-op (no rows match). + """ + v, client = _make_vector() + col = _collection(client) + + v.delete_by_ids(["doc-aaa-1"]) + + col.delete_by_id.assert_called_once_with("doc-aaa-1") + # The metadata filter always runs, but it's a no-op when the row + # is already gone (Weaviate's ``delete_many`` returns an empty + # result set against the cleared UUID). + col.delete_many.assert_called_once() + where = col.delete_many.call_args.kwargs["where"] + assert "doc_id" in str(where) + assert "doc-aaa-1" in str(where) + + def test_non_404_error_propagates(self) -> None: + """A non-404 (e.g. 500) error on direct delete must propagate, + not be swallowed as if it were a 404. + """ + import pytest + + v, client = _make_vector() + col = _collection(client) + col.delete_by_id.side_effect = _FakeUnexpectedStatusCodeError(500) + + with ( + patch.object(weaviate_vector_module, "UnexpectedStatusCodeError", _FakeUnexpectedStatusCodeError), + pytest.raises(_FakeUnexpectedStatusCodeError), + ): + v.delete_by_ids(["doc-aaa-1"]) + + def test_delete_many_404_is_swallowed(self) -> None: + """The metadata-filter 404 (column doesn't exist) is also + swallowed -- the schema may not have been migrated yet. + """ + v, client = _make_vector() + col = _collection(client) + col.delete_by_id.side_effect = _FakeUnexpectedStatusCodeError(404) + col.delete_many.side_effect = _FakeUnexpectedStatusCodeError(404) + + # Neither path raises. + with patch.object(weaviate_vector_module, "UnexpectedStatusCodeError", _FakeUnexpectedStatusCodeError): + v.delete_by_ids(["doc-aaa-1"]) + + +def test_init_client_with_valid_config(): + """Test successful client initialization with valid configuration.""" + config = WeaviateConfig( + endpoint="http://localhost:8080", + api_key="WVF5YThaHlkYwhGUSmCRgsX3tD5ngdN8pkih", + ) + + with patch("weaviate.connect_to_custom") as mock_connect: + mock_client = MagicMock() + mock_client.is_ready.return_value = True + mock_connect.return_value = mock_client + + vector = WeaviateVector( + collection_name="test_collection", + config=config, + attributes=["doc_id"], + ) + + assert vector._client == mock_client + mock_connect.assert_called_once() + call_kwargs = mock_connect.call_args[1] + assert call_kwargs["http_host"] == "localhost" + assert call_kwargs["http_port"] == 8080 + assert call_kwargs["http_secure"] is False + assert call_kwargs["grpc_host"] == "localhost" + assert call_kwargs["grpc_port"] == 50051 + assert call_kwargs["grpc_secure"] is False + assert call_kwargs["auth_credentials"] is not None + + +"""Unit tests for the Weaviate client integration telemetry header. + +Verifies that Dify registers the ``X-Weaviate-Client-Integration`` header on the Weaviate +client so server-side telemetry can attribute traffic to Dify, mirroring the behavior of other +official Weaviate integrations. +""" + + +class TestIntegrationVersion(unittest.TestCase): + """Tests for resolving the version reported in the integration header.""" + + def test_uses_dify_project_version_when_present(self): + with patch.object(weaviate_vector_module.dify_config.project, "version", "1.15.0"): + assert _integration_version() == "1.15.0" + + def test_strips_whitespace_from_version(self): + with patch.object(weaviate_vector_module.dify_config.project, "version", " 1.15.0 "): + assert _integration_version() == "1.15.0" + + def test_falls_back_to_package_metadata_when_version_empty(self): + with ( + patch.object(weaviate_vector_module.dify_config.project, "version", ""), + patch("dify_vdb_weaviate.weaviate_vector.importlib.metadata.version", return_value="9.9.9"), + ): + assert _integration_version() == "9.9.9" + + def test_falls_back_to_unknown_when_everything_fails(self): + with ( + patch.object(weaviate_vector_module.dify_config.project, "version", ""), + patch( + "dify_vdb_weaviate.weaviate_vector.importlib.metadata.version", + side_effect=Exception("no metadata"), + ), + ): + assert _integration_version() == "unknown" + + +class TestRegisterIntegration(unittest.TestCase): + """Tests for registering the integration header on the client.""" + + def test_configures_header_with_dify_name_and_version(self): + client = MagicMock() + with patch.object(weaviate_vector_module.dify_config.project, "version", "1.15.0"): + _register_integration(client) + + client.integrations.configure.assert_called_once() + config = client.integrations.configure.call_args.args[0] + header = config._to_header() + assert header == {"X-Weaviate-Client-Integration": "dify/1.15.0"} + assert _INTEGRATION_NAME == "dify" + + def test_swallows_errors_so_init_never_breaks(self): + client = MagicMock() + client.integrations.configure.side_effect = Exception("old client without integrations API") + + # Must not raise. + with patch.object(weaviate_vector_module.dify_config.project, "version", "1.15.0"): + _register_integration(client) + + +class TestInitClientRegistersIntegration(unittest.TestCase): + """Ensures the integration header is registered during client initialization.""" + + @override + def setUp(self): + weaviate_vector_module._weaviate_client = None + self.config = WeaviateConfig(endpoint="http://localhost:8080", api_key="test-key") + + @override + def tearDown(self): + weaviate_vector_module._weaviate_client = None + + @patch("dify_vdb_weaviate.weaviate_vector.weaviate") + def test_init_client_registers_integration(self, mock_weaviate_module): + mock_client = MagicMock() + mock_client.is_ready.return_value = True + mock_weaviate_module.connect_to_custom.return_value = mock_client + + with patch("dify_vdb_weaviate.weaviate_vector._register_integration") as mock_register: + WeaviateVector( + collection_name="Test_Collection_Node", + config=self.config, + attributes=["doc_id"], + ) + + mock_register.assert_called_once_with(mock_client) + + +"""Unit tests for Weaviate vector database implementation. + +Focuses on verifying that doc_type is properly handled in: +- Collection schema creation (_create_collection) +- Property migration (_ensure_properties) +- Vector search result metadata (search_by_vector) +- Full-text search result metadata (search_by_full_text) +""" + class TestWeaviateVector(unittest.TestCase): """Tests for WeaviateVector class with focus on doc_type metadata handling.""" @@ -933,3 +1260,7 @@ class TestWeaviateVectorFactory(unittest.TestCase): if __name__ == "__main__": unittest.main() + + +if __name__ == "__main__": + unittest.main() diff --git a/api/providers/vdb/vdb-weaviate/tests/unit_tests/test_doc_id_cleanup_41714.py b/api/providers/vdb/vdb-weaviate/tests/unit_tests/test_doc_id_cleanup_41714.py deleted file mode 100644 index dfc43d58b88..00000000000 --- a/api/providers/vdb/vdb-weaviate/tests/unit_tests/test_doc_id_cleanup_41714.py +++ /dev/null @@ -1,208 +0,0 @@ -"""Regression tests for issue #41714. - -The pre-fix code generated Weaviate object ids via -``uuid5(URL_NAMESPACE, page_content)`` (UUID v5 from content). The -cleanup path (``batch_clean_document_task`` -> ``index_processor.clean`` --> ``vector.delete_by_ids``) passes the segment's ``index_node_id`` -(random UUID v4 from the database), so the two never agreed and -``delete_by_id`` silently no-op'd on every cleanup pass. - -The fix: - -1. ``_get_uuids`` now uses the segment's ``doc_id`` (== ``index_node_id``) - from each document's metadata. New objects line up with the - cleanup path one-for-one. -2. Documents without a ``doc_id`` keep a stable slot via a freshly - generated ``uuid4`` so the parallel ``objs`` list in ``add_texts`` - never goes out of alignment with the input documents. -3. ``delete_by_ids`` keeps the best-effort direct-UUID delete (the - fresh-write path), then falls through to a ``doc_id``-metadata - ``contains_any`` filter so the legacy UUID5 objects are reaped - too -- backwards-compatible with the rows already in production. -""" - -from unittest.mock import MagicMock, patch - -from dify_vdb_weaviate import weaviate_vector as weaviate_vector_module -from dify_vdb_weaviate.weaviate_vector import WeaviateVector - -from core.rag.models.document import Document - - -class _FakeUnexpectedStatusCodeError(Exception): - """Stand-in for ``weaviate.exceptions.UnexpectedStatusCodeError``. - - The real class has a read-only ``status_code`` property; the - production code only reads ``.status_code`` off the exception, so - a plain attribute is enough. We patch the module-level reference - in the production module so the ``except`` clause catches ours. - """ - - def __init__(self, status_code: int) -> None: - super().__init__(f"status={status_code}") - self.status_code = status_code - - -def _make_vector() -> tuple[WeaviateVector, MagicMock]: - """Build a ``WeaviateVector`` without invoking ``__init__`` so we - never reach ``_init_client`` (which would try to connect). - - Only ``_client`` and ``_collection_name`` are touched by the - production code paths exercised here (``_get_uuids`` and - ``delete_by_ids``), so the rest of the instance dict is left bare. - - Returns the ``WeaviateVector`` together with its mocked client so - tests can reach ``client.collections.use(...).data`` (the same - path the production code walks) without pyrefly tripping over - chained ``.return_value`` access through the dynamic ``_client`` - instance attribute. - """ - # pyrefly cannot resolve chained ``.return_value`` access through - # dynamically-attached attributes (it infers ``v._client`` as the - # missing-on-class attribute and falls through to the stdlib - # ``collections`` module), so keep the mocked client in a local - # variable whose type pyrefly can see. - client: MagicMock = MagicMock() - client.collections.exists.return_value = True - client.collections.use.return_value = client.data - - v = WeaviateVector.__new__(WeaviateVector) - v._collection_name = "Test_Collection" - v._client = client - return v, client - - -def _collection(client: MagicMock) -> MagicMock: - """Resolve ``client.collections.use().data`` as a typed MagicMock. - - Centralised so the same ``client.data`` instance is reused -- the - production code reaches the Weaviate collection's ``data`` object - via ``v._client.collections.use(name).data``; we mirror that path - so ``delete_by_id`` / ``delete_many`` side effects wired on the - returned mock are observed by the production calls. - """ - col: MagicMock = client.collections.use.return_value - return col.data - - -class TestGetUuidsUsesDocId: - """_get_uuids must return one id per input document, positionally aligned.""" - - def test_doc_id_passes_through(self) -> None: - v, _client = _make_vector() - docs = [Document(page_content="x", metadata={"doc_id": "doc-aaa-1"})] - - uuids = v._get_uuids(docs) - - assert uuids == ["doc-aaa-1"] - - def test_missing_doc_id_falls_back_to_fresh_uuid4(self) -> None: - """A document without ``doc_id`` still needs a slot in the - returned list so the parallel ``objs`` list in ``add_texts`` - never goes out of alignment with the input documents. - """ - v, _client = _make_vector() - docs = [ - Document(page_content="x", metadata={"doc_id": "doc-aaa-1"}), - Document(page_content="y", metadata={}), - Document(page_content="z", metadata={"doc_id": "doc-ccc-3"}), - ] - - uuids = v._get_uuids(docs) - - assert uuids[0] == "doc-aaa-1" - assert uuids[2] == "doc-ccc-3" - # The middle slot is filled, not dropped. - assert len(uuids) == 3 - assert uuids[1] != "" - # And the fallback is a valid UUID4. - import uuid as _uuid - - parsed = _uuid.UUID(uuids[1]) - assert parsed.version == 4 - - def test_all_doc_ids_present(self) -> None: - v, _client = _make_vector() - docs = [ - Document(page_content="x", metadata={"doc_id": "a"}), - Document(page_content="y", metadata={"doc_id": "b"}), - Document(page_content="z", metadata={"doc_id": "c"}), - ] - - assert v._get_uuids(docs) == ["a", "b", "c"] - - -class TestDeleteByIdsBackwardCompatible: - """#41714 backwards-compatibility: legacy UUID5 objects whose Weaviate - UUID no longer matches ``index_node_id`` are reaped via the - ``doc_id``-metadata filter.""" - - def test_legacy_object_with_matching_doc_id_is_reaped(self) -> None: - """A legacy object whose ``doc_id`` matches the supplied - ``index_node_id`` must be reaped even though its Weaviate UUID - no longer matches. - """ - v, client = _make_vector() - col = _collection(client) - col.delete_by_id.side_effect = _FakeUnexpectedStatusCodeError(404) - - with patch.object(weaviate_vector_module, "UnexpectedStatusCodeError", _FakeUnexpectedStatusCodeError): - v.delete_by_ids(["doc-aaa-1"]) - - # Best-effort direct delete ran for each id. - assert col.delete_by_id.call_count == 1 - # And the metadata catch-up filter ran with the same ids. - col.delete_many.assert_called_once() - kwargs = col.delete_many.call_args.kwargs - where = kwargs["where"] - # Filter.by_property("doc_id").contains_any(["doc-aaa-1"]) - assert "doc_id" in str(where) - assert "doc-aaa-1" in str(where) - - def test_fresh_object_is_reaped_by_direct_delete(self) -> None: - """A fresh object whose Weaviate UUID equals the supplied - ``index_node_id`` is reaped by the direct delete, and the - metadata-filter pass is a redundant no-op (no rows match). - """ - v, client = _make_vector() - col = _collection(client) - - v.delete_by_ids(["doc-aaa-1"]) - - col.delete_by_id.assert_called_once_with("doc-aaa-1") - # The metadata filter always runs, but it's a no-op when the row - # is already gone (Weaviate's ``delete_many`` returns an empty - # result set against the cleared UUID). - col.delete_many.assert_called_once() - where = col.delete_many.call_args.kwargs["where"] - assert "doc_id" in str(where) - assert "doc-aaa-1" in str(where) - - def test_non_404_error_propagates(self) -> None: - """A non-404 (e.g. 500) error on direct delete must propagate, - not be swallowed as if it were a 404. - """ - import pytest - - v, client = _make_vector() - col = _collection(client) - col.delete_by_id.side_effect = _FakeUnexpectedStatusCodeError(500) - - with ( - patch.object(weaviate_vector_module, "UnexpectedStatusCodeError", _FakeUnexpectedStatusCodeError), - pytest.raises(_FakeUnexpectedStatusCodeError), - ): - v.delete_by_ids(["doc-aaa-1"]) - - def test_delete_many_404_is_swallowed(self) -> None: - """The metadata-filter 404 (column doesn't exist) is also - swallowed -- the schema may not have been migrated yet. - """ - v, client = _make_vector() - col = _collection(client) - col.delete_by_id.side_effect = _FakeUnexpectedStatusCodeError(404) - col.delete_many.side_effect = _FakeUnexpectedStatusCodeError(404) - - # Neither path raises. - with patch.object(weaviate_vector_module, "UnexpectedStatusCodeError", _FakeUnexpectedStatusCodeError): - v.delete_by_ids(["doc-aaa-1"]) diff --git a/api/providers/vdb/vdb-weaviate/tests/unit_tests/test_weavaite.py b/api/providers/vdb/vdb-weaviate/tests/unit_tests/test_weavaite.py deleted file mode 100644 index c773e4d5529..00000000000 --- a/api/providers/vdb/vdb-weaviate/tests/unit_tests/test_weavaite.py +++ /dev/null @@ -1,33 +0,0 @@ -from unittest.mock import MagicMock, patch - -from dify_vdb_weaviate.weaviate_vector import WeaviateConfig, WeaviateVector - - -def test_init_client_with_valid_config(): - """Test successful client initialization with valid configuration.""" - config = WeaviateConfig( - endpoint="http://localhost:8080", - api_key="WVF5YThaHlkYwhGUSmCRgsX3tD5ngdN8pkih", - ) - - with patch("weaviate.connect_to_custom") as mock_connect: - mock_client = MagicMock() - mock_client.is_ready.return_value = True - mock_connect.return_value = mock_client - - vector = WeaviateVector( - collection_name="test_collection", - config=config, - attributes=["doc_id"], - ) - - assert vector._client == mock_client - mock_connect.assert_called_once() - call_kwargs = mock_connect.call_args[1] - assert call_kwargs["http_host"] == "localhost" - assert call_kwargs["http_port"] == 8080 - assert call_kwargs["http_secure"] is False - assert call_kwargs["grpc_host"] == "localhost" - assert call_kwargs["grpc_port"] == 50051 - assert call_kwargs["grpc_secure"] is False - assert call_kwargs["auth_credentials"] is not None diff --git a/api/providers/vdb/vdb-weaviate/tests/unit_tests/test_weaviate_integration_header.py b/api/providers/vdb/vdb-weaviate/tests/unit_tests/test_weaviate_integration_header.py deleted file mode 100644 index 927ed344100..00000000000 --- a/api/providers/vdb/vdb-weaviate/tests/unit_tests/test_weaviate_integration_header.py +++ /dev/null @@ -1,103 +0,0 @@ -"""Unit tests for the Weaviate client integration telemetry header. - -Verifies that Dify registers the ``X-Weaviate-Client-Integration`` header on the Weaviate -client so server-side telemetry can attribute traffic to Dify, mirroring the behavior of other -official Weaviate integrations. -""" - -import unittest -from typing import override -from unittest.mock import MagicMock, patch - -from dify_vdb_weaviate import weaviate_vector as weaviate_vector_module -from dify_vdb_weaviate.weaviate_vector import ( - _INTEGRATION_NAME, - WeaviateConfig, - WeaviateVector, - _integration_version, - _register_integration, -) - - -class TestIntegrationVersion(unittest.TestCase): - """Tests for resolving the version reported in the integration header.""" - - def test_uses_dify_project_version_when_present(self): - with patch.object(weaviate_vector_module.dify_config.project, "version", "1.15.0"): - assert _integration_version() == "1.15.0" - - def test_strips_whitespace_from_version(self): - with patch.object(weaviate_vector_module.dify_config.project, "version", " 1.15.0 "): - assert _integration_version() == "1.15.0" - - def test_falls_back_to_package_metadata_when_version_empty(self): - with ( - patch.object(weaviate_vector_module.dify_config.project, "version", ""), - patch("dify_vdb_weaviate.weaviate_vector.importlib.metadata.version", return_value="9.9.9"), - ): - assert _integration_version() == "9.9.9" - - def test_falls_back_to_unknown_when_everything_fails(self): - with ( - patch.object(weaviate_vector_module.dify_config.project, "version", ""), - patch( - "dify_vdb_weaviate.weaviate_vector.importlib.metadata.version", - side_effect=Exception("no metadata"), - ), - ): - assert _integration_version() == "unknown" - - -class TestRegisterIntegration(unittest.TestCase): - """Tests for registering the integration header on the client.""" - - def test_configures_header_with_dify_name_and_version(self): - client = MagicMock() - with patch.object(weaviate_vector_module.dify_config.project, "version", "1.15.0"): - _register_integration(client) - - client.integrations.configure.assert_called_once() - config = client.integrations.configure.call_args.args[0] - header = config._to_header() - assert header == {"X-Weaviate-Client-Integration": "dify/1.15.0"} - assert _INTEGRATION_NAME == "dify" - - def test_swallows_errors_so_init_never_breaks(self): - client = MagicMock() - client.integrations.configure.side_effect = Exception("old client without integrations API") - - # Must not raise. - with patch.object(weaviate_vector_module.dify_config.project, "version", "1.15.0"): - _register_integration(client) - - -class TestInitClientRegistersIntegration(unittest.TestCase): - """Ensures the integration header is registered during client initialization.""" - - @override - def setUp(self): - weaviate_vector_module._weaviate_client = None - self.config = WeaviateConfig(endpoint="http://localhost:8080", api_key="test-key") - - @override - def tearDown(self): - weaviate_vector_module._weaviate_client = None - - @patch("dify_vdb_weaviate.weaviate_vector.weaviate") - def test_init_client_registers_integration(self, mock_weaviate_module): - mock_client = MagicMock() - mock_client.is_ready.return_value = True - mock_weaviate_module.connect_to_custom.return_value = mock_client - - with patch("dify_vdb_weaviate.weaviate_vector._register_integration") as mock_register: - WeaviateVector( - collection_name="Test_Collection_Node", - config=self.config, - attributes=["doc_id"], - ) - - mock_register.assert_called_once_with(mock_client) - - -if __name__ == "__main__": - unittest.main() diff --git a/api/pyrefly-local-excludes.txt b/api/pyrefly-local-excludes.txt index 1e625d6ffe7..37be2bedc50 100644 --- a/api/pyrefly-local-excludes.txt +++ b/api/pyrefly-local-excludes.txt @@ -1,20 +1,19 @@ -providers/vdb/vdb-analyticdb/tests/integration_tests/test_analyticdb.py -providers/vdb/vdb-analyticdb/tests/unit_tests/test_analyticdb_vector_sql.py -providers/vdb/vdb-elasticsearch/tests/unit_tests/test_elasticsearch_vector.py -providers/vdb/vdb-hologres/tests/integration_tests/conftest.py -providers/vdb/vdb-iris/tests/unit_tests/test_iris_vector.py -providers/vdb/vdb-lindorm/tests/integration_tests/test_lindorm.py -providers/vdb/vdb-lindorm/tests/unit_tests/test_lindorm_vector.py -providers/vdb/vdb-milvus/tests/unit_tests/test_milvus.py -providers/vdb/vdb-opengauss/tests/integration_tests/test_opengauss.py -providers/vdb/vdb-opensearch/tests/unit_tests/test_opensearch.py -providers/vdb/vdb-opensearch/tests/unit_tests/test_opensearch_vector.py -providers/vdb/vdb-oracle/tests/unit_tests/test_oraclevector.py -providers/vdb/vdb-pgvector/tests/unit_tests/test_pgvector.py -providers/vdb/vdb-qdrant/tests/unit_tests/test_qdrant_vector.py -providers/vdb/vdb-relyt/tests/unit_tests/test_relyt_vector.py -providers/vdb/vdb-tencent/tests/integration_tests/conftest.py +providers/vdb/vdb-analyticdb/tests/e2e_tests/test_analyticdb.py +providers/vdb/vdb-analyticdb/tests/unit_tests/src/dify_vdb_analyticdb/test_analyticdb_vector_sql.py +providers/vdb/vdb-elasticsearch/tests/unit_tests/src/dify_vdb_elasticsearch/test_elasticsearch_vector.py +providers/vdb/vdb-hologres/tests/e2e_tests/conftest.py +providers/vdb/vdb-iris/tests/unit_tests/src/dify_vdb_iris/test_iris_vector.py +providers/vdb/vdb-lindorm/tests/e2e_tests/test_lindorm.py +providers/vdb/vdb-lindorm/tests/unit_tests/src/dify_vdb_lindorm/test_lindorm_vector.py +providers/vdb/vdb-milvus/tests/unit_tests/src/dify_vdb_milvus/test_milvus_vector.py +providers/vdb/vdb-opengauss/tests/e2e_tests/test_opengauss.py +providers/vdb/vdb-opensearch/tests/unit_tests/src/dify_vdb_opensearch/test_opensearch_vector.py +providers/vdb/vdb-oracle/tests/unit_tests/src/dify_vdb_oracle/test_oraclevector.py +providers/vdb/vdb-pgvector/tests/unit_tests/src/dify_vdb_pgvector/test_pgvector.py +providers/vdb/vdb-qdrant/tests/unit_tests/src/dify_vdb_qdrant/test_qdrant_vector.py +providers/vdb/vdb-relyt/tests/unit_tests/src/dify_vdb_relyt/test_relyt_vector.py +providers/vdb/vdb-tencent/tests/e2e_tests/conftest.py providers/vdb/vdb-tidb-on-qdrant/src/dify_vdb_tidb_on_qdrant/tidb_on_qdrant_vector.py providers/vdb/vdb-tidb-on-qdrant/src/dify_vdb_tidb_on_qdrant/tidb_service.py -providers/vdb/vdb-tidb-vector/tests/unit_tests/test_tidb_vector.py -providers/vdb/vdb-weaviate/tests/unit_tests/test_weaviate_vector.py +providers/vdb/vdb-tidb-vector/tests/unit_tests/src/dify_vdb_tidb_vector/test_tidb_vector.py +providers/vdb/vdb-weaviate/tests/unit_tests/src/dify_vdb_weaviate/test_weaviate_vector.py diff --git a/api/tests/__init__.py b/api/tests/__init__.py index ced6188ce8e..528db521065 100644 --- a/api/tests/__init__.py +++ b/api/tests/__init__.py @@ -1 +1 @@ -"""Test suite root package (enables ``import tests.integration_tests...`` with ``pythonpath = .``).""" +"""Test suite root package (enables ``import tests.e2e_tests.middleware...`` with ``pythonpath = .``).""" diff --git a/api/tests/integration_tests/controllers/console/__init__.py b/api/tests/e2e_tests/__init__.py similarity index 100% rename from api/tests/integration_tests/controllers/console/__init__.py rename to api/tests/e2e_tests/__init__.py diff --git a/api/tests/test_containers_integration_tests/.ruff.toml b/api/tests/e2e_tests/containers/.ruff.toml similarity index 98% rename from api/tests/test_containers_integration_tests/.ruff.toml rename to api/tests/e2e_tests/containers/.ruff.toml index 5a63b0df947..b11d58d645e 100644 --- a/api/tests/test_containers_integration_tests/.ruff.toml +++ b/api/tests/e2e_tests/containers/.ruff.toml @@ -1,5 +1,5 @@ -extend = "../../.ruff.toml" -src = ["../.."] +extend = "../../../.ruff.toml" +src = ["../../.."] [lint] extend-select = ["ANN401", "ARG"] diff --git a/api/tests/integration_tests/controllers/console/app/__init__.py b/api/tests/e2e_tests/containers/__init__.py similarity index 100% rename from api/tests/integration_tests/controllers/console/app/__init__.py rename to api/tests/e2e_tests/containers/__init__.py diff --git a/api/tests/test_containers_integration_tests/commands/test_legacy_model_type_migration.py b/api/tests/e2e_tests/containers/commands/test_legacy_model_type_migration.py similarity index 99% rename from api/tests/test_containers_integration_tests/commands/test_legacy_model_type_migration.py rename to api/tests/e2e_tests/containers/commands/test_legacy_model_type_migration.py index 69ffc70dbda..830cdae3e58 100644 --- a/api/tests/test_containers_integration_tests/commands/test_legacy_model_type_migration.py +++ b/api/tests/e2e_tests/containers/commands/test_legacy_model_type_migration.py @@ -21,7 +21,7 @@ from tests.helpers.legacy_model_type_migration import ( ) _ALEMBIC_MIGRATION_PATH = ( - Path(__file__).resolve().parents[3] + Path(__file__).resolve().parents[4] / "migrations/versions/2026_08_27_1200-5578e028b2f2_migrate_legacy_model_types.py" ) diff --git a/api/tests/test_containers_integration_tests/conftest.py b/api/tests/e2e_tests/containers/conftest.py similarity index 99% rename from api/tests/test_containers_integration_tests/conftest.py rename to api/tests/e2e_tests/containers/conftest.py index 6890d7fa029..c62805f5fdb 100644 --- a/api/tests/test_containers_integration_tests/conftest.py +++ b/api/tests/e2e_tests/containers/conftest.py @@ -295,7 +295,7 @@ _container_manager = DifyTestContainers() def _get_migration_dir() -> Path: conftest_dir = Path(__file__).parent - return conftest_dir.parent.parent / "migrations" + return conftest_dir.parent.parent.parent / "migrations" def _get_engine_url(engine: Engine) -> str: diff --git a/api/tests/integration_tests/controllers/openapi/__init__.py b/api/tests/e2e_tests/containers/controllers/__init__.py similarity index 100% rename from api/tests/integration_tests/controllers/openapi/__init__.py rename to api/tests/e2e_tests/containers/controllers/__init__.py diff --git a/api/tests/test_containers_integration_tests/controllers/console/app/test_app_apis.py b/api/tests/e2e_tests/containers/controllers/console/app/test_app_apis.py similarity index 93% rename from api/tests/test_containers_integration_tests/controllers/console/app/test_app_apis.py rename to api/tests/e2e_tests/containers/controllers/console/app/test_app_apis.py index 4a7a9750b5a..65024005006 100644 --- a/api/tests/test_containers_integration_tests/controllers/console/app/test_app_apis.py +++ b/api/tests/e2e_tests/containers/controllers/console/app/test_app_apis.py @@ -5,7 +5,7 @@ from flask.testing import FlaskClient from sqlalchemy.orm import Session from models.model import AppMode -from tests.test_containers_integration_tests.controllers.console.helpers import ( +from tests.e2e_tests.containers.controllers.console.helpers import ( authenticate_console_client, create_console_account_and_tenant, create_console_app, diff --git a/api/tests/test_containers_integration_tests/controllers/console/app/test_app_detail_config_integrity.py b/api/tests/e2e_tests/containers/controllers/console/app/test_app_detail_config_integrity.py similarity index 97% rename from api/tests/test_containers_integration_tests/controllers/console/app/test_app_detail_config_integrity.py rename to api/tests/e2e_tests/containers/controllers/console/app/test_app_detail_config_integrity.py index 7ce218bba19..a92928d2daf 100644 --- a/api/tests/test_containers_integration_tests/controllers/console/app/test_app_detail_config_integrity.py +++ b/api/tests/e2e_tests/containers/controllers/console/app/test_app_detail_config_integrity.py @@ -6,7 +6,7 @@ from flask.testing import FlaskClient from sqlalchemy.orm import Session from models.model import AppMode, AppModelConfig -from tests.test_containers_integration_tests.controllers.console.helpers import ( +from tests.e2e_tests.containers.controllers.console.helpers import ( authenticate_console_client, create_console_account_and_tenant, create_console_app, diff --git a/api/tests/test_containers_integration_tests/controllers/console/app/test_chat_conversation_status_count_api.py b/api/tests/e2e_tests/containers/controllers/console/app/test_chat_conversation_status_count_api.py similarity index 98% rename from api/tests/test_containers_integration_tests/controllers/console/app/test_chat_conversation_status_count_api.py rename to api/tests/e2e_tests/containers/controllers/console/app/test_chat_conversation_status_count_api.py index 0d5deadb3b9..5f17ecf8df0 100644 --- a/api/tests/test_containers_integration_tests/controllers/console/app/test_chat_conversation_status_count_api.py +++ b/api/tests/e2e_tests/containers/controllers/console/app/test_chat_conversation_status_count_api.py @@ -16,7 +16,7 @@ from models.account import AccountStatus, TenantAccountRole, TenantStatus from models.enums import ConversationFromSource, CreatorUserRole from models.model import App, AppMode, Conversation, Message from models.workflow import WorkflowRun -from tests.test_containers_integration_tests.controllers.console.helpers import ensure_dify_setup +from tests.e2e_tests.containers.controllers.console.helpers import ensure_dify_setup def _create_account_and_tenant(db_session: Session) -> tuple[Account, Tenant]: diff --git a/api/tests/test_containers_integration_tests/controllers/console/app/test_conversation_read_timestamp.py b/api/tests/e2e_tests/containers/controllers/console/app/test_conversation_read_timestamp.py similarity index 96% rename from api/tests/test_containers_integration_tests/controllers/console/app/test_conversation_read_timestamp.py rename to api/tests/e2e_tests/containers/controllers/console/app/test_conversation_read_timestamp.py index a15f8991875..0559f14698f 100644 --- a/api/tests/test_containers_integration_tests/controllers/console/app/test_conversation_read_timestamp.py +++ b/api/tests/e2e_tests/containers/controllers/console/app/test_conversation_read_timestamp.py @@ -8,7 +8,7 @@ from werkzeug.exceptions import NotFound from controllers.console.app.conversation import _get_conversation from models.enums import ConversationFromSource from models.model import AppMode, Conversation -from tests.test_containers_integration_tests.controllers.console.helpers import ( +from tests.e2e_tests.containers.controllers.console.helpers import ( create_console_account_and_tenant, create_console_app, ) diff --git a/api/tests/test_containers_integration_tests/controllers/console/app/test_message.py b/api/tests/e2e_tests/containers/controllers/console/app/test_message.py similarity index 99% rename from api/tests/test_containers_integration_tests/controllers/console/app/test_message.py rename to api/tests/e2e_tests/containers/controllers/console/app/test_message.py index bcfc4a5981c..949ec8dd066 100644 --- a/api/tests/test_containers_integration_tests/controllers/console/app/test_message.py +++ b/api/tests/e2e_tests/containers/controllers/console/app/test_message.py @@ -18,7 +18,7 @@ from models.enums import ConversationFromSource, ConversationStatus, FeedbackRat from models.model import AppMode, Conversation, Message, MessageAnnotation, MessageFeedback from services.errors.conversation import ConversationNotExistsError from services.errors.message import MessageNotExistsError, SuggestedQuestionsAfterAnswerDisabledError -from tests.test_containers_integration_tests.controllers.console.helpers import ( +from tests.e2e_tests.containers.controllers.console.helpers import ( authenticate_console_client, create_console_account_and_tenant, create_console_app, diff --git a/api/tests/test_containers_integration_tests/controllers/console/app/test_statistic.py b/api/tests/e2e_tests/containers/controllers/console/app/test_statistic.py similarity index 99% rename from api/tests/test_containers_integration_tests/controllers/console/app/test_statistic.py rename to api/tests/e2e_tests/containers/controllers/console/app/test_statistic.py index ddf2974d62e..841d25f02bb 100644 --- a/api/tests/test_containers_integration_tests/controllers/console/app/test_statistic.py +++ b/api/tests/e2e_tests/containers/controllers/console/app/test_statistic.py @@ -11,7 +11,7 @@ from sqlalchemy.orm import Session from libs.datetime_utils import naive_utc_now from models.enums import ConversationFromSource, ConversationStatus, FeedbackFromSource, FeedbackRating, InvokeFrom from models.model import AppMode, Conversation, Message, MessageFeedback -from tests.test_containers_integration_tests.controllers.console.helpers import ( +from tests.e2e_tests.containers.controllers.console.helpers import ( authenticate_console_client, create_console_account_and_tenant, create_console_app, diff --git a/api/tests/test_containers_integration_tests/controllers/console/app/test_workflow_draft_variable.py b/api/tests/e2e_tests/containers/controllers/console/app/test_workflow_draft_variable.py similarity index 99% rename from api/tests/test_containers_integration_tests/controllers/console/app/test_workflow_draft_variable.py rename to api/tests/e2e_tests/containers/controllers/console/app/test_workflow_draft_variable.py index a071d22ee91..fbaf21c6dc0 100644 --- a/api/tests/test_containers_integration_tests/controllers/console/app/test_workflow_draft_variable.py +++ b/api/tests/e2e_tests/containers/controllers/console/app/test_workflow_draft_variable.py @@ -12,7 +12,7 @@ from graphon.variables.segments import StringSegment from models import Workflow from models.model import AppMode from models.workflow import WorkflowDraftVariable -from tests.test_containers_integration_tests.controllers.console.helpers import ( +from tests.e2e_tests.containers.controllers.console.helpers import ( authenticate_console_client, create_console_account_and_tenant, create_console_app, diff --git a/api/tests/test_containers_integration_tests/controllers/console/auth/test_data_source_bearer_auth.py b/api/tests/e2e_tests/containers/controllers/console/auth/test_data_source_bearer_auth.py similarity index 98% rename from api/tests/test_containers_integration_tests/controllers/console/auth/test_data_source_bearer_auth.py rename to api/tests/e2e_tests/containers/controllers/console/auth/test_data_source_bearer_auth.py index 1723f00bcb2..7bfdc50ba78 100644 --- a/api/tests/test_containers_integration_tests/controllers/console/auth/test_data_source_bearer_auth.py +++ b/api/tests/e2e_tests/containers/controllers/console/auth/test_data_source_bearer_auth.py @@ -13,7 +13,7 @@ from services.data_source.entities.api_key_auth import ( DataSourceApiKeyAuthBindingCreate, DataSourceApiKeyAuthCredentials, ) -from tests.test_containers_integration_tests.controllers.console.helpers import ( +from tests.e2e_tests.containers.controllers.console.helpers import ( authenticate_console_client, create_console_account_and_tenant, ) diff --git a/api/tests/test_containers_integration_tests/controllers/console/auth/test_data_source_oauth.py b/api/tests/e2e_tests/containers/controllers/console/auth/test_data_source_oauth.py similarity index 98% rename from api/tests/test_containers_integration_tests/controllers/console/auth/test_data_source_oauth.py rename to api/tests/e2e_tests/containers/controllers/console/auth/test_data_source_oauth.py index 307f80d62cd..bcab384b127 100644 --- a/api/tests/test_containers_integration_tests/controllers/console/auth/test_data_source_oauth.py +++ b/api/tests/e2e_tests/containers/controllers/console/auth/test_data_source_oauth.py @@ -9,7 +9,7 @@ from extensions.ext_application_services import ApplicationServices from models.source import DataSourceOauthBinding from services.data_source.entities.oauth import DataSourceOAuthCallback from services.data_source.oauth_service import InvalidDataSourceOAuthProviderError -from tests.test_containers_integration_tests.controllers.console.helpers import ( +from tests.e2e_tests.containers.controllers.console.helpers import ( authenticate_console_client, create_console_account_and_tenant, ) diff --git a/api/tests/test_containers_integration_tests/controllers/console/auth/test_oauth_server.py b/api/tests/e2e_tests/containers/controllers/console/auth/test_oauth_server.py similarity index 99% rename from api/tests/test_containers_integration_tests/controllers/console/auth/test_oauth_server.py rename to api/tests/e2e_tests/containers/controllers/console/auth/test_oauth_server.py index 9e87503c726..8b7a159cdf2 100644 --- a/api/tests/test_containers_integration_tests/controllers/console/auth/test_oauth_server.py +++ b/api/tests/e2e_tests/containers/controllers/console/auth/test_oauth_server.py @@ -20,7 +20,7 @@ from services.oauth_server_service import ( OAuthServerRequestError, OAuthServerUnauthorizedError, ) -from tests.test_containers_integration_tests.controllers.console.helpers import ( +from tests.e2e_tests.containers.controllers.console.helpers import ( authenticate_console_client, create_console_account_and_tenant, ensure_dify_setup, diff --git a/api/tests/test_containers_integration_tests/controllers/console/datasets/rag_pipeline/test_rag_pipeline.py b/api/tests/e2e_tests/containers/controllers/console/datasets/rag_pipeline/test_rag_pipeline.py similarity index 100% rename from api/tests/test_containers_integration_tests/controllers/console/datasets/rag_pipeline/test_rag_pipeline.py rename to api/tests/e2e_tests/containers/controllers/console/datasets/rag_pipeline/test_rag_pipeline.py diff --git a/api/tests/test_containers_integration_tests/controllers/console/datasets/rag_pipeline/test_rag_pipeline_datasets.py b/api/tests/e2e_tests/containers/controllers/console/datasets/rag_pipeline/test_rag_pipeline_datasets.py similarity index 100% rename from api/tests/test_containers_integration_tests/controllers/console/datasets/rag_pipeline/test_rag_pipeline_datasets.py rename to api/tests/e2e_tests/containers/controllers/console/datasets/rag_pipeline/test_rag_pipeline_datasets.py diff --git a/api/tests/test_containers_integration_tests/controllers/console/datasets/test_data_source.py b/api/tests/e2e_tests/containers/controllers/console/datasets/test_data_source.py similarity index 100% rename from api/tests/test_containers_integration_tests/controllers/console/datasets/test_data_source.py rename to api/tests/e2e_tests/containers/controllers/console/datasets/test_data_source.py diff --git a/api/tests/test_containers_integration_tests/controllers/console/datasets/test_external.py b/api/tests/e2e_tests/containers/controllers/console/datasets/test_external.py similarity index 97% rename from api/tests/test_containers_integration_tests/controllers/console/datasets/test_external.py rename to api/tests/e2e_tests/containers/controllers/console/datasets/test_external.py index d6b7e9e636f..4d21f3b37bb 100644 --- a/api/tests/test_containers_integration_tests/controllers/console/datasets/test_external.py +++ b/api/tests/e2e_tests/containers/controllers/console/datasets/test_external.py @@ -8,7 +8,7 @@ from flask.testing import FlaskClient from sqlalchemy.orm import Session from models.dataset import ExternalKnowledgeApis -from tests.test_containers_integration_tests.controllers.console.helpers import ( +from tests.e2e_tests.containers.controllers.console.helpers import ( authenticate_console_client, create_console_account_and_tenant, ) diff --git a/api/tests/test_containers_integration_tests/controllers/console/datasets/test_segments.py b/api/tests/e2e_tests/containers/controllers/console/datasets/test_segments.py similarity index 97% rename from api/tests/test_containers_integration_tests/controllers/console/datasets/test_segments.py rename to api/tests/e2e_tests/containers/controllers/console/datasets/test_segments.py index 6b9a4f2c961..033198c1bce 100644 --- a/api/tests/test_containers_integration_tests/controllers/console/datasets/test_segments.py +++ b/api/tests/e2e_tests/containers/controllers/console/datasets/test_segments.py @@ -10,7 +10,7 @@ from sqlalchemy.orm import Session from core.rag.index_processor.constant.index_type import IndexStructureType, IndexTechniqueType from models.dataset import Dataset, Document, DocumentSegment, DocumentSegmentSummary from models.enums import DataSourceType, DocumentCreatedFrom, IndexingStatus, SegmentStatus, SummaryStatus -from tests.test_containers_integration_tests.controllers.console.helpers import ( +from tests.e2e_tests.containers.controllers.console.helpers import ( authenticate_console_client, create_console_account_and_tenant, ) diff --git a/api/tests/test_containers_integration_tests/controllers/console/helpers.py b/api/tests/e2e_tests/containers/controllers/console/helpers.py similarity index 100% rename from api/tests/test_containers_integration_tests/controllers/console/helpers.py rename to api/tests/e2e_tests/containers/controllers/console/helpers.py diff --git a/api/tests/test_containers_integration_tests/controllers/console/test_api_based_extension.py b/api/tests/e2e_tests/containers/controllers/console/test_api_based_extension.py similarity index 98% rename from api/tests/test_containers_integration_tests/controllers/console/test_api_based_extension.py rename to api/tests/e2e_tests/containers/controllers/console/test_api_based_extension.py index 349ec8dceb5..38f2400034c 100644 --- a/api/tests/test_containers_integration_tests/controllers/console/test_api_based_extension.py +++ b/api/tests/e2e_tests/containers/controllers/console/test_api_based_extension.py @@ -13,7 +13,7 @@ from constants import HIDDEN_VALUE from core.helper.encrypter import encrypt_token from libs.rsa import generate_key_pair from models.api_based_extension import APIBasedExtension -from tests.test_containers_integration_tests.controllers.console.helpers import ( +from tests.e2e_tests.containers.controllers.console.helpers import ( authenticate_console_client, create_console_account_and_tenant, ) diff --git a/api/tests/test_containers_integration_tests/controllers/console/test_apikey.py b/api/tests/e2e_tests/containers/controllers/console/test_apikey.py similarity index 99% rename from api/tests/test_containers_integration_tests/controllers/console/test_apikey.py rename to api/tests/e2e_tests/containers/controllers/console/test_apikey.py index 791c9c7a340..09cbcccdb96 100644 --- a/api/tests/test_containers_integration_tests/controllers/console/test_apikey.py +++ b/api/tests/e2e_tests/containers/controllers/console/test_apikey.py @@ -18,7 +18,7 @@ from models.enums import ApiTokenType, DataSourceType from models.model import ApiToken, App, AppMode, DatasetApiTokenBinding from repositories.knowledge.dataset_api_key_repository import DatasetApiKeyRepository from services.auth.api_key_contracts import ApiKeyLimitExceededError -from tests.test_containers_integration_tests.controllers.console.helpers import ( +from tests.e2e_tests.containers.controllers.console.helpers import ( authenticate_console_client, create_console_account_and_tenant, create_console_app, diff --git a/api/tests/test_containers_integration_tests/controllers/console/test_feature.py b/api/tests/e2e_tests/containers/controllers/console/test_feature.py similarity index 96% rename from api/tests/test_containers_integration_tests/controllers/console/test_feature.py rename to api/tests/e2e_tests/containers/controllers/console/test_feature.py index e0e2a46ffc5..cb192a92b57 100644 --- a/api/tests/test_containers_integration_tests/controllers/console/test_feature.py +++ b/api/tests/e2e_tests/containers/controllers/console/test_feature.py @@ -10,7 +10,7 @@ from sqlalchemy.orm import Session from services.entities.feature_entities import FeatureModel, LimitationModel from services.feature_service import FeatureService -from tests.test_containers_integration_tests.controllers.console.helpers import ( +from tests.e2e_tests.containers.controllers.console.helpers import ( authenticate_console_client, create_console_account_and_tenant, ) diff --git a/api/tests/test_containers_integration_tests/controllers/console/test_files.py b/api/tests/e2e_tests/containers/controllers/console/test_files.py similarity index 98% rename from api/tests/test_containers_integration_tests/controllers/console/test_files.py rename to api/tests/e2e_tests/containers/controllers/console/test_files.py index e8f1f3ec778..fe38c25e23c 100644 --- a/api/tests/test_containers_integration_tests/controllers/console/test_files.py +++ b/api/tests/e2e_tests/containers/controllers/console/test_files.py @@ -10,7 +10,7 @@ from sqlalchemy.orm import Session from configs import dify_config from models.model import UploadFile -from tests.test_containers_integration_tests.controllers.console.helpers import ( +from tests.e2e_tests.containers.controllers.console.helpers import ( authenticate_console_client, create_console_account_and_tenant, ) diff --git a/api/tests/test_containers_integration_tests/controllers/console/test_setup.py b/api/tests/e2e_tests/containers/controllers/console/test_setup.py similarity index 98% rename from api/tests/test_containers_integration_tests/controllers/console/test_setup.py rename to api/tests/e2e_tests/containers/controllers/console/test_setup.py index f488a8860cd..d2dc449d997 100644 --- a/api/tests/test_containers_integration_tests/controllers/console/test_setup.py +++ b/api/tests/e2e_tests/containers/controllers/console/test_setup.py @@ -16,7 +16,7 @@ from models.model import DifySetup from services.account.contracts import SetupInput from services.account.service import AccountSetupProvisioner from services.setup_adapters import RedisSetupLock -from tests.test_containers_integration_tests.helpers import generate_valid_password +from tests.e2e_tests.containers.helpers import generate_valid_password @pytest.fixture diff --git a/api/tests/test_containers_integration_tests/controllers/console/test_spec.py b/api/tests/e2e_tests/containers/controllers/console/test_spec.py similarity index 93% rename from api/tests/test_containers_integration_tests/controllers/console/test_spec.py rename to api/tests/e2e_tests/containers/controllers/console/test_spec.py index 798cf852ebd..c02c962f0e5 100644 --- a/api/tests/test_containers_integration_tests/controllers/console/test_spec.py +++ b/api/tests/e2e_tests/containers/controllers/console/test_spec.py @@ -1,7 +1,7 @@ from flask.testing import FlaskClient from sqlalchemy.orm import Session -from tests.test_containers_integration_tests.controllers.console.helpers import ( +from tests.e2e_tests.containers.controllers.console.helpers import ( authenticate_console_client, create_console_account_and_tenant, ) diff --git a/api/tests/test_containers_integration_tests/controllers/console/workspace/test_members.py b/api/tests/e2e_tests/containers/controllers/console/workspace/test_members.py similarity index 100% rename from api/tests/test_containers_integration_tests/controllers/console/workspace/test_members.py rename to api/tests/e2e_tests/containers/controllers/console/workspace/test_members.py diff --git a/api/tests/test_containers_integration_tests/controllers/console/workspace/test_tool_provider.py b/api/tests/e2e_tests/containers/controllers/console/workspace/test_tool_provider.py similarity index 97% rename from api/tests/test_containers_integration_tests/controllers/console/workspace/test_tool_provider.py rename to api/tests/e2e_tests/containers/controllers/console/workspace/test_tool_provider.py index 35c8eb684f7..16d14435808 100644 --- a/api/tests/test_containers_integration_tests/controllers/console/workspace/test_tool_provider.py +++ b/api/tests/e2e_tests/containers/controllers/console/workspace/test_tool_provider.py @@ -9,7 +9,7 @@ from flask.testing import FlaskClient from sqlalchemy.orm import Session from core.tools.entities.api_entities import ToolProviderApiEntity -from tests.test_containers_integration_tests.controllers.console.helpers import ( +from tests.e2e_tests.containers.controllers.console.helpers import ( authenticate_console_client, create_console_account_and_tenant, ) diff --git a/api/tests/test_containers_integration_tests/controllers/console/workspace/test_workspace_wraps.py b/api/tests/e2e_tests/containers/controllers/console/workspace/test_workspace_wraps.py similarity index 100% rename from api/tests/test_containers_integration_tests/controllers/console/workspace/test_workspace_wraps.py rename to api/tests/e2e_tests/containers/controllers/console/workspace/test_workspace_wraps.py diff --git a/api/tests/integration_tests/factories/__init__.py b/api/tests/e2e_tests/containers/controllers/openapi/__init__.py similarity index 100% rename from api/tests/integration_tests/factories/__init__.py rename to api/tests/e2e_tests/containers/controllers/openapi/__init__.py diff --git a/api/tests/test_containers_integration_tests/controllers/openapi/conftest.py b/api/tests/e2e_tests/containers/controllers/openapi/conftest.py similarity index 96% rename from api/tests/test_containers_integration_tests/controllers/openapi/conftest.py rename to api/tests/e2e_tests/containers/controllers/openapi/conftest.py index f05a5f5ce39..c1ed36ee311 100644 --- a/api/tests/test_containers_integration_tests/controllers/openapi/conftest.py +++ b/api/tests/e2e_tests/containers/controllers/openapi/conftest.py @@ -16,8 +16,8 @@ from controllers.openapi.auth.requirements import Requirement, ResolveCaller from controllers.openapi.auth.subjects import subject_from_auth from libs.oauth_bearer import AuthContext from models import Account, Tenant -from tests.test_containers_integration_tests.helpers import accounts as account_fixtures -from tests.test_containers_integration_tests.helpers import generate_valid_password +from tests.e2e_tests.containers.helpers import accounts as account_fixtures +from tests.e2e_tests.containers.helpers import generate_valid_password _CLIENT_ID = "integration-cli" diff --git a/api/tests/test_containers_integration_tests/controllers/openapi/test_account.py b/api/tests/e2e_tests/containers/controllers/openapi/test_account.py similarity index 95% rename from api/tests/test_containers_integration_tests/controllers/openapi/test_account.py rename to api/tests/e2e_tests/containers/controllers/openapi/test_account.py index 3211b6d9cd3..26203f120bb 100644 --- a/api/tests/test_containers_integration_tests/controllers/openapi/test_account.py +++ b/api/tests/e2e_tests/containers/controllers/openapi/test_account.py @@ -8,7 +8,7 @@ from sqlalchemy.orm import Session from controllers.openapi.account import AccountApi from models import Account from models.account import TenantAccountRole -from tests.test_containers_integration_tests.controllers.openapi.conftest import add_tenant_for_account, context_for +from tests.e2e_tests.containers.controllers.openapi.conftest import add_tenant_for_account, context_for class TestAccountInfo: diff --git a/api/tests/test_containers_integration_tests/controllers/openapi/test_account_sessions.py b/api/tests/e2e_tests/containers/controllers/openapi/test_account_sessions.py similarity index 98% rename from api/tests/test_containers_integration_tests/controllers/openapi/test_account_sessions.py rename to api/tests/e2e_tests/containers/controllers/openapi/test_account_sessions.py index c7284c2898a..92924630922 100644 --- a/api/tests/test_containers_integration_tests/controllers/openapi/test_account_sessions.py +++ b/api/tests/e2e_tests/containers/controllers/openapi/test_account_sessions.py @@ -19,7 +19,7 @@ from controllers.openapi.account import ( ) from models import Account from models.oauth import OAuthAccessToken -from tests.test_containers_integration_tests.controllers.openapi.conftest import context_for +from tests.e2e_tests.containers.controllers.openapi.conftest import context_for @dataclass(frozen=True, slots=True) diff --git a/api/tests/test_containers_integration_tests/controllers/openapi/test_app_dsl.py b/api/tests/e2e_tests/containers/controllers/openapi/test_app_dsl.py similarity index 98% rename from api/tests/test_containers_integration_tests/controllers/openapi/test_app_dsl.py rename to api/tests/e2e_tests/containers/controllers/openapi/test_app_dsl.py index 5a41f79544b..a66d15593a6 100644 --- a/api/tests/test_containers_integration_tests/controllers/openapi/test_app_dsl.py +++ b/api/tests/e2e_tests/containers/controllers/openapi/test_app_dsl.py @@ -27,8 +27,8 @@ from services.app_dsl_service import CURRENT_DSL_VERSION from services.app_service import AppService, CreateAppParams from services.entities.dsl_entities import CheckDependenciesResult, Import, ImportStatus from services.errors.app import WorkflowNotFoundError -from tests.test_containers_integration_tests.helpers import accounts as account_fixtures -from tests.test_containers_integration_tests.helpers import generate_valid_password +from tests.e2e_tests.containers.helpers import accounts as account_fixtures +from tests.e2e_tests.containers.helpers import generate_valid_password class _Endpoint[T](Protocol): diff --git a/api/tests/test_containers_integration_tests/controllers/openapi/test_app_run.py b/api/tests/e2e_tests/containers/controllers/openapi/test_app_run.py similarity index 94% rename from api/tests/test_containers_integration_tests/controllers/openapi/test_app_run.py rename to api/tests/e2e_tests/containers/controllers/openapi/test_app_run.py index 92b7f91f7b2..ef6077983fb 100644 --- a/api/tests/test_containers_integration_tests/controllers/openapi/test_app_run.py +++ b/api/tests/e2e_tests/containers/controllers/openapi/test_app_run.py @@ -9,7 +9,7 @@ from sqlalchemy.orm import Session from controllers.openapi.app_run import AppRunTaskStopApi from models import Account, App from services.app_service import AppService, CreateAppParams -from tests.test_containers_integration_tests.controllers.openapi.conftest import context_for +from tests.e2e_tests.containers.controllers.openapi.conftest import context_for def _create_app(db_session: Session, account: Account, *, name: str = "Runner") -> App: diff --git a/api/tests/test_containers_integration_tests/controllers/openapi/test_apps.py b/api/tests/e2e_tests/containers/controllers/openapi/test_apps.py similarity index 100% rename from api/tests/test_containers_integration_tests/controllers/openapi/test_apps.py rename to api/tests/e2e_tests/containers/controllers/openapi/test_apps.py diff --git a/api/tests/test_containers_integration_tests/controllers/openapi/test_files.py b/api/tests/e2e_tests/containers/controllers/openapi/test_files.py similarity index 96% rename from api/tests/test_containers_integration_tests/controllers/openapi/test_files.py rename to api/tests/e2e_tests/containers/controllers/openapi/test_files.py index a74bbd8a7bd..d32ab95873e 100644 --- a/api/tests/test_containers_integration_tests/controllers/openapi/test_files.py +++ b/api/tests/e2e_tests/containers/controllers/openapi/test_files.py @@ -12,7 +12,7 @@ from controllers.openapi.auth.requirements import CheckWorkspaceMember from controllers.openapi.files import AppFileUploadApi from models import Account, App from services.app_service import AppService, CreateAppParams -from tests.test_containers_integration_tests.controllers.openapi.conftest import context_for +from tests.e2e_tests.containers.controllers.openapi.conftest import context_for def _create_app(db_session: Session, account: Account, *, name: str = "Uploader") -> App: diff --git a/api/tests/test_containers_integration_tests/controllers/openapi/test_workspaces.py b/api/tests/e2e_tests/containers/controllers/openapi/test_workspaces.py similarity index 96% rename from api/tests/test_containers_integration_tests/controllers/openapi/test_workspaces.py rename to api/tests/e2e_tests/containers/controllers/openapi/test_workspaces.py index bde2c3da93c..2b99781e49b 100644 --- a/api/tests/test_containers_integration_tests/controllers/openapi/test_workspaces.py +++ b/api/tests/e2e_tests/containers/controllers/openapi/test_workspaces.py @@ -11,7 +11,7 @@ from controllers.openapi._models import WorkspaceListQuery from controllers.openapi.workspaces import WorkspaceByIdApi, WorkspacesApi from models import Account from models.account import TenantAccountRole -from tests.test_containers_integration_tests.controllers.openapi.conftest import add_tenant_for_account, context_for +from tests.e2e_tests.containers.controllers.openapi.conftest import add_tenant_for_account, context_for class TestWorkspacesList: diff --git a/api/tests/integration_tests/services/__init__.py b/api/tests/e2e_tests/containers/controllers/service_api/__init__.py similarity index 100% rename from api/tests/integration_tests/services/__init__.py rename to api/tests/e2e_tests/containers/controllers/service_api/__init__.py diff --git a/api/tests/integration_tests/services/plugin/__init__.py b/api/tests/e2e_tests/containers/controllers/service_api/dataset/__init__.py similarity index 100% rename from api/tests/integration_tests/services/plugin/__init__.py rename to api/tests/e2e_tests/containers/controllers/service_api/dataset/__init__.py diff --git a/api/tests/test_containers_integration_tests/controllers/service_api/dataset/test_dataset.py b/api/tests/e2e_tests/containers/controllers/service_api/dataset/test_dataset.py similarity index 93% rename from api/tests/test_containers_integration_tests/controllers/service_api/dataset/test_dataset.py rename to api/tests/e2e_tests/containers/controllers/service_api/dataset/test_dataset.py index d72b6b736cc..c47f4236e48 100644 --- a/api/tests/test_containers_integration_tests/controllers/service_api/dataset/test_dataset.py +++ b/api/tests/e2e_tests/containers/controllers/service_api/dataset/test_dataset.py @@ -9,7 +9,7 @@ from sqlalchemy.orm import Session from models.account import Account from models.enums import TagType from models.model import Tag -from tests.test_containers_integration_tests.controllers.console.helpers import create_console_account_and_tenant +from tests.e2e_tests.containers.controllers.console.helpers import create_console_account_and_tenant @pytest.fixture diff --git a/api/tests/test_containers_integration_tests/controllers/service_api/dataset/test_segment.py b/api/tests/e2e_tests/containers/controllers/service_api/dataset/test_segment.py similarity index 97% rename from api/tests/test_containers_integration_tests/controllers/service_api/dataset/test_segment.py rename to api/tests/e2e_tests/containers/controllers/service_api/dataset/test_segment.py index 050149aec1b..9804cf2a902 100644 --- a/api/tests/test_containers_integration_tests/controllers/service_api/dataset/test_segment.py +++ b/api/tests/e2e_tests/containers/controllers/service_api/dataset/test_segment.py @@ -19,7 +19,7 @@ from models.enums import ( SummaryStatus, ) from models.model import ApiToken -from tests.test_containers_integration_tests.controllers.console.helpers import create_console_account_and_tenant +from tests.e2e_tests.containers.controllers.console.helpers import create_console_account_and_tenant def _create_dataset_graph(db_session: Session) -> tuple[Dataset, Document, DocumentSegment]: diff --git a/api/tests/test_containers_integration_tests/controllers/web/test_conversation.py b/api/tests/e2e_tests/containers/controllers/web/test_conversation.py similarity index 100% rename from api/tests/test_containers_integration_tests/controllers/web/test_conversation.py rename to api/tests/e2e_tests/containers/controllers/web/test_conversation.py diff --git a/api/tests/test_containers_integration_tests/controllers/web/test_human_input_form.py b/api/tests/e2e_tests/containers/controllers/web/test_human_input_form.py similarity index 100% rename from api/tests/test_containers_integration_tests/controllers/web/test_human_input_form.py rename to api/tests/e2e_tests/containers/controllers/web/test_human_input_form.py diff --git a/api/tests/test_containers_integration_tests/controllers/web/test_login.py b/api/tests/e2e_tests/containers/controllers/web/test_login.py similarity index 100% rename from api/tests/test_containers_integration_tests/controllers/web/test_login.py rename to api/tests/e2e_tests/containers/controllers/web/test_login.py diff --git a/api/tests/test_containers_integration_tests/controllers/web/test_site.py b/api/tests/e2e_tests/containers/controllers/web/test_site.py similarity index 100% rename from api/tests/test_containers_integration_tests/controllers/web/test_site.py rename to api/tests/e2e_tests/containers/controllers/web/test_site.py diff --git a/api/tests/test_containers_integration_tests/controllers/web/test_wraps.py b/api/tests/e2e_tests/containers/controllers/web/test_wraps.py similarity index 100% rename from api/tests/test_containers_integration_tests/controllers/web/test_wraps.py rename to api/tests/e2e_tests/containers/controllers/web/test_wraps.py diff --git a/api/tests/test_containers_integration_tests/core/__init__.py b/api/tests/e2e_tests/containers/core/__init__.py similarity index 100% rename from api/tests/test_containers_integration_tests/core/__init__.py rename to api/tests/e2e_tests/containers/core/__init__.py diff --git a/api/tests/test_containers_integration_tests/core/app/__init__.py b/api/tests/e2e_tests/containers/core/app/__init__.py similarity index 100% rename from api/tests/test_containers_integration_tests/core/app/__init__.py rename to api/tests/e2e_tests/containers/core/app/__init__.py diff --git a/api/tests/test_containers_integration_tests/core/app/layers/__init__.py b/api/tests/e2e_tests/containers/core/app/layers/__init__.py similarity index 100% rename from api/tests/test_containers_integration_tests/core/app/layers/__init__.py rename to api/tests/e2e_tests/containers/core/app/layers/__init__.py diff --git a/api/tests/test_containers_integration_tests/core/app/layers/test_pause_state_persist_layer.py b/api/tests/e2e_tests/containers/core/app/layers/test_pause_state_persist_layer.py similarity index 100% rename from api/tests/test_containers_integration_tests/core/app/layers/test_pause_state_persist_layer.py rename to api/tests/e2e_tests/containers/core/app/layers/test_pause_state_persist_layer.py diff --git a/api/tests/test_containers_integration_tests/core/rag/__init__.py b/api/tests/e2e_tests/containers/core/rag/__init__.py similarity index 100% rename from api/tests/test_containers_integration_tests/core/rag/__init__.py rename to api/tests/e2e_tests/containers/core/rag/__init__.py diff --git a/api/tests/integration_tests/services/retention/__init__.py b/api/tests/e2e_tests/containers/core/rag/pipeline/__init__.py similarity index 100% rename from api/tests/integration_tests/services/retention/__init__.py rename to api/tests/e2e_tests/containers/core/rag/pipeline/__init__.py diff --git a/api/tests/test_containers_integration_tests/core/rag/pipeline/test_queue_integration.py b/api/tests/e2e_tests/containers/core/rag/pipeline/test_queue_integration.py similarity index 100% rename from api/tests/test_containers_integration_tests/core/rag/pipeline/test_queue_integration.py rename to api/tests/e2e_tests/containers/core/rag/pipeline/test_queue_integration.py diff --git a/api/tests/test_containers_integration_tests/core/rag/retrieval/test_dataset_retrieval_integration.py b/api/tests/e2e_tests/containers/core/rag/retrieval/test_dataset_retrieval_integration.py similarity index 99% rename from api/tests/test_containers_integration_tests/core/rag/retrieval/test_dataset_retrieval_integration.py rename to api/tests/e2e_tests/containers/core/rag/retrieval/test_dataset_retrieval_integration.py index e67f23d4473..3976ea9fec4 100644 --- a/api/tests/test_containers_integration_tests/core/rag/retrieval/test_dataset_retrieval_integration.py +++ b/api/tests/e2e_tests/containers/core/rag/retrieval/test_dataset_retrieval_integration.py @@ -10,8 +10,8 @@ from core.rag.retrieval.dataset_retrieval import DatasetRetrieval from core.workflow.nodes.knowledge_retrieval.retrieval import KnowledgeRetrievalRequest from models.dataset import Dataset, Document from models.enums import DataSourceType, DocumentCreatedFrom, IndexingStatus -from tests.test_containers_integration_tests.helpers import accounts as account_fixtures -from tests.test_containers_integration_tests.helpers import generate_valid_password +from tests.e2e_tests.containers.helpers import accounts as account_fixtures +from tests.e2e_tests.containers.helpers import generate_valid_password class TestGetAvailableDatasetsIntegration: diff --git a/api/tests/test_containers_integration_tests/core/repositories/test_human_input_form_repository_impl.py b/api/tests/e2e_tests/containers/core/repositories/test_human_input_form_repository_impl.py similarity index 100% rename from api/tests/test_containers_integration_tests/core/repositories/test_human_input_form_repository_impl.py rename to api/tests/e2e_tests/containers/core/repositories/test_human_input_form_repository_impl.py diff --git a/api/tests/test_containers_integration_tests/core/workflow/test_human_input_resume_node_execution.py b/api/tests/e2e_tests/containers/core/workflow/test_human_input_resume_node_execution.py similarity index 100% rename from api/tests/test_containers_integration_tests/core/workflow/test_human_input_resume_node_execution.py rename to api/tests/e2e_tests/containers/core/workflow/test_human_input_resume_node_execution.py diff --git a/api/tests/integration_tests/tasks/__init__.py b/api/tests/e2e_tests/containers/factories/__init__.py similarity index 100% rename from api/tests/integration_tests/tasks/__init__.py rename to api/tests/e2e_tests/containers/factories/__init__.py diff --git a/api/tests/test_containers_integration_tests/factories/test_storage_key_loader.py b/api/tests/e2e_tests/containers/factories/test_storage_key_loader.py similarity index 100% rename from api/tests/test_containers_integration_tests/factories/test_storage_key_loader.py rename to api/tests/e2e_tests/containers/factories/test_storage_key_loader.py diff --git a/api/tests/test_containers_integration_tests/helpers/__init__.py b/api/tests/e2e_tests/containers/helpers/__init__.py similarity index 100% rename from api/tests/test_containers_integration_tests/helpers/__init__.py rename to api/tests/e2e_tests/containers/helpers/__init__.py diff --git a/api/tests/test_containers_integration_tests/helpers/accounts.py b/api/tests/e2e_tests/containers/helpers/accounts.py similarity index 100% rename from api/tests/test_containers_integration_tests/helpers/accounts.py rename to api/tests/e2e_tests/containers/helpers/accounts.py diff --git a/api/tests/test_containers_integration_tests/helpers/execution_extra_content.py b/api/tests/e2e_tests/containers/helpers/execution_extra_content.py similarity index 100% rename from api/tests/test_containers_integration_tests/helpers/execution_extra_content.py rename to api/tests/e2e_tests/containers/helpers/execution_extra_content.py diff --git a/api/tests/test_containers_integration_tests/libs/broadcast_channel/redis/test_channel.py b/api/tests/e2e_tests/containers/libs/broadcast_channel/redis/test_channel.py similarity index 100% rename from api/tests/test_containers_integration_tests/libs/broadcast_channel/redis/test_channel.py rename to api/tests/e2e_tests/containers/libs/broadcast_channel/redis/test_channel.py diff --git a/api/tests/test_containers_integration_tests/libs/broadcast_channel/redis/test_sharded_channel.py b/api/tests/e2e_tests/containers/libs/broadcast_channel/redis/test_sharded_channel.py similarity index 100% rename from api/tests/test_containers_integration_tests/libs/broadcast_channel/redis/test_sharded_channel.py rename to api/tests/e2e_tests/containers/libs/broadcast_channel/redis/test_sharded_channel.py diff --git a/api/tests/test_containers_integration_tests/libs/broadcast_channel/redis/test_streams_channel.py b/api/tests/e2e_tests/containers/libs/broadcast_channel/redis/test_streams_channel.py similarity index 100% rename from api/tests/test_containers_integration_tests/libs/broadcast_channel/redis/test_streams_channel.py rename to api/tests/e2e_tests/containers/libs/broadcast_channel/redis/test_streams_channel.py diff --git a/api/tests/test_containers_integration_tests/libs/test_api_token_cache_integration.py b/api/tests/e2e_tests/containers/libs/test_api_token_cache_integration.py similarity index 100% rename from api/tests/test_containers_integration_tests/libs/test_api_token_cache_integration.py rename to api/tests/e2e_tests/containers/libs/test_api_token_cache_integration.py diff --git a/api/tests/test_containers_integration_tests/libs/test_auto_renew_redis_lock_integration.py b/api/tests/e2e_tests/containers/libs/test_auto_renew_redis_lock_integration.py similarity index 100% rename from api/tests/test_containers_integration_tests/libs/test_auto_renew_redis_lock_integration.py rename to api/tests/e2e_tests/containers/libs/test_auto_renew_redis_lock_integration.py diff --git a/api/tests/test_containers_integration_tests/libs/test_rate_limiter_integration.py b/api/tests/e2e_tests/containers/libs/test_rate_limiter_integration.py similarity index 100% rename from api/tests/test_containers_integration_tests/libs/test_rate_limiter_integration.py rename to api/tests/e2e_tests/containers/libs/test_rate_limiter_integration.py diff --git a/api/tests/test_containers_integration_tests/models/test_account.py b/api/tests/e2e_tests/containers/models/test_account.py similarity index 100% rename from api/tests/test_containers_integration_tests/models/test_account.py rename to api/tests/e2e_tests/containers/models/test_account.py diff --git a/api/tests/test_containers_integration_tests/models/test_app_model_config.py b/api/tests/e2e_tests/containers/models/test_app_model_config.py similarity index 100% rename from api/tests/test_containers_integration_tests/models/test_app_model_config.py rename to api/tests/e2e_tests/containers/models/test_app_model_config.py diff --git a/api/tests/test_containers_integration_tests/models/test_conversation_message_inputs.py b/api/tests/e2e_tests/containers/models/test_conversation_message_inputs.py similarity index 100% rename from api/tests/test_containers_integration_tests/models/test_conversation_message_inputs.py rename to api/tests/e2e_tests/containers/models/test_conversation_message_inputs.py diff --git a/api/tests/test_containers_integration_tests/models/test_conversation_status_count.py b/api/tests/e2e_tests/containers/models/test_conversation_status_count.py similarity index 100% rename from api/tests/test_containers_integration_tests/models/test_conversation_status_count.py rename to api/tests/e2e_tests/containers/models/test_conversation_status_count.py diff --git a/api/tests/test_containers_integration_tests/models/test_dataset_models.py b/api/tests/e2e_tests/containers/models/test_dataset_models.py similarity index 100% rename from api/tests/test_containers_integration_tests/models/test_dataset_models.py rename to api/tests/e2e_tests/containers/models/test_dataset_models.py diff --git a/api/tests/test_containers_integration_tests/models/test_types_enum_text.py b/api/tests/e2e_tests/containers/models/test_types_enum_text.py similarity index 100% rename from api/tests/test_containers_integration_tests/models/test_types_enum_text.py rename to api/tests/e2e_tests/containers/models/test_types_enum_text.py diff --git a/api/tests/test_containers_integration_tests/models/test_workflow_node_execution_model.py b/api/tests/e2e_tests/containers/models/test_workflow_node_execution_model.py similarity index 100% rename from api/tests/test_containers_integration_tests/models/test_workflow_node_execution_model.py rename to api/tests/e2e_tests/containers/models/test_workflow_node_execution_model.py diff --git a/api/tests/test_containers_integration_tests/pyrefly.toml b/api/tests/e2e_tests/containers/pyrefly.toml similarity index 99% rename from api/tests/test_containers_integration_tests/pyrefly.toml rename to api/tests/e2e_tests/containers/pyrefly.toml index 1c9ab99cd11..419a66163c9 100644 --- a/api/tests/test_containers_integration_tests/pyrefly.toml +++ b/api/tests/e2e_tests/containers/pyrefly.toml @@ -1,6 +1,6 @@ preset = "strict" project-includes = ["."] -search-path = ["../.."] +search-path = ["../../.."] python-platform = "linux" python-version = "3.13.0" infer-with-first-use = true diff --git a/api/tests/test_containers_integration_tests/repositories/test_dataset_api_key_binding_deletion.py b/api/tests/e2e_tests/containers/repositories/test_dataset_api_key_binding_deletion.py similarity index 100% rename from api/tests/test_containers_integration_tests/repositories/test_dataset_api_key_binding_deletion.py rename to api/tests/e2e_tests/containers/repositories/test_dataset_api_key_binding_deletion.py diff --git a/api/tests/test_containers_integration_tests/repositories/test_recommended_app_catalog_repository.py b/api/tests/e2e_tests/containers/repositories/test_recommended_app_catalog_repository.py similarity index 100% rename from api/tests/test_containers_integration_tests/repositories/test_recommended_app_catalog_repository.py rename to api/tests/e2e_tests/containers/repositories/test_recommended_app_catalog_repository.py diff --git a/api/tests/test_containers_integration_tests/repositories/test_sqlalchemy_api_workflow_node_execution_repository.py b/api/tests/e2e_tests/containers/repositories/test_sqlalchemy_api_workflow_node_execution_repository.py similarity index 100% rename from api/tests/test_containers_integration_tests/repositories/test_sqlalchemy_api_workflow_node_execution_repository.py rename to api/tests/e2e_tests/containers/repositories/test_sqlalchemy_api_workflow_node_execution_repository.py diff --git a/api/tests/test_containers_integration_tests/repositories/test_sqlalchemy_api_workflow_run_repository.py b/api/tests/e2e_tests/containers/repositories/test_sqlalchemy_api_workflow_run_repository.py similarity index 100% rename from api/tests/test_containers_integration_tests/repositories/test_sqlalchemy_api_workflow_run_repository.py rename to api/tests/e2e_tests/containers/repositories/test_sqlalchemy_api_workflow_run_repository.py diff --git a/api/tests/test_containers_integration_tests/repositories/test_sqlalchemy_execution_extra_content_repository.py b/api/tests/e2e_tests/containers/repositories/test_sqlalchemy_execution_extra_content_repository.py similarity index 100% rename from api/tests/test_containers_integration_tests/repositories/test_sqlalchemy_execution_extra_content_repository.py rename to api/tests/e2e_tests/containers/repositories/test_sqlalchemy_execution_extra_content_repository.py diff --git a/api/tests/test_containers_integration_tests/repositories/test_sqlalchemy_workflow_node_execution_repository.py b/api/tests/e2e_tests/containers/repositories/test_sqlalchemy_workflow_node_execution_repository.py similarity index 100% rename from api/tests/test_containers_integration_tests/repositories/test_sqlalchemy_workflow_node_execution_repository.py rename to api/tests/e2e_tests/containers/repositories/test_sqlalchemy_workflow_node_execution_repository.py diff --git a/api/tests/test_containers_integration_tests/repositories/test_sqlalchemy_workflow_run_cleanup_repository.py b/api/tests/e2e_tests/containers/repositories/test_sqlalchemy_workflow_run_cleanup_repository.py similarity index 100% rename from api/tests/test_containers_integration_tests/repositories/test_sqlalchemy_workflow_run_cleanup_repository.py rename to api/tests/e2e_tests/containers/repositories/test_sqlalchemy_workflow_run_cleanup_repository.py diff --git a/api/tests/test_containers_integration_tests/repositories/test_sqlalchemy_workflow_trigger_log_repository.py b/api/tests/e2e_tests/containers/repositories/test_sqlalchemy_workflow_trigger_log_repository.py similarity index 100% rename from api/tests/test_containers_integration_tests/repositories/test_sqlalchemy_workflow_trigger_log_repository.py rename to api/tests/e2e_tests/containers/repositories/test_sqlalchemy_workflow_trigger_log_repository.py diff --git a/api/tests/test_containers_integration_tests/repositories/test_workflow_run_repository.py b/api/tests/e2e_tests/containers/repositories/test_workflow_run_repository.py similarity index 100% rename from api/tests/test_containers_integration_tests/repositories/test_workflow_run_repository.py rename to api/tests/e2e_tests/containers/repositories/test_workflow_run_repository.py diff --git a/api/tests/integration_tests/tools/__init__.py b/api/tests/e2e_tests/containers/services/__init__.py similarity index 100% rename from api/tests/integration_tests/tools/__init__.py rename to api/tests/e2e_tests/containers/services/__init__.py diff --git a/api/tests/test_containers_integration_tests/services/account/test_authentication.py b/api/tests/e2e_tests/containers/services/account/test_authentication.py similarity index 100% rename from api/tests/test_containers_integration_tests/services/account/test_authentication.py rename to api/tests/e2e_tests/containers/services/account/test_authentication.py diff --git a/api/tests/integration_tests/tools/api_tool/__init__.py b/api/tests/e2e_tests/containers/services/auth/__init__.py similarity index 100% rename from api/tests/integration_tests/tools/api_tool/__init__.py rename to api/tests/e2e_tests/containers/services/auth/__init__.py diff --git a/api/tests/test_containers_integration_tests/services/auth/test_auth_integration.py b/api/tests/e2e_tests/containers/services/auth/test_auth_integration.py similarity index 100% rename from api/tests/test_containers_integration_tests/services/auth/test_auth_integration.py rename to api/tests/e2e_tests/containers/services/auth/test_auth_integration.py diff --git a/api/tests/test_containers_integration_tests/services/dataset_collection_binding.py b/api/tests/e2e_tests/containers/services/dataset_collection_binding.py similarity index 100% rename from api/tests/test_containers_integration_tests/services/dataset_collection_binding.py rename to api/tests/e2e_tests/containers/services/dataset_collection_binding.py diff --git a/api/tests/test_containers_integration_tests/services/dataset_service_update_delete.py b/api/tests/e2e_tests/containers/services/dataset_service_update_delete.py similarity index 100% rename from api/tests/test_containers_integration_tests/services/dataset_service_update_delete.py rename to api/tests/e2e_tests/containers/services/dataset_service_update_delete.py diff --git a/api/tests/test_containers_integration_tests/services/document_service_status.py b/api/tests/e2e_tests/containers/services/document_service_status.py similarity index 100% rename from api/tests/test_containers_integration_tests/services/document_service_status.py rename to api/tests/e2e_tests/containers/services/document_service_status.py diff --git a/api/tests/integration_tests/tools/code/__init__.py b/api/tests/e2e_tests/containers/services/enterprise/__init__.py similarity index 100% rename from api/tests/integration_tests/tools/code/__init__.py rename to api/tests/e2e_tests/containers/services/enterprise/__init__.py diff --git a/api/tests/test_containers_integration_tests/services/enterprise/test_account_deletion_sync.py b/api/tests/e2e_tests/containers/services/enterprise/test_account_deletion_sync.py similarity index 100% rename from api/tests/test_containers_integration_tests/services/enterprise/test_account_deletion_sync.py rename to api/tests/e2e_tests/containers/services/enterprise/test_account_deletion_sync.py diff --git a/api/tests/integration_tests/vdb/__init__.py b/api/tests/e2e_tests/containers/services/plugin/__init__.py similarity index 100% rename from api/tests/integration_tests/vdb/__init__.py rename to api/tests/e2e_tests/containers/services/plugin/__init__.py diff --git a/api/tests/test_containers_integration_tests/services/plugin/test_plugin_permission_service.py b/api/tests/e2e_tests/containers/services/plugin/test_plugin_permission_service.py similarity index 100% rename from api/tests/test_containers_integration_tests/services/plugin/test_plugin_permission_service.py rename to api/tests/e2e_tests/containers/services/plugin/test_plugin_permission_service.py diff --git a/api/tests/integration_tests/workflow/__init__.py b/api/tests/e2e_tests/containers/services/rag_pipeline/__init__.py similarity index 100% rename from api/tests/integration_tests/workflow/__init__.py rename to api/tests/e2e_tests/containers/services/rag_pipeline/__init__.py diff --git a/api/tests/test_containers_integration_tests/services/rag_pipeline/test_rag_pipeline_service_db.py b/api/tests/e2e_tests/containers/services/rag_pipeline/test_rag_pipeline_service_db.py similarity index 100% rename from api/tests/test_containers_integration_tests/services/rag_pipeline/test_rag_pipeline_service_db.py rename to api/tests/e2e_tests/containers/services/rag_pipeline/test_rag_pipeline_service_db.py diff --git a/api/tests/test_containers_integration_tests/services/test_account_oauth_identity_lock.py b/api/tests/e2e_tests/containers/services/test_account_oauth_identity_lock.py similarity index 100% rename from api/tests/test_containers_integration_tests/services/test_account_oauth_identity_lock.py rename to api/tests/e2e_tests/containers/services/test_account_oauth_identity_lock.py diff --git a/api/tests/test_containers_integration_tests/services/test_advanced_prompt_template_service.py b/api/tests/e2e_tests/containers/services/test_advanced_prompt_template_service.py similarity index 100% rename from api/tests/test_containers_integration_tests/services/test_advanced_prompt_template_service.py rename to api/tests/e2e_tests/containers/services/test_advanced_prompt_template_service.py diff --git a/api/tests/test_containers_integration_tests/services/test_agent_service.py b/api/tests/e2e_tests/containers/services/test_agent_service.py similarity index 99% rename from api/tests/test_containers_integration_tests/services/test_agent_service.py rename to api/tests/e2e_tests/containers/services/test_agent_service.py index 96e7375d6af..615126cffd5 100644 --- a/api/tests/test_containers_integration_tests/services/test_agent_service.py +++ b/api/tests/e2e_tests/containers/services/test_agent_service.py @@ -11,8 +11,8 @@ from models.enums import ConversationFromSource, EndUserType, MessageFileBelongs from models.model import AppModelConfig, Conversation, EndUser, Message, MessageAgentThought from services.agent_service import AgentService from services.app_service import AppService, CreateAppParams -from tests.test_containers_integration_tests.helpers import accounts as account_fixtures -from tests.test_containers_integration_tests.helpers import generate_valid_password +from tests.e2e_tests.containers.helpers import accounts as account_fixtures +from tests.e2e_tests.containers.helpers import generate_valid_password class TestAgentService: diff --git a/api/tests/test_containers_integration_tests/services/test_annotation_service.py b/api/tests/e2e_tests/containers/services/test_annotation_service.py similarity index 99% rename from api/tests/test_containers_integration_tests/services/test_annotation_service.py rename to api/tests/e2e_tests/containers/services/test_annotation_service.py index 6cab51ba01c..b37d76da46f 100644 --- a/api/tests/test_containers_integration_tests/services/test_annotation_service.py +++ b/api/tests/e2e_tests/containers/services/test_annotation_service.py @@ -12,8 +12,8 @@ from models.model import MessageAnnotation from services.annotation_service import AppAnnotationService from services.app_ref_service import AnnotationRef, AppRef from services.app_service import AppService, CreateAppParams -from tests.test_containers_integration_tests.helpers import generate_valid_password -from tests.unit_tests.config_override import config_overrides_context +from tests.e2e_tests.containers.helpers import generate_valid_password +from tests.helpers.unit.config_override import config_overrides_context class TestAnnotationService: @@ -74,7 +74,7 @@ class TestAnnotationService: mock_external_service_dependencies["account_feature_service"].is_registration_allowed.return_value = True # Create account and tenant first - from tests.test_containers_integration_tests.helpers import accounts as account_fixtures + from tests.e2e_tests.containers.helpers import accounts as account_fixtures account = account_fixtures.create_account( email=fake.email(), diff --git a/api/tests/test_containers_integration_tests/services/test_app_dsl_service.py b/api/tests/e2e_tests/containers/services/test_app_dsl_service.py similarity index 99% rename from api/tests/test_containers_integration_tests/services/test_app_dsl_service.py rename to api/tests/e2e_tests/containers/services/test_app_dsl_service.py index 5ffe0a16994..0ecfc528a00 100644 --- a/api/tests/test_containers_integration_tests/services/test_app_dsl_service.py +++ b/api/tests/e2e_tests/containers/services/test_app_dsl_service.py @@ -56,8 +56,8 @@ from services.app_dsl_service import ( from services.app_service import AppService, CreateAppParams from services.dsl_version import check_version_compatibility from services.errors.app import WorkflowNotFoundError -from tests.test_containers_integration_tests.helpers import accounts as account_fixtures -from tests.test_containers_integration_tests.helpers import generate_valid_password +from tests.e2e_tests.containers.helpers import accounts as account_fixtures +from tests.e2e_tests.containers.helpers import generate_valid_password _DEFAULT_TENANT_ID = "00000000-0000-0000-0000-000000000001" _DEFAULT_ACCOUNT_ID = "00000000-0000-0000-0000-000000000002" diff --git a/api/tests/test_containers_integration_tests/services/test_app_generate_service.py b/api/tests/e2e_tests/containers/services/test_app_generate_service.py similarity index 99% rename from api/tests/test_containers_integration_tests/services/test_app_generate_service.py rename to api/tests/e2e_tests/containers/services/test_app_generate_service.py index 2a90e8a3273..9c39c8cb1bd 100644 --- a/api/tests/test_containers_integration_tests/services/test_app_generate_service.py +++ b/api/tests/e2e_tests/containers/services/test_app_generate_service.py @@ -16,7 +16,7 @@ from models.model import EndUser from models.workflow import Workflow from services.app_generate_service import AppGenerateService from services.errors.app import WorkflowIdFormatError, WorkflowNotFoundError -from tests.test_containers_integration_tests.helpers import generate_valid_password +from tests.e2e_tests.containers.helpers import generate_valid_password class TestAppGenerateService: @@ -168,7 +168,7 @@ class TestAppGenerateService: mock_external_service_dependencies["account_feature_service"].is_registration_allowed.return_value = True # Create account and tenant - from tests.test_containers_integration_tests.helpers import accounts as account_fixtures + from tests.e2e_tests.containers.helpers import accounts as account_fixtures account = account_fixtures.create_account( email=fake.email(), diff --git a/api/tests/test_containers_integration_tests/services/test_app_service.py b/api/tests/e2e_tests/containers/services/test_app_service.py similarity index 99% rename from api/tests/test_containers_integration_tests/services/test_app_service.py rename to api/tests/e2e_tests/containers/services/test_app_service.py index eaa6c999979..79b6e5e5a9a 100644 --- a/api/tests/test_containers_integration_tests/services/test_app_service.py +++ b/api/tests/e2e_tests/containers/services/test_app_service.py @@ -11,8 +11,8 @@ from constants.model_template import default_app_templates from models import Account from models.model import App, IconType from repositories.app.console_repository import ConsoleAppRepository -from tests.test_containers_integration_tests.helpers import accounts as account_fixtures -from tests.test_containers_integration_tests.helpers import generate_valid_password +from tests.e2e_tests.containers.helpers import accounts as account_fixtures +from tests.e2e_tests.containers.helpers import generate_valid_password # Delay import of AppService to avoid circular dependency # from services.app_service import AppService, AppListParams, CreateAppParams diff --git a/api/tests/test_containers_integration_tests/services/test_attachment_service.py b/api/tests/e2e_tests/containers/services/test_attachment_service.py similarity index 100% rename from api/tests/test_containers_integration_tests/services/test_attachment_service.py rename to api/tests/e2e_tests/containers/services/test_attachment_service.py diff --git a/api/tests/test_containers_integration_tests/services/test_audio_service_db.py b/api/tests/e2e_tests/containers/services/test_audio_service_db.py similarity index 99% rename from api/tests/test_containers_integration_tests/services/test_audio_service_db.py rename to api/tests/e2e_tests/containers/services/test_audio_service_db.py index 110d2cf100a..7356b7fcc06 100644 --- a/api/tests/test_containers_integration_tests/services/test_audio_service_db.py +++ b/api/tests/e2e_tests/containers/services/test_audio_service_db.py @@ -26,7 +26,7 @@ from models.enums import ConversationFromSource, MessageStatus from models.model import App, AppMode, Conversation, Message from services.app_ref_service import AppRef, MessageRef from services.audio_service import AudioService -from tests.test_containers_integration_tests.controllers.console.helpers import ( +from tests.e2e_tests.containers.controllers.console.helpers import ( create_console_account_and_tenant, create_console_app, ) diff --git a/api/tests/test_containers_integration_tests/services/test_billing_service.py b/api/tests/e2e_tests/containers/services/test_billing_service.py similarity index 100% rename from api/tests/test_containers_integration_tests/services/test_billing_service.py rename to api/tests/e2e_tests/containers/services/test_billing_service.py diff --git a/api/tests/test_containers_integration_tests/services/test_conversation_service.py b/api/tests/e2e_tests/containers/services/test_conversation_service.py similarity index 100% rename from api/tests/test_containers_integration_tests/services/test_conversation_service.py rename to api/tests/e2e_tests/containers/services/test_conversation_service.py diff --git a/api/tests/test_containers_integration_tests/services/test_conversation_service_variables.py b/api/tests/e2e_tests/containers/services/test_conversation_service_variables.py similarity index 100% rename from api/tests/test_containers_integration_tests/services/test_conversation_service_variables.py rename to api/tests/e2e_tests/containers/services/test_conversation_service_variables.py diff --git a/api/tests/test_containers_integration_tests/services/test_conversation_variable_updater.py b/api/tests/e2e_tests/containers/services/test_conversation_variable_updater.py similarity index 100% rename from api/tests/test_containers_integration_tests/services/test_conversation_variable_updater.py rename to api/tests/e2e_tests/containers/services/test_conversation_variable_updater.py diff --git a/api/tests/test_containers_integration_tests/services/test_credit_pool_service.py b/api/tests/e2e_tests/containers/services/test_credit_pool_service.py similarity index 100% rename from api/tests/test_containers_integration_tests/services/test_credit_pool_service.py rename to api/tests/e2e_tests/containers/services/test_credit_pool_service.py diff --git a/api/tests/test_containers_integration_tests/services/test_dataset_permission_service.py b/api/tests/e2e_tests/containers/services/test_dataset_permission_service.py similarity index 100% rename from api/tests/test_containers_integration_tests/services/test_dataset_permission_service.py rename to api/tests/e2e_tests/containers/services/test_dataset_permission_service.py diff --git a/api/tests/test_containers_integration_tests/services/test_dataset_service.py b/api/tests/e2e_tests/containers/services/test_dataset_service.py similarity index 100% rename from api/tests/test_containers_integration_tests/services/test_dataset_service.py rename to api/tests/e2e_tests/containers/services/test_dataset_service.py diff --git a/api/tests/test_containers_integration_tests/services/test_dataset_service_batch_update_document_status.py b/api/tests/e2e_tests/containers/services/test_dataset_service_batch_update_document_status.py similarity index 100% rename from api/tests/test_containers_integration_tests/services/test_dataset_service_batch_update_document_status.py rename to api/tests/e2e_tests/containers/services/test_dataset_service_batch_update_document_status.py diff --git a/api/tests/test_containers_integration_tests/services/test_dataset_service_create_dataset.py b/api/tests/e2e_tests/containers/services/test_dataset_service_create_dataset.py similarity index 100% rename from api/tests/test_containers_integration_tests/services/test_dataset_service_create_dataset.py rename to api/tests/e2e_tests/containers/services/test_dataset_service_create_dataset.py diff --git a/api/tests/test_containers_integration_tests/services/test_dataset_service_delete_dataset.py b/api/tests/e2e_tests/containers/services/test_dataset_service_delete_dataset.py similarity index 100% rename from api/tests/test_containers_integration_tests/services/test_dataset_service_delete_dataset.py rename to api/tests/e2e_tests/containers/services/test_dataset_service_delete_dataset.py diff --git a/api/tests/test_containers_integration_tests/services/test_dataset_service_document.py b/api/tests/e2e_tests/containers/services/test_dataset_service_document.py similarity index 100% rename from api/tests/test_containers_integration_tests/services/test_dataset_service_document.py rename to api/tests/e2e_tests/containers/services/test_dataset_service_document.py diff --git a/api/tests/test_containers_integration_tests/services/test_dataset_service_get_segments.py b/api/tests/e2e_tests/containers/services/test_dataset_service_get_segments.py similarity index 100% rename from api/tests/test_containers_integration_tests/services/test_dataset_service_get_segments.py rename to api/tests/e2e_tests/containers/services/test_dataset_service_get_segments.py diff --git a/api/tests/test_containers_integration_tests/services/test_dataset_service_permissions.py b/api/tests/e2e_tests/containers/services/test_dataset_service_permissions.py similarity index 99% rename from api/tests/test_containers_integration_tests/services/test_dataset_service_permissions.py rename to api/tests/e2e_tests/containers/services/test_dataset_service_permissions.py index 194bca1d8a7..39a4371d5eb 100644 --- a/api/tests/test_containers_integration_tests/services/test_dataset_service_permissions.py +++ b/api/tests/e2e_tests/containers/services/test_dataset_service_permissions.py @@ -24,7 +24,7 @@ from models.enums import DataSourceType from services.errors.base import NoPermissionError from services.knowledge.dataset_service import DatasetCollectionBindingService, DatasetPermissionService, DatasetService from services.knowledge.resource_scope import DatasetRef -from tests.unit_tests.config_override import config_overrides_context +from tests.helpers.unit.config_override import config_overrides_context class DatasetPermissionIntegrationFactory: diff --git a/api/tests/test_containers_integration_tests/services/test_dataset_service_retrieval.py b/api/tests/e2e_tests/containers/services/test_dataset_service_retrieval.py similarity index 100% rename from api/tests/test_containers_integration_tests/services/test_dataset_service_retrieval.py rename to api/tests/e2e_tests/containers/services/test_dataset_service_retrieval.py diff --git a/api/tests/test_containers_integration_tests/services/test_dataset_service_update_dataset.py b/api/tests/e2e_tests/containers/services/test_dataset_service_update_dataset.py similarity index 100% rename from api/tests/test_containers_integration_tests/services/test_dataset_service_update_dataset.py rename to api/tests/e2e_tests/containers/services/test_dataset_service_update_dataset.py diff --git a/api/tests/test_containers_integration_tests/services/test_delete_archived_workflow_run.py b/api/tests/e2e_tests/containers/services/test_delete_archived_workflow_run.py similarity index 100% rename from api/tests/test_containers_integration_tests/services/test_delete_archived_workflow_run.py rename to api/tests/e2e_tests/containers/services/test_delete_archived_workflow_run.py diff --git a/api/tests/test_containers_integration_tests/services/test_document_service_display_status.py b/api/tests/e2e_tests/containers/services/test_document_service_display_status.py similarity index 100% rename from api/tests/test_containers_integration_tests/services/test_document_service_display_status.py rename to api/tests/e2e_tests/containers/services/test_document_service_display_status.py diff --git a/api/tests/test_containers_integration_tests/services/test_document_service_rename_document.py b/api/tests/e2e_tests/containers/services/test_document_service_rename_document.py similarity index 100% rename from api/tests/test_containers_integration_tests/services/test_document_service_rename_document.py rename to api/tests/e2e_tests/containers/services/test_document_service_rename_document.py diff --git a/api/tests/test_containers_integration_tests/services/test_end_user_service.py b/api/tests/e2e_tests/containers/services/test_end_user_service.py similarity index 100% rename from api/tests/test_containers_integration_tests/services/test_end_user_service.py rename to api/tests/e2e_tests/containers/services/test_end_user_service.py diff --git a/api/tests/test_containers_integration_tests/services/test_feature_service.py b/api/tests/e2e_tests/containers/services/test_feature_service.py similarity index 100% rename from api/tests/test_containers_integration_tests/services/test_feature_service.py rename to api/tests/e2e_tests/containers/services/test_feature_service.py diff --git a/api/tests/test_containers_integration_tests/services/test_file_service.py b/api/tests/e2e_tests/containers/services/test_file_service.py similarity index 100% rename from api/tests/test_containers_integration_tests/services/test_file_service.py rename to api/tests/e2e_tests/containers/services/test_file_service.py diff --git a/api/tests/test_containers_integration_tests/services/test_file_service_zip_and_lookup.py b/api/tests/e2e_tests/containers/services/test_file_service_zip_and_lookup.py similarity index 100% rename from api/tests/test_containers_integration_tests/services/test_file_service_zip_and_lookup.py rename to api/tests/e2e_tests/containers/services/test_file_service_zip_and_lookup.py diff --git a/api/tests/test_containers_integration_tests/services/test_hit_testing_service.py b/api/tests/e2e_tests/containers/services/test_hit_testing_service.py similarity index 100% rename from api/tests/test_containers_integration_tests/services/test_hit_testing_service.py rename to api/tests/e2e_tests/containers/services/test_hit_testing_service.py diff --git a/api/tests/test_containers_integration_tests/services/test_human_input_delivery_test.py b/api/tests/e2e_tests/containers/services/test_human_input_delivery_test.py similarity index 100% rename from api/tests/test_containers_integration_tests/services/test_human_input_delivery_test.py rename to api/tests/e2e_tests/containers/services/test_human_input_delivery_test.py diff --git a/api/tests/test_containers_integration_tests/services/test_human_input_file_upload_service.py b/api/tests/e2e_tests/containers/services/test_human_input_file_upload_service.py similarity index 100% rename from api/tests/test_containers_integration_tests/services/test_human_input_file_upload_service.py rename to api/tests/e2e_tests/containers/services/test_human_input_file_upload_service.py diff --git a/api/tests/test_containers_integration_tests/services/test_message_export_service.py b/api/tests/e2e_tests/containers/services/test_message_export_service.py similarity index 100% rename from api/tests/test_containers_integration_tests/services/test_message_export_service.py rename to api/tests/e2e_tests/containers/services/test_message_export_service.py diff --git a/api/tests/test_containers_integration_tests/services/test_message_service.py b/api/tests/e2e_tests/containers/services/test_message_service.py similarity index 99% rename from api/tests/test_containers_integration_tests/services/test_message_service.py rename to api/tests/e2e_tests/containers/services/test_message_service.py index 427d81dbd29..46dc0c59929 100644 --- a/api/tests/test_containers_integration_tests/services/test_message_service.py +++ b/api/tests/e2e_tests/containers/services/test_message_service.py @@ -14,7 +14,7 @@ from services.errors.message import ( SuggestedQuestionsAfterAnswerDisabledError, ) from services.message_service import MessageService -from tests.test_containers_integration_tests.helpers import generate_valid_password +from tests.e2e_tests.containers.helpers import generate_valid_password class TestMessageService: @@ -88,7 +88,7 @@ class TestMessageService: mock_external_service_dependencies["account_feature_service"].is_registration_allowed.return_value = True # Create account and tenant first - from tests.test_containers_integration_tests.helpers import accounts as account_fixtures + from tests.e2e_tests.containers.helpers import accounts as account_fixtures account = account_fixtures.create_account( email=fake.email(), @@ -699,7 +699,7 @@ class TestMessageService: message = self._create_test_message(db_session_with_containers, app, conversation, account, fake) # Create another account - from tests.test_containers_integration_tests.helpers import accounts as account_fixtures + from tests.e2e_tests.containers.helpers import accounts as account_fixtures other_account = account_fixtures.create_account( email=fake.email(), diff --git a/api/tests/test_containers_integration_tests/services/test_message_service_execution_extra_content.py b/api/tests/e2e_tests/containers/services/test_message_service_execution_extra_content.py similarity index 97% rename from api/tests/test_containers_integration_tests/services/test_message_service_execution_extra_content.py rename to api/tests/e2e_tests/containers/services/test_message_service_execution_extra_content.py index 9da20d96f06..c914daf770e 100644 --- a/api/tests/test_containers_integration_tests/services/test_message_service_execution_extra_content.py +++ b/api/tests/e2e_tests/containers/services/test_message_service_execution_extra_content.py @@ -5,7 +5,7 @@ from sqlalchemy.orm import Session from models.human_input import HumanInputFormStatus from services.message_service import MessageService -from tests.test_containers_integration_tests.helpers.execution_extra_content import ( +from tests.e2e_tests.containers.helpers.execution_extra_content import ( create_human_input_message_fixture, ) diff --git a/api/tests/test_containers_integration_tests/services/test_message_service_extra_contents.py b/api/tests/e2e_tests/containers/services/test_message_service_extra_contents.py similarity index 97% rename from api/tests/test_containers_integration_tests/services/test_message_service_extra_contents.py rename to api/tests/e2e_tests/containers/services/test_message_service_extra_contents.py index 78a2e344ae0..0d2a8751ef5 100644 --- a/api/tests/test_containers_integration_tests/services/test_message_service_extra_contents.py +++ b/api/tests/e2e_tests/containers/services/test_message_service_extra_contents.py @@ -8,7 +8,7 @@ from libs.helper import to_timestamp from models.enums import ConversationFromSource from models.model import Message from services import message_service -from tests.test_containers_integration_tests.helpers.execution_extra_content import ( +from tests.e2e_tests.containers.helpers.execution_extra_content import ( create_human_input_message_fixture, ) diff --git a/api/tests/test_containers_integration_tests/services/test_messages_clean_service.py b/api/tests/e2e_tests/containers/services/test_messages_clean_service.py similarity index 100% rename from api/tests/test_containers_integration_tests/services/test_messages_clean_service.py rename to api/tests/e2e_tests/containers/services/test_messages_clean_service.py diff --git a/api/tests/test_containers_integration_tests/services/test_metadata_partial_update.py b/api/tests/e2e_tests/containers/services/test_metadata_partial_update.py similarity index 100% rename from api/tests/test_containers_integration_tests/services/test_metadata_partial_update.py rename to api/tests/e2e_tests/containers/services/test_metadata_partial_update.py diff --git a/api/tests/test_containers_integration_tests/services/test_metadata_service.py b/api/tests/e2e_tests/containers/services/test_metadata_service.py similarity index 100% rename from api/tests/test_containers_integration_tests/services/test_metadata_service.py rename to api/tests/e2e_tests/containers/services/test_metadata_service.py diff --git a/api/tests/test_containers_integration_tests/services/test_model_load_balancing_service.py b/api/tests/e2e_tests/containers/services/test_model_load_balancing_service.py similarity index 100% rename from api/tests/test_containers_integration_tests/services/test_model_load_balancing_service.py rename to api/tests/e2e_tests/containers/services/test_model_load_balancing_service.py diff --git a/api/tests/test_containers_integration_tests/services/test_model_provider_service.py b/api/tests/e2e_tests/containers/services/test_model_provider_service.py similarity index 100% rename from api/tests/test_containers_integration_tests/services/test_model_provider_service.py rename to api/tests/e2e_tests/containers/services/test_model_provider_service.py diff --git a/api/tests/test_containers_integration_tests/services/test_oauth_device_flow.py b/api/tests/e2e_tests/containers/services/test_oauth_device_flow.py similarity index 100% rename from api/tests/test_containers_integration_tests/services/test_oauth_device_flow.py rename to api/tests/e2e_tests/containers/services/test_oauth_device_flow.py diff --git a/api/tests/test_containers_integration_tests/services/test_restore_archived_workflow_run.py b/api/tests/e2e_tests/containers/services/test_restore_archived_workflow_run.py similarity index 100% rename from api/tests/test_containers_integration_tests/services/test_restore_archived_workflow_run.py rename to api/tests/e2e_tests/containers/services/test_restore_archived_workflow_run.py diff --git a/api/tests/test_containers_integration_tests/services/test_schedule_service.py b/api/tests/e2e_tests/containers/services/test_schedule_service.py similarity index 100% rename from api/tests/test_containers_integration_tests/services/test_schedule_service.py rename to api/tests/e2e_tests/containers/services/test_schedule_service.py diff --git a/api/tests/test_containers_integration_tests/services/test_tag_service.py b/api/tests/e2e_tests/containers/services/test_tag_service.py similarity index 100% rename from api/tests/test_containers_integration_tests/services/test_tag_service.py rename to api/tests/e2e_tests/containers/services/test_tag_service.py diff --git a/api/tests/test_containers_integration_tests/services/test_trigger_provider_service.py b/api/tests/e2e_tests/containers/services/test_trigger_provider_service.py similarity index 99% rename from api/tests/test_containers_integration_tests/services/test_trigger_provider_service.py rename to api/tests/e2e_tests/containers/services/test_trigger_provider_service.py index f6a90934a92..c42187afac6 100644 --- a/api/tests/test_containers_integration_tests/services/test_trigger_provider_service.py +++ b/api/tests/e2e_tests/containers/services/test_trigger_provider_service.py @@ -12,7 +12,7 @@ from models.account import Account, Tenant from models.provider_ids import TriggerProviderID from models.trigger import TriggerSubscription from services.trigger.trigger_provider_service import TriggerProviderService -from tests.test_containers_integration_tests.helpers import generate_valid_password +from tests.e2e_tests.containers.helpers import generate_valid_password MockExternalServiceDependencies = dict[str, MagicMock] @@ -68,7 +68,7 @@ class TestTriggerProviderService: """ fake = Faker() - from tests.test_containers_integration_tests.helpers import accounts as account_fixtures + from tests.e2e_tests.containers.helpers import accounts as account_fixtures # Setup mocks for account creation mock_external_service_dependencies["account_feature_service"].is_registration_allowed.return_value = True diff --git a/api/tests/test_containers_integration_tests/services/test_web_conversation_service.py b/api/tests/e2e_tests/containers/services/test_web_conversation_service.py similarity index 99% rename from api/tests/test_containers_integration_tests/services/test_web_conversation_service.py rename to api/tests/e2e_tests/containers/services/test_web_conversation_service.py index 3c703e0203c..928fa6344d0 100644 --- a/api/tests/test_containers_integration_tests/services/test_web_conversation_service.py +++ b/api/tests/e2e_tests/containers/services/test_web_conversation_service.py @@ -12,8 +12,8 @@ from models.model import Conversation, EndUser from models.web import PinnedConversation from services.app_service import AppService, CreateAppParams from services.web_conversation_service import WebConversationService -from tests.test_containers_integration_tests.helpers import accounts as account_fixtures -from tests.test_containers_integration_tests.helpers import generate_valid_password +from tests.e2e_tests.containers.helpers import accounts as account_fixtures +from tests.e2e_tests.containers.helpers import generate_valid_password class TestWebConversationService: diff --git a/api/tests/test_containers_integration_tests/services/test_webapp_auth_service.py b/api/tests/e2e_tests/containers/services/test_webapp_auth_service.py similarity index 100% rename from api/tests/test_containers_integration_tests/services/test_webapp_auth_service.py rename to api/tests/e2e_tests/containers/services/test_webapp_auth_service.py diff --git a/api/tests/test_containers_integration_tests/services/test_webhook_service.py b/api/tests/e2e_tests/containers/services/test_webhook_service.py similarity index 96% rename from api/tests/test_containers_integration_tests/services/test_webhook_service.py rename to api/tests/e2e_tests/containers/services/test_webhook_service.py index 74dbee453d4..f544acf43ca 100644 --- a/api/tests/test_containers_integration_tests/services/test_webhook_service.py +++ b/api/tests/e2e_tests/containers/services/test_webhook_service.py @@ -14,8 +14,8 @@ from models.model import App from models.trigger import AppTrigger, WorkflowWebhookTrigger from models.workflow import Workflow from services.trigger.webhook_service import WebhookService -from tests.test_containers_integration_tests.helpers import accounts as account_fixtures -from tests.test_containers_integration_tests.helpers import generate_valid_password +from tests.e2e_tests.containers.helpers import accounts as account_fixtures +from tests.e2e_tests.containers.helpers import generate_valid_password class WebhookIntegrationData(TypedDict): diff --git a/api/tests/test_containers_integration_tests/services/test_webhook_service_relationships.py b/api/tests/e2e_tests/containers/services/test_webhook_service_relationships.py similarity index 100% rename from api/tests/test_containers_integration_tests/services/test_webhook_service_relationships.py rename to api/tests/e2e_tests/containers/services/test_webhook_service_relationships.py diff --git a/api/tests/test_containers_integration_tests/services/test_workflow_app_service.py b/api/tests/e2e_tests/containers/services/test_workflow_app_service.py similarity index 99% rename from api/tests/test_containers_integration_tests/services/test_workflow_app_service.py rename to api/tests/e2e_tests/containers/services/test_workflow_app_service.py index 4576b1ec736..d8fb636057f 100644 --- a/api/tests/test_containers_integration_tests/services/test_workflow_app_service.py +++ b/api/tests/e2e_tests/containers/services/test_workflow_app_service.py @@ -15,8 +15,8 @@ from models.enums import CreatorUserRole, EndUserType from models.workflow import WorkflowAppLogCreatedFrom from repositories.workflow_app_log_query_repository import WorkflowAppLogQueryRepository from services.workflow_app_log_query_service import WorkflowAppLogQueryService -from tests.test_containers_integration_tests.helpers import accounts as account_fixtures -from tests.test_containers_integration_tests.helpers import generate_valid_password +from tests.e2e_tests.containers.helpers import accounts as account_fixtures +from tests.e2e_tests.containers.helpers import generate_valid_password class _WorkflowAppLogTestClient: diff --git a/api/tests/test_containers_integration_tests/services/test_workflow_draft_variable_service.py b/api/tests/e2e_tests/containers/services/test_workflow_draft_variable_service.py similarity index 100% rename from api/tests/test_containers_integration_tests/services/test_workflow_draft_variable_service.py rename to api/tests/e2e_tests/containers/services/test_workflow_draft_variable_service.py diff --git a/api/tests/test_containers_integration_tests/services/test_workflow_run_service.py b/api/tests/e2e_tests/containers/services/test_workflow_run_service.py similarity index 99% rename from api/tests/test_containers_integration_tests/services/test_workflow_run_service.py rename to api/tests/e2e_tests/containers/services/test_workflow_run_service.py index 74158b79bd7..8f78bd98b01 100644 --- a/api/tests/test_containers_integration_tests/services/test_workflow_run_service.py +++ b/api/tests/e2e_tests/containers/services/test_workflow_run_service.py @@ -20,8 +20,8 @@ from repositories.factory import DifyAPIRepositoryFactory from repositories.sqlalchemy_api_workflow_run_repository import DifyAPISQLAlchemyWorkflowRunRepository from services.app_service import AppService, CreateAppParams from services.workflow_run_service import WorkflowRunService -from tests.test_containers_integration_tests.helpers import accounts as account_fixtures -from tests.test_containers_integration_tests.helpers import generate_valid_password +from tests.e2e_tests.containers.helpers import accounts as account_fixtures +from tests.e2e_tests.containers.helpers import generate_valid_password class TestWorkflowRunService: diff --git a/api/tests/test_containers_integration_tests/services/test_workflow_service.py b/api/tests/e2e_tests/containers/services/test_workflow_service.py similarity index 99% rename from api/tests/test_containers_integration_tests/services/test_workflow_service.py rename to api/tests/e2e_tests/containers/services/test_workflow_service.py index 4817f445e23..86945178042 100644 --- a/api/tests/test_containers_integration_tests/services/test_workflow_service.py +++ b/api/tests/e2e_tests/containers/services/test_workflow_service.py @@ -19,7 +19,7 @@ from models.model import AppMode from models.workflow import WorkflowType from services.workflow_ref_service import WorkflowRef from services.workflow_service import WorkflowService -from tests.unit_tests.core.model_fixtures import make_model_instance +from tests.helpers.unit.core.model_fixtures import make_model_instance @dataclass(frozen=True) diff --git a/api/tests/integration_tests/workflow/nodes/__init__.py b/api/tests/e2e_tests/containers/services/tools/__init__.py similarity index 100% rename from api/tests/integration_tests/workflow/nodes/__init__.py rename to api/tests/e2e_tests/containers/services/tools/__init__.py diff --git a/api/tests/test_containers_integration_tests/services/tools/test_api_tools_manage_service.py b/api/tests/e2e_tests/containers/services/tools/test_api_tools_manage_service.py similarity index 100% rename from api/tests/test_containers_integration_tests/services/tools/test_api_tools_manage_service.py rename to api/tests/e2e_tests/containers/services/tools/test_api_tools_manage_service.py diff --git a/api/tests/test_containers_integration_tests/services/tools/test_mcp_tools_manage_service.py b/api/tests/e2e_tests/containers/services/tools/test_mcp_tools_manage_service.py similarity index 100% rename from api/tests/test_containers_integration_tests/services/tools/test_mcp_tools_manage_service.py rename to api/tests/e2e_tests/containers/services/tools/test_mcp_tools_manage_service.py diff --git a/api/tests/test_containers_integration_tests/services/tools/test_tools_transform_service.py b/api/tests/e2e_tests/containers/services/tools/test_tools_transform_service.py similarity index 100% rename from api/tests/test_containers_integration_tests/services/tools/test_tools_transform_service.py rename to api/tests/e2e_tests/containers/services/tools/test_tools_transform_service.py diff --git a/api/tests/test_containers_integration_tests/services/tools/test_workflow_tools_manage_service.py b/api/tests/e2e_tests/containers/services/tools/test_workflow_tools_manage_service.py similarity index 99% rename from api/tests/test_containers_integration_tests/services/tools/test_workflow_tools_manage_service.py rename to api/tests/e2e_tests/containers/services/tools/test_workflow_tools_manage_service.py index 4c7c1d666e2..d1ad8a7b03c 100644 --- a/api/tests/test_containers_integration_tests/services/tools/test_workflow_tools_manage_service.py +++ b/api/tests/e2e_tests/containers/services/tools/test_workflow_tools_manage_service.py @@ -12,8 +12,8 @@ from models.tools import WorkflowToolProvider from models.workflow import Workflow as WorkflowModel from services.app_service import AppService, CreateAppParams from services.tools.workflow_tools_manage_service import WorkflowToolManageService -from tests.test_containers_integration_tests.helpers import accounts as account_fixtures -from tests.test_containers_integration_tests.helpers import generate_valid_password +from tests.e2e_tests.containers.helpers import accounts as account_fixtures +from tests.e2e_tests.containers.helpers import generate_valid_password class TestWorkflowToolManageService: diff --git a/api/tests/integration_tests/workflow/nodes/code_executor/__init__.py b/api/tests/e2e_tests/containers/services/workflow/__init__.py similarity index 100% rename from api/tests/integration_tests/workflow/nodes/code_executor/__init__.py rename to api/tests/e2e_tests/containers/services/workflow/__init__.py diff --git a/api/tests/test_containers_integration_tests/services/workflow/test_workflow_converter.py b/api/tests/e2e_tests/containers/services/workflow/test_workflow_converter.py similarity index 100% rename from api/tests/test_containers_integration_tests/services/workflow/test_workflow_converter.py rename to api/tests/e2e_tests/containers/services/workflow/test_workflow_converter.py diff --git a/api/tests/test_containers_integration_tests/services/workflow/test_workflow_deletion.py b/api/tests/e2e_tests/containers/services/workflow/test_workflow_deletion.py similarity index 100% rename from api/tests/test_containers_integration_tests/services/workflow/test_workflow_deletion.py rename to api/tests/e2e_tests/containers/services/workflow/test_workflow_deletion.py diff --git a/api/tests/test_containers_integration_tests/services/workflow/test_workflow_event_snapshot_service.py b/api/tests/e2e_tests/containers/services/workflow/test_workflow_event_snapshot_service.py similarity index 100% rename from api/tests/test_containers_integration_tests/services/workflow/test_workflow_event_snapshot_service.py rename to api/tests/e2e_tests/containers/services/workflow/test_workflow_event_snapshot_service.py diff --git a/api/tests/test_containers_integration_tests/services/workflow/test_workflow_node_execution_service_repository.py b/api/tests/e2e_tests/containers/services/workflow/test_workflow_node_execution_service_repository.py similarity index 100% rename from api/tests/test_containers_integration_tests/services/workflow/test_workflow_node_execution_service_repository.py rename to api/tests/e2e_tests/containers/services/workflow/test_workflow_node_execution_service_repository.py diff --git a/api/tests/test_containers_integration_tests/services/workspace/test_invitations.py b/api/tests/e2e_tests/containers/services/workspace/test_invitations.py similarity index 100% rename from api/tests/test_containers_integration_tests/services/workspace/test_invitations.py rename to api/tests/e2e_tests/containers/services/workspace/test_invitations.py diff --git a/api/tests/integration_tests/workflow/nodes/knowledge_index/__init__.py b/api/tests/e2e_tests/containers/tasks/__init__.py similarity index 100% rename from api/tests/integration_tests/workflow/nodes/knowledge_index/__init__.py rename to api/tests/e2e_tests/containers/tasks/__init__.py diff --git a/api/tests/test_containers_integration_tests/tasks/test_add_document_to_index_task.py b/api/tests/e2e_tests/containers/tasks/test_add_document_to_index_task.py similarity index 100% rename from api/tests/test_containers_integration_tests/tasks/test_add_document_to_index_task.py rename to api/tests/e2e_tests/containers/tasks/test_add_document_to_index_task.py diff --git a/api/tests/test_containers_integration_tests/tasks/test_batch_clean_document_task.py b/api/tests/e2e_tests/containers/tasks/test_batch_clean_document_task.py similarity index 100% rename from api/tests/test_containers_integration_tests/tasks/test_batch_clean_document_task.py rename to api/tests/e2e_tests/containers/tasks/test_batch_clean_document_task.py diff --git a/api/tests/test_containers_integration_tests/tasks/test_batch_create_segment_to_index_task.py b/api/tests/e2e_tests/containers/tasks/test_batch_create_segment_to_index_task.py similarity index 100% rename from api/tests/test_containers_integration_tests/tasks/test_batch_create_segment_to_index_task.py rename to api/tests/e2e_tests/containers/tasks/test_batch_create_segment_to_index_task.py diff --git a/api/tests/test_containers_integration_tests/tasks/test_clean_dataset_task.py b/api/tests/e2e_tests/containers/tasks/test_clean_dataset_task.py similarity index 100% rename from api/tests/test_containers_integration_tests/tasks/test_clean_dataset_task.py rename to api/tests/e2e_tests/containers/tasks/test_clean_dataset_task.py diff --git a/api/tests/test_containers_integration_tests/tasks/test_clean_notion_document_task.py b/api/tests/e2e_tests/containers/tasks/test_clean_notion_document_task.py similarity index 99% rename from api/tests/test_containers_integration_tests/tasks/test_clean_notion_document_task.py rename to api/tests/e2e_tests/containers/tasks/test_clean_notion_document_task.py index ed1c37c7317..911cd35dca5 100644 --- a/api/tests/test_containers_integration_tests/tasks/test_clean_notion_document_task.py +++ b/api/tests/e2e_tests/containers/tasks/test_clean_notion_document_task.py @@ -18,8 +18,8 @@ from core.rag.index_processor.constant.index_type import IndexStructureType from models.dataset import Dataset, Document, DocumentSegment from models.enums import DataSourceType, DocumentCreatedFrom, IndexingStatus, SegmentStatus from tasks.clean_notion_document_task import clean_notion_document_task -from tests.test_containers_integration_tests.helpers import accounts as account_fixtures -from tests.test_containers_integration_tests.helpers import generate_valid_password +from tests.e2e_tests.containers.helpers import accounts as account_fixtures +from tests.e2e_tests.containers.helpers import generate_valid_password def _count_documents(session: Session, condition: ColumnElement[bool]) -> int: diff --git a/api/tests/test_containers_integration_tests/tasks/test_create_segment_to_index_task.py b/api/tests/e2e_tests/containers/tasks/test_create_segment_to_index_task.py similarity index 100% rename from api/tests/test_containers_integration_tests/tasks/test_create_segment_to_index_task.py rename to api/tests/e2e_tests/containers/tasks/test_create_segment_to_index_task.py diff --git a/api/tests/test_containers_integration_tests/tasks/test_dataset_indexing_task.py b/api/tests/e2e_tests/containers/tasks/test_dataset_indexing_task.py similarity index 99% rename from api/tests/test_containers_integration_tests/tasks/test_dataset_indexing_task.py rename to api/tests/e2e_tests/containers/tasks/test_dataset_indexing_task.py index fa5496ee03c..3006d7adecb 100644 --- a/api/tests/test_containers_integration_tests/tasks/test_dataset_indexing_task.py +++ b/api/tests/e2e_tests/containers/tasks/test_dataset_indexing_task.py @@ -23,7 +23,7 @@ from tasks.document_indexing_task import ( normal_document_indexing_task, priority_document_indexing_task, ) -from tests.unit_tests.config_override import config_overrides_context +from tests.helpers.unit.config_override import config_overrides_context class _TrackedSessionContext: diff --git a/api/tests/test_containers_integration_tests/tasks/test_deal_dataset_vector_index_task.py b/api/tests/e2e_tests/containers/tasks/test_deal_dataset_vector_index_task.py similarity index 99% rename from api/tests/test_containers_integration_tests/tasks/test_deal_dataset_vector_index_task.py rename to api/tests/e2e_tests/containers/tasks/test_deal_dataset_vector_index_task.py index b301da469df..39999aff3cd 100644 --- a/api/tests/test_containers_integration_tests/tasks/test_deal_dataset_vector_index_task.py +++ b/api/tests/e2e_tests/containers/tasks/test_deal_dataset_vector_index_task.py @@ -18,8 +18,8 @@ from core.rag.index_processor.constant.index_type import IndexStructureType from models.dataset import Dataset, Document, DocumentSegment from models.enums import DataSourceType, DocumentCreatedFrom, IndexingStatus, SegmentStatus from tasks.deal_dataset_vector_index_task import deal_dataset_vector_index_task -from tests.test_containers_integration_tests.helpers import accounts as account_fixtures -from tests.test_containers_integration_tests.helpers import generate_valid_password +from tests.e2e_tests.containers.helpers import accounts as account_fixtures +from tests.e2e_tests.containers.helpers import generate_valid_password class TestDealDatasetVectorIndexTask: diff --git a/api/tests/test_containers_integration_tests/tasks/test_delete_account_task.py b/api/tests/e2e_tests/containers/tasks/test_delete_account_task.py similarity index 100% rename from api/tests/test_containers_integration_tests/tasks/test_delete_account_task.py rename to api/tests/e2e_tests/containers/tasks/test_delete_account_task.py diff --git a/api/tests/test_containers_integration_tests/tasks/test_delete_segment_from_index_task.py b/api/tests/e2e_tests/containers/tasks/test_delete_segment_from_index_task.py similarity index 100% rename from api/tests/test_containers_integration_tests/tasks/test_delete_segment_from_index_task.py rename to api/tests/e2e_tests/containers/tasks/test_delete_segment_from_index_task.py diff --git a/api/tests/test_containers_integration_tests/tasks/test_disable_segment_from_index_task.py b/api/tests/e2e_tests/containers/tasks/test_disable_segment_from_index_task.py similarity index 100% rename from api/tests/test_containers_integration_tests/tasks/test_disable_segment_from_index_task.py rename to api/tests/e2e_tests/containers/tasks/test_disable_segment_from_index_task.py diff --git a/api/tests/test_containers_integration_tests/tasks/test_disable_segments_from_index_task.py b/api/tests/e2e_tests/containers/tasks/test_disable_segments_from_index_task.py similarity index 100% rename from api/tests/test_containers_integration_tests/tasks/test_disable_segments_from_index_task.py rename to api/tests/e2e_tests/containers/tasks/test_disable_segments_from_index_task.py diff --git a/api/tests/test_containers_integration_tests/tasks/test_document_indexing_sync_task.py b/api/tests/e2e_tests/containers/tasks/test_document_indexing_sync_task.py similarity index 100% rename from api/tests/test_containers_integration_tests/tasks/test_document_indexing_sync_task.py rename to api/tests/e2e_tests/containers/tasks/test_document_indexing_sync_task.py diff --git a/api/tests/test_containers_integration_tests/tasks/test_document_indexing_task.py b/api/tests/e2e_tests/containers/tasks/test_document_indexing_task.py similarity index 99% rename from api/tests/test_containers_integration_tests/tasks/test_document_indexing_task.py rename to api/tests/e2e_tests/containers/tasks/test_document_indexing_task.py index a0cdf1980ea..47d84c0e72b 100644 --- a/api/tests/test_containers_integration_tests/tasks/test_document_indexing_task.py +++ b/api/tests/e2e_tests/containers/tasks/test_document_indexing_task.py @@ -19,7 +19,7 @@ from tasks.document_indexing_task import ( normal_document_indexing_task, # New normal task priority_document_indexing_task, # New priority task ) -from tests.unit_tests.config_override import config_overrides_context +from tests.helpers.unit.config_override import config_overrides_context class TestDocumentIndexingTasks: diff --git a/api/tests/test_containers_integration_tests/tasks/test_document_indexing_update_task.py b/api/tests/e2e_tests/containers/tasks/test_document_indexing_update_task.py similarity index 100% rename from api/tests/test_containers_integration_tests/tasks/test_document_indexing_update_task.py rename to api/tests/e2e_tests/containers/tasks/test_document_indexing_update_task.py diff --git a/api/tests/test_containers_integration_tests/tasks/test_duplicate_document_indexing_task.py b/api/tests/e2e_tests/containers/tasks/test_duplicate_document_indexing_task.py similarity index 99% rename from api/tests/test_containers_integration_tests/tasks/test_duplicate_document_indexing_task.py rename to api/tests/e2e_tests/containers/tasks/test_duplicate_document_indexing_task.py index 8b6062f3877..bd99949137c 100644 --- a/api/tests/test_containers_integration_tests/tasks/test_duplicate_document_indexing_task.py +++ b/api/tests/e2e_tests/containers/tasks/test_duplicate_document_indexing_task.py @@ -19,7 +19,7 @@ from tasks.duplicate_document_indexing_task import ( normal_duplicate_document_indexing_task, # New normal task priority_duplicate_document_indexing_task, # New priority task ) -from tests.unit_tests.config_override import config_overrides_context +from tests.helpers.unit.config_override import config_overrides_context class TestDuplicateDocumentIndexingTasks: diff --git a/api/tests/test_containers_integration_tests/tasks/test_enable_segments_to_index_task.py b/api/tests/e2e_tests/containers/tasks/test_enable_segments_to_index_task.py similarity index 100% rename from api/tests/test_containers_integration_tests/tasks/test_enable_segments_to_index_task.py rename to api/tests/e2e_tests/containers/tasks/test_enable_segments_to_index_task.py diff --git a/api/tests/test_containers_integration_tests/tasks/test_mail_account_deletion_task.py b/api/tests/e2e_tests/containers/tasks/test_mail_account_deletion_task.py similarity index 100% rename from api/tests/test_containers_integration_tests/tasks/test_mail_account_deletion_task.py rename to api/tests/e2e_tests/containers/tasks/test_mail_account_deletion_task.py diff --git a/api/tests/test_containers_integration_tests/tasks/test_mail_change_mail_task.py b/api/tests/e2e_tests/containers/tasks/test_mail_change_mail_task.py similarity index 100% rename from api/tests/test_containers_integration_tests/tasks/test_mail_change_mail_task.py rename to api/tests/e2e_tests/containers/tasks/test_mail_change_mail_task.py diff --git a/api/tests/test_containers_integration_tests/tasks/test_mail_email_code_login_task.py b/api/tests/e2e_tests/containers/tasks/test_mail_email_code_login_task.py similarity index 100% rename from api/tests/test_containers_integration_tests/tasks/test_mail_email_code_login_task.py rename to api/tests/e2e_tests/containers/tasks/test_mail_email_code_login_task.py diff --git a/api/tests/test_containers_integration_tests/tasks/test_mail_human_input_delivery_task.py b/api/tests/e2e_tests/containers/tasks/test_mail_human_input_delivery_task.py similarity index 100% rename from api/tests/test_containers_integration_tests/tasks/test_mail_human_input_delivery_task.py rename to api/tests/e2e_tests/containers/tasks/test_mail_human_input_delivery_task.py diff --git a/api/tests/test_containers_integration_tests/tasks/test_mail_inner_task.py b/api/tests/e2e_tests/containers/tasks/test_mail_inner_task.py similarity index 100% rename from api/tests/test_containers_integration_tests/tasks/test_mail_inner_task.py rename to api/tests/e2e_tests/containers/tasks/test_mail_inner_task.py diff --git a/api/tests/test_containers_integration_tests/tasks/test_mail_invite_member_task.py b/api/tests/e2e_tests/containers/tasks/test_mail_invite_member_task.py similarity index 100% rename from api/tests/test_containers_integration_tests/tasks/test_mail_invite_member_task.py rename to api/tests/e2e_tests/containers/tasks/test_mail_invite_member_task.py diff --git a/api/tests/test_containers_integration_tests/tasks/test_mail_owner_transfer_task.py b/api/tests/e2e_tests/containers/tasks/test_mail_owner_transfer_task.py similarity index 100% rename from api/tests/test_containers_integration_tests/tasks/test_mail_owner_transfer_task.py rename to api/tests/e2e_tests/containers/tasks/test_mail_owner_transfer_task.py diff --git a/api/tests/test_containers_integration_tests/tasks/test_mail_register_task.py b/api/tests/e2e_tests/containers/tasks/test_mail_register_task.py similarity index 100% rename from api/tests/test_containers_integration_tests/tasks/test_mail_register_task.py rename to api/tests/e2e_tests/containers/tasks/test_mail_register_task.py diff --git a/api/tests/test_containers_integration_tests/tasks/test_rag_pipeline_run_tasks.py b/api/tests/e2e_tests/containers/tasks/test_rag_pipeline_run_tasks.py similarity index 100% rename from api/tests/test_containers_integration_tests/tasks/test_rag_pipeline_run_tasks.py rename to api/tests/e2e_tests/containers/tasks/test_rag_pipeline_run_tasks.py diff --git a/api/tests/test_containers_integration_tests/tasks/test_remove_app_and_related_data_task.py b/api/tests/e2e_tests/containers/tasks/test_remove_app_and_related_data_task.py similarity index 100% rename from api/tests/test_containers_integration_tests/tasks/test_remove_app_and_related_data_task.py rename to api/tests/e2e_tests/containers/tasks/test_remove_app_and_related_data_task.py diff --git a/api/tests/test_containers_integration_tests/test_container_state_isolation.py b/api/tests/e2e_tests/containers/test_container_state_isolation.py similarity index 100% rename from api/tests/test_containers_integration_tests/test_container_state_isolation.py rename to api/tests/e2e_tests/containers/test_container_state_isolation.py diff --git a/api/tests/test_containers_integration_tests/test_workflow_pause_integration.py b/api/tests/e2e_tests/containers/test_workflow_pause_integration.py similarity index 100% rename from api/tests/test_containers_integration_tests/test_workflow_pause_integration.py rename to api/tests/e2e_tests/containers/test_workflow_pause_integration.py diff --git a/api/tests/test_containers_integration_tests/trigger/__init__.py b/api/tests/e2e_tests/containers/trigger/__init__.py similarity index 100% rename from api/tests/test_containers_integration_tests/trigger/__init__.py rename to api/tests/e2e_tests/containers/trigger/__init__.py diff --git a/api/tests/test_containers_integration_tests/trigger/conftest.py b/api/tests/e2e_tests/containers/trigger/conftest.py similarity index 100% rename from api/tests/test_containers_integration_tests/trigger/conftest.py rename to api/tests/e2e_tests/containers/trigger/conftest.py diff --git a/api/tests/test_containers_integration_tests/trigger/test_trigger_e2e.py b/api/tests/e2e_tests/containers/trigger/test_trigger_e2e.py similarity index 100% rename from api/tests/test_containers_integration_tests/trigger/test_trigger_e2e.py rename to api/tests/e2e_tests/containers/trigger/test_trigger_e2e.py diff --git a/api/tests/integration_tests/workflow/nodes/knowledge_retrieval/__init__.py b/api/tests/e2e_tests/containers/workflow/__init__.py similarity index 100% rename from api/tests/integration_tests/workflow/nodes/knowledge_retrieval/__init__.py rename to api/tests/e2e_tests/containers/workflow/__init__.py diff --git a/api/tests/test_containers_integration_tests/__init__.py b/api/tests/e2e_tests/containers/workflow/nodes/__init__.py similarity index 100% rename from api/tests/test_containers_integration_tests/__init__.py rename to api/tests/e2e_tests/containers/workflow/nodes/__init__.py diff --git a/api/tests/test_containers_integration_tests/controllers/openapi/__init__.py b/api/tests/e2e_tests/containers/workflow/nodes/code_executor/__init__.py similarity index 100% rename from api/tests/test_containers_integration_tests/controllers/openapi/__init__.py rename to api/tests/e2e_tests/containers/workflow/nodes/code_executor/__init__.py diff --git a/api/tests/test_containers_integration_tests/workflow/nodes/code_executor/test_code_javascript.py b/api/tests/e2e_tests/containers/workflow/nodes/code_executor/test_code_javascript.py similarity index 100% rename from api/tests/test_containers_integration_tests/workflow/nodes/code_executor/test_code_javascript.py rename to api/tests/e2e_tests/containers/workflow/nodes/code_executor/test_code_javascript.py diff --git a/api/tests/test_containers_integration_tests/workflow/nodes/code_executor/test_code_jinja2.py b/api/tests/e2e_tests/containers/workflow/nodes/code_executor/test_code_jinja2.py similarity index 100% rename from api/tests/test_containers_integration_tests/workflow/nodes/code_executor/test_code_jinja2.py rename to api/tests/e2e_tests/containers/workflow/nodes/code_executor/test_code_jinja2.py diff --git a/api/tests/test_containers_integration_tests/workflow/nodes/code_executor/test_code_python3.py b/api/tests/e2e_tests/containers/workflow/nodes/code_executor/test_code_python3.py similarity index 100% rename from api/tests/test_containers_integration_tests/workflow/nodes/code_executor/test_code_python3.py rename to api/tests/e2e_tests/containers/workflow/nodes/code_executor/test_code_python3.py diff --git a/api/tests/test_containers_integration_tests/workflow/nodes/code_executor/test_utils.py b/api/tests/e2e_tests/containers/workflow/nodes/code_executor/test_utils.py similarity index 100% rename from api/tests/test_containers_integration_tests/workflow/nodes/code_executor/test_utils.py rename to api/tests/e2e_tests/containers/workflow/nodes/code_executor/test_utils.py diff --git a/api/tests/unit_tests/commands/test_check_no_new_base_model.py b/api/tests/e2e_tests/contracts/commands/test_check_no_new_base_model.py similarity index 98% rename from api/tests/unit_tests/commands/test_check_no_new_base_model.py rename to api/tests/e2e_tests/contracts/commands/test_check_no_new_base_model.py index f4537b1d17a..bdbe395baea 100644 --- a/api/tests/unit_tests/commands/test_check_no_new_base_model.py +++ b/api/tests/e2e_tests/contracts/commands/test_check_no_new_base_model.py @@ -8,7 +8,7 @@ from pathlib import Path import pytest -SCRIPT = Path(__file__).resolve().parents[4] / "scripts" / "check_no_new_base_model.py" +SCRIPT = Path(__file__).resolve().parents[5] / "scripts" / "check_no_new_base_model.py" SOURCE_PATH = "api/models/example.py" diff --git a/api/tests/unit_tests/commands/test_check_no_new_getattr.py b/api/tests/e2e_tests/contracts/commands/test_check_no_new_getattr.py similarity index 93% rename from api/tests/unit_tests/commands/test_check_no_new_getattr.py rename to api/tests/e2e_tests/contracts/commands/test_check_no_new_getattr.py index a846955a770..4086c5a39ff 100644 --- a/api/tests/unit_tests/commands/test_check_no_new_getattr.py +++ b/api/tests/e2e_tests/contracts/commands/test_check_no_new_getattr.py @@ -12,7 +12,7 @@ from pathlib import Path import pytest -REPO_ROOT = Path(__file__).resolve().parents[4] +REPO_ROOT = Path(__file__).resolve().parents[5] SCRIPTS_DIR = REPO_ROOT / "scripts" SCRIPT_PATH = SCRIPTS_DIR / "check_no_new_getattr.py" GUARD_HELPER_PATH = SCRIPTS_DIR / "ast_grep_guard.py" @@ -805,6 +805,52 @@ def test_modified_hunk_with_increased_getattr_count_fails(tmp_path: Path) -> Non assert "net-new getattr" in result.stderr +@pytest.mark.parametrize("mode", ["staged", "base"]) +@pytest.mark.parametrize("change", ["move", "copy", "duplicate", "different"]) +@pytest.mark.parametrize( + ("script", "statement"), + [ + ("check_no_new_getattr.py", 'return getattr(obj, "legacy", None)'), + ("check_no_new_session_mock.py", 'session = MagicMock(name="legacy")'), + ], +) +def test_consolidation_preserves_only_removed_matching_calls( + tmp_path: Path, mode: str, change: str, script: str, statement: str +) -> None: + init_repo(tmp_path) + source_path = "api/tests/unit_tests/test_legacy.py" + destination = "api/tests/unit_tests/test_consolidated.py" + source = f"def legacy(obj):\n {statement}\n" + write_repo_file(tmp_path, source_path, source) + commit_all(tmp_path, "baseline") + baseline = main_branch_rev(tmp_path) + + if change != "copy": + (tmp_path / source_path).unlink() + moved = source.replace('"legacy"', '"new"') if change == "different" else source + consolidated = "class TestConsolidated:\n" + textwrap.indent(moved, " ") + if change == "duplicate": + consolidated += textwrap.indent(source.replace("def legacy", "def duplicate"), " ") + write_repo_file(tmp_path, destination, consolidated) + git(tmp_path, "add", ".") + arguments = ["--staged"] + if mode == "base": + commit_all(tmp_path, "consolidate") + arguments = ["--base-rev", baseline] + + result = subprocess.run( + [sys.executable, str(SCRIPTS_DIR / script), *arguments], + cwd=tmp_path, + text=True, + capture_output=True, + check=False, + ) + + assert result.returncode == int(change != "move"), result.stderr + if change != "move": + assert destination in result.stderr + + def test_inline_guard_ignore_with_explanatory_text_skips_added_getattr(tmp_path: Path) -> None: init_repo(tmp_path) write_repo_file( diff --git a/api/tests/unit_tests/commands/test_check_no_new_test_doubles.py b/api/tests/e2e_tests/contracts/commands/test_check_no_new_test_doubles.py similarity index 83% rename from api/tests/unit_tests/commands/test_check_no_new_test_doubles.py rename to api/tests/e2e_tests/contracts/commands/test_check_no_new_test_doubles.py index 526b3b1b38c..460879270af 100644 --- a/api/tests/unit_tests/commands/test_check_no_new_test_doubles.py +++ b/api/tests/e2e_tests/contracts/commands/test_check_no_new_test_doubles.py @@ -10,7 +10,7 @@ from pathlib import Path import pytest -REPO_ROOT = Path(__file__).resolve().parents[4] +REPO_ROOT = Path(__file__).resolve().parents[5] SCRIPT = REPO_ROOT / "scripts" / "check_no_new_test_doubles.py" TEST_PATH = "api/tests/unit_tests/test_example.py" @@ -209,6 +209,37 @@ def test_replacing_one_mock_with_another_is_not_grandfathered(repo: Path) -> Non assert result.returncode == 1, result.stderr +def test_patched_method_invocations_are_not_mock_constructors(repo: Path) -> None: + write_source(repo, "def setup(client):\n client.run = Mock()\n") + baseline = commit(repo) + write_source(repo, "def setup(client):\n client.run = Mock()\n\ndef invoke(client):\n client.run()\n") + commit(repo) + + result = run_guard(repo, "--base-rev", baseline) + + assert result.returncode == 0, result.stderr + + +@pytest.mark.parametrize("change_value", [False, True]) +def test_consolidation_preserves_import_aliases_but_rejects_changed_arguments(repo: Path, change_value: bool) -> None: + source = "import datetime\nfrom unittest.mock import Mock\nvalue = Mock(return_value=datetime.timedelta(days=1))\n" + write_source(repo, source) + baseline = commit(repo) + (repo / TEST_PATH).unlink() + moved = source.replace("import datetime", "import datetime as dates").replace( + "datetime.timedelta", "dates.timedelta" + ) + moved = moved.replace("import Mock", "import Mock as Factory").replace("= Mock(", "= Factory(") + if change_value: + moved = moved.replace("days=1", "days=2") + write_source(repo, moved, "api/tests/unit_tests/test_moved.py") + commit(repo) + + result = run_guard(repo, "--base-rev", baseline) + + assert result.returncode == int(change_value), result.stderr + + @pytest.mark.parametrize(("reason", "expected"), [(" -- abstract adapter contract", 0), ("", 1)]) def test_subclass_exception_requires_a_reason(repo: Path, reason: str, expected: int) -> None: source = ( @@ -271,6 +302,37 @@ def test_invalid_revision_fails_closed(repo: Path) -> None: assert result.stderr +@pytest.mark.parametrize("mode", ["staged", "base"]) +@pytest.mark.parametrize("change", ["move", "copy", "duplicate", "different"]) +@pytest.mark.parametrize("kind", ["mock", "subclass"]) +def test_consolidation_keeps_only_removed_matching_baselines(repo: Path, mode: str, change: str, kind: str) -> None: + source = ( + 'def legacy():\n return Mock(\n name="legacy",\n )\n' + if kind == "mock" + else "class Legacy(ModelInstance):\n def __init__(self): self.calls = []\n" + ) + write_source(repo, source) + baseline = commit(repo) + if change != "copy": + (repo / TEST_PATH).unlink() + moved = source.replace("legacy", "new").replace("Legacy", "New") if change == "different" else source + consolidated = "class TestConsolidated:\n" + textwrap.indent(moved, " ") + if change == "duplicate": + consolidated += textwrap.indent(source, " ") + destination = "api/tests/unit_tests/test_consolidated.py" + write_source(repo, consolidated, destination) + git(repo, "add", ".") + if mode == "base": + commit(repo) + result = run_guard(repo, "--base-rev", baseline) + else: + result = run_guard(repo, "--staged") + + assert result.returncode == int(change != "move"), result.stderr + if change != "move": + assert destination in result.stderr + + def test_requires_an_explicit_diff_source(repo: Path) -> None: assert run_guard(repo).returncode == 2 assert run_guard(repo, "--staged", "--base-rev", "HEAD").returncode == 2 diff --git a/api/tests/unit_tests/commands/test_check_no_spec_mock.py b/api/tests/e2e_tests/contracts/commands/test_check_no_spec_mock.py similarity index 98% rename from api/tests/unit_tests/commands/test_check_no_spec_mock.py rename to api/tests/e2e_tests/contracts/commands/test_check_no_spec_mock.py index b55851a6250..0ec2e07fe7d 100644 --- a/api/tests/unit_tests/commands/test_check_no_spec_mock.py +++ b/api/tests/e2e_tests/contracts/commands/test_check_no_spec_mock.py @@ -8,7 +8,7 @@ from pathlib import Path import pytest -REPO_ROOT = Path(__file__).resolve().parents[4] +REPO_ROOT = Path(__file__).resolve().parents[5] SCRIPT = REPO_ROOT / "scripts" / "check_no_spec_mock.py" diff --git a/api/tests/unit_tests/configs/test_env_consistency.py b/api/tests/e2e_tests/contracts/configs/test_env_consistency.py similarity index 98% rename from api/tests/unit_tests/configs/test_env_consistency.py rename to api/tests/e2e_tests/contracts/configs/test_env_consistency.py index 372ef970d59..ccbf66dd4fb 100644 --- a/api/tests/unit_tests/configs/test_env_consistency.py +++ b/api/tests/e2e_tests/contracts/configs/test_env_consistency.py @@ -92,7 +92,7 @@ BASE_API_AND_DOCKER_COMPOSE_CONFIG_SET_DIFF: frozenset[str] = frozenset( ) ) -REPO_ROOT = Path(__file__).resolve().parents[4] +REPO_ROOT = Path(__file__).resolve().parents[5] def _api_config_set() -> set[str]: diff --git a/api/tests/unit_tests/configs/test_knowledge_fs_worker_deployment.py b/api/tests/e2e_tests/contracts/configs/test_knowledge_fs_worker_deployment.py similarity index 98% rename from api/tests/unit_tests/configs/test_knowledge_fs_worker_deployment.py rename to api/tests/e2e_tests/contracts/configs/test_knowledge_fs_worker_deployment.py index 1a96f5afe90..210805636b7 100644 --- a/api/tests/unit_tests/configs/test_knowledge_fs_worker_deployment.py +++ b/api/tests/e2e_tests/contracts/configs/test_knowledge_fs_worker_deployment.py @@ -6,7 +6,7 @@ from pathlib import Path import pytest import yaml -_ROOT = Path(__file__).resolve().parents[4] +_ROOT = Path(__file__).resolve().parents[5] def _worker_configuration(environment: dict[str, str]) -> subprocess.CompletedProcess[str]: diff --git a/api/tests/e2e_tests/contracts/conftest.py b/api/tests/e2e_tests/contracts/conftest.py new file mode 100644 index 00000000000..9b80dfc983a --- /dev/null +++ b/api/tests/e2e_tests/contracts/conftest.py @@ -0,0 +1,25 @@ +from tests.unit_tests.conftest import ( # noqa: F401 + _patch_redis_clients, + _provide_app_context, + _sqlite_database_template, + _sqlite_engine, + _sqlite_session_factory, + _unbound_session_factory, + account_application_services, + account_domain, + app, + app_services, + application_tags, + config_overrides, + pytest_plugins, + rbac_domain, + redis_transport, + reset_redis_mock, + reset_secret_key, + sqlite_engine, + sqlite_session, + sqlite_session_factory, + tenant_queue_commands, + unbound_session, + unbound_session_factory, +) diff --git a/api/tests/unit_tests/controllers/service_api/test_conftest.py b/api/tests/e2e_tests/contracts/controllers/service_api/test_conftest.py similarity index 51% rename from api/tests/unit_tests/controllers/service_api/test_conftest.py rename to api/tests/e2e_tests/contracts/controllers/service_api/test_conftest.py index 014d99a0636..f26ba315481 100644 --- a/api/tests/unit_tests/controllers/service_api/test_conftest.py +++ b/api/tests/e2e_tests/contracts/controllers/service_api/test_conftest.py @@ -9,7 +9,27 @@ from tests.unit_tests.conftest import ( persist_service_api_dataset_owner, persist_service_api_tenant_owner, ) +from tests.unit_tests.controllers.conftest import app_query_services as app_query_services # noqa: PLC0414 +from tests.unit_tests.controllers.conftest import ( + datasource_application_dependencies as datasource_application_dependencies, # noqa: PLC0414 +) +from tests.unit_tests.controllers.conftest import error_response_contract as error_response_contract # noqa: PLC0414 +from tests.unit_tests.controllers.conftest import pipeline_application as pipeline_application # noqa: PLC0414 from tests.unit_tests.controllers.service_api.conftest import ServiceApiIdentity +from tests.unit_tests.controllers.service_api.conftest import app as app # noqa: PLC0414 +from tests.unit_tests.controllers.service_api.conftest import mock_account as mock_account # noqa: PLC0414 +from tests.unit_tests.controllers.service_api.conftest import mock_api_token as mock_api_token # noqa: PLC0414 +from tests.unit_tests.controllers.service_api.conftest import mock_app_id as mock_app_id # noqa: PLC0414 +from tests.unit_tests.controllers.service_api.conftest import mock_app_model as mock_app_model # noqa: PLC0414 +from tests.unit_tests.controllers.service_api.conftest import mock_child_chunk as mock_child_chunk # noqa: PLC0414 +from tests.unit_tests.controllers.service_api.conftest import mock_dataset as mock_dataset # noqa: PLC0414 +from tests.unit_tests.controllers.service_api.conftest import mock_dataset_api_token as mock_dataset_api_token # noqa: PLC0414 +from tests.unit_tests.controllers.service_api.conftest import mock_document as mock_document # noqa: PLC0414 +from tests.unit_tests.controllers.service_api.conftest import mock_end_user as mock_end_user # noqa: PLC0414 +from tests.unit_tests.controllers.service_api.conftest import mock_segment as mock_segment # noqa: PLC0414 +from tests.unit_tests.controllers.service_api.conftest import mock_tenant as mock_tenant # noqa: PLC0414 +from tests.unit_tests.controllers.service_api.conftest import mock_tenant_id as mock_tenant_id # noqa: PLC0414 +from tests.unit_tests.controllers.service_api.conftest import service_api_identity as service_api_identity # noqa: PLC0414 def test_service_api_identity_persists_tenant_scoped_owner(service_api_identity: ServiceApiIdentity) -> None: diff --git a/api/tests/unit_tests/core/datasource/test_file_upload.py b/api/tests/e2e_tests/contracts/core/datasource/test_file_upload.py similarity index 99% rename from api/tests/unit_tests/core/datasource/test_file_upload.py rename to api/tests/e2e_tests/contracts/core/datasource/test_file_upload.py index 1dd5ddd5680..b5395d418b0 100644 --- a/api/tests/unit_tests/core/datasource/test_file_upload.py +++ b/api/tests/e2e_tests/contracts/core/datasource/test_file_upload.py @@ -111,9 +111,10 @@ IMPORTANT NOTES: RUNNING TESTS: ============== -Run all tests: pytest api/tests/unit_tests/core/datasource/test_file_upload.py -v -Run specific test class: pytest api/tests/unit_tests/core/datasource/test_file_upload.py::TestFileTypeValidation -v -Run with coverage: pytest api/tests/unit_tests/core/datasource/test_file_upload.py --cov=services.file_service +Run all tests: pytest api/tests/e2e_tests/contracts/core/datasource/test_file_upload.py -v +Run specific test class: +pytest api/tests/e2e_tests/contracts/core/datasource/test_file_upload.py::TestFileTypeValidation -v +Run with coverage: pytest api/tests/e2e_tests/contracts/core/datasource/test_file_upload.py --cov=services.file_service """ # Standard library imports diff --git a/api/tests/unit_tests/core/file/test_models.py b/api/tests/e2e_tests/contracts/core/file/test_models.py similarity index 100% rename from api/tests/unit_tests/core/file/test_models.py rename to api/tests/e2e_tests/contracts/core/file/test_models.py diff --git a/api/tests/unit_tests/core/test_file.py b/api/tests/e2e_tests/contracts/core/test_file.py similarity index 100% rename from api/tests/unit_tests/core/test_file.py rename to api/tests/e2e_tests/contracts/core/test_file.py diff --git a/api/tests/unit_tests/core/variables/test_segment.py b/api/tests/e2e_tests/contracts/core/variables/test_segment.py similarity index 100% rename from api/tests/unit_tests/core/variables/test_segment.py rename to api/tests/e2e_tests/contracts/core/variables/test_segment.py diff --git a/api/tests/unit_tests/core/variables/test_segment_type.py b/api/tests/e2e_tests/contracts/core/variables/test_segment_type.py similarity index 100% rename from api/tests/unit_tests/core/variables/test_segment_type.py rename to api/tests/e2e_tests/contracts/core/variables/test_segment_type.py diff --git a/api/tests/unit_tests/core/variables/test_segment_type_validation.py b/api/tests/e2e_tests/contracts/core/variables/test_segment_type_validation.py similarity index 100% rename from api/tests/unit_tests/core/variables/test_segment_type_validation.py rename to api/tests/e2e_tests/contracts/core/variables/test_segment_type_validation.py diff --git a/api/tests/unit_tests/core/variables/test_variables.py b/api/tests/e2e_tests/contracts/core/variables/test_variables.py similarity index 100% rename from api/tests/unit_tests/core/variables/test_variables.py rename to api/tests/e2e_tests/contracts/core/variables/test_variables.py diff --git a/api/tests/unit_tests/core/workflow/entities/test_private_workflow_pause.py b/api/tests/e2e_tests/contracts/core/workflow/entities/test_private_workflow_pause.py similarity index 100% rename from api/tests/unit_tests/core/workflow/entities/test_private_workflow_pause.py rename to api/tests/e2e_tests/contracts/core/workflow/entities/test_private_workflow_pause.py diff --git a/api/tests/unit_tests/core/workflow/graph_engine/README.md b/api/tests/e2e_tests/contracts/core/workflow/graph_engine/README.md similarity index 67% rename from api/tests/unit_tests/core/workflow/graph_engine/README.md rename to api/tests/e2e_tests/contracts/core/workflow/graph_engine/README.md index dd419f0810f..a0b901914d3 100644 --- a/api/tests/unit_tests/core/workflow/graph_engine/README.md +++ b/api/tests/e2e_tests/contracts/core/workflow/graph_engine/README.md @@ -24,7 +24,7 @@ Examples: ```bash uv run --project api --dev pytest api/tests/unit_tests/core/workflow/graph_engine/layers/test_llm_quota.py -uv run --project api --dev pytest api/tests/unit_tests/core/workflow/graph_engine/layers/test_observability.py -uv run --project api --dev pytest api/tests/unit_tests/core/workflow/graph_engine/test_parallel_human_input_join_resume.py -uv run --project api --dev pytest api/tests/unit_tests/core/workflow/graph_engine/test_tool_in_chatflow.py +uv run --project api --dev pytest api/tests/unit_tests/core/app/workflow/layers/test_observability.py +uv run --project api --dev pytest api/tests/e2e_tests/contracts/core/workflow/graph_engine/test_parallel_human_input_join_resume.py +uv run --project api --dev pytest api/tests/e2e_tests/contracts/core/workflow/graph_engine/test_tool_in_chatflow.py ``` diff --git a/api/tests/unit_tests/core/workflow/graph_engine/test_parallel_human_input_join_resume.py b/api/tests/e2e_tests/contracts/core/workflow/graph_engine/test_parallel_human_input_join_resume.py similarity index 100% rename from api/tests/unit_tests/core/workflow/graph_engine/test_parallel_human_input_join_resume.py rename to api/tests/e2e_tests/contracts/core/workflow/graph_engine/test_parallel_human_input_join_resume.py diff --git a/api/tests/unit_tests/core/workflow/graph_engine/test_tool_in_chatflow.py b/api/tests/e2e_tests/contracts/core/workflow/graph_engine/test_tool_in_chatflow.py similarity index 95% rename from api/tests/unit_tests/core/workflow/graph_engine/test_tool_in_chatflow.py rename to api/tests/e2e_tests/contracts/core/workflow/graph_engine/test_tool_in_chatflow.py index 00910502df8..5cf9b1e518c 100644 --- a/api/tests/unit_tests/core/workflow/graph_engine/test_tool_in_chatflow.py +++ b/api/tests/e2e_tests/contracts/core/workflow/graph_engine/test_tool_in_chatflow.py @@ -5,9 +5,8 @@ from graphon.engine_events import ( GraphRunSucceededEvent, NodeRunStreamChunkEvent, ) - -from .test_mock_config import MockConfigBuilder -from .test_table_runner import TableTestRunner +from tests.helpers.workflow.mock_config import MockConfigBuilder +from tests.helpers.workflow.table_runner import TableTestRunner def test_tool_in_chatflow(): diff --git a/api/tests/unit_tests/core/workflow/nodes/answer/test_answer.py b/api/tests/e2e_tests/contracts/core/workflow/nodes/answer/test_answer.py similarity index 100% rename from api/tests/unit_tests/core/workflow/nodes/answer/test_answer.py rename to api/tests/e2e_tests/contracts/core/workflow/nodes/answer/test_answer.py diff --git a/api/tests/unit_tests/core/workflow/nodes/base/test_base_node.py b/api/tests/e2e_tests/contracts/core/workflow/nodes/base/test_base_node.py similarity index 100% rename from api/tests/unit_tests/core/workflow/nodes/base/test_base_node.py rename to api/tests/e2e_tests/contracts/core/workflow/nodes/base/test_base_node.py diff --git a/api/tests/unit_tests/core/workflow/nodes/code/test_code_node_output_validation.py b/api/tests/e2e_tests/contracts/core/workflow/nodes/code/test_code_node_output_validation.py similarity index 100% rename from api/tests/unit_tests/core/workflow/nodes/code/test_code_node_output_validation.py rename to api/tests/e2e_tests/contracts/core/workflow/nodes/code/test_code_node_output_validation.py diff --git a/api/tests/unit_tests/core/workflow/nodes/http_request/test_http_request_executor.py b/api/tests/e2e_tests/contracts/core/workflow/nodes/http_request/test_http_request_executor.py similarity index 100% rename from api/tests/unit_tests/core/workflow/nodes/http_request/test_http_request_executor.py rename to api/tests/e2e_tests/contracts/core/workflow/nodes/http_request/test_http_request_executor.py diff --git a/api/tests/unit_tests/core/workflow/nodes/http_request/test_http_request_node.py b/api/tests/e2e_tests/contracts/core/workflow/nodes/http_request/test_http_request_node.py similarity index 100% rename from api/tests/unit_tests/core/workflow/nodes/http_request/test_http_request_node.py rename to api/tests/e2e_tests/contracts/core/workflow/nodes/http_request/test_http_request_node.py diff --git a/api/tests/unit_tests/core/workflow/nodes/llm/test_llm_utils.py b/api/tests/e2e_tests/contracts/core/workflow/nodes/llm/test_llm_utils.py similarity index 99% rename from api/tests/unit_tests/core/workflow/nodes/llm/test_llm_utils.py rename to api/tests/e2e_tests/contracts/core/workflow/nodes/llm/test_llm_utils.py index 141053d24d5..aeec0709be5 100644 --- a/api/tests/unit_tests/core/workflow/nodes/llm/test_llm_utils.py +++ b/api/tests/e2e_tests/contracts/core/workflow/nodes/llm/test_llm_utils.py @@ -36,7 +36,7 @@ from graphon.nodes.llm.exc import ( from graphon.runtime import VariablePool from graphon.template_rendering import TemplateRenderError from graphon.variables import ArrayAnySegment, ArrayFileSegment, NoneSegment -from tests.unit_tests.core.model_fixtures import make_model_instance +from tests.helpers.unit.core.model_fixtures import make_model_instance def _build_model_schema( diff --git a/api/tests/unit_tests/core/workflow/nodes/llm/test_node.py b/api/tests/e2e_tests/contracts/core/workflow/nodes/llm/test_node.py similarity index 99% rename from api/tests/unit_tests/core/workflow/nodes/llm/test_node.py rename to api/tests/e2e_tests/contracts/core/workflow/nodes/llm/test_node.py index 8b8119c2d47..d15f83dc7fb 100644 --- a/api/tests/unit_tests/core/workflow/nodes/llm/test_node.py +++ b/api/tests/e2e_tests/contracts/core/workflow/nodes/llm/test_node.py @@ -269,7 +269,7 @@ def llm_node(llm_node_data: LLMNodeData, graph_init_params: InitParams, graph_ru @pytest.fixture def model_config(monkeypatch: pytest.MonkeyPatch): - from tests.integration_tests.model_runtime.__mock.plugin_model import MockModelClass + from tests.e2e_tests.middleware.model_runtime.__mock.plugin_model import MockModelClass def mock_model_providers(_self): providers = [] diff --git a/api/tests/unit_tests/core/workflow/nodes/parameter_extractor/test_parameter_extractor_node.py b/api/tests/e2e_tests/contracts/core/workflow/nodes/parameter_extractor/test_parameter_extractor_node.py similarity index 100% rename from api/tests/unit_tests/core/workflow/nodes/parameter_extractor/test_parameter_extractor_node.py rename to api/tests/e2e_tests/contracts/core/workflow/nodes/parameter_extractor/test_parameter_extractor_node.py diff --git a/api/tests/unit_tests/core/workflow/nodes/template_transform/test_template_transform_execution.py b/api/tests/e2e_tests/contracts/core/workflow/nodes/template_transform/test_template_transform_execution.py similarity index 100% rename from api/tests/unit_tests/core/workflow/nodes/template_transform/test_template_transform_execution.py rename to api/tests/e2e_tests/contracts/core/workflow/nodes/template_transform/test_template_transform_execution.py diff --git a/api/tests/unit_tests/core/workflow/nodes/template_transform/test_template_transform_node.py b/api/tests/e2e_tests/contracts/core/workflow/nodes/template_transform/test_template_transform_node.py similarity index 85% rename from api/tests/unit_tests/core/workflow/nodes/template_transform/test_template_transform_node.py rename to api/tests/e2e_tests/contracts/core/workflow/nodes/template_transform/test_template_transform_node.py index 8f53ec0f25a..feb45995bf5 100644 --- a/api/tests/unit_tests/core/workflow/nodes/template_transform/test_template_transform_node.py +++ b/api/tests/e2e_tests/contracts/core/workflow/nodes/template_transform/test_template_transform_node.py @@ -2,8 +2,9 @@ from graphon.nodes.base.entities import VariableSelector from graphon.nodes.template_transform.template_transform_node import ( TemplateTransformNode, ) - -from .template_transform_node_spec import TestTemplateTransformNode # noqa: F401 +from tests.helpers.unit.core.workflow.nodes.template_transform.template_transform_node_spec import ( + TestTemplateTransformNode, # noqa: F401 +) def test_extract_variable_selector_to_variable_mapping_accepts_mixed_valid_entries(): diff --git a/api/tests/unit_tests/core/workflow/nodes/test_base_node.py b/api/tests/e2e_tests/contracts/core/workflow/nodes/test_base_node.py similarity index 100% rename from api/tests/unit_tests/core/workflow/nodes/test_base_node.py rename to api/tests/e2e_tests/contracts/core/workflow/nodes/test_base_node.py diff --git a/api/tests/unit_tests/core/workflow/nodes/test_document_extractor_node.py b/api/tests/e2e_tests/contracts/core/workflow/nodes/test_document_extractor_node.py similarity index 100% rename from api/tests/unit_tests/core/workflow/nodes/test_document_extractor_node.py rename to api/tests/e2e_tests/contracts/core/workflow/nodes/test_document_extractor_node.py diff --git a/api/tests/unit_tests/core/workflow/nodes/test_if_else.py b/api/tests/e2e_tests/contracts/core/workflow/nodes/test_if_else.py similarity index 100% rename from api/tests/unit_tests/core/workflow/nodes/test_if_else.py rename to api/tests/e2e_tests/contracts/core/workflow/nodes/test_if_else.py diff --git a/api/tests/unit_tests/core/workflow/nodes/test_list_operator.py b/api/tests/e2e_tests/contracts/core/workflow/nodes/test_list_operator.py similarity index 100% rename from api/tests/unit_tests/core/workflow/nodes/test_list_operator.py rename to api/tests/e2e_tests/contracts/core/workflow/nodes/test_list_operator.py diff --git a/api/tests/unit_tests/core/workflow/nodes/test_start_node_json_object.py b/api/tests/e2e_tests/contracts/core/workflow/nodes/test_start_node_json_object.py similarity index 100% rename from api/tests/unit_tests/core/workflow/nodes/test_start_node_json_object.py rename to api/tests/e2e_tests/contracts/core/workflow/nodes/test_start_node_json_object.py diff --git a/api/tests/unit_tests/core/workflow/nodes/tool/test_tool_node.py b/api/tests/e2e_tests/contracts/core/workflow/nodes/tool/test_tool_node.py similarity index 100% rename from api/tests/unit_tests/core/workflow/nodes/tool/test_tool_node.py rename to api/tests/e2e_tests/contracts/core/workflow/nodes/tool/test_tool_node.py diff --git a/api/tests/unit_tests/core/workflow/test_node_mapping_bootstrap.py b/api/tests/e2e_tests/contracts/core/workflow/test_node_mapping_bootstrap.py similarity index 98% rename from api/tests/unit_tests/core/workflow/test_node_mapping_bootstrap.py rename to api/tests/e2e_tests/contracts/core/workflow/test_node_mapping_bootstrap.py index c512eb04637..0c798e29fb2 100644 --- a/api/tests/unit_tests/core/workflow/test_node_mapping_bootstrap.py +++ b/api/tests/e2e_tests/contracts/core/workflow/test_node_mapping_bootstrap.py @@ -6,7 +6,7 @@ from pathlib import Path def test_moved_core_nodes_resolve_after_importing_production_entrypoints(): - api_root = Path(__file__).resolve().parents[4] + api_root = Path(__file__).resolve().parents[5] # `PYTHONSAFEPATH=1` enables Python's safe-path mode, which suppresses the # usual implicit insertion of the working directory into `sys.path`. diff --git a/api/tests/unit_tests/core/workflow/test_variable_pool.py b/api/tests/e2e_tests/contracts/core/workflow/test_variable_pool.py similarity index 100% rename from api/tests/unit_tests/core/workflow/test_variable_pool.py rename to api/tests/e2e_tests/contracts/core/workflow/test_variable_pool.py diff --git a/api/tests/unit_tests/dev/test_knowledge_fs_ci_workflow.py b/api/tests/e2e_tests/contracts/dev/test_knowledge_fs_ci_workflow.py similarity index 96% rename from api/tests/unit_tests/dev/test_knowledge_fs_ci_workflow.py rename to api/tests/e2e_tests/contracts/dev/test_knowledge_fs_ci_workflow.py index 3b3df821fc6..af95abe25c1 100644 --- a/api/tests/unit_tests/dev/test_knowledge_fs_ci_workflow.py +++ b/api/tests/e2e_tests/contracts/dev/test_knowledge_fs_ci_workflow.py @@ -14,7 +14,7 @@ import yaml from dev.check_knowledge_fs_coverage import discover_core_coverage_paths -ROOT = Path(__file__).resolve().parents[4] +ROOT = Path(__file__).resolve().parents[5] class _Yaml: @@ -90,9 +90,9 @@ def test_dify_gate_target_collection_includes_all_owned_integration_files(tmp_pa "api/models/model.py", } <= production assert { - "api/tests/unit_tests/dev/test_knowledge_fs_ci_workflow.py", - "api/tests/unit_tests/services/test_knowledge_fs_service_query_images.py", - "api/tests/unit_tests/migrations/test_knowledge_fs_migration_chain.py", + "api/tests/e2e_tests/contracts/dev/test_knowledge_fs_ci_workflow.py", + "api/tests/unit_tests/services/knowledge_fs/test_service_query_image_upload.py", + "api/tests/e2e_tests/contracts/migrations/test_knowledge_fs_migration_chain.py", } <= tests assert all((ROOT / path).is_file() for path in production | tests) for name in ( diff --git a/api/tests/unit_tests/libs/_human_input/test_form_service.py b/api/tests/e2e_tests/contracts/libs/_human_input/test_form_service.py similarity index 99% rename from api/tests/unit_tests/libs/_human_input/test_form_service.py rename to api/tests/e2e_tests/contracts/libs/_human_input/test_form_service.py index 6a2265b3ab9..e7ad85b8e3b 100644 --- a/api/tests/unit_tests/libs/_human_input/test_form_service.py +++ b/api/tests/e2e_tests/contracts/libs/_human_input/test_form_service.py @@ -14,8 +14,7 @@ from core.workflow.nodes.human_input.enums import ( TimeoutUnit, ) from libs.datetime_utils import naive_utc_now - -from .support import ( +from tests.helpers.unit.libs._human_input.support import ( FormAlreadySubmittedError, FormExpiredError, FormNotFoundError, diff --git a/api/tests/unit_tests/libs/_human_input/test_models.py b/api/tests/e2e_tests/contracts/libs/_human_input/test_models.py similarity index 98% rename from api/tests/unit_tests/libs/_human_input/test_models.py rename to api/tests/e2e_tests/contracts/libs/_human_input/test_models.py index 727e94ca00b..5e9862ef864 100644 --- a/api/tests/unit_tests/libs/_human_input/test_models.py +++ b/api/tests/e2e_tests/contracts/libs/_human_input/test_models.py @@ -14,8 +14,7 @@ from core.workflow.nodes.human_input.enums import ( TimeoutUnit, ) from libs.datetime_utils import naive_utc_now - -from .support import FormSubmissionData, FormSubmissionRequest, HumanInputForm +from tests.helpers.unit.libs._human_input.support import FormSubmissionData, FormSubmissionRequest, HumanInputForm class TestHumanInputForm: diff --git a/api/tests/unit_tests/libs/test_jwt_imports.py b/api/tests/e2e_tests/contracts/libs/test_jwt_imports.py similarity index 100% rename from api/tests/unit_tests/libs/test_jwt_imports.py rename to api/tests/e2e_tests/contracts/libs/test_jwt_imports.py diff --git a/api/tests/unit_tests/libs/test_pandas.py b/api/tests/e2e_tests/contracts/libs/test_pandas.py similarity index 100% rename from api/tests/unit_tests/libs/test_pandas.py rename to api/tests/e2e_tests/contracts/libs/test_pandas.py diff --git a/api/tests/unit_tests/libs/test_yarl.py b/api/tests/e2e_tests/contracts/libs/test_yarl.py similarity index 100% rename from api/tests/unit_tests/libs/test_yarl.py rename to api/tests/e2e_tests/contracts/libs/test_yarl.py diff --git a/api/tests/unit_tests/migrations/test_knowledge_fs_migration_chain.py b/api/tests/e2e_tests/contracts/migrations/test_knowledge_fs_migration_chain.py similarity index 96% rename from api/tests/unit_tests/migrations/test_knowledge_fs_migration_chain.py rename to api/tests/e2e_tests/contracts/migrations/test_knowledge_fs_migration_chain.py index adf3c715a47..4d6a6c3bf22 100644 --- a/api/tests/unit_tests/migrations/test_knowledge_fs_migration_chain.py +++ b/api/tests/e2e_tests/contracts/migrations/test_knowledge_fs_migration_chain.py @@ -4,7 +4,7 @@ from pathlib import Path from alembic.script import ScriptDirectory -_VERSIONS = Path(__file__).resolve().parents[3] / "migrations/versions" +_VERSIONS = Path(__file__).resolve().parents[4] / "migrations/versions" _MAIN_HEAD = "d6b9e2f1a407" # KnowledgeFS revisions run in this order after main's head, without merge revisions. _KNOWLEDGE_FS_CHAIN = [ diff --git a/api/tests/unit_tests/test_config_overrides.py b/api/tests/e2e_tests/contracts/test_config_overrides.py similarity index 100% rename from api/tests/unit_tests/test_config_overrides.py rename to api/tests/e2e_tests/contracts/test_config_overrides.py diff --git a/api/tests/unit_tests/test_integration_account_fixtures.py b/api/tests/e2e_tests/contracts/test_integration_account_fixtures.py similarity index 91% rename from api/tests/unit_tests/test_integration_account_fixtures.py rename to api/tests/e2e_tests/contracts/test_integration_account_fixtures.py index 6e09594b00e..84782ca825b 100644 --- a/api/tests/unit_tests/test_integration_account_fixtures.py +++ b/api/tests/e2e_tests/contracts/test_integration_account_fixtures.py @@ -6,8 +6,8 @@ from sqlalchemy.orm import Session from extensions.ext_application_services import ApplicationServices from models.account import Account, TenantAccountJoin, TenantAccountRole -from tests.test_containers_integration_tests.helpers import accounts as account_fixtures -from tests.unit_tests.account_domain import AccountDomain +from tests.e2e_tests.containers.helpers import accounts as account_fixtures +from tests.helpers.unit.account_domain import AccountDomain def test_owner_workspace_fixture_seeds_once_when_self_service_creation_is_disabled( diff --git a/api/tests/unit_tests/test_makefile_backend_tests.py b/api/tests/e2e_tests/contracts/test_makefile_backend_tests.py similarity index 58% rename from api/tests/unit_tests/test_makefile_backend_tests.py rename to api/tests/e2e_tests/contracts/test_makefile_backend_tests.py index 4f62e45948d..c4fc0b2e619 100644 --- a/api/tests/unit_tests/test_makefile_backend_tests.py +++ b/api/tests/e2e_tests/contracts/test_makefile_backend_tests.py @@ -3,7 +3,7 @@ from pathlib import Path def test_default_make_test_runs_backend_unit_suites(): - repo_root = Path(__file__).resolve().parents[3] + repo_root = Path(__file__).resolve().parents[4] completed = subprocess.run(["make", "-n", "test"], cwd=repo_root, check=True, capture_output=True, text=True) dry_run_output = completed.stdout @@ -14,14 +14,14 @@ def test_default_make_test_runs_backend_unit_suites(): assert "-p no:benchmark" in dry_run_output assert "api/tests/unit_tests/controllers" in dry_run_output assert "--start-middleware" not in dry_run_output - assert "api/tests/integration_tests/workflow" not in dry_run_output - assert "api/tests/test_containers_integration_tests" not in dry_run_output + assert "api/tests/e2e_tests/middleware/workflow" not in dry_run_output + assert "api/tests/e2e_tests/containers" not in dry_run_output assert "--start-vdb" not in dry_run_output - assert "api/providers/vdb/vdb-chroma/tests/integration_tests" not in dry_run_output + assert "api/providers/vdb/vdb-chroma/tests/e2e_tests" not in dry_run_output def test_make_test_all_runs_backend_pytest_suites(): - repo_root = Path(__file__).resolve().parents[3] + repo_root = Path(__file__).resolve().parents[4] completed = subprocess.run(["make", "-n", "test-all"], cwd=repo_root, check=True, capture_output=True, text=True) dry_run_output = completed.stdout @@ -31,11 +31,11 @@ def test_make_test_all_runs_backend_pytest_suites(): assert "api/providers/trace/*/tests/unit_tests" in dry_run_output assert "-p no:benchmark" in dry_run_output assert "--start-middleware" in dry_run_output - assert "api/tests/integration_tests/workflow" in dry_run_output - assert "api/tests/integration_tests/tools" in dry_run_output - assert "api/tests/test_containers_integration_tests" in dry_run_output + assert "api/tests/e2e_tests/middleware/workflow" in dry_run_output + assert "api/tests/e2e_tests/middleware/tools" in dry_run_output + assert "api/tests/e2e_tests/containers" in dry_run_output assert "--start-vdb" in dry_run_output - assert "api/providers/vdb/vdb-chroma/tests/integration_tests" in dry_run_output - assert "api/providers/vdb/vdb-pgvector/tests/integration_tests" in dry_run_output - assert "api/providers/vdb/vdb-qdrant/tests/integration_tests" in dry_run_output - assert "api/providers/vdb/vdb-weaviate/tests/integration_tests" in dry_run_output + assert "api/providers/vdb/vdb-chroma/tests/e2e_tests" in dry_run_output + assert "api/providers/vdb/vdb-pgvector/tests/e2e_tests" in dry_run_output + assert "api/providers/vdb/vdb-qdrant/tests/e2e_tests" in dry_run_output + assert "api/providers/vdb/vdb-weaviate/tests/e2e_tests" in dry_run_output diff --git a/api/tests/unit_tests/test_pytest_dify.py b/api/tests/e2e_tests/contracts/test_pytest_dify.py similarity index 98% rename from api/tests/unit_tests/test_pytest_dify.py rename to api/tests/e2e_tests/contracts/test_pytest_dify.py index c41229d0b5c..59cde61f3c2 100644 --- a/api/tests/unit_tests/test_pytest_dify.py +++ b/api/tests/e2e_tests/contracts/test_pytest_dify.py @@ -22,7 +22,7 @@ def test_ensure_backend_test_environment_uses_example_env_and_stable_logging( monkeypatch, ): repo_root = tmp_path - integration_tests_dir = repo_root / "api" / "tests" / "integration_tests" + integration_tests_dir = repo_root / "api" / "tests" / "e2e_tests" / "middleware" integration_tests_dir.mkdir(parents=True) env_example = integration_tests_dir / ".env.example" env_example.write_text("LOG_LEVEL=INFO\n") diff --git a/api/tests/unit_tests/test_pytest_timing.py b/api/tests/e2e_tests/contracts/test_pytest_timing.py similarity index 100% rename from api/tests/unit_tests/test_pytest_timing.py rename to api/tests/e2e_tests/contracts/test_pytest_timing.py diff --git a/api/tests/unit_tests/test_sqlite_fixtures.py b/api/tests/e2e_tests/contracts/test_sqlite_fixtures.py similarity index 100% rename from api/tests/unit_tests/test_sqlite_fixtures.py rename to api/tests/e2e_tests/contracts/test_sqlite_fixtures.py diff --git a/api/tests/knowledge_fs_vector_integration/test_backends.py b/api/tests/e2e_tests/knowledge_fs_vector/test_backends.py similarity index 100% rename from api/tests/knowledge_fs_vector_integration/test_backends.py rename to api/tests/e2e_tests/knowledge_fs_vector/test_backends.py diff --git a/api/tests/knowledge_fs_vector_integration/test_fulltext.py b/api/tests/e2e_tests/knowledge_fs_vector/test_fulltext.py similarity index 100% rename from api/tests/knowledge_fs_vector_integration/test_fulltext.py rename to api/tests/e2e_tests/knowledge_fs_vector/test_fulltext.py diff --git a/api/tests/integration_tests/.env.example b/api/tests/e2e_tests/middleware/.env.example similarity index 100% rename from api/tests/integration_tests/.env.example rename to api/tests/e2e_tests/middleware/.env.example diff --git a/api/tests/integration_tests/.gitignore b/api/tests/e2e_tests/middleware/.gitignore similarity index 100% rename from api/tests/integration_tests/.gitignore rename to api/tests/e2e_tests/middleware/.gitignore diff --git a/api/tests/integration_tests/__init__.py b/api/tests/e2e_tests/middleware/__init__.py similarity index 100% rename from api/tests/integration_tests/__init__.py rename to api/tests/e2e_tests/middleware/__init__.py diff --git a/api/tests/integration_tests/conftest.py b/api/tests/e2e_tests/middleware/conftest.py similarity index 100% rename from api/tests/integration_tests/conftest.py rename to api/tests/e2e_tests/middleware/conftest.py diff --git a/api/tests/test_containers_integration_tests/controllers/service_api/__init__.py b/api/tests/e2e_tests/middleware/controllers/__init__.py similarity index 100% rename from api/tests/test_containers_integration_tests/controllers/service_api/__init__.py rename to api/tests/e2e_tests/middleware/controllers/__init__.py diff --git a/api/tests/test_containers_integration_tests/controllers/service_api/dataset/__init__.py b/api/tests/e2e_tests/middleware/controllers/console/__init__.py similarity index 100% rename from api/tests/test_containers_integration_tests/controllers/service_api/dataset/__init__.py rename to api/tests/e2e_tests/middleware/controllers/console/__init__.py diff --git a/api/tests/test_containers_integration_tests/core/rag/pipeline/__init__.py b/api/tests/e2e_tests/middleware/controllers/console/app/__init__.py similarity index 100% rename from api/tests/test_containers_integration_tests/core/rag/pipeline/__init__.py rename to api/tests/e2e_tests/middleware/controllers/console/app/__init__.py diff --git a/api/tests/integration_tests/controllers/console/app/test_chat_message_permissions.py b/api/tests/e2e_tests/middleware/controllers/console/app/test_chat_message_permissions.py similarity index 100% rename from api/tests/integration_tests/controllers/console/app/test_chat_message_permissions.py rename to api/tests/e2e_tests/middleware/controllers/console/app/test_chat_message_permissions.py diff --git a/api/tests/integration_tests/controllers/console/app/test_feedback_api_basic.py b/api/tests/e2e_tests/middleware/controllers/console/app/test_feedback_api_basic.py similarity index 100% rename from api/tests/integration_tests/controllers/console/app/test_feedback_api_basic.py rename to api/tests/e2e_tests/middleware/controllers/console/app/test_feedback_api_basic.py diff --git a/api/tests/integration_tests/controllers/console/app/test_model_config_permissions.py b/api/tests/e2e_tests/middleware/controllers/console/app/test_model_config_permissions.py similarity index 100% rename from api/tests/integration_tests/controllers/console/app/test_model_config_permissions.py rename to api/tests/e2e_tests/middleware/controllers/console/app/test_model_config_permissions.py diff --git a/api/tests/test_containers_integration_tests/factories/__init__.py b/api/tests/e2e_tests/middleware/controllers/openapi/__init__.py similarity index 100% rename from api/tests/test_containers_integration_tests/factories/__init__.py rename to api/tests/e2e_tests/middleware/controllers/openapi/__init__.py diff --git a/api/tests/integration_tests/controllers/openapi/conftest.py b/api/tests/e2e_tests/middleware/controllers/openapi/conftest.py similarity index 100% rename from api/tests/integration_tests/controllers/openapi/conftest.py rename to api/tests/e2e_tests/middleware/controllers/openapi/conftest.py diff --git a/api/tests/integration_tests/controllers/openapi/test_app_run.py b/api/tests/e2e_tests/middleware/controllers/openapi/test_app_run.py similarity index 100% rename from api/tests/integration_tests/controllers/openapi/test_app_run.py rename to api/tests/e2e_tests/middleware/controllers/openapi/test_app_run.py diff --git a/api/tests/integration_tests/controllers/openapi/test_apps.py b/api/tests/e2e_tests/middleware/controllers/openapi/test_apps.py similarity index 100% rename from api/tests/integration_tests/controllers/openapi/test_apps.py rename to api/tests/e2e_tests/middleware/controllers/openapi/test_apps.py diff --git a/api/tests/test_containers_integration_tests/services/__init__.py b/api/tests/e2e_tests/middleware/factories/__init__.py similarity index 100% rename from api/tests/test_containers_integration_tests/services/__init__.py rename to api/tests/e2e_tests/middleware/factories/__init__.py diff --git a/api/tests/integration_tests/model_runtime/__mock/plugin_model.py b/api/tests/e2e_tests/middleware/model_runtime/__mock/plugin_model.py similarity index 52% rename from api/tests/integration_tests/model_runtime/__mock/plugin_model.py rename to api/tests/e2e_tests/middleware/model_runtime/__mock/plugin_model.py index 7e8c24a9f4e..fa015ffee5e 100644 --- a/api/tests/integration_tests/model_runtime/__mock/plugin_model.py +++ b/api/tests/e2e_tests/middleware/model_runtime/__mock/plugin_model.py @@ -1,5 +1,5 @@ """Compatibility import for unit tests that predate the integration-test migration.""" -from tests.unit_tests.core.workflow.nodes.parameter_extractor.fixtures.plugin_model import MockModelClass +from tests.helpers.unit.core.workflow.nodes.parameter_extractor.fixtures.plugin_model import MockModelClass __all__ = ["MockModelClass"] diff --git a/api/tests/test_containers_integration_tests/services/auth/__init__.py b/api/tests/e2e_tests/middleware/services/__init__.py similarity index 100% rename from api/tests/test_containers_integration_tests/services/auth/__init__.py rename to api/tests/e2e_tests/middleware/services/__init__.py diff --git a/api/tests/test_containers_integration_tests/services/enterprise/__init__.py b/api/tests/e2e_tests/middleware/services/plugin/__init__.py similarity index 100% rename from api/tests/test_containers_integration_tests/services/enterprise/__init__.py rename to api/tests/e2e_tests/middleware/services/plugin/__init__.py diff --git a/api/tests/test_containers_integration_tests/services/plugin/__init__.py b/api/tests/e2e_tests/middleware/services/retention/__init__.py similarity index 100% rename from api/tests/test_containers_integration_tests/services/plugin/__init__.py rename to api/tests/e2e_tests/middleware/services/retention/__init__.py diff --git a/api/tests/integration_tests/storage/test_clickzetta_volume.py b/api/tests/e2e_tests/middleware/storage/test_clickzetta_volume.py similarity index 100% rename from api/tests/integration_tests/storage/test_clickzetta_volume.py rename to api/tests/e2e_tests/middleware/storage/test_clickzetta_volume.py diff --git a/api/tests/test_containers_integration_tests/services/rag_pipeline/__init__.py b/api/tests/e2e_tests/middleware/tasks/__init__.py similarity index 100% rename from api/tests/test_containers_integration_tests/services/rag_pipeline/__init__.py rename to api/tests/e2e_tests/middleware/tasks/__init__.py diff --git a/api/tests/integration_tests/tasks/test_remove_app_and_related_data_task.py b/api/tests/e2e_tests/middleware/tasks/test_remove_app_and_related_data_task.py similarity index 100% rename from api/tests/integration_tests/tasks/test_remove_app_and_related_data_task.py rename to api/tests/e2e_tests/middleware/tasks/test_remove_app_and_related_data_task.py diff --git a/api/tests/test_containers_integration_tests/services/tools/__init__.py b/api/tests/e2e_tests/middleware/tools/__init__.py similarity index 100% rename from api/tests/test_containers_integration_tests/services/tools/__init__.py rename to api/tests/e2e_tests/middleware/tools/__init__.py diff --git a/api/tests/integration_tests/tools/__mock_server/openapi_todo.py b/api/tests/e2e_tests/middleware/tools/__mock_server/openapi_todo.py similarity index 100% rename from api/tests/integration_tests/tools/__mock_server/openapi_todo.py rename to api/tests/e2e_tests/middleware/tools/__mock_server/openapi_todo.py diff --git a/api/tests/test_containers_integration_tests/services/workflow/__init__.py b/api/tests/e2e_tests/middleware/tools/api_tool/__init__.py similarity index 100% rename from api/tests/test_containers_integration_tests/services/workflow/__init__.py rename to api/tests/e2e_tests/middleware/tools/api_tool/__init__.py diff --git a/api/tests/test_containers_integration_tests/tasks/__init__.py b/api/tests/e2e_tests/middleware/tools/code/__init__.py similarity index 100% rename from api/tests/test_containers_integration_tests/tasks/__init__.py rename to api/tests/e2e_tests/middleware/tools/code/__init__.py diff --git a/api/tests/test_containers_integration_tests/workflow/__init__.py b/api/tests/e2e_tests/middleware/vdb/__init__.py similarity index 100% rename from api/tests/test_containers_integration_tests/workflow/__init__.py rename to api/tests/e2e_tests/middleware/vdb/__init__.py diff --git a/api/tests/test_containers_integration_tests/workflow/nodes/__init__.py b/api/tests/e2e_tests/middleware/workflow/__init__.py similarity index 100% rename from api/tests/test_containers_integration_tests/workflow/nodes/__init__.py rename to api/tests/e2e_tests/middleware/workflow/__init__.py diff --git a/api/tests/test_containers_integration_tests/workflow/nodes/code_executor/__init__.py b/api/tests/e2e_tests/middleware/workflow/nodes/__init__.py similarity index 100% rename from api/tests/test_containers_integration_tests/workflow/nodes/code_executor/__init__.py rename to api/tests/e2e_tests/middleware/workflow/nodes/__init__.py diff --git a/api/tests/integration_tests/workflow/nodes/__mock/code_executor.py b/api/tests/e2e_tests/middleware/workflow/nodes/__mock/code_executor.py similarity index 100% rename from api/tests/integration_tests/workflow/nodes/__mock/code_executor.py rename to api/tests/e2e_tests/middleware/workflow/nodes/__mock/code_executor.py diff --git a/api/tests/integration_tests/workflow/nodes/__mock/http.py b/api/tests/e2e_tests/middleware/workflow/nodes/__mock/http.py similarity index 100% rename from api/tests/integration_tests/workflow/nodes/__mock/http.py rename to api/tests/e2e_tests/middleware/workflow/nodes/__mock/http.py diff --git a/api/tests/unit_tests/core/rag/indexing/__init__.py b/api/tests/e2e_tests/middleware/workflow/nodes/code_executor/__init__.py similarity index 100% rename from api/tests/unit_tests/core/rag/indexing/__init__.py rename to api/tests/e2e_tests/middleware/workflow/nodes/code_executor/__init__.py diff --git a/api/tests/unit_tests/core/workflow/graph_engine/__init__.py b/api/tests/e2e_tests/middleware/workflow/nodes/knowledge_index/__init__.py similarity index 100% rename from api/tests/unit_tests/core/workflow/graph_engine/__init__.py rename to api/tests/e2e_tests/middleware/workflow/nodes/knowledge_index/__init__.py diff --git a/api/tests/unit_tests/core/workflow/graph_engine/layers/__init__.py b/api/tests/e2e_tests/middleware/workflow/nodes/knowledge_retrieval/__init__.py similarity index 100% rename from api/tests/unit_tests/core/workflow/graph_engine/layers/__init__.py rename to api/tests/e2e_tests/middleware/workflow/nodes/knowledge_retrieval/__init__.py diff --git a/api/tests/integration_tests/workflow/nodes/test_code.py b/api/tests/e2e_tests/middleware/workflow/nodes/test_code.py similarity index 98% rename from api/tests/integration_tests/workflow/nodes/test_code.py rename to api/tests/e2e_tests/middleware/workflow/nodes/test_code.py index 8bc74f721c3..901cf3aa7db 100644 --- a/api/tests/integration_tests/workflow/nodes/test_code.py +++ b/api/tests/e2e_tests/middleware/workflow/nodes/test_code.py @@ -18,7 +18,7 @@ from graphon.nodes.code.limits import CodeNodeLimits from graphon.runtime import RuntimeState, VariablePool from tests.workflow_test_utils import build_test_graph_init_params -pytest_plugins = ("tests.integration_tests.workflow.nodes.__mock.code_executor",) +pytest_plugins = ("tests.e2e_tests.middleware.workflow.nodes.__mock.code_executor",) def _init_code_node(code_config: dict) -> CodeNode: diff --git a/api/tests/integration_tests/workflow/nodes/test_http.py b/api/tests/e2e_tests/middleware/workflow/nodes/test_http.py similarity index 99% rename from api/tests/integration_tests/workflow/nodes/test_http.py rename to api/tests/e2e_tests/middleware/workflow/nodes/test_http.py index 66f464f18c1..092f305506a 100644 --- a/api/tests/integration_tests/workflow/nodes/test_http.py +++ b/api/tests/e2e_tests/middleware/workflow/nodes/test_http.py @@ -18,7 +18,7 @@ from graphon.nodes.http_request import HttpRequestNode, HttpRequestNodeConfig, H from graphon.runtime import RuntimeState, VariablePool from tests.workflow_test_utils import build_test_graph_init_params -pytest_plugins = ("tests.integration_tests.workflow.nodes.__mock.http",) +pytest_plugins = ("tests.e2e_tests.middleware.workflow.nodes.__mock.http",) HTTP_REQUEST_CONFIG = HttpRequestNodeConfig( max_connect_timeout=dify_config.HTTP_REQUEST_MAX_CONNECT_TIMEOUT, diff --git a/api/tests/unit_tests/account_domain.py b/api/tests/helpers/unit/account_domain.py similarity index 100% rename from api/tests/unit_tests/account_domain.py rename to api/tests/helpers/unit/account_domain.py diff --git a/api/tests/unit_tests/agent_runtime_fixtures.py b/api/tests/helpers/unit/agent_runtime_fixtures.py similarity index 98% rename from api/tests/unit_tests/agent_runtime_fixtures.py rename to api/tests/helpers/unit/agent_runtime_fixtures.py index 818546ab33a..f7a37ff3757 100644 --- a/api/tests/unit_tests/agent_runtime_fixtures.py +++ b/api/tests/helpers/unit/agent_runtime_fixtures.py @@ -26,7 +26,7 @@ from graphon.model_runtime.model_providers.base.large_language_model import Larg from models.enums import ConversationFromSource, ConversationStatus from models.enums import InvokeFrom as ModelInvokeFrom from models.model import App, AppMode, Conversation, Message -from tests.unit_tests.core.model_fixtures import make_model_config +from tests.helpers.unit.core.model_fixtures import make_model_config PROVIDER = "langgenius/openai/openai" diff --git a/api/tests/unit_tests/audio_runtime_fixtures.py b/api/tests/helpers/unit/audio_runtime_fixtures.py similarity index 100% rename from api/tests/unit_tests/audio_runtime_fixtures.py rename to api/tests/helpers/unit/audio_runtime_fixtures.py diff --git a/api/tests/unit_tests/config_override.py b/api/tests/helpers/unit/config_override.py similarity index 100% rename from api/tests/unit_tests/config_override.py rename to api/tests/helpers/unit/config_override.py diff --git a/api/tests/unit_tests/configs/_isolated_settings.py b/api/tests/helpers/unit/configs/_isolated_settings.py similarity index 100% rename from api/tests/unit_tests/configs/_isolated_settings.py rename to api/tests/helpers/unit/configs/_isolated_settings.py diff --git a/api/tests/unit_tests/controllers/error_body.py b/api/tests/helpers/unit/controllers/error_body.py similarity index 100% rename from api/tests/unit_tests/controllers/error_body.py rename to api/tests/helpers/unit/controllers/error_body.py diff --git a/api/tests/unit_tests/controllers/error_contract.py b/api/tests/helpers/unit/controllers/error_contract.py similarity index 100% rename from api/tests/unit_tests/controllers/error_contract.py rename to api/tests/helpers/unit/controllers/error_contract.py diff --git a/api/tests/unit_tests/controllers/knowledge_fs_controller_fakes.py b/api/tests/helpers/unit/controllers/knowledge_fs_controller_fakes.py similarity index 100% rename from api/tests/unit_tests/controllers/knowledge_fs_controller_fakes.py rename to api/tests/helpers/unit/controllers/knowledge_fs_controller_fakes.py diff --git a/api/tests/unit_tests/controllers/openapi/_mode_constants.py b/api/tests/helpers/unit/controllers/openapi/_mode_constants.py similarity index 100% rename from api/tests/unit_tests/controllers/openapi/_mode_constants.py rename to api/tests/helpers/unit/controllers/openapi/_mode_constants.py diff --git a/api/tests/unit_tests/controllers/openapi/auth/_world.py b/api/tests/helpers/unit/controllers/openapi/auth/_world.py similarity index 100% rename from api/tests/unit_tests/controllers/openapi/auth/_world.py rename to api/tests/helpers/unit/controllers/openapi/auth/_world.py diff --git a/api/tests/unit_tests/controllers/rbac_introspection.py b/api/tests/helpers/unit/controllers/rbac_introspection.py similarity index 100% rename from api/tests/unit_tests/controllers/rbac_introspection.py rename to api/tests/helpers/unit/controllers/rbac_introspection.py diff --git a/api/tests/unit_tests/core/datasource/factories.py b/api/tests/helpers/unit/core/datasource/factories.py similarity index 100% rename from api/tests/unit_tests/core/datasource/factories.py rename to api/tests/helpers/unit/core/datasource/factories.py diff --git a/api/tests/unit_tests/core/mcp/fixtures.py b/api/tests/helpers/unit/core/mcp/fixtures.py similarity index 100% rename from api/tests/unit_tests/core/mcp/fixtures.py rename to api/tests/helpers/unit/core/mcp/fixtures.py diff --git a/api/tests/unit_tests/core/model_fixtures.py b/api/tests/helpers/unit/core/model_fixtures.py similarity index 100% rename from api/tests/unit_tests/core/model_fixtures.py rename to api/tests/helpers/unit/core/model_fixtures.py diff --git a/api/tests/unit_tests/core/ops/trace_fixtures.py b/api/tests/helpers/unit/core/ops/trace_fixtures.py similarity index 100% rename from api/tests/unit_tests/core/ops/trace_fixtures.py rename to api/tests/helpers/unit/core/ops/trace_fixtures.py diff --git a/api/tests/unit_tests/core/tools/custom_tool/api_tool_http_fixture.py b/api/tests/helpers/unit/core/tools/custom_tool/api_tool_http_fixture.py similarity index 100% rename from api/tests/unit_tests/core/tools/custom_tool/api_tool_http_fixture.py rename to api/tests/helpers/unit/core/tools/custom_tool/api_tool_http_fixture.py diff --git a/api/tests/unit_tests/core/workflow/context/__init__.py b/api/tests/helpers/unit/core/workflow/context/__init__.py similarity index 100% rename from api/tests/unit_tests/core/workflow/context/__init__.py rename to api/tests/helpers/unit/core/workflow/context/__init__.py diff --git a/api/tests/unit_tests/core/workflow/nodes/answer/__init__.py b/api/tests/helpers/unit/core/workflow/graph_engine/__init__.py similarity index 100% rename from api/tests/unit_tests/core/workflow/nodes/answer/__init__.py rename to api/tests/helpers/unit/core/workflow/graph_engine/__init__.py diff --git a/api/tests/unit_tests/core/workflow/nodes/code/__init__.py b/api/tests/helpers/unit/core/workflow/graph_engine/layers/__init__.py similarity index 100% rename from api/tests/unit_tests/core/workflow/nodes/code/__init__.py rename to api/tests/helpers/unit/core/workflow/graph_engine/layers/__init__.py diff --git a/api/tests/unit_tests/core/workflow/nodes/iteration/__init__.py b/api/tests/helpers/unit/core/workflow/nodes/answer/__init__.py similarity index 100% rename from api/tests/unit_tests/core/workflow/nodes/iteration/__init__.py rename to api/tests/helpers/unit/core/workflow/nodes/answer/__init__.py diff --git a/api/tests/unit_tests/core/workflow/nodes/parameter_extractor/__init__.py b/api/tests/helpers/unit/core/workflow/nodes/code/__init__.py similarity index 100% rename from api/tests/unit_tests/core/workflow/nodes/parameter_extractor/__init__.py rename to api/tests/helpers/unit/core/workflow/nodes/code/__init__.py diff --git a/api/tests/unit_tests/core/workflow/nodes/code/code_node_spec.py b/api/tests/helpers/unit/core/workflow/nodes/code/code_node_spec.py similarity index 100% rename from api/tests/unit_tests/core/workflow/nodes/code/code_node_spec.py rename to api/tests/helpers/unit/core/workflow/nodes/code/code_node_spec.py diff --git a/api/tests/unit_tests/core/workflow/nodes/webhook/__init__.py b/api/tests/helpers/unit/core/workflow/nodes/iteration/__init__.py similarity index 100% rename from api/tests/unit_tests/core/workflow/nodes/webhook/__init__.py rename to api/tests/helpers/unit/core/workflow/nodes/iteration/__init__.py diff --git a/api/tests/unit_tests/core/workflow/nodes/list_operator/__init__.py b/api/tests/helpers/unit/core/workflow/nodes/list_operator/__init__.py similarity index 100% rename from api/tests/unit_tests/core/workflow/nodes/list_operator/__init__.py rename to api/tests/helpers/unit/core/workflow/nodes/list_operator/__init__.py diff --git a/api/tests/unit_tests/core/workflow/nodes/list_operator/node_spec.py b/api/tests/helpers/unit/core/workflow/nodes/list_operator/node_spec.py similarity index 100% rename from api/tests/unit_tests/core/workflow/nodes/list_operator/node_spec.py rename to api/tests/helpers/unit/core/workflow/nodes/list_operator/node_spec.py diff --git a/api/tests/unit_tests/oss/__init__.py b/api/tests/helpers/unit/core/workflow/nodes/parameter_extractor/__init__.py similarity index 100% rename from api/tests/unit_tests/oss/__init__.py rename to api/tests/helpers/unit/core/workflow/nodes/parameter_extractor/__init__.py diff --git a/api/tests/unit_tests/core/workflow/nodes/parameter_extractor/fixtures/model.py b/api/tests/helpers/unit/core/workflow/nodes/parameter_extractor/fixtures/model.py similarity index 100% rename from api/tests/unit_tests/core/workflow/nodes/parameter_extractor/fixtures/model.py rename to api/tests/helpers/unit/core/workflow/nodes/parameter_extractor/fixtures/model.py diff --git a/api/tests/unit_tests/core/workflow/nodes/parameter_extractor/fixtures/plugin_daemon.py b/api/tests/helpers/unit/core/workflow/nodes/parameter_extractor/fixtures/plugin_daemon.py similarity index 88% rename from api/tests/unit_tests/core/workflow/nodes/parameter_extractor/fixtures/plugin_daemon.py rename to api/tests/helpers/unit/core/workflow/nodes/parameter_extractor/fixtures/plugin_daemon.py index a7bf449136f..749ca81934f 100644 --- a/api/tests/unit_tests/core/workflow/nodes/parameter_extractor/fixtures/plugin_daemon.py +++ b/api/tests/helpers/unit/core/workflow/nodes/parameter_extractor/fixtures/plugin_daemon.py @@ -3,7 +3,7 @@ from collections.abc import Callable, Generator import pytest from core.plugin.impl.model import PluginModelClient -from tests.unit_tests.core.workflow.nodes.parameter_extractor.fixtures.plugin_model import MockModelClass +from tests.helpers.unit.core.workflow.nodes.parameter_extractor.fixtures.plugin_model import MockModelClass def mock_plugin_daemon( diff --git a/api/tests/unit_tests/core/workflow/nodes/parameter_extractor/fixtures/plugin_model.py b/api/tests/helpers/unit/core/workflow/nodes/parameter_extractor/fixtures/plugin_model.py similarity index 100% rename from api/tests/unit_tests/core/workflow/nodes/parameter_extractor/fixtures/plugin_model.py rename to api/tests/helpers/unit/core/workflow/nodes/parameter_extractor/fixtures/plugin_model.py diff --git a/api/tests/unit_tests/core/workflow/nodes/template_transform/__init__.py b/api/tests/helpers/unit/core/workflow/nodes/template_transform/__init__.py similarity index 100% rename from api/tests/unit_tests/core/workflow/nodes/template_transform/__init__.py rename to api/tests/helpers/unit/core/workflow/nodes/template_transform/__init__.py diff --git a/api/tests/unit_tests/core/workflow/nodes/template_transform/template_transform_node_spec.py b/api/tests/helpers/unit/core/workflow/nodes/template_transform/template_transform_node_spec.py similarity index 100% rename from api/tests/unit_tests/core/workflow/nodes/template_transform/template_transform_node_spec.py rename to api/tests/helpers/unit/core/workflow/nodes/template_transform/template_transform_node_spec.py diff --git a/api/tests/unit_tests/core/workflow/nodes/tool/__init__.py b/api/tests/helpers/unit/core/workflow/nodes/tool/__init__.py similarity index 100% rename from api/tests/unit_tests/core/workflow/nodes/tool/__init__.py rename to api/tests/helpers/unit/core/workflow/nodes/tool/__init__.py diff --git a/api/tests/unit_tests/oss/__mock/__init__.py b/api/tests/helpers/unit/core/workflow/nodes/webhook/__init__.py similarity index 100% rename from api/tests/unit_tests/oss/__mock/__init__.py rename to api/tests/helpers/unit/core/workflow/nodes/webhook/__init__.py diff --git a/api/tests/unit_tests/credential_query_fakes.py b/api/tests/helpers/unit/credential_query_fakes.py similarity index 100% rename from api/tests/unit_tests/credential_query_fakes.py rename to api/tests/helpers/unit/credential_query_fakes.py diff --git a/api/tests/unit_tests/file_grant_test_utils.py b/api/tests/helpers/unit/file_grant_test_utils.py similarity index 100% rename from api/tests/unit_tests/file_grant_test_utils.py rename to api/tests/helpers/unit/file_grant_test_utils.py diff --git a/api/tests/unit_tests/knowledge_fs_core_fakes.py b/api/tests/helpers/unit/knowledge_fs_core_fakes.py similarity index 100% rename from api/tests/unit_tests/knowledge_fs_core_fakes.py rename to api/tests/helpers/unit/knowledge_fs_core_fakes.py diff --git a/api/tests/unit_tests/libs/_human_input/__init__.py b/api/tests/helpers/unit/libs/_human_input/__init__.py similarity index 100% rename from api/tests/unit_tests/libs/_human_input/__init__.py rename to api/tests/helpers/unit/libs/_human_input/__init__.py diff --git a/api/tests/unit_tests/libs/_human_input/support.py b/api/tests/helpers/unit/libs/_human_input/support.py similarity index 100% rename from api/tests/unit_tests/libs/_human_input/support.py rename to api/tests/helpers/unit/libs/_human_input/support.py diff --git a/api/tests/unit_tests/model_factories.py b/api/tests/helpers/unit/model_factories.py similarity index 100% rename from api/tests/unit_tests/model_factories.py rename to api/tests/helpers/unit/model_factories.py diff --git a/api/tests/unit_tests/oss/aliyun_oss/aliyun_oss/__init__.py b/api/tests/helpers/unit/oss/__init__.py similarity index 100% rename from api/tests/unit_tests/oss/aliyun_oss/aliyun_oss/__init__.py rename to api/tests/helpers/unit/oss/__init__.py diff --git a/api/tests/unit_tests/oss/opendal/__init__.py b/api/tests/helpers/unit/oss/__mock/__init__.py similarity index 100% rename from api/tests/unit_tests/oss/opendal/__init__.py rename to api/tests/helpers/unit/oss/__mock/__init__.py diff --git a/api/tests/unit_tests/oss/__mock/aliyun_oss.py b/api/tests/helpers/unit/oss/__mock/aliyun_oss.py similarity index 98% rename from api/tests/unit_tests/oss/__mock/aliyun_oss.py rename to api/tests/helpers/unit/oss/__mock/aliyun_oss.py index e47644802bf..bf7d49ddca1 100644 --- a/api/tests/unit_tests/oss/__mock/aliyun_oss.py +++ b/api/tests/helpers/unit/oss/__mock/aliyun_oss.py @@ -8,7 +8,7 @@ from oss2 import Bucket from oss2.http import Response as OssResponse from oss2.models import GetObjectResult, PutObjectResult -from tests.unit_tests.oss.__mock.base import ( +from tests.helpers.unit.oss.__mock.base import ( get_example_bucket, get_example_data, get_example_filename, diff --git a/api/tests/unit_tests/oss/__mock/baidu_obs.py b/api/tests/helpers/unit/oss/__mock/baidu_obs.py similarity index 98% rename from api/tests/unit_tests/oss/__mock/baidu_obs.py rename to api/tests/helpers/unit/oss/__mock/baidu_obs.py index 3a9d68f9db7..3a22f8da39c 100644 --- a/api/tests/unit_tests/oss/__mock/baidu_obs.py +++ b/api/tests/helpers/unit/oss/__mock/baidu_obs.py @@ -8,7 +8,7 @@ from types import SimpleNamespace import pytest from baidubce.services.bos.bos_client import BosClient -from tests.unit_tests.oss.__mock.base import ( +from tests.helpers.unit.oss.__mock.base import ( get_example_bucket, get_example_data, get_example_filename, diff --git a/api/tests/unit_tests/oss/__mock/base.py b/api/tests/helpers/unit/oss/__mock/base.py similarity index 100% rename from api/tests/unit_tests/oss/__mock/base.py rename to api/tests/helpers/unit/oss/__mock/base.py diff --git a/api/tests/unit_tests/oss/__mock/local.py b/api/tests/helpers/unit/oss/__mock/local.py similarity index 97% rename from api/tests/unit_tests/oss/__mock/local.py rename to api/tests/helpers/unit/oss/__mock/local.py index 92f42bb2ef2..971dacdae75 100644 --- a/api/tests/unit_tests/oss/__mock/local.py +++ b/api/tests/helpers/unit/oss/__mock/local.py @@ -6,7 +6,7 @@ from unittest.mock import MagicMock, mock_open, patch import pytest -from tests.unit_tests.oss.__mock.base import ( +from tests.helpers.unit.oss.__mock.base import ( get_example_data, get_example_filename, get_example_filepath, diff --git a/api/tests/unit_tests/oss/__mock/tencent_cos.py b/api/tests/helpers/unit/oss/__mock/tencent_cos.py similarity index 98% rename from api/tests/unit_tests/oss/__mock/tencent_cos.py rename to api/tests/helpers/unit/oss/__mock/tencent_cos.py index 8fcf7666ce4..373625efb6d 100644 --- a/api/tests/unit_tests/oss/__mock/tencent_cos.py +++ b/api/tests/helpers/unit/oss/__mock/tencent_cos.py @@ -5,7 +5,7 @@ import pytest from qcloud_cos import CosS3Client from qcloud_cos.streambody import StreamBody -from tests.unit_tests.oss.__mock.base import ( +from tests.helpers.unit.oss.__mock.base import ( get_example_bucket, get_example_data, get_example_filename, diff --git a/api/tests/unit_tests/oss/__mock/volcengine_tos.py b/api/tests/helpers/unit/oss/__mock/volcengine_tos.py similarity index 98% rename from api/tests/unit_tests/oss/__mock/volcengine_tos.py rename to api/tests/helpers/unit/oss/__mock/volcengine_tos.py index ceb6fd463c4..5c32d3ab8c4 100644 --- a/api/tests/unit_tests/oss/__mock/volcengine_tos.py +++ b/api/tests/helpers/unit/oss/__mock/volcengine_tos.py @@ -7,7 +7,7 @@ from tos import TosClientV2 from tos.clientv2 import DeleteObjectOutput, GetObjectOutput, HeadObjectOutput, PutObjectOutput from tos.http import Response as TosResponse -from tests.unit_tests.oss.__mock.base import ( +from tests.helpers.unit.oss.__mock.base import ( get_example_bucket, get_example_data, get_example_filename, diff --git a/api/tests/unit_tests/oss/tencent_cos/__init__.py b/api/tests/helpers/unit/oss/aliyun_oss/aliyun_oss/__init__.py similarity index 100% rename from api/tests/unit_tests/oss/tencent_cos/__init__.py rename to api/tests/helpers/unit/oss/aliyun_oss/aliyun_oss/__init__.py diff --git a/api/tests/unit_tests/oss/baidu_obs/__init__.py b/api/tests/helpers/unit/oss/baidu_obs/__init__.py similarity index 100% rename from api/tests/unit_tests/oss/baidu_obs/__init__.py rename to api/tests/helpers/unit/oss/baidu_obs/__init__.py diff --git a/api/tests/unit_tests/oss/volcengine_tos/__init__.py b/api/tests/helpers/unit/oss/opendal/__init__.py similarity index 100% rename from api/tests/unit_tests/oss/volcengine_tos/__init__.py rename to api/tests/helpers/unit/oss/opendal/__init__.py diff --git a/api/tests/unit_tests/tools/__init__.py b/api/tests/helpers/unit/oss/tencent_cos/__init__.py similarity index 100% rename from api/tests/unit_tests/tools/__init__.py rename to api/tests/helpers/unit/oss/tencent_cos/__init__.py diff --git a/api/tests/unit_tests/utils/__init__.py b/api/tests/helpers/unit/oss/volcengine_tos/__init__.py similarity index 100% rename from api/tests/unit_tests/utils/__init__.py rename to api/tests/helpers/unit/oss/volcengine_tos/__init__.py diff --git a/api/tests/unit_tests/rbac_fakes.py b/api/tests/helpers/unit/rbac_fakes.py similarity index 100% rename from api/tests/unit_tests/rbac_fakes.py rename to api/tests/helpers/unit/rbac_fakes.py diff --git a/api/tests/unit_tests/repositories/app/console_visibility.py b/api/tests/helpers/unit/repositories/app/console_visibility.py similarity index 100% rename from api/tests/unit_tests/repositories/app/console_visibility.py rename to api/tests/helpers/unit/repositories/app/console_visibility.py diff --git a/api/tests/unit_tests/services/controller_api.py b/api/tests/helpers/unit/services/controller_api.py similarity index 100% rename from api/tests/unit_tests/services/controller_api.py rename to api/tests/helpers/unit/services/controller_api.py diff --git a/api/tests/unit_tests/services/dataset_service_test_helpers.py b/api/tests/helpers/unit/services/dataset_service_test_helpers.py similarity index 100% rename from api/tests/unit_tests/services/dataset_service_test_helpers.py rename to api/tests/helpers/unit/services/dataset_service_test_helpers.py diff --git a/api/tests/unit_tests/services/document_service_validation.py b/api/tests/helpers/unit/services/document_service_validation.py similarity index 100% rename from api/tests/unit_tests/services/document_service_validation.py rename to api/tests/helpers/unit/services/document_service_validation.py diff --git a/api/tests/unit_tests/services/hit_service.py b/api/tests/helpers/unit/services/hit_service.py similarity index 100% rename from api/tests/unit_tests/services/hit_service.py rename to api/tests/helpers/unit/services/hit_service.py diff --git a/api/tests/unit_tests/services/knowledge_fs_fakes.py b/api/tests/helpers/unit/services/knowledge_fs_fakes.py similarity index 100% rename from api/tests/unit_tests/services/knowledge_fs_fakes.py rename to api/tests/helpers/unit/services/knowledge_fs_fakes.py diff --git a/api/tests/unit_tests/services/knowledge_fs_services_a_fakes.py b/api/tests/helpers/unit/services/knowledge_fs_services_a_fakes.py similarity index 100% rename from api/tests/unit_tests/services/knowledge_fs_services_a_fakes.py rename to api/tests/helpers/unit/services/knowledge_fs_services_a_fakes.py diff --git a/api/tests/unit_tests/services/knowledge_fs_services_b_fakes.py b/api/tests/helpers/unit/services/knowledge_fs_services_b_fakes.py similarity index 100% rename from api/tests/unit_tests/services/knowledge_fs_services_b_fakes.py rename to api/tests/helpers/unit/services/knowledge_fs_services_b_fakes.py diff --git a/api/tests/unit_tests/services/workflow_collaboration_fixtures.py b/api/tests/helpers/unit/services/workflow_collaboration_fixtures.py similarity index 100% rename from api/tests/unit_tests/services/workflow_collaboration_fixtures.py rename to api/tests/helpers/unit/services/workflow_collaboration_fixtures.py diff --git a/api/tests/unit_tests/tasks/knowledge_fs_task_fakes.py b/api/tests/helpers/unit/tasks/knowledge_fs_task_fakes.py similarity index 100% rename from api/tests/unit_tests/tasks/knowledge_fs_task_fakes.py rename to api/tests/helpers/unit/tasks/knowledge_fs_task_fakes.py diff --git a/api/tests/unit_tests/tasks/task_options.py b/api/tests/helpers/unit/tasks/task_options.py similarity index 100% rename from api/tests/unit_tests/tasks/task_options.py rename to api/tests/helpers/unit/tasks/task_options.py diff --git a/api/tests/unit_tests/utils/structured_output_parser/__init__.py b/api/tests/helpers/unit/utils/__init__.py similarity index 100% rename from api/tests/unit_tests/utils/structured_output_parser/__init__.py rename to api/tests/helpers/unit/utils/__init__.py diff --git a/api/tests/unit_tests/utils/module_import_fixtures/child_class.py b/api/tests/helpers/unit/utils/module_import_fixtures/child_class.py similarity index 60% rename from api/tests/unit_tests/utils/module_import_fixtures/child_class.py rename to api/tests/helpers/unit/utils/module_import_fixtures/child_class.py index af05af1bd9c..a6c35e9e72d 100644 --- a/api/tests/unit_tests/utils/module_import_fixtures/child_class.py +++ b/api/tests/helpers/unit/utils/module_import_fixtures/child_class.py @@ -1,4 +1,4 @@ -from tests.unit_tests.utils.module_import_fixtures.parent_class import ParentClass +from tests.helpers.unit.utils.module_import_fixtures.parent_class import ParentClass class ChildClass(ParentClass): diff --git a/api/tests/unit_tests/utils/module_import_fixtures/lazy_load_class.py b/api/tests/helpers/unit/utils/module_import_fixtures/lazy_load_class.py similarity index 62% rename from api/tests/unit_tests/utils/module_import_fixtures/lazy_load_class.py rename to api/tests/helpers/unit/utils/module_import_fixtures/lazy_load_class.py index 8f4454b9639..0b0888eb904 100644 --- a/api/tests/unit_tests/utils/module_import_fixtures/lazy_load_class.py +++ b/api/tests/helpers/unit/utils/module_import_fixtures/lazy_load_class.py @@ -1,4 +1,4 @@ -from tests.unit_tests.utils.module_import_fixtures.parent_class import ParentClass +from tests.helpers.unit.utils.module_import_fixtures.parent_class import ParentClass class LazyLoadChildClass(ParentClass): diff --git a/api/tests/unit_tests/utils/module_import_fixtures/parent_class.py b/api/tests/helpers/unit/utils/module_import_fixtures/parent_class.py similarity index 100% rename from api/tests/unit_tests/utils/module_import_fixtures/parent_class.py rename to api/tests/helpers/unit/utils/module_import_fixtures/parent_class.py diff --git a/api/tests/unit_tests/utils/yaml/__init__.py b/api/tests/helpers/unit/utils/structured_output_parser/__init__.py similarity index 100% rename from api/tests/unit_tests/utils/yaml/__init__.py rename to api/tests/helpers/unit/utils/structured_output_parser/__init__.py diff --git a/api/tests/unit_tests/core/workflow/graph_engine/test_mock_config.py b/api/tests/helpers/workflow/mock_config.py similarity index 100% rename from api/tests/unit_tests/core/workflow/graph_engine/test_mock_config.py rename to api/tests/helpers/workflow/mock_config.py diff --git a/api/tests/unit_tests/core/workflow/graph_engine/test_mock_factory.py b/api/tests/helpers/workflow/mock_factory.py similarity index 98% rename from api/tests/unit_tests/core/workflow/graph_engine/test_mock_factory.py rename to api/tests/helpers/workflow/mock_factory.py index 48296df3a05..18ae1d0edd1 100644 --- a/api/tests/unit_tests/core/workflow/graph_engine/test_mock_factory.py +++ b/api/tests/helpers/workflow/mock_factory.py @@ -13,8 +13,7 @@ from graphon.entities.base_node_data import BaseNodeData from graphon.entities.graph_config import NodeConfigDict, NodeConfigDictAdapter from graphon.enums import BuiltinNodeTypes, NodeType from graphon.nodes.base.node import Node - -from .test_mock_nodes import ( +from tests.helpers.workflow.mock_nodes import ( MockAgentNode, MockCodeNode, MockDocumentExtractorNode, @@ -31,8 +30,7 @@ from .test_mock_nodes import ( if TYPE_CHECKING: from graphon.runtime import InitParams, RuntimeState - - from .test_mock_config import MockConfig + from tests.helpers.workflow.mock_config import MockConfig class MockNodeFactory(DifyNodeFactory): diff --git a/api/tests/unit_tests/core/workflow/graph_engine/test_mock_nodes.py b/api/tests/helpers/workflow/mock_nodes.py similarity index 99% rename from api/tests/unit_tests/core/workflow/graph_engine/test_mock_nodes.py rename to api/tests/helpers/workflow/mock_nodes.py index 8e77e23d8f6..377494360ec 100644 --- a/api/tests/unit_tests/core/workflow/graph_engine/test_mock_nodes.py +++ b/api/tests/helpers/workflow/mock_nodes.py @@ -37,12 +37,11 @@ from graphon.nodes.question_classifier import QuestionClassifierNode from graphon.nodes.template_transform import TemplateTransformNode from graphon.nodes.tool import ToolNode from graphon.template_rendering import Jinja2TemplateRenderer, TemplateRenderError -from tests.unit_tests.core.model_fixtures import make_model_instance +from tests.helpers.unit.core.model_fixtures import make_model_instance if TYPE_CHECKING: from graphon.runtime import InitParams, RuntimeState - - from .test_mock_config import MockConfig + from tests.helpers.workflow.mock_config import MockConfig class _TestJinja2Renderer(Jinja2TemplateRenderer): diff --git a/api/tests/unit_tests/core/workflow/graph_engine/test_table_runner.py b/api/tests/helpers/workflow/table_runner.py similarity index 99% rename from api/tests/unit_tests/core/workflow/graph_engine/test_table_runner.py rename to api/tests/helpers/workflow/table_runner.py index cd2b061f792..bbe53521c82 100644 --- a/api/tests/unit_tests/core/workflow/graph_engine/test_table_runner.py +++ b/api/tests/helpers/workflow/table_runner.py @@ -43,9 +43,8 @@ from graphon.variables import ( ObjectVariable, StringVariable, ) - -from .test_mock_config import MockConfig -from .test_mock_factory import MockNodeFactory +from tests.helpers.workflow.mock_config import MockConfig +from tests.helpers.workflow.mock_factory import MockNodeFactory logger = logging.getLogger(__name__) diff --git a/api/tests/unit_tests/core/workflow/test_variable_pool_conver.py b/api/tests/helpers/workflow/variable_pool_conver.py similarity index 100% rename from api/tests/unit_tests/core/workflow/test_variable_pool_conver.py rename to api/tests/helpers/workflow/variable_pool_conver.py diff --git a/api/tests/pytest_dify.py b/api/tests/pytest_dify.py index 4451dfdc2d3..a57899d34f2 100644 --- a/api/tests/pytest_dify.py +++ b/api/tests/pytest_dify.py @@ -41,7 +41,7 @@ def parse_services(value: str) -> list[str]: def ensure_backend_test_environment(repo_root: Path) -> None: """Set deterministic defaults needed before test conftests import application config.""" - integration_tests_dir = repo_root / "api" / "tests" / "integration_tests" + integration_tests_dir = repo_root / "api" / "tests" / "e2e_tests" / "middleware" test_env_file = integration_tests_dir / ".env" test_env_example_file = integration_tests_dir / ".env.example" vdb_env_file = integration_tests_dir / "vdb.env" diff --git a/api/tests/unit_tests/.ruff.toml b/api/tests/unit_tests/.ruff.toml index ede09ab84c2..f28b9056483 100644 --- a/api/tests/unit_tests/.ruff.toml +++ b/api/tests/unit_tests/.ruff.toml @@ -7,34 +7,31 @@ extend-select = ["ANN401", "ARG"] # Existing strict-mode debt. Remove a file entry when bringing it under strict checking. [lint.per-file-ignores] "clients/agent_backend/test_request_builder.py" = ["TID251"] -"commands/test_archive_workflow_runs.py" = ["ARG005"] -"commands/test_legacy_model_type_migration.py" = ["ARG001", "ARG002", "ARG005"] -"repositories/test_app_definition_query_repository_agent.py" = ["ARG005", "TID251"] -"controllers/console/agent/test_agent_controllers.py" = ["ARG001", "ARG002", "ARG003", "ARG005", "TID251"] +"commands/test_retention.py" = ["ARG005"] +"commands/test_data_migration.py" = ["ARG001", "ARG002", "ARG005"] +"repositories/test_app_definition_query_repository.py" = ["ARG005", "TID251"] +"controllers/console/agent/test_roster.py" = ["ARG001", "ARG002", "ARG003", "ARG005", "TID251"] "controllers/console/app/test_agent_config_inspector.py" = ["ARG005"] "controllers/console/app/test_agent_manage_guard.py" = ["ARG001"] -"controllers/console/app/test_annotation_security.py" = ["ARG002"] -"controllers/console/app/test_app_apis.py" = ["ARG001", "ARG002"] -"controllers/console/app/test_app_import_api.py" = ["ARG001", "ARG002", "ARG005"] -"controllers/console/app/test_app_response_models.py" = ["ARG002", "ARG004", "ARG005"] -"controllers/console/app/test_conversation_api.py" = ["ARG001"] -"controllers/console/app/test_generator_api.py" = ["ARG001"] -"controllers/console/app/test_mcp_server_response.py" = ["ARG002"] -"controllers/console/app/test_message_api.py" = ["ARG001"] -"controllers/console/app/test_statistic_api.py" = ["ANN401", "ARG001", "TID251"] -"controllers/console/app/test_workflow.py" = ["ARG001", "ARG005"] -"controllers/console/app/test_workflow_convert_api.py" = ["ARG005"] +"controllers/console/app/test_annotation.py" = ["ARG002"] +"controllers/console/app/test_app.py" = ["ARG001", "ARG002", "ARG004", "ARG005"] +"controllers/console/app/test_app_import.py" = ["ARG001", "ARG002", "ARG005"] +"controllers/console/app/test_conversation.py" = ["ARG001"] +"controllers/console/app/test_generator.py" = ["ARG001"] +"controllers/console/app/test_mcp_server.py" = ["ARG002"] +"controllers/console/app/test_message.py" = ["ARG001", "ARG004"] +"controllers/console/app/test_statistic.py" = ["ANN401", "ARG001", "TID251"] +"controllers/console/app/test_workflow.py" = ["ARG001", "ARG004", "ARG005"] "controllers/console/app/test_workflow_node_output_inspector.py" = ["ANN401", "ARG001", "TID251"] -"controllers/console/app/test_workflow_run_api.py" = ["ANN401", "TID251"] -"controllers/console/app/workflow_draft_variables_test.py" = ["TID251"] -"controllers/console/auth/test_account_activation.py" = ["ARG002"] -"controllers/console/auth/test_login_logout.py" = ["ARG002"] +"controllers/console/app/test_workflow_run.py" = ["ANN401", "TID251"] +"fields/test_workflow_draft_variable_fields.py" = ["TID251"] +"controllers/console/auth/test_activate.py" = ["ARG002"] +"controllers/console/auth/test_login.py" = ["ARG002"] "controllers/console/auth/test_oauth.py" = ["ARG002"] "controllers/console/auth/test_oauth_timezone.py" = ["ARG001"] "controllers/console/billing/test_billing.py" = ["ARG002"] "controllers/console/datasets/test_datasets.py" = ["ARG005"] "controllers/console/datasets/test_datasets_document.py" = ["ARG002", "ARG005"] -"controllers/console/datasets/test_datasets_document_download.py" = ["ARG005"] "controllers/console/datasets/test_datasets_segments.py" = ["TID251"] "controllers/console/datasets/test_external.py" = ["TID251"] "controllers/console/explore/test_trial.py" = ["ANN401", "ARG001", "ARG005", "TID251"] @@ -46,27 +43,25 @@ extend-select = ["ANN401", "ARG"] "controllers/console/test_wraps.py" = ["ARG001", "ARG002"] "controllers/console/workspace/test_plugin.py" = ["ARG002", "TID251"] "controllers/console/workspace/test_snippets.py" = ["ARG002"] -"controllers/console/workspace/test_tool_providers.py" = ["ARG001", "ARG005"] +"controllers/console/workspace/test_tool_providers.py" = ["ARG001", "ARG004", "ARG005"] "controllers/inner_api/plugin/test_plugin.py" = ["ARG002"] -"controllers/inner_api/plugin/test_plugin_wraps.py" = ["ARG001", "ARG002", "ARG003", "TID251"] +"controllers/inner_api/plugin/test_wraps.py" = ["ARG001", "ARG002", "ARG003", "TID251"] "controllers/inner_api/test_runtime_credentials.py" = ["ARG001"] "controllers/mcp/test_mcp.py" = ["ARG002"] "controllers/openapi/test_account.py" = ["ARG005"] -"controllers/openapi/test_app_run_streaming.py" = ["ARG001"] -"controllers/openapi/test_contract.py" = ["ARG001", "TID251"] -"controllers/openapi/test_error_contract.py" = ["ARG002"] +"controllers/openapi/test_app_run.py" = ["ARG001"] +"controllers/openapi/test__contract.py" = ["ARG001", "TID251"] +"controllers/openapi/test__errors.py" = ["ARG002"] "controllers/openapi/test_human_input_form.py" = ["ARG002"] -"controllers/openapi/test_workflow_events_openapi.py" = ["ARG002", "ARG005"] -"controllers/openapi/test_workspaces_members.py" = ["ARG001"] +"controllers/openapi/test_workflow_events.py" = ["ARG002", "ARG005"] +"controllers/openapi/test_workspaces.py" = ["ARG001"] "controllers/service_api/app/test_app.py" = ["ARG002"] -"controllers/service_api/app/test_completion.py" = ["ARG002"] +"controllers/service_api/app/test_completion.py" = ["ARG001", "ARG002", "ARG004"] "controllers/service_api/app/test_file.py" = ["ARG002"] -"controllers/service_api/app/test_hitl_service_api.py" = ["ARG002", "ARG005"] -"controllers/service_api/app/test_workflow_events.py" = ["ARG005"] +"controllers/service_api/app/test_workflow_events.py" = ["ARG002", "ARG005"] "controllers/service_api/dataset/rag_pipeline/test_rag_pipeline_workflow.py" = ["ARG002"] -"controllers/service_api/dataset/test_dataset_segment.py" = ["ARG002"] +"controllers/service_api/dataset/test_segment.py" = ["ARG002"] "controllers/service_api/dataset/test_document.py" = ["ARG001", "ARG002"] -"controllers/service_api/test_trace_session_id_parsing.py" = ["ARG001"] "controllers/service_api/test_wraps.py" = ["ARG001", "ARG002"] "controllers/trigger/test_trigger.py" = ["ARG002"] "controllers/trigger/test_webhook.py" = ["ARG002"] @@ -74,245 +69,219 @@ extend-select = ["ANN401", "ARG"] "controllers/web/test_audio.py" = ["ARG002"] "controllers/web/test_completion.py" = ["ARG002"] "controllers/web/test_human_input_form.py" = ["ARG001", "ARG002", "ARG005"] -"controllers/web/test_message_endpoints.py" = ["ARG002"] -"controllers/web/test_web_login.py" = ["ARG002"] +"controllers/web/test_message.py" = ["ARG002"] +"controllers/web/test_login.py" = ["ARG002"] "controllers/web/test_workflow.py" = ["ARG002"] "core/agent/test_base_agent_runner.py" = ["ARG002"] "core/agent/test_cot_agent_runner.py" = ["ARG001"] "core/agent/test_cot_chat_agent_runner.py" = ["ARG002"] "core/agent/test_fc_agent_runner.py" = ["TID251"] -"core/app/app_config/easy_ui_based_app/test_model_config_converter.py" = ["ARG002"] -"core/app/app_config/easy_ui_based_app/test_variables_manager.py" = ["ARG002"] -"core/app/apps/advanced_chat/test_app_generator.py" = ["ARG001", "ARG002", "ARG005"] -"core/app/apps/advanced_chat/test_app_runner_input_moderation.py" = ["ARG001", "ARG005"] -"core/app/apps/advanced_chat/test_generate_task_pipeline.py" = ["ARG005"] -"core/app/apps/advanced_chat/test_generate_task_pipeline_core.py" = ["ARG001", "ARG002", "ARG005"] -"core/app/apps/agent_app/test_app_generator.py" = ["ARG001", "ARG005"] +"core/app/app_config/easy_ui_based_app/model_config/test_converter.py" = ["ARG002"] +"core/app/app_config/easy_ui_based_app/variables/test_manager.py" = ["ARG002"] +"core/app/apps/advanced_chat/test_app_generator.py" = ["ANN401", "ARG001", "ARG002", "ARG005", "TID251"] +"core/app/apps/advanced_chat/test_app_runner.py" = ["ARG001", "ARG005"] +"core/app/apps/advanced_chat/test_generate_task_pipeline.py" = ["ARG001", "ARG002", "ARG005"] +"core/app/apps/agent_app/test_app_generator.py" = ["ANN401", "ARG001", "ARG002", "ARG005", "TID251"] "core/app/apps/agent_app/test_app_runner.py" = ["ANN401", "ARG002", "ARG005", "TID251"] -"core/app/apps/agent_app/test_input_guards.py" = ["ANN401", "ARG002", "TID251"] -"core/app/apps/agent_app/test_resolve_agent.py" = ["ANN401", "TID251"] "core/app/apps/agent_app/test_runtime_request_builder.py" = ["ARG002", "TID251"] -"core/app/apps/agent_chat/test_agent_chat_app_config_manager.py" = ["ARG005"] -"core/app/apps/agent_chat/test_agent_chat_app_generator.py" = ["ARG001"] +"core/app/apps/agent_chat/test_app_config_manager.py" = ["ARG005"] +"core/app/apps/agent_chat/test_app_generator.py" = ["ARG001"] "core/app/apps/chat/test_app_config_manager.py" = ["ARG001"] -"core/app/apps/chat/test_app_generator_and_runner.py" = ["ARG001", "ARG002", "ARG005"] -"core/app/apps/common/test_workflow_response_converter_truncation.py" = ["TID251"] +"core/app/apps/chat/test_app_generator.py" = ["ARG001", "ARG002", "ARG005"] +"core/app/apps/common/test_workflow_response_converter.py" = ["ARG005", "TID251"] "core/app/apps/completion/test_app_runner.py" = ["ARG001", "ARG002", "ARG005"] "core/app/apps/pipeline/test_pipeline_generator.py" = ["ARG001", "ARG002", "ARG005"] "core/app/apps/pipeline/test_pipeline_runner.py" = ["ARG001", "ARG002"] -"core/app/apps/test_advanced_chat_app_generator.py" = ["ARG001"] "core/app/apps/test_base_app_generator.py" = ["ARG005"] "core/app/apps/test_base_app_runner.py" = ["ARG001", "ARG002", "ARG005"] -"core/app/apps/test_pause_resume.py" = ["ANN401", "TID251"] "core/app/apps/test_streaming_utils.py" = ["ARG001"] -"core/app/apps/test_workflow_app_generator.py" = ["ARG005"] -"core/app/apps/test_workflow_app_runner_core.py" = ["ARG001", "ARG002", "ARG004", "ARG005"] -"core/app/apps/test_workflow_app_runner_single_node.py" = ["ANN401", "TID251"] -"core/app/apps/test_workflow_pause_events.py" = ["ARG005"] -"core/app/apps/workflow/test_app_generator_extra.py" = ["ARG005"] -"core/app/apps/workflow/test_generate_task_pipeline_core.py" = ["ARG002", "ARG005"] +"core/app/apps/workflow/test_app_generator.py" = ["ARG005"] +"core/app/apps/test_workflow_app_runner.py" = ["ANN401", "ARG001", "ARG002", "ARG004", "ARG005", "TID251"] +"core/app/apps/workflow/test_generate_task_pipeline.py" = ["ARG002", "ARG005"] "core/app/features/rate_limiting/test_rate_limit.py" = ["ARG001"] -"core/app/task_pipeline/test_easy_ui_based_generate_task_pipeline.py" = ["ARG002"] -"core/app/task_pipeline/test_easy_ui_based_generate_task_pipeline_core.py" = ["ARG001", "ARG002", "ARG005"] -"core/app/task_pipeline/test_message_cycle_manager_optimization.py" = ["ARG002"] -"core/app/test_easy_ui_model_config_manager.py" = ["ARG005"] -"core/app/workflow/layers/test_persistence_inspector_publish.py" = ["ANN401", "ARG005", "TID251"] +"core/app/task_pipeline/test_easy_ui_based_generate_task_pipeline.py" = ["ARG001", "ARG002", "ARG005"] +"core/app/task_pipeline/test_message_cycle_manager.py" = ["ARG002"] +"core/app/app_config/easy_ui_based_app/model_config/test_manager.py" = ["ARG005"] +"services/test_workflow_persistence.py" = ["ANN401", "ARG001", "ARG005", "TID251"] "core/app/workflow/test_file_runtime.py" = ["ARG001", "ARG005"] -"core/app/workflow/test_observability_layer_extra.py" = ["ARG005"] -"core/app/workflow/test_persistence_layer.py" = ["ARG001"] -"core/base/test_app_generator_tts_publisher.py" = ["ARG002"] +"core/app/workflow/layers/test_observability.py" = ["ARG002", "ARG005"] +"core/base/tts/test_app_generator_tts_publisher.py" = ["ARG002"] "core/callback_handler/test_agent_tool_callback_handler.py" = ["ARG002"] "core/callback_handler/test_workflow_tool_callback_handler.py" = ["ARG002"] -"core/entities/test_entities_mcp_provider.py" = ["ARG001"] -"core/entities/test_entities_provider_configuration.py" = ["ANN401", "ARG001", "ARG005", "TID251"] +"core/entities/test_mcp_provider.py" = ["ARG001"] +"core/entities/test_provider_configuration.py" = ["ANN401", "ARG001", "ARG005", "TID251"] "core/external_data_tool/api/test_api.py" = ["ARG001"] "core/external_data_tool/test_base.py" = ["TID251"] "core/external_data_tool/test_external_data_fetch.py" = ["ARG001"] "core/helper/code_executor/test_code_executor.py" = ["TID251"] "core/llm_generator/test_llm_generator.py" = ["ARG002"] "core/mcp/auth/test_auth_flow.py" = ["ARG002"] -"core/mcp/client/test_session.py" = ["ARG001", "TID251"] -"core/mcp/client/test_sse.py" = ["ARG001", "TID251"] -"core/mcp/client/test_streamable_http.py" = ["ARG001", "ARG005", "S110", "TID251"] +"core/mcp/session/test_client_session.py" = ["ARG001", "ARG005", "TID251"] +"core/mcp/client/test_sse_client.py" = ["ARG001", "TID251"] +"core/mcp/client/test_streamable_client.py" = ["ARG001", "ARG005", "S110", "TID251"] "core/mcp/session/test_base_session.py" = ["S110"] -"core/mcp/session/test_client_session.py" = ["ARG005"] "core/mcp/test_mcp_client.py" = ["ARG002"] "core/memory/test_token_buffer_memory.py" = ["ARG002"] -"core/ops/test_base_trace_instance.py" = ["ARG001"] -"core/ops/test_lookup_helpers.py" = ["ARG002"] -"core/ops/test_ops_trace_manager.py" = ["ARG001", "ARG002", "ARG005"] -"core/ops/test_trace_queue_manager.py" = ["ARG004"] -"core/ops/test_trace_session_metadata.py" = ["ARG001", "ARG005"] -"core/plugin/impl/test_agent_client.py" = ["ARG001"] -"core/plugin/impl/test_datasource_manager.py" = ["ARG001"] -"core/plugin/impl/test_oauth_handler.py" = ["ARG001"] -"core/plugin/impl/test_tool_manager.py" = ["ARG001"] -"core/plugin/impl/test_trigger_client.py" = ["ARG001"] -"core/plugin/test_endpoint_client.py" = ["ARG002"] -"core/plugin/test_model_runtime_adapter.py" = ["ARG002"] -"core/plugin/test_plugin_runtime.py" = ["ARG001", "ARG002", "TID251"] +"core/ops/test_base_trace_instance.py" = ["ARG001", "ARG004"] +"core/ops/test_ops_trace_manager.py" = ["ARG001", "ARG002", "ARG004", "ARG005"] +"core/ops/entities/test_trace_entity.py" = ["ARG001", "ARG005"] +"core/plugin/impl/test_agent.py" = ["ARG001"] +"core/plugin/impl/test_datasource.py" = ["ARG001"] +"core/plugin/impl/test_oauth.py" = ["ARG001"] +"core/plugin/impl/test_tool.py" = ["ARG001"] +"core/plugin/impl/test_trigger.py" = ["ARG001"] +"core/plugin/impl/test_endpoint.py" = ["ARG002"] +"core/plugin/impl/test_model_runtime.py" = ["ARG002"] +"core/plugin/impl/test_plugin.py" = ["ARG001", "ARG002", "TID251"] "core/prompt/test_advanced_prompt_transform.py" = ["ARG005"] "core/rag/datasource/keyword/jieba/test_jieba.py" = ["ARG001", "TID251"] "core/rag/datasource/keyword/jieba/test_jieba_keyword_table_handler.py" = ["ARG004"] -"core/rag/datasource/test_datasource_retrieval.py" = ["ARG001", "ARG002", "ARG005", "TID251"] -"core/rag/datasource/test_retrieval_attachment_access.py" = ["ARG005"] +"core/rag/datasource/test_retrieval_service.py" = ["ARG001", "ARG002", "ARG005", "TID251"] "core/rag/datasource/vdb/test_vector_factory.py" = ["ARG002"] -"core/rag/embedding/test_embedding_service.py" = ["ARG001"] +"core/rag/embedding/test_cached_embedding.py" = ["ARG001"] "core/rag/extractor/test_csv_extractor.py" = ["ARG001", "ARG005"] "core/rag/extractor/test_excel_extractor.py" = ["ARG002", "ARG005"] "core/rag/extractor/test_extract_processor.py" = ["ARG001", "ARG005"] "core/rag/extractor/test_helpers.py" = ["ARG002"] "core/rag/extractor/test_markdown_extractor.py" = ["ARG001"] "core/rag/extractor/test_word_extractor.py" = ["ARG001", "ARG002", "ARG005"] -"core/rag/extractor/unstructured/test_unstructured_extractors.py" = ["ARG001", "ARG005"] -"core/rag/extractor/watercrawl/test_watercrawl.py" = ["ARG001", "ARG005", "TID251"] -"core/rag/indexing/processor/test_paragraph_index_processor.py" = ["ARG002", "TID251"] -"core/rag/indexing/processor/test_qa_index_processor.py" = ["ARG001", "TID251"] -"core/rag/indexing/test_index_processor.py" = ["ARG005"] +"core/rag/extractor/unstructured/test_unstructured_doc_extractor.py" = ["ARG001", "ARG005"] +"core/rag/extractor/watercrawl/test_client.py" = ["ARG001", "ARG005", "TID251"] +"core/rag/index_processor/processor/test_paragraph_index_processor.py" = ["ARG002", "TID251"] +"core/rag/index_processor/processor/test_qa_index_processor.py" = ["ARG001", "TID251"] +"core/rag/index_processor/test_index_processor.py" = ["ARG005"] "core/rag/retrieval/test_dataset_retrieval.py" = ["ARG001", "ARG002", "ARG005", "TID251"] "core/rag/splitter/test_text_splitter.py" = ["ARG002", "ARG005"] "core/repositories/test_celery_workflow_execution_repository.py" = ["ARG002"] "core/repositories/test_celery_workflow_node_execution_repository.py" = ["ARG002"] -"core/repositories/test_human_input_form_repository_impl.py" = ["ARG001", "ARG005"] -"core/repositories/test_human_input_repository.py" = ["ANN401", "ARG001", "ARG005", "TID251"] +"core/repositories/test_human_input_repository.py" = ["ANN401", "ARG001", "ARG004", "ARG005", "TID251"] "core/repositories/test_sqlalchemy_workflow_node_execution_repository.py" = ["ANN401", "ARG002", "ARG005", "TID251"] -"core/repositories/test_workflow_node_execution_truncation.py" = ["TID251"] "core/schemas/test_resolver.py" = ["ARG005", "T201"] -"core/telemetry/test_facade.py" = ["ARG002", "ARG004"] -"core/telemetry/test_gateway_integration.py" = ["ARG002"] +"core/telemetry/test___init__.py" = ["ARG002", "ARG004"] +"core/telemetry/test_gateway.py" = ["ARG002"] "core/test_model_manager.py" = ["ARG001"] -"core/test_trigger_debug_event_selectors.py" = ["ARG002"] -"core/tools/test_base_tool.py" = ["ANN401", "ARG002", "TID251"] -"core/tools/test_builtin_tool_base.py" = ["ARG001", "ARG002", "TID251"] -"core/tools/test_builtin_tool_provider.py" = ["ARG001", "ARG005", "TID251"] -"core/tools/test_builtin_tools_extra.py" = ["ARG005"] -"core/tools/test_custom_tool.py" = ["ARG001", "ARG005", "TID251"] -"core/tools/test_dataset_retriever_tool.py" = ["ARG005"] -"core/tools/test_mcp_tool.py" = ["S110"] +"core/trigger/debug/test_event_selectors.py" = ["ARG002"] +"core/tools/__base/test_tool.py" = ["ANN401", "ARG002", "TID251"] +"core/tools/builtin_tool/test_tool.py" = ["ARG001", "ARG002", "TID251"] +"core/tools/builtin_tool/test_provider.py" = ["ARG001", "ARG005", "TID251"] +"core/tools/builtin_tool/providers/time/tools/test_localtime_to_timestamp.py" = ["ARG005"] +"core/tools/custom_tool/test_tool.py" = ["ARG001", "ARG005", "TID251"] +"core/tools/utils/test_dataset_retriever_tool.py" = ["ARG005"] +"core/tools/mcp_tool/test_tool.py" = ["S110", "TID251"] "core/tools/test_tool_engine.py" = ["ANN401", "ARG002", "ARG005", "TID251"] "core/tools/test_tool_file_manager.py" = ["ARG001"] "core/tools/test_tool_label_manager.py" = ["TID251"] "core/tools/test_tool_manager.py" = ["ANN401", "ARG001", "ARG005", "TID251"] -"core/tools/test_tool_provider_controller.py" = ["TID251"] +"core/tools/__base/test_tool_provider.py" = ["TID251"] "core/tools/utils/test_configuration.py" = ["ARG001", "ARG002", "TID251"] "core/tools/utils/test_message_transformer.py" = ["TID251"] -"core/tools/utils/test_misc_utils_extra.py" = ["ARG002"] +"core/tools/utils/dataset_retriever/test_dataset_multi_retriever_tool.py" = ["ARG002"] "core/tools/utils/test_model_invocation_utils.py" = ["ARG005", "TID251"] "core/tools/utils/test_parser.py" = ["TID251"] "core/tools/utils/test_web_reader_tool.py" = ["ARG001", "ARG002", "ARG005"] "core/tools/workflow_as_tool/test_provider.py" = ["TID251"] "core/tools/workflow_as_tool/test_tool.py" = ["ANN401", "ARG001", "ARG005", "TID251"] -"core/workflow/context/test_execution_context.py" = ["ANN401", "ARG002", "S110", "TID251"] -"core/workflow/context/test_flask_app_context.py" = ["ARG002"] -"core/workflow/generator/test_runner.py" = ["ARG001", "ARG002", "TID251"] -"core/workflow/generator/test_runner_missing.py" = ["ARG003"] +"context/test_execution_context.py" = ["ANN401", "ARG002", "S110", "TID251"] +"context/test_flask_app_context.py" = ["ARG002"] +"core/workflow/generator/test_runner.py" = ["ARG001", "ARG002", "ARG003", "TID251"] "core/workflow/generator/test_tool_catalogue.py" = ["ARG002"] -"core/workflow/graph_engine/layers/test_observability.py" = ["ARG002"] -"core/workflow/graph_engine/test_mock_config.py" = ["TID251"] -"core/workflow/graph_engine/test_mock_factory.py" = ["TID251"] -"core/workflow/graph_engine/test_mock_nodes.py" = ["ANN401", "S110", "TID251"] -"core/workflow/graph_engine/test_parallel_human_input_join_resume.py" = ["ARG002", "TID251"] -"core/workflow/graph_engine/test_table_runner.py" = ["ARG001", "TID251"] +"../helpers/workflow/mock_config.py" = ["TID251"] +"../helpers/workflow/mock_factory.py" = ["TID251"] +"../helpers/workflow/mock_nodes.py" = ["ANN401", "S110", "TID251"] +"../e2e_tests/contracts/core/workflow/graph_engine/test_parallel_human_input_join_resume.py" = ["ARG002", "TID251"] +"../helpers/workflow/table_runner.py" = ["ARG001", "TID251"] "core/workflow/nodes/agent_v2/test_agent_node.py" = ["ARG001", "ARG002", "ARG005"] "core/workflow/nodes/agent_v2/test_dify_tools_builder.py" = ["ANN401", "ARG001", "ARG002", "ARG005", "TID251"] "core/workflow/nodes/agent_v2/test_output_adapter.py" = ["ARG005"] "core/workflow/nodes/agent_v2/test_runtime_request_builder.py" = ["ARG002"] "core/workflow/nodes/agent_v2/test_validators.py" = ["ARG001"] -"core/workflow/nodes/http_request/test_http_request_node.py" = ["ANN401", "ARG002", "TID251"] +"../e2e_tests/contracts/core/workflow/nodes/http_request/test_http_request_node.py" = ["ANN401", "ARG002", "TID251"] "core/workflow/nodes/human_input/test_entities.py" = ["TID251"] -"core/workflow/nodes/human_input/test_human_input_form_filled_event.py" = ["TID251"] "core/workflow/nodes/knowledge_index/test_knowledge_index_node.py" = ["ARG002"] "core/workflow/nodes/knowledge_retrieval/test_knowledge_retrieval_node.py" = ["ARG002"] -"core/workflow/nodes/llm/test_node.py" = ["ARG002"] -"core/workflow/nodes/parameter_extractor/test_parameter_extractor_node.py" = ["TID251"] -"core/workflow/nodes/test_document_extractor_node.py" = ["ARG001"] -"core/workflow/nodes/tool/test_tool_node.py" = ["ANN401", "ARG002", "TID251"] -"core/workflow/nodes/webhook/test_webhook_file_conversion.py" = ["TID251"] -"core/workflow/nodes/webhook/test_webhook_node.py" = ["TID251"] -"core/workflow/test_form_input_serialization_compat.py" = ["ANN401", "TID251"] +"../e2e_tests/contracts/core/workflow/nodes/llm/test_node.py" = ["ARG002"] +"../e2e_tests/contracts/core/workflow/nodes/parameter_extractor/test_parameter_extractor_node.py" = ["TID251"] +"../e2e_tests/contracts/core/workflow/nodes/test_document_extractor_node.py" = ["ARG001"] +"../e2e_tests/contracts/core/workflow/nodes/tool/test_tool_node.py" = ["ANN401", "ARG002", "TID251"] +"core/workflow/nodes/trigger_webhook/test_node.py" = ["TID251"] +"core/workflow/nodes/human_input/test_pause_reason.py" = ["ANN401", "TID251"] "core/workflow/test_human_input_adapter.py" = ["ARG005"] "core/workflow/test_node_factory.py" = ["ARG002"] "core/workflow/test_workflow_entry.py" = ["ARG001"] "enterprise/telemetry/test_enterprise_trace.py" = ["ARG002", "TID251"] -"enterprise/telemetry/test_gateway.py" = ["ARG002"] "events/event_handlers/test_delete_tool_parameters_cache_when_sync_draft_workflow.py" = ["ARG005"] "extensions/logstore/test_sql_escape.py" = ["ARG001", "ARG002"] "extensions/otel/decorators/handlers/test_generate_handler.py" = ["ARG001", "ARG002"] "extensions/otel/decorators/handlers/test_workflow_app_runner_handler.py" = ["ARG001"] -"extensions/otel/test_retrieval_tracing.py" = ["ARG001"] -"factories/test_build_from_mapping.py" = ["ARG001"] -"factories/test_file_factory.py" = ["ARG001"] +"core/workflow/nodes/knowledge_retrieval/test_retrieval.py" = ["ARG001"] +"factories/file_factory/test_builders.py" = ["ARG001"] +"factories/file_factory/test_remote.py" = ["ARG001"] "fields/test_file_fields.py" = ["ARG005"] -"libs/_human_input/support.py" = ["TID251"] -"libs/broadcast_channel/redis/test_channel_unit_tests.py" = ["ARG002", "ARG005"] -"libs/broadcast_channel/redis/test_streams_channel_unit_tests.py" = ["ARG001", "ARG002", "TID251"] -"libs/test_cron_compatibility.py" = ["S110"] +"../helpers/unit/libs/_human_input/support.py" = ["TID251"] +"libs/broadcast_channel/redis/test_pubsub_channel.py" = ["ARG002", "ARG005"] +"libs/broadcast_channel/redis/test_streams_channel.py" = ["ARG001", "ARG002", "TID251"] +"libs/test_schedule_utils.py" = ["S110"] "libs/test_email_i18n.py" = ["ANN401", "TID251"] -"libs/test_schedule_utils_enhanced.py" = ["S110"] -"models/test_dataset_models.py" = ["ARG005"] -"services/agent/test_agent_observability_service.py" = ["ARG002", "ARG005"] -"services/agent/test_agent_services.py" = ["ARG001", "ARG002", "ARG003", "ARG005"] +"models/test_dataset.py" = ["ARG005"] +"services/agent/test_observability_service.py" = ["ARG002", "ARG005"] +"services/agent/test_roster_service.py" = ["ARG001", "ARG002", "ARG003", "ARG005"] "services/agent/test_prompt_mentions.py" = ["ARG005"] -"services/controller_api.py" = ["ARG002"] +"../helpers/unit/services/controller_api.py" = ["ARG002"] "services/data_migration/test_import_service.py" = ["ARG002", "ARG005"] -"services/dataset_service_test_helpers.py" = ["TID251"] +"../helpers/unit/services/dataset_service_test_helpers.py" = ["TID251"] "services/enterprise/test_account_deletion_sync.py" = ["ARG001"] -"services/enterprise/test_traceparent_propagation.py" = ["ARG002"] -"services/hit_service.py" = ["TID251"] +"services/enterprise/test_base.py" = ["ARG002"] +"../helpers/unit/services/hit_service.py" = ["TID251"] "services/plugin/test_plugin_parameter_service.py" = ["ARG002"] -"services/rag_pipeline/pipeline_template/test_built_in_retrieval.py" = ["ARG001"] +"services/rag_pipeline/pipeline_template/built_in/test_built_in_retrieval.py" = ["ARG001"] "services/rag_pipeline/test_rag_pipeline_dsl_service.py" = ["ARG001", "ARG005", "T201", "TID251"] -"services/rag_pipeline/test_rag_pipeline_service.py" = ["ARG001", "ARG005"] -"services/rag_pipeline/test_rag_pipeline_task_proxy.py" = ["ARG001", "ARG005"] +"services/rag_pipeline/test_rag_pipeline.py" = ["ARG001", "ARG004", "ARG005"] +"services/rag_pipeline/test_rag_pipeline_task_proxy.py" = ["ARG001", "ARG002", "ARG004", "ARG005"] "services/rag_pipeline/test_rag_pipeline_transform_service.py" = ["ARG001"] "services/retention/workflow_run/test_archive_download_preparation.py" = ["ARG002"] "services/retention/workflow_run/test_archive_log_service.py" = ["ARG001", "ARG002"] "services/retention/workflow_run/test_bundle_archive_maintenance.py" = ["TID251"] "services/retention/workflow_run/test_restore_archived_workflow_run.py" = ["ARG002"] -"services/test_annotation_service.py" = ["ANN401", "TID251"] +"services/test_annotation_service.py" = ["ANN401", "ARG002", "TID251"] "services/test_api_token_service.py" = ["ARG002"] -"services/test_app_generate_service.py" = ["ARG001", "ARG002", "ARG004"] -"services/test_app_generate_service_streaming_integration.py" = ["ARG002", "TID251"] +"services/test_app_generate_service.py" = ["ARG001", "ARG002", "ARG004", "TID251"] "services/test_audio_service.py" = ["ARG002", "TID251"] "services/test_billing_service.py" = ["ARG001", "ARG002"] -"services/test_clear_free_plan_expired_workflow_run_logs.py" = ["ANN401", "ARG002", "ARG005", "TID251"] +"services/retention/workflow_run/test_clear_free_plan_expired_workflow_run_logs.py" = ["ANN401", "ARG002", "ARG005", "TID251"] "services/test_clear_free_plan_tenant_expired_logs.py" = ["ARG002", "ARG003"] -"services/test_dataset_service_document.py" = ["ARG002"] -"services/test_dataset_service_segment.py" = ["ARG002"] -"services/test_datasource_provider_service.py" = ["ARG002"] -"services/test_external_dataset_service.py" = ["ARG002", "TID251"] +"services/knowledge/test_dataset_service.py" = ["ARG002"] +"services/data_source/test_provider_service.py" = ["ARG002"] +"services/knowledge/external/test_service.py" = ["ARG002", "TID251"] "services/test_human_input_delivery_test_service.py" = ["ARG005"] "services/test_message_service.py" = ["ARG002"] -"services/test_messages_clean_service.py" = ["TID251"] -"services/test_model_provider_service.py" = ["ANN401", "TID251"] +"services/retention/conversation/test_messages_clean_service.py" = ["TID251"] +"services/model_provider/test_service.py" = ["ANN401", "TID251"] "services/test_operation_service.py" = ["TID251"] -"services/test_rag_pipeline_task_proxy.py" = ["ARG002"] -"services/test_schedule_service.py" = ["ANN401", "TID251"] +"services/trigger/test_schedule_service.py" = ["ANN401", "TID251"] "services/test_snippet_service.py" = ["ARG001", "ARG002"] -"services/test_summary_index_service.py" = ["ARG001"] +"services/knowledge/summaries/test_adapters.py" = ["ARG001"] "services/test_variable_truncator.py" = ["ARG002", "TID251"] -"services/test_webhook_service_additional.py" = ["ANN401", "ARG002", "TID251"] -"services/test_website_service.py" = ["TID251"] +"services/trigger/test_webhook_service.py" = ["ANN401", "ARG002", "TID251"] +"services/data_source/test_website_service.py" = ["TID251"] "services/test_workflow_run_service.py" = ["ANN401", "ARG002", "TID251"] "services/test_workflow_service.py" = ["ANN401", "ARG002", "TID251"] "services/tools/test_builtin_tools_manage_service.py" = ["ARG001", "ARG002"] "services/tools/test_tools_manage_service.py" = ["ARG002"] "services/workflow/test_node_output_inspector_service.py" = ["TID251"] -"services/workflow/test_workflow_converter_additional.py" = ["ANN401", "ARG001", "ARG005", "TID251"] -"services/workflow/test_workflow_event_snapshot_service.py" = ["ANN401", "ARG002", "ARG005", "TID251"] +"services/workflow/test_workflow_converter.py" = ["ANN401", "ARG001", "ARG005", "TID251"] +"services/test_workflow_event_snapshot_service.py" = ["ANN401", "ARG002", "ARG005", "TID251"] "services/workflow/test_workflow_event_snapshot_service_additional.py" = ["ANN401", "ARG002", "ARG005", "TID251"] "tasks/test_clean_dataset_task.py" = ["ARG002"] "tasks/test_clean_document_task.py" = ["ARG002"] -"tasks/test_dataset_indexing_task.py" = ["ARG001", "ARG002"] +"tasks/test_document_indexing_task.py" = ["ARG001", "ARG002"] "tasks/test_document_indexing_sync_task.py" = ["ARG002"] "tasks/test_human_input_timeout_tasks.py" = ["ARG001", "ARG002", "ARG005", "TID251"] "tasks/test_initialize_created_app_rbac_access_task.py" = ["ARG005"] "tasks/test_ops_trace_task.py" = ["ARG004"] "tasks/test_process_tenant_plugin_autoupgrade_check_task.py" = ["ARG001"] "tasks/test_remove_app_and_related_data_task.py" = ["ARG002"] -"tasks/test_workflow_execute_task.py" = ["ARG005"] -"test_pytest_dify.py" = ["ARG001"] -"tools/test_mcp_tool.py" = ["TID251"] +"tasks/app_generate/test_workflow_execute_task.py" = ["ARG005"] +"../e2e_tests/contracts/test_pytest_dify.py" = ["ARG001"] [lint.flake8-tidy-imports.banned-api."flask_restx.reqparse"] msg = "Use Pydantic payload/query models instead of reqparse." diff --git a/api/tests/unit_tests/clients/agent_backend/test_factory.py b/api/tests/unit_tests/clients/agent_backend/test_factory.py index f289b73e013..516dbbbdbf3 100644 --- a/api/tests/unit_tests/clients/agent_backend/test_factory.py +++ b/api/tests/unit_tests/clients/agent_backend/test_factory.py @@ -9,7 +9,7 @@ from clients.agent_backend.factory import create_agent_backend_client, create_ag from configs import dify_config from services.agent import home_snapshot_service, workspace_service from services.app import agent_sandbox_file_gateway -from tests.unit_tests.config_override import apply_config_overrides +from tests.helpers.unit.config_override import apply_config_overrides @pytest.mark.parametrize( diff --git a/api/tests/unit_tests/commands/test_data_migration_cli_registration.py b/api/tests/unit_tests/commands/test___init__.py similarity index 100% rename from api/tests/unit_tests/commands/test_data_migration_cli_registration.py rename to api/tests/unit_tests/commands/test___init__.py diff --git a/api/tests/unit_tests/commands/test_account_commands.py b/api/tests/unit_tests/commands/test_account.py similarity index 100% rename from api/tests/unit_tests/commands/test_account_commands.py rename to api/tests/unit_tests/commands/test_account.py diff --git a/api/tests/unit_tests/commands/test_archive_workflow_runs.py b/api/tests/unit_tests/commands/test_archive_workflow_runs.py deleted file mode 100644 index c92cd8657da..00000000000 --- a/api/tests/unit_tests/commands/test_archive_workflow_runs.py +++ /dev/null @@ -1,631 +0,0 @@ -"""Tests for workflow-run archive command database boundaries. - -Planning deliberately creates a fresh session for every tenant prefix and for -every database retry. SQLite-backed tests keep the query, filtering, counting, -and session lifecycle real while clocks, billing lookup, and command failures -remain narrow external-boundary substitutions. -""" - -import datetime -import logging -from dataclasses import dataclass -from types import SimpleNamespace -from unittest.mock import MagicMock, Mock -from uuid import uuid4 - -import click -import pytest -from click.testing import CliRunner -from sqlalchemy import Engine, event, text -from sqlalchemy.exc import OperationalError -from sqlalchemy.orm import Session, scoped_session, sessionmaker - -from commands import retention -from graphon.enums import WorkflowExecutionStatus, WorkflowNodeExecutionStatus -from models.base import TypeBase -from models.enums import CreatorUserRole, WorkflowRunTriggeredFrom -from models.workflow import ( - WorkflowNodeExecutionModel, - WorkflowNodeExecutionTriggeredFrom, - WorkflowRun, - WorkflowType, -) -from services.retention.workflow_run import bundle_archive_maintenance -from services.retention.workflow_run.bundle_archive_maintenance import ( - BundleOperationResult, - BundleOperationSummary, -) - -_CURSOR_0 = "00000000-0000-0000-0000-000000000000" -_CURSOR_1 = "00000000-0000-0000-0000-000000000001" -_CURSOR_2 = "00000000-0000-0000-0000-000000000002" - - -def _db_disconnect_error() -> OperationalError: - return OperationalError( - "select 1", - {}, - RuntimeError("server closed the connection unexpectedly"), - connection_invalidated=True, - ) - - -@dataclass(frozen=True) -class ArchiveDatabase: - """Creates archive candidates and their node executions in SQLite.""" - - session_maker: sessionmaker[Session] - end_before: datetime.datetime - - def add_run( - self, - tenant_id: str, - *, - created_at: datetime.datetime | None = None, - status: WorkflowExecutionStatus = WorkflowExecutionStatus.SUCCEEDED, - run_type: WorkflowType = WorkflowType.WORKFLOW, - ) -> str: - run_id = str(uuid4()) - with self.session_maker.begin() as session: - session.add( - WorkflowRun( - id=run_id, - tenant_id=tenant_id, - app_id=str(uuid4()), - workflow_id=str(uuid4()), - type=run_type, - triggered_from=WorkflowRunTriggeredFrom.APP_RUN, - version="1", - graph="{}", - inputs="{}", - status=status, - outputs="{}", - error=None, - elapsed_time=0, - total_tokens=0, - total_steps=1, - created_by_role=CreatorUserRole.ACCOUNT, - created_by=str(uuid4()), - created_at=created_at or self.end_before - datetime.timedelta(days=1), - ) - ) - return run_id - - def add_node(self, run_id: str, tenant_id: str, *, index: int) -> None: - with self.session_maker.begin() as session: - session.add( - WorkflowNodeExecutionModel( - id=str(uuid4()), - tenant_id=tenant_id, - app_id=str(uuid4()), - workflow_id=str(uuid4()), - triggered_from=WorkflowNodeExecutionTriggeredFrom.WORKFLOW_RUN, - workflow_run_id=run_id, - index=index, - predecessor_node_id=None, - node_execution_id=None, - node_id=f"node-{index}", - node_type="start", - title="Start", - inputs="{}", - process_data="{}", - outputs="{}", - status=WorkflowNodeExecutionStatus.SUCCEEDED, - error=None, - elapsed_time=0, - execution_metadata="{}", - created_by_role=CreatorUserRole.ACCOUNT, - created_by=str(uuid4()), - ) - ) - - -def _delete_summary( - *, - processed: int, - succeeded: int = 0, - failed: int = 0, - next_catalog_id: str | None = None, - preview_next_catalog_id: str | None = None, - results: list[BundleOperationResult] | None = None, -) -> BundleOperationSummary: - return BundleOperationSummary( - operation="delete", - bundles_processed=processed, - bundles_succeeded=succeeded, - bundles_failed=failed, - next_catalog_id=next_catalog_id, - preview_next_catalog_id=preview_next_catalog_id, - results=results or [], - ) - - -def _patch_bundle_deleter( - monkeypatch: pytest.MonkeyPatch, - summaries: list[BundleOperationSummary], -) -> MagicMock: - deleter = MagicMock() - deleter.delete_batch.side_effect = summaries - monkeypatch.setattr( - bundle_archive_maintenance, - "WorkflowRunBundleArchiveMaintenance", - MagicMock(return_value=deleter), - ) - return deleter - - -@pytest.mark.parametrize( - "command", - [retention.restore_workflow_runs, retention.delete_archived_workflow_runs], -) -def test_v2_archive_maintenance_rejects_explicitly_empty_tenant_ids(command): - result = CliRunner().invoke( - command, - ["--tenant-ids", "", "--target-month", "2025-03"], - ) - - assert result.exit_code == 2 - assert "tenant-ids must not be empty" in result.output - - -def test_archive_tenant_id_parser_keeps_omitted_scope_unset(): - assert retention._parse_comma_separated_ids(None, param_name="tenant-ids") is None - - -@pytest.fixture -def archive_db(sqlite_engine: Engine) -> ArchiveDatabase: - """Create only the workflow tables used by archive planning.""" - - TypeBase.metadata.create_all( - sqlite_engine, - tables=[WorkflowRun.__table__, WorkflowNodeExecutionModel.__table__], - ) - return ArchiveDatabase( - session_maker=sessionmaker(bind=sqlite_engine, expire_on_commit=False), - end_before=datetime.datetime(2025, 4, 1, tzinfo=datetime.UTC), - ) - - -def _tenant_id(prefix: str, suffix: int) -> str: - return f"{prefix}{suffix:07x}-0000-0000-0000-000000000000" - - -def test_resolve_archive_tenant_ids_from_plan_uses_fresh_real_sessions( - archive_db: ArchiveDatabase, monkeypatch: pytest.MonkeyPatch -) -> None: - paid_a = _tenant_id("a", 1) - free_a = _tenant_id("a", 2) - paid_b = _tenant_id("b", 1) - free_b = _tenant_id("b", 2) - for tenant_id in (paid_a, free_a, paid_b, free_b): - archive_db.add_run(tenant_id) - - # These decoys verify the real candidate query's time, status, and type filters. - archive_db.add_run( - _tenant_id("a", 3), - created_at=archive_db.end_before + datetime.timedelta(seconds=1), - ) - archive_db.add_run(_tenant_id("a", 4), status=WorkflowExecutionStatus.RUNNING) - archive_db.add_run(_tenant_id("a", 5), run_type=WorkflowType.CHAT) - - opened_sessions: list[Session] = [] - - def record_session(session: Session, _transaction: object, _connection: object) -> None: - opened_sessions.append(session) - - event.listen(archive_db.session_maker.class_, "after_begin", record_session) - monkeypatch.setattr( - retention, - "_filter_paid_workflow_archive_tenant_ids", - lambda tenant_ids: ([paid_a, paid_b], sorted(set(tenant_ids) - {paid_a, paid_b})), - ) - try: - tenant_plan = retention._resolve_archive_tenant_ids_from_plan( - session_maker=archive_db.session_maker, - tenant_ids=None, - tenant_prefixes=["a", "b"], - start_from=None, - end_before=archive_db.end_before, - ) - finally: - event.remove(archive_db.session_maker.class_, "after_begin", record_session) - - assert tenant_plan == { - "archive_tenant_ids": [paid_a, paid_b], - "paid_tenant_ids": [paid_a, paid_b], - "unpaid_tenant_ids": [free_a, free_b], - } - assert len(opened_sessions) == 2 - assert opened_sessions[0] is not opened_sessions[1] - - -def test_safe_remove_scoped_session_recovers_from_real_closed_connection( - sqlite_engine: Engine, - monkeypatch: pytest.MonkeyPatch, - caplog: pytest.LogCaptureFixture, -) -> None: - maker = sessionmaker(bind=sqlite_engine) - registry = scoped_session(maker) - registry().execute(text("select 1")) - dbapi_connection = registry().connection().connection.dbapi_connection - assert dbapi_connection is not None - dbapi_connection.close() - monkeypatch.setattr(retention, "db", SimpleNamespace(session=registry, engine=sqlite_engine)) - - with caplog.at_level(logging.WARNING, logger="commands.retention"): - retention._safe_remove_scoped_session("archive workflow run command") - - assert not registry.registry.has() - assert any("Ignoring DB scoped-session cleanup error" in message for message in caplog.messages) - - -def test_archive_command_db_retry_retries_retryable_db_disconnect(monkeypatch: pytest.MonkeyPatch) -> None: - attempts = iter([_db_disconnect_error(), "ok"]) - sleep = Mock() - monkeypatch.setattr("services.retention.workflow_run.db_retry.time.sleep", sleep) - - def operation() -> str: - result = next(attempts) - if isinstance(result, Exception): - raise result - return result - - assert retention._run_archive_command_db_retry("archive plan", operation) == "ok" - sleep.assert_called_once_with(1.0) - - -def test_archive_plan_prefix_stats_retries_with_fresh_session_and_real_counts( - archive_db: ArchiveDatabase, sqlite_engine: Engine, monkeypatch: pytest.MonkeyPatch -) -> None: - tenant_id = _tenant_id("a", 1) - run_ids = [archive_db.add_run(tenant_id) for _ in range(7)] - for index in range(9): - archive_db.add_node(run_ids[index % len(run_ids)], tenant_id, index=index) - - # Decoys outside the selected prefix and archive window must not affect counts. - decoy_run_id = archive_db.add_run(_tenant_id("b", 1)) - archive_db.add_node(decoy_run_id, _tenant_id("b", 1), index=99) - archive_db.add_run( - tenant_id, - created_at=archive_db.end_before + datetime.timedelta(seconds=1), - ) - - fail_next_query = True - - def disconnect_once( - _connection: object, - _cursor: object, - _statement: str, - _parameters: object, - _context: object, - _executemany: bool, - ) -> None: - nonlocal fail_next_query - if fail_next_query: - fail_next_query = False - raise _db_disconnect_error() - - opened_sessions: list[Session] = [] - - def record_session(session: Session, _transaction: object, _connection: object) -> None: - opened_sessions.append(session) - - sleep = Mock() - monkeypatch.setattr("services.retention.workflow_run.db_retry.time.sleep", sleep) - event.listen(sqlite_engine, "before_cursor_execute", disconnect_once) - event.listen(archive_db.session_maker.class_, "after_begin", record_session) - try: - stats = retention._get_archive_plan_prefix_stats( - archive_db.session_maker, - "a", - start_from=None, - end_before=archive_db.end_before, - ) - finally: - event.remove(sqlite_engine, "before_cursor_execute", disconnect_once) - event.remove(archive_db.session_maker.class_, "after_begin", record_session) - - assert stats == { - "tenant_ids": [tenant_id], - "workflow_runs": 7, - "workflow_node_executions": 9, - } - assert len(opened_sessions) == 2 - assert opened_sessions[0] is not opened_sessions[1] - sleep.assert_called_once_with(1.0) - - -def test_archive_workflow_runs_raises_click_exception_when_tenant_plan_fails( - sqlite_engine: Engine, monkeypatch: pytest.MonkeyPatch -) -> None: - registry = scoped_session(sessionmaker(bind=sqlite_engine)) - monkeypatch.setattr(retention, "db", SimpleNamespace(engine=sqlite_engine, session=registry)) - monkeypatch.setattr( - retention, - "_resolve_archive_tenant_ids_from_plan", - Mock(side_effect=RuntimeError("tenant plan failed")), - ) - - with pytest.raises(click.ClickException, match="Failed to resolve workflow archive tenant plan"): - retention.archive_workflow_runs.callback( - tenant_ids="tenant-1", - tenant_prefixes=None, - before_days=90, - from_days_ago=None, - to_days_ago=None, - start_from=None, - end_before=None, - batch_size=10000, - workers=1, - run_shard_index=None, - run_shard_total=None, - limit=None, - dry_run=True, - delete_after_archive=False, - ) - - -def test_delete_archived_workflow_runs_keeps_single_page_behavior_without_all_pages(monkeypatch: pytest.MonkeyPatch): - deleter = _patch_bundle_deleter( - monkeypatch, - [_delete_summary(processed=2, succeeded=2, next_catalog_id=_CURSOR_1)], - ) - - result = CliRunner().invoke( - retention.delete_archived_workflow_runs, - ["--target-month", "2025-03", "--limit", "2"], - ) - - assert result.exit_code == 0 - deleter.delete_batch.assert_called_once() - assert deleter.delete_batch.call_args.kwargs["target_year"] == 2025 - assert deleter.delete_batch.call_args.kwargs["target_month"] == 3 - assert deleter.delete_batch.call_args.kwargs["after_catalog_id"] is None - assert deleter.delete_batch.call_args.kwargs["limit"] == 2 - - -def test_delete_archived_workflow_runs_all_pages_continues_until_empty_page(monkeypatch: pytest.MonkeyPatch): - deleter = _patch_bundle_deleter( - monkeypatch, - [ - _delete_summary(processed=2, succeeded=2, next_catalog_id=_CURSOR_1), - _delete_summary(processed=1, succeeded=1, next_catalog_id=_CURSOR_2), - _delete_summary(processed=0), - ], - ) - - result = CliRunner().invoke( - retention.delete_archived_workflow_runs, - ["--target-month", "2025-03", "--all-pages", "--limit", "2"], - ) - - assert result.exit_code == 0 - assert [call.kwargs["after_catalog_id"] for call in deleter.delete_batch.call_args_list] == [ - None, - _CURSOR_1, - _CURSOR_2, - ] - - -def test_delete_archived_workflow_runs_all_pages_fetches_empty_page_after_exact_full_page( - monkeypatch: pytest.MonkeyPatch, -): - deleter = _patch_bundle_deleter( - monkeypatch, - [ - _delete_summary(processed=2, succeeded=2, next_catalog_id=_CURSOR_1), - _delete_summary(processed=0), - ], - ) - - result = CliRunner().invoke( - retention.delete_archived_workflow_runs, - ["--target-month", "2025-03", "--all-pages", "--limit", "2"], - ) - - assert result.exit_code == 0 - assert deleter.delete_batch.call_count == 2 - assert deleter.delete_batch.call_args_list[1].kwargs["after_catalog_id"] == _CURSOR_1 - - -def test_delete_archived_workflow_runs_all_pages_stops_at_first_failed_page(monkeypatch: pytest.MonkeyPatch): - failed_result = BundleOperationResult( - catalog_id=_CURSOR_2, - bundle_id="bundle-failed", - tenant_id="tenant-1", - object_prefix="workflow-runs/v2/tenant-1/2025/03/00-of-16/bundle-failed", - error="archive checksum mismatch", - ) - deleter = _patch_bundle_deleter( - monkeypatch, - [ - _delete_summary(processed=1, succeeded=1, next_catalog_id=_CURSOR_1), - _delete_summary(processed=1, failed=1, results=[failed_result]), - _delete_summary(processed=0), - ], - ) - - result = CliRunner().invoke( - retention.delete_archived_workflow_runs, - ["--target-month", "2025-03", "--all-pages"], - ) - - assert result.exit_code == 1 - assert deleter.delete_batch.call_count == 2 - assert "target_month=2025-03" in result.output - assert f"failed_catalog_id={_CURSOR_2}" in result.output - assert f"resume_after_catalog_id={_CURSOR_1}" in result.output - - -def test_delete_archived_workflow_runs_all_pages_fails_when_cursor_does_not_advance(monkeypatch: pytest.MonkeyPatch): - deleter = _patch_bundle_deleter( - monkeypatch, - [_delete_summary(processed=1, succeeded=1, next_catalog_id=None)], - ) - - result = CliRunner().invoke( - retention.delete_archived_workflow_runs, - ["--target-month", "2025-03", "--all-pages"], - ) - - assert result.exit_code == 1 - deleter.delete_batch.assert_called_once() - assert "cursor did not advance" in result.output.lower() - - -def test_delete_archived_workflow_runs_all_pages_uses_preview_cursor_for_dry_run(monkeypatch: pytest.MonkeyPatch): - deleter = _patch_bundle_deleter( - monkeypatch, - [ - _delete_summary(processed=1, succeeded=1, preview_next_catalog_id=_CURSOR_1), - _delete_summary(processed=0), - ], - ) - - result = CliRunner().invoke( - retention.delete_archived_workflow_runs, - ["--target-month", "2025-03", "--all-pages", "--dry-run"], - ) - - assert result.exit_code == 0 - assert [call.kwargs["after_catalog_id"] for call in deleter.delete_batch.call_args_list] == [ - None, - _CURSOR_1, - ] - - -def test_delete_archived_workflow_runs_dry_run_failure_separates_preview_and_destructive_cursors( - monkeypatch: pytest.MonkeyPatch, -): - failed_result = BundleOperationResult( - catalog_id=_CURSOR_2, - bundle_id="bundle-failed", - tenant_id="tenant-1", - object_prefix="workflow-runs/v2/tenant-1/2025/03/00-of-16/bundle-failed", - error="archive checksum mismatch", - ) - deleter = _patch_bundle_deleter( - monkeypatch, - [ - _delete_summary( - processed=1, - succeeded=1, - preview_next_catalog_id=_CURSOR_1, - ), - _delete_summary( - processed=1, - failed=1, - results=[failed_result], - ), - ], - ) - - result = CliRunner().invoke( - retention.delete_archived_workflow_runs, - [ - "--target-month", - "2025-03", - "--after-catalog-id", - _CURSOR_0, - "--all-pages", - "--dry-run", - ], - ) - - assert result.exit_code == 1 - assert deleter.delete_batch.call_count == 2 - assert f"failed_catalog_id={_CURSOR_2}" in result.output - assert f"preview_after_catalog_id={_CURSOR_1}" in result.output - assert f"destructive_retry_after_catalog_id={_CURSOR_0}" in result.output - - -def test_delete_archived_workflow_runs_all_pages_starts_after_explicit_cursor(monkeypatch: pytest.MonkeyPatch): - deleter = _patch_bundle_deleter(monkeypatch, [_delete_summary(processed=0)]) - - result = CliRunner().invoke( - retention.delete_archived_workflow_runs, - [ - "--target-month", - "2025-03", - "--after-catalog-id", - _CURSOR_0, - "--all-pages", - ], - ) - - assert result.exit_code == 0 - deleter.delete_batch.assert_called_once() - assert deleter.delete_batch.call_args.kwargs["after_catalog_id"] == _CURSOR_0 - - -@pytest.mark.parametrize( - "shard_args", - [ - ["--run-shard-index", "0"], - ["--run-shard-total", "16"], - ["--run-shard-index", "16", "--run-shard-total", "16"], - ["--run-shard-index", "-1", "--run-shard-total", "16"], - ["--run-shard-index", "0", "--run-shard-total", "0"], - ["--run-shard-index", "0", "--run-shard-total", "17"], - ], -) -def test_delete_archived_workflow_runs_rejects_invalid_run_shard_options(monkeypatch, shard_args): - deleter = _patch_bundle_deleter(monkeypatch, [_delete_summary(processed=0)]) - - result = CliRunner().invoke( - retention.delete_archived_workflow_runs, - ["--target-month", "2025-03", *shard_args], - ) - - assert result.exit_code == 2 - deleter.delete_batch.assert_not_called() - - -def test_delete_archived_workflow_runs_passes_formatted_run_shard_to_service(monkeypatch: pytest.MonkeyPatch): - deleter = _patch_bundle_deleter(monkeypatch, [_delete_summary(processed=0)]) - - result = CliRunner().invoke( - retention.delete_archived_workflow_runs, - [ - "--target-month", - "2025-03", - "--tenant-ids", - "tenant-1", - "--run-shard-index", - "3", - "--run-shard-total", - "16", - ], - ) - - assert result.exit_code == 0 - deleter.validate_catalog_shards.assert_called_once_with( - target_year=2025, - target_month=3, - shard_total=16, - tenant_ids=["tenant-1"], - ) - deleter.delete_batch.assert_called_once() - assert deleter.delete_batch.call_args.kwargs["shard"] == "03-of-16" - - -def test_delete_archived_workflow_runs_rejects_mixed_catalog_shards_before_delete(monkeypatch: pytest.MonkeyPatch): - deleter = _patch_bundle_deleter(monkeypatch, [_delete_summary(processed=0)]) - deleter.validate_catalog_shards.side_effect = ValueError("unexpected shards: 00-of-01") - - result = CliRunner().invoke( - retention.delete_archived_workflow_runs, - [ - "--target-month", - "2025-03", - "--run-shard-index", - "3", - "--run-shard-total", - "16", - ], - ) - - assert result.exit_code == 1 - assert "shard preflight failed" in result.output.lower() - assert "00-of-01" in result.output - deleter.delete_batch.assert_not_called() diff --git a/api/tests/unit_tests/commands/test_clean_expired_messages.py b/api/tests/unit_tests/commands/test_clean_expired_messages.py deleted file mode 100644 index b89b637bd6a..00000000000 --- a/api/tests/unit_tests/commands/test_clean_expired_messages.py +++ /dev/null @@ -1,190 +0,0 @@ -import datetime -import re -from unittest.mock import MagicMock, patch - -import click -import pytest - -from commands import clean_expired_messages - - -def _mock_service() -> MagicMock: - service = MagicMock() - service.run.return_value = { - "batches": 1, - "total_messages": 10, - "filtered_messages": 5, - "total_deleted": 5, - } - return service - - -def _invoke_clean_expired_messages(**kwargs: object) -> None: - callback = clean_expired_messages.callback - assert callback is not None - callback(**kwargs) - - -def test_absolute_mode_calls_from_time_range() -> None: - policy = object() - service = _mock_service() - start_from = datetime.datetime(2024, 1, 1, 0, 0, 0) - end_before = datetime.datetime(2024, 2, 1, 0, 0, 0) - - with ( - patch("commands.retention.create_message_clean_policy", return_value=policy), - patch("commands.retention.MessagesCleanService.from_time_range", return_value=service) as mock_from_time_range, - patch("commands.retention.MessagesCleanService.from_days") as mock_from_days, - ): - _invoke_clean_expired_messages( - batch_size=200, - graceful_period=21, - start_from=start_from, - end_before=end_before, - from_days_ago=None, - before_days=None, - dry_run=True, - ) - - mock_from_time_range.assert_called_once_with( - policy=policy, - start_from=start_from, - end_before=end_before, - batch_size=200, - dry_run=True, - task_label="custom", - ) - mock_from_days.assert_not_called() - - -def test_relative_mode_before_days_only_calls_from_days() -> None: - policy = object() - service = _mock_service() - - with ( - patch("commands.retention.create_message_clean_policy", return_value=policy), - patch("commands.retention.MessagesCleanService.from_days", return_value=service) as mock_from_days, - patch("commands.retention.MessagesCleanService.from_time_range") as mock_from_time_range, - ): - _invoke_clean_expired_messages( - batch_size=500, - graceful_period=14, - start_from=None, - end_before=None, - from_days_ago=None, - before_days=30, - dry_run=False, - ) - - mock_from_days.assert_called_once_with( - policy=policy, - days=30, - batch_size=500, - dry_run=False, - task_label="before-30", - ) - mock_from_time_range.assert_not_called() - - -def test_relative_mode_with_from_days_ago_calls_from_time_range() -> None: - policy = object() - service = _mock_service() - fixed_now = datetime.datetime(2024, 8, 20, 12, 0, 0) - - with ( - patch("commands.retention.create_message_clean_policy", return_value=policy), - patch("commands.retention.MessagesCleanService.from_time_range", return_value=service) as mock_from_time_range, - patch("commands.retention.MessagesCleanService.from_days") as mock_from_days, - patch("commands.retention.naive_utc_now", return_value=fixed_now), - ): - _invoke_clean_expired_messages( - batch_size=1000, - graceful_period=21, - start_from=None, - end_before=None, - from_days_ago=60, - before_days=30, - dry_run=False, - ) - - mock_from_time_range.assert_called_once_with( - policy=policy, - start_from=fixed_now - datetime.timedelta(days=60), - end_before=fixed_now - datetime.timedelta(days=30), - batch_size=1000, - dry_run=False, - task_label="60to30", - ) - mock_from_days.assert_not_called() - - -@pytest.mark.parametrize( - ("kwargs", "message"), - [ - ( - { - "start_from": datetime.datetime(2024, 1, 1), - "end_before": datetime.datetime(2024, 2, 1), - "from_days_ago": None, - "before_days": 30, - }, - "mutually exclusive", - ), - ( - { - "start_from": datetime.datetime(2024, 1, 1), - "end_before": None, - "from_days_ago": None, - "before_days": None, - }, - "Both --start-from and --end-before are required", - ), - ( - { - "start_from": None, - "end_before": None, - "from_days_ago": 10, - "before_days": None, - }, - "--from-days-ago must be used together with --before-days", - ), - ( - { - "start_from": None, - "end_before": None, - "from_days_ago": None, - "before_days": -1, - }, - "--before-days must be >= 0", - ), - ( - { - "start_from": None, - "end_before": None, - "from_days_ago": 30, - "before_days": 30, - }, - "--from-days-ago must be greater than --before-days", - ), - ( - { - "start_from": None, - "end_before": None, - "from_days_ago": None, - "before_days": None, - }, - "You must provide either (--start-from,--end-before) or (--before-days [--from-days-ago])", - ), - ], -) -def test_invalid_inputs_raise_usage_error(kwargs: dict[str, object], message: str) -> None: - with pytest.raises(click.UsageError, match=re.escape(message)): - _invoke_clean_expired_messages( - batch_size=1000, - graceful_period=21, - start_from=kwargs["start_from"], - end_before=kwargs["end_before"], - from_days_ago=kwargs["from_days_ago"], - before_days=kwargs["before_days"], - dry_run=False, - ) diff --git a/api/tests/unit_tests/commands/test_data_migration.py b/api/tests/unit_tests/commands/test_data_migration.py new file mode 100644 index 00000000000..92f9df31884 --- /dev/null +++ b/api/tests/unit_tests/commands/test_data_migration.py @@ -0,0 +1,2971 @@ +from __future__ import annotations + +import importlib +import io +import json +import os +import threading +import time +from collections.abc import Iterator, Sequence +from datetime import datetime, timedelta +from pathlib import Path +from types import ModuleType, SimpleNamespace +from typing import cast, override + +import pytest +import sqlalchemy as sa +from click.testing import CliRunner +from sqlalchemy import event +from sqlalchemy.engine import Engine +from sqlalchemy.exc import OperationalError +from sqlalchemy.orm import Session, SessionTransaction, sessionmaker + +from commands import data_migration +from commands.data_migration import ( + CONFLICT_STRATEGY_CHOICES, + ID_STRATEGY_CHOICES, + _confirm_wizard_summary, + _print_auto_tools, + _print_final_tool_selection, + _print_wizard_step, + _prompt_additional_tools, + _prompt_output_file, + _prompt_tool_category, + _resolve_mcp_tool_names, + export_migration_data, + export_migration_data_template, + import_migration_data, + migration_data_wizard, + parse_index_selection, +) +from graphon.model_runtime.entities.model_entities import ModelType +from models import Dataset, DatasetPermission, DatasetPermissionEnum, TenantAccountJoin, TenantAccountRole +from models.account import Tenant +from models.base import TypeBase +from models.enums import CredentialSourceType +from models.provider import LoadBalancingModelConfig, ProviderModel +from models.tools import MCPToolProvider +from services import legacy_model_type_migration as migration_module +from services.data_migration.entities import ( + ConflictStrategy, + ExportResult, + ImportOptions, + ImportResult, + MigrationPackage, + ReportContext, +) +from services.enterprise.rbac_service import ReplaceMemberBindings, ReplaceUserAccessPolicies +from tests.helpers.legacy_model_type_migration import ( + ALL_TABLE_NAMES, + LEGACY_TO_CANONICAL, + DirtyDataFixture, + assert_tenant_rows_use_only_canonical_model_types, + count_rows, + create_minimal_legacy_model_type_schema, + fetch_table_rows, + seed_legacy_model_type_dirty_data, + snapshot_legacy_model_type_state, +) + +"""CLI migration tests using real command-owned SQLite sessions. + +Migration services and package I/O remain fakes so these cases can focus on +CLI parsing and session lifecycle without fabricating the SQLAlchemy boundary. +""" + + +class TestDataMigrationCommandsCases: + @staticmethod + def test_export_command_requires_input_and_output(): + result = CliRunner().invoke(export_migration_data, []) + + assert result.exit_code != 0 + assert export_migration_data.name == "export-app-migration" + assert "--input" in result.output + assert "--output" in result.output + + @staticmethod + def test_import_command_requires_input_and_target_tenant_or_package_metadata(): + result = CliRunner().invoke(import_migration_data, []) + + assert result.exit_code != 0 + assert import_migration_data.name == "import-app-migration" + assert "--input" in result.output + + @staticmethod + def test_import_command_does_not_expose_unimplemented_map_id_strategy(): + assert ID_STRATEGY_CHOICES == ["preserve-id", "generate-new-id"] + + @staticmethod + def test_export_template_command_prints_scripted_json_template(): + result = CliRunner().invoke(export_migration_data_template, []) + + assert result.exit_code == 0 + assert export_migration_data_template.name == "app-migration-template" + template = json.loads(result.output) + assert template == { + "source_tenant": {"mode": "single", "id": "", "name": "admin's Workspace"}, + "apps": {"modes": ["workflow", "advanced-chat"], "ids": [], "all": True}, + "include_referenced_tools": True, + "additional_tools": {"api_tools": [], "workflow_tools": [], "mcp_tools": []}, + "include_secrets": False, + "import_options": { + "create_app_api_token_on_import": False, + "id_strategy": "preserve-id", + "conflict_strategy": "fail", + }, + } + + @staticmethod + def test_export_template_command_writes_output_file(tmp_path: Path): + output_file = tmp_path / "export-template.json" + + result = CliRunner().invoke(export_migration_data_template, ["--output", str(output_file)]) + + assert result.exit_code == 0 + assert f"Output written to {output_file}" in result.output + assert json.loads(output_file.read_text())["apps"]["all"] is True + + @staticmethod + def test_export_template_command_requires_overwrite_for_existing_output(tmp_path: Path): + output_file = tmp_path / "export-template.json" + output_file.write_text("{}") + + result = CliRunner().invoke(export_migration_data_template, ["--output", str(output_file)]) + + assert result.exit_code != 0 + assert "already exists" in result.output + + @staticmethod + def test_export_command_uses_cli_owned_session(monkeypatch, tmp_path: Path, sqlite_engine: Engine): + captured: dict[str, object] = {} + input_file = tmp_path / "export-config.json" + output_file = tmp_path / "migration-package.json" + input_file.write_text(json.dumps({"source_tenant": {"name": "source"}, "apps": {"all": True}})) + package = MigrationPackage.from_mapping({"metadata": {"version": "1", "source_scope": "single"}}) + + class FakeMigrationExportService: + def export(self, selection, *, session): + captured["session"] = session + captured["selection"] = selection + return ExportResult(package=package, report_items=[], report_context=ReportContext()) + + class FakeMigrationPackageService: + def save_package(self, package_to_save, path, *, overwrite): + captured["package"] = package_to_save + captured["path"] = path + captured["overwrite"] = overwrite + + monkeypatch.setattr(data_migration, "MigrationExportService", FakeMigrationExportService) + monkeypatch.setattr(data_migration, "MigrationPackageService", FakeMigrationPackageService) + + result = CliRunner().invoke( + export_migration_data, + ["--input", str(input_file), "--output", str(output_file)], + ) + + assert result.exit_code == 0 + captured_session = captured["session"] + assert isinstance(captured_session, Session) + assert captured_session.get_bind() is sqlite_engine + assert captured_session.in_transaction() is False + assert captured["package"] is package + assert captured["path"] == str(output_file) + assert captured["overwrite"] is False + + @staticmethod + def test_import_command_uses_cli_owned_session(monkeypatch, tmp_path: Path, sqlite_engine: Engine): + captured: dict[str, object] = {} + input_file = tmp_path / "migration-package.json" + input_file.write_text("{}") + package = MigrationPackage.from_mapping( + { + "metadata": { + "version": "1", + "source_scope": "single", + "target_tenant": {"name": "target"}, + "import_options": {"conflict_strategy": "fail"}, + } + } + ) + + class FakeMigrationImportService: + def import_package(self, request, *, session): + captured["session"] = session + captured["request"] = request + return ImportResult(report_items=[], report_context=ReportContext(target_tenant="target")) + + class FakeMigrationPackageService: + def load_package(self, path): + captured["path"] = path + return package + + monkeypatch.setattr(data_migration, "MigrationImportService", FakeMigrationImportService) + monkeypatch.setattr(data_migration, "MigrationPackageService", FakeMigrationPackageService) + + result = CliRunner().invoke( + import_migration_data, + ["--input", str(input_file), "--conflict-strategy", "skip"], + ) + + assert result.exit_code == 0 + captured_session = captured["session"] + assert isinstance(captured_session, Session) + assert captured_session.get_bind() is sqlite_engine + assert captured_session.in_transaction() is False + assert captured["path"] == str(input_file) + request = captured["request"] + assert request.package is package + assert request.options_override == ImportOptions(conflict_strategy=ConflictStrategy.SKIP) + + +class TestDataMigrationWizardCases: + @staticmethod + def test_parse_index_selection_supports_all() -> None: + assert parse_index_selection("all", ["a", "b", "c"]) == ["a", "b", "c"] + + @staticmethod + def test_wizard_command_uses_app_migration_name() -> None: + assert migration_data_wizard.name == "app-migration-wizard" + + @staticmethod + def test_parse_index_selection_supports_comma_indexes() -> None: + assert parse_index_selection("1, 3", ["a", "b", "c"]) == ["a", "c"] + + @staticmethod + def test_print_wizard_step_adds_separator(monkeypatch: pytest.MonkeyPatch) -> None: + output_lines: list[str] = [] + + monkeypatch.setattr("commands.data_migration.click.echo", output_lines.append) + + _print_wizard_step("App Selection") + + assert output_lines == ["", "==== App Selection ===="] + + @staticmethod + def test_conflict_strategy_choices_exclude_replace() -> None: + assert CONFLICT_STRATEGY_CHOICES == ["fail", "skip", "update"] + + @staticmethod + def test_prompt_app_ids_explains_comma_selection_and_default(monkeypatch: pytest.MonkeyPatch) -> None: + from commands.data_migration import _prompt_app_ids + + prompts = [] + output_lines: list[str] = [] + apps = [ + type("App", (), {"id": "app-1", "name": "embedded", "mode": "workflow"})(), + type("App", (), {"id": "app-2", "name": "main", "mode": "advanced-chat"})(), + ] + + def capture_prompt(text: str, **kwargs: object) -> str: + prompts.append((text, kwargs)) + return "1,2" + + monkeypatch.setattr("commands.data_migration.click.echo", output_lines.append) + monkeypatch.setattr("commands.data_migration.click.prompt", capture_prompt) + + assert _prompt_app_ids(apps) == ["app-1", "app-2"] + assert prompts == [("Select apps by number, comma-separated numbers, or all", {"default": "all"})] + assert "Currently supported app types: workflow and chatflow." in output_lines + + @staticmethod + def test_prompt_tool_category_marks_auto_discovered_tools(monkeypatch: pytest.MonkeyPatch) -> None: + output_lines: list[str] = [] + + monkeypatch.setattr("commands.data_migration.click.echo", output_lines.append) + monkeypatch.setattr("commands.data_migration.click.prompt", lambda *_args, **_kwargs: "") + + selected = _prompt_tool_category( + "Custom API tools", + [("weather", "weather", "tool-id"), ("calendar", "calendar", "calendar-id")], + auto_tools={"weather": "tool-id"}, + ) + + assert selected == [] + assert "1. [auto] weather (tool-id)" in output_lines + assert "2. [ ] calendar (calendar-id)" in output_lines + assert output_lines[:2] == ["", "==== Custom API tools ===="] + + @staticmethod + def test_prompt_tool_category_explains_comma_selection_and_default(monkeypatch: pytest.MonkeyPatch) -> None: + prompts = [] + + def capture_prompt(text: str, **kwargs: object) -> str: + prompts.append((text, kwargs)) + return "" + + monkeypatch.setattr("commands.data_migration.click.echo", lambda *_args, **_kwargs: None) + monkeypatch.setattr("commands.data_migration.click.prompt", capture_prompt) + + selected = _prompt_tool_category( + "Custom API tools", + [("weather", "weather", "tool-id")], + auto_tools={}, + ) + + assert selected == [] + assert prompts == [ + ( + "Select custom api tools by number, comma-separated numbers, all, or empty", + {"default": "", "show_default": "empty"}, + ) + ] + + @staticmethod + def test_prompt_output_file_shows_default(monkeypatch: pytest.MonkeyPatch) -> None: + prompts = [] + + def capture_prompt(text: str, **kwargs: object) -> str: + prompts.append((text, kwargs)) + return "migration-data.json" + + monkeypatch.setattr("commands.data_migration.click.prompt", capture_prompt) + + assert _prompt_output_file() == ("migration-data.json", False) + assert prompts[0][0] == "Output path" + assert prompts[0][1]["show_default"] is True + + @staticmethod + def test_prompt_tool_category_marks_auto_by_detail_and_supports_multi_select( + monkeypatch: pytest.MonkeyPatch, + ) -> None: + monkeypatch.setattr("commands.data_migration.click.echo", lambda *_args, **_kwargs: None) + monkeypatch.setattr("commands.data_migration.click.prompt", lambda *_args, **_kwargs: "1,2") + + selected = _prompt_tool_category( + "Workflow tools", + [("tool-1", "embedded", "app-1"), ("tool-2", "other", "app-2")], + auto_tools={"embedded": "app-1"}, + ) + + assert selected == ["tool-1", "tool-2"] + + @staticmethod + def test_prompt_tool_category_marks_auto_by_value() -> None: + output_lines: list[str] = [] + + from commands import data_migration + + original_echo = data_migration.click.echo + original_prompt = data_migration.click.prompt + data_migration.click.echo = output_lines.append # pyrefly: ignore[bad-assignment] + data_migration.click.prompt = lambda *_args, **_kwargs: "" + try: + _prompt_tool_category( + "Workflow tools", + [("tool-1", "embedded_workflow_as_tool", "tool-1")], + auto_tools={"embedded_workflow_as_tool": "tool-1"}, + ) + finally: + data_migration.click.echo = original_echo + data_migration.click.prompt = original_prompt + + assert "1. [auto] embedded_workflow_as_tool (tool-1)" in output_lines + + @staticmethod + def test_print_auto_tools_lists_each_category(monkeypatch: pytest.MonkeyPatch) -> None: + output_lines: list[str] = [] + + monkeypatch.setattr("commands.data_migration.click.echo", output_lines.append) + + _print_auto_tools( + { + "api_tools": {"weather": "3bac3aa9-dd87-4351-9459-a7099137b028"}, + "workflow_tools": {"embedded_workflow_as_tool": "e6024578-41b7-4fb5-a81f-9201358e5835"}, + "mcp_tools": {}, + } + ) + + assert "Automatically discovered tools:" in output_lines + assert "Custom API tools" in output_lines + assert "- weather: 3bac3aa9-dd87-4351-9459-a7099137b028" in output_lines + assert "Workflow tools" in output_lines + assert "- embedded_workflow_as_tool: e6024578-41b7-4fb5-a81f-9201358e5835" in output_lines + assert "MCP tools" in output_lines + assert "- none" in output_lines + + @staticmethod + @pytest.mark.parametrize("sqlite_session", [(MCPToolProvider,)], indirect=True) + def test_resolve_mcp_tool_names_does_not_compare_non_uuid_identifier_to_uuid_id( + monkeypatch: pytest.MonkeyPatch, + sqlite_session: Session, + ) -> None: + tenant_id = "49a99e46-bc2c-4885-91fa-47615f6192b5" + provider = MCPToolProvider( + name="resolved-provider", + server_identifier="my-test-mcp", + server_url="encrypted-url", + server_url_hash="server-url-hash", + icon=None, + tenant_id=tenant_id, + user_id="d13f1f2a-5c68-4a2f-8d76-cb59cefe95b8", + encrypted_credentials=None, + authed=False, + tools="[]", + ) + sqlite_session.add(provider) + sqlite_session.commit() + monkeypatch.setattr(data_migration.db, "session", sqlite_session) + + resolved = _resolve_mcp_tool_names(tenant_id, {"alias": provider.server_identifier}) + + assert resolved == {provider.name: provider.id} + + @staticmethod + def test_prompt_additional_tools_prints_final_selection_when_skipped(monkeypatch: pytest.MonkeyPatch) -> None: + output_lines: list[str] = [] + confirm_prompts: list[tuple[str, dict[str, object]]] = [] + + def capture_confirm(prompt: str, **kwargs: object) -> bool: + confirm_prompts.append((prompt, kwargs)) + return False + + monkeypatch.setattr("commands.data_migration.click.confirm", capture_confirm) + monkeypatch.setattr("commands.data_migration.click.echo", output_lines.append) + + selected = _prompt_additional_tools( + "tenant-id", + { + "api_tools": {"weather": "3bac3aa9-dd87-4351-9459-a7099137b028"}, + "workflow_tools": {}, + "mcp_tools": {}, + }, + ) + + assert selected == {"api_tools": [], "workflow_tools": [], "mcp_tools": []} + assert confirm_prompts == [ + ("Export additional tools manually? [y/n, default: n]", {"default": False, "show_default": False}) + ] + assert "Final tools to export:" in output_lines + assert "- [auto] weather: 3bac3aa9-dd87-4351-9459-a7099137b028" in output_lines + + @staticmethod + def test_final_tool_selection_deduplicates_manual_tool_already_auto(monkeypatch: pytest.MonkeyPatch) -> None: + output_lines: list[str] = [] + + monkeypatch.setattr("commands.data_migration.click.echo", output_lines.append) + + _print_final_tool_selection( + { + "api_tools": {}, + "workflow_tools": {"embedded_workflow_as_tool": "e6024578-41b7-4fb5-a81f-9201358e5835"}, + "mcp_tools": {}, + }, + { + "api_tools": [], + "workflow_tools": ["e6024578-41b7-4fb5-a81f-9201358e5835"], + "mcp_tools": [], + }, + {"e6024578-41b7-4fb5-a81f-9201358e5835": "embedded_workflow_as_tool: e6024578"}, + ) + + assert "- [auto] embedded_workflow_as_tool: e6024578-41b7-4fb5-a81f-9201358e5835" in output_lines + assert not any(line.startswith("- [manual]") for line in output_lines) + + @staticmethod + def test_prompt_output_file_rejects_yes_no_typo(monkeypatch: pytest.MonkeyPatch) -> None: + import click + import pytest + + monkeypatch.setattr("commands.data_migration.click.prompt", lambda *_args, **_kwargs: "y") + + with pytest.raises(click.ClickException, match="Output path must be a file path"): + _prompt_output_file() + + @staticmethod + def test_confirm_wizard_summary_shows_conflict_strategy(monkeypatch: pytest.MonkeyPatch) -> None: + output_lines: list[str] = [] + confirm_prompts: list[tuple[str, dict[str, object]]] = [] + + monkeypatch.setattr("commands.data_migration.click.echo", output_lines.append) + monkeypatch.setattr( + "commands.data_migration.click.confirm", + lambda prompt, **kwargs: confirm_prompts.append((prompt, kwargs)) or True, + ) + + _confirm_wizard_summary( + tenant_name="admin's Workspace", + app_names=["main_chatflow"], + auto_tools={"api_tools": {}, "workflow_tools": {}, "mcp_tools": {}}, + additional_tools={"api_tools": [], "workflow_tools": [], "mcp_tools": []}, + manual_labels={}, + include_referenced_tools=True, + include_secrets=False, + create_tokens=True, + id_strategy="preserve-id", + conflict_strategy="fail", + output_file="migration-data.json", + ) + + assert "id strategy: preserve-id" in output_lines + assert "conflict strategy: fail" in output_lines + assert confirm_prompts == [ + ("Write migration package? [y/n, default: y]", {"default": True, "show_default": False}) + ] + + @staticmethod + def test_confirm_wizard_summary_shows_final_deduplicated_tool_selection(monkeypatch: pytest.MonkeyPatch) -> None: + output_lines: list[str] = [] + + monkeypatch.setattr("commands.data_migration.click.echo", output_lines.append) + monkeypatch.setattr("commands.data_migration.click.confirm", lambda *_args, **_kwargs: True) + + _confirm_wizard_summary( + tenant_name="admin's Workspace", + app_names=["main_chatflow"], + auto_tools={ + "api_tools": {"weather": "weather-id"}, + "workflow_tools": {"embedded_workflow_as_tool": "workflow-tool-id"}, + "mcp_tools": {}, + }, + additional_tools={ + "api_tools": ["weather-id", "calendar"], + "workflow_tools": [], + "mcp_tools": ["mcp-id"], + }, + manual_labels={ + "calendar": "calendar: calendar-id", + "mcp-id": "my-test-mcp: mcp-id", + }, + include_referenced_tools=True, + include_secrets=False, + create_tokens=False, + id_strategy="preserve-id", + conflict_strategy="update", + output_file="migration-data.json", + ) + + assert "Final tools to export:" in output_lines + assert "Custom API tools" in output_lines + assert "- [auto] weather: weather-id" in output_lines + assert "- [manual] calendar: calendar-id" in output_lines + assert "Workflow tools" in output_lines + assert "- [auto] embedded_workflow_as_tool: workflow-tool-id" in output_lines + assert "MCP tools" in output_lines + assert "- [manual] my-test-mcp: mcp-id" in output_lines + assert not any(line.startswith("additional api tools:") for line in output_lines) + assert not any(line.startswith("additional workflow tools:") for line in output_lines) + assert not any(line.startswith("additional mcp tools:") for line in output_lines) + assert "- [manual] weather-id" not in output_lines + + @staticmethod + def test_import_options_prompts_explain_secrets_reuse_and_conflicts(monkeypatch: pytest.MonkeyPatch) -> None: + import click + + from commands.data_migration import _prompt_import_options + + output_lines: list[str] = [] + confirm_prompts: list[tuple[str, dict[str, object]]] = [] + prompt_calls = [] + + def capture_confirm(prompt: str, **kwargs: object) -> bool: + confirm_prompts.append((prompt, kwargs)) + return False + + def capture_prompt(prompt: str, **kwargs: object) -> object: + prompt_calls.append((prompt, kwargs)) + return kwargs["default"] + + monkeypatch.setattr("commands.data_migration.click.echo", output_lines.append) + monkeypatch.setattr("commands.data_migration.click.confirm", capture_confirm) + monkeypatch.setattr("commands.data_migration.click.prompt", capture_prompt) + + include_secrets, create_tokens, id_strategy, conflict_strategy = _prompt_import_options() + + assert include_secrets is False + assert create_tokens is False + assert id_strategy == "preserve-id" + assert conflict_strategy == "update" + assert "Secrets include workflow/app DSL secret values, custom API tool credentials," in output_lines + assert "-- Secrets --" in output_lines + assert "If you choose no, credentials are omitted or masked," in output_lines + assert "-- App API Tokens --" in output_lines + assert "When enabled, import will create an app API token if the imported app has none," in output_lines + assert "or reuse an existing app API token if one already exists." in output_lines + assert "-- ID Strategy --" in output_lines + assert "ID strategy controls whether imported app and tool IDs preserve source IDs" in output_lines + assert "or use target-generated IDs." in output_lines + assert "preserve-id: keep source IDs where the target service supports it." in output_lines + assert ( + "generate-new-id: let the target environment generate new IDs and rewrite references via mapping." + in output_lines + ) + assert "-- Conflict Strategy --" in output_lines + assert "Conflict strategy controls what import does when a target resource already exists." in output_lines + assert "fail: stop at the first conflict; previously committed resources are not rolled back." in output_lines + assert "skip: keep the existing target resource and skip importing that resource." in output_lines + assert "update: update the existing target resource in place." in output_lines + assert confirm_prompts == [ + ("Include secrets in output JSON? [y/n, default: n]", {"default": False, "show_default": False}), + ( + "Create or reuse app API tokens during import? [y/n, default: n]", + {"default": False, "show_default": False}, + ), + ] + assert prompt_calls[0][0] == "Import ID strategy. Enter one of: preserve-id, generate-new-id" + assert prompt_calls[0][1]["default"] == "preserve-id" + assert prompt_calls[0][1]["show_default"] is True + id_strategy_type = prompt_calls[0][1]["type"] + assert isinstance(id_strategy_type, click.Choice) + assert list(id_strategy_type.choices) == ID_STRATEGY_CHOICES + assert prompt_calls[1][0] == "Import conflict strategy. Enter one of: fail, skip, update" + assert prompt_calls[1][1]["default"] == "update" + assert prompt_calls[1][1]["show_default"] is True + conflict_strategy_type = prompt_calls[1][1]["type"] + assert isinstance(conflict_strategy_type, click.Choice) + assert list(conflict_strategy_type.choices) == CONFLICT_STRATEGY_CHOICES + + +class _DatabaseError(Exception): + def __init__( + self, + *, + pgcode: str | None = None, + sqlstate: str | None = None, + errno: int | None = None, + ) -> None: + super().__init__() + self.pgcode = pgcode + self.sqlstate = sqlstate + self.errno = errno + + +def _persist_dataset( + session: Session, + *, + dataset_id: str = "dataset-1", + tenant_id: str = "tenant-1", + permission: DatasetPermissionEnum = DatasetPermissionEnum.ONLY_ME, + created_by: str = "creator-account-1", + maintainer: str | None = None, +) -> Dataset: + dataset = Dataset( + id=dataset_id, + tenant_id=tenant_id, + name=f"Dataset {dataset_id}", + permission=permission, + created_by=created_by, + maintainer=maintainer, + ) + session.add(dataset) + session.commit() + return dataset + + +def _parse_json_lines(output: io.StringIO) -> list[dict[str, object]]: + return [json.loads(line) for line in output.getvalue().splitlines() if line.strip()] + + +def _json_key(value: object) -> str: + return json.dumps(value, sort_keys=True) + + +def _event_signature(line: dict[str, object]) -> tuple[object, ...] | None: + event = line.get("event") + attrs = line.get("attrs") + if not isinstance(attrs, dict): + return None + + if event == "row_updated": + return ( + event, + attrs.get("table_name"), + attrs.get("id"), + _json_key(attrs.get("business_key")), + _json_key(attrs.get("old_values")), + _json_key(attrs.get("new_values")), + _json_key(attrs.get("rewrite_source")), + ) + if event == "row_deleted": + return ( + event, + attrs.get("table_name"), + attrs.get("id"), + _json_key(attrs.get("business_key")), + attrs.get("merge_winner_id"), + ) + if event == "group_processed": + return ( + event, + attrs.get("table_name"), + _json_key(attrs.get("business_key")), + tuple(cast(list[str], attrs.get("group_row_ids", []))), + ) + return None + + +def _collect_processing_signatures(lines: list[dict[str, object]]) -> set[tuple[object, ...]]: + signatures: set[tuple[object, ...]] = set() + for line in lines: + signature = _event_signature(line) + if signature is not None: + signatures.add(signature) + return signatures + + +def _cache_event_row_ids( + lines: list[dict[str, object]], + *, + table_name: str, + row_ids: set[str], + event_name: str, +) -> set[str]: + matching_row_ids: set[str] = set() + for line in lines: + if line.get("event") != event_name: + continue + attrs = line.get("attrs") + if not isinstance(attrs, dict): + continue + if attrs.get("table_name") != table_name: + continue + row_id = str(attrs.get("id")) + if row_id in row_ids: + matching_row_ids.add(row_id) + return matching_row_ids + + +def _patch_batch_size( + monkeypatch: pytest.MonkeyPatch, + *, + batch_size: int, +) -> None: + original_init = migration_module.Migration.__init__ + + def _patched_init( + self: migration_module.Migration, + tenant_id: str, + engine: sa.Engine, + apply: bool, + output: io.TextIOBase, + model_types: Sequence[ModelType], + orm_models: Sequence[migration_module.ORMModel], + ) -> None: + original_init(self, tenant_id, engine, apply, output, model_types, orm_models) + self._batch_size = batch_size + + monkeypatch.setattr(migration_module.Migration, "__init__", _patched_init) + + +def _insert_provider_model( + engine: sa.Engine, + *, + row_id: str, + tenant_id: str, + provider_name: str, + model_name: str, + model_type: str, + credential_id: str | None, + created_at: datetime, + updated_at: datetime, +) -> None: + with engine.begin() as conn: + conn.execute( + sa.text( + """ + INSERT INTO provider_models + ( + id, tenant_id, provider_name, model_name, model_type, + credential_id, is_valid, created_at, updated_at + ) + VALUES + ( + :id, :tenant_id, :provider_name, :model_name, :model_type, + :credential_id, :is_valid, :created_at, :updated_at + ) + """ + ), + { + "id": row_id, + "tenant_id": tenant_id, + "provider_name": provider_name, + "model_name": model_name, + "model_type": model_type, + "credential_id": credential_id, + "is_valid": True, + "created_at": created_at, + "updated_at": updated_at, + }, + ) + + +def _insert_tenant(engine: sa.Engine, *, tenant_id: str) -> None: + with engine.begin() as conn: + conn.execute( + sa.insert(Tenant).values( + id=tenant_id, + name=f"Tenant {tenant_id}", + plan="basic", + status="normal", + ) + ) + + +def _insert_tenant_default_model( + engine: sa.Engine, + *, + row_id: str, + tenant_id: str, + provider_name: str, + model_name: str, + model_type: str, + created_at: datetime, + updated_at: datetime, +) -> None: + with engine.begin() as conn: + conn.execute( + sa.text( + """ + INSERT INTO tenant_default_models + (id, tenant_id, provider_name, model_name, model_type, created_at, updated_at) + VALUES + (:id, :tenant_id, :provider_name, :model_name, :model_type, :created_at, :updated_at) + """ + ), + { + "id": row_id, + "tenant_id": tenant_id, + "provider_name": provider_name, + "model_name": model_name, + "model_type": model_type, + "created_at": created_at, + "updated_at": updated_at, + }, + ) + + +def _insert_provider_model_setting( + engine: sa.Engine, + *, + row_id: str, + tenant_id: str, + provider_name: str, + model_name: str, + model_type: str, + enabled: bool, + load_balancing_enabled: bool, + created_at: datetime, + updated_at: datetime, +) -> None: + with engine.begin() as conn: + conn.execute( + sa.text( + """ + INSERT INTO provider_model_settings + ( + id, tenant_id, provider_name, model_name, model_type, + enabled, load_balancing_enabled, + created_at, updated_at + ) + VALUES + ( + :id, :tenant_id, :provider_name, :model_name, :model_type, + :enabled, :load_balancing_enabled, + :created_at, :updated_at + ) + """ + ), + { + "id": row_id, + "tenant_id": tenant_id, + "provider_name": provider_name, + "model_name": model_name, + "model_type": model_type, + "enabled": enabled, + "load_balancing_enabled": load_balancing_enabled, + "created_at": created_at, + "updated_at": updated_at, + }, + ) + + +def _insert_load_balancing_model_config( + engine: sa.Engine, + *, + row_id: str, + tenant_id: str, + provider_name: str, + model_name: str, + model_type: str, + name: str, + encrypted_config: str, + credential_id: str, + enabled: bool, + created_at: datetime, + updated_at: datetime, +) -> None: + with engine.begin() as conn: + conn.execute( + sa.text( + """ + INSERT INTO load_balancing_model_configs + ( + id, tenant_id, provider_name, model_name, model_type, name, + encrypted_config, credential_id, credential_source_type, enabled, created_at, updated_at + ) + VALUES + ( + :id, :tenant_id, :provider_name, :model_name, :model_type, :name, + :encrypted_config, :credential_id, :credential_source_type, :enabled, :created_at, :updated_at + ) + """ + ), + { + "id": row_id, + "tenant_id": tenant_id, + "provider_name": provider_name, + "model_name": model_name, + "model_type": model_type, + "name": name, + "encrypted_config": encrypted_config, + "credential_id": credential_id, + "credential_source_type": CredentialSourceType.CUSTOM_MODEL.value, + "enabled": enabled, + "created_at": created_at, + "updated_at": updated_at, + }, + ) + + +class TestLegacyModelTypeMigrationCases: + @staticmethod + @pytest.fixture + def dirty_fixture(sqlite_engine: sa.Engine) -> DirtyDataFixture: + return seed_legacy_model_type_dirty_data(sqlite_engine) + + @staticmethod + @pytest.fixture + def command_module() -> ModuleType: + try: + return importlib.import_module("commands.data_migration") + except ModuleNotFoundError: # pragma: no cover - explicit TDD failure path + pytest.fail( + "commands.data_migration is missing. " + "Implement the `flask data-migrate legacy-model-types` command group before running these tests." + ) + + @staticmethod + @pytest.fixture + def rbac_session(sqlite_engine: sa.Engine, monkeypatch: pytest.MonkeyPatch) -> Iterator[Session]: + """Bind RBAC command reads to persisted SQLite dataset rows.""" + + TypeBase.metadata.create_all( + sqlite_engine, + tables=[ + TypeBase.metadata.tables[model.__tablename__] + for model in (Dataset, DatasetPermission, TenantAccountJoin) + ], + ) + factory = sessionmaker(bind=sqlite_engine, expire_on_commit=False) + monkeypatch.setattr("commands.rbac.session_factory.create_session", factory) + with factory() as session: + yield session + + @staticmethod + def test_data_migrate_group_registers_dataset_permission_rbac_migration(command_module: ModuleType) -> None: + command = command_module.data_migrate.commands["rbac-migrate-dataset-permissions"] + + assert command is command_module.migrate_dataset_permissions_to_rbac + assert "operator_account_id" not in {param.name for param in command.params} + + @staticmethod + def test_data_migrate_group_registers_resource_whitelist_scope_migration(command_module: ModuleType) -> None: + command = command_module.data_migrate.commands["rbac-migrate-resource-whitelist-scopes"] + + assert command is command_module.migrate_only_me_resource_whitelist_scopes_to_automatic_include + + @staticmethod + def test_dataset_permission_rbac_migration_help_mentions_binding_clear_side_effect( + command_module: ModuleType, + ) -> None: + result = CliRunner().invoke( + command_module.data_migrate, + ["rbac-migrate-dataset-permissions", "--help"], + ) + + assert result.exit_code == 0 + normalized_output = " ".join(result.output.split()) + assert "clears existing per-user policy bindings" in normalized_output + assert "recreates legacy partial-member default bindings" in normalized_output + + @staticmethod + def test_dataset_permission_rbac_migration_maps_legacy_permissions_to_enum_scopes() -> None: + rbac_module = importlib.import_module("commands.rbac") + + assert ( + rbac_module._rbac_dataset_scope_for_legacy_permission(rbac_module.DatasetPermissionEnum.ALL_TEAM) + is rbac_module.RBACResourceWhitelistScope.ALL + ) + assert ( + rbac_module._rbac_dataset_scope_for_legacy_permission(rbac_module.DatasetPermissionEnum.PARTIAL_TEAM) + is rbac_module.RBACResourceWhitelistScope.SPECIFIC + ) + assert rbac_module._dataset_permission_enum("partial_members") is rbac_module.DatasetPermissionEnum.PARTIAL_TEAM + assert rbac_module._dataset_permission_enum(None) is rbac_module.DatasetPermissionEnum.ONLY_ME + + @staticmethod + def test_dataset_permission_rbac_migration_uses_dataset_creator_as_operator( + command_module: ModuleType, + rbac_session: Session, + monkeypatch: pytest.MonkeyPatch, + ) -> None: + rbac_module = importlib.import_module("commands.rbac") + _persist_dataset(rbac_session) + calls: list[dict[str, object]] = [] + read_transaction_ended = False + + def fake_replace_whitelist(**kwargs: object) -> None: + assert read_transaction_ended + calls.append(kwargs) + + def _record_transaction_end(session: Session, transaction: object) -> None: + nonlocal read_transaction_ended + del transaction + if session.get_bind() is rbac_session.get_bind(): + read_transaction_ended = True + + event.listen(Session, "after_transaction_end", _record_transaction_end) + monkeypatch.setattr(rbac_module.RBACService.DatasetAccess, "replace_whitelist", fake_replace_whitelist) + try: + command_module.migrate_dataset_permissions_to_rbac.callback( + tenant_id=None, + dataset_id=None, + batch_size=500, + dry_run=False, + ) + finally: + event.remove(Session, "after_transaction_end", _record_transaction_end) + + assert calls[0]["tenant_id"] == "tenant-1" + assert calls[0]["account_id"] == "creator-account-1" + assert calls[0]["dataset_id"] == "dataset-1" + payload = calls[0]["payload"] + assert isinstance(payload, ReplaceMemberBindings) + assert payload.automatic_include_workspace_members is False + + @staticmethod + def test_dataset_permission_rbac_migration_dry_run_outputs_structured_proposed_changes( + command_module: ModuleType, + rbac_session: Session, + monkeypatch: pytest.MonkeyPatch, + ) -> None: + rbac_module = importlib.import_module("commands.rbac") + dataset = _persist_dataset(rbac_session, permission=DatasetPermissionEnum.PARTIAL_TEAM) + rbac_session.add( + DatasetPermission( + dataset_id=dataset.id, + account_id="member-account-1", + tenant_id=dataset.tenant_id, + ) + ) + rbac_session.commit() + monkeypatch.setattr( + rbac_module.RBACService.DatasetAccess, + "replace_whitelist", + lambda **kwargs: pytest.fail("dry-run must not replace whitelist"), + ) + monkeypatch.setattr( + rbac_module.RBACService.DatasetAccess, + "replace_user_access_policies", + lambda **kwargs: pytest.fail("dry-run must not replace user access policies"), + ) + + result = CliRunner().invoke( + command_module.data_migrate, + ["rbac-migrate-dataset-permissions", "--dry-run"], + ) + + assert result.exit_code == 0 + events = [json.loads(line) for line in result.output.splitlines() if line.startswith("{")] + assert [event["action"] for event in events] == ["replace_whitelist", "replace_user_access_policies"] + assert events[0]["before"] == { + "legacy_dataset_permission": "partial_members", + "legacy_partial_member_ids": ["member-account-1"], + } + assert events[0]["after"] == {"rbac_whitelist_scope": "specific"} + assert events[0]["call"] == { + "method": "RBACService.DatasetAccess.replace_whitelist", + "kwargs": { + "tenant_id": "tenant-1", + "account_id": "creator-account-1", + "dataset_id": "dataset-1", + "payload": {"automatic_include_workspace_members": False}, + }, + } + assert events[1]["target_account_id"] == "member-account-1" + assert events[1]["after"] == {"rbac_user_access_policy_ids": ["default"]} + assert events[1]["call"]["kwargs"]["payload"] == {"access_policy_ids": ["default"]} + + @staticmethod + def test_resource_whitelist_scope_migration_specific_preserves_existing_members( + command_module: ModuleType, + rbac_session: Session, + monkeypatch: pytest.MonkeyPatch, + ) -> None: + rbac_module = importlib.import_module("commands.rbac") + _persist_dataset(rbac_session, maintainer="maintainer-account-1") + replace_whitelist_calls: list[dict[str, object]] = [] + replace_policy_calls: list[dict[str, object]] = [] + + monkeypatch.setattr( + rbac_module.RBACService.DatasetAccess, + "legacy_whitelist_config", + lambda **kwargs: SimpleNamespace( + rbac_whitelist_scope="specific", + account_ids=["member-account-2", "member-account-1", "member-account-1"], + ), + ) + monkeypatch.setattr( + rbac_module.RBACService.DatasetAccess, + "replace_whitelist", + lambda **kwargs: replace_whitelist_calls.append(kwargs), + ) + monkeypatch.setattr( + rbac_module.RBACService.DatasetAccess, + "replace_user_access_policies", + lambda **kwargs: replace_policy_calls.append(kwargs), + ) + monkeypatch.setattr(rbac_module.RBACService.MemberRoles, "batch_get", lambda **kwargs: []) + + command_module.migrate_resource_whitelist_scopes_to_automatic_include.callback( + tenant_id=None, + resource_type="dataset", + resource_id=None, + batch_size=500, + member_batch_size=500, + dry_run=False, + ) + + whitelist_payload = replace_whitelist_calls[0]["payload"] + policy_payload = replace_policy_calls[0]["payload"] + assert isinstance(whitelist_payload, ReplaceMemberBindings) + assert isinstance(policy_payload, ReplaceUserAccessPolicies) + assert whitelist_payload.automatic_include_workspace_members is False + assert policy_payload.access_policy_ids == ["default"] + assert policy_payload.account_ids == ["member-account-1", "member-account-2"] + + @staticmethod + def test_resource_whitelist_scope_migration_all_syncs_workspace_members( + command_module: ModuleType, + rbac_session: Session, + monkeypatch: pytest.MonkeyPatch, + ) -> None: + rbac_module = importlib.import_module("commands.rbac") + _persist_dataset(rbac_session, maintainer="maintainer-account-1") + rbac_session.add_all( + [ + TenantAccountJoin( + tenant_id="tenant-1", + account_id="maintainer-account-1", + role=TenantAccountRole.OWNER, + ), + TenantAccountJoin( + tenant_id="tenant-1", + account_id="member-account-1", + role=TenantAccountRole.NORMAL, + ), + ] + ) + rbac_session.commit() + replace_whitelist_calls: list[dict[str, object]] = [] + replace_policy_calls: list[dict[str, object]] = [] + + monkeypatch.setattr( + rbac_module.RBACService.DatasetAccess, + "legacy_whitelist_config", + lambda **kwargs: SimpleNamespace(rbac_whitelist_scope="all", account_ids=[]), + ) + monkeypatch.setattr( + rbac_module.RBACService.DatasetAccess, + "replace_whitelist", + lambda **kwargs: replace_whitelist_calls.append(kwargs), + ) + monkeypatch.setattr( + rbac_module.RBACService.DatasetAccess, + "replace_user_access_policies", + lambda **kwargs: replace_policy_calls.append(kwargs), + ) + monkeypatch.setattr(rbac_module.RBACService.MemberRoles, "batch_get", lambda **kwargs: []) + + command_module.migrate_resource_whitelist_scopes_to_automatic_include.callback( + tenant_id=None, + resource_type="dataset", + resource_id=None, + batch_size=500, + member_batch_size=500, + dry_run=False, + ) + + whitelist_payload = replace_whitelist_calls[0]["payload"] + policy_payload = replace_policy_calls[0]["payload"] + assert isinstance(whitelist_payload, ReplaceMemberBindings) + assert isinstance(policy_payload, ReplaceUserAccessPolicies) + assert whitelist_payload.automatic_include_workspace_members is True + assert policy_payload.access_policy_ids == ["default"] + assert set(policy_payload.account_ids) == {"maintainer-account-1", "member-account-1"} + + @staticmethod + def test_only_me_resource_whitelist_scope_migration_syncs_workspace_members( + command_module: ModuleType, + rbac_session: Session, + monkeypatch: pytest.MonkeyPatch, + ) -> None: + rbac_module = importlib.import_module("commands.rbac") + _persist_dataset(rbac_session, maintainer="maintainer-account-1") + rbac_session.add_all( + [ + TenantAccountJoin( + tenant_id="tenant-1", + account_id="maintainer-account-1", + role=TenantAccountRole.OWNER, + ), + TenantAccountJoin( + tenant_id="tenant-1", + account_id="member-account-1", + role=TenantAccountRole.NORMAL, + ), + ] + ) + rbac_session.commit() + replace_whitelist_calls: list[dict[str, object]] = [] + replace_policy_calls: list[dict[str, object]] = [] + + monkeypatch.setattr( + rbac_module.RBACService.DatasetAccess, + "legacy_whitelist_config", + lambda **kwargs: SimpleNamespace(rbac_whitelist_scope="only_me", account_ids=[]), + ) + monkeypatch.setattr( + rbac_module.RBACService.DatasetAccess, + "replace_whitelist", + lambda **kwargs: replace_whitelist_calls.append(kwargs), + ) + monkeypatch.setattr( + rbac_module.RBACService.DatasetAccess, + "replace_user_access_policies", + lambda **kwargs: replace_policy_calls.append(kwargs), + ) + monkeypatch.setattr(rbac_module.RBACService.MemberRoles, "batch_get", lambda **kwargs: []) + + command_module.migrate_only_me_resource_whitelist_scopes_to_automatic_include.callback( + tenant_id=None, + resource_type="dataset", + resource_id=None, + batch_size=500, + member_batch_size=500, + dry_run=False, + ) + + whitelist_payload = replace_whitelist_calls[0]["payload"] + policy_payload = replace_policy_calls[0]["payload"] + assert isinstance(whitelist_payload, ReplaceMemberBindings) + assert isinstance(policy_payload, ReplaceUserAccessPolicies) + assert whitelist_payload.automatic_include_workspace_members is True + assert policy_payload.access_policy_ids == ["default"] + assert set(policy_payload.account_ids) == {"maintainer-account-1", "member-account-1"} + + @staticmethod + def test_data_migrate_command_defaults_output_to_stdout_stream( + command_module: ModuleType, + monkeypatch: pytest.MonkeyPatch, + tmp_path: Path, + ) -> None: + service_calls: list[dict[str, object]] = [] + fake_stdout = io.StringIO() + + class FakeService: + def __init__( + self, + *, + engine: sa.Engine, + apply: bool, + concurrency: int, + output: io.TextIOBase | None = None, + tables: tuple[str, ...] | None, + model_types: tuple[ModelType, ...], + tenant_ids: tuple[str, ...] | None, + ) -> None: + service_calls.append( + { + "engine": engine, + "apply": apply, + "concurrency": concurrency, + "output": output, + "tables": tables, + "model_types": model_types, + "tenant_ids": tenant_ids, + } + ) + + def migrate(self) -> None: + service_calls.append({"migrated": True}) + + monkeypatch.setattr(command_module, "LegacyModelTypeMigrationService", FakeService) + monkeypatch.setattr(command_module, "db", SimpleNamespace(engine=object())) + monkeypatch.setattr(command_module.sys, "stdout", fake_stdout) + tenant_id_file = tmp_path / "tenant_ids.txt" + tenant_id_file.write_text("tenant-alpha\n", encoding="utf-8") + + data_migrate = command_module.data_migrate + legacy_model_types = data_migrate.commands["legacy-model-types"] + assert legacy_model_types.callback is not None + + legacy_model_types.callback( + apply=True, + tables=("provider_models",), + model_types=("llm", "text-embedding"), + tenant_id_file=str(tenant_id_file), + output=None, + concurrency=7, + ) + + assert service_calls[0]["apply"] is True + assert service_calls[0]["concurrency"] == 7 + assert service_calls[0]["output"] is fake_stdout + assert service_calls[0]["tables"] == ("provider_models",) + assert tuple(cast(list[str], service_calls[0]["tenant_ids"])) == ("tenant-alpha",) + assert service_calls[0]["model_types"] == (ModelType.LLM, ModelType.TEXT_EMBEDDING) + assert service_calls[1] == {"migrated": True} + + @staticmethod + def test_data_migrate_command_opens_output_file_and_closes_stream( + command_module: ModuleType, + monkeypatch: pytest.MonkeyPatch, + tmp_path: Path, + ) -> None: + service_calls: list[dict[str, object]] = [] + + class FakeService: + def __init__( + self, + *, + engine: sa.Engine, + apply: bool, + concurrency: int, + output: io.TextIOBase | None = None, + tables: tuple[str, ...] | None, + model_types: tuple[ModelType, ...], + tenant_ids: tuple[str, ...] | None, + ) -> None: + service_calls.append( + { + "engine": engine, + "apply": apply, + "concurrency": concurrency, + "output": output, + "tables": tables, + "model_types": model_types, + "tenant_ids": tenant_ids, + } + ) + + def migrate(self) -> None: + output = cast(io.TextIOBase, service_calls[0]["output"]) + output.write('{"event":"test"}\n') + service_calls.append({"migrated": True}) + + monkeypatch.setattr(command_module, "LegacyModelTypeMigrationService", FakeService) + monkeypatch.setattr(command_module, "db", SimpleNamespace(engine=object())) + output_path = tmp_path / "migration.jsonl" + + data_migrate = command_module.data_migrate + legacy_model_types = data_migrate.commands["legacy-model-types"] + assert legacy_model_types.callback is not None + + legacy_model_types.callback( + apply=False, + tables=(), + model_types=(), + tenant_id_file=None, + output=output_path, + concurrency=3, + ) + + output_stream = cast(io.TextIOBase, service_calls[0]["output"]) + assert service_calls[0]["concurrency"] == 3 + assert output_stream is not output_path + assert isinstance(output_stream, io.TextIOWrapper) + assert Path(output_stream.name) == output_path + assert output_stream.closed is True + assert output_path.read_text(encoding="utf-8") == '{"event":"test"}\n' + assert service_calls[1] == {"migrated": True} + + @staticmethod + @pytest.mark.parametrize( + ("cpu_count", "expected_concurrency"), + [ + (8, 8), + (None, 1), + ], + ) + def test_data_migrate_command_defaults_concurrency_from_cpu_count_or_falls_back_to_one( + monkeypatch: pytest.MonkeyPatch, + cpu_count: int | None, + expected_concurrency: int, + ) -> None: + service_calls: list[dict[str, object]] = [] + command_module = importlib.import_module("commands.data_migration") + + class FakeService: + def __init__( + self, + *, + engine: sa.Engine, + apply: bool, + concurrency: int, + output: io.TextIOBase | None = None, + tables: tuple[str, ...] | None, + model_types: tuple[ModelType, ...], + tenant_ids: tuple[str, ...] | None, + ) -> None: + service_calls.append( + { + "engine": engine, + "apply": apply, + "concurrency": concurrency, + "output": output, + "tables": tables, + "model_types": model_types, + "tenant_ids": tenant_ids, + } + ) + + def migrate(self) -> None: + service_calls.append({"migrated": True}) + + monkeypatch.setattr(os, "cpu_count", lambda: cpu_count) + importlib.reload(command_module) + try: + monkeypatch.setattr(command_module, "LegacyModelTypeMigrationService", FakeService) + monkeypatch.setattr(command_module, "db", SimpleNamespace(engine=object())) + + result = CliRunner().invoke(command_module.data_migrate, ["legacy-model-types"]) + + assert result.exit_code == 0, result.output + assert expected_concurrency == command_module._DEFAULT_CONCURRENCY + assert service_calls[0]["concurrency"] == expected_concurrency + assert service_calls[1] == {"migrated": True} + finally: + monkeypatch.undo() + importlib.reload(command_module) + + @staticmethod + def test_service_migrate_batches_by_tenant_respects_selected_tables_without_reverse_dependency_expansion( + sqlite_engine: sa.Engine, + ) -> None: + seen_runs: list[tuple[str, tuple[str, ...], tuple[ModelType, ...]]] = [] + + class FakeMigration: + def __init__( + self, + *, + tenant_id: str, + engine: sa.Engine, + apply: bool, + output: io.TextIOBase, + model_types: tuple[ModelType, ...], + orm_models: tuple[migration_module.ORMModel, ...], + ) -> None: + assert engine is sqlite_engine + assert apply is False + seen_runs.append((tenant_id, tuple(model.__tablename__ for model in orm_models), model_types)) + + def run(self) -> None: + return None + + monkeypatch = pytest.MonkeyPatch() + try: + monkeypatch.setattr(migration_module, "Migration", FakeMigration) + service = migration_module.LegacyModelTypeMigrationService( + engine=sqlite_engine, + apply=False, + concurrency=1, + tables=("provider_models", "tenant_default_models"), + model_types=(ModelType.LLM,), + tenant_ids=("tenant-alpha", "tenant-beta"), + ) + + service.migrate() + finally: + monkeypatch.undo() + + assert seen_runs == [ + ("tenant-alpha", ("provider_models", "tenant_default_models"), (ModelType.LLM,)), + ("tenant-beta", ("provider_models", "tenant_default_models"), (ModelType.LLM,)), + ] + + @staticmethod + def test_service_migrate_without_tenant_ids_discovers_tenants_per_selected_table_without_querying_tenants( + sqlite_engine: sa.Engine, + monkeypatch: pytest.MonkeyPatch, + ) -> None: + create_minimal_legacy_model_type_schema(sqlite_engine) + provider_tenant_id = "00000000-0000-0000-0000-000000000111" + default_tenant_id = "00000000-0000-0000-0000-000000000222" + empty_tenant_id = "00000000-0000-0000-0000-000000000333" + for tenant_id in (provider_tenant_id, default_tenant_id, empty_tenant_id): + _insert_tenant(sqlite_engine, tenant_id=tenant_id) + + created_at = datetime(2025, 1, 1, 12, 0, 0) + updated_at = created_at + timedelta(minutes=1) + _insert_provider_model( + sqlite_engine, + row_id="10000000-0000-0000-0000-000000000111", + tenant_id=provider_tenant_id, + provider_name="openai", + model_name="gpt-4o-mini", + model_type="text-generation", + credential_id=None, + created_at=created_at, + updated_at=updated_at, + ) + _insert_tenant_default_model( + sqlite_engine, + row_id="20000000-0000-0000-0000-000000000222", + tenant_id=default_tenant_id, + provider_name="openai", + model_name="gpt-4o-mini", + model_type="text-generation", + created_at=created_at, + updated_at=updated_at, + ) + + seen_runs: list[tuple[str, tuple[str, ...], tuple[ModelType, ...]]] = [] + executed_sql: list[str] = [] + + class FakeMigration: + def __init__( + self, + *, + tenant_id: str, + engine: sa.Engine, + apply: bool, + output: io.TextIOBase, + model_types: tuple[ModelType, ...], + orm_models: tuple[migration_module.ORMModel, ...], + ) -> None: + assert engine is sqlite_engine + assert apply is False + seen_runs.append((tenant_id, tuple(model.__tablename__ for model in orm_models), model_types)) + + def run(self) -> None: + return None + + def _record_sql( + conn: sa.engine.Connection, + cursor: object, + statement: str, + parameters: object, + context: object, + executemany: bool, + ) -> None: + del conn, cursor, parameters, context, executemany + executed_sql.append(statement) + + event.listen(sqlite_engine, "before_cursor_execute", _record_sql) + try: + monkeypatch.setattr(migration_module, "Migration", FakeMigration) + service = migration_module.LegacyModelTypeMigrationService( + engine=sqlite_engine, + apply=False, + tables=("provider_models", "tenant_default_models"), + model_types=(ModelType.LLM,), + ) + + service.migrate() + finally: + event.remove(sqlite_engine, "before_cursor_execute", _record_sql) + + assert seen_runs == [ + (provider_tenant_id, ("provider_models",), (ModelType.LLM,)), + (default_tenant_id, ("tenant_default_models",), (ModelType.LLM,)), + ] + normalized_statements = [" ".join(statement.lower().split()) for statement in executed_sql] + discovery_statements = [statement for statement in normalized_statements if statement.startswith("select")] + table_names = ("provider_models", "tenant_default_models") + table_discovery_statements = [ + statement + for statement in discovery_statements + if any(f" from {table_name} " in f" {statement} " for table_name in table_names) + ] + + assert [statement for statement in discovery_statements if " from tenants " in f" {statement} "] == [] + assert [statement for statement in discovery_statements if " union " in f" {statement} "] == [] + assert [ + next(table_name for table_name in table_names if f" from {table_name} " in f" {statement} ") + for statement in table_discovery_statements + ] == list(table_names) + + @staticmethod + def test_service_migrate_without_tenant_ids_filters_provider_model_tenants_by_selected_model_types( + sqlite_engine: sa.Engine, + monkeypatch: pytest.MonkeyPatch, + ) -> None: + create_minimal_legacy_model_type_schema(sqlite_engine) + llm_tenant_id = "00000000-0000-0000-0000-000000000411" + embedding_tenant_id = "00000000-0000-0000-0000-000000000422" + empty_tenant_id = "00000000-0000-0000-0000-000000000433" + for tenant_id in (llm_tenant_id, embedding_tenant_id, empty_tenant_id): + _insert_tenant(sqlite_engine, tenant_id=tenant_id) + + created_at = datetime(2025, 1, 2, 12, 0, 0) + updated_at = created_at + timedelta(minutes=1) + _insert_provider_model( + sqlite_engine, + row_id="30000000-0000-0000-0000-000000000411", + tenant_id=llm_tenant_id, + provider_name="openai", + model_name="gpt-4o-mini", + model_type="text-generation", + credential_id=None, + created_at=created_at, + updated_at=updated_at, + ) + _insert_provider_model( + sqlite_engine, + row_id="30000000-0000-0000-0000-000000000422", + tenant_id=embedding_tenant_id, + provider_name="openai", + model_name="text-embedding-3-large", + model_type="embeddings", + credential_id=None, + created_at=created_at, + updated_at=updated_at, + ) + + seen_runs: list[tuple[str, tuple[str, ...], tuple[ModelType, ...]]] = [] + + class FakeMigration: + def __init__( + self, + *, + tenant_id: str, + engine: sa.Engine, + apply: bool, + output: io.TextIOBase, + model_types: tuple[ModelType, ...], + orm_models: tuple[migration_module.ORMModel, ...], + ) -> None: + assert engine is sqlite_engine + assert apply is False + seen_runs.append((tenant_id, tuple(model.__tablename__ for model in orm_models), model_types)) + + def run(self) -> None: + return None + + monkeypatch.setattr(migration_module, "Migration", FakeMigration) + service = migration_module.LegacyModelTypeMigrationService( + engine=sqlite_engine, + apply=False, + tables=("provider_models",), + model_types=(ModelType.LLM,), + ) + + service.migrate() + + assert seen_runs == [ + (llm_tenant_id, ("provider_models",), (ModelType.LLM,)), + ] + + @staticmethod + def test_service_migrate_without_tenant_ids_discovers_all_load_balancing_tenants_for_simpler_table_scoped_query( + sqlite_engine: sa.Engine, + monkeypatch: pytest.MonkeyPatch, + ) -> None: + create_minimal_legacy_model_type_schema(sqlite_engine) + inherit_llm_tenant_id = "00000000-0000-0000-0000-000000000511" + inherit_embedding_tenant_id = "00000000-0000-0000-0000-000000000522" + empty_tenant_id = "00000000-0000-0000-0000-000000000533" + for tenant_id in (inherit_llm_tenant_id, inherit_embedding_tenant_id, empty_tenant_id): + _insert_tenant(sqlite_engine, tenant_id=tenant_id) + + created_at = datetime(2025, 1, 3, 12, 0, 0) + updated_at = created_at + timedelta(minutes=1) + _insert_load_balancing_model_config( + sqlite_engine, + row_id="40000000-0000-0000-0000-000000000511", + tenant_id=inherit_llm_tenant_id, + provider_name="openai", + model_name="gpt-4o-mini", + model_type=ModelType.LLM.value, + name="__inherit__", + encrypted_config=json.dumps({"api_key": "inherit-llm"}), + credential_id="50000000-0000-0000-0000-000000000511", + enabled=True, + created_at=created_at, + updated_at=updated_at, + ) + _insert_load_balancing_model_config( + sqlite_engine, + row_id="40000000-0000-0000-0000-000000000522", + tenant_id=inherit_embedding_tenant_id, + provider_name="openai", + model_name="text-embedding-3-large", + model_type=ModelType.TEXT_EMBEDDING.value, + name="__inherit__", + encrypted_config=json.dumps({"api_key": "inherit-embedding"}), + credential_id="50000000-0000-0000-0000-000000000522", + enabled=True, + created_at=created_at, + updated_at=updated_at, + ) + + seen_runs: list[tuple[str, tuple[str, ...], tuple[ModelType, ...]]] = [] + + class FakeMigration: + def __init__( + self, + *, + tenant_id: str, + engine: sa.Engine, + apply: bool, + output: io.TextIOBase, + model_types: tuple[ModelType, ...], + orm_models: tuple[migration_module.ORMModel, ...], + ) -> None: + assert engine is sqlite_engine + assert apply is False + seen_runs.append((tenant_id, tuple(model.__tablename__ for model in orm_models), model_types)) + + def run(self) -> None: + return None + + monkeypatch.setattr(migration_module, "Migration", FakeMigration) + # Load-balancing tenant discovery is a deliberate exception: it scans the + # whole table so the discovery query stays easy to understand, even when + # the scheduled tenant set is wider than the selected model types. + service = migration_module.LegacyModelTypeMigrationService( + engine=sqlite_engine, + apply=False, + tables=("load_balancing_model_configs",), + model_types=(ModelType.LLM,), + ) + + service.migrate() + + assert seen_runs == [ + (inherit_llm_tenant_id, ("load_balancing_model_configs",), (ModelType.LLM,)), + (inherit_embedding_tenant_id, ("load_balancing_model_configs",), (ModelType.LLM,)), + ] + + @staticmethod + def test_service_migrate_with_concurrency_greater_than_one_runs_tenants_in_parallel_without_changing_migration_scope( # noqa: E501 + sqlite_engine: sa.Engine, + ) -> None: + init_calls: list[dict[str, object]] = [] + started_tenants: list[str] = [] + worker_errors: list[BaseException] = [] + release_runs = threading.Event() + all_started = threading.Event() + active_runs = 0 + max_active_runs = 0 + state_lock = threading.Lock() + + class FakeMigration: + def __init__( + self, + *, + tenant_id: str, + engine: sa.Engine, + apply: bool, + output: io.TextIOBase, + model_types: tuple[ModelType, ...], + orm_models: tuple[migration_module.ORMModel, ...], + ) -> None: + self._tenant_id = tenant_id + init_calls.append( + { + "tenant_id": tenant_id, + "engine": engine, + "apply": apply, + "model_types": model_types, + "table_names": tuple(model.__tablename__ for model in orm_models), + } + ) + + def run(self) -> None: + nonlocal active_runs, max_active_runs + with state_lock: + active_runs += 1 + max_active_runs = max(max_active_runs, active_runs) + started_tenants.append(self._tenant_id) + if len(started_tenants) == 2: + all_started.set() + + release_runs.wait(timeout=1) + + with state_lock: + active_runs -= 1 + + monkeypatch = pytest.MonkeyPatch() + try: + monkeypatch.setattr(migration_module, "Migration", FakeMigration) + service = migration_module.LegacyModelTypeMigrationService( + engine=sqlite_engine, + apply=False, + concurrency=2, + tables=("provider_models",), + model_types=(ModelType.LLM,), + tenant_ids=("tenant-alpha", "tenant-beta"), + ) + + def _run_service() -> None: + try: + service.migrate() + except BaseException as exc: # pragma: no cover - test harness + worker_errors.append(exc) + + worker = threading.Thread(target=_run_service) + worker.start() + started_in_parallel = all_started.wait(timeout=0.5) + release_runs.set() + worker.join(timeout=1) + finally: + monkeypatch.undo() + + assert worker_errors == [] + assert started_in_parallel is True + assert worker.is_alive() is False + assert max_active_runs == 2 + assert {call["tenant_id"] for call in init_calls} == {"tenant-alpha", "tenant-beta"} + for call in init_calls: + assert tuple(cast(tuple[str, ...], call["table_names"])) == ("provider_models",) + assert call["model_types"] == (ModelType.LLM,) + + @staticmethod + def test_service_parallel_migrate_serializes_shared_output_by_line( + sqlite_engine: sa.Engine, + ) -> None: + worker_errors: list[BaseException] = [] + start_barrier = threading.Barrier(2) + + class SlowLineOutput(io.StringIO): + def __init__(self) -> None: + super().__init__() + self.overlap_count = 0 + self._in_write = False + self._state_lock = threading.Lock() + + @override + def write(self, s: str) -> int: + with self._state_lock: + if self._in_write: + self.overlap_count += 1 + self._in_write = True + try: + time.sleep(0.01) + return super().write(s) + finally: + with self._state_lock: + self._in_write = False + + class FakeMigration: + def __init__( + self, + *, + tenant_id: str, + engine: sa.Engine, + apply: bool, + output: io.TextIOBase, + model_types: tuple[ModelType, ...], + orm_models: tuple[migration_module.ORMModel, ...], + ) -> None: + self._tenant_id = tenant_id + self._output = output + + def run(self) -> None: + try: + start_barrier.wait(timeout=1) + except threading.BrokenBarrierError as exc: + raise AssertionError("parallel migrate should schedule both tenant runs together") from exc + + for index in range(3): + self._output.write(f"{self._tenant_id}:line-{index}\n") + + monkeypatch = pytest.MonkeyPatch() + output = SlowLineOutput() + try: + monkeypatch.setattr(migration_module, "Migration", FakeMigration) + service = migration_module.LegacyModelTypeMigrationService( + engine=sqlite_engine, + apply=False, + concurrency=2, + output=output, + tables=("provider_models",), + model_types=(ModelType.LLM,), + tenant_ids=("tenant-alpha", "tenant-beta"), + ) + + def _run_service() -> None: + try: + service.migrate() + except BaseException as exc: # pragma: no cover - test harness + worker_errors.append(exc) + + worker = threading.Thread(target=_run_service) + worker.start() + worker.join(timeout=2) + finally: + monkeypatch.undo() + + assert worker.is_alive() is False + assert worker_errors == [] + assert output.overlap_count == 0 + assert sorted(output.getvalue().splitlines()) == sorted( + [ + "tenant-alpha:line-0", + "tenant-alpha:line-1", + "tenant-alpha:line-2", + "tenant-beta:line-0", + "tenant-beta:line-1", + "tenant-beta:line-2", + ] + ) + + @staticmethod + def test_migration_dry_run_emits_json_lines_without_db_or_cache_mutation( + sqlite_engine: sa.Engine, + dirty_fixture: DirtyDataFixture, + monkeypatch: pytest.MonkeyPatch, + ) -> None: + before = snapshot_legacy_model_type_state(sqlite_engine) + deleted_cache_keys: list[str] = [] + + def _record_delete(self: migration_module.ProviderCredentialsCache) -> None: + deleted_cache_keys.append(self.cache_key) + + monkeypatch.setattr(migration_module.ProviderCredentialsCache, "delete", _record_delete) + output = io.StringIO() + + service = migration_module.LegacyModelTypeMigrationService( + engine=sqlite_engine, + apply=False, + output=output, + tenant_ids=(dirty_fixture.primary.tenant_id,), + ) + + service.migrate() + + after = snapshot_legacy_model_type_state(sqlite_engine) + assert after == before + assert deleted_cache_keys == [] + + lines = [json.loads(line) for line in output.getvalue().splitlines() if line.strip()] + assert lines, "dry-run should emit JSON lines" + assert all({"event", "message", "attrs", "ts"} <= set(line) for line in lines) + rendered_output = output.getvalue() + assert dirty_fixture.primary.loser_credential_id in rendered_output + assert dirty_fixture.primary.loser_credential_name in rendered_output + assert dirty_fixture.primary.loser_encrypted_config in rendered_output + + @staticmethod + def test_dry_run_and_apply_share_processing_scope_and_differ_only_on_side_effects( + sqlite_engine: sa.Engine, + dirty_fixture: DirtyDataFixture, + monkeypatch: pytest.MonkeyPatch, + ) -> None: + before = snapshot_legacy_model_type_state(sqlite_engine) + deleted_cache_keys: list[str] = [] + + def _record_delete(self: migration_module.ProviderCredentialsCache) -> None: + deleted_cache_keys.append(self.cache_key) + + monkeypatch.setattr(migration_module.ProviderCredentialsCache, "delete", _record_delete) + + dry_run_output = io.StringIO() + migration_module.LegacyModelTypeMigrationService( + engine=sqlite_engine, + apply=False, + output=dry_run_output, + tenant_ids=(dirty_fixture.primary.tenant_id,), + ).migrate() + after_dry_run = snapshot_legacy_model_type_state(sqlite_engine) + dry_run_lines = _parse_json_lines(dry_run_output) + + apply_output = io.StringIO() + migration_module.LegacyModelTypeMigrationService( + engine=sqlite_engine, + apply=True, + output=apply_output, + tenant_ids=(dirty_fixture.primary.tenant_id,), + ).migrate() + after_apply = snapshot_legacy_model_type_state(sqlite_engine) + apply_lines = _parse_json_lines(apply_output) + + assert after_dry_run == before + assert after_apply != before + + dry_run_signatures = _collect_processing_signatures(dry_run_lines) + apply_signatures = _collect_processing_signatures(apply_lines) + assert apply_signatures == dry_run_signatures + + dry_run_cache_events = [line["event"] for line in dry_run_lines if str(line.get("event")).startswith("cache_")] + apply_cache_events = [line["event"] for line in apply_lines if str(line.get("event")).startswith("cache_")] + assert "cache_deleted" not in dry_run_cache_events + assert "cache_delete_planned" in dry_run_cache_events + assert "cache_deleted" in apply_cache_events + assert deleted_cache_keys + + dry_run_rewrite_signatures = { + signature + for signature in dry_run_signatures + if signature[0] == "row_updated" + and signature[1] in {"provider_models", "load_balancing_model_configs"} + and signature[-1] != _json_key(None) + } + apply_rewrite_signatures = { + signature + for signature in apply_signatures + if signature[0] == "row_updated" + and signature[1] in {"provider_models", "load_balancing_model_configs"} + and signature[-1] != _json_key(None) + } + assert apply_rewrite_signatures == dry_run_rewrite_signatures + + dry_run_lb_signatures = { + signature + for signature in dry_run_signatures + if signature[0] == "row_updated" and signature[1] == "load_balancing_model_configs" + } + apply_lb_signatures = { + signature + for signature in apply_signatures + if signature[0] == "row_updated" and signature[1] == "load_balancing_model_configs" + } + assert apply_lb_signatures == dry_run_lb_signatures + + @staticmethod + def test_provider_models_processing_uses_same_plan_locking_and_transaction_entry_for_dry_run_and_apply( + sqlite_engine: sa.Engine, + dirty_fixture: DirtyDataFixture, + monkeypatch: pytest.MonkeyPatch, + ) -> None: + dry_migration = migration_module.Migration( + tenant_id=dirty_fixture.primary.tenant_id, + engine=sqlite_engine, + apply=False, + output=io.StringIO(), + model_types=(ModelType.LLM,), + orm_models=(ProviderModel,), + ) + candidate = dry_migration._load_provider_model_candidates(None)[0] + business_key = migration_module._ProviderModelBusinessKey( + tenant_id=candidate.row.tenant_id, + provider_name=candidate.row.provider_name, + model_name=candidate.row.model_name, + model_type=candidate.canonical_model_type, + ) + + apply_migration = migration_module.Migration( + tenant_id=dirty_fixture.primary.tenant_id, + engine=sqlite_engine, + apply=True, + output=io.StringIO(), + model_types=(ModelType.LLM,), + orm_models=(ProviderModel,), + ) + + current_phase = {"name": "dry"} + lock_rows_seen: list[tuple[str, bool]] = [] + begin_calls: list[str] = [] + configure_calls: list[str] = [] + + def _record_begin(session: Session, transaction: SessionTransaction) -> None: + if session.get_bind() is sqlite_engine and transaction.parent is None: + begin_calls.append(current_phase["name"]) + + def _fake_build_plan( + self: migration_module.Migration, + session: Session, + candidate: migration_module._RowWithRawModelType[ProviderModel], + *, + lock_rows: bool, + ) -> migration_module._ProviderModelGroupPlan: + assert session.get_bind() is sqlite_engine + lock_rows_seen.append((current_phase["name"], lock_rows)) + return migration_module._ProviderModelGroupPlan( + group_row_ids=[candidate.row.id], + winner=None, + loser_rows=[], + ) + + def _fake_emit_plan( + self: migration_module.Migration, + plan: migration_module._ProviderModelGroupPlan, + *, + session: Session, + tx_id: str, + business_key: migration_module._BusinessKey, + ) -> None: + return None + + def _fake_configure(self: migration_module.Migration, session: Session) -> None: + assert session.get_bind() is sqlite_engine + configure_calls.append(current_phase["name"]) + + monkeypatch.setattr(migration_module.Migration, "_build_provider_model_group_plan", _fake_build_plan) + monkeypatch.setattr(migration_module.Migration, "_emit_provider_model_group_plan", _fake_emit_plan) + monkeypatch.setattr(migration_module.Migration, "_configure_lock_timeout", _fake_configure) + event.listen(Session, "after_transaction_create", _record_begin) + try: + dry_migration._process_provider_model_group(candidate, business_key) + current_phase["name"] = "apply" + apply_migration._process_provider_model_group(candidate, business_key) + finally: + event.remove(Session, "after_transaction_create", _record_begin) + + assert [phase for phase, _ in lock_rows_seen] == ["dry", "apply"] + assert lock_rows_seen[0][1] == lock_rows_seen[1][1] + assert begin_calls == ["dry", "apply"] + assert configure_calls == ["dry", "apply"] + + @staticmethod + @pytest.mark.parametrize( + ("orig", "expected"), + [ + (_DatabaseError(pgcode="55P03"), True), + (_DatabaseError(sqlstate="55P03"), True), + (_DatabaseError(errno=1205), True), + (RuntimeError("canceling statement due to lock timeout"), True), + (_DatabaseError(pgcode="23505"), False), + (_DatabaseError(errno=1213), False), + ], + ) + def test_is_lock_timeout_error_prefers_structured_backend_codes( + sqlite_engine: sa.Engine, + orig: BaseException, + expected: bool, + ) -> None: + migration = migration_module.Migration( + tenant_id="tenant-1", + engine=sqlite_engine, + apply=True, + output=io.StringIO(), + model_types=(ModelType.LLM,), + orm_models=(), + ) + exc = OperationalError("SELECT 1", {}, orig) + + assert migration._is_lock_timeout_error(exc) is expected + + @staticmethod + def test_process_load_balancing_model_config_row_logs_stacktrace_for_lock_timeout( + sqlite_engine: sa.Engine, + monkeypatch: pytest.MonkeyPatch, + ) -> None: + create_minimal_legacy_model_type_schema(sqlite_engine) + created_at = datetime(2025, 1, 1, 12, 0, 0) + _insert_load_balancing_model_config( + sqlite_engine, + row_id="40000000-0000-0000-0000-000000000001", + tenant_id="tenant-1", + provider_name="openai", + model_name="gpt-4o-mini", + model_type="text-generation", + name="credential", + encrypted_config="{}", + credential_id="50000000-0000-0000-0000-000000000001", + enabled=True, + created_at=created_at, + updated_at=created_at, + ) + output = io.StringIO() + migration = migration_module.Migration( + tenant_id="tenant-1", + engine=sqlite_engine, + apply=True, + output=output, + model_types=(ModelType.LLM,), + orm_models=(migration_module.LoadBalancingModelConfig,), + ) + candidate = migration._load_load_balancing_model_config_candidates(None)[0] + lock_timeout_exc = OperationalError("SELECT 1", {}, _DatabaseError(pgcode="55P03")) + transaction_begins = 0 + + def _record_begin(session: Session, transaction: SessionTransaction) -> None: + nonlocal transaction_begins + if session.get_bind() is sqlite_engine and transaction.parent is None: + transaction_begins += 1 + + def _fake_reload( + self: migration_module.Migration, + session: Session, + original_candidate: migration_module._RowWithRawModelType[LoadBalancingModelConfig], + *, + lock_rows: bool, + ) -> migration_module._RowWithRawModelType[LoadBalancingModelConfig] | None: + raise lock_timeout_exc + + monkeypatch.setattr(migration_module.Migration, "_configure_lock_timeout", lambda self, session: None) + monkeypatch.setattr( + migration_module.Migration, + "_reload_load_balancing_model_config_candidate", + _fake_reload, + ) + + event.listen(Session, "after_transaction_create", _record_begin) + try: + migration._process_load_balancing_model_config_row(candidate) + finally: + event.remove(Session, "after_transaction_create", _record_begin) + + lines = _parse_json_lines(output) + assert len(lines) == 1 + assert lines[0]["event"] == "lock_timeout_skipped" + attrs = cast(dict[str, object], lines[0]["attrs"]) + assert attrs["table_name"] == "load_balancing_model_configs" + assert attrs["id"] == candidate.row.id + assert attrs["error"] == str(lock_timeout_exc) + assert isinstance(attrs["stacktrace"], str) + assert "OperationalError" in attrs["stacktrace"] + assert transaction_begins == 1 + + @staticmethod + def test_process_load_balancing_model_config_row_logs_update_after_sql_execution( + sqlite_engine: sa.Engine, + monkeypatch: pytest.MonkeyPatch, + ) -> None: + create_minimal_legacy_model_type_schema(sqlite_engine) + created_at = datetime(2025, 1, 1, 12, 0, 0) + row_id = "40000000-0000-0000-0000-000000000002" + _insert_load_balancing_model_config( + sqlite_engine, + row_id=row_id, + tenant_id="tenant-1", + provider_name="openai", + model_name="gpt-4o-mini", + model_type="text-generation", + name="credential", + encrypted_config="{}", + credential_id="50000000-0000-0000-0000-000000000002", + enabled=True, + created_at=created_at, + updated_at=created_at, + ) + migration = migration_module.Migration( + tenant_id="tenant-1", + engine=sqlite_engine, + apply=True, + output=io.StringIO(), + model_types=(ModelType.LLM,), + orm_models=(migration_module.LoadBalancingModelConfig,), + ) + candidate = migration._load_load_balancing_model_config_candidates(None)[0] + action_log: list[str] = [] + + def _record_begin(session: Session, transaction: SessionTransaction) -> None: + if session.get_bind() is sqlite_engine and transaction.parent is None: + action_log.append("begin") + + def _record_sql( + connection: sa.Connection, + cursor: object, + statement: str, + parameters: object, + context: object, + executemany: bool, + ) -> None: + del connection, cursor, parameters, context, executemany + if statement.lstrip().upper().startswith("UPDATE"): + action_log.append("sql_execute") + + def _fake_configure(self: migration_module.Migration, session: Session) -> None: + action_log.append("configure_lock_timeout") + + original_reload = migration_module.Migration._reload_load_balancing_model_config_candidate + + def _record_reload( + self: migration_module.Migration, + session: Session, + original_candidate: migration_module._RowWithRawModelType[LoadBalancingModelConfig], + *, + lock_rows: bool, + ) -> migration_module._RowWithRawModelType[LoadBalancingModelConfig] | None: + action_log.append(f"reload_candidate:{lock_rows}") + return original_reload(self, session, original_candidate, lock_rows=lock_rows) + + def _fake_log_row_updated(self: migration_module.Migration, *args: object, **kwargs: object) -> None: + action_log.append("log_row_updated") + + def _fake_cache_cleanup(self: migration_module.Migration, *, row_id: str, tx_id: str) -> None: + action_log.append("cache_cleanup") + + monkeypatch.setattr(migration_module.Migration, "_configure_lock_timeout", _fake_configure) + monkeypatch.setattr( + migration_module.Migration, + "_reload_load_balancing_model_config_candidate", + _record_reload, + ) + monkeypatch.setattr(migration_module.Migration, "_log_row_updated", _fake_log_row_updated) + monkeypatch.setattr( + migration_module.Migration, + "_log_load_balancing_model_config_cache_cleanup", + _fake_cache_cleanup, + ) + + event.listen(Session, "after_transaction_create", _record_begin) + event.listen(sqlite_engine, "before_cursor_execute", _record_sql) + try: + migration._process_load_balancing_model_config_row(candidate) + finally: + event.remove(sqlite_engine, "before_cursor_execute", _record_sql) + event.remove(Session, "after_transaction_create", _record_begin) + + assert action_log == [ + "begin", + "configure_lock_timeout", + "reload_candidate:True", + "sql_execute", + "log_row_updated", + "cache_cleanup", + ] + with Session(sqlite_engine) as session: + persisted = session.get(migration_module.LoadBalancingModelConfig, row_id) + assert persisted is not None + assert persisted.model_type == ModelType.LLM + + @staticmethod + def test_load_balancing_model_config_cache_delete_failure_logs_stacktrace( + sqlite_engine: sa.Engine, + dirty_fixture: DirtyDataFixture, + monkeypatch: pytest.MonkeyPatch, + ) -> None: + def _raise_delete_failure(self: migration_module.ProviderCredentialsCache) -> None: + raise RuntimeError("cache delete boom") + + monkeypatch.setattr(migration_module.ProviderCredentialsCache, "delete", _raise_delete_failure) + + output = io.StringIO() + migration_module.LegacyModelTypeMigrationService( + engine=sqlite_engine, + apply=True, + output=output, + tables=("load_balancing_model_configs",), + model_types=(ModelType.LLM,), + tenant_ids=(dirty_fixture.primary.tenant_id,), + ).migrate() + + failed_events = [ + cast(dict[str, object], line["attrs"]) + for line in _parse_json_lines(output) + if line.get("event") == "cache_delete_failed" + and isinstance(line.get("attrs"), dict) + and cast(dict[str, object], line["attrs"]).get("table_name") == "load_balancing_model_configs" + ] + + assert len(failed_events) == 1 + assert failed_events[0]["error"] == "cache delete boom" + assert isinstance(failed_events[0]["stacktrace"], str) + assert "RuntimeError: cache delete boom" in failed_events[0]["stacktrace"] + + @staticmethod + def test_group_completed_logs_exist_for_all_grouped_tables_and_use_canonical_model_type( + sqlite_engine: sa.Engine, + dirty_fixture: DirtyDataFixture, + ) -> None: + output = io.StringIO() + + service = migration_module.LegacyModelTypeMigrationService( + engine=sqlite_engine, + apply=False, + output=output, + tenant_ids=(dirty_fixture.primary.tenant_id,), + ) + + service.migrate() + + lines = _parse_json_lines(output) + group_completed_records = [ + line + for line in lines + if isinstance(line.get("attrs"), dict) and "group_row_ids" in cast(dict[str, object], line["attrs"]) + ] + grouped_table_names = { + cast(dict[str, object], record["attrs"]).get("table_name") for record in group_completed_records + } + + assert grouped_table_names >= { + "provider_models", + "tenant_default_models", + "provider_model_settings", + "provider_model_credentials", + } + + for record in group_completed_records: + attrs = cast(dict[str, object], record["attrs"]) + business_key = cast(dict[str, object], attrs["business_key"]) + assert isinstance(attrs["group_row_ids"], list) + assert attrs["group_row_ids"] + if "model_type" in business_key: + assert business_key["model_type"] in { + ModelType.LLM.value, + ModelType.TEXT_EMBEDDING.value, + ModelType.RERANK.value, + } + assert business_key["model_type"] not in LEGACY_TO_CANONICAL + + @staticmethod + def test_provider_models_group_completed_log_includes_related_canonical_row_ids( + sqlite_engine: sa.Engine, + dirty_fixture: DirtyDataFixture, + monkeypatch: pytest.MonkeyPatch, + ) -> None: + _patch_batch_size(monkeypatch, batch_size=1) + inserted_row_id = "00000000-0000-0000-0000-00000000aa01" + created_at = datetime(2025, 1, 1, 10, 0, 0) + updated_at = created_at + timedelta(minutes=5) + _insert_provider_model( + sqlite_engine, + row_id=inserted_row_id, + tenant_id=dirty_fixture.primary.tenant_id, + provider_name="openai", + model_name="gpt-4o-mini", + model_type=ModelType.LLM.value, + credential_id=dirty_fixture.primary.distinct_credential_id, + created_at=created_at, + updated_at=updated_at, + ) + + output = io.StringIO() + migration_module.LegacyModelTypeMigrationService( + engine=sqlite_engine, + apply=False, + output=output, + tables=("provider_models",), + model_types=(ModelType.LLM,), + tenant_ids=(dirty_fixture.primary.tenant_id,), + ).migrate() + + lines = _parse_json_lines(output) + matching_records = [] + for line in lines: + attrs = line.get("attrs") + if not isinstance(attrs, dict): + continue + business_key = attrs.get("business_key") + if not isinstance(business_key, dict): + continue + if ( + attrs.get("table_name") == "provider_models" + and business_key.get("tenant_id") == dirty_fixture.primary.tenant_id + and business_key.get("provider_name") == "openai" + and business_key.get("model_name") == "gpt-4o-mini" + and business_key.get("model_type") == ModelType.LLM.value + and "group_row_ids" in attrs + ): + matching_records.append(attrs) + + assert len(matching_records) == 1 + assert set(cast(list[str], matching_records[0]["group_row_ids"])) == { + dirty_fixture.primary.provider_model_id, + inserted_row_id, + } + + @staticmethod + def test_provider_model_settings_group_crossing_batches_is_completed_once_with_all_group_row_ids( + sqlite_engine: sa.Engine, + dirty_fixture: DirtyDataFixture, + monkeypatch: pytest.MonkeyPatch, + ) -> None: + _patch_batch_size(monkeypatch, batch_size=1) + inserted_row_id = "00000000-0000-0000-0000-00000000cc01" + created_at = datetime(2025, 1, 1, 9, 0, 0) + updated_at = created_at + timedelta(minutes=10) + _insert_provider_model_setting( + sqlite_engine, + row_id=inserted_row_id, + tenant_id=dirty_fixture.primary.tenant_id, + provider_name="openai", + model_name="gpt-4o-mini", + model_type="text-generation", + enabled=True, + load_balancing_enabled=False, + created_at=created_at, + updated_at=updated_at, + ) + + output = io.StringIO() + migration_module.LegacyModelTypeMigrationService( + engine=sqlite_engine, + apply=False, + output=output, + tables=("provider_model_settings",), + model_types=(ModelType.LLM,), + tenant_ids=(dirty_fixture.primary.tenant_id,), + ).migrate() + + lines = _parse_json_lines(output) + matching_records = [] + for line in lines: + attrs = line.get("attrs") + if not isinstance(attrs, dict): + continue + business_key = attrs.get("business_key") + if not isinstance(business_key, dict): + continue + if ( + attrs.get("table_name") == "provider_model_settings" + and business_key.get("tenant_id") == dirty_fixture.primary.tenant_id + and business_key.get("provider_name") == "openai" + and business_key.get("model_name") == "gpt-4o-mini" + and business_key.get("model_type") == ModelType.LLM.value + and "group_row_ids" in attrs + ): + matching_records.append(attrs) + + assert len(matching_records) == 1 + assert set(cast(list[str], matching_records[0]["group_row_ids"])) == { + dirty_fixture.primary.provider_model_setting_id, + inserted_row_id, + } + + @staticmethod + def test_load_balancing_inherit_rows_are_deduplicated_by_normalized_model_type_before_canonicalization( + sqlite_engine: sa.Engine, + dirty_fixture: DirtyDataFixture, + monkeypatch: pytest.MonkeyPatch, + ) -> None: + older_canonical_row_id = "00000000-0000-0000-0000-00000000dd01" + newer_legacy_row_id = "00000000-0000-0000-0000-00000000dd02" + created_at = datetime(2025, 1, 1, 8, 0, 0) + older_updated_at = created_at + timedelta(minutes=15) + newer_updated_at = created_at + timedelta(minutes=30) + _insert_load_balancing_model_config( + sqlite_engine, + row_id=older_canonical_row_id, + tenant_id=dirty_fixture.primary.tenant_id, + provider_name="openai", + model_name="gpt-4o-mini", + model_type=ModelType.LLM.value, + name="__inherit__", + encrypted_config='{"api_key":"older-inherit"}', + credential_id=dirty_fixture.primary.winner_credential_id, + enabled=True, + created_at=created_at, + updated_at=older_updated_at, + ) + _insert_load_balancing_model_config( + sqlite_engine, + row_id=newer_legacy_row_id, + tenant_id=dirty_fixture.primary.tenant_id, + provider_name="openai", + model_name="gpt-4o-mini", + model_type="text-generation", + name="__inherit__", + encrypted_config='{"api_key":"newer-inherit"}', + credential_id=dirty_fixture.primary.distinct_credential_id, + enabled=True, + created_at=created_at, + updated_at=newer_updated_at, + ) + + deleted_cache_keys: list[str] = [] + + def _record_delete(self: migration_module.ProviderCredentialsCache) -> None: + deleted_cache_keys.append(self.cache_key) + + monkeypatch.setattr(migration_module.ProviderCredentialsCache, "delete", _record_delete) + + tenant_id = dirty_fixture.primary.tenant_id + table_name = "load_balancing_model_configs" + expected_row_ids = {older_canonical_row_id, newer_legacy_row_id} + + dry_run_output = io.StringIO() + migration_module.LegacyModelTypeMigrationService( + engine=sqlite_engine, + apply=False, + output=dry_run_output, + tables=(table_name,), + model_types=(ModelType.LLM,), + tenant_ids=(tenant_id,), + ).migrate() + + dry_run_lines = _parse_json_lines(dry_run_output) + dry_run_signatures = { + signature + for signature in _collect_processing_signatures(dry_run_lines) + if signature[1] == table_name and signature[2] in expected_row_ids + } + dry_run_row_updates = [ + cast(dict[str, object], line["attrs"]) + for line in dry_run_lines + if line.get("event") == "row_updated" + and isinstance(line.get("attrs"), dict) + and cast(dict[str, object], line["attrs"]).get("table_name") == table_name + and str(cast(dict[str, object], line["attrs"]).get("id")) in expected_row_ids + ] + assert len(dry_run_row_updates) == 1 + assert str(dry_run_row_updates[0]["id"]) == newer_legacy_row_id + assert dry_run_row_updates[0]["old_values"] == {"model_type": "text-generation"} + assert dry_run_row_updates[0]["new_values"] == {"model_type": ModelType.LLM.value} + assert all("rewrite_source" not in attrs for attrs in dry_run_row_updates) + + dry_run_row_deletes = [ + cast(dict[str, object], line["attrs"]) + for line in dry_run_lines + if line.get("event") == "row_deleted" + and isinstance(line.get("attrs"), dict) + and cast(dict[str, object], line["attrs"]).get("table_name") == table_name + and str(cast(dict[str, object], line["attrs"]).get("id")) in expected_row_ids + ] + assert len(dry_run_row_deletes) == 1 + assert dry_run_row_deletes[0]["business_key"] == { + "tenant_id": tenant_id, + "provider_name": "openai", + "model_name": "gpt-4o-mini", + "model_type": ModelType.LLM.value, + } + assert dry_run_row_deletes[0]["merge_winner_id"] == newer_legacy_row_id + assert dry_run_row_deletes[0]["row"] == { + "id": older_canonical_row_id, + "tenant_id": tenant_id, + "provider_name": "openai", + "model_name": "gpt-4o-mini", + "model_type": ModelType.LLM.value, + "name": "__inherit__", + "encrypted_config": {"api_key": "older-inherit"}, + "credential_id": dirty_fixture.primary.winner_credential_id, + "credential_source_type": CredentialSourceType.CUSTOM_MODEL.value, + "enabled": True, + "created_at": created_at.isoformat(), + "updated_at": older_updated_at.isoformat(), + } + + dry_run_deleted_index = next( + index + for index, line in enumerate(dry_run_lines) + if line.get("event") == "row_deleted" + and isinstance(line.get("attrs"), dict) + and cast(dict[str, object], line["attrs"]).get("id") == older_canonical_row_id + ) + dry_run_updated_index = next( + index + for index, line in enumerate(dry_run_lines) + if line.get("event") == "row_updated" + and isinstance(line.get("attrs"), dict) + and cast(dict[str, object], line["attrs"]).get("id") == newer_legacy_row_id + ) + assert dry_run_deleted_index < dry_run_updated_index + + dry_run_cache_plan_ids = _cache_event_row_ids( + dry_run_lines, + table_name=table_name, + row_ids=expected_row_ids, + event_name="cache_delete_planned", + ) + assert newer_legacy_row_id in dry_run_cache_plan_ids + + apply_output = io.StringIO() + migration_module.LegacyModelTypeMigrationService( + engine=sqlite_engine, + apply=True, + output=apply_output, + tables=(table_name,), + model_types=(ModelType.LLM,), + tenant_ids=(tenant_id,), + ).migrate() + + apply_lines = _parse_json_lines(apply_output) + apply_signatures = { + signature + for signature in _collect_processing_signatures(apply_lines) + if signature[1] == table_name and signature[2] in expected_row_ids + } + apply_row_updates = [ + cast(dict[str, object], line["attrs"]) + for line in apply_lines + if line.get("event") == "row_updated" + and isinstance(line.get("attrs"), dict) + and cast(dict[str, object], line["attrs"]).get("table_name") == table_name + and str(cast(dict[str, object], line["attrs"]).get("id")) in expected_row_ids + ] + assert len(apply_row_updates) == 1 + assert str(apply_row_updates[0]["id"]) == newer_legacy_row_id + assert apply_signatures == dry_run_signatures + + apply_cache_delete_ids = _cache_event_row_ids( + apply_lines, + table_name=table_name, + row_ids=expected_row_ids, + event_name="cache_deleted", + ) + assert apply_cache_delete_ids == dry_run_cache_plan_ids + assert deleted_cache_keys + + lb_rows = fetch_table_rows(sqlite_engine, table_name, tenant_id=tenant_id) + surviving_rows = [row for row in lb_rows if str(row["id"]) in expected_row_ids] + assert len(surviving_rows) == 1 + surviving_row = surviving_rows[0] + assert surviving_row["id"] == newer_legacy_row_id + assert surviving_row["tenant_id"] == tenant_id + assert surviving_row["provider_name"] == "openai" + assert surviving_row["model_name"] == "gpt-4o-mini" + assert surviving_row["model_type"] == ModelType.LLM.value + assert surviving_row["name"] == "__inherit__" + assert surviving_row["encrypted_config"] == '{"api_key":"newer-inherit"}' + assert surviving_row["credential_id"] == dirty_fixture.primary.distinct_credential_id + assert surviving_row["credential_source_type"] == CredentialSourceType.CUSTOM_MODEL.value + + @staticmethod + def test_load_balancing_non_inherit_rows_do_not_participate_in_normalized_model_type_deduplication( + sqlite_engine: sa.Engine, + dirty_fixture: DirtyDataFixture, + ) -> None: + inserted_row_id = "00000000-0000-0000-0000-00000000dd03" + created_at = datetime(2025, 1, 1, 8, 0, 0) + updated_at = created_at + timedelta(minutes=15) + _insert_load_balancing_model_config( + sqlite_engine, + row_id=inserted_row_id, + tenant_id=dirty_fixture.primary.tenant_id, + provider_name="openai", + model_name="gpt-4o-mini", + model_type=ModelType.LLM.value, + name=dirty_fixture.primary.loser_credential_name, + encrypted_config='{"api_key":"second-lb"}', + credential_id=dirty_fixture.primary.distinct_credential_id, + enabled=True, + created_at=created_at, + updated_at=updated_at, + ) + + output = io.StringIO() + migration_module.LegacyModelTypeMigrationService( + engine=sqlite_engine, + apply=True, + output=output, + tables=("load_balancing_model_configs",), + model_types=(ModelType.LLM,), + tenant_ids=(dirty_fixture.primary.tenant_id,), + ).migrate() + + lines = _parse_json_lines(output) + row_deleted_events = [ + cast(dict[str, object], line["attrs"]) + for line in lines + if line.get("event") == "row_deleted" + and isinstance(line.get("attrs"), dict) + and cast(dict[str, object], line["attrs"]).get("table_name") == "load_balancing_model_configs" + ] + assert row_deleted_events == [] + + lb_rows = fetch_table_rows( + sqlite_engine, + "load_balancing_model_configs", + tenant_id=dirty_fixture.primary.tenant_id, + ) + matching_rows = [ + row + for row in lb_rows + if str(row["id"]) in {dirty_fixture.primary.load_balancing_config_id, inserted_row_id} + ] + assert len(matching_rows) == 2 + assert all(row["model_type"] == ModelType.LLM.value for row in matching_rows) + + @staticmethod + def test_migration_apply_updates_all_five_tables_and_rewrites_credential_references( + sqlite_engine: sa.Engine, + dirty_fixture: DirtyDataFixture, + monkeypatch: pytest.MonkeyPatch, + ) -> None: + deleted_cache_keys: list[str] = [] + + def _record_delete(self: migration_module.ProviderCredentialsCache) -> None: + deleted_cache_keys.append(self.cache_key) + + monkeypatch.setattr(migration_module.ProviderCredentialsCache, "delete", _record_delete) + output = io.StringIO() + + service = migration_module.LegacyModelTypeMigrationService( + engine=sqlite_engine, + apply=True, + output=output, + tenant_ids=(dirty_fixture.primary.tenant_id,), + ) + + service.migrate() + + assert_tenant_rows_use_only_canonical_model_types(sqlite_engine, dirty_fixture.primary.tenant_id) + + provider_model_rows = fetch_table_rows( + sqlite_engine, "provider_models", tenant_id=dirty_fixture.primary.tenant_id + ) + provider_model_row = next( + row for row in provider_model_rows if row["id"] == dirty_fixture.primary.provider_model_id + ) + assert provider_model_row["model_type"] == LEGACY_TO_CANONICAL["text-generation"] + assert provider_model_row["credential_id"] == dirty_fixture.primary.winner_credential_id + + lb_rows = fetch_table_rows( + sqlite_engine, "load_balancing_model_configs", tenant_id=dirty_fixture.primary.tenant_id + ) + lb_row = next(row for row in lb_rows if row["id"] == dirty_fixture.primary.load_balancing_config_id) + assert lb_row["model_type"] == LEGACY_TO_CANONICAL["text-generation"] + assert lb_row["credential_id"] == dirty_fixture.primary.winner_credential_id + assert lb_row["encrypted_config"] == dirty_fixture.primary.winner_encrypted_config + + credential_rows = fetch_table_rows( + sqlite_engine, "provider_model_credentials", tenant_id=dirty_fixture.primary.tenant_id + ) + assert ( + count_rows( + sqlite_engine, + "provider_model_credentials", + tenant_id=dirty_fixture.primary.tenant_id, + ) + == 2 + ) + credential_ids = {str(row["id"]) for row in credential_rows} + assert credential_ids == { + dirty_fixture.primary.winner_credential_id, + dirty_fixture.primary.distinct_credential_id, + } + distinct_row = next(row for row in credential_rows if row["id"] == dirty_fixture.primary.distinct_credential_id) + assert distinct_row["credential_name"] == dirty_fixture.primary.distinct_credential_name + assert distinct_row["model_type"] == LEGACY_TO_CANONICAL["text-generation"] + + rendered_output = output.getvalue() + assert dirty_fixture.primary.loser_credential_id in rendered_output + assert dirty_fixture.primary.loser_encrypted_config in rendered_output + assert any("load_balancing_provider_model_credentials" in key for key in deleted_cache_keys) or any( + "load_balancing_provider_model" in key for key in deleted_cache_keys + ) + + @staticmethod + def test_migration_filters_by_tenant_model_types_and_tables( + sqlite_engine: sa.Engine, + dirty_fixture: DirtyDataFixture, + ) -> None: + before_primary_credentials = fetch_table_rows( + sqlite_engine, + "provider_model_credentials", + tenant_id=dirty_fixture.primary.tenant_id, + ) + before_secondary = { + table_name: fetch_table_rows(sqlite_engine, table_name, tenant_id=dirty_fixture.secondary.tenant_id) + for table_name in ALL_TABLE_NAMES + } + + service = migration_module.LegacyModelTypeMigrationService( + engine=sqlite_engine, + apply=True, + output=io.StringIO(), + tables=("provider_models",), + model_types=(ModelType.LLM,), + tenant_ids=(dirty_fixture.primary.tenant_id,), + ) + + service.migrate() + + assert ( + count_rows( + sqlite_engine, + "provider_model_credentials", + tenant_id=dirty_fixture.primary.tenant_id, + ) + == 3 + ) + credential_rows = fetch_table_rows( + sqlite_engine, + "provider_model_credentials", + tenant_id=dirty_fixture.primary.tenant_id, + ) + assert credential_rows == before_primary_credentials + provider_model_rows = fetch_table_rows( + sqlite_engine, + "provider_models", + tenant_id=dirty_fixture.primary.tenant_id, + ) + provider_model_row = next( + row for row in provider_model_rows if row["id"] == dirty_fixture.primary.provider_model_id + ) + embedding_provider_model_row = next( + row for row in provider_model_rows if row["id"] == dirty_fixture.primary.embedding_provider_model_id + ) + assert provider_model_row["model_type"] == LEGACY_TO_CANONICAL["text-generation"] + assert embedding_provider_model_row["model_type"] == "embeddings" + + tenant_default_row = fetch_table_rows( + sqlite_engine, + "tenant_default_models", + tenant_id=dirty_fixture.primary.tenant_id, + )[0] + assert tenant_default_row["model_type"] == "text-generation" + + provider_model_setting_rows = fetch_table_rows( + sqlite_engine, + "provider_model_settings", + tenant_id=dirty_fixture.primary.tenant_id, + ) + llm_setting_row = next( + row for row in provider_model_setting_rows if row["id"] == dirty_fixture.primary.provider_model_setting_id + ) + embedding_setting_row = next( + row for row in provider_model_setting_rows if row["id"] == dirty_fixture.primary.embedding_setting_id + ) + assert llm_setting_row["model_type"] == "text-generation" + assert embedding_setting_row["model_type"] == "embeddings" + + lb_row = fetch_table_rows( + sqlite_engine, + "load_balancing_model_configs", + tenant_id=dirty_fixture.primary.tenant_id, + )[0] + assert lb_row["model_type"] == "text-generation" + + after_secondary = { + table_name: fetch_table_rows(sqlite_engine, table_name, tenant_id=dirty_fixture.secondary.tenant_id) + for table_name in ALL_TABLE_NAMES + } + assert after_secondary == before_secondary + + @staticmethod + def test_migration_does_not_merge_credentials_with_different_credential_name( + sqlite_engine: sa.Engine, + dirty_fixture: DirtyDataFixture, + ) -> None: + service = migration_module.LegacyModelTypeMigrationService( + engine=sqlite_engine, + apply=True, + output=io.StringIO(), + tenant_ids=(dirty_fixture.primary.tenant_id,), + ) + + service.migrate() + + credential_rows = fetch_table_rows( + sqlite_engine, + "provider_model_credentials", + tenant_id=dirty_fixture.primary.tenant_id, + ) + distinct_row = next(row for row in credential_rows if row["id"] == dirty_fixture.primary.distinct_credential_id) + assert distinct_row["credential_name"] == dirty_fixture.primary.distinct_credential_name + assert distinct_row["model_type"] == LEGACY_TO_CANONICAL["text-generation"] + assert ( + count_rows( + sqlite_engine, + "provider_model_credentials", + tenant_id=dirty_fixture.primary.tenant_id, + ) + == 2 + ) + + @staticmethod + def test_migration_is_idempotent_on_second_apply( + sqlite_engine: sa.Engine, + dirty_fixture: DirtyDataFixture, + ) -> None: + service = migration_module.LegacyModelTypeMigrationService( + engine=sqlite_engine, + apply=True, + output=io.StringIO(), + tenant_ids=(dirty_fixture.primary.tenant_id,), + ) + + service.migrate() + after_first = snapshot_legacy_model_type_state(sqlite_engine) + + service.migrate() + after_second = snapshot_legacy_model_type_state(sqlite_engine) + + assert after_second == after_first diff --git a/api/tests/unit_tests/commands/test_data_migration_commands.py b/api/tests/unit_tests/commands/test_data_migration_commands.py deleted file mode 100644 index ad53964d135..00000000000 --- a/api/tests/unit_tests/commands/test_data_migration_commands.py +++ /dev/null @@ -1,171 +0,0 @@ -"""CLI migration tests using real command-owned SQLite sessions. - -Migration services and package I/O remain fakes so these cases can focus on -CLI parsing and session lifecycle without fabricating the SQLAlchemy boundary. -""" - -import json -from pathlib import Path - -from click.testing import CliRunner -from sqlalchemy.engine import Engine -from sqlalchemy.orm import Session - -from commands import data_migration -from commands.data_migration import ( - ID_STRATEGY_CHOICES, - export_migration_data, - export_migration_data_template, - import_migration_data, -) -from services.data_migration.entities import ( - ConflictStrategy, - ExportResult, - ImportOptions, - ImportResult, - MigrationPackage, - ReportContext, -) - - -def test_export_command_requires_input_and_output(): - result = CliRunner().invoke(export_migration_data, []) - - assert result.exit_code != 0 - assert export_migration_data.name == "export-app-migration" - assert "--input" in result.output - assert "--output" in result.output - - -def test_import_command_requires_input_and_target_tenant_or_package_metadata(): - result = CliRunner().invoke(import_migration_data, []) - - assert result.exit_code != 0 - assert import_migration_data.name == "import-app-migration" - assert "--input" in result.output - - -def test_import_command_does_not_expose_unimplemented_map_id_strategy(): - assert ID_STRATEGY_CHOICES == ["preserve-id", "generate-new-id"] - - -def test_export_template_command_prints_scripted_json_template(): - result = CliRunner().invoke(export_migration_data_template, []) - - assert result.exit_code == 0 - assert export_migration_data_template.name == "app-migration-template" - template = json.loads(result.output) - assert template == { - "source_tenant": {"mode": "single", "id": "", "name": "admin's Workspace"}, - "apps": {"modes": ["workflow", "advanced-chat"], "ids": [], "all": True}, - "include_referenced_tools": True, - "additional_tools": {"api_tools": [], "workflow_tools": [], "mcp_tools": []}, - "include_secrets": False, - "import_options": { - "create_app_api_token_on_import": False, - "id_strategy": "preserve-id", - "conflict_strategy": "fail", - }, - } - - -def test_export_template_command_writes_output_file(tmp_path: Path): - output_file = tmp_path / "export-template.json" - - result = CliRunner().invoke(export_migration_data_template, ["--output", str(output_file)]) - - assert result.exit_code == 0 - assert f"Output written to {output_file}" in result.output - assert json.loads(output_file.read_text())["apps"]["all"] is True - - -def test_export_template_command_requires_overwrite_for_existing_output(tmp_path: Path): - output_file = tmp_path / "export-template.json" - output_file.write_text("{}") - - result = CliRunner().invoke(export_migration_data_template, ["--output", str(output_file)]) - - assert result.exit_code != 0 - assert "already exists" in result.output - - -def test_export_command_uses_cli_owned_session(monkeypatch, tmp_path: Path, sqlite_engine: Engine): - captured: dict[str, object] = {} - input_file = tmp_path / "export-config.json" - output_file = tmp_path / "migration-package.json" - input_file.write_text(json.dumps({"source_tenant": {"name": "source"}, "apps": {"all": True}})) - package = MigrationPackage.from_mapping({"metadata": {"version": "1", "source_scope": "single"}}) - - class FakeMigrationExportService: - def export(self, selection, *, session): - captured["session"] = session - captured["selection"] = selection - return ExportResult(package=package, report_items=[], report_context=ReportContext()) - - class FakeMigrationPackageService: - def save_package(self, package_to_save, path, *, overwrite): - captured["package"] = package_to_save - captured["path"] = path - captured["overwrite"] = overwrite - - monkeypatch.setattr(data_migration, "MigrationExportService", FakeMigrationExportService) - monkeypatch.setattr(data_migration, "MigrationPackageService", FakeMigrationPackageService) - - result = CliRunner().invoke( - export_migration_data, - ["--input", str(input_file), "--output", str(output_file)], - ) - - assert result.exit_code == 0 - captured_session = captured["session"] - assert isinstance(captured_session, Session) - assert captured_session.get_bind() is sqlite_engine - assert captured_session.in_transaction() is False - assert captured["package"] is package - assert captured["path"] == str(output_file) - assert captured["overwrite"] is False - - -def test_import_command_uses_cli_owned_session(monkeypatch, tmp_path: Path, sqlite_engine: Engine): - captured: dict[str, object] = {} - input_file = tmp_path / "migration-package.json" - input_file.write_text("{}") - package = MigrationPackage.from_mapping( - { - "metadata": { - "version": "1", - "source_scope": "single", - "target_tenant": {"name": "target"}, - "import_options": {"conflict_strategy": "fail"}, - } - } - ) - - class FakeMigrationImportService: - def import_package(self, request, *, session): - captured["session"] = session - captured["request"] = request - return ImportResult(report_items=[], report_context=ReportContext(target_tenant="target")) - - class FakeMigrationPackageService: - def load_package(self, path): - captured["path"] = path - return package - - monkeypatch.setattr(data_migration, "MigrationImportService", FakeMigrationImportService) - monkeypatch.setattr(data_migration, "MigrationPackageService", FakeMigrationPackageService) - - result = CliRunner().invoke( - import_migration_data, - ["--input", str(input_file), "--conflict-strategy", "skip"], - ) - - assert result.exit_code == 0 - captured_session = captured["session"] - assert isinstance(captured_session, Session) - assert captured_session.get_bind() is sqlite_engine - assert captured_session.in_transaction() is False - assert captured["path"] == str(input_file) - request = captured["request"] - assert request.package is package - assert request.options_override == ImportOptions(conflict_strategy=ConflictStrategy.SKIP) diff --git a/api/tests/unit_tests/commands/test_data_migration_wizard.py b/api/tests/unit_tests/commands/test_data_migration_wizard.py deleted file mode 100644 index 8d524cb914d..00000000000 --- a/api/tests/unit_tests/commands/test_data_migration_wizard.py +++ /dev/null @@ -1,407 +0,0 @@ -import pytest -from sqlalchemy.orm import Session - -from commands import data_migration -from commands.data_migration import ( - CONFLICT_STRATEGY_CHOICES, - ID_STRATEGY_CHOICES, - _confirm_wizard_summary, - _print_auto_tools, - _print_final_tool_selection, - _print_wizard_step, - _prompt_additional_tools, - _prompt_output_file, - _prompt_tool_category, - _resolve_mcp_tool_names, - migration_data_wizard, - parse_index_selection, -) -from models.tools import MCPToolProvider - - -def test_parse_index_selection_supports_all() -> None: - assert parse_index_selection("all", ["a", "b", "c"]) == ["a", "b", "c"] - - -def test_wizard_command_uses_app_migration_name() -> None: - assert migration_data_wizard.name == "app-migration-wizard" - - -def test_parse_index_selection_supports_comma_indexes() -> None: - assert parse_index_selection("1, 3", ["a", "b", "c"]) == ["a", "c"] - - -def test_print_wizard_step_adds_separator(monkeypatch: pytest.MonkeyPatch) -> None: - output_lines: list[str] = [] - - monkeypatch.setattr("commands.data_migration.click.echo", output_lines.append) - - _print_wizard_step("App Selection") - - assert output_lines == ["", "==== App Selection ===="] - - -def test_conflict_strategy_choices_exclude_replace() -> None: - assert CONFLICT_STRATEGY_CHOICES == ["fail", "skip", "update"] - - -def test_prompt_app_ids_explains_comma_selection_and_default(monkeypatch: pytest.MonkeyPatch) -> None: - from commands.data_migration import _prompt_app_ids - - prompts = [] - output_lines: list[str] = [] - apps = [ - type("App", (), {"id": "app-1", "name": "embedded", "mode": "workflow"})(), - type("App", (), {"id": "app-2", "name": "main", "mode": "advanced-chat"})(), - ] - - def capture_prompt(text: str, **kwargs: object) -> str: - prompts.append((text, kwargs)) - return "1,2" - - monkeypatch.setattr("commands.data_migration.click.echo", output_lines.append) - monkeypatch.setattr("commands.data_migration.click.prompt", capture_prompt) - - assert _prompt_app_ids(apps) == ["app-1", "app-2"] - assert prompts == [("Select apps by number, comma-separated numbers, or all", {"default": "all"})] - assert "Currently supported app types: workflow and chatflow." in output_lines - - -def test_prompt_tool_category_marks_auto_discovered_tools(monkeypatch: pytest.MonkeyPatch) -> None: - output_lines: list[str] = [] - - monkeypatch.setattr("commands.data_migration.click.echo", output_lines.append) - monkeypatch.setattr("commands.data_migration.click.prompt", lambda *_args, **_kwargs: "") - - selected = _prompt_tool_category( - "Custom API tools", - [("weather", "weather", "tool-id"), ("calendar", "calendar", "calendar-id")], - auto_tools={"weather": "tool-id"}, - ) - - assert selected == [] - assert "1. [auto] weather (tool-id)" in output_lines - assert "2. [ ] calendar (calendar-id)" in output_lines - assert output_lines[:2] == ["", "==== Custom API tools ===="] - - -def test_prompt_tool_category_explains_comma_selection_and_default(monkeypatch: pytest.MonkeyPatch) -> None: - prompts = [] - - def capture_prompt(text: str, **kwargs: object) -> str: - prompts.append((text, kwargs)) - return "" - - monkeypatch.setattr("commands.data_migration.click.echo", lambda *_args, **_kwargs: None) - monkeypatch.setattr("commands.data_migration.click.prompt", capture_prompt) - - selected = _prompt_tool_category( - "Custom API tools", - [("weather", "weather", "tool-id")], - auto_tools={}, - ) - - assert selected == [] - assert prompts == [ - ( - "Select custom api tools by number, comma-separated numbers, all, or empty", - {"default": "", "show_default": "empty"}, - ) - ] - - -def test_prompt_output_file_shows_default(monkeypatch: pytest.MonkeyPatch) -> None: - prompts = [] - - def capture_prompt(text: str, **kwargs: object) -> str: - prompts.append((text, kwargs)) - return "migration-data.json" - - monkeypatch.setattr("commands.data_migration.click.prompt", capture_prompt) - - assert _prompt_output_file() == ("migration-data.json", False) - assert prompts[0][0] == "Output path" - assert prompts[0][1]["show_default"] is True - - -def test_prompt_tool_category_marks_auto_by_detail_and_supports_multi_select(monkeypatch: pytest.MonkeyPatch) -> None: - monkeypatch.setattr("commands.data_migration.click.echo", lambda *_args, **_kwargs: None) - monkeypatch.setattr("commands.data_migration.click.prompt", lambda *_args, **_kwargs: "1,2") - - selected = _prompt_tool_category( - "Workflow tools", - [("tool-1", "embedded", "app-1"), ("tool-2", "other", "app-2")], - auto_tools={"embedded": "app-1"}, - ) - - assert selected == ["tool-1", "tool-2"] - - -def test_prompt_tool_category_marks_auto_by_value() -> None: - output_lines: list[str] = [] - - from commands import data_migration - - original_echo = data_migration.click.echo - original_prompt = data_migration.click.prompt - data_migration.click.echo = output_lines.append # pyrefly: ignore[bad-assignment] - data_migration.click.prompt = lambda *_args, **_kwargs: "" - try: - _prompt_tool_category( - "Workflow tools", - [("tool-1", "embedded_workflow_as_tool", "tool-1")], - auto_tools={"embedded_workflow_as_tool": "tool-1"}, - ) - finally: - data_migration.click.echo = original_echo - data_migration.click.prompt = original_prompt - - assert "1. [auto] embedded_workflow_as_tool (tool-1)" in output_lines - - -def test_print_auto_tools_lists_each_category(monkeypatch: pytest.MonkeyPatch) -> None: - output_lines: list[str] = [] - - monkeypatch.setattr("commands.data_migration.click.echo", output_lines.append) - - _print_auto_tools( - { - "api_tools": {"weather": "3bac3aa9-dd87-4351-9459-a7099137b028"}, - "workflow_tools": {"embedded_workflow_as_tool": "e6024578-41b7-4fb5-a81f-9201358e5835"}, - "mcp_tools": {}, - } - ) - - assert "Automatically discovered tools:" in output_lines - assert "Custom API tools" in output_lines - assert "- weather: 3bac3aa9-dd87-4351-9459-a7099137b028" in output_lines - assert "Workflow tools" in output_lines - assert "- embedded_workflow_as_tool: e6024578-41b7-4fb5-a81f-9201358e5835" in output_lines - assert "MCP tools" in output_lines - assert "- none" in output_lines - - -@pytest.mark.parametrize("sqlite_session", [(MCPToolProvider,)], indirect=True) -def test_resolve_mcp_tool_names_does_not_compare_non_uuid_identifier_to_uuid_id( - monkeypatch: pytest.MonkeyPatch, - sqlite_session: Session, -) -> None: - tenant_id = "49a99e46-bc2c-4885-91fa-47615f6192b5" - provider = MCPToolProvider( - name="resolved-provider", - server_identifier="my-test-mcp", - server_url="encrypted-url", - server_url_hash="server-url-hash", - icon=None, - tenant_id=tenant_id, - user_id="d13f1f2a-5c68-4a2f-8d76-cb59cefe95b8", - encrypted_credentials=None, - authed=False, - tools="[]", - ) - sqlite_session.add(provider) - sqlite_session.commit() - monkeypatch.setattr(data_migration.db, "session", sqlite_session) - - resolved = _resolve_mcp_tool_names(tenant_id, {"alias": provider.server_identifier}) - - assert resolved == {provider.name: provider.id} - - -def test_prompt_additional_tools_prints_final_selection_when_skipped(monkeypatch: pytest.MonkeyPatch) -> None: - output_lines: list[str] = [] - confirm_prompts: list[tuple[str, dict[str, object]]] = [] - - def capture_confirm(prompt: str, **kwargs: object) -> bool: - confirm_prompts.append((prompt, kwargs)) - return False - - monkeypatch.setattr("commands.data_migration.click.confirm", capture_confirm) - monkeypatch.setattr("commands.data_migration.click.echo", output_lines.append) - - selected = _prompt_additional_tools( - "tenant-id", - { - "api_tools": {"weather": "3bac3aa9-dd87-4351-9459-a7099137b028"}, - "workflow_tools": {}, - "mcp_tools": {}, - }, - ) - - assert selected == {"api_tools": [], "workflow_tools": [], "mcp_tools": []} - assert confirm_prompts == [ - ("Export additional tools manually? [y/n, default: n]", {"default": False, "show_default": False}) - ] - assert "Final tools to export:" in output_lines - assert "- [auto] weather: 3bac3aa9-dd87-4351-9459-a7099137b028" in output_lines - - -def test_final_tool_selection_deduplicates_manual_tool_already_auto(monkeypatch: pytest.MonkeyPatch) -> None: - output_lines: list[str] = [] - - monkeypatch.setattr("commands.data_migration.click.echo", output_lines.append) - - _print_final_tool_selection( - { - "api_tools": {}, - "workflow_tools": {"embedded_workflow_as_tool": "e6024578-41b7-4fb5-a81f-9201358e5835"}, - "mcp_tools": {}, - }, - { - "api_tools": [], - "workflow_tools": ["e6024578-41b7-4fb5-a81f-9201358e5835"], - "mcp_tools": [], - }, - {"e6024578-41b7-4fb5-a81f-9201358e5835": "embedded_workflow_as_tool: e6024578"}, - ) - - assert "- [auto] embedded_workflow_as_tool: e6024578-41b7-4fb5-a81f-9201358e5835" in output_lines - assert not any(line.startswith("- [manual]") for line in output_lines) - - -def test_prompt_output_file_rejects_yes_no_typo(monkeypatch: pytest.MonkeyPatch) -> None: - import click - import pytest - - monkeypatch.setattr("commands.data_migration.click.prompt", lambda *_args, **_kwargs: "y") - - with pytest.raises(click.ClickException, match="Output path must be a file path"): - _prompt_output_file() - - -def test_confirm_wizard_summary_shows_conflict_strategy(monkeypatch: pytest.MonkeyPatch) -> None: - output_lines: list[str] = [] - confirm_prompts: list[tuple[str, dict[str, object]]] = [] - - monkeypatch.setattr("commands.data_migration.click.echo", output_lines.append) - monkeypatch.setattr( - "commands.data_migration.click.confirm", - lambda prompt, **kwargs: confirm_prompts.append((prompt, kwargs)) or True, - ) - - _confirm_wizard_summary( - tenant_name="admin's Workspace", - app_names=["main_chatflow"], - auto_tools={"api_tools": {}, "workflow_tools": {}, "mcp_tools": {}}, - additional_tools={"api_tools": [], "workflow_tools": [], "mcp_tools": []}, - manual_labels={}, - include_referenced_tools=True, - include_secrets=False, - create_tokens=True, - id_strategy="preserve-id", - conflict_strategy="fail", - output_file="migration-data.json", - ) - - assert "id strategy: preserve-id" in output_lines - assert "conflict strategy: fail" in output_lines - assert confirm_prompts == [("Write migration package? [y/n, default: y]", {"default": True, "show_default": False})] - - -def test_confirm_wizard_summary_shows_final_deduplicated_tool_selection(monkeypatch: pytest.MonkeyPatch) -> None: - output_lines: list[str] = [] - - monkeypatch.setattr("commands.data_migration.click.echo", output_lines.append) - monkeypatch.setattr("commands.data_migration.click.confirm", lambda *_args, **_kwargs: True) - - _confirm_wizard_summary( - tenant_name="admin's Workspace", - app_names=["main_chatflow"], - auto_tools={ - "api_tools": {"weather": "weather-id"}, - "workflow_tools": {"embedded_workflow_as_tool": "workflow-tool-id"}, - "mcp_tools": {}, - }, - additional_tools={ - "api_tools": ["weather-id", "calendar"], - "workflow_tools": [], - "mcp_tools": ["mcp-id"], - }, - manual_labels={ - "calendar": "calendar: calendar-id", - "mcp-id": "my-test-mcp: mcp-id", - }, - include_referenced_tools=True, - include_secrets=False, - create_tokens=False, - id_strategy="preserve-id", - conflict_strategy="update", - output_file="migration-data.json", - ) - - assert "Final tools to export:" in output_lines - assert "Custom API tools" in output_lines - assert "- [auto] weather: weather-id" in output_lines - assert "- [manual] calendar: calendar-id" in output_lines - assert "Workflow tools" in output_lines - assert "- [auto] embedded_workflow_as_tool: workflow-tool-id" in output_lines - assert "MCP tools" in output_lines - assert "- [manual] my-test-mcp: mcp-id" in output_lines - assert not any(line.startswith("additional api tools:") for line in output_lines) - assert not any(line.startswith("additional workflow tools:") for line in output_lines) - assert not any(line.startswith("additional mcp tools:") for line in output_lines) - assert "- [manual] weather-id" not in output_lines - - -def test_import_options_prompts_explain_secrets_reuse_and_conflicts(monkeypatch: pytest.MonkeyPatch) -> None: - import click - - from commands.data_migration import _prompt_import_options - - output_lines: list[str] = [] - confirm_prompts: list[tuple[str, dict[str, object]]] = [] - prompt_calls = [] - - def capture_confirm(prompt: str, **kwargs: object) -> bool: - confirm_prompts.append((prompt, kwargs)) - return False - - def capture_prompt(prompt: str, **kwargs: object) -> object: - prompt_calls.append((prompt, kwargs)) - return kwargs["default"] - - monkeypatch.setattr("commands.data_migration.click.echo", output_lines.append) - monkeypatch.setattr("commands.data_migration.click.confirm", capture_confirm) - monkeypatch.setattr("commands.data_migration.click.prompt", capture_prompt) - - include_secrets, create_tokens, id_strategy, conflict_strategy = _prompt_import_options() - - assert include_secrets is False - assert create_tokens is False - assert id_strategy == "preserve-id" - assert conflict_strategy == "update" - assert "Secrets include workflow/app DSL secret values, custom API tool credentials," in output_lines - assert "-- Secrets --" in output_lines - assert "If you choose no, credentials are omitted or masked," in output_lines - assert "-- App API Tokens --" in output_lines - assert "When enabled, import will create an app API token if the imported app has none," in output_lines - assert "or reuse an existing app API token if one already exists." in output_lines - assert "-- ID Strategy --" in output_lines - assert "ID strategy controls whether imported app and tool IDs preserve source IDs" in output_lines - assert "or use target-generated IDs." in output_lines - assert "preserve-id: keep source IDs where the target service supports it." in output_lines - assert ( - "generate-new-id: let the target environment generate new IDs and rewrite references via mapping." - in output_lines - ) - assert "-- Conflict Strategy --" in output_lines - assert "Conflict strategy controls what import does when a target resource already exists." in output_lines - assert "fail: stop at the first conflict; previously committed resources are not rolled back." in output_lines - assert "skip: keep the existing target resource and skip importing that resource." in output_lines - assert "update: update the existing target resource in place." in output_lines - assert confirm_prompts == [ - ("Include secrets in output JSON? [y/n, default: n]", {"default": False, "show_default": False}), - ("Create or reuse app API tokens during import? [y/n, default: n]", {"default": False, "show_default": False}), - ] - assert prompt_calls[0][0] == "Import ID strategy. Enter one of: preserve-id, generate-new-id" - assert prompt_calls[0][1]["default"] == "preserve-id" - assert prompt_calls[0][1]["show_default"] is True - id_strategy_type = prompt_calls[0][1]["type"] - assert isinstance(id_strategy_type, click.Choice) - assert list(id_strategy_type.choices) == ID_STRATEGY_CHOICES - assert prompt_calls[1][0] == "Import conflict strategy. Enter one of: fail, skip, update" - assert prompt_calls[1][1]["default"] == "update" - assert prompt_calls[1][1]["show_default"] is True - conflict_strategy_type = prompt_calls[1][1]["type"] - assert isinstance(conflict_strategy_type, click.Choice) - assert list(conflict_strategy_type.choices) == CONFLICT_STRATEGY_CHOICES diff --git a/api/tests/unit_tests/commands/test_fix_app_site_missing.py b/api/tests/unit_tests/commands/test_fix_app_site_missing.py deleted file mode 100644 index 8666429fc97..00000000000 --- a/api/tests/unit_tests/commands/test_fix_app_site_missing.py +++ /dev/null @@ -1,139 +0,0 @@ -import uuid -from types import SimpleNamespace -from unittest.mock import MagicMock - -import pytest -from sqlalchemy import event -from sqlalchemy.engine import Engine -from sqlalchemy.orm import Session, scoped_session, sessionmaker - -from commands import system as system_commands -from models.account import Account, Tenant, TenantAccountJoin, TenantAccountRole -from models.enums import CustomizeTokenStrategy -from models.model import App, AppMode, IconType, Site - - -def _persist_missing_site_owner(session: Session) -> tuple[Account, App]: - """Persist an app without a Site and its complete tenant owner chain.""" - tenant = Tenant(name="Command workspace") - account = Account(name="Owner", email=f"owner-{uuid.uuid4()}@example.com") - membership = TenantAccountJoin( - tenant_id=tenant.id, - account_id=account.id, - current=True, - role=TenantAccountRole.OWNER, - ) - app = App( - id=str(uuid.uuid4()), - tenant_id=tenant.id, - name="Missing Site App", - mode=AppMode.CHAT, - icon_type=IconType.EMOJI, - icon="chat", - icon_background="#FFFFFF", - enable_site=True, - enable_api=False, - created_by=account.id, - ) - session.add_all([tenant, account, membership, app]) - session.commit() - return account, app - - -def _site_for(app: App) -> Site: - return Site( - app_id=app.id, - title=app.name, - default_language="en-US", - customize_token_strategy=CustomizeTokenStrategy.UUID, - code=f"site-{app.id}", - ) - - -def _bind_command_database( - monkeypatch: pytest.MonkeyPatch, - sqlite_engine: Engine, - sqlite_session_factory: sessionmaker[Session], -) -> tuple[scoped_session[Session], Session]: - """Expose a callable real scoped session through the Flask-SQLAlchemy shape.""" - command_sessions = scoped_session(sqlite_session_factory) - command_session = command_sessions() - monkeypatch.setattr( - system_commands, - "db", - SimpleNamespace(engine=sqlite_engine, session=command_sessions), - ) - return command_sessions, command_session - - -def _run_fix_app_site_missing() -> None: - """Invoke the command callback through a narrowed, non-optional reference.""" - callback = system_commands.fix_app_site_missing.callback - assert callback is not None - callback() - - -def test_fix_app_site_missing_passes_loaded_session_to_signal( - monkeypatch: pytest.MonkeyPatch, - sqlite_engine: Engine, - sqlite_session: Session, - sqlite_session_factory: sessionmaker[Session], -) -> None: - account, app = _persist_missing_site_owner(sqlite_session) - command_sessions, command_session = _bind_command_database(monkeypatch, sqlite_engine, sqlite_session_factory) - phase_events: list[str] = [] - event.listen(command_session, "after_commit", lambda _session: phase_events.append("commit")) - - def create_site(sender: App, *, account: Account, session: Session) -> None: - phase_events.append("signal") - assert sender.id == app.id - assert account.id == account_id - assert session is command_session - session.add(_site_for(sender)) - - account_id = account.id - send = MagicMock(side_effect=create_site) - monkeypatch.setattr(system_commands.app_was_created, "send", send) - - try: - _run_fix_app_site_missing() - finally: - command_sessions.remove() - - send.assert_called_once() - assert phase_events == ["signal", "commit"] - sqlite_session.expire_all() - persisted_site = sqlite_session.query(Site).filter_by(app_id=app.id).one() - assert persisted_site.title == app.name - - -def test_fix_app_site_missing_rolls_back_when_signal_fails( - monkeypatch: pytest.MonkeyPatch, - sqlite_engine: Engine, - sqlite_session: Session, - sqlite_session_factory: sessionmaker[Session], -) -> None: - _account, app = _persist_missing_site_owner(sqlite_session) - command_sessions, command_session = _bind_command_database(monkeypatch, sqlite_engine, sqlite_session_factory) - phase_events: list[str] = [] - event.listen(command_session, "after_rollback", lambda _session: phase_events.append("rollback")) - - def fail_signal(sender: App, **_kwargs: object) -> None: - phase_events.append("signal") - # Ensure the command's next raw scan terminates while its own transaction - # still exercises the rollback path. - with sqlite_session_factory() as observer: - observer.add(_site_for(sender)) - observer.commit() - raise RuntimeError("failed") - - monkeypatch.setattr(system_commands.app_was_created, "send", MagicMock(side_effect=fail_signal)) - - try: - _run_fix_app_site_missing() - finally: - command_sessions.remove() - - assert phase_events == ["signal", "rollback"] - sqlite_session.expire_all() - assert sqlite_session.get(App, app.id) is not None diff --git a/api/tests/unit_tests/commands/test_knowledge_fs_commands.py b/api/tests/unit_tests/commands/test_knowledge_fs.py similarity index 100% rename from api/tests/unit_tests/commands/test_knowledge_fs_commands.py rename to api/tests/unit_tests/commands/test_knowledge_fs.py diff --git a/api/tests/unit_tests/commands/test_legacy_model_type_migration.py b/api/tests/unit_tests/commands/test_legacy_model_type_migration.py deleted file mode 100644 index 70e5774d92d..00000000000 --- a/api/tests/unit_tests/commands/test_legacy_model_type_migration.py +++ /dev/null @@ -1,2384 +0,0 @@ -from __future__ import annotations - -import importlib -import io -import json -import os -import threading -import time -from collections.abc import Iterator, Sequence -from datetime import datetime, timedelta -from pathlib import Path -from types import ModuleType, SimpleNamespace -from typing import cast, override - -import pytest -import sqlalchemy as sa -from click.testing import CliRunner -from sqlalchemy import event -from sqlalchemy.exc import OperationalError -from sqlalchemy.orm import Session, SessionTransaction, sessionmaker - -from graphon.model_runtime.entities.model_entities import ModelType -from models import Dataset, DatasetPermission, DatasetPermissionEnum, TenantAccountJoin, TenantAccountRole -from models.account import Tenant -from models.base import TypeBase -from models.enums import CredentialSourceType -from models.provider import LoadBalancingModelConfig, ProviderModel -from services import legacy_model_type_migration as migration_module -from services.enterprise.rbac_service import ReplaceMemberBindings, ReplaceUserAccessPolicies -from tests.helpers.legacy_model_type_migration import ( - ALL_TABLE_NAMES, - LEGACY_TO_CANONICAL, - DirtyDataFixture, - assert_tenant_rows_use_only_canonical_model_types, - count_rows, - create_minimal_legacy_model_type_schema, - fetch_table_rows, - seed_legacy_model_type_dirty_data, - snapshot_legacy_model_type_state, -) - - -class _DatabaseError(Exception): - def __init__( - self, - *, - pgcode: str | None = None, - sqlstate: str | None = None, - errno: int | None = None, - ) -> None: - super().__init__() - self.pgcode = pgcode - self.sqlstate = sqlstate - self.errno = errno - - -@pytest.fixture -def dirty_fixture(sqlite_engine: sa.Engine) -> DirtyDataFixture: - return seed_legacy_model_type_dirty_data(sqlite_engine) - - -@pytest.fixture -def command_module() -> ModuleType: - try: - return importlib.import_module("commands.data_migration") - except ModuleNotFoundError: # pragma: no cover - explicit TDD failure path - pytest.fail( - "commands.data_migration is missing. " - "Implement the `flask data-migrate legacy-model-types` command group before running these tests." - ) - - -@pytest.fixture -def rbac_session(sqlite_engine: sa.Engine, monkeypatch: pytest.MonkeyPatch) -> Iterator[Session]: - """Bind RBAC command reads to persisted SQLite dataset rows.""" - - TypeBase.metadata.create_all( - sqlite_engine, - tables=[ - TypeBase.metadata.tables[model.__tablename__] for model in (Dataset, DatasetPermission, TenantAccountJoin) - ], - ) - factory = sessionmaker(bind=sqlite_engine, expire_on_commit=False) - monkeypatch.setattr("commands.rbac.session_factory.create_session", factory) - with factory() as session: - yield session - - -def _persist_dataset( - session: Session, - *, - dataset_id: str = "dataset-1", - tenant_id: str = "tenant-1", - permission: DatasetPermissionEnum = DatasetPermissionEnum.ONLY_ME, - created_by: str = "creator-account-1", - maintainer: str | None = None, -) -> Dataset: - dataset = Dataset( - id=dataset_id, - tenant_id=tenant_id, - name=f"Dataset {dataset_id}", - permission=permission, - created_by=created_by, - maintainer=maintainer, - ) - session.add(dataset) - session.commit() - return dataset - - -def _parse_json_lines(output: io.StringIO) -> list[dict[str, object]]: - return [json.loads(line) for line in output.getvalue().splitlines() if line.strip()] - - -def _json_key(value: object) -> str: - return json.dumps(value, sort_keys=True) - - -def _event_signature(line: dict[str, object]) -> tuple[object, ...] | None: - event = line.get("event") - attrs = line.get("attrs") - if not isinstance(attrs, dict): - return None - - if event == "row_updated": - return ( - event, - attrs.get("table_name"), - attrs.get("id"), - _json_key(attrs.get("business_key")), - _json_key(attrs.get("old_values")), - _json_key(attrs.get("new_values")), - _json_key(attrs.get("rewrite_source")), - ) - if event == "row_deleted": - return ( - event, - attrs.get("table_name"), - attrs.get("id"), - _json_key(attrs.get("business_key")), - attrs.get("merge_winner_id"), - ) - if event == "group_processed": - return ( - event, - attrs.get("table_name"), - _json_key(attrs.get("business_key")), - tuple(cast(list[str], attrs.get("group_row_ids", []))), - ) - return None - - -def _collect_processing_signatures(lines: list[dict[str, object]]) -> set[tuple[object, ...]]: - signatures: set[tuple[object, ...]] = set() - for line in lines: - signature = _event_signature(line) - if signature is not None: - signatures.add(signature) - return signatures - - -def _cache_event_row_ids( - lines: list[dict[str, object]], - *, - table_name: str, - row_ids: set[str], - event_name: str, -) -> set[str]: - matching_row_ids: set[str] = set() - for line in lines: - if line.get("event") != event_name: - continue - attrs = line.get("attrs") - if not isinstance(attrs, dict): - continue - if attrs.get("table_name") != table_name: - continue - row_id = str(attrs.get("id")) - if row_id in row_ids: - matching_row_ids.add(row_id) - return matching_row_ids - - -def _patch_batch_size( - monkeypatch: pytest.MonkeyPatch, - *, - batch_size: int, -) -> None: - original_init = migration_module.Migration.__init__ - - def _patched_init( - self: migration_module.Migration, - tenant_id: str, - engine: sa.Engine, - apply: bool, - output: io.TextIOBase, - model_types: Sequence[ModelType], - orm_models: Sequence[migration_module.ORMModel], - ) -> None: - original_init(self, tenant_id, engine, apply, output, model_types, orm_models) - self._batch_size = batch_size - - monkeypatch.setattr(migration_module.Migration, "__init__", _patched_init) - - -def _insert_provider_model( - engine: sa.Engine, - *, - row_id: str, - tenant_id: str, - provider_name: str, - model_name: str, - model_type: str, - credential_id: str | None, - created_at: datetime, - updated_at: datetime, -) -> None: - with engine.begin() as conn: - conn.execute( - sa.text( - """ - INSERT INTO provider_models - ( - id, tenant_id, provider_name, model_name, model_type, - credential_id, is_valid, created_at, updated_at - ) - VALUES - ( - :id, :tenant_id, :provider_name, :model_name, :model_type, - :credential_id, :is_valid, :created_at, :updated_at - ) - """ - ), - { - "id": row_id, - "tenant_id": tenant_id, - "provider_name": provider_name, - "model_name": model_name, - "model_type": model_type, - "credential_id": credential_id, - "is_valid": True, - "created_at": created_at, - "updated_at": updated_at, - }, - ) - - -def _insert_tenant(engine: sa.Engine, *, tenant_id: str) -> None: - with engine.begin() as conn: - conn.execute( - sa.insert(Tenant).values( - id=tenant_id, - name=f"Tenant {tenant_id}", - plan="basic", - status="normal", - ) - ) - - -def _insert_tenant_default_model( - engine: sa.Engine, - *, - row_id: str, - tenant_id: str, - provider_name: str, - model_name: str, - model_type: str, - created_at: datetime, - updated_at: datetime, -) -> None: - with engine.begin() as conn: - conn.execute( - sa.text( - """ - INSERT INTO tenant_default_models - (id, tenant_id, provider_name, model_name, model_type, created_at, updated_at) - VALUES - (:id, :tenant_id, :provider_name, :model_name, :model_type, :created_at, :updated_at) - """ - ), - { - "id": row_id, - "tenant_id": tenant_id, - "provider_name": provider_name, - "model_name": model_name, - "model_type": model_type, - "created_at": created_at, - "updated_at": updated_at, - }, - ) - - -def _insert_provider_model_setting( - engine: sa.Engine, - *, - row_id: str, - tenant_id: str, - provider_name: str, - model_name: str, - model_type: str, - enabled: bool, - load_balancing_enabled: bool, - created_at: datetime, - updated_at: datetime, -) -> None: - with engine.begin() as conn: - conn.execute( - sa.text( - """ - INSERT INTO provider_model_settings - ( - id, tenant_id, provider_name, model_name, model_type, - enabled, load_balancing_enabled, - created_at, updated_at - ) - VALUES - ( - :id, :tenant_id, :provider_name, :model_name, :model_type, - :enabled, :load_balancing_enabled, - :created_at, :updated_at - ) - """ - ), - { - "id": row_id, - "tenant_id": tenant_id, - "provider_name": provider_name, - "model_name": model_name, - "model_type": model_type, - "enabled": enabled, - "load_balancing_enabled": load_balancing_enabled, - "created_at": created_at, - "updated_at": updated_at, - }, - ) - - -def _insert_load_balancing_model_config( - engine: sa.Engine, - *, - row_id: str, - tenant_id: str, - provider_name: str, - model_name: str, - model_type: str, - name: str, - encrypted_config: str, - credential_id: str, - enabled: bool, - created_at: datetime, - updated_at: datetime, -) -> None: - with engine.begin() as conn: - conn.execute( - sa.text( - """ - INSERT INTO load_balancing_model_configs - ( - id, tenant_id, provider_name, model_name, model_type, name, - encrypted_config, credential_id, credential_source_type, enabled, created_at, updated_at - ) - VALUES - ( - :id, :tenant_id, :provider_name, :model_name, :model_type, :name, - :encrypted_config, :credential_id, :credential_source_type, :enabled, :created_at, :updated_at - ) - """ - ), - { - "id": row_id, - "tenant_id": tenant_id, - "provider_name": provider_name, - "model_name": model_name, - "model_type": model_type, - "name": name, - "encrypted_config": encrypted_config, - "credential_id": credential_id, - "credential_source_type": CredentialSourceType.CUSTOM_MODEL.value, - "enabled": enabled, - "created_at": created_at, - "updated_at": updated_at, - }, - ) - - -def test_data_migrate_group_registers_dataset_permission_rbac_migration(command_module: ModuleType) -> None: - command = command_module.data_migrate.commands["rbac-migrate-dataset-permissions"] - - assert command is command_module.migrate_dataset_permissions_to_rbac - assert "operator_account_id" not in {param.name for param in command.params} - - -def test_data_migrate_group_registers_resource_whitelist_scope_migration(command_module: ModuleType) -> None: - command = command_module.data_migrate.commands["rbac-migrate-resource-whitelist-scopes"] - - assert command is command_module.migrate_only_me_resource_whitelist_scopes_to_automatic_include - - -def test_dataset_permission_rbac_migration_help_mentions_binding_clear_side_effect(command_module: ModuleType) -> None: - result = CliRunner().invoke( - command_module.data_migrate, - ["rbac-migrate-dataset-permissions", "--help"], - ) - - assert result.exit_code == 0 - normalized_output = " ".join(result.output.split()) - assert "clears existing per-user policy bindings" in normalized_output - assert "recreates legacy partial-member default bindings" in normalized_output - - -def test_dataset_permission_rbac_migration_maps_legacy_permissions_to_enum_scopes() -> None: - rbac_module = importlib.import_module("commands.rbac") - - assert ( - rbac_module._rbac_dataset_scope_for_legacy_permission(rbac_module.DatasetPermissionEnum.ALL_TEAM) - is rbac_module.RBACResourceWhitelistScope.ALL - ) - assert ( - rbac_module._rbac_dataset_scope_for_legacy_permission(rbac_module.DatasetPermissionEnum.PARTIAL_TEAM) - is rbac_module.RBACResourceWhitelistScope.SPECIFIC - ) - assert rbac_module._dataset_permission_enum("partial_members") is rbac_module.DatasetPermissionEnum.PARTIAL_TEAM - assert rbac_module._dataset_permission_enum(None) is rbac_module.DatasetPermissionEnum.ONLY_ME - - -def test_dataset_permission_rbac_migration_uses_dataset_creator_as_operator( - command_module: ModuleType, - rbac_session: Session, - monkeypatch: pytest.MonkeyPatch, -) -> None: - rbac_module = importlib.import_module("commands.rbac") - _persist_dataset(rbac_session) - calls: list[dict[str, object]] = [] - read_transaction_ended = False - - def fake_replace_whitelist(**kwargs: object) -> None: - assert read_transaction_ended - calls.append(kwargs) - - def _record_transaction_end(session: Session, transaction: object) -> None: - nonlocal read_transaction_ended - del transaction - if session.get_bind() is rbac_session.get_bind(): - read_transaction_ended = True - - event.listen(Session, "after_transaction_end", _record_transaction_end) - monkeypatch.setattr(rbac_module.RBACService.DatasetAccess, "replace_whitelist", fake_replace_whitelist) - try: - command_module.migrate_dataset_permissions_to_rbac.callback( - tenant_id=None, - dataset_id=None, - batch_size=500, - dry_run=False, - ) - finally: - event.remove(Session, "after_transaction_end", _record_transaction_end) - - assert calls[0]["tenant_id"] == "tenant-1" - assert calls[0]["account_id"] == "creator-account-1" - assert calls[0]["dataset_id"] == "dataset-1" - payload = calls[0]["payload"] - assert isinstance(payload, ReplaceMemberBindings) - assert payload.automatic_include_workspace_members is False - - -def test_dataset_permission_rbac_migration_dry_run_outputs_structured_proposed_changes( - command_module: ModuleType, - rbac_session: Session, - monkeypatch: pytest.MonkeyPatch, -) -> None: - rbac_module = importlib.import_module("commands.rbac") - dataset = _persist_dataset(rbac_session, permission=DatasetPermissionEnum.PARTIAL_TEAM) - rbac_session.add( - DatasetPermission( - dataset_id=dataset.id, - account_id="member-account-1", - tenant_id=dataset.tenant_id, - ) - ) - rbac_session.commit() - monkeypatch.setattr( - rbac_module.RBACService.DatasetAccess, - "replace_whitelist", - lambda **kwargs: pytest.fail("dry-run must not replace whitelist"), - ) - monkeypatch.setattr( - rbac_module.RBACService.DatasetAccess, - "replace_user_access_policies", - lambda **kwargs: pytest.fail("dry-run must not replace user access policies"), - ) - - result = CliRunner().invoke( - command_module.data_migrate, - ["rbac-migrate-dataset-permissions", "--dry-run"], - ) - - assert result.exit_code == 0 - events = [json.loads(line) for line in result.output.splitlines() if line.startswith("{")] - assert [event["action"] for event in events] == ["replace_whitelist", "replace_user_access_policies"] - assert events[0]["before"] == { - "legacy_dataset_permission": "partial_members", - "legacy_partial_member_ids": ["member-account-1"], - } - assert events[0]["after"] == {"rbac_whitelist_scope": "specific"} - assert events[0]["call"] == { - "method": "RBACService.DatasetAccess.replace_whitelist", - "kwargs": { - "tenant_id": "tenant-1", - "account_id": "creator-account-1", - "dataset_id": "dataset-1", - "payload": {"automatic_include_workspace_members": False}, - }, - } - assert events[1]["target_account_id"] == "member-account-1" - assert events[1]["after"] == {"rbac_user_access_policy_ids": ["default"]} - assert events[1]["call"]["kwargs"]["payload"] == {"access_policy_ids": ["default"]} - - -def test_resource_whitelist_scope_migration_specific_preserves_existing_members( - command_module: ModuleType, - rbac_session: Session, - monkeypatch: pytest.MonkeyPatch, -) -> None: - rbac_module = importlib.import_module("commands.rbac") - _persist_dataset(rbac_session, maintainer="maintainer-account-1") - replace_whitelist_calls: list[dict[str, object]] = [] - replace_policy_calls: list[dict[str, object]] = [] - - monkeypatch.setattr( - rbac_module.RBACService.DatasetAccess, - "legacy_whitelist_config", - lambda **kwargs: SimpleNamespace( - rbac_whitelist_scope="specific", - account_ids=["member-account-2", "member-account-1", "member-account-1"], - ), - ) - monkeypatch.setattr( - rbac_module.RBACService.DatasetAccess, - "replace_whitelist", - lambda **kwargs: replace_whitelist_calls.append(kwargs), - ) - monkeypatch.setattr( - rbac_module.RBACService.DatasetAccess, - "replace_user_access_policies", - lambda **kwargs: replace_policy_calls.append(kwargs), - ) - monkeypatch.setattr(rbac_module.RBACService.MemberRoles, "batch_get", lambda **kwargs: []) - - command_module.migrate_resource_whitelist_scopes_to_automatic_include.callback( - tenant_id=None, - resource_type="dataset", - resource_id=None, - batch_size=500, - member_batch_size=500, - dry_run=False, - ) - - whitelist_payload = replace_whitelist_calls[0]["payload"] - policy_payload = replace_policy_calls[0]["payload"] - assert isinstance(whitelist_payload, ReplaceMemberBindings) - assert isinstance(policy_payload, ReplaceUserAccessPolicies) - assert whitelist_payload.automatic_include_workspace_members is False - assert policy_payload.access_policy_ids == ["default"] - assert policy_payload.account_ids == ["member-account-1", "member-account-2"] - - -def test_resource_whitelist_scope_migration_all_syncs_workspace_members( - command_module: ModuleType, - rbac_session: Session, - monkeypatch: pytest.MonkeyPatch, -) -> None: - rbac_module = importlib.import_module("commands.rbac") - _persist_dataset(rbac_session, maintainer="maintainer-account-1") - rbac_session.add_all( - [ - TenantAccountJoin( - tenant_id="tenant-1", - account_id="maintainer-account-1", - role=TenantAccountRole.OWNER, - ), - TenantAccountJoin( - tenant_id="tenant-1", - account_id="member-account-1", - role=TenantAccountRole.NORMAL, - ), - ] - ) - rbac_session.commit() - replace_whitelist_calls: list[dict[str, object]] = [] - replace_policy_calls: list[dict[str, object]] = [] - - monkeypatch.setattr( - rbac_module.RBACService.DatasetAccess, - "legacy_whitelist_config", - lambda **kwargs: SimpleNamespace(rbac_whitelist_scope="all", account_ids=[]), - ) - monkeypatch.setattr( - rbac_module.RBACService.DatasetAccess, - "replace_whitelist", - lambda **kwargs: replace_whitelist_calls.append(kwargs), - ) - monkeypatch.setattr( - rbac_module.RBACService.DatasetAccess, - "replace_user_access_policies", - lambda **kwargs: replace_policy_calls.append(kwargs), - ) - monkeypatch.setattr(rbac_module.RBACService.MemberRoles, "batch_get", lambda **kwargs: []) - - command_module.migrate_resource_whitelist_scopes_to_automatic_include.callback( - tenant_id=None, - resource_type="dataset", - resource_id=None, - batch_size=500, - member_batch_size=500, - dry_run=False, - ) - - whitelist_payload = replace_whitelist_calls[0]["payload"] - policy_payload = replace_policy_calls[0]["payload"] - assert isinstance(whitelist_payload, ReplaceMemberBindings) - assert isinstance(policy_payload, ReplaceUserAccessPolicies) - assert whitelist_payload.automatic_include_workspace_members is True - assert policy_payload.access_policy_ids == ["default"] - assert set(policy_payload.account_ids) == {"maintainer-account-1", "member-account-1"} - - -def test_only_me_resource_whitelist_scope_migration_syncs_workspace_members( - command_module: ModuleType, - rbac_session: Session, - monkeypatch: pytest.MonkeyPatch, -) -> None: - rbac_module = importlib.import_module("commands.rbac") - _persist_dataset(rbac_session, maintainer="maintainer-account-1") - rbac_session.add_all( - [ - TenantAccountJoin( - tenant_id="tenant-1", - account_id="maintainer-account-1", - role=TenantAccountRole.OWNER, - ), - TenantAccountJoin( - tenant_id="tenant-1", - account_id="member-account-1", - role=TenantAccountRole.NORMAL, - ), - ] - ) - rbac_session.commit() - replace_whitelist_calls: list[dict[str, object]] = [] - replace_policy_calls: list[dict[str, object]] = [] - - monkeypatch.setattr( - rbac_module.RBACService.DatasetAccess, - "legacy_whitelist_config", - lambda **kwargs: SimpleNamespace(rbac_whitelist_scope="only_me", account_ids=[]), - ) - monkeypatch.setattr( - rbac_module.RBACService.DatasetAccess, - "replace_whitelist", - lambda **kwargs: replace_whitelist_calls.append(kwargs), - ) - monkeypatch.setattr( - rbac_module.RBACService.DatasetAccess, - "replace_user_access_policies", - lambda **kwargs: replace_policy_calls.append(kwargs), - ) - monkeypatch.setattr(rbac_module.RBACService.MemberRoles, "batch_get", lambda **kwargs: []) - - command_module.migrate_only_me_resource_whitelist_scopes_to_automatic_include.callback( - tenant_id=None, - resource_type="dataset", - resource_id=None, - batch_size=500, - member_batch_size=500, - dry_run=False, - ) - - whitelist_payload = replace_whitelist_calls[0]["payload"] - policy_payload = replace_policy_calls[0]["payload"] - assert isinstance(whitelist_payload, ReplaceMemberBindings) - assert isinstance(policy_payload, ReplaceUserAccessPolicies) - assert whitelist_payload.automatic_include_workspace_members is True - assert policy_payload.access_policy_ids == ["default"] - assert set(policy_payload.account_ids) == {"maintainer-account-1", "member-account-1"} - - -def test_data_migrate_command_defaults_output_to_stdout_stream( - command_module: ModuleType, - monkeypatch: pytest.MonkeyPatch, - tmp_path: Path, -) -> None: - service_calls: list[dict[str, object]] = [] - fake_stdout = io.StringIO() - - class FakeService: - def __init__( - self, - *, - engine: sa.Engine, - apply: bool, - concurrency: int, - output: io.TextIOBase | None = None, - tables: tuple[str, ...] | None, - model_types: tuple[ModelType, ...], - tenant_ids: tuple[str, ...] | None, - ) -> None: - service_calls.append( - { - "engine": engine, - "apply": apply, - "concurrency": concurrency, - "output": output, - "tables": tables, - "model_types": model_types, - "tenant_ids": tenant_ids, - } - ) - - def migrate(self) -> None: - service_calls.append({"migrated": True}) - - monkeypatch.setattr(command_module, "LegacyModelTypeMigrationService", FakeService) - monkeypatch.setattr(command_module, "db", SimpleNamespace(engine=object())) - monkeypatch.setattr(command_module.sys, "stdout", fake_stdout) - tenant_id_file = tmp_path / "tenant_ids.txt" - tenant_id_file.write_text("tenant-alpha\n", encoding="utf-8") - - data_migrate = command_module.data_migrate - legacy_model_types = data_migrate.commands["legacy-model-types"] - assert legacy_model_types.callback is not None - - legacy_model_types.callback( - apply=True, - tables=("provider_models",), - model_types=("llm", "text-embedding"), - tenant_id_file=str(tenant_id_file), - output=None, - concurrency=7, - ) - - assert service_calls[0]["apply"] is True - assert service_calls[0]["concurrency"] == 7 - assert service_calls[0]["output"] is fake_stdout - assert service_calls[0]["tables"] == ("provider_models",) - assert tuple(cast(list[str], service_calls[0]["tenant_ids"])) == ("tenant-alpha",) - assert service_calls[0]["model_types"] == (ModelType.LLM, ModelType.TEXT_EMBEDDING) - assert service_calls[1] == {"migrated": True} - - -def test_data_migrate_command_opens_output_file_and_closes_stream( - command_module: ModuleType, - monkeypatch: pytest.MonkeyPatch, - tmp_path: Path, -) -> None: - service_calls: list[dict[str, object]] = [] - - class FakeService: - def __init__( - self, - *, - engine: sa.Engine, - apply: bool, - concurrency: int, - output: io.TextIOBase | None = None, - tables: tuple[str, ...] | None, - model_types: tuple[ModelType, ...], - tenant_ids: tuple[str, ...] | None, - ) -> None: - service_calls.append( - { - "engine": engine, - "apply": apply, - "concurrency": concurrency, - "output": output, - "tables": tables, - "model_types": model_types, - "tenant_ids": tenant_ids, - } - ) - - def migrate(self) -> None: - output = cast(io.TextIOBase, service_calls[0]["output"]) - output.write('{"event":"test"}\n') - service_calls.append({"migrated": True}) - - monkeypatch.setattr(command_module, "LegacyModelTypeMigrationService", FakeService) - monkeypatch.setattr(command_module, "db", SimpleNamespace(engine=object())) - output_path = tmp_path / "migration.jsonl" - - data_migrate = command_module.data_migrate - legacy_model_types = data_migrate.commands["legacy-model-types"] - assert legacy_model_types.callback is not None - - legacy_model_types.callback( - apply=False, - tables=(), - model_types=(), - tenant_id_file=None, - output=output_path, - concurrency=3, - ) - - output_stream = cast(io.TextIOBase, service_calls[0]["output"]) - assert service_calls[0]["concurrency"] == 3 - assert output_stream is not output_path - assert isinstance(output_stream, io.TextIOWrapper) - assert Path(output_stream.name) == output_path - assert output_stream.closed is True - assert output_path.read_text(encoding="utf-8") == '{"event":"test"}\n' - assert service_calls[1] == {"migrated": True} - - -@pytest.mark.parametrize( - ("cpu_count", "expected_concurrency"), - [ - (8, 8), - (None, 1), - ], -) -def test_data_migrate_command_defaults_concurrency_from_cpu_count_or_falls_back_to_one( - monkeypatch: pytest.MonkeyPatch, - cpu_count: int | None, - expected_concurrency: int, -) -> None: - service_calls: list[dict[str, object]] = [] - command_module = importlib.import_module("commands.data_migration") - - class FakeService: - def __init__( - self, - *, - engine: sa.Engine, - apply: bool, - concurrency: int, - output: io.TextIOBase | None = None, - tables: tuple[str, ...] | None, - model_types: tuple[ModelType, ...], - tenant_ids: tuple[str, ...] | None, - ) -> None: - service_calls.append( - { - "engine": engine, - "apply": apply, - "concurrency": concurrency, - "output": output, - "tables": tables, - "model_types": model_types, - "tenant_ids": tenant_ids, - } - ) - - def migrate(self) -> None: - service_calls.append({"migrated": True}) - - monkeypatch.setattr(os, "cpu_count", lambda: cpu_count) - importlib.reload(command_module) - try: - monkeypatch.setattr(command_module, "LegacyModelTypeMigrationService", FakeService) - monkeypatch.setattr(command_module, "db", SimpleNamespace(engine=object())) - - result = CliRunner().invoke(command_module.data_migrate, ["legacy-model-types"]) - - assert result.exit_code == 0, result.output - assert expected_concurrency == command_module._DEFAULT_CONCURRENCY - assert service_calls[0]["concurrency"] == expected_concurrency - assert service_calls[1] == {"migrated": True} - finally: - monkeypatch.undo() - importlib.reload(command_module) - - -def test_service_migrate_batches_by_tenant_respects_selected_tables_without_reverse_dependency_expansion( - sqlite_engine: sa.Engine, -) -> None: - seen_runs: list[tuple[str, tuple[str, ...], tuple[ModelType, ...]]] = [] - - class FakeMigration: - def __init__( - self, - *, - tenant_id: str, - engine: sa.Engine, - apply: bool, - output: io.TextIOBase, - model_types: tuple[ModelType, ...], - orm_models: tuple[migration_module.ORMModel, ...], - ) -> None: - assert engine is sqlite_engine - assert apply is False - seen_runs.append((tenant_id, tuple(model.__tablename__ for model in orm_models), model_types)) - - def run(self) -> None: - return None - - monkeypatch = pytest.MonkeyPatch() - try: - monkeypatch.setattr(migration_module, "Migration", FakeMigration) - service = migration_module.LegacyModelTypeMigrationService( - engine=sqlite_engine, - apply=False, - concurrency=1, - tables=("provider_models", "tenant_default_models"), - model_types=(ModelType.LLM,), - tenant_ids=("tenant-alpha", "tenant-beta"), - ) - - service.migrate() - finally: - monkeypatch.undo() - - assert seen_runs == [ - ("tenant-alpha", ("provider_models", "tenant_default_models"), (ModelType.LLM,)), - ("tenant-beta", ("provider_models", "tenant_default_models"), (ModelType.LLM,)), - ] - - -def test_service_migrate_without_tenant_ids_discovers_tenants_per_selected_table_without_querying_tenants( - sqlite_engine: sa.Engine, - monkeypatch: pytest.MonkeyPatch, -) -> None: - create_minimal_legacy_model_type_schema(sqlite_engine) - provider_tenant_id = "00000000-0000-0000-0000-000000000111" - default_tenant_id = "00000000-0000-0000-0000-000000000222" - empty_tenant_id = "00000000-0000-0000-0000-000000000333" - for tenant_id in (provider_tenant_id, default_tenant_id, empty_tenant_id): - _insert_tenant(sqlite_engine, tenant_id=tenant_id) - - created_at = datetime(2025, 1, 1, 12, 0, 0) - updated_at = created_at + timedelta(minutes=1) - _insert_provider_model( - sqlite_engine, - row_id="10000000-0000-0000-0000-000000000111", - tenant_id=provider_tenant_id, - provider_name="openai", - model_name="gpt-4o-mini", - model_type="text-generation", - credential_id=None, - created_at=created_at, - updated_at=updated_at, - ) - _insert_tenant_default_model( - sqlite_engine, - row_id="20000000-0000-0000-0000-000000000222", - tenant_id=default_tenant_id, - provider_name="openai", - model_name="gpt-4o-mini", - model_type="text-generation", - created_at=created_at, - updated_at=updated_at, - ) - - seen_runs: list[tuple[str, tuple[str, ...], tuple[ModelType, ...]]] = [] - executed_sql: list[str] = [] - - class FakeMigration: - def __init__( - self, - *, - tenant_id: str, - engine: sa.Engine, - apply: bool, - output: io.TextIOBase, - model_types: tuple[ModelType, ...], - orm_models: tuple[migration_module.ORMModel, ...], - ) -> None: - assert engine is sqlite_engine - assert apply is False - seen_runs.append((tenant_id, tuple(model.__tablename__ for model in orm_models), model_types)) - - def run(self) -> None: - return None - - def _record_sql( - conn: sa.engine.Connection, - cursor: object, - statement: str, - parameters: object, - context: object, - executemany: bool, - ) -> None: - del conn, cursor, parameters, context, executemany - executed_sql.append(statement) - - event.listen(sqlite_engine, "before_cursor_execute", _record_sql) - try: - monkeypatch.setattr(migration_module, "Migration", FakeMigration) - service = migration_module.LegacyModelTypeMigrationService( - engine=sqlite_engine, - apply=False, - tables=("provider_models", "tenant_default_models"), - model_types=(ModelType.LLM,), - ) - - service.migrate() - finally: - event.remove(sqlite_engine, "before_cursor_execute", _record_sql) - - assert seen_runs == [ - (provider_tenant_id, ("provider_models",), (ModelType.LLM,)), - (default_tenant_id, ("tenant_default_models",), (ModelType.LLM,)), - ] - normalized_statements = [" ".join(statement.lower().split()) for statement in executed_sql] - discovery_statements = [statement for statement in normalized_statements if statement.startswith("select")] - table_names = ("provider_models", "tenant_default_models") - table_discovery_statements = [ - statement - for statement in discovery_statements - if any(f" from {table_name} " in f" {statement} " for table_name in table_names) - ] - - assert [statement for statement in discovery_statements if " from tenants " in f" {statement} "] == [] - assert [statement for statement in discovery_statements if " union " in f" {statement} "] == [] - assert [ - next(table_name for table_name in table_names if f" from {table_name} " in f" {statement} ") - for statement in table_discovery_statements - ] == list(table_names) - - -def test_service_migrate_without_tenant_ids_filters_provider_model_tenants_by_selected_model_types( - sqlite_engine: sa.Engine, - monkeypatch: pytest.MonkeyPatch, -) -> None: - create_minimal_legacy_model_type_schema(sqlite_engine) - llm_tenant_id = "00000000-0000-0000-0000-000000000411" - embedding_tenant_id = "00000000-0000-0000-0000-000000000422" - empty_tenant_id = "00000000-0000-0000-0000-000000000433" - for tenant_id in (llm_tenant_id, embedding_tenant_id, empty_tenant_id): - _insert_tenant(sqlite_engine, tenant_id=tenant_id) - - created_at = datetime(2025, 1, 2, 12, 0, 0) - updated_at = created_at + timedelta(minutes=1) - _insert_provider_model( - sqlite_engine, - row_id="30000000-0000-0000-0000-000000000411", - tenant_id=llm_tenant_id, - provider_name="openai", - model_name="gpt-4o-mini", - model_type="text-generation", - credential_id=None, - created_at=created_at, - updated_at=updated_at, - ) - _insert_provider_model( - sqlite_engine, - row_id="30000000-0000-0000-0000-000000000422", - tenant_id=embedding_tenant_id, - provider_name="openai", - model_name="text-embedding-3-large", - model_type="embeddings", - credential_id=None, - created_at=created_at, - updated_at=updated_at, - ) - - seen_runs: list[tuple[str, tuple[str, ...], tuple[ModelType, ...]]] = [] - - class FakeMigration: - def __init__( - self, - *, - tenant_id: str, - engine: sa.Engine, - apply: bool, - output: io.TextIOBase, - model_types: tuple[ModelType, ...], - orm_models: tuple[migration_module.ORMModel, ...], - ) -> None: - assert engine is sqlite_engine - assert apply is False - seen_runs.append((tenant_id, tuple(model.__tablename__ for model in orm_models), model_types)) - - def run(self) -> None: - return None - - monkeypatch.setattr(migration_module, "Migration", FakeMigration) - service = migration_module.LegacyModelTypeMigrationService( - engine=sqlite_engine, - apply=False, - tables=("provider_models",), - model_types=(ModelType.LLM,), - ) - - service.migrate() - - assert seen_runs == [ - (llm_tenant_id, ("provider_models",), (ModelType.LLM,)), - ] - - -def test_service_migrate_without_tenant_ids_discovers_all_load_balancing_tenants_for_simpler_table_scoped_query( - sqlite_engine: sa.Engine, - monkeypatch: pytest.MonkeyPatch, -) -> None: - create_minimal_legacy_model_type_schema(sqlite_engine) - inherit_llm_tenant_id = "00000000-0000-0000-0000-000000000511" - inherit_embedding_tenant_id = "00000000-0000-0000-0000-000000000522" - empty_tenant_id = "00000000-0000-0000-0000-000000000533" - for tenant_id in (inherit_llm_tenant_id, inherit_embedding_tenant_id, empty_tenant_id): - _insert_tenant(sqlite_engine, tenant_id=tenant_id) - - created_at = datetime(2025, 1, 3, 12, 0, 0) - updated_at = created_at + timedelta(minutes=1) - _insert_load_balancing_model_config( - sqlite_engine, - row_id="40000000-0000-0000-0000-000000000511", - tenant_id=inherit_llm_tenant_id, - provider_name="openai", - model_name="gpt-4o-mini", - model_type=ModelType.LLM.value, - name="__inherit__", - encrypted_config=json.dumps({"api_key": "inherit-llm"}), - credential_id="50000000-0000-0000-0000-000000000511", - enabled=True, - created_at=created_at, - updated_at=updated_at, - ) - _insert_load_balancing_model_config( - sqlite_engine, - row_id="40000000-0000-0000-0000-000000000522", - tenant_id=inherit_embedding_tenant_id, - provider_name="openai", - model_name="text-embedding-3-large", - model_type=ModelType.TEXT_EMBEDDING.value, - name="__inherit__", - encrypted_config=json.dumps({"api_key": "inherit-embedding"}), - credential_id="50000000-0000-0000-0000-000000000522", - enabled=True, - created_at=created_at, - updated_at=updated_at, - ) - - seen_runs: list[tuple[str, tuple[str, ...], tuple[ModelType, ...]]] = [] - - class FakeMigration: - def __init__( - self, - *, - tenant_id: str, - engine: sa.Engine, - apply: bool, - output: io.TextIOBase, - model_types: tuple[ModelType, ...], - orm_models: tuple[migration_module.ORMModel, ...], - ) -> None: - assert engine is sqlite_engine - assert apply is False - seen_runs.append((tenant_id, tuple(model.__tablename__ for model in orm_models), model_types)) - - def run(self) -> None: - return None - - monkeypatch.setattr(migration_module, "Migration", FakeMigration) - # Load-balancing tenant discovery is a deliberate exception: it scans the - # whole table so the discovery query stays easy to understand, even when - # the scheduled tenant set is wider than the selected model types. - service = migration_module.LegacyModelTypeMigrationService( - engine=sqlite_engine, - apply=False, - tables=("load_balancing_model_configs",), - model_types=(ModelType.LLM,), - ) - - service.migrate() - - assert seen_runs == [ - (inherit_llm_tenant_id, ("load_balancing_model_configs",), (ModelType.LLM,)), - (inherit_embedding_tenant_id, ("load_balancing_model_configs",), (ModelType.LLM,)), - ] - - -def test_service_migrate_with_concurrency_greater_than_one_runs_tenants_in_parallel_without_changing_migration_scope( - sqlite_engine: sa.Engine, -) -> None: - init_calls: list[dict[str, object]] = [] - started_tenants: list[str] = [] - worker_errors: list[BaseException] = [] - release_runs = threading.Event() - all_started = threading.Event() - active_runs = 0 - max_active_runs = 0 - state_lock = threading.Lock() - - class FakeMigration: - def __init__( - self, - *, - tenant_id: str, - engine: sa.Engine, - apply: bool, - output: io.TextIOBase, - model_types: tuple[ModelType, ...], - orm_models: tuple[migration_module.ORMModel, ...], - ) -> None: - self._tenant_id = tenant_id - init_calls.append( - { - "tenant_id": tenant_id, - "engine": engine, - "apply": apply, - "model_types": model_types, - "table_names": tuple(model.__tablename__ for model in orm_models), - } - ) - - def run(self) -> None: - nonlocal active_runs, max_active_runs - with state_lock: - active_runs += 1 - max_active_runs = max(max_active_runs, active_runs) - started_tenants.append(self._tenant_id) - if len(started_tenants) == 2: - all_started.set() - - release_runs.wait(timeout=1) - - with state_lock: - active_runs -= 1 - - monkeypatch = pytest.MonkeyPatch() - try: - monkeypatch.setattr(migration_module, "Migration", FakeMigration) - service = migration_module.LegacyModelTypeMigrationService( - engine=sqlite_engine, - apply=False, - concurrency=2, - tables=("provider_models",), - model_types=(ModelType.LLM,), - tenant_ids=("tenant-alpha", "tenant-beta"), - ) - - def _run_service() -> None: - try: - service.migrate() - except BaseException as exc: # pragma: no cover - test harness - worker_errors.append(exc) - - worker = threading.Thread(target=_run_service) - worker.start() - started_in_parallel = all_started.wait(timeout=0.5) - release_runs.set() - worker.join(timeout=1) - finally: - monkeypatch.undo() - - assert worker_errors == [] - assert started_in_parallel is True - assert worker.is_alive() is False - assert max_active_runs == 2 - assert {call["tenant_id"] for call in init_calls} == {"tenant-alpha", "tenant-beta"} - for call in init_calls: - assert tuple(cast(tuple[str, ...], call["table_names"])) == ("provider_models",) - assert call["model_types"] == (ModelType.LLM,) - - -def test_service_parallel_migrate_serializes_shared_output_by_line( - sqlite_engine: sa.Engine, -) -> None: - worker_errors: list[BaseException] = [] - start_barrier = threading.Barrier(2) - - class SlowLineOutput(io.StringIO): - def __init__(self) -> None: - super().__init__() - self.overlap_count = 0 - self._in_write = False - self._state_lock = threading.Lock() - - @override - def write(self, s: str) -> int: - with self._state_lock: - if self._in_write: - self.overlap_count += 1 - self._in_write = True - try: - time.sleep(0.01) - return super().write(s) - finally: - with self._state_lock: - self._in_write = False - - class FakeMigration: - def __init__( - self, - *, - tenant_id: str, - engine: sa.Engine, - apply: bool, - output: io.TextIOBase, - model_types: tuple[ModelType, ...], - orm_models: tuple[migration_module.ORMModel, ...], - ) -> None: - self._tenant_id = tenant_id - self._output = output - - def run(self) -> None: - try: - start_barrier.wait(timeout=1) - except threading.BrokenBarrierError as exc: - raise AssertionError("parallel migrate should schedule both tenant runs together") from exc - - for index in range(3): - self._output.write(f"{self._tenant_id}:line-{index}\n") - - monkeypatch = pytest.MonkeyPatch() - output = SlowLineOutput() - try: - monkeypatch.setattr(migration_module, "Migration", FakeMigration) - service = migration_module.LegacyModelTypeMigrationService( - engine=sqlite_engine, - apply=False, - concurrency=2, - output=output, - tables=("provider_models",), - model_types=(ModelType.LLM,), - tenant_ids=("tenant-alpha", "tenant-beta"), - ) - - def _run_service() -> None: - try: - service.migrate() - except BaseException as exc: # pragma: no cover - test harness - worker_errors.append(exc) - - worker = threading.Thread(target=_run_service) - worker.start() - worker.join(timeout=2) - finally: - monkeypatch.undo() - - assert worker.is_alive() is False - assert worker_errors == [] - assert output.overlap_count == 0 - assert sorted(output.getvalue().splitlines()) == sorted( - [ - "tenant-alpha:line-0", - "tenant-alpha:line-1", - "tenant-alpha:line-2", - "tenant-beta:line-0", - "tenant-beta:line-1", - "tenant-beta:line-2", - ] - ) - - -def test_migration_dry_run_emits_json_lines_without_db_or_cache_mutation( - sqlite_engine: sa.Engine, - dirty_fixture: DirtyDataFixture, - monkeypatch: pytest.MonkeyPatch, -) -> None: - before = snapshot_legacy_model_type_state(sqlite_engine) - deleted_cache_keys: list[str] = [] - - def _record_delete(self: migration_module.ProviderCredentialsCache) -> None: - deleted_cache_keys.append(self.cache_key) - - monkeypatch.setattr(migration_module.ProviderCredentialsCache, "delete", _record_delete) - output = io.StringIO() - - service = migration_module.LegacyModelTypeMigrationService( - engine=sqlite_engine, - apply=False, - output=output, - tenant_ids=(dirty_fixture.primary.tenant_id,), - ) - - service.migrate() - - after = snapshot_legacy_model_type_state(sqlite_engine) - assert after == before - assert deleted_cache_keys == [] - - lines = [json.loads(line) for line in output.getvalue().splitlines() if line.strip()] - assert lines, "dry-run should emit JSON lines" - assert all({"event", "message", "attrs", "ts"} <= set(line) for line in lines) - rendered_output = output.getvalue() - assert dirty_fixture.primary.loser_credential_id in rendered_output - assert dirty_fixture.primary.loser_credential_name in rendered_output - assert dirty_fixture.primary.loser_encrypted_config in rendered_output - - -def test_dry_run_and_apply_share_processing_scope_and_differ_only_on_side_effects( - sqlite_engine: sa.Engine, - dirty_fixture: DirtyDataFixture, - monkeypatch: pytest.MonkeyPatch, -) -> None: - before = snapshot_legacy_model_type_state(sqlite_engine) - deleted_cache_keys: list[str] = [] - - def _record_delete(self: migration_module.ProviderCredentialsCache) -> None: - deleted_cache_keys.append(self.cache_key) - - monkeypatch.setattr(migration_module.ProviderCredentialsCache, "delete", _record_delete) - - dry_run_output = io.StringIO() - migration_module.LegacyModelTypeMigrationService( - engine=sqlite_engine, - apply=False, - output=dry_run_output, - tenant_ids=(dirty_fixture.primary.tenant_id,), - ).migrate() - after_dry_run = snapshot_legacy_model_type_state(sqlite_engine) - dry_run_lines = _parse_json_lines(dry_run_output) - - apply_output = io.StringIO() - migration_module.LegacyModelTypeMigrationService( - engine=sqlite_engine, - apply=True, - output=apply_output, - tenant_ids=(dirty_fixture.primary.tenant_id,), - ).migrate() - after_apply = snapshot_legacy_model_type_state(sqlite_engine) - apply_lines = _parse_json_lines(apply_output) - - assert after_dry_run == before - assert after_apply != before - - dry_run_signatures = _collect_processing_signatures(dry_run_lines) - apply_signatures = _collect_processing_signatures(apply_lines) - assert apply_signatures == dry_run_signatures - - dry_run_cache_events = [line["event"] for line in dry_run_lines if str(line.get("event")).startswith("cache_")] - apply_cache_events = [line["event"] for line in apply_lines if str(line.get("event")).startswith("cache_")] - assert "cache_deleted" not in dry_run_cache_events - assert "cache_delete_planned" in dry_run_cache_events - assert "cache_deleted" in apply_cache_events - assert deleted_cache_keys - - dry_run_rewrite_signatures = { - signature - for signature in dry_run_signatures - if signature[0] == "row_updated" - and signature[1] in {"provider_models", "load_balancing_model_configs"} - and signature[-1] != _json_key(None) - } - apply_rewrite_signatures = { - signature - for signature in apply_signatures - if signature[0] == "row_updated" - and signature[1] in {"provider_models", "load_balancing_model_configs"} - and signature[-1] != _json_key(None) - } - assert apply_rewrite_signatures == dry_run_rewrite_signatures - - dry_run_lb_signatures = { - signature - for signature in dry_run_signatures - if signature[0] == "row_updated" and signature[1] == "load_balancing_model_configs" - } - apply_lb_signatures = { - signature - for signature in apply_signatures - if signature[0] == "row_updated" and signature[1] == "load_balancing_model_configs" - } - assert apply_lb_signatures == dry_run_lb_signatures - - -def test_provider_models_processing_uses_same_plan_locking_and_transaction_entry_for_dry_run_and_apply( - sqlite_engine: sa.Engine, - dirty_fixture: DirtyDataFixture, - monkeypatch: pytest.MonkeyPatch, -) -> None: - dry_migration = migration_module.Migration( - tenant_id=dirty_fixture.primary.tenant_id, - engine=sqlite_engine, - apply=False, - output=io.StringIO(), - model_types=(ModelType.LLM,), - orm_models=(ProviderModel,), - ) - candidate = dry_migration._load_provider_model_candidates(None)[0] - business_key = migration_module._ProviderModelBusinessKey( - tenant_id=candidate.row.tenant_id, - provider_name=candidate.row.provider_name, - model_name=candidate.row.model_name, - model_type=candidate.canonical_model_type, - ) - - apply_migration = migration_module.Migration( - tenant_id=dirty_fixture.primary.tenant_id, - engine=sqlite_engine, - apply=True, - output=io.StringIO(), - model_types=(ModelType.LLM,), - orm_models=(ProviderModel,), - ) - - current_phase = {"name": "dry"} - lock_rows_seen: list[tuple[str, bool]] = [] - begin_calls: list[str] = [] - configure_calls: list[str] = [] - - def _record_begin(session: Session, transaction: SessionTransaction) -> None: - if session.get_bind() is sqlite_engine and transaction.parent is None: - begin_calls.append(current_phase["name"]) - - def _fake_build_plan( - self: migration_module.Migration, - session: Session, - candidate: migration_module._RowWithRawModelType[ProviderModel], - *, - lock_rows: bool, - ) -> migration_module._ProviderModelGroupPlan: - assert session.get_bind() is sqlite_engine - lock_rows_seen.append((current_phase["name"], lock_rows)) - return migration_module._ProviderModelGroupPlan( - group_row_ids=[candidate.row.id], - winner=None, - loser_rows=[], - ) - - def _fake_emit_plan( - self: migration_module.Migration, - plan: migration_module._ProviderModelGroupPlan, - *, - session: Session, - tx_id: str, - business_key: migration_module._BusinessKey, - ) -> None: - return None - - def _fake_configure(self: migration_module.Migration, session: Session) -> None: - assert session.get_bind() is sqlite_engine - configure_calls.append(current_phase["name"]) - - monkeypatch.setattr(migration_module.Migration, "_build_provider_model_group_plan", _fake_build_plan) - monkeypatch.setattr(migration_module.Migration, "_emit_provider_model_group_plan", _fake_emit_plan) - monkeypatch.setattr(migration_module.Migration, "_configure_lock_timeout", _fake_configure) - event.listen(Session, "after_transaction_create", _record_begin) - try: - dry_migration._process_provider_model_group(candidate, business_key) - current_phase["name"] = "apply" - apply_migration._process_provider_model_group(candidate, business_key) - finally: - event.remove(Session, "after_transaction_create", _record_begin) - - assert [phase for phase, _ in lock_rows_seen] == ["dry", "apply"] - assert lock_rows_seen[0][1] == lock_rows_seen[1][1] - assert begin_calls == ["dry", "apply"] - assert configure_calls == ["dry", "apply"] - - -@pytest.mark.parametrize( - ("orig", "expected"), - [ - (_DatabaseError(pgcode="55P03"), True), - (_DatabaseError(sqlstate="55P03"), True), - (_DatabaseError(errno=1205), True), - (RuntimeError("canceling statement due to lock timeout"), True), - (_DatabaseError(pgcode="23505"), False), - (_DatabaseError(errno=1213), False), - ], -) -def test_is_lock_timeout_error_prefers_structured_backend_codes( - sqlite_engine: sa.Engine, - orig: BaseException, - expected: bool, -) -> None: - migration = migration_module.Migration( - tenant_id="tenant-1", - engine=sqlite_engine, - apply=True, - output=io.StringIO(), - model_types=(ModelType.LLM,), - orm_models=(), - ) - exc = OperationalError("SELECT 1", {}, orig) - - assert migration._is_lock_timeout_error(exc) is expected - - -def test_process_load_balancing_model_config_row_logs_stacktrace_for_lock_timeout( - sqlite_engine: sa.Engine, - monkeypatch: pytest.MonkeyPatch, -) -> None: - create_minimal_legacy_model_type_schema(sqlite_engine) - created_at = datetime(2025, 1, 1, 12, 0, 0) - _insert_load_balancing_model_config( - sqlite_engine, - row_id="40000000-0000-0000-0000-000000000001", - tenant_id="tenant-1", - provider_name="openai", - model_name="gpt-4o-mini", - model_type="text-generation", - name="credential", - encrypted_config="{}", - credential_id="50000000-0000-0000-0000-000000000001", - enabled=True, - created_at=created_at, - updated_at=created_at, - ) - output = io.StringIO() - migration = migration_module.Migration( - tenant_id="tenant-1", - engine=sqlite_engine, - apply=True, - output=output, - model_types=(ModelType.LLM,), - orm_models=(migration_module.LoadBalancingModelConfig,), - ) - candidate = migration._load_load_balancing_model_config_candidates(None)[0] - lock_timeout_exc = OperationalError("SELECT 1", {}, _DatabaseError(pgcode="55P03")) - transaction_begins = 0 - - def _record_begin(session: Session, transaction: SessionTransaction) -> None: - nonlocal transaction_begins - if session.get_bind() is sqlite_engine and transaction.parent is None: - transaction_begins += 1 - - def _fake_reload( - self: migration_module.Migration, - session: Session, - original_candidate: migration_module._RowWithRawModelType[LoadBalancingModelConfig], - *, - lock_rows: bool, - ) -> migration_module._RowWithRawModelType[LoadBalancingModelConfig] | None: - raise lock_timeout_exc - - monkeypatch.setattr(migration_module.Migration, "_configure_lock_timeout", lambda self, session: None) - monkeypatch.setattr( - migration_module.Migration, - "_reload_load_balancing_model_config_candidate", - _fake_reload, - ) - - event.listen(Session, "after_transaction_create", _record_begin) - try: - migration._process_load_balancing_model_config_row(candidate) - finally: - event.remove(Session, "after_transaction_create", _record_begin) - - lines = _parse_json_lines(output) - assert len(lines) == 1 - assert lines[0]["event"] == "lock_timeout_skipped" - attrs = cast(dict[str, object], lines[0]["attrs"]) - assert attrs["table_name"] == "load_balancing_model_configs" - assert attrs["id"] == candidate.row.id - assert attrs["error"] == str(lock_timeout_exc) - assert isinstance(attrs["stacktrace"], str) - assert "OperationalError" in attrs["stacktrace"] - assert transaction_begins == 1 - - -def test_process_load_balancing_model_config_row_logs_update_after_sql_execution( - sqlite_engine: sa.Engine, - monkeypatch: pytest.MonkeyPatch, -) -> None: - create_minimal_legacy_model_type_schema(sqlite_engine) - created_at = datetime(2025, 1, 1, 12, 0, 0) - row_id = "40000000-0000-0000-0000-000000000002" - _insert_load_balancing_model_config( - sqlite_engine, - row_id=row_id, - tenant_id="tenant-1", - provider_name="openai", - model_name="gpt-4o-mini", - model_type="text-generation", - name="credential", - encrypted_config="{}", - credential_id="50000000-0000-0000-0000-000000000002", - enabled=True, - created_at=created_at, - updated_at=created_at, - ) - migration = migration_module.Migration( - tenant_id="tenant-1", - engine=sqlite_engine, - apply=True, - output=io.StringIO(), - model_types=(ModelType.LLM,), - orm_models=(migration_module.LoadBalancingModelConfig,), - ) - candidate = migration._load_load_balancing_model_config_candidates(None)[0] - action_log: list[str] = [] - - def _record_begin(session: Session, transaction: SessionTransaction) -> None: - if session.get_bind() is sqlite_engine and transaction.parent is None: - action_log.append("begin") - - def _record_sql( - connection: sa.Connection, - cursor: object, - statement: str, - parameters: object, - context: object, - executemany: bool, - ) -> None: - del connection, cursor, parameters, context, executemany - if statement.lstrip().upper().startswith("UPDATE"): - action_log.append("sql_execute") - - def _fake_configure(self: migration_module.Migration, session: Session) -> None: - action_log.append("configure_lock_timeout") - - original_reload = migration_module.Migration._reload_load_balancing_model_config_candidate - - def _record_reload( - self: migration_module.Migration, - session: Session, - original_candidate: migration_module._RowWithRawModelType[LoadBalancingModelConfig], - *, - lock_rows: bool, - ) -> migration_module._RowWithRawModelType[LoadBalancingModelConfig] | None: - action_log.append(f"reload_candidate:{lock_rows}") - return original_reload(self, session, original_candidate, lock_rows=lock_rows) - - def _fake_log_row_updated(self: migration_module.Migration, *args: object, **kwargs: object) -> None: - action_log.append("log_row_updated") - - def _fake_cache_cleanup(self: migration_module.Migration, *, row_id: str, tx_id: str) -> None: - action_log.append("cache_cleanup") - - monkeypatch.setattr(migration_module.Migration, "_configure_lock_timeout", _fake_configure) - monkeypatch.setattr( - migration_module.Migration, - "_reload_load_balancing_model_config_candidate", - _record_reload, - ) - monkeypatch.setattr(migration_module.Migration, "_log_row_updated", _fake_log_row_updated) - monkeypatch.setattr( - migration_module.Migration, - "_log_load_balancing_model_config_cache_cleanup", - _fake_cache_cleanup, - ) - - event.listen(Session, "after_transaction_create", _record_begin) - event.listen(sqlite_engine, "before_cursor_execute", _record_sql) - try: - migration._process_load_balancing_model_config_row(candidate) - finally: - event.remove(sqlite_engine, "before_cursor_execute", _record_sql) - event.remove(Session, "after_transaction_create", _record_begin) - - assert action_log == [ - "begin", - "configure_lock_timeout", - "reload_candidate:True", - "sql_execute", - "log_row_updated", - "cache_cleanup", - ] - with Session(sqlite_engine) as session: - persisted = session.get(migration_module.LoadBalancingModelConfig, row_id) - assert persisted is not None - assert persisted.model_type == ModelType.LLM - - -def test_load_balancing_model_config_cache_delete_failure_logs_stacktrace( - sqlite_engine: sa.Engine, - dirty_fixture: DirtyDataFixture, - monkeypatch: pytest.MonkeyPatch, -) -> None: - def _raise_delete_failure(self: migration_module.ProviderCredentialsCache) -> None: - raise RuntimeError("cache delete boom") - - monkeypatch.setattr(migration_module.ProviderCredentialsCache, "delete", _raise_delete_failure) - - output = io.StringIO() - migration_module.LegacyModelTypeMigrationService( - engine=sqlite_engine, - apply=True, - output=output, - tables=("load_balancing_model_configs",), - model_types=(ModelType.LLM,), - tenant_ids=(dirty_fixture.primary.tenant_id,), - ).migrate() - - failed_events = [ - cast(dict[str, object], line["attrs"]) - for line in _parse_json_lines(output) - if line.get("event") == "cache_delete_failed" - and isinstance(line.get("attrs"), dict) - and cast(dict[str, object], line["attrs"]).get("table_name") == "load_balancing_model_configs" - ] - - assert len(failed_events) == 1 - assert failed_events[0]["error"] == "cache delete boom" - assert isinstance(failed_events[0]["stacktrace"], str) - assert "RuntimeError: cache delete boom" in failed_events[0]["stacktrace"] - - -def test_group_completed_logs_exist_for_all_grouped_tables_and_use_canonical_model_type( - sqlite_engine: sa.Engine, - dirty_fixture: DirtyDataFixture, -) -> None: - output = io.StringIO() - - service = migration_module.LegacyModelTypeMigrationService( - engine=sqlite_engine, - apply=False, - output=output, - tenant_ids=(dirty_fixture.primary.tenant_id,), - ) - - service.migrate() - - lines = _parse_json_lines(output) - group_completed_records = [ - line - for line in lines - if isinstance(line.get("attrs"), dict) and "group_row_ids" in cast(dict[str, object], line["attrs"]) - ] - grouped_table_names = { - cast(dict[str, object], record["attrs"]).get("table_name") for record in group_completed_records - } - - assert grouped_table_names >= { - "provider_models", - "tenant_default_models", - "provider_model_settings", - "provider_model_credentials", - } - - for record in group_completed_records: - attrs = cast(dict[str, object], record["attrs"]) - business_key = cast(dict[str, object], attrs["business_key"]) - assert isinstance(attrs["group_row_ids"], list) - assert attrs["group_row_ids"] - if "model_type" in business_key: - assert business_key["model_type"] in { - ModelType.LLM.value, - ModelType.TEXT_EMBEDDING.value, - ModelType.RERANK.value, - } - assert business_key["model_type"] not in LEGACY_TO_CANONICAL - - -def test_provider_models_group_completed_log_includes_related_canonical_row_ids( - sqlite_engine: sa.Engine, - dirty_fixture: DirtyDataFixture, - monkeypatch: pytest.MonkeyPatch, -) -> None: - _patch_batch_size(monkeypatch, batch_size=1) - inserted_row_id = "00000000-0000-0000-0000-00000000aa01" - created_at = datetime(2025, 1, 1, 10, 0, 0) - updated_at = created_at + timedelta(minutes=5) - _insert_provider_model( - sqlite_engine, - row_id=inserted_row_id, - tenant_id=dirty_fixture.primary.tenant_id, - provider_name="openai", - model_name="gpt-4o-mini", - model_type=ModelType.LLM.value, - credential_id=dirty_fixture.primary.distinct_credential_id, - created_at=created_at, - updated_at=updated_at, - ) - - output = io.StringIO() - migration_module.LegacyModelTypeMigrationService( - engine=sqlite_engine, - apply=False, - output=output, - tables=("provider_models",), - model_types=(ModelType.LLM,), - tenant_ids=(dirty_fixture.primary.tenant_id,), - ).migrate() - - lines = _parse_json_lines(output) - matching_records = [] - for line in lines: - attrs = line.get("attrs") - if not isinstance(attrs, dict): - continue - business_key = attrs.get("business_key") - if not isinstance(business_key, dict): - continue - if ( - attrs.get("table_name") == "provider_models" - and business_key.get("tenant_id") == dirty_fixture.primary.tenant_id - and business_key.get("provider_name") == "openai" - and business_key.get("model_name") == "gpt-4o-mini" - and business_key.get("model_type") == ModelType.LLM.value - and "group_row_ids" in attrs - ): - matching_records.append(attrs) - - assert len(matching_records) == 1 - assert set(cast(list[str], matching_records[0]["group_row_ids"])) == { - dirty_fixture.primary.provider_model_id, - inserted_row_id, - } - - -def test_provider_model_settings_group_crossing_batches_is_completed_once_with_all_group_row_ids( - sqlite_engine: sa.Engine, - dirty_fixture: DirtyDataFixture, - monkeypatch: pytest.MonkeyPatch, -) -> None: - _patch_batch_size(monkeypatch, batch_size=1) - inserted_row_id = "00000000-0000-0000-0000-00000000cc01" - created_at = datetime(2025, 1, 1, 9, 0, 0) - updated_at = created_at + timedelta(minutes=10) - _insert_provider_model_setting( - sqlite_engine, - row_id=inserted_row_id, - tenant_id=dirty_fixture.primary.tenant_id, - provider_name="openai", - model_name="gpt-4o-mini", - model_type="text-generation", - enabled=True, - load_balancing_enabled=False, - created_at=created_at, - updated_at=updated_at, - ) - - output = io.StringIO() - migration_module.LegacyModelTypeMigrationService( - engine=sqlite_engine, - apply=False, - output=output, - tables=("provider_model_settings",), - model_types=(ModelType.LLM,), - tenant_ids=(dirty_fixture.primary.tenant_id,), - ).migrate() - - lines = _parse_json_lines(output) - matching_records = [] - for line in lines: - attrs = line.get("attrs") - if not isinstance(attrs, dict): - continue - business_key = attrs.get("business_key") - if not isinstance(business_key, dict): - continue - if ( - attrs.get("table_name") == "provider_model_settings" - and business_key.get("tenant_id") == dirty_fixture.primary.tenant_id - and business_key.get("provider_name") == "openai" - and business_key.get("model_name") == "gpt-4o-mini" - and business_key.get("model_type") == ModelType.LLM.value - and "group_row_ids" in attrs - ): - matching_records.append(attrs) - - assert len(matching_records) == 1 - assert set(cast(list[str], matching_records[0]["group_row_ids"])) == { - dirty_fixture.primary.provider_model_setting_id, - inserted_row_id, - } - - -def test_load_balancing_inherit_rows_are_deduplicated_by_normalized_model_type_before_canonicalization( - sqlite_engine: sa.Engine, - dirty_fixture: DirtyDataFixture, - monkeypatch: pytest.MonkeyPatch, -) -> None: - older_canonical_row_id = "00000000-0000-0000-0000-00000000dd01" - newer_legacy_row_id = "00000000-0000-0000-0000-00000000dd02" - created_at = datetime(2025, 1, 1, 8, 0, 0) - older_updated_at = created_at + timedelta(minutes=15) - newer_updated_at = created_at + timedelta(minutes=30) - _insert_load_balancing_model_config( - sqlite_engine, - row_id=older_canonical_row_id, - tenant_id=dirty_fixture.primary.tenant_id, - provider_name="openai", - model_name="gpt-4o-mini", - model_type=ModelType.LLM.value, - name="__inherit__", - encrypted_config='{"api_key":"older-inherit"}', - credential_id=dirty_fixture.primary.winner_credential_id, - enabled=True, - created_at=created_at, - updated_at=older_updated_at, - ) - _insert_load_balancing_model_config( - sqlite_engine, - row_id=newer_legacy_row_id, - tenant_id=dirty_fixture.primary.tenant_id, - provider_name="openai", - model_name="gpt-4o-mini", - model_type="text-generation", - name="__inherit__", - encrypted_config='{"api_key":"newer-inherit"}', - credential_id=dirty_fixture.primary.distinct_credential_id, - enabled=True, - created_at=created_at, - updated_at=newer_updated_at, - ) - - deleted_cache_keys: list[str] = [] - - def _record_delete(self: migration_module.ProviderCredentialsCache) -> None: - deleted_cache_keys.append(self.cache_key) - - monkeypatch.setattr(migration_module.ProviderCredentialsCache, "delete", _record_delete) - - tenant_id = dirty_fixture.primary.tenant_id - table_name = "load_balancing_model_configs" - expected_row_ids = {older_canonical_row_id, newer_legacy_row_id} - - dry_run_output = io.StringIO() - migration_module.LegacyModelTypeMigrationService( - engine=sqlite_engine, - apply=False, - output=dry_run_output, - tables=(table_name,), - model_types=(ModelType.LLM,), - tenant_ids=(tenant_id,), - ).migrate() - - dry_run_lines = _parse_json_lines(dry_run_output) - dry_run_signatures = { - signature - for signature in _collect_processing_signatures(dry_run_lines) - if signature[1] == table_name and signature[2] in expected_row_ids - } - dry_run_row_updates = [ - cast(dict[str, object], line["attrs"]) - for line in dry_run_lines - if line.get("event") == "row_updated" - and isinstance(line.get("attrs"), dict) - and cast(dict[str, object], line["attrs"]).get("table_name") == table_name - and str(cast(dict[str, object], line["attrs"]).get("id")) in expected_row_ids - ] - assert len(dry_run_row_updates) == 1 - assert str(dry_run_row_updates[0]["id"]) == newer_legacy_row_id - assert dry_run_row_updates[0]["old_values"] == {"model_type": "text-generation"} - assert dry_run_row_updates[0]["new_values"] == {"model_type": ModelType.LLM.value} - assert all("rewrite_source" not in attrs for attrs in dry_run_row_updates) - - dry_run_row_deletes = [ - cast(dict[str, object], line["attrs"]) - for line in dry_run_lines - if line.get("event") == "row_deleted" - and isinstance(line.get("attrs"), dict) - and cast(dict[str, object], line["attrs"]).get("table_name") == table_name - and str(cast(dict[str, object], line["attrs"]).get("id")) in expected_row_ids - ] - assert len(dry_run_row_deletes) == 1 - assert dry_run_row_deletes[0]["business_key"] == { - "tenant_id": tenant_id, - "provider_name": "openai", - "model_name": "gpt-4o-mini", - "model_type": ModelType.LLM.value, - } - assert dry_run_row_deletes[0]["merge_winner_id"] == newer_legacy_row_id - assert dry_run_row_deletes[0]["row"] == { - "id": older_canonical_row_id, - "tenant_id": tenant_id, - "provider_name": "openai", - "model_name": "gpt-4o-mini", - "model_type": ModelType.LLM.value, - "name": "__inherit__", - "encrypted_config": {"api_key": "older-inherit"}, - "credential_id": dirty_fixture.primary.winner_credential_id, - "credential_source_type": CredentialSourceType.CUSTOM_MODEL.value, - "enabled": True, - "created_at": created_at.isoformat(), - "updated_at": older_updated_at.isoformat(), - } - - dry_run_deleted_index = next( - index - for index, line in enumerate(dry_run_lines) - if line.get("event") == "row_deleted" - and isinstance(line.get("attrs"), dict) - and cast(dict[str, object], line["attrs"]).get("id") == older_canonical_row_id - ) - dry_run_updated_index = next( - index - for index, line in enumerate(dry_run_lines) - if line.get("event") == "row_updated" - and isinstance(line.get("attrs"), dict) - and cast(dict[str, object], line["attrs"]).get("id") == newer_legacy_row_id - ) - assert dry_run_deleted_index < dry_run_updated_index - - dry_run_cache_plan_ids = _cache_event_row_ids( - dry_run_lines, - table_name=table_name, - row_ids=expected_row_ids, - event_name="cache_delete_planned", - ) - assert newer_legacy_row_id in dry_run_cache_plan_ids - - apply_output = io.StringIO() - migration_module.LegacyModelTypeMigrationService( - engine=sqlite_engine, - apply=True, - output=apply_output, - tables=(table_name,), - model_types=(ModelType.LLM,), - tenant_ids=(tenant_id,), - ).migrate() - - apply_lines = _parse_json_lines(apply_output) - apply_signatures = { - signature - for signature in _collect_processing_signatures(apply_lines) - if signature[1] == table_name and signature[2] in expected_row_ids - } - apply_row_updates = [ - cast(dict[str, object], line["attrs"]) - for line in apply_lines - if line.get("event") == "row_updated" - and isinstance(line.get("attrs"), dict) - and cast(dict[str, object], line["attrs"]).get("table_name") == table_name - and str(cast(dict[str, object], line["attrs"]).get("id")) in expected_row_ids - ] - assert len(apply_row_updates) == 1 - assert str(apply_row_updates[0]["id"]) == newer_legacy_row_id - assert apply_signatures == dry_run_signatures - - apply_cache_delete_ids = _cache_event_row_ids( - apply_lines, - table_name=table_name, - row_ids=expected_row_ids, - event_name="cache_deleted", - ) - assert apply_cache_delete_ids == dry_run_cache_plan_ids - assert deleted_cache_keys - - lb_rows = fetch_table_rows(sqlite_engine, table_name, tenant_id=tenant_id) - surviving_rows = [row for row in lb_rows if str(row["id"]) in expected_row_ids] - assert len(surviving_rows) == 1 - surviving_row = surviving_rows[0] - assert surviving_row["id"] == newer_legacy_row_id - assert surviving_row["tenant_id"] == tenant_id - assert surviving_row["provider_name"] == "openai" - assert surviving_row["model_name"] == "gpt-4o-mini" - assert surviving_row["model_type"] == ModelType.LLM.value - assert surviving_row["name"] == "__inherit__" - assert surviving_row["encrypted_config"] == '{"api_key":"newer-inherit"}' - assert surviving_row["credential_id"] == dirty_fixture.primary.distinct_credential_id - assert surviving_row["credential_source_type"] == CredentialSourceType.CUSTOM_MODEL.value - - -def test_load_balancing_non_inherit_rows_do_not_participate_in_normalized_model_type_deduplication( - sqlite_engine: sa.Engine, - dirty_fixture: DirtyDataFixture, -) -> None: - inserted_row_id = "00000000-0000-0000-0000-00000000dd03" - created_at = datetime(2025, 1, 1, 8, 0, 0) - updated_at = created_at + timedelta(minutes=15) - _insert_load_balancing_model_config( - sqlite_engine, - row_id=inserted_row_id, - tenant_id=dirty_fixture.primary.tenant_id, - provider_name="openai", - model_name="gpt-4o-mini", - model_type=ModelType.LLM.value, - name=dirty_fixture.primary.loser_credential_name, - encrypted_config='{"api_key":"second-lb"}', - credential_id=dirty_fixture.primary.distinct_credential_id, - enabled=True, - created_at=created_at, - updated_at=updated_at, - ) - - output = io.StringIO() - migration_module.LegacyModelTypeMigrationService( - engine=sqlite_engine, - apply=True, - output=output, - tables=("load_balancing_model_configs",), - model_types=(ModelType.LLM,), - tenant_ids=(dirty_fixture.primary.tenant_id,), - ).migrate() - - lines = _parse_json_lines(output) - row_deleted_events = [ - cast(dict[str, object], line["attrs"]) - for line in lines - if line.get("event") == "row_deleted" - and isinstance(line.get("attrs"), dict) - and cast(dict[str, object], line["attrs"]).get("table_name") == "load_balancing_model_configs" - ] - assert row_deleted_events == [] - - lb_rows = fetch_table_rows( - sqlite_engine, - "load_balancing_model_configs", - tenant_id=dirty_fixture.primary.tenant_id, - ) - matching_rows = [ - row for row in lb_rows if str(row["id"]) in {dirty_fixture.primary.load_balancing_config_id, inserted_row_id} - ] - assert len(matching_rows) == 2 - assert all(row["model_type"] == ModelType.LLM.value for row in matching_rows) - - -def test_migration_apply_updates_all_five_tables_and_rewrites_credential_references( - sqlite_engine: sa.Engine, - dirty_fixture: DirtyDataFixture, - monkeypatch: pytest.MonkeyPatch, -) -> None: - deleted_cache_keys: list[str] = [] - - def _record_delete(self: migration_module.ProviderCredentialsCache) -> None: - deleted_cache_keys.append(self.cache_key) - - monkeypatch.setattr(migration_module.ProviderCredentialsCache, "delete", _record_delete) - output = io.StringIO() - - service = migration_module.LegacyModelTypeMigrationService( - engine=sqlite_engine, - apply=True, - output=output, - tenant_ids=(dirty_fixture.primary.tenant_id,), - ) - - service.migrate() - - assert_tenant_rows_use_only_canonical_model_types(sqlite_engine, dirty_fixture.primary.tenant_id) - - provider_model_rows = fetch_table_rows(sqlite_engine, "provider_models", tenant_id=dirty_fixture.primary.tenant_id) - provider_model_row = next( - row for row in provider_model_rows if row["id"] == dirty_fixture.primary.provider_model_id - ) - assert provider_model_row["model_type"] == LEGACY_TO_CANONICAL["text-generation"] - assert provider_model_row["credential_id"] == dirty_fixture.primary.winner_credential_id - - lb_rows = fetch_table_rows(sqlite_engine, "load_balancing_model_configs", tenant_id=dirty_fixture.primary.tenant_id) - lb_row = next(row for row in lb_rows if row["id"] == dirty_fixture.primary.load_balancing_config_id) - assert lb_row["model_type"] == LEGACY_TO_CANONICAL["text-generation"] - assert lb_row["credential_id"] == dirty_fixture.primary.winner_credential_id - assert lb_row["encrypted_config"] == dirty_fixture.primary.winner_encrypted_config - - credential_rows = fetch_table_rows( - sqlite_engine, "provider_model_credentials", tenant_id=dirty_fixture.primary.tenant_id - ) - assert ( - count_rows( - sqlite_engine, - "provider_model_credentials", - tenant_id=dirty_fixture.primary.tenant_id, - ) - == 2 - ) - credential_ids = {str(row["id"]) for row in credential_rows} - assert credential_ids == { - dirty_fixture.primary.winner_credential_id, - dirty_fixture.primary.distinct_credential_id, - } - distinct_row = next(row for row in credential_rows if row["id"] == dirty_fixture.primary.distinct_credential_id) - assert distinct_row["credential_name"] == dirty_fixture.primary.distinct_credential_name - assert distinct_row["model_type"] == LEGACY_TO_CANONICAL["text-generation"] - - rendered_output = output.getvalue() - assert dirty_fixture.primary.loser_credential_id in rendered_output - assert dirty_fixture.primary.loser_encrypted_config in rendered_output - assert any("load_balancing_provider_model_credentials" in key for key in deleted_cache_keys) or any( - "load_balancing_provider_model" in key for key in deleted_cache_keys - ) - - -def test_migration_filters_by_tenant_model_types_and_tables( - sqlite_engine: sa.Engine, - dirty_fixture: DirtyDataFixture, -) -> None: - before_primary_credentials = fetch_table_rows( - sqlite_engine, - "provider_model_credentials", - tenant_id=dirty_fixture.primary.tenant_id, - ) - before_secondary = { - table_name: fetch_table_rows(sqlite_engine, table_name, tenant_id=dirty_fixture.secondary.tenant_id) - for table_name in ALL_TABLE_NAMES - } - - service = migration_module.LegacyModelTypeMigrationService( - engine=sqlite_engine, - apply=True, - output=io.StringIO(), - tables=("provider_models",), - model_types=(ModelType.LLM,), - tenant_ids=(dirty_fixture.primary.tenant_id,), - ) - - service.migrate() - - assert ( - count_rows( - sqlite_engine, - "provider_model_credentials", - tenant_id=dirty_fixture.primary.tenant_id, - ) - == 3 - ) - credential_rows = fetch_table_rows( - sqlite_engine, - "provider_model_credentials", - tenant_id=dirty_fixture.primary.tenant_id, - ) - assert credential_rows == before_primary_credentials - provider_model_rows = fetch_table_rows( - sqlite_engine, - "provider_models", - tenant_id=dirty_fixture.primary.tenant_id, - ) - provider_model_row = next( - row for row in provider_model_rows if row["id"] == dirty_fixture.primary.provider_model_id - ) - embedding_provider_model_row = next( - row for row in provider_model_rows if row["id"] == dirty_fixture.primary.embedding_provider_model_id - ) - assert provider_model_row["model_type"] == LEGACY_TO_CANONICAL["text-generation"] - assert embedding_provider_model_row["model_type"] == "embeddings" - - tenant_default_row = fetch_table_rows( - sqlite_engine, - "tenant_default_models", - tenant_id=dirty_fixture.primary.tenant_id, - )[0] - assert tenant_default_row["model_type"] == "text-generation" - - provider_model_setting_rows = fetch_table_rows( - sqlite_engine, - "provider_model_settings", - tenant_id=dirty_fixture.primary.tenant_id, - ) - llm_setting_row = next( - row for row in provider_model_setting_rows if row["id"] == dirty_fixture.primary.provider_model_setting_id - ) - embedding_setting_row = next( - row for row in provider_model_setting_rows if row["id"] == dirty_fixture.primary.embedding_setting_id - ) - assert llm_setting_row["model_type"] == "text-generation" - assert embedding_setting_row["model_type"] == "embeddings" - - lb_row = fetch_table_rows( - sqlite_engine, - "load_balancing_model_configs", - tenant_id=dirty_fixture.primary.tenant_id, - )[0] - assert lb_row["model_type"] == "text-generation" - - after_secondary = { - table_name: fetch_table_rows(sqlite_engine, table_name, tenant_id=dirty_fixture.secondary.tenant_id) - for table_name in ALL_TABLE_NAMES - } - assert after_secondary == before_secondary - - -def test_migration_does_not_merge_credentials_with_different_credential_name( - sqlite_engine: sa.Engine, - dirty_fixture: DirtyDataFixture, -) -> None: - service = migration_module.LegacyModelTypeMigrationService( - engine=sqlite_engine, - apply=True, - output=io.StringIO(), - tenant_ids=(dirty_fixture.primary.tenant_id,), - ) - - service.migrate() - - credential_rows = fetch_table_rows( - sqlite_engine, - "provider_model_credentials", - tenant_id=dirty_fixture.primary.tenant_id, - ) - distinct_row = next(row for row in credential_rows if row["id"] == dirty_fixture.primary.distinct_credential_id) - assert distinct_row["credential_name"] == dirty_fixture.primary.distinct_credential_name - assert distinct_row["model_type"] == LEGACY_TO_CANONICAL["text-generation"] - assert ( - count_rows( - sqlite_engine, - "provider_model_credentials", - tenant_id=dirty_fixture.primary.tenant_id, - ) - == 2 - ) - - -def test_migration_is_idempotent_on_second_apply( - sqlite_engine: sa.Engine, - dirty_fixture: DirtyDataFixture, -) -> None: - service = migration_module.LegacyModelTypeMigrationService( - engine=sqlite_engine, - apply=True, - output=io.StringIO(), - tenant_ids=(dirty_fixture.primary.tenant_id,), - ) - - service.migrate() - after_first = snapshot_legacy_model_type_state(sqlite_engine) - - service.migrate() - after_second = snapshot_legacy_model_type_state(sqlite_engine) - - assert after_second == after_first diff --git a/api/tests/unit_tests/commands/test_rbac.py b/api/tests/unit_tests/commands/test_rbac.py new file mode 100644 index 00000000000..d9828740955 --- /dev/null +++ b/api/tests/unit_tests/commands/test_rbac.py @@ -0,0 +1,177 @@ +import json +from collections.abc import Iterator + +import pytest +from click.testing import CliRunner + +from commands import rbac +from enums.account import TenantAccountRole +from models import TenantAccountJoin +from models.agent import Agent, AgentScope, AgentSource +from tests.helpers.unit.model_factories import make_tenant +from tests.helpers.unit.rbac_fakes import RBACDomain + +"""Exercise migration queries against SQLite and the shared RBAC transport fake.""" + + +def _events(output: str) -> list[dict[str, object]]: + return [json.loads(line) for line in output.splitlines() if line.startswith("{")] + + +def _seed_agent(domain: RBACDomain) -> None: + with domain.sessions.begin() as session: + session.add(make_tenant(tenant_id="t1")) + session.add( + Agent( + id="ag1", + tenant_id="t1", + name="Agent", + scope=AgentScope.ROSTER, + source=AgentSource.ROSTER, + created_by="c1", + ) + ) + for index, account_id in enumerate(("owner-1", "m1", "admin-1", "m2", "m3")): + membership = TenantAccountJoin(tenant_id="t1", account_id=account_id, role=TenantAccountRole.NORMAL) + membership.id = f"membership-{index}" + session.add(membership) + session.add(TenantAccountJoin(tenant_id="other", account_id="foreign", role=TenantAccountRole.NORMAL)) + domain.transport.responses["/rbac/agents/whitelist"] = {"rbac_whitelist_scope": "all"} + # Remote RBAC bindings protect administrators even when the local role is normal. + domain.transport.responses["/rbac/members/rbac-roles/batch"] = { + f"{role_tag}-1": [ + { + "id": f"role-{role_tag}", + "type": "workspace", + "category": "global_system_default", + "name": role_tag, + "is_builtin": True, + "role_tag": role_tag, + } + ] + for role_tag in ("owner", "admin") + } + + +class TestRbacAgentMigrationCases: + @staticmethod + @pytest.fixture(autouse=True) + def migration_database(monkeypatch: pytest.MonkeyPatch, rbac_domain: RBACDomain) -> None: + monkeypatch.setattr(rbac.session_factory, "create_session", rbac_domain.sessions) + + @staticmethod + def test_apply_flag_writes_and_reports_applied(rbac_domain: RBACDomain) -> None: + with rbac_domain.sessions.begin() as session: + session.add_all([make_tenant(tenant_id="t1"), make_tenant(tenant_id="t2")]) + rbac_domain.transport.response = { + "roles": [{"role_id": "r1", "role_name": "ops", "added_keys": ["agent.create"]}] + } + + result = CliRunner().invoke(rbac.migrate_agent_permissions_to_rbac, ["--apply"]) + + assert result.exit_code == 0, result.output + assert len(rbac_domain.transport.requests) == 2 + for request in rbac_domain.transport.requests: + assert (request.method, request.endpoint) == ("POST", "/rbac/migrations/agent-manage-roles") + assert request.json == {"apply": True} + events = _events(result.output) + assert [event["event"] for event in events] == ["agent_manage_role_migration_applied"] * 2 + assert {event["tenant_id"] for event in events} == {"t1", "t2"} + assert "changed" in result.output + + @staticmethod + def test_agent_bootstrap_apply_writes_whitelist_member_batches_and_creator_sync(rbac_domain: RBACDomain) -> None: + _seed_agent(rbac_domain) + + result = CliRunner().invoke(rbac.migrate_agent_permissions_to_rbac, ["--apply", "--member-batch-size", "2"]) + + assert result.exit_code == 0, result.output + assert [event["event"] for event in _events(result.output)] == ["agent_access_bootstrap_applied"] + assert _events(result.output)[0]["dry_run"] is False + role_queries = [ + request + for request in rbac_domain.transport.requests + if request.endpoint == "/rbac/members/rbac-roles/batch" + ] + assert [request.json for request in role_queries] == [ + {"member_ids": ["owner-1", "m1"]}, + {"member_ids": ["admin-1", "m2"]}, + {"member_ids": ["m3"]}, + ] + for request in role_queries: + assert (request.method, request.tenant_id, request.account_id) == ("POST", "t1", "c1") + writes = [request for request in rbac_domain.transport.requests if request.method == "PUT"] + assert [request.endpoint for request in writes] == [ + "/rbac/agents/user-access-policies", + "/rbac/agents/user-access-policies", + "/rbac/agents/user-access-policies", + "/rbac/access-policies/creator-member-bindings", + "/rbac/agents/whitelist", + ] + first, second, third, creator, whitelist = writes + for request in writes: + assert (request.tenant_id, request.account_id) == ("t1", "c1") + for request, account_ids in ((first, ["m1"]), (second, ["m2"]), (third, ["m3"])): + assert request.params == {"agent_id": "ag1", "account_id": None} + assert request.json == {"account_ids": account_ids, "access_policy_ids": ["default"]} + assert creator.params == {"resource_type": "agent", "agent_id": "ag1"} + assert whitelist.params == {"agent_id": "ag1"} + assert whitelist.json == {"automatic_include_workspace_members": True} + assert "1 agent(s) changed, 0 already initialised" in result.output + + @staticmethod + def test_agent_bootstrap_is_idempotent_on_a_second_apply(rbac_domain: RBACDomain) -> None: + _seed_agent(rbac_domain) + first = CliRunner().invoke(rbac.migrate_agent_permissions_to_rbac, ["--apply"]) + assert first.exit_code == 0, first.output + assert any(request.method == "PUT" for request in rbac_domain.transport.requests) + rbac_domain.transport.requests.clear() + rbac_domain.transport.responses["/rbac/migrations/agent-access-state"] = {"configured_agent_ids": ["ag1"]} + + result = CliRunner().invoke(rbac.migrate_agent_permissions_to_rbac, ["--apply"]) + + assert result.exit_code == 0, result.output + events = _events(result.output) + assert [event["event"] for event in events] == ["agent_access_bootstrap_skipped"] + assert events[0]["reason"] == "already_initialized" + assert not any(request.method == "PUT" for request in rbac_domain.transport.requests) + assert "0 agent(s) changed, 1 already initialised" in result.output + + +class TestRbacMemberMigrationCases: + @staticmethod + @pytest.mark.parametrize("workers", [1, 2]) + def test_member_migration_uses_shared_remote_adapter_without_write_sessions( + monkeypatch: pytest.MonkeyPatch, rbac_domain: RBACDomain, workers: int + ) -> None: + def batches( + tenant_id: str | None, *, db_batch_size: int, api_batch_size: int + ) -> Iterator[tuple[str, str, list[tuple[str, str]]]]: + assert tenant_id is None + assert db_batch_size > 0 + assert api_batch_size > 0 + yield "workspace", "owner", [("first", "editor"), ("second", "normal")] + + def unexpected_session() -> None: + raise AssertionError("A remote role replacement must not open a database session") + + monkeypatch.setattr(rbac, "_iter_tenant_member_batches", batches) + monkeypatch.setattr( + rbac, "_resolve_builtin_role_ids", lambda *_args: {"editor": "role-editor", "normal": "role-normal"} + ) + monkeypatch.setattr(rbac.RBACService.MemberRoles, "batch_get", lambda **_kwargs: []) + monkeypatch.setattr(rbac.session_factory, "create_session", unexpected_session) + rbac_domain.transport.response = {"account_id": "member", "roles": []} + + result = CliRunner().invoke(rbac.migrate_member_roles_to_rbac, ["--workers", str(workers)]) + + assert result.exit_code == 0, result.output + sent = rbac_domain.transport.requests + assert len(sent) == 2 + for request in sent: + assert (request.method, request.endpoint) == ("PUT", "/rbac/members/rbac-roles") + assert (request.tenant_id, request.account_id) == ("workspace", "owner") + assert request.params is not None + expected_role = "role-editor" if request.params["account_id"] == "first" else "role-normal" + assert request.json == {"role_ids": [expected_role]} + assert {request.params["account_id"] for request in sent if request.params} == {"first", "second"} diff --git a/api/tests/unit_tests/commands/test_rbac_agent_migration.py b/api/tests/unit_tests/commands/test_rbac_agent_migration.py deleted file mode 100644 index 81e929c7f7d..00000000000 --- a/api/tests/unit_tests/commands/test_rbac_agent_migration.py +++ /dev/null @@ -1,131 +0,0 @@ -"""Exercise migration queries against SQLite and the shared RBAC transport fake.""" - -import json - -import pytest -from click.testing import CliRunner - -from commands import rbac -from enums.account import TenantAccountRole -from models import TenantAccountJoin -from models.agent import Agent, AgentScope, AgentSource -from tests.unit_tests.model_factories import make_tenant -from tests.unit_tests.rbac_fakes import RBACDomain - - -def _events(output: str) -> list[dict[str, object]]: - return [json.loads(line) for line in output.splitlines() if line.startswith("{")] - - -@pytest.fixture(autouse=True) -def migration_database(monkeypatch: pytest.MonkeyPatch, rbac_domain: RBACDomain) -> None: - monkeypatch.setattr(rbac.session_factory, "create_session", rbac_domain.sessions) - - -def test_apply_flag_writes_and_reports_applied(rbac_domain: RBACDomain) -> None: - with rbac_domain.sessions.begin() as session: - session.add_all([make_tenant(tenant_id="t1"), make_tenant(tenant_id="t2")]) - rbac_domain.transport.response = {"roles": [{"role_id": "r1", "role_name": "ops", "added_keys": ["agent.create"]}]} - - result = CliRunner().invoke(rbac.migrate_agent_permissions_to_rbac, ["--apply"]) - - assert result.exit_code == 0, result.output - assert len(rbac_domain.transport.requests) == 2 - for request in rbac_domain.transport.requests: - assert (request.method, request.endpoint) == ("POST", "/rbac/migrations/agent-manage-roles") - assert request.json == {"apply": True} - events = _events(result.output) - assert [event["event"] for event in events] == ["agent_manage_role_migration_applied"] * 2 - assert {event["tenant_id"] for event in events} == {"t1", "t2"} - assert "changed" in result.output - - -def _seed_agent(domain: RBACDomain) -> None: - with domain.sessions.begin() as session: - session.add(make_tenant(tenant_id="t1")) - session.add( - Agent( - id="ag1", - tenant_id="t1", - name="Agent", - scope=AgentScope.ROSTER, - source=AgentSource.ROSTER, - created_by="c1", - ) - ) - for index, account_id in enumerate(("owner-1", "m1", "admin-1", "m2", "m3")): - membership = TenantAccountJoin(tenant_id="t1", account_id=account_id, role=TenantAccountRole.NORMAL) - membership.id = f"membership-{index}" - session.add(membership) - session.add(TenantAccountJoin(tenant_id="other", account_id="foreign", role=TenantAccountRole.NORMAL)) - domain.transport.responses["/rbac/agents/whitelist"] = {"rbac_whitelist_scope": "all"} - # Remote RBAC bindings protect administrators even when the local role is normal. - domain.transport.responses["/rbac/members/rbac-roles/batch"] = { - f"{role_tag}-1": [ - { - "id": f"role-{role_tag}", - "type": "workspace", - "category": "global_system_default", - "name": role_tag, - "is_builtin": True, - "role_tag": role_tag, - } - ] - for role_tag in ("owner", "admin") - } - - -def test_agent_bootstrap_apply_writes_whitelist_member_batches_and_creator_sync(rbac_domain: RBACDomain) -> None: - _seed_agent(rbac_domain) - - result = CliRunner().invoke(rbac.migrate_agent_permissions_to_rbac, ["--apply", "--member-batch-size", "2"]) - - assert result.exit_code == 0, result.output - assert [event["event"] for event in _events(result.output)] == ["agent_access_bootstrap_applied"] - assert _events(result.output)[0]["dry_run"] is False - role_queries = [ - request for request in rbac_domain.transport.requests if request.endpoint == "/rbac/members/rbac-roles/batch" - ] - assert [request.json for request in role_queries] == [ - {"member_ids": ["owner-1", "m1"]}, - {"member_ids": ["admin-1", "m2"]}, - {"member_ids": ["m3"]}, - ] - for request in role_queries: - assert (request.method, request.tenant_id, request.account_id) == ("POST", "t1", "c1") - writes = [request for request in rbac_domain.transport.requests if request.method == "PUT"] - assert [request.endpoint for request in writes] == [ - "/rbac/agents/user-access-policies", - "/rbac/agents/user-access-policies", - "/rbac/agents/user-access-policies", - "/rbac/access-policies/creator-member-bindings", - "/rbac/agents/whitelist", - ] - first, second, third, creator, whitelist = writes - for request in writes: - assert (request.tenant_id, request.account_id) == ("t1", "c1") - for request, account_ids in ((first, ["m1"]), (second, ["m2"]), (third, ["m3"])): - assert request.params == {"agent_id": "ag1", "account_id": None} - assert request.json == {"account_ids": account_ids, "access_policy_ids": ["default"]} - assert creator.params == {"resource_type": "agent", "agent_id": "ag1"} - assert whitelist.params == {"agent_id": "ag1"} - assert whitelist.json == {"automatic_include_workspace_members": True} - assert "1 agent(s) changed, 0 already initialised" in result.output - - -def test_agent_bootstrap_is_idempotent_on_a_second_apply(rbac_domain: RBACDomain) -> None: - _seed_agent(rbac_domain) - first = CliRunner().invoke(rbac.migrate_agent_permissions_to_rbac, ["--apply"]) - assert first.exit_code == 0, first.output - assert any(request.method == "PUT" for request in rbac_domain.transport.requests) - rbac_domain.transport.requests.clear() - rbac_domain.transport.responses["/rbac/migrations/agent-access-state"] = {"configured_agent_ids": ["ag1"]} - - result = CliRunner().invoke(rbac.migrate_agent_permissions_to_rbac, ["--apply"]) - - assert result.exit_code == 0, result.output - events = _events(result.output) - assert [event["event"] for event in events] == ["agent_access_bootstrap_skipped"] - assert events[0]["reason"] == "already_initialized" - assert not any(request.method == "PUT" for request in rbac_domain.transport.requests) - assert "0 agent(s) changed, 1 already initialised" in result.output diff --git a/api/tests/unit_tests/commands/test_rbac_member_migration.py b/api/tests/unit_tests/commands/test_rbac_member_migration.py deleted file mode 100644 index 693a0b9b0d1..00000000000 --- a/api/tests/unit_tests/commands/test_rbac_member_migration.py +++ /dev/null @@ -1,44 +0,0 @@ -from collections.abc import Iterator - -import pytest -from click.testing import CliRunner - -from commands import rbac -from tests.unit_tests.rbac_fakes import RBACDomain - - -@pytest.mark.parametrize("workers", [1, 2]) -def test_member_migration_uses_shared_remote_adapter_without_write_sessions( - monkeypatch: pytest.MonkeyPatch, rbac_domain: RBACDomain, workers: int -) -> None: - def batches( - tenant_id: str | None, *, db_batch_size: int, api_batch_size: int - ) -> Iterator[tuple[str, str, list[tuple[str, str]]]]: - assert tenant_id is None - assert db_batch_size > 0 - assert api_batch_size > 0 - yield "workspace", "owner", [("first", "editor"), ("second", "normal")] - - def unexpected_session() -> None: - raise AssertionError("A remote role replacement must not open a database session") - - monkeypatch.setattr(rbac, "_iter_tenant_member_batches", batches) - monkeypatch.setattr( - rbac, "_resolve_builtin_role_ids", lambda *_args: {"editor": "role-editor", "normal": "role-normal"} - ) - monkeypatch.setattr(rbac.RBACService.MemberRoles, "batch_get", lambda **_kwargs: []) - monkeypatch.setattr(rbac.session_factory, "create_session", unexpected_session) - rbac_domain.transport.response = {"account_id": "member", "roles": []} - - result = CliRunner().invoke(rbac.migrate_member_roles_to_rbac, ["--workers", str(workers)]) - - assert result.exit_code == 0, result.output - sent = rbac_domain.transport.requests - assert len(sent) == 2 - for request in sent: - assert (request.method, request.endpoint) == ("PUT", "/rbac/members/rbac-roles") - assert (request.tenant_id, request.account_id) == ("workspace", "owner") - assert request.params is not None - expected_role = "role-editor" if request.params["account_id"] == "first" else "role-normal" - assert request.json == {"role_ids": [expected_role]} - assert {request.params["account_id"] for request in sent if request.params} == {"first", "second"} diff --git a/api/tests/unit_tests/commands/test_reset_encrypt_key_pair.py b/api/tests/unit_tests/commands/test_reset_encrypt_key_pair.py deleted file mode 100644 index 24afa1cb44a..00000000000 --- a/api/tests/unit_tests/commands/test_reset_encrypt_key_pair.py +++ /dev/null @@ -1,175 +0,0 @@ -"""SQLite-backed tests for the reset-encrypt-key-pair CLI command (#35396). - -The command must purge every table that stores ciphertext encrypted with the -tenant's asymmetric key, otherwise stale rows cause downstream API failures -such as `/console/api/workspaces/current/tool-providers` returning 500. -Tests bind the command-owned transaction to the fixture engine and assert the -committed state rather than inspecting fabricated ``Session.execute`` calls. -""" - -from types import SimpleNamespace - -import pytest -from sqlalchemy import select -from sqlalchemy.orm import Session - -import commands -from commands import system as system_commands -from core.tools.entities.tool_entities import ApiProviderSchemaType -from enums import DeploymentEdition -from graphon.model_runtime.entities.model_entities import ModelType -from models import Tenant -from models.provider import Provider, ProviderModel, ProviderType -from models.tools import ApiToolProvider, BuiltinToolProvider, MCPToolProvider -from tests.unit_tests.config_override import apply_config_overrides -from tests.unit_tests.model_factories import make_tenant - - -def _invoke_reset() -> int: - callback = commands.reset_encrypt_key_pair.callback - assert callback is not None - try: - callback() - except SystemExit as e: - return int(e.code or 0) - return 0 - - -TENANT_ID = "11111111-1111-1111-1111-111111111111" -OTHER_TENANT_ID = "11111111-1111-1111-1111-111111111112" -USER_ID = "22222222-2222-2222-2222-222222222222" - - -def _tenant(tenant_id: str, *, name: str = "Test tenant") -> Tenant: - return make_tenant(tenant_id=tenant_id, name=name, encrypt_public_key="old-key") - - -def _encrypted_rows(tenant_id: str, *, suffix: str = "1") -> tuple[object, ...]: - """Build one persisted credential-bearing row for every purge target.""" - return ( - Provider(tenant_id=tenant_id, provider_name=f"provider-{suffix}"), - ProviderModel( - tenant_id=tenant_id, - provider_name=f"provider-{suffix}", - model_name=f"model-{suffix}", - model_type=ModelType.LLM, - ), - BuiltinToolProvider( - name=f"builtin-credential-{suffix}", - tenant_id=tenant_id, - user_id=USER_ID, - provider=f"builtin-{suffix}", - encrypted_credentials="ciphertext", - ), - ApiToolProvider( - name=f"api-{suffix}", - icon="icon", - schema="{}", - schema_type_str=ApiProviderSchemaType.OPENAPI, - user_id=USER_ID, - tenant_id=tenant_id, - description="description", - tools_str="[]", - credentials_str="{}", - ), - MCPToolProvider( - name=f"mcp-{suffix}", - server_identifier=f"server-{suffix}", - server_url="ciphertext", - server_url_hash=f"hash-{suffix}", - icon=None, - tenant_id=tenant_id, - user_id=USER_ID, - encrypted_credentials="ciphertext", - ), - ) - - -def _bind_command_to_sqlite(monkeypatch: pytest.MonkeyPatch, session: Session) -> None: - monkeypatch.setattr(system_commands, "db", SimpleNamespace(engine=session.get_bind())) - - -def test_reset_aborts_when_not_self_hosted(monkeypatch: pytest.MonkeyPatch, capsys: pytest.CaptureFixture[str]) -> None: - apply_config_overrides(monkeypatch, DEPLOYMENT_EDITION=DeploymentEdition.CLOUD) - - exit_code = _invoke_reset() - captured = capsys.readouterr() - - assert exit_code == 0 - assert "only for SELF_HOSTED" in captured.out - - -@pytest.mark.parametrize( - "sqlite_session", - [(Tenant, Provider, ProviderModel, BuiltinToolProvider, ApiToolProvider, MCPToolProvider)], - indirect=True, -) -def test_reset_purges_provider_and_tool_tables_for_each_tenant( - monkeypatch: pytest.MonkeyPatch, capsys: pytest.CaptureFixture[str], sqlite_session: Session -) -> None: - """The command must purge LLM provider rows AND every tool provider table - that stores ciphertext encrypted under the tenant key (#35396).""" - apply_config_overrides(monkeypatch, DEPLOYMENT_EDITION=DeploymentEdition.COMMUNITY) - monkeypatch.setattr(system_commands, "generate_key_pair", lambda tenant_id: f"new-key-{tenant_id}") - _bind_command_to_sqlite(monkeypatch, sqlite_session) - - tenant = _tenant(TENANT_ID) - other_tenant = _tenant(OTHER_TENANT_ID, name="Other tenant") - system_provider = Provider( - tenant_id=TENANT_ID, - provider_name="system-provider", - provider_type=ProviderType.SYSTEM, - ) - sqlite_session.add_all((tenant, other_tenant, system_provider, *_encrypted_rows(TENANT_ID))) - sqlite_session.commit() - - exit_code = _invoke_reset() - - captured = capsys.readouterr() - assert exit_code == 0 - assert TENANT_ID in captured.out - - sqlite_session.expire_all() - tenant_after = sqlite_session.get(Tenant, TENANT_ID) - other_tenant_after = sqlite_session.get(Tenant, OTHER_TENANT_ID) - assert tenant_after is not None - assert other_tenant_after is not None - assert tenant_after.encrypt_public_key == f"new-key-{TENANT_ID}" - assert other_tenant_after.encrypt_public_key == f"new-key-{OTHER_TENANT_ID}" - assert sqlite_session.scalars(select(Provider).where(Provider.provider_type == ProviderType.CUSTOM)).all() == [] - assert sqlite_session.scalars(select(ProviderModel)).all() == [] - assert sqlite_session.scalars(select(BuiltinToolProvider)).all() == [] - assert sqlite_session.scalars(select(ApiToolProvider)).all() == [] - assert sqlite_session.scalars(select(MCPToolProvider)).all() == [] - assert ( - sqlite_session.scalar(select(Provider).where(Provider.provider_type == ProviderType.SYSTEM)) is system_provider - ) - - -@pytest.mark.parametrize( - "sqlite_session", - [(Tenant, Provider, ProviderModel, BuiltinToolProvider, ApiToolProvider, MCPToolProvider)], - indirect=True, -) -def test_reset_iterates_all_tenants(monkeypatch: pytest.MonkeyPatch, sqlite_session: Session) -> None: - """Multi-tenant deployments must purge every tenant, not just the first.""" - apply_config_overrides(monkeypatch, DEPLOYMENT_EDITION=DeploymentEdition.COMMUNITY) - monkeypatch.setattr(system_commands, "generate_key_pair", lambda tenant_id: f"new-key-{tenant_id}") - - _bind_command_to_sqlite(monkeypatch, sqlite_session) - tenant_ids = [f"11111111-1111-1111-1111-{index:012d}" for index in range(3)] - tenants = [_tenant(tenant_id, name=f"Tenant {index}") for index, tenant_id in enumerate(tenant_ids)] - for index, tenant in enumerate(tenants): - sqlite_session.add(tenant) - sqlite_session.add_all(_encrypted_rows(tenant.id, suffix=str(index))) - sqlite_session.commit() - - assert _invoke_reset() == 0 - - sqlite_session.expire_all() - persisted_tenants = sqlite_session.scalars(select(Tenant).order_by(Tenant.id)).all() - assert [tenant.encrypt_public_key for tenant in persisted_tenants] == [ - f"new-key-{tenant_id}" for tenant_id in tenant_ids - ] - for model in (Provider, ProviderModel, BuiltinToolProvider, ApiToolProvider, MCPToolProvider): - assert sqlite_session.scalars(select(model)).all() == [] diff --git a/api/tests/unit_tests/commands/test_retention.py b/api/tests/unit_tests/commands/test_retention.py new file mode 100644 index 00000000000..f19f2cb59a8 --- /dev/null +++ b/api/tests/unit_tests/commands/test_retention.py @@ -0,0 +1,827 @@ +import datetime +import logging +import re +from dataclasses import dataclass +from types import SimpleNamespace +from unittest.mock import MagicMock, Mock, patch +from uuid import uuid4 + +import click +import pytest +from click.testing import CliRunner +from sqlalchemy import Engine, event, text +from sqlalchemy.exc import OperationalError +from sqlalchemy.orm import Session, scoped_session, sessionmaker + +from commands import clean_expired_messages, retention +from graphon.enums import WorkflowExecutionStatus, WorkflowNodeExecutionStatus +from models.base import TypeBase +from models.enums import CreatorUserRole, WorkflowRunTriggeredFrom +from models.workflow import ( + WorkflowNodeExecutionModel, + WorkflowNodeExecutionTriggeredFrom, + WorkflowRun, + WorkflowType, +) +from services.retention.workflow_run import bundle_archive_maintenance +from services.retention.workflow_run.bundle_archive_maintenance import ( + BundleOperationResult, + BundleOperationSummary, +) + +"""Tests for workflow-run archive command database boundaries. + +Planning deliberately creates a fresh session for every tenant prefix and for +every database retry. SQLite-backed tests keep the query, filtering, counting, +and session lifecycle real while clocks, billing lookup, and command failures +remain narrow external-boundary substitutions. +""" + + +_CURSOR_0 = "00000000-0000-0000-0000-000000000000" +_CURSOR_1 = "00000000-0000-0000-0000-000000000001" +_CURSOR_2 = "00000000-0000-0000-0000-000000000002" + + +def _db_disconnect_error() -> OperationalError: + return OperationalError( + "select 1", + {}, + RuntimeError("server closed the connection unexpectedly"), + connection_invalidated=True, + ) + + +@dataclass(frozen=True) +class ArchiveDatabase: + """Creates archive candidates and their node executions in SQLite.""" + + session_maker: sessionmaker[Session] + end_before: datetime.datetime + + def add_run( + self, + tenant_id: str, + *, + created_at: datetime.datetime | None = None, + status: WorkflowExecutionStatus = WorkflowExecutionStatus.SUCCEEDED, + run_type: WorkflowType = WorkflowType.WORKFLOW, + ) -> str: + run_id = str(uuid4()) + with self.session_maker.begin() as session: + session.add( + WorkflowRun( + id=run_id, + tenant_id=tenant_id, + app_id=str(uuid4()), + workflow_id=str(uuid4()), + type=run_type, + triggered_from=WorkflowRunTriggeredFrom.APP_RUN, + version="1", + graph="{}", + inputs="{}", + status=status, + outputs="{}", + error=None, + elapsed_time=0, + total_tokens=0, + total_steps=1, + created_by_role=CreatorUserRole.ACCOUNT, + created_by=str(uuid4()), + created_at=created_at or self.end_before - datetime.timedelta(days=1), + ) + ) + return run_id + + def add_node(self, run_id: str, tenant_id: str, *, index: int) -> None: + with self.session_maker.begin() as session: + session.add( + WorkflowNodeExecutionModel( + id=str(uuid4()), + tenant_id=tenant_id, + app_id=str(uuid4()), + workflow_id=str(uuid4()), + triggered_from=WorkflowNodeExecutionTriggeredFrom.WORKFLOW_RUN, + workflow_run_id=run_id, + index=index, + predecessor_node_id=None, + node_execution_id=None, + node_id=f"node-{index}", + node_type="start", + title="Start", + inputs="{}", + process_data="{}", + outputs="{}", + status=WorkflowNodeExecutionStatus.SUCCEEDED, + error=None, + elapsed_time=0, + execution_metadata="{}", + created_by_role=CreatorUserRole.ACCOUNT, + created_by=str(uuid4()), + ) + ) + + +def _delete_summary( + *, + processed: int, + succeeded: int = 0, + failed: int = 0, + next_catalog_id: str | None = None, + preview_next_catalog_id: str | None = None, + results: list[BundleOperationResult] | None = None, +) -> BundleOperationSummary: + return BundleOperationSummary( + operation="delete", + bundles_processed=processed, + bundles_succeeded=succeeded, + bundles_failed=failed, + next_catalog_id=next_catalog_id, + preview_next_catalog_id=preview_next_catalog_id, + results=results or [], + ) + + +def _patch_bundle_deleter( + monkeypatch: pytest.MonkeyPatch, + summaries: list[BundleOperationSummary], +) -> MagicMock: + deleter = MagicMock() + deleter.delete_batch.side_effect = summaries + monkeypatch.setattr( + bundle_archive_maintenance, + "WorkflowRunBundleArchiveMaintenance", + MagicMock(return_value=deleter), + ) + return deleter + + +def _tenant_id(prefix: str, suffix: int) -> str: + return f"{prefix}{suffix:07x}-0000-0000-0000-000000000000" + + +class TestArchiveWorkflowRunsCases: + @staticmethod + @pytest.mark.parametrize( + "command", + [retention.restore_workflow_runs, retention.delete_archived_workflow_runs], + ) + def test_v2_archive_maintenance_rejects_explicitly_empty_tenant_ids(command): + result = CliRunner().invoke( + command, + ["--tenant-ids", "", "--target-month", "2025-03"], + ) + + assert result.exit_code == 2 + assert "tenant-ids must not be empty" in result.output + + @staticmethod + def test_archive_tenant_id_parser_keeps_omitted_scope_unset(): + assert retention._parse_comma_separated_ids(None, param_name="tenant-ids") is None + + @staticmethod + @pytest.fixture + def archive_db(sqlite_engine: Engine) -> ArchiveDatabase: + """Create only the workflow tables used by archive planning.""" + + TypeBase.metadata.create_all( + sqlite_engine, + tables=[WorkflowRun.__table__, WorkflowNodeExecutionModel.__table__], + ) + return ArchiveDatabase( + session_maker=sessionmaker(bind=sqlite_engine, expire_on_commit=False), + end_before=datetime.datetime(2025, 4, 1, tzinfo=datetime.UTC), + ) + + @staticmethod + def test_resolve_archive_tenant_ids_from_plan_uses_fresh_real_sessions( + archive_db: ArchiveDatabase, monkeypatch: pytest.MonkeyPatch + ) -> None: + paid_a = _tenant_id("a", 1) + free_a = _tenant_id("a", 2) + paid_b = _tenant_id("b", 1) + free_b = _tenant_id("b", 2) + for tenant_id in (paid_a, free_a, paid_b, free_b): + archive_db.add_run(tenant_id) + + # These decoys verify the real candidate query's time, status, and type filters. + archive_db.add_run( + _tenant_id("a", 3), + created_at=archive_db.end_before + datetime.timedelta(seconds=1), + ) + archive_db.add_run(_tenant_id("a", 4), status=WorkflowExecutionStatus.RUNNING) + archive_db.add_run(_tenant_id("a", 5), run_type=WorkflowType.CHAT) + + opened_sessions: list[Session] = [] + + def record_session(session: Session, _transaction: object, _connection: object) -> None: + opened_sessions.append(session) + + event.listen(archive_db.session_maker.class_, "after_begin", record_session) + monkeypatch.setattr( + retention, + "_filter_paid_workflow_archive_tenant_ids", + lambda tenant_ids: ([paid_a, paid_b], sorted(set(tenant_ids) - {paid_a, paid_b})), + ) + try: + tenant_plan = retention._resolve_archive_tenant_ids_from_plan( + session_maker=archive_db.session_maker, + tenant_ids=None, + tenant_prefixes=["a", "b"], + start_from=None, + end_before=archive_db.end_before, + ) + finally: + event.remove(archive_db.session_maker.class_, "after_begin", record_session) + + assert tenant_plan == { + "archive_tenant_ids": [paid_a, paid_b], + "paid_tenant_ids": [paid_a, paid_b], + "unpaid_tenant_ids": [free_a, free_b], + } + assert len(opened_sessions) == 2 + assert opened_sessions[0] is not opened_sessions[1] + + @staticmethod + def test_safe_remove_scoped_session_recovers_from_real_closed_connection( + sqlite_engine: Engine, + monkeypatch: pytest.MonkeyPatch, + caplog: pytest.LogCaptureFixture, + ) -> None: + maker = sessionmaker(bind=sqlite_engine) + registry = scoped_session(maker) + registry().execute(text("select 1")) + dbapi_connection = registry().connection().connection.dbapi_connection + assert dbapi_connection is not None + dbapi_connection.close() + monkeypatch.setattr(retention, "db", SimpleNamespace(session=registry, engine=sqlite_engine)) + + with caplog.at_level(logging.WARNING, logger="commands.retention"): + retention._safe_remove_scoped_session("archive workflow run command") + + assert not registry.registry.has() + assert any("Ignoring DB scoped-session cleanup error" in message for message in caplog.messages) + + @staticmethod + def test_archive_command_db_retry_retries_retryable_db_disconnect(monkeypatch: pytest.MonkeyPatch) -> None: + attempts = iter([_db_disconnect_error(), "ok"]) + sleep = Mock() + monkeypatch.setattr("services.retention.workflow_run.db_retry.time.sleep", sleep) + + def operation() -> str: + result = next(attempts) + if isinstance(result, Exception): + raise result + return result + + assert retention._run_archive_command_db_retry("archive plan", operation) == "ok" + sleep.assert_called_once_with(1.0) + + @staticmethod + def test_archive_plan_prefix_stats_retries_with_fresh_session_and_real_counts( + archive_db: ArchiveDatabase, sqlite_engine: Engine, monkeypatch: pytest.MonkeyPatch + ) -> None: + tenant_id = _tenant_id("a", 1) + run_ids = [archive_db.add_run(tenant_id) for _ in range(7)] + for index in range(9): + archive_db.add_node(run_ids[index % len(run_ids)], tenant_id, index=index) + + # Decoys outside the selected prefix and archive window must not affect counts. + decoy_run_id = archive_db.add_run(_tenant_id("b", 1)) + archive_db.add_node(decoy_run_id, _tenant_id("b", 1), index=99) + archive_db.add_run( + tenant_id, + created_at=archive_db.end_before + datetime.timedelta(seconds=1), + ) + + fail_next_query = True + + def disconnect_once( + _connection: object, + _cursor: object, + _statement: str, + _parameters: object, + _context: object, + _executemany: bool, + ) -> None: + nonlocal fail_next_query + if fail_next_query: + fail_next_query = False + raise _db_disconnect_error() + + opened_sessions: list[Session] = [] + + def record_session(session: Session, _transaction: object, _connection: object) -> None: + opened_sessions.append(session) + + sleep = Mock() + monkeypatch.setattr("services.retention.workflow_run.db_retry.time.sleep", sleep) + event.listen(sqlite_engine, "before_cursor_execute", disconnect_once) + event.listen(archive_db.session_maker.class_, "after_begin", record_session) + try: + stats = retention._get_archive_plan_prefix_stats( + archive_db.session_maker, + "a", + start_from=None, + end_before=archive_db.end_before, + ) + finally: + event.remove(sqlite_engine, "before_cursor_execute", disconnect_once) + event.remove(archive_db.session_maker.class_, "after_begin", record_session) + + assert stats == { + "tenant_ids": [tenant_id], + "workflow_runs": 7, + "workflow_node_executions": 9, + } + assert len(opened_sessions) == 2 + assert opened_sessions[0] is not opened_sessions[1] + sleep.assert_called_once_with(1.0) + + @staticmethod + def test_archive_workflow_runs_raises_click_exception_when_tenant_plan_fails( + sqlite_engine: Engine, monkeypatch: pytest.MonkeyPatch + ) -> None: + registry = scoped_session(sessionmaker(bind=sqlite_engine)) + monkeypatch.setattr(retention, "db", SimpleNamespace(engine=sqlite_engine, session=registry)) + monkeypatch.setattr( + retention, + "_resolve_archive_tenant_ids_from_plan", + Mock(side_effect=RuntimeError("tenant plan failed")), + ) + + with pytest.raises(click.ClickException, match="Failed to resolve workflow archive tenant plan"): + retention.archive_workflow_runs.callback( + tenant_ids="tenant-1", + tenant_prefixes=None, + before_days=90, + from_days_ago=None, + to_days_ago=None, + start_from=None, + end_before=None, + batch_size=10000, + workers=1, + run_shard_index=None, + run_shard_total=None, + limit=None, + dry_run=True, + delete_after_archive=False, + ) + + @staticmethod + def test_delete_archived_workflow_runs_keeps_single_page_behavior_without_all_pages( + monkeypatch: pytest.MonkeyPatch, + ): + deleter = _patch_bundle_deleter( + monkeypatch, + [_delete_summary(processed=2, succeeded=2, next_catalog_id=_CURSOR_1)], + ) + + result = CliRunner().invoke( + retention.delete_archived_workflow_runs, + ["--target-month", "2025-03", "--limit", "2"], + ) + + assert result.exit_code == 0 + deleter.delete_batch.assert_called_once() + assert deleter.delete_batch.call_args.kwargs["target_year"] == 2025 + assert deleter.delete_batch.call_args.kwargs["target_month"] == 3 + assert deleter.delete_batch.call_args.kwargs["after_catalog_id"] is None + assert deleter.delete_batch.call_args.kwargs["limit"] == 2 + + @staticmethod + def test_delete_archived_workflow_runs_all_pages_continues_until_empty_page(monkeypatch: pytest.MonkeyPatch): + deleter = _patch_bundle_deleter( + monkeypatch, + [ + _delete_summary(processed=2, succeeded=2, next_catalog_id=_CURSOR_1), + _delete_summary(processed=1, succeeded=1, next_catalog_id=_CURSOR_2), + _delete_summary(processed=0), + ], + ) + + result = CliRunner().invoke( + retention.delete_archived_workflow_runs, + ["--target-month", "2025-03", "--all-pages", "--limit", "2"], + ) + + assert result.exit_code == 0 + assert [call.kwargs["after_catalog_id"] for call in deleter.delete_batch.call_args_list] == [ + None, + _CURSOR_1, + _CURSOR_2, + ] + + @staticmethod + def test_delete_archived_workflow_runs_all_pages_fetches_empty_page_after_exact_full_page( + monkeypatch: pytest.MonkeyPatch, + ): + deleter = _patch_bundle_deleter( + monkeypatch, + [ + _delete_summary(processed=2, succeeded=2, next_catalog_id=_CURSOR_1), + _delete_summary(processed=0), + ], + ) + + result = CliRunner().invoke( + retention.delete_archived_workflow_runs, + ["--target-month", "2025-03", "--all-pages", "--limit", "2"], + ) + + assert result.exit_code == 0 + assert deleter.delete_batch.call_count == 2 + assert deleter.delete_batch.call_args_list[1].kwargs["after_catalog_id"] == _CURSOR_1 + + @staticmethod + def test_delete_archived_workflow_runs_all_pages_stops_at_first_failed_page(monkeypatch: pytest.MonkeyPatch): + failed_result = BundleOperationResult( + catalog_id=_CURSOR_2, + bundle_id="bundle-failed", + tenant_id="tenant-1", + object_prefix="workflow-runs/v2/tenant-1/2025/03/00-of-16/bundle-failed", + error="archive checksum mismatch", + ) + deleter = _patch_bundle_deleter( + monkeypatch, + [ + _delete_summary(processed=1, succeeded=1, next_catalog_id=_CURSOR_1), + _delete_summary(processed=1, failed=1, results=[failed_result]), + _delete_summary(processed=0), + ], + ) + + result = CliRunner().invoke( + retention.delete_archived_workflow_runs, + ["--target-month", "2025-03", "--all-pages"], + ) + + assert result.exit_code == 1 + assert deleter.delete_batch.call_count == 2 + assert "target_month=2025-03" in result.output + assert f"failed_catalog_id={_CURSOR_2}" in result.output + assert f"resume_after_catalog_id={_CURSOR_1}" in result.output + + @staticmethod + def test_delete_archived_workflow_runs_all_pages_fails_when_cursor_does_not_advance( + monkeypatch: pytest.MonkeyPatch, + ): + deleter = _patch_bundle_deleter( + monkeypatch, + [_delete_summary(processed=1, succeeded=1, next_catalog_id=None)], + ) + + result = CliRunner().invoke( + retention.delete_archived_workflow_runs, + ["--target-month", "2025-03", "--all-pages"], + ) + + assert result.exit_code == 1 + deleter.delete_batch.assert_called_once() + assert "cursor did not advance" in result.output.lower() + + @staticmethod + def test_delete_archived_workflow_runs_all_pages_uses_preview_cursor_for_dry_run(monkeypatch: pytest.MonkeyPatch): + deleter = _patch_bundle_deleter( + monkeypatch, + [ + _delete_summary(processed=1, succeeded=1, preview_next_catalog_id=_CURSOR_1), + _delete_summary(processed=0), + ], + ) + + result = CliRunner().invoke( + retention.delete_archived_workflow_runs, + ["--target-month", "2025-03", "--all-pages", "--dry-run"], + ) + + assert result.exit_code == 0 + assert [call.kwargs["after_catalog_id"] for call in deleter.delete_batch.call_args_list] == [ + None, + _CURSOR_1, + ] + + @staticmethod + def test_delete_archived_workflow_runs_dry_run_failure_separates_preview_and_destructive_cursors( + monkeypatch: pytest.MonkeyPatch, + ): + failed_result = BundleOperationResult( + catalog_id=_CURSOR_2, + bundle_id="bundle-failed", + tenant_id="tenant-1", + object_prefix="workflow-runs/v2/tenant-1/2025/03/00-of-16/bundle-failed", + error="archive checksum mismatch", + ) + deleter = _patch_bundle_deleter( + monkeypatch, + [ + _delete_summary( + processed=1, + succeeded=1, + preview_next_catalog_id=_CURSOR_1, + ), + _delete_summary( + processed=1, + failed=1, + results=[failed_result], + ), + ], + ) + + result = CliRunner().invoke( + retention.delete_archived_workflow_runs, + [ + "--target-month", + "2025-03", + "--after-catalog-id", + _CURSOR_0, + "--all-pages", + "--dry-run", + ], + ) + + assert result.exit_code == 1 + assert deleter.delete_batch.call_count == 2 + assert f"failed_catalog_id={_CURSOR_2}" in result.output + assert f"preview_after_catalog_id={_CURSOR_1}" in result.output + assert f"destructive_retry_after_catalog_id={_CURSOR_0}" in result.output + + @staticmethod + def test_delete_archived_workflow_runs_all_pages_starts_after_explicit_cursor(monkeypatch: pytest.MonkeyPatch): + deleter = _patch_bundle_deleter(monkeypatch, [_delete_summary(processed=0)]) + + result = CliRunner().invoke( + retention.delete_archived_workflow_runs, + [ + "--target-month", + "2025-03", + "--after-catalog-id", + _CURSOR_0, + "--all-pages", + ], + ) + + assert result.exit_code == 0 + deleter.delete_batch.assert_called_once() + assert deleter.delete_batch.call_args.kwargs["after_catalog_id"] == _CURSOR_0 + + @staticmethod + @pytest.mark.parametrize( + "shard_args", + [ + ["--run-shard-index", "0"], + ["--run-shard-total", "16"], + ["--run-shard-index", "16", "--run-shard-total", "16"], + ["--run-shard-index", "-1", "--run-shard-total", "16"], + ["--run-shard-index", "0", "--run-shard-total", "0"], + ["--run-shard-index", "0", "--run-shard-total", "17"], + ], + ) + def test_delete_archived_workflow_runs_rejects_invalid_run_shard_options(monkeypatch, shard_args): + deleter = _patch_bundle_deleter(monkeypatch, [_delete_summary(processed=0)]) + + result = CliRunner().invoke( + retention.delete_archived_workflow_runs, + ["--target-month", "2025-03", *shard_args], + ) + + assert result.exit_code == 2 + deleter.delete_batch.assert_not_called() + + @staticmethod + def test_delete_archived_workflow_runs_passes_formatted_run_shard_to_service(monkeypatch: pytest.MonkeyPatch): + deleter = _patch_bundle_deleter(monkeypatch, [_delete_summary(processed=0)]) + + result = CliRunner().invoke( + retention.delete_archived_workflow_runs, + [ + "--target-month", + "2025-03", + "--tenant-ids", + "tenant-1", + "--run-shard-index", + "3", + "--run-shard-total", + "16", + ], + ) + + assert result.exit_code == 0 + deleter.validate_catalog_shards.assert_called_once_with( + target_year=2025, + target_month=3, + shard_total=16, + tenant_ids=["tenant-1"], + ) + deleter.delete_batch.assert_called_once() + assert deleter.delete_batch.call_args.kwargs["shard"] == "03-of-16" + + @staticmethod + def test_delete_archived_workflow_runs_rejects_mixed_catalog_shards_before_delete(monkeypatch: pytest.MonkeyPatch): + deleter = _patch_bundle_deleter(monkeypatch, [_delete_summary(processed=0)]) + deleter.validate_catalog_shards.side_effect = ValueError("unexpected shards: 00-of-01") + + result = CliRunner().invoke( + retention.delete_archived_workflow_runs, + [ + "--target-month", + "2025-03", + "--run-shard-index", + "3", + "--run-shard-total", + "16", + ], + ) + + assert result.exit_code == 1 + assert "shard preflight failed" in result.output.lower() + assert "00-of-01" in result.output + deleter.delete_batch.assert_not_called() + + +def _mock_service() -> MagicMock: + service = MagicMock() + service.run.return_value = { + "batches": 1, + "total_messages": 10, + "filtered_messages": 5, + "total_deleted": 5, + } + return service + + +def _invoke_clean_expired_messages(**kwargs: object) -> None: + callback = clean_expired_messages.callback + assert callback is not None + callback(**kwargs) + + +class TestCleanExpiredMessagesCases: + @staticmethod + def test_absolute_mode_calls_from_time_range() -> None: + policy = object() + service = _mock_service() + start_from = datetime.datetime(2024, 1, 1, 0, 0, 0) + end_before = datetime.datetime(2024, 2, 1, 0, 0, 0) + + with ( + patch("commands.retention.create_message_clean_policy", return_value=policy), + patch( + "commands.retention.MessagesCleanService.from_time_range", return_value=service + ) as mock_from_time_range, + patch("commands.retention.MessagesCleanService.from_days") as mock_from_days, + ): + _invoke_clean_expired_messages( + batch_size=200, + graceful_period=21, + start_from=start_from, + end_before=end_before, + from_days_ago=None, + before_days=None, + dry_run=True, + ) + + mock_from_time_range.assert_called_once_with( + policy=policy, + start_from=start_from, + end_before=end_before, + batch_size=200, + dry_run=True, + task_label="custom", + ) + mock_from_days.assert_not_called() + + @staticmethod + def test_relative_mode_before_days_only_calls_from_days() -> None: + policy = object() + service = _mock_service() + + with ( + patch("commands.retention.create_message_clean_policy", return_value=policy), + patch("commands.retention.MessagesCleanService.from_days", return_value=service) as mock_from_days, + patch("commands.retention.MessagesCleanService.from_time_range") as mock_from_time_range, + ): + _invoke_clean_expired_messages( + batch_size=500, + graceful_period=14, + start_from=None, + end_before=None, + from_days_ago=None, + before_days=30, + dry_run=False, + ) + + mock_from_days.assert_called_once_with( + policy=policy, + days=30, + batch_size=500, + dry_run=False, + task_label="before-30", + ) + mock_from_time_range.assert_not_called() + + @staticmethod + def test_relative_mode_with_from_days_ago_calls_from_time_range() -> None: + policy = object() + service = _mock_service() + fixed_now = datetime.datetime(2024, 8, 20, 12, 0, 0) + + with ( + patch("commands.retention.create_message_clean_policy", return_value=policy), + patch( + "commands.retention.MessagesCleanService.from_time_range", return_value=service + ) as mock_from_time_range, + patch("commands.retention.MessagesCleanService.from_days") as mock_from_days, + patch("commands.retention.naive_utc_now", return_value=fixed_now), + ): + _invoke_clean_expired_messages( + batch_size=1000, + graceful_period=21, + start_from=None, + end_before=None, + from_days_ago=60, + before_days=30, + dry_run=False, + ) + + mock_from_time_range.assert_called_once_with( + policy=policy, + start_from=fixed_now - datetime.timedelta(days=60), + end_before=fixed_now - datetime.timedelta(days=30), + batch_size=1000, + dry_run=False, + task_label="60to30", + ) + mock_from_days.assert_not_called() + + @staticmethod + @pytest.mark.parametrize( + ("kwargs", "message"), + [ + ( + { + "start_from": datetime.datetime(2024, 1, 1), + "end_before": datetime.datetime(2024, 2, 1), + "from_days_ago": None, + "before_days": 30, + }, + "mutually exclusive", + ), + ( + { + "start_from": datetime.datetime(2024, 1, 1), + "end_before": None, + "from_days_ago": None, + "before_days": None, + }, + "Both --start-from and --end-before are required", + ), + ( + { + "start_from": None, + "end_before": None, + "from_days_ago": 10, + "before_days": None, + }, + "--from-days-ago must be used together with --before-days", + ), + ( + { + "start_from": None, + "end_before": None, + "from_days_ago": None, + "before_days": -1, + }, + "--before-days must be >= 0", + ), + ( + { + "start_from": None, + "end_before": None, + "from_days_ago": 30, + "before_days": 30, + }, + "--from-days-ago must be greater than --before-days", + ), + ( + { + "start_from": None, + "end_before": None, + "from_days_ago": None, + "before_days": None, + }, + "You must provide either (--start-from,--end-before) or (--before-days [--from-days-ago])", + ), + ], + ) + def test_invalid_inputs_raise_usage_error(kwargs: dict[str, object], message: str) -> None: + with pytest.raises(click.UsageError, match=re.escape(message)): + _invoke_clean_expired_messages( + batch_size=1000, + graceful_period=21, + start_from=kwargs["start_from"], + end_before=kwargs["end_before"], + from_days_ago=kwargs["from_days_ago"], + before_days=kwargs["before_days"], + dry_run=False, + ) diff --git a/api/tests/unit_tests/commands/test_system.py b/api/tests/unit_tests/commands/test_system.py new file mode 100644 index 00000000000..8a788cb8f1a --- /dev/null +++ b/api/tests/unit_tests/commands/test_system.py @@ -0,0 +1,455 @@ +import sys +import threading +import types +import uuid +from types import SimpleNamespace +from unittest.mock import MagicMock + +import pytest +from sqlalchemy import event, select +from sqlalchemy.engine import Engine +from sqlalchemy.orm import Session, scoped_session, sessionmaker + +import commands +from commands import system as system_commands +from core.tools.entities.tool_entities import ApiProviderSchemaType +from enums import DeploymentEdition +from graphon.model_runtime.entities.model_entities import ModelType +from libs.db_migration_lock import LockNotOwnedError, RedisError +from models import Tenant as _reset_encrypt_key_pair_Tenant +from models.account import Account, TenantAccountJoin, TenantAccountRole +from models.account import Tenant as _fix_app_site_missing_Tenant +from models.enums import CustomizeTokenStrategy +from models.model import App, AppMode, IconType, Site +from models.provider import Provider, ProviderModel, ProviderType +from models.tools import ApiToolProvider, BuiltinToolProvider, MCPToolProvider +from tests.helpers.unit.config_override import apply_config_overrides +from tests.helpers.unit.model_factories import make_tenant + + +def _persist_missing_site_owner(session: Session) -> tuple[Account, App]: + """Persist an app without a Site and its complete tenant owner chain.""" + tenant = _fix_app_site_missing_Tenant(name="Command workspace") + account = Account(name="Owner", email=f"owner-{uuid.uuid4()}@example.com") + membership = TenantAccountJoin( + tenant_id=tenant.id, + account_id=account.id, + current=True, + role=TenantAccountRole.OWNER, + ) + app = App( + id=str(uuid.uuid4()), + tenant_id=tenant.id, + name="Missing Site App", + mode=AppMode.CHAT, + icon_type=IconType.EMOJI, + icon="chat", + icon_background="#FFFFFF", + enable_site=True, + enable_api=False, + created_by=account.id, + ) + session.add_all([tenant, account, membership, app]) + session.commit() + return account, app + + +def _site_for(app: App) -> Site: + return Site( + app_id=app.id, + title=app.name, + default_language="en-US", + customize_token_strategy=CustomizeTokenStrategy.UUID, + code=f"site-{app.id}", + ) + + +def _bind_command_database( + monkeypatch: pytest.MonkeyPatch, + sqlite_engine: Engine, + sqlite_session_factory: sessionmaker[Session], +) -> tuple[scoped_session[Session], Session]: + """Expose a callable real scoped session through the Flask-SQLAlchemy shape.""" + command_sessions = scoped_session(sqlite_session_factory) + command_session = command_sessions() + monkeypatch.setattr( + system_commands, + "db", + SimpleNamespace(engine=sqlite_engine, session=command_sessions), + ) + return command_sessions, command_session + + +def _run_fix_app_site_missing() -> None: + """Invoke the command callback through a narrowed, non-optional reference.""" + callback = system_commands.fix_app_site_missing.callback + assert callback is not None + callback() + + +def test_fix_app_site_missing_passes_loaded_session_to_signal( + monkeypatch: pytest.MonkeyPatch, + sqlite_engine: Engine, + sqlite_session: Session, + sqlite_session_factory: sessionmaker[Session], +) -> None: + account, app = _persist_missing_site_owner(sqlite_session) + command_sessions, command_session = _bind_command_database(monkeypatch, sqlite_engine, sqlite_session_factory) + phase_events: list[str] = [] + event.listen(command_session, "after_commit", lambda _session: phase_events.append("commit")) + + def create_site(sender: App, *, account: Account, session: Session) -> None: + phase_events.append("signal") + assert sender.id == app.id + assert account.id == account_id + assert session is command_session + session.add(_site_for(sender)) + + account_id = account.id + send = MagicMock(side_effect=create_site) + monkeypatch.setattr(system_commands.app_was_created, "send", send) + + try: + _run_fix_app_site_missing() + finally: + command_sessions.remove() + + send.assert_called_once() + assert phase_events == ["signal", "commit"] + sqlite_session.expire_all() + persisted_site = sqlite_session.query(Site).filter_by(app_id=app.id).one() + assert persisted_site.title == app.name + + +def test_fix_app_site_missing_rolls_back_when_signal_fails( + monkeypatch: pytest.MonkeyPatch, + sqlite_engine: Engine, + sqlite_session: Session, + sqlite_session_factory: sessionmaker[Session], +) -> None: + _account, app = _persist_missing_site_owner(sqlite_session) + command_sessions, command_session = _bind_command_database(monkeypatch, sqlite_engine, sqlite_session_factory) + phase_events: list[str] = [] + event.listen(command_session, "after_rollback", lambda _session: phase_events.append("rollback")) + + def fail_signal(sender: App, **_kwargs: object) -> None: + phase_events.append("signal") + # Ensure the command's next raw scan terminates while its own transaction + # still exercises the rollback path. + with sqlite_session_factory() as observer: + observer.add(_site_for(sender)) + observer.commit() + raise RuntimeError("failed") + + monkeypatch.setattr(system_commands.app_was_created, "send", MagicMock(side_effect=fail_signal)) + + try: + _run_fix_app_site_missing() + finally: + command_sessions.remove() + + assert phase_events == ["signal", "rollback"] + sqlite_session.expire_all() + assert sqlite_session.get(App, app.id) is not None + + +"""SQLite-backed tests for the reset-encrypt-key-pair CLI command (#35396). + +The command must purge every table that stores ciphertext encrypted with the +tenant's asymmetric key, otherwise stale rows cause downstream API failures +such as `/console/api/workspaces/current/tool-providers` returning 500. +Tests bind the command-owned transaction to the fixture engine and assert the +committed state rather than inspecting fabricated ``Session.execute`` calls. +""" + + +def _invoke_reset() -> int: + callback = commands.reset_encrypt_key_pair.callback + assert callback is not None + try: + callback() + except SystemExit as e: + return int(e.code or 0) + return 0 + + +TENANT_ID = "11111111-1111-1111-1111-111111111111" +OTHER_TENANT_ID = "11111111-1111-1111-1111-111111111112" +USER_ID = "22222222-2222-2222-2222-222222222222" + + +def _tenant(tenant_id: str, *, name: str = "Test tenant") -> _reset_encrypt_key_pair_Tenant: + return make_tenant(tenant_id=tenant_id, name=name, encrypt_public_key="old-key") + + +def _encrypted_rows(tenant_id: str, *, suffix: str = "1") -> tuple[object, ...]: + """Build one persisted credential-bearing row for every purge target.""" + return ( + Provider(tenant_id=tenant_id, provider_name=f"provider-{suffix}"), + ProviderModel( + tenant_id=tenant_id, + provider_name=f"provider-{suffix}", + model_name=f"model-{suffix}", + model_type=ModelType.LLM, + ), + BuiltinToolProvider( + name=f"builtin-credential-{suffix}", + tenant_id=tenant_id, + user_id=USER_ID, + provider=f"builtin-{suffix}", + encrypted_credentials="ciphertext", + ), + ApiToolProvider( + name=f"api-{suffix}", + icon="icon", + schema="{}", + schema_type_str=ApiProviderSchemaType.OPENAPI, + user_id=USER_ID, + tenant_id=tenant_id, + description="description", + tools_str="[]", + credentials_str="{}", + ), + MCPToolProvider( + name=f"mcp-{suffix}", + server_identifier=f"server-{suffix}", + server_url="ciphertext", + server_url_hash=f"hash-{suffix}", + icon=None, + tenant_id=tenant_id, + user_id=USER_ID, + encrypted_credentials="ciphertext", + ), + ) + + +def _bind_command_to_sqlite(monkeypatch: pytest.MonkeyPatch, session: Session) -> None: + monkeypatch.setattr(system_commands, "db", SimpleNamespace(engine=session.get_bind())) + + +def test_reset_aborts_when_not_self_hosted(monkeypatch: pytest.MonkeyPatch, capsys: pytest.CaptureFixture[str]) -> None: + apply_config_overrides(monkeypatch, DEPLOYMENT_EDITION=DeploymentEdition.CLOUD) + + exit_code = _invoke_reset() + captured = capsys.readouterr() + + assert exit_code == 0 + assert "only for SELF_HOSTED" in captured.out + + +@pytest.mark.parametrize( + "sqlite_session", + [(_reset_encrypt_key_pair_Tenant, Provider, ProviderModel, BuiltinToolProvider, ApiToolProvider, MCPToolProvider)], + indirect=True, +) +def test_reset_purges_provider_and_tool_tables_for_each_tenant( + monkeypatch: pytest.MonkeyPatch, capsys: pytest.CaptureFixture[str], sqlite_session: Session +) -> None: + """The command must purge LLM provider rows AND every tool provider table + that stores ciphertext encrypted under the tenant key (#35396).""" + apply_config_overrides(monkeypatch, DEPLOYMENT_EDITION=DeploymentEdition.COMMUNITY) + monkeypatch.setattr(system_commands, "generate_key_pair", lambda tenant_id: f"new-key-{tenant_id}") + _bind_command_to_sqlite(monkeypatch, sqlite_session) + + tenant = _tenant(TENANT_ID) + other_tenant = _tenant(OTHER_TENANT_ID, name="Other tenant") + system_provider = Provider( + tenant_id=TENANT_ID, + provider_name="system-provider", + provider_type=ProviderType.SYSTEM, + ) + sqlite_session.add_all((tenant, other_tenant, system_provider, *_encrypted_rows(TENANT_ID))) + sqlite_session.commit() + + exit_code = _invoke_reset() + + captured = capsys.readouterr() + assert exit_code == 0 + assert TENANT_ID in captured.out + + sqlite_session.expire_all() + tenant_after = sqlite_session.get(_reset_encrypt_key_pair_Tenant, TENANT_ID) + other_tenant_after = sqlite_session.get(_reset_encrypt_key_pair_Tenant, OTHER_TENANT_ID) + assert tenant_after is not None + assert other_tenant_after is not None + assert tenant_after.encrypt_public_key == f"new-key-{TENANT_ID}" + assert other_tenant_after.encrypt_public_key == f"new-key-{OTHER_TENANT_ID}" + assert sqlite_session.scalars(select(Provider).where(Provider.provider_type == ProviderType.CUSTOM)).all() == [] + assert sqlite_session.scalars(select(ProviderModel)).all() == [] + assert sqlite_session.scalars(select(BuiltinToolProvider)).all() == [] + assert sqlite_session.scalars(select(ApiToolProvider)).all() == [] + assert sqlite_session.scalars(select(MCPToolProvider)).all() == [] + assert ( + sqlite_session.scalar(select(Provider).where(Provider.provider_type == ProviderType.SYSTEM)) is system_provider + ) + + +@pytest.mark.parametrize( + "sqlite_session", + [(_reset_encrypt_key_pair_Tenant, Provider, ProviderModel, BuiltinToolProvider, ApiToolProvider, MCPToolProvider)], + indirect=True, +) +def test_reset_iterates_all_tenants(monkeypatch: pytest.MonkeyPatch, sqlite_session: Session) -> None: + """Multi-tenant deployments must purge every tenant, not just the first.""" + apply_config_overrides(monkeypatch, DEPLOYMENT_EDITION=DeploymentEdition.COMMUNITY) + monkeypatch.setattr(system_commands, "generate_key_pair", lambda tenant_id: f"new-key-{tenant_id}") + + _bind_command_to_sqlite(monkeypatch, sqlite_session) + tenant_ids = [f"11111111-1111-1111-1111-{index:012d}" for index in range(3)] + tenants = [_tenant(tenant_id, name=f"Tenant {index}") for index, tenant_id in enumerate(tenant_ids)] + for index, tenant in enumerate(tenants): + sqlite_session.add(tenant) + sqlite_session.add_all(_encrypted_rows(tenant.id, suffix=str(index))) + sqlite_session.commit() + + assert _invoke_reset() == 0 + + sqlite_session.expire_all() + persisted_tenants = sqlite_session.scalars( + select(_reset_encrypt_key_pair_Tenant).order_by(_reset_encrypt_key_pair_Tenant.id) + ).all() + assert [tenant.encrypt_public_key for tenant in persisted_tenants] == [ + f"new-key-{tenant_id}" for tenant_id in tenant_ids + ] + for model in (Provider, ProviderModel, BuiltinToolProvider, ApiToolProvider, MCPToolProvider): + assert sqlite_session.scalars(select(model)).all() == [] + + +HEARTBEAT_WAIT_TIMEOUT_SECONDS = 5.0 + + +def _install_fake_flask_migrate(monkeypatch, upgrade_impl) -> None: + module = types.ModuleType("flask_migrate") + module.upgrade = upgrade_impl + monkeypatch.setitem(sys.modules, "flask_migrate", module) + + +def _invoke_upgrade_db() -> int: + try: + commands.upgrade_db.callback() + except SystemExit as e: + return int(e.code or 0) + return 0 + + +def test_upgrade_db_skips_when_lock_not_acquired(monkeypatch, capsys): + monkeypatch.setattr(system_commands, "DB_UPGRADE_LOCK_TTL_SECONDS", 1234) + + lock = MagicMock() + lock.acquire.return_value = False + system_commands.redis_client.lock.return_value = lock + + exit_code = _invoke_upgrade_db() + captured = capsys.readouterr() + + assert exit_code == 0 + assert "Database migration skipped" in captured.out + + system_commands.redis_client.lock.assert_called_once_with(name="db_upgrade_lock", timeout=1234, thread_local=False) + lock.acquire.assert_called_once_with(blocking=False) + lock.release.assert_not_called() + + +def test_upgrade_db_failure_not_masked_by_lock_release(monkeypatch, capsys): + monkeypatch.setattr(system_commands, "DB_UPGRADE_LOCK_TTL_SECONDS", 321) + + lock = MagicMock() + lock.acquire.return_value = True + lock.release.side_effect = LockNotOwnedError("simulated") + system_commands.redis_client.lock.return_value = lock + + def _upgrade(): + raise RuntimeError("boom") + + _install_fake_flask_migrate(monkeypatch, _upgrade) + + exit_code = _invoke_upgrade_db() + captured = capsys.readouterr() + + assert exit_code == 1 + assert "Database migration failed: boom" in captured.out + + system_commands.redis_client.lock.assert_called_once_with(name="db_upgrade_lock", timeout=321, thread_local=False) + lock.acquire.assert_called_once_with(blocking=False) + lock.release.assert_called_once() + + +def test_upgrade_db_success_ignores_lock_not_owned_on_release(monkeypatch, capsys): + monkeypatch.setattr(system_commands, "DB_UPGRADE_LOCK_TTL_SECONDS", 999) + + lock = MagicMock() + lock.acquire.return_value = True + lock.release.side_effect = LockNotOwnedError("simulated") + system_commands.redis_client.lock.return_value = lock + + _install_fake_flask_migrate(monkeypatch, lambda: None) + + exit_code = _invoke_upgrade_db() + captured = capsys.readouterr() + + assert exit_code == 0 + assert "Database migration successful!" in captured.out + + system_commands.redis_client.lock.assert_called_once_with(name="db_upgrade_lock", timeout=999, thread_local=False) + lock.acquire.assert_called_once_with(blocking=False) + lock.release.assert_called_once() + + +def test_upgrade_db_renews_lock_during_migration(monkeypatch, capsys): + """ + Ensure the lock is renewed while migrations are running, so the base TTL can stay short. + """ + + # Use a small TTL so the heartbeat interval triggers quickly. + monkeypatch.setattr(system_commands, "DB_UPGRADE_LOCK_TTL_SECONDS", 0.3) + + lock = MagicMock() + lock.acquire.return_value = True + system_commands.redis_client.lock.return_value = lock + + renewed = threading.Event() + + def _reacquire(): + renewed.set() + return True + + lock.reacquire.side_effect = _reacquire + + def _upgrade(): + assert renewed.wait(HEARTBEAT_WAIT_TIMEOUT_SECONDS) + + _install_fake_flask_migrate(monkeypatch, _upgrade) + + exit_code = _invoke_upgrade_db() + _ = capsys.readouterr() + + assert exit_code == 0 + assert lock.reacquire.call_count >= 1 + + +def test_upgrade_db_ignores_reacquire_errors(monkeypatch, capsys): + # Use a small TTL so heartbeat runs during the upgrade call. + monkeypatch.setattr(system_commands, "DB_UPGRADE_LOCK_TTL_SECONDS", 0.3) + + lock = MagicMock() + lock.acquire.return_value = True + system_commands.redis_client.lock.return_value = lock + + attempted = threading.Event() + + def _reacquire(): + attempted.set() + raise RedisError("simulated") + + lock.reacquire.side_effect = _reacquire + + def _upgrade(): + assert attempted.wait(HEARTBEAT_WAIT_TIMEOUT_SECONDS) + + _install_fake_flask_migrate(monkeypatch, _upgrade) + + exit_code = _invoke_upgrade_db() + _ = capsys.readouterr() + + assert exit_code == 0 + assert lock.reacquire.call_count >= 1 diff --git a/api/tests/unit_tests/commands/test_upgrade_db.py b/api/tests/unit_tests/commands/test_upgrade_db.py deleted file mode 100644 index 5aa03134299..00000000000 --- a/api/tests/unit_tests/commands/test_upgrade_db.py +++ /dev/null @@ -1,147 +0,0 @@ -import sys -import threading -import types -from unittest.mock import MagicMock - -import commands -from commands import system as system_commands -from libs.db_migration_lock import LockNotOwnedError, RedisError - -HEARTBEAT_WAIT_TIMEOUT_SECONDS = 5.0 - - -def _install_fake_flask_migrate(monkeypatch, upgrade_impl) -> None: - module = types.ModuleType("flask_migrate") - module.upgrade = upgrade_impl - monkeypatch.setitem(sys.modules, "flask_migrate", module) - - -def _invoke_upgrade_db() -> int: - try: - commands.upgrade_db.callback() - except SystemExit as e: - return int(e.code or 0) - return 0 - - -def test_upgrade_db_skips_when_lock_not_acquired(monkeypatch, capsys): - monkeypatch.setattr(system_commands, "DB_UPGRADE_LOCK_TTL_SECONDS", 1234) - - lock = MagicMock() - lock.acquire.return_value = False - system_commands.redis_client.lock.return_value = lock - - exit_code = _invoke_upgrade_db() - captured = capsys.readouterr() - - assert exit_code == 0 - assert "Database migration skipped" in captured.out - - system_commands.redis_client.lock.assert_called_once_with(name="db_upgrade_lock", timeout=1234, thread_local=False) - lock.acquire.assert_called_once_with(blocking=False) - lock.release.assert_not_called() - - -def test_upgrade_db_failure_not_masked_by_lock_release(monkeypatch, capsys): - monkeypatch.setattr(system_commands, "DB_UPGRADE_LOCK_TTL_SECONDS", 321) - - lock = MagicMock() - lock.acquire.return_value = True - lock.release.side_effect = LockNotOwnedError("simulated") - system_commands.redis_client.lock.return_value = lock - - def _upgrade(): - raise RuntimeError("boom") - - _install_fake_flask_migrate(monkeypatch, _upgrade) - - exit_code = _invoke_upgrade_db() - captured = capsys.readouterr() - - assert exit_code == 1 - assert "Database migration failed: boom" in captured.out - - system_commands.redis_client.lock.assert_called_once_with(name="db_upgrade_lock", timeout=321, thread_local=False) - lock.acquire.assert_called_once_with(blocking=False) - lock.release.assert_called_once() - - -def test_upgrade_db_success_ignores_lock_not_owned_on_release(monkeypatch, capsys): - monkeypatch.setattr(system_commands, "DB_UPGRADE_LOCK_TTL_SECONDS", 999) - - lock = MagicMock() - lock.acquire.return_value = True - lock.release.side_effect = LockNotOwnedError("simulated") - system_commands.redis_client.lock.return_value = lock - - _install_fake_flask_migrate(monkeypatch, lambda: None) - - exit_code = _invoke_upgrade_db() - captured = capsys.readouterr() - - assert exit_code == 0 - assert "Database migration successful!" in captured.out - - system_commands.redis_client.lock.assert_called_once_with(name="db_upgrade_lock", timeout=999, thread_local=False) - lock.acquire.assert_called_once_with(blocking=False) - lock.release.assert_called_once() - - -def test_upgrade_db_renews_lock_during_migration(monkeypatch, capsys): - """ - Ensure the lock is renewed while migrations are running, so the base TTL can stay short. - """ - - # Use a small TTL so the heartbeat interval triggers quickly. - monkeypatch.setattr(system_commands, "DB_UPGRADE_LOCK_TTL_SECONDS", 0.3) - - lock = MagicMock() - lock.acquire.return_value = True - system_commands.redis_client.lock.return_value = lock - - renewed = threading.Event() - - def _reacquire(): - renewed.set() - return True - - lock.reacquire.side_effect = _reacquire - - def _upgrade(): - assert renewed.wait(HEARTBEAT_WAIT_TIMEOUT_SECONDS) - - _install_fake_flask_migrate(monkeypatch, _upgrade) - - exit_code = _invoke_upgrade_db() - _ = capsys.readouterr() - - assert exit_code == 0 - assert lock.reacquire.call_count >= 1 - - -def test_upgrade_db_ignores_reacquire_errors(monkeypatch, capsys): - # Use a small TTL so heartbeat runs during the upgrade call. - monkeypatch.setattr(system_commands, "DB_UPGRADE_LOCK_TTL_SECONDS", 0.3) - - lock = MagicMock() - lock.acquire.return_value = True - system_commands.redis_client.lock.return_value = lock - - attempted = threading.Event() - - def _reacquire(): - attempted.set() - raise RedisError("simulated") - - lock.reacquire.side_effect = _reacquire - - def _upgrade(): - assert attempted.wait(HEARTBEAT_WAIT_TIMEOUT_SECONDS) - - _install_fake_flask_migrate(monkeypatch, _upgrade) - - exit_code = _invoke_upgrade_db() - _ = capsys.readouterr() - - assert exit_code == 0 - assert lock.reacquire.call_count >= 1 diff --git a/api/tests/unit_tests/configs/test_agent_backend_config.py b/api/tests/unit_tests/configs/extra/test_agent_backend_config.py similarity index 97% rename from api/tests/unit_tests/configs/test_agent_backend_config.py rename to api/tests/unit_tests/configs/extra/test_agent_backend_config.py index 542e2934ced..b6e1bc51c42 100644 --- a/api/tests/unit_tests/configs/test_agent_backend_config.py +++ b/api/tests/unit_tests/configs/extra/test_agent_backend_config.py @@ -7,7 +7,7 @@ from pydantic import ValidationError from configs.extra.agent_backend_config import AgentBackendConfig -_REPOSITORY_ROOT = Path(__file__).resolve().parents[4] +_REPOSITORY_ROOT = Path(__file__).resolve().parents[5] _API_TIMEOUT_ENV = "AGENT_BACKEND_BINDING_FILE_DOWNLOAD_TIMEOUT_SECONDS" _AGENT_TIMEOUT_ENV = "DIFY_AGENT_BINDING_FILE_DOWNLOAD_COMMAND_TIMEOUT_SECONDS" diff --git a/api/tests/unit_tests/configs/test_knowledge_fs_config.py b/api/tests/unit_tests/configs/extra/test_knowledge_fs_config.py similarity index 98% rename from api/tests/unit_tests/configs/test_knowledge_fs_config.py rename to api/tests/unit_tests/configs/extra/test_knowledge_fs_config.py index c60a4e00c1e..1723de1be8c 100644 --- a/api/tests/unit_tests/configs/test_knowledge_fs_config.py +++ b/api/tests/unit_tests/configs/extra/test_knowledge_fs_config.py @@ -7,7 +7,7 @@ from pydantic import ValidationError from configs.app_config import DifyConfig from configs.extra.knowledge_fs_config import KnowledgeFSConfig as _KnowledgeFSConfig -from tests.unit_tests.configs._isolated_settings import InitSettingsOnly +from tests.helpers.unit.configs._isolated_settings import InitSettingsOnly class KnowledgeFSConfig(InitSettingsOnly, _KnowledgeFSConfig): @@ -18,7 +18,7 @@ class _DeployedKnowledgeFSConfig(InitSettingsOnly, DifyConfig): DEPLOY_ENV: str = "PRODUCTION" -_REPOSITORY_ROOT = Path(__file__).resolve().parents[4] +_REPOSITORY_ROOT = Path(__file__).resolve().parents[5] _KNOWLEDGE_FS_DOCKER_VARIABLES = ( "KNOWLEDGE_FS_ENABLED", "KNOWLEDGE_FS_COMMUNITY_DEV_ENABLED", diff --git a/api/tests/unit_tests/configs/test_logstore_config.py b/api/tests/unit_tests/configs/extra/test_logstore_config.py similarity index 90% rename from api/tests/unit_tests/configs/test_logstore_config.py rename to api/tests/unit_tests/configs/extra/test_logstore_config.py index 960e50b5a3a..9001f5db04e 100644 --- a/api/tests/unit_tests/configs/test_logstore_config.py +++ b/api/tests/unit_tests/configs/extra/test_logstore_config.py @@ -1,7 +1,7 @@ import pytest from configs.extra.logstore_config import LogStoreConfig -from tests.unit_tests.configs._isolated_settings import InitSettingsOnly +from tests.helpers.unit.configs._isolated_settings import InitSettingsOnly class _IsolatedLogStoreConfig(InitSettingsOnly, LogStoreConfig): diff --git a/api/tests/unit_tests/configs/feature/test___init__.py b/api/tests/unit_tests/configs/feature/test___init__.py new file mode 100644 index 00000000000..d21c2cbcd75 --- /dev/null +++ b/api/tests/unit_tests/configs/feature/test___init__.py @@ -0,0 +1,45 @@ +import pytest + +from configs.feature import AuthConfig, FileUploadConfig +from tests.helpers.unit.configs._isolated_settings import InitSettingsOnly + + +class _IsolatedAuthConfig(InitSettingsOnly, AuthConfig): + pass + + +def test_oauth_ttl_days_defaults_to_fourteen() -> None: + config = _IsolatedAuthConfig() + + assert config.OAUTH_TTL_DAYS == 14 + + +@pytest.mark.parametrize( + ("configured", "expected"), + [("30", 30), ("1", 1), ("365", 365), ("0", 1), ("-5", 1), ("500", 365), ("invalid", 14), ("", 14)], +) +def test_oauth_ttl_days_clamps_or_falls_back_instead_of_failing(configured: str, expected: int) -> None: + config = _IsolatedAuthConfig(OAUTH_TTL_DAYS=configured) + + assert expected == config.OAUTH_TTL_DAYS + + +class _IsolatedFileUploadConfig(InitSettingsOnly, FileUploadConfig): + pass + + +def test_paid_plan_file_size_limit_uses_its_default() -> None: + config = _IsolatedFileUploadConfig(UPLOAD_FILE_SIZE_LIMIT="23") + + assert config.UPLOAD_FILE_SIZE_LIMIT == 23 + assert config.KNOWLEDGE_UPLOAD_FILE_SIZE_LIMIT_FOR_PAID_PLAN == 15 + + +def test_paid_plan_file_size_limit_can_be_configured_separately() -> None: + config = _IsolatedFileUploadConfig( + UPLOAD_FILE_SIZE_LIMIT="23", + KNOWLEDGE_UPLOAD_FILE_SIZE_LIMIT_FOR_PAID_PLAN="50", + ) + + assert config.UPLOAD_FILE_SIZE_LIMIT == 23 + assert config.KNOWLEDGE_UPLOAD_FILE_SIZE_LIMIT_FOR_PAID_PLAN == 50 diff --git a/api/tests/unit_tests/configs/test_tidb_on_qdrant_config.py b/api/tests/unit_tests/configs/middleware/vdb/test_tidb_on_qdrant_config.py similarity index 89% rename from api/tests/unit_tests/configs/test_tidb_on_qdrant_config.py rename to api/tests/unit_tests/configs/middleware/vdb/test_tidb_on_qdrant_config.py index 7c2ae5c1f84..e3da65a6c98 100644 --- a/api/tests/unit_tests/configs/test_tidb_on_qdrant_config.py +++ b/api/tests/unit_tests/configs/middleware/vdb/test_tidb_on_qdrant_config.py @@ -1,5 +1,5 @@ from configs.middleware.vdb.tidb_on_qdrant_config import TidbOnQdrantConfig -from tests.unit_tests.configs._isolated_settings import InitSettingsOnly +from tests.helpers.unit.configs._isolated_settings import InitSettingsOnly class _IsolatedTidbOnQdrantConfig(InitSettingsOnly, TidbOnQdrantConfig): diff --git a/api/tests/unit_tests/configs/test_remote_settings_sources.py b/api/tests/unit_tests/configs/remote_settings_sources/apollo/test___init__.py similarity index 100% rename from api/tests/unit_tests/configs/test_remote_settings_sources.py rename to api/tests/unit_tests/configs/remote_settings_sources/apollo/test___init__.py diff --git a/api/tests/unit_tests/configs/test_nacos_http_client.py b/api/tests/unit_tests/configs/remote_settings_sources/nacos/test_http_request.py similarity index 100% rename from api/tests/unit_tests/configs/test_nacos_http_client.py rename to api/tests/unit_tests/configs/remote_settings_sources/nacos/test_http_request.py diff --git a/api/tests/unit_tests/configs/test_dify_config.py b/api/tests/unit_tests/configs/test_app_config.py similarity index 100% rename from api/tests/unit_tests/configs/test_dify_config.py rename to api/tests/unit_tests/configs/test_app_config.py diff --git a/api/tests/unit_tests/configs/test_auth_config.py b/api/tests/unit_tests/configs/test_auth_config.py deleted file mode 100644 index d310161bb4a..00000000000 --- a/api/tests/unit_tests/configs/test_auth_config.py +++ /dev/null @@ -1,24 +0,0 @@ -import pytest - -from configs.feature import AuthConfig -from tests.unit_tests.configs._isolated_settings import InitSettingsOnly - - -class _IsolatedAuthConfig(InitSettingsOnly, AuthConfig): - pass - - -def test_oauth_ttl_days_defaults_to_fourteen() -> None: - config = _IsolatedAuthConfig() - - assert config.OAUTH_TTL_DAYS == 14 - - -@pytest.mark.parametrize( - ("configured", "expected"), - [("30", 30), ("1", 1), ("365", 365), ("0", 1), ("-5", 1), ("500", 365), ("invalid", 14), ("", 14)], -) -def test_oauth_ttl_days_clamps_or_falls_back_instead_of_failing(configured: str, expected: int) -> None: - config = _IsolatedAuthConfig(OAUTH_TTL_DAYS=configured) - - assert expected == config.OAUTH_TTL_DAYS diff --git a/api/tests/unit_tests/configs/test_file_upload_config.py b/api/tests/unit_tests/configs/test_file_upload_config.py deleted file mode 100644 index 64c76cb2fb5..00000000000 --- a/api/tests/unit_tests/configs/test_file_upload_config.py +++ /dev/null @@ -1,23 +0,0 @@ -from configs.feature import FileUploadConfig -from tests.unit_tests.configs._isolated_settings import InitSettingsOnly - - -class _IsolatedFileUploadConfig(InitSettingsOnly, FileUploadConfig): - pass - - -def test_paid_plan_file_size_limit_uses_its_default() -> None: - config = _IsolatedFileUploadConfig(UPLOAD_FILE_SIZE_LIMIT="23") - - assert config.UPLOAD_FILE_SIZE_LIMIT == 23 - assert config.KNOWLEDGE_UPLOAD_FILE_SIZE_LIMIT_FOR_PAID_PLAN == 15 - - -def test_paid_plan_file_size_limit_can_be_configured_separately() -> None: - config = _IsolatedFileUploadConfig( - UPLOAD_FILE_SIZE_LIMIT="23", - KNOWLEDGE_UPLOAD_FILE_SIZE_LIMIT_FOR_PAID_PLAN="50", - ) - - assert config.UPLOAD_FILE_SIZE_LIMIT == 23 - assert config.KNOWLEDGE_UPLOAD_FILE_SIZE_LIMIT_FOR_PAID_PLAN == 50 diff --git a/api/tests/unit_tests/conftest.py b/api/tests/unit_tests/conftest.py index d9a765d5746..3c9b2131d9d 100644 --- a/api/tests/unit_tests/conftest.py +++ b/api/tests/unit_tests/conftest.py @@ -24,12 +24,12 @@ import core.db.session_factory as session_factory_module from extensions import ext_redis from models.account import Account, Tenant, TenantAccountJoin, TenantAccountRole from models.base import TypeBase -from tests.unit_tests.config_override import apply_config_overrides -from tests.unit_tests.rbac_fakes import RBACDomain, build_rbac_domain +from tests.helpers.unit.config_override import apply_config_overrides +from tests.helpers.unit.rbac_fakes import RBACDomain, build_rbac_domain if TYPE_CHECKING: from extensions.ext_application_services import ApplicationServices - from tests.unit_tests.account_domain import AccountDomain + from tests.helpers.unit.account_domain import AccountDomain # Getting the absolute path of the current file's directory ABS_PATH = os.path.dirname(os.path.abspath(__file__)) @@ -39,7 +39,7 @@ PROJECT_DIR = os.path.abspath(os.path.join(ABS_PATH, os.pardir, os.pardir)) CACHED_APP = Flask(__name__) -pytest_plugins = ("tests.unit_tests.audio_runtime_fixtures", "tests.unit_tests.agent_runtime_fixtures") +pytest_plugins = ("tests.helpers.unit.audio_runtime_fixtures", "tests.helpers.unit.agent_runtime_fixtures") # set global mock for Redis client redis_mock = MagicMock() @@ -305,7 +305,7 @@ def account_domain( from enums import DeploymentEdition from services.workspace import gateways - from tests.unit_tests.account_domain import build_account_domain + from tests.helpers.unit.account_domain import build_account_domain config_overrides(DEPLOYMENT_EDITION=DeploymentEdition.COMMUNITY, RBAC_ENABLED=False) monkeypatch.setattr(gateways, "generate_key_pair", lambda _workspace_id: "public-key") diff --git a/api/tests/unit_tests/test_constants.py b/api/tests/unit_tests/constants/test___init__.py similarity index 92% rename from api/tests/unit_tests/test_constants.py rename to api/tests/unit_tests/constants/test___init__.py index d66cdff65db..0bbdc605097 100644 --- a/api/tests/unit_tests/test_constants.py +++ b/api/tests/unit_tests/constants/test___init__.py @@ -4,7 +4,7 @@ import pytest import constants from configs import dify_config -from tests.unit_tests.config_override import apply_config_overrides +from tests.helpers.unit.config_override import apply_config_overrides @pytest.mark.parametrize("etl_type", ["SelfHosted", "Unstructured"]) diff --git a/api/tests/unit_tests/core/workflow/context/test_execution_context.py b/api/tests/unit_tests/context/test_execution_context.py similarity index 100% rename from api/tests/unit_tests/core/workflow/context/test_execution_context.py rename to api/tests/unit_tests/context/test_execution_context.py diff --git a/api/tests/unit_tests/core/workflow/context/test_flask_app_context.py b/api/tests/unit_tests/context/test_flask_app_context.py similarity index 100% rename from api/tests/unit_tests/core/workflow/context/test_flask_app_context.py rename to api/tests/unit_tests/context/test_flask_app_context.py diff --git a/api/tests/unit_tests/controllers/test_rbac_route_resource_contract.py b/api/tests/unit_tests/controllers/common/rbac/test_checks.py similarity index 100% rename from api/tests/unit_tests/controllers/test_rbac_route_resource_contract.py rename to api/tests/unit_tests/controllers/common/rbac/test_checks.py diff --git a/api/tests/unit_tests/controllers/common/test_rbac_locators.py b/api/tests/unit_tests/controllers/common/rbac/test_locators.py similarity index 97% rename from api/tests/unit_tests/controllers/common/test_rbac_locators.py rename to api/tests/unit_tests/controllers/common/rbac/test_locators.py index 84bcdaf2234..99268e64c49 100644 --- a/api/tests/unit_tests/controllers/common/test_rbac_locators.py +++ b/api/tests/unit_tests/controllers/common/rbac/test_locators.py @@ -10,8 +10,8 @@ from controllers.common.rbac import locators from core.rbac import RBACResourceScope from models.agent import Agent, AgentScope, AgentSource, AgentStatus from models.model import AppMode -from tests.unit_tests.model_factories import make_app, make_dataset, make_document -from tests.unit_tests.rbac_fakes import RBACDomain +from tests.helpers.unit.model_factories import make_app, make_dataset, make_document +from tests.helpers.unit.rbac_fakes import RBACDomain @pytest.fixture(autouse=True) diff --git a/api/tests/unit_tests/controllers/common/test_access_response.py b/api/tests/unit_tests/controllers/common/test_access_response.py index 0081f5f5c00..46eef9bb07d 100644 --- a/api/tests/unit_tests/controllers/common/test_access_response.py +++ b/api/tests/unit_tests/controllers/common/test_access_response.py @@ -34,7 +34,7 @@ from libs.external_api import ExternalApi from models.engine import db from models.model import ApiToken, App, AppMCPServer from services.app.mcp_server_service import AppMCPServerStatus -from tests.unit_tests.config_override import apply_config_overrides +from tests.helpers.unit.config_override import apply_config_overrides type TokenClient = tuple[FlaskClient, MagicMock, MagicMock, MagicMock] diff --git a/api/tests/unit_tests/controllers/common/test_app_access_error.py b/api/tests/unit_tests/controllers/common/test_app_access_error.py index f0eb6666b72..a076b0179e0 100644 --- a/api/tests/unit_tests/controllers/common/test_app_access_error.py +++ b/api/tests/unit_tests/controllers/common/test_app_access_error.py @@ -14,7 +14,7 @@ from werkzeug.middleware.proxy_fix import ProxyFix from controllers.common.app_access_error import register_app_access_error_metadata from libs.exception import BaseHTTPException from libs.external_api import ExternalApi -from tests.unit_tests.config_override import apply_config_overrides +from tests.helpers.unit.config_override import apply_config_overrides APP_ID = "9d6e3fb3-94ce-48f9-a958-58ca80b1c02e" TRUSTED = "172.18.0.0/16,10.0.0.0/8,2400:cb00::/32" diff --git a/api/tests/unit_tests/controllers/common/test_schema.py b/api/tests/unit_tests/controllers/common/test_schema.py index 014b396ff05..b73a645c3eb 100644 --- a/api/tests/unit_tests/controllers/common/test_schema.py +++ b/api/tests/unit_tests/controllers/common/test_schema.py @@ -1,15 +1,28 @@ +import importlib.util +import json as _schema_json +import re import sys +from collections.abc import Iterator from datetime import datetime from enum import StrEnum -from types import SimpleNamespace -from typing import Annotated, Literal +from pathlib import Path +from types import ModuleType, SimpleNamespace +from typing import Annotated, Literal, cast from unittest.mock import patch import pytest -from flask import Flask -from flask_restx import Namespace, fields +from flask import Blueprint, Flask, Response, request_finished +from flask import json as _error_response_contract_check_json +from flask_restx import Namespace, Resource, fields +from jsonschema import Draft202012Validator, ValidationError from pydantic import BaseModel, ConfigDict, Field, create_model +from controllers.common.schema import register_response_schema_models +from libs.exception import BaseHTTPException, RelayedHTTPError +from libs.external_api import ExternalApi, error_responses +from tests.helpers.unit.controllers import error_contract +from tests.helpers.unit.controllers.error_contract import check_response, violation + class UserModel(BaseModel): id: int @@ -75,390 +88,878 @@ class ResponseAliasModel(BaseModel): public_name: str = Field(validation_alias="internal_name") -@pytest.fixture(autouse=True) -def console_ns(): - """Isolate the console import while exercising a real namespace.""" - namespace = Namespace("console") - with patch.dict(sys.modules, {"controllers.console": SimpleNamespace(console_ns=namespace)}): - yield namespace - - -def test_default_ref_template_value(): - from controllers.common.schema import DEFAULT_REF_TEMPLATE_OPENAPI_3_0 - - assert DEFAULT_REF_TEMPLATE_OPENAPI_3_0 == "#/components/schemas/{model}" - - -def test_register_schema_model_calls_namespace_schema_model(): - from controllers.common.schema import register_schema_model - - namespace = Namespace("test") - - register_schema_model(namespace, UserModel) - - assert list(namespace.models) == ["UserModel"] - schema = namespace.models["UserModel"].__schema__ - assert isinstance(schema, dict) - assert "properties" in schema - - -def test_register_schema_model_passes_schema_from_pydantic(): - from controllers.common.schema import DEFAULT_REF_TEMPLATE_OPENAPI_3_0, register_schema_model - - namespace = Namespace("test") - - register_schema_model(namespace, UserModel) - - schema = namespace.models["UserModel"].__schema__ - - expected_schema = UserModel.model_json_schema(ref_template=DEFAULT_REF_TEMPLATE_OPENAPI_3_0) - - assert schema == expected_schema - - -def test_register_schema_model_promotes_nested_pydantic_definitions(): - from controllers.common.schema import DEFAULT_REF_TEMPLATE_OPENAPI_3_0, register_schema_model - - namespace = Namespace("test") - - register_schema_model(namespace, ParentModel) - - called_schemas = {name: model.__schema__ for name, model in namespace.models.items()} - parent_schema = ParentModel.model_json_schema(ref_template=DEFAULT_REF_TEMPLATE_OPENAPI_3_0) - - assert set(called_schemas) == {"ParentModel", "ChildModel"} - assert "$defs" not in called_schemas["ParentModel"] - assert called_schemas["ParentModel"]["properties"]["child"]["$ref"] == "#/components/schemas/ChildModel" - assert called_schemas["ChildModel"] == parent_schema["$defs"]["ChildModel"] - - def _make_same_named_model() -> type[BaseModel]: """Build a model whose qualified name differs from UserModel but whose bare name matches.""" return create_model("UserModel", __module__="other_module", email=(str, ...)) -def test_register_schema_model_allows_reregistering_the_same_model(): - from controllers.common.schema import register_schema_model +class TestSchemaCases: + @staticmethod + @pytest.fixture(autouse=True) + def console_ns(): + """Isolate the console import while exercising a real namespace.""" + namespace = Namespace("console") + with patch.dict(sys.modules, {"controllers.console": SimpleNamespace(console_ns=namespace)}): + yield namespace - namespace = Namespace("test") + @staticmethod + def test_default_ref_template_value(): + from controllers.common.schema import DEFAULT_REF_TEMPLATE_OPENAPI_3_0 - register_schema_model(namespace, UserModel) - register_schema_model(namespace, UserModel) + assert DEFAULT_REF_TEMPLATE_OPENAPI_3_0 == "#/components/schemas/{model}" - assert list(namespace.models) == ["UserModel"] + @staticmethod + def test_register_schema_model_calls_namespace_schema_model(): + from controllers.common.schema import register_schema_model + namespace = Namespace("test") -def test_register_schema_model_rejects_different_model_with_same_name(): - from controllers.common.schema import SchemaNameCollisionError, register_schema_model + register_schema_model(namespace, UserModel) - namespace = Namespace("test") - register_schema_model(namespace, UserModel) + assert list(namespace.models) == ["UserModel"] + schema = namespace.models["UserModel"].__schema__ + assert isinstance(schema, dict) + assert "properties" in schema - with pytest.raises(SchemaNameCollisionError, match="'UserModel'.*produced by both"): - register_schema_model(namespace, _make_same_named_model()) + @staticmethod + def test_register_schema_model_passes_schema_from_pydantic(): + from controllers.common.schema import DEFAULT_REF_TEMPLATE_OPENAPI_3_0, register_schema_model + namespace = Namespace("test") -def test_register_schema_model_rejects_nested_definition_with_same_name(): - from controllers.common.schema import SchemaNameCollisionError, register_schema_model + register_schema_model(namespace, UserModel) - other_child = create_model("ChildModel", __module__="other_module", other=(int, ...)) - other_parent = create_model("OtherParentModel", child=(other_child, ...)) - namespace = Namespace("test") - register_schema_model(namespace, ParentModel) + schema = namespace.models["UserModel"].__schema__ - with pytest.raises(SchemaNameCollisionError, match="'ChildModel'"): - register_schema_model(namespace, other_parent) + expected_schema = UserModel.model_json_schema(ref_template=DEFAULT_REF_TEMPLATE_OPENAPI_3_0) + assert schema == expected_schema -def test_register_schema_model_rejects_same_named_types_inside_one_schema(): - from controllers.common.schema import SchemaNameCollisionError, register_schema_model + @staticmethod + def test_register_schema_model_promotes_nested_pydantic_definitions(): + from controllers.common.schema import DEFAULT_REF_TEMPLATE_OPENAPI_3_0, register_schema_model - other_child = create_model("ChildModel", __module__="other_module", other=(int, ...)) - both_children = create_model("BothChildrenModel", first=(ChildModel, ...), second=(other_child, ...)) + namespace = Namespace("test") - with pytest.raises(SchemaNameCollisionError, match="'ChildModel'"): - register_schema_model(Namespace("test"), both_children) + register_schema_model(namespace, ParentModel) + called_schemas = {name: model.__schema__ for name, model in namespace.models.items()} + parent_schema = ParentModel.model_json_schema(ref_template=DEFAULT_REF_TEMPLATE_OPENAPI_3_0) -def test_register_schema_model_rejects_model_registered_in_both_modes_with_different_shapes(): - from controllers.common.schema import ( - SchemaNameCollisionError, - register_response_schema_model, - register_schema_model, - ) + assert set(called_schemas) == {"ParentModel", "ChildModel"} + assert "$defs" not in called_schemas["ParentModel"] + assert called_schemas["ParentModel"]["properties"]["child"]["$ref"] == "#/components/schemas/ChildModel" + assert called_schemas["ChildModel"] == parent_schema["$defs"]["ChildModel"] - namespace = Namespace("test") - register_response_schema_model(namespace, ResponseAliasModel) + @staticmethod + def test_register_schema_model_allows_reregistering_the_same_model(): + from controllers.common.schema import register_schema_model - with pytest.raises(SchemaNameCollisionError, match="two different shapes"): - register_schema_model(namespace, ResponseAliasModel) + namespace = Namespace("test") + register_schema_model(namespace, UserModel) + register_schema_model(namespace, UserModel) -def test_register_schema_model_rejects_names_that_collide_after_normalization(): - from controllers.common.schema import SchemaNameCollisionError, register_schema_model + assert list(namespace.models) == ["UserModel"] - namespace = Namespace("test") - register_schema_model(namespace, create_model("JSONPayload", value=(str, ...))) + @staticmethod + def test_register_schema_model_rejects_different_model_with_same_name(): + from controllers.common.schema import SchemaNameCollisionError, register_schema_model - with pytest.raises(SchemaNameCollisionError, match="once normalized"): - register_schema_model(namespace, create_model("JsonPayload", value=(str, ...))) + namespace = Namespace("test") + register_schema_model(namespace, UserModel) + with pytest.raises(SchemaNameCollisionError, match="'UserModel'.*produced by both"): + register_schema_model(namespace, _make_same_named_model()) -def test_register_schema_model_keeps_namespaces_independent(): - from controllers.common.schema import register_schema_model + @staticmethod + def test_register_schema_model_rejects_nested_definition_with_same_name(): + from controllers.common.schema import SchemaNameCollisionError, register_schema_model - register_schema_model(Namespace("first"), UserModel) - register_schema_model(Namespace("second"), _make_same_named_model()) + other_child = create_model("ChildModel", __module__="other_module", other=(int, ...)) + other_parent = create_model("OtherParentModel", child=(other_child, ...)) + namespace = Namespace("test") + register_schema_model(namespace, ParentModel) + with pytest.raises(SchemaNameCollisionError, match="'ChildModel'"): + register_schema_model(namespace, other_parent) -def test_register_schema_model_applies_schema_name_overrides(monkeypatch: pytest.MonkeyPatch): - from controllers.common import schema as schema_module + @staticmethod + def test_register_schema_model_rejects_same_named_types_inside_one_schema(): + from controllers.common.schema import SchemaNameCollisionError, register_schema_model - monkeypatch.setattr( - schema_module, - "SCHEMA_NAME_OVERRIDES", - {f"{ChildModel.__module__}.{ChildModel.__qualname__}": "RenamedChildModel"}, - ) - namespace = Namespace("test") + other_child = create_model("ChildModel", __module__="other_module", other=(int, ...)) + both_children = create_model("BothChildrenModel", first=(ChildModel, ...), second=(other_child, ...)) - schema_module.register_schema_model(namespace, ParentModel) + with pytest.raises(SchemaNameCollisionError, match="'ChildModel'"): + register_schema_model(Namespace("test"), both_children) - assert set(namespace.models) == {"ParentModel", "RenamedChildModel"} - parent_schema = namespace.models["ParentModel"].__schema__ - assert parent_schema["properties"]["child"]["$ref"] == "#/components/schemas/RenamedChildModel" - - -def test_register_schema_models_registers_multiple_models(): - from controllers.common.schema import register_schema_models - - namespace = Namespace("test") - - register_schema_models(namespace, UserModel, ProductModel) - - called_names = list(namespace.models) - assert called_names == ["UserModel", "ProductModel"] - - -def test_register_schema_models_calls_register_schema_model(monkeypatch: pytest.MonkeyPatch): - from controllers.common.schema import register_schema_models - - namespace = Namespace("test") - - calls = [] - - def fake_register(ns, model): - calls.append((ns, model)) - - monkeypatch.setattr( - "controllers.common.schema.register_schema_model", - fake_register, - ) - - register_schema_models(namespace, UserModel, ProductModel) - - assert calls == [ - (namespace, UserModel), - (namespace, ProductModel), - ] - - -def test_register_response_schema_model_uses_serialized_field_names(): - from controllers.common.schema import register_response_schema_model - - namespace = Namespace("test") - - register_response_schema_model(namespace, ResponseAliasModel) - - assert list(namespace.models) == ["ResponseAliasModel"] - schema = namespace.models["ResponseAliasModel"].__schema__ - assert "public_name" in schema["properties"] - assert "internal_name" not in schema["properties"] - - -def test_register_schema_model_preserves_openapi_nullable_unions(): - from controllers.common.schema import register_schema_model - - namespace = Namespace("test") - - register_schema_model(namespace, NullableSchemaModel) - - called_schemas = {name: model.__schema__ for name, model in namespace.models.items()} - properties = called_schemas["NullableSchemaModel"]["properties"] - - assert properties["name"]["anyOf"] == [{"type": "string"}, {"type": "null"}] - assert properties["tags"]["anyOf"] == [{"items": {"type": "string"}, "type": "array"}, {"type": "null"}] - assert properties["owner"]["anyOf"] == [{"$ref": "#/components/schemas/UserModel"}, {"type": "null"}] - assert "anyOf" in properties["ambiguous"] - - -def test_get_or_create_model_returns_existing_model(console_ns): - from controllers.common.schema import get_or_create_model - - existing_model = console_ns.model("TestModel", {"name": fields.String}) - - result = get_or_create_model("TestModel", {"key": "value"}) - - assert result is existing_model - assert len(console_ns.models) == 1 - - -def test_get_or_create_model_creates_new_model_when_not_exists(console_ns): - from controllers.common.schema import get_or_create_model - - field_def = {"name": fields.String} - - result = get_or_create_model("NewModel", field_def) - - assert result is console_ns.models["NewModel"] - assert result.__schema__["properties"] == {"name": {"type": "string"}} - - -def test_get_or_create_model_does_not_call_model_if_exists(console_ns): - from controllers.common.schema import get_or_create_model - - existing_model = console_ns.model("ExistingModel", {"name": fields.String}) - - result = get_or_create_model("ExistingModel", {"key": "value"}) - - assert result is existing_model - assert len(console_ns.models) == 1 - - -def test_register_enum_models_registers_single_enum(): - from controllers.common.schema import register_enum_models - - namespace = Namespace("test") - - register_enum_models(namespace, StatusEnum) - - assert list(namespace.models) == ["StatusEnum"] - schema = namespace.models["StatusEnum"].__schema__ - assert isinstance(schema, dict) - - -def test_register_enum_models_registers_multiple_enums(): - from controllers.common.schema import register_enum_models - - namespace = Namespace("test") - - register_enum_models(namespace, StatusEnum, PriorityEnum) - - called_names = list(namespace.models) - assert called_names == ["StatusEnum", "PriorityEnum"] - - -def test_register_enum_models_uses_correct_ref_template(): - from controllers.common.schema import register_enum_models - - namespace = Namespace("test") - - register_enum_models(namespace, StatusEnum) - - schema = namespace.models["StatusEnum"].__schema__ - - # Verify the schema contains enum values - assert "enum" in schema or "anyOf" in schema - - -def test_query_params_from_model_builds_flask_restx_doc_params(): - from controllers.common.schema import query_params_from_model - - params = query_params_from_model(QueryModel) - - assert params["page"] == { - "in": "query", - "required": False, - "description": "Page number", - "type": "integer", - "default": 1, - "minimum": 1, - "maximum": 100, - } - assert params["keyword"] == { - "in": "query", - "required": False, - "description": "Search keyword", - "type": "string", - "minLength": 1, - "maxLength": 50, - } - assert params["status"] == { - "in": "query", - "required": False, - "description": "Status filter", - "type": "string", - "enum": ["active", "inactive"], - } - assert params["enum_status"] == { - "in": "query", - "required": False, - "description": "Enum status filter", - "type": "string", - "enum": ["active", "inactive"], - } - assert params["created_at"] == { - "in": "query", - "required": False, - "description": "Creation time", - "type": "string", - "format": "date-time", - } - assert params["appId"] == { - "in": "query", - "required": True, - "description": "Application ID", - "type": "string", - } - assert params["tag_ids"] == { - "in": "query", - "required": False, - "description": "Tag IDs", - "type": "array", - "items": {"type": "string", "minLength": 1, "maxLength": 255}, - "minItems": 1, - "maxItems": 3, - } - assert params["ambiguous"] == { - "in": "query", - "required": False, - "description": "Ambiguous query parameter", - } - - -def test_query_params_from_request_preserves_repeated_list_params(): - from controllers.common.schema import query_params_from_request - - app = Flask(__name__) - with app.test_request_context("/?page=2&limit=30&status=active&status=inactive&keyword=hello"): - query = query_params_from_request(HelperQueryModel, list_fields=("status",)) - - assert query.page == 2 - assert query.limit == 30 - assert query.status == ["active", "inactive"] - assert query.keyword == "hello" - - -def test_query_params_from_request_raises_for_malformed_ints_by_default(): - from controllers.common.schema import query_params_from_request - - app = Flask(__name__) - with app.test_request_context("/?page=bad&limit="): - with pytest.raises(ValueError, match="2 validation errors for HelperQueryModel\\\npage"): - query_params_from_request(HelperQueryModel, list_fields=("status",)) - - -def test_query_params_from_request_can_use_model_default_for_malformed_defaulted_ints(): - from controllers.common.schema import query_params_from_request - - app = Flask(__name__) - with app.test_request_context("/?page=bad&limit="): - query = query_params_from_request( - HelperQueryModel, - list_fields=("status",), - use_defaults_for_malformed_ints=True, + @staticmethod + def test_register_schema_model_rejects_model_registered_in_both_modes_with_different_shapes(): + from controllers.common.schema import ( + SchemaNameCollisionError, + register_response_schema_model, + register_schema_model, ) - assert query.page == 1 - assert query.limit == 20 - assert query.status == [] + namespace = Namespace("test") + register_response_schema_model(namespace, ResponseAliasModel) + + with pytest.raises(SchemaNameCollisionError, match="two different shapes"): + register_schema_model(namespace, ResponseAliasModel) + + @staticmethod + def test_register_schema_model_rejects_names_that_collide_after_normalization(): + from controllers.common.schema import SchemaNameCollisionError, register_schema_model + + namespace = Namespace("test") + register_schema_model(namespace, create_model("JSONPayload", value=(str, ...))) + + with pytest.raises(SchemaNameCollisionError, match="once normalized"): + register_schema_model(namespace, create_model("JsonPayload", value=(str, ...))) + + @staticmethod + def test_register_schema_model_keeps_namespaces_independent(): + from controllers.common.schema import register_schema_model + + register_schema_model(Namespace("first"), UserModel) + register_schema_model(Namespace("second"), _make_same_named_model()) + + @staticmethod + def test_register_schema_model_applies_schema_name_overrides(monkeypatch: pytest.MonkeyPatch): + from controllers.common import schema as schema_module + + monkeypatch.setattr( + schema_module, + "SCHEMA_NAME_OVERRIDES", + {f"{ChildModel.__module__}.{ChildModel.__qualname__}": "RenamedChildModel"}, + ) + namespace = Namespace("test") + + schema_module.register_schema_model(namespace, ParentModel) + + assert set(namespace.models) == {"ParentModel", "RenamedChildModel"} + parent_schema = namespace.models["ParentModel"].__schema__ + assert parent_schema["properties"]["child"]["$ref"] == "#/components/schemas/RenamedChildModel" + + @staticmethod + def test_register_schema_models_registers_multiple_models(): + from controllers.common.schema import register_schema_models + + namespace = Namespace("test") + + register_schema_models(namespace, UserModel, ProductModel) + + called_names = list(namespace.models) + assert called_names == ["UserModel", "ProductModel"] + + @staticmethod + def test_register_schema_models_calls_register_schema_model(monkeypatch: pytest.MonkeyPatch): + from controllers.common.schema import register_schema_models + + namespace = Namespace("test") + + calls = [] + + def fake_register(ns, model): + calls.append((ns, model)) + + monkeypatch.setattr( + "controllers.common.schema.register_schema_model", + fake_register, + ) + + register_schema_models(namespace, UserModel, ProductModel) + + assert calls == [ + (namespace, UserModel), + (namespace, ProductModel), + ] + + @staticmethod + def test_register_response_schema_model_uses_serialized_field_names(): + from controllers.common.schema import register_response_schema_model + + namespace = Namespace("test") + + register_response_schema_model(namespace, ResponseAliasModel) + + assert list(namespace.models) == ["ResponseAliasModel"] + schema = namespace.models["ResponseAliasModel"].__schema__ + assert "public_name" in schema["properties"] + assert "internal_name" not in schema["properties"] + + @staticmethod + def test_register_schema_model_preserves_openapi_nullable_unions(): + from controllers.common.schema import register_schema_model + + namespace = Namespace("test") + + register_schema_model(namespace, NullableSchemaModel) + + called_schemas = {name: model.__schema__ for name, model in namespace.models.items()} + properties = called_schemas["NullableSchemaModel"]["properties"] + + assert properties["name"]["anyOf"] == [{"type": "string"}, {"type": "null"}] + assert properties["tags"]["anyOf"] == [{"items": {"type": "string"}, "type": "array"}, {"type": "null"}] + assert properties["owner"]["anyOf"] == [{"$ref": "#/components/schemas/UserModel"}, {"type": "null"}] + assert "anyOf" in properties["ambiguous"] + + @staticmethod + def test_get_or_create_model_returns_existing_model(console_ns): + from controllers.common.schema import get_or_create_model + + existing_model = console_ns.model("TestModel", {"name": fields.String}) + + result = get_or_create_model("TestModel", {"key": "value"}) + + assert result is existing_model + assert len(console_ns.models) == 1 + + @staticmethod + def test_get_or_create_model_creates_new_model_when_not_exists(console_ns): + from controllers.common.schema import get_or_create_model + + field_def = {"name": fields.String} + + result = get_or_create_model("NewModel", field_def) + + assert result is console_ns.models["NewModel"] + assert result.__schema__["properties"] == {"name": {"type": "string"}} + + @staticmethod + def test_get_or_create_model_does_not_call_model_if_exists(console_ns): + from controllers.common.schema import get_or_create_model + + existing_model = console_ns.model("ExistingModel", {"name": fields.String}) + + result = get_or_create_model("ExistingModel", {"key": "value"}) + + assert result is existing_model + assert len(console_ns.models) == 1 + + @staticmethod + def test_register_enum_models_registers_single_enum(): + from controllers.common.schema import register_enum_models + + namespace = Namespace("test") + + register_enum_models(namespace, StatusEnum) + + assert list(namespace.models) == ["StatusEnum"] + schema = namespace.models["StatusEnum"].__schema__ + assert isinstance(schema, dict) + + @staticmethod + def test_register_enum_models_registers_multiple_enums(): + from controllers.common.schema import register_enum_models + + namespace = Namespace("test") + + register_enum_models(namespace, StatusEnum, PriorityEnum) + + called_names = list(namespace.models) + assert called_names == ["StatusEnum", "PriorityEnum"] + + @staticmethod + def test_register_enum_models_uses_correct_ref_template(): + from controllers.common.schema import register_enum_models + + namespace = Namespace("test") + + register_enum_models(namespace, StatusEnum) + + schema = namespace.models["StatusEnum"].__schema__ + + # Verify the schema contains enum values + assert "enum" in schema or "anyOf" in schema + + @staticmethod + def test_query_params_from_model_builds_flask_restx_doc_params(): + from controllers.common.schema import query_params_from_model + + params = query_params_from_model(QueryModel) + + assert params["page"] == { + "in": "query", + "required": False, + "description": "Page number", + "type": "integer", + "default": 1, + "minimum": 1, + "maximum": 100, + } + assert params["keyword"] == { + "in": "query", + "required": False, + "description": "Search keyword", + "type": "string", + "minLength": 1, + "maxLength": 50, + } + assert params["status"] == { + "in": "query", + "required": False, + "description": "Status filter", + "type": "string", + "enum": ["active", "inactive"], + } + assert params["enum_status"] == { + "in": "query", + "required": False, + "description": "Enum status filter", + "type": "string", + "enum": ["active", "inactive"], + } + assert params["created_at"] == { + "in": "query", + "required": False, + "description": "Creation time", + "type": "string", + "format": "date-time", + } + assert params["appId"] == { + "in": "query", + "required": True, + "description": "Application ID", + "type": "string", + } + assert params["tag_ids"] == { + "in": "query", + "required": False, + "description": "Tag IDs", + "type": "array", + "items": {"type": "string", "minLength": 1, "maxLength": 255}, + "minItems": 1, + "maxItems": 3, + } + assert params["ambiguous"] == { + "in": "query", + "required": False, + "description": "Ambiguous query parameter", + } + + @staticmethod + def test_query_params_from_request_preserves_repeated_list_params(): + from controllers.common.schema import query_params_from_request + + app = Flask(__name__) + with app.test_request_context("/?page=2&limit=30&status=active&status=inactive&keyword=hello"): + query = query_params_from_request(HelperQueryModel, list_fields=("status",)) + + assert query.page == 2 + assert query.limit == 30 + assert query.status == ["active", "inactive"] + assert query.keyword == "hello" + + @staticmethod + def test_query_params_from_request_raises_for_malformed_ints_by_default(): + from controllers.common.schema import query_params_from_request + + app = Flask(__name__) + with app.test_request_context("/?page=bad&limit="): + with pytest.raises(ValueError, match="2 validation errors for HelperQueryModel\\\npage"): + query_params_from_request(HelperQueryModel, list_fields=("status",)) + + @staticmethod + def test_query_params_from_request_can_use_model_default_for_malformed_defaulted_ints(): + from controllers.common.schema import query_params_from_request + + app = Flask(__name__) + with app.test_request_context("/?page=bad&limit="): + query = query_params_from_request( + HelperQueryModel, + list_fields=("status",), + use_defaults_for_malformed_ints=True, + ) + + assert query.page == 1 + assert query.limit == 20 + assert query.status == [] + + +"""Coverage regressions for the enterprise inner API OpenAPI document.""" + + +type Json = dict[str, object] + +HTTP_METHODS = {"delete", "get", "head", "options", "patch", "post", "put", "trace"} +WRITE_METHODS = {"patch", "post", "put"} +ENTERPRISE_PREFIX = "/enterprise/" + +EXPECTED_ENTERPRISE_OPERATION_IDS = { + "create_enterprise_workspace", + "create_enterprise_workspace_ownerless", + "delete_enterprise_workspace_resources", + "enterprise_app_detail", + "enterprise_app_dsl_export", + "enterprise_app_dsl_import", + "enterprise_app_list", + "enterprise_mint_file_grant", + "enterprise_runtime_credentials_resolve", + "join_enterprise_workspace_member", + "send_enterprise_mail", +} + + +def _obj(value: object) -> Json: + return cast(Json, value) + + +def _at(value: object, *keys: str) -> Json: + """Walk nested JSON objects by key; each step is cast because the spec is untyped JSON.""" + node = _obj(value) + for key in keys: + node = _obj(node[key]) + return node + + +def _load_generate_swagger_specs_module() -> ModuleType: + api_dir = Path(__file__).resolve().parents[4] + script_path = api_dir / "dev" / "generate_swagger_specs.py" + + spec = importlib.util.spec_from_file_location("generate_swagger_specs", script_path) + assert spec + assert spec.loader + + module = importlib.util.module_from_spec(spec) + sys.modules[spec.name] = module + spec.loader.exec_module(module) # type: ignore[attr-defined] + return module + + +def _enterprise_operations(spec: Json) -> Iterator[tuple[str, str, Json, Json]]: + for path, path_item in _obj(spec["paths"]).items(): + if not path.startswith(ENTERPRISE_PREFIX): + continue + item = _obj(path_item) + for method, operation in item.items(): + if method in HTTP_METHODS: + yield path, method, item, _obj(operation) + + +def _operation(spec: Json, path: str, method: str) -> Json: + return _at(spec, "paths", path, method) + + +def _schema(spec: Json, name: str) -> Json: + return _at(spec, "components", "schemas", name) + + +def _refs(value: object) -> Iterator[str]: + if isinstance(value, dict): + mapping = _obj(value) + ref = mapping.get("$ref") + if isinstance(ref, str): + yield ref + for child in mapping.values(): + yield from _refs(child) + elif isinstance(value, list): + for child in cast(list[object], value): + yield from _refs(child) + + +def _validator(spec: Json, name: str) -> Draft202012Validator: + return Draft202012Validator({"$ref": f"#/components/schemas/{name}", "components": spec["components"]}) + + +ERROR_RESPONSE_REF = {"$ref": "#/components/schemas/ErrorResponse"} +LEGACY_SIMPLE_MESSAGE_RESPONSES = { + "create_enterprise_workspace": ("404",), + "enterprise_app_dsl_import": ("404",), + "enterprise_runtime_credentials_resolve": ("400", "404"), + "join_enterprise_workspace_member": ("400", "404"), +} +FILE_GRANT_BUSINESS_COMPONENTS = { + "400": {"InvalidRequestError", "GrantTtlTooLongError", "InvalidSubjectError"}, + "404": {"AppNotFoundError", "FileNotFoundError"}, +} + + +def _response_schema(operation: Json, code: str) -> Json: + return _at(operation, "responses", code, "content", "application/json", "schema") + + +def _alternatives(schema: Json) -> list[Json]: + return [_obj(item) for item in cast(list[object], schema["anyOf"])] if "anyOf" in schema else [schema] + + +def _operations_by_id(spec: Json) -> dict[str, Json]: + return {str(operation["operationId"]): operation for _, _, _, operation in _enterprise_operations(spec)} + + +class TestSchemaCasesInnerApi: + @staticmethod + @pytest.fixture(scope="module") + def inner_spec(tmp_path_factory: pytest.TempPathFactory) -> Json: + module = _load_generate_swagger_specs_module() + written_paths = module.generate_specs(tmp_path_factory.mktemp("inner_api_spec")) + inner_path = next(path for path in written_paths if path.name == "inner-openapi.json") + return _obj(_schema_json.loads(inner_path.read_text(encoding="utf-8"))) + + @staticmethod + def test_enterprise_operations_declare_typed_request_bodies(inner_spec: Json) -> None: + operations = list(_enterprise_operations(inner_spec)) + assert operations + + for path, method, _, operation in operations: + if method in WRITE_METHODS: + media = _at(operation, "requestBody", "content", "application/json") + assert "schema" in media, f"{method} {path} request body has no schema" + if method == "get": + assert "requestBody" not in operation, f"GET {path} must not declare a request body" + + @staticmethod + def test_enterprise_success_responses_are_typed(inner_spec: Json) -> None: + for path, method, _, operation in _enterprise_operations(inner_spec): + for code, response in _obj(operation["responses"]).items(): + if not code.startswith("2"): + continue + content = _obj(_obj(response).get("content")) + assert content, f"{method} {path} {code} has no content" + for media_type, media in content.items(): + assert "schema" in _obj(media), f"{method} {path} {code} {media_type} has no schema" + + @staticmethod + def test_enterprise_error_responses_are_typed(inner_spec: Json) -> None: + for path, method, _, operation in _enterprise_operations(inner_spec): + for code, response in _obj(operation["responses"]).items(): + if code.startswith("2"): + continue + content = _obj(_obj(response).get("content")) + assert content, f"{method} {path} {code} is not typed" + json_media = _obj(content.get("application/json")) + assert "schema" in json_media, f"{method} {path} {code} has no JSON schema" + + @staticmethod + def test_enterprise_path_placeholders_have_required_path_parameters(inner_spec: Json) -> None: + for path, method, path_item, operation in _enterprise_operations(inner_spec): + parameters = [ + *cast(list[Json], path_item.get("parameters", [])), + *cast(list[Json], operation.get("parameters", [])), + ] + declared = { + parameter["name"] + for parameter in parameters + if parameter.get("in") == "path" and parameter.get("required") is True + } + for placeholder in re.findall(r"{(\w+)}", path): + assert placeholder in declared, f"{method} {path} missing required path parameter {placeholder}" + + @staticmethod + def test_enterprise_schema_references_resolve(inner_spec: Json) -> None: + schemas = _at(inner_spec, "components", "schemas") + for _, _, _, operation in _enterprise_operations(inner_spec): + for ref in _refs(operation): + assert ref.removeprefix("#/components/schemas/") in schemas, ref + + @staticmethod + def test_enterprise_operation_ids_are_unchanged(inner_spec: Json) -> None: + operation_ids = [operation["operationId"] for _, _, _, operation in _enterprise_operations(inner_spec)] + + assert len(operation_ids) == len(set(operation_ids)) + assert set(operation_ids) == EXPECTED_ENTERPRISE_OPERATION_IDS + + @staticmethod + def test_enterprise_mail_documents_simple_message_response(inner_spec: Json) -> None: + media = _at( + _operation(inner_spec, "/enterprise/mail", "post"), "responses", "200", "content", "application/json" + ) + + assert media["schema"] == {"$ref": "#/components/schemas/SimpleMessageResponse"} + + @staticmethod + def test_ownerless_workspace_custom_config_is_object(inner_spec: Json) -> None: + schema = _schema(inner_spec, "WorkspaceOwnerlessTenantResponse") + custom_config = _at(schema, "properties", "custom_config") + + assert custom_config["type"] == "object" + assert custom_config["additionalProperties"] is True + assert "custom_config" in cast(list[str], schema["required"]) + + @staticmethod + def test_owner_workspace_variant_excludes_custom_config_and_public_key(inner_spec: Json) -> None: + owner = _obj(_schema(inner_spec, "WorkspaceResponse")["properties"]) + + assert {"custom_config", "encrypt_public_key"}.isdisjoint(owner) + + @staticmethod + def test_runtime_credential_values_accept_arbitrary_fields(inner_spec: Json) -> None: + values = _at(_schema(inner_spec, "InnerRuntimeCredentialResponse"), "properties", "values") + + assert values["type"] == "object" + assert values["additionalProperties"] is True + + @staticmethod + def test_import_response_excludes_app_created(inner_spec: Json) -> None: + assert "app_created" not in _obj(_schema(inner_spec, "Import")["properties"]) + + @staticmethod + def test_dsl_export_documents_yaml_data_string(inner_spec: Json) -> None: + assert _at(_schema(inner_spec, "AppExportResponse"), "properties", "data")["type"] == "string" + + @staticmethod + def test_ownerless_create_sample_validates_custom_config_object(inner_spec: Json) -> None: + tenant = { + "id": "t", + "name": "n", + "plan": "p", + "status": "normal", + "encrypt_public_key": None, + "custom_config": {"remove_webapp_brand": True}, + "created_at": None, + "updated_at": None, + } + validator = _validator(inner_spec, "WorkspaceOwnerlessCreateResponse") + + validator.validate({"message": "ok", "tenant": tenant}) + validator.validate({"message": "ok", "tenant": {**tenant, "custom_config": {}}}) + with pytest.raises(ValidationError): + validator.validate({"message": "ok", "tenant": {**tenant, "custom_config": _schema_json.dumps({"a": 1})}}) + + @staticmethod + def test_credentials_sample_accepts_arbitrary_values(inner_spec: Json) -> None: + _validator(inner_spec, "InnerRuntimeCredentialsResolveResponse").validate( + { + "credentials": [ + {"credential_id": "c", "provider": "p", "kind": "model", "values": {"api_key": "k", "x": 1}} + ] + } + ) + + @staticmethod + def test_enterprise_operations_document_default_error_response(inner_spec: Json) -> None: + for path, method, _, operation in _enterprise_operations(inner_spec): + assert _response_schema(operation, "default") == ERROR_RESPONSE_REF, f"{method} {path}" + + @staticmethod + def test_enterprise_documented_errors_include_standard_error_body(inner_spec: Json) -> None: + for path, method, _, operation in _enterprise_operations(inner_spec): + for code in _obj(operation["responses"]): + if code == "default" or not code.startswith(("4", "5")): + continue + assert ERROR_RESPONSE_REF in _alternatives(_response_schema(operation, code)), f"{method} {path} {code}" + + @staticmethod + def test_legacy_message_statuses_document_simple_message(inner_spec: Json) -> None: + operations = _operations_by_id(inner_spec) + for operation_id, codes in LEGACY_SIMPLE_MESSAGE_RESPONSES.items(): + for code in codes: + assert {"$ref": "#/components/schemas/SimpleMessageResponse"} in _alternatives( + _response_schema(operations[operation_id], code) + ), f"{operation_id} {code}" + + @staticmethod + def test_dsl_export_not_found_union_documents_simple_message(inner_spec: Json) -> None: + operation = _operations_by_id(inner_spec)["enterprise_app_dsl_export"] + assert _alternatives(_response_schema(operation, "404"))[0] == { + "$ref": "#/components/schemas/EnterpriseAppDSLExportNotFoundResponse" + } + union = _alternatives(_schema(inner_spec, "EnterpriseAppDSLExportNotFoundResponse")) + assert {"$ref": "#/components/schemas/SimpleMessageResponse"} in union + + @staticmethod + def test_file_grants_document_business_error_components(inner_spec: Json) -> None: + operation = _operations_by_id(inner_spec)["enterprise_mint_file_grant"] + for code, expected in FILE_GRANT_BUSINESS_COMPONENTS.items(): + refs = { + str(item["$ref"]).rsplit("/", 1)[-1] + for item in _alternatives(_response_schema(operation, code)) + if "$ref" in item + } + assert expected <= refs, (code, refs) + for name in expected: + error_body = cast(list[object], _schema(inner_spec, name)["allOf"])[1] + assert _at(error_body, "properties", "status")["const"] == int(code), name + + +class _ItemGoneError(BaseHTTPException): + error_code = "item_gone" + code = 404 + + +class _ItemLockedDetails(BaseModel): + locked_by: str + + +class _ItemLockedError(BaseHTTPException): + error_code = "item_locked" + code = 404 + details_model = _ItemLockedDetails + + +class _LegacyConflict(BaseModel): + error: str + + +def _api() -> tuple[Flask, ExternalApi]: + app = Flask(__name__) + blueprint = Blueprint("contract_check", __name__, url_prefix="/check") + api = ExternalApi(blueprint) + ns = Namespace("items", path="/") + register_response_schema_models(ns, _LegacyConflict) + + @ns.route("/items") + class Items(Resource): + @ns.response(409, "Legacy conflict", ns.models[_LegacyConflict.__name__]) + @error_responses(_ItemGoneError, _ItemLockedError) + def get(self) -> dict[str, str]: + return {} + + @ns.route("/relayed") + class Relayed(Resource): + @error_responses(_ItemGoneError, RelayedHTTPError) + def get(self) -> dict[str, str]: + return {} + + api.add_namespace(ns) + app.register_blueprint(blueprint) + return app, api + + +def _check(body: object, status: int, path: str = "/check/items") -> str | None: + app, api = _api() + with app.test_request_context(path): + return violation( + api, Response(_error_response_contract_check_json.dumps(body), status=status, mimetype="application/json") + ) + + +class TestErrorResponseContractCheckCases: + @staticmethod + def test_rejects_a_body_without_the_standard_fields() -> None: + found = _check({"error": "gone"}, 404) + assert found is not None + assert "lacks" in found + + @staticmethod + def test_rejects_a_code_the_operation_does_not_document_for_the_status() -> None: + found = _check({"code": "surprise", "message": "m", "status": 404}, 404) + assert found is not None + assert "'surprise' is not documented" in found + + @staticmethod + def test_accepts_documented_and_implied_codes() -> None: + assert _check({"code": "item_gone", "message": "m", "status": 404}, 404) is None + assert _check({"code": "not_found", "message": "m", "status": 404}, 404) is None + + @staticmethod + def test_generic_codes_satisfy_the_full_schema_of_a_status_that_documents_business_errors() -> None: + # A generic 404 answers an operation that documents its own 404 codes: the status's complete + # schema (not only ``default``) must admit it, so the open error body stays in the ``anyOf``. + app, api = _api() + with app.test_request_context("/check/items"): + schema = error_contract._document(api)["paths"]["/items"]["get"]["responses"]["404"]["content"][ # type: ignore[index] + "application/json" + ]["schema"] + assert schema["anyOf"][-1] == {"$ref": "#/components/schemas/ErrorResponse"} + assert _check({"code": "not_found", "message": "m", "status": 404}, 404) is None + + @staticmethod + def test_validates_a_documented_code_against_its_component() -> None: + assert ( + _check({"code": "item_locked", "message": "m", "status": 404, "details": {"locked_by": "a"}}, 404) is None + ) + + found = _check({"code": "item_locked", "message": "m", "status": 404, "details": {}}, 404) + assert found is not None + assert "'item_locked' component" in found + assert "locked_by" in found + + @staticmethod + def test_accepts_a_body_matching_a_legacy_model_the_status_declares() -> None: + assert _check({"error": "legacy"}, 409) is None + assert _check({"code": "conflict", "message": "m", "status": 409}, 409) is None + + @staticmethod + def test_rejects_a_documented_code_served_with_another_status(monkeypatch: pytest.MonkeyPatch) -> None: + # On the surfaces the lint covers every served code must be documented or generic. + monkeypatch.setattr(error_contract, "_documents_business_errors", lambda _api: True) + found = _check({"code": "item_gone", "message": "m", "status": 409}, 409) + assert found is not None + assert "'item_gone' is not documented" in found + + @staticmethod + def test_accepts_any_code_from_an_operation_that_relays_errors(monkeypatch: pytest.MonkeyPatch) -> None: + monkeypatch.setattr(error_contract, "_documents_business_errors", lambda _api: True) + assert _check({"code": "downstream_busy", "message": "m", "status": 503}, 503, "/check/relayed") is None + assert _check({"code": "item_gone", "message": "m", "status": 404}, 404, "/check/relayed") is None + # Relaying never excuses a body that is not the error body. + found = _check({"error": "x"}, 503, "/check/relayed") + assert found is not None + assert "lacks" in found + + @staticmethod + def test_only_checks_documented_statuses_on_apis_without_business_errors() -> None: + # A test-built API (like /openapi/v1) documents nothing, so an undocumented status stays open. + assert _check({"code": "item_gone", "message": "m", "status": 409}, 409) is None + + @staticmethod + def test_rejects_a_body_matching_neither_the_legacy_nor_the_standard_model() -> None: + found = _check({"code": "conflict", "message": "m"}, 409) + assert found is not None + assert "lacks ['status']" in found + + @staticmethod + @pytest.mark.parametrize( + ("body", "field"), + [ + ({"code": "item_gone", "message": {"nested": True}, "status": 404}, "message"), + ({"code": ["item_gone"], "message": "m", "status": 404}, "code"), + ({"code": "item_gone", "message": "m", "status": "404"}, "status"), + ], + ) + def test_rejects_standard_fields_of_the_wrong_type(body: dict[str, object], field: str) -> None: + found = _check(body, 404) + assert found is not None + assert "does not match ErrorResponse" in found + assert f"{field}:" in found + + @staticmethod + @pytest.mark.allow_error_contract_violations + def test_checks_responses_served_by_an_api_mounted_on_the_app() -> None: + app = Flask(__name__) + api = ExternalApi(app) + + @api.route("/items") + class Items(Resource): + def get(self) -> tuple[dict[str, str], int]: + return {"wrong": "shape"}, 400 + + violations: list[str] = [] + + def receiver(sender: Flask, response: Response, **_: object) -> None: + check_response(sender, response, violations) + + with request_finished.connected_to(receiver, app): + app.test_client().get("/items") + + assert len(violations) == 1 + assert "lacks" in violations[0] + + @staticmethod + @pytest.mark.allow_error_contract_violations + def test_reports_responses_it_cannot_check() -> None: + app = Flask(__name__) + api = ExternalApi(app) + + class _Unregistered(BaseModel): + value: str + + @api.route("/items") + class Items(Resource): + @api.response(200, "OK", _Unregistered.__name__) + def get(self) -> tuple[dict[str, str], int]: + return {"wrong": "shape"}, 400 + + violations: list[str] = [] + + def receiver(sender: Flask, response: Response, **_: object) -> None: + check_response(sender, response, violations) + + with request_finished.connected_to(receiver, app): + app.test_client().get("/items") + + assert len(violations) == 1 + assert "cannot check the error contract" in violations[0] diff --git a/api/tests/unit_tests/controllers/conftest.py b/api/tests/unit_tests/controllers/conftest.py index cd8c5e8b1dd..d150e267d39 100644 --- a/api/tests/unit_tests/controllers/conftest.py +++ b/api/tests/unit_tests/controllers/conftest.py @@ -29,7 +29,7 @@ from services.tag_application_service import TagApplicationService from services.webapp_access_adapters import EnterpriseWebAppAccessPolicyGateway from services.webapp_access_query_service import WebAppAccessQueryService from services.workflow.dsl_import import DslImportService -from tests.unit_tests.controllers.error_contract import check_response, track_external_apis +from tests.helpers.unit.controllers.error_contract import check_response, track_external_apis track_external_apis() diff --git a/api/tests/unit_tests/controllers/console/agent/test_app_helpers.py b/api/tests/unit_tests/controllers/console/agent/test_app_helpers.py index 0c7eb73889b..e93972f63f6 100644 --- a/api/tests/unit_tests/controllers/console/agent/test_app_helpers.py +++ b/api/tests/unit_tests/controllers/console/agent/test_app_helpers.py @@ -6,7 +6,7 @@ from sqlalchemy.orm import Session from controllers.console.agent import app_helpers from models.model import App, AppMode -from tests.unit_tests.model_factories import make_app +from tests.helpers.unit.model_factories import make_app def _app() -> App: diff --git a/api/tests/unit_tests/controllers/console/agent/test_agent_controllers.py b/api/tests/unit_tests/controllers/console/agent/test_roster.py similarity index 99% rename from api/tests/unit_tests/controllers/console/agent/test_agent_controllers.py rename to api/tests/unit_tests/controllers/console/agent/test_roster.py index 68808e5dd34..fdf02f432e9 100644 --- a/api/tests/unit_tests/controllers/console/agent/test_agent_controllers.py +++ b/api/tests/unit_tests/controllers/console/agent/test_roster.py @@ -78,9 +78,9 @@ from services.entities.agent_entities import ( WorkflowComposerCopyFromRosterPayload, ) from services.entities.app_entities import AgentAppPublicationCounts, AppListParams, CreateAppParams -from tests.unit_tests.config_override import apply_config_overrides -from tests.unit_tests.model_factories import make_account -from tests.unit_tests.rbac_fakes import RBACDomain +from tests.helpers.unit.config_override import apply_config_overrides +from tests.helpers.unit.model_factories import make_account +from tests.helpers.unit.rbac_fakes import RBACDomain @pytest.fixture(autouse=True) diff --git a/api/tests/unit_tests/controllers/console/agent/test_tracing.py b/api/tests/unit_tests/controllers/console/agent/test_tracing.py index ca1c59e3440..cac1ebba869 100644 --- a/api/tests/unit_tests/controllers/console/agent/test_tracing.py +++ b/api/tests/unit_tests/controllers/console/agent/test_tracing.py @@ -1,26 +1,45 @@ -"""Agent tracing Console contracts and authorization stay separate from PlainApp.""" - +import importlib as _test_importlib import inspect +import json from collections.abc import Callable +from copy import deepcopy from dataclasses import dataclass, field -from uuid import UUID +from uuid import UUID, uuid4 import pytest +from dify_trace_arize_phoenix.arize_phoenix_trace import ArizePhoenixDataTrace from flask import Flask +from flask_restx import Api from pydantic import ValidationError +from sqlalchemy import select +from sqlalchemy.orm import Session, sessionmaker from werkzeug.exceptions import Forbidden, NotFound from controllers.common.rbac import AgentId, RBACCheck -from controllers.console import flask_admission +from controllers.console import flask_admission, wraps from controllers.console.agent import tracing +from controllers.console.agent.tracing import AgentTraceConfigApi from controllers.console.app.app import AppTracePayload +from controllers.console.app.ops_trace import TraceAppConfigApi +from core.helper import ssrf_proxy +from core.ops.ops_trace_manager import OpsTraceManager +from enums import DeploymentEdition from extensions.ext_application_services import ApplicationServices from libs.login import AccountWithTenant from machinery.context import RequestContext from models.account import Account, AccountStatus, TenantAccountRole +from models.agent import Agent, AgentScope, AgentSource, AgentStatus +from models.model import App, AppMode, TraceAppConfig +from repositories.app.tracing_config_repository import SQLAlchemyAppTracingConfigRepository from services.app.agent_app_contracts import AgentAppNotFoundError from services.app_tracing_config_service import AppTraceSettings, AppTracingConfigRecord -from tests.unit_tests.config_override import apply_config_overrides +from tests.helpers.unit.config_override import apply_config_overrides +from tests.helpers.unit.model_factories import make_app + +pytest_plugins = ("tests.unit_tests.core.ops.test_ops_trace_manager",) + +"""Agent tracing Console contracts and authorization stay separate from PlainApp.""" + AGENT_ID = UUID("1fbc2f02-f7e0-4910-9032-cbe5bf6c65d5") CONTEXT = RequestContext("request", None, "account", "tenant") @@ -34,239 +53,445 @@ class ConfigCase: failure: Exception | None = None -@pytest.fixture -def service(monkeypatch: pytest.MonkeyPatch, account_application_services: ApplicationServices) -> ConfigCase: - case = ConfigCase(account_application_services) - target = account_application_services.agent_tracing_configs +class TestTracingCases: + @staticmethod + @pytest.fixture + def service(monkeypatch: pytest.MonkeyPatch, account_application_services: ApplicationServices) -> ConfigCase: + case = ConfigCase(account_application_services) + target = account_application_services.agent_tracing_configs - def get_trace(context: RequestContext, agent_id: str) -> AppTraceSettings: - case.calls.append(("get_trace", context, agent_id)) - if case.failure is not None: - raise case.failure - return AppTraceSettings() + def get_trace(context: RequestContext, agent_id: str) -> AppTraceSettings: + case.calls.append(("get_trace", context, agent_id)) + if case.failure is not None: + raise case.failure + return AppTraceSettings() - def set_trace(context: RequestContext, agent_id: str, settings: AppTraceSettings) -> None: - case.calls.append(("set_trace", context, agent_id, settings)) + def set_trace(context: RequestContext, agent_id: str, settings: AppTraceSettings) -> None: + case.calls.append(("set_trace", context, agent_id, settings)) - def get(context: RequestContext, agent_id: str, provider: str) -> AppTracingConfigRecord | None: - case.calls.append(("get", context, agent_id, provider)) - return case.record + def get(context: RequestContext, agent_id: str, provider: str) -> AppTracingConfigRecord | None: + case.calls.append(("get", context, agent_id, provider)) + return case.record - def delete(context: RequestContext, agent_id: str, provider: str) -> None: - case.calls.append(("delete", context, agent_id, provider)) + def delete(context: RequestContext, agent_id: str, provider: str) -> None: + case.calls.append(("delete", context, agent_id, provider)) - monkeypatch.setattr(target, "get_trace", get_trace) - monkeypatch.setattr(target, "set_trace", set_trace) - monkeypatch.setattr(target, "get", get) - monkeypatch.setattr(target, "delete", delete) - monkeypatch.setattr(tracing, "application_services", lambda: account_application_services) - return case + monkeypatch.setattr(target, "get_trace", get_trace) + monkeypatch.setattr(target, "set_trace", set_trace) + monkeypatch.setattr(target, "get", get) + monkeypatch.setattr(target, "delete", delete) + monkeypatch.setattr(tracing, "application_services", lambda: account_application_services) + return case - -def test_status_round_trips_content_policy(app: Flask, service: ConfigCase) -> None: - with app.test_request_context( - json={"enabled": True, "tracing_provider": "langfuse", "capture_content": True, "capture_agent": False} - ): - result = inspect.unwrap(tracing.AgentTraceApi.post)(tracing.AgentTraceApi(), CONTEXT, AGENT_ID) - assert result == {"enabled": True, "tracing_provider": "langfuse", "capture_content": True, "capture_agent": False} - assert service.calls == [("set_trace", CONTEXT, str(AGENT_ID), AppTraceSettings(True, "langfuse", True, False))] - - -def test_enable_requires_provider_even_when_field_omitted() -> None: - with pytest.raises(ValidationError, match="tracing_provider is required"): - AppTracePayload.model_validate({"enabled": True}) - payload = AppTracePayload.model_validate({"enabled": False}) - assert payload.capture_content is False - assert payload.capture_agent is True - - -@pytest.mark.parametrize("enabled", [False, True]) -def test_capabilities_reports_current_api_gate_without_claiming_runtime_delivery( - app: Flask, service: ConfigCase, monkeypatch: pytest.MonkeyPatch, enabled: bool -) -> None: - apply_config_overrides(monkeypatch, OPS_AGENT_TRACE_ENABLED=enabled) - with app.test_request_context("/?tracing_provider=langfuse"): - result = inspect.unwrap(tracing.AgentTraceCapabilitiesApi.get)( - tracing.AgentTraceCapabilitiesApi(), CONTEXT, AGENT_ID - ) - assert service.calls == [("get_trace", CONTEXT, str(AGENT_ID))] - assert result["api_collection_enabled"] is enabled - assert result["agent_spans"] is True - assert result["export_ack"] == "accepted" - assert "runtime_enabled" not in result - assert "queryable" not in result - - -@pytest.mark.parametrize("enabled", [False, True]) -def test_app_capabilities_reports_api_gate( - app: Flask, monkeypatch: pytest.MonkeyPatch, enabled: bool, service: ConfigCase -) -> None: - calls: list[tuple[RequestContext, str]] = [] - - def get_trace(context: RequestContext, app_id: str) -> AppTraceSettings: - calls.append((context, app_id)) - return AppTraceSettings() - - monkeypatch.setattr(service.services.apps.console, "get_trace", get_trace) - apply_config_overrides(monkeypatch, OPS_AGENT_TRACE_ENABLED=enabled) - with app.test_request_context("/?tracing_provider=mlflow"): - result = inspect.unwrap(tracing.AppTraceCapabilitiesApi.get)( - tracing.AppTraceCapabilitiesApi(), CONTEXT, AGENT_ID - ) - assert result["api_collection_enabled"] is enabled - assert result["provider"] == "mlflow" - assert calls == [(CONTEXT, str(AGENT_ID))] - - -def test_get_config_absent_and_delete_has_no_body(app: Flask, service: ConfigCase) -> None: - with app.test_request_context("/?tracing_provider=langfuse"): - assert inspect.unwrap(tracing.AgentTraceConfigApi.get)(tracing.AgentTraceConfigApi(), CONTEXT, AGENT_ID) == { - "has_not_configured": True + @staticmethod + def test_status_round_trips_content_policy(app: Flask, service: ConfigCase) -> None: + with app.test_request_context( + json={"enabled": True, "tracing_provider": "langfuse", "capture_content": True, "capture_agent": False} + ): + result = inspect.unwrap(tracing.AgentTraceApi.post)(tracing.AgentTraceApi(), CONTEXT, AGENT_ID) + assert result == { + "enabled": True, + "tracing_provider": "langfuse", + "capture_content": True, + "capture_agent": False, } - assert inspect.unwrap(tracing.AgentTraceConfigApi.delete)(tracing.AgentTraceConfigApi(), CONTEXT, AGENT_ID) == ( - "", - 204, - ) - assert service.calls == [ - ("get", CONTEXT, str(AGENT_ID), "langfuse"), - ("delete", CONTEXT, str(AGENT_ID), "langfuse"), - ] + assert service.calls == [("set_trace", CONTEXT, str(AGENT_ID), AppTraceSettings(True, "langfuse", True, False))] + @staticmethod + def test_enable_requires_provider_even_when_field_omitted() -> None: + with pytest.raises(ValidationError, match="tracing_provider is required"): + AppTracePayload.model_validate({"enabled": True}) + payload = AppTracePayload.model_validate({"enabled": False}) + assert payload.capture_content is False + assert payload.capture_agent is True -def test_invalid_agent_returns_not_found(app: Flask, service: ConfigCase) -> None: - service.failure = AgentAppNotFoundError() - with app.test_request_context(), pytest.raises(NotFound): - inspect.unwrap(tracing.AgentTraceApi.get)(tracing.AgentTraceApi(), CONTEXT, AGENT_ID) - - -@pytest.mark.parametrize( - "method", - [ - tracing.AgentTraceApi.post, - tracing.AgentTraceConfigApi.get, - tracing.AgentTraceConfigApi.post, - tracing.AgentTraceConfigApi.patch, - tracing.AgentTraceConfigApi.delete, - ], -) -@pytest.mark.parametrize("rbac_enabled", [False, True]) -def test_config_requires_editor_role_or_agent_edit_permission( - app: Flask, monkeypatch: pytest.MonkeyPatch, method: Callable[..., object], rbac_enabled: bool -) -> None: - account = Account(name="Reader", email="reader@example.com", status=AccountStatus.ACTIVE) - account.id, account.role = "account", TenantAccountRole.NORMAL - apply_config_overrides(monkeypatch, RBAC_ENABLED=rbac_enabled) - monkeypatch.setattr( - flask_admission, "current_account_with_tenant", lambda: AccountWithTenant(account=account, tenant_id="tenant") - ) - denied_calls: list[dict[str, object]] = [] - - def denied(**kwargs: object) -> None: - denied_calls.append(kwargs) - raise Forbidden() - - monkeypatch.setattr(flask_admission, "enforce_rbac_checks", denied) - admission = inspect.unwrap( - method, stop=lambda candidate: "allowed_roles" in inspect.getclosurevars(candidate).nonlocals - ) - with app.test_request_context(), pytest.raises(Forbidden): - admission(None, agent_id=AGENT_ID) - if rbac_enabled: - checks = denied_calls[0]["checks"] - assert isinstance(checks, (list, tuple)) - (check,) = checks - assert isinstance(check, RBACCheck) - assert check.scene == tracing.RBACPermission.AGENT_EDIT - assert isinstance(check.locator, AgentId) - assert denied_calls[0]["path_args"] == {"agent_id": AGENT_ID} - else: - assert denied_calls == [] - - -def test_read_settings_and_masked_config_keep_agent_identity(app: Flask, service: ConfigCase) -> None: - from datetime import UTC, datetime - - now = datetime(2026, 1, 1, tzinfo=UTC) - service.record = AppTracingConfigRecord( - id="config", - app_id="backing-app", - tracing_provider="langfuse", - tracing_config={"secret_key": "[__HIDDEN__]", "project_url": "https://example.test/project"}, - is_active=True, - created_at=now, - updated_at=now, - ) - with app.test_request_context("/?tracing_provider=langfuse"): - settings = inspect.unwrap(tracing.AgentTraceApi.get)(tracing.AgentTraceApi(), CONTEXT, AGENT_ID) - config = inspect.unwrap(tracing.AgentTraceConfigApi.get)(tracing.AgentTraceConfigApi(), CONTEXT, AGENT_ID) - assert settings == {"enabled": False, "tracing_provider": None, "capture_content": False, "capture_agent": True} - assert config["tracing_config"] == service.record.tracing_config - assert config["app_id"] == "backing-app" - assert service.calls == [("get_trace", CONTEXT, str(AGENT_ID)), ("get", CONTEXT, str(AGENT_ID), "langfuse")] - - -@pytest.mark.parametrize("create", [True, False]) -@pytest.mark.parametrize("activate", [False, True]) -def test_config_write_forwards_credentials_but_response_does_not_echo_them( - app: Flask, service: ConfigCase, monkeypatch: pytest.MonkeyPatch, create: bool, activate: bool -) -> None: - calls: list[tuple[RequestContext, str, str, dict[str, object], bool]] = [] - - def write( - context: RequestContext, agent_id: str, provider: str, config: dict[str, object], *, activate: bool + @staticmethod + @pytest.mark.parametrize("enabled", [False, True]) + def test_capabilities_reports_current_api_gate_without_claiming_runtime_delivery( + app: Flask, service: ConfigCase, monkeypatch: pytest.MonkeyPatch, enabled: bool ) -> None: - calls.append((context, agent_id, provider, config, activate)) + apply_config_overrides(monkeypatch, OPS_AGENT_TRACE_ENABLED=enabled) + with app.test_request_context("/?tracing_provider=langfuse"): + result = inspect.unwrap(tracing.AgentTraceCapabilitiesApi.get)( + tracing.AgentTraceCapabilitiesApi(), CONTEXT, AGENT_ID + ) + assert service.calls == [("get_trace", CONTEXT, str(AGENT_ID))] + assert result["api_collection_enabled"] is enabled + assert result["agent_spans"] is True + assert result["export_ack"] == "accepted" + assert "runtime_enabled" not in result + assert "queryable" not in result - target = service.services.agent_tracing_configs - monkeypatch.setattr(target, "create" if create else "update", write) - credentials = {"secret_key": "never-echo-this-secret"} - with app.test_request_context( - method="POST" if create else "PATCH", - json={"tracing_provider": "langfuse", "tracing_config": credentials, "activate": activate}, - ): - api = tracing.AgentTraceConfigApi() - result = inspect.unwrap(api.post if create else api.patch)(api, CONTEXT, AGENT_ID) - assert calls == [(CONTEXT, str(AGENT_ID), "langfuse", credentials, activate)] - assert result == (({"result": "success"}, 201) if create else {"result": "success"}) - assert "never-echo-this-secret" not in str(result) + @staticmethod + @pytest.mark.parametrize("enabled", [False, True]) + def test_app_capabilities_reports_api_gate( + app: Flask, monkeypatch: pytest.MonkeyPatch, enabled: bool, service: ConfigCase + ) -> None: + calls: list[tuple[RequestContext, str]] = [] + + def get_trace(context: RequestContext, app_id: str) -> AppTraceSettings: + calls.append((context, app_id)) + return AppTraceSettings() + + monkeypatch.setattr(service.services.apps.console, "get_trace", get_trace) + apply_config_overrides(monkeypatch, OPS_AGENT_TRACE_ENABLED=enabled) + with app.test_request_context("/?tracing_provider=mlflow"): + result = inspect.unwrap(tracing.AppTraceCapabilitiesApi.get)( + tracing.AppTraceCapabilitiesApi(), CONTEXT, AGENT_ID + ) + assert result["api_collection_enabled"] is enabled + assert result["provider"] == "mlflow" + assert calls == [(CONTEXT, str(AGENT_ID))] + + @staticmethod + def test_get_config_absent_and_delete_has_no_body(app: Flask, service: ConfigCase) -> None: + with app.test_request_context("/?tracing_provider=langfuse"): + assert inspect.unwrap(tracing.AgentTraceConfigApi.get)( + tracing.AgentTraceConfigApi(), CONTEXT, AGENT_ID + ) == {"has_not_configured": True} + assert inspect.unwrap(tracing.AgentTraceConfigApi.delete)( + tracing.AgentTraceConfigApi(), CONTEXT, AGENT_ID + ) == ( + "", + 204, + ) + assert service.calls == [ + ("get", CONTEXT, str(AGENT_ID), "langfuse"), + ("delete", CONTEXT, str(AGENT_ID), "langfuse"), + ] + + @staticmethod + def test_invalid_agent_returns_not_found(app: Flask, service: ConfigCase) -> None: + service.failure = AgentAppNotFoundError() + with app.test_request_context(), pytest.raises(NotFound): + inspect.unwrap(tracing.AgentTraceApi.get)(tracing.AgentTraceApi(), CONTEXT, AGENT_ID) + + @staticmethod + @pytest.mark.parametrize( + "method", + [ + tracing.AgentTraceApi.post, + tracing.AgentTraceConfigApi.get, + tracing.AgentTraceConfigApi.post, + tracing.AgentTraceConfigApi.patch, + tracing.AgentTraceConfigApi.delete, + ], + ) + @pytest.mark.parametrize("rbac_enabled", [False, True]) + def test_config_requires_editor_role_or_agent_edit_permission( + app: Flask, monkeypatch: pytest.MonkeyPatch, method: Callable[..., object], rbac_enabled: bool + ) -> None: + account = Account(name="Reader", email="reader@example.com", status=AccountStatus.ACTIVE) + account.id, account.role = "account", TenantAccountRole.NORMAL + apply_config_overrides(monkeypatch, RBAC_ENABLED=rbac_enabled) + monkeypatch.setattr( + flask_admission, + "current_account_with_tenant", + lambda: AccountWithTenant(account=account, tenant_id="tenant"), + ) + denied_calls: list[dict[str, object]] = [] + + def denied(**kwargs: object) -> None: + denied_calls.append(kwargs) + raise Forbidden() + + monkeypatch.setattr(flask_admission, "enforce_rbac_checks", denied) + admission = inspect.unwrap( + method, stop=lambda candidate: "allowed_roles" in inspect.getclosurevars(candidate).nonlocals + ) + with app.test_request_context(), pytest.raises(Forbidden): + admission(None, agent_id=AGENT_ID) + if rbac_enabled: + checks = denied_calls[0]["checks"] + assert isinstance(checks, (list, tuple)) + (check,) = checks + assert isinstance(check, RBACCheck) + assert check.scene == tracing.RBACPermission.AGENT_EDIT + assert isinstance(check.locator, AgentId) + assert denied_calls[0]["path_args"] == {"agent_id": AGENT_ID} + else: + assert denied_calls == [] + + @staticmethod + def test_read_settings_and_masked_config_keep_agent_identity(app: Flask, service: ConfigCase) -> None: + from datetime import UTC, datetime + + now = datetime(2026, 1, 1, tzinfo=UTC) + service.record = AppTracingConfigRecord( + id="config", + app_id="backing-app", + tracing_provider="langfuse", + tracing_config={"secret_key": "[__HIDDEN__]", "project_url": "https://example.test/project"}, + is_active=True, + created_at=now, + updated_at=now, + ) + with app.test_request_context("/?tracing_provider=langfuse"): + settings = inspect.unwrap(tracing.AgentTraceApi.get)(tracing.AgentTraceApi(), CONTEXT, AGENT_ID) + config = inspect.unwrap(tracing.AgentTraceConfigApi.get)(tracing.AgentTraceConfigApi(), CONTEXT, AGENT_ID) + assert settings == {"enabled": False, "tracing_provider": None, "capture_content": False, "capture_agent": True} + assert config["tracing_config"] == service.record.tracing_config + assert config["app_id"] == "backing-app" + assert service.calls == [("get_trace", CONTEXT, str(AGENT_ID)), ("get", CONTEXT, str(AGENT_ID), "langfuse")] + + @staticmethod + @pytest.mark.parametrize("create", [True, False]) + @pytest.mark.parametrize("activate", [False, True]) + def test_config_write_forwards_credentials_but_response_does_not_echo_them( + app: Flask, service: ConfigCase, monkeypatch: pytest.MonkeyPatch, create: bool, activate: bool + ) -> None: + calls: list[tuple[RequestContext, str, str, dict[str, object], bool]] = [] + + def write( + context: RequestContext, agent_id: str, provider: str, config: dict[str, object], *, activate: bool + ) -> None: + calls.append((context, agent_id, provider, config, activate)) + + target = service.services.agent_tracing_configs + monkeypatch.setattr(target, "create" if create else "update", write) + credentials = {"secret_key": "never-echo-this-secret"} + with app.test_request_context( + method="POST" if create else "PATCH", + json={"tracing_provider": "langfuse", "tracing_config": credentials, "activate": activate}, + ): + api = tracing.AgentTraceConfigApi() + result = inspect.unwrap(api.post if create else api.patch)(api, CONTEXT, AGENT_ID) + assert calls == [(CONTEXT, str(AGENT_ID), "langfuse", credentials, activate)] + assert result == (({"result": "success"}, 201) if create else {"result": "success"}) + assert "never-echo-this-secret" not in str(result) + + @staticmethod + @pytest.mark.parametrize( + ("error", "expected"), + [ + (tracing.ConsoleAppNotFoundError(), NotFound), + (tracing.AppTracingConfigAppNotFoundError(), tracing.AppNotFoundError), + (tracing.AppTracingConfigAlreadyExistsError(), tracing.TracingConfigAlreadyExistsError), + (tracing.AppTracingConfigNotFoundError(), tracing.TracingConfigNotFoundError), + (tracing.AppTracingConfigConflictError(), tracing.TracingConfigConflictError), + ( + tracing.AppTracingConfigInvalidProviderError("private-provider-secret"), + tracing.UnsupportedTracingProviderError, + ), + (tracing.AppTracingConfigProviderUnavailableError(), tracing.TracingProviderUnavailableError), + (tracing.AppTracingConfigInvalidConfigurationError(), tracing.InvalidTracingConfigError), + (tracing.AppTracingConfigVerificationFailedError(), tracing.TracingConfigVerificationFailedError), + (tracing.AppTracingConfigProcessingError(), tracing.TracingConfigProcessingError), + ], + ) + def test_config_domain_errors_use_public_contract_without_provider_details( + app: Flask, service: ConfigCase, error: Exception, expected: type[Exception] + ) -> None: + service.failure = error + with app.test_request_context(), pytest.raises(expected) as caught: + inspect.unwrap(tracing.AgentTraceApi.get)(tracing.AgentTraceApi(), CONTEXT, AGENT_ID) + assert "private-provider-secret" not in str(caught.value) + + @staticmethod + def test_unknown_capability_and_missing_app_return_public_errors( + app: Flask, service: ConfigCase, monkeypatch: pytest.MonkeyPatch + ) -> None: + with ( + app.test_request_context("/?tracing_provider=unknown"), + pytest.raises(tracing.UnsupportedTracingProviderError), + ): + inspect.unwrap(tracing.AgentTraceCapabilitiesApi.get)( + tracing.AgentTraceCapabilitiesApi(), CONTEXT, AGENT_ID + ) + + def missing(_context: RequestContext, _app_id: str) -> AppTraceSettings: + raise tracing.ConsoleAppNotFoundError() + + monkeypatch.setattr(service.services.apps.console, "get_trace", missing) + with app.test_request_context("/?tracing_provider=langfuse"), pytest.raises(NotFound): + inspect.unwrap(tracing.AppTraceCapabilitiesApi.get)(tracing.AppTraceCapabilitiesApi(), CONTEXT, AGENT_ID) -@pytest.mark.parametrize( - ("error", "expected"), - [ - (tracing.ConsoleAppNotFoundError(), NotFound), - (tracing.AppTracingConfigAppNotFoundError(), tracing.AppNotFoundError), - (tracing.AppTracingConfigAlreadyExistsError(), tracing.TracingConfigAlreadyExistsError), - (tracing.AppTracingConfigNotFoundError(), tracing.TracingConfigNotFoundError), - (tracing.AppTracingConfigConflictError(), tracing.TracingConfigConflictError), - ( - tracing.AppTracingConfigInvalidProviderError("private-provider-secret"), - tracing.UnsupportedTracingProviderError, - ), - (tracing.AppTracingConfigProviderUnavailableError(), tracing.TracingProviderUnavailableError), - (tracing.AppTracingConfigInvalidConfigurationError(), tracing.InvalidTracingConfigError), - (tracing.AppTracingConfigVerificationFailedError(), tracing.TracingConfigVerificationFailedError), - (tracing.AppTracingConfigProcessingError(), tracing.TracingConfigProcessingError), - ], -) -def test_config_domain_errors_use_public_contract_without_provider_details( - app: Flask, service: ConfigCase, error: Exception, expected: type[Exception] -) -> None: - service.failure = error - with app.test_request_context(), pytest.raises(expected) as caught: - inspect.unwrap(tracing.AgentTraceApi.get)(tracing.AgentTraceApi(), CONTEXT, AGENT_ID) - assert "private-provider-secret" not in str(caught.value) +"""Both real config Resources reach the same ACK gate before atomic SQL updates. + +HTTP dispatch, payload validation, owner resolution, DI services, encryption and +SQLite persistence are real. Login admission supplies a local owner account and +the existing RSA fixture supplies private-key storage; no external account is used. +""" -def test_unknown_capability_and_missing_app_return_public_errors( - app: Flask, service: ConfigCase, monkeypatch: pytest.MonkeyPatch -) -> None: - with app.test_request_context("/?tracing_provider=unknown"), pytest.raises(tracing.UnsupportedTracingProviderError): - inspect.unwrap(tracing.AgentTraceCapabilitiesApi.get)(tracing.AgentTraceCapabilitiesApi(), CONTEXT, AGENT_ID) +_324fde08_TestArizeConfigurationAckCases = _test_importlib.import_module( + "tests.unit_tests.providers.trace.trace-arize-phoenix.src.dify_trace_arize_phoenix.test_arize_configuration" +).TestArizeConfigurationAckCases - def missing(_context: RequestContext, _app_id: str) -> AppTraceSettings: - raise tracing.ConsoleAppNotFoundError() - monkeypatch.setattr(service.services.apps.console, "get_trace", missing) - with app.test_request_context("/?tracing_provider=langfuse"), pytest.raises(NotFound): - inspect.unwrap(tracing.AppTraceCapabilitiesApi.get)(tracing.AppTraceCapabilitiesApi(), CONTEXT, AGENT_ID) +class TestArizeTracingAckRoutesCases: + collector = staticmethod(_324fde08_TestArizeConfigurationAckCases.collector) + + @staticmethod + @pytest.mark.parametrize("surface", ["app", "agent"]) + @pytest.mark.parametrize("method", ["POST", "PATCH"]) + def test_config_write_rejects_bad_ack_without_changing_old_rows_policy_or_cache( + workspace: str, + sqlite_session_factory: sessionmaker[Session], + account_application_services: ApplicationServices, + monkeypatch: pytest.MonkeyPatch, + collector: tuple[str, list[tuple[bytes, dict[str, str]]], dict[str, object]], + surface: str, + method: str, + ) -> None: + endpoint, calls, _ = collector + app_id, agent_id, account_id = (str(uuid4()) for _ in range(3)) + old_policy = json.dumps( + {"enabled": True, "tracing_provider": "langsmith", "capture_content": False, "capture_agent": False} + ) + with sqlite_session_factory.begin() as session: + session.add( + make_app( + app_id=app_id, + tenant_id=workspace, + mode=AppMode.AGENT if surface == "agent" else AppMode.WORKFLOW, + tracing=old_policy, + ) + ) + if surface == "agent": + session.add( + Agent( + id=agent_id, + tenant_id=workspace, + name="ACK test", + scope=AgentScope.ROSTER, + source=AgentSource.AGENT_APP, + app_id=app_id, + status=AgentStatus.ACTIVE, + created_by=account_id, + updated_by=account_id, + ) + ) + store = SQLAlchemyAppTracingConfigRepository(session_factory=sqlite_session_factory) + original = OpsTraceManager.encrypt_tracing_config( + workspace, "langsmith", {"api_key": "other-provider-key", "project": "other"} + ) + assert store.create( + workspace_id=workspace, app_id=app_id, tracing_provider="langsmith", tracing_config=original + ) + if method == "PATCH": + old = OpsTraceManager.encrypt_tracing_config( + workspace, + "arize", + { + "endpoint": endpoint, + "project": "old-project", + "api_key": "synthetic-good", + "space_id": "synthetic-space", + }, + ) + old.update( + { + "project_url": "https://app.arize.com/?redirect_project_name=old-project", + "arize_project_destination": {"legacy": "unchanged"}, + } + ) + assert store.create(workspace_id=workspace, app_id=app_id, tracing_provider="arize", tracing_config=old) + with sqlite_session_factory() as session: + before = deepcopy( + list( + session.execute( + select(TraceAppConfig.__table__) + .where(TraceAppConfig.app_id == app_id) + .order_by(TraceAppConfig.id) + ).mappings() + ) + ) + + account = Account(name="Owner", email="owner@example.test", status=AccountStatus.ACTIVE) + account.id, account.role = account_id, TenantAccountRole.OWNER + identity = AccountWithTenant(account, workspace) + monkeypatch.setattr(flask_admission, "current_account_with_tenant", lambda: identity) + monkeypatch.setattr(wraps, "current_account_with_tenant", lambda: identity) + apply_config_overrides( + monkeypatch, DEPLOYMENT_EDITION=DeploymentEdition.CLOUD, LOGIN_DISABLED=True, RBAC_ENABLED=False + ) + + def forbidden(*_args: object, **_kwargs: object) -> None: + pytest.fail("Local custom collector verification must not perform optional Cloud HTTP reads") + + monkeypatch.setattr(ssrf_proxy, "get", forbidden) + monkeypatch.setattr(ArizePhoenixDataTrace, "__init__", forbidden) + app = Flask(__name__) + app.extensions["application_services"] = account_application_services + api = Api(app) + api.add_resource(TraceAppConfigApi, "/console/api/apps//trace-config") + api.add_resource(AgentTraceConfigApi, "/console/api/agent//trace-config") + url = ( + f"/console/api/apps/{app_id}/trace-config" + if surface == "app" + else f"/console/api/agent/{agent_id}/trace-config" + ) + values = { + "endpoint": endpoint, + "project": "not-created-by-verification", + "api_key": "synthetic-invalid", + "space_id": "synthetic-space", + } + with app.test_client() as client: + failed = client.open( + url, method=method, json={"tracing_provider": "arize", "tracing_config": values, "activate": True} + ) + assert failed.status_code == 400 + body = failed.json + assert isinstance(body, dict) + assert body["code"] == "tracing_config_verification_failed" + assert "private-vendor-details" not in failed.get_data(as_text=True) + assert "synthetic-invalid" not in failed.get_data(as_text=True) + with sqlite_session_factory() as session: + after = list( + session.execute( + select(TraceAppConfig.__table__) + .where(TraceAppConfig.app_id == app_id) + .order_by(TraceAppConfig.id) + ).mappings() + ) + assert after == before + assert session.scalar(select(App.tracing).where(App.id == app_id)) == old_policy + values["api_key"] = "synthetic-good" + accepted = client.open( + url, method=method, json={"tracing_provider": "arize", "tracing_config": values, "activate": True} + ) + assert accepted.status_code == (201 if method == "POST" else 200) + saved = store.get(workspace_id=workspace, app_id=app_id, tracing_provider="arize") + assert saved is not None + assert saved.tracing_config is not None + assert ( + OpsTraceManager.decrypt_tracing_config(workspace, "arize", saved.tracing_config)["api_key"] + == "synthetic-good" + ) + assert "project_url" not in saved.tracing_config + assert "arize_project_destination" not in saved.tracing_config + assert len(calls) == 2 + with sqlite_session_factory() as session: + policy = json.loads(session.scalar(select(App.tracing).where(App.id == app_id)) or "{}") + assert policy == { + "enabled": True, + "tracing_provider": "arize", + "capture_content": False, + "capture_agent": False, + } + if method == "PATCH": + with app.test_client() as client: + masked = client.patch( + url, + json={ + "tracing_provider": "arize", + "tracing_config": { + **values, + "project": "next-project", + "api_key": "******", + "space_id": "******", + }, + }, + ) + assert masked.status_code == 200 + updated = store.get(workspace_id=workspace, app_id=app_id, tracing_provider="arize") + assert updated is not None + assert updated.tracing_config is not None + assert updated.tracing_config["api_key"] == saved.tracing_config["api_key"] + assert updated.tracing_config["space_id"] == saved.tracing_config["space_id"] + assert calls[-1][1]["api_key"] == "synthetic-good" + assert calls[-1][1]["space_id"] == "synthetic-space" + assert len(calls) == 3 diff --git a/api/tests/unit_tests/controllers/console/agent/test_tracing_runs.py b/api/tests/unit_tests/controllers/console/agent/test_tracing_runs.py index 7fb163b5bf0..c817de50d46 100644 --- a/api/tests/unit_tests/controllers/console/agent/test_tracing_runs.py +++ b/api/tests/unit_tests/controllers/console/agent/test_tracing_runs.py @@ -25,7 +25,7 @@ from models.account import Account, AccountStatus, TenantAccountRole from repositories.application_trace_repository import ApplicationTraceRepository, TraceCollectionError from services.app.agent_app_contracts import AgentAppNotFoundError from services.tracing.service import ApplicationTracingService -from tests.unit_tests.config_override import apply_config_overrides +from tests.helpers.unit.config_override import apply_config_overrides CONTEXT = RequestContext("request", None, "account", "tenant") AGENT_ID = UUID("01a11111-1111-7111-8111-111111111111") diff --git a/api/tests/unit_tests/controllers/console/app/test_advanced_prompt_template_response.py b/api/tests/unit_tests/controllers/console/app/test_advanced_prompt_template.py similarity index 100% rename from api/tests/unit_tests/controllers/console/app/test_advanced_prompt_template_response.py rename to api/tests/unit_tests/controllers/console/app/test_advanced_prompt_template.py diff --git a/api/tests/unit_tests/controllers/console/app/test_agent_app_boundaries.py b/api/tests/unit_tests/controllers/console/app/test_agent_app_access.py similarity index 96% rename from api/tests/unit_tests/controllers/console/app/test_agent_app_boundaries.py rename to api/tests/unit_tests/controllers/console/app/test_agent_app_access.py index e72706a9d99..a0b2c944af6 100644 --- a/api/tests/unit_tests/controllers/console/app/test_agent_app_boundaries.py +++ b/api/tests/unit_tests/controllers/console/app/test_agent_app_access.py @@ -16,9 +16,9 @@ from models.account import TenantAccountRole from services.agent.errors import AgentNotFoundError from services.app.advanced_prompt_template_service import AdvancedPromptTemplateService from services.app.agent_app_contracts import AgentAppNotFoundError, AgentReferencingWorkflow -from tests.unit_tests.config_override import apply_config_overrides -from tests.unit_tests.controllers.rbac_introspection import rbac_checks -from tests.unit_tests.model_factories import make_account +from tests.helpers.unit.config_override import apply_config_overrides +from tests.helpers.unit.controllers.rbac_introspection import rbac_checks +from tests.helpers.unit.model_factories import make_account CONTEXT = RequestContext("request", None, "actor", "tenant") AGENT_ID = UUID("11111111-1111-1111-1111-111111111111") diff --git a/api/tests/unit_tests/controllers/console/app/test_agent_app_sandbox.py b/api/tests/unit_tests/controllers/console/app/test_agent_app_sandbox.py index 66c37fd0b18..a688c2a823d 100644 --- a/api/tests/unit_tests/controllers/console/app/test_agent_app_sandbox.py +++ b/api/tests/unit_tests/controllers/console/app/test_agent_app_sandbox.py @@ -26,8 +26,8 @@ from services.app.agent_app_contracts import ( WorkflowSandboxAppNotFoundError, WorkflowSandboxCaller, ) -from tests.unit_tests.controllers.error_body import raised_error -from tests.unit_tests.controllers.rbac_introspection import rbac_checks +from tests.helpers.unit.controllers.error_body import raised_error +from tests.helpers.unit.controllers.rbac_introspection import rbac_checks CONTEXT = RequestContext("request", "trace", "actor", "workspace") RESOURCE_ID = UUID("11111111-1111-1111-1111-111111111111") diff --git a/api/tests/unit_tests/controllers/console/app/test_agent_audio_readonly.py b/api/tests/unit_tests/controllers/console/app/test_agent_audio_readonly.py deleted file mode 100644 index 439f4309e23..00000000000 --- a/api/tests/unit_tests/controllers/console/app/test_agent_audio_readonly.py +++ /dev/null @@ -1,215 +0,0 @@ -from collections.abc import Callable, Iterator -from uuid import UUID - -import pytest -from flask import Flask -from sqlalchemy import Engine, event -from sqlalchemy.orm import Session -from werkzeug.exceptions import InternalServerError - -from controllers.console import wraps -from controllers.console.app import audio -from controllers.console.app.error import ( - CompletionRequestError, - ProviderModelCurrentlyNotSupportError, - ProviderNotInitializeError, - ProviderQuotaExceededError, -) -from core.errors.error import ModelCurrentlyNotSupportError, ProviderTokenNotInitError, QuotaExceededError -from graphon.model_runtime.errors.invoke import InvokeError -from models.agent import Agent, AgentKind, AgentScope, AgentSource, AgentStatus -from models.model import App, AppMode -from services.agent.errors import AgentNotFoundError -from tests.unit_tests.audio_runtime_fixtures import AudioRuntimeObservations -from tests.unit_tests.model_factories import make_account, make_app - -AGENT_ID = UUID("019ef3d2-b24c-7803-b428-18b5ee8fb853") - - -@pytest.fixture(autouse=True) -def audio_identity(monkeypatch: pytest.MonkeyPatch, config_overrides: Callable[..., None]) -> None: - config_overrides(LOGIN_DISABLED=True, RBAC_ENABLED=True) - monkeypatch.setattr(wraps, "_is_setup_completed", lambda: True) - monkeypatch.setattr(wraps, "current_account_with_tenant", lambda: (make_account(), "tenant-1")) - monkeypatch.setattr(audio, "enforce_rbac_checks", lambda **_kwargs: None) - - -@pytest.fixture -def persistence_events(sqlite_engine: Engine) -> Iterator[list[str]]: - observed: list[str] = [] - - def record_statement( - _connection: object, - _cursor: object, - statement: str, - _parameters: object, - _context: object, - _executemany: bool, - ) -> None: - operation = statement.lstrip().split(maxsplit=1)[0].upper() - if operation in {"INSERT", "UPDATE", "DELETE"}: - observed.append(operation) - - def record_commit(_session: Session) -> None: - observed.append("COMMIT") - - event.listen(sqlite_engine, "before_cursor_execute", record_statement) - event.listen(Session, "before_commit", record_commit) - try: - yield observed - finally: - event.remove(sqlite_engine, "before_cursor_execute", record_statement) - event.remove(Session, "before_commit", record_commit) - - -@pytest.fixture(params=["voices", "preview"]) -def request_audio( - request: pytest.FixtureRequest, app: Flask, audio_runtime: AudioRuntimeObservations -) -> tuple[Callable[[], object], AudioRuntimeObservations]: - if request.param == "voices": - audio_runtime.voices = [{"name": "Echo", "value": "echo"}] - - def invoke() -> object: - with app.test_request_context(f"/console/api/agent/{AGENT_ID}/text-to-audio/voices?language=en-US"): - return audio.AgentTextToSpeechVoicesApi().get(agent_id=AGENT_ID) - - else: - audio_runtime.tts_result = b"ID3audio" - - def invoke() -> object: - with app.test_request_context( - f"/console/api/agent/{AGENT_ID}/text-to-audio", - method="POST", - json={"text": "Preview this voice", "voice": "echo"}, - ): - response = audio.AgentChatMessageTextApi().post(agent_id=AGENT_ID) - assert response is not None - try: - assert response.get_data() == b"ID3audio" - assert response.content_type == "audio/mpeg" - return response - finally: - response.close() - - return invoke, audio_runtime - - -def _agent(scope: AgentScope = AgentScope.WORKFLOW_ONLY) -> Agent: - workflow_only = scope == AgentScope.WORKFLOW_ONLY - return Agent( - id=str(AGENT_ID), - tenant_id="tenant-1", - name="Audio agent", - agent_kind=AgentKind.DIFY_AGENT, - scope=scope, - source=AgentSource.WORKFLOW if workflow_only else AgentSource.AGENT_APP, - status=AgentStatus.ACTIVE, - app_id="parent-app-1" if workflow_only else "runtime-app-1", - backing_app_id="runtime-app-1" if workflow_only else None, - workflow_id="workflow-1" if workflow_only else None, - workflow_node_id="node-1" if workflow_only else None, - created_by="account-1", - updated_by="account-1", - ) - - -@pytest.mark.parametrize("scope", [AgentScope.ROSTER, AgentScope.WORKFLOW_ONLY]) -def test_audio_uses_existing_runtime_state_without_writes( - sqlite_session: Session, - request_audio: tuple[Callable[[], object], AudioRuntimeObservations], - persistence_events: list[str], - scope: AgentScope, -) -> None: - agent = _agent(scope) - runtime_app = make_app(app_id="runtime-app-1", mode=AppMode.AGENT) - sqlite_session.add_all([agent, runtime_app]) - sqlite_session.commit() - persistence_events.clear() - invoke, provider = request_audio - - response = invoke() - - assert len(provider.resolutions) == 1 - assert response is not None - assert persistence_events == [] - - -def test_audio_does_not_materialize_a_missing_workflow_backing_app( - sqlite_session: Session, - request_audio: tuple[Callable[[], object], AudioRuntimeObservations], - persistence_events: list[str], -) -> None: - agent = _agent() - agent.backing_app_id = None - parent_app = make_app(app_id="parent-app-1", mode=AppMode.WORKFLOW) - sqlite_session.add_all([agent, parent_app]) - sqlite_session.commit() - persistence_events.clear() - invoke, provider = request_audio - - with pytest.raises(AgentNotFoundError): - invoke() - - assert provider.resolutions == [] - assert provider.daemon_calls == [] - assert persistence_events == [] - sqlite_session.refresh(agent) - assert agent.backing_app_id is None - assert sqlite_session.get(App, "runtime-app-1") is None - - -@pytest.mark.parametrize("foreign_owner", ["agent", "app"]) -def test_audio_scopes_both_agent_and_runtime_app_to_the_current_tenant( - sqlite_session: Session, - request_audio: tuple[Callable[[], object], AudioRuntimeObservations], - persistence_events: list[str], - foreign_owner: str, -) -> None: - agent = _agent() - runtime_app = make_app(app_id="runtime-app-1", mode=AppMode.AGENT) - if foreign_owner == "agent": - agent.tenant_id = "tenant-2" - else: - runtime_app.tenant_id = "tenant-2" - sqlite_session.add_all([agent, runtime_app]) - sqlite_session.commit() - persistence_events.clear() - invoke, provider = request_audio - - with pytest.raises(AgentNotFoundError): - invoke() - - assert provider.resolutions == [] - assert provider.daemon_calls == [] - assert persistence_events == [] - - -@pytest.mark.parametrize( - ("provider_error", "expected_error"), - [ - (ProviderTokenNotInitError("No TTS model configured"), ProviderNotInitializeError), - (QuotaExceededError(), ProviderQuotaExceededError), - (ModelCurrentlyNotSupportError(), ProviderModelCurrentlyNotSupportError), - (InvokeError("TTS provider unavailable"), CompletionRequestError), - (ValueError("Unknown voice"), ValueError), - (RuntimeError("Unexpected provider failure"), InternalServerError), - ], -) -def test_provider_failures_preserve_error_contract_and_readonly_state( - sqlite_session: Session, - request_audio: tuple[Callable[[], object], AudioRuntimeObservations], - persistence_events: list[str], - provider_error: Exception, - expected_error: type[Exception], -) -> None: - sqlite_session.add_all([_agent(AgentScope.ROSTER), make_app(app_id="runtime-app-1", mode=AppMode.AGENT)]) - sqlite_session.commit() - persistence_events.clear() - invoke, provider = request_audio - provider.lookup_error = provider_error - - with pytest.raises(expected_error): - invoke() - - assert len(provider.resolutions) == 1 - assert persistence_events == [] diff --git a/api/tests/unit_tests/controllers/console/app/test_agent_config_inspector.py b/api/tests/unit_tests/controllers/console/app/test_agent_config_inspector.py index 39970a11b9d..ad19503a611 100644 --- a/api/tests/unit_tests/controllers/console/app/test_agent_config_inspector.py +++ b/api/tests/unit_tests/controllers/console/app/test_agent_config_inspector.py @@ -32,7 +32,7 @@ from controllers.console.app.agent_config_inspector import ( from models.account import Account from models.model import App, AppMode from services.agent_config_service import AgentConfigServiceError -from tests.unit_tests.controllers.error_body import raised_error +from tests.helpers.unit.controllers.error_body import raised_error _MOD = "controllers.console.app.agent_config_inspector" app = Flask(__name__) diff --git a/api/tests/unit_tests/controllers/console/app/test_annotation.py b/api/tests/unit_tests/controllers/console/app/test_annotation.py new file mode 100644 index 00000000000..021b811d977 --- /dev/null +++ b/api/tests/unit_tests/controllers/console/app/test_annotation.py @@ -0,0 +1,753 @@ +from __future__ import annotations + +import io +from datetime import datetime +from inspect import unwrap +from unittest.mock import Mock, patch + +import pytest +from flask import Flask +from pandas.errors import ParserError +from sqlalchemy.orm import Session +from werkzeug.datastructures import FileStorage +from werkzeug.exceptions import NotFound + +from configs import dify_config +from controllers.console.app import annotation as annotation_module +from controllers.console.wraps import annotation_import_concurrency_limit, annotation_import_rate_limit +from enums import DeploymentEdition +from models.account import Account +from models.model import App, AppAnnotationHitHistory, AppMode, IconType, MessageAnnotation +from services.annotation_service import AppAnnotationService +from services.app_ref_service import AnnotationRef, AppRef +from tasks.annotation.batch_import_annotations_task import batch_import_annotations_task +from tests.helpers.unit.config_override import config_overrides_context +from tests.helpers.unit.model_factories import make_account + + +def _annotation_api_persist_app(session: Session) -> App: + app = App( + id="app-1", + tenant_id="tenant-1", + name="Annotation app", + mode=AppMode.CHAT, + icon_type=IconType.EMOJI, + icon="chat", + icon_background="#ffffff", + enable_site=False, + enable_api=True, + ) + session.add(app) + session.commit() + return app + + +def _annotation_api_account() -> Account: + return make_account(name="Owner", email="owner@example.com") + + +def _annotation_model(annotation_id: str = "ann-1") -> MessageAnnotation: + annotation = MessageAnnotation( + app_id="app-1", + question="q", + content="a", + account_id="account-1", + ) + annotation.id = annotation_id + annotation.hit_count = 0 + annotation.created_at = datetime(2026, 1, 1) + return annotation + + +def _annotation_hit_history() -> AppAnnotationHitHistory: + history = AppAnnotationHitHistory( + app_id="app-1", + annotation_id="ann-1", + source="hit-testing", + score=0.9, + question="q", + annotation_question="q", + annotation_content="a", + account_id="account-1", + message_id="message-1", + ) + history.id = "history-1" + history.created_at = datetime(2026, 1, 1) + return history + + +def test_annotation_reply_payload_valid(): + """Test AnnotationReplyPayload with valid data.""" + payload = annotation_module.AnnotationReplyPayload( + score_threshold=0.5, + embedding_provider_name="openai", + embedding_model_name="text-embedding-3-small", + ) + assert payload.score_threshold == 0.5 + assert payload.embedding_provider_name == "openai" + assert payload.embedding_model_name == "text-embedding-3-small" + + +def test_annotation_setting_update_payload_valid(): + """Test AnnotationSettingUpdatePayload with valid data.""" + payload = annotation_module.AnnotationSettingUpdatePayload( + score_threshold=0.75, + ) + assert payload.score_threshold == 0.75 + + +def test_annotation_list_query_defaults(): + """Test AnnotationListQuery with default parameters.""" + query = annotation_module.AnnotationListQuery() + assert query.page == 1 + assert query.limit == 20 + assert query.keyword == "" + + +def test_annotation_list_query_custom_page(): + """Test AnnotationListQuery with custom page.""" + query = annotation_module.AnnotationListQuery(page=3, limit=50) + assert query.page == 3 + assert query.limit == 50 + + +def test_annotation_list_query_with_keyword(): + """Test AnnotationListQuery with keyword.""" + query = annotation_module.AnnotationListQuery(keyword="test") + assert query.keyword == "test" + + +def test_create_annotation_payload_with_message_id(): + """Test CreateAnnotationPayload with message ID.""" + payload = annotation_module.CreateAnnotationPayload( + message_id="550e8400-e29b-41d4-a716-446655440000", + question="What is AI?", + ) + assert payload.message_id == "550e8400-e29b-41d4-a716-446655440000" + assert payload.question == "What is AI?" + + +def test_create_annotation_payload_with_text(): + """Test CreateAnnotationPayload with text content.""" + payload = annotation_module.CreateAnnotationPayload( + question="What is ML?", + answer="Machine learning is...", + ) + assert payload.question == "What is ML?" + assert payload.answer == "Machine learning is..." + + +def test_update_annotation_payload(): + """Test UpdateAnnotationPayload.""" + payload = annotation_module.UpdateAnnotationPayload( + question="Updated question", + answer="Updated answer", + ) + assert payload.question == "Updated question" + assert payload.answer == "Updated answer" + + +def test_annotation_reply_status_query_enable(): + """Test AnnotationReplyStatusQuery with enable action.""" + query = annotation_module.AnnotationReplyStatusQuery(action="enable") + assert query.action == "enable" + + +def test_annotation_reply_status_query_disable(): + """Test AnnotationReplyStatusQuery with disable action.""" + query = annotation_module.AnnotationReplyStatusQuery(action="disable") + assert query.action == "disable" + + +def test_annotation_file_payload_valid(): + """Test AnnotationFilePayload with valid message ID.""" + payload = annotation_module.AnnotationFilePayload(message_id="550e8400-e29b-41d4-a716-446655440000") + assert payload.message_id == "550e8400-e29b-41d4-a716-446655440000" + + +def test_get_app_ref_raises_not_found_when_app_is_not_in_current_tenant(sqlite_session: Session): + _annotation_api_persist_app(sqlite_session) + with ( + patch.object( + annotation_module, + "current_account_with_tenant", + return_value=(_annotation_api_account(), "tenant-2"), + ), + ): + with pytest.raises(NotFound): + annotation_module._get_app_ref(sqlite_session, "app-1") + + +class TestConsoleAnnotationRefBoundaries: + def test_batch_delete_uses_app_ref(self, app: Flask, sqlite_session: Session): + api = annotation_module.AnnotationApi() + handler = unwrap(api.delete) + delete_mock = Mock() + _annotation_api_persist_app(sqlite_session) + + with ( + app.test_request_context("/?annotation_id=ann-1&annotation_id=ann-2", method="DELETE"), + patch.object( + annotation_module, + "current_account_with_tenant", + return_value=(_annotation_api_account(), "tenant-1"), + ), + patch.object(annotation_module.AppAnnotationService, "delete_app_annotations_in_batch", delete_mock), + ): + response, status = handler(api, sqlite_session, "app-1") + + assert response == "" + assert status == 204 + delete_mock.assert_called_once_with(AppRef("tenant-1", "app-1"), ["ann-1", "ann-2"], sqlite_session) + + def test_update_uses_annotation_ref(self, app: Flask, sqlite_session: Session): + api = annotation_module.AnnotationUpdateDeleteApi() + handler = unwrap(api.post) + update_mock = Mock(return_value=_annotation_model()) + _annotation_api_persist_app(sqlite_session) + + with ( + app.test_request_context("/annotations/ann-1", method="POST", json={"question": "updated"}), + patch.object( + annotation_module, + "current_account_with_tenant", + return_value=(_annotation_api_account(), "tenant-1"), + ), + patch.object(annotation_module.AppAnnotationService, "update_app_annotation_directly", update_mock), + ): + response = handler( + api, + annotation_module.UpdateAnnotationPayload(question="updated"), + sqlite_session, + "app-1", + "ann-1", + ) + + assert response["question"] == "q" + update_mock.assert_called_once() + assert update_mock.call_args.args[1] == AnnotationRef(AppRef("tenant-1", "app-1"), "ann-1") + assert update_mock.call_args.args[2] is sqlite_session + + def test_delete_uses_annotation_ref(self, app: Flask, sqlite_session: Session): + api = annotation_module.AnnotationUpdateDeleteApi() + handler = unwrap(api.delete) + delete_mock = Mock() + _annotation_api_persist_app(sqlite_session) + + with ( + app.test_request_context("/annotations/ann-1", method="DELETE"), + patch.object( + annotation_module, + "current_account_with_tenant", + return_value=(_annotation_api_account(), "tenant-1"), + ), + patch.object(annotation_module.AppAnnotationService, "delete_app_annotation", delete_mock), + ): + response, status = handler(api, sqlite_session, "app-1", "ann-1") + + assert response == "" + assert status == 204 + delete_mock.assert_called_once() + assert delete_mock.call_args.args[0] == AnnotationRef(AppRef("tenant-1", "app-1"), "ann-1") + assert delete_mock.call_args.args[1] is sqlite_session + + def test_hit_history_uses_annotation_ref(self, app: Flask, sqlite_session: Session): + api = annotation_module.AnnotationHitHistoryListApi() + handler = unwrap(api.get) + history = _annotation_hit_history() + hit_history_mock = Mock(return_value=([history], 1)) + _annotation_api_persist_app(sqlite_session) + with ( + app.test_request_context("/hit-histories?page=2&limit=5", method="GET"), + patch.object( + annotation_module, + "current_account_with_tenant", + return_value=(_annotation_api_account(), "tenant-1"), + ), + patch.object(annotation_module.AppAnnotationService, "get_annotation_hit_histories", hit_history_mock), + ): + response = handler(api, sqlite_session, "app-1", "ann-1") + assert response["total"] == 1 + hit_history_mock.assert_called_once_with( + AnnotationRef(AppRef("tenant-1", "app-1"), "ann-1"), 2, 5, sqlite_session + ) + + def test_hit_history_pager_ends_for_a_non_positive_limit(self, app: Flask, sqlite_session: Session): + """This route reads page and limit straight off the query string with no bounds. + + `paginate_query` floors both at 1, so `limit=0` is served one row. Reporting + `has_more` from the requested 0 made `page * limit < total` true on every page, + including the pages past the end that return nothing, so a client walking until + `has_more` is false never stopped. + """ + api = annotation_module.AnnotationHitHistoryListApi() + handler = unwrap(api.get) + hit_history_mock = Mock(return_value=([_annotation_hit_history()], 1)) + _annotation_api_persist_app(sqlite_session) + with ( + app.test_request_context("/hit-histories?page=1&limit=0", method="GET"), + patch.object( + annotation_module, + "current_account_with_tenant", + return_value=(_annotation_api_account(), "tenant-1"), + ), + patch.object(annotation_module.AppAnnotationService, "get_annotation_hit_histories", hit_history_mock), + ): + response = handler(api, sqlite_session, "app-1", "ann-1") + assert response["limit"] == 1 + assert response["page"] == 1 + assert response["has_more"] is False + hit_history_mock.assert_called_once_with( + AnnotationRef(AppRef("tenant-1", "app-1"), "ann-1"), 1, 1, sqlite_session + ) + + def test_hit_history_clamps_a_non_positive_page(self, app: Flask, sqlite_session: Session): + """A page below 1 is served as page 1, so the response has to say page 1.""" + api = annotation_module.AnnotationHitHistoryListApi() + handler = unwrap(api.get) + hit_history_mock = Mock(return_value=([_annotation_hit_history()], 9)) + _annotation_api_persist_app(sqlite_session) + with ( + app.test_request_context("/hit-histories?page=0&limit=5", method="GET"), + patch.object( + annotation_module, + "current_account_with_tenant", + return_value=(_annotation_api_account(), "tenant-1"), + ), + patch.object(annotation_module.AppAnnotationService, "get_annotation_hit_histories", hit_history_mock), + ): + response = handler(api, sqlite_session, "app-1", "ann-1") + assert response["page"] == 1 + assert response["has_more"] is True + hit_history_mock.assert_called_once_with( + AnnotationRef(AppRef("tenant-1", "app-1"), "ann-1"), 1, 5, sqlite_session + ) + + +class TestAnnotationListHasMore: + def test_list_annotations_has_more_false_on_last_page_exact_limit( + self, app: Flask, monkeypatch: pytest.MonkeyPatch + ) -> None: + """A full last page must set has_more false instead of forcing another fetch.""" + page_size = 20 + annotation = _annotation_model() + get_mock = Mock(return_value=([annotation] * page_size, page_size)) + monkeypatch.setattr(annotation_module.AppAnnotationService, "get_annotation_list_by_app_id", get_mock) + api = annotation_module.AnnotationApi() + handler = unwrap(api.get) + with app.test_request_context(f"/apps/app-1/annotations?page=1&limit={page_size}", method="GET"): + response, status = handler( + api, annotation_module.AnnotationListQuery(page=1, limit=page_size), Mock(), "app-1" + ) + assert status == 200 + assert response["has_more"] is False + assert response["limit"] == page_size + assert response["total"] == page_size + assert response["page"] == 1 + + def test_list_annotations_has_more_true_when_limit_exceeds_cap( + self, app: Flask, monkeypatch: pytest.MonkeyPatch + ) -> None: + """limit>100 still reports remaining rows after the server cap of 100.""" + annotation = _annotation_model() + get_mock = Mock(return_value=([annotation] * 100, 150)) + monkeypatch.setattr(annotation_module.AppAnnotationService, "get_annotation_list_by_app_id", get_mock) + api = annotation_module.AnnotationApi() + handler = unwrap(api.get) + with app.test_request_context("/apps/app-1/annotations?page=1&limit=200", method="GET"): + response, status = handler(api, annotation_module.AnnotationListQuery(page=1, limit=200), Mock(), "app-1") + assert status == 200 + assert response["has_more"] is True + assert response["limit"] == 100 + assert response["total"] == 150 + assert response["page"] == 1 + + def test_hit_histories_has_more_false_on_last_page_exact_limit(self, app: Flask, sqlite_session: Session): + api = annotation_module.AnnotationHitHistoryListApi() + handler = unwrap(api.get) + history = _annotation_hit_history() + hit_history_mock = Mock(return_value=([history] * 20, 20)) + _annotation_api_persist_app(sqlite_session) + + with ( + app.test_request_context("/hit-histories?page=1&limit=20", method="GET"), + patch.object( + annotation_module, + "current_account_with_tenant", + return_value=(_annotation_api_account(), "tenant-1"), + ), + patch.object(annotation_module.AppAnnotationService, "get_annotation_hit_histories", hit_history_mock), + ): + response = handler(api, sqlite_session, "app-1", "ann-1") + + assert response["has_more"] is False + assert response["limit"] == 20 + assert response["total"] == 20 + assert response["page"] == 1 + + def test_hit_histories_has_more_true_when_limit_exceeds_cap(self, app: Flask, sqlite_session: Session): + api = annotation_module.AnnotationHitHistoryListApi() + handler = unwrap(api.get) + history = _annotation_hit_history() + hit_history_mock = Mock(return_value=([history] * 100, 150)) + _annotation_api_persist_app(sqlite_session) + + with ( + app.test_request_context("/hit-histories?page=1&limit=200", method="GET"), + patch.object( + annotation_module, + "current_account_with_tenant", + return_value=(_annotation_api_account(), "tenant-1"), + ), + patch.object(annotation_module.AppAnnotationService, "get_annotation_hit_histories", hit_history_mock), + ): + response = handler(api, sqlite_session, "app-1", "ann-1") + + assert response["has_more"] is True + assert response["limit"] == 100 + assert response["total"] == 150 + assert response["page"] == 1 + + +""" +Unit tests for annotation import security features. + +Tests rate limiting, concurrency control, file validation, and other +security features added to prevent DoS attacks on the annotation import endpoint. +""" + + +def _annotation_security_account() -> Account: + return make_account(account_id="user_id", name="Annotation Tester", email="annotation-tester@example.com") + + +def _annotation_security_persist_app(session: Session) -> App: + app = App( + id="app_id", + tenant_id="tenant_id", + name="Annotation App", + description="", + mode=AppMode.CHAT, + icon_type=IconType.EMOJI, + icon="chat", + icon_background="#ffffff", + enable_site=False, + enable_api=False, + ) + session.add(app) + session.commit() + return app + + +class TestAnnotationImportRateLimiting: + """Test rate limiting for annotation import operations.""" + + @pytest.fixture + def mock_redis(self): + """Mock Redis client for testing.""" + with patch("controllers.console.wraps.redis_client") as mock: + yield mock + + @pytest.fixture + def mock_current_account(self): + """Mock current account with tenant.""" + with patch("controllers.console.wraps.current_account_with_tenant") as mock: + mock.return_value = (_annotation_security_account(), "test_tenant_id") + yield mock + + def test_rate_limit_per_minute_enforced(self, mock_redis, mock_current_account): + """Test that per-minute rate limit is enforced.""" + # Simulate exceeding per-minute limit + mock_redis.zcard.side_effect = [ + dify_config.ANNOTATION_IMPORT_RATE_LIMIT_PER_MINUTE + 1, # Minute check + 10, # Hour check + ] + + @annotation_import_rate_limit + def dummy_view(): + return "success" + + # Should abort with 429 + with pytest.raises( + Exception, + match=( + "429 Too Many Requests: Too many annotation import requests\\. Maximum " + "5 requests per minute allowed\\. Please try again later\\." + ), + ) as exc_info: + dummy_view() + + # Verify it's a rate limit error + assert "429" in str(exc_info.value) or "Too many" in str(exc_info.value) + + def test_rate_limit_per_hour_enforced(self, mock_redis, mock_current_account): + """Test that per-hour rate limit is enforced.""" + + # Simulate exceeding per-hour limit + mock_redis.zcard.side_effect = [ + 3, # Minute check (under limit) + dify_config.ANNOTATION_IMPORT_RATE_LIMIT_PER_HOUR + 1, # Hour check (over limit) + ] + + @annotation_import_rate_limit + def dummy_view(): + return "success" + + # Should abort with 429 + with pytest.raises( + Exception, + match=( + "429 Too Many Requests: Too many annotation import requests\\. Maximum " + "20 requests per hour allowed\\. Please try again later\\." + ), + ) as exc_info: + dummy_view() + + assert "429" in str(exc_info.value) or "Too many" in str(exc_info.value) + + def test_rate_limit_within_limits_passes(self, mock_redis, mock_current_account): + """Test that requests within limits are allowed.""" + + # Simulate being under both limits + mock_redis.zcard.return_value = 2 + + @annotation_import_rate_limit + def dummy_view(): + return "success" + + # Should succeed + result = dummy_view() + assert result == "success" + + # Verify Redis operations were called + assert mock_redis.zadd.called + assert mock_redis.zremrangebyscore.called + + +class TestAnnotationImportConcurrencyControl: + """Test concurrency control for annotation import operations.""" + + @pytest.fixture + def mock_redis(self): + """Mock Redis client for testing.""" + with patch("controllers.console.wraps.redis_client") as mock: + yield mock + + @pytest.fixture + def mock_current_account(self): + """Mock current account with tenant.""" + with patch("controllers.console.wraps.current_account_with_tenant") as mock: + mock.return_value = (_annotation_security_account(), "test_tenant_id") + yield mock + + def test_concurrency_limit_enforced(self, mock_redis, mock_current_account): + """Test that concurrent task limit is enforced.""" + + # Simulate max concurrent tasks already running + mock_redis.zcard.return_value = dify_config.ANNOTATION_IMPORT_MAX_CONCURRENT + + @annotation_import_concurrency_limit + def dummy_view(): + return "success" + + # Should abort with 429 + with pytest.raises( + Exception, + match=( + "429 Too Many Requests: Too many concurrent import tasks\\. Maximum 2 " + "concurrent imports allowed per workspace\\. Please wait for existing " + "imports to complete\\." + ), + ) as exc_info: + dummy_view() + + assert "429" in str(exc_info.value) or "concurrent" in str(exc_info.value).lower() + + def test_concurrency_within_limit_passes(self, mock_redis, mock_current_account): + """Test that requests within concurrency limits are allowed.""" + + # Simulate being under concurrent task limit + mock_redis.zcard.return_value = 1 + + @annotation_import_concurrency_limit + def dummy_view(): + return "success" + + # Should succeed + result = dummy_view() + assert result == "success" + + def test_stale_jobs_are_cleaned_up(self, mock_redis, mock_current_account): + """Test that old/stale job entries are removed.""" + + mock_redis.zcard.return_value = 0 + + @annotation_import_concurrency_limit + def dummy_view(): + return "success" + + dummy_view() + + # Verify cleanup was called + assert mock_redis.zremrangebyscore.called + + +class TestAnnotationImportFileValidation: + """Test file validation in annotation import.""" + + def test_file_size_limit_enforced(self): + """Test that files exceeding size limit are rejected.""" + # Create a file larger than the limit + max_size = dify_config.ANNOTATION_IMPORT_FILE_SIZE_LIMIT * 1024 * 1024 + large_content = b"x" * (max_size + 1024) # Exceed by 1KB + + FileStorage(stream=io.BytesIO(large_content), filename="test.csv", content_type="text/csv") + + # Should be rejected in controller + # This would be tested in integration tests with actual endpoint + + def test_empty_file_rejected(self): + """Test that empty files are rejected.""" + FileStorage(stream=io.BytesIO(b""), filename="test.csv", content_type="text/csv") + + # Should be rejected + # This would be tested in integration tests + + def test_non_csv_file_rejected(self): + """Test that non-CSV files are rejected.""" + FileStorage(stream=io.BytesIO(b"test"), filename="test.txt", content_type="text/plain") + + # Should be rejected based on extension + # This would be tested in integration tests + + +@pytest.mark.parametrize("sqlite_session", [(App,)], indirect=True) +class TestAnnotationImportServiceValidation: + """Test service layer validation for annotation import.""" + + @config_overrides_context(DEPLOYMENT_EDITION=DeploymentEdition.COMMUNITY) + def test_max_records_limit_enforced(self, sqlite_session: Session): + """Test that files with too many records are rejected.""" + + # Create CSV with too many records + max_records = dify_config.ANNOTATION_IMPORT_MAX_RECORDS + csv_content = "question,answer\n" + for i in range(max_records + 100): + csv_content += f"Question {i},Answer {i}\n" + + file = FileStorage(stream=io.BytesIO(csv_content.encode()), filename="test.csv", content_type="text/csv") + + _annotation_security_persist_app(sqlite_session) + with patch("services.annotation_service.current_account_with_tenant") as mock_auth: + mock_auth.return_value = (_annotation_security_account(), "tenant_id") + + with patch("services.annotation_service.FeatureService"): + result = AppAnnotationService.batch_import_app_annotations("app_id", file, sqlite_session) + + # Should return error about too many records + assert "error_msg" in result + assert "too many" in result["error_msg"].lower() or "maximum" in result["error_msg"].lower() + + def test_min_records_limit_enforced(self, sqlite_session: Session): + """Test that files with too few valid records are rejected.""" + + # Create CSV with only header (no data rows) + csv_content = "question,answer\n" + + file = FileStorage(stream=io.BytesIO(csv_content.encode()), filename="test.csv", content_type="text/csv") + + _annotation_security_persist_app(sqlite_session) + with patch("services.annotation_service.current_account_with_tenant") as mock_auth: + mock_auth.return_value = (_annotation_security_account(), "tenant_id") + + result = AppAnnotationService.batch_import_app_annotations("app_id", file, sqlite_session) + + # Should return error about insufficient records + assert "error_msg" in result + assert "at least" in result["error_msg"].lower() or "minimum" in result["error_msg"].lower() + + def test_invalid_csv_format_handled(self, sqlite_session: Session): + """Test that invalid CSV format is handled gracefully.""" + + # Any content is fine once we force ParserError + csv_content = 'invalid,csv,format\nwith,unbalanced,quotes,and"stuff' + file = FileStorage(stream=io.BytesIO(csv_content.encode()), filename="test.csv", content_type="text/csv") + + _annotation_security_persist_app(sqlite_session) + with ( + patch("services.annotation_service.current_account_with_tenant") as mock_auth, + patch("services.annotation_service.pd.read_csv", side_effect=ParserError("malformed CSV")), + ): + mock_auth.return_value = (_annotation_security_account(), "tenant_id") + + result = AppAnnotationService.batch_import_app_annotations("app_id", file, sqlite_session) + + assert "error_msg" in result + assert "malformed" in result["error_msg"].lower() + + @config_overrides_context(DEPLOYMENT_EDITION=DeploymentEdition.COMMUNITY) + def test_valid_import_succeeds(self, sqlite_session: Session): + """Test that valid import request succeeds.""" + + # Create valid CSV + csv_content = "question,answer\nWhat is AI?,Artificial Intelligence\nWhat is ML?,Machine Learning\n" + + file = FileStorage(stream=io.BytesIO(csv_content.encode()), filename="test.csv", content_type="text/csv") + + _annotation_security_persist_app(sqlite_session) + with patch("services.annotation_service.current_account_with_tenant") as mock_auth: + mock_auth.return_value = (_annotation_security_account(), "tenant_id") + + with patch("services.annotation_service.FeatureService"): + with patch("services.annotation_service.batch_import_annotations_task"): + with patch("services.annotation_service.redis_client"): + result = AppAnnotationService.batch_import_app_annotations("app_id", file, sqlite_session) + + # Should return success response + assert "job_id" in result + assert "job_status" in result + assert result["job_status"] == "waiting" + assert "record_count" in result + assert result["record_count"] == 2 + + +class TestAnnotationImportTaskOptimization: + """Test optimizations in batch import task.""" + + def test_task_is_registered_with_queue(self): + """Test that task is registered with the correct queue.""" + assert hasattr(batch_import_annotations_task, "apply_async") + assert hasattr(batch_import_annotations_task, "delay") + + +class TestConfigurationValues: + """Test that security configuration values are properly set.""" + + def test_rate_limit_configs_exist(self): + """Test that rate limit configurations are defined.""" + assert hasattr(dify_config, "ANNOTATION_IMPORT_RATE_LIMIT_PER_MINUTE") + assert hasattr(dify_config, "ANNOTATION_IMPORT_RATE_LIMIT_PER_HOUR") + + assert dify_config.ANNOTATION_IMPORT_RATE_LIMIT_PER_MINUTE > 0 + assert dify_config.ANNOTATION_IMPORT_RATE_LIMIT_PER_HOUR > 0 + + def test_file_size_limit_config_exists(self): + """Test that file size limit configuration is defined.""" + assert hasattr(dify_config, "ANNOTATION_IMPORT_FILE_SIZE_LIMIT") + assert dify_config.ANNOTATION_IMPORT_FILE_SIZE_LIMIT > 0 + assert dify_config.ANNOTATION_IMPORT_FILE_SIZE_LIMIT <= 10 # Reasonable max (10MB) + + def test_record_limit_configs_exist(self): + """Test that record limit configurations are defined.""" + assert hasattr(dify_config, "ANNOTATION_IMPORT_MAX_RECORDS") + assert hasattr(dify_config, "ANNOTATION_IMPORT_MIN_RECORDS") + + assert dify_config.ANNOTATION_IMPORT_MAX_RECORDS > 0 + assert dify_config.ANNOTATION_IMPORT_MIN_RECORDS > 0 + assert dify_config.ANNOTATION_IMPORT_MIN_RECORDS < dify_config.ANNOTATION_IMPORT_MAX_RECORDS + + def test_concurrency_limit_config_exists(self): + """Test that concurrency limit configuration is defined.""" + assert hasattr(dify_config, "ANNOTATION_IMPORT_MAX_CONCURRENT") + assert dify_config.ANNOTATION_IMPORT_MAX_CONCURRENT > 0 + assert dify_config.ANNOTATION_IMPORT_MAX_CONCURRENT <= 10 # Reasonable upper bound diff --git a/api/tests/unit_tests/controllers/console/app/test_annotation_api.py b/api/tests/unit_tests/controllers/console/app/test_annotation_api.py deleted file mode 100644 index 5094bb1fab9..00000000000 --- a/api/tests/unit_tests/controllers/console/app/test_annotation_api.py +++ /dev/null @@ -1,400 +0,0 @@ -from __future__ import annotations - -from datetime import datetime -from inspect import unwrap -from unittest.mock import Mock, patch - -import pytest -from flask import Flask -from sqlalchemy.orm import Session -from werkzeug.exceptions import NotFound - -from controllers.console.app import annotation as annotation_module -from models.account import Account -from models.model import App, AppAnnotationHitHistory, AppMode, IconType, MessageAnnotation -from services.app_ref_service import AnnotationRef, AppRef -from tests.unit_tests.model_factories import make_account - - -def _persist_app(session: Session) -> App: - app = App( - id="app-1", - tenant_id="tenant-1", - name="Annotation app", - mode=AppMode.CHAT, - icon_type=IconType.EMOJI, - icon="chat", - icon_background="#ffffff", - enable_site=False, - enable_api=True, - ) - session.add(app) - session.commit() - return app - - -def _account() -> Account: - return make_account(name="Owner", email="owner@example.com") - - -def _annotation_model(annotation_id: str = "ann-1") -> MessageAnnotation: - annotation = MessageAnnotation( - app_id="app-1", - question="q", - content="a", - account_id="account-1", - ) - annotation.id = annotation_id - annotation.hit_count = 0 - annotation.created_at = datetime(2026, 1, 1) - return annotation - - -def _annotation_hit_history() -> AppAnnotationHitHistory: - history = AppAnnotationHitHistory( - app_id="app-1", - annotation_id="ann-1", - source="hit-testing", - score=0.9, - question="q", - annotation_question="q", - annotation_content="a", - account_id="account-1", - message_id="message-1", - ) - history.id = "history-1" - history.created_at = datetime(2026, 1, 1) - return history - - -def test_annotation_reply_payload_valid(): - """Test AnnotationReplyPayload with valid data.""" - payload = annotation_module.AnnotationReplyPayload( - score_threshold=0.5, - embedding_provider_name="openai", - embedding_model_name="text-embedding-3-small", - ) - assert payload.score_threshold == 0.5 - assert payload.embedding_provider_name == "openai" - assert payload.embedding_model_name == "text-embedding-3-small" - - -def test_annotation_setting_update_payload_valid(): - """Test AnnotationSettingUpdatePayload with valid data.""" - payload = annotation_module.AnnotationSettingUpdatePayload( - score_threshold=0.75, - ) - assert payload.score_threshold == 0.75 - - -def test_annotation_list_query_defaults(): - """Test AnnotationListQuery with default parameters.""" - query = annotation_module.AnnotationListQuery() - assert query.page == 1 - assert query.limit == 20 - assert query.keyword == "" - - -def test_annotation_list_query_custom_page(): - """Test AnnotationListQuery with custom page.""" - query = annotation_module.AnnotationListQuery(page=3, limit=50) - assert query.page == 3 - assert query.limit == 50 - - -def test_annotation_list_query_with_keyword(): - """Test AnnotationListQuery with keyword.""" - query = annotation_module.AnnotationListQuery(keyword="test") - assert query.keyword == "test" - - -def test_create_annotation_payload_with_message_id(): - """Test CreateAnnotationPayload with message ID.""" - payload = annotation_module.CreateAnnotationPayload( - message_id="550e8400-e29b-41d4-a716-446655440000", - question="What is AI?", - ) - assert payload.message_id == "550e8400-e29b-41d4-a716-446655440000" - assert payload.question == "What is AI?" - - -def test_create_annotation_payload_with_text(): - """Test CreateAnnotationPayload with text content.""" - payload = annotation_module.CreateAnnotationPayload( - question="What is ML?", - answer="Machine learning is...", - ) - assert payload.question == "What is ML?" - assert payload.answer == "Machine learning is..." - - -def test_update_annotation_payload(): - """Test UpdateAnnotationPayload.""" - payload = annotation_module.UpdateAnnotationPayload( - question="Updated question", - answer="Updated answer", - ) - assert payload.question == "Updated question" - assert payload.answer == "Updated answer" - - -def test_annotation_reply_status_query_enable(): - """Test AnnotationReplyStatusQuery with enable action.""" - query = annotation_module.AnnotationReplyStatusQuery(action="enable") - assert query.action == "enable" - - -def test_annotation_reply_status_query_disable(): - """Test AnnotationReplyStatusQuery with disable action.""" - query = annotation_module.AnnotationReplyStatusQuery(action="disable") - assert query.action == "disable" - - -def test_annotation_file_payload_valid(): - """Test AnnotationFilePayload with valid message ID.""" - payload = annotation_module.AnnotationFilePayload(message_id="550e8400-e29b-41d4-a716-446655440000") - assert payload.message_id == "550e8400-e29b-41d4-a716-446655440000" - - -def test_get_app_ref_raises_not_found_when_app_is_not_in_current_tenant(sqlite_session: Session): - _persist_app(sqlite_session) - with ( - patch.object( - annotation_module, - "current_account_with_tenant", - return_value=(_account(), "tenant-2"), - ), - ): - with pytest.raises(NotFound): - annotation_module._get_app_ref(sqlite_session, "app-1") - - -class TestConsoleAnnotationRefBoundaries: - def test_batch_delete_uses_app_ref(self, app: Flask, sqlite_session: Session): - api = annotation_module.AnnotationApi() - handler = unwrap(api.delete) - delete_mock = Mock() - _persist_app(sqlite_session) - - with ( - app.test_request_context("/?annotation_id=ann-1&annotation_id=ann-2", method="DELETE"), - patch.object( - annotation_module, - "current_account_with_tenant", - return_value=(_account(), "tenant-1"), - ), - patch.object(annotation_module.AppAnnotationService, "delete_app_annotations_in_batch", delete_mock), - ): - response, status = handler(api, sqlite_session, "app-1") - - assert response == "" - assert status == 204 - delete_mock.assert_called_once_with(AppRef("tenant-1", "app-1"), ["ann-1", "ann-2"], sqlite_session) - - def test_update_uses_annotation_ref(self, app: Flask, sqlite_session: Session): - api = annotation_module.AnnotationUpdateDeleteApi() - handler = unwrap(api.post) - update_mock = Mock(return_value=_annotation_model()) - _persist_app(sqlite_session) - - with ( - app.test_request_context("/annotations/ann-1", method="POST", json={"question": "updated"}), - patch.object( - annotation_module, - "current_account_with_tenant", - return_value=(_account(), "tenant-1"), - ), - patch.object(annotation_module.AppAnnotationService, "update_app_annotation_directly", update_mock), - ): - response = handler( - api, - annotation_module.UpdateAnnotationPayload(question="updated"), - sqlite_session, - "app-1", - "ann-1", - ) - - assert response["question"] == "q" - update_mock.assert_called_once() - assert update_mock.call_args.args[1] == AnnotationRef(AppRef("tenant-1", "app-1"), "ann-1") - assert update_mock.call_args.args[2] is sqlite_session - - def test_delete_uses_annotation_ref(self, app: Flask, sqlite_session: Session): - api = annotation_module.AnnotationUpdateDeleteApi() - handler = unwrap(api.delete) - delete_mock = Mock() - _persist_app(sqlite_session) - - with ( - app.test_request_context("/annotations/ann-1", method="DELETE"), - patch.object( - annotation_module, - "current_account_with_tenant", - return_value=(_account(), "tenant-1"), - ), - patch.object(annotation_module.AppAnnotationService, "delete_app_annotation", delete_mock), - ): - response, status = handler(api, sqlite_session, "app-1", "ann-1") - - assert response == "" - assert status == 204 - delete_mock.assert_called_once() - assert delete_mock.call_args.args[0] == AnnotationRef(AppRef("tenant-1", "app-1"), "ann-1") - assert delete_mock.call_args.args[1] is sqlite_session - - def test_hit_history_uses_annotation_ref(self, app: Flask, sqlite_session: Session): - api = annotation_module.AnnotationHitHistoryListApi() - handler = unwrap(api.get) - history = _annotation_hit_history() - hit_history_mock = Mock(return_value=([history], 1)) - _persist_app(sqlite_session) - with ( - app.test_request_context("/hit-histories?page=2&limit=5", method="GET"), - patch.object( - annotation_module, - "current_account_with_tenant", - return_value=(_account(), "tenant-1"), - ), - patch.object(annotation_module.AppAnnotationService, "get_annotation_hit_histories", hit_history_mock), - ): - response = handler(api, sqlite_session, "app-1", "ann-1") - assert response["total"] == 1 - hit_history_mock.assert_called_once_with( - AnnotationRef(AppRef("tenant-1", "app-1"), "ann-1"), 2, 5, sqlite_session - ) - - def test_hit_history_pager_ends_for_a_non_positive_limit(self, app: Flask, sqlite_session: Session): - """This route reads page and limit straight off the query string with no bounds. - - `paginate_query` floors both at 1, so `limit=0` is served one row. Reporting - `has_more` from the requested 0 made `page * limit < total` true on every page, - including the pages past the end that return nothing, so a client walking until - `has_more` is false never stopped. - """ - api = annotation_module.AnnotationHitHistoryListApi() - handler = unwrap(api.get) - hit_history_mock = Mock(return_value=([_annotation_hit_history()], 1)) - _persist_app(sqlite_session) - with ( - app.test_request_context("/hit-histories?page=1&limit=0", method="GET"), - patch.object( - annotation_module, - "current_account_with_tenant", - return_value=(_account(), "tenant-1"), - ), - patch.object(annotation_module.AppAnnotationService, "get_annotation_hit_histories", hit_history_mock), - ): - response = handler(api, sqlite_session, "app-1", "ann-1") - assert response["limit"] == 1 - assert response["page"] == 1 - assert response["has_more"] is False - hit_history_mock.assert_called_once_with( - AnnotationRef(AppRef("tenant-1", "app-1"), "ann-1"), 1, 1, sqlite_session - ) - - def test_hit_history_clamps_a_non_positive_page(self, app: Flask, sqlite_session: Session): - """A page below 1 is served as page 1, so the response has to say page 1.""" - api = annotation_module.AnnotationHitHistoryListApi() - handler = unwrap(api.get) - hit_history_mock = Mock(return_value=([_annotation_hit_history()], 9)) - _persist_app(sqlite_session) - with ( - app.test_request_context("/hit-histories?page=0&limit=5", method="GET"), - patch.object( - annotation_module, - "current_account_with_tenant", - return_value=(_account(), "tenant-1"), - ), - patch.object(annotation_module.AppAnnotationService, "get_annotation_hit_histories", hit_history_mock), - ): - response = handler(api, sqlite_session, "app-1", "ann-1") - assert response["page"] == 1 - assert response["has_more"] is True - hit_history_mock.assert_called_once_with( - AnnotationRef(AppRef("tenant-1", "app-1"), "ann-1"), 1, 5, sqlite_session - ) - - -class TestAnnotationListHasMore: - def test_list_annotations_has_more_false_on_last_page_exact_limit( - self, app: Flask, monkeypatch: pytest.MonkeyPatch - ) -> None: - """A full last page must set has_more false instead of forcing another fetch.""" - page_size = 20 - annotation = _annotation_model() - get_mock = Mock(return_value=([annotation] * page_size, page_size)) - monkeypatch.setattr(annotation_module.AppAnnotationService, "get_annotation_list_by_app_id", get_mock) - api = annotation_module.AnnotationApi() - handler = unwrap(api.get) - with app.test_request_context(f"/apps/app-1/annotations?page=1&limit={page_size}", method="GET"): - response, status = handler( - api, annotation_module.AnnotationListQuery(page=1, limit=page_size), Mock(), "app-1" - ) - assert status == 200 - assert response["has_more"] is False - assert response["limit"] == page_size - assert response["total"] == page_size - assert response["page"] == 1 - - def test_list_annotations_has_more_true_when_limit_exceeds_cap( - self, app: Flask, monkeypatch: pytest.MonkeyPatch - ) -> None: - """limit>100 still reports remaining rows after the server cap of 100.""" - annotation = _annotation_model() - get_mock = Mock(return_value=([annotation] * 100, 150)) - monkeypatch.setattr(annotation_module.AppAnnotationService, "get_annotation_list_by_app_id", get_mock) - api = annotation_module.AnnotationApi() - handler = unwrap(api.get) - with app.test_request_context("/apps/app-1/annotations?page=1&limit=200", method="GET"): - response, status = handler(api, annotation_module.AnnotationListQuery(page=1, limit=200), Mock(), "app-1") - assert status == 200 - assert response["has_more"] is True - assert response["limit"] == 100 - assert response["total"] == 150 - assert response["page"] == 1 - - def test_hit_histories_has_more_false_on_last_page_exact_limit(self, app: Flask, sqlite_session: Session): - api = annotation_module.AnnotationHitHistoryListApi() - handler = unwrap(api.get) - history = _annotation_hit_history() - hit_history_mock = Mock(return_value=([history] * 20, 20)) - _persist_app(sqlite_session) - - with ( - app.test_request_context("/hit-histories?page=1&limit=20", method="GET"), - patch.object( - annotation_module, - "current_account_with_tenant", - return_value=(_account(), "tenant-1"), - ), - patch.object(annotation_module.AppAnnotationService, "get_annotation_hit_histories", hit_history_mock), - ): - response = handler(api, sqlite_session, "app-1", "ann-1") - - assert response["has_more"] is False - assert response["limit"] == 20 - assert response["total"] == 20 - assert response["page"] == 1 - - def test_hit_histories_has_more_true_when_limit_exceeds_cap(self, app: Flask, sqlite_session: Session): - api = annotation_module.AnnotationHitHistoryListApi() - handler = unwrap(api.get) - history = _annotation_hit_history() - hit_history_mock = Mock(return_value=([history] * 100, 150)) - _persist_app(sqlite_session) - - with ( - app.test_request_context("/hit-histories?page=1&limit=200", method="GET"), - patch.object( - annotation_module, - "current_account_with_tenant", - return_value=(_account(), "tenant-1"), - ), - patch.object(annotation_module.AppAnnotationService, "get_annotation_hit_histories", hit_history_mock), - ): - response = handler(api, sqlite_session, "app-1", "ann-1") - - assert response["has_more"] is True - assert response["limit"] == 100 - assert response["total"] == 150 - assert response["page"] == 1 diff --git a/api/tests/unit_tests/controllers/console/app/test_annotation_security.py b/api/tests/unit_tests/controllers/console/app/test_annotation_security.py deleted file mode 100644 index 715a90e0f75..00000000000 --- a/api/tests/unit_tests/controllers/console/app/test_annotation_security.py +++ /dev/null @@ -1,360 +0,0 @@ -""" -Unit tests for annotation import security features. - -Tests rate limiting, concurrency control, file validation, and other -security features added to prevent DoS attacks on the annotation import endpoint. -""" - -import io -from unittest.mock import patch - -import pytest -from pandas.errors import ParserError -from sqlalchemy.orm import Session -from werkzeug.datastructures import FileStorage - -from configs import dify_config -from controllers.console.wraps import annotation_import_concurrency_limit, annotation_import_rate_limit -from enums import DeploymentEdition -from models.account import Account -from models.model import App, AppMode, IconType -from services.annotation_service import AppAnnotationService -from tasks.annotation.batch_import_annotations_task import batch_import_annotations_task -from tests.unit_tests.config_override import config_overrides_context -from tests.unit_tests.model_factories import make_account - - -def _account() -> Account: - return make_account(account_id="user_id", name="Annotation Tester", email="annotation-tester@example.com") - - -def _persist_app(session: Session) -> App: - app = App( - id="app_id", - tenant_id="tenant_id", - name="Annotation App", - description="", - mode=AppMode.CHAT, - icon_type=IconType.EMOJI, - icon="chat", - icon_background="#ffffff", - enable_site=False, - enable_api=False, - ) - session.add(app) - session.commit() - return app - - -class TestAnnotationImportRateLimiting: - """Test rate limiting for annotation import operations.""" - - @pytest.fixture - def mock_redis(self): - """Mock Redis client for testing.""" - with patch("controllers.console.wraps.redis_client") as mock: - yield mock - - @pytest.fixture - def mock_current_account(self): - """Mock current account with tenant.""" - with patch("controllers.console.wraps.current_account_with_tenant") as mock: - mock.return_value = (_account(), "test_tenant_id") - yield mock - - def test_rate_limit_per_minute_enforced(self, mock_redis, mock_current_account): - """Test that per-minute rate limit is enforced.""" - # Simulate exceeding per-minute limit - mock_redis.zcard.side_effect = [ - dify_config.ANNOTATION_IMPORT_RATE_LIMIT_PER_MINUTE + 1, # Minute check - 10, # Hour check - ] - - @annotation_import_rate_limit - def dummy_view(): - return "success" - - # Should abort with 429 - with pytest.raises( - Exception, - match=( - "429 Too Many Requests: Too many annotation import requests\\. Maximum " - "5 requests per minute allowed\\. Please try again later\\." - ), - ) as exc_info: - dummy_view() - - # Verify it's a rate limit error - assert "429" in str(exc_info.value) or "Too many" in str(exc_info.value) - - def test_rate_limit_per_hour_enforced(self, mock_redis, mock_current_account): - """Test that per-hour rate limit is enforced.""" - - # Simulate exceeding per-hour limit - mock_redis.zcard.side_effect = [ - 3, # Minute check (under limit) - dify_config.ANNOTATION_IMPORT_RATE_LIMIT_PER_HOUR + 1, # Hour check (over limit) - ] - - @annotation_import_rate_limit - def dummy_view(): - return "success" - - # Should abort with 429 - with pytest.raises( - Exception, - match=( - "429 Too Many Requests: Too many annotation import requests\\. Maximum " - "20 requests per hour allowed\\. Please try again later\\." - ), - ) as exc_info: - dummy_view() - - assert "429" in str(exc_info.value) or "Too many" in str(exc_info.value) - - def test_rate_limit_within_limits_passes(self, mock_redis, mock_current_account): - """Test that requests within limits are allowed.""" - - # Simulate being under both limits - mock_redis.zcard.return_value = 2 - - @annotation_import_rate_limit - def dummy_view(): - return "success" - - # Should succeed - result = dummy_view() - assert result == "success" - - # Verify Redis operations were called - assert mock_redis.zadd.called - assert mock_redis.zremrangebyscore.called - - -class TestAnnotationImportConcurrencyControl: - """Test concurrency control for annotation import operations.""" - - @pytest.fixture - def mock_redis(self): - """Mock Redis client for testing.""" - with patch("controllers.console.wraps.redis_client") as mock: - yield mock - - @pytest.fixture - def mock_current_account(self): - """Mock current account with tenant.""" - with patch("controllers.console.wraps.current_account_with_tenant") as mock: - mock.return_value = (_account(), "test_tenant_id") - yield mock - - def test_concurrency_limit_enforced(self, mock_redis, mock_current_account): - """Test that concurrent task limit is enforced.""" - - # Simulate max concurrent tasks already running - mock_redis.zcard.return_value = dify_config.ANNOTATION_IMPORT_MAX_CONCURRENT - - @annotation_import_concurrency_limit - def dummy_view(): - return "success" - - # Should abort with 429 - with pytest.raises( - Exception, - match=( - "429 Too Many Requests: Too many concurrent import tasks\\. Maximum 2 " - "concurrent imports allowed per workspace\\. Please wait for existing " - "imports to complete\\." - ), - ) as exc_info: - dummy_view() - - assert "429" in str(exc_info.value) or "concurrent" in str(exc_info.value).lower() - - def test_concurrency_within_limit_passes(self, mock_redis, mock_current_account): - """Test that requests within concurrency limits are allowed.""" - - # Simulate being under concurrent task limit - mock_redis.zcard.return_value = 1 - - @annotation_import_concurrency_limit - def dummy_view(): - return "success" - - # Should succeed - result = dummy_view() - assert result == "success" - - def test_stale_jobs_are_cleaned_up(self, mock_redis, mock_current_account): - """Test that old/stale job entries are removed.""" - - mock_redis.zcard.return_value = 0 - - @annotation_import_concurrency_limit - def dummy_view(): - return "success" - - dummy_view() - - # Verify cleanup was called - assert mock_redis.zremrangebyscore.called - - -class TestAnnotationImportFileValidation: - """Test file validation in annotation import.""" - - def test_file_size_limit_enforced(self): - """Test that files exceeding size limit are rejected.""" - # Create a file larger than the limit - max_size = dify_config.ANNOTATION_IMPORT_FILE_SIZE_LIMIT * 1024 * 1024 - large_content = b"x" * (max_size + 1024) # Exceed by 1KB - - FileStorage(stream=io.BytesIO(large_content), filename="test.csv", content_type="text/csv") - - # Should be rejected in controller - # This would be tested in integration tests with actual endpoint - - def test_empty_file_rejected(self): - """Test that empty files are rejected.""" - FileStorage(stream=io.BytesIO(b""), filename="test.csv", content_type="text/csv") - - # Should be rejected - # This would be tested in integration tests - - def test_non_csv_file_rejected(self): - """Test that non-CSV files are rejected.""" - FileStorage(stream=io.BytesIO(b"test"), filename="test.txt", content_type="text/plain") - - # Should be rejected based on extension - # This would be tested in integration tests - - -@pytest.mark.parametrize("sqlite_session", [(App,)], indirect=True) -class TestAnnotationImportServiceValidation: - """Test service layer validation for annotation import.""" - - @config_overrides_context(DEPLOYMENT_EDITION=DeploymentEdition.COMMUNITY) - def test_max_records_limit_enforced(self, sqlite_session: Session): - """Test that files with too many records are rejected.""" - - # Create CSV with too many records - max_records = dify_config.ANNOTATION_IMPORT_MAX_RECORDS - csv_content = "question,answer\n" - for i in range(max_records + 100): - csv_content += f"Question {i},Answer {i}\n" - - file = FileStorage(stream=io.BytesIO(csv_content.encode()), filename="test.csv", content_type="text/csv") - - _persist_app(sqlite_session) - with patch("services.annotation_service.current_account_with_tenant") as mock_auth: - mock_auth.return_value = (_account(), "tenant_id") - - with patch("services.annotation_service.FeatureService"): - result = AppAnnotationService.batch_import_app_annotations("app_id", file, sqlite_session) - - # Should return error about too many records - assert "error_msg" in result - assert "too many" in result["error_msg"].lower() or "maximum" in result["error_msg"].lower() - - def test_min_records_limit_enforced(self, sqlite_session: Session): - """Test that files with too few valid records are rejected.""" - - # Create CSV with only header (no data rows) - csv_content = "question,answer\n" - - file = FileStorage(stream=io.BytesIO(csv_content.encode()), filename="test.csv", content_type="text/csv") - - _persist_app(sqlite_session) - with patch("services.annotation_service.current_account_with_tenant") as mock_auth: - mock_auth.return_value = (_account(), "tenant_id") - - result = AppAnnotationService.batch_import_app_annotations("app_id", file, sqlite_session) - - # Should return error about insufficient records - assert "error_msg" in result - assert "at least" in result["error_msg"].lower() or "minimum" in result["error_msg"].lower() - - def test_invalid_csv_format_handled(self, sqlite_session: Session): - """Test that invalid CSV format is handled gracefully.""" - - # Any content is fine once we force ParserError - csv_content = 'invalid,csv,format\nwith,unbalanced,quotes,and"stuff' - file = FileStorage(stream=io.BytesIO(csv_content.encode()), filename="test.csv", content_type="text/csv") - - _persist_app(sqlite_session) - with ( - patch("services.annotation_service.current_account_with_tenant") as mock_auth, - patch("services.annotation_service.pd.read_csv", side_effect=ParserError("malformed CSV")), - ): - mock_auth.return_value = (_account(), "tenant_id") - - result = AppAnnotationService.batch_import_app_annotations("app_id", file, sqlite_session) - - assert "error_msg" in result - assert "malformed" in result["error_msg"].lower() - - @config_overrides_context(DEPLOYMENT_EDITION=DeploymentEdition.COMMUNITY) - def test_valid_import_succeeds(self, sqlite_session: Session): - """Test that valid import request succeeds.""" - - # Create valid CSV - csv_content = "question,answer\nWhat is AI?,Artificial Intelligence\nWhat is ML?,Machine Learning\n" - - file = FileStorage(stream=io.BytesIO(csv_content.encode()), filename="test.csv", content_type="text/csv") - - _persist_app(sqlite_session) - with patch("services.annotation_service.current_account_with_tenant") as mock_auth: - mock_auth.return_value = (_account(), "tenant_id") - - with patch("services.annotation_service.FeatureService"): - with patch("services.annotation_service.batch_import_annotations_task"): - with patch("services.annotation_service.redis_client"): - result = AppAnnotationService.batch_import_app_annotations("app_id", file, sqlite_session) - - # Should return success response - assert "job_id" in result - assert "job_status" in result - assert result["job_status"] == "waiting" - assert "record_count" in result - assert result["record_count"] == 2 - - -class TestAnnotationImportTaskOptimization: - """Test optimizations in batch import task.""" - - def test_task_is_registered_with_queue(self): - """Test that task is registered with the correct queue.""" - assert hasattr(batch_import_annotations_task, "apply_async") - assert hasattr(batch_import_annotations_task, "delay") - - -class TestConfigurationValues: - """Test that security configuration values are properly set.""" - - def test_rate_limit_configs_exist(self): - """Test that rate limit configurations are defined.""" - assert hasattr(dify_config, "ANNOTATION_IMPORT_RATE_LIMIT_PER_MINUTE") - assert hasattr(dify_config, "ANNOTATION_IMPORT_RATE_LIMIT_PER_HOUR") - - assert dify_config.ANNOTATION_IMPORT_RATE_LIMIT_PER_MINUTE > 0 - assert dify_config.ANNOTATION_IMPORT_RATE_LIMIT_PER_HOUR > 0 - - def test_file_size_limit_config_exists(self): - """Test that file size limit configuration is defined.""" - assert hasattr(dify_config, "ANNOTATION_IMPORT_FILE_SIZE_LIMIT") - assert dify_config.ANNOTATION_IMPORT_FILE_SIZE_LIMIT > 0 - assert dify_config.ANNOTATION_IMPORT_FILE_SIZE_LIMIT <= 10 # Reasonable max (10MB) - - def test_record_limit_configs_exist(self): - """Test that record limit configurations are defined.""" - assert hasattr(dify_config, "ANNOTATION_IMPORT_MAX_RECORDS") - assert hasattr(dify_config, "ANNOTATION_IMPORT_MIN_RECORDS") - - assert dify_config.ANNOTATION_IMPORT_MAX_RECORDS > 0 - assert dify_config.ANNOTATION_IMPORT_MIN_RECORDS > 0 - assert dify_config.ANNOTATION_IMPORT_MIN_RECORDS < dify_config.ANNOTATION_IMPORT_MAX_RECORDS - - def test_concurrency_limit_config_exists(self): - """Test that concurrency limit configuration is defined.""" - assert hasattr(dify_config, "ANNOTATION_IMPORT_MAX_CONCURRENT") - assert dify_config.ANNOTATION_IMPORT_MAX_CONCURRENT > 0 - assert dify_config.ANNOTATION_IMPORT_MAX_CONCURRENT <= 10 # Reasonable upper bound diff --git a/api/tests/unit_tests/controllers/console/app/test_app.py b/api/tests/unit_tests/controllers/console/app/test_app.py new file mode 100644 index 00000000000..ee34f58eed1 --- /dev/null +++ b/api/tests/unit_tests/controllers/console/app/test_app.py @@ -0,0 +1,1812 @@ +from __future__ import annotations + +import inspect +import json +import uuid +from collections.abc import Callable, Iterator +from dataclasses import dataclass, replace +from dataclasses import field as app_field +from datetime import datetime +from inspect import unwrap +from io import BytesIO +from pathlib import Path +from types import SimpleNamespace +from typing import cast +from unittest.mock import MagicMock, patch +from uuid import UUID + +import pytest +import yaml +from flask import Flask, Response +from flask_restx import Resource +from pydantic import ValidationError +from sqlalchemy import Engine +from sqlalchemy.orm import Session +from werkzeug.datastructures import MultiDict +from werkzeug.exceptions import BadRequest, Forbidden, NotFound + +from controllers.console import flask_admission +from controllers.console.app import ( + annotation as annotation_module, +) +from controllers.console.app import ( + app as app_module, +) +from controllers.console.app import app as controller +from controllers.console.app import ( + completion as completion_module, +) +from controllers.console.app import ( + message as message_module, +) +from controllers.console.app import ( + ops_trace as ops_trace_module, +) +from controllers.console.app import ( + site as site_module, +) +from controllers.console.app import ( + statistic as statistic_module, +) +from controllers.console.app import ( + workflow_app_log as workflow_app_log_module, +) +from controllers.console.app import ( + workflow_draft_variable as workflow_draft_variable_module, +) +from controllers.console.app import ( + workflow_statistic as workflow_statistic_module, +) +from controllers.console.app import ( + wraps as wraps_module, +) +from controllers.console.app.app import CreateAppPayload +from controllers.console.app.completion import ChatMessagePayload, CompletionMessagePayload +from controllers.console.app.error import ( + AppNotFoundError, + TracingConfigNotFoundError, + TracingConfigVerificationFailedError, +) +from controllers.console.app.mcp_server import MCPServerCreatePayload, MCPServerUpdatePayload +from controllers.console.app.ops_trace import TraceConfigPayload, TraceProviderQuery +from controllers.console.app.site import AppSiteUpdatePayload +from controllers.console.app.workflow import AdvancedChatWorkflowRunPayload, SyncDraftWorkflowPayload +from controllers.console.app.workflow_app_log import WorkflowAppLogQuery +from controllers.console.app.workflow_draft_variable import ( + EnvironmentVariableUpdatePayload, + WorkflowDraftVariableListQuery, + WorkflowDraftVariableUpdatePayload, +) +from controllers.console.app.workflow_statistic import WorkflowStatisticQuery +from fields import app_fields +from fields.app_model_config_response import AppModelConfigResponse +from libs.login import AccountWithTenant +from libs.pagination import PaginatedResult +from machinery.context import RequestContext +from models import App as _app_apis_App +from models import Site as _app_apis_Site +from models.account import Account, AccountStatus, TenantAccountRole +from models.engine import db +from models.enums import CustomizeTokenStrategy, TagType +from models.model import App as _app_response_models_App +from models.model import AppMode, AppModelConfig, IconType, Tag, TagBinding +from models.model import Site as _app_response_models_Site +from models.workflow import Workflow, WorkflowType +from repositories.app.response import app_record +from services.agent.roster_package_entities import RosterAgentPackageExport +from services.app.console_service import ( + AppExportPaidPlanRequiredError, + ConsoleAppNotFoundError, + CreatorsPlatformDisabledError, + InvalidAppExportError, +) +from services.app_service import AppResponseView +from services.app_site_service import ( + AppSiteAppNotFoundError, + AppSiteChanges, + AppSiteCommandResult, + AppSiteNotFoundError, +) +from services.app_tracing_config_service import ( + AppTraceSettings, + AppTracingConfigNotFoundError, + AppTracingConfigVerificationFailedError, +) +from services.entities.app_entities import ( + AppExportOptions, + AppListParams, + AppPage, + AppRecord, + CopyAppParams, + CreateAppParams, + StarredAppListParams, + UpdateAppParams, +) +from services.entities.dsl_entities import Import, ImportStatus +from tests.helpers.unit.config_override import apply_config_overrides + +"""Unit coverage for console app controller contracts and response mapping.""" + + +_APP_APIS_APP_ID = "11111111-1111-1111-1111-111111111111" +_APP_APIS_TENANT_ID = "22222222-2222-2222-2222-222222222222" +USER_ID = "33333333-3333-3333-3333-333333333333" + + +def _make_account() -> Account: + account = Account( + name="tester", + email="tester@example.com", + status=AccountStatus.ACTIVE, + ) + account.id = USER_ID + return account + + +def _make_app( + app_id: str = _APP_APIS_APP_ID, + *, + tenant_id: str = _APP_APIS_TENANT_ID, + icon_type: app_module.IconType | None = None, +) -> _app_apis_App: + app = _app_apis_App() + app.id = app_id + app.tenant_id = tenant_id + app.icon_type = icon_type + return app + + +def _make_request_context() -> RequestContext: + return RequestContext( + request_id="request-1", + trace_id="trace-1", + account_id=USER_ID, + active_workspace_id=_APP_APIS_TENANT_ID, + ) + + +class TestAppApisCases: + @staticmethod + @pytest.fixture + def database_app() -> Iterator[Flask]: + app = Flask(__name__) + app.config["SQLALCHEMY_DATABASE_URI"] = "sqlite:///:memory:" + db.init_app(app) + + with app.app_context(): + _app_apis_Site.__table__.create(db.engine) + yield app + + class TestCompletionEndpoints: + def test_completion_create_payload(self): + payload = CompletionMessagePayload(inputs={"prompt": "test"}, model_config={}) + assert payload.inputs == {"prompt": "test"} + + def test_chat_message_payload_uuid_validation(self): + payload = ChatMessagePayload( + inputs={}, + model_config={}, + query="hi", + conversation_id=str(uuid.uuid4()), + parent_message_id=str(uuid.uuid4()), + ) + assert payload.query == "hi" + + def test_completion_api_success( + self, app: Flask, monkeypatch: pytest.MonkeyPatch, sqlite_engine: Engine + ) -> None: + api = completion_module.CompletionMessageApi() + method = unwrap(api.post) + + monkeypatch.setattr( + completion_module.AppGenerateService, + "generate", + lambda **_kwargs: {"text": "ok"}, + ) + monkeypatch.setattr( + completion_module.helper, + "compact_generate_response", + lambda response: {"result": response}, + ) + + with ( + Session(sqlite_engine) as session, + app.test_request_context("/", json={"inputs": {}, "model_config": {}, "query": "hi"}), + ): + resp = method( + api, + CompletionMessagePayload(inputs={}, model_config={}, query="hi"), + session, + _make_account(), + app_model=_make_app(), + ) + + assert resp == {"result": {"text": "ok"}} + + def test_completion_api_conversation_not_exists( + self, app: Flask, monkeypatch: pytest.MonkeyPatch, sqlite_engine: Engine + ) -> None: + api = completion_module.CompletionMessageApi() + method = unwrap(api.post) + + monkeypatch.setattr( + completion_module.AppGenerateService, + "generate", + lambda **_kwargs: (_ for _ in ()).throw( + completion_module.services.errors.conversation.ConversationNotExistsError() + ), + ) + + with ( + Session(sqlite_engine) as session, + app.test_request_context("/", json={"inputs": {}, "model_config": {}, "query": "hi"}), + pytest.raises(NotFound), + ): + method( + api, + CompletionMessagePayload(inputs={}, model_config={}, query="hi"), + session, + _make_account(), + app_model=_make_app(), + ) + + def test_completion_api_provider_not_initialized( + self, app: Flask, monkeypatch: pytest.MonkeyPatch, sqlite_engine: Engine + ) -> None: + api = completion_module.CompletionMessageApi() + method = unwrap(api.post) + + monkeypatch.setattr( + completion_module.AppGenerateService, + "generate", + lambda **_kwargs: (_ for _ in ()).throw(completion_module.ProviderTokenNotInitError("x")), + ) + + with ( + Session(sqlite_engine) as session, + app.test_request_context("/", json={"inputs": {}, "model_config": {}, "query": "hi"}), + pytest.raises(completion_module.ProviderNotInitializeError), + ): + method( + api, + CompletionMessagePayload(inputs={}, model_config={}, query="hi"), + session, + _make_account(), + app_model=_make_app(), + ) + + def test_completion_api_quota_exceeded( + self, app: Flask, monkeypatch: pytest.MonkeyPatch, sqlite_engine: Engine + ) -> None: + api = completion_module.CompletionMessageApi() + method = unwrap(api.post) + + monkeypatch.setattr( + completion_module.AppGenerateService, + "generate", + lambda **_kwargs: (_ for _ in ()).throw(completion_module.QuotaExceededError()), + ) + + with ( + Session(sqlite_engine) as session, + app.test_request_context("/", json={"inputs": {}, "model_config": {}, "query": "hi"}), + pytest.raises(completion_module.ProviderQuotaExceededError), + ): + method( + api, + CompletionMessagePayload(inputs={}, model_config={}, query="hi"), + session, + _make_account(), + app_model=_make_app(), + ) + + class TestOpsTraceEndpoints: + def test_ops_trace_query_basic(self): + query = TraceProviderQuery(tracing_provider="langfuse") + assert query.tracing_provider == "langfuse" + + def test_ops_trace_config_payload(self): + payload = TraceConfigPayload(tracing_provider="langfuse", tracing_config={"api_key": "k"}) + assert payload.tracing_config["api_key"] == "k" + + def test_trace_app_config_get_empty(self, app: Flask, monkeypatch: pytest.MonkeyPatch): + api = ops_trace_module.TraceAppConfigApi() + method = unwrap(api.get) + tracing_configs = MagicMock() + tracing_configs.get.return_value = None + monkeypatch.setattr( + ops_trace_module, + "application_services", + lambda: SimpleNamespace(app_tracing_configs=tracing_configs), + ) + + with app.test_request_context("/?tracing_provider=langfuse"): + result = method( + api, + TraceProviderQuery(tracing_provider="langfuse"), + _make_request_context(), + uuid.UUID(_APP_APIS_APP_ID), + ) + + assert result == {"has_not_configured": True} + tracing_configs.get.assert_called_once_with( + context=_make_request_context(), + app_id=_APP_APIS_APP_ID, + tracing_provider="langfuse", + ) + + def test_trace_app_config_post_invalid(self, app: Flask, monkeypatch: pytest.MonkeyPatch): + api = ops_trace_module.TraceAppConfigApi() + method = unwrap(api.post) + tracing_configs = MagicMock() + tracing_configs.create.side_effect = AppTracingConfigVerificationFailedError() + monkeypatch.setattr( + ops_trace_module, + "application_services", + lambda: SimpleNamespace(app_tracing_configs=tracing_configs), + ) + + with app.test_request_context( + "/", + json={"tracing_provider": "langfuse", "tracing_config": {"api_key": "k"}}, + ): + with pytest.raises(TracingConfigVerificationFailedError): + method( + api, + TraceConfigPayload(tracing_provider="langfuse", tracing_config={"api_key": "k"}), + _make_request_context(), + uuid.UUID(_APP_APIS_APP_ID), + ) + + def test_trace_app_config_delete_not_found(self, app: Flask, monkeypatch: pytest.MonkeyPatch): + api = ops_trace_module.TraceAppConfigApi() + method = unwrap(api.delete) + tracing_configs = MagicMock() + tracing_configs.delete.side_effect = AppTracingConfigNotFoundError() + monkeypatch.setattr( + ops_trace_module, + "application_services", + lambda: SimpleNamespace(app_tracing_configs=tracing_configs), + ) + + with app.test_request_context("/?tracing_provider=langfuse"): + with pytest.raises(TracingConfigNotFoundError): + method( + api, + TraceProviderQuery(tracing_provider="langfuse"), + _make_request_context(), + uuid.UUID(_APP_APIS_APP_ID), + ) + + class TestSiteEndpoints: + @staticmethod + def _command_result(*, code: str = "test-code") -> AppSiteCommandResult: + return AppSiteCommandResult( + app_id=_APP_APIS_APP_ID, + title="My Site", + description="Test site", + default_language="en-US", + input_placeholder="Ask me anything", + code=code, + icon=None, + icon_background=None, + customize_domain=None, + copyright=None, + privacy_policy=None, + custom_disclaimer="", + customize_token_strategy="not_allow", + prompt_public=False, + show_workflow_steps=True, + use_icon_as_answer_icon=False, + ) + + def test_site_payload_maps_to_application_changes(self): + payload = AppSiteUpdatePayload( + title="My Site", + description="Test site", + input_placeholder="Ask me anything", + ) + assert payload.to_changes() == AppSiteChanges( + title="My Site", + description="Test site", + input_placeholder="Ask me anything", + ) + + def test_site_default_language_validation(self): + payload = AppSiteUpdatePayload(default_language="en-US") + assert payload.default_language == "en-US" + + def test_app_site_update_post( + self, + ) -> None: + api = site_module.AppSite() + method = unwrap(api.post) + services = MagicMock() + services.app_sites.update.return_value = self._command_result() + context = RequestContext("request-1", None, USER_ID, _APP_APIS_TENANT_ID) + + with patch.object(site_module, "application_services", return_value=services): + result = method( + api, + AppSiteUpdatePayload(title="My Site", input_placeholder="Ask me anything"), + context, + app_id=uuid.UUID(_APP_APIS_APP_ID), + ) + + assert isinstance(result, dict) + assert result["title"] == "My Site" + assert result["input_placeholder"] == "Ask me anything" + assert result["access_token"] == "test-code" + assert result["code"] == "test-code" + services.app_sites.update.assert_called_once_with( + context, + _APP_APIS_APP_ID, + AppSiteChanges(title="My Site", input_placeholder="Ask me anything"), + ) + + def test_app_site_access_token_reset( + self, + ) -> None: + api = site_module.AppSiteAccessTokenReset() + method = unwrap(api.post) + services = MagicMock() + services.app_sites.reset_access_token.return_value = self._command_result(code="new-code") + context = RequestContext("request-1", None, USER_ID, _APP_APIS_TENANT_ID) + + with patch.object(site_module, "application_services", return_value=services): + result = method(api, context, app_id=uuid.UUID(_APP_APIS_APP_ID)) + + assert isinstance(result, dict) + assert result["access_token"] == "new-code" + assert result["code"] == "new-code" + services.app_sites.reset_access_token.assert_called_once_with(context, _APP_APIS_APP_ID) + + def test_app_site_update_maps_missing_app(self) -> None: + api = site_module.AppSite() + method = unwrap(api.post) + services = MagicMock() + services.app_sites.update.side_effect = AppSiteAppNotFoundError() + context = RequestContext("request-1", None, USER_ID, _APP_APIS_TENANT_ID) + + with ( + patch.object(site_module, "application_services", return_value=services), + pytest.raises(AppNotFoundError), + ): + method(api, AppSiteUpdatePayload(), context, app_id=uuid.UUID(_APP_APIS_APP_ID)) + + def test_app_site_reset_maps_missing_site(self) -> None: + api = site_module.AppSiteAccessTokenReset() + method = unwrap(api.post) + services = MagicMock() + services.app_sites.reset_access_token.side_effect = AppSiteNotFoundError() + context = RequestContext("request-1", None, USER_ID, _APP_APIS_TENANT_ID) + + with ( + patch.object(site_module, "application_services", return_value=services), + pytest.raises(NotFound), + ): + method(api, context, app_id=uuid.UUID(_APP_APIS_APP_ID)) + + class TestWorkflowEndpoints: + def test_workflow_copy_payload(self): + payload = SyncDraftWorkflowPayload(graph={}, features={}) + assert payload.graph == {} + assert payload.model_dump()["is_collaborative"] is False + + def test_workflow_sync_payload_accepts_collaboration_marker(self): + payload = SyncDraftWorkflowPayload.model_validate({"graph": {}, "features": {}, "_is_collaborative": True}) + assert payload.is_collaborative is True + assert payload.model_dump()["is_collaborative"] is True + + def test_workflow_mode_query(self): + payload = AdvancedChatWorkflowRunPayload(inputs={}, query="hi") + assert payload.query == "hi" + + class TestWorkflowAppLogEndpoints: + def test_workflow_app_log_query(self): + query = WorkflowAppLogQuery(keyword="test", page=1, limit=20) + assert query.keyword == "test" + + def test_workflow_app_log_query_detail_bool(self): + query = WorkflowAppLogQuery(detail="true") + assert query.detail is True + + def test_workflow_app_log_api_get(self, database_app: Flask, monkeypatch: pytest.MonkeyPatch) -> None: + api = workflow_app_log_module.WorkflowAppLogApi() + method = unwrap(api.get) + workflow_app_logs = MagicMock() + workflow_app_logs.list_logs.return_value = { + "page": 1, + "limit": 20, + "total": 0, + "has_more": False, + "data": [], + } + services = MagicMock(workflow_app_logs=workflow_app_logs) + monkeypatch.setattr(workflow_app_log_module, "application_services", lambda: services) + context = RequestContext("request-1", None, USER_ID, _APP_APIS_TENANT_ID) + app_model = _make_app("app-1") + + with database_app.test_request_context("/?page=1&limit=20"): + result = method(api, WorkflowAppLogQuery(page=1, limit=20), context, app_model=app_model) + + assert result == {"page": 1, "limit": 20, "total": 0, "has_more": False, "data": []} + workflow_app_logs.list_logs.assert_called_once_with( + tenant_id=app_model.tenant_id, + app_id=app_model.id, + keyword=None, + status=None, + created_at_before=None, + created_at_after=None, + page=1, + limit=20, + detail=False, + created_by_end_user_session_id=None, + created_by_account=None, + ) + + class TestWorkflowDraftVariableEndpoints: + def test_workflow_variable_creation(self): + payload = WorkflowDraftVariableUpdatePayload(name="var1", value="test") + assert payload.name == "var1" + + def test_workflow_variable_collection_get(self, database_app: Flask, monkeypatch: pytest.MonkeyPatch) -> None: + api = workflow_draft_variable_module.WorkflowVariableCollectionApi() + method = unwrap(api.get) + + class DummyDraftService: + def __init__(self, session: Session): + self.session = session + + def list_variables_without_values(self, **_kwargs): + assert self.session.get_bind() is db.engine + return {"items": [], "total": 0} + + class DummyWorkflowService: + def is_workflow_exist(self, *args, **kwargs): + return True + + monkeypatch.setattr(workflow_draft_variable_module, "WorkflowDraftVariableService", DummyDraftService) + monkeypatch.setattr(workflow_draft_variable_module, "WorkflowService", DummyWorkflowService) + + with database_app.test_request_context("/?page=1&limit=20"): + result = method( + api, + WorkflowDraftVariableListQuery(page=1, limit=20), + _make_account(), + app_model=_make_app("app-1"), + ) + + assert result == {"items": [], "total": 0} + + def test_environment_variable_update_payload_preserves_full_replace_default(self) -> None: + payload = EnvironmentVariableUpdatePayload(environment_variables=[]) + + assert payload.patch is False + assert payload.deleted_environment_variable_ids == [] + + @pytest.mark.parametrize( + "payload", + [ + {"environment_variables": [], "deleted_environment_variable_ids": ["env-a"]}, + { + "environment_variables": [{"name": "a", "value_type": "string", "value": "a"}], + "patch": True, + }, + { + "environment_variables": [{"id": "env-a", "name": "a", "value_type": "string", "value": "a"}], + "patch": True, + "deleted_environment_variable_ids": ["env-a"], + }, + ], + ) + def test_environment_variable_patch_payload_rejects_ambiguous_mutations(self, payload: dict) -> None: + with pytest.raises(ValidationError): + EnvironmentVariableUpdatePayload.model_validate(payload) + + def test_environment_variable_collection_post_routes_patch_to_service( + self, database_app: Flask, monkeypatch: pytest.MonkeyPatch + ) -> None: + api = workflow_draft_variable_module.EnvironmentVariableCollectionApi() + method = unwrap(api.post) + captured: dict = {} + + class DummyWorkflowService: + def patch_draft_workflow_environment_variables(self, **kwargs) -> None: + captured.update(kwargs) + + def update_draft_workflow_environment_variables(self, **_kwargs) -> None: + raise AssertionError("patch request must not use full replacement") + + monkeypatch.setattr(workflow_draft_variable_module, "WorkflowService", DummyWorkflowService) + + with database_app.test_request_context( + "/", + json={ + "environment_variables": [{"id": "env-a", "name": "a", "value_type": "string", "value": "new-a"}], + "patch": True, + "deleted_environment_variable_ids": ["env-b"], + }, + ): + result = method( + api, + EnvironmentVariableUpdatePayload( + environment_variables=[{"id": "env-a", "name": "a", "value_type": "string", "value": "new-a"}], + patch=True, + deleted_environment_variable_ids=["env-b"], + ), + _make_account(), + app_model=_make_app(), + ) + + assert result == {"result": "success"} + assert [(variable.id, variable.value) for variable in captured["environment_variables"]] == [ + ("env-a", "new-a") + ] + assert captured["deleted_environment_variable_ids"] == ["env-b"] + assert captured["app_model"].id == _APP_APIS_APP_ID + assert captured["account"].id == USER_ID + assert captured["session"].get_bind() is db.engine + + class TestWorkflowStatisticEndpoints: + def test_workflow_statistic_time_range(self): + query = WorkflowStatisticQuery(start="2024-01-01", end="2024-12-31") + assert query.start == "2024-01-01" + + def test_workflow_statistic_blank_to_none(self): + query = WorkflowStatisticQuery(start="", end="") + assert query.start is None + assert query.end is None + + def test_workflow_daily_runs_statistic(self, database_app: Flask, monkeypatch: pytest.MonkeyPatch) -> None: + service = MagicMock() + service.get_daily_runs.return_value = [{"date": "2024-01-01", "runs": 2}] + monkeypatch.setattr( + workflow_statistic_module, + "application_services", + lambda: SimpleNamespace(workflow_statistics=service), + ) + self._patch_statistic_time_range(monkeypatch) + + api = workflow_statistic_module.WorkflowDailyRunsStatistic() + method = unwrap(api.get) + request_context = _make_request_context() + + with database_app.test_request_context("/"): + response = method( + api, + WorkflowStatisticQuery(), + request_context, + app_model=_make_app("app-1", tenant_id="t1"), + ) + + assert response == {"data": [{"date": "2024-01-01", "runs": 2}]} + service.get_daily_runs.assert_called_once_with( + request_context, + app_id="app-1", + start_date=None, + end_date=None, + timezone="UTC", + ) + + def test_workflow_daily_terminals_statistic(self, database_app: Flask, monkeypatch: pytest.MonkeyPatch) -> None: + service = MagicMock() + service.get_daily_terminals.return_value = [{"date": "2024-01-02", "terminal_count": 3}] + monkeypatch.setattr( + workflow_statistic_module, + "application_services", + lambda: SimpleNamespace(workflow_statistics=service), + ) + self._patch_statistic_time_range(monkeypatch) + + api = workflow_statistic_module.WorkflowDailyTerminalsStatistic() + method = unwrap(api.get) + request_context = _make_request_context() + + with database_app.test_request_context("/"): + response = method( + api, + WorkflowStatisticQuery(), + request_context, + app_model=_make_app("app-1", tenant_id="t1"), + ) + + assert response == {"data": [{"date": "2024-01-02", "terminal_count": 3}]} + service.get_daily_terminals.assert_called_once_with( + request_context, + app_id="app-1", + start_date=None, + end_date=None, + timezone="UTC", + ) + + def test_workflow_daily_token_cost_statistic( + self, + database_app: Flask, + monkeypatch: pytest.MonkeyPatch, + ) -> None: + service = MagicMock() + service.get_daily_token_costs.return_value = [{"date": "2024-01-03", "token_count": 4}] + monkeypatch.setattr( + workflow_statistic_module, + "application_services", + lambda: SimpleNamespace(workflow_statistics=service), + ) + self._patch_statistic_time_range(monkeypatch) + + api = workflow_statistic_module.WorkflowDailyTokenCostStatistic() + method = unwrap(api.get) + request_context = _make_request_context() + + with database_app.test_request_context("/"): + response = method( + api, + WorkflowStatisticQuery(), + request_context, + app_model=_make_app("app-1", tenant_id="t1"), + ) + + assert response == {"data": [{"date": "2024-01-03", "token_count": 4}]} + service.get_daily_token_costs.assert_called_once_with( + request_context, + app_id="app-1", + start_date=None, + end_date=None, + timezone="UTC", + ) + + def test_workflow_average_app_interaction_statistic( + self, + database_app: Flask, + monkeypatch: pytest.MonkeyPatch, + ) -> None: + service = MagicMock() + service.get_average_app_interactions.return_value = [{"date": "2024-01-04", "interactions": 2.5}] + monkeypatch.setattr( + workflow_statistic_module, + "application_services", + lambda: SimpleNamespace(workflow_statistics=service), + ) + self._patch_statistic_time_range(monkeypatch) + + api = workflow_statistic_module.WorkflowAverageAppInteractionStatistic() + method = unwrap(api.get) + request_context = _make_request_context() + + with database_app.test_request_context("/"): + response = method( + api, + WorkflowStatisticQuery(), + request_context, + app_model=_make_app("app-1", tenant_id="t1"), + ) + + assert response == {"data": [{"date": "2024-01-04", "interactions": 2.5}]} + service.get_average_app_interactions.assert_called_once_with( + request_context, + app_id="app-1", + start_date=None, + end_date=None, + timezone="UTC", + ) + + def test_workflow_statistic_invalid_time_range(self, monkeypatch: pytest.MonkeyPatch) -> None: + account = _make_account() + account.timezone = "UTC" + monkeypatch.setattr( + workflow_statistic_module, + "current_account_with_tenant", + lambda: SimpleNamespace(account=account), + ) + monkeypatch.setattr( + workflow_statistic_module, + "parse_time_range", + MagicMock(side_effect=ValueError("invalid range")), + ) + + with pytest.raises(BadRequest, match="invalid range"): + workflow_statistic_module._resolve_statistic_time_range(WorkflowStatisticQuery()) + + @staticmethod + def _patch_statistic_time_range(monkeypatch: pytest.MonkeyPatch) -> None: + account = _make_account() + account.timezone = "UTC" + monkeypatch.setattr( + workflow_statistic_module, + "current_account_with_tenant", + lambda: SimpleNamespace(account=account), + ) + monkeypatch.setattr( + workflow_statistic_module, + "parse_time_range", + lambda *_args, **_kwargs: (None, None), + ) + + class TestWrapsEndpoints: + def test_get_app_model_context(self): + assert hasattr(wraps_module, "get_app_model") + + class TestMCPServerEndpoints: + def test_mcp_server_connection(self): + payload = MCPServerCreatePayload(parameters={"url": "http://localhost:3000"}) + assert payload.parameters["url"] == "http://localhost:3000" + + def test_mcp_server_update_payload(self): + payload = MCPServerUpdatePayload( + id="abcdefab-1234-4567-89ab-abcdefabcdef", parameters={"timeout": 30}, status="active" + ) + assert payload.status == "active" + + class TestErrorHandling: + def test_annotation_list_query_validation(self): + with pytest.raises(ValueError, match="1 validation error for AnnotationListQuery\\\npage"): + annotation_module.AnnotationListQuery(page=0) + + class TestPayloadIntegration: + def test_multiple_payload_types(self): + payloads = [ + annotation_module.AnnotationReplyPayload( + score_threshold=0.5, embedding_provider_name="openai", embedding_model_name="text-embedding-3-small" + ), + message_module.AppMessageFeedbackPayload(message_id=str(uuid.uuid4()), rating="like"), + statistic_module.StatisticTimeRangeQuery(start="2024-01-01"), + ] + assert len(payloads) == 3 + assert all(p is not None for p in payloads) + + +"""App request contracts and serialization of materialized response data.""" + + +def _ts(hour: int = 12) -> datetime: + return datetime(2024, 1, 1, hour, 0, 0) + + +_APP_RESPONSE_MODELS_TENANT_ID = "00000000-0000-0000-0000-000000000001" +_APP_RESPONSE_MODELS_APP_ID = "00000000-0000-0000-0000-000000000101" +ACCOUNT_ID = "00000000-0000-0000-0000-000000000201" +CONFIG_ID = "00000000-0000-0000-0000-000000000301" +WORKFLOW_ID = "00000000-0000-0000-0000-000000000401" +SITE_ID = "00000000-0000-0000-0000-000000000501" +TAG_ID = "00000000-0000-0000-0000-000000000601" + + +def _account(*, account_id: str = ACCOUNT_ID) -> Account: + account = Account(name="Creator", email=f"{account_id}@example.com") + account.id = account_id + return account + + +def _app( + *, + app_id: str = _APP_RESPONSE_MODELS_APP_ID, + tenant_id: str = _APP_RESPONSE_MODELS_TENANT_ID, + name: str = "My App", + description: str = "Summary", + mode: AppMode = AppMode.CHAT, + icon_type: IconType | None = IconType.IMAGE, + icon: str | None = "icon-key", + created_at: datetime | None = None, +) -> _app_response_models_App: + timestamp = created_at or _ts() + app = _app_response_models_App( + id=app_id, + tenant_id=tenant_id, + name=name, + description=description, + mode=mode, + icon_type=icon_type, + icon=icon, + icon_background="#fff", + enable_site=True, + enable_api=True, + max_active_requests=0, + ) + app.created_at = timestamp + app.updated_at = timestamp + return app + + +def _workflow( + *, app_id: str = _APP_RESPONSE_MODELS_APP_ID, tenant_id: str = _APP_RESPONSE_MODELS_TENANT_ID +) -> Workflow: + return Workflow( + id=WORKFLOW_ID, + tenant_id=tenant_id, + app_id=app_id, + type=WorkflowType.CHAT, + version=Workflow.VERSION_DRAFT, + graph=json.dumps({"nodes": [], "edges": []}), + features=json.dumps({}), + created_by=ACCOUNT_ID, + created_at=_ts(8), + updated_by=ACCOUNT_ID, + updated_at=_ts(9), + environment_variables=[], + conversation_variables=[], + ) + + +def _persist_response_graph(session: Session) -> _app_response_models_App: + app = _app(description="Description") + app.created_by = ACCOUNT_ID + app.app_model_config_id = CONFIG_ID + app.workflow_id = WORKFLOW_ID + app.access_mode = "private" + app.create_user_name = "Creator" + app.has_draft_trigger = True + app.permission_keys = ["app.acl.view_layout"] + + model_config = AppModelConfig( + app_id=_APP_RESPONSE_MODELS_APP_ID, + model=json.dumps({"provider": "openai", "name": "gpt-4o"}), + pre_prompt="hello", + created_by=ACCOUNT_ID, + updated_by=ACCOUNT_ID, + ) + model_config.id = CONFIG_ID + model_config.created_at = _ts(9) + model_config.updated_at = _ts(10) + + site = _app_response_models_Site( + id=SITE_ID, + app_id=_APP_RESPONSE_MODELS_APP_ID, + code="site-code", + title="Public Site", + icon_type=IconType.IMAGE, + icon="site-icon", + icon_background="#fff", + description="Site description", + default_language="en-US", + input_placeholder="Ask anything", + customize_token_strategy=CustomizeTokenStrategy.NOT_ALLOW, + ) + site.created_at = _ts(14) + site.updated_at = _ts(14) + tag = Tag(tenant_id=_APP_RESPONSE_MODELS_TENANT_ID, type=TagType.APP, name="Utilities", created_by=ACCOUNT_ID) + tag.id = TAG_ID + binding = TagBinding( + tenant_id=_APP_RESPONSE_MODELS_TENANT_ID, + tag_id=TAG_ID, + target_id=_APP_RESPONSE_MODELS_APP_ID, + created_by=ACCOUNT_ID, + ) + + session.add_all([_account(), app, model_config, _workflow(), site, tag, binding]) + session.commit() + return app + + +class TestAppResponseModelsCases: + @staticmethod + @pytest.fixture + def app_module(): + return controller + + @staticmethod + @pytest.fixture + def app_models(monkeypatch): + monkeypatch.setattr( + app_fields, "build_icon_url", lambda kind, icon: f"signed:{icon}" if kind == IconType.IMAGE else None + ) + return app_fields + + @staticmethod + def test_app_list_query_reads_repeated_tag_ids(app_module): + first_tag_id = "8c4ef3d1-58a1-4d94-8a1c-1c171d889e08" + second_tag_id = "3c39395b-6d1f-4030-8b17-eaa7cc85221c" + query_args = MultiDict( + [ + ("page", "1"), + ("limit", "30"), + ("tag_ids", first_tag_id), + ("tag_ids", second_tag_id), + ] + ) + + query = app_module.query_params_from_request( + app_module.AppListQuery, + list_fields=app_module.APP_LIST_QUERY_ARRAY_FIELDS, + args=query_args, + ) + + assert query.tag_ids == [first_tag_id, second_tag_id] + + @staticmethod + def test_app_list_query_reads_repeated_creator_ids(app_module): + first_creator_id = "9e8959cf-a67b-4d34-9906-1d687517b248" + second_creator_id = "1886f96a-5bf0-42bf-961d-8d2129049076" + query_args = MultiDict( + [ + ("page", "1"), + ("limit", "30"), + ("creator_ids", first_creator_id), + ("creator_ids", second_creator_id), + ] + ) + + query = app_module.query_params_from_request( + app_module.AppListQuery, + list_fields=app_module.APP_LIST_QUERY_ARRAY_FIELDS, + args=query_args, + ) + + assert query.creator_ids == [first_creator_id, second_creator_id] + + @staticmethod + def test_app_list_query_preserves_regular_query_params(app_module): + query_args = MultiDict( + [ + ("page", "2"), + ("limit", "50"), + ("mode", "chat"), + ("name", "Sales Copilot"), + ("is_created_by_me", "true"), + ] + ) + + query = app_module.query_params_from_request( + app_module.AppListQuery, + list_fields=app_module.APP_LIST_QUERY_ARRAY_FIELDS, + args=query_args, + ) + + assert query.page == 2 + assert query.limit == 50 + assert query.mode == "chat" + assert query.name == "Sales Copilot" + assert query.is_created_by_me is True + assert query.tag_ids is None + + @staticmethod + def test_app_list_query_normalizes_empty_repeated_tag_ids_to_none(app_module): + query_args = MultiDict( + [ + ("tag_ids", ""), + ("tag_ids", " "), + ] + ) + + query = app_module.query_params_from_request( + app_module.AppListQuery, + list_fields=app_module.APP_LIST_QUERY_ARRAY_FIELDS, + args=query_args, + ) + + assert query.tag_ids is None + + @staticmethod + def test_app_list_query_rejects_invalid_repeated_tag_id(app_module): + with pytest.raises(ValidationError): + app_module.query_params_from_request( + app_module.AppListQuery, + list_fields=app_module.APP_LIST_QUERY_ARRAY_FIELDS, + args=MultiDict([("tag_ids", "not-a-uuid")]), + ) + + @staticmethod + def test_app_list_query_rejects_invalid_repeated_creator_id(app_module): + with pytest.raises(ValidationError): + app_module.query_params_from_request( + app_module.AppListQuery, + list_fields=app_module.APP_LIST_QUERY_ARRAY_FIELDS, + args=MultiDict([("creator_ids", "not-a-uuid")]), + ) + + @staticmethod + def test_app_list_query_ignores_indexed_tag_ids(app_module): + tag_id = "8c4ef3d1-58a1-4d94-8a1c-1c171d889e08" + query_args = MultiDict( + [ + ("tag_ids[0]", tag_id), + ] + ) + + query = app_module.query_params_from_request( + app_module.AppListQuery, + list_fields=app_module.APP_LIST_QUERY_ARRAY_FIELDS, + args=query_args, + ) + + assert query.tag_ids is None + + @staticmethod + def test_app_list_query_accepts_single_repeated_tag_id(app_module): + tag_id = "8c4ef3d1-58a1-4d94-8a1c-1c171d889e08" + query = app_module.query_params_from_request( + app_module.AppListQuery, + list_fields=app_module.APP_LIST_QUERY_ARRAY_FIELDS, + args=MultiDict([("tag_ids", tag_id)]), + ) + + assert query.tag_ids == [tag_id] + + @staticmethod + def test_create_app_endpoint_rejects_agent_mode(app_module): + with pytest.raises(ValidationError): + app_module.CreateAppPayload.model_validate({"name": "Iris", "mode": "agent", "description": "Agent app"}) + + @staticmethod + def test_app_partial_serialization_uses_aliases(app_models, sqlite_session: Session): + app_partial = app_models.AppPartial + app_obj = _persist_response_graph(sqlite_session) + app_obj.description = "Prompt snippet" + + serialized = app_partial.model_validate( + replace(app_record(app_obj, session=sqlite_session), permission_keys=app_obj.permission_keys), + from_attributes=True, + ).model_dump(mode="json") + + assert serialized["description"] == "Prompt snippet" + assert serialized["mode"] == "chat" + assert serialized["icon_url"] == "signed:icon-key" + assert serialized["created_at"] == int(app_obj.created_at.timestamp()) + assert serialized["updated_at"] == int(app_obj.updated_at.timestamp()) + assert serialized["model_config"]["model"] == {"provider": "openai", "name": "gpt-4o"} + assert serialized["workflow"]["id"] == WORKFLOW_ID + assert serialized["tags"][0]["name"] == "Utilities" + assert serialized["permission_keys"] == ["app.acl.view_layout"] + assert "role" not in serialized + + @staticmethod + def test_app_detail_with_site_includes_nested_serialization( + app: Flask, + app_models, + monkeypatch: pytest.MonkeyPatch, + sqlite_session: Session, + ): + app_obj = _persist_response_graph(sqlite_session) + app_obj.name = "Detailed App" + app_obj.description = "Desc" + app_obj.mode = AppMode.ADVANCED_CHAT + app_obj.icon = "detail-icon" + app_obj.icon_background = "#123456" + app_obj.use_icon_as_answer_icon = True + app_obj.max_active_requests = 5 + app_obj.access_mode = "public" + app_obj.permission_keys = ["app.acl.view_layout", "app.acl.edit"] + model_config = sqlite_session.get(AppModelConfig, CONFIG_ID) + assert model_config is not None + model_config.opening_statement = "hi" + model_config.retriever_resource = json.dumps({"enabled": True}) + monkeypatch.setattr( + "repositories.app.response.load_annotation_reply_config", lambda _session, _app_id: {"enabled": False} + ) + + with app.test_request_context("/"): + response = controller._app_detail_response( + replace( + app_record(app_obj, session=sqlite_session, projection="detail-with-site"), + permission_keys=app_obj.permission_keys, + ), + ) + assert isinstance(response, app_models.AppDetailWithSite) + serialized = response.model_dump(mode="json") + + assert serialized["icon_url"] == "signed:detail-icon" + assert serialized["model_config"]["retriever_resource"] == {"enabled": True} + assert serialized["deleted_tools"] == [] + assert serialized["site"]["icon_url"] == "signed:site-icon" + assert serialized["site"]["input_placeholder"] == "Ask anything" + assert serialized["site"]["created_at"] == int(_ts(14).timestamp()) + assert serialized["permission_keys"] == ["app.acl.view_layout", "app.acl.edit"] + assert serialized["bound_agent_id"] is None + assert "role" not in serialized + + @staticmethod + def test_app_response_view_uses_the_caller_session_for_query_backed_fields( + monkeypatch: pytest.MonkeyPatch, sqlite_session: Session + ): + app_obj = _persist_response_graph(sqlite_session) + decoy_tenant_id = "00000000-0000-0000-0000-000000000002" + decoy_tag = Tag(tenant_id=decoy_tenant_id, type=TagType.APP, name="Decoy", created_by=ACCOUNT_ID) + decoy_tag.id = "00000000-0000-0000-0000-000000000602" + sqlite_session.add_all( + [ + decoy_tag, + TagBinding( + tenant_id=decoy_tenant_id, + tag_id=decoy_tag.id, + target_id=_APP_RESPONSE_MODELS_APP_ID, + created_by=ACCOUNT_ID, + ), + ] + ) + sqlite_session.commit() + calls = [] + + def load_annotation_reply(session, app_id): + calls.append((session, app_id)) + return {"enabled": False} + + monkeypatch.setattr("services.app_service.load_annotation_reply_config", load_annotation_reply) + + view = AppResponseView(app_obj, session=sqlite_session) + site = view.site + workflow = view.workflow + model_config = view.app_model_config + + assert view.desc_or_prompt == "Description" + assert site is not None + assert site.id == SITE_ID + assert workflow is not None + assert workflow.id == WORKFLOW_ID + assert view.bound_agent_id is None + assert view.mode_compatible_with_agent == "chat" + assert view.deleted_tools == [] + assert [tag.name for tag in view.tags] == ["Utilities"] + assert view.author_name == "Creator" + assert model_config is not None + assert model_config.annotation_reply_dict == {"enabled": False} + assert calls == [(sqlite_session, _APP_RESPONSE_MODELS_APP_ID)] + + @staticmethod + def test_app_pagination_aliases_per_page_and_has_next(app_models, sqlite_session: Session): + app_pagination = app_models.AppPagination + item_one = _app( + app_id="00000000-0000-0000-0000-000000000110", + name="Paginated One", + icon="first-icon", + created_at=_ts(15), + ) + item_one.permission_keys = ["app.acl.edit"] + item_two = _app( + app_id="00000000-0000-0000-0000-000000000111", + name="Paginated Two", + mode=AppMode.AGENT_CHAT, + icon_type=IconType.EMOJI, + icon="🙂", + created_at=_ts(16), + ) + pagination = PaginatedResult( + page=2, + per_page=10, + total=50, + items=[app_record(item, session=sqlite_session) for item in (item_one, item_two)], + ) + + serialized = app_pagination.model_validate( + pagination, + from_attributes=True, + ).model_dump(mode="json") + + assert serialized["page"] == 2 + assert serialized["limit"] == 10 + assert serialized["has_more"] is True + assert len(serialized["data"]) == 2 + assert serialized["data"][0]["icon_url"] == "signed:first-icon" + assert serialized["data"][1]["icon_url"] is None + + @staticmethod + @pytest.mark.parametrize("strategy", [None, "cot", "function-calling"]) + def test_app_detail_preserves_historical_agent_strategy( + app: Flask, monkeypatch: pytest.MonkeyPatch, sqlite_session: Session, strategy: str | None + ): + app_obj = _persist_response_graph(sqlite_session) + config = sqlite_session.get(AppModelConfig, CONFIG_ID) + assert config is not None + agent_mode = {"enabled": True, "tools": []} + if strategy is not None: + agent_mode["strategy"] = strategy + config.agent_mode = json.dumps(agent_mode) + sqlite_session.commit() + monkeypatch.setattr( + "repositories.app.response.load_annotation_reply_config", lambda _session, _app_id: {"enabled": False} + ) + + with app.test_request_context("/"): + response = controller._app_detail_response( + app_record(app_obj, session=sqlite_session, projection="detail-with-site") + ).model_dump(mode="json") + + assert response["model_config"]["agent_mode"] == agent_mode + + @staticmethod + def test_app_detail_preserves_historical_agent_mode_without_tools( + app: Flask, monkeypatch: pytest.MonkeyPatch, sqlite_session: Session + ): + app_obj = _persist_response_graph(sqlite_session) + app_obj.mode = AppMode.AGENT_CHAT + config = sqlite_session.get(AppModelConfig, CONFIG_ID) + assert config is not None + agent_mode = {"enabled": True, "max_iteration": 5, "strategy": "function_call"} + config.agent_mode = json.dumps(agent_mode) + sqlite_session.commit() + monkeypatch.setattr( + "repositories.app.response.load_annotation_reply_config", lambda _session, _app_id: {"enabled": False} + ) + + with app.test_request_context("/"): + response = controller._app_detail_response( + app_record(app_obj, session=sqlite_session, projection="detail-with-site") + ).model_dump(mode="json") + + assert response["model_config"]["agent_mode"] == agent_mode + + @staticmethod + def test_app_model_config_response_round_trips_persisted_dataset_and_external_tool_fields( + app: Flask, monkeypatch: pytest.MonkeyPatch, sqlite_session: Session + ): + app_obj = _persist_response_graph(sqlite_session) + config = sqlite_session.get(AppModelConfig, CONFIG_ID) + assert config is not None + dataset_configs = { + "retrieval_model": "multiple", + "reranking_enable": True, + "reranking_enabled": True, + "metadata_filtering_conditions": { + "logical_operator": "and", + "conditions": [ + { + "id": "condition-1", + "metadata_id": "metadata-1", + "name": "category", + "comparison_operator": "is", + "value": "books", + } + ], + }, + "extension_setting": {"value": 1}, + } + external_form = { + "external_data_tool": { + "variable": "customer", + "label": "Customer", + "type": "api", + "enabled": True, + "icon": "icon-key", + "icon_background": "#fff", + "config": {"provider": "crm"}, + "extension_setting": {"value": 2}, + } + } + config.dataset_configs = json.dumps(dataset_configs) + config.user_input_form = json.dumps([external_form]) + sqlite_session.commit() + monkeypatch.setattr( + "repositories.app.response.load_annotation_reply_config", lambda _session, _app_id: {"enabled": False} + ) + + with app.test_request_context("/"): + response = controller._app_detail_response( + app_record(app_obj, session=sqlite_session, projection="detail-with-site") + ).model_dump(mode="json") + + model_config = response["model_config"] + assert model_config["dataset_configs"] == dataset_configs + assert model_config["user_input_form"] == [external_form] + + reimported = AppModelConfig(app_id=_APP_RESPONSE_MODELS_APP_ID) + reimported.from_model_config_dict(model_config) + assert reimported.dataset_configs_dict == dataset_configs + assert reimported.user_input_form_list == [external_form] + + @staticmethod + def test_recommended_app_model_configs_keep_all_published_fields(): + catalog_path = Path(__file__).resolve().parents[5] / "constants" / "recommended_apps.json" + catalog = json.loads(catalog_path.read_text()) + + def assert_fields_preserved(source, serialized): + if isinstance(source, dict): + for key, value in source.items(): + assert key in serialized + assert_fields_preserved(value, serialized[key]) + elif isinstance(source, list): + assert len(source) == len(serialized) + for value, serialized_value in zip(source, serialized, strict=True): + assert_fields_preserved(value, serialized_value) + else: + assert serialized == source + + for entry in catalog["app_details"].values(): + source = yaml.safe_load(entry["export_data"]).get("model_config") or {} + if not source: + continue + payload = {**source, "created_by": None, "created_at": 1, "updated_by": None, "updated_at": 1} + response = AppModelConfigResponse.model_validate(payload).model_dump(mode="json") + assert_fields_preserved(source, response) + + @staticmethod + def test_app_detail_rejects_invalid_known_configuration_fields( + app: Flask, monkeypatch: pytest.MonkeyPatch, sqlite_session: Session + ): + app_obj = _persist_response_graph(sqlite_session) + monkeypatch.setattr( + "repositories.app.response.load_annotation_reply_config", lambda _session, _app_id: {"enabled": False} + ) + with app.test_request_context("/"): + detail = controller._app_detail_response( + app_record(app_obj, session=sqlite_session, projection="detail-with-site") + ).model_dump(mode="json") + + for field in ("mode", "permission_keys", "site", "model_config"): + with pytest.raises(ValidationError): + app_fields.AppDetailWithSite.model_validate( + {key: value for key, value in detail.items() if key != field} + ) + + with pytest.raises(ValidationError): + app_fields.AppDetailWithSite.model_validate({**detail, "mode": "unsupported"}) + + config = detail["model_config"] + for invalid in ( + {"speech_to_text": {"enabled": "not-a-boolean"}}, + {"agent_mode": {"enabled": True, "strategy": "unsupported", "tools": []}}, + {"agent_mode": {"enabled": True, "tools": "not-a-list"}}, + {"agent_mode": {"enabled": True, "tools": [{"provider_id": "only-id"}]}}, + {"dataset_configs": {"retrieval_model": "unsupported"}}, + {"dataset_configs": {"retrieval_model": "multiple", "reranking_mode": "unsupported"}}, + {"model": {"provider": "openai", "name": "gpt-4o", "mode": "unsupported"}}, + {"user_input_form": [{"unsupported": {"label": "Name", "variable": "name"}}]}, + ): + with pytest.raises(ValidationError): + app_fields.AppDetailWithSite.model_validate({**detail, "model_config": {**config, **invalid}}) + + +"""Transport contracts for the Console app application boundary.""" + + +_CONSOLE_APP_API_APP_ID = UUID("11111111-1111-1111-1111-111111111111") +CONTEXT = RequestContext("request", "trace", "actor", "workspace") +RECORD = AppRecord( + str(_CONSOLE_APP_API_APP_ID), + "Example", + "chat", + description="", + created_at=datetime(2024, 1, 1), + updated_at=datetime(2024, 1, 1), +) +SITE = { + "code": "site", + "title": "Example", + "icon_type": None, + "icon": None, + "icon_background": None, + "description": None, + "default_language": "en-US", + "chat_color_theme": None, + "chat_color_theme_inverted": False, + "customize_domain": None, + "copyright": None, + "privacy_policy": None, + "input_placeholder": None, + "custom_disclaimer": "", + "customize_token_strategy": CustomizeTokenStrategy.NOT_ALLOW, + "prompt_public": False, + "show_workflow_steps": False, + "use_icon_as_answer_icon": False, + "created_by": None, + "created_at": datetime(2024, 1, 1), + "updated_by": None, + "updated_at": datetime(2024, 1, 1), +} + + +@dataclass +class AppBoundary: + calls: list[tuple[object, ...]] = app_field(default_factory=list) + status: ImportStatus = ImportStatus.COMPLETED + exported: str | RosterAgentPackageExport = "dsl" + record: AppRecord = RECORD + + def list_apps(self, context: RequestContext, params: AppListParams | StarredAppListParams) -> AppPage: + self.calls.append(("list", context, params)) + return AppPage(params.page, params.limit, 1, False, [RECORD]) + + def create(self, context: RequestContext, params: CreateAppParams) -> AppRecord: + self.calls.append(("create", context, params)) + return RECORD + + def get(self, context: RequestContext, app_id: str) -> AppRecord: + self.calls.append(("get", context, app_id)) + return self.record + + def update(self, context: RequestContext, app_id: str, params: UpdateAppParams) -> AppRecord: + self.calls.append(("update", context, app_id, params)) + return RECORD + + def update_icon(self, context: RequestContext, app_id: str, **kwargs: object) -> AppRecord: + self.calls.append(("icon", context, app_id, kwargs)) + return RECORD + + def copy(self, context: RequestContext, app_id: str, params: CopyAppParams) -> tuple[Import, AppRecord | None]: + self.calls.append(("copy", context, app_id, params)) + return Import(id="import", status=self.status), RECORD + + def export(self, context: RequestContext, app_id: str, options: AppExportOptions) -> str | RosterAgentPackageExport: + self.calls.append(("export", context, app_id, options)) + return self.exported + + def publish(self, context: RequestContext, app_id: str) -> str: + self.calls.append(("publish", context, app_id)) + return "https://creators.example.com" + + def delete(self, context: RequestContext, app_id: str) -> None: + self.calls.append(("delete", context, app_id)) + + def set_trace(self, context: RequestContext, app_id: str, settings: AppTraceSettings) -> None: + self.calls.append(("trace", context, app_id, settings)) + + +@dataclass +class AppServices: + console: AppBoundary + + +@dataclass +class Services: + apps: AppServices + + +def invoke(resource: type[Resource], method: Callable[..., object], **kwargs: object) -> object: + return inspect.unwrap(method)(resource(), CONTEXT, **kwargs) + + +class TestConsoleAppApiCases: + @staticmethod + @pytest.fixture + def boundary(monkeypatch: pytest.MonkeyPatch) -> AppBoundary: + boundary = AppBoundary() + monkeypatch.setattr(controller, "application_services", lambda: Services(AppServices(boundary))) + return boundary + + @staticmethod + def test_list_parses_repeated_filters_and_serializes_materialized_page(app: Flask, boundary: AppBoundary) -> None: + tag = str(_CONSOLE_APP_API_APP_ID) + with app.test_request_context(f"/?page=2&limit=5&tag_ids={tag}&is_created_by_me=true"): + data, status = cast(tuple[dict[str, object], int], invoke(controller.AppListApi, controller.AppListApi.get)) + assert status == 200 + assert data["page"] == 2 + assert data["limit"] == 5 + assert boundary.calls[0][1] is CONTEXT + params = boundary.calls[0][2] + assert isinstance(params, AppListParams) + assert params.tag_ids == [tag] + assert params.is_created_by_me is True + + @staticmethod + @pytest.mark.parametrize("configured", [True, False]) + def test_detail_resolves_urls_at_the_http_boundary( + app: Flask, boundary: AppBoundary, config_overrides: Callable[..., None], configured: bool + ) -> None: + config_overrides( + SERVICE_API_URL="https://api.example.com" if configured else "", + APP_WEB_URL="https://apps.example.com" if configured else "", + ) + boundary.record = replace(RECORD, site=SITE) + with app.test_request_context("/", base_url="https://console.example.com/"): + response = cast( + dict[str, object], invoke(controller.AppApi, controller.AppApi.get, app_id=_CONSOLE_APP_API_APP_ID) + ) + assert response["api_base_url"] == ( + "https://api.example.com/v1" if configured else "https://console.example.com/v1" + ) + site = response["site"] + assert isinstance(site, dict) + assert site["app_base_url"] == ("https://apps.example.com" if configured else "https://console.example.com") + + @staticmethod + @pytest.mark.parametrize("icon_type", [None, "image"]) + def test_update_preserves_icon_type_omission_and_forwards_explicit_identity( + app: Flask, boundary: AppBoundary, icon_type: str | None + ) -> None: + payload = {"name": "Updated", "icon": "image-id"} + if icon_type is not None: + payload["icon_type"] = icon_type + with app.test_request_context("/", method="POST", json=payload): + data = cast( + dict[str, object], invoke(controller.AppApi, controller.AppApi.put, app_id=_CONSOLE_APP_API_APP_ID) + ) + assert data["id"] == str(_CONSOLE_APP_API_APP_ID) + assert boundary.calls[0][1:3] == (CONTEXT, str(_CONSOLE_APP_API_APP_ID)) + params = boundary.calls[0][3] + assert isinstance(params, UpdateAppParams) + assert params.icon_type == icon_type + + @staticmethod + def test_icon_body_and_trace_body_are_parsed(app: Flask, boundary: AppBoundary) -> None: + with app.test_request_context("/", method="POST", json={"icon_type": "image", "icon": "image-id"}): + invoke(controller.AppIconApi, controller.AppIconApi.post, app_id=_CONSOLE_APP_API_APP_ID) + assert boundary.calls[-1][-1] == {"icon_type": "image", "icon": "image-id", "icon_background": ""} + with app.test_request_context("/", method="POST", json={"enabled": True, "tracing_provider": "langfuse"}): + assert invoke(controller.AppTraceApi, controller.AppTraceApi.post, app_id=_CONSOLE_APP_API_APP_ID) == { + "result": "success" + } + assert boundary.calls[-1] == ( + "trace", + CONTEXT, + str(_CONSOLE_APP_API_APP_ID), + AppTraceSettings(True, "langfuse"), + ) + + @staticmethod + @pytest.mark.parametrize( + ("status", "expected"), + [ + (ImportStatus.FAILED, 400), + (ImportStatus.PENDING, 202), + (ImportStatus.COMPLETED, 201), + (ImportStatus.COMPLETED_WITH_WARNINGS, 201), + ], + ) + def test_copy_preserves_import_status_responses( + app: Flask, boundary: AppBoundary, status: ImportStatus, expected: int + ) -> None: + boundary.status = status + with app.test_request_context("/", method="POST", json={"name": "Copy"}): + data, code = cast( + tuple[dict[str, object], int], + invoke(controller.AppCopyApi, controller.AppCopyApi.post, app_id=_CONSOLE_APP_API_APP_ID), + ) + assert code == expected + assert data["id"] == ("import" if expected != 201 else str(_CONSOLE_APP_API_APP_ID)) + + @staticmethod + def test_export_serializes_yaml_and_closes_archive_on_response_close(app: Flask, boundary: AppBoundary) -> None: + with app.test_request_context("/?format=yaml&include_secret=true"): + assert invoke(controller.AppExportApi, controller.AppExportApi.get, app_id=_CONSOLE_APP_API_APP_ID) == { + "data": "dsl" + } + options = boundary.calls[-1][-1] + assert isinstance(options, AppExportOptions) + assert options.include_secret is True + archive = BytesIO(b"zip") + boundary.exported = RosterAgentPackageExport(archive=archive, filename="example.ifpkg", size=3) + with app.test_request_context("/"): + response = invoke(controller.AppExportApi, controller.AppExportApi.get, app_id=_CONSOLE_APP_API_APP_ID) + assert isinstance(response, Response) + assert response.mimetype == "application/zip" + assert "example.ifpkg" in response.headers["Content-Disposition"] + response.close() + assert archive.closed + + @staticmethod + def test_export_closes_archive_when_send_file_fails( + app: Flask, boundary: AppBoundary, monkeypatch: pytest.MonkeyPatch + ) -> None: + archive = BytesIO(b"zip") + boundary.exported = RosterAgentPackageExport(archive=archive, filename="example.ifpkg", size=3) + + def fail(*_args: object, **_kwargs: object) -> None: + raise OSError("delivery failed") + + monkeypatch.setattr(controller, "send_file", fail) + with app.test_request_context("/"), pytest.raises(OSError, match="delivery failed"): + invoke(controller.AppExportApi, controller.AppExportApi.get, app_id=_CONSOLE_APP_API_APP_ID) + assert archive.closed + + @staticmethod + def test_publish_delegates_and_delete_has_no_body(app: Flask, boundary: AppBoundary) -> None: + with app.test_request_context("/"): + assert invoke( + controller.AppPublishToCreatorsPlatformApi, + controller.AppPublishToCreatorsPlatformApi.post, + app_id=_CONSOLE_APP_API_APP_ID, + ) == { + "redirect_url": "https://creators.example.com", + } + assert invoke(controller.AppApi, controller.AppApi.delete, app_id=_CONSOLE_APP_API_APP_ID) == ("", 204) + assert boundary.calls == [ + ("publish", CONTEXT, str(_CONSOLE_APP_API_APP_ID)), + ("delete", CONTEXT, str(_CONSOLE_APP_API_APP_ID)), + ] + + @staticmethod + @pytest.mark.parametrize( + ("error", "expected"), + [ + (ConsoleAppNotFoundError(), AppNotFoundError), + (InvalidAppExportError("invalid version"), BadRequest), + (AppExportPaidPlanRequiredError("paid plan required"), Forbidden), + ], + ) + def test_application_errors_are_translated_at_transport_boundary( + monkeypatch: pytest.MonkeyPatch, error: Exception, expected: type[Exception] + ) -> None: + def fail(*_args: object, **_kwargs: object) -> None: + raise error + + monkeypatch.setattr(Resource, "dispatch_request", fail) + with pytest.raises(expected): + controller.AppResource().dispatch_request() + + @staticmethod + def test_disabled_creators_preserves_error_body(monkeypatch: pytest.MonkeyPatch) -> None: + def fail(*_args: object, **_kwargs: object) -> None: + raise CreatorsPlatformDisabledError("Creators Platform features are not enabled") + + monkeypatch.setattr(Resource, "dispatch_request", fail) + assert controller.AppResource().dispatch_request() == ( + {"error": "Creators Platform features are not enabled"}, + 403, + ) + + @staticmethod + @pytest.mark.parametrize( + "method", + [ + controller.AppApi.put, + controller.AppApi.delete, + controller.AppCopyApi.post, + controller.AppExportApi.get, + controller.AppNameApi.post, + controller.AppIconApi.post, + controller.AppSiteStatus.post, + controller.AppApiStatus.post, + controller.AppTraceApi.post, + ], + ) + def test_mutations_reject_read_only_role_at_admission( + app: Flask, monkeypatch: pytest.MonkeyPatch, method: Callable[..., object] + ) -> None: + account = Account(name="Reader", email="reader@example.com") + account.role = TenantAccountRole.NORMAL + apply_config_overrides(monkeypatch, RBAC_ENABLED=False) + monkeypatch.setattr( + flask_admission, + "current_account_with_tenant", + lambda: AccountWithTenant(account=account, tenant_id="workspace"), + ) + injector = inspect.unwrap(method, stop=lambda f: "allowed_roles" in inspect.getclosurevars(f).nonlocals) + with app.test_request_context(), pytest.raises(Forbidden): + injector(None, app_id=_CONSOLE_APP_API_APP_ID) + + @staticmethod + def test_create_quota_is_part_of_admission_and_precedes_application_call( + app: Flask, boundary: AppBoundary, monkeypatch: pytest.MonkeyPatch + ) -> None: + from services.entities.feature_entities import FeatureModel, LimitationModel + + account = Account(name="Editor", email="editor@example.com") + account.role = TenantAccountRole.EDITOR + account.id = CONTEXT.account_id + identity = AccountWithTenant(account=account, tenant_id=CONTEXT.active_workspace_id) + apply_config_overrides(monkeypatch, DEPLOYMENT_EDITION="CLOUD", RBAC_ENABLED=False) + monkeypatch.setattr(flask_admission, "current_account_with_tenant", lambda: identity) + monkeypatch.setattr("controllers.console.wraps.current_account_with_tenant", lambda: identity) + monkeypatch.setattr(flask_admission, "enforce_rbac_checks", lambda **_kwargs: None) + monkeypatch.setattr( + "controllers.console.wraps.FeatureService.get_features", + lambda *_args, **_kwargs: FeatureModel(apps=LimitationModel(size=1, limit=1)), + ) + admission = inspect.unwrap( + controller.AppListApi.post, stop=lambda f: inspect.getclosurevars(f).nonlocals.get("resource") == "apps" + ) + with ( + app.test_request_context(method="POST", json={"name": "Example", "mode": "chat"}), + pytest.raises(Forbidden), + ): + admission(controller.AppListApi()) + assert boundary.calls == [] + + @staticmethod + def test_request_validation_retains_422_and_redacts_submitted_values(app: Flask, boundary: AppBoundary) -> None: + from werkzeug.exceptions import UnprocessableEntity + + with app.test_request_context(method="PUT", json={"name": "", "icon": "private-input"}): + with pytest.raises(UnprocessableEntity) as error: + invoke(controller.AppApi, controller.AppApi.put, app_id=_CONSOLE_APP_API_APP_ID) + assert "private-input" not in (error.value.description or "") + assert boundary.calls == [] + + +"""Regression tests for CreateAppPayload mode validation.""" + + +class TestCreateAppPayloadCases: + class TestCreateAppPayloadMode: + @pytest.mark.parametrize( + "mode", + ["chat", "agent-chat", "advanced-chat", "workflow", "completion"], + ) + def test_accepts_supported_modes(self, mode: str): + payload = CreateAppPayload.model_validate({"name": "X", "mode": mode}) + assert payload.mode == mode + + def test_rejects_agent_mode(self): + with pytest.raises(ValidationError): + CreateAppPayload.model_validate({"name": "X", "mode": "agent"}) + + def test_rejects_unknown_mode(self): + with pytest.raises(ValidationError): + CreateAppPayload.model_validate({"name": "X", "mode": "not-a-mode"}) diff --git a/api/tests/unit_tests/controllers/console/app/test_app_apis.py b/api/tests/unit_tests/controllers/console/app/test_app_apis.py deleted file mode 100644 index b68a3fb717d..00000000000 --- a/api/tests/unit_tests/controllers/console/app/test_app_apis.py +++ /dev/null @@ -1,812 +0,0 @@ -"""Unit coverage for console app controller contracts and response mapping.""" - -from __future__ import annotations - -import uuid -from collections.abc import Iterator -from inspect import unwrap -from types import SimpleNamespace -from unittest.mock import MagicMock, patch - -import pytest -from flask import Flask -from pydantic import ValidationError -from sqlalchemy import Engine -from sqlalchemy.orm import Session -from werkzeug.exceptions import BadRequest, NotFound - -from controllers.console.app import ( - annotation as annotation_module, -) -from controllers.console.app import ( - app as app_module, -) -from controllers.console.app import ( - completion as completion_module, -) -from controllers.console.app import ( - message as message_module, -) -from controllers.console.app import ( - ops_trace as ops_trace_module, -) -from controllers.console.app import ( - site as site_module, -) -from controllers.console.app import ( - statistic as statistic_module, -) -from controllers.console.app import ( - workflow_app_log as workflow_app_log_module, -) -from controllers.console.app import ( - workflow_draft_variable as workflow_draft_variable_module, -) -from controllers.console.app import ( - workflow_statistic as workflow_statistic_module, -) -from controllers.console.app import ( - wraps as wraps_module, -) -from controllers.console.app.completion import ChatMessagePayload, CompletionMessagePayload -from controllers.console.app.error import ( - AppNotFoundError, - TracingConfigNotFoundError, - TracingConfigVerificationFailedError, -) -from controllers.console.app.mcp_server import MCPServerCreatePayload, MCPServerUpdatePayload -from controllers.console.app.ops_trace import TraceConfigPayload, TraceProviderQuery -from controllers.console.app.site import AppSiteUpdatePayload -from controllers.console.app.workflow import AdvancedChatWorkflowRunPayload, SyncDraftWorkflowPayload -from controllers.console.app.workflow_app_log import WorkflowAppLogQuery -from controllers.console.app.workflow_draft_variable import ( - EnvironmentVariableUpdatePayload, - WorkflowDraftVariableListQuery, - WorkflowDraftVariableUpdatePayload, -) -from controllers.console.app.workflow_statistic import WorkflowStatisticQuery -from machinery.context import RequestContext -from models import App, Site -from models.account import Account, AccountStatus -from models.engine import db -from services.app_site_service import ( - AppSiteAppNotFoundError, - AppSiteChanges, - AppSiteCommandResult, - AppSiteNotFoundError, -) -from services.app_tracing_config_service import ( - AppTracingConfigNotFoundError, - AppTracingConfigVerificationFailedError, -) - -APP_ID = "11111111-1111-1111-1111-111111111111" -TENANT_ID = "22222222-2222-2222-2222-222222222222" -USER_ID = "33333333-3333-3333-3333-333333333333" - - -def _make_account() -> Account: - account = Account( - name="tester", - email="tester@example.com", - status=AccountStatus.ACTIVE, - ) - account.id = USER_ID - return account - - -def _make_app( - app_id: str = APP_ID, - *, - tenant_id: str = TENANT_ID, - icon_type: app_module.IconType | None = None, -) -> App: - app = App() - app.id = app_id - app.tenant_id = tenant_id - app.icon_type = icon_type - return app - - -def _make_request_context() -> RequestContext: - return RequestContext( - request_id="request-1", - trace_id="trace-1", - account_id=USER_ID, - active_workspace_id=TENANT_ID, - ) - - -@pytest.fixture -def database_app() -> Iterator[Flask]: - app = Flask(__name__) - app.config["SQLALCHEMY_DATABASE_URI"] = "sqlite:///:memory:" - db.init_app(app) - - with app.app_context(): - Site.__table__.create(db.engine) - yield app - - -class TestCompletionEndpoints: - def test_completion_create_payload(self): - payload = CompletionMessagePayload(inputs={"prompt": "test"}, model_config={}) - assert payload.inputs == {"prompt": "test"} - - def test_chat_message_payload_uuid_validation(self): - payload = ChatMessagePayload( - inputs={}, - model_config={}, - query="hi", - conversation_id=str(uuid.uuid4()), - parent_message_id=str(uuid.uuid4()), - ) - assert payload.query == "hi" - - def test_completion_api_success(self, app: Flask, monkeypatch: pytest.MonkeyPatch, sqlite_engine: Engine) -> None: - api = completion_module.CompletionMessageApi() - method = unwrap(api.post) - - monkeypatch.setattr( - completion_module.AppGenerateService, - "generate", - lambda **_kwargs: {"text": "ok"}, - ) - monkeypatch.setattr( - completion_module.helper, - "compact_generate_response", - lambda response: {"result": response}, - ) - - with ( - Session(sqlite_engine) as session, - app.test_request_context("/", json={"inputs": {}, "model_config": {}, "query": "hi"}), - ): - resp = method( - api, - CompletionMessagePayload(inputs={}, model_config={}, query="hi"), - session, - _make_account(), - app_model=_make_app(), - ) - - assert resp == {"result": {"text": "ok"}} - - def test_completion_api_conversation_not_exists( - self, app: Flask, monkeypatch: pytest.MonkeyPatch, sqlite_engine: Engine - ) -> None: - api = completion_module.CompletionMessageApi() - method = unwrap(api.post) - - monkeypatch.setattr( - completion_module.AppGenerateService, - "generate", - lambda **_kwargs: (_ for _ in ()).throw( - completion_module.services.errors.conversation.ConversationNotExistsError() - ), - ) - - with ( - Session(sqlite_engine) as session, - app.test_request_context("/", json={"inputs": {}, "model_config": {}, "query": "hi"}), - pytest.raises(NotFound), - ): - method( - api, - CompletionMessagePayload(inputs={}, model_config={}, query="hi"), - session, - _make_account(), - app_model=_make_app(), - ) - - def test_completion_api_provider_not_initialized( - self, app: Flask, monkeypatch: pytest.MonkeyPatch, sqlite_engine: Engine - ) -> None: - api = completion_module.CompletionMessageApi() - method = unwrap(api.post) - - monkeypatch.setattr( - completion_module.AppGenerateService, - "generate", - lambda **_kwargs: (_ for _ in ()).throw(completion_module.ProviderTokenNotInitError("x")), - ) - - with ( - Session(sqlite_engine) as session, - app.test_request_context("/", json={"inputs": {}, "model_config": {}, "query": "hi"}), - pytest.raises(completion_module.ProviderNotInitializeError), - ): - method( - api, - CompletionMessagePayload(inputs={}, model_config={}, query="hi"), - session, - _make_account(), - app_model=_make_app(), - ) - - def test_completion_api_quota_exceeded( - self, app: Flask, monkeypatch: pytest.MonkeyPatch, sqlite_engine: Engine - ) -> None: - api = completion_module.CompletionMessageApi() - method = unwrap(api.post) - - monkeypatch.setattr( - completion_module.AppGenerateService, - "generate", - lambda **_kwargs: (_ for _ in ()).throw(completion_module.QuotaExceededError()), - ) - - with ( - Session(sqlite_engine) as session, - app.test_request_context("/", json={"inputs": {}, "model_config": {}, "query": "hi"}), - pytest.raises(completion_module.ProviderQuotaExceededError), - ): - method( - api, - CompletionMessagePayload(inputs={}, model_config={}, query="hi"), - session, - _make_account(), - app_model=_make_app(), - ) - - -class TestOpsTraceEndpoints: - def test_ops_trace_query_basic(self): - query = TraceProviderQuery(tracing_provider="langfuse") - assert query.tracing_provider == "langfuse" - - def test_ops_trace_config_payload(self): - payload = TraceConfigPayload(tracing_provider="langfuse", tracing_config={"api_key": "k"}) - assert payload.tracing_config["api_key"] == "k" - - def test_trace_app_config_get_empty(self, app: Flask, monkeypatch: pytest.MonkeyPatch): - api = ops_trace_module.TraceAppConfigApi() - method = unwrap(api.get) - tracing_configs = MagicMock() - tracing_configs.get.return_value = None - monkeypatch.setattr( - ops_trace_module, - "application_services", - lambda: SimpleNamespace(app_tracing_configs=tracing_configs), - ) - - with app.test_request_context("/?tracing_provider=langfuse"): - result = method( - api, - TraceProviderQuery(tracing_provider="langfuse"), - _make_request_context(), - uuid.UUID(APP_ID), - ) - - assert result == {"has_not_configured": True} - tracing_configs.get.assert_called_once_with( - context=_make_request_context(), - app_id=APP_ID, - tracing_provider="langfuse", - ) - - def test_trace_app_config_post_invalid(self, app: Flask, monkeypatch: pytest.MonkeyPatch): - api = ops_trace_module.TraceAppConfigApi() - method = unwrap(api.post) - tracing_configs = MagicMock() - tracing_configs.create.side_effect = AppTracingConfigVerificationFailedError() - monkeypatch.setattr( - ops_trace_module, - "application_services", - lambda: SimpleNamespace(app_tracing_configs=tracing_configs), - ) - - with app.test_request_context( - "/", - json={"tracing_provider": "langfuse", "tracing_config": {"api_key": "k"}}, - ): - with pytest.raises(TracingConfigVerificationFailedError): - method( - api, - TraceConfigPayload(tracing_provider="langfuse", tracing_config={"api_key": "k"}), - _make_request_context(), - uuid.UUID(APP_ID), - ) - - def test_trace_app_config_delete_not_found(self, app: Flask, monkeypatch: pytest.MonkeyPatch): - api = ops_trace_module.TraceAppConfigApi() - method = unwrap(api.delete) - tracing_configs = MagicMock() - tracing_configs.delete.side_effect = AppTracingConfigNotFoundError() - monkeypatch.setattr( - ops_trace_module, - "application_services", - lambda: SimpleNamespace(app_tracing_configs=tracing_configs), - ) - - with app.test_request_context("/?tracing_provider=langfuse"): - with pytest.raises(TracingConfigNotFoundError): - method( - api, - TraceProviderQuery(tracing_provider="langfuse"), - _make_request_context(), - uuid.UUID(APP_ID), - ) - - -class TestSiteEndpoints: - @staticmethod - def _command_result(*, code: str = "test-code") -> AppSiteCommandResult: - return AppSiteCommandResult( - app_id=APP_ID, - title="My Site", - description="Test site", - default_language="en-US", - input_placeholder="Ask me anything", - code=code, - icon=None, - icon_background=None, - customize_domain=None, - copyright=None, - privacy_policy=None, - custom_disclaimer="", - customize_token_strategy="not_allow", - prompt_public=False, - show_workflow_steps=True, - use_icon_as_answer_icon=False, - ) - - def test_site_payload_maps_to_application_changes(self): - payload = AppSiteUpdatePayload( - title="My Site", - description="Test site", - input_placeholder="Ask me anything", - ) - assert payload.to_changes() == AppSiteChanges( - title="My Site", - description="Test site", - input_placeholder="Ask me anything", - ) - - def test_site_default_language_validation(self): - payload = AppSiteUpdatePayload(default_language="en-US") - assert payload.default_language == "en-US" - - def test_app_site_update_post( - self, - ) -> None: - api = site_module.AppSite() - method = unwrap(api.post) - services = MagicMock() - services.app_sites.update.return_value = self._command_result() - context = RequestContext("request-1", None, USER_ID, TENANT_ID) - - with patch.object(site_module, "application_services", return_value=services): - result = method( - api, - AppSiteUpdatePayload(title="My Site", input_placeholder="Ask me anything"), - context, - app_id=uuid.UUID(APP_ID), - ) - - assert isinstance(result, dict) - assert result["title"] == "My Site" - assert result["input_placeholder"] == "Ask me anything" - assert result["access_token"] == "test-code" - assert result["code"] == "test-code" - services.app_sites.update.assert_called_once_with( - context, - APP_ID, - AppSiteChanges(title="My Site", input_placeholder="Ask me anything"), - ) - - def test_app_site_access_token_reset( - self, - ) -> None: - api = site_module.AppSiteAccessTokenReset() - method = unwrap(api.post) - services = MagicMock() - services.app_sites.reset_access_token.return_value = self._command_result(code="new-code") - context = RequestContext("request-1", None, USER_ID, TENANT_ID) - - with patch.object(site_module, "application_services", return_value=services): - result = method(api, context, app_id=uuid.UUID(APP_ID)) - - assert isinstance(result, dict) - assert result["access_token"] == "new-code" - assert result["code"] == "new-code" - services.app_sites.reset_access_token.assert_called_once_with(context, APP_ID) - - def test_app_site_update_maps_missing_app(self) -> None: - api = site_module.AppSite() - method = unwrap(api.post) - services = MagicMock() - services.app_sites.update.side_effect = AppSiteAppNotFoundError() - context = RequestContext("request-1", None, USER_ID, TENANT_ID) - - with ( - patch.object(site_module, "application_services", return_value=services), - pytest.raises(AppNotFoundError), - ): - method(api, AppSiteUpdatePayload(), context, app_id=uuid.UUID(APP_ID)) - - def test_app_site_reset_maps_missing_site(self) -> None: - api = site_module.AppSiteAccessTokenReset() - method = unwrap(api.post) - services = MagicMock() - services.app_sites.reset_access_token.side_effect = AppSiteNotFoundError() - context = RequestContext("request-1", None, USER_ID, TENANT_ID) - - with ( - patch.object(site_module, "application_services", return_value=services), - pytest.raises(NotFound), - ): - method(api, context, app_id=uuid.UUID(APP_ID)) - - -class TestWorkflowEndpoints: - def test_workflow_copy_payload(self): - payload = SyncDraftWorkflowPayload(graph={}, features={}) - assert payload.graph == {} - assert payload.model_dump()["is_collaborative"] is False - - def test_workflow_sync_payload_accepts_collaboration_marker(self): - payload = SyncDraftWorkflowPayload.model_validate({"graph": {}, "features": {}, "_is_collaborative": True}) - assert payload.is_collaborative is True - assert payload.model_dump()["is_collaborative"] is True - - def test_workflow_mode_query(self): - payload = AdvancedChatWorkflowRunPayload(inputs={}, query="hi") - assert payload.query == "hi" - - -class TestWorkflowAppLogEndpoints: - def test_workflow_app_log_query(self): - query = WorkflowAppLogQuery(keyword="test", page=1, limit=20) - assert query.keyword == "test" - - def test_workflow_app_log_query_detail_bool(self): - query = WorkflowAppLogQuery(detail="true") - assert query.detail is True - - def test_workflow_app_log_api_get(self, database_app: Flask, monkeypatch: pytest.MonkeyPatch) -> None: - api = workflow_app_log_module.WorkflowAppLogApi() - method = unwrap(api.get) - workflow_app_logs = MagicMock() - workflow_app_logs.list_logs.return_value = { - "page": 1, - "limit": 20, - "total": 0, - "has_more": False, - "data": [], - } - services = MagicMock(workflow_app_logs=workflow_app_logs) - monkeypatch.setattr(workflow_app_log_module, "application_services", lambda: services) - context = RequestContext("request-1", None, USER_ID, TENANT_ID) - app_model = _make_app("app-1") - - with database_app.test_request_context("/?page=1&limit=20"): - result = method(api, WorkflowAppLogQuery(page=1, limit=20), context, app_model=app_model) - - assert result == {"page": 1, "limit": 20, "total": 0, "has_more": False, "data": []} - workflow_app_logs.list_logs.assert_called_once_with( - tenant_id=app_model.tenant_id, - app_id=app_model.id, - keyword=None, - status=None, - created_at_before=None, - created_at_after=None, - page=1, - limit=20, - detail=False, - created_by_end_user_session_id=None, - created_by_account=None, - ) - - -class TestWorkflowDraftVariableEndpoints: - def test_workflow_variable_creation(self): - payload = WorkflowDraftVariableUpdatePayload(name="var1", value="test") - assert payload.name == "var1" - - def test_workflow_variable_collection_get(self, database_app: Flask, monkeypatch: pytest.MonkeyPatch) -> None: - api = workflow_draft_variable_module.WorkflowVariableCollectionApi() - method = unwrap(api.get) - - class DummyDraftService: - def __init__(self, session: Session): - self.session = session - - def list_variables_without_values(self, **_kwargs): - assert self.session.get_bind() is db.engine - return {"items": [], "total": 0} - - class DummyWorkflowService: - def is_workflow_exist(self, *args, **kwargs): - return True - - monkeypatch.setattr(workflow_draft_variable_module, "WorkflowDraftVariableService", DummyDraftService) - monkeypatch.setattr(workflow_draft_variable_module, "WorkflowService", DummyWorkflowService) - - with database_app.test_request_context("/?page=1&limit=20"): - result = method( - api, - WorkflowDraftVariableListQuery(page=1, limit=20), - _make_account(), - app_model=_make_app("app-1"), - ) - - assert result == {"items": [], "total": 0} - - def test_environment_variable_update_payload_preserves_full_replace_default(self) -> None: - payload = EnvironmentVariableUpdatePayload(environment_variables=[]) - - assert payload.patch is False - assert payload.deleted_environment_variable_ids == [] - - @pytest.mark.parametrize( - "payload", - [ - {"environment_variables": [], "deleted_environment_variable_ids": ["env-a"]}, - { - "environment_variables": [{"name": "a", "value_type": "string", "value": "a"}], - "patch": True, - }, - { - "environment_variables": [{"id": "env-a", "name": "a", "value_type": "string", "value": "a"}], - "patch": True, - "deleted_environment_variable_ids": ["env-a"], - }, - ], - ) - def test_environment_variable_patch_payload_rejects_ambiguous_mutations(self, payload: dict) -> None: - with pytest.raises(ValidationError): - EnvironmentVariableUpdatePayload.model_validate(payload) - - def test_environment_variable_collection_post_routes_patch_to_service( - self, database_app: Flask, monkeypatch: pytest.MonkeyPatch - ) -> None: - api = workflow_draft_variable_module.EnvironmentVariableCollectionApi() - method = unwrap(api.post) - captured: dict = {} - - class DummyWorkflowService: - def patch_draft_workflow_environment_variables(self, **kwargs) -> None: - captured.update(kwargs) - - def update_draft_workflow_environment_variables(self, **_kwargs) -> None: - raise AssertionError("patch request must not use full replacement") - - monkeypatch.setattr(workflow_draft_variable_module, "WorkflowService", DummyWorkflowService) - - with database_app.test_request_context( - "/", - json={ - "environment_variables": [{"id": "env-a", "name": "a", "value_type": "string", "value": "new-a"}], - "patch": True, - "deleted_environment_variable_ids": ["env-b"], - }, - ): - result = method( - api, - EnvironmentVariableUpdatePayload( - environment_variables=[{"id": "env-a", "name": "a", "value_type": "string", "value": "new-a"}], - patch=True, - deleted_environment_variable_ids=["env-b"], - ), - _make_account(), - app_model=_make_app(), - ) - - assert result == {"result": "success"} - assert [(variable.id, variable.value) for variable in captured["environment_variables"]] == [("env-a", "new-a")] - assert captured["deleted_environment_variable_ids"] == ["env-b"] - assert captured["app_model"].id == APP_ID - assert captured["account"].id == USER_ID - assert captured["session"].get_bind() is db.engine - - -class TestWorkflowStatisticEndpoints: - def test_workflow_statistic_time_range(self): - query = WorkflowStatisticQuery(start="2024-01-01", end="2024-12-31") - assert query.start == "2024-01-01" - - def test_workflow_statistic_blank_to_none(self): - query = WorkflowStatisticQuery(start="", end="") - assert query.start is None - assert query.end is None - - def test_workflow_daily_runs_statistic(self, database_app: Flask, monkeypatch: pytest.MonkeyPatch) -> None: - service = MagicMock() - service.get_daily_runs.return_value = [{"date": "2024-01-01", "runs": 2}] - monkeypatch.setattr( - workflow_statistic_module, - "application_services", - lambda: SimpleNamespace(workflow_statistics=service), - ) - self._patch_statistic_time_range(monkeypatch) - - api = workflow_statistic_module.WorkflowDailyRunsStatistic() - method = unwrap(api.get) - request_context = _make_request_context() - - with database_app.test_request_context("/"): - response = method( - api, - WorkflowStatisticQuery(), - request_context, - app_model=_make_app("app-1", tenant_id="t1"), - ) - - assert response == {"data": [{"date": "2024-01-01", "runs": 2}]} - service.get_daily_runs.assert_called_once_with( - request_context, - app_id="app-1", - start_date=None, - end_date=None, - timezone="UTC", - ) - - def test_workflow_daily_terminals_statistic(self, database_app: Flask, monkeypatch: pytest.MonkeyPatch) -> None: - service = MagicMock() - service.get_daily_terminals.return_value = [{"date": "2024-01-02", "terminal_count": 3}] - monkeypatch.setattr( - workflow_statistic_module, - "application_services", - lambda: SimpleNamespace(workflow_statistics=service), - ) - self._patch_statistic_time_range(monkeypatch) - - api = workflow_statistic_module.WorkflowDailyTerminalsStatistic() - method = unwrap(api.get) - request_context = _make_request_context() - - with database_app.test_request_context("/"): - response = method( - api, - WorkflowStatisticQuery(), - request_context, - app_model=_make_app("app-1", tenant_id="t1"), - ) - - assert response == {"data": [{"date": "2024-01-02", "terminal_count": 3}]} - service.get_daily_terminals.assert_called_once_with( - request_context, - app_id="app-1", - start_date=None, - end_date=None, - timezone="UTC", - ) - - def test_workflow_daily_token_cost_statistic( - self, - database_app: Flask, - monkeypatch: pytest.MonkeyPatch, - ) -> None: - service = MagicMock() - service.get_daily_token_costs.return_value = [{"date": "2024-01-03", "token_count": 4}] - monkeypatch.setattr( - workflow_statistic_module, - "application_services", - lambda: SimpleNamespace(workflow_statistics=service), - ) - self._patch_statistic_time_range(monkeypatch) - - api = workflow_statistic_module.WorkflowDailyTokenCostStatistic() - method = unwrap(api.get) - request_context = _make_request_context() - - with database_app.test_request_context("/"): - response = method( - api, - WorkflowStatisticQuery(), - request_context, - app_model=_make_app("app-1", tenant_id="t1"), - ) - - assert response == {"data": [{"date": "2024-01-03", "token_count": 4}]} - service.get_daily_token_costs.assert_called_once_with( - request_context, - app_id="app-1", - start_date=None, - end_date=None, - timezone="UTC", - ) - - def test_workflow_average_app_interaction_statistic( - self, - database_app: Flask, - monkeypatch: pytest.MonkeyPatch, - ) -> None: - service = MagicMock() - service.get_average_app_interactions.return_value = [{"date": "2024-01-04", "interactions": 2.5}] - monkeypatch.setattr( - workflow_statistic_module, - "application_services", - lambda: SimpleNamespace(workflow_statistics=service), - ) - self._patch_statistic_time_range(monkeypatch) - - api = workflow_statistic_module.WorkflowAverageAppInteractionStatistic() - method = unwrap(api.get) - request_context = _make_request_context() - - with database_app.test_request_context("/"): - response = method( - api, - WorkflowStatisticQuery(), - request_context, - app_model=_make_app("app-1", tenant_id="t1"), - ) - - assert response == {"data": [{"date": "2024-01-04", "interactions": 2.5}]} - service.get_average_app_interactions.assert_called_once_with( - request_context, - app_id="app-1", - start_date=None, - end_date=None, - timezone="UTC", - ) - - def test_workflow_statistic_invalid_time_range(self, monkeypatch: pytest.MonkeyPatch) -> None: - account = _make_account() - account.timezone = "UTC" - monkeypatch.setattr( - workflow_statistic_module, - "current_account_with_tenant", - lambda: SimpleNamespace(account=account), - ) - monkeypatch.setattr( - workflow_statistic_module, - "parse_time_range", - MagicMock(side_effect=ValueError("invalid range")), - ) - - with pytest.raises(BadRequest, match="invalid range"): - workflow_statistic_module._resolve_statistic_time_range(WorkflowStatisticQuery()) - - @staticmethod - def _patch_statistic_time_range(monkeypatch: pytest.MonkeyPatch) -> None: - account = _make_account() - account.timezone = "UTC" - monkeypatch.setattr( - workflow_statistic_module, - "current_account_with_tenant", - lambda: SimpleNamespace(account=account), - ) - monkeypatch.setattr( - workflow_statistic_module, - "parse_time_range", - lambda *_args, **_kwargs: (None, None), - ) - - -class TestWrapsEndpoints: - def test_get_app_model_context(self): - assert hasattr(wraps_module, "get_app_model") - - -class TestMCPServerEndpoints: - def test_mcp_server_connection(self): - payload = MCPServerCreatePayload(parameters={"url": "http://localhost:3000"}) - assert payload.parameters["url"] == "http://localhost:3000" - - def test_mcp_server_update_payload(self): - payload = MCPServerUpdatePayload( - id="abcdefab-1234-4567-89ab-abcdefabcdef", parameters={"timeout": 30}, status="active" - ) - assert payload.status == "active" - - -class TestErrorHandling: - def test_annotation_list_query_validation(self): - with pytest.raises(ValueError, match="1 validation error for AnnotationListQuery\\\npage"): - annotation_module.AnnotationListQuery(page=0) - - -class TestPayloadIntegration: - def test_multiple_payload_types(self): - payloads = [ - annotation_module.AnnotationReplyPayload( - score_threshold=0.5, embedding_provider_name="openai", embedding_model_name="text-embedding-3-small" - ), - message_module.AppMessageFeedbackPayload(message_id=str(uuid.uuid4()), rating="like"), - statistic_module.StatisticTimeRangeQuery(start="2024-01-01"), - ] - assert len(payloads) == 3 - assert all(p is not None for p in payloads) diff --git a/api/tests/unit_tests/controllers/console/app/test_app_import.py b/api/tests/unit_tests/controllers/console/app/test_app_import.py new file mode 100644 index 00000000000..2531648168d --- /dev/null +++ b/api/tests/unit_tests/controllers/console/app/test_app_import.py @@ -0,0 +1,881 @@ +import io +import zipfile +from collections.abc import Callable, Iterator +from dataclasses import dataclass, field +from inspect import unwrap +from typing import BinaryIO, Protocol, cast +from unittest.mock import Mock +from uuid import UUID, uuid4 + +import httpx +import pytest +from flask import Flask +from flask_restx import Resource +from pydantic import ValidationError +from sqlalchemy import Engine, event, select +from sqlalchemy.orm import Session, sessionmaker +from werkzeug.exceptions import Forbidden + +from constants.dsl_version import CURRENT_APP_DSL_VERSION +from controllers.console.app import app as app_module +from controllers.console.app import app_import as console +from controllers.console.app import app_import as controller +from controllers.console.app import app_import as import_module +from controllers.console.app.error import AppNotFoundError, DraftWorkflowNotSync +from controllers.inner_api.app import dsl as inner +from controllers.openapi import app_dsl as openapi +from controllers.openapi._errors import OpenApiErrorFormatter +from controllers.openapi._models import AppDslImportPayload +from core.rbac import RBACPermission +from enums import DeploymentEdition +from extensions.ext_redis import redis_client +from libs.external_api import ExternalApi +from machinery.context import RequestContext +from models.account import Account, Tenant, TenantAccountJoin, TenantAccountRole +from models.model import App, AppMode +from services import app_import_source +from services.agent.errors import InvalidRosterAgentPackageError +from services.agent.roster_package_importer import RosterAgentPackageImportResult +from services.app.console_gateway import AppTransferGateway, EnterpriseConsoleAppAccess +from services.app_dsl_service import AppDslService, PendingData +from services.app_package_service import AppPackageService, PreparedAppPackage +from services.billing_service import BillingService, _BillingHTTPStatusError +from services.enterprise.enterprise_service import EnterpriseService +from services.entities.dsl_entities import AppImportPackage, AppImportParams, DslImportWarning, Import, ImportStatus +from services.entities.feature_entities import FeatureModel, LimitationModel +from services.errors.app import DraftChangedError +from services.errors.base import NoPermissionError +from services.feature_service import FeatureService +from services.plugin.dependencies_analysis import DependenciesAnalysisService +from services.recommended_app_package_service import RecommendedAgentPackageSource, RecommendedAppPackageService +from services.system_feature_service import SystemFeatureService +from tests.helpers.unit.config_override import apply_config_overrides +from tests.unit_tests.controllers.conftest import ControllerTestServices + +"""Console imports finalize persistence before publishing permissions and access settings.""" + + +class TestAppImportApiCases: + @staticmethod + @pytest.fixture + def import_context(sqlite_session_factory: sessionmaker[Session]) -> RequestContext: + account_id, tenant_id = str(uuid4()), str(uuid4()) + account = Account(name="Importer", email="importer@example.com") + account.id = account_id + tenant = Tenant(name="Workspace") + tenant.id = tenant_id + with sqlite_session_factory.begin() as session: + session.add_all( + [ + account, + tenant, + TenantAccountJoin(tenant_id=tenant_id, account_id=account_id, role=TenantAccountRole.OWNER), + ] + ) + return RequestContext("import-request", None, account_id, tenant_id) + + @staticmethod + @pytest.fixture + def connections(sqlite_engine: Engine) -> Iterator[set[object]]: + checked_out: set[object] = set() + + def checkout(connection: object, *_args: object) -> None: + checked_out.add(connection) + + def checkin(connection: object, *_args: object) -> None: + checked_out.remove(connection) + + event.listen(sqlite_engine, "checkout", checkout) + event.listen(sqlite_engine, "checkin", checkin) + try: + yield checked_out + assert not checked_out + finally: + event.remove(sqlite_engine, "checkout", checkout) + event.remove(sqlite_engine, "checkin", checkin) + + @staticmethod + @pytest.mark.usefixtures("app_query_services") + @pytest.mark.parametrize("confirm", [False, True]) + @pytest.mark.parametrize("overwrite", [False, True]) + @pytest.mark.parametrize("status", list(ImportStatus)) + def test_import_transaction_and_response_contract( + app: Flask, + sqlite_session_factory: sessionmaker[Session], + import_context: RequestContext, + connections: set[object], + monkeypatch: pytest.MonkeyPatch, + config_overrides, + status: ImportStatus, + overwrite: bool, + confirm: bool, + ) -> None: + config_overrides(RBAC_ENABLED=True, DEPLOYMENT_EDITION="COMMUNITY") + monkeypatch.setattr( + "services.app.console_gateway.rbac_service.RBACService.CheckAccess.check", lambda *a, **k: True + ) + app_id = str(uuid4()) + result = Import(id="import-1", status=status, app_id=app_id, app_created=not overwrite) + pending = PendingData( + tenant_id=import_context.active_workspace_id, + account_id=import_context.account_id, + import_mode="yaml-content", + yaml_content="app: {}", + app_id=app_id if overwrite else None, + ) + monkeypatch.setattr(redis_client, "get", lambda _key: pending.model_dump_json()) + + def persist(dsl: AppDslService, *, account: Account, **_kwargs) -> Import: + assert not connections, "Actor lookup must release its connection before starting the DSL operation" + assert account.id == import_context.account_id + assert account.current_tenant_id == import_context.active_workspace_id + dsl._session.add( + App( + id=app_id, + tenant_id=account.current_tenant_id, + name="Imported", + mode=AppMode.WORKFLOW, + enable_site=True, + enable_api=True, + ) + ) + dsl._session.flush() + return result + + monkeypatch.setattr(AppDslService, "confirm_import" if confirm else "import_app", persist) + permissions: list[str] = [] + access_updates: list[tuple[str, str]] = [] + + def get_permissions(_self, context: RequestContext, imported_id: str) -> list[str]: + assert not connections + assert context == import_context + with sqlite_session_factory() as session: + assert session.get(App, imported_id) is not None + permissions.append(imported_id) + return ["app.acl.view_layout"] + + def update_access(imported_id: str, access_mode: str) -> None: + assert not connections, "WebApp access initialization must follow transaction completion" + access_updates.append((imported_id, access_mode)) + + monkeypatch.setattr(EnterpriseConsoleAppAccess, "created_permissions", get_permissions) + monkeypatch.setattr(SystemFeatureService, "is_webapp_auth_enabled", lambda: True) + monkeypatch.setattr(EnterpriseService.WebAppAuth, "update_app_access_mode", update_access) + with app.test_request_context( + method="POST", + json={ + "mode": "yaml-content", + "yaml_content": "app: {}", + "app_id": app_id if overwrite else None, + }, + ): + if confirm: + api = controller.AppImportConfirmApi() + response, status_code = unwrap(api.post)(api, import_context, import_id="import-1") + else: + api = controller.AppImportApi() + response, status_code = unwrap(api.post)(api, import_context) + with sqlite_session_factory() as session: + assert (session.get(App, app_id) is not None) is (status != ImportStatus.FAILED) + completed = status in {ImportStatus.COMPLETED, ImportStatus.COMPLETED_WITH_WARNINGS} + assert permissions == ([app_id] if completed and not overwrite else []) + assert response["permission_keys"] == (["app.acl.view_layout"] if permissions else []) + assert access_updates == ([] if confirm else [(app_id, "private")]) + assert status_code == ( + 400 if status == ImportStatus.FAILED else 202 if status == ImportStatus.PENDING and not confirm else 200 + ) + + @staticmethod + @pytest.mark.usefixtures("app_query_services") + @pytest.mark.parametrize("missing", [False, True]) + def test_check_dependencies_releases_database_before_plugin_io( + app: Flask, + import_context: RequestContext, + sqlite_session_factory: sessionmaker[Session], + connections: set[object], + monkeypatch: pytest.MonkeyPatch, + missing: bool, + ) -> None: + app_id = str(uuid4()) + if not missing: + with sqlite_session_factory.begin() as session: + session.add( + App( + id=app_id, + tenant_id=import_context.active_workspace_id, + name="Imported", + mode=AppMode.CHAT, + enable_site=True, + enable_api=True, + ) + ) + monkeypatch.setattr(redis_client, "get", lambda _key: '{"dependencies":[]}') + called = [] + + def check(*, tenant_id: str, dependencies): + assert not connections + assert tenant_id == import_context.active_workspace_id + called.append(dependencies) + return [] + + monkeypatch.setattr(DependenciesAnalysisService, "get_leaked_dependencies", check) + api = controller.AppImportCheckDependenciesApi() + with app.test_request_context(): + if missing: + with pytest.raises(AppNotFoundError): + unwrap(api.get)(api, import_context, app_id=app_id) + assert called == [] + else: + response, code = unwrap(api.get)(api, import_context, app_id=app_id) + assert code == 200 + assert response == {"leaked_dependencies": []} + assert called == [[]] + + @staticmethod + @pytest.mark.parametrize("confirm", [False, True]) + def test_permission_errors_are_mapped_at_http_boundary( + app_query_services, app: Flask, monkeypatch, confirm: bool + ) -> None: + def denied(*_args, **_kwargs): + raise NoPermissionError("Import denied") + + monkeypatch.setattr(app_query_services.apps.console, "confirm_import" if confirm else "import_app", denied) + context = RequestContext("request", None, "actor", "workspace") + with app.test_request_context(method="POST", json={"mode": "yaml-content"}): + if confirm: + api = controller.AppImportConfirmApi() + with pytest.raises(Forbidden, match="Import denied"): + unwrap(api.post)(api, context, import_id="import-1") + else: + api = controller.AppImportApi() + with pytest.raises(Forbidden, match="Import denied"): + unwrap(api.post)(api, context) + + @staticmethod + def test_confirming_an_import_whose_draft_changed_is_a_conflict( + app_query_services, app: Flask, monkeypatch + ) -> None: + """A pending openapi import carries a draft_hash; a console confirm after the draft changed must not 500.""" + + def changed(*_args, **_kwargs): + raise DraftChangedError() + + monkeypatch.setattr(app_query_services.apps.console, "confirm_import", changed) + context = RequestContext("request", None, "actor", "workspace") + with app.test_request_context(method="POST"): + api = controller.AppImportConfirmApi() + with pytest.raises(DraftWorkflowNotSync) as exc_info: + unwrap(api.post)(api, context, import_id="import-1") + + assert exc_info.value.code == 409 + + @staticmethod + def test_import_response_schema_hides_creation_marker() -> None: + assert "app_created" not in controller.console_ns.models[Import.__name__].__schema__["properties"] + + +def _roster_archive() -> io.BytesIO: + source = io.BytesIO() + with zipfile.ZipFile(source, "w") as archive: + # Fixed timestamp: passing a filename string stamps the current local + # time (2-second ZIP resolution), so two archives built across a + # timestamp boundary differ byte-for-byte and flake the assertion below. + info = zipfile.ZipInfo("manifest.yaml", date_time=(1980, 1, 1, 0, 0, 0)) + archive.writestr(info, "format: dify.roster-agent") + source.seek(0) + return source + + +def _mock_download(monkeypatch: pytest.MonkeyPatch, content: bytes | None = None) -> list[str]: + calls: list[str] = [] + + def fetch(method: str, url: str, **_kwargs: object) -> httpx.Response: + calls.append(url) + return httpx.Response( + 200, + content=content if content is not None else _roster_archive().getvalue(), + request=httpx.Request(method, url), + ) + + monkeypatch.setattr(app_import_source.remote_fetcher, "make_request", fetch) + return calls + + +@dataclass(frozen=True) +class AgentImportCall: + tenant_id: str + account: Account + content: bytes + + +@dataclass +class Imports: + context: RequestContext + dsl_calls: list[AppImportParams] = field(default_factory=list) + packages: list[PreparedAppPackage] = field(default_factory=list) + agent_calls: list[AgentImportCall] = field(default_factory=list) + status: ImportStatus = ImportStatus.COMPLETED + warnings: list[DslImportWarning] = field(default_factory=list) + + def dsl( + self, context: RequestContext, params: AppImportParams, *, package: AppImportPackage | None = None + ) -> Import: + assert context == self.context + self.dsl_calls.append(params) + if package is not None: + assert isinstance(package, PreparedAppPackage) + assert not package.archive.closed + self.packages.append(package) + return Import(id="import-1", status=self.status, app_id="app-1") + + def agent(self, *, source: BinaryIO, tenant_id: str, account: Account) -> RosterAgentPackageImportResult: + self.agent_calls.append(AgentImportCall(tenant_id, account, source.read())) + return RosterAgentPackageImportResult(app_id="app-1", agent_id="agent-1", warnings=self.warnings) + + +def _post( + app: Flask, + imports: Imports, + *, + json: dict[str, object] | None = None, + data: dict[str, object] | None = None, +) -> tuple[dict[str, object], int]: + api = import_module.AppImportApi() + with app.test_request_context("/console/api/apps/imports", method="POST", json=json, data=data): + return cast(tuple[dict[str, object], int], unwrap(api.post)(api, imports.context)) + + +class TestRosterPackageRoutesCases: + @staticmethod + @pytest.fixture + def imports( + app_query_services: ControllerTestServices, + sqlite_session_factory: sessionmaker[Session], + monkeypatch: pytest.MonkeyPatch, + config_overrides: Callable[..., None], + ) -> Imports: + config_overrides(RBAC_ENABLED=False, DEPLOYMENT_EDITION="COMMUNITY") + actor, tenant = str(uuid4()), str(uuid4()) + account = Account(name="Importer", email="importer@example.com") + account.id = actor + workspace = Tenant(name="Workspace") + workspace.id = tenant + with sqlite_session_factory.begin() as session: + session.add_all( + [ + account, + workspace, + TenantAccountJoin(tenant_id=tenant, account_id=actor, role=TenantAccountRole.OWNER), + ] + ) + result = Imports(RequestContext("request", None, actor, tenant)) + transfers = app_query_services.apps.console._transfers + assert isinstance(transfers, AppTransferGateway) + monkeypatch.setattr(transfers, "import_dsl", result.dsl) + monkeypatch.setattr(transfers._agent_importer, "import_package", result.agent) + monkeypatch.setattr("services.app.console_gateway.SystemFeatureService.is_webapp_auth_enabled", lambda: False) + return result + + @staticmethod + def test_existing_import_route_dispatches_json_without_changing_payload(app: Flask, imports: Imports) -> None: + data, status = _post(app, imports, json={"mode": "yaml-content", "yaml_content": "app: {}"}) + assert status == 200 + assert data["status"] == "completed" + assert [params.model_dump() for params in imports.dsl_calls] == [ + AppImportParams(mode="yaml-content", yaml_content="app: {}").model_dump() + ] + assert imports.agent_calls == [] + + @staticmethod + @pytest.mark.parametrize( + "url", + [ + "https://example.com/download?id=123", + "https://example.com/agent.ifpkg", + " https://example.com/agent.IFPKG?token=secret#download ", + ], + ) + def test_package_url_uses_agent_import( + app: Flask, imports: Imports, monkeypatch: pytest.MonkeyPatch, url: str + ) -> None: + fetched = _mock_download(monkeypatch) + data, status = _post(app, imports, json={"mode": "yaml-url", "yaml_url": url}) + assert status == 200 + assert data["app_mode"] == "agent" + assert data["status"] == "completed" + assert len(fetched) == 1 + assert imports.dsl_calls == [] + (call,) = imports.agent_calls + assert call.tenant_id == imports.context.active_workspace_id + assert call.account.id == imports.context.account_id + assert call.account.is_admin_or_owner is True + + @staticmethod + @pytest.mark.parametrize( + "url", ["https://example.com/app.yaml", "https://example.com/download", "https://example.com/agent.ifpkg"] + ) + def test_yaml_url_keeps_dsl_import_without_package_retry( + app: Flask, imports: Imports, monkeypatch: pytest.MonkeyPatch, url: str + ) -> None: + imports.status = ImportStatus.FAILED + fetched = _mock_download(monkeypatch, b"app: {}") + _, status = _post( + app, imports, json={"mode": "yaml-url", "yaml_url": url, "app_id": "existing", "name": "Renamed"} + ) + assert status == 400 + assert ( + imports.dsl_calls[0].model_dump() + == AppImportParams( + mode="yaml-content", yaml_content="app: {}", app_id="existing", name="Renamed" + ).model_dump() + ) + assert len(fetched) == 1 + assert imports.agent_calls == [] + + @staticmethod + def test_url_without_any_import_permission_does_not_download( + app: Flask, imports: Imports, monkeypatch: pytest.MonkeyPatch, config_overrides: Callable[..., None] + ) -> None: + config_overrides(RBAC_ENABLED=True) + monkeypatch.setattr( + "services.app.console_gateway.rbac_service.RBACService.CheckAccess.check", lambda *_a, **_k: False + ) + fetched = _mock_download(monkeypatch) + with pytest.raises(Forbidden): + _post(app, imports, json={"mode": "yaml-url", "yaml_url": "https://example.com/download"}) + assert fetched == [] + + @staticmethod + @pytest.mark.parametrize("is_yaml", [False, True]) + def test_url_quota_is_enforced_after_content_detection( + app: Flask, + imports: Imports, + app_query_services: ControllerTestServices, + monkeypatch: pytest.MonkeyPatch, + config_overrides: Callable[..., None], + is_yaml: bool, + ) -> None: + config_overrides(DEPLOYMENT_EDITION="CLOUD") + transfers = app_query_services.apps.console._transfers + assert isinstance(transfers, AppTransferGateway) + monkeypatch.setattr(transfers, "import_dsl", AppTransferGateway.import_dsl.__get__(transfers)) + features = FeatureModel(apps=LimitationModel(size=1, limit=1)) + monkeypatch.setattr("services.app_dsl_service.FeatureService.get_features", lambda *_a, **_k: features) + _mock_download(monkeypatch, b"app: {mode: workflow}" if is_yaml else _roster_archive().getvalue()) + if is_yaml: + with pytest.raises(Forbidden, match="number of apps"): + _post(app, imports, json={"mode": "yaml-url", "yaml_url": "https://example.com/download"}) + assert imports.agent_calls == [] + assert imports.dsl_calls == [] + else: + data, status = _post(app, imports, json={"mode": "yaml-url", "yaml_url": "https://example.com/download"}) + assert status == 200 + assert data["app_mode"] == "agent" + + @staticmethod + def test_url_rejects_content_that_is_neither_yaml_nor_package( + app: Flask, imports: Imports, monkeypatch: pytest.MonkeyPatch + ) -> None: + fetched = _mock_download(monkeypatch, b"[invalid") + with pytest.raises(InvalidRosterAgentPackageError): + _post(app, imports, json={"mode": "yaml-url", "yaml_url": "https://example.com/download"}) + assert len(fetched) == 1 + + @staticmethod + @pytest.mark.parametrize("has_warning", [False, True]) + def test_existing_import_route_accepts_package_and_preserves_import_response( + app: Flask, imports: Imports, has_warning: bool + ) -> None: + imports.warnings = ( + [DslImportWarning(code="agent_skill_missing", path="skills.s_000001", message="Missing Skill")] + if has_warning + else [] + ) + data, status = _post(app, imports, data={"file": (_roster_archive(), "agent.ifpkg")}) + assert status == 200 + assert data["app_id"] == "app-1" + assert data["app_mode"] == "agent" + assert data["status"] == ("completed-with-warnings" if has_warning else "completed") + assert data["warnings"] == [warning.model_dump(mode="json") for warning in imports.warnings] + assert data["id"] + assert imports.agent_calls[0].content == _roster_archive().getvalue() + + @staticmethod + @pytest.mark.parametrize("from_url", [False, True]) + @pytest.mark.parametrize("denied", [RBACPermission.AGENT_CREATE, RBACPermission.AGENT_IMPORT_EXPORT_DSL]) + def test_package_import_checks_agent_permissions_before_materializing_resources( + app: Flask, + imports: Imports, + monkeypatch: pytest.MonkeyPatch, + config_overrides: Callable[..., None], + denied: RBACPermission, + from_url: bool, + ) -> None: + config_overrides(RBAC_ENABLED=True) + _mock_download(monkeypatch) + scenes: list[RBACPermission] = [] + + def check(*_args: object, scene: RBACPermission, **_kwargs: object) -> bool: + scenes.append(scene) + return scene != denied + + monkeypatch.setattr("services.app.console_gateway.rbac_service.RBACService.CheckAccess.check", check) + kwargs: dict[str, dict[str, object]] = ( + {"json": {"mode": "yaml-url", "yaml_url": "https://example.com/agent.ifpkg"}} + if from_url + else {"data": {"file": (_roster_archive(), "agent.ifpkg")}} + ) + with pytest.raises(Forbidden): + _post(app, imports, **kwargs) + assert denied in scenes + if not from_url: + assert RBACPermission.APP_IMPORT_EXPORT_DSL not in scenes + assert imports.agent_calls == [] + + @staticmethod + @pytest.mark.parametrize("from_url", [False, True]) + def test_package_import_rejects_overwrite( + app: Flask, imports: Imports, monkeypatch: pytest.MonkeyPatch, from_url: bool + ) -> None: + _mock_download(monkeypatch) + kwargs: dict[str, dict[str, object]] = ( + {"json": {"mode": "yaml-url", "yaml_url": "https://example.com/agent.ifpkg", "app_id": "existing"}} + if from_url + else {"data": {"app_id": "existing", "file": (_roster_archive(), "agent.ifpkg")}} + ) + with pytest.raises(InvalidRosterAgentPackageError, match="overwriting"): + _post(app, imports, **kwargs) + + @staticmethod + @pytest.mark.parametrize("from_url", [False, True]) + @pytest.mark.parametrize("app_id", [None, "existing"]) + @pytest.mark.parametrize("status", [ImportStatus.COMPLETED, ImportStatus.PENDING, ImportStatus.FAILED]) + def test_ordinary_package_import_uses_dsl_permissions_and_confirmation( + app: Flask, + imports: Imports, + monkeypatch: pytest.MonkeyPatch, + app_id: str | None, + status: ImportStatus, + from_url: bool, + ) -> None: + imports.status = status + dsl = "kind: app\nversion: 99.0.0\napp:\n mode: workflow\n" + with AppPackageService().export(dsl=dsl, name="Workflow") as package: + form: dict[str, object] = {"file": (package.archive, "workflow.ifpkg"), "name": "Renamed"} + if app_id: + form["app_id"] = app_id + if from_url: + _mock_download(monkeypatch, package.archive.read()) + data, code = _post( + app, + imports, + json={ + "mode": "yaml-url", + "yaml_url": "https://example.com/workflow.ifpkg", + "name": "Renamed", + "app_id": app_id, + }, + ) + else: + data, code = _post(app, imports, data=form) + assert code == {ImportStatus.FAILED: 400, ImportStatus.PENDING: 202}.get(status, 200) + assert len(imports.packages) == 1 + assert imports.packages[0].archive.closed + assert data["status"] == status + assert len(imports.dsl_calls) == 1 + assert ( + imports.dsl_calls[0].model_dump() + == AppImportParams(mode="yaml-content", yaml_content=dsl, name="Renamed", app_id=app_id).model_dump() + ) + assert imports.agent_calls == [] + + @staticmethod + def test_export_query_rejects_conflicting_version_selectors() -> None: + with pytest.raises(ValidationError, match="version_id and workflow_id cannot be used together"): + app_module.AppExportQuery.model_validate( + {"version_id": "11111111-1111-4111-8111-111111111111", "workflow_id": "workflow-1"} + ) + + @staticmethod + def test_template_import_resolves_catalog_source_and_preserves_overrides( + app: Flask, imports: Imports, app_query_services: ControllerTestServices, monkeypatch: pytest.MonkeyPatch + ) -> None: + transfers = app_query_services.apps.console._transfers + assert isinstance(transfers, AppTransferGateway) + template_id = UUID("11111111-1111-4111-8111-111111111111") + version_id = UUID("22222222-2222-4222-8222-222222222222") + source = RecommendedAgentPackageSource("source-tenant", "source-agent", version_id) + catalog = Mock() + catalog.get_package_source.return_value = source + transfers._recommended_packages = RecommendedAppPackageService(sources=catalog, exporter=Mock()) + import_template = Mock(return_value=RosterAgentPackageImportResult("new-app", "new-agent", [])) + monkeypatch.setattr(transfers._agent_importer, "import_template", import_template) + + data, status = _post( + app, + imports, + json={ + "mode": "template", + "template_id": str(template_id), + "version_id": str(version_id), + "name": "My Agent", + "app_id": "", + }, + ) + + assert status == 200 + assert data["app_id"] == "new-app" + assert data["app_mode"] == "agent" + catalog.get_package_source.assert_called_once_with(str(template_id), version_id) + import_template.assert_called_once() + assert import_template.call_args.kwargs["source"] == source + assert import_template.call_args.kwargs["tenant_id"] == imports.context.active_workspace_id + assert import_template.call_args.kwargs["account"].id == imports.context.account_id + assert import_template.call_args.kwargs["name"] == "My Agent" + + @staticmethod + def test_template_import_rejects_unlisted_source( + app: Flask, imports: Imports, app_query_services: ControllerTestServices, monkeypatch: pytest.MonkeyPatch + ) -> None: + transfers = app_query_services.apps.console._transfers + assert isinstance(transfers, AppTransferGateway) + catalog = Mock() + catalog.get_package_source.return_value = None + transfers._recommended_packages = RecommendedAppPackageService(sources=catalog, exporter=Mock()) + import_template = Mock() + monkeypatch.setattr(transfers._agent_importer, "import_template", import_template) + + with pytest.raises(import_module.RecommendedAppNotFoundHttpError): + _post( + app, + imports, + json={ + "mode": "template", + "template_id": "11111111-1111-4111-8111-111111111111", + "version_id": "22222222-2222-4222-8222-222222222222", + }, + ) + import_template.assert_not_called() + + @staticmethod + def test_package_url_import_preserves_overrides( + app: Flask, imports: Imports, app_query_services: ControllerTestServices, monkeypatch: pytest.MonkeyPatch + ) -> None: + transfers = app_query_services.apps.console._transfers + assert isinstance(transfers, AppTransferGateway) + import_url = Mock(return_value=RosterAgentPackageImportResult("new-app", "new-agent", [])) + monkeypatch.setattr(transfers._agent_importer, "import_package_url", import_url) + + data, status = _post( + app, + imports, + json={"mode": "ifpkg-url", "package_url": "https://example.com/agent.ifpkg", "name": "My Agent"}, + ) + + assert status == 200 + assert data["app_id"] == "new-app" + import_url.assert_called_once() + assert import_url.call_args.kwargs["url"] == "https://example.com/agent.ifpkg" + assert import_url.call_args.kwargs["name"] == "My Agent" + + @staticmethod + def test_package_url_retains_legacy_optional_selectors( + app: Flask, imports: Imports, app_query_services: ControllerTestServices, monkeypatch: pytest.MonkeyPatch + ) -> None: + transfers = app_query_services.apps.console._transfers + assert isinstance(transfers, AppTransferGateway) + import_url = Mock(return_value=RosterAgentPackageImportResult("new-app", "new-agent", [])) + monkeypatch.setattr(transfers._agent_importer, "import_package_url", import_url) + + data, status = _post( + app, + imports, + json={ + "mode": "ifpkg-url", + "package_url": "https://example.com/agent.ifpkg", + "app_id": "", + "template_id": "11111111-1111-4111-8111-111111111111", + "version_id": "22222222-2222-4222-8222-222222222222", + }, + ) + + assert status == 200 + assert data["app_id"] == "new-app" + import_url.assert_called_once() + assert import_url.call_args.kwargs["url"] == "https://example.com/agent.ifpkg" + + @staticmethod + @pytest.mark.parametrize("mode", ["template", "ifpkg-url"]) + def test_template_import_checks_permissions_before_accessing_source( + app: Flask, + imports: Imports, + app_query_services: ControllerTestServices, + monkeypatch: pytest.MonkeyPatch, + config_overrides: Callable[..., None], + mode: str, + ) -> None: + config_overrides(RBAC_ENABLED=True) + monkeypatch.setattr( + "services.app.console_gateway.rbac_service.RBACService.CheckAccess.check", lambda *_a, **_k: False + ) + transfers = app_query_services.apps.console._transfers + assert isinstance(transfers, AppTransferGateway) + importer = Mock() + monkeypatch.setattr(transfers, "_agent_importer", importer) + payload: dict[str, object] = ( + { + "mode": "template", + "template_id": "11111111-1111-4111-8111-111111111111", + "version_id": "22222222-2222-4222-8222-222222222222", + } + if mode == "template" + else {"mode": "ifpkg-url", "package_url": "https://example.com/agent.ifpkg"} + ) + with pytest.raises(Forbidden): + _post(app, imports, json=payload) + importer.assert_not_called() + + @staticmethod + @pytest.mark.parametrize( + "extra", + [ + {"app_id": "existing"}, + {"yaml_content": "app: {}"}, + ], + ) + def test_package_url_rejects_ambiguous_sources(app: Flask, imports: Imports, extra: dict[str, str]) -> None: + with pytest.raises(InvalidRosterAgentPackageError): + _post(app, imports, json={"mode": "ifpkg-url", "package_url": "https://example.com/agent.ifpkg", **extra}) + + @staticmethod + @pytest.mark.parametrize("quota", [False, True]) + def test_ordinary_package_import_checks_permission_and_quota( + app: Flask, + imports: Imports, + app_query_services: ControllerTestServices, + monkeypatch: pytest.MonkeyPatch, + config_overrides: Callable[..., None], + quota: bool, + ) -> None: + config_overrides(RBAC_ENABLED=not quota, DEPLOYMENT_EDITION="CLOUD" if quota else "COMMUNITY") + transfers = app_query_services.apps.console._transfers + assert isinstance(transfers, AppTransferGateway) + monkeypatch.setattr(transfers, "import_dsl", AppTransferGateway.import_dsl.__get__(transfers)) + monkeypatch.setattr( + "services.app.console_gateway.rbac_service.RBACService.CheckAccess.check", lambda *_a, **_k: False + ) + features = FeatureModel(apps=LimitationModel(size=1, limit=1)) + monkeypatch.setattr("services.app_dsl_service.FeatureService.get_features", lambda *_a, **_k: features) + with AppPackageService().export(dsl="kind: app\napp: {mode: workflow}\n", name="Workflow") as package: + with pytest.raises(Forbidden): + _post(app, imports, data={"file": (package.archive, "workflow.ifpkg")}) + assert imports.dsl_calls == [] + + +"""Quota failures retain their HTTP semantics across both app import surfaces.""" + + +class _EndpointView(Protocol): + __handler__: Callable[..., tuple[object, int]] + + +class TestAppImportQuotaCases: + @staticmethod + @pytest.mark.parametrize( + ("surface", "confirm"), + [("console", False), ("console", True), ("openapi", False), ("openapi", True), ("inner", False)], + ) + @pytest.mark.parametrize("failure", ["timeout", "upstream_503", "invalid_response", "quota_exhausted"]) + def test_import_quota_failure_http_status( + app_query_services: ControllerTestServices, + sqlite_session: Session, + monkeypatch: pytest.MonkeyPatch, + surface: str, + confirm: bool, + failure: str, + ) -> None: + apply_config_overrides( + monkeypatch, DEPLOYMENT_EDITION=DeploymentEdition.CLOUD, INNER_API=True, INNER_API_KEY="test-inner-api-key" + ) + monkeypatch.setattr("controllers.console.wraps._is_setup_completed", lambda: True) + account = Account(name="Importer", email="importer@example.com") + tenant = Tenant(name="Workspace") + account.id, tenant.id = str(uuid4()), str(uuid4()) + sqlite_session.add_all( + [ + account, + tenant, + TenantAccountJoin(tenant_id=tenant.id, account_id=account.id, role=TenantAccountRole.OWNER), + ] + ) + sqlite_session.commit() + context = RequestContext("request-1", None, account.id, tenant.id) + content = ( + f"version: {CURRENT_APP_DSL_VERSION}\napp: {{name: Imported, mode: workflow}}\n" + "workflow: {graph: {nodes: [], edges: []}, features: {}}\n" + ) + pending = PendingData( + tenant_id=tenant.id, account_id=account.id, import_mode="yaml-content", yaml_content=content + ) + monkeypatch.setattr(redis_client, "get", lambda _: pending.model_dump_json()) + + def unexpected_side_effect(*_args: object, **_kwargs: object) -> None: + pytest.fail("A failed quota check must preserve pending data and avoid database writes") + + monkeypatch.setattr(redis_client, "delete", unexpected_side_effect) + monkeypatch.setattr(AppDslService, "_create_or_update_app", unexpected_side_effect) + if failure == "quota_exhausted": + monkeypatch.setattr( + FeatureService, + "get_features", + lambda *_args, **_kwargs: FeatureModel(apps=LimitationModel(size=5, limit=5)), + ) + else: + + def failed_billing(*_args: object, **_kwargs: object) -> dict[str, str]: + if failure == "timeout": + raise httpx.ReadTimeout("Private billing timeout details") + if failure == "upstream_503": + raise _BillingHTTPStatusError("Private billing failure details", 503) + return {"subscription": "Private invalid billing response"} + + monkeypatch.setattr(BillingService, "_send_request", failed_billing) + + app = Flask(__name__) + app.config["TESTING"] = True + app.extensions["application_services"] = app_query_services + api = ExternalApi(app, error_body_formatter=OpenApiErrorFormatter() if surface == "openapi" else None) + + class AdmittedImport(Resource): + def post(self) -> tuple[object, int]: + if surface == "inner": + return inner.EnterpriseAppDSLImport().post(workspace_id=tenant.id) + if surface == "console": + if confirm: + controller = console.AppImportConfirmApi() + return unwrap(controller.post)(controller, context, import_id="pending-1") + controller = console.AppImportApi() + return unwrap(controller.post)(controller, context) + if confirm: + controller = openapi.AppDslImportConfirmApi() + return cast(_EndpointView, controller.post).__handler__( + controller, context, workspace_id=tenant.id, import_id="pending-1" + ) + controller = openapi.AppDslImportApi() + return cast(_EndpointView, controller.post).__handler__( + controller, + context, + workspace_id=tenant.id, + body=AppDslImportPayload(mode="yaml-content", yaml_content=content), + ) + + api.add_resource(AdmittedImport, "/imports") + response = app.test_client().post( + "/imports", + json={"mode": "yaml-content", "yaml_content": content, "creator_email": account.email}, + headers={"X-Inner-Api-Key": "test-inner-api-key"}, + ) + + assert response.status_code == (403 if failure == "quota_exhausted" else 500) + assert "Private" not in response.get_data(as_text=True) + assert sqlite_session.scalar(select(App.id).where(App.tenant_id == tenant.id)) is None diff --git a/api/tests/unit_tests/controllers/console/app/test_app_import_api.py b/api/tests/unit_tests/controllers/console/app/test_app_import_api.py deleted file mode 100644 index 5485070f512..00000000000 --- a/api/tests/unit_tests/controllers/console/app/test_app_import_api.py +++ /dev/null @@ -1,239 +0,0 @@ -"""Console imports finalize persistence before publishing permissions and access settings.""" - -from collections.abc import Iterator -from inspect import unwrap -from uuid import uuid4 - -import pytest -from flask import Flask -from sqlalchemy import Engine, event -from sqlalchemy.orm import Session, sessionmaker -from werkzeug.exceptions import Forbidden - -from controllers.console.app import app_import as controller -from controllers.console.app.error import AppNotFoundError, DraftWorkflowNotSync -from extensions.ext_redis import redis_client -from machinery.context import RequestContext -from models.account import Account, Tenant, TenantAccountJoin, TenantAccountRole -from models.model import App, AppMode -from services.app.console_gateway import EnterpriseConsoleAppAccess -from services.app_dsl_service import AppDslService, PendingData -from services.enterprise.enterprise_service import EnterpriseService -from services.entities.dsl_entities import Import, ImportStatus -from services.errors.app import DraftChangedError -from services.errors.base import NoPermissionError -from services.plugin.dependencies_analysis import DependenciesAnalysisService -from services.system_feature_service import SystemFeatureService - - -@pytest.fixture -def import_context(sqlite_session_factory: sessionmaker[Session]) -> RequestContext: - account_id, tenant_id = str(uuid4()), str(uuid4()) - account = Account(name="Importer", email="importer@example.com") - account.id = account_id - tenant = Tenant(name="Workspace") - tenant.id = tenant_id - with sqlite_session_factory.begin() as session: - session.add_all( - [ - account, - tenant, - TenantAccountJoin(tenant_id=tenant_id, account_id=account_id, role=TenantAccountRole.OWNER), - ] - ) - return RequestContext("import-request", None, account_id, tenant_id) - - -@pytest.fixture -def connections(sqlite_engine: Engine) -> Iterator[set[object]]: - checked_out: set[object] = set() - - def checkout(connection: object, *_args: object) -> None: - checked_out.add(connection) - - def checkin(connection: object, *_args: object) -> None: - checked_out.remove(connection) - - event.listen(sqlite_engine, "checkout", checkout) - event.listen(sqlite_engine, "checkin", checkin) - try: - yield checked_out - assert not checked_out - finally: - event.remove(sqlite_engine, "checkout", checkout) - event.remove(sqlite_engine, "checkin", checkin) - - -@pytest.mark.usefixtures("app_query_services") -@pytest.mark.parametrize("confirm", [False, True]) -@pytest.mark.parametrize("overwrite", [False, True]) -@pytest.mark.parametrize("status", list(ImportStatus)) -def test_import_transaction_and_response_contract( - app: Flask, - sqlite_session_factory: sessionmaker[Session], - import_context: RequestContext, - connections: set[object], - monkeypatch: pytest.MonkeyPatch, - config_overrides, - status: ImportStatus, - overwrite: bool, - confirm: bool, -) -> None: - config_overrides(RBAC_ENABLED=True, DEPLOYMENT_EDITION="COMMUNITY") - monkeypatch.setattr("services.app.console_gateway.rbac_service.RBACService.CheckAccess.check", lambda *a, **k: True) - app_id = str(uuid4()) - result = Import(id="import-1", status=status, app_id=app_id, app_created=not overwrite) - pending = PendingData( - tenant_id=import_context.active_workspace_id, - account_id=import_context.account_id, - import_mode="yaml-content", - yaml_content="app: {}", - app_id=app_id if overwrite else None, - ) - monkeypatch.setattr(redis_client, "get", lambda _key: pending.model_dump_json()) - - def persist(dsl: AppDslService, *, account: Account, **_kwargs) -> Import: - assert not connections, "Actor lookup must release its connection before starting the DSL operation" - assert account.id == import_context.account_id - assert account.current_tenant_id == import_context.active_workspace_id - dsl._session.add( - App( - id=app_id, - tenant_id=account.current_tenant_id, - name="Imported", - mode=AppMode.WORKFLOW, - enable_site=True, - enable_api=True, - ) - ) - dsl._session.flush() - return result - - monkeypatch.setattr(AppDslService, "confirm_import" if confirm else "import_app", persist) - permissions: list[str] = [] - access_updates: list[tuple[str, str]] = [] - - def get_permissions(_self, context: RequestContext, imported_id: str) -> list[str]: - assert not connections - assert context == import_context - with sqlite_session_factory() as session: - assert session.get(App, imported_id) is not None - permissions.append(imported_id) - return ["app.acl.view_layout"] - - def update_access(imported_id: str, access_mode: str) -> None: - assert not connections, "WebApp access initialization must follow transaction completion" - access_updates.append((imported_id, access_mode)) - - monkeypatch.setattr(EnterpriseConsoleAppAccess, "created_permissions", get_permissions) - monkeypatch.setattr(SystemFeatureService, "is_webapp_auth_enabled", lambda: True) - monkeypatch.setattr(EnterpriseService.WebAppAuth, "update_app_access_mode", update_access) - with app.test_request_context( - method="POST", - json={ - "mode": "yaml-content", - "yaml_content": "app: {}", - "app_id": app_id if overwrite else None, - }, - ): - if confirm: - api = controller.AppImportConfirmApi() - response, status_code = unwrap(api.post)(api, import_context, import_id="import-1") - else: - api = controller.AppImportApi() - response, status_code = unwrap(api.post)(api, import_context) - with sqlite_session_factory() as session: - assert (session.get(App, app_id) is not None) is (status != ImportStatus.FAILED) - completed = status in {ImportStatus.COMPLETED, ImportStatus.COMPLETED_WITH_WARNINGS} - assert permissions == ([app_id] if completed and not overwrite else []) - assert response["permission_keys"] == (["app.acl.view_layout"] if permissions else []) - assert access_updates == ([] if confirm else [(app_id, "private")]) - assert status_code == ( - 400 if status == ImportStatus.FAILED else 202 if status == ImportStatus.PENDING and not confirm else 200 - ) - - -@pytest.mark.usefixtures("app_query_services") -@pytest.mark.parametrize("missing", [False, True]) -def test_check_dependencies_releases_database_before_plugin_io( - app: Flask, - import_context: RequestContext, - sqlite_session_factory: sessionmaker[Session], - connections: set[object], - monkeypatch: pytest.MonkeyPatch, - missing: bool, -) -> None: - app_id = str(uuid4()) - if not missing: - with sqlite_session_factory.begin() as session: - session.add( - App( - id=app_id, - tenant_id=import_context.active_workspace_id, - name="Imported", - mode=AppMode.CHAT, - enable_site=True, - enable_api=True, - ) - ) - monkeypatch.setattr(redis_client, "get", lambda _key: '{"dependencies":[]}') - called = [] - - def check(*, tenant_id: str, dependencies): - assert not connections - assert tenant_id == import_context.active_workspace_id - called.append(dependencies) - return [] - - monkeypatch.setattr(DependenciesAnalysisService, "get_leaked_dependencies", check) - api = controller.AppImportCheckDependenciesApi() - with app.test_request_context(): - if missing: - with pytest.raises(AppNotFoundError): - unwrap(api.get)(api, import_context, app_id=app_id) - assert called == [] - else: - response, code = unwrap(api.get)(api, import_context, app_id=app_id) - assert code == 200 - assert response == {"leaked_dependencies": []} - assert called == [[]] - - -@pytest.mark.parametrize("confirm", [False, True]) -def test_permission_errors_are_mapped_at_http_boundary( - app_query_services, app: Flask, monkeypatch, confirm: bool -) -> None: - def denied(*_args, **_kwargs): - raise NoPermissionError("Import denied") - - monkeypatch.setattr(app_query_services.apps.console, "confirm_import" if confirm else "import_app", denied) - context = RequestContext("request", None, "actor", "workspace") - with app.test_request_context(method="POST", json={"mode": "yaml-content"}): - if confirm: - api = controller.AppImportConfirmApi() - with pytest.raises(Forbidden, match="Import denied"): - unwrap(api.post)(api, context, import_id="import-1") - else: - api = controller.AppImportApi() - with pytest.raises(Forbidden, match="Import denied"): - unwrap(api.post)(api, context) - - -def test_confirming_an_import_whose_draft_changed_is_a_conflict(app_query_services, app: Flask, monkeypatch) -> None: - """A pending openapi import carries a draft_hash; a console confirm after the draft changed must not 500.""" - - def changed(*_args, **_kwargs): - raise DraftChangedError() - - monkeypatch.setattr(app_query_services.apps.console, "confirm_import", changed) - context = RequestContext("request", None, "actor", "workspace") - with app.test_request_context(method="POST"): - api = controller.AppImportConfirmApi() - with pytest.raises(DraftWorkflowNotSync) as exc_info: - unwrap(api.post)(api, context, import_id="import-1") - - assert exc_info.value.code == 409 - - -def test_import_response_schema_hides_creation_marker() -> None: - assert "app_created" not in controller.console_ns.models[Import.__name__].__schema__["properties"] diff --git a/api/tests/unit_tests/controllers/console/app/test_app_response_models.py b/api/tests/unit_tests/controllers/console/app/test_app_response_models.py deleted file mode 100644 index 50fd30c65de..00000000000 --- a/api/tests/unit_tests/controllers/console/app/test_app_response_models.py +++ /dev/null @@ -1,598 +0,0 @@ -"""App request contracts and serialization of materialized response data.""" - -import json -from dataclasses import replace -from datetime import datetime -from pathlib import Path - -import pytest -import yaml -from flask import Flask -from pydantic import ValidationError -from sqlalchemy.orm import Session -from werkzeug.datastructures import MultiDict - -from controllers.console.app import app as controller -from fields import app_fields -from fields.app_model_config_response import AppModelConfigResponse -from libs.pagination import PaginatedResult -from models.account import Account -from models.enums import CustomizeTokenStrategy, TagType -from models.model import App, AppMode, AppModelConfig, IconType, Site, Tag, TagBinding -from models.workflow import Workflow, WorkflowType -from repositories.app.response import app_record -from services.app_service import AppResponseView - - -@pytest.fixture -def app_module(): - return controller - - -@pytest.fixture -def app_models(monkeypatch): - monkeypatch.setattr( - app_fields, "build_icon_url", lambda kind, icon: f"signed:{icon}" if kind == IconType.IMAGE else None - ) - return app_fields - - -def _ts(hour: int = 12) -> datetime: - return datetime(2024, 1, 1, hour, 0, 0) - - -TENANT_ID = "00000000-0000-0000-0000-000000000001" -APP_ID = "00000000-0000-0000-0000-000000000101" -ACCOUNT_ID = "00000000-0000-0000-0000-000000000201" -CONFIG_ID = "00000000-0000-0000-0000-000000000301" -WORKFLOW_ID = "00000000-0000-0000-0000-000000000401" -SITE_ID = "00000000-0000-0000-0000-000000000501" -TAG_ID = "00000000-0000-0000-0000-000000000601" - - -def _account(*, account_id: str = ACCOUNT_ID) -> Account: - account = Account(name="Creator", email=f"{account_id}@example.com") - account.id = account_id - return account - - -def _app( - *, - app_id: str = APP_ID, - tenant_id: str = TENANT_ID, - name: str = "My App", - description: str = "Summary", - mode: AppMode = AppMode.CHAT, - icon_type: IconType | None = IconType.IMAGE, - icon: str | None = "icon-key", - created_at: datetime | None = None, -) -> App: - timestamp = created_at or _ts() - app = App( - id=app_id, - tenant_id=tenant_id, - name=name, - description=description, - mode=mode, - icon_type=icon_type, - icon=icon, - icon_background="#fff", - enable_site=True, - enable_api=True, - max_active_requests=0, - ) - app.created_at = timestamp - app.updated_at = timestamp - return app - - -def _workflow(*, app_id: str = APP_ID, tenant_id: str = TENANT_ID) -> Workflow: - return Workflow( - id=WORKFLOW_ID, - tenant_id=tenant_id, - app_id=app_id, - type=WorkflowType.CHAT, - version=Workflow.VERSION_DRAFT, - graph=json.dumps({"nodes": [], "edges": []}), - features=json.dumps({}), - created_by=ACCOUNT_ID, - created_at=_ts(8), - updated_by=ACCOUNT_ID, - updated_at=_ts(9), - environment_variables=[], - conversation_variables=[], - ) - - -def _persist_response_graph(session: Session) -> App: - app = _app(description="Description") - app.created_by = ACCOUNT_ID - app.app_model_config_id = CONFIG_ID - app.workflow_id = WORKFLOW_ID - app.access_mode = "private" - app.create_user_name = "Creator" - app.has_draft_trigger = True - app.permission_keys = ["app.acl.view_layout"] - - model_config = AppModelConfig( - app_id=APP_ID, - model=json.dumps({"provider": "openai", "name": "gpt-4o"}), - pre_prompt="hello", - created_by=ACCOUNT_ID, - updated_by=ACCOUNT_ID, - ) - model_config.id = CONFIG_ID - model_config.created_at = _ts(9) - model_config.updated_at = _ts(10) - - site = Site( - id=SITE_ID, - app_id=APP_ID, - code="site-code", - title="Public Site", - icon_type=IconType.IMAGE, - icon="site-icon", - icon_background="#fff", - description="Site description", - default_language="en-US", - input_placeholder="Ask anything", - customize_token_strategy=CustomizeTokenStrategy.NOT_ALLOW, - ) - site.created_at = _ts(14) - site.updated_at = _ts(14) - tag = Tag(tenant_id=TENANT_ID, type=TagType.APP, name="Utilities", created_by=ACCOUNT_ID) - tag.id = TAG_ID - binding = TagBinding(tenant_id=TENANT_ID, tag_id=TAG_ID, target_id=APP_ID, created_by=ACCOUNT_ID) - - session.add_all([_account(), app, model_config, _workflow(), site, tag, binding]) - session.commit() - return app - - -def test_app_list_query_reads_repeated_tag_ids(app_module): - first_tag_id = "8c4ef3d1-58a1-4d94-8a1c-1c171d889e08" - second_tag_id = "3c39395b-6d1f-4030-8b17-eaa7cc85221c" - query_args = MultiDict( - [ - ("page", "1"), - ("limit", "30"), - ("tag_ids", first_tag_id), - ("tag_ids", second_tag_id), - ] - ) - - query = app_module.query_params_from_request( - app_module.AppListQuery, - list_fields=app_module.APP_LIST_QUERY_ARRAY_FIELDS, - args=query_args, - ) - - assert query.tag_ids == [first_tag_id, second_tag_id] - - -def test_app_list_query_reads_repeated_creator_ids(app_module): - first_creator_id = "9e8959cf-a67b-4d34-9906-1d687517b248" - second_creator_id = "1886f96a-5bf0-42bf-961d-8d2129049076" - query_args = MultiDict( - [ - ("page", "1"), - ("limit", "30"), - ("creator_ids", first_creator_id), - ("creator_ids", second_creator_id), - ] - ) - - query = app_module.query_params_from_request( - app_module.AppListQuery, - list_fields=app_module.APP_LIST_QUERY_ARRAY_FIELDS, - args=query_args, - ) - - assert query.creator_ids == [first_creator_id, second_creator_id] - - -def test_app_list_query_preserves_regular_query_params(app_module): - query_args = MultiDict( - [ - ("page", "2"), - ("limit", "50"), - ("mode", "chat"), - ("name", "Sales Copilot"), - ("is_created_by_me", "true"), - ] - ) - - query = app_module.query_params_from_request( - app_module.AppListQuery, - list_fields=app_module.APP_LIST_QUERY_ARRAY_FIELDS, - args=query_args, - ) - - assert query.page == 2 - assert query.limit == 50 - assert query.mode == "chat" - assert query.name == "Sales Copilot" - assert query.is_created_by_me is True - assert query.tag_ids is None - - -def test_app_list_query_normalizes_empty_repeated_tag_ids_to_none(app_module): - query_args = MultiDict( - [ - ("tag_ids", ""), - ("tag_ids", " "), - ] - ) - - query = app_module.query_params_from_request( - app_module.AppListQuery, - list_fields=app_module.APP_LIST_QUERY_ARRAY_FIELDS, - args=query_args, - ) - - assert query.tag_ids is None - - -def test_app_list_query_rejects_invalid_repeated_tag_id(app_module): - with pytest.raises(ValidationError): - app_module.query_params_from_request( - app_module.AppListQuery, - list_fields=app_module.APP_LIST_QUERY_ARRAY_FIELDS, - args=MultiDict([("tag_ids", "not-a-uuid")]), - ) - - -def test_app_list_query_rejects_invalid_repeated_creator_id(app_module): - with pytest.raises(ValidationError): - app_module.query_params_from_request( - app_module.AppListQuery, - list_fields=app_module.APP_LIST_QUERY_ARRAY_FIELDS, - args=MultiDict([("creator_ids", "not-a-uuid")]), - ) - - -def test_app_list_query_ignores_indexed_tag_ids(app_module): - tag_id = "8c4ef3d1-58a1-4d94-8a1c-1c171d889e08" - query_args = MultiDict( - [ - ("tag_ids[0]", tag_id), - ] - ) - - query = app_module.query_params_from_request( - app_module.AppListQuery, - list_fields=app_module.APP_LIST_QUERY_ARRAY_FIELDS, - args=query_args, - ) - - assert query.tag_ids is None - - -def test_app_list_query_accepts_single_repeated_tag_id(app_module): - tag_id = "8c4ef3d1-58a1-4d94-8a1c-1c171d889e08" - query = app_module.query_params_from_request( - app_module.AppListQuery, - list_fields=app_module.APP_LIST_QUERY_ARRAY_FIELDS, - args=MultiDict([("tag_ids", tag_id)]), - ) - - assert query.tag_ids == [tag_id] - - -def test_create_app_endpoint_rejects_agent_mode(app_module): - with pytest.raises(ValidationError): - app_module.CreateAppPayload.model_validate({"name": "Iris", "mode": "agent", "description": "Agent app"}) - - -def test_app_partial_serialization_uses_aliases(app_models, sqlite_session: Session): - app_partial = app_models.AppPartial - app_obj = _persist_response_graph(sqlite_session) - app_obj.description = "Prompt snippet" - - serialized = app_partial.model_validate( - replace(app_record(app_obj, session=sqlite_session), permission_keys=app_obj.permission_keys), - from_attributes=True, - ).model_dump(mode="json") - - assert serialized["description"] == "Prompt snippet" - assert serialized["mode"] == "chat" - assert serialized["icon_url"] == "signed:icon-key" - assert serialized["created_at"] == int(app_obj.created_at.timestamp()) - assert serialized["updated_at"] == int(app_obj.updated_at.timestamp()) - assert serialized["model_config"]["model"] == {"provider": "openai", "name": "gpt-4o"} - assert serialized["workflow"]["id"] == WORKFLOW_ID - assert serialized["tags"][0]["name"] == "Utilities" - assert serialized["permission_keys"] == ["app.acl.view_layout"] - assert "role" not in serialized - - -def test_app_detail_with_site_includes_nested_serialization( - app: Flask, - app_models, - monkeypatch: pytest.MonkeyPatch, - sqlite_session: Session, -): - app_obj = _persist_response_graph(sqlite_session) - app_obj.name = "Detailed App" - app_obj.description = "Desc" - app_obj.mode = AppMode.ADVANCED_CHAT - app_obj.icon = "detail-icon" - app_obj.icon_background = "#123456" - app_obj.use_icon_as_answer_icon = True - app_obj.max_active_requests = 5 - app_obj.access_mode = "public" - app_obj.permission_keys = ["app.acl.view_layout", "app.acl.edit"] - model_config = sqlite_session.get(AppModelConfig, CONFIG_ID) - assert model_config is not None - model_config.opening_statement = "hi" - model_config.retriever_resource = json.dumps({"enabled": True}) - monkeypatch.setattr( - "repositories.app.response.load_annotation_reply_config", lambda _session, _app_id: {"enabled": False} - ) - - with app.test_request_context("/"): - response = controller._app_detail_response( - replace( - app_record(app_obj, session=sqlite_session, projection="detail-with-site"), - permission_keys=app_obj.permission_keys, - ), - ) - assert isinstance(response, app_models.AppDetailWithSite) - serialized = response.model_dump(mode="json") - - assert serialized["icon_url"] == "signed:detail-icon" - assert serialized["model_config"]["retriever_resource"] == {"enabled": True} - assert serialized["deleted_tools"] == [] - assert serialized["site"]["icon_url"] == "signed:site-icon" - assert serialized["site"]["input_placeholder"] == "Ask anything" - assert serialized["site"]["created_at"] == int(_ts(14).timestamp()) - assert serialized["permission_keys"] == ["app.acl.view_layout", "app.acl.edit"] - assert serialized["bound_agent_id"] is None - assert "role" not in serialized - - -def test_app_response_view_uses_the_caller_session_for_query_backed_fields( - monkeypatch: pytest.MonkeyPatch, sqlite_session: Session -): - app_obj = _persist_response_graph(sqlite_session) - decoy_tenant_id = "00000000-0000-0000-0000-000000000002" - decoy_tag = Tag(tenant_id=decoy_tenant_id, type=TagType.APP, name="Decoy", created_by=ACCOUNT_ID) - decoy_tag.id = "00000000-0000-0000-0000-000000000602" - sqlite_session.add_all( - [ - decoy_tag, - TagBinding( - tenant_id=decoy_tenant_id, - tag_id=decoy_tag.id, - target_id=APP_ID, - created_by=ACCOUNT_ID, - ), - ] - ) - sqlite_session.commit() - calls = [] - - def load_annotation_reply(session, app_id): - calls.append((session, app_id)) - return {"enabled": False} - - monkeypatch.setattr("services.app_service.load_annotation_reply_config", load_annotation_reply) - - view = AppResponseView(app_obj, session=sqlite_session) - site = view.site - workflow = view.workflow - model_config = view.app_model_config - - assert view.desc_or_prompt == "Description" - assert site is not None - assert site.id == SITE_ID - assert workflow is not None - assert workflow.id == WORKFLOW_ID - assert view.bound_agent_id is None - assert view.mode_compatible_with_agent == "chat" - assert view.deleted_tools == [] - assert [tag.name for tag in view.tags] == ["Utilities"] - assert view.author_name == "Creator" - assert model_config is not None - assert model_config.annotation_reply_dict == {"enabled": False} - assert calls == [(sqlite_session, APP_ID)] - - -def test_app_pagination_aliases_per_page_and_has_next(app_models, sqlite_session: Session): - app_pagination = app_models.AppPagination - item_one = _app( - app_id="00000000-0000-0000-0000-000000000110", - name="Paginated One", - icon="first-icon", - created_at=_ts(15), - ) - item_one.permission_keys = ["app.acl.edit"] - item_two = _app( - app_id="00000000-0000-0000-0000-000000000111", - name="Paginated Two", - mode=AppMode.AGENT_CHAT, - icon_type=IconType.EMOJI, - icon="🙂", - created_at=_ts(16), - ) - pagination = PaginatedResult( - page=2, - per_page=10, - total=50, - items=[app_record(item, session=sqlite_session) for item in (item_one, item_two)], - ) - - serialized = app_pagination.model_validate( - pagination, - from_attributes=True, - ).model_dump(mode="json") - - assert serialized["page"] == 2 - assert serialized["limit"] == 10 - assert serialized["has_more"] is True - assert len(serialized["data"]) == 2 - assert serialized["data"][0]["icon_url"] == "signed:first-icon" - assert serialized["data"][1]["icon_url"] is None - - -@pytest.mark.parametrize("strategy", [None, "cot", "function-calling"]) -def test_app_detail_preserves_historical_agent_strategy( - app: Flask, monkeypatch: pytest.MonkeyPatch, sqlite_session: Session, strategy: str | None -): - app_obj = _persist_response_graph(sqlite_session) - config = sqlite_session.get(AppModelConfig, CONFIG_ID) - assert config is not None - agent_mode = {"enabled": True, "tools": []} - if strategy is not None: - agent_mode["strategy"] = strategy - config.agent_mode = json.dumps(agent_mode) - sqlite_session.commit() - monkeypatch.setattr( - "repositories.app.response.load_annotation_reply_config", lambda _session, _app_id: {"enabled": False} - ) - - with app.test_request_context("/"): - response = controller._app_detail_response( - app_record(app_obj, session=sqlite_session, projection="detail-with-site") - ).model_dump(mode="json") - - assert response["model_config"]["agent_mode"] == agent_mode - - -def test_app_detail_preserves_historical_agent_mode_without_tools( - app: Flask, monkeypatch: pytest.MonkeyPatch, sqlite_session: Session -): - app_obj = _persist_response_graph(sqlite_session) - app_obj.mode = AppMode.AGENT_CHAT - config = sqlite_session.get(AppModelConfig, CONFIG_ID) - assert config is not None - agent_mode = {"enabled": True, "max_iteration": 5, "strategy": "function_call"} - config.agent_mode = json.dumps(agent_mode) - sqlite_session.commit() - monkeypatch.setattr( - "repositories.app.response.load_annotation_reply_config", lambda _session, _app_id: {"enabled": False} - ) - - with app.test_request_context("/"): - response = controller._app_detail_response( - app_record(app_obj, session=sqlite_session, projection="detail-with-site") - ).model_dump(mode="json") - - assert response["model_config"]["agent_mode"] == agent_mode - - -def test_app_model_config_response_round_trips_persisted_dataset_and_external_tool_fields( - app: Flask, monkeypatch: pytest.MonkeyPatch, sqlite_session: Session -): - app_obj = _persist_response_graph(sqlite_session) - config = sqlite_session.get(AppModelConfig, CONFIG_ID) - assert config is not None - dataset_configs = { - "retrieval_model": "multiple", - "reranking_enable": True, - "reranking_enabled": True, - "metadata_filtering_conditions": { - "logical_operator": "and", - "conditions": [ - { - "id": "condition-1", - "metadata_id": "metadata-1", - "name": "category", - "comparison_operator": "is", - "value": "books", - } - ], - }, - "extension_setting": {"value": 1}, - } - external_form = { - "external_data_tool": { - "variable": "customer", - "label": "Customer", - "type": "api", - "enabled": True, - "icon": "icon-key", - "icon_background": "#fff", - "config": {"provider": "crm"}, - "extension_setting": {"value": 2}, - } - } - config.dataset_configs = json.dumps(dataset_configs) - config.user_input_form = json.dumps([external_form]) - sqlite_session.commit() - monkeypatch.setattr( - "repositories.app.response.load_annotation_reply_config", lambda _session, _app_id: {"enabled": False} - ) - - with app.test_request_context("/"): - response = controller._app_detail_response( - app_record(app_obj, session=sqlite_session, projection="detail-with-site") - ).model_dump(mode="json") - - model_config = response["model_config"] - assert model_config["dataset_configs"] == dataset_configs - assert model_config["user_input_form"] == [external_form] - - reimported = AppModelConfig(app_id=APP_ID) - reimported.from_model_config_dict(model_config) - assert reimported.dataset_configs_dict == dataset_configs - assert reimported.user_input_form_list == [external_form] - - -def test_recommended_app_model_configs_keep_all_published_fields(): - catalog_path = Path(__file__).resolve().parents[5] / "constants" / "recommended_apps.json" - catalog = json.loads(catalog_path.read_text()) - - def assert_fields_preserved(source, serialized): - if isinstance(source, dict): - for key, value in source.items(): - assert key in serialized - assert_fields_preserved(value, serialized[key]) - elif isinstance(source, list): - assert len(source) == len(serialized) - for value, serialized_value in zip(source, serialized, strict=True): - assert_fields_preserved(value, serialized_value) - else: - assert serialized == source - - for entry in catalog["app_details"].values(): - source = yaml.safe_load(entry["export_data"]).get("model_config") or {} - if not source: - continue - payload = {**source, "created_by": None, "created_at": 1, "updated_by": None, "updated_at": 1} - response = AppModelConfigResponse.model_validate(payload).model_dump(mode="json") - assert_fields_preserved(source, response) - - -def test_app_detail_rejects_invalid_known_configuration_fields( - app: Flask, monkeypatch: pytest.MonkeyPatch, sqlite_session: Session -): - app_obj = _persist_response_graph(sqlite_session) - monkeypatch.setattr( - "repositories.app.response.load_annotation_reply_config", lambda _session, _app_id: {"enabled": False} - ) - with app.test_request_context("/"): - detail = controller._app_detail_response( - app_record(app_obj, session=sqlite_session, projection="detail-with-site") - ).model_dump(mode="json") - - for field in ("mode", "permission_keys", "site", "model_config"): - with pytest.raises(ValidationError): - app_fields.AppDetailWithSite.model_validate({key: value for key, value in detail.items() if key != field}) - - with pytest.raises(ValidationError): - app_fields.AppDetailWithSite.model_validate({**detail, "mode": "unsupported"}) - - config = detail["model_config"] - for invalid in ( - {"speech_to_text": {"enabled": "not-a-boolean"}}, - {"agent_mode": {"enabled": True, "strategy": "unsupported", "tools": []}}, - {"agent_mode": {"enabled": True, "tools": "not-a-list"}}, - {"agent_mode": {"enabled": True, "tools": [{"provider_id": "only-id"}]}}, - {"dataset_configs": {"retrieval_model": "unsupported"}}, - {"dataset_configs": {"retrieval_model": "multiple", "reranking_mode": "unsupported"}}, - {"model": {"provider": "openai", "name": "gpt-4o", "mode": "unsupported"}}, - {"user_input_form": [{"unsupported": {"label": "Name", "variable": "name"}}]}, - ): - with pytest.raises(ValidationError): - app_fields.AppDetailWithSite.model_validate({**detail, "model_config": {**config, **invalid}}) diff --git a/api/tests/unit_tests/controllers/console/app/test_audio.py b/api/tests/unit_tests/controllers/console/app/test_audio.py index 246f17f77c3..ce06e715f0d 100644 --- a/api/tests/unit_tests/controllers/console/app/test_audio.py +++ b/api/tests/unit_tests/controllers/console/app/test_audio.py @@ -1,16 +1,20 @@ from __future__ import annotations import io +from collections.abc import Callable, Iterator from inspect import unwrap from unittest.mock import patch from uuid import UUID import pytest from flask import Flask, Response +from sqlalchemy import Engine, event from sqlalchemy.orm import Session from werkzeug.datastructures import FileStorage from werkzeug.exceptions import Forbidden, InternalServerError, UnprocessableEntity +from controllers.console import wraps +from controllers.console.app import audio from controllers.console.app import audio as audio_module from controllers.console.app.audio import ( AgentChatMessageAudioApi, @@ -36,9 +40,13 @@ from controllers.console.app.error import ( ) from core.errors.error import ModelCurrentlyNotSupportError, ProviderTokenNotInitError, QuotaExceededError from graphon.model_runtime.errors.invoke import InvokeError -from models import Account, App, AppMode -from models.agent import AgentConfigDraftType +from models import Account +from models import App as _audio_App +from models import AppMode as _audio_AppMode +from models.agent import Agent, AgentConfigDraftType, AgentKind, AgentScope, AgentSource, AgentStatus from models.agent_config_entities import AgentSoulConfig +from models.model import App as _agent_audio_readonly_App +from models.model import AppMode as _agent_audio_readonly_AppMode from services.agent.composer_service import AgentComposerService from services.agent.errors import AgentNotFoundError, AgentVersionNotFoundError from services.app_ref_service import AppRef, MessageRef @@ -51,14 +59,210 @@ from services.errors.audio import ( SpeechToTextDisabledServiceError, UnsupportedAudioTypeServiceError, ) -from tests.unit_tests.model_factories import make_account, make_app +from tests.helpers.unit.audio_runtime_fixtures import AudioRuntimeObservations +from tests.helpers.unit.model_factories import make_account, make_app + +AGENT_ID = UUID("019ef3d2-b24c-7803-b428-18b5ee8fb853") + + +def _agent(scope: AgentScope = AgentScope.WORKFLOW_ONLY) -> Agent: + workflow_only = scope == AgentScope.WORKFLOW_ONLY + return Agent( + id=str(AGENT_ID), + tenant_id="tenant-1", + name="Audio agent", + agent_kind=AgentKind.DIFY_AGENT, + scope=scope, + source=AgentSource.WORKFLOW if workflow_only else AgentSource.AGENT_APP, + status=AgentStatus.ACTIVE, + app_id="parent-app-1" if workflow_only else "runtime-app-1", + backing_app_id="runtime-app-1" if workflow_only else None, + workflow_id="workflow-1" if workflow_only else None, + workflow_node_id="node-1" if workflow_only else None, + created_by="account-1", + updated_by="account-1", + ) + + +class TestAgentAudioReadonlyCases: + @staticmethod + @pytest.fixture(autouse=True) + def audio_identity(monkeypatch: pytest.MonkeyPatch, config_overrides: Callable[..., None]) -> None: + config_overrides(LOGIN_DISABLED=True, RBAC_ENABLED=True) + monkeypatch.setattr(wraps, "_is_setup_completed", lambda: True) + monkeypatch.setattr(wraps, "current_account_with_tenant", lambda: (make_account(), "tenant-1")) + monkeypatch.setattr(audio, "enforce_rbac_checks", lambda **_kwargs: None) + + @staticmethod + @pytest.fixture + def persistence_events(sqlite_engine: Engine) -> Iterator[list[str]]: + observed: list[str] = [] + + def record_statement( + _connection: object, + _cursor: object, + statement: str, + _parameters: object, + _context: object, + _executemany: bool, + ) -> None: + operation = statement.lstrip().split(maxsplit=1)[0].upper() + if operation in {"INSERT", "UPDATE", "DELETE"}: + observed.append(operation) + + def record_commit(_session: Session) -> None: + observed.append("COMMIT") + + event.listen(sqlite_engine, "before_cursor_execute", record_statement) + event.listen(Session, "before_commit", record_commit) + try: + yield observed + finally: + event.remove(sqlite_engine, "before_cursor_execute", record_statement) + event.remove(Session, "before_commit", record_commit) + + @staticmethod + @pytest.fixture(params=["voices", "preview"]) + def request_audio( + request: pytest.FixtureRequest, app: Flask, audio_runtime: AudioRuntimeObservations + ) -> tuple[Callable[[], object], AudioRuntimeObservations]: + if request.param == "voices": + audio_runtime.voices = [{"name": "Echo", "value": "echo"}] + + def invoke() -> object: + with app.test_request_context(f"/console/api/agent/{AGENT_ID}/text-to-audio/voices?language=en-US"): + return audio.AgentTextToSpeechVoicesApi().get(agent_id=AGENT_ID) + + else: + audio_runtime.tts_result = b"ID3audio" + + def invoke() -> object: + with app.test_request_context( + f"/console/api/agent/{AGENT_ID}/text-to-audio", + method="POST", + json={"text": "Preview this voice", "voice": "echo"}, + ): + response = audio.AgentChatMessageTextApi().post(agent_id=AGENT_ID) + assert response is not None + try: + assert response.get_data() == b"ID3audio" + assert response.content_type == "audio/mpeg" + return response + finally: + response.close() + + return invoke, audio_runtime + + @staticmethod + @pytest.mark.parametrize("scope", [AgentScope.ROSTER, AgentScope.WORKFLOW_ONLY]) + def test_audio_uses_existing_runtime_state_without_writes( + sqlite_session: Session, + request_audio: tuple[Callable[[], object], AudioRuntimeObservations], + persistence_events: list[str], + scope: AgentScope, + ) -> None: + agent = _agent(scope) + runtime_app = make_app(app_id="runtime-app-1", mode=_agent_audio_readonly_AppMode.AGENT) + sqlite_session.add_all([agent, runtime_app]) + sqlite_session.commit() + persistence_events.clear() + invoke, provider = request_audio + + response = invoke() + + assert len(provider.resolutions) == 1 + assert response is not None + assert persistence_events == [] + + @staticmethod + def test_audio_does_not_materialize_a_missing_workflow_backing_app( + sqlite_session: Session, + request_audio: tuple[Callable[[], object], AudioRuntimeObservations], + persistence_events: list[str], + ) -> None: + agent = _agent() + agent.backing_app_id = None + parent_app = make_app(app_id="parent-app-1", mode=_agent_audio_readonly_AppMode.WORKFLOW) + sqlite_session.add_all([agent, parent_app]) + sqlite_session.commit() + persistence_events.clear() + invoke, provider = request_audio + + with pytest.raises(AgentNotFoundError): + invoke() + + assert provider.resolutions == [] + assert provider.daemon_calls == [] + assert persistence_events == [] + sqlite_session.refresh(agent) + assert agent.backing_app_id is None + assert sqlite_session.get(_agent_audio_readonly_App, "runtime-app-1") is None + + @staticmethod + @pytest.mark.parametrize("foreign_owner", ["agent", "app"]) + def test_audio_scopes_both_agent_and_runtime_app_to_the_current_tenant( + sqlite_session: Session, + request_audio: tuple[Callable[[], object], AudioRuntimeObservations], + persistence_events: list[str], + foreign_owner: str, + ) -> None: + agent = _agent() + runtime_app = make_app(app_id="runtime-app-1", mode=_agent_audio_readonly_AppMode.AGENT) + if foreign_owner == "agent": + agent.tenant_id = "tenant-2" + else: + runtime_app.tenant_id = "tenant-2" + sqlite_session.add_all([agent, runtime_app]) + sqlite_session.commit() + persistence_events.clear() + invoke, provider = request_audio + + with pytest.raises(AgentNotFoundError): + invoke() + + assert provider.resolutions == [] + assert provider.daemon_calls == [] + assert persistence_events == [] + + @staticmethod + @pytest.mark.parametrize( + ("provider_error", "expected_error"), + [ + (ProviderTokenNotInitError("No TTS model configured"), ProviderNotInitializeError), + (QuotaExceededError(), ProviderQuotaExceededError), + (ModelCurrentlyNotSupportError(), ProviderModelCurrentlyNotSupportError), + (InvokeError("TTS provider unavailable"), CompletionRequestError), + (ValueError("Unknown voice"), ValueError), + (RuntimeError("Unexpected provider failure"), InternalServerError), + ], + ) + def test_provider_failures_preserve_error_contract_and_readonly_state( + sqlite_session: Session, + request_audio: tuple[Callable[[], object], AudioRuntimeObservations], + persistence_events: list[str], + provider_error: Exception, + expected_error: type[Exception], + ) -> None: + sqlite_session.add_all( + [_agent(AgentScope.ROSTER), make_app(app_id="runtime-app-1", mode=_agent_audio_readonly_AppMode.AGENT)] + ) + sqlite_session.commit() + persistence_events.clear() + invoke, provider = request_audio + provider.lookup_error = provider_error + + with pytest.raises(expected_error): + invoke() + + assert len(provider.resolutions) == 1 + assert persistence_events == [] def _file_data(): return FileStorage(stream=io.BytesIO(b"audio"), filename="audio.wav", content_type="audio/wav") -def _app(*, app_id: str = "a1", tenant_id: str = "tenant-1") -> App: +def _app(*, app_id: str = "a1", tenant_id: str = "tenant-1") -> _audio_App: return make_app(app_id=app_id, tenant_id=tenant_id, name="Audio app", icon_type=None, max_active_requests=0) @@ -66,159 +270,122 @@ def _account(account_id: str = "account-1") -> Account: return make_account(account_id=account_id, name="Audio account", email=f"{account_id}@example.com") -def test_console_audio_api_success(app: Flask, monkeypatch: pytest.MonkeyPatch) -> None: - monkeypatch.setattr(AudioService, "transcript_asr", lambda **_kwargs: {"text": "ok"}) - api = ChatMessageAudioApi() - handler = unwrap(api.post) - app_model = _app() +class TestAudioCases: + @staticmethod + def test_console_audio_api_success(app: Flask, monkeypatch: pytest.MonkeyPatch) -> None: + monkeypatch.setattr(AudioService, "transcript_asr", lambda **_kwargs: {"text": "ok"}) + api = ChatMessageAudioApi() + handler = unwrap(api.post) + app_model = _app() - with app.test_request_context("/console/api/apps/app/audio-to-text", method="POST", data={"file": _file_data()}): - response = handler(api, app_model=app_model) + with app.test_request_context( + "/console/api/apps/app/audio-to-text", method="POST", data={"file": _file_data()} + ): + response = handler(api, app_model=app_model) - assert response == {"text": "ok"} + assert response == {"text": "ok"} + @staticmethod + def test_console_audio_api_accepts_published_agent_apps() -> None: + assert _audio_AppMode.AGENT in audio_module._CONSOLE_AUDIO_TRANSCRIPT_APP_MODES -def test_console_audio_api_accepts_published_agent_apps() -> None: - assert AppMode.AGENT in audio_module._CONSOLE_AUDIO_TRANSCRIPT_APP_MODES + @staticmethod + def test_agent_console_audio_api_uses_agent_draft( + app: Flask, monkeypatch: pytest.MonkeyPatch, unbound_session: Session + ) -> None: + agent_id = UUID("019ef3d2-b24c-7803-b428-18b5ee8fb853") + app_model = _app(app_id="backing-app-1") + agent_soul = AgentSoulConfig.model_validate({"app_features": {"speech_to_text": {"enabled": True}}}) + calls: dict[str, object] = {} + def resolve_agent_runtime_app_model(**kwargs): + calls["resolver"] = kwargs + return app_model -def test_agent_console_audio_api_uses_agent_draft( - app: Flask, monkeypatch: pytest.MonkeyPatch, unbound_session: Session -) -> None: - agent_id = UUID("019ef3d2-b24c-7803-b428-18b5ee8fb853") - app_model = _app(app_id="backing-app-1") - agent_soul = AgentSoulConfig.model_validate({"app_features": {"speech_to_text": {"enabled": True}}}) - calls: dict[str, object] = {} + def load_agent_soul_for_debug(**kwargs): + calls["draft"] = kwargs + return agent_soul - def resolve_agent_runtime_app_model(**kwargs): - calls["resolver"] = kwargs - return app_model + def transcript_agent_asr(**kwargs): + calls["asr"] = kwargs + return {"text": "agent transcript"} - def load_agent_soul_for_debug(**kwargs): - calls["draft"] = kwargs - return agent_soul + def enforce_rbac_checks(**kwargs): + calls["rbac"] = kwargs - def transcript_agent_asr(**kwargs): - calls["asr"] = kwargs - return {"text": "agent transcript"} + monkeypatch.setattr(audio_module, "resolve_agent_runtime_app_model", resolve_agent_runtime_app_model) + monkeypatch.setattr(audio_module, "enforce_rbac_checks", enforce_rbac_checks) + monkeypatch.setattr(AgentComposerService, "load_agent_soul_for_debug", load_agent_soul_for_debug) + monkeypatch.setattr(AudioService, "transcript_agent_asr", transcript_agent_asr) - def enforce_rbac_checks(**kwargs): - calls["rbac"] = kwargs + api = AgentChatMessageAudioApi() + handler = unwrap(api.post) + session = unbound_session + current_user = _account() + with app.test_request_context( + f"/console/api/agent/{agent_id}/audio-to-text", + method="POST", + data={"file": _file_data(), "draft_type": "debug_build"}, + ): + response = handler( + api, + session=session, + current_tenant_id="tenant-1", + current_user=current_user, + agent_id=agent_id, + ) - monkeypatch.setattr(audio_module, "resolve_agent_runtime_app_model", resolve_agent_runtime_app_model) - monkeypatch.setattr(audio_module, "enforce_rbac_checks", enforce_rbac_checks) - monkeypatch.setattr(AgentComposerService, "load_agent_soul_for_debug", load_agent_soul_for_debug) - monkeypatch.setattr(AudioService, "transcript_agent_asr", transcript_agent_asr) + assert response == {"text": "agent transcript"} + assert calls["resolver"] == {"session": session, "tenant_id": "tenant-1", "agent_id": agent_id} + rbac_call = calls["rbac"] + assert rbac_call["tenant_id"] == "tenant-1" + assert rbac_call["account_id"] == "account-1" + assert rbac_call["path_args"] == {"app_id": "backing-app-1"} + (rbac_check,) = rbac_call["checks"] + assert rbac_check.scene is audio_module.RBACPermission.APP_TEST_AND_RUN + assert isinstance(rbac_check.locator, audio_module.PlainApp) + assert calls["draft"] == { + "tenant_id": "tenant-1", + "agent_id": str(agent_id), + "account_id": "account-1", + "draft_type": AgentConfigDraftType.DEBUG_BUILD, + "session": session, + } + assert calls["asr"] == { + "app_model": app_model, + "agent_soul": agent_soul, + "file": calls["asr"]["file"], + "session": session, + "end_user": None, + } - api = AgentChatMessageAudioApi() - handler = unwrap(api.post) - session = unbound_session - current_user = _account() - with app.test_request_context( - f"/console/api/agent/{agent_id}/audio-to-text", - method="POST", - data={"file": _file_data(), "draft_type": "debug_build"}, - ): - response = handler( - api, - session=session, - current_tenant_id="tenant-1", - current_user=current_user, - agent_id=agent_id, + @staticmethod + def test_agent_console_audio_api_defaults_to_normal_draft( + app: Flask, monkeypatch: pytest.MonkeyPatch, unbound_session: Session + ) -> None: + agent_id = UUID("019ef3d2-b24c-7803-b428-18b5ee8fb853") + captured: dict[str, object] = {} + monkeypatch.setattr( + audio_module, + "resolve_agent_runtime_app_model", + lambda **_kwargs: _app(app_id="backing-app-1"), ) - assert response == {"text": "agent transcript"} - assert calls["resolver"] == {"session": session, "tenant_id": "tenant-1", "agent_id": agent_id} - rbac_call = calls["rbac"] - assert rbac_call["tenant_id"] == "tenant-1" - assert rbac_call["account_id"] == "account-1" - assert rbac_call["path_args"] == {"app_id": "backing-app-1"} - (rbac_check,) = rbac_call["checks"] - assert rbac_check.scene is audio_module.RBACPermission.APP_TEST_AND_RUN - assert isinstance(rbac_check.locator, audio_module.PlainApp) - assert calls["draft"] == { - "tenant_id": "tenant-1", - "agent_id": str(agent_id), - "account_id": "account-1", - "draft_type": AgentConfigDraftType.DEBUG_BUILD, - "session": session, - } - assert calls["asr"] == { - "app_model": app_model, - "agent_soul": agent_soul, - "file": calls["asr"]["file"], - "session": session, - "end_user": None, - } + def load_agent_soul_for_debug(**kwargs): + captured.update(kwargs) + return AgentSoulConfig.model_validate({"app_features": {"speech_to_text": {"enabled": True}}}) + monkeypatch.setattr(AgentComposerService, "load_agent_soul_for_debug", load_agent_soul_for_debug) + monkeypatch.setattr(AudioService, "transcript_agent_asr", lambda **_kwargs: {"text": "ok"}) -def test_agent_console_audio_api_defaults_to_normal_draft( - app: Flask, monkeypatch: pytest.MonkeyPatch, unbound_session: Session -) -> None: - agent_id = UUID("019ef3d2-b24c-7803-b428-18b5ee8fb853") - captured: dict[str, object] = {} - monkeypatch.setattr( - audio_module, - "resolve_agent_runtime_app_model", - lambda **_kwargs: _app(app_id="backing-app-1"), - ) - - def load_agent_soul_for_debug(**kwargs): - captured.update(kwargs) - return AgentSoulConfig.model_validate({"app_features": {"speech_to_text": {"enabled": True}}}) - - monkeypatch.setattr(AgentComposerService, "load_agent_soul_for_debug", load_agent_soul_for_debug) - monkeypatch.setattr(AudioService, "transcript_agent_asr", lambda **_kwargs: {"text": "ok"}) - - api = AgentChatMessageAudioApi() - handler = unwrap(api.post) - with app.test_request_context( - f"/console/api/agent/{agent_id}/audio-to-text", - method="POST", - data={"file": _file_data()}, - ): - response = handler( - api, - session=unbound_session, - current_tenant_id="tenant-1", - current_user=_account(), - agent_id=agent_id, - ) - - assert response == {"text": "ok"} - assert captured["draft_type"] == AgentConfigDraftType.DRAFT - - -def test_agent_console_audio_api_checks_rbac_with_backing_app_id( - app: Flask, monkeypatch: pytest.MonkeyPatch, unbound_session: Session -) -> None: - agent_id = UUID("019ef3d2-b24c-7803-b428-18b5ee8fb853") - app_model = _app(app_id="backing-app-1") - soul_loaded = False - - monkeypatch.setattr(audio_module, "resolve_agent_runtime_app_model", lambda **_kwargs: app_model) - - def deny_access(**kwargs): - assert kwargs["path_args"] == {"app_id": "backing-app-1"} - raise Forbidden() - - def load_agent_soul_for_debug(**_kwargs): - nonlocal soul_loaded - soul_loaded = True - return AgentSoulConfig() - - monkeypatch.setattr(audio_module, "enforce_rbac_checks", deny_access) - monkeypatch.setattr(AgentComposerService, "load_agent_soul_for_debug", load_agent_soul_for_debug) - - api = AgentChatMessageAudioApi() - handler = unwrap(api.post) - with app.test_request_context( - f"/console/api/agent/{agent_id}/audio-to-text", - method="POST", - data={"file": _file_data()}, - ): - with pytest.raises(Forbidden): - handler( + api = AgentChatMessageAudioApi() + handler = unwrap(api.post) + with app.test_request_context( + f"/console/api/agent/{agent_id}/audio-to-text", + method="POST", + data={"file": _file_data()}, + ): + response = handler( api, session=unbound_session, current_tenant_id="tenant-1", @@ -226,467 +393,520 @@ def test_agent_console_audio_api_checks_rbac_with_backing_app_id( agent_id=agent_id, ) - assert soul_loaded is False + assert response == {"text": "ok"} + assert captured["draft_type"] == AgentConfigDraftType.DRAFT + @staticmethod + def test_agent_console_audio_api_checks_rbac_with_backing_app_id( + app: Flask, monkeypatch: pytest.MonkeyPatch, unbound_session: Session + ) -> None: + agent_id = UUID("019ef3d2-b24c-7803-b428-18b5ee8fb853") + app_model = _app(app_id="backing-app-1") + soul_loaded = False -def test_agent_console_audio_api_preserves_missing_build_draft_404( - app: Flask, monkeypatch: pytest.MonkeyPatch, unbound_session: Session -) -> None: - agent_id = UUID("019ef3d2-b24c-7803-b428-18b5ee8fb853") - monkeypatch.setattr( - audio_module, - "resolve_agent_runtime_app_model", - lambda **_kwargs: _app(app_id="backing-app-1"), - ) - monkeypatch.setattr( - AgentComposerService, - "load_agent_soul_for_debug", - lambda **_kwargs: (_ for _ in ()).throw(AgentVersionNotFoundError()), + monkeypatch.setattr(audio_module, "resolve_agent_runtime_app_model", lambda **_kwargs: app_model) + + def deny_access(**kwargs): + assert kwargs["path_args"] == {"app_id": "backing-app-1"} + raise Forbidden() + + def load_agent_soul_for_debug(**_kwargs): + nonlocal soul_loaded + soul_loaded = True + return AgentSoulConfig() + + monkeypatch.setattr(audio_module, "enforce_rbac_checks", deny_access) + monkeypatch.setattr(AgentComposerService, "load_agent_soul_for_debug", load_agent_soul_for_debug) + + api = AgentChatMessageAudioApi() + handler = unwrap(api.post) + with app.test_request_context( + f"/console/api/agent/{agent_id}/audio-to-text", + method="POST", + data={"file": _file_data()}, + ): + with pytest.raises(Forbidden): + handler( + api, + session=unbound_session, + current_tenant_id="tenant-1", + current_user=_account(), + agent_id=agent_id, + ) + + assert soul_loaded is False + + @staticmethod + def test_agent_console_audio_api_preserves_missing_build_draft_404( + app: Flask, monkeypatch: pytest.MonkeyPatch, unbound_session: Session + ) -> None: + agent_id = UUID("019ef3d2-b24c-7803-b428-18b5ee8fb853") + monkeypatch.setattr( + audio_module, + "resolve_agent_runtime_app_model", + lambda **_kwargs: _app(app_id="backing-app-1"), + ) + monkeypatch.setattr( + AgentComposerService, + "load_agent_soul_for_debug", + lambda **_kwargs: (_ for _ in ()).throw(AgentVersionNotFoundError()), + ) + + api = AgentChatMessageAudioApi() + handler = unwrap(api.post) + with app.test_request_context( + f"/console/api/agent/{agent_id}/audio-to-text", + method="POST", + data={"file": _file_data(), "draft_type": "debug_build"}, + ): + with pytest.raises(AgentVersionNotFoundError): + handler( + api, + session=unbound_session, + current_tenant_id="tenant-1", + current_user=_account(), + agent_id=agent_id, + ) + + @staticmethod + @pytest.mark.parametrize( + ("exc", "expected"), + [ + (AppModelConfigBrokenError(), AppUnavailableError), + (NoAudioUploadedServiceError(), NoAudioUploadedError), + (AudioTooLargeServiceError("too big"), AudioTooLargeError), + (UnsupportedAudioTypeServiceError(), UnsupportedAudioTypeError), + (ProviderNotSupportSpeechToTextServiceError(), ProviderNotSupportSpeechToTextError), + (SpeechToTextDisabledServiceError(), SpeechToTextDisabledError), + (ProviderTokenNotInitError("token"), ProviderNotInitializeError), + (QuotaExceededError(), ProviderQuotaExceededError), + (ModelCurrentlyNotSupportError(), ProviderModelCurrentlyNotSupportError), + (InvokeError("invoke"), CompletionRequestError), + ], ) + def test_console_audio_api_error_mapping(app: Flask, monkeypatch: pytest.MonkeyPatch, exc, expected) -> None: + monkeypatch.setattr(AudioService, "transcript_asr", lambda **_kwargs: (_ for _ in ()).throw(exc)) + api = ChatMessageAudioApi() + handler = unwrap(api.post) + app_model = _app() - api = AgentChatMessageAudioApi() - handler = unwrap(api.post) - with app.test_request_context( - f"/console/api/agent/{agent_id}/audio-to-text", - method="POST", - data={"file": _file_data(), "draft_type": "debug_build"}, - ): - with pytest.raises(AgentVersionNotFoundError): - handler( - api, - session=unbound_session, - current_tenant_id="tenant-1", - current_user=_account(), - agent_id=agent_id, - ) + with app.test_request_context( + "/console/api/apps/app/audio-to-text", method="POST", data={"file": _file_data()} + ): + with pytest.raises(expected): + handler(api, app_model=app_model) + @staticmethod + def test_console_audio_api_unhandled_error(app: Flask, monkeypatch: pytest.MonkeyPatch) -> None: + monkeypatch.setattr( + AudioService, "transcript_asr", lambda **_kwargs: (_ for _ in ()).throw(RuntimeError("boom")) + ) + api = ChatMessageAudioApi() + handler = unwrap(api.post) + app_model = _app() -@pytest.mark.parametrize( - ("exc", "expected"), - [ - (AppModelConfigBrokenError(), AppUnavailableError), - (NoAudioUploadedServiceError(), NoAudioUploadedError), - (AudioTooLargeServiceError("too big"), AudioTooLargeError), - (UnsupportedAudioTypeServiceError(), UnsupportedAudioTypeError), - (ProviderNotSupportSpeechToTextServiceError(), ProviderNotSupportSpeechToTextError), - (SpeechToTextDisabledServiceError(), SpeechToTextDisabledError), - (ProviderTokenNotInitError("token"), ProviderNotInitializeError), - (QuotaExceededError(), ProviderQuotaExceededError), - (ModelCurrentlyNotSupportError(), ProviderModelCurrentlyNotSupportError), - (InvokeError("invoke"), CompletionRequestError), - ], -) -def test_console_audio_api_error_mapping(app: Flask, monkeypatch: pytest.MonkeyPatch, exc, expected) -> None: - monkeypatch.setattr(AudioService, "transcript_asr", lambda **_kwargs: (_ for _ in ()).throw(exc)) - api = ChatMessageAudioApi() - handler = unwrap(api.post) - app_model = _app() + with app.test_request_context( + "/console/api/apps/app/audio-to-text", method="POST", data={"file": _file_data()} + ): + with pytest.raises(InternalServerError): + handler(api, app_model=app_model) - with app.test_request_context("/console/api/apps/app/audio-to-text", method="POST", data={"file": _file_data()}): - with pytest.raises(expected): - handler(api, app_model=app_model) + @staticmethod + def test_console_text_api_success(app: Flask, monkeypatch: pytest.MonkeyPatch) -> None: + monkeypatch.setattr(AudioService, "transcript_tts", lambda **_kwargs: {"audio": "ok"}) + api = ChatMessageTextApi() + handler = unwrap(api.post) + app_model = _app() -def test_console_audio_api_unhandled_error(app: Flask, monkeypatch: pytest.MonkeyPatch) -> None: - monkeypatch.setattr(AudioService, "transcript_asr", lambda **_kwargs: (_ for _ in ()).throw(RuntimeError("boom"))) - api = ChatMessageAudioApi() - handler = unwrap(api.post) - app_model = _app() + with app.test_request_context( + "/console/api/apps/app/text-to-audio", + method="POST", + json={"text": "hello", "voice": "v"}, + ): + response = handler(api, TextToSpeechPayload(text="hello"), app_model=app_model) - with app.test_request_context("/console/api/apps/app/audio-to-text", method="POST", data={"file": _file_data()}): - with pytest.raises(InternalServerError): - handler(api, app_model=app_model) + assert response == {"audio": "ok"} + @staticmethod + def test_console_text_api_builds_message_ref(app: Flask, monkeypatch: pytest.MonkeyPatch) -> None: + api = ChatMessageTextApi() + handler = unwrap(api.post) + app_model = _app(app_id="app-1") + calls = {} -def test_console_text_api_success(app: Flask, monkeypatch: pytest.MonkeyPatch) -> None: - monkeypatch.setattr(AudioService, "transcript_tts", lambda **_kwargs: {"audio": "ok"}) + def fake_transcript_tts(**kwargs): + calls.update(kwargs) + return {"audio": "ok"} - api = ChatMessageTextApi() - handler = unwrap(api.post) - app_model = _app() + monkeypatch.setattr(AudioService, "transcript_tts", fake_transcript_tts) - with app.test_request_context( - "/console/api/apps/app/text-to-audio", - method="POST", - json={"text": "hello", "voice": "v"}, - ): - response = handler(api, TextToSpeechPayload(text="hello"), app_model=app_model) + with ( + app.test_request_context( + "/console/api/apps/app-1/text-to-audio", + method="POST", + json={"text": "hello", "message_id": "message-1"}, + ), + patch("controllers.console.app.audio.current_user", _account()), + ): + response = handler(api, TextToSpeechPayload(text="hello", message_id="message-1"), app_model=app_model) - assert response == {"audio": "ok"} + assert response == {"audio": "ok"} + assert calls["message_ref"] == MessageRef(AppRef("tenant-1", "app-1"), "message-1", account_id="account-1") + @staticmethod + def test_console_text_api_error_mapping(app: Flask, monkeypatch: pytest.MonkeyPatch) -> None: + monkeypatch.setattr( + AudioService, "transcript_tts", lambda **_kwargs: (_ for _ in ()).throw(QuotaExceededError()) + ) -def test_console_text_api_builds_message_ref(app: Flask, monkeypatch: pytest.MonkeyPatch) -> None: - api = ChatMessageTextApi() - handler = unwrap(api.post) - app_model = _app(app_id="app-1") - calls = {} + api = ChatMessageTextApi() + handler = unwrap(api.post) + app_model = _app() - def fake_transcript_tts(**kwargs): - calls.update(kwargs) - return {"audio": "ok"} + with app.test_request_context( + "/console/api/apps/app/text-to-audio", + method="POST", + json={"text": "hello"}, + ): + with pytest.raises(ProviderQuotaExceededError): + handler(api, TextToSpeechPayload(text="hello"), app_model=app_model) - monkeypatch.setattr(AudioService, "transcript_tts", fake_transcript_tts) + @staticmethod + def test_console_text_modes_success(app: Flask, monkeypatch: pytest.MonkeyPatch) -> None: + expected_voices = [{"name": "Voice 1", "value": "voice-1"}] + monkeypatch.setattr(AudioService, "transcript_tts_voices", lambda **_kwargs: expected_voices) - with ( - app.test_request_context( + api = TextModesApi() + handler = unwrap(api.get) + app_model = _app(tenant_id="t1") + + with app.test_request_context("/console/api/apps/app/text-to-audio/voices?language=en", method="GET"): + response = handler(api, TextToSpeechVoiceQuery(language="en-US"), app_model=app_model) + + assert response == expected_voices + + @staticmethod + def test_audio_to_text_success(app: Flask, monkeypatch: pytest.MonkeyPatch) -> None: + api = ChatMessageAudioApi() + method = unwrap(api.post) + + response_payload = {"text": "hello"} + monkeypatch.setattr(AudioService, "transcript_asr", lambda **_kwargs: response_payload) + + app_model = _app(app_id="app-1") + + data = {"file": (io.BytesIO(b"x"), "sample.wav")} + with app.test_request_context( + "/console/api/apps/app-1/audio-to-text", + method="POST", + data=data, + content_type="multipart/form-data", + ): + response = method(api, app_model=app_model) + + assert response == response_payload + + @staticmethod + def test_audio_to_text_maps_audio_too_large(app: Flask, monkeypatch: pytest.MonkeyPatch) -> None: + api = ChatMessageAudioApi() + method = unwrap(api.post) + + monkeypatch.setattr( + AudioService, + "transcript_asr", + lambda **_kwargs: (_ for _ in ()).throw(AudioTooLargeServiceError("too large")), + ) + + app_model = _app(app_id="app-1") + + data = {"file": (io.BytesIO(b"x"), "sample.wav")} + with app.test_request_context( + "/console/api/apps/app-1/audio-to-text", + method="POST", + data=data, + content_type="multipart/form-data", + ): + with pytest.raises(AudioTooLargeError): + method(api, app_model=app_model) + + @staticmethod + def test_text_to_audio_success(app: Flask, monkeypatch: pytest.MonkeyPatch) -> None: + api = ChatMessageTextApi() + method = unwrap(api.post) + + monkeypatch.setattr(AudioService, "transcript_tts", lambda **_kwargs: {"audio": "ok"}) + + app_model = _app(app_id="app-1") + + with app.test_request_context( "/console/api/apps/app-1/text-to-audio", method="POST", - json={"text": "hello", "message_id": "message-1"}, - ), - patch("controllers.console.app.audio.current_user", _account()), - ): - response = handler(api, TextToSpeechPayload(text="hello", message_id="message-1"), app_model=app_model) + json={"text": "hello"}, + ): + response = method(api, TextToSpeechPayload(text="hello"), app_model=app_model) - assert response == {"audio": "ok"} - assert calls["message_ref"] == MessageRef(AppRef("tenant-1", "app-1"), "message-1", account_id="account-1") + assert response == {"audio": "ok"} + @staticmethod + def test_text_to_audio_voices_success(app: Flask, monkeypatch: pytest.MonkeyPatch) -> None: + api = TextModesApi() + method = unwrap(api.get) -def test_console_text_api_error_mapping(app: Flask, monkeypatch: pytest.MonkeyPatch) -> None: - monkeypatch.setattr(AudioService, "transcript_tts", lambda **_kwargs: (_ for _ in ()).throw(QuotaExceededError())) + expected_voices = [{"name": "Voice 1", "value": "voice-1"}] + monkeypatch.setattr(AudioService, "transcript_tts_voices", lambda **_kwargs: expected_voices) - api = ChatMessageTextApi() - handler = unwrap(api.post) - app_model = _app() + app_model = _app() - with app.test_request_context( - "/console/api/apps/app/text-to-audio", - method="POST", - json={"text": "hello"}, - ): - with pytest.raises(ProviderQuotaExceededError): - handler(api, TextToSpeechPayload(text="hello"), app_model=app_model) + with app.test_request_context( + "/console/api/apps/app-1/text-to-audio/voices", + method="GET", + query_string={"language": "en-US"}, + ): + response = method(api, TextToSpeechVoiceQuery(language="en-US"), app_model=app_model) + assert response == expected_voices -def test_console_text_modes_success(app: Flask, monkeypatch: pytest.MonkeyPatch) -> None: - expected_voices = [{"name": "Voice 1", "value": "voice-1"}] - monkeypatch.setattr(AudioService, "transcript_tts_voices", lambda **_kwargs: expected_voices) + @staticmethod + def test_audio_to_text_with_invalid_file(app: Flask, monkeypatch: pytest.MonkeyPatch) -> None: + api = ChatMessageAudioApi() + method = unwrap(api.post) - api = TextModesApi() - handler = unwrap(api.get) - app_model = _app(tenant_id="t1") + monkeypatch.setattr(AudioService, "transcript_asr", lambda **_kwargs: {"text": "test"}) - with app.test_request_context("/console/api/apps/app/text-to-audio/voices?language=en", method="GET"): - response = handler(api, TextToSpeechVoiceQuery(language="en-US"), app_model=app_model) + app_model = _app(app_id="app-1") - assert response == expected_voices + data = {"file": (io.BytesIO(b"invalid"), "sample.xyz")} + with app.test_request_context( + "/console/api/apps/app-1/audio-to-text", + method="POST", + data=data, + content_type="multipart/form-data", + ): + # Should not raise, AudioService is mocked + response = method(api, app_model=app_model) + assert response == {"text": "test"} + @staticmethod + def test_text_to_audio_with_language_param(app: Flask, monkeypatch: pytest.MonkeyPatch) -> None: + api = ChatMessageTextApi() + method = unwrap(api.post) -def test_audio_to_text_success(app: Flask, monkeypatch: pytest.MonkeyPatch) -> None: - api = ChatMessageAudioApi() - method = unwrap(api.post) + monkeypatch.setattr(AudioService, "transcript_tts", lambda **_kwargs: {"audio": "test"}) - response_payload = {"text": "hello"} - monkeypatch.setattr(AudioService, "transcript_asr", lambda **_kwargs: response_payload) + app_model = _app(app_id="app-1") - app_model = _app(app_id="app-1") + with app.test_request_context( + "/console/api/apps/app-1/text-to-audio", + method="POST", + json={"text": "hello", "language": "en-US"}, + ): + response = method(api, TextToSpeechPayload(text="hello"), app_model=app_model) + assert response == {"audio": "test"} - data = {"file": (io.BytesIO(b"x"), "sample.wav")} - with app.test_request_context( - "/console/api/apps/app-1/audio-to-text", - method="POST", - data=data, - content_type="multipart/form-data", - ): - response = method(api, app_model=app_model) + @staticmethod + def test_text_to_audio_voices_with_language_filter(app: Flask, monkeypatch: pytest.MonkeyPatch) -> None: + api = TextModesApi() + method = unwrap(api.get) - assert response == response_payload - - -def test_audio_to_text_maps_audio_too_large(app: Flask, monkeypatch: pytest.MonkeyPatch) -> None: - api = ChatMessageAudioApi() - method = unwrap(api.post) - - monkeypatch.setattr( - AudioService, - "transcript_asr", - lambda **_kwargs: (_ for _ in ()).throw(AudioTooLargeServiceError("too large")), - ) - - app_model = _app(app_id="app-1") - - data = {"file": (io.BytesIO(b"x"), "sample.wav")} - with app.test_request_context( - "/console/api/apps/app-1/audio-to-text", - method="POST", - data=data, - content_type="multipart/form-data", - ): - with pytest.raises(AudioTooLargeError): - method(api, app_model=app_model) - - -def test_text_to_audio_success(app: Flask, monkeypatch: pytest.MonkeyPatch) -> None: - api = ChatMessageTextApi() - method = unwrap(api.post) - - monkeypatch.setattr(AudioService, "transcript_tts", lambda **_kwargs: {"audio": "ok"}) - - app_model = _app(app_id="app-1") - - with app.test_request_context( - "/console/api/apps/app-1/text-to-audio", - method="POST", - json={"text": "hello"}, - ): - response = method(api, TextToSpeechPayload(text="hello"), app_model=app_model) - - assert response == {"audio": "ok"} - - -def test_text_to_audio_voices_success(app: Flask, monkeypatch: pytest.MonkeyPatch) -> None: - api = TextModesApi() - method = unwrap(api.get) - - expected_voices = [{"name": "Voice 1", "value": "voice-1"}] - monkeypatch.setattr(AudioService, "transcript_tts_voices", lambda **_kwargs: expected_voices) - - app_model = _app() - - with app.test_request_context( - "/console/api/apps/app-1/text-to-audio/voices", - method="GET", - query_string={"language": "en-US"}, - ): - response = method(api, TextToSpeechVoiceQuery(language="en-US"), app_model=app_model) - - assert response == expected_voices - - -def test_audio_to_text_with_invalid_file(app: Flask, monkeypatch: pytest.MonkeyPatch) -> None: - api = ChatMessageAudioApi() - method = unwrap(api.post) - - monkeypatch.setattr(AudioService, "transcript_asr", lambda **_kwargs: {"text": "test"}) - - app_model = _app(app_id="app-1") - - data = {"file": (io.BytesIO(b"invalid"), "sample.xyz")} - with app.test_request_context( - "/console/api/apps/app-1/audio-to-text", - method="POST", - data=data, - content_type="multipart/form-data", - ): - # Should not raise, AudioService is mocked - response = method(api, app_model=app_model) - assert response == {"text": "test"} - - -def test_text_to_audio_with_language_param(app: Flask, monkeypatch: pytest.MonkeyPatch) -> None: - api = ChatMessageTextApi() - method = unwrap(api.post) - - monkeypatch.setattr(AudioService, "transcript_tts", lambda **_kwargs: {"audio": "test"}) - - app_model = _app(app_id="app-1") - - with app.test_request_context( - "/console/api/apps/app-1/text-to-audio", - method="POST", - json={"text": "hello", "language": "en-US"}, - ): - response = method(api, TextToSpeechPayload(text="hello"), app_model=app_model) - assert response == {"audio": "test"} - - -def test_text_to_audio_voices_with_language_filter(app: Flask, monkeypatch: pytest.MonkeyPatch) -> None: - api = TextModesApi() - method = unwrap(api.get) - - monkeypatch.setattr( - AudioService, - "transcript_tts_voices", - lambda **_kwargs: [{"name": "Voice 1", "value": "voice-1"}], - ) - - app_model = _app() - - with app.test_request_context( - "/console/api/apps/app-1/text-to-audio/voices?language=en-US", - method="GET", - ): - response = method(api, TextToSpeechVoiceQuery(language="en-US"), app_model=app_model) - assert isinstance(response, list) - - -def test_agent_text_to_speech_voices_uses_backing_app_and_language(app: Flask, unbound_session: Session) -> None: - agent_id = UUID("019ef3d2-b24c-7803-b428-18b5ee8fb853") - app_model = _app(app_id="backing-app-1") - current_user = _account() - voices = [{"name": "Voice 1", "value": "voice-1"}] - api = AgentTextToSpeechVoicesApi() - with ( - patch.object(audio_module, "resolve_existing_agent_runtime_app_model", return_value=app_model) as resolve_app, - patch.object(audio_module, "enforce_rbac_checks") as check_access, - patch.object(AudioService, "transcript_tts_voices", return_value=voices) as get_voices, - app.test_request_context(f"/console/api/agent/{agent_id}/text-to-audio/voices?language=en-US"), - ): - response = unwrap(api.get)( - api, - TextToSpeechVoiceQuery(language="en-US"), - session=unbound_session, - current_tenant_id="tenant-1", - current_user=current_user, - agent_id=agent_id, + monkeypatch.setattr( + AudioService, + "transcript_tts_voices", + lambda **_kwargs: [{"name": "Voice 1", "value": "voice-1"}], ) - assert response == voices - resolve_app.assert_called_once_with(session=unbound_session, tenant_id="tenant-1", agent_id=agent_id) - check_access.assert_called_once() - rbac_call = check_access.call_args.kwargs - assert rbac_call["tenant_id"] == "tenant-1" - assert rbac_call["account_id"] == current_user.id - assert rbac_call["path_args"] == {"agent_id": str(agent_id)} - (rbac_check,) = rbac_call["checks"] - assert rbac_check.scene is audio_module.RBACPermission.AGENT_PREVIEW - assert isinstance(rbac_check.locator, audio_module.AgentId) - get_voices.assert_called_once_with(tenant_id="tenant-1", language="en-US") + app_model = _app() + with app.test_request_context( + "/console/api/apps/app-1/text-to-audio/voices?language=en-US", + method="GET", + ): + response = method(api, TextToSpeechVoiceQuery(language="en-US"), app_model=app_model) + assert isinstance(response, list) -def test_agent_text_to_speech_voices_does_not_query_provider_without_permission( - app: Flask, unbound_session: Session -) -> None: - agent_id = UUID("019ef3d2-b24c-7803-b428-18b5ee8fb853") - api = AgentTextToSpeechVoicesApi() - with ( - patch.object(audio_module, "resolve_existing_agent_runtime_app_model") as resolve_app, - patch.object(audio_module, "enforce_rbac_checks", side_effect=Forbidden()), - patch.object(AudioService, "transcript_tts_voices") as get_voices, - app.test_request_context(f"/console/api/agent/{agent_id}/text-to-audio/voices?language=en-US"), - pytest.raises(Forbidden), - ): - unwrap(api.get)( - api, - TextToSpeechVoiceQuery(language="en-US"), + @staticmethod + def test_agent_text_to_speech_voices_uses_backing_app_and_language(app: Flask, unbound_session: Session) -> None: + agent_id = UUID("019ef3d2-b24c-7803-b428-18b5ee8fb853") + app_model = _app(app_id="backing-app-1") + current_user = _account() + voices = [{"name": "Voice 1", "value": "voice-1"}] + api = AgentTextToSpeechVoicesApi() + with ( + patch.object( + audio_module, "resolve_existing_agent_runtime_app_model", return_value=app_model + ) as resolve_app, + patch.object(audio_module, "enforce_rbac_checks") as check_access, + patch.object(AudioService, "transcript_tts_voices", return_value=voices) as get_voices, + app.test_request_context(f"/console/api/agent/{agent_id}/text-to-audio/voices?language=en-US"), + ): + response = unwrap(api.get)( + api, + TextToSpeechVoiceQuery(language="en-US"), + session=unbound_session, + current_tenant_id="tenant-1", + current_user=current_user, + agent_id=agent_id, + ) + + assert response == voices + resolve_app.assert_called_once_with(session=unbound_session, tenant_id="tenant-1", agent_id=agent_id) + check_access.assert_called_once() + rbac_call = check_access.call_args.kwargs + assert rbac_call["tenant_id"] == "tenant-1" + assert rbac_call["account_id"] == current_user.id + assert rbac_call["path_args"] == {"agent_id": str(agent_id)} + (rbac_check,) = rbac_call["checks"] + assert rbac_check.scene is audio_module.RBACPermission.AGENT_PREVIEW + assert isinstance(rbac_check.locator, audio_module.AgentId) + get_voices.assert_called_once_with(tenant_id="tenant-1", language="en-US") + + @staticmethod + def test_agent_text_to_speech_voices_does_not_query_provider_without_permission( + app: Flask, unbound_session: Session + ) -> None: + agent_id = UUID("019ef3d2-b24c-7803-b428-18b5ee8fb853") + api = AgentTextToSpeechVoicesApi() + with ( + patch.object(audio_module, "resolve_existing_agent_runtime_app_model") as resolve_app, + patch.object(audio_module, "enforce_rbac_checks", side_effect=Forbidden()), + patch.object(AudioService, "transcript_tts_voices") as get_voices, + app.test_request_context(f"/console/api/agent/{agent_id}/text-to-audio/voices?language=en-US"), + pytest.raises(Forbidden), + ): + unwrap(api.get)( + api, + TextToSpeechVoiceQuery(language="en-US"), + session=unbound_session, + current_tenant_id="tenant-1", + current_user=_account(), + agent_id=agent_id, + ) + + get_voices.assert_not_called() + resolve_app.assert_not_called() + + @staticmethod + def test_agent_text_to_speech_voices_rejects_missing_language(app: Flask, unbound_session: Session) -> None: + agent_id = UUID("019ef3d2-b24c-7803-b428-18b5ee8fb853") + api = AgentTextToSpeechVoicesApi() + view = unwrap(api.get) + # Stop at the innermost decorator so request validation runs without the auth and session layers. + handler = unwrap(api.get, stop=lambda f: f.__dict__.get("__wrapped__") is view) + with ( + patch.object(audio_module, "resolve_existing_agent_runtime_app_model") as resolve_app, + patch.object(audio_module, "enforce_rbac_checks") as check_access, + app.test_request_context(f"/console/api/agent/{agent_id}/text-to-audio/voices", method="GET"), + pytest.raises(UnprocessableEntity), + ): + handler( + api, + session=unbound_session, + current_tenant_id="tenant-1", + current_user=_account(), + agent_id=agent_id, + ) + + check_access.assert_not_called() + resolve_app.assert_not_called() + + @staticmethod + @pytest.mark.parametrize("message_id", [None, "019ef3d2-b24c-7803-b428-18b5ee8fb854"]) + def test_agent_text_to_speech_returns_audio_with_scoped_message( + app: Flask, unbound_session: Session, message_id: str | None + ) -> None: + agent_id = UUID("019ef3d2-b24c-7803-b428-18b5ee8fb853") + app_model = _app(app_id="backing-app-1") + current_user = _account() + payload = TextToSpeechPayload(text="Preview this voice", voice="voice-1", message_id=message_id) + audio = Response(b"RIFF\x00\x00\x00\x00WAVE", content_type="audio/wav") + api = AgentChatMessageTextApi() + with ( + patch.object( + audio_module, "resolve_existing_agent_runtime_app_model", return_value=app_model + ) as resolve_app, + patch.object(audio_module, "enforce_rbac_checks") as check_access, + patch.object(AudioService, "transcript_tts", return_value=audio) as synthesize, + app.test_request_context(f"/console/api/agent/{agent_id}/text-to-audio", method="POST"), + ): + response = unwrap(api.post)( + api, + req_data=payload, + session=unbound_session, + current_tenant_id="tenant-1", + current_user=current_user, + agent_id=agent_id, + ) + + assert response.get_data() == audio.get_data() + assert response.content_type == "audio/wav" + resolve_app.assert_called_once_with(session=unbound_session, tenant_id="tenant-1", agent_id=agent_id) + check_access.assert_called_once() + rbac_call = check_access.call_args.kwargs + assert rbac_call["tenant_id"] == "tenant-1" + assert rbac_call["account_id"] == current_user.id + assert rbac_call["path_args"] == {"agent_id": str(agent_id)} + (rbac_check,) = rbac_call["checks"] + assert rbac_check.scene is audio_module.RBACPermission.AGENT_TEST_AND_RUN + assert isinstance(rbac_check.locator, audio_module.AgentId) + synthesize.assert_called_once_with( + app_model=app_model, session=unbound_session, - current_tenant_id="tenant-1", - current_user=_account(), - agent_id=agent_id, + text=payload.text, + voice=payload.voice, + message_ref=MessageRef(AppRef("tenant-1", "backing-app-1"), message_id, account_id=current_user.id) + if message_id + else None, + is_draft=True, ) - get_voices.assert_not_called() - resolve_app.assert_not_called() + @staticmethod + def test_agent_text_to_speech_denies_access_before_resolving_app(app: Flask, unbound_session: Session) -> None: + agent_id = UUID("019ef3d2-b24c-7803-b428-18b5ee8fb853") + api = AgentChatMessageTextApi() + with ( + patch.object(audio_module, "resolve_existing_agent_runtime_app_model") as resolve_app, + patch.object(audio_module, "enforce_rbac_checks", side_effect=Forbidden()), + patch.object(AudioService, "transcript_tts") as synthesize, + app.test_request_context(f"/console/api/agent/{agent_id}/text-to-audio", method="POST"), + pytest.raises(Forbidden), + ): + unwrap(api.post)( + api, + req_data=TextToSpeechPayload(text="Preview"), + session=unbound_session, + current_tenant_id="tenant-1", + current_user=_account(), + agent_id=agent_id, + ) + resolve_app.assert_not_called() + synthesize.assert_not_called() -def test_agent_text_to_speech_voices_rejects_missing_language(app: Flask, unbound_session: Session) -> None: - agent_id = UUID("019ef3d2-b24c-7803-b428-18b5ee8fb853") - api = AgentTextToSpeechVoicesApi() - view = unwrap(api.get) - # Stop at the innermost decorator so request validation runs without the auth and session layers. - handler = unwrap(api.get, stop=lambda f: f.__dict__.get("__wrapped__") is view) - with ( - patch.object(audio_module, "resolve_existing_agent_runtime_app_model") as resolve_app, - patch.object(audio_module, "enforce_rbac_checks") as check_access, - app.test_request_context(f"/console/api/agent/{agent_id}/text-to-audio/voices", method="GET"), - pytest.raises(UnprocessableEntity), - ): - handler( - api, - session=unbound_session, - current_tenant_id="tenant-1", - current_user=_account(), - agent_id=agent_id, - ) + @staticmethod + @pytest.mark.parametrize("voices", [False, True]) + def test_agent_tts_preserves_missing_agent_error(app: Flask, unbound_session: Session, voices: bool) -> None: + agent_id = UUID("019ef3d2-b24c-7803-b428-18b5ee8fb853") + api = AgentTextToSpeechVoicesApi() if voices else AgentChatMessageTextApi() + handler = unwrap(api.get if isinstance(api, AgentTextToSpeechVoicesApi) else api.post) + args: dict[str, object] = { + "session": unbound_session, + "current_tenant_id": "tenant-1", + "current_user": _account(), + "agent_id": agent_id, + } + if voices: + args["req_data"] = TextToSpeechVoiceQuery(language="en-US") + else: + args["req_data"] = TextToSpeechPayload(text="Preview") + with ( + patch.object(audio_module, "resolve_existing_agent_runtime_app_model", side_effect=AgentNotFoundError()), + patch.object(audio_module, "enforce_rbac_checks"), + patch.object(AudioService, "transcript_tts") as synthesize, + patch.object(AudioService, "transcript_tts_voices") as get_voices, + app.test_request_context(f"/console/api/agent/{agent_id}/text-to-audio?language=en-US"), + pytest.raises(AgentNotFoundError), + ): + handler(api, **args) - check_access.assert_not_called() - resolve_app.assert_not_called() - - -@pytest.mark.parametrize("message_id", [None, "019ef3d2-b24c-7803-b428-18b5ee8fb854"]) -def test_agent_text_to_speech_returns_audio_with_scoped_message( - app: Flask, unbound_session: Session, message_id: str | None -) -> None: - agent_id = UUID("019ef3d2-b24c-7803-b428-18b5ee8fb853") - app_model = _app(app_id="backing-app-1") - current_user = _account() - payload = TextToSpeechPayload(text="Preview this voice", voice="voice-1", message_id=message_id) - audio = Response(b"RIFF\x00\x00\x00\x00WAVE", content_type="audio/wav") - api = AgentChatMessageTextApi() - with ( - patch.object(audio_module, "resolve_existing_agent_runtime_app_model", return_value=app_model) as resolve_app, - patch.object(audio_module, "enforce_rbac_checks") as check_access, - patch.object(AudioService, "transcript_tts", return_value=audio) as synthesize, - app.test_request_context(f"/console/api/agent/{agent_id}/text-to-audio", method="POST"), - ): - response = unwrap(api.post)( - api, - req_data=payload, - session=unbound_session, - current_tenant_id="tenant-1", - current_user=current_user, - agent_id=agent_id, - ) - - assert response.get_data() == audio.get_data() - assert response.content_type == "audio/wav" - resolve_app.assert_called_once_with(session=unbound_session, tenant_id="tenant-1", agent_id=agent_id) - check_access.assert_called_once() - rbac_call = check_access.call_args.kwargs - assert rbac_call["tenant_id"] == "tenant-1" - assert rbac_call["account_id"] == current_user.id - assert rbac_call["path_args"] == {"agent_id": str(agent_id)} - (rbac_check,) = rbac_call["checks"] - assert rbac_check.scene is audio_module.RBACPermission.AGENT_TEST_AND_RUN - assert isinstance(rbac_check.locator, audio_module.AgentId) - synthesize.assert_called_once_with( - app_model=app_model, - session=unbound_session, - text=payload.text, - voice=payload.voice, - message_ref=MessageRef(AppRef("tenant-1", "backing-app-1"), message_id, account_id=current_user.id) - if message_id - else None, - is_draft=True, - ) - - -def test_agent_text_to_speech_denies_access_before_resolving_app(app: Flask, unbound_session: Session) -> None: - agent_id = UUID("019ef3d2-b24c-7803-b428-18b5ee8fb853") - api = AgentChatMessageTextApi() - with ( - patch.object(audio_module, "resolve_existing_agent_runtime_app_model") as resolve_app, - patch.object(audio_module, "enforce_rbac_checks", side_effect=Forbidden()), - patch.object(AudioService, "transcript_tts") as synthesize, - app.test_request_context(f"/console/api/agent/{agent_id}/text-to-audio", method="POST"), - pytest.raises(Forbidden), - ): - unwrap(api.post)( - api, - req_data=TextToSpeechPayload(text="Preview"), - session=unbound_session, - current_tenant_id="tenant-1", - current_user=_account(), - agent_id=agent_id, - ) - - resolve_app.assert_not_called() - synthesize.assert_not_called() - - -@pytest.mark.parametrize("voices", [False, True]) -def test_agent_tts_preserves_missing_agent_error(app: Flask, unbound_session: Session, voices: bool) -> None: - agent_id = UUID("019ef3d2-b24c-7803-b428-18b5ee8fb853") - api = AgentTextToSpeechVoicesApi() if voices else AgentChatMessageTextApi() - handler = unwrap(api.get if isinstance(api, AgentTextToSpeechVoicesApi) else api.post) - args: dict[str, object] = { - "session": unbound_session, - "current_tenant_id": "tenant-1", - "current_user": _account(), - "agent_id": agent_id, - } - if voices: - args["req_data"] = TextToSpeechVoiceQuery(language="en-US") - else: - args["req_data"] = TextToSpeechPayload(text="Preview") - with ( - patch.object(audio_module, "resolve_existing_agent_runtime_app_model", side_effect=AgentNotFoundError()), - patch.object(audio_module, "enforce_rbac_checks"), - patch.object(AudioService, "transcript_tts") as synthesize, - patch.object(AudioService, "transcript_tts_voices") as get_voices, - app.test_request_context(f"/console/api/agent/{agent_id}/text-to-audio?language=en-US"), - pytest.raises(AgentNotFoundError), - ): - handler(api, **args) - - synthesize.assert_not_called() - get_voices.assert_not_called() + synthesize.assert_not_called() + get_voices.assert_not_called() diff --git a/api/tests/unit_tests/controllers/console/app/test_console_app_api.py b/api/tests/unit_tests/controllers/console/app/test_console_app_api.py deleted file mode 100644 index 8dd0edd30b4..00000000000 --- a/api/tests/unit_tests/controllers/console/app/test_console_app_api.py +++ /dev/null @@ -1,353 +0,0 @@ -"""Transport contracts for the Console app application boundary.""" - -import inspect -from collections.abc import Callable -from dataclasses import dataclass, field, replace -from datetime import datetime -from io import BytesIO -from typing import cast -from uuid import UUID - -import pytest -from flask import Flask, Response -from flask_restx import Resource -from werkzeug.exceptions import BadRequest, Forbidden - -from controllers.console import flask_admission -from controllers.console.app import app as controller -from controllers.console.app.error import AppNotFoundError -from libs.login import AccountWithTenant -from machinery.context import RequestContext -from models.account import Account, TenantAccountRole -from models.enums import CustomizeTokenStrategy -from services.agent.roster_package_entities import RosterAgentPackageExport -from services.app.console_service import ( - AppExportPaidPlanRequiredError, - ConsoleAppNotFoundError, - CreatorsPlatformDisabledError, - InvalidAppExportError, -) -from services.app_tracing_config_service import AppTraceSettings -from services.entities.app_entities import ( - AppExportOptions, - AppListParams, - AppPage, - AppRecord, - CopyAppParams, - CreateAppParams, - StarredAppListParams, - UpdateAppParams, -) -from services.entities.dsl_entities import Import, ImportStatus -from tests.unit_tests.config_override import apply_config_overrides - -APP_ID = UUID("11111111-1111-1111-1111-111111111111") -CONTEXT = RequestContext("request", "trace", "actor", "workspace") -RECORD = AppRecord( - str(APP_ID), "Example", "chat", description="", created_at=datetime(2024, 1, 1), updated_at=datetime(2024, 1, 1) -) -SITE = { - "code": "site", - "title": "Example", - "icon_type": None, - "icon": None, - "icon_background": None, - "description": None, - "default_language": "en-US", - "chat_color_theme": None, - "chat_color_theme_inverted": False, - "customize_domain": None, - "copyright": None, - "privacy_policy": None, - "input_placeholder": None, - "custom_disclaimer": "", - "customize_token_strategy": CustomizeTokenStrategy.NOT_ALLOW, - "prompt_public": False, - "show_workflow_steps": False, - "use_icon_as_answer_icon": False, - "created_by": None, - "created_at": datetime(2024, 1, 1), - "updated_by": None, - "updated_at": datetime(2024, 1, 1), -} - - -@dataclass -class AppBoundary: - calls: list[tuple[object, ...]] = field(default_factory=list) - status: ImportStatus = ImportStatus.COMPLETED - exported: str | RosterAgentPackageExport = "dsl" - record: AppRecord = RECORD - - def list_apps(self, context: RequestContext, params: AppListParams | StarredAppListParams) -> AppPage: - self.calls.append(("list", context, params)) - return AppPage(params.page, params.limit, 1, False, [RECORD]) - - def create(self, context: RequestContext, params: CreateAppParams) -> AppRecord: - self.calls.append(("create", context, params)) - return RECORD - - def get(self, context: RequestContext, app_id: str) -> AppRecord: - self.calls.append(("get", context, app_id)) - return self.record - - def update(self, context: RequestContext, app_id: str, params: UpdateAppParams) -> AppRecord: - self.calls.append(("update", context, app_id, params)) - return RECORD - - def update_icon(self, context: RequestContext, app_id: str, **kwargs: object) -> AppRecord: - self.calls.append(("icon", context, app_id, kwargs)) - return RECORD - - def copy(self, context: RequestContext, app_id: str, params: CopyAppParams) -> tuple[Import, AppRecord | None]: - self.calls.append(("copy", context, app_id, params)) - return Import(id="import", status=self.status), RECORD - - def export(self, context: RequestContext, app_id: str, options: AppExportOptions) -> str | RosterAgentPackageExport: - self.calls.append(("export", context, app_id, options)) - return self.exported - - def publish(self, context: RequestContext, app_id: str) -> str: - self.calls.append(("publish", context, app_id)) - return "https://creators.example.com" - - def delete(self, context: RequestContext, app_id: str) -> None: - self.calls.append(("delete", context, app_id)) - - def set_trace(self, context: RequestContext, app_id: str, settings: AppTraceSettings) -> None: - self.calls.append(("trace", context, app_id, settings)) - - -@dataclass -class AppServices: - console: AppBoundary - - -@dataclass -class Services: - apps: AppServices - - -@pytest.fixture -def boundary(monkeypatch: pytest.MonkeyPatch) -> AppBoundary: - boundary = AppBoundary() - monkeypatch.setattr(controller, "application_services", lambda: Services(AppServices(boundary))) - return boundary - - -def invoke(resource: type[Resource], method: Callable[..., object], **kwargs: object) -> object: - return inspect.unwrap(method)(resource(), CONTEXT, **kwargs) - - -def test_list_parses_repeated_filters_and_serializes_materialized_page(app: Flask, boundary: AppBoundary) -> None: - tag = str(APP_ID) - with app.test_request_context(f"/?page=2&limit=5&tag_ids={tag}&is_created_by_me=true"): - data, status = cast(tuple[dict[str, object], int], invoke(controller.AppListApi, controller.AppListApi.get)) - assert status == 200 - assert data["page"] == 2 - assert data["limit"] == 5 - assert boundary.calls[0][1] is CONTEXT - params = boundary.calls[0][2] - assert isinstance(params, AppListParams) - assert params.tag_ids == [tag] - assert params.is_created_by_me is True - - -@pytest.mark.parametrize("configured", [True, False]) -def test_detail_resolves_urls_at_the_http_boundary( - app: Flask, boundary: AppBoundary, config_overrides: Callable[..., None], configured: bool -) -> None: - config_overrides( - SERVICE_API_URL="https://api.example.com" if configured else "", - APP_WEB_URL="https://apps.example.com" if configured else "", - ) - boundary.record = replace(RECORD, site=SITE) - with app.test_request_context("/", base_url="https://console.example.com/"): - response = cast(dict[str, object], invoke(controller.AppApi, controller.AppApi.get, app_id=APP_ID)) - assert response["api_base_url"] == ( - "https://api.example.com/v1" if configured else "https://console.example.com/v1" - ) - site = response["site"] - assert isinstance(site, dict) - assert site["app_base_url"] == ("https://apps.example.com" if configured else "https://console.example.com") - - -@pytest.mark.parametrize("icon_type", [None, "image"]) -def test_update_preserves_icon_type_omission_and_forwards_explicit_identity( - app: Flask, boundary: AppBoundary, icon_type: str | None -) -> None: - payload = {"name": "Updated", "icon": "image-id"} - if icon_type is not None: - payload["icon_type"] = icon_type - with app.test_request_context("/", method="POST", json=payload): - data = cast(dict[str, object], invoke(controller.AppApi, controller.AppApi.put, app_id=APP_ID)) - assert data["id"] == str(APP_ID) - assert boundary.calls[0][1:3] == (CONTEXT, str(APP_ID)) - params = boundary.calls[0][3] - assert isinstance(params, UpdateAppParams) - assert params.icon_type == icon_type - - -def test_icon_body_and_trace_body_are_parsed(app: Flask, boundary: AppBoundary) -> None: - with app.test_request_context("/", method="POST", json={"icon_type": "image", "icon": "image-id"}): - invoke(controller.AppIconApi, controller.AppIconApi.post, app_id=APP_ID) - assert boundary.calls[-1][-1] == {"icon_type": "image", "icon": "image-id", "icon_background": ""} - with app.test_request_context("/", method="POST", json={"enabled": True, "tracing_provider": "langfuse"}): - assert invoke(controller.AppTraceApi, controller.AppTraceApi.post, app_id=APP_ID) == {"result": "success"} - assert boundary.calls[-1] == ("trace", CONTEXT, str(APP_ID), AppTraceSettings(True, "langfuse")) - - -@pytest.mark.parametrize( - ("status", "expected"), - [ - (ImportStatus.FAILED, 400), - (ImportStatus.PENDING, 202), - (ImportStatus.COMPLETED, 201), - (ImportStatus.COMPLETED_WITH_WARNINGS, 201), - ], -) -def test_copy_preserves_import_status_responses( - app: Flask, boundary: AppBoundary, status: ImportStatus, expected: int -) -> None: - boundary.status = status - with app.test_request_context("/", method="POST", json={"name": "Copy"}): - data, code = cast( - tuple[dict[str, object], int], invoke(controller.AppCopyApi, controller.AppCopyApi.post, app_id=APP_ID) - ) - assert code == expected - assert data["id"] == ("import" if expected != 201 else str(APP_ID)) - - -def test_export_serializes_yaml_and_closes_archive_on_response_close(app: Flask, boundary: AppBoundary) -> None: - with app.test_request_context("/?format=yaml&include_secret=true"): - assert invoke(controller.AppExportApi, controller.AppExportApi.get, app_id=APP_ID) == {"data": "dsl"} - options = boundary.calls[-1][-1] - assert isinstance(options, AppExportOptions) - assert options.include_secret is True - archive = BytesIO(b"zip") - boundary.exported = RosterAgentPackageExport(archive=archive, filename="example.ifpkg", size=3) - with app.test_request_context("/"): - response = invoke(controller.AppExportApi, controller.AppExportApi.get, app_id=APP_ID) - assert isinstance(response, Response) - assert response.mimetype == "application/zip" - assert "example.ifpkg" in response.headers["Content-Disposition"] - response.close() - assert archive.closed - - -def test_export_closes_archive_when_send_file_fails( - app: Flask, boundary: AppBoundary, monkeypatch: pytest.MonkeyPatch -) -> None: - archive = BytesIO(b"zip") - boundary.exported = RosterAgentPackageExport(archive=archive, filename="example.ifpkg", size=3) - - def fail(*_args: object, **_kwargs: object) -> None: - raise OSError("delivery failed") - - monkeypatch.setattr(controller, "send_file", fail) - with app.test_request_context("/"), pytest.raises(OSError, match="delivery failed"): - invoke(controller.AppExportApi, controller.AppExportApi.get, app_id=APP_ID) - assert archive.closed - - -def test_publish_delegates_and_delete_has_no_body(app: Flask, boundary: AppBoundary) -> None: - with app.test_request_context("/"): - assert invoke( - controller.AppPublishToCreatorsPlatformApi, controller.AppPublishToCreatorsPlatformApi.post, app_id=APP_ID - ) == { - "redirect_url": "https://creators.example.com", - } - assert invoke(controller.AppApi, controller.AppApi.delete, app_id=APP_ID) == ("", 204) - assert boundary.calls == [("publish", CONTEXT, str(APP_ID)), ("delete", CONTEXT, str(APP_ID))] - - -@pytest.mark.parametrize( - ("error", "expected"), - [ - (ConsoleAppNotFoundError(), AppNotFoundError), - (InvalidAppExportError("invalid version"), BadRequest), - (AppExportPaidPlanRequiredError("paid plan required"), Forbidden), - ], -) -def test_application_errors_are_translated_at_transport_boundary( - monkeypatch: pytest.MonkeyPatch, error: Exception, expected: type[Exception] -) -> None: - def fail(*_args: object, **_kwargs: object) -> None: - raise error - - monkeypatch.setattr(Resource, "dispatch_request", fail) - with pytest.raises(expected): - controller.AppResource().dispatch_request() - - -def test_disabled_creators_preserves_error_body(monkeypatch: pytest.MonkeyPatch) -> None: - def fail(*_args: object, **_kwargs: object) -> None: - raise CreatorsPlatformDisabledError("Creators Platform features are not enabled") - - monkeypatch.setattr(Resource, "dispatch_request", fail) - assert controller.AppResource().dispatch_request() == ({"error": "Creators Platform features are not enabled"}, 403) - - -@pytest.mark.parametrize( - "method", - [ - controller.AppApi.put, - controller.AppApi.delete, - controller.AppCopyApi.post, - controller.AppExportApi.get, - controller.AppNameApi.post, - controller.AppIconApi.post, - controller.AppSiteStatus.post, - controller.AppApiStatus.post, - controller.AppTraceApi.post, - ], -) -def test_mutations_reject_read_only_role_at_admission( - app: Flask, monkeypatch: pytest.MonkeyPatch, method: Callable[..., object] -) -> None: - account = Account(name="Reader", email="reader@example.com") - account.role = TenantAccountRole.NORMAL - apply_config_overrides(monkeypatch, RBAC_ENABLED=False) - monkeypatch.setattr( - flask_admission, - "current_account_with_tenant", - lambda: AccountWithTenant(account=account, tenant_id="workspace"), - ) - injector = inspect.unwrap(method, stop=lambda f: "allowed_roles" in inspect.getclosurevars(f).nonlocals) - with app.test_request_context(), pytest.raises(Forbidden): - injector(None, app_id=APP_ID) - - -def test_create_quota_is_part_of_admission_and_precedes_application_call( - app: Flask, boundary: AppBoundary, monkeypatch: pytest.MonkeyPatch -) -> None: - from services.entities.feature_entities import FeatureModel, LimitationModel - - account = Account(name="Editor", email="editor@example.com") - account.role = TenantAccountRole.EDITOR - account.id = CONTEXT.account_id - identity = AccountWithTenant(account=account, tenant_id=CONTEXT.active_workspace_id) - apply_config_overrides(monkeypatch, DEPLOYMENT_EDITION="CLOUD", RBAC_ENABLED=False) - monkeypatch.setattr(flask_admission, "current_account_with_tenant", lambda: identity) - monkeypatch.setattr("controllers.console.wraps.current_account_with_tenant", lambda: identity) - monkeypatch.setattr(flask_admission, "enforce_rbac_checks", lambda **_kwargs: None) - monkeypatch.setattr( - "controllers.console.wraps.FeatureService.get_features", - lambda *_args, **_kwargs: FeatureModel(apps=LimitationModel(size=1, limit=1)), - ) - admission = inspect.unwrap( - controller.AppListApi.post, stop=lambda f: inspect.getclosurevars(f).nonlocals.get("resource") == "apps" - ) - with app.test_request_context(method="POST", json={"name": "Example", "mode": "chat"}), pytest.raises(Forbidden): - admission(controller.AppListApi()) - assert boundary.calls == [] - - -def test_request_validation_retains_422_and_redacts_submitted_values(app: Flask, boundary: AppBoundary) -> None: - from werkzeug.exceptions import UnprocessableEntity - - with app.test_request_context(method="PUT", json={"name": "", "icon": "private-input"}): - with pytest.raises(UnprocessableEntity) as error: - invoke(controller.AppApi, controller.AppApi.put, app_id=APP_ID) - assert "private-input" not in (error.value.description or "") - assert boundary.calls == [] diff --git a/api/tests/unit_tests/controllers/console/app/test_conversation_api.py b/api/tests/unit_tests/controllers/console/app/test_conversation.py similarity index 99% rename from api/tests/unit_tests/controllers/console/app/test_conversation_api.py rename to api/tests/unit_tests/controllers/console/app/test_conversation.py index 26b0286e030..eaf406da334 100644 --- a/api/tests/unit_tests/controllers/console/app/test_conversation_api.py +++ b/api/tests/unit_tests/controllers/console/app/test_conversation.py @@ -26,7 +26,7 @@ from models.enums import ( from models.model import AppMode, Conversation from models.workflow import WorkflowRun, WorkflowType from services.errors.conversation import ConversationNotExistsError -from tests.unit_tests.model_factories import make_app, make_conversation +from tests.helpers.unit.model_factories import make_app, make_conversation def _make_account() -> Account: diff --git a/api/tests/unit_tests/controllers/console/app/test_conversation_variables_api.py b/api/tests/unit_tests/controllers/console/app/test_conversation_variables.py similarity index 99% rename from api/tests/unit_tests/controllers/console/app/test_conversation_variables_api.py rename to api/tests/unit_tests/controllers/console/app/test_conversation_variables.py index 5fcca2eae67..eb21ffba737 100644 --- a/api/tests/unit_tests/controllers/console/app/test_conversation_variables_api.py +++ b/api/tests/unit_tests/controllers/console/app/test_conversation_variables.py @@ -26,7 +26,7 @@ from models import AppMode, ConversationVariable from models.account import Account, AccountStatus, TenantAccountRole from models.agent import Agent, AgentScope, AgentSource, AgentStatus from services.enterprise.rbac_service import RBACService -from tests.unit_tests.model_factories import make_account, make_app, make_tenant +from tests.helpers.unit.model_factories import make_account, make_app, make_tenant APP_ID = "11111111-1111-1111-1111-111111111111" OTHER_APP_ID = "22222222-2222-2222-2222-222222222222" diff --git a/api/tests/unit_tests/controllers/console/app/test_create_app_payload.py b/api/tests/unit_tests/controllers/console/app/test_create_app_payload.py deleted file mode 100644 index dfb2a4c1398..00000000000 --- a/api/tests/unit_tests/controllers/console/app/test_create_app_payload.py +++ /dev/null @@ -1,24 +0,0 @@ -"""Regression tests for CreateAppPayload mode validation.""" - -import pytest -from pydantic import ValidationError - -from controllers.console.app.app import CreateAppPayload - - -class TestCreateAppPayloadMode: - @pytest.mark.parametrize( - "mode", - ["chat", "agent-chat", "advanced-chat", "workflow", "completion"], - ) - def test_accepts_supported_modes(self, mode: str): - payload = CreateAppPayload.model_validate({"name": "X", "mode": mode}) - assert payload.mode == mode - - def test_rejects_agent_mode(self): - with pytest.raises(ValidationError): - CreateAppPayload.model_validate({"name": "X", "mode": "agent"}) - - def test_rejects_unknown_mode(self): - with pytest.raises(ValidationError): - CreateAppPayload.model_validate({"name": "X", "mode": "not-a-mode"}) diff --git a/api/tests/unit_tests/controllers/console/app/test_generator_api.py b/api/tests/unit_tests/controllers/console/app/test_generator.py similarity index 77% rename from api/tests/unit_tests/controllers/console/app/test_generator_api.py rename to api/tests/unit_tests/controllers/console/app/test_generator.py index 338cc8ef6a0..2456439778d 100644 --- a/api/tests/unit_tests/controllers/console/app/test_generator_api.py +++ b/api/tests/unit_tests/controllers/console/app/test_generator.py @@ -1,7 +1,7 @@ from __future__ import annotations import json -from inspect import unwrap +from inspect import unwrap as _generator_api_unwrap from unittest.mock import MagicMock import pytest @@ -15,13 +15,15 @@ from controllers.console.app.generator import ( InstructionTemplatePayload, RuleCodeGeneratePayload, RuleGeneratePayload, + RuleStructuredOutputPayload, WorkflowGeneratePayload, WorkflowInstructionSuggestionsPayload, ) -from core.errors.error import ProviderTokenNotInitError +from core.errors.error import ModelCurrentlyNotSupportError, ProviderTokenNotInitError, QuotaExceededError +from graphon.model_runtime.errors.invoke import InvokeError from models.model import App, AppMode from models.workflow import Workflow -from tests.unit_tests.model_factories import make_workflow +from tests.helpers.unit.model_factories import make_workflow def _persist_app(session: Session, *, tenant_id: str = "t1") -> App: @@ -40,7 +42,7 @@ def _persist_app(session: Session, *, tenant_id: str = "t1") -> App: return app_model -def _model_config_payload(): +def _generator_api_model_config_payload(): return {"provider": "openai", "name": "gpt-4o", "mode": "chat", "completion_params": {}} @@ -65,7 +67,7 @@ def _install_workflow_service(monkeypatch: pytest.MonkeyPatch, workflow): def test_rule_generate_success(app: Flask, monkeypatch: pytest.MonkeyPatch) -> None: api = generator_module.RuleGenerateApi() - method = unwrap(api.post) + method = _generator_api_unwrap(api.post) rule_config = { "prompt": "You are a helpful assistant.", @@ -78,7 +80,7 @@ def test_rule_generate_success(app: Flask, monkeypatch: pytest.MonkeyPatch) -> N with app.test_request_context( "/console/api/rule-generate", method="POST", - json={"instruction": "do it", "model_config": _model_config_payload()}, + json={"instruction": "do it", "model_config": _generator_api_model_config_payload()}, ): response = method(api, RuleGeneratePayload.model_validate(request.get_json()), "t1") @@ -87,7 +89,7 @@ def test_rule_generate_success(app: Flask, monkeypatch: pytest.MonkeyPatch) -> N def test_rule_code_generate_maps_token_error(app: Flask, monkeypatch: pytest.MonkeyPatch) -> None: api = generator_module.RuleCodeGenerateApi() - method = unwrap(api.post) + method = _generator_api_unwrap(api.post) def _raise(*_args, **_kwargs): raise ProviderTokenNotInitError("missing token") @@ -97,7 +99,7 @@ def test_rule_code_generate_maps_token_error(app: Flask, monkeypatch: pytest.Mon with app.test_request_context( "/console/api/rule-code-generate", method="POST", - json={"instruction": "do it", "model_config": _model_config_payload()}, + json={"instruction": "do it", "model_config": _generator_api_model_config_payload()}, ): with pytest.raises(ProviderNotInitializeError): method(api, RuleCodeGeneratePayload.model_validate(request.get_json()), "t1") @@ -106,7 +108,7 @@ def test_rule_code_generate_maps_token_error(app: Flask, monkeypatch: pytest.Mon @pytest.mark.parametrize("sqlite_session", [(App,)], indirect=True) def test_instruction_generate_app_not_found(app: Flask, sqlite_session: Session) -> None: api = generator_module.InstructionGenerateApi() - method = unwrap(api.post) + method = _generator_api_unwrap(api.post) _persist_app(sqlite_session, tenant_id="other-tenant") with app.test_request_context( @@ -116,7 +118,7 @@ def test_instruction_generate_app_not_found(app: Flask, sqlite_session: Session) "flow_id": "app-1", "node_id": "node-1", "instruction": "do", - "model_config": _model_config_payload(), + "model_config": _generator_api_model_config_payload(), }, ): response, status = method( @@ -133,7 +135,7 @@ def test_instruction_generate_workflow_not_found( app: Flask, monkeypatch: pytest.MonkeyPatch, sqlite_session: Session ) -> None: api = generator_module.InstructionGenerateApi() - method = unwrap(api.post) + method = _generator_api_unwrap(api.post) _persist_app(sqlite_session) _install_workflow_service(monkeypatch, workflow=None) @@ -145,7 +147,7 @@ def test_instruction_generate_workflow_not_found( "flow_id": "app-1", "node_id": "node-1", "instruction": "do", - "model_config": _model_config_payload(), + "model_config": _generator_api_model_config_payload(), }, ): response, status = method( @@ -161,7 +163,7 @@ def test_instruction_generate_node_missing( app: Flask, monkeypatch: pytest.MonkeyPatch, sqlite_session: Session ) -> None: api = generator_module.InstructionGenerateApi() - method = unwrap(api.post) + method = _generator_api_unwrap(api.post) _persist_app(sqlite_session) @@ -175,7 +177,7 @@ def test_instruction_generate_node_missing( "flow_id": "app-1", "node_id": "node-1", "instruction": "do", - "model_config": _model_config_payload(), + "model_config": _generator_api_model_config_payload(), }, ): response, status = method( @@ -189,7 +191,7 @@ def test_instruction_generate_node_missing( @pytest.mark.parametrize("sqlite_session", [(App,)], indirect=True) def test_instruction_generate_code_node(app: Flask, monkeypatch: pytest.MonkeyPatch, sqlite_session: Session) -> None: api = generator_module.InstructionGenerateApi() - method = unwrap(api.post) + method = _generator_api_unwrap(api.post) app_model = _persist_app(sqlite_session) @@ -210,7 +212,7 @@ def test_instruction_generate_code_node(app: Flask, monkeypatch: pytest.MonkeyPa "flow_id": "app-1", "node_id": "node-1", "instruction": "do", - "model_config": _model_config_payload(), + "model_config": _generator_api_model_config_payload(), }, ): response = method(api, InstructionGeneratePayload.model_validate(request.get_json()), sqlite_session, "t1") @@ -225,7 +227,7 @@ def test_instruction_generate_legacy_modify( app: Flask, monkeypatch: pytest.MonkeyPatch, sqlite_session: Session ) -> None: api = generator_module.InstructionGenerateApi() - method = unwrap(api.post) + method = _generator_api_unwrap(api.post) captured: dict = {} def _capture(**kwargs): @@ -242,7 +244,7 @@ def test_instruction_generate_legacy_modify( "node_id": "", "current": "old", "instruction": "do", - "model_config": _model_config_payload(), + "model_config": _generator_api_model_config_payload(), }, ): response = method(api, InstructionGeneratePayload.model_validate(request.get_json()), sqlite_session, "t1") @@ -254,7 +256,7 @@ def test_instruction_generate_legacy_modify( @pytest.mark.parametrize("sqlite_session", [(App,)], indirect=True) def test_instruction_generate_incompatible_params(app: Flask, sqlite_session: Session) -> None: api = generator_module.InstructionGenerateApi() - method = unwrap(api.post) + method = _generator_api_unwrap(api.post) with app.test_request_context( "/console/api/instruction-generate", method="POST", @@ -263,7 +265,7 @@ def test_instruction_generate_incompatible_params(app: Flask, sqlite_session: Se "node_id": "", "current": "", "instruction": "do", - "model_config": _model_config_payload(), + "model_config": _generator_api_model_config_payload(), }, ): response, status = method( @@ -276,7 +278,7 @@ def test_instruction_generate_incompatible_params(app: Flask, sqlite_session: Se def test_instruction_template_prompt(app: Flask) -> None: api = generator_module.InstructionGenerationTemplateApi() - method = unwrap(api.post) + method = _generator_api_unwrap(api.post) with app.test_request_context( "/console/api/instruction-generate/template", @@ -290,7 +292,7 @@ def test_instruction_template_prompt(app: Flask) -> None: def test_instruction_template_invalid_type(app: Flask) -> None: api = generator_module.InstructionGenerationTemplateApi() - method = unwrap(api.post) + method = _generator_api_unwrap(api.post) with app.test_request_context( "/console/api/instruction-generate/template", @@ -309,7 +311,7 @@ def _workflow_generate_payload() -> dict: "mode": "workflow", "instruction": "Summarize a URL", "ideal_output": "A 3-sentence summary.", - "model_config": _model_config_payload(), + "model_config": _generator_api_model_config_payload(), } @@ -342,7 +344,7 @@ def _stub_workflow_service(monkeypatch: pytest.MonkeyPatch, returns=None, raises def test_workflow_generate_returns_service_result(app: Flask, monkeypatch: pytest.MonkeyPatch) -> None: api = generator_module.WorkflowGenerateApi() - method = unwrap(api.post) + method = _generator_api_unwrap(api.post) expected = { "graph": _workflow_graph(node_id="node-1"), @@ -381,7 +383,7 @@ def test_workflow_generate_maps_provider_token_error(app: Flask, monkeypatch: py """ProviderTokenNotInitError → ProviderNotInitializeError so the frontend can render the same "provider missing" UX as /rule-generate.""" api = generator_module.WorkflowGenerateApi() - method = unwrap(api.post) + method = _generator_api_unwrap(api.post) _stub_workflow_service(monkeypatch, raises=ProviderTokenNotInitError("missing token")) @@ -399,7 +401,7 @@ def test_workflow_generate_maps_quota_error(app: Flask, monkeypatch: pytest.Monk from core.errors.error import QuotaExceededError api = generator_module.WorkflowGenerateApi() - method = unwrap(api.post) + method = _generator_api_unwrap(api.post) _stub_workflow_service(monkeypatch, raises=QuotaExceededError()) @@ -417,7 +419,7 @@ def test_workflow_generate_maps_model_not_support_error(app: Flask, monkeypatch: from core.errors.error import ModelCurrentlyNotSupportError api = generator_module.WorkflowGenerateApi() - method = unwrap(api.post) + method = _generator_api_unwrap(api.post) _stub_workflow_service(monkeypatch, raises=ModelCurrentlyNotSupportError("not supported")) @@ -435,7 +437,7 @@ def test_workflow_generate_maps_invoke_error(app: Flask, monkeypatch: pytest.Mon from graphon.model_runtime.errors.invoke import InvokeError api = generator_module.WorkflowGenerateApi() - method = unwrap(api.post) + method = _generator_api_unwrap(api.post) _stub_workflow_service(monkeypatch, raises=InvokeError("LLM unreachable")) @@ -451,7 +453,7 @@ def test_workflow_generate_maps_invoke_error(app: Flask, monkeypatch: pytest.Mon def test_workflow_generate_accepts_advanced_chat_mode(app: Flask, monkeypatch: pytest.MonkeyPatch) -> None: """The payload Literal must accept advanced-chat as well as workflow.""" api = generator_module.WorkflowGenerateApi() - method = unwrap(api.post) + method = _generator_api_unwrap(api.post) captured: dict = {} @@ -482,7 +484,7 @@ def test_workflow_generate_accepts_advanced_chat_mode(app: Flask, monkeypatch: p def test_workflow_generate_forwards_current_graph_for_refine(app: Flask, monkeypatch: pytest.MonkeyPatch) -> None: """cmd+k `/refine`: the optional current_graph field reaches the service.""" api = generator_module.WorkflowGenerateApi() - method = unwrap(api.post) + method = _generator_api_unwrap(api.post) captured: dict = {} @@ -512,7 +514,7 @@ def test_workflow_generate_forwards_current_graph_for_refine(app: Flask, monkeyp def test_workflow_generate_current_graph_defaults_to_none(app: Flask, monkeypatch: pytest.MonkeyPatch) -> None: """Omitting current_graph (the `/create` path) forwards None to the service.""" api = generator_module.WorkflowGenerateApi() - method = unwrap(api.post) + method = _generator_api_unwrap(api.post) captured: dict = {} @@ -540,7 +542,7 @@ def test_workflow_generate_accepts_auto_mode(app: Flask, monkeypatch: pytest.Mon """Task 3: the payload Literal must accept ``auto``; the controller forwards it unchanged (the service resolves it) and returns the resolved ``mode``.""" api = generator_module.WorkflowGenerateApi() - method = unwrap(api.post) + method = _generator_api_unwrap(api.post) captured: dict = {} @@ -623,7 +625,7 @@ def test_generate_instruction_suggestions_no_default_model_returns_empty(monkeyp def test_workflow_instruction_suggestions_route_returns_list(app: Flask, monkeypatch: pytest.MonkeyPatch) -> None: """Task 1c (iii): the route wraps the generator output in {"suggestions": [...]}.""" api = generator_module.WorkflowInstructionSuggestionsApi() - method = unwrap(api.post) + method = _generator_api_unwrap(api.post) captured: dict = {} @@ -649,7 +651,7 @@ def test_workflow_instruction_suggestions_route_returns_list(app: Flask, monkeyp def test_workflow_instruction_suggestions_route_empty_is_valid_200(app: Flask, monkeypatch: pytest.MonkeyPatch) -> None: """Task 1c (iii): an empty list is a valid soft-fail response.""" api = generator_module.WorkflowInstructionSuggestionsApi() - method = unwrap(api.post) + method = _generator_api_unwrap(api.post) monkeypatch.setattr( generator_module.LLMGenerator, @@ -684,7 +686,7 @@ def _read_sse_frames(response) -> list[dict]: def test_workflow_generate_stream_emits_plan_then_result(app: Flask, monkeypatch: pytest.MonkeyPatch) -> None: """Task 2c: the stream endpoint writes one SSE frame per service event.""" api = generator_module.WorkflowGenerateStreamApi() - method = unwrap(api.post) + method = _generator_api_unwrap(api.post) def _stream(**_kwargs): yield ("plan", {"title": "Summarizer", "mode": "workflow", "nodes": []}) @@ -711,7 +713,7 @@ def test_workflow_generate_stream_provider_error_emits_result_event( ) -> None: """Task 2c: a provider-init error becomes a single MODEL_ERROR result frame, not a non-SSE error.""" api = generator_module.WorkflowGenerateStreamApi() - method = unwrap(api.post) + method = _generator_api_unwrap(api.post) def _stream(**_kwargs): raise ProviderTokenNotInitError("missing token") @@ -736,7 +738,7 @@ def test_workflow_generate_stream_provider_error_emits_result_event( def test_workflow_generate_stream_rejects_empty_instruction(app: Flask, monkeypatch: pytest.MonkeyPatch) -> None: """Task 2c: empty instructions get a normal 400 JSON BEFORE the stream opens.""" api = generator_module.WorkflowGenerateStreamApi() - method = unwrap(api.post) + method = _generator_api_unwrap(api.post) payload = _workflow_generate_payload() payload["instruction"] = " " @@ -745,3 +747,137 @@ def test_workflow_generate_stream_rejects_empty_instruction(app: Flask, monkeypa assert status == 400 assert response["errors"][0]["code"] == "EMPTY_INSTRUCTION" + + +def _generator_api_missing_unwrap(func): + """Unwrap a decorated function to test it directly.""" + while hasattr(func, "__wrapped__"): + func = func.__wrapped__ + return func + + +def _generator_api_missing_model_config_payload(): + return { + "provider": "test_provider", + "name": "test_model", + "mode": "chat", + "completion_params": {}, + } + + +def test_rule_generate_exceptions(app: Flask, monkeypatch: pytest.MonkeyPatch) -> None: + api = generator_module.RuleGenerateApi() + method = _generator_api_missing_unwrap(api.post) + + exceptions_to_test = [ + (ProviderTokenNotInitError("token error"), generator_module.ProviderNotInitializeError), + (QuotaExceededError("quota error"), generator_module.ProviderQuotaExceededError), + (ModelCurrentlyNotSupportError("model error"), generator_module.ProviderModelCurrentlyNotSupportError), + (InvokeError("invoke error"), generator_module.CompletionRequestError), + ] + + for err_to_raise, expected_exception in exceptions_to_test: + + def _raise(*_args, _err=err_to_raise, **_kwargs): + raise _err + + monkeypatch.setattr(generator_module.LLMGenerator, "generate_rule_config", _raise) + + with app.test_request_context( + "/console/api/rule-generate", + method="POST", + json={"instruction": "do it", "model_config": _generator_api_missing_model_config_payload()}, + ): + with pytest.raises(expected_exception): + method(api, RuleGeneratePayload.model_validate(request.get_json()), "t1") + + +def test_rule_code_generate_exceptions(app: Flask, monkeypatch: pytest.MonkeyPatch) -> None: + api = generator_module.RuleCodeGenerateApi() + method = _generator_api_missing_unwrap(api.post) + + exceptions_to_test = [ + (QuotaExceededError("quota error"), generator_module.ProviderQuotaExceededError), + (ModelCurrentlyNotSupportError("model error"), generator_module.ProviderModelCurrentlyNotSupportError), + (InvokeError("invoke error"), generator_module.CompletionRequestError), + ] + + for err_to_raise, expected_exception in exceptions_to_test: + + def _raise(*_args, _err=err_to_raise, **_kwargs): + raise _err + + monkeypatch.setattr(generator_module.LLMGenerator, "generate_code", _raise) + + with app.test_request_context( + "/console/api/rule-code-generate", + method="POST", + json={"instruction": "do it", "model_config": _generator_api_missing_model_config_payload()}, + ): + with pytest.raises(expected_exception): + method(api, RuleCodeGeneratePayload.model_validate(request.get_json()), "t1") + + +def test_structured_output_generate_exceptions(app: Flask, monkeypatch: pytest.MonkeyPatch) -> None: + api = generator_module.RuleStructuredOutputGenerateApi() + method = _generator_api_missing_unwrap(api.post) + + exceptions_to_test = [ + (ProviderTokenNotInitError("token error"), generator_module.ProviderNotInitializeError), + (QuotaExceededError("quota error"), generator_module.ProviderQuotaExceededError), + (ModelCurrentlyNotSupportError("model error"), generator_module.ProviderModelCurrentlyNotSupportError), + (InvokeError("invoke error"), generator_module.CompletionRequestError), + ] + + for err_to_raise, expected_exception in exceptions_to_test: + + def _raise(*_args, _err=err_to_raise, **_kwargs): + raise _err + + monkeypatch.setattr(generator_module.LLMGenerator, "generate_structured_output", _raise) + + with app.test_request_context( + "/console/api/structured-output-generate", + method="POST", + json={"instruction": "do it", "model_config": _generator_api_missing_model_config_payload()}, + ): + with pytest.raises(expected_exception): + method(api, RuleStructuredOutputPayload.model_validate(request.get_json()), "t1") + + +@pytest.mark.parametrize("sqlite_session", [()], indirect=True) +def test_instruction_generate_exceptions( + app: Flask, + monkeypatch: pytest.MonkeyPatch, + sqlite_session: Session, +) -> None: + api = generator_module.InstructionGenerateApi() + method = _generator_api_missing_unwrap(api.post) + + exceptions_to_test = [ + (ProviderTokenNotInitError("token error"), generator_module.ProviderNotInitializeError), + (QuotaExceededError("quota error"), generator_module.ProviderQuotaExceededError), + (ModelCurrentlyNotSupportError("model error"), generator_module.ProviderModelCurrentlyNotSupportError), + (InvokeError("invoke error"), generator_module.CompletionRequestError), + ] + + for err_to_raise, expected_exception in exceptions_to_test: + + def _raise(*_args, _err=err_to_raise, **_kwargs): + raise _err + + monkeypatch.setattr(generator_module.LLMGenerator, "instruction_modify_legacy", _raise) + + with app.test_request_context( + "/console/api/instruction-generate", + method="POST", + json={ + "flow_id": "app-1", + "node_id": "", + "current": "old", + "instruction": "do it", + "model_config": _generator_api_missing_model_config_payload(), + }, + ): + with pytest.raises(expected_exception): + method(api, InstructionGeneratePayload.model_validate(request.get_json()), sqlite_session, "t1") diff --git a/api/tests/unit_tests/controllers/console/app/test_generator_api_missing.py b/api/tests/unit_tests/controllers/console/app/test_generator_api_missing.py deleted file mode 100644 index 6ce20305bdc..00000000000 --- a/api/tests/unit_tests/controllers/console/app/test_generator_api_missing.py +++ /dev/null @@ -1,147 +0,0 @@ -import pytest -from flask import Flask, request -from sqlalchemy.orm import Session - -from controllers.console.app import generator as generator_module -from controllers.console.app.generator import ( - InstructionGeneratePayload, - RuleCodeGeneratePayload, - RuleGeneratePayload, - RuleStructuredOutputPayload, -) -from core.errors.error import ModelCurrentlyNotSupportError, ProviderTokenNotInitError, QuotaExceededError -from graphon.model_runtime.errors.invoke import InvokeError - - -def unwrap(func): - """Unwrap a decorated function to test it directly.""" - while hasattr(func, "__wrapped__"): - func = func.__wrapped__ - return func - - -def _model_config_payload(): - return { - "provider": "test_provider", - "name": "test_model", - "mode": "chat", - "completion_params": {}, - } - - -def test_rule_generate_exceptions(app: Flask, monkeypatch: pytest.MonkeyPatch) -> None: - api = generator_module.RuleGenerateApi() - method = unwrap(api.post) - - exceptions_to_test = [ - (ProviderTokenNotInitError("token error"), generator_module.ProviderNotInitializeError), - (QuotaExceededError("quota error"), generator_module.ProviderQuotaExceededError), - (ModelCurrentlyNotSupportError("model error"), generator_module.ProviderModelCurrentlyNotSupportError), - (InvokeError("invoke error"), generator_module.CompletionRequestError), - ] - - for err_to_raise, expected_exception in exceptions_to_test: - - def _raise(*_args, _err=err_to_raise, **_kwargs): - raise _err - - monkeypatch.setattr(generator_module.LLMGenerator, "generate_rule_config", _raise) - - with app.test_request_context( - "/console/api/rule-generate", - method="POST", - json={"instruction": "do it", "model_config": _model_config_payload()}, - ): - with pytest.raises(expected_exception): - method(api, RuleGeneratePayload.model_validate(request.get_json()), "t1") - - -def test_rule_code_generate_exceptions(app: Flask, monkeypatch: pytest.MonkeyPatch) -> None: - api = generator_module.RuleCodeGenerateApi() - method = unwrap(api.post) - - exceptions_to_test = [ - (QuotaExceededError("quota error"), generator_module.ProviderQuotaExceededError), - (ModelCurrentlyNotSupportError("model error"), generator_module.ProviderModelCurrentlyNotSupportError), - (InvokeError("invoke error"), generator_module.CompletionRequestError), - ] - - for err_to_raise, expected_exception in exceptions_to_test: - - def _raise(*_args, _err=err_to_raise, **_kwargs): - raise _err - - monkeypatch.setattr(generator_module.LLMGenerator, "generate_code", _raise) - - with app.test_request_context( - "/console/api/rule-code-generate", - method="POST", - json={"instruction": "do it", "model_config": _model_config_payload()}, - ): - with pytest.raises(expected_exception): - method(api, RuleCodeGeneratePayload.model_validate(request.get_json()), "t1") - - -def test_structured_output_generate_exceptions(app: Flask, monkeypatch: pytest.MonkeyPatch) -> None: - api = generator_module.RuleStructuredOutputGenerateApi() - method = unwrap(api.post) - - exceptions_to_test = [ - (ProviderTokenNotInitError("token error"), generator_module.ProviderNotInitializeError), - (QuotaExceededError("quota error"), generator_module.ProviderQuotaExceededError), - (ModelCurrentlyNotSupportError("model error"), generator_module.ProviderModelCurrentlyNotSupportError), - (InvokeError("invoke error"), generator_module.CompletionRequestError), - ] - - for err_to_raise, expected_exception in exceptions_to_test: - - def _raise(*_args, _err=err_to_raise, **_kwargs): - raise _err - - monkeypatch.setattr(generator_module.LLMGenerator, "generate_structured_output", _raise) - - with app.test_request_context( - "/console/api/structured-output-generate", - method="POST", - json={"instruction": "do it", "model_config": _model_config_payload()}, - ): - with pytest.raises(expected_exception): - method(api, RuleStructuredOutputPayload.model_validate(request.get_json()), "t1") - - -@pytest.mark.parametrize("sqlite_session", [()], indirect=True) -def test_instruction_generate_exceptions( - app: Flask, - monkeypatch: pytest.MonkeyPatch, - sqlite_session: Session, -) -> None: - api = generator_module.InstructionGenerateApi() - method = unwrap(api.post) - - exceptions_to_test = [ - (ProviderTokenNotInitError("token error"), generator_module.ProviderNotInitializeError), - (QuotaExceededError("quota error"), generator_module.ProviderQuotaExceededError), - (ModelCurrentlyNotSupportError("model error"), generator_module.ProviderModelCurrentlyNotSupportError), - (InvokeError("invoke error"), generator_module.CompletionRequestError), - ] - - for err_to_raise, expected_exception in exceptions_to_test: - - def _raise(*_args, _err=err_to_raise, **_kwargs): - raise _err - - monkeypatch.setattr(generator_module.LLMGenerator, "instruction_modify_legacy", _raise) - - with app.test_request_context( - "/console/api/instruction-generate", - method="POST", - json={ - "flow_id": "app-1", - "node_id": "", - "current": "old", - "instruction": "do it", - "model_config": _model_config_payload(), - }, - ): - with pytest.raises(expected_exception): - method(api, InstructionGeneratePayload.model_validate(request.get_json()), sqlite_session, "t1") diff --git a/api/tests/unit_tests/controllers/console/app/test_mcp_server.py b/api/tests/unit_tests/controllers/console/app/test_mcp_server.py new file mode 100644 index 00000000000..d4346e598e5 --- /dev/null +++ b/api/tests/unit_tests/controllers/console/app/test_mcp_server.py @@ -0,0 +1,374 @@ +import datetime as _mcp_server_response_datetime +from collections.abc import Callable, Sequence +from dataclasses import dataclass, field +from datetime import datetime as _mcp_server_admission_datetime +from inspect import unwrap +from types import SimpleNamespace +from uuid import UUID + +import pytest +from flask import Flask +from werkzeug.exceptions import Conflict, Forbidden, NotFound, UnprocessableEntity + +from controllers.common.rbac import PlainApp, RBACCheck +from controllers.console import console_ns, flask_admission +from controllers.console.app import mcp_server +from controllers.console.app import mcp_server as module +from controllers.console.app.error import AppNotFoundError +from controllers.console.app.mcp_server import ( + AppMCPServerController, + AppMCPServerRefreshController, + AppMCPServerResponse, +) +from controllers.console.wraps import RBACPermission +from libs.login import AccountWithTenant +from machinery.context import RequestContext +from models.account import TenantAccountRole +from services.app.mcp_server_service import ( + AppMCPServerAlreadyExistsError, + AppMCPServerAppNotFoundError, + AppMCPServerNotFoundError, + AppMCPServerRecord, + AppMCPServerStatus, +) +from tests.helpers.unit.config_override import apply_config_overrides +from tests.helpers.unit.model_factories import make_account + +"""App MCP Server routes retain their role and RBAC policy under unified admission.""" + + +_MCP_SERVER_ADMISSION_APP_ID = UUID("11111111-1111-1111-1111-111111111111") +_MCP_SERVER_ADMISSION_RECORD = AppMCPServerRecord( + id="server", + name="name", + server_code="code", + description="description", + status=AppMCPServerStatus.ACTIVE, + parameters="{}", + created_at=_mcp_server_admission_datetime(2024, 1, 1), + updated_at=_mcp_server_admission_datetime(2024, 1, 1), +) +MUTATIONS = [ + AppMCPServerController.post, + AppMCPServerController.put, + AppMCPServerRefreshController.post, +] + + +def admission(method: Callable[..., object]) -> Callable[..., object]: + return unwrap(method, stop=lambda view: "inject_request_context" in view.__code__.co_qualname) + + +class TestMcpServerAdmissionCases: + @staticmethod + @pytest.mark.parametrize("method", MUTATIONS) + @pytest.mark.parametrize("role", [TenantAccountRole.NORMAL, TenantAccountRole.DATASET_OPERATOR]) + def test_readonly_roles_cannot_mutate( + app: Flask, monkeypatch: pytest.MonkeyPatch, method: Callable[..., object], role: TenantAccountRole + ) -> None: + account = make_account(role=role) + apply_config_overrides(monkeypatch, RBAC_ENABLED=False) + monkeypatch.setattr( + flask_admission, "current_account_with_tenant", lambda: AccountWithTenant(account, "workspace") + ) + with app.test_request_context(), pytest.raises(Forbidden): + admission(method)(None, app_id=_MCP_SERVER_ADMISSION_APP_ID) + + @staticmethod + @pytest.mark.parametrize( + ("method", "permission"), + [ + (AppMCPServerController.get, RBACPermission.APP_VIEW_LAYOUT), + (AppMCPServerController.post, RBACPermission.APP_EDIT), + (AppMCPServerController.put, RBACPermission.APP_EDIT), + (AppMCPServerRefreshController.post, RBACPermission.APP_EDIT), + ], + ) + def test_rbac_checks_keep_app_scope( + app: Flask, monkeypatch: pytest.MonkeyPatch, method: Callable[..., object], permission: RBACPermission + ) -> None: + account = make_account(role=TenantAccountRole.NORMAL) + apply_config_overrides(monkeypatch, RBAC_ENABLED=True) + monkeypatch.setattr( + flask_admission, "current_account_with_tenant", lambda: AccountWithTenant(account, "workspace") + ) + called = [] + + def denied( + *, tenant_id: str, account_id: str, checks: Sequence[RBACCheck], path_args: dict[str, object] + ) -> None: + assert (tenant_id, account_id) == ("workspace", account.id) + assert path_args == {"app_id": _MCP_SERVER_ADMISSION_APP_ID} + assert len(checks) == 1 + assert checks[0].scene == permission + assert isinstance(checks[0].locator, PlainApp) + called.append(True) + raise Forbidden + + monkeypatch.setattr(flask_admission, "enforce_rbac_checks", denied) + with app.test_request_context(), pytest.raises(Forbidden): + admission(method)(None, app_id=_MCP_SERVER_ADMISSION_APP_ID) + assert called == [True] + + @staticmethod + @pytest.mark.parametrize("role", [TenantAccountRole.OWNER, TenantAccountRole.ADMIN, TenantAccountRole.EDITOR]) + def test_edit_roles_reach_service_with_request_context( + app: Flask, monkeypatch: pytest.MonkeyPatch, role: TenantAccountRole + ) -> None: + account = make_account(role=role) + apply_config_overrides(monkeypatch, RBAC_ENABLED=False) + monkeypatch.setattr( + flask_admission, "current_account_with_tenant", lambda: AccountWithTenant(account, "workspace") + ) + monkeypatch.setattr(flask_admission, "enforce_rbac_checks", lambda **_kwargs: None) + monkeypatch.setattr(flask_admission, "get_trace_id", lambda: None) + contexts: list[RequestContext] = [] + + def refresh(context: RequestContext, _app_id: str) -> AppMCPServerRecord: + contexts.append(context) + return _MCP_SERVER_ADMISSION_RECORD + + monkeypatch.setattr( + mcp_server, + "application_services", + lambda: SimpleNamespace(app_mcp_servers=SimpleNamespace(refresh=refresh)), + ) + with app.test_request_context(method="POST", headers={"X-Trace-Id": "trace"}): + response = admission(AppMCPServerRefreshController.post)(None, app_id=_MCP_SERVER_ADMISSION_APP_ID) + assert isinstance(response, dict) + assert response["id"] == _MCP_SERVER_ADMISSION_RECORD.id + assert [(c.account_id, c.active_workspace_id, c.trace_id) for c in contexts] == [ + (account.id, "workspace", "trace") + ] + + @staticmethod + @pytest.mark.parametrize("role", [TenantAccountRole.NORMAL, TenantAccountRole.DATASET_OPERATOR]) + def test_readonly_roles_can_read(app: Flask, monkeypatch: pytest.MonkeyPatch, role: TenantAccountRole) -> None: + account = make_account(role=role) + apply_config_overrides(monkeypatch, RBAC_ENABLED=False) + monkeypatch.setattr( + flask_admission, "current_account_with_tenant", lambda: AccountWithTenant(account, "workspace") + ) + monkeypatch.setattr(flask_admission, "enforce_rbac_checks", lambda **_kwargs: None) + monkeypatch.setattr( + mcp_server, + "application_services", + lambda: SimpleNamespace(app_mcp_servers=SimpleNamespace(get=lambda _context, _app_id: None)), + ) + with app.test_request_context(): + assert admission(AppMCPServerController.get)(None, app_id=_MCP_SERVER_ADMISSION_APP_ID) == {} + + +CONTEXT = RequestContext("request", None, "account", "workspace") +SERVER_ID = "abcdefab-1234-4567-89ab-abcdefabcdef" +_MCP_SERVER_RESPONSE_APP_ID = UUID("11111111-1111-1111-1111-111111111111") +_MCP_SERVER_RESPONSE_RECORD = AppMCPServerRecord( + id="server", + name="name", + server_code="code", + description="description", + status=AppMCPServerStatus.ACTIVE, + parameters='{"timeout": 30}', + created_at=_mcp_server_response_datetime.datetime(2024, 1, 1, tzinfo=_mcp_server_response_datetime.UTC), + updated_at=_mcp_server_response_datetime.datetime(2024, 1, 2, tzinfo=_mcp_server_response_datetime.UTC), +) + + +@dataclass +class Servers: + result: AppMCPServerRecord | None = _MCP_SERVER_RESPONSE_RECORD + error: Exception | None = None + calls: list[tuple[str, str, dict[str, object]]] = field(default_factory=list) + + def get(self, context: RequestContext, app_id: str) -> AppMCPServerRecord | None: + return self._call("get", app_id) + + def create(self, context: RequestContext, app_id: str, **kwargs: object) -> AppMCPServerRecord | None: + return self._call("create", app_id, **kwargs) + + def update(self, context: RequestContext, app_id: str, **kwargs: object) -> AppMCPServerRecord | None: + return self._call("update", app_id, **kwargs) + + def refresh(self, context: RequestContext, app_id: str) -> AppMCPServerRecord | None: + return self._call("refresh", app_id) + + def _call(self, method: str, app_id: str, **kwargs: object) -> AppMCPServerRecord | None: + self.calls.append((method, app_id, kwargs)) + if self.error is not None: + raise self.error + return self.result + + +class TestMcpServerResponseCases: + @staticmethod + @pytest.fixture + def servers(monkeypatch: pytest.MonkeyPatch) -> Servers: + servers = Servers() + monkeypatch.setattr(module, "application_services", lambda: SimpleNamespace(app_mcp_servers=servers)) + return servers + + class TestAppMCPServerResponse: + def test_parameters_json_string_parsed(self) -> None: + data: dict[str, object] = { + "id": "s1", + "name": "test", + "server_code": "code", + "description": "desc", + "status": "active", + "parameters": '{"key": "value"}', + } + resp = AppMCPServerResponse.model_validate(data) + assert resp.parameters == {"key": "value"} + + def test_parameters_invalid_json_returns_original(self) -> None: + data: dict[str, object] = { + "id": "s1", + "name": "test", + "server_code": "code", + "description": "desc", + "status": "active", + "parameters": "not-valid-json", + } + resp = AppMCPServerResponse.model_validate(data) + assert resp.parameters == "not-valid-json" + + def test_parameters_dict_passthrough(self) -> None: + data: dict[str, object] = { + "id": "s1", + "name": "test", + "server_code": "code", + "description": "desc", + "status": "active", + "parameters": {"already": "parsed"}, + } + resp = AppMCPServerResponse.model_validate(data) + assert resp.parameters == {"already": "parsed"} + + def test_parameters_json_array_parsed(self) -> None: + data: dict[str, object] = { + "id": "s1", + "name": "test", + "server_code": "code", + "description": "desc", + "status": "active", + "parameters": '["a", "b"]', + } + resp = AppMCPServerResponse.model_validate(data) + assert resp.parameters == ["a", "b"] + + def test_timestamps_normalized(self) -> None: + dt = _mcp_server_response_datetime.datetime(2024, 1, 1, 0, 0, 0, tzinfo=_mcp_server_response_datetime.UTC) + data: dict[str, object] = { + "id": "s1", + "name": "test", + "server_code": "code", + "description": "desc", + "status": "active", + "parameters": {}, + "created_at": dt, + "updated_at": dt, + } + resp = AppMCPServerResponse.model_validate(data) + assert resp.created_at == int(dt.timestamp()) + assert resp.updated_at == int(dt.timestamp()) + + def test_timestamps_none(self) -> None: + data: dict[str, object] = { + "id": "s1", + "name": "test", + "server_code": "code", + "description": "desc", + "status": "active", + "parameters": {}, + } + resp = AppMCPServerResponse.model_validate(data) + assert resp.created_at is None + assert resp.updated_at is None + + @staticmethod + def test_get_serializes_record_and_missing_server(servers: Servers) -> None: + method = unwrap(module.AppMCPServerController.get) + response = method(None, CONTEXT, _MCP_SERVER_RESPONSE_APP_ID) + assert response["parameters"] == {"timeout": 30} + assert response["created_at"] == 1704067200 + servers.result = None + assert method(None, CONTEXT, _MCP_SERVER_RESPONSE_APP_ID) == {} + + @staticmethod + def test_post_parses_input_and_returns_201(app: Flask, servers: Servers) -> None: + with app.test_request_context(method="POST", json={"parameters": {"timeout": 30}}): + response, status = unwrap(module.AppMCPServerController.post)(None, CONTEXT, _MCP_SERVER_RESPONSE_APP_ID) + assert status == 201 + assert response["id"] == "server" + assert servers.calls == [ + ("create", str(_MCP_SERVER_RESPONSE_APP_ID), {"description": None, "parameters": {"timeout": 30}}) + ] + + @staticmethod + @pytest.mark.parametrize("candidate", [SERVER_ID, SERVER_ID.upper(), UUID(SERVER_ID).hex, "{" + SERVER_ID + "}"]) + def test_put_normalizes_server_id_and_passes_status(app: Flask, servers: Servers, candidate: str) -> None: + body = {"id": candidate, "parameters": {}, "description": "updated", "status": "inactive"} + with app.test_request_context(method="PUT", json=body): + unwrap(module.AppMCPServerController.put)(None, CONTEXT, _MCP_SERVER_RESPONSE_APP_ID) + assert servers.calls == [ + ( + "update", + str(_MCP_SERVER_RESPONSE_APP_ID), + {"server_id": SERVER_ID, "description": "updated", "parameters": {}, "status": "inactive"}, + ) + ] + + @staticmethod + @pytest.mark.parametrize( + "body", + [ + pytest.param({"id": "not-a-uuid", "parameters": {}}, id="malformed-id"), + pytest.param({"id": SERVER_ID, "parameters": {}, "status": "paused"}, id="unknown-status"), + pytest.param({"id": SERVER_ID, "parameters": {}, "status": ""}, id="blank-status"), + ], + ) + def test_put_rejects_invalid_payload_before_calling_service( + app: Flask, servers: Servers, body: dict[str, object] + ) -> None: + with app.test_request_context(method="PUT", json=body), pytest.raises(UnprocessableEntity): + unwrap(module.AppMCPServerController.put)(None, CONTEXT, _MCP_SERVER_RESPONSE_APP_ID) + assert servers.calls == [] + + @staticmethod + @pytest.mark.parametrize( + "method", + [ + module.AppMCPServerController.get, + module.AppMCPServerController.post, + module.AppMCPServerController.put, + module.AppMCPServerRefreshController.post, + ], + ) + def test_unavailable_app_becomes_app_404(app: Flask, servers: Servers, method: Callable[..., object]) -> None: + servers.error = AppMCPServerAppNotFoundError() + with ( + app.test_request_context(method="POST", json={"id": SERVER_ID, "parameters": {}}), + pytest.raises(AppNotFoundError) as error, + ): + unwrap(method)(None, CONTEXT, _MCP_SERVER_RESPONSE_APP_ID) + assert error.value.description == AppNotFoundError.description + + @staticmethod + @pytest.mark.parametrize("method", [module.AppMCPServerController.put, module.AppMCPServerRefreshController.post]) + def test_missing_server_becomes_404(app: Flask, servers: Servers, method: Callable[..., object]) -> None: + servers.error = AppMCPServerNotFoundError() + with app.test_request_context(method="POST", json={"id": SERVER_ID, "parameters": {}}), pytest.raises(NotFound): + unwrap(method)(None, CONTEXT, _MCP_SERVER_RESPONSE_APP_ID) + + @staticmethod + def test_duplicate_create_becomes_409(app: Flask, servers: Servers) -> None: + servers.error = AppMCPServerAlreadyExistsError() + with app.test_request_context(method="POST", json={"parameters": {}}), pytest.raises(Conflict): + unwrap(module.AppMCPServerController.post)(None, CONTEXT, _MCP_SERVER_RESPONSE_APP_ID) + + @staticmethod + def test_routes_stay_app_scoped() -> None: + route_map = {resource.__name__: urls for resource, urls, _, _ in console_ns.resources} + assert route_map["AppMCPServerController"] == ("/apps//server",) + assert route_map["AppMCPServerRefreshController"] == ("/apps//server/refresh",) + assert not hasattr(module.AppMCPServerRefreshController, "get") diff --git a/api/tests/unit_tests/controllers/console/app/test_mcp_server_admission.py b/api/tests/unit_tests/controllers/console/app/test_mcp_server_admission.py deleted file mode 100644 index 64183af3ed3..00000000000 --- a/api/tests/unit_tests/controllers/console/app/test_mcp_server_admission.py +++ /dev/null @@ -1,128 +0,0 @@ -"""App MCP Server routes retain their role and RBAC policy under unified admission.""" - -from collections.abc import Callable, Sequence -from datetime import datetime -from inspect import unwrap -from types import SimpleNamespace -from uuid import UUID - -import pytest -from flask import Flask -from werkzeug.exceptions import Forbidden - -from controllers.common.rbac import PlainApp, RBACCheck -from controllers.console import flask_admission -from controllers.console.app import mcp_server -from controllers.console.app.mcp_server import AppMCPServerController, AppMCPServerRefreshController -from controllers.console.wraps import RBACPermission -from libs.login import AccountWithTenant -from machinery.context import RequestContext -from models.account import TenantAccountRole -from services.app.mcp_server_service import AppMCPServerRecord, AppMCPServerStatus -from tests.unit_tests.config_override import apply_config_overrides -from tests.unit_tests.model_factories import make_account - -APP_ID = UUID("11111111-1111-1111-1111-111111111111") -RECORD = AppMCPServerRecord( - id="server", - name="name", - server_code="code", - description="description", - status=AppMCPServerStatus.ACTIVE, - parameters="{}", - created_at=datetime(2024, 1, 1), - updated_at=datetime(2024, 1, 1), -) -MUTATIONS = [ - AppMCPServerController.post, - AppMCPServerController.put, - AppMCPServerRefreshController.post, -] - - -def admission(method: Callable[..., object]) -> Callable[..., object]: - return unwrap(method, stop=lambda view: "inject_request_context" in view.__code__.co_qualname) - - -@pytest.mark.parametrize("method", MUTATIONS) -@pytest.mark.parametrize("role", [TenantAccountRole.NORMAL, TenantAccountRole.DATASET_OPERATOR]) -def test_readonly_roles_cannot_mutate( - app: Flask, monkeypatch: pytest.MonkeyPatch, method: Callable[..., object], role: TenantAccountRole -) -> None: - account = make_account(role=role) - apply_config_overrides(monkeypatch, RBAC_ENABLED=False) - monkeypatch.setattr(flask_admission, "current_account_with_tenant", lambda: AccountWithTenant(account, "workspace")) - with app.test_request_context(), pytest.raises(Forbidden): - admission(method)(None, app_id=APP_ID) - - -@pytest.mark.parametrize( - ("method", "permission"), - [ - (AppMCPServerController.get, RBACPermission.APP_VIEW_LAYOUT), - (AppMCPServerController.post, RBACPermission.APP_EDIT), - (AppMCPServerController.put, RBACPermission.APP_EDIT), - (AppMCPServerRefreshController.post, RBACPermission.APP_EDIT), - ], -) -def test_rbac_checks_keep_app_scope( - app: Flask, monkeypatch: pytest.MonkeyPatch, method: Callable[..., object], permission: RBACPermission -) -> None: - account = make_account(role=TenantAccountRole.NORMAL) - apply_config_overrides(monkeypatch, RBAC_ENABLED=True) - monkeypatch.setattr(flask_admission, "current_account_with_tenant", lambda: AccountWithTenant(account, "workspace")) - called = [] - - def denied(*, tenant_id: str, account_id: str, checks: Sequence[RBACCheck], path_args: dict[str, object]) -> None: - assert (tenant_id, account_id) == ("workspace", account.id) - assert path_args == {"app_id": APP_ID} - assert len(checks) == 1 - assert checks[0].scene == permission - assert isinstance(checks[0].locator, PlainApp) - called.append(True) - raise Forbidden - - monkeypatch.setattr(flask_admission, "enforce_rbac_checks", denied) - with app.test_request_context(), pytest.raises(Forbidden): - admission(method)(None, app_id=APP_ID) - assert called == [True] - - -@pytest.mark.parametrize("role", [TenantAccountRole.OWNER, TenantAccountRole.ADMIN, TenantAccountRole.EDITOR]) -def test_edit_roles_reach_service_with_request_context( - app: Flask, monkeypatch: pytest.MonkeyPatch, role: TenantAccountRole -) -> None: - account = make_account(role=role) - apply_config_overrides(monkeypatch, RBAC_ENABLED=False) - monkeypatch.setattr(flask_admission, "current_account_with_tenant", lambda: AccountWithTenant(account, "workspace")) - monkeypatch.setattr(flask_admission, "enforce_rbac_checks", lambda **_kwargs: None) - monkeypatch.setattr(flask_admission, "get_trace_id", lambda: None) - contexts: list[RequestContext] = [] - - def refresh(context: RequestContext, _app_id: str) -> AppMCPServerRecord: - contexts.append(context) - return RECORD - - monkeypatch.setattr( - mcp_server, "application_services", lambda: SimpleNamespace(app_mcp_servers=SimpleNamespace(refresh=refresh)) - ) - with app.test_request_context(method="POST", headers={"X-Trace-Id": "trace"}): - response = admission(AppMCPServerRefreshController.post)(None, app_id=APP_ID) - assert isinstance(response, dict) - assert response["id"] == RECORD.id - assert [(c.account_id, c.active_workspace_id, c.trace_id) for c in contexts] == [(account.id, "workspace", "trace")] - - -@pytest.mark.parametrize("role", [TenantAccountRole.NORMAL, TenantAccountRole.DATASET_OPERATOR]) -def test_readonly_roles_can_read(app: Flask, monkeypatch: pytest.MonkeyPatch, role: TenantAccountRole) -> None: - account = make_account(role=role) - apply_config_overrides(monkeypatch, RBAC_ENABLED=False) - monkeypatch.setattr(flask_admission, "current_account_with_tenant", lambda: AccountWithTenant(account, "workspace")) - monkeypatch.setattr(flask_admission, "enforce_rbac_checks", lambda **_kwargs: None) - monkeypatch.setattr( - mcp_server, - "application_services", - lambda: SimpleNamespace(app_mcp_servers=SimpleNamespace(get=lambda _context, _app_id: None)), - ) - with app.test_request_context(): - assert admission(AppMCPServerController.get)(None, app_id=APP_ID) == {} diff --git a/api/tests/unit_tests/controllers/console/app/test_mcp_server_response.py b/api/tests/unit_tests/controllers/console/app/test_mcp_server_response.py deleted file mode 100644 index f2e6e5c48e8..00000000000 --- a/api/tests/unit_tests/controllers/console/app/test_mcp_server_response.py +++ /dev/null @@ -1,234 +0,0 @@ -import datetime -from collections.abc import Callable -from dataclasses import dataclass, field -from inspect import unwrap -from types import SimpleNamespace -from uuid import UUID - -import pytest -from flask import Flask -from werkzeug.exceptions import Conflict, NotFound, UnprocessableEntity - -from controllers.console import console_ns -from controllers.console.app import mcp_server as module -from controllers.console.app.error import AppNotFoundError -from controllers.console.app.mcp_server import AppMCPServerResponse -from machinery.context import RequestContext -from services.app.mcp_server_service import ( - AppMCPServerAlreadyExistsError, - AppMCPServerAppNotFoundError, - AppMCPServerNotFoundError, - AppMCPServerRecord, - AppMCPServerStatus, -) - -CONTEXT = RequestContext("request", None, "account", "workspace") -SERVER_ID = "abcdefab-1234-4567-89ab-abcdefabcdef" -APP_ID = UUID("11111111-1111-1111-1111-111111111111") -RECORD = AppMCPServerRecord( - id="server", - name="name", - server_code="code", - description="description", - status=AppMCPServerStatus.ACTIVE, - parameters='{"timeout": 30}', - created_at=datetime.datetime(2024, 1, 1, tzinfo=datetime.UTC), - updated_at=datetime.datetime(2024, 1, 2, tzinfo=datetime.UTC), -) - - -@dataclass -class Servers: - result: AppMCPServerRecord | None = RECORD - error: Exception | None = None - calls: list[tuple[str, str, dict[str, object]]] = field(default_factory=list) - - def get(self, context: RequestContext, app_id: str) -> AppMCPServerRecord | None: - return self._call("get", app_id) - - def create(self, context: RequestContext, app_id: str, **kwargs: object) -> AppMCPServerRecord | None: - return self._call("create", app_id, **kwargs) - - def update(self, context: RequestContext, app_id: str, **kwargs: object) -> AppMCPServerRecord | None: - return self._call("update", app_id, **kwargs) - - def refresh(self, context: RequestContext, app_id: str) -> AppMCPServerRecord | None: - return self._call("refresh", app_id) - - def _call(self, method: str, app_id: str, **kwargs: object) -> AppMCPServerRecord | None: - self.calls.append((method, app_id, kwargs)) - if self.error is not None: - raise self.error - return self.result - - -@pytest.fixture -def servers(monkeypatch: pytest.MonkeyPatch) -> Servers: - servers = Servers() - monkeypatch.setattr(module, "application_services", lambda: SimpleNamespace(app_mcp_servers=servers)) - return servers - - -class TestAppMCPServerResponse: - def test_parameters_json_string_parsed(self) -> None: - data: dict[str, object] = { - "id": "s1", - "name": "test", - "server_code": "code", - "description": "desc", - "status": "active", - "parameters": '{"key": "value"}', - } - resp = AppMCPServerResponse.model_validate(data) - assert resp.parameters == {"key": "value"} - - def test_parameters_invalid_json_returns_original(self) -> None: - data: dict[str, object] = { - "id": "s1", - "name": "test", - "server_code": "code", - "description": "desc", - "status": "active", - "parameters": "not-valid-json", - } - resp = AppMCPServerResponse.model_validate(data) - assert resp.parameters == "not-valid-json" - - def test_parameters_dict_passthrough(self) -> None: - data: dict[str, object] = { - "id": "s1", - "name": "test", - "server_code": "code", - "description": "desc", - "status": "active", - "parameters": {"already": "parsed"}, - } - resp = AppMCPServerResponse.model_validate(data) - assert resp.parameters == {"already": "parsed"} - - def test_parameters_json_array_parsed(self) -> None: - data: dict[str, object] = { - "id": "s1", - "name": "test", - "server_code": "code", - "description": "desc", - "status": "active", - "parameters": '["a", "b"]', - } - resp = AppMCPServerResponse.model_validate(data) - assert resp.parameters == ["a", "b"] - - def test_timestamps_normalized(self) -> None: - dt = datetime.datetime(2024, 1, 1, 0, 0, 0, tzinfo=datetime.UTC) - data: dict[str, object] = { - "id": "s1", - "name": "test", - "server_code": "code", - "description": "desc", - "status": "active", - "parameters": {}, - "created_at": dt, - "updated_at": dt, - } - resp = AppMCPServerResponse.model_validate(data) - assert resp.created_at == int(dt.timestamp()) - assert resp.updated_at == int(dt.timestamp()) - - def test_timestamps_none(self) -> None: - data: dict[str, object] = { - "id": "s1", - "name": "test", - "server_code": "code", - "description": "desc", - "status": "active", - "parameters": {}, - } - resp = AppMCPServerResponse.model_validate(data) - assert resp.created_at is None - assert resp.updated_at is None - - -def test_get_serializes_record_and_missing_server(servers: Servers) -> None: - method = unwrap(module.AppMCPServerController.get) - response = method(None, CONTEXT, APP_ID) - assert response["parameters"] == {"timeout": 30} - assert response["created_at"] == 1704067200 - servers.result = None - assert method(None, CONTEXT, APP_ID) == {} - - -def test_post_parses_input_and_returns_201(app: Flask, servers: Servers) -> None: - with app.test_request_context(method="POST", json={"parameters": {"timeout": 30}}): - response, status = unwrap(module.AppMCPServerController.post)(None, CONTEXT, APP_ID) - assert status == 201 - assert response["id"] == "server" - assert servers.calls == [("create", str(APP_ID), {"description": None, "parameters": {"timeout": 30}})] - - -@pytest.mark.parametrize("candidate", [SERVER_ID, SERVER_ID.upper(), UUID(SERVER_ID).hex, "{" + SERVER_ID + "}"]) -def test_put_normalizes_server_id_and_passes_status(app: Flask, servers: Servers, candidate: str) -> None: - body = {"id": candidate, "parameters": {}, "description": "updated", "status": "inactive"} - with app.test_request_context(method="PUT", json=body): - unwrap(module.AppMCPServerController.put)(None, CONTEXT, APP_ID) - assert servers.calls == [ - ( - "update", - str(APP_ID), - {"server_id": SERVER_ID, "description": "updated", "parameters": {}, "status": "inactive"}, - ) - ] - - -@pytest.mark.parametrize( - "body", - [ - pytest.param({"id": "not-a-uuid", "parameters": {}}, id="malformed-id"), - pytest.param({"id": SERVER_ID, "parameters": {}, "status": "paused"}, id="unknown-status"), - pytest.param({"id": SERVER_ID, "parameters": {}, "status": ""}, id="blank-status"), - ], -) -def test_put_rejects_invalid_payload_before_calling_service( - app: Flask, servers: Servers, body: dict[str, object] -) -> None: - with app.test_request_context(method="PUT", json=body), pytest.raises(UnprocessableEntity): - unwrap(module.AppMCPServerController.put)(None, CONTEXT, APP_ID) - assert servers.calls == [] - - -@pytest.mark.parametrize( - "method", - [ - module.AppMCPServerController.get, - module.AppMCPServerController.post, - module.AppMCPServerController.put, - module.AppMCPServerRefreshController.post, - ], -) -def test_unavailable_app_becomes_app_404(app: Flask, servers: Servers, method: Callable[..., object]) -> None: - servers.error = AppMCPServerAppNotFoundError() - with ( - app.test_request_context(method="POST", json={"id": SERVER_ID, "parameters": {}}), - pytest.raises(AppNotFoundError) as error, - ): - unwrap(method)(None, CONTEXT, APP_ID) - assert error.value.description == AppNotFoundError.description - - -@pytest.mark.parametrize("method", [module.AppMCPServerController.put, module.AppMCPServerRefreshController.post]) -def test_missing_server_becomes_404(app: Flask, servers: Servers, method: Callable[..., object]) -> None: - servers.error = AppMCPServerNotFoundError() - with app.test_request_context(method="POST", json={"id": SERVER_ID, "parameters": {}}), pytest.raises(NotFound): - unwrap(method)(None, CONTEXT, APP_ID) - - -def test_duplicate_create_becomes_409(app: Flask, servers: Servers) -> None: - servers.error = AppMCPServerAlreadyExistsError() - with app.test_request_context(method="POST", json={"parameters": {}}), pytest.raises(Conflict): - unwrap(module.AppMCPServerController.post)(None, CONTEXT, APP_ID) - - -def test_routes_stay_app_scoped() -> None: - route_map = {resource.__name__: urls for resource, urls, _, _ in console_ns.resources} - assert route_map["AppMCPServerController"] == ("/apps//server",) - assert route_map["AppMCPServerRefreshController"] == ("/apps//server/refresh",) - assert not hasattr(module.AppMCPServerRefreshController, "get") diff --git a/api/tests/unit_tests/controllers/console/app/test_message.py b/api/tests/unit_tests/controllers/console/app/test_message.py new file mode 100644 index 00000000000..72d39692a8a --- /dev/null +++ b/api/tests/unit_tests/controllers/console/app/test_message.py @@ -0,0 +1,945 @@ +from __future__ import annotations + +import json +from collections.abc import Iterator, Sequence +from dataclasses import dataclass, field +from datetime import UTC, datetime +from decimal import Decimal +from http import HTTPStatus +from inspect import unwrap +from typing import cast +from unittest.mock import MagicMock +from uuid import uuid4 + +import pytest +from flask import Flask +from sqlalchemy import Connection, Engine, event, func, select, text, update +from sqlalchemy.orm import Session, SessionTransaction, sessionmaker +from werkzeug.test import TestResponse + +import controllers.console.wraps as console_wraps +from controllers.common.rbac import RBACPermission, RBACResourceScope +from controllers.common.rbac.checks import RBACService +from controllers.console.app import message as controller +from controllers.console.app import message as message_module +from core.app.entities.app_invoke_entities import InvokeFrom +from core.errors.error import ModelCurrentlyNotSupportError, ProviderTokenNotInitError, QuotaExceededError +from core.model_manager import ModelInstance +from core.ops.ops_trace_manager import TraceTask +from extensions.application_services.app import AppServices +from extensions.application_services.rbac import RBACServices +from extensions.ext_database import db +from graphon.model_runtime.entities.llm_entities import LLMResult, LLMUsage +from graphon.model_runtime.entities.message_entities import AssistantPromptMessage, PromptMessage +from graphon.model_runtime.entities.model_entities import ModelType +from graphon.model_runtime.errors.invoke import InvokeError +from libs import login +from libs.external_api import ExternalApi +from models.account import Account, Tenant +from models.agent import ( + Agent, + AgentConfigDraft, + AgentConfigDraftType, + AgentConfigVersionKind, + AgentDebugConversation, + AgentScope, + AgentSource, + AgentStatus, + AgentWorkspaceBinding, +) +from models.agent_config_entities import AgentSoulConfig +from models.enums import ConversationFromSource, CreatorUserRole, FeedbackFromSource, FeedbackRating +from models.model import ( + App, + AppAnnotationHitHistory, + AppMode, + AppModelConfig, + Conversation, + Message, + MessageAgentThought, + MessageAnnotation, + MessageFeedback, +) +from repositories.app.agent_app_repository import AgentAppRepository +from services import message_service +from services.app.agent_app_service import AgentAppAccessService +from services.app_definition_query_service import AppDefinitionUnavailableError +from services.errors.conversation import ConversationNotExistsError +from services.errors.message import MessageNotExistsError, SuggestedQuestionsAfterAnswerDisabledError +from services.message_suggested_questions_adapters import MessageSuggestedQuestionsRuntime +from services.message_suggested_questions_service import ( + MessageSuggestedQuestions, + SuggestedQuestionsAccount, + SuggestedQuestionsActor, + SuggestedQuestionsActorNotFoundError, +) +from tests.helpers.unit.config_override import apply_config_overrides +from tests.helpers.unit.model_factories import make_account, make_app, make_conversation, make_message +from tests.helpers.unit.rbac_fakes import RBACDomain + + +def _account() -> Account: + account = Account(name="Account", email="account@example.com") + account.id = "account-1" + return account + + +def _app(*, app_id: str = "app-1") -> App: + return App( + id=app_id, + tenant_id="tenant-1", + name="Chat App", + description="", + mode=AppMode.CHAT, + enable_site=True, + enable_api=True, + max_active_requests=0, + ) + + +def _persist_message(session: Session, *, message_id: str, app_id: str = "app-1") -> Message: + conversation = Conversation( + app_id=app_id, + app_model_config_id=None, + model_provider=None, + override_model_configs=None, + model_id=None, + mode=AppMode.CHAT, + name="Conversation", + _inputs={}, + introduction="", + system_instruction="", + system_instruction_tokens=0, + status="normal", + invoke_from=InvokeFrom.DEBUGGER, + from_source=ConversationFromSource.CONSOLE, + from_end_user_id=None, + from_account_id="account-1", + ) + conversation.id = "conversation-1" + message = Message( + app_id=app_id, + conversation_id=conversation.id, + _inputs={}, + query="query", + message="", + message_tokens=0, + message_unit_price=0, + message_price_unit=0, + answer="answer", + answer_tokens=0, + answer_unit_price=0, + answer_price_unit=0, + provider_response_latency=0, + total_price=0, + currency="USD", + invoke_from=InvokeFrom.DEBUGGER, + from_source=ConversationFromSource.CONSOLE, + from_end_user_id=None, + from_account_id="account-1", + app_mode=AppMode.CHAT, + ) + message.id = message_id + session.add_all([_account(), _app(app_id=app_id), conversation, message]) + session.flush() + return message + + +class TestMessageApiCases: + @staticmethod + def test_app_message_routes_pass_injected_session( + app: Flask, monkeypatch: pytest.MonkeyPatch, unbound_session: Session + ) -> None: + session = unbound_session + current_user = _account() + app_model = _app() + message_id = "550e8400-e29b-41d4-a716-446655440000" + list_messages = MagicMock(return_value={"data": []}) + update_feedback = MagicMock(return_value={"result": "success"}) + get_message_detail = MagicMock(return_value={"id": message_id}) + monkeypatch.setattr(message_module, "_list_chat_messages", list_messages) + monkeypatch.setattr(message_module, "_update_message_feedback", update_feedback) + monkeypatch.setattr(message_module, "_get_message_detail", get_message_detail) + + assert unwrap(message_module.ChatMessageListApi.get)( + message_module.ChatMessageListApi(), + message_module.ChatMessagesQuery(conversation_id="550e8400-e29b-41d4-a716-446655440001"), + session, + current_user, + app_model, + ) == {"data": []} + assert unwrap(message_module.MessageFeedbackApi.post)( + message_module.MessageFeedbackApi(), + message_module.AppMessageFeedbackPayload(message_id=message_id, rating="like"), + session, + current_user, + app_model, + ) == {"result": "success"} + assert unwrap(message_module.MessageApi.get)(message_module.MessageApi(), session, app_model, message_id) == { + "id": message_id + } + + assert list_messages.call_args.kwargs["session"] is session + assert update_feedback.call_args.kwargs["session"] is session + assert get_message_detail.call_args.kwargs["session"] is session + + @staticmethod + def test_update_message_feedback_commits_injected_session(app: Flask, sqlite_session: Session) -> None: + message_id = "550e8400-e29b-41d4-a716-446655440000" + message = _persist_message(sqlite_session, message_id=message_id) + feedback = MessageFeedback( + app_id=message.app_id, + conversation_id=message.conversation_id, + message_id=message.id, + rating=FeedbackRating.DISLIKE, + from_source=FeedbackFromSource.ADMIN, + from_account_id="account-1", + ) + sqlite_session.add(feedback) + sqlite_session.commit() + session = sqlite_session + commits: list[str] = [] + event.listen(session, "after_commit", lambda _session: commits.append("commit")) + + with app.test_request_context(json={"message_id": message_id, "rating": "like", "content": "helpful"}): + result = message_module._update_message_feedback( + args=message_module.AppMessageFeedbackPayload( + message_id=message_id, + rating="like", + content="helpful", + ), + session=session, + current_user=_account(), + app_model=_app(), + ) + + assert result == {"result": "success"} + updated_feedback = sqlite_session.get(MessageFeedback, feedback.id) + assert updated_feedback is not None + assert updated_feedback.rating == FeedbackRating.LIKE + assert updated_feedback.content == "helpful" + assert commits == ["commit"] + + @staticmethod + def test_get_message_detail_uses_injected_session(monkeypatch: pytest.MonkeyPatch, sqlite_session: Session) -> None: + message_id = "550e8400-e29b-41d4-a716-446655440000" + message = _persist_message(sqlite_session, message_id=message_id) + response_source = object() + response_source_factory = MagicMock(return_value=response_source) + session = sqlite_session + monkeypatch.setattr(message_module, "attach_message_extra_contents", MagicMock()) + monkeypatch.setattr(message_module, "MessageResponseSource", response_source_factory) + monkeypatch.setattr(message_module, "dump_response", lambda _model, value: value) + + result = message_module._get_message_detail( + session=session, + app_model=_app(), + message_id=message_id, + ) + + assert result is response_source + response_source_factory.assert_called_once_with(message, session=session) + + @staticmethod + def test_message_response_source_uses_caller_session_for_nested_fields(sqlite_session: Session) -> None: + session = sqlite_session + message_id = "550e8400-e29b-41d4-a716-446655440000" + message = _persist_message(sqlite_session, message_id=message_id) + message.inputs = {"topic": "support"} + feedback = MessageFeedback( + app_id=message.app_id, + conversation_id=message.conversation_id, + message_id=message.id, + rating=FeedbackRating.LIKE, + from_source=FeedbackFromSource.USER, + from_account_id="account-1", + ) + annotation = MessageAnnotation( + app_id=message.app_id, + question="Question", + content="Answer", + account_id="account-1", + conversation_id=message.conversation_id, + message_id=message.id, + ) + annotation.id = "annotation-1" + annotation_hit = AppAnnotationHitHistory( + app_id=message.app_id, + annotation_id=annotation.id, + source="annotation", + question="Question", + account_id="account-1", + score=0.0, + message_id=message.id, + annotation_question="Question", + annotation_content="Answer", + ) + thought = MessageAgentThought( + message_id=message.id, + position=1, + created_by_role=CreatorUserRole.ACCOUNT, + created_by="account-1", + thought="Consider the request", + tool_labels_str="{}", + tool_meta_str="{}", + ) + sqlite_session.add_all([feedback, annotation, annotation_hit, thought]) + sqlite_session.commit() + + source = message_module.MessageResponseSource(message, session=session) + + assert source.inputs == {"topic": "support"} + assert source.user_feedback is feedback + assert source.feedbacks[0].from_account.id == "account-1" + annotation_source = source.annotation + assert annotation_source is not None + assert annotation_source.account.id == "account-1" + annotation_hit_history_source = source.annotation_hit_history + assert annotation_hit_history_source is not None + assert annotation_hit_history_source.annotation_create_account.id == "account-1" + assert source.agent_thoughts == [thought] + assert source.message_files == [] + + @staticmethod + def test_chat_messages_query_valid(app: Flask, monkeypatch: pytest.MonkeyPatch) -> None: + """Test valid ChatMessagesQuery with all fields.""" + query = message_module.ChatMessagesQuery( + conversation_id="550e8400-e29b-41d4-a716-446655440000", + first_id="550e8400-e29b-41d4-a716-446655440001", + limit=50, + ) + assert query.limit == 50 + + @staticmethod + def test_chat_messages_query_defaults(app: Flask, monkeypatch: pytest.MonkeyPatch) -> None: + """Test ChatMessagesQuery with defaults.""" + query = message_module.ChatMessagesQuery(conversation_id="550e8400-e29b-41d4-a716-446655440000") + assert query.first_id is None + assert query.limit == 20 + + @staticmethod + def test_chat_messages_query_empty_first_id(app: Flask, monkeypatch: pytest.MonkeyPatch) -> None: + """Test ChatMessagesQuery converts empty first_id to None.""" + query = message_module.ChatMessagesQuery( + conversation_id="550e8400-e29b-41d4-a716-446655440000", + first_id="", + ) + assert query.first_id is None + + @staticmethod + def test_message_feedback_payload_valid_like(app: Flask, monkeypatch: pytest.MonkeyPatch) -> None: + """Test AppMessageFeedbackPayload with like rating.""" + payload = message_module.AppMessageFeedbackPayload( + message_id="550e8400-e29b-41d4-a716-446655440000", + rating="like", + content="Good answer", + ) + assert payload.rating == "like" + assert payload.content == "Good answer" + + @staticmethod + def test_message_feedback_payload_valid_dislike(app: Flask, monkeypatch: pytest.MonkeyPatch) -> None: + """Test AppMessageFeedbackPayload with dislike rating.""" + payload = message_module.AppMessageFeedbackPayload( + message_id="550e8400-e29b-41d4-a716-446655440000", + rating="dislike", + ) + assert payload.rating == "dislike" + + @staticmethod + def test_message_feedback_payload_no_rating(app: Flask, monkeypatch: pytest.MonkeyPatch) -> None: + """Test AppMessageFeedbackPayload without rating.""" + payload = message_module.AppMessageFeedbackPayload(message_id="550e8400-e29b-41d4-a716-446655440000") + assert payload.rating is None + + @staticmethod + def test_annotation_count_response(app: Flask, monkeypatch: pytest.MonkeyPatch) -> None: + """Test AnnotationCountResponse creation.""" + response = message_module.AnnotationCountResponse(count=10) + assert response.count == 10 + + @staticmethod + def test_suggested_questions_response(app: Flask, monkeypatch: pytest.MonkeyPatch) -> None: + """Test SuggestedQuestionsResponse creation.""" + response = message_module.SuggestedQuestionsResponse(data=["What is AI?", "How does ML work?"]) + assert len(response.data) == 2 + assert response.data[0] == "What is AI?" + + @staticmethod + def test_message_detail_response_normalizes_aliases_and_timestamp( + app: Flask, monkeypatch: pytest.MonkeyPatch + ) -> None: + """Test MessageDetailResponse normalizes alias fields and datetime timestamps.""" + created_at = datetime(2026, 1, 2, 3, 4, 5, tzinfo=UTC) + response = message_module.MessageDetailResponse.model_validate( + { + "id": "550e8400-e29b-41d4-a716-446655440000", + "conversation_id": "550e8400-e29b-41d4-a716-446655440001", + "inputs": {"foo": "bar"}, + "query": "hello", + "message": [{"text": "hello"}], + "message_tokens": 7, + "answer": "world", + "answer_tokens": 11, + "provider_response_latency": 1.25, + "from_source": "user", + "from_end_user_id": None, + "from_account_id": "550e8400-e29b-41d4-a716-446655440002", + "feedbacks": [], + "workflow_run_id": None, + "annotation": None, + "annotation_hit_history": None, + "status": "normal", + "created_at": created_at, + "agent_thoughts": [], + "message_files": [], + "message_metadata_dict": {"token_usage": 3}, + "error": None, + "parent_message_id": None, + "extra_contents": [], + } + ) + assert response.answer == "world" + assert response.message_tokens == 7 + assert response.answer_tokens == 11 + assert response.provider_response_latency == 1.25 + assert response.metadata == {"token_usage": 3} + assert response.created_at == int(created_at.timestamp()) + assert response.model_dump(mode="json") == { + "id": "550e8400-e29b-41d4-a716-446655440000", + "conversation_id": "550e8400-e29b-41d4-a716-446655440001", + "inputs": {"foo": "bar"}, + "query": "hello", + "message": [{"text": "hello"}], + "message_tokens": 7, + "answer": "world", + "answer_tokens": 11, + "provider_response_latency": 1.25, + "from_source": "user", + "from_end_user_id": None, + "from_account_id": "550e8400-e29b-41d4-a716-446655440002", + "feedbacks": [], + "workflow_run_id": None, + "annotation": None, + "annotation_hit_history": None, + "created_at": int(created_at.timestamp()), + "agent_thoughts": [], + "message_files": [], + "metadata": {"token_usage": 3}, + "retriever_resources": [], + "status": "normal", + "error": None, + "parent_message_id": None, + "extra_contents": [], + } + + +"""Console suggested-question admission and actor/configuration isolation over SQLite.""" + + +@dataclass +class _Questions: + questions: list[str] = field(default_factory=lambda: ["Next?"]) + failure: Exception | None = None + calls: list[tuple[str, str, str, SuggestedQuestionsActor, str]] = field(default_factory=list) + + def get_suggested_questions( + self, + *, + app_id: str, + app_owner_tenant_id: str, + expected_app_mode: str, + actor: SuggestedQuestionsActor, + message_id: str, + ) -> list[str]: + self.calls.append((app_id, app_owner_tenant_id, expected_app_mode, actor, message_id)) + if self.failure is not None: + raise self.failure + return self.questions + + +@dataclass +class _Agents: + access: AgentAppAccessService + + +@dataclass +class _Services: + apps: AppServices + agent_apps: _Agents + message_suggested_questions: MessageSuggestedQuestions + rbac: RBACServices + + +@dataclass +class _Provider: + tenant_id: str + admission_sessions: list[Session] + prompts: list[str] = field(default_factory=list) + traces: list[TraceTask] = field(default_factory=list) + + def get_default_model_instance(self, *, tenant_id: str, model_type: ModelType) -> ModelInstance: + assert tenant_id == self.tenant_id + assert model_type == ModelType.LLM + assert self.admission_sessions + assert all(not session.in_transaction() and not session.identity_map for session in self.admission_sessions) + return cast(ModelInstance, self) + + def get_llm_num_tokens(self, prompt_messages: Sequence[PromptMessage]) -> int: + return len(prompt_messages) + + def get_model_schema(self) -> None: + raise NotImplementedError + + def invoke_llm( + self, + *, + prompt_messages: list[PromptMessage], + model_parameters: dict[str, object], + stop: list[str], + stream: bool, + ) -> LLMResult: + assert model_parameters == {"max_tokens": 256, "temperature": 0.0} + assert stop == [] + assert stream is False + self.prompts.append(prompt_messages[0].get_text_content()) + return LLMResult( + model="test-model", message=AssistantPromptMessage(content='["Next?"]'), usage=LLMUsage.empty_usage() + ) + + def add_trace_task(self, task: TraceTask) -> None: + self.traces.append(task) + + +@dataclass +class _Harness: + flask_app: Flask + target: App + account: Account + conversation: Conversation + message: Message + agent_id: str + factory: sessionmaker[Session] + services: _Services + provider: _Provider + permission_calls: list[tuple[str, str, RBACPermission, RBACResourceScope | None, str | None]] + allowed: bool = True + + def get(self, route: str = "app") -> TestResponse: + resource = f"apps/{self.target.id}" if route == "app" else f"agent/{self.agent_id}" + return self.flask_app.test_client().get(f"/{resource}/chat-messages/{self.message.id}/suggested-questions") + + def add_agent(self, *, scope: AgentScope = AgentScope.ROSTER, backing: bool = True) -> None: + with self.factory.begin() as session: + session.execute(update(App).where(App.id == self.target.id).values(mode=AppMode.AGENT)) + session.add( + Agent( + id=self.agent_id, + tenant_id=self.target.tenant_id, + name="Agent", + scope=scope, + source=AgentSource.AGENT_APP if scope == AgentScope.ROSTER else AgentSource.WORKFLOW, + status=AgentStatus.ACTIVE, + app_id=self.target.id if scope == AgentScope.ROSTER else str(uuid4()), + backing_app_id=self.target.id if scope == AgentScope.WORKFLOW_ONLY and backing else None, + workflow_id=str(uuid4()) if scope == AgentScope.WORKFLOW_ONLY else None, + workflow_node_id="node" if scope == AgentScope.WORKFLOW_ONLY else None, + ) + ) + + +class TestMessageSuggestedQuestionsCases: + @staticmethod + @pytest.fixture + def harness( + monkeypatch: pytest.MonkeyPatch, + sqlite_engine: Engine, + sqlite_session_factory: sessionmaker[Session], + app_services: AppServices, + rbac_domain: RBACDomain, + ) -> Iterator[_Harness]: + target = make_app(app_id=str(uuid4()), tenant_id=str(uuid4())) + account = make_account(account_id=str(uuid4())) + tenant = Tenant(name="Workspace") + tenant.id = target.tenant_id + account._current_tenant = tenant + conversation = make_conversation( + conversation_id=str(uuid4()), + app_id=target.id, + from_source=ConversationFromSource.CONSOLE, + from_account_id=account.id, + inputs={}, + ) + config = AppModelConfig( + app_id=target.id, + suggested_questions_after_answer=json.dumps({"enabled": True, "prompt": "Console questions"}), + ) + message = make_message( + message_id=str(uuid4()), + app_id=target.id, + conversation_id=conversation.id, + inputs={}, + message={}, + query="How?", + answer="Like this.", + message_unit_price=Decimal(0), + answer_unit_price=Decimal(0), + currency="USD", + from_source=ConversationFromSource.CONSOLE, + from_account_id=account.id, + ) + with sqlite_session_factory.begin() as session: + session.add_all([target, account, conversation, config, message]) + session.flush() + target.app_model_config_id = config.id + conversation.app_model_config_id = config.id + sessions: list[Session] = [] + + def track_session(session: Session, _transaction: SessionTransaction, _connection: Connection) -> None: + sessions.append(session) + + event.listen(sqlite_session_factory, "after_begin", track_session) + provider = _Provider(target.tenant_id, sessions) + + def manager(*, tenant_id: str) -> _Provider: + assert tenant_id == target.tenant_id + return provider + + def trace_manager(*, app_id: str) -> _Provider: + assert app_id == target.id + return provider + + monkeypatch.setattr(message_service.ModelManager, "for_tenant", manager) + monkeypatch.setattr(message_service, "TraceQueueManager", trace_manager) + monkeypatch.setattr(login, "current_user", account) + monkeypatch.setattr(console_wraps, "_is_setup_completed", lambda: True) + apply_config_overrides(monkeypatch, LOGIN_DISABLED=True, RBAC_ENABLED=False) + flask_app = Flask(__name__) + flask_app.config.update( + TESTING=True, RESTX_ERROR_404_HELP=False, SQLALCHEMY_DATABASE_URI=str(sqlite_engine.url) + ) + db.init_app(flask_app) + services = _Services( + app_services, + _Agents(AgentAppAccessService(references=AgentAppRepository(session_factory=sqlite_session_factory))), + MessageSuggestedQuestionsRuntime(session_factory=sqlite_session_factory), + rbac_domain.rbac, + ) + flask_app.extensions["application_services"] = services + api = ExternalApi(flask_app) + api.add_resource( + controller.MessageSuggestedQuestionApi, + "/apps//chat-messages//suggested-questions", + ) + api.add_resource( + controller.AgentMessageSuggestedQuestionApi, + "/agent//chat-messages//suggested-questions", + ) + result = _Harness( + flask_app, + target, + account, + conversation, + message, + str(uuid4()), + sqlite_session_factory, + services, + provider, + [], + ) + + def permission( + tenant_id: str, + account_id: str, + *, + scene: RBACPermission, + resource_type: RBACResourceScope | None, + resource_id: str | None, + ) -> bool: + result.permission_calls.append((tenant_id, account_id, scene, resource_type, resource_id)) + return result.allowed + + monkeypatch.setattr(RBACService.CheckAccess, "check", permission) + yield result + event.remove(sqlite_session_factory, "after_begin", track_session) + with flask_app.app_context(): + db.session.remove() + db.engine.dispose() + + @staticmethod + @pytest.mark.parametrize("mode", [AppMode.CHAT, AppMode.AGENT_CHAT, AppMode.ADVANCED_CHAT, AppMode.AGENT]) + @pytest.mark.parametrize("questions", [[], ["Next?", "More?"]]) + def test_app_reference_and_account_are_passed_to_shared_runtime( + harness: _Harness, mode: AppMode, questions: list[str] + ) -> None: + with harness.factory.begin() as session: + session.execute(update(App).where(App.id == harness.target.id).values(mode=mode)) + port = _Questions(questions=questions) + harness.services.message_suggested_questions = port + response = harness.get() + assert response.status_code == HTTPStatus.OK + assert response.headers["Content-Type"] == "application/json" + assert response.get_json() == {"data": questions} + assert port.calls == [ + ( + harness.target.id, + harness.target.tenant_id, + mode, + SuggestedQuestionsAccount(harness.account.id, "debugger"), + harness.message.id, + ) + ] + assert all( + not session.in_transaction() and not session.identity_map for session in harness.provider.admission_sessions + ) + + @staticmethod + @pytest.mark.parametrize("change", ["tenant", "status", "missing", "workflow", "completion", "hidden"]) + def test_app_admission_rejects_unavailable_or_unsupported_apps(harness: _Harness, change: str) -> None: + if change == "hidden": + harness.add_agent(scope=AgentScope.WORKFLOW_ONLY) + else: + field_name, value = { + "tenant": ("tenant_id", str(uuid4())), + "status": ("status", "archived"), + "missing": ("id", str(uuid4())), + "workflow": ("mode", "workflow"), + "completion": ("mode", "completion"), + }[change] + with harness.factory.begin() as session: + if change == "status": + session.execute( + text("UPDATE apps SET status = :status WHERE id = :id"), + {"status": value, "id": harness.target.id}, + ) + else: + session.execute(update(App).where(App.id == harness.target.id).values({field_name: value})) + port = _Questions() + harness.services.message_suggested_questions = port + response = harness.get() + assert response.status_code == HTTPStatus.NOT_FOUND + assert response.get_json()["code"] == "app_not_found" + assert not port.calls + + @staticmethod + @pytest.mark.parametrize("scope", [AgentScope.ROSTER, AgentScope.WORKFLOW_ONLY]) + def test_agent_resolves_existing_backing_app(harness: _Harness, scope: AgentScope) -> None: + harness.add_agent(scope=scope) + port = _Questions() + harness.services.message_suggested_questions = port + response = harness.get("agent") + assert response.status_code == HTTPStatus.OK + assert port.calls == [ + ( + harness.target.id, + harness.target.tenant_id, + AppMode.AGENT, + SuggestedQuestionsAccount(harness.account.id, "debugger"), + harness.message.id, + ) + ] + assert all( + not session.in_transaction() and not session.identity_map for session in harness.provider.admission_sessions + ) + + @staticmethod + @pytest.mark.parametrize("change", ["missing_backing", "tenant", "agent_status", "app_status", "app_mode"]) + def test_agent_rejects_unavailable_runtime_without_creating_app(harness: _Harness, change: str) -> None: + harness.add_agent(scope=AgentScope.WORKFLOW_ONLY, backing=change != "missing_backing") + if change != "missing_backing": + model, key, value = { + "tenant": (Agent, "tenant_id", str(uuid4())), + "agent_status": (Agent, "status", AgentStatus.ARCHIVED), + "app_status": (App, "status", "archived"), + "app_mode": (App, "mode", AppMode.CHAT), + }[change] + with harness.factory.begin() as session: + if change == "app_status": + session.execute( + text("UPDATE apps SET status = :status WHERE id = :id"), + {"status": value, "id": harness.target.id}, + ) + else: + session.execute(update(model).values({key: value})) + response = harness.get("agent") + assert response.status_code == HTTPStatus.NOT_FOUND + assert response.get_json()["code"] == "agent_not_found_error" + with harness.factory() as session: + assert session.scalar(select(func.count()).select_from(App)) == 1 + assert session.scalar(select(func.count()).select_from(AppModelConfig)) == 1 + if change == "missing_backing": + assert session.scalar(select(Agent.backing_app_id)) is None + + @staticmethod + @pytest.mark.parametrize("route", ["app", "agent"]) + @pytest.mark.parametrize("allowed", [True, False]) + def test_real_rbac_admission_keeps_resource_scene( + harness: _Harness, monkeypatch: pytest.MonkeyPatch, route: str, allowed: bool + ) -> None: + apply_config_overrides(monkeypatch, RBAC_ENABLED=True) + if route == "agent": + harness.add_agent() + harness.allowed = allowed + port = _Questions() + harness.services.message_suggested_questions = port + response = harness.get(route) + assert response.status_code == (HTTPStatus.OK if allowed else HTTPStatus.FORBIDDEN) + assert harness.permission_calls == [ + ( + harness.target.tenant_id, + harness.account.id, + RBACPermission.APP_VIEW_LAYOUT if route == "app" else RBACPermission.AGENT_TEST_AND_RUN, + RBACResourceScope.APP if route == "app" else RBACResourceScope.AGENT, + harness.target.id if route == "app" else harness.agent_id, + ) + ] + assert len(port.calls) == int(allowed) + + @staticmethod + @pytest.mark.parametrize( + ("error", "status", "code"), + [ + (AppDefinitionUnavailableError("stale app"), 400, "app_unavailable"), + (SuggestedQuestionsActorNotFoundError("deleted account"), 401, "unauthorized"), + (MessageNotExistsError(), 404, "not_found"), + (ConversationNotExistsError(), 404, "not_found"), + (ProviderTokenNotInitError(), 400, "provider_not_initialize"), + (QuotaExceededError(), 400, "provider_quota_exceeded"), + (ModelCurrentlyNotSupportError(), 400, "model_currently_not_support"), + (InvokeError("provider failed"), 400, "completion_request_error"), + (SuggestedQuestionsAfterAnswerDisabledError(), 403, "app_suggested_questions_after_answer_disabled"), + (RuntimeError("private credentials"), 500, "internal_server_error"), + ], + ) + @pytest.mark.parametrize("route", ["app", "agent"]) + def test_precise_errors_are_translated_at_http_boundary( + harness: _Harness, error: Exception, status: int, code: str, route: str + ) -> None: + if route == "agent": + harness.add_agent() + harness.services.message_suggested_questions = _Questions(failure=error) + response = harness.get(route) + assert response.status_code == status + assert response.get_json()["code"] == code + assert "private credentials" not in response.get_data(as_text=True) + + @staticmethod + def test_chat_runtime_uses_console_history_after_reference_session_closed(harness: _Harness) -> None: + with harness.flask_app.app_context(): + caller_session = db.session() + caller_app = caller_session.get(App, harness.target.id) + assert caller_app is not None + caller_app.name = "Pending caller edit" + response = harness.get() + assert db.session() is caller_session + assert caller_session.in_transaction() + assert caller_app in caller_session.dirty + assert response.status_code == HTTPStatus.OK + assert response.get_json() == {"data": ["Next?"]} + assert len(harness.provider.prompts) == 1 + assert "Console questions" in harness.provider.prompts[0] + assert "How?" in harness.provider.prompts[0] + assert len(harness.provider.traces) == 1 + + @staticmethod + @pytest.mark.parametrize("entity", [Message, Conversation]) + @pytest.mark.parametrize("change", ["app_id", "from_account_id", "from_source"]) + def test_runtime_cannot_read_another_accounts_message_or_conversation( + harness: _Harness, entity: type[Message] | type[Conversation], change: str + ) -> None: + with harness.factory.begin() as session: + session.execute(update(entity).values({change: "api" if change == "from_source" else str(uuid4())})) + response = harness.get() + assert response.status_code == HTTPStatus.NOT_FOUND + assert response.get_json()["message"] == ( + "Message not found" if entity is Message else "Conversation not found" + ) + assert not harness.provider.prompts + + @staticmethod + def test_agent_runtime_uses_current_accounts_personal_debug_draft(harness: _Harness) -> None: + harness.add_agent() + with harness.factory.begin() as session: + session.execute(update(App).where(App.id == harness.target.id).values({App.app_model_config_id: None})) + session.add( + AgentDebugConversation( + tenant_id=harness.target.tenant_id, + agent_id=harness.agent_id, + app_id=harness.target.id, + account_id=harness.account.id, + draft_type=AgentConfigDraftType.DEBUG_BUILD, + conversation_id=harness.conversation.id, + ) + ) + for owner, prompt in [ + (harness.account.id, "My personal draft"), + (str(uuid4()), "Other accounts secret draft"), + ]: + session.add( + AgentConfigDraft( + tenant_id=harness.target.tenant_id, + agent_id=harness.agent_id, + account_id=owner, + draft_owner_key=owner, + draft_type=AgentConfigDraftType.DEBUG_BUILD, + config_snapshot=AgentSoulConfig.model_validate( + {"app_features": {"suggested_questions_after_answer": {"enabled": True, "prompt": prompt}}} + ), + ) + ) + response = harness.get("agent") + assert response.status_code == HTTPStatus.OK + assert response.get_json() == {"data": ["Next?"]} + assert "My personal draft" in harness.provider.prompts[0] + assert "Other accounts secret draft" not in harness.provider.prompts[0] + + @staticmethod + @pytest.mark.parametrize("same_tenant", [True, False]) + def test_agent_runtime_preserves_conversation_binding_scope(harness: _Harness, same_tenant: bool) -> None: + harness.add_agent() + draft_id, binding_id = str(uuid4()), str(uuid4()) + with harness.factory.begin() as session: + session.execute(update(App).where(App.id == harness.target.id).values({App.app_model_config_id: None})) + session.add( + AgentConfigDraft( + id=draft_id, + tenant_id=harness.target.tenant_id, + agent_id=harness.agent_id, + draft_type=AgentConfigDraftType.DRAFT, + config_snapshot=AgentSoulConfig.model_validate( + { + "app_features": { + "suggested_questions_after_answer": { + "enabled": True, + "prompt": "Conversation bound questions", + } + } + } + ), + ) + ) + session.add( + AgentWorkspaceBinding( + id=binding_id, + tenant_id=harness.target.tenant_id if same_tenant else str(uuid4()), + app_id=harness.target.id, + workspace_id=str(uuid4()), + agent_id=harness.agent_id, + agent_config_version_id=draft_id, + agent_config_version_kind=AgentConfigVersionKind.DRAFT, + backend_binding_ref="existing-binding", + ) + ) + session.execute( + update(Conversation) + .where(Conversation.id == harness.conversation.id) + .values({Conversation.agent_workspace_binding_id: binding_id}) + ) + response = harness.get("agent") + if same_tenant: + assert response.status_code == HTTPStatus.OK + assert "Conversation bound questions" in harness.provider.prompts[0] + else: + assert response.status_code == HTTPStatus.NOT_FOUND + assert response.get_json()["code"] == "agent_version_not_found_error" + assert not harness.provider.prompts diff --git a/api/tests/unit_tests/controllers/console/app/test_message_api.py b/api/tests/unit_tests/controllers/console/app/test_message_api.py deleted file mode 100644 index 26125d6c3c3..00000000000 --- a/api/tests/unit_tests/controllers/console/app/test_message_api.py +++ /dev/null @@ -1,377 +0,0 @@ -from __future__ import annotations - -from datetime import UTC, datetime -from inspect import unwrap -from unittest.mock import MagicMock - -import pytest -from flask import Flask -from sqlalchemy import event -from sqlalchemy.orm import Session - -from controllers.console.app import message as message_module -from core.app.entities.app_invoke_entities import InvokeFrom -from models.account import Account -from models.enums import ConversationFromSource, CreatorUserRole, FeedbackFromSource, FeedbackRating -from models.model import ( - App, - AppAnnotationHitHistory, - AppMode, - Conversation, - Message, - MessageAgentThought, - MessageAnnotation, - MessageFeedback, -) - - -def _account() -> Account: - account = Account(name="Account", email="account@example.com") - account.id = "account-1" - return account - - -def _app(*, app_id: str = "app-1") -> App: - return App( - id=app_id, - tenant_id="tenant-1", - name="Chat App", - description="", - mode=AppMode.CHAT, - enable_site=True, - enable_api=True, - max_active_requests=0, - ) - - -def _persist_message(session: Session, *, message_id: str, app_id: str = "app-1") -> Message: - conversation = Conversation( - app_id=app_id, - app_model_config_id=None, - model_provider=None, - override_model_configs=None, - model_id=None, - mode=AppMode.CHAT, - name="Conversation", - _inputs={}, - introduction="", - system_instruction="", - system_instruction_tokens=0, - status="normal", - invoke_from=InvokeFrom.DEBUGGER, - from_source=ConversationFromSource.CONSOLE, - from_end_user_id=None, - from_account_id="account-1", - ) - conversation.id = "conversation-1" - message = Message( - app_id=app_id, - conversation_id=conversation.id, - _inputs={}, - query="query", - message="", - message_tokens=0, - message_unit_price=0, - message_price_unit=0, - answer="answer", - answer_tokens=0, - answer_unit_price=0, - answer_price_unit=0, - provider_response_latency=0, - total_price=0, - currency="USD", - invoke_from=InvokeFrom.DEBUGGER, - from_source=ConversationFromSource.CONSOLE, - from_end_user_id=None, - from_account_id="account-1", - app_mode=AppMode.CHAT, - ) - message.id = message_id - session.add_all([_account(), _app(app_id=app_id), conversation, message]) - session.flush() - return message - - -def test_app_message_routes_pass_injected_session( - app: Flask, monkeypatch: pytest.MonkeyPatch, unbound_session: Session -) -> None: - session = unbound_session - current_user = _account() - app_model = _app() - message_id = "550e8400-e29b-41d4-a716-446655440000" - list_messages = MagicMock(return_value={"data": []}) - update_feedback = MagicMock(return_value={"result": "success"}) - get_message_detail = MagicMock(return_value={"id": message_id}) - monkeypatch.setattr(message_module, "_list_chat_messages", list_messages) - monkeypatch.setattr(message_module, "_update_message_feedback", update_feedback) - monkeypatch.setattr(message_module, "_get_message_detail", get_message_detail) - - assert unwrap(message_module.ChatMessageListApi.get)( - message_module.ChatMessageListApi(), - message_module.ChatMessagesQuery(conversation_id="550e8400-e29b-41d4-a716-446655440001"), - session, - current_user, - app_model, - ) == {"data": []} - assert unwrap(message_module.MessageFeedbackApi.post)( - message_module.MessageFeedbackApi(), - message_module.AppMessageFeedbackPayload(message_id=message_id, rating="like"), - session, - current_user, - app_model, - ) == {"result": "success"} - assert unwrap(message_module.MessageApi.get)(message_module.MessageApi(), session, app_model, message_id) == { - "id": message_id - } - - assert list_messages.call_args.kwargs["session"] is session - assert update_feedback.call_args.kwargs["session"] is session - assert get_message_detail.call_args.kwargs["session"] is session - - -def test_update_message_feedback_commits_injected_session(app: Flask, sqlite_session: Session) -> None: - message_id = "550e8400-e29b-41d4-a716-446655440000" - message = _persist_message(sqlite_session, message_id=message_id) - feedback = MessageFeedback( - app_id=message.app_id, - conversation_id=message.conversation_id, - message_id=message.id, - rating=FeedbackRating.DISLIKE, - from_source=FeedbackFromSource.ADMIN, - from_account_id="account-1", - ) - sqlite_session.add(feedback) - sqlite_session.commit() - session = sqlite_session - commits: list[str] = [] - event.listen(session, "after_commit", lambda _session: commits.append("commit")) - - with app.test_request_context(json={"message_id": message_id, "rating": "like", "content": "helpful"}): - result = message_module._update_message_feedback( - args=message_module.AppMessageFeedbackPayload( - message_id=message_id, - rating="like", - content="helpful", - ), - session=session, - current_user=_account(), - app_model=_app(), - ) - - assert result == {"result": "success"} - updated_feedback = sqlite_session.get(MessageFeedback, feedback.id) - assert updated_feedback is not None - assert updated_feedback.rating == FeedbackRating.LIKE - assert updated_feedback.content == "helpful" - assert commits == ["commit"] - - -def test_get_message_detail_uses_injected_session(monkeypatch: pytest.MonkeyPatch, sqlite_session: Session) -> None: - message_id = "550e8400-e29b-41d4-a716-446655440000" - message = _persist_message(sqlite_session, message_id=message_id) - response_source = object() - response_source_factory = MagicMock(return_value=response_source) - session = sqlite_session - monkeypatch.setattr(message_module, "attach_message_extra_contents", MagicMock()) - monkeypatch.setattr(message_module, "MessageResponseSource", response_source_factory) - monkeypatch.setattr(message_module, "dump_response", lambda _model, value: value) - - result = message_module._get_message_detail( - session=session, - app_model=_app(), - message_id=message_id, - ) - - assert result is response_source - response_source_factory.assert_called_once_with(message, session=session) - - -def test_message_response_source_uses_caller_session_for_nested_fields(sqlite_session: Session) -> None: - session = sqlite_session - message_id = "550e8400-e29b-41d4-a716-446655440000" - message = _persist_message(sqlite_session, message_id=message_id) - message.inputs = {"topic": "support"} - feedback = MessageFeedback( - app_id=message.app_id, - conversation_id=message.conversation_id, - message_id=message.id, - rating=FeedbackRating.LIKE, - from_source=FeedbackFromSource.USER, - from_account_id="account-1", - ) - annotation = MessageAnnotation( - app_id=message.app_id, - question="Question", - content="Answer", - account_id="account-1", - conversation_id=message.conversation_id, - message_id=message.id, - ) - annotation.id = "annotation-1" - annotation_hit = AppAnnotationHitHistory( - app_id=message.app_id, - annotation_id=annotation.id, - source="annotation", - question="Question", - account_id="account-1", - score=0.0, - message_id=message.id, - annotation_question="Question", - annotation_content="Answer", - ) - thought = MessageAgentThought( - message_id=message.id, - position=1, - created_by_role=CreatorUserRole.ACCOUNT, - created_by="account-1", - thought="Consider the request", - tool_labels_str="{}", - tool_meta_str="{}", - ) - sqlite_session.add_all([feedback, annotation, annotation_hit, thought]) - sqlite_session.commit() - - source = message_module.MessageResponseSource(message, session=session) - - assert source.inputs == {"topic": "support"} - assert source.user_feedback is feedback - assert source.feedbacks[0].from_account.id == "account-1" - annotation_source = source.annotation - assert annotation_source is not None - assert annotation_source.account.id == "account-1" - annotation_hit_history_source = source.annotation_hit_history - assert annotation_hit_history_source is not None - assert annotation_hit_history_source.annotation_create_account.id == "account-1" - assert source.agent_thoughts == [thought] - assert source.message_files == [] - - -def test_chat_messages_query_valid(app: Flask, monkeypatch: pytest.MonkeyPatch) -> None: - """Test valid ChatMessagesQuery with all fields.""" - query = message_module.ChatMessagesQuery( - conversation_id="550e8400-e29b-41d4-a716-446655440000", - first_id="550e8400-e29b-41d4-a716-446655440001", - limit=50, - ) - assert query.limit == 50 - - -def test_chat_messages_query_defaults(app: Flask, monkeypatch: pytest.MonkeyPatch) -> None: - """Test ChatMessagesQuery with defaults.""" - query = message_module.ChatMessagesQuery(conversation_id="550e8400-e29b-41d4-a716-446655440000") - assert query.first_id is None - assert query.limit == 20 - - -def test_chat_messages_query_empty_first_id(app: Flask, monkeypatch: pytest.MonkeyPatch) -> None: - """Test ChatMessagesQuery converts empty first_id to None.""" - query = message_module.ChatMessagesQuery( - conversation_id="550e8400-e29b-41d4-a716-446655440000", - first_id="", - ) - assert query.first_id is None - - -def test_message_feedback_payload_valid_like(app: Flask, monkeypatch: pytest.MonkeyPatch) -> None: - """Test AppMessageFeedbackPayload with like rating.""" - payload = message_module.AppMessageFeedbackPayload( - message_id="550e8400-e29b-41d4-a716-446655440000", - rating="like", - content="Good answer", - ) - assert payload.rating == "like" - assert payload.content == "Good answer" - - -def test_message_feedback_payload_valid_dislike(app: Flask, monkeypatch: pytest.MonkeyPatch) -> None: - """Test AppMessageFeedbackPayload with dislike rating.""" - payload = message_module.AppMessageFeedbackPayload( - message_id="550e8400-e29b-41d4-a716-446655440000", - rating="dislike", - ) - assert payload.rating == "dislike" - - -def test_message_feedback_payload_no_rating(app: Flask, monkeypatch: pytest.MonkeyPatch) -> None: - """Test AppMessageFeedbackPayload without rating.""" - payload = message_module.AppMessageFeedbackPayload(message_id="550e8400-e29b-41d4-a716-446655440000") - assert payload.rating is None - - -def test_annotation_count_response(app: Flask, monkeypatch: pytest.MonkeyPatch) -> None: - """Test AnnotationCountResponse creation.""" - response = message_module.AnnotationCountResponse(count=10) - assert response.count == 10 - - -def test_suggested_questions_response(app: Flask, monkeypatch: pytest.MonkeyPatch) -> None: - """Test SuggestedQuestionsResponse creation.""" - response = message_module.SuggestedQuestionsResponse(data=["What is AI?", "How does ML work?"]) - assert len(response.data) == 2 - assert response.data[0] == "What is AI?" - - -def test_message_detail_response_normalizes_aliases_and_timestamp(app: Flask, monkeypatch: pytest.MonkeyPatch) -> None: - """Test MessageDetailResponse normalizes alias fields and datetime timestamps.""" - created_at = datetime(2026, 1, 2, 3, 4, 5, tzinfo=UTC) - response = message_module.MessageDetailResponse.model_validate( - { - "id": "550e8400-e29b-41d4-a716-446655440000", - "conversation_id": "550e8400-e29b-41d4-a716-446655440001", - "inputs": {"foo": "bar"}, - "query": "hello", - "message": [{"text": "hello"}], - "message_tokens": 7, - "answer": "world", - "answer_tokens": 11, - "provider_response_latency": 1.25, - "from_source": "user", - "from_end_user_id": None, - "from_account_id": "550e8400-e29b-41d4-a716-446655440002", - "feedbacks": [], - "workflow_run_id": None, - "annotation": None, - "annotation_hit_history": None, - "status": "normal", - "created_at": created_at, - "agent_thoughts": [], - "message_files": [], - "message_metadata_dict": {"token_usage": 3}, - "error": None, - "parent_message_id": None, - "extra_contents": [], - } - ) - assert response.answer == "world" - assert response.message_tokens == 7 - assert response.answer_tokens == 11 - assert response.provider_response_latency == 1.25 - assert response.metadata == {"token_usage": 3} - assert response.created_at == int(created_at.timestamp()) - assert response.model_dump(mode="json") == { - "id": "550e8400-e29b-41d4-a716-446655440000", - "conversation_id": "550e8400-e29b-41d4-a716-446655440001", - "inputs": {"foo": "bar"}, - "query": "hello", - "message": [{"text": "hello"}], - "message_tokens": 7, - "answer": "world", - "answer_tokens": 11, - "provider_response_latency": 1.25, - "from_source": "user", - "from_end_user_id": None, - "from_account_id": "550e8400-e29b-41d4-a716-446655440002", - "feedbacks": [], - "workflow_run_id": None, - "annotation": None, - "annotation_hit_history": None, - "created_at": int(created_at.timestamp()), - "agent_thoughts": [], - "message_files": [], - "metadata": {"token_usage": 3}, - "retriever_resources": [], - "status": "normal", - "error": None, - "parent_message_id": None, - "extra_contents": [], - } diff --git a/api/tests/unit_tests/controllers/console/app/test_message_suggested_questions.py b/api/tests/unit_tests/controllers/console/app/test_message_suggested_questions.py deleted file mode 100644 index 44014ae4964..00000000000 --- a/api/tests/unit_tests/controllers/console/app/test_message_suggested_questions.py +++ /dev/null @@ -1,550 +0,0 @@ -"""Console suggested-question admission and actor/configuration isolation over SQLite.""" - -import json -from collections.abc import Iterator, Sequence -from dataclasses import dataclass, field -from decimal import Decimal -from http import HTTPStatus -from typing import cast -from uuid import uuid4 - -import pytest -from flask import Flask -from sqlalchemy import Connection, Engine, event, func, select, text, update -from sqlalchemy.orm import Session, SessionTransaction, sessionmaker -from werkzeug.test import TestResponse - -import controllers.console.wraps as console_wraps -from controllers.common.rbac import RBACPermission, RBACResourceScope -from controllers.common.rbac.checks import RBACService -from controllers.console.app import message as controller -from core.errors.error import ModelCurrentlyNotSupportError, ProviderTokenNotInitError, QuotaExceededError -from core.model_manager import ModelInstance -from core.ops.ops_trace_manager import TraceTask -from extensions.application_services.app import AppServices -from extensions.application_services.rbac import RBACServices -from extensions.ext_database import db -from graphon.model_runtime.entities.llm_entities import LLMResult, LLMUsage -from graphon.model_runtime.entities.message_entities import AssistantPromptMessage, PromptMessage -from graphon.model_runtime.entities.model_entities import ModelType -from graphon.model_runtime.errors.invoke import InvokeError -from libs import login -from libs.external_api import ExternalApi -from models.account import Account, Tenant -from models.agent import ( - Agent, - AgentConfigDraft, - AgentConfigDraftType, - AgentConfigVersionKind, - AgentDebugConversation, - AgentScope, - AgentSource, - AgentStatus, - AgentWorkspaceBinding, -) -from models.agent_config_entities import AgentSoulConfig -from models.enums import ConversationFromSource -from models.model import App, AppMode, AppModelConfig, Conversation, Message -from repositories.app.agent_app_repository import AgentAppRepository -from services import message_service -from services.app.agent_app_service import AgentAppAccessService -from services.app_definition_query_service import AppDefinitionUnavailableError -from services.errors.conversation import ConversationNotExistsError -from services.errors.message import MessageNotExistsError, SuggestedQuestionsAfterAnswerDisabledError -from services.message_suggested_questions_adapters import MessageSuggestedQuestionsRuntime -from services.message_suggested_questions_service import ( - MessageSuggestedQuestions, - SuggestedQuestionsAccount, - SuggestedQuestionsActor, - SuggestedQuestionsActorNotFoundError, -) -from tests.unit_tests.config_override import apply_config_overrides -from tests.unit_tests.model_factories import make_account, make_app, make_conversation, make_message -from tests.unit_tests.rbac_fakes import RBACDomain - - -@dataclass -class _Questions: - questions: list[str] = field(default_factory=lambda: ["Next?"]) - failure: Exception | None = None - calls: list[tuple[str, str, str, SuggestedQuestionsActor, str]] = field(default_factory=list) - - def get_suggested_questions( - self, - *, - app_id: str, - app_owner_tenant_id: str, - expected_app_mode: str, - actor: SuggestedQuestionsActor, - message_id: str, - ) -> list[str]: - self.calls.append((app_id, app_owner_tenant_id, expected_app_mode, actor, message_id)) - if self.failure is not None: - raise self.failure - return self.questions - - -@dataclass -class _Agents: - access: AgentAppAccessService - - -@dataclass -class _Services: - apps: AppServices - agent_apps: _Agents - message_suggested_questions: MessageSuggestedQuestions - rbac: RBACServices - - -@dataclass -class _Provider: - tenant_id: str - admission_sessions: list[Session] - prompts: list[str] = field(default_factory=list) - traces: list[TraceTask] = field(default_factory=list) - - def get_default_model_instance(self, *, tenant_id: str, model_type: ModelType) -> ModelInstance: - assert tenant_id == self.tenant_id - assert model_type == ModelType.LLM - assert self.admission_sessions - assert all(not session.in_transaction() and not session.identity_map for session in self.admission_sessions) - return cast(ModelInstance, self) - - def get_llm_num_tokens(self, prompt_messages: Sequence[PromptMessage]) -> int: - return len(prompt_messages) - - def get_model_schema(self) -> None: - raise NotImplementedError - - def invoke_llm( - self, - *, - prompt_messages: list[PromptMessage], - model_parameters: dict[str, object], - stop: list[str], - stream: bool, - ) -> LLMResult: - assert model_parameters == {"max_tokens": 256, "temperature": 0.0} - assert stop == [] - assert stream is False - self.prompts.append(prompt_messages[0].get_text_content()) - return LLMResult( - model="test-model", message=AssistantPromptMessage(content='["Next?"]'), usage=LLMUsage.empty_usage() - ) - - def add_trace_task(self, task: TraceTask) -> None: - self.traces.append(task) - - -@dataclass -class _Harness: - flask_app: Flask - target: App - account: Account - conversation: Conversation - message: Message - agent_id: str - factory: sessionmaker[Session] - services: _Services - provider: _Provider - permission_calls: list[tuple[str, str, RBACPermission, RBACResourceScope | None, str | None]] - allowed: bool = True - - def get(self, route: str = "app") -> TestResponse: - resource = f"apps/{self.target.id}" if route == "app" else f"agent/{self.agent_id}" - return self.flask_app.test_client().get(f"/{resource}/chat-messages/{self.message.id}/suggested-questions") - - def add_agent(self, *, scope: AgentScope = AgentScope.ROSTER, backing: bool = True) -> None: - with self.factory.begin() as session: - session.execute(update(App).where(App.id == self.target.id).values(mode=AppMode.AGENT)) - session.add( - Agent( - id=self.agent_id, - tenant_id=self.target.tenant_id, - name="Agent", - scope=scope, - source=AgentSource.AGENT_APP if scope == AgentScope.ROSTER else AgentSource.WORKFLOW, - status=AgentStatus.ACTIVE, - app_id=self.target.id if scope == AgentScope.ROSTER else str(uuid4()), - backing_app_id=self.target.id if scope == AgentScope.WORKFLOW_ONLY and backing else None, - workflow_id=str(uuid4()) if scope == AgentScope.WORKFLOW_ONLY else None, - workflow_node_id="node" if scope == AgentScope.WORKFLOW_ONLY else None, - ) - ) - - -@pytest.fixture -def harness( - monkeypatch: pytest.MonkeyPatch, - sqlite_engine: Engine, - sqlite_session_factory: sessionmaker[Session], - app_services: AppServices, - rbac_domain: RBACDomain, -) -> Iterator[_Harness]: - target = make_app(app_id=str(uuid4()), tenant_id=str(uuid4())) - account = make_account(account_id=str(uuid4())) - tenant = Tenant(name="Workspace") - tenant.id = target.tenant_id - account._current_tenant = tenant - conversation = make_conversation( - conversation_id=str(uuid4()), - app_id=target.id, - from_source=ConversationFromSource.CONSOLE, - from_account_id=account.id, - inputs={}, - ) - config = AppModelConfig( - app_id=target.id, suggested_questions_after_answer=json.dumps({"enabled": True, "prompt": "Console questions"}) - ) - message = make_message( - message_id=str(uuid4()), - app_id=target.id, - conversation_id=conversation.id, - inputs={}, - message={}, - query="How?", - answer="Like this.", - message_unit_price=Decimal(0), - answer_unit_price=Decimal(0), - currency="USD", - from_source=ConversationFromSource.CONSOLE, - from_account_id=account.id, - ) - with sqlite_session_factory.begin() as session: - session.add_all([target, account, conversation, config, message]) - session.flush() - target.app_model_config_id = config.id - conversation.app_model_config_id = config.id - sessions: list[Session] = [] - - def track_session(session: Session, _transaction: SessionTransaction, _connection: Connection) -> None: - sessions.append(session) - - event.listen(sqlite_session_factory, "after_begin", track_session) - provider = _Provider(target.tenant_id, sessions) - - def manager(*, tenant_id: str) -> _Provider: - assert tenant_id == target.tenant_id - return provider - - def trace_manager(*, app_id: str) -> _Provider: - assert app_id == target.id - return provider - - monkeypatch.setattr(message_service.ModelManager, "for_tenant", manager) - monkeypatch.setattr(message_service, "TraceQueueManager", trace_manager) - monkeypatch.setattr(login, "current_user", account) - monkeypatch.setattr(console_wraps, "_is_setup_completed", lambda: True) - apply_config_overrides(monkeypatch, LOGIN_DISABLED=True, RBAC_ENABLED=False) - flask_app = Flask(__name__) - flask_app.config.update(TESTING=True, RESTX_ERROR_404_HELP=False, SQLALCHEMY_DATABASE_URI=str(sqlite_engine.url)) - db.init_app(flask_app) - services = _Services( - app_services, - _Agents(AgentAppAccessService(references=AgentAppRepository(session_factory=sqlite_session_factory))), - MessageSuggestedQuestionsRuntime(session_factory=sqlite_session_factory), - rbac_domain.rbac, - ) - flask_app.extensions["application_services"] = services - api = ExternalApi(flask_app) - api.add_resource( - controller.MessageSuggestedQuestionApi, - "/apps//chat-messages//suggested-questions", - ) - api.add_resource( - controller.AgentMessageSuggestedQuestionApi, - "/agent//chat-messages//suggested-questions", - ) - result = _Harness( - flask_app, target, account, conversation, message, str(uuid4()), sqlite_session_factory, services, provider, [] - ) - - def permission( - tenant_id: str, - account_id: str, - *, - scene: RBACPermission, - resource_type: RBACResourceScope | None, - resource_id: str | None, - ) -> bool: - result.permission_calls.append((tenant_id, account_id, scene, resource_type, resource_id)) - return result.allowed - - monkeypatch.setattr(RBACService.CheckAccess, "check", permission) - yield result - event.remove(sqlite_session_factory, "after_begin", track_session) - with flask_app.app_context(): - db.session.remove() - db.engine.dispose() - - -@pytest.mark.parametrize("mode", [AppMode.CHAT, AppMode.AGENT_CHAT, AppMode.ADVANCED_CHAT, AppMode.AGENT]) -@pytest.mark.parametrize("questions", [[], ["Next?", "More?"]]) -def test_app_reference_and_account_are_passed_to_shared_runtime( - harness: _Harness, mode: AppMode, questions: list[str] -) -> None: - with harness.factory.begin() as session: - session.execute(update(App).where(App.id == harness.target.id).values(mode=mode)) - port = _Questions(questions=questions) - harness.services.message_suggested_questions = port - response = harness.get() - assert response.status_code == HTTPStatus.OK - assert response.headers["Content-Type"] == "application/json" - assert response.get_json() == {"data": questions} - assert port.calls == [ - ( - harness.target.id, - harness.target.tenant_id, - mode, - SuggestedQuestionsAccount(harness.account.id, "debugger"), - harness.message.id, - ) - ] - assert all( - not session.in_transaction() and not session.identity_map for session in harness.provider.admission_sessions - ) - - -@pytest.mark.parametrize("change", ["tenant", "status", "missing", "workflow", "completion", "hidden"]) -def test_app_admission_rejects_unavailable_or_unsupported_apps(harness: _Harness, change: str) -> None: - if change == "hidden": - harness.add_agent(scope=AgentScope.WORKFLOW_ONLY) - else: - field_name, value = { - "tenant": ("tenant_id", str(uuid4())), - "status": ("status", "archived"), - "missing": ("id", str(uuid4())), - "workflow": ("mode", "workflow"), - "completion": ("mode", "completion"), - }[change] - with harness.factory.begin() as session: - if change == "status": - session.execute( - text("UPDATE apps SET status = :status WHERE id = :id"), {"status": value, "id": harness.target.id} - ) - else: - session.execute(update(App).where(App.id == harness.target.id).values({field_name: value})) - port = _Questions() - harness.services.message_suggested_questions = port - response = harness.get() - assert response.status_code == HTTPStatus.NOT_FOUND - assert response.get_json()["code"] == "app_not_found" - assert not port.calls - - -@pytest.mark.parametrize("scope", [AgentScope.ROSTER, AgentScope.WORKFLOW_ONLY]) -def test_agent_resolves_existing_backing_app(harness: _Harness, scope: AgentScope) -> None: - harness.add_agent(scope=scope) - port = _Questions() - harness.services.message_suggested_questions = port - response = harness.get("agent") - assert response.status_code == HTTPStatus.OK - assert port.calls == [ - ( - harness.target.id, - harness.target.tenant_id, - AppMode.AGENT, - SuggestedQuestionsAccount(harness.account.id, "debugger"), - harness.message.id, - ) - ] - assert all( - not session.in_transaction() and not session.identity_map for session in harness.provider.admission_sessions - ) - - -@pytest.mark.parametrize("change", ["missing_backing", "tenant", "agent_status", "app_status", "app_mode"]) -def test_agent_rejects_unavailable_runtime_without_creating_app(harness: _Harness, change: str) -> None: - harness.add_agent(scope=AgentScope.WORKFLOW_ONLY, backing=change != "missing_backing") - if change != "missing_backing": - model, key, value = { - "tenant": (Agent, "tenant_id", str(uuid4())), - "agent_status": (Agent, "status", AgentStatus.ARCHIVED), - "app_status": (App, "status", "archived"), - "app_mode": (App, "mode", AppMode.CHAT), - }[change] - with harness.factory.begin() as session: - if change == "app_status": - session.execute( - text("UPDATE apps SET status = :status WHERE id = :id"), {"status": value, "id": harness.target.id} - ) - else: - session.execute(update(model).values({key: value})) - response = harness.get("agent") - assert response.status_code == HTTPStatus.NOT_FOUND - assert response.get_json()["code"] == "agent_not_found_error" - with harness.factory() as session: - assert session.scalar(select(func.count()).select_from(App)) == 1 - assert session.scalar(select(func.count()).select_from(AppModelConfig)) == 1 - if change == "missing_backing": - assert session.scalar(select(Agent.backing_app_id)) is None - - -@pytest.mark.parametrize("route", ["app", "agent"]) -@pytest.mark.parametrize("allowed", [True, False]) -def test_real_rbac_admission_keeps_resource_scene( - harness: _Harness, monkeypatch: pytest.MonkeyPatch, route: str, allowed: bool -) -> None: - apply_config_overrides(monkeypatch, RBAC_ENABLED=True) - if route == "agent": - harness.add_agent() - harness.allowed = allowed - port = _Questions() - harness.services.message_suggested_questions = port - response = harness.get(route) - assert response.status_code == (HTTPStatus.OK if allowed else HTTPStatus.FORBIDDEN) - assert harness.permission_calls == [ - ( - harness.target.tenant_id, - harness.account.id, - RBACPermission.APP_VIEW_LAYOUT if route == "app" else RBACPermission.AGENT_TEST_AND_RUN, - RBACResourceScope.APP if route == "app" else RBACResourceScope.AGENT, - harness.target.id if route == "app" else harness.agent_id, - ) - ] - assert len(port.calls) == int(allowed) - - -@pytest.mark.parametrize( - ("error", "status", "code"), - [ - (AppDefinitionUnavailableError("stale app"), 400, "app_unavailable"), - (SuggestedQuestionsActorNotFoundError("deleted account"), 401, "unauthorized"), - (MessageNotExistsError(), 404, "not_found"), - (ConversationNotExistsError(), 404, "not_found"), - (ProviderTokenNotInitError(), 400, "provider_not_initialize"), - (QuotaExceededError(), 400, "provider_quota_exceeded"), - (ModelCurrentlyNotSupportError(), 400, "model_currently_not_support"), - (InvokeError("provider failed"), 400, "completion_request_error"), - (SuggestedQuestionsAfterAnswerDisabledError(), 403, "app_suggested_questions_after_answer_disabled"), - (RuntimeError("private credentials"), 500, "internal_server_error"), - ], -) -@pytest.mark.parametrize("route", ["app", "agent"]) -def test_precise_errors_are_translated_at_http_boundary( - harness: _Harness, error: Exception, status: int, code: str, route: str -) -> None: - if route == "agent": - harness.add_agent() - harness.services.message_suggested_questions = _Questions(failure=error) - response = harness.get(route) - assert response.status_code == status - assert response.get_json()["code"] == code - assert "private credentials" not in response.get_data(as_text=True) - - -def test_chat_runtime_uses_console_history_after_reference_session_closed(harness: _Harness) -> None: - with harness.flask_app.app_context(): - caller_session = db.session() - caller_app = caller_session.get(App, harness.target.id) - assert caller_app is not None - caller_app.name = "Pending caller edit" - response = harness.get() - assert db.session() is caller_session - assert caller_session.in_transaction() - assert caller_app in caller_session.dirty - assert response.status_code == HTTPStatus.OK - assert response.get_json() == {"data": ["Next?"]} - assert len(harness.provider.prompts) == 1 - assert "Console questions" in harness.provider.prompts[0] - assert "How?" in harness.provider.prompts[0] - assert len(harness.provider.traces) == 1 - - -@pytest.mark.parametrize("entity", [Message, Conversation]) -@pytest.mark.parametrize("change", ["app_id", "from_account_id", "from_source"]) -def test_runtime_cannot_read_another_accounts_message_or_conversation( - harness: _Harness, entity: type[Message] | type[Conversation], change: str -) -> None: - with harness.factory.begin() as session: - session.execute(update(entity).values({change: "api" if change == "from_source" else str(uuid4())})) - response = harness.get() - assert response.status_code == HTTPStatus.NOT_FOUND - assert response.get_json()["message"] == ("Message not found" if entity is Message else "Conversation not found") - assert not harness.provider.prompts - - -def test_agent_runtime_uses_current_accounts_personal_debug_draft(harness: _Harness) -> None: - harness.add_agent() - with harness.factory.begin() as session: - session.execute(update(App).where(App.id == harness.target.id).values({App.app_model_config_id: None})) - session.add( - AgentDebugConversation( - tenant_id=harness.target.tenant_id, - agent_id=harness.agent_id, - app_id=harness.target.id, - account_id=harness.account.id, - draft_type=AgentConfigDraftType.DEBUG_BUILD, - conversation_id=harness.conversation.id, - ) - ) - for owner, prompt in [(harness.account.id, "My personal draft"), (str(uuid4()), "Other accounts secret draft")]: - session.add( - AgentConfigDraft( - tenant_id=harness.target.tenant_id, - agent_id=harness.agent_id, - account_id=owner, - draft_owner_key=owner, - draft_type=AgentConfigDraftType.DEBUG_BUILD, - config_snapshot=AgentSoulConfig.model_validate( - {"app_features": {"suggested_questions_after_answer": {"enabled": True, "prompt": prompt}}} - ), - ) - ) - response = harness.get("agent") - assert response.status_code == HTTPStatus.OK - assert response.get_json() == {"data": ["Next?"]} - assert "My personal draft" in harness.provider.prompts[0] - assert "Other accounts secret draft" not in harness.provider.prompts[0] - - -@pytest.mark.parametrize("same_tenant", [True, False]) -def test_agent_runtime_preserves_conversation_binding_scope(harness: _Harness, same_tenant: bool) -> None: - harness.add_agent() - draft_id, binding_id = str(uuid4()), str(uuid4()) - with harness.factory.begin() as session: - session.execute(update(App).where(App.id == harness.target.id).values({App.app_model_config_id: None})) - session.add( - AgentConfigDraft( - id=draft_id, - tenant_id=harness.target.tenant_id, - agent_id=harness.agent_id, - draft_type=AgentConfigDraftType.DRAFT, - config_snapshot=AgentSoulConfig.model_validate( - { - "app_features": { - "suggested_questions_after_answer": { - "enabled": True, - "prompt": "Conversation bound questions", - } - } - } - ), - ) - ) - session.add( - AgentWorkspaceBinding( - id=binding_id, - tenant_id=harness.target.tenant_id if same_tenant else str(uuid4()), - app_id=harness.target.id, - workspace_id=str(uuid4()), - agent_id=harness.agent_id, - agent_config_version_id=draft_id, - agent_config_version_kind=AgentConfigVersionKind.DRAFT, - backend_binding_ref="existing-binding", - ) - ) - session.execute( - update(Conversation) - .where(Conversation.id == harness.conversation.id) - .values({Conversation.agent_workspace_binding_id: binding_id}) - ) - response = harness.get("agent") - if same_tenant: - assert response.status_code == HTTPStatus.OK - assert "Conversation bound questions" in harness.provider.prompts[0] - else: - assert response.status_code == HTTPStatus.NOT_FOUND - assert response.get_json()["code"] == "agent_version_not_found_error" - assert not harness.provider.prompts diff --git a/api/tests/unit_tests/controllers/console/app/test_model_config_api.py b/api/tests/unit_tests/controllers/console/app/test_model_config.py similarity index 100% rename from api/tests/unit_tests/controllers/console/app/test_model_config_api.py rename to api/tests/unit_tests/controllers/console/app/test_model_config.py diff --git a/api/tests/unit_tests/controllers/console/app/test_ops_trace_api.py b/api/tests/unit_tests/controllers/console/app/test_ops_trace.py similarity index 99% rename from api/tests/unit_tests/controllers/console/app/test_ops_trace_api.py rename to api/tests/unit_tests/controllers/console/app/test_ops_trace.py index dc28f6eee73..9d0e9847a14 100644 --- a/api/tests/unit_tests/controllers/console/app/test_ops_trace_api.py +++ b/api/tests/unit_tests/controllers/console/app/test_ops_trace.py @@ -44,7 +44,7 @@ from services.app_tracing_config_service import ( AppTracingConfigService, AppTracingConfigVerificationFailedError, ) -from tests.unit_tests.config_override import apply_config_overrides +from tests.helpers.unit.config_override import apply_config_overrides APP_ID = "11111111-1111-1111-1111-111111111111" WORKSPACE_ID = "22222222-2222-2222-2222-222222222222" diff --git a/api/tests/unit_tests/controllers/console/app/test_roster_package_routes.py b/api/tests/unit_tests/controllers/console/app/test_roster_package_routes.py deleted file mode 100644 index a400023d9a3..00000000000 --- a/api/tests/unit_tests/controllers/console/app/test_roster_package_routes.py +++ /dev/null @@ -1,520 +0,0 @@ -import io -import zipfile -from collections.abc import Callable -from dataclasses import dataclass, field -from inspect import unwrap -from typing import BinaryIO, cast -from unittest.mock import Mock -from uuid import UUID, uuid4 - -import httpx -import pytest -from flask import Flask -from pydantic import ValidationError -from sqlalchemy.orm import Session, sessionmaker -from werkzeug.exceptions import Forbidden - -from controllers.console.app import app as app_module -from controllers.console.app import app_import as import_module -from core.rbac import RBACPermission -from machinery.context import RequestContext -from models.account import Account, Tenant, TenantAccountJoin, TenantAccountRole -from services import app_import_source -from services.agent.errors import InvalidRosterAgentPackageError -from services.agent.roster_package_importer import RosterAgentPackageImportResult -from services.app.console_gateway import AppTransferGateway -from services.app_package_service import AppPackageService, PreparedAppPackage -from services.entities.dsl_entities import AppImportPackage, AppImportParams, DslImportWarning, Import, ImportStatus -from services.entities.feature_entities import FeatureModel, LimitationModel -from services.recommended_app_package_service import RecommendedAgentPackageSource, RecommendedAppPackageService -from tests.unit_tests.controllers.conftest import ControllerTestServices - - -def _roster_archive() -> io.BytesIO: - source = io.BytesIO() - with zipfile.ZipFile(source, "w") as archive: - # Fixed timestamp: passing a filename string stamps the current local - # time (2-second ZIP resolution), so two archives built across a - # timestamp boundary differ byte-for-byte and flake the assertion below. - info = zipfile.ZipInfo("manifest.yaml", date_time=(1980, 1, 1, 0, 0, 0)) - archive.writestr(info, "format: dify.roster-agent") - source.seek(0) - return source - - -def _mock_download(monkeypatch: pytest.MonkeyPatch, content: bytes | None = None) -> list[str]: - calls: list[str] = [] - - def fetch(method: str, url: str, **_kwargs: object) -> httpx.Response: - calls.append(url) - return httpx.Response( - 200, - content=content if content is not None else _roster_archive().getvalue(), - request=httpx.Request(method, url), - ) - - monkeypatch.setattr(app_import_source.remote_fetcher, "make_request", fetch) - return calls - - -@dataclass(frozen=True) -class AgentImportCall: - tenant_id: str - account: Account - content: bytes - - -@dataclass -class Imports: - context: RequestContext - dsl_calls: list[AppImportParams] = field(default_factory=list) - packages: list[PreparedAppPackage] = field(default_factory=list) - agent_calls: list[AgentImportCall] = field(default_factory=list) - status: ImportStatus = ImportStatus.COMPLETED - warnings: list[DslImportWarning] = field(default_factory=list) - - def dsl( - self, context: RequestContext, params: AppImportParams, *, package: AppImportPackage | None = None - ) -> Import: - assert context == self.context - self.dsl_calls.append(params) - if package is not None: - assert isinstance(package, PreparedAppPackage) - assert not package.archive.closed - self.packages.append(package) - return Import(id="import-1", status=self.status, app_id="app-1") - - def agent(self, *, source: BinaryIO, tenant_id: str, account: Account) -> RosterAgentPackageImportResult: - self.agent_calls.append(AgentImportCall(tenant_id, account, source.read())) - return RosterAgentPackageImportResult(app_id="app-1", agent_id="agent-1", warnings=self.warnings) - - -@pytest.fixture -def imports( - app_query_services: ControllerTestServices, - sqlite_session_factory: sessionmaker[Session], - monkeypatch: pytest.MonkeyPatch, - config_overrides: Callable[..., None], -) -> Imports: - config_overrides(RBAC_ENABLED=False, DEPLOYMENT_EDITION="COMMUNITY") - actor, tenant = str(uuid4()), str(uuid4()) - account = Account(name="Importer", email="importer@example.com") - account.id = actor - workspace = Tenant(name="Workspace") - workspace.id = tenant - with sqlite_session_factory.begin() as session: - session.add_all( - [ - account, - workspace, - TenantAccountJoin(tenant_id=tenant, account_id=actor, role=TenantAccountRole.OWNER), - ] - ) - result = Imports(RequestContext("request", None, actor, tenant)) - transfers = app_query_services.apps.console._transfers - assert isinstance(transfers, AppTransferGateway) - monkeypatch.setattr(transfers, "import_dsl", result.dsl) - monkeypatch.setattr(transfers._agent_importer, "import_package", result.agent) - monkeypatch.setattr("services.app.console_gateway.SystemFeatureService.is_webapp_auth_enabled", lambda: False) - return result - - -def _post( - app: Flask, - imports: Imports, - *, - json: dict[str, object] | None = None, - data: dict[str, object] | None = None, -) -> tuple[dict[str, object], int]: - api = import_module.AppImportApi() - with app.test_request_context("/console/api/apps/imports", method="POST", json=json, data=data): - return cast(tuple[dict[str, object], int], unwrap(api.post)(api, imports.context)) - - -def test_existing_import_route_dispatches_json_without_changing_payload(app: Flask, imports: Imports) -> None: - data, status = _post(app, imports, json={"mode": "yaml-content", "yaml_content": "app: {}"}) - assert status == 200 - assert data["status"] == "completed" - assert [params.model_dump() for params in imports.dsl_calls] == [ - AppImportParams(mode="yaml-content", yaml_content="app: {}").model_dump() - ] - assert imports.agent_calls == [] - - -@pytest.mark.parametrize( - "url", - [ - "https://example.com/download?id=123", - "https://example.com/agent.ifpkg", - " https://example.com/agent.IFPKG?token=secret#download ", - ], -) -def test_package_url_uses_agent_import(app: Flask, imports: Imports, monkeypatch: pytest.MonkeyPatch, url: str) -> None: - fetched = _mock_download(monkeypatch) - data, status = _post(app, imports, json={"mode": "yaml-url", "yaml_url": url}) - assert status == 200 - assert data["app_mode"] == "agent" - assert data["status"] == "completed" - assert len(fetched) == 1 - assert imports.dsl_calls == [] - (call,) = imports.agent_calls - assert call.tenant_id == imports.context.active_workspace_id - assert call.account.id == imports.context.account_id - assert call.account.is_admin_or_owner is True - - -@pytest.mark.parametrize( - "url", ["https://example.com/app.yaml", "https://example.com/download", "https://example.com/agent.ifpkg"] -) -def test_yaml_url_keeps_dsl_import_without_package_retry( - app: Flask, imports: Imports, monkeypatch: pytest.MonkeyPatch, url: str -) -> None: - imports.status = ImportStatus.FAILED - fetched = _mock_download(monkeypatch, b"app: {}") - _, status = _post(app, imports, json={"mode": "yaml-url", "yaml_url": url, "app_id": "existing", "name": "Renamed"}) - assert status == 400 - assert ( - imports.dsl_calls[0].model_dump() - == AppImportParams(mode="yaml-content", yaml_content="app: {}", app_id="existing", name="Renamed").model_dump() - ) - assert len(fetched) == 1 - assert imports.agent_calls == [] - - -def test_url_without_any_import_permission_does_not_download( - app: Flask, imports: Imports, monkeypatch: pytest.MonkeyPatch, config_overrides: Callable[..., None] -) -> None: - config_overrides(RBAC_ENABLED=True) - monkeypatch.setattr( - "services.app.console_gateway.rbac_service.RBACService.CheckAccess.check", lambda *_a, **_k: False - ) - fetched = _mock_download(monkeypatch) - with pytest.raises(Forbidden): - _post(app, imports, json={"mode": "yaml-url", "yaml_url": "https://example.com/download"}) - assert fetched == [] - - -@pytest.mark.parametrize("is_yaml", [False, True]) -def test_url_quota_is_enforced_after_content_detection( - app: Flask, - imports: Imports, - app_query_services: ControllerTestServices, - monkeypatch: pytest.MonkeyPatch, - config_overrides: Callable[..., None], - is_yaml: bool, -) -> None: - config_overrides(DEPLOYMENT_EDITION="CLOUD") - transfers = app_query_services.apps.console._transfers - assert isinstance(transfers, AppTransferGateway) - monkeypatch.setattr(transfers, "import_dsl", AppTransferGateway.import_dsl.__get__(transfers)) - features = FeatureModel(apps=LimitationModel(size=1, limit=1)) - monkeypatch.setattr("services.app_dsl_service.FeatureService.get_features", lambda *_a, **_k: features) - _mock_download(monkeypatch, b"app: {mode: workflow}" if is_yaml else _roster_archive().getvalue()) - if is_yaml: - with pytest.raises(Forbidden, match="number of apps"): - _post(app, imports, json={"mode": "yaml-url", "yaml_url": "https://example.com/download"}) - assert imports.agent_calls == [] - assert imports.dsl_calls == [] - else: - data, status = _post(app, imports, json={"mode": "yaml-url", "yaml_url": "https://example.com/download"}) - assert status == 200 - assert data["app_mode"] == "agent" - - -def test_url_rejects_content_that_is_neither_yaml_nor_package( - app: Flask, imports: Imports, monkeypatch: pytest.MonkeyPatch -) -> None: - fetched = _mock_download(monkeypatch, b"[invalid") - with pytest.raises(InvalidRosterAgentPackageError): - _post(app, imports, json={"mode": "yaml-url", "yaml_url": "https://example.com/download"}) - assert len(fetched) == 1 - - -@pytest.mark.parametrize("has_warning", [False, True]) -def test_existing_import_route_accepts_package_and_preserves_import_response( - app: Flask, imports: Imports, has_warning: bool -) -> None: - imports.warnings = ( - [DslImportWarning(code="agent_skill_missing", path="skills.s_000001", message="Missing Skill")] - if has_warning - else [] - ) - data, status = _post(app, imports, data={"file": (_roster_archive(), "agent.ifpkg")}) - assert status == 200 - assert data["app_id"] == "app-1" - assert data["app_mode"] == "agent" - assert data["status"] == ("completed-with-warnings" if has_warning else "completed") - assert data["warnings"] == [warning.model_dump(mode="json") for warning in imports.warnings] - assert data["id"] - assert imports.agent_calls[0].content == _roster_archive().getvalue() - - -@pytest.mark.parametrize("from_url", [False, True]) -@pytest.mark.parametrize("denied", [RBACPermission.AGENT_CREATE, RBACPermission.AGENT_IMPORT_EXPORT_DSL]) -def test_package_import_checks_agent_permissions_before_materializing_resources( - app: Flask, - imports: Imports, - monkeypatch: pytest.MonkeyPatch, - config_overrides: Callable[..., None], - denied: RBACPermission, - from_url: bool, -) -> None: - config_overrides(RBAC_ENABLED=True) - _mock_download(monkeypatch) - scenes: list[RBACPermission] = [] - - def check(*_args: object, scene: RBACPermission, **_kwargs: object) -> bool: - scenes.append(scene) - return scene != denied - - monkeypatch.setattr("services.app.console_gateway.rbac_service.RBACService.CheckAccess.check", check) - kwargs: dict[str, dict[str, object]] = ( - {"json": {"mode": "yaml-url", "yaml_url": "https://example.com/agent.ifpkg"}} - if from_url - else {"data": {"file": (_roster_archive(), "agent.ifpkg")}} - ) - with pytest.raises(Forbidden): - _post(app, imports, **kwargs) - assert denied in scenes - if not from_url: - assert RBACPermission.APP_IMPORT_EXPORT_DSL not in scenes - assert imports.agent_calls == [] - - -@pytest.mark.parametrize("from_url", [False, True]) -def test_package_import_rejects_overwrite( - app: Flask, imports: Imports, monkeypatch: pytest.MonkeyPatch, from_url: bool -) -> None: - _mock_download(monkeypatch) - kwargs: dict[str, dict[str, object]] = ( - {"json": {"mode": "yaml-url", "yaml_url": "https://example.com/agent.ifpkg", "app_id": "existing"}} - if from_url - else {"data": {"app_id": "existing", "file": (_roster_archive(), "agent.ifpkg")}} - ) - with pytest.raises(InvalidRosterAgentPackageError, match="overwriting"): - _post(app, imports, **kwargs) - - -@pytest.mark.parametrize("from_url", [False, True]) -@pytest.mark.parametrize("app_id", [None, "existing"]) -@pytest.mark.parametrize("status", [ImportStatus.COMPLETED, ImportStatus.PENDING, ImportStatus.FAILED]) -def test_ordinary_package_import_uses_dsl_permissions_and_confirmation( - app: Flask, - imports: Imports, - monkeypatch: pytest.MonkeyPatch, - app_id: str | None, - status: ImportStatus, - from_url: bool, -) -> None: - imports.status = status - dsl = "kind: app\nversion: 99.0.0\napp:\n mode: workflow\n" - with AppPackageService().export(dsl=dsl, name="Workflow") as package: - form: dict[str, object] = {"file": (package.archive, "workflow.ifpkg"), "name": "Renamed"} - if app_id: - form["app_id"] = app_id - if from_url: - _mock_download(monkeypatch, package.archive.read()) - data, code = _post( - app, - imports, - json={ - "mode": "yaml-url", - "yaml_url": "https://example.com/workflow.ifpkg", - "name": "Renamed", - "app_id": app_id, - }, - ) - else: - data, code = _post(app, imports, data=form) - assert code == {ImportStatus.FAILED: 400, ImportStatus.PENDING: 202}.get(status, 200) - assert len(imports.packages) == 1 - assert imports.packages[0].archive.closed - assert data["status"] == status - assert len(imports.dsl_calls) == 1 - assert ( - imports.dsl_calls[0].model_dump() - == AppImportParams(mode="yaml-content", yaml_content=dsl, name="Renamed", app_id=app_id).model_dump() - ) - assert imports.agent_calls == [] - - -def test_export_query_rejects_conflicting_version_selectors() -> None: - with pytest.raises(ValidationError, match="version_id and workflow_id cannot be used together"): - app_module.AppExportQuery.model_validate( - {"version_id": "11111111-1111-4111-8111-111111111111", "workflow_id": "workflow-1"} - ) - - -def test_template_import_resolves_catalog_source_and_preserves_overrides( - app: Flask, imports: Imports, app_query_services: ControllerTestServices, monkeypatch: pytest.MonkeyPatch -) -> None: - transfers = app_query_services.apps.console._transfers - assert isinstance(transfers, AppTransferGateway) - template_id = UUID("11111111-1111-4111-8111-111111111111") - version_id = UUID("22222222-2222-4222-8222-222222222222") - source = RecommendedAgentPackageSource("source-tenant", "source-agent", version_id) - catalog = Mock() - catalog.get_package_source.return_value = source - transfers._recommended_packages = RecommendedAppPackageService(sources=catalog, exporter=Mock()) - import_template = Mock(return_value=RosterAgentPackageImportResult("new-app", "new-agent", [])) - monkeypatch.setattr(transfers._agent_importer, "import_template", import_template) - - data, status = _post( - app, - imports, - json={ - "mode": "template", - "template_id": str(template_id), - "version_id": str(version_id), - "name": "My Agent", - "app_id": "", - }, - ) - - assert status == 200 - assert data["app_id"] == "new-app" - assert data["app_mode"] == "agent" - catalog.get_package_source.assert_called_once_with(str(template_id), version_id) - import_template.assert_called_once() - assert import_template.call_args.kwargs["source"] == source - assert import_template.call_args.kwargs["tenant_id"] == imports.context.active_workspace_id - assert import_template.call_args.kwargs["account"].id == imports.context.account_id - assert import_template.call_args.kwargs["name"] == "My Agent" - - -def test_template_import_rejects_unlisted_source( - app: Flask, imports: Imports, app_query_services: ControllerTestServices, monkeypatch: pytest.MonkeyPatch -) -> None: - transfers = app_query_services.apps.console._transfers - assert isinstance(transfers, AppTransferGateway) - catalog = Mock() - catalog.get_package_source.return_value = None - transfers._recommended_packages = RecommendedAppPackageService(sources=catalog, exporter=Mock()) - import_template = Mock() - monkeypatch.setattr(transfers._agent_importer, "import_template", import_template) - - with pytest.raises(import_module.RecommendedAppNotFoundHttpError): - _post( - app, - imports, - json={ - "mode": "template", - "template_id": "11111111-1111-4111-8111-111111111111", - "version_id": "22222222-2222-4222-8222-222222222222", - }, - ) - import_template.assert_not_called() - - -def test_package_url_import_preserves_overrides( - app: Flask, imports: Imports, app_query_services: ControllerTestServices, monkeypatch: pytest.MonkeyPatch -) -> None: - transfers = app_query_services.apps.console._transfers - assert isinstance(transfers, AppTransferGateway) - import_url = Mock(return_value=RosterAgentPackageImportResult("new-app", "new-agent", [])) - monkeypatch.setattr(transfers._agent_importer, "import_package_url", import_url) - - data, status = _post( - app, - imports, - json={"mode": "ifpkg-url", "package_url": "https://example.com/agent.ifpkg", "name": "My Agent"}, - ) - - assert status == 200 - assert data["app_id"] == "new-app" - import_url.assert_called_once() - assert import_url.call_args.kwargs["url"] == "https://example.com/agent.ifpkg" - assert import_url.call_args.kwargs["name"] == "My Agent" - - -def test_package_url_retains_legacy_optional_selectors( - app: Flask, imports: Imports, app_query_services: ControllerTestServices, monkeypatch: pytest.MonkeyPatch -) -> None: - transfers = app_query_services.apps.console._transfers - assert isinstance(transfers, AppTransferGateway) - import_url = Mock(return_value=RosterAgentPackageImportResult("new-app", "new-agent", [])) - monkeypatch.setattr(transfers._agent_importer, "import_package_url", import_url) - - data, status = _post( - app, - imports, - json={ - "mode": "ifpkg-url", - "package_url": "https://example.com/agent.ifpkg", - "app_id": "", - "template_id": "11111111-1111-4111-8111-111111111111", - "version_id": "22222222-2222-4222-8222-222222222222", - }, - ) - - assert status == 200 - assert data["app_id"] == "new-app" - import_url.assert_called_once() - assert import_url.call_args.kwargs["url"] == "https://example.com/agent.ifpkg" - - -@pytest.mark.parametrize("mode", ["template", "ifpkg-url"]) -def test_template_import_checks_permissions_before_accessing_source( - app: Flask, - imports: Imports, - app_query_services: ControllerTestServices, - monkeypatch: pytest.MonkeyPatch, - config_overrides: Callable[..., None], - mode: str, -) -> None: - config_overrides(RBAC_ENABLED=True) - monkeypatch.setattr( - "services.app.console_gateway.rbac_service.RBACService.CheckAccess.check", lambda *_a, **_k: False - ) - transfers = app_query_services.apps.console._transfers - assert isinstance(transfers, AppTransferGateway) - importer = Mock() - monkeypatch.setattr(transfers, "_agent_importer", importer) - payload: dict[str, object] = ( - { - "mode": "template", - "template_id": "11111111-1111-4111-8111-111111111111", - "version_id": "22222222-2222-4222-8222-222222222222", - } - if mode == "template" - else {"mode": "ifpkg-url", "package_url": "https://example.com/agent.ifpkg"} - ) - with pytest.raises(Forbidden): - _post(app, imports, json=payload) - importer.assert_not_called() - - -@pytest.mark.parametrize( - "extra", - [ - {"app_id": "existing"}, - {"yaml_content": "app: {}"}, - ], -) -def test_package_url_rejects_ambiguous_sources(app: Flask, imports: Imports, extra: dict[str, str]) -> None: - with pytest.raises(InvalidRosterAgentPackageError): - _post(app, imports, json={"mode": "ifpkg-url", "package_url": "https://example.com/agent.ifpkg", **extra}) - - -@pytest.mark.parametrize("quota", [False, True]) -def test_ordinary_package_import_checks_permission_and_quota( - app: Flask, - imports: Imports, - app_query_services: ControllerTestServices, - monkeypatch: pytest.MonkeyPatch, - config_overrides: Callable[..., None], - quota: bool, -) -> None: - config_overrides(RBAC_ENABLED=not quota, DEPLOYMENT_EDITION="CLOUD" if quota else "COMMUNITY") - transfers = app_query_services.apps.console._transfers - assert isinstance(transfers, AppTransferGateway) - monkeypatch.setattr(transfers, "import_dsl", AppTransferGateway.import_dsl.__get__(transfers)) - monkeypatch.setattr( - "services.app.console_gateway.rbac_service.RBACService.CheckAccess.check", lambda *_a, **_k: False - ) - features = FeatureModel(apps=LimitationModel(size=1, limit=1)) - monkeypatch.setattr("services.app_dsl_service.FeatureService.get_features", lambda *_a, **_k: features) - with AppPackageService().export(dsl="kind: app\napp: {mode: workflow}\n", name="Workflow") as package: - with pytest.raises(Forbidden): - _post(app, imports, data={"file": (package.archive, "workflow.ifpkg")}) - assert imports.dsl_calls == [] diff --git a/api/tests/unit_tests/controllers/console/app/test_statistic_api.py b/api/tests/unit_tests/controllers/console/app/test_statistic.py similarity index 100% rename from api/tests/unit_tests/controllers/console/app/test_statistic_api.py rename to api/tests/unit_tests/controllers/console/app/test_statistic.py diff --git a/api/tests/unit_tests/controllers/console/app/test_workflow.py b/api/tests/unit_tests/controllers/console/app/test_workflow.py index 55117245293..e66027d8f24 100644 --- a/api/tests/unit_tests/controllers/console/app/test_workflow.py +++ b/api/tests/unit_tests/controllers/console/app/test_workflow.py @@ -3,48 +3,55 @@ from __future__ import annotations import inspect import json from contextlib import contextmanager, nullcontext +from dataclasses import dataclass from datetime import datetime +from inspect import unwrap from types import SimpleNamespace -from unittest.mock import Mock +from unittest.mock import ANY, MagicMock, Mock import pytest from flask import Flask from pydantic import ValidationError from sqlalchemy.orm import Session -from werkzeug.exceptions import HTTPException, NotFound +from werkzeug.exceptions import HTTPException, NotFound, UnprocessableEntity from controllers.common.errors import InvalidArgumentError +from controllers.console import wraps as console_wraps from controllers.console.app import workflow as workflow_module +from controllers.console.app import wraps as app_wraps from controllers.console.app.error import DraftWorkflowNotExist, DraftWorkflowNotSync +from controllers.console.app.workflow import ConvertToWorkflowApi from core.workflow.llm_environment_variable import LLMEnvironmentVariable +from enums import DeploymentEdition from graphon.file import File, FileTransferMethod, FileType from graphon.variables import SecretVariable, StringVariable from graphon.variables.variables import RAGPipelineVariable -from models.account import Account -from models.model import App, AppMode +from libs import login as login_lib +from models import Account as _workflow_convert_api_Account +from models import App as _workflow_convert_api_App +from models import App as _workflow_human_input_debug_api_App +from models import Tenant +from models.account import Account as _workflow_Account +from models.account import Account as _workflow_human_input_debug_api_Account +from models.account import AccountStatus, TenantAccountRole +from models.model import App as _workflow_App +from models.model import AppMode, IconType from models.workflow import Workflow, WorkflowType from services import workflow_variable_reference_validator as variable_reference_validator from services.app.access import AppAccessFilter -from tests.unit_tests.config_override import apply_config_overrides -from tests.unit_tests.controllers.error_body import raised_error +from tests.helpers.unit.config_override import apply_config_overrides +from tests.helpers.unit.controllers.error_body import raised_error +from tests.helpers.unit.model_factories import make_app -@pytest.fixture(autouse=True) -def _identity_workflow_encryption(monkeypatch: pytest.MonkeyPatch) -> None: - """Keep variable serialization focused on ORM behavior, not the external key provider.""" - - monkeypatch.setattr(workflow_module.encrypter, "encrypt_token", lambda *, tenant_id, token: token) - monkeypatch.setattr(workflow_module.encrypter, "decrypt_token", lambda *, tenant_id, token: token) - - -def _account() -> Account: - account = Account(name="Alice", email="alice@example.com") +def _account() -> _workflow_Account: + account = _workflow_Account(name="Alice", email="alice@example.com") account.id = "user-1" return account -def _app(*, app_id: str = "app", tenant_id: str = "t1") -> App: - return App( +def _workflow_app(*, app_id: str = "app", tenant_id: str = "t1") -> _workflow_App: + return _workflow_App( id=app_id, tenant_id=tenant_id, name="Workflow App", @@ -137,851 +144,885 @@ def _make_workflow(**overrides) -> Workflow: return workflow -@pytest.mark.parametrize( - "advisory", ["clean", "warning", "checker-error", "formatter-error", "empty", "non-object", "invalid-json"] -) -def test_publish_workflow_returns_success( - app: Flask, - monkeypatch: pytest.MonkeyPatch, - advisory: str, -) -> None: - current_user = SimpleNamespace(id="account-1") - app_model = SimpleNamespace(id="app-1", tenant_id="tenant-1") - graph = { - "nodes": [ - {"id": "start", "data": {"type": "start"}}, - {"id": "branch", "data": {"type": "if-else"}}, - {"id": "producer", "data": {"type": "code", "title": "Producer"}}, - { - "id": "consumer", - "data": {"type": "answer", "title": "Consumer", "answer": "{{#producer.text#}}"}, - }, - ], - "edges": [ - {"source": "start", "target": "branch"}, - {"source": "branch", "target": "producer", "sourceHandle": "true"}, - {"source": "branch", "target": "consumer", "sourceHandle": "false"}, - ], - } - workflow = SimpleNamespace( - id="published-workflow", - created_at=datetime(2026, 8, 17, 12, 0, 0), - graph={"clean": "{}", "empty": None, "non-object": "[]", "invalid-json": "{"}.get(advisory, json.dumps(graph)), +class TestWorkflowCases: + @staticmethod + @pytest.fixture(autouse=True) + def _identity_workflow_encryption(monkeypatch: pytest.MonkeyPatch) -> None: + """Keep variable serialization focused on ORM behavior, not the external key provider.""" + + monkeypatch.setattr(workflow_module.encrypter, "encrypt_token", lambda *, tenant_id, token: token) + monkeypatch.setattr(workflow_module.encrypter, "decrypt_token", lambda *, tenant_id, token: token) + + @staticmethod + @pytest.mark.parametrize( + "advisory", ["clean", "warning", "checker-error", "formatter-error", "empty", "non-object", "invalid-json"] ) - if advisory == "checker-error": + def test_publish_workflow_returns_success( + app: Flask, + monkeypatch: pytest.MonkeyPatch, + advisory: str, + ) -> None: + current_user = SimpleNamespace(id="account-1") + app_model = SimpleNamespace(id="app-1", tenant_id="tenant-1") + graph = { + "nodes": [ + {"id": "start", "data": {"type": "start"}}, + {"id": "branch", "data": {"type": "if-else"}}, + {"id": "producer", "data": {"type": "code", "title": "Producer"}}, + { + "id": "consumer", + "data": {"type": "answer", "title": "Consumer", "answer": "{{#producer.text#}}"}, + }, + ], + "edges": [ + {"source": "start", "target": "branch"}, + {"source": "branch", "target": "producer", "sourceHandle": "true"}, + {"source": "branch", "target": "consumer", "sourceHandle": "false"}, + ], + } + workflow = SimpleNamespace( + id="published-workflow", + created_at=datetime(2026, 8, 17, 12, 0, 0), + graph={"clean": "{}", "empty": None, "non-object": "[]", "invalid-json": "{"}.get( + advisory, json.dumps(graph) + ), + ) + if advisory == "checker-error": + monkeypatch.setattr( + variable_reference_validator, "validate_variable_references", Mock(side_effect=RuntimeError("checker")) + ) + elif advisory == "formatter-error": + monkeypatch.setattr( + variable_reference_validator, + "format_variable_reference_errors", + Mock(side_effect=RuntimeError("format")), + ) + session = Mock() + session.get.return_value = app_model monkeypatch.setattr( - variable_reference_validator, "validate_variable_references", Mock(side_effect=RuntimeError("checker")) + workflow_module, + "WorkflowService", + Mock(return_value=SimpleNamespace(publish_workflow=Mock(return_value=workflow))), ) - elif advisory == "formatter-error": monkeypatch.setattr( - variable_reference_validator, "format_variable_reference_errors", Mock(side_effect=RuntimeError("format")) + workflow_module, + "sessionmaker", + lambda _engine: SimpleNamespace(begin=lambda: nullcontext(session)), ) - session = Mock() - session.get.return_value = app_model - monkeypatch.setattr( - workflow_module, - "WorkflowService", - Mock(return_value=SimpleNamespace(publish_workflow=Mock(return_value=workflow))), - ) - monkeypatch.setattr( - workflow_module, - "sessionmaker", - lambda _engine: SimpleNamespace(begin=lambda: nullcontext(session)), - ) - monkeypatch.setattr(workflow_module, "db", SimpleNamespace(engine=object())) - with app.test_request_context("/apps/app-1/workflows/publish", method="POST", json={}): - response = inspect.unwrap(workflow_module.PublishedWorkflowApi.post)( - workflow_module.PublishedWorkflowApi(), - workflow_module.PublishWorkflowPayload.model_validate({}), - current_user, - app_model, + monkeypatch.setattr(workflow_module, "db", SimpleNamespace(engine=object())) + with app.test_request_context("/apps/app-1/workflows/publish", method="POST", json={}): + response = inspect.unwrap(workflow_module.PublishedWorkflowApi.post)( + workflow_module.PublishedWorkflowApi(), + workflow_module.PublishWorkflowPayload.model_validate({}), + current_user, + app_model, + ) + + assert response["result"] == "success" + assert app_model.workflow_id == workflow.id + assert isinstance(response["created_at"], int) + if advisory == "warning": + assert "Consumer" in response["warning"] + assert "Producer" in response["warning"] + assert "skipped branch" in response["warning"] + else: + assert "warning" not in response + + @staticmethod + @pytest.mark.parametrize("transaction_fails", [False, True], ids=["commit-succeeds", "commit-fails"]) + def test_delete_workflow_retires_candidates_only_after_transaction_exit( + app: Flask, + monkeypatch: pytest.MonkeyPatch, + transaction_fails: bool, + ) -> None: + current_user = SimpleNamespace(id="account-1") + app_model = SimpleNamespace(id="app-1", tenant_id="tenant-1") + session = Mock() + events: list[str] = [] + error = RuntimeError("commit failed") + workflow_service = SimpleNamespace( + delete_workflow=Mock(side_effect=lambda **_kwargs: events.append("delete") or ["inline-agent"]) ) - assert response["result"] == "success" - assert app_model.workflow_id == workflow.id - assert isinstance(response["created_at"], int) - if advisory == "warning": - assert "Consumer" in response["warning"] - assert "Producer" in response["warning"] - assert "skipped branch" in response["warning"] - else: - assert "warning" not in response + @contextmanager + def transaction(): + events.append("transaction-enter") + yield session + events.append("transaction-exit") + if transaction_fails: + raise error + monkeypatch.setattr(workflow_module, "WorkflowService", Mock(return_value=workflow_service)) + monkeypatch.setattr( + workflow_module, + "sessionmaker", + lambda _engine: SimpleNamespace(begin=transaction), + ) + monkeypatch.setattr(workflow_module, "db", SimpleNamespace(engine=object())) + retire_unowned = Mock(side_effect=lambda **_kwargs: events.append("retire")) + monkeypatch.setattr(workflow_module.WorkflowAgentRetirementService, "retire_unowned", retire_unowned) -@pytest.mark.parametrize("transaction_fails", [False, True], ids=["commit-succeeds", "commit-fails"]) -def test_delete_workflow_retires_candidates_only_after_transaction_exit( - app: Flask, - monkeypatch: pytest.MonkeyPatch, - transaction_fails: bool, -) -> None: - current_user = SimpleNamespace(id="account-1") - app_model = SimpleNamespace(id="app-1", tenant_id="tenant-1") - session = Mock() - events: list[str] = [] - error = RuntimeError("commit failed") - workflow_service = SimpleNamespace( - delete_workflow=Mock(side_effect=lambda **_kwargs: events.append("delete") or ["inline-agent"]) - ) - - @contextmanager - def transaction(): - events.append("transaction-enter") - yield session - events.append("transaction-exit") - if transaction_fails: - raise error - - monkeypatch.setattr(workflow_module, "WorkflowService", Mock(return_value=workflow_service)) - monkeypatch.setattr( - workflow_module, - "sessionmaker", - lambda _engine: SimpleNamespace(begin=transaction), - ) - monkeypatch.setattr(workflow_module, "db", SimpleNamespace(engine=object())) - retire_unowned = Mock(side_effect=lambda **_kwargs: events.append("retire")) - monkeypatch.setattr(workflow_module.WorkflowAgentRetirementService, "retire_unowned", retire_unowned) - - with app.test_request_context("/apps/app-1/workflows/workflow-1", method="DELETE"): - if transaction_fails: - with pytest.raises(RuntimeError) as exc_info: - inspect.unwrap(workflow_module.WorkflowByIdApi.delete)( + with app.test_request_context("/apps/app-1/workflows/workflow-1", method="DELETE"): + if transaction_fails: + with pytest.raises(RuntimeError) as exc_info: + inspect.unwrap(workflow_module.WorkflowByIdApi.delete)( + workflow_module.WorkflowByIdApi(), + current_user, + app_model, + "workflow-1", + ) + assert exc_info.value is error + else: + response = inspect.unwrap(workflow_module.WorkflowByIdApi.delete)( workflow_module.WorkflowByIdApi(), current_user, app_model, "workflow-1", ) - assert exc_info.value is error + assert response == (None, 204) + + assert events == ["transaction-enter", "delete", "transaction-exit"] + ([] if transaction_fails else ["retire"]) + if transaction_fails: + retire_unowned.assert_not_called() else: - response = inspect.unwrap(workflow_module.WorkflowByIdApi.delete)( - workflow_module.WorkflowByIdApi(), - current_user, - app_model, - "workflow-1", + retire_unowned.assert_called_once_with( + tenant_id=app_model.tenant_id, + agent_ids=["inline-agent"], + account_id=current_user.id, ) - assert response == (None, 204) - assert events == ["transaction-enter", "delete", "transaction-exit"] + ([] if transaction_fails else ["retire"]) - if transaction_fails: - retire_unowned.assert_not_called() - else: - retire_unowned.assert_called_once_with( - tenant_id=app_model.tenant_id, - agent_ids=["inline-agent"], - account_id=current_user.id, + @staticmethod + def test_parse_file_no_config(monkeypatch: pytest.MonkeyPatch) -> None: + monkeypatch.setattr(workflow_module.FileUploadConfigManager, "convert", lambda *_args, **_kwargs: None) + workflow = _make_workflow(features_dict={}) + + assert workflow_module._parse_file(workflow, files=[{"id": "f"}]) == [] + + @staticmethod + def test_parse_file_with_config(monkeypatch: pytest.MonkeyPatch) -> None: + config = object() + file_list = [ + File( + tenant_id="t1", + file_type=FileType.IMAGE, + transfer_method=FileTransferMethod.REMOTE_URL, + remote_url="http://u", + ) + ] + build_mock = Mock(return_value=file_list) + monkeypatch.setattr(workflow_module.FileUploadConfigManager, "convert", lambda *_args, **_kwargs: config) + monkeypatch.setattr(workflow_module.file_factory, "build_from_mappings", build_mock) + + workflow = _make_workflow(features_dict={}) + result = workflow_module._parse_file(workflow, files=[{"id": "f"}]) + + assert result == file_list + build_mock.assert_called_once() + + @staticmethod + def test_sync_draft_workflow_invalid_content_type(app: Flask, monkeypatch: pytest.MonkeyPatch) -> None: + api = workflow_module.DraftWorkflowApi() + handler = inspect.unwrap(api.post) + + with app.test_request_context("/apps/app/workflows/draft", method="POST", data="x", content_type="text/html"): + with pytest.raises(HTTPException) as exc: + handler(api, _account(), app_model=_workflow_app()) + + assert exc.value.code == 415 + + @staticmethod + def test_sync_draft_workflow_invalid_json(app: Flask, monkeypatch: pytest.MonkeyPatch) -> None: + api = workflow_module.DraftWorkflowApi() + handler = inspect.unwrap(api.post) + + with app.test_request_context( + "/apps/app/workflows/draft", + method="POST", + data="[]", + content_type="application/json", + ): + response, status = raised_error(lambda: handler(api, _account(), app_model=_workflow_app())) + + assert status == 400 + assert response["message"] == "Invalid JSON data" + + @staticmethod + def test_sync_draft_workflow_success(app: Flask, monkeypatch: pytest.MonkeyPatch) -> None: + workflow = _make_workflow(updated_at=None, created_at=datetime(2024, 1, 1)) + + monkeypatch.setattr( + workflow_module.variable_factory, "build_environment_variable_from_mapping", lambda *_args: "env" + ) + monkeypatch.setattr( + workflow_module.variable_factory, "build_conversation_variable_from_mapping", lambda *_args: "conv" ) + sync_draft_workflow = Mock(return_value=workflow) + service = SimpleNamespace(sync_draft_workflow=sync_draft_workflow) + monkeypatch.setattr(workflow_module, "WorkflowService", lambda: service) -def test_parse_file_no_config(monkeypatch: pytest.MonkeyPatch) -> None: - monkeypatch.setattr(workflow_module.FileUploadConfigManager, "convert", lambda *_args, **_kwargs: None) - workflow = _make_workflow(features_dict={}) + api = workflow_module.DraftWorkflowApi() + handler = inspect.unwrap(api.post) - assert workflow_module._parse_file(workflow, files=[{"id": "f"}]) == [] + with app.test_request_context( + "/apps/app/workflows/draft", + method="POST", + json={"graph": {}, "features": {}, "hash": "h"}, + ): + response = handler(api, _account(), app_model=_workflow_app()) + assert response["result"] == "success" + assert sync_draft_workflow.call_args.kwargs["environment_variables"] == [] + assert sync_draft_workflow.call_args.kwargs["preserve_environment_variables"] is True -def test_parse_file_with_config(monkeypatch: pytest.MonkeyPatch) -> None: - config = object() - file_list = [ - File( - tenant_id="t1", - file_type=FileType.IMAGE, - transfer_method=FileTransferMethod.REMOTE_URL, - remote_url="http://u", + @staticmethod + @pytest.mark.parametrize("content_type", ["application/json", "text/plain"]) + @pytest.mark.parametrize("confirmation", [{}, {"force": False}], ids=["missing-force", "force-false"]) + def test_sync_draft_workflow_rejects_empty_graph_without_confirmation( + app: Flask, + monkeypatch: pytest.MonkeyPatch, + content_type: str, + confirmation: dict[str, bool], + ) -> None: + workflow_service = Mock() + monkeypatch.setattr(workflow_module, "WorkflowService", workflow_service) + api = workflow_module.DraftWorkflowApi() + handler = inspect.unwrap(api.post) + + with app.test_request_context( + "/apps/app/workflows/draft", + method="POST", + data=json.dumps({"graph": {"nodes": [], "edges": []}, "features": {}, "hash": "h", **confirmation}), + content_type=content_type, + ): + with pytest.raises(InvalidArgumentError) as exc: + handler(api, _account(), app_model=_workflow_app()) + + assert exc.value.code == 400 + assert exc.value.description == "Saving an empty workflow requires force=true." + workflow_service.assert_not_called() + + @staticmethod + @pytest.mark.parametrize("content_type", ["application/json", "text/plain"]) + @pytest.mark.parametrize( + ("graph", "confirmation"), + [ + ({"nodes": [], "edges": []}, {"force": True}), + ({"nodes": [{"id": "start"}], "edges": []}, {}), + ({"nodes": [], "edges": [{"id": "edge-1", "source": "start", "target": "end"}]}, {}), + ], + ids=["confirmed-empty-graph", "nodes-without-edges", "edges-without-nodes"], + ) + def test_sync_draft_workflow_saves_confirmed_empty_or_nonempty_graph( + app: Flask, + monkeypatch: pytest.MonkeyPatch, + content_type: str, + graph: dict[str, list[dict[str, str]]], + confirmation: dict[str, bool], + ) -> None: + workflow = SimpleNamespace(unique_hash="next-hash", updated_at=None, created_at=datetime(2024, 1, 1)) + sync_draft_workflow = Mock(return_value=workflow) + monkeypatch.setattr( + workflow_module, + "WorkflowService", + lambda: SimpleNamespace(sync_draft_workflow=sync_draft_workflow), ) - ] - build_mock = Mock(return_value=file_list) - monkeypatch.setattr(workflow_module.FileUploadConfigManager, "convert", lambda *_args, **_kwargs: config) - monkeypatch.setattr(workflow_module.file_factory, "build_from_mappings", build_mock) + api = workflow_module.DraftWorkflowApi() + handler = inspect.unwrap(api.post) - workflow = _make_workflow(features_dict={}) - result = workflow_module._parse_file(workflow, files=[{"id": "f"}]) + with app.test_request_context( + "/apps/app/workflows/draft", + method="POST", + data=json.dumps({"graph": graph, "features": {}, "hash": "current-hash", **confirmation}), + content_type=content_type, + ): + response = handler(api, _account(), app_model=_workflow_app()) - assert result == file_list - build_mock.assert_called_once() + assert response["result"] == "success" + assert response["hash"] == "next-hash" + sync_draft_workflow.assert_called_once() + assert sync_draft_workflow.call_args.kwargs["graph"] == graph + assert sync_draft_workflow.call_args.kwargs["unique_hash"] == "current-hash" + @staticmethod + @pytest.mark.parametrize("content_type", ["application/json", "text/plain"]) + def test_sync_draft_workflow_force_does_not_bypass_hash_conflict( + app: Flask, + monkeypatch: pytest.MonkeyPatch, + content_type: str, + ) -> None: + sync_draft_workflow = Mock(side_effect=workflow_module.WorkflowHashNotEqualError()) + monkeypatch.setattr( + workflow_module, + "WorkflowService", + lambda: SimpleNamespace(sync_draft_workflow=sync_draft_workflow), + ) + api = workflow_module.DraftWorkflowApi() + handler = inspect.unwrap(api.post) -def test_sync_draft_workflow_invalid_content_type(app: Flask, monkeypatch: pytest.MonkeyPatch) -> None: - api = workflow_module.DraftWorkflowApi() - handler = inspect.unwrap(api.post) + with app.test_request_context( + "/apps/app/workflows/draft", + method="POST", + data=json.dumps({"graph": {"nodes": [], "edges": []}, "features": {}, "hash": "stale-hash", "force": True}), + content_type=content_type, + ): + with pytest.raises(DraftWorkflowNotSync) as exc: + handler(api, _account(), app_model=_workflow_app()) - with app.test_request_context("/apps/app/workflows/draft", method="POST", data="x", content_type="text/html"): - with pytest.raises(HTTPException) as exc: - handler(api, _account(), app_model=_app()) + assert exc.value.code == 409 + sync_draft_workflow.assert_called_once() + assert sync_draft_workflow.call_args.kwargs["unique_hash"] == "stale-hash" - assert exc.value.code == 415 + @staticmethod + @pytest.mark.parametrize("content_type", ["application/json", "text/plain"]) + @pytest.mark.parametrize("force", ["true", 1, None], ids=["string", "integer", "null"]) + def test_sync_draft_workflow_rejects_non_boolean_force( + app: Flask, + monkeypatch: pytest.MonkeyPatch, + content_type: str, + force: str | int | None, + ) -> None: + workflow_service = Mock() + monkeypatch.setattr(workflow_module, "WorkflowService", workflow_service) + api = workflow_module.DraftWorkflowApi() + handler = inspect.unwrap(api.post) + with app.test_request_context( + "/apps/app/workflows/draft", + method="POST", + data=json.dumps({"graph": {"nodes": [], "edges": []}, "features": {}, "force": force}), + content_type=content_type, + ): + if content_type == "application/json": + with pytest.raises(ValidationError, match="Input should be a valid boolean"): + handler(api, _account(), app_model=_workflow_app()) + else: + response, status = raised_error(lambda: handler(api, _account(), app_model=_workflow_app())) + assert status == 400 + assert response["message"] == "Invalid JSON data" -def test_sync_draft_workflow_invalid_json(app: Flask, monkeypatch: pytest.MonkeyPatch) -> None: - api = workflow_module.DraftWorkflowApi() - handler = inspect.unwrap(api.post) + workflow_service.assert_not_called() - with app.test_request_context( - "/apps/app/workflows/draft", - method="POST", - data="[]", - content_type="application/json", - ): - response, status = raised_error(lambda: handler(api, _account(), app_model=_app())) + @staticmethod + def test_sync_draft_workflow_passes_environment_patch(app: Flask, monkeypatch: pytest.MonkeyPatch) -> None: + workflow = _make_workflow(updated_at=None, created_at=datetime(2024, 1, 1)) + patched_variable = StringVariable( + id="env-model", + name="shared_model", + value="model", + selector=["env", "shared_model"], + ) + build_environment_variable = Mock(return_value=patched_variable) + sync_draft_workflow = Mock(return_value=workflow) + monkeypatch.setattr( + workflow_module.variable_factory, + "build_environment_variable_from_mapping", + build_environment_variable, + ) + monkeypatch.setattr( + workflow_module, + "WorkflowService", + lambda: SimpleNamespace(sync_draft_workflow=sync_draft_workflow), + ) - assert status == 400 - assert response["message"] == "Invalid JSON data" - - -def test_sync_draft_workflow_success(app: Flask, monkeypatch: pytest.MonkeyPatch) -> None: - workflow = _make_workflow(updated_at=None, created_at=datetime(2024, 1, 1)) - - monkeypatch.setattr( - workflow_module.variable_factory, "build_environment_variable_from_mapping", lambda *_args: "env" - ) - monkeypatch.setattr( - workflow_module.variable_factory, "build_conversation_variable_from_mapping", lambda *_args: "conv" - ) - - sync_draft_workflow = Mock(return_value=workflow) - service = SimpleNamespace(sync_draft_workflow=sync_draft_workflow) - monkeypatch.setattr(workflow_module, "WorkflowService", lambda: service) - - api = workflow_module.DraftWorkflowApi() - handler = inspect.unwrap(api.post) - - with app.test_request_context( - "/apps/app/workflows/draft", - method="POST", - json={"graph": {}, "features": {}, "hash": "h"}, - ): - response = handler(api, _account(), app_model=_app()) - - assert response["result"] == "success" - assert sync_draft_workflow.call_args.kwargs["environment_variables"] == [] - assert sync_draft_workflow.call_args.kwargs["preserve_environment_variables"] is True - - -@pytest.mark.parametrize("content_type", ["application/json", "text/plain"]) -@pytest.mark.parametrize("confirmation", [{}, {"force": False}], ids=["missing-force", "force-false"]) -def test_sync_draft_workflow_rejects_empty_graph_without_confirmation( - app: Flask, - monkeypatch: pytest.MonkeyPatch, - content_type: str, - confirmation: dict[str, bool], -) -> None: - workflow_service = Mock() - monkeypatch.setattr(workflow_module, "WorkflowService", workflow_service) - api = workflow_module.DraftWorkflowApi() - handler = inspect.unwrap(api.post) - - with app.test_request_context( - "/apps/app/workflows/draft", - method="POST", - data=json.dumps({"graph": {"nodes": [], "edges": []}, "features": {}, "hash": "h", **confirmation}), - content_type=content_type, - ): - with pytest.raises(InvalidArgumentError) as exc: - handler(api, _account(), app_model=_app()) - - assert exc.value.code == 400 - assert exc.value.description == "Saving an empty workflow requires force=true." - workflow_service.assert_not_called() - - -@pytest.mark.parametrize("content_type", ["application/json", "text/plain"]) -@pytest.mark.parametrize( - ("graph", "confirmation"), - [ - ({"nodes": [], "edges": []}, {"force": True}), - ({"nodes": [{"id": "start"}], "edges": []}, {}), - ({"nodes": [], "edges": [{"id": "edge-1", "source": "start", "target": "end"}]}, {}), - ], - ids=["confirmed-empty-graph", "nodes-without-edges", "edges-without-nodes"], -) -def test_sync_draft_workflow_saves_confirmed_empty_or_nonempty_graph( - app: Flask, - monkeypatch: pytest.MonkeyPatch, - content_type: str, - graph: dict[str, list[dict[str, str]]], - confirmation: dict[str, bool], -) -> None: - workflow = SimpleNamespace(unique_hash="next-hash", updated_at=None, created_at=datetime(2024, 1, 1)) - sync_draft_workflow = Mock(return_value=workflow) - monkeypatch.setattr( - workflow_module, - "WorkflowService", - lambda: SimpleNamespace(sync_draft_workflow=sync_draft_workflow), - ) - api = workflow_module.DraftWorkflowApi() - handler = inspect.unwrap(api.post) - - with app.test_request_context( - "/apps/app/workflows/draft", - method="POST", - data=json.dumps({"graph": graph, "features": {}, "hash": "current-hash", **confirmation}), - content_type=content_type, - ): - response = handler(api, _account(), app_model=_app()) - - assert response["result"] == "success" - assert response["hash"] == "next-hash" - sync_draft_workflow.assert_called_once() - assert sync_draft_workflow.call_args.kwargs["graph"] == graph - assert sync_draft_workflow.call_args.kwargs["unique_hash"] == "current-hash" - - -@pytest.mark.parametrize("content_type", ["application/json", "text/plain"]) -def test_sync_draft_workflow_force_does_not_bypass_hash_conflict( - app: Flask, - monkeypatch: pytest.MonkeyPatch, - content_type: str, -) -> None: - sync_draft_workflow = Mock(side_effect=workflow_module.WorkflowHashNotEqualError()) - monkeypatch.setattr( - workflow_module, - "WorkflowService", - lambda: SimpleNamespace(sync_draft_workflow=sync_draft_workflow), - ) - api = workflow_module.DraftWorkflowApi() - handler = inspect.unwrap(api.post) - - with app.test_request_context( - "/apps/app/workflows/draft", - method="POST", - data=json.dumps({"graph": {"nodes": [], "edges": []}, "features": {}, "hash": "stale-hash", "force": True}), - content_type=content_type, - ): - with pytest.raises(DraftWorkflowNotSync) as exc: - handler(api, _account(), app_model=_app()) - - assert exc.value.code == 409 - sync_draft_workflow.assert_called_once() - assert sync_draft_workflow.call_args.kwargs["unique_hash"] == "stale-hash" - - -@pytest.mark.parametrize("content_type", ["application/json", "text/plain"]) -@pytest.mark.parametrize("force", ["true", 1, None], ids=["string", "integer", "null"]) -def test_sync_draft_workflow_rejects_non_boolean_force( - app: Flask, - monkeypatch: pytest.MonkeyPatch, - content_type: str, - force: str | int | None, -) -> None: - workflow_service = Mock() - monkeypatch.setattr(workflow_module, "WorkflowService", workflow_service) - api = workflow_module.DraftWorkflowApi() - handler = inspect.unwrap(api.post) - - with app.test_request_context( - "/apps/app/workflows/draft", - method="POST", - data=json.dumps({"graph": {"nodes": [], "edges": []}, "features": {}, "force": force}), - content_type=content_type, - ): - if content_type == "application/json": - with pytest.raises(ValidationError, match="Input should be a valid boolean"): - handler(api, _account(), app_model=_app()) - else: - response, status = raised_error(lambda: handler(api, _account(), app_model=_app())) - assert status == 400 - assert response["message"] == "Invalid JSON data" - - workflow_service.assert_not_called() - - -def test_sync_draft_workflow_passes_environment_patch(app: Flask, monkeypatch: pytest.MonkeyPatch) -> None: - workflow = _make_workflow(updated_at=None, created_at=datetime(2024, 1, 1)) - patched_variable = StringVariable( - id="env-model", - name="shared_model", - value="model", - selector=["env", "shared_model"], - ) - build_environment_variable = Mock(return_value=patched_variable) - sync_draft_workflow = Mock(return_value=workflow) - monkeypatch.setattr( - workflow_module.variable_factory, - "build_environment_variable_from_mapping", - build_environment_variable, - ) - monkeypatch.setattr( - workflow_module, - "WorkflowService", - lambda: SimpleNamespace(sync_draft_workflow=sync_draft_workflow), - ) - - api = workflow_module.DraftWorkflowApi() - handler = inspect.unwrap(api.post) - with app.test_request_context( - "/apps/app/workflows/draft", - method="POST", - json={ - "graph": {}, - "features": {}, - "hash": "current-hash", - "environment_variable_patch": { - "environment_variables": [ - { - "id": "env-model", - "name": "shared_model", - "value": "model", - "value_type": "string", - } - ], - "deleted_environment_variable_ids": ["env-old"], - }, - }, - ): - response = handler(api, _account(), app_model=_app()) - - assert response["result"] == "success" - assert sync_draft_workflow.call_args.kwargs["preserve_environment_variables"] is True - assert sync_draft_workflow.call_args.kwargs["environment_variable_upserts"] == [patched_variable] - assert sync_draft_workflow.call_args.kwargs["deleted_environment_variable_ids"] == ["env-old"] - build_environment_variable.assert_called_once() - - -def test_sync_draft_workflow_rejects_overlapping_environment_patch_ids() -> None: - with pytest.raises(ValidationError, match="cannot be upserted and deleted"): - workflow_module.SyncDraftWorkflowPayload.model_validate( - { + api = workflow_module.DraftWorkflowApi() + handler = inspect.unwrap(api.post) + with app.test_request_context( + "/apps/app/workflows/draft", + method="POST", + json={ "graph": {}, "features": {}, + "hash": "current-hash", "environment_variable_patch": { - "environment_variables": [{"id": "env-model"}], - "deleted_environment_variable_ids": ["env-model"], + "environment_variables": [ + { + "id": "env-model", + "name": "shared_model", + "value": "model", + "value_type": "string", + } + ], + "deleted_environment_variable_ids": ["env-old"], }, - } - ) + }, + ): + response = handler(api, _account(), app_model=_workflow_app()) + assert response["result"] == "success" + assert sync_draft_workflow.call_args.kwargs["preserve_environment_variables"] is True + assert sync_draft_workflow.call_args.kwargs["environment_variable_upserts"] == [patched_variable] + assert sync_draft_workflow.call_args.kwargs["deleted_environment_variable_ids"] == ["env-old"] + build_environment_variable.assert_called_once() -def test_sync_draft_workflow_rejects_legacy_environment_variables() -> None: - with pytest.raises(ValidationError, match="Extra inputs are not permitted"): - workflow_module.SyncDraftWorkflowPayload.model_validate( - { - "graph": {}, - "features": {}, - "environment_variables": [], - } - ) - - -def test_sync_draft_workflow_hash_mismatch(app: Flask, monkeypatch: pytest.MonkeyPatch) -> None: - def _raise(*_args, **_kwargs): - raise workflow_module.WorkflowHashNotEqualError() - - service = SimpleNamespace(sync_draft_workflow=_raise) - monkeypatch.setattr(workflow_module, "WorkflowService", lambda: service) - - api = workflow_module.DraftWorkflowApi() - handler = inspect.unwrap(api.post) - - with app.test_request_context( - "/apps/app/workflows/draft", - method="POST", - json={"graph": {}, "features": {}, "hash": "h"}, - ): - with pytest.raises(DraftWorkflowNotSync): - handler(api, _account(), app_model=_app()) - - -def test_sync_draft_workflow_variable_validation_error(app: Flask, monkeypatch: pytest.MonkeyPatch) -> None: - def _raise(*_args, **_kwargs): - raise workflow_module.VariableError("description too long") - - monkeypatch.setattr(workflow_module.variable_factory, "build_conversation_variable_from_mapping", _raise) - monkeypatch.setattr( - workflow_module, "WorkflowService", lambda: SimpleNamespace(sync_draft_workflow=lambda **_kwargs: None) - ) - - api = workflow_module.DraftWorkflowApi() - handler = inspect.unwrap(api.post) - - with app.test_request_context( - "/apps/app/workflows/draft", - method="POST", - json={"graph": {}, "features": {}, "hash": "h", "conversation_variables": [{"name": "topic"}]}, - ): - with pytest.raises(InvalidArgumentError) as exc: - handler(api, _account(), app_model=_app()) - - assert exc.value.description == "description too long" - - -def test_restore_published_workflow_to_draft_success(app: Flask, monkeypatch: pytest.MonkeyPatch) -> None: - workflow = _make_workflow(updated_at=None, created_at=datetime(2024, 1, 1)) - - monkeypatch.setattr( - workflow_module, - "WorkflowService", - lambda: SimpleNamespace(restore_published_workflow_to_draft=lambda **_kwargs: workflow), - ) - - api = workflow_module.DraftWorkflowRestoreApi() - handler = inspect.unwrap(api.post) - - with app.test_request_context( - "/apps/app/workflows/published-workflow/restore", - method="POST", - ): - response = handler( - api, - _account(), - app_model=_app(tenant_id="tenant-1"), - workflow_id="published-workflow", - ) - - assert response["result"] == "success" - assert response["hash"] == workflow.unique_hash - - -def test_restore_published_workflow_to_draft_not_found(app: Flask, monkeypatch: pytest.MonkeyPatch) -> None: - monkeypatch.setattr( - workflow_module, - "WorkflowService", - lambda: SimpleNamespace( - restore_published_workflow_to_draft=lambda **_kwargs: (_ for _ in ()).throw( - workflow_module.WorkflowNotFoundError("Workflow not found") + @staticmethod + def test_sync_draft_workflow_rejects_overlapping_environment_patch_ids() -> None: + with pytest.raises(ValidationError, match="cannot be upserted and deleted"): + workflow_module.SyncDraftWorkflowPayload.model_validate( + { + "graph": {}, + "features": {}, + "environment_variable_patch": { + "environment_variables": [{"id": "env-model"}], + "deleted_environment_variable_ids": ["env-model"], + }, + } ) - ), - ) - api = workflow_module.DraftWorkflowRestoreApi() - handler = inspect.unwrap(api.post) + @staticmethod + def test_sync_draft_workflow_rejects_legacy_environment_variables() -> None: + with pytest.raises(ValidationError, match="Extra inputs are not permitted"): + workflow_module.SyncDraftWorkflowPayload.model_validate( + { + "graph": {}, + "features": {}, + "environment_variables": [], + } + ) - with app.test_request_context( - "/apps/app/workflows/published-workflow/restore", - method="POST", - ): - with pytest.raises(NotFound): - handler( + @staticmethod + def test_sync_draft_workflow_hash_mismatch(app: Flask, monkeypatch: pytest.MonkeyPatch) -> None: + def _raise(*_args, **_kwargs): + raise workflow_module.WorkflowHashNotEqualError() + + service = SimpleNamespace(sync_draft_workflow=_raise) + monkeypatch.setattr(workflow_module, "WorkflowService", lambda: service) + + api = workflow_module.DraftWorkflowApi() + handler = inspect.unwrap(api.post) + + with app.test_request_context( + "/apps/app/workflows/draft", + method="POST", + json={"graph": {}, "features": {}, "hash": "h"}, + ): + with pytest.raises(DraftWorkflowNotSync): + handler(api, _account(), app_model=_workflow_app()) + + @staticmethod + def test_sync_draft_workflow_variable_validation_error(app: Flask, monkeypatch: pytest.MonkeyPatch) -> None: + def _raise(*_args, **_kwargs): + raise workflow_module.VariableError("description too long") + + monkeypatch.setattr(workflow_module.variable_factory, "build_conversation_variable_from_mapping", _raise) + monkeypatch.setattr( + workflow_module, "WorkflowService", lambda: SimpleNamespace(sync_draft_workflow=lambda **_kwargs: None) + ) + + api = workflow_module.DraftWorkflowApi() + handler = inspect.unwrap(api.post) + + with app.test_request_context( + "/apps/app/workflows/draft", + method="POST", + json={"graph": {}, "features": {}, "hash": "h", "conversation_variables": [{"name": "topic"}]}, + ): + with pytest.raises(InvalidArgumentError) as exc: + handler(api, _account(), app_model=_workflow_app()) + + assert exc.value.description == "description too long" + + @staticmethod + def test_restore_published_workflow_to_draft_success(app: Flask, monkeypatch: pytest.MonkeyPatch) -> None: + workflow = _make_workflow(updated_at=None, created_at=datetime(2024, 1, 1)) + + monkeypatch.setattr( + workflow_module, + "WorkflowService", + lambda: SimpleNamespace(restore_published_workflow_to_draft=lambda **_kwargs: workflow), + ) + + api = workflow_module.DraftWorkflowRestoreApi() + handler = inspect.unwrap(api.post) + + with app.test_request_context( + "/apps/app/workflows/published-workflow/restore", + method="POST", + ): + response = handler( api, _account(), - app_model=_app(tenant_id="tenant-1"), + app_model=_workflow_app(tenant_id="tenant-1"), workflow_id="published-workflow", ) + assert response["result"] == "success" + assert response["hash"] == workflow.unique_hash -def test_restore_published_workflow_to_draft_returns_400_for_draft_source( - app: Flask, monkeypatch: pytest.MonkeyPatch -) -> None: - monkeypatch.setattr( - workflow_module, - "WorkflowService", - lambda: SimpleNamespace( - restore_published_workflow_to_draft=lambda **_kwargs: (_ for _ in ()).throw( - workflow_module.IsDraftWorkflowError( - "Cannot use draft workflow version. Workflow ID: draft-workflow. " - "Please use a published workflow version or leave workflow_id empty." + @staticmethod + def test_restore_published_workflow_to_draft_not_found(app: Flask, monkeypatch: pytest.MonkeyPatch) -> None: + monkeypatch.setattr( + workflow_module, + "WorkflowService", + lambda: SimpleNamespace( + restore_published_workflow_to_draft=lambda **_kwargs: (_ for _ in ()).throw( + workflow_module.WorkflowNotFoundError("Workflow not found") ) - ) - ), - ) - - api = workflow_module.DraftWorkflowRestoreApi() - handler = inspect.unwrap(api.post) - - with app.test_request_context( - "/apps/app/workflows/draft-workflow/restore", - method="POST", - ): - with pytest.raises(HTTPException) as exc: - handler( - api, - _account(), - app_model=_app(tenant_id="tenant-1"), - workflow_id="draft-workflow", - ) - - assert exc.value.code == 400 - assert exc.value.description == workflow_module.RESTORE_SOURCE_WORKFLOW_MUST_BE_PUBLISHED_MESSAGE - - -def test_restore_published_workflow_to_draft_returns_400_for_invalid_structure( - app: Flask, monkeypatch: pytest.MonkeyPatch -) -> None: - monkeypatch.setattr( - workflow_module, - "WorkflowService", - lambda: SimpleNamespace( - restore_published_workflow_to_draft=lambda **_kwargs: (_ for _ in ()).throw( - ValueError("invalid workflow graph") - ) - ), - ) - - api = workflow_module.DraftWorkflowRestoreApi() - handler = inspect.unwrap(api.post) - - with app.test_request_context( - "/apps/app/workflows/published-workflow/restore", - method="POST", - ): - with pytest.raises(HTTPException) as exc: - handler( - api, - _account(), - app_model=_app(tenant_id="tenant-1"), - workflow_id="published-workflow", - ) - - assert exc.value.code == 400 - assert exc.value.description == "invalid workflow graph" - - -def test_get_published_workflows_uses_the_request_session( - app: Flask, monkeypatch: pytest.MonkeyPatch, sqlite_session: Session -) -> None: - api = workflow_module.PublishedAllWorkflowApi() - handler = inspect.unwrap(api.get) - workflow = _make_workflow() - sqlite_session.add_all([_account(), workflow]) - sqlite_session.commit() - - def get_all_published_workflow(*, session: Session, **_kwargs): - assert session is sqlite_session - return [workflow], False - - monkeypatch.setattr( - workflow_module, - "WorkflowService", - lambda: SimpleNamespace(get_all_published_workflow=get_all_published_workflow), - ) - - with app.test_request_context( - "/apps/app/workflows", - method="GET", - query_string={"page": 1, "limit": 10, "user_id": "", "named_only": "false"}, - ): - query = workflow_module.WorkflowListQuery.model_validate( - {"page": "1", "limit": "10", "user_id": "", "named_only": "false"} + ), ) - response = handler(api, query, sqlite_session, _account(), app_model=_app()) - assert response["items"][0]["id"] == "workflow-1" - assert response["page"] == 1 - assert response["limit"] == 10 - assert response["has_more"] is False + api = workflow_module.DraftWorkflowRestoreApi() + handler = inspect.unwrap(api.post) + with app.test_request_context( + "/apps/app/workflows/published-workflow/restore", + method="POST", + ): + with pytest.raises(NotFound): + handler( + api, + _account(), + app_model=_workflow_app(tenant_id="tenant-1"), + workflow_id="published-workflow", + ) -def test_draft_workflow_get_serializes_response_model(monkeypatch: pytest.MonkeyPatch, sqlite_session: Session) -> None: - workflow = _make_workflow() - sqlite_session.add_all([_account(), workflow]) - sqlite_session.commit() - monkeypatch.setattr( - workflow_module, "WorkflowService", lambda: SimpleNamespace(get_draft_workflow=lambda **_kwargs: workflow) - ) + @staticmethod + def test_restore_published_workflow_to_draft_returns_400_for_draft_source( + app: Flask, monkeypatch: pytest.MonkeyPatch + ) -> None: + monkeypatch.setattr( + workflow_module, + "WorkflowService", + lambda: SimpleNamespace( + restore_published_workflow_to_draft=lambda **_kwargs: (_ for _ in ()).throw( + workflow_module.IsDraftWorkflowError( + "Cannot use draft workflow version. Workflow ID: draft-workflow. " + "Please use a published workflow version or leave workflow_id empty." + ) + ) + ), + ) - api = workflow_module.DraftWorkflowApi() - handler = inspect.unwrap(api.get) + api = workflow_module.DraftWorkflowRestoreApi() + handler = inspect.unwrap(api.post) - response = handler(api, sqlite_session, app_model=_app()) + with app.test_request_context( + "/apps/app/workflows/draft-workflow/restore", + method="POST", + ): + with pytest.raises(HTTPException) as exc: + handler( + api, + _account(), + app_model=_workflow_app(tenant_id="tenant-1"), + workflow_id="draft-workflow", + ) - assert response["id"] == "workflow-1" - assert response["graph"] == {"nodes": [], "edges": []} - assert response["features"] == {"file_upload": {"enabled": False}} - assert response["hash"] == workflow.unique_hash - assert response["created_by"] == {"id": "user-1", "name": "Alice", "email": "alice@example.com"} - assert response["updated_by"] is None - assert response["created_at"] == int(datetime(2024, 1, 1, 12, 0, 0).timestamp()) - assert response["updated_at"] == int(datetime(2024, 1, 1, 12, 1, 0).timestamp()) - assert response["environment_variables"] == [ - { - "id": "env-1", - "name": "API_KEY", - "value": workflow_module.encrypter.full_mask_token(), - "value_type": "secret", - "description": "API key", - } - ] - assert response["conversation_variables"] == [ - { - "id": "conv-1", - "name": "topic", - "value": "hello", - "value_type": "string", - "description": "Topic", - } - ] - assert response["rag_pipeline_variables"] == [ - { - "label": "Query", - "variable": "query", - "type": "text-input", - "belong_to_node_id": "shared", - "max_length": 0, - "required": False, - "unit": "", - "default_value": "", - "options": [], - "placeholder": "", - "tooltips": "", - "allowed_file_types": ["custom"], - "allowed_file_extensions": [".pdf"], - "allowed_file_upload_methods": ["local_file"], - } - ] + assert exc.value.code == 400 + assert exc.value.description == workflow_module.RESTORE_SOURCE_WORKFLOW_MUST_BE_PUBLISHED_MESSAGE + @staticmethod + def test_restore_published_workflow_to_draft_returns_400_for_invalid_structure( + app: Flask, monkeypatch: pytest.MonkeyPatch + ) -> None: + monkeypatch.setattr( + workflow_module, + "WorkflowService", + lambda: SimpleNamespace( + restore_published_workflow_to_draft=lambda **_kwargs: (_ for _ in ()).throw( + ValueError("invalid workflow graph") + ) + ), + ) -def test_published_workflow_get_uses_session_aware_response_source( - monkeypatch: pytest.MonkeyPatch, sqlite_session: Session -) -> None: - sqlite_session.add(_account()) - sqlite_session.commit() - workflow = _make_workflow() - monkeypatch.setattr(workflow_module.db, "session", lambda: sqlite_session) - monkeypatch.setattr( - workflow_module, "WorkflowService", lambda: SimpleNamespace(get_published_workflow=lambda **_kwargs: workflow) - ) + api = workflow_module.DraftWorkflowRestoreApi() + handler = inspect.unwrap(api.post) - api = workflow_module.PublishedWorkflowApi() - handler = inspect.unwrap(api.get) + with app.test_request_context( + "/apps/app/workflows/published-workflow/restore", + method="POST", + ): + with pytest.raises(HTTPException) as exc: + handler( + api, + _account(), + app_model=_workflow_app(tenant_id="tenant-1"), + workflow_id="published-workflow", + ) - response = handler(api, app_model=SimpleNamespace(id="app")) + assert exc.value.code == 400 + assert exc.value.description == "invalid workflow graph" - assert response["id"] == "workflow-1" - assert response["created_by"] == {"id": "user-1", "name": "Alice", "email": "alice@example.com"} - assert response["updated_by"] is None - assert response["tool_published"] is False + @staticmethod + def test_get_published_workflows_uses_the_request_session( + app: Flask, monkeypatch: pytest.MonkeyPatch, sqlite_session: Session + ) -> None: + api = workflow_module.PublishedAllWorkflowApi() + handler = inspect.unwrap(api.get) + workflow = _make_workflow() + sqlite_session.add_all([_account(), workflow]) + sqlite_session.commit() + def get_all_published_workflow(*, session: Session, **_kwargs): + assert session is sqlite_session + return [workflow], False -def test_pipeline_variable_response_accepts_legacy_file_field_names() -> None: - response = workflow_module.PipelineVariableResponse.model_validate( - { - "label": "Query", - "variable": "query", - "type": "single-file", - "belong_to_node_id": "shared", - "max_length": 0, - "required": False, - "unit": "", - "default_value": "", - "options": [], - "placeholder": "", - "tooltips": "", - "allowed_file_types": [], - "allow_file_extension": [".txt"], - "allow_file_upload_methods": ["remote_url"], - } - ).model_dump(mode="json") + monkeypatch.setattr( + workflow_module, + "WorkflowService", + lambda: SimpleNamespace(get_all_published_workflow=get_all_published_workflow), + ) - assert response["allowed_file_extensions"] == [".txt"] - assert response["allowed_file_upload_methods"] == ["remote_url"] - - -def test_pipeline_variable_response_accepts_explicit_null_optional_fields() -> None: - pipeline_variable = RAGPipelineVariable.model_validate( - { - "label": "Query", - "variable": "query", - "type": "text-input", - "belong_to_node_id": "shared", - "max_length": None, - "unit": None, - "default_value": None, - "options": None, - "placeholder": None, - "tooltips": None, - "allowed_file_types": None, - "allowed_file_extensions": None, - "allowed_file_upload_methods": None, - } - ).model_dump(mode="json") - - response = workflow_module.PipelineVariableResponse.model_validate(pipeline_variable).model_dump(mode="json") - - assert response["max_length"] is None - assert response["allowed_file_types"] is None - assert response["allowed_file_extensions"] is None - assert response["allowed_file_upload_methods"] is None - - -def test_workflow_response_masks_secret_environment_variables(sqlite_session: Session) -> None: - workflow = _make_workflow( - environment_variables=[ - SecretVariable(id="env-secret", name="API_KEY", value="plain-token", selector=["env", "API_KEY"]), - StringVariable(id="env-string", name="REGION", value="us-east-1", selector=["env", "REGION"]), - ] - ) - sqlite_session.add_all([_account(), workflow]) - sqlite_session.commit() - - response = workflow_module.WorkflowResponse.model_validate( - workflow_module.WorkflowResponseSource(workflow, session=sqlite_session), - from_attributes=True, - ).model_dump(mode="json") - - assert response["environment_variables"] == [ - { - "id": "env-secret", - "name": "API_KEY", - "value": workflow_module.encrypter.full_mask_token(), - "value_type": "secret", - "description": "", - }, - { - "id": "env-string", - "name": "REGION", - "value": "us-east-1", - "value_type": "string", - "description": "", - }, - ] - - -def test_workflow_response_preserves_llm_environment_variable_type(sqlite_session: Session) -> None: - workflow = _make_workflow( - environment_variables=[ - LLMEnvironmentVariable( - id="env-llm", - name="for_summarize", - value={"provider": "provider", "name": "model", "mode": "chat"}, - selector=["env", "for_summarize"], + with app.test_request_context( + "/apps/app/workflows", + method="GET", + query_string={"page": 1, "limit": 10, "user_id": "", "named_only": "false"}, + ): + query = workflow_module.WorkflowListQuery.model_validate( + {"page": "1", "limit": "10", "user_id": "", "named_only": "false"} ) + response = handler(api, query, sqlite_session, _account(), app_model=_workflow_app()) + + assert response["items"][0]["id"] == "workflow-1" + assert response["page"] == 1 + assert response["limit"] == 10 + assert response["has_more"] is False + + @staticmethod + def test_draft_workflow_get_serializes_response_model( + monkeypatch: pytest.MonkeyPatch, sqlite_session: Session + ) -> None: + workflow = _make_workflow() + sqlite_session.add_all([_account(), workflow]) + sqlite_session.commit() + monkeypatch.setattr( + workflow_module, "WorkflowService", lambda: SimpleNamespace(get_draft_workflow=lambda **_kwargs: workflow) + ) + + api = workflow_module.DraftWorkflowApi() + handler = inspect.unwrap(api.get) + + response = handler(api, sqlite_session, app_model=_workflow_app()) + + assert response["id"] == "workflow-1" + assert response["graph"] == {"nodes": [], "edges": []} + assert response["features"] == {"file_upload": {"enabled": False}} + assert response["hash"] == workflow.unique_hash + assert response["created_by"] == {"id": "user-1", "name": "Alice", "email": "alice@example.com"} + assert response["updated_by"] is None + assert response["created_at"] == int(datetime(2024, 1, 1, 12, 0, 0).timestamp()) + assert response["updated_at"] == int(datetime(2024, 1, 1, 12, 1, 0).timestamp()) + assert response["environment_variables"] == [ + { + "id": "env-1", + "name": "API_KEY", + "value": workflow_module.encrypter.full_mask_token(), + "value_type": "secret", + "description": "API key", + } + ] + assert response["conversation_variables"] == [ + { + "id": "conv-1", + "name": "topic", + "value": "hello", + "value_type": "string", + "description": "Topic", + } + ] + assert response["rag_pipeline_variables"] == [ + { + "label": "Query", + "variable": "query", + "type": "text-input", + "belong_to_node_id": "shared", + "max_length": 0, + "required": False, + "unit": "", + "default_value": "", + "options": [], + "placeholder": "", + "tooltips": "", + "allowed_file_types": ["custom"], + "allowed_file_extensions": [".pdf"], + "allowed_file_upload_methods": ["local_file"], + } ] - ) - sqlite_session.add_all([_account(), workflow]) - sqlite_session.commit() - response = workflow_module.WorkflowResponse.model_validate( - workflow_module.WorkflowResponseSource(workflow, session=sqlite_session), - from_attributes=True, - ).model_dump(mode="json") + @staticmethod + def test_published_workflow_get_uses_session_aware_response_source( + monkeypatch: pytest.MonkeyPatch, sqlite_session: Session + ) -> None: + sqlite_session.add(_account()) + sqlite_session.commit() + workflow = _make_workflow() + monkeypatch.setattr(workflow_module.db, "session", lambda: sqlite_session) + monkeypatch.setattr( + workflow_module, + "WorkflowService", + lambda: SimpleNamespace(get_published_workflow=lambda **_kwargs: workflow), + ) - assert response["environment_variables"] == [ - { - "id": "env-llm", - "name": "for_summarize", - "value": {"provider": "provider", "name": "model", "mode": "chat"}, - "value_type": "llm", - "description": "", - } - ] + api = workflow_module.PublishedWorkflowApi() + handler = inspect.unwrap(api.get) + response = handler(api, app_model=SimpleNamespace(id="app")) -def test_workflow_response_rejects_invalid_environment_variable_dict(sqlite_session: Session) -> None: - workflow = _make_workflow(environment_variables=[{"value_type": "not-a-segment-type"}]) - sqlite_session.add_all([_account(), workflow]) - sqlite_session.commit() + assert response["id"] == "workflow-1" + assert response["created_by"] == {"id": "user-1", "name": "Alice", "email": "alice@example.com"} + assert response["updated_by"] is None + assert response["tool_published"] is False - with pytest.raises(ValidationError): - workflow_module.WorkflowResponse.model_validate( + @staticmethod + def test_pipeline_variable_response_accepts_legacy_file_field_names() -> None: + response = workflow_module.PipelineVariableResponse.model_validate( + { + "label": "Query", + "variable": "query", + "type": "single-file", + "belong_to_node_id": "shared", + "max_length": 0, + "required": False, + "unit": "", + "default_value": "", + "options": [], + "placeholder": "", + "tooltips": "", + "allowed_file_types": [], + "allow_file_extension": [".txt"], + "allow_file_upload_methods": ["remote_url"], + } + ).model_dump(mode="json") + + assert response["allowed_file_extensions"] == [".txt"] + assert response["allowed_file_upload_methods"] == ["remote_url"] + + @staticmethod + def test_pipeline_variable_response_accepts_explicit_null_optional_fields() -> None: + pipeline_variable = RAGPipelineVariable.model_validate( + { + "label": "Query", + "variable": "query", + "type": "text-input", + "belong_to_node_id": "shared", + "max_length": None, + "unit": None, + "default_value": None, + "options": None, + "placeholder": None, + "tooltips": None, + "allowed_file_types": None, + "allowed_file_extensions": None, + "allowed_file_upload_methods": None, + } + ).model_dump(mode="json") + + response = workflow_module.PipelineVariableResponse.model_validate(pipeline_variable).model_dump(mode="json") + + assert response["max_length"] is None + assert response["allowed_file_types"] is None + assert response["allowed_file_extensions"] is None + assert response["allowed_file_upload_methods"] is None + + @staticmethod + def test_workflow_response_masks_secret_environment_variables(sqlite_session: Session) -> None: + workflow = _make_workflow( + environment_variables=[ + SecretVariable(id="env-secret", name="API_KEY", value="plain-token", selector=["env", "API_KEY"]), + StringVariable(id="env-string", name="REGION", value="us-east-1", selector=["env", "REGION"]), + ] + ) + sqlite_session.add_all([_account(), workflow]) + sqlite_session.commit() + + response = workflow_module.WorkflowResponse.model_validate( workflow_module.WorkflowResponseSource(workflow, session=sqlite_session), from_attributes=True, + ).model_dump(mode="json") + + assert response["environment_variables"] == [ + { + "id": "env-secret", + "name": "API_KEY", + "value": workflow_module.encrypter.full_mask_token(), + "value_type": "secret", + "description": "", + }, + { + "id": "env-string", + "name": "REGION", + "value": "us-east-1", + "value_type": "string", + "description": "", + }, + ] + + @staticmethod + def test_workflow_response_preserves_llm_environment_variable_type(sqlite_session: Session) -> None: + workflow = _make_workflow( + environment_variables=[ + LLMEnvironmentVariable( + id="env-llm", + name="for_summarize", + value={"provider": "provider", "name": "model", "mode": "chat"}, + selector=["env", "for_summarize"], + ) + ] + ) + sqlite_session.add_all([_account(), workflow]) + sqlite_session.commit() + + response = workflow_module.WorkflowResponse.model_validate( + workflow_module.WorkflowResponseSource(workflow, session=sqlite_session), + from_attributes=True, + ).model_dump(mode="json") + + assert response["environment_variables"] == [ + { + "id": "env-llm", + "name": "for_summarize", + "value": {"provider": "provider", "name": "model", "mode": "chat"}, + "value_type": "llm", + "description": "", + } + ] + + @staticmethod + def test_workflow_response_rejects_invalid_environment_variable_dict(sqlite_session: Session) -> None: + workflow = _make_workflow(environment_variables=[{"value_type": "not-a-segment-type"}]) + sqlite_session.add_all([_account(), workflow]) + sqlite_session.commit() + + with pytest.raises(ValidationError): + workflow_module.WorkflowResponse.model_validate( + workflow_module.WorkflowResponseSource(workflow, session=sqlite_session), + from_attributes=True, + ) + + @staticmethod + def test_draft_workflow_get_not_found(monkeypatch: pytest.MonkeyPatch, unbound_session: Session) -> None: + monkeypatch.setattr( + workflow_module, "WorkflowService", lambda: SimpleNamespace(get_draft_workflow=lambda **_k: None) ) + api = workflow_module.DraftWorkflowApi() + handler = inspect.unwrap(api.get) -def test_draft_workflow_get_not_found(monkeypatch: pytest.MonkeyPatch, unbound_session: Session) -> None: - monkeypatch.setattr( - workflow_module, "WorkflowService", lambda: SimpleNamespace(get_draft_workflow=lambda **_k: None) - ) + with pytest.raises(DraftWorkflowNotExist): + handler(api, unbound_session, app_model=_workflow_app()) - api = workflow_module.DraftWorkflowApi() - handler = inspect.unwrap(api.get) - - with pytest.raises(DraftWorkflowNotExist): - handler(api, unbound_session, app_model=_app()) - - -def test_draft_workflow_get_projects_agent_node_job_to_graph( - monkeypatch: pytest.MonkeyPatch, sqlite_session: Session -) -> None: - workflow = _make_workflow( - graph_dict={ + @staticmethod + def test_draft_workflow_get_projects_agent_node_job_to_graph( + monkeypatch: pytest.MonkeyPatch, sqlite_session: Session + ) -> None: + workflow = _make_workflow( + graph_dict={ + "nodes": [ + { + "id": "agent-node", + "data": { + "type": "agent", + "version": "2", + "agent_node_kind": "dify_agent", + }, + } + ], + "edges": [], + } + ) + sqlite_session.add_all([_account(), workflow]) + sqlite_session.commit() + projected_graph = { "nodes": [ { "id": "agent-node", @@ -989,300 +1030,606 @@ def test_draft_workflow_get_projects_agent_node_job_to_graph( "type": "agent", "version": "2", "agent_node_kind": "dify_agent", + "agent_task": "Summarize it.", + "agent_declared_outputs": [{"name": "summary", "type": "string"}], }, } ], "edges": [], } - ) - sqlite_session.add_all([_account(), workflow]) - sqlite_session.commit() - projected_graph = { - "nodes": [ - { - "id": "agent-node", - "data": { - "type": "agent", - "version": "2", - "agent_node_kind": "dify_agent", - "agent_task": "Summarize it.", - "agent_declared_outputs": [{"name": "summary", "type": "string"}], - }, - } + + monkeypatch.setattr( + workflow_module, + "WorkflowService", + lambda: SimpleNamespace(get_draft_workflow=lambda **_k: workflow), + ) + + from services.agent.workflow_publish_service import WorkflowAgentPublishService + + monkeypatch.setattr( + WorkflowAgentPublishService, + "project_draft_bindings_to_graph", + lambda **_k: projected_graph, + ) + + api = workflow_module.DraftWorkflowApi() + handler = inspect.unwrap(api.get) + + response = handler(api, sqlite_session, app_model=_workflow_app()) + + assert response["graph"] == projected_graph + + @staticmethod + def test_advanced_chat_run_conversation_not_exists( + app: Flask, monkeypatch: pytest.MonkeyPatch, unbound_session: Session + ) -> None: + monkeypatch.setattr( + workflow_module.AppGenerateService, + "generate", + lambda *_args, **_kwargs: (_ for _ in ()).throw( + workflow_module.services.errors.conversation.ConversationNotExistsError() + ), + ) + + api = workflow_module.AdvancedChatDraftWorkflowRunApi() + handler = inspect.unwrap(api.post) + + with app.test_request_context( + "/apps/app/advanced-chat/workflows/draft/run", + method="POST", + json={"inputs": {}}, + ): + payload = workflow_module.AdvancedChatWorkflowRunPayload.model_validate({"inputs": {}}) + with pytest.raises(NotFound): + handler(api, payload, unbound_session, _account(), app_model=_workflow_app()) + + @staticmethod + @pytest.mark.parametrize( + ("resource", "payload_model", "payload"), + [ + ( + workflow_module.DraftWorkflowTriggerRunApi, + workflow_module.DraftWorkflowTriggerRunPayload, + {"node_id": "node-1"}, + ), + ( + workflow_module.DraftWorkflowTriggerRunAllApi, + workflow_module.DraftWorkflowTriggerRunAllPayload, + {"node_ids": ["node-1"]}, + ), ], - "edges": [], - } - - monkeypatch.setattr( - workflow_module, - "WorkflowService", - lambda: SimpleNamespace(get_draft_workflow=lambda **_k: workflow), ) + def test_trigger_run_loads_draft_with_request_session( + app: Flask, + monkeypatch: pytest.MonkeyPatch, + unbound_session: Session, + resource: type, + payload_model: type, + payload: dict[str, object], + ) -> None: + get_draft_workflow = Mock(return_value=None) + monkeypatch.setattr( + workflow_module, + "WorkflowService", + lambda: SimpleNamespace(get_draft_workflow=get_draft_workflow), + ) + session = unbound_session + app_model = _workflow_app(app_id="app-1") + handler = inspect.unwrap(resource.post) - from services.agent.workflow_publish_service import WorkflowAgentPublishService + with app.test_request_context("/", method="POST", json=payload): + with pytest.raises(ValueError, match="Workflow not found"): + handler( + resource(), + payload_model.model_validate(payload), + session, + _account(), + app_model, + ) - monkeypatch.setattr( - WorkflowAgentPublishService, - "project_draft_bindings_to_graph", - lambda **_k: projected_graph, - ) + get_draft_workflow.assert_called_once_with(app_model, session=session) - api = workflow_module.DraftWorkflowApi() - handler = inspect.unwrap(api.get) + @staticmethod + def test_workflow_online_users_filters_inaccessible_workflow( + app: Flask, monkeypatch: pytest.MonkeyPatch, unbound_session: Session + ) -> None: + app_id_1 = "11111111-1111-1111-1111-111111111111" + app_id_2 = "22222222-2222-2222-2222-222222222222" + signed_avatar_url = "https://files.example.com/signed/avatar-1" + signed_ids: list[str] = [] + queries: list[tuple[list[str], str, Session]] = [] + actors: list[tuple[str, str]] = [] - response = handler(api, sqlite_session, app_model=_app()) + def sign_avatar(file_id: str) -> str: + signed_ids.append(file_id) + return signed_avatar_url - assert response["graph"] == projected_graph + class WorkflowQueries: + def get_tenant_app_maintainers( + self, app_ids: list[str], tenant_id: str, *, session: Session + ) -> dict[str, str]: + queries.append((app_ids, tenant_id, session)) + return {app_id_1: "owner-1", app_id_2: "owner-2"} + def resolve_access(tenant_id: str, account_id: str) -> AppAccessFilter: + actors.append((tenant_id, account_id)) + return AppAccessFilter(accessible_app_ids={app_id_1}, can_manage_own_apps=False) -def test_advanced_chat_run_conversation_not_exists( - app: Flask, monkeypatch: pytest.MonkeyPatch, unbound_session: Session -) -> None: - monkeypatch.setattr( - workflow_module.AppGenerateService, - "generate", - lambda *_args, **_kwargs: (_ for _ in ()).throw( - workflow_module.services.errors.conversation.ConversationNotExistsError() - ), - ) + monkeypatch.setattr(workflow_module, "WorkflowService", WorkflowQueries) + monkeypatch.setattr(workflow_module, "resolve_app_access_filter", resolve_access) + apply_config_overrides(monkeypatch, RBAC_ENABLED=True) + monkeypatch.setattr(workflow_module.file_helpers, "get_signed_file_url", sign_avatar) + monkeypatch.setattr(workflow_module.session_factory, "create_session", lambda: nullcontext(unbound_session)) - api = workflow_module.AdvancedChatDraftWorkflowRunApi() - handler = inspect.unwrap(api.post) - - with app.test_request_context( - "/apps/app/advanced-chat/workflows/draft/run", - method="POST", - json={"inputs": {}}, - ): - payload = workflow_module.AdvancedChatWorkflowRunPayload.model_validate({"inputs": {}}) - with pytest.raises(NotFound): - handler(api, payload, unbound_session, _account(), app_model=_app()) - - -@pytest.mark.parametrize( - ("resource", "payload_model", "payload"), - [ - ( - workflow_module.DraftWorkflowTriggerRunApi, - workflow_module.DraftWorkflowTriggerRunPayload, - {"node_id": "node-1"}, - ), - ( - workflow_module.DraftWorkflowTriggerRunAllApi, - workflow_module.DraftWorkflowTriggerRunAllPayload, - {"node_ids": ["node-1"]}, - ), - ], -) -def test_trigger_run_loads_draft_with_request_session( - app: Flask, - monkeypatch: pytest.MonkeyPatch, - unbound_session: Session, - resource: type, - payload_model: type, - payload: dict[str, object], -) -> None: - get_draft_workflow = Mock(return_value=None) - monkeypatch.setattr( - workflow_module, - "WorkflowService", - lambda: SimpleNamespace(get_draft_workflow=get_draft_workflow), - ) - session = unbound_session - app_model = _app(app_id="app-1") - handler = inspect.unwrap(resource.post) - - with app.test_request_context("/", method="POST", json=payload): - with pytest.raises(ValueError, match="Workflow not found"): - handler( - resource(), - payload_model.model_validate(payload), - session, - _account(), - app_model, - ) - - get_draft_workflow.assert_called_once_with(app_model, session=session) - - -def test_workflow_online_users_filters_inaccessible_workflow( - app: Flask, monkeypatch: pytest.MonkeyPatch, unbound_session: Session -) -> None: - app_id_1 = "11111111-1111-1111-1111-111111111111" - app_id_2 = "22222222-2222-2222-2222-222222222222" - signed_avatar_url = "https://files.example.com/signed/avatar-1" - signed_ids: list[str] = [] - queries: list[tuple[list[str], str, Session]] = [] - actors: list[tuple[str, str]] = [] - - def sign_avatar(file_id: str) -> str: - signed_ids.append(file_id) - return signed_avatar_url - - class WorkflowQueries: - def get_tenant_app_maintainers(self, app_ids: list[str], tenant_id: str, *, session: Session) -> dict[str, str]: - queries.append((app_ids, tenant_id, session)) - return {app_id_1: "owner-1", app_id_2: "owner-2"} - - def resolve_access(tenant_id: str, account_id: str) -> AppAccessFilter: - actors.append((tenant_id, account_id)) - return AppAccessFilter(accessible_app_ids={app_id_1}, can_manage_own_apps=False) - - monkeypatch.setattr(workflow_module, "WorkflowService", WorkflowQueries) - monkeypatch.setattr(workflow_module, "resolve_app_access_filter", resolve_access) - apply_config_overrides(monkeypatch, RBAC_ENABLED=True) - monkeypatch.setattr(workflow_module.file_helpers, "get_signed_file_url", sign_avatar) - monkeypatch.setattr(workflow_module.session_factory, "create_session", lambda: nullcontext(unbound_session)) - - redis_result = [ - { - b"sid-1": json.dumps( - { - "user_id": "u-1", - "username": "Alice", - "avatar": "avatar-file-id", - "sid": "sid-1", - } - ), - b"sid-malformed": json.dumps({"avatar": "avatar-file-id", "sid": "sid-malformed"}), - b"sid-invalid-avatar": json.dumps( - { - "user_id": "u-2", - "username": "Bob", - "avatar": {"file_id": "avatar-file-id"}, - } - ), - b"sid-invalid-user-id": json.dumps( - { - "user_id": 42, - "username": "Carol", - "avatar": "avatar-file-id", - } - ), - b"sid-invalid-username": json.dumps( - { - "user_id": "u-4", - "username": ["Dave"], - "avatar": "avatar-file-id", - } - ), - } - ] - redis_keys: list[str] = [] - pipeline_transactions: list[bool] = [] - executions: list[bool] = [] - - class PresencePipeline: - def hgetall(self, key: str) -> None: - redis_keys.append(key) - - def execute(self) -> list[dict[bytes, str]]: - executions.append(True) - return redis_result - - def pipeline(*, transaction: bool) -> PresencePipeline: - pipeline_transactions.append(transaction) - return PresencePipeline() - - monkeypatch.setattr(workflow_module.redis_client, "pipeline", pipeline) - - api = workflow_module.WorkflowOnlineUsersApi() - handler = inspect.unwrap(api.post) - - with app.test_request_context( - "/apps/workflows/online-users", - method="POST", - json={"app_ids": [app_id_1, app_id_2]}, - ): - args = workflow_module.WorkflowOnlineUsersPayload.model_validate({"app_ids": [app_id_1, app_id_2]}) - actor = _account() - actor.id = "account-1" - response = handler(api, args, "tenant-1", actor) - - assert response == { - "data": [ + redis_result = [ { - "app_id": app_id_1, - "users": [ + b"sid-1": json.dumps( { "user_id": "u-1", "username": "Alice", - "avatar": signed_avatar_url, - }, + "avatar": "avatar-file-id", + "sid": "sid-1", + } + ), + b"sid-malformed": json.dumps({"avatar": "avatar-file-id", "sid": "sid-malformed"}), + b"sid-invalid-avatar": json.dumps( { "user_id": "u-2", "username": "Bob", - "avatar": None, - }, - ], + "avatar": {"file_id": "avatar-file-id"}, + } + ), + b"sid-invalid-user-id": json.dumps( + { + "user_id": 42, + "username": "Carol", + "avatar": "avatar-file-id", + } + ), + b"sid-invalid-username": json.dumps( + { + "user_id": "u-4", + "username": ["Dave"], + "avatar": "avatar-file-id", + } + ), } ] - } - assert pipeline_transactions == [False] - assert redis_keys == [f"{workflow_module.WORKFLOW_ONLINE_USERS_PREFIX}{app_id_1}"] - assert executions == [True] - assert signed_ids == ["avatar-file-id"] - assert queries == [([app_id_1, app_id_2], "tenant-1", unbound_session)] - assert actors == [("tenant-1", "account-1")] + redis_keys: list[str] = [] + pipeline_transactions: list[bool] = [] + executions: list[bool] = [] + + class PresencePipeline: + def hgetall(self, key: str) -> None: + redis_keys.append(key) + + def execute(self) -> list[dict[bytes, str]]: + executions.append(True) + return redis_result + + def pipeline(*, transaction: bool) -> PresencePipeline: + pipeline_transactions.append(transaction) + return PresencePipeline() + + monkeypatch.setattr(workflow_module.redis_client, "pipeline", pipeline) + + api = workflow_module.WorkflowOnlineUsersApi() + handler = inspect.unwrap(api.post) + + with app.test_request_context( + "/apps/workflows/online-users", + method="POST", + json={"app_ids": [app_id_1, app_id_2]}, + ): + args = workflow_module.WorkflowOnlineUsersPayload.model_validate({"app_ids": [app_id_1, app_id_2]}) + actor = _account() + actor.id = "account-1" + response = handler(api, args, "tenant-1", actor) + + assert response == { + "data": [ + { + "app_id": app_id_1, + "users": [ + { + "user_id": "u-1", + "username": "Alice", + "avatar": signed_avatar_url, + }, + { + "user_id": "u-2", + "username": "Bob", + "avatar": None, + }, + ], + } + ] + } + assert pipeline_transactions == [False] + assert redis_keys == [f"{workflow_module.WORKFLOW_ONLINE_USERS_PREFIX}{app_id_1}"] + assert executions == [True] + assert signed_ids == ["avatar-file-id"] + assert queries == [([app_id_1, app_id_2], "tenant-1", unbound_session)] + assert actors == [("tenant-1", "account-1")] + + @staticmethod + def test_workflow_online_users_batches_redis_reads(app: Flask, monkeypatch: pytest.MonkeyPatch) -> None: + app_ids = [f"wf-{index}" for index in range(workflow_module.WORKFLOW_ONLINE_USERS_REDIS_BATCH_SIZE + 1)] + monkeypatch.setattr( + workflow_module, + "WorkflowService", + lambda: SimpleNamespace( + get_tenant_app_maintainers=lambda app_ids, tenant_id, session: dict.fromkeys(app_ids) + ), + ) + apply_config_overrides(monkeypatch, RBAC_ENABLED=False) + monkeypatch.setattr(workflow_module.session_factory, "create_session", lambda: nullcontext(Mock())) + + first_pipeline = Mock() + first_pipeline.execute.return_value = [ + {} for _ in range(workflow_module.WORKFLOW_ONLINE_USERS_REDIS_BATCH_SIZE) + ] + second_pipeline = Mock() + second_pipeline.execute.return_value = [{}] + redis_pipeline_factory = Mock(side_effect=[first_pipeline, second_pipeline]) + monkeypatch.setattr(workflow_module.redis_client, "pipeline", redis_pipeline_factory) + + api = workflow_module.WorkflowOnlineUsersApi() + handler = inspect.unwrap(api.post) + + with app.test_request_context( + "/apps/workflows/online-users", + method="POST", + json={"app_ids": app_ids}, + ): + args = workflow_module.WorkflowOnlineUsersPayload.model_validate({"app_ids": app_ids}) + response = handler(api, args, "tenant-1", SimpleNamespace(id="account-1")) + + assert len(response["data"]) == len(app_ids) + assert redis_pipeline_factory.call_count == 2 + assert first_pipeline.hgetall.call_count == workflow_module.WORKFLOW_ONLINE_USERS_REDIS_BATCH_SIZE + assert second_pipeline.hgetall.call_count == 1 + + @staticmethod + def test_workflow_online_users_rejects_excessive_workflow_ids(app: Flask, monkeypatch: pytest.MonkeyPatch) -> None: + get_tenant_app_maintainers = Mock(return_value={}) + monkeypatch.setattr( + workflow_module, + "WorkflowService", + lambda: SimpleNamespace(get_tenant_app_maintainers=get_tenant_app_maintainers), + ) + + excessive_ids = [f"wf-{index}" for index in range(workflow_module.MAX_WORKFLOW_ONLINE_USERS_REQUEST_IDS + 1)] + + api = workflow_module.WorkflowOnlineUsersApi() + handler = inspect.unwrap(api.post) + + with app.test_request_context( + "/apps/workflows/online-users", + method="POST", + json={"app_ids": excessive_ids}, + ): + args = workflow_module.WorkflowOnlineUsersPayload.model_validate({"app_ids": excessive_ids}) + with pytest.raises(HTTPException) as exc: + handler(api, args, "tenant-1", SimpleNamespace(id="account-1")) + + assert exc.value.code == 400 + assert exc.value.description is not None + assert "Maximum" in exc.value.description + get_tenant_app_maintainers.assert_not_called() -def test_workflow_online_users_batches_redis_reads(app: Flask, monkeypatch: pytest.MonkeyPatch) -> None: - app_ids = [f"wf-{index}" for index in range(workflow_module.WORKFLOW_ONLINE_USERS_REDIS_BATCH_SIZE + 1)] - monkeypatch.setattr( - workflow_module, - "WorkflowService", - lambda: SimpleNamespace(get_tenant_app_maintainers=lambda app_ids, tenant_id, session: dict.fromkeys(app_ids)), - ) - apply_config_overrides(monkeypatch, RBAC_ENABLED=False) - monkeypatch.setattr(workflow_module.session_factory, "create_session", lambda: nullcontext(Mock())) - - first_pipeline = Mock() - first_pipeline.execute.return_value = [{} for _ in range(workflow_module.WORKFLOW_ONLINE_USERS_REDIS_BATCH_SIZE)] - second_pipeline = Mock() - second_pipeline.execute.return_value = [{}] - redis_pipeline_factory = Mock(side_effect=[first_pipeline, second_pipeline]) - monkeypatch.setattr(workflow_module.redis_client, "pipeline", redis_pipeline_factory) - - api = workflow_module.WorkflowOnlineUsersApi() - handler = inspect.unwrap(api.post) - - with app.test_request_context( - "/apps/workflows/online-users", - method="POST", - json={"app_ids": app_ids}, - ): - args = workflow_module.WorkflowOnlineUsersPayload.model_validate({"app_ids": app_ids}) - response = handler(api, args, "tenant-1", SimpleNamespace(id="account-1")) - - assert len(response["data"]) == len(app_ids) - assert redis_pipeline_factory.call_count == 2 - assert first_pipeline.hgetall.call_count == workflow_module.WORKFLOW_ONLINE_USERS_REDIS_BATCH_SIZE - assert second_pipeline.hgetall.call_count == 1 +"""Unit tests for convert-to-workflow endpoint.""" -def test_workflow_online_users_rejects_excessive_workflow_ids(app: Flask, monkeypatch: pytest.MonkeyPatch) -> None: - get_tenant_app_maintainers = Mock(return_value={}) - monkeypatch.setattr( - workflow_module, - "WorkflowService", - lambda: SimpleNamespace(get_tenant_app_maintainers=get_tenant_app_maintainers), +def _workflow_convert_api_app(app_id: str) -> _workflow_convert_api_App: + return make_app(app_id=app_id, name=f"App {app_id}") + + +class TestWorkflowConvertApiCases: + class TestConvertToWorkflowApi: + @pytest.fixture + def api(self): + return workflow_module.ConvertToWorkflowApi() + + def test_convert_to_workflow_attaches_permission_keys_when_rbac_enabled( + self, api: ConvertToWorkflowApi, app: Flask, monkeypatch: pytest.MonkeyPatch + ) -> None: + method = unwrap(api.post) + new_app = _workflow_convert_api_app("new-app-1") + + monkeypatch.setattr( + workflow_module, + "WorkflowService", + lambda: type("WorkflowServiceStub", (), {"convert_to_workflow": lambda self, **_kwargs: new_app})(), + ) + monkeypatch.setattr( + workflow_module, + "get_app_permission_keys", + lambda tenant_id, account_id, app_id: ["app.acl.view_layout", "app.acl.edit"], + ) + + with app.test_request_context( + "/console/api/apps/app-1/convert-to-workflow", + method="POST", + json={}, + ): + current_user = _workflow_convert_api_Account(name="Current user", email="user@example.com") + current_user.id = "u1" + response = method( + api, + workflow_module.ConvertToWorkflowPayload.model_validate({}), + current_tenant_id="tenant-1", + current_user=current_user, + app_model=_workflow_convert_api_app("app-1"), + ) + + assert response["new_app_id"] == "new-app-1" + assert response["permission_keys"] == ["app.acl.view_layout", "app.acl.edit"] + + +def _make_account() -> _workflow_human_input_debug_api_Account: + account = _workflow_human_input_debug_api_Account(name="tester", email="tester@example.com") + account.status = AccountStatus.ACTIVE + account.role = TenantAccountRole.OWNER + account.id = "account-123" # type: ignore[assignment] + tenant = Tenant(name="Test tenant") + tenant.id = "tenant-123" + account._current_tenant = tenant + account._get_current_object = lambda: account # type: ignore[attr-defined] + return account + + +def _make_app(mode: AppMode) -> _workflow_human_input_debug_api_App: + return _workflow_human_input_debug_api_App( + id="app-123", + tenant_id="tenant-123", + name="Human input app", + description="", + mode=mode, + icon_type=IconType.EMOJI, + icon="robot", + icon_background="#FFFFFF", + enable_site=True, + enable_api=True, + max_active_requests=None, ) - excessive_ids = [f"wf-{index}" for index in range(workflow_module.MAX_WORKFLOW_ONLINE_USERS_REQUEST_IDS + 1)] - api = workflow_module.WorkflowOnlineUsersApi() - handler = inspect.unwrap(api.post) +def _patch_console_guards( + monkeypatch: pytest.MonkeyPatch, + account: _workflow_human_input_debug_api_Account, + app_model: _workflow_human_input_debug_api_App, +) -> None: + # Skip setup and auth guardrails + apply_config_overrides( + monkeypatch, + DEPLOYMENT_EDITION=DeploymentEdition.CLOUD, + LOGIN_DISABLED=True, + INIT_PASSWORD="", + ) + monkeypatch.setattr(login_lib, "current_user", account) + monkeypatch.setattr(login_lib, "current_account_with_tenant", lambda: (account, account.current_tenant_id)) + monkeypatch.setattr(login_lib, "check_csrf_token", lambda *_, **__: None) + monkeypatch.setattr(console_wraps, "current_account_with_tenant", lambda: (account, account.current_tenant_id)) + monkeypatch.setattr(app_wraps, "current_account_with_tenant", lambda: (account, account.current_tenant_id)) - with app.test_request_context( - "/apps/workflows/online-users", - method="POST", - json={"app_ids": excessive_ids}, - ): - args = workflow_module.WorkflowOnlineUsersPayload.model_validate({"app_ids": excessive_ids}) - with pytest.raises(HTTPException) as exc: - handler(api, args, "tenant-1", SimpleNamespace(id="account-1")) + # Avoid hitting the database when resolving the app model + monkeypatch.setattr(app_wraps, "_load_app_model_from_scoped_session", lambda _app_id: app_model) - assert exc.value.code == 400 - assert exc.value.description is not None - assert "Maximum" in exc.value.description - get_tenant_app_maintainers.assert_not_called() + +@dataclass +class PreviewCase: + resource_cls: type + path: str + mode: AppMode + + +@dataclass +class SubmitCase: + resource_cls: type + path: str + mode: AppMode + + +@dataclass +class DeliveryTestCase: + resource_cls: type + path: str + mode: AppMode + + +class TestWorkflowHumanInputDebugApiCases: + @staticmethod + @pytest.mark.parametrize( + "case", + [ + PreviewCase( + resource_cls=workflow_module.AdvancedChatDraftHumanInputFormPreviewApi, + path="/console/api/apps/app-123/advanced-chat/workflows/draft/human-input/nodes/node-42/form/preview", + mode=AppMode.ADVANCED_CHAT, + ), + PreviewCase( + resource_cls=workflow_module.WorkflowDraftHumanInputFormPreviewApi, + path="/console/api/apps/app-123/workflows/draft/human-input/nodes/node-42/form/preview", + mode=AppMode.WORKFLOW, + ), + ], + ) + def test_human_input_preview_delegates_to_service( + app: Flask, monkeypatch: pytest.MonkeyPatch, case: PreviewCase + ) -> None: + account = _make_account() + app_model = _make_app(case.mode) + _patch_console_guards(monkeypatch, account, app_model) + + preview_payload = { + "form_id": "node-42", + "node_id": "node-42", + "node_title": "Human input", + "form_content": "
example
", + "inputs": [ + {"type": "paragraph", "output_variable_name": "topic", "default": None}, + { + "type": "select", + "output_variable_name": "choice", + "option_source": {"type": "constant", "selector": list[str](), "value": ["one", "two"]}, + }, + { + "type": "file", + "output_variable_name": "attachment", + "allowed_file_types": ["document"], + "allowed_file_extensions": list[str](), + "allowed_file_upload_methods": ["local_file", "remote_url"], + }, + { + "type": "file-list", + "output_variable_name": "attachments", + "allowed_file_types": ["image"], + "allowed_file_extensions": list[str](), + "allowed_file_upload_methods": ["local_file"], + "number_limits": 0, + }, + ], + "actions": [{"id": "continue", "title": "Continue", "button_style": "default"}], + "resolved_default_values": {"topic": ""}, + } + service_instance = MagicMock() + service_instance.get_human_input_form_preview.return_value = preview_payload + monkeypatch.setattr(workflow_module, "WorkflowService", MagicMock(return_value=service_instance)) + + with app.test_request_context(case.path, method="POST", json={"inputs": {"topic": "tech"}}): + response = case.resource_cls().post(app_id=app_model.id, node_id="node-42") + + assert response == { + **preview_payload, + "display_in_ui": None, + "form_token": None, + "expiration_time": None, + } + service_instance.get_human_input_form_preview.assert_called_once_with( + app_model=app_model, + account=account, + node_id="node-42", + inputs={"topic": "tech"}, + session=ANY, + ) + + @staticmethod + @pytest.mark.parametrize( + "case", + [ + SubmitCase( + resource_cls=workflow_module.AdvancedChatDraftHumanInputFormRunApi, + path="/console/api/apps/app-123/advanced-chat/workflows/draft/human-input/nodes/node-99/form/run", + mode=AppMode.ADVANCED_CHAT, + ), + SubmitCase( + resource_cls=workflow_module.WorkflowDraftHumanInputFormRunApi, + path="/console/api/apps/app-123/workflows/draft/human-input/nodes/node-99/form/run", + mode=AppMode.WORKFLOW, + ), + ], + ) + def test_human_input_submit_forwards_payload(app: Flask, monkeypatch: pytest.MonkeyPatch, case: SubmitCase) -> None: + account = _make_account() + app_model = _make_app(case.mode) + _patch_console_guards(monkeypatch, account, app_model) + + result_payload = {"node_id": "node-99", "outputs": {"__rendered_content": "

done

"}, "action": "approve"} + service_instance = MagicMock() + service_instance.submit_human_input_form_preview.return_value = result_payload + monkeypatch.setattr(workflow_module, "WorkflowService", MagicMock(return_value=service_instance)) + + with app.test_request_context( + case.path, + method="POST", + json={"form_inputs": {"answer": "42"}, "inputs": {"#node-1.result#": "LLM output"}, "action": "approve"}, + ): + response = case.resource_cls().post(app_id=app_model.id, node_id="node-99") + + assert response == result_payload + service_instance.submit_human_input_form_preview.assert_called_once_with( + app_model=app_model, + account=account, + node_id="node-99", + form_inputs={"answer": "42"}, + inputs={"#node-1.result#": "LLM output"}, + action="approve", + session=ANY, + ) + + @staticmethod + @pytest.mark.parametrize( + "case", + [ + DeliveryTestCase( + resource_cls=workflow_module.WorkflowDraftHumanInputDeliveryTestApi, + path="/console/api/apps/app-123/workflows/draft/human-input/nodes/node-7/delivery-test", + mode=AppMode.ADVANCED_CHAT, + ), + DeliveryTestCase( + resource_cls=workflow_module.WorkflowDraftHumanInputDeliveryTestApi, + path="/console/api/apps/app-123/workflows/draft/human-input/nodes/node-7/delivery-test", + mode=AppMode.WORKFLOW, + ), + ], + ) + def test_human_input_delivery_test_calls_service( + app: Flask, monkeypatch: pytest.MonkeyPatch, case: DeliveryTestCase + ) -> None: + account = _make_account() + app_model = _make_app(case.mode) + _patch_console_guards(monkeypatch, account, app_model) + + service_instance = MagicMock() + monkeypatch.setattr(workflow_module, "WorkflowService", MagicMock(return_value=service_instance)) + + with app.test_request_context( + case.path, + method="POST", + json={"delivery_method_id": "delivery-123"}, + ): + response = case.resource_cls().post(app_id=app_model.id, node_id="node-7") + + assert response == {} + service_instance.test_human_input_delivery.assert_called_once_with( + app_model=app_model, + account=account, + node_id="node-7", + delivery_method_id="delivery-123", + inputs={}, + session=ANY, + ) + + @staticmethod + def test_human_input_delivery_test_maps_validation_error(app: Flask, monkeypatch: pytest.MonkeyPatch) -> None: + account = _make_account() + app_model = _make_app(AppMode.ADVANCED_CHAT) + _patch_console_guards(monkeypatch, account, app_model) + + service_instance = MagicMock() + service_instance.test_human_input_delivery.side_effect = ValueError("bad delivery method") + monkeypatch.setattr(workflow_module, "WorkflowService", MagicMock(return_value=service_instance)) + + with app.test_request_context( + "/console/api/apps/app-123/workflows/draft/human-input/nodes/node-1/delivery-test", + method="POST", + json={"delivery_method_id": "bad"}, + ): + with pytest.raises(ValueError, match="bad delivery method"): + workflow_module.WorkflowDraftHumanInputDeliveryTestApi().post(app_id=app_model.id, node_id="node-1") + + @staticmethod + def test_human_input_preview_rejects_non_mapping(app: Flask, monkeypatch: pytest.MonkeyPatch) -> None: + account = _make_account() + app_model = _make_app(AppMode.ADVANCED_CHAT) + _patch_console_guards(monkeypatch, account, app_model) + + with app.test_request_context( + "/console/api/apps/app-123/advanced-chat/workflows/draft/human-input/nodes/node-1/form/preview", + method="POST", + json={"inputs": ["not-a-dict"]}, + ): + with pytest.raises(UnprocessableEntity): + workflow_module.AdvancedChatDraftHumanInputFormPreviewApi().post(app_id=app_model.id, node_id="node-1") diff --git a/api/tests/unit_tests/controllers/console/app/test_workflow_app_log_api.py b/api/tests/unit_tests/controllers/console/app/test_workflow_app_log.py similarity index 100% rename from api/tests/unit_tests/controllers/console/app/test_workflow_app_log_api.py rename to api/tests/unit_tests/controllers/console/app/test_workflow_app_log.py diff --git a/api/tests/unit_tests/controllers/console/app/test_workflow_comment_api.py b/api/tests/unit_tests/controllers/console/app/test_workflow_comment.py similarity index 99% rename from api/tests/unit_tests/controllers/console/app/test_workflow_comment_api.py rename to api/tests/unit_tests/controllers/console/app/test_workflow_comment.py index 4712ee48294..34ac563f187 100644 --- a/api/tests/unit_tests/controllers/console/app/test_workflow_comment_api.py +++ b/api/tests/unit_tests/controllers/console/app/test_workflow_comment.py @@ -49,8 +49,8 @@ from services.app.workflow_comment_service import ( WorkflowCommentSummary, WorkflowCommentUpdated, ) -from tests.unit_tests.config_override import apply_config_overrides -from tests.unit_tests.model_factories import make_account +from tests.helpers.unit.config_override import apply_config_overrides +from tests.helpers.unit.model_factories import make_account WORKSPACE_ID = "tenant-123" APP_ID = UUID("11111111-1111-1111-1111-111111111111") diff --git a/api/tests/unit_tests/controllers/console/app/test_workflow_convert_api.py b/api/tests/unit_tests/controllers/console/app/test_workflow_convert_api.py deleted file mode 100644 index f9a9bc80d06..00000000000 --- a/api/tests/unit_tests/controllers/console/app/test_workflow_convert_api.py +++ /dev/null @@ -1,58 +0,0 @@ -"""Unit tests for convert-to-workflow endpoint.""" - -from __future__ import annotations - -from inspect import unwrap - -import pytest -from flask import Flask - -from controllers.console.app import workflow as workflow_module -from controllers.console.app.workflow import ConvertToWorkflowApi -from models import Account, App -from tests.unit_tests.model_factories import make_app - - -def _app(app_id: str) -> App: - return make_app(app_id=app_id, name=f"App {app_id}") - - -class TestConvertToWorkflowApi: - @pytest.fixture - def api(self): - return workflow_module.ConvertToWorkflowApi() - - def test_convert_to_workflow_attaches_permission_keys_when_rbac_enabled( - self, api: ConvertToWorkflowApi, app: Flask, monkeypatch: pytest.MonkeyPatch - ) -> None: - method = unwrap(api.post) - new_app = _app("new-app-1") - - monkeypatch.setattr( - workflow_module, - "WorkflowService", - lambda: type("WorkflowServiceStub", (), {"convert_to_workflow": lambda self, **_kwargs: new_app})(), - ) - monkeypatch.setattr( - workflow_module, - "get_app_permission_keys", - lambda tenant_id, account_id, app_id: ["app.acl.view_layout", "app.acl.edit"], - ) - - with app.test_request_context( - "/console/api/apps/app-1/convert-to-workflow", - method="POST", - json={}, - ): - current_user = Account(name="Current user", email="user@example.com") - current_user.id = "u1" - response = method( - api, - workflow_module.ConvertToWorkflowPayload.model_validate({}), - current_tenant_id="tenant-1", - current_user=current_user, - app_model=_app("app-1"), - ) - - assert response["new_app_id"] == "new-app-1" - assert response["permission_keys"] == ["app.acl.view_layout", "app.acl.edit"] diff --git a/api/tests/unit_tests/controllers/console/app/test_workflow_human_input_debug_api.py b/api/tests/unit_tests/controllers/console/app/test_workflow_human_input_debug_api.py deleted file mode 100644 index ea1dc087a1d..00000000000 --- a/api/tests/unit_tests/controllers/console/app/test_workflow_human_input_debug_api.py +++ /dev/null @@ -1,279 +0,0 @@ -from __future__ import annotations - -from dataclasses import dataclass -from unittest.mock import ANY, MagicMock - -import pytest -from flask import Flask -from werkzeug.exceptions import UnprocessableEntity - -from controllers.console import wraps as console_wraps -from controllers.console.app import workflow as workflow_module -from controllers.console.app import wraps as app_wraps -from enums import DeploymentEdition -from libs import login as login_lib -from models import App, Tenant -from models.account import Account, AccountStatus, TenantAccountRole -from models.model import AppMode, IconType -from tests.unit_tests.config_override import apply_config_overrides - - -def _make_account() -> Account: - account = Account(name="tester", email="tester@example.com") - account.status = AccountStatus.ACTIVE - account.role = TenantAccountRole.OWNER - account.id = "account-123" # type: ignore[assignment] - tenant = Tenant(name="Test tenant") - tenant.id = "tenant-123" - account._current_tenant = tenant - account._get_current_object = lambda: account # type: ignore[attr-defined] - return account - - -def _make_app(mode: AppMode) -> App: - return App( - id="app-123", - tenant_id="tenant-123", - name="Human input app", - description="", - mode=mode, - icon_type=IconType.EMOJI, - icon="robot", - icon_background="#FFFFFF", - enable_site=True, - enable_api=True, - max_active_requests=None, - ) - - -def _patch_console_guards(monkeypatch: pytest.MonkeyPatch, account: Account, app_model: App) -> None: - # Skip setup and auth guardrails - apply_config_overrides( - monkeypatch, - DEPLOYMENT_EDITION=DeploymentEdition.CLOUD, - LOGIN_DISABLED=True, - INIT_PASSWORD="", - ) - monkeypatch.setattr(login_lib, "current_user", account) - monkeypatch.setattr(login_lib, "current_account_with_tenant", lambda: (account, account.current_tenant_id)) - monkeypatch.setattr(login_lib, "check_csrf_token", lambda *_, **__: None) - monkeypatch.setattr(console_wraps, "current_account_with_tenant", lambda: (account, account.current_tenant_id)) - monkeypatch.setattr(app_wraps, "current_account_with_tenant", lambda: (account, account.current_tenant_id)) - - # Avoid hitting the database when resolving the app model - monkeypatch.setattr(app_wraps, "_load_app_model_from_scoped_session", lambda _app_id: app_model) - - -@dataclass -class PreviewCase: - resource_cls: type - path: str - mode: AppMode - - -@pytest.mark.parametrize( - "case", - [ - PreviewCase( - resource_cls=workflow_module.AdvancedChatDraftHumanInputFormPreviewApi, - path="/console/api/apps/app-123/advanced-chat/workflows/draft/human-input/nodes/node-42/form/preview", - mode=AppMode.ADVANCED_CHAT, - ), - PreviewCase( - resource_cls=workflow_module.WorkflowDraftHumanInputFormPreviewApi, - path="/console/api/apps/app-123/workflows/draft/human-input/nodes/node-42/form/preview", - mode=AppMode.WORKFLOW, - ), - ], -) -def test_human_input_preview_delegates_to_service( - app: Flask, monkeypatch: pytest.MonkeyPatch, case: PreviewCase -) -> None: - account = _make_account() - app_model = _make_app(case.mode) - _patch_console_guards(monkeypatch, account, app_model) - - preview_payload = { - "form_id": "node-42", - "node_id": "node-42", - "node_title": "Human input", - "form_content": "
example
", - "inputs": [ - {"type": "paragraph", "output_variable_name": "topic", "default": None}, - { - "type": "select", - "output_variable_name": "choice", - "option_source": {"type": "constant", "selector": list[str](), "value": ["one", "two"]}, - }, - { - "type": "file", - "output_variable_name": "attachment", - "allowed_file_types": ["document"], - "allowed_file_extensions": list[str](), - "allowed_file_upload_methods": ["local_file", "remote_url"], - }, - { - "type": "file-list", - "output_variable_name": "attachments", - "allowed_file_types": ["image"], - "allowed_file_extensions": list[str](), - "allowed_file_upload_methods": ["local_file"], - "number_limits": 0, - }, - ], - "actions": [{"id": "continue", "title": "Continue", "button_style": "default"}], - "resolved_default_values": {"topic": ""}, - } - service_instance = MagicMock() - service_instance.get_human_input_form_preview.return_value = preview_payload - monkeypatch.setattr(workflow_module, "WorkflowService", MagicMock(return_value=service_instance)) - - with app.test_request_context(case.path, method="POST", json={"inputs": {"topic": "tech"}}): - response = case.resource_cls().post(app_id=app_model.id, node_id="node-42") - - assert response == { - **preview_payload, - "display_in_ui": None, - "form_token": None, - "expiration_time": None, - } - service_instance.get_human_input_form_preview.assert_called_once_with( - app_model=app_model, - account=account, - node_id="node-42", - inputs={"topic": "tech"}, - session=ANY, - ) - - -@dataclass -class SubmitCase: - resource_cls: type - path: str - mode: AppMode - - -@pytest.mark.parametrize( - "case", - [ - SubmitCase( - resource_cls=workflow_module.AdvancedChatDraftHumanInputFormRunApi, - path="/console/api/apps/app-123/advanced-chat/workflows/draft/human-input/nodes/node-99/form/run", - mode=AppMode.ADVANCED_CHAT, - ), - SubmitCase( - resource_cls=workflow_module.WorkflowDraftHumanInputFormRunApi, - path="/console/api/apps/app-123/workflows/draft/human-input/nodes/node-99/form/run", - mode=AppMode.WORKFLOW, - ), - ], -) -def test_human_input_submit_forwards_payload(app: Flask, monkeypatch: pytest.MonkeyPatch, case: SubmitCase) -> None: - account = _make_account() - app_model = _make_app(case.mode) - _patch_console_guards(monkeypatch, account, app_model) - - result_payload = {"node_id": "node-99", "outputs": {"__rendered_content": "

done

"}, "action": "approve"} - service_instance = MagicMock() - service_instance.submit_human_input_form_preview.return_value = result_payload - monkeypatch.setattr(workflow_module, "WorkflowService", MagicMock(return_value=service_instance)) - - with app.test_request_context( - case.path, - method="POST", - json={"form_inputs": {"answer": "42"}, "inputs": {"#node-1.result#": "LLM output"}, "action": "approve"}, - ): - response = case.resource_cls().post(app_id=app_model.id, node_id="node-99") - - assert response == result_payload - service_instance.submit_human_input_form_preview.assert_called_once_with( - app_model=app_model, - account=account, - node_id="node-99", - form_inputs={"answer": "42"}, - inputs={"#node-1.result#": "LLM output"}, - action="approve", - session=ANY, - ) - - -@dataclass -class DeliveryTestCase: - resource_cls: type - path: str - mode: AppMode - - -@pytest.mark.parametrize( - "case", - [ - DeliveryTestCase( - resource_cls=workflow_module.WorkflowDraftHumanInputDeliveryTestApi, - path="/console/api/apps/app-123/workflows/draft/human-input/nodes/node-7/delivery-test", - mode=AppMode.ADVANCED_CHAT, - ), - DeliveryTestCase( - resource_cls=workflow_module.WorkflowDraftHumanInputDeliveryTestApi, - path="/console/api/apps/app-123/workflows/draft/human-input/nodes/node-7/delivery-test", - mode=AppMode.WORKFLOW, - ), - ], -) -def test_human_input_delivery_test_calls_service( - app: Flask, monkeypatch: pytest.MonkeyPatch, case: DeliveryTestCase -) -> None: - account = _make_account() - app_model = _make_app(case.mode) - _patch_console_guards(monkeypatch, account, app_model) - - service_instance = MagicMock() - monkeypatch.setattr(workflow_module, "WorkflowService", MagicMock(return_value=service_instance)) - - with app.test_request_context( - case.path, - method="POST", - json={"delivery_method_id": "delivery-123"}, - ): - response = case.resource_cls().post(app_id=app_model.id, node_id="node-7") - - assert response == {} - service_instance.test_human_input_delivery.assert_called_once_with( - app_model=app_model, - account=account, - node_id="node-7", - delivery_method_id="delivery-123", - inputs={}, - session=ANY, - ) - - -def test_human_input_delivery_test_maps_validation_error(app: Flask, monkeypatch: pytest.MonkeyPatch) -> None: - account = _make_account() - app_model = _make_app(AppMode.ADVANCED_CHAT) - _patch_console_guards(monkeypatch, account, app_model) - - service_instance = MagicMock() - service_instance.test_human_input_delivery.side_effect = ValueError("bad delivery method") - monkeypatch.setattr(workflow_module, "WorkflowService", MagicMock(return_value=service_instance)) - - with app.test_request_context( - "/console/api/apps/app-123/workflows/draft/human-input/nodes/node-1/delivery-test", - method="POST", - json={"delivery_method_id": "bad"}, - ): - with pytest.raises(ValueError, match="bad delivery method"): - workflow_module.WorkflowDraftHumanInputDeliveryTestApi().post(app_id=app_model.id, node_id="node-1") - - -def test_human_input_preview_rejects_non_mapping(app: Flask, monkeypatch: pytest.MonkeyPatch) -> None: - account = _make_account() - app_model = _make_app(AppMode.ADVANCED_CHAT) - _patch_console_guards(monkeypatch, account, app_model) - - with app.test_request_context( - "/console/api/apps/app-123/advanced-chat/workflows/draft/human-input/nodes/node-1/form/preview", - method="POST", - json={"inputs": ["not-a-dict"]}, - ): - with pytest.raises(UnprocessableEntity): - workflow_module.AdvancedChatDraftHumanInputFormPreviewApi().post(app_id=app_model.id, node_id="node-1") diff --git a/api/tests/unit_tests/controllers/console/app/test_workflow_node_output_inspector.py b/api/tests/unit_tests/controllers/console/app/test_workflow_node_output_inspector.py index cdb76b1aa6c..0b4e51e896b 100644 --- a/api/tests/unit_tests/controllers/console/app/test_workflow_node_output_inspector.py +++ b/api/tests/unit_tests/controllers/console/app/test_workflow_node_output_inspector.py @@ -475,7 +475,7 @@ def test_serve_output_preview_translates_inspector_error(patch_service, app_mode # ``@login_required`` / ``@account_initialization_required`` / # ``@get_app_model``) needs a real Flask request context + DB-backed account. # The integration test in -# ``tests/integration_tests/services/test_node_output_inspector_service.py`` +# ``tests/e2e_tests/middleware/services/test_node_output_inspector_service.py`` # (and the E2E driver in /tmp/e2e_inspector_sse_published.py) exercise them # through the HTTP stack. # ────────────────────────────────────────────────────────────────────────────── diff --git a/api/tests/unit_tests/controllers/console/app/test_workflow_pause_details_api.py b/api/tests/unit_tests/controllers/console/app/test_workflow_pause_details_api.py deleted file mode 100644 index c5872bbf434..00000000000 --- a/api/tests/unit_tests/controllers/console/app/test_workflow_pause_details_api.py +++ /dev/null @@ -1,115 +0,0 @@ -"""Controller tests for Console workflow pause details.""" - -from __future__ import annotations - -from datetime import datetime -from inspect import unwrap -from types import SimpleNamespace -from unittest.mock import Mock - -import pytest -from flask import Flask - -from controllers.common.errors import NotFoundError -from controllers.console.app import workflow_run as workflow_run_module -from machinery.context import RequestContext -from services.workflow_run_service import WorkflowRunPauseDetails, WorkflowRunPausedNode -from tests.unit_tests.config_override import apply_config_overrides - - -def _request_context(*, workspace_id: str = "tenant-1") -> RequestContext: - return RequestContext( - request_id="request-1", - trace_id="trace-1", - account_id="account-1", - active_workspace_id=workspace_id, - ) - - -def _mock_application_services(monkeypatch: pytest.MonkeyPatch, workflow_runs: Mock) -> None: - monkeypatch.setattr( - workflow_run_module, - "application_services", - lambda: SimpleNamespace(workflow_runs=workflow_runs), - ) - - -def test_pause_details_returns_backstage_input_url(app: Flask, monkeypatch: pytest.MonkeyPatch) -> None: - apply_config_overrides(monkeypatch, APP_WEB_URL="https://web.example.com/") - workflow_runs = Mock() - workflow_runs.get_pause_details.return_value = WorkflowRunPauseDetails( - paused_at=datetime(2024, 1, 1, 12, 0, 0), - paused_nodes=( - WorkflowRunPausedNode( - node_id="node-1", - node_title="Ask Name", - form_id="form-1", - form_token="backstage-token", - ), - ), - ) - _mock_application_services(monkeypatch, workflow_runs) - request_context = _request_context() - - api = workflow_run_module.ConsoleWorkflowPauseDetailsApi() - handler = unwrap(api.get) - with app.test_request_context("/console/api/workflow/run-1/pause-details", method="GET"): - response, status = handler(api, request_context, workflow_run_id="run-1") - - assert status == 200 - assert response == { - "paused_at": "2024-01-01T12:00:00Z", - "paused_nodes": [ - { - "node_id": "node-1", - "node_title": "Ask Name", - "pause_type": { - "type": "human_input", - "form_id": "form-1", - "backstage_input_url": "https://web.example.com/form/backstage-token", - }, - } - ], - } - workflow_runs.get_pause_details.assert_called_once_with( - request_context, - workflow_run_id="run-1", - ) - - -def test_pause_details_maps_missing_or_inaccessible_run_to_not_found( - app: Flask, monkeypatch: pytest.MonkeyPatch -) -> None: - workflow_runs = Mock() - workflow_runs.get_pause_details.return_value = None - _mock_application_services(monkeypatch, workflow_runs) - request_context = _request_context(workspace_id="other-tenant") - api = workflow_run_module.ConsoleWorkflowPauseDetailsApi() - handler = unwrap(api.get) - - with app.test_request_context("/console/api/workflow/run-1/pause-details", method="GET"): - with pytest.raises(NotFoundError, match="Workflow run not found"): - handler(api, request_context, workflow_run_id="run-1") - - workflow_runs.get_pause_details.assert_called_once_with( - request_context, - workflow_run_id="run-1", - ) - - -def test_pause_details_returns_empty_response_for_non_paused_run(app: Flask, monkeypatch: pytest.MonkeyPatch) -> None: - workflow_runs = Mock() - workflow_runs.get_pause_details.return_value = WorkflowRunPauseDetails(paused_at=None, paused_nodes=()) - _mock_application_services(monkeypatch, workflow_runs) - - api = workflow_run_module.ConsoleWorkflowPauseDetailsApi() - handler = unwrap(api.get) - with app.test_request_context("/console/api/workflow/run-1/pause-details", method="GET"): - response, status = handler(api, _request_context(), workflow_run_id="run-1") - - assert status == 200 - assert response == {"paused_at": None, "paused_nodes": []} - - -def test_pause_details_response_schema_is_registered() -> None: - assert workflow_run_module.WorkflowPauseDetailsResponse.__name__ in workflow_run_module.console_ns.models diff --git a/api/tests/unit_tests/controllers/console/app/test_workflow_run.py b/api/tests/unit_tests/controllers/console/app/test_workflow_run.py new file mode 100644 index 00000000000..8a11d913c30 --- /dev/null +++ b/api/tests/unit_tests/controllers/console/app/test_workflow_run.py @@ -0,0 +1,632 @@ +from __future__ import annotations + +import json +from collections.abc import Callable +from datetime import UTC, datetime +from inspect import unwrap +from types import SimpleNamespace +from typing import Any +from unittest.mock import Mock + +import pytest +from flask import Flask, request +from flask_restx import marshal +from pytest_mock import MockerFixture +from sqlalchemy.orm import Session, sessionmaker +from werkzeug.exceptions import Forbidden + +from controllers.common.errors import NotFoundError +from controllers.common.rbac import locators as rbac_locators +from controllers.console.app import workflow_run as workflow_run_module +from core.rbac import RBACPermission, RBACResourceScope +from extensions.ext_database import db +from fields.workflow_run_fields import node_execution_response_source +from graphon.enums import WorkflowExecutionStatus, WorkflowNodeExecutionStatus +from machinery.context import RequestContext +from models import Account, App, AppMode +from models.enums import CreatorUserRole, WorkflowRunTriggeredFrom +from models.workflow import ( + WorkflowNodeExecutionModel, + WorkflowNodeExecutionTriggeredFrom, + WorkflowRun, + WorkflowType, +) +from repositories.sqlalchemy_api_workflow_node_execution_repository import ( + DifyAPISQLAlchemyWorkflowNodeExecutionRepository, +) +from repositories.sqlalchemy_api_workflow_run_repository import DifyAPISQLAlchemyWorkflowRunRepository +from services.enterprise.rbac_service import RBACService +from services.workflow_node_execution_trace_service import WorkflowNodeExecutionTrace +from services.workflow_run_service import WorkflowRunPauseDetails, WorkflowRunPausedNode, WorkflowRunService +from tests.helpers.unit.config_override import apply_config_overrides +from tests.helpers.unit.model_factories import make_account, make_app +from tests.helpers.unit.rbac_fakes import RBACDomain + +"""Controller tests for Console workflow pause details.""" + + +def _workflow_pause_details_api_request_context(*, workspace_id: str = "tenant-1") -> RequestContext: + return RequestContext( + request_id="request-1", + trace_id="trace-1", + account_id="account-1", + active_workspace_id=workspace_id, + ) + + +def _workflow_pause_details_api_mock_application_services(monkeypatch: pytest.MonkeyPatch, workflow_runs: Mock) -> None: + monkeypatch.setattr( + workflow_run_module, + "application_services", + lambda: SimpleNamespace(workflow_runs=workflow_runs), + ) + + +class TestWorkflowPauseDetailsApiCases: + @staticmethod + def test_pause_details_returns_backstage_input_url(app: Flask, monkeypatch: pytest.MonkeyPatch) -> None: + apply_config_overrides(monkeypatch, APP_WEB_URL="https://web.example.com/") + workflow_runs = Mock() + workflow_runs.get_pause_details.return_value = WorkflowRunPauseDetails( + paused_at=datetime(2024, 1, 1, 12, 0, 0), + paused_nodes=( + WorkflowRunPausedNode( + node_id="node-1", + node_title="Ask Name", + form_id="form-1", + form_token="backstage-token", + ), + ), + ) + _workflow_pause_details_api_mock_application_services(monkeypatch, workflow_runs) + request_context = _workflow_pause_details_api_request_context() + + api = workflow_run_module.ConsoleWorkflowPauseDetailsApi() + handler = unwrap(api.get) + with app.test_request_context("/console/api/workflow/run-1/pause-details", method="GET"): + response, status = handler(api, request_context, workflow_run_id="run-1") + + assert status == 200 + assert response == { + "paused_at": "2024-01-01T12:00:00Z", + "paused_nodes": [ + { + "node_id": "node-1", + "node_title": "Ask Name", + "pause_type": { + "type": "human_input", + "form_id": "form-1", + "backstage_input_url": "https://web.example.com/form/backstage-token", + }, + } + ], + } + workflow_runs.get_pause_details.assert_called_once_with( + request_context, + workflow_run_id="run-1", + ) + + @staticmethod + def test_pause_details_maps_missing_or_inaccessible_run_to_not_found( + app: Flask, monkeypatch: pytest.MonkeyPatch + ) -> None: + workflow_runs = Mock() + workflow_runs.get_pause_details.return_value = None + _workflow_pause_details_api_mock_application_services(monkeypatch, workflow_runs) + request_context = _workflow_pause_details_api_request_context(workspace_id="other-tenant") + api = workflow_run_module.ConsoleWorkflowPauseDetailsApi() + handler = unwrap(api.get) + + with app.test_request_context("/console/api/workflow/run-1/pause-details", method="GET"): + with pytest.raises(NotFoundError, match="Workflow run not found"): + handler(api, request_context, workflow_run_id="run-1") + + workflow_runs.get_pause_details.assert_called_once_with( + request_context, + workflow_run_id="run-1", + ) + + @staticmethod + def test_pause_details_returns_empty_response_for_non_paused_run( + app: Flask, monkeypatch: pytest.MonkeyPatch + ) -> None: + workflow_runs = Mock() + workflow_runs.get_pause_details.return_value = WorkflowRunPauseDetails(paused_at=None, paused_nodes=()) + _workflow_pause_details_api_mock_application_services(monkeypatch, workflow_runs) + + api = workflow_run_module.ConsoleWorkflowPauseDetailsApi() + handler = unwrap(api.get) + with app.test_request_context("/console/api/workflow/run-1/pause-details", method="GET"): + response, status = handler(api, _workflow_pause_details_api_request_context(), workflow_run_id="run-1") + + assert status == 200 + assert response == {"paused_at": None, "paused_nodes": []} + + @staticmethod + def test_pause_details_response_schema_is_registered() -> None: + assert workflow_run_module.WorkflowPauseDetailsResponse.__name__ in workflow_run_module.console_ns.models + + +def _serialize_200_response(handler, payload: Any) -> Any: + response_doc = getattr(handler, "__apidoc__", {}).get("responses", {}).get("200") + if response_doc is None: + return payload + + response_model = response_doc[1] + if isinstance(response_model, dict): + return marshal(payload, response_model) + return payload + + +def _account(session: Session) -> Account: + account = make_account(name="Alice", email="alice@example.com") + session.add(account) + session.commit() + return account + + +def _app() -> App: + return make_app(name="Workflow run app", mode=AppMode.WORKFLOW) + + +def _workflow_run_api_request_context(*, workspace_id: str = "tenant-1") -> RequestContext: + return RequestContext( + request_id="request-1", + trace_id="trace-1", + account_id="account-1", + active_workspace_id=workspace_id, + ) + + +def _workflow_run_api_mock_application_services( + monkeypatch: pytest.MonkeyPatch, workflow_runs: Mock | WorkflowRunService +) -> None: + monkeypatch.setattr( + workflow_run_module, + "application_services", + lambda: SimpleNamespace(workflow_runs=workflow_runs), + ) + + +def _workflow_run_summary(session: Session, **overrides: object) -> WorkflowRun: + created_at = datetime(2026, 1, 2, 3, 4, 5, tzinfo=UTC) + workflow_run = WorkflowRun( + id="run-1", + tenant_id="tenant-1", + app_id="app-1", + workflow_id="workflow-1", + type=WorkflowType.WORKFLOW, + triggered_from=WorkflowRunTriggeredFrom.DEBUGGING, + version="v1", + graph='{"nodes": []}', + inputs='{"query": "hello"}', + status=WorkflowExecutionStatus.SUCCEEDED, + outputs='{"answer": "world"}', + error=None, + elapsed_time=1.5, + total_tokens=10, + total_steps=2, + created_by_role=CreatorUserRole.ACCOUNT, + created_by="account-1", + created_at=created_at, + finished_at=created_at, + exceptions_count=0, + ) + for name, value in overrides.items(): + setattr(workflow_run, name, value) + workflow_run.retry_index = 0 + session.add(workflow_run) + session.commit() + return workflow_run + + +def _workflow_run_node_execution(session: Session) -> WorkflowNodeExecutionModel: + created_at = datetime(2026, 1, 2, 3, 4, 5, tzinfo=UTC) + execution = WorkflowNodeExecutionModel( + id="node-exec-1", + tenant_id="tenant-1", + app_id="app-1", + workflow_id="workflow-1", + triggered_from=WorkflowNodeExecutionTriggeredFrom.WORKFLOW_RUN, + workflow_run_id="run-1", + index=1, + predecessor_node_id=None, + node_execution_id="node-execution-1", + node_id="node-1", + node_type="start", + title="Start", + agent_workspace_binding_id=None, + inputs=json.dumps({"query": "hello"}), + process_data=json.dumps({"step": "prepared"}), + outputs=json.dumps({"answer": "world"}), + status=WorkflowNodeExecutionStatus.SUCCEEDED, + error=None, + elapsed_time=1.0, + execution_metadata=json.dumps({"total_tokens": 3}), + created_at=created_at, + created_by_role=CreatorUserRole.ACCOUNT, + created_by="account-1", + finished_at=created_at, + ) + execution.offload_data = [] + session.add(execution) + session.commit() + return execution + + +class TestWorkflowRunApiCases: + @staticmethod + def test_workflow_run_list_returns_frontend_history_contract( + app: Flask, monkeypatch: pytest.MonkeyPatch, sqlite_session: Session + ) -> None: + _account(sqlite_session) + workflow_run = _workflow_run_summary(sqlite_session) + workflow_runs = Mock() + workflow_runs.get_paginate_workflow_runs.return_value = SimpleNamespace( + limit=10, + has_more=False, + data=[workflow_run], + ) + _workflow_run_api_mock_application_services(monkeypatch, workflow_runs) + request_context = _workflow_run_api_request_context() + + api = workflow_run_module.WorkflowRunListApi() + handler = unwrap(api.get) + + with app.test_request_context("/apps/app-1/workflow-runs?limit=10", method="GET"): + payload = handler( + api, + sqlite_session, + workflow_run_module.WorkflowRunListQuery(limit=10), + request_context, + app_model=_app(), + ) + + response = _serialize_200_response(api.get, payload) + + assert response["limit"] == 10 + assert response["has_more"] is False + assert response["data"][0] == { + "id": "run-1", + "version": "v1", + "status": "succeeded", + "elapsed_time": 1.5, + "total_tokens": 10, + "total_steps": 2, + "created_by_account": {"id": "account-1", "name": "Alice", "email": "alice@example.com"}, + "created_at": 1767323045, + "finished_at": 1767323045, + "exceptions_count": 0, + "retry_index": 0, + } + workflow_runs.get_paginate_workflow_runs.assert_called_once_with( + request_context, + app_id="app-1", + args={"limit": 10}, + triggered_from=WorkflowRunTriggeredFrom.DEBUGGING, + ) + + @staticmethod + def test_advanced_chat_workflow_run_list_keeps_message_fields( + app: Flask, monkeypatch: pytest.MonkeyPatch, sqlite_session: Session + ) -> None: + _account(sqlite_session) + workflow_run = _workflow_run_summary( + sqlite_session, + conversation_id="conversation-1", + message_id="message-1", + ) + workflow_runs = Mock() + workflow_runs.get_paginate_advanced_chat_workflow_runs.return_value = SimpleNamespace( + limit=1, + has_more=True, + data=[workflow_run], + ) + _workflow_run_api_mock_application_services(monkeypatch, workflow_runs) + request_context = _workflow_run_api_request_context() + + api = workflow_run_module.AdvancedChatAppWorkflowRunListApi() + handler = unwrap(api.get) + + with app.test_request_context("/apps/app-1/advanced-chat/workflow-runs?limit=1", method="GET"): + payload = handler( + api, + sqlite_session, + workflow_run_module.WorkflowRunListQuery(limit=1), + request_context, + app_model=_app(), + ) + + response = _serialize_200_response(api.get, payload) + + assert response["data"][0]["conversation_id"] == "conversation-1" + assert response["data"][0]["message_id"] == "message-1" + workflow_runs.get_paginate_advanced_chat_workflow_runs.assert_called_once_with( + request_context, + app_id="app-1", + args={"limit": 1}, + triggered_from=WorkflowRunTriggeredFrom.DEBUGGING, + ) + + @staticmethod + def test_workflow_run_count_passes_filters_to_application_service( + app: Flask, monkeypatch: pytest.MonkeyPatch + ) -> None: + workflow_runs = Mock() + workflow_runs.get_workflow_runs_count.return_value = { + "total": 2, + "running": 0, + "succeeded": 2, + "failed": 0, + "stopped": 0, + "partial-succeeded": 0, + } + _workflow_run_api_mock_application_services(monkeypatch, workflow_runs) + request_context = _workflow_run_api_request_context() + query = workflow_run_module.WorkflowRunCountQuery( + status="succeeded", + time_range="7d", + triggered_from="app-run", + ) + + api = workflow_run_module.WorkflowRunCountApi() + handler = unwrap(api.get) + with app.test_request_context("/apps/app-1/workflow-runs/count", method="GET"): + payload = handler(api, query, request_context, app_model=_app()) + + assert payload == { + "total": 2, + "running": 0, + "succeeded": 2, + "failed": 0, + "stopped": 0, + "partial_succeeded": 0, + } + workflow_runs.get_workflow_runs_count.assert_called_once_with( + request_context, + app_id="app-1", + status="succeeded", + time_range="7d", + triggered_from=WorkflowRunTriggeredFrom.APP_RUN, + ) + + @staticmethod + def test_workflow_run_detail_returns_frontend_detail_contract( + app: Flask, monkeypatch: pytest.MonkeyPatch, sqlite_session: Session + ) -> None: + _account(sqlite_session) + workflow_run = _workflow_run_summary(sqlite_session) + workflow_runs = Mock() + workflow_runs.get_workflow_run.return_value = workflow_run + _workflow_run_api_mock_application_services(monkeypatch, workflow_runs) + request_context = _workflow_run_api_request_context() + + api = workflow_run_module.WorkflowRunDetailApi() + handler = unwrap(api.get) + + with app.test_request_context("/apps/app-1/workflow-runs/run-1", method="GET"): + payload = handler(api, sqlite_session, request_context, app_model=_app(), run_id="run-1") + + response = _serialize_200_response(api.get, payload) + + assert response == { + "id": "run-1", + "version": "v1", + "graph": {"nodes": []}, + "inputs": {"query": "hello"}, + "status": "succeeded", + "outputs": {"answer": "world"}, + "error": None, + "elapsed_time": 1.5, + "total_tokens": 10, + "total_steps": 2, + "created_by_role": "account", + "created_by_account": {"id": "account-1", "name": "Alice", "email": "alice@example.com"}, + "created_by_end_user": None, + "created_at": 1767323045, + "finished_at": 1767323045, + "exceptions_count": 0, + } + workflow_runs.get_workflow_run.assert_called_once_with( + request_context, + app_id="app-1", + run_id="run-1", + ) + + @staticmethod + def test_workflow_run_detail_maps_missing_run_to_not_found( + app: Flask, monkeypatch: pytest.MonkeyPatch, sqlite_session: Session + ) -> None: + workflow_runs = Mock() + workflow_runs.get_workflow_run.return_value = None + _workflow_run_api_mock_application_services(monkeypatch, workflow_runs) + request_context = _workflow_run_api_request_context(workspace_id="tenant-2") + api = workflow_run_module.WorkflowRunDetailApi() + handler = unwrap(api.get) + + with app.test_request_context("/apps/app-1/workflow-runs/run-1", method="GET"): + with pytest.raises(NotFoundError, match="Workflow run not found"): + handler(api, sqlite_session, request_context, app_model=_app(), run_id="run-1") + + workflow_runs.get_workflow_run.assert_called_once_with( + request_context, + app_id="app-1", + run_id="run-1", + ) + + @staticmethod + def test_workflow_run_node_executions_return_frontend_trace_contract( + app: Flask, monkeypatch: pytest.MonkeyPatch, sqlite_session: Session + ) -> None: + _account(sqlite_session) + execution = _workflow_run_node_execution(sqlite_session) + workflow_runs = Mock() + workflow_runs.get_workflow_run_node_executions.return_value = [ + node_execution_response_source(execution, session=sqlite_session) + ] + _workflow_run_api_mock_application_services(monkeypatch, workflow_runs) + monkeypatch.setattr(db, "session", lambda: sqlite_session) + request_context = _workflow_run_api_request_context() + + api = workflow_run_module.WorkflowRunNodeExecutionListApi() + handler = unwrap(api.get) + + with app.test_request_context("/apps/app-1/workflow-runs/run-1/node-executions", method="GET"): + payload = handler(api, request_context, app_model=_app(), run_id="run-1") + + response = _serialize_200_response(api.get, payload) + + assert response == { + "data": [ + { + "id": "node-exec-1", + "index": 1, + "predecessor_node_id": None, + "node_id": "node-1", + "node_type": "start", + "title": "Start", + "inputs": {"query": "hello"}, + "process_data": {"step": "prepared"}, + "outputs": {"answer": "world"}, + "status": "succeeded", + "error": None, + "elapsed_time": 1.0, + "execution_metadata": {"total_tokens": 3}, + "extras": {}, + "created_at": 1767323045, + "created_by_role": "account", + "created_by_account": {"id": "account-1", "name": "Alice", "email": "alice@example.com"}, + "created_by_end_user": None, + "finished_at": 1767323045, + "inputs_truncated": False, + "outputs_truncated": False, + "process_data_truncated": False, + } + ] + } + workflow_runs.get_workflow_run_node_executions.assert_called_once_with( + request_context, + app_id="app-1", + run_id="run-1", + ) + + @staticmethod + @pytest.fixture + def workflow_run_service(sqlite_session_factory: sessionmaker[Session]) -> WorkflowRunService: + return WorkflowRunService( + workflow_runs=DifyAPISQLAlchemyWorkflowRunRepository(sqlite_session_factory), + node_executions=DifyAPISQLAlchemyWorkflowNodeExecutionRepository(sqlite_session_factory), + ) + + @staticmethod + @pytest.mark.parametrize("source_access", ["granted", "maintainer", "rbac-disabled"]) + def test_workflow_tool_children_preserve_source_identity( + app: Flask, + monkeypatch: pytest.MonkeyPatch, + sqlite_session: Session, + rbac_domain: RBACDomain, + config_overrides: Callable[..., None], + source_access: str, + workflow_run_service: WorkflowRunService, + mocker: MockerFixture, + ) -> None: + config_overrides(RBAC_ENABLED=source_access != "rbac-disabled") + monkeypatch.setattr(rbac_locators, "application_services", lambda: rbac_domain) + sqlite_session.add( + make_app(app_id="source-app", maintainer="account-1" if source_access == "maintainer" else None) + ) + sqlite_session.commit() + monkeypatch.setattr(db, "session", sqlite_session) + check_access = mocker.patch.object(RBACService.CheckAccess, "check", return_value=source_access == "granted") + workflow_runs = workflow_run_service + child_traces = [ + WorkflowNodeExecutionTrace( + id="child-row", + tenant_id="tenant-1", + app_id="source-app", + workflow_id="source-workflow", + triggered_from="workflow-tool", + node_execution_id="source-execution", + created_by="account-1", + outputs={"answer": "world"}, + ) + ] + get_children = mocker.patch.object( + workflow_runs, "get_workflow_tool_node_executions", return_value=child_traces + ) + _workflow_run_api_mock_application_services(monkeypatch, workflow_runs) + context = _workflow_run_api_request_context() + api = workflow_run_module.WorkflowToolNodeExecutionListApi() + handler = unwrap(api.get) + with app.test_request_context("/apps/app-1/workflow-runs/run-1/node-executions/parent/children"): + response = handler(api, context, app_model=_app(), run_id="run-1", node_execution_id="parent") + get_children.return_value = None + with pytest.raises(NotFoundError, match="Workflow run not found"): + handler(api, context, app_model=_app(), run_id="run-1", node_execution_id="parent") + assert response["data"][0]["node_execution_id"] == "source-execution" + assert response["data"][0]["outputs"] == {"answer": "world"} + if source_access == "granted": + check_access.assert_called_once_with( + "tenant-1", + "account-1", + scene=RBACPermission.APP_CREATE_AND_MANAGEMENT, + resource_type=RBACResourceScope.APP, + resource_id="source-app", + ) + else: + check_access.assert_not_called() + get_children.assert_called_with(context, app_id="app-1", run_id="run-1", node_execution_id="parent") + + @staticmethod + @pytest.mark.parametrize("source_app_ids", [("source-app",), ("allowed-source", "source-app")]) + def test_workflow_tool_children_deny_any_unauthorized_source( + app: Flask, + monkeypatch: pytest.MonkeyPatch, + sqlite_session: Session, + rbac_domain: RBACDomain, + config_overrides: Callable[..., None], + source_app_ids: tuple[str, ...], + workflow_run_service: WorkflowRunService, + mocker: MockerFixture, + ) -> None: + config_overrides(RBAC_ENABLED=True) + monkeypatch.setattr(rbac_locators, "application_services", lambda: rbac_domain) + sqlite_session.add_all(make_app(app_id=app_id) for app_id in ("app-1", "allowed-source", "source-app")) + sqlite_session.commit() + monkeypatch.setattr(db, "session", sqlite_session) + check_access = mocker.patch.object( + RBACService.CheckAccess, + "check", + side_effect=lambda *_args, resource_id, **_kwargs: resource_id != "source-app", + ) + workflow_runs = workflow_run_service + child_traces = [ + WorkflowNodeExecutionTrace( + id=f"{app_id}-row", + tenant_id="tenant-1", + app_id=app_id, + workflow_id=f"{app_id}-workflow", + triggered_from="workflow-tool", + created_by="account-1", + inputs={"secret": "source input"}, + process_data={"secret": "source internals"}, + outputs={"secret": "source output"}, + ) + for app_id in source_app_ids + ] + get_children = mocker.patch.object( + workflow_runs, "get_workflow_tool_node_executions", return_value=child_traces + ) + _workflow_run_api_mock_application_services(monkeypatch, workflow_runs) + context = _workflow_run_api_request_context() + api = workflow_run_module.WorkflowToolNodeExecutionListApi() + handler = unwrap(api.get) + + with app.test_request_context("/apps/app-1/workflow-runs/run-1/node-executions/parent/children"): + request.view_args = {"app_id": "app-1"} + with pytest.raises(Forbidden, match="permission to view this tool's internal execution details") as error: + handler(api, context, app_model=_app(), run_id="run-1", node_execution_id="parent") + + assert error.value.code == 403 + assert check_access.call_args.kwargs["resource_id"] == "source-app" + get_children.assert_called_once_with(context, app_id="app-1", run_id="run-1", node_execution_id="parent") diff --git a/api/tests/unit_tests/controllers/console/app/test_workflow_run_api.py b/api/tests/unit_tests/controllers/console/app/test_workflow_run_api.py deleted file mode 100644 index ef35887ae30..00000000000 --- a/api/tests/unit_tests/controllers/console/app/test_workflow_run_api.py +++ /dev/null @@ -1,513 +0,0 @@ -from __future__ import annotations - -import json -from collections.abc import Callable -from datetime import UTC, datetime -from inspect import unwrap -from types import SimpleNamespace -from typing import Any -from unittest.mock import Mock - -import pytest -from flask import Flask, request -from flask_restx import marshal -from pytest_mock import MockerFixture -from sqlalchemy.orm import Session, sessionmaker -from werkzeug.exceptions import Forbidden - -from controllers.common.errors import NotFoundError -from controllers.common.rbac import locators as rbac_locators -from controllers.console.app import workflow_run as workflow_run_module -from core.rbac import RBACPermission, RBACResourceScope -from extensions.ext_database import db -from fields.workflow_run_fields import node_execution_response_source -from graphon.enums import WorkflowExecutionStatus, WorkflowNodeExecutionStatus -from machinery.context import RequestContext -from models import Account, App, AppMode -from models.enums import CreatorUserRole, WorkflowRunTriggeredFrom -from models.workflow import ( - WorkflowNodeExecutionModel, - WorkflowNodeExecutionTriggeredFrom, - WorkflowRun, - WorkflowType, -) -from repositories.sqlalchemy_api_workflow_node_execution_repository import ( - DifyAPISQLAlchemyWorkflowNodeExecutionRepository, -) -from repositories.sqlalchemy_api_workflow_run_repository import DifyAPISQLAlchemyWorkflowRunRepository -from services.enterprise.rbac_service import RBACService -from services.workflow_node_execution_trace_service import WorkflowNodeExecutionTrace -from services.workflow_run_service import WorkflowRunService -from tests.unit_tests.model_factories import make_account, make_app -from tests.unit_tests.rbac_fakes import RBACDomain - - -def _serialize_200_response(handler, payload: Any) -> Any: - response_doc = getattr(handler, "__apidoc__", {}).get("responses", {}).get("200") - if response_doc is None: - return payload - - response_model = response_doc[1] - if isinstance(response_model, dict): - return marshal(payload, response_model) - return payload - - -def _account(session: Session) -> Account: - account = make_account(name="Alice", email="alice@example.com") - session.add(account) - session.commit() - return account - - -def _app() -> App: - return make_app(name="Workflow run app", mode=AppMode.WORKFLOW) - - -def _request_context(*, workspace_id: str = "tenant-1") -> RequestContext: - return RequestContext( - request_id="request-1", - trace_id="trace-1", - account_id="account-1", - active_workspace_id=workspace_id, - ) - - -def _mock_application_services(monkeypatch: pytest.MonkeyPatch, workflow_runs: Mock | WorkflowRunService) -> None: - monkeypatch.setattr( - workflow_run_module, - "application_services", - lambda: SimpleNamespace(workflow_runs=workflow_runs), - ) - - -def _workflow_run_summary(session: Session, **overrides: object) -> WorkflowRun: - created_at = datetime(2026, 1, 2, 3, 4, 5, tzinfo=UTC) - workflow_run = WorkflowRun( - id="run-1", - tenant_id="tenant-1", - app_id="app-1", - workflow_id="workflow-1", - type=WorkflowType.WORKFLOW, - triggered_from=WorkflowRunTriggeredFrom.DEBUGGING, - version="v1", - graph='{"nodes": []}', - inputs='{"query": "hello"}', - status=WorkflowExecutionStatus.SUCCEEDED, - outputs='{"answer": "world"}', - error=None, - elapsed_time=1.5, - total_tokens=10, - total_steps=2, - created_by_role=CreatorUserRole.ACCOUNT, - created_by="account-1", - created_at=created_at, - finished_at=created_at, - exceptions_count=0, - ) - for name, value in overrides.items(): - setattr(workflow_run, name, value) - workflow_run.retry_index = 0 - session.add(workflow_run) - session.commit() - return workflow_run - - -def _workflow_run_node_execution(session: Session) -> WorkflowNodeExecutionModel: - created_at = datetime(2026, 1, 2, 3, 4, 5, tzinfo=UTC) - execution = WorkflowNodeExecutionModel( - id="node-exec-1", - tenant_id="tenant-1", - app_id="app-1", - workflow_id="workflow-1", - triggered_from=WorkflowNodeExecutionTriggeredFrom.WORKFLOW_RUN, - workflow_run_id="run-1", - index=1, - predecessor_node_id=None, - node_execution_id="node-execution-1", - node_id="node-1", - node_type="start", - title="Start", - agent_workspace_binding_id=None, - inputs=json.dumps({"query": "hello"}), - process_data=json.dumps({"step": "prepared"}), - outputs=json.dumps({"answer": "world"}), - status=WorkflowNodeExecutionStatus.SUCCEEDED, - error=None, - elapsed_time=1.0, - execution_metadata=json.dumps({"total_tokens": 3}), - created_at=created_at, - created_by_role=CreatorUserRole.ACCOUNT, - created_by="account-1", - finished_at=created_at, - ) - execution.offload_data = [] - session.add(execution) - session.commit() - return execution - - -def test_workflow_run_list_returns_frontend_history_contract( - app: Flask, monkeypatch: pytest.MonkeyPatch, sqlite_session: Session -) -> None: - _account(sqlite_session) - workflow_run = _workflow_run_summary(sqlite_session) - workflow_runs = Mock() - workflow_runs.get_paginate_workflow_runs.return_value = SimpleNamespace( - limit=10, - has_more=False, - data=[workflow_run], - ) - _mock_application_services(monkeypatch, workflow_runs) - request_context = _request_context() - - api = workflow_run_module.WorkflowRunListApi() - handler = unwrap(api.get) - - with app.test_request_context("/apps/app-1/workflow-runs?limit=10", method="GET"): - payload = handler( - api, - sqlite_session, - workflow_run_module.WorkflowRunListQuery(limit=10), - request_context, - app_model=_app(), - ) - - response = _serialize_200_response(api.get, payload) - - assert response["limit"] == 10 - assert response["has_more"] is False - assert response["data"][0] == { - "id": "run-1", - "version": "v1", - "status": "succeeded", - "elapsed_time": 1.5, - "total_tokens": 10, - "total_steps": 2, - "created_by_account": {"id": "account-1", "name": "Alice", "email": "alice@example.com"}, - "created_at": 1767323045, - "finished_at": 1767323045, - "exceptions_count": 0, - "retry_index": 0, - } - workflow_runs.get_paginate_workflow_runs.assert_called_once_with( - request_context, - app_id="app-1", - args={"limit": 10}, - triggered_from=WorkflowRunTriggeredFrom.DEBUGGING, - ) - - -def test_advanced_chat_workflow_run_list_keeps_message_fields( - app: Flask, monkeypatch: pytest.MonkeyPatch, sqlite_session: Session -) -> None: - _account(sqlite_session) - workflow_run = _workflow_run_summary( - sqlite_session, - conversation_id="conversation-1", - message_id="message-1", - ) - workflow_runs = Mock() - workflow_runs.get_paginate_advanced_chat_workflow_runs.return_value = SimpleNamespace( - limit=1, - has_more=True, - data=[workflow_run], - ) - _mock_application_services(monkeypatch, workflow_runs) - request_context = _request_context() - - api = workflow_run_module.AdvancedChatAppWorkflowRunListApi() - handler = unwrap(api.get) - - with app.test_request_context("/apps/app-1/advanced-chat/workflow-runs?limit=1", method="GET"): - payload = handler( - api, - sqlite_session, - workflow_run_module.WorkflowRunListQuery(limit=1), - request_context, - app_model=_app(), - ) - - response = _serialize_200_response(api.get, payload) - - assert response["data"][0]["conversation_id"] == "conversation-1" - assert response["data"][0]["message_id"] == "message-1" - workflow_runs.get_paginate_advanced_chat_workflow_runs.assert_called_once_with( - request_context, - app_id="app-1", - args={"limit": 1}, - triggered_from=WorkflowRunTriggeredFrom.DEBUGGING, - ) - - -def test_workflow_run_count_passes_filters_to_application_service(app: Flask, monkeypatch: pytest.MonkeyPatch) -> None: - workflow_runs = Mock() - workflow_runs.get_workflow_runs_count.return_value = { - "total": 2, - "running": 0, - "succeeded": 2, - "failed": 0, - "stopped": 0, - "partial-succeeded": 0, - } - _mock_application_services(monkeypatch, workflow_runs) - request_context = _request_context() - query = workflow_run_module.WorkflowRunCountQuery( - status="succeeded", - time_range="7d", - triggered_from="app-run", - ) - - api = workflow_run_module.WorkflowRunCountApi() - handler = unwrap(api.get) - with app.test_request_context("/apps/app-1/workflow-runs/count", method="GET"): - payload = handler(api, query, request_context, app_model=_app()) - - assert payload == { - "total": 2, - "running": 0, - "succeeded": 2, - "failed": 0, - "stopped": 0, - "partial_succeeded": 0, - } - workflow_runs.get_workflow_runs_count.assert_called_once_with( - request_context, - app_id="app-1", - status="succeeded", - time_range="7d", - triggered_from=WorkflowRunTriggeredFrom.APP_RUN, - ) - - -def test_workflow_run_detail_returns_frontend_detail_contract( - app: Flask, monkeypatch: pytest.MonkeyPatch, sqlite_session: Session -) -> None: - _account(sqlite_session) - workflow_run = _workflow_run_summary(sqlite_session) - workflow_runs = Mock() - workflow_runs.get_workflow_run.return_value = workflow_run - _mock_application_services(monkeypatch, workflow_runs) - request_context = _request_context() - - api = workflow_run_module.WorkflowRunDetailApi() - handler = unwrap(api.get) - - with app.test_request_context("/apps/app-1/workflow-runs/run-1", method="GET"): - payload = handler(api, sqlite_session, request_context, app_model=_app(), run_id="run-1") - - response = _serialize_200_response(api.get, payload) - - assert response == { - "id": "run-1", - "version": "v1", - "graph": {"nodes": []}, - "inputs": {"query": "hello"}, - "status": "succeeded", - "outputs": {"answer": "world"}, - "error": None, - "elapsed_time": 1.5, - "total_tokens": 10, - "total_steps": 2, - "created_by_role": "account", - "created_by_account": {"id": "account-1", "name": "Alice", "email": "alice@example.com"}, - "created_by_end_user": None, - "created_at": 1767323045, - "finished_at": 1767323045, - "exceptions_count": 0, - } - workflow_runs.get_workflow_run.assert_called_once_with( - request_context, - app_id="app-1", - run_id="run-1", - ) - - -def test_workflow_run_detail_maps_missing_run_to_not_found( - app: Flask, monkeypatch: pytest.MonkeyPatch, sqlite_session: Session -) -> None: - workflow_runs = Mock() - workflow_runs.get_workflow_run.return_value = None - _mock_application_services(monkeypatch, workflow_runs) - request_context = _request_context(workspace_id="tenant-2") - api = workflow_run_module.WorkflowRunDetailApi() - handler = unwrap(api.get) - - with app.test_request_context("/apps/app-1/workflow-runs/run-1", method="GET"): - with pytest.raises(NotFoundError, match="Workflow run not found"): - handler(api, sqlite_session, request_context, app_model=_app(), run_id="run-1") - - workflow_runs.get_workflow_run.assert_called_once_with( - request_context, - app_id="app-1", - run_id="run-1", - ) - - -def test_workflow_run_node_executions_return_frontend_trace_contract( - app: Flask, monkeypatch: pytest.MonkeyPatch, sqlite_session: Session -) -> None: - _account(sqlite_session) - execution = _workflow_run_node_execution(sqlite_session) - workflow_runs = Mock() - workflow_runs.get_workflow_run_node_executions.return_value = [ - node_execution_response_source(execution, session=sqlite_session) - ] - _mock_application_services(monkeypatch, workflow_runs) - monkeypatch.setattr(db, "session", lambda: sqlite_session) - request_context = _request_context() - - api = workflow_run_module.WorkflowRunNodeExecutionListApi() - handler = unwrap(api.get) - - with app.test_request_context("/apps/app-1/workflow-runs/run-1/node-executions", method="GET"): - payload = handler(api, request_context, app_model=_app(), run_id="run-1") - - response = _serialize_200_response(api.get, payload) - - assert response == { - "data": [ - { - "id": "node-exec-1", - "index": 1, - "predecessor_node_id": None, - "node_id": "node-1", - "node_type": "start", - "title": "Start", - "inputs": {"query": "hello"}, - "process_data": {"step": "prepared"}, - "outputs": {"answer": "world"}, - "status": "succeeded", - "error": None, - "elapsed_time": 1.0, - "execution_metadata": {"total_tokens": 3}, - "extras": {}, - "created_at": 1767323045, - "created_by_role": "account", - "created_by_account": {"id": "account-1", "name": "Alice", "email": "alice@example.com"}, - "created_by_end_user": None, - "finished_at": 1767323045, - "inputs_truncated": False, - "outputs_truncated": False, - "process_data_truncated": False, - } - ] - } - workflow_runs.get_workflow_run_node_executions.assert_called_once_with( - request_context, - app_id="app-1", - run_id="run-1", - ) - - -@pytest.fixture -def workflow_run_service(sqlite_session_factory: sessionmaker[Session]) -> WorkflowRunService: - return WorkflowRunService( - workflow_runs=DifyAPISQLAlchemyWorkflowRunRepository(sqlite_session_factory), - node_executions=DifyAPISQLAlchemyWorkflowNodeExecutionRepository(sqlite_session_factory), - ) - - -@pytest.mark.parametrize("source_access", ["granted", "maintainer", "rbac-disabled"]) -def test_workflow_tool_children_preserve_source_identity( - app: Flask, - monkeypatch: pytest.MonkeyPatch, - sqlite_session: Session, - rbac_domain: RBACDomain, - config_overrides: Callable[..., None], - source_access: str, - workflow_run_service: WorkflowRunService, - mocker: MockerFixture, -) -> None: - config_overrides(RBAC_ENABLED=source_access != "rbac-disabled") - monkeypatch.setattr(rbac_locators, "application_services", lambda: rbac_domain) - sqlite_session.add(make_app(app_id="source-app", maintainer="account-1" if source_access == "maintainer" else None)) - sqlite_session.commit() - monkeypatch.setattr(db, "session", sqlite_session) - check_access = mocker.patch.object(RBACService.CheckAccess, "check", return_value=source_access == "granted") - workflow_runs = workflow_run_service - child_traces = [ - WorkflowNodeExecutionTrace( - id="child-row", - tenant_id="tenant-1", - app_id="source-app", - workflow_id="source-workflow", - triggered_from="workflow-tool", - node_execution_id="source-execution", - created_by="account-1", - outputs={"answer": "world"}, - ) - ] - get_children = mocker.patch.object(workflow_runs, "get_workflow_tool_node_executions", return_value=child_traces) - _mock_application_services(monkeypatch, workflow_runs) - context = _request_context() - api = workflow_run_module.WorkflowToolNodeExecutionListApi() - handler = unwrap(api.get) - with app.test_request_context("/apps/app-1/workflow-runs/run-1/node-executions/parent/children"): - response = handler(api, context, app_model=_app(), run_id="run-1", node_execution_id="parent") - get_children.return_value = None - with pytest.raises(NotFoundError, match="Workflow run not found"): - handler(api, context, app_model=_app(), run_id="run-1", node_execution_id="parent") - assert response["data"][0]["node_execution_id"] == "source-execution" - assert response["data"][0]["outputs"] == {"answer": "world"} - if source_access == "granted": - check_access.assert_called_once_with( - "tenant-1", - "account-1", - scene=RBACPermission.APP_CREATE_AND_MANAGEMENT, - resource_type=RBACResourceScope.APP, - resource_id="source-app", - ) - else: - check_access.assert_not_called() - get_children.assert_called_with(context, app_id="app-1", run_id="run-1", node_execution_id="parent") - - -@pytest.mark.parametrize("source_app_ids", [("source-app",), ("allowed-source", "source-app")]) -def test_workflow_tool_children_deny_any_unauthorized_source( - app: Flask, - monkeypatch: pytest.MonkeyPatch, - sqlite_session: Session, - rbac_domain: RBACDomain, - config_overrides: Callable[..., None], - source_app_ids: tuple[str, ...], - workflow_run_service: WorkflowRunService, - mocker: MockerFixture, -) -> None: - config_overrides(RBAC_ENABLED=True) - monkeypatch.setattr(rbac_locators, "application_services", lambda: rbac_domain) - sqlite_session.add_all(make_app(app_id=app_id) for app_id in ("app-1", "allowed-source", "source-app")) - sqlite_session.commit() - monkeypatch.setattr(db, "session", sqlite_session) - check_access = mocker.patch.object( - RBACService.CheckAccess, "check", side_effect=lambda *_args, resource_id, **_kwargs: resource_id != "source-app" - ) - workflow_runs = workflow_run_service - child_traces = [ - WorkflowNodeExecutionTrace( - id=f"{app_id}-row", - tenant_id="tenant-1", - app_id=app_id, - workflow_id=f"{app_id}-workflow", - triggered_from="workflow-tool", - created_by="account-1", - inputs={"secret": "source input"}, - process_data={"secret": "source internals"}, - outputs={"secret": "source output"}, - ) - for app_id in source_app_ids - ] - get_children = mocker.patch.object(workflow_runs, "get_workflow_tool_node_executions", return_value=child_traces) - _mock_application_services(monkeypatch, workflow_runs) - context = _request_context() - api = workflow_run_module.WorkflowToolNodeExecutionListApi() - handler = unwrap(api.get) - - with app.test_request_context("/apps/app-1/workflow-runs/run-1/node-executions/parent/children"): - request.view_args = {"app_id": "app-1"} - with pytest.raises(Forbidden, match="permission to view this tool's internal execution details") as error: - handler(api, context, app_model=_app(), run_id="run-1", node_execution_id="parent") - - assert error.value.code == 403 - assert check_access.call_args.kwargs["resource_id"] == "source-app" - get_children.assert_called_once_with(context, app_id="app-1", run_id="run-1", node_execution_id="parent") diff --git a/api/tests/unit_tests/controllers/console/app/test_workflow_trigger_api.py b/api/tests/unit_tests/controllers/console/app/test_workflow_trigger.py similarity index 99% rename from api/tests/unit_tests/controllers/console/app/test_workflow_trigger_api.py rename to api/tests/unit_tests/controllers/console/app/test_workflow_trigger.py index 6a26681f11a..d9ca56a28b0 100644 --- a/api/tests/unit_tests/controllers/console/app/test_workflow_trigger_api.py +++ b/api/tests/unit_tests/controllers/console/app/test_workflow_trigger.py @@ -34,7 +34,7 @@ from models.engine import db from models.enums import AppTriggerStatus, AppTriggerType from models.model import App, AppMode from models.trigger import AppTrigger, WorkflowWebhookTrigger -from tests.unit_tests.config_override import config_overrides_context +from tests.helpers.unit.config_override import config_overrides_context _CREATED_AT = datetime(2026, 1, 2, 3, 4, 5) diff --git a/api/tests/unit_tests/controllers/console/auth/test_account_activation.py b/api/tests/unit_tests/controllers/console/auth/test_activate.py similarity index 100% rename from api/tests/unit_tests/controllers/console/auth/test_account_activation.py rename to api/tests/unit_tests/controllers/console/auth/test_activate.py diff --git a/api/tests/unit_tests/controllers/console/auth/test_data_source_oauth.py b/api/tests/unit_tests/controllers/console/auth/test_data_source_oauth.py index e47876a7c2e..1864f2191a1 100644 --- a/api/tests/unit_tests/controllers/console/auth/test_data_source_oauth.py +++ b/api/tests/unit_tests/controllers/console/auth/test_data_source_oauth.py @@ -19,7 +19,7 @@ from services.data_source.oauth_service import ( InvalidDataSourceOAuthCodeError, InvalidDataSourceOAuthProviderError, ) -from tests.unit_tests.config_override import config_overrides_context +from tests.helpers.unit.config_override import config_overrides_context def _request_context() -> RequestContext: diff --git a/api/tests/unit_tests/controllers/console/auth/test_forgot_password.py b/api/tests/unit_tests/controllers/console/auth/test_forgot_password.py index f2d2a9ec4ef..0e797b992aa 100644 --- a/api/tests/unit_tests/controllers/console/auth/test_forgot_password.py +++ b/api/tests/unit_tests/controllers/console/auth/test_forgot_password.py @@ -20,7 +20,7 @@ from controllers.console.error import AccountNotFound, EmailSendIpLimitError from enums import DeploymentEdition from services import account_errors from services.entities.account_entities import ForgotPasswordVerification -from tests.unit_tests.config_override import apply_config_overrides +from tests.helpers.unit.config_override import apply_config_overrides class FakeForgotPasswordService: diff --git a/api/tests/unit_tests/controllers/console/auth/test_login_logout.py b/api/tests/unit_tests/controllers/console/auth/test_login.py similarity index 100% rename from api/tests/unit_tests/controllers/console/auth/test_login_logout.py rename to api/tests/unit_tests/controllers/console/auth/test_login.py diff --git a/api/tests/unit_tests/controllers/console/datasets/rag_pipeline/test_datasource_auth.py b/api/tests/unit_tests/controllers/console/datasets/rag_pipeline/test_datasource_auth.py index 4f4e850d29b..6045b0b1816 100644 --- a/api/tests/unit_tests/controllers/console/datasets/rag_pipeline/test_datasource_auth.py +++ b/api/tests/unit_tests/controllers/console/datasets/rag_pipeline/test_datasource_auth.py @@ -36,7 +36,7 @@ from repositories.credentials.query_repository import CredentialQueryRepository from services.data_source.provider_service import DatasourceProviderService from services.plugin.dispatch_identity import PluginOAuthClientConfig from services.plugin.oauth_service import OAuthProxyService -from tests.unit_tests.model_factories import make_account +from tests.helpers.unit.model_factories import make_account _PROVIDER_ID = "langgenius/notion_datasource/notion" diff --git a/api/tests/unit_tests/controllers/console/datasets/rag_pipeline/test_datasource_catalog.py b/api/tests/unit_tests/controllers/console/datasets/rag_pipeline/test_datasource_catalog.py deleted file mode 100644 index cc2029a4a75..00000000000 --- a/api/tests/unit_tests/controllers/console/datasets/rag_pipeline/test_datasource_catalog.py +++ /dev/null @@ -1,204 +0,0 @@ -from inspect import unwrap -from unittest.mock import patch - -import pytest -from flask import Flask -from pydantic import ValidationError - -from controllers.console.datasets.rag_pipeline import rag_pipeline_workflow as workflow_module -from controllers.console.datasets.rag_pipeline.rag_pipeline_workflow import ( - DatasourceListApi, - RagPipelineDatasourceListResponse, - RagPipelineDatasourceProviderResponse, -) -from core.datasource.entities.datasource_entities import DatasourceParameter -from core.plugin.entities import OAuthSchema -from core.plugin.entities.plugin_daemon import PluginDatasourceProviderEntity -from core.plugin.impl.datasource import PluginDatasourceManager -from services.rag_pipeline.rag_pipeline_manage_service import RagPipelineManageService - - -def _provider() -> PluginDatasourceProviderEntity: - return PluginDatasourceProviderEntity.model_validate( - { - "provider": "documents", - "plugin_id": "example/documents", - "plugin_unique_identifier": "example/documents:0.0.1@checksum", - "declaration": { - "identity": { - "author": "example", - "name": "documents", - "label": {"en_US": "Documents"}, - "description": {"en_US": "Import documents"}, - "icon": "icon.svg", - "tags": None, - }, - "provider_type": "online_document", - "credentials_schema": [{"type": "secret-input", "name": "token", "required": True}], - "datasources": [ - { - "identity": { - "author": "example", - "name": "pages", - "provider": "documents", - "label": {"en_US": "Pages"}, - }, - "description": {"en_US": "Import pages"}, - "parameters": [ - { - "name": "limit", - "type": "number", - "label": {"en_US": "Limit"}, - "description": {"en_US": "Maximum pages"}, - "required": True, - "default": 0, - "min": 0, - "max": 100, - }, - { - "name": "archived", - "type": "boolean", - "label": {"en_US": "Archived"}, - "description": {"en_US": "Include archived pages"}, - "required": True, - "default": False, - }, - ], - "output_schema": { - "type": "object", - "properties": {"page": {"anyOf": [{"type": "string"}, {"type": "null"}]}}, - "additionalProperties": False, - }, - } - ], - }, - } - ) - - -@pytest.mark.parametrize("credentials", [None, {"token": "secret"}]) -def test_catalog_serializes_real_plugin_and_local_file_providers( - app: Flask, credentials: dict[str, str] | None -) -> None: - provider = _provider() - api = DatasourceListApi() - - with ( - app.test_request_context("/rag/pipelines/datasource-plugins"), - patch.object( - PluginDatasourceManager, "_request_with_plugin_daemon_response", return_value=[provider] - ) as request, - patch.object( - workflow_module.application_services().data_sources.providers, - "get_datasource_credentials", - return_value=credentials, - ) as get_credentials, - ): - result = unwrap(api.get)(api, "tenant-1") - - local_file, remote = result - assert local_file["provider"] == "file" - assert local_file["declaration"]["provider_type"] == "local_file" - assert local_file["declaration"]["identity"]["name"] == "langgenius/file/file" - assert local_file["is_authorized"] is True - assert local_file["declaration"]["datasources"][0]["identity"]["provider"] == "langgenius/file/file" - assert local_file["declaration"]["datasources"][0]["output_schema"] is None - assert remote == provider.model_dump(mode="json") - assert remote["provider"] == "documents" - assert remote["plugin_unique_identifier"] == "example/documents:0.0.1@checksum" - assert remote["is_authorized"] is bool(credentials) - identity = remote["declaration"]["identity"] - assert identity["name"] == "example/documents/documents" - assert identity["icon"].endswith("/console/api/workspaces/current/plugin/icon?tenant_id=tenant-1&filename=icon.svg") - assert identity["tags"] is None - assert identity["label"]["pt_BR"] == "Documents" - datasource = remote["declaration"]["datasources"][0] - assert datasource["identity"]["provider"] == "example/documents/documents" - assert datasource["identity"]["icon"] is None - assert datasource["parameters"][0]["default"] == 0 - assert type(datasource["parameters"][0]["default"]) is int - assert datasource["parameters"][0]["min"] == 0 - assert datasource["parameters"][0]["max"] == 100 - assert datasource["parameters"][1]["default"] is False - assert datasource["output_schema"]["properties"]["page"] == {"anyOf": [{"type": "string"}, {"type": "null"}]} - assert datasource["output_schema"]["additionalProperties"] is False - request.assert_called_once() - assert request.call_args.args[:3] == ( - "GET", - "plugin/tenant-1/management/datasources", - list[PluginDatasourceProviderEntity], - ) - get_credentials.assert_called_once_with(tenant_id="tenant-1", provider="documents", plugin_id="example/documents") - - -def test_catalog_keeps_oauth_only_provider_unauthorized_without_credentials(app: Flask) -> None: - provider = _provider() - provider.declaration.credentials_schema = [] - provider.declaration.oauth_schema = OAuthSchema() - api = DatasourceListApi() - - with ( - app.test_request_context("/rag/pipelines/datasource-plugins"), - patch.object(PluginDatasourceManager, "fetch_datasource_providers", return_value=[provider]), - patch.object( - workflow_module.application_services().data_sources.providers, - "get_datasource_credentials", - return_value=None, - ) as get_credentials, - ): - result = unwrap(api.get)(api, "tenant-1") - - assert result[0]["is_authorized"] is False - assert result[0]["declaration"]["oauth_schema"] == {"client_schema": [], "credentials_schema": []} - get_credentials.assert_called_once_with(tenant_id="tenant-1", provider="documents", plugin_id="example/documents") - - -def test_catalog_keeps_provider_visible_after_credential_failure(app: Flask) -> None: - provider = _provider() - api = DatasourceListApi() - - with ( - app.test_request_context("/rag/pipelines/datasource-plugins"), - patch.object(PluginDatasourceManager, "fetch_datasource_providers", return_value=[provider]), - patch.object( - workflow_module.application_services().data_sources.providers, - "get_datasource_credentials", - side_effect=ValueError("decrypt failed"), - ), - ): - result = unwrap(api.get)(api, "tenant-1") - - assert result == [provider.model_dump(mode="json")] - assert result[0]["is_authorized"] is False - - -def test_catalog_serializes_empty_list(app: Flask) -> None: - api = DatasourceListApi() - - with ( - app.test_request_context("/rag/pipelines/datasource-plugins"), - patch.object(RagPipelineManageService, "list_rag_pipeline_datasources", return_value=[]), - ): - result = unwrap(api.get)(api, "tenant-1") - - assert result == [] - - -def test_catalog_list_schema_reuses_datasource_domain_entities() -> None: - schema = RagPipelineDatasourceListResponse.model_json_schema(mode="serialization") - - assert schema["type"] == "array" - assert schema["items"] == {"$ref": "#/$defs/RagPipelineDatasourceProviderResponse"} - declaration = schema["$defs"]["RagPipelineDatasourceProviderResponse"]["properties"]["declaration"] - assert declaration == {"$ref": "#/$defs/DatasourceProviderEntityWithPlugin"} - assert schema["$defs"]["DatasourceParameterType"]["enum"] == [ - value.value for value in DatasourceParameter.DatasourceParameterType - ] - - -def test_catalog_response_rejects_tool_only_parameter_types() -> None: - payload = _provider().model_dump(mode="json") - payload["declaration"]["datasources"][0]["parameters"][0]["type"] = "dynamic-select" - - with pytest.raises(ValidationError, match="string.*number.*boolean"): - RagPipelineDatasourceProviderResponse.model_validate(payload) diff --git a/api/tests/unit_tests/controllers/console/datasets/rag_pipeline/test_rag_pipeline.py b/api/tests/unit_tests/controllers/console/datasets/rag_pipeline/test_rag_pipeline.py index 6e00620aebd..8b9d692e757 100644 --- a/api/tests/unit_tests/controllers/console/datasets/rag_pipeline/test_rag_pipeline.py +++ b/api/tests/unit_tests/controllers/console/datasets/rag_pipeline/test_rag_pipeline.py @@ -27,9 +27,9 @@ from models.engine import db from services.entities.knowledge_entities.rag_pipeline_entities import PipelineTemplateInfoEntity from services.errors.base import NoPermissionError from services.errors.rag_pipeline import RagPipelineResourceNotFoundError -from tests.unit_tests.config_override import config_overrides_context -from tests.unit_tests.controllers.rbac_introspection import rbac_checks -from tests.unit_tests.model_factories import make_account +from tests.helpers.unit.config_override import config_overrides_context +from tests.helpers.unit.controllers.rbac_introspection import rbac_checks +from tests.helpers.unit.model_factories import make_account def _template_item() -> dict[str, object]: diff --git a/api/tests/unit_tests/controllers/console/datasets/rag_pipeline/test_rag_pipeline_datasets.py b/api/tests/unit_tests/controllers/console/datasets/rag_pipeline/test_rag_pipeline_datasets.py index 2b28b23d9c7..eef7b21225e 100644 --- a/api/tests/unit_tests/controllers/console/datasets/rag_pipeline/test_rag_pipeline_datasets.py +++ b/api/tests/unit_tests/controllers/console/datasets/rag_pipeline/test_rag_pipeline_datasets.py @@ -20,8 +20,8 @@ from controllers.console.datasets.rag_pipeline.rag_pipeline_datasets import ( from models.account import Account, TenantAccountRole from services.entities.dsl_entities import ImportStatus from services.rbac import contracts as rbac_contracts -from tests.unit_tests.config_override import config_overrides_context -from tests.unit_tests.model_factories import make_account +from tests.helpers.unit.config_override import config_overrides_context +from tests.helpers.unit.model_factories import make_account def _account(*, editor: bool) -> Account: diff --git a/api/tests/unit_tests/controllers/console/datasets/rag_pipeline/test_rag_pipeline_draft_variable.py b/api/tests/unit_tests/controllers/console/datasets/rag_pipeline/test_rag_pipeline_draft_variable.py index c5f43180cee..e9011cde1e4 100644 --- a/api/tests/unit_tests/controllers/console/datasets/rag_pipeline/test_rag_pipeline_draft_variable.py +++ b/api/tests/unit_tests/controllers/console/datasets/rag_pipeline/test_rag_pipeline_draft_variable.py @@ -31,8 +31,8 @@ from graphon.variables import IntegerVariable, StringVariable from models.account import Account, TenantAccountRole from models.dataset import Pipeline from models.workflow import Workflow, WorkflowDraftVariable, WorkflowType -from tests.unit_tests.controllers.rbac_introspection import rbac_checks -from tests.unit_tests.model_factories import make_workflow +from tests.helpers.unit.controllers.rbac_introspection import rbac_checks +from tests.helpers.unit.model_factories import make_workflow @pytest.fixture diff --git a/api/tests/unit_tests/controllers/console/datasets/rag_pipeline/test_rag_pipeline_import.py b/api/tests/unit_tests/controllers/console/datasets/rag_pipeline/test_rag_pipeline_import.py index 0e576120c46..7d1b885cc7e 100644 --- a/api/tests/unit_tests/controllers/console/datasets/rag_pipeline/test_rag_pipeline_import.py +++ b/api/tests/unit_tests/controllers/console/datasets/rag_pipeline/test_rag_pipeline_import.py @@ -27,7 +27,7 @@ from models.engine import db from services.enterprise import rbac_service from services.entities.dsl_entities import CheckDependenciesResult, ImportStatus from services.rag_pipeline.rag_pipeline_dsl_service import RagPipelineImportInfo -from tests.unit_tests.model_factories import make_account, make_tenant +from tests.helpers.unit.model_factories import make_account, make_tenant def _account() -> Account: diff --git a/api/tests/unit_tests/controllers/console/datasets/rag_pipeline/test_rag_pipeline_workflow.py b/api/tests/unit_tests/controllers/console/datasets/rag_pipeline/test_rag_pipeline_workflow.py index 5d1705d3696..9aa06e3b934 100644 --- a/api/tests/unit_tests/controllers/console/datasets/rag_pipeline/test_rag_pipeline_workflow.py +++ b/api/tests/unit_tests/controllers/console/datasets/rag_pipeline/test_rag_pipeline_workflow.py @@ -1,46 +1,285 @@ -"""Unit coverage for RAG workflow controllers using real models and disposable SQLite state.""" - from __future__ import annotations import json -from collections.abc import Iterator +from collections.abc import Generator, Iterator +from contextlib import contextmanager +from dataclasses import dataclass from datetime import datetime +from inspect import unwrap from inspect import unwrap as unwrap_all -from unittest.mock import patch -from uuid import UUID +from typing import TypedDict, Unpack +from unittest.mock import MagicMock, patch +from uuid import UUID, uuid4 import pytest from flask import Flask +from pydantic import ValidationError from sqlalchemy import Engine -from sqlalchemy.orm import Session -from werkzeug.exceptions import Forbidden, NotFound +from sqlalchemy.orm import Session, scoped_session, sessionmaker +from werkzeug.exceptions import BadRequest, Forbidden, HTTPException, NotFound +import models.workflow as workflow_models +import services +import services.rag_pipeline.rag_pipeline as rag_pipeline_service +from controllers.console.app.error import DraftWorkflowNotExist, DraftWorkflowNotSync from controllers.console.datasets.rag_pipeline import rag_pipeline_workflow as module +from controllers.console.datasets.rag_pipeline import rag_pipeline_workflow as workflow_controller +from controllers.console.datasets.rag_pipeline import rag_pipeline_workflow as workflow_module from controllers.console.datasets.rag_pipeline.rag_pipeline_workflow import ( + DatasourceListApi, + DatasourceVariablesPayload, + DefaultBlockConfigQuery, + DefaultRagPipelineBlockConfigApi, + DraftRagPipelineApi, NodeIdQuery, + NodeRunPayload, + NodeRunRequiredPayload, + PublishedAllRagPipelineApi, + PublishedRagPipelineApi, PublishedWorkflowRunPayload, + RagPipelineByIdApi, + RagPipelineDatasourceListResponse, + RagPipelineDatasourceProviderResponse, + RagPipelineDatasourceVariableApi, + RagPipelineDraftNodeRunApi, + RagPipelineDraftRunIterationNodeApi, + RagPipelineDraftRunLoopNodeApi, + RagPipelineDraftWorkflowRestoreApi, RagPipelineDraftWorkflowRunPayload, + RagPipelineRecommendedPluginApi, RagPipelineRecommendedPluginQuery, + RagPipelineTaskStopApi, + RagPipelineWorkflowLastRunApi, + RagPipelineWorkflowRunNodeExecutionListApi, WorkflowListQuery, WorkflowUpdatePayload, ) from controllers.web.error import InvokeRateLimitError as InvokeRateLimitHttpError from core.app.apps.pipeline.pipeline_generator import PipelineGenerator +from core.datasource.entities.datasource_entities import DatasourceParameter +from core.plugin.entities import OAuthSchema +from core.plugin.entities.plugin_daemon import PluginDatasourceProviderEntity +from core.plugin.impl.datasource import PluginDatasourceManager +from fields.workflow_run_fields import node_execution_response_source +from graphon.enums import WorkflowNodeExecutionStatus +from libs.datetime_utils import naive_utc_now from models.account import Account, TenantAccountRole from models.dataset import Dataset, Pipeline from models.engine import db -from models.enums import PermissionEnum +from models.enums import CreatorUserRole, PermissionEnum from models.tools import WorkflowToolProvider -from models.workflow import Workflow, WorkflowType +from models.workflow import Workflow, WorkflowNodeExecutionModel, WorkflowNodeExecutionTriggeredFrom, WorkflowType +from services.errors.app import IsDraftWorkflowError, WorkflowHashNotEqualError, WorkflowNotFoundError from services.errors.llm import InvokeRateLimitError from services.errors.rag_pipeline import RagPipelineResourceNotFoundError from services.rag_pipeline.rag_pipeline import RagPipelineService -from tests.unit_tests.config_override import config_overrides_context -from tests.unit_tests.model_factories import make_account, make_dataset, make_tenant +from services.rag_pipeline.rag_pipeline_manage_service import RagPipelineManageService +from tests.helpers.unit.config_override import config_overrides_context +from tests.helpers.unit.controllers.error_body import raised_error +from tests.helpers.unit.model_factories import make_account as _rag_pipeline_workflow_make_account +from tests.helpers.unit.model_factories import make_dataset, make_tenant -DEFAULT_WORKFLOW_TENANT_ID = "00000000-0000-0000-0000-000000000001" -DEFAULT_WORKFLOW_APP_ID = "00000000-0000-0000-0000-000000000002" -DEFAULT_WORKFLOW_CREATED_BY = "00000000-0000-0000-0000-000000000003" + +def _provider() -> PluginDatasourceProviderEntity: + return PluginDatasourceProviderEntity.model_validate( + { + "provider": "documents", + "plugin_id": "example/documents", + "plugin_unique_identifier": "example/documents:0.0.1@checksum", + "declaration": { + "identity": { + "author": "example", + "name": "documents", + "label": {"en_US": "Documents"}, + "description": {"en_US": "Import documents"}, + "icon": "icon.svg", + "tags": None, + }, + "provider_type": "online_document", + "credentials_schema": [{"type": "secret-input", "name": "token", "required": True}], + "datasources": [ + { + "identity": { + "author": "example", + "name": "pages", + "provider": "documents", + "label": {"en_US": "Pages"}, + }, + "description": {"en_US": "Import pages"}, + "parameters": [ + { + "name": "limit", + "type": "number", + "label": {"en_US": "Limit"}, + "description": {"en_US": "Maximum pages"}, + "required": True, + "default": 0, + "min": 0, + "max": 100, + }, + { + "name": "archived", + "type": "boolean", + "label": {"en_US": "Archived"}, + "description": {"en_US": "Include archived pages"}, + "required": True, + "default": False, + }, + ], + "output_schema": { + "type": "object", + "properties": {"page": {"anyOf": [{"type": "string"}, {"type": "null"}]}}, + "additionalProperties": False, + }, + } + ], + }, + } + ) + + +class TestDatasourceCatalogCases: + @staticmethod + @pytest.mark.parametrize("credentials", [None, {"token": "secret"}]) + def test_catalog_serializes_real_plugin_and_local_file_providers( + app: Flask, credentials: dict[str, str] | None + ) -> None: + provider = _provider() + api = DatasourceListApi() + + with ( + app.test_request_context("/rag/pipelines/datasource-plugins"), + patch.object( + PluginDatasourceManager, "_request_with_plugin_daemon_response", return_value=[provider] + ) as request, + patch.object( + workflow_module.application_services().data_sources.providers, + "get_datasource_credentials", + return_value=credentials, + ) as get_credentials, + ): + result = unwrap(api.get)(api, "tenant-1") + + local_file, remote = result + assert local_file["provider"] == "file" + assert local_file["declaration"]["provider_type"] == "local_file" + assert local_file["declaration"]["identity"]["name"] == "langgenius/file/file" + assert local_file["is_authorized"] is True + assert local_file["declaration"]["datasources"][0]["identity"]["provider"] == "langgenius/file/file" + assert local_file["declaration"]["datasources"][0]["output_schema"] is None + assert remote == provider.model_dump(mode="json") + assert remote["provider"] == "documents" + assert remote["plugin_unique_identifier"] == "example/documents:0.0.1@checksum" + assert remote["is_authorized"] is bool(credentials) + identity = remote["declaration"]["identity"] + assert identity["name"] == "example/documents/documents" + assert identity["icon"].endswith( + "/console/api/workspaces/current/plugin/icon?tenant_id=tenant-1&filename=icon.svg" + ) + assert identity["tags"] is None + assert identity["label"]["pt_BR"] == "Documents" + datasource = remote["declaration"]["datasources"][0] + assert datasource["identity"]["provider"] == "example/documents/documents" + assert datasource["identity"]["icon"] is None + assert datasource["parameters"][0]["default"] == 0 + assert type(datasource["parameters"][0]["default"]) is int + assert datasource["parameters"][0]["min"] == 0 + assert datasource["parameters"][0]["max"] == 100 + assert datasource["parameters"][1]["default"] is False + assert datasource["output_schema"]["properties"]["page"] == {"anyOf": [{"type": "string"}, {"type": "null"}]} + assert datasource["output_schema"]["additionalProperties"] is False + request.assert_called_once() + assert request.call_args.args[:3] == ( + "GET", + "plugin/tenant-1/management/datasources", + list[PluginDatasourceProviderEntity], + ) + get_credentials.assert_called_once_with( + tenant_id="tenant-1", provider="documents", plugin_id="example/documents" + ) + + @staticmethod + def test_catalog_keeps_oauth_only_provider_unauthorized_without_credentials(app: Flask) -> None: + provider = _provider() + provider.declaration.credentials_schema = [] + provider.declaration.oauth_schema = OAuthSchema() + api = DatasourceListApi() + + with ( + app.test_request_context("/rag/pipelines/datasource-plugins"), + patch.object(PluginDatasourceManager, "fetch_datasource_providers", return_value=[provider]), + patch.object( + workflow_module.application_services().data_sources.providers, + "get_datasource_credentials", + return_value=None, + ) as get_credentials, + ): + result = unwrap(api.get)(api, "tenant-1") + + assert result[0]["is_authorized"] is False + assert result[0]["declaration"]["oauth_schema"] == {"client_schema": [], "credentials_schema": []} + get_credentials.assert_called_once_with( + tenant_id="tenant-1", provider="documents", plugin_id="example/documents" + ) + + @staticmethod + def test_catalog_keeps_provider_visible_after_credential_failure(app: Flask) -> None: + provider = _provider() + api = DatasourceListApi() + + with ( + app.test_request_context("/rag/pipelines/datasource-plugins"), + patch.object(PluginDatasourceManager, "fetch_datasource_providers", return_value=[provider]), + patch.object( + workflow_module.application_services().data_sources.providers, + "get_datasource_credentials", + side_effect=ValueError("decrypt failed"), + ), + ): + result = unwrap(api.get)(api, "tenant-1") + + assert result == [provider.model_dump(mode="json")] + assert result[0]["is_authorized"] is False + + @staticmethod + def test_catalog_serializes_empty_list(app: Flask) -> None: + api = DatasourceListApi() + + with ( + app.test_request_context("/rag/pipelines/datasource-plugins"), + patch.object(RagPipelineManageService, "list_rag_pipeline_datasources", return_value=[]), + ): + result = unwrap(api.get)(api, "tenant-1") + + assert result == [] + + @staticmethod + def test_catalog_list_schema_reuses_datasource_domain_entities() -> None: + schema = RagPipelineDatasourceListResponse.model_json_schema(mode="serialization") + + assert schema["type"] == "array" + assert schema["items"] == {"$ref": "#/$defs/RagPipelineDatasourceProviderResponse"} + declaration = schema["$defs"]["RagPipelineDatasourceProviderResponse"]["properties"]["declaration"] + assert declaration == {"$ref": "#/$defs/DatasourceProviderEntityWithPlugin"} + assert schema["$defs"]["DatasourceParameterType"]["enum"] == [ + value.value for value in DatasourceParameter.DatasourceParameterType + ] + + @staticmethod + def test_catalog_response_rejects_tool_only_parameter_types() -> None: + payload = _provider().model_dump(mode="json") + payload["declaration"]["datasources"][0]["parameters"][0]["type"] = "dynamic-select" + + with pytest.raises(ValidationError, match="string.*number.*boolean"): + RagPipelineDatasourceProviderResponse.model_validate(payload) + + +"""Unit coverage for RAG workflow controllers using real models and disposable SQLite state.""" + + +_RAG_PIPELINE_WORKFLOW_DEFAULT_WORKFLOW_TENANT_ID = "00000000-0000-0000-0000-000000000001" +_RAG_PIPELINE_WORKFLOW_DEFAULT_WORKFLOW_APP_ID = "00000000-0000-0000-0000-000000000002" +_RAG_PIPELINE_WORKFLOW_DEFAULT_WORKFLOW_CREATED_BY = "00000000-0000-0000-0000-000000000003" DEFAULT_WORKFLOW_ID = "00000000-0000-0000-0000-000000000004" DEFAULT_DATASET_ID = "44444444-4444-4444-4444-444444444444" @@ -48,15 +287,15 @@ DEFAULT_DATASET_ID = "44444444-4444-4444-4444-444444444444" def _make_workflow(**overrides: object) -> Workflow: workflow = Workflow( id=DEFAULT_WORKFLOW_ID, - tenant_id=DEFAULT_WORKFLOW_TENANT_ID, - app_id=DEFAULT_WORKFLOW_APP_ID, + tenant_id=_RAG_PIPELINE_WORKFLOW_DEFAULT_WORKFLOW_TENANT_ID, + app_id=_RAG_PIPELINE_WORKFLOW_DEFAULT_WORKFLOW_APP_ID, type=WorkflowType.WORKFLOW, version=Workflow.VERSION_DRAFT, marked_name="Release 1", marked_comment="Initial release", graph=json.dumps({"nodes": [], "edges": []}), features=json.dumps({"file_upload": {"enabled": False}}), - created_by=DEFAULT_WORKFLOW_CREATED_BY, + created_by=_RAG_PIPELINE_WORKFLOW_DEFAULT_WORKFLOW_CREATED_BY, created_at=datetime(2024, 1, 1, 12, 0, 0), updated_by=None, updated_at=datetime(2024, 1, 1, 12, 1, 0), @@ -70,22 +309,28 @@ def _make_workflow(**overrides: object) -> Workflow: def _account() -> Account: - return make_account( - account_id=DEFAULT_WORKFLOW_CREATED_BY, + return _rag_pipeline_workflow_make_account( + account_id=_RAG_PIPELINE_WORKFLOW_DEFAULT_WORKFLOW_CREATED_BY, name="Alice", email="alice@example.com", role=TenantAccountRole.EDITOR, - tenant=make_tenant(tenant_id=DEFAULT_WORKFLOW_TENANT_ID, name="Tenant"), + tenant=make_tenant(tenant_id=_RAG_PIPELINE_WORKFLOW_DEFAULT_WORKFLOW_TENANT_ID, name="Tenant"), ) def _pipeline() -> Pipeline: - pipeline = Pipeline(tenant_id=DEFAULT_WORKFLOW_TENANT_ID, name="Pipeline", description="desc") - pipeline.id = DEFAULT_WORKFLOW_APP_ID + pipeline = Pipeline( + tenant_id=_RAG_PIPELINE_WORKFLOW_DEFAULT_WORKFLOW_TENANT_ID, name="Pipeline", description="desc" + ) + pipeline.id = _RAG_PIPELINE_WORKFLOW_DEFAULT_WORKFLOW_APP_ID return pipeline -def _dataset(*, tenant_id: str = DEFAULT_WORKFLOW_TENANT_ID, maintainer: str = DEFAULT_WORKFLOW_CREATED_BY) -> Dataset: +def _dataset( + *, + tenant_id: str = _RAG_PIPELINE_WORKFLOW_DEFAULT_WORKFLOW_TENANT_ID, + maintainer: str = _RAG_PIPELINE_WORKFLOW_DEFAULT_WORKFLOW_CREATED_BY, +) -> Dataset: return make_dataset( dataset_id=DEFAULT_DATASET_ID, tenant_id=tenant_id, @@ -102,371 +347,1273 @@ def _persist_workflow(workflow: Workflow) -> None: db.session.expunge(workflow) -@pytest.fixture -def database_app() -> Iterator[Flask]: - app = Flask(__name__) - app.config["SQLALCHEMY_DATABASE_URI"] = "sqlite:///:memory:" - db.init_app(app) +class TestRagPipelineWorkflowCases: + @staticmethod + @pytest.fixture + def database_app() -> Iterator[Flask]: + app = Flask(__name__) + app.config["SQLALCHEMY_DATABASE_URI"] = "sqlite:///:memory:" + db.init_app(app) - with app.app_context(): - db.metadata.tables[Account.__tablename__].create(db.engine) - db.metadata.tables[WorkflowToolProvider.__tablename__].create(db.engine) - db.metadata.tables[Workflow.__tablename__].create(db.engine) - db.session.add(_account()) - db.session.commit() + with app.app_context(): + db.metadata.tables[Account.__tablename__].create(db.engine) + db.metadata.tables[WorkflowToolProvider.__tablename__].create(db.engine) + db.metadata.tables[Workflow.__tablename__].create(db.engine) + db.session.add(_account()) + db.session.commit() - try: - yield app - finally: - db.session.remove() + try: + yield app + finally: + db.session.remove() + @staticmethod + @pytest.mark.usefixtures("database_app") + def test_draft_rag_pipeline_workflow_get_serializes_response_model() -> None: + workflow = _make_workflow() + expected_hash = workflow.unique_hash + _persist_workflow(workflow) -@pytest.mark.usefixtures("database_app") -def test_draft_rag_pipeline_workflow_get_serializes_response_model() -> None: - workflow = _make_workflow() - expected_hash = workflow.unique_hash - _persist_workflow(workflow) + api = module.DraftRagPipelineApi() + handler = unwrap_all(api.get) - api = module.DraftRagPipelineApi() - handler = unwrap_all(api.get) + response = handler(api, db.session(), _pipeline()) - response = handler(api, db.session(), _pipeline()) - - assert response["id"] == DEFAULT_WORKFLOW_ID - assert response["graph"] == {"nodes": [], "edges": []} - assert response["features"] == {"file_upload": {"enabled": False}} - assert response["hash"] == expected_hash - assert response["created_by"] == { - "id": DEFAULT_WORKFLOW_CREATED_BY, - "name": "Alice", - "email": "alice@example.com", - } - assert response["updated_by"] is None - assert response["created_at"] == int(datetime(2024, 1, 1, 12, 0, 0).timestamp()) - assert response["updated_at"] == int(datetime(2024, 1, 1, 12, 1, 0).timestamp()) - - -def test_published_rag_pipeline_workflows_serialize_items_before_session_closes( - database_app: Flask, -) -> None: - api = module.PublishedAllRagPipelineApi() - handler = unwrap_all(api.get) - workflow = _make_workflow(version="1") - _persist_workflow(workflow) - pipeline = _pipeline() - pipeline.workflow_id = DEFAULT_WORKFLOW_ID - - with database_app.test_request_context( - "/rag/pipelines/pipeline-1/workflows", - method="GET", - query_string={"page": 1, "limit": 10, "user_id": "", "named_only": "false"}, - ): - response = handler( - api, WorkflowListQuery(page=1, limit=10, user_id="", named_only=False), _account(), pipeline=pipeline - ) - - assert response["items"][0]["id"] == DEFAULT_WORKFLOW_ID - assert response["page"] == 1 - assert response["limit"] == 10 - assert response["has_more"] is False - - -def test_rag_pipeline_workflow_patch_serializes_response_model( - database_app: Flask, -) -> None: - workflow = _make_workflow(marked_name="Updated release") - expected_hash = workflow.unique_hash - _persist_workflow(workflow) - payload: dict[str, object] = {"marked_name": "Updated release"} - - api = module.RagPipelineByIdApi() - handler = unwrap_all(api.patch) - - with database_app.test_request_context( - f"/rag/pipelines/{DEFAULT_WORKFLOW_APP_ID}/workflows/{DEFAULT_WORKFLOW_ID}", method="PATCH", json=payload - ): - response = handler( - api, - WorkflowUpdatePayload.model_validate(payload), - _account(), - pipeline=_pipeline(), - workflow_id=DEFAULT_WORKFLOW_ID, - ) - - assert response["id"] == DEFAULT_WORKFLOW_ID - assert response["marked_name"] == "Updated release" - assert response["hash"] == expected_hash - - -@pytest.mark.usefixtures("database_app") -def test_default_rag_pipeline_block_configs_serializes_root_response() -> None: - block_configs = [{"type": "start", "config": {"title": "Start"}}] - - api = module.DefaultRagPipelineBlockConfigsApi() - handler = unwrap_all(api.get) - - with patch.object(RagPipelineService, "get_default_block_configs", return_value=block_configs): - response = handler(api, _pipeline()) - - assert response == block_configs - - -def test_draft_rag_pipeline_second_step_parameters_serializes_variables(database_app: Flask) -> None: - variables = [ - { - "belong_to_node_id": "shared", - "type": "number", - "label": "Chunk size", - "variable": "chunk_size", - "default_value": 1024, - "required": True, + assert response["id"] == DEFAULT_WORKFLOW_ID + assert response["graph"] == {"nodes": [], "edges": []} + assert response["features"] == {"file_upload": {"enabled": False}} + assert response["hash"] == expected_hash + assert response["created_by"] == { + "id": _RAG_PIPELINE_WORKFLOW_DEFAULT_WORKFLOW_CREATED_BY, + "name": "Alice", + "email": "alice@example.com", } - ] - api = module.DraftRagPipelineSecondStepApi() - handler = unwrap_all(api.get) + assert response["updated_by"] is None + assert response["created_at"] == int(datetime(2024, 1, 1, 12, 0, 0).timestamp()) + assert response["updated_at"] == int(datetime(2024, 1, 1, 12, 1, 0).timestamp()) - with ( - database_app.test_request_context("/?node_id=node-1"), - patch.object(RagPipelineService, "get_second_step_parameters", return_value=variables), - ): - response = handler(api, NodeIdQuery(node_id="node-1"), _pipeline()) + @staticmethod + def test_published_rag_pipeline_workflows_serialize_items_before_session_closes( + database_app: Flask, + ) -> None: + api = module.PublishedAllRagPipelineApi() + handler = unwrap_all(api.get) + workflow = _make_workflow(version="1") + _persist_workflow(workflow) + pipeline = _pipeline() + pipeline.workflow_id = DEFAULT_WORKFLOW_ID - assert response["variables"] == variables + with database_app.test_request_context( + "/rag/pipelines/pipeline-1/workflows", + method="GET", + query_string={"page": 1, "limit": 10, "user_id": "", "named_only": "false"}, + ): + response = handler( + api, WorkflowListQuery(page=1, limit=10, user_id="", named_only=False), _account(), pipeline=pipeline + ) + assert response["items"][0]["id"] == DEFAULT_WORKFLOW_ID + assert response["page"] == 1 + assert response["limit"] == 10 + assert response["has_more"] is False -def test_rag_pipeline_recommended_plugins_serializes_known_envelope(database_app: Flask) -> None: - recommended_plugins = { - "installed_recommended_plugins": [ + @staticmethod + def test_rag_pipeline_workflow_patch_serializes_response_model( + database_app: Flask, + ) -> None: + workflow = _make_workflow(marked_name="Updated release") + expected_hash = workflow.unique_hash + _persist_workflow(workflow) + payload: dict[str, object] = {"marked_name": "Updated release"} + + api = module.RagPipelineByIdApi() + handler = unwrap_all(api.patch) + + with database_app.test_request_context( + f"/rag/pipelines/{_RAG_PIPELINE_WORKFLOW_DEFAULT_WORKFLOW_APP_ID}/workflows/{DEFAULT_WORKFLOW_ID}", + method="PATCH", + json=payload, + ): + response = handler( + api, + WorkflowUpdatePayload.model_validate(payload), + _account(), + pipeline=_pipeline(), + workflow_id=DEFAULT_WORKFLOW_ID, + ) + + assert response["id"] == DEFAULT_WORKFLOW_ID + assert response["marked_name"] == "Updated release" + assert response["hash"] == expected_hash + + @staticmethod + @pytest.mark.usefixtures("database_app") + def test_default_rag_pipeline_block_configs_serializes_root_response() -> None: + block_configs = [{"type": "start", "config": {"title": "Start"}}] + + api = module.DefaultRagPipelineBlockConfigsApi() + handler = unwrap_all(api.get) + + with patch.object(RagPipelineService, "get_default_block_configs", return_value=block_configs): + response = handler(api, _pipeline()) + + assert response == block_configs + + @staticmethod + def test_draft_rag_pipeline_second_step_parameters_serializes_variables(database_app: Flask) -> None: + variables = [ { - "id": "langgenius/dify_extractor/dify_extractor", - "author": "langgenius", - "name": "langgenius/dify_extractor/dify_extractor", - "description": { - "en_US": "Extract text", - "zh_Hans": "Extract text", - "pt_BR": "Extract text", - "ja_JP": "Extract text", - }, - "icon": "/icon.svg", - "label": { - "en_US": "Dify Extractor", - "zh_Hans": "Dify Extractor", - "pt_BR": "Dify Extractor", - "ja_JP": "Dify Extractor", - }, - "type": "builtin", + "belong_to_node_id": "shared", + "type": "number", + "label": "Chunk size", + "variable": "chunk_size", + "default_value": 1024, + "required": True, } - ], - "uninstalled_recommended_plugins": [{"plugin_id": "langgenius/notion_datasource"}], - } - api = module.RagPipelineRecommendedPluginApi() - handler = unwrap_all(api.get) - - with ( - database_app.test_request_context("/?type=tool"), - patch.object(RagPipelineService, "get_recommended_plugins", return_value=recommended_plugins), - ): - response = handler(api, RagPipelineRecommendedPluginQuery(type="tool"), DEFAULT_WORKFLOW_TENANT_ID, _account()) - - assert response == recommended_plugins - - -def test_rag_pipeline_transform_rejects_read_only_member(sqlite_engine: Engine) -> None: - account = _account() - account.role = TenantAccountRole.NORMAL - api = module.RagPipelineTransformApi() - handler = unwrap_all(api.post) - - with Session(sqlite_engine) as session: - session.add(_dataset()) + ] + api = module.DraftRagPipelineSecondStepApi() + handler = unwrap_all(api.get) with ( - config_overrides_context(RBAC_ENABLED=False), - pytest.raises(Forbidden), + database_app.test_request_context("/?node_id=node-1"), + patch.object(RagPipelineService, "get_second_step_parameters", return_value=variables), ): - handler(api, session, DEFAULT_WORKFLOW_TENANT_ID, account, UUID(DEFAULT_DATASET_ID)) + response = handler(api, NodeIdQuery(node_id="node-1"), _pipeline()) + assert response["variables"] == variables -def test_rag_pipeline_transform_rejects_dataset_from_another_tenant_before_service_call( - sqlite_engine: Engine, -) -> None: - api = module.RagPipelineTransformApi() - handler = unwrap_all(api.post) - - with Session(sqlite_engine) as session: - session.add(_dataset(tenant_id="00000000-0000-0000-0000-000000000099")) + @staticmethod + def test_rag_pipeline_recommended_plugins_serializes_known_envelope(database_app: Flask) -> None: + recommended_plugins = { + "installed_recommended_plugins": [ + { + "id": "langgenius/dify_extractor/dify_extractor", + "author": "langgenius", + "name": "langgenius/dify_extractor/dify_extractor", + "description": { + "en_US": "Extract text", + "zh_Hans": "Extract text", + "pt_BR": "Extract text", + "ja_JP": "Extract text", + }, + "icon": "/icon.svg", + "label": { + "en_US": "Dify Extractor", + "zh_Hans": "Dify Extractor", + "pt_BR": "Dify Extractor", + "ja_JP": "Dify Extractor", + }, + "type": "builtin", + } + ], + "uninstalled_recommended_plugins": [{"plugin_id": "langgenius/notion_datasource"}], + } + api = module.RagPipelineRecommendedPluginApi() + handler = unwrap_all(api.get) with ( - patch.object(module.RagPipelineTransformService, "transform_dataset") as transform_dataset, - pytest.raises(NotFound), + database_app.test_request_context("/?type=tool"), + patch.object(RagPipelineService, "get_recommended_plugins", return_value=recommended_plugins), ): - handler(api, session, DEFAULT_WORKFLOW_TENANT_ID, _account(), UUID(DEFAULT_DATASET_ID)) + response = handler( + api, + RagPipelineRecommendedPluginQuery(type="tool"), + _RAG_PIPELINE_WORKFLOW_DEFAULT_WORKFLOW_TENANT_ID, + _account(), + ) - transform_dataset.assert_not_called() + assert response == recommended_plugins + @staticmethod + def test_rag_pipeline_transform_rejects_read_only_member(sqlite_engine: Engine) -> None: + account = _account() + account.role = TenantAccountRole.NORMAL + api = module.RagPipelineTransformApi() + handler = unwrap_all(api.post) -def test_rag_pipeline_transform_enforces_legacy_dataset_permission_before_service_call( - sqlite_engine: Engine, -) -> None: - api = module.RagPipelineTransformApi() - handler = unwrap_all(api.post) + with Session(sqlite_engine) as session: + session.add(_dataset()) - with Session(sqlite_engine) as session: - session.add(_dataset(maintainer="00000000-0000-0000-0000-000000000099")) + with ( + config_overrides_context(RBAC_ENABLED=False), + pytest.raises(Forbidden), + ): + handler( + api, session, _RAG_PIPELINE_WORKFLOW_DEFAULT_WORKFLOW_TENANT_ID, account, UUID(DEFAULT_DATASET_ID) + ) - with ( - config_overrides_context(RBAC_ENABLED=False), - patch.object(module.RagPipelineTransformService, "transform_dataset") as transform_dataset, - pytest.raises(Forbidden), - ): - handler(api, session, DEFAULT_WORKFLOW_TENANT_ID, _account(), UUID(DEFAULT_DATASET_ID)) + @staticmethod + def test_rag_pipeline_transform_rejects_dataset_from_another_tenant_before_service_call( + sqlite_engine: Engine, + ) -> None: + api = module.RagPipelineTransformApi() + handler = unwrap_all(api.post) - transform_dataset.assert_not_called() + with Session(sqlite_engine) as session: + session.add(_dataset(tenant_id="00000000-0000-0000-0000-000000000099")) + with ( + patch.object(module.RagPipelineTransformService, "transform_dataset") as transform_dataset, + pytest.raises(NotFound), + ): + handler( + api, + session, + _RAG_PIPELINE_WORKFLOW_DEFAULT_WORKFLOW_TENANT_ID, + _account(), + UUID(DEFAULT_DATASET_ID), + ) -def test_rag_pipeline_transform_passes_authorized_dataset_and_account_to_service( - sqlite_engine: Engine, -) -> None: - api = module.RagPipelineTransformApi() - handler = unwrap_all(api.post) - account = _account() - expected = {"pipeline_id": "pipeline-1", "dataset_id": DEFAULT_DATASET_ID, "status": "success"} + transform_dataset.assert_not_called() - with Session(sqlite_engine) as session: - dataset = _dataset() - session.add(dataset) + @staticmethod + def test_rag_pipeline_transform_enforces_legacy_dataset_permission_before_service_call( + sqlite_engine: Engine, + ) -> None: + api = module.RagPipelineTransformApi() + handler = unwrap_all(api.post) - with ( - config_overrides_context(RBAC_ENABLED=False), - patch.object(module.RagPipelineTransformService, "transform_dataset", return_value=expected) as transform, - ): - response = handler(api, session, DEFAULT_WORKFLOW_TENANT_ID, account, UUID(DEFAULT_DATASET_ID)) + with Session(sqlite_engine) as session: + session.add(_dataset(maintainer="00000000-0000-0000-0000-000000000099")) - transform.assert_called_once_with(dataset, account.id, session) + with ( + config_overrides_context(RBAC_ENABLED=False), + patch.object(module.RagPipelineTransformService, "transform_dataset") as transform_dataset, + pytest.raises(Forbidden), + ): + handler( + api, + session, + _RAG_PIPELINE_WORKFLOW_DEFAULT_WORKFLOW_TENANT_ID, + _account(), + UUID(DEFAULT_DATASET_ID), + ) - assert response == expected + transform_dataset.assert_not_called() + @staticmethod + def test_rag_pipeline_transform_passes_authorized_dataset_and_account_to_service( + sqlite_engine: Engine, + ) -> None: + api = module.RagPipelineTransformApi() + handler = unwrap_all(api.post) + account = _account() + expected = {"pipeline_id": "pipeline-1", "dataset_id": DEFAULT_DATASET_ID, "status": "success"} -def test_rag_pipeline_transform_maps_missing_pipeline_to_not_found(sqlite_engine: Engine) -> None: - api = module.RagPipelineTransformApi() - handler = unwrap_all(api.post) + with Session(sqlite_engine) as session: + dataset = _dataset() + session.add(dataset) - with Session(sqlite_engine) as session: - session.add(_dataset()) + with ( + config_overrides_context(RBAC_ENABLED=False), + patch.object( + module.RagPipelineTransformService, "transform_dataset", return_value=expected + ) as transform, + ): + response = handler( + api, session, _RAG_PIPELINE_WORKFLOW_DEFAULT_WORKFLOW_TENANT_ID, account, UUID(DEFAULT_DATASET_ID) + ) - with ( - config_overrides_context(RBAC_ENABLED=False), - patch.object( - module.RagPipelineTransformService, - "transform_dataset", - side_effect=RagPipelineResourceNotFoundError("Pipeline not found"), + transform.assert_called_once_with(dataset, account.id, session) + + assert response == expected + + @staticmethod + def test_rag_pipeline_transform_maps_missing_pipeline_to_not_found(sqlite_engine: Engine) -> None: + api = module.RagPipelineTransformApi() + handler = unwrap_all(api.post) + + with Session(sqlite_engine) as session: + session.add(_dataset()) + + with ( + config_overrides_context(RBAC_ENABLED=False), + patch.object( + module.RagPipelineTransformService, + "transform_dataset", + side_effect=RagPipelineResourceNotFoundError("Pipeline not found"), + ), + pytest.raises(NotFound, match="Pipeline not found"), + ): + handler( + api, + session, + _RAG_PIPELINE_WORKFLOW_DEFAULT_WORKFLOW_TENANT_ID, + _account(), + UUID(DEFAULT_DATASET_ID), + ) + + @staticmethod + def test_rag_pipeline_transform_skips_legacy_acl_when_rbac_is_enabled(sqlite_engine: Engine) -> None: + api = module.RagPipelineTransformApi() + handler = unwrap_all(api.post) + account = _account() + account.role = TenantAccountRole.NORMAL + expected = {"pipeline_id": "pipeline-1", "dataset_id": DEFAULT_DATASET_ID, "status": "success"} + + with Session(sqlite_engine) as session: + session.add(_dataset(maintainer="00000000-0000-0000-0000-000000000099")) + + with ( + config_overrides_context(RBAC_ENABLED=True), + patch.object( + module.RagPipelineTransformService, "transform_dataset", return_value=expected + ) as transform, + ): + response = handler( + api, session, _RAG_PIPELINE_WORKFLOW_DEFAULT_WORKFLOW_TENANT_ID, account, UUID(DEFAULT_DATASET_ID) + ) + + assert response == expected + transform.assert_called_once() + + @staticmethod + @pytest.mark.parametrize( + ("api_type", "payload"), + [ + ( + module.DraftRagPipelineRunApi, + {"inputs": {}, "datasource_type": "x", "datasource_info_list": [], "start_node_id": "node-1"}, ), - pytest.raises(NotFound, match="Pipeline not found"), - ): - handler(api, session, DEFAULT_WORKFLOW_TENANT_ID, _account(), UUID(DEFAULT_DATASET_ID)) - - -def test_rag_pipeline_transform_skips_legacy_acl_when_rbac_is_enabled(sqlite_engine: Engine) -> None: - api = module.RagPipelineTransformApi() - handler = unwrap_all(api.post) - account = _account() - account.role = TenantAccountRole.NORMAL - expected = {"pipeline_id": "pipeline-1", "dataset_id": DEFAULT_DATASET_ID, "status": "success"} - - with Session(sqlite_engine) as session: - session.add(_dataset(maintainer="00000000-0000-0000-0000-000000000099")) + ( + module.PublishedRagPipelineRunApi, + { + "inputs": {}, + "datasource_type": "x", + "datasource_info_list": [], + "start_node_id": "node-1", + "response_mode": "blocking", + }, + ), + ], + ) + def test_rag_pipeline_run_uses_sqlite_session( + app: Flask, + sqlite_engine: Engine, + pipeline_application: PipelineGenerator, + api_type: type, + payload: dict[str, object], + ) -> None: + api = api_type() + handler = unwrap_all(api.post) + pipeline = _pipeline() with ( - config_overrides_context(RBAC_ENABLED=True), - patch.object(module.RagPipelineTransformService, "transform_dataset", return_value=expected) as transform, + Session(sqlite_engine) as session, + app.test_request_context("/", json=payload), + patch.object(module, "load_rag_pipeline", return_value=pipeline) as load_pipeline, + patch.object(module.PipelineGenerateService, "generate", return_value={"event": "done"}) as generate, + patch.object(module.helper, "compact_generate_response", return_value={"ok": True}), ): - response = handler(api, session, DEFAULT_WORKFLOW_TENANT_ID, account, UUID(DEFAULT_DATASET_ID)) + req_data = ( + RagPipelineDraftWorkflowRunPayload.model_validate(payload) + if api_type is module.DraftRagPipelineRunApi + else PublishedWorkflowRunPayload.model_validate(payload) + ) + response = handler(api, req_data, session, _account(), pipeline.id) - assert response == expected - transform.assert_called_once() + assert response == {"ok": True} + load_pipeline.assert_called_once_with(session, pipeline.id) + assert generate.call_args.kwargs["session"] is session + assert generate.call_args.kwargs["generator"] is pipeline_application + assert session.get_bind() is sqlite_engine - -@pytest.mark.parametrize( - ("api_type", "payload"), - [ - ( - module.DraftRagPipelineRunApi, - {"inputs": {}, "datasource_type": "x", "datasource_info_list": [], "start_node_id": "node-1"}, - ), - ( - module.PublishedRagPipelineRunApi, - { - "inputs": {}, - "datasource_type": "x", - "datasource_info_list": [], - "start_node_id": "node-1", - "response_mode": "blocking", - }, - ), - ], -) -def test_rag_pipeline_run_uses_sqlite_session( - app: Flask, - sqlite_engine: Engine, - pipeline_application: PipelineGenerator, - api_type: type, - payload: dict[str, object], -) -> None: - api = api_type() - handler = unwrap_all(api.post) - pipeline = _pipeline() - - with ( - Session(sqlite_engine) as session, - app.test_request_context("/", json=payload), - patch.object(module, "load_rag_pipeline", return_value=pipeline) as load_pipeline, - patch.object(module.PipelineGenerateService, "generate", return_value={"event": "done"}) as generate, - patch.object(module.helper, "compact_generate_response", return_value={"ok": True}), - ): + @staticmethod + @pytest.mark.parametrize("api_type", [module.DraftRagPipelineRunApi, module.PublishedRagPipelineRunApi]) + def test_rag_pipeline_run_translates_rate_limit( + app: Flask, + sqlite_engine: Engine, + pipeline_application: PipelineGenerator, + api_type: type, + ) -> None: + payload: dict[str, object] = { + "inputs": {}, + "datasource_type": "x", + "datasource_info_list": [], + "start_node_id": "node-1", + } + api = api_type() + handler = unwrap_all(api.post) + pipeline = _pipeline() req_data = ( RagPipelineDraftWorkflowRunPayload.model_validate(payload) if api_type is module.DraftRagPipelineRunApi else PublishedWorkflowRunPayload.model_validate(payload) ) - response = handler(api, req_data, session, _account(), pipeline.id) - assert response == {"ok": True} - load_pipeline.assert_called_once_with(session, pipeline.id) - assert generate.call_args.kwargs["session"] is session - assert generate.call_args.kwargs["generator"] is pipeline_application - assert session.get_bind() is sqlite_engine + with ( + Session(sqlite_engine) as session, + app.test_request_context("/", json=payload), + patch.object(module, "load_rag_pipeline", return_value=pipeline), + patch.object( + module.PipelineGenerateService, "generate", side_effect=InvokeRateLimitError("limit") + ) as generate, + pytest.raises(InvokeRateLimitHttpError), + ): + handler(api, req_data, session, _account(), pipeline.id) + + assert generate.call_args.kwargs["session"] is session + assert generate.call_args.kwargs["generator"] is pipeline_application -@pytest.mark.parametrize("api_type", [module.DraftRagPipelineRunApi, module.PublishedRagPipelineRunApi]) -def test_rag_pipeline_run_translates_rate_limit( - app: Flask, - sqlite_engine: Engine, - pipeline_application: PipelineGenerator, - api_type: type, -) -> None: - payload: dict[str, object] = { - "inputs": {}, - "datasource_type": "x", - "datasource_info_list": [], - "start_node_id": "node-1", +"""Unit tests for rag_pipeline_workflow controller endpoints.""" + + +_RAG_PIPELINE_WORKFLOW_APIS_DEFAULT_WORKFLOW_TENANT_ID = "00000000-0000-0000-0000-000000000001" +_RAG_PIPELINE_WORKFLOW_APIS_DEFAULT_WORKFLOW_APP_ID = "00000000-0000-0000-0000-000000000002" +_RAG_PIPELINE_WORKFLOW_APIS_DEFAULT_WORKFLOW_CREATED_BY = "00000000-0000-0000-0000-000000000003" +type WorkflowVariablePayload = dict[str, object] + + +@dataclass(frozen=True) +class SQLiteDatabase: + """Expose the concrete SQLite engine and scoped session interface used by controller code.""" + + engine: Engine + session: scoped_session[Session] + + +def empty_mapping() -> dict[str, object]: + return {} + + +def empty_list() -> list[object]: + return [] + + +class WorkflowFactoryPayload(TypedDict): + id: str + tenant_id: str + app_id: str + type: str + version: str + marked_name: str + marked_comment: str + graph: str + features: str + created_by: str + created_at: datetime + updated_by: str | None + updated_at: datetime | None + environment_variables: list[WorkflowVariablePayload] + conversation_variables: list[WorkflowVariablePayload] + rag_pipeline_variables: list[WorkflowVariablePayload] + + +class WorkflowFactoryOverrides(TypedDict, total=False): + id: str + tenant_id: str + app_id: str + type: str + version: str + marked_name: str + marked_comment: str + graph: str + features: str + created_by: str + created_at: datetime + updated_by: str | None + updated_at: datetime | None + environment_variables: list[WorkflowVariablePayload] + conversation_variables: list[WorkflowVariablePayload] + rag_pipeline_variables: list[WorkflowVariablePayload] + + +class NodeExecutionOverrides(TypedDict, total=False): + id: str + tenant_id: str + app_id: str + workflow_id: str + workflow_run_id: str | None + index: int + predecessor_node_id: str | None + node_execution_id: str | None + node_id: str + node_type: str + title: str + inputs: str | None + process_data: str | None + outputs: str | None + status: WorkflowNodeExecutionStatus + error: str | None + elapsed_time: float + execution_metadata: str | None + created_at: datetime + created_by_role: CreatorUserRole + created_by: str + finished_at: datetime | None + + +def make_node_execution(**overrides: Unpack[NodeExecutionOverrides]) -> WorkflowNodeExecutionModel: + payload: NodeExecutionOverrides = { + "id": "node-exec-1", + "tenant_id": _RAG_PIPELINE_WORKFLOW_APIS_DEFAULT_WORKFLOW_TENANT_ID, + "app_id": _RAG_PIPELINE_WORKFLOW_APIS_DEFAULT_WORKFLOW_APP_ID, + "workflow_id": "workflow-1", + "workflow_run_id": None, + "index": 1, + "predecessor_node_id": None, + "node_execution_id": None, + "node_id": "node1", + "node_type": "start", + "title": "Start", + "inputs": json.dumps({"query": "hello"}), + "process_data": json.dumps({}), + "outputs": json.dumps({"answer": "world"}), + "status": WorkflowNodeExecutionStatus.SUCCEEDED, + "error": None, + "elapsed_time": 1.0, + "execution_metadata": json.dumps({}), + "created_at": datetime(2026, 1, 1, 0, 0, 0), + "created_by_role": CreatorUserRole.ACCOUNT, + "created_by": _RAG_PIPELINE_WORKFLOW_APIS_DEFAULT_WORKFLOW_CREATED_BY, + "finished_at": datetime(2026, 1, 1, 0, 0, 1), } - api = api_type() - handler = unwrap_all(api.post) - pipeline = _pipeline() - req_data = ( - RagPipelineDraftWorkflowRunPayload.model_validate(payload) - if api_type is module.DraftRagPipelineRunApi - else PublishedWorkflowRunPayload.model_validate(payload) + payload.update(overrides) + execution = WorkflowNodeExecutionModel( + triggered_from=WorkflowNodeExecutionTriggeredFrom.RAG_PIPELINE_RUN, + **payload, ) + execution.offload_data = [] + return execution - with ( - Session(sqlite_engine) as session, - app.test_request_context("/", json=payload), - patch.object(module, "load_rag_pipeline", return_value=pipeline), - patch.object(module.PipelineGenerateService, "generate", side_effect=InvokeRateLimitError("limit")) as generate, - pytest.raises(InvokeRateLimitHttpError), - ): - handler(api, req_data, session, _account(), pipeline.id) - assert generate.call_args.kwargs["session"] is session - assert generate.call_args.kwargs["generator"] is pipeline_application +def default_workflow_payload() -> WorkflowFactoryPayload: + return { + "id": "workflow-1", + "tenant_id": _RAG_PIPELINE_WORKFLOW_APIS_DEFAULT_WORKFLOW_TENANT_ID, + "app_id": _RAG_PIPELINE_WORKFLOW_APIS_DEFAULT_WORKFLOW_APP_ID, + "type": "workflow", + "version": "1", + "marked_name": "Release 1", + "marked_comment": "Initial release", + "graph": json.dumps({"nodes": [], "edges": []}), + "features": json.dumps({"file_upload": {"enabled": False}}), + "created_by": _RAG_PIPELINE_WORKFLOW_APIS_DEFAULT_WORKFLOW_CREATED_BY, + "created_at": datetime(2024, 1, 1, 12, 0, 0), + "updated_by": None, + "updated_at": datetime(2024, 1, 1, 12, 1, 0), + "environment_variables": [], + "conversation_variables": [], + "rag_pipeline_variables": [], + } + + +def make_workflow(**overrides: Unpack[WorkflowFactoryOverrides]) -> Workflow: + payload = default_workflow_payload() + payload.update(overrides) + return Workflow(**payload) + + +def _rag_pipeline_workflow_apis_make_account( + *, id: str = "account-1", role: TenantAccountRole = TenantAccountRole.EDITOR +) -> Account: + account = Account(name="Alice", email=f"{id}@example.com") + account.id = id + account.role = role + return account + + +def make_pipeline( + *, + id: str = "pipeline-1", + tenant_id: str = "tenant-1", + workflow_id: str | None = None, + is_published: bool = False, +) -> Pipeline: + pipeline = Pipeline(tenant_id=tenant_id, name="test-pipeline", description="test") + pipeline.id = id + pipeline.workflow_id = workflow_id + pipeline.is_published = is_published + return pipeline + + +class TestRagPipelineWorkflowApisCases: + pytestmark = pytest.mark.usefixtures("pipeline_application") + + @staticmethod + @pytest.fixture(autouse=True) + def sqlite_database( + monkeypatch: pytest.MonkeyPatch, + sqlite_engine: Engine, + ) -> Iterator[scoped_session[Session]]: + """Route controller transactions and model author lookups through SQLite.""" + + database_session = scoped_session(sessionmaker(bind=sqlite_engine, expire_on_commit=False)) + database = SQLiteDatabase(engine=sqlite_engine, session=database_session) + monkeypatch.setattr(workflow_controller, "db", database) + monkeypatch.setattr(workflow_models, "db", database) + monkeypatch.setattr(rag_pipeline_service, "db", database) + + with database_session() as session: + default_author = Account(name="Default Author", email="default-author@example.com") + default_author.id = _RAG_PIPELINE_WORKFLOW_APIS_DEFAULT_WORKFLOW_CREATED_BY + session.add(default_author) + session.commit() + + try: + yield database_session + finally: + database_session.remove() + + @staticmethod + @pytest.fixture + def workflow_author(sqlite_database: scoped_session[Session]) -> Account: + account = Account(name="Alice", email=f"alice-{uuid4()}@example.com") + account.id = str(uuid4()) + sqlite_database.add(account) + sqlite_database.commit() + return account + + class TestDraftWorkflowApi: + def test_get_draft_success(self, app: Flask, workflow_author: Account, sqlite_session: Session) -> None: + api = DraftRagPipelineApi() + method = unwrap(api.get) + + pipeline = make_pipeline() + workflow = make_workflow(created_by=workflow_author.id) + + service = MagicMock() + service.get_draft_workflow.return_value = workflow + + with ( + app.test_request_context("/"), + patch( + "controllers.console.datasets.rag_pipeline.rag_pipeline_workflow.RagPipelineService", + return_value=service, + ), + ): + result = method(api, sqlite_session, pipeline) + + assert result["id"] == "workflow-1" + assert result["graph"] == {"nodes": [], "edges": []} + assert result["features"] == {"file_upload": {"enabled": False}} + assert result["hash"] == workflow.unique_hash + assert result["created_by"] == { + "id": workflow_author.id, + "name": workflow_author.name, + "email": workflow_author.email, + } + assert result["updated_by"] is None + + def test_get_draft_not_exist(self, app: Flask, sqlite_session: Session) -> None: + api = DraftRagPipelineApi() + method = unwrap(api.get) + + pipeline = make_pipeline() + service = MagicMock() + service.get_draft_workflow.return_value = None + + with ( + app.test_request_context("/"), + patch( + "controllers.console.datasets.rag_pipeline.rag_pipeline_workflow.RagPipelineService", + return_value=service, + ), + ): + with pytest.raises(DraftWorkflowNotExist): + method(api, sqlite_session, pipeline) + + def test_sync_hash_not_match(self, app: Flask) -> None: + api = DraftRagPipelineApi() + method = unwrap(api.post) + + pipeline = make_pipeline() + user = _rag_pipeline_workflow_apis_make_account() + + service = MagicMock() + service.sync_draft_workflow.side_effect = WorkflowHashNotEqualError() + + with ( + app.test_request_context("/", json={"graph": empty_mapping(), "features": empty_mapping()}), + patch( + "controllers.console.datasets.rag_pipeline.rag_pipeline_workflow.RagPipelineService", + return_value=service, + ), + ): + with pytest.raises(DraftWorkflowNotSync): + method(api, user, pipeline) + + def test_sync_invalid_text_plain(self, app: Flask) -> None: + api = DraftRagPipelineApi() + method = unwrap(api.post) + + pipeline = make_pipeline() + user = _rag_pipeline_workflow_apis_make_account() + + with ( + app.test_request_context("/", data="bad-json", headers={"Content-Type": "text/plain"}), + ): + response, status = raised_error(lambda: method(api, user, pipeline)) + assert status == 400 + + def test_restore_published_workflow_to_draft_success(self, app: Flask) -> None: + api = RagPipelineDraftWorkflowRestoreApi() + method = unwrap(api.post) + + pipeline = make_pipeline() + user = _rag_pipeline_workflow_apis_make_account(id="account-1") + workflow = make_workflow( + graph=json.dumps({"nodes": [{"id": "restored"}], "edges": []}), + created_at=datetime(2024, 1, 1), + ) + + service = MagicMock() + service.restore_published_workflow_to_draft.return_value = workflow + + with ( + app.test_request_context("/", method="POST"), + patch( + "controllers.console.datasets.rag_pipeline.rag_pipeline_workflow.RagPipelineService", + return_value=service, + ), + ): + result = method(api, user, pipeline, "published-workflow") + + assert result["result"] == "success" + assert result["hash"] == workflow.unique_hash + + def test_restore_published_workflow_to_draft_not_found(self, app: Flask) -> None: + api = RagPipelineDraftWorkflowRestoreApi() + method = unwrap(api.post) + + pipeline = make_pipeline() + user = _rag_pipeline_workflow_apis_make_account(id="account-1") + + service = MagicMock() + service.restore_published_workflow_to_draft.side_effect = WorkflowNotFoundError("Workflow not found") + + with ( + app.test_request_context("/", method="POST"), + patch( + "controllers.console.datasets.rag_pipeline.rag_pipeline_workflow.RagPipelineService", + return_value=service, + ), + ): + with pytest.raises(NotFound): + method(api, user, pipeline, "published-workflow") + + def test_restore_published_workflow_to_draft_returns_400_for_draft_source(self, app: Flask) -> None: + api = RagPipelineDraftWorkflowRestoreApi() + method = unwrap(api.post) + + pipeline = make_pipeline() + user = _rag_pipeline_workflow_apis_make_account(id="account-1") + + service = MagicMock() + service.restore_published_workflow_to_draft.side_effect = IsDraftWorkflowError( + "source workflow must be published" + ) + + with ( + app.test_request_context("/", method="POST"), + patch( + "controllers.console.datasets.rag_pipeline.rag_pipeline_workflow.RagPipelineService", + return_value=service, + ), + ): + with pytest.raises(HTTPException) as exc: + method(api, user, pipeline, "draft-workflow") + + assert exc.value.code == 400 + assert exc.value.description == "source workflow must be published" + + class TestDraftRunNodes: + def test_iteration_node_success(self, app: Flask) -> None: + api = RagPipelineDraftRunIterationNodeApi() + method = unwrap(api.post) + + pipeline = make_pipeline() + user = _rag_pipeline_workflow_apis_make_account() + + with ( + app.test_request_context("/", json={"inputs": empty_mapping()}), + patch( + "controllers.console.datasets.rag_pipeline.rag_pipeline_workflow.PipelineGenerateService.generate_single_iteration", + return_value=MagicMock(), + ), + patch( + "controllers.console.datasets.rag_pipeline.rag_pipeline_workflow.helper.compact_generate_response", + return_value={"ok": True}, + ), + ): + result = method(api, NodeRunPayload(), user, pipeline, "node") + assert result == {"ok": True} + + def test_iteration_node_conversation_not_exists(self, app: Flask) -> None: + api = RagPipelineDraftRunIterationNodeApi() + method = unwrap(api.post) + + pipeline = make_pipeline() + user = _rag_pipeline_workflow_apis_make_account() + + with ( + app.test_request_context("/", json={"inputs": empty_mapping()}), + patch( + "controllers.console.datasets.rag_pipeline.rag_pipeline_workflow.PipelineGenerateService.generate_single_iteration", + side_effect=services.errors.conversation.ConversationNotExistsError(), + ), + ): + with pytest.raises(NotFound): + method(api, NodeRunPayload(), user, pipeline, "node") + + def test_loop_node_success(self, app: Flask) -> None: + api = RagPipelineDraftRunLoopNodeApi() + method = unwrap(api.post) + + pipeline = make_pipeline() + user = _rag_pipeline_workflow_apis_make_account() + + with ( + app.test_request_context("/", json={"inputs": empty_mapping()}), + patch( + "controllers.console.datasets.rag_pipeline.rag_pipeline_workflow.PipelineGenerateService.generate_single_loop", + return_value=MagicMock(), + ), + patch( + "controllers.console.datasets.rag_pipeline.rag_pipeline_workflow.helper.compact_generate_response", + return_value={"ok": True}, + ), + ): + assert method(api, NodeRunPayload(), user, pipeline, "node") == {"ok": True} + + class TestDraftNodeRun: + def test_execution_not_found(self, app: Flask) -> None: + api = RagPipelineDraftNodeRunApi() + method = unwrap(api.post) + + pipeline = make_pipeline() + user = _rag_pipeline_workflow_apis_make_account() + + service = MagicMock() + service.run_draft_workflow_node.return_value = None + + with ( + app.test_request_context("/", json={"inputs": empty_mapping()}), + patch( + "controllers.console.datasets.rag_pipeline.rag_pipeline_workflow.RagPipelineService", + return_value=service, + ), + ): + with pytest.raises(ValueError, match="Workflow node execution not found"): + method(api, NodeRunRequiredPayload(inputs={}), user, pipeline, "node") + + class TestPublishedPipelineApis: + def test_get_published_workflow_preserves_response_fields( + self, app: Flask, workflow_author: Account, sqlite_database: scoped_session[Session] + ) -> None: + workflow = make_workflow(id=str(uuid4()), created_by=workflow_author.id) + sqlite_database.add(workflow) + sqlite_database.commit() + pipeline = make_pipeline( + id=workflow.app_id, + tenant_id=workflow.tenant_id, + workflow_id=workflow.id, + is_published=True, + ) + + with app.test_request_context("/"): + result = unwrap(PublishedRagPipelineApi.get)(PublishedRagPipelineApi(), pipeline) + + assert result["id"] == workflow.id + assert result["graph"] == {"nodes": [], "edges": []} + assert result["hash"] == workflow.unique_hash + assert result["created_by"]["id"] == workflow_author.id + + @pytest.mark.parametrize("is_published", [False, True]) + def test_get_without_published_workflow_returns_null(self, app: Flask, is_published: bool) -> None: + pipeline = make_pipeline( + id=_RAG_PIPELINE_WORKFLOW_APIS_DEFAULT_WORKFLOW_APP_ID, + tenant_id=_RAG_PIPELINE_WORKFLOW_APIS_DEFAULT_WORKFLOW_TENANT_ID, + workflow_id=str(uuid4()), + is_published=is_published, + ) + + with app.test_request_context("/"): + result = unwrap(PublishedRagPipelineApi.get)(PublishedRagPipelineApi(), pipeline) + + assert result is None + + def test_publish_success(self, app: Flask) -> None: + api = PublishedRagPipelineApi() + method = unwrap(api.post) + + tenant_id = str(uuid4()) + pipeline = Pipeline( + tenant_id=tenant_id, + name="test-pipeline", + description="test", + created_by=str(uuid4()), + ) + user = _rag_pipeline_workflow_apis_make_account(id="u1") + + workflow = make_workflow(id=str(uuid4()), created_at=naive_utc_now()) + + service = MagicMock() + service.publish_workflow.return_value = workflow + + with ( + app.test_request_context("/"), + patch( + "controllers.console.datasets.rag_pipeline.rag_pipeline_workflow.RagPipelineService", + return_value=service, + ), + ): + result = method(api, user, pipeline) + + assert result["result"] == "success" + assert "created_at" in result + + class TestMiscApis: + def test_task_stop(self, app: Flask) -> None: + api = RagPipelineTaskStopApi() + method = unwrap(api.post) + + pipeline = make_pipeline() + user = _rag_pipeline_workflow_apis_make_account(id="u1") + + with ( + app.test_request_context("/"), + patch( + "controllers.console.datasets.rag_pipeline.rag_pipeline_workflow.AppQueueManager.set_stop_flag" + ) as stop_mock, + ): + result = method(api, user, pipeline, "task-1") + stop_mock.assert_called_once() + assert result["result"] == "success" + + def test_recommended_plugins(self, app: Flask) -> None: + api = RagPipelineRecommendedPluginApi() + method = unwrap(api.get) + + service = MagicMock() + recommended_plugins = { + "installed_recommended_plugins": [ + { + "id": "p1", + "author": "langgenius", + "name": "p1", + "description": { + "en_US": "Plugin one", + "zh_Hans": "Plugin one", + "pt_BR": "Plugin one", + "ja_JP": "Plugin one", + }, + "icon": "/icon.svg", + "label": { + "en_US": "Plugin one", + "zh_Hans": "Plugin one", + "pt_BR": "Plugin one", + "ja_JP": "Plugin one", + }, + "type": "builtin", + } + ], + "uninstalled_recommended_plugins": [{"id": "p2"}], + } + service.get_recommended_plugins.return_value = recommended_plugins + user = _rag_pipeline_workflow_apis_make_account() + tenant_id = "tenant-1" + + with ( + app.test_request_context("/?type=all"), + patch( + "controllers.console.datasets.rag_pipeline.rag_pipeline_workflow.RagPipelineService", + return_value=service, + ), + ): + result = method(api, RagPipelineRecommendedPluginQuery(type="all"), tenant_id, user) + assert result == recommended_plugins + service.get_recommended_plugins.assert_called_once_with("all", user, tenant_id) + + class TestDefaultBlockConfigApi: + def test_get_block_config_success(self, app: Flask) -> None: + api = DefaultRagPipelineBlockConfigApi() + method = unwrap(api.get) + + pipeline = make_pipeline() + + service = MagicMock() + service.get_default_block_config.return_value = {"k": "v"} + + with ( + app.test_request_context("/?q={}"), + patch( + "controllers.console.datasets.rag_pipeline.rag_pipeline_workflow.RagPipelineService", + return_value=service, + ), + ): + result = method(api, DefaultBlockConfigQuery(q="{}"), pipeline, "llm") + assert result == {"k": "v"} + + def test_get_block_config_invalid_json(self, app: Flask) -> None: + api = DefaultRagPipelineBlockConfigApi() + method = unwrap(api.get) + + pipeline = make_pipeline() + + with app.test_request_context("/?q=bad-json"): + with pytest.raises(ValueError, match="Invalid filters"): + method(api, DefaultBlockConfigQuery(q="bad-json"), pipeline, "llm") + + class TestPublishedAllRagPipelineApi: + def test_get_published_workflows_success(self, app: Flask) -> None: + api = PublishedAllRagPipelineApi() + method = unwrap(api.get) + + pipeline = make_pipeline() + user = _rag_pipeline_workflow_apis_make_account(id="u1") + + service = MagicMock() + service.get_all_published_workflow.return_value = ([make_workflow(id="w1")], False) + + with ( + app.test_request_context("/"), + patch( + "controllers.console.datasets.rag_pipeline.rag_pipeline_workflow.RagPipelineService", + return_value=service, + ), + ): + result = method(api, WorkflowListQuery(), user, pipeline) + + assert result["items"][0]["id"] == "w1" + assert result["items"][0]["graph"] == {"nodes": [], "edges": []} + assert result["has_more"] is False + + def test_get_published_workflows_forbidden(self, app: Flask) -> None: + api = PublishedAllRagPipelineApi() + method = unwrap(api.get) + + pipeline = make_pipeline() + user = _rag_pipeline_workflow_apis_make_account(id="u1") + + with ( + app.test_request_context("/?user_id=u2"), + ): + with pytest.raises(Forbidden): + method(api, WorkflowListQuery(user_id="u2"), user, pipeline) + + class TestRagPipelineByIdApi: + def test_patch_success(self, app: Flask) -> None: + api = RagPipelineByIdApi() + method = unwrap(api.patch) + + pipeline = make_pipeline(tenant_id="t1") + user = _rag_pipeline_workflow_apis_make_account(id="u1") + + workflow = make_workflow(id="w1", marked_name="test") + + service = MagicMock() + service.update_workflow.return_value = workflow + + payload = {"marked_name": "test"} + + with ( + app.test_request_context("/", json=payload), + patch( + "controllers.console.datasets.rag_pipeline.rag_pipeline_workflow.RagPipelineService", + return_value=service, + ), + ): + result = method(api, WorkflowUpdatePayload.model_validate(payload), user, pipeline, "w1") + + assert result["id"] == "w1" + assert result["marked_name"] == "test" + assert result["hash"] == workflow.unique_hash + + def test_patch_no_fields(self, app: Flask) -> None: + api = RagPipelineByIdApi() + method = unwrap(api.patch) + + pipeline = make_pipeline() + user = _rag_pipeline_workflow_apis_make_account() + + with app.test_request_context("/", json={}): + result, status = raised_error(lambda: method(api, WorkflowUpdatePayload(), user, pipeline, "w1")) + assert status == 400 + assert result == {"code": "invalid_param", "message": "No valid fields to update"} + + @pytest.mark.parametrize("transaction_fails", [False, True], ids=["commit-succeeds", "commit-fails"]) + def test_delete_retires_candidates_only_after_transaction_exit( + self, + app: Flask, + transaction_fails: bool, + ) -> None: + api = RagPipelineByIdApi() + method = unwrap(api.delete) + + pipeline = make_pipeline(tenant_id="t1", workflow_id="active-workflow") + user = _rag_pipeline_workflow_apis_make_account() + + events: list[str] = [] + error = RuntimeError("commit failed") + workflow_service = MagicMock() + workflow_service.delete_workflow.side_effect = lambda **_kwargs: events.append("delete") or ["inline-agent"] + transaction_factory = MagicMock() + + @contextmanager + def transaction() -> Generator[Session]: + events.append("transaction-enter") + with Session() as session: + yield session + events.append("transaction-exit") + if transaction_fails: + raise error + + transaction_factory.begin.side_effect = transaction + + with ( + app.test_request_context("/", method="DELETE"), + patch( + "controllers.console.datasets.rag_pipeline.rag_pipeline_workflow.WorkflowService", + return_value=workflow_service, + ), + patch( + "controllers.console.datasets.rag_pipeline.rag_pipeline_workflow.sessionmaker", + return_value=transaction_factory, + ), + patch( + "controllers.console.datasets.rag_pipeline.rag_pipeline_workflow." + "WorkflowAgentRetirementService.retire_unowned" + ) as retire_unowned, + ): + retire_unowned.side_effect = lambda **_kwargs: events.append("retire") + if transaction_fails: + with pytest.raises(RuntimeError) as exc_info: + method(api, user, pipeline, "old-workflow") + assert exc_info.value is error + else: + result = method(api, user, pipeline, "old-workflow") + assert result == (None, 204) + + workflow_service.delete_workflow.assert_called_once() + assert events == ["transaction-enter", "delete", "transaction-exit"] + ( + [] if transaction_fails else ["retire"] + ) + if transaction_fails: + retire_unowned.assert_not_called() + else: + retire_unowned.assert_called_once_with( + tenant_id=pipeline.tenant_id, + agent_ids=["inline-agent"], + account_id=user.id, + ) + + def test_delete_active_workflow_rejected(self, app: Flask) -> None: + api = RagPipelineByIdApi() + method = unwrap(api.delete) + + pipeline = make_pipeline(tenant_id="t1", workflow_id="active-workflow") + user = _rag_pipeline_workflow_apis_make_account() + + with app.test_request_context("/", method="DELETE"): + with pytest.raises(BadRequest, match="currently in use by pipeline"): + method(api, user, pipeline, "active-workflow") + + class TestRagPipelineWorkflowLastRunApi: + def test_last_run_success(self, app: Flask) -> None: + api = RagPipelineWorkflowLastRunApi() + method = unwrap(api.get) + + pipeline = make_pipeline() + workflow = make_workflow() + node_exec = make_node_execution() + + service = MagicMock() + service.get_draft_workflow.return_value = workflow + service.get_node_last_run.return_value = node_exec + + with ( + app.test_request_context("/"), + patch( + "controllers.console.datasets.rag_pipeline.rag_pipeline_workflow.RagPipelineService", + return_value=service, + ), + ): + result = method(api, pipeline, "node1") + assert result["id"] == "node-exec-1" + assert result["inputs"] == {"query": "hello"} + assert result["outputs"] == {"answer": "world"} + + def test_last_run_not_found(self, app: Flask) -> None: + api = RagPipelineWorkflowLastRunApi() + method = unwrap(api.get) + + pipeline = make_pipeline() + + service = MagicMock() + service.get_draft_workflow.return_value = None + + with ( + app.test_request_context("/"), + patch( + "controllers.console.datasets.rag_pipeline.rag_pipeline_workflow.RagPipelineService", + return_value=service, + ), + ): + with pytest.raises(NotFound): + method(api, pipeline, "node1") + + class TestRagPipelineWorkflowRunNodeExecutionListApi: + def test_get_node_executions_passes_current_user(self, app: Flask) -> None: + api = RagPipelineWorkflowRunNodeExecutionListApi() + method = unwrap(api.get) + + user = _rag_pipeline_workflow_apis_make_account() + pipeline = make_pipeline() + run_id = uuid4() + node_exec = make_node_execution(workflow_run_id=str(run_id)) + + session_stub = MagicMock() + session_stub.scalar.return_value = None + service = MagicMock() + service.get_rag_pipeline_workflow_run_node_executions.return_value = [ + node_execution_response_source(node_exec, session=session_stub) + ] + + with ( + app.test_request_context("/"), + patch( + "controllers.console.datasets.rag_pipeline.rag_pipeline_workflow.RagPipelineService", + return_value=service, + ), + ): + result = method(api, user, pipeline, run_id) + + service.get_rag_pipeline_workflow_run_node_executions.assert_called_once_with( + pipeline=pipeline, + run_id=str(run_id), + user=user, + ) + assert result["data"][0]["id"] == "node-exec-1" + assert result["data"][0]["inputs"] == {"query": "hello"} + assert result["data"][0]["outputs"] == {"answer": "world"} + + class TestRagPipelineDatasourceVariableApi: + def test_set_datasource_variables_success(self, app: Flask) -> None: + api = RagPipelineDatasourceVariableApi() + method = unwrap(api.post) + + pipeline = make_pipeline() + user = _rag_pipeline_workflow_apis_make_account() + + payload = { + "datasource_type": "db", + "datasource_info": empty_mapping(), + "start_node_id": "n1", + "start_node_title": "Node", + } + + service = MagicMock() + service.set_datasource_variables.return_value = make_node_execution(node_id="n1") + + with ( + app.test_request_context("/", json=payload), + patch( + "controllers.console.datasets.rag_pipeline.rag_pipeline_workflow.RagPipelineService", + return_value=service, + ), + ): + result = method(api, DatasourceVariablesPayload.model_validate(payload), user, pipeline) + assert result["node_id"] == "n1" + assert result["process_data"] == {} diff --git a/api/tests/unit_tests/controllers/console/datasets/rag_pipeline/test_rag_pipeline_workflow_apis.py b/api/tests/unit_tests/controllers/console/datasets/rag_pipeline/test_rag_pipeline_workflow_apis.py deleted file mode 100644 index 8b8c9302ef0..00000000000 --- a/api/tests/unit_tests/controllers/console/datasets/rag_pipeline/test_rag_pipeline_workflow_apis.py +++ /dev/null @@ -1,922 +0,0 @@ -"""Unit tests for rag_pipeline_workflow controller endpoints.""" - -from __future__ import annotations - -import json -from collections.abc import Generator, Iterator -from contextlib import contextmanager -from dataclasses import dataclass -from datetime import datetime -from inspect import unwrap -from typing import TypedDict, Unpack -from unittest.mock import MagicMock, patch -from uuid import uuid4 - -import pytest -from flask import Flask -from sqlalchemy import Engine -from sqlalchemy.orm import Session, scoped_session, sessionmaker -from werkzeug.exceptions import BadRequest, Forbidden, HTTPException, NotFound - -import models.workflow as workflow_models -import services -import services.rag_pipeline.rag_pipeline as rag_pipeline_service -from controllers.console.app.error import DraftWorkflowNotExist, DraftWorkflowNotSync -from controllers.console.datasets.rag_pipeline import rag_pipeline_workflow as workflow_controller -from controllers.console.datasets.rag_pipeline.rag_pipeline_workflow import ( - DatasourceVariablesPayload, - DefaultBlockConfigQuery, - DefaultRagPipelineBlockConfigApi, - DraftRagPipelineApi, - NodeRunPayload, - NodeRunRequiredPayload, - PublishedAllRagPipelineApi, - PublishedRagPipelineApi, - RagPipelineByIdApi, - RagPipelineDatasourceVariableApi, - RagPipelineDraftNodeRunApi, - RagPipelineDraftRunIterationNodeApi, - RagPipelineDraftRunLoopNodeApi, - RagPipelineDraftWorkflowRestoreApi, - RagPipelineRecommendedPluginApi, - RagPipelineRecommendedPluginQuery, - RagPipelineTaskStopApi, - RagPipelineWorkflowLastRunApi, - RagPipelineWorkflowRunNodeExecutionListApi, - WorkflowListQuery, - WorkflowUpdatePayload, -) -from fields.workflow_run_fields import node_execution_response_source -from graphon.enums import WorkflowNodeExecutionStatus -from libs.datetime_utils import naive_utc_now -from models.account import Account, TenantAccountRole -from models.dataset import Pipeline -from models.enums import CreatorUserRole -from models.workflow import Workflow, WorkflowNodeExecutionModel, WorkflowNodeExecutionTriggeredFrom -from services.errors.app import IsDraftWorkflowError, WorkflowHashNotEqualError, WorkflowNotFoundError -from tests.unit_tests.controllers.error_body import raised_error - -pytestmark = pytest.mark.usefixtures("pipeline_application") - - -DEFAULT_WORKFLOW_TENANT_ID = "00000000-0000-0000-0000-000000000001" -DEFAULT_WORKFLOW_APP_ID = "00000000-0000-0000-0000-000000000002" -DEFAULT_WORKFLOW_CREATED_BY = "00000000-0000-0000-0000-000000000003" -type WorkflowVariablePayload = dict[str, object] - - -@dataclass(frozen=True) -class SQLiteDatabase: - """Expose the concrete SQLite engine and scoped session interface used by controller code.""" - - engine: Engine - session: scoped_session[Session] - - -@pytest.fixture(autouse=True) -def sqlite_database( - monkeypatch: pytest.MonkeyPatch, - sqlite_engine: Engine, -) -> Iterator[scoped_session[Session]]: - """Route controller transactions and model author lookups through SQLite.""" - - database_session = scoped_session(sessionmaker(bind=sqlite_engine, expire_on_commit=False)) - database = SQLiteDatabase(engine=sqlite_engine, session=database_session) - monkeypatch.setattr(workflow_controller, "db", database) - monkeypatch.setattr(workflow_models, "db", database) - monkeypatch.setattr(rag_pipeline_service, "db", database) - - with database_session() as session: - default_author = Account(name="Default Author", email="default-author@example.com") - default_author.id = DEFAULT_WORKFLOW_CREATED_BY - session.add(default_author) - session.commit() - - try: - yield database_session - finally: - database_session.remove() - - -def empty_mapping() -> dict[str, object]: - return {} - - -def empty_list() -> list[object]: - return [] - - -class WorkflowFactoryPayload(TypedDict): - id: str - tenant_id: str - app_id: str - type: str - version: str - marked_name: str - marked_comment: str - graph: str - features: str - created_by: str - created_at: datetime - updated_by: str | None - updated_at: datetime | None - environment_variables: list[WorkflowVariablePayload] - conversation_variables: list[WorkflowVariablePayload] - rag_pipeline_variables: list[WorkflowVariablePayload] - - -class WorkflowFactoryOverrides(TypedDict, total=False): - id: str - tenant_id: str - app_id: str - type: str - version: str - marked_name: str - marked_comment: str - graph: str - features: str - created_by: str - created_at: datetime - updated_by: str | None - updated_at: datetime | None - environment_variables: list[WorkflowVariablePayload] - conversation_variables: list[WorkflowVariablePayload] - rag_pipeline_variables: list[WorkflowVariablePayload] - - -class NodeExecutionOverrides(TypedDict, total=False): - id: str - tenant_id: str - app_id: str - workflow_id: str - workflow_run_id: str | None - index: int - predecessor_node_id: str | None - node_execution_id: str | None - node_id: str - node_type: str - title: str - inputs: str | None - process_data: str | None - outputs: str | None - status: WorkflowNodeExecutionStatus - error: str | None - elapsed_time: float - execution_metadata: str | None - created_at: datetime - created_by_role: CreatorUserRole - created_by: str - finished_at: datetime | None - - -def make_node_execution(**overrides: Unpack[NodeExecutionOverrides]) -> WorkflowNodeExecutionModel: - payload: NodeExecutionOverrides = { - "id": "node-exec-1", - "tenant_id": DEFAULT_WORKFLOW_TENANT_ID, - "app_id": DEFAULT_WORKFLOW_APP_ID, - "workflow_id": "workflow-1", - "workflow_run_id": None, - "index": 1, - "predecessor_node_id": None, - "node_execution_id": None, - "node_id": "node1", - "node_type": "start", - "title": "Start", - "inputs": json.dumps({"query": "hello"}), - "process_data": json.dumps({}), - "outputs": json.dumps({"answer": "world"}), - "status": WorkflowNodeExecutionStatus.SUCCEEDED, - "error": None, - "elapsed_time": 1.0, - "execution_metadata": json.dumps({}), - "created_at": datetime(2026, 1, 1, 0, 0, 0), - "created_by_role": CreatorUserRole.ACCOUNT, - "created_by": DEFAULT_WORKFLOW_CREATED_BY, - "finished_at": datetime(2026, 1, 1, 0, 0, 1), - } - payload.update(overrides) - execution = WorkflowNodeExecutionModel( - triggered_from=WorkflowNodeExecutionTriggeredFrom.RAG_PIPELINE_RUN, - **payload, - ) - execution.offload_data = [] - return execution - - -def default_workflow_payload() -> WorkflowFactoryPayload: - return { - "id": "workflow-1", - "tenant_id": DEFAULT_WORKFLOW_TENANT_ID, - "app_id": DEFAULT_WORKFLOW_APP_ID, - "type": "workflow", - "version": "1", - "marked_name": "Release 1", - "marked_comment": "Initial release", - "graph": json.dumps({"nodes": [], "edges": []}), - "features": json.dumps({"file_upload": {"enabled": False}}), - "created_by": DEFAULT_WORKFLOW_CREATED_BY, - "created_at": datetime(2024, 1, 1, 12, 0, 0), - "updated_by": None, - "updated_at": datetime(2024, 1, 1, 12, 1, 0), - "environment_variables": [], - "conversation_variables": [], - "rag_pipeline_variables": [], - } - - -def make_workflow(**overrides: Unpack[WorkflowFactoryOverrides]) -> Workflow: - payload = default_workflow_payload() - payload.update(overrides) - return Workflow(**payload) - - -def make_account(*, id: str = "account-1", role: TenantAccountRole = TenantAccountRole.EDITOR) -> Account: - account = Account(name="Alice", email=f"{id}@example.com") - account.id = id - account.role = role - return account - - -def make_pipeline( - *, - id: str = "pipeline-1", - tenant_id: str = "tenant-1", - workflow_id: str | None = None, - is_published: bool = False, -) -> Pipeline: - pipeline = Pipeline(tenant_id=tenant_id, name="test-pipeline", description="test") - pipeline.id = id - pipeline.workflow_id = workflow_id - pipeline.is_published = is_published - return pipeline - - -@pytest.fixture -def workflow_author(sqlite_database: scoped_session[Session]) -> Account: - account = Account(name="Alice", email=f"alice-{uuid4()}@example.com") - account.id = str(uuid4()) - sqlite_database.add(account) - sqlite_database.commit() - return account - - -class TestDraftWorkflowApi: - def test_get_draft_success(self, app: Flask, workflow_author: Account, sqlite_session: Session) -> None: - api = DraftRagPipelineApi() - method = unwrap(api.get) - - pipeline = make_pipeline() - workflow = make_workflow(created_by=workflow_author.id) - - service = MagicMock() - service.get_draft_workflow.return_value = workflow - - with ( - app.test_request_context("/"), - patch( - "controllers.console.datasets.rag_pipeline.rag_pipeline_workflow.RagPipelineService", - return_value=service, - ), - ): - result = method(api, sqlite_session, pipeline) - - assert result["id"] == "workflow-1" - assert result["graph"] == {"nodes": [], "edges": []} - assert result["features"] == {"file_upload": {"enabled": False}} - assert result["hash"] == workflow.unique_hash - assert result["created_by"] == { - "id": workflow_author.id, - "name": workflow_author.name, - "email": workflow_author.email, - } - assert result["updated_by"] is None - - def test_get_draft_not_exist(self, app: Flask, sqlite_session: Session) -> None: - api = DraftRagPipelineApi() - method = unwrap(api.get) - - pipeline = make_pipeline() - service = MagicMock() - service.get_draft_workflow.return_value = None - - with ( - app.test_request_context("/"), - patch( - "controllers.console.datasets.rag_pipeline.rag_pipeline_workflow.RagPipelineService", - return_value=service, - ), - ): - with pytest.raises(DraftWorkflowNotExist): - method(api, sqlite_session, pipeline) - - def test_sync_hash_not_match(self, app: Flask) -> None: - api = DraftRagPipelineApi() - method = unwrap(api.post) - - pipeline = make_pipeline() - user = make_account() - - service = MagicMock() - service.sync_draft_workflow.side_effect = WorkflowHashNotEqualError() - - with ( - app.test_request_context("/", json={"graph": empty_mapping(), "features": empty_mapping()}), - patch( - "controllers.console.datasets.rag_pipeline.rag_pipeline_workflow.RagPipelineService", - return_value=service, - ), - ): - with pytest.raises(DraftWorkflowNotSync): - method(api, user, pipeline) - - def test_sync_invalid_text_plain(self, app: Flask) -> None: - api = DraftRagPipelineApi() - method = unwrap(api.post) - - pipeline = make_pipeline() - user = make_account() - - with ( - app.test_request_context("/", data="bad-json", headers={"Content-Type": "text/plain"}), - ): - response, status = raised_error(lambda: method(api, user, pipeline)) - assert status == 400 - - def test_restore_published_workflow_to_draft_success(self, app: Flask) -> None: - api = RagPipelineDraftWorkflowRestoreApi() - method = unwrap(api.post) - - pipeline = make_pipeline() - user = make_account(id="account-1") - workflow = make_workflow( - graph=json.dumps({"nodes": [{"id": "restored"}], "edges": []}), - created_at=datetime(2024, 1, 1), - ) - - service = MagicMock() - service.restore_published_workflow_to_draft.return_value = workflow - - with ( - app.test_request_context("/", method="POST"), - patch( - "controllers.console.datasets.rag_pipeline.rag_pipeline_workflow.RagPipelineService", - return_value=service, - ), - ): - result = method(api, user, pipeline, "published-workflow") - - assert result["result"] == "success" - assert result["hash"] == workflow.unique_hash - - def test_restore_published_workflow_to_draft_not_found(self, app: Flask) -> None: - api = RagPipelineDraftWorkflowRestoreApi() - method = unwrap(api.post) - - pipeline = make_pipeline() - user = make_account(id="account-1") - - service = MagicMock() - service.restore_published_workflow_to_draft.side_effect = WorkflowNotFoundError("Workflow not found") - - with ( - app.test_request_context("/", method="POST"), - patch( - "controllers.console.datasets.rag_pipeline.rag_pipeline_workflow.RagPipelineService", - return_value=service, - ), - ): - with pytest.raises(NotFound): - method(api, user, pipeline, "published-workflow") - - def test_restore_published_workflow_to_draft_returns_400_for_draft_source(self, app: Flask) -> None: - api = RagPipelineDraftWorkflowRestoreApi() - method = unwrap(api.post) - - pipeline = make_pipeline() - user = make_account(id="account-1") - - service = MagicMock() - service.restore_published_workflow_to_draft.side_effect = IsDraftWorkflowError( - "source workflow must be published" - ) - - with ( - app.test_request_context("/", method="POST"), - patch( - "controllers.console.datasets.rag_pipeline.rag_pipeline_workflow.RagPipelineService", - return_value=service, - ), - ): - with pytest.raises(HTTPException) as exc: - method(api, user, pipeline, "draft-workflow") - - assert exc.value.code == 400 - assert exc.value.description == "source workflow must be published" - - -class TestDraftRunNodes: - def test_iteration_node_success(self, app: Flask) -> None: - api = RagPipelineDraftRunIterationNodeApi() - method = unwrap(api.post) - - pipeline = make_pipeline() - user = make_account() - - with ( - app.test_request_context("/", json={"inputs": empty_mapping()}), - patch( - "controllers.console.datasets.rag_pipeline.rag_pipeline_workflow.PipelineGenerateService.generate_single_iteration", - return_value=MagicMock(), - ), - patch( - "controllers.console.datasets.rag_pipeline.rag_pipeline_workflow.helper.compact_generate_response", - return_value={"ok": True}, - ), - ): - result = method(api, NodeRunPayload(), user, pipeline, "node") - assert result == {"ok": True} - - def test_iteration_node_conversation_not_exists(self, app: Flask) -> None: - api = RagPipelineDraftRunIterationNodeApi() - method = unwrap(api.post) - - pipeline = make_pipeline() - user = make_account() - - with ( - app.test_request_context("/", json={"inputs": empty_mapping()}), - patch( - "controllers.console.datasets.rag_pipeline.rag_pipeline_workflow.PipelineGenerateService.generate_single_iteration", - side_effect=services.errors.conversation.ConversationNotExistsError(), - ), - ): - with pytest.raises(NotFound): - method(api, NodeRunPayload(), user, pipeline, "node") - - def test_loop_node_success(self, app: Flask) -> None: - api = RagPipelineDraftRunLoopNodeApi() - method = unwrap(api.post) - - pipeline = make_pipeline() - user = make_account() - - with ( - app.test_request_context("/", json={"inputs": empty_mapping()}), - patch( - "controllers.console.datasets.rag_pipeline.rag_pipeline_workflow.PipelineGenerateService.generate_single_loop", - return_value=MagicMock(), - ), - patch( - "controllers.console.datasets.rag_pipeline.rag_pipeline_workflow.helper.compact_generate_response", - return_value={"ok": True}, - ), - ): - assert method(api, NodeRunPayload(), user, pipeline, "node") == {"ok": True} - - -class TestDraftNodeRun: - def test_execution_not_found(self, app: Flask) -> None: - api = RagPipelineDraftNodeRunApi() - method = unwrap(api.post) - - pipeline = make_pipeline() - user = make_account() - - service = MagicMock() - service.run_draft_workflow_node.return_value = None - - with ( - app.test_request_context("/", json={"inputs": empty_mapping()}), - patch( - "controllers.console.datasets.rag_pipeline.rag_pipeline_workflow.RagPipelineService", - return_value=service, - ), - ): - with pytest.raises(ValueError, match="Workflow node execution not found"): - method(api, NodeRunRequiredPayload(inputs={}), user, pipeline, "node") - - -class TestPublishedPipelineApis: - def test_get_published_workflow_preserves_response_fields( - self, app: Flask, workflow_author: Account, sqlite_database: scoped_session[Session] - ) -> None: - workflow = make_workflow(id=str(uuid4()), created_by=workflow_author.id) - sqlite_database.add(workflow) - sqlite_database.commit() - pipeline = make_pipeline( - id=workflow.app_id, - tenant_id=workflow.tenant_id, - workflow_id=workflow.id, - is_published=True, - ) - - with app.test_request_context("/"): - result = unwrap(PublishedRagPipelineApi.get)(PublishedRagPipelineApi(), pipeline) - - assert result["id"] == workflow.id - assert result["graph"] == {"nodes": [], "edges": []} - assert result["hash"] == workflow.unique_hash - assert result["created_by"]["id"] == workflow_author.id - - @pytest.mark.parametrize("is_published", [False, True]) - def test_get_without_published_workflow_returns_null(self, app: Flask, is_published: bool) -> None: - pipeline = make_pipeline( - id=DEFAULT_WORKFLOW_APP_ID, - tenant_id=DEFAULT_WORKFLOW_TENANT_ID, - workflow_id=str(uuid4()), - is_published=is_published, - ) - - with app.test_request_context("/"): - result = unwrap(PublishedRagPipelineApi.get)(PublishedRagPipelineApi(), pipeline) - - assert result is None - - def test_publish_success(self, app: Flask) -> None: - api = PublishedRagPipelineApi() - method = unwrap(api.post) - - tenant_id = str(uuid4()) - pipeline = Pipeline( - tenant_id=tenant_id, - name="test-pipeline", - description="test", - created_by=str(uuid4()), - ) - user = make_account(id="u1") - - workflow = make_workflow(id=str(uuid4()), created_at=naive_utc_now()) - - service = MagicMock() - service.publish_workflow.return_value = workflow - - with ( - app.test_request_context("/"), - patch( - "controllers.console.datasets.rag_pipeline.rag_pipeline_workflow.RagPipelineService", - return_value=service, - ), - ): - result = method(api, user, pipeline) - - assert result["result"] == "success" - assert "created_at" in result - - -class TestMiscApis: - def test_task_stop(self, app: Flask) -> None: - api = RagPipelineTaskStopApi() - method = unwrap(api.post) - - pipeline = make_pipeline() - user = make_account(id="u1") - - with ( - app.test_request_context("/"), - patch( - "controllers.console.datasets.rag_pipeline.rag_pipeline_workflow.AppQueueManager.set_stop_flag" - ) as stop_mock, - ): - result = method(api, user, pipeline, "task-1") - stop_mock.assert_called_once() - assert result["result"] == "success" - - def test_recommended_plugins(self, app: Flask) -> None: - api = RagPipelineRecommendedPluginApi() - method = unwrap(api.get) - - service = MagicMock() - recommended_plugins = { - "installed_recommended_plugins": [ - { - "id": "p1", - "author": "langgenius", - "name": "p1", - "description": { - "en_US": "Plugin one", - "zh_Hans": "Plugin one", - "pt_BR": "Plugin one", - "ja_JP": "Plugin one", - }, - "icon": "/icon.svg", - "label": { - "en_US": "Plugin one", - "zh_Hans": "Plugin one", - "pt_BR": "Plugin one", - "ja_JP": "Plugin one", - }, - "type": "builtin", - } - ], - "uninstalled_recommended_plugins": [{"id": "p2"}], - } - service.get_recommended_plugins.return_value = recommended_plugins - user = make_account() - tenant_id = "tenant-1" - - with ( - app.test_request_context("/?type=all"), - patch( - "controllers.console.datasets.rag_pipeline.rag_pipeline_workflow.RagPipelineService", - return_value=service, - ), - ): - result = method(api, RagPipelineRecommendedPluginQuery(type="all"), tenant_id, user) - assert result == recommended_plugins - service.get_recommended_plugins.assert_called_once_with("all", user, tenant_id) - - -class TestDefaultBlockConfigApi: - def test_get_block_config_success(self, app: Flask) -> None: - api = DefaultRagPipelineBlockConfigApi() - method = unwrap(api.get) - - pipeline = make_pipeline() - - service = MagicMock() - service.get_default_block_config.return_value = {"k": "v"} - - with ( - app.test_request_context("/?q={}"), - patch( - "controllers.console.datasets.rag_pipeline.rag_pipeline_workflow.RagPipelineService", - return_value=service, - ), - ): - result = method(api, DefaultBlockConfigQuery(q="{}"), pipeline, "llm") - assert result == {"k": "v"} - - def test_get_block_config_invalid_json(self, app: Flask) -> None: - api = DefaultRagPipelineBlockConfigApi() - method = unwrap(api.get) - - pipeline = make_pipeline() - - with app.test_request_context("/?q=bad-json"): - with pytest.raises(ValueError, match="Invalid filters"): - method(api, DefaultBlockConfigQuery(q="bad-json"), pipeline, "llm") - - -class TestPublishedAllRagPipelineApi: - def test_get_published_workflows_success(self, app: Flask) -> None: - api = PublishedAllRagPipelineApi() - method = unwrap(api.get) - - pipeline = make_pipeline() - user = make_account(id="u1") - - service = MagicMock() - service.get_all_published_workflow.return_value = ([make_workflow(id="w1")], False) - - with ( - app.test_request_context("/"), - patch( - "controllers.console.datasets.rag_pipeline.rag_pipeline_workflow.RagPipelineService", - return_value=service, - ), - ): - result = method(api, WorkflowListQuery(), user, pipeline) - - assert result["items"][0]["id"] == "w1" - assert result["items"][0]["graph"] == {"nodes": [], "edges": []} - assert result["has_more"] is False - - def test_get_published_workflows_forbidden(self, app: Flask) -> None: - api = PublishedAllRagPipelineApi() - method = unwrap(api.get) - - pipeline = make_pipeline() - user = make_account(id="u1") - - with ( - app.test_request_context("/?user_id=u2"), - ): - with pytest.raises(Forbidden): - method(api, WorkflowListQuery(user_id="u2"), user, pipeline) - - -class TestRagPipelineByIdApi: - def test_patch_success(self, app: Flask) -> None: - api = RagPipelineByIdApi() - method = unwrap(api.patch) - - pipeline = make_pipeline(tenant_id="t1") - user = make_account(id="u1") - - workflow = make_workflow(id="w1", marked_name="test") - - service = MagicMock() - service.update_workflow.return_value = workflow - - payload = {"marked_name": "test"} - - with ( - app.test_request_context("/", json=payload), - patch( - "controllers.console.datasets.rag_pipeline.rag_pipeline_workflow.RagPipelineService", - return_value=service, - ), - ): - result = method(api, WorkflowUpdatePayload.model_validate(payload), user, pipeline, "w1") - - assert result["id"] == "w1" - assert result["marked_name"] == "test" - assert result["hash"] == workflow.unique_hash - - def test_patch_no_fields(self, app: Flask) -> None: - api = RagPipelineByIdApi() - method = unwrap(api.patch) - - pipeline = make_pipeline() - user = make_account() - - with app.test_request_context("/", json={}): - result, status = raised_error(lambda: method(api, WorkflowUpdatePayload(), user, pipeline, "w1")) - assert status == 400 - assert result == {"code": "invalid_param", "message": "No valid fields to update"} - - @pytest.mark.parametrize("transaction_fails", [False, True], ids=["commit-succeeds", "commit-fails"]) - def test_delete_retires_candidates_only_after_transaction_exit( - self, - app: Flask, - transaction_fails: bool, - ) -> None: - api = RagPipelineByIdApi() - method = unwrap(api.delete) - - pipeline = make_pipeline(tenant_id="t1", workflow_id="active-workflow") - user = make_account() - - events: list[str] = [] - error = RuntimeError("commit failed") - workflow_service = MagicMock() - workflow_service.delete_workflow.side_effect = lambda **_kwargs: events.append("delete") or ["inline-agent"] - transaction_factory = MagicMock() - - @contextmanager - def transaction() -> Generator[Session]: - events.append("transaction-enter") - with Session() as session: - yield session - events.append("transaction-exit") - if transaction_fails: - raise error - - transaction_factory.begin.side_effect = transaction - - with ( - app.test_request_context("/", method="DELETE"), - patch( - "controllers.console.datasets.rag_pipeline.rag_pipeline_workflow.WorkflowService", - return_value=workflow_service, - ), - patch( - "controllers.console.datasets.rag_pipeline.rag_pipeline_workflow.sessionmaker", - return_value=transaction_factory, - ), - patch( - "controllers.console.datasets.rag_pipeline.rag_pipeline_workflow." - "WorkflowAgentRetirementService.retire_unowned" - ) as retire_unowned, - ): - retire_unowned.side_effect = lambda **_kwargs: events.append("retire") - if transaction_fails: - with pytest.raises(RuntimeError) as exc_info: - method(api, user, pipeline, "old-workflow") - assert exc_info.value is error - else: - result = method(api, user, pipeline, "old-workflow") - assert result == (None, 204) - - workflow_service.delete_workflow.assert_called_once() - assert events == ["transaction-enter", "delete", "transaction-exit"] + ([] if transaction_fails else ["retire"]) - if transaction_fails: - retire_unowned.assert_not_called() - else: - retire_unowned.assert_called_once_with( - tenant_id=pipeline.tenant_id, - agent_ids=["inline-agent"], - account_id=user.id, - ) - - def test_delete_active_workflow_rejected(self, app: Flask) -> None: - api = RagPipelineByIdApi() - method = unwrap(api.delete) - - pipeline = make_pipeline(tenant_id="t1", workflow_id="active-workflow") - user = make_account() - - with app.test_request_context("/", method="DELETE"): - with pytest.raises(BadRequest, match="currently in use by pipeline"): - method(api, user, pipeline, "active-workflow") - - -class TestRagPipelineWorkflowLastRunApi: - def test_last_run_success(self, app: Flask) -> None: - api = RagPipelineWorkflowLastRunApi() - method = unwrap(api.get) - - pipeline = make_pipeline() - workflow = make_workflow() - node_exec = make_node_execution() - - service = MagicMock() - service.get_draft_workflow.return_value = workflow - service.get_node_last_run.return_value = node_exec - - with ( - app.test_request_context("/"), - patch( - "controllers.console.datasets.rag_pipeline.rag_pipeline_workflow.RagPipelineService", - return_value=service, - ), - ): - result = method(api, pipeline, "node1") - assert result["id"] == "node-exec-1" - assert result["inputs"] == {"query": "hello"} - assert result["outputs"] == {"answer": "world"} - - def test_last_run_not_found(self, app: Flask) -> None: - api = RagPipelineWorkflowLastRunApi() - method = unwrap(api.get) - - pipeline = make_pipeline() - - service = MagicMock() - service.get_draft_workflow.return_value = None - - with ( - app.test_request_context("/"), - patch( - "controllers.console.datasets.rag_pipeline.rag_pipeline_workflow.RagPipelineService", - return_value=service, - ), - ): - with pytest.raises(NotFound): - method(api, pipeline, "node1") - - -class TestRagPipelineWorkflowRunNodeExecutionListApi: - def test_get_node_executions_passes_current_user(self, app: Flask) -> None: - api = RagPipelineWorkflowRunNodeExecutionListApi() - method = unwrap(api.get) - - user = make_account() - pipeline = make_pipeline() - run_id = uuid4() - node_exec = make_node_execution(workflow_run_id=str(run_id)) - - session_stub = MagicMock() - session_stub.scalar.return_value = None - service = MagicMock() - service.get_rag_pipeline_workflow_run_node_executions.return_value = [ - node_execution_response_source(node_exec, session=session_stub) - ] - - with ( - app.test_request_context("/"), - patch( - "controllers.console.datasets.rag_pipeline.rag_pipeline_workflow.RagPipelineService", - return_value=service, - ), - ): - result = method(api, user, pipeline, run_id) - - service.get_rag_pipeline_workflow_run_node_executions.assert_called_once_with( - pipeline=pipeline, - run_id=str(run_id), - user=user, - ) - assert result["data"][0]["id"] == "node-exec-1" - assert result["data"][0]["inputs"] == {"query": "hello"} - assert result["data"][0]["outputs"] == {"answer": "world"} - - -class TestRagPipelineDatasourceVariableApi: - def test_set_datasource_variables_success(self, app: Flask) -> None: - api = RagPipelineDatasourceVariableApi() - method = unwrap(api.post) - - pipeline = make_pipeline() - user = make_account() - - payload = { - "datasource_type": "db", - "datasource_info": empty_mapping(), - "start_node_id": "n1", - "start_node_title": "Node", - } - - service = MagicMock() - service.set_datasource_variables.return_value = make_node_execution(node_id="n1") - - with ( - app.test_request_context("/", json=payload), - patch( - "controllers.console.datasets.rag_pipeline.rag_pipeline_workflow.RagPipelineService", - return_value=service, - ), - ): - result = method(api, DatasourceVariablesPayload.model_validate(payload), user, pipeline) - assert result["node_id"] == "n1" - assert result["process_data"] == {} diff --git a/api/tests/unit_tests/controllers/console/datasets/test_application_boundaries.py b/api/tests/unit_tests/controllers/console/datasets/test_application_boundaries.py deleted file mode 100644 index 09841912ed0..00000000000 --- a/api/tests/unit_tests/controllers/console/datasets/test_application_boundaries.py +++ /dev/null @@ -1,62 +0,0 @@ -"""Guard all migrated dataset endpoints against transport/persistence regressions.""" - -import ast -from pathlib import Path - -import pytest - - -@pytest.mark.parametrize( - ("filename", "count"), - [("data_source.py", 7), ("datasets.py", 25), ("datasets_document.py", 24), ("external.py", 8)], -) -def test_every_endpoint_declares_admission_context_and_application_service(filename: str, count: int) -> None: - root = Path(__file__).resolve().parents[5] - tree = ast.parse((root / "controllers/console/datasets" / filename).read_text()) - # This namespace package is not traversed by import-linter's package graph. - forbidden = ( - "sqlalchemy", - "repositories", - "controllers.common.session", - "libs.login", - "services.knowledge.dataset_service", - "services.knowledge.external.service", - "services.hit_testing_service", - "services.knowledge.dataset_read_service", - "services.knowledge.datasets.adapters", - "services.knowledge.documents.adapters", - "services.knowledge.external.adapters", - ) - for node in ast.walk(tree): - modules = ( - [node.module or ""] - if isinstance(node, ast.ImportFrom) - else [item.name for item in node.names] - if isinstance(node, ast.Import) - else [] - ) - for module in modules: - assert not any(module == name or module.startswith(name + ".") for name in forbidden), module - endpoints = [ - method - for resource in tree.body - if isinstance(resource, ast.ClassDef) - for method in resource.body - if isinstance(method, ast.FunctionDef) and method.name in {"get", "post", "patch", "put", "delete"} - ] - assert len(endpoints) == count - for endpoint in endpoints: - assert any( - isinstance(d, ast.Call) and ast.unparse(d.func) == "console_account_admission" - for d in endpoint.decorator_list - ), endpoint.name - assert any( - arg.arg == "request_context" - and arg.annotation is not None - and ast.unparse(arg.annotation) == "RequestContext" - for arg in endpoint.args.args - ), endpoint.name - assert any( - isinstance(n, ast.Call) and ast.unparse(n.func) == "application_services" for n in ast.walk(endpoint) - ), endpoint.name - assert not {"session", "current_user", "current_tenant_id"}.intersection(arg.arg for arg in endpoint.args.args) diff --git a/api/tests/unit_tests/controllers/console/datasets/test_data_source.py b/api/tests/unit_tests/controllers/console/datasets/test_data_source.py index b5463712cbf..927f68c674d 100644 --- a/api/tests/unit_tests/controllers/console/datasets/test_data_source.py +++ b/api/tests/unit_tests/controllers/console/datasets/test_data_source.py @@ -1,9 +1,11 @@ from __future__ import annotations +import ast from collections.abc import Callable from dataclasses import dataclass, field from datetime import UTC, datetime from inspect import unwrap +from pathlib import Path from typing import TYPE_CHECKING, cast import pytest @@ -15,8 +17,6 @@ from controllers.common.errors import NotFoundError from controllers.common.rbac import DatasetId from controllers.console import console_ns from controllers.console.app.error import ProviderNotInitializeError - -# Importing the module registers only the resources under test on console_ns. from controllers.console.datasets import data_source as data_source_controller from controllers.console.datasets.error import IndexingEstimateError from core.entities.knowledge_entities import IndexingEstimate @@ -51,7 +51,72 @@ from services.knowledge.indexing.estimate import ( IndexingEstimateProviderUnavailableError, ) from services.knowledge.resource_scope import DatasetRef, DocumentRef -from tests.unit_tests.controllers.rbac_introspection import rbac_checks +from tests.helpers.unit.controllers.rbac_introspection import rbac_checks + +"""Guard all migrated dataset endpoints against transport/persistence regressions.""" + + +class TestApplicationBoundariesCases: + @staticmethod + @pytest.mark.parametrize( + ("filename", "count"), + [("data_source.py", 7), ("datasets.py", 25), ("datasets_document.py", 24), ("external.py", 8)], + ) + def test_every_endpoint_declares_admission_context_and_application_service(filename: str, count: int) -> None: + root = Path(__file__).resolve().parents[5] + tree = ast.parse((root / "controllers/console/datasets" / filename).read_text()) + # This namespace package is not traversed by import-linter's package graph. + forbidden = ( + "sqlalchemy", + "repositories", + "controllers.common.session", + "libs.login", + "services.knowledge.dataset_service", + "services.knowledge.external.service", + "services.hit_testing_service", + "services.knowledge.dataset_read_service", + "services.knowledge.datasets.adapters", + "services.knowledge.documents.adapters", + "services.knowledge.external.adapters", + ) + for node in ast.walk(tree): + modules = ( + [node.module or ""] + if isinstance(node, ast.ImportFrom) + else [item.name for item in node.names] + if isinstance(node, ast.Import) + else [] + ) + for module in modules: + assert not any(module == name or module.startswith(name + ".") for name in forbidden), module + endpoints = [ + method + for resource in tree.body + if isinstance(resource, ast.ClassDef) + for method in resource.body + if isinstance(method, ast.FunctionDef) and method.name in {"get", "post", "patch", "put", "delete"} + ] + assert len(endpoints) == count + for endpoint in endpoints: + assert any( + isinstance(d, ast.Call) and ast.unparse(d.func) == "console_account_admission" + for d in endpoint.decorator_list + ), endpoint.name + assert any( + arg.arg == "request_context" + and arg.annotation is not None + and ast.unparse(arg.annotation) == "RequestContext" + for arg in endpoint.args.args + ), endpoint.name + assert any( + isinstance(n, ast.Call) and ast.unparse(n.func) == "application_services" for n in ast.walk(endpoint) + ), endpoint.name + assert not {"session", "current_user", "current_tenant_id"}.intersection( + arg.arg for arg in endpoint.args.args + ) + + +# Importing the module registers only the resources under test on console_ns. if TYPE_CHECKING: from models import Account @@ -248,350 +313,354 @@ class AdmittedAccount: status: str -@pytest.fixture -def collaborators( - config_overrides: Callable[..., None], - monkeypatch: pytest.MonkeyPatch, -) -> ControllerCollaborators: - config_overrides( - DEPLOYMENT_EDITION=DeploymentEdition.CLOUD, - LOGIN_DISABLED=True, - RBAC_ENABLED=False, - CONSOLE_API_URL="https://console-api.example.com", - ) - app = Flask(__name__) - app.config.update(TESTING=False, PROPAGATE_EXCEPTIONS=False) - Api(app).add_namespace(console_ns, path="/console/api") +class TestDataSourceCases: + @staticmethod + @pytest.fixture + def collaborators( + config_overrides: Callable[..., None], + monkeypatch: pytest.MonkeyPatch, + ) -> ControllerCollaborators: + config_overrides( + DEPLOYMENT_EDITION=DeploymentEdition.CLOUD, + LOGIN_DISABLED=True, + RBAC_ENABLED=False, + CONSOLE_API_URL="https://console-api.example.com", + ) + app = Flask(__name__) + app.config.update(TESTING=False, PROPAGATE_EXCEPTIONS=False) + Api(app).add_namespace(console_ns, path="/console/api") - account = AdmittedAccount(id=ACCOUNT_ID, role="owner", status="active") - account_with_tenant = AccountWithTenant(account=cast("Account", account), tenant_id=TENANT_ID) - monkeypatch.setattr( - "controllers.console.flask_admission.current_account_with_tenant", - lambda: account_with_tenant, - ) - monkeypatch.setattr( - "controllers.console.wraps.current_account_with_tenant", - lambda: account_with_tenant, - ) + account = AdmittedAccount(id=ACCOUNT_ID, role="owner", status="active") + account_with_tenant = AccountWithTenant(account=cast("Account", account), tenant_id=TENANT_ID) + monkeypatch.setattr( + "controllers.console.flask_admission.current_account_with_tenant", + lambda: account_with_tenant, + ) + monkeypatch.setattr( + "controllers.console.wraps.current_account_with_tenant", + lambda: account_with_tenant, + ) - dataset = DatasetAccessRecord( - id=DATASET_ID, - workspace_id=TENANT_ID, - maintainer_id=ACCOUNT_ID, - permission="only_me", - ) - dataset_catalog = DatasetCatalog(dataset) - dataset_access = DatasetAccessService( - datasets=dataset_catalog, - workspace_roles=OwnerWorkspaceRoles(), - legacy_permissions_enabled=True, - ) - documents = KnowledgeDocuments( - bound_page_ids=frozenset({PAGE_ID}), - active_refs=(DatasetRef(TENANT_ID, DATASET_ID).document(DOCUMENT_ID),), - ) - bindings = RecordingBindingStore() - notion_source = RecordingNotionSource( - workspaces=( - NotionWorkspace( - workspace_id="notion-workspace-1", - workspace_name="One", - workspace_icon=None, - pages=( - AuthorizedNotionPage( - page_id=PAGE_ID, - page_name="Page one", - page_icon=NotionPageIcon(type="emoji", emoji="📄"), - parent_id=None, - page_type=NotionPageType.PAGE, + dataset = DatasetAccessRecord( + id=DATASET_ID, + workspace_id=TENANT_ID, + maintainer_id=ACCOUNT_ID, + permission="only_me", + ) + dataset_catalog = DatasetCatalog(dataset) + dataset_access = DatasetAccessService( + datasets=dataset_catalog, + workspace_roles=OwnerWorkspaceRoles(), + legacy_permissions_enabled=True, + ) + documents = KnowledgeDocuments( + bound_page_ids=frozenset({PAGE_ID}), + active_refs=(DatasetRef(TENANT_ID, DATASET_ID).document(DOCUMENT_ID),), + ) + bindings = RecordingBindingStore() + notion_source = RecordingNotionSource( + workspaces=( + NotionWorkspace( + workspace_id="notion-workspace-1", + workspace_name="One", + workspace_icon=None, + pages=( + AuthorizedNotionPage( + page_id=PAGE_ID, + page_name="Page one", + page_icon=NotionPageIcon(type="emoji", emoji="📄"), + parent_id=None, + page_type=NotionPageType.PAGE, + ), ), ), + NotionWorkspace( + workspace_id="notion-workspace-2", + workspace_name="Two", + workspace_icon=None, + pages=(), + ), + ) + ) + estimates = RecordingEstimateService() + dispatcher = RecordingDispatcher() + app.extensions["application_services"] = ApplicationServiceRegistry( + data_sources=DataSourceApplicationServices( + bindings=DataSourceBindingApplicationService(bindings=bindings), + notion_imports=NotionImportApplicationService( + dataset_access=dataset_access, + datasets=dataset_catalog, + documents=documents, + source=notion_source, + ), ), - NotionWorkspace( - workspace_id="notion-workspace-2", - workspace_name="Two", - workspace_icon=None, - pages=(), + knowledge=KnowledgeApplicationServices( + document_sync=DocumentSyncApplicationService( + dataset_access=dataset_access, + documents=documents, + dispatcher=dispatcher, + ), + indexing_estimates=estimates, ), ) + return ControllerCollaborators( + client=app.test_client(), + account=account, + bindings=bindings, + notion_source=notion_source, + estimates=estimates, + dispatcher=dispatcher, + ) + + @staticmethod + def test_integration_collection_only_exposes_get(collaborators: ControllerCollaborators) -> None: + client = collaborators.client + assert client.get("/console/api/data-source/integrates").status_code == 200 + assert client.patch("/console/api/data-source/integrates").status_code == 405 + + @staticmethod + @pytest.mark.parametrize( + ("method", "path"), + [ + ("get", "/console/api/data-source/integrates"), + ("patch", f"/console/api/data-source/integrates/{BINDING_ID}/enable"), + ], ) - estimates = RecordingEstimateService() - dispatcher = RecordingDispatcher() - app.extensions["application_services"] = ApplicationServiceRegistry( - data_sources=DataSourceApplicationServices( - bindings=DataSourceBindingApplicationService(bindings=bindings), - notion_imports=NotionImportApplicationService( - dataset_access=dataset_access, - datasets=dataset_catalog, - documents=documents, - source=notion_source, + def test_integration_management_rejects_non_admin_role( + collaborators: ControllerCollaborators, + method: str, + path: str, + ) -> None: + collaborators.account.role = "editor" + + assert collaborators.client.open(path, method=method.upper()).status_code == 403 + + @staticmethod + def test_integration_management_enforces_workspace_credential_permission( + collaborators: ControllerCollaborators, + config_overrides: Callable[..., None], + monkeypatch: pytest.MonkeyPatch, + ) -> None: + config_overrides(RBAC_ENABLED=True) + calls: list[dict[str, object]] = [] + + def record_access(tenant_id: str, account_id: str, **requirement: object) -> bool: + calls.append({"tenant_id": tenant_id, "account_id": account_id, **requirement}) + return True + + monkeypatch.setattr("controllers.common.rbac.checks.RBACService.CheckAccess.check", record_access) + + list_response = collaborators.client.get("/console/api/data-source/integrates") + mutation_response = collaborators.client.patch(f"/console/api/data-source/integrates/{BINDING_ID}/enable") + + assert list_response.status_code == 200 + assert mutation_response.status_code == 200 + assert calls == [ + { + "tenant_id": TENANT_ID, + "account_id": ACCOUNT_ID, + "scene": RBACPermission.CREDENTIAL_MANAGE, + "resource_type": None, + "resource_id": None, + }, + { + "tenant_id": TENANT_ID, + "account_id": ACCOUNT_ID, + "scene": RBACPermission.CREDENTIAL_MANAGE, + "resource_type": None, + "resource_id": None, + }, + ] + + @staticmethod + @pytest.mark.parametrize(("action", "disabled"), [("enable", False), ("disable", True)]) + def test_integration_item_only_exposes_patch_and_rejects_unknown_action( + collaborators: ControllerCollaborators, + action: str, + disabled: bool, + ) -> None: + client = collaborators.client + item_path = f"/console/api/data-source/integrates/{BINDING_ID}/{action}" + assert client.get(item_path).status_code == 405 + response = client.patch(item_path) + assert response.status_code == 200 + assert response.get_json() == {"result": "success"} + assert collaborators.bindings.mutations == [(TENANT_ID, BINDING_ID, disabled)] + invalid = client.patch(f"/console/api/data-source/integrates/{BINDING_ID}/remove") + assert invalid.status_code == 400 + assert invalid.get_json()["code"] == "invalid_action" + assert len(collaborators.bindings.mutations) == 1 + + @staticmethod + def test_integration_list_serializes_application_result_and_http_link( + collaborators: ControllerCollaborators, + ) -> None: + collaborators.bindings.bindings = ( + DataSourceBindingSummary( + id=BINDING_ID, + provider="notion", + created_at=datetime(2026, 5, 25, 1, 2, 3, tzinfo=UTC), + disabled=False, + source_info={ + "workspace_name": "Workspace", + "workspace_id": "workspace-1", + "workspace_icon": None, + "pages": [], + "total": 0, + }, ), - ), - knowledge=KnowledgeApplicationServices( - document_sync=DocumentSyncApplicationService( - dataset_access=dataset_access, - documents=documents, - dispatcher=dispatcher, - ), - indexing_estimates=estimates, - ), - ) - return ControllerCollaborators( - client=app.test_client(), - account=account, - bindings=bindings, - notion_source=notion_source, - estimates=estimates, - dispatcher=dispatcher, - ) + ) + response = collaborators.client.get("/console/api/data-source/integrates") -def test_integration_collection_only_exposes_get(collaborators: ControllerCollaborators) -> None: - client = collaborators.client - assert client.get("/console/api/data-source/integrates").status_code == 200 - assert client.patch("/console/api/data-source/integrates").status_code == 405 - - -@pytest.mark.parametrize( - ("method", "path"), - [ - ("get", "/console/api/data-source/integrates"), - ("patch", f"/console/api/data-source/integrates/{BINDING_ID}/enable"), - ], -) -def test_integration_management_rejects_non_admin_role( - collaborators: ControllerCollaborators, - method: str, - path: str, -) -> None: - collaborators.account.role = "editor" - - assert collaborators.client.open(path, method=method.upper()).status_code == 403 - - -def test_integration_management_enforces_workspace_credential_permission( - collaborators: ControllerCollaborators, - config_overrides: Callable[..., None], - monkeypatch: pytest.MonkeyPatch, -) -> None: - config_overrides(RBAC_ENABLED=True) - calls: list[dict[str, object]] = [] - - def record_access(tenant_id: str, account_id: str, **requirement: object) -> bool: - calls.append({"tenant_id": tenant_id, "account_id": account_id, **requirement}) - return True - - monkeypatch.setattr("controllers.common.rbac.checks.RBACService.CheckAccess.check", record_access) - - list_response = collaborators.client.get("/console/api/data-source/integrates") - mutation_response = collaborators.client.patch(f"/console/api/data-source/integrates/{BINDING_ID}/enable") - - assert list_response.status_code == 200 - assert mutation_response.status_code == 200 - assert calls == [ - { - "tenant_id": TENANT_ID, - "account_id": ACCOUNT_ID, - "scene": RBACPermission.CREDENTIAL_MANAGE, - "resource_type": None, - "resource_id": None, - }, - { - "tenant_id": TENANT_ID, - "account_id": ACCOUNT_ID, - "scene": RBACPermission.CREDENTIAL_MANAGE, - "resource_type": None, - "resource_id": None, - }, - ] - - -@pytest.mark.parametrize(("action", "disabled"), [("enable", False), ("disable", True)]) -def test_integration_item_only_exposes_patch_and_rejects_unknown_action( - collaborators: ControllerCollaborators, - action: str, - disabled: bool, -) -> None: - client = collaborators.client - item_path = f"/console/api/data-source/integrates/{BINDING_ID}/{action}" - assert client.get(item_path).status_code == 405 - response = client.patch(item_path) - assert response.status_code == 200 - assert response.get_json() == {"result": "success"} - assert collaborators.bindings.mutations == [(TENANT_ID, BINDING_ID, disabled)] - invalid = client.patch(f"/console/api/data-source/integrates/{BINDING_ID}/remove") - assert invalid.status_code == 400 - assert invalid.get_json()["code"] == "invalid_action" - assert len(collaborators.bindings.mutations) == 1 - - -def test_integration_list_serializes_application_result_and_http_link( - collaborators: ControllerCollaborators, -) -> None: - collaborators.bindings.bindings = ( - DataSourceBindingSummary( - id=BINDING_ID, - provider="notion", - created_at=datetime(2026, 5, 25, 1, 2, 3, tzinfo=UTC), - disabled=False, - source_info={ + assert response.status_code == 200 + assert response.get_json()["data"][0] == { + "id": BINDING_ID, + "provider": "notion", + "created_at": 1779670923, + "is_bound": True, + "disabled": False, + "source_info": { "workspace_name": "Workspace", "workspace_id": "workspace-1", "workspace_icon": None, "pages": [], "total": 0, }, - ), - ) - - response = collaborators.client.get("/console/api/data-source/integrates") - - assert response.status_code == 200 - assert response.get_json()["data"][0] == { - "id": BINDING_ID, - "provider": "notion", - "created_at": 1779670923, - "is_bound": True, - "disabled": False, - "source_info": { - "workspace_name": "Workspace", - "workspace_id": "workspace-1", - "workspace_icon": None, - "pages": [], - "total": 0, - }, - "link": "https://console-api.example.com/console/api/oauth/data-source/notion", - } - - -def test_notion_list_preserves_workspace_groups_and_bound_state(collaborators: ControllerCollaborators) -> None: - response = collaborators.client.get( - "/console/api/notion/pre-import/pages", - query_string={"credential_id": "credential-1", "dataset_id": DATASET_ID}, - ) - - assert response.status_code == 200 - payload = response.get_json()["notion_info"] - assert [workspace["workspace_id"] for workspace in payload] == [ - "notion-workspace-1", - "notion-workspace-2", - ] - assert payload[0]["pages"][0]["is_bound"] is True - assert payload[0]["pages"][0]["type"] == "page" - assert "page_type" not in payload[0]["pages"][0] - assert collaborators.notion_source.list_calls == [(TENANT_ID, ACCOUNT_ID, "credential-1")] - - -@pytest.mark.parametrize("page_type", ["page", "database"]) -def test_notion_preview_validates_page_type_and_delegates( - collaborators: ControllerCollaborators, page_type: str -) -> None: - response = collaborators.client.get( - f"/console/api/notion/pages/{PAGE_ID}/{page_type}/preview", - query_string={"credential_id": "credential-1"}, - ) - invalid = collaborators.client.get( - f"/console/api/notion/pages/{PAGE_ID}/collection/preview", - query_string={"credential_id": "credential-1"}, - ) - - assert response.status_code == 200 - assert response.get_json() == {"content": "preview text"} - assert invalid.status_code == 400 - assert invalid.get_json()["code"] == "invalid_param" - assert collaborators.notion_source.preview_calls == [(TENANT_ID, ACCOUNT_ID, "credential-1", PAGE_ID, page_type)] - - -def test_notion_estimate_maps_nested_payload_to_shared_command(collaborators: ControllerCollaborators) -> None: - response = collaborators.client.post( - "/console/api/datasets/notion-indexing-estimate", - json={ - "notion_info_list": [ - { - "workspace_id": "notion-workspace-1", - "credential_id": "credential-1", - "pages": [{"page_id": PAGE_ID, "type": "page"}], - } - ], - "process_rule": {"mode": "automatic"}, - "doc_form": "text_model", - "doc_language": "English", - }, - ) - - assert response.status_code == 200 - assert response.get_json()["total_segments"] == 3 - context, command = collaborators.estimates.calls[0] - assert context.active_workspace_id == TENANT_ID - assert context.account_id == ACCOUNT_ID - assert command.process_rule == {"mode": "automatic"} - source = command.sources[0] - assert isinstance(source, NotionEstimateSource) - assert source.credential_id == "credential-1" - assert source.page_id == PAGE_ID - - -@pytest.mark.parametrize( - ("error", "expected_http_error"), - [ - (IndexingEstimateCredentialUnavailableError(), NotFoundError), - (EstimateSourceNotFoundError("page-1"), NotFoundError), - (IndexingEstimateProviderUnavailableError(), ProviderNotInitializeError), - (IndexingEstimateExecutionError(), IndexingEstimateError), - ], -) -def test_notion_estimate_maps_application_errors( - error: Exception, - expected_http_error: type[Exception], - collaborators: ControllerCollaborators, -) -> None: - collaborators.estimates.error = error - api = data_source_controller.DataSourceNotionIndexingEstimateApi() - method = unwrap(api.post) - payload = data_source_controller.NotionEstimatePayload.model_validate( - { - "notion_info_list": [ - { - "workspace_id": "notion-workspace", - "credential_id": "credential-1", - "pages": [{"page_id": "page-1", "type": "page"}], - } - ], - "process_rule": {"mode": "automatic"}, + "link": "https://console-api.example.com/console/api/oauth/data-source/notion", } + + @staticmethod + def test_notion_list_preserves_workspace_groups_and_bound_state(collaborators: ControllerCollaborators) -> None: + response = collaborators.client.get( + "/console/api/notion/pre-import/pages", + query_string={"credential_id": "credential-1", "dataset_id": DATASET_ID}, + ) + + assert response.status_code == 200 + payload = response.get_json()["notion_info"] + assert [workspace["workspace_id"] for workspace in payload] == [ + "notion-workspace-1", + "notion-workspace-2", + ] + assert payload[0]["pages"][0]["is_bound"] is True + assert payload[0]["pages"][0]["type"] == "page" + assert "page_type" not in payload[0]["pages"][0] + assert collaborators.notion_source.list_calls == [(TENANT_ID, ACCOUNT_ID, "credential-1")] + + @staticmethod + @pytest.mark.parametrize("page_type", ["page", "database"]) + def test_notion_preview_validates_page_type_and_delegates( + collaborators: ControllerCollaborators, page_type: str + ) -> None: + response = collaborators.client.get( + f"/console/api/notion/pages/{PAGE_ID}/{page_type}/preview", + query_string={"credential_id": "credential-1"}, + ) + invalid = collaborators.client.get( + f"/console/api/notion/pages/{PAGE_ID}/collection/preview", + query_string={"credential_id": "credential-1"}, + ) + + assert response.status_code == 200 + assert response.get_json() == {"content": "preview text"} + assert invalid.status_code == 400 + assert invalid.get_json()["code"] == "invalid_param" + assert collaborators.notion_source.preview_calls == [ + (TENANT_ID, ACCOUNT_ID, "credential-1", PAGE_ID, page_type) + ] + + @staticmethod + def test_notion_estimate_maps_nested_payload_to_shared_command(collaborators: ControllerCollaborators) -> None: + response = collaborators.client.post( + "/console/api/datasets/notion-indexing-estimate", + json={ + "notion_info_list": [ + { + "workspace_id": "notion-workspace-1", + "credential_id": "credential-1", + "pages": [{"page_id": PAGE_ID, "type": "page"}], + } + ], + "process_rule": {"mode": "automatic"}, + "doc_form": "text_model", + "doc_language": "English", + }, + ) + + assert response.status_code == 200 + assert response.get_json()["total_segments"] == 3 + context, command = collaborators.estimates.calls[0] + assert context.active_workspace_id == TENANT_ID + assert context.account_id == ACCOUNT_ID + assert command.process_rule == {"mode": "automatic"} + source = command.sources[0] + assert isinstance(source, NotionEstimateSource) + assert source.credential_id == "credential-1" + assert source.page_id == PAGE_ID + + @staticmethod + @pytest.mark.parametrize( + ("error", "expected_http_error"), + [ + (IndexingEstimateCredentialUnavailableError(), NotFoundError), + (EstimateSourceNotFoundError("page-1"), NotFoundError), + (IndexingEstimateProviderUnavailableError(), ProviderNotInitializeError), + (IndexingEstimateExecutionError(), IndexingEstimateError), + ], ) - context = RequestContext("request-1", None, "account-1", "workspace-1") + def test_notion_estimate_maps_application_errors( + error: Exception, + expected_http_error: type[Exception], + collaborators: ControllerCollaborators, + ) -> None: + collaborators.estimates.error = error + api = data_source_controller.DataSourceNotionIndexingEstimateApi() + method = unwrap(api.post) + payload = data_source_controller.NotionEstimatePayload.model_validate( + { + "notion_info_list": [ + { + "workspace_id": "notion-workspace", + "credential_id": "credential-1", + "pages": [{"page_id": "page-1", "type": "page"}], + } + ], + "process_rule": {"mode": "automatic"}, + } + ) + context = RequestContext("request-1", None, "account-1", "workspace-1") - with collaborators.client.application.app_context(), pytest.raises(expected_http_error): - method(api, payload, context) + with collaborators.client.application.app_context(), pytest.raises(expected_http_error): + method(api, payload, context) + @staticmethod + def test_sync_routes_dispatch_only_application_selected_documents(collaborators: ControllerCollaborators) -> None: + dataset_response = collaborators.client.get(f"/console/api/datasets/{DATASET_ID}/notion/sync") + document_response = collaborators.client.get( + f"/console/api/datasets/{DATASET_ID}/documents/{DOCUMENT_ID}/notion/sync" + ) -def test_sync_routes_dispatch_only_application_selected_documents(collaborators: ControllerCollaborators) -> None: - dataset_response = collaborators.client.get(f"/console/api/datasets/{DATASET_ID}/notion/sync") - document_response = collaborators.client.get( - f"/console/api/datasets/{DATASET_ID}/documents/{DOCUMENT_ID}/notion/sync" + assert dataset_response.status_code == 200 + assert dataset_response.get_json() == {"result": "success"} + assert document_response.status_code == 200 + assert document_response.get_json() == {"result": "success"} + assert collaborators.dispatcher.calls == [(DATASET_ID, DOCUMENT_ID), (DATASET_ID, DOCUMENT_ID)] + + @staticmethod + def test_sync_route_does_not_dispatch_cross_workspace_dataset(collaborators: ControllerCollaborators) -> None: + response = collaborators.client.get(f"/console/api/datasets/{FOREIGN_DATASET_ID}/notion/sync") + + assert response.status_code == 404 + assert response.get_json()["code"] == "not_found" + assert collaborators.dispatcher.calls == [] + + @staticmethod + @pytest.mark.parametrize( + "resource", + [data_source_controller.DataSourceNotionDatasetSyncApi, data_source_controller.DataSourceNotionDocumentSyncApi], ) - - assert dataset_response.status_code == 200 - assert dataset_response.get_json() == {"result": "success"} - assert document_response.status_code == 200 - assert document_response.get_json() == {"result": "success"} - assert collaborators.dispatcher.calls == [(DATASET_ID, DOCUMENT_ID), (DATASET_ID, DOCUMENT_ID)] - - -def test_sync_route_does_not_dispatch_cross_workspace_dataset(collaborators: ControllerCollaborators) -> None: - response = collaborators.client.get(f"/console/api/datasets/{FOREIGN_DATASET_ID}/notion/sync") - - assert response.status_code == 404 - assert response.get_json()["code"] == "not_found" - assert collaborators.dispatcher.calls == [] - - -@pytest.mark.parametrize( - "resource", - [data_source_controller.DataSourceNotionDatasetSyncApi, data_source_controller.DataSourceNotionDocumentSyncApi], -) -def test_notion_sync_requires_dataset_edit_permission(resource): - [check] = rbac_checks(resource.get) - assert check.scene is RBACPermission.DATASET_EDIT - assert isinstance(check.locator, DatasetId) + def test_notion_sync_requires_dataset_edit_permission(resource): + [check] = rbac_checks(resource.get) + assert check.scene is RBACPermission.DATASET_EDIT + assert isinstance(check.locator, DatasetId) diff --git a/api/tests/unit_tests/controllers/console/datasets/test_datasets.py b/api/tests/unit_tests/controllers/console/datasets/test_datasets.py index 995f99e4b7b..18e48bcdc4c 100644 --- a/api/tests/unit_tests/controllers/console/datasets/test_datasets.py +++ b/api/tests/unit_tests/controllers/console/datasets/test_datasets.py @@ -68,7 +68,7 @@ from services.knowledge.indexing.estimate import ( UnsupportedEstimateSourceError, ) from services.knowledge_fs.product_dto import KnowledgeFSUpgradeDiscoveryResponse -from tests.unit_tests.controllers.rbac_introspection import rbac_checks +from tests.helpers.unit.controllers.rbac_introspection import rbac_checks CONTEXT = RequestContext("request-1", None, "account-1", "tenant-1") DATASET_ID = UUID(int=1) diff --git a/api/tests/unit_tests/controllers/console/datasets/test_datasets_document.py b/api/tests/unit_tests/controllers/console/datasets/test_datasets_document.py index b4b33fe7528..ff2e0e1d100 100644 --- a/api/tests/unit_tests/controllers/console/datasets/test_datasets_document.py +++ b/api/tests/unit_tests/controllers/console/datasets/test_datasets_document.py @@ -1,10 +1,17 @@ -import datetime +import datetime as _datasets_document_datetime +import json +from datetime import datetime as _datasets_document_download_datetime from inspect import unwrap +from io import BytesIO +from pathlib import Path from types import SimpleNamespace from unittest.mock import MagicMock, create_autospec, patch from uuid import UUID +from zipfile import ZipFile import pytest +from flask import send_file +from sqlalchemy import event from werkzeug.exceptions import Forbidden, NotFound from controllers.common.errors import InvalidArgumentError, NotFoundError @@ -48,7 +55,15 @@ from controllers.console.datasets.error import ( InvalidMetadataError, ) from machinery.context import RequestContext -from services.knowledge.dataset_access import DatasetAccessDeniedError, DatasetNotFoundError +from models.dataset import Dataset, Document +from models.model import UploadFile +from services.knowledge.dataset_access import ( + AccessibleDataset, + DatasetAccess, + DatasetAccessDeniedError, + DatasetNotFoundError, +) +from services.knowledge.documents.adapters import SQLAlchemyDocumentOperations from services.knowledge.documents.application import ( DatasetDocumentApplicationService, DocumentArchivedError, @@ -67,503 +82,703 @@ from services.knowledge.indexing.estimate import ( IndexingEstimateProviderUnavailableError, UnsupportedEstimateSourceError, ) -from tests.unit_tests.controllers.rbac_introspection import rbac_checks +from tests.helpers.unit.controllers.rbac_introspection import rbac_checks -CONTEXT = RequestContext("request-1", None, "account-1", "tenant-1") +_DATASETS_DOCUMENT_CONTEXT = RequestContext("request-1", None, "account-1", "tenant-1") DS = UUID(int=1) DOC = UUID(int=2) -@pytest.fixture -def documents(monkeypatch): - documents = create_autospec(DatasetDocumentApplicationService, instance=True, spec_set=True) - monkeypatch.setattr( - controller, "application_services", lambda: SimpleNamespace(knowledge=SimpleNamespace(documents=documents)) - ) - monkeypatch.setattr(controller, "check_knowledge_rate_limit", lambda: None) - return documents - - -def test_process_rule_authorizes_document_dataset_before_read(app, documents): - with ( - app.test_request_context(f"/?document_id={DOC}"), - patch.object(controller, "enforce_rbac_checks", side_effect=Forbidden) as enforce_checks, - pytest.raises(Forbidden), - ): - unwrap(GetProcessRuleApi.get)(GetProcessRuleApi(), CONTEXT) - - enforce_checks.assert_called_once() - kwargs = enforce_checks.call_args.kwargs - assert kwargs["tenant_id"] == CONTEXT.active_workspace_id - assert kwargs["account_id"] == CONTEXT.account_id - assert kwargs["path_args"] == {"document_id": str(DOC)} - [check] = kwargs["checks"] - assert check.scene is RBACPermission.DATASET_READONLY - assert isinstance(check.locator, DatasetByDocument) - documents.get_process_rule.assert_not_called() - - -def test_default_process_rule_does_not_require_dataset_permission(app, documents): - documents.get_process_rule.return_value = { - "mode": "automatic", - "rules": {}, - "limits": {"indexing_max_segmentation_tokens_length": 4000}, - } - with app.test_request_context("/"), patch.object(controller, "enforce_rbac_checks") as enforce_checks: - unwrap(GetProcessRuleApi.get)(GetProcessRuleApi(), CONTEXT) - enforce_checks.assert_not_called() - documents.get_process_rule.assert_called_once_with(CONTEXT, document_id=None) - - -@pytest.mark.parametrize( - ("resource", "verb", "operation", "kwargs"), - [ - (GetProcessRuleApi, GetProcessRuleApi.get, lambda service: service.get_process_rule, {}), - ( - DatasetDocumentListApi, - DatasetDocumentListApi.get, - lambda service: service.list_documents, - {"dataset_id": DS}, - ), - ( - DatasetDocumentListApi, - DatasetDocumentListApi.post, - lambda service: service.create_documents, - {"dataset_id": DS}, - ), - ( - DatasetDocumentListApi, - DatasetDocumentListApi.delete, - lambda service: service.delete_documents, - {"dataset_id": DS}, - ), - (DatasetInitApi, DatasetInitApi.post, lambda service: service.initialize_dataset, {}), - ( - DocumentBatchIndexingStatusApi, - DocumentBatchIndexingStatusApi.get, - lambda service: service.get_batch_indexing_status, - {"dataset_id": DS, "batch": "batch-1"}, - ), - ( - DocumentIndexingStatusApi, - DocumentIndexingStatusApi.get, - lambda service: service.get_indexing_status, - {"dataset_id": DS, "document_id": DOC}, - ), - (DocumentApi, DocumentApi.get, lambda service: service.get_document, {"dataset_id": DS, "document_id": DOC}), - ( - DocumentApi, - DocumentApi.delete, - lambda service: service.delete_document, - {"dataset_id": DS, "document_id": DOC}, - ), - ( - DocumentDownloadApi, - DocumentDownloadApi.get, - lambda service: service.get_download_url, - {"dataset_id": DS, "document_id": DOC}, - ), - ( - DocumentBatchDownloadZipApi, - DocumentBatchDownloadZipApi.post, - lambda service: service.build_download_zip, - {"dataset_id": DS}, - ), - ( - DocumentProcessingApi, - DocumentProcessingApi.patch, - lambda service: service.update_processing, - {"dataset_id": DS, "document_id": DOC, "action": "pause"}, - ), - ( - DocumentMetadataApi, - DocumentMetadataApi.put, - lambda service: service.update_metadata, - { - "dataset_id": DS, - "document_id": DOC, - "req_data": DocumentMetadataUpdatePayload(doc_type="book", doc_metadata={"title": "Book"}), - }, - ), - ( - DocumentStatusApi, - DocumentStatusApi.patch, - lambda service: service.change_status, - {"dataset_id": DS, "action": "enable"}, - ), - ( - DocumentPauseApi, - DocumentPauseApi.patch, - lambda service: service.pause_document, - {"dataset_id": DS, "document_id": DOC}, - ), - ( - DocumentRecoverApi, - DocumentRecoverApi.patch, - lambda service: service.recover_document, - {"dataset_id": DS, "document_id": DOC}, - ), - ( - DocumentRetryApi, - DocumentRetryApi.post, - lambda service: service.retry_documents, - {"dataset_id": DS, "req_data": DocumentRetryPayload(document_ids=[str(DOC)])}, - ), - ( - DocumentRenameApi, - DocumentRenameApi.post, - lambda service: service.rename_document, - {"dataset_id": DS, "document_id": DOC, "req_data": DocumentRenamePayload(name="New name")}, - ), - ( - WebsiteDocumentSyncApi, - WebsiteDocumentSyncApi.get, - lambda service: service.sync_website, - {"dataset_id": DS, "document_id": DOC}, - ), - ( - DocumentPipelineExecutionLogApi, - DocumentPipelineExecutionLogApi.get, - lambda service: service.get_execution_log, - {"dataset_id": DS, "document_id": DOC}, - ), - ( - DocumentGenerateSummaryApi, - DocumentGenerateSummaryApi.post, - lambda service: service.generate_summary, - {"dataset_id": DS, "req_data": GenerateSummaryPayload(document_list=[str(DOC)])}, - ), - ( - DocumentSummaryStatusApi, - DocumentSummaryStatusApi.get, - lambda service: service.get_summary_status, - {"dataset_id": DS, "document_id": DOC}, - ), - ], -) -def test_endpoints_pass_request_context_and_map_application_errors(app, documents, resource, verb, operation, kwargs): - operation(documents).side_effect = DocumentNotFoundError("Document not found.") - with app.test_request_context( - "/?document_id=" + str(DOC), json={"indexing_technique": "economy", "document_ids": [str(DOC)]} - ): - with pytest.raises(NotFound): - unwrap(verb)(resource(), request_context=CONTEXT, **kwargs) - called = operation(documents) - assert called.call_count == 1 - assert called.call_args.args == (CONTEXT,) - if "dataset_id" in kwargs: - assert called.call_args.kwargs["dataset_id"] == str(DS) - if "document_id" in kwargs: - assert called.call_args.kwargs["document_id"] == str(DOC) - - -@pytest.mark.parametrize(("fetch", "expected"), [("true", True), ("YES", True), ("false", False), ("invalid", False)]) -def test_list_preserves_filters_pagination_and_fetch_parsing(app, documents, fetch, expected): - documents.list_documents.return_value = {"data": [], "has_more": False, "page": 2, "limit": 3, "total": 3} - with app.test_request_context(f"/?page=2&limit=3&keyword=needle&status=available&sort=hit_count&fetch={fetch}"): - response = unwrap(DatasetDocumentListApi.get)(DatasetDocumentListApi(), CONTEXT, DS) - documents.list_documents.assert_called_once_with( - CONTEXT, - dataset_id=str(DS), - query=DocumentListFilter( - page=2, limit=3, search="needle", status="available", sort="hit_count", fetch=expected - ), - ) - assert response == {"data": [], "has_more": False, "page": 2, "limit": 3, "total": 3} - - -@pytest.mark.parametrize("metadata", ["only", "without", "all"]) -def test_document_detail_metadata_projection_and_raw_source(app, documents, metadata): - documents.get_document.return_value = { - "id": str(DOC), - "doc_type": "others", - "doc_metadata": [], - "name": "Name", - "data_source_info": {"file_path": "/path/to/local/file"}, - "data_source_detail_dict": {}, - } - with app.test_request_context(f"/?metadata={metadata}"): - response, status = unwrap(DocumentApi.get)(DocumentApi(), CONTEXT, DS, DOC) - assert status == 200 - documents.get_document.assert_called_once_with( - CONTEXT, dataset_id=str(DS), document_id=str(DOC), metadata_only=metadata == "only" - ) - if metadata == "only": - assert response == {"id": str(DOC), "doc_type": "others", "doc_metadata": []} - else: - assert response["data_source_info"] == {"file_path": "/path/to/local/file"} - assert ("doc_metadata" in response) is (metadata == "all") - - -def test_invalid_metadata_fails_before_application_call(app, documents): - with app.test_request_context("/?metadata=invalid"), pytest.raises(InvalidMetadataError): - unwrap(DocumentApi.get)(DocumentApi(), CONTEXT, DS, DOC) - documents.get_document.assert_not_called() - - -@pytest.mark.parametrize( - ("resource", "verb", "kwargs"), - [ - (DocumentApi, DocumentApi.delete, {"document_id": DOC}), - (DatasetDocumentListApi, DatasetDocumentListApi.delete, {}), - (DocumentPauseApi, DocumentPauseApi.patch, {"document_id": DOC}), - (DocumentRecoverApi, DocumentRecoverApi.patch, {"document_id": DOC}), - (DocumentRetryApi, DocumentRetryApi.post, {"req_data": DocumentRetryPayload(document_ids=[str(DOC)])}), - ], -) -@pytest.mark.usefixtures("documents") -def test_no_content_mutations_return_empty_204(app, resource, verb, kwargs): - with app.test_request_context("/?document_id=" + str(DOC)): - result = unwrap(verb)(resource(), request_context=CONTEXT, dataset_id=DS, **kwargs) - assert result == ("", 204) - - -@pytest.mark.parametrize( - ("resource", "verb", "operation", "kwargs"), - [ - ( - DocumentProcessingApi, - DocumentProcessingApi.patch, - lambda service: service.update_processing, - {"document_id": DOC, "action": "resume"}, - ), - (DocumentStatusApi, DocumentStatusApi.patch, lambda service: service.change_status, {"action": "disable"}), - ( - WebsiteDocumentSyncApi, - WebsiteDocumentSyncApi.get, - lambda service: service.sync_website, - {"document_id": DOC}, - ), - ( - DocumentGenerateSummaryApi, - DocumentGenerateSummaryApi.post, - lambda service: service.generate_summary, - {"req_data": GenerateSummaryPayload(document_list=[str(DOC)])}, - ), - ], -) -def test_success_mutations_return_success_response(app, documents, resource, verb, operation, kwargs): - with app.test_request_context("/?document_id=" + str(DOC)): - result = unwrap(verb)(resource(), request_context=CONTEXT, dataset_id=DS, **kwargs) - assert result == ({"result": "success"}, 200) - assert operation(documents).call_count == 1 - - -@pytest.mark.parametrize("resource", [DatasetInitApi, DatasetDocumentListApi]) -def test_creation_serializes_materialized_dataset_and_documents(app, documents, resource): - operation = documents.initialize_dataset if resource is DatasetInitApi else documents.create_documents - operation.return_value = { - "dataset": {"id": str(DS), "name": "Dataset", "created_at": datetime.datetime(2024, 1, 1, tzinfo=datetime.UTC)}, - "documents": [ - { - "id": str(DOC), - "name": "Document", - "data_source_detail_dict": {}, - "data_source_info_dict": {}, - "doc_metadata_details": [], - } - ], - "batch": "batch-1", - } - with app.test_request_context("/", json={"indexing_technique": "economy"}): - response = unwrap(resource.post)( - resource(), request_context=CONTEXT, **({} if resource is DatasetInitApi else {"dataset_id": DS}) - ) - assert response["dataset"]["created_at"] == 1704067200 - assert response["documents"][0]["data_source_info"] == {} - assert response["documents"][0]["doc_metadata"] == [] - assert response["batch"] == "batch-1" - - -@pytest.mark.parametrize("resource", [DocumentIndexingStatusApi, DocumentBatchIndexingStatusApi]) -def test_indexing_status_keeps_timestamps_counts_and_error_fields(app, documents, resource): - value = { - "id": str(DOC), - "indexing_status": "paused", - "completed_segments": 2, - "total_segments": 3, - **dict.fromkeys( - ( - "processing_started_at", - "parsing_completed_at", - "cleaning_completed_at", - "splitting_completed_at", - "completed_at", - "paused_at", - "error", - "stopped_at", - ) - ), - } - value["processing_started_at"] = datetime.datetime(2024, 1, 1, tzinfo=datetime.UTC) - batch = resource is DocumentBatchIndexingStatusApi - operation = documents.get_batch_indexing_status if batch else documents.get_indexing_status - operation.return_value = {"data": [value]} if batch else value - with app.test_request_context("/"): - response = unwrap(resource.get)( - resource(), - request_context=CONTEXT, - dataset_id=DS, - **({"batch": "batch-1"} if batch else {"document_id": DOC}), - ) - row = response["data"][0] if batch else response - assert row["indexing_status"] == "paused" - assert row["processing_started_at"] == 1704067200 - assert row["completed_segments"] == 2 - assert row["total_segments"] == 3 - - -@pytest.mark.parametrize( - ("translate", "error", "http_error"), - [ - (controller._raise_document_error, DatasetNotFoundError(), NotFound), - (controller._raise_document_error, DocumentNotFoundError("missing"), NotFound), - (controller._raise_document_error, DatasetAccessDeniedError(), Forbidden), - (controller._raise_document_indexing_error, DocumentIndexingStateError("indexing"), DocumentIndexingError), - ( - controller._raise_document_indexing_error, - DocumentProviderError("provider", kind="uninitialized"), - ProviderNotInitializeError, - ), - (controller._raise_document_action_error, DocumentInvalidActionError("invalid"), InvalidActionError), - (controller._raise_document_status_error, DocumentInvalidActionError("invalid"), InvalidActionError), - (controller._raise_document_status_error, DocumentIndexingStateError("indexing"), DocumentIndexingError), - (controller._raise_document_archive_error, DocumentArchivedError(), ArchivedDocumentImmutableError), - (controller._raise_document_archive_error, DocumentIndexingStateError("indexing"), DocumentIndexingError), - # Every narrow translator falls back to the shared scope mapping. - (controller._raise_document_archive_error, DatasetNotFoundError(), NotFound), - ], -) -def test_document_errors_are_mapped_at_transport_boundary(translate, error, http_error): - with pytest.raises(http_error): - translate(error) - - -def test_scope_translator_does_not_map_other_operations_errors(): - error = DocumentInvalidActionError("invalid") - with pytest.raises(DocumentInvalidActionError) as raised: - controller._raise_document_error(error) - assert raised.value is error - - -def test_unexpected_error_is_not_hidden(): - error = RuntimeError("unexpected") - with pytest.raises(RuntimeError) as raised: - controller._raise_document_error(error) - assert raised.value is error - - -class TestIndexingEstimateExceptionMapping: +class TestDatasetsDocumentCases: @staticmethod - def _registry(method_name: str, error: Exception) -> SimpleNamespace: - estimates = MagicMock() - method = estimates.estimate_document if method_name == "estimate_document" else estimates.estimate_batch - method.side_effect = error - return SimpleNamespace(knowledge=SimpleNamespace(indexing_estimates=estimates)) + @pytest.fixture + def documents(monkeypatch): + documents = create_autospec(DatasetDocumentApplicationService, instance=True, spec_set=True) + monkeypatch.setattr( + controller, "application_services", lambda: SimpleNamespace(knowledge=SimpleNamespace(documents=documents)) + ) + monkeypatch.setattr(controller, "check_knowledge_rate_limit", lambda: None) + return documents + @staticmethod + def test_process_rule_authorizes_document_dataset_before_read(app, documents): + with ( + app.test_request_context(f"/?document_id={DOC}"), + patch.object(controller, "enforce_rbac_checks", side_effect=Forbidden) as enforce_checks, + pytest.raises(Forbidden), + ): + unwrap(GetProcessRuleApi.get)(GetProcessRuleApi(), _DATASETS_DOCUMENT_CONTEXT) + + enforce_checks.assert_called_once() + kwargs = enforce_checks.call_args.kwargs + assert kwargs["tenant_id"] == _DATASETS_DOCUMENT_CONTEXT.active_workspace_id + assert kwargs["account_id"] == _DATASETS_DOCUMENT_CONTEXT.account_id + assert kwargs["path_args"] == {"document_id": str(DOC)} + [check] = kwargs["checks"] + assert check.scene is RBACPermission.DATASET_READONLY + assert isinstance(check.locator, DatasetByDocument) + documents.get_process_rule.assert_not_called() + + @staticmethod + def test_default_process_rule_does_not_require_dataset_permission(app, documents): + documents.get_process_rule.return_value = { + "mode": "automatic", + "rules": {}, + "limits": {"indexing_max_segmentation_tokens_length": 4000}, + } + with app.test_request_context("/"), patch.object(controller, "enforce_rbac_checks") as enforce_checks: + unwrap(GetProcessRuleApi.get)(GetProcessRuleApi(), _DATASETS_DOCUMENT_CONTEXT) + enforce_checks.assert_not_called() + documents.get_process_rule.assert_called_once_with(_DATASETS_DOCUMENT_CONTEXT, document_id=None) + + @staticmethod @pytest.mark.parametrize( - ("error", "expected_http_error"), + ("resource", "verb", "operation", "kwargs"), [ - (DatasetNotFoundError(), NotFoundError), - (EstimateDocumentNotFoundError(), NotFoundError), - (EstimateSourceNotFoundError("source-1"), NotFoundError), - (IndexingEstimateCredentialUnavailableError(), NotFoundError), - (DatasetAccessDeniedError(), DatasetAccessDeniedRequestError), - (EstimateDocumentAlreadyFinishedError(), DocumentAlreadyFinishedError), - (UnsupportedEstimateSourceError("unsupported"), InvalidArgumentError), - (IndexingEstimateProviderUnavailableError(), ProviderNotInitializeError), - (IndexingEstimateExecutionError(), IndexingEstimateError), + (GetProcessRuleApi, GetProcessRuleApi.get, lambda service: service.get_process_rule, {}), + ( + DatasetDocumentListApi, + DatasetDocumentListApi.get, + lambda service: service.list_documents, + {"dataset_id": DS}, + ), + ( + DatasetDocumentListApi, + DatasetDocumentListApi.post, + lambda service: service.create_documents, + {"dataset_id": DS}, + ), + ( + DatasetDocumentListApi, + DatasetDocumentListApi.delete, + lambda service: service.delete_documents, + {"dataset_id": DS}, + ), + (DatasetInitApi, DatasetInitApi.post, lambda service: service.initialize_dataset, {}), + ( + DocumentBatchIndexingStatusApi, + DocumentBatchIndexingStatusApi.get, + lambda service: service.get_batch_indexing_status, + {"dataset_id": DS, "batch": "batch-1"}, + ), + ( + DocumentIndexingStatusApi, + DocumentIndexingStatusApi.get, + lambda service: service.get_indexing_status, + {"dataset_id": DS, "document_id": DOC}, + ), + ( + DocumentApi, + DocumentApi.get, + lambda service: service.get_document, + {"dataset_id": DS, "document_id": DOC}, + ), + ( + DocumentApi, + DocumentApi.delete, + lambda service: service.delete_document, + {"dataset_id": DS, "document_id": DOC}, + ), + ( + DocumentDownloadApi, + DocumentDownloadApi.get, + lambda service: service.get_download_url, + {"dataset_id": DS, "document_id": DOC}, + ), + ( + DocumentBatchDownloadZipApi, + DocumentBatchDownloadZipApi.post, + lambda service: service.build_download_zip, + {"dataset_id": DS}, + ), + ( + DocumentProcessingApi, + DocumentProcessingApi.patch, + lambda service: service.update_processing, + {"dataset_id": DS, "document_id": DOC, "action": "pause"}, + ), + ( + DocumentMetadataApi, + DocumentMetadataApi.put, + lambda service: service.update_metadata, + { + "dataset_id": DS, + "document_id": DOC, + "req_data": DocumentMetadataUpdatePayload(doc_type="book", doc_metadata={"title": "Book"}), + }, + ), + ( + DocumentStatusApi, + DocumentStatusApi.patch, + lambda service: service.change_status, + {"dataset_id": DS, "action": "enable"}, + ), + ( + DocumentPauseApi, + DocumentPauseApi.patch, + lambda service: service.pause_document, + {"dataset_id": DS, "document_id": DOC}, + ), + ( + DocumentRecoverApi, + DocumentRecoverApi.patch, + lambda service: service.recover_document, + {"dataset_id": DS, "document_id": DOC}, + ), + ( + DocumentRetryApi, + DocumentRetryApi.post, + lambda service: service.retry_documents, + {"dataset_id": DS, "req_data": DocumentRetryPayload(document_ids=[str(DOC)])}, + ), + ( + DocumentRenameApi, + DocumentRenameApi.post, + lambda service: service.rename_document, + {"dataset_id": DS, "document_id": DOC, "req_data": DocumentRenamePayload(name="New name")}, + ), + ( + WebsiteDocumentSyncApi, + WebsiteDocumentSyncApi.get, + lambda service: service.sync_website, + {"dataset_id": DS, "document_id": DOC}, + ), + ( + DocumentPipelineExecutionLogApi, + DocumentPipelineExecutionLogApi.get, + lambda service: service.get_execution_log, + {"dataset_id": DS, "document_id": DOC}, + ), + ( + DocumentGenerateSummaryApi, + DocumentGenerateSummaryApi.post, + lambda service: service.generate_summary, + {"dataset_id": DS, "req_data": GenerateSummaryPayload(document_list=[str(DOC)])}, + ), + ( + DocumentSummaryStatusApi, + DocumentSummaryStatusApi.get, + lambda service: service.get_summary_status, + {"dataset_id": DS, "document_id": DOC}, + ), ], ) - def test_document_estimate_maps_application_errors( - self, - error: Exception, - expected_http_error: type[Exception], - ) -> None: - api = DocumentIndexingEstimateApi() - method = unwrap(api.get) - context = RequestContext("request-1", None, "account-1", "workspace-1") - registry = self._registry("estimate_document", error) - - with patch( - "controllers.console.datasets.datasets_document.application_services", - return_value=registry, + def test_endpoints_pass_request_context_and_map_application_errors( + app, documents, resource, verb, operation, kwargs + ): + operation(documents).side_effect = DocumentNotFoundError("Document not found.") + with app.test_request_context( + "/?document_id=" + str(DOC), json={"indexing_technique": "economy", "document_ids": [str(DOC)]} ): - with pytest.raises(expected_http_error): - method(api, context, UUID(int=1), UUID(int=2)) + with pytest.raises(NotFound): + unwrap(verb)(resource(), request_context=_DATASETS_DOCUMENT_CONTEXT, **kwargs) + called = operation(documents) + assert called.call_count == 1 + assert called.call_args.args == (_DATASETS_DOCUMENT_CONTEXT,) + if "dataset_id" in kwargs: + assert called.call_args.kwargs["dataset_id"] == str(DS) + if "document_id" in kwargs: + assert called.call_args.kwargs["document_id"] == str(DOC) + @staticmethod @pytest.mark.parametrize( - ("error", "expected_http_error"), + ("fetch", "expected"), [("true", True), ("YES", True), ("false", False), ("invalid", False)] + ) + def test_list_preserves_filters_pagination_and_fetch_parsing(app, documents, fetch, expected): + documents.list_documents.return_value = {"data": [], "has_more": False, "page": 2, "limit": 3, "total": 3} + with app.test_request_context(f"/?page=2&limit=3&keyword=needle&status=available&sort=hit_count&fetch={fetch}"): + response = unwrap(DatasetDocumentListApi.get)(DatasetDocumentListApi(), _DATASETS_DOCUMENT_CONTEXT, DS) + documents.list_documents.assert_called_once_with( + _DATASETS_DOCUMENT_CONTEXT, + dataset_id=str(DS), + query=DocumentListFilter( + page=2, limit=3, search="needle", status="available", sort="hit_count", fetch=expected + ), + ) + assert response == {"data": [], "has_more": False, "page": 2, "limit": 3, "total": 3} + + @staticmethod + @pytest.mark.parametrize("metadata", ["only", "without", "all"]) + def test_document_detail_metadata_projection_and_raw_source(app, documents, metadata): + documents.get_document.return_value = { + "id": str(DOC), + "doc_type": "others", + "doc_metadata": [], + "name": "Name", + "data_source_info": {"file_path": "/path/to/local/file"}, + "data_source_detail_dict": {}, + } + with app.test_request_context(f"/?metadata={metadata}"): + response, status = unwrap(DocumentApi.get)(DocumentApi(), _DATASETS_DOCUMENT_CONTEXT, DS, DOC) + assert status == 200 + documents.get_document.assert_called_once_with( + _DATASETS_DOCUMENT_CONTEXT, dataset_id=str(DS), document_id=str(DOC), metadata_only=metadata == "only" + ) + if metadata == "only": + assert response == {"id": str(DOC), "doc_type": "others", "doc_metadata": []} + else: + assert response["data_source_info"] == {"file_path": "/path/to/local/file"} + assert ("doc_metadata" in response) is (metadata == "all") + + @staticmethod + def test_invalid_metadata_fails_before_application_call(app, documents): + with app.test_request_context("/?metadata=invalid"), pytest.raises(InvalidMetadataError): + unwrap(DocumentApi.get)(DocumentApi(), _DATASETS_DOCUMENT_CONTEXT, DS, DOC) + documents.get_document.assert_not_called() + + @staticmethod + @pytest.mark.parametrize( + ("resource", "verb", "kwargs"), [ - (DatasetNotFoundError(), NotFoundError), - (EstimateDocumentNotFoundError(), NotFoundError), - (EstimateSourceNotFoundError("source-1"), NotFoundError), - (IndexingEstimateCredentialUnavailableError(), NotFoundError), - (DatasetAccessDeniedError(), DatasetAccessDeniedRequestError), - (EstimateDocumentAlreadyFinishedError(), DocumentAlreadyFinishedError), - (UnsupportedEstimateSourceError("unsupported"), InvalidArgumentError), - (IndexingEstimateProviderUnavailableError(), ProviderNotInitializeError), - (IndexingEstimateExecutionError(), IndexingEstimateError), + (DocumentApi, DocumentApi.delete, {"document_id": DOC}), + (DatasetDocumentListApi, DatasetDocumentListApi.delete, {}), + (DocumentPauseApi, DocumentPauseApi.patch, {"document_id": DOC}), + (DocumentRecoverApi, DocumentRecoverApi.patch, {"document_id": DOC}), + (DocumentRetryApi, DocumentRetryApi.post, {"req_data": DocumentRetryPayload(document_ids=[str(DOC)])}), ], ) - def test_batch_estimate_maps_application_errors( - self, - error: Exception, - expected_http_error: type[Exception], - ) -> None: - api = DocumentBatchIndexingEstimateApi() - method = unwrap(api.get) - context = RequestContext("request-1", None, "account-1", "workspace-1") - registry = self._registry("estimate_batch", error) + @pytest.mark.usefixtures("documents") + def test_no_content_mutations_return_empty_204(app, resource, verb, kwargs): + with app.test_request_context("/?document_id=" + str(DOC)): + result = unwrap(verb)(resource(), request_context=_DATASETS_DOCUMENT_CONTEXT, dataset_id=DS, **kwargs) + assert result == ("", 204) - with patch( - "controllers.console.datasets.datasets_document.application_services", - return_value=registry, + @staticmethod + @pytest.mark.parametrize( + ("resource", "verb", "operation", "kwargs"), + [ + ( + DocumentProcessingApi, + DocumentProcessingApi.patch, + lambda service: service.update_processing, + {"document_id": DOC, "action": "resume"}, + ), + (DocumentStatusApi, DocumentStatusApi.patch, lambda service: service.change_status, {"action": "disable"}), + ( + WebsiteDocumentSyncApi, + WebsiteDocumentSyncApi.get, + lambda service: service.sync_website, + {"document_id": DOC}, + ), + ( + DocumentGenerateSummaryApi, + DocumentGenerateSummaryApi.post, + lambda service: service.generate_summary, + {"req_data": GenerateSummaryPayload(document_list=[str(DOC)])}, + ), + ], + ) + def test_success_mutations_return_success_response(app, documents, resource, verb, operation, kwargs): + with app.test_request_context("/?document_id=" + str(DOC)): + result = unwrap(verb)(resource(), request_context=_DATASETS_DOCUMENT_CONTEXT, dataset_id=DS, **kwargs) + assert result == ({"result": "success"}, 200) + assert operation(documents).call_count == 1 + + @staticmethod + @pytest.mark.parametrize("resource", [DatasetInitApi, DatasetDocumentListApi]) + def test_creation_serializes_materialized_dataset_and_documents(app, documents, resource): + operation = documents.initialize_dataset if resource is DatasetInitApi else documents.create_documents + operation.return_value = { + "dataset": { + "id": str(DS), + "name": "Dataset", + "created_at": _datasets_document_datetime.datetime(2024, 1, 1, tzinfo=_datasets_document_datetime.UTC), + }, + "documents": [ + { + "id": str(DOC), + "name": "Document", + "data_source_detail_dict": {}, + "data_source_info_dict": {}, + "doc_metadata_details": [], + } + ], + "batch": "batch-1", + } + with app.test_request_context("/", json={"indexing_technique": "economy"}): + response = unwrap(resource.post)( + resource(), + request_context=_DATASETS_DOCUMENT_CONTEXT, + **({} if resource is DatasetInitApi else {"dataset_id": DS}), + ) + assert response["dataset"]["created_at"] == 1704067200 + assert response["documents"][0]["data_source_info"] == {} + assert response["documents"][0]["doc_metadata"] == [] + assert response["batch"] == "batch-1" + + @staticmethod + @pytest.mark.parametrize("resource", [DocumentIndexingStatusApi, DocumentBatchIndexingStatusApi]) + def test_indexing_status_keeps_timestamps_counts_and_error_fields(app, documents, resource): + value = { + "id": str(DOC), + "indexing_status": "paused", + "completed_segments": 2, + "total_segments": 3, + **dict.fromkeys( + ( + "processing_started_at", + "parsing_completed_at", + "cleaning_completed_at", + "splitting_completed_at", + "completed_at", + "paused_at", + "error", + "stopped_at", + ) + ), + } + value["processing_started_at"] = _datasets_document_datetime.datetime( + 2024, 1, 1, tzinfo=_datasets_document_datetime.UTC + ) + batch = resource is DocumentBatchIndexingStatusApi + operation = documents.get_batch_indexing_status if batch else documents.get_indexing_status + operation.return_value = {"data": [value]} if batch else value + with app.test_request_context("/"): + response = unwrap(resource.get)( + resource(), + request_context=_DATASETS_DOCUMENT_CONTEXT, + dataset_id=DS, + **({"batch": "batch-1"} if batch else {"document_id": DOC}), + ) + row = response["data"][0] if batch else response + assert row["indexing_status"] == "paused" + assert row["processing_started_at"] == 1704067200 + assert row["completed_segments"] == 2 + assert row["total_segments"] == 3 + + @staticmethod + @pytest.mark.parametrize( + ("translate", "error", "http_error"), + [ + (controller._raise_document_error, DatasetNotFoundError(), NotFound), + (controller._raise_document_error, DocumentNotFoundError("missing"), NotFound), + (controller._raise_document_error, DatasetAccessDeniedError(), Forbidden), + (controller._raise_document_indexing_error, DocumentIndexingStateError("indexing"), DocumentIndexingError), + ( + controller._raise_document_indexing_error, + DocumentProviderError("provider", kind="uninitialized"), + ProviderNotInitializeError, + ), + (controller._raise_document_action_error, DocumentInvalidActionError("invalid"), InvalidActionError), + (controller._raise_document_status_error, DocumentInvalidActionError("invalid"), InvalidActionError), + (controller._raise_document_status_error, DocumentIndexingStateError("indexing"), DocumentIndexingError), + (controller._raise_document_archive_error, DocumentArchivedError(), ArchivedDocumentImmutableError), + (controller._raise_document_archive_error, DocumentIndexingStateError("indexing"), DocumentIndexingError), + # Every narrow translator falls back to the shared scope mapping. + (controller._raise_document_archive_error, DatasetNotFoundError(), NotFound), + ], + ) + def test_document_errors_are_mapped_at_transport_boundary(translate, error, http_error): + with pytest.raises(http_error): + translate(error) + + @staticmethod + def test_scope_translator_does_not_map_other_operations_errors(): + error = DocumentInvalidActionError("invalid") + with pytest.raises(DocumentInvalidActionError) as raised: + controller._raise_document_error(error) + assert raised.value is error + + @staticmethod + def test_unexpected_error_is_not_hidden(): + error = RuntimeError("unexpected") + with pytest.raises(RuntimeError) as raised: + controller._raise_document_error(error) + assert raised.value is error + + class TestIndexingEstimateExceptionMapping: + @staticmethod + def _registry(method_name: str, error: Exception) -> SimpleNamespace: + estimates = MagicMock() + method = estimates.estimate_document if method_name == "estimate_document" else estimates.estimate_batch + method.side_effect = error + return SimpleNamespace(knowledge=SimpleNamespace(indexing_estimates=estimates)) + + @pytest.mark.parametrize( + ("error", "expected_http_error"), + [ + (DatasetNotFoundError(), NotFoundError), + (EstimateDocumentNotFoundError(), NotFoundError), + (EstimateSourceNotFoundError("source-1"), NotFoundError), + (IndexingEstimateCredentialUnavailableError(), NotFoundError), + (DatasetAccessDeniedError(), DatasetAccessDeniedRequestError), + (EstimateDocumentAlreadyFinishedError(), DocumentAlreadyFinishedError), + (UnsupportedEstimateSourceError("unsupported"), InvalidArgumentError), + (IndexingEstimateProviderUnavailableError(), ProviderNotInitializeError), + (IndexingEstimateExecutionError(), IndexingEstimateError), + ], + ) + def test_document_estimate_maps_application_errors( + self, + error: Exception, + expected_http_error: type[Exception], + ) -> None: + api = DocumentIndexingEstimateApi() + method = unwrap(api.get) + context = RequestContext("request-1", None, "account-1", "workspace-1") + registry = self._registry("estimate_document", error) + + with patch( + "controllers.console.datasets.datasets_document.application_services", + return_value=registry, + ): + with pytest.raises(expected_http_error): + method(api, context, UUID(int=1), UUID(int=2)) + + @pytest.mark.parametrize( + ("error", "expected_http_error"), + [ + (DatasetNotFoundError(), NotFoundError), + (EstimateDocumentNotFoundError(), NotFoundError), + (EstimateSourceNotFoundError("source-1"), NotFoundError), + (IndexingEstimateCredentialUnavailableError(), NotFoundError), + (DatasetAccessDeniedError(), DatasetAccessDeniedRequestError), + (EstimateDocumentAlreadyFinishedError(), DocumentAlreadyFinishedError), + (UnsupportedEstimateSourceError("unsupported"), InvalidArgumentError), + (IndexingEstimateProviderUnavailableError(), ProviderNotInitializeError), + (IndexingEstimateExecutionError(), IndexingEstimateError), + ], + ) + def test_batch_estimate_maps_application_errors( + self, + error: Exception, + expected_http_error: type[Exception], + ) -> None: + api = DocumentBatchIndexingEstimateApi() + method = unwrap(api.get) + context = RequestContext("request-1", None, "account-1", "workspace-1") + registry = self._registry("estimate_batch", error) + + with patch( + "controllers.console.datasets.datasets_document.application_services", + return_value=registry, + ): + with pytest.raises(expected_http_error): + method(api, context, UUID(int=1), "unknown-batch") + + @staticmethod + @pytest.mark.parametrize( + "method", + [ + DatasetDocumentListApi.get, + DocumentBatchIndexingStatusApi.get, + DocumentIndexingStatusApi.get, + DocumentApi.get, + DocumentPipelineExecutionLogApi.get, + DocumentSummaryStatusApi.get, + ], + ) + def test_document_read_routes_require_dataset_readonly_permission(method) -> None: + [check] = rbac_checks(method) + + assert check.scene is RBACPermission.DATASET_READONLY + assert isinstance(check.locator, DatasetId) + + @staticmethod + @pytest.mark.parametrize( + "method", + [DocumentIndexingEstimateApi.get, DocumentBatchIndexingEstimateApi.get], + ) + def test_document_indexing_estimates_require_dataset_use_permission(method) -> None: + [check] = rbac_checks(method) + + assert check.scene is RBACPermission.DATASET_USE + assert isinstance(check.locator, DatasetId) + + @staticmethod + @pytest.mark.parametrize( + ("method", "expected_scene"), + [ + (DatasetDocumentListApi.post, RBACPermission.DATASET_USE), + (DatasetDocumentListApi.delete, RBACPermission.DATASET_DELETE_FILE), + (DocumentApi.delete, RBACPermission.DATASET_DELETE_FILE), + (DocumentBatchDownloadZipApi.post, RBACPermission.DATASET_DOCUMENT_DOWNLOAD), + ], + ) + def test_document_mutation_routes_require_operation_specific_permissions(method, expected_scene) -> None: + [check] = rbac_checks(method) + + assert check.scene is expected_scene + assert isinstance(check.locator, DatasetId) + + +"""Document downloads through the controller, application service, and SQLite adapter.""" + + +_DATASETS_DOCUMENT_DOWNLOAD_CONTEXT = RequestContext("request-1", None, "actor-1", "tenant-1") +DOC_ID = str(UUID(int=1)) +FILE_ID = str(UUID(int=2)) + + +def _document(**values): + return Document( + **{ + "id": DOC_ID, + "tenant_id": "tenant-1", + "dataset_id": "dataset-1", + "position": 1, + "data_source_type": "upload_file", + "data_source_info": json.dumps({"upload_file_id": FILE_ID}), + "batch": "batch-1", + "name": "document.txt", + "created_from": "web", + "created_by": "actor-1", + **values, + } + ) + + +class TestDatasetsDocumentDownloadCases: + @staticmethod + @pytest.fixture(autouse=True) + def downloads(monkeypatch, sqlite_session_factory): + with sqlite_session_factory.begin() as session: + session.add(Dataset(id="dataset-1", tenant_id="tenant-1", name="Dataset", created_by="actor-1")) + session.add(_document()) + upload = UploadFile( + tenant_id="tenant-1", + storage_type="local", + key="key-1", + name="document.txt", + size=5, + extension="txt", + mime_type="text/plain", + created_by_role="account", + created_by="actor-1", + created_at=_datasets_document_download_datetime(2024, 1, 1), + used=False, + ) + upload.id = FILE_ID + session.add(upload) + access = create_autospec(DatasetAccess, instance=True, spec_set=True) + access.require_accessible.return_value = AccessibleDataset("dataset-1", "tenant-1") + documents = DatasetDocumentApplicationService( + dataset_access=access, + operations=SQLAlchemyDocumentOperations(session_factory=sqlite_session_factory), + metadata_schema={}, + ) + monkeypatch.setattr( + controller, "application_services", lambda: SimpleNamespace(knowledge=SimpleNamespace(documents=documents)) + ) + return documents + + @staticmethod + def test_document_download_returns_signed_url(app, monkeypatch): + from services.knowledge import dataset_service + + signed = [] + + def sign(**kwargs): + signed.append(kwargs) + return "https://files.example/download" + + monkeypatch.setattr(dataset_service.file_helpers, "get_signed_file_url", sign) + with app.test_request_context("/"): + result = unwrap(controller.DocumentDownloadApi.get)( + controller.DocumentDownloadApi(), _DATASETS_DOCUMENT_DOWNLOAD_CONTEXT, "dataset-1", DOC_ID + ) + assert result == {"url": "https://files.example/download"} + assert signed == [{"upload_file_id": FILE_ID, "as_attachment": True}] + + @staticmethod + @pytest.mark.parametrize( + "values", + [ + {"data_source_type": "website_crawl"}, + {"data_source_info": "{}"}, + {"data_source_info": json.dumps({"upload_file_id": "missing"})}, + {"tenant_id": "foreign"}, + {"dataset_id": "foreign"}, + ], + ) + def test_invalid_or_foreign_document_is_not_downloadable(app, sqlite_session_factory, values): + with sqlite_session_factory.begin() as session: + document = session.get(Document, DOC_ID) + for key, value in values.items(): + setattr(document, key, value) + with app.test_request_context("/"), pytest.raises(NotFound): + unwrap(controller.DocumentDownloadApi.get)( + controller.DocumentDownloadApi(), _DATASETS_DOCUMENT_DOWNLOAD_CONTEXT, "dataset-1", DOC_ID + ) + + @staticmethod + def test_cross_tenant_upload_file_is_not_signed(app, sqlite_session_factory): + with sqlite_session_factory.begin() as session: + session.get(UploadFile, FILE_ID).tenant_id = "foreign" + with app.test_request_context("/"), pytest.raises(NotFound): + unwrap(controller.DocumentDownloadApi.get)( + controller.DocumentDownloadApi(), _DATASETS_DOCUMENT_DOWNLOAD_CONTEXT, "dataset-1", DOC_ID + ) + + @staticmethod + def test_zip_is_readable_and_cleanup_waits_for_response_close(app, monkeypatch, sqlite_engine): + from services import file_service + + paths = [] + checked_out = set() + + def checkout(_connection, record, _proxy): + checked_out.add(record) + + def checkin(_connection, record): + checked_out.discard(record) + + def load(key, *, stream): + assert not checked_out, "storage I/O must run after the read transaction closes" + assert key == "key-1" + assert stream is True + return iter([b"hello"]) + + def send(path, **kwargs): + paths.append(Path(path)) + return send_file(path, **kwargs) + + monkeypatch.setattr(file_service.storage, "load", load) + monkeypatch.setattr(controller, "send_file", send) + event.listen(sqlite_engine, "checkout", checkout) + event.listen(sqlite_engine, "checkin", checkin) + try: + with app.test_request_context("/", json={"document_ids": [DOC_ID]}): + response = unwrap(controller.DocumentBatchDownloadZipApi.post)( + controller.DocumentBatchDownloadZipApi(), _DATASETS_DOCUMENT_DOWNLOAD_CONTEXT, "dataset-1" + ) + assert paths[0].exists() + response.direct_passthrough = False + with ZipFile(BytesIO(response.get_data())) as archive: + assert archive.read("document.txt") == b"hello" + response.close() + assert not paths[0].exists() + finally: + event.remove(sqlite_engine, "checkout", checkout) + event.remove(sqlite_engine, "checkin", checkin) + + @staticmethod + def test_zip_cleanup_when_response_creation_fails(app, monkeypatch): + from services import file_service + + paths = [] + monkeypatch.setattr(file_service.storage, "load", lambda *_args, **_kwargs: iter([b"hello"])) + + def fail_send(path, **_kwargs): + paths.append(Path(path)) + raise RuntimeError("response failed") + + monkeypatch.setattr(controller, "send_file", fail_send) + with ( + app.test_request_context("/", json={"document_ids": [DOC_ID]}), + pytest.raises(RuntimeError, match="response failed"), ): - with pytest.raises(expected_http_error): - method(api, context, UUID(int=1), "unknown-batch") + unwrap(controller.DocumentBatchDownloadZipApi.post)( + controller.DocumentBatchDownloadZipApi(), _DATASETS_DOCUMENT_DOWNLOAD_CONTEXT, "dataset-1" + ) + assert paths + assert not paths[0].exists() - -@pytest.mark.parametrize( - "method", - [ - DatasetDocumentListApi.get, - DocumentBatchIndexingStatusApi.get, - DocumentIndexingStatusApi.get, - DocumentApi.get, - DocumentPipelineExecutionLogApi.get, - DocumentSummaryStatusApi.get, - ], -) -def test_document_read_routes_require_dataset_readonly_permission(method) -> None: - [check] = rbac_checks(method) - - assert check.scene is RBACPermission.DATASET_READONLY - assert isinstance(check.locator, DatasetId) - - -@pytest.mark.parametrize( - "method", - [DocumentIndexingEstimateApi.get, DocumentBatchIndexingEstimateApi.get], -) -def test_document_indexing_estimates_require_dataset_use_permission(method) -> None: - [check] = rbac_checks(method) - - assert check.scene is RBACPermission.DATASET_USE - assert isinstance(check.locator, DatasetId) - - -@pytest.mark.parametrize( - ("method", "expected_scene"), - [ - (DatasetDocumentListApi.post, RBACPermission.DATASET_USE), - (DatasetDocumentListApi.delete, RBACPermission.DATASET_DELETE_FILE), - (DocumentApi.delete, RBACPermission.DATASET_DELETE_FILE), - (DocumentBatchDownloadZipApi.post, RBACPermission.DATASET_DOCUMENT_DOWNLOAD), - ], -) -def test_document_mutation_routes_require_operation_specific_permissions(method, expected_scene) -> None: - [check] = rbac_checks(method) - - assert check.scene is expected_scene - assert isinstance(check.locator, DatasetId) + @staticmethod + def test_zip_rejects_non_upload_document(app, sqlite_session_factory): + with sqlite_session_factory.begin() as session: + session.get(Document, DOC_ID).data_source_type = "website_crawl" + with app.test_request_context("/", json={"document_ids": [DOC_ID]}), pytest.raises(NotFound): + unwrap(controller.DocumentBatchDownloadZipApi.post)( + controller.DocumentBatchDownloadZipApi(), _DATASETS_DOCUMENT_DOWNLOAD_CONTEXT, "dataset-1" + ) diff --git a/api/tests/unit_tests/controllers/console/datasets/test_datasets_document_download.py b/api/tests/unit_tests/controllers/console/datasets/test_datasets_document_download.py deleted file mode 100644 index 92b64d4875f..00000000000 --- a/api/tests/unit_tests/controllers/console/datasets/test_datasets_document_download.py +++ /dev/null @@ -1,196 +0,0 @@ -"""Document downloads through the controller, application service, and SQLite adapter.""" - -import json -from datetime import datetime -from inspect import unwrap -from io import BytesIO -from pathlib import Path -from types import SimpleNamespace -from unittest.mock import create_autospec -from uuid import UUID -from zipfile import ZipFile - -import pytest -from flask import send_file -from sqlalchemy import event -from werkzeug.exceptions import NotFound - -from controllers.console.datasets import datasets_document as controller -from machinery.context import RequestContext -from models.dataset import Dataset, Document -from models.model import UploadFile -from services.knowledge.dataset_access import AccessibleDataset, DatasetAccess -from services.knowledge.documents.adapters import SQLAlchemyDocumentOperations -from services.knowledge.documents.application import DatasetDocumentApplicationService - -CONTEXT = RequestContext("request-1", None, "actor-1", "tenant-1") -DOC_ID = str(UUID(int=1)) -FILE_ID = str(UUID(int=2)) - - -def _document(**values): - return Document( - **{ - "id": DOC_ID, - "tenant_id": "tenant-1", - "dataset_id": "dataset-1", - "position": 1, - "data_source_type": "upload_file", - "data_source_info": json.dumps({"upload_file_id": FILE_ID}), - "batch": "batch-1", - "name": "document.txt", - "created_from": "web", - "created_by": "actor-1", - **values, - } - ) - - -@pytest.fixture(autouse=True) -def downloads(monkeypatch, sqlite_session_factory): - with sqlite_session_factory.begin() as session: - session.add(Dataset(id="dataset-1", tenant_id="tenant-1", name="Dataset", created_by="actor-1")) - session.add(_document()) - upload = UploadFile( - tenant_id="tenant-1", - storage_type="local", - key="key-1", - name="document.txt", - size=5, - extension="txt", - mime_type="text/plain", - created_by_role="account", - created_by="actor-1", - created_at=datetime(2024, 1, 1), - used=False, - ) - upload.id = FILE_ID - session.add(upload) - access = create_autospec(DatasetAccess, instance=True, spec_set=True) - access.require_accessible.return_value = AccessibleDataset("dataset-1", "tenant-1") - documents = DatasetDocumentApplicationService( - dataset_access=access, - operations=SQLAlchemyDocumentOperations(session_factory=sqlite_session_factory), - metadata_schema={}, - ) - monkeypatch.setattr( - controller, "application_services", lambda: SimpleNamespace(knowledge=SimpleNamespace(documents=documents)) - ) - return documents - - -def test_document_download_returns_signed_url(app, monkeypatch): - from services.knowledge import dataset_service - - signed = [] - - def sign(**kwargs): - signed.append(kwargs) - return "https://files.example/download" - - monkeypatch.setattr(dataset_service.file_helpers, "get_signed_file_url", sign) - with app.test_request_context("/"): - result = unwrap(controller.DocumentDownloadApi.get)( - controller.DocumentDownloadApi(), CONTEXT, "dataset-1", DOC_ID - ) - assert result == {"url": "https://files.example/download"} - assert signed == [{"upload_file_id": FILE_ID, "as_attachment": True}] - - -@pytest.mark.parametrize( - "values", - [ - {"data_source_type": "website_crawl"}, - {"data_source_info": "{}"}, - {"data_source_info": json.dumps({"upload_file_id": "missing"})}, - {"tenant_id": "foreign"}, - {"dataset_id": "foreign"}, - ], -) -def test_invalid_or_foreign_document_is_not_downloadable(app, sqlite_session_factory, values): - with sqlite_session_factory.begin() as session: - document = session.get(Document, DOC_ID) - for key, value in values.items(): - setattr(document, key, value) - with app.test_request_context("/"), pytest.raises(NotFound): - unwrap(controller.DocumentDownloadApi.get)(controller.DocumentDownloadApi(), CONTEXT, "dataset-1", DOC_ID) - - -def test_cross_tenant_upload_file_is_not_signed(app, sqlite_session_factory): - with sqlite_session_factory.begin() as session: - session.get(UploadFile, FILE_ID).tenant_id = "foreign" - with app.test_request_context("/"), pytest.raises(NotFound): - unwrap(controller.DocumentDownloadApi.get)(controller.DocumentDownloadApi(), CONTEXT, "dataset-1", DOC_ID) - - -def test_zip_is_readable_and_cleanup_waits_for_response_close(app, monkeypatch, sqlite_engine): - from services import file_service - - paths = [] - checked_out = set() - - def checkout(_connection, record, _proxy): - checked_out.add(record) - - def checkin(_connection, record): - checked_out.discard(record) - - def load(key, *, stream): - assert not checked_out, "storage I/O must run after the read transaction closes" - assert key == "key-1" - assert stream is True - return iter([b"hello"]) - - def send(path, **kwargs): - paths.append(Path(path)) - return send_file(path, **kwargs) - - monkeypatch.setattr(file_service.storage, "load", load) - monkeypatch.setattr(controller, "send_file", send) - event.listen(sqlite_engine, "checkout", checkout) - event.listen(sqlite_engine, "checkin", checkin) - try: - with app.test_request_context("/", json={"document_ids": [DOC_ID]}): - response = unwrap(controller.DocumentBatchDownloadZipApi.post)( - controller.DocumentBatchDownloadZipApi(), CONTEXT, "dataset-1" - ) - assert paths[0].exists() - response.direct_passthrough = False - with ZipFile(BytesIO(response.get_data())) as archive: - assert archive.read("document.txt") == b"hello" - response.close() - assert not paths[0].exists() - finally: - event.remove(sqlite_engine, "checkout", checkout) - event.remove(sqlite_engine, "checkin", checkin) - - -def test_zip_cleanup_when_response_creation_fails(app, monkeypatch): - from services import file_service - - paths = [] - monkeypatch.setattr(file_service.storage, "load", lambda *_args, **_kwargs: iter([b"hello"])) - - def fail_send(path, **_kwargs): - paths.append(Path(path)) - raise RuntimeError("response failed") - - monkeypatch.setattr(controller, "send_file", fail_send) - with ( - app.test_request_context("/", json={"document_ids": [DOC_ID]}), - pytest.raises(RuntimeError, match="response failed"), - ): - unwrap(controller.DocumentBatchDownloadZipApi.post)( - controller.DocumentBatchDownloadZipApi(), CONTEXT, "dataset-1" - ) - assert paths - assert not paths[0].exists() - - -def test_zip_rejects_non_upload_document(app, sqlite_session_factory): - with sqlite_session_factory.begin() as session: - session.get(Document, DOC_ID).data_source_type = "website_crawl" - with app.test_request_context("/", json={"document_ids": [DOC_ID]}), pytest.raises(NotFound): - unwrap(controller.DocumentBatchDownloadZipApi.post)( - controller.DocumentBatchDownloadZipApi(), CONTEXT, "dataset-1" - ) diff --git a/api/tests/unit_tests/controllers/console/datasets/test_metadata.py b/api/tests/unit_tests/controllers/console/datasets/test_metadata.py index 55200cc8805..670135ae8cf 100644 --- a/api/tests/unit_tests/controllers/console/datasets/test_metadata.py +++ b/api/tests/unit_tests/controllers/console/datasets/test_metadata.py @@ -21,7 +21,7 @@ from services.errors.metadata import MetadataResourceNotFoundError from services.knowledge.dataset_access import DatasetAccessDeniedError, DatasetNotFoundError from services.knowledge.metadata.application import MetadataService from services.knowledge.resource_scope import DatasetRef -from tests.unit_tests.controllers.rbac_introspection import rbac_checks +from tests.helpers.unit.controllers.rbac_introspection import rbac_checks @pytest.fixture diff --git a/api/tests/unit_tests/controllers/console/explore/test_completion.py b/api/tests/unit_tests/controllers/console/explore/test_completion.py index 1dd8ce3c7f9..e842be4ac0c 100644 --- a/api/tests/unit_tests/controllers/console/explore/test_completion.py +++ b/api/tests/unit_tests/controllers/console/explore/test_completion.py @@ -1,14 +1,17 @@ import json -from collections.abc import Generator +from collections.abc import Generator, Mapping +from dataclasses import dataclass, field +from datetime import datetime from uuid import uuid4 import pytest from sqlalchemy.orm import Session, sessionmaker import controllers.console.explore.completion as completion_module +import services.installed_app_generation_service as completion_service_module from core.errors.error import ModelCurrentlyNotSupportError, ProviderTokenNotInitError, QuotaExceededError from graphon.model_runtime.errors.invoke import InvokeError -from models import AppMode, Conversation, Tenant +from models import App, AppMode, Conversation, InstalledApp, Tenant from repositories.installed_app_repository import SQLAlchemyInstalledAppRepository from services.account_errors import AccountNotFoundError from services.app_generate_service import AppGenerateService @@ -21,363 +24,830 @@ from tests.unit_tests.controllers.console.explore.test_installed_app_admission i _assert_json_response, _Harness, _set_app_mode, - harness, -) -from tests.unit_tests.controllers.console.explore.test_installed_app_completion import ( - _USED_AT, - _InstalledAppServices, - _last_used_at, - _Runtime, - _RuntimeCall, - _Services, - runtime, ) -__all__ = ["harness", "runtime"] +_installed_app_completion_all__ = ["harness"] + +_INSTALLED_APP_COMPLETION_USED_AT = datetime(2026, 9, 6, 12, 30, 0) +_BLOCKING_RESPONSE: dict[str, object] = {"event": "message", "answer": "Hello", "metadata": {}, "usage": None} + + +@dataclass(frozen=True) +class _InstalledAppCompletionRuntimecall: + app_id: str + account_id: str + args: Mapping[str, object] + streaming: bool + + +@dataclass +class _InstalledAppCompletionRuntime: + session_factory: sessionmaker[Session] + installed_app_id: str + response: GenerationResponse = field(default_factory=lambda: dict(_BLOCKING_RESPONSE)) + error: Exception | None = None + calls: list[_InstalledAppCompletionRuntimecall] = field(default_factory=list) + + def generate( + self, *, app_id: str, account_id: str, args: Mapping[str, object], streaming: bool + ) -> GenerationResponse: + # The runtime sees a committed usage update through an independent + # connection, including when generation immediately fails. + with self.session_factory() as session: + installation = session.get(InstalledApp, self.installed_app_id) + assert installation is not None + assert installation.last_used_at == _INSTALLED_APP_COMPLETION_USED_AT + self.calls.append(_InstalledAppCompletionRuntimecall(app_id, account_id, dict(args), streaming)) + if self.error is not None: + raise self.error + return self.response + + def generate_more_like_this( + self, *, app_id: str, account_id: str, message_id: str, streaming: bool + ) -> GenerationResponse: + pytest.fail(f"Unexpected more-like-this call: {app_id=}, {account_id=}, {message_id=}, {streaming=}") + + +@dataclass(frozen=True) +class _InstalledAppCompletionInstalledappservices: + generation: InstalledAppGenerationService + + +@dataclass(frozen=True) +class _InstalledAppCompletionServices: + installed_apps: _InstalledAppCompletionInstalledappservices + + +def _installed_app_completion_url(harness: _Harness) -> str: + return f"/installed-apps/{harness.installed_app.id}/completion-messages" + + +def _installed_app_completion_last_used_at( + harness: _Harness, session_factory: sessionmaker[Session] +) -> datetime | None: + with session_factory() as session: + installation = session.get(InstalledApp, harness.installed_app.id) + assert installation is not None + return installation.last_used_at + + +class TestInstalledAppCompletionCases: + from tests.unit_tests.controllers.console.explore.test_installed_app_admission import harness + + harness = staticmethod(harness) + + @staticmethod + @pytest.fixture + def runtime( + harness: _Harness, + monkeypatch: pytest.MonkeyPatch, + sqlite_session_factory: sessionmaker[Session], + ) -> _InstalledAppCompletionRuntime: + runtime = _InstalledAppCompletionRuntime(sqlite_session_factory, harness.installed_app.id) + service = InstalledAppGenerationService( + usage=SQLAlchemyInstalledAppRepository(session_factory=sqlite_session_factory), + runtime=runtime, + ) + services = _InstalledAppCompletionServices( + installed_apps=_InstalledAppCompletionInstalledappservices(generation=service) + ) + monkeypatch.setattr(completion_module, "application_services", lambda: services) + monkeypatch.setattr(completion_service_module, "naive_utc_now", lambda: _INSTALLED_APP_COMPLETION_USED_AT) + harness.api.add_resource( + completion_module.CompletionApi, + "/installed-apps//completion-messages", + ) + return runtime + + @staticmethod + @pytest.mark.parametrize( + ("payload", "expected_args"), + [ + ( + {"inputs": {"zero": 0, "enabled": False, "empty": [], "nullable": None}}, + { + "inputs": {"zero": 0, "enabled": False, "empty": [], "nullable": None}, + "query": "", + "retriever_from": "explore_app", + "auto_generate_name": False, + }, + ), + ( + {"inputs": {}, "query": "Hi", "files": None, "response_mode": None}, + {"inputs": {}, "query": "Hi", "retriever_from": "explore_app", "auto_generate_name": False}, + ), + ( + {"inputs": {}, "files": [], "response_mode": "blocking", "auto_generate_name": True}, + { + "inputs": {}, + "query": "", + "files": [], + "response_mode": "blocking", + "retriever_from": "explore_app", + "auto_generate_name": False, + }, + ), + ( + { + "inputs": {}, + "files": [ + {"type": "image", "transfer_method": "remote_url", "url": "https://example.com/image.png"} + ], + "retriever_from": "custom-source", + }, + { + "inputs": {}, + "query": "", + "files": [ + {"type": "image", "transfer_method": "remote_url", "url": "https://example.com/image.png"} + ], + "retriever_from": "custom-source", + "auto_generate_name": False, + }, + ), + ], + ) + def test_blocking_completion_preserves_payload_defaults_and_response( + harness: _Harness, + runtime: _InstalledAppCompletionRuntime, + sqlite_session_factory: sessionmaker[Session], + payload: dict[str, object], + expected_args: dict[str, object], + ) -> None: + response = harness.app.test_client().post(_installed_app_completion_url(harness), json=payload) + + assert response.status_code == 200 + assert response.get_json() == _BLOCKING_RESPONSE + assert dict(response.headers) == { + "Content-Type": "application/json; charset=utf-8", + "Content-Length": str(len(response.data)), + } + assert runtime.calls == [ + _InstalledAppCompletionRuntimecall(harness.target_app.id, harness.account.id, expected_args, False) + ] + assert ( + _installed_app_completion_last_used_at(harness, sqlite_session_factory) == _INSTALLED_APP_COMPLETION_USED_AT + ) + assert harness.installed_app.tenant_id != harness.target_app.tenant_id + + @staticmethod + @pytest.mark.parametrize("consume_all", [False, True]) + def test_streaming_completion_preserves_bytes_headers_and_closes_runtime_iterator( + harness: _Harness, + runtime: _InstalledAppCompletionRuntime, + consume_all: bool, + ) -> None: + closed: list[bool] = [] + + def chunks() -> Generator[str]: + try: + yield 'data: {"answer":"你好"}\n\n' + yield "data: [DONE]\n\n" + finally: + closed.append(True) + + runtime.response = chunks() + response = harness.app.test_client().post( + _installed_app_completion_url(harness), json={"inputs": {}, "response_mode": "streaming"}, buffered=False + ) + + assert response.status_code == 200 + assert dict(response.headers) == {"Content-Type": "text/event-stream; charset=utf-8"} + if consume_all: + assert response.data == 'data: {"answer":"你好"}\n\ndata: [DONE]\n\n'.encode() + else: + assert next(iter(response.response)) == 'data: {"answer":"你好"}\n\n'.encode() + assert closed == [] + response.close() + assert closed == [True] + assert runtime.calls == [ + _InstalledAppCompletionRuntimecall( + harness.target_app.id, + harness.account.id, + { + "inputs": {}, + "query": "", + "response_mode": "streaming", + "retriever_from": "explore_app", + "auto_generate_name": False, + }, + True, + ) + ] + + @staticmethod + @pytest.mark.parametrize("response_mode", ["blocking", "streaming"]) + @pytest.mark.parametrize( + ("failure", "status", "code", "message"), + [ + (ConversationNotExistsError(), 404, "not_found", "Conversation Not Exists."), + # Nothing raises ConversationCompletedError today, so the contract does not list its code. + pytest.param( + ConversationCompletedError(), + 400, + "conversation_completed", + "The conversation has ended. Please start a new conversation.", + marks=pytest.mark.allow_error_contract_violations, + ), + ( + AppModelConfigBrokenError(), + 400, + "app_unavailable", + "App unavailable, please check your app configurations.", + ), + ( + ProviderTokenNotInitError("Missing provider credentials"), + 400, + "provider_not_initialize", + "Missing provider credentials", + ), + ( + QuotaExceededError(), + 400, + "provider_quota_exceeded", + "Your quota for Dify Hosted Model Provider has been exhausted. " + "Please go to Settings -> Model Provider to complete your own provider credentials.", + ), + ( + ModelCurrentlyNotSupportError(), + 400, + "model_currently_not_support", + "Dify Hosted OpenAI trial currently not support the GPT-4 model.", + ), + (InvokeError("Provider rejected input"), 400, "completion_request_error", "Provider rejected input"), + (ValueError("Invalid runtime arguments"), 400, "invalid_param", "Invalid runtime arguments"), + (AccountNotFoundError(), 401, "unauthorized", "Account no longer exists."), + ( + RuntimeError("Unexpected runtime failure"), + 500, + "internal_server_error", + "The server encountered an internal error and was unable to complete your request. " + "Either the server is overloaded or there is an error in the application.", + ), + ], + ) + def test_completion_preserves_error_contract_and_committed_usage_on_runtime_failure( + harness: _Harness, + runtime: _InstalledAppCompletionRuntime, + sqlite_session_factory: sessionmaker[Session], + failure: Exception, + status: int, + code: str, + message: str, + response_mode: str, + ) -> None: + runtime.error = failure + + response = harness.app.test_client().post( + _installed_app_completion_url(harness), json={"inputs": {}, "response_mode": response_mode} + ) + + _assert_json_response(response, status=status, body={"code": code, "message": message, "status": status}) + assert len(runtime.calls) == 1 + assert runtime.calls[0].streaming == (response_mode == "streaming") + assert ( + _installed_app_completion_last_used_at(harness, sqlite_session_factory) == _INSTALLED_APP_COMPLETION_USED_AT + ) + if status == 401: + assert response.headers["WWW-Authenticate"] == 'Bearer realm="api"' + + @staticmethod + @pytest.mark.parametrize( + ("payload", "error"), + [ + ({}, {"type": "missing", "loc": ["inputs"], "msg": "Field required"}), + ({"inputs": None}, {"type": "dict_type", "loc": ["inputs"], "msg": "Input should be a valid dictionary"}), + ( + {"inputs": {}, "response_mode": "invalid"}, + {"type": "literal_error", "loc": ["response_mode"], "msg": "Input should be 'blocking' or 'streaming'"}, + ), + ], + ) + def test_invalid_payload_does_not_record_usage_or_start_runtime( + harness: _Harness, + runtime: _InstalledAppCompletionRuntime, + sqlite_session_factory: sessionmaker[Session], + payload: dict[str, object], + error: dict[str, object], + ) -> None: + response = harness.app.test_client().post(_installed_app_completion_url(harness), json=payload) + + _assert_json_response( + response, + status=422, + body={"code": "unprocessable_entity", "message": json.dumps([error]), "status": 422}, + ) + assert runtime.calls == [] + assert _installed_app_completion_last_used_at(harness, sqlite_session_factory) is None + + @staticmethod + def test_wrong_mode_does_not_record_usage_or_start_runtime( + harness: _Harness, + runtime: _InstalledAppCompletionRuntime, + sqlite_session_factory: sessionmaker[Session], + ) -> None: + with sqlite_session_factory.begin() as session: + app = session.get(App, harness.target_app.id) + assert app is not None + app.mode = AppMode.CHAT + + response = harness.app.test_client().post(_installed_app_completion_url(harness), json={"inputs": {}}) + + _assert_json_response( + response, status=400, body={"code": "not_completion_app", "message": "Not Completion App", "status": 400} + ) + assert runtime.calls == [] + assert _installed_app_completion_last_used_at(harness, sqlite_session_factory) is None + + @staticmethod + def test_completion_uses_the_mode_captured_at_admission( + harness: _Harness, + runtime: _InstalledAppCompletionRuntime, + sqlite_session_factory: sessionmaker[Session], + ) -> None: + harness.state.permission_action = lambda: _set_app_mode(harness, sqlite_session_factory, AppMode.CHAT) + + response = harness.app.test_client().post(_installed_app_completion_url(harness), json={"inputs": {}}) + + assert response.status_code == 200 + assert response.get_json() == _BLOCKING_RESPONSE + assert len(runtime.calls) == 1 + assert ( + _installed_app_completion_last_used_at(harness, sqlite_session_factory) == _INSTALLED_APP_COMPLETION_USED_AT + ) + + @staticmethod + @pytest.mark.parametrize("rejection", ["permission", "missing", "tenant"]) + def test_completion_requires_installed_app_admission_before_payload_validation( + harness: _Harness, + runtime: _InstalledAppCompletionRuntime, + sqlite_session_factory: sessionmaker[Session], + rejection: str, + ) -> None: + if rejection == "permission": + harness.state.allowed = False + elif rejection == "tenant": + harness.account._current_tenant = Tenant(name="Other workspace") + else: + with sqlite_session_factory.begin() as session: + installation = session.get(InstalledApp, harness.installed_app.id) + assert installation is not None + session.delete(installation) + + response = harness.app.test_client().post(_installed_app_completion_url(harness), json={}) + + if rejection == "permission": + _assert_json_response( + response, status=403, body={"code": "access_denied", "message": "App access denied.", "status": 403} + ) + else: + _assert_json_response( + response, + status=404, + body={ + "code": "installed_app_not_found", + "message": "The app was not found in this workspace.", + "status": 404, + "details": {"request_id": "request-1"}, + }, + ) + assert runtime.calls == [] + if rejection != "missing": + assert _installed_app_completion_last_used_at(harness, sqlite_session_factory) is None + + @staticmethod + @pytest.mark.parametrize("deleted_resource", ["app", "installation"]) + @pytest.mark.usefixtures("runtime") + def test_resource_removed_after_admission_is_revalidated_before_generation( + harness: _Harness, + sqlite_session_factory: sessionmaker[Session], + monkeypatch: pytest.MonkeyPatch, + deleted_resource: str, + ) -> None: + def delete_resource() -> None: + with sqlite_session_factory.begin() as session: + if deleted_resource == "app": + resource = session.get(App, harness.target_app.id) + else: + resource = session.get(InstalledApp, harness.installed_app.id) + assert resource is not None + session.delete(resource) + + harness.state.permission_action = delete_resource + services = _InstalledAppCompletionServices( + installed_apps=_InstalledAppCompletionInstalledappservices( + generation=InstalledAppGenerationService( + usage=SQLAlchemyInstalledAppRepository(session_factory=sqlite_session_factory), + runtime=AppGenerateServiceRuntime(session_factory=sqlite_session_factory), + ) + ) + ) + monkeypatch.setattr(completion_module, "application_services", lambda: services) + + def generate(**_kwargs: object) -> GenerationResponse: + pytest.fail("Removed app or installation must be rejected before generation") + + monkeypatch.setattr(AppGenerateService, "generate", generate) + + response = harness.app.test_client().post(_installed_app_completion_url(harness), json={"inputs": {}}) + + if deleted_resource == "app": + _assert_json_response( + response, + status=400, + body={ + "code": "app_unavailable", + "message": "App unavailable, please check your app configurations.", + "status": 400, + }, + ) + assert ( + _installed_app_completion_last_used_at(harness, sqlite_session_factory) + == _INSTALLED_APP_COMPLETION_USED_AT + ) + else: + _assert_json_response( + response, + status=404, + body={"code": "not_found", "message": "Installed app not found", "status": 404}, + ) + + +_COMPLETION_USED_AT = _INSTALLED_APP_COMPLETION_USED_AT +_CompletionInstalledAppServices = _InstalledAppCompletionInstalledappservices +_completion_last_used_at = _installed_app_completion_last_used_at +_CompletionRuntime = _InstalledAppCompletionRuntime +_CompletionRuntimeCall = _InstalledAppCompletionRuntimecall +_CompletionServices = _InstalledAppCompletionServices + +_completion_all__ = ["harness", "runtime"] _CONVERSATION_ID = "aaaaaaaa-aaaa-4aaa-8aaa-aaaaaaaaaaaa" _PARENT_MESSAGE_ID = "bbbbbbbb-bbbb-4bbb-8bbb-bbbbbbbbbbbb" -@pytest.fixture -def chat_runtime( - harness: _Harness, - runtime: _Runtime, - sqlite_session_factory: sessionmaker[Session], -) -> _Runtime: - _set_app_mode(harness, sqlite_session_factory, AppMode.CHAT) - harness.api.add_resource(completion_module.ChatApi, "/installed-apps//chat-messages") - return runtime - - -def _url(harness: _Harness) -> str: +def _completion_url(harness: _Harness) -> str: return f"/installed-apps/{harness.installed_app.id}/chat-messages" -@pytest.mark.parametrize( - ("payload", "expected_args"), - [ - ( - {"inputs": {"zero": 0, "enabled": False, "empty": [], "nullable": None}, "query": ""}, - {"inputs": {"zero": 0, "enabled": False, "empty": [], "nullable": None}, "query": ""}, - ), - ( - {"inputs": {}, "query": "Hi", "files": None, "conversation_id": None, "parent_message_id": None}, - {"inputs": {}, "query": "Hi"}, - ), - ( - { - "inputs": {}, - "query": "Hi", - "files": [], - "conversation_id": "", - "parent_message_id": "", - "response_mode": "blocking", - "auto_generate_name": True, - }, - {"inputs": {}, "query": "Hi", "files": []}, - ), - ( - { - "inputs": {}, - "query": "Hi", - "files": [{"type": "image", "transfer_method": "remote_url", "url": "https://example.com/image.png"}], - "conversation_id": _CONVERSATION_ID.upper(), - "parent_message_id": _PARENT_MESSAGE_ID.replace("-", ""), - "retriever_from": "custom-source", - }, - { - "inputs": {}, - "query": "Hi", - "files": [{"type": "image", "transfer_method": "remote_url", "url": "https://example.com/image.png"}], - "conversation_id": _CONVERSATION_ID, - "parent_message_id": _PARENT_MESSAGE_ID, - "retriever_from": "custom-source", - }, - ), - ], -) -def test_chat_preserves_payload_defaults_and_always_requests_streaming( - harness: _Harness, - chat_runtime: _Runtime, - payload: dict[str, object], - expected_args: dict[str, object], -) -> None: - def chunks() -> Generator[str]: - yield 'data: {"answer":"你好","metadata":{},"usage":null}\n\n' +class TestCompletionCases: + from tests.unit_tests.controllers.console.explore.test_installed_app_admission import harness - chat_runtime.response = chunks() - response = harness.app.test_client().post(_url(harness), json=payload) + harness = staticmethod(harness) - assert response.status_code == 200 - assert dict(response.headers) == {"Content-Type": "text/event-stream; charset=utf-8"} - assert response.data == 'data: {"answer":"你好","metadata":{},"usage":null}\n\n'.encode() - response.close() - expected_args = {"retriever_from": "explore_app", **expected_args, "auto_generate_name": False} - assert chat_runtime.calls == [_RuntimeCall(harness.target_app.id, harness.account.id, expected_args, True)] - assert harness.installed_app.tenant_id != harness.target_app.tenant_id + runtime = staticmethod(TestInstalledAppCompletionCases.runtime) + @staticmethod + @pytest.fixture + def chat_runtime( + harness: _Harness, + runtime: _CompletionRuntime, + sqlite_session_factory: sessionmaker[Session], + ) -> _CompletionRuntime: + _set_app_mode(harness, sqlite_session_factory, AppMode.CHAT) + harness.api.add_resource(completion_module.ChatApi, "/installed-apps//chat-messages") + return runtime -@pytest.mark.parametrize("mode", [AppMode.CHAT, AppMode.AGENT_CHAT, AppMode.ADVANCED_CHAT]) -@pytest.mark.parametrize("consume_all", [False, True]) -def test_chat_modes_preserve_stream_bytes_headers_and_close_on_completion_or_disconnect( - harness: _Harness, - chat_runtime: _Runtime, - sqlite_session_factory: sessionmaker[Session], - mode: AppMode, - consume_all: bool, -) -> None: - _set_app_mode(harness, sqlite_session_factory, mode) - closed: list[bool] = [] - - def chunks() -> Generator[str]: - try: - yield 'data: {"answer":"Hello"}\n\n' - yield "data: [DONE]\n\n" - finally: - closed.append(True) - - chat_runtime.response = chunks() - response = harness.app.test_client().post(_url(harness), json={"inputs": {}, "query": "Hi"}, buffered=False) - - assert response.status_code == 200 - assert dict(response.headers) == {"Content-Type": "text/event-stream; charset=utf-8"} - if consume_all: - assert response.data == b'data: {"answer":"Hello"}\n\ndata: [DONE]\n\n' - else: - assert next(iter(response.response)) == b'data: {"answer":"Hello"}\n\n' - assert closed == [] - response.close() - assert closed == [True] - assert len(chat_runtime.calls) == 1 - assert chat_runtime.calls[0].streaming is True - assert _last_used_at(harness, sqlite_session_factory) == _USED_AT - - -@pytest.mark.parametrize( - ("failure", "status", "code", "message"), - [ - (ConversationNotExistsError(), 404, "not_found", "Conversation Not Exists."), - # Nothing raises ConversationCompletedError today, so the contract does not list its code. - pytest.param( - ConversationCompletedError(), - 400, - "conversation_completed", - "The conversation has ended. Please start a new conversation.", - marks=pytest.mark.allow_error_contract_violations, - ), - (AppModelConfigBrokenError(), 400, "app_unavailable", "App unavailable, please check your app configurations."), - (ProviderTokenNotInitError("Missing credentials"), 400, "provider_not_initialize", "Missing credentials"), - ( - QuotaExceededError(), - 400, - "provider_quota_exceeded", - "Your quota for Dify Hosted Model Provider has been exhausted. " - "Please go to Settings -> Model Provider to complete your own provider credentials.", - ), - ( - ModelCurrentlyNotSupportError(), - 400, - "model_currently_not_support", - "Dify Hosted OpenAI trial currently not support the GPT-4 model.", - ), - (InvokeError("Provider rejected input"), 400, "completion_request_error", "Provider rejected input"), - ( - InvokeRateLimitError("Concurrent request limit exceeded"), - 429, - "rate_limit_error", - "Concurrent request limit exceeded", - ), - (ValueError("Invalid runtime arguments"), 400, "invalid_param", "Invalid runtime arguments"), - (AccountNotFoundError(), 401, "unauthorized", "Account no longer exists."), - ( - RuntimeError("Unexpected runtime failure"), - 500, - "internal_server_error", - "The server encountered an internal error and was unable to complete your request. " - "Either the server is overloaded or there is an error in the application.", - ), - ], -) -def test_chat_preserves_precise_errors_and_committed_usage_when_generation_fails( - harness: _Harness, - chat_runtime: _Runtime, - sqlite_session_factory: sessionmaker[Session], - failure: Exception, - status: int, - code: str, - message: str, -) -> None: - chat_runtime.error = failure - - response = harness.app.test_client().post(_url(harness), json={"inputs": {}, "query": "Hi"}) - - _assert_json_response(response, status=status, body={"code": code, "message": message, "status": status}) - assert len(chat_runtime.calls) == 1 - assert chat_runtime.calls[0].streaming is True - assert _last_used_at(harness, sqlite_session_factory) == _USED_AT - if status == 401: - assert response.headers["WWW-Authenticate"] == 'Bearer realm="api"' - - -@pytest.mark.parametrize( - ("payload", "error"), - [ - ({"inputs": {}}, {"type": "missing", "loc": ["query"], "msg": "Field required"}), - ( - {"inputs": {}, "query": None}, - {"type": "string_type", "loc": ["query"], "msg": "Input should be a valid string"}, - ), - ( - {"inputs": None, "query": "Hi"}, - {"type": "dict_type", "loc": ["inputs"], "msg": "Input should be a valid dictionary"}, - ), - ( - {"inputs": {}, "query": "Hi", "conversation_id": "invalid"}, - {"type": "value_error", "loc": ["conversation_id"], "msg": "Value error, must be a valid UUID"}, - ), - ( - {"inputs": {}, "query": "Hi", "parent_message_id": "invalid"}, - {"type": "value_error", "loc": ["parent_message_id"], "msg": "Value error, must be a valid UUID"}, - ), - ], -) -def test_chat_validates_payload_before_mode_and_never_records_usage_for_invalid_input( - harness: _Harness, - chat_runtime: _Runtime, - sqlite_session_factory: sessionmaker[Session], - payload: dict[str, object], - error: dict[str, object], -) -> None: - _set_app_mode(harness, sqlite_session_factory, AppMode.COMPLETION) - - response = harness.app.test_client().post(_url(harness), json=payload) - - _assert_json_response( - response, - status=422, - body={"code": "unprocessable_entity", "message": json.dumps([error]), "status": 422}, + @staticmethod + @pytest.mark.parametrize( + ("payload", "expected_args"), + [ + ( + {"inputs": {"zero": 0, "enabled": False, "empty": [], "nullable": None}, "query": ""}, + {"inputs": {"zero": 0, "enabled": False, "empty": [], "nullable": None}, "query": ""}, + ), + ( + {"inputs": {}, "query": "Hi", "files": None, "conversation_id": None, "parent_message_id": None}, + {"inputs": {}, "query": "Hi"}, + ), + ( + { + "inputs": {}, + "query": "Hi", + "files": [], + "conversation_id": "", + "parent_message_id": "", + "response_mode": "blocking", + "auto_generate_name": True, + }, + {"inputs": {}, "query": "Hi", "files": []}, + ), + ( + { + "inputs": {}, + "query": "Hi", + "files": [ + {"type": "image", "transfer_method": "remote_url", "url": "https://example.com/image.png"} + ], + "conversation_id": _CONVERSATION_ID.upper(), + "parent_message_id": _PARENT_MESSAGE_ID.replace("-", ""), + "retriever_from": "custom-source", + }, + { + "inputs": {}, + "query": "Hi", + "files": [ + {"type": "image", "transfer_method": "remote_url", "url": "https://example.com/image.png"} + ], + "conversation_id": _CONVERSATION_ID, + "parent_message_id": _PARENT_MESSAGE_ID, + "retriever_from": "custom-source", + }, + ), + ], ) - assert chat_runtime.calls == [] - assert _last_used_at(harness, sqlite_session_factory) is None + def test_chat_preserves_payload_defaults_and_always_requests_streaming( + harness: _Harness, + chat_runtime: _CompletionRuntime, + payload: dict[str, object], + expected_args: dict[str, object], + ) -> None: + def chunks() -> Generator[str]: + yield 'data: {"answer":"你好","metadata":{},"usage":null}\n\n' + chat_runtime.response = chunks() + response = harness.app.test_client().post(_completion_url(harness), json=payload) -@pytest.mark.parametrize( - "mode", [AppMode.COMPLETION, AppMode.WORKFLOW, AppMode.AGENT, AppMode.CHANNEL, AppMode.RAG_PIPELINE] -) -def test_chat_rejects_other_modes_before_recording_usage( - harness: _Harness, - chat_runtime: _Runtime, - sqlite_session_factory: sessionmaker[Session], - mode: AppMode, -) -> None: - _set_app_mode(harness, sqlite_session_factory, mode) + assert response.status_code == 200 + assert dict(response.headers) == {"Content-Type": "text/event-stream; charset=utf-8"} + assert response.data == 'data: {"answer":"你好","metadata":{},"usage":null}\n\n'.encode() + response.close() + expected_args = {"retriever_from": "explore_app", **expected_args, "auto_generate_name": False} + assert chat_runtime.calls == [ + _CompletionRuntimeCall(harness.target_app.id, harness.account.id, expected_args, True) + ] + assert harness.installed_app.tenant_id != harness.target_app.tenant_id - response = harness.app.test_client().post(_url(harness), json={"inputs": {}, "query": "Hi"}) + @staticmethod + @pytest.mark.parametrize("mode", [AppMode.CHAT, AppMode.AGENT_CHAT, AppMode.ADVANCED_CHAT]) + @pytest.mark.parametrize("consume_all", [False, True]) + def test_chat_modes_preserve_stream_bytes_headers_and_close_on_completion_or_disconnect( + harness: _Harness, + chat_runtime: _CompletionRuntime, + sqlite_session_factory: sessionmaker[Session], + mode: AppMode, + consume_all: bool, + ) -> None: + _set_app_mode(harness, sqlite_session_factory, mode) + closed: list[bool] = [] - _assert_json_response( - response, status=400, body={"code": "not_chat_app", "message": "App mode is invalid.", "status": 400} - ) - assert chat_runtime.calls == [] - assert _last_used_at(harness, sqlite_session_factory) is None + def chunks() -> Generator[str]: + try: + yield 'data: {"answer":"Hello"}\n\n' + yield "data: [DONE]\n\n" + finally: + closed.append(True) - -def test_chat_uses_the_mode_captured_at_admission( - harness: _Harness, - chat_runtime: _Runtime, - sqlite_session_factory: sessionmaker[Session], -) -> None: - harness.state.permission_action = lambda: _set_app_mode(harness, sqlite_session_factory, AppMode.COMPLETION) - - response = harness.app.test_client().post(_url(harness), json={"inputs": {}, "query": "Hi"}) - - assert response.status_code == 200 - assert len(chat_runtime.calls) == 1 - assert chat_runtime.calls[0].streaming is True - assert _last_used_at(harness, sqlite_session_factory) == _USED_AT - - -@pytest.mark.parametrize("rejection", ["permission", "tenant"]) -def test_chat_enforces_admission_before_payload_validation( - harness: _Harness, - chat_runtime: _Runtime, - sqlite_session_factory: sessionmaker[Session], - rejection: str, -) -> None: - if rejection == "permission": - harness.state.allowed = False - else: - harness.account._current_tenant = Tenant(name="Other workspace") - - response = harness.app.test_client().post(_url(harness), json={}) - - if rejection == "permission": - _assert_json_response( - response, status=403, body={"code": "access_denied", "message": "App access denied.", "status": 403} + chat_runtime.response = chunks() + response = harness.app.test_client().post( + _completion_url(harness), json={"inputs": {}, "query": "Hi"}, buffered=False ) - else: + + assert response.status_code == 200 + assert dict(response.headers) == {"Content-Type": "text/event-stream; charset=utf-8"} + if consume_all: + assert response.data == b'data: {"answer":"Hello"}\n\ndata: [DONE]\n\n' + else: + assert next(iter(response.response)) == b'data: {"answer":"Hello"}\n\n' + assert closed == [] + response.close() + assert closed == [True] + assert len(chat_runtime.calls) == 1 + assert chat_runtime.calls[0].streaming is True + assert _completion_last_used_at(harness, sqlite_session_factory) == _COMPLETION_USED_AT + + @staticmethod + @pytest.mark.parametrize( + ("failure", "status", "code", "message"), + [ + (ConversationNotExistsError(), 404, "not_found", "Conversation Not Exists."), + # Nothing raises ConversationCompletedError today, so the contract does not list its code. + pytest.param( + ConversationCompletedError(), + 400, + "conversation_completed", + "The conversation has ended. Please start a new conversation.", + marks=pytest.mark.allow_error_contract_violations, + ), + ( + AppModelConfigBrokenError(), + 400, + "app_unavailable", + "App unavailable, please check your app configurations.", + ), + (ProviderTokenNotInitError("Missing credentials"), 400, "provider_not_initialize", "Missing credentials"), + ( + QuotaExceededError(), + 400, + "provider_quota_exceeded", + "Your quota for Dify Hosted Model Provider has been exhausted. " + "Please go to Settings -> Model Provider to complete your own provider credentials.", + ), + ( + ModelCurrentlyNotSupportError(), + 400, + "model_currently_not_support", + "Dify Hosted OpenAI trial currently not support the GPT-4 model.", + ), + (InvokeError("Provider rejected input"), 400, "completion_request_error", "Provider rejected input"), + ( + InvokeRateLimitError("Concurrent request limit exceeded"), + 429, + "rate_limit_error", + "Concurrent request limit exceeded", + ), + (ValueError("Invalid runtime arguments"), 400, "invalid_param", "Invalid runtime arguments"), + (AccountNotFoundError(), 401, "unauthorized", "Account no longer exists."), + ( + RuntimeError("Unexpected runtime failure"), + 500, + "internal_server_error", + "The server encountered an internal error and was unable to complete your request. " + "Either the server is overloaded or there is an error in the application.", + ), + ], + ) + def test_chat_preserves_precise_errors_and_committed_usage_when_generation_fails( + harness: _Harness, + chat_runtime: _CompletionRuntime, + sqlite_session_factory: sessionmaker[Session], + failure: Exception, + status: int, + code: str, + message: str, + ) -> None: + chat_runtime.error = failure + + response = harness.app.test_client().post(_completion_url(harness), json={"inputs": {}, "query": "Hi"}) + + _assert_json_response(response, status=status, body={"code": code, "message": message, "status": status}) + assert len(chat_runtime.calls) == 1 + assert chat_runtime.calls[0].streaming is True + assert _completion_last_used_at(harness, sqlite_session_factory) == _COMPLETION_USED_AT + if status == 401: + assert response.headers["WWW-Authenticate"] == 'Bearer realm="api"' + + @staticmethod + @pytest.mark.parametrize( + ("payload", "error"), + [ + ({"inputs": {}}, {"type": "missing", "loc": ["query"], "msg": "Field required"}), + ( + {"inputs": {}, "query": None}, + {"type": "string_type", "loc": ["query"], "msg": "Input should be a valid string"}, + ), + ( + {"inputs": None, "query": "Hi"}, + {"type": "dict_type", "loc": ["inputs"], "msg": "Input should be a valid dictionary"}, + ), + ( + {"inputs": {}, "query": "Hi", "conversation_id": "invalid"}, + {"type": "value_error", "loc": ["conversation_id"], "msg": "Value error, must be a valid UUID"}, + ), + ( + {"inputs": {}, "query": "Hi", "parent_message_id": "invalid"}, + {"type": "value_error", "loc": ["parent_message_id"], "msg": "Value error, must be a valid UUID"}, + ), + ], + ) + def test_chat_validates_payload_before_mode_and_never_records_usage_for_invalid_input( + harness: _Harness, + chat_runtime: _CompletionRuntime, + sqlite_session_factory: sessionmaker[Session], + payload: dict[str, object], + error: dict[str, object], + ) -> None: + _set_app_mode(harness, sqlite_session_factory, AppMode.COMPLETION) + + response = harness.app.test_client().post(_completion_url(harness), json=payload) + _assert_json_response( response, - status=404, - body={ - "code": "installed_app_not_found", - "message": "The app was not found in this workspace.", - "status": 404, - "details": {"request_id": "request-1"}, - }, + status=422, + body={"code": "unprocessable_entity", "message": json.dumps([error]), "status": 422}, ) - assert chat_runtime.calls == [] - assert _last_used_at(harness, sqlite_session_factory) is None + assert chat_runtime.calls == [] + assert _completion_last_used_at(harness, sqlite_session_factory) is None + @staticmethod + @pytest.mark.parametrize( + "mode", [AppMode.COMPLETION, AppMode.WORKFLOW, AppMode.AGENT, AppMode.CHANNEL, AppMode.RAG_PIPELINE] + ) + def test_chat_rejects_other_modes_before_recording_usage( + harness: _Harness, + chat_runtime: _CompletionRuntime, + sqlite_session_factory: sessionmaker[Session], + mode: AppMode, + ) -> None: + _set_app_mode(harness, sqlite_session_factory, mode) -@pytest.mark.parametrize("rejection", ["missing", "other_app", "other_account", "api", "deleted"]) -@pytest.mark.usefixtures("chat_runtime") -def test_chat_conversation_preflight_returns_404_before_starting_generation_and_preserves_usage( - harness: _Harness, - sqlite_session_factory: sessionmaker[Session], - monkeypatch: pytest.MonkeyPatch, - rejection: str, -) -> None: - conversation_id = str(uuid4()) - with sqlite_session_factory.begin() as session: - session.add(harness.account) - if rejection != "missing": - session.add( - Conversation( - id=conversation_id, - app_id=str(uuid4()) if rejection == "other_app" else harness.target_app.id, - from_account_id=str(uuid4()) if rejection == "other_account" else harness.account.id, - from_source="api" if rejection == "api" else "console", - mode=AppMode.CHAT, - name="Existing conversation", - _inputs={}, - is_deleted=rejection == "deleted", + response = harness.app.test_client().post(_completion_url(harness), json={"inputs": {}, "query": "Hi"}) + + _assert_json_response( + response, status=400, body={"code": "not_chat_app", "message": "App mode is invalid.", "status": 400} + ) + assert chat_runtime.calls == [] + assert _completion_last_used_at(harness, sqlite_session_factory) is None + + @staticmethod + def test_chat_uses_the_mode_captured_at_admission( + harness: _Harness, + chat_runtime: _CompletionRuntime, + sqlite_session_factory: sessionmaker[Session], + ) -> None: + harness.state.permission_action = lambda: _set_app_mode(harness, sqlite_session_factory, AppMode.COMPLETION) + + response = harness.app.test_client().post(_completion_url(harness), json={"inputs": {}, "query": "Hi"}) + + assert response.status_code == 200 + assert len(chat_runtime.calls) == 1 + assert chat_runtime.calls[0].streaming is True + assert _completion_last_used_at(harness, sqlite_session_factory) == _COMPLETION_USED_AT + + @staticmethod + @pytest.mark.parametrize("rejection", ["permission", "tenant"]) + def test_chat_enforces_admission_before_payload_validation( + harness: _Harness, + chat_runtime: _CompletionRuntime, + sqlite_session_factory: sessionmaker[Session], + rejection: str, + ) -> None: + if rejection == "permission": + harness.state.allowed = False + else: + harness.account._current_tenant = Tenant(name="Other workspace") + + response = harness.app.test_client().post(_completion_url(harness), json={}) + + if rejection == "permission": + _assert_json_response( + response, status=403, body={"code": "access_denied", "message": "App access denied.", "status": 403} + ) + else: + _assert_json_response( + response, + status=404, + body={ + "code": "installed_app_not_found", + "message": "The app was not found in this workspace.", + "status": 404, + "details": {"request_id": "request-1"}, + }, + ) + assert chat_runtime.calls == [] + assert _completion_last_used_at(harness, sqlite_session_factory) is None + + @staticmethod + @pytest.mark.parametrize("rejection", ["missing", "other_app", "other_account", "api", "deleted"]) + @pytest.mark.usefixtures("chat_runtime") + def test_chat_conversation_preflight_returns_404_before_starting_generation_and_preserves_usage( + harness: _Harness, + sqlite_session_factory: sessionmaker[Session], + monkeypatch: pytest.MonkeyPatch, + rejection: str, + ) -> None: + conversation_id = str(uuid4()) + with sqlite_session_factory.begin() as session: + session.add(harness.account) + if rejection != "missing": + session.add( + Conversation( + id=conversation_id, + app_id=str(uuid4()) if rejection == "other_app" else harness.target_app.id, + from_account_id=str(uuid4()) if rejection == "other_account" else harness.account.id, + from_source="api" if rejection == "api" else "console", + mode=AppMode.CHAT, + name="Existing conversation", + _inputs={}, + is_deleted=rejection == "deleted", + ) + ) + + services = _CompletionServices( + installed_apps=_CompletionInstalledAppServices( + generation=InstalledAppGenerationService( + usage=SQLAlchemyInstalledAppRepository(session_factory=sqlite_session_factory), + runtime=AppGenerateServiceRuntime(session_factory=sqlite_session_factory), ) ) - - services = _Services( - installed_apps=_InstalledAppServices( - generation=InstalledAppGenerationService( - usage=SQLAlchemyInstalledAppRepository(session_factory=sqlite_session_factory), - runtime=AppGenerateServiceRuntime(session_factory=sqlite_session_factory), - ) ) - ) - monkeypatch.setattr(completion_module, "application_services", lambda: services) - generation_started: list[bool] = [] + monkeypatch.setattr(completion_module, "application_services", lambda: services) + generation_started: list[bool] = [] - def generate(**_kwargs: object) -> GenerationResponse: - generation_started.append(True) - pytest.fail("The generation runtime must not start for an inaccessible conversation") + def generate(**_kwargs: object) -> GenerationResponse: + generation_started.append(True) + pytest.fail("The generation runtime must not start for an inaccessible conversation") - monkeypatch.setattr(AppGenerateService, "generate", generate) + monkeypatch.setattr(AppGenerateService, "generate", generate) - response = harness.app.test_client().post( - _url(harness), json={"inputs": {}, "query": "Hi", "conversation_id": conversation_id}, buffered=False - ) + response = harness.app.test_client().post( + _completion_url(harness), + json={"inputs": {}, "query": "Hi", "conversation_id": conversation_id}, + buffered=False, + ) - _assert_json_response( - response, status=404, body={"code": "not_found", "message": "Conversation Not Exists.", "status": 404} - ) - assert generation_started == [] - assert _last_used_at(harness, sqlite_session_factory) == _USED_AT + _assert_json_response( + response, status=404, body={"code": "not_found", "message": "Conversation Not Exists.", "status": 404} + ) + assert generation_started == [] + assert _completion_last_used_at(harness, sqlite_session_factory) == _COMPLETION_USED_AT diff --git a/api/tests/unit_tests/controllers/console/explore/test_installed_app_admission.py b/api/tests/unit_tests/controllers/console/explore/test_installed_app_admission.py index 49a29544119..91f93fd9ef9 100644 --- a/api/tests/unit_tests/controllers/console/explore/test_installed_app_admission.py +++ b/api/tests/unit_tests/controllers/console/explore/test_installed_app_admission.py @@ -39,7 +39,7 @@ from services.app_task_service import AppTaskControlService from services.installed_app_access_service import InstalledAppAccessService, InstalledAppRef from services.saved_message_service import SavedMessageActor, SavedMessagePage, SavedMessageRecord, SavedMessageService from services.webapp_access_query_service import WebAppAccessUnavailableError -from tests.unit_tests.services.test_app_task_service import _StopRedis +from tests.unit_tests.services.test_app_task_service import _AppTaskServiceStopredis as _StopRedis @dataclass diff --git a/api/tests/unit_tests/controllers/console/explore/test_installed_app_completion.py b/api/tests/unit_tests/controllers/console/explore/test_installed_app_completion.py deleted file mode 100644 index e58836e1b1e..00000000000 --- a/api/tests/unit_tests/controllers/console/explore/test_installed_app_completion.py +++ /dev/null @@ -1,440 +0,0 @@ -import json -from collections.abc import Generator, Mapping -from dataclasses import dataclass, field -from datetime import datetime - -import pytest -from sqlalchemy.orm import Session, sessionmaker - -import controllers.console.explore.completion as completion_module -import services.installed_app_generation_service as completion_service_module -from core.errors.error import ModelCurrentlyNotSupportError, ProviderTokenNotInitError, QuotaExceededError -from graphon.model_runtime.errors.invoke import InvokeError -from models import App, AppMode, InstalledApp, Tenant -from repositories.installed_app_repository import SQLAlchemyInstalledAppRepository -from services.account_errors import AccountNotFoundError -from services.app_generate_service import AppGenerateService -from services.errors.app_model_config import AppModelConfigBrokenError -from services.errors.conversation import ConversationCompletedError, ConversationNotExistsError -from services.installed_app_generation_adapters import AppGenerateServiceRuntime -from services.installed_app_generation_service import GenerationResponse, InstalledAppGenerationService -from tests.unit_tests.controllers.console.explore.test_installed_app_admission import ( - _assert_json_response, - _Harness, - _set_app_mode, - harness, -) - -__all__ = ["harness"] - -_USED_AT = datetime(2026, 9, 6, 12, 30, 0) -_BLOCKING_RESPONSE: dict[str, object] = {"event": "message", "answer": "Hello", "metadata": {}, "usage": None} - - -@dataclass(frozen=True) -class _RuntimeCall: - app_id: str - account_id: str - args: Mapping[str, object] - streaming: bool - - -@dataclass -class _Runtime: - session_factory: sessionmaker[Session] - installed_app_id: str - response: GenerationResponse = field(default_factory=lambda: dict(_BLOCKING_RESPONSE)) - error: Exception | None = None - calls: list[_RuntimeCall] = field(default_factory=list) - - def generate( - self, *, app_id: str, account_id: str, args: Mapping[str, object], streaming: bool - ) -> GenerationResponse: - # The runtime sees a committed usage update through an independent - # connection, including when generation immediately fails. - with self.session_factory() as session: - installation = session.get(InstalledApp, self.installed_app_id) - assert installation is not None - assert installation.last_used_at == _USED_AT - self.calls.append(_RuntimeCall(app_id, account_id, dict(args), streaming)) - if self.error is not None: - raise self.error - return self.response - - def generate_more_like_this( - self, *, app_id: str, account_id: str, message_id: str, streaming: bool - ) -> GenerationResponse: - pytest.fail(f"Unexpected more-like-this call: {app_id=}, {account_id=}, {message_id=}, {streaming=}") - - -@dataclass(frozen=True) -class _InstalledAppServices: - generation: InstalledAppGenerationService - - -@dataclass(frozen=True) -class _Services: - installed_apps: _InstalledAppServices - - -@pytest.fixture -def runtime( - harness: _Harness, - monkeypatch: pytest.MonkeyPatch, - sqlite_session_factory: sessionmaker[Session], -) -> _Runtime: - runtime = _Runtime(sqlite_session_factory, harness.installed_app.id) - service = InstalledAppGenerationService( - usage=SQLAlchemyInstalledAppRepository(session_factory=sqlite_session_factory), - runtime=runtime, - ) - services = _Services(installed_apps=_InstalledAppServices(generation=service)) - monkeypatch.setattr(completion_module, "application_services", lambda: services) - monkeypatch.setattr(completion_service_module, "naive_utc_now", lambda: _USED_AT) - harness.api.add_resource( - completion_module.CompletionApi, - "/installed-apps//completion-messages", - ) - return runtime - - -def _url(harness: _Harness) -> str: - return f"/installed-apps/{harness.installed_app.id}/completion-messages" - - -def _last_used_at(harness: _Harness, session_factory: sessionmaker[Session]) -> datetime | None: - with session_factory() as session: - installation = session.get(InstalledApp, harness.installed_app.id) - assert installation is not None - return installation.last_used_at - - -@pytest.mark.parametrize( - ("payload", "expected_args"), - [ - ( - {"inputs": {"zero": 0, "enabled": False, "empty": [], "nullable": None}}, - { - "inputs": {"zero": 0, "enabled": False, "empty": [], "nullable": None}, - "query": "", - "retriever_from": "explore_app", - "auto_generate_name": False, - }, - ), - ( - {"inputs": {}, "query": "Hi", "files": None, "response_mode": None}, - {"inputs": {}, "query": "Hi", "retriever_from": "explore_app", "auto_generate_name": False}, - ), - ( - {"inputs": {}, "files": [], "response_mode": "blocking", "auto_generate_name": True}, - { - "inputs": {}, - "query": "", - "files": [], - "response_mode": "blocking", - "retriever_from": "explore_app", - "auto_generate_name": False, - }, - ), - ( - { - "inputs": {}, - "files": [{"type": "image", "transfer_method": "remote_url", "url": "https://example.com/image.png"}], - "retriever_from": "custom-source", - }, - { - "inputs": {}, - "query": "", - "files": [{"type": "image", "transfer_method": "remote_url", "url": "https://example.com/image.png"}], - "retriever_from": "custom-source", - "auto_generate_name": False, - }, - ), - ], -) -def test_blocking_completion_preserves_payload_defaults_and_response( - harness: _Harness, - runtime: _Runtime, - sqlite_session_factory: sessionmaker[Session], - payload: dict[str, object], - expected_args: dict[str, object], -) -> None: - response = harness.app.test_client().post(_url(harness), json=payload) - - assert response.status_code == 200 - assert response.get_json() == _BLOCKING_RESPONSE - assert dict(response.headers) == { - "Content-Type": "application/json; charset=utf-8", - "Content-Length": str(len(response.data)), - } - assert runtime.calls == [_RuntimeCall(harness.target_app.id, harness.account.id, expected_args, False)] - assert _last_used_at(harness, sqlite_session_factory) == _USED_AT - assert harness.installed_app.tenant_id != harness.target_app.tenant_id - - -@pytest.mark.parametrize("consume_all", [False, True]) -def test_streaming_completion_preserves_bytes_headers_and_closes_runtime_iterator( - harness: _Harness, - runtime: _Runtime, - consume_all: bool, -) -> None: - closed: list[bool] = [] - - def chunks() -> Generator[str]: - try: - yield 'data: {"answer":"你好"}\n\n' - yield "data: [DONE]\n\n" - finally: - closed.append(True) - - runtime.response = chunks() - response = harness.app.test_client().post( - _url(harness), json={"inputs": {}, "response_mode": "streaming"}, buffered=False - ) - - assert response.status_code == 200 - assert dict(response.headers) == {"Content-Type": "text/event-stream; charset=utf-8"} - if consume_all: - assert response.data == 'data: {"answer":"你好"}\n\ndata: [DONE]\n\n'.encode() - else: - assert next(iter(response.response)) == 'data: {"answer":"你好"}\n\n'.encode() - assert closed == [] - response.close() - assert closed == [True] - assert runtime.calls == [ - _RuntimeCall( - harness.target_app.id, - harness.account.id, - { - "inputs": {}, - "query": "", - "response_mode": "streaming", - "retriever_from": "explore_app", - "auto_generate_name": False, - }, - True, - ) - ] - - -@pytest.mark.parametrize("response_mode", ["blocking", "streaming"]) -@pytest.mark.parametrize( - ("failure", "status", "code", "message"), - [ - (ConversationNotExistsError(), 404, "not_found", "Conversation Not Exists."), - # Nothing raises ConversationCompletedError today, so the contract does not list its code. - pytest.param( - ConversationCompletedError(), - 400, - "conversation_completed", - "The conversation has ended. Please start a new conversation.", - marks=pytest.mark.allow_error_contract_violations, - ), - (AppModelConfigBrokenError(), 400, "app_unavailable", "App unavailable, please check your app configurations."), - ( - ProviderTokenNotInitError("Missing provider credentials"), - 400, - "provider_not_initialize", - "Missing provider credentials", - ), - ( - QuotaExceededError(), - 400, - "provider_quota_exceeded", - "Your quota for Dify Hosted Model Provider has been exhausted. " - "Please go to Settings -> Model Provider to complete your own provider credentials.", - ), - ( - ModelCurrentlyNotSupportError(), - 400, - "model_currently_not_support", - "Dify Hosted OpenAI trial currently not support the GPT-4 model.", - ), - (InvokeError("Provider rejected input"), 400, "completion_request_error", "Provider rejected input"), - (ValueError("Invalid runtime arguments"), 400, "invalid_param", "Invalid runtime arguments"), - (AccountNotFoundError(), 401, "unauthorized", "Account no longer exists."), - ( - RuntimeError("Unexpected runtime failure"), - 500, - "internal_server_error", - "The server encountered an internal error and was unable to complete your request. " - "Either the server is overloaded or there is an error in the application.", - ), - ], -) -def test_completion_preserves_error_contract_and_committed_usage_on_runtime_failure( - harness: _Harness, - runtime: _Runtime, - sqlite_session_factory: sessionmaker[Session], - failure: Exception, - status: int, - code: str, - message: str, - response_mode: str, -) -> None: - runtime.error = failure - - response = harness.app.test_client().post(_url(harness), json={"inputs": {}, "response_mode": response_mode}) - - _assert_json_response(response, status=status, body={"code": code, "message": message, "status": status}) - assert len(runtime.calls) == 1 - assert runtime.calls[0].streaming == (response_mode == "streaming") - assert _last_used_at(harness, sqlite_session_factory) == _USED_AT - if status == 401: - assert response.headers["WWW-Authenticate"] == 'Bearer realm="api"' - - -@pytest.mark.parametrize( - ("payload", "error"), - [ - ({}, {"type": "missing", "loc": ["inputs"], "msg": "Field required"}), - ({"inputs": None}, {"type": "dict_type", "loc": ["inputs"], "msg": "Input should be a valid dictionary"}), - ( - {"inputs": {}, "response_mode": "invalid"}, - {"type": "literal_error", "loc": ["response_mode"], "msg": "Input should be 'blocking' or 'streaming'"}, - ), - ], -) -def test_invalid_payload_does_not_record_usage_or_start_runtime( - harness: _Harness, - runtime: _Runtime, - sqlite_session_factory: sessionmaker[Session], - payload: dict[str, object], - error: dict[str, object], -) -> None: - response = harness.app.test_client().post(_url(harness), json=payload) - - _assert_json_response( - response, - status=422, - body={"code": "unprocessable_entity", "message": json.dumps([error]), "status": 422}, - ) - assert runtime.calls == [] - assert _last_used_at(harness, sqlite_session_factory) is None - - -def test_wrong_mode_does_not_record_usage_or_start_runtime( - harness: _Harness, - runtime: _Runtime, - sqlite_session_factory: sessionmaker[Session], -) -> None: - with sqlite_session_factory.begin() as session: - app = session.get(App, harness.target_app.id) - assert app is not None - app.mode = AppMode.CHAT - - response = harness.app.test_client().post(_url(harness), json={"inputs": {}}) - - _assert_json_response( - response, status=400, body={"code": "not_completion_app", "message": "Not Completion App", "status": 400} - ) - assert runtime.calls == [] - assert _last_used_at(harness, sqlite_session_factory) is None - - -def test_completion_uses_the_mode_captured_at_admission( - harness: _Harness, - runtime: _Runtime, - sqlite_session_factory: sessionmaker[Session], -) -> None: - harness.state.permission_action = lambda: _set_app_mode(harness, sqlite_session_factory, AppMode.CHAT) - - response = harness.app.test_client().post(_url(harness), json={"inputs": {}}) - - assert response.status_code == 200 - assert response.get_json() == _BLOCKING_RESPONSE - assert len(runtime.calls) == 1 - assert _last_used_at(harness, sqlite_session_factory) == _USED_AT - - -@pytest.mark.parametrize("rejection", ["permission", "missing", "tenant"]) -def test_completion_requires_installed_app_admission_before_payload_validation( - harness: _Harness, - runtime: _Runtime, - sqlite_session_factory: sessionmaker[Session], - rejection: str, -) -> None: - if rejection == "permission": - harness.state.allowed = False - elif rejection == "tenant": - harness.account._current_tenant = Tenant(name="Other workspace") - else: - with sqlite_session_factory.begin() as session: - installation = session.get(InstalledApp, harness.installed_app.id) - assert installation is not None - session.delete(installation) - - response = harness.app.test_client().post(_url(harness), json={}) - - if rejection == "permission": - _assert_json_response( - response, status=403, body={"code": "access_denied", "message": "App access denied.", "status": 403} - ) - else: - _assert_json_response( - response, - status=404, - body={ - "code": "installed_app_not_found", - "message": "The app was not found in this workspace.", - "status": 404, - "details": {"request_id": "request-1"}, - }, - ) - assert runtime.calls == [] - if rejection != "missing": - assert _last_used_at(harness, sqlite_session_factory) is None - - -@pytest.mark.parametrize("deleted_resource", ["app", "installation"]) -@pytest.mark.usefixtures("runtime") -def test_resource_removed_after_admission_is_revalidated_before_generation( - harness: _Harness, - sqlite_session_factory: sessionmaker[Session], - monkeypatch: pytest.MonkeyPatch, - deleted_resource: str, -) -> None: - def delete_resource() -> None: - with sqlite_session_factory.begin() as session: - if deleted_resource == "app": - resource = session.get(App, harness.target_app.id) - else: - resource = session.get(InstalledApp, harness.installed_app.id) - assert resource is not None - session.delete(resource) - - harness.state.permission_action = delete_resource - services = _Services( - installed_apps=_InstalledAppServices( - generation=InstalledAppGenerationService( - usage=SQLAlchemyInstalledAppRepository(session_factory=sqlite_session_factory), - runtime=AppGenerateServiceRuntime(session_factory=sqlite_session_factory), - ) - ) - ) - monkeypatch.setattr(completion_module, "application_services", lambda: services) - - def generate(**_kwargs: object) -> GenerationResponse: - pytest.fail("Removed app or installation must be rejected before generation") - - monkeypatch.setattr(AppGenerateService, "generate", generate) - - response = harness.app.test_client().post(_url(harness), json={"inputs": {}}) - - if deleted_resource == "app": - _assert_json_response( - response, - status=400, - body={ - "code": "app_unavailable", - "message": "App unavailable, please check your app configurations.", - "status": 400, - }, - ) - assert _last_used_at(harness, sqlite_session_factory) == _USED_AT - else: - _assert_json_response( - response, - status=404, - body={"code": "not_found", "message": "Installed app not found", "status": 404}, - ) diff --git a/api/tests/unit_tests/controllers/console/explore/test_saved_message.py b/api/tests/unit_tests/controllers/console/explore/test_saved_message.py index 4f971feb178..da9f5cd40fd 100644 --- a/api/tests/unit_tests/controllers/console/explore/test_saved_message.py +++ b/api/tests/unit_tests/controllers/console/explore/test_saved_message.py @@ -23,7 +23,7 @@ from services.saved_message_service import ( SavedMessagePage, SavedMessageRecord, ) -from tests.unit_tests.model_factories import make_account, make_tenant +from tests.helpers.unit.model_factories import make_account, make_tenant _INPUT_FILE_URL = "https://example.com/input.pdf" _CREATED_AT = datetime(2026, 1, 2, 3, 4, 5, tzinfo=UTC) diff --git a/api/tests/unit_tests/controllers/console/explore/test_trial.py b/api/tests/unit_tests/controllers/console/explore/test_trial.py index 7af229dee80..d79d76d14c9 100644 --- a/api/tests/unit_tests/controllers/console/explore/test_trial.py +++ b/api/tests/unit_tests/controllers/console/explore/test_trial.py @@ -1,47 +1,4074 @@ -import controllers.console.explore.trial as module +import hashlib +import io +import json +import zipfile +from collections.abc import Callable, Iterator, Mapping, Sequence +from dataclasses import dataclass, field +from datetime import datetime +from decimal import Decimal +from io import BytesIO +from types import SimpleNamespace +from typing import Literal, cast, override +from unittest.mock import Mock +from urllib.parse import parse_qs, urlsplit +from uuid import UUID, uuid4 + +import httpx +import pytest +from flask import Flask, Request, got_request_exception +from pydantic import JsonValue +from sqlalchemy import Connection, Engine, delete, event, func, select, text, update +from sqlalchemy.engine import ExecutionContext +from sqlalchemy.orm import Mapper, Session, SessionTransaction, sessionmaker +from sqlalchemy.pool import QueuePool +from werkzeug.exceptions import Forbidden, Unauthorized +from werkzeug.test import TestResponse + +import controllers.common.fields as fields_module +import controllers.console.explore.trial as trial_module +import controllers.console.explore.trial_app_admission as _trial_audio_admission_module +import controllers.console.explore.trial_app_admission as _trial_file_upload_admission_module +import controllers.console.explore.trial_app_admission as _trial_generation_admission_module +import controllers.console.explore.trial_app_admission as _trial_suggested_questions_admission_module +import controllers.console.files as files_module +import controllers.console.remote_files as remote_module +import controllers.console.wraps as console_wraps +import libs.login as login_module +import services.app_service as app_module +import services.file_service as file_module +import services.message_service as message_module +from constants import COOKIE_NAME_CSRF_TOKEN, HEADER_NAME_CSRF_TOKEN +from controllers.common.errors import UnauthorizedError +from controllers.console import api as console_api from controllers.console import console_ns +from controllers.console.app import agent_config_inspector as inspector +from controllers.console.app import preview_admission +from controllers.console.app import preview_admission as _trial_preview_admission_module +from controllers.console.app import preview_admission as _trial_preview_details_admission_module +from controllers.console.explore import trial +from core.agent.entities import AgentToolEntity +from core.app.apps.agent_app.errors import AgentAppNotPublishedError +from core.app.entities.app_invoke_entities import InvokeFrom +from core.errors.error import ( + AppInvokeQuotaExceededError, + ModelCurrentlyNotSupportError, + ProviderTokenNotInitError, + QuotaExceededError, +) +from core.file import remote_fetcher +from core.helper import encrypter +from core.model_context import get_credit_usage_metadata +from core.model_manager import ModelInstance +from core.ops.ops_trace_manager import TraceTask +from core.plugin.plugin_service import PluginService +from core.tools.__base.tool import Tool +from core.tools.entities.tool_entities import ToolParameter +from core.tools.errors import ToolSSRFError +from core.tools.tool_manager import ToolManager +from core.workflow.llm_environment_variable import LLMEnvironmentVariable +from enums import DeploymentEdition +from extensions.ext_database import db +from extensions.ext_login import DifyLoginManager, unauthorized_handler +from extensions.ext_storage import StorageType +from graphon.model_runtime.entities import AssistantPromptMessage, PromptMessage +from graphon.model_runtime.entities.llm_entities import LLMResult, LLMUsage +from graphon.model_runtime.entities.model_entities import AIModelEntity, ModelType +from graphon.model_runtime.errors.invoke import InvokeBadRequestError, InvokeError +from graphon.variables import SecretVariable, StringVariable +from libs.external_api import ExternalApi +from libs.token import generate_csrf_token +from machinery.errors import ActiveWorkspaceRequiredError +from models import Account, AccountTrialAppRecord, Conversation, Message, Tenant, TrialApp +from models import App as _trial_audio_App +from models import App as _trial_file_upload_App +from models import App as _trial_generation_App +from models import App as _trial_preview_App +from models import App as _trial_preview_details_App +from models import App as _trial_suggested_questions_App +from models import AppMode as _trial_audio_AppMode +from models import AppMode as _trial_file_upload_AppMode +from models import AppMode as _trial_generation_AppMode +from models import AppMode as _trial_preview_AppMode +from models import AppMode as _trial_preview_details_AppMode +from models import AppMode as _trial_suggested_questions_AppMode +from models import UploadFile as _trial_file_upload_UploadFile +from models.account import AccountStatus, TenantAccountRole, TenantStatus +from models.agent import ( + Agent, + AgentConfigDraft, + AgentConfigDraftType, + AgentConfigRevision, + AgentConfigRevisionOperation, + AgentConfigSnapshot, + AgentScope, + AgentSource, +) +from models.agent_config_entities import AgentSoulConfig +from models.dataset import Dataset +from models.enums import ConversationFromSource, CreatorUserRole, CustomizeTokenStrategy, TagType +from models.model import App as _trial_agent_config_App +from models.model import AppMode as _trial_agent_config_AppMode +from models.model import AppModelConfig, IconType, RecommendedApp, Site, Tag, TagBinding +from models.model import UploadFile as _trial_agent_config_UploadFile +from models.provider_ids import GenericProviderID +from models.tools import ToolFile, WorkflowToolProvider +from models.workflow import Workflow +from repositories.app_definition_query_repository import AppDefinitionQueryRepository +from repositories.app_preview_query_repository import AppPreviewQueryRepository +from repositories.recommended_app_catalog_repository import DatabaseRecommendedAppCatalogRepository +from repositories.trial_app_repository import TrialAppRepository +from services.account_errors import AccountNotFoundError +from services.agent.errors import AgentVersionNotFoundError +from services.agent_config_service import AgentConfigService +from services.app_definition_query_service import AppDefinitionQueryService, AppDefinitionUnavailableError +from services.app_preview_details_adapters import AppPreviewDetailsRuntime +from services.app_preview_query_service import AppPreviewQueryService +from services.app_service import AppResponseView +from services.audio_types import AudioAppRef, AudioOutput, AudioUpload +from services.entities.feature_entities import FeatureModel +from services.errors.app_model_config import AppModelConfigBrokenError +from services.errors.audio import ( + AudioTooLargeServiceError, + NoAudioUploadedServiceError, + ProviderNotSupportSpeechToTextServiceError, + ProviderNotSupportTextToSpeechServiceError, + SpeechToTextDisabledServiceError, + UnsupportedAudioTypeServiceError, +) +from services.errors.conversation import ConversationCompletedError, ConversationNotExistsError +from services.errors.llm import InvokeRateLimitError +from services.feature_service import FeatureService +from services.file_service import FileService +from services.message_suggested_questions_adapters import MessageSuggestedQuestionsRuntime +from services.message_suggested_questions_service import ( + SuggestedQuestionsAccount, + SuggestedQuestionsActor, + SuggestedQuestionsActorNotFoundError, +) +from services.recommended_app_query_service import ( + RecommendedAppCatalogPage, + RecommendedAppDetailRecord, + RecommendedAppQueryService, +) +from services.remote_file_service import RemoteFileService +from services.trial_app_access_service import TrialAppAccessService, TrialAppRef +from services.trial_app_generation_service import GenerationResponse, TrialAppGenerationService -def test_trial_workflow_uses_trial_scoped_simple_account_model() -> None: - assert module.simple_account_model.name == "TrialSimpleAccount" - assert module.simple_account_model.__schema__["properties"].keys() >= {"id", "name", "email"} +class TestTrialCases: + @staticmethod + def test_trial_workflow_uses_trial_scoped_simple_account_model() -> None: + assert trial_module.simple_account_model.name == "TrialSimpleAccount" + assert trial_module.simple_account_model.__schema__["properties"].keys() >= {"id", "name", "email"} + + @staticmethod + def test_trial_routes_keep_paths_and_endpoints() -> None: + expected = { + "TrialAppFileUploadApi": (("/trial-apps//files/upload",), "trial_app_file_upload"), + "TrialAppRemoteFileUploadApi": ( + ("/trial-apps//remote-files/upload",), + "trial_app_remote_file_upload", + ), + "TrialAppWorkflowRunApi": (("/trial-apps//workflows/run",), "trial_app_workflow_run"), + "TrialAppWorkflowTaskStopApi": ( + ("/trial-apps//workflows/tasks//stop",), + None, + ), + "TrialChatApi": (("/trial-apps//chat-messages",), "trial_app_chat_completion"), + "TrialChatTaskStopApi": ( + ("/trial-apps//chat-messages//stop",), + "trial_app_stop_chat_completion", + ), + "TrialAgentComposerApi": (("/trial-apps//agent-composer",), None), + "TrialMessageSuggestedQuestionApi": ( + ("/trial-apps//messages//suggested-questions",), + "trial_app_suggested_question", + ), + "TrialChatAudioApi": (("/trial-apps//audio-to-text",), "trial_app_audio"), + "TrialChatTextApi": (("/trial-apps//text-to-audio",), "trial_app_text"), + "TrialCompletionApi": (("/trial-apps//completion-messages",), "trial_app_completion"), + "TrialSitApi": (("/trial-apps//site",), None), + "TrialAppParameterApi": (("/trial-apps//parameters",), "trial_app_parameters"), + "AppApi": (("/trial-apps/",), "trial_app"), + "AppWorkflowApi": (("/trial-apps//workflows",), "trial_app_workflow"), + "DatasetListApi": (("/trial-apps//datasets",), "trial_app_datasets"), + } + actual = { + resource.__name__: (urls, kwargs.get("endpoint")) + for resource, urls, _route_doc, kwargs in console_ns.resources + if resource.__module__ == trial_module.__name__ + } + + assert actual == expected -def test_trial_routes_keep_paths_and_endpoints() -> None: - expected = { - "TrialAppFileUploadApi": (("/trial-apps//files/upload",), "trial_app_file_upload"), - "TrialAppRemoteFileUploadApi": ( - ("/trial-apps//remote-files/upload",), - "trial_app_remote_file_upload", - ), - "TrialAppWorkflowRunApi": (("/trial-apps//workflows/run",), "trial_app_workflow_run"), - "TrialAppWorkflowTaskStopApi": ( - ("/trial-apps//workflows/tasks//stop",), - None, - ), - "TrialChatApi": (("/trial-apps//chat-messages",), "trial_app_chat_completion"), - "TrialChatTaskStopApi": ( - ("/trial-apps//chat-messages//stop",), - "trial_app_stop_chat_completion", - ), - "TrialAgentComposerApi": (("/trial-apps//agent-composer",), None), - "TrialMessageSuggestedQuestionApi": ( - ("/trial-apps//messages//suggested-questions",), - "trial_app_suggested_question", - ), - "TrialChatAudioApi": (("/trial-apps//audio-to-text",), "trial_app_audio"), - "TrialChatTextApi": (("/trial-apps//text-to-audio",), "trial_app_text"), - "TrialCompletionApi": (("/trial-apps//completion-messages",), "trial_app_completion"), - "TrialSitApi": (("/trial-apps//site",), None), - "TrialAppParameterApi": (("/trial-apps//parameters",), "trial_app_parameters"), - "AppApi": (("/trial-apps/",), "trial_app"), - "AppWorkflowApi": (("/trial-apps//workflows",), "trial_app_workflow"), - "DatasetListApi": (("/trial-apps//datasets",), "trial_app_datasets"), - } - actual = { - resource.__name__: (urls, kwargs.get("endpoint")) - for resource, urls, _route_doc, kwargs in console_ns.resources - if resource.__module__ == module.__name__ - } +"""Exercise public template assets through real admission, snapshots, and file services.""" - assert actual == expected + +@dataclass +class Harness: + app: Flask + factory: sessionmaker[Session] + source: _trial_agent_config_App + agent: Agent + snapshot: AgentConfigSnapshot + upload_id: str + skill_id: str + reads: list[str] + + def url(self, suffix: str) -> str: + return f"/console/api/trial-apps/{self.source.id}/agent/config/{suffix}" + + def get(self, suffix: str, **query: str) -> TestResponse: + return self.app.test_client().get(self.url(suffix), query_string=query) + + +class TestTrialAgentConfigCases: + @staticmethod + @pytest.fixture + def harness( + sqlite_session_factory: sessionmaker[Session], + monkeypatch: pytest.MonkeyPatch, + config_overrides: Callable[..., None], + ) -> Harness: + config_overrides(SECRET_KEY="trial-config-test", FILES_URL="http://localhost") + tenant_id, agent_id, snapshot_id, account_id = (str(uuid4()) for _ in range(4)) + source = _trial_agent_config_App( + id=str(uuid4()), + tenant_id=tenant_id, + name="Template", + mode=_trial_agent_config_AppMode.AGENT, + is_public=True, + enable_site=False, + enable_api=False, + ) + agent = Agent( + id=agent_id, + tenant_id=tenant_id, + app_id=source.id, + name="Template", + scope=AgentScope.ROSTER, + source=AgentSource.AGENT_APP, + active_config_snapshot_id=snapshot_id, + ) + skill_id = str(uuid4()) + skill_zip = io.BytesIO() + with zipfile.ZipFile(skill_zip, "w") as archive: + archive.writestr("SKILL.md", "---\nname: research\ndescription: Research\n---\nPublished skill") + archive.writestr("scripts/run.py", "print('published')") + payloads = {"skill.zip": skill_zip.getvalue(), "guide.txt": b"Published attachment"} + upload = _trial_agent_config_UploadFile( + tenant_id=tenant_id, + storage_type=StorageType.LOCAL, + key="guide.txt", + name="guide.txt", + size=len(payloads["guide.txt"]), + extension="txt", + mime_type="text/plain", + created_by=account_id, + created_by_role=CreatorUserRole.ACCOUNT, + created_at=datetime(2026, 1, 1), + used=False, + ) + snapshot = AgentConfigSnapshot( + id=snapshot_id, + tenant_id=tenant_id, + agent_id=agent_id, + version=1, + config_snapshot=AgentSoulConfig.model_validate( + { + "config_skills": [{"name": "research", "file_id": skill_id}], + "config_files": [{"name": "guide.txt", "file_kind": "upload_file", "file_id": upload.id}], + } + ), + ) + skill_file = ToolFile( + tenant_id=tenant_id, + user_id=account_id, + conversation_id=None, + file_key="skill.zip", + mimetype="application/zip", + name="research.zip", + size=len(payloads["skill.zip"]), + ) + skill_file.id = skill_id + with sqlite_session_factory.begin() as session: + session.add_all( + [ + source, + agent, + snapshot, + upload, + skill_file, + RecommendedApp( + app_id=source.id, + description={}, + copyright="", + privacy_policy="", + category="Agent", + categories=["Agent"], + custom_disclaimer="", + position=1, + is_listed=True, + language="en-US", + ), + AgentConfigRevision( + tenant_id=tenant_id, + agent_id=agent_id, + current_snapshot_id=snapshot_id, + revision=1, + operation=AgentConfigRevisionOperation.PUBLISH_DRAFT, + ), + AgentConfigDraft( + tenant_id=tenant_id, + agent_id=agent_id, + draft_type=AgentConfigDraftType.DRAFT, + draft_owner_key="", + base_snapshot_id=snapshot_id, + config_snapshot=AgentSoulConfig.model_validate( + { + "config_files": [ + {"name": "private.txt", "file_kind": "upload_file", "file_id": upload.id} + ], + } + ), + ), + ] + ) + + reads: list[str] = [] + + def load_once(key: str) -> bytes: + reads.append(key) + return payloads[key] + + monkeypatch.setattr("services.agent_config_service.storage", SimpleNamespace(load_once=load_once)) + catalog = DatabaseRecommendedAppCatalogRepository(sqlite_session_factory, redis=Mock()) + services = SimpleNamespace( + app_previews=AppPreviewQueryService( + apps=AppPreviewQueryRepository(session_factory=sqlite_session_factory), + is_previewable=catalog.contains, + ) + ) + monkeypatch.setattr(preview_admission, "application_services", lambda: services) + monkeypatch.setattr(inspector, "application_services", lambda: services) + monkeypatch.setattr(trial, "application_services", lambda: services) + monkeypatch.setattr(inspector, "_service", lambda: AgentConfigService(session_factory=sqlite_session_factory)) + app = Flask(__name__) + api = ExternalApi(app) + api.add_resource(trial.TrialAgentComposerApi, "/console/api/trial-apps//agent-composer") + routes = [ + (inspector.TrialAgentConfigSkillsApi, "skills"), + (inspector.TrialAgentConfigFilesApi, "files"), + (inspector.TrialAgentConfigSkillInspectApi, "skills//inspect"), + (inspector.TrialAgentConfigSkillDownloadApi, "skills//download"), + (inspector.TrialAgentConfigSkillFilePreviewApi, "skills//files/preview"), + (inspector.TrialAgentConfigSkillFileDownloadApi, "skills//files/download"), + (inspector.TrialAgentConfigFilePreviewApi, "files//preview"), + (inspector.TrialAgentConfigFileDownloadApi, "files//download"), + ] + prefix = "/console/api/trial-apps//agent/config/" + for resource, suffix in routes: + api.add_resource(resource, prefix + suffix) + api.add_resource( + inspector.TrialAgentConfigSkillFileContentApi, + prefix + "skills//files/content", + endpoint="console.trial_agent_config_skill_file_content", + ) + return Harness(app, sqlite_session_factory, source, agent, snapshot, upload.id, skill_id, reads) + + @staticmethod + def test_published_template_configuration_is_readable_without_workspace_membership(harness: Harness) -> None: + composer_url = f"/console/api/trial-apps/{harness.source.id}/agent-composer" + composer = harness.app.test_client().get(composer_url) + assert composer.status_code == 200 + assert composer.json is not None + assert composer.json["agent_soul"]["config_files"][0]["name"] == "guide.txt" + assert composer.json["active_config_snapshot"]["id"] == harness.snapshot.id + assert composer.json["draft"] is None + assert composer.json["save_options"] == [] + for kind, name in [("skills", "research"), ("files", "guide.txt")]: + response = harness.get(kind) + assert response.status_code == 200 + assert response.json is not None + assert response.json["items"][0]["name"] == name + assert response.json["config_version"] == {"id": harness.snapshot.id, "kind": "snapshot", "writable": False} + response = harness.get("skills/research/inspect") + assert response.status_code == 200 + assert response.json is not None + assert "Published skill" in response.json["skill_md"]["text"] + response = harness.get("skills/research/files/preview", path="scripts/run.py") + assert response.status_code == 200 + assert response.json is not None + assert response.json["text"] == "print('published')" + response = harness.get("files/guide.txt/preview", version_id=harness.snapshot.id) + assert response.status_code == 200 + assert response.json is not None + assert response.json["text"] == "Published attachment" + assert response.headers["Cache-Control"] == "no-store" + for path in ["files/guide.txt/download", "skills/research/download"]: + response = harness.get(path) + assert response.status_code == 200 + assert response.json is not None + assert {"timestamp", "nonce", "sign"} <= parse_qs(urlsplit(response.json["url"]).query).keys() + response = harness.get("skills/research/files/download", path="scripts/run.py") + assert response.json is not None + url = response.json["url"] + download = harness.app.test_client().get(url) + assert download.status_code == 200 + assert download.data == b"print('published')" + assert download.headers["Content-Disposition"].startswith("attachment;") + with harness.factory() as session: + assert session.scalar(select(func.count()).select_from(AgentConfigDraft)) == 1 + assert session.scalar(select(func.count()).select_from(_trial_agent_config_UploadFile)) == 1 + with harness.factory.begin() as session: + app = session.get(_trial_agent_config_App, harness.source.id) + assert app is not None + app.is_public = False + reads_before = list(harness.reads) + assert harness.app.test_client().get(url).status_code == 404 + assert harness.app.test_client().get(composer_url).status_code == 404 + assert harness.reads == reads_before + + @staticmethod + @pytest.mark.parametrize( + ("query", "status"), + [ + ({"draft_type": "draft"}, 422), + ({"file_id": "other"}, 422), + ({"version_id": str(uuid4())}, 404), + ], + ) + def test_version_and_resource_selectors_cannot_bypass_admission( + harness: Harness, + query: dict[str, str], + status: int, + ) -> None: + assert harness.get("files/guide.txt/preview", **query).status_code == status + assert harness.reads == [] + + @staticmethod + @pytest.mark.parametrize("path", ["../SKILL.md", "scripts/../../SKILL.md"]) + def test_skill_member_paths_cannot_escape_archive(harness: Harness, path: str) -> None: + assert harness.get("skills/research/files/preview", path=path).status_code == 400 + assert harness.reads == [] + + @staticmethod + @pytest.mark.parametrize("kind", ["file", "skill"]) + def test_foreign_tenant_payloads_are_not_read(harness: Harness, kind: str) -> None: + with harness.factory.begin() as session: + if kind == "file": + upload = session.get(_trial_agent_config_UploadFile, harness.upload_id) + assert upload is not None + upload.tenant_id = str(uuid4()) + else: + skill = session.get(ToolFile, harness.skill_id) + assert skill is not None + skill.tenant_id = str(uuid4()) + suffix = "files/guide.txt/preview" if kind == "file" else "skills/research/files/preview" + assert harness.get(suffix, **({"path": "SKILL.md"} if kind == "skill" else {})).status_code == 404 + assert harness.reads == [] + + @staticmethod + def test_draft_only_files_and_writes_are_unavailable(harness: Harness) -> None: + assert harness.get("files/private.txt/preview").status_code == 404 + for method in ["POST", "PUT", "PATCH", "DELETE"]: + assert harness.app.test_client().open(harness.url("files"), method=method).status_code == 405 + assert harness.reads == [] + + @staticmethod + @pytest.mark.parametrize( + ("state", "status"), + [ + ("unlisted", 404), + ("unpublished", 400), + ("foreign_snapshot", 400), + ], + ) + def test_unavailable_template_never_reads_storage(harness: Harness, state: str, status: int) -> None: + with harness.factory.begin() as session: + if state == "unlisted": + listing = session.scalar(select(RecommendedApp).where(RecommendedApp.app_id == harness.source.id)) + assert listing is not None + listing.is_listed = False + elif state == "unpublished": + revision = session.scalar( + select(AgentConfigRevision).where(AgentConfigRevision.agent_id == harness.agent.id) + ) + assert revision is not None + revision.operation = AgentConfigRevisionOperation.IMPORT_PACKAGE + else: + snapshot = session.get(AgentConfigSnapshot, harness.snapshot.id) + assert snapshot is not None + snapshot.tenant_id = str(uuid4()) + preview = harness.get("files/guide.txt/preview") + assert preview.status_code == status + assert preview.json is not None + assert preview.json["status"] == status + if status == 400: + # Raised inside the template guard, so its error answer is built by the surface and not cached. + assert preview.headers["Cache-Control"] == "no-store" + composer = harness.app.test_client().get(f"/console/api/trial-apps/{harness.source.id}/agent-composer") + assert composer.status_code == status + if status == 400: + assert composer.json is not None + assert composer.json["code"] == "app_unavailable" + assert harness.reads == [] + + +"""Trial audio admission, response and usage contracts through HTTP and SQLite.""" + + +_ASR = "audio-to-text" +_TTS = "text-to-audio" +_WAV = b"RIFF\x24\x00\x00\x00WAVEfmt " + b"a" * 32 + + +@dataclass +class _TrialAudioFeaturestate: + enabled: bool = True + setup_completed: bool = True + events: list[str] = field(default_factory=list) + + def is_trial_enabled(self) -> bool: + self.events.append("feature") + return self.enabled + + +@dataclass(frozen=True) +class _AsrCall: + app: AudioAppRef + content: bytes | None + mime_type: str | None + + +@dataclass(frozen=True) +class _TtsCall: + app: AudioAppRef + account_id: str + text: str | None + voice: str | None + message_id: str | None + + +@dataclass +class _TrialAudioStream(Iterator[bytes]): + chunks: Iterator[bytes] = field(default_factory=lambda: iter([b"a" * 32, b"remaining-audio"])) + reads: int = 0 + close_calls: int = 0 + close_error: Exception | None = None + + @override + def __next__(self) -> bytes: + chunk = next(self.chunks) + self.reads += 1 + return chunk + + def close(self) -> None: + self.close_calls += 1 + if self.close_error is not None: + raise self.close_error + + +@dataclass +class _TrialAudioRuntime: + sessions: list[Session] + output: AudioOutput | None = field(default_factory=lambda: AudioOutput(data=_WAV, mime_type=None)) + error: Exception | None = None + calls: list[_AsrCall | _TtsCall] = field(default_factory=list) + + def _assert_admission_closed(self) -> None: + assert self.sessions + assert all(not session.in_transaction() and not session.identity_map for session in self.sessions) + + def transcript_asr(self, *, app: AudioAppRef, audio: AudioUpload | None) -> dict[str, str]: + self._assert_admission_closed() + self.calls.append(_AsrCall(app, audio.stream.read() if audio else None, audio.mime_type if audio else None)) + if self.error is not None: + raise self.error + if audio is None: + raise NoAudioUploadedServiceError() + return {"text": " transcript "} + + def transcript_tts( + self, + *, + app: AudioAppRef, + account_id: str, + text: str | None, + voice: str | None, + message_id: str | None, + ) -> AudioOutput | None: + self._assert_admission_closed() + self.calls.append(_TtsCall(app, account_id, text, voice, message_id)) + if self.error is not None: + raise self.error + return self.output + + +@dataclass(frozen=True) +class _TrialAudioTrialappservices: + access: TrialAppAccessService + usage: TrialAppRepository + + +@dataclass(frozen=True) +class _TrialAudioApplicationservices: + trial_apps: _TrialAudioTrialappservices + app_audio: _TrialAudioRuntime + recommended_app_queries: _TrialAudioFeaturestate + + +@dataclass(frozen=True) +class _TrialAudioHarness: + app: Flask + account: Account + target: _trial_audio_App + trial: TrialApp + factory: sessionmaker[Session] + runtime: _TrialAudioRuntime + state: _TrialAudioFeaturestate + + def post(self, endpoint: str, *, payload: Mapping[str, object] | None = None) -> TestResponse: + client = self.app.test_client() + url = f"/trial-apps/{self.target.id}/{endpoint}" + if endpoint == _ASR: + return client.post(url, data={"file": (BytesIO(b"audio input"), "recording.m4a", "audio/x-m4a")}) + return client.post(url, json=dict(payload) if payload is not None else {"text": "Hello"}) + + def usage(self) -> int | None: + with self.factory() as session: + return session.scalar( + select(AccountTrialAppRecord.count).where( + AccountTrialAppRecord.app_id == self.target.id, + AccountTrialAppRecord.account_id == self.account.id, + ) + ) + + +def _trial_audio_assert_error(response: TestResponse, status: int, code: str) -> Mapping[str, object]: + assert response.status_code == status + body = response.get_json() + assert isinstance(body, dict) + assert body["code"] == code + assert body["status"] == status + assert response.headers["Content-Type"] == "application/json" + assert isinstance(body["message"], str) + assert body["message"] + return body + + +class TestTrialAudioCases: + @staticmethod + @pytest.fixture + def harness( + monkeypatch: pytest.MonkeyPatch, + config_overrides: Callable[..., None], + sqlite_session_factory: sessionmaker[Session], + ) -> _TrialAudioHarness: + config_overrides(DEPLOYMENT_EDITION=DeploymentEdition.COMMUNITY, INIT_PASSWORD="", LOGIN_DISABLED=False) + state = _TrialAudioFeaturestate() + account = Account(name="Viewer", email="viewer@example.com") + account.status = AccountStatus.ACTIVE + account._current_tenant = Tenant(name="Viewer workspace") + target = _trial_audio_App( + id=str(uuid4()), + tenant_id=str(uuid4()), + name="Trial app", + mode=_trial_audio_AppMode.CHAT, + enable_site=True, + enable_api=True, + ) + trial = TrialApp(app_id=target.id, tenant_id=target.tenant_id, trial_limit=3) + with sqlite_session_factory.begin() as session: + session.add_all([target, trial]) + + sessions: list[Session] = [] + repository_factory = sessionmaker(bind=sqlite_session_factory.kw["bind"], expire_on_commit=False) + + @event.listens_for(repository_factory, "after_begin") + def track_session(session: Session, _transaction: SessionTransaction, _connection: Connection) -> None: + sessions.append(session) + + runtime = _TrialAudioRuntime(sessions=sessions) + repository = TrialAppRepository(session_factory=repository_factory) + services = _TrialAudioApplicationservices( + trial_apps=_TrialAudioTrialappservices(access=TrialAppAccessService(apps=repository), usage=repository), + app_audio=runtime, + recommended_app_queries=state, + ) + + def setup_completed() -> bool: + state.events.append("setup") + return state.setup_completed + + def csrf(_request: Request, account_id: str) -> None: + assert account_id == account.id + state.events.append("csrf") + + for module in (trial_module, _trial_audio_admission_module): + monkeypatch.setattr(module, "application_services", lambda: services) + monkeypatch.setattr(console_wraps, "_is_setup_completed", setup_completed) + monkeypatch.setattr(login_module, "current_user", account) + monkeypatch.setattr(login_module, "check_csrf_token", csrf) + app = Flask(__name__) + app.config.update(TESTING=True, RESTX_ERROR_404_HELP=False) + login_manager = DifyLoginManager() + login_manager.init_app(app) + login_manager.unauthorized_handler(unauthorized_handler) + api = ExternalApi(app) + api.add_resource(trial_module.TrialChatAudioApi, "/trial-apps//audio-to-text") + api.add_resource(trial_module.TrialChatTextApi, "/trial-apps//text-to-audio") + return _TrialAudioHarness(app, account, target, trial, sqlite_session_factory, runtime, state) + + @staticmethod + def test_asr_upload_keeps_transcript_and_owner_tenant(harness: _TrialAudioHarness) -> None: + response = harness.post(_ASR) + + assert response.status_code == 200 + assert response.get_json() == {"text": " transcript "} + assert response.headers["Content-Type"] == "application/json" + assert int(response.headers["Content-Length"]) == len(response.data) + assert harness.runtime.calls == [ + _AsrCall(AudioAppRef(harness.target.id, harness.target.tenant_id, "chat"), b"audio input", "audio/x-m4a") + ] + assert harness.target.tenant_id != harness.account.current_tenant_id + assert harness.state.events == ["setup", "csrf", "feature"] + assert harness.usage() == 1 + + @staticmethod + @pytest.mark.parametrize("streaming", [None, False, True]) + def test_tts_preserves_payload_actor_and_binary_response_ignoring_streaming_flag( + harness: _TrialAudioHarness, streaming: bool | None + ) -> None: + message_id = str(uuid4()) + response = harness.post( + _TTS, payload={"text": " hello ", "voice": "voice-1", "message_id": message_id, "streaming": streaming} + ) + + assert response.status_code == 200 + assert response.data == _WAV + assert response.headers["Content-Type"] == "audio/wav" + assert int(response.headers["Content-Length"]) == len(_WAV) + assert harness.runtime.calls == [ + _TtsCall( + AudioAppRef(harness.target.id, harness.target.tenant_id, "chat"), + harness.account.id, + " hello ", + "voice-1", + message_id, + ) + ] + assert harness.usage() == 1 + + @staticmethod + @pytest.mark.parametrize("null_fields", [False, True]) + def test_missing_message_null_response_still_records_usage(harness: _TrialAudioHarness, null_fields: bool) -> None: + harness.runtime.output = None + payload: dict[str, object] = {"text": None, "voice": None, "message_id": None} if null_fields else {} + response = harness.post(_TTS, payload=payload) + + assert response.status_code == 200 + assert response.get_json() is None + assert response.headers["Content-Type"] == "application/json" + assert harness.runtime.calls == [ + _TtsCall( + AudioAppRef(harness.target.id, harness.target.tenant_id, "chat"), harness.account.id, None, None, None + ) + ] + assert harness.usage() == 1 + + @staticmethod + def test_missing_upload_passes_none_and_does_not_record_usage(harness: _TrialAudioHarness) -> None: + response = harness.app.test_client().post(f"/trial-apps/{harness.target.id}/{_ASR}", data={}) + + _trial_audio_assert_error(response, 400, "no_audio_uploaded") + assert harness.runtime.calls == [ + _AsrCall(AudioAppRef(harness.target.id, harness.target.tenant_id, "chat"), None, None) + ] + assert harness.usage() is None + + @staticmethod + def test_oversized_multipart_request_keeps_transport_error(harness: _TrialAudioHarness) -> None: + harness.app.config["MAX_CONTENT_LENGTH"] = 1 + + response = harness.post(_ASR) + + _trial_audio_assert_error(response, 413, "request_entity_too_large") + assert harness.runtime.calls == [] + assert harness.usage() is None + + @staticmethod + @pytest.mark.parametrize("endpoint", [_ASR, _TTS]) + @pytest.mark.parametrize( + ("error", "status", "code"), + [ + (AppDefinitionUnavailableError("app removed"), 400, "app_unavailable"), + (AppModelConfigBrokenError(), 400, "app_unavailable"), + (ProviderTokenNotInitError("provider key missing"), 400, "provider_not_initialize"), + (QuotaExceededError(), 400, "provider_quota_exceeded"), + (ModelCurrentlyNotSupportError(), 400, "model_currently_not_support"), + (InvokeError("provider rejected audio"), 400, "completion_request_error"), + (ValueError("Text is required"), 400, "invalid_param"), + (UnauthorizedError("session expired"), 401, "unauthorized"), + (AgentVersionNotFoundError(), 404, "agent_version_not_found_error"), + (RuntimeError("private provider failure"), 500, "internal_server_error"), + (Forbidden("private provider failure"), 500, "internal_server_error"), + ], + ) + def test_runtime_errors_keep_specific_codes_without_recording_usage( + harness: _TrialAudioHarness, endpoint: str, error: Exception, status: int, code: str + ) -> None: + harness.runtime.error = error + response = harness.post(endpoint) + + body = _trial_audio_assert_error(response, status, code) + if isinstance(error, (ProviderTokenNotInitError, InvokeError)): + assert body["message"] == error.description + elif code in {"audio_too_large", "invalid_param"}: + assert str(error) in str(body["message"]) + elif status == 401: + assert body["message"] == "session expired" + assert response.headers["WWW-Authenticate"] == 'Bearer realm="api"' + elif status == 404: + assert body["message"] == "Agent config version not found." + elif status == 500: + assert body["message"] == ( + "The server encountered an internal error and was unable to complete your request. " + "Either the server is overloaded or there is an error in the application." + ) + assert harness.usage() is None + + @staticmethod + @pytest.mark.parametrize( + ("error", "status", "code"), + [ + (NoAudioUploadedServiceError(), 400, "no_audio_uploaded"), + (AudioTooLargeServiceError("Audio size larger than 30 mb"), 413, "audio_too_large"), + (UnsupportedAudioTypeServiceError(), 415, "unsupported_audio_type"), + (ProviderNotSupportSpeechToTextServiceError(), 400, "provider_not_support_speech_to_text"), + (SpeechToTextDisabledServiceError(), 400, "speech_to_text_disabled"), + ], + ) + def test_transcription_errors_keep_specific_codes_without_recording_usage( + harness: _TrialAudioHarness, error: Exception, status: int, code: str + ) -> None: + # Only transcription validates the upload and needs a speech-to-text model; synthesis documents none of these. + harness.runtime.error = error + body = _trial_audio_assert_error(harness.post(_ASR), status, code) + if code == "audio_too_large": + assert str(error) in str(body["message"]) + assert harness.usage() is None + + @staticmethod + def test_unsupported_tts_provider_has_specific_error(harness: _TrialAudioHarness) -> None: + harness.runtime.error = ProviderNotSupportTextToSpeechServiceError() + + _trial_audio_assert_error(harness.post(_TTS), 400, "provider_not_support_text_to_speech") + assert harness.usage() is None + + @staticmethod + @pytest.mark.parametrize("field", ["text", "voice", "message_id", "streaming"]) + def test_invalid_tts_payload_keeps_field_context(harness: _TrialAudioHarness, field: str) -> None: + response = harness.post(_TTS, payload={field: ["invalid"]}) + + body = _trial_audio_assert_error(response, 422, "unprocessable_entity") + assert field in str(body) + assert harness.runtime.calls == [] + assert harness.usage() is None + + @staticmethod + @pytest.mark.parametrize("endpoint", [_ASR, _TTS]) + @pytest.mark.parametrize( + ("gate", "status", "code"), + [ + ("setup", 401, "not_setup"), + ("login", 401, "unauthorized"), + ("csrf", 401, "unauthorized"), + ("initialization", 400, "account_not_initialized"), + ("feature", 403, "trial_app_feature_disabled"), + ], + ) + def test_console_and_trial_gates_run_before_app_lookup( + harness: _TrialAudioHarness, monkeypatch: pytest.MonkeyPatch, endpoint: str, gate: str, status: int, code: str + ) -> None: + if gate == "setup": + harness.state.setup_completed = False + elif gate == "login": + monkeypatch.setattr(login_module, "current_user", None) + elif gate == "csrf": + + def reject_csrf(_request: Request, _account_id: str) -> None: + raise Unauthorized("invalid csrf token") + + monkeypatch.setattr(login_module, "check_csrf_token", reject_csrf) + elif gate == "initialization": + harness.account.status = AccountStatus.UNINITIALIZED + else: + harness.state.enabled = False + response = harness.post(endpoint) + + if gate == "login": + assert response.status_code == 401 + assert response.get_json() == {"code": "unauthorized", "message": "Unauthorized.", "status": 401} + assert response.headers["Content-Type"] == "application/json" + else: + _trial_audio_assert_error(response, status, code) + assert harness.runtime.calls == [] + assert harness.runtime.sessions == [] + assert harness.usage() is None + + @staticmethod + @pytest.mark.parametrize("endpoint", [_ASR, _TTS]) + @pytest.mark.parametrize("gate", ["membership", "missing-app", "quota"]) + def test_trial_membership_and_account_quota_block_audio( + harness: _TrialAudioHarness, endpoint: str, gate: str + ) -> None: + with harness.factory.begin() as session: + if gate == "membership": + trial = session.get(TrialApp, harness.trial.id) + assert trial is not None + session.delete(trial) + elif gate == "missing-app": + app = session.get(_trial_audio_App, harness.target.id) + assert app is not None + session.delete(app) + else: + session.add(AccountTrialAppRecord(app_id=harness.target.id, account_id=harness.account.id, count=3)) + response = harness.post(endpoint) + + _trial_audio_assert_error( + response, 403, "trial_app_limit_exceeded" if gate == "quota" else "trial_app_not_allowed" + ) + assert harness.runtime.calls == [] + assert harness.usage() == (3 if gate == "quota" else None) + + @staticmethod + @pytest.mark.parametrize("endpoint", [_ASR, _TTS]) + def test_usage_is_isolated_to_current_account(harness: _TrialAudioHarness, endpoint: str) -> None: + other_id = str(uuid4()) + with harness.factory.begin() as session: + session.add_all( + [ + AccountTrialAppRecord(app_id=harness.target.id, account_id=harness.account.id, count=1), + AccountTrialAppRecord(app_id=harness.target.id, account_id=other_id, count=99), + ] + ) + response = harness.post(endpoint) + + assert response.status_code == 200 + assert harness.usage() == 2 + with harness.factory() as session: + assert ( + session.scalar(select(AccountTrialAppRecord.count).where(AccountTrialAppRecord.account_id == other_id)) + == 99 + ) + + @staticmethod + @pytest.mark.parametrize("failure", ["mime", "first-chunk"]) + def test_stream_preparation_failure_closes_provider_without_consuming_trial( + harness: _TrialAudioHarness, failure: str + ) -> None: + def broken_chunks() -> Iterator[bytes]: + raise InvokeBadRequestError("first audio chunk failed") + yield b"unreachable" + + stream = _TrialAudioStream(chunks=iter([_WAV]) if failure == "mime" else broken_chunks()) + harness.runtime.output = AudioOutput(data=stream, mime_type="audio/mpeg") + response = harness.post(_TTS) + + body = _trial_audio_assert_error(response, 400, "completion_request_error") + assert ("MIME" if failure == "mime" else "first audio chunk failed") in str(body["message"]) + assert stream.close_calls == 1 + assert harness.usage() is None + + @staticmethod + @pytest.mark.parametrize("cleanup_fails", [False, True]) + def test_usage_write_failure_closes_stream_without_hiding_database_error( + harness: _TrialAudioHarness, cleanup_fails: bool, caplog: pytest.LogCaptureFixture + ) -> None: + stream = _TrialAudioStream(close_error=RuntimeError("cleanup failed") if cleanup_fails else None) + harness.runtime.output = AudioOutput(data=stream, mime_type="audio/mpeg") + failure = RuntimeError("usage database unavailable") + + def reject_usage_insert( + _mapper: Mapper[AccountTrialAppRecord], _connection: Connection, _record: AccountTrialAppRecord + ) -> None: + raise failure + + event.listen(AccountTrialAppRecord, "before_insert", reject_usage_insert) + try: + response = harness.post(_TTS) + finally: + event.remove(AccountTrialAppRecord, "before_insert", reject_usage_insert) + + _trial_audio_assert_error(response, 500, "internal_server_error") + assert stream.reads == 1 + assert stream.close_calls == 1 + assert harness.usage() is None + assert any(record.exc_info is not None and record.exc_info[1] is failure for record in caplog.records) + + @staticmethod + def test_partial_http_consumption_closes_provider_stream(harness: _TrialAudioHarness) -> None: + stream = _TrialAudioStream() + harness.runtime.output = AudioOutput(data=stream, mime_type="audio/mpeg") + response = harness.post(_TTS) + + assert response.status_code == 200 + assert response.headers["Content-Type"] == "audio/mpeg" + assert "Content-Length" not in response.headers + assert next(iter(response.response)) == b"a" * 32 + assert stream.reads == 1 + assert stream.close_calls == 0 + response.close() + assert stream.close_calls == 1 + assert harness.usage() == 1 + + +"""Trial uploads through HTTP, real file services, and SQLite admission/persistence.""" + + +_TrialFileUploadEndpoint = Literal["files/upload", "remote-files/upload"] +_TRIAL_FILE_UPLOAD_CREATED_AT = datetime(2024, 1, 1) +_REMOTE_URL = "https://example.com/trial.txt" +_CONTENT = b"trial upload" + + +@dataclass +class _TrialFileUploadFeatures: + enabled: bool = True + events: list[str] = field(default_factory=list) + + def is_trial_enabled(self) -> bool: + self.events.append("feature") + return self.enabled + + +@dataclass +class _TrialFileUploadExternalio: + admission_sessions: list[Session] + saved: dict[str, bytes] = field(default_factory=dict) + requests: list[tuple[str, str, Mapping[str, object]]] = field(default_factory=list) + head_status: int = 200 + get_status: int = 200 + content_length: str = str(len(_CONTENT)) + error: Exception | None = None + + def assert_admission_closed(self) -> None: + assert self.admission_sessions + assert all(not session.in_transaction() and not session.identity_map for session in self.admission_sessions) + + def save(self, key: str, content: bytes) -> None: + self.assert_admission_closed() + self.saved[key] = content + + def request(self, method: str, *, url: str, **kwargs: object) -> httpx.Response: + self.assert_admission_closed() + self.requests.append((method, url, kwargs)) + if self.error is not None: + raise self.error + return httpx.Response( + self.head_status if method == "HEAD" else self.get_status, + headers={"Content-Type": "text/plain", "Content-Length": self.content_length}, + content=_CONTENT if method == "GET" else b"", + request=httpx.Request(method, url), + ) + + +@dataclass(frozen=True) +class _TrialFileUploadTrialappservices: + access: TrialAppAccessService + + +@dataclass(frozen=True) +class _TrialFileUploadApplicationservices: + trial_apps: _TrialFileUploadTrialappservices + recommended_app_queries: _TrialFileUploadFeatures + files: FileService + remote_files: RemoteFileService + + +@dataclass(frozen=True) +class _TrialFileUploadHarness: + app: Flask + account: Account + target: _trial_file_upload_App + trial: TrialApp + factory: sessionmaker[Session] + features: _TrialFileUploadFeatures + io: _TrialFileUploadExternalio + + def url(self, endpoint: _TrialFileUploadEndpoint) -> str: + return f"/trial-apps/{self.target.id}/{endpoint}" + + def post(self, endpoint: _TrialFileUploadEndpoint, *, query: str = "") -> TestResponse: + if endpoint == "remote-files/upload": + return self.app.test_client().post(self.url(endpoint) + query, json={"url": _REMOTE_URL}) + return self.app.test_client().post( + self.url(endpoint) + query, + data={"file": (BytesIO(_CONTENT), "trial.txt", "text/plain")}, + ) + + def uploaded(self) -> list[_trial_file_upload_UploadFile]: + with self.factory() as session: + return list(session.scalars(select(_trial_file_upload_UploadFile))) + + def usage(self) -> int | None: + with self.factory() as session: + return session.scalar( + select(AccountTrialAppRecord.count).where( + AccountTrialAppRecord.app_id == self.target.id, + AccountTrialAppRecord.account_id == self.account.id, + ) + ) + + +def _trial_file_upload_assert_error(response: TestResponse, status: int, code: str) -> None: + assert response.status_code == status, response.get_json() + body = response.get_json() + assert isinstance(body, dict) + assert body["code"] == code + assert body["status"] == status + assert isinstance(body["message"], str) + assert body["message"] + assert response.headers["Content-Type"] == "application/json" + + +def _assert_signed_url(url: object, file_id: str) -> None: + assert isinstance(url, str) + parsed = urlsplit(url) + assert parsed.path == f"/files/{file_id}/file-preview" + assert parse_qs(parsed.query).keys() == {"timestamp", "nonce", "sign"} + + +class TestTrialFileUploadCases: + @staticmethod + @pytest.fixture + def harness( + monkeypatch: pytest.MonkeyPatch, + config_overrides: Callable[..., None], + sqlite_session_factory: sessionmaker[Session], + ) -> _TrialFileUploadHarness: + config_overrides( + DEPLOYMENT_EDITION=DeploymentEdition.COMMUNITY, + INIT_PASSWORD="", + LOGIN_DISABLED=False, + RBAC_ENABLED=False, + UPLOAD_FILE_SIZE_LIMIT=1, + inner_UPLOAD_FILE_EXTENSION_BLACKLIST="exe", + ) + features = _TrialFileUploadFeatures() + account = Account(name="Trial viewer", email="trial@example.com") + account._current_tenant = Tenant(name="Viewer workspace") + account.role = TenantAccountRole.OWNER + target = _trial_file_upload_App( + tenant_id=str(uuid4()), + name="Trial app", + mode=_trial_file_upload_AppMode.CHAT, + enable_site=True, + enable_api=False, + ) + target.id = str(uuid4()) + trial = TrialApp(app_id=target.id, tenant_id=str(uuid4()), trial_limit=3) + with sqlite_session_factory.begin() as session: + session.add_all([target, trial]) + + sessions: list[Session] = [] + repository_factory = sessionmaker(bind=sqlite_session_factory.kw["bind"], expire_on_commit=False) + + @event.listens_for(repository_factory, "after_begin") + def track_session(session: Session, _transaction: SessionTransaction, _connection: Connection) -> None: + sessions.append(session) + + io = _TrialFileUploadExternalio(sessions) + files = FileService(session_factory=sqlite_session_factory) + services = _TrialFileUploadApplicationservices( + trial_apps=_TrialFileUploadTrialappservices( + access=TrialAppAccessService(apps=TrialAppRepository(session_factory=repository_factory)) + ), + recommended_app_queries=features, + files=files, + remote_files=RemoteFileService(files=files), + ) + + def setup_completed() -> bool: + features.events.append("setup") + return True + + def csrf(_request: Request, account_id: str) -> None: + assert account_id == account.id + features.events.append("csrf") + + for module in (trial_module, _trial_file_upload_admission_module, files_module, remote_module): + monkeypatch.setattr(module, "application_services", lambda: services) + monkeypatch.setattr(console_wraps, "_is_setup_completed", setup_completed) + monkeypatch.setattr(login_module, "current_user", account) + monkeypatch.setattr(login_module, "check_csrf_token", csrf) + monkeypatch.setattr(file_module.storage, "save", io.save) + monkeypatch.setattr(file_module, "naive_utc_now", lambda: _TRIAL_FILE_UPLOAD_CREATED_AT) + monkeypatch.setattr(remote_fetcher, "make_request", io.request) + + app = Flask(__name__) + app.config.update(TESTING=True, RESTX_ERROR_404_HELP=False) + login_manager = DifyLoginManager() + login_manager.init_app(app) + login_manager.unauthorized_handler(unauthorized_handler) + api = ExternalApi(app) + api.add_resource(trial_module.TrialAppFileUploadApi, "/trial-apps//files/upload") + api.add_resource(trial_module.TrialAppRemoteFileUploadApi, "/trial-apps//remote-files/upload") + return _TrialFileUploadHarness(app, account, target, trial, sqlite_session_factory, features, io) + + @staticmethod + @pytest.mark.parametrize("content", [_CONTENT, b""]) + def test_local_upload_preserves_complete_response_and_owner_tenant( + harness: _TrialFileUploadHarness, content: bytes + ) -> None: + response = harness.app.test_client().post( + harness.url("files/upload"), data={"file": (BytesIO(content), "trial.txt", "text/plain")} + ) + + assert response.status_code == 201 + assert response.headers["Content-Type"] == "application/json" + assert int(response.headers["Content-Length"]) == len(response.data) + body = response.get_json() + uploaded = harness.uploaded() + assert len(uploaded) == 1 + file = uploaded[0] + assert body == { + "id": file.id, + "reference": None, + "name": "trial.txt", + "size": len(content), + "extension": "txt", + "mime_type": "text/plain", + "created_by": harness.account.id, + "created_at": int(_TRIAL_FILE_UPLOAD_CREATED_AT.timestamp()), + "preview_url": None, + "source_url": body["source_url"], + "original_url": None, + "user_id": None, + "tenant_id": harness.target.tenant_id, + "conversation_id": None, + "file_key": None, + } + assert str(UUID(file.id)) == file.id + _assert_signed_url(body["source_url"], file.id) + assert file.created_by_role == CreatorUserRole.ACCOUNT + assert file.created_by == harness.account.id + assert file.tenant_id == harness.target.tenant_id + assert file.tenant_id not in {harness.trial.tenant_id, harness.account.current_tenant_id} + assert file.key.startswith(f"upload_files/{file.tenant_id}/") + assert file.hash == hashlib.sha3_256(content).hexdigest() + assert file.used is False + assert harness.io.saved == {file.key: content} + assert harness.features.events == ["setup", "csrf", "feature"] + assert harness.usage() is None + + @staticmethod + @pytest.mark.parametrize("head_status", [200, 405]) + def test_remote_upload_fetches_once_and_preserves_response_and_actor( + harness: _TrialFileUploadHarness, head_status: int + ) -> None: + harness.io.head_status = head_status + response = harness.post("remote-files/upload") + + assert response.status_code == 201 + assert response.headers["Content-Type"] == "application/json" + assert int(response.headers["Content-Length"]) == len(response.data) + body = response.get_json() + uploaded = harness.uploaded() + assert len(uploaded) == 1 + file = uploaded[0] + assert body == { + "id": file.id, + "name": "trial.txt", + "size": len(_CONTENT), + "extension": "txt", + "url": body["url"], + "mime_type": "text/plain", + "created_by": harness.account.id, + "created_at": int(_TRIAL_FILE_UPLOAD_CREATED_AT.timestamp()), + } + _assert_signed_url(body["url"], file.id) + assert file.source_url == _REMOTE_URL + assert file.tenant_id == harness.target.tenant_id + assert file.tenant_id not in {harness.trial.tenant_id, harness.account.current_tenant_id} + assert file.created_by == harness.account.id + assert file.created_by_role == CreatorUserRole.ACCOUNT + assert harness.io.saved == {file.key: _CONTENT} + assert harness.io.requests == [ + ("HEAD", _REMOTE_URL, {}), + ("GET", _REMOTE_URL, {} if head_status == 200 else {"timeout": 3, "follow_redirects": True}), + ] + assert harness.usage() is None + + @staticmethod + @pytest.mark.parametrize("endpoint", ["files/upload", "remote-files/upload"]) + @pytest.mark.parametrize( + ("denial", "code"), + [ + ("feature", "trial_app_feature_disabled"), + ("missing_trial", "trial_app_not_allowed"), + ("missing_app", "trial_app_not_allowed"), + ("quota", "trial_app_limit_exceeded"), + ], + ) + def test_trial_admission_rejects_before_parsing_or_external_io( + harness: _TrialFileUploadHarness, endpoint: _TrialFileUploadEndpoint, denial: str, code: str + ) -> None: + if denial == "feature": + harness.features.enabled = False + else: + with harness.factory.begin() as session: + if denial == "missing_trial": + session.execute(delete(TrialApp).where(TrialApp.app_id == harness.target.id)) + elif denial == "missing_app": + session.execute( + delete(_trial_file_upload_App).where(_trial_file_upload_App.id == harness.target.id) + ) + else: + session.add(AccountTrialAppRecord(app_id=harness.target.id, account_id=harness.account.id, count=3)) + + response = harness.app.test_client().post(harness.url(endpoint), json={}) + + _trial_file_upload_assert_error(response, 403, code) + assert harness.io.saved == {} + assert harness.io.requests == [] + assert harness.uploaded() == [] + assert harness.usage() == (3 if denial == "quota" else None) + if denial == "feature": + assert harness.io.admission_sessions == [] + else: + harness.io.assert_admission_closed() + + @staticmethod + @pytest.mark.parametrize("endpoint", ["files/upload", "remote-files/upload"]) + def test_upload_does_not_charge_usage_or_inherit_another_accounts_exhausted_quota( + harness: _TrialFileUploadHarness, endpoint: _TrialFileUploadEndpoint + ) -> None: + other_account_id = str(uuid4()) + with harness.factory.begin() as session: + session.add_all( + [ + AccountTrialAppRecord(app_id=harness.target.id, account_id=harness.account.id, count=1), + AccountTrialAppRecord(app_id=harness.target.id, account_id=other_account_id, count=3), + ] + ) + + response = harness.post(endpoint) + + assert response.status_code == 201 + assert harness.usage() == 1 + with harness.factory() as session: + assert ( + session.scalar( + select(AccountTrialAppRecord.count).where(AccountTrialAppRecord.account_id == other_account_id) + ) + == 3 + ) + + @staticmethod + @pytest.mark.parametrize( + ("shape", "code"), + [("missing", "no_file_uploaded"), ("multiple", "too_many_files"), ("empty_name", "filename_not_exists_error")], + ) + def test_multipart_shape_errors_precede_dataset_permissions( + harness: _TrialFileUploadHarness, shape: str, code: str + ) -> None: + harness.account.role = TenantAccountRole.NORMAL + data = {} + if shape != "missing": + data["file"] = (BytesIO(_CONTENT), "" if shape == "empty_name" else "trial.txt", "text/plain") + if shape == "multiple": + data["second_file"] = (BytesIO(_CONTENT), "second.txt", "text/plain") + + response = harness.app.test_client().post(harness.url("files/upload") + "?source=datasets", data=data) + + _trial_file_upload_assert_error(response, 400, code) + assert harness.io.saved == {} + assert harness.uploaded() == [] + + @staticmethod + @pytest.mark.parametrize( + ("query_source", "form_source", "status"), + [("datasets", "unknown", 403), ("unknown", "datasets", 201), ("", "datasets", 403)], + ) + def test_query_source_takes_precedence_over_form_before_active_workspace_permission( + harness: _TrialFileUploadHarness, query_source: str, form_source: str, status: int + ) -> None: + harness.account.role = TenantAccountRole.NORMAL + response = harness.app.test_client().post( + harness.url("files/upload") + f"?source={query_source}", + data={"file": (BytesIO(_CONTENT), "trial.txt", "text/plain"), "source": form_source}, + ) + + assert response.status_code == status + if status == 403: + _trial_file_upload_assert_error(response, status, "forbidden") + assert harness.io.saved == {} + else: + assert len(harness.uploaded()) == 1 + + @staticmethod + @pytest.mark.parametrize( + ("role", "rbac", "status"), + [ + (TenantAccountRole.OWNER, False, 201), + (TenantAccountRole.DATASET_OPERATOR, False, 201), + (TenantAccountRole.NORMAL, False, 403), + (TenantAccountRole.NORMAL, True, 201), + ], + ) + def test_dataset_upload_uses_active_workspace_role_and_retains_rbac_bypass( + harness: _TrialFileUploadHarness, + config_overrides: Callable[..., None], + role: TenantAccountRole, + rbac: bool, + status: int, + ) -> None: + harness.account.role = role + config_overrides(RBAC_ENABLED=rbac) + + response = harness.post("files/upload", query="?source=datasets") + + assert response.status_code == status + assert len(harness.uploaded()) == (1 if status == 201 else 0) + + @staticmethod + @pytest.mark.parametrize("endpoint", ["files/upload", "remote-files/upload"]) + def test_cloud_document_quota_uses_active_workspace_before_upload( + harness: _TrialFileUploadHarness, + monkeypatch: pytest.MonkeyPatch, + config_overrides: Callable[..., None], + endpoint: _TrialFileUploadEndpoint, + ) -> None: + config_overrides(DEPLOYMENT_EDITION=DeploymentEdition.CLOUD) + checked_tenants: list[str] = [] + + def get_features(tenant_id: str, exclude_vector_space: bool = False) -> FeatureModel: + checked_tenants.append(tenant_id) + assert exclude_vector_space is True + features = FeatureModel() + features.documents_upload_quota.limit = 1 + features.documents_upload_quota.size = 1 + return features + + monkeypatch.setattr(FeatureService, "get_features", get_features) + + response = harness.post(endpoint, query="?source=datasets") + + _trial_file_upload_assert_error(response, 403, "forbidden") + assert checked_tenants == [harness.account.current_tenant_id] + assert harness.io.saved == {} + assert harness.io.requests == [] + assert harness.usage() is None + + @staticmethod + def test_dataset_file_size_limit_comes_from_app_owner_tenant( + harness: _TrialFileUploadHarness, monkeypatch: pytest.MonkeyPatch, config_overrides: Callable[..., None] + ) -> None: + config_overrides(DEPLOYMENT_EDITION=DeploymentEdition.CLOUD) + quota_tenants: list[str] = [] + size_limit_tenants: list[str | None] = [] + + def get_features(tenant_id: str, exclude_vector_space: bool = False) -> FeatureModel: + quota_tenants.append(tenant_id) + assert exclude_vector_space is True + return FeatureModel() + + def get_size_limit(tenant_id: str | None) -> int: + size_limit_tenants.append(tenant_id) + return 0 + + monkeypatch.setattr(FeatureService, "get_features", get_features) + monkeypatch.setattr(FeatureService, "get_knowledge_file_size_limit", get_size_limit) + + response = harness.post("files/upload", query="?source=datasets") + + _trial_file_upload_assert_error(response, 413, "file_too_large") + assert response.get_json()["message"] == "File size exceeded." + assert quota_tenants == [harness.account.current_tenant_id] + assert size_limit_tenants == [harness.target.tenant_id] + assert harness.io.saved == {} + assert harness.uploaded() == [] + + @staticmethod + @pytest.mark.parametrize( + ("filename", "query", "status", "code"), + [ + ("trial.exe", "", 400, "file_extension_blocked"), + ("trial.png", "?source=datasets", 415, "unsupported_file_type"), + ], + ) + def test_local_file_policy_errors_keep_specific_codes( + harness: _TrialFileUploadHarness, filename: str, query: str, status: int, code: str + ) -> None: + response = harness.app.test_client().post( + harness.url("files/upload") + query, + data={"file": (BytesIO(_CONTENT), filename, "application/octet-stream")}, + ) + + _trial_file_upload_assert_error(response, status, code) + assert harness.io.saved == {} + assert harness.uploaded() == [] + + @staticmethod + @pytest.mark.parametrize( + ("get_status", "status", "code"), + [ + (404, 404, "remote_file_not_found"), + (403, 400, "remote_file_access_denied"), + (503, 502, "remote_file_unavailable"), + ], + ) + def test_remote_http_failures_keep_precise_error_codes( + harness: _TrialFileUploadHarness, get_status: int, status: int, code: str + ) -> None: + harness.io.head_status = 405 + harness.io.get_status = get_status + + response = harness.post("remote-files/upload") + + _trial_file_upload_assert_error(response, status, code) + assert [method for method, _url, _kwargs in harness.io.requests] == ["HEAD", "GET"] + assert harness.io.saved == {} + assert harness.uploaded() == [] + assert harness.usage() is None + + @staticmethod + @pytest.mark.parametrize( + ("failure", "status", "code"), + [ + ("blocked", 400, "remote_file_url_blocked"), + ("timeout", 502, "remote_file_unavailable"), + ("metadata", 502, "remote_file_invalid_response"), + ("size", 413, "file_too_large"), + ], + ) + def test_remote_transport_metadata_and_size_failures_prevent_storage( + harness: _TrialFileUploadHarness, failure: str, status: int, code: str + ) -> None: + if failure == "blocked": + harness.io.error = ToolSSRFError("Blocked by SSRF proxy") + elif failure == "timeout": + harness.io.error = httpx.ReadTimeout("Timed out") + elif failure == "metadata": + harness.io.content_length = "invalid" + else: + harness.io.content_length = str(1024 * 1024 + 1) + + response = harness.post("remote-files/upload") + + _trial_file_upload_assert_error(response, status, code) + assert len(harness.io.requests) == 1 + assert harness.io.saved == {} + assert harness.uploaded() == [] + assert harness.usage() is None + + @staticmethod + def test_remote_invalid_url_fails_before_network(harness: _TrialFileUploadHarness) -> None: + response = harness.app.test_client().post( + harness.url("remote-files/upload"), json={"url": "file:///etc/passwd"} + ) + + _trial_file_upload_assert_error(response, 400, "remote_file_invalid_url") + assert harness.io.requests == [] + assert harness.io.saved == {} + + @staticmethod + def test_remote_payload_validation_stays_on_http_boundary(harness: _TrialFileUploadHarness) -> None: + response = harness.app.test_client().post(harness.url("remote-files/upload"), json={}) + + assert response.status_code == 422 + assert "url" in response.get_json()["message"] + assert harness.io.requests == [] + assert harness.io.saved == {} + + +"""Trial admission and generation through real HTTP handlers and SQLite repositories.""" + + +@dataclass +class _TrialGenerationFeaturestate: + enabled: bool = True + setup_completed: bool = True + events: list[str] = field(default_factory=list) + + def is_trial_enabled(self) -> bool: + self.events.append("feature") + return self.enabled + + +@dataclass(frozen=True) +class _GenerationCall: + app: TrialAppRef + account_id: str + args: Mapping[str, object] + streaming: bool + + +@dataclass +class _TrialGenerationStream(Iterator[str]): + chunks: Iterator[str] = field(default_factory=lambda: iter(['data: {"answer":"first"}\n\n', "data: [DONE]\n\n"])) + reads: int = 0 + close_calls: int = 0 + close_error: Exception | None = None + on_read: Callable[[], None] | None = None + + @override + def __next__(self) -> str: + if self.on_read is not None: + self.on_read() + chunk = next(self.chunks) + self.reads += 1 + return chunk + + def close(self) -> None: + self.close_calls += 1 + if self.close_error is not None: + raise self.close_error + + +@dataclass +class _TrialGenerationRuntime: + sessions: list[Session] + response: GenerationResponse = field( + default_factory=lambda: {"answer": "", "metadata": dict[str, object](), "created_at": 0} + ) + error: Exception | None = None + calls: list[_GenerationCall] = field(default_factory=list) + before_generate: Callable[[], None] | None = None + + def generate( + self, *, app: TrialAppRef, account_id: str, args: Mapping[str, object], streaming: bool + ) -> GenerationResponse: + # Database admission must release its connection before model invocation. + assert self.sessions + assert all(not session.in_transaction() and not session.identity_map for session in self.sessions) + self.calls.append(_GenerationCall(app, account_id, dict(args), streaming)) + if self.before_generate is not None: + self.before_generate() + if self.error is not None: + raise self.error + return self.response + + +@dataclass +class _Tasks: + calls: list[str] = field(default_factory=list) + chat_calls: list[tuple[str, str]] = field(default_factory=list) + + def stop_task(self, *, task_id: str, invoke_from: InvokeFrom, user_id: str) -> None: + assert invoke_from == InvokeFrom.EXPLORE + self.chat_calls.append((task_id, user_id)) + + def stop_workflow_task( + self, + *, + task_id: str, + tenant_id: str, + app_id: str, + app_mode: _trial_generation_AppMode, + owner: tuple[CreatorUserRole, str] | None, + ) -> None: + assert tenant_id + assert app_id + if app_mode == _trial_generation_AppMode.ADVANCED_CHAT: + assert owner is not None + assert owner[0] == CreatorUserRole.ACCOUNT + self.chat_calls.append((task_id, owner[1])) + else: + assert owner is None + assert app_mode == _trial_generation_AppMode.WORKFLOW + self.calls.append(task_id) + + +@dataclass(frozen=True) +class _TrialGenerationTrialappservices: + access: TrialAppAccessService + generation: TrialAppGenerationService + + +@dataclass(frozen=True) +class _TrialGenerationApplicationservices: + trial_apps: _TrialGenerationTrialappservices + recommended_app_queries: _TrialGenerationFeaturestate + app_tasks: _Tasks + + +@dataclass(frozen=True) +class _TrialGenerationHarness: + app: Flask + account: Account + target: _trial_generation_App + trial: TrialApp + factory: sessionmaker[Session] + services: _TrialGenerationApplicationservices + runtime: _TrialGenerationRuntime + state: _TrialGenerationFeaturestate + + def url(self, endpoint: str, *, app_id: str | None = None) -> str: + return f"/trial-apps/{app_id or self.target.id}/{endpoint}" + + def post(self, endpoint: str, *, payload: Mapping[str, object] | None = None) -> TestResponse: + return self.app.test_client().post(self.url(endpoint), json=dict(payload or {"inputs": {}})) + + def set_mode(self, mode: _trial_generation_AppMode) -> None: + with self.factory.begin() as session: + app = session.get(_trial_generation_App, self.target.id) + assert app is not None + app.mode = mode + + def usage(self) -> int | None: + with self.factory() as session: + return session.scalar( + select(AccountTrialAppRecord.count).where( + AccountTrialAppRecord.app_id == self.target.id, + AccountTrialAppRecord.account_id == self.account.id, + ) + ) + + +def _trial_generation_assert_error(response: TestResponse, status: int, code: str) -> Mapping[str, object]: + assert response.status_code == status + body = response.get_json() + assert isinstance(body, dict) + assert body["code"] == code + assert body["status"] == status + assert response.headers["Content-Type"] == "application/json" + assert isinstance(body["message"], str) + assert body["message"] + return body + + +class TestTrialGenerationCases: + @staticmethod + @pytest.fixture + def harness( + monkeypatch: pytest.MonkeyPatch, + config_overrides: Callable[..., None], + sqlite_session_factory: sessionmaker[Session], + ) -> _TrialGenerationHarness: + config_overrides(DEPLOYMENT_EDITION=DeploymentEdition.COMMUNITY, INIT_PASSWORD="", LOGIN_DISABLED=False) + state = _TrialGenerationFeaturestate() + account = Account(name="Viewer", email="viewer@example.com") + account._current_tenant = Tenant(name="Viewer workspace") + target = _trial_generation_App( + tenant_id=str(uuid4()), + name="Trial app", + mode=_trial_generation_AppMode.COMPLETION, + enable_site=True, + enable_api=True, + ) + target.id = str(uuid4()) + trial = TrialApp(app_id=target.id, tenant_id=target.tenant_id, trial_limit=3) + with sqlite_session_factory.begin() as session: + session.add_all([target, trial]) + + sessions: list[Session] = [] + repository_factory = sessionmaker(bind=sqlite_session_factory.kw["bind"], expire_on_commit=False) + + @event.listens_for(repository_factory, "after_begin") + def track_session(session: Session, _transaction: SessionTransaction, _connection: Connection) -> None: + sessions.append(session) + + runtime = _TrialGenerationRuntime(sessions=sessions) + repository = TrialAppRepository(session_factory=repository_factory) + services = _TrialGenerationApplicationservices( + trial_apps=_TrialGenerationTrialappservices( + access=TrialAppAccessService(apps=repository), + generation=TrialAppGenerationService(runtime=runtime, usage=repository), + ), + recommended_app_queries=state, + app_tasks=_Tasks(), + ) + + def setup_completed() -> bool: + state.events.append("setup") + return state.setup_completed + + def csrf(_request: Request, account_id: str) -> None: + assert account_id == account.id + state.events.append("csrf") + + for module in (trial_module, _trial_generation_admission_module): + monkeypatch.setattr(module, "application_services", lambda: services) + monkeypatch.setattr(console_wraps, "_is_setup_completed", setup_completed) + monkeypatch.setattr(login_module, "current_user", account) + monkeypatch.setattr(login_module, "check_csrf_token", csrf) + + app = Flask(__name__) + app.config.update(TESTING=True, RESTX_ERROR_404_HELP=False) + login_manager = DifyLoginManager() + login_manager.init_app(app) + login_manager.unauthorized_handler(unauthorized_handler) + api = ExternalApi(app) + for resource, endpoint in ( + (trial_module.TrialAppWorkflowRunApi, "workflows/run"), + (trial_module.TrialChatApi, "chat-messages"), + (trial_module.TrialChatTaskStopApi, "chat-messages//stop"), + (trial_module.TrialCompletionApi, "completion-messages"), + (trial_module.TrialAppWorkflowTaskStopApi, "workflows/tasks//stop"), + ): + api.add_resource(resource, f"/trial-apps//{endpoint}") + return _TrialGenerationHarness(app, account, target, trial, sqlite_session_factory, services, runtime, state) + + @staticmethod + @pytest.mark.parametrize("response_mode", [None, "blocking"]) + def test_completion_blocking_preserves_response_and_uses_app_owner( + harness: _TrialGenerationHarness, response_mode: str | None + ) -> None: + response = harness.post("completion-messages", payload={"inputs": {}, "response_mode": response_mode}) + + assert response.status_code == 200 + assert response.get_json() == {"answer": "", "metadata": {}, "created_at": 0} + assert response.headers["Content-Type"] == "application/json; charset=utf-8" + assert int(response.headers["Content-Length"]) == len(response.data) + assert harness.usage() == 1 + assert harness.state.events == ["setup", "csrf", "feature"] + call = harness.runtime.calls[0] + assert call.app == TrialAppRef( + app_id=harness.target.id, tenant_id=harness.target.tenant_id, app_mode=_trial_generation_AppMode.COMPLETION + ) + assert call.app.tenant_id != harness.account.current_tenant_id + assert call.account_id == harness.account.id + assert call.streaming is False + assert call.args == { + "inputs": {}, + "query": "", + "files": None, + "response_mode": response_mode, + "retriever_from": "explore_app", + "auto_generate_name": False, + } + + @staticmethod + @pytest.mark.parametrize( + ("endpoint", "mode", "payload"), + [ + ("workflows/run", _trial_generation_AppMode.WORKFLOW, {"inputs": {"number": 0}, "files": []}), + ("chat-messages", _trial_generation_AppMode.CHAT, {"inputs": {}, "query": "hello"}), + ("chat-messages", _trial_generation_AppMode.AGENT, {"inputs": {}, "query": "hello"}), + ("chat-messages", _trial_generation_AppMode.ADVANCED_CHAT, {"inputs": {}, "query": "hello"}), + ("chat-messages", _trial_generation_AppMode.AGENT_CHAT, {"inputs": {}, "query": "hello"}), + ("completion-messages", _trial_generation_AppMode.COMPLETION, {"inputs": {}, "response_mode": "streaming"}), + ], + ) + def test_stream_records_usage_after_runtime_returns_before_consumption( + harness: _TrialGenerationHarness, endpoint: str, mode: _trial_generation_AppMode, payload: Mapping[str, object] + ) -> None: + harness.set_mode(mode) + + def before_generate() -> None: + assert harness.usage() is None + + def before_read() -> None: + assert harness.usage() == 1 + + stream = _TrialGenerationStream(on_read=before_read) + harness.runtime.response = stream + harness.runtime.before_generate = before_generate + response = harness.post(endpoint, payload=payload) + + assert response.status_code == 200 + assert response.headers["Content-Type"] == "text/event-stream; charset=utf-8" + assert "Content-Length" not in response.headers + assert response.data == b'data: {"answer":"first"}\n\ndata: [DONE]\n\n' + response.close() + assert harness.usage() == 1 + assert harness.runtime.calls[0].streaming is True + if endpoint == "chat-messages": + assert harness.runtime.calls[0].args["auto_generate_name"] is False + elif endpoint == "workflows/run": + assert harness.runtime.calls[0].args == payload + + @staticmethod + @pytest.mark.parametrize( + "endpoint", + [ + "completion-messages", + "chat-messages", + "chat-messages/task/stop", + "workflows/run", + "workflows/tasks/task/stop", + ], + ) + @pytest.mark.parametrize( + ("state", "status", "code"), [("setup", 401, "not_setup"), ("feature", 403, "trial_app_feature_disabled")] + ) + def test_bootstrap_and_feature_admission_apply_before_trial_lookup( + harness: _TrialGenerationHarness, endpoint: str, state: str, status: int, code: str + ) -> None: + if state == "setup": + harness.state.setup_completed = False + else: + harness.state.enabled = False + response = harness.app.test_client().post(harness.url(endpoint, app_id=str(uuid4())), json={}) + + _trial_generation_assert_error(response, status, code) + assert harness.state.events == (["setup"] if state == "setup" else ["setup", "csrf", "feature"]) + assert harness.runtime.calls == [] + assert harness.services.app_tasks.calls == [] + + @staticmethod + @pytest.mark.parametrize( + ("case", "code"), + [ + ("missing_trial", "trial_app_not_allowed"), + ("missing_app", "trial_app_not_allowed"), + ("limit", "trial_app_limit_exceeded"), + ], + ) + def test_trial_lookup_and_account_limit_block_generation( + harness: _TrialGenerationHarness, case: str, code: str + ) -> None: + with harness.factory.begin() as session: + if case == "missing_trial": + trial = session.get(TrialApp, harness.trial.id) + assert trial is not None + session.delete(trial) + elif case == "missing_app": + target = session.get(_trial_generation_App, harness.target.id) + assert target is not None + session.delete(target) + else: + session.add(AccountTrialAppRecord(app_id=harness.target.id, account_id=harness.account.id, count=3)) + response = harness.post("completion-messages") + + _trial_generation_assert_error(response, 403, code) + assert harness.runtime.calls == [] + + @staticmethod + @pytest.mark.parametrize("count", [0, 2]) + def test_usage_increments_existing_record_and_ignores_other_account( + harness: _TrialGenerationHarness, count: int + ) -> None: + other_id = str(uuid4()) + with harness.factory.begin() as session: + session.add_all( + [ + AccountTrialAppRecord(app_id=harness.target.id, account_id=harness.account.id, count=count), + AccountTrialAppRecord(app_id=harness.target.id, account_id=other_id, count=99), + ] + ) + response = harness.post("completion-messages") + + assert response.status_code == 200 + assert harness.usage() == count + 1 + with harness.factory() as session: + assert ( + session.scalar(select(AccountTrialAppRecord.count).where(AccountTrialAppRecord.account_id == other_id)) + == 99 + ) + + @staticmethod + @pytest.mark.parametrize( + ("endpoint", "mode", "code"), + [ + ("completion-messages", _trial_generation_AppMode.CHAT, "not_completion_app"), + ("chat-messages", _trial_generation_AppMode.COMPLETION, "not_chat_app"), + ("workflows/run", _trial_generation_AppMode.CHAT, "not_workflow_app"), + ("workflows/tasks/task/stop", _trial_generation_AppMode.CHAT, "not_workflow_app"), + ], + ) + def test_wrong_mode_returns_specific_error_without_invocation( + harness: _TrialGenerationHarness, endpoint: str, mode: _trial_generation_AppMode, code: str + ) -> None: + harness.set_mode(mode) + response = harness.post(endpoint, payload={"inputs": {}, "query": "hello"}) + + _trial_generation_assert_error(response, 400, code) + assert harness.runtime.calls == [] + assert harness.services.app_tasks.calls == [] + assert harness.usage() is None + + @staticmethod + @pytest.mark.parametrize("field", ["conversation_id", "parent_message_id"]) + def test_chat_rejects_invalid_uuid_before_runtime(harness: _TrialGenerationHarness, field: str) -> None: + harness.set_mode(_trial_generation_AppMode.CHAT) + response = harness.post("chat-messages", payload={"inputs": {}, "query": "hello", field: "not-a-uuid"}) + + _trial_generation_assert_error(response, 400, "invalid_param") + assert harness.runtime.calls == [] + assert harness.usage() is None + + @staticmethod + def test_unpublished_agent_remains_unavailable_after_service_migration(harness: _TrialGenerationHarness) -> None: + harness.set_mode(_trial_generation_AppMode.AGENT) + harness.runtime.error = AgentAppNotPublishedError("Agent has not been published") + response = harness.post("chat-messages", payload={"inputs": {}, "query": "hello"}) + _trial_generation_assert_error(response, 400, "app_unavailable") + assert harness.usage() is None + + @staticmethod + def test_stop_chat_rejects_workflow(harness: _TrialGenerationHarness) -> None: + harness.set_mode(_trial_generation_AppMode.WORKFLOW) + _trial_generation_assert_error(harness.post("chat-messages/task-1/stop"), 400, "not_chat_app") + assert harness.services.app_tasks.chat_calls == [] + + @staticmethod + @pytest.mark.parametrize( + "mode", + [ + _trial_generation_AppMode.CHAT, + _trial_generation_AppMode.AGENT_CHAT, + _trial_generation_AppMode.AGENT, + _trial_generation_AppMode.ADVANCED_CHAT, + ], + ) + def test_last_trial_can_be_stopped_without_allowing_another_generation( + harness: _TrialGenerationHarness, mode: _trial_generation_AppMode + ) -> None: + harness.set_mode(mode) + with harness.factory.begin() as session: + session.add(AccountTrialAppRecord(app_id=harness.target.id, account_id=harness.account.id, count=2)) + + response = harness.post("chat-messages", payload={"inputs": {}, "query": "hello"}) + assert response.status_code == 200 + assert harness.usage() == 3 + + stopped = harness.post("chat-messages/task-last/stop") + assert stopped.status_code == 200 + assert harness.services.app_tasks.chat_calls == [("task-last", harness.account.id)] + assert harness.usage() == 3 + unavailable = harness.app.test_client().post(harness.url("chat-messages/task-last/stop", app_id=str(uuid4()))) + _trial_generation_assert_error(unavailable, 403, "trial_app_not_allowed") + assert len(harness.services.app_tasks.chat_calls) == 1 + _trial_generation_assert_error( + harness.post("chat-messages", payload={"inputs": {}, "query": "again"}), + 403, + "trial_app_limit_exceeded", + ) + + @staticmethod + def test_chat_normalizes_uuid_and_preserves_empty_ids(harness: _TrialGenerationHarness) -> None: + harness.set_mode(_trial_generation_AppMode.CHAT) + conversation_id = uuid4() + response = harness.post( + "chat-messages", + payload={ + "inputs": {}, + "query": "hello", + "conversation_id": conversation_id.hex.upper(), + "parent_message_id": "", + }, + ) + + assert response.status_code == 200 + assert harness.runtime.calls[0].args["conversation_id"] == str(conversation_id) + assert harness.runtime.calls[0].args["parent_message_id"] == "" + + @staticmethod + @pytest.mark.parametrize( + ("error", "status", "code"), + [ + (ProviderTokenNotInitError("missing key for provider"), 400, "provider_not_initialize"), + (QuotaExceededError(), 400, "provider_quota_exceeded"), + (ModelCurrentlyNotSupportError(), 400, "model_currently_not_support"), + (InvokeError("provider rejected request"), 400, "completion_request_error"), + (InvokeRateLimitError("too many requests"), 429, "rate_limit_error"), + (AppInvokeQuotaExceededError("active requests exceeded"), 429, "rate_limit_error"), + (AppDefinitionUnavailableError("app removed"), 400, "app_unavailable"), + (AccountNotFoundError("account removed"), 401, "unauthorized"), + (AppModelConfigBrokenError(), 400, "app_unavailable"), + ], + ) + @pytest.mark.parametrize( + ("endpoint", "mode"), + [ + ("completion-messages", _trial_generation_AppMode.COMPLETION), + ("chat-messages", _trial_generation_AppMode.CHAT), + ("workflows/run", _trial_generation_AppMode.WORKFLOW), + ], + ) + def test_runtime_errors_keep_specific_codes_and_do_not_record_usage( + harness: _TrialGenerationHarness, + endpoint: str, + mode: _trial_generation_AppMode, + error: Exception, + status: int, + code: str, + ) -> None: + harness.set_mode(mode) + harness.runtime.error = error + response = harness.post(endpoint, payload={"inputs": {}, "query": "hello"}) + + body = _trial_generation_assert_error(response, status, code) + if isinstance(error, (ProviderTokenNotInitError, InvokeError, InvokeRateLimitError)): + assert body["message"] == error.description + assert harness.usage() is None + + @staticmethod + @pytest.mark.parametrize( + ("error", "status", "code"), + [ + (ConversationNotExistsError(), 404, "not_found"), + # Nothing raises ConversationCompletedError today, so the contract does not list its code. + pytest.param( + ConversationCompletedError(), + 400, + "conversation_completed", + marks=pytest.mark.allow_error_contract_violations, + ), + (AppModelConfigBrokenError(), 400, "app_unavailable"), + ], + ) + @pytest.mark.parametrize( + ("endpoint", "mode"), + [ + ("completion-messages", _trial_generation_AppMode.COMPLETION), + ("chat-messages", _trial_generation_AppMode.CHAT), + ], + ) + def test_conversation_and_config_errors_remain_actionable( + harness: _TrialGenerationHarness, + endpoint: str, + mode: _trial_generation_AppMode, + error: Exception, + status: int, + code: str, + ) -> None: + harness.set_mode(mode) + harness.runtime.error = error + response = harness.post(endpoint, payload={"inputs": {}, "query": "hello"}) + + _trial_generation_assert_error(response, status, code) + assert harness.usage() is None + + @staticmethod + def test_stop_retains_trial_quota_admission_without_recording_usage(harness: _TrialGenerationHarness) -> None: + harness.set_mode(_trial_generation_AppMode.WORKFLOW) + response = harness.post("workflows/tasks/task-123/stop") + + assert response.status_code == 200 + assert response.get_json() == {"result": "success"} + assert harness.services.app_tasks.calls == ["task-123"] + assert harness.usage() is None + with harness.factory.begin() as session: + session.add(AccountTrialAppRecord(app_id=harness.target.id, account_id=harness.account.id, count=3)) + blocked = harness.post("workflows/tasks/task-456/stop") + _trial_generation_assert_error(blocked, 403, "trial_app_limit_exceeded") + assert harness.services.app_tasks.calls == ["task-123"] + + @staticmethod + def test_unauthenticated_account_cannot_enter_trial( + harness: _TrialGenerationHarness, monkeypatch: pytest.MonkeyPatch + ) -> None: + monkeypatch.setattr(login_module, "current_user", None) + response = harness.post("completion-messages") + + assert response.status_code == 401 + assert harness.state.events == ["setup"] + assert harness.runtime.calls == [] + + @staticmethod + def test_uninitialized_account_cannot_enter_trial(harness: _TrialGenerationHarness) -> None: + harness.account.status = AccountStatus.UNINITIALIZED + response = harness.post("completion-messages") + + _trial_generation_assert_error(response, 400, "account_not_initialized") + assert harness.state.events == ["setup", "csrf"] + assert harness.runtime.calls == [] + + @staticmethod + def test_failed_csrf_cannot_enter_trial(harness: _TrialGenerationHarness, monkeypatch: pytest.MonkeyPatch) -> None: + def reject_csrf(_request: Request, _account_id: str) -> None: + raise Unauthorized("invalid csrf token") + + monkeypatch.setattr(login_module, "check_csrf_token", reject_csrf) + response = harness.post("completion-messages") + + _trial_generation_assert_error(response, 401, "unauthorized") + assert harness.runtime.calls == [] + + @staticmethod + def test_payload_validation_keeps_field_context(harness: _TrialGenerationHarness) -> None: + response = harness.post("completion-messages", payload={"inputs": [], "response_mode": "invalid"}) + + body = _trial_generation_assert_error(response, 422, "unprocessable_entity") + assert "inputs" in str(body) + assert harness.runtime.calls == [] + assert harness.usage() is None + + @staticmethod + @pytest.mark.parametrize("cleanup_fails", [False, True]) + @pytest.mark.parametrize( + ("endpoint", "mode"), + [ + ("completion-messages", _trial_generation_AppMode.COMPLETION), + ("chat-messages", _trial_generation_AppMode.CHAT), + ("workflows/run", _trial_generation_AppMode.WORKFLOW), + ], + ) + def test_usage_write_failure_closes_unconsumed_stream_without_hiding_database_error( + harness: _TrialGenerationHarness, + endpoint: str, + mode: _trial_generation_AppMode, + cleanup_fails: bool, + caplog: pytest.LogCaptureFixture, + ) -> None: + harness.set_mode(mode) + stream = _TrialGenerationStream(close_error=RuntimeError("cleanup failed") if cleanup_fails else None) + harness.runtime.response = stream + failure = RuntimeError("usage database unavailable") + + def reject_usage_insert( + _mapper: Mapper[AccountTrialAppRecord], _connection: Connection, _record: AccountTrialAppRecord + ) -> None: + raise failure + + # Fail the persistence boundary, retaining the real service, unit of work, + # SQLAlchemy rollback and HTTP error handler. + event.listen(AccountTrialAppRecord, "before_insert", reject_usage_insert) + try: + response = harness.post(endpoint, payload={"inputs": {}, "query": "hello", "response_mode": "streaming"}) + finally: + event.remove(AccountTrialAppRecord, "before_insert", reject_usage_insert) + + _trial_generation_assert_error(response, 500, "internal_server_error") + assert stream.close_calls == 1 + assert stream.reads == 0 + assert harness.usage() is None + assert any(record.exc_info is not None and record.exc_info[1] is failure for record in caplog.records) + + @staticmethod + def test_later_stream_failure_does_not_refund_trial_usage(harness: _TrialGenerationHarness) -> None: + def chunks() -> Iterator[str]: + yield 'data: {"answer":"partial"}\n\n' + raise RuntimeError("provider disconnected") + + harness.runtime.response = _TrialGenerationStream(chunks=chunks()) + response = harness.post("completion-messages", payload={"inputs": {}, "response_mode": "streaming"}) + + chunks = iter(response.response) + assert next(chunks) == b'data: {"answer":"partial"}\n\n' + with pytest.raises(RuntimeError, match="provider disconnected"): + next(chunks) + response.close() + assert harness.usage() == 1 + + +"""Anonymous app previews through HTTP and real SQLite query boundaries.""" + + +_TrialPreviewEndpoint = Literal["parameters", "site", "datasets"] +_TRIAL_PREVIEW_CREATED_AT = datetime(2024, 1, 1) +_INPUT_FORM = [{"number": {"label": "Count", "variable": "count", "required": False, "default": 0}}] + + +@dataclass +class _TrialPreviewCatalog: + engine: Engine + ids: set[str] = field(default_factory=set) + calls: list[str] = field(default_factory=list) + + def contains(self, app_id: str) -> bool: + assert isinstance(self.engine.pool, QueuePool) + assert self.engine.pool.checkedout() == 0 + self.calls.append(app_id) + return app_id in self.ids + + def list_recommended(self, language: str) -> RecommendedAppCatalogPage: + raise AssertionError(f"Unexpected list request for {language}") + + def list_learn_dify(self, language: str) -> RecommendedAppCatalogPage: + raise AssertionError(f"Unexpected list request for {language}") + + def get_detail(self, app_id: str) -> RecommendedAppDetailRecord | None: + raise AssertionError(f"Unexpected detail request for {app_id}") + + +@dataclass(frozen=True) +class _TrialPreviewApplicationservices: + app_previews: AppPreviewQueryService + app_definitions: AppDefinitionQueryService + recommended_app_queries: RecommendedAppQueryService + + +@dataclass(frozen=True) +class _TrialPreviewHarness: + app: Flask + owner: Tenant + target: _trial_preview_App + listing: TrialApp + site: Site + config: AppModelConfig + dataset: Dataset + creator_id: str + factory: sessionmaker[Session] + engine: Engine + catalog: _TrialPreviewCatalog + sessions: list[Session] + signed_icons: list[str] + + def get(self, endpoint: _TrialPreviewEndpoint, *, query: str = "", app_id: str | None = None) -> TestResponse: + response = self.app.test_client().get(f"/trial-apps/{app_id or self.target.id}/{endpoint}{query}") + assert response.headers["Content-Type"] == "application/json" + assert int(response.headers["Content-Length"]) == len(response.data) + assert isinstance(self.engine.pool, QueuePool) + assert self.engine.pool.checkedout() == 0 + assert all(not session.in_transaction() and not session.identity_map for session in self.sessions) + return response + + def dataset_query(self) -> str: + return f"?ids={self.dataset.id}" + + +def _trial_preview_assert_error(response: TestResponse, status: int, code: str) -> None: + assert response.status_code == status, response.get_json() + body = response.get_json() + assert body["code"] == code + assert body["status"] == status + assert isinstance(body["message"], str) + assert body["message"] + + +class TestTrialPreviewCases: + @staticmethod + @pytest.fixture + def harness( + monkeypatch: pytest.MonkeyPatch, + config_overrides: Callable[..., None], + sqlite_engine: Engine, + sqlite_session_factory: sessionmaker[Session], + ) -> _TrialPreviewHarness: + config_overrides( + LOGIN_DISABLED=False, + INIT_PASSWORD="preview-does-not-require-setup", + UPLOAD_IMAGE_FILE_SIZE_LIMIT=1, + UPLOAD_VIDEO_FILE_SIZE_LIMIT=2, + UPLOAD_AUDIO_FILE_SIZE_LIMIT=3, + UPLOAD_FILE_SIZE_LIMIT=4, + WORKFLOW_FILE_UPLOAD_LIMIT=5, + ) + owner = Tenant(name="App owner") + creator_id = str(uuid4()) + target = _trial_preview_App( + id=str(uuid4()), + tenant_id=owner.id, + name="Preview", + mode=_trial_preview_AppMode.CHAT, + enable_site=False, + enable_api=False, + ) + listing = TrialApp(app_id=target.id, tenant_id=str(uuid4()), trial_limit=0) + config = AppModelConfig( + app_id=target.id, + opening_statement="", + suggested_questions="[]", + suggested_questions_after_answer='{"enabled":false}', + speech_to_text='{"enabled":false}', + text_to_speech='{"enabled":false,"voice":"","autoPlay":"disabled"}', + retriever_resource='{"enabled":false}', + more_like_this='{"enabled":false}', + sensitive_word_avoidance='{"enabled":false,"type":"","configs":[]}', + file_upload='{"enabled":false,"number_limits":0}', + user_input_form=json.dumps(_INPUT_FORM), + ) + site = Site( + app_id=target.id, + title="Preview site", + default_language="en-US", + customize_token_strategy=CustomizeTokenStrategy.UUID, + icon_type=IconType.IMAGE, + icon=str(uuid4()), + icon_background="", + chat_color_theme="", + chat_color_theme_inverted=False, + description=None, + copyright="", + privacy_policy=None, + input_placeholder="", + custom_disclaimer="", + show_workflow_steps=False, + use_icon_as_answer_icon=False, + prompt_public=False, + ) + dataset = Dataset( + tenant_id=owner.id, + name="Owner dataset", + description=None, + permission="only_me", + data_source_type="upload_file", + indexing_technique=None, + created_by=creator_id, + ) + dataset.created_at = _TRIAL_PREVIEW_CREATED_AT + with sqlite_session_factory.begin() as session: + session.add_all([owner, target, listing, config, site, dataset]) + session.flush() + target.app_model_config_id = config.id + session.add(AccountTrialAppRecord(app_id=target.id, account_id=creator_id, count=100)) + + factory = sessionmaker(bind=sqlite_engine, expire_on_commit=False) + sessions: list[Session] = [] + + @event.listens_for(factory, "after_begin") + def track_session(session: Session, _transaction: SessionTransaction, _connection: Connection) -> None: + assert all(not previous.in_transaction() for previous in sessions) + assert isinstance(sqlite_engine.pool, QueuePool) + assert sqlite_engine.pool.checkedout() == 1 + sessions.append(session) + + catalog = _TrialPreviewCatalog(sqlite_engine) + recommendations = RecommendedAppQueryService( + catalog=catalog, trial_apps=TrialAppRepository(factory), trial_enabled=False + ) + services = _TrialPreviewApplicationservices( + app_previews=AppPreviewQueryService( + apps=AppPreviewQueryRepository(session_factory=factory), is_previewable=recommendations.is_previewable + ), + app_definitions=AppDefinitionQueryService( + definitions=AppDefinitionQueryRepository(session_factory=factory), builtin_icon_url_prefix="/tools/" + ), + recommended_app_queries=recommendations, + ) + signed_icons: list[str] = [] + + def sign_icon(file_id: str) -> str: + assert isinstance(sqlite_engine.pool, QueuePool) + assert sqlite_engine.pool.checkedout() == 0 + signed_icons.append(file_id) + return f"https://files.example/{file_id}?sign=preview" + + for module in (trial_module, _trial_preview_admission_module): + monkeypatch.setattr(module, "application_services", lambda: services) + monkeypatch.setattr(login_module, "current_user", None) + monkeypatch.setattr(console_wraps, "_is_setup_completed", lambda: False) + monkeypatch.setattr(fields_module.file_helpers, "get_signed_file_url", sign_icon) + app = Flask(__name__) + app.config.update(TESTING=True, RESTX_ERROR_404_HELP=False) + api = ExternalApi(app) + api.add_resource(trial_module.TrialAppParameterApi, "/trial-apps//parameters") + api.add_resource(trial_module.TrialSitApi, "/trial-apps//site") + api.add_resource(trial_module.DatasetListApi, "/trial-apps//datasets") + return _TrialPreviewHarness( + app, + owner, + target, + listing, + site, + config, + dataset, + creator_id, + sqlite_session_factory, + sqlite_engine, + catalog, + sessions, + signed_icons, + ) + + @staticmethod + @pytest.mark.parametrize("endpoint", ["parameters", "site", "datasets"]) + @pytest.mark.parametrize("membership", ["trial", "catalog"]) + def test_anonymous_preview_ignores_trial_feature_quota_and_setup( + harness: _TrialPreviewHarness, endpoint: _TrialPreviewEndpoint, membership: str + ) -> None: + if membership == "catalog": + with harness.factory.begin() as session: + session.execute(delete(TrialApp).where(TrialApp.app_id == harness.target.id)) + harness.catalog.ids.add(harness.target.id) + + response = harness.get(endpoint, query=harness.dataset_query() if endpoint == "datasets" else "") + + assert response.status_code == 200, response.get_json() + assert harness.catalog.calls == ([] if membership == "trial" else [harness.target.id]) + with harness.factory() as session: + assert session.scalar(select(AccountTrialAppRecord.count)) == 100 + + @staticmethod + @pytest.mark.parametrize("endpoint", ["parameters", "site", "datasets"]) + @pytest.mark.parametrize("unavailable", ["unlisted", "disabled", "missing"]) + def test_preview_rejects_unavailable_apps_before_loading_content( + harness: _TrialPreviewHarness, endpoint: _TrialPreviewEndpoint, unavailable: str + ) -> None: + with harness.factory.begin() as session: + if unavailable == "unlisted": + session.execute(delete(TrialApp).where(TrialApp.app_id == harness.target.id)) + elif unavailable == "disabled": + session.execute( + text("UPDATE apps SET status = 'disabled' WHERE id = :app_id"), {"app_id": harness.target.id} + ) + else: + session.execute(delete(_trial_preview_App).where(_trial_preview_App.id == harness.target.id)) + + _trial_preview_assert_error(harness.get(endpoint, query=harness.dataset_query()), 404, "app_not_found") + assert harness.signed_icons == [] + + @staticmethod + def test_parameters_preserve_complete_shape_and_zero_false_empty_values(harness: _TrialPreviewHarness) -> None: + response = harness.get("parameters") + assert response.status_code == 200, response.get_json() + assert response.get_json() == { + "opening_statement": "", + "suggested_questions": [], + "suggested_questions_after_answer": {"enabled": False}, + "speech_to_text": {"enabled": False}, + "text_to_speech": {"enabled": False, "voice": "", "autoPlay": "disabled"}, + "retriever_resource": {"enabled": False}, + "annotation_reply": {"enabled": False}, + "more_like_this": {"enabled": False}, + "user_input_form": _INPUT_FORM, + "sensitive_word_avoidance": {"enabled": False, "type": "", "configs": []}, + "file_upload": {"enabled": False, "number_limits": 0}, + "system_parameters": { + "image_file_size_limit": 1, + "video_file_size_limit": 2, + "audio_file_size_limit": 3, + "file_size_limit": 4, + "workflow_file_upload_limit": 5, + }, + } + + @staticmethod + def test_missing_parameter_configuration_is_app_unavailable(harness: _TrialPreviewHarness) -> None: + with harness.factory.begin() as session: + session.execute(delete(AppModelConfig).where(AppModelConfig.id == harness.config.id)) + _trial_preview_assert_error(harness.get("parameters"), 400, "app_unavailable") + + @staticmethod + def test_site_preserves_complete_fields_and_signs_after_session_closes(harness: _TrialPreviewHarness) -> None: + response = harness.get("site") + assert response.status_code == 200, response.get_json() + assert harness.target.enable_site is False + assert response.get_json() == { + "title": "Preview site", + "chat_color_theme": "", + "chat_color_theme_inverted": False, + "icon_type": "image", + "icon": harness.site.icon, + "icon_background": "", + "description": None, + "copyright": "", + "privacy_policy": None, + "input_placeholder": "", + "custom_disclaimer": "", + "default_language": "en-US", + "show_workflow_steps": False, + "use_icon_as_answer_icon": False, + "icon_url": f"https://files.example/{harness.site.icon}?sign=preview", + } + assert harness.signed_icons == [harness.site.icon] + + @staticmethod + @pytest.mark.parametrize("icon_type", [IconType.EMOJI, None]) + def test_non_image_site_icon_has_no_signed_url(harness: _TrialPreviewHarness, icon_type: IconType | None) -> None: + with harness.factory.begin() as session: + session.execute(update(Site).where(Site.id == harness.site.id).values(icon_type=icon_type, icon="")) + response = harness.get("site") + assert response.status_code == 200 + assert response.get_json()["icon_url"] is None + assert harness.signed_icons == [] + + @staticmethod + @pytest.mark.parametrize( + ("missing", "code"), + [ + ("site", "app_site_unavailable"), + ("owner", "app_owner_unavailable"), + ("archived_owner", "app_owner_unavailable"), + ], + ) + def test_site_errors_distinguish_site_and_owner_unavailability( + harness: _TrialPreviewHarness, missing: str, code: str + ) -> None: + with harness.factory.begin() as session: + if missing == "site": + session.execute(delete(Site).where(Site.app_id == harness.target.id)) + elif missing == "owner": + session.execute(delete(Tenant).where(Tenant.id == harness.owner.id)) + else: + session.execute(update(Tenant).where(Tenant.id == harness.owner.id).values(status=TenantStatus.ARCHIVE)) + _trial_preview_assert_error(harness.get("site"), 403, code) + assert harness.signed_icons == [] + + @staticmethod + def test_dataset_ids_filter_by_actual_app_owner_without_requiring_binding(harness: _TrialPreviewHarness) -> None: + with harness.factory.begin() as session: + other_owner = Dataset( + tenant_id=harness.listing.tenant_id, name="Listing tenant dataset", created_by=harness.creator_id + ) + unrequested = Dataset(tenant_id=harness.owner.id, name="Not requested", created_by=harness.creator_id) + session.add_all([other_owner, unrequested]) + response = harness.get( + "datasets", query=f"?ids={harness.dataset.id}&ids={harness.dataset.id}&ids={other_owner.id}&ids={uuid4()}" + ) + assert response.status_code == 200, response.get_json() + assert response.get_json() == { + "data": [ + { + "id": harness.dataset.id, + "name": "Owner dataset", + "description": None, + "permission": "only_me", + "data_source_type": "upload_file", + "indexing_technique": None, + "created_by": harness.creator_id, + "created_at": int(_TRIAL_PREVIEW_CREATED_AT.timestamp()), + "permission_keys": [], + } + ], + "has_more": False, + "limit": 20, + "total": 1, + "page": 1, + } + + @staticmethod + @pytest.mark.parametrize( + ("query", "page", "limit", "has_more"), + [ + ("&page=0&limit=0", 0, 0, False), + ("&page=-2&limit=-1", -2, -1, False), + ("&page=invalid&limit=invalid", 1, 20, False), + ("&page=99&limit=1", 99, 1, False), + ], + ) + def test_dataset_page_and_limit_are_metadata_without_slicing( + harness: _TrialPreviewHarness, query: str, page: int, limit: int, has_more: bool + ) -> None: + response = harness.get("datasets", query=harness.dataset_query() + query) + assert response.status_code == 200, response.get_json() + body = response.get_json() + assert [item["id"] for item in body["data"]] == [harness.dataset.id] + assert {key: body[key] for key in ("page", "limit", "total", "has_more")} == { + "page": page, + "limit": limit, + "total": 1, + "has_more": has_more, + } + + @staticmethod + @pytest.mark.parametrize("ids", ["missing", "empty"]) + @pytest.mark.parametrize("limit", [20, 0]) + def test_unmatched_dataset_ids_return_empty_result(harness: _TrialPreviewHarness, ids: str, limit: int) -> None: + response = harness.get("datasets", query=f"?ids={uuid4() if ids == 'missing' else ''}&limit={limit}") + assert response.status_code == 200 + assert response.get_json() == {"data": [], "has_more": False, "limit": limit, "total": 0, "page": 1} + + @staticmethod + def test_dataset_list_requires_ids(harness: _TrialPreviewHarness) -> None: + _trial_preview_assert_error(harness.get("datasets"), 400, "need_add_ids") + + +"""Preview detail contracts through real HTTP, services, and SQLite queries.""" + + +_TrialPreviewDetailsEndpoint = Literal["detail", "workflows"] +_TRIAL_PREVIEW_DETAILS_CREATED_AT = datetime(2024, 1, 1) +_UPDATED_AT = datetime(2024, 1, 2) + + +@dataclass +class _TrialPreviewDetailsCatalog: + engine: Engine + ids: set[str] = field(default_factory=set) + calls: list[str] = field(default_factory=list) + + def contains(self, app_id: str) -> bool: + assert isinstance(self.engine.pool, QueuePool) + assert self.engine.pool.checkedout() == 0 + self.calls.append(app_id) + return app_id in self.ids + + def list_recommended(self, language: str) -> RecommendedAppCatalogPage: + raise AssertionError(f"Unexpected listing for {language}") + + def list_learn_dify(self, language: str) -> RecommendedAppCatalogPage: + raise AssertionError(f"Unexpected listing for {language}") + + def get_detail(self, app_id: str) -> RecommendedAppDetailRecord | None: + raise AssertionError(f"Unexpected catalog detail for {app_id}") + + +@dataclass +class _TrialPreviewDetailsExternalio: + engine: Engine + owner_id: str + viewer_workspace_id: str + decryptions: list[tuple[str, str]] = field(default_factory=list) + tool_calls: list[tuple[str, str, str]] = field(default_factory=list) + + def encrypt(self, tenant_id: str, token: str) -> str: + assert tenant_id == self.owner_id + assert token == "workflow-secret" + return "cipher-workflow" + + def decrypt(self, tenant_id: str, token: str) -> str: + assert isinstance(self.engine.pool, QueuePool) + assert self.engine.pool.checkedout() == 0 + self.decryptions.append((tenant_id, token)) + if token == "cipher-workflow": + assert tenant_id == self.owner_id + return "workflow-secret" + assert token == "cipher-tool" + assert tenant_id == self.viewer_workspace_id + return "tool-secret" + + def tool_runtime(self, *, tenant_id: str, app_id: str, agent_tool: AgentToolEntity, user_id: str) -> Tool: + assert isinstance(self.engine.pool, QueuePool) + assert self.engine.pool.checkedout() == 0 + assert tenant_id == self.viewer_workspace_id + assert agent_tool.tool_name == "search" + self.tool_calls.append((tenant_id, app_id, user_id)) + parameter = ToolParameter.model_validate( + {"name": "api_key", "label": {"en_US": "Key"}, "type": "secret-input", "form": "form"} + ) + return cast( + Tool, _ToolRuntime(SimpleNamespace(parameters=[parameter], identity=SimpleNamespace(name="search"))) + ) + + +@dataclass +class _ToolRuntime: + entity: SimpleNamespace + + def get_runtime_parameters(self) -> list[ToolParameter]: + return [] + + +@dataclass(frozen=True) +class _TrialPreviewDetailsApplicationservices: + app_previews: AppPreviewQueryService + app_preview_details: AppPreviewDetailsRuntime + recommended_app_queries: RecommendedAppQueryService + + +@dataclass(frozen=True) +class _TrialPreviewDetailsHarness: + app: Flask + owner: Tenant + viewer_workspace: Tenant + viewer: Account + author: Account + updater: Account + target: _trial_preview_details_App + config: AppModelConfig + site: Site + workflow: Workflow + tag: Tag + factory: sessionmaker[Session] + engine: Engine + catalog: _TrialPreviewDetailsCatalog + io: _TrialPreviewDetailsExternalio + sessions: list[Session] + + def get( + self, + endpoint: _TrialPreviewDetailsEndpoint = "detail", + *, + method: str = "GET", + csrf_token: str | None = None, + ) -> TestResponse: + suffix = "" if endpoint == "detail" else "/workflows" + statements: list[str] = [] + token = generate_csrf_token(self.viewer.id) if csrf_token is None else csrf_token + client = self.app.test_client() + client.set_cookie(COOKIE_NAME_CSRF_TOKEN, token) + + def record_statement( + _connection: Connection, + _cursor: object, + statement: str, + _parameters: object, + _context: ExecutionContext, + _executemany: bool, + ) -> None: + statements.append(statement) + + event.listen(self.engine, "before_cursor_execute", record_statement) + try: + response = client.open( + f"/trial-apps/{self.target.id}{suffix}", method=method, headers={HEADER_NAME_CSRF_TOKEN: token} + ) + finally: + event.remove(self.engine, "before_cursor_execute", record_statement) + assert not [sql for sql in statements if sql.lstrip().upper().startswith(("INSERT ", "UPDATE ", "DELETE "))] + assert response.headers["Content-Type"] == "application/json" + if method != "HEAD": + assert int(response.headers["Content-Length"]) == len(response.data) + assert isinstance(self.engine.pool, QueuePool) + assert self.engine.pool.checkedout() == 0 + assert all(not session.in_transaction() and not session.identity_map for session in self.sessions) + return response + + def assert_app_not_queried(self) -> None: + assert self.sessions == [] + assert self.catalog.calls == [] + assert self.io.tool_calls == [] + assert self.io.decryptions == [] + + def legacy_detail(self) -> dict[str, object]: + """The unchanged ORM response view is the pre-migration serialization oracle.""" + with self.app.test_request_context(f"/trial-apps/{self.target.id}"), self.factory() as session: + app = session.get(_trial_preview_details_App, self.target.id) + assert app is not None + return trial_module.TrialAppDetailResponse.model_validate( + AppResponseView(app, session=session, account=self.viewer), from_attributes=True + ).model_dump(mode="json") + + +def _trial_preview_details_assert_error(response: TestResponse, status: int, code: str) -> None: + assert response.status_code == status, response.get_json() + body = response.get_json() + assert body["status"] == status + assert body["code"] == code + assert body["message"] + + +class TestTrialPreviewDetailsCases: + @staticmethod + @pytest.fixture + def harness( + monkeypatch: pytest.MonkeyPatch, + config_overrides: Callable[..., None], + sqlite_engine: Engine, + sqlite_session_factory: sessionmaker[Session], + ) -> _TrialPreviewDetailsHarness: + config_overrides( + LOGIN_DISABLED=False, + DEPLOYMENT_EDITION=DeploymentEdition.COMMUNITY, + INIT_PASSWORD="", + SECRET_KEY="preview-admission-test-secret-key", + CONSOLE_WEB_URL="", + CONSOLE_API_URL="", + SERVICE_API_URL="", + APP_WEB_URL="", + ) + owner = Tenant(name="App owner") + viewer_workspace = Tenant(name="Viewer workspace") + viewer = Account(name="Viewer", email="viewer@example.com", status=AccountStatus.ACTIVE) + viewer._current_tenant = viewer_workspace + author = Account(name="Author", email="author@example.com") + updater = Account(name="Updater", email="updater@example.com") + target = _trial_preview_details_App( + id=str(uuid4()), + tenant_id=owner.id, + name="Preview app", + description="", + mode=_trial_preview_details_AppMode.CHAT, + icon_type=IconType.EMOJI, + icon="", + icon_background=None, + enable_site=False, + enable_api=False, + max_active_requests=0, + ) + target.created_at = _TRIAL_PREVIEW_DETAILS_CREATED_AT + target.updated_at = _UPDATED_AT + config = AppModelConfig( + app_id=target.id, + opening_statement="", + suggested_questions="[]", + model='{"provider":"example/model","name":"model","mode":"chat","completion_params":{"temperature":0}}', + agent_mode='{"enabled":false,"strategy":"function_call","tools":[]}', + user_input_form='[{"number":{"label":"Count","variable":"count","default":0,"required":false}}]', + ) + config.created_at = _TRIAL_PREVIEW_DETAILS_CREATED_AT + config.updated_at = _UPDATED_AT + site = Site( + app_id=target.id, + title="Preview site", + default_language="en-US", + customize_token_strategy=CustomizeTokenStrategy.UUID, + code="preview-code", + icon_type=IconType.EMOJI, + icon="", + icon_background=None, + chat_color_theme="", + description=None, + custom_disclaimer="", + show_workflow_steps=False, + ) + site.created_at = _TRIAL_PREVIEW_DETAILS_CREATED_AT + site.updated_at = _UPDATED_AT + io = _TrialPreviewDetailsExternalio(sqlite_engine, owner.id, viewer_workspace.id) + monkeypatch.setattr(encrypter, "encrypt_token", io.encrypt) + monkeypatch.setattr(encrypter, "decrypt_token", io.decrypt) + monkeypatch.setattr(ToolManager, "get_agent_tool_runtime", io.tool_runtime) + monkeypatch.setattr(ToolManager, "get_hardcoded_provider", lambda _provider_id: object()) + with sqlite_session_factory.begin() as session: + session.add_all([owner, viewer_workspace, viewer, author, updater, target, config, site]) + session.flush() + target.created_by = author.id + target.updated_by = updater.id + target.app_model_config_id = config.id + workflow = Workflow.new( + tenant_id=owner.id, + app_id=target.id, + type="workflow", + version="draft", + graph='{"nodes":[],"edges":[]}', + features='{"suggested_questions_after_answer":{"enabled":false}}', + created_by=author.id, + environment_variables=[ + SecretVariable(id="secret", name="api_key", value="workflow-secret"), + LLMEnvironmentVariable( + id="llm", name="shared_model", value={"provider": "provider", "name": "model", "mode": "chat"} + ), + ], + conversation_variables=[ + StringVariable(id="topic", name="topic", value="sqlite", selector=["conversation", "topic"]) + ], + rag_pipeline_variables=[], + ) + workflow.created_at = _TRIAL_PREVIEW_DETAILS_CREATED_AT + workflow.updated_at = _UPDATED_AT + workflow.updated_by = updater.id + target.workflow_id = workflow.id + tag = Tag(tenant_id=owner.id, type=TagType.APP, name="Owner tag", created_by=author.id) + session.add_all([workflow, tag]) + session.flush() + session.add_all( + [ + TrialApp(app_id=target.id, tenant_id=str(uuid4()), trial_limit=0), + AccountTrialAppRecord(app_id=target.id, account_id=viewer.id, count=100), + TagBinding(tenant_id=owner.id, tag_id=tag.id, target_id=target.id, created_by=author.id), + WorkflowToolProvider( + name="trial-workflow", + label="Trial workflow", + icon="icon", + app_id=target.id, + version="1.0.0", + user_id=author.id, + tenant_id=owner.id, + description="Workflow provider", + parameter_configuration="[]", + ), + ] + ) + factory = sessionmaker(bind=sqlite_engine, expire_on_commit=False) + sessions: list[Session] = [] + + @event.listens_for(factory, "after_begin") + def track_session(session: Session, _transaction: SessionTransaction, _connection: Connection) -> None: + assert all(previous is session or not previous.in_transaction() for previous in sessions) + sessions.append(session) + + catalog = _TrialPreviewDetailsCatalog(sqlite_engine) + recommendations = RecommendedAppQueryService( + catalog=catalog, trial_apps=TrialAppRepository(factory), trial_enabled=False + ) + services = _TrialPreviewDetailsApplicationservices( + app_previews=AppPreviewQueryService( + apps=AppPreviewQueryRepository(session_factory=factory), is_previewable=recommendations.is_previewable + ), + app_preview_details=AppPreviewDetailsRuntime(details=AppPreviewQueryRepository(session_factory=factory)), + recommended_app_queries=recommendations, + ) + for module in (trial_module, _trial_preview_details_admission_module): + monkeypatch.setattr(module, "application_services", lambda: services) + monkeypatch.setattr(login_module, "current_user", viewer) + monkeypatch.setattr(app_module, "current_user", viewer) + monkeypatch.setattr(console_wraps, "_is_setup_completed", lambda: True) + app = Flask(__name__) + app.config.update(TESTING=True, RESTX_ERROR_404_HELP=False) + login_manager = DifyLoginManager() + login_manager.init_app(app) + login_manager.unauthorized_handler(unauthorized_handler) + api = ExternalApi(app) + api.error_handlers = console_api.error_handlers.copy() + api.add_resource(trial_module.AppApi, "/trial-apps/") + api.add_resource(trial_module.AppWorkflowApi, "/trial-apps//workflows") + return _TrialPreviewDetailsHarness( + app, + owner, + viewer_workspace, + viewer, + author, + updater, + target, + config, + site, + workflow, + tag, + sqlite_session_factory, + sqlite_engine, + catalog, + io, + sessions, + ) + + @staticmethod + @pytest.mark.parametrize( + "mode", + [ + _trial_preview_details_AppMode.CHAT, + _trial_preview_details_AppMode.COMPLETION, + _trial_preview_details_AppMode.ADVANCED_CHAT, + _trial_preview_details_AppMode.WORKFLOW, + _trial_preview_details_AppMode.AGENT_CHAT, + ], + ) + def test_app_detail_keeps_original_response_fields_and_stored_audit_times( + harness: _TrialPreviewDetailsHarness, mode: _trial_preview_details_AppMode + ) -> None: + with harness.factory.begin() as session: + session.execute( + update(_trial_preview_details_App) + .where(_trial_preview_details_App.id == harness.target.id) + .values(mode=mode) + ) + expected = harness.legacy_detail() + if mode == _trial_preview_details_AppMode.AGENT_CHAT: + expected_config = expected["model_config"] + assert isinstance(expected_config, dict) + # Legacy masking could autoflush a response-only config mutation and bump + # this timestamp. A read-only preview returns the actual stored audit value. + expected_config["updated_at"] = int(_UPDATED_AT.timestamp()) + + response = harness.get() + + assert response.status_code == 200, response.get_json() + assert response.get_json() == expected + body = response.get_json() + assert body["max_active_requests"] == 0 + assert body["enable_site"] is False + assert body["tags"] == [{"id": harness.tag.id, "name": "Owner tag", "type": "app"}] + assert body["api_base_url"] == "http://localhost/v1" + assert body["site"]["app_base_url"] == "http://localhost" + assert body["site"]["access_token"] == "preview-code" + with harness.factory() as session: + assert session.scalar(select(AccountTrialAppRecord.count)) == 100 + assert ( + session.scalar(select(AppModelConfig.updated_at).where(AppModelConfig.id == harness.config.id)) + == _UPDATED_AT + ) + + @staticmethod + def test_legacy_agent_preview_masks_tools_without_upgrading_mode_or_mutating_config( + harness: _TrialPreviewDetailsHarness, + ) -> None: + agent_mode = { + "enabled": True, + "strategy": "function_call", + "tools": [ + { + "provider_type": "builtin", + "provider_id": "example", + "tool_name": "search", + "tool_parameters": {"api_key": "cipher-tool"}, + } + ], + } + with harness.factory.begin() as session: + session.execute( + update(AppModelConfig) + .where(AppModelConfig.id == harness.config.id) + .values(agent_mode=json.dumps(agent_mode)) + ) + app_updated_at = session.scalar( + select(_trial_preview_details_App.updated_at).where(_trial_preview_details_App.id == harness.target.id) + ) + config_updated_at = session.scalar( + select(AppModelConfig.updated_at).where(AppModelConfig.id == harness.config.id) + ) + + response = harness.get() + + assert response.status_code == 200, response.get_json() + body = response.get_json() + assert body["mode"] == "agent-chat" + assert body["model_config"]["agent_mode"]["tools"][0]["tool_parameters"] == {"api_key": "to*******et"} + assert harness.io.tool_calls == [(harness.viewer_workspace.id, harness.target.id, harness.viewer.id)] + assert harness.io.decryptions == [(harness.viewer_workspace.id, "cipher-tool")] + assert harness.viewer_workspace.id != harness.owner.id + with harness.factory() as session: + app = session.get(_trial_preview_details_App, harness.target.id) + config = session.get(AppModelConfig, harness.config.id) + assert app is not None + assert config is not None + assert app.mode == _trial_preview_details_AppMode.CHAT + assert app.updated_at == app_updated_at + assert config.agent_mode == json.dumps(agent_mode) + assert config.updated_at == config_updated_at + + @staticmethod + @pytest.mark.parametrize( + "provider_fields", + [{}, {"provider_id": None}, {"provider_id": 42}, {"provider_id": ""}, {"provider_id": "invalid/provider"}], + ids=["missing", "null", "non-string", "empty", "invalid-format"], + ) + def test_invalid_builtin_config_does_not_break_detail_or_valid_tool_masking( + harness: _TrialPreviewDetailsHarness, monkeypatch: pytest.MonkeyPatch, provider_fields: dict[str, JsonValue] + ) -> None: + agent_mode: dict[str, JsonValue] = { + "enabled": True, + "strategy": "function_call", + "tools": [ + { + "provider_type": "builtin", + "tool_name": "broken", + "tool_parameters": {"api_key": "cipher-secret"}, + "credential_id": None, + **provider_fields, + }, + { + "provider_type": "builtin", + "provider_id": "example", + "tool_name": "search", + "tool_parameters": {"api_key": "cipher-tool"}, + }, + ], + } + with harness.factory.begin() as session: + session.execute( + update(AppModelConfig) + .where(AppModelConfig.id == harness.config.id) + .values(agent_mode=json.dumps(agent_mode)) + ) + + def hardcoded_provider(provider_id: str) -> object: + if provider_id == "example": + return object() + raise KeyError(provider_id) + + def tool_runtime(*, tenant_id: str, app_id: str, agent_tool: AgentToolEntity, user_id: str) -> Tool: + if agent_tool.provider_id != "example": + raise ValueError("Invalid historical provider ID") + return harness.io.tool_runtime(tenant_id=tenant_id, app_id=app_id, agent_tool=agent_tool, user_id=user_id) + + plugin_queries: list[str] = [] + + def check_plugins(tenant_id: str, provider_ids: Sequence[GenericProviderID]) -> Sequence[bool]: + assert tenant_id == harness.owner.id + plugin_queries.extend(str(provider_id) for provider_id in provider_ids) + return [False] * len(provider_ids) + + monkeypatch.setattr(ToolManager, "get_hardcoded_provider", hardcoded_provider) + monkeypatch.setattr(ToolManager, "get_agent_tool_runtime", tool_runtime) + monkeypatch.setattr(PluginService, "check_tools_existence", check_plugins) + + response = harness.get() + + assert response.status_code == 200, response.get_json() + body = response.get_json() + assert body["deleted_tools"] == [] + assert body["model_config"]["agent_mode"]["tools"][0]["tool_parameters"] == {} + assert body["model_config"]["agent_mode"]["tools"][1]["tool_parameters"] == {"api_key": "to*******et"} + assert b"cipher-secret" not in response.data + assert plugin_queries == [] + assert harness.io.tool_calls == [(harness.viewer_workspace.id, harness.target.id, harness.viewer.id)] + with harness.factory() as session: + assert session.scalar( + select(AppModelConfig.agent_mode).where(AppModelConfig.id == harness.config.id) + ) == json.dumps(agent_mode) + + @staticmethod + @pytest.mark.parametrize("failure", ["short-config", "runtime", "masking"]) + def test_unmaskable_tool_parameters_are_removed_from_detail_without_changing_storage( + harness: _TrialPreviewDetailsHarness, monkeypatch: pytest.MonkeyPatch, failure: str + ) -> None: + broken_tool: dict[str, JsonValue] = { + "provider_type": "builtin", + "tool_parameters": {"api_key": "cipher-secret"}, + } + if failure != "short-config": + broken_tool.update(provider_id="broken", tool_name="broken") + agent_mode: dict[str, JsonValue] = { + "enabled": True, + "strategy": "function_call", + "tools": [ + broken_tool, + { + "provider_type": "builtin", + "provider_id": "example", + "tool_name": "search", + "tool_parameters": {"api_key": "cipher-tool"}, + }, + ], + } + with harness.factory.begin() as session: + session.execute( + update(AppModelConfig) + .where(AppModelConfig.id == harness.config.id) + .values(agent_mode=json.dumps(agent_mode)) + ) + + class MaskingFailureRuntime(_ToolRuntime): + parameter_queries: int = 0 + + @override + def get_runtime_parameters(self) -> list[ToolParameter]: + self.parameter_queries += 1 + if self.parameter_queries == 2: + raise RuntimeError("Tool metadata became unavailable during masking") + return [] + + parameter = ToolParameter.model_validate( + {"name": "api_key", "label": {"en_US": "Key"}, "type": "secret-input", "form": "form"} + ) + broken_runtime = MaskingFailureRuntime( + SimpleNamespace(parameters=[parameter], identity=SimpleNamespace(name="broken")) + ) + + def tool_runtime(*, tenant_id: str, app_id: str, agent_tool: AgentToolEntity, user_id: str) -> Tool: + if agent_tool.provider_id == "example": + return harness.io.tool_runtime( + tenant_id=tenant_id, app_id=app_id, agent_tool=agent_tool, user_id=user_id + ) + if failure == "runtime": + raise RuntimeError("Tool provider is unavailable") + return cast(Tool, broken_runtime) + + def decrypt(tenant_id: str, token: str) -> str: + if token == "cipher-secret": + assert tenant_id == harness.viewer_workspace.id + return "plain-secret" + return harness.io.decrypt(tenant_id, token) + + monkeypatch.setattr(ToolManager, "get_agent_tool_runtime", tool_runtime) + monkeypatch.setattr(encrypter, "decrypt_token", decrypt) + + response = harness.get() + + assert response.status_code == 200, response.get_json() + tools = response.get_json()["model_config"]["agent_mode"]["tools"] + assert tools[0]["tool_parameters"] == {} + assert tools[1]["tool_parameters"] == {"api_key": "to*******et"} + assert b"cipher-secret" not in response.data + assert b"plain-secret" not in response.data + assert broken_runtime.parameter_queries == (2 if failure == "masking" else 0) + with harness.factory() as session: + assert session.scalar( + select(AppModelConfig.agent_mode).where(AppModelConfig.id == harness.config.id) + ) == json.dumps(agent_mode) + + @staticmethod + def test_workflow_is_anonymous_and_masks_secrets_after_query_session_closes( + harness: _TrialPreviewDetailsHarness, monkeypatch: pytest.MonkeyPatch + ) -> None: + monkeypatch.setattr(login_module, "current_user", None) + monkeypatch.setattr(console_wraps, "_is_setup_completed", lambda: False) + response = harness.get("workflows", csrf_token="") + assert response.status_code == 200, response.get_json() + assert response.get_json() == { + "id": harness.workflow.id, + "graph": {"nodes": [], "edges": []}, + "features": {"suggested_questions_after_answer": {"enabled": False}}, + "hash": harness.workflow.unique_hash, + "version": "draft", + "marked_name": "", + "marked_comment": "", + "created_by": {"id": harness.author.id, "name": "Author", "email": "author@example.com"}, + "created_at": int(_TRIAL_PREVIEW_DETAILS_CREATED_AT.timestamp()), + "updated_by": {"id": harness.updater.id, "name": "Updater", "email": "updater@example.com"}, + "updated_at": int(_UPDATED_AT.timestamp()), + "tool_published": True, + "environment_variables": [ + { + "value_type": "secret", + "value": encrypter.full_mask_token(), + "id": "secret", + "name": "api_key", + "description": "", + "selector": ["env", "api_key"], + }, + { + "value_type": "llm", + "value": {"provider": "provider", "name": "model", "mode": "chat"}, + "id": "llm", + "name": "shared_model", + "description": "", + "selector": ["env", "shared_model"], + }, + ], + "conversation_variables": [ + {"id": "topic", "name": "topic", "value_type": "string", "value": "sqlite", "description": ""} + ], + "rag_pipeline_variables": [], + } + assert harness.io.decryptions == [(harness.owner.id, "cipher-workflow")] + assert b"workflow-secret" not in response.data + assert b"cipher-workflow" not in response.data + + @staticmethod + @pytest.mark.parametrize("endpoint", ["detail", "workflows"]) + def test_catalog_only_admission_keeps_preview_available( + harness: _TrialPreviewDetailsHarness, endpoint: _TrialPreviewDetailsEndpoint + ) -> None: + with harness.factory.begin() as session: + session.execute(delete(TrialApp).where(TrialApp.app_id == harness.target.id)) + harness.catalog.ids.add(harness.target.id) + response = harness.get(endpoint) + assert response.status_code == 200, response.get_json() + assert harness.catalog.calls == [harness.target.id] + + @staticmethod + @pytest.mark.parametrize("endpoint", ["detail", "workflows"]) + @pytest.mark.parametrize("unavailable", ["unlisted", "disabled", "missing"]) + def test_unavailable_apps_fail_admission( + harness: _TrialPreviewDetailsHarness, endpoint: _TrialPreviewDetailsEndpoint, unavailable: str + ) -> None: + with harness.factory.begin() as session: + if unavailable == "unlisted": + session.execute(delete(TrialApp).where(TrialApp.app_id == harness.target.id)) + elif unavailable == "disabled": + session.execute( + text("UPDATE apps SET status='disabled' WHERE id=:app_id"), {"app_id": harness.target.id} + ) + else: + session.execute( + delete(_trial_preview_details_App).where(_trial_preview_details_App.id == harness.target.id) + ) + _trial_preview_details_assert_error(harness.get(endpoint), 404, "app_not_found") + assert harness.io.decryptions == [] + + @staticmethod + def test_anonymous_app_detail_requires_account( + harness: _TrialPreviewDetailsHarness, monkeypatch: pytest.MonkeyPatch + ) -> None: + monkeypatch.setattr(login_module, "current_user", None) + response = harness.get() + assert response.status_code == 401 + assert response.get_json() == {"code": "unauthorized", "message": "Unauthorized.", "status": 401} + harness.assert_app_not_queried() + + @staticmethod + @pytest.mark.parametrize( + ("init_password", "error_code"), [("", "not_setup"), ("setup-password", "not_init_validated")] + ) + def test_app_detail_requires_setup_before_account_or_app_queries( + harness: _TrialPreviewDetailsHarness, + monkeypatch: pytest.MonkeyPatch, + config_overrides: Callable[..., None], + init_password: str, + error_code: str, + ) -> None: + config_overrides(INIT_PASSWORD=init_password) + monkeypatch.setattr(console_wraps, "_is_setup_completed", lambda: False) + monkeypatch.setattr(login_module, "current_user", None) + _trial_preview_details_assert_error(harness.get(csrf_token=""), 401, error_code) + harness.assert_app_not_queried() + + @staticmethod + @pytest.mark.parametrize("csrf_token", ["", "invalid-signed-token"], ids=["missing", "invalid"]) + def test_app_detail_rejects_invalid_csrf_before_app_queries( + harness: _TrialPreviewDetailsHarness, csrf_token: str + ) -> None: + response = harness.get(csrf_token=csrf_token) + _trial_preview_details_assert_error(response, 401, "unauthorized") + assert response.get_json()["message"] == "CSRF token is missing or invalid." + harness.assert_app_not_queried() + + @staticmethod + def test_app_detail_rejects_csrf_token_for_another_account(harness: _TrialPreviewDetailsHarness) -> None: + response = harness.get(csrf_token=generate_csrf_token(harness.author.id)) + _trial_preview_details_assert_error(response, 401, "unauthorized") + assert response.get_json()["message"] == "CSRF token is missing or invalid." + harness.assert_app_not_queried() + + @staticmethod + def test_app_detail_rejects_uninitialized_account_before_app_queries(harness: _TrialPreviewDetailsHarness) -> None: + harness.viewer.status = AccountStatus.UNINITIALIZED + _trial_preview_details_assert_error(harness.get(), 400, "account_not_initialized") + harness.assert_app_not_queried() + + @staticmethod + def test_missing_active_workspace_uses_shared_console_error_and_monitoring( + harness: _TrialPreviewDetailsHarness, + ) -> None: + harness.viewer._current_tenant = None + errors: list[BaseException] = [] + + def record_error(_sender: object, *, exception: BaseException) -> None: + errors.append(exception) + + with got_request_exception.connected_to(record_error): + response = harness.get() + + _trial_preview_details_assert_error(response, 500, "active_workspace_required") + assert response.get_json()["message"] == "Internal Server Error" + assert len(errors) == 1 + assert isinstance(errors[0], ActiveWorkspaceRequiredError) + harness.assert_app_not_queried() + + @staticmethod + def test_deleted_admitted_account_is_unauthorized(harness: _TrialPreviewDetailsHarness) -> None: + with harness.factory.begin() as session: + session.execute(delete(Account).where(Account.id == harness.viewer.id)) + _trial_preview_details_assert_error(harness.get(), 401, "unauthorized") + + @staticmethod + def test_missing_site_is_specific_error(harness: _TrialPreviewDetailsHarness) -> None: + with harness.factory.begin() as session: + session.execute(delete(Site).where(Site.app_id == harness.target.id)) + _trial_preview_details_assert_error(harness.get(), 403, "app_site_unavailable") + + @staticmethod + def test_missing_model_config_retains_nullable_detail(harness: _TrialPreviewDetailsHarness) -> None: + with harness.factory.begin() as session: + session.execute(delete(AppModelConfig).where(AppModelConfig.id == harness.config.id)) + response = harness.get() + assert response.status_code == 200, response.get_json() + assert response.get_json()["model_config"] is None + + @staticmethod + def test_agent_detail_keeps_model_less_presentation_config(harness: _TrialPreviewDetailsHarness) -> None: + with harness.factory.begin() as session: + session.execute( + update(_trial_preview_details_App) + .where(_trial_preview_details_App.id == harness.target.id) + .values(mode=_trial_preview_details_AppMode.AGENT) + ) + session.execute(update(AppModelConfig).where(AppModelConfig.id == harness.config.id).values(model=None)) + agent_id, snapshot_id = str(uuid4()), str(uuid4()) + session.add_all( + [ + Agent( + id=agent_id, + tenant_id=harness.owner.id, + app_id=harness.target.id, + name="Preview agent", + scope=AgentScope.ROSTER, + source=AgentSource.AGENT_APP, + active_config_snapshot_id=snapshot_id, + ), + AgentConfigSnapshot( + id=snapshot_id, + tenant_id=harness.owner.id, + agent_id=agent_id, + version=1, + config_snapshot=AgentSoulConfig(), + ), + AgentConfigRevision( + tenant_id=harness.owner.id, + agent_id=agent_id, + current_snapshot_id=snapshot_id, + revision=1, + operation=AgentConfigRevisionOperation.PUBLISH_DRAFT, + ), + ] + ) + + response = harness.get() + + assert response.status_code == 200, response.get_json() + body = response.get_json() + assert body["mode"] == "agent" + assert body["model_config"]["model"] is None + assert body["model_config"]["opening_statement"] == "" + + @staticmethod + @pytest.mark.parametrize("missing", ["reference", "workflow"]) + def test_missing_workflow_is_app_unavailable(harness: _TrialPreviewDetailsHarness, missing: str) -> None: + with harness.factory.begin() as session: + if missing == "reference": + session.execute( + update(_trial_preview_details_App) + .where(_trial_preview_details_App.id == harness.target.id) + .values({_trial_preview_details_App.workflow_id: None}) + ) + else: + session.execute(delete(Workflow).where(Workflow.id == harness.workflow.id)) + _trial_preview_details_assert_error(harness.get("workflows"), 400, "app_unavailable") + + @staticmethod + @pytest.mark.parametrize("endpoint", ["detail", "workflows"]) + def test_head_keeps_get_status_and_headers_without_response_body( + harness: _TrialPreviewDetailsHarness, endpoint: _TrialPreviewDetailsEndpoint + ) -> None: + response = harness.get(endpoint) + head = harness.get(endpoint, method="HEAD") + assert response.status_code == head.status_code == 200 + assert head.data == b"" + assert head.headers == response.headers + + @staticmethod + @pytest.mark.parametrize( + ("endpoint", "relation"), [("detail", "config"), ("detail", "workflow"), ("workflows", "workflow")] + ) + def test_cross_owner_configuration_is_not_exposed( + harness: _TrialPreviewDetailsHarness, endpoint: _TrialPreviewDetailsEndpoint, relation: str + ) -> None: + with harness.factory.begin() as session: + if relation == "config": + session.execute( + update(AppModelConfig).where(AppModelConfig.id == harness.config.id).values(app_id=str(uuid4())) + ) + else: + session.execute( + update(Workflow).where(Workflow.id == harness.workflow.id).values(tenant_id=str(uuid4())) + ) + _trial_preview_details_assert_error(harness.get(endpoint), 400, "app_unavailable") + assert harness.io.decryptions == [] + + +"""Trial suggested questions through HTTP, real ownership queries, and SQLite.""" + + +@dataclass +class _TrialSuggestedQuestionsFeatures: + enabled: bool = True + setup_completed: bool = True + events: list[str] = field(default_factory=list) + + def is_trial_enabled(self) -> bool: + self.events.append("feature") + return self.enabled + + +@dataclass +class _Provider: + read_sessions: list[Session] + app_id: str + model_sessions: list[Session] = field(default_factory=list) + token_calls: int = 0 + tenant_ids: list[str] = field(default_factory=list) + prompts: list[str] = field(default_factory=list) + traces: list[TraceTask] = field(default_factory=list) + questions: list[str] = field(default_factory=lambda: ["Next question?", "Another question?"]) + failure: Exception | None = None + missing_model: bool = False + app_type: str = "chatbot" + + def get_default_model_instance(self, *, tenant_id: str, model_type: ModelType) -> ModelInstance: + assert model_type == ModelType.LLM + self.tenant_ids.append(tenant_id) + session = db.session() + assert session.get(_trial_suggested_questions_App, self.app_id) is not None + self.model_sessions.append(session) + if self.missing_model: + raise ProviderTokenNotInitError("No default model") + return cast(ModelInstance, self) + + def assert_queries_closed(self) -> None: + assert all( + not session.in_transaction() and not session.identity_map + for session in self.read_sessions + self.model_sessions + ) + + def get_llm_num_tokens(self, prompt_messages: Sequence[PromptMessage]) -> int: + self.assert_queries_closed() + self.token_calls += 1 + if self.failure is not None: + raise self.failure + return len(prompt_messages) + + def get_model_schema(self) -> AIModelEntity: + self.assert_queries_closed() + return AIModelEntity.model_construct(parameter_rules=[]) + + def invoke_llm( + self, + *, + prompt_messages: list[PromptMessage], + model_parameters: Mapping[str, object], + stop: list[str], + stream: bool, + ) -> LLMResult: + self.assert_queries_closed() + assert self.tenant_ids == [self.tenant_ids[0]] * 2 + assert get_credit_usage_metadata() == {"app_type": self.app_type, "created_by": "suggested_questions"} + assert model_parameters == {"max_tokens": 256, "temperature": 0.0} + assert stop == [] + assert stream is False + prompt = prompt_messages[0].get_text_content() + assert "Suggest concise follow-ups" in prompt + self.prompts.append(prompt) + return LLMResult( + model="question-model", + message=AssistantPromptMessage(content=json.dumps(self.questions)), + usage=LLMUsage.empty_usage(), + ) + + def add_trace_task(self, task: TraceTask) -> None: + self.traces.append(task) + + +@dataclass(frozen=True) +class _TrialSuggestedQuestionsTrialappservices: + access: TrialAppAccessService + + +@dataclass(frozen=True) +class _TrialSuggestedQuestionsApplicationservices: + trial_apps: _TrialSuggestedQuestionsTrialappservices + recommended_app_queries: _TrialSuggestedQuestionsFeatures + message_suggested_questions: MessageSuggestedQuestionsRuntime + + +@dataclass(frozen=True) +class _TrialSuggestedQuestionsHarness: + app: Flask + account: Account + target: _trial_suggested_questions_App + trial: TrialApp + conversation: Conversation + message: Message + config: AppModelConfig + factory: sessionmaker[Session] + features: _TrialSuggestedQuestionsFeatures + provider: _Provider + legacy_sessions: list[Session] + + def get(self, *, app_id: str | None = None, message_id: str | None = None) -> TestResponse: + return self.app.test_client().get( + f"/trial-apps/{app_id or self.target.id}/messages/{message_id or self.message.id}/suggested-questions" + ) + + def assert_closed(self) -> None: + assert all( + not session.in_transaction() and not session.identity_map + for session in self.provider.read_sessions + self.legacy_sessions + ) + + def usage(self) -> int | None: + with self.factory() as session: + return session.scalar( + select(AccountTrialAppRecord.count).where( + AccountTrialAppRecord.app_id == self.target.id, + AccountTrialAppRecord.account_id == self.account.id, + ) + ) + + +def _trial_suggested_questions_assert_error( + response: TestResponse, status: int, code: str, message: str | None = None +) -> None: + assert response.status_code == status, response.get_json() + body = response.get_json() + assert body["code"] == code + assert body["status"] == status + assert body["message"] + if message is not None: + assert body["message"] == message + assert response.headers["Content-Type"] == "application/json" + + +class TestTrialSuggestedQuestionsCases: + @staticmethod + @pytest.fixture + def harness( + monkeypatch: pytest.MonkeyPatch, + config_overrides: Callable[..., None], + sqlite_session_factory: sessionmaker[Session], + sqlite_engine: Engine, + ) -> Iterator[_TrialSuggestedQuestionsHarness]: + config_overrides(DEPLOYMENT_EDITION=DeploymentEdition.COMMUNITY, INIT_PASSWORD="", LOGIN_DISABLED=False) + features = _TrialSuggestedQuestionsFeatures() + account = Account(name="Trial viewer", email="trial@example.com") + account._current_tenant = Tenant(name="Viewer workspace") + target = _trial_suggested_questions_App( + tenant_id=str(uuid4()), + name="Trial app", + mode=_trial_suggested_questions_AppMode.CHAT, + enable_site=True, + enable_api=False, + ) + target.id = str(uuid4()) + trial = TrialApp(app_id=target.id, tenant_id=str(uuid4()), trial_limit=3) + config = AppModelConfig( + app_id=target.id, + suggested_questions_after_answer='{"enabled":true,"prompt":"Suggest concise follow-ups"}', + ) + with sqlite_session_factory.begin() as session: + session.add_all([account, target, trial, config]) + session.flush() + target.app_model_config_id = config.id + conversation = Conversation( + app_id=target.id, + app_model_config_id=config.id, + mode=_trial_suggested_questions_AppMode.CHAT, + name="Trial conversation", + _inputs={}, + from_source=ConversationFromSource.CONSOLE, + from_account_id=account.id, + ) + session.add(conversation) + session.flush() + message = Message( + app_id=target.id, + conversation_id=conversation.id, + _inputs={}, + query="What is a trial?", + message={}, + message_unit_price=Decimal(0), + answer="A way to try an app.", + answer_unit_price=Decimal(0), + currency="USD", + from_source=ConversationFromSource.CONSOLE, + from_account_id=account.id, + ) + session.add(message) + + read_sessions: list[Session] = [] + read_factory = sessionmaker(bind=sqlite_engine, expire_on_commit=False) + + @event.listens_for(read_factory, "after_begin") + def track_read(session: Session, _transaction: SessionTransaction, _connection: Connection) -> None: + read_sessions.append(session) + + provider = _Provider(read_sessions, target.id) + services = _TrialSuggestedQuestionsApplicationservices( + trial_apps=_TrialSuggestedQuestionsTrialappservices( + access=TrialAppAccessService(apps=TrialAppRepository(session_factory=read_factory)) + ), + recommended_app_queries=features, + message_suggested_questions=MessageSuggestedQuestionsRuntime(session_factory=read_factory), + ) + + def setup_completed() -> bool: + features.events.append("setup") + return features.setup_completed + + def csrf(_request: Request, account_id: str) -> None: + assert account_id == account.id + features.events.append("csrf") + + for module in (trial_module, _trial_suggested_questions_admission_module): + monkeypatch.setattr(module, "application_services", lambda: services) + monkeypatch.setattr(console_wraps, "_is_setup_completed", setup_completed) + monkeypatch.setattr(login_module, "current_user", account) + monkeypatch.setattr(login_module, "check_csrf_token", csrf) + + def model_manager(*, tenant_id: str) -> _Provider: + assert tenant_id == target.tenant_id + return provider + + def trace_manager(*, app_id: str) -> _Provider: + assert app_id == target.id + return provider + + monkeypatch.setattr(message_module.ModelManager, "for_tenant", model_manager) + monkeypatch.setattr(message_module, "TraceQueueManager", trace_manager) + + app = Flask(__name__) + app.config.update(TESTING=True, RESTX_ERROR_404_HELP=False, SQLALCHEMY_DATABASE_URI=str(sqlite_engine.url)) + db.init_app(app) + legacy_sessions: list[Session] = [] + + def track_legacy(session: Session, _transaction: SessionTransaction, _connection: Connection) -> None: + legacy_sessions.append(session) + + event.listen(db.session.session_factory, "after_begin", track_legacy) + login_manager = DifyLoginManager() + login_manager.init_app(app) + login_manager.unauthorized_handler(unauthorized_handler) + api = ExternalApi(app) + api.add_resource( + trial_module.TrialMessageSuggestedQuestionApi, + "/trial-apps//messages//suggested-questions", + ) + yield _TrialSuggestedQuestionsHarness( + app, + account, + target, + trial, + conversation, + message, + config, + sqlite_session_factory, + features, + provider, + legacy_sessions, + ) + event.remove(db.session.session_factory, "after_begin", track_legacy) + with app.app_context(): + db.session.remove() + db.engine.dispose() + + @staticmethod + @pytest.mark.parametrize( + "mode", + [ + _trial_suggested_questions_AppMode.CHAT, + _trial_suggested_questions_AppMode.AGENT_CHAT, + _trial_suggested_questions_AppMode.AGENT, + ], + ) + @pytest.mark.parametrize("questions", [["Next question?", "Another question?"], []]) + def test_questions_keep_response_history_owner_and_usage( + harness: _TrialSuggestedQuestionsHarness, questions: list[str], mode: _trial_suggested_questions_AppMode + ) -> None: + harness.provider.questions = questions + harness.provider.app_type = { + _trial_suggested_questions_AppMode.CHAT: "chatbot", + _trial_suggested_questions_AppMode.AGENT_CHAT: "agent", + _trial_suggested_questions_AppMode.AGENT: "agent_v2", + }[mode] + with harness.factory.begin() as session: + session.execute( + update(_trial_suggested_questions_App) + .where(_trial_suggested_questions_App.id == harness.target.id) + .values(mode=mode) + ) + session.execute(update(Conversation).where(Conversation.id == harness.conversation.id).values(mode=mode)) + with harness.factory.begin() as session: + session.add_all( + [ + AccountTrialAppRecord(app_id=harness.target.id, account_id=harness.account.id, count=1), + AccountTrialAppRecord(app_id=harness.target.id, account_id=str(uuid4()), count=3), + ] + ) + + response = harness.get() + + assert response.status_code == 200, response.get_json() + assert response.get_json() == {"data": questions} + assert response.headers["Content-Type"] == "application/json" + assert int(response.headers["Content-Length"]) == len(response.data) + assert len(harness.provider.prompts) == 1 + assert "Human: What is a trial?\nAssistant: A way to try an app." in harness.provider.prompts[0] + assert harness.provider.tenant_ids == [harness.target.tenant_id] * 2 + assert harness.target.tenant_id not in {harness.trial.tenant_id, harness.account.current_tenant_id} + assert len(harness.provider.traces) == 1 + assert harness.features.events == ["setup", "csrf", "feature"] + assert harness.usage() == 1 + assert len(harness.provider.model_sessions) == 2 + harness.assert_closed() + + @staticmethod + @pytest.mark.parametrize("entity", ["message", "conversation"]) + @pytest.mark.parametrize("mismatch", ["missing", "app", "account", "source", "end-user"]) + def test_message_and_conversation_require_complete_owner_chain( + harness: _TrialSuggestedQuestionsHarness, entity: str, mismatch: str + ) -> None: + message_id = harness.message.id + with harness.factory.begin() as session: + record = ( + session.get(Message, harness.message.id) + if entity == "message" + else session.get(Conversation, harness.conversation.id) + ) + assert record is not None + match mismatch: + case "missing": + if entity == "message": + message_id = str(uuid4()) + else: + session.delete(record) + case "app": + record.app_id = str(uuid4()) + case "account": + record.from_account_id = str(uuid4()) + case "source": + record.from_source = ConversationFromSource.API + case "end-user": + record.from_end_user_id = str(uuid4()) + + response = harness.get(message_id=message_id) + + _trial_suggested_questions_assert_error( + response, 404, "not_found", "Message not found" if entity == "message" else "Conversation not found" + ) + assert harness.provider.prompts == [] + assert harness.usage() is None + harness.assert_closed() + + @staticmethod + def test_deleted_conversation_is_not_visible(harness: _TrialSuggestedQuestionsHarness) -> None: + with harness.factory.begin() as session: + session.execute( + update(Conversation).where(Conversation.id == harness.conversation.id).values(is_deleted=True) + ) + _trial_suggested_questions_assert_error(harness.get(), 404, "not_found", "Conversation not found") + assert harness.provider.prompts == [] + harness.assert_closed() + + @staticmethod + @pytest.mark.parametrize("config", [None, '{"enabled":false}']) + def test_disabled_or_unset_suggestions_have_specific_error( + harness: _TrialSuggestedQuestionsHarness, config: str | None + ) -> None: + with harness.factory.begin() as session: + session.execute( + update(AppModelConfig) + .where(AppModelConfig.id == harness.config.id) + .values(suggested_questions_after_answer=config) + ) + _trial_suggested_questions_assert_error(harness.get(), 403, "app_suggested_questions_after_answer_disabled") + assert harness.provider.prompts == [] + harness.assert_closed() + + @staticmethod + def test_absent_default_model_keeps_empty_success(harness: _TrialSuggestedQuestionsHarness) -> None: + harness.provider.missing_model = True + response = harness.get() + assert response.status_code == 200 + assert response.get_json() == {"data": []} + assert harness.provider.prompts == [] + assert harness.provider.traces == [] + harness.assert_closed() + + @staticmethod + @pytest.mark.parametrize( + "mode", [_trial_suggested_questions_AppMode.COMPLETION, _trial_suggested_questions_AppMode.WORKFLOW] + ) + def test_unsupported_app_modes_fail_before_runtime( + harness: _TrialSuggestedQuestionsHarness, mode: _trial_suggested_questions_AppMode + ) -> None: + with harness.factory.begin() as session: + session.execute( + update(_trial_suggested_questions_App) + .where(_trial_suggested_questions_App.id == harness.target.id) + .values(mode=mode) + ) + _trial_suggested_questions_assert_error(harness.get(), 400, "not_chat_app") + assert len(harness.provider.read_sessions) == 1 + assert harness.legacy_sessions == [] + harness.assert_closed() + + @staticmethod + @pytest.mark.parametrize( + ("denial", "status", "code"), + [ + ("setup", 401, "not_setup"), + ("feature", 403, "trial_app_feature_disabled"), + ("missing_trial", 403, "trial_app_not_allowed"), + ("missing_app", 403, "trial_app_not_allowed"), + ("quota", 403, "trial_app_limit_exceeded"), + ("uninitialized", 400, "account_not_initialized"), + ("unauthenticated", 401, "unauthorized"), + ], + ) + def test_admission_blocks_before_runtime( + harness: _TrialSuggestedQuestionsHarness, monkeypatch: pytest.MonkeyPatch, denial: str, status: int, code: str + ) -> None: + if denial == "setup": + harness.features.setup_completed = False + elif denial == "feature": + harness.features.enabled = False + elif denial == "uninitialized": + harness.account.status = AccountStatus.UNINITIALIZED + elif denial == "unauthenticated": + monkeypatch.setattr(login_module, "current_user", None) + else: + with harness.factory.begin() as session: + if denial == "missing_trial": + session.execute(delete(TrialApp).where(TrialApp.app_id == harness.target.id)) + elif denial == "missing_app": + session.execute( + delete(_trial_suggested_questions_App).where( + _trial_suggested_questions_App.id == harness.target.id + ) + ) + else: + session.add(AccountTrialAppRecord(app_id=harness.target.id, account_id=harness.account.id, count=3)) + response = harness.get() + if denial == "unauthenticated": + assert response.status_code == 401 + assert response.get_json() == {"code": "unauthorized", "message": "Unauthorized.", "status": 401} + else: + _trial_suggested_questions_assert_error(response, status, code) + assert harness.legacy_sessions == [] + assert harness.provider.prompts == [] + harness.assert_closed() + + @staticmethod + @pytest.mark.parametrize( + ("failure", "code", "message"), + [ + ( + ProviderTokenNotInitError("Missing tenant credentials"), + "provider_not_initialize", + "Missing tenant credentials", + ), + (QuotaExceededError(), "provider_quota_exceeded", None), + (ModelCurrentlyNotSupportError(), "model_currently_not_support", None), + (InvokeError("Provider unavailable"), "completion_request_error", "Provider unavailable"), + ], + ) + def test_token_count_failures_preserve_specific_errors_and_close_session( + harness: _TrialSuggestedQuestionsHarness, failure: Exception, code: str, message: str | None + ) -> None: + harness.provider.failure = failure + _trial_suggested_questions_assert_error(harness.get(), 400, code, message) + assert harness.provider.token_calls == 1 + assert harness.provider.prompts == [] + assert harness.provider.traces == [] + assert harness.usage() is None + harness.assert_closed() + + @staticmethod + @pytest.mark.parametrize("failure", [None, InvokeError("Provider unavailable")]) + def test_runtime_preserves_outer_request_session( + harness: _TrialSuggestedQuestionsHarness, failure: Exception | None + ) -> None: + harness.provider.failure = failure + with harness.app.app_context(): + outer_session = db.session() + app_model = outer_session.get(_trial_suggested_questions_App, harness.target.id) + assert app_model is not None + app_model.name = "Pending caller change" + + response = harness.get() + + assert response.status_code == (200 if failure is None else 400) + assert db.session() is outer_session + assert outer_session.in_transaction() + assert app_model in outer_session.dirty + assert app_model.name == "Pending caller change" + assert harness.provider.model_sessions + for inner_session in harness.provider.model_sessions: + assert inner_session is not outer_session + assert not inner_session.in_transaction() + assert not inner_session.identity_map + harness.assert_closed() + with harness.factory() as session: + saved_app = session.get(_trial_suggested_questions_App, harness.target.id) + assert saved_app is not None + assert saved_app.name == "Trial app" + + @staticmethod + def test_advanced_chat_without_published_workflow_keeps_empty_success( + harness: _TrialSuggestedQuestionsHarness, + ) -> None: + with harness.factory.begin() as session: + session.execute( + update(_trial_suggested_questions_App) + .where(_trial_suggested_questions_App.id == harness.target.id) + .values(mode=_trial_suggested_questions_AppMode.ADVANCED_CHAT) + ) + response = harness.get() + assert response.status_code == 200 + assert response.get_json() == {"data": []} + assert len(harness.provider.read_sessions) == 2 + assert harness.provider.prompts == [] + assert harness.usage() is None + harness.assert_closed() + + @staticmethod + @pytest.mark.parametrize( + ("failure", "status", "code", "message"), + [ + (AppDefinitionUnavailableError("App changed after admission"), 400, "app_unavailable", None), + ( + SuggestedQuestionsActorNotFoundError("Account disappeared"), + 401, + "unauthorized", + "Account no longer exists.", + ), + ], + ) + def test_reload_errors_have_explicit_http_mapping( + harness: _TrialSuggestedQuestionsHarness, + monkeypatch: pytest.MonkeyPatch, + failure: Exception, + status: int, + code: str, + message: str | None, + ) -> None: + def reject_reload( + _self: MessageSuggestedQuestionsRuntime, + *, + app_id: str, + app_owner_tenant_id: str, + expected_app_mode: str, + actor: SuggestedQuestionsActor, + message_id: str, + ) -> list[str]: + assert (app_id, app_owner_tenant_id, expected_app_mode, message_id) == ( + harness.target.id, + harness.target.tenant_id, + harness.target.mode, + harness.message.id, + ) + assert actor == SuggestedQuestionsAccount(account_id=harness.account.id, invoke_from="explore") + raise failure + + monkeypatch.setattr(MessageSuggestedQuestionsRuntime, "get_suggested_questions", reject_reload) + _trial_suggested_questions_assert_error(harness.get(), status, code, message) + assert harness.provider.prompts == [] + assert harness.usage() is None + harness.assert_closed() diff --git a/api/tests/unit_tests/controllers/console/explore/test_trial_agent_config.py b/api/tests/unit_tests/controllers/console/explore/test_trial_agent_config.py deleted file mode 100644 index 92d4e7b4699..00000000000 --- a/api/tests/unit_tests/controllers/console/explore/test_trial_agent_config.py +++ /dev/null @@ -1,348 +0,0 @@ -"""Exercise public template assets through real admission, snapshots, and file services.""" - -import io -import zipfile -from collections.abc import Callable -from dataclasses import dataclass -from datetime import datetime -from types import SimpleNamespace -from unittest.mock import Mock -from urllib.parse import parse_qs, urlsplit -from uuid import uuid4 - -import pytest -from flask import Flask -from sqlalchemy import func, select -from sqlalchemy.orm import Session, sessionmaker -from werkzeug.test import TestResponse - -from controllers.console.app import agent_config_inspector as inspector -from controllers.console.app import preview_admission -from controllers.console.explore import trial -from extensions.ext_storage import StorageType -from libs.external_api import ExternalApi -from models.agent import ( - Agent, - AgentConfigDraft, - AgentConfigDraftType, - AgentConfigRevision, - AgentConfigRevisionOperation, - AgentConfigSnapshot, - AgentScope, - AgentSource, -) -from models.agent_config_entities import AgentSoulConfig -from models.enums import CreatorUserRole -from models.model import App, AppMode, RecommendedApp, UploadFile -from models.tools import ToolFile -from repositories.app_preview_query_repository import AppPreviewQueryRepository -from repositories.recommended_app_catalog_repository import DatabaseRecommendedAppCatalogRepository -from services.agent_config_service import AgentConfigService -from services.app_preview_query_service import AppPreviewQueryService - - -@dataclass -class Harness: - app: Flask - factory: sessionmaker[Session] - source: App - agent: Agent - snapshot: AgentConfigSnapshot - upload_id: str - skill_id: str - reads: list[str] - - def url(self, suffix: str) -> str: - return f"/console/api/trial-apps/{self.source.id}/agent/config/{suffix}" - - def get(self, suffix: str, **query: str) -> TestResponse: - return self.app.test_client().get(self.url(suffix), query_string=query) - - -@pytest.fixture -def harness( - sqlite_session_factory: sessionmaker[Session], - monkeypatch: pytest.MonkeyPatch, - config_overrides: Callable[..., None], -) -> Harness: - config_overrides(SECRET_KEY="trial-config-test", FILES_URL="http://localhost") - tenant_id, agent_id, snapshot_id, account_id = (str(uuid4()) for _ in range(4)) - source = App( - id=str(uuid4()), - tenant_id=tenant_id, - name="Template", - mode=AppMode.AGENT, - is_public=True, - enable_site=False, - enable_api=False, - ) - agent = Agent( - id=agent_id, - tenant_id=tenant_id, - app_id=source.id, - name="Template", - scope=AgentScope.ROSTER, - source=AgentSource.AGENT_APP, - active_config_snapshot_id=snapshot_id, - ) - skill_id = str(uuid4()) - skill_zip = io.BytesIO() - with zipfile.ZipFile(skill_zip, "w") as archive: - archive.writestr("SKILL.md", "---\nname: research\ndescription: Research\n---\nPublished skill") - archive.writestr("scripts/run.py", "print('published')") - payloads = {"skill.zip": skill_zip.getvalue(), "guide.txt": b"Published attachment"} - upload = UploadFile( - tenant_id=tenant_id, - storage_type=StorageType.LOCAL, - key="guide.txt", - name="guide.txt", - size=len(payloads["guide.txt"]), - extension="txt", - mime_type="text/plain", - created_by=account_id, - created_by_role=CreatorUserRole.ACCOUNT, - created_at=datetime(2026, 1, 1), - used=False, - ) - snapshot = AgentConfigSnapshot( - id=snapshot_id, - tenant_id=tenant_id, - agent_id=agent_id, - version=1, - config_snapshot=AgentSoulConfig.model_validate( - { - "config_skills": [{"name": "research", "file_id": skill_id}], - "config_files": [{"name": "guide.txt", "file_kind": "upload_file", "file_id": upload.id}], - } - ), - ) - skill_file = ToolFile( - tenant_id=tenant_id, - user_id=account_id, - conversation_id=None, - file_key="skill.zip", - mimetype="application/zip", - name="research.zip", - size=len(payloads["skill.zip"]), - ) - skill_file.id = skill_id - with sqlite_session_factory.begin() as session: - session.add_all( - [ - source, - agent, - snapshot, - upload, - skill_file, - RecommendedApp( - app_id=source.id, - description={}, - copyright="", - privacy_policy="", - category="Agent", - categories=["Agent"], - custom_disclaimer="", - position=1, - is_listed=True, - language="en-US", - ), - AgentConfigRevision( - tenant_id=tenant_id, - agent_id=agent_id, - current_snapshot_id=snapshot_id, - revision=1, - operation=AgentConfigRevisionOperation.PUBLISH_DRAFT, - ), - AgentConfigDraft( - tenant_id=tenant_id, - agent_id=agent_id, - draft_type=AgentConfigDraftType.DRAFT, - draft_owner_key="", - base_snapshot_id=snapshot_id, - config_snapshot=AgentSoulConfig.model_validate( - { - "config_files": [{"name": "private.txt", "file_kind": "upload_file", "file_id": upload.id}], - } - ), - ), - ] - ) - - reads: list[str] = [] - - def load_once(key: str) -> bytes: - reads.append(key) - return payloads[key] - - monkeypatch.setattr("services.agent_config_service.storage", SimpleNamespace(load_once=load_once)) - catalog = DatabaseRecommendedAppCatalogRepository(sqlite_session_factory, redis=Mock()) - services = SimpleNamespace( - app_previews=AppPreviewQueryService( - apps=AppPreviewQueryRepository(session_factory=sqlite_session_factory), - is_previewable=catalog.contains, - ) - ) - monkeypatch.setattr(preview_admission, "application_services", lambda: services) - monkeypatch.setattr(inspector, "application_services", lambda: services) - monkeypatch.setattr(trial, "application_services", lambda: services) - monkeypatch.setattr(inspector, "_service", lambda: AgentConfigService(session_factory=sqlite_session_factory)) - app = Flask(__name__) - api = ExternalApi(app) - api.add_resource(trial.TrialAgentComposerApi, "/console/api/trial-apps//agent-composer") - routes = [ - (inspector.TrialAgentConfigSkillsApi, "skills"), - (inspector.TrialAgentConfigFilesApi, "files"), - (inspector.TrialAgentConfigSkillInspectApi, "skills//inspect"), - (inspector.TrialAgentConfigSkillDownloadApi, "skills//download"), - (inspector.TrialAgentConfigSkillFilePreviewApi, "skills//files/preview"), - (inspector.TrialAgentConfigSkillFileDownloadApi, "skills//files/download"), - (inspector.TrialAgentConfigFilePreviewApi, "files//preview"), - (inspector.TrialAgentConfigFileDownloadApi, "files//download"), - ] - prefix = "/console/api/trial-apps//agent/config/" - for resource, suffix in routes: - api.add_resource(resource, prefix + suffix) - api.add_resource( - inspector.TrialAgentConfigSkillFileContentApi, - prefix + "skills//files/content", - endpoint="console.trial_agent_config_skill_file_content", - ) - return Harness(app, sqlite_session_factory, source, agent, snapshot, upload.id, skill_id, reads) - - -def test_published_template_configuration_is_readable_without_workspace_membership(harness: Harness) -> None: - composer_url = f"/console/api/trial-apps/{harness.source.id}/agent-composer" - composer = harness.app.test_client().get(composer_url) - assert composer.status_code == 200 - assert composer.json is not None - assert composer.json["agent_soul"]["config_files"][0]["name"] == "guide.txt" - assert composer.json["active_config_snapshot"]["id"] == harness.snapshot.id - assert composer.json["draft"] is None - assert composer.json["save_options"] == [] - for kind, name in [("skills", "research"), ("files", "guide.txt")]: - response = harness.get(kind) - assert response.status_code == 200 - assert response.json is not None - assert response.json["items"][0]["name"] == name - assert response.json["config_version"] == {"id": harness.snapshot.id, "kind": "snapshot", "writable": False} - response = harness.get("skills/research/inspect") - assert response.status_code == 200 - assert response.json is not None - assert "Published skill" in response.json["skill_md"]["text"] - response = harness.get("skills/research/files/preview", path="scripts/run.py") - assert response.status_code == 200 - assert response.json is not None - assert response.json["text"] == "print('published')" - response = harness.get("files/guide.txt/preview", version_id=harness.snapshot.id) - assert response.status_code == 200 - assert response.json is not None - assert response.json["text"] == "Published attachment" - assert response.headers["Cache-Control"] == "no-store" - for path in ["files/guide.txt/download", "skills/research/download"]: - response = harness.get(path) - assert response.status_code == 200 - assert response.json is not None - assert {"timestamp", "nonce", "sign"} <= parse_qs(urlsplit(response.json["url"]).query).keys() - response = harness.get("skills/research/files/download", path="scripts/run.py") - assert response.json is not None - url = response.json["url"] - download = harness.app.test_client().get(url) - assert download.status_code == 200 - assert download.data == b"print('published')" - assert download.headers["Content-Disposition"].startswith("attachment;") - with harness.factory() as session: - assert session.scalar(select(func.count()).select_from(AgentConfigDraft)) == 1 - assert session.scalar(select(func.count()).select_from(UploadFile)) == 1 - with harness.factory.begin() as session: - app = session.get(App, harness.source.id) - assert app is not None - app.is_public = False - reads_before = list(harness.reads) - assert harness.app.test_client().get(url).status_code == 404 - assert harness.app.test_client().get(composer_url).status_code == 404 - assert harness.reads == reads_before - - -@pytest.mark.parametrize( - ("query", "status"), - [ - ({"draft_type": "draft"}, 422), - ({"file_id": "other"}, 422), - ({"version_id": str(uuid4())}, 404), - ], -) -def test_version_and_resource_selectors_cannot_bypass_admission( - harness: Harness, - query: dict[str, str], - status: int, -) -> None: - assert harness.get("files/guide.txt/preview", **query).status_code == status - assert harness.reads == [] - - -@pytest.mark.parametrize("path", ["../SKILL.md", "scripts/../../SKILL.md"]) -def test_skill_member_paths_cannot_escape_archive(harness: Harness, path: str) -> None: - assert harness.get("skills/research/files/preview", path=path).status_code == 400 - assert harness.reads == [] - - -@pytest.mark.parametrize("kind", ["file", "skill"]) -def test_foreign_tenant_payloads_are_not_read(harness: Harness, kind: str) -> None: - with harness.factory.begin() as session: - if kind == "file": - upload = session.get(UploadFile, harness.upload_id) - assert upload is not None - upload.tenant_id = str(uuid4()) - else: - skill = session.get(ToolFile, harness.skill_id) - assert skill is not None - skill.tenant_id = str(uuid4()) - suffix = "files/guide.txt/preview" if kind == "file" else "skills/research/files/preview" - assert harness.get(suffix, **({"path": "SKILL.md"} if kind == "skill" else {})).status_code == 404 - assert harness.reads == [] - - -def test_draft_only_files_and_writes_are_unavailable(harness: Harness) -> None: - assert harness.get("files/private.txt/preview").status_code == 404 - for method in ["POST", "PUT", "PATCH", "DELETE"]: - assert harness.app.test_client().open(harness.url("files"), method=method).status_code == 405 - assert harness.reads == [] - - -@pytest.mark.parametrize( - ("state", "status"), - [ - ("unlisted", 404), - ("unpublished", 400), - ("foreign_snapshot", 400), - ], -) -def test_unavailable_template_never_reads_storage(harness: Harness, state: str, status: int) -> None: - with harness.factory.begin() as session: - if state == "unlisted": - listing = session.scalar(select(RecommendedApp).where(RecommendedApp.app_id == harness.source.id)) - assert listing is not None - listing.is_listed = False - elif state == "unpublished": - revision = session.scalar( - select(AgentConfigRevision).where(AgentConfigRevision.agent_id == harness.agent.id) - ) - assert revision is not None - revision.operation = AgentConfigRevisionOperation.IMPORT_PACKAGE - else: - snapshot = session.get(AgentConfigSnapshot, harness.snapshot.id) - assert snapshot is not None - snapshot.tenant_id = str(uuid4()) - preview = harness.get("files/guide.txt/preview") - assert preview.status_code == status - assert preview.json is not None - assert preview.json["status"] == status - if status == 400: - # Raised inside the template guard, so its error answer is built by the surface and not cached. - assert preview.headers["Cache-Control"] == "no-store" - composer = harness.app.test_client().get(f"/console/api/trial-apps/{harness.source.id}/agent-composer") - assert composer.status_code == status - if status == 400: - assert composer.json is not None - assert composer.json["code"] == "app_unavailable" - assert harness.reads == [] diff --git a/api/tests/unit_tests/controllers/console/explore/test_trial_audio.py b/api/tests/unit_tests/controllers/console/explore/test_trial_audio.py deleted file mode 100644 index a818b3d4a7d..00000000000 --- a/api/tests/unit_tests/controllers/console/explore/test_trial_audio.py +++ /dev/null @@ -1,531 +0,0 @@ -"""Trial audio admission, response and usage contracts through HTTP and SQLite.""" - -from collections.abc import Callable, Iterator, Mapping -from dataclasses import dataclass, field -from io import BytesIO -from typing import override -from uuid import uuid4 - -import pytest -from flask import Flask, Request -from sqlalchemy import Connection, event, select -from sqlalchemy.orm import Mapper, Session, SessionTransaction, sessionmaker -from werkzeug.exceptions import Forbidden, Unauthorized -from werkzeug.test import TestResponse - -import controllers.console.explore.trial as trial_module -import controllers.console.explore.trial_app_admission as admission_module -import controllers.console.wraps as console_wraps -import libs.login as login_module -from controllers.common.errors import UnauthorizedError -from core.errors.error import ModelCurrentlyNotSupportError, ProviderTokenNotInitError, QuotaExceededError -from enums import DeploymentEdition -from extensions.ext_login import DifyLoginManager, unauthorized_handler -from graphon.model_runtime.errors.invoke import InvokeBadRequestError, InvokeError -from libs.external_api import ExternalApi -from models import Account, AccountTrialAppRecord, App, AppMode, Tenant, TrialApp -from models.account import AccountStatus -from repositories.trial_app_repository import TrialAppRepository -from services.agent.errors import AgentVersionNotFoundError -from services.app_definition_query_service import AppDefinitionUnavailableError -from services.audio_types import AudioAppRef, AudioOutput, AudioUpload -from services.errors.app_model_config import AppModelConfigBrokenError -from services.errors.audio import ( - AudioTooLargeServiceError, - NoAudioUploadedServiceError, - ProviderNotSupportSpeechToTextServiceError, - ProviderNotSupportTextToSpeechServiceError, - SpeechToTextDisabledServiceError, - UnsupportedAudioTypeServiceError, -) -from services.trial_app_access_service import TrialAppAccessService - -_ASR = "audio-to-text" -_TTS = "text-to-audio" -_WAV = b"RIFF\x24\x00\x00\x00WAVEfmt " + b"a" * 32 - - -@dataclass -class _FeatureState: - enabled: bool = True - setup_completed: bool = True - events: list[str] = field(default_factory=list) - - def is_trial_enabled(self) -> bool: - self.events.append("feature") - return self.enabled - - -@dataclass(frozen=True) -class _AsrCall: - app: AudioAppRef - content: bytes | None - mime_type: str | None - - -@dataclass(frozen=True) -class _TtsCall: - app: AudioAppRef - account_id: str - text: str | None - voice: str | None - message_id: str | None - - -@dataclass -class _Stream(Iterator[bytes]): - chunks: Iterator[bytes] = field(default_factory=lambda: iter([b"a" * 32, b"remaining-audio"])) - reads: int = 0 - close_calls: int = 0 - close_error: Exception | None = None - - @override - def __next__(self) -> bytes: - chunk = next(self.chunks) - self.reads += 1 - return chunk - - def close(self) -> None: - self.close_calls += 1 - if self.close_error is not None: - raise self.close_error - - -@dataclass -class _Runtime: - sessions: list[Session] - output: AudioOutput | None = field(default_factory=lambda: AudioOutput(data=_WAV, mime_type=None)) - error: Exception | None = None - calls: list[_AsrCall | _TtsCall] = field(default_factory=list) - - def _assert_admission_closed(self) -> None: - assert self.sessions - assert all(not session.in_transaction() and not session.identity_map for session in self.sessions) - - def transcript_asr(self, *, app: AudioAppRef, audio: AudioUpload | None) -> dict[str, str]: - self._assert_admission_closed() - self.calls.append(_AsrCall(app, audio.stream.read() if audio else None, audio.mime_type if audio else None)) - if self.error is not None: - raise self.error - if audio is None: - raise NoAudioUploadedServiceError() - return {"text": " transcript "} - - def transcript_tts( - self, - *, - app: AudioAppRef, - account_id: str, - text: str | None, - voice: str | None, - message_id: str | None, - ) -> AudioOutput | None: - self._assert_admission_closed() - self.calls.append(_TtsCall(app, account_id, text, voice, message_id)) - if self.error is not None: - raise self.error - return self.output - - -@dataclass(frozen=True) -class _TrialAppServices: - access: TrialAppAccessService - usage: TrialAppRepository - - -@dataclass(frozen=True) -class _ApplicationServices: - trial_apps: _TrialAppServices - app_audio: _Runtime - recommended_app_queries: _FeatureState - - -@dataclass(frozen=True) -class _Harness: - app: Flask - account: Account - target: App - trial: TrialApp - factory: sessionmaker[Session] - runtime: _Runtime - state: _FeatureState - - def post(self, endpoint: str, *, payload: Mapping[str, object] | None = None) -> TestResponse: - client = self.app.test_client() - url = f"/trial-apps/{self.target.id}/{endpoint}" - if endpoint == _ASR: - return client.post(url, data={"file": (BytesIO(b"audio input"), "recording.m4a", "audio/x-m4a")}) - return client.post(url, json=dict(payload) if payload is not None else {"text": "Hello"}) - - def usage(self) -> int | None: - with self.factory() as session: - return session.scalar( - select(AccountTrialAppRecord.count).where( - AccountTrialAppRecord.app_id == self.target.id, - AccountTrialAppRecord.account_id == self.account.id, - ) - ) - - -@pytest.fixture -def harness( - monkeypatch: pytest.MonkeyPatch, - config_overrides: Callable[..., None], - sqlite_session_factory: sessionmaker[Session], -) -> _Harness: - config_overrides(DEPLOYMENT_EDITION=DeploymentEdition.COMMUNITY, INIT_PASSWORD="", LOGIN_DISABLED=False) - state = _FeatureState() - account = Account(name="Viewer", email="viewer@example.com") - account.status = AccountStatus.ACTIVE - account._current_tenant = Tenant(name="Viewer workspace") - target = App( - id=str(uuid4()), tenant_id=str(uuid4()), name="Trial app", mode=AppMode.CHAT, enable_site=True, enable_api=True - ) - trial = TrialApp(app_id=target.id, tenant_id=target.tenant_id, trial_limit=3) - with sqlite_session_factory.begin() as session: - session.add_all([target, trial]) - - sessions: list[Session] = [] - repository_factory = sessionmaker(bind=sqlite_session_factory.kw["bind"], expire_on_commit=False) - - @event.listens_for(repository_factory, "after_begin") - def track_session(session: Session, _transaction: SessionTransaction, _connection: Connection) -> None: - sessions.append(session) - - runtime = _Runtime(sessions=sessions) - repository = TrialAppRepository(session_factory=repository_factory) - services = _ApplicationServices( - trial_apps=_TrialAppServices(access=TrialAppAccessService(apps=repository), usage=repository), - app_audio=runtime, - recommended_app_queries=state, - ) - - def setup_completed() -> bool: - state.events.append("setup") - return state.setup_completed - - def csrf(_request: Request, account_id: str) -> None: - assert account_id == account.id - state.events.append("csrf") - - for module in (trial_module, admission_module): - monkeypatch.setattr(module, "application_services", lambda: services) - monkeypatch.setattr(console_wraps, "_is_setup_completed", setup_completed) - monkeypatch.setattr(login_module, "current_user", account) - monkeypatch.setattr(login_module, "check_csrf_token", csrf) - app = Flask(__name__) - app.config.update(TESTING=True, RESTX_ERROR_404_HELP=False) - login_manager = DifyLoginManager() - login_manager.init_app(app) - login_manager.unauthorized_handler(unauthorized_handler) - api = ExternalApi(app) - api.add_resource(trial_module.TrialChatAudioApi, "/trial-apps//audio-to-text") - api.add_resource(trial_module.TrialChatTextApi, "/trial-apps//text-to-audio") - return _Harness(app, account, target, trial, sqlite_session_factory, runtime, state) - - -def _assert_error(response: TestResponse, status: int, code: str) -> Mapping[str, object]: - assert response.status_code == status - body = response.get_json() - assert isinstance(body, dict) - assert body["code"] == code - assert body["status"] == status - assert response.headers["Content-Type"] == "application/json" - assert isinstance(body["message"], str) - assert body["message"] - return body - - -def test_asr_upload_keeps_transcript_and_owner_tenant(harness: _Harness) -> None: - response = harness.post(_ASR) - - assert response.status_code == 200 - assert response.get_json() == {"text": " transcript "} - assert response.headers["Content-Type"] == "application/json" - assert int(response.headers["Content-Length"]) == len(response.data) - assert harness.runtime.calls == [ - _AsrCall(AudioAppRef(harness.target.id, harness.target.tenant_id, "chat"), b"audio input", "audio/x-m4a") - ] - assert harness.target.tenant_id != harness.account.current_tenant_id - assert harness.state.events == ["setup", "csrf", "feature"] - assert harness.usage() == 1 - - -@pytest.mark.parametrize("streaming", [None, False, True]) -def test_tts_preserves_payload_actor_and_binary_response_ignoring_streaming_flag( - harness: _Harness, streaming: bool | None -) -> None: - message_id = str(uuid4()) - response = harness.post( - _TTS, payload={"text": " hello ", "voice": "voice-1", "message_id": message_id, "streaming": streaming} - ) - - assert response.status_code == 200 - assert response.data == _WAV - assert response.headers["Content-Type"] == "audio/wav" - assert int(response.headers["Content-Length"]) == len(_WAV) - assert harness.runtime.calls == [ - _TtsCall( - AudioAppRef(harness.target.id, harness.target.tenant_id, "chat"), - harness.account.id, - " hello ", - "voice-1", - message_id, - ) - ] - assert harness.usage() == 1 - - -@pytest.mark.parametrize("null_fields", [False, True]) -def test_missing_message_null_response_still_records_usage(harness: _Harness, null_fields: bool) -> None: - harness.runtime.output = None - payload: dict[str, object] = {"text": None, "voice": None, "message_id": None} if null_fields else {} - response = harness.post(_TTS, payload=payload) - - assert response.status_code == 200 - assert response.get_json() is None - assert response.headers["Content-Type"] == "application/json" - assert harness.runtime.calls == [ - _TtsCall(AudioAppRef(harness.target.id, harness.target.tenant_id, "chat"), harness.account.id, None, None, None) - ] - assert harness.usage() == 1 - - -def test_missing_upload_passes_none_and_does_not_record_usage(harness: _Harness) -> None: - response = harness.app.test_client().post(f"/trial-apps/{harness.target.id}/{_ASR}", data={}) - - _assert_error(response, 400, "no_audio_uploaded") - assert harness.runtime.calls == [ - _AsrCall(AudioAppRef(harness.target.id, harness.target.tenant_id, "chat"), None, None) - ] - assert harness.usage() is None - - -def test_oversized_multipart_request_keeps_transport_error(harness: _Harness) -> None: - harness.app.config["MAX_CONTENT_LENGTH"] = 1 - - response = harness.post(_ASR) - - _assert_error(response, 413, "request_entity_too_large") - assert harness.runtime.calls == [] - assert harness.usage() is None - - -@pytest.mark.parametrize("endpoint", [_ASR, _TTS]) -@pytest.mark.parametrize( - ("error", "status", "code"), - [ - (AppDefinitionUnavailableError("app removed"), 400, "app_unavailable"), - (AppModelConfigBrokenError(), 400, "app_unavailable"), - (ProviderTokenNotInitError("provider key missing"), 400, "provider_not_initialize"), - (QuotaExceededError(), 400, "provider_quota_exceeded"), - (ModelCurrentlyNotSupportError(), 400, "model_currently_not_support"), - (InvokeError("provider rejected audio"), 400, "completion_request_error"), - (ValueError("Text is required"), 400, "invalid_param"), - (UnauthorizedError("session expired"), 401, "unauthorized"), - (AgentVersionNotFoundError(), 404, "agent_version_not_found_error"), - (RuntimeError("private provider failure"), 500, "internal_server_error"), - (Forbidden("private provider failure"), 500, "internal_server_error"), - ], -) -def test_runtime_errors_keep_specific_codes_without_recording_usage( - harness: _Harness, endpoint: str, error: Exception, status: int, code: str -) -> None: - harness.runtime.error = error - response = harness.post(endpoint) - - body = _assert_error(response, status, code) - if isinstance(error, (ProviderTokenNotInitError, InvokeError)): - assert body["message"] == error.description - elif code in {"audio_too_large", "invalid_param"}: - assert str(error) in str(body["message"]) - elif status == 401: - assert body["message"] == "session expired" - assert response.headers["WWW-Authenticate"] == 'Bearer realm="api"' - elif status == 404: - assert body["message"] == "Agent config version not found." - elif status == 500: - assert body["message"] == ( - "The server encountered an internal error and was unable to complete your request. " - "Either the server is overloaded or there is an error in the application." - ) - assert harness.usage() is None - - -@pytest.mark.parametrize( - ("error", "status", "code"), - [ - (NoAudioUploadedServiceError(), 400, "no_audio_uploaded"), - (AudioTooLargeServiceError("Audio size larger than 30 mb"), 413, "audio_too_large"), - (UnsupportedAudioTypeServiceError(), 415, "unsupported_audio_type"), - (ProviderNotSupportSpeechToTextServiceError(), 400, "provider_not_support_speech_to_text"), - (SpeechToTextDisabledServiceError(), 400, "speech_to_text_disabled"), - ], -) -def test_transcription_errors_keep_specific_codes_without_recording_usage( - harness: _Harness, error: Exception, status: int, code: str -) -> None: - # Only transcription validates the upload and needs a speech-to-text model; synthesis documents none of these. - harness.runtime.error = error - body = _assert_error(harness.post(_ASR), status, code) - if code == "audio_too_large": - assert str(error) in str(body["message"]) - assert harness.usage() is None - - -def test_unsupported_tts_provider_has_specific_error(harness: _Harness) -> None: - harness.runtime.error = ProviderNotSupportTextToSpeechServiceError() - - _assert_error(harness.post(_TTS), 400, "provider_not_support_text_to_speech") - assert harness.usage() is None - - -@pytest.mark.parametrize("field", ["text", "voice", "message_id", "streaming"]) -def test_invalid_tts_payload_keeps_field_context(harness: _Harness, field: str) -> None: - response = harness.post(_TTS, payload={field: ["invalid"]}) - - body = _assert_error(response, 422, "unprocessable_entity") - assert field in str(body) - assert harness.runtime.calls == [] - assert harness.usage() is None - - -@pytest.mark.parametrize("endpoint", [_ASR, _TTS]) -@pytest.mark.parametrize( - ("gate", "status", "code"), - [ - ("setup", 401, "not_setup"), - ("login", 401, "unauthorized"), - ("csrf", 401, "unauthorized"), - ("initialization", 400, "account_not_initialized"), - ("feature", 403, "trial_app_feature_disabled"), - ], -) -def test_console_and_trial_gates_run_before_app_lookup( - harness: _Harness, monkeypatch: pytest.MonkeyPatch, endpoint: str, gate: str, status: int, code: str -) -> None: - if gate == "setup": - harness.state.setup_completed = False - elif gate == "login": - monkeypatch.setattr(login_module, "current_user", None) - elif gate == "csrf": - - def reject_csrf(_request: Request, _account_id: str) -> None: - raise Unauthorized("invalid csrf token") - - monkeypatch.setattr(login_module, "check_csrf_token", reject_csrf) - elif gate == "initialization": - harness.account.status = AccountStatus.UNINITIALIZED - else: - harness.state.enabled = False - response = harness.post(endpoint) - - if gate == "login": - assert response.status_code == 401 - assert response.get_json() == {"code": "unauthorized", "message": "Unauthorized.", "status": 401} - assert response.headers["Content-Type"] == "application/json" - else: - _assert_error(response, status, code) - assert harness.runtime.calls == [] - assert harness.runtime.sessions == [] - assert harness.usage() is None - - -@pytest.mark.parametrize("endpoint", [_ASR, _TTS]) -@pytest.mark.parametrize("gate", ["membership", "missing-app", "quota"]) -def test_trial_membership_and_account_quota_block_audio(harness: _Harness, endpoint: str, gate: str) -> None: - with harness.factory.begin() as session: - if gate == "membership": - trial = session.get(TrialApp, harness.trial.id) - assert trial is not None - session.delete(trial) - elif gate == "missing-app": - app = session.get(App, harness.target.id) - assert app is not None - session.delete(app) - else: - session.add(AccountTrialAppRecord(app_id=harness.target.id, account_id=harness.account.id, count=3)) - response = harness.post(endpoint) - - _assert_error(response, 403, "trial_app_limit_exceeded" if gate == "quota" else "trial_app_not_allowed") - assert harness.runtime.calls == [] - assert harness.usage() == (3 if gate == "quota" else None) - - -@pytest.mark.parametrize("endpoint", [_ASR, _TTS]) -def test_usage_is_isolated_to_current_account(harness: _Harness, endpoint: str) -> None: - other_id = str(uuid4()) - with harness.factory.begin() as session: - session.add_all( - [ - AccountTrialAppRecord(app_id=harness.target.id, account_id=harness.account.id, count=1), - AccountTrialAppRecord(app_id=harness.target.id, account_id=other_id, count=99), - ] - ) - response = harness.post(endpoint) - - assert response.status_code == 200 - assert harness.usage() == 2 - with harness.factory() as session: - assert ( - session.scalar(select(AccountTrialAppRecord.count).where(AccountTrialAppRecord.account_id == other_id)) - == 99 - ) - - -@pytest.mark.parametrize("failure", ["mime", "first-chunk"]) -def test_stream_preparation_failure_closes_provider_without_consuming_trial(harness: _Harness, failure: str) -> None: - def broken_chunks() -> Iterator[bytes]: - raise InvokeBadRequestError("first audio chunk failed") - yield b"unreachable" - - stream = _Stream(chunks=iter([_WAV]) if failure == "mime" else broken_chunks()) - harness.runtime.output = AudioOutput(data=stream, mime_type="audio/mpeg") - response = harness.post(_TTS) - - body = _assert_error(response, 400, "completion_request_error") - assert ("MIME" if failure == "mime" else "first audio chunk failed") in str(body["message"]) - assert stream.close_calls == 1 - assert harness.usage() is None - - -@pytest.mark.parametrize("cleanup_fails", [False, True]) -def test_usage_write_failure_closes_stream_without_hiding_database_error( - harness: _Harness, cleanup_fails: bool, caplog: pytest.LogCaptureFixture -) -> None: - stream = _Stream(close_error=RuntimeError("cleanup failed") if cleanup_fails else None) - harness.runtime.output = AudioOutput(data=stream, mime_type="audio/mpeg") - failure = RuntimeError("usage database unavailable") - - def reject_usage_insert( - _mapper: Mapper[AccountTrialAppRecord], _connection: Connection, _record: AccountTrialAppRecord - ) -> None: - raise failure - - event.listen(AccountTrialAppRecord, "before_insert", reject_usage_insert) - try: - response = harness.post(_TTS) - finally: - event.remove(AccountTrialAppRecord, "before_insert", reject_usage_insert) - - _assert_error(response, 500, "internal_server_error") - assert stream.reads == 1 - assert stream.close_calls == 1 - assert harness.usage() is None - assert any(record.exc_info is not None and record.exc_info[1] is failure for record in caplog.records) - - -def test_partial_http_consumption_closes_provider_stream(harness: _Harness) -> None: - stream = _Stream() - harness.runtime.output = AudioOutput(data=stream, mime_type="audio/mpeg") - response = harness.post(_TTS) - - assert response.status_code == 200 - assert response.headers["Content-Type"] == "audio/mpeg" - assert "Content-Length" not in response.headers - assert next(iter(response.response)) == b"a" * 32 - assert stream.reads == 1 - assert stream.close_calls == 0 - response.close() - assert stream.close_calls == 1 - assert harness.usage() == 1 diff --git a/api/tests/unit_tests/controllers/console/explore/test_trial_file_upload.py b/api/tests/unit_tests/controllers/console/explore/test_trial_file_upload.py deleted file mode 100644 index 402c1526608..00000000000 --- a/api/tests/unit_tests/controllers/console/explore/test_trial_file_upload.py +++ /dev/null @@ -1,566 +0,0 @@ -"""Trial uploads through HTTP, real file services, and SQLite admission/persistence.""" - -import hashlib -from collections.abc import Callable, Mapping -from dataclasses import dataclass, field -from datetime import datetime -from io import BytesIO -from typing import Literal -from urllib.parse import parse_qs, urlsplit -from uuid import UUID, uuid4 - -import httpx -import pytest -from flask import Flask, Request -from sqlalchemy import Connection, delete, event, select -from sqlalchemy.orm import Session, SessionTransaction, sessionmaker -from werkzeug.test import TestResponse - -import controllers.console.explore.trial as trial_module -import controllers.console.explore.trial_app_admission as admission_module -import controllers.console.files as files_module -import controllers.console.remote_files as remote_module -import controllers.console.wraps as console_wraps -import libs.login as login_module -import services.file_service as file_module -from core.file import remote_fetcher -from core.tools.errors import ToolSSRFError -from enums import DeploymentEdition -from extensions.ext_login import DifyLoginManager, unauthorized_handler -from libs.external_api import ExternalApi -from models import Account, AccountTrialAppRecord, App, AppMode, Tenant, TrialApp, UploadFile -from models.account import TenantAccountRole -from models.enums import CreatorUserRole -from repositories.trial_app_repository import TrialAppRepository -from services.entities.feature_entities import FeatureModel -from services.feature_service import FeatureService -from services.file_service import FileService -from services.remote_file_service import RemoteFileService -from services.trial_app_access_service import TrialAppAccessService - -_Endpoint = Literal["files/upload", "remote-files/upload"] -_CREATED_AT = datetime(2024, 1, 1) -_REMOTE_URL = "https://example.com/trial.txt" -_CONTENT = b"trial upload" - - -@dataclass -class _Features: - enabled: bool = True - events: list[str] = field(default_factory=list) - - def is_trial_enabled(self) -> bool: - self.events.append("feature") - return self.enabled - - -@dataclass -class _ExternalIO: - admission_sessions: list[Session] - saved: dict[str, bytes] = field(default_factory=dict) - requests: list[tuple[str, str, Mapping[str, object]]] = field(default_factory=list) - head_status: int = 200 - get_status: int = 200 - content_length: str = str(len(_CONTENT)) - error: Exception | None = None - - def assert_admission_closed(self) -> None: - assert self.admission_sessions - assert all(not session.in_transaction() and not session.identity_map for session in self.admission_sessions) - - def save(self, key: str, content: bytes) -> None: - self.assert_admission_closed() - self.saved[key] = content - - def request(self, method: str, *, url: str, **kwargs: object) -> httpx.Response: - self.assert_admission_closed() - self.requests.append((method, url, kwargs)) - if self.error is not None: - raise self.error - return httpx.Response( - self.head_status if method == "HEAD" else self.get_status, - headers={"Content-Type": "text/plain", "Content-Length": self.content_length}, - content=_CONTENT if method == "GET" else b"", - request=httpx.Request(method, url), - ) - - -@dataclass(frozen=True) -class _TrialAppServices: - access: TrialAppAccessService - - -@dataclass(frozen=True) -class _ApplicationServices: - trial_apps: _TrialAppServices - recommended_app_queries: _Features - files: FileService - remote_files: RemoteFileService - - -@dataclass(frozen=True) -class _Harness: - app: Flask - account: Account - target: App - trial: TrialApp - factory: sessionmaker[Session] - features: _Features - io: _ExternalIO - - def url(self, endpoint: _Endpoint) -> str: - return f"/trial-apps/{self.target.id}/{endpoint}" - - def post(self, endpoint: _Endpoint, *, query: str = "") -> TestResponse: - if endpoint == "remote-files/upload": - return self.app.test_client().post(self.url(endpoint) + query, json={"url": _REMOTE_URL}) - return self.app.test_client().post( - self.url(endpoint) + query, - data={"file": (BytesIO(_CONTENT), "trial.txt", "text/plain")}, - ) - - def uploaded(self) -> list[UploadFile]: - with self.factory() as session: - return list(session.scalars(select(UploadFile))) - - def usage(self) -> int | None: - with self.factory() as session: - return session.scalar( - select(AccountTrialAppRecord.count).where( - AccountTrialAppRecord.app_id == self.target.id, - AccountTrialAppRecord.account_id == self.account.id, - ) - ) - - -@pytest.fixture -def harness( - monkeypatch: pytest.MonkeyPatch, - config_overrides: Callable[..., None], - sqlite_session_factory: sessionmaker[Session], -) -> _Harness: - config_overrides( - DEPLOYMENT_EDITION=DeploymentEdition.COMMUNITY, - INIT_PASSWORD="", - LOGIN_DISABLED=False, - RBAC_ENABLED=False, - UPLOAD_FILE_SIZE_LIMIT=1, - inner_UPLOAD_FILE_EXTENSION_BLACKLIST="exe", - ) - features = _Features() - account = Account(name="Trial viewer", email="trial@example.com") - account._current_tenant = Tenant(name="Viewer workspace") - account.role = TenantAccountRole.OWNER - target = App(tenant_id=str(uuid4()), name="Trial app", mode=AppMode.CHAT, enable_site=True, enable_api=False) - target.id = str(uuid4()) - trial = TrialApp(app_id=target.id, tenant_id=str(uuid4()), trial_limit=3) - with sqlite_session_factory.begin() as session: - session.add_all([target, trial]) - - sessions: list[Session] = [] - repository_factory = sessionmaker(bind=sqlite_session_factory.kw["bind"], expire_on_commit=False) - - @event.listens_for(repository_factory, "after_begin") - def track_session(session: Session, _transaction: SessionTransaction, _connection: Connection) -> None: - sessions.append(session) - - io = _ExternalIO(sessions) - files = FileService(session_factory=sqlite_session_factory) - services = _ApplicationServices( - trial_apps=_TrialAppServices( - access=TrialAppAccessService(apps=TrialAppRepository(session_factory=repository_factory)) - ), - recommended_app_queries=features, - files=files, - remote_files=RemoteFileService(files=files), - ) - - def setup_completed() -> bool: - features.events.append("setup") - return True - - def csrf(_request: Request, account_id: str) -> None: - assert account_id == account.id - features.events.append("csrf") - - for module in (trial_module, admission_module, files_module, remote_module): - monkeypatch.setattr(module, "application_services", lambda: services) - monkeypatch.setattr(console_wraps, "_is_setup_completed", setup_completed) - monkeypatch.setattr(login_module, "current_user", account) - monkeypatch.setattr(login_module, "check_csrf_token", csrf) - monkeypatch.setattr(file_module.storage, "save", io.save) - monkeypatch.setattr(file_module, "naive_utc_now", lambda: _CREATED_AT) - monkeypatch.setattr(remote_fetcher, "make_request", io.request) - - app = Flask(__name__) - app.config.update(TESTING=True, RESTX_ERROR_404_HELP=False) - login_manager = DifyLoginManager() - login_manager.init_app(app) - login_manager.unauthorized_handler(unauthorized_handler) - api = ExternalApi(app) - api.add_resource(trial_module.TrialAppFileUploadApi, "/trial-apps//files/upload") - api.add_resource(trial_module.TrialAppRemoteFileUploadApi, "/trial-apps//remote-files/upload") - return _Harness(app, account, target, trial, sqlite_session_factory, features, io) - - -def _assert_error(response: TestResponse, status: int, code: str) -> None: - assert response.status_code == status, response.get_json() - body = response.get_json() - assert isinstance(body, dict) - assert body["code"] == code - assert body["status"] == status - assert isinstance(body["message"], str) - assert body["message"] - assert response.headers["Content-Type"] == "application/json" - - -def _assert_signed_url(url: object, file_id: str) -> None: - assert isinstance(url, str) - parsed = urlsplit(url) - assert parsed.path == f"/files/{file_id}/file-preview" - assert parse_qs(parsed.query).keys() == {"timestamp", "nonce", "sign"} - - -@pytest.mark.parametrize("content", [_CONTENT, b""]) -def test_local_upload_preserves_complete_response_and_owner_tenant(harness: _Harness, content: bytes) -> None: - response = harness.app.test_client().post( - harness.url("files/upload"), data={"file": (BytesIO(content), "trial.txt", "text/plain")} - ) - - assert response.status_code == 201 - assert response.headers["Content-Type"] == "application/json" - assert int(response.headers["Content-Length"]) == len(response.data) - body = response.get_json() - uploaded = harness.uploaded() - assert len(uploaded) == 1 - file = uploaded[0] - assert body == { - "id": file.id, - "reference": None, - "name": "trial.txt", - "size": len(content), - "extension": "txt", - "mime_type": "text/plain", - "created_by": harness.account.id, - "created_at": int(_CREATED_AT.timestamp()), - "preview_url": None, - "source_url": body["source_url"], - "original_url": None, - "user_id": None, - "tenant_id": harness.target.tenant_id, - "conversation_id": None, - "file_key": None, - } - assert str(UUID(file.id)) == file.id - _assert_signed_url(body["source_url"], file.id) - assert file.created_by_role == CreatorUserRole.ACCOUNT - assert file.created_by == harness.account.id - assert file.tenant_id == harness.target.tenant_id - assert file.tenant_id not in {harness.trial.tenant_id, harness.account.current_tenant_id} - assert file.key.startswith(f"upload_files/{file.tenant_id}/") - assert file.hash == hashlib.sha3_256(content).hexdigest() - assert file.used is False - assert harness.io.saved == {file.key: content} - assert harness.features.events == ["setup", "csrf", "feature"] - assert harness.usage() is None - - -@pytest.mark.parametrize("head_status", [200, 405]) -def test_remote_upload_fetches_once_and_preserves_response_and_actor(harness: _Harness, head_status: int) -> None: - harness.io.head_status = head_status - response = harness.post("remote-files/upload") - - assert response.status_code == 201 - assert response.headers["Content-Type"] == "application/json" - assert int(response.headers["Content-Length"]) == len(response.data) - body = response.get_json() - uploaded = harness.uploaded() - assert len(uploaded) == 1 - file = uploaded[0] - assert body == { - "id": file.id, - "name": "trial.txt", - "size": len(_CONTENT), - "extension": "txt", - "url": body["url"], - "mime_type": "text/plain", - "created_by": harness.account.id, - "created_at": int(_CREATED_AT.timestamp()), - } - _assert_signed_url(body["url"], file.id) - assert file.source_url == _REMOTE_URL - assert file.tenant_id == harness.target.tenant_id - assert file.tenant_id not in {harness.trial.tenant_id, harness.account.current_tenant_id} - assert file.created_by == harness.account.id - assert file.created_by_role == CreatorUserRole.ACCOUNT - assert harness.io.saved == {file.key: _CONTENT} - assert harness.io.requests == [ - ("HEAD", _REMOTE_URL, {}), - ("GET", _REMOTE_URL, {} if head_status == 200 else {"timeout": 3, "follow_redirects": True}), - ] - assert harness.usage() is None - - -@pytest.mark.parametrize("endpoint", ["files/upload", "remote-files/upload"]) -@pytest.mark.parametrize( - ("denial", "code"), - [ - ("feature", "trial_app_feature_disabled"), - ("missing_trial", "trial_app_not_allowed"), - ("missing_app", "trial_app_not_allowed"), - ("quota", "trial_app_limit_exceeded"), - ], -) -def test_trial_admission_rejects_before_parsing_or_external_io( - harness: _Harness, endpoint: _Endpoint, denial: str, code: str -) -> None: - if denial == "feature": - harness.features.enabled = False - else: - with harness.factory.begin() as session: - if denial == "missing_trial": - session.execute(delete(TrialApp).where(TrialApp.app_id == harness.target.id)) - elif denial == "missing_app": - session.execute(delete(App).where(App.id == harness.target.id)) - else: - session.add(AccountTrialAppRecord(app_id=harness.target.id, account_id=harness.account.id, count=3)) - - response = harness.app.test_client().post(harness.url(endpoint), json={}) - - _assert_error(response, 403, code) - assert harness.io.saved == {} - assert harness.io.requests == [] - assert harness.uploaded() == [] - assert harness.usage() == (3 if denial == "quota" else None) - if denial == "feature": - assert harness.io.admission_sessions == [] - else: - harness.io.assert_admission_closed() - - -@pytest.mark.parametrize("endpoint", ["files/upload", "remote-files/upload"]) -def test_upload_does_not_charge_usage_or_inherit_another_accounts_exhausted_quota( - harness: _Harness, endpoint: _Endpoint -) -> None: - other_account_id = str(uuid4()) - with harness.factory.begin() as session: - session.add_all( - [ - AccountTrialAppRecord(app_id=harness.target.id, account_id=harness.account.id, count=1), - AccountTrialAppRecord(app_id=harness.target.id, account_id=other_account_id, count=3), - ] - ) - - response = harness.post(endpoint) - - assert response.status_code == 201 - assert harness.usage() == 1 - with harness.factory() as session: - assert ( - session.scalar( - select(AccountTrialAppRecord.count).where(AccountTrialAppRecord.account_id == other_account_id) - ) - == 3 - ) - - -@pytest.mark.parametrize( - ("shape", "code"), - [("missing", "no_file_uploaded"), ("multiple", "too_many_files"), ("empty_name", "filename_not_exists_error")], -) -def test_multipart_shape_errors_precede_dataset_permissions(harness: _Harness, shape: str, code: str) -> None: - harness.account.role = TenantAccountRole.NORMAL - data = {} - if shape != "missing": - data["file"] = (BytesIO(_CONTENT), "" if shape == "empty_name" else "trial.txt", "text/plain") - if shape == "multiple": - data["second_file"] = (BytesIO(_CONTENT), "second.txt", "text/plain") - - response = harness.app.test_client().post(harness.url("files/upload") + "?source=datasets", data=data) - - _assert_error(response, 400, code) - assert harness.io.saved == {} - assert harness.uploaded() == [] - - -@pytest.mark.parametrize( - ("query_source", "form_source", "status"), - [("datasets", "unknown", 403), ("unknown", "datasets", 201), ("", "datasets", 403)], -) -def test_query_source_takes_precedence_over_form_before_active_workspace_permission( - harness: _Harness, query_source: str, form_source: str, status: int -) -> None: - harness.account.role = TenantAccountRole.NORMAL - response = harness.app.test_client().post( - harness.url("files/upload") + f"?source={query_source}", - data={"file": (BytesIO(_CONTENT), "trial.txt", "text/plain"), "source": form_source}, - ) - - assert response.status_code == status - if status == 403: - _assert_error(response, status, "forbidden") - assert harness.io.saved == {} - else: - assert len(harness.uploaded()) == 1 - - -@pytest.mark.parametrize( - ("role", "rbac", "status"), - [ - (TenantAccountRole.OWNER, False, 201), - (TenantAccountRole.DATASET_OPERATOR, False, 201), - (TenantAccountRole.NORMAL, False, 403), - (TenantAccountRole.NORMAL, True, 201), - ], -) -def test_dataset_upload_uses_active_workspace_role_and_retains_rbac_bypass( - harness: _Harness, config_overrides: Callable[..., None], role: TenantAccountRole, rbac: bool, status: int -) -> None: - harness.account.role = role - config_overrides(RBAC_ENABLED=rbac) - - response = harness.post("files/upload", query="?source=datasets") - - assert response.status_code == status - assert len(harness.uploaded()) == (1 if status == 201 else 0) - - -@pytest.mark.parametrize("endpoint", ["files/upload", "remote-files/upload"]) -def test_cloud_document_quota_uses_active_workspace_before_upload( - harness: _Harness, monkeypatch: pytest.MonkeyPatch, config_overrides: Callable[..., None], endpoint: _Endpoint -) -> None: - config_overrides(DEPLOYMENT_EDITION=DeploymentEdition.CLOUD) - checked_tenants: list[str] = [] - - def get_features(tenant_id: str, exclude_vector_space: bool = False) -> FeatureModel: - checked_tenants.append(tenant_id) - assert exclude_vector_space is True - features = FeatureModel() - features.documents_upload_quota.limit = 1 - features.documents_upload_quota.size = 1 - return features - - monkeypatch.setattr(FeatureService, "get_features", get_features) - - response = harness.post(endpoint, query="?source=datasets") - - _assert_error(response, 403, "forbidden") - assert checked_tenants == [harness.account.current_tenant_id] - assert harness.io.saved == {} - assert harness.io.requests == [] - assert harness.usage() is None - - -def test_dataset_file_size_limit_comes_from_app_owner_tenant( - harness: _Harness, monkeypatch: pytest.MonkeyPatch, config_overrides: Callable[..., None] -) -> None: - config_overrides(DEPLOYMENT_EDITION=DeploymentEdition.CLOUD) - quota_tenants: list[str] = [] - size_limit_tenants: list[str | None] = [] - - def get_features(tenant_id: str, exclude_vector_space: bool = False) -> FeatureModel: - quota_tenants.append(tenant_id) - assert exclude_vector_space is True - return FeatureModel() - - def get_size_limit(tenant_id: str | None) -> int: - size_limit_tenants.append(tenant_id) - return 0 - - monkeypatch.setattr(FeatureService, "get_features", get_features) - monkeypatch.setattr(FeatureService, "get_knowledge_file_size_limit", get_size_limit) - - response = harness.post("files/upload", query="?source=datasets") - - _assert_error(response, 413, "file_too_large") - assert response.get_json()["message"] == "File size exceeded." - assert quota_tenants == [harness.account.current_tenant_id] - assert size_limit_tenants == [harness.target.tenant_id] - assert harness.io.saved == {} - assert harness.uploaded() == [] - - -@pytest.mark.parametrize( - ("filename", "query", "status", "code"), - [("trial.exe", "", 400, "file_extension_blocked"), ("trial.png", "?source=datasets", 415, "unsupported_file_type")], -) -def test_local_file_policy_errors_keep_specific_codes( - harness: _Harness, filename: str, query: str, status: int, code: str -) -> None: - response = harness.app.test_client().post( - harness.url("files/upload") + query, data={"file": (BytesIO(_CONTENT), filename, "application/octet-stream")} - ) - - _assert_error(response, status, code) - assert harness.io.saved == {} - assert harness.uploaded() == [] - - -@pytest.mark.parametrize( - ("get_status", "status", "code"), - [ - (404, 404, "remote_file_not_found"), - (403, 400, "remote_file_access_denied"), - (503, 502, "remote_file_unavailable"), - ], -) -def test_remote_http_failures_keep_precise_error_codes( - harness: _Harness, get_status: int, status: int, code: str -) -> None: - harness.io.head_status = 405 - harness.io.get_status = get_status - - response = harness.post("remote-files/upload") - - _assert_error(response, status, code) - assert [method for method, _url, _kwargs in harness.io.requests] == ["HEAD", "GET"] - assert harness.io.saved == {} - assert harness.uploaded() == [] - assert harness.usage() is None - - -@pytest.mark.parametrize( - ("failure", "status", "code"), - [ - ("blocked", 400, "remote_file_url_blocked"), - ("timeout", 502, "remote_file_unavailable"), - ("metadata", 502, "remote_file_invalid_response"), - ("size", 413, "file_too_large"), - ], -) -def test_remote_transport_metadata_and_size_failures_prevent_storage( - harness: _Harness, failure: str, status: int, code: str -) -> None: - if failure == "blocked": - harness.io.error = ToolSSRFError("Blocked by SSRF proxy") - elif failure == "timeout": - harness.io.error = httpx.ReadTimeout("Timed out") - elif failure == "metadata": - harness.io.content_length = "invalid" - else: - harness.io.content_length = str(1024 * 1024 + 1) - - response = harness.post("remote-files/upload") - - _assert_error(response, status, code) - assert len(harness.io.requests) == 1 - assert harness.io.saved == {} - assert harness.uploaded() == [] - assert harness.usage() is None - - -def test_remote_invalid_url_fails_before_network(harness: _Harness) -> None: - response = harness.app.test_client().post(harness.url("remote-files/upload"), json={"url": "file:///etc/passwd"}) - - _assert_error(response, 400, "remote_file_invalid_url") - assert harness.io.requests == [] - assert harness.io.saved == {} - - -def test_remote_payload_validation_stays_on_http_boundary(harness: _Harness) -> None: - response = harness.app.test_client().post(harness.url("remote-files/upload"), json={}) - - assert response.status_code == 422 - assert "url" in response.get_json()["message"] - assert harness.io.requests == [] - assert harness.io.saved == {} diff --git a/api/tests/unit_tests/controllers/console/explore/test_trial_generation.py b/api/tests/unit_tests/controllers/console/explore/test_trial_generation.py deleted file mode 100644 index 3ad37acf1c7..00000000000 --- a/api/tests/unit_tests/controllers/console/explore/test_trial_generation.py +++ /dev/null @@ -1,649 +0,0 @@ -"""Trial admission and generation through real HTTP handlers and SQLite repositories.""" - -from collections.abc import Callable, Iterator, Mapping -from dataclasses import dataclass, field -from typing import override -from uuid import uuid4 - -import pytest -from flask import Flask, Request -from sqlalchemy import Connection, event, select -from sqlalchemy.orm import Mapper, Session, SessionTransaction, sessionmaker -from werkzeug.exceptions import Unauthorized -from werkzeug.test import TestResponse - -import controllers.console.explore.trial as trial_module -import controllers.console.explore.trial_app_admission as admission_module -import controllers.console.wraps as console_wraps -import libs.login as login_module -from core.app.apps.agent_app.errors import AgentAppNotPublishedError -from core.app.entities.app_invoke_entities import InvokeFrom -from core.errors.error import ( - AppInvokeQuotaExceededError, - ModelCurrentlyNotSupportError, - ProviderTokenNotInitError, - QuotaExceededError, -) -from enums import DeploymentEdition -from extensions.ext_login import DifyLoginManager, unauthorized_handler -from graphon.model_runtime.errors.invoke import InvokeError -from libs.external_api import ExternalApi -from models import Account, AccountTrialAppRecord, App, AppMode, Tenant, TrialApp -from models.account import AccountStatus -from models.enums import CreatorUserRole -from repositories.trial_app_repository import TrialAppRepository -from services.account_errors import AccountNotFoundError -from services.app_definition_query_service import AppDefinitionUnavailableError -from services.errors.app_model_config import AppModelConfigBrokenError -from services.errors.conversation import ConversationCompletedError, ConversationNotExistsError -from services.errors.llm import InvokeRateLimitError -from services.trial_app_access_service import TrialAppAccessService, TrialAppRef -from services.trial_app_generation_service import GenerationResponse, TrialAppGenerationService - - -@dataclass -class _FeatureState: - enabled: bool = True - setup_completed: bool = True - events: list[str] = field(default_factory=list) - - def is_trial_enabled(self) -> bool: - self.events.append("feature") - return self.enabled - - -@dataclass(frozen=True) -class _GenerationCall: - app: TrialAppRef - account_id: str - args: Mapping[str, object] - streaming: bool - - -@dataclass -class _Stream(Iterator[str]): - chunks: Iterator[str] = field(default_factory=lambda: iter(['data: {"answer":"first"}\n\n', "data: [DONE]\n\n"])) - reads: int = 0 - close_calls: int = 0 - close_error: Exception | None = None - on_read: Callable[[], None] | None = None - - @override - def __next__(self) -> str: - if self.on_read is not None: - self.on_read() - chunk = next(self.chunks) - self.reads += 1 - return chunk - - def close(self) -> None: - self.close_calls += 1 - if self.close_error is not None: - raise self.close_error - - -@dataclass -class _Runtime: - sessions: list[Session] - response: GenerationResponse = field( - default_factory=lambda: {"answer": "", "metadata": dict[str, object](), "created_at": 0} - ) - error: Exception | None = None - calls: list[_GenerationCall] = field(default_factory=list) - before_generate: Callable[[], None] | None = None - - def generate( - self, *, app: TrialAppRef, account_id: str, args: Mapping[str, object], streaming: bool - ) -> GenerationResponse: - # Database admission must release its connection before model invocation. - assert self.sessions - assert all(not session.in_transaction() and not session.identity_map for session in self.sessions) - self.calls.append(_GenerationCall(app, account_id, dict(args), streaming)) - if self.before_generate is not None: - self.before_generate() - if self.error is not None: - raise self.error - return self.response - - -@dataclass -class _Tasks: - calls: list[str] = field(default_factory=list) - chat_calls: list[tuple[str, str]] = field(default_factory=list) - - def stop_task(self, *, task_id: str, invoke_from: InvokeFrom, user_id: str) -> None: - assert invoke_from == InvokeFrom.EXPLORE - self.chat_calls.append((task_id, user_id)) - - def stop_workflow_task( - self, - *, - task_id: str, - tenant_id: str, - app_id: str, - app_mode: AppMode, - owner: tuple[CreatorUserRole, str] | None, - ) -> None: - assert tenant_id - assert app_id - if app_mode == AppMode.ADVANCED_CHAT: - assert owner is not None - assert owner[0] == CreatorUserRole.ACCOUNT - self.chat_calls.append((task_id, owner[1])) - else: - assert owner is None - assert app_mode == AppMode.WORKFLOW - self.calls.append(task_id) - - -@dataclass(frozen=True) -class _TrialAppServices: - access: TrialAppAccessService - generation: TrialAppGenerationService - - -@dataclass(frozen=True) -class _ApplicationServices: - trial_apps: _TrialAppServices - recommended_app_queries: _FeatureState - app_tasks: _Tasks - - -@dataclass(frozen=True) -class _Harness: - app: Flask - account: Account - target: App - trial: TrialApp - factory: sessionmaker[Session] - services: _ApplicationServices - runtime: _Runtime - state: _FeatureState - - def url(self, endpoint: str, *, app_id: str | None = None) -> str: - return f"/trial-apps/{app_id or self.target.id}/{endpoint}" - - def post(self, endpoint: str, *, payload: Mapping[str, object] | None = None) -> TestResponse: - return self.app.test_client().post(self.url(endpoint), json=dict(payload or {"inputs": {}})) - - def set_mode(self, mode: AppMode) -> None: - with self.factory.begin() as session: - app = session.get(App, self.target.id) - assert app is not None - app.mode = mode - - def usage(self) -> int | None: - with self.factory() as session: - return session.scalar( - select(AccountTrialAppRecord.count).where( - AccountTrialAppRecord.app_id == self.target.id, - AccountTrialAppRecord.account_id == self.account.id, - ) - ) - - -@pytest.fixture -def harness( - monkeypatch: pytest.MonkeyPatch, - config_overrides: Callable[..., None], - sqlite_session_factory: sessionmaker[Session], -) -> _Harness: - config_overrides(DEPLOYMENT_EDITION=DeploymentEdition.COMMUNITY, INIT_PASSWORD="", LOGIN_DISABLED=False) - state = _FeatureState() - account = Account(name="Viewer", email="viewer@example.com") - account._current_tenant = Tenant(name="Viewer workspace") - target = App(tenant_id=str(uuid4()), name="Trial app", mode=AppMode.COMPLETION, enable_site=True, enable_api=True) - target.id = str(uuid4()) - trial = TrialApp(app_id=target.id, tenant_id=target.tenant_id, trial_limit=3) - with sqlite_session_factory.begin() as session: - session.add_all([target, trial]) - - sessions: list[Session] = [] - repository_factory = sessionmaker(bind=sqlite_session_factory.kw["bind"], expire_on_commit=False) - - @event.listens_for(repository_factory, "after_begin") - def track_session(session: Session, _transaction: SessionTransaction, _connection: Connection) -> None: - sessions.append(session) - - runtime = _Runtime(sessions=sessions) - repository = TrialAppRepository(session_factory=repository_factory) - services = _ApplicationServices( - trial_apps=_TrialAppServices( - access=TrialAppAccessService(apps=repository), - generation=TrialAppGenerationService(runtime=runtime, usage=repository), - ), - recommended_app_queries=state, - app_tasks=_Tasks(), - ) - - def setup_completed() -> bool: - state.events.append("setup") - return state.setup_completed - - def csrf(_request: Request, account_id: str) -> None: - assert account_id == account.id - state.events.append("csrf") - - for module in (trial_module, admission_module): - monkeypatch.setattr(module, "application_services", lambda: services) - monkeypatch.setattr(console_wraps, "_is_setup_completed", setup_completed) - monkeypatch.setattr(login_module, "current_user", account) - monkeypatch.setattr(login_module, "check_csrf_token", csrf) - - app = Flask(__name__) - app.config.update(TESTING=True, RESTX_ERROR_404_HELP=False) - login_manager = DifyLoginManager() - login_manager.init_app(app) - login_manager.unauthorized_handler(unauthorized_handler) - api = ExternalApi(app) - for resource, endpoint in ( - (trial_module.TrialAppWorkflowRunApi, "workflows/run"), - (trial_module.TrialChatApi, "chat-messages"), - (trial_module.TrialChatTaskStopApi, "chat-messages//stop"), - (trial_module.TrialCompletionApi, "completion-messages"), - (trial_module.TrialAppWorkflowTaskStopApi, "workflows/tasks//stop"), - ): - api.add_resource(resource, f"/trial-apps//{endpoint}") - return _Harness(app, account, target, trial, sqlite_session_factory, services, runtime, state) - - -def _assert_error(response: TestResponse, status: int, code: str) -> Mapping[str, object]: - assert response.status_code == status - body = response.get_json() - assert isinstance(body, dict) - assert body["code"] == code - assert body["status"] == status - assert response.headers["Content-Type"] == "application/json" - assert isinstance(body["message"], str) - assert body["message"] - return body - - -@pytest.mark.parametrize("response_mode", [None, "blocking"]) -def test_completion_blocking_preserves_response_and_uses_app_owner( - harness: _Harness, response_mode: str | None -) -> None: - response = harness.post("completion-messages", payload={"inputs": {}, "response_mode": response_mode}) - - assert response.status_code == 200 - assert response.get_json() == {"answer": "", "metadata": {}, "created_at": 0} - assert response.headers["Content-Type"] == "application/json; charset=utf-8" - assert int(response.headers["Content-Length"]) == len(response.data) - assert harness.usage() == 1 - assert harness.state.events == ["setup", "csrf", "feature"] - call = harness.runtime.calls[0] - assert call.app == TrialAppRef( - app_id=harness.target.id, tenant_id=harness.target.tenant_id, app_mode=AppMode.COMPLETION - ) - assert call.app.tenant_id != harness.account.current_tenant_id - assert call.account_id == harness.account.id - assert call.streaming is False - assert call.args == { - "inputs": {}, - "query": "", - "files": None, - "response_mode": response_mode, - "retriever_from": "explore_app", - "auto_generate_name": False, - } - - -@pytest.mark.parametrize( - ("endpoint", "mode", "payload"), - [ - ("workflows/run", AppMode.WORKFLOW, {"inputs": {"number": 0}, "files": []}), - ("chat-messages", AppMode.CHAT, {"inputs": {}, "query": "hello"}), - ("chat-messages", AppMode.AGENT, {"inputs": {}, "query": "hello"}), - ("chat-messages", AppMode.ADVANCED_CHAT, {"inputs": {}, "query": "hello"}), - ("chat-messages", AppMode.AGENT_CHAT, {"inputs": {}, "query": "hello"}), - ("completion-messages", AppMode.COMPLETION, {"inputs": {}, "response_mode": "streaming"}), - ], -) -def test_stream_records_usage_after_runtime_returns_before_consumption( - harness: _Harness, endpoint: str, mode: AppMode, payload: Mapping[str, object] -) -> None: - harness.set_mode(mode) - - def before_generate() -> None: - assert harness.usage() is None - - def before_read() -> None: - assert harness.usage() == 1 - - stream = _Stream(on_read=before_read) - harness.runtime.response = stream - harness.runtime.before_generate = before_generate - response = harness.post(endpoint, payload=payload) - - assert response.status_code == 200 - assert response.headers["Content-Type"] == "text/event-stream; charset=utf-8" - assert "Content-Length" not in response.headers - assert response.data == b'data: {"answer":"first"}\n\ndata: [DONE]\n\n' - response.close() - assert harness.usage() == 1 - assert harness.runtime.calls[0].streaming is True - if endpoint == "chat-messages": - assert harness.runtime.calls[0].args["auto_generate_name"] is False - elif endpoint == "workflows/run": - assert harness.runtime.calls[0].args == payload - - -@pytest.mark.parametrize( - "endpoint", - ["completion-messages", "chat-messages", "chat-messages/task/stop", "workflows/run", "workflows/tasks/task/stop"], -) -@pytest.mark.parametrize( - ("state", "status", "code"), [("setup", 401, "not_setup"), ("feature", 403, "trial_app_feature_disabled")] -) -def test_bootstrap_and_feature_admission_apply_before_trial_lookup( - harness: _Harness, endpoint: str, state: str, status: int, code: str -) -> None: - if state == "setup": - harness.state.setup_completed = False - else: - harness.state.enabled = False - response = harness.app.test_client().post(harness.url(endpoint, app_id=str(uuid4())), json={}) - - _assert_error(response, status, code) - assert harness.state.events == (["setup"] if state == "setup" else ["setup", "csrf", "feature"]) - assert harness.runtime.calls == [] - assert harness.services.app_tasks.calls == [] - - -@pytest.mark.parametrize( - ("case", "code"), - [ - ("missing_trial", "trial_app_not_allowed"), - ("missing_app", "trial_app_not_allowed"), - ("limit", "trial_app_limit_exceeded"), - ], -) -def test_trial_lookup_and_account_limit_block_generation(harness: _Harness, case: str, code: str) -> None: - with harness.factory.begin() as session: - if case == "missing_trial": - trial = session.get(TrialApp, harness.trial.id) - assert trial is not None - session.delete(trial) - elif case == "missing_app": - target = session.get(App, harness.target.id) - assert target is not None - session.delete(target) - else: - session.add(AccountTrialAppRecord(app_id=harness.target.id, account_id=harness.account.id, count=3)) - response = harness.post("completion-messages") - - _assert_error(response, 403, code) - assert harness.runtime.calls == [] - - -@pytest.mark.parametrize("count", [0, 2]) -def test_usage_increments_existing_record_and_ignores_other_account(harness: _Harness, count: int) -> None: - other_id = str(uuid4()) - with harness.factory.begin() as session: - session.add_all( - [ - AccountTrialAppRecord(app_id=harness.target.id, account_id=harness.account.id, count=count), - AccountTrialAppRecord(app_id=harness.target.id, account_id=other_id, count=99), - ] - ) - response = harness.post("completion-messages") - - assert response.status_code == 200 - assert harness.usage() == count + 1 - with harness.factory() as session: - assert ( - session.scalar(select(AccountTrialAppRecord.count).where(AccountTrialAppRecord.account_id == other_id)) - == 99 - ) - - -@pytest.mark.parametrize( - ("endpoint", "mode", "code"), - [ - ("completion-messages", AppMode.CHAT, "not_completion_app"), - ("chat-messages", AppMode.COMPLETION, "not_chat_app"), - ("workflows/run", AppMode.CHAT, "not_workflow_app"), - ("workflows/tasks/task/stop", AppMode.CHAT, "not_workflow_app"), - ], -) -def test_wrong_mode_returns_specific_error_without_invocation( - harness: _Harness, endpoint: str, mode: AppMode, code: str -) -> None: - harness.set_mode(mode) - response = harness.post(endpoint, payload={"inputs": {}, "query": "hello"}) - - _assert_error(response, 400, code) - assert harness.runtime.calls == [] - assert harness.services.app_tasks.calls == [] - assert harness.usage() is None - - -@pytest.mark.parametrize("field", ["conversation_id", "parent_message_id"]) -def test_chat_rejects_invalid_uuid_before_runtime(harness: _Harness, field: str) -> None: - harness.set_mode(AppMode.CHAT) - response = harness.post("chat-messages", payload={"inputs": {}, "query": "hello", field: "not-a-uuid"}) - - _assert_error(response, 400, "invalid_param") - assert harness.runtime.calls == [] - assert harness.usage() is None - - -def test_unpublished_agent_remains_unavailable_after_service_migration(harness: _Harness) -> None: - harness.set_mode(AppMode.AGENT) - harness.runtime.error = AgentAppNotPublishedError("Agent has not been published") - response = harness.post("chat-messages", payload={"inputs": {}, "query": "hello"}) - _assert_error(response, 400, "app_unavailable") - assert harness.usage() is None - - -def test_stop_chat_rejects_workflow(harness: _Harness) -> None: - harness.set_mode(AppMode.WORKFLOW) - _assert_error(harness.post("chat-messages/task-1/stop"), 400, "not_chat_app") - assert harness.services.app_tasks.chat_calls == [] - - -@pytest.mark.parametrize("mode", [AppMode.CHAT, AppMode.AGENT_CHAT, AppMode.AGENT, AppMode.ADVANCED_CHAT]) -def test_last_trial_can_be_stopped_without_allowing_another_generation(harness: _Harness, mode: AppMode) -> None: - harness.set_mode(mode) - with harness.factory.begin() as session: - session.add(AccountTrialAppRecord(app_id=harness.target.id, account_id=harness.account.id, count=2)) - - response = harness.post("chat-messages", payload={"inputs": {}, "query": "hello"}) - assert response.status_code == 200 - assert harness.usage() == 3 - - stopped = harness.post("chat-messages/task-last/stop") - assert stopped.status_code == 200 - assert harness.services.app_tasks.chat_calls == [("task-last", harness.account.id)] - assert harness.usage() == 3 - unavailable = harness.app.test_client().post(harness.url("chat-messages/task-last/stop", app_id=str(uuid4()))) - _assert_error(unavailable, 403, "trial_app_not_allowed") - assert len(harness.services.app_tasks.chat_calls) == 1 - _assert_error( - harness.post("chat-messages", payload={"inputs": {}, "query": "again"}), - 403, - "trial_app_limit_exceeded", - ) - - -def test_chat_normalizes_uuid_and_preserves_empty_ids(harness: _Harness) -> None: - harness.set_mode(AppMode.CHAT) - conversation_id = uuid4() - response = harness.post( - "chat-messages", - payload={ - "inputs": {}, - "query": "hello", - "conversation_id": conversation_id.hex.upper(), - "parent_message_id": "", - }, - ) - - assert response.status_code == 200 - assert harness.runtime.calls[0].args["conversation_id"] == str(conversation_id) - assert harness.runtime.calls[0].args["parent_message_id"] == "" - - -@pytest.mark.parametrize( - ("error", "status", "code"), - [ - (ProviderTokenNotInitError("missing key for provider"), 400, "provider_not_initialize"), - (QuotaExceededError(), 400, "provider_quota_exceeded"), - (ModelCurrentlyNotSupportError(), 400, "model_currently_not_support"), - (InvokeError("provider rejected request"), 400, "completion_request_error"), - (InvokeRateLimitError("too many requests"), 429, "rate_limit_error"), - (AppInvokeQuotaExceededError("active requests exceeded"), 429, "rate_limit_error"), - (AppDefinitionUnavailableError("app removed"), 400, "app_unavailable"), - (AccountNotFoundError("account removed"), 401, "unauthorized"), - (AppModelConfigBrokenError(), 400, "app_unavailable"), - ], -) -@pytest.mark.parametrize( - ("endpoint", "mode"), - [ - ("completion-messages", AppMode.COMPLETION), - ("chat-messages", AppMode.CHAT), - ("workflows/run", AppMode.WORKFLOW), - ], -) -def test_runtime_errors_keep_specific_codes_and_do_not_record_usage( - harness: _Harness, endpoint: str, mode: AppMode, error: Exception, status: int, code: str -) -> None: - harness.set_mode(mode) - harness.runtime.error = error - response = harness.post(endpoint, payload={"inputs": {}, "query": "hello"}) - - body = _assert_error(response, status, code) - if isinstance(error, (ProviderTokenNotInitError, InvokeError, InvokeRateLimitError)): - assert body["message"] == error.description - assert harness.usage() is None - - -@pytest.mark.parametrize( - ("error", "status", "code"), - [ - (ConversationNotExistsError(), 404, "not_found"), - # Nothing raises ConversationCompletedError today, so the contract does not list its code. - pytest.param( - ConversationCompletedError(), - 400, - "conversation_completed", - marks=pytest.mark.allow_error_contract_violations, - ), - (AppModelConfigBrokenError(), 400, "app_unavailable"), - ], -) -@pytest.mark.parametrize( - ("endpoint", "mode"), [("completion-messages", AppMode.COMPLETION), ("chat-messages", AppMode.CHAT)] -) -def test_conversation_and_config_errors_remain_actionable( - harness: _Harness, endpoint: str, mode: AppMode, error: Exception, status: int, code: str -) -> None: - harness.set_mode(mode) - harness.runtime.error = error - response = harness.post(endpoint, payload={"inputs": {}, "query": "hello"}) - - _assert_error(response, status, code) - assert harness.usage() is None - - -def test_stop_retains_trial_quota_admission_without_recording_usage(harness: _Harness) -> None: - harness.set_mode(AppMode.WORKFLOW) - response = harness.post("workflows/tasks/task-123/stop") - - assert response.status_code == 200 - assert response.get_json() == {"result": "success"} - assert harness.services.app_tasks.calls == ["task-123"] - assert harness.usage() is None - with harness.factory.begin() as session: - session.add(AccountTrialAppRecord(app_id=harness.target.id, account_id=harness.account.id, count=3)) - blocked = harness.post("workflows/tasks/task-456/stop") - _assert_error(blocked, 403, "trial_app_limit_exceeded") - assert harness.services.app_tasks.calls == ["task-123"] - - -def test_unauthenticated_account_cannot_enter_trial(harness: _Harness, monkeypatch: pytest.MonkeyPatch) -> None: - monkeypatch.setattr(login_module, "current_user", None) - response = harness.post("completion-messages") - - assert response.status_code == 401 - assert harness.state.events == ["setup"] - assert harness.runtime.calls == [] - - -def test_uninitialized_account_cannot_enter_trial(harness: _Harness) -> None: - harness.account.status = AccountStatus.UNINITIALIZED - response = harness.post("completion-messages") - - _assert_error(response, 400, "account_not_initialized") - assert harness.state.events == ["setup", "csrf"] - assert harness.runtime.calls == [] - - -def test_failed_csrf_cannot_enter_trial(harness: _Harness, monkeypatch: pytest.MonkeyPatch) -> None: - def reject_csrf(_request: Request, _account_id: str) -> None: - raise Unauthorized("invalid csrf token") - - monkeypatch.setattr(login_module, "check_csrf_token", reject_csrf) - response = harness.post("completion-messages") - - _assert_error(response, 401, "unauthorized") - assert harness.runtime.calls == [] - - -def test_payload_validation_keeps_field_context(harness: _Harness) -> None: - response = harness.post("completion-messages", payload={"inputs": [], "response_mode": "invalid"}) - - body = _assert_error(response, 422, "unprocessable_entity") - assert "inputs" in str(body) - assert harness.runtime.calls == [] - assert harness.usage() is None - - -@pytest.mark.parametrize("cleanup_fails", [False, True]) -@pytest.mark.parametrize( - ("endpoint", "mode"), - [("completion-messages", AppMode.COMPLETION), ("chat-messages", AppMode.CHAT), ("workflows/run", AppMode.WORKFLOW)], -) -def test_usage_write_failure_closes_unconsumed_stream_without_hiding_database_error( - harness: _Harness, endpoint: str, mode: AppMode, cleanup_fails: bool, caplog: pytest.LogCaptureFixture -) -> None: - harness.set_mode(mode) - stream = _Stream(close_error=RuntimeError("cleanup failed") if cleanup_fails else None) - harness.runtime.response = stream - failure = RuntimeError("usage database unavailable") - - def reject_usage_insert( - _mapper: Mapper[AccountTrialAppRecord], _connection: Connection, _record: AccountTrialAppRecord - ) -> None: - raise failure - - # Fail the persistence boundary, retaining the real service, unit of work, - # SQLAlchemy rollback and HTTP error handler. - event.listen(AccountTrialAppRecord, "before_insert", reject_usage_insert) - try: - response = harness.post(endpoint, payload={"inputs": {}, "query": "hello", "response_mode": "streaming"}) - finally: - event.remove(AccountTrialAppRecord, "before_insert", reject_usage_insert) - - _assert_error(response, 500, "internal_server_error") - assert stream.close_calls == 1 - assert stream.reads == 0 - assert harness.usage() is None - assert any(record.exc_info is not None and record.exc_info[1] is failure for record in caplog.records) - - -def test_later_stream_failure_does_not_refund_trial_usage(harness: _Harness) -> None: - def chunks() -> Iterator[str]: - yield 'data: {"answer":"partial"}\n\n' - raise RuntimeError("provider disconnected") - - harness.runtime.response = _Stream(chunks=chunks()) - response = harness.post("completion-messages", payload={"inputs": {}, "response_mode": "streaming"}) - - chunks = iter(response.response) - assert next(chunks) == b'data: {"answer":"partial"}\n\n' - with pytest.raises(RuntimeError, match="provider disconnected"): - next(chunks) - response.close() - assert harness.usage() == 1 diff --git a/api/tests/unit_tests/controllers/console/explore/test_trial_preview.py b/api/tests/unit_tests/controllers/console/explore/test_trial_preview.py deleted file mode 100644 index 3db174ef0c1..00000000000 --- a/api/tests/unit_tests/controllers/console/explore/test_trial_preview.py +++ /dev/null @@ -1,423 +0,0 @@ -"""Anonymous app previews through HTTP and real SQLite query boundaries.""" - -import json -from collections.abc import Callable -from dataclasses import dataclass, field -from datetime import datetime -from typing import Literal -from uuid import uuid4 - -import pytest -from flask import Flask -from sqlalchemy import Connection, Engine, delete, event, select, text, update -from sqlalchemy.orm import Session, SessionTransaction, sessionmaker -from sqlalchemy.pool import QueuePool -from werkzeug.test import TestResponse - -import controllers.common.fields as fields_module -import controllers.console.explore.trial as trial_module -import controllers.console.wraps as console_wraps -import libs.login as login_module -from controllers.console.app import preview_admission as admission_module -from libs.external_api import ExternalApi -from models import AccountTrialAppRecord, App, AppMode, Tenant, TrialApp -from models.account import TenantStatus -from models.dataset import Dataset -from models.enums import CustomizeTokenStrategy -from models.model import AppModelConfig, IconType, Site -from repositories.app_definition_query_repository import AppDefinitionQueryRepository -from repositories.app_preview_query_repository import AppPreviewQueryRepository -from repositories.trial_app_repository import TrialAppRepository -from services.app_definition_query_service import AppDefinitionQueryService -from services.app_preview_query_service import AppPreviewQueryService -from services.recommended_app_query_service import ( - RecommendedAppCatalogPage, - RecommendedAppDetailRecord, - RecommendedAppQueryService, -) - -_Endpoint = Literal["parameters", "site", "datasets"] -_CREATED_AT = datetime(2024, 1, 1) -_INPUT_FORM = [{"number": {"label": "Count", "variable": "count", "required": False, "default": 0}}] - - -@dataclass -class _Catalog: - engine: Engine - ids: set[str] = field(default_factory=set) - calls: list[str] = field(default_factory=list) - - def contains(self, app_id: str) -> bool: - assert isinstance(self.engine.pool, QueuePool) - assert self.engine.pool.checkedout() == 0 - self.calls.append(app_id) - return app_id in self.ids - - def list_recommended(self, language: str) -> RecommendedAppCatalogPage: - raise AssertionError(f"Unexpected list request for {language}") - - def list_learn_dify(self, language: str) -> RecommendedAppCatalogPage: - raise AssertionError(f"Unexpected list request for {language}") - - def get_detail(self, app_id: str) -> RecommendedAppDetailRecord | None: - raise AssertionError(f"Unexpected detail request for {app_id}") - - -@dataclass(frozen=True) -class _ApplicationServices: - app_previews: AppPreviewQueryService - app_definitions: AppDefinitionQueryService - recommended_app_queries: RecommendedAppQueryService - - -@dataclass(frozen=True) -class _Harness: - app: Flask - owner: Tenant - target: App - listing: TrialApp - site: Site - config: AppModelConfig - dataset: Dataset - creator_id: str - factory: sessionmaker[Session] - engine: Engine - catalog: _Catalog - sessions: list[Session] - signed_icons: list[str] - - def get(self, endpoint: _Endpoint, *, query: str = "", app_id: str | None = None) -> TestResponse: - response = self.app.test_client().get(f"/trial-apps/{app_id or self.target.id}/{endpoint}{query}") - assert response.headers["Content-Type"] == "application/json" - assert int(response.headers["Content-Length"]) == len(response.data) - assert isinstance(self.engine.pool, QueuePool) - assert self.engine.pool.checkedout() == 0 - assert all(not session.in_transaction() and not session.identity_map for session in self.sessions) - return response - - def dataset_query(self) -> str: - return f"?ids={self.dataset.id}" - - -@pytest.fixture -def harness( - monkeypatch: pytest.MonkeyPatch, - config_overrides: Callable[..., None], - sqlite_engine: Engine, - sqlite_session_factory: sessionmaker[Session], -) -> _Harness: - config_overrides( - LOGIN_DISABLED=False, - INIT_PASSWORD="preview-does-not-require-setup", - UPLOAD_IMAGE_FILE_SIZE_LIMIT=1, - UPLOAD_VIDEO_FILE_SIZE_LIMIT=2, - UPLOAD_AUDIO_FILE_SIZE_LIMIT=3, - UPLOAD_FILE_SIZE_LIMIT=4, - WORKFLOW_FILE_UPLOAD_LIMIT=5, - ) - owner = Tenant(name="App owner") - creator_id = str(uuid4()) - target = App( - id=str(uuid4()), tenant_id=owner.id, name="Preview", mode=AppMode.CHAT, enable_site=False, enable_api=False - ) - listing = TrialApp(app_id=target.id, tenant_id=str(uuid4()), trial_limit=0) - config = AppModelConfig( - app_id=target.id, - opening_statement="", - suggested_questions="[]", - suggested_questions_after_answer='{"enabled":false}', - speech_to_text='{"enabled":false}', - text_to_speech='{"enabled":false,"voice":"","autoPlay":"disabled"}', - retriever_resource='{"enabled":false}', - more_like_this='{"enabled":false}', - sensitive_word_avoidance='{"enabled":false,"type":"","configs":[]}', - file_upload='{"enabled":false,"number_limits":0}', - user_input_form=json.dumps(_INPUT_FORM), - ) - site = Site( - app_id=target.id, - title="Preview site", - default_language="en-US", - customize_token_strategy=CustomizeTokenStrategy.UUID, - icon_type=IconType.IMAGE, - icon=str(uuid4()), - icon_background="", - chat_color_theme="", - chat_color_theme_inverted=False, - description=None, - copyright="", - privacy_policy=None, - input_placeholder="", - custom_disclaimer="", - show_workflow_steps=False, - use_icon_as_answer_icon=False, - prompt_public=False, - ) - dataset = Dataset( - tenant_id=owner.id, - name="Owner dataset", - description=None, - permission="only_me", - data_source_type="upload_file", - indexing_technique=None, - created_by=creator_id, - ) - dataset.created_at = _CREATED_AT - with sqlite_session_factory.begin() as session: - session.add_all([owner, target, listing, config, site, dataset]) - session.flush() - target.app_model_config_id = config.id - session.add(AccountTrialAppRecord(app_id=target.id, account_id=creator_id, count=100)) - - factory = sessionmaker(bind=sqlite_engine, expire_on_commit=False) - sessions: list[Session] = [] - - @event.listens_for(factory, "after_begin") - def track_session(session: Session, _transaction: SessionTransaction, _connection: Connection) -> None: - assert all(not previous.in_transaction() for previous in sessions) - assert isinstance(sqlite_engine.pool, QueuePool) - assert sqlite_engine.pool.checkedout() == 1 - sessions.append(session) - - catalog = _Catalog(sqlite_engine) - recommendations = RecommendedAppQueryService( - catalog=catalog, trial_apps=TrialAppRepository(factory), trial_enabled=False - ) - services = _ApplicationServices( - app_previews=AppPreviewQueryService( - apps=AppPreviewQueryRepository(session_factory=factory), is_previewable=recommendations.is_previewable - ), - app_definitions=AppDefinitionQueryService( - definitions=AppDefinitionQueryRepository(session_factory=factory), builtin_icon_url_prefix="/tools/" - ), - recommended_app_queries=recommendations, - ) - signed_icons: list[str] = [] - - def sign_icon(file_id: str) -> str: - assert isinstance(sqlite_engine.pool, QueuePool) - assert sqlite_engine.pool.checkedout() == 0 - signed_icons.append(file_id) - return f"https://files.example/{file_id}?sign=preview" - - for module in (trial_module, admission_module): - monkeypatch.setattr(module, "application_services", lambda: services) - monkeypatch.setattr(login_module, "current_user", None) - monkeypatch.setattr(console_wraps, "_is_setup_completed", lambda: False) - monkeypatch.setattr(fields_module.file_helpers, "get_signed_file_url", sign_icon) - app = Flask(__name__) - app.config.update(TESTING=True, RESTX_ERROR_404_HELP=False) - api = ExternalApi(app) - api.add_resource(trial_module.TrialAppParameterApi, "/trial-apps//parameters") - api.add_resource(trial_module.TrialSitApi, "/trial-apps//site") - api.add_resource(trial_module.DatasetListApi, "/trial-apps//datasets") - return _Harness( - app, - owner, - target, - listing, - site, - config, - dataset, - creator_id, - sqlite_session_factory, - sqlite_engine, - catalog, - sessions, - signed_icons, - ) - - -def _assert_error(response: TestResponse, status: int, code: str) -> None: - assert response.status_code == status, response.get_json() - body = response.get_json() - assert body["code"] == code - assert body["status"] == status - assert isinstance(body["message"], str) - assert body["message"] - - -@pytest.mark.parametrize("endpoint", ["parameters", "site", "datasets"]) -@pytest.mark.parametrize("membership", ["trial", "catalog"]) -def test_anonymous_preview_ignores_trial_feature_quota_and_setup( - harness: _Harness, endpoint: _Endpoint, membership: str -) -> None: - if membership == "catalog": - with harness.factory.begin() as session: - session.execute(delete(TrialApp).where(TrialApp.app_id == harness.target.id)) - harness.catalog.ids.add(harness.target.id) - - response = harness.get(endpoint, query=harness.dataset_query() if endpoint == "datasets" else "") - - assert response.status_code == 200, response.get_json() - assert harness.catalog.calls == ([] if membership == "trial" else [harness.target.id]) - with harness.factory() as session: - assert session.scalar(select(AccountTrialAppRecord.count)) == 100 - - -@pytest.mark.parametrize("endpoint", ["parameters", "site", "datasets"]) -@pytest.mark.parametrize("unavailable", ["unlisted", "disabled", "missing"]) -def test_preview_rejects_unavailable_apps_before_loading_content( - harness: _Harness, endpoint: _Endpoint, unavailable: str -) -> None: - with harness.factory.begin() as session: - if unavailable == "unlisted": - session.execute(delete(TrialApp).where(TrialApp.app_id == harness.target.id)) - elif unavailable == "disabled": - session.execute( - text("UPDATE apps SET status = 'disabled' WHERE id = :app_id"), {"app_id": harness.target.id} - ) - else: - session.execute(delete(App).where(App.id == harness.target.id)) - - _assert_error(harness.get(endpoint, query=harness.dataset_query()), 404, "app_not_found") - assert harness.signed_icons == [] - - -def test_parameters_preserve_complete_shape_and_zero_false_empty_values(harness: _Harness) -> None: - response = harness.get("parameters") - assert response.status_code == 200, response.get_json() - assert response.get_json() == { - "opening_statement": "", - "suggested_questions": [], - "suggested_questions_after_answer": {"enabled": False}, - "speech_to_text": {"enabled": False}, - "text_to_speech": {"enabled": False, "voice": "", "autoPlay": "disabled"}, - "retriever_resource": {"enabled": False}, - "annotation_reply": {"enabled": False}, - "more_like_this": {"enabled": False}, - "user_input_form": _INPUT_FORM, - "sensitive_word_avoidance": {"enabled": False, "type": "", "configs": []}, - "file_upload": {"enabled": False, "number_limits": 0}, - "system_parameters": { - "image_file_size_limit": 1, - "video_file_size_limit": 2, - "audio_file_size_limit": 3, - "file_size_limit": 4, - "workflow_file_upload_limit": 5, - }, - } - - -def test_missing_parameter_configuration_is_app_unavailable(harness: _Harness) -> None: - with harness.factory.begin() as session: - session.execute(delete(AppModelConfig).where(AppModelConfig.id == harness.config.id)) - _assert_error(harness.get("parameters"), 400, "app_unavailable") - - -def test_site_preserves_complete_fields_and_signs_after_session_closes(harness: _Harness) -> None: - response = harness.get("site") - assert response.status_code == 200, response.get_json() - assert harness.target.enable_site is False - assert response.get_json() == { - "title": "Preview site", - "chat_color_theme": "", - "chat_color_theme_inverted": False, - "icon_type": "image", - "icon": harness.site.icon, - "icon_background": "", - "description": None, - "copyright": "", - "privacy_policy": None, - "input_placeholder": "", - "custom_disclaimer": "", - "default_language": "en-US", - "show_workflow_steps": False, - "use_icon_as_answer_icon": False, - "icon_url": f"https://files.example/{harness.site.icon}?sign=preview", - } - assert harness.signed_icons == [harness.site.icon] - - -@pytest.mark.parametrize("icon_type", [IconType.EMOJI, None]) -def test_non_image_site_icon_has_no_signed_url(harness: _Harness, icon_type: IconType | None) -> None: - with harness.factory.begin() as session: - session.execute(update(Site).where(Site.id == harness.site.id).values(icon_type=icon_type, icon="")) - response = harness.get("site") - assert response.status_code == 200 - assert response.get_json()["icon_url"] is None - assert harness.signed_icons == [] - - -@pytest.mark.parametrize( - ("missing", "code"), - [("site", "app_site_unavailable"), ("owner", "app_owner_unavailable"), ("archived_owner", "app_owner_unavailable")], -) -def test_site_errors_distinguish_site_and_owner_unavailability(harness: _Harness, missing: str, code: str) -> None: - with harness.factory.begin() as session: - if missing == "site": - session.execute(delete(Site).where(Site.app_id == harness.target.id)) - elif missing == "owner": - session.execute(delete(Tenant).where(Tenant.id == harness.owner.id)) - else: - session.execute(update(Tenant).where(Tenant.id == harness.owner.id).values(status=TenantStatus.ARCHIVE)) - _assert_error(harness.get("site"), 403, code) - assert harness.signed_icons == [] - - -def test_dataset_ids_filter_by_actual_app_owner_without_requiring_binding(harness: _Harness) -> None: - with harness.factory.begin() as session: - other_owner = Dataset( - tenant_id=harness.listing.tenant_id, name="Listing tenant dataset", created_by=harness.creator_id - ) - unrequested = Dataset(tenant_id=harness.owner.id, name="Not requested", created_by=harness.creator_id) - session.add_all([other_owner, unrequested]) - response = harness.get( - "datasets", query=f"?ids={harness.dataset.id}&ids={harness.dataset.id}&ids={other_owner.id}&ids={uuid4()}" - ) - assert response.status_code == 200, response.get_json() - assert response.get_json() == { - "data": [ - { - "id": harness.dataset.id, - "name": "Owner dataset", - "description": None, - "permission": "only_me", - "data_source_type": "upload_file", - "indexing_technique": None, - "created_by": harness.creator_id, - "created_at": int(_CREATED_AT.timestamp()), - "permission_keys": [], - } - ], - "has_more": False, - "limit": 20, - "total": 1, - "page": 1, - } - - -@pytest.mark.parametrize( - ("query", "page", "limit", "has_more"), - [ - ("&page=0&limit=0", 0, 0, False), - ("&page=-2&limit=-1", -2, -1, False), - ("&page=invalid&limit=invalid", 1, 20, False), - ("&page=99&limit=1", 99, 1, False), - ], -) -def test_dataset_page_and_limit_are_metadata_without_slicing( - harness: _Harness, query: str, page: int, limit: int, has_more: bool -) -> None: - response = harness.get("datasets", query=harness.dataset_query() + query) - assert response.status_code == 200, response.get_json() - body = response.get_json() - assert [item["id"] for item in body["data"]] == [harness.dataset.id] - assert {key: body[key] for key in ("page", "limit", "total", "has_more")} == { - "page": page, - "limit": limit, - "total": 1, - "has_more": has_more, - } - - -@pytest.mark.parametrize("ids", ["missing", "empty"]) -@pytest.mark.parametrize("limit", [20, 0]) -def test_unmatched_dataset_ids_return_empty_result(harness: _Harness, ids: str, limit: int) -> None: - response = harness.get("datasets", query=f"?ids={uuid4() if ids == 'missing' else ''}&limit={limit}") - assert response.status_code == 200 - assert response.get_json() == {"data": [], "has_more": False, "limit": limit, "total": 0, "page": 1} - - -def test_dataset_list_requires_ids(harness: _Harness) -> None: - _assert_error(harness.get("datasets"), 400, "need_add_ids") diff --git a/api/tests/unit_tests/controllers/console/explore/test_trial_preview_details.py b/api/tests/unit_tests/controllers/console/explore/test_trial_preview_details.py deleted file mode 100644 index e618179e66c..00000000000 --- a/api/tests/unit_tests/controllers/console/explore/test_trial_preview_details.py +++ /dev/null @@ -1,853 +0,0 @@ -"""Preview detail contracts through real HTTP, services, and SQLite queries.""" - -import json -from collections.abc import Callable, Sequence -from dataclasses import dataclass, field -from datetime import datetime -from types import SimpleNamespace -from typing import Literal, cast, override -from uuid import uuid4 - -import pytest -from flask import Flask, got_request_exception -from pydantic import JsonValue -from sqlalchemy import Connection, Engine, delete, event, select, text, update -from sqlalchemy.engine import ExecutionContext -from sqlalchemy.orm import Session, SessionTransaction, sessionmaker -from sqlalchemy.pool import QueuePool -from werkzeug.test import TestResponse - -import controllers.console.explore.trial as trial_module -import controllers.console.wraps as console_wraps -import libs.login as login_module -import services.app_service as app_module -from constants import COOKIE_NAME_CSRF_TOKEN, HEADER_NAME_CSRF_TOKEN -from controllers.console import api as console_api -from controllers.console.app import preview_admission as admission_module -from core.agent.entities import AgentToolEntity -from core.helper import encrypter -from core.plugin.plugin_service import PluginService -from core.tools.__base.tool import Tool -from core.tools.entities.tool_entities import ToolParameter -from core.tools.tool_manager import ToolManager -from core.workflow.llm_environment_variable import LLMEnvironmentVariable -from enums import DeploymentEdition -from extensions.ext_login import DifyLoginManager, unauthorized_handler -from graphon.variables import SecretVariable, StringVariable -from libs.external_api import ExternalApi -from libs.token import generate_csrf_token -from machinery.errors import ActiveWorkspaceRequiredError -from models import Account, AccountTrialAppRecord, App, AppMode, Tenant, TrialApp -from models.account import AccountStatus -from models.agent import ( - Agent, - AgentConfigRevision, - AgentConfigRevisionOperation, - AgentConfigSnapshot, - AgentScope, - AgentSource, -) -from models.agent_config_entities import AgentSoulConfig -from models.enums import CustomizeTokenStrategy, TagType -from models.model import AppModelConfig, IconType, Site, Tag, TagBinding -from models.provider_ids import GenericProviderID -from models.tools import WorkflowToolProvider -from models.workflow import Workflow -from repositories.app_preview_query_repository import AppPreviewQueryRepository -from repositories.trial_app_repository import TrialAppRepository -from services.app_preview_details_adapters import AppPreviewDetailsRuntime -from services.app_preview_query_service import AppPreviewQueryService -from services.app_service import AppResponseView -from services.recommended_app_query_service import ( - RecommendedAppCatalogPage, - RecommendedAppDetailRecord, - RecommendedAppQueryService, -) - -_Endpoint = Literal["detail", "workflows"] -_CREATED_AT = datetime(2024, 1, 1) -_UPDATED_AT = datetime(2024, 1, 2) - - -@dataclass -class _Catalog: - engine: Engine - ids: set[str] = field(default_factory=set) - calls: list[str] = field(default_factory=list) - - def contains(self, app_id: str) -> bool: - assert isinstance(self.engine.pool, QueuePool) - assert self.engine.pool.checkedout() == 0 - self.calls.append(app_id) - return app_id in self.ids - - def list_recommended(self, language: str) -> RecommendedAppCatalogPage: - raise AssertionError(f"Unexpected listing for {language}") - - def list_learn_dify(self, language: str) -> RecommendedAppCatalogPage: - raise AssertionError(f"Unexpected listing for {language}") - - def get_detail(self, app_id: str) -> RecommendedAppDetailRecord | None: - raise AssertionError(f"Unexpected catalog detail for {app_id}") - - -@dataclass -class _ExternalIO: - engine: Engine - owner_id: str - viewer_workspace_id: str - decryptions: list[tuple[str, str]] = field(default_factory=list) - tool_calls: list[tuple[str, str, str]] = field(default_factory=list) - - def encrypt(self, tenant_id: str, token: str) -> str: - assert tenant_id == self.owner_id - assert token == "workflow-secret" - return "cipher-workflow" - - def decrypt(self, tenant_id: str, token: str) -> str: - assert isinstance(self.engine.pool, QueuePool) - assert self.engine.pool.checkedout() == 0 - self.decryptions.append((tenant_id, token)) - if token == "cipher-workflow": - assert tenant_id == self.owner_id - return "workflow-secret" - assert token == "cipher-tool" - assert tenant_id == self.viewer_workspace_id - return "tool-secret" - - def tool_runtime(self, *, tenant_id: str, app_id: str, agent_tool: AgentToolEntity, user_id: str) -> Tool: - assert isinstance(self.engine.pool, QueuePool) - assert self.engine.pool.checkedout() == 0 - assert tenant_id == self.viewer_workspace_id - assert agent_tool.tool_name == "search" - self.tool_calls.append((tenant_id, app_id, user_id)) - parameter = ToolParameter.model_validate( - {"name": "api_key", "label": {"en_US": "Key"}, "type": "secret-input", "form": "form"} - ) - return cast( - Tool, _ToolRuntime(SimpleNamespace(parameters=[parameter], identity=SimpleNamespace(name="search"))) - ) - - -@dataclass -class _ToolRuntime: - entity: SimpleNamespace - - def get_runtime_parameters(self) -> list[ToolParameter]: - return [] - - -@dataclass(frozen=True) -class _ApplicationServices: - app_previews: AppPreviewQueryService - app_preview_details: AppPreviewDetailsRuntime - recommended_app_queries: RecommendedAppQueryService - - -@dataclass(frozen=True) -class _Harness: - app: Flask - owner: Tenant - viewer_workspace: Tenant - viewer: Account - author: Account - updater: Account - target: App - config: AppModelConfig - site: Site - workflow: Workflow - tag: Tag - factory: sessionmaker[Session] - engine: Engine - catalog: _Catalog - io: _ExternalIO - sessions: list[Session] - - def get( - self, - endpoint: _Endpoint = "detail", - *, - method: str = "GET", - csrf_token: str | None = None, - ) -> TestResponse: - suffix = "" if endpoint == "detail" else "/workflows" - statements: list[str] = [] - token = generate_csrf_token(self.viewer.id) if csrf_token is None else csrf_token - client = self.app.test_client() - client.set_cookie(COOKIE_NAME_CSRF_TOKEN, token) - - def record_statement( - _connection: Connection, - _cursor: object, - statement: str, - _parameters: object, - _context: ExecutionContext, - _executemany: bool, - ) -> None: - statements.append(statement) - - event.listen(self.engine, "before_cursor_execute", record_statement) - try: - response = client.open( - f"/trial-apps/{self.target.id}{suffix}", method=method, headers={HEADER_NAME_CSRF_TOKEN: token} - ) - finally: - event.remove(self.engine, "before_cursor_execute", record_statement) - assert not [sql for sql in statements if sql.lstrip().upper().startswith(("INSERT ", "UPDATE ", "DELETE "))] - assert response.headers["Content-Type"] == "application/json" - if method != "HEAD": - assert int(response.headers["Content-Length"]) == len(response.data) - assert isinstance(self.engine.pool, QueuePool) - assert self.engine.pool.checkedout() == 0 - assert all(not session.in_transaction() and not session.identity_map for session in self.sessions) - return response - - def assert_app_not_queried(self) -> None: - assert self.sessions == [] - assert self.catalog.calls == [] - assert self.io.tool_calls == [] - assert self.io.decryptions == [] - - def legacy_detail(self) -> dict[str, object]: - """The unchanged ORM response view is the pre-migration serialization oracle.""" - with self.app.test_request_context(f"/trial-apps/{self.target.id}"), self.factory() as session: - app = session.get(App, self.target.id) - assert app is not None - return trial_module.TrialAppDetailResponse.model_validate( - AppResponseView(app, session=session, account=self.viewer), from_attributes=True - ).model_dump(mode="json") - - -@pytest.fixture -def harness( - monkeypatch: pytest.MonkeyPatch, - config_overrides: Callable[..., None], - sqlite_engine: Engine, - sqlite_session_factory: sessionmaker[Session], -) -> _Harness: - config_overrides( - LOGIN_DISABLED=False, - DEPLOYMENT_EDITION=DeploymentEdition.COMMUNITY, - INIT_PASSWORD="", - SECRET_KEY="preview-admission-test-secret-key", - CONSOLE_WEB_URL="", - CONSOLE_API_URL="", - SERVICE_API_URL="", - APP_WEB_URL="", - ) - owner = Tenant(name="App owner") - viewer_workspace = Tenant(name="Viewer workspace") - viewer = Account(name="Viewer", email="viewer@example.com", status=AccountStatus.ACTIVE) - viewer._current_tenant = viewer_workspace - author = Account(name="Author", email="author@example.com") - updater = Account(name="Updater", email="updater@example.com") - target = App( - id=str(uuid4()), - tenant_id=owner.id, - name="Preview app", - description="", - mode=AppMode.CHAT, - icon_type=IconType.EMOJI, - icon="", - icon_background=None, - enable_site=False, - enable_api=False, - max_active_requests=0, - ) - target.created_at = _CREATED_AT - target.updated_at = _UPDATED_AT - config = AppModelConfig( - app_id=target.id, - opening_statement="", - suggested_questions="[]", - model='{"provider":"example/model","name":"model","mode":"chat","completion_params":{"temperature":0}}', - agent_mode='{"enabled":false,"strategy":"function_call","tools":[]}', - user_input_form='[{"number":{"label":"Count","variable":"count","default":0,"required":false}}]', - ) - config.created_at = _CREATED_AT - config.updated_at = _UPDATED_AT - site = Site( - app_id=target.id, - title="Preview site", - default_language="en-US", - customize_token_strategy=CustomizeTokenStrategy.UUID, - code="preview-code", - icon_type=IconType.EMOJI, - icon="", - icon_background=None, - chat_color_theme="", - description=None, - custom_disclaimer="", - show_workflow_steps=False, - ) - site.created_at = _CREATED_AT - site.updated_at = _UPDATED_AT - io = _ExternalIO(sqlite_engine, owner.id, viewer_workspace.id) - monkeypatch.setattr(encrypter, "encrypt_token", io.encrypt) - monkeypatch.setattr(encrypter, "decrypt_token", io.decrypt) - monkeypatch.setattr(ToolManager, "get_agent_tool_runtime", io.tool_runtime) - monkeypatch.setattr(ToolManager, "get_hardcoded_provider", lambda _provider_id: object()) - with sqlite_session_factory.begin() as session: - session.add_all([owner, viewer_workspace, viewer, author, updater, target, config, site]) - session.flush() - target.created_by = author.id - target.updated_by = updater.id - target.app_model_config_id = config.id - workflow = Workflow.new( - tenant_id=owner.id, - app_id=target.id, - type="workflow", - version="draft", - graph='{"nodes":[],"edges":[]}', - features='{"suggested_questions_after_answer":{"enabled":false}}', - created_by=author.id, - environment_variables=[ - SecretVariable(id="secret", name="api_key", value="workflow-secret"), - LLMEnvironmentVariable( - id="llm", name="shared_model", value={"provider": "provider", "name": "model", "mode": "chat"} - ), - ], - conversation_variables=[ - StringVariable(id="topic", name="topic", value="sqlite", selector=["conversation", "topic"]) - ], - rag_pipeline_variables=[], - ) - workflow.created_at = _CREATED_AT - workflow.updated_at = _UPDATED_AT - workflow.updated_by = updater.id - target.workflow_id = workflow.id - tag = Tag(tenant_id=owner.id, type=TagType.APP, name="Owner tag", created_by=author.id) - session.add_all([workflow, tag]) - session.flush() - session.add_all( - [ - TrialApp(app_id=target.id, tenant_id=str(uuid4()), trial_limit=0), - AccountTrialAppRecord(app_id=target.id, account_id=viewer.id, count=100), - TagBinding(tenant_id=owner.id, tag_id=tag.id, target_id=target.id, created_by=author.id), - WorkflowToolProvider( - name="trial-workflow", - label="Trial workflow", - icon="icon", - app_id=target.id, - version="1.0.0", - user_id=author.id, - tenant_id=owner.id, - description="Workflow provider", - parameter_configuration="[]", - ), - ] - ) - factory = sessionmaker(bind=sqlite_engine, expire_on_commit=False) - sessions: list[Session] = [] - - @event.listens_for(factory, "after_begin") - def track_session(session: Session, _transaction: SessionTransaction, _connection: Connection) -> None: - assert all(previous is session or not previous.in_transaction() for previous in sessions) - sessions.append(session) - - catalog = _Catalog(sqlite_engine) - recommendations = RecommendedAppQueryService( - catalog=catalog, trial_apps=TrialAppRepository(factory), trial_enabled=False - ) - services = _ApplicationServices( - app_previews=AppPreviewQueryService( - apps=AppPreviewQueryRepository(session_factory=factory), is_previewable=recommendations.is_previewable - ), - app_preview_details=AppPreviewDetailsRuntime(details=AppPreviewQueryRepository(session_factory=factory)), - recommended_app_queries=recommendations, - ) - for module in (trial_module, admission_module): - monkeypatch.setattr(module, "application_services", lambda: services) - monkeypatch.setattr(login_module, "current_user", viewer) - monkeypatch.setattr(app_module, "current_user", viewer) - monkeypatch.setattr(console_wraps, "_is_setup_completed", lambda: True) - app = Flask(__name__) - app.config.update(TESTING=True, RESTX_ERROR_404_HELP=False) - login_manager = DifyLoginManager() - login_manager.init_app(app) - login_manager.unauthorized_handler(unauthorized_handler) - api = ExternalApi(app) - api.error_handlers = console_api.error_handlers.copy() - api.add_resource(trial_module.AppApi, "/trial-apps/") - api.add_resource(trial_module.AppWorkflowApi, "/trial-apps//workflows") - return _Harness( - app, - owner, - viewer_workspace, - viewer, - author, - updater, - target, - config, - site, - workflow, - tag, - sqlite_session_factory, - sqlite_engine, - catalog, - io, - sessions, - ) - - -def _assert_error(response: TestResponse, status: int, code: str) -> None: - assert response.status_code == status, response.get_json() - body = response.get_json() - assert body["status"] == status - assert body["code"] == code - assert body["message"] - - -@pytest.mark.parametrize( - "mode", [AppMode.CHAT, AppMode.COMPLETION, AppMode.ADVANCED_CHAT, AppMode.WORKFLOW, AppMode.AGENT_CHAT] -) -def test_app_detail_keeps_original_response_fields_and_stored_audit_times(harness: _Harness, mode: AppMode) -> None: - with harness.factory.begin() as session: - session.execute(update(App).where(App.id == harness.target.id).values(mode=mode)) - expected = harness.legacy_detail() - if mode == AppMode.AGENT_CHAT: - expected_config = expected["model_config"] - assert isinstance(expected_config, dict) - # Legacy masking could autoflush a response-only config mutation and bump - # this timestamp. A read-only preview returns the actual stored audit value. - expected_config["updated_at"] = int(_UPDATED_AT.timestamp()) - - response = harness.get() - - assert response.status_code == 200, response.get_json() - assert response.get_json() == expected - body = response.get_json() - assert body["max_active_requests"] == 0 - assert body["enable_site"] is False - assert body["tags"] == [{"id": harness.tag.id, "name": "Owner tag", "type": "app"}] - assert body["api_base_url"] == "http://localhost/v1" - assert body["site"]["app_base_url"] == "http://localhost" - assert body["site"]["access_token"] == "preview-code" - with harness.factory() as session: - assert session.scalar(select(AccountTrialAppRecord.count)) == 100 - assert ( - session.scalar(select(AppModelConfig.updated_at).where(AppModelConfig.id == harness.config.id)) - == _UPDATED_AT - ) - - -def test_legacy_agent_preview_masks_tools_without_upgrading_mode_or_mutating_config(harness: _Harness) -> None: - agent_mode = { - "enabled": True, - "strategy": "function_call", - "tools": [ - { - "provider_type": "builtin", - "provider_id": "example", - "tool_name": "search", - "tool_parameters": {"api_key": "cipher-tool"}, - } - ], - } - with harness.factory.begin() as session: - session.execute( - update(AppModelConfig) - .where(AppModelConfig.id == harness.config.id) - .values(agent_mode=json.dumps(agent_mode)) - ) - app_updated_at = session.scalar(select(App.updated_at).where(App.id == harness.target.id)) - config_updated_at = session.scalar( - select(AppModelConfig.updated_at).where(AppModelConfig.id == harness.config.id) - ) - - response = harness.get() - - assert response.status_code == 200, response.get_json() - body = response.get_json() - assert body["mode"] == "agent-chat" - assert body["model_config"]["agent_mode"]["tools"][0]["tool_parameters"] == {"api_key": "to*******et"} - assert harness.io.tool_calls == [(harness.viewer_workspace.id, harness.target.id, harness.viewer.id)] - assert harness.io.decryptions == [(harness.viewer_workspace.id, "cipher-tool")] - assert harness.viewer_workspace.id != harness.owner.id - with harness.factory() as session: - app = session.get(App, harness.target.id) - config = session.get(AppModelConfig, harness.config.id) - assert app is not None - assert config is not None - assert app.mode == AppMode.CHAT - assert app.updated_at == app_updated_at - assert config.agent_mode == json.dumps(agent_mode) - assert config.updated_at == config_updated_at - - -@pytest.mark.parametrize( - "provider_fields", - [{}, {"provider_id": None}, {"provider_id": 42}, {"provider_id": ""}, {"provider_id": "invalid/provider"}], - ids=["missing", "null", "non-string", "empty", "invalid-format"], -) -def test_invalid_builtin_config_does_not_break_detail_or_valid_tool_masking( - harness: _Harness, monkeypatch: pytest.MonkeyPatch, provider_fields: dict[str, JsonValue] -) -> None: - agent_mode: dict[str, JsonValue] = { - "enabled": True, - "strategy": "function_call", - "tools": [ - { - "provider_type": "builtin", - "tool_name": "broken", - "tool_parameters": {"api_key": "cipher-secret"}, - "credential_id": None, - **provider_fields, - }, - { - "provider_type": "builtin", - "provider_id": "example", - "tool_name": "search", - "tool_parameters": {"api_key": "cipher-tool"}, - }, - ], - } - with harness.factory.begin() as session: - session.execute( - update(AppModelConfig) - .where(AppModelConfig.id == harness.config.id) - .values(agent_mode=json.dumps(agent_mode)) - ) - - def hardcoded_provider(provider_id: str) -> object: - if provider_id == "example": - return object() - raise KeyError(provider_id) - - def tool_runtime(*, tenant_id: str, app_id: str, agent_tool: AgentToolEntity, user_id: str) -> Tool: - if agent_tool.provider_id != "example": - raise ValueError("Invalid historical provider ID") - return harness.io.tool_runtime(tenant_id=tenant_id, app_id=app_id, agent_tool=agent_tool, user_id=user_id) - - plugin_queries: list[str] = [] - - def check_plugins(tenant_id: str, provider_ids: Sequence[GenericProviderID]) -> Sequence[bool]: - assert tenant_id == harness.owner.id - plugin_queries.extend(str(provider_id) for provider_id in provider_ids) - return [False] * len(provider_ids) - - monkeypatch.setattr(ToolManager, "get_hardcoded_provider", hardcoded_provider) - monkeypatch.setattr(ToolManager, "get_agent_tool_runtime", tool_runtime) - monkeypatch.setattr(PluginService, "check_tools_existence", check_plugins) - - response = harness.get() - - assert response.status_code == 200, response.get_json() - body = response.get_json() - assert body["deleted_tools"] == [] - assert body["model_config"]["agent_mode"]["tools"][0]["tool_parameters"] == {} - assert body["model_config"]["agent_mode"]["tools"][1]["tool_parameters"] == {"api_key": "to*******et"} - assert b"cipher-secret" not in response.data - assert plugin_queries == [] - assert harness.io.tool_calls == [(harness.viewer_workspace.id, harness.target.id, harness.viewer.id)] - with harness.factory() as session: - assert session.scalar( - select(AppModelConfig.agent_mode).where(AppModelConfig.id == harness.config.id) - ) == json.dumps(agent_mode) - - -@pytest.mark.parametrize("failure", ["short-config", "runtime", "masking"]) -def test_unmaskable_tool_parameters_are_removed_from_detail_without_changing_storage( - harness: _Harness, monkeypatch: pytest.MonkeyPatch, failure: str -) -> None: - broken_tool: dict[str, JsonValue] = { - "provider_type": "builtin", - "tool_parameters": {"api_key": "cipher-secret"}, - } - if failure != "short-config": - broken_tool.update(provider_id="broken", tool_name="broken") - agent_mode: dict[str, JsonValue] = { - "enabled": True, - "strategy": "function_call", - "tools": [ - broken_tool, - { - "provider_type": "builtin", - "provider_id": "example", - "tool_name": "search", - "tool_parameters": {"api_key": "cipher-tool"}, - }, - ], - } - with harness.factory.begin() as session: - session.execute( - update(AppModelConfig) - .where(AppModelConfig.id == harness.config.id) - .values(agent_mode=json.dumps(agent_mode)) - ) - - class MaskingFailureRuntime(_ToolRuntime): - parameter_queries: int = 0 - - @override - def get_runtime_parameters(self) -> list[ToolParameter]: - self.parameter_queries += 1 - if self.parameter_queries == 2: - raise RuntimeError("Tool metadata became unavailable during masking") - return [] - - parameter = ToolParameter.model_validate( - {"name": "api_key", "label": {"en_US": "Key"}, "type": "secret-input", "form": "form"} - ) - broken_runtime = MaskingFailureRuntime( - SimpleNamespace(parameters=[parameter], identity=SimpleNamespace(name="broken")) - ) - - def tool_runtime(*, tenant_id: str, app_id: str, agent_tool: AgentToolEntity, user_id: str) -> Tool: - if agent_tool.provider_id == "example": - return harness.io.tool_runtime(tenant_id=tenant_id, app_id=app_id, agent_tool=agent_tool, user_id=user_id) - if failure == "runtime": - raise RuntimeError("Tool provider is unavailable") - return cast(Tool, broken_runtime) - - def decrypt(tenant_id: str, token: str) -> str: - if token == "cipher-secret": - assert tenant_id == harness.viewer_workspace.id - return "plain-secret" - return harness.io.decrypt(tenant_id, token) - - monkeypatch.setattr(ToolManager, "get_agent_tool_runtime", tool_runtime) - monkeypatch.setattr(encrypter, "decrypt_token", decrypt) - - response = harness.get() - - assert response.status_code == 200, response.get_json() - tools = response.get_json()["model_config"]["agent_mode"]["tools"] - assert tools[0]["tool_parameters"] == {} - assert tools[1]["tool_parameters"] == {"api_key": "to*******et"} - assert b"cipher-secret" not in response.data - assert b"plain-secret" not in response.data - assert broken_runtime.parameter_queries == (2 if failure == "masking" else 0) - with harness.factory() as session: - assert session.scalar( - select(AppModelConfig.agent_mode).where(AppModelConfig.id == harness.config.id) - ) == json.dumps(agent_mode) - - -def test_workflow_is_anonymous_and_masks_secrets_after_query_session_closes( - harness: _Harness, monkeypatch: pytest.MonkeyPatch -) -> None: - monkeypatch.setattr(login_module, "current_user", None) - monkeypatch.setattr(console_wraps, "_is_setup_completed", lambda: False) - response = harness.get("workflows", csrf_token="") - assert response.status_code == 200, response.get_json() - assert response.get_json() == { - "id": harness.workflow.id, - "graph": {"nodes": [], "edges": []}, - "features": {"suggested_questions_after_answer": {"enabled": False}}, - "hash": harness.workflow.unique_hash, - "version": "draft", - "marked_name": "", - "marked_comment": "", - "created_by": {"id": harness.author.id, "name": "Author", "email": "author@example.com"}, - "created_at": int(_CREATED_AT.timestamp()), - "updated_by": {"id": harness.updater.id, "name": "Updater", "email": "updater@example.com"}, - "updated_at": int(_UPDATED_AT.timestamp()), - "tool_published": True, - "environment_variables": [ - { - "value_type": "secret", - "value": encrypter.full_mask_token(), - "id": "secret", - "name": "api_key", - "description": "", - "selector": ["env", "api_key"], - }, - { - "value_type": "llm", - "value": {"provider": "provider", "name": "model", "mode": "chat"}, - "id": "llm", - "name": "shared_model", - "description": "", - "selector": ["env", "shared_model"], - }, - ], - "conversation_variables": [ - {"id": "topic", "name": "topic", "value_type": "string", "value": "sqlite", "description": ""} - ], - "rag_pipeline_variables": [], - } - assert harness.io.decryptions == [(harness.owner.id, "cipher-workflow")] - assert b"workflow-secret" not in response.data - assert b"cipher-workflow" not in response.data - - -@pytest.mark.parametrize("endpoint", ["detail", "workflows"]) -def test_catalog_only_admission_keeps_preview_available(harness: _Harness, endpoint: _Endpoint) -> None: - with harness.factory.begin() as session: - session.execute(delete(TrialApp).where(TrialApp.app_id == harness.target.id)) - harness.catalog.ids.add(harness.target.id) - response = harness.get(endpoint) - assert response.status_code == 200, response.get_json() - assert harness.catalog.calls == [harness.target.id] - - -@pytest.mark.parametrize("endpoint", ["detail", "workflows"]) -@pytest.mark.parametrize("unavailable", ["unlisted", "disabled", "missing"]) -def test_unavailable_apps_fail_admission(harness: _Harness, endpoint: _Endpoint, unavailable: str) -> None: - with harness.factory.begin() as session: - if unavailable == "unlisted": - session.execute(delete(TrialApp).where(TrialApp.app_id == harness.target.id)) - elif unavailable == "disabled": - session.execute(text("UPDATE apps SET status='disabled' WHERE id=:app_id"), {"app_id": harness.target.id}) - else: - session.execute(delete(App).where(App.id == harness.target.id)) - _assert_error(harness.get(endpoint), 404, "app_not_found") - assert harness.io.decryptions == [] - - -def test_anonymous_app_detail_requires_account(harness: _Harness, monkeypatch: pytest.MonkeyPatch) -> None: - monkeypatch.setattr(login_module, "current_user", None) - response = harness.get() - assert response.status_code == 401 - assert response.get_json() == {"code": "unauthorized", "message": "Unauthorized.", "status": 401} - harness.assert_app_not_queried() - - -@pytest.mark.parametrize(("init_password", "error_code"), [("", "not_setup"), ("setup-password", "not_init_validated")]) -def test_app_detail_requires_setup_before_account_or_app_queries( - harness: _Harness, - monkeypatch: pytest.MonkeyPatch, - config_overrides: Callable[..., None], - init_password: str, - error_code: str, -) -> None: - config_overrides(INIT_PASSWORD=init_password) - monkeypatch.setattr(console_wraps, "_is_setup_completed", lambda: False) - monkeypatch.setattr(login_module, "current_user", None) - _assert_error(harness.get(csrf_token=""), 401, error_code) - harness.assert_app_not_queried() - - -@pytest.mark.parametrize("csrf_token", ["", "invalid-signed-token"], ids=["missing", "invalid"]) -def test_app_detail_rejects_invalid_csrf_before_app_queries(harness: _Harness, csrf_token: str) -> None: - response = harness.get(csrf_token=csrf_token) - _assert_error(response, 401, "unauthorized") - assert response.get_json()["message"] == "CSRF token is missing or invalid." - harness.assert_app_not_queried() - - -def test_app_detail_rejects_csrf_token_for_another_account(harness: _Harness) -> None: - response = harness.get(csrf_token=generate_csrf_token(harness.author.id)) - _assert_error(response, 401, "unauthorized") - assert response.get_json()["message"] == "CSRF token is missing or invalid." - harness.assert_app_not_queried() - - -def test_app_detail_rejects_uninitialized_account_before_app_queries(harness: _Harness) -> None: - harness.viewer.status = AccountStatus.UNINITIALIZED - _assert_error(harness.get(), 400, "account_not_initialized") - harness.assert_app_not_queried() - - -def test_missing_active_workspace_uses_shared_console_error_and_monitoring(harness: _Harness) -> None: - harness.viewer._current_tenant = None - errors: list[BaseException] = [] - - def record_error(_sender: object, *, exception: BaseException) -> None: - errors.append(exception) - - with got_request_exception.connected_to(record_error): - response = harness.get() - - _assert_error(response, 500, "active_workspace_required") - assert response.get_json()["message"] == "Internal Server Error" - assert len(errors) == 1 - assert isinstance(errors[0], ActiveWorkspaceRequiredError) - harness.assert_app_not_queried() - - -def test_deleted_admitted_account_is_unauthorized(harness: _Harness) -> None: - with harness.factory.begin() as session: - session.execute(delete(Account).where(Account.id == harness.viewer.id)) - _assert_error(harness.get(), 401, "unauthorized") - - -def test_missing_site_is_specific_error(harness: _Harness) -> None: - with harness.factory.begin() as session: - session.execute(delete(Site).where(Site.app_id == harness.target.id)) - _assert_error(harness.get(), 403, "app_site_unavailable") - - -def test_missing_model_config_retains_nullable_detail(harness: _Harness) -> None: - with harness.factory.begin() as session: - session.execute(delete(AppModelConfig).where(AppModelConfig.id == harness.config.id)) - response = harness.get() - assert response.status_code == 200, response.get_json() - assert response.get_json()["model_config"] is None - - -def test_agent_detail_keeps_model_less_presentation_config(harness: _Harness) -> None: - with harness.factory.begin() as session: - session.execute(update(App).where(App.id == harness.target.id).values(mode=AppMode.AGENT)) - session.execute(update(AppModelConfig).where(AppModelConfig.id == harness.config.id).values(model=None)) - agent_id, snapshot_id = str(uuid4()), str(uuid4()) - session.add_all( - [ - Agent( - id=agent_id, - tenant_id=harness.owner.id, - app_id=harness.target.id, - name="Preview agent", - scope=AgentScope.ROSTER, - source=AgentSource.AGENT_APP, - active_config_snapshot_id=snapshot_id, - ), - AgentConfigSnapshot( - id=snapshot_id, - tenant_id=harness.owner.id, - agent_id=agent_id, - version=1, - config_snapshot=AgentSoulConfig(), - ), - AgentConfigRevision( - tenant_id=harness.owner.id, - agent_id=agent_id, - current_snapshot_id=snapshot_id, - revision=1, - operation=AgentConfigRevisionOperation.PUBLISH_DRAFT, - ), - ] - ) - - response = harness.get() - - assert response.status_code == 200, response.get_json() - body = response.get_json() - assert body["mode"] == "agent" - assert body["model_config"]["model"] is None - assert body["model_config"]["opening_statement"] == "" - - -@pytest.mark.parametrize("missing", ["reference", "workflow"]) -def test_missing_workflow_is_app_unavailable(harness: _Harness, missing: str) -> None: - with harness.factory.begin() as session: - if missing == "reference": - session.execute(update(App).where(App.id == harness.target.id).values({App.workflow_id: None})) - else: - session.execute(delete(Workflow).where(Workflow.id == harness.workflow.id)) - _assert_error(harness.get("workflows"), 400, "app_unavailable") - - -@pytest.mark.parametrize("endpoint", ["detail", "workflows"]) -def test_head_keeps_get_status_and_headers_without_response_body(harness: _Harness, endpoint: _Endpoint) -> None: - response = harness.get(endpoint) - head = harness.get(endpoint, method="HEAD") - assert response.status_code == head.status_code == 200 - assert head.data == b"" - assert head.headers == response.headers - - -@pytest.mark.parametrize( - ("endpoint", "relation"), [("detail", "config"), ("detail", "workflow"), ("workflows", "workflow")] -) -def test_cross_owner_configuration_is_not_exposed(harness: _Harness, endpoint: _Endpoint, relation: str) -> None: - with harness.factory.begin() as session: - if relation == "config": - session.execute( - update(AppModelConfig).where(AppModelConfig.id == harness.config.id).values(app_id=str(uuid4())) - ) - else: - session.execute(update(Workflow).where(Workflow.id == harness.workflow.id).values(tenant_id=str(uuid4()))) - _assert_error(harness.get(endpoint), 400, "app_unavailable") - assert harness.io.decryptions == [] diff --git a/api/tests/unit_tests/controllers/console/explore/test_trial_suggested_questions.py b/api/tests/unit_tests/controllers/console/explore/test_trial_suggested_questions.py deleted file mode 100644 index df23cd50b9b..00000000000 --- a/api/tests/unit_tests/controllers/console/explore/test_trial_suggested_questions.py +++ /dev/null @@ -1,564 +0,0 @@ -"""Trial suggested questions through HTTP, real ownership queries, and SQLite.""" - -import json -from collections.abc import Callable, Iterator, Mapping, Sequence -from dataclasses import dataclass, field -from decimal import Decimal -from typing import cast -from uuid import uuid4 - -import pytest -from flask import Flask, Request -from sqlalchemy import Connection, Engine, delete, event, select, update -from sqlalchemy.orm import Session, SessionTransaction, sessionmaker -from werkzeug.test import TestResponse - -import controllers.console.explore.trial as trial_module -import controllers.console.explore.trial_app_admission as admission_module -import controllers.console.wraps as console_wraps -import libs.login as login_module -import services.message_service as message_module -from core.errors.error import ModelCurrentlyNotSupportError, ProviderTokenNotInitError, QuotaExceededError -from core.model_context import get_credit_usage_metadata -from core.model_manager import ModelInstance -from core.ops.ops_trace_manager import TraceTask -from enums import DeploymentEdition -from extensions.ext_database import db -from extensions.ext_login import DifyLoginManager, unauthorized_handler -from graphon.model_runtime.entities import AssistantPromptMessage, PromptMessage -from graphon.model_runtime.entities.llm_entities import LLMResult, LLMUsage -from graphon.model_runtime.entities.model_entities import AIModelEntity, ModelType -from graphon.model_runtime.errors.invoke import InvokeError -from libs.external_api import ExternalApi -from models import Account, AccountTrialAppRecord, App, AppMode, Conversation, Message, Tenant, TrialApp -from models.account import AccountStatus -from models.enums import ConversationFromSource -from models.model import AppModelConfig -from repositories.trial_app_repository import TrialAppRepository -from services.app_definition_query_service import AppDefinitionUnavailableError -from services.message_suggested_questions_adapters import MessageSuggestedQuestionsRuntime -from services.message_suggested_questions_service import ( - SuggestedQuestionsAccount, - SuggestedQuestionsActor, - SuggestedQuestionsActorNotFoundError, -) -from services.trial_app_access_service import TrialAppAccessService - - -@dataclass -class _Features: - enabled: bool = True - setup_completed: bool = True - events: list[str] = field(default_factory=list) - - def is_trial_enabled(self) -> bool: - self.events.append("feature") - return self.enabled - - -@dataclass -class _Provider: - read_sessions: list[Session] - app_id: str - model_sessions: list[Session] = field(default_factory=list) - token_calls: int = 0 - tenant_ids: list[str] = field(default_factory=list) - prompts: list[str] = field(default_factory=list) - traces: list[TraceTask] = field(default_factory=list) - questions: list[str] = field(default_factory=lambda: ["Next question?", "Another question?"]) - failure: Exception | None = None - missing_model: bool = False - app_type: str = "chatbot" - - def get_default_model_instance(self, *, tenant_id: str, model_type: ModelType) -> ModelInstance: - assert model_type == ModelType.LLM - self.tenant_ids.append(tenant_id) - session = db.session() - assert session.get(App, self.app_id) is not None - self.model_sessions.append(session) - if self.missing_model: - raise ProviderTokenNotInitError("No default model") - return cast(ModelInstance, self) - - def assert_queries_closed(self) -> None: - assert all( - not session.in_transaction() and not session.identity_map - for session in self.read_sessions + self.model_sessions - ) - - def get_llm_num_tokens(self, prompt_messages: Sequence[PromptMessage]) -> int: - self.assert_queries_closed() - self.token_calls += 1 - if self.failure is not None: - raise self.failure - return len(prompt_messages) - - def get_model_schema(self) -> AIModelEntity: - self.assert_queries_closed() - return AIModelEntity.model_construct(parameter_rules=[]) - - def invoke_llm( - self, - *, - prompt_messages: list[PromptMessage], - model_parameters: Mapping[str, object], - stop: list[str], - stream: bool, - ) -> LLMResult: - self.assert_queries_closed() - assert self.tenant_ids == [self.tenant_ids[0]] * 2 - assert get_credit_usage_metadata() == {"app_type": self.app_type, "created_by": "suggested_questions"} - assert model_parameters == {"max_tokens": 256, "temperature": 0.0} - assert stop == [] - assert stream is False - prompt = prompt_messages[0].get_text_content() - assert "Suggest concise follow-ups" in prompt - self.prompts.append(prompt) - return LLMResult( - model="question-model", - message=AssistantPromptMessage(content=json.dumps(self.questions)), - usage=LLMUsage.empty_usage(), - ) - - def add_trace_task(self, task: TraceTask) -> None: - self.traces.append(task) - - -@dataclass(frozen=True) -class _TrialAppServices: - access: TrialAppAccessService - - -@dataclass(frozen=True) -class _ApplicationServices: - trial_apps: _TrialAppServices - recommended_app_queries: _Features - message_suggested_questions: MessageSuggestedQuestionsRuntime - - -@dataclass(frozen=True) -class _Harness: - app: Flask - account: Account - target: App - trial: TrialApp - conversation: Conversation - message: Message - config: AppModelConfig - factory: sessionmaker[Session] - features: _Features - provider: _Provider - legacy_sessions: list[Session] - - def get(self, *, app_id: str | None = None, message_id: str | None = None) -> TestResponse: - return self.app.test_client().get( - f"/trial-apps/{app_id or self.target.id}/messages/{message_id or self.message.id}/suggested-questions" - ) - - def assert_closed(self) -> None: - assert all( - not session.in_transaction() and not session.identity_map - for session in self.provider.read_sessions + self.legacy_sessions - ) - - def usage(self) -> int | None: - with self.factory() as session: - return session.scalar( - select(AccountTrialAppRecord.count).where( - AccountTrialAppRecord.app_id == self.target.id, - AccountTrialAppRecord.account_id == self.account.id, - ) - ) - - -@pytest.fixture -def harness( - monkeypatch: pytest.MonkeyPatch, - config_overrides: Callable[..., None], - sqlite_session_factory: sessionmaker[Session], - sqlite_engine: Engine, -) -> Iterator[_Harness]: - config_overrides(DEPLOYMENT_EDITION=DeploymentEdition.COMMUNITY, INIT_PASSWORD="", LOGIN_DISABLED=False) - features = _Features() - account = Account(name="Trial viewer", email="trial@example.com") - account._current_tenant = Tenant(name="Viewer workspace") - target = App(tenant_id=str(uuid4()), name="Trial app", mode=AppMode.CHAT, enable_site=True, enable_api=False) - target.id = str(uuid4()) - trial = TrialApp(app_id=target.id, tenant_id=str(uuid4()), trial_limit=3) - config = AppModelConfig( - app_id=target.id, - suggested_questions_after_answer='{"enabled":true,"prompt":"Suggest concise follow-ups"}', - ) - with sqlite_session_factory.begin() as session: - session.add_all([account, target, trial, config]) - session.flush() - target.app_model_config_id = config.id - conversation = Conversation( - app_id=target.id, - app_model_config_id=config.id, - mode=AppMode.CHAT, - name="Trial conversation", - _inputs={}, - from_source=ConversationFromSource.CONSOLE, - from_account_id=account.id, - ) - session.add(conversation) - session.flush() - message = Message( - app_id=target.id, - conversation_id=conversation.id, - _inputs={}, - query="What is a trial?", - message={}, - message_unit_price=Decimal(0), - answer="A way to try an app.", - answer_unit_price=Decimal(0), - currency="USD", - from_source=ConversationFromSource.CONSOLE, - from_account_id=account.id, - ) - session.add(message) - - read_sessions: list[Session] = [] - read_factory = sessionmaker(bind=sqlite_engine, expire_on_commit=False) - - @event.listens_for(read_factory, "after_begin") - def track_read(session: Session, _transaction: SessionTransaction, _connection: Connection) -> None: - read_sessions.append(session) - - provider = _Provider(read_sessions, target.id) - services = _ApplicationServices( - trial_apps=_TrialAppServices( - access=TrialAppAccessService(apps=TrialAppRepository(session_factory=read_factory)) - ), - recommended_app_queries=features, - message_suggested_questions=MessageSuggestedQuestionsRuntime(session_factory=read_factory), - ) - - def setup_completed() -> bool: - features.events.append("setup") - return features.setup_completed - - def csrf(_request: Request, account_id: str) -> None: - assert account_id == account.id - features.events.append("csrf") - - for module in (trial_module, admission_module): - monkeypatch.setattr(module, "application_services", lambda: services) - monkeypatch.setattr(console_wraps, "_is_setup_completed", setup_completed) - monkeypatch.setattr(login_module, "current_user", account) - monkeypatch.setattr(login_module, "check_csrf_token", csrf) - - def model_manager(*, tenant_id: str) -> _Provider: - assert tenant_id == target.tenant_id - return provider - - def trace_manager(*, app_id: str) -> _Provider: - assert app_id == target.id - return provider - - monkeypatch.setattr(message_module.ModelManager, "for_tenant", model_manager) - monkeypatch.setattr(message_module, "TraceQueueManager", trace_manager) - - app = Flask(__name__) - app.config.update(TESTING=True, RESTX_ERROR_404_HELP=False, SQLALCHEMY_DATABASE_URI=str(sqlite_engine.url)) - db.init_app(app) - legacy_sessions: list[Session] = [] - - def track_legacy(session: Session, _transaction: SessionTransaction, _connection: Connection) -> None: - legacy_sessions.append(session) - - event.listen(db.session.session_factory, "after_begin", track_legacy) - login_manager = DifyLoginManager() - login_manager.init_app(app) - login_manager.unauthorized_handler(unauthorized_handler) - api = ExternalApi(app) - api.add_resource( - trial_module.TrialMessageSuggestedQuestionApi, - "/trial-apps//messages//suggested-questions", - ) - yield _Harness( - app, - account, - target, - trial, - conversation, - message, - config, - sqlite_session_factory, - features, - provider, - legacy_sessions, - ) - event.remove(db.session.session_factory, "after_begin", track_legacy) - with app.app_context(): - db.session.remove() - db.engine.dispose() - - -def _assert_error(response: TestResponse, status: int, code: str, message: str | None = None) -> None: - assert response.status_code == status, response.get_json() - body = response.get_json() - assert body["code"] == code - assert body["status"] == status - assert body["message"] - if message is not None: - assert body["message"] == message - assert response.headers["Content-Type"] == "application/json" - - -@pytest.mark.parametrize("mode", [AppMode.CHAT, AppMode.AGENT_CHAT, AppMode.AGENT]) -@pytest.mark.parametrize("questions", [["Next question?", "Another question?"], []]) -def test_questions_keep_response_history_owner_and_usage( - harness: _Harness, questions: list[str], mode: AppMode -) -> None: - harness.provider.questions = questions - harness.provider.app_type = {AppMode.CHAT: "chatbot", AppMode.AGENT_CHAT: "agent", AppMode.AGENT: "agent_v2"}[mode] - with harness.factory.begin() as session: - session.execute(update(App).where(App.id == harness.target.id).values(mode=mode)) - session.execute(update(Conversation).where(Conversation.id == harness.conversation.id).values(mode=mode)) - with harness.factory.begin() as session: - session.add_all( - [ - AccountTrialAppRecord(app_id=harness.target.id, account_id=harness.account.id, count=1), - AccountTrialAppRecord(app_id=harness.target.id, account_id=str(uuid4()), count=3), - ] - ) - - response = harness.get() - - assert response.status_code == 200, response.get_json() - assert response.get_json() == {"data": questions} - assert response.headers["Content-Type"] == "application/json" - assert int(response.headers["Content-Length"]) == len(response.data) - assert len(harness.provider.prompts) == 1 - assert "Human: What is a trial?\nAssistant: A way to try an app." in harness.provider.prompts[0] - assert harness.provider.tenant_ids == [harness.target.tenant_id] * 2 - assert harness.target.tenant_id not in {harness.trial.tenant_id, harness.account.current_tenant_id} - assert len(harness.provider.traces) == 1 - assert harness.features.events == ["setup", "csrf", "feature"] - assert harness.usage() == 1 - assert len(harness.provider.model_sessions) == 2 - harness.assert_closed() - - -@pytest.mark.parametrize("entity", ["message", "conversation"]) -@pytest.mark.parametrize("mismatch", ["missing", "app", "account", "source", "end-user"]) -def test_message_and_conversation_require_complete_owner_chain(harness: _Harness, entity: str, mismatch: str) -> None: - message_id = harness.message.id - with harness.factory.begin() as session: - record = ( - session.get(Message, harness.message.id) - if entity == "message" - else session.get(Conversation, harness.conversation.id) - ) - assert record is not None - match mismatch: - case "missing": - if entity == "message": - message_id = str(uuid4()) - else: - session.delete(record) - case "app": - record.app_id = str(uuid4()) - case "account": - record.from_account_id = str(uuid4()) - case "source": - record.from_source = ConversationFromSource.API - case "end-user": - record.from_end_user_id = str(uuid4()) - - response = harness.get(message_id=message_id) - - _assert_error(response, 404, "not_found", "Message not found" if entity == "message" else "Conversation not found") - assert harness.provider.prompts == [] - assert harness.usage() is None - harness.assert_closed() - - -def test_deleted_conversation_is_not_visible(harness: _Harness) -> None: - with harness.factory.begin() as session: - session.execute(update(Conversation).where(Conversation.id == harness.conversation.id).values(is_deleted=True)) - _assert_error(harness.get(), 404, "not_found", "Conversation not found") - assert harness.provider.prompts == [] - harness.assert_closed() - - -@pytest.mark.parametrize("config", [None, '{"enabled":false}']) -def test_disabled_or_unset_suggestions_have_specific_error(harness: _Harness, config: str | None) -> None: - with harness.factory.begin() as session: - session.execute( - update(AppModelConfig) - .where(AppModelConfig.id == harness.config.id) - .values(suggested_questions_after_answer=config) - ) - _assert_error(harness.get(), 403, "app_suggested_questions_after_answer_disabled") - assert harness.provider.prompts == [] - harness.assert_closed() - - -def test_absent_default_model_keeps_empty_success(harness: _Harness) -> None: - harness.provider.missing_model = True - response = harness.get() - assert response.status_code == 200 - assert response.get_json() == {"data": []} - assert harness.provider.prompts == [] - assert harness.provider.traces == [] - harness.assert_closed() - - -@pytest.mark.parametrize("mode", [AppMode.COMPLETION, AppMode.WORKFLOW]) -def test_unsupported_app_modes_fail_before_runtime(harness: _Harness, mode: AppMode) -> None: - with harness.factory.begin() as session: - session.execute(update(App).where(App.id == harness.target.id).values(mode=mode)) - _assert_error(harness.get(), 400, "not_chat_app") - assert len(harness.provider.read_sessions) == 1 - assert harness.legacy_sessions == [] - harness.assert_closed() - - -@pytest.mark.parametrize( - ("denial", "status", "code"), - [ - ("setup", 401, "not_setup"), - ("feature", 403, "trial_app_feature_disabled"), - ("missing_trial", 403, "trial_app_not_allowed"), - ("missing_app", 403, "trial_app_not_allowed"), - ("quota", 403, "trial_app_limit_exceeded"), - ("uninitialized", 400, "account_not_initialized"), - ("unauthenticated", 401, "unauthorized"), - ], -) -def test_admission_blocks_before_runtime( - harness: _Harness, monkeypatch: pytest.MonkeyPatch, denial: str, status: int, code: str -) -> None: - if denial == "setup": - harness.features.setup_completed = False - elif denial == "feature": - harness.features.enabled = False - elif denial == "uninitialized": - harness.account.status = AccountStatus.UNINITIALIZED - elif denial == "unauthenticated": - monkeypatch.setattr(login_module, "current_user", None) - else: - with harness.factory.begin() as session: - if denial == "missing_trial": - session.execute(delete(TrialApp).where(TrialApp.app_id == harness.target.id)) - elif denial == "missing_app": - session.execute(delete(App).where(App.id == harness.target.id)) - else: - session.add(AccountTrialAppRecord(app_id=harness.target.id, account_id=harness.account.id, count=3)) - response = harness.get() - if denial == "unauthenticated": - assert response.status_code == 401 - assert response.get_json() == {"code": "unauthorized", "message": "Unauthorized.", "status": 401} - else: - _assert_error(response, status, code) - assert harness.legacy_sessions == [] - assert harness.provider.prompts == [] - harness.assert_closed() - - -@pytest.mark.parametrize( - ("failure", "code", "message"), - [ - ( - ProviderTokenNotInitError("Missing tenant credentials"), - "provider_not_initialize", - "Missing tenant credentials", - ), - (QuotaExceededError(), "provider_quota_exceeded", None), - (ModelCurrentlyNotSupportError(), "model_currently_not_support", None), - (InvokeError("Provider unavailable"), "completion_request_error", "Provider unavailable"), - ], -) -def test_token_count_failures_preserve_specific_errors_and_close_session( - harness: _Harness, failure: Exception, code: str, message: str | None -) -> None: - harness.provider.failure = failure - _assert_error(harness.get(), 400, code, message) - assert harness.provider.token_calls == 1 - assert harness.provider.prompts == [] - assert harness.provider.traces == [] - assert harness.usage() is None - harness.assert_closed() - - -@pytest.mark.parametrize("failure", [None, InvokeError("Provider unavailable")]) -def test_runtime_preserves_outer_request_session(harness: _Harness, failure: Exception | None) -> None: - harness.provider.failure = failure - with harness.app.app_context(): - outer_session = db.session() - app_model = outer_session.get(App, harness.target.id) - assert app_model is not None - app_model.name = "Pending caller change" - - response = harness.get() - - assert response.status_code == (200 if failure is None else 400) - assert db.session() is outer_session - assert outer_session.in_transaction() - assert app_model in outer_session.dirty - assert app_model.name == "Pending caller change" - assert harness.provider.model_sessions - for inner_session in harness.provider.model_sessions: - assert inner_session is not outer_session - assert not inner_session.in_transaction() - assert not inner_session.identity_map - harness.assert_closed() - with harness.factory() as session: - saved_app = session.get(App, harness.target.id) - assert saved_app is not None - assert saved_app.name == "Trial app" - - -def test_advanced_chat_without_published_workflow_keeps_empty_success(harness: _Harness) -> None: - with harness.factory.begin() as session: - session.execute(update(App).where(App.id == harness.target.id).values(mode=AppMode.ADVANCED_CHAT)) - response = harness.get() - assert response.status_code == 200 - assert response.get_json() == {"data": []} - assert len(harness.provider.read_sessions) == 2 - assert harness.provider.prompts == [] - assert harness.usage() is None - harness.assert_closed() - - -@pytest.mark.parametrize( - ("failure", "status", "code", "message"), - [ - (AppDefinitionUnavailableError("App changed after admission"), 400, "app_unavailable", None), - (SuggestedQuestionsActorNotFoundError("Account disappeared"), 401, "unauthorized", "Account no longer exists."), - ], -) -def test_reload_errors_have_explicit_http_mapping( - harness: _Harness, - monkeypatch: pytest.MonkeyPatch, - failure: Exception, - status: int, - code: str, - message: str | None, -) -> None: - def reject_reload( - _self: MessageSuggestedQuestionsRuntime, - *, - app_id: str, - app_owner_tenant_id: str, - expected_app_mode: str, - actor: SuggestedQuestionsActor, - message_id: str, - ) -> list[str]: - assert (app_id, app_owner_tenant_id, expected_app_mode, message_id) == ( - harness.target.id, - harness.target.tenant_id, - harness.target.mode, - harness.message.id, - ) - assert actor == SuggestedQuestionsAccount(account_id=harness.account.id, invoke_from="explore") - raise failure - - monkeypatch.setattr(MessageSuggestedQuestionsRuntime, "get_suggested_questions", reject_reload) - _assert_error(harness.get(), status, code, message) - assert harness.provider.prompts == [] - assert harness.usage() is None - harness.assert_closed() diff --git a/api/tests/unit_tests/controllers/console/explore/test_workflow.py b/api/tests/unit_tests/controllers/console/explore/test_workflow.py index 1896d35f85e..0290978f50a 100644 --- a/api/tests/unit_tests/controllers/console/explore/test_workflow.py +++ b/api/tests/unit_tests/controllers/console/explore/test_workflow.py @@ -31,7 +31,7 @@ from tests.unit_tests.controllers.console.explore.test_installed_app_admission i _stop_global_redis, harness, ) -from tests.unit_tests.services.test_app_task_service import _StopRedis +from tests.unit_tests.services.test_app_task_service import _AppTaskServiceStopredis as _StopRedis __all__ = ["_stop_global_redis", "harness"] diff --git a/api/tests/unit_tests/controllers/console/test_knowledge_fs_features.py b/api/tests/unit_tests/controllers/console/knowledge_fs/test_features.py similarity index 100% rename from api/tests/unit_tests/controllers/console/test_knowledge_fs_features.py rename to api/tests/unit_tests/controllers/console/knowledge_fs/test_features.py diff --git a/api/tests/unit_tests/controllers/console/knowledge_fs/test_resources.py b/api/tests/unit_tests/controllers/console/knowledge_fs/test_resources.py new file mode 100644 index 00000000000..b53bcd99050 --- /dev/null +++ b/api/tests/unit_tests/controllers/console/knowledge_fs/test_resources.py @@ -0,0 +1,4166 @@ +from __future__ import annotations + +import ast +import inspect +from collections.abc import Callable, Generator +from contextlib import contextmanager +from dataclasses import dataclass, field +from datetime import UTC, datetime +from hashlib import sha256 +from http import HTTPStatus +from inspect import unwrap +from io import BytesIO +from pathlib import Path +from types import FunctionType +from typing import Literal, Never, cast + +import pytest +from flask import Flask +from flask_restx import Resource +from pydantic import BaseModel +from werkzeug.exceptions import Forbidden, NotFound, ServiceUnavailable + +from controllers.common import wraps as common_wraps +from controllers.console import console_ns, flask_admission, wraps +from controllers.console import wraps as console_wraps +from controllers.console.knowledge_fs import resources +from controllers.console.knowledge_fs import resources as console_resources +from controllers.console.knowledge_fs.error import ( + KnowledgeFSAccessDeniedHTTPError, + KnowledgeFSConflictHTTPError, + KnowledgeFSInvalidRequestHTTPError, + KnowledgeFSOperationUnavailableHTTPError, + KnowledgeFSRateLimitHTTPError, + KnowledgeFSRequestRejectedHTTPError, + KnowledgeFSRequestTooLargeHTTPError, + KnowledgeFSResourceNotFoundHTTPError, + KnowledgeFSSpaceNotFoundHTTPError, + KnowledgeFSTimeoutHTTPError, + KnowledgeFSUpstreamUnavailableHTTPError, +) +from controllers.console.wraps import RBACPermission as _knowledge_fs_product_controllers_RBACPermission +from controllers.service_api import service_api_ns +from controllers.service_api.knowledge_fs import resources as service_resources +from controllers.service_api.knowledge_fs.error import ( + KnowledgeFSInvalidCredentialHTTPError, + KnowledgeFSServiceAccessDeniedHTTPError, + KnowledgeFSServiceConflictHTTPError, + KnowledgeFSServiceInvalidRequestHTTPError, + KnowledgeFSServiceOperationUnavailableHTTPError, + KnowledgeFSServiceRateLimitHTTPError, + KnowledgeFSServiceRequestRejectedHTTPError, + KnowledgeFSServiceRequestTooLargeHTTPError, + KnowledgeFSServiceResourceNotFoundHTTPError, + KnowledgeFSServiceUpstreamUnavailableHTTPError, +) +from core.rbac.entities import RBACPermission as _knowledge_fs_creation_RBACPermission +from enums import DeploymentEdition +from libs import login +from libs.exception import BaseHTTPException +from libs.login import AccountWithTenant +from machinery.context import RequestContext +from models import Account +from models.account import TenantAccountRole +from models.enums import ApiTokenType +from models.knowledge_fs import KnowledgeFSAppSpaceJoinType, KnowledgeFSControlSpaceState +from models.model import ApiToken +from services.enterprise.rbac_service import RBACService +from services.knowledge_fs import source_edit_service +from services.knowledge_fs.capability_broker import KnowledgeFSIssuedProductCapability +from services.knowledge_fs.control_space_commands import KnowledgeFSControlSpaceIntentConflictError +from services.knowledge_fs.download_service import KnowledgeFSDownloadUnavailableError +from services.knowledge_fs.object_storage import KnowledgeFSObjectMetadata +from services.knowledge_fs.product_dto import ( + KnowledgeFSAdmittedQueryRequest, + KnowledgeFSAtomicSourceUpdatePayload, + KnowledgeFSDocumentDownloadDescriptor, + KnowledgeFSDocumentMultimodalManifest, + KnowledgeFSDocumentStagedUploadAcceptedResponse, + KnowledgeFSDocumentStagedUploadPayload, + KnowledgeFSDocumentUploadAcceptedResponse, + KnowledgeFSDurableDeletionAcceptedResponse, + KnowledgeFSLogicalDocumentDeletePayload, + KnowledgeFSNamespacePreviewJobResponse, + KnowledgeFSNamespacePreviewPageResponse, + KnowledgeFSOverviewBaseStatsResponse, + KnowledgeFSOverviewQueryOutcomesResponse, + KnowledgeFSPublicFailureResponse, + KnowledgeFSQueryCreatePayload, + KnowledgeFSSmallFileUploadResponse, + KnowledgeFSSourceImportConfigurationPayload, + KnowledgeFSSourceResponse, + KnowledgeFSSourceSyncPolicyResponse, + KnowledgeFSSourceUpdatePayload, + KnowledgeFSSourceWorkflowResponse, + KnowledgeFSSpaceCreatePayload, + KnowledgeFSSpaceCreateResponse, + KnowledgeFSStagedUploadResponse, + KnowledgeFSStreamCapabilityPayload, +) +from services.knowledge_fs.product_remote import ( + KnowledgeFSOperationUnavailableError, + KnowledgeFSProductRequestRejectedError, + KnowledgeFSProductResourceNotFoundError, + KnowledgeFSRemoteMultipartFile, + KnowledgeFSRemoteSSEResponse, +) +from services.knowledge_fs.runtime import KnowledgeFSRuntime +from services.knowledge_fs.service_api_authorization import ( + KnowledgeFSServiceApiAuthorizationError, + KnowledgeFSServiceApiProfile, + KnowledgeFSServiceApiScopeError, +) +from tests.helpers.unit.config_override import apply_config_overrides +from tests.helpers.unit.controllers.knowledge_fs_controller_fakes import ( + FakeComponent, + FakeRuntime, + Recorder, + invocation, + narrowed, + raising, +) +from tests.helpers.unit.model_factories import make_account + +"""KnowledgeFS creation must preserve legacy roles and Enterprise RBAC admission.""" + + +class TestKnowledgeFsCreationCases: + @staticmethod + @pytest.mark.parametrize( + ("role", "rbac_enabled", "rbac_allowed", "allowed"), + [ + (TenantAccountRole.NORMAL, False, True, False), + (None, False, True, False), + (TenantAccountRole.OWNER, False, True, True), + (TenantAccountRole.ADMIN, False, True, True), + (TenantAccountRole.EDITOR, False, True, True), + (TenantAccountRole.DATASET_OPERATOR, False, True, True), + (TenantAccountRole.NORMAL, True, True, True), + (TenantAccountRole.ADMIN, True, False, False), + ], + ) + def test_create_space_enforces_admission_before_provisioning( + app: Flask, + monkeypatch: pytest.MonkeyPatch, + config_overrides: Callable[..., None], + role: TenantAccountRole | None, + rbac_enabled: bool, + rbac_allowed: bool, + allowed: bool, + ) -> None: + config_overrides( + DEPLOYMENT_EDITION=DeploymentEdition.CLOUD, + KNOWLEDGE_FS_ENABLED=True, + RBAC_ENABLED=rbac_enabled, + LOGIN_DISABLED=False, + ) + account = make_account(role=role) + identity = AccountWithTenant(account, "workspace") + monkeypatch.setattr(login, "_resolve_current_user", lambda: account) + monkeypatch.setattr(login, "check_csrf_token", lambda _request, _user_id: None) + monkeypatch.setattr(wraps, "current_account_with_tenant", lambda: identity) + monkeypatch.setattr(flask_admission, "current_account_with_tenant", lambda: identity) + monkeypatch.setattr(resources, "current_account_with_tenant", lambda: identity) + monkeypatch.setattr(wraps, "check_knowledge_rate_limit", lambda: None) + check_access = Recorder(rbac_allowed) + monkeypatch.setattr(RBACService.CheckAccess, "check", check_access) + create_space = Recorder( + KnowledgeFSSpaceCreateResponse( + control_space_id="space", + state=KnowledgeFSControlSpaceState.PROVISIONING, + operation_id="operation", + model_setup_required=True, + ) + ) + runtime = FakeRuntime(application=FakeComponent(create_space=create_space)) + monkeypatch.setattr(resources, "_console_services", lambda: runtime) + + with app.test_request_context("/knowledge-fs/spaces", method="POST", json={"name": "Test", "slug": "test"}): + if allowed: + result, status = resources.KnowledgeFSSpacesApi().post() + assert status == HTTPStatus.ACCEPTED + assert result["control_space_id"] == "space" + else: + with pytest.raises(Forbidden): + resources.KnowledgeFSSpacesApi().post() + + if allowed: + assert create_space.only_call.kwargs["tenant_id"] == "workspace" + assert create_space.only_call.kwargs["account_id"] == account.id + else: + assert not create_space.calls + if rbac_enabled: + assert check_access.only_call == invocation( + "workspace", + account.id, + scene=_knowledge_fs_creation_RBACPermission.DATASET_CREATE_AND_MANAGEMENT.value, + resource_type=None, + resource_id=None, + ) + else: + assert not check_access.calls + + +_API_ROOT = Path(__file__).resolve().parents[5] + + +@dataclass(frozen=True) +class _FacadeRuntime[T]: + facade: T + + +@dataclass(frozen=True) +class _BrokerRuntime[T]: + broker: T + + +def _knowledge_fs_product_controllers_account(account_id: str = "account-1") -> Account: + account = Account(name="Test User", email="user@example.com") + account.id = account_id + return account + + +def _dataset_api_token(tenant_id: str | None = "tenant-1") -> ApiToken: + api_token = ApiToken() + api_token.id = "token-1" + api_token.tenant_id = tenant_id + api_token.type = ApiTokenType.DATASET + api_token.token = "dataset-test-key" + return api_token + + +def _issued_capability( + *, token: str, operation_id: str, knowledge_space_id: str = "space-1" +) -> KnowledgeFSIssuedProductCapability: + return KnowledgeFSIssuedProductCapability( + token=token, + expires_at=datetime(2030, 1, 1, tzinfo=UTC), + operation_id=operation_id, + knowledge_space_id=knowledge_space_id, + knowledge_space_revision=1, + trace_id="trace-1", + ) + + +def _service_profile(control_space_id: str) -> KnowledgeFSServiceApiProfile: + return KnowledgeFSServiceApiProfile( + tenant_id="tenant-1", + control_space_id=control_space_id, + api_token_id="token-1", + principal_id="token-1", + knowledge_space_id="space-1", + knowledge_space_revision=1, + membership_epoch=0, + space_acl_epoch=0, + external_access_epoch=0, + content_policy_revision=0, + ) + + +def _rbac_wrapper(view: FunctionType) -> FunctionType: + current = inspect.unwrap( + view, + stop=lambda candidate: ( + isinstance(candidate, FunctionType) and "rbac_permission_required" in candidate.__code__.co_qualname + ), + ) + assert isinstance(current, FunctionType) + assert "rbac_permission_required" in current.__code__.co_qualname, "RBAC permission wrapper is missing" + return current + + +def _document_multimodal_manifest(*, object_key: str | None = None) -> KnowledgeFSDocumentMultimodalManifest: + asset_ref = ( + { + "contentType": "image/png", + "objectKey": object_key, + "sha256": "a" * 64, + "variants": { + "thumbnail": { + "contentType": "image/webp", + "objectKey": object_key.replace("full.png", "thumbnail.webp"), + "sha256": "b" * 64, + } + }, + } + if object_key is not None + else None + ) + return KnowledgeFSDocumentMultimodalManifest.model_validate( + { + "artifactHash": "c" * 64, + "createdAt": "2026-08-14T12:00:00.000Z", + "documentAssetId": "018f0d60-7a49-7cc2-9c1b-5b36f18f2c44", + "id": "018f0d60-7a49-7cc2-9c1b-5b36f18f2c45", + "items": [ + { + "assetRef": asset_ref, + "caption": "Architecture", + "endOffset": 42, + "id": "figure:1", + "modality": "image", + "ocrText": "Service A to Service B", + "pageNumber": 2, + "parseElementId": "parse-element-1", + "sectionPath": ["Architecture"], + "startOffset": 42, + "title": "System diagram", + } + ], + "knowledgeSpaceId": "018f0d60-7a49-7cc2-9c1b-5b36f18f2c43", + "manifestVersion": "document-multimodal@1", + "updatedAt": "2026-08-14T12:01:00.000Z", + "version": 3, + } + ) + + +class TestKnowledgeFsProductControllersCases: + @staticmethod + @pytest.fixture(autouse=True) + def _enabled_knowledge_fs(config_overrides: Callable[..., None]) -> None: + config_overrides( + DEPLOYMENT_EDITION=DeploymentEdition.COMMUNITY, + KNOWLEDGE_FS_ENABLED=True, + KNOWLEDGE_FS_COMMUNITY_DEV_ENABLED=True, + ) + + @staticmethod + def test_disabled_knowledge_fs_rejects_console_work_before_preview_upload_or_query_admission( + config_overrides: Callable[..., None], + ) -> None: + config_overrides(KNOWLEDGE_FS_ENABLED=False) + routes = ( + ( + console_resources.KnowledgeFSInitialSourcePreviewApi, + console_resources.KnowledgeFSInitialSourcePreviewApi.post, + (), + ), + ( + console_resources.KnowledgeFSInitialSourcePreviewJobsApi, + console_resources.KnowledgeFSInitialSourcePreviewJobsApi.post, + (), + ), + (console_resources.KnowledgeFSStagedUploadsApi, console_resources.KnowledgeFSStagedUploadsApi.post, ()), + ( + console_resources.KnowledgeFSSpaceQueryAdmissionApi, + console_resources.KnowledgeFSSpaceQueryAdmissionApi.post, + ("space-1",), + ), + ) + for resource_type, view, args in routes: + handler = inspect.unwrap( + view, + stop=lambda view: "_knowledge_fs_errors..decorated" in view.__code__.co_qualname, + ) + with pytest.raises(NotFound): + handler(resource_type(), *args) + + @staticmethod + def test_disabled_knowledge_fs_hides_jwks_even_when_capability_v2_is_configured( + config_overrides: Callable[..., None], + monkeypatch: pytest.MonkeyPatch, + ) -> None: + config_overrides(KNOWLEDGE_FS_ENABLED=False, KNOWLEDGE_FS_CAPABILITY_V2_ENABLED=True) + monkeypatch.setattr( + console_resources.session_factory, + "get_session_maker", + lambda: pytest.fail("disabled JWKS must not access the database"), + ) + with pytest.raises(NotFound): + console_resources.KnowledgeFSJWKSApi().get() + + @staticmethod + def test_disabled_knowledge_fs_rejects_service_api_before_token_validation( + config_overrides: Callable[..., None], + ) -> None: + config_overrides(KNOWLEDGE_FS_ENABLED=False) + handler = inspect.unwrap( + service_resources.KnowledgeFSServiceSourceProvidersApi.get, + stop=lambda view: "_service_api_errors..decorated" in view.__code__.co_qualname, + ) + with pytest.raises(NotFound): + handler(service_resources.KnowledgeFSServiceSourceProvidersApi(), "space-1") + + @staticmethod + def test_console_and_service_api_routes_are_registered() -> None: + console_urls = {url for route in console_ns.resources for url in route.urls if url.startswith("/knowledge-fs/")} + service_urls = { + url for route in service_api_ns.resources for url in route.urls if url.startswith("/knowledge-fs/") + } + + assert { + "/knowledge-fs/spaces", + "/knowledge-fs/uploads", + "/knowledge-fs/uploads/", + "/knowledge-fs/spaces/", + "/knowledge-fs/spaces//tags", + "/knowledge-fs/spaces//permissions", + "/knowledge-fs/spaces//members", + "/knowledge-fs/spaces//app-bindings", + ("/knowledge-fs/spaces//app-bindings//"), + "/knowledge-fs/spaces//external-access", + "/knowledge-fs/spaces//settings", + "/knowledge-fs/spaces//documents", + "/knowledge-fs/spaces//documents//multimodal", + ( + "/knowledge-fs/spaces//documents//multimodal/" + "/asset" + ), + "/knowledge-fs/spaces//logical-documents", + "/knowledge-fs/spaces//logical-documents/download-zip", + "/knowledge-fs/spaces//logical-documents/", + "/knowledge-fs/spaces//logical-documents//download", + "/knowledge-fs/spaces//sources", + "/knowledge-fs/spaces//source-connections", + "/knowledge-fs/spaces//source-providers", + ("/knowledge-fs/spaces//source-connections//refresh"), + "/knowledge-fs/spaces//sources//sync", + "/knowledge-fs/spaces//sources//crawl-preview", + "/knowledge-fs/spaces//sources//workflow-imports", + "/knowledge-fs/spaces//sources//sync-policy", + "/knowledge-fs/spaces//source-workflows/", + "/knowledge-fs/spaces//source-workflows//cancel", + "/knowledge-fs/spaces//source-workflows//retry", + "/knowledge-fs/spaces//source-workflows//pages", + "/knowledge-fs/spaces//source-workflows//selection", + "/knowledge-fs/spaces//sources//async-import", + "/knowledge-fs/spaces//queries", + "/knowledge-fs/spaces//research-tasks", + "/knowledge-fs/spaces//traces", + "/knowledge-fs/spaces//golden-questions/bulk-import", + "/knowledge-fs/spaces//golden-questions/evidence-matches", + "/knowledge-fs/spaces//upload-sessions", + ( + "/knowledge-fs/spaces//upload-sessions/" + "/parts//presign" + ), + ("/knowledge-fs/spaces//upload-sessions//complete"), + ("/knowledge-fs/spaces//upload-sessions//abort"), + ("/knowledge-fs/spaces//upload-sessions//small-file"), + "/knowledge-fs/spaces//query-stream-capability", + "/knowledge-fs/tasks//stream-capability", + "/knowledge-fs/query-stream", + "/knowledge-fs/research-tasks//events", + "/knowledge-fs/.well-known/jwks.json", + }.issubset(console_urls) + assert service_urls == { + "/knowledge-fs/spaces//logical-documents", + "/knowledge-fs/spaces//logical-documents/", + "/knowledge-fs/spaces//document-references/resolve", + "/knowledge-fs/spaces//deletion-jobs/", + "/knowledge-fs/spaces//deletion-jobs//retry", + "/knowledge-fs/spaces//deletion-batches/", + "/knowledge-fs/spaces//research-tasks//resume", + "/knowledge-fs/spaces//background-tasks", + "/knowledge-fs/spaces//settings/migrations/", + "/knowledge-fs/spaces//sources//sync", + "/knowledge-fs/spaces//sources//crawl-preview", + "/knowledge-fs/spaces//sources//crawl-import", + "/knowledge-fs/spaces//sources//workflow-imports", + "/knowledge-fs/spaces//sources//sync-policy", + "/knowledge-fs/spaces//source-connections", + "/knowledge-fs/spaces//source-providers", + "/knowledge-fs/spaces//query-images", + "/knowledge-fs/spaces//source-connections//refresh", + "/knowledge-fs/spaces//source-workflows", + "/knowledge-fs/spaces//source-workflows/", + "/knowledge-fs/spaces//source-workflows//cancel", + "/knowledge-fs/spaces//source-workflows//retry", + "/knowledge-fs/spaces//source-workflows//pages", + "/knowledge-fs/spaces//source-workflows//selection", + "/knowledge-fs/spaces//logical-documents//processing-tasks", + "/knowledge-fs/spaces//logical-documents//processing-tasks/", + "/knowledge-fs/spaces//bulk-jobs/", + "/knowledge-fs/spaces//documents", + "/knowledge-fs/spaces//documents/", + "/knowledge-fs/spaces//documents//outline", + "/knowledge-fs/spaces//documents//revisions", + "/knowledge-fs/spaces//documents//revisions//chunks", + "/knowledge-fs/spaces//documents//revisions//chunks/", + "/knowledge-fs/spaces//documents/bulk", + "/knowledge-fs/spaces//documents/reindex", + "/knowledge-fs/spaces//jobs/", + "/knowledge-fs/spaces//jobs//retry", + "/knowledge-fs/spaces//queries/admission", + "/knowledge-fs/query-stream", + "/knowledge-fs/spaces//research-tasks", + "/knowledge-fs/spaces//research-tasks/", + "/knowledge-fs/spaces//research-tasks//partials", + "/knowledge-fs/spaces//research-tasks/plan", + "/knowledge-fs/spaces//settings", + "/knowledge-fs/spaces//sources", + "/knowledge-fs/spaces//sources/", + "/knowledge-fs/spaces//sources//crawl", + "/knowledge-fs/spaces//sources//files", + "/knowledge-fs/spaces//sources//import", + "/knowledge-fs/spaces//sources//import-files", + "/knowledge-fs/spaces//sources//pages", + "/knowledge-fs/spaces//sources//test", + "/knowledge-fs/spaces//traces", + "/knowledge-fs/spaces//traces/", + "/knowledge-fs/spaces//traces//conflicts", + "/knowledge-fs/spaces//traces//evidence", + "/knowledge-fs/spaces//traces//missing", + } + + @staticmethod + def test_initial_website_preview_job_routes_delegate_to_namespace_capabilities( + monkeypatch: pytest.MonkeyPatch, + ) -> None: + account = _knowledge_fs_product_controllers_account() + preview_service = FakeComponent(require_visible_credential=Recorder()) + facade = FakeComponent( + create_namespace_source_preview=Recorder( + KnowledgeFSNamespacePreviewJobResponse( + job_id="job-1", + status="queued", + configuration_fingerprint="f" * 64, + expires_at="2030-01-01T00:00:00Z", + ) + ), + get_namespace_source_preview=Recorder( + KnowledgeFSNamespacePreviewJobResponse( + job_id="job-1", + status="completed", + configuration_fingerprint="f" * 64, + expires_at="2030-01-01T00:00:00Z", + pages=[ + KnowledgeFSNamespacePreviewPageResponse( + page_id="page-1", + source_url="https://example.com/page", + title="Page", + ) + ], + ) + ), + cancel_namespace_source_preview=Recorder( + KnowledgeFSNamespacePreviewJobResponse( + job_id="job-1", + status="consumed", + configuration_fingerprint="f" * 64, + expires_at="2030-01-01T00:00:00Z", + ) + ), + ) + monkeypatch.setattr(console_resources, "current_account_with_tenant", lambda: (account, "tenant-1")) + monkeypatch.setattr(console_resources, "KnowledgeFSInitialSourcePreviewService", lambda _: preview_service) + monkeypatch.setattr(console_resources, "_console_services", lambda: FakeRuntime(facade=facade)) + app = Flask(__name__) + + with app.test_request_context( + method="POST", + json={ + "kind": "website_crawl", + "credentialId": "credential-1", + "datasource": "website", + "pluginId": "plugin-1", + "provider": "firecrawl", + "parameters": {"url": "https://example.com"}, + }, + ): + response, status = inspect.unwrap(console_resources.KnowledgeFSInitialSourcePreviewJobsApi.post)( + console_resources.KnowledgeFSInitialSourcePreviewJobsApi() + ) + assert status == 202 + assert response == {"job_id": "job-1", "status": "pending"} + assert len(preview_service.require_visible_credential.calls) == 1 + assert len(facade.create_namespace_source_preview.calls) == 1 + + with app.test_request_context(): + response = inspect.unwrap(console_resources.KnowledgeFSInitialSourcePreviewJobApi.get)( + console_resources.KnowledgeFSInitialSourcePreviewJobApi(), "job-1" + ) + assert response["status"] == "completed" + assert response["result"]["pages"] == [ + { + "description": None, + "page_id": "page-1", + "source_url": "https://example.com/page", + "title": "Page", + } + ] + + facade.get_namespace_source_preview.return_value = KnowledgeFSNamespacePreviewJobResponse( + job_id="job-1", + status="queued", + configuration_fingerprint="f" * 64, + expires_at="2030-01-01T00:00:00Z", + ) + with app.test_request_context(): + response = inspect.unwrap(console_resources.KnowledgeFSInitialSourcePreviewJobApi.get)( + console_resources.KnowledgeFSInitialSourcePreviewJobApi(), "job-1" + ) + assert response["status"] == "pending" + assert response["result"] is None + + with app.test_request_context(method="DELETE"): + response = inspect.unwrap(console_resources.KnowledgeFSInitialSourcePreviewJobApi.delete)( + console_resources.KnowledgeFSInitialSourcePreviewJobApi(), "job-1" + ) + assert response == {"job_id": "job-1", "result": None, "status": "completed"} + + @staticmethod + @pytest.mark.parametrize( + ("api_class", "method_name", "path_args"), + [ + ( + console_resources.KnowledgeFSSpaceLogicalDocumentDownloadApi, + console_resources.KnowledgeFSSpaceLogicalDocumentDownloadApi.get, + {"control_space_id": "control-1", "document_id": "document-1"}, + ), + ( + console_resources.KnowledgeFSSpaceLogicalDocumentsDownloadApi, + console_resources.KnowledgeFSSpaceLogicalDocumentsDownloadApi.post, + {"control_space_id": "control-1"}, + ), + ], + ) + def test_document_download_routes_deny_callers_without_dataset_download_permission( + monkeypatch: pytest.MonkeyPatch, + api_class: type, + method_name: FunctionType, + path_args: dict[str, str], + ) -> None: + permission_gate = Recorder(raises=Forbidden()) + apply_config_overrides(monkeypatch, RBAC_ENABLED=True) + monkeypatch.setattr( + common_wraps, + "current_account_with_tenant", + lambda: (_knowledge_fs_product_controllers_account(), "tenant-1"), + ) + monkeypatch.setattr("controllers.common.rbac.checks.RBACService.CheckAccess.check", permission_gate) + # The download routes carry the knowledge rate limit too; it needs a logged-in account. + monkeypatch.setattr(console_wraps, "check_knowledge_rate_limit", lambda: None) + monkeypatch.setattr( + console_resources, + "_console_services", + lambda: pytest.fail("permission denial must happen before KnowledgeFS access"), + ) + + permission_wrapper = _rbac_wrapper(method_name) + with pytest.raises(Forbidden): + permission_wrapper(api_class(), **path_args) + + assert permission_gate.calls == [ + invocation( + "tenant-1", + "account-1", + scene=_knowledge_fs_product_controllers_RBACPermission.DATASET_DOCUMENT_DOWNLOAD, + resource_type=None, + resource_id=None, + ) + ] + + @staticmethod + def test_single_document_download_allows_dataset_download_permission_and_keeps_tenant_scope( + monkeypatch: pytest.MonkeyPatch, + ) -> None: + descriptor = KnowledgeFSDocumentDownloadDescriptor( + document_id="document-1", + filename="guide.md", + mime_type="text/markdown", + object_key="namespaces/tenant-1/spaces/space-1/documents/guide.md", + sha256="a" * 64, + size_bytes=4, + ) + facade = FakeComponent(prepare_logical_document_download=Recorder(descriptor)) + download_service = FakeComponent(load_stream=Recorder(iter((b"body",)))) + permission_gate = Recorder(True) + apply_config_overrides(monkeypatch, RBAC_ENABLED=True) + monkeypatch.setattr( + common_wraps, + "current_account_with_tenant", + lambda: (_knowledge_fs_product_controllers_account(), "tenant-1"), + ) + monkeypatch.setattr("controllers.common.rbac.checks.RBACService.CheckAccess.check", permission_gate) + # The download routes carry the knowledge rate limit too; it needs a logged-in account. + monkeypatch.setattr(console_wraps, "check_knowledge_rate_limit", lambda: None) + monkeypatch.setattr(console_resources, "_actor", lambda: ("account-1", "tenant-1")) + monkeypatch.setattr(console_resources, "_console_services", lambda: FakeRuntime(facade=facade)) + monkeypatch.setattr(console_resources, "KnowledgeFSDownloadService", lambda: download_service) + permission_wrapper = _rbac_wrapper(console_resources.KnowledgeFSSpaceLogicalDocumentDownloadApi.get) + app = Flask(__name__) + + with app.test_request_context(): + response = permission_wrapper( + console_resources.KnowledgeFSSpaceLogicalDocumentDownloadApi(), + control_space_id="control-1", + document_id="document-1", + ) + + assert response.get_data() == b"body" + assert response.headers["Content-Disposition"] == "attachment; filename*=UTF-8''guide.md" + assert permission_gate.calls == [ + invocation( + "tenant-1", + "account-1", + scene=_knowledge_fs_product_controllers_RBACPermission.DATASET_DOCUMENT_DOWNLOAD, + resource_type=None, + resource_id=None, + ) + ] + assert facade.prepare_logical_document_download.calls == [ + invocation( + tenant_id="tenant-1", + account_id="account-1", + control_space_id="control-1", + document_id="document-1", + ) + ] + + @staticmethod + def test_single_document_download_preserves_resource_not_found_after_permission_allows( + monkeypatch: pytest.MonkeyPatch, + ) -> None: + facade = FakeComponent( + prepare_logical_document_download=Recorder( + raises=KnowledgeFSProductResourceNotFoundError("document not found") + ) + ) + permission_gate = Recorder(True) + download_service_factory = Recorder() + apply_config_overrides(monkeypatch, RBAC_ENABLED=True) + monkeypatch.setattr( + common_wraps, + "current_account_with_tenant", + lambda: (_knowledge_fs_product_controllers_account(), "tenant-1"), + ) + monkeypatch.setattr("controllers.common.rbac.checks.RBACService.CheckAccess.check", permission_gate) + # The download routes carry the knowledge rate limit too; it needs a logged-in account. + monkeypatch.setattr(console_wraps, "check_knowledge_rate_limit", lambda: None) + monkeypatch.setattr(console_resources, "_actor", lambda: ("account-1", "tenant-1")) + monkeypatch.setattr(console_resources, "_console_services", lambda: FakeRuntime(facade=facade)) + monkeypatch.setattr(console_resources, "KnowledgeFSDownloadService", download_service_factory) + permission_wrapper = _rbac_wrapper(console_resources.KnowledgeFSSpaceLogicalDocumentDownloadApi.get) + app = Flask(__name__) + + with app.test_request_context(), pytest.raises(KnowledgeFSResourceNotFoundHTTPError): + permission_wrapper( + console_resources.KnowledgeFSSpaceLogicalDocumentDownloadApi(), + control_space_id="foreign-control", + document_id="document-1", + ) + + assert len(permission_gate.calls) == 1 + assert facade.prepare_logical_document_download.calls == [ + invocation( + tenant_id="tenant-1", + account_id="account-1", + control_space_id="foreign-control", + document_id="document-1", + ) + ] + assert download_service_factory.calls == [] + + @staticmethod + def test_single_document_download_maps_storage_unavailable_after_permission_allows( + monkeypatch: pytest.MonkeyPatch, + ) -> None: + descriptor = KnowledgeFSDocumentDownloadDescriptor( + document_id="document-1", + filename="guide.md", + mime_type="text/markdown", + object_key="namespaces/tenant-1/spaces/space-1/documents/guide.md", + sha256="a" * 64, + size_bytes=4, + ) + facade = FakeComponent(prepare_logical_document_download=Recorder(descriptor)) + download_service = FakeComponent( + load_stream=Recorder(raises=KnowledgeFSDownloadUnavailableError("storage unavailable")) + ) + permission_gate = Recorder(True) + apply_config_overrides(monkeypatch, RBAC_ENABLED=True) + monkeypatch.setattr( + common_wraps, + "current_account_with_tenant", + lambda: (_knowledge_fs_product_controllers_account(), "tenant-1"), + ) + monkeypatch.setattr("controllers.common.rbac.checks.RBACService.CheckAccess.check", permission_gate) + # The download routes carry the knowledge rate limit too; it needs a logged-in account. + monkeypatch.setattr(console_wraps, "check_knowledge_rate_limit", lambda: None) + monkeypatch.setattr(console_resources, "_actor", lambda: ("account-1", "tenant-1")) + monkeypatch.setattr(console_resources, "_console_services", lambda: FakeRuntime(facade=facade)) + monkeypatch.setattr(console_resources, "KnowledgeFSDownloadService", lambda: download_service) + permission_wrapper = _rbac_wrapper(console_resources.KnowledgeFSSpaceLogicalDocumentDownloadApi.get) + app = Flask(__name__) + + with app.test_request_context(), pytest.raises(ServiceUnavailable): + permission_wrapper( + console_resources.KnowledgeFSSpaceLogicalDocumentDownloadApi(), + control_space_id="control-1", + document_id="document-1", + ) + + assert len(permission_gate.calls) == 1 + assert download_service.load_stream.calls == [invocation(descriptor)] + + @staticmethod + def test_batch_document_download_maps_storage_unavailable_after_permission_allows( + monkeypatch: pytest.MonkeyPatch, + ) -> None: + descriptor = KnowledgeFSDocumentDownloadDescriptor( + document_id="document-1", + filename="guide.md", + mime_type="text/markdown", + object_key="namespaces/tenant-1/spaces/space-1/documents/guide.md", + sha256="a" * 64, + size_bytes=4, + ) + facade = FakeComponent(prepare_logical_document_download=Recorder(descriptor)) + + @contextmanager + def unavailable_zip() -> Generator[Never]: + raise KnowledgeFSDownloadUnavailableError("storage unavailable") + yield + + download_service = FakeComponent(build_zip_tempfile=Recorder(unavailable_zip())) + permission_gate = Recorder(True) + apply_config_overrides(monkeypatch, RBAC_ENABLED=True) + monkeypatch.setattr( + common_wraps, + "current_account_with_tenant", + lambda: (_knowledge_fs_product_controllers_account(), "tenant-1"), + ) + monkeypatch.setattr("controllers.common.rbac.checks.RBACService.CheckAccess.check", permission_gate) + # The download routes carry the knowledge rate limit too; it needs a logged-in account. + monkeypatch.setattr(console_wraps, "check_knowledge_rate_limit", lambda: None) + monkeypatch.setattr(console_resources, "_actor", lambda: ("account-1", "tenant-1")) + monkeypatch.setattr(console_resources, "_console_services", lambda: FakeRuntime(facade=facade)) + monkeypatch.setattr(console_resources, "KnowledgeFSDownloadService", lambda: download_service) + permission_wrapper = _rbac_wrapper(console_resources.KnowledgeFSSpaceLogicalDocumentsDownloadApi.post) + app = Flask(__name__) + + with ( + app.test_request_context(json={"document_ids": ["document-1"]}), + pytest.raises(ServiceUnavailable), + ): + permission_wrapper( + console_resources.KnowledgeFSSpaceLogicalDocumentsDownloadApi(), + control_space_id="control-1", + ) + + assert len(permission_gate.calls) == 1 + assert download_service.build_zip_tempfile.calls == [invocation([descriptor])] + + @staticmethod + def test_knowledge_fs_request_and_response_schemas_are_registered() -> None: + assert { + "KnowledgeFSSpaceCreatePayload", + "KnowledgeFSAppBindingPayload", + "KnowledgeFSAppBindingListResponse", + "KnowledgeFSSpaceUpdatePayload", + "KnowledgeFSMembersReplacePayload", + "KnowledgeFSExternalAccessPayload", + "KnowledgeFSSettingsPayload", + "KnowledgeFSSourceCreatePayload", + "KnowledgeFSQueryCreatePayload", + "KnowledgeFSQueryStreamCapabilityResponse", + "KnowledgeFSResearchTaskCreatePayload", + "KnowledgeFSStreamCapabilityPayload", + "KnowledgeFSSpaceListResponse", + "KnowledgeFSSpaceDetailResponse", + "KnowledgeFSDocumentStagedUploadPayload", + "KnowledgeFSDocumentStagedUploadAcceptedResponse", + "KnowledgeFSDocumentMultimodalAssetQuery", + "KnowledgeFSDocumentMultimodalManifestResponse", + "KnowledgeFSStreamCapabilityResponse", + "KnowledgeFSJWKSResponse", + "KnowledgeFSSmallFileUploadResponse", + "KnowledgeFSCrawlPreviewPageListQuery", + "KnowledgeFSCrawlPreviewPageListResponse", + "KnowledgeFSCrawlPreviewSelectionPayload", + "KnowledgeFSSourceConnectionCreatePayload", + "KnowledgeFSSourceConnectionListQuery", + "KnowledgeFSSourceConnectionListResponse", + "KnowledgeFSSourceConnectionRefreshPayload", + "KnowledgeFSSourceProviderListResponse", + "KnowledgeFSSourceSyncPolicyPayload", + "KnowledgeFSSourceSyncPolicyResponse", + "KnowledgeFSSourceWorkflowImportPayload", + "KnowledgeFSSourceWorkflowCancelPayload", + "KnowledgeFSSourceWorkflowResponse", + "KnowledgeFSUploadPartPresignPayload", + "KnowledgeFSUploadSessionAbortPayload", + "KnowledgeFSUploadSessionCompletePayload", + "KnowledgeFSUploadSessionCreatePayload", + "KnowledgeFSPresignedUploadResponse", + "KnowledgeFSUploadSessionCreateResponse", + "KnowledgeFSUploadSessionMutationResponse", + "KnowledgeFSGoldenQuestionBulkImportPayload", + "KnowledgeFSGoldenQuestionBulkImportResponse", + "KnowledgeFSGoldenQuestionEvidenceMatchPayload", + "KnowledgeFSGoldenQuestionEvidenceMatchResponse", + }.issubset(console_ns.models) + assert { + "KnowledgeFSQueryCreatePayload", + "KnowledgeFSDocumentListResponse", + "KnowledgeFSResearchTaskListResponse", + "KnowledgeFSSettingsPayload", + "KnowledgeFSSettingsResponse", + "KnowledgeFSSourceCreatePayload", + "KnowledgeFSSourceListResponse", + "KnowledgeFSSourceResponse", + "KnowledgeFSTraceListResponse", + }.issubset(service_api_ns.models) + assert console_ns.models["KnowledgeFSSpaceCreatePayload"].__schema__["additionalProperties"] is False + query_capability_schema = console_ns.models["KnowledgeFSQueryStreamCapabilityResponse"].__schema__ + assert set(query_capability_schema["required"]) == {"expires_at", "operation_id", "token", "url"} + assert query_capability_schema["properties"]["operation_id"]["const"] == "createQuery" + + @staticmethod + def test_document_multimodal_manifest_console_bff_exposes_only_authorized_asset_urls( + monkeypatch: pytest.MonkeyPatch, + ) -> None: + document_id = "018f0d60-7a49-7cc2-9c1b-5b36f18f2c44" + manifest = _document_multimodal_manifest( + object_key=( + f"tenant-1/spaces/018f0d60-7a49-7cc2-9c1b-5b36f18f2c43/documents/{document_id}/assets/figure-1-full.png" + ) + ) + facade = FakeComponent(get_document_multimodal_manifest=Recorder(manifest)) + monkeypatch.setattr(console_resources, "_actor", lambda: ("account-1", "tenant-1")) + monkeypatch.setattr(console_resources, "_console_services", lambda: FakeRuntime(facade=facade)) + apply_config_overrides(monkeypatch, CONSOLE_API_URL="") + app = Flask(__name__) + + with app.test_request_context(): + get = inspect.unwrap(console_resources.KnowledgeFSSpaceDocumentMultimodalManifestApi.get) + response = get( + console_resources.KnowledgeFSSpaceDocumentMultimodalManifestApi(), + control_space_id="control-1", + document_id=document_id, + ) + + item = response["items"][0] + assert item["asset_url"] == ( + f"/console/api/knowledge-fs/spaces/control-1/documents/{document_id}/multimodal/figure%3A1/asset" + ) + assert item["thumbnail_url"] == f"{item['asset_url']}?variant=thumbnail" + assert item["parse_element_id"] == "parse-element-1" + assert "objectKey" not in str(response) + assert "sha256" not in str(response) + assert facade.get_document_multimodal_manifest.calls == [ + invocation( + tenant_id="tenant-1", + account_id="account-1", + control_space_id="control-1", + document_id=document_id, + ) + ] + + @staticmethod + def test_document_multimodal_asset_console_bff_reauthorizes_and_streams_scoped_images( + monkeypatch: pytest.MonkeyPatch, + ) -> None: + document_id = "018f0d60-7a49-7cc2-9c1b-5b36f18f2c44" + object_key = ( + f"tenant-1/spaces/018f0d60-7a49-7cc2-9c1b-5b36f18f2c43/documents/{document_id}/assets/figure-1-full.png" + ) + manifest = _document_multimodal_manifest(object_key=object_key) + facade = FakeComponent(get_document_multimodal_manifest=Recorder(manifest)) + object_storage = FakeComponent( + head_object=Recorder( + KnowledgeFSObjectMetadata( + checksum_sha256_base64="checksum", + content_type="image/png", + key=object_key, + metadata={}, + size_bytes=4, + ) + ) + ) + download_service = FakeComponent(load_stream=Recorder(iter((b"im", b"age")))) + download_service_factory = Recorder(download_service) + monkeypatch.setattr(console_resources, "_actor", lambda: ("account-1", "tenant-1")) + monkeypatch.setattr(console_resources, "_console_services", lambda: FakeRuntime(facade=facade)) + monkeypatch.setattr(console_resources, "KnowledgeFSObjectStorageService", lambda: object_storage) + monkeypatch.setattr(console_resources, "KnowledgeFSDownloadService", download_service_factory) + app = Flask(__name__) + + with app.test_request_context(): + get = inspect.unwrap(console_resources.KnowledgeFSSpaceDocumentMultimodalAssetApi.get) + response = get( + console_resources.KnowledgeFSSpaceDocumentMultimodalAssetApi(), + control_space_id="control-1", + document_id=document_id, + item_id="figure:1", + ) + + assert b"".join(response.response) == b"image" + assert response.content_type == "image/png" + assert response.headers["Content-Disposition"] == "inline" + assert response.headers["Content-Security-Policy"] == "default-src 'none'; sandbox" + assert response.headers["X-Content-Type-Options"] == "nosniff" + assert object_storage.head_object.calls == [invocation(key=object_key)] + assert download_service_factory.calls == [invocation(object_storage=object_storage)] + assert len(download_service.load_stream.calls) == 1 + + @staticmethod + def test_document_multimodal_asset_console_bff_rejects_cross_document_object_keys( + monkeypatch: pytest.MonkeyPatch, + ) -> None: + document_id = "018f0d60-7a49-7cc2-9c1b-5b36f18f2c44" + manifest = _document_multimodal_manifest( + object_key=( + "tenant-1/spaces/018f0d60-7a49-7cc2-9c1b-5b36f18f2c43/" + "documents/018f0d60-7a49-7cc2-9c1b-5b36f18f2cff/assets/foreign.png" + ) + ) + facade = FakeComponent(get_document_multimodal_manifest=Recorder(manifest)) + object_storage_factory = Recorder() + monkeypatch.setattr(console_resources, "_actor", lambda: ("account-1", "tenant-1")) + monkeypatch.setattr(console_resources, "_console_services", lambda: FakeRuntime(facade=facade)) + monkeypatch.setattr(console_resources, "KnowledgeFSObjectStorageService", object_storage_factory) + app = Flask(__name__) + get = inspect.unwrap(console_resources.KnowledgeFSSpaceDocumentMultimodalAssetApi.get) + + with app.test_request_context(), pytest.raises(NotFound): + get( + console_resources.KnowledgeFSSpaceDocumentMultimodalAssetApi(), + control_space_id="control-1", + document_id=document_id, + item_id="figure:1", + ) + + assert object_storage_factory.calls == [] + + @staticmethod + def test_source_workflow_import_contract_is_discriminated_idempotent_and_accepted() -> None: + import_schema = console_ns.models["KnowledgeFSSourceWorkflowImportPayload"].__schema__ + assert import_schema["discriminator"] == { + "propertyName": "kind", + "mapping": { + "online-document-import": "#/components/schemas/KnowledgeFSOnlineDocumentWorkflowImportPayload", + "online-drive-import": "#/components/schemas/KnowledgeFSOnlineDriveWorkflowImportPayload", + }, + } + assert import_schema["oneOf"] == [ + {"$ref": "#/components/schemas/KnowledgeFSOnlineDocumentWorkflowImportPayload"}, + {"$ref": "#/components/schemas/KnowledgeFSOnlineDriveWorkflowImportPayload"}, + ] + + api_doc = console_resources.KnowledgeFSSpaceSourceWorkflowImportApi.post.__apidoc__ + assert api_doc["params"]["Idempotency-Key"]["required"] is True + assert api_doc["params"]["Idempotency-Key"]["minLength"] == 8 + assert api_doc["params"]["Idempotency-Key"]["maxLength"] == 255 + assert set(api_doc["responses"]) == {"202"} + + @staticmethod + def test_small_file_console_bff_reads_only_through_facade_and_returns_no_capability_material( + monkeypatch: pytest.MonkeyPatch, + ) -> None: + calls: list[dict[str, object]] = [] + + class Facade: + def upload_small_file(self, **kwargs: object) -> KnowledgeFSSmallFileUploadResponse: + calls.append({name: value for name, value in kwargs.items() if name != "body_reader"}) + body_reader = kwargs["body_reader"] + assert callable(body_reader) + assert body_reader(8 * 1024 * 1024) == b"tiny" + return KnowledgeFSSmallFileUploadResponse.model_validate( + { + "session": { + "compilationJobId": "compilation-1", + "completedAt": 2_000_000, + "documentAssetId": "asset-1", + "expectedSizeBytes": 4, + "expiresAt": 2_060_000, + "id": "session-1", + "mode": "small_fallback", + "status": "completed", + } + } + ) + + monkeypatch.setattr(console_resources, "_actor", lambda: ("account-1", "tenant-1")) + monkeypatch.setattr(console_resources, "_console_services", lambda: _FacadeRuntime(facade=Facade())) + app = Flask(__name__) + + with app.test_request_context( + method="POST", + data={"file": (BytesIO(b"tiny"), "small.txt")}, + content_type="multipart/form-data", + ): + post = inspect.unwrap(console_resources.KnowledgeFSSpaceSmallFileUploadApi.post) + response = post(console_resources.KnowledgeFSSpaceSmallFileUploadApi(), "control-1", "session-1") + + assert response["session"]["status"] == "completed" + assert "token" not in str(response).lower() + assert "url" not in str(response).lower() + assert calls == [ + { + "tenant_id": "tenant-1", + "account_id": "account-1", + "control_space_id": "control-1", + "upload_session_id": "session-1", + } + ] + + @staticmethod + def test_document_upload_console_bff_reads_only_through_facade_and_returns_accepted_response( + monkeypatch: pytest.MonkeyPatch, + ) -> None: + calls: list[dict[str, object]] = [] + + class Facade: + def create_document(self, **kwargs: object) -> KnowledgeFSDocumentUploadAcceptedResponse: + calls.append({name: value for name, value in kwargs.items() if name != "body_reader"}) + body_reader = kwargs["body_reader"] + assert callable(body_reader) + assert body_reader(15 * 1024 * 1024) == KnowledgeFSRemoteMultipartFile( + filename="notes.md", + content_type="text/markdown", + body=b"# Notes", + ) + return KnowledgeFSDocumentUploadAcceptedResponse.model_validate( + { + "asset": { + "createdAt": "2030-01-01T00:00:00Z", + "filename": "notes.md", + "id": "asset-1", + "knowledgeSpaceId": "space-1", + "metadata": {}, + "mimeType": "text/markdown", + "objectKey": "documents/asset-1/notes.md", + "parserStatus": "pending", + "sha256": "sha256", + "sizeBytes": 7, + "sourceId": None, + "updatedAt": None, + "version": 1, + }, + "assetStatusUrl": "/knowledge-spaces/space-1/documents/asset-1", + "compilationJob": {"id": "job-1", "stage": "queued"}, + "documentRevision": 1, + "logicalDocument": {"id": "document-1", "revision": 1}, + "logicalDocumentId": "document-1", + "statusUrl": "/knowledge-spaces/space-1/logical-documents/document-1/tasks/job-1", + } + ) + + monkeypatch.setattr(console_resources, "_actor", lambda: ("account-1", "tenant-1")) + monkeypatch.setattr(console_resources, "_console_services", lambda: _FacadeRuntime(facade=Facade())) + app = Flask(__name__) + + with app.test_request_context( + method="POST", + data={"file": (BytesIO(b"# Notes"), "notes.md", "text/markdown")}, + content_type="multipart/form-data", + ): + post = inspect.unwrap(console_resources.KnowledgeFSSpaceDocumentsApi.post) + response, status = post(console_resources.KnowledgeFSSpaceDocumentsApi(), "control-1") + + assert status == 202 + assert response["logical_document_id"] == "document-1" + assert calls == [ + { + "tenant_id": "tenant-1", + "account_id": "account-1", + "control_space_id": "control-1", + } + ] + + @staticmethod + def test_staged_document_claim_passes_only_upload_id_to_the_admission_service( + monkeypatch: pytest.MonkeyPatch, + ) -> None: + staged_uploads = FakeComponent( + claim=Recorder( + KnowledgeFSDocumentStagedUploadAcceptedResponse( + upload_id="staged-1", + document_asset_id="asset-1", + compilation_job_id="job-1", + ) + ) + ) + monkeypatch.setattr(console_resources, "_actor", lambda: ("account-1", "tenant-1")) + monkeypatch.setattr(console_resources, "_staged_uploads", lambda: staged_uploads) + app = Flask(__name__) + + with app.test_request_context(method="POST", json={"upload_id": "staged-1"}): + post = inspect.unwrap(console_resources.KnowledgeFSSpaceDocumentsApi.post) + response, status = post(console_resources.KnowledgeFSSpaceDocumentsApi(), "control-1") + + assert status == 202 + assert response == { + "status": "accepted", + "upload_id": "staged-1", + "document_asset_id": "asset-1", + "compilation_job_id": "job-1", + } + assert len(staged_uploads.claim.calls) == 1 + claimed = narrowed(staged_uploads.claim.last_call.kwargs["payload"], KnowledgeFSDocumentStagedUploadPayload) + assert claimed.upload_id == "staged-1" + + @staticmethod + @pytest.mark.parametrize("file_size_limit_mb", [15, 50]) + def test_workspace_staging_upload_persists_before_space_creation( + monkeypatch: pytest.MonkeyPatch, + file_size_limit_mb: int, + ) -> None: + admission_events: list[str] = [] + + class RecordingAdmission: + active = False + + @contextmanager + def admit(self, *, reserved_bytes: int) -> Generator[None]: + admission_events.append(f"admit:{reserved_bytes}") + self.active = True + try: + yield + finally: + self.active = False + admission_events.append("release") + + admission = RecordingAdmission() + staged = KnowledgeFSStagedUploadResponse( + id="staged-1", + file_name="notes.md", + content_type="text/markdown", + size_bytes=7, + status="uploaded", + expires_at=datetime(2030, 1, 1, tzinfo=UTC), + ) + + def stage(**_: object) -> KnowledgeFSStagedUploadResponse: + assert admission.active + admission_events.append("stage") + return staged + + staged_uploads = FakeComponent(stage=Recorder(effect=stage)) + account = _knowledge_fs_product_controllers_account() + monkeypatch.setattr(console_resources, "current_account_with_tenant", lambda: (account, "tenant-1")) + monkeypatch.setattr(console_resources, "_staged_uploads", lambda: staged_uploads) + monkeypatch.setattr(console_resources, "DEFAULT_KNOWLEDGE_FS_BUFFERED_UPLOAD_ADMISSION", admission) + monkeypatch.setattr( + console_resources.FeatureService, + "get_knowledge_file_size_limit", + lambda _tenant_id: file_size_limit_mb, + ) + app = Flask(__name__) + + with app.test_request_context( + method="POST", + data={"file": (BytesIO(b"# Notes"), "notes.md", "text/markdown")}, + content_type="multipart/form-data", + ): + post = inspect.unwrap(console_resources.KnowledgeFSStagedUploadsApi.post) + response, status = post(console_resources.KnowledgeFSStagedUploadsApi()) + + assert status == 201 + assert response["id"] == "staged-1" + assert admission_events == [f"admit:{file_size_limit_mb * 1024 * 1024}", "stage", "release"] + assert staged_uploads.stage.calls == [ + invocation( + tenant_id="tenant-1", + account=account, + file_name="notes.md", + content_type="text/markdown", + body=b"# Notes", + file_size_limit_mb=file_size_limit_mb, + ) + ] + + @staticmethod + def test_workspace_staging_zero_byte_plan_limit_returns_request_too_large_before_admission( + monkeypatch: pytest.MonkeyPatch, + ) -> None: + class FailingAdmission: + def admit(self, *, reserved_bytes: int) -> Never: + _ = reserved_bytes + raise AssertionError("a disabled upload limit must fail before admission") + + staged_uploads = FakeComponent(stage=Recorder()) + account = _knowledge_fs_product_controllers_account() + monkeypatch.setattr(console_resources, "current_account_with_tenant", lambda: (account, "tenant-1")) + monkeypatch.setattr(console_resources, "_staged_uploads", lambda: staged_uploads) + monkeypatch.setattr(console_resources, "DEFAULT_KNOWLEDGE_FS_BUFFERED_UPLOAD_ADMISSION", FailingAdmission()) + monkeypatch.setattr( + console_resources.FeatureService, + "get_knowledge_file_size_limit", + lambda _tenant_id: 0, + ) + app = Flask(__name__) + + with app.test_request_context( + method="POST", + data={"file": (BytesIO(b"# Notes"), "notes.md", "text/markdown")}, + content_type="multipart/form-data", + ): + post = inspect.unwrap(console_resources.KnowledgeFSStagedUploadsApi.post) + with pytest.raises(KnowledgeFSProductRequestRejectedError) as rejected: + post(console_resources.KnowledgeFSStagedUploadsApi()) + + assert rejected.value.status_code == 413 + assert staged_uploads.stage.calls == [] + + @staticmethod + def test_logical_document_delete_accepts_initial_row_version( + monkeypatch: pytest.MonkeyPatch, + ) -> None: + calls: list[dict[str, object]] = [] + + class Facade: + def delete_logical_document(self, **kwargs: object) -> KnowledgeFSDurableDeletionAcceptedResponse: + calls.append(kwargs) + payload = kwargs["payload"] + assert isinstance(payload, KnowledgeFSLogicalDocumentDeletePayload) + assert payload.expected_revision == 0 + return KnowledgeFSDurableDeletionAcceptedResponse.model_validate( + { + "job": { + "checkpoint": "requested", + "createdAt": "2030-01-01T00:00:00Z", + "id": "00000000-0000-4000-8000-000000000001", + "knowledgeSpaceId": "space-1", + "runState": "queued", + "targetId": "00000000-0000-4000-8000-000000000002", + "targetType": "logical_document", + "updatedAt": "2030-01-01T00:00:00Z", + }, + "statusUrl": "/deletion-jobs/job-1", + } + ) + + monkeypatch.setattr(console_resources, "_actor", lambda: ("account-1", "tenant-1")) + monkeypatch.setattr(console_resources, "_console_services", lambda: _FacadeRuntime(facade=Facade())) + app = Flask(__name__) + + with app.test_request_context( + method="DELETE", + json={"expectedRevision": 0}, + headers={"Idempotency-Key": "delete-logical-document-once"}, + ): + delete = inspect.unwrap(console_resources.KnowledgeFSSpaceLogicalDocumentApi.delete) + response, status = delete( + console_resources.KnowledgeFSSpaceLogicalDocumentApi(), + "control-1", + "document-1", + ) + + assert status == 202 + assert response["job"]["target_type"] == "logical_document" + assert len(calls) == 1 + assert {name: value for name, value in calls[0].items() if name != "payload"} == { + "tenant_id": "tenant-1", + "account_id": "account-1", + "control_space_id": "control-1", + "document_id": "document-1", + "idempotency_key": "delete-logical-document-once", + } + + @staticmethod + def test_small_file_console_bff_maps_oversize_to_413() -> None: + app = Flask(__name__) + with app.test_request_context( + method="POST", + data={"file": (BytesIO(b"large"), "large.bin")}, + content_type="multipart/form-data", + ): + with pytest.raises(KnowledgeFSProductRequestRejectedError) as oversized: + console_resources._read_small_file_body(4) + assert oversized.value.status_code == 413 + + @console_resources._knowledge_fs_errors + def reject() -> Never: + raise KnowledgeFSProductRequestRejectedError(status_code=413) + + with pytest.raises(KnowledgeFSRequestTooLargeHTTPError): + reject() + + @staticmethod + def test_console_knowledge_rate_limit_covers_every_write_download_and_retrieval_entrypoint() -> None: + """Mirror the legacy dataset console: every mutating, download, or retrieval route is rate limited. + + Pure reads stay unlimited. The exemptions are workspace-level source previews (like legacy + data source browsing), a staged-upload discard, the deprecated buffered query route that + fails closed anyway, and the internal capability-token stream transport. + """ + space = "/knowledge-fs/spaces/" + exempt_writes = { + ("post", "/knowledge-fs/source-provider-preview"), + ("post", "/knowledge-fs/source-provider-preview/jobs"), + ("delete", "/knowledge-fs/source-provider-preview/jobs/"), + ("delete", "/knowledge-fs/uploads/"), + ("post", f"{space}/queries"), + ("post", "_QUERY_STREAM_PROXY_PATH"), + } + rate_limited_reads = {("get", f"{space}/logical-documents//download")} + tree = ast.parse( + Path(console_resources.__file__).read_text(encoding="utf-8"), + filename=console_resources.__file__, + ) + seen: set[tuple[str, str]] = set() + + for class_node in (node for node in tree.body if isinstance(node, ast.ClassDef)): + routes = [ + ast.unparse(decorator.args[0]).strip("'\"") + for decorator in class_node.decorator_list + if isinstance(decorator, ast.Call) + and isinstance(decorator.func, ast.Attribute) + and decorator.func.attr == "route" + ] + if not routes: + continue + for method_node in ( + node + for node in class_node.body + if isinstance(node, ast.FunctionDef) and node.name in {"delete", "get", "patch", "post", "put"} + ): + key = (method_node.name, routes[0]) + seen.add(key) + rate_limited = any( + isinstance(decorator, ast.Call) + and isinstance(decorator.func, ast.Name) + and decorator.func.id == "cloud_edition_billing_rate_limit_check" + and len(decorator.args) == 1 + and isinstance(decorator.args[0], ast.Constant) + and decorator.args[0].value == "knowledge" + for decorator in method_node.decorator_list + ) + expected = key in rate_limited_reads if method_node.name == "get" else key not in exempt_writes + assert rate_limited == expected, f"{key} rate limited={rate_limited}, expected {expected}" + + assert exempt_writes <= seen + assert rate_limited_reads <= seen + + @staticmethod + def test_space_update_and_delete_publish_their_actual_http_status_contracts() -> None: + patch_responses = console_resources.KnowledgeFSSpaceApi.patch.__apidoc__["responses"] + delete_responses = console_resources.KnowledgeFSSpaceApi.delete.__apidoc__["responses"] + + assert set(patch_responses) == {"200"} + assert set(delete_responses) == {"204"} + + @staticmethod + def test_space_create_profile_intent_matches_the_exact_kfs_pending_configuration_shape() -> None: + payload = KnowledgeFSSpaceCreatePayload.model_validate( + { + "name": "Technical space", + "slug": "technical-space", + "embedding": { + "pluginId": "langgenius/openai", + "provider": "openai", + "model": "text-embedding-3-small", + }, + "retrieval": { + "defaultMode": "deep", + "reasoningModel": { + "pluginId": "langgenius/openai", + "provider": "openai", + "model": "gpt-4.1-mini", + }, + "rerank": { + "enabled": True, + "model": { + "pluginId": "langgenius/cohere", + "provider": "cohere", + "model": "rerank-v3.5", + }, + }, + "scoreThreshold": {"enabled": True, "stage": "mode-final", "value": 0.6}, + "topK": 20, + }, + } + ) + + assert payload.embedding is not None + assert payload.embedding is not None + assert payload.embedding.model_dump(mode="json", by_alias=True) == { + "pluginId": "langgenius/openai", + "provider": "openai", + "model": "text-embedding-3-small", + } + assert payload.retrieval is not None + assert payload.retrieval is not None + assert payload.retrieval.model_dump(mode="json", by_alias=True) == { + "defaultMode": "deep", + "reasoningModel": { + "pluginId": "langgenius/openai", + "provider": "openai", + "model": "gpt-4.1-mini", + }, + "rerank": { + "enabled": True, + "model": { + "pluginId": "langgenius/cohere", + "provider": "cohere", + "model": "rerank-v3.5", + }, + }, + "scoreThreshold": {"enabled": True, "stage": "mode-final", "value": 0.6}, + "topK": 20, + } + + with pytest.raises(ValueError, match="retrieval.defaultMode"): + KnowledgeFSSpaceCreatePayload.model_validate( + { + "name": "Invalid preset", + "slug": "invalid-preset", + "embedding": {"pluginId": "plugin", "provider": "provider", "model": "embedding"}, + "retrieval": { + "defaultMode": "auto", + "reasoningModel": {"pluginId": "plugin", "provider": "provider", "model": "reasoning"}, + "rerank": { + "enabled": True, + "model": {"pluginId": "plugin", "provider": "provider", "model": "rerank"}, + }, + "scoreThreshold": {"enabled": False, "stage": "mode-final"}, + "topK": 10, + }, + } + ) + + @staticmethod + def test_space_create_allows_deferred_model_setup_but_rejects_incomplete_fast_profile() -> None: + payload = KnowledgeFSSpaceCreatePayload.model_validate( + { + "name": "Setup later", + "slug": "setup-later", + } + ) + + assert payload.embedding is None + assert payload.retrieval is None + + with pytest.raises(ValueError, match="requires an embedding model"): + KnowledgeFSSpaceCreatePayload.model_validate( + { + "name": "Incomplete fast profile", + "slug": "incomplete-fast-profile", + "retrieval": { + "defaultMode": "fast", + "reasoningModel": { + "pluginId": "plugin", + "provider": "provider", + "model": "reasoning", + }, + "rerank": { + "enabled": True, + "model": {"pluginId": "plugin", "provider": "provider", "model": "rerank"}, + }, + "scoreThreshold": {"enabled": False, "stage": "mode-final"}, + "topK": 10, + }, + } + ) + + @staticmethod + def test_jwks_http_resource_returns_only_public_keys(monkeypatch: pytest.MonkeyPatch) -> None: + public_jwks = { + "keys": [ + {"alg": "RS256", "e": "AQAB", "kid": "current", "kty": "RSA", "n": "modulus-1", "use": "sig"}, + {"alg": "RS256", "e": "AQAB", "kid": "previous", "kty": "RSA", "n": "modulus-2", "use": "sig"}, + ] + } + issuer = FakeComponent(public_jwks=Recorder(public_jwks)) + apply_config_overrides( + monkeypatch, + KNOWLEDGE_FS_CAPABILITY_V2_ENABLED=True, + KNOWLEDGE_FS_JWKS_CACHE_MAX_AGE_SECONDS=123, + ) + monkeypatch.setattr(console_resources.session_factory, "get_session_maker", lambda: object()) + monkeypatch.setattr(console_resources, "create_configured_knowledge_fs_capability_issuer", lambda **_: issuer) + app = Flask(__name__) + + with app.app_context(): + response = console_resources.KnowledgeFSJWKSApi().get() + + payload = response.get_json() + assert payload == public_jwks + assert response.headers["Cache-Control"] == "public, max-age=123, must-revalidate" + assert not {"d", "p", "q", "dp", "dq", "qi"} & set().union(*(key.keys() for key in payload["keys"])) + + @staticmethod + def test_service_profile_authorizes_the_dataset_key_for_the_route_space(monkeypatch: pytest.MonkeyPatch) -> None: + profile = KnowledgeFSServiceApiProfile( + tenant_id="tenant-1", + control_space_id="control-2", + api_token_id="token-1", + principal_id="token-1", + knowledge_space_id="space-1", + knowledge_space_revision=1, + membership_epoch=0, + space_acl_epoch=0, + external_access_epoch=0, + content_policy_revision=0, + ) + authorization = FakeComponent(authorize=Recorder(profile)) + runtime = FakeRuntime(service_api_authorization=authorization) + monkeypatch.setattr(service_resources, "validate_and_get_api_token", Recorder(_dataset_api_token())) + app = Flask(__name__) + with app.test_request_context(headers={"Authorization": "Bearer dataset-test-key"}): + result = service_resources._profile( + runtime, # type: ignore[arg-type] + operation_id="listDocuments", + control_space_id="control-2", + ) + + assert result is profile + assert authorization.authorize.calls == [ + invocation( + api_token_id="token-1", + tenant_id="tenant-1", + control_space_id="control-2", + ) + ] + + @staticmethod + def test_service_profile_applies_the_knowledge_rate_limit_before_authorizing( + monkeypatch: pytest.MonkeyPatch, + ) -> None: + """Every KnowledgeFS service route shares the legacy dataset service API rate limit.""" + authorization = FakeComponent(authorize=Recorder()) + runtime = FakeRuntime(service_api_authorization=authorization) + api_token = _dataset_api_token() + monkeypatch.setattr(service_resources, "validate_and_get_api_token", Recorder(api_token)) + rate_limit = Recorder(raises=Forbidden("rate limited")) + monkeypatch.setattr(service_resources, "check_knowledge_rate_limit", rate_limit) + + with pytest.raises(Forbidden, match="rate limited"): + service_resources._profile( + runtime, # type: ignore[arg-type] + operation_id="listDocuments", + control_space_id="control-2", + ) + + assert rate_limit.calls == [invocation(api_token)] + assert authorization.authorize.calls == [] + + @staticmethod + def test_service_profile_rejects_a_dataset_key_without_a_workspace(monkeypatch: pytest.MonkeyPatch) -> None: + authorization = FakeComponent(authorize=Recorder()) + runtime = FakeRuntime(service_api_authorization=authorization) + monkeypatch.setattr( + service_resources, "validate_and_get_api_token", Recorder(_dataset_api_token(tenant_id=None)) + ) + + with pytest.raises(KnowledgeFSServiceApiAuthorizationError, match="not workspace-scoped"): + service_resources._profile( + runtime, # type: ignore[arg-type] + operation_id="listDocuments", + control_space_id="control-2", + ) + + assert authorization.authorize.calls == [] + + @staticmethod + def test_product_modules_do_not_import_dify_dataset_or_document_services() -> None: + paths = [ + *(path for path in _API_ROOT.glob("services/knowledge_fs/*.py") if not path.name.startswith("upgrade_")), + *_API_ROOT.glob("controllers/console/knowledge_fs/*.py"), + *_API_ROOT.glob("controllers/service_api/knowledge_fs/*.py"), + ] + forbidden_modules = { + "models.dataset", + "services.dataset_service", + "services.document_service", + "controllers.console.datasets", + "controllers.service_api.dataset", + } + + for path in paths: + tree = ast.parse(path.read_text(encoding="utf-8"), filename=str(path)) + imported_modules = set() + for node in ast.walk(tree): + if not isinstance(node, ast.ImportFrom) or node.module is None: + continue + # Native vector credentials share the legacy models.dataset module. + # This bridge may read the tenant binding, never Dataset/Document data. + if ( + path == _API_ROOT / "services/knowledge_fs/vector_store.py" + and node.module == "models.dataset" + and {alias.name for alias in node.names} == {"TidbAuthBinding"} + ): + continue + imported_modules.add(node.module) + assert not any( + module == forbidden or module.startswith(f"{forbidden}.") + for module in imported_modules + for forbidden in forbidden_modules + ), path + + @staticmethod + def test_research_task_stream_url_binds_task_and_parent_space_without_token( + monkeypatch: pytest.MonkeyPatch, + ) -> None: + apply_config_overrides(monkeypatch, CONSOLE_API_URL="https://dify.test") + url = console_resources._research_task_events_url( + task_id="task/one", + knowledge_space_id="space one", + ) + + assert ( + url + == "https://dify.test/console/api/knowledge-fs/research-tasks/task%2Fone/events?knowledgeSpaceId=space+one" + ) + assert "token" not in url.lower() + + @staticmethod + def test_query_stream_capability_issues_exact_space_grant_without_token_in_url( + monkeypatch: pytest.MonkeyPatch, + ) -> None: + calls: list[dict[str, object]] = [] + + class Broker: + def issue_interactive(self, **kwargs: object) -> KnowledgeFSIssuedProductCapability: + calls.append(kwargs) + return _issued_capability(token="query-capability", operation_id="createQuery") + + apply_config_overrides(monkeypatch, CONSOLE_API_URL="https://dify.test") + monkeypatch.setattr(console_resources, "_actor", lambda: ("account-1", "tenant-1")) + monkeypatch.setattr( + console_resources, + "_console_services", + lambda: _BrokerRuntime(broker=Broker()), + ) + app = Flask(__name__) + + with app.test_request_context(): + post = inspect.unwrap(console_resources.KnowledgeFSSpaceQueryStreamCapabilityApi.post) + response = post(console_resources.KnowledgeFSSpaceQueryStreamCapabilityApi(), "control-1") + + assert response == { + "expires_at": "2030-01-01T00:00:00Z", + "operation_id": "createQuery", + "token": "query-capability", + "url": "https://dify.test/console/api/knowledge-fs/query-stream", + } + assert calls == [ + { + "account_id": "account-1", + "control_space_id": "control-1", + "operation_id": "createQuery", + "tenant_id": "tenant-1", + } + ] + assert "token" not in response["url"].lower() + assert console_resources.KnowledgeFSSpaceQueryStreamCapabilityApi.post.__apidoc__["deprecated"] is True + + @staticmethod + def test_query_admission_binds_validated_mode_to_resolved_kfs_space(monkeypatch: pytest.MonkeyPatch) -> None: + calls: list[dict[str, object]] = [] + + class Broker: + def issue_interactive(self, **kwargs: object) -> KnowledgeFSIssuedProductCapability: + calls.append(kwargs) + return _issued_capability(token="query-capability", operation_id="createQuery") + + apply_config_overrides(monkeypatch, CONSOLE_API_URL="https://dify.test") + monkeypatch.setattr(console_resources, "_actor", lambda: ("account-1", "tenant-1")) + monkeypatch.setattr( + console_resources, + "_console_services", + lambda: _BrokerRuntime(broker=Broker()), + ) + app = Flask(__name__) + + with app.test_request_context(json={"query": "What changed?", "mode": "auto"}): + post = inspect.unwrap(console_resources.KnowledgeFSSpaceQueryAdmissionApi.post) + response = post(console_resources.KnowledgeFSSpaceQueryAdmissionApi(), "control-1") + + assert response["operation_id"] == "createQuery" + assert response["trace_id"] == "trace-1" + assert response["request"] == { + "activeDocumentIds": [], + "activeEntityIds": [], + "knowledgeSpaceId": "space-1", + "mode": "auto", + "query": "What changed?", + } + assert response["url"] == "https://dify.test/console/api/knowledge-fs/query-stream" + assert calls == [ + { + "account_id": "account-1", + "control_space_id": "control-1", + "operation_id": "createQuery", + "tenant_id": "tenant-1", + } + ] + + @staticmethod + def test_console_query_stream_proxy_forwards_only_the_admitted_capability( + monkeypatch: pytest.MonkeyPatch, + ) -> None: + closed: list[str] = [] + + def close() -> None: + closed.append("closed") + + facade = FakeComponent( + stream_query=Recorder( + KnowledgeFSRemoteSSEResponse( + status_code=200, + headers=( + ("content-type", "text/event-stream"), + ("x-query-run-id", "query-1"), + ), + chunks=iter((b"event: answer\n", b'data: {"answer":"ok"}\n\n')), + close=close, + ) + ) + ) + monkeypatch.setattr(console_resources, "_console_services", lambda: FakeRuntime(facade=facade)) + app = Flask(__name__) + + with app.test_request_context( + json={ + "activeDocumentIds": [], + "activeEntityIds": [], + "knowledgeSpaceId": "space-1", + "mode": "fast", + "query": "hello", + }, + headers={ + "Authorization": "Bearer capability-token", + "Cookie": "session=must-not-forward", + "X-Trace-ID": "trace-1", + }, + ): + post = inspect.unwrap(console_resources.KnowledgeFSQueryStreamProxyApi.post) + response = post(console_resources.KnowledgeFSQueryStreamProxyApi()) + + assert response.status_code == 200 + assert response.headers["Content-Type"] == "text/event-stream" + assert response.headers["X-Accel-Buffering"] == "no" + assert b"".join(response.response) == b'event: answer\ndata: {"answer":"ok"}\n\n' + assert closed == ["closed"] + call = facade.stream_query.last_call.kwargs + assert call["capability_token"] == "capability-token" + assert call["trace_id"] == "trace-1" + assert narrowed(call["payload"], KnowledgeFSAdmittedQueryRequest).knowledge_space_id == "space-1" + assert set(call) == {"capability_token", "trace_id", "payload"} + + @staticmethod + def test_service_query_stream_proxy_forwards_only_the_admitted_capability( + monkeypatch: pytest.MonkeyPatch, + ) -> None: + closed: list[str] = [] + + def close() -> None: + closed.append("closed") + + facade = FakeComponent( + stream_query=Recorder( + KnowledgeFSRemoteSSEResponse( + status_code=200, + headers=( + ("content-type", "text/event-stream"), + ("x-query-run-id", "query-1"), + ), + chunks=iter((b"event: answer\n", b'data: {"answer":"ok"}\n\n')), + close=close, + ) + ) + ) + monkeypatch.setattr(service_resources, "_runtime", lambda: FakeRuntime(facade=facade)) + app = Flask(__name__) + + with app.test_request_context( + json={ + "activeDocumentIds": [], + "activeEntityIds": [], + "knowledgeSpaceId": "space-1", + "mode": "fast", + "query": "hello", + }, + headers={ + "Authorization": "Bearer capability-token", + "Cookie": "session=must-not-forward", + "X-Trace-ID": "trace-1", + }, + ): + post = inspect.unwrap(service_resources.KnowledgeFSServiceQueryStreamProxyApi.post) + response = post(service_resources.KnowledgeFSServiceQueryStreamProxyApi()) + + assert response.status_code == 200 + assert response.headers["Content-Type"] == "text/event-stream" + assert response.headers["X-Accel-Buffering"] == "no" + assert b"".join(response.response) == b'event: answer\ndata: {"answer":"ok"}\n\n' + assert closed == ["closed"] + call = facade.stream_query.last_call.kwargs + assert call["capability_token"] == "capability-token" + assert call["trace_id"] == "trace-1" + assert narrowed(call["payload"], KnowledgeFSAdmittedQueryRequest).knowledge_space_id == "space-1" + assert set(call) == {"capability_token", "trace_id", "payload"} + + @staticmethod + @pytest.mark.parametrize( + ("headers", "error_type"), + [ + ({"X-Trace-ID": "trace-1"}, service_resources.KnowledgeFSInvalidCredentialHTTPError), + ( + {"Authorization": "Bearer capability-token"}, + service_resources.KnowledgeFSServiceInvalidRequestHTTPError, + ), + ], + ) + def test_service_query_stream_proxy_rejects_invalid_internal_transport_headers( + headers: dict[str, str], + error_type: type[Exception], + ) -> None: + app = Flask(__name__) + + with app.test_request_context(headers=headers), pytest.raises(error_type): + service_resources._stream_capability() + + @staticmethod + def test_task_stream_capability_uses_broker( + monkeypatch: pytest.MonkeyPatch, + ) -> None: + calls: list[dict[str, object]] = [] + + class Broker: + def issue_interactive(self, **kwargs: object) -> KnowledgeFSIssuedProductCapability: + calls.append(kwargs) + return _issued_capability(token="direct-capability", operation_id="streamResearchTask") + + runtime = _BrokerRuntime(broker=Broker()) + apply_config_overrides(monkeypatch, CONSOLE_API_URL="https://dify.test") + monkeypatch.setattr(console_resources, "_actor", lambda: ("account-1", "tenant-1")) + monkeypatch.setattr(console_resources, "_console_services", lambda: runtime) + app = Flask(__name__) + + with app.test_request_context(json={"control_space_id": "control-1"}): + stream_post = inspect.unwrap(console_resources.KnowledgeFSTaskStreamCapabilityApi.post) + stream_response = stream_post(console_resources.KnowledgeFSTaskStreamCapabilityApi(), "task-1") + + assert stream_response["operation_id"] == "streamResearchTask" + assert calls == [ + { + "tenant_id": "tenant-1", + "account_id": "account-1", + "control_space_id": "control-1", + "operation_id": "streamResearchTask", + "resource_id": "task-1", + }, + ] + + @staticmethod + def test_service_query_admission_uses_broker(monkeypatch: pytest.MonkeyPatch) -> None: + calls: list[dict[str, object]] = [] + profile = _service_profile("control-1") + authorization = FakeComponent(authorize=Recorder(profile)) + + class Broker: + def issue_service(self, **kwargs: object) -> KnowledgeFSIssuedProductCapability: + calls.append(kwargs) + return _issued_capability(token="service-capability", operation_id="createQuery") + + @dataclass(frozen=True) + class ServiceRuntime: + service_api_authorization: FakeComponent + broker: Broker + + runtime = ServiceRuntime(service_api_authorization=authorization, broker=Broker()) + monkeypatch.setattr(service_resources, "validate_and_get_api_token", Recorder(_dataset_api_token())) + apply_config_overrides(monkeypatch, SERVICE_API_URL="https://api.dify.test") + monkeypatch.setattr(service_resources, "_runtime", lambda: runtime) + app = Flask(__name__) + + with app.test_request_context( + json={"query": "What changed?", "mode": "fast"}, + headers={"Authorization": "Bearer dataset-test-key"}, + ): + post = inspect.unwrap(service_resources.KnowledgeFSServiceQueryAdmissionApi.post) + response = post(service_resources.KnowledgeFSServiceQueryAdmissionApi(), "control-1") + + assert response["operation_id"] == "createQuery" + assert response["trace_id"] == "trace-1" + assert response["request"]["knowledgeSpaceId"] == "space-1" + assert response["url"] == "https://api.dify.test/v1/knowledge-fs/query-stream" + assert authorization.authorize.calls == [ + invocation( + api_token_id="token-1", + tenant_id="tenant-1", + control_space_id="control-1", + ) + ] + assert calls == [{"profile": profile, "operation_id": "createQuery"}] + + @staticmethod + @pytest.mark.parametrize( + ("raised", "expected"), + [ + # An unknown or wrong-kind key is an invalid credential (401). + ( + KnowledgeFSServiceApiAuthorizationError("bad key"), + service_resources.KnowledgeFSInvalidCredentialHTTPError, + ), + # A real key bound to other knowledge bases is out of scope for this space (403). + ( + KnowledgeFSServiceApiScopeError("out of scope"), + service_resources.KnowledgeFSServiceAccessDeniedHTTPError, + ), + ], + ) + def test_service_api_errors_distinguish_unknown_key_from_out_of_scope_key( + raised: Exception, + expected: type[Exception], + ) -> None: + @service_resources._service_api_errors + def view() -> None: + raise raised + + with pytest.raises(expected): + view() + + @staticmethod + def test_space_create_member_schema_requires_viewer_role() -> None: + payload_schema = console_ns.models["KnowledgeFSSpaceCreatePayload"].__schema__ + member_schema = console_ns.models["KnowledgeFSSpaceCreateMemberPayload"].__schema__ + + assert payload_schema["properties"]["members"]["maxItems"] == 1_000 + assert member_schema["properties"]["role"]["const"] == "viewer" + assert set(member_schema["required"]) == {"account_id", "role"} + assert member_schema["additionalProperties"] is False + + @staticmethod + def test_space_create_initial_access_idempotency_conflict_returns_409() -> None: + @console_resources._knowledge_fs_errors + def conflicting_creation() -> None: + raise KnowledgeFSControlSpaceIntentConflictError("different initial access") + + with pytest.raises(KnowledgeFSConflictHTTPError) as exc: + conflicting_creation() + + assert exc.value.code == 409 + + +_RAW_RESULT = object() + + +def source_response(**overrides: object) -> KnowledgeFSSourceResponse: + return KnowledgeFSSourceResponse.model_validate( + { + "id": "source-1", + "created_at": datetime(2026, 1, 1, tzinfo=UTC), + "updated_at": datetime(2026, 1, 1, tzinfo=UTC), + "knowledge_space_id": "space-1", + "name": "Source", + "permission_scope": [], + "status": "active", + "metadata": {}, + "type": "web", + "uri": "https://example.com", + "version": 1, + **overrides, + } + ) + + +def sync_policy_response(revision: int) -> KnowledgeFSSourceSyncPolicyResponse: + return KnowledgeFSSourceSyncPolicyResponse.model_validate( + { + "created_at": datetime(2026, 1, 1, tzinfo=UTC), + "enabled": True, + "expected_source_version": 1, + "id": "policy-1", + "knowledge_space_id": "space-1", + "mode": "interval", + "revision": revision, + "source_id": "source-1", + "updated_at": datetime(2026, 1, 1, tzinfo=UTC), + } + ) + + +def issued_capability() -> KnowledgeFSIssuedProductCapability: + return KnowledgeFSIssuedProductCapability( + token="capability-token", + expires_at=datetime(2026, 7, 21, tzinfo=UTC), + operation_id="createQuery", + knowledge_space_id="knowledge-space-1", + knowledge_space_revision=1, + trace_id="trace-1", + ) + + +def _knowledge_fs_resource_delegation_account(account_id: str) -> Account: + account = Account(name="Test User", email="user@example.com") + account.id = account_id + return account + + +@dataclass(frozen=True) +class _SettingsUpdateResult: + """The only field the settings route reads before dumping the facade result.""" + + migration: None = None + + +@dataclass(frozen=True) +class _MembersPayload: + members: tuple[str, ...] + query_images: list[str] = field(default_factory=list) + + +def _invoke[R](resource_type: type[Resource], handler: Callable[..., R], *args: object) -> R: + return unwrap(handler)(resource_type(), *args) + + +# Each case names the public resource, HTTP method, route arguments, and the one +# application boundary that owns the operation. Route-specific identifiers are +# asserted below so accidental delegation to a sibling resource is observable. +_CONSOLE_DELEGATION_CASES = ( + ( + console_resources.KnowledgeFSSpacesApi, + console_resources.KnowledgeFSSpacesApi.get, + (), + "application", + "list_spaces", + dict[str, object](), + ), + ( + console_resources.KnowledgeFSSpacesApi, + console_resources.KnowledgeFSSpacesApi.post, + ( + RequestContext( + request_id="request-1", trace_id=None, account_id="account-1", active_workspace_id="tenant-1" + ), + ), + "application", + "create_space", + dict[str, object](), + ), + ( + console_resources.KnowledgeFSSpaceApi, + console_resources.KnowledgeFSSpaceApi.get, + ("space-1",), + "application", + "get_space", + {"control_space_id": "space-1"}, + ), + ( + console_resources.KnowledgeFSSpaceApi, + console_resources.KnowledgeFSSpaceApi.patch, + ("space-1",), + "application", + "update_space", + {"control_space_id": "space-1"}, + ), + ( + console_resources.KnowledgeFSSpaceApi, + console_resources.KnowledgeFSSpaceApi.delete, + ("space-1",), + "application", + "delete_space", + {"control_space_id": "space-1"}, + ), + ( + console_resources.KnowledgeFSSpacePermissionsApi, + console_resources.KnowledgeFSSpacePermissionsApi.get, + ("space-1",), + "control_plane", + "list_permissions", + {"control_space_id": "space-1"}, + ), + ( + console_resources.KnowledgeFSSpaceMembersApi, + console_resources.KnowledgeFSSpaceMembersApi.put, + ("space-1",), + "control_plane", + "replace_members", + {"control_space_id": "space-1"}, + ), + ( + console_resources.KnowledgeFSSpaceExternalAccessApi, + console_resources.KnowledgeFSSpaceExternalAccessApi.get, + ("space-1",), + "control_plane", + "get_external_access", + {"control_space_id": "space-1"}, + ), + ( + console_resources.KnowledgeFSSpaceExternalAccessApi, + console_resources.KnowledgeFSSpaceExternalAccessApi.put, + ("space-1",), + "control_plane", + "update_external_access", + {"control_space_id": "space-1"}, + ), + ( + console_resources.KnowledgeFSSpaceAppBindingsApi, + console_resources.KnowledgeFSSpaceAppBindingsApi.get, + ("space-1",), + "app_bindings", + "list", + {"control_space_id": "space-1"}, + ), + ( + console_resources.KnowledgeFSSpaceAppBindingsApi, + console_resources.KnowledgeFSSpaceAppBindingsApi.put, + ("space-1",), + "app_bindings", + "upsert", + {"control_space_id": "space-1"}, + ), + ( + console_resources.KnowledgeFSSpaceAppBindingApi, + console_resources.KnowledgeFSSpaceAppBindingApi.delete, + ("space-1", KnowledgeFSAppSpaceJoinType.AGENT.value, "app-1"), + "app_bindings", + "revoke", + { + "control_space_id": "space-1", + "app_id": "app-1", + "caller_kind": KnowledgeFSAppSpaceJoinType.AGENT, + }, + ), + ( + console_resources.KnowledgeFSSpaceSettingsApi, + console_resources.KnowledgeFSSpaceSettingsApi.get, + ("space-1",), + "facade", + "get_settings", + {"control_space_id": "space-1"}, + ), + ( + console_resources.KnowledgeFSSpaceSettingsApi, + console_resources.KnowledgeFSSpaceSettingsApi.patch, + ("space-1",), + "facade", + "update_settings", + {"control_space_id": "space-1"}, + ), + ( + console_resources.KnowledgeFSSpaceOverviewQueryOutcomesApi, + console_resources.KnowledgeFSSpaceOverviewQueryOutcomesApi.get, + ("space-1",), + "facade", + "get_overview_query_outcomes", + {"control_space_id": "space-1", "window": "24h"}, + ), + ( + console_resources.KnowledgeFSSpaceOverviewInventoryApi, + console_resources.KnowledgeFSSpaceOverviewInventoryApi.get, + ("space-1",), + "facade", + "get_overview_inventory", + {"control_space_id": "space-1"}, + ), + ( + console_resources.KnowledgeFSSpaceOverviewAttentionApi, + console_resources.KnowledgeFSSpaceOverviewAttentionApi.get, + ("space-1",), + "facade", + "list_overview_attention", + {"control_space_id": "space-1", "include_dismissed": False, "limit": 50}, + ), + ( + console_resources.KnowledgeFSSpaceOverviewActivityApi, + console_resources.KnowledgeFSSpaceOverviewActivityApi.get, + ("space-1",), + "facade", + "list_overview_activity", + { + "control_space_id": "space-1", + "action": None, + "actor_id": None, + "actor_type": None, + "cursor": None, + "from_at": None, + "limit": 50, + "resource_type": None, + "result": None, + "to_at": None, + }, + ), + ( + console_resources.KnowledgeFSSpaceOverviewHealthApi, + console_resources.KnowledgeFSSpaceOverviewHealthApi.get, + ("space-1",), + "facade", + "get_overview_health", + {"control_space_id": "space-1"}, + ), + ( + console_resources.KnowledgeFSSpaceDocumentsApi, + console_resources.KnowledgeFSSpaceDocumentsApi.get, + ("space-1",), + "facade", + "list_documents", + {"control_space_id": "space-1"}, + ), + ( + console_resources.KnowledgeFSSpaceBulkDocumentsApi, + console_resources.KnowledgeFSSpaceBulkDocumentsApi.delete, + ("space-1",), + "facade", + "bulk_delete_documents", + {"control_space_id": "space-1"}, + ), + ( + console_resources.KnowledgeFSSpaceDocumentReindexApi, + console_resources.KnowledgeFSSpaceDocumentReindexApi.post, + ("space-1",), + "facade", + "reindex_documents", + {"control_space_id": "space-1"}, + ), + ( + console_resources.KnowledgeFSSpaceDocumentApi, + console_resources.KnowledgeFSSpaceDocumentApi.get, + ("space-1", "document-1"), + "facade", + "get_document", + {"control_space_id": "space-1", "document_id": "document-1"}, + ), + ( + console_resources.KnowledgeFSSpaceDocumentApi, + console_resources.KnowledgeFSSpaceDocumentApi.patch, + ("space-1", "document-1"), + "facade", + "update_document_metadata", + {"control_space_id": "space-1", "document_id": "document-1"}, + ), + ( + console_resources.KnowledgeFSSpaceMetadataApi, + console_resources.KnowledgeFSSpaceMetadataApi.get, + ("space-1",), + "facade", + "list_metadata_fields", + {"control_space_id": "space-1", "cursor": None, "limit": 100}, + ), + ( + console_resources.KnowledgeFSSpaceMetadataApi, + console_resources.KnowledgeFSSpaceMetadataApi.post, + ("space-1",), + "facade", + "create_metadata_field", + {"control_space_id": "space-1"}, + ), + ( + console_resources.KnowledgeFSSpaceMetadataFieldApi, + console_resources.KnowledgeFSSpaceMetadataFieldApi.patch, + ("space-1", "field-1"), + "facade", + "update_metadata_field", + {"control_space_id": "space-1", "field_id": "field-1"}, + ), + ( + console_resources.KnowledgeFSSpaceDocumentApi, + console_resources.KnowledgeFSSpaceDocumentApi.delete, + ("space-1", "document-1"), + "facade", + "delete_document", + {"control_space_id": "space-1", "document_id": "document-1"}, + ), + ( + console_resources.KnowledgeFSSpaceLogicalDocumentApi, + console_resources.KnowledgeFSSpaceLogicalDocumentApi.delete, + ("space-1", "document-1"), + "facade", + "delete_logical_document", + {"control_space_id": "space-1", "document_id": "document-1"}, + ), + ( + console_resources.KnowledgeFSSpaceDocumentOutlineApi, + console_resources.KnowledgeFSSpaceDocumentOutlineApi.get, + ("space-1", "document-1"), + "facade", + "get_document_outline", + {"control_space_id": "space-1", "document_id": "document-1"}, + ), + ( + console_resources.KnowledgeFSSpaceDocumentRevisionsApi, + console_resources.KnowledgeFSSpaceDocumentRevisionsApi.get, + ("space-1", "document-1"), + "facade", + "list_document_revisions", + {"control_space_id": "space-1", "document_id": "document-1"}, + ), + ( + console_resources.KnowledgeFSSpaceDocumentChunksApi, + console_resources.KnowledgeFSSpaceDocumentChunksApi.get, + ("space-1", "document-1", 3), + "facade", + "list_document_chunks", + {"control_space_id": "space-1", "document_id": "document-1", "revision": 3}, + ), + ( + console_resources.KnowledgeFSSpaceDocumentChunkApi, + console_resources.KnowledgeFSSpaceDocumentChunkApi.get, + ("space-1", "document-1", 3, "chunk-1"), + "facade", + "get_document_chunk", + { + "control_space_id": "space-1", + "document_id": "document-1", + "revision": 3, + "chunk_id": "chunk-1", + }, + ), + ( + console_resources.KnowledgeFSSpaceCompilationJobApi, + console_resources.KnowledgeFSSpaceCompilationJobApi.get, + ("space-1", "job-1"), + "facade", + "get_compilation_job", + {"control_space_id": "space-1", "job_id": "job-1"}, + ), + ( + console_resources.KnowledgeFSSpaceCompilationJobApi, + console_resources.KnowledgeFSSpaceCompilationJobApi.delete, + ("space-1", "job-1"), + "facade", + "cancel_compilation_job", + {"control_space_id": "space-1", "job_id": "job-1"}, + ), + ( + console_resources.KnowledgeFSSpaceCompilationJobRetryApi, + console_resources.KnowledgeFSSpaceCompilationJobRetryApi.post, + ("space-1", "job-1"), + "facade", + "retry_compilation_job", + {"control_space_id": "space-1", "job_id": "job-1"}, + ), + ( + console_resources.KnowledgeFSSpaceBulkJobApi, + console_resources.KnowledgeFSSpaceBulkJobApi.get, + ("space-1", "job-1"), + "facade", + "get_bulk_job", + {"control_space_id": "space-1", "job_id": "job-1"}, + ), + ( + console_resources.KnowledgeFSSpaceBackgroundTasksApi, + console_resources.KnowledgeFSSpaceBackgroundTasksApi.get, + ("space-1",), + "facade", + "list_background_tasks", + {"control_space_id": "space-1", "cursor": None, "limit": 50}, + ), + ( + console_resources.KnowledgeFSSpaceBackgroundTaskCancelApi, + console_resources.KnowledgeFSSpaceBackgroundTaskCancelApi.post, + ("space-1", "source", "task-1"), + "facade", + "cancel_background_task", + {"control_space_id": "space-1", "task_kind": "source", "task_id": "task-1"}, + ), + ( + console_resources.KnowledgeFSSpaceBackgroundTaskRetryApi, + console_resources.KnowledgeFSSpaceBackgroundTaskRetryApi.post, + ("space-1", "document_bulk", "task-1"), + "facade", + "retry_background_task", + { + "control_space_id": "space-1", + "task_kind": "document_bulk", + "task_id": "task-1", + }, + ), + ( + console_resources.KnowledgeFSSpaceSourcesApi, + console_resources.KnowledgeFSSpaceSourcesApi.get, + ("space-1",), + "facade", + "list_sources", + {"control_space_id": "space-1", "limit": 50}, + ), + ( + console_resources.KnowledgeFSSpaceSourcesApi, + console_resources.KnowledgeFSSpaceSourcesApi.post, + ("space-1",), + "facade", + "create_source", + {"control_space_id": "space-1"}, + ), + ( + console_resources.KnowledgeFSSpaceSourceApi, + console_resources.KnowledgeFSSpaceSourceApi.get, + ("space-1", "source-1"), + "facade", + "get_source", + {"control_space_id": "space-1", "source_id": "source-1"}, + ), + ( + console_resources.KnowledgeFSSpaceSourceApi, + console_resources.KnowledgeFSSpaceSourceApi.patch, + ("space-1", "source-1"), + "facade", + "update_source", + {"control_space_id": "space-1", "source_id": "source-1"}, + ), + ( + console_resources.KnowledgeFSSpaceSourceApi, + console_resources.KnowledgeFSSpaceSourceApi.delete, + ("space-1", "source-1"), + "facade", + "delete_source", + {"control_space_id": "space-1", "source_id": "source-1"}, + ), + ( + console_resources.KnowledgeFSSpaceSourceTestApi, + console_resources.KnowledgeFSSpaceSourceTestApi.post, + ("space-1", "source-1"), + "facade", + "test_source", + {"control_space_id": "space-1", "source_id": "source-1"}, + ), + ( + console_resources.KnowledgeFSSpaceSourceSyncApi, + console_resources.KnowledgeFSSpaceSourceSyncApi.post, + ("space-1", "source-1"), + "facade", + "sync_source", + {"control_space_id": "space-1", "source_id": "source-1"}, + ), + ( + console_resources.KnowledgeFSSpaceSourceWorkflowImportApi, + console_resources.KnowledgeFSSpaceSourceWorkflowImportApi.post, + ("space-1", "source-1"), + "facade", + "import_source_workflow", + {"control_space_id": "space-1", "source_id": "source-1"}, + ), + ( + console_resources.KnowledgeFSSpaceSourcePagesApi, + console_resources.KnowledgeFSSpaceSourcePagesApi.get, + ("space-1", "source-1"), + "facade", + "list_source_pages", + {"control_space_id": "space-1", "source_id": "source-1"}, + ), + ( + console_resources.KnowledgeFSSpaceSourcePageImportApi, + console_resources.KnowledgeFSSpaceSourcePageImportApi.post, + ("space-1", "source-1"), + "facade", + "import_source_pages", + {"control_space_id": "space-1", "source_id": "source-1"}, + ), + ( + console_resources.KnowledgeFSSpaceSourceFilesApi, + console_resources.KnowledgeFSSpaceSourceFilesApi.get, + ("space-1", "source-1"), + "facade", + "list_source_files", + {"control_space_id": "space-1", "source_id": "source-1"}, + ), + ( + console_resources.KnowledgeFSSpaceSourceFileImportApi, + console_resources.KnowledgeFSSpaceSourceFileImportApi.post, + ("space-1", "source-1"), + "facade", + "import_source_files", + {"control_space_id": "space-1", "source_id": "source-1"}, + ), + ( + console_resources.KnowledgeFSSpaceResearchTasksApi, + console_resources.KnowledgeFSSpaceResearchTasksApi.get, + ("space-1",), + "facade", + "list_research_tasks", + {"control_space_id": "space-1"}, + ), + ( + console_resources.KnowledgeFSSpaceResearchTasksApi, + console_resources.KnowledgeFSSpaceResearchTasksApi.post, + ("space-1",), + "facade", + "create_research_task", + {"control_space_id": "space-1"}, + ), + ( + console_resources.KnowledgeFSSpaceResearchTaskPlanApi, + console_resources.KnowledgeFSSpaceResearchTaskPlanApi.post, + ("space-1",), + "facade", + "plan_research_task", + {"control_space_id": "space-1"}, + ), + ( + console_resources.KnowledgeFSSpaceResearchTaskApi, + console_resources.KnowledgeFSSpaceResearchTaskApi.get, + ("space-1", "task-1"), + "facade", + "get_research_task", + {"control_space_id": "space-1", "task_id": "task-1"}, + ), + ( + console_resources.KnowledgeFSSpaceResearchTaskApi, + console_resources.KnowledgeFSSpaceResearchTaskApi.delete, + ("space-1", "task-1"), + "facade", + "cancel_research_task", + {"control_space_id": "space-1", "task_id": "task-1"}, + ), + ( + console_resources.KnowledgeFSSpaceResearchTaskPartialsApi, + console_resources.KnowledgeFSSpaceResearchTaskPartialsApi.get, + ("space-1", "task-1"), + "facade", + "list_research_task_partials", + {"control_space_id": "space-1", "task_id": "task-1"}, + ), + ( + console_resources.KnowledgeFSSpaceTracesApi, + console_resources.KnowledgeFSSpaceTracesApi.get, + ("space-1",), + "facade", + "list_traces", + {"control_space_id": "space-1"}, + ), + ( + console_resources.KnowledgeFSSpaceTraceApi, + console_resources.KnowledgeFSSpaceTraceApi.get, + ("space-1", "trace-1"), + "facade", + "get_trace", + {"control_space_id": "space-1", "trace_id": "trace-1"}, + ), + ( + console_resources.KnowledgeFSSpaceTraceEvidenceApi, + console_resources.KnowledgeFSSpaceTraceEvidenceApi.get, + ("space-1", "trace-1"), + "facade", + "list_trace_entries", + {"control_space_id": "space-1", "trace_id": "trace-1", "kind": "evidence"}, + ), + ( + console_resources.KnowledgeFSSpaceTraceConflictsApi, + console_resources.KnowledgeFSSpaceTraceConflictsApi.get, + ("space-1", "trace-1"), + "facade", + "list_trace_entries", + {"control_space_id": "space-1", "trace_id": "trace-1", "kind": "conflicts"}, + ), + ( + console_resources.KnowledgeFSSpaceTraceMissingApi, + console_resources.KnowledgeFSSpaceTraceMissingApi.get, + ("space-1", "trace-1"), + "facade", + "list_trace_entries", + {"control_space_id": "space-1", "trace_id": "trace-1", "kind": "missing"}, + ), + ( + console_resources.KnowledgeFSSpaceUploadSessionsApi, + console_resources.KnowledgeFSSpaceUploadSessionsApi.post, + ("space-1",), + "facade", + "create_upload_session", + {"control_space_id": "space-1"}, + ), + ( + console_resources.KnowledgeFSSpaceUploadSessionPartPresignApi, + console_resources.KnowledgeFSSpaceUploadSessionPartPresignApi.post, + ("space-1", "session-1", 3), + "facade", + "presign_upload_session_part", + {"control_space_id": "space-1", "upload_session_id": "session-1", "part_number": 3}, + ), + ( + console_resources.KnowledgeFSSpaceUploadSessionCompleteApi, + console_resources.KnowledgeFSSpaceUploadSessionCompleteApi.post, + ("space-1", "session-1"), + "facade", + "complete_upload_session", + {"control_space_id": "space-1", "upload_session_id": "session-1"}, + ), + ( + console_resources.KnowledgeFSSpaceUploadSessionAbortApi, + console_resources.KnowledgeFSSpaceUploadSessionAbortApi.post, + ("space-1", "session-1"), + "facade", + "abort_upload_session", + {"control_space_id": "space-1", "upload_session_id": "session-1"}, + ), + ( + console_resources.KnowledgeFSSpaceSmallFileUploadApi, + console_resources.KnowledgeFSSpaceSmallFileUploadApi.post, + ("space-1", "session-1"), + "facade", + "upload_small_file", + {"control_space_id": "space-1", "upload_session_id": "session-1"}, + ), +) + + +_SERVICE_DELEGATION_CASES = ( + ( + service_resources.KnowledgeFSServiceBulkDocumentsApi, + service_resources.KnowledgeFSServiceBulkDocumentsApi.delete, + ("space-1",), + "bulkDeleteDocuments", + dict[str, object](), + ), + ( + service_resources.KnowledgeFSServiceDocumentReindexApi, + service_resources.KnowledgeFSServiceDocumentReindexApi.post, + ("space-1",), + "reindexDocuments", + dict[str, object](), + ), + ( + service_resources.KnowledgeFSServiceDocumentApi, + service_resources.KnowledgeFSServiceDocumentApi.get, + ("space-1", "document-1"), + "getDocument", + {"resource_id": "document-1", "path_parameters": (("documentId", "document-1"),)}, + ), + ( + service_resources.KnowledgeFSServiceDocumentApi, + service_resources.KnowledgeFSServiceDocumentApi.patch, + ("space-1", "document-1"), + "updateDocumentMetadata", + {"resource_id": "document-1", "path_parameters": (("documentId", "document-1"),)}, + ), + ( + service_resources.KnowledgeFSServiceDocumentApi, + service_resources.KnowledgeFSServiceDocumentApi.delete, + ("space-1", "document-1"), + "deleteDocument", + {"resource_id": "document-1", "path_parameters": (("documentId", "document-1"),)}, + ), + ( + service_resources.KnowledgeFSServiceDocumentOutlineApi, + service_resources.KnowledgeFSServiceDocumentOutlineApi.get, + ("space-1", "document-1"), + "getDocumentOutline", + {"resource_id": "document-1", "path_parameters": (("documentId", "document-1"),)}, + ), + ( + service_resources.KnowledgeFSServiceDocumentRevisionsApi, + service_resources.KnowledgeFSServiceDocumentRevisionsApi.get, + ("space-1", "document-1"), + "listDocumentRevisions", + {"resource_id": "document-1", "path_parameters": (("documentId", "document-1"),)}, + ), + ( + service_resources.KnowledgeFSServiceDocumentChunksApi, + service_resources.KnowledgeFSServiceDocumentChunksApi.get, + ("space-1", "document-1", 3), + "listDocumentChunks", + { + "resource_id": "document-1", + "path_parameters": (("documentId", "document-1"), ("revision", "3")), + }, + ), + ( + service_resources.KnowledgeFSServiceDocumentChunkApi, + service_resources.KnowledgeFSServiceDocumentChunkApi.get, + ("space-1", "document-1", 3, "chunk-1"), + "getDocumentChunk", + { + "resource_id": "document-1", + "path_parameters": (("documentId", "document-1"), ("revision", "3"), ("chunkId", "chunk-1")), + }, + ), + ( + service_resources.KnowledgeFSServiceCompilationJobApi, + service_resources.KnowledgeFSServiceCompilationJobApi.get, + ("space-1", "job-1"), + "getCompilationJob", + {"resource_id": "job-1"}, + ), + ( + service_resources.KnowledgeFSServiceCompilationJobApi, + service_resources.KnowledgeFSServiceCompilationJobApi.delete, + ("space-1", "job-1"), + "cancelCompilationJob", + {"resource_id": "job-1"}, + ), + ( + service_resources.KnowledgeFSServiceCompilationJobRetryApi, + service_resources.KnowledgeFSServiceCompilationJobRetryApi.post, + ("space-1", "job-1"), + "retryCompilationJob", + {"resource_id": "job-1"}, + ), + ( + service_resources.KnowledgeFSServiceBulkJobApi, + service_resources.KnowledgeFSServiceBulkJobApi.get, + ("space-1", "job-1"), + "getBulkJob", + {"resource_id": "job-1"}, + ), + ( + service_resources.KnowledgeFSServiceSourceApi, + service_resources.KnowledgeFSServiceSourceApi.get, + ("space-1", "source-1"), + "getSource", + {"resource_id": "source-1", "path_parameters": (("sourceId", "source-1"),)}, + ), + ( + service_resources.KnowledgeFSServiceSourceApi, + service_resources.KnowledgeFSServiceSourceApi.patch, + ("space-1", "source-1"), + "updateSource", + {"resource_id": "source-1", "path_parameters": (("sourceId", "source-1"),)}, + ), + ( + service_resources.KnowledgeFSServiceSourceApi, + service_resources.KnowledgeFSServiceSourceApi.delete, + ("space-1", "source-1"), + "deleteSource", + {"resource_id": "source-1", "path_parameters": (("sourceId", "source-1"),)}, + ), + ( + service_resources.KnowledgeFSServiceSourceTestApi, + service_resources.KnowledgeFSServiceSourceTestApi.post, + ("space-1", "source-1"), + "testSource", + {"resource_id": "source-1", "path_parameters": (("sourceId", "source-1"),)}, + ), + ( + service_resources.KnowledgeFSServiceSourceCrawlApi, + service_resources.KnowledgeFSServiceSourceCrawlApi.post, + ("space-1", "source-1"), + "previewSourceCrawl", + {"resource_id": "source-1", "path_parameters": (("sourceId", "source-1"),)}, + ), + ( + service_resources.KnowledgeFSServiceSourcePagesApi, + service_resources.KnowledgeFSServiceSourcePagesApi.get, + ("space-1", "source-1"), + "listSourcePages", + {"resource_id": "source-1", "path_parameters": (("sourceId", "source-1"),)}, + ), + ( + service_resources.KnowledgeFSServiceSourcePageImportApi, + service_resources.KnowledgeFSServiceSourcePageImportApi.post, + ("space-1", "source-1"), + "importSourceWorkflow", + {"resource_id": "source-1", "path_parameters": (("sourceId", "source-1"),)}, + ), + ( + service_resources.KnowledgeFSServiceSourceFilesApi, + service_resources.KnowledgeFSServiceSourceFilesApi.get, + ("space-1", "source-1"), + "listSourceFiles", + {"resource_id": "source-1", "path_parameters": (("sourceId", "source-1"),)}, + ), + ( + service_resources.KnowledgeFSServiceSourceFileImportApi, + service_resources.KnowledgeFSServiceSourceFileImportApi.post, + ("space-1", "source-1"), + "importSourceWorkflow", + {"resource_id": "source-1", "path_parameters": (("sourceId", "source-1"),)}, + ), + ( + service_resources.KnowledgeFSServiceResearchTasksApi, + service_resources.KnowledgeFSServiceResearchTasksApi.post, + ("space-1",), + "createResearchTask", + {"bind_space_in_body": True}, + ), + ( + service_resources.KnowledgeFSServiceResearchTaskPlanApi, + service_resources.KnowledgeFSServiceResearchTaskPlanApi.post, + ("space-1",), + "planResearchTask", + {"bind_space_in_body": True}, + ), + ( + service_resources.KnowledgeFSServiceResearchTaskApi, + service_resources.KnowledgeFSServiceResearchTaskApi.get, + ("space-1", "task-1"), + "getResearchTask", + {"resource_id": "task-1"}, + ), + ( + service_resources.KnowledgeFSServiceResearchTaskApi, + service_resources.KnowledgeFSServiceResearchTaskApi.delete, + ("space-1", "task-1"), + "cancelResearchTask", + {"resource_id": "task-1"}, + ), + ( + service_resources.KnowledgeFSServiceResearchTaskPartialsApi, + service_resources.KnowledgeFSServiceResearchTaskPartialsApi.get, + ("space-1", "task-1"), + "listResearchTaskPartials", + {"resource_id": "task-1"}, + ), + ( + service_resources.KnowledgeFSServiceTraceApi, + service_resources.KnowledgeFSServiceTraceApi.get, + ("space-1", "trace-1"), + "getTrace", + {"resource_id": "trace-1", "path_parameters": (("traceId", "trace-1"),)}, + ), + ( + service_resources.KnowledgeFSServiceTraceEvidenceApi, + service_resources.KnowledgeFSServiceTraceEvidenceApi.get, + ("space-1", "trace-1"), + "listTraceEvidence", + {"resource_id": "trace-1", "path_parameters": (("traceId", "trace-1"),)}, + ), + ( + service_resources.KnowledgeFSServiceTraceConflictsApi, + service_resources.KnowledgeFSServiceTraceConflictsApi.get, + ("space-1", "trace-1"), + "listTraceConflicts", + {"resource_id": "trace-1", "path_parameters": (("traceId", "trace-1"),)}, + ), + ( + service_resources.KnowledgeFSServiceTraceMissingApi, + service_resources.KnowledgeFSServiceTraceMissingApi.get, + ("space-1", "trace-1"), + "listTraceMissing", + {"resource_id": "trace-1", "path_parameters": (("traceId", "trace-1"),)}, + ), +) + + +class TestKnowledgeFsResourceDelegationCases: + @staticmethod + @pytest.fixture(autouse=True) + def _enabled_knowledge_fs(config_overrides: Callable[..., None]) -> None: + config_overrides(DEPLOYMENT_EDITION=DeploymentEdition.CLOUD, KNOWLEDGE_FS_ENABLED=True) + + @staticmethod + @pytest.mark.parametrize( + ("class_name", "method_name", "route_args", "component_name", "delegate_name", "expected_fields"), + _CONSOLE_DELEGATION_CASES, + ) + def test_console_resources_delegate_one_tenant_scoped_product_operation( + monkeypatch: pytest.MonkeyPatch, + class_name: type[Resource], + method_name: Callable[..., object], + route_args: tuple[object, ...], + component_name: str, + delegate_name: str, + expected_fields: dict[str, object], + ) -> None: + payload = ( + KnowledgeFSSourceUpdatePayload(name="Renamed") + if class_name.__name__ == "KnowledgeFSSpaceSourceApi" and method_name.__name__ == "patch" + else _MembersPayload(members=("member-1",)) + ) + delegate = Recorder(_RAW_RESULT) + components = {name: FakeComponent() for name in ("application", "control_plane", "app_bindings", "facade")} + components[component_name] = FakeComponent(**{delegate_name: delegate}) + runtime = FakeRuntime( + application=components["application"], + control_plane=components["control_plane"], + app_bindings=components["app_bindings"], + facade=components["facade"], + ) + dump_response = Recorder(effect=lambda schema, raw: (schema.__name__, raw)) + monkeypatch.setattr(console_resources, "_actor", lambda: ("account-1", "tenant-1")) + monkeypatch.setattr(console_resources, "_console_services", lambda: runtime) + monkeypatch.setattr(console_resources, "_payload", lambda _: payload) + monkeypatch.setattr(console_resources, "_idempotency_key", lambda: "idempotency-1") + monkeypatch.setattr(console_resources, "_query_pairs", lambda _: (("normalized", "true"),)) + monkeypatch.setattr(console_resources, "dump_response", dump_response) + app = Flask(__name__) + + with app.test_request_context("/", method="POST"): + result = _invoke(class_name, method_name, *route_args) + + assert len(delegate.calls) == 1 + call_fields = delegate.last_call.kwargs + assert call_fields["tenant_id"] == "tenant-1" + assert call_fields.get("account_id", call_fields.get("actor_account_id")) == "account-1" + for field_name, expected in expected_fields.items(): + assert call_fields[field_name] == expected + if "payload" in call_fields: + if class_name.__name__ == "KnowledgeFSSpaceSourceApi" and method_name.__name__ == "patch": + forwarded = narrowed(call_fields["payload"], BaseModel) + source_update = narrowed(payload, KnowledgeFSSourceUpdatePayload) + assert forwarded.model_dump(exclude_none=True) == source_update.model_dump(exclude_none=True) + else: + assert call_fields["payload"] is payload + if "members" in call_fields: + assert call_fields["members"] == ("member-1",) + if "idempotency_key" in call_fields: + assert call_fields["idempotency_key"] == "idempotency-1" + if dump_response.calls: + assert dump_response.last_call.args[1] is _RAW_RESULT + else: + assert result == ("", HTTPStatus.NO_CONTENT) + + @staticmethod + def test_console_document_reference_delegates_exact_asset_identity(monkeypatch: pytest.MonkeyPatch) -> None: + facade = FakeComponent(resolve_document_reference=Recorder(_RAW_RESULT)) + runtime = FakeRuntime(facade=facade) + monkeypatch.setattr(console_resources, "_actor", lambda: ("account-1", "tenant-1")) + monkeypatch.setattr(console_resources, "_console_services", lambda: runtime) + monkeypatch.setattr(console_resources, "dump_response", lambda _schema, raw: raw) + app = Flask(__name__) + + with app.test_request_context("/?document_asset_id=asset-1&document_asset_version=7"): + result = _invoke( + console_resources.KnowledgeFSSpaceDocumentReferenceApi, + console_resources.KnowledgeFSSpaceDocumentReferenceApi.get, + "space-1", + ) + + assert result is _RAW_RESULT + assert facade.resolve_document_reference.calls == [ + invocation( + tenant_id="tenant-1", + account_id="account-1", + control_space_id="space-1", + document_asset_id="asset-1", + document_asset_version=7, + ) + ] + + @staticmethod + def test_console_space_list_preserves_repeated_creator_filters(monkeypatch: pytest.MonkeyPatch) -> None: + application = FakeComponent(list_spaces=Recorder(_RAW_RESULT)) + runtime = FakeRuntime(application=application) + monkeypatch.setattr(console_resources, "_actor", lambda: ("account-1", "tenant-1")) + monkeypatch.setattr(console_resources, "_console_services", lambda: runtime) + monkeypatch.setattr(console_resources, "dump_response", lambda _schema, raw: raw) + app = Flask(__name__) + + with app.test_request_context( + "/?page=2&limit=10&creator_ids=creator-1&creator_ids=creator-2&tag_ids=tag-1&tag_ids=tag-2&query=%20Support%20" + ): + result = _invoke(console_resources.KnowledgeFSSpacesApi, console_resources.KnowledgeFSSpacesApi.get) + + assert application.list_spaces.calls == [ + invocation( + tenant_id="tenant-1", + account_id="account-1", + page=2, + limit=10, + creator_ids=["creator-1", "creator-2"], + tag_ids=["tag-1", "tag-2"], + query="Support", + ) + ] + assert result is _RAW_RESULT + + @staticmethod + def test_console_source_update_does_not_start_sync_before_resource_selection( + monkeypatch: pytest.MonkeyPatch, + ) -> None: + facade = FakeComponent(update_source=Recorder(_RAW_RESULT)) + runtime = FakeRuntime(facade=facade) + payload = KnowledgeFSSourceUpdatePayload.model_validate( + { + "expectedVersion": 3, + "status": "disabled", + "uri": "https://new.example.com", + } + ) + monkeypatch.setattr(console_resources, "_actor", lambda: ("account-1", "tenant-1")) + monkeypatch.setattr(console_resources, "_console_services", lambda: runtime) + monkeypatch.setattr(console_resources, "_payload", lambda _: payload) + monkeypatch.setattr(console_resources, "dump_response", lambda _schema, raw: raw) + app = Flask(__name__) + + with app.test_request_context("/", method="PATCH"): + result = _invoke( + console_resources.KnowledgeFSSpaceSourceApi, + console_resources.KnowledgeFSSpaceSourceApi.patch, + "space-1", + "source-1", + ) + + forwarded = narrowed(facade.update_source.last_call.kwargs["payload"], KnowledgeFSAtomicSourceUpdatePayload) + assert forwarded is not payload + assert forwarded.status == "disabled" + assert forwarded.sync_after_update is None + assert result is _RAW_RESULT + + @staticmethod + @pytest.mark.parametrize( + ("selection", "expected_import_kind"), + [ + ( + {"kind": "website_crawl", "sourceUrls": ["https://new.example.com/a"]}, + "website-crawl-import", + ), + ( + { + "kind": "online_document", + "items": [ + { + "pageId": "page-1", + "providerItemId": "provider-page-1", + "type": "page", + "workspaceId": "workspace-1", + } + ], + }, + "online-document-import", + ), + ( + { + "kind": "online_drive", + "items": [{"id": "file-1", "name": "Plan.pdf", "providerItemId": "provider-file-1"}], + }, + "online-drive-import", + ), + ], + ) + def test_console_source_update_commits_complete_edit_selection( + monkeypatch: pytest.MonkeyPatch, + selection: dict[str, object], + expected_import_kind: str, + ) -> None: + selection_kind = selection["kind"] + assert isinstance(selection_kind, str) + provider_kind = {"online_document": "online-document", "online_drive": "online-drive"}.get(selection_kind) + original = source_response( + metadata={"providerKind": provider_kind} if provider_kind is not None else {}, + type="web" if selection_kind == "website_crawl" else "connector", + uri="https://old.example.com", + ) + updated = source_response(sync_workflow=None) + facade = FakeComponent( + get_source=Recorder(results=[original, updated]), + update_source=Recorder(source_response(version=4)), + ) + workflow = KnowledgeFSSourceWorkflowResponse.model_validate( + { + "checkpoint": "queued", + "created_at": datetime(2026, 1, 1, tzinfo=UTC), + "execution_attempts": 0, + "id": "workflow-1", + "knowledge_space_id": "space-1", + "kind": "source-import", + "max_execution_attempts": 3, + "progress_completed": 0, + "progress_failed": 0, + "progress_skipped": 0, + "state": "queued", + "updated_at": datetime(2026, 1, 1, tzinfo=UTC), + } + ) + runtime = FakeRuntime(facade=facade) + payload = KnowledgeFSSourceUpdatePayload.model_validate( + { + "expectedVersion": 3, + "name": "Edited source", + "selection": selection, + "syncPolicy": {"enabled": True, "mode": "interval"}, + } + ) + commit = Recorder(workflow) + monkeypatch.setattr(console_resources, "_actor", lambda: ("account-1", "tenant-1")) + monkeypatch.setattr(console_resources, "_console_services", lambda: runtime) + monkeypatch.setattr(console_resources, "_payload", lambda _: payload) + monkeypatch.setattr(source_edit_service, "commit_source_import", commit) + monkeypatch.setattr(console_resources, "dump_response", lambda _schema, raw: raw) + app = Flask(__name__) + + with app.test_request_context("/", method="PATCH"): + result = _invoke( + console_resources.KnowledgeFSSpaceSourceApi, + console_resources.KnowledgeFSSpaceSourceApi.patch, + "space-1", + "source-1", + ) + + assert facade.update_source.calls == [] + source_payload = narrowed(commit.last_call.kwargs["source_update"], KnowledgeFSSourceImportConfigurationPayload) + assert source_payload.status == "disabled" + import_payload = narrowed(commit.last_call.kwargs["payload"], BaseModel) + assert import_payload.model_dump()["kind"] == expected_import_kind + assert commit.last_call.kwargs["idempotency_key"] == "source-edit:source-1:3" + assert updated.sync_workflow is workflow + assert result is updated + + @staticmethod + def test_console_source_update_does_not_import_an_unchanged_complete_selection( + monkeypatch: pytest.MonkeyPatch, + ) -> None: + source = source_response( + metadata={"crawled": {"https://example.com/a": {}}}, + sync_policy=None, + type="web", + uri="https://example.com", + version=3, + ) + facade = FakeComponent( + get_source=Recorder(source), + get_source_sync_policy=Recorder(sync_policy_response(revision=2)), + update_source=Recorder(source), + update_source_sync_policy=Recorder(), + ) + runtime = FakeRuntime(facade=facade) + payload = KnowledgeFSSourceUpdatePayload.model_validate( + { + "selection": {"kind": "website_crawl", "sourceUrls": ["https://example.com/a"]}, + "syncPolicy": {"enabled": True, "mode": "interval"}, + } + ) + commit = Recorder() + monkeypatch.setattr(console_resources, "_actor", lambda: ("account-1", "tenant-1")) + monkeypatch.setattr(console_resources, "_console_services", lambda: runtime) + monkeypatch.setattr(console_resources, "_payload", lambda _: payload) + monkeypatch.setattr(source_edit_service, "commit_source_import", commit) + monkeypatch.setattr(console_resources, "dump_response", lambda _schema, raw: raw) + app = Flask(__name__) + + with app.test_request_context("/", method="PATCH"): + result = _invoke( + console_resources.KnowledgeFSSpaceSourceApi, + console_resources.KnowledgeFSSpaceSourceApi.patch, + "space-1", + "source-1", + ) + + assert len(facade.update_source.calls) == 1 + assert commit.calls == [] + assert facade.update_source_sync_policy.calls == [] + command = narrowed(facade.update_source.last_call.kwargs["payload"], KnowledgeFSAtomicSourceUpdatePayload) + assert command.expected_version == 3 + assert command.expected_policy_revision == 2 + assert command.sync_policy is not None + assert command.sync_policy.mode == "interval" + assert result is source + + @staticmethod + def test_source_edit_reimports_when_selection_adds_a_canonical_duplicate() -> None: + source = source_response( + metadata={"crawled": {"https://docs.dify.ai/a": {}}}, + status="active", + type="web", + uri="https://docs.dify.ai", + ) + payload = KnowledgeFSSourceUpdatePayload.model_validate( + { + "selection": { + "kind": "website_crawl", + "sourceUrls": [ + "https://docs.dify.ai/a/", + "https://DOCS.dify.ai:443/a#section", + ], + }, + "syncPolicy": {"enabled": True, "mode": "interval"}, + } + ) + + assert source_edit_service.source_edit_requires_import(source, payload) is True + + @staticmethod + def test_source_edit_ignores_order_when_canonical_url_multiplicity_is_unchanged() -> None: + source = source_response( + metadata={ + "crawled": { + "https://docs.dify.ai/a/": {}, + "https://DOCS.dify.ai:443/a#stored": {}, + } + }, + status="active", + type="web", + uri="https://docs.dify.ai", + ) + payload = KnowledgeFSSourceUpdatePayload.model_validate( + { + "selection": { + "kind": "website_crawl", + "sourceUrls": [ + "https://docs.dify.ai/a#selected", + "https://docs.dify.ai/a/", + ], + }, + "syncPolicy": {"enabled": True, "mode": "interval"}, + } + ) + + assert source_edit_service.source_edit_requires_import(source, payload) is False + + @staticmethod + def test_source_edit_uses_latest_effective_website_selection_instead_of_stale_crawl_metadata() -> None: + source = source_response( + metadata={ + "crawled": {"https://docs.dify.ai/old": {}}, + "initialPreview": {"canonicalSourceUrls": ["https://docs.dify.ai/old"]}, + "__knowledgeFsWebsiteSelection": { + "sourceUrls": ["https://docs.dify.ai/current/"], + "version": 1, + }, + }, + status="active", + type="web", + uri="https://docs.dify.ai", + ) + unchanged = KnowledgeFSSourceUpdatePayload.model_validate( + { + "selection": { + "kind": "website_crawl", + "sourceUrls": ["https://DOCS.dify.ai:443/current#selected"], + }, + "syncPolicy": {"enabled": True, "mode": "interval"}, + } + ) + changed = KnowledgeFSSourceUpdatePayload.model_validate( + { + "selection": { + "kind": "website_crawl", + "sourceUrls": ["https://docs.dify.ai/old"], + }, + "syncPolicy": {"enabled": True, "mode": "interval"}, + } + ) + + assert source_edit_service.source_edit_requires_import(source, unchanged) is False + assert source_edit_service.source_edit_requires_import(source, changed) is True + + @staticmethod + def test_source_edit_reimports_an_unchanged_selection_when_source_is_in_error() -> None: + provider_item_id = '["workspace-a","page-a"]' + identity_hash = sha256(f"online-document\0{provider_item_id}".encode()).hexdigest() + source = source_response( + metadata={ + "providerKind": "online-document", + "__knowledgeFsProviderSelection": {"identityHashes": [identity_hash]}, + }, + status="error", + type="connector", + uri="notion://connection-a", + ) + payload = KnowledgeFSSourceUpdatePayload.model_validate( + { + "selection": { + "kind": "online_document", + "items": [ + { + "pageId": "page-a", + "providerItemId": provider_item_id, + "type": "page", + "workspaceId": "workspace-a", + } + ], + }, + "syncPolicy": {"enabled": True, "mode": "interval"}, + } + ) + + assert source_edit_service.source_edit_requires_import(source, payload) is True + + @staticmethod + def test_console_source_update_applies_policy_without_creating_import_work( + monkeypatch: pytest.MonkeyPatch, + ) -> None: + source = source_response(version=5, sync_policy=None) + policy = sync_policy_response(revision=3) + source.sync_policy = policy + facade = FakeComponent( + get_source=Recorder(source), + get_source_sync_policy=Recorder(sync_policy_response(revision=2)), + update_source=Recorder(source), + update_source_sync_policy=Recorder(), + ) + runtime = FakeRuntime(facade=facade) + payload = KnowledgeFSSourceUpdatePayload.model_validate( + {"syncPolicy": {"customIntervalSeconds": 7200, "enabled": True, "mode": "custom"}} + ) + commit = Recorder() + monkeypatch.setattr(console_resources, "_actor", lambda: ("account-1", "tenant-1")) + monkeypatch.setattr(console_resources, "_console_services", lambda: runtime) + monkeypatch.setattr(console_resources, "_payload", lambda _: payload) + monkeypatch.setattr(source_edit_service, "commit_source_import", commit) + monkeypatch.setattr(console_resources, "dump_response", lambda _schema, raw: raw) + app = Flask(__name__) + + with app.test_request_context("/", method="PATCH"): + result = _invoke( + console_resources.KnowledgeFSSpaceSourceApi, + console_resources.KnowledgeFSSpaceSourceApi.patch, + "space-1", + "source-1", + ) + + assert len(facade.update_source.calls) == 1 + assert commit.calls == [] + assert facade.update_source_sync_policy.calls == [] + policy_payload = narrowed( + facade.update_source.last_call.kwargs["payload"], KnowledgeFSAtomicSourceUpdatePayload + ) + assert policy_payload.expected_policy_revision == 2 + assert policy_payload.expected_version == 5 + assert policy_payload.sync_policy is not None + assert policy_payload.sync_policy.custom_interval_seconds == 7200 + assert source.sync_policy is policy + assert result is source + + @staticmethod + def test_console_metadata_delete_forwards_row_version_cas(monkeypatch: pytest.MonkeyPatch) -> None: + facade = FakeComponent(delete_metadata_field=Recorder(_RAW_RESULT)) + runtime = FakeRuntime(facade=facade) + monkeypatch.setattr(console_resources, "_actor", lambda: ("account-1", "tenant-1")) + monkeypatch.setattr(console_resources, "_console_services", lambda: runtime) + monkeypatch.setattr(console_resources, "dump_response", lambda _schema, raw: raw) + app = Flask(__name__) + + with app.test_request_context("/?expectedRowVersion=7", method="DELETE"): + result = _invoke( + console_resources.KnowledgeFSSpaceMetadataFieldApi, + console_resources.KnowledgeFSSpaceMetadataFieldApi.delete, + "space-1", + "field-1", + ) + + assert facade.delete_metadata_field.calls == [ + invocation( + tenant_id="tenant-1", + account_id="account-1", + control_space_id="space-1", + field_id="field-1", + expected_row_version=7, + ) + ] + assert result is _RAW_RESULT + + @staticmethod + def test_console_overview_stats_composes_kfs_metrics_with_dify_app_bindings( + monkeypatch: pytest.MonkeyPatch, + ) -> None: + stats = KnowledgeFSOverviewBaseStatsResponse.model_validate( + { + "current": { + "freshSourceCount": 2, + "knowledgeCount": 13, + "latestSourceSyncAt": "2026-07-23T11:00:00.000Z", + "linkedAppCount": 0, + "sourceCount": 3, + "staleSourceCount": 1, + }, + "generatedAt": "2026-07-23T11:59:59.000Z", + "knowledgeSpaceId": "knowledge-space-1", + "windows": { + "24h": { + "answerRate": 0.8, + "answeredQueryCount": 8, + "queryCount": 10, + "since": "2026-07-22T12:00:00.000Z", + }, + "7d": { + "answerRate": 0.75, + "answeredQueryCount": 75, + "queryCount": 100, + "since": "2026-07-16T12:00:00.000Z", + }, + "30d": { + "answerRate": 0.7, + "answeredQueryCount": 210, + "queryCount": 300, + "since": "2026-06-23T12:00:00.000Z", + }, + }, + } + ) + outcomes = KnowledgeFSOverviewQueryOutcomesResponse.model_validate( + { + "buckets": [], + "current": { + "answerRate": 0.8, + "answered": 96, + "lowConfidence": 12, + "noEvidence": 12, + "queryCount": 120, + }, + "generatedAt": "2026-07-23T12:00:00.000Z", + "knowledgeSpaceId": "knowledge-space-1", + "previous": { + "answerRate": 0.85, + "answered": 85, + "lowConfidence": 8, + "noEvidence": 7, + "queryCount": 100, + }, + "previousSince": "2026-07-09T12:00:00.000Z", + "since": "2026-07-16T12:00:00.000Z", + "window": "7d", + } + ) + runtime = FakeRuntime( + facade=FakeComponent( + get_overview_query_outcomes=Recorder(outcomes), + get_overview_stats=Recorder(stats), + ), + app_bindings=FakeComponent(count_active=Recorder(7)), + ) + dump_response = Recorder(effect=lambda _schema, raw: raw) + monkeypatch.setattr(console_resources, "_actor", lambda: ("account-1", "tenant-1")) + monkeypatch.setattr(console_resources, "_console_services", lambda: runtime) + monkeypatch.setattr(console_resources, "dump_response", dump_response) + app = Flask(__name__) + + with app.test_request_context("/?window=7d"): + result = _invoke( + console_resources.KnowledgeFSSpaceOverviewStatsApi, + console_resources.KnowledgeFSSpaceOverviewStatsApi.get, + "control-space-1", + ) + + assert runtime.facade.get_overview_stats.calls == [ + invocation( + tenant_id="tenant-1", + account_id="account-1", + control_space_id="control-space-1", + ) + ] + assert runtime.facade.get_overview_query_outcomes.calls == [ + invocation( + tenant_id="tenant-1", + account_id="account-1", + control_space_id="control-space-1", + window="7d", + ) + ] + assert runtime.app_bindings.count_active.calls == [ + invocation( + tenant_id="tenant-1", + actor_account_id="account-1", + control_space_id="control-space-1", + ) + ] + assert result.documents == 13 + assert result.linked_apps == 7 + assert result.queries.value == 120 + assert result.queries.previous_value == 100 + assert result.queries.change_rate == pytest.approx(0.2) + assert result.answer_rate.value == pytest.approx(0.8) + assert result.answer_rate.change_percentage_points == pytest.approx(-5) + assert result.freshness_seconds == 3600 + assert result.generated_at == outcomes.generated_at + assert dump_response.last_call.args[1] is result + + @staticmethod + @pytest.mark.parametrize( + ("class_name", "method_name", "route_args", "operation_id", "expected_fields"), + _SERVICE_DELEGATION_CASES, + ) + def test_service_resources_bind_route_identifiers_to_one_declared_operation( + monkeypatch: pytest.MonkeyPatch, + class_name: type[Resource], + method_name: Callable[..., object], + route_args: tuple[object, ...], + operation_id: str, + expected_fields: dict[str, object], + ) -> None: + payload = object() + execute = Recorder(_RAW_RESULT) + dump_response = Recorder(effect=lambda schema, raw: (schema.__name__, raw)) + monkeypatch.setattr(service_resources, "_execute_service_operation", execute) + monkeypatch.setattr(service_resources, "_payload", lambda _: payload) + monkeypatch.setattr(service_resources, "_idempotency_key", lambda: "idempotency-1") + monkeypatch.setattr(service_resources, "_query_pairs", lambda _: (("normalized", "true"),)) + monkeypatch.setattr(service_resources, "dump_response", dump_response) + monkeypatch.setattr(service_resources, "durable_page_import", lambda value: value) + monkeypatch.setattr(service_resources, "durable_file_import", lambda value: value) + monkeypatch.setattr( + service_resources, + "_deletion_accepted", + lambda _space, raw: dump_response(service_resources.KnowledgeFSDurableDeletionAcceptedResponse, raw), + ) + monkeypatch.setattr( + service_resources, + "_bulk_deletion_accepted", + lambda _space, raw: dump_response(service_resources.KnowledgeFSDurableDeletionAcceptedResponse, raw), + ) + app = Flask(__name__) + + with app.test_request_context("/", method="POST", headers={"Idempotency-Key": "idempotency-1"}): + _invoke(class_name, method_name, *route_args) + + assert len(execute.calls) == 1 + call_fields = execute.last_call.kwargs + assert call_fields["control_space_id"] == "space-1" + assert call_fields["operation_id"] == operation_id + for field_name, expected in expected_fields.items(): + assert call_fields[field_name] == expected + if "payload" in call_fields: + assert call_fields["payload"] is payload + if "headers" in call_fields: + assert call_fields["headers"] == (("Idempotency-Key", "idempotency-1"),) + if "query" in call_fields: + assert call_fields["query"] == (("normalized", "true"),) + assert dump_response.last_call.args[1] is _RAW_RESULT + + @staticmethod + @pytest.mark.parametrize( + ("class_name", "method_name", "operation_id"), + [ + ( + service_resources.KnowledgeFSServiceDocumentsApi, + service_resources.KnowledgeFSServiceDocumentsApi.get, + "listDocuments", + ), + ( + service_resources.KnowledgeFSServiceSettingsApi, + service_resources.KnowledgeFSServiceSettingsApi.get, + "getSettings", + ), + ( + service_resources.KnowledgeFSServiceSettingsApi, + service_resources.KnowledgeFSServiceSettingsApi.patch, + "updateSettings", + ), + ( + service_resources.KnowledgeFSServiceSourcesApi, + service_resources.KnowledgeFSServiceSourcesApi.get, + "listSources", + ), + ( + service_resources.KnowledgeFSServiceSourcesApi, + service_resources.KnowledgeFSServiceSourcesApi.post, + "createSource", + ), + ( + service_resources.KnowledgeFSServiceResearchTasksApi, + service_resources.KnowledgeFSServiceResearchTasksApi.get, + "listResearchTasks", + ), + ( + service_resources.KnowledgeFSServiceTracesApi, + service_resources.KnowledgeFSServiceTracesApi.get, + "listTraces", + ), + ], + ) + @pytest.mark.parametrize("cursor", [None, "cursor-2"]) + def test_service_credential_routes_validate_profile_before_facade_delegation( + monkeypatch: pytest.MonkeyPatch, + class_name: type[Resource], + method_name: Callable[..., object], + operation_id: str, + cursor: str | None, + ) -> None: + facade = FakeComponent( + execute_service=Recorder(_RAW_RESULT), + update_service_settings=Recorder(_SettingsUpdateResult()), + ) + runtime = FakeRuntime(facade=facade) + profile = object() + validate_profile = Recorder(profile) + monkeypatch.setattr(service_resources, "_runtime", lambda: runtime) + monkeypatch.setattr(service_resources, "_profile", validate_profile) + monkeypatch.setattr(service_resources, "_payload", lambda _: object()) + monkeypatch.setattr(service_resources, "dump_response", lambda schema, raw: (schema.__name__, raw)) + app = Flask(__name__) + query_string = {"cursor": cursor} if cursor is not None else None + + with app.test_request_context("/", method="POST", query_string=query_string): + _invoke(class_name, method_name, "space-1") + + assert validate_profile.calls == [invocation(runtime, operation_id=operation_id, control_space_id="space-1")] + if operation_id == "updateSettings": + assert facade.update_service_settings.last_call.kwargs["profile"] is profile + assert facade.execute_service.calls == [] + return + assert facade.execute_service.last_call.kwargs["profile"] is profile + assert facade.execute_service.last_call.kwargs["operation_id"] == operation_id + if "query" in facade.execute_service.last_call.kwargs: + expected_query = (("limit", "50"),) + ((("cursor", cursor),) if cursor else ()) + assert facade.execute_service.last_call.kwargs["query"] == expected_query + + @staticmethod + @pytest.mark.parametrize( + "class_name", + [ + console_resources.KnowledgeFSSpaceQueriesApi, + ], + ) + def test_deprecated_buffered_routes_fail_closed( + class_name: type[console_resources.KnowledgeFSSpaceQueriesApi], + ) -> None: + with pytest.raises(KnowledgeFSOperationUnavailableError, match="deprecated"): + _invoke(class_name, class_name.post, "space-1") + + @staticmethod + def test_console_stream_capabilities_bind_the_authorized_resource(monkeypatch: pytest.MonkeyPatch) -> None: + broker = FakeComponent(issue_interactive=Recorder(issued_capability())) + payloads = iter( + [ + KnowledgeFSQueryCreatePayload(query="question", mode="deep"), + KnowledgeFSStreamCapabilityPayload(control_space_id="space-2"), + ] + ) + apply_config_overrides(monkeypatch, CONSOLE_API_URL="https://dify.example") + monkeypatch.setattr(console_resources, "_actor", lambda: ("account-1", "tenant-1")) + monkeypatch.setattr( + console_resources, + "_console_services", + lambda: FakeRuntime(broker=broker), + ) + monkeypatch.setattr(console_resources, "_payload", lambda _: next(payloads)) + monkeypatch.setattr(console_resources, "dump_response", lambda _, response: response) + app = Flask(__name__) + + with app.test_request_context("/", method="POST"): + query = _invoke( + console_resources.KnowledgeFSSpaceQueryAdmissionApi, + console_resources.KnowledgeFSSpaceQueryAdmissionApi.post, + "space-1", + ) + task = _invoke( + console_resources.KnowledgeFSTaskStreamCapabilityApi, + console_resources.KnowledgeFSTaskStreamCapabilityApi.post, + "task/1", + ) + legacy_query = _invoke( + console_resources.KnowledgeFSSpaceQueryStreamCapabilityApi, + console_resources.KnowledgeFSSpaceQueryStreamCapabilityApi.post, + "space-1", + ) + + assert query.request.knowledge_space_id == "knowledge-space-1" + assert query.url == "https://dify.example/console/api/knowledge-fs/query-stream" + assert ( + task.url == "https://dify.example/console/api/knowledge-fs/research-tasks/task%2F1/events" + "?knowledgeSpaceId=knowledge-space-1" + ) + assert legacy_query.url == "https://dify.example/console/api/knowledge-fs/query-stream" + assert broker.issue_interactive.calls[0].kwargs["operation_id"] == "createQuery" + assert broker.issue_interactive.calls[1].kwargs == { + "tenant_id": "tenant-1", + "account_id": "account-1", + "control_space_id": "space-2", + "operation_id": "streamResearchTask", + "resource_id": "task/1", + } + assert broker.issue_interactive.calls[2].kwargs["operation_id"] == "createQuery" + + @staticmethod + def test_service_query_admission_binds_profile_space_and_payload(monkeypatch: pytest.MonkeyPatch) -> None: + broker = FakeComponent(issue_service=Recorder(issued_capability())) + runtime = FakeRuntime(broker=broker) + profile = object() + apply_config_overrides(monkeypatch, SERVICE_API_URL="https://api.dify.example") + monkeypatch.setattr(service_resources, "_runtime", lambda: runtime) + monkeypatch.setattr(service_resources, "_profile", lambda *_args, **_kwargs: profile) + monkeypatch.setattr(service_resources, "_payload", lambda _: KnowledgeFSQueryCreatePayload(query="question")) + monkeypatch.setattr(service_resources, "dump_response", lambda _, response: response) + app = Flask(__name__) + + with app.test_request_context("/", method="POST"): + response = _invoke( + service_resources.KnowledgeFSServiceQueryAdmissionApi, + service_resources.KnowledgeFSServiceQueryAdmissionApi.post, + "space-1", + ) + + assert broker.issue_service.calls == [invocation(profile=profile, operation_id="createQuery")] + assert response.request.knowledge_space_id == "knowledge-space-1" + assert response.request.query == "question" + assert response.url == "https://api.dify.example/v1/knowledge-fs/query-stream" + + @staticmethod + def test_console_resource_helpers_validate_feature_payload_headers_and_query_pairs( + monkeypatch: pytest.MonkeyPatch, + ) -> None: + runtime = FakeRuntime() + session_maker = object() + get_runtime = Recorder(runtime) + apply_config_overrides(monkeypatch, DEPLOYMENT_EDITION=DeploymentEdition.CLOUD, KNOWLEDGE_FS_ENABLED=True) + monkeypatch.setattr(console_resources.session_factory, "get_session_maker", lambda: session_maker) + monkeypatch.setattr(console_resources, "get_knowledge_fs_runtime", get_runtime) + monkeypatch.setattr( + console_resources, + "current_account_with_tenant", + lambda: (_knowledge_fs_resource_delegation_account("account-1"), "tenant-1"), + ) + app = Flask(__name__) + + assert console_resources._console_services() is runtime + assert get_runtime.calls == [invocation(session_maker)] + assert console_resources._actor() == ("account-1", "tenant-1") + query = KnowledgeFSQueryCreatePayload(query="question", mode="fast", active_document_ids=["doc-1"]) + assert console_resources._query_pairs(query) == ( + ("query", "question"), + ("mode", "fast"), + ("activeDocumentIds", "['doc-1']"), + ("activeEntityIds", "[]"), + ) + with app.test_request_context( + "/", + method="POST", + json={"query": "question"}, + headers={"Idempotency-Key": "request-key"}, + ): + assert console_resources._payload(KnowledgeFSQueryCreatePayload).query == "question" + assert console_resources._idempotency_key() == "request-key" + + apply_config_overrides(monkeypatch, DEPLOYMENT_EDITION=DeploymentEdition.COMMUNITY) + with pytest.raises(NotFound): + console_resources._console_services() + + apply_config_overrides(monkeypatch, KNOWLEDGE_FS_COMMUNITY_DEV_ENABLED=True) + assert console_resources._console_services() is runtime + + apply_config_overrides(monkeypatch, KNOWLEDGE_FS_ENABLED=False) + with pytest.raises(NotFound): + console_resources._console_services() + + @staticmethod + def test_service_resource_helpers_validate_feature_bearer_headers_and_boolean_queries( + monkeypatch: pytest.MonkeyPatch, + ) -> None: + runtime = FakeRuntime() + session_maker = object() + get_runtime = Recorder(runtime) + apply_config_overrides(monkeypatch, DEPLOYMENT_EDITION=DeploymentEdition.CLOUD, KNOWLEDGE_FS_ENABLED=True) + monkeypatch.setattr(service_resources.session_factory, "get_session_maker", lambda: session_maker) + monkeypatch.setattr(service_resources, "get_knowledge_fs_runtime", get_runtime) + app = Flask(__name__) + + assert service_resources._runtime() is runtime + assert get_runtime.calls == [invocation(session_maker)] + + class Query(BaseModel): + enabled: bool = True + disabled: bool = False + count: int = 2 + + query = Query() + assert service_resources._query_pairs(query) == ( + ("enabled", "true"), + ("disabled", "false"), + ("count", "2"), + ) + with app.test_request_context( + "/", + method="POST", + json={"query": "question"}, + headers={"Authorization": "bEaReR credential-value ", "Idempotency-Key": "request-key"}, + ): + assert service_resources._payload(KnowledgeFSQueryCreatePayload).query == "question" + assert service_resources._idempotency_key() == "request-key" + + apply_config_overrides(monkeypatch, DEPLOYMENT_EDITION=DeploymentEdition.COMMUNITY) + with pytest.raises(NotFound): + service_resources._runtime() + + apply_config_overrides(monkeypatch, KNOWLEDGE_FS_COMMUNITY_DEV_ENABLED=True) + assert service_resources._runtime() is runtime + + apply_config_overrides(monkeypatch, KNOWLEDGE_FS_ENABLED=False) + with pytest.raises(NotFound): + service_resources._runtime() + + @staticmethod + def test_console_request_rejections_preserve_http_status_contracts() -> None: + from services.knowledge_fs.product_remote import KnowledgeFSProductRequestRejectedError + + expected: dict[Literal[400, 403, 409, 413, 422, 429, 503, 504], type[BaseHTTPException]] = { + 400: KnowledgeFSInvalidRequestHTTPError, + 403: KnowledgeFSAccessDeniedHTTPError, + 409: KnowledgeFSConflictHTTPError, + 413: KnowledgeFSRequestTooLargeHTTPError, + 422: KnowledgeFSRequestRejectedHTTPError, + 429: KnowledgeFSRateLimitHTTPError, + 503: KnowledgeFSOperationUnavailableHTTPError, + 504: KnowledgeFSTimeoutHTTPError, + } + for status, http_error in expected.items(): + guarded = console_resources._knowledge_fs_errors( + raising(KnowledgeFSProductRequestRejectedError(status_code=status)) + ) + + with pytest.raises(http_error) as raised: + guarded() + + assert raised.value.code == status + assert raised.value.data == { + "code": http_error.error_code, + "message": http_error.description, + "status": status, + } + + @staticmethod + @pytest.mark.parametrize( + ("status", "failure_code", "category", "safe_message"), + [ + ( + 503, + "DOCUMENT_COUNT_QUOTA_UNAVAILABLE", + "dependency", + "The workspace document quota could not be verified. Try again later.", + ), + ( + 503, + "VECTOR_SPACE_QUOTA_UNAVAILABLE", + "dependency", + "The workspace vector storage quota could not be verified. Try again later.", + ), + (504, "KNOWLEDGE_FS_TIMEOUT", "timeout", "The Agent Knowledge Base operation timed out. Try again later."), + ], + ) + def test_console_unavailable_and_timeout_rejections_preserve_safe_failure_metadata( + status: Literal[503, 504], failure_code: str, category: str, safe_message: str + ) -> None: + from services.knowledge_fs.product_remote import KnowledgeFSProductRequestRejectedError + + expected_failure = { + "action": "retry", + "category": category, + "code": failure_code, + "message": safe_message, + "retryPolicy": "manual", + "traceId": "trace-quota", + "parameters": {"retryAfterSeconds": 30}, + } + failure = KnowledgeFSPublicFailureResponse.model_validate( + {**expected_failure, "message": "Private upstream diagnostic that must not reach the client"} + ) + guarded = console_resources._knowledge_fs_errors( + raising(KnowledgeFSProductRequestRejectedError(status_code=status, failure=failure)) + ) + http_error = KnowledgeFSOperationUnavailableHTTPError if status == 503 else KnowledgeFSTimeoutHTTPError + + with pytest.raises(http_error) as raised: + guarded() + + assert raised.value.code == status + assert raised.value.data == { + "code": http_error.error_code, + "message": safe_message, + "status": status, + "failure": expected_failure, + } + + @staticmethod + def test_console_request_rejection_preserves_safe_failure_metadata() -> None: + from services.knowledge_fs.product_remote import KnowledgeFSProductRequestRejectedError + + failure = KnowledgeFSPublicFailureResponse.model_validate( + { + "action": "configure_model", + "category": "configuration", + "code": "MODEL_SELECTION_NOT_FOUND", + "message": "Select another model before retrying.", + "retryPolicy": "after_configuration", + "traceId": "trace-model", + } + ) + guarded = console_resources._knowledge_fs_errors( + raising(KnowledgeFSProductRequestRejectedError(status_code=422, failure=failure)) + ) + + with pytest.raises(KnowledgeFSRequestRejectedHTTPError) as raised: + guarded() + + assert raised.value.data == { + "code": "knowledge_fs_request_rejected", + "message": "The Agent Knowledge Base operation requires a configuration change before it can continue.", + "status": 422, + "failure": { + "action": "configure_model", + "category": "configuration", + "code": "MODEL_SELECTION_NOT_FOUND", + "message": "The Agent Knowledge Base operation requires a configuration change before it can continue.", + "retryPolicy": "after_configuration", + "traceId": "trace-model", + }, + } + + @staticmethod + def test_jwks_resource_fails_closed_for_disabled_missing_and_misconfigured_issuers( + monkeypatch: pytest.MonkeyPatch, + ) -> None: + from services.knowledge_fs_capability import KnowledgeFSCapabilityConfigurationError + + app = Flask(__name__) + apply_config_overrides(monkeypatch, KNOWLEDGE_FS_CAPABILITY_V2_ENABLED=False) + with app.app_context(), pytest.raises(NotFound): + console_resources.KnowledgeFSJWKSApi().get() + + apply_config_overrides(monkeypatch, KNOWLEDGE_FS_CAPABILITY_V2_ENABLED=True) + monkeypatch.setattr(console_resources.session_factory, "get_session_maker", lambda: object()) + monkeypatch.setattr(console_resources, "create_configured_knowledge_fs_capability_issuer", lambda **_: None) + with app.app_context(), pytest.raises(NotFound): + console_resources.KnowledgeFSJWKSApi().get() + + def misconfigured(**_: object) -> None: + raise KnowledgeFSCapabilityConfigurationError("missing signing key") + + monkeypatch.setattr(console_resources, "create_configured_knowledge_fs_capability_issuer", misconfigured) + with app.app_context(), pytest.raises(ServiceUnavailable, match="not configured"): + console_resources.KnowledgeFSJWKSApi().get() + + @staticmethod + def test_console_error_adapter_maps_every_domain_boundary_to_the_stable_http_contract() -> None: + from pydantic import ValidationError + + from services.knowledge_fs.app_binding_management import KnowledgeFSAppBindingManagementError + from services.knowledge_fs.control_plane_service import KnowledgeFSControlPlaneInvariantError + from services.knowledge_fs.product_authorization import KnowledgeFSProductNotFoundError + from services.knowledge_fs.product_remote import ( + KnowledgeFSProductRemoteError, + KnowledgeFSProductResourceNotFoundError, + ) + from services.knowledge_fs_capability import KnowledgeFSCapabilityConfigurationError + + with pytest.raises(ValidationError) as raised_validation: + KnowledgeFSQueryCreatePayload.model_validate({"query": ""}) + validation_error = raised_validation.value + mappings = ( + (KnowledgeFSProductNotFoundError("hidden"), KnowledgeFSSpaceNotFoundHTTPError), + (KnowledgeFSProductResourceNotFoundError("missing child"), KnowledgeFSResourceNotFoundHTTPError), + (KnowledgeFSOperationUnavailableError("manifest mismatch"), KnowledgeFSOperationUnavailableHTTPError), + (KnowledgeFSProductRemoteError("upstream unavailable"), KnowledgeFSUpstreamUnavailableHTTPError), + (KnowledgeFSAppBindingManagementError("invalid binding"), KnowledgeFSInvalidRequestHTTPError), + (KnowledgeFSControlPlaneInvariantError("missing revision"), KnowledgeFSInvalidRequestHTTPError), + (validation_error, KnowledgeFSInvalidRequestHTTPError), + (PermissionError("forbidden"), KnowledgeFSAccessDeniedHTTPError), + (KnowledgeFSCapabilityConfigurationError("missing issuer"), KnowledgeFSOperationUnavailableHTTPError), + ) + + for domain_error, http_error in mappings: + guarded = console_resources._knowledge_fs_errors(raising(domain_error)) + + with pytest.raises(http_error): + guarded() + + assert KnowledgeFSUpstreamUnavailableHTTPError.code == HTTPStatus.SERVICE_UNAVAILABLE + + @staticmethod + def test_service_error_adapter_maps_every_domain_boundary_to_the_stable_http_contract() -> None: + from pydantic import ValidationError + + from services.knowledge_fs.product_remote import ( + KnowledgeFSProductRemoteError, + KnowledgeFSProductRequestRejectedError, + KnowledgeFSProductResourceNotFoundError, + ) + from services.knowledge_fs.service_api_authorization import KnowledgeFSServiceApiAuthorizationError + + with pytest.raises(ValidationError) as raised_validation: + KnowledgeFSQueryCreatePayload.model_validate({"query": ""}) + validation_error = raised_validation.value + mappings = ( + (KnowledgeFSServiceApiAuthorizationError("revoked"), KnowledgeFSInvalidCredentialHTTPError), + ( + KnowledgeFSOperationUnavailableError("manifest mismatch"), + KnowledgeFSServiceOperationUnavailableHTTPError, + ), + ( + KnowledgeFSProductResourceNotFoundError("missing child"), + KnowledgeFSServiceResourceNotFoundHTTPError, + ), + (KnowledgeFSProductRemoteError("upstream unavailable"), KnowledgeFSServiceUpstreamUnavailableHTTPError), + (validation_error, KnowledgeFSServiceInvalidRequestHTTPError), + ) + + for domain_error, http_error in mappings: + guarded = service_resources._service_api_errors(raising(domain_error)) + + with pytest.raises(http_error): + guarded() + + rejected_mappings = ( + (400, KnowledgeFSServiceInvalidRequestHTTPError), + (403, KnowledgeFSServiceAccessDeniedHTTPError), + (409, KnowledgeFSServiceConflictHTTPError), + (413, KnowledgeFSServiceRequestTooLargeHTTPError), + (422, KnowledgeFSServiceRequestRejectedHTTPError), + (429, KnowledgeFSServiceRateLimitHTTPError), + ) + for status, http_error in rejected_mappings: + guarded = service_resources._service_api_errors( + raising(KnowledgeFSProductRequestRejectedError(status_code=status)) + ) + with pytest.raises(http_error): + guarded() + + @staticmethod + def test_console_small_file_reader_rejects_malformed_and_empty_multipart_bodies() -> None: + from io import BytesIO + + from services.knowledge_fs.product_remote import KnowledgeFSProductRequestRejectedError + + app = Flask(__name__) + with app.test_request_context( + "/", + method="POST", + data=b"x", + environ_overrides={"CONTENT_LENGTH": str(64 * 1024 + 2)}, + ): + with pytest.raises(KnowledgeFSProductRequestRejectedError) as too_large: + console_resources._read_small_file_body(1) + assert too_large.value.status_code == HTTPStatus.REQUEST_ENTITY_TOO_LARGE + + malformed_payloads = ( + dict[str, object](), + {"file": (BytesIO(b"content"), "")}, + {"file": (BytesIO(b""), "empty.txt")}, + ) + for payload in malformed_payloads: + with app.test_request_context("/", method="POST", data=payload, content_type="multipart/form-data"): + with pytest.raises(KnowledgeFSProductRequestRejectedError) as rejected: + console_resources._read_small_file_body(10) + assert rejected.value.status_code == HTTPStatus.UNPROCESSABLE_ENTITY + + @staticmethod + def test_console_app_binding_route_rejects_unknown_caller_kind_before_revoke( + monkeypatch: pytest.MonkeyPatch, + ) -> None: + from services.knowledge_fs.app_binding_management import KnowledgeFSAppBindingManagementError + + revoke = Recorder() + monkeypatch.setattr(console_resources, "_actor", lambda: ("account-1", "tenant-1")) + monkeypatch.setattr( + console_resources, + "_console_services", + lambda: FakeRuntime(app_bindings=FakeComponent(revoke=revoke)), + ) + + with pytest.raises(KnowledgeFSAppBindingManagementError, match="caller kind"): + _invoke( + console_resources.KnowledgeFSSpaceAppBindingApi, + console_resources.KnowledgeFSSpaceAppBindingApi.delete, + "space-1", + "unknown-caller", + "app-1", + ) + + assert revoke.calls == [] + + @staticmethod + def test_service_profile_hides_unknown_operations_before_dataset_key_authorization( + monkeypatch: pytest.MonkeyPatch, + ) -> None: + authorization = FakeComponent(authorize=Recorder()) + # The profile check must fail before it touches any runtime component but authorization. + runtime = cast(KnowledgeFSRuntime, FakeRuntime(service_api_authorization=authorization)) + monkeypatch.setattr(service_resources, "product_operation_action", raising(KeyError("unknown"))) + + with pytest.raises(KnowledgeFSOperationUnavailableError, match="unknownOperation"): + service_resources._profile( + runtime, + operation_id="unknownOperation", + control_space_id="space-1", + ) + + assert authorization.authorize.calls == [] + + @staticmethod + def test_service_operation_helper_forwards_the_complete_validated_request( + monkeypatch: pytest.MonkeyPatch, + ) -> None: + profile = object() + payload = KnowledgeFSQueryCreatePayload(query="question") + facade = FakeComponent(execute_service=Recorder({"id": "result-1"})) + runtime = FakeRuntime(facade=facade) + monkeypatch.setattr(service_resources, "_runtime", lambda: runtime) + monkeypatch.setattr(service_resources, "_profile", Recorder(profile)) + + result = service_resources._execute_service_operation( + control_space_id="space-1", + operation_id="createResearchTask", + payload=payload, + query=(("cursor", "next"),), + bind_space_in_body=True, + resource_id="task-1", + path_parameters=(("taskId", "task-1"),), + headers=(("Idempotency-Key", "request-key"),), + ) + + assert result == {"id": "result-1"} + assert facade.execute_service.calls == [ + invocation( + profile=profile, + operation_id="createResearchTask", + payload=payload, + query=(("cursor", "next"),), + bind_space_in_body=True, + resource_id="task-1", + path_parameters=(("taskId", "task-1"),), + headers=(("Idempotency-Key", "request-key"),), + ) + ] diff --git a/api/tests/unit_tests/controllers/console/snippets/test_snippet_list_query.py b/api/tests/unit_tests/controllers/console/snippets/test_payloads.py similarity index 100% rename from api/tests/unit_tests/controllers/console/snippets/test_snippet_list_query.py rename to api/tests/unit_tests/controllers/console/snippets/test_payloads.py diff --git a/api/tests/unit_tests/controllers/console/snippets/test_snippet_workflow.py b/api/tests/unit_tests/controllers/console/snippets/test_snippet_workflow.py index a3b3ddfd15f..eef3224942c 100644 --- a/api/tests/unit_tests/controllers/console/snippets/test_snippet_workflow.py +++ b/api/tests/unit_tests/controllers/console/snippets/test_snippet_workflow.py @@ -13,8 +13,8 @@ from controllers.console.snippets import snippet_workflow as snippet_workflow_mo from models.account import Account, TenantAccountRole from models.snippet import CustomizedSnippet from models.workflow import Workflow -from tests.unit_tests.controllers.error_body import raised_error -from tests.unit_tests.model_factories import make_account +from tests.helpers.unit.controllers.error_body import raised_error +from tests.helpers.unit.model_factories import make_account def _account(account_id: str = "account-1") -> Account: diff --git a/api/tests/unit_tests/controllers/console/tag/test_tags.py b/api/tests/unit_tests/controllers/console/tag/test_tags.py index 41aa923a471..3efc2ff985f 100644 --- a/api/tests/unit_tests/controllers/console/tag/test_tags.py +++ b/api/tests/unit_tests/controllers/console/tag/test_tags.py @@ -31,8 +31,8 @@ from services.tag_application_service import ( TagSummary, UpdateTagInput, ) -from tests.unit_tests.config_override import config_overrides_context -from tests.unit_tests.model_factories import make_account +from tests.helpers.unit.config_override import config_overrides_context +from tests.helpers.unit.model_factories import make_account def unwrap(func): diff --git a/api/tests/unit_tests/controllers/console/test_arize_tracing_ack_routes.py b/api/tests/unit_tests/controllers/console/test_arize_tracing_ack_routes.py deleted file mode 100644 index cf158494655..00000000000 --- a/api/tests/unit_tests/controllers/console/test_arize_tracing_ack_routes.py +++ /dev/null @@ -1,193 +0,0 @@ -"""Both real config Resources reach the same ACK gate before atomic SQL updates. - -HTTP dispatch, payload validation, owner resolution, DI services, encryption and -SQLite persistence are real. Login admission supplies a local owner account and -the existing RSA fixture supplies private-key storage; no external account is used. -""" - -import json -from copy import deepcopy -from uuid import uuid4 - -import pytest -from dify_trace_arize_phoenix.arize_phoenix_trace import ArizePhoenixDataTrace -from flask import Flask -from flask_restx import Api -from sqlalchemy import select -from sqlalchemy.orm import Session, sessionmaker - -from controllers.console import flask_admission, wraps -from controllers.console.agent.tracing import AgentTraceConfigApi -from controllers.console.app.ops_trace import TraceAppConfigApi -from core.helper import ssrf_proxy -from core.ops.ops_trace_manager import OpsTraceManager -from enums import DeploymentEdition -from extensions.ext_application_services import ApplicationServices -from libs.login import AccountWithTenant -from models.account import Account, AccountStatus, TenantAccountRole -from models.agent import Agent, AgentScope, AgentSource, AgentStatus -from models.model import App, AppMode, TraceAppConfig -from repositories.app.tracing_config_repository import SQLAlchemyAppTracingConfigRepository -from tests.unit_tests.config_override import apply_config_overrides -from tests.unit_tests.model_factories import make_app -from tests.unit_tests.services.tracing.test_arize_configuration_ack import ( - collector as collector, # noqa: PLC0414 - Re-export the real local gRPC collector fixture. -) - -pytest_plugins = ("tests.unit_tests.services.tracing.test_nullable_tracing_credentials",) - - -@pytest.mark.parametrize("surface", ["app", "agent"]) -@pytest.mark.parametrize("method", ["POST", "PATCH"]) -def test_config_write_rejects_bad_ack_without_changing_old_rows_policy_or_cache( - workspace: str, - sqlite_session_factory: sessionmaker[Session], - account_application_services: ApplicationServices, - monkeypatch: pytest.MonkeyPatch, - collector: tuple[str, list[tuple[bytes, dict[str, str]]], dict[str, object]], - surface: str, - method: str, -) -> None: - endpoint, calls, _ = collector - app_id, agent_id, account_id = (str(uuid4()) for _ in range(3)) - old_policy = json.dumps( - {"enabled": True, "tracing_provider": "langsmith", "capture_content": False, "capture_agent": False} - ) - with sqlite_session_factory.begin() as session: - session.add( - make_app( - app_id=app_id, - tenant_id=workspace, - mode=AppMode.AGENT if surface == "agent" else AppMode.WORKFLOW, - tracing=old_policy, - ) - ) - if surface == "agent": - session.add( - Agent( - id=agent_id, - tenant_id=workspace, - name="ACK test", - scope=AgentScope.ROSTER, - source=AgentSource.AGENT_APP, - app_id=app_id, - status=AgentStatus.ACTIVE, - created_by=account_id, - updated_by=account_id, - ) - ) - store = SQLAlchemyAppTracingConfigRepository(session_factory=sqlite_session_factory) - original = OpsTraceManager.encrypt_tracing_config( - workspace, "langsmith", {"api_key": "other-provider-key", "project": "other"} - ) - assert store.create(workspace_id=workspace, app_id=app_id, tracing_provider="langsmith", tracing_config=original) - if method == "PATCH": - old = OpsTraceManager.encrypt_tracing_config( - workspace, - "arize", - { - "endpoint": endpoint, - "project": "old-project", - "api_key": "synthetic-good", - "space_id": "synthetic-space", - }, - ) - old.update( - { - "project_url": "https://app.arize.com/?redirect_project_name=old-project", - "arize_project_destination": {"legacy": "unchanged"}, - } - ) - assert store.create(workspace_id=workspace, app_id=app_id, tracing_provider="arize", tracing_config=old) - with sqlite_session_factory() as session: - before = deepcopy( - list( - session.execute( - select(TraceAppConfig.__table__).where(TraceAppConfig.app_id == app_id).order_by(TraceAppConfig.id) - ).mappings() - ) - ) - - account = Account(name="Owner", email="owner@example.test", status=AccountStatus.ACTIVE) - account.id, account.role = account_id, TenantAccountRole.OWNER - identity = AccountWithTenant(account, workspace) - monkeypatch.setattr(flask_admission, "current_account_with_tenant", lambda: identity) - monkeypatch.setattr(wraps, "current_account_with_tenant", lambda: identity) - apply_config_overrides( - monkeypatch, DEPLOYMENT_EDITION=DeploymentEdition.CLOUD, LOGIN_DISABLED=True, RBAC_ENABLED=False - ) - - def forbidden(*_args: object, **_kwargs: object) -> None: - pytest.fail("Local custom collector verification must not perform optional Cloud HTTP reads") - - monkeypatch.setattr(ssrf_proxy, "get", forbidden) - monkeypatch.setattr(ArizePhoenixDataTrace, "__init__", forbidden) - app = Flask(__name__) - app.extensions["application_services"] = account_application_services - api = Api(app) - api.add_resource(TraceAppConfigApi, "/console/api/apps//trace-config") - api.add_resource(AgentTraceConfigApi, "/console/api/agent//trace-config") - url = ( - f"/console/api/apps/{app_id}/trace-config" - if surface == "app" - else f"/console/api/agent/{agent_id}/trace-config" - ) - values = { - "endpoint": endpoint, - "project": "not-created-by-verification", - "api_key": "synthetic-invalid", - "space_id": "synthetic-space", - } - with app.test_client() as client: - failed = client.open( - url, method=method, json={"tracing_provider": "arize", "tracing_config": values, "activate": True} - ) - assert failed.status_code == 400 - body = failed.json - assert isinstance(body, dict) - assert body["code"] == "tracing_config_verification_failed" - assert "private-vendor-details" not in failed.get_data(as_text=True) - assert "synthetic-invalid" not in failed.get_data(as_text=True) - with sqlite_session_factory() as session: - after = list( - session.execute( - select(TraceAppConfig.__table__).where(TraceAppConfig.app_id == app_id).order_by(TraceAppConfig.id) - ).mappings() - ) - assert after == before - assert session.scalar(select(App.tracing).where(App.id == app_id)) == old_policy - values["api_key"] = "synthetic-good" - accepted = client.open( - url, method=method, json={"tracing_provider": "arize", "tracing_config": values, "activate": True} - ) - assert accepted.status_code == (201 if method == "POST" else 200) - saved = store.get(workspace_id=workspace, app_id=app_id, tracing_provider="arize") - assert saved is not None - assert saved.tracing_config is not None - assert ( - OpsTraceManager.decrypt_tracing_config(workspace, "arize", saved.tracing_config)["api_key"] == "synthetic-good" - ) - assert "project_url" not in saved.tracing_config - assert "arize_project_destination" not in saved.tracing_config - assert len(calls) == 2 - with sqlite_session_factory() as session: - policy = json.loads(session.scalar(select(App.tracing).where(App.id == app_id)) or "{}") - assert policy == {"enabled": True, "tracing_provider": "arize", "capture_content": False, "capture_agent": False} - if method == "PATCH": - with app.test_client() as client: - masked = client.patch( - url, - json={ - "tracing_provider": "arize", - "tracing_config": {**values, "project": "next-project", "api_key": "******", "space_id": "******"}, - }, - ) - assert masked.status_code == 200 - updated = store.get(workspace_id=workspace, app_id=app_id, tracing_provider="arize") - assert updated is not None - assert updated.tracing_config is not None - assert updated.tracing_config["api_key"] == saved.tracing_config["api_key"] - assert updated.tracing_config["space_id"] == saved.tracing_config["space_id"] - assert calls[-1][1]["api_key"] == "synthetic-good" - assert calls[-1][1]["space_id"] == "synthetic-space" - assert len(calls) == 3 diff --git a/api/tests/unit_tests/controllers/console/test_extension.py b/api/tests/unit_tests/controllers/console/test_extension.py index 574adca81a8..9527c59a493 100644 --- a/api/tests/unit_tests/controllers/console/test_extension.py +++ b/api/tests/unit_tests/controllers/console/test_extension.py @@ -30,7 +30,7 @@ from services.api_based_extension_application_service import ( APIBasedExtensionRecord, APIBasedExtensionUpdate, ) -from tests.unit_tests.config_override import apply_config_overrides +from tests.helpers.unit.config_override import apply_config_overrides _NEEDS_METHOD_VIEW_CLEANUP = False if not hasattr(builtins, "MethodView"): diff --git a/api/tests/unit_tests/controllers/console/test_fastopenapi_init_validate.py b/api/tests/unit_tests/controllers/console/test_fastopenapi_init_validate.py deleted file mode 100644 index dc8604a88f7..00000000000 --- a/api/tests/unit_tests/controllers/console/test_fastopenapi_init_validate.py +++ /dev/null @@ -1,135 +0,0 @@ -"""HTTP contract tests for the FastOpenAPI initialization routes.""" - -from types import SimpleNamespace -from unittest.mock import Mock, create_autospec - -import pytest - -from controllers.console import init_validate -from dify_app import DifyApp -from enums import DeploymentEdition -from extensions import ext_fastopenapi -from services.init_validation_service import ( - AlreadyInitializedError, - InitValidationService, - InvalidInitializationPasswordError, -) -from tests.unit_tests.config_override import apply_config_overrides - - -@pytest.fixture -def init_validation(monkeypatch: pytest.MonkeyPatch) -> Mock: - service = create_autospec(InitValidationService, instance=True, spec_set=True) - services = SimpleNamespace(init_validation=service) - monkeypatch.setattr(init_validate, "application_services", lambda: services) - return service - - -@pytest.fixture -def app() -> DifyApp: - app = DifyApp(__name__) - app.config.update(TESTING=True, SECRET_KEY="test-secret") - ext_fastopenapi.init_app(app) - return app - - -@pytest.mark.parametrize( - ("validated", "expected_status"), - [ - pytest.param(True, "finished", id="finished"), - pytest.param(False, "not_started", id="not-started"), - ], -) -def test_get_init_status( - app: DifyApp, - init_validation: Mock, - validated: bool, - expected_status: str, -) -> None: - init_validation.is_validated.return_value = validated - - response = app.test_client().get("/console/api/init") - - assert response.status_code == 200 - assert response.get_json() == {"status": expected_status} - init_validation.is_validated.assert_called_once_with(session_validated=False) - - -def test_validate_init_password_success( - app: DifyApp, - init_validation: Mock, - monkeypatch: pytest.MonkeyPatch, -) -> None: - apply_config_overrides(monkeypatch, DEPLOYMENT_EDITION=DeploymentEdition.COMMUNITY) - client = app.test_client() - - response = client.post("/console/api/init", json={"password": "expected"}) - - assert response.status_code == 201 - assert response.get_json() == {"result": "success"} - init_validation.validate_password.assert_called_once_with("expected") - with client.session_transaction() as browser_session: - assert browser_session["is_init_validated"] is True - - -def test_validate_init_password_rejects_a_mismatch( - app: DifyApp, - init_validation: Mock, - monkeypatch: pytest.MonkeyPatch, -) -> None: - apply_config_overrides(monkeypatch, DEPLOYMENT_EDITION=DeploymentEdition.COMMUNITY) - init_validation.validate_password.side_effect = InvalidInitializationPasswordError - client = app.test_client() - - response = client.post("/console/api/init", json={"password": "wrong"}) - - assert response.status_code == 401 - with client.session_transaction() as browser_session: - assert browser_session["is_init_validated"] is False - - -def test_validate_init_password_rejects_an_initialized_installation( - app: DifyApp, - init_validation: Mock, - monkeypatch: pytest.MonkeyPatch, -) -> None: - apply_config_overrides(monkeypatch, DEPLOYMENT_EDITION=DeploymentEdition.COMMUNITY) - init_validation.validate_password.side_effect = AlreadyInitializedError - - response = app.test_client().post("/console/api/init", json={"password": "expected"}) - - assert response.status_code == 403 - - -@pytest.mark.parametrize( - "payload", - [ - pytest.param({}, id="missing-password"), - pytest.param({"password": "x" * 31}, id="password-too-long"), - ], -) -def test_validate_init_password_rejects_an_invalid_payload( - app: DifyApp, - init_validation: Mock, - monkeypatch: pytest.MonkeyPatch, - payload: dict[str, str], -) -> None: - apply_config_overrides(monkeypatch, DEPLOYMENT_EDITION=DeploymentEdition.COMMUNITY) - - response = app.test_client().post("/console/api/init", json=payload) - - assert response.status_code == 422 - init_validation.validate_password.assert_not_called() - - -def test_validate_init_password_is_not_available_in_cloud( - app: DifyApp, - init_validation: Mock, - monkeypatch: pytest.MonkeyPatch, -) -> None: - apply_config_overrides(monkeypatch, DEPLOYMENT_EDITION=DeploymentEdition.CLOUD) - - response = app.test_client().post("/console/api/init", json={"password": "expected"}) - - assert response.status_code == 404 - init_validation.validate_password.assert_not_called() diff --git a/api/tests/unit_tests/controllers/console/test_fastopenapi_system.py b/api/tests/unit_tests/controllers/console/test_fastopenapi_system.py deleted file mode 100644 index e82eca70956..00000000000 --- a/api/tests/unit_tests/controllers/console/test_fastopenapi_system.py +++ /dev/null @@ -1,41 +0,0 @@ -import builtins - -import pytest -from flask.views import MethodView - -from configs import dify_config -from dify_app import DifyApp -from extensions import ext_fastopenapi -from tests.unit_tests.config_override import config_overrides_context - -if not hasattr(builtins, "MethodView"): - builtins.MethodView = MethodView # type: ignore[attr-defined] - - -@pytest.fixture -def app() -> DifyApp: - app = DifyApp(__name__) - app.config["TESTING"] = True - return app - - -def test_console_ping_fastopenapi_returns_pong(app: DifyApp) -> None: - ext_fastopenapi.init_app(app) - - response = app.test_client().get("/console/api/ping") - - assert response.status_code == 200 - assert response.get_json() == {"result": "pong"} - - -def test_console_version_fastopenapi_returns_current_version(app: DifyApp) -> None: - ext_fastopenapi.init_app(app) - - with config_overrides_context(CHECK_UPDATE_URL=None): - response = app.test_client().get("/console/api/version", query_string={"current_version": "0.0.0"}) - - assert response.status_code == 200 - assert response.get_json() == { - "version": dify_config.project.version, - "release_notes": "", - } diff --git a/api/tests/unit_tests/controllers/console/test_feature.py b/api/tests/unit_tests/controllers/console/test_feature.py index df58c5efec7..6d9327721b7 100644 --- a/api/tests/unit_tests/controllers/console/test_feature.py +++ b/api/tests/unit_tests/controllers/console/test_feature.py @@ -1,11 +1,22 @@ from inspect import unwrap from types import SimpleNamespace +from typing import Never from unittest.mock import MagicMock, create_autospec +import pytest +from flask import Flask +from flask_restx import Api +from pydantic import ValidationError from pytest_mock import MockerFixture +from controllers.console import feature, wraps +from controllers.web.feature import SystemFeatureApi as WebSystemFeatureApi from enums import DeploymentEdition +from extensions.ext_application_services import ApplicationServices +from extensions.ext_login import DifyLoginManager +from libs import login from machinery.context import RequestContext +from models.account import AccountStatus from services.entities.feature_entities import ( FeatureModel, LicenseLimitationModel, @@ -16,6 +27,9 @@ from services.entities.feature_entities import ( VectorSpaceLimitationModel, ) from services.feature_query_service import FeatureQueryService +from services.feature_service import FeatureService +from tests.helpers.unit.config_override import apply_config_overrides +from tests.helpers.unit.model_factories import make_account, make_tenant def _request_context() -> RequestContext: @@ -188,3 +202,106 @@ class TestSystemFeatureLicenseApi: assert result == license_model.model_dump() assert result["seats"] == {"enabled": True, "limit": 5, "size": 2} get_license.assert_called_once_with() + + +"""Public deployment discovery is independent of workspace tracing configuration.""" + + +@pytest.mark.parametrize("enabled", [False, True]) +@pytest.mark.parametrize("surface", ["console", "web"]) +def test_public_system_features_exposes_gate_without_account_app_or_provider( + monkeypatch: pytest.MonkeyPatch, + account_application_services: ApplicationServices, + enabled: bool, + surface: str, +) -> None: + apply_config_overrides( + monkeypatch, + DEPLOYMENT_EDITION=DeploymentEdition.COMMUNITY, + OPS_AGENT_TRACE_ENABLED=enabled, + LOGIN_DISABLED=False, + ) + + def forbidden_account() -> Never: + pytest.fail("public deployment discovery must not resolve a Console account") + + def forbidden_workspace(_workspace_id: str, **_kwargs: object) -> Never: + pytest.fail("public deployment discovery must not query workspace features") + + monkeypatch.setattr(login, "_resolve_current_user", forbidden_account) + monkeypatch.setattr(FeatureService, "get_features", forbidden_workspace) + app = Flask(__name__) + app.extensions["application_services"] = account_application_services + resource = feature.SystemFeatureApi if surface == "console" else WebSystemFeatureApi + path = "/console/api/system-features" if surface == "console" else "/api/system-features" + Api(app).add_resource(resource, path) + + response = app.test_client().get(path) + + assert response.status_code == 200 + payload = response.get_json() + assert payload["agent_tracing_api_collection_enabled"] is enabled + assert "tracing_provider" not in payload + assert "runtime_enabled" not in payload + assert "queryable" not in payload + schema = SystemFeatureModel.model_json_schema(mode="serialization") + gate = schema["properties"]["agent_tracing_api_collection_enabled"] + assert gate["type"] == "boolean" + assert "default" not in gate + assert "agent_tracing_api_collection_enabled" in schema["required"] + + +@pytest.mark.parametrize("enabled", [False, True]) +def test_authenticated_workspace_features_omits_deployment_gate( + monkeypatch: pytest.MonkeyPatch, + account_application_services: ApplicationServices, + enabled: bool, +) -> None: + apply_config_overrides( + monkeypatch, + DEPLOYMENT_EDITION=DeploymentEdition.COMMUNITY, + OPS_AGENT_TRACE_ENABLED=enabled, + LOGIN_DISABLED=False, + ) + account = make_account(status=AccountStatus.ACTIVE) + account._current_tenant = make_tenant() + monkeypatch.setattr(login, "_resolve_current_user", lambda: account) + monkeypatch.setattr(login, "check_csrf_token", lambda _request, _account_id: None) + monkeypatch.setattr(wraps, "_is_setup_completed", lambda: True) + monkeypatch.setattr(feature, "application_services", lambda: account_application_services) + app = Flask(__name__) + Api(app).add_resource(feature.FeatureApi, "/console/api/features") + + response = app.test_client().get("/console/api/features") + + assert response.status_code == 200 + payload = response.get_json() + assert "agent_tracing_api_collection_enabled" not in payload + schema = FeatureModel.model_json_schema(mode="serialization") + assert "agent_tracing_api_collection_enabled" not in schema["properties"] + + +def test_workspace_features_still_require_authentication(monkeypatch: pytest.MonkeyPatch) -> None: + apply_config_overrides(monkeypatch, LOGIN_DISABLED=False, OPS_AGENT_TRACE_ENABLED=True) + monkeypatch.setattr(wraps, "_is_setup_completed", lambda: True) + monkeypatch.setattr(login, "_resolve_current_user", lambda: None) + + def forbidden() -> ApplicationServices: + pytest.fail("anonymous requests must not reach feature queries") + + monkeypatch.setattr(feature, "application_services", forbidden) + app = Flask(__name__) + DifyLoginManager(app) + Api(app).add_resource(feature.FeatureApi, "/console/api/features") + + response = app.test_client().get("/console/api/features") + + assert response.status_code == 401 + assert b"agent_tracing_api_collection_enabled" not in response.data + + +def test_missing_public_gate_cannot_default_to_disabled() -> None: + with pytest.raises(ValidationError) as error: + SystemFeatureModel.model_validate({"deployment_edition": DeploymentEdition.COMMUNITY}) + + assert error.value.errors(include_input=False)[0]["loc"] == ("agent_tracing_api_collection_enabled",) diff --git a/api/tests/unit_tests/controllers/console/test_feature_agent_tracing.py b/api/tests/unit_tests/controllers/console/test_feature_agent_tracing.py deleted file mode 100644 index 3a65d7b8511..00000000000 --- a/api/tests/unit_tests/controllers/console/test_feature_agent_tracing.py +++ /dev/null @@ -1,120 +0,0 @@ -"""Public deployment discovery is independent of workspace tracing configuration.""" - -from typing import Never - -import pytest -from flask import Flask -from flask_restx import Api -from pydantic import ValidationError - -from controllers.console import feature, wraps -from controllers.web.feature import SystemFeatureApi as WebSystemFeatureApi -from enums import DeploymentEdition -from extensions.ext_application_services import ApplicationServices -from extensions.ext_login import DifyLoginManager -from libs import login -from models.account import AccountStatus -from services.entities.feature_entities import FeatureModel, SystemFeatureModel -from services.feature_service import FeatureService -from tests.unit_tests.config_override import apply_config_overrides -from tests.unit_tests.model_factories import make_account, make_tenant - - -@pytest.mark.parametrize("enabled", [False, True]) -@pytest.mark.parametrize("surface", ["console", "web"]) -def test_public_system_features_exposes_gate_without_account_app_or_provider( - monkeypatch: pytest.MonkeyPatch, - account_application_services: ApplicationServices, - enabled: bool, - surface: str, -) -> None: - apply_config_overrides( - monkeypatch, - DEPLOYMENT_EDITION=DeploymentEdition.COMMUNITY, - OPS_AGENT_TRACE_ENABLED=enabled, - LOGIN_DISABLED=False, - ) - - def forbidden_account() -> Never: - pytest.fail("public deployment discovery must not resolve a Console account") - - def forbidden_workspace(_workspace_id: str, **_kwargs: object) -> Never: - pytest.fail("public deployment discovery must not query workspace features") - - monkeypatch.setattr(login, "_resolve_current_user", forbidden_account) - monkeypatch.setattr(FeatureService, "get_features", forbidden_workspace) - app = Flask(__name__) - app.extensions["application_services"] = account_application_services - resource = feature.SystemFeatureApi if surface == "console" else WebSystemFeatureApi - path = "/console/api/system-features" if surface == "console" else "/api/system-features" - Api(app).add_resource(resource, path) - - response = app.test_client().get(path) - - assert response.status_code == 200 - payload = response.get_json() - assert payload["agent_tracing_api_collection_enabled"] is enabled - assert "tracing_provider" not in payload - assert "runtime_enabled" not in payload - assert "queryable" not in payload - schema = SystemFeatureModel.model_json_schema(mode="serialization") - gate = schema["properties"]["agent_tracing_api_collection_enabled"] - assert gate["type"] == "boolean" - assert "default" not in gate - assert "agent_tracing_api_collection_enabled" in schema["required"] - - -@pytest.mark.parametrize("enabled", [False, True]) -def test_authenticated_workspace_features_omits_deployment_gate( - monkeypatch: pytest.MonkeyPatch, - account_application_services: ApplicationServices, - enabled: bool, -) -> None: - apply_config_overrides( - monkeypatch, - DEPLOYMENT_EDITION=DeploymentEdition.COMMUNITY, - OPS_AGENT_TRACE_ENABLED=enabled, - LOGIN_DISABLED=False, - ) - account = make_account(status=AccountStatus.ACTIVE) - account._current_tenant = make_tenant() - monkeypatch.setattr(login, "_resolve_current_user", lambda: account) - monkeypatch.setattr(login, "check_csrf_token", lambda _request, _account_id: None) - monkeypatch.setattr(wraps, "_is_setup_completed", lambda: True) - monkeypatch.setattr(feature, "application_services", lambda: account_application_services) - app = Flask(__name__) - Api(app).add_resource(feature.FeatureApi, "/console/api/features") - - response = app.test_client().get("/console/api/features") - - assert response.status_code == 200 - payload = response.get_json() - assert "agent_tracing_api_collection_enabled" not in payload - schema = FeatureModel.model_json_schema(mode="serialization") - assert "agent_tracing_api_collection_enabled" not in schema["properties"] - - -def test_workspace_features_still_require_authentication(monkeypatch: pytest.MonkeyPatch) -> None: - apply_config_overrides(monkeypatch, LOGIN_DISABLED=False, OPS_AGENT_TRACE_ENABLED=True) - monkeypatch.setattr(wraps, "_is_setup_completed", lambda: True) - monkeypatch.setattr(login, "_resolve_current_user", lambda: None) - - def forbidden() -> ApplicationServices: - pytest.fail("anonymous requests must not reach feature queries") - - monkeypatch.setattr(feature, "application_services", forbidden) - app = Flask(__name__) - DifyLoginManager(app) - Api(app).add_resource(feature.FeatureApi, "/console/api/features") - - response = app.test_client().get("/console/api/features") - - assert response.status_code == 401 - assert b"agent_tracing_api_collection_enabled" not in response.data - - -def test_missing_public_gate_cannot_default_to_disabled() -> None: - with pytest.raises(ValidationError) as error: - SystemFeatureModel.model_validate({"deployment_edition": DeploymentEdition.COMMUNITY}) - - assert error.value.errors(include_input=False)[0]["loc"] == ("agent_tracing_api_collection_enabled",) diff --git a/api/tests/unit_tests/controllers/console/test_files.py b/api/tests/unit_tests/controllers/console/test_files.py index bf2b3da2671..e06f1d781fb 100644 --- a/api/tests/unit_tests/controllers/console/test_files.py +++ b/api/tests/unit_tests/controllers/console/test_files.py @@ -1,8 +1,10 @@ +import builtins import io from unittest.mock import patch import pytest from flask import Flask +from flask.views import MethodView from werkzeug.exceptions import Forbidden from configs import dify_config @@ -25,7 +27,7 @@ from machinery.context import RequestContext from models import Account from models.account import AccountStatus, TenantAccountRole from models.model import UploadFile -from tests.unit_tests.model_factories import make_upload_file +from tests.helpers.unit.model_factories import make_upload_file def unwrap(func): @@ -56,290 +58,523 @@ def _request_context(*, workspace_id: str = "tenant-1") -> RequestContext: ) -@pytest.fixture -def app(): - app = Flask(__name__) - app.testing = True - return app - - -@pytest.fixture(autouse=True) -def mock_decorators(): - """ - Make decorators no-ops so logic is directly testable - """ - with ( - patch("controllers.console.files.setup_required", new=lambda f: f), - patch("controllers.console.files.login_required", new=lambda f: f), - patch("controllers.console.files.account_initialization_required", new=lambda f: f), - patch("controllers.console.files.cloud_edition_billing_resource_check", return_value=lambda f: f), - ): - yield - - -@pytest.fixture -def mock_current_user(): - user = Account(name="Test User", email="user-1@example.com", status=AccountStatus.ACTIVE) - user.id = "user-1" - user.role = TenantAccountRole.OWNER - return user - - -@pytest.fixture -def mock_account_context(mock_current_user): - return mock_current_user - - -@pytest.fixture -def mock_file_service(): - with patch("controllers.console.files.application_services") as services: - yield services.return_value.files - - -class TestFileApiGet: - def test_get_upload_config(self, app: Flask): - api = FileApi() - get_method = unwrap(api.get) +class TestFilesCases: + @staticmethod + @pytest.fixture + def app(): + app = Flask(__name__) + app.testing = True + return app + @staticmethod + @pytest.fixture(autouse=True) + def mock_decorators(): + """ + Make decorators no-ops so logic is directly testable + """ with ( - app.test_request_context(), - patch( - "controllers.console.files.FeatureService.get_knowledge_file_size_limit", - return_value=50, - ) as get_knowledge_file_size_limit, + patch("controllers.console.files.setup_required", new=lambda f: f), + patch("controllers.console.files.login_required", new=lambda f: f), + patch("controllers.console.files.account_initialization_required", new=lambda f: f), + patch("controllers.console.files.cloud_edition_billing_resource_check", return_value=lambda f: f), ): - data, status = get_method(api, _request_context()) + yield - assert status == 200 - assert "file_size_limit" in data - assert data["knowledge_file_size_limit"] == 50 - assert "batch_count_limit" in data - get_knowledge_file_size_limit.assert_called_once_with("tenant-1") - assert data["skill_file_size_limit"] == dify_config.UPLOAD_SKILL_FILE_SIZE_LIMIT + @staticmethod + @pytest.fixture + def mock_current_user(): + user = Account(name="Test User", email="user-1@example.com", status=AccountStatus.ACTIVE) + user.id = "user-1" + user.role = TenantAccountRole.OWNER + return user + @staticmethod + @pytest.fixture + def mock_account_context(mock_current_user): + return mock_current_user -class TestFileApiPost: - def test_no_file_uploaded(self, app: Flask, mock_account_context): - api = FileApi() - post_method = unwrap(api.post) + @staticmethod + @pytest.fixture + def mock_file_service(): + with patch("controllers.console.files.application_services") as services: + yield services.return_value.files - with app.test_request_context(method="POST", data={}): - with pytest.raises(NoFileUploadedError): - post_method(api, mock_account_context) + class TestFileApiGet: + def test_get_upload_config(self, app: Flask): + api = FileApi() + get_method = unwrap(api.get) - def test_too_many_files(self, app: Flask, mock_account_context): - api = FileApi() - post_method = unwrap(api.post) + with ( + app.test_request_context(), + patch( + "controllers.console.files.FeatureService.get_knowledge_file_size_limit", + return_value=50, + ) as get_knowledge_file_size_limit, + ): + data, status = get_method(api, _request_context()) - with app.test_request_context(method="POST"): - from unittest.mock import MagicMock, patch + assert status == 200 + assert "file_size_limit" in data + assert data["knowledge_file_size_limit"] == 50 + assert "batch_count_limit" in data + get_knowledge_file_size_limit.assert_called_once_with("tenant-1") + assert data["skill_file_size_limit"] == dify_config.UPLOAD_SKILL_FILE_SIZE_LIMIT - with patch("controllers.console.files.request") as mock_request: - mock_request.files = MagicMock() - mock_request.files.__len__.return_value = 2 - mock_request.files.__contains__.return_value = True - mock_request.form = MagicMock() - mock_request.form.get.return_value = None + class TestFileApiPost: + def test_no_file_uploaded(self, app: Flask, mock_account_context): + api = FileApi() + post_method = unwrap(api.post) - with pytest.raises(TooManyFilesError): + with app.test_request_context(method="POST", data={}): + with pytest.raises(NoFileUploadedError): post_method(api, mock_account_context) - def test_filename_missing(self, app: Flask, mock_account_context): - api = FileApi() - post_method = unwrap(api.post) + def test_too_many_files(self, app: Flask, mock_account_context): + api = FileApi() + post_method = unwrap(api.post) - data = { - "file": (io.BytesIO(b"abc"), ""), - } + with app.test_request_context(method="POST"): + from unittest.mock import MagicMock, patch - with app.test_request_context(method="POST", data=data): - with pytest.raises(FilenameNotExistsError): - post_method(api, mock_account_context) + with patch("controllers.console.files.request") as mock_request: + mock_request.files = MagicMock() + mock_request.files.__len__.return_value = 2 + mock_request.files.__contains__.return_value = True + mock_request.form = MagicMock() + mock_request.form.get.return_value = None - def test_dataset_upload_without_permission(self, app: Flask, mock_current_user): - mock_current_user.role = TenantAccountRole.NORMAL + with pytest.raises(TooManyFilesError): + post_method(api, mock_account_context) - api = FileApi() - post_method = unwrap(api.post) + def test_filename_missing(self, app: Flask, mock_account_context): + api = FileApi() + post_method = unwrap(api.post) - data = { - "file": (io.BytesIO(b"abc"), "test.txt"), - "source": "datasets", - } + data = { + "file": (io.BytesIO(b"abc"), ""), + } - with app.test_request_context(method="POST", data=data): - with pytest.raises(Forbidden): - post_method(api, mock_current_user) + with app.test_request_context(method="POST", data=data): + with pytest.raises(FilenameNotExistsError): + post_method(api, mock_account_context) - def test_successful_upload(self, app: Flask, mock_account_context, mock_file_service): - api = FileApi() - post_method = unwrap(api.post) + def test_dataset_upload_without_permission(self, app: Flask, mock_current_user): + mock_current_user.role = TenantAccountRole.NORMAL - mock_file_service.upload_file.return_value = _upload_file() + api = FileApi() + post_method = unwrap(api.post) - data = { - "file": (io.BytesIO(b"hello"), "test.txt"), - } + data = { + "file": (io.BytesIO(b"abc"), "test.txt"), + "source": "datasets", + } - with app.test_request_context(method="POST", data=data): - response, status = post_method(api, mock_account_context) + with app.test_request_context(method="POST", data=data): + with pytest.raises(Forbidden): + post_method(api, mock_current_user) - assert status == 201 - assert response["id"] == "file-id-123" - assert response["name"] == "test.txt" - mock_file_service.upload_file.assert_called_once_with( - filename="test.txt", - content=b"hello", - mimetype="text/plain", - user=mock_account_context, - tenant_id=None, - source=None, - default_file_size_limit=None, - ) + def test_successful_upload(self, app: Flask, mock_account_context, mock_file_service): + api = FileApi() + post_method = unwrap(api.post) - def test_upload_with_resource_tenant(self, app: Flask, mock_account_context, mock_file_service): - upload_file = _upload_file() - mock_file_service.upload_file.return_value = upload_file + mock_file_service.upload_file.return_value = _upload_file() - with app.test_request_context( - method="POST", - data={"file": (io.BytesIO(b"hello"), "test.txt")}, - ): - result = upload_file_from_request( - current_user=mock_account_context, - resource_tenant_id="app-tenant-id", + data = { + "file": (io.BytesIO(b"hello"), "test.txt"), + } + + with app.test_request_context(method="POST", data=data): + response, status = post_method(api, mock_account_context) + + assert status == 201 + assert response["id"] == "file-id-123" + assert response["name"] == "test.txt" + mock_file_service.upload_file.assert_called_once_with( + filename="test.txt", + content=b"hello", + mimetype="text/plain", + user=mock_account_context, + tenant_id=None, + source=None, + default_file_size_limit=None, ) - assert result is upload_file - assert mock_file_service.upload_file.call_args.kwargs["tenant_id"] == "app-tenant-id" + def test_upload_with_resource_tenant(self, app: Flask, mock_account_context, mock_file_service): + upload_file = _upload_file() + mock_file_service.upload_file.return_value = upload_file - def test_dataset_source_from_query_uses_knowledge_limit( - self, - app: Flask, - mock_account_context, - mock_file_service, - ): - upload_file = _upload_file() - mock_file_service.upload_file.return_value = upload_file - - with ( - app.test_request_context( - "/?source=datasets", + with app.test_request_context( method="POST", data={"file": (io.BytesIO(b"hello"), "test.txt")}, - ), - patch( - "controllers.console.files.FeatureService.get_knowledge_file_size_limit", - return_value=50, - ) as get_knowledge_file_size_limit, + ): + result = upload_file_from_request( + current_user=mock_account_context, + resource_tenant_id="app-tenant-id", + ) + + assert result is upload_file + assert mock_file_service.upload_file.call_args.kwargs["tenant_id"] == "app-tenant-id" + + def test_dataset_source_from_query_uses_knowledge_limit( + self, + app: Flask, + mock_account_context, + mock_file_service, ): - result = upload_file_from_request(current_user=mock_account_context) + upload_file = _upload_file() + mock_file_service.upload_file.return_value = upload_file - assert result is upload_file - assert mock_file_service.upload_file.call_args.kwargs["source"] == "datasets" - assert mock_file_service.upload_file.call_args.kwargs["default_file_size_limit"] == 50 - get_knowledge_file_size_limit.assert_called_once_with(mock_account_context.current_tenant_id) + with ( + app.test_request_context( + "/?source=datasets", + method="POST", + data={"file": (io.BytesIO(b"hello"), "test.txt")}, + ), + patch( + "controllers.console.files.FeatureService.get_knowledge_file_size_limit", + return_value=50, + ) as get_knowledge_file_size_limit, + ): + result = upload_file_from_request(current_user=mock_account_context) - def test_upload_with_invalid_source(self, app: Flask, mock_account_context, mock_file_service): - """Test that invalid source parameter gets normalized to None""" - api = FileApi() - post_method = unwrap(api.post) + assert result is upload_file + assert mock_file_service.upload_file.call_args.kwargs["source"] == "datasets" + assert mock_file_service.upload_file.call_args.kwargs["default_file_size_limit"] == 50 + get_knowledge_file_size_limit.assert_called_once_with(mock_account_context.current_tenant_id) - mock_file_service.upload_file.return_value = _upload_file(file_id="file-id-456", size=512) + def test_upload_with_invalid_source(self, app: Flask, mock_account_context, mock_file_service): + """Test that invalid source parameter gets normalized to None""" + api = FileApi() + post_method = unwrap(api.post) - data = { - "file": (io.BytesIO(b"content"), "test.txt"), - "source": "invalid_source", # Should be normalized to None - } + mock_file_service.upload_file.return_value = _upload_file(file_id="file-id-456", size=512) - with app.test_request_context(method="POST", data=data): - response, status = post_method(api, mock_account_context) + data = { + "file": (io.BytesIO(b"content"), "test.txt"), + "source": "invalid_source", # Should be normalized to None + } - assert status == 201 - assert response["id"] == "file-id-456" - # Verify that FileService was called with source=None - mock_file_service.upload_file.assert_called_once() - call_kwargs = mock_file_service.upload_file.call_args[1] - assert call_kwargs["source"] is None + with app.test_request_context(method="POST", data=data): + response, status = post_method(api, mock_account_context) - def test_file_too_large_error(self, app: Flask, mock_account_context, mock_file_service): - api = FileApi() - post_method = unwrap(api.post) + assert status == 201 + assert response["id"] == "file-id-456" + # Verify that FileService was called with source=None + mock_file_service.upload_file.assert_called_once() + call_kwargs = mock_file_service.upload_file.call_args[1] + assert call_kwargs["source"] is None - from services.errors.file import FileTooLargeError as ServiceFileTooLargeError + def test_file_too_large_error(self, app: Flask, mock_account_context, mock_file_service): + api = FileApi() + post_method = unwrap(api.post) - error = ServiceFileTooLargeError("File is too large") - mock_file_service.upload_file.side_effect = error + from services.errors.file import FileTooLargeError as ServiceFileTooLargeError - data = { - "file": (io.BytesIO(b"x" * 1000000), "big.txt"), - } + error = ServiceFileTooLargeError("File is too large") + mock_file_service.upload_file.side_effect = error - with app.test_request_context(method="POST", data=data): - with pytest.raises(FileTooLargeError) as error_info: - post_method(api, mock_account_context) + data = { + "file": (io.BytesIO(b"x" * 1000000), "big.txt"), + } - assert error_info.value.__cause__ is error + with app.test_request_context(method="POST", data=data): + with pytest.raises(FileTooLargeError) as error_info: + post_method(api, mock_account_context) - def test_unsupported_file_type(self, app: Flask, mock_account_context, mock_file_service): - api = FileApi() - post_method = unwrap(api.post) + assert error_info.value.__cause__ is error - from services.errors.file import UnsupportedFileTypeError as ServiceUnsupportedFileTypeError + def test_unsupported_file_type(self, app: Flask, mock_account_context, mock_file_service): + api = FileApi() + post_method = unwrap(api.post) - error = ServiceUnsupportedFileTypeError() - mock_file_service.upload_file.side_effect = error + from services.errors.file import UnsupportedFileTypeError as ServiceUnsupportedFileTypeError - data = { - "file": (io.BytesIO(b"x"), "bad.exe"), - } + error = ServiceUnsupportedFileTypeError() + mock_file_service.upload_file.side_effect = error - with app.test_request_context(method="POST", data=data): - with pytest.raises(UnsupportedFileTypeError) as error_info: - post_method(api, mock_account_context) + data = { + "file": (io.BytesIO(b"x"), "bad.exe"), + } - assert error_info.value.__cause__ is error + with app.test_request_context(method="POST", data=data): + with pytest.raises(UnsupportedFileTypeError) as error_info: + post_method(api, mock_account_context) - def test_blocked_extension(self, app: Flask, mock_account_context, mock_file_service): - api = FileApi() - post_method = unwrap(api.post) + assert error_info.value.__cause__ is error - from services.errors.file import BlockedFileExtensionError as ServiceBlockedFileExtensionError + def test_blocked_extension(self, app: Flask, mock_account_context, mock_file_service): + api = FileApi() + post_method = unwrap(api.post) - error = ServiceBlockedFileExtensionError("File extension is blocked") - mock_file_service.upload_file.side_effect = error + from services.errors.file import BlockedFileExtensionError as ServiceBlockedFileExtensionError - data = { - "file": (io.BytesIO(b"x"), "blocked.txt"), - } + error = ServiceBlockedFileExtensionError("File extension is blocked") + mock_file_service.upload_file.side_effect = error - with app.test_request_context(method="POST", data=data): - with pytest.raises(BlockedFileExtensionError) as error_info: - post_method(api, mock_account_context) + data = { + "file": (io.BytesIO(b"x"), "blocked.txt"), + } - assert error_info.value.description == error.description - assert error_info.value.__cause__ is error + with app.test_request_context(method="POST", data=data): + with pytest.raises(BlockedFileExtensionError) as error_info: + post_method(api, mock_account_context) + + assert error_info.value.description == error.description + assert error_info.value.__cause__ is error + + class TestFilePreviewApi: + def test_get_preview(self, app: Flask, mock_account_context, mock_file_service): + api = FilePreviewApi() + get_method = unwrap(api.get) + mock_file_service.get_file_preview.return_value = "preview text" + + with app.test_request_context(): + result = get_method(api, _request_context(workspace_id="tenant-123"), "1234") + + assert result == {"content": "preview text"} + mock_file_service.get_file_preview.assert_called_once_with(file_id="1234", tenant_id="tenant-123") + + class TestFileSupportTypeApi: + def test_get_supported_types(self, app: Flask): + api = FileSupportTypeApi() + get_method = unwrap(api.get) + + with app.test_request_context(): + result = get_method(api, _request_context()) + + assert result == {"allowed_extensions": list(DOCUMENT_EXTENSIONS)} -class TestFilePreviewApi: - def test_get_preview(self, app: Flask, mock_account_context, mock_file_service): - api = FilePreviewApi() - get_method = unwrap(api.get) - mock_file_service.get_file_preview.return_value = "preview text" - - with app.test_request_context(): - result = get_method(api, _request_context(workspace_id="tenant-123"), "1234") - - assert result == {"content": "preview text"} - mock_file_service.get_file_preview.assert_called_once_with(file_id="1234", tenant_id="tenant-123") +if not hasattr(builtins, "MethodView"): + builtins.MethodView = MethodView # type: ignore[attr-defined] -class TestFileSupportTypeApi: - def test_get_supported_types(self, app: Flask): - api = FileSupportTypeApi() - get_method = unwrap(api.get) +class TestFilesSecurityCases: + class TestFileUploadSecurity: + """Test file upload security logic without complex framework setup""" - with app.test_request_context(): - result = get_method(api, _request_context()) + # Test 1: Basic file validation + def test_should_validate_file_presence(self) -> None: + """Test that missing file is detected""" + from flask import Flask, request - assert result == {"allowed_extensions": list(DOCUMENT_EXTENSIONS)} + app = Flask(__name__) + + with app.test_request_context(method="POST", data={}): + # Simulate the check in FileApi.post() + if "file" not in request.files: + with pytest.raises(NoFileUploadedError): + raise NoFileUploadedError() + + def test_should_validate_multiple_files(self) -> None: + """Test that multiple files are rejected""" + from flask import Flask, request + + app = Flask(__name__) + + file_data = { + "file": (io.BytesIO(b"content1"), "file1.txt", "text/plain"), + "file2": (io.BytesIO(b"content2"), "file2.txt", "text/plain"), + } + + with app.test_request_context(method="POST", data=file_data, content_type="multipart/form-data"): + # Simulate the check in FileApi.post() + if len(request.files) > 1: + with pytest.raises(TooManyFilesError): + raise TooManyFilesError() + + def test_should_validate_empty_filename(self) -> None: + """Test that empty filename is rejected""" + from flask import Flask, request + + app = Flask(__name__) + + file_data = {"file": (io.BytesIO(b"content"), "", "text/plain")} + + with app.test_request_context(method="POST", data=file_data, content_type="multipart/form-data"): + file = request.files["file"] + if not file.filename: + with pytest.raises(FilenameNotExistsError): + raise FilenameNotExistsError + + # Test 2: Security - Filename sanitization + def test_should_detect_path_traversal_in_filename(self) -> None: + """Test protection against directory traversal attacks""" + dangerous_filenames = [ + "../../../etc/passwd", + "..\\..\\windows\\system32\\config\\sam", + "../../../../etc/shadow", + "./../../../sensitive.txt", + ] + + for filename in dangerous_filenames: + # Any filename containing .. should be considered dangerous + assert ".." in filename, f"Filename {filename} should be detected as path traversal" + + def test_should_detect_null_byte_injection(self) -> None: + """Test protection against null byte injection""" + dangerous_filenames = [ + "file.jpg\x00.php", + "document.pdf\x00.exe", + "image.png\x00.sh", + ] + + for filename in dangerous_filenames: + # Null bytes should be detected + assert "\x00" in filename, f"Filename {filename} should be detected as null byte injection" + + def test_should_sanitize_special_characters(self) -> None: + """Test that special characters in filenames are handled safely""" + # Characters that could be problematic in various contexts + dangerous_chars = ["/", "\\", ":", "*", "?", '"', "<", ">", "|", "\x00"] + + for char in dangerous_chars: + filename = f"file{char}name.txt" + # These characters should be detected or sanitized + assert any(c in filename for c in dangerous_chars) + + # Test 3: Permission validation + def test_should_validate_dataset_permissions(self) -> None: + """Test dataset upload permission logic""" + user = Account(name="Viewer", email="viewer@example.com", status=AccountStatus.ACTIVE) + user.role = TenantAccountRole.NORMAL + source = "datasets" + + # Simulate the permission check in FileApi.post() + if source == "datasets" and not user.is_dataset_editor: + with pytest.raises(Forbidden): + raise Forbidden() + + def test_should_allow_general_upload_without_permission(self) -> None: + """Test general upload doesn't require dataset permission""" + user = Account(name="Viewer", email="viewer@example.com", status=AccountStatus.ACTIVE) + user.role = TenantAccountRole.NORMAL + source = None # General upload + + # This should not raise an exception + if source == "datasets" and not user.is_dataset_editor: + raise Forbidden() + # Test passes if no exception is raised + + # Test 4: File type security + def test_should_identify_dangerous_file_extensions(self) -> None: + """Test detection of potentially dangerous file extensions""" + dangerous_extensions = [ + ".php", + ".PHP", + ".pHp", # PHP files (case variations) + ".exe", + ".EXE", # Executables + ".sh", + ".SH", # Shell scripts + ".bat", + ".BAT", # Batch files + ".cmd", + ".CMD", # Command files + ".ps1", + ".PS1", # PowerShell + ".jar", + ".JAR", # Java archives + ".vbs", + ".VBS", # VBScript + ] + + safe_extensions = [".txt", ".pdf", ".jpg", ".png", ".doc", ".docx"] + + # Just verify our test data is correct + for ext in dangerous_extensions: + assert ext.lower() in [".php", ".exe", ".sh", ".bat", ".cmd", ".ps1", ".jar", ".vbs"] + + for ext in safe_extensions: + assert ext.lower() not in [".php", ".exe", ".sh", ".bat", ".cmd", ".ps1", ".jar", ".vbs"] + + def test_should_detect_double_extensions(self) -> None: + """Test detection of double extension attacks""" + suspicious_filenames = [ + "image.jpg.php", + "document.pdf.exe", + "photo.png.sh", + "file.txt.bat", + ] + + for filename in suspicious_filenames: + # Check that these have multiple extensions + parts = filename.split(".") + assert len(parts) > 2, f"Filename {filename} should have multiple extensions" + + # Test 5: Configuration validation + def test_upload_configuration_structure(self) -> None: + """Test that upload configuration has correct structure""" + # Simulate the configuration returned by FileApi.get() + config = { + "file_size_limit": 15, + "batch_count_limit": 5, + "image_file_size_limit": 10, + "video_file_size_limit": 500, + "audio_file_size_limit": 50, + "workflow_file_upload_limit": 10, + } + + # Verify all required fields are present + required_fields = [ + "file_size_limit", + "batch_count_limit", + "image_file_size_limit", + "video_file_size_limit", + "audio_file_size_limit", + "workflow_file_upload_limit", + ] + + for field in required_fields: + assert field in config, f"Missing required field: {field}" + assert isinstance(config[field], int), f"Field {field} should be an integer" + assert config[field] > 0, f"Field {field} should be positive" + + # Test 7: Source parameter handling + def test_source_parameter_normalization(self) -> None: + """Test that source parameter is properly normalized""" + test_cases = [ + ("datasets", "datasets"), + ("other", None), + ("", None), + (None, None), + ] + + for input_source, expected in test_cases: + # Simulate the source normalization in FileApi.post() + source = "datasets" if input_source == "datasets" else None + if source not in ("datasets", None): + source = None + assert source == expected + + # Test 8: Boundary conditions + def test_should_handle_edge_case_file_sizes(self) -> None: + """Test handling of boundary file sizes""" + test_cases = [ + (0, "Empty file"), # 0 bytes + (1, "Single byte"), # 1 byte + (15 * 1024 * 1024 - 1, "Just under limit"), # Just under 15MB + (15 * 1024 * 1024, "At limit"), # Exactly 15MB + (15 * 1024 * 1024 + 1, "Just over limit"), # Just over 15MB + ] + + for size, description in test_cases: + # Just verify our test data + assert isinstance(size, int), f"{description}: Size should be integer" + assert size >= 0, f"{description}: Size should be non-negative" + + def test_should_handle_special_mime_types(self) -> None: + """Test handling of various MIME types""" + mime_type_tests = [ + ("application/octet-stream", "Generic binary"), + ("text/plain", "Plain text"), + ("image/jpeg", "JPEG image"), + ("application/pdf", "PDF document"), + ("", "Empty MIME type"), + (None, "None MIME type"), + ] + + for mime_type, description in mime_type_tests: + # Verify test data structure + if mime_type is not None: + assert isinstance(mime_type, str), f"{description}: MIME type should be string or None" diff --git a/api/tests/unit_tests/controllers/console/test_files_security.py b/api/tests/unit_tests/controllers/console/test_files_security.py deleted file mode 100644 index 3f120feb4ae..00000000000 --- a/api/tests/unit_tests/controllers/console/test_files_security.py +++ /dev/null @@ -1,247 +0,0 @@ -import builtins -import io - -import pytest -from flask.views import MethodView -from werkzeug.exceptions import Forbidden - -from controllers.common.errors import ( - FilenameNotExistsError, - NoFileUploadedError, - TooManyFilesError, -) -from models import Account -from models.account import AccountStatus, TenantAccountRole - -if not hasattr(builtins, "MethodView"): - builtins.MethodView = MethodView # type: ignore[attr-defined] - - -class TestFileUploadSecurity: - """Test file upload security logic without complex framework setup""" - - # Test 1: Basic file validation - def test_should_validate_file_presence(self) -> None: - """Test that missing file is detected""" - from flask import Flask, request - - app = Flask(__name__) - - with app.test_request_context(method="POST", data={}): - # Simulate the check in FileApi.post() - if "file" not in request.files: - with pytest.raises(NoFileUploadedError): - raise NoFileUploadedError() - - def test_should_validate_multiple_files(self) -> None: - """Test that multiple files are rejected""" - from flask import Flask, request - - app = Flask(__name__) - - file_data = { - "file": (io.BytesIO(b"content1"), "file1.txt", "text/plain"), - "file2": (io.BytesIO(b"content2"), "file2.txt", "text/plain"), - } - - with app.test_request_context(method="POST", data=file_data, content_type="multipart/form-data"): - # Simulate the check in FileApi.post() - if len(request.files) > 1: - with pytest.raises(TooManyFilesError): - raise TooManyFilesError() - - def test_should_validate_empty_filename(self) -> None: - """Test that empty filename is rejected""" - from flask import Flask, request - - app = Flask(__name__) - - file_data = {"file": (io.BytesIO(b"content"), "", "text/plain")} - - with app.test_request_context(method="POST", data=file_data, content_type="multipart/form-data"): - file = request.files["file"] - if not file.filename: - with pytest.raises(FilenameNotExistsError): - raise FilenameNotExistsError - - # Test 2: Security - Filename sanitization - def test_should_detect_path_traversal_in_filename(self) -> None: - """Test protection against directory traversal attacks""" - dangerous_filenames = [ - "../../../etc/passwd", - "..\\..\\windows\\system32\\config\\sam", - "../../../../etc/shadow", - "./../../../sensitive.txt", - ] - - for filename in dangerous_filenames: - # Any filename containing .. should be considered dangerous - assert ".." in filename, f"Filename {filename} should be detected as path traversal" - - def test_should_detect_null_byte_injection(self) -> None: - """Test protection against null byte injection""" - dangerous_filenames = [ - "file.jpg\x00.php", - "document.pdf\x00.exe", - "image.png\x00.sh", - ] - - for filename in dangerous_filenames: - # Null bytes should be detected - assert "\x00" in filename, f"Filename {filename} should be detected as null byte injection" - - def test_should_sanitize_special_characters(self) -> None: - """Test that special characters in filenames are handled safely""" - # Characters that could be problematic in various contexts - dangerous_chars = ["/", "\\", ":", "*", "?", '"', "<", ">", "|", "\x00"] - - for char in dangerous_chars: - filename = f"file{char}name.txt" - # These characters should be detected or sanitized - assert any(c in filename for c in dangerous_chars) - - # Test 3: Permission validation - def test_should_validate_dataset_permissions(self) -> None: - """Test dataset upload permission logic""" - user = Account(name="Viewer", email="viewer@example.com", status=AccountStatus.ACTIVE) - user.role = TenantAccountRole.NORMAL - source = "datasets" - - # Simulate the permission check in FileApi.post() - if source == "datasets" and not user.is_dataset_editor: - with pytest.raises(Forbidden): - raise Forbidden() - - def test_should_allow_general_upload_without_permission(self) -> None: - """Test general upload doesn't require dataset permission""" - user = Account(name="Viewer", email="viewer@example.com", status=AccountStatus.ACTIVE) - user.role = TenantAccountRole.NORMAL - source = None # General upload - - # This should not raise an exception - if source == "datasets" and not user.is_dataset_editor: - raise Forbidden() - # Test passes if no exception is raised - - # Test 4: File type security - def test_should_identify_dangerous_file_extensions(self) -> None: - """Test detection of potentially dangerous file extensions""" - dangerous_extensions = [ - ".php", - ".PHP", - ".pHp", # PHP files (case variations) - ".exe", - ".EXE", # Executables - ".sh", - ".SH", # Shell scripts - ".bat", - ".BAT", # Batch files - ".cmd", - ".CMD", # Command files - ".ps1", - ".PS1", # PowerShell - ".jar", - ".JAR", # Java archives - ".vbs", - ".VBS", # VBScript - ] - - safe_extensions = [".txt", ".pdf", ".jpg", ".png", ".doc", ".docx"] - - # Just verify our test data is correct - for ext in dangerous_extensions: - assert ext.lower() in [".php", ".exe", ".sh", ".bat", ".cmd", ".ps1", ".jar", ".vbs"] - - for ext in safe_extensions: - assert ext.lower() not in [".php", ".exe", ".sh", ".bat", ".cmd", ".ps1", ".jar", ".vbs"] - - def test_should_detect_double_extensions(self) -> None: - """Test detection of double extension attacks""" - suspicious_filenames = [ - "image.jpg.php", - "document.pdf.exe", - "photo.png.sh", - "file.txt.bat", - ] - - for filename in suspicious_filenames: - # Check that these have multiple extensions - parts = filename.split(".") - assert len(parts) > 2, f"Filename {filename} should have multiple extensions" - - # Test 5: Configuration validation - def test_upload_configuration_structure(self) -> None: - """Test that upload configuration has correct structure""" - # Simulate the configuration returned by FileApi.get() - config = { - "file_size_limit": 15, - "batch_count_limit": 5, - "image_file_size_limit": 10, - "video_file_size_limit": 500, - "audio_file_size_limit": 50, - "workflow_file_upload_limit": 10, - } - - # Verify all required fields are present - required_fields = [ - "file_size_limit", - "batch_count_limit", - "image_file_size_limit", - "video_file_size_limit", - "audio_file_size_limit", - "workflow_file_upload_limit", - ] - - for field in required_fields: - assert field in config, f"Missing required field: {field}" - assert isinstance(config[field], int), f"Field {field} should be an integer" - assert config[field] > 0, f"Field {field} should be positive" - - # Test 7: Source parameter handling - def test_source_parameter_normalization(self) -> None: - """Test that source parameter is properly normalized""" - test_cases = [ - ("datasets", "datasets"), - ("other", None), - ("", None), - (None, None), - ] - - for input_source, expected in test_cases: - # Simulate the source normalization in FileApi.post() - source = "datasets" if input_source == "datasets" else None - if source not in ("datasets", None): - source = None - assert source == expected - - # Test 8: Boundary conditions - def test_should_handle_edge_case_file_sizes(self) -> None: - """Test handling of boundary file sizes""" - test_cases = [ - (0, "Empty file"), # 0 bytes - (1, "Single byte"), # 1 byte - (15 * 1024 * 1024 - 1, "Just under limit"), # Just under 15MB - (15 * 1024 * 1024, "At limit"), # Exactly 15MB - (15 * 1024 * 1024 + 1, "Just over limit"), # Just over 15MB - ] - - for size, description in test_cases: - # Just verify our test data - assert isinstance(size, int), f"{description}: Size should be integer" - assert size >= 0, f"{description}: Size should be non-negative" - - def test_should_handle_special_mime_types(self) -> None: - """Test handling of various MIME types""" - mime_type_tests = [ - ("application/octet-stream", "Generic binary"), - ("text/plain", "Plain text"), - ("image/jpeg", "JPEG image"), - ("application/pdf", "PDF document"), - ("", "Empty MIME type"), - (None, "None MIME type"), - ] - - for mime_type, description in mime_type_tests: - # Verify test data structure - if mime_type is not None: - assert isinstance(mime_type, str), f"{description}: MIME type should be string or None" diff --git a/api/tests/unit_tests/controllers/console/test_human_input_form.py b/api/tests/unit_tests/controllers/console/test_human_input_form.py index daf8b04c002..d44c6c81585 100644 --- a/api/tests/unit_tests/controllers/console/test_human_input_form.py +++ b/api/tests/unit_tests/controllers/console/test_human_input_form.py @@ -35,9 +35,9 @@ from models.model import App, AppMode from models.workflow import WorkflowRun from services.enterprise.rbac_service import RBACService from services.entities.app_entities import AppAgentBinding -from tests.unit_tests.config_override import apply_config_overrides -from tests.unit_tests.model_factories import make_account, make_app -from tests.unit_tests.rbac_fakes import RBACDomain +from tests.helpers.unit.config_override import apply_config_overrides +from tests.helpers.unit.model_factories import make_account, make_app +from tests.helpers.unit.rbac_fakes import RBACDomain @pytest.fixture(autouse=True) diff --git a/api/tests/unit_tests/controllers/console/test_init_validate.py b/api/tests/unit_tests/controllers/console/test_init_validate.py index cfd6efd230b..1083708a5db 100644 --- a/api/tests/unit_tests/controllers/console/test_init_validate.py +++ b/api/tests/unit_tests/controllers/console/test_init_validate.py @@ -1,5 +1,3 @@ -"""Tests for the Flask adapter around initialization validation.""" - from types import SimpleNamespace from unittest.mock import Mock, create_autospec @@ -8,108 +6,237 @@ from flask import Flask from controllers.console import init_validate from controllers.console.error import AlreadySetupError, InitValidateFailedError +from dify_app import DifyApp from enums import DeploymentEdition +from extensions import ext_fastopenapi from services.init_validation_service import ( AlreadyInitializedError, InitValidationService, InvalidInitializationPasswordError, ) -from tests.unit_tests.config_override import apply_config_overrides +from tests.helpers.unit.config_override import apply_config_overrides + +"""HTTP contract tests for the FastOpenAPI initialization routes.""" -@pytest.fixture -def init_validation(monkeypatch: pytest.MonkeyPatch) -> Mock: - service = create_autospec(InitValidationService, instance=True, spec_set=True) - application_services = SimpleNamespace(init_validation=service) - monkeypatch.setattr(init_validate, "application_services", lambda: application_services) - return service +class TestFastopenapiInitValidateCases: + @staticmethod + @pytest.fixture + def init_validation(monkeypatch: pytest.MonkeyPatch) -> Mock: + service = create_autospec(InitValidationService, instance=True, spec_set=True) + services = SimpleNamespace(init_validation=service) + monkeypatch.setattr(init_validate, "application_services", lambda: services) + return service + @staticmethod + @pytest.fixture + def app() -> DifyApp: + app = DifyApp(__name__) + app.config.update(TESTING=True, SECRET_KEY="test-secret") + ext_fastopenapi.init_app(app) + return app -def test_get_init_status_finished(app: Flask, init_validation: Mock) -> None: - init_validation.is_validated.return_value = True - app.secret_key = "test-secret" + @staticmethod + @pytest.mark.parametrize( + ("validated", "expected_status"), + [ + pytest.param(True, "finished", id="finished"), + pytest.param(False, "not_started", id="not-started"), + ], + ) + def test_get_init_status( + app: DifyApp, + init_validation: Mock, + validated: bool, + expected_status: str, + ) -> None: + init_validation.is_validated.return_value = validated - with app.test_request_context("/console/api/init", method="GET"): - result = init_validate.get_init_status() + response = app.test_client().get("/console/api/init") - assert result.status == "finished" + assert response.status_code == 200 + assert response.get_json() == {"status": expected_status} + init_validation.is_validated.assert_called_once_with(session_validated=False) + @staticmethod + def test_validate_init_password_success( + app: DifyApp, + init_validation: Mock, + monkeypatch: pytest.MonkeyPatch, + ) -> None: + apply_config_overrides(monkeypatch, DEPLOYMENT_EDITION=DeploymentEdition.COMMUNITY) + client = app.test_client() -def test_get_init_status_not_started(app: Flask, init_validation: Mock) -> None: - init_validation.is_validated.return_value = False - app.secret_key = "test-secret" + response = client.post("/console/api/init", json={"password": "expected"}) - with app.test_request_context("/console/api/init", method="GET"): - result = init_validate.get_init_status() - - assert result.status == "not_started" - - -def test_validate_init_password_already_setup( - app: Flask, - init_validation: Mock, - monkeypatch: pytest.MonkeyPatch, -) -> None: - apply_config_overrides(monkeypatch, DEPLOYMENT_EDITION=DeploymentEdition.COMMUNITY) - init_validation.validate_password.side_effect = AlreadyInitializedError - app.secret_key = "test-secret" - - with app.test_request_context("/console/api/init", method="POST"): - with pytest.raises(AlreadySetupError): - init_validate.validate_init_password(init_validate.InitValidatePayload(password="pw")) - - -def test_validate_init_password_wrong_password( - app: Flask, - init_validation: Mock, - monkeypatch: pytest.MonkeyPatch, -) -> None: - apply_config_overrides(monkeypatch, DEPLOYMENT_EDITION=DeploymentEdition.COMMUNITY) - init_validation.validate_password.side_effect = InvalidInitializationPasswordError - app.secret_key = "test-secret" - - with app.test_request_context("/console/api/init", method="POST"): - with pytest.raises(InitValidateFailedError): - init_validate.validate_init_password(init_validate.InitValidatePayload(password="wrong")) - assert init_validate.session.get("is_init_validated") is False - - -def test_validate_init_password_success( - app: Flask, - init_validation: Mock, - monkeypatch: pytest.MonkeyPatch, -) -> None: - apply_config_overrides(monkeypatch, DEPLOYMENT_EDITION=DeploymentEdition.COMMUNITY) - app.secret_key = "test-secret" - - with app.test_request_context("/console/api/init", method="POST"): - result = init_validate.validate_init_password(init_validate.InitValidatePayload(password="expected")) - - assert result.result == "success" - assert init_validate.session.get("is_init_validated") is True + assert response.status_code == 201 + assert response.get_json() == {"result": "success"} init_validation.validate_password.assert_called_once_with("expected") + with client.session_transaction() as browser_session: + assert browser_session["is_init_validated"] is True + + @staticmethod + def test_validate_init_password_rejects_a_mismatch( + app: DifyApp, + init_validation: Mock, + monkeypatch: pytest.MonkeyPatch, + ) -> None: + apply_config_overrides(monkeypatch, DEPLOYMENT_EDITION=DeploymentEdition.COMMUNITY) + init_validation.validate_password.side_effect = InvalidInitializationPasswordError + client = app.test_client() + + response = client.post("/console/api/init", json={"password": "wrong"}) + + assert response.status_code == 401 + with client.session_transaction() as browser_session: + assert browser_session["is_init_validated"] is False + + @staticmethod + def test_validate_init_password_rejects_an_initialized_installation( + app: DifyApp, + init_validation: Mock, + monkeypatch: pytest.MonkeyPatch, + ) -> None: + apply_config_overrides(monkeypatch, DEPLOYMENT_EDITION=DeploymentEdition.COMMUNITY) + init_validation.validate_password.side_effect = AlreadyInitializedError + + response = app.test_client().post("/console/api/init", json={"password": "expected"}) + + assert response.status_code == 403 + + @staticmethod + @pytest.mark.parametrize( + "payload", + [ + pytest.param({}, id="missing-password"), + pytest.param({"password": "x" * 31}, id="password-too-long"), + ], + ) + def test_validate_init_password_rejects_an_invalid_payload( + app: DifyApp, + init_validation: Mock, + monkeypatch: pytest.MonkeyPatch, + payload: dict[str, str], + ) -> None: + apply_config_overrides(monkeypatch, DEPLOYMENT_EDITION=DeploymentEdition.COMMUNITY) + + response = app.test_client().post("/console/api/init", json=payload) + + assert response.status_code == 422 + init_validation.validate_password.assert_not_called() + + @staticmethod + def test_validate_init_password_is_not_available_in_cloud( + app: DifyApp, + init_validation: Mock, + monkeypatch: pytest.MonkeyPatch, + ) -> None: + apply_config_overrides(monkeypatch, DEPLOYMENT_EDITION=DeploymentEdition.CLOUD) + + response = app.test_client().post("/console/api/init", json={"password": "expected"}) + + assert response.status_code == 404 + init_validation.validate_password.assert_not_called() -@pytest.mark.parametrize( - ("session_value", "expected"), - [ - pytest.param(None, False, id="missing"), - pytest.param(False, False, id="not-validated"), - pytest.param(True, True, id="validated"), - ], -) -def test_is_init_validated_passes_session_state( - app: Flask, - init_validation: Mock, - session_value: bool | None, - expected: bool, -) -> None: - init_validation.is_validated.return_value = True - app.secret_key = "test-secret" +"""Tests for the Flask adapter around initialization validation.""" - with app.test_request_context("/console/api/init", method="GET"): - if session_value is not None: - init_validate.session["is_init_validated"] = session_value - assert init_validate.is_init_validated() is True - init_validation.is_validated.assert_called_once_with(session_validated=expected) +class TestInitValidateCases: + @staticmethod + @pytest.fixture + def init_validation(monkeypatch: pytest.MonkeyPatch) -> Mock: + service = create_autospec(InitValidationService, instance=True, spec_set=True) + application_services = SimpleNamespace(init_validation=service) + monkeypatch.setattr(init_validate, "application_services", lambda: application_services) + return service + + @staticmethod + def test_get_init_status_finished(app: Flask, init_validation: Mock) -> None: + init_validation.is_validated.return_value = True + app.secret_key = "test-secret" + + with app.test_request_context("/console/api/init", method="GET"): + result = init_validate.get_init_status() + + assert result.status == "finished" + + @staticmethod + def test_get_init_status_not_started(app: Flask, init_validation: Mock) -> None: + init_validation.is_validated.return_value = False + app.secret_key = "test-secret" + + with app.test_request_context("/console/api/init", method="GET"): + result = init_validate.get_init_status() + + assert result.status == "not_started" + + @staticmethod + def test_validate_init_password_already_setup( + app: Flask, + init_validation: Mock, + monkeypatch: pytest.MonkeyPatch, + ) -> None: + apply_config_overrides(monkeypatch, DEPLOYMENT_EDITION=DeploymentEdition.COMMUNITY) + init_validation.validate_password.side_effect = AlreadyInitializedError + app.secret_key = "test-secret" + + with app.test_request_context("/console/api/init", method="POST"): + with pytest.raises(AlreadySetupError): + init_validate.validate_init_password(init_validate.InitValidatePayload(password="pw")) + + @staticmethod + def test_validate_init_password_wrong_password( + app: Flask, + init_validation: Mock, + monkeypatch: pytest.MonkeyPatch, + ) -> None: + apply_config_overrides(monkeypatch, DEPLOYMENT_EDITION=DeploymentEdition.COMMUNITY) + init_validation.validate_password.side_effect = InvalidInitializationPasswordError + app.secret_key = "test-secret" + + with app.test_request_context("/console/api/init", method="POST"): + with pytest.raises(InitValidateFailedError): + init_validate.validate_init_password(init_validate.InitValidatePayload(password="wrong")) + assert init_validate.session.get("is_init_validated") is False + + @staticmethod + def test_validate_init_password_success( + app: Flask, + init_validation: Mock, + monkeypatch: pytest.MonkeyPatch, + ) -> None: + apply_config_overrides(monkeypatch, DEPLOYMENT_EDITION=DeploymentEdition.COMMUNITY) + app.secret_key = "test-secret" + + with app.test_request_context("/console/api/init", method="POST"): + result = init_validate.validate_init_password(init_validate.InitValidatePayload(password="expected")) + + assert result.result == "success" + assert init_validate.session.get("is_init_validated") is True + init_validation.validate_password.assert_called_once_with("expected") + + @staticmethod + @pytest.mark.parametrize( + ("session_value", "expected"), + [ + pytest.param(None, False, id="missing"), + pytest.param(False, False, id="not-validated"), + pytest.param(True, True, id="validated"), + ], + ) + def test_is_init_validated_passes_session_state( + app: Flask, + init_validation: Mock, + session_value: bool | None, + expected: bool, + ) -> None: + init_validation.is_validated.return_value = True + app.secret_key = "test-secret" + + with app.test_request_context("/console/api/init", method="GET"): + if session_value is not None: + init_validate.session["is_init_validated"] = session_value + + assert init_validate.is_init_validated() is True + init_validation.is_validated.assert_called_once_with(session_validated=expected) diff --git a/api/tests/unit_tests/controllers/console/test_knowledge_fs_creation.py b/api/tests/unit_tests/controllers/console/test_knowledge_fs_creation.py deleted file mode 100644 index c2d6fbf7fd4..00000000000 --- a/api/tests/unit_tests/controllers/console/test_knowledge_fs_creation.py +++ /dev/null @@ -1,96 +0,0 @@ -"""KnowledgeFS creation must preserve legacy roles and Enterprise RBAC admission.""" - -from collections.abc import Callable -from http import HTTPStatus - -import pytest -from flask import Flask -from werkzeug.exceptions import Forbidden - -from controllers.console import flask_admission, wraps -from controllers.console.knowledge_fs import resources -from core.rbac.entities import RBACPermission -from enums import DeploymentEdition -from libs import login -from libs.login import AccountWithTenant -from models.account import TenantAccountRole -from models.knowledge_fs import KnowledgeFSControlSpaceState -from services.enterprise.rbac_service import RBACService -from services.knowledge_fs.product_dto import KnowledgeFSSpaceCreateResponse -from tests.unit_tests.controllers.knowledge_fs_controller_fakes import FakeComponent, FakeRuntime, Recorder, invocation -from tests.unit_tests.model_factories import make_account - - -@pytest.mark.parametrize( - ("role", "rbac_enabled", "rbac_allowed", "allowed"), - [ - (TenantAccountRole.NORMAL, False, True, False), - (None, False, True, False), - (TenantAccountRole.OWNER, False, True, True), - (TenantAccountRole.ADMIN, False, True, True), - (TenantAccountRole.EDITOR, False, True, True), - (TenantAccountRole.DATASET_OPERATOR, False, True, True), - (TenantAccountRole.NORMAL, True, True, True), - (TenantAccountRole.ADMIN, True, False, False), - ], -) -def test_create_space_enforces_admission_before_provisioning( - app: Flask, - monkeypatch: pytest.MonkeyPatch, - config_overrides: Callable[..., None], - role: TenantAccountRole | None, - rbac_enabled: bool, - rbac_allowed: bool, - allowed: bool, -) -> None: - config_overrides( - DEPLOYMENT_EDITION=DeploymentEdition.CLOUD, - KNOWLEDGE_FS_ENABLED=True, - RBAC_ENABLED=rbac_enabled, - LOGIN_DISABLED=False, - ) - account = make_account(role=role) - identity = AccountWithTenant(account, "workspace") - monkeypatch.setattr(login, "_resolve_current_user", lambda: account) - monkeypatch.setattr(login, "check_csrf_token", lambda _request, _user_id: None) - monkeypatch.setattr(wraps, "current_account_with_tenant", lambda: identity) - monkeypatch.setattr(flask_admission, "current_account_with_tenant", lambda: identity) - monkeypatch.setattr(resources, "current_account_with_tenant", lambda: identity) - monkeypatch.setattr(wraps, "check_knowledge_rate_limit", lambda: None) - check_access = Recorder(rbac_allowed) - monkeypatch.setattr(RBACService.CheckAccess, "check", check_access) - create_space = Recorder( - KnowledgeFSSpaceCreateResponse( - control_space_id="space", - state=KnowledgeFSControlSpaceState.PROVISIONING, - operation_id="operation", - model_setup_required=True, - ) - ) - runtime = FakeRuntime(application=FakeComponent(create_space=create_space)) - monkeypatch.setattr(resources, "_console_services", lambda: runtime) - - with app.test_request_context("/knowledge-fs/spaces", method="POST", json={"name": "Test", "slug": "test"}): - if allowed: - result, status = resources.KnowledgeFSSpacesApi().post() - assert status == HTTPStatus.ACCEPTED - assert result["control_space_id"] == "space" - else: - with pytest.raises(Forbidden): - resources.KnowledgeFSSpacesApi().post() - - if allowed: - assert create_space.only_call.kwargs["tenant_id"] == "workspace" - assert create_space.only_call.kwargs["account_id"] == account.id - else: - assert not create_space.calls - if rbac_enabled: - assert check_access.only_call == invocation( - "workspace", - account.id, - scene=RBACPermission.DATASET_CREATE_AND_MANAGEMENT.value, - resource_type=None, - resource_id=None, - ) - else: - assert not check_access.calls diff --git a/api/tests/unit_tests/controllers/console/test_notification.py b/api/tests/unit_tests/controllers/console/test_notification.py index 129d27668f7..1b8a57f5e36 100644 --- a/api/tests/unit_tests/controllers/console/test_notification.py +++ b/api/tests/unit_tests/controllers/console/test_notification.py @@ -15,8 +15,8 @@ from enums import DeploymentEdition from libs.login import AccountWithTenant from machinery.context import RequestContext from services.entities.notification_entities import NotificationItem, NotificationResult -from tests.unit_tests.config_override import apply_config_overrides -from tests.unit_tests.model_factories import make_account +from tests.helpers.unit.config_override import apply_config_overrides +from tests.helpers.unit.model_factories import make_account def _request_context() -> RequestContext: diff --git a/api/tests/unit_tests/controllers/console/test_oauth_custom_client_permissions.py b/api/tests/unit_tests/controllers/console/test_oauth_custom_client_permissions.py deleted file mode 100644 index d53b4572526..00000000000 --- a/api/tests/unit_tests/controllers/console/test_oauth_custom_client_permissions.py +++ /dev/null @@ -1,27 +0,0 @@ -from inspect import unwrap -from types import FunctionType - -import pytest - -from controllers.common.rbac import RBACPermission, Workspace -from controllers.console.datasets.rag_pipeline.datasource_auth import DatasourceAuthOauthCustomClient -from controllers.console.workspace.tool_providers import ToolOAuthCustomClient -from tests.unit_tests.controllers.rbac_introspection import rbac_checks - - -@pytest.mark.parametrize( - ("method", "legacy_gate"), - [ - (ToolOAuthCustomClient.delete, "is_admin_or_owner_required"), - (DatasourceAuthOauthCustomClient.delete, "edit_permission_required"), - ], -) -def test_custom_oauth_client_delete_requires_management_permission(method: FunctionType, legacy_gate: str) -> None: - legacy_wrapper = unwrap(method, stop=lambda wrapper: legacy_gate in wrapper.__code__.co_qualname) - assert legacy_gate in legacy_wrapper.__code__.co_qualname - - # Both routes pass resource_required=False, so their bundle resolves to a Workspace - # locator regardless of the resource_type they declared (see _single_check). - [check] = rbac_checks(method) - assert check.scene == RBACPermission.CREDENTIAL_MANAGE - assert isinstance(check.locator, Workspace) diff --git a/api/tests/unit_tests/controllers/console/test_resource_access_token.py b/api/tests/unit_tests/controllers/console/test_resource_access_token.py index 184a21cff50..46051b8c231 100644 --- a/api/tests/unit_tests/controllers/console/test_resource_access_token.py +++ b/api/tests/unit_tests/controllers/console/test_resource_access_token.py @@ -15,8 +15,8 @@ from controllers.console import flask_admission from controllers.console import resource_access_token as controller from libs.login import AccountWithTenant from models.account import TenantAccountRole -from tests.unit_tests.config_override import apply_config_overrides -from tests.unit_tests.model_factories import make_account +from tests.helpers.unit.config_override import apply_config_overrides +from tests.helpers.unit.model_factories import make_account _ENDPOINTS = [ controller.ResourceAccessTokenListApi.get, diff --git a/api/tests/unit_tests/controllers/console/test_fastopenapi_setup.py b/api/tests/unit_tests/controllers/console/test_setup.py similarity index 100% rename from api/tests/unit_tests/controllers/console/test_fastopenapi_setup.py rename to api/tests/unit_tests/controllers/console/test_setup.py diff --git a/api/tests/unit_tests/controllers/console/test_system.py b/api/tests/unit_tests/controllers/console/test_system.py index 8eeb5e48ab1..b50b257285b 100644 --- a/api/tests/unit_tests/controllers/console/test_system.py +++ b/api/tests/unit_tests/controllers/console/test_system.py @@ -1,119 +1,182 @@ +import builtins import logging from unittest.mock import MagicMock, patch import pytest +from flask.views import MethodView import controllers.console.system as system_module -from tests.unit_tests.config_override import config_overrides_context +from configs import dify_config +from controllers.console.system import _has_new_version +from dify_app import DifyApp +from extensions import ext_fastopenapi +from tests.helpers.unit.config_override import config_overrides_context + +if not hasattr(builtins, "MethodView"): + builtins.MethodView = MethodView # type: ignore[attr-defined] -class TestHasNewVersion: - def test_has_new_version_true(self) -> None: - result = system_module._has_new_version( - latest_version="1.2.0", - current_version="1.1.0", - ) - assert result is True +class TestFastopenapiSystemCases: + @staticmethod + @pytest.fixture + def app() -> DifyApp: + app = DifyApp(__name__) + app.config["TESTING"] = True + return app - def test_has_new_version_false(self) -> None: - result = system_module._has_new_version( - latest_version="1.0.0", - current_version="1.1.0", - ) - assert result is False + @staticmethod + def test_console_ping_fastopenapi_returns_pong(app: DifyApp) -> None: + ext_fastopenapi.init_app(app) - def test_has_new_version_invalid_version(self, caplog: pytest.LogCaptureFixture) -> None: - with caplog.at_level(logging.WARNING, logger="controllers.console.system"): + response = app.test_client().get("/console/api/ping") + + assert response.status_code == 200 + assert response.get_json() == {"result": "pong"} + + @staticmethod + def test_console_version_fastopenapi_returns_current_version(app: DifyApp) -> None: + ext_fastopenapi.init_app(app) + + with config_overrides_context(CHECK_UPDATE_URL=None): + response = app.test_client().get("/console/api/version", query_string={"current_version": "0.0.0"}) + + assert response.status_code == 200 + assert response.get_json() == { + "version": dify_config.project.version, + "release_notes": "", + } + + +class TestSystemCases: + class TestHasNewVersion: + def test_has_new_version_true(self) -> None: result = system_module._has_new_version( - latest_version="invalid", - current_version="1.0.0", + latest_version="1.2.0", + current_version="1.1.0", ) + assert result is True - assert result is False - assert "Invalid version format" in caplog.text + def test_has_new_version_false(self) -> None: + result = system_module._has_new_version( + latest_version="1.0.0", + current_version="1.1.0", + ) + assert result is False + + def test_has_new_version_invalid_version(self, caplog: pytest.LogCaptureFixture) -> None: + with caplog.at_level(logging.WARNING, logger="controllers.console.system"): + result = system_module._has_new_version( + latest_version="invalid", + current_version="1.0.0", + ) + + assert result is False + assert "Invalid version format" in caplog.text + + class TestCheckVersionUpdate: + def test_no_check_update_url(self) -> None: + query = system_module.VersionQuery(current_version="1.0.0") + + with ( + config_overrides_context(CHECK_UPDATE_URL=""), + patch.object( + system_module.dify_config.project, + "version", + "1.0.0", + ), + ): + result = system_module.check_version_update(query) + + assert result == system_module.VersionResponse(version="1.0.0", release_notes="") + + def test_http_error_fallback(self, caplog: pytest.LogCaptureFixture) -> None: + query = system_module.VersionQuery(current_version="1.0.0") + + with ( + config_overrides_context(CHECK_UPDATE_URL="http://example.com"), + patch.object( + system_module.httpx, + "get", + side_effect=Exception("boom"), + ), + caplog.at_level(logging.WARNING, logger="controllers.console.system"), + ): + result = system_module.check_version_update(query) + + assert result.version == "1.0.0" + assert "Check update version error" in caplog.text + + def test_new_version_available(self) -> None: + query = system_module.VersionQuery(current_version="1.0.0") + + response = MagicMock() + response.json.return_value = { + "version": "1.2.0", + "releaseNotes": "New features", + } + + with ( + config_overrides_context(CHECK_UPDATE_URL="http://example.com"), + patch.object( + system_module.httpx, + "get", + return_value=response, + ), + patch.object( + system_module.dify_config.project, + "version", + "1.0.0", + ), + ): + result = system_module.check_version_update(query) + + assert result.version == "1.2.0" + assert result.release_notes == "New features" + + def test_no_new_version(self) -> None: + query = system_module.VersionQuery(current_version="1.2.0") + + response = MagicMock() + response.json.return_value = { + "version": "1.1.0", + } + + with ( + config_overrides_context(CHECK_UPDATE_URL="http://example.com"), + patch.object( + system_module.httpx, + "get", + return_value=response, + ), + patch.object( + system_module.dify_config.project, + "version", + "1.2.0", + ), + ): + result = system_module.check_version_update(query) + + assert result == system_module.VersionResponse(version="1.2.0", release_notes="") -class TestCheckVersionUpdate: - def test_no_check_update_url(self) -> None: - query = system_module.VersionQuery(current_version="1.0.0") - - with ( - config_overrides_context(CHECK_UPDATE_URL=""), - patch.object( - system_module.dify_config.project, - "version", - "1.0.0", - ), - ): - result = system_module.check_version_update(query) - - assert result == system_module.VersionResponse(version="1.0.0", release_notes="") - - def test_http_error_fallback(self, caplog: pytest.LogCaptureFixture) -> None: - query = system_module.VersionQuery(current_version="1.0.0") - - with ( - config_overrides_context(CHECK_UPDATE_URL="http://example.com"), - patch.object( - system_module.httpx, - "get", - side_effect=Exception("boom"), - ), - caplog.at_level(logging.WARNING, logger="controllers.console.system"), - ): - result = system_module.check_version_update(query) - - assert result.version == "1.0.0" - assert "Check update version error" in caplog.text - - def test_new_version_available(self) -> None: - query = system_module.VersionQuery(current_version="1.0.0") - - response = MagicMock() - response.json.return_value = { - "version": "1.2.0", - "releaseNotes": "New features", - } - - with ( - config_overrides_context(CHECK_UPDATE_URL="http://example.com"), - patch.object( - system_module.httpx, - "get", - return_value=response, - ), - patch.object( - system_module.dify_config.project, - "version", - "1.0.0", - ), - ): - result = system_module.check_version_update(query) - - assert result.version == "1.2.0" - assert result.release_notes == "New features" - - def test_no_new_version(self) -> None: - query = system_module.VersionQuery(current_version="1.2.0") - - response = MagicMock() - response.json.return_value = { - "version": "1.1.0", - } - - with ( - config_overrides_context(CHECK_UPDATE_URL="http://example.com"), - patch.object( - system_module.httpx, - "get", - return_value=response, - ), - patch.object( - system_module.dify_config.project, - "version", - "1.2.0", - ), - ): - result = system_module.check_version_update(query) - - assert result == system_module.VersionResponse(version="1.2.0", release_notes="") +class TestCompareVersionsCases: + @staticmethod + @pytest.mark.parametrize( + ("latest_version", "current_version", "expected"), + [ + ("1.0.1", "1.0.0", True), + ("1.1.0", "1.0.0", True), + ("2.0.0", "1.9.9", True), + ("1.0.0", "1.0.0", False), + ("1.0.0", "1.0.1", False), + ("1.0.0", "2.0.0", False), + ("1.0.1", "1.0.0-beta", True), + ("1.0.0", "1.0.0-alpha", True), + ("1.0.0-beta", "1.0.0-alpha", True), + ("1.0.0", "1.0.0-rc1", True), + ("1.0.0", "0.9.9", True), + ("1.0.0", "1.0.0-dev", True), + ], + ) + def test_has_new_version(latest_version: str, current_version: str, expected: bool) -> None: + assert _has_new_version(latest_version=latest_version, current_version=current_version) == expected diff --git a/api/tests/unit_tests/controllers/console/test_workflow_run_archive.py b/api/tests/unit_tests/controllers/console/test_workflow_run_archive.py index 1f26fa80609..0cbea8b1528 100644 --- a/api/tests/unit_tests/controllers/console/test_workflow_run_archive.py +++ b/api/tests/unit_tests/controllers/console/test_workflow_run_archive.py @@ -39,8 +39,8 @@ from services.retention.workflow_run.archive_log_service import ( WorkflowRunArchiveNotFoundError, WorkflowRunArchiveService, ) -from tests.unit_tests.config_override import apply_config_overrides -from tests.unit_tests.model_factories import make_account +from tests.helpers.unit.config_override import apply_config_overrides +from tests.helpers.unit.model_factories import make_account _CONTEXT = RequestContext( request_id="request-1", diff --git a/api/tests/unit_tests/controllers/console/test_workspace_account.py b/api/tests/unit_tests/controllers/console/test_workspace_account.py deleted file mode 100644 index 60e0bb821a3..00000000000 --- a/api/tests/unit_tests/controllers/console/test_workspace_account.py +++ /dev/null @@ -1,419 +0,0 @@ -import inspect -from types import SimpleNamespace -from unittest.mock import MagicMock, patch -from uuid import NAMESPACE_URL, uuid5 - -import pytest -from flask import Flask, request -from sqlalchemy.orm import Session -from typing_extensions import TypedDict - -from controllers.console.auth.error import InvalidTokenError -from controllers.console.error import EducationActivateLimitError, EducationVerifyLimitError, EmailDomainSuspendedError -from controllers.console.workspace.account import ( - AccountDeleteUpdateFeedbackApi, - AccountDeletionFeedbackPayload, - ChangeEmailCheckApi, - ChangeEmailResetApi, - ChangeEmailResetPayload, - ChangeEmailSendEmailApi, - ChangeEmailSendPayload, - ChangeEmailValidityPayload, - CheckEmailUnique, - CheckEmailUniquePayload, - EducationActivatePayload, - EducationApi, - EducationVerifyApi, -) -from machinery.context import RequestContext -from models import Account, AccountStatus, Tenant, TenantAccountJoin -from models.account import TenantAccountRole -from services import account_errors -from services.account_email import normalize_email -from services.entities.account_entities import AccountEducationActivation, ChangeEmailVerification -from services.entities.auth_entities import ( - ChangeEmailNewEmailToken, - ChangeEmailNewEmailVerifiedToken, - ChangeEmailOldEmailToken, - ChangeEmailOldEmailVerifiedToken, - ChangeEmailPhase, - ChangeEmailTokenBase, -) - - -@pytest.fixture -def app() -> Flask: - app = Flask(__name__) - app.config["TESTING"] = True - app.config["RESTX_MASK_HEADER"] = "X-Fields" - setattr(app, "login_manager", SimpleNamespace(load_user_from_request_context=lambda: None)) # noqa: B010 - return app - - -def _build_account(email: str, account_id: str = "acc", tenant: Tenant | None = None) -> Account: - if tenant is None: - tenant_obj = Tenant(name="Tenant") - tenant_obj.id = "tenant-id" - else: - tenant_obj = tenant - account = Account(name=account_id, email=email) - account.email = email - account.id = account_id - account.status = AccountStatus.ACTIVE - account._current_tenant = tenant_obj - return account - - -def _stable_uuid(value: str) -> str: - return str(uuid5(NAMESPACE_URL, value)) - - -def _persist_account_with_tenant(session: Session, email: str, account_name: str = "account") -> tuple[Account, Tenant]: - tenant = Tenant(name=f"{account_name} tenant") - tenant.id = _stable_uuid(f"tenant:{account_name}") - account = Account( - name=account_name, email=email, normalized_email=normalize_email(email), status=AccountStatus.ACTIVE - ) - account.id = _stable_uuid(f"account:{account_name}") - membership = TenantAccountJoin( - tenant_id=tenant.id, - account_id=account.id, - current=True, - role=TenantAccountRole.OWNER, - ) - session.add_all([account, tenant, membership]) - session.commit() - account._current_tenant = tenant - account.role = TenantAccountRole.OWNER - return account, tenant - - -class _ChangeEmailTokenKwargs(TypedDict, closed=True): - account_id: str - email: str - old_email: str - code: str - - -def _build_change_email_token( - phase: str, - *, - account_id: str = "acc", - email: str, - old_email: str, - code: str = "1234", -) -> ChangeEmailTokenBase: - token_kwargs: _ChangeEmailTokenKwargs = { - "account_id": account_id, - "email": email, - "old_email": old_email, - "code": code, - } - if phase == "old_email": - return ChangeEmailOldEmailToken(**token_kwargs) - if phase == ChangeEmailPhase.OLD_EMAIL_VERIFIED: - return ChangeEmailOldEmailVerifiedToken(**token_kwargs) - if phase == "new_email": - return ChangeEmailNewEmailToken(**token_kwargs) - if phase == ChangeEmailPhase.NEW_EMAIL_VERIFIED: - return ChangeEmailNewEmailVerifiedToken(**token_kwargs) - raise AssertionError(f"Unsupported phase for test helper: {phase}") - - -class TestEducationApi: - def test_post_activates_education_discount(self, app: Flask) -> None: - education = MagicMock() - education.activate.return_value = AccountEducationActivation(message="success") - request_context = RequestContext( - request_id="request-1", - trace_id=None, - account_id="account-1", - active_workspace_id="workspace-1", - ) - - with ( - app.test_request_context( - "/account/education", - method="POST", - json={"token": "education-token", "institution": "Dify University", "role": "Student"}, - ), - patch( - "controllers.console.workspace.account.application_services", - return_value=SimpleNamespace(accounts=SimpleNamespace(education=education)), - ), - ): - api = EducationApi() - method = inspect.unwrap(api.post) - result = method(api, EducationActivatePayload.model_validate(request.get_json() or {}), request_context) - - assert result == {"message": "success"} - education.activate.assert_called_once_with( - request_context, - token="education-token", - institution="Dify University", - role="Student", - ) - - def test_verify_maps_rate_limit_error(self, app: Flask) -> None: - education = MagicMock() - education.verify.side_effect = account_errors.EducationRateLimitExceededError - request_context = RequestContext( - request_id="request-1", - trace_id=None, - account_id="account-1", - active_workspace_id="workspace-1", - ) - - with ( - app.test_request_context("/account/education/verify", method="GET"), - patch( - "controllers.console.workspace.account.application_services", - return_value=SimpleNamespace(accounts=SimpleNamespace(education=education)), - ), - ): - api = EducationVerifyApi() - with pytest.raises(EducationVerifyLimitError): - inspect.unwrap(api.get)(api, request_context) - - def test_post_maps_rate_limit_error(self, app: Flask) -> None: - education = MagicMock() - education.activate.side_effect = account_errors.EducationRateLimitExceededError - request_context = RequestContext( - request_id="request-1", - trace_id=None, - account_id="account-1", - active_workspace_id="workspace-1", - ) - - with ( - app.test_request_context( - "/account/education", - method="POST", - json={"token": "education-token", "institution": "Dify University", "role": "Student"}, - ), - patch( - "controllers.console.workspace.account.application_services", - return_value=SimpleNamespace(accounts=SimpleNamespace(education=education)), - ), - ): - api = EducationApi() - with pytest.raises(EducationActivateLimitError): - inspect.unwrap(api.post)( - api, EducationActivatePayload.model_validate(request.get_json() or {}), request_context - ) - - -def _change_email_context(account_id: str = "acc") -> RequestContext: - return RequestContext( - request_id="request-1", - trace_id="trace-1", - account_id=account_id, - active_workspace_id="workspace-1", - ) - - -class TestChangeEmailControllers: - def test_send_delegates_parsed_request_to_application_service(self, app: Flask) -> None: - change_email = MagicMock() - change_email.send_code.return_value = "change-token" - context = _change_email_context() - payload = { - "email": "new@example.com", - "language": "zh-Hans", - "phase": "new_email", - "token": "verified-old-token", - } - - with ( - app.test_request_context( - "/account/change-email", - method="POST", - json=payload, - environ_base={"REMOTE_ADDR": "127.0.0.1"}, - ), - patch( - "controllers.console.workspace.account.application_services", - return_value=SimpleNamespace(accounts=SimpleNamespace(change_email=change_email)), - ), - ): - api = ChangeEmailSendEmailApi() - response = inspect.unwrap(api.post)( - api, ChangeEmailSendPayload.model_validate(request.get_json() or {}), context - ) - - assert response == {"result": "success", "data": "change-token"} - change_email.send_code.assert_called_once_with( - context, - requested_email="new@example.com", - language="zh-Hans", - phase="new_email", - predecessor_token="verified-old-token", - ip_address="127.0.0.1", - ) - - def test_send_maps_invalid_state_to_http_error(self, app: Flask) -> None: - change_email = MagicMock() - change_email.send_code.side_effect = account_errors.InvalidChangeEmailTokenError - - with ( - app.test_request_context( - "/account/change-email", - method="POST", - json={"email": "new@example.com", "phase": "new_email"}, - ), - patch( - "controllers.console.workspace.account.application_services", - return_value=SimpleNamespace(accounts=SimpleNamespace(change_email=change_email)), - ), - ): - api = ChangeEmailSendEmailApi() - method = inspect.unwrap(api.post) - with pytest.raises(InvalidTokenError): - method(api, ChangeEmailSendPayload.model_validate(request.get_json() or {}), _change_email_context()) - - def test_validity_serializes_promoted_token(self, app: Flask) -> None: - change_email = MagicMock() - change_email.verify_code.return_value = ChangeEmailVerification( - email="new@example.com", - token="verified-token", - ) - context = _change_email_context() - - with ( - app.test_request_context( - "/account/change-email/validity", - method="POST", - json={"email": "New@Example.com", "code": "123456", "token": "pending-token"}, - ), - patch( - "controllers.console.workspace.account.application_services", - return_value=SimpleNamespace(accounts=SimpleNamespace(change_email=change_email)), - ), - ): - api = ChangeEmailCheckApi() - response = inspect.unwrap(api.post)( - api, ChangeEmailValidityPayload.model_validate(request.get_json() or {}), context - ) - - assert response == {"is_valid": True, "email": "new@example.com", "token": "verified-token"} - change_email.verify_code.assert_called_once_with( - context, - email="New@Example.com", - code="123456", - token="pending-token", - ) - - def test_reset_returns_updated_account(self, app: Flask) -> None: - change_email = MagicMock() - updated_account = _build_account("new@example.com", "acc") - change_email.reset.return_value = updated_account - context = _change_email_context() - - with ( - app.test_request_context( - "/account/change-email/reset", - method="POST", - json={"new_email": "New@Example.com", "token": "verified-token"}, - ), - patch( - "controllers.console.workspace.account.application_services", - return_value=SimpleNamespace(accounts=SimpleNamespace(change_email=change_email)), - ), - ): - api = ChangeEmailResetApi() - response = inspect.unwrap(api.post)( - api, ChangeEmailResetPayload.model_validate(request.get_json() or {}), context - ) - - assert response["email"] == "new@example.com" - change_email.reset.assert_called_once_with( - context, - new_email="New@Example.com", - token="verified-token", - ) - - def test_reset_maps_suspended_email_domain(self, app: Flask) -> None: - change_email = MagicMock() - change_email.reset.side_effect = account_errors.AccountEmailDomainSuspendedError - - with ( - app.test_request_context( - "/account/change-email/reset", - method="POST", - json={"new_email": "user@suspended.example", "token": "verified-token"}, - ), - patch( - "controllers.console.workspace.account.application_services", - return_value=SimpleNamespace(accounts=SimpleNamespace(change_email=change_email)), - ), - ): - api = ChangeEmailResetApi() - with pytest.raises(EmailDomainSuspendedError): - inspect.unwrap(api.post)( - api, ChangeEmailResetPayload.model_validate(request.get_json() or {}), _change_email_context() - ) - - -class TestAccountDeletionFeedback: - def test_delegates_feedback_to_application_service(self, app: Flask) -> None: - deletion_feedback = MagicMock() - with ( - app.test_request_context( - "/account/delete/feedback", - method="POST", - json={"email": "User@Example.com", "feedback": "test"}, - ), - patch( - "controllers.console.workspace.account.application_services", - return_value=SimpleNamespace( - accounts=SimpleNamespace(deletion_feedback=deletion_feedback), - ), - ), - ): - api = AccountDeleteUpdateFeedbackApi() - method = inspect.unwrap(api.post) - response = method(api, AccountDeletionFeedbackPayload.model_validate(request.get_json() or {})) - - assert response == {"result": "success"} - deletion_feedback.submit.assert_called_once_with(email="User@Example.com", feedback="test") - - -class TestCheckEmailUnique: - def test_delegates_to_email_availability_policy(self, app: Flask) -> None: - change_email = MagicMock() - - with ( - app.test_request_context( - "/account/change-email/check-email-unique", - method="POST", - json={"email": "Case@Test.com"}, - ), - patch( - "controllers.console.workspace.account.application_services", - return_value=SimpleNamespace(accounts=SimpleNamespace(change_email=change_email)), - ), - ): - api = CheckEmailUnique() - response = inspect.unwrap(api.post)(api, CheckEmailUniquePayload.model_validate(request.get_json() or {})) - - assert response == {"result": "success"} - change_email.ensure_available.assert_called_once_with("Case@Test.com") - - def test_maps_suspended_email_domain(self, app: Flask) -> None: - change_email = MagicMock() - change_email.ensure_available.side_effect = account_errors.AccountEmailDomainSuspendedError - - with ( - app.test_request_context( - "/account/change-email/check-email-unique", - method="POST", - json={"email": "user@suspended.example"}, - ), - patch( - "controllers.console.workspace.account.application_services", - return_value=SimpleNamespace(accounts=SimpleNamespace(change_email=change_email)), - ), - ): - api = CheckEmailUnique() - with pytest.raises(EmailDomainSuspendedError): - inspect.unwrap(api.post)(api, CheckEmailUniquePayload.model_validate(request.get_json() or {})) diff --git a/api/tests/unit_tests/controllers/console/test_workspace_admission.py b/api/tests/unit_tests/controllers/console/test_workspace_admission.py deleted file mode 100644 index e83e14e4bab..00000000000 --- a/api/tests/unit_tests/controllers/console/test_workspace_admission.py +++ /dev/null @@ -1,115 +0,0 @@ -"""Workspace admission composes existing checks before controller execution.""" - -from collections.abc import Callable -from unittest.mock import Mock - -import pytest -from flask import Flask -from werkzeug.exceptions import Forbidden, NotFound, Unauthorized - -from controllers.console import flask_admission, wraps -from controllers.console.error import NotInitValidateError, NotSetupError -from controllers.console.notification import NotificationApi -from controllers.console.workspace.error import AccountNotInitializedError -from controllers.console.workspace.workspace import WorkspacePermissionApi -from enums import DeploymentEdition -from libs import login as login_adapter -from libs.login import AccountWithTenant -from machinery.context import RequestContext -from models.account import Account, AccountStatus -from services.entities.feature_entities import FeatureModel - - -def test_admin_key_admits_without_account_login(app: Flask, config_overrides: Callable[..., None]) -> None: - config_overrides(DEPLOYMENT_EDITION=DeploymentEdition.CLOUD, ADMIN_API_KEY="admin") - - class Handler: - @flask_admission.console_admin_admission - def get(self, *, page: int) -> int: - return page - - with app.test_request_context(headers={"Authorization": "Bearer admin"}): - assert Handler().get(page=2) == 2 - with app.test_request_context(), pytest.raises(Unauthorized): - Handler().get(page=2) - - -@pytest.mark.parametrize( - ("edition", "enabled", "allowed"), - [ - (DeploymentEdition.CLOUD, True, True), - (DeploymentEdition.CLOUD, False, False), - (DeploymentEdition.COMMUNITY, False, True), - ], -) -def test_custom_brand_billing_admission( - app: Flask, - monkeypatch: pytest.MonkeyPatch, - config_overrides: Callable[..., None], - edition: DeploymentEdition, - enabled: bool, - allowed: bool, -) -> None: - config_overrides(DEPLOYMENT_EDITION=edition) - monkeypatch.setattr(flask_admission, "setup_required", lambda view: view) - monkeypatch.setattr(flask_admission, "login_required", lambda view: view) - monkeypatch.setattr(flask_admission, "account_initialization_required", lambda view: view) - account = Account(name="User", email="user@example.com") - identity = AccountWithTenant(account, "workspace") - monkeypatch.setattr(flask_admission, "current_account_with_tenant", lambda: identity) - monkeypatch.setattr(wraps, "current_account_with_tenant", lambda: identity) - features = Mock(return_value=FeatureModel(can_replace_logo=enabled)) - monkeypatch.setattr(wraps.FeatureService, "get_features", features) - reached = [] - - class Handler: - @flask_admission.console_account_admission(billing_resource="workspace_custom") - def post(self, context: RequestContext) -> RequestContext: - reached.append(context) - return context - - with app.test_request_context(): - if allowed: - context = Handler().post() - assert isinstance(context, RequestContext) - assert context.active_workspace_id == "workspace" - else: - with pytest.raises(Forbidden): - Handler().post() - assert bool(reached) is allowed - assert features.call_count == int(edition == DeploymentEdition.CLOUD) - - -@pytest.mark.parametrize("controller_type", [WorkspacePermissionApi, NotificationApi]) -@pytest.mark.parametrize( - ("setup_complete", "init_password", "status", "error"), - [ - (False, "", None, NotSetupError), - (False, "initial-password", None, NotInitValidateError), - (True, "", None, Unauthorized), - (True, "", AccountStatus.UNINITIALIZED, AccountNotInitializedError), - (True, "", AccountStatus.ACTIVE, NotFound), - ], -) -def test_account_admission_checks_identity_before_edition( - app: Flask, - monkeypatch: pytest.MonkeyPatch, - config_overrides: Callable[..., None], - controller_type: type[WorkspacePermissionApi] | type[NotificationApi], - setup_complete: bool, - init_password: str, - status: AccountStatus | None, - error: type[Exception], -) -> None: - config_overrides(DEPLOYMENT_EDITION=DeploymentEdition.COMMUNITY, INIT_PASSWORD=init_password, LOGIN_DISABLED=False) - monkeypatch.setattr(wraps, "_is_setup_completed", lambda: setup_complete) - account = Account(name="User", email="user@example.com", status=status or AccountStatus.ACTIVE) - identity = Mock(return_value=account, side_effect=Unauthorized() if status is None else None) - monkeypatch.setattr(login_adapter, "_resolve_current_user", identity) - monkeypatch.setattr(login_adapter, "check_csrf_token", Mock()) - monkeypatch.setattr(wraps, "current_account_with_tenant", lambda: AccountWithTenant(account, "workspace")) - - with app.test_request_context(), pytest.raises(error): - controller_type().get() - - assert identity.call_count == int(setup_complete) diff --git a/api/tests/unit_tests/controllers/console/test_workspace_credential_mutation_permissions.py b/api/tests/unit_tests/controllers/console/test_workspace_credential_mutation_permissions.py deleted file mode 100644 index 7e4888e185b..00000000000 --- a/api/tests/unit_tests/controllers/console/test_workspace_credential_mutation_permissions.py +++ /dev/null @@ -1,105 +0,0 @@ -from inspect import unwrap -from types import FunctionType - -import pytest - -from controllers.common.rbac import AgentId, RBACPermission, Workspace -from controllers.console.agent.composer import AgentComposerApi -from controllers.console.agent.roster import AgentAppApi -from controllers.console.datasets.rag_pipeline.datasource_auth import DatasourceAuth -from controllers.console.workspace.model_providers import ModelProviderCredentialApi -from controllers.console.workspace.models import ModelProviderModelCredentialApi -from controllers.console.workspace.tool_providers import ( - ToolBuiltinProviderAddApi, - ToolMCPDetailApi, - ToolOAuthCustomClient, -) -from tests.unit_tests.controllers.rbac_introspection import rbac_checks - - -@pytest.mark.parametrize( - ("method", "permission"), - [ - (ToolBuiltinProviderAddApi.post, RBACPermission.CREDENTIAL_CREATE), - ], -) -def test_workspace_credential_mutations_require_management_permission( - method: FunctionType, permission: RBACPermission -) -> None: - legacy_wrapper = unwrap(method, stop=lambda wrapper: "is_admin_or_owner_required" in wrapper.__code__.co_qualname) - assert "is_admin_or_owner_required" in legacy_wrapper.__code__.co_qualname - - [check] = rbac_checks(method) - assert check.scene == permission - assert isinstance(check.locator, Workspace) - - -@pytest.mark.parametrize( - "method", - [ - ModelProviderCredentialApi.get, - ModelProviderModelCredentialApi.get, - ], -) -def test_model_provider_credential_get_requires_admin_and_rbac( - method: FunctionType, -) -> None: - """GET endpoints that return provider credential details must enforce - the same admin + RBAC gates as their sibling POST/PUT/DELETE methods.""" - legacy_wrapper = unwrap(method, stop=lambda wrapper: "is_admin_or_owner_required" in wrapper.__code__.co_qualname) - assert "is_admin_or_owner_required" in legacy_wrapper.__code__.co_qualname - - [check] = rbac_checks(method) - assert check.scene == RBACPermission.CREDENTIAL_MANAGE - assert isinstance(check.locator, Workspace) - - -def test_tool_oauth_custom_client_get_requires_admin_and_rbac() -> None: - """GET endpoint that returns custom OAuth client params must enforce - the same admin + RBAC gates as its sibling POST and DELETE methods.""" - method = ToolOAuthCustomClient.get - - legacy_wrapper = unwrap(method, stop=lambda wrapper: "is_admin_or_owner_required" in wrapper.__code__.co_qualname) - assert "is_admin_or_owner_required" in legacy_wrapper.__code__.co_qualname - - [check] = rbac_checks(method) - assert check.scene == RBACPermission.CREDENTIAL_MANAGE - assert isinstance(check.locator, Workspace) - - -def test_datasource_auth_get_requires_edit_and_rbac() -> None: - """GET endpoint that lists datasource credentials must enforce - the same edit + RBAC gates as its sibling POST method.""" - method = DatasourceAuth.get - - edit_wrapper = unwrap(method, stop=lambda wrapper: "edit_permission_required" in wrapper.__code__.co_qualname) - assert "edit_permission_required" in edit_wrapper.__code__.co_qualname - - [check] = rbac_checks(method) - assert check.scene == RBACPermission.CREDENTIAL_MANAGE - assert isinstance(check.locator, Workspace) - - -@pytest.mark.parametrize( - "method", - [ - AgentAppApi.get, - AgentComposerApi.get, - ], -) -def test_agent_app_get_requires_rbac(method: FunctionType) -> None: - """GET endpoints that return agent app details or composer state must require - the agent preview scene on the agent resource, same as their sibling routes.""" - [check] = rbac_checks(method) - assert check.scene == RBACPermission.AGENT_PREVIEW - assert isinstance(check.locator, AgentId) - - -def test_tool_mcp_detail_get_requires_rbac() -> None: - """GET endpoint that returns MCP provider details must enforce - the workspace MCP management RBAC gate.""" - method = ToolMCPDetailApi.get - - [check] = rbac_checks(method) - assert check.scene == RBACPermission.MCP_MANAGE - assert isinstance(check.locator, Workspace) diff --git a/api/tests/unit_tests/controllers/console/test_wraps.py b/api/tests/unit_tests/controllers/console/test_wraps.py index c9299baab6e..261d79fff7b 100644 --- a/api/tests/unit_tests/controllers/console/test_wraps.py +++ b/api/tests/unit_tests/controllers/console/test_wraps.py @@ -50,7 +50,7 @@ from models import Account, DifySetup from models.account import AccountStatus, TenantAccountRole from models.dataset import RateLimitLog from services.entities.feature_entities import LicenseStatus -from tests.unit_tests.config_override import config_overrides_context +from tests.helpers.unit.config_override import config_overrides_context @pytest.fixture(autouse=True) diff --git a/api/tests/unit_tests/controllers/test_rbac_access_policy_schema.py b/api/tests/unit_tests/controllers/console/workspace/rbac/test_policies.py similarity index 98% rename from api/tests/unit_tests/controllers/test_rbac_access_policy_schema.py rename to api/tests/unit_tests/controllers/console/workspace/rbac/test_policies.py index 46b96d7db4c..c38a5355bbe 100644 --- a/api/tests/unit_tests/controllers/test_rbac_access_policy_schema.py +++ b/api/tests/unit_tests/controllers/console/workspace/rbac/test_policies.py @@ -14,7 +14,7 @@ from sqlalchemy.orm import Session, sessionmaker from controllers.console.workspace.rbac import policies, resources, roles from machinery.context import RequestContext from services.rbac.contracts import RBACResourceType -from tests.unit_tests.rbac_fakes import build_rbac_domain +from tests.helpers.unit.rbac_fakes import build_rbac_domain class _DeleteResource(Protocol): diff --git a/api/tests/unit_tests/controllers/console/workspace/test_rbac_resource_access_routes.py b/api/tests/unit_tests/controllers/console/workspace/rbac/test_resources.py similarity index 98% rename from api/tests/unit_tests/controllers/console/workspace/test_rbac_resource_access_routes.py rename to api/tests/unit_tests/controllers/console/workspace/rbac/test_resources.py index c717d9b061e..45a458b33e0 100644 --- a/api/tests/unit_tests/controllers/console/workspace/test_rbac_resource_access_routes.py +++ b/api/tests/unit_tests/controllers/console/workspace/rbac/test_resources.py @@ -20,7 +20,7 @@ from controllers.console.workspace.rbac import resources as rbac_mod from enums import DeploymentEdition from machinery.context import RequestContext from services.rbac import contracts as rbac_contracts -from tests.unit_tests.rbac_fakes import Initialization, RBACDomain, build_rbac_domain +from tests.helpers.unit.rbac_fakes import Initialization, RBACDomain, build_rbac_domain RESOURCE_ID = { rbac_contracts.RBACResourceType.APP: "app-1", diff --git a/api/tests/unit_tests/controllers/console/workspace/test_rbac.py b/api/tests/unit_tests/controllers/console/workspace/rbac/test_roles.py similarity index 98% rename from api/tests/unit_tests/controllers/console/workspace/test_rbac.py rename to api/tests/unit_tests/controllers/console/workspace/rbac/test_roles.py index 0d9e785e108..ec69b03e44e 100644 --- a/api/tests/unit_tests/controllers/console/workspace/test_rbac.py +++ b/api/tests/unit_tests/controllers/console/workspace/rbac/test_roles.py @@ -23,8 +23,8 @@ from libs.login import AccountWithTenant from machinery.context import RequestContext from models.account import Account, AccountStatus, TenantAccountJoin from services.rbac import contracts as dto -from tests.unit_tests.model_factories import make_account, make_tenant -from tests.unit_tests.rbac_fakes import RBACDomain, build_rbac_domain +from tests.helpers.unit.model_factories import make_account, make_tenant +from tests.helpers.unit.rbac_fakes import RBACDomain, build_rbac_domain CONTEXT = RequestContext("request", "trace", "actor", "tenant") diff --git a/api/tests/unit_tests/controllers/console/workspace/test_account.py b/api/tests/unit_tests/controllers/console/workspace/test_account.py new file mode 100644 index 00000000000..95de5823d56 --- /dev/null +++ b/api/tests/unit_tests/controllers/console/workspace/test_account.py @@ -0,0 +1,1209 @@ +import inspect +from collections.abc import Iterator +from datetime import datetime +from types import SimpleNamespace +from unittest.mock import MagicMock, PropertyMock, patch +from uuid import NAMESPACE_URL, uuid5 + +import pytest +from flask import Flask, request +from jsonschema import Draft202012Validator +from sqlalchemy.orm import Session +from typing_extensions import TypedDict +from werkzeug.exceptions import NotFound, UnprocessableEntity + +from controllers.console import console_ns +from controllers.console.auth.error import ( + EmailAlreadyInUseError, + EmailCodeAccountDeletionRateLimitExceededError, + EmailCodeError, + InvalidTokenError, +) +from controllers.console.error import ( + AccountInFreezeError, + EducationActivateLimitError, + EducationVerifyLimitError, + EmailDomainSuspendedError, + InvalidAccountPasswordRequestError, +) +from controllers.console.workspace.account import ( + AccountAvatarApi, + AccountAvatarPayload, + AccountAvatarQuery, + AccountDeleteApi, + AccountDeletePayload, + AccountDeleteUpdateFeedbackApi, + AccountDeleteVerifyApi, + AccountDeletionFeedbackPayload, + AccountInitApi, + AccountInitPayload, + AccountIntegrateApi, + AccountInterfaceLanguageApi, + AccountInterfaceLanguagePayload, + AccountInterfaceThemeApi, + AccountInterfaceThemePayload, + AccountNameApi, + AccountNamePayload, + AccountPasswordApi, + AccountPasswordPayload, + AccountProfileApi, + AccountProfilePatchPayload, + AccountTimezoneApi, + AccountTimezonePayload, + ChangeEmailCheckApi, + ChangeEmailResetApi, + ChangeEmailResetPayload, + ChangeEmailSendEmailApi, + ChangeEmailSendPayload, + ChangeEmailValidityPayload, + CheckEmailUnique, + CheckEmailUniquePayload, + EducationActivatePayload, + EducationApi, + EducationVerifyApi, +) +from controllers.console.workspace.error import ( + AccountAlreadyInitedError, + CurrentPasswordIncorrectError, + InvalidAccountDeletionCodeError, + MissingInvitationCodeRequestError, +) +from extensions.application_services.account import AccountServices +from extensions.ext_application_services import ApplicationServices +from libs.login import AccountWithTenant +from machinery.context import RequestContext +from models import Account, Tenant, TenantAccountJoin +from models import AccountStatus as _workspace_account_AccountStatus +from models.account import AccountStatus as _accounts_AccountStatus +from models.account import TenantAccountRole +from services import account_errors +from services.account_email import normalize_email +from services.account_errors import ( + AccountAlreadyInitializedError, + AccountDeletionRateLimitError, + AccountEmailAlreadyInUseError, + AccountEmailDomainSuspendedError, + AccountEmailFrozenError, + AvatarFileNotFoundError, + CurrentAccountPasswordIncorrectError, + InvalidAccountDeletionVerificationError, + InvalidAccountPasswordError, + InvalidChangeEmailCodeError, + MissingInvitationCodeError, +) +from services.entities.account_entities import ( + AccountEducationActivation, + AccountIntegrationStatus, + AccountProfileChanges, + ChangeEmailVerification, +) +from services.entities.auth_entities import ( + ChangeEmailNewEmailToken, + ChangeEmailNewEmailVerifiedToken, + ChangeEmailOldEmailToken, + ChangeEmailOldEmailVerifiedToken, + ChangeEmailPhase, + ChangeEmailTokenBase, +) +from tests.helpers.unit.config_override import config_overrides_context + + +def _build_account(email: str, account_id: str = "acc", tenant: Tenant | None = None) -> Account: + if tenant is None: + tenant_obj = Tenant(name="Tenant") + tenant_obj.id = "tenant-id" + else: + tenant_obj = tenant + account = Account(name=account_id, email=email) + account.email = email + account.id = account_id + account.status = _workspace_account_AccountStatus.ACTIVE + account._current_tenant = tenant_obj + return account + + +def _stable_uuid(value: str) -> str: + return str(uuid5(NAMESPACE_URL, value)) + + +def _persist_account_with_tenant(session: Session, email: str, account_name: str = "account") -> tuple[Account, Tenant]: + tenant = Tenant(name=f"{account_name} tenant") + tenant.id = _stable_uuid(f"tenant:{account_name}") + account = Account( + name=account_name, + email=email, + normalized_email=normalize_email(email), + status=_workspace_account_AccountStatus.ACTIVE, + ) + account.id = _stable_uuid(f"account:{account_name}") + membership = TenantAccountJoin( + tenant_id=tenant.id, + account_id=account.id, + current=True, + role=TenantAccountRole.OWNER, + ) + session.add_all([account, tenant, membership]) + session.commit() + account._current_tenant = tenant + account.role = TenantAccountRole.OWNER + return account, tenant + + +class _ChangeEmailTokenKwargs(TypedDict, closed=True): + account_id: str + email: str + old_email: str + code: str + + +def _build_change_email_token( + phase: str, + *, + account_id: str = "acc", + email: str, + old_email: str, + code: str = "1234", +) -> ChangeEmailTokenBase: + token_kwargs: _ChangeEmailTokenKwargs = { + "account_id": account_id, + "email": email, + "old_email": old_email, + "code": code, + } + if phase == "old_email": + return ChangeEmailOldEmailToken(**token_kwargs) + if phase == ChangeEmailPhase.OLD_EMAIL_VERIFIED: + return ChangeEmailOldEmailVerifiedToken(**token_kwargs) + if phase == "new_email": + return ChangeEmailNewEmailToken(**token_kwargs) + if phase == ChangeEmailPhase.NEW_EMAIL_VERIFIED: + return ChangeEmailNewEmailVerifiedToken(**token_kwargs) + raise AssertionError(f"Unsupported phase for test helper: {phase}") + + +def _change_email_context(account_id: str = "acc") -> RequestContext: + return RequestContext( + request_id="request-1", + trace_id="trace-1", + account_id=account_id, + active_workspace_id="workspace-1", + ) + + +class TestWorkspaceAccountCases: + @staticmethod + @pytest.fixture + def app() -> Flask: + app = Flask(__name__) + app.config["TESTING"] = True + app.config["RESTX_MASK_HEADER"] = "X-Fields" + setattr(app, "login_manager", SimpleNamespace(load_user_from_request_context=lambda: None)) # noqa: B010 + return app + + class TestEducationApi: + def test_post_activates_education_discount(self, app: Flask) -> None: + education = MagicMock() + education.activate.return_value = AccountEducationActivation(message="success") + request_context = RequestContext( + request_id="request-1", + trace_id=None, + account_id="account-1", + active_workspace_id="workspace-1", + ) + + with ( + app.test_request_context( + "/account/education", + method="POST", + json={"token": "education-token", "institution": "Dify University", "role": "Student"}, + ), + patch( + "controllers.console.workspace.account.application_services", + return_value=SimpleNamespace(accounts=SimpleNamespace(education=education)), + ), + ): + api = EducationApi() + method = inspect.unwrap(api.post) + result = method(api, EducationActivatePayload.model_validate(request.get_json() or {}), request_context) + + assert result == {"message": "success"} + education.activate.assert_called_once_with( + request_context, + token="education-token", + institution="Dify University", + role="Student", + ) + + def test_verify_maps_rate_limit_error(self, app: Flask) -> None: + education = MagicMock() + education.verify.side_effect = account_errors.EducationRateLimitExceededError + request_context = RequestContext( + request_id="request-1", + trace_id=None, + account_id="account-1", + active_workspace_id="workspace-1", + ) + + with ( + app.test_request_context("/account/education/verify", method="GET"), + patch( + "controllers.console.workspace.account.application_services", + return_value=SimpleNamespace(accounts=SimpleNamespace(education=education)), + ), + ): + api = EducationVerifyApi() + with pytest.raises(EducationVerifyLimitError): + inspect.unwrap(api.get)(api, request_context) + + def test_post_maps_rate_limit_error(self, app: Flask) -> None: + education = MagicMock() + education.activate.side_effect = account_errors.EducationRateLimitExceededError + request_context = RequestContext( + request_id="request-1", + trace_id=None, + account_id="account-1", + active_workspace_id="workspace-1", + ) + + with ( + app.test_request_context( + "/account/education", + method="POST", + json={"token": "education-token", "institution": "Dify University", "role": "Student"}, + ), + patch( + "controllers.console.workspace.account.application_services", + return_value=SimpleNamespace(accounts=SimpleNamespace(education=education)), + ), + ): + api = EducationApi() + with pytest.raises(EducationActivateLimitError): + inspect.unwrap(api.post)( + api, EducationActivatePayload.model_validate(request.get_json() or {}), request_context + ) + + class TestChangeEmailControllers: + def test_send_delegates_parsed_request_to_application_service(self, app: Flask) -> None: + change_email = MagicMock() + change_email.send_code.return_value = "change-token" + context = _change_email_context() + payload = { + "email": "new@example.com", + "language": "zh-Hans", + "phase": "new_email", + "token": "verified-old-token", + } + + with ( + app.test_request_context( + "/account/change-email", + method="POST", + json=payload, + environ_base={"REMOTE_ADDR": "127.0.0.1"}, + ), + patch( + "controllers.console.workspace.account.application_services", + return_value=SimpleNamespace(accounts=SimpleNamespace(change_email=change_email)), + ), + ): + api = ChangeEmailSendEmailApi() + response = inspect.unwrap(api.post)( + api, ChangeEmailSendPayload.model_validate(request.get_json() or {}), context + ) + + assert response == {"result": "success", "data": "change-token"} + change_email.send_code.assert_called_once_with( + context, + requested_email="new@example.com", + language="zh-Hans", + phase="new_email", + predecessor_token="verified-old-token", + ip_address="127.0.0.1", + ) + + def test_send_maps_invalid_state_to_http_error(self, app: Flask) -> None: + change_email = MagicMock() + change_email.send_code.side_effect = account_errors.InvalidChangeEmailTokenError + + with ( + app.test_request_context( + "/account/change-email", + method="POST", + json={"email": "new@example.com", "phase": "new_email"}, + ), + patch( + "controllers.console.workspace.account.application_services", + return_value=SimpleNamespace(accounts=SimpleNamespace(change_email=change_email)), + ), + ): + api = ChangeEmailSendEmailApi() + method = inspect.unwrap(api.post) + with pytest.raises(InvalidTokenError): + method( + api, ChangeEmailSendPayload.model_validate(request.get_json() or {}), _change_email_context() + ) + + def test_validity_serializes_promoted_token(self, app: Flask) -> None: + change_email = MagicMock() + change_email.verify_code.return_value = ChangeEmailVerification( + email="new@example.com", + token="verified-token", + ) + context = _change_email_context() + + with ( + app.test_request_context( + "/account/change-email/validity", + method="POST", + json={"email": "New@Example.com", "code": "123456", "token": "pending-token"}, + ), + patch( + "controllers.console.workspace.account.application_services", + return_value=SimpleNamespace(accounts=SimpleNamespace(change_email=change_email)), + ), + ): + api = ChangeEmailCheckApi() + response = inspect.unwrap(api.post)( + api, ChangeEmailValidityPayload.model_validate(request.get_json() or {}), context + ) + + assert response == {"is_valid": True, "email": "new@example.com", "token": "verified-token"} + change_email.verify_code.assert_called_once_with( + context, + email="New@Example.com", + code="123456", + token="pending-token", + ) + + def test_reset_returns_updated_account(self, app: Flask) -> None: + change_email = MagicMock() + updated_account = _build_account("new@example.com", "acc") + change_email.reset.return_value = updated_account + context = _change_email_context() + + with ( + app.test_request_context( + "/account/change-email/reset", + method="POST", + json={"new_email": "New@Example.com", "token": "verified-token"}, + ), + patch( + "controllers.console.workspace.account.application_services", + return_value=SimpleNamespace(accounts=SimpleNamespace(change_email=change_email)), + ), + ): + api = ChangeEmailResetApi() + response = inspect.unwrap(api.post)( + api, ChangeEmailResetPayload.model_validate(request.get_json() or {}), context + ) + + assert response["email"] == "new@example.com" + change_email.reset.assert_called_once_with( + context, + new_email="New@Example.com", + token="verified-token", + ) + + def test_reset_maps_suspended_email_domain(self, app: Flask) -> None: + change_email = MagicMock() + change_email.reset.side_effect = account_errors.AccountEmailDomainSuspendedError + + with ( + app.test_request_context( + "/account/change-email/reset", + method="POST", + json={"new_email": "user@suspended.example", "token": "verified-token"}, + ), + patch( + "controllers.console.workspace.account.application_services", + return_value=SimpleNamespace(accounts=SimpleNamespace(change_email=change_email)), + ), + ): + api = ChangeEmailResetApi() + with pytest.raises(EmailDomainSuspendedError): + inspect.unwrap(api.post)( + api, ChangeEmailResetPayload.model_validate(request.get_json() or {}), _change_email_context() + ) + + class TestAccountDeletionFeedback: + def test_delegates_feedback_to_application_service(self, app: Flask) -> None: + deletion_feedback = MagicMock() + with ( + app.test_request_context( + "/account/delete/feedback", + method="POST", + json={"email": "User@Example.com", "feedback": "test"}, + ), + patch( + "controllers.console.workspace.account.application_services", + return_value=SimpleNamespace( + accounts=SimpleNamespace(deletion_feedback=deletion_feedback), + ), + ), + ): + api = AccountDeleteUpdateFeedbackApi() + method = inspect.unwrap(api.post) + response = method(api, AccountDeletionFeedbackPayload.model_validate(request.get_json() or {})) + + assert response == {"result": "success"} + deletion_feedback.submit.assert_called_once_with(email="User@Example.com", feedback="test") + + class TestCheckEmailUnique: + def test_delegates_to_email_availability_policy(self, app: Flask) -> None: + change_email = MagicMock() + + with ( + app.test_request_context( + "/account/change-email/check-email-unique", + method="POST", + json={"email": "Case@Test.com"}, + ), + patch( + "controllers.console.workspace.account.application_services", + return_value=SimpleNamespace(accounts=SimpleNamespace(change_email=change_email)), + ), + ): + api = CheckEmailUnique() + response = inspect.unwrap(api.post)( + api, CheckEmailUniquePayload.model_validate(request.get_json() or {}) + ) + + assert response == {"result": "success"} + change_email.ensure_available.assert_called_once_with("Case@Test.com") + + def test_maps_suspended_email_domain(self, app: Flask) -> None: + change_email = MagicMock() + change_email.ensure_available.side_effect = account_errors.AccountEmailDomainSuspendedError + + with ( + app.test_request_context( + "/account/change-email/check-email-unique", + method="POST", + json={"email": "user@suspended.example"}, + ), + patch( + "controllers.console.workspace.account.application_services", + return_value=SimpleNamespace(accounts=SimpleNamespace(change_email=change_email)), + ), + ): + api = CheckEmailUnique() + with pytest.raises(EmailDomainSuspendedError): + inspect.unwrap(api.post)(api, CheckEmailUniquePayload.model_validate(request.get_json() or {})) + + +def make_account( + account_id: str = "u1", *, status: _accounts_AccountStatus = _accounts_AccountStatus.ACTIVE +) -> Account: + account = Account(name="John", email=f"{account_id}@test.com", status=status) + account.id = str(uuid5(NAMESPACE_URL, f"account:{account_id}")) + account.avatar = "avatar.png" + account.interface_language = "en-US" + account.interface_theme = "light" + account.timezone = "UTC" + account.last_login_ip = "127.0.0.1" + return account + + +def persist_account_with_tenant( + session: Session, + account_id: str = "u1", + *, + status: _accounts_AccountStatus = _accounts_AccountStatus.ACTIVE, + tenant_id: str = "tenant-1", +) -> tuple[Account, Tenant]: + account = make_account(account_id, status=status) + tenant = Tenant(name=tenant_id) + tenant.id = str(uuid5(NAMESPACE_URL, f"tenant:{tenant_id}")) + membership = TenantAccountJoin( + tenant_id=tenant.id, + account_id=account.id, + current=True, + role=TenantAccountRole.OWNER, + ) + session.add_all([account, tenant, membership]) + session.commit() + account._current_tenant = tenant + account.role = TenantAccountRole.OWNER + return account, tenant + + +class TestAccountsCases: + @staticmethod + @pytest.fixture + def account_services( + account_application_services: ApplicationServices, + ) -> Iterator[AccountServices]: + """Use the real application-service composition root in controller tests.""" + with patch( + "controllers.console.workspace.account.application_services", + return_value=account_application_services, + ): + yield account_application_services.accounts + + class TestAccountInitApi: + def test_init_success(self, app: Flask, account_services: AccountServices): + api = AccountInitApi() + method = inspect.unwrap(api.post) + request_context = RequestContext( + request_id="request-1", + trace_id=None, + account_id="account-1", + active_workspace_id="workspace-1", + ) + initialize_calls: list[tuple[RequestContext, str, str, str | None]] = [] + + def initialize( + context: RequestContext, + *, + interface_language: str, + timezone: str, + invitation_code: str | None, + ) -> None: + initialize_calls.append((context, interface_language, timezone, invitation_code)) + + payload = { + "interface_language": "en-US", + "timezone": "UTC", + "invitation_code": "code123", + } + + with ( + app.test_request_context("/account/init", json=payload), + patch.object(account_services.initialization, "initialize", initialize), + ): + resp = method(api, AccountInitPayload.model_validate(payload), request_context) + + assert resp["result"] == "success" + assert initialize_calls == [(request_context, "en-US", "UTC", "code123")] + + def test_init_already_initialized(self, app: Flask, account_services: AccountServices): + api = AccountInitApi() + method = inspect.unwrap(api.post) + + request_context = RequestContext( + request_id="request-1", + trace_id=None, + account_id="account-1", + active_workspace_id="workspace-1", + ) + + def initialize(*_args: object, **_kwargs: object) -> None: + raise AccountAlreadyInitializedError + + payload = {"interface_language": "en-US", "timezone": "UTC"} + + with ( + app.test_request_context("/account/init", json=payload), + patch.object(account_services.initialization, "initialize", initialize), + ): + with pytest.raises(AccountAlreadyInitedError): + method(api, AccountInitPayload.model_validate(payload), request_context) + + def test_init_missing_invitation_code_is_mapped(self, app: Flask, account_services: AccountServices): + api = AccountInitApi() + method = inspect.unwrap(api.post) + request_context = RequestContext( + request_id="request-1", + trace_id=None, + account_id="account-1", + active_workspace_id="workspace-1", + ) + + def initialize(*_args: object, **_kwargs: object) -> None: + raise MissingInvitationCodeError("invitation_code is required") + + payload = {"interface_language": "en-US", "timezone": "UTC"} + + with ( + app.test_request_context("/account/init", json=payload), + patch.object(account_services.initialization, "initialize", initialize), + ): + with pytest.raises(MissingInvitationCodeRequestError) as exc_info: + method(api, AccountInitPayload.model_validate(payload), request_context) + + assert exc_info.value.data == { + "code": "missing_invitation_code", + "message": "Invitation code is required.", + "status": 400, + } + + class TestAccountProfileApi: + def test_get_profile_success(self, app: Flask, account_services: AccountServices): + api = AccountProfileApi() + method = inspect.unwrap(api.get) + user = make_account() + request_context = RequestContext( + request_id="request-1", + trace_id=None, + account_id=user.id, + active_workspace_id="workspace-1", + ) + get_profile = MagicMock(return_value=user) + + with ( + app.test_request_context("/account/profile"), + patch.object(account_services.profile, "get", get_profile), + ): + result = method(api, request_context) + + assert result["id"] == user.id + get_profile.assert_called_once_with(request_context) + + class TestAccountUpdateApis: + @pytest.mark.parametrize( + ("api_cls", "payload_model", "payload", "expected_changes"), + [ + (AccountNameApi, AccountNamePayload, {"name": "test"}, AccountProfileChanges(name="test")), + ( + AccountAvatarApi, + AccountAvatarPayload, + {"avatar": "img.png"}, + AccountProfileChanges(avatar="img.png"), + ), + ( + AccountInterfaceLanguageApi, + AccountInterfaceLanguagePayload, + {"interface_language": "en-US"}, + AccountProfileChanges(interface_language="en-US"), + ), + ( + AccountInterfaceThemeApi, + AccountInterfaceThemePayload, + {"interface_theme": "dark"}, + AccountProfileChanges(interface_theme="dark"), + ), + ( + AccountTimezoneApi, + AccountTimezonePayload, + {"timezone": "UTC"}, + AccountProfileChanges(timezone="UTC"), + ), + ], + ) + def test_deprecated_update_routes_delegate_to_profile_service( + self, + app: Flask, + account_services: AccountServices, + api_cls, + payload_model, + payload, + expected_changes: AccountProfileChanges, + ): + api = api_cls() + method = inspect.unwrap(api.post) + user = make_account() + request_context = RequestContext( + request_id="request-1", + trace_id=None, + account_id=user.id, + active_workspace_id="workspace-1", + ) + update_profile = MagicMock(return_value=user) + + with ( + app.test_request_context("/", json=payload), + patch.object(account_services.profile, "update", update_profile), + ): + result = method(api, payload_model.model_validate(payload), request_context) + + assert result["id"] == user.id + update_profile.assert_called_once_with(request_context, expected_changes) + + def test_deprecated_update_routes_are_marked_deprecated(self): + for api_cls in ( + AccountNameApi, + AccountAvatarApi, + AccountInterfaceLanguageApi, + AccountInterfaceThemeApi, + AccountTimezoneApi, + ): + assert api_cls.post.__apidoc__["deprecated"] is True + + class TestAccountProfilePatchApi: + def test_json_schema_matches_runtime_patch_rules(self): + schema = AccountProfilePatchPayload.model_json_schema() + validator = Draft202012Validator(schema) + + assert schema["type"] == "object" + assert schema["additionalProperties"] is False + assert "required" not in schema + assert set(schema["properties"]) == { + "name", + "avatar", + "interface_language", + "interface_theme", + "timezone", + } + validator.validate({}) + validator.validate({"name": "Jane"}) + validator.validate({"name": "Jane", "interface_language": "en-US", "timezone": "UTC"}) + for payload in ( + {"name": None}, + {"unexpected": "value"}, + {"name": "Jane", "unexpected": "value"}, + ): + assert list(validator.iter_errors(payload)) + + def test_updates_multiple_profile_fields(self, app: Flask, account_services: AccountServices): + api = AccountProfileApi() + method = inspect.unwrap(api.patch) + user = make_account() + request_context = RequestContext( + request_id="request-1", + trace_id="trace-1", + account_id=user.id, + active_workspace_id="workspace-1", + ) + update_profile = MagicMock(return_value=user) + payload = {"name": "Jane", "interface_language": "en-US", "timezone": "UTC"} + args = AccountProfilePatchPayload.model_validate(payload) + + with ( + app.test_request_context("/account/profile", method="PATCH", json=payload), + patch.object(account_services.profile, "update", update_profile), + ): + result = method(api, args, request_context) + + assert result["id"] == user.id + update_profile.assert_called_once_with( + request_context, + AccountProfileChanges(name="Jane", interface_language="en-US", timezone="UTC"), + ) + + def test_empty_patch_is_a_noop(self, app: Flask, account_services: AccountServices): + api = AccountProfileApi() + method = inspect.unwrap(api.patch) + user = make_account() + request_context = RequestContext( + request_id="request-1", + trace_id="trace-1", + account_id=user.id, + active_workspace_id="workspace-1", + ) + update_profile = MagicMock(return_value=user) + args = AccountProfilePatchPayload.model_validate({}) + + with ( + app.test_request_context("/account/profile", method="PATCH", json={}), + patch.object(account_services.profile, "update", update_profile), + ): + result = method(api, args, request_context) + + assert result["id"] == user.id + update_profile.assert_called_once_with(request_context, AccountProfileChanges()) + + @pytest.mark.parametrize("payload", [{"name": None}, {"unexpected": "value"}]) + def test_rejects_null_or_unknown_changes(self, payload: dict[str, object]): + with pytest.raises( + ValueError, + match=( + "(?:1 validation error for AccountProfilePatchPayload\\\nname|1 " + "validation error for AccountProfilePatchPayload\\\nunexpected)" + ), + ): + AccountProfilePatchPayload.model_validate(payload) + + class TestAccountAvatarApiGet: + def test_get_avatar_delegates_to_service(self, app: Flask, account_services: AccountServices): + api = AccountAvatarApi() + method = inspect.unwrap(api.get) + file_id = "550e8400-e29b-41d4-a716-446655440000" + request_context = RequestContext( + request_id="request-1", + trace_id=None, + account_id="account-1", + active_workspace_id="workspace-1", + ) + resolve_avatar = MagicMock(return_value="https://signed/example") + + with ( + app.test_request_context(f"/account/avatar?avatar={file_id}"), + patch.object(account_services.avatar, "resolve", resolve_avatar), + ): + result = method(api, AccountAvatarQuery(avatar=file_id), request_context) + + assert result == {"avatar_url": "https://signed/example"} + resolve_avatar.assert_called_once_with(request_context, file_id) + + def test_get_avatar_maps_not_found(self, app: Flask, account_services: AccountServices): + api = AccountAvatarApi() + method = inspect.unwrap(api.get) + file_id = "550e8400-e29b-41d4-a716-446655440001" + request_context = RequestContext( + request_id="request-1", + trace_id=None, + account_id="account-1", + active_workspace_id="workspace-1", + ) + resolve_avatar = MagicMock(side_effect=AvatarFileNotFoundError) + + with ( + app.test_request_context(f"/account/avatar?avatar={file_id}"), + patch.object(account_services.avatar, "resolve", resolve_avatar), + ): + with pytest.raises(NotFound): + method(api, AccountAvatarQuery(avatar=file_id), request_context) + + def test_get_avatar_missing_query_returns_unprocessable_entity(self, app: Flask): + account = make_account() + + with ( + app.test_request_context("/account/avatar"), + patch("controllers.console.wraps._is_setup_completed", return_value=True), + config_overrides_context(LOGIN_DISABLED=True), + patch( + "controllers.console.wraps.current_account_with_tenant", + return_value=(account, "workspace-1"), + ), + patch( + "controllers.console.flask_admission.current_account_with_tenant", + return_value=AccountWithTenant(account=account, tenant_id="workspace-1"), + ), + ): + with pytest.raises(UnprocessableEntity) as exc_info: + AccountAvatarApi().get() + + assert exc_info.value.code == 422 + + class TestConvertedPostDecorator: + def test_rejects_an_invalid_body_through_the_decorator(self, app: Flask): + """The decorator validates the JSON body before the view runs, for the POST handlers too.""" + account = make_account() + + with ( + app.test_request_context("/account/name", method="POST", json={}), + patch("controllers.console.wraps._is_setup_completed", return_value=True), + config_overrides_context(LOGIN_DISABLED=True), + patch( + "controllers.console.wraps.current_account_with_tenant", + return_value=(account, "workspace-1"), + ), + patch( + "controllers.console.flask_admission.current_account_with_tenant", + return_value=AccountWithTenant(account=account, tenant_id="workspace-1"), + ), + ): + with pytest.raises(UnprocessableEntity) as exc_info: + AccountNameApi().post() + + assert exc_info.value.code == 422 + + class TestAccountPasswordApi: + def test_password_success(self, app: Flask, account_services: AccountServices): + api = AccountPasswordApi() + method = inspect.unwrap(api.post) + + payload = { + "password": "old", + "new_password": "new123", + "repeat_new_password": "new123", + } + + user = make_account() + request_context = RequestContext( + request_id="request-1", + trace_id=None, + account_id=user.id, + active_workspace_id="workspace-1", + ) + change_password = MagicMock(return_value=user) + + with ( + app.test_request_context("/", json=payload), + patch.object(account_services.password, "change", change_password), + ): + result = method(api, AccountPasswordPayload.model_validate(payload), request_context) + + assert result["id"] == user.id + change_password.assert_called_once_with( + request_context, + current_password="old", + new_password="new123", + ) + + def test_password_wrong_current(self, app: Flask, account_services: AccountServices): + api = AccountPasswordApi() + method = inspect.unwrap(api.post) + + payload = { + "password": "bad", + "new_password": "new123", + "repeat_new_password": "new123", + } + user = make_account() + request_context = RequestContext( + request_id="request-1", + trace_id=None, + account_id=user.id, + active_workspace_id="workspace-1", + ) + change_password = MagicMock(side_effect=CurrentAccountPasswordIncorrectError) + + with ( + app.test_request_context("/", json=payload), + patch.object(account_services.password, "change", change_password), + ): + with pytest.raises(CurrentPasswordIncorrectError): + method(api, AccountPasswordPayload.model_validate(payload), request_context) + + def test_password_policy_error_is_mapped(self, app: Flask, account_services: AccountServices): + api = AccountPasswordApi() + method = inspect.unwrap(api.post) + payload = { + "password": "old", + "new_password": "letters-only", + "repeat_new_password": "letters-only", + } + request_context = RequestContext( + request_id="request-1", + trace_id=None, + account_id="account-1", + active_workspace_id="workspace-1", + ) + change_password = MagicMock( + side_effect=InvalidAccountPasswordError( + "Password must contain letters and numbers, and the length must be at least 8 characters." + ) + ) + + with ( + app.test_request_context("/", json=payload), + patch.object(account_services.password, "change", change_password), + ): + with pytest.raises(InvalidAccountPasswordRequestError) as exc_info: + method(api, AccountPasswordPayload.model_validate(payload), request_context) + + assert exc_info.value.data == { + "code": "invalid_account_password", + "message": "Password must contain letters and numbers, and the length must be at least 8 characters.", + "status": 400, + } + + class TestAccountIntegrateApi: + def test_get_integrates(self, app: Flask, account_services: AccountServices): + api = AccountIntegrateApi() + method = inspect.unwrap(api.get) + request_context = RequestContext( + request_id="request-1", + trace_id=None, + account_id="account-1", + active_workspace_id="workspace-1", + ) + list_integrations = MagicMock( + return_value=[ + AccountIntegrationStatus(provider="github", created_at=datetime(2026, 1, 1), is_bound=True), + AccountIntegrationStatus(provider="google", created_at=None, is_bound=False), + ] + ) + + with ( + app.test_request_context("/"), + patch.object(account_services.integrations, "list", list_integrations), + ): + result = method(api, request_context) + + list_integrations.assert_called_once_with(request_context) + assert result["data"][0]["provider"] == "github" + assert result["data"][0]["is_bound"] is True + assert result["data"][0]["link"] is None + assert result["data"][1]["provider"] == "google" + assert result["data"][1]["is_bound"] is False + assert result["data"][1]["link"].endswith("/console/api/oauth/login/google") + + class TestAccountDeleteApi: + def test_delete_verify_success(self, app: Flask, account_services: AccountServices): + api = AccountDeleteVerifyApi() + method = inspect.unwrap(api.get) + request_context = RequestContext( + request_id="request-1", + trace_id=None, + account_id="account-1", + active_workspace_id="workspace-1", + ) + issue_verification = MagicMock(return_value="token") + + with ( + app.test_request_context("/"), + patch.object(account_services.deletion, "issue_verification", issue_verification), + ): + result = method(api, request_context) + + assert result["result"] == "success" + assert result["data"] == "token" + issue_verification.assert_called_once_with(request_context) + + def test_delete_invalid_code(self, app: Flask, account_services: AccountServices): + api = AccountDeleteApi() + method = inspect.unwrap(api.post) + + payload = {"token": "t", "code": "x"} + request_context = RequestContext( + request_id="request-1", + trace_id=None, + account_id="account-1", + active_workspace_id="workspace-1", + ) + request_deletion = MagicMock(side_effect=InvalidAccountDeletionVerificationError) + + with ( + app.test_request_context("/", json=payload), + patch.object(account_services.deletion, "request_deletion", request_deletion), + ): + with pytest.raises(InvalidAccountDeletionCodeError): + method(api, AccountDeletePayload.model_validate(payload), request_context) + + def test_delete_verify_maps_rate_limit(self, app: Flask, account_services: AccountServices): + api = AccountDeleteVerifyApi() + method = inspect.unwrap(api.get) + request_context = RequestContext( + request_id="request-1", + trace_id=None, + account_id="account-1", + active_workspace_id="workspace-1", + ) + issue_verification = MagicMock(side_effect=AccountDeletionRateLimitError(1)) + + with ( + app.test_request_context("/"), + patch.object(account_services.deletion, "issue_verification", issue_verification), + pytest.raises(EmailCodeAccountDeletionRateLimitExceededError), + ): + method(api, request_context) + + def test_delete_success(self, app: Flask, account_services: AccountServices): + api = AccountDeleteApi() + method = inspect.unwrap(api.post) + request_context = RequestContext( + request_id="request-1", + trace_id=None, + account_id="account-1", + active_workspace_id="workspace-1", + ) + request_deletion = MagicMock() + payload = {"token": "token", "code": "123456"} + + with ( + app.test_request_context("/", json=payload), + patch.object(account_services.deletion, "request_deletion", request_deletion), + ): + result = method(api, AccountDeletePayload.model_validate(payload), request_context) + + assert result["result"] == "success" + request_deletion.assert_called_once_with(request_context, token="token", code="123456") + + class TestChangeEmailApis: + def test_check_email_code_invalid(self, app: Flask, account_services: AccountServices): + api = ChangeEmailCheckApi() + method = inspect.unwrap(api.post) + + payload = {"email": "a@test.com", "code": "x", "token": "t"} + user = make_account("acc-1") + request_context = RequestContext( + request_id="request-1", + trace_id=None, + account_id=user.id, + active_workspace_id="workspace-1", + ) + verify_code = MagicMock(side_effect=InvalidChangeEmailCodeError) + + with ( + app.test_request_context("/", json=payload), + patch.object( + type(console_ns), + "payload", + new_callable=PropertyMock, + return_value=payload, + ), + patch.object(account_services.change_email, "verify_code", verify_code), + ): + with pytest.raises(EmailCodeError): + method(api, ChangeEmailValidityPayload.model_validate(payload), request_context) + + def test_reset_email_already_used(self, app: Flask, account_services: AccountServices): + api = ChangeEmailResetApi() + method = inspect.unwrap(api.post) + + payload = {"new_email": "x@test.com", "token": "t"} + user = make_account() + request_context = RequestContext( + request_id="request-1", + trace_id=None, + account_id=user.id, + active_workspace_id="workspace-1", + ) + reset_email = MagicMock(side_effect=AccountEmailAlreadyInUseError) + + with ( + app.test_request_context("/", json=payload), + patch.object( + type(console_ns), + "payload", + new_callable=PropertyMock, + return_value=payload, + ), + patch.object(account_services.change_email, "reset", reset_email), + ): + with pytest.raises(EmailAlreadyInUseError): + method(api, ChangeEmailResetPayload.model_validate(payload), request_context) + + class TestCheckEmailUniqueApi: + def test_email_unique_success(self, app: Flask, account_services: AccountServices): + api = CheckEmailUnique() + method = inspect.unwrap(api.post) + + payload = {"email": "ok@test.com"} + ensure_available = MagicMock() + + with ( + app.test_request_context("/", json=payload), + patch.object( + type(console_ns), + "payload", + new_callable=PropertyMock, + return_value=payload, + ), + patch.object(account_services.change_email, "ensure_available", ensure_available), + ): + result = method(api, CheckEmailUniquePayload.model_validate(payload)) + + assert result["result"] == "success" + + def test_email_in_freeze(self, app: Flask, account_services: AccountServices): + api = CheckEmailUnique() + method = inspect.unwrap(api.post) + + payload = {"email": "x@test.com"} + ensure_available = MagicMock(side_effect=AccountEmailFrozenError) + + with ( + app.test_request_context("/", json=payload), + patch.object( + type(console_ns), + "payload", + new_callable=PropertyMock, + return_value=payload, + ), + patch.object(account_services.change_email, "ensure_available", ensure_available), + ): + with pytest.raises(AccountInFreezeError): + method(api, CheckEmailUniquePayload.model_validate(payload)) + + def test_email_domain_is_suspended(self, app: Flask, account_services: AccountServices): + api = CheckEmailUnique() + method = inspect.unwrap(api.post) + + payload = {"email": "user@suspended.example"} + ensure_available = MagicMock(side_effect=AccountEmailDomainSuspendedError) + + with ( + app.test_request_context("/", json=payload), + patch.object( + type(console_ns), + "payload", + new_callable=PropertyMock, + return_value=payload, + ), + patch.object(account_services.change_email, "ensure_available", ensure_available), + ): + with pytest.raises(EmailDomainSuspendedError): + method(api, CheckEmailUniquePayload.model_validate(payload)) diff --git a/api/tests/unit_tests/controllers/console/workspace/test_accounts.py b/api/tests/unit_tests/controllers/console/workspace/test_accounts.py deleted file mode 100644 index 4b8bbc74a39..00000000000 --- a/api/tests/unit_tests/controllers/console/workspace/test_accounts.py +++ /dev/null @@ -1,794 +0,0 @@ -import inspect -from collections.abc import Iterator -from datetime import datetime -from unittest.mock import MagicMock, PropertyMock, patch -from uuid import NAMESPACE_URL, uuid5 - -import pytest -from flask import Flask -from jsonschema import Draft202012Validator -from sqlalchemy.orm import Session -from werkzeug.exceptions import NotFound, UnprocessableEntity - -from controllers.console import console_ns -from controllers.console.auth.error import ( - EmailAlreadyInUseError, - EmailCodeAccountDeletionRateLimitExceededError, - EmailCodeError, -) -from controllers.console.error import ( - AccountInFreezeError, - EmailDomainSuspendedError, - InvalidAccountPasswordRequestError, -) -from controllers.console.workspace.account import ( - AccountAvatarApi, - AccountAvatarPayload, - AccountAvatarQuery, - AccountDeleteApi, - AccountDeletePayload, - AccountDeleteVerifyApi, - AccountInitApi, - AccountInitPayload, - AccountIntegrateApi, - AccountInterfaceLanguageApi, - AccountInterfaceLanguagePayload, - AccountInterfaceThemeApi, - AccountInterfaceThemePayload, - AccountNameApi, - AccountNamePayload, - AccountPasswordApi, - AccountPasswordPayload, - AccountProfileApi, - AccountProfilePatchPayload, - AccountTimezoneApi, - AccountTimezonePayload, - ChangeEmailCheckApi, - ChangeEmailResetApi, - ChangeEmailResetPayload, - ChangeEmailValidityPayload, - CheckEmailUnique, - CheckEmailUniquePayload, -) -from controllers.console.workspace.error import ( - AccountAlreadyInitedError, - CurrentPasswordIncorrectError, - InvalidAccountDeletionCodeError, - MissingInvitationCodeRequestError, -) -from extensions.application_services.account import AccountServices -from extensions.ext_application_services import ApplicationServices -from libs.login import AccountWithTenant -from machinery.context import RequestContext -from models import Account, Tenant, TenantAccountJoin -from models.account import AccountStatus, TenantAccountRole -from services.account_errors import ( - AccountAlreadyInitializedError, - AccountDeletionRateLimitError, - AccountEmailAlreadyInUseError, - AccountEmailDomainSuspendedError, - AccountEmailFrozenError, - AvatarFileNotFoundError, - CurrentAccountPasswordIncorrectError, - InvalidAccountDeletionVerificationError, - InvalidAccountPasswordError, - InvalidChangeEmailCodeError, - MissingInvitationCodeError, -) -from services.entities.account_entities import AccountIntegrationStatus, AccountProfileChanges -from tests.unit_tests.config_override import config_overrides_context - - -def make_account(account_id: str = "u1", *, status: AccountStatus = AccountStatus.ACTIVE) -> Account: - account = Account(name="John", email=f"{account_id}@test.com", status=status) - account.id = str(uuid5(NAMESPACE_URL, f"account:{account_id}")) - account.avatar = "avatar.png" - account.interface_language = "en-US" - account.interface_theme = "light" - account.timezone = "UTC" - account.last_login_ip = "127.0.0.1" - return account - - -def persist_account_with_tenant( - session: Session, - account_id: str = "u1", - *, - status: AccountStatus = AccountStatus.ACTIVE, - tenant_id: str = "tenant-1", -) -> tuple[Account, Tenant]: - account = make_account(account_id, status=status) - tenant = Tenant(name=tenant_id) - tenant.id = str(uuid5(NAMESPACE_URL, f"tenant:{tenant_id}")) - membership = TenantAccountJoin( - tenant_id=tenant.id, - account_id=account.id, - current=True, - role=TenantAccountRole.OWNER, - ) - session.add_all([account, tenant, membership]) - session.commit() - account._current_tenant = tenant - account.role = TenantAccountRole.OWNER - return account, tenant - - -@pytest.fixture -def account_services( - account_application_services: ApplicationServices, -) -> Iterator[AccountServices]: - """Use the real application-service composition root in controller tests.""" - with patch( - "controllers.console.workspace.account.application_services", - return_value=account_application_services, - ): - yield account_application_services.accounts - - -class TestAccountInitApi: - def test_init_success(self, app: Flask, account_services: AccountServices): - api = AccountInitApi() - method = inspect.unwrap(api.post) - request_context = RequestContext( - request_id="request-1", - trace_id=None, - account_id="account-1", - active_workspace_id="workspace-1", - ) - initialize_calls: list[tuple[RequestContext, str, str, str | None]] = [] - - def initialize( - context: RequestContext, - *, - interface_language: str, - timezone: str, - invitation_code: str | None, - ) -> None: - initialize_calls.append((context, interface_language, timezone, invitation_code)) - - payload = { - "interface_language": "en-US", - "timezone": "UTC", - "invitation_code": "code123", - } - - with ( - app.test_request_context("/account/init", json=payload), - patch.object(account_services.initialization, "initialize", initialize), - ): - resp = method(api, AccountInitPayload.model_validate(payload), request_context) - - assert resp["result"] == "success" - assert initialize_calls == [(request_context, "en-US", "UTC", "code123")] - - def test_init_already_initialized(self, app: Flask, account_services: AccountServices): - api = AccountInitApi() - method = inspect.unwrap(api.post) - - request_context = RequestContext( - request_id="request-1", - trace_id=None, - account_id="account-1", - active_workspace_id="workspace-1", - ) - - def initialize(*_args: object, **_kwargs: object) -> None: - raise AccountAlreadyInitializedError - - payload = {"interface_language": "en-US", "timezone": "UTC"} - - with ( - app.test_request_context("/account/init", json=payload), - patch.object(account_services.initialization, "initialize", initialize), - ): - with pytest.raises(AccountAlreadyInitedError): - method(api, AccountInitPayload.model_validate(payload), request_context) - - def test_init_missing_invitation_code_is_mapped(self, app: Flask, account_services: AccountServices): - api = AccountInitApi() - method = inspect.unwrap(api.post) - request_context = RequestContext( - request_id="request-1", - trace_id=None, - account_id="account-1", - active_workspace_id="workspace-1", - ) - - def initialize(*_args: object, **_kwargs: object) -> None: - raise MissingInvitationCodeError("invitation_code is required") - - payload = {"interface_language": "en-US", "timezone": "UTC"} - - with ( - app.test_request_context("/account/init", json=payload), - patch.object(account_services.initialization, "initialize", initialize), - ): - with pytest.raises(MissingInvitationCodeRequestError) as exc_info: - method(api, AccountInitPayload.model_validate(payload), request_context) - - assert exc_info.value.data == { - "code": "missing_invitation_code", - "message": "Invitation code is required.", - "status": 400, - } - - -class TestAccountProfileApi: - def test_get_profile_success(self, app: Flask, account_services: AccountServices): - api = AccountProfileApi() - method = inspect.unwrap(api.get) - user = make_account() - request_context = RequestContext( - request_id="request-1", - trace_id=None, - account_id=user.id, - active_workspace_id="workspace-1", - ) - get_profile = MagicMock(return_value=user) - - with ( - app.test_request_context("/account/profile"), - patch.object(account_services.profile, "get", get_profile), - ): - result = method(api, request_context) - - assert result["id"] == user.id - get_profile.assert_called_once_with(request_context) - - -class TestAccountUpdateApis: - @pytest.mark.parametrize( - ("api_cls", "payload_model", "payload", "expected_changes"), - [ - (AccountNameApi, AccountNamePayload, {"name": "test"}, AccountProfileChanges(name="test")), - (AccountAvatarApi, AccountAvatarPayload, {"avatar": "img.png"}, AccountProfileChanges(avatar="img.png")), - ( - AccountInterfaceLanguageApi, - AccountInterfaceLanguagePayload, - {"interface_language": "en-US"}, - AccountProfileChanges(interface_language="en-US"), - ), - ( - AccountInterfaceThemeApi, - AccountInterfaceThemePayload, - {"interface_theme": "dark"}, - AccountProfileChanges(interface_theme="dark"), - ), - (AccountTimezoneApi, AccountTimezonePayload, {"timezone": "UTC"}, AccountProfileChanges(timezone="UTC")), - ], - ) - def test_deprecated_update_routes_delegate_to_profile_service( - self, - app: Flask, - account_services: AccountServices, - api_cls, - payload_model, - payload, - expected_changes: AccountProfileChanges, - ): - api = api_cls() - method = inspect.unwrap(api.post) - user = make_account() - request_context = RequestContext( - request_id="request-1", - trace_id=None, - account_id=user.id, - active_workspace_id="workspace-1", - ) - update_profile = MagicMock(return_value=user) - - with ( - app.test_request_context("/", json=payload), - patch.object(account_services.profile, "update", update_profile), - ): - result = method(api, payload_model.model_validate(payload), request_context) - - assert result["id"] == user.id - update_profile.assert_called_once_with(request_context, expected_changes) - - def test_deprecated_update_routes_are_marked_deprecated(self): - for api_cls in ( - AccountNameApi, - AccountAvatarApi, - AccountInterfaceLanguageApi, - AccountInterfaceThemeApi, - AccountTimezoneApi, - ): - assert api_cls.post.__apidoc__["deprecated"] is True - - -class TestAccountProfilePatchApi: - def test_json_schema_matches_runtime_patch_rules(self): - schema = AccountProfilePatchPayload.model_json_schema() - validator = Draft202012Validator(schema) - - assert schema["type"] == "object" - assert schema["additionalProperties"] is False - assert "required" not in schema - assert set(schema["properties"]) == { - "name", - "avatar", - "interface_language", - "interface_theme", - "timezone", - } - validator.validate({}) - validator.validate({"name": "Jane"}) - validator.validate({"name": "Jane", "interface_language": "en-US", "timezone": "UTC"}) - for payload in ( - {"name": None}, - {"unexpected": "value"}, - {"name": "Jane", "unexpected": "value"}, - ): - assert list(validator.iter_errors(payload)) - - def test_updates_multiple_profile_fields(self, app: Flask, account_services: AccountServices): - api = AccountProfileApi() - method = inspect.unwrap(api.patch) - user = make_account() - request_context = RequestContext( - request_id="request-1", - trace_id="trace-1", - account_id=user.id, - active_workspace_id="workspace-1", - ) - update_profile = MagicMock(return_value=user) - payload = {"name": "Jane", "interface_language": "en-US", "timezone": "UTC"} - args = AccountProfilePatchPayload.model_validate(payload) - - with ( - app.test_request_context("/account/profile", method="PATCH", json=payload), - patch.object(account_services.profile, "update", update_profile), - ): - result = method(api, args, request_context) - - assert result["id"] == user.id - update_profile.assert_called_once_with( - request_context, - AccountProfileChanges(name="Jane", interface_language="en-US", timezone="UTC"), - ) - - def test_empty_patch_is_a_noop(self, app: Flask, account_services: AccountServices): - api = AccountProfileApi() - method = inspect.unwrap(api.patch) - user = make_account() - request_context = RequestContext( - request_id="request-1", - trace_id="trace-1", - account_id=user.id, - active_workspace_id="workspace-1", - ) - update_profile = MagicMock(return_value=user) - args = AccountProfilePatchPayload.model_validate({}) - - with ( - app.test_request_context("/account/profile", method="PATCH", json={}), - patch.object(account_services.profile, "update", update_profile), - ): - result = method(api, args, request_context) - - assert result["id"] == user.id - update_profile.assert_called_once_with(request_context, AccountProfileChanges()) - - @pytest.mark.parametrize("payload", [{"name": None}, {"unexpected": "value"}]) - def test_rejects_null_or_unknown_changes(self, payload: dict[str, object]): - with pytest.raises( - ValueError, - match=( - "(?:1 validation error for AccountProfilePatchPayload\\\nname|1 " - "validation error for AccountProfilePatchPayload\\\nunexpected)" - ), - ): - AccountProfilePatchPayload.model_validate(payload) - - -class TestAccountAvatarApiGet: - def test_get_avatar_delegates_to_service(self, app: Flask, account_services: AccountServices): - api = AccountAvatarApi() - method = inspect.unwrap(api.get) - file_id = "550e8400-e29b-41d4-a716-446655440000" - request_context = RequestContext( - request_id="request-1", - trace_id=None, - account_id="account-1", - active_workspace_id="workspace-1", - ) - resolve_avatar = MagicMock(return_value="https://signed/example") - - with ( - app.test_request_context(f"/account/avatar?avatar={file_id}"), - patch.object(account_services.avatar, "resolve", resolve_avatar), - ): - result = method(api, AccountAvatarQuery(avatar=file_id), request_context) - - assert result == {"avatar_url": "https://signed/example"} - resolve_avatar.assert_called_once_with(request_context, file_id) - - def test_get_avatar_maps_not_found(self, app: Flask, account_services: AccountServices): - api = AccountAvatarApi() - method = inspect.unwrap(api.get) - file_id = "550e8400-e29b-41d4-a716-446655440001" - request_context = RequestContext( - request_id="request-1", - trace_id=None, - account_id="account-1", - active_workspace_id="workspace-1", - ) - resolve_avatar = MagicMock(side_effect=AvatarFileNotFoundError) - - with ( - app.test_request_context(f"/account/avatar?avatar={file_id}"), - patch.object(account_services.avatar, "resolve", resolve_avatar), - ): - with pytest.raises(NotFound): - method(api, AccountAvatarQuery(avatar=file_id), request_context) - - def test_get_avatar_missing_query_returns_unprocessable_entity(self, app: Flask): - account = make_account() - - with ( - app.test_request_context("/account/avatar"), - patch("controllers.console.wraps._is_setup_completed", return_value=True), - config_overrides_context(LOGIN_DISABLED=True), - patch( - "controllers.console.wraps.current_account_with_tenant", - return_value=(account, "workspace-1"), - ), - patch( - "controllers.console.flask_admission.current_account_with_tenant", - return_value=AccountWithTenant(account=account, tenant_id="workspace-1"), - ), - ): - with pytest.raises(UnprocessableEntity) as exc_info: - AccountAvatarApi().get() - - assert exc_info.value.code == 422 - - -class TestConvertedPostDecorator: - def test_rejects_an_invalid_body_through_the_decorator(self, app: Flask): - """The decorator validates the JSON body before the view runs, for the POST handlers too.""" - account = make_account() - - with ( - app.test_request_context("/account/name", method="POST", json={}), - patch("controllers.console.wraps._is_setup_completed", return_value=True), - config_overrides_context(LOGIN_DISABLED=True), - patch( - "controllers.console.wraps.current_account_with_tenant", - return_value=(account, "workspace-1"), - ), - patch( - "controllers.console.flask_admission.current_account_with_tenant", - return_value=AccountWithTenant(account=account, tenant_id="workspace-1"), - ), - ): - with pytest.raises(UnprocessableEntity) as exc_info: - AccountNameApi().post() - - assert exc_info.value.code == 422 - - -class TestAccountPasswordApi: - def test_password_success(self, app: Flask, account_services: AccountServices): - api = AccountPasswordApi() - method = inspect.unwrap(api.post) - - payload = { - "password": "old", - "new_password": "new123", - "repeat_new_password": "new123", - } - - user = make_account() - request_context = RequestContext( - request_id="request-1", - trace_id=None, - account_id=user.id, - active_workspace_id="workspace-1", - ) - change_password = MagicMock(return_value=user) - - with ( - app.test_request_context("/", json=payload), - patch.object(account_services.password, "change", change_password), - ): - result = method(api, AccountPasswordPayload.model_validate(payload), request_context) - - assert result["id"] == user.id - change_password.assert_called_once_with( - request_context, - current_password="old", - new_password="new123", - ) - - def test_password_wrong_current(self, app: Flask, account_services: AccountServices): - api = AccountPasswordApi() - method = inspect.unwrap(api.post) - - payload = { - "password": "bad", - "new_password": "new123", - "repeat_new_password": "new123", - } - user = make_account() - request_context = RequestContext( - request_id="request-1", - trace_id=None, - account_id=user.id, - active_workspace_id="workspace-1", - ) - change_password = MagicMock(side_effect=CurrentAccountPasswordIncorrectError) - - with ( - app.test_request_context("/", json=payload), - patch.object(account_services.password, "change", change_password), - ): - with pytest.raises(CurrentPasswordIncorrectError): - method(api, AccountPasswordPayload.model_validate(payload), request_context) - - def test_password_policy_error_is_mapped(self, app: Flask, account_services: AccountServices): - api = AccountPasswordApi() - method = inspect.unwrap(api.post) - payload = { - "password": "old", - "new_password": "letters-only", - "repeat_new_password": "letters-only", - } - request_context = RequestContext( - request_id="request-1", - trace_id=None, - account_id="account-1", - active_workspace_id="workspace-1", - ) - change_password = MagicMock( - side_effect=InvalidAccountPasswordError( - "Password must contain letters and numbers, and the length must be at least 8 characters." - ) - ) - - with ( - app.test_request_context("/", json=payload), - patch.object(account_services.password, "change", change_password), - ): - with pytest.raises(InvalidAccountPasswordRequestError) as exc_info: - method(api, AccountPasswordPayload.model_validate(payload), request_context) - - assert exc_info.value.data == { - "code": "invalid_account_password", - "message": "Password must contain letters and numbers, and the length must be at least 8 characters.", - "status": 400, - } - - -class TestAccountIntegrateApi: - def test_get_integrates(self, app: Flask, account_services: AccountServices): - api = AccountIntegrateApi() - method = inspect.unwrap(api.get) - request_context = RequestContext( - request_id="request-1", - trace_id=None, - account_id="account-1", - active_workspace_id="workspace-1", - ) - list_integrations = MagicMock( - return_value=[ - AccountIntegrationStatus(provider="github", created_at=datetime(2026, 1, 1), is_bound=True), - AccountIntegrationStatus(provider="google", created_at=None, is_bound=False), - ] - ) - - with ( - app.test_request_context("/"), - patch.object(account_services.integrations, "list", list_integrations), - ): - result = method(api, request_context) - - list_integrations.assert_called_once_with(request_context) - assert result["data"][0]["provider"] == "github" - assert result["data"][0]["is_bound"] is True - assert result["data"][0]["link"] is None - assert result["data"][1]["provider"] == "google" - assert result["data"][1]["is_bound"] is False - assert result["data"][1]["link"].endswith("/console/api/oauth/login/google") - - -class TestAccountDeleteApi: - def test_delete_verify_success(self, app: Flask, account_services: AccountServices): - api = AccountDeleteVerifyApi() - method = inspect.unwrap(api.get) - request_context = RequestContext( - request_id="request-1", - trace_id=None, - account_id="account-1", - active_workspace_id="workspace-1", - ) - issue_verification = MagicMock(return_value="token") - - with ( - app.test_request_context("/"), - patch.object(account_services.deletion, "issue_verification", issue_verification), - ): - result = method(api, request_context) - - assert result["result"] == "success" - assert result["data"] == "token" - issue_verification.assert_called_once_with(request_context) - - def test_delete_invalid_code(self, app: Flask, account_services: AccountServices): - api = AccountDeleteApi() - method = inspect.unwrap(api.post) - - payload = {"token": "t", "code": "x"} - request_context = RequestContext( - request_id="request-1", - trace_id=None, - account_id="account-1", - active_workspace_id="workspace-1", - ) - request_deletion = MagicMock(side_effect=InvalidAccountDeletionVerificationError) - - with ( - app.test_request_context("/", json=payload), - patch.object(account_services.deletion, "request_deletion", request_deletion), - ): - with pytest.raises(InvalidAccountDeletionCodeError): - method(api, AccountDeletePayload.model_validate(payload), request_context) - - def test_delete_verify_maps_rate_limit(self, app: Flask, account_services: AccountServices): - api = AccountDeleteVerifyApi() - method = inspect.unwrap(api.get) - request_context = RequestContext( - request_id="request-1", - trace_id=None, - account_id="account-1", - active_workspace_id="workspace-1", - ) - issue_verification = MagicMock(side_effect=AccountDeletionRateLimitError(1)) - - with ( - app.test_request_context("/"), - patch.object(account_services.deletion, "issue_verification", issue_verification), - pytest.raises(EmailCodeAccountDeletionRateLimitExceededError), - ): - method(api, request_context) - - def test_delete_success(self, app: Flask, account_services: AccountServices): - api = AccountDeleteApi() - method = inspect.unwrap(api.post) - request_context = RequestContext( - request_id="request-1", - trace_id=None, - account_id="account-1", - active_workspace_id="workspace-1", - ) - request_deletion = MagicMock() - payload = {"token": "token", "code": "123456"} - - with ( - app.test_request_context("/", json=payload), - patch.object(account_services.deletion, "request_deletion", request_deletion), - ): - result = method(api, AccountDeletePayload.model_validate(payload), request_context) - - assert result["result"] == "success" - request_deletion.assert_called_once_with(request_context, token="token", code="123456") - - -class TestChangeEmailApis: - def test_check_email_code_invalid(self, app: Flask, account_services: AccountServices): - api = ChangeEmailCheckApi() - method = inspect.unwrap(api.post) - - payload = {"email": "a@test.com", "code": "x", "token": "t"} - user = make_account("acc-1") - request_context = RequestContext( - request_id="request-1", - trace_id=None, - account_id=user.id, - active_workspace_id="workspace-1", - ) - verify_code = MagicMock(side_effect=InvalidChangeEmailCodeError) - - with ( - app.test_request_context("/", json=payload), - patch.object( - type(console_ns), - "payload", - new_callable=PropertyMock, - return_value=payload, - ), - patch.object(account_services.change_email, "verify_code", verify_code), - ): - with pytest.raises(EmailCodeError): - method(api, ChangeEmailValidityPayload.model_validate(payload), request_context) - - def test_reset_email_already_used(self, app: Flask, account_services: AccountServices): - api = ChangeEmailResetApi() - method = inspect.unwrap(api.post) - - payload = {"new_email": "x@test.com", "token": "t"} - user = make_account() - request_context = RequestContext( - request_id="request-1", - trace_id=None, - account_id=user.id, - active_workspace_id="workspace-1", - ) - reset_email = MagicMock(side_effect=AccountEmailAlreadyInUseError) - - with ( - app.test_request_context("/", json=payload), - patch.object( - type(console_ns), - "payload", - new_callable=PropertyMock, - return_value=payload, - ), - patch.object(account_services.change_email, "reset", reset_email), - ): - with pytest.raises(EmailAlreadyInUseError): - method(api, ChangeEmailResetPayload.model_validate(payload), request_context) - - -class TestCheckEmailUniqueApi: - def test_email_unique_success(self, app: Flask, account_services: AccountServices): - api = CheckEmailUnique() - method = inspect.unwrap(api.post) - - payload = {"email": "ok@test.com"} - ensure_available = MagicMock() - - with ( - app.test_request_context("/", json=payload), - patch.object( - type(console_ns), - "payload", - new_callable=PropertyMock, - return_value=payload, - ), - patch.object(account_services.change_email, "ensure_available", ensure_available), - ): - result = method(api, CheckEmailUniquePayload.model_validate(payload)) - - assert result["result"] == "success" - - def test_email_in_freeze(self, app: Flask, account_services: AccountServices): - api = CheckEmailUnique() - method = inspect.unwrap(api.post) - - payload = {"email": "x@test.com"} - ensure_available = MagicMock(side_effect=AccountEmailFrozenError) - - with ( - app.test_request_context("/", json=payload), - patch.object( - type(console_ns), - "payload", - new_callable=PropertyMock, - return_value=payload, - ), - patch.object(account_services.change_email, "ensure_available", ensure_available), - ): - with pytest.raises(AccountInFreezeError): - method(api, CheckEmailUniquePayload.model_validate(payload)) - - def test_email_domain_is_suspended(self, app: Flask, account_services: AccountServices): - api = CheckEmailUnique() - method = inspect.unwrap(api.post) - - payload = {"email": "user@suspended.example"} - ensure_available = MagicMock(side_effect=AccountEmailDomainSuspendedError) - - with ( - app.test_request_context("/", json=payload), - patch.object( - type(console_ns), - "payload", - new_callable=PropertyMock, - return_value=payload, - ), - patch.object(account_services.change_email, "ensure_available", ensure_available), - ): - with pytest.raises(EmailDomainSuspendedError): - method(api, CheckEmailUniquePayload.model_validate(payload)) diff --git a/api/tests/unit_tests/controllers/console/workspace/test_agent_providers.py b/api/tests/unit_tests/controllers/console/workspace/test_agent_providers.py index 47d7d00bab4..5b1e137162e 100644 --- a/api/tests/unit_tests/controllers/console/workspace/test_agent_providers.py +++ b/api/tests/unit_tests/controllers/console/workspace/test_agent_providers.py @@ -15,7 +15,7 @@ from core.agent.plugin_entities import AgentStrategyParameter from core.plugin.entities.plugin_daemon import PluginAgentProviderEntity from core.plugin.impl.exc import PluginNotFoundError from models.account import Account -from tests.unit_tests.model_factories import make_account +from tests.helpers.unit.model_factories import make_account def _account() -> Account: diff --git a/api/tests/unit_tests/controllers/console/workspace/test_endpoint.py b/api/tests/unit_tests/controllers/console/workspace/test_endpoint.py index 1f5abb0fc5c..95759dbd286 100644 --- a/api/tests/unit_tests/controllers/console/workspace/test_endpoint.py +++ b/api/tests/unit_tests/controllers/console/workspace/test_endpoint.py @@ -33,7 +33,7 @@ from core.entities.provider_entities import ProviderConfig, ProviderConfigType from core.plugin.entities.endpoint import EndpointEntityWithInstance, EndpointProviderDeclaration from core.plugin.impl.exc import PluginDaemonInternalServerError, PluginPermissionDeniedError from libs.external_api import ExternalApi -from tests.unit_tests.controllers.rbac_introspection import rbac_checks +from tests.helpers.unit.controllers.rbac_introspection import rbac_checks def _endpoint_entity() -> EndpointEntityWithInstance: diff --git a/api/tests/unit_tests/controllers/console/workspace/test_load_balancing_config.py b/api/tests/unit_tests/controllers/console/workspace/test_load_balancing_config.py index b6c5f7651a2..76c5d8f1beb 100644 --- a/api/tests/unit_tests/controllers/console/workspace/test_load_balancing_config.py +++ b/api/tests/unit_tests/controllers/console/workspace/test_load_balancing_config.py @@ -21,7 +21,7 @@ if not hasattr(builtins, "MethodView"): builtins.MethodView = MethodView # type: ignore[attr-defined] from models.account import Account, TenantAccountRole -from tests.unit_tests.model_factories import make_account +from tests.helpers.unit.model_factories import make_account @pytest.fixture diff --git a/api/tests/unit_tests/controllers/console/workspace/test_members.py b/api/tests/unit_tests/controllers/console/workspace/test_members.py index 83f88ad064a..e929060a28b 100644 --- a/api/tests/unit_tests/controllers/console/workspace/test_members.py +++ b/api/tests/unit_tests/controllers/console/workspace/test_members.py @@ -55,7 +55,7 @@ from services.errors.workspace import ( ) from services.workspace.contracts import WorkspaceMemberRecord, WorkspaceMemberRole, WorkspaceMemberSummary from services.workspace.member_service import WorkspaceMemberQueryService, WorkspaceMemberService -from tests.unit_tests.controllers.error_body import raised_error +from tests.helpers.unit.controllers.error_body import raised_error class _RecordingWorkspaceMemberQueryService(WorkspaceMemberQueryService): diff --git a/api/tests/unit_tests/controllers/console/workspace/test_model_providers.py b/api/tests/unit_tests/controllers/console/workspace/test_model_providers.py index d9c6249028c..299cbb5d5c1 100644 --- a/api/tests/unit_tests/controllers/console/workspace/test_model_providers.py +++ b/api/tests/unit_tests/controllers/console/workspace/test_model_providers.py @@ -1,6 +1,6 @@ from collections.abc import Callable from inspect import unwrap -from types import SimpleNamespace +from types import FunctionType, SimpleNamespace from unittest.mock import patch import pytest @@ -9,6 +9,10 @@ from pydantic_core import ValidationError from werkzeug.exceptions import Forbidden, UnprocessableEntity from configs import dify_config +from controllers.common.rbac import AgentId, RBACPermission, Workspace +from controllers.console.agent.composer import AgentComposerApi +from controllers.console.agent.roster import AgentAppApi +from controllers.console.datasets.rag_pipeline.datasource_auth import DatasourceAuth from controllers.console.workspace.model_providers import ( ModelProviderCredentialApi, ModelProviderCredentialSwitchApi, @@ -28,6 +32,12 @@ from controllers.console.workspace.model_providers import ( ParserPreferredProviderType, PreferredProviderTypeUpdateApi, ) +from controllers.console.workspace.models import ModelProviderModelCredentialApi +from controllers.console.workspace.tool_providers import ( + ToolBuiltinProviderAddApi, + ToolMCPDetailApi, + ToolOAuthCustomClient, +) from core.entities.provider_entities import CredentialConfiguration from enums import DeploymentEdition from graphon.model_runtime.entities.common_entities import I18nObject @@ -48,7 +58,97 @@ from services.entities.model_provider_entities import ( SystemConfigurationResponse, ) from services.workspace.contracts import EffectiveCreditPool -from tests.unit_tests.config_override import config_overrides_context +from tests.helpers.unit.config_override import config_overrides_context +from tests.helpers.unit.controllers.rbac_introspection import rbac_checks + + +@pytest.mark.parametrize( + ("method", "permission"), + [ + (ToolBuiltinProviderAddApi.post, RBACPermission.CREDENTIAL_CREATE), + ], +) +def test_workspace_credential_mutations_require_management_permission( + method: FunctionType, permission: RBACPermission +) -> None: + legacy_wrapper = unwrap(method, stop=lambda wrapper: "is_admin_or_owner_required" in wrapper.__code__.co_qualname) + assert "is_admin_or_owner_required" in legacy_wrapper.__code__.co_qualname + + [check] = rbac_checks(method) + assert check.scene == permission + assert isinstance(check.locator, Workspace) + + +@pytest.mark.parametrize( + "method", + [ + ModelProviderCredentialApi.get, + ModelProviderModelCredentialApi.get, + ], +) +def test_model_provider_credential_get_requires_admin_and_rbac( + method: FunctionType, +) -> None: + """GET endpoints that return provider credential details must enforce + the same admin + RBAC gates as their sibling POST/PUT/DELETE methods.""" + legacy_wrapper = unwrap(method, stop=lambda wrapper: "is_admin_or_owner_required" in wrapper.__code__.co_qualname) + assert "is_admin_or_owner_required" in legacy_wrapper.__code__.co_qualname + + [check] = rbac_checks(method) + assert check.scene == RBACPermission.CREDENTIAL_MANAGE + assert isinstance(check.locator, Workspace) + + +def test_tool_oauth_custom_client_get_requires_admin_and_rbac() -> None: + """GET endpoint that returns custom OAuth client params must enforce + the same admin + RBAC gates as its sibling POST and DELETE methods.""" + method = ToolOAuthCustomClient.get + + legacy_wrapper = unwrap(method, stop=lambda wrapper: "is_admin_or_owner_required" in wrapper.__code__.co_qualname) + assert "is_admin_or_owner_required" in legacy_wrapper.__code__.co_qualname + + [check] = rbac_checks(method) + assert check.scene == RBACPermission.CREDENTIAL_MANAGE + assert isinstance(check.locator, Workspace) + + +def test_datasource_auth_get_requires_edit_and_rbac() -> None: + """GET endpoint that lists datasource credentials must enforce + the same edit + RBAC gates as its sibling POST method.""" + method = DatasourceAuth.get + + edit_wrapper = unwrap(method, stop=lambda wrapper: "edit_permission_required" in wrapper.__code__.co_qualname) + assert "edit_permission_required" in edit_wrapper.__code__.co_qualname + + [check] = rbac_checks(method) + assert check.scene == RBACPermission.CREDENTIAL_MANAGE + assert isinstance(check.locator, Workspace) + + +@pytest.mark.parametrize( + "method", + [ + AgentAppApi.get, + AgentComposerApi.get, + ], +) +def test_agent_app_get_requires_rbac(method: FunctionType) -> None: + """GET endpoints that return agent app details or composer state must require + the agent preview scene on the agent resource, same as their sibling routes.""" + [check] = rbac_checks(method) + assert check.scene == RBACPermission.AGENT_PREVIEW + assert isinstance(check.locator, AgentId) + + +def test_tool_mcp_detail_get_requires_rbac() -> None: + """GET endpoint that returns MCP provider details must enforce + the workspace MCP management RBAC gate.""" + method = ToolMCPDetailApi.get + + [check] = rbac_checks(method) + assert check.scene == RBACPermission.MCP_MANAGE + assert isinstance(check.locator, Workspace) + VALID_UUID = "123e4567-e89b-12d3-a456-426614174000" INVALID_UUID = "123" diff --git a/api/tests/unit_tests/controllers/console/workspace/test_models.py b/api/tests/unit_tests/controllers/console/workspace/test_models.py index bb26734f578..1dad21e31cc 100644 --- a/api/tests/unit_tests/controllers/console/workspace/test_models.py +++ b/api/tests/unit_tests/controllers/console/workspace/test_models.py @@ -33,7 +33,7 @@ from graphon.model_runtime.entities.model_entities import ModelType from graphon.model_runtime.errors.validate import CredentialsValidateFailedError from models import Account from models.provider import TenantDefaultModel -from tests.unit_tests.model_factories import make_account +from tests.helpers.unit.model_factories import make_account def _account() -> Account: diff --git a/api/tests/unit_tests/controllers/console/workspace/test_plugin.py b/api/tests/unit_tests/controllers/console/workspace/test_plugin.py index ed2ae031615..0b27fa22601 100644 --- a/api/tests/unit_tests/controllers/console/workspace/test_plugin.py +++ b/api/tests/unit_tests/controllers/console/workspace/test_plugin.py @@ -82,9 +82,9 @@ from models.account import ( TenantPluginDebugPermission, TenantPluginInstallPermission, ) -from tests.unit_tests.config_override import config_overrides_context -from tests.unit_tests.controllers.error_body import raised_error -from tests.unit_tests.model_factories import make_account +from tests.helpers.unit.config_override import config_overrides_context +from tests.helpers.unit.controllers.error_body import raised_error +from tests.helpers.unit.model_factories import make_account def _plugin_category_list_item(category: str = "tool") -> dict[str, Any]: diff --git a/api/tests/unit_tests/controllers/console/workspace/test_skills.py b/api/tests/unit_tests/controllers/console/workspace/test_skills.py index 760998fd0ea..1b4505ae116 100644 --- a/api/tests/unit_tests/controllers/console/workspace/test_skills.py +++ b/api/tests/unit_tests/controllers/console/workspace/test_skills.py @@ -51,7 +51,7 @@ from services.skill_management_service import ( SkillRestorePayload, SkillVersionUpdatePayload, ) -from tests.unit_tests.controllers.error_body import raised_error +from tests.helpers.unit.controllers.error_body import raised_error @dataclass(frozen=True) diff --git a/api/tests/unit_tests/controllers/console/workspace/test_snippets.py b/api/tests/unit_tests/controllers/console/workspace/test_snippets.py index d0da00c5143..ea9a6a57859 100644 --- a/api/tests/unit_tests/controllers/console/workspace/test_snippets.py +++ b/api/tests/unit_tests/controllers/console/workspace/test_snippets.py @@ -13,8 +13,8 @@ from controllers.console.workspace import snippets as snippets_module from models.account import Account, TenantAccountRole from models.snippet import CustomizedSnippet from services.snippet_dsl_service import ImportStatus, SnippetImportInfo -from tests.unit_tests.controllers.error_body import raised_error -from tests.unit_tests.model_factories import make_account +from tests.helpers.unit.controllers.error_body import raised_error +from tests.helpers.unit.model_factories import make_account def _account(account_id: str = "account-1") -> Account: diff --git a/api/tests/unit_tests/controllers/console/workspace/test_tool_provider_apis.py b/api/tests/unit_tests/controllers/console/workspace/test_tool_provider_apis.py deleted file mode 100644 index ebfc88c2a68..00000000000 --- a/api/tests/unit_tests/controllers/console/workspace/test_tool_provider_apis.py +++ /dev/null @@ -1,744 +0,0 @@ -"""Unit tests for controllers.console.workspace.tool_providers endpoints.""" - -from __future__ import annotations - -from collections.abc import Iterator -from inspect import unwrap -from types import SimpleNamespace -from unittest.mock import MagicMock, create_autospec, patch - -import pytest -from flask import Flask -from werkzeug.exceptions import Forbidden - -from controllers.console.workspace.tool_providers import ( - ApiToolProviderAddPayload, - ApiToolProviderDeletePayload, - ApiToolProviderUpdatePayload, - BuiltinProviderDefaultCredentialPayload, - BuiltinToolAddPayload, - BuiltinToolCredentialDeletePayload, - BuiltinToolUpdatePayload, - ToolApiListApi, - ToolApiProviderAddApi, - ToolApiProviderDeleteApi, - ToolApiProviderGetApi, - ToolApiProviderGetRemoteSchemaApi, - ToolApiProviderListToolsApi, - ToolApiProviderUpdateApi, - ToolBuiltinListApi, - ToolBuiltinProviderAddApi, - ToolBuiltinProviderCredentialsSchemaApi, - ToolBuiltinProviderDeleteApi, - ToolBuiltinProviderGetCredentialInfoApi, - ToolBuiltinProviderGetCredentialsApi, - ToolBuiltinProviderGetOauthClientSchemaApi, - ToolBuiltinProviderIconApi, - ToolBuiltinProviderInfoApi, - ToolBuiltinProviderListToolsApi, - ToolBuiltinProviderSetDefaultApi, - ToolBuiltinProviderUpdateApi, - ToolLabelsApi, - ToolOAuthCallback, - ToolOAuthCustomClient, - ToolOAuthCustomClientPayload, - ToolPluginOAuthApi, - ToolProviderListApi, - ToolWorkflowListApi, - ToolWorkflowProviderCreateApi, - ToolWorkflowProviderDeleteApi, - ToolWorkflowProviderGetApi, - ToolWorkflowProviderUpdateApi, - WorkflowToolCreatePayload, - WorkflowToolDeletePayload, - WorkflowToolUpdatePayload, - is_valid_url, -) -from core.tools.entities.api_entities import ToolProviderApiEntity as CoreToolProviderApiEntity -from models.account import Account, TenantAccountRole -from services.credentials.query import CredentialQuery - - -@pytest.fixture -def credential_query() -> Iterator[MagicMock]: - query = create_autospec(CredentialQuery, instance=True, spec_set=True) - with patch( - "controllers.console.workspace.tool_providers.application_services", - return_value=SimpleNamespace(credential_queries=query), - ): - yield query - - -def empty_mapping() -> dict[str, object]: - return {} - - -def empty_list() -> list[object]: - return [] - - -def emoji_icon() -> dict[str, str]: - return {"content": "tool", "background": "#252525"} - - -def i18n(text: str) -> dict[str, str]: - return {"en_US": text} - - -def tool_payload(name: str = "ping") -> dict[str, object]: - return { - "author": "langgenius", - "name": name, - "label": i18n(name.title()), - "description": i18n(f"{name} description"), - "parameters": [], - "labels": ["utilities"], - "output_schema": {}, - } - - -def provider_payload( - *, - provider_id: str = "provider-1", - name: str = "provider", - provider_type: str = "builtin", - tools: list[dict[str, object]] | None = None, -) -> dict[str, object]: - return { - "id": provider_id, - "author": "langgenius", - "name": name, - "description": i18n(f"{name} description"), - "icon": emoji_icon(), - "icon_dark": emoji_icon(), - "label": i18n(name.title()), - "type": provider_type, - "masked_credentials": {"api_key": "[__HIDDEN__]"}, - "original_credentials": {"api_key": "sk-secret"}, - "is_team_authorization": False, - "allow_delete": True, - "plugin_id": "langgenius/provider", - "plugin_unique_identifier": "langgenius/provider:1.0.0", - "tools": tools or [tool_payload()], - "labels": ["utilities"], - "server_url": "", - "updated_at": 1710000000, - "server_identifier": "", - "masked_headers": None, - "original_headers": None, - "authentication": None, - "is_dynamic_registration": True, - "configuration": None, - "identity_mode": "off", - "workflow_app_id": None, - } - - -def provider_entity( - *, - provider_id: str = "provider-1", - name: str = "provider", - provider_type: str = "builtin", - tools: list[dict[str, object]] | None = None, -) -> CoreToolProviderApiEntity: - return CoreToolProviderApiEntity.model_validate( - provider_payload(provider_id=provider_id, name=name, provider_type=provider_type, tools=tools) - ) - - -def credential_payload() -> dict[str, object]: - return { - "id": "credential-1", - "name": "Default credential", - "provider": "provider", - "credential_type": "api-key", - "is_default": True, - "credentials": {"api_key": "masked"}, - "visibility": "all_team_members", - "created_by": "user-1", - "partial_member_list": [], - "from_other_member": False, - } - - -def provider_config_payload() -> dict[str, object]: - return {"type": "secret-input", "name": "api_key", "required": True} - - -def api_tool_bundle_payload() -> dict[str, object]: - return { - "server_url": "https://api.example.com", - "method": "get", - "summary": "Ping", - "operation_id": "ping", - "parameters": [], - "author": "langgenius", - "icon": None, - "openapi": {"operationId": "ping"}, - "output_schema": {}, - } - - -def api_provider_detail_payload() -> dict[str, object]: - return { - "schema_type": "openapi", - "schema": "{}", - "tools": [api_tool_bundle_payload()], - "icon": emoji_icon(), - "description": "API provider", - "credentials": {}, - "privacy_policy": "", - "custom_disclaimer": "", - "labels": ["utilities"], - } - - -def credential_info_payload() -> dict[str, object]: - return { - "supported_credential_types": ["api-key", "oauth2"], - "is_oauth_custom_client_enabled": False, - "credentials": [credential_payload()], - } - - -def oauth_client_schema_payload() -> dict[str, object]: - return { - "schema": [provider_config_payload()], - "is_oauth_custom_client_enabled": False, - "is_system_oauth_params_exists": True, - "client_params": {"client_id": "masked"}, - "redirect_uri": "https://console.example.com/oauth/callback", - } - - -def tool_label_payload() -> dict[str, object]: - return { - "name": "utilities", - "label": i18n("Utilities"), - "icon": "wrench", - } - - -def make_account(*, id: str = "u", role: TenantAccountRole = TenantAccountRole.EDITOR) -> Account: - account = Account(name="Alice", email=f"{id}@example.com") - account.id = id - account.role = role - return account - - -class TestUtils: - def test_is_valid_url(self) -> None: - assert is_valid_url("https://example.com") - assert is_valid_url("http://example.com") - assert not is_valid_url("") - assert not is_valid_url("ftp://example.com") - assert not is_valid_url("not-a-url") - assert not is_valid_url(None) # type: ignore[arg-type] - - -class TestToolProviderListApi: - def test_get_success(self, app: Flask) -> None: - api = ToolProviderListApi() - method = unwrap(api.get) - - with ( - app.test_request_context("/"), - patch( - "controllers.console.workspace.tool_providers.ToolCommonService.list_tool_providers", - return_value=[provider_entity(provider_id="p1").to_dict()], - ), - ): - result = method(api, "t1", make_account(id="u1")) - - assert result[0]["id"] == "p1" - assert result[0]["team_credentials"] == {"api_key": "[__HIDDEN__]"} - assert "masked_credentials" not in result[0] - assert "original_credentials" not in result[0] - assert result[0]["tools"][0]["name"] == "ping" - - -class TestBuiltinProviderApis: - def test_list_tools(self, app: Flask) -> None: - api = ToolBuiltinProviderListToolsApi() - method = unwrap(api.get) - - with ( - app.test_request_context("/"), - patch( - "controllers.console.workspace.tool_providers.BuiltinToolManageService.list_builtin_tool_provider_tools", - return_value=[tool_payload()], - ), - ): - assert method(api, "t1", "provider")[0]["name"] == "ping" - - def test_info(self, app: Flask) -> None: - api = ToolBuiltinProviderInfoApi() - method = unwrap(api.get) - - with ( - app.test_request_context("/"), - patch( - "controllers.console.workspace.tool_providers.BuiltinToolManageService.get_builtin_tool_provider_info", - return_value=provider_entity(), - ), - ): - result = method(api, "t1", "provider") - - assert result["id"] == "provider-1" - assert result["team_credentials"] == {"api_key": "[__HIDDEN__]"} - assert "masked_credentials" not in result - assert "original_credentials" not in result - - def test_delete(self, app: Flask) -> None: - api = ToolBuiltinProviderDeleteApi() - method = unwrap(api.post) - - with ( - app.test_request_context("/", json={"credential_id": "cid"}), - patch( - "controllers.console.workspace.tool_providers.BuiltinToolManageService.delete_builtin_tool_provider", - return_value={"result": "success"}, - ), - ): - req = BuiltinToolCredentialDeletePayload(credential_id="cid") - assert method(api, req, "t1", "provider")["result"] == "success" - - def test_add_invalid_type(self, app: Flask) -> None: - api = ToolBuiltinProviderAddApi() - method = unwrap(api.post) - - with ( - app.test_request_context("/", json={"credentials": empty_mapping(), "type": "invalid"}), - ): - with pytest.raises(ValueError, match="1 validation error for BuiltinToolAddPayload\\\ntype"): - method( - api, - BuiltinToolAddPayload(credentials=empty_mapping(), type="invalid"), - "t", - make_account(), - "provider", - ) - - def test_add_success(self, app: Flask) -> None: - api = ToolBuiltinProviderAddApi() - method = unwrap(api.post) - - payload = {"credentials": empty_mapping(), "type": "oauth2", "name": "n"} - - with ( - app.test_request_context("/", json=payload), - patch( - "controllers.console.workspace.tool_providers.BuiltinToolManageService.add_builtin_tool_provider", - return_value={"result": "success"}, - ), - ): - assert method(api, BuiltinToolAddPayload(**payload), "t", make_account(), "provider")["result"] == "success" - - def test_update(self, app: Flask) -> None: - api = ToolBuiltinProviderUpdateApi() - method = unwrap(api.post) - - payload = {"credential_id": "c1", "credentials": empty_mapping(), "name": "n"} - - with ( - app.test_request_context("/", json=payload), - patch( - "controllers.console.workspace.tool_providers.BuiltinToolManageService.update_builtin_tool_provider", - return_value={"result": "success"}, - ), - ): - req = BuiltinToolUpdatePayload(**payload) - assert method(api, req, "t", make_account(), "provider")["result"] == "success" - - def test_get_credentials(self, app: Flask, credential_query: MagicMock) -> None: - api = ToolBuiltinProviderGetCredentialsApi() - method = unwrap(api.get) - - with ( - app.test_request_context("/"), - patch( - "controllers.console.workspace.tool_providers.BuiltinToolManageService.get_builtin_tool_provider_credentials", - return_value=[credential_payload()], - ) as get_credentials, - ): - assert method(api, "t", make_account(id="user-1"), "provider")[0]["id"] == "credential-1" - assert get_credentials.call_args.kwargs["credential_query"] is credential_query - - def test_icon(self, app: Flask) -> None: - api = ToolBuiltinProviderIconApi() - method = unwrap(api.get) - - with ( - app.test_request_context("/"), - patch( - "controllers.console.workspace.tool_providers.BuiltinToolManageService.get_builtin_tool_provider_icon", - return_value=(b"x", "image/png"), - ), - ): - response = method(api, "provider") - assert response.mimetype == "image/png" - - def test_credentials_schema(self, app: Flask) -> None: - api = ToolBuiltinProviderCredentialsSchemaApi() - method = unwrap(api.get) - - with ( - app.test_request_context("/"), - patch( - "controllers.console.workspace.tool_providers.BuiltinToolManageService.list_builtin_provider_credentials_schema", - return_value=[provider_config_payload()], - ), - ): - assert method(api, "t", "provider", "oauth2")[0]["name"] == "api_key" - - def test_set_default_credential(self, app: Flask) -> None: - api = ToolBuiltinProviderSetDefaultApi() - method = unwrap(api.post) - - with ( - app.test_request_context("/", json={"id": "c1"}), - patch( - "controllers.console.workspace.tool_providers.BuiltinToolManageService.set_default_provider", - return_value={"result": "success"}, - ), - ): - req = BuiltinProviderDefaultCredentialPayload(id="c1") - assert method(api, req, "t", "provider")["result"] == "success" - - def test_get_credential_info(self, app: Flask, credential_query: MagicMock) -> None: - api = ToolBuiltinProviderGetCredentialInfoApi() - method = unwrap(api.get) - - with ( - app.test_request_context("/"), - patch( - "controllers.console.workspace.tool_providers.BuiltinToolManageService.get_builtin_tool_provider_credential_info", - return_value=credential_info_payload(), - ) as get_info, - ): - assert method(api, "t", make_account(), "provider")["credentials"][0]["id"] == "credential-1" - assert get_info.call_args.kwargs["credential_query"] is credential_query - - def test_get_oauth_client_schema(self, app: Flask) -> None: - api = ToolBuiltinProviderGetOauthClientSchemaApi() - method = unwrap(api.get) - - with ( - app.test_request_context("/"), - patch( - "controllers.console.workspace.tool_providers.BuiltinToolManageService.get_builtin_tool_provider_oauth_client_schema", - return_value=oauth_client_schema_payload(), - ), - ): - assert method(api, "t", "provider")["schema"][0]["name"] == "api_key" - - -class TestApiProviderApis: - def test_add(self, app: Flask) -> None: - api = ToolApiProviderAddApi() - method = unwrap(api.post) - - payload = { - "credentials": empty_mapping(), - "schema_type": "openapi", - "schema": "{}", - "provider": "p", - "icon": emoji_icon(), - } - - with ( - app.test_request_context("/", json=payload), - patch( - "controllers.console.workspace.tool_providers.ApiToolManageService.create_api_tool_provider", - return_value={"result": "success"}, - ) as create_api_tool_provider, - ): - assert method(api, ApiToolProviderAddPayload(**payload), "t", make_account()) == {"result": "success"} - - create_api_tool_provider.assert_called_once() - assert create_api_tool_provider.call_args.args[3] == emoji_icon() - - def test_remote_schema(self, app: Flask) -> None: - api = ToolApiProviderGetRemoteSchemaApi() - method = unwrap(api.get) - openapi_schema = '{"openapi":"3.0.0","info":{"title":"Demo API","version":"1.0.0"},"paths":{}}' - - with ( - app.test_request_context("/?url=http://x.com"), - patch( - "controllers.console.workspace.tool_providers.ApiToolManageService.get_api_tool_provider_remote_schema", - return_value={"schema": openapi_schema}, - ), - ): - assert method(api, "t", make_account()) == {"schema": openapi_schema} - - def test_list_tools(self, app: Flask) -> None: - api = ToolApiProviderListToolsApi() - method = unwrap(api.get) - - with ( - app.test_request_context("/?provider=p"), - patch( - "controllers.console.workspace.tool_providers.ApiToolManageService.list_api_tool_provider_tools", - return_value=[tool_payload("api_ping")], - ), - ): - assert method(api, "t", make_account())[0]["name"] == "api_ping" - - def test_update(self, app: Flask) -> None: - api = ToolApiProviderUpdateApi() - method = unwrap(api.post) - - payload = { - "credentials": empty_mapping(), - "schema_type": "openapi", - "schema": "{}", - "provider": "p", - "original_provider": "o", - "icon": emoji_icon(), - "privacy_policy": "", - "custom_disclaimer": "", - } - - with ( - app.test_request_context("/", json=payload), - patch( - "controllers.console.workspace.tool_providers.ApiToolManageService.update_api_tool_provider", - return_value={"result": "success"}, - ) as update_api_tool_provider, - ): - assert method(api, ApiToolProviderUpdatePayload(**payload), "t", make_account()) == {"result": "success"} - - update_api_tool_provider.assert_called_once() - assert update_api_tool_provider.call_args.args[4] == emoji_icon() - - def test_delete(self, app: Flask) -> None: - api = ToolApiProviderDeleteApi() - method = unwrap(api.post) - - with ( - app.test_request_context("/", json={"provider": "p"}), - patch( - "controllers.console.workspace.tool_providers.ApiToolManageService.delete_api_tool_provider", - return_value={"result": "success"}, - ), - ): - req = ApiToolProviderDeletePayload(provider="p") - assert method(api, req, "t", make_account())["result"] == "success" - - def test_get(self, app: Flask) -> None: - api = ToolApiProviderGetApi() - method = unwrap(api.get) - - with ( - app.test_request_context("/?provider=p"), - patch( - "controllers.console.workspace.tool_providers.ApiToolManageService.get_api_tool_provider", - return_value=api_provider_detail_payload(), - ), - ): - assert method(api, "t", make_account())["schema"] == "{}" - - -class TestWorkflowApis: - def test_create(self, app: Flask) -> None: - api = ToolWorkflowProviderCreateApi() - method = unwrap(api.post) - - payload = { - "workflow_app_id": "123e4567-e89b-12d3-a456-426614174000", - "name": "n", - "label": "l", - "description": "d", - "icon": emoji_icon(), - "parameters": empty_list(), - } - - with ( - app.test_request_context("/", json=payload), - patch( - "controllers.console.workspace.tool_providers.WorkflowToolManageService.create_workflow_tool", - return_value={"result": "success"}, - ) as create_workflow_tool, - ): - assert method(api, WorkflowToolCreatePayload(**payload), "t", make_account()) == {"result": "success"} - - create_workflow_tool.assert_called_once() - assert create_workflow_tool.call_args.kwargs["icon"] == emoji_icon() - - def test_update_invalid(self, app: Flask) -> None: - api = ToolWorkflowProviderUpdateApi() - method = unwrap(api.post) - - payload = { - "workflow_tool_id": "123e4567-e89b-12d3-a456-426614174000", - "name": "Tool", - "label": "Tool Label", - "description": "A tool", - "icon": emoji_icon(), - } - - with ( - app.test_request_context("/", json=payload), - patch( - "controllers.console.workspace.tool_providers.WorkflowToolManageService.update_workflow_tool", - return_value={"result": "success"}, - ) as update_workflow_tool, - ): - result = method(api, WorkflowToolUpdatePayload(**payload), "t", make_account()) - assert result == {"result": "success"} - - update_workflow_tool.assert_called_once() - assert update_workflow_tool.call_args.args[5] == emoji_icon() - - def test_delete(self, app: Flask) -> None: - api = ToolWorkflowProviderDeleteApi() - method = unwrap(api.post) - - with ( - app.test_request_context("/", json={"workflow_tool_id": "123e4567-e89b-12d3-a456-426614174000"}), - patch( - "controllers.console.workspace.tool_providers.WorkflowToolManageService.delete_workflow_tool", - return_value={"result": "success"}, - ), - ): - req = WorkflowToolDeletePayload(workflow_tool_id="123e4567-e89b-12d3-a456-426614174000") - assert method(api, req, "t", make_account())["result"] == "success" - - def test_get_error(self, app: Flask) -> None: - api = ToolWorkflowProviderGetApi() - method = unwrap(api.get) - - with ( - app.test_request_context("/"), - ): - with pytest.raises( - ValueError, - match=( - "1 validation error for WorkflowToolGetQuery\\\n Value error, " - "workflow_tool_id or workflow_app_id is required" - ), - ): - method(api, "t", make_account()) - - -class TestLists: - def test_builtin_list(self, app: Flask) -> None: - api = ToolBuiltinListApi() - method = unwrap(api.get) - - with ( - app.test_request_context("/"), - patch( - "controllers.console.workspace.tool_providers.BuiltinToolManageService.list_builtin_tools", - return_value=[provider_entity(provider_id="builtin-1")], - ), - ): - assert method(api, "t", make_account())[0]["id"] == "builtin-1" - - def test_api_list(self, app: Flask) -> None: - api = ToolApiListApi() - method = unwrap(api.get) - - with ( - app.test_request_context("/"), - patch( - "controllers.console.workspace.tool_providers.ApiToolManageService.list_api_tools", - return_value=[provider_entity(provider_id="api-1", provider_type="api")], - ), - ): - assert method(api, "t")[0]["id"] == "api-1" - - def test_workflow_list(self, app: Flask) -> None: - api = ToolWorkflowListApi() - method = unwrap(api.get) - - with ( - app.test_request_context("/"), - patch( - "controllers.console.workspace.tool_providers.WorkflowToolManageService.list_tenant_workflow_tools", - return_value=[provider_entity(provider_id="workflow-1", provider_type="workflow")], - ), - ): - assert method(api, "t", make_account())[0]["id"] == "workflow-1" - - -class TestLabels: - def test_labels(self, app: Flask) -> None: - api = ToolLabelsApi() - method = unwrap(api.get) - - with ( - app.test_request_context("/"), - patch( - "controllers.console.workspace.tool_providers.ToolLabelsService.list_tool_labels", - return_value=[tool_label_payload()], - ), - ): - assert method(api)[0]["name"] == "utilities" - - -class TestOAuth: - def test_oauth_no_client(self, app: Flask) -> None: - api = ToolPluginOAuthApi() - method = unwrap(api.get) - - with ( - app.test_request_context("/"), - patch( - "controllers.console.workspace.tool_providers.BuiltinToolManageService.get_oauth_client", - return_value=None, - ), - ): - with pytest.raises(Forbidden): - method(api, "t", make_account(), "provider") - - def test_oauth_callback_no_cookie(self, app: Flask) -> None: - api = ToolOAuthCallback() - method = unwrap(api.get) - - with app.test_request_context("/"): - with pytest.raises(Forbidden): - method(api, "provider") - - -class TestOAuthCustomClient: - def test_save_custom_client(self, app: Flask) -> None: - api = ToolOAuthCustomClient() - method = unwrap(api.post) - - with ( - app.test_request_context("/", json={"client_params": {"a": 1}}), - patch( - "controllers.console.workspace.tool_providers.BuiltinToolManageService.save_custom_oauth_client_params", - return_value={"result": "success"}, - ), - ): - req = ToolOAuthCustomClientPayload(client_params={"a": 1}) - assert method(api, req, "t", "provider") == {"result": "success"} - - def test_get_custom_client(self, app: Flask) -> None: - api = ToolOAuthCustomClient() - method = unwrap(api.get) - - with ( - app.test_request_context("/"), - patch( - "controllers.console.workspace.tool_providers.BuiltinToolManageService.get_custom_oauth_client_params", - return_value={"client_id": "x"}, - ), - ): - assert method(api, "t", "provider") == {"client_id": "x"} - - def test_delete_custom_client(self, app: Flask) -> None: - api = ToolOAuthCustomClient() - method = unwrap(api.delete) - - with ( - app.test_request_context("/"), - patch( - "controllers.console.workspace.tool_providers.BuiltinToolManageService.delete_custom_oauth_client_params", - return_value={"result": "success"}, - ), - ): - assert method(api, "t", "provider") == {"result": "success"} diff --git a/api/tests/unit_tests/controllers/console/workspace/test_tool_providers.py b/api/tests/unit_tests/controllers/console/workspace/test_tool_providers.py index c7a7ee03f77..10e34fbb32d 100644 --- a/api/tests/unit_tests/controllers/console/workspace/test_tool_providers.py +++ b/api/tests/unit_tests/controllers/console/workspace/test_tool_providers.py @@ -1,14 +1,12 @@ -"""Endpoint tests for controllers.console.workspace.tool_providers.""" - from __future__ import annotations import builtins import importlib -from collections.abc import Callable +from collections.abc import Callable, Iterator from contextlib import ExitStack, contextmanager from inspect import unwrap -from types import ModuleType, SimpleNamespace -from unittest.mock import MagicMock, patch +from types import FunctionType, ModuleType, SimpleNamespace +from unittest.mock import MagicMock, create_autospec, patch from uuid import NAMESPACE_URL, uuid5 import pytest @@ -17,16 +15,775 @@ from flask.views import MethodView from sqlalchemy.orm import Session, scoped_session, sessionmaker from werkzeug.exceptions import Forbidden +from controllers.common.rbac import RBACPermission, Workspace +from controllers.console.datasets.rag_pipeline.datasource_auth import DatasourceAuthOauthCustomClient +from controllers.console.workspace.tool_providers import ( + ApiToolProviderAddPayload, + ApiToolProviderDeletePayload, + ApiToolProviderUpdatePayload, + BuiltinProviderDefaultCredentialPayload, + BuiltinToolAddPayload, + BuiltinToolCredentialDeletePayload, + BuiltinToolUpdatePayload, + ToolApiListApi, + ToolApiProviderAddApi, + ToolApiProviderDeleteApi, + ToolApiProviderGetApi, + ToolApiProviderGetRemoteSchemaApi, + ToolApiProviderListToolsApi, + ToolApiProviderUpdateApi, + ToolBuiltinListApi, + ToolBuiltinProviderAddApi, + ToolBuiltinProviderCredentialsSchemaApi, + ToolBuiltinProviderDeleteApi, + ToolBuiltinProviderGetCredentialInfoApi, + ToolBuiltinProviderGetCredentialsApi, + ToolBuiltinProviderGetOauthClientSchemaApi, + ToolBuiltinProviderIconApi, + ToolBuiltinProviderInfoApi, + ToolBuiltinProviderListToolsApi, + ToolBuiltinProviderSetDefaultApi, + ToolBuiltinProviderUpdateApi, + ToolLabelsApi, + ToolOAuthCallback, + ToolOAuthCustomClient, + ToolOAuthCustomClientPayload, + ToolPluginOAuthApi, + ToolProviderListApi, + ToolWorkflowListApi, + ToolWorkflowProviderCreateApi, + ToolWorkflowProviderDeleteApi, + ToolWorkflowProviderGetApi, + ToolWorkflowProviderUpdateApi, + WorkflowToolCreatePayload, + WorkflowToolDeletePayload, + WorkflowToolUpdatePayload, + is_valid_url, +) from core.tools.entities.api_entities import ToolProviderApiEntity as CoreToolProviderApiEntity from core.tools.entities.common_entities import I18nObject from core.tools.entities.tool_entities import ToolParameter from enums import DeploymentEdition -from models import Account, BuiltinToolProvider, Tenant, TenantAccountJoin +from models import Account as _tool_providers_Account +from models import BuiltinToolProvider, Tenant, TenantAccountJoin +from models.account import Account as _tool_provider_apis_Account from models.account import TenantAccountRole from models.credential_permission import CredentialPermission from models.enums import PermissionEnum from repositories.credentials.query_repository import CredentialQueryRepository +from services.credentials.query import CredentialQuery from services.plugin.dispatch_identity import PluginOAuthClientConfig +from tests.helpers.unit.controllers.rbac_introspection import rbac_checks + + +class TestOauthCustomClientPermissionsCases: + @staticmethod + @pytest.mark.parametrize( + ("method", "legacy_gate"), + [ + (ToolOAuthCustomClient.delete, "is_admin_or_owner_required"), + (DatasourceAuthOauthCustomClient.delete, "edit_permission_required"), + ], + ) + def test_custom_oauth_client_delete_requires_management_permission(method: FunctionType, legacy_gate: str) -> None: + legacy_wrapper = unwrap(method, stop=lambda wrapper: legacy_gate in wrapper.__code__.co_qualname) + assert legacy_gate in legacy_wrapper.__code__.co_qualname + + # Both routes pass resource_required=False, so their bundle resolves to a Workspace + # locator regardless of the resource_type they declared (see _single_check). + [check] = rbac_checks(method) + assert check.scene == RBACPermission.CREDENTIAL_MANAGE + assert isinstance(check.locator, Workspace) + + +"""Unit tests for controllers.console.workspace.tool_providers endpoints.""" + + +def empty_mapping() -> dict[str, object]: + return {} + + +def empty_list() -> list[object]: + return [] + + +def emoji_icon() -> dict[str, str]: + return {"content": "tool", "background": "#252525"} + + +def i18n(text: str) -> dict[str, str]: + return {"en_US": text} + + +def tool_payload(name: str = "ping") -> dict[str, object]: + return { + "author": "langgenius", + "name": name, + "label": i18n(name.title()), + "description": i18n(f"{name} description"), + "parameters": [], + "labels": ["utilities"], + "output_schema": {}, + } + + +def provider_payload( + *, + provider_id: str = "provider-1", + name: str = "provider", + provider_type: str = "builtin", + tools: list[dict[str, object]] | None = None, +) -> dict[str, object]: + return { + "id": provider_id, + "author": "langgenius", + "name": name, + "description": i18n(f"{name} description"), + "icon": emoji_icon(), + "icon_dark": emoji_icon(), + "label": i18n(name.title()), + "type": provider_type, + "masked_credentials": {"api_key": "[__HIDDEN__]"}, + "original_credentials": {"api_key": "sk-secret"}, + "is_team_authorization": False, + "allow_delete": True, + "plugin_id": "langgenius/provider", + "plugin_unique_identifier": "langgenius/provider:1.0.0", + "tools": tools or [tool_payload()], + "labels": ["utilities"], + "server_url": "", + "updated_at": 1710000000, + "server_identifier": "", + "masked_headers": None, + "original_headers": None, + "authentication": None, + "is_dynamic_registration": True, + "configuration": None, + "identity_mode": "off", + "workflow_app_id": None, + } + + +def provider_entity( + *, + provider_id: str = "provider-1", + name: str = "provider", + provider_type: str = "builtin", + tools: list[dict[str, object]] | None = None, +) -> CoreToolProviderApiEntity: + return CoreToolProviderApiEntity.model_validate( + provider_payload(provider_id=provider_id, name=name, provider_type=provider_type, tools=tools) + ) + + +def credential_payload() -> dict[str, object]: + return { + "id": "credential-1", + "name": "Default credential", + "provider": "provider", + "credential_type": "api-key", + "is_default": True, + "credentials": {"api_key": "masked"}, + "visibility": "all_team_members", + "created_by": "user-1", + "partial_member_list": [], + "from_other_member": False, + } + + +def provider_config_payload() -> dict[str, object]: + return {"type": "secret-input", "name": "api_key", "required": True} + + +def api_tool_bundle_payload() -> dict[str, object]: + return { + "server_url": "https://api.example.com", + "method": "get", + "summary": "Ping", + "operation_id": "ping", + "parameters": [], + "author": "langgenius", + "icon": None, + "openapi": {"operationId": "ping"}, + "output_schema": {}, + } + + +def api_provider_detail_payload() -> dict[str, object]: + return { + "schema_type": "openapi", + "schema": "{}", + "tools": [api_tool_bundle_payload()], + "icon": emoji_icon(), + "description": "API provider", + "credentials": {}, + "privacy_policy": "", + "custom_disclaimer": "", + "labels": ["utilities"], + } + + +def credential_info_payload() -> dict[str, object]: + return { + "supported_credential_types": ["api-key", "oauth2"], + "is_oauth_custom_client_enabled": False, + "credentials": [credential_payload()], + } + + +def oauth_client_schema_payload() -> dict[str, object]: + return { + "schema": [provider_config_payload()], + "is_oauth_custom_client_enabled": False, + "is_system_oauth_params_exists": True, + "client_params": {"client_id": "masked"}, + "redirect_uri": "https://console.example.com/oauth/callback", + } + + +def tool_label_payload() -> dict[str, object]: + return { + "name": "utilities", + "label": i18n("Utilities"), + "icon": "wrench", + } + + +def make_account(*, id: str = "u", role: TenantAccountRole = TenantAccountRole.EDITOR) -> _tool_provider_apis_Account: + account = _tool_provider_apis_Account(name="Alice", email=f"{id}@example.com") + account.id = id + account.role = role + return account + + +class TestToolProviderApisCases: + @staticmethod + @pytest.fixture + def credential_query() -> Iterator[MagicMock]: + query = create_autospec(CredentialQuery, instance=True, spec_set=True) + with patch( + "controllers.console.workspace.tool_providers.application_services", + return_value=SimpleNamespace(credential_queries=query), + ): + yield query + + class TestUtils: + def test_is_valid_url(self) -> None: + assert is_valid_url("https://example.com") + assert is_valid_url("http://example.com") + assert not is_valid_url("") + assert not is_valid_url("ftp://example.com") + assert not is_valid_url("not-a-url") + assert not is_valid_url(None) # type: ignore[arg-type] + + class TestToolProviderListApi: + def test_get_success(self, app: Flask) -> None: + api = ToolProviderListApi() + method = unwrap(api.get) + + with ( + app.test_request_context("/"), + patch( + "controllers.console.workspace.tool_providers.ToolCommonService.list_tool_providers", + return_value=[provider_entity(provider_id="p1").to_dict()], + ), + ): + result = method(api, "t1", make_account(id="u1")) + + assert result[0]["id"] == "p1" + assert result[0]["team_credentials"] == {"api_key": "[__HIDDEN__]"} + assert "masked_credentials" not in result[0] + assert "original_credentials" not in result[0] + assert result[0]["tools"][0]["name"] == "ping" + + class TestBuiltinProviderApis: + def test_list_tools(self, app: Flask) -> None: + api = ToolBuiltinProviderListToolsApi() + method = unwrap(api.get) + + with ( + app.test_request_context("/"), + patch( + "controllers.console.workspace.tool_providers.BuiltinToolManageService.list_builtin_tool_provider_tools", + return_value=[tool_payload()], + ), + ): + assert method(api, "t1", "provider")[0]["name"] == "ping" + + def test_info(self, app: Flask) -> None: + api = ToolBuiltinProviderInfoApi() + method = unwrap(api.get) + + with ( + app.test_request_context("/"), + patch( + "controllers.console.workspace.tool_providers.BuiltinToolManageService.get_builtin_tool_provider_info", + return_value=provider_entity(), + ), + ): + result = method(api, "t1", "provider") + + assert result["id"] == "provider-1" + assert result["team_credentials"] == {"api_key": "[__HIDDEN__]"} + assert "masked_credentials" not in result + assert "original_credentials" not in result + + def test_delete(self, app: Flask) -> None: + api = ToolBuiltinProviderDeleteApi() + method = unwrap(api.post) + + with ( + app.test_request_context("/", json={"credential_id": "cid"}), + patch( + "controllers.console.workspace.tool_providers.BuiltinToolManageService.delete_builtin_tool_provider", + return_value={"result": "success"}, + ), + ): + req = BuiltinToolCredentialDeletePayload(credential_id="cid") + assert method(api, req, "t1", "provider")["result"] == "success" + + def test_add_invalid_type(self, app: Flask) -> None: + api = ToolBuiltinProviderAddApi() + method = unwrap(api.post) + + with ( + app.test_request_context("/", json={"credentials": empty_mapping(), "type": "invalid"}), + ): + with pytest.raises(ValueError, match="1 validation error for BuiltinToolAddPayload\\\ntype"): + method( + api, + BuiltinToolAddPayload(credentials=empty_mapping(), type="invalid"), + "t", + make_account(), + "provider", + ) + + def test_add_success(self, app: Flask) -> None: + api = ToolBuiltinProviderAddApi() + method = unwrap(api.post) + + payload = {"credentials": empty_mapping(), "type": "oauth2", "name": "n"} + + with ( + app.test_request_context("/", json=payload), + patch( + "controllers.console.workspace.tool_providers.BuiltinToolManageService.add_builtin_tool_provider", + return_value={"result": "success"}, + ), + ): + assert ( + method(api, BuiltinToolAddPayload(**payload), "t", make_account(), "provider")["result"] + == "success" + ) + + def test_update(self, app: Flask) -> None: + api = ToolBuiltinProviderUpdateApi() + method = unwrap(api.post) + + payload = {"credential_id": "c1", "credentials": empty_mapping(), "name": "n"} + + with ( + app.test_request_context("/", json=payload), + patch( + "controllers.console.workspace.tool_providers.BuiltinToolManageService.update_builtin_tool_provider", + return_value={"result": "success"}, + ), + ): + req = BuiltinToolUpdatePayload(**payload) + assert method(api, req, "t", make_account(), "provider")["result"] == "success" + + def test_get_credentials(self, app: Flask, credential_query: MagicMock) -> None: + api = ToolBuiltinProviderGetCredentialsApi() + method = unwrap(api.get) + + with ( + app.test_request_context("/"), + patch( + "controllers.console.workspace.tool_providers.BuiltinToolManageService.get_builtin_tool_provider_credentials", + return_value=[credential_payload()], + ) as get_credentials, + ): + assert method(api, "t", make_account(id="user-1"), "provider")[0]["id"] == "credential-1" + assert get_credentials.call_args.kwargs["credential_query"] is credential_query + + def test_icon(self, app: Flask) -> None: + api = ToolBuiltinProviderIconApi() + method = unwrap(api.get) + + with ( + app.test_request_context("/"), + patch( + "controllers.console.workspace.tool_providers.BuiltinToolManageService.get_builtin_tool_provider_icon", + return_value=(b"x", "image/png"), + ), + ): + response = method(api, "provider") + assert response.mimetype == "image/png" + + def test_credentials_schema(self, app: Flask) -> None: + api = ToolBuiltinProviderCredentialsSchemaApi() + method = unwrap(api.get) + + with ( + app.test_request_context("/"), + patch( + "controllers.console.workspace.tool_providers.BuiltinToolManageService.list_builtin_provider_credentials_schema", + return_value=[provider_config_payload()], + ), + ): + assert method(api, "t", "provider", "oauth2")[0]["name"] == "api_key" + + def test_set_default_credential(self, app: Flask) -> None: + api = ToolBuiltinProviderSetDefaultApi() + method = unwrap(api.post) + + with ( + app.test_request_context("/", json={"id": "c1"}), + patch( + "controllers.console.workspace.tool_providers.BuiltinToolManageService.set_default_provider", + return_value={"result": "success"}, + ), + ): + req = BuiltinProviderDefaultCredentialPayload(id="c1") + assert method(api, req, "t", "provider")["result"] == "success" + + def test_get_credential_info(self, app: Flask, credential_query: MagicMock) -> None: + api = ToolBuiltinProviderGetCredentialInfoApi() + method = unwrap(api.get) + + with ( + app.test_request_context("/"), + patch( + "controllers.console.workspace.tool_providers.BuiltinToolManageService.get_builtin_tool_provider_credential_info", + return_value=credential_info_payload(), + ) as get_info, + ): + assert method(api, "t", make_account(), "provider")["credentials"][0]["id"] == "credential-1" + assert get_info.call_args.kwargs["credential_query"] is credential_query + + def test_get_oauth_client_schema(self, app: Flask) -> None: + api = ToolBuiltinProviderGetOauthClientSchemaApi() + method = unwrap(api.get) + + with ( + app.test_request_context("/"), + patch( + "controllers.console.workspace.tool_providers.BuiltinToolManageService.get_builtin_tool_provider_oauth_client_schema", + return_value=oauth_client_schema_payload(), + ), + ): + assert method(api, "t", "provider")["schema"][0]["name"] == "api_key" + + class TestApiProviderApis: + def test_add(self, app: Flask) -> None: + api = ToolApiProviderAddApi() + method = unwrap(api.post) + + payload = { + "credentials": empty_mapping(), + "schema_type": "openapi", + "schema": "{}", + "provider": "p", + "icon": emoji_icon(), + } + + with ( + app.test_request_context("/", json=payload), + patch( + "controllers.console.workspace.tool_providers.ApiToolManageService.create_api_tool_provider", + return_value={"result": "success"}, + ) as create_api_tool_provider, + ): + assert method(api, ApiToolProviderAddPayload(**payload), "t", make_account()) == {"result": "success"} + + create_api_tool_provider.assert_called_once() + assert create_api_tool_provider.call_args.args[3] == emoji_icon() + + def test_remote_schema(self, app: Flask) -> None: + api = ToolApiProviderGetRemoteSchemaApi() + method = unwrap(api.get) + openapi_schema = '{"openapi":"3.0.0","info":{"title":"Demo API","version":"1.0.0"},"paths":{}}' + + with ( + app.test_request_context("/?url=http://x.com"), + patch( + "controllers.console.workspace.tool_providers.ApiToolManageService.get_api_tool_provider_remote_schema", + return_value={"schema": openapi_schema}, + ), + ): + assert method(api, "t", make_account()) == {"schema": openapi_schema} + + def test_list_tools(self, app: Flask) -> None: + api = ToolApiProviderListToolsApi() + method = unwrap(api.get) + + with ( + app.test_request_context("/?provider=p"), + patch( + "controllers.console.workspace.tool_providers.ApiToolManageService.list_api_tool_provider_tools", + return_value=[tool_payload("api_ping")], + ), + ): + assert method(api, "t", make_account())[0]["name"] == "api_ping" + + def test_update(self, app: Flask) -> None: + api = ToolApiProviderUpdateApi() + method = unwrap(api.post) + + payload = { + "credentials": empty_mapping(), + "schema_type": "openapi", + "schema": "{}", + "provider": "p", + "original_provider": "o", + "icon": emoji_icon(), + "privacy_policy": "", + "custom_disclaimer": "", + } + + with ( + app.test_request_context("/", json=payload), + patch( + "controllers.console.workspace.tool_providers.ApiToolManageService.update_api_tool_provider", + return_value={"result": "success"}, + ) as update_api_tool_provider, + ): + assert method(api, ApiToolProviderUpdatePayload(**payload), "t", make_account()) == { + "result": "success" + } + + update_api_tool_provider.assert_called_once() + assert update_api_tool_provider.call_args.args[4] == emoji_icon() + + def test_delete(self, app: Flask) -> None: + api = ToolApiProviderDeleteApi() + method = unwrap(api.post) + + with ( + app.test_request_context("/", json={"provider": "p"}), + patch( + "controllers.console.workspace.tool_providers.ApiToolManageService.delete_api_tool_provider", + return_value={"result": "success"}, + ), + ): + req = ApiToolProviderDeletePayload(provider="p") + assert method(api, req, "t", make_account())["result"] == "success" + + def test_get(self, app: Flask) -> None: + api = ToolApiProviderGetApi() + method = unwrap(api.get) + + with ( + app.test_request_context("/?provider=p"), + patch( + "controllers.console.workspace.tool_providers.ApiToolManageService.get_api_tool_provider", + return_value=api_provider_detail_payload(), + ), + ): + assert method(api, "t", make_account())["schema"] == "{}" + + class TestWorkflowApis: + def test_create(self, app: Flask) -> None: + api = ToolWorkflowProviderCreateApi() + method = unwrap(api.post) + + payload = { + "workflow_app_id": "123e4567-e89b-12d3-a456-426614174000", + "name": "n", + "label": "l", + "description": "d", + "icon": emoji_icon(), + "parameters": empty_list(), + } + + with ( + app.test_request_context("/", json=payload), + patch( + "controllers.console.workspace.tool_providers.WorkflowToolManageService.create_workflow_tool", + return_value={"result": "success"}, + ) as create_workflow_tool, + ): + assert method(api, WorkflowToolCreatePayload(**payload), "t", make_account()) == {"result": "success"} + + create_workflow_tool.assert_called_once() + assert create_workflow_tool.call_args.kwargs["icon"] == emoji_icon() + + def test_update_invalid(self, app: Flask) -> None: + api = ToolWorkflowProviderUpdateApi() + method = unwrap(api.post) + + payload = { + "workflow_tool_id": "123e4567-e89b-12d3-a456-426614174000", + "name": "Tool", + "label": "Tool Label", + "description": "A tool", + "icon": emoji_icon(), + } + + with ( + app.test_request_context("/", json=payload), + patch( + "controllers.console.workspace.tool_providers.WorkflowToolManageService.update_workflow_tool", + return_value={"result": "success"}, + ) as update_workflow_tool, + ): + result = method(api, WorkflowToolUpdatePayload(**payload), "t", make_account()) + assert result == {"result": "success"} + + update_workflow_tool.assert_called_once() + assert update_workflow_tool.call_args.args[5] == emoji_icon() + + def test_delete(self, app: Flask) -> None: + api = ToolWorkflowProviderDeleteApi() + method = unwrap(api.post) + + with ( + app.test_request_context("/", json={"workflow_tool_id": "123e4567-e89b-12d3-a456-426614174000"}), + patch( + "controllers.console.workspace.tool_providers.WorkflowToolManageService.delete_workflow_tool", + return_value={"result": "success"}, + ), + ): + req = WorkflowToolDeletePayload(workflow_tool_id="123e4567-e89b-12d3-a456-426614174000") + assert method(api, req, "t", make_account())["result"] == "success" + + def test_get_error(self, app: Flask) -> None: + api = ToolWorkflowProviderGetApi() + method = unwrap(api.get) + + with ( + app.test_request_context("/"), + ): + with pytest.raises( + ValueError, + match=( + "1 validation error for WorkflowToolGetQuery\\\n Value error, " + "workflow_tool_id or workflow_app_id is required" + ), + ): + method(api, "t", make_account()) + + class TestLists: + def test_builtin_list(self, app: Flask) -> None: + api = ToolBuiltinListApi() + method = unwrap(api.get) + + with ( + app.test_request_context("/"), + patch( + "controllers.console.workspace.tool_providers.BuiltinToolManageService.list_builtin_tools", + return_value=[provider_entity(provider_id="builtin-1")], + ), + ): + assert method(api, "t", make_account())[0]["id"] == "builtin-1" + + def test_api_list(self, app: Flask) -> None: + api = ToolApiListApi() + method = unwrap(api.get) + + with ( + app.test_request_context("/"), + patch( + "controllers.console.workspace.tool_providers.ApiToolManageService.list_api_tools", + return_value=[provider_entity(provider_id="api-1", provider_type="api")], + ), + ): + assert method(api, "t")[0]["id"] == "api-1" + + def test_workflow_list(self, app: Flask) -> None: + api = ToolWorkflowListApi() + method = unwrap(api.get) + + with ( + app.test_request_context("/"), + patch( + "controllers.console.workspace.tool_providers.WorkflowToolManageService.list_tenant_workflow_tools", + return_value=[provider_entity(provider_id="workflow-1", provider_type="workflow")], + ), + ): + assert method(api, "t", make_account())[0]["id"] == "workflow-1" + + class TestLabels: + def test_labels(self, app: Flask) -> None: + api = ToolLabelsApi() + method = unwrap(api.get) + + with ( + app.test_request_context("/"), + patch( + "controllers.console.workspace.tool_providers.ToolLabelsService.list_tool_labels", + return_value=[tool_label_payload()], + ), + ): + assert method(api)[0]["name"] == "utilities" + + class TestOAuth: + def test_oauth_no_client(self, app: Flask) -> None: + api = ToolPluginOAuthApi() + method = unwrap(api.get) + + with ( + app.test_request_context("/"), + patch( + "controllers.console.workspace.tool_providers.BuiltinToolManageService.get_oauth_client", + return_value=None, + ), + ): + with pytest.raises(Forbidden): + method(api, "t", make_account(), "provider") + + def test_oauth_callback_no_cookie(self, app: Flask) -> None: + api = ToolOAuthCallback() + method = unwrap(api.get) + + with app.test_request_context("/"): + with pytest.raises(Forbidden): + method(api, "provider") + + class TestOAuthCustomClient: + def test_save_custom_client(self, app: Flask) -> None: + api = ToolOAuthCustomClient() + method = unwrap(api.post) + + with ( + app.test_request_context("/", json={"client_params": {"a": 1}}), + patch( + "controllers.console.workspace.tool_providers.BuiltinToolManageService.save_custom_oauth_client_params", + return_value={"result": "success"}, + ), + ): + req = ToolOAuthCustomClientPayload(client_params={"a": 1}) + assert method(api, req, "t", "provider") == {"result": "success"} + + def test_get_custom_client(self, app: Flask) -> None: + api = ToolOAuthCustomClient() + method = unwrap(api.get) + + with ( + app.test_request_context("/"), + patch( + "controllers.console.workspace.tool_providers.BuiltinToolManageService.get_custom_oauth_client_params", + return_value={"client_id": "x"}, + ), + ): + assert method(api, "t", "provider") == {"client_id": "x"} + + def test_delete_custom_client(self, app: Flask) -> None: + api = ToolOAuthCustomClient() + method = unwrap(api.delete) + + with ( + app.test_request_context("/"), + patch( + "controllers.console.workspace.tool_providers.BuiltinToolManageService.delete_custom_oauth_client_params", + return_value={"result": "success"}, + ), + ): + assert method(api, "t", "provider") == {"result": "success"} + + +"""Endpoint tests for controllers.console.workspace.tool_providers.""" + if not hasattr(builtins, "MethodView"): builtins.MethodView = MethodView # type: ignore[attr-defined] @@ -36,62 +793,8 @@ _CONTROLLER_MODULE: ModuleType | None = None _WRAPS_MODULE: ModuleType | None = None -@pytest.fixture -def app() -> Flask: - flask_app = Flask(__name__) - flask_app.config["TESTING"] = True - return flask_app - - -@pytest.fixture -def controller_module(monkeypatch: pytest.MonkeyPatch, config_overrides: Callable[..., None], sqlite_session_factory): - """ - Import the controller with auth decorators neutralized only during import. - - The imported view classes retain those no-op decorators after import, so we - can restore the original globals immediately and avoid leaking auth patches - into unrelated tests such as libs.login unit coverage. - """ - - module_name = "controllers.console.workspace.tool_providers" - global _CONTROLLER_MODULE - if _CONTROLLER_MODULE is None: - - def _noop(func): - return func - - patch_targets = [ - ("libs.login.login_required", _noop), - ("controllers.console.wraps.setup_required", _noop), - ("controllers.console.wraps.account_initialization_required", _noop), - ("controllers.console.wraps.is_admin_or_owner_required", _noop), - ("controllers.console.wraps.enterprise_license_required", _noop), - ] - with ExitStack() as stack: - for target, value in patch_targets: - stack.enter_context(patch(target, value)) - _CONTROLLER_MODULE = importlib.import_module(module_name) - - module = _CONTROLLER_MODULE - - # Ensure decorators that consult deployment edition do not reach the database. - global _WRAPS_MODULE - wraps_module = importlib.import_module("controllers.console.wraps") - _WRAPS_MODULE = wraps_module - config_overrides(DEPLOYMENT_EDITION=DeploymentEdition.CLOUD) - - login_module = importlib.import_module("libs.login") - monkeypatch.setattr(login_module, "check_csrf_token", lambda *args, **kwargs: None) - monkeypatch.setattr( - module, - "application_services", - lambda: SimpleNamespace(credential_queries=CredentialQueryRepository(session_factory=sqlite_session_factory)), - ) - return module - - -def _mock_account(user_id: str = "user-123") -> Account: - user = Account(name="Test User", email=f"{user_id}@example.com") +def _mock_account(user_id: str = "user-123") -> _tool_providers_Account: + user = _tool_providers_Account(name="Test User", email=f"{user_id}@example.com") user.id = _stable_uuid(f"account:{user_id}") user.role = TenantAccountRole.NORMAL return user @@ -101,7 +804,7 @@ def _stable_uuid(value: str) -> str: return str(uuid5(NAMESPACE_URL, value)) -def _persist_workspace(session: Session, user: Account, tenant_name: str) -> Tenant: +def _persist_workspace(session: Session, user: _tool_providers_Account, tenant_name: str) -> Tenant: tenant = Tenant(name=tenant_name) tenant.id = _stable_uuid(f"tenant:{tenant_name}") membership = TenantAccountJoin( @@ -175,7 +878,7 @@ def _mock_credential_encryption(controller_module: ModuleType): def _set_current_account( monkeypatch: pytest.MonkeyPatch, controller_module: ModuleType, - user: Account, + user: _tool_providers_Account, tenant_id: str, ) -> None: def _getter(): @@ -329,461 +1032,6 @@ def _tool_label_response(controller_module: ModuleType, name: str = "search") -> return label.model_dump(mode="json"), expected -def test_tool_provider_list_calls_service_with_query( - app: Flask, controller_module: ModuleType, monkeypatch: pytest.MonkeyPatch -): - user = _mock_account() - _set_current_account(monkeypatch, controller_module, user, "tenant-456") - - service_payload, expected_response = _provider_list_item(controller_module, "builtin", "builtin") - service_mock = MagicMock(return_value=[service_payload]) - monkeypatch.setattr(controller_module.ToolCommonService, "list_tool_providers", service_mock) - - with app.test_request_context("/workspaces/current/tool-providers?type=builtin"): - response = controller_module.ToolProviderListApi().get() - - assert response == [expected_response] - service_mock.assert_called_once_with(user.id, "tenant-456", "builtin") - - -def test_builtin_provider_add_passes_payload( - app: Flask, controller_module: ModuleType, monkeypatch: pytest.MonkeyPatch -): - user = _mock_account() - user.role = TenantAccountRole.ADMIN - _set_current_account(monkeypatch, controller_module, user, "tenant-456") - - service_mock = MagicMock(return_value={"result": "success"}) - monkeypatch.setattr(controller_module.BuiltinToolManageService, "add_builtin_tool_provider", service_mock) - - payload = { - "credentials": {"api_key": "sk-test"}, - "name": "MyTool", - "type": controller_module.CredentialType.API_KEY, - } - - with app.test_request_context( - "/workspaces/current/tool-provider/builtin/openai/add", - method="POST", - json=payload, - ): - response = controller_module.ToolBuiltinProviderAddApi().post(provider="openai") - - assert response == {"result": "success"} - service_mock.assert_called_once_with( - user_id=user.id, - tenant_id="tenant-456", - provider="openai", - credentials={"api_key": "sk-test"}, - name="MyTool", - api_type=controller_module.CredentialType.API_KEY, - visibility=None, - ) - - -def test_builtin_provider_tools_get(app: Flask, controller_module, monkeypatch: pytest.MonkeyPatch): - user = _mock_account("user-tenant-789") - _set_current_account(monkeypatch, controller_module, user, "tenant-789") - - service_payload, expected_response = _tool_response(controller_module, "tool-a") - service_mock = MagicMock(return_value=[service_payload]) - monkeypatch.setattr(controller_module.BuiltinToolManageService, "list_builtin_tool_provider_tools", service_mock) - - with app.test_request_context( - "/workspaces/current/tool-provider/builtin/my-provider/tools", - method="GET", - ): - response = controller_module.ToolBuiltinProviderListToolsApi().get(provider="my-provider") - - assert response == [expected_response] - service_mock.assert_called_once_with("tenant-789", "my-provider") - - -def test_builtin_provider_info_get(app: Flask, controller_module, monkeypatch: pytest.MonkeyPatch): - user = _mock_account("user-tenant-9") - _set_current_account(monkeypatch, controller_module, user, "tenant-9") - service_payload, expected_response = _provider_entity_response(controller_module, "demo", "builtin") - service_mock = MagicMock(return_value=service_payload) - monkeypatch.setattr(controller_module.BuiltinToolManageService, "get_builtin_tool_provider_info", service_mock) - - with app.test_request_context("/info", method="GET"): - resp = controller_module.ToolBuiltinProviderInfoApi().get(provider="demo") - - assert resp == expected_response - service_mock.assert_called_once_with("tenant-9", "demo") - - -def test_builtin_provider_info_uses_core_to_dict_tool_projection( - app: Flask, controller_module: ModuleType, monkeypatch: pytest.MonkeyPatch -): - user = _mock_account("user-tenant-9") - _set_current_account(monkeypatch, controller_module, user, "tenant-9") - tool_parameter = ToolParameter( - name="system_files", - label=I18nObject(en_US="System Files", zh_Hans="System Files"), - type=ToolParameter.ToolParameterType.SYSTEM_FILES, - form=ToolParameter.ToolParameterForm.LLM, - input_schema=None, - ) - tool = controller_module.ToolApiEntity( - author="Dify", - name="demo-tool", - label=I18nObject(en_US="Demo Tool", zh_Hans="Demo Tool"), - description=I18nObject(en_US="Demo Tool description", zh_Hans="Demo Tool description"), - parameters=[tool_parameter], - labels=[], - output_schema={}, - ) - provider = CoreToolProviderApiEntity( - id="demo-id", - author="Dify", - name="demo", - description=I18nObject(en_US="demo description", zh_Hans="demo description"), - icon="tool.svg", - label=I18nObject(en_US="demo", zh_Hans="demo"), - type=controller_module.ToolProviderType.BUILT_IN, - masked_credentials={"api_key": "[__HIDDEN__]"}, - original_credentials={"api_key": "sk-secret"}, - tools=[tool], - ) - service_mock = MagicMock(return_value=provider) - monkeypatch.setattr(controller_module.BuiltinToolManageService, "get_builtin_tool_provider_info", service_mock) - - with app.test_request_context("/info", method="GET"): - resp = controller_module.ToolBuiltinProviderInfoApi().get(provider="demo") - - parameter = resp["tools"][0]["parameters"][0] - assert parameter["type"] == "files" - assert parameter["input_schema"] is None - assert resp["team_credentials"] == {"api_key": "[__HIDDEN__]"} - assert "masked_credentials" not in resp - assert "original_credentials" not in resp - - -@pytest.mark.parametrize( - "sqlite_session", - [(Account, Tenant, TenantAccountJoin, BuiltinToolProvider, CredentialPermission)], - indirect=True, -) -def test_builtin_provider_credentials_get( - app: Flask, - controller_module: ModuleType, - monkeypatch: pytest.MonkeyPatch, - sqlite_session: Session, -): - user = _mock_account("user-tenant-cred") - tenant = _persist_workspace(sqlite_session, user, "tenant-cred") - credential = _provider_credential( - tenant_id=tenant.id, - user_id=user.id, - credential_name="Credential", - ) - other_tenant = Tenant(name="other-tenant") - other_tenant.id = _stable_uuid("tenant:other-tenant") - other_credential = _provider_credential( - tenant_id=other_tenant.id, - user_id=user.id, - credential_name="Other Tenant Credential", - ) - sqlite_session.add_all([credential, other_tenant, other_credential]) - sqlite_session.commit() - _set_current_account(monkeypatch, controller_module, user, tenant.id) - - with ( - _bind_database_session(sqlite_session), - _mock_credential_encryption(controller_module), - app.test_request_context("/creds", method="GET"), - ): - response = controller_module.ToolBuiltinProviderGetCredentialsApi().get(provider="demo") - - assert [item["id"] for item in response] == [credential.id] - assert response[0]["name"] == "Credential" - assert response[0]["credentials"] == {"api_key": "[__HIDDEN__]"} - assert response[0]["created_by"] == user.id - assert other_credential.id not in {item["id"] for item in response} - - -@pytest.mark.parametrize( - "sqlite_session", - [(Account, Tenant, TenantAccountJoin, BuiltinToolProvider, CredentialPermission)], - indirect=True, -) -def test_builtin_provider_credentials_get_reads_repeated_include_ids( - app: Flask, - controller_module: ModuleType, - monkeypatch: pytest.MonkeyPatch, - sqlite_session: Session, -): - user = _mock_account("user-tenant-cred") - tenant = _persist_workspace(sqlite_session, user, "tenant-cred") - visible_credential = _provider_credential( - tenant_id=tenant.id, - user_id=user.id, - credential_name="Visible Credential", - ) - other_user = _mock_account("other-user") - borrowed_credential = _provider_credential( - tenant_id=tenant.id, - user_id=other_user.id, - credential_name="Borrowed Credential", - visibility=PermissionEnum.ONLY_ME, - ) - other_membership = TenantAccountJoin( - tenant_id=tenant.id, - account_id=other_user.id, - role=TenantAccountRole.NORMAL, - ) - sqlite_session.add_all([visible_credential, other_user, other_membership, borrowed_credential]) - sqlite_session.commit() - - request_path = ( - f"/creds?include_credential_ids={visible_credential.id}&include_credential_ids={borrowed_credential.id}" - ) - with ( - _bind_database_session(sqlite_session), - _mock_credential_encryption(controller_module), - app.test_request_context(request_path, method="GET"), - ): - api = controller_module.ToolBuiltinProviderGetCredentialsApi() - response = unwrap(api.get)(api, tenant.id, user, provider="demo") - - assert [item["id"] for item in response] == [visible_credential.id, borrowed_credential.id] - assert response[0]["from_other_member"] is False - assert response[1]["from_other_member"] is True - - -def test_api_provider_remote_schema_get(app: Flask, controller_module, monkeypatch: pytest.MonkeyPatch): - user = _mock_account() - _set_current_account(monkeypatch, controller_module, user, "tenant-10") - openapi_schema = '{"openapi":"3.0.0","info":{"title":"Demo API","version":"1.0.0"},"paths":{}}' - service_mock = MagicMock(return_value={"schema": openapi_schema}) - monkeypatch.setattr(controller_module.ApiToolManageService, "get_api_tool_provider_remote_schema", service_mock) - - with app.test_request_context("/remote?url=https://example.com/"): - resp = controller_module.ToolApiProviderGetRemoteSchemaApi().get() - - assert resp == {"schema": openapi_schema} - service_mock.assert_called_once_with(user.id, "tenant-10", "https://example.com/") - - -def test_api_provider_list_tools_get(app: Flask, controller_module, monkeypatch: pytest.MonkeyPatch): - user = _mock_account() - _set_current_account(monkeypatch, controller_module, user, "tenant-11") - service_payload, expected_response = _tool_response(controller_module, "t") - service_mock = MagicMock(return_value=[service_payload]) - monkeypatch.setattr(controller_module.ApiToolManageService, "list_api_tool_provider_tools", service_mock) - - with app.test_request_context("/tools?provider=foo"): - resp = controller_module.ToolApiProviderListToolsApi().get() - - assert resp == [expected_response] - service_mock.assert_called_once_with(user.id, "tenant-11", "foo") - - -def test_api_provider_get(app: Flask, controller_module, monkeypatch: pytest.MonkeyPatch): - user = _mock_account() - _set_current_account(monkeypatch, controller_module, user, "tenant-12") - service_payload, expected_response = _api_provider_detail_response(controller_module) - service_mock = MagicMock(return_value=service_payload) - monkeypatch.setattr(controller_module.ApiToolManageService, "get_api_tool_provider", service_mock) - - with app.test_request_context("/get?provider=foo"): - resp = controller_module.ToolApiProviderGetApi().get() - - assert resp == expected_response - service_mock.assert_called_once_with(user.id, "tenant-12", "foo") - - -def test_builtin_provider_credentials_schema_get(app: Flask, controller_module, monkeypatch: pytest.MonkeyPatch): - user = _mock_account("user-tenant-13") - _set_current_account(monkeypatch, controller_module, user, "tenant-13") - service_payload, expected_response = _provider_config_response(controller_module) - service_mock = MagicMock(return_value=[service_payload]) - monkeypatch.setattr( - controller_module.BuiltinToolManageService, - "list_builtin_provider_credentials_schema", - service_mock, - ) - - with app.test_request_context("/schema", method="GET"): - resp = controller_module.ToolBuiltinProviderCredentialsSchemaApi().get( - provider="demo", credential_type="api-key" - ) - - assert resp == [expected_response] - service_mock.assert_called_once() - - -def test_workflow_provider_get_by_tool(app: Flask, controller_module, monkeypatch: pytest.MonkeyPatch): - user = _mock_account() - _set_current_account(monkeypatch, controller_module, user, "tenant-wf") - service_payload, expected_response = _workflow_detail_response(controller_module) - tool_service = MagicMock(return_value=service_payload) - monkeypatch.setattr( - controller_module.WorkflowToolManageService, - "get_workflow_tool_by_tool_id", - tool_service, - ) - - tool_id = "00000000-0000-0000-0000-000000000001" - with app.test_request_context(f"/workflow?workflow_tool_id={tool_id}"): - resp = controller_module.ToolWorkflowProviderGetApi().get() - - assert resp == expected_response - tool_service.assert_called_once_with(user.id, "tenant-wf", tool_id) - - -def test_workflow_provider_get_by_app(app: Flask, controller_module, monkeypatch: pytest.MonkeyPatch): - user = _mock_account() - _set_current_account(monkeypatch, controller_module, user, "tenant-wf2") - service_payload, expected_response = _workflow_detail_response(controller_module) - service_mock = MagicMock(return_value=service_payload) - monkeypatch.setattr( - controller_module.WorkflowToolManageService, - "get_workflow_tool_by_app_id", - service_mock, - ) - - app_id = "00000000-0000-0000-0000-000000000002" - with app.test_request_context(f"/workflow?workflow_app_id={app_id}"): - resp = controller_module.ToolWorkflowProviderGetApi().get() - - assert resp == expected_response - service_mock.assert_called_once_with(user.id, "tenant-wf2", app_id) - - -def test_workflow_provider_list_tools(app: Flask, controller_module, monkeypatch: pytest.MonkeyPatch): - user = _mock_account() - _set_current_account(monkeypatch, controller_module, user, "tenant-wf3") - service_payload, expected_response = _tool_response(controller_module, "workflow-tool") - service_mock = MagicMock(return_value=[service_payload]) - monkeypatch.setattr(controller_module.WorkflowToolManageService, "list_single_workflow_tools", service_mock) - - tool_id = "00000000-0000-0000-0000-000000000003" - with app.test_request_context(f"/workflow/tools?workflow_tool_id={tool_id}"): - resp = controller_module.ToolWorkflowProviderListToolApi().get() - - assert resp == [expected_response] - service_mock.assert_called_once_with(user.id, "tenant-wf3", tool_id) - - -def test_builtin_tools_list(app: Flask, controller_module, monkeypatch: pytest.MonkeyPatch): - user = _mock_account() - _set_current_account(monkeypatch, controller_module, user, "tenant-bt") - - provider, expected_response = _provider_entity_response(controller_module, "builtin", "builtin") - monkeypatch.setattr( - controller_module.BuiltinToolManageService, - "list_builtin_tools", - MagicMock(return_value=[provider]), - ) - - with app.test_request_context("/tools/builtin"): - resp = controller_module.ToolBuiltinListApi().get() - - assert resp == [expected_response] - - -def test_api_tools_list(app: Flask, controller_module, monkeypatch: pytest.MonkeyPatch): - user = _mock_account("user-tenant-api") - _set_current_account(monkeypatch, controller_module, user, "tenant-api") - - provider, expected_response = _provider_entity_response(controller_module, "api", "api") - monkeypatch.setattr( - controller_module.ApiToolManageService, - "list_api_tools", - MagicMock(return_value=[provider]), - ) - - with app.test_request_context("/tools/api"): - resp = controller_module.ToolApiListApi().get() - - assert resp == [expected_response] - - -def test_workflow_tools_list(app: Flask, controller_module, monkeypatch: pytest.MonkeyPatch): - user = _mock_account() - _set_current_account(monkeypatch, controller_module, user, "tenant-wf4") - - provider, expected_response = _provider_entity_response(controller_module, "wf", "workflow") - monkeypatch.setattr( - controller_module.WorkflowToolManageService, - "list_tenant_workflow_tools", - MagicMock(return_value=[provider]), - ) - - with app.test_request_context("/tools/workflow"): - resp = controller_module.ToolWorkflowListApi().get() - - assert resp == [expected_response] - - -def test_tool_labels_list(app: Flask, controller_module, monkeypatch: pytest.MonkeyPatch): - user = _mock_account("user-label") - _set_current_account(monkeypatch, controller_module, user, "tenant-labels") - service_payload, expected_response = _tool_label_response(controller_module, "a") - monkeypatch.setattr(controller_module.ToolLabelsService, "list_tool_labels", lambda: [service_payload]) - - with app.test_request_context("/tool-labels"): - resp = controller_module.ToolLabelsApi().get() - - assert resp == [expected_response] - - -# --- _resolve_identity_mode: gating + None-resolution (PR #36839 review) --- - - -def test_resolve_identity_mode_none_keeps_current_when_enterprise( - controller_module, config_overrides: Callable[..., None] -): - """None means 'leave unchanged' — fall back to the stored mode (update path).""" - identity_mode = importlib.import_module("core.entities.mcp_provider").IdentityMode - config_overrides(DEPLOYMENT_EDITION=DeploymentEdition.ENTERPRISE) - - resolved = controller_module._resolve_identity_mode(None, current=identity_mode.IDP_TOKEN) - - assert resolved == identity_mode.IDP_TOKEN - - -def test_resolve_identity_mode_explicit_value_overrides_current( - controller_module, config_overrides: Callable[..., None] -): - """An explicit value wins over the stored mode.""" - identity_mode = importlib.import_module("core.entities.mcp_provider").IdentityMode - config_overrides(DEPLOYMENT_EDITION=DeploymentEdition.ENTERPRISE) - - resolved = controller_module._resolve_identity_mode(identity_mode.OFF, current=identity_mode.IDP_TOKEN) - - assert resolved == identity_mode.OFF - - -def test_resolve_identity_mode_coerces_non_off_to_off_when_not_enterprise( - controller_module, config_overrides: Callable[..., None] -): - """Gate: a non-EE deployment must never persist a non-OFF mode — the - runtime won't forward, so the stored row must not imply it does.""" - identity_mode = importlib.import_module("core.entities.mcp_provider").IdentityMode - config_overrides(DEPLOYMENT_EDITION=DeploymentEdition.COMMUNITY) - - # Both an explicit idp_token request AND an inherited non-OFF current - # must collapse to OFF. - assert ( - controller_module._resolve_identity_mode(identity_mode.IDP_TOKEN, current=identity_mode.OFF) - == identity_mode.OFF - ) - assert controller_module._resolve_identity_mode(None, current=identity_mode.IDP_TOKEN) == identity_mode.OFF - - -def test_resolve_identity_mode_off_is_passthrough_when_not_enterprise( - controller_module, config_overrides: Callable[..., None] -): - """OFF is always fine — the gate only neutralizes non-OFF values.""" - identity_mode = importlib.import_module("core.entities.mcp_provider").IdentityMode - config_overrides(DEPLOYMENT_EDITION=DeploymentEdition.COMMUNITY) - - assert controller_module._resolve_identity_mode(None, current=identity_mode.OFF) == identity_mode.OFF - - # --- OAuth start + callback: visibility passthrough (PR #39840 review) --- # # The OAuth flow stashes the user-chosen visibility in the proxy context on the @@ -798,7 +1046,7 @@ _OAUTH_PROVIDER_PATH = "langgenius/github/github" def _invoke_oauth_start( controller_module: ModuleType, app: Flask, - user: Account, + user: _tool_providers_Account, tenant_id: str, *, query_string: str = "", @@ -838,67 +1086,6 @@ def _mock_oauth_start_deps(controller_module: ModuleType, monkeypatch: pytest.Mo yield create_proxy_context -@pytest.mark.parametrize( - ("query_visibility", "expected_stored"), - [ - ("only_me", "only_me"), - ("all_team_members", "all_team_members"), - ], -) -def test_tool_plugin_oauth_url_stashes_valid_visibility( - app: Flask, - controller_module: ModuleType, - monkeypatch: pytest.MonkeyPatch, - query_visibility: str, - expected_stored: str, -): - user = _mock_account("user-oauth") - _set_current_account(monkeypatch, controller_module, user, "tenant-oauth") - - with _mock_oauth_start_deps(controller_module, monkeypatch) as create_proxy_context: - _invoke_oauth_start( - controller_module, - app, - user, - "tenant-oauth", - query_string=f"visibility={query_visibility}", - ) - - _, kwargs = create_proxy_context.call_args - assert kwargs["extra_data"] == { - "visibility": expected_stored, - "plugin_unique_identifier": "langgenius/github:1.0.0@aaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaa", - } - - -@pytest.mark.parametrize("query_string", ["", "visibility=partial_members", "visibility=nonsense"]) -def test_tool_plugin_oauth_url_unknown_visibility_falls_back_to_only_me( - app: Flask, - controller_module: ModuleType, - monkeypatch: pytest.MonkeyPatch, - query_string: str, -): - """Missing, partial_members (rejected for plugin creds), and garbage inputs - all collapse to only_me — OAuth tokens are personal by default.""" - user = _mock_account("user-oauth-fallback") - _set_current_account(monkeypatch, controller_module, user, "tenant-oauth") - - with _mock_oauth_start_deps(controller_module, monkeypatch) as create_proxy_context: - _invoke_oauth_start( - controller_module, - app, - user, - "tenant-oauth", - query_string=query_string, - ) - - _, kwargs = create_proxy_context.call_args - assert kwargs["extra_data"] == { - "visibility": "only_me", - "plugin_unique_identifier": "langgenius/github:1.0.0@aaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaa", - } - - def _invoke_oauth_callback( controller_module: ModuleType, app: Flask, @@ -919,32 +1106,6 @@ def _invoke_oauth_callback( return unwrap(api.get)(api, provider=provider) -def test_tool_oauth_callback_rejects_replaced_package_before_dispatch( - app: Flask, controller_module: ModuleType, monkeypatch: pytest.MonkeyPatch -) -> None: - context = { - "user_id": "user-oauth", - "tenant_id": "tenant-oauth", - "plugin_id": "langgenius/github", - "provider": "github", - "plugin_unique_identifier": f"langgenius/github:1.0.0@{'a' * 64}", - } - config = PluginOAuthClientConfig(f"langgenius/github:2.0.0@{'b' * 64}", {"client_secret": "test-secret"}) - monkeypatch.setattr(controller_module.OAuthProxyService, "use_proxy_context", lambda _context_id: context) - monkeypatch.setattr( - controller_module.BuiltinToolManageService, "get_oauth_client", lambda *_args, **_kwargs: config - ) - - def unexpected_dispatch(*_args: object, **_kwargs: object) -> None: - raise AssertionError("Changed packages must not receive OAuth credentials") - - monkeypatch.setattr(controller_module.OAuthHandler, "get_credentials", unexpected_dispatch) - api = controller_module.ToolOAuthCallback() - with app.test_request_context("/", headers={"Cookie": "context_id=ctx-id"}): - with pytest.raises(Forbidden, match="Plugin changed"): - unwrap(api.get)(api, provider=_OAUTH_PROVIDER_PATH) - - @contextmanager def _mock_oauth_callback_deps(controller_module: ModuleType, monkeypatch: pytest.MonkeyPatch): """Mock everything ToolOAuthCallback.get calls out to and capture the @@ -966,39 +1127,642 @@ def _mock_oauth_callback_deps(controller_module: ModuleType, monkeypatch: pytest yield add_provider -@pytest.mark.parametrize( - ("stored", "expected"), - [ - ("only_me", "only_me"), - ("all_team_members", "all_team_members"), - ], -) -def test_tool_oauth_callback_persists_stored_visibility( - app: Flask, - controller_module: ModuleType, - monkeypatch: pytest.MonkeyPatch, - stored: str, - expected: str, -): - with _mock_oauth_callback_deps(controller_module, monkeypatch) as add_provider: - _invoke_oauth_callback(controller_module, app, stored_visibility=stored) +class TestToolProvidersCases: + @staticmethod + @pytest.fixture + def app() -> Flask: + flask_app = Flask(__name__) + flask_app.config["TESTING"] = True + return flask_app - _, kwargs = add_provider.call_args - assert kwargs["visibility"] == expected - assert kwargs["api_type"] == controller_module.CredentialType.OAUTH2 + @staticmethod + @pytest.fixture + def controller_module( + monkeypatch: pytest.MonkeyPatch, config_overrides: Callable[..., None], sqlite_session_factory + ): + """ + Import the controller with auth decorators neutralized only during import. + The imported view classes retain those no-op decorators after import, so we + can restore the original globals immediately and avoid leaking auth patches + into unrelated tests such as libs.login unit coverage. + """ -@pytest.mark.parametrize("stored", [None, "partial_members", "garbage"]) -def test_tool_oauth_callback_unknown_stored_visibility_falls_back_to_only_me( - app: Flask, - controller_module: ModuleType, - monkeypatch: pytest.MonkeyPatch, - stored: str | None, -): - """Older proxy contexts (pre-feature) and anything unexpected must default - to only_me — matches the safe-default we apply on the start endpoint.""" - with _mock_oauth_callback_deps(controller_module, monkeypatch) as add_provider: - _invoke_oauth_callback(controller_module, app, stored_visibility=stored) + module_name = "controllers.console.workspace.tool_providers" + global _CONTROLLER_MODULE + if _CONTROLLER_MODULE is None: - _, kwargs = add_provider.call_args - assert kwargs["visibility"] == "only_me" + def _noop(func): + return func + + patch_targets = [ + ("libs.login.login_required", _noop), + ("controllers.console.wraps.setup_required", _noop), + ("controllers.console.wraps.account_initialization_required", _noop), + ("controllers.console.wraps.is_admin_or_owner_required", _noop), + ("controllers.console.wraps.enterprise_license_required", _noop), + ] + with ExitStack() as stack: + for target, value in patch_targets: + stack.enter_context(patch(target, value)) + _CONTROLLER_MODULE = importlib.import_module(module_name) + + module = _CONTROLLER_MODULE + + # Ensure decorators that consult deployment edition do not reach the database. + global _WRAPS_MODULE + wraps_module = importlib.import_module("controllers.console.wraps") + _WRAPS_MODULE = wraps_module + config_overrides(DEPLOYMENT_EDITION=DeploymentEdition.CLOUD) + + login_module = importlib.import_module("libs.login") + monkeypatch.setattr(login_module, "check_csrf_token", lambda *args, **kwargs: None) + monkeypatch.setattr( + module, + "application_services", + lambda: SimpleNamespace( + credential_queries=CredentialQueryRepository(session_factory=sqlite_session_factory) + ), + ) + return module + + @staticmethod + def test_tool_provider_list_calls_service_with_query( + app: Flask, controller_module: ModuleType, monkeypatch: pytest.MonkeyPatch + ): + user = _mock_account() + _set_current_account(monkeypatch, controller_module, user, "tenant-456") + + service_payload, expected_response = _provider_list_item(controller_module, "builtin", "builtin") + service_mock = MagicMock(return_value=[service_payload]) + monkeypatch.setattr(controller_module.ToolCommonService, "list_tool_providers", service_mock) + + with app.test_request_context("/workspaces/current/tool-providers?type=builtin"): + response = controller_module.ToolProviderListApi().get() + + assert response == [expected_response] + service_mock.assert_called_once_with(user.id, "tenant-456", "builtin") + + @staticmethod + def test_builtin_provider_add_passes_payload( + app: Flask, controller_module: ModuleType, monkeypatch: pytest.MonkeyPatch + ): + user = _mock_account() + user.role = TenantAccountRole.ADMIN + _set_current_account(monkeypatch, controller_module, user, "tenant-456") + + service_mock = MagicMock(return_value={"result": "success"}) + monkeypatch.setattr(controller_module.BuiltinToolManageService, "add_builtin_tool_provider", service_mock) + + payload = { + "credentials": {"api_key": "sk-test"}, + "name": "MyTool", + "type": controller_module.CredentialType.API_KEY, + } + + with app.test_request_context( + "/workspaces/current/tool-provider/builtin/openai/add", + method="POST", + json=payload, + ): + response = controller_module.ToolBuiltinProviderAddApi().post(provider="openai") + + assert response == {"result": "success"} + service_mock.assert_called_once_with( + user_id=user.id, + tenant_id="tenant-456", + provider="openai", + credentials={"api_key": "sk-test"}, + name="MyTool", + api_type=controller_module.CredentialType.API_KEY, + visibility=None, + ) + + @staticmethod + def test_builtin_provider_tools_get(app: Flask, controller_module, monkeypatch: pytest.MonkeyPatch): + user = _mock_account("user-tenant-789") + _set_current_account(monkeypatch, controller_module, user, "tenant-789") + + service_payload, expected_response = _tool_response(controller_module, "tool-a") + service_mock = MagicMock(return_value=[service_payload]) + monkeypatch.setattr( + controller_module.BuiltinToolManageService, "list_builtin_tool_provider_tools", service_mock + ) + + with app.test_request_context( + "/workspaces/current/tool-provider/builtin/my-provider/tools", + method="GET", + ): + response = controller_module.ToolBuiltinProviderListToolsApi().get(provider="my-provider") + + assert response == [expected_response] + service_mock.assert_called_once_with("tenant-789", "my-provider") + + @staticmethod + def test_builtin_provider_info_get(app: Flask, controller_module, monkeypatch: pytest.MonkeyPatch): + user = _mock_account("user-tenant-9") + _set_current_account(monkeypatch, controller_module, user, "tenant-9") + service_payload, expected_response = _provider_entity_response(controller_module, "demo", "builtin") + service_mock = MagicMock(return_value=service_payload) + monkeypatch.setattr(controller_module.BuiltinToolManageService, "get_builtin_tool_provider_info", service_mock) + + with app.test_request_context("/info", method="GET"): + resp = controller_module.ToolBuiltinProviderInfoApi().get(provider="demo") + + assert resp == expected_response + service_mock.assert_called_once_with("tenant-9", "demo") + + @staticmethod + def test_builtin_provider_info_uses_core_to_dict_tool_projection( + app: Flask, controller_module: ModuleType, monkeypatch: pytest.MonkeyPatch + ): + user = _mock_account("user-tenant-9") + _set_current_account(monkeypatch, controller_module, user, "tenant-9") + tool_parameter = ToolParameter( + name="system_files", + label=I18nObject(en_US="System Files", zh_Hans="System Files"), + type=ToolParameter.ToolParameterType.SYSTEM_FILES, + form=ToolParameter.ToolParameterForm.LLM, + input_schema=None, + ) + tool = controller_module.ToolApiEntity( + author="Dify", + name="demo-tool", + label=I18nObject(en_US="Demo Tool", zh_Hans="Demo Tool"), + description=I18nObject(en_US="Demo Tool description", zh_Hans="Demo Tool description"), + parameters=[tool_parameter], + labels=[], + output_schema={}, + ) + provider = CoreToolProviderApiEntity( + id="demo-id", + author="Dify", + name="demo", + description=I18nObject(en_US="demo description", zh_Hans="demo description"), + icon="tool.svg", + label=I18nObject(en_US="demo", zh_Hans="demo"), + type=controller_module.ToolProviderType.BUILT_IN, + masked_credentials={"api_key": "[__HIDDEN__]"}, + original_credentials={"api_key": "sk-secret"}, + tools=[tool], + ) + service_mock = MagicMock(return_value=provider) + monkeypatch.setattr(controller_module.BuiltinToolManageService, "get_builtin_tool_provider_info", service_mock) + + with app.test_request_context("/info", method="GET"): + resp = controller_module.ToolBuiltinProviderInfoApi().get(provider="demo") + + parameter = resp["tools"][0]["parameters"][0] + assert parameter["type"] == "files" + assert parameter["input_schema"] is None + assert resp["team_credentials"] == {"api_key": "[__HIDDEN__]"} + assert "masked_credentials" not in resp + assert "original_credentials" not in resp + + @staticmethod + @pytest.mark.parametrize( + "sqlite_session", + [(_tool_providers_Account, Tenant, TenantAccountJoin, BuiltinToolProvider, CredentialPermission)], + indirect=True, + ) + def test_builtin_provider_credentials_get( + app: Flask, + controller_module: ModuleType, + monkeypatch: pytest.MonkeyPatch, + sqlite_session: Session, + ): + user = _mock_account("user-tenant-cred") + tenant = _persist_workspace(sqlite_session, user, "tenant-cred") + credential = _provider_credential( + tenant_id=tenant.id, + user_id=user.id, + credential_name="Credential", + ) + other_tenant = Tenant(name="other-tenant") + other_tenant.id = _stable_uuid("tenant:other-tenant") + other_credential = _provider_credential( + tenant_id=other_tenant.id, + user_id=user.id, + credential_name="Other Tenant Credential", + ) + sqlite_session.add_all([credential, other_tenant, other_credential]) + sqlite_session.commit() + _set_current_account(monkeypatch, controller_module, user, tenant.id) + + with ( + _bind_database_session(sqlite_session), + _mock_credential_encryption(controller_module), + app.test_request_context("/creds", method="GET"), + ): + response = controller_module.ToolBuiltinProviderGetCredentialsApi().get(provider="demo") + + assert [item["id"] for item in response] == [credential.id] + assert response[0]["name"] == "Credential" + assert response[0]["credentials"] == {"api_key": "[__HIDDEN__]"} + assert response[0]["created_by"] == user.id + assert other_credential.id not in {item["id"] for item in response} + + @staticmethod + @pytest.mark.parametrize( + "sqlite_session", + [(_tool_providers_Account, Tenant, TenantAccountJoin, BuiltinToolProvider, CredentialPermission)], + indirect=True, + ) + def test_builtin_provider_credentials_get_reads_repeated_include_ids( + app: Flask, + controller_module: ModuleType, + monkeypatch: pytest.MonkeyPatch, + sqlite_session: Session, + ): + user = _mock_account("user-tenant-cred") + tenant = _persist_workspace(sqlite_session, user, "tenant-cred") + visible_credential = _provider_credential( + tenant_id=tenant.id, + user_id=user.id, + credential_name="Visible Credential", + ) + other_user = _mock_account("other-user") + borrowed_credential = _provider_credential( + tenant_id=tenant.id, + user_id=other_user.id, + credential_name="Borrowed Credential", + visibility=PermissionEnum.ONLY_ME, + ) + other_membership = TenantAccountJoin( + tenant_id=tenant.id, + account_id=other_user.id, + role=TenantAccountRole.NORMAL, + ) + sqlite_session.add_all([visible_credential, other_user, other_membership, borrowed_credential]) + sqlite_session.commit() + + request_path = ( + f"/creds?include_credential_ids={visible_credential.id}&include_credential_ids={borrowed_credential.id}" + ) + with ( + _bind_database_session(sqlite_session), + _mock_credential_encryption(controller_module), + app.test_request_context(request_path, method="GET"), + ): + api = controller_module.ToolBuiltinProviderGetCredentialsApi() + response = unwrap(api.get)(api, tenant.id, user, provider="demo") + + assert [item["id"] for item in response] == [visible_credential.id, borrowed_credential.id] + assert response[0]["from_other_member"] is False + assert response[1]["from_other_member"] is True + + @staticmethod + def test_api_provider_remote_schema_get(app: Flask, controller_module, monkeypatch: pytest.MonkeyPatch): + user = _mock_account() + _set_current_account(monkeypatch, controller_module, user, "tenant-10") + openapi_schema = '{"openapi":"3.0.0","info":{"title":"Demo API","version":"1.0.0"},"paths":{}}' + service_mock = MagicMock(return_value={"schema": openapi_schema}) + monkeypatch.setattr(controller_module.ApiToolManageService, "get_api_tool_provider_remote_schema", service_mock) + + with app.test_request_context("/remote?url=https://example.com/"): + resp = controller_module.ToolApiProviderGetRemoteSchemaApi().get() + + assert resp == {"schema": openapi_schema} + service_mock.assert_called_once_with(user.id, "tenant-10", "https://example.com/") + + @staticmethod + def test_api_provider_list_tools_get(app: Flask, controller_module, monkeypatch: pytest.MonkeyPatch): + user = _mock_account() + _set_current_account(monkeypatch, controller_module, user, "tenant-11") + service_payload, expected_response = _tool_response(controller_module, "t") + service_mock = MagicMock(return_value=[service_payload]) + monkeypatch.setattr(controller_module.ApiToolManageService, "list_api_tool_provider_tools", service_mock) + + with app.test_request_context("/tools?provider=foo"): + resp = controller_module.ToolApiProviderListToolsApi().get() + + assert resp == [expected_response] + service_mock.assert_called_once_with(user.id, "tenant-11", "foo") + + @staticmethod + def test_api_provider_get(app: Flask, controller_module, monkeypatch: pytest.MonkeyPatch): + user = _mock_account() + _set_current_account(monkeypatch, controller_module, user, "tenant-12") + service_payload, expected_response = _api_provider_detail_response(controller_module) + service_mock = MagicMock(return_value=service_payload) + monkeypatch.setattr(controller_module.ApiToolManageService, "get_api_tool_provider", service_mock) + + with app.test_request_context("/get?provider=foo"): + resp = controller_module.ToolApiProviderGetApi().get() + + assert resp == expected_response + service_mock.assert_called_once_with(user.id, "tenant-12", "foo") + + @staticmethod + def test_builtin_provider_credentials_schema_get(app: Flask, controller_module, monkeypatch: pytest.MonkeyPatch): + user = _mock_account("user-tenant-13") + _set_current_account(monkeypatch, controller_module, user, "tenant-13") + service_payload, expected_response = _provider_config_response(controller_module) + service_mock = MagicMock(return_value=[service_payload]) + monkeypatch.setattr( + controller_module.BuiltinToolManageService, + "list_builtin_provider_credentials_schema", + service_mock, + ) + + with app.test_request_context("/schema", method="GET"): + resp = controller_module.ToolBuiltinProviderCredentialsSchemaApi().get( + provider="demo", credential_type="api-key" + ) + + assert resp == [expected_response] + service_mock.assert_called_once() + + @staticmethod + def test_workflow_provider_get_by_tool(app: Flask, controller_module, monkeypatch: pytest.MonkeyPatch): + user = _mock_account() + _set_current_account(monkeypatch, controller_module, user, "tenant-wf") + service_payload, expected_response = _workflow_detail_response(controller_module) + tool_service = MagicMock(return_value=service_payload) + monkeypatch.setattr( + controller_module.WorkflowToolManageService, + "get_workflow_tool_by_tool_id", + tool_service, + ) + + tool_id = "00000000-0000-0000-0000-000000000001" + with app.test_request_context(f"/workflow?workflow_tool_id={tool_id}"): + resp = controller_module.ToolWorkflowProviderGetApi().get() + + assert resp == expected_response + tool_service.assert_called_once_with(user.id, "tenant-wf", tool_id) + + @staticmethod + def test_workflow_provider_get_by_app(app: Flask, controller_module, monkeypatch: pytest.MonkeyPatch): + user = _mock_account() + _set_current_account(monkeypatch, controller_module, user, "tenant-wf2") + service_payload, expected_response = _workflow_detail_response(controller_module) + service_mock = MagicMock(return_value=service_payload) + monkeypatch.setattr( + controller_module.WorkflowToolManageService, + "get_workflow_tool_by_app_id", + service_mock, + ) + + app_id = "00000000-0000-0000-0000-000000000002" + with app.test_request_context(f"/workflow?workflow_app_id={app_id}"): + resp = controller_module.ToolWorkflowProviderGetApi().get() + + assert resp == expected_response + service_mock.assert_called_once_with(user.id, "tenant-wf2", app_id) + + @staticmethod + def test_workflow_provider_list_tools(app: Flask, controller_module, monkeypatch: pytest.MonkeyPatch): + user = _mock_account() + _set_current_account(monkeypatch, controller_module, user, "tenant-wf3") + service_payload, expected_response = _tool_response(controller_module, "workflow-tool") + service_mock = MagicMock(return_value=[service_payload]) + monkeypatch.setattr(controller_module.WorkflowToolManageService, "list_single_workflow_tools", service_mock) + + tool_id = "00000000-0000-0000-0000-000000000003" + with app.test_request_context(f"/workflow/tools?workflow_tool_id={tool_id}"): + resp = controller_module.ToolWorkflowProviderListToolApi().get() + + assert resp == [expected_response] + service_mock.assert_called_once_with(user.id, "tenant-wf3", tool_id) + + @staticmethod + def test_builtin_tools_list(app: Flask, controller_module, monkeypatch: pytest.MonkeyPatch): + user = _mock_account() + _set_current_account(monkeypatch, controller_module, user, "tenant-bt") + + provider, expected_response = _provider_entity_response(controller_module, "builtin", "builtin") + monkeypatch.setattr( + controller_module.BuiltinToolManageService, + "list_builtin_tools", + MagicMock(return_value=[provider]), + ) + + with app.test_request_context("/tools/builtin"): + resp = controller_module.ToolBuiltinListApi().get() + + assert resp == [expected_response] + + @staticmethod + def test_api_tools_list(app: Flask, controller_module, monkeypatch: pytest.MonkeyPatch): + user = _mock_account("user-tenant-api") + _set_current_account(monkeypatch, controller_module, user, "tenant-api") + + provider, expected_response = _provider_entity_response(controller_module, "api", "api") + monkeypatch.setattr( + controller_module.ApiToolManageService, + "list_api_tools", + MagicMock(return_value=[provider]), + ) + + with app.test_request_context("/tools/api"): + resp = controller_module.ToolApiListApi().get() + + assert resp == [expected_response] + + @staticmethod + def test_workflow_tools_list(app: Flask, controller_module, monkeypatch: pytest.MonkeyPatch): + user = _mock_account() + _set_current_account(monkeypatch, controller_module, user, "tenant-wf4") + + provider, expected_response = _provider_entity_response(controller_module, "wf", "workflow") + monkeypatch.setattr( + controller_module.WorkflowToolManageService, + "list_tenant_workflow_tools", + MagicMock(return_value=[provider]), + ) + + with app.test_request_context("/tools/workflow"): + resp = controller_module.ToolWorkflowListApi().get() + + assert resp == [expected_response] + + @staticmethod + def test_tool_labels_list(app: Flask, controller_module, monkeypatch: pytest.MonkeyPatch): + user = _mock_account("user-label") + _set_current_account(monkeypatch, controller_module, user, "tenant-labels") + service_payload, expected_response = _tool_label_response(controller_module, "a") + monkeypatch.setattr(controller_module.ToolLabelsService, "list_tool_labels", lambda: [service_payload]) + + with app.test_request_context("/tool-labels"): + resp = controller_module.ToolLabelsApi().get() + + assert resp == [expected_response] + + # --- _resolve_identity_mode: gating + None-resolution (PR #36839 review) --- + + @staticmethod + def test_resolve_identity_mode_none_keeps_current_when_enterprise( + controller_module, config_overrides: Callable[..., None] + ): + """None means 'leave unchanged' — fall back to the stored mode (update path).""" + identity_mode = importlib.import_module("core.entities.mcp_provider").IdentityMode + config_overrides(DEPLOYMENT_EDITION=DeploymentEdition.ENTERPRISE) + + resolved = controller_module._resolve_identity_mode(None, current=identity_mode.IDP_TOKEN) + + assert resolved == identity_mode.IDP_TOKEN + + @staticmethod + def test_resolve_identity_mode_explicit_value_overrides_current( + controller_module, config_overrides: Callable[..., None] + ): + """An explicit value wins over the stored mode.""" + identity_mode = importlib.import_module("core.entities.mcp_provider").IdentityMode + config_overrides(DEPLOYMENT_EDITION=DeploymentEdition.ENTERPRISE) + + resolved = controller_module._resolve_identity_mode(identity_mode.OFF, current=identity_mode.IDP_TOKEN) + + assert resolved == identity_mode.OFF + + @staticmethod + def test_resolve_identity_mode_coerces_non_off_to_off_when_not_enterprise( + controller_module, config_overrides: Callable[..., None] + ): + """Gate: a non-EE deployment must never persist a non-OFF mode — the + runtime won't forward, so the stored row must not imply it does.""" + identity_mode = importlib.import_module("core.entities.mcp_provider").IdentityMode + config_overrides(DEPLOYMENT_EDITION=DeploymentEdition.COMMUNITY) + + # Both an explicit idp_token request AND an inherited non-OFF current + # must collapse to OFF. + assert ( + controller_module._resolve_identity_mode(identity_mode.IDP_TOKEN, current=identity_mode.OFF) + == identity_mode.OFF + ) + assert controller_module._resolve_identity_mode(None, current=identity_mode.IDP_TOKEN) == identity_mode.OFF + + @staticmethod + def test_resolve_identity_mode_off_is_passthrough_when_not_enterprise( + controller_module, config_overrides: Callable[..., None] + ): + """OFF is always fine — the gate only neutralizes non-OFF values.""" + identity_mode = importlib.import_module("core.entities.mcp_provider").IdentityMode + config_overrides(DEPLOYMENT_EDITION=DeploymentEdition.COMMUNITY) + + assert controller_module._resolve_identity_mode(None, current=identity_mode.OFF) == identity_mode.OFF + + @staticmethod + @pytest.mark.parametrize( + ("query_visibility", "expected_stored"), + [ + ("only_me", "only_me"), + ("all_team_members", "all_team_members"), + ], + ) + def test_tool_plugin_oauth_url_stashes_valid_visibility( + app: Flask, + controller_module: ModuleType, + monkeypatch: pytest.MonkeyPatch, + query_visibility: str, + expected_stored: str, + ): + user = _mock_account("user-oauth") + _set_current_account(monkeypatch, controller_module, user, "tenant-oauth") + + with _mock_oauth_start_deps(controller_module, monkeypatch) as create_proxy_context: + _invoke_oauth_start( + controller_module, + app, + user, + "tenant-oauth", + query_string=f"visibility={query_visibility}", + ) + + _, kwargs = create_proxy_context.call_args + assert kwargs["extra_data"] == { + "visibility": expected_stored, + "plugin_unique_identifier": "langgenius/github:1.0.0@aaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaa", + } + + @staticmethod + @pytest.mark.parametrize("query_string", ["", "visibility=partial_members", "visibility=nonsense"]) + def test_tool_plugin_oauth_url_unknown_visibility_falls_back_to_only_me( + app: Flask, + controller_module: ModuleType, + monkeypatch: pytest.MonkeyPatch, + query_string: str, + ): + """Missing, partial_members (rejected for plugin creds), and garbage inputs + all collapse to only_me — OAuth tokens are personal by default.""" + user = _mock_account("user-oauth-fallback") + _set_current_account(monkeypatch, controller_module, user, "tenant-oauth") + + with _mock_oauth_start_deps(controller_module, monkeypatch) as create_proxy_context: + _invoke_oauth_start( + controller_module, + app, + user, + "tenant-oauth", + query_string=query_string, + ) + + _, kwargs = create_proxy_context.call_args + assert kwargs["extra_data"] == { + "visibility": "only_me", + "plugin_unique_identifier": "langgenius/github:1.0.0@aaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaa", + } + + @staticmethod + def test_tool_oauth_callback_rejects_replaced_package_before_dispatch( + app: Flask, controller_module: ModuleType, monkeypatch: pytest.MonkeyPatch + ) -> None: + context = { + "user_id": "user-oauth", + "tenant_id": "tenant-oauth", + "plugin_id": "langgenius/github", + "provider": "github", + "plugin_unique_identifier": f"langgenius/github:1.0.0@{'a' * 64}", + } + config = PluginOAuthClientConfig(f"langgenius/github:2.0.0@{'b' * 64}", {"client_secret": "test-secret"}) + monkeypatch.setattr(controller_module.OAuthProxyService, "use_proxy_context", lambda _context_id: context) + monkeypatch.setattr( + controller_module.BuiltinToolManageService, "get_oauth_client", lambda *_args, **_kwargs: config + ) + + def unexpected_dispatch(*_args: object, **_kwargs: object) -> None: + raise AssertionError("Changed packages must not receive OAuth credentials") + + monkeypatch.setattr(controller_module.OAuthHandler, "get_credentials", unexpected_dispatch) + api = controller_module.ToolOAuthCallback() + with app.test_request_context("/", headers={"Cookie": "context_id=ctx-id"}): + with pytest.raises(Forbidden, match="Plugin changed"): + unwrap(api.get)(api, provider=_OAUTH_PROVIDER_PATH) + + @staticmethod + @pytest.mark.parametrize( + ("stored", "expected"), + [ + ("only_me", "only_me"), + ("all_team_members", "all_team_members"), + ], + ) + def test_tool_oauth_callback_persists_stored_visibility( + app: Flask, + controller_module: ModuleType, + monkeypatch: pytest.MonkeyPatch, + stored: str, + expected: str, + ): + with _mock_oauth_callback_deps(controller_module, monkeypatch) as add_provider: + _invoke_oauth_callback(controller_module, app, stored_visibility=stored) + + _, kwargs = add_provider.call_args + assert kwargs["visibility"] == expected + assert kwargs["api_type"] == controller_module.CredentialType.OAUTH2 + + @staticmethod + @pytest.mark.parametrize("stored", [None, "partial_members", "garbage"]) + def test_tool_oauth_callback_unknown_stored_visibility_falls_back_to_only_me( + app: Flask, + controller_module: ModuleType, + monkeypatch: pytest.MonkeyPatch, + stored: str | None, + ): + """Older proxy contexts (pre-feature) and anything unexpected must default + to only_me — matches the safe-default we apply on the start endpoint.""" + with _mock_oauth_callback_deps(controller_module, monkeypatch) as add_provider: + _invoke_oauth_callback(controller_module, app, stored_visibility=stored) + + _, kwargs = add_provider.call_args + assert kwargs["visibility"] == "only_me" diff --git a/api/tests/unit_tests/controllers/console/workspace/test_trigger_provider_apis.py b/api/tests/unit_tests/controllers/console/workspace/test_trigger_provider_apis.py deleted file mode 100644 index e6c949eb3e5..00000000000 --- a/api/tests/unit_tests/controllers/console/workspace/test_trigger_provider_apis.py +++ /dev/null @@ -1,692 +0,0 @@ -"""Unit tests for controllers.console.workspace.trigger_providers endpoints.""" - -from __future__ import annotations - -from collections.abc import Iterator -from datetime import datetime -from importlib import import_module -from inspect import unwrap -from types import SimpleNamespace -from unittest.mock import ANY, MagicMock, create_autospec, patch - -import pytest -from flask import Flask -from werkzeug.exceptions import BadRequest, Forbidden - -from controllers.common.errors import NotFoundError -from controllers.console.workspace.trigger_providers import ( - TriggerOAuthAuthorizeApi, - TriggerOAuthCallbackApi, - TriggerOAuthClientManageApi, - TriggerOAuthClientPayload, - TriggerProviderIconApi, - TriggerProviderInfoApi, - TriggerProviderListApi, - TriggerSubscriptionBuilderBuildApi, - TriggerSubscriptionBuilderCreateApi, - TriggerSubscriptionBuilderCreatePayload, - TriggerSubscriptionBuilderGetApi, - TriggerSubscriptionBuilderLogsApi, - TriggerSubscriptionBuilderUpdateApi, - TriggerSubscriptionBuilderUpdatePayload, - TriggerSubscriptionBuilderVerifyApi, - TriggerSubscriptionBuilderVerifyPayload, - TriggerSubscriptionListApi, - TriggerSubscriptionUpdateApi, - TriggerSubscriptionVerifyApi, -) -from core.plugin.entities.plugin_daemon import CredentialType -from core.trigger.entities.api_entities import SubscriptionBuilderApiEntity, TriggerProviderApiEntity -from core.trigger.entities.entities import RequestLog -from models.account import Account -from services.credentials.query import CredentialQuery -from services.plugin.dispatch_identity import PluginOAuthClientConfig - -_OAUTH_IDENTIFIER = f"langgenius/github:1.0.0@{'a' * 64}" - - -@pytest.fixture -def credential_query() -> Iterator[MagicMock]: - query = create_autospec(CredentialQuery, instance=True, spec_set=True) - with patch( - "controllers.console.workspace.trigger_providers.application_services", - return_value=SimpleNamespace(credential_queries=query), - ): - yield query - - -def mock_user() -> Account: - user = Account(name="User", email="user.com") - user.id = "u1" - return user - - -def trigger_provider() -> TriggerProviderApiEntity: - return TriggerProviderApiEntity( - author="Dify", - name="github", - label={"en_US": "GitHub"}, - description={"en_US": "GitHub trigger provider"}, - icon="icon.svg", - icon_dark=None, - tags=["code"], - plugin_id="plugin", - plugin_unique_identifier="plugin:github", - supported_creation_methods=[], - subscription_constructor=None, - subscription_schema=[], - events=[], - ) - - -def subscription_builder() -> SubscriptionBuilderApiEntity: - return SubscriptionBuilderApiEntity( - id="b1", - name="Builder", - provider="github", - endpoint="b1", - parameters={"repo": "dify"}, - properties={"branch": "main"}, - credentials={"token": "secret"}, - credential_type=CredentialType.UNAUTHORIZED, - ) - - -def request_log() -> RequestLog: - return RequestLog( - id="log1", - endpoint="/hooks/b1", - request={"headers": {}, "body": {"event": "push"}}, - response={"status": 200, "body": {"ok": True}}, - created_at=datetime(2024, 1, 1), - ) - - -class TestTriggerProviderApis: - def test_icon_success(self, app: Flask) -> None: - api = TriggerProviderIconApi() - method = unwrap(api.get) - - with ( - app.test_request_context("/"), - patch( - "controllers.console.workspace.trigger_providers.TriggerManager.get_trigger_plugin_icon", - return_value="icon", - ), - ): - assert method(api, "t1", "github") == "icon" - - def test_list_providers(self, app: Flask) -> None: - api = TriggerProviderListApi() - method = unwrap(api.get) - - with ( - app.test_request_context("/"), - patch( - "controllers.console.workspace.trigger_providers.TriggerProviderService.list_trigger_providers", - return_value=[], - ), - ): - assert method(api, "t1") == [] - - def test_provider_info(self, app: Flask) -> None: - api = TriggerProviderInfoApi() - method = unwrap(api.get) - - with ( - app.test_request_context("/"), - patch( - "controllers.console.workspace.trigger_providers.TriggerProviderService.get_trigger_provider", - return_value=trigger_provider(), - ), - ): - assert method(api, "t1", "github")["name"] == "github" - - -class TestTriggerSubscriptionListApi: - def test_list_success(self, app: Flask, credential_query: MagicMock) -> None: - api = TriggerSubscriptionListApi() - method = unwrap(api.get) - - with ( - app.test_request_context("/"), - patch( - "controllers.console.workspace.trigger_providers.TriggerProviderService.list_trigger_provider_subscriptions", - return_value=[], - ) as list_subscriptions, - ): - assert method(api, "t1", mock_user(), "github") == [] - assert list_subscriptions.call_args.kwargs["credential_query"] is credential_query - - def test_list_invalid_provider(self, app: Flask, credential_query: MagicMock) -> None: - api = TriggerSubscriptionListApi() - method = unwrap(api.get) - - with ( - app.test_request_context("/"), - patch( - "controllers.console.workspace.trigger_providers.TriggerProviderService.list_trigger_provider_subscriptions", - side_effect=ValueError("bad"), - ) as list_subscriptions, - ): - result, status = method(api, "t1", mock_user(), "bad") - assert status == 404 - assert list_subscriptions.call_args.kwargs["credential_query"] is credential_query - - -class TestTriggerSubscriptionBuilderApis: - def test_create_builder(self, app: Flask) -> None: - api = TriggerSubscriptionBuilderCreateApi() - method = unwrap(api.post) - - with ( - app.test_request_context("/", json={"credential_type": "UNAUTHORIZED"}), - patch( - "controllers.console.workspace.trigger_providers.TriggerSubscriptionBuilderService.create_trigger_subscription_builder", - return_value=subscription_builder(), - ), - ): - result = method( - api, - TriggerSubscriptionBuilderCreatePayload(credential_type="UNAUTHORIZED"), - "t1", - mock_user(), - "github", - ) - assert result["subscription_builder"]["id"] == "b1" - - def test_get_builder(self, app: Flask) -> None: - api = TriggerSubscriptionBuilderGetApi() - method = unwrap(api.get) - - with ( - app.test_request_context("/"), - patch( - "controllers.console.workspace.trigger_providers.TriggerSubscriptionBuilderService.get_subscription_builder_by_id", - return_value=subscription_builder(), - ) as mock_get_builder, - ): - assert method(api, "t1", mock_user(), "github", "b1")["id"] == "b1" - mock_get_builder.assert_called_once_with( - tenant_id="t1", - user_id="u1", - provider_id=ANY, - subscription_builder_id="b1", - ) - - def test_verify_builder(self, app: Flask) -> None: - api = TriggerSubscriptionBuilderVerifyApi() - method = unwrap(api.post) - - with ( - app.test_request_context("/", json={"credentials": {"a": 1}}), - patch( - "controllers.console.workspace.trigger_providers.TriggerSubscriptionBuilderService.update_and_verify_builder", - return_value={"verified": True}, - ), - ): - assert method( - api, - TriggerSubscriptionBuilderVerifyPayload(credentials={"a": 1}), - "t1", - mock_user(), - "github", - "b1", - ) == {"verified": True} - - def test_verify_builder_error(self, app: Flask) -> None: - api = TriggerSubscriptionBuilderVerifyApi() - method = unwrap(api.post) - - with ( - app.test_request_context("/", json={"credentials": {}}), - patch( - "controllers.console.workspace.trigger_providers.TriggerSubscriptionBuilderService.update_and_verify_builder", - side_effect=Exception("err"), - ), - ): - with pytest.raises(ValueError, match="err"): - method( - api, - TriggerSubscriptionBuilderVerifyPayload(credentials={}), - "t1", - mock_user(), - "github", - "b1", - ) - - def test_update_builder(self, app: Flask) -> None: - api = TriggerSubscriptionBuilderUpdateApi() - method = unwrap(api.post) - - with ( - app.test_request_context("/", json={"name": "n"}), - patch( - "controllers.console.workspace.trigger_providers.TriggerSubscriptionBuilderService.update_trigger_subscription_builder", - return_value=subscription_builder(), - ) as mock_update_builder, - ): - assert ( - method( - api, - TriggerSubscriptionBuilderUpdatePayload(name="n"), - "t1", - mock_user(), - "github", - "b1", - )["id"] - == "b1" - ) - mock_update_builder.assert_called_once_with( - tenant_id="t1", - user_id="u1", - provider_id=ANY, - subscription_builder_id="b1", - subscription_builder_updater=ANY, - ) - - def test_logs(self, app: Flask) -> None: - api = TriggerSubscriptionBuilderLogsApi() - method = unwrap(api.get) - - with ( - app.test_request_context("/"), - patch( - "controllers.console.workspace.trigger_providers.TriggerSubscriptionBuilderService.list_logs", - return_value=[request_log()], - ) as mock_list_logs, - ): - result = method(api, "t1", mock_user(), "github", "b1") - assert result["logs"][0]["id"] == "log1" - mock_list_logs.assert_called_once_with( - tenant_id="t1", - user_id="u1", - provider_id=ANY, - subscription_builder_id="b1", - ) - - def test_build(self, app: Flask) -> None: - api = TriggerSubscriptionBuilderBuildApi() - method = unwrap(api.post) - - with ( - app.test_request_context("/", json={"name": "x"}), - patch( - "controllers.console.workspace.trigger_providers.TriggerSubscriptionBuilderService.update_and_build_builder", - return_value=None, - ), - ): - assert method( - api, - TriggerSubscriptionBuilderUpdatePayload(name="x"), - "t1", - mock_user(), - "github", - "b1", - ) == {"result": "success"} - - -class TestTriggerSubscriptionCrud: - def test_update_rename_only(self, app: Flask) -> None: - api = TriggerSubscriptionUpdateApi() - method = unwrap(api.post) - - sub = MagicMock() - sub.provider_id = "github" - sub.credential_type = CredentialType.UNAUTHORIZED - - with ( - app.test_request_context("/", json={"name": "x"}), - patch( - "controllers.console.workspace.trigger_providers.TriggerProviderService.get_subscription_by_id", - return_value=sub, - ), - patch("controllers.console.workspace.trigger_providers.TriggerProviderService.update_trigger_subscription"), - ): - assert method( - api, - TriggerSubscriptionBuilderUpdatePayload(name="x"), - "t1", - "s1", - ) == {"result": "success"} - - def test_update_not_found(self, app: Flask) -> None: - api = TriggerSubscriptionUpdateApi() - method = unwrap(api.post) - - with ( - app.test_request_context("/", json={"name": "x"}), - patch( - "controllers.console.workspace.trigger_providers.TriggerProviderService.get_subscription_by_id", - return_value=None, - ), - ): - with pytest.raises(NotFoundError): - method(api, TriggerSubscriptionBuilderUpdatePayload(name="x"), "t1", "x") - - def test_update_rebuild(self, app: Flask) -> None: - api = TriggerSubscriptionUpdateApi() - method = unwrap(api.post) - - sub = MagicMock() - sub.provider_id = "github" - sub.credential_type = CredentialType.OAUTH2 - sub.credentials = {"token": "old"} - sub.parameters = {"repo": "demo"} - - with ( - app.test_request_context("/", json={"credentials": {}}), - patch( - "controllers.console.workspace.trigger_providers.TriggerProviderService.get_subscription_by_id", - return_value=sub, - ), - patch( - "controllers.console.workspace.trigger_providers.TriggerProviderService.rebuild_trigger_subscription" - ), - ): - assert method( - api, - TriggerSubscriptionBuilderUpdatePayload(credentials={}), - "t1", - "s1", - ) == {"result": "success"} - - -class TestTriggerOAuthApis: - def test_oauth_callback_rejects_replaced_package_before_dispatch( - self, app: Flask, monkeypatch: pytest.MonkeyPatch - ) -> None: - context = { - "user_id": "u1", - "tenant_id": "t1", - "subscription_builder_id": "b1", - "plugin_id": "langgenius/github", - "provider": "github", - "plugin_unique_identifier": _OAUTH_IDENTIFIER, - } - config = PluginOAuthClientConfig(f"langgenius/github:2.0.0@{'b' * 64}", {"client_secret": "test-secret"}) - controller_module = import_module("controllers.console.workspace.trigger_providers") - monkeypatch.setattr( - controller_module.OAuthProxyService, - "use_proxy_context", - lambda _context_id: context, - ) - monkeypatch.setattr( - controller_module.TriggerProviderService, - "get_oauth_client", - lambda *_args, **_kwargs: config, - ) - - def unexpected_dispatch(*_args: object, **_kwargs: object) -> None: - raise AssertionError("Changed packages must not receive OAuth credentials") - - monkeypatch.setattr(controller_module.OAuthHandler, "get_credentials", unexpected_dispatch) - api = TriggerOAuthCallbackApi() - with app.test_request_context("/", headers={"Cookie": "context_id=ctx"}): - with pytest.raises(Forbidden, match="Plugin changed"): - unwrap(api.get)(api, "github") - - def test_oauth_authorize_success(self, app: Flask) -> None: - api = TriggerOAuthAuthorizeApi() - method = unwrap(api.get) - - with ( - app.test_request_context("/"), - patch( - "controllers.console.workspace.trigger_providers.TriggerProviderService.get_oauth_client", - return_value=PluginOAuthClientConfig(_OAUTH_IDENTIFIER, {"a": 1}), - ), - patch( - "controllers.console.workspace.trigger_providers.TriggerSubscriptionBuilderService.create_trigger_subscription_builder", - return_value=subscription_builder(), - ), - patch( - "controllers.console.workspace.trigger_providers.OAuthProxyService.create_proxy_context", - return_value="ctx", - ) as create_context, - patch( - "controllers.console.workspace.trigger_providers.OAuthHandler.get_authorization_url", - return_value=MagicMock(authorization_url="url"), - ) as authorize, - ): - resp = method(api, "t1", mock_user(), "github") - assert resp.status_code == 200 - assert create_context.call_args.kwargs["extra_data"]["plugin_unique_identifier"] == _OAUTH_IDENTIFIER - assert authorize.call_args.kwargs["plugin_unique_identifier"] == _OAUTH_IDENTIFIER - assert authorize.call_args.kwargs["system_credentials"] == {"a": 1} - - def test_oauth_authorize_no_client(self, app: Flask) -> None: - api = TriggerOAuthAuthorizeApi() - method = unwrap(api.get) - - with ( - app.test_request_context("/"), - patch( - "controllers.console.workspace.trigger_providers.TriggerProviderService.get_oauth_client", - return_value=None, - ), - ): - with pytest.raises(NotFoundError): - method(api, "t1", mock_user(), "github") - - def test_oauth_callback_forbidden(self, app: Flask) -> None: - api = TriggerOAuthCallbackApi() - method = unwrap(api.get) - - with app.test_request_context("/"): - with pytest.raises(Forbidden): - method(api, "github") - - def test_oauth_callback_success(self, app: Flask) -> None: - api = TriggerOAuthCallbackApi() - method = unwrap(api.get) - - ctx = { - "user_id": "u1", - "tenant_id": "t1", - "subscription_builder_id": "b1", - "plugin_id": "langgenius/github", - "provider": "github", - "plugin_unique_identifier": _OAUTH_IDENTIFIER, - } - - with ( - app.test_request_context("/", headers={"Cookie": "context_id=ctx"}), - patch( - "controllers.console.workspace.trigger_providers.OAuthProxyService.use_proxy_context", return_value=ctx - ), - patch( - "controllers.console.workspace.trigger_providers.TriggerProviderService.get_oauth_client", - return_value=PluginOAuthClientConfig(_OAUTH_IDENTIFIER, {"a": 1}), - ), - patch( - "controllers.console.workspace.trigger_providers.OAuthHandler.get_credentials", - return_value=MagicMock(credentials={"a": 1}, expires_at=1), - ) as get_credentials, - patch( - "controllers.console.workspace.trigger_providers.TriggerSubscriptionBuilderService.update_trigger_subscription_builder" - ) as mock_update_builder, - ): - resp = method(api, "github") - assert resp.status_code == 302 - assert get_credentials.call_args.kwargs["plugin_unique_identifier"] == _OAUTH_IDENTIFIER - assert get_credentials.call_args.kwargs["system_credentials"] == {"a": 1} - mock_update_builder.assert_called_once_with( - tenant_id="t1", - user_id="u1", - provider_id=ANY, - subscription_builder_id="b1", - subscription_builder_updater=ANY, - ) - - def test_oauth_callback_no_oauth_client(self, app: Flask) -> None: - api = TriggerOAuthCallbackApi() - method = unwrap(api.get) - - ctx = { - "user_id": "u1", - "tenant_id": "t1", - "subscription_builder_id": "b1", - } - - with ( - app.test_request_context("/", headers={"Cookie": "context_id=ctx"}), - patch( - "controllers.console.workspace.trigger_providers.OAuthProxyService.use_proxy_context", - return_value=ctx, - ), - patch( - "controllers.console.workspace.trigger_providers.TriggerProviderService.get_oauth_client", - return_value=None, - ), - ): - with pytest.raises(Forbidden): - method(api, "github") - - def test_oauth_callback_empty_credentials(self, app: Flask) -> None: - api = TriggerOAuthCallbackApi() - method = unwrap(api.get) - - ctx = { - "user_id": "u1", - "tenant_id": "t1", - "subscription_builder_id": "b1", - "plugin_id": "langgenius/github", - "provider": "github", - "plugin_unique_identifier": _OAUTH_IDENTIFIER, - } - - with ( - app.test_request_context("/", headers={"Cookie": "context_id=ctx"}), - patch( - "controllers.console.workspace.trigger_providers.OAuthProxyService.use_proxy_context", - return_value=ctx, - ), - patch( - "controllers.console.workspace.trigger_providers.TriggerProviderService.get_oauth_client", - return_value=PluginOAuthClientConfig(_OAUTH_IDENTIFIER, {"a": 1}), - ), - patch( - "controllers.console.workspace.trigger_providers.OAuthHandler.get_credentials", - return_value=MagicMock(credentials=None, expires_at=None), - ), - ): - with pytest.raises(ValueError, match="Failed to get OAuth credentials from the provider\\."): - method(api, "github") - - -class TestTriggerOAuthClientManageApi: - def test_get_client(self, app: Flask) -> None: - api = TriggerOAuthClientManageApi() - method = unwrap(api.get) - - with ( - app.test_request_context("/"), - patch( - "controllers.console.workspace.trigger_providers.TriggerProviderService.get_custom_oauth_client_params", - return_value={}, - ), - patch( - "controllers.console.workspace.trigger_providers.TriggerProviderService.is_oauth_custom_client_enabled", - return_value=False, - ), - patch( - "controllers.console.workspace.trigger_providers.TriggerProviderService.is_oauth_system_client_exists", - return_value=True, - ), - patch( - "controllers.console.workspace.trigger_providers.TriggerManager.get_trigger_provider", - return_value=MagicMock(get_oauth_client_schema=lambda: []), - ), - ): - result = method(api, "t1", "github") - assert "configured" in result - - def test_post_client(self, app: Flask) -> None: - api = TriggerOAuthClientManageApi() - method = unwrap(api.post) - - with ( - app.test_request_context("/", json={"enabled": True}), - patch( - "controllers.console.workspace.trigger_providers.TriggerProviderService.save_custom_oauth_client_params", - return_value={"result": "success"}, - ), - ): - assert method( - api, - TriggerOAuthClientPayload(enabled=True), - "t1", - "github", - ) == {"result": "success"} - - def test_delete_client(self, app: Flask) -> None: - api = TriggerOAuthClientManageApi() - method = unwrap(api.delete) - - with ( - app.test_request_context("/"), - patch( - "controllers.console.workspace.trigger_providers.TriggerProviderService.delete_custom_oauth_client_params", - return_value={"result": "success"}, - ), - ): - assert method(api, "t1", "github") == {"result": "success"} - - def test_oauth_client_post_value_error(self, app: Flask) -> None: - api = TriggerOAuthClientManageApi() - method = unwrap(api.post) - - with ( - app.test_request_context("/", json={"enabled": True}), - patch( - "controllers.console.workspace.trigger_providers.TriggerProviderService.save_custom_oauth_client_params", - side_effect=ValueError("bad"), - ), - ): - with pytest.raises(BadRequest): - method(api, TriggerOAuthClientPayload(enabled=True), "t1", "github") - - -class TestTriggerSubscriptionVerifyApi: - def test_verify_success(self, app: Flask) -> None: - api = TriggerSubscriptionVerifyApi() - method = unwrap(api.post) - - with ( - app.test_request_context("/", json={"credentials": {}}), - patch( - "controllers.console.workspace.trigger_providers.TriggerProviderService.verify_subscription_credentials", - return_value={"verified": True}, - ), - ): - assert method( - api, - TriggerSubscriptionBuilderVerifyPayload(credentials={}), - "t1", - mock_user(), - "github", - "s1", - ) == {"verified": True} - - @pytest.mark.parametrize("raised_exception", [ValueError("bad"), Exception("boom")]) - def test_verify_errors(self, app: Flask, raised_exception: Exception) -> None: - api = TriggerSubscriptionVerifyApi() - method = unwrap(api.post) - - with ( - app.test_request_context("/", json={"credentials": {}}), - patch( - "controllers.console.workspace.trigger_providers.TriggerProviderService.verify_subscription_credentials", - side_effect=raised_exception, - ), - ): - with pytest.raises(BadRequest): - method( - api, - TriggerSubscriptionBuilderVerifyPayload(credentials={}), - "t1", - mock_user(), - "github", - "s1", - ) diff --git a/api/tests/unit_tests/controllers/console/workspace/test_trigger_providers.py b/api/tests/unit_tests/controllers/console/workspace/test_trigger_providers.py index 660b275ce03..4b2eda34495 100644 --- a/api/tests/unit_tests/controllers/console/workspace/test_trigger_providers.py +++ b/api/tests/unit_tests/controllers/console/workspace/test_trigger_providers.py @@ -1,78 +1,761 @@ -"""Unit tests for trigger-provider controller transaction boundaries.""" +from __future__ import annotations from collections.abc import Iterator +from datetime import datetime from importlib import import_module from inspect import unwrap -from unittest.mock import patch +from types import SimpleNamespace +from unittest.mock import ANY, MagicMock, create_autospec, patch import pytest from flask import Flask from pydantic import ValidationError from sqlalchemy.engine import Engine -from werkzeug.exceptions import BadRequest +from werkzeug.exceptions import BadRequest, Forbidden +from controllers.common.errors import NotFoundError from controllers.console.workspace.trigger_providers import ( + TriggerOAuthAuthorizeApi, + TriggerOAuthCallbackApi, + TriggerOAuthClientManageApi, + TriggerOAuthClientPayload, + TriggerProviderIconApi, + TriggerProviderInfoApi, + TriggerProviderListApi, + TriggerSubscriptionBuilderBuildApi, + TriggerSubscriptionBuilderCreateApi, + TriggerSubscriptionBuilderCreatePayload, + TriggerSubscriptionBuilderGetApi, + TriggerSubscriptionBuilderLogsApi, + TriggerSubscriptionBuilderUpdateApi, + TriggerSubscriptionBuilderUpdatePayload, + TriggerSubscriptionBuilderVerifyApi, TriggerSubscriptionBuilderVerifyPayload, TriggerSubscriptionDeleteApi, + TriggerSubscriptionListApi, + TriggerSubscriptionUpdateApi, + TriggerSubscriptionVerifyApi, TriggerSubscriptionVerifyPayload, ) +from core.plugin.entities.plugin_daemon import CredentialType +from core.trigger.entities.api_entities import SubscriptionBuilderApiEntity, TriggerProviderApiEntity +from core.trigger.entities.entities import RequestLog +from models.account import Account from models.engine import db +from services.credentials.query import CredentialQuery +from services.plugin.dispatch_identity import PluginOAuthClientConfig + +"""Unit tests for controllers.console.workspace.trigger_providers endpoints.""" + + +_OAUTH_IDENTIFIER = f"langgenius/github:1.0.0@{'a' * 64}" + + +def mock_user() -> Account: + user = Account(name="User", email="user.com") + user.id = "u1" + return user + + +def trigger_provider() -> TriggerProviderApiEntity: + return TriggerProviderApiEntity( + author="Dify", + name="github", + label={"en_US": "GitHub"}, + description={"en_US": "GitHub trigger provider"}, + icon="icon.svg", + icon_dark=None, + tags=["code"], + plugin_id="plugin", + plugin_unique_identifier="plugin:github", + supported_creation_methods=[], + subscription_constructor=None, + subscription_schema=[], + events=[], + ) + + +def subscription_builder() -> SubscriptionBuilderApiEntity: + return SubscriptionBuilderApiEntity( + id="b1", + name="Builder", + provider="github", + endpoint="b1", + parameters={"repo": "dify"}, + properties={"branch": "main"}, + credentials={"token": "secret"}, + credential_type=CredentialType.UNAUTHORIZED, + ) + + +def request_log() -> RequestLog: + return RequestLog( + id="log1", + endpoint="/hooks/b1", + request={"headers": {}, "body": {"event": "push"}}, + response={"status": 200, "body": {"ok": True}}, + created_at=datetime(2024, 1, 1), + ) + + +class TestTriggerProviderApisCases: + @staticmethod + @pytest.fixture + def credential_query() -> Iterator[MagicMock]: + query = create_autospec(CredentialQuery, instance=True, spec_set=True) + with patch( + "controllers.console.workspace.trigger_providers.application_services", + return_value=SimpleNamespace(credential_queries=query), + ): + yield query + + class TestTriggerProviderApis: + def test_icon_success(self, app: Flask) -> None: + api = TriggerProviderIconApi() + method = unwrap(api.get) + + with ( + app.test_request_context("/"), + patch( + "controllers.console.workspace.trigger_providers.TriggerManager.get_trigger_plugin_icon", + return_value="icon", + ), + ): + assert method(api, "t1", "github") == "icon" + + def test_list_providers(self, app: Flask) -> None: + api = TriggerProviderListApi() + method = unwrap(api.get) + + with ( + app.test_request_context("/"), + patch( + "controllers.console.workspace.trigger_providers.TriggerProviderService.list_trigger_providers", + return_value=[], + ), + ): + assert method(api, "t1") == [] + + def test_provider_info(self, app: Flask) -> None: + api = TriggerProviderInfoApi() + method = unwrap(api.get) + + with ( + app.test_request_context("/"), + patch( + "controllers.console.workspace.trigger_providers.TriggerProviderService.get_trigger_provider", + return_value=trigger_provider(), + ), + ): + assert method(api, "t1", "github")["name"] == "github" + + class TestTriggerSubscriptionListApi: + def test_list_success(self, app: Flask, credential_query: MagicMock) -> None: + api = TriggerSubscriptionListApi() + method = unwrap(api.get) + + with ( + app.test_request_context("/"), + patch( + "controllers.console.workspace.trigger_providers.TriggerProviderService.list_trigger_provider_subscriptions", + return_value=[], + ) as list_subscriptions, + ): + assert method(api, "t1", mock_user(), "github") == [] + assert list_subscriptions.call_args.kwargs["credential_query"] is credential_query + + def test_list_invalid_provider(self, app: Flask, credential_query: MagicMock) -> None: + api = TriggerSubscriptionListApi() + method = unwrap(api.get) + + with ( + app.test_request_context("/"), + patch( + "controllers.console.workspace.trigger_providers.TriggerProviderService.list_trigger_provider_subscriptions", + side_effect=ValueError("bad"), + ) as list_subscriptions, + ): + result, status = method(api, "t1", mock_user(), "bad") + assert status == 404 + assert list_subscriptions.call_args.kwargs["credential_query"] is credential_query + + class TestTriggerSubscriptionBuilderApis: + def test_create_builder(self, app: Flask) -> None: + api = TriggerSubscriptionBuilderCreateApi() + method = unwrap(api.post) + + with ( + app.test_request_context("/", json={"credential_type": "UNAUTHORIZED"}), + patch( + "controllers.console.workspace.trigger_providers.TriggerSubscriptionBuilderService.create_trigger_subscription_builder", + return_value=subscription_builder(), + ), + ): + result = method( + api, + TriggerSubscriptionBuilderCreatePayload(credential_type="UNAUTHORIZED"), + "t1", + mock_user(), + "github", + ) + assert result["subscription_builder"]["id"] == "b1" + + def test_get_builder(self, app: Flask) -> None: + api = TriggerSubscriptionBuilderGetApi() + method = unwrap(api.get) + + with ( + app.test_request_context("/"), + patch( + "controllers.console.workspace.trigger_providers.TriggerSubscriptionBuilderService.get_subscription_builder_by_id", + return_value=subscription_builder(), + ) as mock_get_builder, + ): + assert method(api, "t1", mock_user(), "github", "b1")["id"] == "b1" + mock_get_builder.assert_called_once_with( + tenant_id="t1", + user_id="u1", + provider_id=ANY, + subscription_builder_id="b1", + ) + + def test_verify_builder(self, app: Flask) -> None: + api = TriggerSubscriptionBuilderVerifyApi() + method = unwrap(api.post) + + with ( + app.test_request_context("/", json={"credentials": {"a": 1}}), + patch( + "controllers.console.workspace.trigger_providers.TriggerSubscriptionBuilderService.update_and_verify_builder", + return_value={"verified": True}, + ), + ): + assert method( + api, + TriggerSubscriptionBuilderVerifyPayload(credentials={"a": 1}), + "t1", + mock_user(), + "github", + "b1", + ) == {"verified": True} + + def test_verify_builder_error(self, app: Flask) -> None: + api = TriggerSubscriptionBuilderVerifyApi() + method = unwrap(api.post) + + with ( + app.test_request_context("/", json={"credentials": {}}), + patch( + "controllers.console.workspace.trigger_providers.TriggerSubscriptionBuilderService.update_and_verify_builder", + side_effect=Exception("err"), + ), + ): + with pytest.raises(ValueError, match="err"): + method( + api, + TriggerSubscriptionBuilderVerifyPayload(credentials={}), + "t1", + mock_user(), + "github", + "b1", + ) + + def test_update_builder(self, app: Flask) -> None: + api = TriggerSubscriptionBuilderUpdateApi() + method = unwrap(api.post) + + with ( + app.test_request_context("/", json={"name": "n"}), + patch( + "controllers.console.workspace.trigger_providers.TriggerSubscriptionBuilderService.update_trigger_subscription_builder", + return_value=subscription_builder(), + ) as mock_update_builder, + ): + assert ( + method( + api, + TriggerSubscriptionBuilderUpdatePayload(name="n"), + "t1", + mock_user(), + "github", + "b1", + )["id"] + == "b1" + ) + mock_update_builder.assert_called_once_with( + tenant_id="t1", + user_id="u1", + provider_id=ANY, + subscription_builder_id="b1", + subscription_builder_updater=ANY, + ) + + def test_logs(self, app: Flask) -> None: + api = TriggerSubscriptionBuilderLogsApi() + method = unwrap(api.get) + + with ( + app.test_request_context("/"), + patch( + "controllers.console.workspace.trigger_providers.TriggerSubscriptionBuilderService.list_logs", + return_value=[request_log()], + ) as mock_list_logs, + ): + result = method(api, "t1", mock_user(), "github", "b1") + assert result["logs"][0]["id"] == "log1" + mock_list_logs.assert_called_once_with( + tenant_id="t1", + user_id="u1", + provider_id=ANY, + subscription_builder_id="b1", + ) + + def test_build(self, app: Flask) -> None: + api = TriggerSubscriptionBuilderBuildApi() + method = unwrap(api.post) + + with ( + app.test_request_context("/", json={"name": "x"}), + patch( + "controllers.console.workspace.trigger_providers.TriggerSubscriptionBuilderService.update_and_build_builder", + return_value=None, + ), + ): + assert method( + api, + TriggerSubscriptionBuilderUpdatePayload(name="x"), + "t1", + mock_user(), + "github", + "b1", + ) == {"result": "success"} + + class TestTriggerSubscriptionCrud: + def test_update_rename_only(self, app: Flask) -> None: + api = TriggerSubscriptionUpdateApi() + method = unwrap(api.post) + + sub = MagicMock() + sub.provider_id = "github" + sub.credential_type = CredentialType.UNAUTHORIZED + + with ( + app.test_request_context("/", json={"name": "x"}), + patch( + "controllers.console.workspace.trigger_providers.TriggerProviderService.get_subscription_by_id", + return_value=sub, + ), + patch( + "controllers.console.workspace.trigger_providers.TriggerProviderService.update_trigger_subscription" + ), + ): + assert method( + api, + TriggerSubscriptionBuilderUpdatePayload(name="x"), + "t1", + "s1", + ) == {"result": "success"} + + def test_update_not_found(self, app: Flask) -> None: + api = TriggerSubscriptionUpdateApi() + method = unwrap(api.post) + + with ( + app.test_request_context("/", json={"name": "x"}), + patch( + "controllers.console.workspace.trigger_providers.TriggerProviderService.get_subscription_by_id", + return_value=None, + ), + ): + with pytest.raises(NotFoundError): + method(api, TriggerSubscriptionBuilderUpdatePayload(name="x"), "t1", "x") + + def test_update_rebuild(self, app: Flask) -> None: + api = TriggerSubscriptionUpdateApi() + method = unwrap(api.post) + + sub = MagicMock() + sub.provider_id = "github" + sub.credential_type = CredentialType.OAUTH2 + sub.credentials = {"token": "old"} + sub.parameters = {"repo": "demo"} + + with ( + app.test_request_context("/", json={"credentials": {}}), + patch( + "controllers.console.workspace.trigger_providers.TriggerProviderService.get_subscription_by_id", + return_value=sub, + ), + patch( + "controllers.console.workspace.trigger_providers.TriggerProviderService.rebuild_trigger_subscription" + ), + ): + assert method( + api, + TriggerSubscriptionBuilderUpdatePayload(credentials={}), + "t1", + "s1", + ) == {"result": "success"} + + class TestTriggerOAuthApis: + def test_oauth_callback_rejects_replaced_package_before_dispatch( + self, app: Flask, monkeypatch: pytest.MonkeyPatch + ) -> None: + context = { + "user_id": "u1", + "tenant_id": "t1", + "subscription_builder_id": "b1", + "plugin_id": "langgenius/github", + "provider": "github", + "plugin_unique_identifier": _OAUTH_IDENTIFIER, + } + config = PluginOAuthClientConfig(f"langgenius/github:2.0.0@{'b' * 64}", {"client_secret": "test-secret"}) + controller_module = import_module("controllers.console.workspace.trigger_providers") + monkeypatch.setattr( + controller_module.OAuthProxyService, + "use_proxy_context", + lambda _context_id: context, + ) + monkeypatch.setattr( + controller_module.TriggerProviderService, + "get_oauth_client", + lambda *_args, **_kwargs: config, + ) + + def unexpected_dispatch(*_args: object, **_kwargs: object) -> None: + raise AssertionError("Changed packages must not receive OAuth credentials") + + monkeypatch.setattr(controller_module.OAuthHandler, "get_credentials", unexpected_dispatch) + api = TriggerOAuthCallbackApi() + with app.test_request_context("/", headers={"Cookie": "context_id=ctx"}): + with pytest.raises(Forbidden, match="Plugin changed"): + unwrap(api.get)(api, "github") + + def test_oauth_authorize_success(self, app: Flask) -> None: + api = TriggerOAuthAuthorizeApi() + method = unwrap(api.get) + + with ( + app.test_request_context("/"), + patch( + "controllers.console.workspace.trigger_providers.TriggerProviderService.get_oauth_client", + return_value=PluginOAuthClientConfig(_OAUTH_IDENTIFIER, {"a": 1}), + ), + patch( + "controllers.console.workspace.trigger_providers.TriggerSubscriptionBuilderService.create_trigger_subscription_builder", + return_value=subscription_builder(), + ), + patch( + "controllers.console.workspace.trigger_providers.OAuthProxyService.create_proxy_context", + return_value="ctx", + ) as create_context, + patch( + "controllers.console.workspace.trigger_providers.OAuthHandler.get_authorization_url", + return_value=MagicMock(authorization_url="url"), + ) as authorize, + ): + resp = method(api, "t1", mock_user(), "github") + assert resp.status_code == 200 + assert create_context.call_args.kwargs["extra_data"]["plugin_unique_identifier"] == _OAUTH_IDENTIFIER + assert authorize.call_args.kwargs["plugin_unique_identifier"] == _OAUTH_IDENTIFIER + assert authorize.call_args.kwargs["system_credentials"] == {"a": 1} + + def test_oauth_authorize_no_client(self, app: Flask) -> None: + api = TriggerOAuthAuthorizeApi() + method = unwrap(api.get) + + with ( + app.test_request_context("/"), + patch( + "controllers.console.workspace.trigger_providers.TriggerProviderService.get_oauth_client", + return_value=None, + ), + ): + with pytest.raises(NotFoundError): + method(api, "t1", mock_user(), "github") + + def test_oauth_callback_forbidden(self, app: Flask) -> None: + api = TriggerOAuthCallbackApi() + method = unwrap(api.get) + + with app.test_request_context("/"): + with pytest.raises(Forbidden): + method(api, "github") + + def test_oauth_callback_success(self, app: Flask) -> None: + api = TriggerOAuthCallbackApi() + method = unwrap(api.get) + + ctx = { + "user_id": "u1", + "tenant_id": "t1", + "subscription_builder_id": "b1", + "plugin_id": "langgenius/github", + "provider": "github", + "plugin_unique_identifier": _OAUTH_IDENTIFIER, + } + + with ( + app.test_request_context("/", headers={"Cookie": "context_id=ctx"}), + patch( + "controllers.console.workspace.trigger_providers.OAuthProxyService.use_proxy_context", + return_value=ctx, + ), + patch( + "controllers.console.workspace.trigger_providers.TriggerProviderService.get_oauth_client", + return_value=PluginOAuthClientConfig(_OAUTH_IDENTIFIER, {"a": 1}), + ), + patch( + "controllers.console.workspace.trigger_providers.OAuthHandler.get_credentials", + return_value=MagicMock(credentials={"a": 1}, expires_at=1), + ) as get_credentials, + patch( + "controllers.console.workspace.trigger_providers.TriggerSubscriptionBuilderService.update_trigger_subscription_builder" + ) as mock_update_builder, + ): + resp = method(api, "github") + assert resp.status_code == 302 + assert get_credentials.call_args.kwargs["plugin_unique_identifier"] == _OAUTH_IDENTIFIER + assert get_credentials.call_args.kwargs["system_credentials"] == {"a": 1} + mock_update_builder.assert_called_once_with( + tenant_id="t1", + user_id="u1", + provider_id=ANY, + subscription_builder_id="b1", + subscription_builder_updater=ANY, + ) + + def test_oauth_callback_no_oauth_client(self, app: Flask) -> None: + api = TriggerOAuthCallbackApi() + method = unwrap(api.get) + + ctx = { + "user_id": "u1", + "tenant_id": "t1", + "subscription_builder_id": "b1", + } + + with ( + app.test_request_context("/", headers={"Cookie": "context_id=ctx"}), + patch( + "controllers.console.workspace.trigger_providers.OAuthProxyService.use_proxy_context", + return_value=ctx, + ), + patch( + "controllers.console.workspace.trigger_providers.TriggerProviderService.get_oauth_client", + return_value=None, + ), + ): + with pytest.raises(Forbidden): + method(api, "github") + + def test_oauth_callback_empty_credentials(self, app: Flask) -> None: + api = TriggerOAuthCallbackApi() + method = unwrap(api.get) + + ctx = { + "user_id": "u1", + "tenant_id": "t1", + "subscription_builder_id": "b1", + "plugin_id": "langgenius/github", + "provider": "github", + "plugin_unique_identifier": _OAUTH_IDENTIFIER, + } + + with ( + app.test_request_context("/", headers={"Cookie": "context_id=ctx"}), + patch( + "controllers.console.workspace.trigger_providers.OAuthProxyService.use_proxy_context", + return_value=ctx, + ), + patch( + "controllers.console.workspace.trigger_providers.TriggerProviderService.get_oauth_client", + return_value=PluginOAuthClientConfig(_OAUTH_IDENTIFIER, {"a": 1}), + ), + patch( + "controllers.console.workspace.trigger_providers.OAuthHandler.get_credentials", + return_value=MagicMock(credentials=None, expires_at=None), + ), + ): + with pytest.raises(ValueError, match="Failed to get OAuth credentials from the provider\\."): + method(api, "github") + + class TestTriggerOAuthClientManageApi: + def test_get_client(self, app: Flask) -> None: + api = TriggerOAuthClientManageApi() + method = unwrap(api.get) + + with ( + app.test_request_context("/"), + patch( + "controllers.console.workspace.trigger_providers.TriggerProviderService.get_custom_oauth_client_params", + return_value={}, + ), + patch( + "controllers.console.workspace.trigger_providers.TriggerProviderService.is_oauth_custom_client_enabled", + return_value=False, + ), + patch( + "controllers.console.workspace.trigger_providers.TriggerProviderService.is_oauth_system_client_exists", + return_value=True, + ), + patch( + "controllers.console.workspace.trigger_providers.TriggerManager.get_trigger_provider", + return_value=MagicMock(get_oauth_client_schema=lambda: []), + ), + ): + result = method(api, "t1", "github") + assert "configured" in result + + def test_post_client(self, app: Flask) -> None: + api = TriggerOAuthClientManageApi() + method = unwrap(api.post) + + with ( + app.test_request_context("/", json={"enabled": True}), + patch( + "controllers.console.workspace.trigger_providers.TriggerProviderService.save_custom_oauth_client_params", + return_value={"result": "success"}, + ), + ): + assert method( + api, + TriggerOAuthClientPayload(enabled=True), + "t1", + "github", + ) == {"result": "success"} + + def test_delete_client(self, app: Flask) -> None: + api = TriggerOAuthClientManageApi() + method = unwrap(api.delete) + + with ( + app.test_request_context("/"), + patch( + "controllers.console.workspace.trigger_providers.TriggerProviderService.delete_custom_oauth_client_params", + return_value={"result": "success"}, + ), + ): + assert method(api, "t1", "github") == {"result": "success"} + + def test_oauth_client_post_value_error(self, app: Flask) -> None: + api = TriggerOAuthClientManageApi() + method = unwrap(api.post) + + with ( + app.test_request_context("/", json={"enabled": True}), + patch( + "controllers.console.workspace.trigger_providers.TriggerProviderService.save_custom_oauth_client_params", + side_effect=ValueError("bad"), + ), + ): + with pytest.raises(BadRequest): + method(api, TriggerOAuthClientPayload(enabled=True), "t1", "github") + + class TestTriggerSubscriptionVerifyApi: + def test_verify_success(self, app: Flask) -> None: + api = TriggerSubscriptionVerifyApi() + method = unwrap(api.post) + + with ( + app.test_request_context("/", json={"credentials": {}}), + patch( + "controllers.console.workspace.trigger_providers.TriggerProviderService.verify_subscription_credentials", + return_value={"verified": True}, + ), + ): + assert method( + api, + TriggerSubscriptionBuilderVerifyPayload(credentials={}), + "t1", + mock_user(), + "github", + "s1", + ) == {"verified": True} + + @pytest.mark.parametrize("raised_exception", [ValueError("bad"), Exception("boom")]) + def test_verify_errors(self, app: Flask, raised_exception: Exception) -> None: + api = TriggerSubscriptionVerifyApi() + method = unwrap(api.post) + + with ( + app.test_request_context("/", json={"credentials": {}}), + patch( + "controllers.console.workspace.trigger_providers.TriggerProviderService.verify_subscription_credentials", + side_effect=raised_exception, + ), + ): + with pytest.raises(BadRequest): + method( + api, + TriggerSubscriptionBuilderVerifyPayload(credentials={}), + "t1", + mock_user(), + "github", + "s1", + ) + + +"""Unit tests for trigger-provider controller transaction boundaries.""" + trigger_provider_module = import_module("controllers.console.workspace.trigger_providers") -@pytest.fixture -def flask_sqlite_engine() -> Iterator[Engine]: - app = Flask(__name__) - app.config["SQLALCHEMY_DATABASE_URI"] = "sqlite:///:memory:" - db.init_app(app) +class TestTriggerProvidersCases: + @staticmethod + @pytest.fixture + def flask_sqlite_engine() -> Iterator[Engine]: + app = Flask(__name__) + app.config["SQLALCHEMY_DATABASE_URI"] = "sqlite:///:memory:" + db.init_app(app) - with app.app_context(): - yield db.engine + with app.app_context(): + yield db.engine + @staticmethod + def test_delete_subscription_uses_sqlite_transaction(flask_sqlite_engine: Engine) -> None: + api = TriggerSubscriptionDeleteApi() + method = unwrap(api.post) -def test_delete_subscription_uses_sqlite_transaction(flask_sqlite_engine: Engine) -> None: - api = TriggerSubscriptionDeleteApi() - method = unwrap(api.post) + with ( + patch.object(trigger_provider_module.TriggerProviderService, "delete_trigger_provider") as delete_provider, + patch.object( + trigger_provider_module.TriggerSubscriptionOperatorService, + "delete_plugin_trigger_by_subscription", + ) as delete_triggers, + ): + result = method(api, "t1", "sub1") - with ( - patch.object(trigger_provider_module.TriggerProviderService, "delete_trigger_provider") as delete_provider, - patch.object( - trigger_provider_module.TriggerSubscriptionOperatorService, - "delete_plugin_trigger_by_subscription", - ) as delete_triggers, - ): - result = method(api, "t1", "sub1") + assert result == {"result": "success"} + provider_session = delete_provider.call_args.kwargs["session"] + trigger_session = delete_triggers.call_args.kwargs["session"] + assert provider_session is trigger_session + assert provider_session.get_bind() is flask_sqlite_engine - assert result == {"result": "success"} - provider_session = delete_provider.call_args.kwargs["session"] - trigger_session = delete_triggers.call_args.kwargs["session"] - assert provider_session is trigger_session - assert provider_session.get_bind() is flask_sqlite_engine + @staticmethod + @pytest.mark.usefixtures("flask_sqlite_engine") + def test_delete_subscription_translates_value_error() -> None: + api = TriggerSubscriptionDeleteApi() + method = unwrap(api.post) + with ( + patch.object( + trigger_provider_module.TriggerProviderService, + "delete_trigger_provider", + side_effect=ValueError("bad"), + ), + pytest.raises(BadRequest, match="bad"), + ): + method(api, "t1", "sub1") -@pytest.mark.usefixtures("flask_sqlite_engine") -def test_delete_subscription_translates_value_error() -> None: - api = TriggerSubscriptionDeleteApi() - method = unwrap(api.post) + @staticmethod + def test_builder_verify_payload_allows_credentials_to_be_omitted() -> None: + payload = TriggerSubscriptionBuilderVerifyPayload.model_validate({}) - with ( - patch.object( - trigger_provider_module.TriggerProviderService, - "delete_trigger_provider", - side_effect=ValueError("bad"), - ), - pytest.raises(BadRequest, match="bad"), - ): - method(api, "t1", "sub1") + assert payload.credentials is None - -def test_builder_verify_payload_allows_credentials_to_be_omitted() -> None: - payload = TriggerSubscriptionBuilderVerifyPayload.model_validate({}) - - assert payload.credentials is None - - -def test_subscription_verify_payload_requires_credentials() -> None: - with pytest.raises(ValidationError): - TriggerSubscriptionVerifyPayload.model_validate({}) + @staticmethod + def test_subscription_verify_payload_requires_credentials() -> None: + with pytest.raises(ValidationError): + TriggerSubscriptionVerifyPayload.model_validate({}) diff --git a/api/tests/unit_tests/controllers/console/workspace/test_workspace.py b/api/tests/unit_tests/controllers/console/workspace/test_workspace.py index 5a02662df81..00def9f4310 100644 --- a/api/tests/unit_tests/controllers/console/workspace/test_workspace.py +++ b/api/tests/unit_tests/controllers/console/workspace/test_workspace.py @@ -1,5 +1,3 @@ -"""Controller contracts: parsing, delegation, serialization and HTTP errors.""" - from collections.abc import Callable from dataclasses import dataclass from datetime import datetime @@ -11,7 +9,7 @@ from unittest.mock import Mock, create_autospec import pytest from flask import Flask from flask_restx import Resource -from werkzeug.exceptions import NotFound +from werkzeug.exceptions import Forbidden, NotFound, Unauthorized from controllers.common.errors import ( FilenameNotExistsError, @@ -20,12 +18,18 @@ from controllers.common.errors import ( TooManyFilesError, UnsupportedFileTypeError, ) -from controllers.console import console_ns -from controllers.console.error import AccountNotLinkTenantError +from controllers.console import console_ns, flask_admission, wraps +from controllers.console.error import AccountNotLinkTenantError, NotInitValidateError, NotSetupError +from controllers.console.notification import NotificationApi from controllers.console.workspace import workspace as controller -from controllers.console.workspace.error import CurrentWorkspaceArchivedError -from enums import CloudPlan +from controllers.console.workspace.error import AccountNotInitializedError, CurrentWorkspaceArchivedError +from controllers.console.workspace.workspace import WorkspacePermissionApi +from enums import CloudPlan, DeploymentEdition +from libs import login as login_adapter +from libs.login import AccountWithTenant from machinery.context import RequestContext +from models.account import Account, AccountStatus +from services.entities.feature_entities import FeatureModel from services.errors import file as file_errors from services.errors.workspace import WorkspaceArchivedError, WorkspaceNotFoundError, WorkspaceNotLinkedError from services.workspace.contracts import ( @@ -39,6 +43,111 @@ from services.workspace.contracts import ( from services.workspace.provisioning_service import WorkspaceProvisioningService from services.workspace.service import WorkspaceQueryService, WorkspaceService +"""Workspace admission composes existing checks before controller execution.""" + + +class TestWorkspaceAdmissionCases: + @staticmethod + def test_admin_key_admits_without_account_login(app: Flask, config_overrides: Callable[..., None]) -> None: + config_overrides(DEPLOYMENT_EDITION=DeploymentEdition.CLOUD, ADMIN_API_KEY="admin") + + class Handler: + @flask_admission.console_admin_admission + def get(self, *, page: int) -> int: + return page + + with app.test_request_context(headers={"Authorization": "Bearer admin"}): + assert Handler().get(page=2) == 2 + with app.test_request_context(), pytest.raises(Unauthorized): + Handler().get(page=2) + + @staticmethod + @pytest.mark.parametrize( + ("edition", "enabled", "allowed"), + [ + (DeploymentEdition.CLOUD, True, True), + (DeploymentEdition.CLOUD, False, False), + (DeploymentEdition.COMMUNITY, False, True), + ], + ) + def test_custom_brand_billing_admission( + app: Flask, + monkeypatch: pytest.MonkeyPatch, + config_overrides: Callable[..., None], + edition: DeploymentEdition, + enabled: bool, + allowed: bool, + ) -> None: + config_overrides(DEPLOYMENT_EDITION=edition) + monkeypatch.setattr(flask_admission, "setup_required", lambda view: view) + monkeypatch.setattr(flask_admission, "login_required", lambda view: view) + monkeypatch.setattr(flask_admission, "account_initialization_required", lambda view: view) + account = Account(name="User", email="user@example.com") + identity = AccountWithTenant(account, "workspace") + monkeypatch.setattr(flask_admission, "current_account_with_tenant", lambda: identity) + monkeypatch.setattr(wraps, "current_account_with_tenant", lambda: identity) + features = Mock(return_value=FeatureModel(can_replace_logo=enabled)) + monkeypatch.setattr(wraps.FeatureService, "get_features", features) + reached = [] + + class Handler: + @flask_admission.console_account_admission(billing_resource="workspace_custom") + def post(self, context: RequestContext) -> RequestContext: + reached.append(context) + return context + + with app.test_request_context(): + if allowed: + context = Handler().post() + assert isinstance(context, RequestContext) + assert context.active_workspace_id == "workspace" + else: + with pytest.raises(Forbidden): + Handler().post() + assert bool(reached) is allowed + assert features.call_count == int(edition == DeploymentEdition.CLOUD) + + @staticmethod + @pytest.mark.parametrize("controller_type", [WorkspacePermissionApi, NotificationApi]) + @pytest.mark.parametrize( + ("setup_complete", "init_password", "status", "error"), + [ + (False, "", None, NotSetupError), + (False, "initial-password", None, NotInitValidateError), + (True, "", None, Unauthorized), + (True, "", AccountStatus.UNINITIALIZED, AccountNotInitializedError), + (True, "", AccountStatus.ACTIVE, NotFound), + ], + ) + def test_account_admission_checks_identity_before_edition( + app: Flask, + monkeypatch: pytest.MonkeyPatch, + config_overrides: Callable[..., None], + controller_type: type[WorkspacePermissionApi] | type[NotificationApi], + setup_complete: bool, + init_password: str, + status: AccountStatus | None, + error: type[Exception], + ) -> None: + config_overrides( + DEPLOYMENT_EDITION=DeploymentEdition.COMMUNITY, INIT_PASSWORD=init_password, LOGIN_DISABLED=False + ) + monkeypatch.setattr(wraps, "_is_setup_completed", lambda: setup_complete) + account = Account(name="User", email="user@example.com", status=status or AccountStatus.ACTIVE) + identity = Mock(return_value=account, side_effect=Unauthorized() if status is None else None) + monkeypatch.setattr(login_adapter, "_resolve_current_user", identity) + monkeypatch.setattr(login_adapter, "check_csrf_token", Mock()) + monkeypatch.setattr(wraps, "current_account_with_tenant", lambda: AccountWithTenant(account, "workspace")) + + with app.test_request_context(), pytest.raises(error): + controller_type().get() + + assert identity.call_count == int(setup_complete) + + +"""Controller contracts: parsing, delegation, serialization and HTTP errors.""" + + CONTEXT = RequestContext("request-1", "trace-1", "account-1", "workspace-1") NOW = datetime(2026, 1, 1) @@ -55,201 +164,208 @@ class Services: workspaces: WorkspaceMocks -@pytest.fixture -def workspaces(monkeypatch: pytest.MonkeyPatch) -> WorkspaceMocks: - services = WorkspaceMocks( - management=create_autospec(WorkspaceService, instance=True), - queries=create_autospec(WorkspaceQueryService, instance=True), - provisioning=create_autospec(WorkspaceProvisioningService, instance=True), - ) - monkeypatch.setattr(controller, "application_services", lambda: Services(services)) - return services +class TestWorkspaceCases: + @staticmethod + @pytest.fixture + def workspaces(monkeypatch: pytest.MonkeyPatch) -> WorkspaceMocks: + services = WorkspaceMocks( + management=create_autospec(WorkspaceService, instance=True), + queries=create_autospec(WorkspaceQueryService, instance=True), + provisioning=create_autospec(WorkspaceProvisioningService, instance=True), + ) + monkeypatch.setattr(controller, "application_services", lambda: Services(services)) + return services + @staticmethod + def test_workspace_list_serialization(workspaces: WorkspaceMocks) -> None: + workspaces.queries.list_for_account.return_value = ( + WorkspaceSummary("workspace-1", "Test", "team", "normal", NOW, None, True), + ) + body, status = unwrap(controller.TenantListApi.get)(controller.TenantListApi(), CONTEXT) + assert status == HTTPStatus.OK + assert body["workspaces"] == [ + { + "id": "workspace-1", + "name": "Test", + "plan": "team", + "status": "normal", + "created_at": int(NOW.timestamp()), + "last_opened_at": None, + "current": True, + } + ] + workspaces.queries.list_for_account.assert_called_once_with(CONTEXT) -def test_workspace_list_serialization(workspaces: WorkspaceMocks) -> None: - workspaces.queries.list_for_account.return_value = ( - WorkspaceSummary("workspace-1", "Test", "team", "normal", NOW, None, True), - ) - body, status = unwrap(controller.TenantListApi.get)(controller.TenantListApi(), CONTEXT) - assert status == HTTPStatus.OK - assert body["workspaces"] == [ - { - "id": "workspace-1", - "name": "Test", - "plan": "team", - "status": "normal", - "created_at": int(NOW.timestamp()), - "last_opened_at": None, - "current": True, + @staticmethod + def test_admin_pagination(app: Flask, workspaces: WorkspaceMocks) -> None: + workspaces.management.list_all.return_value = WorkspacePage( + (WorkspaceSnapshot("w1", "Test", "archive", NOW),), 5, 2, 2, True + ) + with app.test_request_context("/all-workspaces?page=2&limit=2"): + body, status = unwrap(controller.WorkspaceListApi.get)(controller.WorkspaceListApi()) + assert status == 200 + assert body == { + "data": [{"id": "w1", "name": "Test", "status": "archive", "created_at": int(NOW.timestamp())}], + "total": 5, + "page": 2, + "limit": 2, + "has_more": True, } - ] - workspaces.queries.list_for_account.assert_called_once_with(CONTEXT) + workspaces.management.list_all.assert_called_once_with(page=2, limit=2) + @staticmethod + def test_summary(workspaces: WorkspaceMocks) -> None: + workspaces.management.current_summary.return_value = { + "id": "w1", + "name": "Test", + "role": "owner", + "plan": CloudPlan.TEAM, + "credits": -1, + } + body, status = unwrap(controller.CurrentWorkspaceSummaryApi.get)( + controller.CurrentWorkspaceSummaryApi(), CONTEXT + ) + assert status == 200 + assert body["credits"] == -1 + workspaces.management.current_summary.assert_called_once_with(CONTEXT) -def test_admin_pagination(app: Flask, workspaces: WorkspaceMocks) -> None: - workspaces.management.list_all.return_value = WorkspacePage( - (WorkspaceSnapshot("w1", "Test", "archive", NOW),), 5, 2, 2, True + @staticmethod + @pytest.mark.parametrize( + ("error", "http_error"), + [(WorkspaceArchivedError(), CurrentWorkspaceArchivedError), (WorkspaceNotFoundError(), NotFound)], ) - with app.test_request_context("/all-workspaces?page=2&limit=2"): - body, status = unwrap(controller.WorkspaceListApi.get)(controller.WorkspaceListApi()) - assert status == 200 - assert body == { - "data": [{"id": "w1", "name": "Test", "status": "archive", "created_at": int(NOW.timestamp())}], - "total": 5, - "page": 2, - "limit": 2, - "has_more": True, - } - workspaces.management.list_all.assert_called_once_with(page=2, limit=2) + def test_summary_errors(workspaces: WorkspaceMocks, error: Exception, http_error: type[Exception]) -> None: + workspaces.management.current_summary.side_effect = error + with pytest.raises(http_error): + unwrap(controller.CurrentWorkspaceSummaryApi.get)(controller.CurrentWorkspaceSummaryApi(), CONTEXT) + @staticmethod + def test_switch_and_rename(app: Flask, workspaces: WorkspaceMocks) -> None: + workspaces.management.switch.return_value = {"id": "target", "name": "Switched"} + workspaces.management.rename.return_value = {"id": "workspace-1", "name": "Renamed"} + with app.test_request_context(json={"tenant_id": "target"}): + body = unwrap(controller.SwitchWorkspaceApi.post)(controller.SwitchWorkspaceApi(), CONTEXT) + assert body["new_tenant"]["id"] == "target" + workspaces.management.switch.assert_called_once_with(CONTEXT, "target") + with app.test_request_context(json={"name": "Renamed"}): + body = unwrap(controller.WorkspaceInfoApi.post)(controller.WorkspaceInfoApi(), CONTEXT) + assert body["tenant"]["name"] == "Renamed" + workspaces.management.rename.assert_called_once_with(CONTEXT, "Renamed") -def test_summary(workspaces: WorkspaceMocks) -> None: - workspaces.management.current_summary.return_value = { - "id": "w1", - "name": "Test", - "role": "owner", - "plan": CloudPlan.TEAM, - "credits": -1, - } - body, status = unwrap(controller.CurrentWorkspaceSummaryApi.get)(controller.CurrentWorkspaceSummaryApi(), CONTEXT) - assert status == 200 - assert body["credits"] == -1 - workspaces.management.current_summary.assert_called_once_with(CONTEXT) + @staticmethod + def test_switch_not_linked(app: Flask, workspaces: WorkspaceMocks) -> None: + workspaces.management.switch.side_effect = WorkspaceNotLinkedError() + with app.test_request_context(json={"tenant_id": "target"}), pytest.raises(AccountNotLinkTenantError): + unwrap(controller.SwitchWorkspaceApi.post)(controller.SwitchWorkspaceApi(), CONTEXT) + @staticmethod + def test_config_response_and_partial_update(app: Flask, workspaces: WorkspaceMocks) -> None: + workspaces.management.custom_config.return_value = WorkspaceCustomConfig(True, "https://files/logo") + assert unwrap(controller.CustomConfigWorkspaceApi.get)(controller.CustomConfigWorkspaceApi(), CONTEXT) == { + "remove_webapp_brand": True, + "replace_webapp_logo": "https://files/logo", + } + workspaces.management.update_custom_config.return_value = {"id": "workspace-1"} + with app.test_request_context(json={"remove_webapp_brand": False}): + body = unwrap(controller.CustomConfigWorkspaceApi.post)(controller.CustomConfigWorkspaceApi(), CONTEXT) + assert body["result"] == "success" + workspaces.management.update_custom_config.assert_called_once_with( + CONTEXT, WorkspaceCustomConfigChanges(False, None) + ) -@pytest.mark.parametrize( - ("error", "http_error"), - [(WorkspaceArchivedError(), CurrentWorkspaceArchivedError), (WorkspaceNotFoundError(), NotFound)], -) -def test_summary_errors(workspaces: WorkspaceMocks, error: Exception, http_error: type[Exception]) -> None: - workspaces.management.current_summary.side_effect = error - with pytest.raises(http_error): - unwrap(controller.CurrentWorkspaceSummaryApi.get)(controller.CurrentWorkspaceSummaryApi(), CONTEXT) - - -def test_switch_and_rename(app: Flask, workspaces: WorkspaceMocks) -> None: - workspaces.management.switch.return_value = {"id": "target", "name": "Switched"} - workspaces.management.rename.return_value = {"id": "workspace-1", "name": "Renamed"} - with app.test_request_context(json={"tenant_id": "target"}): - body = unwrap(controller.SwitchWorkspaceApi.post)(controller.SwitchWorkspaceApi(), CONTEXT) - assert body["new_tenant"]["id"] == "target" - workspaces.management.switch.assert_called_once_with(CONTEXT, "target") - with app.test_request_context(json={"name": "Renamed"}): - body = unwrap(controller.WorkspaceInfoApi.post)(controller.WorkspaceInfoApi(), CONTEXT) - assert body["tenant"]["name"] == "Renamed" - workspaces.management.rename.assert_called_once_with(CONTEXT, "Renamed") - - -def test_switch_not_linked(app: Flask, workspaces: WorkspaceMocks) -> None: - workspaces.management.switch.side_effect = WorkspaceNotLinkedError() - with app.test_request_context(json={"tenant_id": "target"}), pytest.raises(AccountNotLinkTenantError): - unwrap(controller.SwitchWorkspaceApi.post)(controller.SwitchWorkspaceApi(), CONTEXT) - - -def test_config_response_and_partial_update(app: Flask, workspaces: WorkspaceMocks) -> None: - workspaces.management.custom_config.return_value = WorkspaceCustomConfig(True, "https://files/logo") - assert unwrap(controller.CustomConfigWorkspaceApi.get)(controller.CustomConfigWorkspaceApi(), CONTEXT) == { - "remove_webapp_brand": True, - "replace_webapp_logo": "https://files/logo", - } - workspaces.management.update_custom_config.return_value = {"id": "workspace-1"} - with app.test_request_context(json={"remove_webapp_brand": False}): - body = unwrap(controller.CustomConfigWorkspaceApi.post)(controller.CustomConfigWorkspaceApi(), CONTEXT) - assert body["result"] == "success" - workspaces.management.update_custom_config.assert_called_once_with( - CONTEXT, WorkspaceCustomConfigChanges(False, None) + @staticmethod + @pytest.mark.parametrize( + ("resource", "method", "payload", "service_method"), + [ + pytest.param( + controller.CustomConfigWorkspaceApi, + controller.CustomConfigWorkspaceApi.get, + {}, + lambda management: management.custom_config, + id="custom_config", + ), + pytest.param( + controller.CustomConfigWorkspaceApi, + controller.CustomConfigWorkspaceApi.post, + {}, + lambda management: management.update_custom_config, + id="update_custom_config", + ), + pytest.param( + controller.WorkspaceInfoApi, + controller.WorkspaceInfoApi.post, + {"name": "New"}, + lambda management: management.rename, + id="rename", + ), + ], ) + def test_missing_workspace( + app: Flask, + workspaces: WorkspaceMocks, + resource: type[Resource], + method: Callable[..., object], + payload: dict[str, object], + service_method: Callable[[Mock], Mock], + ) -> None: + service_method(workspaces.management).side_effect = WorkspaceNotFoundError() + with app.test_request_context(json=payload), pytest.raises(NotFound): + unwrap(method)(resource(), CONTEXT) - -@pytest.mark.parametrize( - ("resource", "method", "payload", "service_method"), - [ - pytest.param( - controller.CustomConfigWorkspaceApi, - controller.CustomConfigWorkspaceApi.get, - {}, - lambda management: management.custom_config, - id="custom_config", - ), - pytest.param( - controller.CustomConfigWorkspaceApi, - controller.CustomConfigWorkspaceApi.post, - {}, - lambda management: management.update_custom_config, - id="update_custom_config", - ), - pytest.param( - controller.WorkspaceInfoApi, - controller.WorkspaceInfoApi.post, - {"name": "New"}, - lambda management: management.rename, - id="rename", - ), - ], -) -def test_missing_workspace( - app: Flask, - workspaces: WorkspaceMocks, - resource: type[Resource], - method: Callable[..., object], - payload: dict[str, object], - service_method: Callable[[Mock], Mock], -) -> None: - service_method(workspaces.management).side_effect = WorkspaceNotFoundError() - with app.test_request_context(json=payload), pytest.raises(NotFound): - unwrap(method)(resource(), CONTEXT) - - -@pytest.mark.parametrize( - ("files", "error"), - [ - ({}, NoFileUploadedError), - ({"file": (BytesIO(b"x"), "a.png"), "extra": (BytesIO(b"x"), "b.png")}, TooManyFilesError), - ({"file": (BytesIO(b"x"), "")}, FilenameNotExistsError), - ], -) -def test_upload_form_errors( - app: Flask, workspaces: WorkspaceMocks, files: dict[str, tuple[BytesIO, str]], error: type[Exception] -) -> None: - with app.test_request_context(method="POST", data=files), pytest.raises(error): - unwrap(controller.WebappLogoWorkspaceApi.post)(controller.WebappLogoWorkspaceApi(), CONTEXT) - workspaces.management.upload_logo.assert_not_called() - - -def test_upload(app: Flask, workspaces: WorkspaceMocks) -> None: - workspaces.management.upload_logo.return_value = "upload-1" - with app.test_request_context(method="POST", data={"file": (BytesIO(b"image"), "logo.PNG")}): - result = unwrap(controller.WebappLogoWorkspaceApi.post)(controller.WebappLogoWorkspaceApi(), CONTEXT) - assert result == ({"id": "upload-1"}, HTTPStatus.CREATED) - workspaces.management.upload_logo.assert_called_once_with( - CONTEXT, filename="logo.PNG", content=b"image", mimetype="image/png" + @staticmethod + @pytest.mark.parametrize( + ("files", "error"), + [ + ({}, NoFileUploadedError), + ({"file": (BytesIO(b"x"), "a.png"), "extra": (BytesIO(b"x"), "b.png")}, TooManyFilesError), + ({"file": (BytesIO(b"x"), "")}, FilenameNotExistsError), + ], ) + def test_upload_form_errors( + app: Flask, workspaces: WorkspaceMocks, files: dict[str, tuple[BytesIO, str]], error: type[Exception] + ) -> None: + with app.test_request_context(method="POST", data=files), pytest.raises(error): + unwrap(controller.WebappLogoWorkspaceApi.post)(controller.WebappLogoWorkspaceApi(), CONTEXT) + workspaces.management.upload_logo.assert_not_called() + @staticmethod + def test_upload(app: Flask, workspaces: WorkspaceMocks) -> None: + workspaces.management.upload_logo.return_value = "upload-1" + with app.test_request_context(method="POST", data={"file": (BytesIO(b"image"), "logo.PNG")}): + result = unwrap(controller.WebappLogoWorkspaceApi.post)(controller.WebappLogoWorkspaceApi(), CONTEXT) + assert result == ({"id": "upload-1"}, HTTPStatus.CREATED) + workspaces.management.upload_logo.assert_called_once_with( + CONTEXT, filename="logo.PNG", content=b"image", mimetype="image/png" + ) -@pytest.mark.parametrize( - ("error", "http_error"), - [ - (file_errors.FileTooLargeError(), FileTooLargeError), - (file_errors.UnsupportedFileTypeError(), UnsupportedFileTypeError), - ], -) -def test_upload_service_errors( - app: Flask, workspaces: WorkspaceMocks, error: Exception, http_error: type[Exception] -) -> None: - workspaces.management.upload_logo.side_effect = error - with app.test_request_context(method="POST", data={"file": (BytesIO(b"x"), "logo.png")}), pytest.raises(http_error): - unwrap(controller.WebappLogoWorkspaceApi.post)(controller.WebappLogoWorkspaceApi(), CONTEXT) + @staticmethod + @pytest.mark.parametrize( + ("error", "http_error"), + [ + (file_errors.FileTooLargeError(), FileTooLargeError), + (file_errors.UnsupportedFileTypeError(), UnsupportedFileTypeError), + ], + ) + def test_upload_service_errors( + app: Flask, workspaces: WorkspaceMocks, error: Exception, http_error: type[Exception] + ) -> None: + workspaces.management.upload_logo.side_effect = error + with ( + app.test_request_context(method="POST", data={"file": (BytesIO(b"x"), "logo.png")}), + pytest.raises(http_error), + ): + unwrap(controller.WebappLogoWorkspaceApi.post)(controller.WebappLogoWorkspaceApi(), CONTEXT) + @staticmethod + def test_permissions(workspaces: WorkspaceMocks) -> None: + workspaces.management.permission.return_value = WorkspacePermission("workspace-1", True, False) + body, status = unwrap(controller.WorkspacePermissionApi.get)(controller.WorkspacePermissionApi(), CONTEXT) + assert status == 200 + assert body == {"workspace_id": "workspace-1", "allow_member_invite": True, "allow_owner_transfer": False} + workspaces.management.permission.assert_called_once_with(CONTEXT) -def test_permissions(workspaces: WorkspaceMocks) -> None: - workspaces.management.permission.return_value = WorkspacePermission("workspace-1", True, False) - body, status = unwrap(controller.WorkspacePermissionApi.get)(controller.WorkspacePermissionApi(), CONTEXT) - assert status == 200 - assert body == {"workspace_id": "workspace-1", "allow_member_invite": True, "allow_owner_transfer": False} - workspaces.management.permission.assert_called_once_with(CONTEXT) - - -def test_removed_current_workspace_routes_stay_unregistered() -> None: - urls = {url for resource in console_ns.resources for url in resource.urls} - assert "/workspaces/current" not in urls - assert "/workspaces/current/tool-provider" not in urls + @staticmethod + def test_removed_current_workspace_routes_stay_unregistered() -> None: + urls = {url for resource in console_ns.resources for url in resource.urls} + assert "/workspaces/current" not in urls + assert "/workspaces/current/tool-provider" not in urls diff --git a/api/tests/unit_tests/controllers/files/test_appdeploy_files.py b/api/tests/unit_tests/controllers/files/test_appdeploy_files.py index 92ba817c388..c436341b797 100644 --- a/api/tests/unit_tests/controllers/files/test_appdeploy_files.py +++ b/api/tests/unit_tests/controllers/files/test_appdeploy_files.py @@ -42,7 +42,7 @@ from services.entities.file_grant_entities import FileGrantContext, FileGrantSco from services.errors.file import FileTooLargeError as FileTooLargeServiceError from services.file_grant_gateways import FILE_CONTENT_AUDIENCE, FileGrantFileGateway from services.file_grant_service import MAX_FILE_GRANT_REFS, FileGrantService -from tests.unit_tests.file_grant_test_utils import issue_file_grant +from tests.helpers.unit.file_grant_test_utils import issue_file_grant CONTROLLER_MODULE = "controllers.files.appdeploy_files" diff --git a/api/tests/unit_tests/controllers/files/test_file_grant_wraps.py b/api/tests/unit_tests/controllers/files/test_wraps.py similarity index 98% rename from api/tests/unit_tests/controllers/files/test_file_grant_wraps.py rename to api/tests/unit_tests/controllers/files/test_wraps.py index 893001f54df..8d462256109 100644 --- a/api/tests/unit_tests/controllers/files/test_file_grant_wraps.py +++ b/api/tests/unit_tests/controllers/files/test_wraps.py @@ -13,7 +13,7 @@ from controllers.files.wraps import FileGrantInvalidError, FileGrantScopeDeniedE from libs.passport import PassportService from services.entities.file_grant_entities import FileGrantClaims, FileGrantScope from services.file_grant_gateways import FILE_GRANT_AUDIENCE -from tests.unit_tests.file_grant_test_utils import issue_file_grant, token_gateway +from tests.helpers.unit.file_grant_test_utils import issue_file_grant, token_gateway SECRET_KEY = "file-grant-test-secret-long-enough-for-hs256" TENANT_ID = "11111111-1111-4111-8111-111111111111" diff --git a/api/tests/unit_tests/controllers/inner_api/test_agent_files.py b/api/tests/unit_tests/controllers/inner_api/agent/test_files.py similarity index 97% rename from api/tests/unit_tests/controllers/inner_api/test_agent_files.py rename to api/tests/unit_tests/controllers/inner_api/agent/test_files.py index d55b6d5f343..6d337a5f40a 100644 --- a/api/tests/unit_tests/controllers/inner_api/test_agent_files.py +++ b/api/tests/unit_tests/controllers/inner_api/agent/test_files.py @@ -15,8 +15,8 @@ from controllers.inner_api.agent.files import ( from core.workflow.file_reference import build_file_reference from models.account import Account, Tenant from services.file_request_service import DownloadFileRequestResult -from tests.unit_tests.config_override import apply_config_overrides -from tests.unit_tests.model_factories import make_tenant +from tests.helpers.unit.config_override import apply_config_overrides +from tests.helpers.unit.model_factories import make_tenant MODULE = "controllers.inner_api.agent.files" diff --git a/api/tests/unit_tests/controllers/inner_api/test_agent_knowledge_investigations.py b/api/tests/unit_tests/controllers/inner_api/agent/test_knowledge.py similarity index 98% rename from api/tests/unit_tests/controllers/inner_api/test_agent_knowledge_investigations.py rename to api/tests/unit_tests/controllers/inner_api/agent/test_knowledge.py index ea5711a4eea..0d850be20d8 100644 --- a/api/tests/unit_tests/controllers/inner_api/test_agent_knowledge_investigations.py +++ b/api/tests/unit_tests/controllers/inner_api/agent/test_knowledge.py @@ -13,7 +13,7 @@ from flask_restx.swagger import Swagger from controllers.inner_api import api as inner_api from controllers.inner_api import bp as inner_api_bp -from tests.unit_tests.config_override import config_overrides_context +from tests.helpers.unit.config_override import config_overrides_context PATH = "/inner/api/agent/knowledge/investigations" HEADERS = {"X-Inner-Api-Key": "inner-key"} diff --git a/api/tests/unit_tests/controllers/inner_api/test_agent_llm.py b/api/tests/unit_tests/controllers/inner_api/agent/test_llm.py similarity index 99% rename from api/tests/unit_tests/controllers/inner_api/test_agent_llm.py rename to api/tests/unit_tests/controllers/inner_api/agent/test_llm.py index f9eb3006bc5..ff017978e7e 100644 --- a/api/tests/unit_tests/controllers/inner_api/test_agent_llm.py +++ b/api/tests/unit_tests/controllers/inner_api/agent/test_llm.py @@ -15,7 +15,7 @@ from graphon.model_runtime.entities.llm_entities import LLMResultChunk, LLMResul from graphon.model_runtime.entities.message_entities import AssistantPromptMessage, UserPromptMessage from services.agent_llm_inner_service import AgentLLMInnerServiceError, PreparedAgentLLMInvocation from services.entities.agent_llm_inner import AgentLLMInvokeRequest -from tests.unit_tests.config_override import config_overrides_context +from tests.helpers.unit.config_override import config_overrides_context def _payload() -> dict[str, object]: diff --git a/api/tests/unit_tests/controllers/inner_api/test_runtime_usage.py b/api/tests/unit_tests/controllers/inner_api/agent/test_runtime_usage.py similarity index 99% rename from api/tests/unit_tests/controllers/inner_api/test_runtime_usage.py rename to api/tests/unit_tests/controllers/inner_api/agent/test_runtime_usage.py index fc084e63bd4..5ca9641076e 100644 --- a/api/tests/unit_tests/controllers/inner_api/test_runtime_usage.py +++ b/api/tests/unit_tests/controllers/inner_api/agent/test_runtime_usage.py @@ -14,7 +14,7 @@ from sqlalchemy.orm import Session, sessionmaker from controllers.inner_api import bp as inner_api_bp from controllers.inner_api import inner_api_ns from models.agent_sandbox_usage import AgentSandboxUsageEvent -from tests.unit_tests.config_override import config_overrides_context +from tests.helpers.unit.config_override import config_overrides_context PROJECT = "431de237-596f-4d59-8a85-20a9846bf243" diff --git a/api/tests/unit_tests/controllers/inner_api/test_agent_tools.py b/api/tests/unit_tests/controllers/inner_api/agent/test_tools.py similarity index 98% rename from api/tests/unit_tests/controllers/inner_api/test_agent_tools.py rename to api/tests/unit_tests/controllers/inner_api/agent/test_tools.py index 186615f3d90..3723ecf8a10 100644 --- a/api/tests/unit_tests/controllers/inner_api/test_agent_tools.py +++ b/api/tests/unit_tests/controllers/inner_api/agent/test_tools.py @@ -9,7 +9,7 @@ from flask import Flask from controllers.inner_api import bp as inner_api_bp from services.entities.agent_tool_inner import AgentToolInvokeResponse from services.errors.agent_tool_inner import AgentToolInnerServiceError -from tests.unit_tests.config_override import config_overrides_context +from tests.helpers.unit.config_override import config_overrides_context def _headers(api_key: str | None = "inner-key") -> dict[str, str]: diff --git a/api/tests/unit_tests/controllers/inner_api/test_agent_tracing.py b/api/tests/unit_tests/controllers/inner_api/agent/test_tracing.py similarity index 93% rename from api/tests/unit_tests/controllers/inner_api/test_agent_tracing.py rename to api/tests/unit_tests/controllers/inner_api/agent/test_tracing.py index df1b4576ec0..496772c2169 100644 --- a/api/tests/unit_tests/controllers/inner_api/test_agent_tracing.py +++ b/api/tests/unit_tests/controllers/inner_api/agent/test_tracing.py @@ -12,8 +12,9 @@ from sqlalchemy.orm import Session, sessionmaker from controllers.inner_api import bp from models.application_trace import ApplicationTraceBatch from services.tracing.service import ApplicationTracingService -from tests.unit_tests.config_override import config_overrides_context -from tests.unit_tests.services.tracing.test_application_trace_collection import batch_for, seed +from tests.helpers.unit.config_override import config_overrides_context +from tests.unit_tests.services.tracing.test_service import _application_trace_collection_batch_for as batch_for +from tests.unit_tests.services.tracing.test_service import _application_trace_collection_seed as seed @pytest.fixture diff --git a/api/tests/unit_tests/controllers/inner_api/app/test_dsl.py b/api/tests/unit_tests/controllers/inner_api/app/test_dsl.py index 0b86a7167cb..eb5f96d5a0d 100644 --- a/api/tests/unit_tests/controllers/inner_api/app/test_dsl.py +++ b/api/tests/unit_tests/controllers/inner_api/app/test_dsl.py @@ -29,7 +29,7 @@ from services.app_dsl_service import AppDslService from services.entities.app_entities import AppExportOptions from services.entities.dsl_entities import Import, ImportStatus from services.errors.app import AppDiscoveryNotFoundError, IsDraftWorkflowError, WorkflowNotFoundError -from tests.unit_tests.config_override import config_overrides_context +from tests.helpers.unit.config_override import config_overrides_context from tests.unit_tests.controllers.conftest import ControllerTestServices diff --git a/api/tests/unit_tests/controllers/inner_api/test_knowledge_retrieval.py b/api/tests/unit_tests/controllers/inner_api/knowledge/test_retrieval.py similarity index 99% rename from api/tests/unit_tests/controllers/inner_api/test_knowledge_retrieval.py rename to api/tests/unit_tests/controllers/inner_api/knowledge/test_retrieval.py index 805385b09b2..947dd3fe0d0 100644 --- a/api/tests/unit_tests/controllers/inner_api/test_knowledge_retrieval.py +++ b/api/tests/unit_tests/controllers/inner_api/knowledge/test_retrieval.py @@ -16,7 +16,7 @@ from services.errors.knowledge_retrieval import ( InnerKnowledgeRetrieveAppNotFoundError, InnerKnowledgeRetrieveDatasetTenantMismatchError, ) -from tests.unit_tests.config_override import config_overrides_context +from tests.helpers.unit.config_override import config_overrides_context @pytest.fixture diff --git a/api/tests/unit_tests/controllers/inner_api/knowledge_fs/test_text_diagnostics.py b/api/tests/unit_tests/controllers/inner_api/knowledge_fs/test_text_diagnostics.py deleted file mode 100644 index 04d807f839a..00000000000 --- a/api/tests/unit_tests/controllers/inner_api/knowledge_fs/test_text_diagnostics.py +++ /dev/null @@ -1,178 +0,0 @@ -"""Exercise provider failures through the real HTTP bridge and TiDB store.""" - -import json -import logging -import time -from collections.abc import Callable -from unittest.mock import patch -from uuid import UUID - -import pymysql -import pytest -from flask import Flask -from flask.testing import FlaskClient - -from controllers.inner_api import inner_api_ns -from controllers.inner_api.knowledge_fs.texts import KnowledgeFSTextApi -from libs.external_api import ExternalApi -from models.dataset import TidbAuthBinding -from models.enums import TidbAuthBindingStatus -from services.knowledge_fs.text_store import TextIndexPendingError - -REQUEST_ID = "d4103b40-0744-458f-a740-d8047d1cbe8d" -HEADERS = {"X-Inner-Api-Key": "trusted-key", "X-Knowledge-FS-Request-Id": REQUEST_ID} -PAYLOAD = { - "operation": "search", - "scope": { - "tenant_id": "018f0d60-7a49-7cc2-9c1b-5b36f18f0001", - "knowledge_space_id": "018f0d60-7a49-7cc2-9c1b-5b36f18f0002", - "index_version": "private-index-version", - }, - "ids": ["018f0d60-7a49-7cc2-9c1b-5b36f18f0004"], - "query": "private-search-content", -} - - -@pytest.fixture -def http_client(config_overrides: Callable[..., None]) -> FlaskClient: - config_overrides(VECTOR_STORE="tidb_on_qdrant", PLUGIN_DAEMON_KEY="enabled", INNER_API_KEY_FOR_PLUGIN="trusted-key") - app = Flask(__name__) - api = ExternalApi(app) - api.models.update(inner_api_ns.models) - api.add_resource(KnowledgeFSTextApi, "/texts") - return app.test_client() - - -@pytest.mark.parametrize("stage", ["connect", "index_check", "replica_readiness", "search"]) -def test_database_failure_retains_code_stage_and_correlation_without_secrets( - http_client: FlaskClient, caplog: pytest.LogCaptureFixture, stage: str -) -> None: - error = pymysql.err.OperationalError( - 1815, "private-password private-search-content SELECT text FROM private-table at private-host" - ) - binding = TidbAuthBinding( - tenant_id=PAYLOAD["scope"]["tenant_id"], - cluster_id="private-cluster", - cluster_name="private-cluster-name", - account="private-account", - password="private-password", - active=True, - status=TidbAuthBindingStatus.ACTIVE, - ) - with ( - patch("services.tidb_binding_service.resolve_tidb_auth_binding", return_value=binding), - patch("services.knowledge_fs.text_store_tidb._sql_endpoint", return_value=("private-host", 4000)), - patch("pymysql.connect") as connect, - caplog.at_level(logging.WARNING), - ): - connection = connect.return_value - cursor = connection.cursor.return_value.__enter__.return_value - cursor.fetchone.return_value = (1,) - - def execute(sql: str, *_args: object) -> None: - if ( - (stage == "index_check" and "information_schema.TABLES" in sql) - or (stage == "replica_readiness" and "information_schema.TIFLASH_REPLICA" in sql) - or (stage == "search" and "fts_match_word" in sql) - ): - raise error - - cursor.execute.side_effect = execute - if stage == "connect": - connect.side_effect = error - if stage == "search": - # Cleanup must not replace the SQL failure with an unrelated secondary error. - connection.close.side_effect = OSError("private-cleanup-error") - response = http_client.post("/texts", json=PAYLOAD, headers=HEADERS) - - assert response.status_code == 503 - assert response.get_json() == { - "code": "service_unavailable", - "message": "KnowledgeFS full-text backend is unavailable", - "status": 503, - } - assert response.headers["X-Knowledge-FS-Request-Id"] == REQUEST_ID - records = [r for r in caplog.records if "knowledge_fs.text_backend_failure" in r.getMessage()] - assert len(records) == 1 - record = records[0] - attributes = record.__dict__["attributes"] - assert attributes["request_id"] == REQUEST_ID - assert attributes["tenant_id"] == PAYLOAD["scope"]["tenant_id"] - assert attributes["knowledge_space_id"] == PAYLOAD["scope"]["knowledge_space_id"] - assert attributes["provider"] == "tidb_on_qdrant" - assert attributes["operation"] == "search" - assert attributes["stage"] == stage - assert attributes["elapsed_ms"] >= attributes["stage_elapsed_ms"] >= 0 - assert attributes["errors"][0]["type"] == "pymysql.err.OperationalError" - assert attributes["errors"][0]["db_error_code"] == 1815 - assert attributes["errors"][0]["db_error_name"] == "INTERNAL_ERROR" - assert attributes["errors"][0]["frames"] - # Both text and structured JSON logs must contain diagnostics, never raw provider data. - assert "1815" in record.getMessage() - assert stage in record.getMessage() - assert record.exc_info is None - serialized = caplog.text + json.dumps(attributes) + response.get_data(as_text=True) - for secret in ( - "private-password", - "private-search-content", - "private-table", - "private-host", - "private-account", - "private-cluster", - "private-index-version", - "private-cleanup-error", - ): - assert secret not in serialized - - -def test_deadline_failure_is_diagnosable_before_connect( - http_client: FlaskClient, caplog: pytest.LogCaptureFixture -) -> None: - response = http_client.post( - "/texts", - json=PAYLOAD, - headers={**HEADERS, "X-Knowledge-FS-Deadline-Ms": str(int(time.time() * 1000) - 1)}, - ) - assert response.status_code == 503 - attributes = next(r.__dict__["attributes"] for r in caplog.records if "text_backend_failure" in r.getMessage()) - assert attributes["stage"] == "deadline" - assert attributes["errors"][0]["type"].endswith("TextSearchTimeoutError") - - -def test_pending_index_keeps_retry_header_and_warning_diagnostic( - http_client: FlaskClient, caplog: pytest.LogCaptureFixture -) -> None: - with patch( - "controllers.inner_api.knowledge_fs.texts.configured_text_client", - side_effect=TextIndexPendingError("private-pending-message"), - ): - response = http_client.post("/texts", json=PAYLOAD, headers=HEADERS) - assert response.status_code == 503 - assert response.headers["Retry-After"] == "5" - record = next(r for r in caplog.records if "text_backend_failure" in r.getMessage()) - assert record.levelno == logging.WARNING - assert record.__dict__["attributes"]["outcome"] == "pending" - assert "private-pending-message" not in caplog.text - - -@pytest.mark.parametrize("request_id", [None, "invalid sensitive value", "a" * 500, REQUEST_ID]) -def test_correlation_header_is_bounded_and_does_not_change_success( - http_client: FlaskClient, request_id: str | None, caplog: pytest.LogCaptureFixture -) -> None: - headers = {"X-Inner-Api-Key": "trusted-key"} - if request_id is not None: - headers["X-Knowledge-FS-Request-Id"] = request_id - with ( - patch("controllers.inner_api.knowledge_fs.texts.configured_text_client"), - patch( - "controllers.inner_api.knowledge_fs.texts.execute_text_request", return_value={"points": [], "matches": []} - ), - ): - response = http_client.post("/texts", json=PAYLOAD, headers=headers) - assert response.status_code == 200 - assert response.get_json() == {"points": [], "matches": []} - correlation_id = response.headers["X-Knowledge-FS-Request-Id"] - assert str(UUID(correlation_id)) == correlation_id - if request_id == REQUEST_ID: - assert correlation_id == request_id - assert not [r for r in caplog.records if "text_backend_failure" in r.getMessage()] diff --git a/api/tests/unit_tests/controllers/inner_api/knowledge_fs/test_texts.py b/api/tests/unit_tests/controllers/inner_api/knowledge_fs/test_texts.py index 28630a56933..0001bd23d7d 100644 --- a/api/tests/unit_tests/controllers/inner_api/knowledge_fs/test_texts.py +++ b/api/tests/unit_tests/controllers/inner_api/knowledge_fs/test_texts.py @@ -1,14 +1,18 @@ import hashlib import inspect +import json +import logging import time from collections.abc import Callable, Generator from contextlib import contextmanager from unittest.mock import call, patch +from uuid import UUID import gevent import pymysql import pytest from flask import Flask +from flask.testing import FlaskClient from werkzeug.exceptions import BadRequest, NotFound, ServiceUnavailable from controllers.inner_api import inner_api_ns @@ -20,7 +24,174 @@ from services.knowledge_fs.text_store import TextIndexPendingError from services.knowledge_fs.vector_store import VectorResult from services.tidb_binding_service import TidbBindingPendingError -PAYLOAD = { +"""Exercise provider failures through the real HTTP bridge and TiDB store.""" + + +REQUEST_ID = "d4103b40-0744-458f-a740-d8047d1cbe8d" +HEADERS = {"X-Inner-Api-Key": "trusted-key", "X-Knowledge-FS-Request-Id": REQUEST_ID} +_TEXT_DIAGNOSTICS_PAYLOAD = { + "operation": "search", + "scope": { + "tenant_id": "018f0d60-7a49-7cc2-9c1b-5b36f18f0001", + "knowledge_space_id": "018f0d60-7a49-7cc2-9c1b-5b36f18f0002", + "index_version": "private-index-version", + }, + "ids": ["018f0d60-7a49-7cc2-9c1b-5b36f18f0004"], + "query": "private-search-content", +} + + +class TestTextDiagnosticsCases: + @staticmethod + @pytest.fixture + def http_client(config_overrides: Callable[..., None]) -> FlaskClient: + config_overrides( + VECTOR_STORE="tidb_on_qdrant", PLUGIN_DAEMON_KEY="enabled", INNER_API_KEY_FOR_PLUGIN="trusted-key" + ) + app = Flask(__name__) + api = ExternalApi(app) + api.models.update(inner_api_ns.models) + api.add_resource(KnowledgeFSTextApi, "/texts") + return app.test_client() + + @staticmethod + @pytest.mark.parametrize("stage", ["connect", "index_check", "replica_readiness", "search"]) + def test_database_failure_retains_code_stage_and_correlation_without_secrets( + http_client: FlaskClient, caplog: pytest.LogCaptureFixture, stage: str + ) -> None: + error = pymysql.err.OperationalError( + 1815, "private-password private-search-content SELECT text FROM private-table at private-host" + ) + binding = TidbAuthBinding( + tenant_id=_TEXT_DIAGNOSTICS_PAYLOAD["scope"]["tenant_id"], + cluster_id="private-cluster", + cluster_name="private-cluster-name", + account="private-account", + password="private-password", + active=True, + status=TidbAuthBindingStatus.ACTIVE, + ) + with ( + patch("services.tidb_binding_service.resolve_tidb_auth_binding", return_value=binding), + patch("services.knowledge_fs.text_store_tidb._sql_endpoint", return_value=("private-host", 4000)), + patch("pymysql.connect") as connect, + caplog.at_level(logging.WARNING), + ): + connection = connect.return_value + cursor = connection.cursor.return_value.__enter__.return_value + cursor.fetchone.return_value = (1,) + + def execute(sql: str, *_args: object) -> None: + if ( + (stage == "index_check" and "information_schema.TABLES" in sql) + or (stage == "replica_readiness" and "information_schema.TIFLASH_REPLICA" in sql) + or (stage == "search" and "fts_match_word" in sql) + ): + raise error + + cursor.execute.side_effect = execute + if stage == "connect": + connect.side_effect = error + if stage == "search": + # Cleanup must not replace the SQL failure with an unrelated secondary error. + connection.close.side_effect = OSError("private-cleanup-error") + response = http_client.post("/texts", json=_TEXT_DIAGNOSTICS_PAYLOAD, headers=HEADERS) + + assert response.status_code == 503 + assert response.get_json() == { + "code": "service_unavailable", + "message": "KnowledgeFS full-text backend is unavailable", + "status": 503, + } + assert response.headers["X-Knowledge-FS-Request-Id"] == REQUEST_ID + records = [r for r in caplog.records if "knowledge_fs.text_backend_failure" in r.getMessage()] + assert len(records) == 1 + record = records[0] + attributes = record.__dict__["attributes"] + assert attributes["request_id"] == REQUEST_ID + assert attributes["tenant_id"] == _TEXT_DIAGNOSTICS_PAYLOAD["scope"]["tenant_id"] + assert attributes["knowledge_space_id"] == _TEXT_DIAGNOSTICS_PAYLOAD["scope"]["knowledge_space_id"] + assert attributes["provider"] == "tidb_on_qdrant" + assert attributes["operation"] == "search" + assert attributes["stage"] == stage + assert attributes["elapsed_ms"] >= attributes["stage_elapsed_ms"] >= 0 + assert attributes["errors"][0]["type"] == "pymysql.err.OperationalError" + assert attributes["errors"][0]["db_error_code"] == 1815 + assert attributes["errors"][0]["db_error_name"] == "INTERNAL_ERROR" + assert attributes["errors"][0]["frames"] + # Both text and structured JSON logs must contain diagnostics, never raw provider data. + assert "1815" in record.getMessage() + assert stage in record.getMessage() + assert record.exc_info is None + serialized = caplog.text + json.dumps(attributes) + response.get_data(as_text=True) + for secret in ( + "private-password", + "private-search-content", + "private-table", + "private-host", + "private-account", + "private-cluster", + "private-index-version", + "private-cleanup-error", + ): + assert secret not in serialized + + @staticmethod + def test_deadline_failure_is_diagnosable_before_connect( + http_client: FlaskClient, caplog: pytest.LogCaptureFixture + ) -> None: + response = http_client.post( + "/texts", + json=_TEXT_DIAGNOSTICS_PAYLOAD, + headers={**HEADERS, "X-Knowledge-FS-Deadline-Ms": str(int(time.time() * 1000) - 1)}, + ) + assert response.status_code == 503 + attributes = next(r.__dict__["attributes"] for r in caplog.records if "text_backend_failure" in r.getMessage()) + assert attributes["stage"] == "deadline" + assert attributes["errors"][0]["type"].endswith("TextSearchTimeoutError") + + @staticmethod + def test_pending_index_keeps_retry_header_and_warning_diagnostic( + http_client: FlaskClient, caplog: pytest.LogCaptureFixture + ) -> None: + with patch( + "controllers.inner_api.knowledge_fs.texts.configured_text_client", + side_effect=TextIndexPendingError("private-pending-message"), + ): + response = http_client.post("/texts", json=_TEXT_DIAGNOSTICS_PAYLOAD, headers=HEADERS) + assert response.status_code == 503 + assert response.headers["Retry-After"] == "5" + record = next(r for r in caplog.records if "text_backend_failure" in r.getMessage()) + assert record.levelno == logging.WARNING + assert record.__dict__["attributes"]["outcome"] == "pending" + assert "private-pending-message" not in caplog.text + + @staticmethod + @pytest.mark.parametrize("request_id", [None, "invalid sensitive value", "a" * 500, REQUEST_ID]) + def test_correlation_header_is_bounded_and_does_not_change_success( + http_client: FlaskClient, request_id: str | None, caplog: pytest.LogCaptureFixture + ) -> None: + headers = {"X-Inner-Api-Key": "trusted-key"} + if request_id is not None: + headers["X-Knowledge-FS-Request-Id"] = request_id + with ( + patch("controllers.inner_api.knowledge_fs.texts.configured_text_client"), + patch( + "controllers.inner_api.knowledge_fs.texts.execute_text_request", + return_value={"points": [], "matches": []}, + ), + ): + response = http_client.post("/texts", json=_TEXT_DIAGNOSTICS_PAYLOAD, headers=headers) + assert response.status_code == 200 + assert response.get_json() == {"points": [], "matches": []} + correlation_id = response.headers["X-Knowledge-FS-Request-Id"] + assert str(UUID(correlation_id)) == correlation_id + if request_id == REQUEST_ID: + assert correlation_id == request_id + assert not [r for r in caplog.records if "text_backend_failure" in r.getMessage()] + + +_TEXTS_PAYLOAD = { "operation": "get", "scope": { "tenant_id": "018f0d60-7a49-7cc2-9c1b-5b36f18f0001", @@ -31,256 +202,267 @@ PAYLOAD = { } -@pytest.mark.parametrize("deadline", ["invalid", "NaN", "1.5", "-1", "9" * 30]) -def test_search_rejects_invalid_deadline_before_provider_io(app: Flask, deadline: str) -> None: - handler = KnowledgeFSTextApi() - with ( - app.test_request_context( - json={**PAYLOAD, "operation": "search", "query": "refund"}, - headers={"X-Knowledge-FS-Deadline-Ms": deadline}, - ), - patch("controllers.inner_api.knowledge_fs.texts.configured_text_client") as client, - ): - with pytest.raises(BadRequest): - inspect.unwrap(handler.post)(handler) - client.assert_not_called() +class TestTextsCases: + @staticmethod + @pytest.mark.parametrize("deadline", ["invalid", "NaN", "1.5", "-1", "9" * 30]) + def test_search_rejects_invalid_deadline_before_provider_io(app: Flask, deadline: str) -> None: + handler = KnowledgeFSTextApi() + with ( + app.test_request_context( + json={**_TEXTS_PAYLOAD, "operation": "search", "query": "refund"}, + headers={"X-Knowledge-FS-Deadline-Ms": deadline}, + ), + patch("controllers.inner_api.knowledge_fs.texts.configured_text_client") as client, + ): + with pytest.raises(BadRequest): + inspect.unwrap(handler.post)(handler) + client.assert_not_called() + @staticmethod + def test_search_rejects_expired_deadline_before_provider_io(app: Flask) -> None: + handler = KnowledgeFSTextApi() + with ( + app.test_request_context( + json={**_TEXTS_PAYLOAD, "operation": "search", "query": "refund"}, + headers={"X-Knowledge-FS-Deadline-Ms": str(int(time.time() * 1000) - 1)}, + ), + patch("controllers.inner_api.knowledge_fs.texts.configured_text_client") as client, + ): + with pytest.raises(ServiceUnavailable): + inspect.unwrap(handler.post)(handler) + client.assert_not_called() -def test_search_rejects_expired_deadline_before_provider_io(app: Flask) -> None: - handler = KnowledgeFSTextApi() - with ( - app.test_request_context( - json={**PAYLOAD, "operation": "search", "query": "refund"}, - headers={"X-Knowledge-FS-Deadline-Ms": str(int(time.time() * 1000) - 1)}, - ), - patch("controllers.inner_api.knowledge_fs.texts.configured_text_client") as client, - ): - with pytest.raises(ServiceUnavailable): - inspect.unwrap(handler.post)(handler) - client.assert_not_called() + @staticmethod + def test_legacy_search_without_deadline_is_still_bounded(app: Flask) -> None: + handler = KnowledgeFSTextApi() + with ( + app.test_request_context(json={**_TEXTS_PAYLOAD, "operation": "search", "query": "refund"}), + patch("controllers.inner_api.knowledge_fs.texts.configured_text_client") as client, + patch( + "controllers.inner_api.knowledge_fs.texts.execute_text_request", + return_value={"points": [], "matches": []}, + ), + ): + assert inspect.unwrap(handler.post)(handler) == {"points": [], "matches": []} + assert 49 < client.call_args.kwargs["deadline"].remaining() <= 50 + @staticmethod + @pytest.mark.parametrize("stage", ["connect", "query"]) + def test_search_deadline_interrupts_stalled_backend_and_runs_cleanup(app: Flask, stage: str) -> None: + cleaned_up = False -def test_legacy_search_without_deadline_is_still_bounded(app: Flask) -> None: - handler = KnowledgeFSTextApi() - with ( - app.test_request_context(json={**PAYLOAD, "operation": "search", "query": "refund"}), - patch("controllers.inner_api.knowledge_fs.texts.configured_text_client") as client, - patch( - "controllers.inner_api.knowledge_fs.texts.execute_text_request", - return_value={"points": [], "matches": []}, - ), - ): - assert inspect.unwrap(handler.post)(handler) == {"points": [], "matches": []} - assert 49 < client.call_args.kwargs["deadline"].remaining() <= 50 + @contextmanager + def stalled_client(*_args: object, **_kwargs: object) -> Generator[object]: + nonlocal cleaned_up + try: + if stage == "connect": + gevent.sleep(1) + yield object() + finally: + cleaned_up = True + def execute(*_args: object) -> VectorResult: + gevent.sleep(1) + return {"points": [], "matches": []} -@pytest.mark.parametrize("stage", ["connect", "query"]) -def test_search_deadline_interrupts_stalled_backend_and_runs_cleanup(app: Flask, stage: str) -> None: - cleaned_up = False + handler = KnowledgeFSTextApi() + started = time.monotonic() + with ( + app.test_request_context( + json={**_TEXTS_PAYLOAD, "operation": "search", "query": "refund"}, + headers={"X-Knowledge-FS-Deadline-Ms": str(int(time.time() * 1000) + 50)}, + ), + patch("controllers.inner_api.knowledge_fs.texts.configured_text_client", stalled_client), + patch("controllers.inner_api.knowledge_fs.texts.execute_text_request", execute), + ): + with pytest.raises(ServiceUnavailable): + inspect.unwrap(handler.post)(handler) + assert cleaned_up + assert time.monotonic() - started < 0.5 - @contextmanager - def stalled_client(*_args: object, **_kwargs: object) -> Generator[object]: - nonlocal cleaned_up - try: - if stage == "connect": - gevent.sleep(1) - yield object() - finally: - cleaned_up = True + @staticmethod + def test_text_controller_rejects_unauthenticated_requests(app: Flask) -> None: + with app.test_request_context(json=_TEXTS_PAYLOAD): + with pytest.raises(NotFound): + KnowledgeFSTextApi().post() - def execute(*_args: object) -> VectorResult: - gevent.sleep(1) - return {"points": [], "matches": []} + @staticmethod + def test_text_controller_validates_before_opening_backend(app: Flask) -> None: + handler = KnowledgeFSTextApi() + with ( + app.test_request_context(json={**_TEXTS_PAYLOAD, "endpoint": "http://untrusted"}), + patch("controllers.inner_api.knowledge_fs.texts.configured_text_client") as client, + ): + with pytest.raises(BadRequest): + inspect.unwrap(handler.post)(handler) + client.assert_not_called() - handler = KnowledgeFSTextApi() - started = time.monotonic() - with ( - app.test_request_context( - json={**PAYLOAD, "operation": "search", "query": "refund"}, - headers={"X-Knowledge-FS-Deadline-Ms": str(int(time.time() * 1000) + 50)}, - ), - patch("controllers.inner_api.knowledge_fs.texts.configured_text_client", stalled_client), - patch("controllers.inner_api.knowledge_fs.texts.execute_text_request", execute), - ): - with pytest.raises(ServiceUnavailable): - inspect.unwrap(handler.post)(handler) - assert cleaned_up - assert time.monotonic() - started < 0.5 + @staticmethod + def test_text_controller_hides_provider_credentials(app: Flask) -> None: + handler = KnowledgeFSTextApi() + with ( + app.test_request_context(json=_TEXTS_PAYLOAD), + patch( + "controllers.inner_api.knowledge_fs.texts.configured_text_client", + side_effect=RuntimeError("private-api-key"), + ), + ): + with pytest.raises(ServiceUnavailable) as error: + inspect.unwrap(handler.post)(handler) + assert "private-api-key" not in str(error.value) + assert error.value.__suppress_context__ + @staticmethod + def test_text_controller_serializes_successful_response(app: Flask) -> None: + handler = KnowledgeFSTextApi() + result: dict[str, list[dict[str, object]]] = { + "points": [], + "matches": [{"id": _TEXTS_PAYLOAD["ids"][0], "score": 0.75}], + } + with ( + app.test_request_context(json=_TEXTS_PAYLOAD), + patch("controllers.inner_api.knowledge_fs.texts.configured_text_client"), + patch("controllers.inner_api.knowledge_fs.texts.execute_text_request", return_value=result), + ): + assert inspect.unwrap(handler.post)(handler) == result -def test_text_controller_rejects_unauthenticated_requests(app: Flask) -> None: - with app.test_request_context(json=PAYLOAD): - with pytest.raises(NotFound): - KnowledgeFSTextApi().post() + @staticmethod + def test_text_controller_accepts_only_configured_inner_key( + app: Flask, config_overrides: Callable[..., None] + ) -> None: + config_overrides(PLUGIN_DAEMON_KEY="enabled", INNER_API_KEY_FOR_PLUGIN="trusted-key") + with ( + patch("controllers.inner_api.knowledge_fs.texts.configured_text_client") as client, + patch( + "controllers.inner_api.knowledge_fs.texts.execute_text_request", + return_value={"points": [], "matches": []}, + ), + ): + for key in [None, "wrong-key"]: + with app.test_request_context(json=_TEXTS_PAYLOAD, headers={"X-Inner-Api-Key": key} if key else {}): + with pytest.raises(NotFound): + KnowledgeFSTextApi().post() + client.assert_not_called() + with app.test_request_context(json=_TEXTS_PAYLOAD, headers={"X-Inner-Api-Key": "trusted-key"}): + assert KnowledgeFSTextApi().post() == {"points": [], "matches": []} + client.assert_called_once_with(_TEXTS_PAYLOAD["scope"]["tenant_id"], allow_create=False) + @staticmethod + def test_text_controller_reports_pending_index_without_leaking_credentials(app: Flask) -> None: + handler = KnowledgeFSTextApi() + with ( + app.test_request_context(json=_TEXTS_PAYLOAD), + patch( + "controllers.inner_api.knowledge_fs.texts.configured_text_client", + side_effect=TextIndexPendingError("secret"), + ), + ): + with pytest.raises(ServiceUnavailable) as error: + inspect.unwrap(handler.post)(handler) + assert error.value.retry_after == 5 + assert "secret" not in str(error.value) -def test_text_controller_validates_before_opening_backend(app: Flask) -> None: - handler = KnowledgeFSTextApi() - with ( - app.test_request_context(json={**PAYLOAD, "endpoint": "http://untrusted"}), - patch("controllers.inner_api.knowledge_fs.texts.configured_text_client") as client, - ): - with pytest.raises(BadRequest): - inspect.unwrap(handler.post)(handler) - client.assert_not_called() - - -def test_text_controller_hides_provider_credentials(app: Flask) -> None: - handler = KnowledgeFSTextApi() - with ( - app.test_request_context(json=PAYLOAD), - patch( + @staticmethod + @pytest.mark.parametrize("pending_error", [TidbBindingPendingError, TextIndexPendingError]) + def test_http_pending_response_preserves_retry_after_header( + config_overrides: Callable[..., None], pending_error: type[Exception] + ) -> None: + config_overrides(PLUGIN_DAEMON_KEY="enabled", INNER_API_KEY_FOR_PLUGIN="trusted-key") + http_app = Flask(__name__) + api = ExternalApi(http_app) + api.models.update(inner_api_ns.models) + api.add_resource(KnowledgeFSTextApi, "/texts") + with patch( "controllers.inner_api.knowledge_fs.texts.configured_text_client", - side_effect=RuntimeError("private-api-key"), - ), - ): - with pytest.raises(ServiceUnavailable) as error: - inspect.unwrap(handler.post)(handler) - assert "private-api-key" not in str(error.value) - assert error.value.__suppress_context__ + side_effect=pending_error("private endpoint"), + ): + response = http_app.test_client().post( + "/texts", json=_TEXTS_PAYLOAD, headers={"X-Inner-Api-Key": "trusted-key"} + ) - -def test_text_controller_serializes_successful_response(app: Flask) -> None: - handler = KnowledgeFSTextApi() - result: dict[str, list[dict[str, object]]] = {"points": [], "matches": [{"id": PAYLOAD["ids"][0], "score": 0.75}]} - with ( - app.test_request_context(json=PAYLOAD), - patch("controllers.inner_api.knowledge_fs.texts.configured_text_client"), - patch("controllers.inner_api.knowledge_fs.texts.execute_text_request", return_value=result), - ): - assert inspect.unwrap(handler.post)(handler) == result - - -def test_text_controller_accepts_only_configured_inner_key(app: Flask, config_overrides: Callable[..., None]) -> None: - config_overrides(PLUGIN_DAEMON_KEY="enabled", INNER_API_KEY_FOR_PLUGIN="trusted-key") - with ( - patch("controllers.inner_api.knowledge_fs.texts.configured_text_client") as client, - patch( - "controllers.inner_api.knowledge_fs.texts.execute_text_request", - return_value={"points": [], "matches": []}, - ), - ): - for key in [None, "wrong-key"]: - with app.test_request_context(json=PAYLOAD, headers={"X-Inner-Api-Key": key} if key else {}): - with pytest.raises(NotFound): - KnowledgeFSTextApi().post() - client.assert_not_called() - with app.test_request_context(json=PAYLOAD, headers={"X-Inner-Api-Key": "trusted-key"}): - assert KnowledgeFSTextApi().post() == {"points": [], "matches": []} - client.assert_called_once_with(PAYLOAD["scope"]["tenant_id"], allow_create=False) - - -def test_text_controller_reports_pending_index_without_leaking_credentials(app: Flask) -> None: - handler = KnowledgeFSTextApi() - with ( - app.test_request_context(json=PAYLOAD), - patch( - "controllers.inner_api.knowledge_fs.texts.configured_text_client", - side_effect=TextIndexPendingError("secret"), - ), - ): - with pytest.raises(ServiceUnavailable) as error: - inspect.unwrap(handler.post)(handler) - assert error.value.retry_after == 5 - assert "secret" not in str(error.value) - - -@pytest.mark.parametrize("pending_error", [TidbBindingPendingError, TextIndexPendingError]) -def test_http_pending_response_preserves_retry_after_header( - config_overrides: Callable[..., None], pending_error: type[Exception] -) -> None: - config_overrides(PLUGIN_DAEMON_KEY="enabled", INNER_API_KEY_FOR_PLUGIN="trusted-key") - http_app = Flask(__name__) - api = ExternalApi(http_app) - api.models.update(inner_api_ns.models) - api.add_resource(KnowledgeFSTextApi, "/texts") - with patch( - "controllers.inner_api.knowledge_fs.texts.configured_text_client", - side_effect=pending_error("private endpoint"), - ): - response = http_app.test_client().post("/texts", json=PAYLOAD, headers={"X-Inner-Api-Key": "trusted-key"}) - - assert response.status_code == 503 - assert response.headers["Retry-After"] == "5" - assert response.mimetype == "application/json" - assert response.get_json() == { - "code": "service_unavailable", - "message": "KnowledgeFS full-text index is being initialized", - "status": 503, - } - - -def test_http_cleanup_after_interrupted_first_write_does_not_initialize_database( - config_overrides: Callable[..., None], -) -> None: - config_overrides(VECTOR_STORE="tidb_on_qdrant", PLUGIN_DAEMON_KEY="enabled", INNER_API_KEY_FOR_PLUGIN="trusted-key") - http_app = Flask(__name__) - api = ExternalApi(http_app) - api.models.update(inner_api_ns.models) - api.add_resource(KnowledgeFSTextApi, "/texts") - binding = TidbAuthBinding( - tenant_id=PAYLOAD["scope"]["tenant_id"], - cluster_id="created-cluster", - cluster_name="tenant-cluster", - account="tenant.root", - password="test-password", - active=True, - status=TidbAuthBindingStatus.ACTIVE, - ) - text = "refund policy" - upsert = { - "operation": "upsert", - "scope": PAYLOAD["scope"], - "points": [ - { - "id": PAYLOAD["ids"][0], - "generation_id": "018f0d60-7a49-7cc2-9c1b-5b36f18f0003", - "content_hash": hashlib.sha256(text.encode()).hexdigest(), - "text": text, - } - ], - } - with ( - patch( - "services.tidb_binding_service.resolve_tidb_auth_binding", - side_effect=[TidbBindingPendingError(), binding, binding, binding, binding], - ) as resolver, - patch("services.knowledge_fs.text_store_tidb._sql_endpoint", return_value=("tenant.test", 4000)), - patch("pymysql.connect") as connect, - ): - connection = connect.return_value - - def open_connection(**kwargs: object) -> object: - if kwargs.get("database") == "ai": - raise pymysql.err.OperationalError(1049, "Unknown database 'ai'") - return connection - - def read_only_sql(sql: str, *_args: object) -> None: - assert sql.lstrip().upper().startswith("SELECT "), sql - - connect.side_effect = open_connection - cursor = connection.cursor.return_value.__enter__.return_value - cursor.execute.side_effect = read_only_sql - cursor.fetchone.return_value = None - client = http_app.test_client() - headers = {"X-Inner-Api-Key": "trusted-key"} - - # Provisioning was committed, but the writer stopped before initializing SQL storage. - response = client.post("/texts", json=upsert, headers=headers) assert response.status_code == 503 assert response.headers["Retry-After"] == "5" - connect.assert_not_called() + assert response.mimetype == "application/json" + assert response.get_json() == { + "code": "service_unavailable", + "message": "KnowledgeFS full-text index is being initialized", + "status": 503, + } - for operation in ("delete", "get", "delete", "get"): - response = client.post("/texts", json={**PAYLOAD, "operation": operation}, headers=headers) - assert response.status_code == 200 - assert response.get_json() == {"points": [], "matches": []} - assert "Retry-After" not in response.headers + @staticmethod + def test_http_cleanup_after_interrupted_first_write_does_not_initialize_database( + config_overrides: Callable[..., None], + ) -> None: + config_overrides( + VECTOR_STORE="tidb_on_qdrant", PLUGIN_DAEMON_KEY="enabled", INNER_API_KEY_FOR_PLUGIN="trusted-key" + ) + http_app = Flask(__name__) + api = ExternalApi(http_app) + api.models.update(inner_api_ns.models) + api.add_resource(KnowledgeFSTextApi, "/texts") + binding = TidbAuthBinding( + tenant_id=_TEXTS_PAYLOAD["scope"]["tenant_id"], + cluster_id="created-cluster", + cluster_name="tenant-cluster", + account="tenant.root", + password="test-password", + active=True, + status=TidbAuthBindingStatus.ACTIVE, + ) + text = "refund policy" + upsert = { + "operation": "upsert", + "scope": _TEXTS_PAYLOAD["scope"], + "points": [ + { + "id": _TEXTS_PAYLOAD["ids"][0], + "generation_id": "018f0d60-7a49-7cc2-9c1b-5b36f18f0003", + "content_hash": hashlib.sha256(text.encode()).hexdigest(), + "text": text, + } + ], + } + with ( + patch( + "services.tidb_binding_service.resolve_tidb_auth_binding", + side_effect=[TidbBindingPendingError(), binding, binding, binding, binding], + ) as resolver, + patch("services.knowledge_fs.text_store_tidb._sql_endpoint", return_value=("tenant.test", 4000)), + patch("pymysql.connect") as connect, + ): + connection = connect.return_value - assert resolver.call_args_list == [ - call(PAYLOAD["scope"]["tenant_id"], allow_create=True), - *[call(PAYLOAD["scope"]["tenant_id"], allow_create=False)] * 4, - ] - assert connect.call_count == 4 - connection.select_db.assert_not_called() - assert connection.close.call_count == 4 + def open_connection(**kwargs: object) -> object: + if kwargs.get("database") == "ai": + raise pymysql.err.OperationalError(1049, "Unknown database 'ai'") + return connection + + def read_only_sql(sql: str, *_args: object) -> None: + assert sql.lstrip().upper().startswith("SELECT "), sql + + connect.side_effect = open_connection + cursor = connection.cursor.return_value.__enter__.return_value + cursor.execute.side_effect = read_only_sql + cursor.fetchone.return_value = None + client = http_app.test_client() + headers = {"X-Inner-Api-Key": "trusted-key"} + + # Provisioning was committed, but the writer stopped before initializing SQL storage. + response = client.post("/texts", json=upsert, headers=headers) + assert response.status_code == 503 + assert response.headers["Retry-After"] == "5" + connect.assert_not_called() + + for operation in ("delete", "get", "delete", "get"): + response = client.post("/texts", json={**_TEXTS_PAYLOAD, "operation": operation}, headers=headers) + assert response.status_code == 200 + assert response.get_json() == {"points": [], "matches": []} + assert "Retry-After" not in response.headers + + assert resolver.call_args_list == [ + call(_TEXTS_PAYLOAD["scope"]["tenant_id"], allow_create=True), + *[call(_TEXTS_PAYLOAD["scope"]["tenant_id"], allow_create=False)] * 4, + ] + assert connect.call_count == 4 + connection.select_db.assert_not_called() + assert connection.close.call_count == 4 diff --git a/api/tests/unit_tests/controllers/inner_api/plugin/test_plugin.py b/api/tests/unit_tests/controllers/inner_api/plugin/test_plugin.py index 0abe7b6c222..337a023563d 100644 --- a/api/tests/unit_tests/controllers/inner_api/plugin/test_plugin.py +++ b/api/tests/unit_tests/controllers/inner_api/plugin/test_plugin.py @@ -41,8 +41,8 @@ from core.plugin.entities.request import RequestInvokeDatasource, RequestInvokeT from core.workflow.file_reference import build_file_reference from graphon.model_runtime.entities.model_entities import ModelType from models import Account, Tenant -from tests.unit_tests.config_override import apply_config_overrides -from tests.unit_tests.model_factories import make_tenant +from tests.helpers.unit.config_override import apply_config_overrides +from tests.helpers.unit.model_factories import make_tenant def _tenant() -> Tenant: diff --git a/api/tests/unit_tests/controllers/inner_api/plugin/test_plugin_wraps.py b/api/tests/unit_tests/controllers/inner_api/plugin/test_wraps.py similarity index 100% rename from api/tests/unit_tests/controllers/inner_api/plugin/test_plugin_wraps.py rename to api/tests/unit_tests/controllers/inner_api/plugin/test_wraps.py diff --git a/api/tests/unit_tests/controllers/inner_api/test_datasource_credentials.py b/api/tests/unit_tests/controllers/inner_api/test_datasource_credentials.py index 4c8d83f135f..ca2eb9b0b55 100644 --- a/api/tests/unit_tests/controllers/inner_api/test_datasource_credentials.py +++ b/api/tests/unit_tests/controllers/inner_api/test_datasource_credentials.py @@ -8,7 +8,7 @@ from controllers.inner_api.datasource_credentials import ( InnerDatasourceDefaultCredentialPayload, ) from services.credentials.query import DatasourceCredentialListItem -from tests.unit_tests.credential_query_fakes import CredentialServices, DatasourceCredentialQuery +from tests.helpers.unit.credential_query_fakes import CredentialServices, DatasourceCredentialQuery def test_resolves_default_datasource_credential_metadata(monkeypatch: pytest.MonkeyPatch) -> None: diff --git a/api/tests/unit_tests/controllers/inner_api/test_runtime_credentials.py b/api/tests/unit_tests/controllers/inner_api/test_runtime_credentials.py index 01e725f9eac..5772098d265 100644 --- a/api/tests/unit_tests/controllers/inner_api/test_runtime_credentials.py +++ b/api/tests/unit_tests/controllers/inner_api/test_runtime_credentials.py @@ -16,7 +16,7 @@ from controllers.inner_api.runtime_credentials import ( ) from models.provider import ProviderCredential from models.tools import BuiltinToolProvider -from tests.unit_tests.config_override import config_overrides_context +from tests.helpers.unit.config_override import config_overrides_context def test_runtime_credentials_payload_accepts_items(): diff --git a/api/tests/unit_tests/controllers/inner_api/test_schema.py b/api/tests/unit_tests/controllers/inner_api/test_schema.py deleted file mode 100644 index 90b1ccd3fc3..00000000000 --- a/api/tests/unit_tests/controllers/inner_api/test_schema.py +++ /dev/null @@ -1,296 +0,0 @@ -"""Coverage regressions for the enterprise inner API OpenAPI document.""" - -import importlib.util -import json -import re -import sys -from collections.abc import Iterator -from pathlib import Path -from types import ModuleType -from typing import cast - -import pytest -from jsonschema import Draft202012Validator, ValidationError - -type Json = dict[str, object] - -HTTP_METHODS = {"delete", "get", "head", "options", "patch", "post", "put", "trace"} -WRITE_METHODS = {"patch", "post", "put"} -ENTERPRISE_PREFIX = "/enterprise/" - -EXPECTED_ENTERPRISE_OPERATION_IDS = { - "create_enterprise_workspace", - "create_enterprise_workspace_ownerless", - "delete_enterprise_workspace_resources", - "enterprise_app_detail", - "enterprise_app_dsl_export", - "enterprise_app_dsl_import", - "enterprise_app_list", - "enterprise_mint_file_grant", - "enterprise_runtime_credentials_resolve", - "join_enterprise_workspace_member", - "send_enterprise_mail", -} - - -def _obj(value: object) -> Json: - return cast(Json, value) - - -def _at(value: object, *keys: str) -> Json: - """Walk nested JSON objects by key; each step is cast because the spec is untyped JSON.""" - node = _obj(value) - for key in keys: - node = _obj(node[key]) - return node - - -def _load_generate_swagger_specs_module() -> ModuleType: - api_dir = Path(__file__).resolve().parents[4] - script_path = api_dir / "dev" / "generate_swagger_specs.py" - - spec = importlib.util.spec_from_file_location("generate_swagger_specs", script_path) - assert spec - assert spec.loader - - module = importlib.util.module_from_spec(spec) - sys.modules[spec.name] = module - spec.loader.exec_module(module) # type: ignore[attr-defined] - return module - - -@pytest.fixture(scope="module") -def inner_spec(tmp_path_factory: pytest.TempPathFactory) -> Json: - module = _load_generate_swagger_specs_module() - written_paths = module.generate_specs(tmp_path_factory.mktemp("inner_api_spec")) - inner_path = next(path for path in written_paths if path.name == "inner-openapi.json") - return _obj(json.loads(inner_path.read_text(encoding="utf-8"))) - - -def _enterprise_operations(spec: Json) -> Iterator[tuple[str, str, Json, Json]]: - for path, path_item in _obj(spec["paths"]).items(): - if not path.startswith(ENTERPRISE_PREFIX): - continue - item = _obj(path_item) - for method, operation in item.items(): - if method in HTTP_METHODS: - yield path, method, item, _obj(operation) - - -def _operation(spec: Json, path: str, method: str) -> Json: - return _at(spec, "paths", path, method) - - -def _schema(spec: Json, name: str) -> Json: - return _at(spec, "components", "schemas", name) - - -def _refs(value: object) -> Iterator[str]: - if isinstance(value, dict): - mapping = _obj(value) - ref = mapping.get("$ref") - if isinstance(ref, str): - yield ref - for child in mapping.values(): - yield from _refs(child) - elif isinstance(value, list): - for child in cast(list[object], value): - yield from _refs(child) - - -def _validator(spec: Json, name: str) -> Draft202012Validator: - return Draft202012Validator({"$ref": f"#/components/schemas/{name}", "components": spec["components"]}) - - -def test_enterprise_operations_declare_typed_request_bodies(inner_spec: Json) -> None: - operations = list(_enterprise_operations(inner_spec)) - assert operations - - for path, method, _, operation in operations: - if method in WRITE_METHODS: - media = _at(operation, "requestBody", "content", "application/json") - assert "schema" in media, f"{method} {path} request body has no schema" - if method == "get": - assert "requestBody" not in operation, f"GET {path} must not declare a request body" - - -def test_enterprise_success_responses_are_typed(inner_spec: Json) -> None: - for path, method, _, operation in _enterprise_operations(inner_spec): - for code, response in _obj(operation["responses"]).items(): - if not code.startswith("2"): - continue - content = _obj(_obj(response).get("content")) - assert content, f"{method} {path} {code} has no content" - for media_type, media in content.items(): - assert "schema" in _obj(media), f"{method} {path} {code} {media_type} has no schema" - - -def test_enterprise_error_responses_are_typed(inner_spec: Json) -> None: - for path, method, _, operation in _enterprise_operations(inner_spec): - for code, response in _obj(operation["responses"]).items(): - if code.startswith("2"): - continue - content = _obj(_obj(response).get("content")) - assert content, f"{method} {path} {code} is not typed" - json_media = _obj(content.get("application/json")) - assert "schema" in json_media, f"{method} {path} {code} has no JSON schema" - - -def test_enterprise_path_placeholders_have_required_path_parameters(inner_spec: Json) -> None: - for path, method, path_item, operation in _enterprise_operations(inner_spec): - parameters = [ - *cast(list[Json], path_item.get("parameters", [])), - *cast(list[Json], operation.get("parameters", [])), - ] - declared = { - parameter["name"] - for parameter in parameters - if parameter.get("in") == "path" and parameter.get("required") is True - } - for placeholder in re.findall(r"{(\w+)}", path): - assert placeholder in declared, f"{method} {path} missing required path parameter {placeholder}" - - -def test_enterprise_schema_references_resolve(inner_spec: Json) -> None: - schemas = _at(inner_spec, "components", "schemas") - for _, _, _, operation in _enterprise_operations(inner_spec): - for ref in _refs(operation): - assert ref.removeprefix("#/components/schemas/") in schemas, ref - - -def test_enterprise_operation_ids_are_unchanged(inner_spec: Json) -> None: - operation_ids = [operation["operationId"] for _, _, _, operation in _enterprise_operations(inner_spec)] - - assert len(operation_ids) == len(set(operation_ids)) - assert set(operation_ids) == EXPECTED_ENTERPRISE_OPERATION_IDS - - -def test_enterprise_mail_documents_simple_message_response(inner_spec: Json) -> None: - media = _at(_operation(inner_spec, "/enterprise/mail", "post"), "responses", "200", "content", "application/json") - - assert media["schema"] == {"$ref": "#/components/schemas/SimpleMessageResponse"} - - -def test_ownerless_workspace_custom_config_is_object(inner_spec: Json) -> None: - schema = _schema(inner_spec, "WorkspaceOwnerlessTenantResponse") - custom_config = _at(schema, "properties", "custom_config") - - assert custom_config["type"] == "object" - assert custom_config["additionalProperties"] is True - assert "custom_config" in cast(list[str], schema["required"]) - - -def test_owner_workspace_variant_excludes_custom_config_and_public_key(inner_spec: Json) -> None: - owner = _obj(_schema(inner_spec, "WorkspaceResponse")["properties"]) - - assert {"custom_config", "encrypt_public_key"}.isdisjoint(owner) - - -def test_runtime_credential_values_accept_arbitrary_fields(inner_spec: Json) -> None: - values = _at(_schema(inner_spec, "InnerRuntimeCredentialResponse"), "properties", "values") - - assert values["type"] == "object" - assert values["additionalProperties"] is True - - -def test_import_response_excludes_app_created(inner_spec: Json) -> None: - assert "app_created" not in _obj(_schema(inner_spec, "Import")["properties"]) - - -def test_dsl_export_documents_yaml_data_string(inner_spec: Json) -> None: - assert _at(_schema(inner_spec, "AppExportResponse"), "properties", "data")["type"] == "string" - - -def test_ownerless_create_sample_validates_custom_config_object(inner_spec: Json) -> None: - tenant = { - "id": "t", - "name": "n", - "plan": "p", - "status": "normal", - "encrypt_public_key": None, - "custom_config": {"remove_webapp_brand": True}, - "created_at": None, - "updated_at": None, - } - validator = _validator(inner_spec, "WorkspaceOwnerlessCreateResponse") - - validator.validate({"message": "ok", "tenant": tenant}) - validator.validate({"message": "ok", "tenant": {**tenant, "custom_config": {}}}) - with pytest.raises(ValidationError): - validator.validate({"message": "ok", "tenant": {**tenant, "custom_config": json.dumps({"a": 1})}}) - - -def test_credentials_sample_accepts_arbitrary_values(inner_spec: Json) -> None: - _validator(inner_spec, "InnerRuntimeCredentialsResolveResponse").validate( - {"credentials": [{"credential_id": "c", "provider": "p", "kind": "model", "values": {"api_key": "k", "x": 1}}]} - ) - - -ERROR_RESPONSE_REF = {"$ref": "#/components/schemas/ErrorResponse"} -LEGACY_SIMPLE_MESSAGE_RESPONSES = { - "create_enterprise_workspace": ("404",), - "enterprise_app_dsl_import": ("404",), - "enterprise_runtime_credentials_resolve": ("400", "404"), - "join_enterprise_workspace_member": ("400", "404"), -} -FILE_GRANT_BUSINESS_COMPONENTS = { - "400": {"InvalidRequestError", "GrantTtlTooLongError", "InvalidSubjectError"}, - "404": {"AppNotFoundError", "FileNotFoundError"}, -} - - -def _response_schema(operation: Json, code: str) -> Json: - return _at(operation, "responses", code, "content", "application/json", "schema") - - -def _alternatives(schema: Json) -> list[Json]: - return [_obj(item) for item in cast(list[object], schema["anyOf"])] if "anyOf" in schema else [schema] - - -def _operations_by_id(spec: Json) -> dict[str, Json]: - return {str(operation["operationId"]): operation for _, _, _, operation in _enterprise_operations(spec)} - - -def test_enterprise_operations_document_default_error_response(inner_spec: Json) -> None: - for path, method, _, operation in _enterprise_operations(inner_spec): - assert _response_schema(operation, "default") == ERROR_RESPONSE_REF, f"{method} {path}" - - -def test_enterprise_documented_errors_include_standard_error_body(inner_spec: Json) -> None: - for path, method, _, operation in _enterprise_operations(inner_spec): - for code in _obj(operation["responses"]): - if code == "default" or not code.startswith(("4", "5")): - continue - assert ERROR_RESPONSE_REF in _alternatives(_response_schema(operation, code)), f"{method} {path} {code}" - - -def test_legacy_message_statuses_document_simple_message(inner_spec: Json) -> None: - operations = _operations_by_id(inner_spec) - for operation_id, codes in LEGACY_SIMPLE_MESSAGE_RESPONSES.items(): - for code in codes: - assert {"$ref": "#/components/schemas/SimpleMessageResponse"} in _alternatives( - _response_schema(operations[operation_id], code) - ), f"{operation_id} {code}" - - -def test_dsl_export_not_found_union_documents_simple_message(inner_spec: Json) -> None: - operation = _operations_by_id(inner_spec)["enterprise_app_dsl_export"] - assert _alternatives(_response_schema(operation, "404"))[0] == { - "$ref": "#/components/schemas/EnterpriseAppDSLExportNotFoundResponse" - } - union = _alternatives(_schema(inner_spec, "EnterpriseAppDSLExportNotFoundResponse")) - assert {"$ref": "#/components/schemas/SimpleMessageResponse"} in union - - -def test_file_grants_document_business_error_components(inner_spec: Json) -> None: - operation = _operations_by_id(inner_spec)["enterprise_mint_file_grant"] - for code, expected in FILE_GRANT_BUSINESS_COMPONENTS.items(): - refs = { - str(item["$ref"]).rsplit("/", 1)[-1] - for item in _alternatives(_response_schema(operation, code)) - if "$ref" in item - } - assert expected <= refs, (code, refs) - for name in expected: - error_body = cast(list[object], _schema(inner_spec, name)["allOf"])[1] - assert _at(error_body, "properties", "status")["const"] == int(code), name diff --git a/api/tests/unit_tests/controllers/inner_api/test_auth_wraps.py b/api/tests/unit_tests/controllers/inner_api/test_wraps.py similarity index 100% rename from api/tests/unit_tests/controllers/inner_api/test_auth_wraps.py rename to api/tests/unit_tests/controllers/inner_api/test_wraps.py diff --git a/api/tests/unit_tests/controllers/inner_api/workspace/test_resource_deletion.py b/api/tests/unit_tests/controllers/inner_api/workspace/test_resource_deletion.py deleted file mode 100644 index 4c33d4e2d9e..00000000000 --- a/api/tests/unit_tests/controllers/inner_api/workspace/test_resource_deletion.py +++ /dev/null @@ -1,58 +0,0 @@ -from collections.abc import Callable -from inspect import unwrap -from types import SimpleNamespace -from unittest.mock import Mock -from uuid import UUID - -import pytest -from flask import Flask - -from controllers.inner_api.workspace import workspace as controller -from controllers.inner_api.wraps import InnerApiUnauthorizedError - -WORKSPACE_ID = UUID("0dd1d570-cf7b-4bab-9c7f-84ebe6c36985") - - -def test_delete_calls_cleanup_with_the_workspace_id(app: Flask, monkeypatch: pytest.MonkeyPatch) -> None: - service = Mock() - monkeypatch.setattr( - controller, - "application_services", - lambda: SimpleNamespace(workspaces=SimpleNamespace(resource_cleanup=service)), - ) - with app.test_request_context(method="DELETE"): - result = unwrap(controller.EnterpriseWorkspaceResourceDeletion.delete)( - controller.EnterpriseWorkspaceResourceDeletion(), WORKSPACE_ID - ) - assert result == ({"message": "workspace resource cleanup accepted."}, 202) - service.cleanup.assert_called_once_with(str(WORKSPACE_ID)) - - -def test_delete_requires_inner_api_key( - app: Flask, monkeypatch: pytest.MonkeyPatch, config_overrides: Callable[..., None] -) -> None: - config_overrides(INNER_API=True, INNER_API_KEY="secret", DEPLOYMENT_EDITION="CLOUD") - service = Mock() - monkeypatch.setattr( - controller, - "application_services", - lambda: SimpleNamespace(workspaces=SimpleNamespace(resource_cleanup=service)), - ) - with app.test_request_context(method="DELETE"): - with pytest.raises(InnerApiUnauthorizedError): - controller.EnterpriseWorkspaceResourceDeletion().delete(WORKSPACE_ID) - service.cleanup.assert_not_called() - - -def test_cleanup_failure_does_not_return_accepted(app: Flask, monkeypatch: pytest.MonkeyPatch) -> None: - service = Mock() - service.cleanup.side_effect = RuntimeError("broker unavailable") - monkeypatch.setattr( - controller, - "application_services", - lambda: SimpleNamespace(workspaces=SimpleNamespace(resource_cleanup=service)), - ) - with app.test_request_context(method="DELETE"), pytest.raises(RuntimeError, match="broker unavailable"): - unwrap(controller.EnterpriseWorkspaceResourceDeletion.delete)( - controller.EnterpriseWorkspaceResourceDeletion(), WORKSPACE_ID - ) diff --git a/api/tests/unit_tests/controllers/inner_api/workspace/test_workspace.py b/api/tests/unit_tests/controllers/inner_api/workspace/test_workspace.py index 6216747a731..5e8710a0d86 100644 --- a/api/tests/unit_tests/controllers/inner_api/workspace/test_workspace.py +++ b/api/tests/unit_tests/controllers/inner_api/workspace/test_workspace.py @@ -1,10 +1,10 @@ -"""Internal workspace transport behavior and admission.""" - from collections.abc import Callable from dataclasses import dataclass from datetime import datetime from inspect import unwrap +from types import SimpleNamespace from unittest.mock import Mock, create_autospec +from uuid import UUID import pytest from flask import Flask @@ -27,164 +27,226 @@ from services.workspace.contracts import CreatedWorkspace, WorkspaceMembership from services.workspace.provisioning_service import WorkspaceProvisioningService from services.workspace.service import WorkspaceQueryService, WorkspaceService +WORKSPACE_ID = UUID("0dd1d570-cf7b-4bab-9c7f-84ebe6c36985") + + +class TestResourceDeletionCases: + @staticmethod + def test_delete_calls_cleanup_with_the_workspace_id(app: Flask, monkeypatch: pytest.MonkeyPatch) -> None: + service = Mock() + monkeypatch.setattr( + controller, + "application_services", + lambda: SimpleNamespace(workspaces=SimpleNamespace(resource_cleanup=service)), + ) + with app.test_request_context(method="DELETE"): + result = unwrap(controller.EnterpriseWorkspaceResourceDeletion.delete)( + controller.EnterpriseWorkspaceResourceDeletion(), WORKSPACE_ID + ) + assert result == ({"message": "workspace resource cleanup accepted."}, 202) + service.cleanup.assert_called_once_with(str(WORKSPACE_ID)) + + @staticmethod + def test_delete_requires_inner_api_key( + app: Flask, monkeypatch: pytest.MonkeyPatch, config_overrides: Callable[..., None] + ) -> None: + config_overrides(INNER_API=True, INNER_API_KEY="secret", DEPLOYMENT_EDITION="CLOUD") + service = Mock() + monkeypatch.setattr( + controller, + "application_services", + lambda: SimpleNamespace(workspaces=SimpleNamespace(resource_cleanup=service)), + ) + with app.test_request_context(method="DELETE"): + with pytest.raises(InnerApiUnauthorizedError): + controller.EnterpriseWorkspaceResourceDeletion().delete(WORKSPACE_ID) + service.cleanup.assert_not_called() + + @staticmethod + def test_cleanup_failure_does_not_return_accepted(app: Flask, monkeypatch: pytest.MonkeyPatch) -> None: + service = Mock() + service.cleanup.side_effect = RuntimeError("broker unavailable") + monkeypatch.setattr( + controller, + "application_services", + lambda: SimpleNamespace(workspaces=SimpleNamespace(resource_cleanup=service)), + ) + with app.test_request_context(method="DELETE"), pytest.raises(RuntimeError, match="broker unavailable"): + unwrap(controller.EnterpriseWorkspaceResourceDeletion.delete)( + controller.EnterpriseWorkspaceResourceDeletion(), WORKSPACE_ID + ) + + +"""Internal workspace transport behavior and admission.""" + @dataclass class Services: workspaces: WorkspaceServices -@pytest.fixture -def provisioning(monkeypatch: pytest.MonkeyPatch) -> Mock: - provisioning = create_autospec(WorkspaceProvisioningService, instance=True) - services = WorkspaceServices( - queries=create_autospec(WorkspaceQueryService, instance=True), - management=create_autospec(WorkspaceService, instance=True), - provisioning=provisioning, - members=Mock(), - member_queries=Mock(), - owner_transfer=Mock(), - invitations=Mock(), - identity=Mock(), - resource_cleanup=Mock(), - ) - monkeypatch.setattr(controller, "application_services", lambda: Services(services)) - return provisioning - - -@pytest.mark.parametrize( - ("resource", "payload", "ownerless"), - [ - (controller.EnterpriseWorkspace, {"name": "Test", "owner_email": "owner@example.com"}, False), - (controller.EnterpriseWorkspaceNoOwnerEmail, {"name": "Test"}, True), - ], -) -def test_create_serialization( - app: Flask, - provisioning: Mock, - resource: type[ - controller.EnterpriseWorkspace - | controller.EnterpriseWorkspaceNoOwnerEmail - | controller.EnterpriseWorkspaceMember - ], - payload: dict[str, object], - ownerless: bool, -) -> None: - provisioning.create.return_value = CreatedWorkspace( - "w1", "Test", "sandbox", "normal", datetime(2026, 1, 1), None, "public-key", {"remove_webapp_brand": True} - ) - with app.test_request_context(method="POST", json=payload): - body = unwrap(resource.post)(resource()) - assert body["message"] == "enterprise workspace created." - assert body["tenant"]["created_at"] == "2026-01-01T00:00:00Z" - assert body["tenant"]["updated_at"] is None - if ownerless: - provisioning.create.assert_called_once_with(name="Test") - assert body["tenant"]["encrypt_public_key"] == "public-key" - assert body["tenant"]["custom_config"] == {"remove_webapp_brand": True} - else: - provisioning.create.assert_called_once_with(name="Test", owner_email="owner@example.com") - assert "encrypt_public_key" not in body["tenant"] - assert "custom_config" not in body["tenant"] - - -def test_ownerless_create_keeps_empty_custom_config_object(app: Flask, provisioning: Mock) -> None: - provisioning.create.return_value = CreatedWorkspace( - "w1", "Test", "sandbox", "normal", datetime(2026, 1, 1), None, "public-key", {} - ) - with app.test_request_context(method="POST", json={"name": "Test"}): - body = unwrap(controller.EnterpriseWorkspaceNoOwnerEmail.post)(controller.EnterpriseWorkspaceNoOwnerEmail()) - assert body["tenant"]["custom_config"] == {} - - -def test_owner_missing(app: Flask, provisioning: Mock) -> None: - provisioning.create.side_effect = WorkspaceOwnerNotFoundError() - with app.test_request_context(method="POST", json={"name": "Test", "owner_email": "absent@example.com"}): - result = unwrap(controller.EnterpriseWorkspace.post)(controller.EnterpriseWorkspace()) - assert result == ({"message": "owner account not found."}, 404) - - -def test_workspace_limit_remains_a_client_error(provisioning: Mock) -> None: - provisioning.create.side_effect = WorkspacesLimitExceededError("Workspace limit reached") - http_app = Flask(__name__) - http_api = Api(http_app, doc=False) - register_external_error_handlers(http_api) - - class Endpoint(Resource): - def post(self) -> ResponseReturnValue: - return unwrap(controller.EnterpriseWorkspace.post)(controller.EnterpriseWorkspace()) - - http_api.add_resource(Endpoint, "/workspaces") - response = http_app.test_client().post("/workspaces", json={"name": "Test", "owner_email": "owner@example.com"}) - assert response.status_code == 400 - assert response.json == {"code": "invalid_param", "message": "Workspace limit reached", "status": 400} - - MEMBER = {"workspace_id": "w1", "account_id": "a1", "email": "member@example.com", "current": True} -def test_join_member(app: Flask, provisioning: Mock) -> None: - provisioning.join_member.return_value = WorkspaceMembership("w1", "a1", "normal") - with app.test_request_context(method="POST", json=MEMBER): - body = unwrap(controller.EnterpriseWorkspaceMember.post)(controller.EnterpriseWorkspaceMember()) - assert body["member"] == {"workspace_id": "w1", "account_id": "a1", "role": "normal"} - provisioning.join_member.assert_called_once_with( - workspace_id="w1", account_id="a1", email="member@example.com", role="normal", operator_account_id=None +class TestWorkspaceCases: + @staticmethod + @pytest.fixture + def provisioning(monkeypatch: pytest.MonkeyPatch) -> Mock: + provisioning = create_autospec(WorkspaceProvisioningService, instance=True) + services = WorkspaceServices( + queries=create_autospec(WorkspaceQueryService, instance=True), + management=create_autospec(WorkspaceService, instance=True), + provisioning=provisioning, + members=Mock(), + member_queries=Mock(), + owner_transfer=Mock(), + invitations=Mock(), + identity=Mock(), + resource_cleanup=Mock(), + ) + monkeypatch.setattr(controller, "application_services", lambda: Services(services)) + return provisioning + + @staticmethod + @pytest.mark.parametrize( + ("resource", "payload", "ownerless"), + [ + (controller.EnterpriseWorkspace, {"name": "Test", "owner_email": "owner@example.com"}, False), + (controller.EnterpriseWorkspaceNoOwnerEmail, {"name": "Test"}, True), + ], ) + def test_create_serialization( + app: Flask, + provisioning: Mock, + resource: type[ + controller.EnterpriseWorkspace + | controller.EnterpriseWorkspaceNoOwnerEmail + | controller.EnterpriseWorkspaceMember + ], + payload: dict[str, object], + ownerless: bool, + ) -> None: + provisioning.create.return_value = CreatedWorkspace( + "w1", "Test", "sandbox", "normal", datetime(2026, 1, 1), None, "public-key", {"remove_webapp_brand": True} + ) + with app.test_request_context(method="POST", json=payload): + body = unwrap(resource.post)(resource()) + assert body["message"] == "enterprise workspace created." + assert body["tenant"]["created_at"] == "2026-01-01T00:00:00Z" + assert body["tenant"]["updated_at"] is None + if ownerless: + provisioning.create.assert_called_once_with(name="Test") + assert body["tenant"]["encrypt_public_key"] == "public-key" + assert body["tenant"]["custom_config"] == {"remove_webapp_brand": True} + else: + provisioning.create.assert_called_once_with(name="Test", owner_email="owner@example.com") + assert "encrypt_public_key" not in body["tenant"] + assert "custom_config" not in body["tenant"] + @staticmethod + def test_ownerless_create_keeps_empty_custom_config_object(app: Flask, provisioning: Mock) -> None: + provisioning.create.return_value = CreatedWorkspace( + "w1", "Test", "sandbox", "normal", datetime(2026, 1, 1), None, "public-key", {} + ) + with app.test_request_context(method="POST", json={"name": "Test"}): + body = unwrap(controller.EnterpriseWorkspaceNoOwnerEmail.post)(controller.EnterpriseWorkspaceNoOwnerEmail()) + assert body["tenant"]["custom_config"] == {} -@pytest.mark.parametrize( - ("error", "message", "status"), - [ - (AccountNotFoundError(), "account not found.", 404), - (WorkspaceNotFoundError(), "workspace not found.", 404), - (InvalidWorkspaceMemberRoleError("invalid workspace member role."), "invalid workspace member role.", 400), - (InvalidWorkspaceMemberRoleError("cannot join workspace as owner."), "cannot join workspace as owner.", 400), - ], -) -def test_join_errors(app: Flask, provisioning: Mock, error: Exception, message: str, status: int) -> None: - provisioning.join_member.side_effect = error - with app.test_request_context(method="POST", json=MEMBER): - result = unwrap(controller.EnterpriseWorkspaceMember.post)(controller.EnterpriseWorkspaceMember()) - assert result == ({"message": message}, status) + @staticmethod + def test_owner_missing(app: Flask, provisioning: Mock) -> None: + provisioning.create.side_effect = WorkspaceOwnerNotFoundError() + with app.test_request_context(method="POST", json={"name": "Test", "owner_email": "absent@example.com"}): + result = unwrap(controller.EnterpriseWorkspace.post)(controller.EnterpriseWorkspace()) + assert result == ({"message": "owner account not found."}, 404) + @staticmethod + def test_workspace_limit_remains_a_client_error(provisioning: Mock) -> None: + provisioning.create.side_effect = WorkspacesLimitExceededError("Workspace limit reached") + http_app = Flask(__name__) + http_api = Api(http_app, doc=False) + register_external_error_handlers(http_api) -@pytest.mark.parametrize( - "resource", - [controller.EnterpriseWorkspace, controller.EnterpriseWorkspaceNoOwnerEmail, controller.EnterpriseWorkspaceMember], -) -def test_invalid_payload_after_admission( - app: Flask, - provisioning: Mock, - config_overrides: Callable[..., None], - resource: type[ - controller.EnterpriseWorkspace - | controller.EnterpriseWorkspaceNoOwnerEmail - | controller.EnterpriseWorkspaceMember - ], -) -> None: - config_overrides(INNER_API=True, INNER_API_KEY="secret", DEPLOYMENT_EDITION="CLOUD") - with app.test_request_context(method="POST", json={}, headers={"X-Inner-Api-Key": "secret"}): - with pytest.raises(UnprocessableEntity): - resource().post() - provisioning.create.assert_not_called() - provisioning.join_member.assert_not_called() + class Endpoint(Resource): + def post(self) -> ResponseReturnValue: + return unwrap(controller.EnterpriseWorkspace.post)(controller.EnterpriseWorkspace()) + http_api.add_resource(Endpoint, "/workspaces") + response = http_app.test_client().post("/workspaces", json={"name": "Test", "owner_email": "owner@example.com"}) + assert response.status_code == 400 + assert response.json == {"code": "invalid_param", "message": "Workspace limit reached", "status": 400} -def test_inner_api_auth_precedes_payload_parsing( - app: Flask, provisioning: Mock, config_overrides: Callable[..., None] -) -> None: - config_overrides(INNER_API=True, INNER_API_KEY="secret", DEPLOYMENT_EDITION="CLOUD") - with app.test_request_context(method="POST", json={}): - with pytest.raises(InnerApiUnauthorizedError): - controller.EnterpriseWorkspace().post() - provisioning.create.assert_not_called() + @staticmethod + def test_join_member(app: Flask, provisioning: Mock) -> None: + provisioning.join_member.return_value = WorkspaceMembership("w1", "a1", "normal") + with app.test_request_context(method="POST", json=MEMBER): + body = unwrap(controller.EnterpriseWorkspaceMember.post)(controller.EnterpriseWorkspaceMember()) + assert body["member"] == {"workspace_id": "w1", "account_id": "a1", "role": "normal"} + provisioning.join_member.assert_called_once_with( + workspace_id="w1", account_id="a1", email="member@example.com", role="normal", operator_account_id=None + ) + @staticmethod + @pytest.mark.parametrize( + ("error", "message", "status"), + [ + (AccountNotFoundError(), "account not found.", 404), + (WorkspaceNotFoundError(), "workspace not found.", 404), + (InvalidWorkspaceMemberRoleError("invalid workspace member role."), "invalid workspace member role.", 400), + ( + InvalidWorkspaceMemberRoleError("cannot join workspace as owner."), + "cannot join workspace as owner.", + 400, + ), + ], + ) + def test_join_errors(app: Flask, provisioning: Mock, error: Exception, message: str, status: int) -> None: + provisioning.join_member.side_effect = error + with app.test_request_context(method="POST", json=MEMBER): + result = unwrap(controller.EnterpriseWorkspaceMember.post)(controller.EnterpriseWorkspaceMember()) + assert result == ({"message": message}, status) -def test_admitted_request_delegates_to_service( - app: Flask, provisioning: Mock, config_overrides: Callable[..., None] -) -> None: - config_overrides(INNER_API=True, INNER_API_KEY="secret", DEPLOYMENT_EDITION="CLOUD") - provisioning.create.return_value = CreatedWorkspace("w", "Test", "sandbox", "normal", None, None, None, {}) - with app.test_request_context(method="POST", json={"name": "Test"}, headers={"X-Inner-Api-Key": "secret"}): - controller.EnterpriseWorkspaceNoOwnerEmail().post() - provisioning.create.assert_called_once_with(name="Test") + @staticmethod + @pytest.mark.parametrize( + "resource", + [ + controller.EnterpriseWorkspace, + controller.EnterpriseWorkspaceNoOwnerEmail, + controller.EnterpriseWorkspaceMember, + ], + ) + def test_invalid_payload_after_admission( + app: Flask, + provisioning: Mock, + config_overrides: Callable[..., None], + resource: type[ + controller.EnterpriseWorkspace + | controller.EnterpriseWorkspaceNoOwnerEmail + | controller.EnterpriseWorkspaceMember + ], + ) -> None: + config_overrides(INNER_API=True, INNER_API_KEY="secret", DEPLOYMENT_EDITION="CLOUD") + with app.test_request_context(method="POST", json={}, headers={"X-Inner-Api-Key": "secret"}): + with pytest.raises(UnprocessableEntity): + resource().post() + provisioning.create.assert_not_called() + provisioning.join_member.assert_not_called() + + @staticmethod + def test_inner_api_auth_precedes_payload_parsing( + app: Flask, provisioning: Mock, config_overrides: Callable[..., None] + ) -> None: + config_overrides(INNER_API=True, INNER_API_KEY="secret", DEPLOYMENT_EDITION="CLOUD") + with app.test_request_context(method="POST", json={}): + with pytest.raises(InnerApiUnauthorizedError): + controller.EnterpriseWorkspace().post() + provisioning.create.assert_not_called() + + @staticmethod + def test_admitted_request_delegates_to_service( + app: Flask, provisioning: Mock, config_overrides: Callable[..., None] + ) -> None: + config_overrides(INNER_API=True, INNER_API_KEY="secret", DEPLOYMENT_EDITION="CLOUD") + provisioning.create.return_value = CreatedWorkspace("w", "Test", "sandbox", "normal", None, None, None, {}) + with app.test_request_context(method="POST", json={"name": "Test"}, headers={"X-Inner-Api-Key": "secret"}): + controller.EnterpriseWorkspaceNoOwnerEmail().post() + provisioning.create.assert_called_once_with(name="Test") diff --git a/api/tests/unit_tests/controllers/mcp/test_mcp.py b/api/tests/unit_tests/controllers/mcp/test_mcp.py index 088f085aa52..cf64cb19dbb 100644 --- a/api/tests/unit_tests/controllers/mcp/test_mcp.py +++ b/api/tests/unit_tests/controllers/mcp/test_mcp.py @@ -1,5 +1,3 @@ -"""Unit tests for controllers.mcp.mcp endpoints.""" - from __future__ import annotations import json @@ -11,50 +9,23 @@ from uuid import uuid4 import pytest from flask import Flask, Response +from flask.testing import FlaskClient from pydantic import ValidationError -from sqlalchemy import inspect +from sqlalchemy import inspect, select +from sqlalchemy.engine import Engine from sqlalchemy.orm import Session, sessionmaker +import controllers.mcp as mcp_controller import controllers.mcp.mcp as module +from controllers.mcp.mcp import MCPAppApi from models.engine import db from models.enums import EndUserType -from models.model import App, AppAnnotationSetting, AppMCPServer, AppModelConfig, EndUser, IconType +from models.model import App, AppAnnotationSetting, AppMCPServer, AppMode, AppModelConfig, EndUser, IconType from models.workflow import Workflow, WorkflowType -from tests.unit_tests.model_factories import make_end_user +from services.app.mcp_server_service import AppMCPServerStatus +from tests.helpers.unit.model_factories import make_end_user - -@pytest.fixture -def app() -> Iterator[Flask]: - app = Flask(__name__) - app.config["SQLALCHEMY_DATABASE_URI"] = "sqlite:///:memory:" - db.init_app(app) - - with app.app_context(): - for table in ( - App.__table__, - AppAnnotationSetting.__table__, - AppMCPServer.__table__, - AppModelConfig.__table__, - EndUser.__table__, - Workflow.__table__, - ): - table.create(db.engine) - with patch.object(module.session_factory, "get_session_maker", return_value=sessionmaker(db.engine)): - yield app - - -@pytest.fixture(autouse=True) -def mock_mcp_ns(monkeypatch: pytest.MonkeyPatch): - fake_ns = types.SimpleNamespace() - fake_ns.payload = None - fake_ns.models = {} - monkeypatch.setattr(module, "mcp_ns", fake_ns) - - -@pytest.fixture -def flask_req_ctx(app: Flask): - with app.test_request_context("/"): - yield +"""Unit tests for controllers.mcp.mcp endpoints.""" def fake_payload(data): @@ -145,390 +116,6 @@ class DummyResult: return {"jsonrpc": "2.0", "result": "ok", "id": 1} -@pytest.mark.usefixtures("flask_req_ctx") -class TestMCPAppApi: - @patch.object(module, "handle_mcp_request", return_value=DummyResult(), autospec=True) - def test_success_request(self, mock_handle): - fake_payload( - { - "jsonrpc": "2.0", - "method": "initialize", - "id": 1, - "params": { - "protocolVersion": "2024-11-05", - "capabilities": {}, - "clientInfo": {"name": "test-client", "version": "1.0"}, - }, - } - ) - - server = _server(module.AppMCPServerStatus.ACTIVE) - app = _app(module.AppMode.ADVANCED_CHAT, workflow_variables=[]) - - api = module.MCPAppApi() - api._get_mcp_server_and_app = MagicMock(return_value=(server, app)) - - post_fn = unwrap(api.post) - response = post_fn("server-1") - - assert isinstance(response, Response) - mock_handle.assert_called_once() - - def test_notification_initialized(self): - fake_payload( - { - "jsonrpc": "2.0", - "method": "notifications/initialized", - "params": {}, - } - ) - - server = _server(module.AppMCPServerStatus.ACTIVE) - app = _app(module.AppMode.ADVANCED_CHAT, workflow_variables=[]) - - api = module.MCPAppApi() - api._get_mcp_server_and_app = MagicMock(return_value=(server, app)) - - post_fn = unwrap(api.post) - response = post_fn("server-1") - - assert response.status_code == 202 - - def test_invalid_notification_method(self): - fake_payload( - { - "jsonrpc": "2.0", - "method": "notifications/invalid", - "params": {}, - } - ) - - server = _server(module.AppMCPServerStatus.ACTIVE) - app = _app(module.AppMode.ADVANCED_CHAT, workflow_variables=[]) - - api = module.MCPAppApi() - api._get_mcp_server_and_app = MagicMock(return_value=(server, app)) - - post_fn = unwrap(api.post) - - response = post_fn("server-1") - assert response.status_code == 202 - assert response.data == b"" - - def test_inactive_server(self): - fake_payload( - { - "jsonrpc": "2.0", - "method": "ping", - "id": 1, - "params": {}, - } - ) - - server = _server("inactive") - app = _app(module.AppMode.ADVANCED_CHAT, workflow_variables=[]) - - api = module.MCPAppApi() - api._get_mcp_server_and_app = MagicMock(return_value=(server, app)) - - post_fn = unwrap(api.post) - - response = post_fn("server-1") - assert response.status_code == 404 - assert response.get_json()["error"]["message"] == "Server Not Found" - - def test_invalid_payload(self, mocker): - fake_payload({"invalid": "data"}) - - api = module.MCPAppApi() - mocker.patch.object( - api, - "_get_mcp_server_and_app", - return_value=(_server(module.AppMCPServerStatus.ACTIVE), _app(module.AppMode.WORKFLOW)), - ) - post_fn = unwrap(api.post) - - with pytest.raises(ValidationError): - post_fn("server-1") - - def test_unknown_method_without_request_id_is_ignored(self): - fake_payload( - { - "jsonrpc": "2.0", - "method": "test", - "params": {}, - } - ) - - server = _server(module.AppMCPServerStatus.ACTIVE) - app = _app(module.AppMode.WORKFLOW, workflow_variables=[]) - - api = module.MCPAppApi() - api._get_mcp_server_and_app = MagicMock(return_value=(server, app)) - - post_fn = unwrap(api.post) - - response = post_fn("server-1") - assert response.status_code == 202 - assert response.data == b"" - - def test_server_not_found(self, mocker): - """Test when MCP server doesn't exist""" - fake_payload( - { - "jsonrpc": "2.0", - "method": "initialize", - "id": 1, - "params": { - "protocolVersion": "2024-11-05", - "capabilities": {}, - "clientInfo": {"name": "test-client", "version": "1.0"}, - }, - } - ) - - api = module.MCPAppApi() - mocker.patch.object(api, "_get_mcp_server_and_app", side_effect=module.MCPServerNotFoundError()) - - post_fn = unwrap(api.post) - - response = post_fn("server-1") - assert response.status_code == 404 - assert response.get_json()["error"]["message"] == "Server Not Found" - - def test_app_not_found(self, mocker): - """Test when app associated with server doesn't exist""" - fake_payload( - { - "jsonrpc": "2.0", - "method": "initialize", - "id": 1, - "params": { - "protocolVersion": "2024-11-05", - "capabilities": {}, - "clientInfo": {"name": "test-client", "version": "1.0"}, - }, - } - ) - - api = module.MCPAppApi() - mocker.patch.object(api, "_get_mcp_server_and_app", side_effect=module.MCPServerNotFoundError()) - - post_fn = unwrap(api.post) - - response = post_fn("server-1") - assert response.status_code == 404 - assert response.get_json()["error"]["message"] == "Server Not Found" - - def test_app_unavailable_no_workflow(self): - """Test when app has no workflow (ADVANCED_CHAT mode)""" - fake_payload( - { - "jsonrpc": "2.0", - "method": "initialize", - "id": 1, - "params": { - "protocolVersion": "2024-11-05", - "capabilities": {}, - "clientInfo": {"name": "test-client", "version": "1.0"}, - }, - } - ) - - server = _server(module.AppMCPServerStatus.ACTIVE) - app = _app(module.AppMode.ADVANCED_CHAT) - - api = module.MCPAppApi() - api._get_mcp_server_and_app = MagicMock(return_value=(server, app)) - - post_fn = unwrap(api.post) - - with pytest.raises(module.MCPRequestError) as exc_info: - post_fn("server-1") - assert "App is unavailable" in str(exc_info.value) - - def test_app_unavailable_no_model_config(self): - """Test when app has no model config (chat mode)""" - fake_payload( - { - "jsonrpc": "2.0", - "method": "initialize", - "id": 1, - "params": { - "protocolVersion": "2024-11-05", - "capabilities": {}, - "clientInfo": {"name": "test-client", "version": "1.0"}, - }, - } - ) - - server = _server(module.AppMCPServerStatus.ACTIVE) - app = _app(module.AppMode.CHAT) - - api = module.MCPAppApi() - api._get_mcp_server_and_app = MagicMock(return_value=(server, app)) - - post_fn = unwrap(api.post) - - with pytest.raises(module.MCPRequestError) as exc_info: - post_fn("server-1") - assert "App is unavailable" in str(exc_info.value) - - @patch.object(module, "handle_mcp_request", return_value=None, autospec=True) - def test_mcp_request_no_response(self, mock_handle): - """Test when handle_mcp_request returns None""" - fake_payload( - { - "jsonrpc": "2.0", - "method": "initialize", - "id": 1, - "params": { - "protocolVersion": "2024-11-05", - "capabilities": {}, - "clientInfo": {"name": "test-client", "version": "1.0"}, - }, - } - ) - - server = _server(module.AppMCPServerStatus.ACTIVE) - app = _app(module.AppMode.ADVANCED_CHAT, workflow_variables=[]) - - api = module.MCPAppApi() - api._get_mcp_server_and_app = MagicMock(return_value=(server, app)) - - post_fn = unwrap(api.post) - - with pytest.raises(module.MCPRequestError) as exc_info: - post_fn("server-1") - assert "No response generated" in str(exc_info.value) - - def test_workflow_mode_with_user_input_form(self): - """Test WORKFLOW mode app with user input form""" - fake_payload( - { - "jsonrpc": "2.0", - "method": "initialize", - "id": 1, - "params": { - "protocolVersion": "2024-11-05", - "capabilities": {}, - "clientInfo": {"name": "test-client", "version": "1.0"}, - }, - } - ) - - server = _server(module.AppMCPServerStatus.ACTIVE) - app = _app( - module.AppMode.WORKFLOW, - workflow_variables=[{"type": "text-input", "variable": "test_var", "label": "Test", "required": False}], - ) - - api = module.MCPAppApi() - api._get_mcp_server_and_app = MagicMock(return_value=(server, app)) - - with patch.object(module, "handle_mcp_request", return_value=DummyResult(), autospec=True): - post_fn = unwrap(api.post) - response = post_fn("server-1") - assert isinstance(response, Response) - - def test_chat_mode_with_model_config(self): - """Test CHAT mode app with model config""" - fake_payload( - { - "jsonrpc": "2.0", - "method": "initialize", - "id": 1, - "params": { - "protocolVersion": "2024-11-05", - "capabilities": {}, - "clientInfo": {"name": "test-client", "version": "1.0"}, - }, - } - ) - - server = _server(module.AppMCPServerStatus.ACTIVE) - app = _app(module.AppMode.CHAT, with_model_config=True) - - api = module.MCPAppApi() - api._get_mcp_server_and_app = MagicMock(return_value=(server, app)) - - with patch.object(module, "handle_mcp_request", return_value=DummyResult(), autospec=True): - post_fn = unwrap(api.post) - response = post_fn("server-1") - assert isinstance(response, Response) - - def test_invalid_mcp_request_format(self): - """Test invalid MCP request that doesn't match any type""" - fake_payload( - { - "jsonrpc": "2.0", - "method": "invalid_method_xyz", - "id": 1, - "params": {}, - } - ) - - server = _server(module.AppMCPServerStatus.ACTIVE) - app = _app(module.AppMode.ADVANCED_CHAT, workflow_variables=[]) - - api = module.MCPAppApi() - api._get_mcp_server_and_app = MagicMock(return_value=(server, app)) - - post_fn = unwrap(api.post) - - response = post_fn("server-1") - assert response.status_code == 400 - assert response.get_json() == { - "jsonrpc": "2.0", - "id": 1, - "error": {"code": module.mcp_types.METHOD_NOT_FOUND, "message": "Method not found"}, - } - - def test_validate_server_status_active(self): - """Test successful server status validation""" - api = module.MCPAppApi() - server = _server(module.AppMCPServerStatus.ACTIVE) - - # Should not raise an exception - api._validate_server_status(server) - - def test_convert_user_input_form_empty(self): - """Test converting empty user input form""" - api = module.MCPAppApi() - result = api._convert_user_input_form([]) - assert result == [] - - def test_invalid_user_input_form_validation(self): - """Test invalid user input form that fails validation""" - fake_payload( - { - "jsonrpc": "2.0", - "method": "initialize", - "id": 1, - "params": { - "protocolVersion": "2024-11-05", - "capabilities": {}, - "clientInfo": {"name": "test-client", "version": "1.0"}, - }, - } - ) - - server = _server(module.AppMCPServerStatus.ACTIVE) - app = _app( - module.AppMode.WORKFLOW, - workflow_variables=[{"type": "invalid-type", "variable": "test_var"}], - ) - - api = module.MCPAppApi() - api._get_mcp_server_and_app = MagicMock(return_value=(server, app)) - - post_fn = unwrap(api.post) - - with pytest.raises(module.MCPRequestError) as exc_info: - post_fn("server-1") - assert "Invalid user_input_form" in str(exc_info.value) - - _UNSUPPORTED_VERSION = "1999-01-01" @@ -561,237 +148,878 @@ def _tools_call_payload() -> dict[str, object]: } -def test_execution_reloads_attached_orm_objects_after_identity_preflight( - app: Flask, monkeypatch: pytest.MonkeyPatch -) -> None: - app_model = _app(module.AppMode.CHAT, with_model_config=True) - server = _server(module.AppMCPServerStatus.ACTIVE) - db.session.add_all([app_model, server]) - db.session.commit() - api = module.MCPAppApi() - lookup = api._get_mcp_server_and_app - snapshots: list[tuple[module.AppMCPServer, module.App]] = [] +class TestMcpCases: + @staticmethod + @pytest.fixture + def app() -> Iterator[Flask]: + app = Flask(__name__) + app.config["SQLALCHEMY_DATABASE_URI"] = "sqlite:///:memory:" + db.init_app(app) - def observe_lookup(code: str, session: Session) -> tuple[module.AppMCPServer, module.App]: - pair = lookup(code, session) - snapshots.append(pair) - return pair + with app.app_context(): + for table in ( + App.__table__, + AppAnnotationSetting.__table__, + AppMCPServer.__table__, + AppModelConfig.__table__, + EndUser.__table__, + Workflow.__table__, + ): + table.create(db.engine) + with patch.object(module.session_factory, "get_session_maker", return_value=sessionmaker(db.engine)): + yield app - def process( - _mcp_request: module.mcp_types.ClientRequest | module.mcp_types.ClientNotification, - _request_id: int | str | None, - fresh_app: module.App, - fresh_server: module.AppMCPServer, - _form: list[module.VariableEntity], - session: Session, - _protocol: str, - ) -> Response: - old_server, old_app = snapshots[0] - assert inspect(old_server).detached - assert inspect(old_app).detached - assert fresh_server is not old_server - assert fresh_app is not old_app - assert inspect(fresh_server).session is session - assert inspect(fresh_app).session is session - return Response("ok") + @staticmethod + @pytest.fixture(autouse=True) + def mock_mcp_ns(monkeypatch: pytest.MonkeyPatch): + fake_ns = types.SimpleNamespace() + fake_ns.payload = None + fake_ns.models = {} + monkeypatch.setattr(module, "mcp_ns", fake_ns) - monkeypatch.setattr(api, "_get_mcp_server_and_app", observe_lookup) - monkeypatch.setattr(api, "_process_mcp_message", process) - fake_payload({"jsonrpc": "2.0", "method": "ping", "id": 1}) - with app.test_request_context("/"): - response = api.post("server-1") - assert response.status_code == 200 - assert len(snapshots) == 2 + @staticmethod + @pytest.fixture + def flask_req_ctx(app: Flask): + with app.test_request_context("/"): + yield + @pytest.mark.usefixtures("flask_req_ctx") + class TestMCPAppApi: + @patch.object(module, "handle_mcp_request", return_value=DummyResult(), autospec=True) + def test_success_request(self, mock_handle): + fake_payload( + { + "jsonrpc": "2.0", + "method": "initialize", + "id": 1, + "params": { + "protocolVersion": "2024-11-05", + "capabilities": {}, + "clientInfo": {"name": "test-client", "version": "1.0"}, + }, + } + ) -@pytest.mark.usefixtures("app") -def test_identity_query_does_not_accept_a_server_bound_to_another_tenants_app() -> None: - app_model = _app(module.AppMode.CHAT) - server = _server(module.AppMCPServerStatus.ACTIVE) - server.tenant_id = str(uuid4()) - db.session.add_all([app_model, server]) - db.session.commit() - with sessionmaker(db.engine).begin() as session: - with pytest.raises(module.MCPServerNotFoundError): - module.MCPAppApi()._get_mcp_server_and_app("server-1", session) + server = _server(module.AppMCPServerStatus.ACTIVE) + app = _app(module.AppMode.ADVANCED_CHAT, workflow_variables=[]) + api = module.MCPAppApi() + api._get_mcp_server_and_app = MagicMock(return_value=(server, app)) -class TestMCPProtocolVersionNegotiationApi: - """MCP protocol version negotiation exercised through the HTTP controller layer. + post_fn = unwrap(api.post) + response = post_fn("server-1") - Covers the MCP-Protocol-Version header contract (resolution, rejection, threading) - and the serialized JSON responses seen by modern (2025-06-18) vs legacy (2024-11-05) - clients, including the back-compat guarantee that legacy responses carry none of the - structured-output fields. - """ + assert isinstance(response, Response) + mock_handle.assert_called_once() - def _make_api(self) -> module.MCPAppApi: + def test_notification_initialized(self): + fake_payload( + { + "jsonrpc": "2.0", + "method": "notifications/initialized", + "params": {}, + } + ) + + server = _server(module.AppMCPServerStatus.ACTIVE) + app = _app(module.AppMode.ADVANCED_CHAT, workflow_variables=[]) + + api = module.MCPAppApi() + api._get_mcp_server_and_app = MagicMock(return_value=(server, app)) + + post_fn = unwrap(api.post) + response = post_fn("server-1") + + assert response.status_code == 202 + + def test_invalid_notification_method(self): + fake_payload( + { + "jsonrpc": "2.0", + "method": "notifications/invalid", + "params": {}, + } + ) + + server = _server(module.AppMCPServerStatus.ACTIVE) + app = _app(module.AppMode.ADVANCED_CHAT, workflow_variables=[]) + + api = module.MCPAppApi() + api._get_mcp_server_and_app = MagicMock(return_value=(server, app)) + + post_fn = unwrap(api.post) + + response = post_fn("server-1") + assert response.status_code == 202 + assert response.data == b"" + + def test_inactive_server(self): + fake_payload( + { + "jsonrpc": "2.0", + "method": "ping", + "id": 1, + "params": {}, + } + ) + + server = _server("inactive") + app = _app(module.AppMode.ADVANCED_CHAT, workflow_variables=[]) + + api = module.MCPAppApi() + api._get_mcp_server_and_app = MagicMock(return_value=(server, app)) + + post_fn = unwrap(api.post) + + response = post_fn("server-1") + assert response.status_code == 404 + assert response.get_json()["error"]["message"] == "Server Not Found" + + def test_invalid_payload(self, mocker): + fake_payload({"invalid": "data"}) + + api = module.MCPAppApi() + mocker.patch.object( + api, + "_get_mcp_server_and_app", + return_value=(_server(module.AppMCPServerStatus.ACTIVE), _app(module.AppMode.WORKFLOW)), + ) + post_fn = unwrap(api.post) + + with pytest.raises(ValidationError): + post_fn("server-1") + + def test_unknown_method_without_request_id_is_ignored(self): + fake_payload( + { + "jsonrpc": "2.0", + "method": "test", + "params": {}, + } + ) + + server = _server(module.AppMCPServerStatus.ACTIVE) + app = _app(module.AppMode.WORKFLOW, workflow_variables=[]) + + api = module.MCPAppApi() + api._get_mcp_server_and_app = MagicMock(return_value=(server, app)) + + post_fn = unwrap(api.post) + + response = post_fn("server-1") + assert response.status_code == 202 + assert response.data == b"" + + def test_server_not_found(self, mocker): + """Test when MCP server doesn't exist""" + fake_payload( + { + "jsonrpc": "2.0", + "method": "initialize", + "id": 1, + "params": { + "protocolVersion": "2024-11-05", + "capabilities": {}, + "clientInfo": {"name": "test-client", "version": "1.0"}, + }, + } + ) + + api = module.MCPAppApi() + mocker.patch.object(api, "_get_mcp_server_and_app", side_effect=module.MCPServerNotFoundError()) + + post_fn = unwrap(api.post) + + response = post_fn("server-1") + assert response.status_code == 404 + assert response.get_json()["error"]["message"] == "Server Not Found" + + def test_app_not_found(self, mocker): + """Test when app associated with server doesn't exist""" + fake_payload( + { + "jsonrpc": "2.0", + "method": "initialize", + "id": 1, + "params": { + "protocolVersion": "2024-11-05", + "capabilities": {}, + "clientInfo": {"name": "test-client", "version": "1.0"}, + }, + } + ) + + api = module.MCPAppApi() + mocker.patch.object(api, "_get_mcp_server_and_app", side_effect=module.MCPServerNotFoundError()) + + post_fn = unwrap(api.post) + + response = post_fn("server-1") + assert response.status_code == 404 + assert response.get_json()["error"]["message"] == "Server Not Found" + + def test_app_unavailable_no_workflow(self): + """Test when app has no workflow (ADVANCED_CHAT mode)""" + fake_payload( + { + "jsonrpc": "2.0", + "method": "initialize", + "id": 1, + "params": { + "protocolVersion": "2024-11-05", + "capabilities": {}, + "clientInfo": {"name": "test-client", "version": "1.0"}, + }, + } + ) + + server = _server(module.AppMCPServerStatus.ACTIVE) + app = _app(module.AppMode.ADVANCED_CHAT) + + api = module.MCPAppApi() + api._get_mcp_server_and_app = MagicMock(return_value=(server, app)) + + post_fn = unwrap(api.post) + + with pytest.raises(module.MCPRequestError) as exc_info: + post_fn("server-1") + assert "App is unavailable" in str(exc_info.value) + + def test_app_unavailable_no_model_config(self): + """Test when app has no model config (chat mode)""" + fake_payload( + { + "jsonrpc": "2.0", + "method": "initialize", + "id": 1, + "params": { + "protocolVersion": "2024-11-05", + "capabilities": {}, + "clientInfo": {"name": "test-client", "version": "1.0"}, + }, + } + ) + + server = _server(module.AppMCPServerStatus.ACTIVE) + app = _app(module.AppMode.CHAT) + + api = module.MCPAppApi() + api._get_mcp_server_and_app = MagicMock(return_value=(server, app)) + + post_fn = unwrap(api.post) + + with pytest.raises(module.MCPRequestError) as exc_info: + post_fn("server-1") + assert "App is unavailable" in str(exc_info.value) + + @patch.object(module, "handle_mcp_request", return_value=None, autospec=True) + def test_mcp_request_no_response(self, mock_handle): + """Test when handle_mcp_request returns None""" + fake_payload( + { + "jsonrpc": "2.0", + "method": "initialize", + "id": 1, + "params": { + "protocolVersion": "2024-11-05", + "capabilities": {}, + "clientInfo": {"name": "test-client", "version": "1.0"}, + }, + } + ) + + server = _server(module.AppMCPServerStatus.ACTIVE) + app = _app(module.AppMode.ADVANCED_CHAT, workflow_variables=[]) + + api = module.MCPAppApi() + api._get_mcp_server_and_app = MagicMock(return_value=(server, app)) + + post_fn = unwrap(api.post) + + with pytest.raises(module.MCPRequestError) as exc_info: + post_fn("server-1") + assert "No response generated" in str(exc_info.value) + + def test_workflow_mode_with_user_input_form(self): + """Test WORKFLOW mode app with user input form""" + fake_payload( + { + "jsonrpc": "2.0", + "method": "initialize", + "id": 1, + "params": { + "protocolVersion": "2024-11-05", + "capabilities": {}, + "clientInfo": {"name": "test-client", "version": "1.0"}, + }, + } + ) + + server = _server(module.AppMCPServerStatus.ACTIVE) + app = _app( + module.AppMode.WORKFLOW, + workflow_variables=[{"type": "text-input", "variable": "test_var", "label": "Test", "required": False}], + ) + + api = module.MCPAppApi() + api._get_mcp_server_and_app = MagicMock(return_value=(server, app)) + + with patch.object(module, "handle_mcp_request", return_value=DummyResult(), autospec=True): + post_fn = unwrap(api.post) + response = post_fn("server-1") + assert isinstance(response, Response) + + def test_chat_mode_with_model_config(self): + """Test CHAT mode app with model config""" + fake_payload( + { + "jsonrpc": "2.0", + "method": "initialize", + "id": 1, + "params": { + "protocolVersion": "2024-11-05", + "capabilities": {}, + "clientInfo": {"name": "test-client", "version": "1.0"}, + }, + } + ) + + server = _server(module.AppMCPServerStatus.ACTIVE) + app = _app(module.AppMode.CHAT, with_model_config=True) + + api = module.MCPAppApi() + api._get_mcp_server_and_app = MagicMock(return_value=(server, app)) + + with patch.object(module, "handle_mcp_request", return_value=DummyResult(), autospec=True): + post_fn = unwrap(api.post) + response = post_fn("server-1") + assert isinstance(response, Response) + + def test_invalid_mcp_request_format(self): + """Test invalid MCP request that doesn't match any type""" + fake_payload( + { + "jsonrpc": "2.0", + "method": "invalid_method_xyz", + "id": 1, + "params": {}, + } + ) + + server = _server(module.AppMCPServerStatus.ACTIVE) + app = _app(module.AppMode.ADVANCED_CHAT, workflow_variables=[]) + + api = module.MCPAppApi() + api._get_mcp_server_and_app = MagicMock(return_value=(server, app)) + + post_fn = unwrap(api.post) + + response = post_fn("server-1") + assert response.status_code == 400 + assert response.get_json() == { + "jsonrpc": "2.0", + "id": 1, + "error": {"code": module.mcp_types.METHOD_NOT_FOUND, "message": "Method not found"}, + } + + def test_validate_server_status_active(self): + """Test successful server status validation""" + api = module.MCPAppApi() + server = _server(module.AppMCPServerStatus.ACTIVE) + + # Should not raise an exception + api._validate_server_status(server) + + def test_convert_user_input_form_empty(self): + """Test converting empty user input form""" + api = module.MCPAppApi() + result = api._convert_user_input_form([]) + assert result == [] + + def test_invalid_user_input_form_validation(self): + """Test invalid user input form that fails validation""" + fake_payload( + { + "jsonrpc": "2.0", + "method": "initialize", + "id": 1, + "params": { + "protocolVersion": "2024-11-05", + "capabilities": {}, + "clientInfo": {"name": "test-client", "version": "1.0"}, + }, + } + ) + + server = _server(module.AppMCPServerStatus.ACTIVE) + app = _app( + module.AppMode.WORKFLOW, + workflow_variables=[{"type": "invalid-type", "variable": "test_var"}], + ) + + api = module.MCPAppApi() + api._get_mcp_server_and_app = MagicMock(return_value=(server, app)) + + post_fn = unwrap(api.post) + + with pytest.raises(module.MCPRequestError) as exc_info: + post_fn("server-1") + assert "Invalid user_input_form" in str(exc_info.value) + + @staticmethod + def test_execution_reloads_attached_orm_objects_after_identity_preflight( + app: Flask, monkeypatch: pytest.MonkeyPatch + ) -> None: + app_model = _app(module.AppMode.CHAT, with_model_config=True) server = _server(module.AppMCPServerStatus.ACTIVE) - app = _app(module.AppMode.CHAT, with_model_config=True) + db.session.add_all([app_model, server]) + db.session.commit() api = module.MCPAppApi() - api._get_mcp_server_and_app = MagicMock(return_value=(server, app)) - api._retrieve_end_user = MagicMock(return_value=_end_user()) - return api + lookup = api._get_mcp_server_and_app + snapshots: list[tuple[module.AppMCPServer, module.App]] = [] - def _post( - self, flask_app: Flask, api: module.MCPAppApi, payload: dict[str, object], headers: dict[str, str] | None = None - ) -> Response: - fake_payload(payload) - post_fn = unwrap(api.post) - with flask_app.test_request_context(headers=headers): - return post_fn("server-1") + def observe_lookup(code: str, session: Session) -> tuple[module.AppMCPServer, module.App]: + pair = lookup(code, session) + snapshots.append(pair) + return pair - @pytest.mark.parametrize("version", sorted(module.mcp_types.SERVER_SUPPORTED_PROTOCOL_VERSIONS)) - def test_initialize_echoes_supported_body_version(self, app, version): - """Initialize echoes every supported client-requested version back unchanged.""" - api = self._make_api() + def process( + _mcp_request: module.mcp_types.ClientRequest | module.mcp_types.ClientNotification, + _request_id: int | str | None, + fresh_app: module.App, + fresh_server: module.AppMCPServer, + _form: list[module.VariableEntity], + session: Session, + _protocol: str, + ) -> Response: + old_server, old_app = snapshots[0] + assert inspect(old_server).detached + assert inspect(old_app).detached + assert fresh_server is not old_server + assert fresh_app is not old_app + assert inspect(fresh_server).session is session + assert inspect(fresh_app).session is session + return Response("ok") - response = self._post(app, api, _initialize_payload(version)) + monkeypatch.setattr(api, "_get_mcp_server_and_app", observe_lookup) + monkeypatch.setattr(api, "_process_mcp_message", process) + fake_payload({"jsonrpc": "2.0", "method": "ping", "id": 1}) + with app.test_request_context("/"): + response = api.post("server-1") + assert response.status_code == 200 + assert len(snapshots) == 2 - body = response.get_json() - assert body["result"]["protocolVersion"] == version + @staticmethod + @pytest.mark.usefixtures("app") + def test_identity_query_does_not_accept_a_server_bound_to_another_tenants_app() -> None: + app_model = _app(module.AppMode.CHAT) + server = _server(module.AppMCPServerStatus.ACTIVE) + server.tenant_id = str(uuid4()) + db.session.add_all([app_model, server]) + db.session.commit() + with sessionmaker(db.engine).begin() as session: + with pytest.raises(module.MCPServerNotFoundError): + module.MCPAppApi()._get_mcp_server_and_app("server-1", session) - def test_initialize_falls_back_for_unsupported_body_version(self, app): - """An unsupported requested version falls back to the server latest.""" - api = self._make_api() + class TestMCPProtocolVersionNegotiationApi: + """MCP protocol version negotiation exercised through the HTTP controller layer. - response = self._post(app, api, _initialize_payload(_UNSUPPORTED_VERSION)) + Covers the MCP-Protocol-Version header contract (resolution, rejection, threading) + and the serialized JSON responses seen by modern (2025-06-18) vs legacy (2024-11-05) + clients, including the back-compat guarantee that legacy responses carry none of the + structured-output fields. + """ - body = response.get_json() - assert body["result"]["protocolVersion"] == module.mcp_types.SERVER_LATEST_PROTOCOL_VERSION + def _make_api(self) -> module.MCPAppApi: + server = _server(module.AppMCPServerStatus.ACTIVE) + app = _app(module.AppMode.CHAT, with_model_config=True) + api = module.MCPAppApi() + api._get_mcp_server_and_app = MagicMock(return_value=(server, app)) + api._retrieve_end_user = MagicMock(return_value=_end_user()) + return api - def test_initialize_ignores_unsupported_header(self, app): - """Initialize negotiates via the request body, so its header is never rejected.""" - api = self._make_api() + def _post( + self, + flask_app: Flask, + api: module.MCPAppApi, + payload: dict[str, object], + headers: dict[str, str] | None = None, + ) -> Response: + fake_payload(payload) + post_fn = unwrap(api.post) + with flask_app.test_request_context(headers=headers): + return post_fn("server-1") - response = self._post( - app, - api, - _initialize_payload("2024-11-05"), - headers={"MCP-Protocol-Version": _UNSUPPORTED_VERSION}, - ) + @pytest.mark.parametrize("version", sorted(module.mcp_types.SERVER_SUPPORTED_PROTOCOL_VERSIONS)) + def test_initialize_echoes_supported_body_version(self, app, version): + """Initialize echoes every supported client-requested version back unchanged.""" + api = self._make_api() - body = response.get_json() - assert "error" not in body - assert body["result"]["protocolVersion"] == "2024-11-05" + response = self._post(app, api, _initialize_payload(version)) - @pytest.mark.parametrize("request_id", [5, None]) - def test_unsupported_header_returns_invalid_request_error(self, app, request_id): - """An unsupported header gets a JSON-RPC error echoing the request id (missing id -> null).""" - api = self._make_api() + body = response.get_json() + assert body["result"]["protocolVersion"] == version + + def test_initialize_falls_back_for_unsupported_body_version(self, app): + """An unsupported requested version falls back to the server latest.""" + api = self._make_api() + + response = self._post(app, api, _initialize_payload(_UNSUPPORTED_VERSION)) + + body = response.get_json() + assert body["result"]["protocolVersion"] == module.mcp_types.SERVER_LATEST_PROTOCOL_VERSION + + def test_initialize_ignores_unsupported_header(self, app): + """Initialize negotiates via the request body, so its header is never rejected.""" + api = self._make_api() - with patch.object(module, "handle_mcp_request", autospec=True) as mock_handle: response = self._post( app, api, - _tools_list_payload(request_id=request_id), + _initialize_payload("2024-11-05"), headers={"MCP-Protocol-Version": _UNSUPPORTED_VERSION}, ) - body = response.get_json() - assert response.status_code == 200 - assert body["jsonrpc"] == "2.0" - assert body["id"] == request_id - assert body["error"]["code"] == module.mcp_types.INVALID_REQUEST - assert _UNSUPPORTED_VERSION in body["error"]["message"] - mock_handle.assert_not_called() + body = response.get_json() + assert "error" not in body + assert body["result"]["protocolVersion"] == "2024-11-05" - def test_notification_with_unsupported_header_is_accepted(self, app): - """A notification is accepted (202, no body) even with an unsupported header.""" - api = self._make_api() + @pytest.mark.parametrize("request_id", [5, None]) + def test_unsupported_header_returns_invalid_request_error(self, app, request_id): + """An unsupported header gets a JSON-RPC error echoing the request id (missing id -> null).""" + api = self._make_api() - response = self._post( - app, - api, - {"jsonrpc": "2.0", "method": "notifications/initialized", "params": {}}, - headers={"MCP-Protocol-Version": _UNSUPPORTED_VERSION}, - ) + with patch.object(module, "handle_mcp_request", autospec=True) as mock_handle: + response = self._post( + app, + api, + _tools_list_payload(request_id=request_id), + headers={"MCP-Protocol-Version": _UNSUPPORTED_VERSION}, + ) - assert response.status_code == 202 + body = response.get_json() + assert response.status_code == 200 + assert body["jsonrpc"] == "2.0" + assert body["id"] == request_id + assert body["error"]["code"] == module.mcp_types.INVALID_REQUEST + assert _UNSUPPORTED_VERSION in body["error"]["message"] + mock_handle.assert_not_called() - @pytest.mark.parametrize("version", sorted(module.mcp_types.SERVER_SUPPORTED_PROTOCOL_VERSIONS)) - def test_supported_header_is_threaded_to_handler(self, app, version): - """Every supported header value is passed through to handle_mcp_request.""" - api = self._make_api() + def test_notification_with_unsupported_header_is_accepted(self, app): + """A notification is accepted (202, no body) even with an unsupported header.""" + api = self._make_api() - with patch.object(module, "handle_mcp_request", return_value=DummyResult(), autospec=True) as mock_handle: - self._post( + response = self._post( + app, + api, + {"jsonrpc": "2.0", "method": "notifications/initialized", "params": {}}, + headers={"MCP-Protocol-Version": _UNSUPPORTED_VERSION}, + ) + + assert response.status_code == 202 + + @pytest.mark.parametrize("version", sorted(module.mcp_types.SERVER_SUPPORTED_PROTOCOL_VERSIONS)) + def test_supported_header_is_threaded_to_handler(self, app, version): + """Every supported header value is passed through to handle_mcp_request.""" + api = self._make_api() + + with patch.object(module, "handle_mcp_request", return_value=DummyResult(), autospec=True) as mock_handle: + self._post( + app, + api, + _tools_list_payload(), + headers={"MCP-Protocol-Version": version}, + ) + + assert mock_handle.call_args.args[-1] == version + + def test_absent_header_defaults_to_back_compat_version(self, app): + """An absent header resolves to the spec's default version (2025-03-26).""" + api = self._make_api() + + with patch.object(module, "handle_mcp_request", return_value=DummyResult(), autospec=True) as mock_handle: + self._post(app, api, _tools_list_payload()) + + assert mock_handle.call_args.args[-1] == module.mcp_types.DEFAULT_NEGOTIATED_VERSION + + def test_tools_list_json_advertises_structured_output_for_modern_client(self, app): + """A 2025-06-18 client sees outputSchema and title in the serialized tool JSON.""" + api = self._make_api() + + response = self._post( app, api, _tools_list_payload(), - headers={"MCP-Protocol-Version": version}, + headers={"MCP-Protocol-Version": "2025-06-18"}, ) - assert mock_handle.call_args.args[-1] == version + tool = response.get_json()["result"]["tools"][0] + assert tool["outputSchema"] == {"type": "object"} + assert tool["title"] == "test_app" - def test_absent_header_defaults_to_back_compat_version(self, app): - """An absent header resolves to the spec's default version (2025-03-26).""" - api = self._make_api() + def test_tools_list_json_unchanged_for_legacy_client(self, app): + """A 2024-11-05 client sees exactly the pre-upgrade tool JSON keys.""" + api = self._make_api() - with patch.object(module, "handle_mcp_request", return_value=DummyResult(), autospec=True) as mock_handle: - self._post(app, api, _tools_list_payload()) + response = self._post( + app, + api, + _tools_list_payload(), + headers={"MCP-Protocol-Version": "2024-11-05"}, + ) - assert mock_handle.call_args.args[-1] == module.mcp_types.DEFAULT_NEGOTIATED_VERSION + tool = response.get_json()["result"]["tools"][0] + assert set(tool) == {"name", "description", "inputSchema"} - def test_tools_list_json_advertises_structured_output_for_modern_client(self, app): - """A 2025-06-18 client sees outputSchema and title in the serialized tool JSON.""" - api = self._make_api() + @patch("core.mcp.server.streamable_http.AppGenerateService") + def test_tools_call_json_includes_structured_content_for_modern_client(self, mock_app_generate, app): + """A 2025-06-18 client receives structuredContent alongside the text content.""" + api = self._make_api() + mock_app_generate.generate.return_value = {"answer": "test answer"} - response = self._post( - app, - api, - _tools_list_payload(), - headers={"MCP-Protocol-Version": "2025-06-18"}, + response = self._post( + app, + api, + _tools_call_payload(), + headers={"MCP-Protocol-Version": "2025-06-18"}, + ) + + result = response.get_json()["result"] + assert result["structuredContent"] == {"answer": "test answer"} + assert result["content"][0]["text"] == "test answer" + + @patch("core.mcp.server.streamable_http.AppGenerateService") + def test_tools_call_json_omits_structured_content_for_legacy_client(self, mock_app_generate, app): + """A 2024-11-05 client receives the pre-upgrade tools/call JSON without structuredContent.""" + api = self._make_api() + mock_app_generate.generate.return_value = {"answer": "test answer"} + + response = self._post( + app, + api, + _tools_call_payload(), + headers={"MCP-Protocol-Version": "2024-11-05"}, + ) + + result = response.get_json()["result"] + assert "structuredContent" not in result + assert result["content"][0]["text"] == "test answer" + + +"""MCP request parsing errors exercised through the real HTTP blueprint.""" + + +_URL = "/mcp/server/test-server/mcp" +_MODERN_HEADERS = { + "Accept": "application/json, text/event-stream", + "MCP-Protocol-Version": "2026-07-28", + "Mcp-Method": "server/discover", +} + + +class TestMcpRequestErrorsCases: + @staticmethod + @pytest.fixture + def app(sqlite_engine: Engine, monkeypatch: pytest.MonkeyPatch) -> Flask: + app = Flask(__name__) + app.config["SQLALCHEMY_DATABASE_URI"] = str(sqlite_engine.url) + db.init_app(app) + app.register_blueprint(mcp_controller.bp) + monkeypatch.setattr(module, "mcp_ns", mcp_controller.mcp_ns) + return app + + @staticmethod + @pytest.fixture + def client(app: Flask) -> FlaskClient: + return app.test_client() + + @staticmethod + @pytest.fixture(autouse=True) + def published_server(sqlite_session: Session) -> AppMCPServer: + app_id = str(uuid4()) + tenant_id = str(uuid4()) + model_config = AppModelConfig(app_id=app_id, user_input_form="[]") + app = App( + id=app_id, + tenant_id=tenant_id, + name="test_app", + mode=AppMode.CHAT, + enable_site=False, + enable_api=True, + app_model_config_id=model_config.id, + ) + server = AppMCPServer( + tenant_id=tenant_id, + app_id=app_id, + name="test_app", + description="Test MCP app", + server_code="test-server", + status=AppMCPServerStatus.ACTIVE, + parameters="{}", + ) + sqlite_session.add_all([app, model_config, server]) + sqlite_session.commit() + return server + + @staticmethod + @pytest.mark.parametrize("request_id", [1, 0, "discover-1", None]) + def test_server_discover_returns_method_not_found(client: FlaskClient, request_id: int | str | None) -> None: + response = client.post( + _URL, + json={ + "jsonrpc": "2.0", + "id": request_id, + "method": "server/discover", + "params": { + "_meta": { + "io.modelcontextprotocol/protocolVersion": "2026-07-28", + "io.modelcontextprotocol/clientInfo": {"name": "modern-client", "version": "1.0"}, + "io.modelcontextprotocol/clientCapabilities": { + "extensions": {"io.modelcontextprotocol/ui": {"mimeTypes": ["text/html;profile=mcp-app"]}} + }, + } + }, + }, + headers=_MODERN_HEADERS, ) - tool = response.get_json()["result"]["tools"][0] - assert tool["outputSchema"] == {"type": "object"} - assert tool["title"] == "test_app" + assert response.status_code == 400 + body = response.get_json() + assert body == { + "jsonrpc": "2.0", + "id": request_id, + "error": {"code": -32601, "message": "Method not found"}, + } + assert type(body["id"]) is type(request_id) - def test_tools_list_json_unchanged_for_legacy_client(self, app): - """A 2024-11-05 client sees exactly the pre-upgrade tool JSON keys.""" - api = self._make_api() - - response = self._post( - app, - api, - _tools_list_payload(), - headers={"MCP-Protocol-Version": "2024-11-05"}, + @staticmethod + def test_invalid_tool_params_return_safe_jsonrpc_error(client: FlaskClient) -> None: + response = client.post( + _URL, + json={ + "jsonrpc": "2.0", + "id": "tool-call-1", + "method": "tools/call", + "params": {"arguments": {"secret": "private-input-do-not-echo"}}, + }, ) - tool = response.get_json()["result"]["tools"][0] - assert set(tool) == {"name", "description", "inputSchema"} + assert response.status_code == 400 + body = response.get_json() + assert body["jsonrpc"] == "2.0" + assert body["id"] == "tool-call-1" + assert body["error"]["code"] == -32602 + assert "private-input-do-not-echo" not in response.get_data(as_text=True) + assert "validation error" not in body["error"]["message"].lower() - @patch("core.mcp.server.streamable_http.AppGenerateService") - def test_tools_call_json_includes_structured_content_for_modern_client(self, mock_app_generate, app): - """A 2025-06-18 client receives structuredContent alongside the text content.""" - api = self._make_api() - mock_app_generate.generate.return_value = {"answer": "test answer"} - - response = self._post( - app, - api, - _tools_call_payload(), - headers={"MCP-Protocol-Version": "2025-06-18"}, + @staticmethod + @pytest.mark.parametrize("method", ["server/discover", "notifications/unknown", "tools/call"]) + def test_unparseable_notification_has_no_jsonrpc_response(client: FlaskClient, method: str) -> None: + response = client.post( + _URL, + json={"jsonrpc": "2.0", "method": method, "params": {}}, + headers={**_MODERN_HEADERS, "Mcp-Method": method}, ) - result = response.get_json()["result"] - assert result["structuredContent"] == {"answer": "test answer"} - assert result["content"][0]["text"] == "test answer" + assert response.status_code == 202 + assert response.data == b"" - @patch("core.mcp.server.streamable_http.AppGenerateService") - def test_tools_call_json_omits_structured_content_for_legacy_client(self, mock_app_generate, app): - """A 2024-11-05 client receives the pre-upgrade tools/call JSON without structuredContent.""" - api = self._make_api() - mock_app_generate.generate.return_value = {"answer": "test answer"} - - response = self._post( - app, - api, - _tools_call_payload(), - headers={"MCP-Protocol-Version": "2024-11-05"}, + @staticmethod + def test_legacy_initialize_and_list_tools_remain_compatible( + client: FlaskClient, published_server: AppMCPServer, sqlite_session: Session + ) -> None: + response = client.post( + _URL, + json={ + "jsonrpc": "2.0", + "id": 0, + "method": "initialize", + "params": { + "protocolVersion": "2024-11-05", + "capabilities": {}, + "clientInfo": {"name": "legacy-client", "version": "1.0"}, + }, + }, ) - result = response.get_json()["result"] - assert "structuredContent" not in result - assert result["content"][0]["text"] == "test answer" + assert response.status_code == 200 + initialized = response.get_json() + assert initialized["id"] == 0 + assert initialized["result"]["protocolVersion"] == "2024-11-05" + end_user = sqlite_session.scalar(select(EndUser).where(EndUser.session_id == published_server.id)) + assert end_user is not None + assert end_user.tenant_id == published_server.tenant_id + assert end_user.app_id == published_server.app_id + assert end_user.type == EndUserType.MCP + assert end_user.name == "legacy-client@1.0" + + headers = {"MCP-Protocol-Version": "2024-11-05"} + notification = client.post( + _URL, + json={"jsonrpc": "2.0", "method": "notifications/initialized"}, + headers=headers, + ) + assert notification.status_code == 202 + assert notification.data == b"" + + listed = client.post( + _URL, + json={"jsonrpc": "2.0", "id": "list-1", "method": "tools/list"}, + headers=headers, + ) + assert listed.status_code == 200 + body = listed.get_json() + assert body["id"] == "list-1" + assert body["result"]["tools"] == [ + { + "name": "test_app", + "description": "Test MCP app", + "inputSchema": { + "type": "object", + "properties": {"query": {"type": "string", "description": "User Input/Question content"}}, + "required": ["query"], + }, + } + ] + + +"""SQLite-backed unit tests for MCP app controller database helpers.""" + + +TENANT_ID = "11111111-1111-1111-1111-111111111111" +APP_ID = "22222222-2222-2222-2222-222222222222" + + +class TestMcpAppApiCases: + @staticmethod + @pytest.mark.parametrize("sqlite_session", [(App, AppMCPServer)], indirect=True) + def test_get_mcp_server_and_app_returns_persisted_models(sqlite_session: Session) -> None: + app = App( + tenant_id=TENANT_ID, + name="MCP App", + description="", + mode=AppMode.ADVANCED_CHAT, + icon_type=IconType.EMOJI, + icon="🤖", + icon_background="#FFFFFF", + enable_site=False, + enable_api=False, + max_active_requests=0, + ) + app.id = APP_ID + server = AppMCPServer( + tenant_id=TENANT_ID, + app_id=APP_ID, + name="Test MCP Server", + description="Test server", + server_code="server-1", + status=AppMCPServerStatus.ACTIVE, + parameters="{}", + ) + sqlite_session.add_all([app, server]) + sqlite_session.commit() + + result_server, result_app = MCPAppApi()._get_mcp_server_and_app("server-1", sqlite_session) + + assert result_server is server + assert result_app is app diff --git a/api/tests/unit_tests/controllers/mcp/test_mcp_request_errors.py b/api/tests/unit_tests/controllers/mcp/test_mcp_request_errors.py deleted file mode 100644 index c6aa7bb035d..00000000000 --- a/api/tests/unit_tests/controllers/mcp/test_mcp_request_errors.py +++ /dev/null @@ -1,188 +0,0 @@ -"""MCP request parsing errors exercised through the real HTTP blueprint.""" - -from uuid import uuid4 - -import pytest -from flask import Flask -from flask.testing import FlaskClient -from sqlalchemy import select -from sqlalchemy.engine import Engine -from sqlalchemy.orm import Session - -import controllers.mcp as mcp_controller -import controllers.mcp.mcp as module -from models.engine import db -from models.enums import EndUserType -from models.model import App, AppMCPServer, AppMode, AppModelConfig, EndUser -from services.app.mcp_server_service import AppMCPServerStatus - -_URL = "/mcp/server/test-server/mcp" -_MODERN_HEADERS = { - "Accept": "application/json, text/event-stream", - "MCP-Protocol-Version": "2026-07-28", - "Mcp-Method": "server/discover", -} - - -@pytest.fixture -def app(sqlite_engine: Engine, monkeypatch: pytest.MonkeyPatch) -> Flask: - app = Flask(__name__) - app.config["SQLALCHEMY_DATABASE_URI"] = str(sqlite_engine.url) - db.init_app(app) - app.register_blueprint(mcp_controller.bp) - monkeypatch.setattr(module, "mcp_ns", mcp_controller.mcp_ns) - return app - - -@pytest.fixture -def client(app: Flask) -> FlaskClient: - return app.test_client() - - -@pytest.fixture(autouse=True) -def published_server(sqlite_session: Session) -> AppMCPServer: - app_id = str(uuid4()) - tenant_id = str(uuid4()) - model_config = AppModelConfig(app_id=app_id, user_input_form="[]") - app = App( - id=app_id, - tenant_id=tenant_id, - name="test_app", - mode=AppMode.CHAT, - enable_site=False, - enable_api=True, - app_model_config_id=model_config.id, - ) - server = AppMCPServer( - tenant_id=tenant_id, - app_id=app_id, - name="test_app", - description="Test MCP app", - server_code="test-server", - status=AppMCPServerStatus.ACTIVE, - parameters="{}", - ) - sqlite_session.add_all([app, model_config, server]) - sqlite_session.commit() - return server - - -@pytest.mark.parametrize("request_id", [1, 0, "discover-1", None]) -def test_server_discover_returns_method_not_found(client: FlaskClient, request_id: int | str | None) -> None: - response = client.post( - _URL, - json={ - "jsonrpc": "2.0", - "id": request_id, - "method": "server/discover", - "params": { - "_meta": { - "io.modelcontextprotocol/protocolVersion": "2026-07-28", - "io.modelcontextprotocol/clientInfo": {"name": "modern-client", "version": "1.0"}, - "io.modelcontextprotocol/clientCapabilities": { - "extensions": {"io.modelcontextprotocol/ui": {"mimeTypes": ["text/html;profile=mcp-app"]}} - }, - } - }, - }, - headers=_MODERN_HEADERS, - ) - - assert response.status_code == 400 - body = response.get_json() - assert body == { - "jsonrpc": "2.0", - "id": request_id, - "error": {"code": -32601, "message": "Method not found"}, - } - assert type(body["id"]) is type(request_id) - - -def test_invalid_tool_params_return_safe_jsonrpc_error(client: FlaskClient) -> None: - response = client.post( - _URL, - json={ - "jsonrpc": "2.0", - "id": "tool-call-1", - "method": "tools/call", - "params": {"arguments": {"secret": "private-input-do-not-echo"}}, - }, - ) - - assert response.status_code == 400 - body = response.get_json() - assert body["jsonrpc"] == "2.0" - assert body["id"] == "tool-call-1" - assert body["error"]["code"] == -32602 - assert "private-input-do-not-echo" not in response.get_data(as_text=True) - assert "validation error" not in body["error"]["message"].lower() - - -@pytest.mark.parametrize("method", ["server/discover", "notifications/unknown", "tools/call"]) -def test_unparseable_notification_has_no_jsonrpc_response(client: FlaskClient, method: str) -> None: - response = client.post( - _URL, - json={"jsonrpc": "2.0", "method": method, "params": {}}, - headers={**_MODERN_HEADERS, "Mcp-Method": method}, - ) - - assert response.status_code == 202 - assert response.data == b"" - - -def test_legacy_initialize_and_list_tools_remain_compatible( - client: FlaskClient, published_server: AppMCPServer, sqlite_session: Session -) -> None: - response = client.post( - _URL, - json={ - "jsonrpc": "2.0", - "id": 0, - "method": "initialize", - "params": { - "protocolVersion": "2024-11-05", - "capabilities": {}, - "clientInfo": {"name": "legacy-client", "version": "1.0"}, - }, - }, - ) - - assert response.status_code == 200 - initialized = response.get_json() - assert initialized["id"] == 0 - assert initialized["result"]["protocolVersion"] == "2024-11-05" - end_user = sqlite_session.scalar(select(EndUser).where(EndUser.session_id == published_server.id)) - assert end_user is not None - assert end_user.tenant_id == published_server.tenant_id - assert end_user.app_id == published_server.app_id - assert end_user.type == EndUserType.MCP - assert end_user.name == "legacy-client@1.0" - - headers = {"MCP-Protocol-Version": "2024-11-05"} - notification = client.post( - _URL, - json={"jsonrpc": "2.0", "method": "notifications/initialized"}, - headers=headers, - ) - assert notification.status_code == 202 - assert notification.data == b"" - - listed = client.post( - _URL, - json={"jsonrpc": "2.0", "id": "list-1", "method": "tools/list"}, - headers=headers, - ) - assert listed.status_code == 200 - body = listed.get_json() - assert body["id"] == "list-1" - assert body["result"]["tools"] == [ - { - "name": "test_app", - "description": "Test MCP app", - "inputSchema": { - "type": "object", - "properties": {"query": {"type": "string", "description": "User Input/Question content"}}, - "required": ["query"], - }, - } - ] diff --git a/api/tests/unit_tests/controllers/openapi/auth/test_context.py b/api/tests/unit_tests/controllers/openapi/auth/test_context.py index 5e6a350f635..18107fd3970 100644 --- a/api/tests/unit_tests/controllers/openapi/auth/test_context.py +++ b/api/tests/unit_tests/controllers/openapi/auth/test_context.py @@ -1,17 +1,34 @@ -"""The store hands handlers what the route loaded, and nothing else.""" - from __future__ import annotations from collections.abc import Callable import pytest -from sqlalchemy.orm import Session +from sqlalchemy import select +from sqlalchemy.orm import Session, object_session from controllers.openapi.auth.context import Context, RouteContractError +from controllers.openapi.auth.requirements import CheckWorkspaceMember from core.logging.context import clear_request_context, init_request_context from machinery.context import RequestContext +from models import TenantAccountJoin +from models.account import TenantAccountRole from models.model import EndUser -from tests.unit_tests.controllers.openapi.auth._world import ACCOUNT_ID, TENANT_ID, make_ctx, make_tenant +from tests.e2e_tests.containers.controllers.openapi.conftest import context_for +from tests.helpers.unit.controllers.openapi.auth._world import ( + ACCOUNT_ID, + APP_ID, + OTHER_TENANT_ID, + TENANT_ID, + make_account, + make_app, + make_ctx, + make_membership, + make_tenant, + persist, +) + +"""The store hands handlers what the route loaded, and nothing else.""" + READERS: dict[str, Callable[[Context], object]] = { "app": lambda ctx: ctx.app, @@ -52,3 +69,66 @@ def test_application_context_keeps_identity_after_request_and_session_end(sqlite assert ctx.request_context == snapshot assert snapshot == RequestContext(snapshot.request_id, snapshot.trace_id, ACCOUNT_ID, TENANT_ID) + + +"""Keep the integration handler context aligned with request-scoped identity loading.""" + + +@pytest.mark.parametrize("view_args", [{"workspace_id": TENANT_ID}, {"app_id": APP_ID}]) +def test_workspace_requirement_binds_fresh_caller_to_requested_workspace( + sqlite_session: Session, view_args: dict[str, str] +) -> None: + account = make_account() + original_workspace = make_tenant(tenant_id=OTHER_TENANT_ID) + requested_membership = make_membership(TenantAccountRole.ADMIN) + requested_membership.current = False + persist( + sqlite_session, + account, + original_workspace, + make_tenant(), + make_app(), + requested_membership, + TenantAccountJoin( + tenant_id=OTHER_TENANT_ID, + account_id=ACCOUNT_ID, + current=True, + role=TenantAccountRole.OWNER, + ), + ) + account.set_current_tenant_with_session(original_workspace, session=sqlite_session) + + ctx = context_for( + account, + session=sqlite_session, + view_args=view_args, + requirements=(CheckWorkspaceMember(),), + ) + + assert ctx.account is not account + assert object_session(ctx.account) is None + assert ctx.account.current_tenant_id == ctx.workspace.id == TENANT_ID + assert ctx.account.role == ctx.workspace_role == TenantAccountRole.ADMIN + assert account.current_tenant_id == OTHER_TENANT_ID + assert ( + sqlite_session.scalar( + select(TenantAccountJoin.tenant_id).where( + TenantAccountJoin.account_id == ACCOUNT_ID, TenantAccountJoin.current.is_(True) + ) + ) + == OTHER_TENANT_ID + ) + if "app_id" in view_args: + assert ctx.app.id == APP_ID + + +def test_app_context_without_workspace_requirement_does_not_bind_workspace(sqlite_session: Session) -> None: + account = make_account() + persist(sqlite_session, account, make_tenant(), make_app(), make_membership()) + + ctx = context_for(account, session=sqlite_session, view_args={"app_id": APP_ID}) + + assert ctx.app.id == APP_ID + assert ctx.account.current_tenant_id is None + with pytest.raises(RouteContractError, match="workspace was not loaded"): + _ = ctx.workspace diff --git a/api/tests/unit_tests/controllers/openapi/auth/test_loaders.py b/api/tests/unit_tests/controllers/openapi/auth/test_loaders.py index c16bebd7e1d..fb525d3cea1 100644 --- a/api/tests/unit_tests/controllers/openapi/auth/test_loaders.py +++ b/api/tests/unit_tests/controllers/openapi/auth/test_loaders.py @@ -20,8 +20,7 @@ from controllers.openapi.auth.subjects import Subject from models import App from models.account import TenantAccountRole, TenantStatus from services.app_service import AppService - -from ._world import ( +from tests.helpers.unit.controllers.openapi.auth._world import ( ACCOUNT_ID, APP_ID, OTHER_TENANT_ID, diff --git a/api/tests/unit_tests/controllers/openapi/auth/test_pipelines.py b/api/tests/unit_tests/controllers/openapi/auth/test_pipelines.py index 36bd76cc894..9b9194d40f5 100644 --- a/api/tests/unit_tests/controllers/openapi/auth/test_pipelines.py +++ b/api/tests/unit_tests/controllers/openapi/auth/test_pipelines.py @@ -37,8 +37,7 @@ from libs.oauth_bearer import AuthContext, try_get_auth_ctx from machinery.context import AppRequestContext, RequestContext from services.app_service import AppService from services.enterprise.enterprise_service import WebAppAccessMode - -from ._world import ( +from tests.helpers.unit.controllers.openapi.auth._world import ( ACCOUNT_ID, APP_ID, TENANT_ID, diff --git a/api/tests/unit_tests/controllers/openapi/auth/test_requirements.py b/api/tests/unit_tests/controllers/openapi/auth/test_requirements.py index ac89a860edb..be751f9abff 100644 --- a/api/tests/unit_tests/controllers/openapi/auth/test_requirements.py +++ b/api/tests/unit_tests/controllers/openapi/auth/test_requirements.py @@ -1,37 +1,81 @@ from __future__ import annotations -from collections.abc import Callable +import contextvars +import uuid +from collections.abc import Callable, Iterator +from contextlib import ExitStack +from dataclasses import dataclass, field +from datetime import UTC, datetime, timedelta +from enum import StrEnum, auto +from functools import singledispatch +from typing import Final from unittest.mock import patch import pytest from flask import Flask -from sqlalchemy.orm import Session -from werkzeug.exceptions import Forbidden, NotFound +from flask_login import LoginManager, user_logged_in +from sqlalchemy.orm import Session, sessionmaker +from werkzeug.exceptions import Forbidden, ImATeapot, NotFound +from werkzeug.routing import Rule +from werkzeug.test import TestResponse -from controllers.common.rbac import PlainApp, RBACCheck, RBACPermission +import libs.oauth_bearer as oauth_bearer_module +import libs.rate_limit as rate_limit_module +from app_factory import create_flask_app_with_configs +from constants.oauth_bearer import Scope, TokenType +from controllers.common.rbac import AgentBehindApp, PlainApp, RBACCheck, RBACPermission, Workspace +from controllers.openapi import bp as openapi_bp +from controllers.openapi._app_settings import REGULAR_MODES +from controllers.openapi._catalog import CATALOG_HEADER, catalog_for from controllers.openapi._contract import op_of -from controllers.openapi._errors import OpenApiErrorCode, OpenApiErrorFormatter, ServiceApiDisabled +from controllers.openapi._errors import ( + OpenApiErrorCode, + OpenApiErrorFormatter, + ServiceApiDisabled, + WebAppAccessRequiresEE, +) from controllers.openapi._models import ServiceApiPatch from controllers.openapi.app_info import AgentServiceApiApi, ServiceApiApi +from controllers.openapi.auth import subjects from controllers.openapi.auth.context import RouteContractError from controllers.openapi.auth.requirements import ( CheckAppAccess, CheckAppApiEnabled, + CheckAppMode, + CheckAppQuota, + CheckPluginInstallSetting, CheckRBACPermission, + CheckScope, CheckSubject, + CheckWebAppAuthEnterprise, + CheckWorkspaceInvitationQuota, + CheckWorkspaceMember, CheckWorkspaceRole, + Requirement, assert_license_valid, ) -from controllers.openapi.auth.subjects import AccountSubject +from controllers.openapi.auth.spec import EndpointSpec +from controllers.openapi.auth.subjects import AccountSubject, ExternalSsoSubject, ResourceAccessSubject +from controllers.openapi.human_input_form import CheckFormSurface from enums import DeploymentEdition -from models import AppMode -from models.account import AccountStatus, TenantAccountRole -from services.enterprise.enterprise_service import WebAppAccessMode, WebAppSettings +from enums import WebAppAccessMode as _auth_matrix_WebAppAccessMode +from extensions.ext_application_services import ApplicationServices +from libs.oauth_bearer import BearerAuthenticator, ResolvedRow, sha256_hex +from models import AppMode as _requirements_AppMode +from models.account import Account, AccountStatus, Tenant, TenantAccountJoin, TenantAccountRole +from models.enums import EndUserType +from models.model import App, EndUser +from models.model import AppMode as _auth_matrix_AppMode +from models.oauth import OAuthAccessToken +from services.enterprise.enterprise_service import EnterpriseService, WebAppSettings +from services.enterprise.enterprise_service import WebAppAccessMode as _requirements_WebAppAccessMode from services.entities.feature_entities import ( LicenseStatus, + SystemFeatureModel, ) - -from ._world import ( +from services.system_feature_service import SystemFeatureService +from tests.helpers.unit.config_override import apply_config_overrides +from tests.helpers.unit.controllers.openapi.auth._world import ( APP_ID, CLIENT_ID, TOKEN_ID, @@ -55,212 +99,2423 @@ ENFORCE_RBAC = "controllers.openapi.auth.requirements.enforce_rbac_checks" APP_FETCH = "controllers.openapi.auth.loaders.AppService.get_app_by_id" -def test_subject_check_emits_the_wrong_surface_audit(app: Flask, sqlite_session: Session) -> None: - subject = sso_subject() - - with app.test_request_context("/openapi/v1/account"): - with patch("controllers.openapi.auth.requirements.emit_wrong_surface") as emit: - requirement = CheckSubject(allowed=[AccountSubject]) - with pytest.raises(Forbidden, match="unsupported_token_type"): - requirement.run(subject, make_ctx(sqlite_session, subject=subject), sqlite_session) - - emit.assert_called_once_with( - subject_type="external_sso", - attempted_path="/openapi/v1/account", - client_id=CLIENT_ID, - token_id=TOKEN_ID, - ) - - -@pytest.mark.parametrize( - ("status", "denied"), - [ - (LicenseStatus.INACTIVE, True), - (LicenseStatus.LOST, True), - (LicenseStatus.NONE, False), - ], -) -def test_assert_license_valid_denies_only_dead_licences(status: LicenseStatus, denied: bool) -> None: - with patch(FEATURES, return_value=system_features(license_status=status)): - if denied: - with pytest.raises(Forbidden, match="license_invalid"): - assert_license_valid() - else: - assert_license_valid() - - -def test_an_app_requirement_off_an_app_route_is_a_wiring_bug( - sqlite_session: Session, monkeypatch: pytest.MonkeyPatch -) -> None: - """App requirements are declared per endpoint, so a route with no `app_id` - can only carry one by mistake. It raises rather than passing quietly: - skipping would turn a misdeclaration into a check that silently never runs. - """ - monkeypatch.setattr(APP_FETCH, never_reached) - subject = account_subject() - - with pytest.raises(RouteContractError, match="app_id is not a path parameter"): - CheckAppApiEnabled().run(subject, make_ctx(sqlite_session, subject=subject), sqlite_session) - - -@pytest.mark.parametrize( - ("mode", "view"), - [ - (AppMode.CHAT, ServiceApiApi.patch), - (AppMode.WORKFLOW, ServiceApiApi.patch), - (AppMode.AGENT, AgentServiceApiApi.patch), - ], -) -def test_a_disabled_service_api_hints_the_op_that_turns_it_on( - sqlite_session: Session, mode: AppMode, view: Callable[..., object] -) -> None: - persist(sqlite_session, make_app(enable_api=False, mode=mode)) - subject = account_subject() - - with pytest.raises(ServiceApiDisabled) as raised: - CheckAppApiEnabled().run(subject, make_ctx(sqlite_session, subject=subject, app_id=APP_ID), sqlite_session) - - body = OpenApiErrorFormatter().finalize(raised.value, {}, raised.value.code or 0) - assert (body["status"], body["code"]) == (403, OpenApiErrorCode.SERVICE_API_DISABLED) - [hint] = body["hints"] - assert hint["op"] == op_of(view) - assert hint["input"] == {"app_id": APP_ID, "enabled": True} - ServiceApiPatch.model_validate({k: v for k, v in hint["input"].items() if k != "app_id"}) - - -def test_an_enabled_service_api_passes(sqlite_session: Session) -> None: - persist(sqlite_session, make_app(enable_api=True)) - subject = account_subject() - - CheckAppApiEnabled().run(subject, make_ctx(sqlite_session, subject=subject, app_id=APP_ID), sqlite_session) - - -class TestCheckRBACPermission: +class TestRequirementsCases: @staticmethod - def _requirement() -> CheckRBACPermission: - return CheckRBACPermission(RBACCheck(RBACPermission.APP_VIEW_LAYOUT, PlainApp())) - - def test_skips_a_non_account_caller(self, sqlite_session: Session, config_overrides: Callable[..., None]) -> None: - """No matrix row reaches this: on the routes an SSO token can address, - RBAC permission either does not apply or no permission is declared, so an SSO caller falling - into account-scoped RBAC is invisible there. - """ - config_overrides(RBAC_ENABLED=True) + def test_subject_check_emits_the_wrong_surface_audit(app: Flask, sqlite_session: Session) -> None: subject = sso_subject() - with patch(ENFORCE_RBAC) as enforce: - self._requirement().run(subject, make_ctx(sqlite_session, subject=subject, app_id=APP_ID), sqlite_session) + with app.test_request_context("/openapi/v1/account"): + with patch("controllers.openapi.auth.requirements.emit_wrong_surface") as emit: + requirement = CheckSubject(allowed=[AccountSubject]) + with pytest.raises(Forbidden, match="unsupported_token_type"): + requirement.run(subject, make_ctx(sqlite_session, subject=subject), sqlite_session) - enforce.assert_not_called() + emit.assert_called_once_with( + subject_type="external_sso", + attempted_path="/openapi/v1/account", + client_id=CLIENT_ID, + token_id=TOKEN_ID, + ) - def test_is_inert_wherever_rbac_is_off( - self, app: Flask, sqlite_session: Session, config_overrides: Callable[..., None] + @staticmethod + @pytest.mark.parametrize( + ("status", "denied"), + [ + (LicenseStatus.INACTIVE, True), + (LicenseStatus.LOST, True), + (LicenseStatus.NONE, False), + ], + ) + def test_assert_license_valid_denies_only_dead_licences(status: LicenseStatus, denied: bool) -> None: + with patch(FEATURES, return_value=system_features(license_status=status)): + if denied: + with pytest.raises(Forbidden, match="license_invalid"): + assert_license_valid() + else: + assert_license_valid() + + @staticmethod + def test_an_app_requirement_off_an_app_route_is_a_wiring_bug( + sqlite_session: Session, monkeypatch: pytest.MonkeyPatch ) -> None: - """No matrix row reaches this: every row runs against a stubbed RBAC - backend, so a permission that enforced where RBAC is switched off would still - be admitted there. Standing down is what leaves the `CheckWorkspaceRole` - beside it as the only arm there. - """ - config_overrides(RBAC_ENABLED=False) - persist(sqlite_session, make_app(), make_tenant(), make_account()) - subject = account_subject() - ctx = make_ctx(sqlite_session, subject=subject, app_id=APP_ID) - - with app.test_request_context(f"/openapi/v1/apps/{APP_ID}"): - with patch(ENFORCE_RBAC) as enforce: - self._requirement().run(subject, ctx, sqlite_session) - - enforce.assert_not_called() - - -class TestCheckWorkspaceRole: - admin_only = frozenset({TenantAccountRole.OWNER, TenantAccountRole.ADMIN}) - - def test_skips_a_non_account_caller(self, sqlite_session: Session, monkeypatch: pytest.MonkeyPatch) -> None: - """No matrix row reaches this either — and the workspace role check must not even fetch - the app to decide it. + """App requirements are declared per endpoint, so a route with no `app_id` + can only carry one by mistake. It raises rather than passing quietly: + skipping would turn a misdeclaration into a check that silently never runs. """ monkeypatch.setattr(APP_FETCH, never_reached) - subject = sso_subject() - - CheckWorkspaceRole(self.admin_only).run( - subject, make_ctx(sqlite_session, subject=subject, app_id=APP_ID), sqlite_session - ) - - def test_the_workspace_role_reads_the_role_through_the_loader( - self, app: Flask, sqlite_session: Session, config_overrides: Callable[..., None] - ) -> None: - """The answer is `load_workspace_role`'s, not a direct role read's: an - account that is not `ACTIVE` is a non-member, and hears 404 rather than - the workspace role check's own 403. - - No matrix row reaches this. An `EARLY` membership check pre-empts the - `NORMAL` workspace role check on every shipped route, and the matrix mints only `ACTIVE` - accounts — so this is the only thing standing between the workspace role check and a - refactor that calls `get_account_role_in_tenant` itself, which would - re-admit a banned admin and answer 403 where the surface answers 404. - """ - config_overrides(RBAC_ENABLED=False) - persist( - sqlite_session, - make_app(), - make_tenant(), - make_account(status=AccountStatus.BANNED), - make_membership(TenantAccountRole.ADMIN), - ) - subject = account_subject() - ctx = make_ctx(sqlite_session, subject=subject, app_id=APP_ID) - - with app.test_request_context(f"/openapi/v1/apps/{APP_ID}"): - with pytest.raises(NotFound, match="workspace not found"): - CheckWorkspaceRole(self.admin_only).run(subject, ctx, sqlite_session) - - -class TestCheckAppAccess: - @pytest.fixture(autouse=True) - def _enterprise(self, config_overrides: Callable[..., None]) -> None: - config_overrides(DEPLOYMENT_EDITION=DeploymentEdition.ENTERPRISE) - - def test_no_ops_outside_enterprise( - self, sqlite_session: Session, config_overrides: Callable[..., None], monkeypatch: pytest.MonkeyPatch - ) -> None: - """The named guard for the CE contract: the access-mode service is never - reached at all. Dropping the early return is otherwise only caught - incidentally, by an unrelated test that happens not to stub it. - """ - config_overrides(DEPLOYMENT_EDITION=DeploymentEdition.COMMUNITY) - monkeypatch.setattr(ACCESS_MODE, never_reached) - persist(sqlite_session, make_app()) subject = account_subject() - CheckAppAccess().run(subject, make_ctx(sqlite_session, subject=subject, app_id=APP_ID), sqlite_session) + with pytest.raises(RouteContractError, match="app_id is not a path parameter"): + CheckAppApiEnabled().run(subject, make_ctx(sqlite_session, subject=subject), sqlite_session) + @staticmethod @pytest.mark.parametrize( - ("settings", "failure"), + ("mode", "view"), [ - (webapp_settings("a-mode-this-build-has-never-heard-of"), None), - (None, ValueError("enterprise said no")), + (_requirements_AppMode.CHAT, ServiceApiApi.patch), + (_requirements_AppMode.WORKFLOW, ServiceApiApi.patch), + (_requirements_AppMode.AGENT, AgentServiceApiApi.patch), ], - ids=["unknown mode", "service error"], ) - def test_rejects_an_access_mode_it_could_not_load( - self, sqlite_session: Session, settings: WebAppSettings | None, failure: Exception | None + def test_a_disabled_service_api_hints_the_op_that_turns_it_on( + sqlite_session: Session, mode: _requirements_AppMode, view: Callable[..., object] ) -> None: - persist(sqlite_session, make_app()) + persist(sqlite_session, make_app(enable_api=False, mode=mode)) subject = account_subject() - ctx = make_ctx(sqlite_session, subject=subject, app_id=APP_ID) - with patch(FEATURES, return_value=system_features(webapp_auth=True)): - with patch(ACCESS_MODE, return_value=settings, side_effect=failure): - with pytest.raises(Forbidden, match="app or access mode not loaded"): - CheckAppAccess().run(subject, ctx, sqlite_session) + with pytest.raises(ServiceApiDisabled) as raised: + CheckAppApiEnabled().run(subject, make_ctx(sqlite_session, subject=subject, app_id=APP_ID), sqlite_session) - def test_refuses_a_private_app_when_the_user_cannot_be_resolved(self, sqlite_session: Session) -> None: - persist(sqlite_session, make_app()) - subject = account_subject(account_id=None) - ctx = make_ctx(sqlite_session, subject=subject, app_id=APP_ID) + body = OpenApiErrorFormatter().finalize(raised.value, {}, raised.value.code or 0) + assert (body["status"], body["code"]) == (403, OpenApiErrorCode.SERVICE_API_DISABLED) + [hint] = body["hints"] + assert hint["op"] == op_of(view) + assert hint["input"] == {"app_id": APP_ID, "enabled": True} + ServiceApiPatch.model_validate({k: v for k, v in hint["input"].items() if k != "app_id"}) - with patch(FEATURES, return_value=system_features(webapp_auth=True)): - with patch(ACCESS_MODE, return_value=webapp_settings(WebAppAccessMode.PRIVATE.value)): - with pytest.raises(Forbidden, match="cannot resolve user for private app check"): - CheckAppAccess().run(subject, ctx, sqlite_session) + @staticmethod + def test_an_enabled_service_api_passes(sqlite_session: Session) -> None: + persist(sqlite_session, make_app(enable_api=True)) + subject = account_subject() + + CheckAppApiEnabled().run(subject, make_ctx(sqlite_session, subject=subject, app_id=APP_ID), sqlite_session) + + class TestCheckRBACPermission: + @staticmethod + def _requirement() -> CheckRBACPermission: + return CheckRBACPermission(RBACCheck(RBACPermission.APP_VIEW_LAYOUT, PlainApp())) + + def test_skips_a_non_account_caller( + self, sqlite_session: Session, config_overrides: Callable[..., None] + ) -> None: + """No matrix row reaches this: on the routes an SSO token can address, + RBAC permission either does not apply or no permission is declared, so an SSO caller falling + into account-scoped RBAC is invisible there. + """ + config_overrides(RBAC_ENABLED=True) + subject = sso_subject() + + with patch(ENFORCE_RBAC) as enforce: + self._requirement().run( + subject, make_ctx(sqlite_session, subject=subject, app_id=APP_ID), sqlite_session + ) + + enforce.assert_not_called() + + def test_is_inert_wherever_rbac_is_off( + self, app: Flask, sqlite_session: Session, config_overrides: Callable[..., None] + ) -> None: + """No matrix row reaches this: every row runs against a stubbed RBAC + backend, so a permission that enforced where RBAC is switched off would still + be admitted there. Standing down is what leaves the `CheckWorkspaceRole` + beside it as the only arm there. + """ + config_overrides(RBAC_ENABLED=False) + persist(sqlite_session, make_app(), make_tenant(), make_account()) + subject = account_subject() + ctx = make_ctx(sqlite_session, subject=subject, app_id=APP_ID) + + with app.test_request_context(f"/openapi/v1/apps/{APP_ID}"): + with patch(ENFORCE_RBAC) as enforce: + self._requirement().run(subject, ctx, sqlite_session) + + enforce.assert_not_called() + + class TestCheckWorkspaceRole: + admin_only = frozenset({TenantAccountRole.OWNER, TenantAccountRole.ADMIN}) + + def test_skips_a_non_account_caller(self, sqlite_session: Session, monkeypatch: pytest.MonkeyPatch) -> None: + """No matrix row reaches this either — and the workspace role check must not even fetch + the app to decide it. + """ + monkeypatch.setattr(APP_FETCH, never_reached) + subject = sso_subject() + + CheckWorkspaceRole(self.admin_only).run( + subject, make_ctx(sqlite_session, subject=subject, app_id=APP_ID), sqlite_session + ) + + def test_the_workspace_role_reads_the_role_through_the_loader( + self, app: Flask, sqlite_session: Session, config_overrides: Callable[..., None] + ) -> None: + """The answer is `load_workspace_role`'s, not a direct role read's: an + account that is not `ACTIVE` is a non-member, and hears 404 rather than + the workspace role check's own 403. + + No matrix row reaches this. An `EARLY` membership check pre-empts the + `NORMAL` workspace role check on every shipped route, and the matrix mints only `ACTIVE` + accounts — so this is the only thing standing between the workspace role check and a + refactor that calls `get_account_role_in_tenant` itself, which would + re-admit a banned admin and answer 403 where the surface answers 404. + """ + config_overrides(RBAC_ENABLED=False) + persist( + sqlite_session, + make_app(), + make_tenant(), + make_account(status=AccountStatus.BANNED), + make_membership(TenantAccountRole.ADMIN), + ) + subject = account_subject() + ctx = make_ctx(sqlite_session, subject=subject, app_id=APP_ID) + + with app.test_request_context(f"/openapi/v1/apps/{APP_ID}"): + with pytest.raises(NotFound, match="workspace not found"): + CheckWorkspaceRole(self.admin_only).run(subject, ctx, sqlite_session) + + class TestCheckAppAccess: + @pytest.fixture(autouse=True) + def _enterprise(self, config_overrides: Callable[..., None]) -> None: + config_overrides(DEPLOYMENT_EDITION=DeploymentEdition.ENTERPRISE) + + def test_no_ops_outside_enterprise( + self, sqlite_session: Session, config_overrides: Callable[..., None], monkeypatch: pytest.MonkeyPatch + ) -> None: + """The named guard for the CE contract: the access-mode service is never + reached at all. Dropping the early return is otherwise only caught + incidentally, by an unrelated test that happens not to stub it. + """ + config_overrides(DEPLOYMENT_EDITION=DeploymentEdition.COMMUNITY) + monkeypatch.setattr(ACCESS_MODE, never_reached) + persist(sqlite_session, make_app()) + subject = account_subject() + + CheckAppAccess().run(subject, make_ctx(sqlite_session, subject=subject, app_id=APP_ID), sqlite_session) + + @pytest.mark.parametrize( + ("settings", "failure"), + [ + (webapp_settings("a-mode-this-build-has-never-heard-of"), None), + (None, ValueError("enterprise said no")), + ], + ids=["unknown mode", "service error"], + ) + def test_rejects_an_access_mode_it_could_not_load( + self, sqlite_session: Session, settings: WebAppSettings | None, failure: Exception | None + ) -> None: + persist(sqlite_session, make_app()) + subject = account_subject() + ctx = make_ctx(sqlite_session, subject=subject, app_id=APP_ID) + + with patch(FEATURES, return_value=system_features(webapp_auth=True)): + with patch(ACCESS_MODE, return_value=settings, side_effect=failure): + with pytest.raises(Forbidden, match="app or access mode not loaded"): + CheckAppAccess().run(subject, ctx, sqlite_session) + + def test_refuses_a_private_app_when_the_user_cannot_be_resolved(self, sqlite_session: Session) -> None: + persist(sqlite_session, make_app()) + subject = account_subject(account_id=None) + ctx = make_ctx(sqlite_session, subject=subject, app_id=APP_ID) + + with patch(FEATURES, return_value=system_features(webapp_auth=True)): + with patch(ACCESS_MODE, return_value=webapp_settings(_requirements_WebAppAccessMode.PRIVATE.value)): + with pytest.raises(Forbidden, match="cannot resolve user for private app check"): + CheckAppAccess().run(subject, ctx, sqlite_session) + + +"""The /openapi/v1 allow/deny matrix, and a snapshot of the generated document. + +The matrix is derived from what actually runs, not from what a route declares: a +route's own requirements are merged with the fixed ones its subject's pipeline +carries, a `CheckWorkspaceRole` applies only where RBAC is off and a `CheckRBACPermission` only where it is on, +and `CheckAppAccess` applies the private-app check whether +or not `webapp_auth.enabled` gates the ACL. Reading declarations alone under-counts +every one of those. + +Requests run through the real Flask blueprint with the real router, the real +pipeline and real database rows. The only substitutions are at process seams the +router already treats as pluggable: the bound `BearerAuthenticator` (so a token's +subject kind and scope set are chosen per row), the enterprise HTTP clients, and +the flask-login mount. The RBAC backend and the enterprise web-app service are +stubbed per row, so those rows pin how *this* code handles each answer they can +give — never what the real services decide. + +Cases are never gated on whether a route declares a check. A row that asserts +`ADMIT` where RBAC, a workspace role check or a licence gate does not fire is pinning the +check's *absence*: losing a check is caught by the deny rows, and gaining one is +caught by these. `files.upload` is the route the design says must never acquire an +RBAC permission, and `files.upload-rbac_on_denied` is the row that says so. + +Admission is observed as HTTP 418: a `user_logged_in` receiver raises `ImATeapot` +at the moment the pipeline mounts the caller, which is after every requirement has +passed and before the view body runs. That keeps a row from depending on whether a +view could complete (several stream SSE or invoke a workflow). One route mounts no +caller — `permitted_external.list` resolves no end user because it carries no +`app_id` — so its admission row says so and asserts the view's own 200 instead. + +Two answers in the table are worth reading twice, because they are what the code +does rather than what a route's declaration suggests: + +* `workspaces.describe` declares no membership requirement, so a non-member is + *admitted* by auth and refused by the view's own lookup. `workspaces.switch`, + one path segment away, declares one and is refused by auth. +* No token the shipped registry mints can fail `CheckScope` — `dfoa_` carries + `Scope.FULL` and `dfoe_` carries exactly the two scopes its routes ask for — so + the scope rows mint from a deliberately narrowed registry. Those rows are not + dead weight and must not be simplified away: `CheckScope` is live code on every + request, and the day a narrower token kind is minted it is the only thing + standing between that token and every route it was not scoped for. + +`Case.NON_MEMBER_AND_INSUFFICIENT_SCOPE` is the one case that combines a real +membership failure with a real scope failure on the same request. Neither +`INSUFFICIENT_SCOPE` nor `NON_MEMBER` alone proves which one a doubly-broken request +hears; membership runs EARLY and scope NORMAL, so it hears membership. + +Routes that declare the same requirement tuple answer every case the same way, +so `MATRIX` runs the cases on one representative per tuple and `DECLARED` pins +which tuple each route carries. A route that changes its tuple fails +the wiring test; a tuple that changes its answers fails its representative's rows. +""" + + +ADMITTED = 418 + + +class Trait(StrEnum): + """Structural facts about a route, used only to decide which cases can reach it. + + Every member here is a fact about the *shape* of the request — is there an + `app_id` in the path, can a `dfoe_` token address this route at all. None is + read off a route's declared requirements, deliberately: gating a + case on a declared check means a route that *gains* that check during migration + simply loses the row, and gaining a check is the hazard this plan cares about + most. Cases like `RBAC_ON_DENIED` therefore run on every route an account + caller can reach, and assert `ADMIT` where the check does not exist today. + """ + + APP_SCOPED = auto() + ACCOUNT_PRIMARY = auto() + EXTERNAL_REACHABLE = auto() + ENTERPRISE_ONLY = auto() + + +class Case(StrEnum): + NO_BEARER = auto() + MEMBER = auto() + WRONG_SUBJECT = auto() + INSUFFICIENT_SCOPE = auto() + NON_MEMBER = auto() + NON_MEMBER_AND_INSUFFICIENT_SCOPE = auto() + LOW_ROLE = auto() + APP_API_DISABLED = auto() + UNKNOWN_APP = auto() + FOREIGN_WORKSPACE_QUERY = auto() + EDITION_NOT_ENTERPRISE = auto() + LICENSE_INVALID = auto() + EE_LICENSE_INVALID = auto() + EE_ACCOUNT_PUBLIC = auto() + EE_ACCOUNT_SSO_VERIFIED = auto() + EE_ACCOUNT_PRIVATE_ALL = auto() + EE_ACCOUNT_PRIVATE_PERMITTED = auto() + EE_ACCOUNT_PRIVATE_REFUSED = auto() + EE_ACCOUNT_PRIVATE_REFUSED_WEBAPP_AUTH_OFF = auto() + EE_ACCOUNT_MODE_UNRESOLVED = auto() + EE_EXTERNAL_PUBLIC = auto() + EE_EXTERNAL_SSO_VERIFIED = auto() + EE_EXTERNAL_PRIVATE_ALL = auto() + EE_EXTERNAL_PRIVATE = auto() + EE_EXTERNAL_PRIVATE_REFUSED_WEBAPP_AUTH_OFF = auto() + EE_EXTERNAL_MODE_UNRESOLVED = auto() + RBAC_ON_LOW_ROLE = auto() + RBAC_ON_DENIED = auto() + + +class Bearer(StrEnum): + NONE = auto() + ACCOUNT_MEMBER = auto() + ACCOUNT_LOW_ROLE = auto() + ACCOUNT_OUTSIDER = auto() + EXTERNAL = auto() + OTHER_SUBJECT = auto() + PRIMARY = auto() + + +@dataclass(frozen=True, slots=True) +class Route: + id: str + method: str + path: str + traits: frozenset[Trait] + query: str = "" + + +@dataclass(frozen=True, slots=True) +class Expect: + """`message` is the canonical ErrorBody message; `None` means it is not pinned.""" + + status: int + message: str | None = None + note: str = "" + + +@dataclass(frozen=True, slots=True) +class Scenario: + """The world a case runs in. `edition=None` means the route's native edition.""" + + bearer: Bearer + edition: DeploymentEdition | None = None + narrow_scopes: bool = False + app_api_enabled: bool = True + unknown_app: bool = False + foreign_workspace_query: bool = False + webapp_auth: bool = False + access_mode: _auth_matrix_WebAppAccessMode | None = _auth_matrix_WebAppAccessMode.PUBLIC + private_app_permitted: bool = True + license_status: LicenseStatus = LicenseStatus.ACTIVE + rbac_enabled: bool = False + rbac_allows: bool = True + + +ADMIT = Expect(ADMITTED) +DENY_NO_BEARER = Expect(401, "bearer required") +DENY_WRONG_SUBJECT = Expect(403, "unsupported_token_type") +DENY_SSO_NEEDS_EE = Expect(403, "external_sso_requires_ee") +DENY_SCOPE = Expect(403, "insufficient_scope") +DENY_NON_MEMBER = Expect(404, "workspace not found") +DENY_ROLE = Expect(403, "insufficient workspace role") +DENY_API_DISABLED = Expect(403, ServiceApiDisabled.description) +DENY_UNKNOWN_APP = Expect(404, "app not found") +DENY_ACCESS_MODE = Expect(403, "subject_not_allowed_for_access_mode") +DENY_MODE_UNRESOLVED = Expect(403, "app or access mode not loaded") +DENY_PRIVATE_APP = Expect(403, "user_not_allowed_for_private_app") +DENY_EDITION = Expect(404, note="endpoint-level edition gate, raised before the bearer is read") +DENY_LICENSE = Expect(403, "license_invalid") +DENY_RBAC = Expect(403, note="bare werkzeug Forbidden from enforce_rbac_checks") +ADMIT_NO_RBAC_PERMISSION = Expect( + ADMITTED, + note=( + "pins a check this route does NOT have: RBAC is on and denying, and the route still " + "admits because it declares no scene. Attaching one during migration turns this 403. " + "files.upload is the route the design says must never gain a scene" + ), +) +ADMIT_NO_WORKSPACE_ROLE = Expect( + ADMITTED, + note="pins a check this route does NOT have: RBAC is on and a NORMAL member is still admitted", +) +ADMIT_NO_LICENCE_GATE = Expect( + ADMITTED, + note="pins a check this route does NOT have: the licence is dead and the route still admits", +) +ADMIT_NO_MOUNT = Expect( + 200, + note="external subject mounts no caller on an app-less route, so admission shows as the view's own 200", +) + + +_RUN_TRAITS = frozenset({Trait.ACCOUNT_PRIMARY, Trait.APP_SCOPED, Trait.EXTERNAL_REACHABLE}) + +ROUTES: tuple[Route, ...] = ( + Route( + "knowledge_fs.ls", + "GET", + "/workspaces/{workspace_id}/knowledge-fs/knowledge-spaces/{knowledge_space_id}/fs:ls", + frozenset({Trait.ACCOUNT_PRIMARY}), + ), + Route( + "knowledge_fs.tree", + "GET", + "/workspaces/{workspace_id}/knowledge-fs/knowledge-spaces/{knowledge_space_id}/fs:tree", + frozenset({Trait.ACCOUNT_PRIMARY}), + ), + Route( + "knowledge_fs.grep", + "GET", + "/workspaces/{workspace_id}/knowledge-fs/knowledge-spaces/{knowledge_space_id}/fs:grep", + frozenset({Trait.ACCOUNT_PRIMARY}), + ), + Route( + "knowledge_fs.find", + "GET", + "/workspaces/{workspace_id}/knowledge-fs/knowledge-spaces/{knowledge_space_id}/fs:find", + frozenset({Trait.ACCOUNT_PRIMARY}), + ), + Route( + "knowledge_fs.diff", + "POST", + "/workspaces/{workspace_id}/knowledge-fs/knowledge-spaces/{knowledge_space_id}/fs:diff", + frozenset({Trait.ACCOUNT_PRIMARY}), + ), + Route( + "knowledge_fs.cat", + "GET", + "/workspaces/{workspace_id}/knowledge-fs/knowledge-spaces/{knowledge_space_id}/fs:cat", + frozenset({Trait.ACCOUNT_PRIMARY}), + ), + Route( + "knowledge_fs.stat", + "GET", + "/workspaces/{workspace_id}/knowledge-fs/knowledge-spaces/{knowledge_space_id}/fs:stat", + frozenset({Trait.ACCOUNT_PRIMARY}), + ), + Route("describe.account", "GET", "/account", frozenset({Trait.ACCOUNT_PRIMARY})), + Route("account.sessions.revoke_self", "DELETE", "/account/sessions/self", frozenset({Trait.ACCOUNT_PRIMARY})), + Route("get.account.session", "GET", "/account/sessions", frozenset({Trait.ACCOUNT_PRIMARY})), + Route( + "account.sessions.revoke_one", + "DELETE", + "/account/sessions/{session_id}", + frozenset({Trait.ACCOUNT_PRIMARY}), + ), + Route( + "apps.describe", + "GET", + "/apps/{app_id}", + frozenset({Trait.ACCOUNT_PRIMARY, Trait.APP_SCOPED}), + ), + Route("apps.list", "GET", "/apps", frozenset({Trait.ACCOUNT_PRIMARY}), query="workspace_id={workspace_id}"), + Route("workspaces.list", "GET", "/workspaces", frozenset({Trait.ACCOUNT_PRIMARY})), + Route("workspaces.describe", "GET", "/workspaces/{workspace_id}", frozenset({Trait.ACCOUNT_PRIMARY})), + Route("workspaces.members.list", "GET", "/workspaces/{workspace_id}/members", frozenset({Trait.ACCOUNT_PRIMARY})), + Route( + "workspaces.members.invite", + "POST", + "/workspaces/{workspace_id}/members", + frozenset({Trait.ACCOUNT_PRIMARY}), + ), + Route( + "workspaces.members.remove", + "DELETE", + "/workspaces/{workspace_id}/members/{member_id}", + frozenset({Trait.ACCOUNT_PRIMARY}), + ), + Route( + "workspaces.members.update_role", + "PATCH", + "/workspaces/{workspace_id}/members/{member_id}", + frozenset({Trait.ACCOUNT_PRIMARY}), + ), + Route( + "plugins.marketplace", + "GET", + "/workspaces/{workspace_id}/marketplace/plugins", + frozenset({Trait.ACCOUNT_PRIMARY}), + ), + Route("plugins.list", "GET", "/workspaces/{workspace_id}/plugins", frozenset({Trait.ACCOUNT_PRIMARY})), + Route("plugins.install", "POST", "/workspaces/{workspace_id}/plugins:install", frozenset({Trait.ACCOUNT_PRIMARY})), + Route( + "plugins.task", + "GET", + "/workspaces/{workspace_id}/plugin-tasks/{task_id}", + frozenset({Trait.ACCOUNT_PRIMARY}), + ), + Route( + "plugins.delete", + "DELETE", + "/workspaces/{workspace_id}/plugins/{plugin_id}", + frozenset({Trait.ACCOUNT_PRIMARY}), + ), + Route( + "model_providers.describe", + "GET", + "/workspaces/{workspace_id}/model-providers/{provider}", + frozenset({Trait.ACCOUNT_PRIMARY}), + ), + Route( + "model_providers.credential.create", + "POST", + "/workspaces/{workspace_id}/model-providers/{provider}/credentials", + frozenset({Trait.ACCOUNT_PRIMARY}), + ), + Route( + "model_providers.credential.set", + "PATCH", + "/workspaces/{workspace_id}/model-providers/{provider}/credentials/{credential_id}", + frozenset({Trait.ACCOUNT_PRIMARY}), + ), + Route("tools.list", "GET", "/workspaces/{workspace_id}/tools", frozenset({Trait.ACCOUNT_PRIMARY})), + Route( + "tool_providers.describe", + "GET", + "/workspaces/{workspace_id}/tool-providers/{provider}", + frozenset({Trait.ACCOUNT_PRIMARY}), + ), + Route( + "tool_providers.credential.create", + "POST", + "/workspaces/{workspace_id}/tool-providers/{provider}/credentials", + frozenset({Trait.ACCOUNT_PRIMARY}), + ), + Route( + "tool_providers.credential.set", + "PATCH", + "/workspaces/{workspace_id}/tool-providers/{provider}/credentials/{credential_id}", + frozenset({Trait.ACCOUNT_PRIMARY}), + ), + Route( + "knowledge_bases.list", + "GET", + "/workspaces/{workspace_id}/knowledge-bases", + frozenset({Trait.ACCOUNT_PRIMARY}), + ), + Route( + "models.list", + "GET", + "/workspaces/{workspace_id}/models", + frozenset({Trait.ACCOUNT_PRIMARY}), + ), + Route( + "models.credential.create", + "POST", + "/workspaces/{workspace_id}/model-providers/{provider}/models/credentials", + frozenset({Trait.ACCOUNT_PRIMARY}), + ), + Route( + "models.credential.set", + "PATCH", + "/workspaces/{workspace_id}/model-providers/{provider}/models/credentials/{credential_id}", + frozenset({Trait.ACCOUNT_PRIMARY}), + ), + Route( + "app_dsl.import", + "POST", + "/workspaces/{workspace_id}/apps/imports", + frozenset({Trait.ACCOUNT_PRIMARY}), + ), + Route( + "app_dsl.check", + "POST", + "/workspaces/{workspace_id}/apps/imports:check", + frozenset({Trait.ACCOUNT_PRIMARY}), + ), + Route( + "app_dsl.import_confirm", + "POST", + "/workspaces/{workspace_id}/apps/imports/{import_id}:confirm", + frozenset({Trait.ACCOUNT_PRIMARY}), + ), + Route( + "app_dsl.export", + "GET", + "/apps/{app_id}/dsl", + frozenset({Trait.ACCOUNT_PRIMARY, Trait.APP_SCOPED}), + ), + Route( + "app_dsl.check_dependencies", + "GET", + "/apps/{app_id}/dependencies:check", + frozenset({Trait.ACCOUNT_PRIMARY, Trait.APP_SCOPED}), + ), + Route("app_run.workflow", "POST", "/apps/{app_id}/workflow:run", _RUN_TRAITS), + Route("app_run.chat", "POST", "/apps/{app_id}/chat:run", _RUN_TRAITS), + Route("app_run.advanced_chat", "POST", "/apps/{app_id}/advanced-chat:run", _RUN_TRAITS), + Route("app_run.completion", "POST", "/apps/{app_id}/completion:run", _RUN_TRAITS), + Route( + "app_run.draft.workflow", + "POST", + "/apps/{app_id}/draft/workflow:run", + frozenset({Trait.ACCOUNT_PRIMARY, Trait.APP_SCOPED}), + ), + Route( + "app_run.draft.advanced_chat", + "POST", + "/apps/{app_id}/draft/advanced-chat:run", + frozenset({Trait.ACCOUNT_PRIMARY, Trait.APP_SCOPED}), + ), + Route("app_workflow.run.list", "GET", "/apps/{app_id}/runs", frozenset({Trait.ACCOUNT_PRIMARY, Trait.APP_SCOPED})), + Route( + "app_workflow.run.describe", + "GET", + "/apps/{app_id}/runs/{run_id}", + frozenset({Trait.ACCOUNT_PRIMARY, Trait.APP_SCOPED}), + ), + Route( + "app_workflow.run.nodes", + "GET", + "/apps/{app_id}/runs/{run_id}/nodes", + frozenset({Trait.ACCOUNT_PRIMARY, Trait.APP_SCOPED}), + ), + Route( + "app_workflow.publish", "POST", "/apps/{app_id}:publish", frozenset({Trait.ACCOUNT_PRIMARY, Trait.APP_SCOPED}) + ), + Route( + "app_workflow.version.list", + "GET", + "/apps/{app_id}/versions", + frozenset({Trait.ACCOUNT_PRIMARY, Trait.APP_SCOPED}), + ), + Route( + "app_workflow.version.restore", + "POST", + "/apps/{app_id}/versions/{version_id}:restore", + frozenset({Trait.ACCOUNT_PRIMARY, Trait.APP_SCOPED}), + ), + Route( + "app_workflow.release_check", + "GET", + "/apps/{app_id}/release:check", + frozenset({Trait.ACCOUNT_PRIMARY, Trait.APP_SCOPED}), + ), + Route("app_workflow.env.list", "GET", "/apps/{app_id}/env", frozenset({Trait.ACCOUNT_PRIMARY, Trait.APP_SCOPED})), + Route( + "app_workflow.env.set", + "PUT", + "/apps/{app_id}/env/{env_id}", + frozenset({Trait.ACCOUNT_PRIMARY, Trait.APP_SCOPED}), + ), + Route( + "app_workflow.env.delete", + "DELETE", + "/apps/{app_id}/env/{env_id}", + frozenset({Trait.ACCOUNT_PRIMARY, Trait.APP_SCOPED}), + ), + Route("node_types.list", "GET", "/node-types", frozenset({Trait.ACCOUNT_PRIMARY, Trait.EXTERNAL_REACHABLE})), + Route( + "node_types.describe", + "GET", + "/node-types/{node_type}", + frozenset({Trait.ACCOUNT_PRIMARY, Trait.EXTERNAL_REACHABLE}), + ), + Route( + "app_create.workflow", "POST", "/workspaces/{workspace_id}/apps/workflow", frozenset({Trait.ACCOUNT_PRIMARY}) + ), + Route( + "app_create.advanced_chat", + "POST", + "/workspaces/{workspace_id}/apps/advanced-chat", + frozenset({Trait.ACCOUNT_PRIMARY}), + ), + Route( + "app_workflow.node_run.workflow", + "POST", + "/apps/{app_id}/draft/workflow/nodes/{node_id}:run", + frozenset({Trait.ACCOUNT_PRIMARY, Trait.APP_SCOPED}), + ), + Route( + "app_workflow.node_run.advanced_chat", + "POST", + "/apps/{app_id}/draft/advanced-chat/nodes/{node_id}:run", + frozenset({Trait.ACCOUNT_PRIMARY, Trait.APP_SCOPED}), + ), + Route( + "app_info.describe.workflow", + "GET", + "/apps/{app_id}/app-info/workflow", + frozenset({Trait.ACCOUNT_PRIMARY, Trait.APP_SCOPED}), + ), + Route( + "app_info.set.workflow", + "PATCH", + "/apps/{app_id}/app-info/workflow", + frozenset({Trait.ACCOUNT_PRIMARY, Trait.APP_SCOPED}), + ), + Route( + "service_api.describe", + "GET", + "/apps/{app_id}/service-api", + frozenset({Trait.ACCOUNT_PRIMARY, Trait.APP_SCOPED}), + ), + Route( + "service_api.set", + "PATCH", + "/apps/{app_id}/service-api", + frozenset({Trait.ACCOUNT_PRIMARY, Trait.APP_SCOPED}), + ), + Route( + "app_info.describe.advanced_chat", + "GET", + "/apps/{app_id}/app-info/advanced-chat", + frozenset({Trait.ACCOUNT_PRIMARY, Trait.APP_SCOPED}), + ), + Route( + "app_info.set.advanced_chat", + "PATCH", + "/apps/{app_id}/app-info/advanced-chat", + frozenset({Trait.ACCOUNT_PRIMARY, Trait.APP_SCOPED}), + ), + Route( + "app_info.describe.chat", + "GET", + "/apps/{app_id}/app-info/chat", + frozenset({Trait.ACCOUNT_PRIMARY, Trait.APP_SCOPED}), + ), + Route( + "app_info.set.chat", + "PATCH", + "/apps/{app_id}/app-info/chat", + frozenset({Trait.ACCOUNT_PRIMARY, Trait.APP_SCOPED}), + ), + Route( + "app_info.describe.agent_chat", + "GET", + "/apps/{app_id}/app-info/agent-chat", + frozenset({Trait.ACCOUNT_PRIMARY, Trait.APP_SCOPED}), + ), + Route( + "app_info.set.agent_chat", + "PATCH", + "/apps/{app_id}/app-info/agent-chat", + frozenset({Trait.ACCOUNT_PRIMARY, Trait.APP_SCOPED}), + ), + Route( + "app_info.describe.completion", + "GET", + "/apps/{app_id}/app-info/completion", + frozenset({Trait.ACCOUNT_PRIMARY, Trait.APP_SCOPED}), + ), + Route( + "app_info.set.completion", + "PATCH", + "/apps/{app_id}/app-info/completion", + frozenset({Trait.ACCOUNT_PRIMARY, Trait.APP_SCOPED}), + ), + Route( + "app_info.describe.agent", + "GET", + "/apps/{app_id}/app-info/agent", + frozenset({Trait.ACCOUNT_PRIMARY, Trait.APP_SCOPED}), + ), + Route( + "app_info.set.agent", + "PATCH", + "/apps/{app_id}/app-info/agent", + frozenset({Trait.ACCOUNT_PRIMARY, Trait.APP_SCOPED}), + ), + Route( + "service_api.describe.agent", + "GET", + "/apps/{app_id}/service-api/agent", + frozenset({Trait.ACCOUNT_PRIMARY, Trait.APP_SCOPED}), + ), + Route( + "service_api.set.agent", + "PATCH", + "/apps/{app_id}/service-api/agent", + frozenset({Trait.ACCOUNT_PRIMARY, Trait.APP_SCOPED}), + ), + Route( + "webapp.describe.workflow", + "GET", + "/apps/{app_id}/webapp/workflow", + frozenset({Trait.ACCOUNT_PRIMARY, Trait.APP_SCOPED}), + ), + Route( + "webapp.set.workflow", + "PATCH", + "/apps/{app_id}/webapp/workflow", + frozenset({Trait.ACCOUNT_PRIMARY, Trait.APP_SCOPED}), + ), + Route( + "webapp.reset", + "POST", + "/apps/{app_id}/webapp:reset", + frozenset({Trait.ACCOUNT_PRIMARY, Trait.APP_SCOPED}), + ), + Route( + "webapp_access.describe", + "GET", + "/apps/{app_id}/webapp-access", + frozenset({Trait.ACCOUNT_PRIMARY, Trait.APP_SCOPED}), + ), + Route( + "webapp_access.set", + "PUT", + "/apps/{app_id}/webapp-access", + frozenset({Trait.ACCOUNT_PRIMARY, Trait.APP_SCOPED}), + ), + Route( + "webapp.describe.advanced_chat", + "GET", + "/apps/{app_id}/webapp/advanced-chat", + frozenset({Trait.ACCOUNT_PRIMARY, Trait.APP_SCOPED}), + ), + Route( + "webapp.set.advanced_chat", + "PATCH", + "/apps/{app_id}/webapp/advanced-chat", + frozenset({Trait.ACCOUNT_PRIMARY, Trait.APP_SCOPED}), + ), + Route( + "webapp.describe.chat", + "GET", + "/apps/{app_id}/webapp/chat", + frozenset({Trait.ACCOUNT_PRIMARY, Trait.APP_SCOPED}), + ), + Route( + "webapp.set.chat", + "PATCH", + "/apps/{app_id}/webapp/chat", + frozenset({Trait.ACCOUNT_PRIMARY, Trait.APP_SCOPED}), + ), + Route( + "webapp.describe.agent_chat", + "GET", + "/apps/{app_id}/webapp/agent-chat", + frozenset({Trait.ACCOUNT_PRIMARY, Trait.APP_SCOPED}), + ), + Route( + "webapp.set.agent_chat", + "PATCH", + "/apps/{app_id}/webapp/agent-chat", + frozenset({Trait.ACCOUNT_PRIMARY, Trait.APP_SCOPED}), + ), + Route( + "webapp.describe.completion", + "GET", + "/apps/{app_id}/webapp/completion", + frozenset({Trait.ACCOUNT_PRIMARY, Trait.APP_SCOPED}), + ), + Route( + "webapp.set.completion", + "PATCH", + "/apps/{app_id}/webapp/completion", + frozenset({Trait.ACCOUNT_PRIMARY, Trait.APP_SCOPED}), + ), + Route( + "webapp.describe.agent", + "GET", + "/apps/{app_id}/webapp/agent", + frozenset({Trait.ACCOUNT_PRIMARY, Trait.APP_SCOPED}), + ), + Route( + "webapp.set.agent", + "PATCH", + "/apps/{app_id}/webapp/agent", + frozenset({Trait.ACCOUNT_PRIMARY, Trait.APP_SCOPED}), + ), + Route( + "webapp.reset.agent", + "POST", + "/apps/{app_id}/webapp/agent:reset", + frozenset({Trait.ACCOUNT_PRIMARY, Trait.APP_SCOPED}), + ), + Route( + "webapp_access.describe.agent", + "GET", + "/apps/{app_id}/webapp-access/agent", + frozenset({Trait.ACCOUNT_PRIMARY, Trait.APP_SCOPED}), + ), + Route( + "webapp_access.set.agent", + "PUT", + "/apps/{app_id}/webapp-access/agent", + frozenset({Trait.ACCOUNT_PRIMARY, Trait.APP_SCOPED}), + ), + Route( + "app_run.stop", + "POST", + "/apps/{app_id}/tasks/{task_id}:stop", + frozenset({Trait.ACCOUNT_PRIMARY, Trait.APP_SCOPED, Trait.EXTERNAL_REACHABLE}), + ), + Route( + "files.upload", + "POST", + "/apps/{app_id}/files", + frozenset({Trait.ACCOUNT_PRIMARY, Trait.APP_SCOPED, Trait.EXTERNAL_REACHABLE}), + ), + Route( + "human_input_form.get", + "GET", + "/apps/{app_id}/human-input-forms/{form_token}", + frozenset({Trait.ACCOUNT_PRIMARY, Trait.APP_SCOPED, Trait.EXTERNAL_REACHABLE}), + ), + Route( + "human_input_form.submit", + "POST", + "/apps/{app_id}/human-input-forms/{form_token}:submit", + frozenset({Trait.ACCOUNT_PRIMARY, Trait.APP_SCOPED, Trait.EXTERNAL_REACHABLE}), + ), + Route( + "workflow_events.stream", + "GET", + "/apps/{app_id}/tasks/{task_id}/events", + frozenset({Trait.ACCOUNT_PRIMARY, Trait.APP_SCOPED, Trait.EXTERNAL_REACHABLE}), + ), + Route( + "permitted_external.list", + "GET", + "/permitted-external-apps", + frozenset({Trait.EXTERNAL_REACHABLE, Trait.ENTERPRISE_ONLY}), + ), + Route( + "permitted_external.describe", + "GET", + "/permitted-external-apps/{app_id}", + frozenset({Trait.APP_SCOPED, Trait.EXTERNAL_REACHABLE, Trait.ENTERPRISE_ONLY}), + ), +) + + +CASE_REQUIRES: dict[Case, frozenset[Trait]] = { + Case.NO_BEARER: frozenset(), + Case.MEMBER: frozenset(), + Case.WRONG_SUBJECT: frozenset(), + Case.INSUFFICIENT_SCOPE: frozenset(), + Case.NON_MEMBER: frozenset({Trait.ACCOUNT_PRIMARY}), + # frozenset() rather than {ACCOUNT_PRIMARY}: this one also has to reach the + # SSO-only routes, to pin that they answer on the subject axis, not this one. + Case.NON_MEMBER_AND_INSUFFICIENT_SCOPE: frozenset(), + Case.LOW_ROLE: frozenset({Trait.ACCOUNT_PRIMARY}), + Case.APP_API_DISABLED: frozenset({Trait.APP_SCOPED}), + Case.UNKNOWN_APP: frozenset({Trait.APP_SCOPED}), + Case.FOREIGN_WORKSPACE_QUERY: frozenset({Trait.APP_SCOPED}), + Case.EDITION_NOT_ENTERPRISE: frozenset({Trait.ENTERPRISE_ONLY}), + Case.LICENSE_INVALID: frozenset(), + Case.EE_LICENSE_INVALID: frozenset(), + Case.EE_ACCOUNT_PUBLIC: frozenset({Trait.APP_SCOPED, Trait.ACCOUNT_PRIMARY}), + Case.EE_ACCOUNT_SSO_VERIFIED: frozenset({Trait.APP_SCOPED, Trait.ACCOUNT_PRIMARY}), + Case.EE_ACCOUNT_PRIVATE_ALL: frozenset({Trait.APP_SCOPED, Trait.ACCOUNT_PRIMARY}), + Case.EE_ACCOUNT_PRIVATE_PERMITTED: frozenset({Trait.APP_SCOPED, Trait.ACCOUNT_PRIMARY}), + Case.EE_ACCOUNT_PRIVATE_REFUSED: frozenset({Trait.APP_SCOPED, Trait.ACCOUNT_PRIMARY}), + Case.EE_ACCOUNT_PRIVATE_REFUSED_WEBAPP_AUTH_OFF: frozenset({Trait.APP_SCOPED, Trait.ACCOUNT_PRIMARY}), + Case.EE_ACCOUNT_MODE_UNRESOLVED: frozenset({Trait.APP_SCOPED, Trait.ACCOUNT_PRIMARY}), + Case.EE_EXTERNAL_PUBLIC: frozenset({Trait.APP_SCOPED, Trait.EXTERNAL_REACHABLE}), + Case.EE_EXTERNAL_SSO_VERIFIED: frozenset({Trait.APP_SCOPED, Trait.EXTERNAL_REACHABLE}), + Case.EE_EXTERNAL_PRIVATE_ALL: frozenset({Trait.APP_SCOPED, Trait.EXTERNAL_REACHABLE}), + Case.EE_EXTERNAL_PRIVATE: frozenset({Trait.APP_SCOPED, Trait.EXTERNAL_REACHABLE}), + Case.EE_EXTERNAL_PRIVATE_REFUSED_WEBAPP_AUTH_OFF: frozenset({Trait.APP_SCOPED, Trait.EXTERNAL_REACHABLE}), + Case.EE_EXTERNAL_MODE_UNRESOLVED: frozenset({Trait.APP_SCOPED, Trait.EXTERNAL_REACHABLE}), + Case.RBAC_ON_LOW_ROLE: frozenset({Trait.ACCOUNT_PRIMARY}), + Case.RBAC_ON_DENIED: frozenset({Trait.ACCOUNT_PRIMARY}), +} + + +SCENARIOS: dict[Case, Scenario] = { + Case.NO_BEARER: Scenario(bearer=Bearer.NONE), + Case.MEMBER: Scenario(bearer=Bearer.PRIMARY), + Case.WRONG_SUBJECT: Scenario(bearer=Bearer.OTHER_SUBJECT), + Case.INSUFFICIENT_SCOPE: Scenario(bearer=Bearer.PRIMARY, narrow_scopes=True), + Case.NON_MEMBER: Scenario(bearer=Bearer.ACCOUNT_OUTSIDER), + Case.NON_MEMBER_AND_INSUFFICIENT_SCOPE: Scenario(bearer=Bearer.ACCOUNT_OUTSIDER, narrow_scopes=True), + Case.LOW_ROLE: Scenario(bearer=Bearer.ACCOUNT_LOW_ROLE), + Case.APP_API_DISABLED: Scenario(bearer=Bearer.PRIMARY, app_api_enabled=False), + Case.UNKNOWN_APP: Scenario(bearer=Bearer.PRIMARY, unknown_app=True), + Case.FOREIGN_WORKSPACE_QUERY: Scenario(bearer=Bearer.PRIMARY, foreign_workspace_query=True), + Case.EDITION_NOT_ENTERPRISE: Scenario(bearer=Bearer.PRIMARY, edition=DeploymentEdition.COMMUNITY), + Case.LICENSE_INVALID: Scenario(bearer=Bearer.PRIMARY, license_status=LicenseStatus.EXPIRED), + Case.EE_LICENSE_INVALID: Scenario( + bearer=Bearer.PRIMARY, edition=DeploymentEdition.ENTERPRISE, license_status=LicenseStatus.EXPIRED + ), + Case.EE_ACCOUNT_PUBLIC: Scenario( + bearer=Bearer.ACCOUNT_MEMBER, + edition=DeploymentEdition.ENTERPRISE, + webapp_auth=True, + access_mode=_auth_matrix_WebAppAccessMode.PUBLIC, + ), + Case.EE_ACCOUNT_SSO_VERIFIED: Scenario( + bearer=Bearer.ACCOUNT_MEMBER, + edition=DeploymentEdition.ENTERPRISE, + webapp_auth=True, + access_mode=_auth_matrix_WebAppAccessMode.SSO_VERIFIED, + ), + Case.EE_ACCOUNT_PRIVATE_ALL: Scenario( + bearer=Bearer.ACCOUNT_MEMBER, + edition=DeploymentEdition.ENTERPRISE, + webapp_auth=True, + access_mode=_auth_matrix_WebAppAccessMode.PRIVATE_ALL, + ), + Case.EE_ACCOUNT_PRIVATE_PERMITTED: Scenario( + bearer=Bearer.ACCOUNT_MEMBER, + edition=DeploymentEdition.ENTERPRISE, + webapp_auth=True, + access_mode=_auth_matrix_WebAppAccessMode.PRIVATE, + private_app_permitted=True, + ), + Case.EE_ACCOUNT_PRIVATE_REFUSED: Scenario( + bearer=Bearer.ACCOUNT_MEMBER, + edition=DeploymentEdition.ENTERPRISE, + webapp_auth=True, + access_mode=_auth_matrix_WebAppAccessMode.PRIVATE, + private_app_permitted=False, + ), + Case.EE_ACCOUNT_PRIVATE_REFUSED_WEBAPP_AUTH_OFF: Scenario( + bearer=Bearer.ACCOUNT_MEMBER, + edition=DeploymentEdition.ENTERPRISE, + webapp_auth=False, + access_mode=_auth_matrix_WebAppAccessMode.PRIVATE, + private_app_permitted=False, + ), + Case.EE_ACCOUNT_MODE_UNRESOLVED: Scenario( + bearer=Bearer.ACCOUNT_MEMBER, + edition=DeploymentEdition.ENTERPRISE, + webapp_auth=True, + access_mode=None, + ), + Case.EE_EXTERNAL_PUBLIC: Scenario( + bearer=Bearer.EXTERNAL, + edition=DeploymentEdition.ENTERPRISE, + webapp_auth=True, + access_mode=_auth_matrix_WebAppAccessMode.PUBLIC, + ), + Case.EE_EXTERNAL_SSO_VERIFIED: Scenario( + bearer=Bearer.EXTERNAL, + edition=DeploymentEdition.ENTERPRISE, + webapp_auth=True, + access_mode=_auth_matrix_WebAppAccessMode.SSO_VERIFIED, + ), + Case.EE_EXTERNAL_PRIVATE_ALL: Scenario( + bearer=Bearer.EXTERNAL, + edition=DeploymentEdition.ENTERPRISE, + webapp_auth=True, + access_mode=_auth_matrix_WebAppAccessMode.PRIVATE_ALL, + ), + Case.EE_EXTERNAL_PRIVATE: Scenario( + bearer=Bearer.EXTERNAL, + edition=DeploymentEdition.ENTERPRISE, + webapp_auth=True, + access_mode=_auth_matrix_WebAppAccessMode.PRIVATE, + ), + Case.EE_EXTERNAL_PRIVATE_REFUSED_WEBAPP_AUTH_OFF: Scenario( + bearer=Bearer.EXTERNAL, + edition=DeploymentEdition.ENTERPRISE, + webapp_auth=False, + access_mode=_auth_matrix_WebAppAccessMode.PRIVATE, + private_app_permitted=False, + ), + Case.EE_EXTERNAL_MODE_UNRESOLVED: Scenario( + bearer=Bearer.EXTERNAL, + edition=DeploymentEdition.ENTERPRISE, + webapp_auth=True, + access_mode=None, + ), + Case.RBAC_ON_LOW_ROLE: Scenario(bearer=Bearer.ACCOUNT_LOW_ROLE, rbac_enabled=True, rbac_allows=True), + Case.RBAC_ON_DENIED: Scenario(bearer=Bearer.ACCOUNT_MEMBER, rbac_enabled=True, rbac_allows=False), +} + + +ROUTER_CASE_ROUTE = "describe.account" +ROUTER_CASES: dict[Case, Expect] = { + Case.NO_BEARER: DENY_NO_BEARER, + Case.LICENSE_INVALID: ADMIT_NO_LICENCE_GATE, + Case.EE_LICENSE_INVALID: DENY_LICENSE, +} +"""Answered by the router before any route-specific requirement runs, so pinned once.""" + + +_ACCOUNT_ONLY_NO_WORKSPACE: dict[Case, Expect] = { + Case.MEMBER: ADMIT, + Case.WRONG_SUBJECT: DENY_WRONG_SUBJECT, + Case.INSUFFICIENT_SCOPE: DENY_SCOPE, + Case.NON_MEMBER: ADMIT, + # No live membership check on this route, so the outsider's only failure is + # scope — same answer pre- and post-PR. + Case.NON_MEMBER_AND_INSUFFICIENT_SCOPE: DENY_SCOPE, + Case.LOW_ROLE: ADMIT, + Case.RBAC_ON_LOW_ROLE: ADMIT_NO_WORKSPACE_ROLE, + Case.RBAC_ON_DENIED: ADMIT_NO_RBAC_PERMISSION, +} + +_ACCOUNT_MEMBER_NO_ROLE: dict[Case, Expect] = { + Case.MEMBER: ADMIT, + Case.WRONG_SUBJECT: DENY_WRONG_SUBJECT, + Case.INSUFFICIENT_SCOPE: DENY_SCOPE, + Case.NON_MEMBER: DENY_NON_MEMBER, + Case.NON_MEMBER_AND_INSUFFICIENT_SCOPE: DENY_NON_MEMBER, + Case.LOW_ROLE: ADMIT, + Case.RBAC_ON_LOW_ROLE: ADMIT_NO_WORKSPACE_ROLE, + Case.RBAC_ON_DENIED: ADMIT_NO_RBAC_PERMISSION, +} + +_ACCOUNT_MEMBER_WITH_ROLE: dict[Case, Expect] = { + Case.MEMBER: ADMIT, + Case.WRONG_SUBJECT: DENY_WRONG_SUBJECT, + Case.INSUFFICIENT_SCOPE: DENY_SCOPE, + Case.NON_MEMBER: DENY_NON_MEMBER, + Case.NON_MEMBER_AND_INSUFFICIENT_SCOPE: DENY_NON_MEMBER, + Case.LOW_ROLE: DENY_ROLE, + Case.RBAC_ON_LOW_ROLE: ADMIT, + Case.RBAC_ON_DENIED: DENY_RBAC, +} + +_DUAL_SUBJECT_RUN: dict[Case, Expect] = { + Case.MEMBER: ADMIT, + Case.WRONG_SUBJECT: DENY_SSO_NEEDS_EE, + Case.INSUFFICIENT_SCOPE: DENY_SCOPE, + Case.NON_MEMBER: DENY_NON_MEMBER, + Case.NON_MEMBER_AND_INSUFFICIENT_SCOPE: DENY_NON_MEMBER, + Case.LOW_ROLE: ADMIT, + Case.APP_API_DISABLED: DENY_API_DISABLED, + Case.UNKNOWN_APP: DENY_UNKNOWN_APP, + Case.FOREIGN_WORKSPACE_QUERY: ADMIT, + Case.EE_ACCOUNT_PUBLIC: ADMIT, + Case.EE_ACCOUNT_SSO_VERIFIED: ADMIT, + Case.EE_ACCOUNT_PRIVATE_ALL: ADMIT, + Case.EE_ACCOUNT_PRIVATE_PERMITTED: ADMIT, + Case.EE_ACCOUNT_PRIVATE_REFUSED: DENY_PRIVATE_APP, + Case.EE_ACCOUNT_PRIVATE_REFUSED_WEBAPP_AUTH_OFF: DENY_PRIVATE_APP, + Case.EE_ACCOUNT_MODE_UNRESOLVED: DENY_MODE_UNRESOLVED, + Case.EE_EXTERNAL_PUBLIC: ADMIT, + Case.EE_EXTERNAL_SSO_VERIFIED: ADMIT, + Case.EE_EXTERNAL_PRIVATE_ALL: DENY_ACCESS_MODE, + Case.EE_EXTERNAL_PRIVATE: DENY_ACCESS_MODE, + Case.EE_EXTERNAL_PRIVATE_REFUSED_WEBAPP_AUTH_OFF: DENY_PRIVATE_APP, + Case.EE_EXTERNAL_MODE_UNRESOLVED: DENY_MODE_UNRESOLVED, + Case.RBAC_ON_DENIED: DENY_RBAC, + Case.RBAC_ON_LOW_ROLE: ADMIT_NO_WORKSPACE_ROLE, +} + + +_ACCOUNT_READER_APP: dict[Case, Expect] = { + Case.MEMBER: ADMIT, + Case.WRONG_SUBJECT: DENY_WRONG_SUBJECT, + Case.INSUFFICIENT_SCOPE: DENY_SCOPE, + Case.NON_MEMBER: DENY_NON_MEMBER, + Case.NON_MEMBER_AND_INSUFFICIENT_SCOPE: DENY_NON_MEMBER, + Case.LOW_ROLE: ADMIT, + Case.APP_API_DISABLED: DENY_API_DISABLED, + Case.UNKNOWN_APP: DENY_UNKNOWN_APP, + Case.FOREIGN_WORKSPACE_QUERY: ADMIT, + Case.EE_ACCOUNT_PUBLIC: ADMIT, + Case.EE_ACCOUNT_SSO_VERIFIED: ADMIT, + Case.EE_ACCOUNT_PRIVATE_ALL: ADMIT, + Case.EE_ACCOUNT_PRIVATE_PERMITTED: ADMIT, + Case.EE_ACCOUNT_PRIVATE_REFUSED: ADMIT, + Case.EE_ACCOUNT_PRIVATE_REFUSED_WEBAPP_AUTH_OFF: ADMIT, + Case.EE_ACCOUNT_MODE_UNRESOLVED: ADMIT, + Case.RBAC_ON_DENIED: DENY_RBAC, + Case.RBAC_ON_LOW_ROLE: ADMIT_NO_WORKSPACE_ROLE, +} + +_ACCOUNT_EDITOR_APP: dict[Case, Expect] = { + **_ACCOUNT_READER_APP, + Case.LOW_ROLE: DENY_ROLE, + Case.RBAC_ON_LOW_ROLE: ADMIT, +} + +_ACCOUNT_READER_SETTINGS: dict[Case, Expect] = {**_ACCOUNT_READER_APP, Case.APP_API_DISABLED: ADMIT} +_ACCOUNT_EDITOR_SETTINGS: dict[Case, Expect] = {**_ACCOUNT_EDITOR_APP, Case.APP_API_DISABLED: ADMIT} + +DENY_NEEDS_WEBAPP_EE = Expect(403, WebAppAccessRequiresEE.description) + +_WEBAPP_ACCESS_EDITOR: dict[Case, Expect] = { + **_ACCOUNT_EDITOR_SETTINGS, + Case.MEMBER: DENY_NEEDS_WEBAPP_EE, + Case.APP_API_DISABLED: DENY_NEEDS_WEBAPP_EE, + Case.FOREIGN_WORKSPACE_QUERY: DENY_NEEDS_WEBAPP_EE, + Case.EE_ACCOUNT_PRIVATE_REFUSED_WEBAPP_AUTH_OFF: DENY_NEEDS_WEBAPP_EE, + Case.RBAC_ON_LOW_ROLE: DENY_NEEDS_WEBAPP_EE, +} + + +_ANY_BEARER: dict[Case, Expect] = { + **{case: ADMIT for case in Case if CASE_REQUIRES[case] <= {Trait.ACCOUNT_PRIMARY, Trait.EXTERNAL_REACHABLE}}, + Case.WRONG_SUBJECT: DENY_SSO_NEEDS_EE, +} +"""A route with no requirements admits every case that carries a valid bearer. The one +exception is the router's own answer: an external-SSO bearer is refused outside enterprise.""" + +MATRIX: dict[str, dict[Case, Expect]] = { + "describe.account": dict(_ACCOUNT_ONLY_NO_WORKSPACE), + "workspaces.list": dict(_ACCOUNT_ONLY_NO_WORKSPACE), + "apps.list": dict(_ACCOUNT_MEMBER_NO_ROLE), + "workspaces.switch": dict(_ACCOUNT_MEMBER_NO_ROLE), + # invite/remove check WORKSPACE_MEMBER_MANAGE; update_role checks WORKSPACE_ROLE_MANAGE + # (matches check_member_permission's own add/remove vs update split) + "workspaces.members.invite": dict(_ACCOUNT_MEMBER_WITH_ROLE), + "workspaces.members.update_role": dict(_ACCOUNT_MEMBER_WITH_ROLE), + "plugins.install": {**_ACCOUNT_MEMBER_NO_ROLE, Case.RBAC_ON_DENIED: DENY_RBAC}, + "model_providers.credential.create": dict(_ACCOUNT_MEMBER_WITH_ROLE), + "app_dsl.import": dict(_ACCOUNT_MEMBER_WITH_ROLE), + "app_dsl.check": dict(_ACCOUNT_MEMBER_WITH_ROLE), + "app_create.workflow": dict(_ACCOUNT_MEMBER_WITH_ROLE), + "node_types.list": dict(_ANY_BEARER), + "app_workflow.node_run.workflow": dict(_ACCOUNT_EDITOR_APP), + "app_info.describe.workflow": dict(_ACCOUNT_READER_SETTINGS), + "app_info.set.workflow": dict(_ACCOUNT_EDITOR_SETTINGS), + "service_api.describe": dict(_ACCOUNT_EDITOR_SETTINGS), + "service_api.set": dict(_ACCOUNT_EDITOR_SETTINGS), + "webapp.describe.workflow": dict(_ACCOUNT_READER_SETTINGS), + "webapp.set.workflow": dict(_ACCOUNT_EDITOR_SETTINGS), + "webapp.reset": dict(_ACCOUNT_EDITOR_SETTINGS), + "webapp_access.describe": dict(_WEBAPP_ACCESS_EDITOR), + "apps.describe": dict(_ACCOUNT_READER_APP), + "app_dsl.export": dict(_ACCOUNT_EDITOR_APP), + "app_run.draft.workflow": dict(_ACCOUNT_EDITOR_APP), + "app_workflow.run.list": dict(_ACCOUNT_READER_APP), + "app_workflow.run.describe": dict(_ACCOUNT_READER_APP), + "app_workflow.run.nodes": dict(_ACCOUNT_READER_APP), + "app_workflow.publish": dict(_ACCOUNT_EDITOR_APP), + "app_workflow.version.list": dict(_ACCOUNT_EDITOR_APP), + "app_workflow.version.restore": dict(_ACCOUNT_EDITOR_APP), + "app_workflow.release_check": dict(_ACCOUNT_EDITOR_APP), + "app_workflow.env.list": dict(_ACCOUNT_EDITOR_APP), + "app_workflow.env.set": dict(_ACCOUNT_EDITOR_APP), + "app_workflow.env.delete": dict(_ACCOUNT_EDITOR_APP), + "app_run.workflow": dict(_DUAL_SUBJECT_RUN), + "app_run.chat": dict(_DUAL_SUBJECT_RUN), + "app_run.advanced_chat": dict(_DUAL_SUBJECT_RUN), + "app_run.completion": dict(_DUAL_SUBJECT_RUN), + "human_input_form.get": dict(_DUAL_SUBJECT_RUN), + "files.upload": { + **_DUAL_SUBJECT_RUN, + Case.RBAC_ON_DENIED: ADMIT_NO_RBAC_PERMISSION, + }, + "permitted_external.list": { + Case.MEMBER: ADMIT_NO_MOUNT, + Case.WRONG_SUBJECT: DENY_WRONG_SUBJECT, + Case.INSUFFICIENT_SCOPE: DENY_SCOPE, + # An account bearer answers on the subject axis here, before scope or + # membership: this route's CheckSubject allows only ExternalSsoSubject. + Case.NON_MEMBER_AND_INSUFFICIENT_SCOPE: DENY_WRONG_SUBJECT, + Case.EDITION_NOT_ENTERPRISE: DENY_EDITION, + }, + "permitted_external.describe": { + Case.MEMBER: ADMIT, + Case.WRONG_SUBJECT: DENY_WRONG_SUBJECT, + Case.INSUFFICIENT_SCOPE: DENY_SCOPE, + Case.NON_MEMBER_AND_INSUFFICIENT_SCOPE: DENY_WRONG_SUBJECT, + Case.EDITION_NOT_ENTERPRISE: DENY_EDITION, + Case.APP_API_DISABLED: DENY_API_DISABLED, + Case.UNKNOWN_APP: DENY_UNKNOWN_APP, + Case.FOREIGN_WORKSPACE_QUERY: Expect( + ADMITTED, + note="a foreign ?workspace_id= is ignored on every app-scoped route (spec 2.9)", + ), + Case.EE_EXTERNAL_PUBLIC: ADMIT, + Case.EE_EXTERNAL_SSO_VERIFIED: ADMIT, + Case.EE_EXTERNAL_PRIVATE_ALL: DENY_ACCESS_MODE, + Case.EE_EXTERNAL_PRIVATE: DENY_ACCESS_MODE, + Case.EE_EXTERNAL_PRIVATE_REFUSED_WEBAPP_AUTH_OFF: DENY_PRIVATE_APP, + Case.EE_EXTERNAL_MODE_UNRESOLVED: DENY_MODE_UNRESOLVED, + }, +} +"""One representative per distinct declared requirement tuple; `DECLARED` pins the rest.""" + + +_ACCOUNT = (AccountSubject,) +_ACCOUNT_OR_EXTERNAL = (AccountSubject, ExternalSsoSubject) +_OWNER_ADMIN = frozenset({TenantAccountRole.OWNER, TenantAccountRole.ADMIN}) +_EDITOR_UP = frozenset({TenantAccountRole.EDITOR, TenantAccountRole.ADMIN, TenantAccountRole.OWNER}) + +_REQ_ACCOUNT_FULL = (CheckSubject(allowed=_ACCOUNT), CheckScope(Scope.FULL)) +_REQ_ACCOUNT_WORKSPACE_READ = (CheckSubject(allowed=_ACCOUNT), CheckScope(Scope.WORKSPACE_READ)) +_REQ_ACCOUNT_APPS_READ_MEMBER = ( + CheckSubject(allowed=_ACCOUNT), + CheckScope(Scope.APPS_READ), + CheckWorkspaceMember(), +) +_REQ_ACCOUNT_WORKSPACE_READ_MEMBER = ( + CheckSubject(allowed=_ACCOUNT), + CheckScope(Scope.WORKSPACE_READ), + CheckWorkspaceMember(), +) +_REQ_MEMBER_MANAGE = ( + CheckSubject(allowed=_ACCOUNT), + CheckScope(Scope.WORKSPACE_WRITE), + CheckWorkspaceMember(), + CheckRBACPermission(RBACCheck(RBACPermission.WORKSPACE_MEMBER_MANAGE, Workspace())), + CheckWorkspaceRole(_OWNER_ADMIN), +) +_REQ_ROLE_MANAGE = ( + CheckSubject(allowed=_ACCOUNT), + CheckScope(Scope.WORKSPACE_WRITE), + CheckWorkspaceMember(), + CheckRBACPermission(RBACCheck(RBACPermission.WORKSPACE_ROLE_MANAGE, Workspace())), + CheckWorkspaceRole(_OWNER_ADMIN), +) +_REQ_DSL_WORKSPACE = ( + CheckSubject(allowed=_ACCOUNT), + CheckScope(Scope.WORKSPACE_WRITE), + CheckWorkspaceMember(), + CheckRBACPermission(RBACCheck(RBACPermission.APP_IMPORT_EXPORT_DSL, Workspace())), + CheckWorkspaceRole(_EDITOR_UP), +) +_REQ_DSL_APP = ( + CheckSubject(allowed=_ACCOUNT), + CheckAppApiEnabled(), + CheckWorkspaceMember(), + CheckScope(Scope.APPS_READ), + CheckRBACPermission(RBACCheck(RBACPermission.APP_IMPORT_EXPORT_DSL, PlainApp())), + CheckWorkspaceRole(_EDITOR_UP), +) +_REQ_APP_DESCRIBE = ( + CheckSubject(allowed=_ACCOUNT), + CheckAppApiEnabled(), + CheckWorkspaceMember(), + CheckScope(Scope.APPS_READ), + CheckRBACPermission(RBACCheck(RBACPermission.APP_VIEW_LAYOUT, PlainApp())), +) +_REQ_RUN = ( + CheckSubject(allowed=_ACCOUNT_OR_EXTERNAL), + CheckAppApiEnabled(), + CheckWorkspaceMember(), + CheckScope(Scope.APPS_RUN), + CheckRBACPermission(RBACCheck(RBACPermission.APP_TEST_AND_RUN, PlainApp())), + CheckAppAccess(), +) +_REQ_RUN_FORM = (*_REQ_RUN, CheckFormSurface()) +_REQ_RELEASE = ( + CheckSubject(allowed=_ACCOUNT), + CheckAppApiEnabled(), + CheckWorkspaceMember(), + CheckAppMode(_auth_matrix_AppMode.WORKFLOW, _auth_matrix_AppMode.ADVANCED_CHAT), + CheckScope(Scope.WORKSPACE_WRITE), + CheckRBACPermission(RBACCheck(RBACPermission.APP_RELEASE_AND_VERSION, PlainApp())), + CheckWorkspaceRole(_EDITOR_UP), +) +_REQ_VERSION_READ = ( + CheckSubject(allowed=_ACCOUNT), + CheckAppApiEnabled(), + CheckWorkspaceMember(), + CheckAppMode(_auth_matrix_AppMode.WORKFLOW, _auth_matrix_AppMode.ADVANCED_CHAT), + CheckScope(Scope.APPS_READ), + CheckRBACPermission(RBACCheck(RBACPermission.APP_VIEW_LAYOUT, PlainApp())), + CheckWorkspaceRole(_EDITOR_UP), +) +_REQ_ENV_WRITE = ( + CheckSubject(allowed=_ACCOUNT), + CheckAppApiEnabled(), + CheckWorkspaceMember(), + CheckAppMode(_auth_matrix_AppMode.WORKFLOW, _auth_matrix_AppMode.ADVANCED_CHAT), + CheckScope(Scope.WORKSPACE_WRITE), + CheckRBACPermission(RBACCheck(RBACPermission.APP_EDIT, PlainApp())), + CheckWorkspaceRole(_EDITOR_UP), +) +_REQ_RUN_HISTORY = ( + CheckSubject(allowed=_ACCOUNT), + CheckAppApiEnabled(), + CheckWorkspaceMember(), + CheckAppMode(_auth_matrix_AppMode.WORKFLOW, _auth_matrix_AppMode.ADVANCED_CHAT), + CheckScope(Scope.APPS_READ), + CheckRBACPermission(RBACCheck(RBACPermission.APP_CREATE_AND_MANAGEMENT, PlainApp())), +) +_REQ_DRAFT_RUN = ( + CheckSubject(allowed=_ACCOUNT), + CheckAppApiEnabled(), + CheckWorkspaceMember(), + CheckScope(Scope.APPS_RUN), + CheckRBACPermission(RBACCheck(RBACPermission.APP_TEST_AND_RUN, PlainApp())), + CheckWorkspaceRole(_EDITOR_UP), +) + + +def _draft_req(mode: _auth_matrix_AppMode) -> tuple[Requirement, ...]: + return (*_REQ_DRAFT_RUN[:3], CheckAppMode(mode), *_REQ_DRAFT_RUN[3:]) + + +_REQ_NODE_TYPES: tuple[Requirement, ...] = () +_REQ_APP_CREATE = ( + CheckSubject(allowed=_ACCOUNT), + CheckScope(Scope.WORKSPACE_WRITE), + CheckWorkspaceMember(), + CheckRBACPermission(RBACCheck(RBACPermission.APP_CREATE_AND_MANAGEMENT, Workspace())), + CheckWorkspaceRole(_EDITOR_UP), + CheckAppQuota(), +) +_REQ_FILES = ( + CheckSubject(allowed=_ACCOUNT_OR_EXTERNAL), + CheckAppApiEnabled(), + CheckWorkspaceMember(), + CheckScope(Scope.APPS_RUN), + CheckAppAccess(), +) +_REQ_EXTERNAL_LIST = ( + CheckSubject(allowed=(ExternalSsoSubject,)), + CheckScope(Scope.APPS_READ_PERMITTED_EXTERNAL), +) +_REQ_EXTERNAL_DESCRIBE = ( + CheckSubject(allowed=(ExternalSsoSubject,)), + CheckAppApiEnabled(), + CheckScope(Scope.APPS_READ_PERMITTED_EXTERNAL), + CheckAppAccess(), +) + + +def _req_plugin_write(permission: RBACPermission) -> tuple[Requirement, ...]: + return ( + CheckSubject(allowed=_ACCOUNT), + CheckScope(Scope.WORKSPACE_WRITE), + CheckWorkspaceMember(), + CheckRBACPermission(RBACCheck(permission, Workspace())), + CheckPluginInstallSetting(), + ) + + +def _req_admin_write(permission: RBACPermission) -> tuple[Requirement, ...]: + return ( + CheckSubject(allowed=_ACCOUNT), + CheckScope(Scope.WORKSPACE_WRITE), + CheckWorkspaceMember(), + CheckRBACPermission(RBACCheck(permission, Workspace())), + CheckWorkspaceRole(_OWNER_ADMIN), + ) + + +_REQ_PLUGIN_TASK = ( + CheckSubject(allowed=_ACCOUNT), + CheckScope(Scope.WORKSPACE_READ), + CheckWorkspaceMember(), + CheckPluginInstallSetting(), +) + + +def _settings_req( + modes: tuple[_auth_matrix_AppMode, ...], + perm: RBACPermission, + locator: PlainApp | AgentBehindApp, + scope: Scope, + roles: frozenset[TenantAccountRole] | None, +) -> tuple[Requirement, ...]: + base: tuple[Requirement, ...] = ( + CheckSubject(allowed=_ACCOUNT), + CheckWorkspaceMember(), + CheckAppMode(*modes), + CheckScope(scope), + CheckRBACPermission(RBACCheck(perm, locator)), + ) + return (*base, CheckWorkspaceRole(roles)) if roles is not None else base + + +_RESOURCE_RUN = (CheckSubject(allowed=(*_ACCOUNT_OR_EXTERNAL, ResourceAccessSubject)), *_REQ_RUN[1:]) + +DECLARED: dict[str, tuple[Requirement, ...]] = { + "knowledge_fs.ls": _REQ_ACCOUNT_WORKSPACE_READ_MEMBER, + "knowledge_fs.tree": _REQ_ACCOUNT_WORKSPACE_READ_MEMBER, + "knowledge_fs.grep": _REQ_ACCOUNT_WORKSPACE_READ_MEMBER, + "knowledge_fs.find": _REQ_ACCOUNT_WORKSPACE_READ_MEMBER, + "knowledge_fs.diff": _REQ_ACCOUNT_WORKSPACE_READ_MEMBER, + "knowledge_fs.cat": _REQ_ACCOUNT_WORKSPACE_READ_MEMBER, + "knowledge_fs.stat": _REQ_ACCOUNT_WORKSPACE_READ_MEMBER, + "describe.account": _REQ_ACCOUNT_FULL, + "account.sessions.revoke_self": _REQ_ACCOUNT_FULL, + "get.account.session": _REQ_ACCOUNT_FULL, + "account.sessions.revoke_one": _REQ_ACCOUNT_FULL, + "apps.describe": (CheckSubject(allowed=(AccountSubject, ResourceAccessSubject)), *_REQ_APP_DESCRIBE[1:]), + "apps.list": (CheckSubject(allowed=(AccountSubject, ResourceAccessSubject)), *_REQ_ACCOUNT_APPS_READ_MEMBER[1:]), + "workspaces.list": ( + CheckSubject(allowed=(AccountSubject, ResourceAccessSubject)), + *_REQ_ACCOUNT_WORKSPACE_READ[1:], + ), + "workspaces.describe": _REQ_ACCOUNT_WORKSPACE_READ, + "workspaces.switch": _REQ_ACCOUNT_WORKSPACE_READ_MEMBER, + "workspaces.members.list": _REQ_ACCOUNT_WORKSPACE_READ_MEMBER, + "workspaces.members.invite": (*_REQ_MEMBER_MANAGE, CheckWorkspaceInvitationQuota()), + "workspaces.members.remove": _REQ_MEMBER_MANAGE, + "workspaces.members.update_role": _REQ_ROLE_MANAGE, + "plugins.marketplace": _REQ_ACCOUNT_WORKSPACE_READ_MEMBER, + "plugins.list": _REQ_ACCOUNT_WORKSPACE_READ_MEMBER, + "plugins.install": _req_plugin_write(RBACPermission.PLUGIN_INSTALL), + "plugins.task": _REQ_PLUGIN_TASK, + "plugins.delete": _req_plugin_write(RBACPermission.PLUGIN_DELETE), + "model_providers.describe": _REQ_ACCOUNT_WORKSPACE_READ_MEMBER, + "model_providers.credential.create": _req_admin_write(RBACPermission.CREDENTIAL_CREATE), + "model_providers.credential.set": _req_admin_write(RBACPermission.CREDENTIAL_MANAGE), + "tools.list": _REQ_ACCOUNT_WORKSPACE_READ_MEMBER, + "tool_providers.describe": _REQ_ACCOUNT_WORKSPACE_READ_MEMBER, + "tool_providers.credential.create": _req_admin_write(RBACPermission.CREDENTIAL_CREATE), + "tool_providers.credential.set": _req_admin_write(RBACPermission.CREDENTIAL_MANAGE), + "knowledge_bases.list": _REQ_ACCOUNT_WORKSPACE_READ_MEMBER, + "models.list": _REQ_ACCOUNT_WORKSPACE_READ_MEMBER, + "models.credential.create": _req_admin_write(RBACPermission.CREDENTIAL_CREATE), + "models.credential.set": _req_admin_write(RBACPermission.CREDENTIAL_MANAGE), + "app_dsl.import": _REQ_DSL_WORKSPACE, + "app_dsl.check": _REQ_DSL_WORKSPACE, + "app_dsl.import_confirm": _REQ_DSL_WORKSPACE, + "app_dsl.export": _REQ_DSL_APP, + "app_dsl.check_dependencies": _REQ_DSL_APP, + "app_run.workflow": _RESOURCE_RUN, + "app_run.chat": _RESOURCE_RUN, + "app_run.advanced_chat": _RESOURCE_RUN, + "app_run.completion": _RESOURCE_RUN, + "app_run.draft.workflow": _draft_req(_auth_matrix_AppMode.WORKFLOW), + "app_run.draft.advanced_chat": _draft_req(_auth_matrix_AppMode.ADVANCED_CHAT), + "app_workflow.run.list": _REQ_RUN_HISTORY, + "app_workflow.run.describe": _REQ_RUN_HISTORY, + "app_workflow.run.nodes": _REQ_RUN_HISTORY, + "app_workflow.publish": _REQ_RELEASE, + "app_workflow.version.list": _REQ_VERSION_READ, + "app_workflow.version.restore": _REQ_RELEASE, + "app_workflow.release_check": _REQ_VERSION_READ, + "app_workflow.env.list": _REQ_VERSION_READ, + "app_workflow.env.set": _REQ_ENV_WRITE, + "app_workflow.env.delete": _REQ_ENV_WRITE, + "node_types.list": _REQ_NODE_TYPES, + "node_types.describe": _REQ_NODE_TYPES, + "app_create.workflow": _REQ_APP_CREATE, + "app_create.advanced_chat": _REQ_APP_CREATE, + "app_workflow.node_run.workflow": _draft_req(_auth_matrix_AppMode.WORKFLOW), + "app_workflow.node_run.advanced_chat": _draft_req(_auth_matrix_AppMode.ADVANCED_CHAT), + "app_info.describe.workflow": _settings_req( + (_auth_matrix_AppMode.WORKFLOW,), RBACPermission.APP_VIEW_LAYOUT, PlainApp(), Scope.APPS_READ, None + ), + "app_info.set.workflow": _settings_req( + (_auth_matrix_AppMode.WORKFLOW,), RBACPermission.APP_EDIT, PlainApp(), Scope.WORKSPACE_WRITE, _EDITOR_UP + ), + "service_api.describe": _settings_req( + REGULAR_MODES, RBACPermission.APP_RELEASE_AND_VERSION, PlainApp(), Scope.APPS_READ, _EDITOR_UP + ), + "service_api.set": _settings_req( + REGULAR_MODES, RBACPermission.APP_RELEASE_AND_VERSION, PlainApp(), Scope.WORKSPACE_WRITE, _OWNER_ADMIN + ), + "app_info.describe.advanced_chat": _settings_req( + (_auth_matrix_AppMode.ADVANCED_CHAT,), RBACPermission.APP_VIEW_LAYOUT, PlainApp(), Scope.APPS_READ, None + ), + "app_info.set.advanced_chat": _settings_req( + (_auth_matrix_AppMode.ADVANCED_CHAT,), RBACPermission.APP_EDIT, PlainApp(), Scope.WORKSPACE_WRITE, _EDITOR_UP + ), + "app_info.describe.chat": _settings_req( + (_auth_matrix_AppMode.CHAT,), RBACPermission.APP_VIEW_LAYOUT, PlainApp(), Scope.APPS_READ, None + ), + "app_info.set.chat": _settings_req( + (_auth_matrix_AppMode.CHAT,), RBACPermission.APP_EDIT, PlainApp(), Scope.WORKSPACE_WRITE, _EDITOR_UP + ), + "app_info.describe.agent_chat": _settings_req( + (_auth_matrix_AppMode.AGENT_CHAT,), RBACPermission.APP_VIEW_LAYOUT, PlainApp(), Scope.APPS_READ, None + ), + "app_info.set.agent_chat": _settings_req( + (_auth_matrix_AppMode.AGENT_CHAT,), RBACPermission.APP_EDIT, PlainApp(), Scope.WORKSPACE_WRITE, _EDITOR_UP + ), + "app_info.describe.completion": _settings_req( + (_auth_matrix_AppMode.COMPLETION,), RBACPermission.APP_VIEW_LAYOUT, PlainApp(), Scope.APPS_READ, None + ), + "app_info.set.completion": _settings_req( + (_auth_matrix_AppMode.COMPLETION,), RBACPermission.APP_EDIT, PlainApp(), Scope.WORKSPACE_WRITE, _EDITOR_UP + ), + "app_info.describe.agent": _settings_req( + (_auth_matrix_AppMode.AGENT,), RBACPermission.AGENT_PREVIEW, AgentBehindApp(), Scope.APPS_READ, None + ), + "app_info.set.agent": _settings_req( + (_auth_matrix_AppMode.AGENT,), RBACPermission.AGENT_EDIT, AgentBehindApp(), Scope.WORKSPACE_WRITE, _EDITOR_UP + ), + "service_api.describe.agent": _settings_req( + (_auth_matrix_AppMode.AGENT,), + RBACPermission.AGENT_ACCESS_POINT_VIEW, + AgentBehindApp(), + Scope.APPS_READ, + _EDITOR_UP, + ), + "service_api.set.agent": _settings_req( + (_auth_matrix_AppMode.AGENT,), + RBACPermission.AGENT_ACCESS_POINT_MANAGE, + AgentBehindApp(), + Scope.WORKSPACE_WRITE, + _OWNER_ADMIN, + ), + "webapp.describe.workflow": _settings_req( + (_auth_matrix_AppMode.WORKFLOW,), RBACPermission.APP_VIEW_LAYOUT, PlainApp(), Scope.APPS_READ, None + ), + "webapp.set.workflow": _settings_req( + (_auth_matrix_AppMode.WORKFLOW,), + RBACPermission.APP_RELEASE_AND_VERSION, + PlainApp(), + Scope.WORKSPACE_WRITE, + _EDITOR_UP, + ), + "webapp.reset": _settings_req( + REGULAR_MODES, RBACPermission.APP_RELEASE_AND_VERSION, PlainApp(), Scope.WORKSPACE_WRITE, _OWNER_ADMIN + ), + "webapp_access.describe": ( + *_settings_req(REGULAR_MODES, RBACPermission.APP_ACCESS_CONFIG, PlainApp(), Scope.APPS_READ, _EDITOR_UP), + CheckWebAppAuthEnterprise(), + ), + "webapp_access.set": ( + *_settings_req(REGULAR_MODES, RBACPermission.APP_ACCESS_CONFIG, PlainApp(), Scope.WORKSPACE_WRITE, _EDITOR_UP), + CheckWebAppAuthEnterprise(), + ), + "webapp.describe.advanced_chat": _settings_req( + (_auth_matrix_AppMode.ADVANCED_CHAT,), RBACPermission.APP_VIEW_LAYOUT, PlainApp(), Scope.APPS_READ, None + ), + "webapp.set.advanced_chat": _settings_req( + (_auth_matrix_AppMode.ADVANCED_CHAT,), + RBACPermission.APP_RELEASE_AND_VERSION, + PlainApp(), + Scope.WORKSPACE_WRITE, + _EDITOR_UP, + ), + "webapp.describe.chat": _settings_req( + (_auth_matrix_AppMode.CHAT,), RBACPermission.APP_VIEW_LAYOUT, PlainApp(), Scope.APPS_READ, None + ), + "webapp.set.chat": _settings_req( + (_auth_matrix_AppMode.CHAT,), + RBACPermission.APP_RELEASE_AND_VERSION, + PlainApp(), + Scope.WORKSPACE_WRITE, + _EDITOR_UP, + ), + "webapp.describe.agent_chat": _settings_req( + (_auth_matrix_AppMode.AGENT_CHAT,), RBACPermission.APP_VIEW_LAYOUT, PlainApp(), Scope.APPS_READ, None + ), + "webapp.set.agent_chat": _settings_req( + (_auth_matrix_AppMode.AGENT_CHAT,), + RBACPermission.APP_RELEASE_AND_VERSION, + PlainApp(), + Scope.WORKSPACE_WRITE, + _EDITOR_UP, + ), + "webapp.describe.completion": _settings_req( + (_auth_matrix_AppMode.COMPLETION,), RBACPermission.APP_VIEW_LAYOUT, PlainApp(), Scope.APPS_READ, None + ), + "webapp.set.completion": _settings_req( + (_auth_matrix_AppMode.COMPLETION,), + RBACPermission.APP_RELEASE_AND_VERSION, + PlainApp(), + Scope.WORKSPACE_WRITE, + _EDITOR_UP, + ), + "webapp.describe.agent": _settings_req( + (_auth_matrix_AppMode.AGENT,), RBACPermission.AGENT_ACCESS_POINT_VIEW, AgentBehindApp(), Scope.APPS_READ, None + ), + "webapp.set.agent": _settings_req( + (_auth_matrix_AppMode.AGENT,), + RBACPermission.AGENT_ACCESS_POINT_MANAGE, + AgentBehindApp(), + Scope.WORKSPACE_WRITE, + _EDITOR_UP, + ), + "webapp.reset.agent": _settings_req( + (_auth_matrix_AppMode.AGENT,), + RBACPermission.AGENT_ACCESS_POINT_MANAGE, + AgentBehindApp(), + Scope.WORKSPACE_WRITE, + _OWNER_ADMIN, + ), + "webapp_access.describe.agent": ( + *_settings_req( + (_auth_matrix_AppMode.AGENT,), + RBACPermission.AGENT_ACCESS_CONFIG, + AgentBehindApp(), + Scope.APPS_READ, + _EDITOR_UP, + ), + CheckWebAppAuthEnterprise(), + ), + "webapp_access.set.agent": ( + *_settings_req( + (_auth_matrix_AppMode.AGENT,), + RBACPermission.AGENT_ACCESS_CONFIG, + AgentBehindApp(), + Scope.WORKSPACE_WRITE, + _EDITOR_UP, + ), + CheckWebAppAuthEnterprise(), + ), + "app_run.stop": _RESOURCE_RUN, + "files.upload": _REQ_FILES, + "human_input_form.get": _REQ_RUN_FORM, + "human_input_form.submit": _REQ_RUN_FORM, + "workflow_events.stream": _REQ_RUN, + "permitted_external.list": _REQ_EXTERNAL_LIST, + "permitted_external.describe": _REQ_EXTERNAL_DESCRIBE, +} +"""Every route's `@endpoint(requirements=...)`, in declared order. A route that +shares a tuple with a `MATRIX` representative is covered by that representative's +rows only while this table says it still declares the same tuple.""" + + +ROUTES_BY_ID = {route.id: route for route in ROUTES} + + +def _applicable(route: Route, case: Case) -> bool: + return CASE_REQUIRES[case] <= route.traits + + +def _rows(route: Route) -> Iterator[tuple[Route, Case, Expect]]: + for case in Case: + if not _applicable(route, case): + continue + if case not in ROUTER_CASES: + yield route, case, MATRIX[route.id][case] + elif route.id == ROUTER_CASE_ROUTE: + yield route, case, ROUTER_CASES[case] + + +ROWS: tuple[tuple[Route, Case, Expect], ...] = tuple( + row for route in ROUTES if route.id in MATRIX for row in _rows(route) +) + + +@dataclass(slots=True) +class World: + """Persisted rows and the tokens that address them.""" + + workspace_id: str + other_workspace_id: str + app_id: str + disabled_app_id: str + member_account_id: str + low_role_account_id: str + outsider_account_id: str + tokens: dict[Bearer, str] = field(default_factory=dict) + session_ids: dict[Bearer, str] = field(default_factory=dict) + + +def _admit_on_mount(*_args: object, **_kwargs: object) -> None: + raise ImATeapot("admitted") + + +class _MemoryResolver: + """Stands in for the DB-backed token resolver. + + Token resolution is an input to the matrix, not part of it, and the shipped + resolver compares a timezone-aware expiry against a column SQLite hands back + naive. + """ + + def __init__(self, rows: dict[str, ResolvedRow]) -> None: + self._rows = rows + + def resolve(self, token_hash: str) -> ResolvedRow | None: + return self._rows.get(token_hash) + + +def _authenticator(rows: dict[str, ResolvedRow]) -> BearerAuthenticator: + resolver = _MemoryResolver(rows) + return BearerAuthenticator({TokenType.OAUTH_ACCOUNT: resolver, TokenType.OAUTH_EXTERNAL_SSO: resolver}) + + +def _system_features( + *, edition: DeploymentEdition, webapp_auth: bool, license_status: LicenseStatus +) -> SystemFeatureModel: + features = SystemFeatureModel(agent_tracing_api_collection_enabled=False, deployment_edition=edition) + features.webapp_auth.enabled = webapp_auth + features.license.status = license_status + return features + + +@dataclass(frozen=True, slots=True) +class _WebAppSettings: + access_mode: str + + +def _access_mode_settings(access_mode: _auth_matrix_WebAppAccessMode | None) -> _WebAppSettings | None: + if access_mode is None: + return None + return _WebAppSettings(access_mode=access_mode.value) + + +def _webapp_account() -> Account: + """The account an external-SSO subject's email resolves to for the private-app check.""" + row = Account(name="external", email="external@example.com", avatar="", status=AccountStatus.ACTIVE) + row.id = str(uuid.uuid4()) + return row + + +def _end_user(_type: EndUserType, tenant_id: str, app_id: str, user_id: str | None = None) -> EndUser: + row = EndUser( + tenant_id=tenant_id, + app_id=app_id, + type=EndUserType.OPENAPI, + _is_anonymous=False, + session_id=user_id or "", + ) + row.external_user_id = user_id + return row + + +def _bearer_for(route: Route, scenario: Scenario) -> Bearer | None: + account_primary = Trait.ACCOUNT_PRIMARY in route.traits + match scenario.bearer: + case Bearer.NONE: + return None + case Bearer.PRIMARY: + return Bearer.ACCOUNT_MEMBER if account_primary else Bearer.EXTERNAL + case Bearer.OTHER_SUBJECT: + return Bearer.EXTERNAL if account_primary else Bearer.ACCOUNT_MEMBER + case _: + return scenario.bearer + + +def _url(route: Route, world: World, scenario: Scenario, bearer: Bearer | None) -> str: + if scenario.unknown_app: + app_id = str(uuid.uuid4()) + elif not scenario.app_api_enabled: + app_id = world.disabled_app_id + else: + app_id = world.app_id + ids = { + "app_id": app_id, + "workspace_id": world.workspace_id, + # The caller's own session: the access service 404s any other id. + "session_id": world.session_ids.get(bearer, str(uuid.uuid4())) if bearer else str(uuid.uuid4()), + "member_id": str(uuid.uuid4()), + "import_id": str(uuid.uuid4()), + "task_id": str(uuid.uuid4()), + "form_token": uuid.uuid4().hex, + "run_id": str(uuid.uuid4()), + "version_id": str(uuid.uuid4()), + "env_id": str(uuid.uuid4()), + "node_type": "llm", + "node_id": "node-1", + "plugin_id": "langgenius/openai", + "provider": "langgenius/openai/openai", + "credential_id": str(uuid.uuid4()), + "model_type": "llm", + } + query = route.query.format(**ids) + if scenario.foreign_workspace_query: + extra = f"workspace_id={world.other_workspace_id}" + query = f"{query}&{extra}" if query else extra + path = route.path.format(**ids) + return f"/openapi/v1{path}?{query}" if query else f"/openapi/v1{path}" + + +def _run_case( + *, + route: Route, + scenario: Scenario, + app: Flask, + world: World, + token_rows: dict[str, ResolvedRow], + monkeypatch: pytest.MonkeyPatch, +) -> TestResponse: + edition = scenario.edition or ( + DeploymentEdition.ENTERPRISE if Trait.ENTERPRISE_ONLY in route.traits else DeploymentEdition.COMMUNITY + ) + apply_config_overrides(monkeypatch, DEPLOYMENT_EDITION=edition, RBAC_ENABLED=scenario.rbac_enabled) + monkeypatch.setattr(oauth_bearer_module, "_authenticator", _authenticator(token_rows)) + if scenario.narrow_scopes: + # No shipped token can fail `CheckScope`: `dfoa_` carries `Scope.FULL` and + # `dfoe_` exactly the scopes its routes ask for. Scopes are a constant of + # the subject, so the only way to reach the refusal is to blank them here. + # Keep these rows: they are what catches a third token kind routed + # somewhere it was not scoped for. + monkeypatch.setattr(oauth_bearer_module.AuthContext, "scopes", property(lambda _self: frozenset())) + + features = _system_features( + edition=edition, webapp_auth=scenario.webapp_auth, license_status=scenario.license_status + ) + settings = _access_mode_settings(scenario.access_mode) + + headers: dict[str, str] = {CATALOG_HEADER: catalog_for(app)[1]} + bearer = _bearer_for(route, scenario) + if bearer is not None: + headers["Authorization"] = f"Bearer {world.tokens[bearer]}" + + with ExitStack() as stack: + stack.enter_context(patch.object(SystemFeatureService, "get_public_system_features", return_value=features)) + stack.enter_context( + patch.object(EnterpriseService.WebAppAuth, "get_app_access_mode_by_id", return_value=settings) + ) + stack.enter_context( + patch.object( + EnterpriseService.WebAppAuth, + "is_user_allowed_to_access_webapp", + return_value=scenario.private_app_permitted, + ) + ) + stack.enter_context( + patch.object( + EnterpriseService.WebAppAuth, + "list_externally_accessible_apps", + return_value={"data": [], "total": 0, "hasMore": False}, + ) + ) + services = subjects.application_services() + stack.enter_context( + patch.object( + services.app_scoped_end_users.commands, + "get_or_create_end_user_by_type", + side_effect=_end_user, + ) + ) + stack.enter_context( + patch( + "services.enterprise.rbac_service.RBACService.CheckAccess.check", + return_value=scenario.rbac_allows, + ) + ) + stack.enter_context( + patch.object( + services.accounts.identity, + "get_account_by_email", + return_value=_webapp_account(), + ) + ) + client = app.test_client() + return contextvars.copy_context().run( + lambda: client.open(_url(route, world, scenario, bearer), method=route.method, headers=headers) + ) + + +_PATH_PARAMS: Final = frozenset({"plugin_id", "provider"}) + + +def _rule_path(route: Route) -> str: + path = route.path + for name in _PATH_PARAMS: + path = path.replace(f"{{{name}}}", f"") + return "/openapi/v1" + path.replace("{", "") + + +def _endpoint_spec(app: Flask, rule: Rule, method: str) -> EndpointSpec | None: + """The `EndpointSpec` `@endpoint` attached to this rule's handler, or None. + + Carrying a spec *is* being guarded: `@endpoint` is the only thing that sets + it, and it sets it on the same object it wraps in `subject_router.guard`. + """ + view = app.view_functions.get(rule.endpoint) + resource = view.view_class if hasattr(view, "view_class") else None + if resource is None: + return None + handler = getattr(resource, method.lower(), None) # guard-ignore: no-new-getattr -- HTTP verb selects the handler + return handler.__spec__ if hasattr(handler, "__spec__") else None + + +@singledispatch +def _config(_requirement: Requirement) -> object: + return None + + +@_config.register +def _(requirement: CheckSubject) -> object: + return requirement.allowed + + +@_config.register +def _(requirement: CheckAppMode) -> object: + return requirement.modes + + +@_config.register +def _(requirement: CheckScope) -> object: + return requirement.scope + + +@_config.register +def _(requirement: CheckRBACPermission) -> object: + return tuple((check.scene, repr(check.locator)) for check in requirement.checks) + + +@_config.register +def _(requirement: CheckWorkspaceRole) -> object: + return requirement.allowed_roles + + +def _canonical(requirement: Requirement) -> tuple[type[Requirement], object]: + """`Requirement` has no `__eq__`, so compare by type plus the config that sets it apart.""" + return type(requirement), _config(requirement) + + +def _spec_for(app: Flask, route: Route) -> EndpointSpec | None: + for rule in app.url_map.iter_rules(): + if str(rule) == _rule_path(route) and route.method in (rule.methods or set()): + return _endpoint_spec(app, rule, route.method) + return None + + +ERROR_DEFAULT_RESPONSE: dict[str, object] = { + "description": "Error", + "content": {"application/json": {"schema": {"$ref": "#/components/schemas/ErrorBody"}}}, +} +"""Exactly what `@returns` registers as `("default", "Error", ErrorBody)`, and what +`ExternalApi` adds to every other operation.""" + +EXPECTED_RESPONSE_CODES: dict[tuple[str, str], frozenset[str]] = { + ("get", "/workspaces/{workspace_id}/knowledge-fs/knowledge-spaces/{knowledge_space_id}/fs:ls"): frozenset( + {"200", "400", "401", "403", "404", "409", "413", "422", "503", "default"} + ), + ("get", "/workspaces/{workspace_id}/knowledge-fs/knowledge-spaces/{knowledge_space_id}/fs:tree"): frozenset( + {"200", "400", "401", "403", "404", "409", "413", "422", "503", "default"} + ), + ("get", "/workspaces/{workspace_id}/knowledge-fs/knowledge-spaces/{knowledge_space_id}/fs:grep"): frozenset( + {"200", "400", "401", "403", "404", "409", "413", "422", "503", "default"} + ), + ("get", "/workspaces/{workspace_id}/knowledge-fs/knowledge-spaces/{knowledge_space_id}/fs:find"): frozenset( + {"200", "400", "401", "403", "404", "409", "413", "422", "503", "default"} + ), + ("post", "/workspaces/{workspace_id}/knowledge-fs/knowledge-spaces/{knowledge_space_id}/fs:diff"): frozenset( + {"200", "400", "401", "403", "404", "409", "413", "422", "503", "default"} + ), + ("get", "/workspaces/{workspace_id}/knowledge-fs/knowledge-spaces/{knowledge_space_id}/fs:cat"): frozenset( + {"200", "400", "401", "403", "404", "409", "413", "422", "503", "default"} + ), + ("get", "/workspaces/{workspace_id}/knowledge-fs/knowledge-spaces/{knowledge_space_id}/fs:stat"): frozenset( + {"200", "400", "401", "403", "404", "409", "413", "422", "503", "default"} + ), + ("get", "/_catalog"): frozenset({"200", "default"}), + ("get", "/workspaces/{workspace_id}/marketplace/plugins"): frozenset({"200", "422", "default"}), + ("get", "/workspaces/{workspace_id}/plugins"): frozenset({"200", "422", "default"}), + ("post", "/workspaces/{workspace_id}/plugins:install"): frozenset({"200", "422", "default"}), + ("get", "/workspaces/{workspace_id}/plugin-tasks/{task_id}"): frozenset({"200", "default"}), + ("delete", "/workspaces/{workspace_id}/plugins/{plugin_id}"): frozenset({"200", "default"}), + ("get", "/workspaces/{workspace_id}/tools"): frozenset({"200", "422", "default"}), + ("get", "/workspaces/{workspace_id}/tool-providers/{provider}"): frozenset({"200", "default"}), + ("post", "/workspaces/{workspace_id}/tool-providers/{provider}/credentials"): frozenset({"201", "422", "default"}), + ("patch", "/workspaces/{workspace_id}/tool-providers/{provider}/credentials/{credential_id}"): frozenset( + {"200", "422", "default"} + ), + ("get", "/workspaces/{workspace_id}/model-providers/{provider}"): frozenset({"200", "default"}), + ("post", "/workspaces/{workspace_id}/model-providers/{provider}/credentials"): frozenset({"201", "422", "default"}), + ("patch", "/workspaces/{workspace_id}/model-providers/{provider}/credentials/{credential_id}"): frozenset( + {"200", "422", "default"} + ), + ("get", "/workspaces/{workspace_id}/knowledge-bases"): frozenset({"200", "422", "default"}), + ("get", "/workspaces/{workspace_id}/models"): frozenset({"200", "422", "default"}), + ("post", "/workspaces/{workspace_id}/model-providers/{provider}/models/credentials"): frozenset( + {"201", "422", "default"} + ), + ("patch", "/workspaces/{workspace_id}/model-providers/{provider}/models/credentials/{credential_id}"): frozenset( + {"200", "422", "default"} + ), + ("get", "/_health"): frozenset({"200", "default"}), + ("get", "/_version"): frozenset({"200", "default"}), + ("get", "/account"): frozenset({"200", "default"}), + ("get", "/account/sessions"): frozenset({"200", "422", "default"}), + ("delete", "/account/sessions/self"): frozenset({"200", "default"}), + ("delete", "/account/sessions/{session_id}"): frozenset({"200", "default"}), + ("get", "/apps"): frozenset({"200", "422", "default"}), + ("get", "/apps/{app_id}"): frozenset({"200", "422", "default"}), + ("get", "/apps/{app_id}/app-info/workflow"): frozenset({"200", "default"}), + ("patch", "/apps/{app_id}/app-info/workflow"): frozenset({"200", "422", "default"}), + ("get", "/apps/{app_id}/service-api"): frozenset({"200", "default"}), + ("patch", "/apps/{app_id}/service-api"): frozenset({"200", "422", "default"}), + ("get", "/apps/{app_id}/app-info/advanced-chat"): frozenset({"200", "default"}), + ("patch", "/apps/{app_id}/app-info/advanced-chat"): frozenset({"200", "422", "default"}), + ("get", "/apps/{app_id}/app-info/chat"): frozenset({"200", "default"}), + ("patch", "/apps/{app_id}/app-info/chat"): frozenset({"200", "422", "default"}), + ("get", "/apps/{app_id}/app-info/agent-chat"): frozenset({"200", "default"}), + ("patch", "/apps/{app_id}/app-info/agent-chat"): frozenset({"200", "422", "default"}), + ("get", "/apps/{app_id}/app-info/completion"): frozenset({"200", "default"}), + ("patch", "/apps/{app_id}/app-info/completion"): frozenset({"200", "422", "default"}), + ("get", "/apps/{app_id}/app-info/agent"): frozenset({"200", "default"}), + ("patch", "/apps/{app_id}/app-info/agent"): frozenset({"200", "422", "default"}), + ("get", "/apps/{app_id}/service-api/agent"): frozenset({"200", "default"}), + ("patch", "/apps/{app_id}/service-api/agent"): frozenset({"200", "422", "default"}), + ("get", "/apps/{app_id}/dependencies:check"): frozenset({"200", "default"}), + ("get", "/apps/{app_id}/dsl"): frozenset({"200", "422", "default"}), + ("post", "/apps/{app_id}/files"): frozenset({"201", "400", "401", "413", "415", "422", "default"}), + ("get", "/apps/{app_id}/human-input-forms/{form_token}"): frozenset({"200", "default"}), + ("post", "/apps/{app_id}/human-input-forms/{form_token}:submit"): frozenset({"200", "422", "default"}), + ("get", "/apps/{app_id}/tasks/{task_id}/events"): frozenset({"200", "422", "default"}), + ("post", "/apps/{app_id}/tasks/{task_id}:stop"): frozenset({"200", "default"}), + ("post", "/apps/{app_id}/workflow:run"): frozenset({"200", "422", "default"}), + ("post", "/apps/{app_id}/chat:run"): frozenset({"200", "422", "default"}), + ("post", "/apps/{app_id}/advanced-chat:run"): frozenset({"200", "422", "default"}), + ("post", "/apps/{app_id}/completion:run"): frozenset({"200", "422", "default"}), + ("post", "/apps/{app_id}/draft/workflow:run"): frozenset({"200", "422", "default"}), + ("post", "/apps/{app_id}/draft/advanced-chat:run"): frozenset({"200", "422", "default"}), + ("get", "/apps/{app_id}/runs"): frozenset({"200", "422", "default"}), + ("get", "/apps/{app_id}/runs/{run_id}"): frozenset({"200", "default"}), + ("get", "/apps/{app_id}/runs/{run_id}/nodes"): frozenset({"200", "default"}), + ("post", "/apps/{app_id}:publish"): frozenset({"200", "422", "default"}), + ("get", "/apps/{app_id}/versions"): frozenset({"200", "422", "default"}), + ("post", "/apps/{app_id}/versions/{version_id}:restore"): frozenset({"200", "default"}), + ("get", "/apps/{app_id}/release:check"): frozenset({"200", "default"}), + ("get", "/apps/{app_id}/env"): frozenset({"200", "default"}), + ("put", "/apps/{app_id}/env/{env_id}"): frozenset({"200", "422", "default"}), + ("delete", "/apps/{app_id}/env/{env_id}"): frozenset({"200", "default"}), + # The catalog and the five device-flow rows document their 200 with a raw + # `openapi_ns.response` rather than `@returns`; `ExternalApi` still gives them the + # `default` error every operation documents, with the same `ErrorBody` schema. + # Approve and deny run Console account admission; its setup and account-initialization + # errors are generic on that surface (handled centrally by the client), so the guards + # document nothing and the open `default` error covers them. Code, token, approve and + # deny declare the RFC 8628-style `DeviceFlowErrorResponse` for the statuses their flow + # errors use. + ("post", "/oauth/device/approve"): frozenset({"200", "404", "409", "503", "default"}), + ("post", "/oauth/device/code"): frozenset({"200", "400", "default"}), + ("post", "/oauth/device/deny"): frozenset({"200", "404", "409", "503", "default"}), + ("get", "/oauth/device/lookup"): frozenset({"200", "default"}), + ("post", "/oauth/device/token"): frozenset({"200", "400", "default"}), + ("get", "/node-types"): frozenset({"200", "default"}), + ("get", "/node-types/{node_type}"): frozenset({"200", "default"}), + ("post", "/apps/{app_id}/draft/workflow/nodes/{node_id}:run"): frozenset({"200", "422", "default"}), + ("post", "/apps/{app_id}/draft/advanced-chat/nodes/{node_id}:run"): frozenset({"200", "422", "default"}), + ("post", "/workspaces/{workspace_id}/apps/workflow"): frozenset({"201", "422", "default"}), + ("post", "/workspaces/{workspace_id}/apps/advanced-chat"): frozenset({"201", "422", "default"}), + ("get", "/permitted-external-apps"): frozenset({"200", "422", "default"}), + ("get", "/permitted-external-apps/{app_id}"): frozenset({"200", "422", "default"}), + ("get", "/workspaces"): frozenset({"200", "422", "default"}), + ("get", "/workspaces/{workspace_id}"): frozenset({"200", "default"}), + ("post", "/workspaces/{workspace_id}/apps/imports"): frozenset({"200", "202", "400", "422", "default"}), + ("post", "/workspaces/{workspace_id}/apps/imports:check"): frozenset({"200", "422", "default"}), + ("post", "/workspaces/{workspace_id}/apps/imports/{import_id}:confirm"): frozenset({"200", "400", "default"}), + ("get", "/workspaces/{workspace_id}/members"): frozenset({"200", "422", "default"}), + ("post", "/workspaces/{workspace_id}/members"): frozenset({"201", "422", "default"}), + ("delete", "/workspaces/{workspace_id}/members/{member_id}"): frozenset({"200", "default"}), + ("patch", "/workspaces/{workspace_id}/members/{member_id}"): frozenset({"200", "422", "default"}), + ("get", "/apps/{app_id}/webapp/workflow"): frozenset({"200", "default"}), + ("patch", "/apps/{app_id}/webapp/workflow"): frozenset({"200", "422", "default"}), + ("post", "/apps/{app_id}/webapp:reset"): frozenset({"200", "default"}), + ("get", "/apps/{app_id}/webapp-access"): frozenset({"200", "default"}), + ("put", "/apps/{app_id}/webapp-access"): frozenset({"200", "422", "default"}), + ("get", "/apps/{app_id}/webapp/advanced-chat"): frozenset({"200", "default"}), + ("patch", "/apps/{app_id}/webapp/advanced-chat"): frozenset({"200", "422", "default"}), + ("get", "/apps/{app_id}/webapp/chat"): frozenset({"200", "default"}), + ("patch", "/apps/{app_id}/webapp/chat"): frozenset({"200", "422", "default"}), + ("get", "/apps/{app_id}/webapp/agent-chat"): frozenset({"200", "default"}), + ("patch", "/apps/{app_id}/webapp/agent-chat"): frozenset({"200", "422", "default"}), + ("get", "/apps/{app_id}/webapp/completion"): frozenset({"200", "default"}), + ("patch", "/apps/{app_id}/webapp/completion"): frozenset({"200", "422", "default"}), + ("get", "/apps/{app_id}/webapp/agent"): frozenset({"200", "default"}), + ("patch", "/apps/{app_id}/webapp/agent"): frozenset({"200", "422", "default"}), + ("post", "/apps/{app_id}/webapp/agent:reset"): frozenset({"200", "default"}), + ("get", "/apps/{app_id}/webapp-access/agent"): frozenset({"200", "default"}), + ("put", "/apps/{app_id}/webapp-access/agent"): frozenset({"200", "422", "default"}), +} + +_HTTP_METHODS = frozenset({"get", "post", "put", "patch", "delete", "head", "options", "trace"}) + + +def _operations(document: dict[str, object]) -> Iterator[tuple[tuple[str, str], dict[str, object]]]: + paths = document["paths"] + assert isinstance(paths, dict) + for path, item in paths.items(): + assert isinstance(item, dict) + for method, operation in item.items(): + if method in _HTTP_METHODS: + assert isinstance(operation, dict) + yield (method, path), operation + + +class TestAuthMatrixCases: + @staticmethod + @pytest.fixture + def token_rows() -> dict[str, ResolvedRow]: + return {} + + @staticmethod + @pytest.fixture(scope="module") + def _matrix_app() -> Flask: + """Register the matrix's unchanging routes once; each case gets a fresh client.""" + app = create_flask_app_with_configs() + app.config["TESTING"] = True + app.secret_key = "openapi-auth-matrix" + LoginManager(app) + app.register_blueprint(openapi_bp) + return app + + @staticmethod + @pytest.fixture + def matrix_app( + _matrix_app: Flask, monkeypatch: pytest.MonkeyPatch, account_application_services: ApplicationServices + ) -> Iterator[Flask]: + """Keep rate-limit overrides and the admission probe scoped to each case.""" + + monkeypatch.setitem(_matrix_app.extensions, "application_services", account_application_services) + monkeypatch.setattr( + rate_limit_module, + "LIMIT_BEARER_PER_TOKEN", + rate_limit_module.RateLimit( + 0, + rate_limit_module.LIMIT_BEARER_PER_TOKEN.window, + rate_limit_module.LIMIT_BEARER_PER_TOKEN.scopes, + ), + ) + + user_logged_in.connect(_admit_on_mount) + try: + yield _matrix_app + finally: + user_logged_in.disconnect(_admit_on_mount) + + @staticmethod + @pytest.fixture + def world(sqlite_session_factory: sessionmaker[Session], token_rows: dict[str, ResolvedRow]) -> World: + """One workspace with an owner, a normal member and an outsider; two apps.""" + built = World( + workspace_id=str(uuid.uuid4()), + other_workspace_id=str(uuid.uuid4()), + app_id=str(uuid.uuid4()), + disabled_app_id=str(uuid.uuid4()), + member_account_id=str(uuid.uuid4()), + low_role_account_id=str(uuid.uuid4()), + outsider_account_id=str(uuid.uuid4()), + ) + + def account(account_id: str, email: str) -> Account: + row = Account(name=email, email=email, avatar="", status=AccountStatus.ACTIVE) + row.id = account_id + return row + + def application(app_id: str, *, enable_api: bool) -> App: + return App( + id=app_id, + tenant_id=built.workspace_id, + name="matrix app", + description="", + mode=_auth_matrix_AppMode.WORKFLOW, + enable_site=False, + enable_api=enable_api, + ) + + workspace = Tenant(name="matrix workspace") + workspace.id = built.workspace_id + other_workspace = Tenant(name="other workspace") + other_workspace.id = built.other_workspace_id + + rows: list[object] = [ + workspace, + other_workspace, + account(built.member_account_id, "owner@example.com"), + account(built.low_role_account_id, "normal@example.com"), + account(built.outsider_account_id, "outsider@example.com"), + application(built.app_id, enable_api=True), + application(built.disabled_app_id, enable_api=False), + TenantAccountJoin( + tenant_id=built.workspace_id, + account_id=built.member_account_id, + role=TenantAccountRole.OWNER, + ), + TenantAccountJoin( + tenant_id=built.workspace_id, + account_id=built.low_role_account_id, + role=TenantAccountRole.NORMAL, + ), + ] + + def mint(bearer: Bearer, prefix: str, *, account_id: str | None, email: str) -> None: + """Register a bearer with the fake resolver *and* persist the session row it + names, so the access service's ownership check has something real to match against. + """ + raw = prefix + uuid.uuid4().hex + token_id = uuid.uuid4() + issuer = "dify:account" if account_id else "https://idp.example" + token_rows[sha256_hex(raw)] = ResolvedRow( + subject_email=email, + subject_issuer=issuer, + account_id=uuid.UUID(account_id) if account_id else None, + client_id="difyctl", + token_id=token_id, + expires_at=None, + ) + row = OAuthAccessToken( + token_hash=sha256_hex(raw), + prefix=prefix, + account_id=account_id, + subject_email=email, + subject_issuer=issuer, + client_id="difyctl", + device_label="matrix", + expires_at=datetime.now(UTC) + timedelta(days=365), + ) + row.id = str(token_id) + rows.append(row) + built.tokens[bearer] = raw + built.session_ids[bearer] = str(token_id) + + mint( + Bearer.ACCOUNT_MEMBER, + TokenType.OAUTH_ACCOUNT.prefix, + account_id=built.member_account_id, + email="owner@example.com", + ) + mint( + Bearer.ACCOUNT_LOW_ROLE, + TokenType.OAUTH_ACCOUNT.prefix, + account_id=built.low_role_account_id, + email="normal@example.com", + ) + mint( + Bearer.ACCOUNT_OUTSIDER, + TokenType.OAUTH_ACCOUNT.prefix, + account_id=built.outsider_account_id, + email="outsider@example.com", + ) + mint(Bearer.EXTERNAL, TokenType.OAUTH_EXTERNAL_SSO.prefix, account_id=None, email="external@example.com") + with sqlite_session_factory.begin() as session: + session.add_all(rows) + return built + + @staticmethod + @pytest.mark.parametrize( + ("route", "case", "expected"), + ROWS, + ids=[f"{route.id}-{case.value}" for route, case, _ in ROWS], + ) + def test_allow_deny_matrix( + route: Route, + case: Case, + expected: Expect, + matrix_app: Flask, + world: World, + token_rows: dict[str, ResolvedRow], + monkeypatch: pytest.MonkeyPatch, + ) -> None: + response = _run_case( + route=route, + scenario=SCENARIOS[case], + app=matrix_app, + world=world, + token_rows=token_rows, + monkeypatch=monkeypatch, + ) + + body = response.get_json(silent=True) or {} + assert response.status_code == expected.status, f"{route.id}/{case.value}: body={body}" + if expected.message is not None: + assert body.get("message") == expected.message + + @staticmethod + def test_registered_openapi_routes_match_the_matrix(matrix_app: Flask) -> None: + """Every guarded /openapi/v1 route is in the table, and nothing else is. + + Guarded-ness is derived from `view.__spec__`, not from a hand-kept list of + exemptions — so a route that gains a guard without gaining a row fails here, + and so does a guarded route that quietly stops carrying a spec. Neither can + be silenced by editing this file. + """ + guarded: set[tuple[str, str]] = set() + unguarded: set[tuple[str, str]] = set() + for rule in matrix_app.url_map.iter_rules(): + if not str(rule).startswith("/openapi/v1") or str(rule) == "/openapi/v1/": + continue + for method in rule.methods or set(): + if method in {"HEAD", "OPTIONS"}: + continue + entry = (method, str(rule)) + target = guarded if _endpoint_spec(matrix_app, rule, method) is not None else unguarded + target.add(entry) + + expected = {(route.method, _rule_path(route)) for route in ROUTES} + assert guarded == expected + # The remainder is the device-flow, documentation and catalog surface. Its size is + # pinned rather than enumerated: a list of exemptions rots silently — the one this + # replaced still named `swagger.json`, a route that is not registered at all. + assert len(unguarded) == 14 + + @staticmethod + @pytest.mark.parametrize("route", ROUTES, ids=[route.id for route in ROUTES]) + def test_route_declares_expected_requirements(route: Route, matrix_app: Flask) -> None: + """The guard that lets `MATRIX` run one representative per tuple: a route + that drifts from the tuple its representative was chosen for fails here. + """ + spec = _spec_for(matrix_app, route) + assert spec is not None, route.id + assert [*map(_canonical, spec.requirements)] == [*map(_canonical, DECLARED[route.id])] + + @staticmethod + @pytest.fixture + def openapi_document(config_overrides: Callable[..., None]) -> dict[str, object]: + config_overrides(SWAGGER_UI_ENABLED=True) + app = Flask(__name__) + app.config["TESTING"] = True + app.register_blueprint(openapi_bp) + response = app.test_client().get("/openapi/v1/openapi.json") + assert response.status_code == 200 + return response.get_json() + + @staticmethod + def test_file_bearing_bodies_are_documented_as_multipart(openapi_document: dict[str, object]) -> None: + """The exported contract names the wire form `_multipart.py` accepts: a body of only + files is multipart alone, a run body is JSON or multipart with its JSON-text parts marked.""" + paths = openapi_document["paths"] + assert isinstance(paths, dict) + upload = paths["/apps/{app_id}/files"]["post"]["requestBody"]["content"] + assert set(upload) == {"multipart/form-data"} + assert "encoding" not in upload["multipart/form-data"] + run = paths["/apps/{app_id}/chat:run"]["post"]["requestBody"]["content"] + assert set(run) == {"application/json", "multipart/form-data"} + assert run["multipart/form-data"]["schema"] == run["application/json"]["schema"] + encoding = run["multipart/form-data"]["encoding"] + assert {"inputs", "query"} <= set(encoding) + assert {"files", "attachments"}.isdisjoint(encoding) + assert encoding["inputs"] == {"contentType": "application/json"} + + @staticmethod + def test_openapi_document_operations_and_response_codes(openapi_document: dict[str, object]) -> None: + """Every operation's response codes are pinned exactly. + + `@endpoint` routes through `@returns`, which registers `default` -> `ErrorBody`, + so a route that stops emitting the error schema (or one that starts) moves a row + here and has to be re-pinned deliberately. + """ + seen = dict(_operations(openapi_document)) + assert set(seen) == set(EXPECTED_RESPONSE_CODES) + for key, operation in seen.items(): + responses = operation.get("responses", {}) + assert isinstance(responses, dict) + assert frozenset(responses) == EXPECTED_RESPONSE_CODES[key], ( + f"{key}: {sorted(responses)} != {sorted(EXPECTED_RESPONSE_CODES[key])}" + ) + + @staticmethod + def test_every_default_response_is_the_one_returns_registers(openapi_document: dict[str, object]) -> None: + """A hand-written `default` pointing at some other schema would still satisfy + the code sets above, which compare codes and not bodies. + """ + for key, operation in _operations(openapi_document): + responses = operation["responses"] + assert isinstance(responses, dict) + if "default" in responses: + assert responses["default"] == ERROR_DEFAULT_RESPONSE, key diff --git a/api/tests/unit_tests/controllers/openapi/auth/test_resource_access.py b/api/tests/unit_tests/controllers/openapi/auth/test_resource_access.py index 3c3ff9fc858..0e98ac53a28 100644 --- a/api/tests/unit_tests/controllers/openapi/auth/test_resource_access.py +++ b/api/tests/unit_tests/controllers/openapi/auth/test_resource_access.py @@ -267,7 +267,7 @@ def test_stop_task_rejects_other_token_owner( data = call(flask_app, app_id) end_user.return_value.id = "machine-user" from services.app_task_service import AppTaskControlService - from tests.unit_tests.services.test_app_task_service import _StopRedis + from tests.unit_tests.services.test_app_task_service import _AppTaskServiceStopredis as _StopRedis redis = _StopRedis(values={"generate_task_belong:task": b"end-user-other-user"}) with patch( diff --git a/api/tests/unit_tests/controllers/openapi/auth/test_router.py b/api/tests/unit_tests/controllers/openapi/auth/test_router.py index 4359a9e065c..581c73ba9bb 100644 --- a/api/tests/unit_tests/controllers/openapi/auth/test_router.py +++ b/api/tests/unit_tests/controllers/openapi/auth/test_router.py @@ -26,8 +26,7 @@ from models.oauth import OAuthAccessToken from services.entities.feature_entities import ( LicenseStatus, ) - -from ._world import ( +from tests.helpers.unit.controllers.openapi.auth._world import ( ACCOUNT_ID, TOKEN_ID, make_account, diff --git a/api/tests/unit_tests/controllers/openapi/auth/test_subjects.py b/api/tests/unit_tests/controllers/openapi/auth/test_subjects.py index 63c87b0d42b..46dc2ea3074 100644 --- a/api/tests/unit_tests/controllers/openapi/auth/test_subjects.py +++ b/api/tests/unit_tests/controllers/openapi/auth/test_subjects.py @@ -14,8 +14,7 @@ from models.account import TenantAccountRole from models.enums import EndUserType from repositories.app_scoped_end_user_repository import AppScopedEndUserRepo from services.app_scoped_end_user_service import AppScopedEndUserService - -from ._world import ( +from tests.helpers.unit.controllers.openapi.auth._world import ( ACCOUNT_ID, APP_ID, SSO_EMAIL, diff --git a/api/tests/unit_tests/controllers/openapi/test_app_describe_builder.py b/api/tests/unit_tests/controllers/openapi/test__app_response.py similarity index 85% rename from api/tests/unit_tests/controllers/openapi/test_app_describe_builder.py rename to api/tests/unit_tests/controllers/openapi/test__app_response.py index 79298358e4b..24a9ad89a32 100644 --- a/api/tests/unit_tests/controllers/openapi/test_app_describe_builder.py +++ b/api/tests/unit_tests/controllers/openapi/test__app_response.py @@ -1,3 +1,5 @@ +from __future__ import annotations + from collections.abc import Callable from datetime import datetime from typing import Protocol, cast @@ -14,7 +16,7 @@ from machinery.context import AppRequestContext from models.model import AppMode from services.app.query_service import AppDescription from services.app_definition_query_service import AppParameterConfig -from services.entities.app_entities import AppSummary +from services.entities.app_entities import SUPPORTED_APP_TYPES, AppSummary, SupportedAppType from tests.unit_tests.controllers.conftest import ControllerTestServices @@ -125,3 +127,24 @@ def test_account_describe_unknown_mode_has_no_hints( resp = _describe(app, app_query_services, monkeypatch, "mode-from-a-newer-version") assert resp.hints == [] assert resp.info is not None + + +"""Unit tests for SupportedAppType — the listable subset of AppMode that the +openapi `app` face (`get app`) exposes and the CLI `--mode` whitelist derives from. +""" + + +def test_supported_app_type_is_the_listable_subset_of_app_mode() -> None: + """SupportedAppType (and the derived SUPPORTED_APP_TYPES tuple) is exactly the + curated, listable subset of AppMode; non-app/runtime modes stay out.""" + assert {t.value for t in SupportedAppType} == { + "completion", + "chat", + "advanced-chat", + "workflow", + "agent-chat", + } + assert set(SUPPORTED_APP_TYPES) <= set(AppMode) + assert AppMode.AGENT not in SUPPORTED_APP_TYPES + assert AppMode.RAG_PIPELINE not in SUPPORTED_APP_TYPES + assert AppMode.CHANNEL not in SUPPORTED_APP_TYPES diff --git a/api/tests/unit_tests/controllers/openapi/test_app_settings.py b/api/tests/unit_tests/controllers/openapi/test__app_settings.py similarity index 99% rename from api/tests/unit_tests/controllers/openapi/test_app_settings.py rename to api/tests/unit_tests/controllers/openapi/test__app_settings.py index 38e091356da..a8ca2b1a415 100644 --- a/api/tests/unit_tests/controllers/openapi/test_app_settings.py +++ b/api/tests/unit_tests/controllers/openapi/test__app_settings.py @@ -21,7 +21,7 @@ from services.enterprise.enterprise_service import EnterpriseService, WebAppSett from services.errors.enterprise import EnterpriseServiceError from services.webapp_access_adapters import EnterpriseWebAppAccessPolicyGateway from services.webapp_access_query_service import WebAppAccessQueryService -from tests.unit_tests.model_factories import make_app +from tests.helpers.unit.model_factories import make_app def _context(session: Session) -> Context: diff --git a/api/tests/unit_tests/controllers/openapi/test_audit_app_run.py b/api/tests/unit_tests/controllers/openapi/test__audit.py similarity index 100% rename from api/tests/unit_tests/controllers/openapi/test_audit_app_run.py rename to api/tests/unit_tests/controllers/openapi/test__audit.py diff --git a/api/tests/unit_tests/controllers/openapi/test_catalog.py b/api/tests/unit_tests/controllers/openapi/test__catalog.py similarity index 100% rename from api/tests/unit_tests/controllers/openapi/test_catalog.py rename to api/tests/unit_tests/controllers/openapi/test__catalog.py diff --git a/api/tests/unit_tests/controllers/openapi/test_contract.py b/api/tests/unit_tests/controllers/openapi/test__contract.py similarity index 100% rename from api/tests/unit_tests/controllers/openapi/test_contract.py rename to api/tests/unit_tests/controllers/openapi/test__contract.py diff --git a/api/tests/unit_tests/controllers/openapi/test_error_contract.py b/api/tests/unit_tests/controllers/openapi/test__errors.py similarity index 100% rename from api/tests/unit_tests/controllers/openapi/test_error_contract.py rename to api/tests/unit_tests/controllers/openapi/test__errors.py diff --git a/api/tests/unit_tests/controllers/openapi/test_files_merge.py b/api/tests/unit_tests/controllers/openapi/test__files.py similarity index 100% rename from api/tests/unit_tests/controllers/openapi/test_files_merge.py rename to api/tests/unit_tests/controllers/openapi/test__files.py diff --git a/api/tests/unit_tests/controllers/openapi/test_hints.py b/api/tests/unit_tests/controllers/openapi/test__hints.py similarity index 100% rename from api/tests/unit_tests/controllers/openapi/test_hints.py rename to api/tests/unit_tests/controllers/openapi/test__hints.py diff --git a/api/tests/unit_tests/controllers/openapi/test_input_schema.py b/api/tests/unit_tests/controllers/openapi/test__input_schema.py similarity index 100% rename from api/tests/unit_tests/controllers/openapi/test_input_schema.py rename to api/tests/unit_tests/controllers/openapi/test__input_schema.py diff --git a/api/tests/unit_tests/controllers/openapi/test_models.py b/api/tests/unit_tests/controllers/openapi/test__models.py similarity index 100% rename from api/tests/unit_tests/controllers/openapi/test_models.py rename to api/tests/unit_tests/controllers/openapi/test__models.py diff --git a/api/tests/unit_tests/controllers/openapi/test_multipart.py b/api/tests/unit_tests/controllers/openapi/test__multipart.py similarity index 100% rename from api/tests/unit_tests/controllers/openapi/test_multipart.py rename to api/tests/unit_tests/controllers/openapi/test__multipart.py diff --git a/api/tests/unit_tests/controllers/openapi/test_app_describe_query.py b/api/tests/unit_tests/controllers/openapi/test_app_describe_query.py deleted file mode 100644 index a6abdc95eb4..00000000000 --- a/api/tests/unit_tests/controllers/openapi/test_app_describe_query.py +++ /dev/null @@ -1,48 +0,0 @@ -"""Unit tests for AppDescribeQuery (`?fields=` allow-list).""" - -from __future__ import annotations - -import pytest -from pydantic import ValidationError - -from controllers.openapi.apps import AppDescribeQuery - - -def test_no_fields_returns_none() -> None: - q = AppDescribeQuery.model_validate({}) - assert q.fields is None - - -def test_empty_string_returns_none() -> None: - q = AppDescribeQuery.model_validate({"fields": ""}) - assert q.fields is None - - -def test_single_field() -> None: - q = AppDescribeQuery.model_validate({"fields": "info"}) - assert q.fields == {"info"} - - -def test_comma_list() -> None: - q = AppDescribeQuery.model_validate({"fields": "info,parameters"}) - assert q.fields == {"info", "parameters"} - - -def test_whitespace_tolerant() -> None: - q = AppDescribeQuery.model_validate({"fields": " info , input_schema "}) - assert q.fields == {"info", "input_schema"} - - -def test_unknown_member_rejected() -> None: - with pytest.raises(ValidationError): - AppDescribeQuery.model_validate({"fields": "garbage"}) - - -def test_unknown_among_known_rejected() -> None: - with pytest.raises(ValidationError): - AppDescribeQuery.model_validate({"fields": "info,garbage"}) - - -def test_extra_param_forbidden() -> None: - with pytest.raises(ValidationError): - AppDescribeQuery.model_validate({"fields": "info", "page": "1"}) diff --git a/api/tests/unit_tests/controllers/openapi/test_app_discovery.py b/api/tests/unit_tests/controllers/openapi/test_app_discovery.py deleted file mode 100644 index d208b1be024..00000000000 --- a/api/tests/unit_tests/controllers/openapi/test_app_discovery.py +++ /dev/null @@ -1,87 +0,0 @@ -"""Transport mappings over stable admitted contexts.""" - -from collections.abc import Callable -from typing import Protocol, cast - -import pytest -from flask import Flask -from werkzeug.exceptions import NotFound, ServiceUnavailable - -from controllers.openapi._models import ( - AppDescribeQuery, - AppDescribeResponse, - PermittedExternalAppsListQuery, - PermittedExternalAppsListResponse, -) -from controllers.openapi.apps import AppDescribeApi -from controllers.openapi.apps_permitted_external import PermittedExternalAppDescribeApi, PermittedExternalAppsListApi -from machinery.context import AppRequestContext -from models.model import AppMode -from services.app.query_service import AppDescription -from services.entities.app_entities import AppSummary -from services.errors.app import AppDiscoveryNotFoundError, PermittedAppsUnavailableError -from tests.unit_tests.controllers.conftest import ControllerTestServices - - -class Endpoint[T](Protocol): - __handler__: Callable[..., T] - - -@pytest.mark.parametrize("api", [AppDescribeApi(), PermittedExternalAppDescribeApi()]) -def test_describe_passes_admitted_scope_and_projects_response( - app: Flask, - app_query_services: ControllerTestServices, - monkeypatch: pytest.MonkeyPatch, - api: AppDescribeApi | PermittedExternalAppDescribeApi, -) -> None: - context = AppRequestContext(tenant_id="workspace-id", app_id="app-id") - calls: list[tuple[AppRequestContext, set[str] | None]] = [] - - def describe(ctx: AppRequestContext, fields: set[str] | None) -> AppDescription: - calls.append((ctx, fields)) - return AppDescription( - AppSummary("app-id", "workspace-id", "App", "", AppMode.CHAT, "normal", None, None), True, None - ) - - monkeypatch.setattr(app_query_services.apps.discovery, "describe", describe) - with app.test_request_context(): - result = cast(Endpoint[AppDescribeResponse], api.get).__handler__( - api, context, "app-id", query=AppDescribeQuery.model_validate({"fields": "info"}) - ) - assert calls == [(context, {"info"})] - assert result.info is not None - assert result.info.id == "app-id" - assert result.parameters is None - assert result.input_schema is None - - -@pytest.mark.parametrize("api", [AppDescribeApi(), PermittedExternalAppDescribeApi()]) -def test_describe_translates_disappearing_app_to_404( - app: Flask, - app_query_services: ControllerTestServices, - monkeypatch: pytest.MonkeyPatch, - api: AppDescribeApi | PermittedExternalAppDescribeApi, -) -> None: - def missing(*_args: object) -> AppDescription: - raise AppDiscoveryNotFoundError("app not found") - - monkeypatch.setattr(app_query_services.apps.discovery, "describe", missing) - context = AppRequestContext(tenant_id="workspace-id", app_id="app-id") - with app.test_request_context(), pytest.raises(NotFound, match="app not found"): - cast(Endpoint[AppDescribeResponse], api.get).__handler__(api, context, "app-id", query=AppDescribeQuery()) - - -def test_external_failure_is_translated_at_http_boundary( - app: Flask, - app_query_services: ControllerTestServices, - monkeypatch: pytest.MonkeyPatch, -) -> None: - def unavailable(*_args: object) -> None: - raise PermittedAppsUnavailableError("permitted_apps_unavailable") - - monkeypatch.setattr(app_query_services.apps.discovery, "list_permitted_apps", unavailable) - api = PermittedExternalAppsListApi() - with app.test_request_context(), pytest.raises(ServiceUnavailable, match="permitted_apps_unavailable"): - cast(Endpoint[PermittedExternalAppsListResponse], api.get).__handler__( - api, query=PermittedExternalAppsListQuery() - ) diff --git a/api/tests/unit_tests/controllers/openapi/test_app_list_query.py b/api/tests/unit_tests/controllers/openapi/test_app_list_query.py deleted file mode 100644 index c6559e11684..00000000000 --- a/api/tests/unit_tests/controllers/openapi/test_app_list_query.py +++ /dev/null @@ -1,115 +0,0 @@ -"""Unit tests for AppListQuery — the /apps query-param validator. - -Runs against the model directly, not the HTTP layer. Pins: -- defaults match the plan (page=1, limit=20). -- workspace_id is required. -- numeric bounds enforced (page >= 1, limit in [1, MAX_PAGE_LIMIT]). -- mode validates against the SupportedAppType enum (listable app types only). -- name has a length cap. -""" - -from __future__ import annotations - -import pytest -from pydantic import ValidationError - -from controllers.openapi._models import MAX_PAGE_LIMIT -from controllers.openapi.apps import AppListQuery -from services.app_service import AppListParams - -from ._mode_constants import LISTABLE_MODES, NON_LISTABLE_MODES - -WS_ID = "00000000-0000-0000-0000-000000000001" - - -def test_defaults() -> None: - q = AppListQuery.model_validate({"workspace_id": WS_ID}) - assert q.workspace_id == WS_ID - assert q.page == 1 - assert q.limit == 20 - assert q.mode is None - assert q.name is None - - -def test_workspace_id_required() -> None: - with pytest.raises(ValidationError): - AppListQuery.model_validate({}) - - -def test_page_must_be_positive() -> None: - with pytest.raises(ValidationError): - AppListQuery.model_validate({"workspace_id": WS_ID, "page": 0}) - with pytest.raises(ValidationError): - AppListQuery.model_validate({"workspace_id": WS_ID, "page": -1}) - - -def test_page_rejects_non_integer_string() -> None: - with pytest.raises(ValidationError): - AppListQuery.model_validate({"workspace_id": WS_ID, "page": "abc"}) - - -def test_limit_must_be_positive() -> None: - with pytest.raises(ValidationError): - AppListQuery.model_validate({"workspace_id": WS_ID, "limit": 0}) - with pytest.raises(ValidationError): - AppListQuery.model_validate({"workspace_id": WS_ID, "limit": -1}) - - -def test_limit_caps_at_max_page_limit() -> None: - # Boundary accepts. - q = AppListQuery.model_validate({"workspace_id": WS_ID, "limit": MAX_PAGE_LIMIT}) - assert q.limit == MAX_PAGE_LIMIT - - # Just over rejects. - with pytest.raises(ValidationError): - AppListQuery.model_validate({"workspace_id": WS_ID, "limit": MAX_PAGE_LIMIT + 1}) - - -@pytest.mark.parametrize("mode", LISTABLE_MODES) -def test_mode_accepts_listable_app_types(mode: str) -> None: - q = AppListQuery.model_validate({"workspace_id": WS_ID, "mode": mode}) - assert q.mode is not None - assert q.mode.value == mode - - -def test_mode_rejects_unknown_value() -> None: - with pytest.raises(ValidationError): - AppListQuery.model_validate({"workspace_id": WS_ID, "mode": "not-a-mode"}) - - -@pytest.mark.parametrize("mode", NON_LISTABLE_MODES) -def test_mode_rejects_non_listable_app_modes(mode: str) -> None: - """rag-pipeline (a knowledge Pipeline), channel (unused) and agent (roster-owned) - are AppMode members but not standalone listable apps — the `app` face rejects them.""" - with pytest.raises(ValidationError): - AppListQuery.model_validate({"workspace_id": WS_ID, "mode": mode}) - - -def test_name_length_capped() -> None: - AppListQuery.model_validate({"workspace_id": WS_ID, "name": "x" * 200}) - with pytest.raises(ValidationError): - AppListQuery.model_validate({"workspace_id": WS_ID, "name": "x" * 201}) - - -def test_all_fields_accept_valid_values() -> None: - """Pin the happy-path acceptance for every field in one place.""" - q = AppListQuery.model_validate( - { - "workspace_id": WS_ID, - "page": 5, - "limit": 50, - "mode": "workflow", - "name": "search", - } - ) - assert q.workspace_id == WS_ID - assert q.page == 5 - assert q.limit == 50 - assert q.mode is not None - assert q.mode.value == "workflow" - assert q.name == "search" - - -def test_max_page_limit_is_accepted_by_the_app_service() -> None: - # The catalog advertises MAX_PAGE_LIMIT; the service model must not reject it. - assert AppListParams(limit=MAX_PAGE_LIMIT).limit == MAX_PAGE_LIMIT diff --git a/api/tests/unit_tests/controllers/openapi/test_app_run_streaming.py b/api/tests/unit_tests/controllers/openapi/test_app_run.py similarity index 83% rename from api/tests/unit_tests/controllers/openapi/test_app_run_streaming.py rename to api/tests/unit_tests/controllers/openapi/test_app_run.py index dbfc7c875ea..aadf1af3ced 100644 --- a/api/tests/unit_tests/controllers/openapi/test_app_run_streaming.py +++ b/api/tests/unit_tests/controllers/openapi/test_app_run.py @@ -1,5 +1,3 @@ -"""Tests for the /openapi/v1 run routes: per-mode bodies and handlers, task stop.""" - from __future__ import annotations import json @@ -16,7 +14,7 @@ from pydantic import BaseModel, ValidationError from sqlalchemy import event from sqlalchemy.orm import Session from werkzeug.datastructures import FileStorage -from werkzeug.exceptions import NotFound +from werkzeug.exceptions import BadRequest, NotFound, TooManyRequests from constants.oauth_bearer import TokenType from controllers.openapi import app_run @@ -28,9 +26,18 @@ from controllers.openapi._models import ( TaskStopResponse, WorkflowRunPayload, ) -from controllers.openapi.app_run import AdvancedChatRunApi, AppRunTaskStopApi, ChatRunApi, CompletionRunApi +from controllers.openapi.app_run import ( + AdvancedChatRunApi, + AppRunTaskStopApi, + ChatRunApi, + CompletionRunApi, + _translate_service_errors, +) from controllers.openapi.auth.subjects import Subject, subject_from_auth +from controllers.service_api.app.error import TriggerWorkflowServiceModeUnavailableError +from controllers.web.error import InvokeRateLimitError as InvokeRateLimitHttpError from core.app.entities.app_invoke_entities import InvokeFrom +from core.errors.error import AppInvokeQuotaExceededError from extensions.ext_application_services import ApplicationServices from graphon.enums import WorkflowExecutionStatus, WorkflowType from graphon.file import FileType @@ -40,8 +47,56 @@ from models.enums import CreatorUserRole, WorkflowRunTriggeredFrom from models.model import App, AppMode, EndUser, UploadFile from models.workflow import WorkflowRun from services.app_task_service import AppTaskControlService -from tests.unit_tests.model_factories import make_upload_file -from tests.unit_tests.services.test_app_task_service import _StopRedis +from services.errors.app import ( + TriggerWorkflowServiceModeUnavailableError as TriggerWorkflowServiceModeUnavailableServiceError, +) +from services.errors.llm import InvokeRateLimitError +from tests.helpers.unit.model_factories import make_upload_file +from tests.unit_tests.services.test_app_task_service import _AppTaskServiceStopredis as _StopRedis + +"""The openapi run boundary maps internal rate-limit exceptions to canonical 429s. + +Both render through the ErrorBody formatter: TooManyRequests -> code "too_many_requests" +(retryable throttle), InvokeRateLimitError -> code "rate_limit_error" (quota). +""" + + +def test_translate_maps_app_concurrency_to_too_many_requests(): + # Regression guard: this used to fall through to a 500 (it was not caught here). + with pytest.raises(TooManyRequests) as exc: + with _translate_service_errors(): + raise AppInvokeQuotaExceededError("internal: client_id=abc max=10") + assert exc.value.code == 429 + + +def test_translate_maps_workflow_quota_to_rate_limit_error(): + with pytest.raises(InvokeRateLimitHttpError) as exc: + with _translate_service_errors(): + raise InvokeRateLimitError("workflow quota exhausted") + assert exc.value.error_code == "rate_limit_error" + assert exc.value.code == 429 + + +def test_translate_maps_trigger_workflow_to_stable_unavailable_error(): + with pytest.raises(TriggerWorkflowServiceModeUnavailableError) as exc: + with _translate_service_errors(): + raise TriggerWorkflowServiceModeUnavailableServiceError() + assert exc.value.error_code == "trigger_workflow_service_mode_unavailable" + assert exc.value.code == 403 + + +def test_translate_maps_value_error_to_a_fixed_bad_request(): + # Regression guard: an unpublished workflow used to surface as a detail-less 500. + # The service message names internals, so the client gets a fixed text. + with pytest.raises(BadRequest) as exc: + with _translate_service_errors(): + raise ValueError("variable 'secret_key' of node 42 is missing") + assert exc.value.code == 400 + assert exc.value.description == "invalid run input" + + +"""Tests for the /openapi/v1 run routes: per-mode bodies and handlers, task stop.""" + _TEST_APP_ID = str(uuid.uuid4()) _TEST_TENANT_ID = str(uuid.uuid4()) diff --git a/api/tests/unit_tests/controllers/openapi/test_app_run_rate_limit.py b/api/tests/unit_tests/controllers/openapi/test_app_run_rate_limit.py deleted file mode 100644 index 7d94bae5fab..00000000000 --- a/api/tests/unit_tests/controllers/openapi/test_app_run_rate_limit.py +++ /dev/null @@ -1,51 +0,0 @@ -"""The openapi run boundary maps internal rate-limit exceptions to canonical 429s. - -Both render through the ErrorBody formatter: TooManyRequests -> code "too_many_requests" -(retryable throttle), InvokeRateLimitError -> code "rate_limit_error" (quota). -""" - -import pytest -from werkzeug.exceptions import BadRequest, TooManyRequests - -from controllers.openapi.app_run import _translate_service_errors -from controllers.service_api.app.error import TriggerWorkflowServiceModeUnavailableError -from controllers.web.error import InvokeRateLimitError as InvokeRateLimitHttpError -from core.errors.error import AppInvokeQuotaExceededError -from services.errors.app import ( - TriggerWorkflowServiceModeUnavailableError as TriggerWorkflowServiceModeUnavailableServiceError, -) -from services.errors.llm import InvokeRateLimitError - - -def test_translate_maps_app_concurrency_to_too_many_requests(): - # Regression guard: this used to fall through to a 500 (it was not caught here). - with pytest.raises(TooManyRequests) as exc: - with _translate_service_errors(): - raise AppInvokeQuotaExceededError("internal: client_id=abc max=10") - assert exc.value.code == 429 - - -def test_translate_maps_workflow_quota_to_rate_limit_error(): - with pytest.raises(InvokeRateLimitHttpError) as exc: - with _translate_service_errors(): - raise InvokeRateLimitError("workflow quota exhausted") - assert exc.value.error_code == "rate_limit_error" - assert exc.value.code == 429 - - -def test_translate_maps_trigger_workflow_to_stable_unavailable_error(): - with pytest.raises(TriggerWorkflowServiceModeUnavailableError) as exc: - with _translate_service_errors(): - raise TriggerWorkflowServiceModeUnavailableServiceError() - assert exc.value.error_code == "trigger_workflow_service_mode_unavailable" - assert exc.value.code == 403 - - -def test_translate_maps_value_error_to_a_fixed_bad_request(): - # Regression guard: an unpublished workflow used to surface as a detail-less 500. - # The service message names internals, so the client gets a fixed text. - with pytest.raises(BadRequest) as exc: - with _translate_service_errors(): - raise ValueError("variable 'secret_key' of node 42 is missing") - assert exc.value.code == 400 - assert exc.value.description == "invalid run input" diff --git a/api/tests/unit_tests/controllers/openapi/test_apps.py b/api/tests/unit_tests/controllers/openapi/test_apps.py new file mode 100644 index 00000000000..a43b493f122 --- /dev/null +++ b/api/tests/unit_tests/controllers/openapi/test_apps.py @@ -0,0 +1,248 @@ +from __future__ import annotations + +from collections.abc import Callable +from typing import Protocol, cast + +import pytest +from flask import Flask +from pydantic import ValidationError +from werkzeug.exceptions import NotFound, ServiceUnavailable + +from controllers.openapi._models import ( + MAX_PAGE_LIMIT, + AppDescribeResponse, + PermittedExternalAppsListQuery, + PermittedExternalAppsListResponse, +) +from controllers.openapi._models import ( + AppDescribeQuery as _app_discovery_AppDescribeQuery, +) +from controllers.openapi.apps import AppDescribeApi, AppListQuery +from controllers.openapi.apps import AppDescribeQuery as _app_describe_query_AppDescribeQuery +from controllers.openapi.apps_permitted_external import PermittedExternalAppDescribeApi, PermittedExternalAppsListApi +from machinery.context import AppRequestContext +from models.model import AppMode +from services.app.query_service import AppDescription +from services.app_service import AppListParams +from services.entities.app_entities import AppSummary +from services.errors.app import AppDiscoveryNotFoundError, PermittedAppsUnavailableError +from tests.helpers.unit.controllers.openapi._mode_constants import LISTABLE_MODES, NON_LISTABLE_MODES +from tests.unit_tests.controllers.conftest import ControllerTestServices + +"""Unit tests for AppDescribeQuery (`?fields=` allow-list).""" + + +def test_no_fields_returns_none() -> None: + q = _app_describe_query_AppDescribeQuery.model_validate({}) + assert q.fields is None + + +def test_empty_string_returns_none() -> None: + q = _app_describe_query_AppDescribeQuery.model_validate({"fields": ""}) + assert q.fields is None + + +def test_single_field() -> None: + q = _app_describe_query_AppDescribeQuery.model_validate({"fields": "info"}) + assert q.fields == {"info"} + + +def test_comma_list() -> None: + q = _app_describe_query_AppDescribeQuery.model_validate({"fields": "info,parameters"}) + assert q.fields == {"info", "parameters"} + + +def test_whitespace_tolerant() -> None: + q = _app_describe_query_AppDescribeQuery.model_validate({"fields": " info , input_schema "}) + assert q.fields == {"info", "input_schema"} + + +def test_unknown_member_rejected() -> None: + with pytest.raises(ValidationError): + _app_describe_query_AppDescribeQuery.model_validate({"fields": "garbage"}) + + +def test_unknown_among_known_rejected() -> None: + with pytest.raises(ValidationError): + _app_describe_query_AppDescribeQuery.model_validate({"fields": "info,garbage"}) + + +def test_extra_param_forbidden() -> None: + with pytest.raises(ValidationError): + _app_describe_query_AppDescribeQuery.model_validate({"fields": "info", "page": "1"}) + + +"""Transport mappings over stable admitted contexts.""" + + +class Endpoint[T](Protocol): + __handler__: Callable[..., T] + + +@pytest.mark.parametrize("api", [AppDescribeApi(), PermittedExternalAppDescribeApi()]) +def test_describe_passes_admitted_scope_and_projects_response( + app: Flask, + app_query_services: ControllerTestServices, + monkeypatch: pytest.MonkeyPatch, + api: AppDescribeApi | PermittedExternalAppDescribeApi, +) -> None: + context = AppRequestContext(tenant_id="workspace-id", app_id="app-id") + calls: list[tuple[AppRequestContext, set[str] | None]] = [] + + def describe(ctx: AppRequestContext, fields: set[str] | None) -> AppDescription: + calls.append((ctx, fields)) + return AppDescription( + AppSummary("app-id", "workspace-id", "App", "", AppMode.CHAT, "normal", None, None), True, None + ) + + monkeypatch.setattr(app_query_services.apps.discovery, "describe", describe) + with app.test_request_context(): + result = cast(Endpoint[AppDescribeResponse], api.get).__handler__( + api, context, "app-id", query=_app_discovery_AppDescribeQuery.model_validate({"fields": "info"}) + ) + assert calls == [(context, {"info"})] + assert result.info is not None + assert result.info.id == "app-id" + assert result.parameters is None + assert result.input_schema is None + + +@pytest.mark.parametrize("api", [AppDescribeApi(), PermittedExternalAppDescribeApi()]) +def test_describe_translates_disappearing_app_to_404( + app: Flask, + app_query_services: ControllerTestServices, + monkeypatch: pytest.MonkeyPatch, + api: AppDescribeApi | PermittedExternalAppDescribeApi, +) -> None: + def missing(*_args: object) -> AppDescription: + raise AppDiscoveryNotFoundError("app not found") + + monkeypatch.setattr(app_query_services.apps.discovery, "describe", missing) + context = AppRequestContext(tenant_id="workspace-id", app_id="app-id") + with app.test_request_context(), pytest.raises(NotFound, match="app not found"): + cast(Endpoint[AppDescribeResponse], api.get).__handler__( + api, context, "app-id", query=_app_discovery_AppDescribeQuery() + ) + + +def test_external_failure_is_translated_at_http_boundary( + app: Flask, + app_query_services: ControllerTestServices, + monkeypatch: pytest.MonkeyPatch, +) -> None: + def unavailable(*_args: object) -> None: + raise PermittedAppsUnavailableError("permitted_apps_unavailable") + + monkeypatch.setattr(app_query_services.apps.discovery, "list_permitted_apps", unavailable) + api = PermittedExternalAppsListApi() + with app.test_request_context(), pytest.raises(ServiceUnavailable, match="permitted_apps_unavailable"): + cast(Endpoint[PermittedExternalAppsListResponse], api.get).__handler__( + api, query=PermittedExternalAppsListQuery() + ) + + +"""Unit tests for AppListQuery — the /apps query-param validator. + +Runs against the model directly, not the HTTP layer. Pins: +- defaults match the plan (page=1, limit=20). +- workspace_id is required. +- numeric bounds enforced (page >= 1, limit in [1, MAX_PAGE_LIMIT]). +- mode validates against the SupportedAppType enum (listable app types only). +- name has a length cap. +""" + + +WS_ID = "00000000-0000-0000-0000-000000000001" + + +def test_defaults() -> None: + q = AppListQuery.model_validate({"workspace_id": WS_ID}) + assert q.workspace_id == WS_ID + assert q.page == 1 + assert q.limit == 20 + assert q.mode is None + assert q.name is None + + +def test_workspace_id_required() -> None: + with pytest.raises(ValidationError): + AppListQuery.model_validate({}) + + +def test_page_must_be_positive() -> None: + with pytest.raises(ValidationError): + AppListQuery.model_validate({"workspace_id": WS_ID, "page": 0}) + with pytest.raises(ValidationError): + AppListQuery.model_validate({"workspace_id": WS_ID, "page": -1}) + + +def test_page_rejects_non_integer_string() -> None: + with pytest.raises(ValidationError): + AppListQuery.model_validate({"workspace_id": WS_ID, "page": "abc"}) + + +def test_limit_must_be_positive() -> None: + with pytest.raises(ValidationError): + AppListQuery.model_validate({"workspace_id": WS_ID, "limit": 0}) + with pytest.raises(ValidationError): + AppListQuery.model_validate({"workspace_id": WS_ID, "limit": -1}) + + +def test_limit_caps_at_max_page_limit() -> None: + # Boundary accepts. + q = AppListQuery.model_validate({"workspace_id": WS_ID, "limit": MAX_PAGE_LIMIT}) + assert q.limit == MAX_PAGE_LIMIT + + # Just over rejects. + with pytest.raises(ValidationError): + AppListQuery.model_validate({"workspace_id": WS_ID, "limit": MAX_PAGE_LIMIT + 1}) + + +@pytest.mark.parametrize("mode", LISTABLE_MODES) +def test_mode_accepts_listable_app_types(mode: str) -> None: + q = AppListQuery.model_validate({"workspace_id": WS_ID, "mode": mode}) + assert q.mode is not None + assert q.mode.value == mode + + +def test_mode_rejects_unknown_value() -> None: + with pytest.raises(ValidationError): + AppListQuery.model_validate({"workspace_id": WS_ID, "mode": "not-a-mode"}) + + +@pytest.mark.parametrize("mode", NON_LISTABLE_MODES) +def test_mode_rejects_non_listable_app_modes(mode: str) -> None: + """rag-pipeline (a knowledge Pipeline), channel (unused) and agent (roster-owned) + are AppMode members but not standalone listable apps — the `app` face rejects them.""" + with pytest.raises(ValidationError): + AppListQuery.model_validate({"workspace_id": WS_ID, "mode": mode}) + + +def test_name_length_capped() -> None: + AppListQuery.model_validate({"workspace_id": WS_ID, "name": "x" * 200}) + with pytest.raises(ValidationError): + AppListQuery.model_validate({"workspace_id": WS_ID, "name": "x" * 201}) + + +def test_all_fields_accept_valid_values() -> None: + """Pin the happy-path acceptance for every field in one place.""" + q = AppListQuery.model_validate( + { + "workspace_id": WS_ID, + "page": 5, + "limit": 50, + "mode": "workflow", + "name": "search", + } + ) + assert q.workspace_id == WS_ID + assert q.page == 5 + assert q.limit == 50 + assert q.mode is not None + assert q.mode.value == "workflow" + assert q.name == "search" + + +def test_max_page_limit_is_accepted_by_the_app_service() -> None: + # The catalog advertises MAX_PAGE_LIMIT; the service model must not reject it. + assert AppListParams(limit=MAX_PAGE_LIMIT).limit == MAX_PAGE_LIMIT diff --git a/api/tests/unit_tests/controllers/openapi/test_apps_permitted_external_query.py b/api/tests/unit_tests/controllers/openapi/test_apps_permitted_external.py similarity index 100% rename from api/tests/unit_tests/controllers/openapi/test_apps_permitted_external_query.py rename to api/tests/unit_tests/controllers/openapi/test_apps_permitted_external.py diff --git a/api/tests/unit_tests/controllers/openapi/test_auth_matrix.py b/api/tests/unit_tests/controllers/openapi/test_auth_matrix.py deleted file mode 100644 index c4e5900f581..00000000000 --- a/api/tests/unit_tests/controllers/openapi/test_auth_matrix.py +++ /dev/null @@ -1,2228 +0,0 @@ -"""The /openapi/v1 allow/deny matrix, and a snapshot of the generated document. - -The matrix is derived from what actually runs, not from what a route declares: a -route's own requirements are merged with the fixed ones its subject's pipeline -carries, a `CheckWorkspaceRole` applies only where RBAC is off and a `CheckRBACPermission` only where it is on, -and `CheckAppAccess` applies the private-app check whether -or not `webapp_auth.enabled` gates the ACL. Reading declarations alone under-counts -every one of those. - -Requests run through the real Flask blueprint with the real router, the real -pipeline and real database rows. The only substitutions are at process seams the -router already treats as pluggable: the bound `BearerAuthenticator` (so a token's -subject kind and scope set are chosen per row), the enterprise HTTP clients, and -the flask-login mount. The RBAC backend and the enterprise web-app service are -stubbed per row, so those rows pin how *this* code handles each answer they can -give — never what the real services decide. - -Cases are never gated on whether a route declares a check. A row that asserts -`ADMIT` where RBAC, a workspace role check or a licence gate does not fire is pinning the -check's *absence*: losing a check is caught by the deny rows, and gaining one is -caught by these. `files.upload` is the route the design says must never acquire an -RBAC permission, and `files.upload-rbac_on_denied` is the row that says so. - -Admission is observed as HTTP 418: a `user_logged_in` receiver raises `ImATeapot` -at the moment the pipeline mounts the caller, which is after every requirement has -passed and before the view body runs. That keeps a row from depending on whether a -view could complete (several stream SSE or invoke a workflow). One route mounts no -caller — `permitted_external.list` resolves no end user because it carries no -`app_id` — so its admission row says so and asserts the view's own 200 instead. - -Two answers in the table are worth reading twice, because they are what the code -does rather than what a route's declaration suggests: - -* `workspaces.describe` declares no membership requirement, so a non-member is - *admitted* by auth and refused by the view's own lookup. `workspaces.switch`, - one path segment away, declares one and is refused by auth. -* No token the shipped registry mints can fail `CheckScope` — `dfoa_` carries - `Scope.FULL` and `dfoe_` carries exactly the two scopes its routes ask for — so - the scope rows mint from a deliberately narrowed registry. Those rows are not - dead weight and must not be simplified away: `CheckScope` is live code on every - request, and the day a narrower token kind is minted it is the only thing - standing between that token and every route it was not scoped for. - -`Case.NON_MEMBER_AND_INSUFFICIENT_SCOPE` is the one case that combines a real -membership failure with a real scope failure on the same request. Neither -`INSUFFICIENT_SCOPE` nor `NON_MEMBER` alone proves which one a doubly-broken request -hears; membership runs EARLY and scope NORMAL, so it hears membership. - -Routes that declare the same requirement tuple answer every case the same way, -so `MATRIX` runs the cases on one representative per tuple and `DECLARED` pins -which tuple each route carries. A route that changes its tuple fails -the wiring test; a tuple that changes its answers fails its representative's rows. -""" - -from __future__ import annotations - -import contextvars -import uuid -from collections.abc import Callable, Iterator -from contextlib import ExitStack -from dataclasses import dataclass, field -from datetime import UTC, datetime, timedelta -from enum import StrEnum, auto -from functools import singledispatch -from typing import Final -from unittest.mock import patch - -import pytest -from flask import Flask -from flask_login import LoginManager, user_logged_in -from sqlalchemy.orm import Session, sessionmaker -from werkzeug.exceptions import ImATeapot -from werkzeug.routing import Rule -from werkzeug.test import TestResponse - -import libs.oauth_bearer as oauth_bearer_module -import libs.rate_limit as rate_limit_module -from app_factory import create_flask_app_with_configs -from constants.oauth_bearer import Scope, TokenType -from controllers.common.rbac import AgentBehindApp, PlainApp, RBACCheck, RBACPermission, Workspace -from controllers.openapi import bp as openapi_bp -from controllers.openapi._app_settings import REGULAR_MODES -from controllers.openapi._catalog import CATALOG_HEADER, catalog_for -from controllers.openapi._errors import ServiceApiDisabled, WebAppAccessRequiresEE -from controllers.openapi.auth import subjects -from controllers.openapi.auth.requirements import ( - CheckAppAccess, - CheckAppApiEnabled, - CheckAppMode, - CheckAppQuota, - CheckPluginInstallSetting, - CheckRBACPermission, - CheckScope, - CheckSubject, - CheckWebAppAuthEnterprise, - CheckWorkspaceInvitationQuota, - CheckWorkspaceMember, - CheckWorkspaceRole, - Requirement, -) -from controllers.openapi.auth.spec import EndpointSpec -from controllers.openapi.auth.subjects import AccountSubject, ExternalSsoSubject, ResourceAccessSubject -from controllers.openapi.human_input_form import CheckFormSurface -from enums import DeploymentEdition, WebAppAccessMode -from extensions.ext_application_services import ApplicationServices -from libs.oauth_bearer import BearerAuthenticator, ResolvedRow, sha256_hex -from models.account import Account, AccountStatus, Tenant, TenantAccountJoin, TenantAccountRole -from models.enums import EndUserType -from models.model import App, AppMode, EndUser -from models.oauth import OAuthAccessToken -from services.enterprise.enterprise_service import EnterpriseService -from services.entities.feature_entities import LicenseStatus, SystemFeatureModel -from services.system_feature_service import SystemFeatureService -from tests.unit_tests.config_override import apply_config_overrides - -ADMITTED = 418 - - -class Trait(StrEnum): - """Structural facts about a route, used only to decide which cases can reach it. - - Every member here is a fact about the *shape* of the request — is there an - `app_id` in the path, can a `dfoe_` token address this route at all. None is - read off a route's declared requirements, deliberately: gating a - case on a declared check means a route that *gains* that check during migration - simply loses the row, and gaining a check is the hazard this plan cares about - most. Cases like `RBAC_ON_DENIED` therefore run on every route an account - caller can reach, and assert `ADMIT` where the check does not exist today. - """ - - APP_SCOPED = auto() - ACCOUNT_PRIMARY = auto() - EXTERNAL_REACHABLE = auto() - ENTERPRISE_ONLY = auto() - - -class Case(StrEnum): - NO_BEARER = auto() - MEMBER = auto() - WRONG_SUBJECT = auto() - INSUFFICIENT_SCOPE = auto() - NON_MEMBER = auto() - NON_MEMBER_AND_INSUFFICIENT_SCOPE = auto() - LOW_ROLE = auto() - APP_API_DISABLED = auto() - UNKNOWN_APP = auto() - FOREIGN_WORKSPACE_QUERY = auto() - EDITION_NOT_ENTERPRISE = auto() - LICENSE_INVALID = auto() - EE_LICENSE_INVALID = auto() - EE_ACCOUNT_PUBLIC = auto() - EE_ACCOUNT_SSO_VERIFIED = auto() - EE_ACCOUNT_PRIVATE_ALL = auto() - EE_ACCOUNT_PRIVATE_PERMITTED = auto() - EE_ACCOUNT_PRIVATE_REFUSED = auto() - EE_ACCOUNT_PRIVATE_REFUSED_WEBAPP_AUTH_OFF = auto() - EE_ACCOUNT_MODE_UNRESOLVED = auto() - EE_EXTERNAL_PUBLIC = auto() - EE_EXTERNAL_SSO_VERIFIED = auto() - EE_EXTERNAL_PRIVATE_ALL = auto() - EE_EXTERNAL_PRIVATE = auto() - EE_EXTERNAL_PRIVATE_REFUSED_WEBAPP_AUTH_OFF = auto() - EE_EXTERNAL_MODE_UNRESOLVED = auto() - RBAC_ON_LOW_ROLE = auto() - RBAC_ON_DENIED = auto() - - -class Bearer(StrEnum): - NONE = auto() - ACCOUNT_MEMBER = auto() - ACCOUNT_LOW_ROLE = auto() - ACCOUNT_OUTSIDER = auto() - EXTERNAL = auto() - OTHER_SUBJECT = auto() - PRIMARY = auto() - - -@dataclass(frozen=True, slots=True) -class Route: - id: str - method: str - path: str - traits: frozenset[Trait] - query: str = "" - - -@dataclass(frozen=True, slots=True) -class Expect: - """`message` is the canonical ErrorBody message; `None` means it is not pinned.""" - - status: int - message: str | None = None - note: str = "" - - -@dataclass(frozen=True, slots=True) -class Scenario: - """The world a case runs in. `edition=None` means the route's native edition.""" - - bearer: Bearer - edition: DeploymentEdition | None = None - narrow_scopes: bool = False - app_api_enabled: bool = True - unknown_app: bool = False - foreign_workspace_query: bool = False - webapp_auth: bool = False - access_mode: WebAppAccessMode | None = WebAppAccessMode.PUBLIC - private_app_permitted: bool = True - license_status: LicenseStatus = LicenseStatus.ACTIVE - rbac_enabled: bool = False - rbac_allows: bool = True - - -ADMIT = Expect(ADMITTED) -DENY_NO_BEARER = Expect(401, "bearer required") -DENY_WRONG_SUBJECT = Expect(403, "unsupported_token_type") -DENY_SSO_NEEDS_EE = Expect(403, "external_sso_requires_ee") -DENY_SCOPE = Expect(403, "insufficient_scope") -DENY_NON_MEMBER = Expect(404, "workspace not found") -DENY_ROLE = Expect(403, "insufficient workspace role") -DENY_API_DISABLED = Expect(403, ServiceApiDisabled.description) -DENY_UNKNOWN_APP = Expect(404, "app not found") -DENY_ACCESS_MODE = Expect(403, "subject_not_allowed_for_access_mode") -DENY_MODE_UNRESOLVED = Expect(403, "app or access mode not loaded") -DENY_PRIVATE_APP = Expect(403, "user_not_allowed_for_private_app") -DENY_EDITION = Expect(404, note="endpoint-level edition gate, raised before the bearer is read") -DENY_LICENSE = Expect(403, "license_invalid") -DENY_RBAC = Expect(403, note="bare werkzeug Forbidden from enforce_rbac_checks") -ADMIT_NO_RBAC_PERMISSION = Expect( - ADMITTED, - note=( - "pins a check this route does NOT have: RBAC is on and denying, and the route still " - "admits because it declares no scene. Attaching one during migration turns this 403. " - "files.upload is the route the design says must never gain a scene" - ), -) -ADMIT_NO_WORKSPACE_ROLE = Expect( - ADMITTED, - note="pins a check this route does NOT have: RBAC is on and a NORMAL member is still admitted", -) -ADMIT_NO_LICENCE_GATE = Expect( - ADMITTED, - note="pins a check this route does NOT have: the licence is dead and the route still admits", -) -ADMIT_NO_MOUNT = Expect( - 200, - note="external subject mounts no caller on an app-less route, so admission shows as the view's own 200", -) - - -_RUN_TRAITS = frozenset({Trait.ACCOUNT_PRIMARY, Trait.APP_SCOPED, Trait.EXTERNAL_REACHABLE}) - -ROUTES: tuple[Route, ...] = ( - Route( - "knowledge_fs.ls", - "GET", - "/workspaces/{workspace_id}/knowledge-fs/knowledge-spaces/{knowledge_space_id}/fs:ls", - frozenset({Trait.ACCOUNT_PRIMARY}), - ), - Route( - "knowledge_fs.tree", - "GET", - "/workspaces/{workspace_id}/knowledge-fs/knowledge-spaces/{knowledge_space_id}/fs:tree", - frozenset({Trait.ACCOUNT_PRIMARY}), - ), - Route( - "knowledge_fs.grep", - "GET", - "/workspaces/{workspace_id}/knowledge-fs/knowledge-spaces/{knowledge_space_id}/fs:grep", - frozenset({Trait.ACCOUNT_PRIMARY}), - ), - Route( - "knowledge_fs.find", - "GET", - "/workspaces/{workspace_id}/knowledge-fs/knowledge-spaces/{knowledge_space_id}/fs:find", - frozenset({Trait.ACCOUNT_PRIMARY}), - ), - Route( - "knowledge_fs.diff", - "POST", - "/workspaces/{workspace_id}/knowledge-fs/knowledge-spaces/{knowledge_space_id}/fs:diff", - frozenset({Trait.ACCOUNT_PRIMARY}), - ), - Route( - "knowledge_fs.cat", - "GET", - "/workspaces/{workspace_id}/knowledge-fs/knowledge-spaces/{knowledge_space_id}/fs:cat", - frozenset({Trait.ACCOUNT_PRIMARY}), - ), - Route( - "knowledge_fs.stat", - "GET", - "/workspaces/{workspace_id}/knowledge-fs/knowledge-spaces/{knowledge_space_id}/fs:stat", - frozenset({Trait.ACCOUNT_PRIMARY}), - ), - Route("describe.account", "GET", "/account", frozenset({Trait.ACCOUNT_PRIMARY})), - Route("account.sessions.revoke_self", "DELETE", "/account/sessions/self", frozenset({Trait.ACCOUNT_PRIMARY})), - Route("get.account.session", "GET", "/account/sessions", frozenset({Trait.ACCOUNT_PRIMARY})), - Route( - "account.sessions.revoke_one", - "DELETE", - "/account/sessions/{session_id}", - frozenset({Trait.ACCOUNT_PRIMARY}), - ), - Route( - "apps.describe", - "GET", - "/apps/{app_id}", - frozenset({Trait.ACCOUNT_PRIMARY, Trait.APP_SCOPED}), - ), - Route("apps.list", "GET", "/apps", frozenset({Trait.ACCOUNT_PRIMARY}), query="workspace_id={workspace_id}"), - Route("workspaces.list", "GET", "/workspaces", frozenset({Trait.ACCOUNT_PRIMARY})), - Route("workspaces.describe", "GET", "/workspaces/{workspace_id}", frozenset({Trait.ACCOUNT_PRIMARY})), - Route("workspaces.members.list", "GET", "/workspaces/{workspace_id}/members", frozenset({Trait.ACCOUNT_PRIMARY})), - Route( - "workspaces.members.invite", - "POST", - "/workspaces/{workspace_id}/members", - frozenset({Trait.ACCOUNT_PRIMARY}), - ), - Route( - "workspaces.members.remove", - "DELETE", - "/workspaces/{workspace_id}/members/{member_id}", - frozenset({Trait.ACCOUNT_PRIMARY}), - ), - Route( - "workspaces.members.update_role", - "PATCH", - "/workspaces/{workspace_id}/members/{member_id}", - frozenset({Trait.ACCOUNT_PRIMARY}), - ), - Route( - "plugins.marketplace", - "GET", - "/workspaces/{workspace_id}/marketplace/plugins", - frozenset({Trait.ACCOUNT_PRIMARY}), - ), - Route("plugins.list", "GET", "/workspaces/{workspace_id}/plugins", frozenset({Trait.ACCOUNT_PRIMARY})), - Route("plugins.install", "POST", "/workspaces/{workspace_id}/plugins:install", frozenset({Trait.ACCOUNT_PRIMARY})), - Route( - "plugins.task", - "GET", - "/workspaces/{workspace_id}/plugin-tasks/{task_id}", - frozenset({Trait.ACCOUNT_PRIMARY}), - ), - Route( - "plugins.delete", - "DELETE", - "/workspaces/{workspace_id}/plugins/{plugin_id}", - frozenset({Trait.ACCOUNT_PRIMARY}), - ), - Route( - "model_providers.describe", - "GET", - "/workspaces/{workspace_id}/model-providers/{provider}", - frozenset({Trait.ACCOUNT_PRIMARY}), - ), - Route( - "model_providers.credential.create", - "POST", - "/workspaces/{workspace_id}/model-providers/{provider}/credentials", - frozenset({Trait.ACCOUNT_PRIMARY}), - ), - Route( - "model_providers.credential.set", - "PATCH", - "/workspaces/{workspace_id}/model-providers/{provider}/credentials/{credential_id}", - frozenset({Trait.ACCOUNT_PRIMARY}), - ), - Route("tools.list", "GET", "/workspaces/{workspace_id}/tools", frozenset({Trait.ACCOUNT_PRIMARY})), - Route( - "tool_providers.describe", - "GET", - "/workspaces/{workspace_id}/tool-providers/{provider}", - frozenset({Trait.ACCOUNT_PRIMARY}), - ), - Route( - "tool_providers.credential.create", - "POST", - "/workspaces/{workspace_id}/tool-providers/{provider}/credentials", - frozenset({Trait.ACCOUNT_PRIMARY}), - ), - Route( - "tool_providers.credential.set", - "PATCH", - "/workspaces/{workspace_id}/tool-providers/{provider}/credentials/{credential_id}", - frozenset({Trait.ACCOUNT_PRIMARY}), - ), - Route( - "knowledge_bases.list", - "GET", - "/workspaces/{workspace_id}/knowledge-bases", - frozenset({Trait.ACCOUNT_PRIMARY}), - ), - Route( - "models.list", - "GET", - "/workspaces/{workspace_id}/models", - frozenset({Trait.ACCOUNT_PRIMARY}), - ), - Route( - "models.credential.create", - "POST", - "/workspaces/{workspace_id}/model-providers/{provider}/models/credentials", - frozenset({Trait.ACCOUNT_PRIMARY}), - ), - Route( - "models.credential.set", - "PATCH", - "/workspaces/{workspace_id}/model-providers/{provider}/models/credentials/{credential_id}", - frozenset({Trait.ACCOUNT_PRIMARY}), - ), - Route( - "app_dsl.import", - "POST", - "/workspaces/{workspace_id}/apps/imports", - frozenset({Trait.ACCOUNT_PRIMARY}), - ), - Route( - "app_dsl.check", - "POST", - "/workspaces/{workspace_id}/apps/imports:check", - frozenset({Trait.ACCOUNT_PRIMARY}), - ), - Route( - "app_dsl.import_confirm", - "POST", - "/workspaces/{workspace_id}/apps/imports/{import_id}:confirm", - frozenset({Trait.ACCOUNT_PRIMARY}), - ), - Route( - "app_dsl.export", - "GET", - "/apps/{app_id}/dsl", - frozenset({Trait.ACCOUNT_PRIMARY, Trait.APP_SCOPED}), - ), - Route( - "app_dsl.check_dependencies", - "GET", - "/apps/{app_id}/dependencies:check", - frozenset({Trait.ACCOUNT_PRIMARY, Trait.APP_SCOPED}), - ), - Route("app_run.workflow", "POST", "/apps/{app_id}/workflow:run", _RUN_TRAITS), - Route("app_run.chat", "POST", "/apps/{app_id}/chat:run", _RUN_TRAITS), - Route("app_run.advanced_chat", "POST", "/apps/{app_id}/advanced-chat:run", _RUN_TRAITS), - Route("app_run.completion", "POST", "/apps/{app_id}/completion:run", _RUN_TRAITS), - Route( - "app_run.draft.workflow", - "POST", - "/apps/{app_id}/draft/workflow:run", - frozenset({Trait.ACCOUNT_PRIMARY, Trait.APP_SCOPED}), - ), - Route( - "app_run.draft.advanced_chat", - "POST", - "/apps/{app_id}/draft/advanced-chat:run", - frozenset({Trait.ACCOUNT_PRIMARY, Trait.APP_SCOPED}), - ), - Route("app_workflow.run.list", "GET", "/apps/{app_id}/runs", frozenset({Trait.ACCOUNT_PRIMARY, Trait.APP_SCOPED})), - Route( - "app_workflow.run.describe", - "GET", - "/apps/{app_id}/runs/{run_id}", - frozenset({Trait.ACCOUNT_PRIMARY, Trait.APP_SCOPED}), - ), - Route( - "app_workflow.run.nodes", - "GET", - "/apps/{app_id}/runs/{run_id}/nodes", - frozenset({Trait.ACCOUNT_PRIMARY, Trait.APP_SCOPED}), - ), - Route( - "app_workflow.publish", "POST", "/apps/{app_id}:publish", frozenset({Trait.ACCOUNT_PRIMARY, Trait.APP_SCOPED}) - ), - Route( - "app_workflow.version.list", - "GET", - "/apps/{app_id}/versions", - frozenset({Trait.ACCOUNT_PRIMARY, Trait.APP_SCOPED}), - ), - Route( - "app_workflow.version.restore", - "POST", - "/apps/{app_id}/versions/{version_id}:restore", - frozenset({Trait.ACCOUNT_PRIMARY, Trait.APP_SCOPED}), - ), - Route( - "app_workflow.release_check", - "GET", - "/apps/{app_id}/release:check", - frozenset({Trait.ACCOUNT_PRIMARY, Trait.APP_SCOPED}), - ), - Route("app_workflow.env.list", "GET", "/apps/{app_id}/env", frozenset({Trait.ACCOUNT_PRIMARY, Trait.APP_SCOPED})), - Route( - "app_workflow.env.set", - "PUT", - "/apps/{app_id}/env/{env_id}", - frozenset({Trait.ACCOUNT_PRIMARY, Trait.APP_SCOPED}), - ), - Route( - "app_workflow.env.delete", - "DELETE", - "/apps/{app_id}/env/{env_id}", - frozenset({Trait.ACCOUNT_PRIMARY, Trait.APP_SCOPED}), - ), - Route("node_types.list", "GET", "/node-types", frozenset({Trait.ACCOUNT_PRIMARY, Trait.EXTERNAL_REACHABLE})), - Route( - "node_types.describe", - "GET", - "/node-types/{node_type}", - frozenset({Trait.ACCOUNT_PRIMARY, Trait.EXTERNAL_REACHABLE}), - ), - Route( - "app_create.workflow", "POST", "/workspaces/{workspace_id}/apps/workflow", frozenset({Trait.ACCOUNT_PRIMARY}) - ), - Route( - "app_create.advanced_chat", - "POST", - "/workspaces/{workspace_id}/apps/advanced-chat", - frozenset({Trait.ACCOUNT_PRIMARY}), - ), - Route( - "app_workflow.node_run.workflow", - "POST", - "/apps/{app_id}/draft/workflow/nodes/{node_id}:run", - frozenset({Trait.ACCOUNT_PRIMARY, Trait.APP_SCOPED}), - ), - Route( - "app_workflow.node_run.advanced_chat", - "POST", - "/apps/{app_id}/draft/advanced-chat/nodes/{node_id}:run", - frozenset({Trait.ACCOUNT_PRIMARY, Trait.APP_SCOPED}), - ), - Route( - "app_info.describe.workflow", - "GET", - "/apps/{app_id}/app-info/workflow", - frozenset({Trait.ACCOUNT_PRIMARY, Trait.APP_SCOPED}), - ), - Route( - "app_info.set.workflow", - "PATCH", - "/apps/{app_id}/app-info/workflow", - frozenset({Trait.ACCOUNT_PRIMARY, Trait.APP_SCOPED}), - ), - Route( - "service_api.describe", - "GET", - "/apps/{app_id}/service-api", - frozenset({Trait.ACCOUNT_PRIMARY, Trait.APP_SCOPED}), - ), - Route( - "service_api.set", - "PATCH", - "/apps/{app_id}/service-api", - frozenset({Trait.ACCOUNT_PRIMARY, Trait.APP_SCOPED}), - ), - Route( - "app_info.describe.advanced_chat", - "GET", - "/apps/{app_id}/app-info/advanced-chat", - frozenset({Trait.ACCOUNT_PRIMARY, Trait.APP_SCOPED}), - ), - Route( - "app_info.set.advanced_chat", - "PATCH", - "/apps/{app_id}/app-info/advanced-chat", - frozenset({Trait.ACCOUNT_PRIMARY, Trait.APP_SCOPED}), - ), - Route( - "app_info.describe.chat", - "GET", - "/apps/{app_id}/app-info/chat", - frozenset({Trait.ACCOUNT_PRIMARY, Trait.APP_SCOPED}), - ), - Route( - "app_info.set.chat", - "PATCH", - "/apps/{app_id}/app-info/chat", - frozenset({Trait.ACCOUNT_PRIMARY, Trait.APP_SCOPED}), - ), - Route( - "app_info.describe.agent_chat", - "GET", - "/apps/{app_id}/app-info/agent-chat", - frozenset({Trait.ACCOUNT_PRIMARY, Trait.APP_SCOPED}), - ), - Route( - "app_info.set.agent_chat", - "PATCH", - "/apps/{app_id}/app-info/agent-chat", - frozenset({Trait.ACCOUNT_PRIMARY, Trait.APP_SCOPED}), - ), - Route( - "app_info.describe.completion", - "GET", - "/apps/{app_id}/app-info/completion", - frozenset({Trait.ACCOUNT_PRIMARY, Trait.APP_SCOPED}), - ), - Route( - "app_info.set.completion", - "PATCH", - "/apps/{app_id}/app-info/completion", - frozenset({Trait.ACCOUNT_PRIMARY, Trait.APP_SCOPED}), - ), - Route( - "app_info.describe.agent", - "GET", - "/apps/{app_id}/app-info/agent", - frozenset({Trait.ACCOUNT_PRIMARY, Trait.APP_SCOPED}), - ), - Route( - "app_info.set.agent", - "PATCH", - "/apps/{app_id}/app-info/agent", - frozenset({Trait.ACCOUNT_PRIMARY, Trait.APP_SCOPED}), - ), - Route( - "service_api.describe.agent", - "GET", - "/apps/{app_id}/service-api/agent", - frozenset({Trait.ACCOUNT_PRIMARY, Trait.APP_SCOPED}), - ), - Route( - "service_api.set.agent", - "PATCH", - "/apps/{app_id}/service-api/agent", - frozenset({Trait.ACCOUNT_PRIMARY, Trait.APP_SCOPED}), - ), - Route( - "webapp.describe.workflow", - "GET", - "/apps/{app_id}/webapp/workflow", - frozenset({Trait.ACCOUNT_PRIMARY, Trait.APP_SCOPED}), - ), - Route( - "webapp.set.workflow", - "PATCH", - "/apps/{app_id}/webapp/workflow", - frozenset({Trait.ACCOUNT_PRIMARY, Trait.APP_SCOPED}), - ), - Route( - "webapp.reset", - "POST", - "/apps/{app_id}/webapp:reset", - frozenset({Trait.ACCOUNT_PRIMARY, Trait.APP_SCOPED}), - ), - Route( - "webapp_access.describe", - "GET", - "/apps/{app_id}/webapp-access", - frozenset({Trait.ACCOUNT_PRIMARY, Trait.APP_SCOPED}), - ), - Route( - "webapp_access.set", - "PUT", - "/apps/{app_id}/webapp-access", - frozenset({Trait.ACCOUNT_PRIMARY, Trait.APP_SCOPED}), - ), - Route( - "webapp.describe.advanced_chat", - "GET", - "/apps/{app_id}/webapp/advanced-chat", - frozenset({Trait.ACCOUNT_PRIMARY, Trait.APP_SCOPED}), - ), - Route( - "webapp.set.advanced_chat", - "PATCH", - "/apps/{app_id}/webapp/advanced-chat", - frozenset({Trait.ACCOUNT_PRIMARY, Trait.APP_SCOPED}), - ), - Route( - "webapp.describe.chat", - "GET", - "/apps/{app_id}/webapp/chat", - frozenset({Trait.ACCOUNT_PRIMARY, Trait.APP_SCOPED}), - ), - Route( - "webapp.set.chat", - "PATCH", - "/apps/{app_id}/webapp/chat", - frozenset({Trait.ACCOUNT_PRIMARY, Trait.APP_SCOPED}), - ), - Route( - "webapp.describe.agent_chat", - "GET", - "/apps/{app_id}/webapp/agent-chat", - frozenset({Trait.ACCOUNT_PRIMARY, Trait.APP_SCOPED}), - ), - Route( - "webapp.set.agent_chat", - "PATCH", - "/apps/{app_id}/webapp/agent-chat", - frozenset({Trait.ACCOUNT_PRIMARY, Trait.APP_SCOPED}), - ), - Route( - "webapp.describe.completion", - "GET", - "/apps/{app_id}/webapp/completion", - frozenset({Trait.ACCOUNT_PRIMARY, Trait.APP_SCOPED}), - ), - Route( - "webapp.set.completion", - "PATCH", - "/apps/{app_id}/webapp/completion", - frozenset({Trait.ACCOUNT_PRIMARY, Trait.APP_SCOPED}), - ), - Route( - "webapp.describe.agent", - "GET", - "/apps/{app_id}/webapp/agent", - frozenset({Trait.ACCOUNT_PRIMARY, Trait.APP_SCOPED}), - ), - Route( - "webapp.set.agent", - "PATCH", - "/apps/{app_id}/webapp/agent", - frozenset({Trait.ACCOUNT_PRIMARY, Trait.APP_SCOPED}), - ), - Route( - "webapp.reset.agent", - "POST", - "/apps/{app_id}/webapp/agent:reset", - frozenset({Trait.ACCOUNT_PRIMARY, Trait.APP_SCOPED}), - ), - Route( - "webapp_access.describe.agent", - "GET", - "/apps/{app_id}/webapp-access/agent", - frozenset({Trait.ACCOUNT_PRIMARY, Trait.APP_SCOPED}), - ), - Route( - "webapp_access.set.agent", - "PUT", - "/apps/{app_id}/webapp-access/agent", - frozenset({Trait.ACCOUNT_PRIMARY, Trait.APP_SCOPED}), - ), - Route( - "app_run.stop", - "POST", - "/apps/{app_id}/tasks/{task_id}:stop", - frozenset({Trait.ACCOUNT_PRIMARY, Trait.APP_SCOPED, Trait.EXTERNAL_REACHABLE}), - ), - Route( - "files.upload", - "POST", - "/apps/{app_id}/files", - frozenset({Trait.ACCOUNT_PRIMARY, Trait.APP_SCOPED, Trait.EXTERNAL_REACHABLE}), - ), - Route( - "human_input_form.get", - "GET", - "/apps/{app_id}/human-input-forms/{form_token}", - frozenset({Trait.ACCOUNT_PRIMARY, Trait.APP_SCOPED, Trait.EXTERNAL_REACHABLE}), - ), - Route( - "human_input_form.submit", - "POST", - "/apps/{app_id}/human-input-forms/{form_token}:submit", - frozenset({Trait.ACCOUNT_PRIMARY, Trait.APP_SCOPED, Trait.EXTERNAL_REACHABLE}), - ), - Route( - "workflow_events.stream", - "GET", - "/apps/{app_id}/tasks/{task_id}/events", - frozenset({Trait.ACCOUNT_PRIMARY, Trait.APP_SCOPED, Trait.EXTERNAL_REACHABLE}), - ), - Route( - "permitted_external.list", - "GET", - "/permitted-external-apps", - frozenset({Trait.EXTERNAL_REACHABLE, Trait.ENTERPRISE_ONLY}), - ), - Route( - "permitted_external.describe", - "GET", - "/permitted-external-apps/{app_id}", - frozenset({Trait.APP_SCOPED, Trait.EXTERNAL_REACHABLE, Trait.ENTERPRISE_ONLY}), - ), -) - - -CASE_REQUIRES: dict[Case, frozenset[Trait]] = { - Case.NO_BEARER: frozenset(), - Case.MEMBER: frozenset(), - Case.WRONG_SUBJECT: frozenset(), - Case.INSUFFICIENT_SCOPE: frozenset(), - Case.NON_MEMBER: frozenset({Trait.ACCOUNT_PRIMARY}), - # frozenset() rather than {ACCOUNT_PRIMARY}: this one also has to reach the - # SSO-only routes, to pin that they answer on the subject axis, not this one. - Case.NON_MEMBER_AND_INSUFFICIENT_SCOPE: frozenset(), - Case.LOW_ROLE: frozenset({Trait.ACCOUNT_PRIMARY}), - Case.APP_API_DISABLED: frozenset({Trait.APP_SCOPED}), - Case.UNKNOWN_APP: frozenset({Trait.APP_SCOPED}), - Case.FOREIGN_WORKSPACE_QUERY: frozenset({Trait.APP_SCOPED}), - Case.EDITION_NOT_ENTERPRISE: frozenset({Trait.ENTERPRISE_ONLY}), - Case.LICENSE_INVALID: frozenset(), - Case.EE_LICENSE_INVALID: frozenset(), - Case.EE_ACCOUNT_PUBLIC: frozenset({Trait.APP_SCOPED, Trait.ACCOUNT_PRIMARY}), - Case.EE_ACCOUNT_SSO_VERIFIED: frozenset({Trait.APP_SCOPED, Trait.ACCOUNT_PRIMARY}), - Case.EE_ACCOUNT_PRIVATE_ALL: frozenset({Trait.APP_SCOPED, Trait.ACCOUNT_PRIMARY}), - Case.EE_ACCOUNT_PRIVATE_PERMITTED: frozenset({Trait.APP_SCOPED, Trait.ACCOUNT_PRIMARY}), - Case.EE_ACCOUNT_PRIVATE_REFUSED: frozenset({Trait.APP_SCOPED, Trait.ACCOUNT_PRIMARY}), - Case.EE_ACCOUNT_PRIVATE_REFUSED_WEBAPP_AUTH_OFF: frozenset({Trait.APP_SCOPED, Trait.ACCOUNT_PRIMARY}), - Case.EE_ACCOUNT_MODE_UNRESOLVED: frozenset({Trait.APP_SCOPED, Trait.ACCOUNT_PRIMARY}), - Case.EE_EXTERNAL_PUBLIC: frozenset({Trait.APP_SCOPED, Trait.EXTERNAL_REACHABLE}), - Case.EE_EXTERNAL_SSO_VERIFIED: frozenset({Trait.APP_SCOPED, Trait.EXTERNAL_REACHABLE}), - Case.EE_EXTERNAL_PRIVATE_ALL: frozenset({Trait.APP_SCOPED, Trait.EXTERNAL_REACHABLE}), - Case.EE_EXTERNAL_PRIVATE: frozenset({Trait.APP_SCOPED, Trait.EXTERNAL_REACHABLE}), - Case.EE_EXTERNAL_PRIVATE_REFUSED_WEBAPP_AUTH_OFF: frozenset({Trait.APP_SCOPED, Trait.EXTERNAL_REACHABLE}), - Case.EE_EXTERNAL_MODE_UNRESOLVED: frozenset({Trait.APP_SCOPED, Trait.EXTERNAL_REACHABLE}), - Case.RBAC_ON_LOW_ROLE: frozenset({Trait.ACCOUNT_PRIMARY}), - Case.RBAC_ON_DENIED: frozenset({Trait.ACCOUNT_PRIMARY}), -} - - -SCENARIOS: dict[Case, Scenario] = { - Case.NO_BEARER: Scenario(bearer=Bearer.NONE), - Case.MEMBER: Scenario(bearer=Bearer.PRIMARY), - Case.WRONG_SUBJECT: Scenario(bearer=Bearer.OTHER_SUBJECT), - Case.INSUFFICIENT_SCOPE: Scenario(bearer=Bearer.PRIMARY, narrow_scopes=True), - Case.NON_MEMBER: Scenario(bearer=Bearer.ACCOUNT_OUTSIDER), - Case.NON_MEMBER_AND_INSUFFICIENT_SCOPE: Scenario(bearer=Bearer.ACCOUNT_OUTSIDER, narrow_scopes=True), - Case.LOW_ROLE: Scenario(bearer=Bearer.ACCOUNT_LOW_ROLE), - Case.APP_API_DISABLED: Scenario(bearer=Bearer.PRIMARY, app_api_enabled=False), - Case.UNKNOWN_APP: Scenario(bearer=Bearer.PRIMARY, unknown_app=True), - Case.FOREIGN_WORKSPACE_QUERY: Scenario(bearer=Bearer.PRIMARY, foreign_workspace_query=True), - Case.EDITION_NOT_ENTERPRISE: Scenario(bearer=Bearer.PRIMARY, edition=DeploymentEdition.COMMUNITY), - Case.LICENSE_INVALID: Scenario(bearer=Bearer.PRIMARY, license_status=LicenseStatus.EXPIRED), - Case.EE_LICENSE_INVALID: Scenario( - bearer=Bearer.PRIMARY, edition=DeploymentEdition.ENTERPRISE, license_status=LicenseStatus.EXPIRED - ), - Case.EE_ACCOUNT_PUBLIC: Scenario( - bearer=Bearer.ACCOUNT_MEMBER, - edition=DeploymentEdition.ENTERPRISE, - webapp_auth=True, - access_mode=WebAppAccessMode.PUBLIC, - ), - Case.EE_ACCOUNT_SSO_VERIFIED: Scenario( - bearer=Bearer.ACCOUNT_MEMBER, - edition=DeploymentEdition.ENTERPRISE, - webapp_auth=True, - access_mode=WebAppAccessMode.SSO_VERIFIED, - ), - Case.EE_ACCOUNT_PRIVATE_ALL: Scenario( - bearer=Bearer.ACCOUNT_MEMBER, - edition=DeploymentEdition.ENTERPRISE, - webapp_auth=True, - access_mode=WebAppAccessMode.PRIVATE_ALL, - ), - Case.EE_ACCOUNT_PRIVATE_PERMITTED: Scenario( - bearer=Bearer.ACCOUNT_MEMBER, - edition=DeploymentEdition.ENTERPRISE, - webapp_auth=True, - access_mode=WebAppAccessMode.PRIVATE, - private_app_permitted=True, - ), - Case.EE_ACCOUNT_PRIVATE_REFUSED: Scenario( - bearer=Bearer.ACCOUNT_MEMBER, - edition=DeploymentEdition.ENTERPRISE, - webapp_auth=True, - access_mode=WebAppAccessMode.PRIVATE, - private_app_permitted=False, - ), - Case.EE_ACCOUNT_PRIVATE_REFUSED_WEBAPP_AUTH_OFF: Scenario( - bearer=Bearer.ACCOUNT_MEMBER, - edition=DeploymentEdition.ENTERPRISE, - webapp_auth=False, - access_mode=WebAppAccessMode.PRIVATE, - private_app_permitted=False, - ), - Case.EE_ACCOUNT_MODE_UNRESOLVED: Scenario( - bearer=Bearer.ACCOUNT_MEMBER, - edition=DeploymentEdition.ENTERPRISE, - webapp_auth=True, - access_mode=None, - ), - Case.EE_EXTERNAL_PUBLIC: Scenario( - bearer=Bearer.EXTERNAL, - edition=DeploymentEdition.ENTERPRISE, - webapp_auth=True, - access_mode=WebAppAccessMode.PUBLIC, - ), - Case.EE_EXTERNAL_SSO_VERIFIED: Scenario( - bearer=Bearer.EXTERNAL, - edition=DeploymentEdition.ENTERPRISE, - webapp_auth=True, - access_mode=WebAppAccessMode.SSO_VERIFIED, - ), - Case.EE_EXTERNAL_PRIVATE_ALL: Scenario( - bearer=Bearer.EXTERNAL, - edition=DeploymentEdition.ENTERPRISE, - webapp_auth=True, - access_mode=WebAppAccessMode.PRIVATE_ALL, - ), - Case.EE_EXTERNAL_PRIVATE: Scenario( - bearer=Bearer.EXTERNAL, - edition=DeploymentEdition.ENTERPRISE, - webapp_auth=True, - access_mode=WebAppAccessMode.PRIVATE, - ), - Case.EE_EXTERNAL_PRIVATE_REFUSED_WEBAPP_AUTH_OFF: Scenario( - bearer=Bearer.EXTERNAL, - edition=DeploymentEdition.ENTERPRISE, - webapp_auth=False, - access_mode=WebAppAccessMode.PRIVATE, - private_app_permitted=False, - ), - Case.EE_EXTERNAL_MODE_UNRESOLVED: Scenario( - bearer=Bearer.EXTERNAL, - edition=DeploymentEdition.ENTERPRISE, - webapp_auth=True, - access_mode=None, - ), - Case.RBAC_ON_LOW_ROLE: Scenario(bearer=Bearer.ACCOUNT_LOW_ROLE, rbac_enabled=True, rbac_allows=True), - Case.RBAC_ON_DENIED: Scenario(bearer=Bearer.ACCOUNT_MEMBER, rbac_enabled=True, rbac_allows=False), -} - - -ROUTER_CASE_ROUTE = "describe.account" -ROUTER_CASES: dict[Case, Expect] = { - Case.NO_BEARER: DENY_NO_BEARER, - Case.LICENSE_INVALID: ADMIT_NO_LICENCE_GATE, - Case.EE_LICENSE_INVALID: DENY_LICENSE, -} -"""Answered by the router before any route-specific requirement runs, so pinned once.""" - - -_ACCOUNT_ONLY_NO_WORKSPACE: dict[Case, Expect] = { - Case.MEMBER: ADMIT, - Case.WRONG_SUBJECT: DENY_WRONG_SUBJECT, - Case.INSUFFICIENT_SCOPE: DENY_SCOPE, - Case.NON_MEMBER: ADMIT, - # No live membership check on this route, so the outsider's only failure is - # scope — same answer pre- and post-PR. - Case.NON_MEMBER_AND_INSUFFICIENT_SCOPE: DENY_SCOPE, - Case.LOW_ROLE: ADMIT, - Case.RBAC_ON_LOW_ROLE: ADMIT_NO_WORKSPACE_ROLE, - Case.RBAC_ON_DENIED: ADMIT_NO_RBAC_PERMISSION, -} - -_ACCOUNT_MEMBER_NO_ROLE: dict[Case, Expect] = { - Case.MEMBER: ADMIT, - Case.WRONG_SUBJECT: DENY_WRONG_SUBJECT, - Case.INSUFFICIENT_SCOPE: DENY_SCOPE, - Case.NON_MEMBER: DENY_NON_MEMBER, - Case.NON_MEMBER_AND_INSUFFICIENT_SCOPE: DENY_NON_MEMBER, - Case.LOW_ROLE: ADMIT, - Case.RBAC_ON_LOW_ROLE: ADMIT_NO_WORKSPACE_ROLE, - Case.RBAC_ON_DENIED: ADMIT_NO_RBAC_PERMISSION, -} - -_ACCOUNT_MEMBER_WITH_ROLE: dict[Case, Expect] = { - Case.MEMBER: ADMIT, - Case.WRONG_SUBJECT: DENY_WRONG_SUBJECT, - Case.INSUFFICIENT_SCOPE: DENY_SCOPE, - Case.NON_MEMBER: DENY_NON_MEMBER, - Case.NON_MEMBER_AND_INSUFFICIENT_SCOPE: DENY_NON_MEMBER, - Case.LOW_ROLE: DENY_ROLE, - Case.RBAC_ON_LOW_ROLE: ADMIT, - Case.RBAC_ON_DENIED: DENY_RBAC, -} - -_DUAL_SUBJECT_RUN: dict[Case, Expect] = { - Case.MEMBER: ADMIT, - Case.WRONG_SUBJECT: DENY_SSO_NEEDS_EE, - Case.INSUFFICIENT_SCOPE: DENY_SCOPE, - Case.NON_MEMBER: DENY_NON_MEMBER, - Case.NON_MEMBER_AND_INSUFFICIENT_SCOPE: DENY_NON_MEMBER, - Case.LOW_ROLE: ADMIT, - Case.APP_API_DISABLED: DENY_API_DISABLED, - Case.UNKNOWN_APP: DENY_UNKNOWN_APP, - Case.FOREIGN_WORKSPACE_QUERY: ADMIT, - Case.EE_ACCOUNT_PUBLIC: ADMIT, - Case.EE_ACCOUNT_SSO_VERIFIED: ADMIT, - Case.EE_ACCOUNT_PRIVATE_ALL: ADMIT, - Case.EE_ACCOUNT_PRIVATE_PERMITTED: ADMIT, - Case.EE_ACCOUNT_PRIVATE_REFUSED: DENY_PRIVATE_APP, - Case.EE_ACCOUNT_PRIVATE_REFUSED_WEBAPP_AUTH_OFF: DENY_PRIVATE_APP, - Case.EE_ACCOUNT_MODE_UNRESOLVED: DENY_MODE_UNRESOLVED, - Case.EE_EXTERNAL_PUBLIC: ADMIT, - Case.EE_EXTERNAL_SSO_VERIFIED: ADMIT, - Case.EE_EXTERNAL_PRIVATE_ALL: DENY_ACCESS_MODE, - Case.EE_EXTERNAL_PRIVATE: DENY_ACCESS_MODE, - Case.EE_EXTERNAL_PRIVATE_REFUSED_WEBAPP_AUTH_OFF: DENY_PRIVATE_APP, - Case.EE_EXTERNAL_MODE_UNRESOLVED: DENY_MODE_UNRESOLVED, - Case.RBAC_ON_DENIED: DENY_RBAC, - Case.RBAC_ON_LOW_ROLE: ADMIT_NO_WORKSPACE_ROLE, -} - - -_ACCOUNT_READER_APP: dict[Case, Expect] = { - Case.MEMBER: ADMIT, - Case.WRONG_SUBJECT: DENY_WRONG_SUBJECT, - Case.INSUFFICIENT_SCOPE: DENY_SCOPE, - Case.NON_MEMBER: DENY_NON_MEMBER, - Case.NON_MEMBER_AND_INSUFFICIENT_SCOPE: DENY_NON_MEMBER, - Case.LOW_ROLE: ADMIT, - Case.APP_API_DISABLED: DENY_API_DISABLED, - Case.UNKNOWN_APP: DENY_UNKNOWN_APP, - Case.FOREIGN_WORKSPACE_QUERY: ADMIT, - Case.EE_ACCOUNT_PUBLIC: ADMIT, - Case.EE_ACCOUNT_SSO_VERIFIED: ADMIT, - Case.EE_ACCOUNT_PRIVATE_ALL: ADMIT, - Case.EE_ACCOUNT_PRIVATE_PERMITTED: ADMIT, - Case.EE_ACCOUNT_PRIVATE_REFUSED: ADMIT, - Case.EE_ACCOUNT_PRIVATE_REFUSED_WEBAPP_AUTH_OFF: ADMIT, - Case.EE_ACCOUNT_MODE_UNRESOLVED: ADMIT, - Case.RBAC_ON_DENIED: DENY_RBAC, - Case.RBAC_ON_LOW_ROLE: ADMIT_NO_WORKSPACE_ROLE, -} - -_ACCOUNT_EDITOR_APP: dict[Case, Expect] = { - **_ACCOUNT_READER_APP, - Case.LOW_ROLE: DENY_ROLE, - Case.RBAC_ON_LOW_ROLE: ADMIT, -} - -_ACCOUNT_READER_SETTINGS: dict[Case, Expect] = {**_ACCOUNT_READER_APP, Case.APP_API_DISABLED: ADMIT} -_ACCOUNT_EDITOR_SETTINGS: dict[Case, Expect] = {**_ACCOUNT_EDITOR_APP, Case.APP_API_DISABLED: ADMIT} - -DENY_NEEDS_WEBAPP_EE = Expect(403, WebAppAccessRequiresEE.description) - -_WEBAPP_ACCESS_EDITOR: dict[Case, Expect] = { - **_ACCOUNT_EDITOR_SETTINGS, - Case.MEMBER: DENY_NEEDS_WEBAPP_EE, - Case.APP_API_DISABLED: DENY_NEEDS_WEBAPP_EE, - Case.FOREIGN_WORKSPACE_QUERY: DENY_NEEDS_WEBAPP_EE, - Case.EE_ACCOUNT_PRIVATE_REFUSED_WEBAPP_AUTH_OFF: DENY_NEEDS_WEBAPP_EE, - Case.RBAC_ON_LOW_ROLE: DENY_NEEDS_WEBAPP_EE, -} - - -_ANY_BEARER: dict[Case, Expect] = { - **{case: ADMIT for case in Case if CASE_REQUIRES[case] <= {Trait.ACCOUNT_PRIMARY, Trait.EXTERNAL_REACHABLE}}, - Case.WRONG_SUBJECT: DENY_SSO_NEEDS_EE, -} -"""A route with no requirements admits every case that carries a valid bearer. The one -exception is the router's own answer: an external-SSO bearer is refused outside enterprise.""" - -MATRIX: dict[str, dict[Case, Expect]] = { - "describe.account": dict(_ACCOUNT_ONLY_NO_WORKSPACE), - "workspaces.list": dict(_ACCOUNT_ONLY_NO_WORKSPACE), - "apps.list": dict(_ACCOUNT_MEMBER_NO_ROLE), - "workspaces.switch": dict(_ACCOUNT_MEMBER_NO_ROLE), - # invite/remove check WORKSPACE_MEMBER_MANAGE; update_role checks WORKSPACE_ROLE_MANAGE - # (matches check_member_permission's own add/remove vs update split) - "workspaces.members.invite": dict(_ACCOUNT_MEMBER_WITH_ROLE), - "workspaces.members.update_role": dict(_ACCOUNT_MEMBER_WITH_ROLE), - "plugins.install": {**_ACCOUNT_MEMBER_NO_ROLE, Case.RBAC_ON_DENIED: DENY_RBAC}, - "model_providers.credential.create": dict(_ACCOUNT_MEMBER_WITH_ROLE), - "app_dsl.import": dict(_ACCOUNT_MEMBER_WITH_ROLE), - "app_dsl.check": dict(_ACCOUNT_MEMBER_WITH_ROLE), - "app_create.workflow": dict(_ACCOUNT_MEMBER_WITH_ROLE), - "node_types.list": dict(_ANY_BEARER), - "app_workflow.node_run.workflow": dict(_ACCOUNT_EDITOR_APP), - "app_info.describe.workflow": dict(_ACCOUNT_READER_SETTINGS), - "app_info.set.workflow": dict(_ACCOUNT_EDITOR_SETTINGS), - "service_api.describe": dict(_ACCOUNT_EDITOR_SETTINGS), - "service_api.set": dict(_ACCOUNT_EDITOR_SETTINGS), - "webapp.describe.workflow": dict(_ACCOUNT_READER_SETTINGS), - "webapp.set.workflow": dict(_ACCOUNT_EDITOR_SETTINGS), - "webapp.reset": dict(_ACCOUNT_EDITOR_SETTINGS), - "webapp_access.describe": dict(_WEBAPP_ACCESS_EDITOR), - "apps.describe": dict(_ACCOUNT_READER_APP), - "app_dsl.export": dict(_ACCOUNT_EDITOR_APP), - "app_run.draft.workflow": dict(_ACCOUNT_EDITOR_APP), - "app_workflow.run.list": dict(_ACCOUNT_READER_APP), - "app_workflow.run.describe": dict(_ACCOUNT_READER_APP), - "app_workflow.run.nodes": dict(_ACCOUNT_READER_APP), - "app_workflow.publish": dict(_ACCOUNT_EDITOR_APP), - "app_workflow.version.list": dict(_ACCOUNT_EDITOR_APP), - "app_workflow.version.restore": dict(_ACCOUNT_EDITOR_APP), - "app_workflow.release_check": dict(_ACCOUNT_EDITOR_APP), - "app_workflow.env.list": dict(_ACCOUNT_EDITOR_APP), - "app_workflow.env.set": dict(_ACCOUNT_EDITOR_APP), - "app_workflow.env.delete": dict(_ACCOUNT_EDITOR_APP), - "app_run.workflow": dict(_DUAL_SUBJECT_RUN), - "app_run.chat": dict(_DUAL_SUBJECT_RUN), - "app_run.advanced_chat": dict(_DUAL_SUBJECT_RUN), - "app_run.completion": dict(_DUAL_SUBJECT_RUN), - "human_input_form.get": dict(_DUAL_SUBJECT_RUN), - "files.upload": { - **_DUAL_SUBJECT_RUN, - Case.RBAC_ON_DENIED: ADMIT_NO_RBAC_PERMISSION, - }, - "permitted_external.list": { - Case.MEMBER: ADMIT_NO_MOUNT, - Case.WRONG_SUBJECT: DENY_WRONG_SUBJECT, - Case.INSUFFICIENT_SCOPE: DENY_SCOPE, - # An account bearer answers on the subject axis here, before scope or - # membership: this route's CheckSubject allows only ExternalSsoSubject. - Case.NON_MEMBER_AND_INSUFFICIENT_SCOPE: DENY_WRONG_SUBJECT, - Case.EDITION_NOT_ENTERPRISE: DENY_EDITION, - }, - "permitted_external.describe": { - Case.MEMBER: ADMIT, - Case.WRONG_SUBJECT: DENY_WRONG_SUBJECT, - Case.INSUFFICIENT_SCOPE: DENY_SCOPE, - Case.NON_MEMBER_AND_INSUFFICIENT_SCOPE: DENY_WRONG_SUBJECT, - Case.EDITION_NOT_ENTERPRISE: DENY_EDITION, - Case.APP_API_DISABLED: DENY_API_DISABLED, - Case.UNKNOWN_APP: DENY_UNKNOWN_APP, - Case.FOREIGN_WORKSPACE_QUERY: Expect( - ADMITTED, - note="a foreign ?workspace_id= is ignored on every app-scoped route (spec 2.9)", - ), - Case.EE_EXTERNAL_PUBLIC: ADMIT, - Case.EE_EXTERNAL_SSO_VERIFIED: ADMIT, - Case.EE_EXTERNAL_PRIVATE_ALL: DENY_ACCESS_MODE, - Case.EE_EXTERNAL_PRIVATE: DENY_ACCESS_MODE, - Case.EE_EXTERNAL_PRIVATE_REFUSED_WEBAPP_AUTH_OFF: DENY_PRIVATE_APP, - Case.EE_EXTERNAL_MODE_UNRESOLVED: DENY_MODE_UNRESOLVED, - }, -} -"""One representative per distinct declared requirement tuple; `DECLARED` pins the rest.""" - - -_ACCOUNT = (AccountSubject,) -_ACCOUNT_OR_EXTERNAL = (AccountSubject, ExternalSsoSubject) -_OWNER_ADMIN = frozenset({TenantAccountRole.OWNER, TenantAccountRole.ADMIN}) -_EDITOR_UP = frozenset({TenantAccountRole.EDITOR, TenantAccountRole.ADMIN, TenantAccountRole.OWNER}) - -_REQ_ACCOUNT_FULL = (CheckSubject(allowed=_ACCOUNT), CheckScope(Scope.FULL)) -_REQ_ACCOUNT_WORKSPACE_READ = (CheckSubject(allowed=_ACCOUNT), CheckScope(Scope.WORKSPACE_READ)) -_REQ_ACCOUNT_APPS_READ_MEMBER = ( - CheckSubject(allowed=_ACCOUNT), - CheckScope(Scope.APPS_READ), - CheckWorkspaceMember(), -) -_REQ_ACCOUNT_WORKSPACE_READ_MEMBER = ( - CheckSubject(allowed=_ACCOUNT), - CheckScope(Scope.WORKSPACE_READ), - CheckWorkspaceMember(), -) -_REQ_MEMBER_MANAGE = ( - CheckSubject(allowed=_ACCOUNT), - CheckScope(Scope.WORKSPACE_WRITE), - CheckWorkspaceMember(), - CheckRBACPermission(RBACCheck(RBACPermission.WORKSPACE_MEMBER_MANAGE, Workspace())), - CheckWorkspaceRole(_OWNER_ADMIN), -) -_REQ_ROLE_MANAGE = ( - CheckSubject(allowed=_ACCOUNT), - CheckScope(Scope.WORKSPACE_WRITE), - CheckWorkspaceMember(), - CheckRBACPermission(RBACCheck(RBACPermission.WORKSPACE_ROLE_MANAGE, Workspace())), - CheckWorkspaceRole(_OWNER_ADMIN), -) -_REQ_DSL_WORKSPACE = ( - CheckSubject(allowed=_ACCOUNT), - CheckScope(Scope.WORKSPACE_WRITE), - CheckWorkspaceMember(), - CheckRBACPermission(RBACCheck(RBACPermission.APP_IMPORT_EXPORT_DSL, Workspace())), - CheckWorkspaceRole(_EDITOR_UP), -) -_REQ_DSL_APP = ( - CheckSubject(allowed=_ACCOUNT), - CheckAppApiEnabled(), - CheckWorkspaceMember(), - CheckScope(Scope.APPS_READ), - CheckRBACPermission(RBACCheck(RBACPermission.APP_IMPORT_EXPORT_DSL, PlainApp())), - CheckWorkspaceRole(_EDITOR_UP), -) -_REQ_APP_DESCRIBE = ( - CheckSubject(allowed=_ACCOUNT), - CheckAppApiEnabled(), - CheckWorkspaceMember(), - CheckScope(Scope.APPS_READ), - CheckRBACPermission(RBACCheck(RBACPermission.APP_VIEW_LAYOUT, PlainApp())), -) -_REQ_RUN = ( - CheckSubject(allowed=_ACCOUNT_OR_EXTERNAL), - CheckAppApiEnabled(), - CheckWorkspaceMember(), - CheckScope(Scope.APPS_RUN), - CheckRBACPermission(RBACCheck(RBACPermission.APP_TEST_AND_RUN, PlainApp())), - CheckAppAccess(), -) -_REQ_RUN_FORM = (*_REQ_RUN, CheckFormSurface()) -_REQ_RELEASE = ( - CheckSubject(allowed=_ACCOUNT), - CheckAppApiEnabled(), - CheckWorkspaceMember(), - CheckAppMode(AppMode.WORKFLOW, AppMode.ADVANCED_CHAT), - CheckScope(Scope.WORKSPACE_WRITE), - CheckRBACPermission(RBACCheck(RBACPermission.APP_RELEASE_AND_VERSION, PlainApp())), - CheckWorkspaceRole(_EDITOR_UP), -) -_REQ_VERSION_READ = ( - CheckSubject(allowed=_ACCOUNT), - CheckAppApiEnabled(), - CheckWorkspaceMember(), - CheckAppMode(AppMode.WORKFLOW, AppMode.ADVANCED_CHAT), - CheckScope(Scope.APPS_READ), - CheckRBACPermission(RBACCheck(RBACPermission.APP_VIEW_LAYOUT, PlainApp())), - CheckWorkspaceRole(_EDITOR_UP), -) -_REQ_ENV_WRITE = ( - CheckSubject(allowed=_ACCOUNT), - CheckAppApiEnabled(), - CheckWorkspaceMember(), - CheckAppMode(AppMode.WORKFLOW, AppMode.ADVANCED_CHAT), - CheckScope(Scope.WORKSPACE_WRITE), - CheckRBACPermission(RBACCheck(RBACPermission.APP_EDIT, PlainApp())), - CheckWorkspaceRole(_EDITOR_UP), -) -_REQ_RUN_HISTORY = ( - CheckSubject(allowed=_ACCOUNT), - CheckAppApiEnabled(), - CheckWorkspaceMember(), - CheckAppMode(AppMode.WORKFLOW, AppMode.ADVANCED_CHAT), - CheckScope(Scope.APPS_READ), - CheckRBACPermission(RBACCheck(RBACPermission.APP_CREATE_AND_MANAGEMENT, PlainApp())), -) -_REQ_DRAFT_RUN = ( - CheckSubject(allowed=_ACCOUNT), - CheckAppApiEnabled(), - CheckWorkspaceMember(), - CheckScope(Scope.APPS_RUN), - CheckRBACPermission(RBACCheck(RBACPermission.APP_TEST_AND_RUN, PlainApp())), - CheckWorkspaceRole(_EDITOR_UP), -) - - -def _draft_req(mode: AppMode) -> tuple[Requirement, ...]: - return (*_REQ_DRAFT_RUN[:3], CheckAppMode(mode), *_REQ_DRAFT_RUN[3:]) - - -_REQ_NODE_TYPES: tuple[Requirement, ...] = () -_REQ_APP_CREATE = ( - CheckSubject(allowed=_ACCOUNT), - CheckScope(Scope.WORKSPACE_WRITE), - CheckWorkspaceMember(), - CheckRBACPermission(RBACCheck(RBACPermission.APP_CREATE_AND_MANAGEMENT, Workspace())), - CheckWorkspaceRole(_EDITOR_UP), - CheckAppQuota(), -) -_REQ_FILES = ( - CheckSubject(allowed=_ACCOUNT_OR_EXTERNAL), - CheckAppApiEnabled(), - CheckWorkspaceMember(), - CheckScope(Scope.APPS_RUN), - CheckAppAccess(), -) -_REQ_EXTERNAL_LIST = ( - CheckSubject(allowed=(ExternalSsoSubject,)), - CheckScope(Scope.APPS_READ_PERMITTED_EXTERNAL), -) -_REQ_EXTERNAL_DESCRIBE = ( - CheckSubject(allowed=(ExternalSsoSubject,)), - CheckAppApiEnabled(), - CheckScope(Scope.APPS_READ_PERMITTED_EXTERNAL), - CheckAppAccess(), -) - - -def _req_plugin_write(permission: RBACPermission) -> tuple[Requirement, ...]: - return ( - CheckSubject(allowed=_ACCOUNT), - CheckScope(Scope.WORKSPACE_WRITE), - CheckWorkspaceMember(), - CheckRBACPermission(RBACCheck(permission, Workspace())), - CheckPluginInstallSetting(), - ) - - -def _req_admin_write(permission: RBACPermission) -> tuple[Requirement, ...]: - return ( - CheckSubject(allowed=_ACCOUNT), - CheckScope(Scope.WORKSPACE_WRITE), - CheckWorkspaceMember(), - CheckRBACPermission(RBACCheck(permission, Workspace())), - CheckWorkspaceRole(_OWNER_ADMIN), - ) - - -_REQ_PLUGIN_TASK = ( - CheckSubject(allowed=_ACCOUNT), - CheckScope(Scope.WORKSPACE_READ), - CheckWorkspaceMember(), - CheckPluginInstallSetting(), -) - - -def _settings_req( - modes: tuple[AppMode, ...], - perm: RBACPermission, - locator: PlainApp | AgentBehindApp, - scope: Scope, - roles: frozenset[TenantAccountRole] | None, -) -> tuple[Requirement, ...]: - base: tuple[Requirement, ...] = ( - CheckSubject(allowed=_ACCOUNT), - CheckWorkspaceMember(), - CheckAppMode(*modes), - CheckScope(scope), - CheckRBACPermission(RBACCheck(perm, locator)), - ) - return (*base, CheckWorkspaceRole(roles)) if roles is not None else base - - -_RESOURCE_RUN = (CheckSubject(allowed=(*_ACCOUNT_OR_EXTERNAL, ResourceAccessSubject)), *_REQ_RUN[1:]) - -DECLARED: dict[str, tuple[Requirement, ...]] = { - "knowledge_fs.ls": _REQ_ACCOUNT_WORKSPACE_READ_MEMBER, - "knowledge_fs.tree": _REQ_ACCOUNT_WORKSPACE_READ_MEMBER, - "knowledge_fs.grep": _REQ_ACCOUNT_WORKSPACE_READ_MEMBER, - "knowledge_fs.find": _REQ_ACCOUNT_WORKSPACE_READ_MEMBER, - "knowledge_fs.diff": _REQ_ACCOUNT_WORKSPACE_READ_MEMBER, - "knowledge_fs.cat": _REQ_ACCOUNT_WORKSPACE_READ_MEMBER, - "knowledge_fs.stat": _REQ_ACCOUNT_WORKSPACE_READ_MEMBER, - "describe.account": _REQ_ACCOUNT_FULL, - "account.sessions.revoke_self": _REQ_ACCOUNT_FULL, - "get.account.session": _REQ_ACCOUNT_FULL, - "account.sessions.revoke_one": _REQ_ACCOUNT_FULL, - "apps.describe": (CheckSubject(allowed=(AccountSubject, ResourceAccessSubject)), *_REQ_APP_DESCRIBE[1:]), - "apps.list": (CheckSubject(allowed=(AccountSubject, ResourceAccessSubject)), *_REQ_ACCOUNT_APPS_READ_MEMBER[1:]), - "workspaces.list": ( - CheckSubject(allowed=(AccountSubject, ResourceAccessSubject)), - *_REQ_ACCOUNT_WORKSPACE_READ[1:], - ), - "workspaces.describe": _REQ_ACCOUNT_WORKSPACE_READ, - "workspaces.switch": _REQ_ACCOUNT_WORKSPACE_READ_MEMBER, - "workspaces.members.list": _REQ_ACCOUNT_WORKSPACE_READ_MEMBER, - "workspaces.members.invite": (*_REQ_MEMBER_MANAGE, CheckWorkspaceInvitationQuota()), - "workspaces.members.remove": _REQ_MEMBER_MANAGE, - "workspaces.members.update_role": _REQ_ROLE_MANAGE, - "plugins.marketplace": _REQ_ACCOUNT_WORKSPACE_READ_MEMBER, - "plugins.list": _REQ_ACCOUNT_WORKSPACE_READ_MEMBER, - "plugins.install": _req_plugin_write(RBACPermission.PLUGIN_INSTALL), - "plugins.task": _REQ_PLUGIN_TASK, - "plugins.delete": _req_plugin_write(RBACPermission.PLUGIN_DELETE), - "model_providers.describe": _REQ_ACCOUNT_WORKSPACE_READ_MEMBER, - "model_providers.credential.create": _req_admin_write(RBACPermission.CREDENTIAL_CREATE), - "model_providers.credential.set": _req_admin_write(RBACPermission.CREDENTIAL_MANAGE), - "tools.list": _REQ_ACCOUNT_WORKSPACE_READ_MEMBER, - "tool_providers.describe": _REQ_ACCOUNT_WORKSPACE_READ_MEMBER, - "tool_providers.credential.create": _req_admin_write(RBACPermission.CREDENTIAL_CREATE), - "tool_providers.credential.set": _req_admin_write(RBACPermission.CREDENTIAL_MANAGE), - "knowledge_bases.list": _REQ_ACCOUNT_WORKSPACE_READ_MEMBER, - "models.list": _REQ_ACCOUNT_WORKSPACE_READ_MEMBER, - "models.credential.create": _req_admin_write(RBACPermission.CREDENTIAL_CREATE), - "models.credential.set": _req_admin_write(RBACPermission.CREDENTIAL_MANAGE), - "app_dsl.import": _REQ_DSL_WORKSPACE, - "app_dsl.check": _REQ_DSL_WORKSPACE, - "app_dsl.import_confirm": _REQ_DSL_WORKSPACE, - "app_dsl.export": _REQ_DSL_APP, - "app_dsl.check_dependencies": _REQ_DSL_APP, - "app_run.workflow": _RESOURCE_RUN, - "app_run.chat": _RESOURCE_RUN, - "app_run.advanced_chat": _RESOURCE_RUN, - "app_run.completion": _RESOURCE_RUN, - "app_run.draft.workflow": _draft_req(AppMode.WORKFLOW), - "app_run.draft.advanced_chat": _draft_req(AppMode.ADVANCED_CHAT), - "app_workflow.run.list": _REQ_RUN_HISTORY, - "app_workflow.run.describe": _REQ_RUN_HISTORY, - "app_workflow.run.nodes": _REQ_RUN_HISTORY, - "app_workflow.publish": _REQ_RELEASE, - "app_workflow.version.list": _REQ_VERSION_READ, - "app_workflow.version.restore": _REQ_RELEASE, - "app_workflow.release_check": _REQ_VERSION_READ, - "app_workflow.env.list": _REQ_VERSION_READ, - "app_workflow.env.set": _REQ_ENV_WRITE, - "app_workflow.env.delete": _REQ_ENV_WRITE, - "node_types.list": _REQ_NODE_TYPES, - "node_types.describe": _REQ_NODE_TYPES, - "app_create.workflow": _REQ_APP_CREATE, - "app_create.advanced_chat": _REQ_APP_CREATE, - "app_workflow.node_run.workflow": _draft_req(AppMode.WORKFLOW), - "app_workflow.node_run.advanced_chat": _draft_req(AppMode.ADVANCED_CHAT), - "app_info.describe.workflow": _settings_req( - (AppMode.WORKFLOW,), RBACPermission.APP_VIEW_LAYOUT, PlainApp(), Scope.APPS_READ, None - ), - "app_info.set.workflow": _settings_req( - (AppMode.WORKFLOW,), RBACPermission.APP_EDIT, PlainApp(), Scope.WORKSPACE_WRITE, _EDITOR_UP - ), - "service_api.describe": _settings_req( - REGULAR_MODES, RBACPermission.APP_RELEASE_AND_VERSION, PlainApp(), Scope.APPS_READ, _EDITOR_UP - ), - "service_api.set": _settings_req( - REGULAR_MODES, RBACPermission.APP_RELEASE_AND_VERSION, PlainApp(), Scope.WORKSPACE_WRITE, _OWNER_ADMIN - ), - "app_info.describe.advanced_chat": _settings_req( - (AppMode.ADVANCED_CHAT,), RBACPermission.APP_VIEW_LAYOUT, PlainApp(), Scope.APPS_READ, None - ), - "app_info.set.advanced_chat": _settings_req( - (AppMode.ADVANCED_CHAT,), RBACPermission.APP_EDIT, PlainApp(), Scope.WORKSPACE_WRITE, _EDITOR_UP - ), - "app_info.describe.chat": _settings_req( - (AppMode.CHAT,), RBACPermission.APP_VIEW_LAYOUT, PlainApp(), Scope.APPS_READ, None - ), - "app_info.set.chat": _settings_req( - (AppMode.CHAT,), RBACPermission.APP_EDIT, PlainApp(), Scope.WORKSPACE_WRITE, _EDITOR_UP - ), - "app_info.describe.agent_chat": _settings_req( - (AppMode.AGENT_CHAT,), RBACPermission.APP_VIEW_LAYOUT, PlainApp(), Scope.APPS_READ, None - ), - "app_info.set.agent_chat": _settings_req( - (AppMode.AGENT_CHAT,), RBACPermission.APP_EDIT, PlainApp(), Scope.WORKSPACE_WRITE, _EDITOR_UP - ), - "app_info.describe.completion": _settings_req( - (AppMode.COMPLETION,), RBACPermission.APP_VIEW_LAYOUT, PlainApp(), Scope.APPS_READ, None - ), - "app_info.set.completion": _settings_req( - (AppMode.COMPLETION,), RBACPermission.APP_EDIT, PlainApp(), Scope.WORKSPACE_WRITE, _EDITOR_UP - ), - "app_info.describe.agent": _settings_req( - (AppMode.AGENT,), RBACPermission.AGENT_PREVIEW, AgentBehindApp(), Scope.APPS_READ, None - ), - "app_info.set.agent": _settings_req( - (AppMode.AGENT,), RBACPermission.AGENT_EDIT, AgentBehindApp(), Scope.WORKSPACE_WRITE, _EDITOR_UP - ), - "service_api.describe.agent": _settings_req( - (AppMode.AGENT,), RBACPermission.AGENT_ACCESS_POINT_VIEW, AgentBehindApp(), Scope.APPS_READ, _EDITOR_UP - ), - "service_api.set.agent": _settings_req( - (AppMode.AGENT,), - RBACPermission.AGENT_ACCESS_POINT_MANAGE, - AgentBehindApp(), - Scope.WORKSPACE_WRITE, - _OWNER_ADMIN, - ), - "webapp.describe.workflow": _settings_req( - (AppMode.WORKFLOW,), RBACPermission.APP_VIEW_LAYOUT, PlainApp(), Scope.APPS_READ, None - ), - "webapp.set.workflow": _settings_req( - (AppMode.WORKFLOW,), RBACPermission.APP_RELEASE_AND_VERSION, PlainApp(), Scope.WORKSPACE_WRITE, _EDITOR_UP - ), - "webapp.reset": _settings_req( - REGULAR_MODES, RBACPermission.APP_RELEASE_AND_VERSION, PlainApp(), Scope.WORKSPACE_WRITE, _OWNER_ADMIN - ), - "webapp_access.describe": ( - *_settings_req(REGULAR_MODES, RBACPermission.APP_ACCESS_CONFIG, PlainApp(), Scope.APPS_READ, _EDITOR_UP), - CheckWebAppAuthEnterprise(), - ), - "webapp_access.set": ( - *_settings_req(REGULAR_MODES, RBACPermission.APP_ACCESS_CONFIG, PlainApp(), Scope.WORKSPACE_WRITE, _EDITOR_UP), - CheckWebAppAuthEnterprise(), - ), - "webapp.describe.advanced_chat": _settings_req( - (AppMode.ADVANCED_CHAT,), RBACPermission.APP_VIEW_LAYOUT, PlainApp(), Scope.APPS_READ, None - ), - "webapp.set.advanced_chat": _settings_req( - (AppMode.ADVANCED_CHAT,), RBACPermission.APP_RELEASE_AND_VERSION, PlainApp(), Scope.WORKSPACE_WRITE, _EDITOR_UP - ), - "webapp.describe.chat": _settings_req( - (AppMode.CHAT,), RBACPermission.APP_VIEW_LAYOUT, PlainApp(), Scope.APPS_READ, None - ), - "webapp.set.chat": _settings_req( - (AppMode.CHAT,), RBACPermission.APP_RELEASE_AND_VERSION, PlainApp(), Scope.WORKSPACE_WRITE, _EDITOR_UP - ), - "webapp.describe.agent_chat": _settings_req( - (AppMode.AGENT_CHAT,), RBACPermission.APP_VIEW_LAYOUT, PlainApp(), Scope.APPS_READ, None - ), - "webapp.set.agent_chat": _settings_req( - (AppMode.AGENT_CHAT,), RBACPermission.APP_RELEASE_AND_VERSION, PlainApp(), Scope.WORKSPACE_WRITE, _EDITOR_UP - ), - "webapp.describe.completion": _settings_req( - (AppMode.COMPLETION,), RBACPermission.APP_VIEW_LAYOUT, PlainApp(), Scope.APPS_READ, None - ), - "webapp.set.completion": _settings_req( - (AppMode.COMPLETION,), RBACPermission.APP_RELEASE_AND_VERSION, PlainApp(), Scope.WORKSPACE_WRITE, _EDITOR_UP - ), - "webapp.describe.agent": _settings_req( - (AppMode.AGENT,), RBACPermission.AGENT_ACCESS_POINT_VIEW, AgentBehindApp(), Scope.APPS_READ, None - ), - "webapp.set.agent": _settings_req( - (AppMode.AGENT,), RBACPermission.AGENT_ACCESS_POINT_MANAGE, AgentBehindApp(), Scope.WORKSPACE_WRITE, _EDITOR_UP - ), - "webapp.reset.agent": _settings_req( - (AppMode.AGENT,), - RBACPermission.AGENT_ACCESS_POINT_MANAGE, - AgentBehindApp(), - Scope.WORKSPACE_WRITE, - _OWNER_ADMIN, - ), - "webapp_access.describe.agent": ( - *_settings_req( - (AppMode.AGENT,), RBACPermission.AGENT_ACCESS_CONFIG, AgentBehindApp(), Scope.APPS_READ, _EDITOR_UP - ), - CheckWebAppAuthEnterprise(), - ), - "webapp_access.set.agent": ( - *_settings_req( - (AppMode.AGENT,), RBACPermission.AGENT_ACCESS_CONFIG, AgentBehindApp(), Scope.WORKSPACE_WRITE, _EDITOR_UP - ), - CheckWebAppAuthEnterprise(), - ), - "app_run.stop": _RESOURCE_RUN, - "files.upload": _REQ_FILES, - "human_input_form.get": _REQ_RUN_FORM, - "human_input_form.submit": _REQ_RUN_FORM, - "workflow_events.stream": _REQ_RUN, - "permitted_external.list": _REQ_EXTERNAL_LIST, - "permitted_external.describe": _REQ_EXTERNAL_DESCRIBE, -} -"""Every route's `@endpoint(requirements=...)`, in declared order. A route that -shares a tuple with a `MATRIX` representative is covered by that representative's -rows only while this table says it still declares the same tuple.""" - - -ROUTES_BY_ID = {route.id: route for route in ROUTES} - - -def _applicable(route: Route, case: Case) -> bool: - return CASE_REQUIRES[case] <= route.traits - - -def _rows(route: Route) -> Iterator[tuple[Route, Case, Expect]]: - for case in Case: - if not _applicable(route, case): - continue - if case not in ROUTER_CASES: - yield route, case, MATRIX[route.id][case] - elif route.id == ROUTER_CASE_ROUTE: - yield route, case, ROUTER_CASES[case] - - -ROWS: tuple[tuple[Route, Case, Expect], ...] = tuple( - row for route in ROUTES if route.id in MATRIX for row in _rows(route) -) - - -@dataclass(slots=True) -class World: - """Persisted rows and the tokens that address them.""" - - workspace_id: str - other_workspace_id: str - app_id: str - disabled_app_id: str - member_account_id: str - low_role_account_id: str - outsider_account_id: str - tokens: dict[Bearer, str] = field(default_factory=dict) - session_ids: dict[Bearer, str] = field(default_factory=dict) - - -def _admit_on_mount(*_args: object, **_kwargs: object) -> None: - raise ImATeapot("admitted") - - -class _MemoryResolver: - """Stands in for the DB-backed token resolver. - - Token resolution is an input to the matrix, not part of it, and the shipped - resolver compares a timezone-aware expiry against a column SQLite hands back - naive. - """ - - def __init__(self, rows: dict[str, ResolvedRow]) -> None: - self._rows = rows - - def resolve(self, token_hash: str) -> ResolvedRow | None: - return self._rows.get(token_hash) - - -def _authenticator(rows: dict[str, ResolvedRow]) -> BearerAuthenticator: - resolver = _MemoryResolver(rows) - return BearerAuthenticator({TokenType.OAUTH_ACCOUNT: resolver, TokenType.OAUTH_EXTERNAL_SSO: resolver}) - - -@pytest.fixture -def token_rows() -> dict[str, ResolvedRow]: - return {} - - -@pytest.fixture(scope="module") -def _matrix_app() -> Flask: - """Register the matrix's unchanging routes once; each case gets a fresh client.""" - app = create_flask_app_with_configs() - app.config["TESTING"] = True - app.secret_key = "openapi-auth-matrix" - LoginManager(app) - app.register_blueprint(openapi_bp) - return app - - -@pytest.fixture -def matrix_app( - _matrix_app: Flask, monkeypatch: pytest.MonkeyPatch, account_application_services: ApplicationServices -) -> Iterator[Flask]: - """Keep rate-limit overrides and the admission probe scoped to each case.""" - - monkeypatch.setitem(_matrix_app.extensions, "application_services", account_application_services) - monkeypatch.setattr( - rate_limit_module, - "LIMIT_BEARER_PER_TOKEN", - rate_limit_module.RateLimit( - 0, - rate_limit_module.LIMIT_BEARER_PER_TOKEN.window, - rate_limit_module.LIMIT_BEARER_PER_TOKEN.scopes, - ), - ) - - user_logged_in.connect(_admit_on_mount) - try: - yield _matrix_app - finally: - user_logged_in.disconnect(_admit_on_mount) - - -@pytest.fixture -def world(sqlite_session_factory: sessionmaker[Session], token_rows: dict[str, ResolvedRow]) -> World: - """One workspace with an owner, a normal member and an outsider; two apps.""" - built = World( - workspace_id=str(uuid.uuid4()), - other_workspace_id=str(uuid.uuid4()), - app_id=str(uuid.uuid4()), - disabled_app_id=str(uuid.uuid4()), - member_account_id=str(uuid.uuid4()), - low_role_account_id=str(uuid.uuid4()), - outsider_account_id=str(uuid.uuid4()), - ) - - def account(account_id: str, email: str) -> Account: - row = Account(name=email, email=email, avatar="", status=AccountStatus.ACTIVE) - row.id = account_id - return row - - def application(app_id: str, *, enable_api: bool) -> App: - return App( - id=app_id, - tenant_id=built.workspace_id, - name="matrix app", - description="", - mode=AppMode.WORKFLOW, - enable_site=False, - enable_api=enable_api, - ) - - workspace = Tenant(name="matrix workspace") - workspace.id = built.workspace_id - other_workspace = Tenant(name="other workspace") - other_workspace.id = built.other_workspace_id - - rows: list[object] = [ - workspace, - other_workspace, - account(built.member_account_id, "owner@example.com"), - account(built.low_role_account_id, "normal@example.com"), - account(built.outsider_account_id, "outsider@example.com"), - application(built.app_id, enable_api=True), - application(built.disabled_app_id, enable_api=False), - TenantAccountJoin( - tenant_id=built.workspace_id, - account_id=built.member_account_id, - role=TenantAccountRole.OWNER, - ), - TenantAccountJoin( - tenant_id=built.workspace_id, - account_id=built.low_role_account_id, - role=TenantAccountRole.NORMAL, - ), - ] - - def mint(bearer: Bearer, prefix: str, *, account_id: str | None, email: str) -> None: - """Register a bearer with the fake resolver *and* persist the session row it - names, so the access service's ownership check has something real to match against. - """ - raw = prefix + uuid.uuid4().hex - token_id = uuid.uuid4() - issuer = "dify:account" if account_id else "https://idp.example" - token_rows[sha256_hex(raw)] = ResolvedRow( - subject_email=email, - subject_issuer=issuer, - account_id=uuid.UUID(account_id) if account_id else None, - client_id="difyctl", - token_id=token_id, - expires_at=None, - ) - row = OAuthAccessToken( - token_hash=sha256_hex(raw), - prefix=prefix, - account_id=account_id, - subject_email=email, - subject_issuer=issuer, - client_id="difyctl", - device_label="matrix", - expires_at=datetime.now(UTC) + timedelta(days=365), - ) - row.id = str(token_id) - rows.append(row) - built.tokens[bearer] = raw - built.session_ids[bearer] = str(token_id) - - mint( - Bearer.ACCOUNT_MEMBER, - TokenType.OAUTH_ACCOUNT.prefix, - account_id=built.member_account_id, - email="owner@example.com", - ) - mint( - Bearer.ACCOUNT_LOW_ROLE, - TokenType.OAUTH_ACCOUNT.prefix, - account_id=built.low_role_account_id, - email="normal@example.com", - ) - mint( - Bearer.ACCOUNT_OUTSIDER, - TokenType.OAUTH_ACCOUNT.prefix, - account_id=built.outsider_account_id, - email="outsider@example.com", - ) - mint(Bearer.EXTERNAL, TokenType.OAUTH_EXTERNAL_SSO.prefix, account_id=None, email="external@example.com") - with sqlite_session_factory.begin() as session: - session.add_all(rows) - return built - - -def _system_features( - *, edition: DeploymentEdition, webapp_auth: bool, license_status: LicenseStatus -) -> SystemFeatureModel: - features = SystemFeatureModel(agent_tracing_api_collection_enabled=False, deployment_edition=edition) - features.webapp_auth.enabled = webapp_auth - features.license.status = license_status - return features - - -@dataclass(frozen=True, slots=True) -class _WebAppSettings: - access_mode: str - - -def _access_mode_settings(access_mode: WebAppAccessMode | None) -> _WebAppSettings | None: - if access_mode is None: - return None - return _WebAppSettings(access_mode=access_mode.value) - - -def _webapp_account() -> Account: - """The account an external-SSO subject's email resolves to for the private-app check.""" - row = Account(name="external", email="external@example.com", avatar="", status=AccountStatus.ACTIVE) - row.id = str(uuid.uuid4()) - return row - - -def _end_user(_type: EndUserType, tenant_id: str, app_id: str, user_id: str | None = None) -> EndUser: - row = EndUser( - tenant_id=tenant_id, - app_id=app_id, - type=EndUserType.OPENAPI, - _is_anonymous=False, - session_id=user_id or "", - ) - row.external_user_id = user_id - return row - - -def _bearer_for(route: Route, scenario: Scenario) -> Bearer | None: - account_primary = Trait.ACCOUNT_PRIMARY in route.traits - match scenario.bearer: - case Bearer.NONE: - return None - case Bearer.PRIMARY: - return Bearer.ACCOUNT_MEMBER if account_primary else Bearer.EXTERNAL - case Bearer.OTHER_SUBJECT: - return Bearer.EXTERNAL if account_primary else Bearer.ACCOUNT_MEMBER - case _: - return scenario.bearer - - -def _url(route: Route, world: World, scenario: Scenario, bearer: Bearer | None) -> str: - if scenario.unknown_app: - app_id = str(uuid.uuid4()) - elif not scenario.app_api_enabled: - app_id = world.disabled_app_id - else: - app_id = world.app_id - ids = { - "app_id": app_id, - "workspace_id": world.workspace_id, - # The caller's own session: the access service 404s any other id. - "session_id": world.session_ids.get(bearer, str(uuid.uuid4())) if bearer else str(uuid.uuid4()), - "member_id": str(uuid.uuid4()), - "import_id": str(uuid.uuid4()), - "task_id": str(uuid.uuid4()), - "form_token": uuid.uuid4().hex, - "run_id": str(uuid.uuid4()), - "version_id": str(uuid.uuid4()), - "env_id": str(uuid.uuid4()), - "node_type": "llm", - "node_id": "node-1", - "plugin_id": "langgenius/openai", - "provider": "langgenius/openai/openai", - "credential_id": str(uuid.uuid4()), - "model_type": "llm", - } - query = route.query.format(**ids) - if scenario.foreign_workspace_query: - extra = f"workspace_id={world.other_workspace_id}" - query = f"{query}&{extra}" if query else extra - path = route.path.format(**ids) - return f"/openapi/v1{path}?{query}" if query else f"/openapi/v1{path}" - - -def _run_case( - *, - route: Route, - scenario: Scenario, - app: Flask, - world: World, - token_rows: dict[str, ResolvedRow], - monkeypatch: pytest.MonkeyPatch, -) -> TestResponse: - edition = scenario.edition or ( - DeploymentEdition.ENTERPRISE if Trait.ENTERPRISE_ONLY in route.traits else DeploymentEdition.COMMUNITY - ) - apply_config_overrides(monkeypatch, DEPLOYMENT_EDITION=edition, RBAC_ENABLED=scenario.rbac_enabled) - monkeypatch.setattr(oauth_bearer_module, "_authenticator", _authenticator(token_rows)) - if scenario.narrow_scopes: - # No shipped token can fail `CheckScope`: `dfoa_` carries `Scope.FULL` and - # `dfoe_` exactly the scopes its routes ask for. Scopes are a constant of - # the subject, so the only way to reach the refusal is to blank them here. - # Keep these rows: they are what catches a third token kind routed - # somewhere it was not scoped for. - monkeypatch.setattr(oauth_bearer_module.AuthContext, "scopes", property(lambda _self: frozenset())) - - features = _system_features( - edition=edition, webapp_auth=scenario.webapp_auth, license_status=scenario.license_status - ) - settings = _access_mode_settings(scenario.access_mode) - - headers: dict[str, str] = {CATALOG_HEADER: catalog_for(app)[1]} - bearer = _bearer_for(route, scenario) - if bearer is not None: - headers["Authorization"] = f"Bearer {world.tokens[bearer]}" - - with ExitStack() as stack: - stack.enter_context(patch.object(SystemFeatureService, "get_public_system_features", return_value=features)) - stack.enter_context( - patch.object(EnterpriseService.WebAppAuth, "get_app_access_mode_by_id", return_value=settings) - ) - stack.enter_context( - patch.object( - EnterpriseService.WebAppAuth, - "is_user_allowed_to_access_webapp", - return_value=scenario.private_app_permitted, - ) - ) - stack.enter_context( - patch.object( - EnterpriseService.WebAppAuth, - "list_externally_accessible_apps", - return_value={"data": [], "total": 0, "hasMore": False}, - ) - ) - services = subjects.application_services() - stack.enter_context( - patch.object( - services.app_scoped_end_users.commands, - "get_or_create_end_user_by_type", - side_effect=_end_user, - ) - ) - stack.enter_context( - patch( - "services.enterprise.rbac_service.RBACService.CheckAccess.check", - return_value=scenario.rbac_allows, - ) - ) - stack.enter_context( - patch.object( - services.accounts.identity, - "get_account_by_email", - return_value=_webapp_account(), - ) - ) - client = app.test_client() - return contextvars.copy_context().run( - lambda: client.open(_url(route, world, scenario, bearer), method=route.method, headers=headers) - ) - - -@pytest.mark.parametrize( - ("route", "case", "expected"), - ROWS, - ids=[f"{route.id}-{case.value}" for route, case, _ in ROWS], -) -def test_allow_deny_matrix( - route: Route, - case: Case, - expected: Expect, - matrix_app: Flask, - world: World, - token_rows: dict[str, ResolvedRow], - monkeypatch: pytest.MonkeyPatch, -) -> None: - response = _run_case( - route=route, - scenario=SCENARIOS[case], - app=matrix_app, - world=world, - token_rows=token_rows, - monkeypatch=monkeypatch, - ) - - body = response.get_json(silent=True) or {} - assert response.status_code == expected.status, f"{route.id}/{case.value}: body={body}" - if expected.message is not None: - assert body.get("message") == expected.message - - -_PATH_PARAMS: Final = frozenset({"plugin_id", "provider"}) - - -def _rule_path(route: Route) -> str: - path = route.path - for name in _PATH_PARAMS: - path = path.replace(f"{{{name}}}", f"") - return "/openapi/v1" + path.replace("{", "") - - -def _endpoint_spec(app: Flask, rule: Rule, method: str) -> EndpointSpec | None: - """The `EndpointSpec` `@endpoint` attached to this rule's handler, or None. - - Carrying a spec *is* being guarded: `@endpoint` is the only thing that sets - it, and it sets it on the same object it wraps in `subject_router.guard`. - """ - view = app.view_functions.get(rule.endpoint) - resource = view.view_class if hasattr(view, "view_class") else None - if resource is None: - return None - handler = getattr(resource, method.lower(), None) # guard-ignore: no-new-getattr -- HTTP verb selects the handler - return handler.__spec__ if hasattr(handler, "__spec__") else None - - -def test_registered_openapi_routes_match_the_matrix(matrix_app: Flask) -> None: - """Every guarded /openapi/v1 route is in the table, and nothing else is. - - Guarded-ness is derived from `view.__spec__`, not from a hand-kept list of - exemptions — so a route that gains a guard without gaining a row fails here, - and so does a guarded route that quietly stops carrying a spec. Neither can - be silenced by editing this file. - """ - guarded: set[tuple[str, str]] = set() - unguarded: set[tuple[str, str]] = set() - for rule in matrix_app.url_map.iter_rules(): - if not str(rule).startswith("/openapi/v1") or str(rule) == "/openapi/v1/": - continue - for method in rule.methods or set(): - if method in {"HEAD", "OPTIONS"}: - continue - entry = (method, str(rule)) - target = guarded if _endpoint_spec(matrix_app, rule, method) is not None else unguarded - target.add(entry) - - expected = {(route.method, _rule_path(route)) for route in ROUTES} - assert guarded == expected - # The remainder is the device-flow, documentation and catalog surface. Its size is - # pinned rather than enumerated: a list of exemptions rots silently — the one this - # replaced still named `swagger.json`, a route that is not registered at all. - assert len(unguarded) == 14 - - -@singledispatch -def _config(_requirement: Requirement) -> object: - return None - - -@_config.register -def _(requirement: CheckSubject) -> object: - return requirement.allowed - - -@_config.register -def _(requirement: CheckAppMode) -> object: - return requirement.modes - - -@_config.register -def _(requirement: CheckScope) -> object: - return requirement.scope - - -@_config.register -def _(requirement: CheckRBACPermission) -> object: - return tuple((check.scene, repr(check.locator)) for check in requirement.checks) - - -@_config.register -def _(requirement: CheckWorkspaceRole) -> object: - return requirement.allowed_roles - - -def _canonical(requirement: Requirement) -> tuple[type[Requirement], object]: - """`Requirement` has no `__eq__`, so compare by type plus the config that sets it apart.""" - return type(requirement), _config(requirement) - - -def _spec_for(app: Flask, route: Route) -> EndpointSpec | None: - for rule in app.url_map.iter_rules(): - if str(rule) == _rule_path(route) and route.method in (rule.methods or set()): - return _endpoint_spec(app, rule, route.method) - return None - - -@pytest.mark.parametrize("route", ROUTES, ids=[route.id for route in ROUTES]) -def test_route_declares_expected_requirements(route: Route, matrix_app: Flask) -> None: - """The guard that lets `MATRIX` run one representative per tuple: a route - that drifts from the tuple its representative was chosen for fails here. - """ - spec = _spec_for(matrix_app, route) - assert spec is not None, route.id - assert [*map(_canonical, spec.requirements)] == [*map(_canonical, DECLARED[route.id])] - - -ERROR_DEFAULT_RESPONSE: dict[str, object] = { - "description": "Error", - "content": {"application/json": {"schema": {"$ref": "#/components/schemas/ErrorBody"}}}, -} -"""Exactly what `@returns` registers as `("default", "Error", ErrorBody)`, and what -`ExternalApi` adds to every other operation.""" - -EXPECTED_RESPONSE_CODES: dict[tuple[str, str], frozenset[str]] = { - ("get", "/workspaces/{workspace_id}/knowledge-fs/knowledge-spaces/{knowledge_space_id}/fs:ls"): frozenset( - {"200", "400", "401", "403", "404", "409", "413", "422", "503", "default"} - ), - ("get", "/workspaces/{workspace_id}/knowledge-fs/knowledge-spaces/{knowledge_space_id}/fs:tree"): frozenset( - {"200", "400", "401", "403", "404", "409", "413", "422", "503", "default"} - ), - ("get", "/workspaces/{workspace_id}/knowledge-fs/knowledge-spaces/{knowledge_space_id}/fs:grep"): frozenset( - {"200", "400", "401", "403", "404", "409", "413", "422", "503", "default"} - ), - ("get", "/workspaces/{workspace_id}/knowledge-fs/knowledge-spaces/{knowledge_space_id}/fs:find"): frozenset( - {"200", "400", "401", "403", "404", "409", "413", "422", "503", "default"} - ), - ("post", "/workspaces/{workspace_id}/knowledge-fs/knowledge-spaces/{knowledge_space_id}/fs:diff"): frozenset( - {"200", "400", "401", "403", "404", "409", "413", "422", "503", "default"} - ), - ("get", "/workspaces/{workspace_id}/knowledge-fs/knowledge-spaces/{knowledge_space_id}/fs:cat"): frozenset( - {"200", "400", "401", "403", "404", "409", "413", "422", "503", "default"} - ), - ("get", "/workspaces/{workspace_id}/knowledge-fs/knowledge-spaces/{knowledge_space_id}/fs:stat"): frozenset( - {"200", "400", "401", "403", "404", "409", "413", "422", "503", "default"} - ), - ("get", "/_catalog"): frozenset({"200", "default"}), - ("get", "/workspaces/{workspace_id}/marketplace/plugins"): frozenset({"200", "422", "default"}), - ("get", "/workspaces/{workspace_id}/plugins"): frozenset({"200", "422", "default"}), - ("post", "/workspaces/{workspace_id}/plugins:install"): frozenset({"200", "422", "default"}), - ("get", "/workspaces/{workspace_id}/plugin-tasks/{task_id}"): frozenset({"200", "default"}), - ("delete", "/workspaces/{workspace_id}/plugins/{plugin_id}"): frozenset({"200", "default"}), - ("get", "/workspaces/{workspace_id}/tools"): frozenset({"200", "422", "default"}), - ("get", "/workspaces/{workspace_id}/tool-providers/{provider}"): frozenset({"200", "default"}), - ("post", "/workspaces/{workspace_id}/tool-providers/{provider}/credentials"): frozenset({"201", "422", "default"}), - ("patch", "/workspaces/{workspace_id}/tool-providers/{provider}/credentials/{credential_id}"): frozenset( - {"200", "422", "default"} - ), - ("get", "/workspaces/{workspace_id}/model-providers/{provider}"): frozenset({"200", "default"}), - ("post", "/workspaces/{workspace_id}/model-providers/{provider}/credentials"): frozenset({"201", "422", "default"}), - ("patch", "/workspaces/{workspace_id}/model-providers/{provider}/credentials/{credential_id}"): frozenset( - {"200", "422", "default"} - ), - ("get", "/workspaces/{workspace_id}/knowledge-bases"): frozenset({"200", "422", "default"}), - ("get", "/workspaces/{workspace_id}/models"): frozenset({"200", "422", "default"}), - ("post", "/workspaces/{workspace_id}/model-providers/{provider}/models/credentials"): frozenset( - {"201", "422", "default"} - ), - ("patch", "/workspaces/{workspace_id}/model-providers/{provider}/models/credentials/{credential_id}"): frozenset( - {"200", "422", "default"} - ), - ("get", "/_health"): frozenset({"200", "default"}), - ("get", "/_version"): frozenset({"200", "default"}), - ("get", "/account"): frozenset({"200", "default"}), - ("get", "/account/sessions"): frozenset({"200", "422", "default"}), - ("delete", "/account/sessions/self"): frozenset({"200", "default"}), - ("delete", "/account/sessions/{session_id}"): frozenset({"200", "default"}), - ("get", "/apps"): frozenset({"200", "422", "default"}), - ("get", "/apps/{app_id}"): frozenset({"200", "422", "default"}), - ("get", "/apps/{app_id}/app-info/workflow"): frozenset({"200", "default"}), - ("patch", "/apps/{app_id}/app-info/workflow"): frozenset({"200", "422", "default"}), - ("get", "/apps/{app_id}/service-api"): frozenset({"200", "default"}), - ("patch", "/apps/{app_id}/service-api"): frozenset({"200", "422", "default"}), - ("get", "/apps/{app_id}/app-info/advanced-chat"): frozenset({"200", "default"}), - ("patch", "/apps/{app_id}/app-info/advanced-chat"): frozenset({"200", "422", "default"}), - ("get", "/apps/{app_id}/app-info/chat"): frozenset({"200", "default"}), - ("patch", "/apps/{app_id}/app-info/chat"): frozenset({"200", "422", "default"}), - ("get", "/apps/{app_id}/app-info/agent-chat"): frozenset({"200", "default"}), - ("patch", "/apps/{app_id}/app-info/agent-chat"): frozenset({"200", "422", "default"}), - ("get", "/apps/{app_id}/app-info/completion"): frozenset({"200", "default"}), - ("patch", "/apps/{app_id}/app-info/completion"): frozenset({"200", "422", "default"}), - ("get", "/apps/{app_id}/app-info/agent"): frozenset({"200", "default"}), - ("patch", "/apps/{app_id}/app-info/agent"): frozenset({"200", "422", "default"}), - ("get", "/apps/{app_id}/service-api/agent"): frozenset({"200", "default"}), - ("patch", "/apps/{app_id}/service-api/agent"): frozenset({"200", "422", "default"}), - ("get", "/apps/{app_id}/dependencies:check"): frozenset({"200", "default"}), - ("get", "/apps/{app_id}/dsl"): frozenset({"200", "422", "default"}), - ("post", "/apps/{app_id}/files"): frozenset({"201", "400", "401", "413", "415", "422", "default"}), - ("get", "/apps/{app_id}/human-input-forms/{form_token}"): frozenset({"200", "default"}), - ("post", "/apps/{app_id}/human-input-forms/{form_token}:submit"): frozenset({"200", "422", "default"}), - ("get", "/apps/{app_id}/tasks/{task_id}/events"): frozenset({"200", "422", "default"}), - ("post", "/apps/{app_id}/tasks/{task_id}:stop"): frozenset({"200", "default"}), - ("post", "/apps/{app_id}/workflow:run"): frozenset({"200", "422", "default"}), - ("post", "/apps/{app_id}/chat:run"): frozenset({"200", "422", "default"}), - ("post", "/apps/{app_id}/advanced-chat:run"): frozenset({"200", "422", "default"}), - ("post", "/apps/{app_id}/completion:run"): frozenset({"200", "422", "default"}), - ("post", "/apps/{app_id}/draft/workflow:run"): frozenset({"200", "422", "default"}), - ("post", "/apps/{app_id}/draft/advanced-chat:run"): frozenset({"200", "422", "default"}), - ("get", "/apps/{app_id}/runs"): frozenset({"200", "422", "default"}), - ("get", "/apps/{app_id}/runs/{run_id}"): frozenset({"200", "default"}), - ("get", "/apps/{app_id}/runs/{run_id}/nodes"): frozenset({"200", "default"}), - ("post", "/apps/{app_id}:publish"): frozenset({"200", "422", "default"}), - ("get", "/apps/{app_id}/versions"): frozenset({"200", "422", "default"}), - ("post", "/apps/{app_id}/versions/{version_id}:restore"): frozenset({"200", "default"}), - ("get", "/apps/{app_id}/release:check"): frozenset({"200", "default"}), - ("get", "/apps/{app_id}/env"): frozenset({"200", "default"}), - ("put", "/apps/{app_id}/env/{env_id}"): frozenset({"200", "422", "default"}), - ("delete", "/apps/{app_id}/env/{env_id}"): frozenset({"200", "default"}), - # The catalog and the five device-flow rows document their 200 with a raw - # `openapi_ns.response` rather than `@returns`; `ExternalApi` still gives them the - # `default` error every operation documents, with the same `ErrorBody` schema. - # Approve and deny run Console account admission; its setup and account-initialization - # errors are generic on that surface (handled centrally by the client), so the guards - # document nothing and the open `default` error covers them. Code, token, approve and - # deny declare the RFC 8628-style `DeviceFlowErrorResponse` for the statuses their flow - # errors use. - ("post", "/oauth/device/approve"): frozenset({"200", "404", "409", "503", "default"}), - ("post", "/oauth/device/code"): frozenset({"200", "400", "default"}), - ("post", "/oauth/device/deny"): frozenset({"200", "404", "409", "503", "default"}), - ("get", "/oauth/device/lookup"): frozenset({"200", "default"}), - ("post", "/oauth/device/token"): frozenset({"200", "400", "default"}), - ("get", "/node-types"): frozenset({"200", "default"}), - ("get", "/node-types/{node_type}"): frozenset({"200", "default"}), - ("post", "/apps/{app_id}/draft/workflow/nodes/{node_id}:run"): frozenset({"200", "422", "default"}), - ("post", "/apps/{app_id}/draft/advanced-chat/nodes/{node_id}:run"): frozenset({"200", "422", "default"}), - ("post", "/workspaces/{workspace_id}/apps/workflow"): frozenset({"201", "422", "default"}), - ("post", "/workspaces/{workspace_id}/apps/advanced-chat"): frozenset({"201", "422", "default"}), - ("get", "/permitted-external-apps"): frozenset({"200", "422", "default"}), - ("get", "/permitted-external-apps/{app_id}"): frozenset({"200", "422", "default"}), - ("get", "/workspaces"): frozenset({"200", "422", "default"}), - ("get", "/workspaces/{workspace_id}"): frozenset({"200", "default"}), - ("post", "/workspaces/{workspace_id}/apps/imports"): frozenset({"200", "202", "400", "422", "default"}), - ("post", "/workspaces/{workspace_id}/apps/imports:check"): frozenset({"200", "422", "default"}), - ("post", "/workspaces/{workspace_id}/apps/imports/{import_id}:confirm"): frozenset({"200", "400", "default"}), - ("get", "/workspaces/{workspace_id}/members"): frozenset({"200", "422", "default"}), - ("post", "/workspaces/{workspace_id}/members"): frozenset({"201", "422", "default"}), - ("delete", "/workspaces/{workspace_id}/members/{member_id}"): frozenset({"200", "default"}), - ("patch", "/workspaces/{workspace_id}/members/{member_id}"): frozenset({"200", "422", "default"}), - ("get", "/apps/{app_id}/webapp/workflow"): frozenset({"200", "default"}), - ("patch", "/apps/{app_id}/webapp/workflow"): frozenset({"200", "422", "default"}), - ("post", "/apps/{app_id}/webapp:reset"): frozenset({"200", "default"}), - ("get", "/apps/{app_id}/webapp-access"): frozenset({"200", "default"}), - ("put", "/apps/{app_id}/webapp-access"): frozenset({"200", "422", "default"}), - ("get", "/apps/{app_id}/webapp/advanced-chat"): frozenset({"200", "default"}), - ("patch", "/apps/{app_id}/webapp/advanced-chat"): frozenset({"200", "422", "default"}), - ("get", "/apps/{app_id}/webapp/chat"): frozenset({"200", "default"}), - ("patch", "/apps/{app_id}/webapp/chat"): frozenset({"200", "422", "default"}), - ("get", "/apps/{app_id}/webapp/agent-chat"): frozenset({"200", "default"}), - ("patch", "/apps/{app_id}/webapp/agent-chat"): frozenset({"200", "422", "default"}), - ("get", "/apps/{app_id}/webapp/completion"): frozenset({"200", "default"}), - ("patch", "/apps/{app_id}/webapp/completion"): frozenset({"200", "422", "default"}), - ("get", "/apps/{app_id}/webapp/agent"): frozenset({"200", "default"}), - ("patch", "/apps/{app_id}/webapp/agent"): frozenset({"200", "422", "default"}), - ("post", "/apps/{app_id}/webapp/agent:reset"): frozenset({"200", "default"}), - ("get", "/apps/{app_id}/webapp-access/agent"): frozenset({"200", "default"}), - ("put", "/apps/{app_id}/webapp-access/agent"): frozenset({"200", "422", "default"}), -} - -_HTTP_METHODS = frozenset({"get", "post", "put", "patch", "delete", "head", "options", "trace"}) - - -@pytest.fixture -def openapi_document(config_overrides: Callable[..., None]) -> dict[str, object]: - config_overrides(SWAGGER_UI_ENABLED=True) - app = Flask(__name__) - app.config["TESTING"] = True - app.register_blueprint(openapi_bp) - response = app.test_client().get("/openapi/v1/openapi.json") - assert response.status_code == 200 - return response.get_json() - - -def test_file_bearing_bodies_are_documented_as_multipart(openapi_document: dict[str, object]) -> None: - """The exported contract names the wire form `_multipart.py` accepts: a body of only - files is multipart alone, a run body is JSON or multipart with its JSON-text parts marked.""" - paths = openapi_document["paths"] - assert isinstance(paths, dict) - upload = paths["/apps/{app_id}/files"]["post"]["requestBody"]["content"] - assert set(upload) == {"multipart/form-data"} - assert "encoding" not in upload["multipart/form-data"] - run = paths["/apps/{app_id}/chat:run"]["post"]["requestBody"]["content"] - assert set(run) == {"application/json", "multipart/form-data"} - assert run["multipart/form-data"]["schema"] == run["application/json"]["schema"] - encoding = run["multipart/form-data"]["encoding"] - assert {"inputs", "query"} <= set(encoding) - assert {"files", "attachments"}.isdisjoint(encoding) - assert encoding["inputs"] == {"contentType": "application/json"} - - -def _operations(document: dict[str, object]) -> Iterator[tuple[tuple[str, str], dict[str, object]]]: - paths = document["paths"] - assert isinstance(paths, dict) - for path, item in paths.items(): - assert isinstance(item, dict) - for method, operation in item.items(): - if method in _HTTP_METHODS: - assert isinstance(operation, dict) - yield (method, path), operation - - -def test_openapi_document_operations_and_response_codes(openapi_document: dict[str, object]) -> None: - """Every operation's response codes are pinned exactly. - - `@endpoint` routes through `@returns`, which registers `default` -> `ErrorBody`, - so a route that stops emitting the error schema (or one that starts) moves a row - here and has to be re-pinned deliberately. - """ - seen = dict(_operations(openapi_document)) - assert set(seen) == set(EXPECTED_RESPONSE_CODES) - for key, operation in seen.items(): - responses = operation.get("responses", {}) - assert isinstance(responses, dict) - assert frozenset(responses) == EXPECTED_RESPONSE_CODES[key], ( - f"{key}: {sorted(responses)} != {sorted(EXPECTED_RESPONSE_CODES[key])}" - ) - - -def test_every_default_response_is_the_one_returns_registers(openapi_document: dict[str, object]) -> None: - """A hand-written `default` pointing at some other schema would still satisfy - the code sets above, which compare codes and not bodies. - """ - for key, operation in _operations(openapi_document): - responses = operation["responses"] - assert isinstance(responses, dict) - if "default" in responses: - assert responses["default"] == ERROR_DEFAULT_RESPONSE, key diff --git a/api/tests/unit_tests/controllers/openapi/test_cors.py b/api/tests/unit_tests/controllers/openapi/test_cors.py deleted file mode 100644 index f96d8d3af5c..00000000000 --- a/api/tests/unit_tests/controllers/openapi/test_cors.py +++ /dev/null @@ -1,154 +0,0 @@ -"""CORS posture for /openapi/v1/* — default empty allowlist (same-origin), -expandable via OPENAPI_CORS_ALLOW_ORIGINS. Cross-origin requests from -disallowed origins do not receive the Access-Control-Allow-Origin -header, which the browser then blocks. - -Tests use a fresh Blueprint + Flask-CORS per case because the production -blueprint is a module-level singleton and can't be reconfigured once -registered. -""" - -import builtins - -from flask import Blueprint, Flask -from flask.views import MethodView -from flask_cors import CORS -from flask_restx import Resource - -from configs import dify_config -from constants import HEADER_NAME_CATALOG -from extensions.ext_blueprints import EXPOSED_HEADERS, OPENAPI_HEADERS, OPENAPI_MAX_AGE_SECONDS -from libs.external_api import ExternalApi - -if not hasattr(builtins, "MethodView"): - builtins.MethodView = MethodView # type: ignore[attr-defined] - - -def _make_app(allowed_origins: list[str], blueprint_name: str) -> Flask: - """Build a Flask app with a fresh openapi-style blueprint mirroring - production CORS settings, parameterised on the origin allowlist. - """ - bp = Blueprint(blueprint_name, __name__, url_prefix="/openapi/v1") - api = ExternalApi(bp, version="1.0", title="OpenAPI Test", description="") - - @api.route("/_health") - class _Health(Resource): - def get(self) -> dict[str, bool]: - return {"ok": True} - - CORS( - bp, - resources={r"/*": {"origins": allowed_origins}}, - supports_credentials=True, - allow_headers=list(OPENAPI_HEADERS), - methods=["GET", "POST", "PATCH", "DELETE", "OPTIONS"], - expose_headers=list(EXPOSED_HEADERS), - max_age=OPENAPI_MAX_AGE_SECONDS, - ) - - app = Flask(__name__) - app.config["TESTING"] = True - app.register_blueprint(bp) - return app - - -def test_default_openapi_cors_allowlist_is_empty() -> None: - """Default config admits no cross-origin until operator opts in.""" - assert dify_config.OPENAPI_CORS_ALLOW_ORIGINS == [] - - -def test_preflight_allowed_origin_returns_cors_headers() -> None: - app = _make_app(["https://app.example.com"], "openapi_t1") - client = app.test_client() - response = client.options( - "/openapi/v1/_health", - headers={ - "Origin": "https://app.example.com", - "Access-Control-Request-Method": "GET", - }, - ) - - assert response.headers.get("Access-Control-Allow-Origin") == "https://app.example.com" - assert response.headers.get("Access-Control-Max-Age") == str(OPENAPI_MAX_AGE_SECONDS) - - -def test_preflight_disallowed_origin_omits_cors_headers() -> None: - app = _make_app(["https://app.example.com"], "openapi_t2") - client = app.test_client() - response = client.options( - "/openapi/v1/_health", - headers={ - "Origin": "https://attacker.example", - "Access-Control-Request-Method": "GET", - }, - ) - - # flask-cors omits Allow-Origin for disallowed origins; browser blocks. - assert "Access-Control-Allow-Origin" not in response.headers - - -def test_preflight_with_default_empty_allowlist_omits_cors_headers() -> None: - app = _make_app([], "openapi_t3") - client = app.test_client() - response = client.options( - "/openapi/v1/_health", - headers={ - "Origin": "https://app.example.com", - "Access-Control-Request-Method": "GET", - }, - ) - - assert "Access-Control-Allow-Origin" not in response.headers - - -def test_same_origin_request_succeeds_without_origin_header() -> None: - app = _make_app(["https://app.example.com"], "openapi_t4") - client = app.test_client() - # Browsers don't send Origin on same-origin GETs. - response = client.get("/openapi/v1/_health") - - assert response.status_code == 200 - assert response.get_json() == {"ok": True} - - -def test_authorization_header_is_in_allow_headers() -> None: - """Bearer-authed routes need Authorization in the preflight response.""" - app = _make_app(["https://app.example.com"], "openapi_t5") - client = app.test_client() - response = client.options( - "/openapi/v1/_health", - headers={ - "Origin": "https://app.example.com", - "Access-Control-Request-Method": "GET", - "Access-Control-Request-Headers": "Authorization", - }, - ) - - allow_headers = response.headers.get("Access-Control-Allow-Headers", "").lower() - assert "authorization" in allow_headers - - -def test_catalog_fingerprint_header_is_exposed() -> None: - app = _make_app(["https://app.example.com"], "openapi_t6") - client = app.test_client() - response = client.get("/openapi/v1/_health", headers={"Origin": "https://app.example.com"}) - - expose_headers = response.headers.get("Access-Control-Expose-Headers", "") - assert HEADER_NAME_CATALOG in expose_headers - - -def test_catalog_fingerprint_header_is_allowed_on_requests() -> None: - """Every guarded route needs the fingerprint on the request, so preflight must admit it.""" - app = _make_app(["https://app.example.com"], "openapi_t7") - client = app.test_client() - response = client.options( - "/openapi/v1/_health", - headers={ - "Origin": "https://app.example.com", - "Access-Control-Request-Method": "GET", - "Access-Control-Request-Headers": HEADER_NAME_CATALOG, - }, - ) - - allow_headers = response.headers.get("Access-Control-Allow-Headers", "").lower() - assert HEADER_NAME_CATALOG.lower() in allow_headers diff --git a/api/tests/unit_tests/controllers/openapi/test_device_approve_deny.py b/api/tests/unit_tests/controllers/openapi/test_device_approve_deny.py deleted file mode 100644 index 6897a2b2e0d..00000000000 --- a/api/tests/unit_tests/controllers/openapi/test_device_approve_deny.py +++ /dev/null @@ -1,60 +0,0 @@ -"""Account-branch device-flow approve/deny under /openapi/v1.""" - -import builtins - -import pytest -from flask import Flask -from flask.views import MethodView - -from controllers.openapi import bp as openapi_bp -from controllers.openapi.oauth_device import DeviceApproveApi, DeviceDenyApi, _error_response -from services.oauth_device_contracts import ApprovalOutcomeUnknownError - -if not hasattr(builtins, "MethodView"): - builtins.MethodView = MethodView # type: ignore[attr-defined] - - -@pytest.fixture -def openapi_app() -> Flask: - app = Flask(__name__) - app.config["TESTING"] = True - app.register_blueprint(openapi_bp) - return app - - -def _rule(app: Flask, path: str): - return next(r for r in app.url_map.iter_rules() if r.rule == path) - - -def test_approve_route_registered(openapi_app: Flask): - rules = {r.rule for r in openapi_app.url_map.iter_rules()} - assert "/openapi/v1/oauth/device/approve" in rules - - -def test_deny_route_registered(openapi_app: Flask): - rules = {r.rule for r in openapi_app.url_map.iter_rules()} - assert "/openapi/v1/oauth/device/deny" in rules - - -def test_approve_dispatches_to_class(openapi_app: Flask): - rule = _rule(openapi_app, "/openapi/v1/oauth/device/approve") - assert openapi_app.view_functions[rule.endpoint].view_class is DeviceApproveApi - - -def test_deny_dispatches_to_class(openapi_app: Flask): - rule = _rule(openapi_app, "/openapi/v1/oauth/device/deny") - assert openapi_app.view_functions[rule.endpoint].view_class is DeviceDenyApi - - -def test_approve_and_deny_methods(openapi_app: Flask): - approve = _rule(openapi_app, "/openapi/v1/oauth/device/approve") - deny = _rule(openapi_app, "/openapi/v1/oauth/device/deny") - assert "POST" in approve.methods - assert "POST" in deny.methods - - -def test_unknown_approval_outcome_is_retryable() -> None: - body, status = _error_response(ApprovalOutcomeUnknownError()) - - assert status == 503 - assert body == {"error": "approval_outcome_unknown"} diff --git a/api/tests/unit_tests/controllers/openapi/test_device_code.py b/api/tests/unit_tests/controllers/openapi/test_device_code.py deleted file mode 100644 index da33292ed2c..00000000000 --- a/api/tests/unit_tests/controllers/openapi/test_device_code.py +++ /dev/null @@ -1,50 +0,0 @@ -"""POST /openapi/v1/oauth/device/code is the canonical RFC 8628 device -authorization endpoint. - -Tests verify URL routing without invoking the handler — invoking would -require Redis, which the unit-test runtime does not initialise. -""" - -import builtins - -import pytest -from flask import Flask -from flask.views import MethodView - -from controllers.openapi import bp as openapi_bp -from controllers.openapi.oauth_device import OAuthDeviceCodeApi - -if not hasattr(builtins, "MethodView"): - builtins.MethodView = MethodView # type: ignore[attr-defined] - - -@pytest.fixture -def openapi_app() -> Flask: - app = Flask(__name__) - app.config["TESTING"] = True - app.register_blueprint(openapi_bp) - return app - - -def test_openapi_route_registered(openapi_app: Flask) -> None: - rules = {r.rule for r in openapi_app.url_map.iter_rules()} - assert "/openapi/v1/oauth/device/code" in rules - - -def test_route_dispatches_to_class(openapi_app: Flask) -> None: - rule = next(r for r in openapi_app.url_map.iter_rules() if r.rule == "/openapi/v1/oauth/device/code") - view = openapi_app.view_functions[rule.endpoint] - assert hasattr(view, "view_class") - assert view.view_class is OAuthDeviceCodeApi - - -def test_route_accepts_post(openapi_app: Flask) -> None: - rule = next(r for r in openapi_app.url_map.iter_rules() if r.rule == "/openapi/v1/oauth/device/code") - assert rule.methods is not None - assert "POST" in rule.methods - - -def test_known_client_ids_default_includes_difyctl() -> None: - from configs import dify_config - - assert "difyctl" in dify_config.OPENAPI_KNOWN_CLIENT_IDS diff --git a/api/tests/unit_tests/controllers/openapi/test_device_lookup.py b/api/tests/unit_tests/controllers/openapi/test_device_lookup.py deleted file mode 100644 index 5af852fbd54..00000000000 --- a/api/tests/unit_tests/controllers/openapi/test_device_lookup.py +++ /dev/null @@ -1,39 +0,0 @@ -"""GET /openapi/v1/oauth/device/lookup is the canonical user-code lookup.""" - -import builtins - -import pytest -from flask import Flask -from flask.views import MethodView - -from controllers.openapi import bp as openapi_bp -from controllers.openapi.oauth_device import OAuthDeviceLookupApi - -if not hasattr(builtins, "MethodView"): - builtins.MethodView = MethodView # type: ignore[attr-defined] - - -@pytest.fixture -def openapi_app() -> Flask: - app = Flask(__name__) - app.config["TESTING"] = True - app.register_blueprint(openapi_bp) - return app - - -def test_openapi_route_registered(openapi_app: Flask) -> None: - rules = {r.rule for r in openapi_app.url_map.iter_rules()} - assert "/openapi/v1/oauth/device/lookup" in rules - - -def test_route_dispatches_to_class(openapi_app: Flask) -> None: - rule = next(r for r in openapi_app.url_map.iter_rules() if r.rule == "/openapi/v1/oauth/device/lookup") - view = openapi_app.view_functions[rule.endpoint] - assert hasattr(view, "view_class") - assert view.view_class is OAuthDeviceLookupApi - - -def test_route_accepts_get(openapi_app: Flask) -> None: - rule = next(r for r in openapi_app.url_map.iter_rules() if r.rule == "/openapi/v1/oauth/device/lookup") - assert rule.methods is not None - assert "GET" in rule.methods diff --git a/api/tests/unit_tests/controllers/openapi/test_device_token.py b/api/tests/unit_tests/controllers/openapi/test_device_token.py deleted file mode 100644 index abb28604523..00000000000 --- a/api/tests/unit_tests/controllers/openapi/test_device_token.py +++ /dev/null @@ -1,33 +0,0 @@ -"""POST /openapi/v1/oauth/device/token is the canonical poll endpoint.""" - -import builtins - -import pytest -from flask import Flask -from flask.views import MethodView - -from controllers.openapi import bp as openapi_bp -from controllers.openapi.oauth_device import OAuthDeviceTokenApi - -if not hasattr(builtins, "MethodView"): - builtins.MethodView = MethodView # type: ignore[attr-defined] - - -@pytest.fixture -def openapi_app() -> Flask: - app = Flask(__name__) - app.config["TESTING"] = True - app.register_blueprint(openapi_bp) - return app - - -def test_openapi_route_registered(openapi_app: Flask) -> None: - rules = {r.rule for r in openapi_app.url_map.iter_rules()} - assert "/openapi/v1/oauth/device/token" in rules - - -def test_route_dispatches_to_class(openapi_app: Flask) -> None: - rule = next(r for r in openapi_app.url_map.iter_rules() if r.rule == "/openapi/v1/oauth/device/token") - view = openapi_app.view_functions[rule.endpoint] - assert hasattr(view, "view_class") - assert view.view_class is OAuthDeviceTokenApi diff --git a/api/tests/unit_tests/controllers/openapi/test_health.py b/api/tests/unit_tests/controllers/openapi/test_health.py deleted file mode 100644 index 83064692bb5..00000000000 --- a/api/tests/unit_tests/controllers/openapi/test_health.py +++ /dev/null @@ -1,33 +0,0 @@ -import builtins - -import pytest -from flask import Flask -from flask.views import MethodView - -from controllers.openapi import bp as openapi_bp - -if not hasattr(builtins, "MethodView"): - builtins.MethodView = MethodView # type: ignore[attr-defined] - - -@pytest.fixture -def app() -> Flask: - app = Flask(__name__) - app.config["TESTING"] = True - app.register_blueprint(openapi_bp) - return app - - -def test_health_returns_ok(app: Flask) -> None: - client = app.test_client() - response = client.get("/openapi/v1/_health") - - assert response.status_code == 200 - assert response.get_json() == {"ok": True} - - -def test_health_path_is_under_openapi_v1_prefix(app: Flask) -> None: - client = app.test_client() - assert client.get("/_health").status_code == 404 - assert client.get("/v1/_health").status_code == 404 - assert client.get("/openapi/v1/_health").status_code == 200 diff --git a/api/tests/unit_tests/controllers/openapi/test_index.py b/api/tests/unit_tests/controllers/openapi/test_index.py new file mode 100644 index 00000000000..48eed18fdfa --- /dev/null +++ b/api/tests/unit_tests/controllers/openapi/test_index.py @@ -0,0 +1,91 @@ +from __future__ import annotations + +import builtins + +import pytest +from flask import Flask +from flask.views import MethodView + +from controllers.openapi import bp as openapi_bp +from enums import DeploymentEdition +from tests.helpers.unit.config_override import apply_config_overrides + +if not hasattr(builtins, "MethodView"): + builtins.MethodView = MethodView # type: ignore[attr-defined] + + +class TestHealthCases: + @staticmethod + @pytest.fixture + def app() -> Flask: + app = Flask(__name__) + app.config["TESTING"] = True + app.register_blueprint(openapi_bp) + return app + + @staticmethod + def test_health_returns_ok(app: Flask) -> None: + client = app.test_client() + response = client.get("/openapi/v1/_health") + + assert response.status_code == 200 + assert response.get_json() == {"ok": True} + + @staticmethod + def test_health_path_is_under_openapi_v1_prefix(app: Flask) -> None: + client = app.test_client() + assert client.get("/_health").status_code == 404 + assert client.get("/v1/_health").status_code == 404 + assert client.get("/openapi/v1/_health").status_code == 200 + + +"""Meta endpoint /openapi/v1/_version — no auth, returns version + edition.""" + + +class TestMetaVersionCases: + @staticmethod + def test_version_endpoint_returns_200_without_auth(openapi_app: Flask) -> None: + client = openapi_app.test_client() + response = client.get("/openapi/v1/_version") + + assert response.status_code == 200 + payload = response.get_json() + assert isinstance(payload, dict) + assert "version" in payload + assert "edition" in payload + assert isinstance(payload["version"], str) + assert payload["edition"] in {edition.value for edition in DeploymentEdition} + + @staticmethod + def test_version_endpoint_ignores_bearer_header(openapi_app: Flask) -> None: + """Endpoint is auth-free — a bogus bearer should not break it.""" + client = openapi_app.test_client() + response = client.get( + "/openapi/v1/_version", + headers={"Authorization": "Bearer total-nonsense"}, + ) + + assert response.status_code == 200 + payload = response.get_json() + assert "version" in payload + assert "edition" in payload + + @staticmethod + def test_version_endpoint_reflects_edition_config(openapi_app: Flask, monkeypatch: pytest.MonkeyPatch) -> None: + apply_config_overrides(monkeypatch, DEPLOYMENT_EDITION=DeploymentEdition.CLOUD) + + client = openapi_app.test_client() + response = client.get("/openapi/v1/_version") + + assert response.status_code == 200 + assert response.get_json()["edition"] == "CLOUD" + + @staticmethod + def test_version_endpoint_reflects_enterprise_edition(openapi_app: Flask, monkeypatch: pytest.MonkeyPatch) -> None: + apply_config_overrides(monkeypatch, DEPLOYMENT_EDITION=DeploymentEdition.ENTERPRISE) + + client = openapi_app.test_client() + response = client.get("/openapi/v1/_version") + + assert response.status_code == 200 + assert response.get_json()["edition"] == "ENTERPRISE" diff --git a/api/tests/unit_tests/controllers/openapi/test_integration_context.py b/api/tests/unit_tests/controllers/openapi/test_integration_context.py deleted file mode 100644 index d9087c902ed..00000000000 --- a/api/tests/unit_tests/controllers/openapi/test_integration_context.py +++ /dev/null @@ -1,82 +0,0 @@ -"""Keep the integration handler context aligned with request-scoped identity loading.""" - -import pytest -from sqlalchemy import select -from sqlalchemy.orm import Session, object_session - -from controllers.openapi.auth.context import RouteContractError -from controllers.openapi.auth.requirements import CheckWorkspaceMember -from models import TenantAccountJoin -from models.account import TenantAccountRole -from tests.test_containers_integration_tests.controllers.openapi.conftest import context_for -from tests.unit_tests.controllers.openapi.auth._world import ( - ACCOUNT_ID, - APP_ID, - OTHER_TENANT_ID, - TENANT_ID, - make_account, - make_app, - make_membership, - make_tenant, - persist, -) - - -@pytest.mark.parametrize("view_args", [{"workspace_id": TENANT_ID}, {"app_id": APP_ID}]) -def test_workspace_requirement_binds_fresh_caller_to_requested_workspace( - sqlite_session: Session, view_args: dict[str, str] -) -> None: - account = make_account() - original_workspace = make_tenant(tenant_id=OTHER_TENANT_ID) - requested_membership = make_membership(TenantAccountRole.ADMIN) - requested_membership.current = False - persist( - sqlite_session, - account, - original_workspace, - make_tenant(), - make_app(), - requested_membership, - TenantAccountJoin( - tenant_id=OTHER_TENANT_ID, - account_id=ACCOUNT_ID, - current=True, - role=TenantAccountRole.OWNER, - ), - ) - account.set_current_tenant_with_session(original_workspace, session=sqlite_session) - - ctx = context_for( - account, - session=sqlite_session, - view_args=view_args, - requirements=(CheckWorkspaceMember(),), - ) - - assert ctx.account is not account - assert object_session(ctx.account) is None - assert ctx.account.current_tenant_id == ctx.workspace.id == TENANT_ID - assert ctx.account.role == ctx.workspace_role == TenantAccountRole.ADMIN - assert account.current_tenant_id == OTHER_TENANT_ID - assert ( - sqlite_session.scalar( - select(TenantAccountJoin.tenant_id).where( - TenantAccountJoin.account_id == ACCOUNT_ID, TenantAccountJoin.current.is_(True) - ) - ) - == OTHER_TENANT_ID - ) - if "app_id" in view_args: - assert ctx.app.id == APP_ID - - -def test_app_context_without_workspace_requirement_does_not_bind_workspace(sqlite_session: Session) -> None: - account = make_account() - persist(sqlite_session, account, make_tenant(), make_app(), make_membership()) - - ctx = context_for(account, session=sqlite_session, view_args={"app_id": APP_ID}) - - assert ctx.app.id == APP_ID - assert ctx.account.current_tenant_id is None - with pytest.raises(RouteContractError, match="workspace was not loaded"): - _ = ctx.workspace diff --git a/api/tests/unit_tests/controllers/openapi/test_knowledge_fs.py b/api/tests/unit_tests/controllers/openapi/test_knowledge_fs.py index 8f1212002d7..8f3a5701b14 100644 --- a/api/tests/unit_tests/controllers/openapi/test_knowledge_fs.py +++ b/api/tests/unit_tests/controllers/openapi/test_knowledge_fs.py @@ -19,7 +19,7 @@ from enums import DeploymentEdition from libs.oauth_bearer import AuthContext, TokenType from services.knowledge_fs.feature_service import KnowledgeFSFeatures from services.knowledge_fs.product_dto import KnowledgeFSListQuery -from tests.unit_tests.controllers.knowledge_fs_controller_fakes import ( +from tests.helpers.unit.controllers.knowledge_fs_controller_fakes import ( FakeComponent, FakeRuntime, Recorder, diff --git a/api/tests/unit_tests/controllers/openapi/test_meta_version.py b/api/tests/unit_tests/controllers/openapi/test_meta_version.py deleted file mode 100644 index 313f09a0ba6..00000000000 --- a/api/tests/unit_tests/controllers/openapi/test_meta_version.py +++ /dev/null @@ -1,56 +0,0 @@ -"""Meta endpoint /openapi/v1/_version — no auth, returns version + edition.""" - -from __future__ import annotations - -import pytest -from flask import Flask - -from enums import DeploymentEdition -from tests.unit_tests.config_override import apply_config_overrides - - -def test_version_endpoint_returns_200_without_auth(openapi_app: Flask) -> None: - client = openapi_app.test_client() - response = client.get("/openapi/v1/_version") - - assert response.status_code == 200 - payload = response.get_json() - assert isinstance(payload, dict) - assert "version" in payload - assert "edition" in payload - assert isinstance(payload["version"], str) - assert payload["edition"] in {edition.value for edition in DeploymentEdition} - - -def test_version_endpoint_ignores_bearer_header(openapi_app: Flask) -> None: - """Endpoint is auth-free — a bogus bearer should not break it.""" - client = openapi_app.test_client() - response = client.get( - "/openapi/v1/_version", - headers={"Authorization": "Bearer total-nonsense"}, - ) - - assert response.status_code == 200 - payload = response.get_json() - assert "version" in payload - assert "edition" in payload - - -def test_version_endpoint_reflects_edition_config(openapi_app: Flask, monkeypatch: pytest.MonkeyPatch) -> None: - apply_config_overrides(monkeypatch, DEPLOYMENT_EDITION=DeploymentEdition.CLOUD) - - client = openapi_app.test_client() - response = client.get("/openapi/v1/_version") - - assert response.status_code == 200 - assert response.get_json()["edition"] == "CLOUD" - - -def test_version_endpoint_reflects_enterprise_edition(openapi_app: Flask, monkeypatch: pytest.MonkeyPatch) -> None: - apply_config_overrides(monkeypatch, DEPLOYMENT_EDITION=DeploymentEdition.ENTERPRISE) - - client = openapi_app.test_client() - response = client.get("/openapi/v1/_version") - - assert response.status_code == 200 - assert response.get_json()["edition"] == "ENTERPRISE" diff --git a/api/tests/unit_tests/controllers/openapi/test_model_providers.py b/api/tests/unit_tests/controllers/openapi/test_model_providers.py index 8202aaf9ddb..2ea982f1798 100644 --- a/api/tests/unit_tests/controllers/openapi/test_model_providers.py +++ b/api/tests/unit_tests/controllers/openapi/test_model_providers.py @@ -46,7 +46,7 @@ from services.entities.model_provider_entities import ( SystemConfigurationResponse, ) from services.model_provider.service import ModelProviderService -from tests.unit_tests.model_factories import make_account, make_tenant +from tests.helpers.unit.model_factories import make_account, make_tenant def _model( diff --git a/api/tests/unit_tests/controllers/openapi/test_oauth_device.py b/api/tests/unit_tests/controllers/openapi/test_oauth_device.py new file mode 100644 index 00000000000..43f226e7ffb --- /dev/null +++ b/api/tests/unit_tests/controllers/openapi/test_oauth_device.py @@ -0,0 +1,589 @@ +import builtins +from collections.abc import Callable +from datetime import UTC, datetime +from http.cookies import SimpleCookie +from types import SimpleNamespace +from unittest.mock import Mock, create_autospec + +import pytest +from flask import Flask +from flask.views import MethodView + +from controllers.console import flask_admission, wraps +from controllers.openapi import bp as openapi_bp +from controllers.openapi.oauth_device import ( + DeviceApproveApi, + DeviceDenyApi, + OAuthDeviceCodeApi, + OAuthDeviceLookupApi, + OAuthDeviceTokenApi, + _error_response, +) +from libs.helper import RateLimiter +from libs.login import AccountWithTenant +from machinery.context import RequestContext +from models.account import Account, AccountStatus +from services.oauth_device_application_service import OAuthDeviceApplicationService +from services.oauth_device_contracts import ( + AccessDeniedError, + AlreadyResolvedError, + ApprovalInProgressError, + ApprovalOutcomeUnknownError, + ApprovalSessionConsumedError, + AuthorizationPendingError, + DeviceApprovalContext, + DeviceAuthorization, + DeviceLookup, + DeviceMutation, + DeviceRequestContext, + DeviceSSOCompletion, + DeviceSSOInitiation, + DeviceStateLostError, + ExpiredOrUnknownError, + ExpiredTokenError, + ExternalApprovalCSRFError, + ExternalApprovalRateLimitError, + ExternalIdentityConflictError, + ExternalUserCodeMismatchError, + ExternalUserCodeNotFoundError, + InvalidApprovalSessionError, + InvalidUserCodeError, + OAuthDeviceError, + OAuthDeviceSSOConfigurationError, + OAuthDeviceSSOInitiationError, + PollTooFastError, + UnsupportedClientError, +) + +"""Account-branch device-flow approve/deny under /openapi/v1.""" + + +if not hasattr(builtins, "MethodView"): + builtins.MethodView = MethodView # type: ignore[attr-defined] + + +def _rule(app: Flask, path: str): + return next(r for r in app.url_map.iter_rules() if r.rule == path) + + +class TestDeviceApproveDenyCases: + @staticmethod + @pytest.fixture + def openapi_app() -> Flask: + app = Flask(__name__) + app.config["TESTING"] = True + app.register_blueprint(openapi_bp) + return app + + @staticmethod + def test_approve_route_registered(openapi_app: Flask): + rules = {r.rule for r in openapi_app.url_map.iter_rules()} + assert "/openapi/v1/oauth/device/approve" in rules + + @staticmethod + def test_deny_route_registered(openapi_app: Flask): + rules = {r.rule for r in openapi_app.url_map.iter_rules()} + assert "/openapi/v1/oauth/device/deny" in rules + + @staticmethod + def test_approve_dispatches_to_class(openapi_app: Flask): + rule = _rule(openapi_app, "/openapi/v1/oauth/device/approve") + assert openapi_app.view_functions[rule.endpoint].view_class is DeviceApproveApi + + @staticmethod + def test_deny_dispatches_to_class(openapi_app: Flask): + rule = _rule(openapi_app, "/openapi/v1/oauth/device/deny") + assert openapi_app.view_functions[rule.endpoint].view_class is DeviceDenyApi + + @staticmethod + def test_approve_and_deny_methods(openapi_app: Flask): + approve = _rule(openapi_app, "/openapi/v1/oauth/device/approve") + deny = _rule(openapi_app, "/openapi/v1/oauth/device/deny") + assert "POST" in approve.methods + assert "POST" in deny.methods + + @staticmethod + def test_unknown_approval_outcome_is_retryable() -> None: + body, status = _error_response(ApprovalOutcomeUnknownError()) + + assert status == 503 + assert body == {"error": "approval_outcome_unknown"} + + +"""POST /openapi/v1/oauth/device/code is the canonical RFC 8628 device +authorization endpoint. + +Tests verify URL routing without invoking the handler — invoking would +require Redis, which the unit-test runtime does not initialise. +""" + + +if not hasattr(builtins, "MethodView"): + builtins.MethodView = MethodView # type: ignore[attr-defined] + + +class TestDeviceCodeCases: + @staticmethod + @pytest.fixture + def openapi_app() -> Flask: + app = Flask(__name__) + app.config["TESTING"] = True + app.register_blueprint(openapi_bp) + return app + + @staticmethod + def test_openapi_route_registered(openapi_app: Flask) -> None: + rules = {r.rule for r in openapi_app.url_map.iter_rules()} + assert "/openapi/v1/oauth/device/code" in rules + + @staticmethod + def test_route_dispatches_to_class(openapi_app: Flask) -> None: + rule = next(r for r in openapi_app.url_map.iter_rules() if r.rule == "/openapi/v1/oauth/device/code") + view = openapi_app.view_functions[rule.endpoint] + assert hasattr(view, "view_class") + assert view.view_class is OAuthDeviceCodeApi + + @staticmethod + def test_route_accepts_post(openapi_app: Flask) -> None: + rule = next(r for r in openapi_app.url_map.iter_rules() if r.rule == "/openapi/v1/oauth/device/code") + assert rule.methods is not None + assert "POST" in rule.methods + + @staticmethod + def test_known_client_ids_default_includes_difyctl() -> None: + from configs import dify_config + + assert "difyctl" in dify_config.OPENAPI_KNOWN_CLIENT_IDS + + +"""GET /openapi/v1/oauth/device/lookup is the canonical user-code lookup.""" + + +if not hasattr(builtins, "MethodView"): + builtins.MethodView = MethodView # type: ignore[attr-defined] + + +class TestDeviceLookupCases: + @staticmethod + @pytest.fixture + def openapi_app() -> Flask: + app = Flask(__name__) + app.config["TESTING"] = True + app.register_blueprint(openapi_bp) + return app + + @staticmethod + def test_openapi_route_registered(openapi_app: Flask) -> None: + rules = {r.rule for r in openapi_app.url_map.iter_rules()} + assert "/openapi/v1/oauth/device/lookup" in rules + + @staticmethod + def test_route_dispatches_to_class(openapi_app: Flask) -> None: + rule = next(r for r in openapi_app.url_map.iter_rules() if r.rule == "/openapi/v1/oauth/device/lookup") + view = openapi_app.view_functions[rule.endpoint] + assert hasattr(view, "view_class") + assert view.view_class is OAuthDeviceLookupApi + + @staticmethod + def test_route_accepts_get(openapi_app: Flask) -> None: + rule = next(r for r in openapi_app.url_map.iter_rules() if r.rule == "/openapi/v1/oauth/device/lookup") + assert rule.methods is not None + assert "GET" in rule.methods + + +"""POST /openapi/v1/oauth/device/token is the canonical poll endpoint.""" + + +if not hasattr(builtins, "MethodView"): + builtins.MethodView = MethodView # type: ignore[attr-defined] + + +class TestDeviceTokenCases: + @staticmethod + @pytest.fixture + def openapi_app() -> Flask: + app = Flask(__name__) + app.config["TESTING"] = True + app.register_blueprint(openapi_bp) + return app + + @staticmethod + def test_openapi_route_registered(openapi_app: Flask) -> None: + rules = {r.rule for r in openapi_app.url_map.iter_rules()} + assert "/openapi/v1/oauth/device/token" in rules + + @staticmethod + def test_route_dispatches_to_class(openapi_app: Flask) -> None: + rule = next(r for r in openapi_app.url_map.iter_rules() if r.rule == "/openapi/v1/oauth/device/token") + view = openapi_app.view_functions[rule.endpoint] + assert hasattr(view, "view_class") + assert view.view_class is OAuthDeviceTokenApi + + +"""Device-flow HTTP contracts with application services isolated at their public boundary.""" + + +_PREFIX = "/openapi/v1/oauth/device" +_COOKIE = "device_approval_grant" + + +class TestOauthDeviceHttpCases: + @staticmethod + @pytest.fixture + def device_service( + openapi_app: Flask, monkeypatch: pytest.MonkeyPatch, config_overrides: Callable[..., None] + ) -> Mock: + service = create_autospec(OAuthDeviceApplicationService, instance=True) + openapi_app.extensions["application_services"] = SimpleNamespace( + oauth_device=service, + feature_queries=SimpleNamespace(has_valid_enterprise_license=lambda: True), + ) + # Supply an authenticated account while retaining the real admission/context wrapper. + account = Account(name="Ada", email="ada@example.com", status=AccountStatus.ACTIVE) + account.id = "account-1" + account_with_tenant = AccountWithTenant(account, "workspace-1") + config_overrides( + LOGIN_DISABLED=True, + ENABLE_OAUTH_BEARER=True, + CONSOLE_API_URL="https://api.example", + CONSOLE_WEB_URL="https://console.example", + ) + monkeypatch.setattr(wraps, "_is_setup_completed", lambda: True) + monkeypatch.setattr(wraps, "current_account_with_tenant", lambda: account_with_tenant) + monkeypatch.setattr(flask_admission, "current_account_with_tenant", lambda: account_with_tenant) + monkeypatch.setattr(RateLimiter, "is_rate_limited", lambda *_args: False) + monkeypatch.setattr(RateLimiter, "increment_rate_limit", lambda *_args: None) + return service + + @staticmethod + def test_device_code_serializes_authorization_and_passes_request_origin( + openapi_app: Flask, device_service: Mock + ) -> None: + device_service.start.return_value = DeviceAuthorization( + device_code="opaque-code", + user_code="ABCD-EFGH", + verification_uri="https://console.example/device", + expires_in=900, + interval=5, + ) + + response = openapi_app.test_client().post( + f"{_PREFIX}/code", + json={"client_id": "difyctl", "device_label": "laptop"}, + base_url="https://api.example", + environ_overrides={"REMOTE_ADDR": "192.0.2.10"}, + ) + + assert response.status_code == 200 + assert response.json == { + "device_code": "opaque-code", + "user_code": "ABCD-EFGH", + "verification_uri": "https://console.example/device", + "expires_in": 900, + "interval": 5, + } + device_service.start.assert_called_once_with( + client_id="difyctl", + device_label="laptop", + created_ip="192.0.2.10", + request_origin="https://api.example/", + ) + + @staticmethod + @pytest.mark.parametrize("path", ["code", "token", "approve", "deny", "approve-external"]) + @pytest.mark.parametrize("body", ["{}", "[]", "{invalid"]) + def test_invalid_device_request_never_reaches_service( + openapi_app: Flask, device_service: Mock, path: str, body: str + ) -> None: + response = openapi_app.test_client().post(f"{_PREFIX}/{path}", data=body, content_type="application/json") + + assert response.status_code == 400 + assert device_service.mock_calls == [] + + @staticmethod + def test_lookup_requires_user_code(openapi_app: Flask, device_service: Mock) -> None: + response = openapi_app.test_client().get(f"{_PREFIX}/lookup") + + assert response.status_code == 400 + device_service.lookup.assert_not_called() + + @staticmethod + def test_lookup_returns_public_validity_without_token_data(openapi_app: Flask, device_service: Mock) -> None: + device_service.lookup.return_value = DeviceLookup(valid=False, expires_in_remaining=0, client_id=None) + + response = openapi_app.test_client().get(f"{_PREFIX}/lookup?user_code=ABCD-EFGH") + + assert response.status_code == 200 + assert response.json == {"valid": False, "expires_in_remaining": 0, "client_id": None} + device_service.lookup.assert_called_once_with(user_code="ABCD-EFGH") + + @staticmethod + @pytest.mark.parametrize( + ("path", "method", "error", "status", "code"), + [ + ("code", "start", UnsupportedClientError(), 400, "unsupported_client"), + ("token", "poll", PollTooFastError(), 400, "slow_down"), + ("token", "poll", ExpiredTokenError(), 400, "expired_token"), + ("token", "poll", AuthorizationPendingError(), 400, "authorization_pending"), + ("token", "poll", AccessDeniedError(), 400, "access_denied"), + ("approve", "approve", ExpiredOrUnknownError(), 404, "expired_or_unknown"), + ("approve", "approve", AlreadyResolvedError(), 409, "already_resolved"), + ("approve", "approve", ApprovalInProgressError(), 409, "approve_in_progress"), + ("approve", "approve", ApprovalOutcomeUnknownError(), 503, "approval_outcome_unknown"), + ("deny", "deny", DeviceStateLostError(), 409, "state_lost"), + ], + ) + def test_device_errors_keep_oauth_response_contract( + openapi_app: Flask, + device_service: Mock, + path: str, + method: str, + error: OAuthDeviceError, + status: int, + code: str, + ) -> None: + methods = { + "start": device_service.start, + "poll": device_service.poll, + "approve": device_service.approve, + "deny": device_service.deny, + } + methods[method].side_effect = error + + response = openapi_app.test_client().post( + f"{_PREFIX}/{path}", + json={ + "client_id": "difyctl", + "device_label": "CLI", + "device_code": "opaque-code", + "user_code": "ABCD-EFGH", + }, + ) + + assert response.status_code == status + assert response.json == {"error": code} + + @staticmethod + @pytest.mark.parametrize("external", [False, True], ids=["account", "external"]) + def test_poll_serializes_subject_specific_token_payload( + openapi_app: Flask, device_service: Mock, external: bool + ) -> None: + payload: dict[str, object] = { + "token": "secret-token", + "expires_at": "2030-01-01T00:00:00Z", + "token_id": "token-1", + "subject_type": "external_sso" if external else "account", + "account": None if external else {"id": "account-1", "email": "ada@example.com", "name": "Ada"}, + "workspaces": [], + "default_workspace_id": None, + } + if external: + payload.update(subject_email="external@example.com", subject_issuer="https://idp.example") + device_service.poll.return_value = payload + + response = openapi_app.test_client().post( + f"{_PREFIX}/token", + json={"device_code": "opaque-code", "client_id": "difyctl"}, + environ_overrides={"REMOTE_ADDR": "192.0.2.10"}, + ) + + assert response.status_code == 200 + assert response.json == payload + device_service.poll.assert_called_once_with(device_code="opaque-code", poll_ip="192.0.2.10") + + @staticmethod + @pytest.mark.parametrize("path", ["approve", "deny"]) + def test_account_decision_uses_admitted_identity(openapi_app: Flask, device_service: Mock, path: str) -> None: + method = device_service.approve if path == "approve" else device_service.deny + method.return_value = DeviceMutation(status="approved" if path == "approve" else "denied") + + response = openapi_app.test_client().post( + f"{_PREFIX}/{path}", + json={"user_code": "ABCD-EFGH", "account_id": "forged", "workspace_id": "foreign"}, + ) + + assert response.status_code == 200 + assert response.json == {"status": method.return_value.status} + assert method.call_args.kwargs == {"user_code": "ABCD-EFGH"} + if path == "approve": + context = method.call_args.args[0] + assert isinstance(context, RequestContext) + assert context.account_id == "account-1" + assert context.active_workspace_id == "workspace-1" + + @staticmethod + @pytest.mark.parametrize("path", ["approve", "deny"]) + def test_disabled_bearer_feature_blocks_account_decisions( + openapi_app: Flask, device_service: Mock, config_overrides: Callable[..., None], path: str + ) -> None: + config_overrides(ENABLE_OAUTH_BEARER=False) + + response = openapi_app.test_client().post(f"{_PREFIX}/{path}", json={"user_code": "ABCD-EFGH"}) + + assert response.status_code == 503 + assert device_service.mock_calls == [] + + @staticmethod + def test_sso_initiation_clears_old_approval_cookie(openapi_app: Flask, device_service: Mock) -> None: + device_service.initiate_sso.return_value = DeviceSSOInitiation("https://idp.example/authorize") + + response = openapi_app.test_client().get(f"{_PREFIX}/sso-initiate?user_code=%20ABCD-EFGH%20") + + assert response.status_code == 302 + assert response.headers["Location"] == "https://idp.example/authorize" + cookie = SimpleCookie(response.headers["Set-Cookie"])[_COOKIE] + assert cookie["max-age"] == "0" + assert cookie["path"] == _PREFIX + assert device_service.initiate_sso.call_args.kwargs == {"user_code": "ABCD-EFGH"} + + @staticmethod + def test_sso_initiation_requires_user_code(openapi_app: Flask, device_service: Mock) -> None: + response = openapi_app.test_client().get(f"{_PREFIX}/sso-initiate?user_code=%20") + + assert response.status_code == 400 + device_service.initiate_sso.assert_not_called() + + @staticmethod + def test_sso_completion_sets_scoped_http_only_cookie(openapi_app: Flask, device_service: Mock) -> None: + device_service.complete_sso.return_value = DeviceSSOCompletion( + user_code="ABCD-EFGH", approval_grant="signed-grant" + ) + + response = openapi_app.test_client().get(f"{_PREFIX}/sso-complete?sso_assertion=signed-assertion") + + assert response.status_code == 302 + assert response.headers["Location"] == "/device?sso_verified=1" + cookie = SimpleCookie(response.headers["Set-Cookie"])[_COOKIE] + assert cookie.value == "signed-grant" + assert cookie["path"] == _PREFIX + assert cookie["httponly"] + assert cookie["secure"] + assert cookie["samesite"] == "Lax" + assert cookie["max-age"] == "300" + assert device_service.complete_sso.call_args.kwargs == { + "inbound_error": None, + "inbound_user_code": None, + "assertion": "signed-assertion", + } + + @staticmethod + @pytest.mark.parametrize("failure", ["exception", "missing-grant", "untrusted-error"]) + def test_sso_completion_failure_redirects_without_exposing_details( + openapi_app: Flask, device_service: Mock, failure: str + ) -> None: + device_service.complete_sso.return_value = DeviceSSOCompletion( + error_code="private-provider-error" if failure == "untrusted-error" else None, + ) + if failure == "exception": + device_service.complete_sso.side_effect = RuntimeError("private-provider-error") + + response = openapi_app.test_client().get(f"{_PREFIX}/sso-complete?sso_assertion=private-assertion") + + assert response.status_code == 302 + assert response.headers["Location"] == "/device?sso_error=sso_failed" + assert "Set-Cookie" not in response.headers + assert b"private-" not in response.data + + @staticmethod + def test_approval_context_reads_cookie_and_serializes_expiry(openapi_app: Flask, device_service: Mock) -> None: + device_service.get_approval_context.return_value = DeviceApprovalContext( + subject_email="external@example.com", + subject_issuer="https://idp.example", + user_code="ABCD-EFGH", + csrf_token="csrf-token", + expires_at=datetime(2030, 1, 1, tzinfo=UTC), + ) + client = openapi_app.test_client() + client.set_cookie(_COOKIE, "signed-grant", path=_PREFIX) + + response = client.get(f"{_PREFIX}/approval-context") + + assert response.status_code == 200 + assert response.json == { + "subject_email": "external@example.com", + "subject_issuer": "https://idp.example", + "user_code": "ABCD-EFGH", + "csrf_token": "csrf-token", + "expires_at": "2030-01-01T00:00:00Z", + } + assert device_service.get_approval_context.call_args.kwargs == {"approval_grant": "signed-grant"} + + @staticmethod + def test_external_approval_uses_cookie_and_csrf_header_then_clears_cookie( + openapi_app: Flask, device_service: Mock + ) -> None: + device_service.approve_external.return_value = DeviceMutation("approved") + client = openapi_app.test_client() + client.set_cookie(_COOKIE, "signed-grant", path=_PREFIX) + + response = client.post( + f"{_PREFIX}/approve-external", + headers={"X-CSRF-Token": "csrf-header"}, + json={"user_code": "ABCD-EFGH", "csrf_token": "forged", "approval_grant": "forged"}, + ) + + assert response.status_code == 200 + assert response.json == {"status": "approved"} + assert isinstance(device_service.approve_external.call_args.args[0], DeviceRequestContext) + assert device_service.approve_external.call_args.kwargs == { + "approval_grant": "signed-grant", + "csrf_token": "csrf-header", + "user_code": "ABCD-EFGH", + } + cookie = SimpleCookie(response.headers["Set-Cookie"])[_COOKIE] + assert cookie["max-age"] == "0" + assert cookie["path"] == _PREFIX + + @staticmethod + @pytest.mark.parametrize( + ("path", "method", "error", "status", "message"), + [ + ("sso-initiate", "initiate_sso", InvalidUserCodeError(), 400, "invalid_user_code"), + ("sso-initiate", "initiate_sso", OAuthDeviceSSOConfigurationError(), 502, "console_api_url_unset"), + ("sso-initiate", "initiate_sso", OAuthDeviceSSOInitiationError(), 502, "sso_initiate_failed"), + ("approval-context", "get_approval_context", InvalidApprovalSessionError(), 401, "no_session"), + ("approve-external", "approve_external", InvalidApprovalSessionError(), 401, "invalid_session"), + ("approve-external", "approve_external", ExternalApprovalRateLimitError(), 429, "rate_limited"), + ("approve-external", "approve_external", ExternalApprovalCSRFError(), 403, "csrf_mismatch"), + ("approve-external", "approve_external", ExternalUserCodeMismatchError(), 400, "user_code_mismatch"), + ("approve-external", "approve_external", ExternalUserCodeNotFoundError(), 404, "user_code_not_pending"), + ("approve-external", "approve_external", AlreadyResolvedError(), 409, "user_code_not_pending"), + ("approve-external", "approve_external", ApprovalInProgressError(), 409, "approve_in_progress"), + ("approve-external", "approve_external", ApprovalOutcomeUnknownError(), 503, "approval_outcome_unknown"), + ( + "approve-external", + "approve_external", + ExternalIdentityConflictError(), + 403, + "email_belongs_to_dify_account", + ), + ("approve-external", "approve_external", ApprovalSessionConsumedError(), 401, "session_already_consumed"), + ("approve-external", "approve_external", DeviceStateLostError(), 409, "state_lost"), + ], + ) + def test_sso_errors_preserve_http_status_and_approval_cookie( + openapi_app: Flask, + device_service: Mock, + path: str, + method: str, + error: OAuthDeviceError, + status: int, + message: str, + ) -> None: + methods = { + "initiate_sso": device_service.initiate_sso, + "get_approval_context": device_service.get_approval_context, + "approve_external": device_service.approve_external, + } + methods[method].side_effect = error + client = openapi_app.test_client() + client.set_cookie(_COOKIE, "signed-grant", path=_PREFIX) + if path == "approve-external": + response = client.post(f"{_PREFIX}/{path}", json={"user_code": "ABCD-EFGH"}) + else: + response = client.get(f"{_PREFIX}/{path}?user_code=ABCD-EFGH") + + assert response.status_code == status + assert message in response.get_data(as_text=True) + assert "Set-Cookie" not in response.headers + assert "signed-grant" not in response.get_data(as_text=True) diff --git a/api/tests/unit_tests/controllers/openapi/test_oauth_device_http.py b/api/tests/unit_tests/controllers/openapi/test_oauth_device_http.py deleted file mode 100644 index 9ba8eb8faae..00000000000 --- a/api/tests/unit_tests/controllers/openapi/test_oauth_device_http.py +++ /dev/null @@ -1,383 +0,0 @@ -"""Device-flow HTTP contracts with application services isolated at their public boundary.""" - -from collections.abc import Callable -from datetime import UTC, datetime -from http.cookies import SimpleCookie -from types import SimpleNamespace -from unittest.mock import Mock, create_autospec - -import pytest -from flask import Flask - -from controllers.console import flask_admission, wraps -from libs.helper import RateLimiter -from libs.login import AccountWithTenant -from machinery.context import RequestContext -from models.account import Account, AccountStatus -from services.oauth_device_application_service import OAuthDeviceApplicationService -from services.oauth_device_contracts import ( - AccessDeniedError, - AlreadyResolvedError, - ApprovalInProgressError, - ApprovalOutcomeUnknownError, - ApprovalSessionConsumedError, - AuthorizationPendingError, - DeviceApprovalContext, - DeviceAuthorization, - DeviceLookup, - DeviceMutation, - DeviceRequestContext, - DeviceSSOCompletion, - DeviceSSOInitiation, - DeviceStateLostError, - ExpiredOrUnknownError, - ExpiredTokenError, - ExternalApprovalCSRFError, - ExternalApprovalRateLimitError, - ExternalIdentityConflictError, - ExternalUserCodeMismatchError, - ExternalUserCodeNotFoundError, - InvalidApprovalSessionError, - InvalidUserCodeError, - OAuthDeviceError, - OAuthDeviceSSOConfigurationError, - OAuthDeviceSSOInitiationError, - PollTooFastError, - UnsupportedClientError, -) - -_PREFIX = "/openapi/v1/oauth/device" -_COOKIE = "device_approval_grant" - - -@pytest.fixture -def device_service(openapi_app: Flask, monkeypatch: pytest.MonkeyPatch, config_overrides: Callable[..., None]) -> Mock: - service = create_autospec(OAuthDeviceApplicationService, instance=True) - openapi_app.extensions["application_services"] = SimpleNamespace( - oauth_device=service, - feature_queries=SimpleNamespace(has_valid_enterprise_license=lambda: True), - ) - # Supply an authenticated account while retaining the real admission/context wrapper. - account = Account(name="Ada", email="ada@example.com", status=AccountStatus.ACTIVE) - account.id = "account-1" - account_with_tenant = AccountWithTenant(account, "workspace-1") - config_overrides( - LOGIN_DISABLED=True, - ENABLE_OAUTH_BEARER=True, - CONSOLE_API_URL="https://api.example", - CONSOLE_WEB_URL="https://console.example", - ) - monkeypatch.setattr(wraps, "_is_setup_completed", lambda: True) - monkeypatch.setattr(wraps, "current_account_with_tenant", lambda: account_with_tenant) - monkeypatch.setattr(flask_admission, "current_account_with_tenant", lambda: account_with_tenant) - monkeypatch.setattr(RateLimiter, "is_rate_limited", lambda *_args: False) - monkeypatch.setattr(RateLimiter, "increment_rate_limit", lambda *_args: None) - return service - - -def test_device_code_serializes_authorization_and_passes_request_origin( - openapi_app: Flask, device_service: Mock -) -> None: - device_service.start.return_value = DeviceAuthorization( - device_code="opaque-code", - user_code="ABCD-EFGH", - verification_uri="https://console.example/device", - expires_in=900, - interval=5, - ) - - response = openapi_app.test_client().post( - f"{_PREFIX}/code", - json={"client_id": "difyctl", "device_label": "laptop"}, - base_url="https://api.example", - environ_overrides={"REMOTE_ADDR": "192.0.2.10"}, - ) - - assert response.status_code == 200 - assert response.json == { - "device_code": "opaque-code", - "user_code": "ABCD-EFGH", - "verification_uri": "https://console.example/device", - "expires_in": 900, - "interval": 5, - } - device_service.start.assert_called_once_with( - client_id="difyctl", - device_label="laptop", - created_ip="192.0.2.10", - request_origin="https://api.example/", - ) - - -@pytest.mark.parametrize("path", ["code", "token", "approve", "deny", "approve-external"]) -@pytest.mark.parametrize("body", ["{}", "[]", "{invalid"]) -def test_invalid_device_request_never_reaches_service( - openapi_app: Flask, device_service: Mock, path: str, body: str -) -> None: - response = openapi_app.test_client().post(f"{_PREFIX}/{path}", data=body, content_type="application/json") - - assert response.status_code == 400 - assert device_service.mock_calls == [] - - -def test_lookup_requires_user_code(openapi_app: Flask, device_service: Mock) -> None: - response = openapi_app.test_client().get(f"{_PREFIX}/lookup") - - assert response.status_code == 400 - device_service.lookup.assert_not_called() - - -def test_lookup_returns_public_validity_without_token_data(openapi_app: Flask, device_service: Mock) -> None: - device_service.lookup.return_value = DeviceLookup(valid=False, expires_in_remaining=0, client_id=None) - - response = openapi_app.test_client().get(f"{_PREFIX}/lookup?user_code=ABCD-EFGH") - - assert response.status_code == 200 - assert response.json == {"valid": False, "expires_in_remaining": 0, "client_id": None} - device_service.lookup.assert_called_once_with(user_code="ABCD-EFGH") - - -@pytest.mark.parametrize( - ("path", "method", "error", "status", "code"), - [ - ("code", "start", UnsupportedClientError(), 400, "unsupported_client"), - ("token", "poll", PollTooFastError(), 400, "slow_down"), - ("token", "poll", ExpiredTokenError(), 400, "expired_token"), - ("token", "poll", AuthorizationPendingError(), 400, "authorization_pending"), - ("token", "poll", AccessDeniedError(), 400, "access_denied"), - ("approve", "approve", ExpiredOrUnknownError(), 404, "expired_or_unknown"), - ("approve", "approve", AlreadyResolvedError(), 409, "already_resolved"), - ("approve", "approve", ApprovalInProgressError(), 409, "approve_in_progress"), - ("approve", "approve", ApprovalOutcomeUnknownError(), 503, "approval_outcome_unknown"), - ("deny", "deny", DeviceStateLostError(), 409, "state_lost"), - ], -) -def test_device_errors_keep_oauth_response_contract( - openapi_app: Flask, device_service: Mock, path: str, method: str, error: OAuthDeviceError, status: int, code: str -) -> None: - methods = { - "start": device_service.start, - "poll": device_service.poll, - "approve": device_service.approve, - "deny": device_service.deny, - } - methods[method].side_effect = error - - response = openapi_app.test_client().post( - f"{_PREFIX}/{path}", - json={"client_id": "difyctl", "device_label": "CLI", "device_code": "opaque-code", "user_code": "ABCD-EFGH"}, - ) - - assert response.status_code == status - assert response.json == {"error": code} - - -@pytest.mark.parametrize("external", [False, True], ids=["account", "external"]) -def test_poll_serializes_subject_specific_token_payload( - openapi_app: Flask, device_service: Mock, external: bool -) -> None: - payload: dict[str, object] = { - "token": "secret-token", - "expires_at": "2030-01-01T00:00:00Z", - "token_id": "token-1", - "subject_type": "external_sso" if external else "account", - "account": None if external else {"id": "account-1", "email": "ada@example.com", "name": "Ada"}, - "workspaces": [], - "default_workspace_id": None, - } - if external: - payload.update(subject_email="external@example.com", subject_issuer="https://idp.example") - device_service.poll.return_value = payload - - response = openapi_app.test_client().post( - f"{_PREFIX}/token", - json={"device_code": "opaque-code", "client_id": "difyctl"}, - environ_overrides={"REMOTE_ADDR": "192.0.2.10"}, - ) - - assert response.status_code == 200 - assert response.json == payload - device_service.poll.assert_called_once_with(device_code="opaque-code", poll_ip="192.0.2.10") - - -@pytest.mark.parametrize("path", ["approve", "deny"]) -def test_account_decision_uses_admitted_identity(openapi_app: Flask, device_service: Mock, path: str) -> None: - method = device_service.approve if path == "approve" else device_service.deny - method.return_value = DeviceMutation(status="approved" if path == "approve" else "denied") - - response = openapi_app.test_client().post( - f"{_PREFIX}/{path}", - json={"user_code": "ABCD-EFGH", "account_id": "forged", "workspace_id": "foreign"}, - ) - - assert response.status_code == 200 - assert response.json == {"status": method.return_value.status} - assert method.call_args.kwargs == {"user_code": "ABCD-EFGH"} - if path == "approve": - context = method.call_args.args[0] - assert isinstance(context, RequestContext) - assert context.account_id == "account-1" - assert context.active_workspace_id == "workspace-1" - - -@pytest.mark.parametrize("path", ["approve", "deny"]) -def test_disabled_bearer_feature_blocks_account_decisions( - openapi_app: Flask, device_service: Mock, config_overrides: Callable[..., None], path: str -) -> None: - config_overrides(ENABLE_OAUTH_BEARER=False) - - response = openapi_app.test_client().post(f"{_PREFIX}/{path}", json={"user_code": "ABCD-EFGH"}) - - assert response.status_code == 503 - assert device_service.mock_calls == [] - - -def test_sso_initiation_clears_old_approval_cookie(openapi_app: Flask, device_service: Mock) -> None: - device_service.initiate_sso.return_value = DeviceSSOInitiation("https://idp.example/authorize") - - response = openapi_app.test_client().get(f"{_PREFIX}/sso-initiate?user_code=%20ABCD-EFGH%20") - - assert response.status_code == 302 - assert response.headers["Location"] == "https://idp.example/authorize" - cookie = SimpleCookie(response.headers["Set-Cookie"])[_COOKIE] - assert cookie["max-age"] == "0" - assert cookie["path"] == _PREFIX - assert device_service.initiate_sso.call_args.kwargs == {"user_code": "ABCD-EFGH"} - - -def test_sso_initiation_requires_user_code(openapi_app: Flask, device_service: Mock) -> None: - response = openapi_app.test_client().get(f"{_PREFIX}/sso-initiate?user_code=%20") - - assert response.status_code == 400 - device_service.initiate_sso.assert_not_called() - - -def test_sso_completion_sets_scoped_http_only_cookie(openapi_app: Flask, device_service: Mock) -> None: - device_service.complete_sso.return_value = DeviceSSOCompletion(user_code="ABCD-EFGH", approval_grant="signed-grant") - - response = openapi_app.test_client().get(f"{_PREFIX}/sso-complete?sso_assertion=signed-assertion") - - assert response.status_code == 302 - assert response.headers["Location"] == "/device?sso_verified=1" - cookie = SimpleCookie(response.headers["Set-Cookie"])[_COOKIE] - assert cookie.value == "signed-grant" - assert cookie["path"] == _PREFIX - assert cookie["httponly"] - assert cookie["secure"] - assert cookie["samesite"] == "Lax" - assert cookie["max-age"] == "300" - assert device_service.complete_sso.call_args.kwargs == { - "inbound_error": None, - "inbound_user_code": None, - "assertion": "signed-assertion", - } - - -@pytest.mark.parametrize("failure", ["exception", "missing-grant", "untrusted-error"]) -def test_sso_completion_failure_redirects_without_exposing_details( - openapi_app: Flask, device_service: Mock, failure: str -) -> None: - device_service.complete_sso.return_value = DeviceSSOCompletion( - error_code="private-provider-error" if failure == "untrusted-error" else None, - ) - if failure == "exception": - device_service.complete_sso.side_effect = RuntimeError("private-provider-error") - - response = openapi_app.test_client().get(f"{_PREFIX}/sso-complete?sso_assertion=private-assertion") - - assert response.status_code == 302 - assert response.headers["Location"] == "/device?sso_error=sso_failed" - assert "Set-Cookie" not in response.headers - assert b"private-" not in response.data - - -def test_approval_context_reads_cookie_and_serializes_expiry(openapi_app: Flask, device_service: Mock) -> None: - device_service.get_approval_context.return_value = DeviceApprovalContext( - subject_email="external@example.com", - subject_issuer="https://idp.example", - user_code="ABCD-EFGH", - csrf_token="csrf-token", - expires_at=datetime(2030, 1, 1, tzinfo=UTC), - ) - client = openapi_app.test_client() - client.set_cookie(_COOKIE, "signed-grant", path=_PREFIX) - - response = client.get(f"{_PREFIX}/approval-context") - - assert response.status_code == 200 - assert response.json == { - "subject_email": "external@example.com", - "subject_issuer": "https://idp.example", - "user_code": "ABCD-EFGH", - "csrf_token": "csrf-token", - "expires_at": "2030-01-01T00:00:00Z", - } - assert device_service.get_approval_context.call_args.kwargs == {"approval_grant": "signed-grant"} - - -def test_external_approval_uses_cookie_and_csrf_header_then_clears_cookie( - openapi_app: Flask, device_service: Mock -) -> None: - device_service.approve_external.return_value = DeviceMutation("approved") - client = openapi_app.test_client() - client.set_cookie(_COOKIE, "signed-grant", path=_PREFIX) - - response = client.post( - f"{_PREFIX}/approve-external", - headers={"X-CSRF-Token": "csrf-header"}, - json={"user_code": "ABCD-EFGH", "csrf_token": "forged", "approval_grant": "forged"}, - ) - - assert response.status_code == 200 - assert response.json == {"status": "approved"} - assert isinstance(device_service.approve_external.call_args.args[0], DeviceRequestContext) - assert device_service.approve_external.call_args.kwargs == { - "approval_grant": "signed-grant", - "csrf_token": "csrf-header", - "user_code": "ABCD-EFGH", - } - cookie = SimpleCookie(response.headers["Set-Cookie"])[_COOKIE] - assert cookie["max-age"] == "0" - assert cookie["path"] == _PREFIX - - -@pytest.mark.parametrize( - ("path", "method", "error", "status", "message"), - [ - ("sso-initiate", "initiate_sso", InvalidUserCodeError(), 400, "invalid_user_code"), - ("sso-initiate", "initiate_sso", OAuthDeviceSSOConfigurationError(), 502, "console_api_url_unset"), - ("sso-initiate", "initiate_sso", OAuthDeviceSSOInitiationError(), 502, "sso_initiate_failed"), - ("approval-context", "get_approval_context", InvalidApprovalSessionError(), 401, "no_session"), - ("approve-external", "approve_external", InvalidApprovalSessionError(), 401, "invalid_session"), - ("approve-external", "approve_external", ExternalApprovalRateLimitError(), 429, "rate_limited"), - ("approve-external", "approve_external", ExternalApprovalCSRFError(), 403, "csrf_mismatch"), - ("approve-external", "approve_external", ExternalUserCodeMismatchError(), 400, "user_code_mismatch"), - ("approve-external", "approve_external", ExternalUserCodeNotFoundError(), 404, "user_code_not_pending"), - ("approve-external", "approve_external", AlreadyResolvedError(), 409, "user_code_not_pending"), - ("approve-external", "approve_external", ApprovalInProgressError(), 409, "approve_in_progress"), - ("approve-external", "approve_external", ApprovalOutcomeUnknownError(), 503, "approval_outcome_unknown"), - ("approve-external", "approve_external", ExternalIdentityConflictError(), 403, "email_belongs_to_dify_account"), - ("approve-external", "approve_external", ApprovalSessionConsumedError(), 401, "session_already_consumed"), - ("approve-external", "approve_external", DeviceStateLostError(), 409, "state_lost"), - ], -) -def test_sso_errors_preserve_http_status_and_approval_cookie( - openapi_app: Flask, device_service: Mock, path: str, method: str, error: OAuthDeviceError, status: int, message: str -) -> None: - methods = { - "initiate_sso": device_service.initiate_sso, - "get_approval_context": device_service.get_approval_context, - "approve_external": device_service.approve_external, - } - methods[method].side_effect = error - client = openapi_app.test_client() - client.set_cookie(_COOKIE, "signed-grant", path=_PREFIX) - if path == "approve-external": - response = client.post(f"{_PREFIX}/{path}", json={"user_code": "ABCD-EFGH"}) - else: - response = client.get(f"{_PREFIX}/{path}?user_code=ABCD-EFGH") - - assert response.status_code == status - assert message in response.get_data(as_text=True) - assert "Set-Cookie" not in response.headers - assert "signed-grant" not in response.get_data(as_text=True) diff --git a/api/tests/unit_tests/controllers/openapi/test_device_sso.py b/api/tests/unit_tests/controllers/openapi/test_oauth_device_sso.py similarity index 100% rename from api/tests/unit_tests/controllers/openapi/test_device_sso.py rename to api/tests/unit_tests/controllers/openapi/test_oauth_device_sso.py diff --git a/api/tests/unit_tests/controllers/openapi/test_supported_app_type.py b/api/tests/unit_tests/controllers/openapi/test_supported_app_type.py deleted file mode 100644 index 5c0ed78f54e..00000000000 --- a/api/tests/unit_tests/controllers/openapi/test_supported_app_type.py +++ /dev/null @@ -1,24 +0,0 @@ -"""Unit tests for SupportedAppType — the listable subset of AppMode that the -openapi `app` face (`get app`) exposes and the CLI `--mode` whitelist derives from. -""" - -from __future__ import annotations - -from models.model import AppMode -from services.entities.app_entities import SUPPORTED_APP_TYPES, SupportedAppType - - -def test_supported_app_type_is_the_listable_subset_of_app_mode() -> None: - """SupportedAppType (and the derived SUPPORTED_APP_TYPES tuple) is exactly the - curated, listable subset of AppMode; non-app/runtime modes stay out.""" - assert {t.value for t in SupportedAppType} == { - "completion", - "chat", - "advanced-chat", - "workflow", - "agent-chat", - } - assert set(SUPPORTED_APP_TYPES) <= set(AppMode) - assert AppMode.AGENT not in SUPPORTED_APP_TYPES - assert AppMode.RAG_PIPELINE not in SUPPORTED_APP_TYPES - assert AppMode.CHANNEL not in SUPPORTED_APP_TYPES diff --git a/api/tests/unit_tests/controllers/openapi/test_workflow_events_openapi.py b/api/tests/unit_tests/controllers/openapi/test_workflow_events.py similarity index 100% rename from api/tests/unit_tests/controllers/openapi/test_workflow_events_openapi.py rename to api/tests/unit_tests/controllers/openapi/test_workflow_events.py diff --git a/api/tests/unit_tests/controllers/openapi/test_workspaces_members.py b/api/tests/unit_tests/controllers/openapi/test_workspaces.py similarity index 99% rename from api/tests/unit_tests/controllers/openapi/test_workspaces_members.py rename to api/tests/unit_tests/controllers/openapi/test_workspaces.py index e4a58e2b1c3..af4258e01e6 100644 --- a/api/tests/unit_tests/controllers/openapi/test_workspaces_members.py +++ b/api/tests/unit_tests/controllers/openapi/test_workspaces.py @@ -60,9 +60,9 @@ from services.errors.workspace import ( WorkspaceInvitationQuotaError, WorkspaceMemberLicenseQuotaError, ) -from tests.unit_tests.account_domain import AccountDomain +from tests.helpers.unit.account_domain import AccountDomain +from tests.helpers.unit.model_factories import make_account, make_tenant from tests.unit_tests.controllers.openapi.conftest import AdmittedWorld -from tests.unit_tests.model_factories import make_account, make_tenant if not hasattr(builtins, "MethodView"): builtins.MethodView = MethodView # type: ignore[attr-defined] diff --git a/api/tests/unit_tests/controllers/service_api/app/test_audio.py b/api/tests/unit_tests/controllers/service_api/app/test_audio.py index 250a5f754c0..a9cfe4220c4 100644 --- a/api/tests/unit_tests/controllers/service_api/app/test_audio.py +++ b/api/tests/unit_tests/controllers/service_api/app/test_audio.py @@ -45,7 +45,7 @@ from services.errors.audio import ( SpeechToTextDisabledServiceError, UnsupportedAudioTypeServiceError, ) -from tests.unit_tests.model_factories import make_app, make_end_user +from tests.helpers.unit.model_factories import make_app, make_end_user def _file_data(): diff --git a/api/tests/unit_tests/controllers/service_api/app/test_chat_request_payload.py b/api/tests/unit_tests/controllers/service_api/app/test_chat_request_payload.py deleted file mode 100644 index 1fb7e7009d7..00000000000 --- a/api/tests/unit_tests/controllers/service_api/app/test_chat_request_payload.py +++ /dev/null @@ -1,25 +0,0 @@ -import uuid - -import pytest -from pydantic import ValidationError - -from controllers.service_api.app.completion import ChatRequestPayload - - -def test_chat_request_payload_accepts_blank_conversation_id(): - payload = ChatRequestPayload.model_validate({"inputs": {}, "query": "hello", "conversation_id": ""}) - - assert payload.conversation_id is None - - -def test_chat_request_payload_validates_uuid(): - conversation_id = str(uuid.uuid4()) - - payload = ChatRequestPayload.model_validate({"inputs": {}, "query": "hello", "conversation_id": conversation_id}) - - assert payload.conversation_id == conversation_id - - -def test_chat_request_payload_rejects_invalid_uuid(): - with pytest.raises(ValidationError): - ChatRequestPayload.model_validate({"inputs": {}, "query": "hello", "conversation_id": "invalid"}) diff --git a/api/tests/unit_tests/controllers/service_api/app/test_completion.py b/api/tests/unit_tests/controllers/service_api/app/test_completion.py index 7f51489398a..7841e8e9ee6 100644 --- a/api/tests/unit_tests/controllers/service_api/app/test_completion.py +++ b/api/tests/unit_tests/controllers/service_api/app/test_completion.py @@ -1,21 +1,8 @@ -""" -Unit tests for Service API Completion controllers. - -Tests coverage for: -- CompletionRequestPayload and ChatRequestPayload Pydantic models -- App mode validation logic -- Error mapping from service layer to HTTP errors - -Focus on: -- Pydantic model validation (especially UUID normalization) -- Error types and their mappings -""" - import sys import uuid from decimal import Decimal from inspect import unwrap -from unittest.mock import Mock, patch +from unittest.mock import MagicMock, Mock, patch import pytest from flask import Flask @@ -25,6 +12,8 @@ from sqlalchemy.orm import Session from werkzeug.exceptions import BadRequest, NotFound import services +from controllers.service_api.app import completion as completion_module +from controllers.service_api.app import workflow as workflow_module from controllers.service_api.app.completion import ( ChatApi, ChatRequestPayload, @@ -44,6 +33,7 @@ from controllers.service_api.app.error import ( from controllers.web.error import InvokeRateLimitError as InvokeRateLimitHttpError from core.app.apps.agent_app.errors import AgentAppNotPublishedError from core.errors.error import QuotaExceededError +from core.helper.trace_id_helper import get_trace_session_id from enums import CloudPlan, DeploymentEdition from graphon.model_runtime.errors.invoke import InvokeError from graphon.model_runtime.errors.invoke import InvokeRateLimitError as ProviderInvokeRateLimitError @@ -57,16 +47,45 @@ from services.conversation_service import ConversationService from services.errors.app import IsDraftWorkflowError, WorkflowIdFormatError, WorkflowNotFoundError from services.errors.conversation import ConversationNotExistsError from services.errors.llm import InvokeRateLimitError -from tests.unit_tests.config_override import apply_config_overrides +from tests.helpers.unit.config_override import apply_config_overrides +from tests.helpers.unit.model_factories import make_app, make_end_user -@pytest.fixture -def orm_session(sqlite_engine: Engine): - models = (App, EndUser, Conversation, Message) - tables = [model.metadata.tables[model.__tablename__] for model in models] - TypeBase.metadata.create_all(sqlite_engine, tables=tables) - with Session(sqlite_engine, expire_on_commit=False) as session: - yield session +class TestChatRequestPayloadCases: + @staticmethod + def test_chat_request_payload_accepts_blank_conversation_id(): + payload = ChatRequestPayload.model_validate({"inputs": {}, "query": "hello", "conversation_id": ""}) + + assert payload.conversation_id is None + + @staticmethod + def test_chat_request_payload_validates_uuid(): + conversation_id = str(uuid.uuid4()) + + payload = ChatRequestPayload.model_validate( + {"inputs": {}, "query": "hello", "conversation_id": conversation_id} + ) + + assert payload.conversation_id == conversation_id + + @staticmethod + def test_chat_request_payload_rejects_invalid_uuid(): + with pytest.raises(ValidationError): + ChatRequestPayload.model_validate({"inputs": {}, "query": "hello", "conversation_id": "invalid"}) + + +""" +Unit tests for Service API Completion controllers. + +Tests coverage for: +- CompletionRequestPayload and ChatRequestPayload Pydantic models +- App mode validation logic +- Error mapping from service layer to HTTP errors + +Focus on: +- Pydantic model validation (especially UUID normalization) +- Error types and their mappings +""" def _persist_completion_state(session: Session, mode: AppMode) -> tuple[App, EndUser, Conversation, Message]: @@ -119,647 +138,824 @@ def _persist_completion_state(session: Session, mode: AppMode) -> tuple[App, End return app_model, end_user, conversation, message -class TestCompletionRequestPayload: - """Test suite for CompletionRequestPayload Pydantic model.""" - - def test_payload_with_required_fields(self): - """Test payload with only required inputs field.""" - payload = CompletionRequestPayload(inputs={"name": "test"}) - assert payload.inputs == {"name": "test"} - assert payload.query == "" - assert payload.files is None - assert payload.response_mode is None - assert payload.retriever_from == "dev" - - def test_payload_with_all_fields(self): - """Test payload with all fields populated.""" - payload = CompletionRequestPayload( - inputs={"user_input": "Hello"}, - query="What is AI?", - files=[{"type": "image", "url": "http://example.com/image.png"}], - response_mode="streaming", - retriever_from="api", - ) - assert payload.inputs == {"user_input": "Hello"} - assert payload.query == "What is AI?" - assert payload.files == [{"type": "image", "url": "http://example.com/image.png"}] - assert payload.response_mode == "streaming" - assert payload.retriever_from == "api" - - def test_payload_response_mode_blocking(self): - """Test payload with blocking response mode.""" - payload = CompletionRequestPayload(inputs={}, response_mode="blocking") - assert payload.response_mode == "blocking" - - def test_payload_empty_inputs(self): - """Test payload with empty inputs dict.""" - payload = CompletionRequestPayload(inputs={}) - assert payload.inputs == {} - - def test_payload_complex_inputs(self): - """Test payload with complex nested inputs.""" - complex_inputs = { - "user": {"name": "Alice", "age": 30}, - "context": ["item1", "item2"], - "settings": {"theme": "dark", "notifications": True}, - } - payload = CompletionRequestPayload(inputs=complex_inputs) - assert payload.inputs == complex_inputs - - -class TestChatRequestPayload: - """Test suite for ChatRequestPayload Pydantic model.""" - - def test_payload_with_required_fields(self): - """Test payload with required fields.""" - payload = ChatRequestPayload(inputs={"key": "value"}, query="Hello") - assert payload.inputs == {"key": "value"} - assert payload.query == "Hello" - assert payload.conversation_id is None - assert payload.auto_generate_name is True - - def test_payload_normalizes_valid_uuid_conversation_id(self): - """Test that valid UUID conversation_id is normalized.""" - valid_uuid = str(uuid.uuid4()) - payload = ChatRequestPayload(inputs={}, query="test", conversation_id=valid_uuid) - assert payload.conversation_id == valid_uuid - - def test_payload_normalizes_empty_string_conversation_id_to_none(self): - """Test that empty string conversation_id becomes None.""" - payload = ChatRequestPayload(inputs={}, query="test", conversation_id="") - assert payload.conversation_id is None - - def test_payload_normalizes_whitespace_conversation_id_to_none(self): - """Test that whitespace-only conversation_id becomes None.""" - payload = ChatRequestPayload(inputs={}, query="test", conversation_id=" ") - assert payload.conversation_id is None - - def test_payload_rejects_invalid_uuid_conversation_id(self): - """Test that invalid UUID format raises ValueError.""" - with pytest.raises( - ValueError, match="1 validation error for ChatRequestPayload\\\nconversation_id" - ) as exc_info: - ChatRequestPayload(inputs={}, query="test", conversation_id="not-a-uuid") - assert "valid UUID" in str(exc_info.value) - - def test_payload_with_workflow_id(self): - """Test payload with workflow_id for advanced chat.""" - payload = ChatRequestPayload(inputs={}, query="test", workflow_id="workflow_123") - assert payload.workflow_id == "workflow_123" - - def test_payload_streaming_mode(self): - """Test payload with streaming response mode.""" - payload = ChatRequestPayload(inputs={}, query="test", response_mode="streaming") - assert payload.response_mode == "streaming" - - def test_payload_auto_generate_name_false(self): - """Test payload with auto_generate_name explicitly false.""" - payload = ChatRequestPayload(inputs={}, query="test", auto_generate_name=False) - assert payload.auto_generate_name is False - - def test_payload_with_files(self): - """Test payload with file attachments.""" - files = [ - {"type": "image", "transfer_method": "remote_url", "url": "http://example.com/img.png"}, - {"type": "document", "transfer_method": "local_file", "upload_file_id": "file_123"}, - ] - payload = ChatRequestPayload(inputs={}, query="test", files=files) - assert payload.files == files - assert len(payload.files) == 2 - - -class TestCompletionErrorMappings: - """Test error type mappings for completion endpoints.""" - - def test_conversation_not_exists_error_exists(self): - """Test ConversationNotExistsError can be raised.""" - error = services.errors.conversation.ConversationNotExistsError() - assert isinstance(error, services.errors.conversation.ConversationNotExistsError) - - def test_conversation_completed_error_exists(self): - """Test ConversationCompletedError can be raised.""" - error = services.errors.conversation.ConversationCompletedError() - assert isinstance(error, services.errors.conversation.ConversationCompletedError) - - api_error = ConversationCompletedError() - assert api_error is not None - - def test_app_model_config_broken_error_exists(self): - """Test AppModelConfigBrokenError can be raised.""" - error = services.errors.app_model_config.AppModelConfigBrokenError() - assert isinstance(error, services.errors.app_model_config.AppModelConfigBrokenError) - - api_error = AppUnavailableError() - assert api_error is not None - - def test_workflow_not_found_error_exists(self): - """Test WorkflowNotFoundError can be raised.""" - error = WorkflowNotFoundError("Workflow not found") - assert isinstance(error, WorkflowNotFoundError) - - def test_is_draft_workflow_error_exists(self): - """Test IsDraftWorkflowError can be raised.""" - error = IsDraftWorkflowError("Workflow is in draft state") - assert isinstance(error, IsDraftWorkflowError) - - def test_workflow_id_format_error_exists(self): - """Test WorkflowIdFormatError can be raised.""" - error = WorkflowIdFormatError("Invalid workflow ID format") - assert isinstance(error, WorkflowIdFormatError) - - def test_invoke_rate_limit_error_exists(self): - """Test InvokeRateLimitError can be raised.""" - error = InvokeRateLimitError("Rate limit exceeded") - assert isinstance(error, InvokeRateLimitError) - - -class TestAppModeValidation: - """Test app mode validation logic patterns.""" - - def test_completion_mode_is_valid_for_completion_endpoint(self): - """Test that COMPLETION mode is valid for completion endpoints.""" - assert AppMode.COMPLETION == AppMode.COMPLETION - - def test_chat_modes_are_distinct_from_completion(self): - """Test that chat modes are distinct from completion mode.""" - chat_modes = {AppMode.CHAT, AppMode.AGENT_CHAT, AppMode.ADVANCED_CHAT, AppMode.AGENT} - assert AppMode.COMPLETION not in chat_modes - - def test_workflow_mode_is_distinct_from_chat_modes(self): - """Test that WORKFLOW mode is not a chat mode.""" - chat_modes = {AppMode.CHAT, AppMode.AGENT_CHAT, AppMode.ADVANCED_CHAT, AppMode.AGENT} - assert AppMode.WORKFLOW not in chat_modes - - def test_not_chat_app_error_can_be_raised(self): - """Test NotChatAppError can be raised for non-chat apps.""" - error = NotChatAppError() - assert error is not None - - def test_all_app_modes_are_defined(self): - """Test that all expected app modes are defined.""" - expected_modes = [ - "COMPLETION", - "CHAT", - "AGENT_CHAT", - "AGENT", - "ADVANCED_CHAT", - "WORKFLOW", - "CHANNEL", - "RAG_PIPELINE", - ] - for mode_name in expected_modes: - assert hasattr(AppMode, mode_name), f"AppMode.{mode_name} should exist" - - -class TestAppGenerateService: - """Test AppGenerateService integration patterns.""" - - def test_generate_method_exists(self): - """Test that AppGenerateService.generate method exists.""" - assert hasattr(AppGenerateService, "generate") - assert callable(AppGenerateService.generate) - - @patch.object(AppGenerateService, "generate") - def test_generate_returns_response(self, mock_generate, orm_session: Session): - """Test that generate returns expected response format.""" - expected = {"answer": "Hello!"} - mock_generate.return_value = expected - app_model, end_user, _, _ = _persist_completion_state(orm_session, AppMode.COMPLETION) - - result = AppGenerateService.generate( - app_model=app_model, - user=end_user, - args={"query": "Hi"}, - invoke_from=Mock(), - session=orm_session, - streaming=False, - ) - - assert result == expected - - @patch.object(AppGenerateService, "generate") - def test_generate_raises_conversation_not_exists(self, mock_generate, orm_session: Session): - """Test generate raises ConversationNotExistsError.""" - mock_generate.side_effect = services.errors.conversation.ConversationNotExistsError() - app_model, end_user, _, _ = _persist_completion_state(orm_session, AppMode.COMPLETION) - - with pytest.raises(services.errors.conversation.ConversationNotExistsError): - AppGenerateService.generate( - app_model=app_model, - user=end_user, - args={}, - invoke_from=Mock(), - session=orm_session, - streaming=False, - ) - - @patch.object(AppGenerateService, "generate") - def test_generate_raises_quota_exceeded(self, mock_generate, orm_session: Session): - """Test generate raises QuotaExceededError.""" - mock_generate.side_effect = QuotaExceededError() - app_model, end_user, _, _ = _persist_completion_state(orm_session, AppMode.COMPLETION) - - with pytest.raises(QuotaExceededError): - AppGenerateService.generate( - app_model=app_model, - user=end_user, - args={}, - invoke_from=Mock(), - session=orm_session, - streaming=False, - ) - - @patch.object(AppGenerateService, "generate") - def test_generate_raises_invoke_error(self, mock_generate, orm_session: Session): - """Test generate raises InvokeError.""" - mock_generate.side_effect = InvokeError("Model invocation failed") - app_model, end_user, _, _ = _persist_completion_state(orm_session, AppMode.COMPLETION) - - with pytest.raises(InvokeError): - AppGenerateService.generate( - app_model=app_model, - user=end_user, - args={}, - invoke_from=Mock(), - session=orm_session, - streaming=False, - ) - - -class TestCompletionControllerLogic: - """Test CompletionApi and ChatApi controller logic directly.""" - +class TestCompletionCases: + @staticmethod @pytest.fixture - def app(self): - """Create Flask test application.""" - from flask import Flask + def orm_session(sqlite_engine: Engine): + models = (App, EndUser, Conversation, Message) + tables = [model.metadata.tables[model.__tablename__] for model in models] + TypeBase.metadata.create_all(sqlite_engine, tables=tables) + with Session(sqlite_engine, expire_on_commit=False) as session: + yield session - app = Flask(__name__) - app.config["TESTING"] = True - return app + class TestCompletionRequestPayload: + """Test suite for CompletionRequestPayload Pydantic model.""" - @patch("controllers.service_api.app.completion.service_api_ns") - @patch("controllers.service_api.app.completion.AppGenerateService") - def test_completion_api_post_success( - self, mock_generate_service, mock_service_api_ns, app: Flask, orm_session: Session - ): - """Test CompletionApi.post success path.""" - from controllers.service_api.app.completion import CompletionApi + def test_payload_with_required_fields(self): + """Test payload with only required inputs field.""" + payload = CompletionRequestPayload(inputs={"name": "test"}) + assert payload.inputs == {"name": "test"} + assert payload.query == "" + assert payload.files is None + assert payload.response_mode is None + assert payload.retriever_from == "dev" - # Setup mocks - app_model, end_user, _, _ = _persist_completion_state(orm_session, AppMode.COMPLETION) + def test_payload_with_all_fields(self): + """Test payload with all fields populated.""" + payload = CompletionRequestPayload( + inputs={"user_input": "Hello"}, + query="What is AI?", + files=[{"type": "image", "url": "http://example.com/image.png"}], + response_mode="streaming", + retriever_from="api", + ) + assert payload.inputs == {"user_input": "Hello"} + assert payload.query == "What is AI?" + assert payload.files == [{"type": "image", "url": "http://example.com/image.png"}] + assert payload.response_mode == "streaming" + assert payload.retriever_from == "api" - payload_dict = {"inputs": {"text": "hello"}, "response_mode": "blocking"} - mock_service_api_ns.payload = payload_dict - mock_generate_service.generate.return_value = {"text": "response"} + def test_payload_response_mode_blocking(self): + """Test payload with blocking response mode.""" + payload = CompletionRequestPayload(inputs={}, response_mode="blocking") + assert payload.response_mode == "blocking" - with app.test_request_context(): - # Helper for compact_generate_response logic check - with patch("controllers.service_api.app.completion.helper.compact_generate_response") as mock_compact: - mock_compact.return_value = {"text": "compacted"} + def test_payload_empty_inputs(self): + """Test payload with empty inputs dict.""" + payload = CompletionRequestPayload(inputs={}) + assert payload.inputs == {} - api = CompletionApi() - response = unwrap(api.post)(api, orm_session, app_model, end_user) + def test_payload_complex_inputs(self): + """Test payload with complex nested inputs.""" + complex_inputs = { + "user": {"name": "Alice", "age": 30}, + "context": ["item1", "item2"], + "settings": {"theme": "dark", "notifications": True}, + } + payload = CompletionRequestPayload(inputs=complex_inputs) + assert payload.inputs == complex_inputs - assert response == {"text": "compacted"} - mock_generate_service.generate.assert_called_once() + class TestChatRequestPayload: + """Test suite for ChatRequestPayload Pydantic model.""" - @patch("controllers.service_api.app.completion.service_api_ns") - def test_completion_api_post_wrong_app_mode(self, mock_service_api_ns, app: Flask, orm_session: Session): - """Test CompletionApi.post with wrong app mode.""" - from controllers.service_api.app.completion import CompletionApi + def test_payload_with_required_fields(self): + """Test payload with required fields.""" + payload = ChatRequestPayload(inputs={"key": "value"}, query="Hello") + assert payload.inputs == {"key": "value"} + assert payload.query == "Hello" + assert payload.conversation_id is None + assert payload.auto_generate_name is True - app_model, end_user, _, _ = _persist_completion_state(orm_session, AppMode.CHAT) + def test_payload_normalizes_valid_uuid_conversation_id(self): + """Test that valid UUID conversation_id is normalized.""" + valid_uuid = str(uuid.uuid4()) + payload = ChatRequestPayload(inputs={}, query="test", conversation_id=valid_uuid) + assert payload.conversation_id == valid_uuid - with app.test_request_context(): - with pytest.raises(AppUnavailableError): - unwrap(CompletionApi().post)(CompletionApi(), orm_session, app_model, end_user) + def test_payload_normalizes_empty_string_conversation_id_to_none(self): + """Test that empty string conversation_id becomes None.""" + payload = ChatRequestPayload(inputs={}, query="test", conversation_id="") + assert payload.conversation_id is None - @patch("controllers.service_api.app.completion.service_api_ns") - @patch("controllers.service_api.app.completion.AppGenerateService") - def test_chat_api_post_success(self, mock_generate_service, mock_service_api_ns, app: Flask, orm_session: Session): - """Test ChatApi.post success path.""" - from controllers.service_api.app.completion import ChatApi + def test_payload_normalizes_whitespace_conversation_id_to_none(self): + """Test that whitespace-only conversation_id becomes None.""" + payload = ChatRequestPayload(inputs={}, query="test", conversation_id=" ") + assert payload.conversation_id is None - app_model, end_user, _, _ = _persist_completion_state(orm_session, AppMode.CHAT) + def test_payload_rejects_invalid_uuid_conversation_id(self): + """Test that invalid UUID format raises ValueError.""" + with pytest.raises( + ValueError, match="1 validation error for ChatRequestPayload\\\nconversation_id" + ) as exc_info: + ChatRequestPayload(inputs={}, query="test", conversation_id="not-a-uuid") + assert "valid UUID" in str(exc_info.value) - payload_dict = {"inputs": {}, "query": "hello", "response_mode": "blocking"} - mock_service_api_ns.payload = payload_dict - mock_generate_service.generate.return_value = {"text": "response"} + def test_payload_with_workflow_id(self): + """Test payload with workflow_id for advanced chat.""" + payload = ChatRequestPayload(inputs={}, query="test", workflow_id="workflow_123") + assert payload.workflow_id == "workflow_123" - with app.test_request_context(): - with patch("controllers.service_api.app.completion.helper.compact_generate_response") as mock_compact: - mock_compact.return_value = {"text": "compacted"} + def test_payload_streaming_mode(self): + """Test payload with streaming response mode.""" + payload = ChatRequestPayload(inputs={}, query="test", response_mode="streaming") + assert payload.response_mode == "streaming" - api = ChatApi() - response = unwrap(api.post)(api, orm_session, app_model, end_user) - assert response == {"text": "compacted"} + def test_payload_auto_generate_name_false(self): + """Test payload with auto_generate_name explicitly false.""" + payload = ChatRequestPayload(inputs={}, query="test", auto_generate_name=False) + assert payload.auto_generate_name is False - @patch("controllers.service_api.app.completion.service_api_ns") - def test_chat_api_post_wrong_app_mode(self, mock_service_api_ns, app: Flask, orm_session: Session): - """Test ChatApi.post with wrong app mode.""" - from controllers.service_api.app.completion import ChatApi + def test_payload_with_files(self): + """Test payload with file attachments.""" + files = [ + {"type": "image", "transfer_method": "remote_url", "url": "http://example.com/img.png"}, + {"type": "document", "transfer_method": "local_file", "upload_file_id": "file_123"}, + ] + payload = ChatRequestPayload(inputs={}, query="test", files=files) + assert payload.files == files + assert len(payload.files) == 2 - app_model, end_user, _, _ = _persist_completion_state(orm_session, AppMode.COMPLETION) + class TestCompletionErrorMappings: + """Test error type mappings for completion endpoints.""" - with app.test_request_context(): - with pytest.raises(NotChatAppError): - unwrap(ChatApi().post)(ChatApi(), orm_session, app_model, end_user) + def test_conversation_not_exists_error_exists(self): + """Test ConversationNotExistsError can be raised.""" + error = services.errors.conversation.ConversationNotExistsError() + assert isinstance(error, services.errors.conversation.ConversationNotExistsError) - @patch("controllers.service_api.app.completion.AppTaskService") - def test_completion_stop_api_success(self, mock_task_service, app: Flask, orm_session: Session): - """Test CompletionStopApi.post success.""" - from controllers.service_api.app.completion import CompletionStopApi + def test_conversation_completed_error_exists(self): + """Test ConversationCompletedError can be raised.""" + error = services.errors.conversation.ConversationCompletedError() + assert isinstance(error, services.errors.conversation.ConversationCompletedError) - app_model, end_user, _, _ = _persist_completion_state(orm_session, AppMode.COMPLETION) + api_error = ConversationCompletedError() + assert api_error is not None - with app.test_request_context(): + def test_app_model_config_broken_error_exists(self): + """Test AppModelConfigBrokenError can be raised.""" + error = services.errors.app_model_config.AppModelConfigBrokenError() + assert isinstance(error, services.errors.app_model_config.AppModelConfigBrokenError) + + api_error = AppUnavailableError() + assert api_error is not None + + def test_workflow_not_found_error_exists(self): + """Test WorkflowNotFoundError can be raised.""" + error = WorkflowNotFoundError("Workflow not found") + assert isinstance(error, WorkflowNotFoundError) + + def test_is_draft_workflow_error_exists(self): + """Test IsDraftWorkflowError can be raised.""" + error = IsDraftWorkflowError("Workflow is in draft state") + assert isinstance(error, IsDraftWorkflowError) + + def test_workflow_id_format_error_exists(self): + """Test WorkflowIdFormatError can be raised.""" + error = WorkflowIdFormatError("Invalid workflow ID format") + assert isinstance(error, WorkflowIdFormatError) + + def test_invoke_rate_limit_error_exists(self): + """Test InvokeRateLimitError can be raised.""" + error = InvokeRateLimitError("Rate limit exceeded") + assert isinstance(error, InvokeRateLimitError) + + class TestAppModeValidation: + """Test app mode validation logic patterns.""" + + def test_completion_mode_is_valid_for_completion_endpoint(self): + """Test that COMPLETION mode is valid for completion endpoints.""" + assert AppMode.COMPLETION == AppMode.COMPLETION + + def test_chat_modes_are_distinct_from_completion(self): + """Test that chat modes are distinct from completion mode.""" + chat_modes = {AppMode.CHAT, AppMode.AGENT_CHAT, AppMode.ADVANCED_CHAT, AppMode.AGENT} + assert AppMode.COMPLETION not in chat_modes + + def test_workflow_mode_is_distinct_from_chat_modes(self): + """Test that WORKFLOW mode is not a chat mode.""" + chat_modes = {AppMode.CHAT, AppMode.AGENT_CHAT, AppMode.ADVANCED_CHAT, AppMode.AGENT} + assert AppMode.WORKFLOW not in chat_modes + + def test_not_chat_app_error_can_be_raised(self): + """Test NotChatAppError can be raised for non-chat apps.""" + error = NotChatAppError() + assert error is not None + + def test_all_app_modes_are_defined(self): + """Test that all expected app modes are defined.""" + expected_modes = [ + "COMPLETION", + "CHAT", + "AGENT_CHAT", + "AGENT", + "ADVANCED_CHAT", + "WORKFLOW", + "CHANNEL", + "RAG_PIPELINE", + ] + for mode_name in expected_modes: + assert hasattr(AppMode, mode_name), f"AppMode.{mode_name} should exist" + + class TestAppGenerateService: + """Test AppGenerateService integration patterns.""" + + def test_generate_method_exists(self): + """Test that AppGenerateService.generate method exists.""" + assert hasattr(AppGenerateService, "generate") + assert callable(AppGenerateService.generate) + + @patch.object(AppGenerateService, "generate") + def test_generate_returns_response(self, mock_generate, orm_session: Session): + """Test that generate returns expected response format.""" + expected = {"answer": "Hello!"} + mock_generate.return_value = expected + app_model, end_user, _, _ = _persist_completion_state(orm_session, AppMode.COMPLETION) + + result = AppGenerateService.generate( + app_model=app_model, + user=end_user, + args={"query": "Hi"}, + invoke_from=Mock(), + session=orm_session, + streaming=False, + ) + + assert result == expected + + @patch.object(AppGenerateService, "generate") + def test_generate_raises_conversation_not_exists(self, mock_generate, orm_session: Session): + """Test generate raises ConversationNotExistsError.""" + mock_generate.side_effect = services.errors.conversation.ConversationNotExistsError() + app_model, end_user, _, _ = _persist_completion_state(orm_session, AppMode.COMPLETION) + + with pytest.raises(services.errors.conversation.ConversationNotExistsError): + AppGenerateService.generate( + app_model=app_model, + user=end_user, + args={}, + invoke_from=Mock(), + session=orm_session, + streaming=False, + ) + + @patch.object(AppGenerateService, "generate") + def test_generate_raises_quota_exceeded(self, mock_generate, orm_session: Session): + """Test generate raises QuotaExceededError.""" + mock_generate.side_effect = QuotaExceededError() + app_model, end_user, _, _ = _persist_completion_state(orm_session, AppMode.COMPLETION) + + with pytest.raises(QuotaExceededError): + AppGenerateService.generate( + app_model=app_model, + user=end_user, + args={}, + invoke_from=Mock(), + session=orm_session, + streaming=False, + ) + + @patch.object(AppGenerateService, "generate") + def test_generate_raises_invoke_error(self, mock_generate, orm_session: Session): + """Test generate raises InvokeError.""" + mock_generate.side_effect = InvokeError("Model invocation failed") + app_model, end_user, _, _ = _persist_completion_state(orm_session, AppMode.COMPLETION) + + with pytest.raises(InvokeError): + AppGenerateService.generate( + app_model=app_model, + user=end_user, + args={}, + invoke_from=Mock(), + session=orm_session, + streaming=False, + ) + + class TestCompletionControllerLogic: + """Test CompletionApi and ChatApi controller logic directly.""" + + @pytest.fixture + def app(self): + """Create Flask test application.""" + from flask import Flask + + app = Flask(__name__) + app.config["TESTING"] = True + return app + + @patch("controllers.service_api.app.completion.service_api_ns") + @patch("controllers.service_api.app.completion.AppGenerateService") + def test_completion_api_post_success( + self, mock_generate_service, mock_service_api_ns, app: Flask, orm_session: Session + ): + """Test CompletionApi.post success path.""" + from controllers.service_api.app.completion import CompletionApi + + # Setup mocks + app_model, end_user, _, _ = _persist_completion_state(orm_session, AppMode.COMPLETION) + + payload_dict = {"inputs": {"text": "hello"}, "response_mode": "blocking"} + mock_service_api_ns.payload = payload_dict + mock_generate_service.generate.return_value = {"text": "response"} + + with app.test_request_context(): + # Helper for compact_generate_response logic check + with patch("controllers.service_api.app.completion.helper.compact_generate_response") as mock_compact: + mock_compact.return_value = {"text": "compacted"} + + api = CompletionApi() + response = unwrap(api.post)(api, orm_session, app_model, end_user) + + assert response == {"text": "compacted"} + mock_generate_service.generate.assert_called_once() + + @patch("controllers.service_api.app.completion.service_api_ns") + def test_completion_api_post_wrong_app_mode(self, mock_service_api_ns, app: Flask, orm_session: Session): + """Test CompletionApi.post with wrong app mode.""" + from controllers.service_api.app.completion import CompletionApi + + app_model, end_user, _, _ = _persist_completion_state(orm_session, AppMode.CHAT) + + with app.test_request_context(): + with pytest.raises(AppUnavailableError): + unwrap(CompletionApi().post)(CompletionApi(), orm_session, app_model, end_user) + + @patch("controllers.service_api.app.completion.service_api_ns") + @patch("controllers.service_api.app.completion.AppGenerateService") + def test_chat_api_post_success( + self, mock_generate_service, mock_service_api_ns, app: Flask, orm_session: Session + ): + """Test ChatApi.post success path.""" + from controllers.service_api.app.completion import ChatApi + + app_model, end_user, _, _ = _persist_completion_state(orm_session, AppMode.CHAT) + + payload_dict = {"inputs": {}, "query": "hello", "response_mode": "blocking"} + mock_service_api_ns.payload = payload_dict + mock_generate_service.generate.return_value = {"text": "response"} + + with app.test_request_context(): + with patch("controllers.service_api.app.completion.helper.compact_generate_response") as mock_compact: + mock_compact.return_value = {"text": "compacted"} + + api = ChatApi() + response = unwrap(api.post)(api, orm_session, app_model, end_user) + assert response == {"text": "compacted"} + + @patch("controllers.service_api.app.completion.service_api_ns") + def test_chat_api_post_wrong_app_mode(self, mock_service_api_ns, app: Flask, orm_session: Session): + """Test ChatApi.post with wrong app mode.""" + from controllers.service_api.app.completion import ChatApi + + app_model, end_user, _, _ = _persist_completion_state(orm_session, AppMode.COMPLETION) + + with app.test_request_context(): + with pytest.raises(NotChatAppError): + unwrap(ChatApi().post)(ChatApi(), orm_session, app_model, end_user) + + @patch("controllers.service_api.app.completion.AppTaskService") + def test_completion_stop_api_success(self, mock_task_service, app: Flask, orm_session: Session): + """Test CompletionStopApi.post success.""" + from controllers.service_api.app.completion import CompletionStopApi + + app_model, end_user, _, _ = _persist_completion_state(orm_session, AppMode.COMPLETION) + + with app.test_request_context(): + api = CompletionStopApi() + response = api.post.__wrapped__(api, app_model, end_user, "task_id") + + assert response == ({"result": "success"}, 200) + mock_task_service.stop_task.assert_called_once() + + @patch("controllers.service_api.app.completion.AppTaskService") + def test_chat_stop_api_success(self, mock_task_service, app: Flask, orm_session: Session): + """Test ChatStopApi.post success.""" + from controllers.service_api.app.completion import ChatStopApi + + app_model, end_user, _, _ = _persist_completion_state(orm_session, AppMode.CHAT) + + with app.test_request_context(): + api = ChatStopApi() + response = api.post.__wrapped__(api, app_model, end_user, "task_id") + + assert response == ({"result": "success"}, 200) + mock_task_service.stop_task.assert_called_once() + + class TestChatRequestPayloadController: + def test_normalizes_conversation_id(self) -> None: + payload = ChatRequestPayload.model_validate( + {"inputs": {}, "query": "hi", "conversation_id": " ", "response_mode": "blocking"} + ) + assert payload.conversation_id is None + + with pytest.raises(ValidationError): + ChatRequestPayload.model_validate({"inputs": {}, "query": "hi", "conversation_id": "bad-id"}) + + class TestCompletionApiController: + def test_wrong_mode(self, app: Flask, orm_session: Session) -> None: + api = CompletionApi() + handler = unwrap(api.post) + app_model, end_user, _, _ = _persist_completion_state(orm_session, AppMode.CHAT) + + with app.test_request_context("/completion-messages", method="POST", json={"inputs": {}}): + with pytest.raises(AppUnavailableError): + handler(api, session=orm_session, app_model=app_model, end_user=end_user) + + def test_conversation_not_found( + self, app: Flask, monkeypatch: pytest.MonkeyPatch, orm_session: Session + ) -> None: + monkeypatch.setattr( + AppGenerateService, + "generate", + lambda *_args, **_kwargs: (_ for _ in ()).throw(ConversationNotExistsError()), + ) + app_model, end_user, _, _ = _persist_completion_state(orm_session, AppMode.COMPLETION) + + api = CompletionApi() + handler = unwrap(api.post) + + with app.test_request_context("/completion-messages", method="POST", json={"inputs": {}}): + with pytest.raises(NotFound): + handler(api, session=orm_session, app_model=app_model, end_user=end_user) + + class TestCompletionStopApiController: + def test_wrong_mode(self, app: Flask, orm_session: Session) -> None: api = CompletionStopApi() - response = api.post.__wrapped__(api, app_model, end_user, "task_id") + handler = unwrap(api.post) + app_model, end_user, _, _ = _persist_completion_state(orm_session, AppMode.CHAT) - assert response == ({"result": "success"}, 200) - mock_task_service.stop_task.assert_called_once() + with app.test_request_context("/completion-messages/1/stop", method="POST"): + with pytest.raises(AppUnavailableError): + handler(api, app_model=app_model, end_user=end_user, task_id="t1") - @patch("controllers.service_api.app.completion.AppTaskService") - def test_chat_stop_api_success(self, mock_task_service, app: Flask, orm_session: Session): - """Test ChatStopApi.post success.""" - from controllers.service_api.app.completion import ChatStopApi + def test_success(self, app: Flask, monkeypatch: pytest.MonkeyPatch, orm_session: Session) -> None: + stop_mock = Mock() + monkeypatch.setattr(AppTaskService, "stop_task", stop_mock) - app_model, end_user, _, _ = _persist_completion_state(orm_session, AppMode.CHAT) + api = CompletionStopApi() + handler = unwrap(api.post) + app_model, end_user, _, _ = _persist_completion_state(orm_session, AppMode.COMPLETION) - with app.test_request_context(): + with app.test_request_context("/completion-messages/1/stop", method="POST"): + response, status = handler(api, app_model=app_model, end_user=end_user, task_id="t1") + + assert status == 200 + assert response == {"result": "success"} + + class TestChatApiController: + @pytest.mark.parametrize( + ("source_error", "http_error", "status_code", "error_code"), + [ + pytest.param(InvokeRateLimitError, InvokeRateLimitHttpError, 429, "rate_limit_error", id="cloud-quota"), + pytest.param( + ProviderInvokeRateLimitError, CompletionRequestError, 400, "completion_request_error", id="provider" + ), + ], + ) + def test_maps_rate_limits_by_source( + self, + app: Flask, + monkeypatch: pytest.MonkeyPatch, + orm_session: Session, + source_error: type[InvokeRateLimitError | ProviderInvokeRateLimitError], + http_error: type[InvokeRateLimitHttpError | CompletionRequestError], + status_code: int, + error_code: str, + ) -> None: + generate = Mock(side_effect=source_error("rate limit reached")) + monkeypatch.setattr(AppGenerateService, "generate", generate) + app_model, end_user, _, _ = _persist_completion_state(orm_session, AppMode.ADVANCED_CHAT) + + api = ChatApi() + handler = unwrap(api.post) + + with app.test_request_context( + "/chat-messages", method="POST", json={"inputs": {}, "query": "hi", "response_mode": "blocking"} + ): + with pytest.raises(http_error) as exc_info: + handler(api, session=orm_session, app_model=app_model, end_user=end_user) + + generate.assert_called_once() + assert exc_info.value.code == status_code + assert exc_info.value.error_code == error_code + assert exc_info.value.description == "rate limit reached" + + def test_rejects_sandbox_plan_workflow_version( + self, app: Flask, monkeypatch: pytest.MonkeyPatch, orm_session: Session + ) -> None: + sys.modules["controllers.service_api.app.completion"] + apply_config_overrides(monkeypatch, DEPLOYMENT_EDITION=DeploymentEdition.CLOUD) + + billing_get_info = Mock(return_value={"subscription": {"plan": CloudPlan.SANDBOX}}) + generate = Mock() + monkeypatch.setattr(BillingService, "get_info", billing_get_info) + monkeypatch.setattr(AppGenerateService, "generate", generate) + + api = ChatApi() + handler = unwrap(api.post) + app_model, end_user, _, _ = _persist_completion_state(orm_session, AppMode.ADVANCED_CHAT) + workflow_id = str(uuid.uuid4()) + + with app.test_request_context( + "/chat-messages", + method="POST", + json={"inputs": {}, "query": "hi", "workflow_id": workflow_id}, + ): + with pytest.raises(WorkflowVersionExecutionNotAllowedError) as exc_info: + handler(api, session=orm_session, app_model=app_model, end_user=end_user) + + billing_get_info.assert_called_once_with(app_model.tenant_id, exclude_vector_space=True) + generate.assert_not_called() + assert exc_info.value.code == 403 + assert exc_info.value.error_code == "workflow_version_execution_not_allowed" + + @pytest.mark.parametrize( + ("deployment_edition", "plan", "workflow_id"), + [ + (DeploymentEdition.COMMUNITY, CloudPlan.SANDBOX, str(uuid.uuid4())), + (DeploymentEdition.ENTERPRISE, CloudPlan.SANDBOX, str(uuid.uuid4())), + (DeploymentEdition.CLOUD, CloudPlan.PROFESSIONAL, str(uuid.uuid4())), + (DeploymentEdition.CLOUD, CloudPlan.SANDBOX, None), + ], + ids=["community-version", "enterprise-version", "cloud-paid-version", "cloud-default-version"], + ) + def test_allows_default_or_entitled_workflow_version_execution( + self, + app: Flask, + monkeypatch: pytest.MonkeyPatch, + orm_session: Session, + deployment_edition: DeploymentEdition, + plan: CloudPlan, + workflow_id: str | None, + ) -> None: + completion_module = sys.modules["controllers.service_api.app.completion"] + apply_config_overrides(monkeypatch, DEPLOYMENT_EDITION=deployment_edition) + + billing_get_info = Mock(return_value={"subscription": {"plan": plan}}) + generate = Mock(return_value={"result": "ok"}) + monkeypatch.setattr(BillingService, "get_info", billing_get_info) + monkeypatch.setattr(AppGenerateService, "generate", generate) + monkeypatch.setattr(completion_module.helper, "compact_generate_response", lambda response: response) + + api = ChatApi() + handler = unwrap(api.post) + app_model, end_user, _, _ = _persist_completion_state(orm_session, AppMode.ADVANCED_CHAT) + request_json = {"inputs": {}, "query": "hi"} + if workflow_id: + request_json["workflow_id"] = workflow_id + + with app.test_request_context("/chat-messages", method="POST", json=request_json): + response = handler(api, session=orm_session, app_model=app_model, end_user=end_user) + + assert response == {"result": "ok"} + generate.assert_called_once() + if deployment_edition == DeploymentEdition.CLOUD and workflow_id: + billing_get_info.assert_called_once_with(app_model.tenant_id, exclude_vector_space=True) + else: + billing_get_info.assert_not_called() + + def test_wrong_mode(self, app: Flask, orm_session: Session) -> None: + api = ChatApi() + handler = unwrap(api.post) + app_model, end_user, _, _ = _persist_completion_state(orm_session, AppMode.COMPLETION) + + with app.test_request_context("/chat-messages", method="POST", json={"inputs": {}, "query": "hi"}): + with pytest.raises(NotChatAppError): + handler(api, session=orm_session, app_model=app_model, end_user=end_user) + + def test_workflow_not_found(self, app: Flask, monkeypatch: pytest.MonkeyPatch, orm_session: Session) -> None: + monkeypatch.setattr( + AppGenerateService, + "generate", + lambda *_args, **_kwargs: (_ for _ in ()).throw(WorkflowNotFoundError("missing")), + ) + + api = ChatApi() + handler = unwrap(api.post) + app_model, end_user, _, _ = _persist_completion_state(orm_session, AppMode.CHAT) + + with app.test_request_context("/chat-messages", method="POST", json={"inputs": {}, "query": "hi"}): + with pytest.raises(NotFound): + handler(api, session=orm_session, app_model=app_model, end_user=end_user) + + def test_draft_workflow(self, app: Flask, monkeypatch: pytest.MonkeyPatch, orm_session: Session) -> None: + monkeypatch.setattr( + AppGenerateService, + "generate", + lambda *_args, **_kwargs: (_ for _ in ()).throw(IsDraftWorkflowError("draft")), + ) + + api = ChatApi() + handler = unwrap(api.post) + app_model, end_user, _, _ = _persist_completion_state(orm_session, AppMode.CHAT) + + with app.test_request_context("/chat-messages", method="POST", json={"inputs": {}, "query": "hi"}): + with pytest.raises(BadRequest): + handler(api, session=orm_session, app_model=app_model, end_user=end_user) + + def test_agent_not_published_error_mapped( + self, app: Flask, monkeypatch: pytest.MonkeyPatch, orm_session: Session + ) -> None: + monkeypatch.setattr( + AppGenerateService, + "generate", + lambda *_args, **_kwargs: (_ for _ in ()).throw( + AgentAppNotPublishedError("Agent has not been published") + ), + ) + + api = ChatApi() + handler = unwrap(api.post) + app_model, end_user, _, _ = _persist_completion_state(orm_session, AppMode.AGENT) + + with app.test_request_context("/chat-messages", method="POST", json={"inputs": {}, "query": "hi"}): + with pytest.raises(AgentNotPublishedError): + handler(api, session=orm_session, app_model=app_model, end_user=end_user) + + def test_invalid_conversation_id_fails_fast_as_not_found( + self, app: Flask, monkeypatch: pytest.MonkeyPatch, orm_session: Session + ) -> None: + # A well-formed but nonexistent conversation_id must fail fast as 404, before the + # streaming generator is created. Previously the lookup only ran inside the generator, + # so an invalid id surfaced as a hang instead of a clean error. + get_conversation_mock = Mock(side_effect=ConversationNotExistsError()) + monkeypatch.setattr(ConversationService, "get_conversation", get_conversation_mock) + + generate_mock = Mock(return_value={"text": "unused"}) + monkeypatch.setattr(AppGenerateService, "generate", generate_mock) + + api = ChatApi() + handler = unwrap(api.post) + app_model, end_user, _, _ = _persist_completion_state(orm_session, AppMode.CHAT) + + with app.test_request_context( + "/chat-messages", + method="POST", + json={"inputs": {}, "query": "hi", "conversation_id": str(uuid.uuid4())}, + ): + with pytest.raises(NotFound): + handler(api, session=orm_session, app_model=app_model, end_user=end_user) + + # The lookup must run before generation, so the generator is never started. + generate_mock.assert_not_called() + assert get_conversation_mock.call_args.kwargs["session"] is orm_session + + class TestChatStopApiController: + def test_wrong_mode(self, app: Flask, orm_session: Session) -> None: api = ChatStopApi() - response = api.post.__wrapped__(api, app_model, end_user, "task_id") + handler = unwrap(api.post) + app_model, end_user, _, _ = _persist_completion_state(orm_session, AppMode.COMPLETION) - assert response == ({"result": "success"}, 200) - mock_task_service.stop_task.assert_called_once() + with app.test_request_context("/chat-messages/1/stop", method="POST"): + with pytest.raises(NotChatAppError): + handler(api, app_model=app_model, end_user=end_user, task_id="t1") -class TestChatRequestPayloadController: - def test_normalizes_conversation_id(self) -> None: - payload = ChatRequestPayload.model_validate( - {"inputs": {}, "query": "hi", "conversation_id": " ", "response_mode": "blocking"} +class _Request: + def __init__(self, *, headers=None, args=None, json=None, is_json=True): + self.headers = headers or {} + self.args = args or {} + self.json = json + self.is_json = is_json + + +def _app(mode: AppMode) -> App: + return make_app(mode=mode, icon_type=None) + + +def _end_user() -> EndUser: + return make_end_user(end_user_id="user-1", end_user_type=EndUserType.SERVICE_API) + + +def _assert_generate_trace_session_id(mock_generate_service: MagicMock, expected: str) -> None: + _, kwargs = mock_generate_service.generate.call_args + assert kwargs["args"]["trace_session_id"] == expected + + +class TestTraceSessionIdParsingCases: + @staticmethod + def test_trace_session_id_header_query_body_priority_matches_service_api_contract(): + req = _Request( + headers={"X-Trace-Session-Id": "header"}, + args={"trace_session_id": "query"}, + json={"trace_session_id": "body"}, ) - assert payload.conversation_id is None - with pytest.raises(ValidationError): - ChatRequestPayload.model_validate({"inputs": {}, "query": "hi", "conversation_id": "bad-id"}) + assert get_trace_session_id(req) == "header" - -class TestCompletionApiController: - def test_wrong_mode(self, app: Flask, orm_session: Session) -> None: - api = CompletionApi() - handler = unwrap(api.post) - app_model, end_user, _, _ = _persist_completion_state(orm_session, AppMode.CHAT) - - with app.test_request_context("/completion-messages", method="POST", json={"inputs": {}}): - with pytest.raises(AppUnavailableError): - handler(api, session=orm_session, app_model=app_model, end_user=end_user) - - def test_conversation_not_found(self, app: Flask, monkeypatch: pytest.MonkeyPatch, orm_session: Session) -> None: - monkeypatch.setattr( - AppGenerateService, - "generate", - lambda *_args, **_kwargs: (_ for _ in ()).throw(ConversationNotExistsError()), + @staticmethod + def test_trace_session_id_invalid_highest_priority_raises_bad_request(): + req = _Request( + headers={"X-Trace-Session-Id": " "}, + args={"trace_session_id": "query"}, + json={"trace_session_id": "body"}, ) - app_model, end_user, _, _ = _persist_completion_state(orm_session, AppMode.COMPLETION) - api = CompletionApi() - handler = unwrap(api.post) + with pytest.raises(BadRequest): + get_trace_session_id(req) - with app.test_request_context("/completion-messages", method="POST", json={"inputs": {}}): - with pytest.raises(NotFound): - handler(api, session=orm_session, app_model=app_model, end_user=end_user) - - -class TestCompletionStopApiController: - def test_wrong_mode(self, app: Flask, orm_session: Session) -> None: - api = CompletionStopApi() - handler = unwrap(api.post) - app_model, end_user, _, _ = _persist_completion_state(orm_session, AppMode.CHAT) - - with app.test_request_context("/completion-messages/1/stop", method="POST"): - with pytest.raises(AppUnavailableError): - handler(api, app_model=app_model, end_user=end_user, task_id="t1") - - def test_success(self, app: Flask, monkeypatch: pytest.MonkeyPatch, orm_session: Session) -> None: - stop_mock = Mock() - monkeypatch.setattr(AppTaskService, "stop_task", stop_mock) - - api = CompletionStopApi() - handler = unwrap(api.post) - app_model, end_user, _, _ = _persist_completion_state(orm_session, AppMode.COMPLETION) - - with app.test_request_context("/completion-messages/1/stop", method="POST"): - response, status = handler(api, app_model=app_model, end_user=end_user, task_id="t1") - - assert status == 200 - assert response == {"result": "success"} - - -class TestChatApiController: - @pytest.mark.parametrize( - ("source_error", "http_error", "status_code", "error_code"), - [ - pytest.param(InvokeRateLimitError, InvokeRateLimitHttpError, 429, "rate_limit_error", id="cloud-quota"), - pytest.param( - ProviderInvokeRateLimitError, CompletionRequestError, 400, "completion_request_error", id="provider" - ), - ], - ) - def test_maps_rate_limits_by_source( - self, + @staticmethod + @patch("controllers.service_api.app.completion.AppGenerateService") + @patch("controllers.service_api.app.completion.service_api_ns") + def test_chat_api_rejects_invalid_highest_priority_query_trace_session_id_without_generating( + mock_service_api_ns: MagicMock, + mock_generate_service: MagicMock, app: Flask, - monkeypatch: pytest.MonkeyPatch, - orm_session: Session, - source_error: type[InvokeRateLimitError | ProviderInvokeRateLimitError], - http_error: type[InvokeRateLimitHttpError | CompletionRequestError], - status_code: int, - error_code: str, - ) -> None: - generate = Mock(side_effect=source_error("rate limit reached")) - monkeypatch.setattr(AppGenerateService, "generate", generate) - app_model, end_user, _, _ = _persist_completion_state(orm_session, AppMode.ADVANCED_CHAT) - - api = ChatApi() - handler = unwrap(api.post) + ): + payload = {"inputs": {}, "query": "hello", "trace_session_id": "body-session"} + mock_service_api_ns.payload = payload with app.test_request_context( - "/chat-messages", method="POST", json={"inputs": {}, "query": "hi", "response_mode": "blocking"} - ): - with pytest.raises(http_error) as exc_info: - handler(api, session=orm_session, app_model=app_model, end_user=end_user) - - generate.assert_called_once() - assert exc_info.value.code == status_code - assert exc_info.value.error_code == error_code - assert exc_info.value.description == "rate limit reached" - - def test_rejects_sandbox_plan_workflow_version( - self, app: Flask, monkeypatch: pytest.MonkeyPatch, orm_session: Session - ) -> None: - sys.modules["controllers.service_api.app.completion"] - apply_config_overrides(monkeypatch, DEPLOYMENT_EDITION=DeploymentEdition.CLOUD) - - billing_get_info = Mock(return_value={"subscription": {"plan": CloudPlan.SANDBOX}}) - generate = Mock() - monkeypatch.setattr(BillingService, "get_info", billing_get_info) - monkeypatch.setattr(AppGenerateService, "generate", generate) - - api = ChatApi() - handler = unwrap(api.post) - app_model, end_user, _, _ = _persist_completion_state(orm_session, AppMode.ADVANCED_CHAT) - workflow_id = str(uuid.uuid4()) - - with app.test_request_context( - "/chat-messages", + "/chat-messages?trace_session_id=%20%20%20", method="POST", - json={"inputs": {}, "query": "hi", "workflow_id": workflow_id}, + json=payload, ): - with pytest.raises(WorkflowVersionExecutionNotAllowedError) as exc_info: - handler(api, session=orm_session, app_model=app_model, end_user=end_user) + with pytest.raises(BadRequest): + completion_module.ChatApi().post.__wrapped__( + completion_module.ChatApi(), + _app(AppMode.CHAT), + _end_user(), + ) - billing_get_info.assert_called_once_with(app_model.tenant_id, exclude_vector_space=True) - generate.assert_not_called() - assert exc_info.value.code == 403 - assert exc_info.value.error_code == "workflow_version_execution_not_allowed" + mock_generate_service.generate.assert_not_called() - @pytest.mark.parametrize( - ("deployment_edition", "plan", "workflow_id"), - [ - (DeploymentEdition.COMMUNITY, CloudPlan.SANDBOX, str(uuid.uuid4())), - (DeploymentEdition.ENTERPRISE, CloudPlan.SANDBOX, str(uuid.uuid4())), - (DeploymentEdition.CLOUD, CloudPlan.PROFESSIONAL, str(uuid.uuid4())), - (DeploymentEdition.CLOUD, CloudPlan.SANDBOX, None), - ], - ids=["community-version", "enterprise-version", "cloud-paid-version", "cloud-default-version"], - ) - def test_allows_default_or_entitled_workflow_version_execution( - self, + @staticmethod + @patch("controllers.service_api.app.workflow.AppGenerateService") + @patch("controllers.service_api.app.workflow.service_api_ns") + def test_workflow_run_api_rejects_invalid_highest_priority_body_trace_session_id_without_generating( + mock_service_api_ns: MagicMock, + mock_generate_service: MagicMock, app: Flask, - monkeypatch: pytest.MonkeyPatch, - orm_session: Session, - deployment_edition: DeploymentEdition, - plan: CloudPlan, - workflow_id: str | None, - ) -> None: - completion_module = sys.modules["controllers.service_api.app.completion"] - apply_config_overrides(monkeypatch, DEPLOYMENT_EDITION=deployment_edition) + ): + payload = {"inputs": {}, "trace_session_id": 123} + mock_service_api_ns.payload = payload - billing_get_info = Mock(return_value={"subscription": {"plan": plan}}) - generate = Mock(return_value={"result": "ok"}) - monkeypatch.setattr(BillingService, "get_info", billing_get_info) - monkeypatch.setattr(AppGenerateService, "generate", generate) - monkeypatch.setattr(completion_module.helper, "compact_generate_response", lambda response: response) + with app.test_request_context("/workflows/run", method="POST", json=payload): + with pytest.raises(BadRequest): + workflow_module.WorkflowRunApi().post.__wrapped__( + workflow_module.WorkflowRunApi(), + _app(AppMode.WORKFLOW), + _end_user(), + ) - api = ChatApi() - handler = unwrap(api.post) - app_model, end_user, _, _ = _persist_completion_state(orm_session, AppMode.ADVANCED_CHAT) - request_json = {"inputs": {}, "query": "hi"} - if workflow_id: - request_json["workflow_id"] = workflow_id + mock_generate_service.generate.assert_not_called() - with app.test_request_context("/chat-messages", method="POST", json=request_json): - response = handler(api, session=orm_session, app_model=app_model, end_user=end_user) + @staticmethod + @patch("controllers.service_api.app.completion.helper.compact_generate_response", return_value={"answer": "ok"}) + @patch("controllers.service_api.app.completion.AppGenerateService") + @patch("controllers.service_api.app.completion.service_api_ns") + def test_completion_api_passes_header_trace_session_id_when_body_value_is_invalid_lower_priority( + mock_service_api_ns: MagicMock, + mock_generate_service: MagicMock, + mock_compact: MagicMock, + app: Flask, + ): + payload = {"inputs": {}, "trace_session_id": 123} + mock_service_api_ns.payload = payload + mock_generate_service.generate.return_value = "response" + + with app.test_request_context( + "/completion-messages", + method="POST", + json=payload, + headers={"X-Trace-Session-Id": " header-session "}, + ): + response = completion_module.CompletionApi().post.__wrapped__( + completion_module.CompletionApi(), + _app(AppMode.COMPLETION), + _end_user(), + ) + + assert response == {"answer": "ok"} + _assert_generate_trace_session_id(mock_generate_service, "header-session") + + @staticmethod + @patch("controllers.service_api.app.completion.helper.compact_generate_response", return_value={"answer": "ok"}) + @patch("controllers.service_api.app.completion.AppGenerateService") + @patch("controllers.service_api.app.completion.service_api_ns") + def test_chat_api_passes_query_trace_session_id_when_body_value_is_invalid_lower_priority( + mock_service_api_ns: MagicMock, + mock_generate_service: MagicMock, + mock_compact: MagicMock, + app: Flask, + ): + payload = {"inputs": {}, "query": "hello", "trace_session_id": 123} + mock_service_api_ns.payload = payload + mock_generate_service.generate.return_value = "response" + + with app.test_request_context( + "/chat-messages?trace_session_id=query-session", + method="POST", + json=payload, + ): + response = completion_module.ChatApi().post.__wrapped__( + completion_module.ChatApi(), + _app(AppMode.CHAT), + _end_user(), + ) + + assert response == {"answer": "ok"} + _assert_generate_trace_session_id(mock_generate_service, "query-session") + + @staticmethod + @patch("controllers.service_api.app.workflow.helper.compact_generate_response", return_value={"result": "ok"}) + @patch("controllers.service_api.app.workflow.AppGenerateService") + @patch("controllers.service_api.app.workflow.service_api_ns") + def test_workflow_run_api_passes_body_trace_session_id( + mock_service_api_ns: MagicMock, + mock_generate_service: MagicMock, + mock_compact: MagicMock, + app: Flask, + ): + payload = {"inputs": {}, "trace_session_id": " body-session "} + mock_service_api_ns.payload = payload + mock_generate_service.generate.return_value = "response" + + with app.test_request_context("/workflows/run", method="POST", json=payload): + response = workflow_module.WorkflowRunApi().post.__wrapped__( + workflow_module.WorkflowRunApi(), + _app(AppMode.WORKFLOW), + _end_user(), + ) assert response == {"result": "ok"} - generate.assert_called_once() - if deployment_edition == DeploymentEdition.CLOUD and workflow_id: - billing_get_info.assert_called_once_with(app_model.tenant_id, exclude_vector_space=True) - else: - billing_get_info.assert_not_called() - - def test_wrong_mode(self, app: Flask, orm_session: Session) -> None: - api = ChatApi() - handler = unwrap(api.post) - app_model, end_user, _, _ = _persist_completion_state(orm_session, AppMode.COMPLETION) - - with app.test_request_context("/chat-messages", method="POST", json={"inputs": {}, "query": "hi"}): - with pytest.raises(NotChatAppError): - handler(api, session=orm_session, app_model=app_model, end_user=end_user) - - def test_workflow_not_found(self, app: Flask, monkeypatch: pytest.MonkeyPatch, orm_session: Session) -> None: - monkeypatch.setattr( - AppGenerateService, - "generate", - lambda *_args, **_kwargs: (_ for _ in ()).throw(WorkflowNotFoundError("missing")), - ) - - api = ChatApi() - handler = unwrap(api.post) - app_model, end_user, _, _ = _persist_completion_state(orm_session, AppMode.CHAT) - - with app.test_request_context("/chat-messages", method="POST", json={"inputs": {}, "query": "hi"}): - with pytest.raises(NotFound): - handler(api, session=orm_session, app_model=app_model, end_user=end_user) - - def test_draft_workflow(self, app: Flask, monkeypatch: pytest.MonkeyPatch, orm_session: Session) -> None: - monkeypatch.setattr( - AppGenerateService, - "generate", - lambda *_args, **_kwargs: (_ for _ in ()).throw(IsDraftWorkflowError("draft")), - ) - - api = ChatApi() - handler = unwrap(api.post) - app_model, end_user, _, _ = _persist_completion_state(orm_session, AppMode.CHAT) - - with app.test_request_context("/chat-messages", method="POST", json={"inputs": {}, "query": "hi"}): - with pytest.raises(BadRequest): - handler(api, session=orm_session, app_model=app_model, end_user=end_user) - - def test_agent_not_published_error_mapped( - self, app: Flask, monkeypatch: pytest.MonkeyPatch, orm_session: Session - ) -> None: - monkeypatch.setattr( - AppGenerateService, - "generate", - lambda *_args, **_kwargs: (_ for _ in ()).throw(AgentAppNotPublishedError("Agent has not been published")), - ) - - api = ChatApi() - handler = unwrap(api.post) - app_model, end_user, _, _ = _persist_completion_state(orm_session, AppMode.AGENT) - - with app.test_request_context("/chat-messages", method="POST", json={"inputs": {}, "query": "hi"}): - with pytest.raises(AgentNotPublishedError): - handler(api, session=orm_session, app_model=app_model, end_user=end_user) - - def test_invalid_conversation_id_fails_fast_as_not_found( - self, app: Flask, monkeypatch: pytest.MonkeyPatch, orm_session: Session - ) -> None: - # A well-formed but nonexistent conversation_id must fail fast as 404, before the - # streaming generator is created. Previously the lookup only ran inside the generator, - # so an invalid id surfaced as a hang instead of a clean error. - get_conversation_mock = Mock(side_effect=ConversationNotExistsError()) - monkeypatch.setattr(ConversationService, "get_conversation", get_conversation_mock) - - generate_mock = Mock(return_value={"text": "unused"}) - monkeypatch.setattr(AppGenerateService, "generate", generate_mock) - - api = ChatApi() - handler = unwrap(api.post) - app_model, end_user, _, _ = _persist_completion_state(orm_session, AppMode.CHAT) - - with app.test_request_context( - "/chat-messages", - method="POST", - json={"inputs": {}, "query": "hi", "conversation_id": str(uuid.uuid4())}, - ): - with pytest.raises(NotFound): - handler(api, session=orm_session, app_model=app_model, end_user=end_user) - - # The lookup must run before generation, so the generator is never started. - generate_mock.assert_not_called() - assert get_conversation_mock.call_args.kwargs["session"] is orm_session - - -class TestChatStopApiController: - def test_wrong_mode(self, app: Flask, orm_session: Session) -> None: - api = ChatStopApi() - handler = unwrap(api.post) - app_model, end_user, _, _ = _persist_completion_state(orm_session, AppMode.COMPLETION) - - with app.test_request_context("/chat-messages/1/stop", method="POST"): - with pytest.raises(NotChatAppError): - handler(api, app_model=app_model, end_user=end_user, task_id="t1") + _assert_generate_trace_session_id(mock_generate_service, "body-session") diff --git a/api/tests/unit_tests/controllers/service_api/app/test_conversation.py b/api/tests/unit_tests/controllers/service_api/app/test_conversation.py index 73f0b98f0de..ef4ac87f78e 100644 --- a/api/tests/unit_tests/controllers/service_api/app/test_conversation.py +++ b/api/tests/unit_tests/controllers/service_api/app/test_conversation.py @@ -1,3 +1,48 @@ +import sys +import uuid +from datetime import UTC, datetime +from inspect import unwrap +from types import SimpleNamespace +from unittest.mock import Mock, patch + +import pytest +from flask import Flask, request +from pydantic import ValidationError +from sqlalchemy.engine import Engine +from sqlalchemy.orm import Session +from werkzeug.exceptions import BadRequest, NotFound + +import services +from controllers.console.explore.conversation import ConversationRenamePayload as ConsolePayload +from controllers.service_api.app.conversation import ( + ConversationApi, + ConversationDetailApi, + ConversationListQuery, + ConversationRenameApi, + ConversationRenamePayload, + ConversationVariableDetailApi, + ConversationVariableInfiniteScrollPaginationResponse, + ConversationVariableResponse, + ConversationVariablesApi, + ConversationVariablesQuery, + ConversationVariableUpdatePayload, +) +from controllers.service_api.app.conversation import ConversationRenamePayload as ServicePayload +from controllers.service_api.app.error import NotChatAppError +from core.app.entities.app_invoke_entities import InvokeFrom +from fields._value_type_serializer import serialize_value_type +from graphon.variables import StringSegment +from graphon.variables.types import SegmentType +from models.enums import ConversationFromSource, EndUserType +from models.model import App, AppMode, Conversation, EndUser +from services.conversation_service import ConversationService +from services.errors.conversation import ( + ConversationNotExistsError, + ConversationVariableNotExistsError, + ConversationVariableTypeMismatchError, +) +from tests.helpers.unit.model_factories import make_conversation + """ Unit tests for Service API Conversation controllers. @@ -13,48 +58,6 @@ Focus on: - Error types and mappings """ -import sys -import uuid -from datetime import UTC, datetime -from inspect import unwrap -from types import SimpleNamespace -from unittest.mock import Mock, patch - -import pytest -from flask import Flask, request -from sqlalchemy.engine import Engine -from sqlalchemy.orm import Session -from werkzeug.exceptions import BadRequest, NotFound - -import services -from controllers.service_api.app.conversation import ( - ConversationApi, - ConversationDetailApi, - ConversationListQuery, - ConversationRenameApi, - ConversationRenamePayload, - ConversationVariableDetailApi, - ConversationVariableInfiniteScrollPaginationResponse, - ConversationVariableResponse, - ConversationVariablesApi, - ConversationVariablesQuery, - ConversationVariableUpdatePayload, -) -from controllers.service_api.app.error import NotChatAppError -from core.app.entities.app_invoke_entities import InvokeFrom -from fields._value_type_serializer import serialize_value_type -from graphon.variables import StringSegment -from graphon.variables.types import SegmentType -from models.enums import ConversationFromSource, EndUserType -from models.model import App, AppMode, Conversation, EndUser -from services.conversation_service import ConversationService -from services.errors.conversation import ( - ConversationNotExistsError, - ConversationVariableNotExistsError, - ConversationVariableTypeMismatchError, -) -from tests.unit_tests.model_factories import make_conversation - def _app(*, app_id: str = "app-1", mode: AppMode = AppMode.CHAT) -> App: return App( @@ -98,782 +101,791 @@ def _conversation( ) -class TestConversationListQuery: - """Test suite for ConversationListQuery Pydantic model.""" +class TestConversationCases: + class TestConversationListQuery: + """Test suite for ConversationListQuery Pydantic model.""" - def test_query_with_defaults(self): - """Test query with default values.""" - query = ConversationListQuery() - assert query.last_id is None - assert query.limit == 20 - assert query.sort_by == "-updated_at" + def test_query_with_defaults(self): + """Test query with default values.""" + query = ConversationListQuery() + assert query.last_id is None + assert query.limit == 20 + assert query.sort_by == "-updated_at" - def test_query_with_last_id(self): - """Test query with pagination last_id.""" - last_id = str(uuid.uuid4()) - query = ConversationListQuery(last_id=last_id) - assert str(query.last_id) == last_id + def test_query_with_last_id(self): + """Test query with pagination last_id.""" + last_id = str(uuid.uuid4()) + query = ConversationListQuery(last_id=last_id) + assert str(query.last_id) == last_id - def test_query_limit_boundaries(self): - """Test query respects limit boundaries.""" - query_min = ConversationListQuery(limit=1) - assert query_min.limit == 1 + def test_query_limit_boundaries(self): + """Test query respects limit boundaries.""" + query_min = ConversationListQuery(limit=1) + assert query_min.limit == 1 - query_max = ConversationListQuery(limit=100) - assert query_max.limit == 100 + query_max = ConversationListQuery(limit=100) + assert query_max.limit == 100 - def test_query_rejects_limit_below_minimum(self): - """Test query rejects limit < 1.""" - with pytest.raises(ValueError, match="1 validation error for ConversationListQuery\\\nlimit"): - ConversationListQuery(limit=0) + def test_query_rejects_limit_below_minimum(self): + """Test query rejects limit < 1.""" + with pytest.raises(ValueError, match="1 validation error for ConversationListQuery\\\nlimit"): + ConversationListQuery(limit=0) - def test_query_rejects_limit_above_maximum(self): - """Test query rejects limit > 100.""" - with pytest.raises(ValueError, match="1 validation error for ConversationListQuery\\\nlimit"): - ConversationListQuery(limit=101) + def test_query_rejects_limit_above_maximum(self): + """Test query rejects limit > 100.""" + with pytest.raises(ValueError, match="1 validation error for ConversationListQuery\\\nlimit"): + ConversationListQuery(limit=101) - @pytest.mark.parametrize( - "sort_by", - [ - "created_at", - "-created_at", - "updated_at", - "-updated_at", - ], - ) - def test_query_valid_sort_options(self, sort_by): - """Test all valid sort_by options.""" - query = ConversationListQuery(sort_by=sort_by) - assert query.sort_by == sort_by + @pytest.mark.parametrize( + "sort_by", + [ + "created_at", + "-created_at", + "updated_at", + "-updated_at", + ], + ) + def test_query_valid_sort_options(self, sort_by): + """Test all valid sort_by options.""" + query = ConversationListQuery(sort_by=sort_by) + assert query.sort_by == sort_by + + class TestConversationRenamePayload: + """Test suite for ConversationRenamePayload Pydantic model.""" + + def test_payload_with_name(self): + """Test payload with explicit name.""" + payload = ConversationRenamePayload(name="My New Chat", auto_generate=False) + assert payload.name == "My New Chat" + assert payload.auto_generate is False + + def test_payload_with_auto_generate(self): + """Test payload with auto_generate enabled.""" + payload = ConversationRenamePayload(auto_generate=True) + assert payload.auto_generate is True + assert payload.name is None + + def test_payload_requires_name_when_auto_generate_false(self): + """Test that name is required when auto_generate is False.""" + with pytest.raises( + ValueError, + match=( + "1 validation error for ConversationRenamePayload\\\n Value error, " + "name is required when auto_generate is false" + ), + ) as exc_info: + ConversationRenamePayload(auto_generate=False) + assert "name is required when auto_generate is false" in str(exc_info.value) + + def test_payload_requires_non_empty_name_when_auto_generate_false(self): + """Test that empty string name is rejected.""" + with pytest.raises( + ValueError, + match=( + "1 validation error for ConversationRenamePayload\\\n Value error, " + "name is required when auto_generate is false" + ), + ): + ConversationRenamePayload(name="", auto_generate=False) + + def test_payload_requires_non_whitespace_name_when_auto_generate_false(self): + """Test that whitespace-only name is rejected.""" + with pytest.raises( + ValueError, + match=( + "1 validation error for ConversationRenamePayload\\\n Value error, " + "name is required when auto_generate is false" + ), + ): + ConversationRenamePayload(name=" ", auto_generate=False) + + def test_payload_name_with_special_characters(self): + """Test payload with name containing special characters.""" + payload = ConversationRenamePayload(name="Chat #1 - (Test) & More!", auto_generate=False) + assert payload.name == "Chat #1 - (Test) & More!" + + def test_payload_name_with_unicode(self): + """Test payload with Unicode characters in name.""" + payload = ConversationRenamePayload(name="对话 📝 Чат", auto_generate=False) + assert payload.name == "对话 📝 Чат" + + class TestConversationVariablesQuery: + """Test suite for ConversationVariablesQuery Pydantic model.""" + + def test_query_with_defaults(self): + """Test query with default values.""" + query = ConversationVariablesQuery() + assert query.last_id is None + assert query.limit == 20 + assert query.variable_name is None + + def test_query_with_variable_name(self): + """Test query with valid variable_name filter.""" + query = ConversationVariablesQuery(variable_name="user_preference") + assert query.variable_name == "user_preference" + + def test_query_allows_hyphen_in_variable_name(self): + """Test that hyphens are allowed in variable names.""" + query = ConversationVariablesQuery(variable_name="my-variable") + assert query.variable_name == "my-variable" + + def test_query_allows_underscore_in_variable_name(self): + """Test that underscores are allowed in variable names.""" + query = ConversationVariablesQuery(variable_name="my_variable") + assert query.variable_name == "my_variable" + + def test_query_allows_period_in_variable_name(self): + """Test that periods are allowed in variable names.""" + query = ConversationVariablesQuery(variable_name="config.setting") + assert query.variable_name == "config.setting" + + def test_query_rejects_sql_injection_single_quote(self): + """Test that single quotes are rejected (SQL injection prevention).""" + with pytest.raises( + ValueError, match="1 validation error for ConversationVariablesQuery\\\nvariable_name" + ) as exc_info: + ConversationVariablesQuery(variable_name="'; DROP TABLE users;--") + assert "can only contain" in str(exc_info.value) + + def test_query_rejects_sql_injection_double_quote(self): + """Test that double quotes are rejected.""" + with pytest.raises( + ValueError, match="1 validation error for ConversationVariablesQuery\\\nvariable_name" + ) as exc_info: + ConversationVariablesQuery(variable_name='name"test') + assert "can only contain" in str(exc_info.value) + + def test_query_rejects_sql_injection_semicolon(self): + """Test that semicolons are rejected.""" + with pytest.raises( + ValueError, match="1 validation error for ConversationVariablesQuery\\\nvariable_name" + ) as exc_info: + ConversationVariablesQuery(variable_name="name;malicious") + assert "can only contain" in str(exc_info.value) + + def test_query_rejects_sql_injection_comment(self): + """Test that SQL comments are rejected.""" + with pytest.raises( + ValueError, match="1 validation error for ConversationVariablesQuery\\\nvariable_name" + ) as exc_info: + ConversationVariablesQuery(variable_name="name--comment") + assert "invalid characters" in str(exc_info.value) + + def test_query_rejects_special_characters(self): + """Test that special characters are rejected.""" + with pytest.raises( + ValueError, match="1 validation error for ConversationVariablesQuery\\\nvariable_name" + ) as exc_info: + ConversationVariablesQuery(variable_name="name@domain") + assert "can only contain" in str(exc_info.value) + + def test_query_rejects_backticks(self): + """Test that backticks are rejected (SQL injection prevention).""" + with pytest.raises( + ValueError, match="1 validation error for ConversationVariablesQuery\\\nvariable_name" + ) as exc_info: + ConversationVariablesQuery(variable_name="`table`") + assert "can only contain" in str(exc_info.value) + + def test_query_pagination_limits(self): + """Test query pagination limit boundaries.""" + query_min = ConversationVariablesQuery(limit=1) + assert query_min.limit == 1 + + query_max = ConversationVariablesQuery(limit=100) + assert query_max.limit == 100 + + class TestConversationVariableUpdatePayload: + """Test suite for ConversationVariableUpdatePayload Pydantic model.""" + + def test_payload_with_string_value(self): + """Test payload with string value.""" + payload = ConversationVariableUpdatePayload(value="hello") + assert payload.value == "hello" + + def test_payload_with_number_value(self): + """Test payload with number value.""" + payload = ConversationVariableUpdatePayload(value=42) + assert payload.value == 42 + + def test_payload_with_float_value(self): + """Test payload with float value.""" + payload = ConversationVariableUpdatePayload(value=float("3.14159")) + assert payload.value == float("3.14159") + + def test_payload_with_list_value(self): + """Test payload with list value.""" + payload = ConversationVariableUpdatePayload(value=["a", "b", "c"]) + assert payload.value == ["a", "b", "c"] + + def test_payload_with_dict_value(self): + """Test payload with dictionary value.""" + payload = ConversationVariableUpdatePayload(value={"key": "value"}) + assert payload.value == {"key": "value"} + + def test_payload_with_none_value(self): + """Test payload with None value.""" + payload = ConversationVariableUpdatePayload(value=None) + assert payload.value is None + + def test_payload_with_boolean_value(self): + """Test payload with boolean value.""" + payload = ConversationVariableUpdatePayload(value=True) + assert payload.value is True + + def test_payload_with_nested_structure(self): + """Test payload with deeply nested structure.""" + nested = {"level1": {"level2": {"level3": ["a", "b", {"c": 123}]}}} + payload = ConversationVariableUpdatePayload(value=nested) + assert payload.value == nested + + class TestConversationVariableResponseModels: + def test_variable_response_normalizes_value_type_and_timestamps(self): + created_at = datetime(2026, 1, 2, 3, 4, 5, tzinfo=UTC) + response = ConversationVariableResponse.model_validate( + { + "id": "550e8400-e29b-41d4-a716-446655440000", + "name": "foo", + "value_type": SegmentType.INTEGER, + "value": 1, + "description": "desc", + "created_at": created_at, + "updated_at": created_at, + } + ) + assert response.value_type == "number" + assert response.value == "1" + assert response.created_at == int(created_at.timestamp()) + assert response.updated_at == int(created_at.timestamp()) + + def test_variable_response_normalizes_string_value_type_alias(self): + response = ConversationVariableResponse.model_validate( + { + "id": "550e8400-e29b-41d4-a716-446655440000", + "name": "foo", + "value_type": SegmentType.INTEGER.value, + } + ) + + assert response.value_type == "number" + + def test_variable_response_normalizes_callable_exposed_type(self): + response = ConversationVariableResponse.model_validate( + { + "id": "550e8400-e29b-41d4-a716-446655440000", + "name": "foo", + "value_type": SimpleNamespace(exposed_type=lambda: SegmentType.STRING.exposed_type()), + } + ) + + assert response.value_type == "string" + + def test_serialize_value_type_supports_segments_and_mappings(self): + assert serialize_value_type(StringSegment(value="hello")) == "string" + assert serialize_value_type({"value_type": SegmentType.INTEGER}) == "number" + + def test_variable_pagination_response(self): + response = ConversationVariableInfiniteScrollPaginationResponse.model_validate( + { + "limit": 1, + "has_more": False, + "data": [ + { + "id": "550e8400-e29b-41d4-a716-446655440000", + "name": "foo", + "value_type": "string", + "value": "bar", + } + ], + } + ) + assert response.limit == 1 + assert response.has_more is False + assert len(response.data) == 1 + assert response.data[0].name == "foo" + + class TestConversationAppModeValidation: + """Test app mode validation for conversation endpoints.""" + + @pytest.mark.parametrize( + "mode", + [ + AppMode.CHAT, + AppMode.AGENT_CHAT.value, + AppMode.AGENT.value, + AppMode.ADVANCED_CHAT.value, + ], + ) + def test_chat_modes_are_valid_for_conversation_endpoints(self, mode): + """Test that all chat modes are valid for conversation endpoints. + + Verifies that CHAT, AGENT_CHAT, AGENT, and ADVANCED_CHAT modes pass + validation without raising NotChatAppError. + """ + app = App( + mode=mode, + ) + + # Validation should pass without raising for chat modes + app_mode = AppMode.value_of(app.mode) + assert app_mode in {AppMode.CHAT, AppMode.AGENT_CHAT, AppMode.ADVANCED_CHAT, AppMode.AGENT} + + def test_completion_mode_is_invalid_for_conversation_endpoints(self): + """Test that COMPLETION mode is invalid for conversation endpoints. + + Verifies that calling a conversation endpoint with a COMPLETION mode + app raises NotChatAppError. + """ + app = App( + mode=AppMode.COMPLETION, + ) + + app_mode = AppMode.value_of(app.mode) + assert app_mode not in {AppMode.CHAT, AppMode.AGENT_CHAT, AppMode.ADVANCED_CHAT, AppMode.AGENT} + with pytest.raises(NotChatAppError): + raise NotChatAppError() + + def test_workflow_mode_is_invalid_for_conversation_endpoints(self): + """Test that WORKFLOW mode is invalid for conversation endpoints. + + Verifies that calling a conversation endpoint with a WORKFLOW mode + app raises NotChatAppError. + """ + app = App( + mode=AppMode.WORKFLOW, + ) + + app_mode = AppMode.value_of(app.mode) + assert app_mode not in {AppMode.CHAT, AppMode.AGENT_CHAT, AppMode.ADVANCED_CHAT, AppMode.AGENT} + with pytest.raises(NotChatAppError): + raise NotChatAppError() + + class TestConversationErrorTypes: + """Test conversation-related error types.""" + + def test_conversation_not_exists_error(self): + """Test ConversationNotExistsError exists and can be raised.""" + error = services.errors.conversation.ConversationNotExistsError() + assert isinstance(error, services.errors.conversation.ConversationNotExistsError) + + def test_conversation_completed_error(self): + """Test ConversationCompletedError exists.""" + error = services.errors.conversation.ConversationCompletedError() + assert isinstance(error, services.errors.conversation.ConversationCompletedError) + + def test_last_conversation_not_exists_error(self): + """Test LastConversationNotExistsError exists.""" + error = services.errors.conversation.LastConversationNotExistsError() + assert isinstance(error, services.errors.conversation.LastConversationNotExistsError) + + def test_conversation_variable_not_exists_error(self): + """Test ConversationVariableNotExistsError exists.""" + error = services.errors.conversation.ConversationVariableNotExistsError() + assert isinstance(error, services.errors.conversation.ConversationVariableNotExistsError) + + def test_conversation_variable_type_mismatch_error(self): + """Test ConversationVariableTypeMismatchError exists.""" + error = services.errors.conversation.ConversationVariableTypeMismatchError("Type mismatch") + assert isinstance(error, services.errors.conversation.ConversationVariableTypeMismatchError) + + class TestConversationService: + """Test ConversationService integration patterns.""" + + def test_pagination_by_last_id_method_exists(self): + """Test that ConversationService.pagination_by_last_id exists.""" + assert hasattr(ConversationService, "pagination_by_last_id") + assert callable(ConversationService.pagination_by_last_id) + + def test_delete_method_exists(self): + """Test that ConversationService.delete exists.""" + assert hasattr(ConversationService, "delete") + assert callable(ConversationService.delete) + + def test_rename_method_exists(self): + """Test that ConversationService.rename exists.""" + assert hasattr(ConversationService, "rename") + assert callable(ConversationService.rename) + + def test_get_conversational_variable_method_exists(self): + """Test that ConversationService.get_conversational_variable exists.""" + assert hasattr(ConversationService, "get_conversational_variable") + assert callable(ConversationService.get_conversational_variable) + + def test_update_conversation_variable_method_exists(self): + """Test that ConversationService.update_conversation_variable exists.""" + assert hasattr(ConversationService, "update_conversation_variable") + assert callable(ConversationService.update_conversation_variable) + + @patch.object(ConversationService, "pagination_by_last_id") + def test_pagination_returns_expected_format(self, mock_pagination): + """Test pagination returns expected data format.""" + mock_result = Mock() + mock_result.data = [] + mock_result.limit = 20 + mock_result.has_more = False + mock_pagination.return_value = mock_result + + result = ConversationService.pagination_by_last_id( + app_model=_app(), + user=_end_user(), + last_id=None, + limit=20, + invoke_from=Mock(), + sort_by="-updated_at", + ) + + assert hasattr(result, "data") + assert hasattr(result, "limit") + assert hasattr(result, "has_more") + + def test_rename_returns_conversation(self, sqlite_session: Session): + """Test rename returns updated conversation.""" + conversation_id = "00000000-0000-0000-0000-000000000001" + conversation = _conversation(conversation_id=conversation_id) + sqlite_session.add(conversation) + sqlite_session.commit() + + result = ConversationService.rename( + app_model=App(id="app-1"), + conversation_id=conversation_id, + user=_end_user(), + name="New Name", + auto_generate=False, + session=sqlite_session, + ) + + assert result.name == "New Name" + sqlite_session.refresh(conversation) + assert conversation.name == "New Name" + + class TestConversationPayloadsController: + def test_rename_requires_name(self) -> None: + with pytest.raises( + ValueError, + match=( + "1 validation error for ConversationRenamePayload\\\n Value error, " + "name is required when auto_generate is false" + ), + ): + ConversationRenamePayload(auto_generate=False, name="") + + def test_variables_query_invalid_name(self) -> None: + with pytest.raises(ValueError, match="1 validation error for ConversationVariablesQuery\\\nvariable_name"): + ConversationVariablesQuery(variable_name="bad;") + + class TestConversationApiController: + def test_list_not_chat(self, app: Flask) -> None: + api = ConversationApi() + handler = unwrap(api.get) + app_model = _app(mode=AppMode.COMPLETION) + end_user = _end_user() + + with app.test_request_context("/conversations", method="GET"): + with pytest.raises(NotChatAppError): + handler( + api, + ConversationListQuery.model_validate(request.args.to_dict(flat=True)), + app_model=app_model, + end_user=end_user, + ) + + def test_list_last_not_found( + self, + app: Flask, + monkeypatch: pytest.MonkeyPatch, + sqlite_engine: Engine, + sqlite_session: Session, + ) -> None: + last_id = "00000000-0000-0000-0000-000000000001" + # The id exists for a different app, proving pagination cannot cross app boundaries. + sqlite_session.add(_conversation(conversation_id=last_id, app_id="other-app")) + sqlite_session.commit() + + conversation_module = sys.modules["controllers.service_api.app.conversation"] + monkeypatch.setattr(conversation_module, "db", SimpleNamespace(engine=sqlite_engine)) + + api = ConversationApi() + handler = unwrap(api.get) + app_model = _app() + end_user = _end_user() + + with app.test_request_context( + f"/conversations?last_id={last_id}&limit=20", + method="GET", + ): + with pytest.raises(NotFound): + handler( + api, + ConversationListQuery.model_validate(request.args.to_dict(flat=True)), + app_model=app_model, + end_user=end_user, + ) + + class TestConversationDetailApiController: + def test_delete_not_chat(self, app: Flask) -> None: + api = ConversationDetailApi() + handler = unwrap(api.delete) + app_model = _app(mode=AppMode.COMPLETION) + end_user = _end_user() + + with app.test_request_context("/conversations/1", method="DELETE"): + with pytest.raises(NotChatAppError): + handler( + api, + app_model=app_model, + end_user=end_user, + conversation_id="00000000-0000-0000-0000-000000000001", + ) + + def test_delete_not_found(self, app: Flask, monkeypatch: pytest.MonkeyPatch) -> None: + monkeypatch.setattr( + ConversationService, + "delete", + lambda *_args, **_kwargs: (_ for _ in ()).throw(ConversationNotExistsError()), + ) + + api = ConversationDetailApi() + handler = unwrap(api.delete) + app_model = _app() + end_user = _end_user() + + with app.test_request_context("/conversations/1", method="DELETE"): + with pytest.raises(NotFound): + handler( + api, + app_model=app_model, + end_user=end_user, + conversation_id="00000000-0000-0000-0000-000000000001", + ) + + class TestConversationRenameApiController: + def test_not_found(self, app: Flask, monkeypatch: pytest.MonkeyPatch) -> None: + monkeypatch.setattr( + ConversationService, + "rename", + lambda *_args, **_kwargs: (_ for _ in ()).throw(ConversationNotExistsError()), + ) + + api = ConversationRenameApi() + handler = unwrap(api.post) + app_model = _app() + end_user = _end_user() + + with app.test_request_context( + "/conversations/1/name", + method="POST", + json={"auto_generate": True}, + ): + payload = ConversationRenamePayload.model_validate(request.get_json() or {}) + with pytest.raises(NotFound): + handler( + api, + payload, + app_model=app_model, + end_user=end_user, + conversation_id="00000000-0000-0000-0000-000000000001", + ) + + class TestConversationVariablesApiController: + def test_not_chat(self, app: Flask) -> None: + api = ConversationVariablesApi() + handler = unwrap(api.get) + app_model = _app(mode=AppMode.COMPLETION) + end_user = _end_user() + + with app.test_request_context("/conversations/1/variables", method="GET"): + with pytest.raises(NotChatAppError): + handler( + api, + ConversationVariablesQuery.model_validate(request.args.to_dict(flat=True)), + app_model=app_model, + end_user=end_user, + conversation_id="00000000-0000-0000-0000-000000000001", + ) + + def test_not_found(self, app: Flask, monkeypatch: pytest.MonkeyPatch) -> None: + monkeypatch.setattr( + ConversationService, + "get_conversational_variable", + lambda *_args, **_kwargs: (_ for _ in ()).throw(ConversationNotExistsError()), + ) + + api = ConversationVariablesApi() + handler = unwrap(api.get) + app_model = _app() + end_user = _end_user() + + with app.test_request_context( + "/conversations/1/variables?limit=20", + method="GET", + ): + with pytest.raises(NotFound): + handler( + api, + ConversationVariablesQuery.model_validate(request.args.to_dict(flat=True)), + app_model=app_model, + end_user=end_user, + conversation_id="00000000-0000-0000-0000-000000000001", + ) + + def test_success_serializes_response(self, app: Flask, monkeypatch: pytest.MonkeyPatch) -> None: + created_at = datetime(2026, 1, 2, 3, 4, 5, tzinfo=UTC) + monkeypatch.setattr( + ConversationService, + "get_conversational_variable", + lambda *_args, **_kwargs: SimpleNamespace( + limit=1, + has_more=False, + data=[ + { + "id": "550e8400-e29b-41d4-a716-446655440000", + "name": "foo", + "value_type": SegmentType.INTEGER, + "value": 1, + "created_at": created_at, + "updated_at": created_at, + } + ], + ), + ) + + api = ConversationVariablesApi() + handler = unwrap(api.get) + app_model = _app() + end_user = _end_user() + + with app.test_request_context( + "/conversations/1/variables?limit=20", + method="GET", + ): + result = handler( + api, + ConversationVariablesQuery.model_validate(request.args.to_dict(flat=True)), + app_model=app_model, + end_user=end_user, + conversation_id="00000000-0000-0000-0000-000000000001", + ) + + assert result["limit"] == 1 + assert result["has_more"] is False + assert result["data"][0]["value_type"] == "number" + assert result["data"][0]["value"] == "1" + assert result["data"][0]["created_at"] == int(created_at.timestamp()) + + class TestConversationVariableDetailApiController: + def test_update_type_mismatch(self, app: Flask, monkeypatch: pytest.MonkeyPatch) -> None: + monkeypatch.setattr( + ConversationService, + "update_conversation_variable", + lambda *_args, **_kwargs: (_ for _ in ()).throw(ConversationVariableTypeMismatchError("bad")), + ) + + api = ConversationVariableDetailApi() + handler = unwrap(api.put) + app_model = _app() + end_user = _end_user() + + with app.test_request_context( + "/conversations/1/variables/2", + method="PUT", + json={"value": "x"}, + ): + payload = ConversationVariableUpdatePayload.model_validate(request.get_json() or {}) + with pytest.raises(BadRequest): + handler( + api, + payload, + app_model=app_model, + end_user=end_user, + conversation_id="00000000-0000-0000-0000-000000000001", + variable_id="00000000-0000-0000-0000-000000000002", + ) + + def test_update_not_found(self, app: Flask, monkeypatch: pytest.MonkeyPatch) -> None: + monkeypatch.setattr( + ConversationService, + "update_conversation_variable", + lambda *_args, **_kwargs: (_ for _ in ()).throw(ConversationVariableNotExistsError()), + ) + + api = ConversationVariableDetailApi() + handler = unwrap(api.put) + app_model = _app() + end_user = _end_user() + + with app.test_request_context( + "/conversations/1/variables/2", + method="PUT", + json={"value": "x"}, + ): + payload = ConversationVariableUpdatePayload.model_validate(request.get_json() or {}) + with pytest.raises(NotFound): + handler( + api, + payload, + app_model=app_model, + end_user=end_user, + conversation_id="00000000-0000-0000-0000-000000000001", + variable_id="00000000-0000-0000-0000-000000000002", + ) + + def test_update_success_serializes_response(self, app: Flask, monkeypatch: pytest.MonkeyPatch) -> None: + created_at = datetime(2026, 1, 2, 3, 4, 5, tzinfo=UTC) + monkeypatch.setattr( + ConversationService, + "update_conversation_variable", + lambda *_args, **_kwargs: { + "id": "550e8400-e29b-41d4-a716-446655440000", + "name": "foo", + "value_type": SegmentType.INTEGER, + "value": 1, + "created_at": created_at, + "updated_at": created_at, + }, + ) + + api = ConversationVariableDetailApi() + handler = unwrap(api.put) + app_model = _app() + end_user = _end_user() + + with app.test_request_context( + "/conversations/1/variables/2", + method="PUT", + json={"value": 1}, + ): + payload = ConversationVariableUpdatePayload.model_validate(request.get_json() or {}) + result = handler( + api, + payload, + app_model=app_model, + end_user=end_user, + conversation_id="00000000-0000-0000-0000-000000000001", + variable_id="00000000-0000-0000-0000-000000000002", + ) + + assert result["id"] == "550e8400-e29b-41d4-a716-446655440000" + assert result["value_type"] == "number" + assert result["value"] == "1" + assert result["created_at"] == int(created_at.timestamp()) -class TestConversationRenamePayload: - """Test suite for ConversationRenamePayload Pydantic model.""" +class TestConversationRenamePayloadCases: + from tests.unit_tests.conftest import app - def test_payload_with_name(self): - """Test payload with explicit name.""" - payload = ConversationRenamePayload(name="My New Chat", auto_generate=False) - assert payload.name == "My New Chat" - assert payload.auto_generate is False + app = staticmethod(app) + + @staticmethod + @pytest.mark.parametrize("payload_cls", [ConsolePayload, ServicePayload]) + def test_payload_allows_auto_generate_without_name(payload_cls): + payload = payload_cls.model_validate({"auto_generate": True}) - def test_payload_with_auto_generate(self): - """Test payload with auto_generate enabled.""" - payload = ConversationRenamePayload(auto_generate=True) assert payload.auto_generate is True assert payload.name is None - def test_payload_requires_name_when_auto_generate_false(self): - """Test that name is required when auto_generate is False.""" - with pytest.raises( - ValueError, - match=( - "1 validation error for ConversationRenamePayload\\\n Value error, " - "name is required when auto_generate is false" - ), - ) as exc_info: - ConversationRenamePayload(auto_generate=False) - assert "name is required when auto_generate is false" in str(exc_info.value) - - def test_payload_requires_non_empty_name_when_auto_generate_false(self): - """Test that empty string name is rejected.""" - with pytest.raises( - ValueError, - match=( - "1 validation error for ConversationRenamePayload\\\n Value error, " - "name is required when auto_generate is false" - ), - ): - ConversationRenamePayload(name="", auto_generate=False) - - def test_payload_requires_non_whitespace_name_when_auto_generate_false(self): - """Test that whitespace-only name is rejected.""" - with pytest.raises( - ValueError, - match=( - "1 validation error for ConversationRenamePayload\\\n Value error, " - "name is required when auto_generate is false" - ), - ): - ConversationRenamePayload(name=" ", auto_generate=False) - - def test_payload_name_with_special_characters(self): - """Test payload with name containing special characters.""" - payload = ConversationRenamePayload(name="Chat #1 - (Test) & More!", auto_generate=False) - assert payload.name == "Chat #1 - (Test) & More!" - - def test_payload_name_with_unicode(self): - """Test payload with Unicode characters in name.""" - payload = ConversationRenamePayload(name="对话 📝 Чат", auto_generate=False) - assert payload.name == "对话 📝 Чат" - - -class TestConversationVariablesQuery: - """Test suite for ConversationVariablesQuery Pydantic model.""" - - def test_query_with_defaults(self): - """Test query with default values.""" - query = ConversationVariablesQuery() - assert query.last_id is None - assert query.limit == 20 - assert query.variable_name is None - - def test_query_with_variable_name(self): - """Test query with valid variable_name filter.""" - query = ConversationVariablesQuery(variable_name="user_preference") - assert query.variable_name == "user_preference" - - def test_query_allows_hyphen_in_variable_name(self): - """Test that hyphens are allowed in variable names.""" - query = ConversationVariablesQuery(variable_name="my-variable") - assert query.variable_name == "my-variable" - - def test_query_allows_underscore_in_variable_name(self): - """Test that underscores are allowed in variable names.""" - query = ConversationVariablesQuery(variable_name="my_variable") - assert query.variable_name == "my_variable" - - def test_query_allows_period_in_variable_name(self): - """Test that periods are allowed in variable names.""" - query = ConversationVariablesQuery(variable_name="config.setting") - assert query.variable_name == "config.setting" - - def test_query_rejects_sql_injection_single_quote(self): - """Test that single quotes are rejected (SQL injection prevention).""" - with pytest.raises( - ValueError, match="1 validation error for ConversationVariablesQuery\\\nvariable_name" - ) as exc_info: - ConversationVariablesQuery(variable_name="'; DROP TABLE users;--") - assert "can only contain" in str(exc_info.value) - - def test_query_rejects_sql_injection_double_quote(self): - """Test that double quotes are rejected.""" - with pytest.raises( - ValueError, match="1 validation error for ConversationVariablesQuery\\\nvariable_name" - ) as exc_info: - ConversationVariablesQuery(variable_name='name"test') - assert "can only contain" in str(exc_info.value) - - def test_query_rejects_sql_injection_semicolon(self): - """Test that semicolons are rejected.""" - with pytest.raises( - ValueError, match="1 validation error for ConversationVariablesQuery\\\nvariable_name" - ) as exc_info: - ConversationVariablesQuery(variable_name="name;malicious") - assert "can only contain" in str(exc_info.value) - - def test_query_rejects_sql_injection_comment(self): - """Test that SQL comments are rejected.""" - with pytest.raises( - ValueError, match="1 validation error for ConversationVariablesQuery\\\nvariable_name" - ) as exc_info: - ConversationVariablesQuery(variable_name="name--comment") - assert "invalid characters" in str(exc_info.value) - - def test_query_rejects_special_characters(self): - """Test that special characters are rejected.""" - with pytest.raises( - ValueError, match="1 validation error for ConversationVariablesQuery\\\nvariable_name" - ) as exc_info: - ConversationVariablesQuery(variable_name="name@domain") - assert "can only contain" in str(exc_info.value) - - def test_query_rejects_backticks(self): - """Test that backticks are rejected (SQL injection prevention).""" - with pytest.raises( - ValueError, match="1 validation error for ConversationVariablesQuery\\\nvariable_name" - ) as exc_info: - ConversationVariablesQuery(variable_name="`table`") - assert "can only contain" in str(exc_info.value) - - def test_query_pagination_limits(self): - """Test query pagination limit boundaries.""" - query_min = ConversationVariablesQuery(limit=1) - assert query_min.limit == 1 - - query_max = ConversationVariablesQuery(limit=100) - assert query_max.limit == 100 - - -class TestConversationVariableUpdatePayload: - """Test suite for ConversationVariableUpdatePayload Pydantic model.""" - - def test_payload_with_string_value(self): - """Test payload with string value.""" - payload = ConversationVariableUpdatePayload(value="hello") - assert payload.value == "hello" - - def test_payload_with_number_value(self): - """Test payload with number value.""" - payload = ConversationVariableUpdatePayload(value=42) - assert payload.value == 42 - - def test_payload_with_float_value(self): - """Test payload with float value.""" - payload = ConversationVariableUpdatePayload(value=float("3.14159")) - assert payload.value == float("3.14159") - - def test_payload_with_list_value(self): - """Test payload with list value.""" - payload = ConversationVariableUpdatePayload(value=["a", "b", "c"]) - assert payload.value == ["a", "b", "c"] - - def test_payload_with_dict_value(self): - """Test payload with dictionary value.""" - payload = ConversationVariableUpdatePayload(value={"key": "value"}) - assert payload.value == {"key": "value"} - - def test_payload_with_none_value(self): - """Test payload with None value.""" - payload = ConversationVariableUpdatePayload(value=None) - assert payload.value is None - - def test_payload_with_boolean_value(self): - """Test payload with boolean value.""" - payload = ConversationVariableUpdatePayload(value=True) - assert payload.value is True - - def test_payload_with_nested_structure(self): - """Test payload with deeply nested structure.""" - nested = {"level1": {"level2": {"level3": ["a", "b", {"c": 123}]}}} - payload = ConversationVariableUpdatePayload(value=nested) - assert payload.value == nested - - -class TestConversationVariableResponseModels: - def test_variable_response_normalizes_value_type_and_timestamps(self): - created_at = datetime(2026, 1, 2, 3, 4, 5, tzinfo=UTC) - response = ConversationVariableResponse.model_validate( - { - "id": "550e8400-e29b-41d4-a716-446655440000", - "name": "foo", - "value_type": SegmentType.INTEGER, - "value": 1, - "description": "desc", - "created_at": created_at, - "updated_at": created_at, - } - ) - assert response.value_type == "number" - assert response.value == "1" - assert response.created_at == int(created_at.timestamp()) - assert response.updated_at == int(created_at.timestamp()) - - def test_variable_response_normalizes_string_value_type_alias(self): - response = ConversationVariableResponse.model_validate( - { - "id": "550e8400-e29b-41d4-a716-446655440000", - "name": "foo", - "value_type": SegmentType.INTEGER.value, - } - ) - - assert response.value_type == "number" - - def test_variable_response_normalizes_callable_exposed_type(self): - response = ConversationVariableResponse.model_validate( - { - "id": "550e8400-e29b-41d4-a716-446655440000", - "name": "foo", - "value_type": SimpleNamespace(exposed_type=lambda: SegmentType.STRING.exposed_type()), - } - ) - - assert response.value_type == "string" - - def test_serialize_value_type_supports_segments_and_mappings(self): - assert serialize_value_type(StringSegment(value="hello")) == "string" - assert serialize_value_type({"value_type": SegmentType.INTEGER}) == "number" - - def test_variable_pagination_response(self): - response = ConversationVariableInfiniteScrollPaginationResponse.model_validate( - { - "limit": 1, - "has_more": False, - "data": [ - { - "id": "550e8400-e29b-41d4-a716-446655440000", - "name": "foo", - "value_type": "string", - "value": "bar", - } - ], - } - ) - assert response.limit == 1 - assert response.has_more is False - assert len(response.data) == 1 - assert response.data[0].name == "foo" - - -class TestConversationAppModeValidation: - """Test app mode validation for conversation endpoints.""" - - @pytest.mark.parametrize( - "mode", - [ - AppMode.CHAT, - AppMode.AGENT_CHAT.value, - AppMode.AGENT.value, - AppMode.ADVANCED_CHAT.value, - ], - ) - def test_chat_modes_are_valid_for_conversation_endpoints(self, mode): - """Test that all chat modes are valid for conversation endpoints. - - Verifies that CHAT, AGENT_CHAT, AGENT, and ADVANCED_CHAT modes pass - validation without raising NotChatAppError. - """ - app = App( - mode=mode, - ) - - # Validation should pass without raising for chat modes - app_mode = AppMode.value_of(app.mode) - assert app_mode in {AppMode.CHAT, AppMode.AGENT_CHAT, AppMode.ADVANCED_CHAT, AppMode.AGENT} - - def test_completion_mode_is_invalid_for_conversation_endpoints(self): - """Test that COMPLETION mode is invalid for conversation endpoints. - - Verifies that calling a conversation endpoint with a COMPLETION mode - app raises NotChatAppError. - """ - app = App( - mode=AppMode.COMPLETION, - ) - - app_mode = AppMode.value_of(app.mode) - assert app_mode not in {AppMode.CHAT, AppMode.AGENT_CHAT, AppMode.ADVANCED_CHAT, AppMode.AGENT} - with pytest.raises(NotChatAppError): - raise NotChatAppError() - - def test_workflow_mode_is_invalid_for_conversation_endpoints(self): - """Test that WORKFLOW mode is invalid for conversation endpoints. - - Verifies that calling a conversation endpoint with a WORKFLOW mode - app raises NotChatAppError. - """ - app = App( - mode=AppMode.WORKFLOW, - ) - - app_mode = AppMode.value_of(app.mode) - assert app_mode not in {AppMode.CHAT, AppMode.AGENT_CHAT, AppMode.ADVANCED_CHAT, AppMode.AGENT} - with pytest.raises(NotChatAppError): - raise NotChatAppError() - - -class TestConversationErrorTypes: - """Test conversation-related error types.""" - - def test_conversation_not_exists_error(self): - """Test ConversationNotExistsError exists and can be raised.""" - error = services.errors.conversation.ConversationNotExistsError() - assert isinstance(error, services.errors.conversation.ConversationNotExistsError) - - def test_conversation_completed_error(self): - """Test ConversationCompletedError exists.""" - error = services.errors.conversation.ConversationCompletedError() - assert isinstance(error, services.errors.conversation.ConversationCompletedError) - - def test_last_conversation_not_exists_error(self): - """Test LastConversationNotExistsError exists.""" - error = services.errors.conversation.LastConversationNotExistsError() - assert isinstance(error, services.errors.conversation.LastConversationNotExistsError) - - def test_conversation_variable_not_exists_error(self): - """Test ConversationVariableNotExistsError exists.""" - error = services.errors.conversation.ConversationVariableNotExistsError() - assert isinstance(error, services.errors.conversation.ConversationVariableNotExistsError) - - def test_conversation_variable_type_mismatch_error(self): - """Test ConversationVariableTypeMismatchError exists.""" - error = services.errors.conversation.ConversationVariableTypeMismatchError("Type mismatch") - assert isinstance(error, services.errors.conversation.ConversationVariableTypeMismatchError) - - -class TestConversationService: - """Test ConversationService integration patterns.""" - - def test_pagination_by_last_id_method_exists(self): - """Test that ConversationService.pagination_by_last_id exists.""" - assert hasattr(ConversationService, "pagination_by_last_id") - assert callable(ConversationService.pagination_by_last_id) - - def test_delete_method_exists(self): - """Test that ConversationService.delete exists.""" - assert hasattr(ConversationService, "delete") - assert callable(ConversationService.delete) - - def test_rename_method_exists(self): - """Test that ConversationService.rename exists.""" - assert hasattr(ConversationService, "rename") - assert callable(ConversationService.rename) - - def test_get_conversational_variable_method_exists(self): - """Test that ConversationService.get_conversational_variable exists.""" - assert hasattr(ConversationService, "get_conversational_variable") - assert callable(ConversationService.get_conversational_variable) - - def test_update_conversation_variable_method_exists(self): - """Test that ConversationService.update_conversation_variable exists.""" - assert hasattr(ConversationService, "update_conversation_variable") - assert callable(ConversationService.update_conversation_variable) - - @patch.object(ConversationService, "pagination_by_last_id") - def test_pagination_returns_expected_format(self, mock_pagination): - """Test pagination returns expected data format.""" - mock_result = Mock() - mock_result.data = [] - mock_result.limit = 20 - mock_result.has_more = False - mock_pagination.return_value = mock_result - - result = ConversationService.pagination_by_last_id( - app_model=_app(), - user=_end_user(), - last_id=None, - limit=20, - invoke_from=Mock(), - sort_by="-updated_at", - ) - - assert hasattr(result, "data") - assert hasattr(result, "limit") - assert hasattr(result, "has_more") - - def test_rename_returns_conversation(self, sqlite_session: Session): - """Test rename returns updated conversation.""" - conversation_id = "00000000-0000-0000-0000-000000000001" - conversation = _conversation(conversation_id=conversation_id) - sqlite_session.add(conversation) - sqlite_session.commit() - - result = ConversationService.rename( - app_model=App(id="app-1"), - conversation_id=conversation_id, - user=_end_user(), - name="New Name", - auto_generate=False, - session=sqlite_session, - ) - - assert result.name == "New Name" - sqlite_session.refresh(conversation) - assert conversation.name == "New Name" - - -class TestConversationPayloadsController: - def test_rename_requires_name(self) -> None: - with pytest.raises( - ValueError, - match=( - "1 validation error for ConversationRenamePayload\\\n Value error, " - "name is required when auto_generate is false" - ), - ): - ConversationRenamePayload(auto_generate=False, name="") - - def test_variables_query_invalid_name(self) -> None: - with pytest.raises(ValueError, match="1 validation error for ConversationVariablesQuery\\\nvariable_name"): - ConversationVariablesQuery(variable_name="bad;") - - -class TestConversationApiController: - def test_list_not_chat(self, app: Flask) -> None: - api = ConversationApi() - handler = unwrap(api.get) - app_model = _app(mode=AppMode.COMPLETION) - end_user = _end_user() - - with app.test_request_context("/conversations", method="GET"): - with pytest.raises(NotChatAppError): - handler( - api, - ConversationListQuery.model_validate(request.args.to_dict(flat=True)), - app_model=app_model, - end_user=end_user, - ) - - def test_list_last_not_found( - self, - app: Flask, - monkeypatch: pytest.MonkeyPatch, - sqlite_engine: Engine, - sqlite_session: Session, - ) -> None: - last_id = "00000000-0000-0000-0000-000000000001" - # The id exists for a different app, proving pagination cannot cross app boundaries. - sqlite_session.add(_conversation(conversation_id=last_id, app_id="other-app")) - sqlite_session.commit() - - conversation_module = sys.modules["controllers.service_api.app.conversation"] - monkeypatch.setattr(conversation_module, "db", SimpleNamespace(engine=sqlite_engine)) - - api = ConversationApi() - handler = unwrap(api.get) - app_model = _app() - end_user = _end_user() - - with app.test_request_context( - f"/conversations?last_id={last_id}&limit=20", - method="GET", - ): - with pytest.raises(NotFound): - handler( - api, - ConversationListQuery.model_validate(request.args.to_dict(flat=True)), - app_model=app_model, - end_user=end_user, - ) - - -class TestConversationDetailApiController: - def test_delete_not_chat(self, app: Flask) -> None: - api = ConversationDetailApi() - handler = unwrap(api.delete) - app_model = _app(mode=AppMode.COMPLETION) - end_user = _end_user() - - with app.test_request_context("/conversations/1", method="DELETE"): - with pytest.raises(NotChatAppError): - handler( - api, - app_model=app_model, - end_user=end_user, - conversation_id="00000000-0000-0000-0000-000000000001", - ) - - def test_delete_not_found(self, app: Flask, monkeypatch: pytest.MonkeyPatch) -> None: - monkeypatch.setattr( - ConversationService, - "delete", - lambda *_args, **_kwargs: (_ for _ in ()).throw(ConversationNotExistsError()), - ) - - api = ConversationDetailApi() - handler = unwrap(api.delete) - app_model = _app() - end_user = _end_user() - - with app.test_request_context("/conversations/1", method="DELETE"): - with pytest.raises(NotFound): - handler( - api, - app_model=app_model, - end_user=end_user, - conversation_id="00000000-0000-0000-0000-000000000001", - ) - - -class TestConversationRenameApiController: - def test_not_found(self, app: Flask, monkeypatch: pytest.MonkeyPatch) -> None: - monkeypatch.setattr( - ConversationService, - "rename", - lambda *_args, **_kwargs: (_ for _ in ()).throw(ConversationNotExistsError()), - ) - - api = ConversationRenameApi() - handler = unwrap(api.post) - app_model = _app() - end_user = _end_user() - - with app.test_request_context( - "/conversations/1/name", - method="POST", - json={"auto_generate": True}, - ): - payload = ConversationRenamePayload.model_validate(request.get_json() or {}) - with pytest.raises(NotFound): - handler( - api, - payload, - app_model=app_model, - end_user=end_user, - conversation_id="00000000-0000-0000-0000-000000000001", - ) - - -class TestConversationVariablesApiController: - def test_not_chat(self, app: Flask) -> None: - api = ConversationVariablesApi() - handler = unwrap(api.get) - app_model = _app(mode=AppMode.COMPLETION) - end_user = _end_user() - - with app.test_request_context("/conversations/1/variables", method="GET"): - with pytest.raises(NotChatAppError): - handler( - api, - ConversationVariablesQuery.model_validate(request.args.to_dict(flat=True)), - app_model=app_model, - end_user=end_user, - conversation_id="00000000-0000-0000-0000-000000000001", - ) - - def test_not_found(self, app: Flask, monkeypatch: pytest.MonkeyPatch) -> None: - monkeypatch.setattr( - ConversationService, - "get_conversational_variable", - lambda *_args, **_kwargs: (_ for _ in ()).throw(ConversationNotExistsError()), - ) - - api = ConversationVariablesApi() - handler = unwrap(api.get) - app_model = _app() - end_user = _end_user() - - with app.test_request_context( - "/conversations/1/variables?limit=20", - method="GET", - ): - with pytest.raises(NotFound): - handler( - api, - ConversationVariablesQuery.model_validate(request.args.to_dict(flat=True)), - app_model=app_model, - end_user=end_user, - conversation_id="00000000-0000-0000-0000-000000000001", - ) - - def test_success_serializes_response(self, app: Flask, monkeypatch: pytest.MonkeyPatch) -> None: - created_at = datetime(2026, 1, 2, 3, 4, 5, tzinfo=UTC) - monkeypatch.setattr( - ConversationService, - "get_conversational_variable", - lambda *_args, **_kwargs: SimpleNamespace( - limit=1, - has_more=False, - data=[ - { - "id": "550e8400-e29b-41d4-a716-446655440000", - "name": "foo", - "value_type": SegmentType.INTEGER, - "value": 1, - "created_at": created_at, - "updated_at": created_at, - } - ], - ), - ) - - api = ConversationVariablesApi() - handler = unwrap(api.get) - app_model = _app() - end_user = _end_user() - - with app.test_request_context( - "/conversations/1/variables?limit=20", - method="GET", - ): - result = handler( - api, - ConversationVariablesQuery.model_validate(request.args.to_dict(flat=True)), - app_model=app_model, - end_user=end_user, - conversation_id="00000000-0000-0000-0000-000000000001", - ) - - assert result["limit"] == 1 - assert result["has_more"] is False - assert result["data"][0]["value_type"] == "number" - assert result["data"][0]["value"] == "1" - assert result["data"][0]["created_at"] == int(created_at.timestamp()) - - -class TestConversationVariableDetailApiController: - def test_update_type_mismatch(self, app: Flask, monkeypatch: pytest.MonkeyPatch) -> None: - monkeypatch.setattr( - ConversationService, - "update_conversation_variable", - lambda *_args, **_kwargs: (_ for _ in ()).throw(ConversationVariableTypeMismatchError("bad")), - ) - - api = ConversationVariableDetailApi() - handler = unwrap(api.put) - app_model = _app() - end_user = _end_user() - - with app.test_request_context( - "/conversations/1/variables/2", - method="PUT", - json={"value": "x"}, - ): - payload = ConversationVariableUpdatePayload.model_validate(request.get_json() or {}) - with pytest.raises(BadRequest): - handler( - api, - payload, - app_model=app_model, - end_user=end_user, - conversation_id="00000000-0000-0000-0000-000000000001", - variable_id="00000000-0000-0000-0000-000000000002", - ) - - def test_update_not_found(self, app: Flask, monkeypatch: pytest.MonkeyPatch) -> None: - monkeypatch.setattr( - ConversationService, - "update_conversation_variable", - lambda *_args, **_kwargs: (_ for _ in ()).throw(ConversationVariableNotExistsError()), - ) - - api = ConversationVariableDetailApi() - handler = unwrap(api.put) - app_model = _app() - end_user = _end_user() - - with app.test_request_context( - "/conversations/1/variables/2", - method="PUT", - json={"value": "x"}, - ): - payload = ConversationVariableUpdatePayload.model_validate(request.get_json() or {}) - with pytest.raises(NotFound): - handler( - api, - payload, - app_model=app_model, - end_user=end_user, - conversation_id="00000000-0000-0000-0000-000000000001", - variable_id="00000000-0000-0000-0000-000000000002", - ) - - def test_update_success_serializes_response(self, app: Flask, monkeypatch: pytest.MonkeyPatch) -> None: - created_at = datetime(2026, 1, 2, 3, 4, 5, tzinfo=UTC) - monkeypatch.setattr( - ConversationService, - "update_conversation_variable", - lambda *_args, **_kwargs: { - "id": "550e8400-e29b-41d4-a716-446655440000", - "name": "foo", - "value_type": SegmentType.INTEGER, - "value": 1, - "created_at": created_at, - "updated_at": created_at, - }, - ) - - api = ConversationVariableDetailApi() - handler = unwrap(api.put) - app_model = _app() - end_user = _end_user() - - with app.test_request_context( - "/conversations/1/variables/2", - method="PUT", - json={"value": 1}, - ): - payload = ConversationVariableUpdatePayload.model_validate(request.get_json() or {}) - result = handler( - api, - payload, - app_model=app_model, - end_user=end_user, - conversation_id="00000000-0000-0000-0000-000000000001", - variable_id="00000000-0000-0000-0000-000000000002", - ) - - assert result["id"] == "550e8400-e29b-41d4-a716-446655440000" - assert result["value_type"] == "number" - assert result["value"] == "1" - assert result["created_at"] == int(created_at.timestamp()) + @staticmethod + @pytest.mark.parametrize("payload_cls", [ConsolePayload, ServicePayload]) + @pytest.mark.parametrize("value", [None, "", " "]) + def test_payload_requires_name_when_not_auto_generate(payload_cls, value): + with pytest.raises(ValidationError): + payload_cls.model_validate({"name": value, "auto_generate": False}) diff --git a/api/tests/unit_tests/controllers/service_api/app/test_hitl_service_api.py b/api/tests/unit_tests/controllers/service_api/app/test_hitl_service_api.py deleted file mode 100644 index b9896a2597b..00000000000 --- a/api/tests/unit_tests/controllers/service_api/app/test_hitl_service_api.py +++ /dev/null @@ -1,788 +0,0 @@ -"""Dedicated tests for HITL behavior exposed through the Service API.""" - -from __future__ import annotations - -import json -import sys -from collections.abc import Sequence -from dataclasses import dataclass -from datetime import UTC, datetime -from inspect import unwrap -from types import SimpleNamespace -from typing import override -from unittest.mock import ANY, MagicMock, Mock - -import pytest -from flask import Flask -from sqlalchemy.engine import Engine -from sqlalchemy.orm import Session, sessionmaker - -import services.app_generate_service as ags_module -from controllers.service_api.app.workflow_events import WorkflowEventsApi -from core.app.app_config.entities import AppAdditionalFeatures, WorkflowUIBasedAppConfig -from core.app.apps.advanced_chat.generate_task_pipeline import ConversationSnapshot, MessageSnapshot, WorkflowSnapshot -from core.app.apps.common import workflow_response_converter -from core.app.apps.common.workflow_response_converter import WorkflowResponseConverter -from core.app.entities.app_invoke_entities import AdvancedChatAppGenerateEntity, InvokeFrom, WorkflowAppGenerateEntity -from core.app.entities.queue_entities import QueueWorkflowPausedEvent -from core.app.entities.task_entities import ( - AdvancedChatPausedBlockingResponse, - HumanInputRequiredResponse, - WorkflowAppPausedBlockingResponse, - WorkflowPauseStreamResponse, -) -from core.app.layers.pause_state_persist_layer import WorkflowResumptionContext, _WorkflowGenerateEntityWrapper -from core.workflow.human_input_policy import FormDisposition, HumanInputSurface -from core.workflow.nodes.human_input.entities import ParagraphInputConfig, UserActionConfig -from core.workflow.nodes.human_input.enums import FormInputType, HumanInputFormKind, HumanInputFormStatus -from core.workflow.nodes.human_input.pause_reason import DifyHITLEventType, HumanInputRequired -from core.workflow.system_variables import build_system_variables -from enums import DeploymentEdition -from graphon.entities import WorkflowStartReason -from graphon.enums import WorkflowExecutionStatus, WorkflowNodeExecutionStatus -from graphon.runtime import RuntimeState, VariablePool -from libs.datetime_utils import to_utc_timestamp -from models.account import Account -from models.enums import CreatorUserRole, EndUserType, MessageStatus -from models.human_input import HumanInputForm -from models.model import App, AppMode, EndUser -from models.workflow import Workflow, WorkflowRun, WorkflowType -from repositories.api_workflow_node_execution_repository import WorkflowNodeExecutionSnapshot -from repositories.entities.workflow_pause import WorkflowPauseEntity -from services.app_generate_service import AppGenerateService -from services.workflow_event_snapshot_service import _build_snapshot_events -from tests.unit_tests.config_override import apply_config_overrides - - -class _DummyRateLimit: - @staticmethod - def gen_request_key() -> str: - return "dummy-request-id" - - def __init__(self, client_id: str, max_active_requests: int) -> None: - self.client_id = client_id - self.max_active_requests = max_active_requests - - def enter(self, request_id: str | None = None) -> str: - return request_id or "dummy-request-id" - - def exit(self, request_id: str) -> None: - return None - - def generate(self, generator, request_id: str): - return generator - - -def _mock_repo_for_run(monkeypatch: pytest.MonkeyPatch, workflow_run, sqlite_engine: Engine): - workflow_events_module = sys.modules["controllers.service_api.app.workflow_events"] - repo = SimpleNamespace(get_workflow_run_by_id_and_tenant_id=lambda **_kwargs: workflow_run) - monkeypatch.setattr( - workflow_events_module.DifyAPIRepositoryFactory, - "create_api_workflow_run_repository", - lambda *_args, **_kwargs: repo, - ) - monkeypatch.setattr(workflow_events_module, "db", SimpleNamespace(engine=sqlite_engine)) - return workflow_events_module - - -def _app(*, app_id: str = "app-1", tenant_id: str = "tenant-1", mode: AppMode = AppMode.WORKFLOW) -> App: - return App( - id=app_id, - tenant_id=tenant_id, - name="Service API app", - description="", - mode=mode, - enable_site=True, - enable_api=True, - max_active_requests=0, - ) - - -def _end_user(*, user_id: str = "end-user-1", app_id: str = "app-1", tenant_id: str = "tenant-1") -> EndUser: - return EndUser( - id=user_id, - tenant_id=tenant_id, - app_id=app_id, - type=EndUserType.SERVICE_API, - external_user_id="external-user-1", - name="Service API user", - session_id="session-1", - ) - - -def _workflow(*, workflow_id: str = "workflow-id", app_id: str = "app-id", tenant_id: str = "tenant-id") -> Workflow: - return Workflow( - id=workflow_id, - tenant_id=tenant_id, - app_id=app_id, - type=WorkflowType.WORKFLOW, - version="1", - graph=json.dumps({"nodes": [], "edges": []}), - features="{}", - created_by="owner-id", - environment_variables=[], - conversation_variables=[], - rag_pipeline_variables=[], - ) - - -def _persist_human_input_form( - sqlite_session: Session, - *, - expiration_time: datetime, -) -> HumanInputForm: - form = HumanInputForm( - id="form-1", - tenant_id="tenant-1", - app_id="app-1", - workflow_run_id="run-1", - conversation_id=None, - form_kind=HumanInputFormKind.RUNTIME, - node_id="node-1", - form_definition=json.dumps({"display_in_ui": True}), - rendered_content="Rendered", - status=HumanInputFormStatus.WAITING, - expiration_time=expiration_time, - ) - sqlite_session.add(form) - sqlite_session.commit() - return form - - -def _build_service_api_pause_converter() -> WorkflowResponseConverter: - application_generate_entity = WorkflowAppGenerateEntity( - task_id="task", - workflow_execution_id="run-id", - user_id="user", - stream=True, - inputs={}, - files=[], - invoke_from=InvokeFrom.SERVICE_API, - app_config=WorkflowUIBasedAppConfig( - app_id="app-id", tenant_id="tenant-id", workflow_id="workflow-id", app_mode=AppMode.WORKFLOW - ), - ) - system_variables = build_system_variables( - user_id="user", - app_id="app-id", - workflow_id="workflow-id", - workflow_execution_id="run-id", - ) - user = Account(name="Tester", email="tester@example.com") - user.id = "account-id" - return WorkflowResponseConverter( - application_generate_entity=application_generate_entity, - user=user, - system_variables=system_variables, - ) - - -def _build_advanced_chat_paused_blocking_response() -> AdvancedChatPausedBlockingResponse: - data = AdvancedChatPausedBlockingResponse.Data( - id="msg-1", - mode="chat", - conversation_id="c1", - message_id="m1", - workflow_run_id="run-1", - answer="partial", - metadata={"usage": {"total_tokens": 1}}, - created_at=1, - paused_nodes=["node-1"], - reasons=[ - { - "TYPE": DifyHITLEventType.HUMAN_INPUT_REQUIRED.value, - "form_id": "form-1", - "expiration_time": 100, - } - ], - status=WorkflowExecutionStatus.PAUSED, - elapsed_time=0.1, - total_tokens=0, - total_steps=0, - ) - return AdvancedChatPausedBlockingResponse(task_id="t1", data=data) - - -def _build_workflow_paused_blocking_response() -> WorkflowAppPausedBlockingResponse: - return WorkflowAppPausedBlockingResponse( - task_id="t1", - workflow_run_id="r1", - data=WorkflowAppPausedBlockingResponse.Data( - id="r1", - workflow_id="wf-1", - status=WorkflowExecutionStatus.PAUSED, - outputs={}, - error=None, - elapsed_time=0.5, - total_tokens=0, - total_steps=2, - created_at=1, - finished_at=None, - paused_nodes=["node-1"], - reasons=[{"TYPE": "human_input_required", "form_id": "form-1", "expiration_time": 100}], - ), - ) - - -@dataclass(frozen=True) -class _FakePauseEntity(WorkflowPauseEntity): - pause_id: str - workflow_run_id: str - paused_at_value: datetime - pause_reasons: Sequence[HumanInputRequired] - - @property - @override - def id(self) -> str: - return self.pause_id - - @property - @override - def workflow_execution_id(self) -> str: - return self.workflow_run_id - - @override - def get_state(self) -> bytes: - raise AssertionError("state is not required for snapshot tests") - - @property - @override - def resumed_at(self) -> datetime | None: - return None - - @property - @override - def paused_at(self) -> datetime: - return self.paused_at_value - - @override - def get_pause_reasons(self) -> Sequence[HumanInputRequired]: - return self.pause_reasons - - -def _build_workflow_run(status: WorkflowExecutionStatus) -> WorkflowRun: - return WorkflowRun( - id="run-1", - tenant_id="tenant-1", - app_id="app-1", - workflow_id="workflow-1", - type="workflow", - triggered_from="app-run", - version="v1", - graph=None, - inputs=json.dumps({"input": "value"}), - status=status, - outputs=json.dumps({}), - error=None, - elapsed_time=0.0, - total_tokens=0, - total_steps=0, - created_by_role=CreatorUserRole.END_USER, - created_by="user-1", - created_at=datetime(2024, 1, 1, tzinfo=UTC), - ) - - -def _build_snapshot(status: WorkflowNodeExecutionStatus) -> WorkflowNodeExecutionSnapshot: - created_at = datetime(2024, 1, 1, tzinfo=UTC) - finished_at = datetime(2024, 1, 1, 0, 0, 5, tzinfo=UTC) - return WorkflowNodeExecutionSnapshot( - execution_id="exec-1", - node_id="node-1", - node_type="human-input", - title="Human Input", - index=1, - status=status.value, - elapsed_time=0.5, - created_at=created_at, - finished_at=finished_at, - iteration_id=None, - loop_id=None, - ) - - -def _build_resumption_context(task_id: str) -> WorkflowResumptionContext: - app_config = WorkflowUIBasedAppConfig( - tenant_id="tenant-1", - app_id="app-1", - app_mode=AppMode.WORKFLOW, - workflow_id="workflow-1", - ) - generate_entity = WorkflowAppGenerateEntity( - task_id=task_id, - app_config=app_config, - inputs={}, - files=[], - user_id="user-1", - stream=True, - invoke_from=InvokeFrom.EXPLORE, - call_depth=0, - workflow_execution_id="run-1", - ) - runtime_state = RuntimeState(workflow_id="test-workflow", variable_pool=VariablePool(), start_at=0.0) - runtime_state.set_output("result", "value") - wrapper = _WorkflowGenerateEntityWrapper(entity=generate_entity) - return WorkflowResumptionContext( - generate_entity=wrapper, - serialized_graph_runtime_state=runtime_state.dumps(), - ) - - -class TestHitlServiceApi: - # Service API event-stream continuation - def test_workflow_events_continue_on_pause_keeps_stream_open( - self, - app: Flask, - monkeypatch: pytest.MonkeyPatch, - sqlite_engine: Engine, - ) -> None: - workflow_run = _build_workflow_run(WorkflowExecutionStatus.RUNNING) - workflow_run.created_by = "end-user-1" - workflow_events_module = _mock_repo_for_run( - monkeypatch, - workflow_run=workflow_run, - sqlite_engine=sqlite_engine, - ) - msg_generator = Mock() - msg_generator.retrieve_events.return_value = [{"event": "streamed"}] - monkeypatch.setattr(workflow_events_module, "MessageGenerator", lambda: msg_generator) - - api = WorkflowEventsApi() - handler = unwrap(api.get) - app_model = _app() - end_user = _end_user() - - with app.test_request_context("/workflow/run-1/events?user=u1&continue_on_pause=true", method="GET"): - response = handler(api, app_model=app_model, end_user=end_user, workflow_run_id="run-1") - - assert response.get_data(as_text=True) == 'data: {"event":"streamed"}\n\n' - msg_generator.retrieve_events.assert_called_once_with( - AppMode.WORKFLOW, - "run-1", - terminal_events=[], - ) - - def test_workflow_events_snapshot_continue_on_pause_keeps_pause_open( - self, - app: Flask, - monkeypatch: pytest.MonkeyPatch, - sqlite_engine: Engine, - ) -> None: - workflow_run = _build_workflow_run(WorkflowExecutionStatus.RUNNING) - workflow_run.created_by = "end-user-1" - workflow_events_module = _mock_repo_for_run( - monkeypatch, - workflow_run=workflow_run, - sqlite_engine=sqlite_engine, - ) - msg_generator = Mock() - snapshot_calls: list[dict[str, object]] = [] - - def snapshot_builder(**kwargs: object) -> list[dict[str, str]]: - snapshot_calls.append(kwargs) - return [{"event": "snapshot"}] - - monkeypatch.setattr(workflow_events_module, "MessageGenerator", lambda: msg_generator) - monkeypatch.setattr(workflow_events_module, "build_workflow_event_stream", snapshot_builder) - - api = WorkflowEventsApi() - handler = unwrap(api.get) - app_model = _app() - end_user = _end_user() - - with app.test_request_context( - "/workflow/run-1/events?user=u1&include_state_snapshot=true&continue_on_pause=true", - method="GET", - ): - response = handler(api, app_model=app_model, end_user=end_user, workflow_run_id="run-1") - - assert response.get_data(as_text=True) == 'data: {"event":"snapshot"}\n\n' - msg_generator.retrieve_events.assert_not_called() - assert snapshot_calls == [ - { - "app_mode": AppMode.WORKFLOW, - "workflow_run": workflow_run, - "tenant_id": "tenant-1", - "app_id": "app-1", - "session_maker": ANY, - "human_input_surface": HumanInputSurface.SERVICE_API, - "close_on_pause": False, - } - ] - snapshot_session_maker = snapshot_calls[0]["session_maker"] - assert isinstance(snapshot_session_maker, sessionmaker) - assert snapshot_session_maker.kw["bind"] is sqlite_engine - - def test_advanced_chat_blocking_injects_pause_state_config( - self, - monkeypatch: pytest.MonkeyPatch, - sqlite_engine: Engine, - ) -> None: - apply_config_overrides(monkeypatch, DEPLOYMENT_EDITION=DeploymentEdition.COMMUNITY) - monkeypatch.setattr(ags_module, "RateLimit", _DummyRateLimit) - - workflow = _workflow() - monkeypatch.setattr(AppGenerateService, "_get_workflow", lambda *args, **kwargs: workflow) - sqlite_session_maker = sessionmaker(bind=sqlite_engine, expire_on_commit=False) - monkeypatch.setattr(ags_module.session_factory, "get_session_maker", lambda: sqlite_session_maker) - - generator_instance = MagicMock() - generator_instance.generate.return_value = {"result": "advanced-blocking"} - generator_instance.convert_to_event_stream.side_effect = lambda payload: payload - monkeypatch.setattr(ags_module, "AdvancedChatAppGenerator", lambda **_kwargs: generator_instance) - - app_model = _app(app_id="app-id", tenant_id="tenant-id", mode=AppMode.ADVANCED_CHAT) - user = _end_user(user_id="user-id", app_id="app-id", tenant_id="tenant-id") - - with sqlite_session_maker() as session: - result = AppGenerateService.generate( - session=session, - app_model=app_model, - user=user, - args={"workflow_id": None, "query": "hi", "inputs": {}}, - invoke_from=InvokeFrom.SERVICE_API, - streaming=False, - ) - - assert result == {"result": "advanced-blocking"} - call_kwargs = generator_instance.generate.call_args.kwargs - assert call_kwargs["streaming"] is False - assert call_kwargs["pause_state_config"] is not None - assert call_kwargs["pause_state_config"].session_factory is sqlite_session_maker - assert call_kwargs["pause_state_config"].state_owner_user_id == "owner-id" - - # Blocking payload contract - def test_advanced_chat_blocking_pause_payload_contract(self) -> None: - from core.app.apps.advanced_chat.generate_response_converter import AdvancedChatAppGenerateResponseConverter - - response = AdvancedChatAppGenerateResponseConverter.convert_blocking_full_response( - _build_advanced_chat_paused_blocking_response() - ) - - assert response["event"] == "workflow_paused" - assert response["workflow_run_id"] == "run-1" - assert response["answer"] == "partial" - assert response["data"]["reasons"][0]["TYPE"] == DifyHITLEventType.HUMAN_INPUT_REQUIRED - assert response["data"]["reasons"][0]["expiration_time"] == 100 - assert "human_input_forms" not in response["data"] - - def test_workflow_blocking_pause_payload_contract(self) -> None: - from core.app.apps.workflow.generate_response_converter import WorkflowAppGenerateResponseConverter - - response = WorkflowAppGenerateResponseConverter.convert_blocking_full_response( - _build_workflow_paused_blocking_response() - ) - - assert response["workflow_run_id"] == "r1" - assert response["data"]["status"] == WorkflowExecutionStatus.PAUSED - assert response["data"]["paused_nodes"] == ["node-1"] - assert response["data"]["reasons"] == [ - {"TYPE": "human_input_required", "form_id": "form-1", "expiration_time": 100} - ] - assert "human_input_forms" not in response["data"] - - def test_advanced_chat_blocking_pipeline_pause_payload_contract(self) -> None: - from core.app.app_config.entities import AppAdditionalFeatures - from core.app.apps.advanced_chat.generate_task_pipeline import AdvancedChatAppGenerateTaskPipeline - - app_config = WorkflowUIBasedAppConfig( - tenant_id="tenant", - app_id="app", - app_mode=AppMode.ADVANCED_CHAT, - additional_features=AppAdditionalFeatures(), - variables=[], - workflow_id="workflow-id", - ) - application_generate_entity = AdvancedChatAppGenerateEntity.model_construct( - task_id="task", - app_config=app_config, - inputs={}, - query="hello", - files=[], - user_id="user", - stream=False, - invoke_from=InvokeFrom.WEB_APP, - extras={}, - trace_manager=None, - workflow_run_id="run-id", - ) - pipeline = AdvancedChatAppGenerateTaskPipeline( - application_generate_entity=application_generate_entity, - workflow=WorkflowSnapshot(id="workflow-id", tenant_id="tenant", features_dict={}), - queue_manager=SimpleNamespace(invoke_from=InvokeFrom.WEB_APP, graph_runtime_state=None), - conversation=ConversationSnapshot(id="conv-id", mode=AppMode.ADVANCED_CHAT), - message=MessageSnapshot( - id="message-id", - query="hello", - created_at=datetime.utcnow(), - status=MessageStatus.NORMAL, - answer="", - ), - user=_end_user(user_id="user", app_id="app", tenant_id="tenant"), - stream=False, - dialogue_count=1, - draft_var_saver_factory=lambda **kwargs: None, - ) - pipeline._task_state.answer = "partial answer" - pipeline._workflow_run_id = "run-id" - - def _gen(): - yield HumanInputRequiredResponse( - task_id="task", - workflow_run_id="run-id", - data=HumanInputRequiredResponse.Data( - form_id="form-1", - node_id="node-1", - node_title="Approval", - form_content="Need approval", - inputs=[], - actions=[UserActionConfig(id="approve", title="Approve")], - display_in_ui=True, - form_token="token-1", - resolved_default_values={}, - expiration_time=123, - ), - ) - yield WorkflowPauseStreamResponse( - task_id="task", - workflow_run_id="run-id", - data=WorkflowPauseStreamResponse.Data( - workflow_run_id="run-id", - paused_nodes=["node-1"], - outputs={}, - reasons=[ - { - "TYPE": DifyHITLEventType.HUMAN_INPUT_REQUIRED.value, - "form_id": "form-1", - "node_id": "node-1", - "expiration_time": 123, - }, - ], - status="paused", - created_at=1, - elapsed_time=0.1, - total_tokens=0, - total_steps=0, - ), - ) - - response = pipeline._to_blocking_response(_gen()) - - assert isinstance(response, AdvancedChatPausedBlockingResponse) - assert response.data.answer == "partial answer" - assert response.data.workflow_run_id == "run-id" - assert response.data.reasons[0]["form_id"] == "form-1" - assert response.data.reasons[0]["expiration_time"] == 123 - - def test_workflow_blocking_pipeline_pause_payload_contract(self) -> None: - from core.app.apps.workflow.generate_task_pipeline import WorkflowAppGenerateTaskPipeline - - app_config = WorkflowUIBasedAppConfig( - tenant_id="tenant", - app_id="app", - app_mode=AppMode.WORKFLOW, - additional_features=AppAdditionalFeatures(), - variables=[], - workflow_id="workflow-id", - ) - application_generate_entity = WorkflowAppGenerateEntity.model_construct( - task_id="task", - app_config=app_config, - inputs={}, - files=[], - user_id="user", - stream=False, - invoke_from=InvokeFrom.WEB_APP, - trace_manager=None, - workflow_execution_id="run-id", - extras={}, - call_depth=0, - ) - pipeline = WorkflowAppGenerateTaskPipeline( - application_generate_entity=application_generate_entity, - workflow=_workflow(workflow_id="workflow-id", app_id="app", tenant_id="tenant"), - queue_manager=SimpleNamespace(invoke_from=InvokeFrom.WEB_APP, graph_runtime_state=None), - user=_end_user(user_id="user", app_id="app", tenant_id="tenant"), - stream=False, - draft_var_saver_factory=lambda **kwargs: None, - ) - - def _gen(): - yield HumanInputRequiredResponse( - task_id="task", - workflow_run_id="run", - data=HumanInputRequiredResponse.Data( - form_id="form-1", - node_id="node-1", - node_title="Human Input", - form_content="content", - expiration_time=1, - ), - ) - yield WorkflowPauseStreamResponse( - task_id="task", - workflow_run_id="run", - data=WorkflowPauseStreamResponse.Data( - workflow_run_id="run", - status=WorkflowExecutionStatus.PAUSED, - outputs={}, - paused_nodes=["node-1"], - reasons=[{"TYPE": "human_input_required", "form_id": "form-1", "expiration_time": 1}], - created_at=1, - elapsed_time=0.1, - total_tokens=0, - total_steps=0, - ), - ) - - response = pipeline._to_blocking_response(_gen()) - - assert isinstance(response, WorkflowAppPausedBlockingResponse) - assert response.data.status == WorkflowExecutionStatus.PAUSED - assert response.data.paused_nodes == ["node-1"] - assert response.data.reasons == [{"TYPE": "human_input_required", "form_id": "form-1", "expiration_time": 1}] - - def test_service_api_pause_event_serializes_hitl_reason( - self, - monkeypatch: pytest.MonkeyPatch, - sqlite_engine: Engine, - sqlite_session: Session, - ) -> None: - converter = _build_service_api_pause_converter() - converter.workflow_start_to_stream_response( - task_id="task", - workflow_run_id="run-id", - workflow_id="workflow-id", - reason=WorkflowStartReason.INITIAL, - ) - - expiration_time = datetime(2024, 1, 1, tzinfo=UTC) - _persist_human_input_form(sqlite_session, expiration_time=expiration_time) - - monkeypatch.setattr(workflow_response_converter, "db", SimpleNamespace(engine=sqlite_engine)) - monkeypatch.setattr( - workflow_response_converter, - "load_form_dispositions_by_form_id", - lambda form_ids, session=None, surface=None: { - "form-1": FormDisposition(form_token="token", approval_channels=[]) - }, - ) - - reason = HumanInputRequired( - form_id="form-1", - form_content="Rendered", - inputs=[ - ParagraphInputConfig( - type=FormInputType.PARAGRAPH, - output_variable_name="field", - default=None, - ), - ], - actions=[UserActionConfig(id="approve", title="Approve")], - display_in_ui=True, - node_id="node-id", - node_title="Human Step", - form_token="token", - ) - queue_event = QueueWorkflowPausedEvent( - reasons=[reason], - outputs={"answer": "value"}, - paused_nodes=["node-id"], - ) - - runtime_state = SimpleNamespace(total_tokens=0, node_run_steps=0, variable_pool=VariablePool()) - responses = converter.workflow_pause_to_stream_response( - event=queue_event, - task_id="task", - graph_runtime_state=runtime_state, - ) - - assert isinstance(responses[-1], WorkflowPauseStreamResponse) - pause_resp = responses[-1] - assert pause_resp.workflow_run_id == "run-id" - assert pause_resp.data.paused_nodes == ["node-id"] - assert pause_resp.data.outputs == {} - assert pause_resp.data.reasons[0]["TYPE"] == "human_input_required" - assert pause_resp.data.reasons[0]["form_id"] == "form-1" - assert pause_resp.data.reasons[0]["form_token"] == "token" - assert pause_resp.data.reasons[0]["expiration_time"] == to_utc_timestamp(expiration_time) - - assert isinstance(responses[0], HumanInputRequiredResponse) - hi_resp = responses[0] - assert hi_resp.data.form_id == "form-1" - assert hi_resp.data.node_id == "node-id" - assert hi_resp.data.node_title == "Human Step" - assert hi_resp.data.inputs[0].output_variable_name == "field" - assert hi_resp.data.actions[0].id == "approve" - assert hi_resp.data.display_in_ui is True - assert hi_resp.data.form_token == "token" - assert hi_resp.data.expiration_time == to_utc_timestamp(expiration_time) - - # Snapshot payload contract - def test_snapshot_events_include_pause_payload_contract( - self, - monkeypatch: pytest.MonkeyPatch, - sqlite_engine: Engine, - sqlite_session: Session, - ) -> None: - workflow_run = _build_workflow_run(WorkflowExecutionStatus.PAUSED) - snapshot = _build_snapshot(WorkflowNodeExecutionStatus.PAUSED) - resumption_context = _build_resumption_context("task-ctx") - expiration_time = datetime(2024, 1, 1, tzinfo=UTC) - _persist_human_input_form(sqlite_session, expiration_time=expiration_time) - monkeypatch.setattr( - "services.workflow_event_snapshot_service.load_form_dispositions_by_form_id", - lambda form_ids, session=None, surface=None: { - "form-1": FormDisposition(form_token="wtok", approval_channels=[]) - }, - ) - - sqlite_session_maker = sessionmaker(bind=sqlite_engine, expire_on_commit=False) - - pause_entity = _FakePauseEntity( - pause_id="pause-1", - workflow_run_id="run-1", - paused_at_value=datetime(2024, 1, 1, tzinfo=UTC), - pause_reasons=[ - HumanInputRequired( - form_id="form-1", - form_content="content", - node_id="node-1", - node_title="Human Input", - form_token="wtok", - ) - ], - ) - - events = _build_snapshot_events( - workflow_run=workflow_run, - node_snapshots=[snapshot], - task_id="task-ctx", - message_context=None, - pause_entity=pause_entity, - resumption_context=resumption_context, - session_maker=sqlite_session_maker, - node_executions={}, - ) - - assert [event["event"] for event in events] == [ - "workflow_started", - "node_started", - "node_finished", - "human_input_required", - "workflow_paused", - ] - assert events[2]["data"]["status"] == WorkflowNodeExecutionStatus.PAUSED.value - assert events[3]["data"]["form_token"] == "wtok" - assert events[3]["data"]["expiration_time"] == to_utc_timestamp(expiration_time) - pause_data = events[-1]["data"] - assert pause_data["paused_nodes"] == ["node-1"] - assert pause_data["outputs"] == {"result": "value"} - assert pause_data["reasons"][0]["TYPE"] == "human_input_required" - assert pause_data["reasons"][0]["form_token"] == "wtok" - assert pause_data["reasons"][0]["expiration_time"] == to_utc_timestamp(expiration_time) - assert pause_data["status"] == WorkflowExecutionStatus.PAUSED.value - assert pause_data["created_at"] == int(workflow_run.created_at.timestamp()) - assert pause_data["elapsed_time"] == workflow_run.elapsed_time - assert pause_data["total_tokens"] == workflow_run.total_tokens - assert pause_data["total_steps"] == workflow_run.total_steps diff --git a/api/tests/unit_tests/controllers/service_api/app/test_workflow.py b/api/tests/unit_tests/controllers/service_api/app/test_workflow.py index 4a0b92f2867..08383852191 100644 --- a/api/tests/unit_tests/controllers/service_api/app/test_workflow.py +++ b/api/tests/unit_tests/controllers/service_api/app/test_workflow.py @@ -1,18 +1,3 @@ -""" -Unit tests for Service API Workflow controllers. - -Tests coverage for: -- WorkflowRunPayload and WorkflowLogQuery Pydantic models -- Workflow execution error handling -- App mode validation for workflow endpoints -- Workflow stop mechanism validation - -Focus on: -- Pydantic model validation -- Error type mappings -- Service method interfaces -""" - import json import sys import uuid @@ -48,6 +33,7 @@ from controllers.web.error import InvokeRateLimitError as InvokeRateLimitHttpErr from enums import CloudPlan, DeploymentEdition from graphon.enums import WorkflowExecutionStatus from graphon.model_runtime.errors.invoke import InvokeRateLimitError as ProviderInvokeRateLimitError +from libs.helper import dump_response from models import Account from models.enums import CreatorUserRole, WorkflowRunTriggeredFrom from models.model import App, AppMode, EndUser @@ -65,6 +51,21 @@ from services.errors.app import ( from services.errors.llm import InvokeRateLimitError from services.workflow_app_log_query_service import WorkflowAppLogQueryService +""" +Unit tests for Service API Workflow controllers. + +Tests coverage for: +- WorkflowRunPayload and WorkflowLogQuery Pydantic models +- Workflow execution error handling +- App mode validation for workflow endpoints +- Workflow stop mechanism validation + +Focus on: +- Pydantic model validation +- Error type mappings +- Service method interfaces +""" + def _default_workflow_inputs() -> dict[str, object]: return {"input": "value"} @@ -233,367 +234,416 @@ def _stub_workflow_app_logs(monkeypatch: pytest.MonkeyPatch) -> MagicMock: return workflow_app_logs -class TestWorkflowRunPayload: - """Test suite for WorkflowRunPayload Pydantic model.""" +class TestWorkflowCases: + class TestWorkflowRunPayload: + """Test suite for WorkflowRunPayload Pydantic model.""" - def test_payload_with_required_inputs(self): - """Test payload with required inputs field.""" - payload = WorkflowRunPayload(inputs={"key": "value"}) - assert payload.inputs == {"key": "value"} - assert payload.files is None - assert payload.response_mode is None + def test_payload_with_required_inputs(self): + """Test payload with required inputs field.""" + payload = WorkflowRunPayload(inputs={"key": "value"}) + assert payload.inputs == {"key": "value"} + assert payload.files is None + assert payload.response_mode is None - def test_payload_with_all_fields(self): - """Test payload with all fields populated.""" - files = [{"type": "image", "url": "http://example.com/img.png"}] - payload = WorkflowRunPayload(inputs={"param1": "value1", "param2": 123}, files=files, response_mode="streaming") - assert payload.inputs == {"param1": "value1", "param2": 123} - assert payload.files == files - assert payload.response_mode == "streaming" + def test_payload_with_all_fields(self): + """Test payload with all fields populated.""" + files = [{"type": "image", "url": "http://example.com/img.png"}] + payload = WorkflowRunPayload( + inputs={"param1": "value1", "param2": 123}, files=files, response_mode="streaming" + ) + assert payload.inputs == {"param1": "value1", "param2": 123} + assert payload.files == files + assert payload.response_mode == "streaming" - def test_payload_response_mode_blocking(self): - """Test payload with blocking response mode.""" - payload = WorkflowRunPayload(inputs={}, response_mode="blocking") - assert payload.response_mode == "blocking" + def test_payload_response_mode_blocking(self): + """Test payload with blocking response mode.""" + payload = WorkflowRunPayload(inputs={}, response_mode="blocking") + assert payload.response_mode == "blocking" - def test_payload_with_complex_inputs(self): - """Test payload with nested complex inputs.""" - complex_inputs = { - "config": {"nested": {"value": 123}}, - "items": ["item1", "item2"], - "metadata": {"key": "value"}, - } - payload = WorkflowRunPayload(inputs=complex_inputs) - assert payload.inputs == complex_inputs + def test_payload_with_complex_inputs(self): + """Test payload with nested complex inputs.""" + complex_inputs = { + "config": {"nested": {"value": 123}}, + "items": ["item1", "item2"], + "metadata": {"key": "value"}, + } + payload = WorkflowRunPayload(inputs=complex_inputs) + assert payload.inputs == complex_inputs - def test_payload_with_empty_inputs(self): - """Test payload with empty inputs dict.""" - payload = WorkflowRunPayload(inputs={}) - assert payload.inputs == {} + def test_payload_with_empty_inputs(self): + """Test payload with empty inputs dict.""" + payload = WorkflowRunPayload(inputs={}) + assert payload.inputs == {} - def test_payload_with_multiple_files(self): - """Test payload with multiple file attachments.""" - files = [ - {"type": "image", "url": "http://example.com/img1.png"}, - {"type": "document", "upload_file_id": "file_123"}, - {"type": "audio", "url": "http://example.com/audio.mp3"}, - ] - payload = WorkflowRunPayload(inputs={}, files=files) - assert payload.files is not None - assert len(payload.files) == 3 + def test_payload_with_multiple_files(self): + """Test payload with multiple file attachments.""" + files = [ + {"type": "image", "url": "http://example.com/img1.png"}, + {"type": "document", "upload_file_id": "file_123"}, + {"type": "audio", "url": "http://example.com/audio.mp3"}, + ] + payload = WorkflowRunPayload(inputs={}, files=files) + assert payload.files is not None + assert len(payload.files) == 3 + class TestWorkflowLogQuery: + """Test suite for WorkflowLogQuery Pydantic model.""" -class TestWorkflowLogQuery: - """Test suite for WorkflowLogQuery Pydantic model.""" + def test_query_with_defaults(self): + """Test query with default values.""" + query = WorkflowLogQuery() + assert query.keyword is None + assert query.status is None + assert query.created_at__before is None + assert query.created_at__after is None + assert query.created_by_end_user_session_id is None + assert query.created_by_account is None + assert query.page == 1 + assert query.limit == 20 - def test_query_with_defaults(self): - """Test query with default values.""" - query = WorkflowLogQuery() - assert query.keyword is None - assert query.status is None - assert query.created_at__before is None - assert query.created_at__after is None - assert query.created_by_end_user_session_id is None - assert query.created_by_account is None - assert query.page == 1 - assert query.limit == 20 + def test_query_with_all_filters(self): + """Test query with all filter fields populated.""" + query = WorkflowLogQuery( + keyword="search term", + status="succeeded", + created_at__before="2024-01-15T10:00:00Z", + created_at__after="2024-01-01T00:00:00Z", + created_by_end_user_session_id="session_123", + created_by_account="user@example.com", + page=2, + limit=50, + ) + assert query.keyword == "search term" + assert query.status == "succeeded" + assert query.created_at__before == "2024-01-15T10:00:00Z" + assert query.created_at__after == "2024-01-01T00:00:00Z" + assert query.created_by_end_user_session_id == "session_123" + assert query.created_by_account == "user@example.com" + assert query.page == 2 + assert query.limit == 50 - def test_query_with_all_filters(self): - """Test query with all filter fields populated.""" - query = WorkflowLogQuery( - keyword="search term", - status="succeeded", - created_at__before="2024-01-15T10:00:00Z", - created_at__after="2024-01-01T00:00:00Z", - created_by_end_user_session_id="session_123", - created_by_account="user@example.com", - page=2, - limit=50, - ) - assert query.keyword == "search term" - assert query.status == "succeeded" - assert query.created_at__before == "2024-01-15T10:00:00Z" - assert query.created_at__after == "2024-01-01T00:00:00Z" - assert query.created_by_end_user_session_id == "session_123" - assert query.created_by_account == "user@example.com" - assert query.page == 2 - assert query.limit == 50 + @pytest.mark.parametrize("status", ["succeeded", "failed", "stopped"]) + def test_query_valid_status_values(self, status): + """Test all valid status values.""" + query = WorkflowLogQuery(status=status) + assert query.status == status - @pytest.mark.parametrize("status", ["succeeded", "failed", "stopped"]) - def test_query_valid_status_values(self, status): - """Test all valid status values.""" - query = WorkflowLogQuery(status=status) - assert query.status == status + def test_query_pagination_limits(self): + """Test query pagination boundaries.""" + query_min_page = WorkflowLogQuery(page=1) + assert query_min_page.page == 1 - def test_query_pagination_limits(self): - """Test query pagination boundaries.""" - query_min_page = WorkflowLogQuery(page=1) - assert query_min_page.page == 1 + query_max_page = WorkflowLogQuery(page=99999) + assert query_max_page.page == 99999 - query_max_page = WorkflowLogQuery(page=99999) - assert query_max_page.page == 99999 + query_min_limit = WorkflowLogQuery(limit=1) + assert query_min_limit.limit == 1 - query_min_limit = WorkflowLogQuery(limit=1) - assert query_min_limit.limit == 1 + query_max_limit = WorkflowLogQuery(limit=100) + assert query_max_limit.limit == 100 - query_max_limit = WorkflowLogQuery(limit=100) - assert query_max_limit.limit == 100 + def test_query_with_keyword_search(self): + """Test query with keyword filter.""" + query = WorkflowLogQuery(keyword="workflow execution") + assert query.keyword == "workflow execution" - def test_query_with_keyword_search(self): - """Test query with keyword filter.""" - query = WorkflowLogQuery(keyword="workflow execution") - assert query.keyword == "workflow execution" + def test_query_with_date_filters(self): + """Test query with before/after date filters.""" + query = WorkflowLogQuery( + created_at__before="2024-12-31T23:59:59Z", created_at__after="2024-01-01T00:00:00Z" + ) + assert query.created_at__before == "2024-12-31T23:59:59Z" + assert query.created_at__after == "2024-01-01T00:00:00Z" - def test_query_with_date_filters(self): - """Test query with before/after date filters.""" - query = WorkflowLogQuery(created_at__before="2024-12-31T23:59:59Z", created_at__after="2024-01-01T00:00:00Z") - assert query.created_at__before == "2024-12-31T23:59:59Z" - assert query.created_at__after == "2024-01-01T00:00:00Z" + class TestWorkflowRunResponse: + def test_validates_workflow_run_object_shape_and_clears_paused_outputs(self): + run = _make_workflow_run(run_id="run-paused") + run.status = WorkflowExecutionStatus.PAUSED + run.outputs = json.dumps({"should": "not leak"}) + result = WorkflowRunResponse.model_validate(run, from_attributes=True).model_dump(mode="json") -class TestWorkflowRunResponse: - def test_validates_workflow_run_object_shape_and_clears_paused_outputs(self): - run = _make_workflow_run(run_id="run-paused") - run.status = WorkflowExecutionStatus.PAUSED - run.outputs = json.dumps({"should": "not leak"}) + assert result == { + "id": "run-paused", + "workflow_id": "wf-1", + "status": "paused", + "inputs": '{"input": "value"}', + "outputs": {}, + "error": None, + "total_steps": 1, + "total_tokens": 10, + "created_at": 1767225600, + "finished_at": 1767225600, + "elapsed_time": 0.1, + } - result = WorkflowRunResponse.model_validate(run, from_attributes=True).model_dump(mode="json") + class TestWorkflowRunRepository: + """Test workflow run repository interface.""" - assert result == { - "id": "run-paused", - "workflow_id": "wf-1", - "status": "paused", - "inputs": '{"input": "value"}', - "outputs": {}, - "error": None, - "total_steps": 1, - "total_tokens": 10, - "created_at": 1767225600, - "finished_at": 1767225600, - "elapsed_time": 0.1, - } + @pytest.mark.parametrize("sqlite_session", [(WorkflowRun,)], indirect=True) + def test_workflow_run_repository_get_by_id(self, sqlite_engine: Engine, sqlite_session: Session): + """Test repository lookup against committed tenant-scoped state.""" + run = _make_workflow_run(run_id=str(uuid.uuid4())) + sqlite_session.add(run) + sqlite_session.commit() + from repositories.factory import DifyAPIRepositoryFactory - -class TestWorkflowRunRepository: - """Test workflow run repository interface.""" - - @pytest.mark.parametrize("sqlite_session", [(WorkflowRun,)], indirect=True) - def test_workflow_run_repository_get_by_id(self, sqlite_engine: Engine, sqlite_session: Session): - """Test repository lookup against committed tenant-scoped state.""" - run = _make_workflow_run(run_id=str(uuid.uuid4())) - sqlite_session.add(run) - sqlite_session.commit() - from repositories.factory import DifyAPIRepositoryFactory - - repo = DifyAPIRepositoryFactory.create_api_workflow_run_repository( - sessionmaker(bind=sqlite_engine, expire_on_commit=False) - ) - - result = repo.get_workflow_run_by_id(tenant_id="tenant-1", app_id="app-1", run_id=run.id) - - assert result is not None - assert result.id == run.id - assert repo.get_workflow_run_by_id(tenant_id="other-tenant", app_id="app-1", run_id=run.id) is None - - -class TestWorkflowRunDetailApi: - def test_not_workflow_app(self, app: Flask) -> None: - api = WorkflowRunDetailApi() - handler = unwrap(api.get) - app_model = _make_app_model(mode=AppMode.CHAT) - - with app.test_request_context("/workflows/run/1", method="GET"): - with pytest.raises(NotWorkflowAppError): - handler(api, app_model=app_model, workflow_run_id="run") - - @pytest.mark.parametrize("sqlite_session", [(WorkflowRun,)], indirect=True) - def test_success( - self, - monkeypatch: pytest.MonkeyPatch, - sqlite_engine: Engine, - sqlite_session: Session, - ) -> None: - run = _make_workflow_run(run_id="run", tenant_id="t1", app_id="a1") - sqlite_session.add(run) - sqlite_session.commit() - _bind_sqlite_database(monkeypatch, sqlite_engine, sqlite_session) - - api = WorkflowRunDetailApi() - handler = unwrap(api.get) - app_model = _make_app_model(app_id="a1", tenant_id="t1") - - result = handler(api, app_model=app_model, workflow_run_id="run") - assert result["id"] == "run" - assert result["workflow_id"] == "wf-1" - assert result["status"] == "succeeded" - - -class TestWorkflowRunApi: - @pytest.mark.parametrize("sqlite_session", [()], indirect=True) - def test_not_workflow_app(self, app: Flask, sqlite_session: Session) -> None: - api = WorkflowRunApi() - handler = unwrap(api.post) - app_model = _make_app_model(mode=AppMode.CHAT) - end_user = _make_end_user() - - with app.test_request_context("/workflows/run", method="POST", json={"inputs": {}}): - with pytest.raises(NotWorkflowAppError): - handler(api, session=sqlite_session, app_model=app_model, end_user=end_user) - - @pytest.mark.parametrize("sqlite_session", [()], indirect=True) - @pytest.mark.parametrize("api_class", [WorkflowRunApi, WorkflowRunByIdApi]) - @pytest.mark.parametrize( - ("source_error", "http_error", "status_code", "error_code"), - [ - pytest.param(InvokeRateLimitError, InvokeRateLimitHttpError, 429, "rate_limit_error", id="cloud-quota"), - pytest.param( - ProviderInvokeRateLimitError, CompletionRequestError, 400, "completion_request_error", id="provider" - ), - ], - ) - def test_maps_rate_limits_by_source( - self, - app: Flask, - monkeypatch: pytest.MonkeyPatch, - sqlite_session: Session, - config_overrides: Callable[..., None], - api_class: type[WorkflowRunApi | WorkflowRunByIdApi], - source_error: type[InvokeRateLimitError | ProviderInvokeRateLimitError], - http_error: type[InvokeRateLimitHttpError | CompletionRequestError], - status_code: int, - error_code: str, - ) -> None: - config_overrides(DEPLOYMENT_EDITION=DeploymentEdition.CLOUD) - monkeypatch.setattr( - BillingService, "get_info", Mock(return_value={"subscription": {"plan": CloudPlan.PROFESSIONAL}}) - ) - generate = Mock(side_effect=source_error("rate limit reached")) - monkeypatch.setattr(AppGenerateService, "generate", generate) - - api = api_class() - handler = unwrap(api.post) - app_model = _make_app_model() - end_user = _make_end_user() - kwargs: dict[str, str] = {"workflow_id": str(uuid.uuid4())} if isinstance(api, WorkflowRunByIdApi) else {} - - with app.test_request_context( - "/workflows/run", method="POST", json={"inputs": {}, "response_mode": "blocking"} - ): - with pytest.raises(http_error) as exc_info: - handler(api, session=sqlite_session, app_model=app_model, end_user=end_user, **kwargs) - - generate.assert_called_once() - assert exc_info.value.code == status_code - assert exc_info.value.error_code == error_code - assert exc_info.value.description == "rate limit reached" - - def test_trigger_workflow_returns_stable_unavailable_error( - self, - app: Flask, - monkeypatch: pytest.MonkeyPatch, - sqlite_session: Session, - ) -> None: - monkeypatch.setattr( - AppGenerateService, - "generate", - Mock(side_effect=TriggerWorkflowServiceModeUnavailableServiceError()), - ) - api = WorkflowRunApi() - handler = unwrap(api.post) - - with app.test_request_context("/workflows/run", method="POST", json={"inputs": {}}): - with pytest.raises(TriggerWorkflowServiceModeUnavailableError) as exc_info: - handler( - api, - session=sqlite_session, - app_model=_make_app_model(), - end_user=_make_end_user(), - ) - - assert exc_info.value.code == 403 - assert exc_info.value.error_code == "trigger_workflow_service_mode_unavailable" - - def test_sandbox_billing_does_not_gate_default_workflow_run( - self, - app: Flask, - monkeypatch: pytest.MonkeyPatch, - sqlite_session: Session, - config_overrides: Callable[..., None], - ) -> None: - config_overrides(DEPLOYMENT_EDITION=DeploymentEdition.CLOUD) - sys.modules["controllers.service_api.app.workflow"] - - billing_get_info = Mock(return_value={"subscription": {"plan": CloudPlan.SANDBOX}}) - generate = Mock(return_value={"result": "ok"}) - monkeypatch.setattr(BillingService, "get_info", billing_get_info) - monkeypatch.setattr(AppGenerateService, "generate", generate) - - api = WorkflowRunApi() - handler = unwrap(api.post) - - with app.test_request_context("/workflows/run", method="POST", json={"inputs": {}}): - response = handler( - api, - session=sqlite_session, - app_model=_make_app_model(), - end_user=_make_end_user(), + repo = DifyAPIRepositoryFactory.create_api_workflow_run_repository( + sessionmaker(bind=sqlite_engine, expire_on_commit=False) ) - assert response.get_json() == {"result": "ok"} - billing_get_info.assert_not_called() - generate.assert_called_once() + result = repo.get_workflow_run_by_id(tenant_id="tenant-1", app_id="app-1", run_id=run.id) + assert result is not None + assert result.id == run.id + assert repo.get_workflow_run_by_id(tenant_id="other-tenant", app_id="app-1", run_id=run.id) is None -class TestWorkflowRunByIdApi: - def test_trigger_workflow_version_returns_stable_unavailable_error( - self, - app: Flask, - monkeypatch: pytest.MonkeyPatch, - sqlite_session: Session, - ) -> None: - monkeypatch.setattr( - AppGenerateService, - "generate", - Mock(side_effect=TriggerWorkflowServiceModeUnavailableServiceError()), + class TestWorkflowRunDetailApi: + def test_not_workflow_app(self, app: Flask) -> None: + api = WorkflowRunDetailApi() + handler = unwrap(api.get) + app_model = _make_app_model(mode=AppMode.CHAT) + + with app.test_request_context("/workflows/run/1", method="GET"): + with pytest.raises(NotWorkflowAppError): + handler(api, app_model=app_model, workflow_run_id="run") + + @pytest.mark.parametrize("sqlite_session", [(WorkflowRun,)], indirect=True) + def test_success( + self, + monkeypatch: pytest.MonkeyPatch, + sqlite_engine: Engine, + sqlite_session: Session, + ) -> None: + run = _make_workflow_run(run_id="run", tenant_id="t1", app_id="a1") + sqlite_session.add(run) + sqlite_session.commit() + _bind_sqlite_database(monkeypatch, sqlite_engine, sqlite_session) + + api = WorkflowRunDetailApi() + handler = unwrap(api.get) + app_model = _make_app_model(app_id="a1", tenant_id="t1") + + result = handler(api, app_model=app_model, workflow_run_id="run") + assert result["id"] == "run" + assert result["workflow_id"] == "wf-1" + assert result["status"] == "succeeded" + + class TestWorkflowRunApi: + @pytest.mark.parametrize("sqlite_session", [()], indirect=True) + def test_not_workflow_app(self, app: Flask, sqlite_session: Session) -> None: + api = WorkflowRunApi() + handler = unwrap(api.post) + app_model = _make_app_model(mode=AppMode.CHAT) + end_user = _make_end_user() + + with app.test_request_context("/workflows/run", method="POST", json={"inputs": {}}): + with pytest.raises(NotWorkflowAppError): + handler(api, session=sqlite_session, app_model=app_model, end_user=end_user) + + @pytest.mark.parametrize("sqlite_session", [()], indirect=True) + @pytest.mark.parametrize("api_class", [WorkflowRunApi, WorkflowRunByIdApi]) + @pytest.mark.parametrize( + ("source_error", "http_error", "status_code", "error_code"), + [ + pytest.param(InvokeRateLimitError, InvokeRateLimitHttpError, 429, "rate_limit_error", id="cloud-quota"), + pytest.param( + ProviderInvokeRateLimitError, CompletionRequestError, 400, "completion_request_error", id="provider" + ), + ], ) - api = WorkflowRunByIdApi() - handler = unwrap(api.post) + def test_maps_rate_limits_by_source( + self, + app: Flask, + monkeypatch: pytest.MonkeyPatch, + sqlite_session: Session, + config_overrides: Callable[..., None], + api_class: type[WorkflowRunApi | WorkflowRunByIdApi], + source_error: type[InvokeRateLimitError | ProviderInvokeRateLimitError], + http_error: type[InvokeRateLimitHttpError | CompletionRequestError], + status_code: int, + error_code: str, + ) -> None: + config_overrides(DEPLOYMENT_EDITION=DeploymentEdition.CLOUD) + monkeypatch.setattr( + BillingService, "get_info", Mock(return_value={"subscription": {"plan": CloudPlan.PROFESSIONAL}}) + ) + generate = Mock(side_effect=source_error("rate limit reached")) + monkeypatch.setattr(AppGenerateService, "generate", generate) - with app.test_request_context("/workflows/w1/run", method="POST", json={"inputs": {}}): - with pytest.raises(TriggerWorkflowServiceModeUnavailableError) as exc_info: - handler( + api = api_class() + handler = unwrap(api.post) + app_model = _make_app_model() + end_user = _make_end_user() + kwargs: dict[str, str] = {"workflow_id": str(uuid.uuid4())} if isinstance(api, WorkflowRunByIdApi) else {} + + with app.test_request_context( + "/workflows/run", method="POST", json={"inputs": {}, "response_mode": "blocking"} + ): + with pytest.raises(http_error) as exc_info: + handler(api, session=sqlite_session, app_model=app_model, end_user=end_user, **kwargs) + + generate.assert_called_once() + assert exc_info.value.code == status_code + assert exc_info.value.error_code == error_code + assert exc_info.value.description == "rate limit reached" + + def test_trigger_workflow_returns_stable_unavailable_error( + self, + app: Flask, + monkeypatch: pytest.MonkeyPatch, + sqlite_session: Session, + ) -> None: + monkeypatch.setattr( + AppGenerateService, + "generate", + Mock(side_effect=TriggerWorkflowServiceModeUnavailableServiceError()), + ) + api = WorkflowRunApi() + handler = unwrap(api.post) + + with app.test_request_context("/workflows/run", method="POST", json={"inputs": {}}): + with pytest.raises(TriggerWorkflowServiceModeUnavailableError) as exc_info: + handler( + api, + session=sqlite_session, + app_model=_make_app_model(), + end_user=_make_end_user(), + ) + + assert exc_info.value.code == 403 + assert exc_info.value.error_code == "trigger_workflow_service_mode_unavailable" + + def test_sandbox_billing_does_not_gate_default_workflow_run( + self, + app: Flask, + monkeypatch: pytest.MonkeyPatch, + sqlite_session: Session, + config_overrides: Callable[..., None], + ) -> None: + config_overrides(DEPLOYMENT_EDITION=DeploymentEdition.CLOUD) + sys.modules["controllers.service_api.app.workflow"] + + billing_get_info = Mock(return_value={"subscription": {"plan": CloudPlan.SANDBOX}}) + generate = Mock(return_value={"result": "ok"}) + monkeypatch.setattr(BillingService, "get_info", billing_get_info) + monkeypatch.setattr(AppGenerateService, "generate", generate) + + api = WorkflowRunApi() + handler = unwrap(api.post) + + with app.test_request_context("/workflows/run", method="POST", json={"inputs": {}}): + response = handler( api, session=sqlite_session, app_model=_make_app_model(), end_user=_make_end_user(), - workflow_id=str(uuid.uuid4()), ) - assert exc_info.value.code == 403 - assert exc_info.value.error_code == "trigger_workflow_service_mode_unavailable" + assert response.get_json() == {"result": "ok"} + billing_get_info.assert_not_called() + generate.assert_called_once() - def test_rejects_sandbox_plan_with_upgrade_error( - self, - app: Flask, - monkeypatch: pytest.MonkeyPatch, - sqlite_session: Session, - config_overrides: Callable[..., None], - ) -> None: - config_overrides(DEPLOYMENT_EDITION=DeploymentEdition.CLOUD) - sys.modules["controllers.service_api.app.workflow"] + class TestWorkflowRunByIdApi: + def test_trigger_workflow_version_returns_stable_unavailable_error( + self, + app: Flask, + monkeypatch: pytest.MonkeyPatch, + sqlite_session: Session, + ) -> None: + monkeypatch.setattr( + AppGenerateService, + "generate", + Mock(side_effect=TriggerWorkflowServiceModeUnavailableServiceError()), + ) + api = WorkflowRunByIdApi() + handler = unwrap(api.post) - billing_get_info = Mock(return_value={"subscription": {"plan": CloudPlan.SANDBOX}}) - generate = Mock() - monkeypatch.setattr(BillingService, "get_info", billing_get_info) - monkeypatch.setattr(AppGenerateService, "generate", generate) + with app.test_request_context("/workflows/w1/run", method="POST", json={"inputs": {}}): + with pytest.raises(TriggerWorkflowServiceModeUnavailableError) as exc_info: + handler( + api, + session=sqlite_session, + app_model=_make_app_model(), + end_user=_make_end_user(), + workflow_id=str(uuid.uuid4()), + ) - api = WorkflowRunByIdApi() - handler = unwrap(api.post) - app_model = _make_app_model() + assert exc_info.value.code == 403 + assert exc_info.value.error_code == "trigger_workflow_service_mode_unavailable" - with app.test_request_context("/workflows/w1/run", method="POST", json={"inputs": {}}): - with pytest.raises(WorkflowVersionExecutionNotAllowedError) as exc_info: - handler( + def test_rejects_sandbox_plan_with_upgrade_error( + self, + app: Flask, + monkeypatch: pytest.MonkeyPatch, + sqlite_session: Session, + config_overrides: Callable[..., None], + ) -> None: + config_overrides(DEPLOYMENT_EDITION=DeploymentEdition.CLOUD) + sys.modules["controllers.service_api.app.workflow"] + + billing_get_info = Mock(return_value={"subscription": {"plan": CloudPlan.SANDBOX}}) + generate = Mock() + monkeypatch.setattr(BillingService, "get_info", billing_get_info) + monkeypatch.setattr(AppGenerateService, "generate", generate) + + api = WorkflowRunByIdApi() + handler = unwrap(api.post) + app_model = _make_app_model() + + with app.test_request_context("/workflows/w1/run", method="POST", json={"inputs": {}}): + with pytest.raises(WorkflowVersionExecutionNotAllowedError) as exc_info: + handler( + api, + session=sqlite_session, + app_model=app_model, + end_user=_make_end_user(), + workflow_id="w1", + ) + + billing_get_info.assert_called_once_with(app_model.tenant_id, exclude_vector_space=True) + generate.assert_not_called() + assert exc_info.value.code == 403 + assert exc_info.value.error_code == "workflow_version_execution_not_allowed" + assert exc_info.value.description == ( + "Workflow version execution is not available on your current plan. Please upgrade to a paid plan." + ) + assert exc_info.value.data == { + "code": "workflow_version_execution_not_allowed", + "message": exc_info.value.description, + "status": 403, + } + + @pytest.mark.parametrize( + ("deployment_edition", "plan"), + [ + (DeploymentEdition.COMMUNITY, CloudPlan.SANDBOX), + (DeploymentEdition.ENTERPRISE, CloudPlan.SANDBOX), + (DeploymentEdition.CLOUD, CloudPlan.PROFESSIONAL), + ], + ) + def test_allows_execution_outside_enabled_sandbox_plan( + self, + app: Flask, + monkeypatch: pytest.MonkeyPatch, + deployment_edition: DeploymentEdition, + plan: CloudPlan, + sqlite_session: Session, + config_overrides: Callable[..., None], + ) -> None: + config_overrides(DEPLOYMENT_EDITION=deployment_edition) + + billing_get_info = Mock(return_value={"subscription": {"plan": plan}}) + generate = Mock(return_value={"result": "ok"}) + monkeypatch.setattr(BillingService, "get_info", billing_get_info) + monkeypatch.setattr(AppGenerateService, "generate", generate) + + api = WorkflowRunByIdApi() + handler = unwrap(api.post) + app_model = _make_app_model() + + with app.test_request_context("/workflows/w1/run", method="POST", json={"inputs": {}}): + response = handler( api, session=sqlite_session, app_model=app_model, @@ -601,272 +651,251 @@ class TestWorkflowRunByIdApi: workflow_id="w1", ) - billing_get_info.assert_called_once_with(app_model.tenant_id, exclude_vector_space=True) - generate.assert_not_called() - assert exc_info.value.code == 403 - assert exc_info.value.error_code == "workflow_version_execution_not_allowed" - assert exc_info.value.description == ( - "Workflow version execution is not available on your current plan. Please upgrade to a paid plan." - ) - assert exc_info.value.data == { - "code": "workflow_version_execution_not_allowed", - "message": exc_info.value.description, - "status": 403, - } + assert response.get_json() == {"result": "ok"} + generate.assert_called_once() + if deployment_edition == DeploymentEdition.CLOUD: + billing_get_info.assert_called_once_with(app_model.tenant_id, exclude_vector_space=True) + else: + billing_get_info.assert_not_called() - @pytest.mark.parametrize( - ("deployment_edition", "plan"), - [ - (DeploymentEdition.COMMUNITY, CloudPlan.SANDBOX), - (DeploymentEdition.ENTERPRISE, CloudPlan.SANDBOX), - (DeploymentEdition.CLOUD, CloudPlan.PROFESSIONAL), - ], - ) - def test_allows_execution_outside_enabled_sandbox_plan( - self, - app: Flask, - monkeypatch: pytest.MonkeyPatch, - deployment_edition: DeploymentEdition, - plan: CloudPlan, - sqlite_session: Session, - config_overrides: Callable[..., None], - ) -> None: - config_overrides(DEPLOYMENT_EDITION=deployment_edition) - - billing_get_info = Mock(return_value={"subscription": {"plan": plan}}) - generate = Mock(return_value={"result": "ok"}) - monkeypatch.setattr(BillingService, "get_info", billing_get_info) - monkeypatch.setattr(AppGenerateService, "generate", generate) - - api = WorkflowRunByIdApi() - handler = unwrap(api.post) - app_model = _make_app_model() - - with app.test_request_context("/workflows/w1/run", method="POST", json={"inputs": {}}): - response = handler( - api, - session=sqlite_session, - app_model=app_model, - end_user=_make_end_user(), - workflow_id="w1", + @pytest.mark.parametrize("sqlite_session", [()], indirect=True) + def test_not_found( + self, + app: Flask, + monkeypatch: pytest.MonkeyPatch, + sqlite_session: Session, + config_overrides: Callable[..., None], + ) -> None: + config_overrides(DEPLOYMENT_EDITION=DeploymentEdition.COMMUNITY) + sys.modules["controllers.service_api.app.workflow"] + monkeypatch.setattr( + AppGenerateService, + "generate", + lambda *_args, **_kwargs: (_ for _ in ()).throw(WorkflowNotFoundError("missing")), ) - assert response.get_json() == {"result": "ok"} - generate.assert_called_once() - if deployment_edition == DeploymentEdition.CLOUD: - billing_get_info.assert_called_once_with(app_model.tenant_id, exclude_vector_space=True) - else: - billing_get_info.assert_not_called() + api = WorkflowRunByIdApi() + handler = unwrap(api.post) + app_model = _make_app_model() + end_user = _make_end_user() - @pytest.mark.parametrize("sqlite_session", [()], indirect=True) - def test_not_found( - self, - app: Flask, - monkeypatch: pytest.MonkeyPatch, - sqlite_session: Session, - config_overrides: Callable[..., None], - ) -> None: - config_overrides(DEPLOYMENT_EDITION=DeploymentEdition.COMMUNITY) - sys.modules["controllers.service_api.app.workflow"] - monkeypatch.setattr( - AppGenerateService, - "generate", - lambda *_args, **_kwargs: (_ for _ in ()).throw(WorkflowNotFoundError("missing")), - ) + with app.test_request_context("/workflows/1/run", method="POST", json={"inputs": {}}): + with pytest.raises(NotFound): + handler(api, session=sqlite_session, app_model=app_model, end_user=end_user, workflow_id="w1") - api = WorkflowRunByIdApi() - handler = unwrap(api.post) - app_model = _make_app_model() - end_user = _make_end_user() + @pytest.mark.parametrize("sqlite_session", [()], indirect=True) + def test_draft_workflow( + self, + app: Flask, + monkeypatch: pytest.MonkeyPatch, + sqlite_session: Session, + config_overrides: Callable[..., None], + ) -> None: + config_overrides(DEPLOYMENT_EDITION=DeploymentEdition.COMMUNITY) + sys.modules["controllers.service_api.app.workflow"] + monkeypatch.setattr( + AppGenerateService, + "generate", + lambda *_args, **_kwargs: (_ for _ in ()).throw(IsDraftWorkflowError("draft")), + ) - with app.test_request_context("/workflows/1/run", method="POST", json={"inputs": {}}): - with pytest.raises(NotFound): - handler(api, session=sqlite_session, app_model=app_model, end_user=end_user, workflow_id="w1") + api = WorkflowRunByIdApi() + handler = unwrap(api.post) + app_model = _make_app_model() + end_user = _make_end_user() - @pytest.mark.parametrize("sqlite_session", [()], indirect=True) - def test_draft_workflow( - self, - app: Flask, - monkeypatch: pytest.MonkeyPatch, - sqlite_session: Session, - config_overrides: Callable[..., None], - ) -> None: - config_overrides(DEPLOYMENT_EDITION=DeploymentEdition.COMMUNITY) - sys.modules["controllers.service_api.app.workflow"] - monkeypatch.setattr( - AppGenerateService, - "generate", - lambda *_args, **_kwargs: (_ for _ in ()).throw(IsDraftWorkflowError("draft")), - ) + with app.test_request_context("/workflows/1/run", method="POST", json={"inputs": {}}): + with pytest.raises(BadRequest): + handler(api, session=sqlite_session, app_model=app_model, end_user=end_user, workflow_id="w1") - api = WorkflowRunByIdApi() - handler = unwrap(api.post) - app_model = _make_app_model() - end_user = _make_end_user() + class TestWorkflowTaskStopApi: + @pytest.mark.parametrize("mode", [AppMode.CHAT, AppMode.COMPLETION]) + def test_wrong_mode(self, app: Flask, mode: AppMode) -> None: + api = WorkflowTaskStopApi() + handler = unwrap(api.post) + app_model = _make_app_model(mode=mode) + end_user = _make_end_user() - with app.test_request_context("/workflows/1/run", method="POST", json={"inputs": {}}): - with pytest.raises(BadRequest): - handler(api, session=sqlite_session, app_model=app_model, end_user=end_user, workflow_id="w1") + with app.test_request_context("/workflows/tasks/1/stop", method="POST"): + with pytest.raises(NotWorkflowAppError): + handler(api, app_model=app_model, end_user=end_user, task_id="t1") + def test_success(self, app: Flask, monkeypatch: pytest.MonkeyPatch) -> None: + stop_mock = Mock() + workflow_module = sys.modules["controllers.service_api.app.workflow"] + monkeypatch.setattr(workflow_module.AppTaskService, "stop_workflow_task", stop_mock) -class TestWorkflowTaskStopApi: - @pytest.mark.parametrize("mode", [AppMode.CHAT, AppMode.COMPLETION]) - def test_wrong_mode(self, app: Flask, mode: AppMode) -> None: - api = WorkflowTaskStopApi() - handler = unwrap(api.post) - app_model = _make_app_model(mode=mode) - end_user = _make_end_user() + api = WorkflowTaskStopApi() + handler = unwrap(api.post) + app_model = _make_app_model() + end_user = _make_end_user(user_id="u1") - with app.test_request_context("/workflows/tasks/1/stop", method="POST"): - with pytest.raises(NotWorkflowAppError): - handler(api, app_model=app_model, end_user=end_user, task_id="t1") + with app.test_request_context("/workflows/tasks/1/stop", method="POST"): + response = handler(api, app_model=app_model, end_user=end_user, task_id="t1") - def test_success(self, app: Flask, monkeypatch: pytest.MonkeyPatch) -> None: - stop_mock = Mock() - workflow_module = sys.modules["controllers.service_api.app.workflow"] - monkeypatch.setattr(workflow_module.AppTaskService, "stop_workflow_task", stop_mock) + assert response == {"result": "success"} + stop_mock.assert_called_once_with( + tenant_id=app_model.tenant_id, + app_id=app_model.id, + task_id="t1", + app_mode=AppMode.WORKFLOW, + owner=None, + ) - api = WorkflowTaskStopApi() - handler = unwrap(api.post) - app_model = _make_app_model() - end_user = _make_end_user(user_id="u1") - - with app.test_request_context("/workflows/tasks/1/stop", method="POST"): - response = handler(api, app_model=app_model, end_user=end_user, task_id="t1") - - assert response == {"result": "success"} - stop_mock.assert_called_once_with( - tenant_id=app_model.tenant_id, - app_id=app_model.id, - task_id="t1", - app_mode=AppMode.WORKFLOW, - owner=None, - ) - - -class TestWorkflowAppLogApi: - def test_success( - self, - app: Flask, - monkeypatch: pytest.MonkeyPatch, - sqlite_session: Session, - sqlite_session_factory: sessionmaker[Session], - ) -> None: - api = WorkflowAppLogApi() - handler = unwrap(api.get) - app_model = _make_app_model(app_id="a1") - _persist_workflow_log(sqlite_session, tenant_id=app_model.tenant_id, app_id=app_model.id) - - workflow_app_logs = WorkflowAppLogQueryService( - logs=WorkflowAppLogQueryRepository(session_factory=sqlite_session_factory), - ) - monkeypatch.setattr( - sys.modules["controllers.service_api.app.workflow"], - "application_services", - lambda: SimpleNamespace(workflow_app_logs=workflow_app_logs), - ) - - with app.test_request_context("/workflows/logs", method="GET"): - response = handler(api, app_model=app_model) - - assert response == _expected_workflow_log_pagination_payload() - - -# ============================================================================= -# API Endpoint Tests -# -# ``WorkflowRunDetailApi``, ``WorkflowTaskStopApi``, and -# ``WorkflowAppLogApi`` use ``@validate_app_token`` which preserves -# ``__wrapped__`` via ``functools.wraps``. We call the unwrapped method -# directly to bypass the decorator. -# ============================================================================= - - -@pytest.fixture -def workflow_app() -> App: - return _make_app_model(app_id=str(uuid.uuid4()), tenant_id=str(uuid.uuid4())) - - -class TestWorkflowRunDetailApiGet: - """Test suite for WorkflowRunDetailApi.get() endpoint. - - ``get`` is wrapped by ``@validate_app_token`` (preserves ``__wrapped__``), - and we call the unwrapped method directly in tests. - """ - - @pytest.mark.parametrize("sqlite_session", [(WorkflowRun,)], indirect=True) - def test_get_workflow_run_success( - self, - app: Flask, - workflow_app: App, - monkeypatch: pytest.MonkeyPatch, - sqlite_engine: Engine, - sqlite_session: Session, - ): - """Test successful workflow run detail retrieval.""" - run = _make_workflow_run( - run_id="run-1", - tenant_id=workflow_app.tenant_id, - app_id=workflow_app.id, - ) - sqlite_session.add(run) - sqlite_session.commit() - _bind_sqlite_database(monkeypatch, sqlite_engine, sqlite_session) - - from controllers.service_api.app.workflow import WorkflowRunDetailApi - - with app.test_request_context( - f"/workflows/run/{run.id}", - method="GET", - ): - api = WorkflowRunDetailApi() - result = unwrap(api.get)(api, app_model=workflow_app, workflow_run_id=run.id) - - assert result == { - "id": "run-1", - "workflow_id": "wf-1", - "status": "succeeded", - "inputs": '{"input": "value"}', - "outputs": {"output": "value"}, - "error": None, - "total_steps": 1, - "total_tokens": 10, - "created_at": int(datetime(2026, 1, 1).timestamp()), - "finished_at": int(datetime(2026, 1, 1).timestamp()), - "elapsed_time": 0.1, - } - - def test_get_workflow_run_wrong_app_mode(self, app: Flask): - """Test NotWorkflowAppError when app mode is not workflow or advanced_chat.""" - from controllers.service_api.app.workflow import WorkflowRunDetailApi - - app_model = _make_app_model(mode=AppMode.CHAT) - - with app.test_request_context("/workflows/run/run-1", method="GET"): - api = WorkflowRunDetailApi() - with pytest.raises(NotWorkflowAppError): - unwrap(api.get)(api, app_model=app_model, workflow_run_id="run-1") - - -class TestWorkflowAppLogApiGet: - """Test suite for WorkflowAppLogApi.get() endpoint. - - ``get`` is wrapped by ``@validate_app_token``. - """ - - def test_get_workflow_logs_success( - self, - app: Flask, - workflow_app: App, - monkeypatch: pytest.MonkeyPatch, - ): - """Test successful workflow log retrieval.""" - _stub_workflow_app_logs(monkeypatch) - - from controllers.service_api.app.workflow import WorkflowAppLogApi - - with app.test_request_context("/workflows/logs?page=1&limit=20", method="GET"): + class TestWorkflowAppLogApi: + def test_success( + self, + app: Flask, + monkeypatch: pytest.MonkeyPatch, + sqlite_session: Session, + sqlite_session_factory: sessionmaker[Session], + ) -> None: api = WorkflowAppLogApi() - result = unwrap(api.get)(api, app_model=workflow_app) + handler = unwrap(api.get) + app_model = _make_app_model(app_id="a1") + _persist_workflow_log(sqlite_session, tenant_id=app_model.tenant_id, app_id=app_model.id) - assert result == _expected_workflow_log_pagination_payload() + workflow_app_logs = WorkflowAppLogQueryService( + logs=WorkflowAppLogQueryRepository(session_factory=sqlite_session_factory), + ) + monkeypatch.setattr( + sys.modules["controllers.service_api.app.workflow"], + "application_services", + lambda: SimpleNamespace(workflow_app_logs=workflow_app_logs), + ) + + with app.test_request_context("/workflows/logs", method="GET"): + response = handler(api, app_model=app_model) + + assert response == _expected_workflow_log_pagination_payload() + + # ============================================================================= + # API Endpoint Tests + # + # ``WorkflowRunDetailApi``, ``WorkflowTaskStopApi``, and + # ``WorkflowAppLogApi`` use ``@validate_app_token`` which preserves + # ``__wrapped__`` via ``functools.wraps``. We call the unwrapped method + # directly to bypass the decorator. + # ============================================================================= + + @staticmethod + @pytest.fixture + def workflow_app() -> App: + return _make_app_model(app_id=str(uuid.uuid4()), tenant_id=str(uuid.uuid4())) + + class TestWorkflowRunDetailApiGet: + """Test suite for WorkflowRunDetailApi.get() endpoint. + + ``get`` is wrapped by ``@validate_app_token`` (preserves ``__wrapped__``), + and we call the unwrapped method directly in tests. + """ + + @pytest.mark.parametrize("sqlite_session", [(WorkflowRun,)], indirect=True) + def test_get_workflow_run_success( + self, + app: Flask, + workflow_app: App, + monkeypatch: pytest.MonkeyPatch, + sqlite_engine: Engine, + sqlite_session: Session, + ): + """Test successful workflow run detail retrieval.""" + run = _make_workflow_run( + run_id="run-1", + tenant_id=workflow_app.tenant_id, + app_id=workflow_app.id, + ) + sqlite_session.add(run) + sqlite_session.commit() + _bind_sqlite_database(monkeypatch, sqlite_engine, sqlite_session) + + from controllers.service_api.app.workflow import WorkflowRunDetailApi + + with app.test_request_context( + f"/workflows/run/{run.id}", + method="GET", + ): + api = WorkflowRunDetailApi() + result = unwrap(api.get)(api, app_model=workflow_app, workflow_run_id=run.id) + + assert result == { + "id": "run-1", + "workflow_id": "wf-1", + "status": "succeeded", + "inputs": '{"input": "value"}', + "outputs": {"output": "value"}, + "error": None, + "total_steps": 1, + "total_tokens": 10, + "created_at": int(datetime(2026, 1, 1).timestamp()), + "finished_at": int(datetime(2026, 1, 1).timestamp()), + "elapsed_time": 0.1, + } + + def test_get_workflow_run_wrong_app_mode(self, app: Flask): + """Test NotWorkflowAppError when app mode is not workflow or advanced_chat.""" + from controllers.service_api.app.workflow import WorkflowRunDetailApi + + app_model = _make_app_model(mode=AppMode.CHAT) + + with app.test_request_context("/workflows/run/run-1", method="GET"): + api = WorkflowRunDetailApi() + with pytest.raises(NotWorkflowAppError): + unwrap(api.get)(api, app_model=app_model, workflow_run_id="run-1") + + class TestWorkflowAppLogApiGet: + """Test suite for WorkflowAppLogApi.get() endpoint. + + ``get`` is wrapped by ``@validate_app_token``. + """ + + def test_get_workflow_logs_success( + self, + app: Flask, + workflow_app: App, + monkeypatch: pytest.MonkeyPatch, + ): + """Test successful workflow log retrieval.""" + _stub_workflow_app_logs(monkeypatch) + + from controllers.service_api.app.workflow import WorkflowAppLogApi + + with app.test_request_context("/workflows/logs?page=1&limit=20", method="GET"): + api = WorkflowAppLogApi() + result = unwrap(api.get)(api, app_model=workflow_app) + + assert result == _expected_workflow_log_pagination_payload() + + +def _workflow_run(status: WorkflowExecutionStatus, outputs: str | None = '{"foo": "bar"}') -> WorkflowRun: + return WorkflowRun( + id="run-id", + workflow_id="workflow-id", + status=status, + inputs="{}", + outputs=outputs, + error=None, + total_steps=1, + total_tokens=2, + elapsed_time=3.5, + ) + + +class TestWorkflowFieldsCases: + @staticmethod + def test_workflow_run_serializer_normalizes_status_enum() -> None: + response = dump_response(WorkflowRunResponse, _workflow_run(WorkflowExecutionStatus.PAUSED)) + + assert response["status"] == "paused" + + @staticmethod + def test_workflow_run_serializer_paused_returns_empty_outputs() -> None: + response = dump_response(WorkflowRunResponse, _workflow_run(WorkflowExecutionStatus.PAUSED)) + + assert response["outputs"] == {} + + @staticmethod + def test_workflow_run_serializer_running_returns_outputs() -> None: + response = dump_response(WorkflowRunResponse, _workflow_run(WorkflowExecutionStatus.RUNNING)) + + assert response["outputs"] == {"foo": "bar"} diff --git a/api/tests/unit_tests/controllers/service_api/app/test_workflow_events.py b/api/tests/unit_tests/controllers/service_api/app/test_workflow_events.py index d9ad0b758ff..d1a8f2277e2 100644 --- a/api/tests/unit_tests/controllers/service_api/app/test_workflow_events.py +++ b/api/tests/unit_tests/controllers/service_api/app/test_workflow_events.py @@ -1,28 +1,803 @@ -"""Unit tests for Service API workflow event stream endpoints.""" - from __future__ import annotations import json import sys +from collections.abc import Sequence +from dataclasses import dataclass from datetime import UTC, datetime from inspect import unwrap from types import SimpleNamespace -from unittest.mock import Mock +from typing import override +from unittest.mock import ANY, MagicMock, Mock import pytest from flask import Flask from sqlalchemy.engine import Engine +from sqlalchemy.orm import Session, sessionmaker from werkzeug.exceptions import NotFound +import services.app_generate_service as ags_module from controllers.service_api.app.error import NotWorkflowAppError from controllers.service_api.app.workflow_events import WorkflowEventsApi -from graphon.enums import WorkflowExecutionStatus -from models.enums import CreatorUserRole, EndUserType, WorkflowRunTriggeredFrom +from core.app.app_config.entities import AppAdditionalFeatures, WorkflowUIBasedAppConfig +from core.app.apps.advanced_chat.generate_task_pipeline import ConversationSnapshot, MessageSnapshot, WorkflowSnapshot +from core.app.apps.common import workflow_response_converter +from core.app.apps.common.workflow_response_converter import WorkflowResponseConverter +from core.app.entities.app_invoke_entities import AdvancedChatAppGenerateEntity, InvokeFrom, WorkflowAppGenerateEntity +from core.app.entities.queue_entities import QueueWorkflowPausedEvent +from core.app.entities.task_entities import ( + AdvancedChatPausedBlockingResponse, + HumanInputRequiredResponse, + WorkflowAppPausedBlockingResponse, + WorkflowPauseStreamResponse, +) +from core.app.layers.pause_state_persist_layer import WorkflowResumptionContext, _WorkflowGenerateEntityWrapper +from core.workflow.human_input_policy import FormDisposition, HumanInputSurface +from core.workflow.nodes.human_input.entities import ParagraphInputConfig, UserActionConfig +from core.workflow.nodes.human_input.enums import FormInputType, HumanInputFormKind, HumanInputFormStatus +from core.workflow.nodes.human_input.pause_reason import DifyHITLEventType, HumanInputRequired +from core.workflow.system_variables import build_system_variables +from enums import DeploymentEdition +from graphon.entities import WorkflowStartReason +from graphon.enums import WorkflowExecutionStatus, WorkflowNodeExecutionStatus +from graphon.runtime import RuntimeState, VariablePool +from libs.datetime_utils import to_utc_timestamp +from models.account import Account +from models.enums import CreatorUserRole, EndUserType, MessageStatus, WorkflowRunTriggeredFrom +from models.human_input import HumanInputForm from models.model import App, AppMode, EndUser -from models.workflow import WorkflowRun, WorkflowType +from models.workflow import Workflow, WorkflowRun, WorkflowType +from repositories.api_workflow_node_execution_repository import WorkflowNodeExecutionSnapshot +from repositories.entities.workflow_pause import WorkflowPauseEntity +from services.app_generate_service import AppGenerateService +from services.workflow_event_snapshot_service import _build_snapshot_events +from tests.helpers.unit.config_override import apply_config_overrides + +"""Dedicated tests for HITL behavior exposed through the Service API.""" -def _app(*, mode: AppMode = AppMode.WORKFLOW) -> App: +class _DummyRateLimit: + @staticmethod + def gen_request_key() -> str: + return "dummy-request-id" + + def __init__(self, client_id: str, max_active_requests: int) -> None: + self.client_id = client_id + self.max_active_requests = max_active_requests + + def enter(self, request_id: str | None = None) -> str: + return request_id or "dummy-request-id" + + def exit(self, request_id: str) -> None: + return None + + def generate(self, generator, request_id: str): + return generator + + +def _hitl_service_api_mock_repo_for_run(monkeypatch: pytest.MonkeyPatch, workflow_run, sqlite_engine: Engine): + workflow_events_module = sys.modules["controllers.service_api.app.workflow_events"] + repo = SimpleNamespace(get_workflow_run_by_id_and_tenant_id=lambda **_kwargs: workflow_run) + monkeypatch.setattr( + workflow_events_module.DifyAPIRepositoryFactory, + "create_api_workflow_run_repository", + lambda *_args, **_kwargs: repo, + ) + monkeypatch.setattr(workflow_events_module, "db", SimpleNamespace(engine=sqlite_engine)) + return workflow_events_module + + +def _hitl_service_api_app( + *, app_id: str = "app-1", tenant_id: str = "tenant-1", mode: AppMode = AppMode.WORKFLOW +) -> App: + return App( + id=app_id, + tenant_id=tenant_id, + name="Service API app", + description="", + mode=mode, + enable_site=True, + enable_api=True, + max_active_requests=0, + ) + + +def _hitl_service_api_end_user( + *, user_id: str = "end-user-1", app_id: str = "app-1", tenant_id: str = "tenant-1" +) -> EndUser: + return EndUser( + id=user_id, + tenant_id=tenant_id, + app_id=app_id, + type=EndUserType.SERVICE_API, + external_user_id="external-user-1", + name="Service API user", + session_id="session-1", + ) + + +def _workflow(*, workflow_id: str = "workflow-id", app_id: str = "app-id", tenant_id: str = "tenant-id") -> Workflow: + return Workflow( + id=workflow_id, + tenant_id=tenant_id, + app_id=app_id, + type=WorkflowType.WORKFLOW, + version="1", + graph=json.dumps({"nodes": [], "edges": []}), + features="{}", + created_by="owner-id", + environment_variables=[], + conversation_variables=[], + rag_pipeline_variables=[], + ) + + +def _persist_human_input_form( + sqlite_session: Session, + *, + expiration_time: datetime, +) -> HumanInputForm: + form = HumanInputForm( + id="form-1", + tenant_id="tenant-1", + app_id="app-1", + workflow_run_id="run-1", + conversation_id=None, + form_kind=HumanInputFormKind.RUNTIME, + node_id="node-1", + form_definition=json.dumps({"display_in_ui": True}), + rendered_content="Rendered", + status=HumanInputFormStatus.WAITING, + expiration_time=expiration_time, + ) + sqlite_session.add(form) + sqlite_session.commit() + return form + + +def _build_service_api_pause_converter() -> WorkflowResponseConverter: + application_generate_entity = WorkflowAppGenerateEntity( + task_id="task", + workflow_execution_id="run-id", + user_id="user", + stream=True, + inputs={}, + files=[], + invoke_from=InvokeFrom.SERVICE_API, + app_config=WorkflowUIBasedAppConfig( + app_id="app-id", tenant_id="tenant-id", workflow_id="workflow-id", app_mode=AppMode.WORKFLOW + ), + ) + system_variables = build_system_variables( + user_id="user", + app_id="app-id", + workflow_id="workflow-id", + workflow_execution_id="run-id", + ) + user = Account(name="Tester", email="tester@example.com") + user.id = "account-id" + return WorkflowResponseConverter( + application_generate_entity=application_generate_entity, + user=user, + system_variables=system_variables, + ) + + +def _build_advanced_chat_paused_blocking_response() -> AdvancedChatPausedBlockingResponse: + data = AdvancedChatPausedBlockingResponse.Data( + id="msg-1", + mode="chat", + conversation_id="c1", + message_id="m1", + workflow_run_id="run-1", + answer="partial", + metadata={"usage": {"total_tokens": 1}}, + created_at=1, + paused_nodes=["node-1"], + reasons=[ + { + "TYPE": DifyHITLEventType.HUMAN_INPUT_REQUIRED.value, + "form_id": "form-1", + "expiration_time": 100, + } + ], + status=WorkflowExecutionStatus.PAUSED, + elapsed_time=0.1, + total_tokens=0, + total_steps=0, + ) + return AdvancedChatPausedBlockingResponse(task_id="t1", data=data) + + +def _build_workflow_paused_blocking_response() -> WorkflowAppPausedBlockingResponse: + return WorkflowAppPausedBlockingResponse( + task_id="t1", + workflow_run_id="r1", + data=WorkflowAppPausedBlockingResponse.Data( + id="r1", + workflow_id="wf-1", + status=WorkflowExecutionStatus.PAUSED, + outputs={}, + error=None, + elapsed_time=0.5, + total_tokens=0, + total_steps=2, + created_at=1, + finished_at=None, + paused_nodes=["node-1"], + reasons=[{"TYPE": "human_input_required", "form_id": "form-1", "expiration_time": 100}], + ), + ) + + +@dataclass(frozen=True) +class _FakePauseEntity(WorkflowPauseEntity): + pause_id: str + workflow_run_id: str + paused_at_value: datetime + pause_reasons: Sequence[HumanInputRequired] + + @property + @override + def id(self) -> str: + return self.pause_id + + @property + @override + def workflow_execution_id(self) -> str: + return self.workflow_run_id + + @override + def get_state(self) -> bytes: + raise AssertionError("state is not required for snapshot tests") + + @property + @override + def resumed_at(self) -> datetime | None: + return None + + @property + @override + def paused_at(self) -> datetime: + return self.paused_at_value + + @override + def get_pause_reasons(self) -> Sequence[HumanInputRequired]: + return self.pause_reasons + + +def _build_workflow_run(status: WorkflowExecutionStatus) -> WorkflowRun: + return WorkflowRun( + id="run-1", + tenant_id="tenant-1", + app_id="app-1", + workflow_id="workflow-1", + type="workflow", + triggered_from="app-run", + version="v1", + graph=None, + inputs=json.dumps({"input": "value"}), + status=status, + outputs=json.dumps({}), + error=None, + elapsed_time=0.0, + total_tokens=0, + total_steps=0, + created_by_role=CreatorUserRole.END_USER, + created_by="user-1", + created_at=datetime(2024, 1, 1, tzinfo=UTC), + ) + + +def _build_snapshot(status: WorkflowNodeExecutionStatus) -> WorkflowNodeExecutionSnapshot: + created_at = datetime(2024, 1, 1, tzinfo=UTC) + finished_at = datetime(2024, 1, 1, 0, 0, 5, tzinfo=UTC) + return WorkflowNodeExecutionSnapshot( + execution_id="exec-1", + node_id="node-1", + node_type="human-input", + title="Human Input", + index=1, + status=status.value, + elapsed_time=0.5, + created_at=created_at, + finished_at=finished_at, + iteration_id=None, + loop_id=None, + ) + + +def _build_resumption_context(task_id: str) -> WorkflowResumptionContext: + app_config = WorkflowUIBasedAppConfig( + tenant_id="tenant-1", + app_id="app-1", + app_mode=AppMode.WORKFLOW, + workflow_id="workflow-1", + ) + generate_entity = WorkflowAppGenerateEntity( + task_id=task_id, + app_config=app_config, + inputs={}, + files=[], + user_id="user-1", + stream=True, + invoke_from=InvokeFrom.EXPLORE, + call_depth=0, + workflow_execution_id="run-1", + ) + runtime_state = RuntimeState(workflow_id="test-workflow", variable_pool=VariablePool(), start_at=0.0) + runtime_state.set_output("result", "value") + wrapper = _WorkflowGenerateEntityWrapper(entity=generate_entity) + return WorkflowResumptionContext( + generate_entity=wrapper, + serialized_graph_runtime_state=runtime_state.dumps(), + ) + + +class TestHitlServiceApi: + # Service API event-stream continuation + def test_workflow_events_continue_on_pause_keeps_stream_open( + self, + app: Flask, + monkeypatch: pytest.MonkeyPatch, + sqlite_engine: Engine, + ) -> None: + workflow_run = _build_workflow_run(WorkflowExecutionStatus.RUNNING) + workflow_run.created_by = "end-user-1" + workflow_events_module = _hitl_service_api_mock_repo_for_run( + monkeypatch, + workflow_run=workflow_run, + sqlite_engine=sqlite_engine, + ) + msg_generator = Mock() + msg_generator.retrieve_events.return_value = [{"event": "streamed"}] + monkeypatch.setattr(workflow_events_module, "MessageGenerator", lambda: msg_generator) + + api = WorkflowEventsApi() + handler = unwrap(api.get) + app_model = _hitl_service_api_app() + end_user = _hitl_service_api_end_user() + + with app.test_request_context("/workflow/run-1/events?user=u1&continue_on_pause=true", method="GET"): + response = handler(api, app_model=app_model, end_user=end_user, workflow_run_id="run-1") + + assert response.get_data(as_text=True) == 'data: {"event":"streamed"}\n\n' + msg_generator.retrieve_events.assert_called_once_with( + AppMode.WORKFLOW, + "run-1", + terminal_events=[], + ) + + def test_workflow_events_snapshot_continue_on_pause_keeps_pause_open( + self, + app: Flask, + monkeypatch: pytest.MonkeyPatch, + sqlite_engine: Engine, + ) -> None: + workflow_run = _build_workflow_run(WorkflowExecutionStatus.RUNNING) + workflow_run.created_by = "end-user-1" + workflow_events_module = _hitl_service_api_mock_repo_for_run( + monkeypatch, + workflow_run=workflow_run, + sqlite_engine=sqlite_engine, + ) + msg_generator = Mock() + snapshot_calls: list[dict[str, object]] = [] + + def snapshot_builder(**kwargs: object) -> list[dict[str, str]]: + snapshot_calls.append(kwargs) + return [{"event": "snapshot"}] + + monkeypatch.setattr(workflow_events_module, "MessageGenerator", lambda: msg_generator) + monkeypatch.setattr(workflow_events_module, "build_workflow_event_stream", snapshot_builder) + + api = WorkflowEventsApi() + handler = unwrap(api.get) + app_model = _hitl_service_api_app() + end_user = _hitl_service_api_end_user() + + with app.test_request_context( + "/workflow/run-1/events?user=u1&include_state_snapshot=true&continue_on_pause=true", + method="GET", + ): + response = handler(api, app_model=app_model, end_user=end_user, workflow_run_id="run-1") + + assert response.get_data(as_text=True) == 'data: {"event":"snapshot"}\n\n' + msg_generator.retrieve_events.assert_not_called() + assert snapshot_calls == [ + { + "app_mode": AppMode.WORKFLOW, + "workflow_run": workflow_run, + "tenant_id": "tenant-1", + "app_id": "app-1", + "session_maker": ANY, + "human_input_surface": HumanInputSurface.SERVICE_API, + "close_on_pause": False, + } + ] + snapshot_session_maker = snapshot_calls[0]["session_maker"] + assert isinstance(snapshot_session_maker, sessionmaker) + assert snapshot_session_maker.kw["bind"] is sqlite_engine + + def test_advanced_chat_blocking_injects_pause_state_config( + self, + monkeypatch: pytest.MonkeyPatch, + sqlite_engine: Engine, + ) -> None: + apply_config_overrides(monkeypatch, DEPLOYMENT_EDITION=DeploymentEdition.COMMUNITY) + monkeypatch.setattr(ags_module, "RateLimit", _DummyRateLimit) + + workflow = _workflow() + monkeypatch.setattr(AppGenerateService, "_get_workflow", lambda *args, **kwargs: workflow) + sqlite_session_maker = sessionmaker(bind=sqlite_engine, expire_on_commit=False) + monkeypatch.setattr(ags_module.session_factory, "get_session_maker", lambda: sqlite_session_maker) + + generator_instance = MagicMock() + generator_instance.generate.return_value = {"result": "advanced-blocking"} + generator_instance.convert_to_event_stream.side_effect = lambda payload: payload + monkeypatch.setattr(ags_module, "AdvancedChatAppGenerator", lambda **_kwargs: generator_instance) + + app_model = _hitl_service_api_app(app_id="app-id", tenant_id="tenant-id", mode=AppMode.ADVANCED_CHAT) + user = _hitl_service_api_end_user(user_id="user-id", app_id="app-id", tenant_id="tenant-id") + + with sqlite_session_maker() as session: + result = AppGenerateService.generate( + session=session, + app_model=app_model, + user=user, + args={"workflow_id": None, "query": "hi", "inputs": {}}, + invoke_from=InvokeFrom.SERVICE_API, + streaming=False, + ) + + assert result == {"result": "advanced-blocking"} + call_kwargs = generator_instance.generate.call_args.kwargs + assert call_kwargs["streaming"] is False + assert call_kwargs["pause_state_config"] is not None + assert call_kwargs["pause_state_config"].session_factory is sqlite_session_maker + assert call_kwargs["pause_state_config"].state_owner_user_id == "owner-id" + + # Blocking payload contract + def test_advanced_chat_blocking_pause_payload_contract(self) -> None: + from core.app.apps.advanced_chat.generate_response_converter import AdvancedChatAppGenerateResponseConverter + + response = AdvancedChatAppGenerateResponseConverter.convert_blocking_full_response( + _build_advanced_chat_paused_blocking_response() + ) + + assert response["event"] == "workflow_paused" + assert response["workflow_run_id"] == "run-1" + assert response["answer"] == "partial" + assert response["data"]["reasons"][0]["TYPE"] == DifyHITLEventType.HUMAN_INPUT_REQUIRED + assert response["data"]["reasons"][0]["expiration_time"] == 100 + assert "human_input_forms" not in response["data"] + + def test_workflow_blocking_pause_payload_contract(self) -> None: + from core.app.apps.workflow.generate_response_converter import WorkflowAppGenerateResponseConverter + + response = WorkflowAppGenerateResponseConverter.convert_blocking_full_response( + _build_workflow_paused_blocking_response() + ) + + assert response["workflow_run_id"] == "r1" + assert response["data"]["status"] == WorkflowExecutionStatus.PAUSED + assert response["data"]["paused_nodes"] == ["node-1"] + assert response["data"]["reasons"] == [ + {"TYPE": "human_input_required", "form_id": "form-1", "expiration_time": 100} + ] + assert "human_input_forms" not in response["data"] + + def test_advanced_chat_blocking_pipeline_pause_payload_contract(self) -> None: + from core.app.app_config.entities import AppAdditionalFeatures + from core.app.apps.advanced_chat.generate_task_pipeline import AdvancedChatAppGenerateTaskPipeline + + app_config = WorkflowUIBasedAppConfig( + tenant_id="tenant", + app_id="app", + app_mode=AppMode.ADVANCED_CHAT, + additional_features=AppAdditionalFeatures(), + variables=[], + workflow_id="workflow-id", + ) + application_generate_entity = AdvancedChatAppGenerateEntity.model_construct( + task_id="task", + app_config=app_config, + inputs={}, + query="hello", + files=[], + user_id="user", + stream=False, + invoke_from=InvokeFrom.WEB_APP, + extras={}, + trace_manager=None, + workflow_run_id="run-id", + ) + pipeline = AdvancedChatAppGenerateTaskPipeline( + application_generate_entity=application_generate_entity, + workflow=WorkflowSnapshot(id="workflow-id", tenant_id="tenant", features_dict={}), + queue_manager=SimpleNamespace(invoke_from=InvokeFrom.WEB_APP, graph_runtime_state=None), + conversation=ConversationSnapshot(id="conv-id", mode=AppMode.ADVANCED_CHAT), + message=MessageSnapshot( + id="message-id", + query="hello", + created_at=datetime.utcnow(), + status=MessageStatus.NORMAL, + answer="", + ), + user=_hitl_service_api_end_user(user_id="user", app_id="app", tenant_id="tenant"), + stream=False, + dialogue_count=1, + draft_var_saver_factory=lambda **kwargs: None, + ) + pipeline._task_state.answer = "partial answer" + pipeline._workflow_run_id = "run-id" + + def _gen(): + yield HumanInputRequiredResponse( + task_id="task", + workflow_run_id="run-id", + data=HumanInputRequiredResponse.Data( + form_id="form-1", + node_id="node-1", + node_title="Approval", + form_content="Need approval", + inputs=[], + actions=[UserActionConfig(id="approve", title="Approve")], + display_in_ui=True, + form_token="token-1", + resolved_default_values={}, + expiration_time=123, + ), + ) + yield WorkflowPauseStreamResponse( + task_id="task", + workflow_run_id="run-id", + data=WorkflowPauseStreamResponse.Data( + workflow_run_id="run-id", + paused_nodes=["node-1"], + outputs={}, + reasons=[ + { + "TYPE": DifyHITLEventType.HUMAN_INPUT_REQUIRED.value, + "form_id": "form-1", + "node_id": "node-1", + "expiration_time": 123, + }, + ], + status="paused", + created_at=1, + elapsed_time=0.1, + total_tokens=0, + total_steps=0, + ), + ) + + response = pipeline._to_blocking_response(_gen()) + + assert isinstance(response, AdvancedChatPausedBlockingResponse) + assert response.data.answer == "partial answer" + assert response.data.workflow_run_id == "run-id" + assert response.data.reasons[0]["form_id"] == "form-1" + assert response.data.reasons[0]["expiration_time"] == 123 + + def test_workflow_blocking_pipeline_pause_payload_contract(self) -> None: + from core.app.apps.workflow.generate_task_pipeline import WorkflowAppGenerateTaskPipeline + + app_config = WorkflowUIBasedAppConfig( + tenant_id="tenant", + app_id="app", + app_mode=AppMode.WORKFLOW, + additional_features=AppAdditionalFeatures(), + variables=[], + workflow_id="workflow-id", + ) + application_generate_entity = WorkflowAppGenerateEntity.model_construct( + task_id="task", + app_config=app_config, + inputs={}, + files=[], + user_id="user", + stream=False, + invoke_from=InvokeFrom.WEB_APP, + trace_manager=None, + workflow_execution_id="run-id", + extras={}, + call_depth=0, + ) + pipeline = WorkflowAppGenerateTaskPipeline( + application_generate_entity=application_generate_entity, + workflow=_workflow(workflow_id="workflow-id", app_id="app", tenant_id="tenant"), + queue_manager=SimpleNamespace(invoke_from=InvokeFrom.WEB_APP, graph_runtime_state=None), + user=_hitl_service_api_end_user(user_id="user", app_id="app", tenant_id="tenant"), + stream=False, + draft_var_saver_factory=lambda **kwargs: None, + ) + + def _gen(): + yield HumanInputRequiredResponse( + task_id="task", + workflow_run_id="run", + data=HumanInputRequiredResponse.Data( + form_id="form-1", + node_id="node-1", + node_title="Human Input", + form_content="content", + expiration_time=1, + ), + ) + yield WorkflowPauseStreamResponse( + task_id="task", + workflow_run_id="run", + data=WorkflowPauseStreamResponse.Data( + workflow_run_id="run", + status=WorkflowExecutionStatus.PAUSED, + outputs={}, + paused_nodes=["node-1"], + reasons=[{"TYPE": "human_input_required", "form_id": "form-1", "expiration_time": 1}], + created_at=1, + elapsed_time=0.1, + total_tokens=0, + total_steps=0, + ), + ) + + response = pipeline._to_blocking_response(_gen()) + + assert isinstance(response, WorkflowAppPausedBlockingResponse) + assert response.data.status == WorkflowExecutionStatus.PAUSED + assert response.data.paused_nodes == ["node-1"] + assert response.data.reasons == [{"TYPE": "human_input_required", "form_id": "form-1", "expiration_time": 1}] + + def test_service_api_pause_event_serializes_hitl_reason( + self, + monkeypatch: pytest.MonkeyPatch, + sqlite_engine: Engine, + sqlite_session: Session, + ) -> None: + converter = _build_service_api_pause_converter() + converter.workflow_start_to_stream_response( + task_id="task", + workflow_run_id="run-id", + workflow_id="workflow-id", + reason=WorkflowStartReason.INITIAL, + ) + + expiration_time = datetime(2024, 1, 1, tzinfo=UTC) + _persist_human_input_form(sqlite_session, expiration_time=expiration_time) + + monkeypatch.setattr(workflow_response_converter, "db", SimpleNamespace(engine=sqlite_engine)) + monkeypatch.setattr( + workflow_response_converter, + "load_form_dispositions_by_form_id", + lambda form_ids, session=None, surface=None: { + "form-1": FormDisposition(form_token="token", approval_channels=[]) + }, + ) + + reason = HumanInputRequired( + form_id="form-1", + form_content="Rendered", + inputs=[ + ParagraphInputConfig( + type=FormInputType.PARAGRAPH, + output_variable_name="field", + default=None, + ), + ], + actions=[UserActionConfig(id="approve", title="Approve")], + display_in_ui=True, + node_id="node-id", + node_title="Human Step", + form_token="token", + ) + queue_event = QueueWorkflowPausedEvent( + reasons=[reason], + outputs={"answer": "value"}, + paused_nodes=["node-id"], + ) + + runtime_state = SimpleNamespace(total_tokens=0, node_run_steps=0, variable_pool=VariablePool()) + responses = converter.workflow_pause_to_stream_response( + event=queue_event, + task_id="task", + graph_runtime_state=runtime_state, + ) + + assert isinstance(responses[-1], WorkflowPauseStreamResponse) + pause_resp = responses[-1] + assert pause_resp.workflow_run_id == "run-id" + assert pause_resp.data.paused_nodes == ["node-id"] + assert pause_resp.data.outputs == {} + assert pause_resp.data.reasons[0]["TYPE"] == "human_input_required" + assert pause_resp.data.reasons[0]["form_id"] == "form-1" + assert pause_resp.data.reasons[0]["form_token"] == "token" + assert pause_resp.data.reasons[0]["expiration_time"] == to_utc_timestamp(expiration_time) + + assert isinstance(responses[0], HumanInputRequiredResponse) + hi_resp = responses[0] + assert hi_resp.data.form_id == "form-1" + assert hi_resp.data.node_id == "node-id" + assert hi_resp.data.node_title == "Human Step" + assert hi_resp.data.inputs[0].output_variable_name == "field" + assert hi_resp.data.actions[0].id == "approve" + assert hi_resp.data.display_in_ui is True + assert hi_resp.data.form_token == "token" + assert hi_resp.data.expiration_time == to_utc_timestamp(expiration_time) + + # Snapshot payload contract + def test_snapshot_events_include_pause_payload_contract( + self, + monkeypatch: pytest.MonkeyPatch, + sqlite_engine: Engine, + sqlite_session: Session, + ) -> None: + workflow_run = _build_workflow_run(WorkflowExecutionStatus.PAUSED) + snapshot = _build_snapshot(WorkflowNodeExecutionStatus.PAUSED) + resumption_context = _build_resumption_context("task-ctx") + expiration_time = datetime(2024, 1, 1, tzinfo=UTC) + _persist_human_input_form(sqlite_session, expiration_time=expiration_time) + monkeypatch.setattr( + "services.workflow_event_snapshot_service.load_form_dispositions_by_form_id", + lambda form_ids, session=None, surface=None: { + "form-1": FormDisposition(form_token="wtok", approval_channels=[]) + }, + ) + + sqlite_session_maker = sessionmaker(bind=sqlite_engine, expire_on_commit=False) + + pause_entity = _FakePauseEntity( + pause_id="pause-1", + workflow_run_id="run-1", + paused_at_value=datetime(2024, 1, 1, tzinfo=UTC), + pause_reasons=[ + HumanInputRequired( + form_id="form-1", + form_content="content", + node_id="node-1", + node_title="Human Input", + form_token="wtok", + ) + ], + ) + + events = _build_snapshot_events( + workflow_run=workflow_run, + node_snapshots=[snapshot], + task_id="task-ctx", + message_context=None, + pause_entity=pause_entity, + resumption_context=resumption_context, + session_maker=sqlite_session_maker, + node_executions={}, + ) + + assert [event["event"] for event in events] == [ + "workflow_started", + "node_started", + "node_finished", + "human_input_required", + "workflow_paused", + ] + assert events[2]["data"]["status"] == WorkflowNodeExecutionStatus.PAUSED.value + assert events[3]["data"]["form_token"] == "wtok" + assert events[3]["data"]["expiration_time"] == to_utc_timestamp(expiration_time) + pause_data = events[-1]["data"] + assert pause_data["paused_nodes"] == ["node-1"] + assert pause_data["outputs"] == {"result": "value"} + assert pause_data["reasons"][0]["TYPE"] == "human_input_required" + assert pause_data["reasons"][0]["form_token"] == "wtok" + assert pause_data["reasons"][0]["expiration_time"] == to_utc_timestamp(expiration_time) + assert pause_data["status"] == WorkflowExecutionStatus.PAUSED.value + assert pause_data["created_at"] == int(workflow_run.created_at.timestamp()) + assert pause_data["elapsed_time"] == workflow_run.elapsed_time + assert pause_data["total_tokens"] == workflow_run.total_tokens + assert pause_data["total_steps"] == workflow_run.total_steps + + +"""Unit tests for Service API workflow event stream endpoints.""" + + +def _workflow_events_app(*, mode: AppMode = AppMode.WORKFLOW) -> App: return App( id="app-1", tenant_id="tenant-1", @@ -35,7 +810,7 @@ def _app(*, mode: AppMode = AppMode.WORKFLOW) -> App: ) -def _end_user() -> EndUser: +def _workflow_events_end_user() -> EndUser: return EndUser( id="end-user-1", tenant_id="tenant-1", @@ -70,7 +845,7 @@ def _workflow_run(*, created_by_role: CreatorUserRole = CreatorUserRole.END_USER ) -def _mock_repo_for_run(monkeypatch: pytest.MonkeyPatch, workflow_run, sqlite_engine: Engine): +def _workflow_events_mock_repo_for_run(monkeypatch: pytest.MonkeyPatch, workflow_run, sqlite_engine: Engine): workflow_events_module = sys.modules["controllers.service_api.app.workflow_events"] repo = SimpleNamespace(get_workflow_run_by_id_and_tenant_id=lambda **_kwargs: workflow_run) monkeypatch.setattr( @@ -86,19 +861,19 @@ class TestWorkflowEventsApi: def test_wrong_app_mode(self, app: Flask) -> None: api = WorkflowEventsApi() handler = unwrap(api.get) - app_model = _app(mode=AppMode.CHAT) - end_user = _end_user() + app_model = _workflow_events_app(mode=AppMode.CHAT) + end_user = _workflow_events_end_user() with app.test_request_context("/workflow/run-1/events?user=u1", method="GET"): with pytest.raises(NotWorkflowAppError): handler(api, app_model=app_model, end_user=end_user, workflow_run_id="run-1") def test_workflow_run_not_found(self, app: Flask, monkeypatch: pytest.MonkeyPatch, sqlite_engine: Engine) -> None: - _mock_repo_for_run(monkeypatch, workflow_run=None, sqlite_engine=sqlite_engine) + _workflow_events_mock_repo_for_run(monkeypatch, workflow_run=None, sqlite_engine=sqlite_engine) api = WorkflowEventsApi() handler = unwrap(api.get) - app_model = _app() - end_user = _end_user() + app_model = _workflow_events_app() + end_user = _workflow_events_end_user() with app.test_request_context("/workflow/run-1/events?user=u1", method="GET"): with pytest.raises(NotFound): @@ -108,11 +883,11 @@ class TestWorkflowEventsApi: self, app: Flask, monkeypatch: pytest.MonkeyPatch, sqlite_engine: Engine ) -> None: workflow_run = _workflow_run(created_by_role=CreatorUserRole.ACCOUNT) - _mock_repo_for_run(monkeypatch, workflow_run=workflow_run, sqlite_engine=sqlite_engine) + _workflow_events_mock_repo_for_run(monkeypatch, workflow_run=workflow_run, sqlite_engine=sqlite_engine) api = WorkflowEventsApi() handler = unwrap(api.get) - app_model = _app() - end_user = _end_user() + app_model = _workflow_events_app() + end_user = _workflow_events_end_user() with app.test_request_context("/workflow/run-1/events?user=u1", method="GET"): with pytest.raises(NotFound): @@ -121,7 +896,9 @@ class TestWorkflowEventsApi: def test_finished_run_returns_sse(self, app: Flask, monkeypatch: pytest.MonkeyPatch, sqlite_engine: Engine) -> None: workflow_run = _workflow_run() workflow_run.finished_at = datetime(2099, 1, 1, tzinfo=UTC) - workflow_events_module = _mock_repo_for_run(monkeypatch, workflow_run=workflow_run, sqlite_engine=sqlite_engine) + workflow_events_module = _workflow_events_mock_repo_for_run( + monkeypatch, workflow_run=workflow_run, sqlite_engine=sqlite_engine + ) monkeypatch.setattr( workflow_events_module.WorkflowResponseConverter, "workflow_run_result_to_finish_response", @@ -133,8 +910,8 @@ class TestWorkflowEventsApi: api = WorkflowEventsApi() handler = unwrap(api.get) - app_model = _app() - end_user = _end_user() + app_model = _workflow_events_app() + end_user = _workflow_events_end_user() with app.test_request_context("/workflow/run-1/events?user=u1", method="GET"): response = handler(api, app_model=app_model, end_user=end_user, workflow_run_id="run-1") @@ -150,15 +927,17 @@ class TestWorkflowEventsApi: self, app: Flask, monkeypatch: pytest.MonkeyPatch, sqlite_engine: Engine ) -> None: workflow_run = _workflow_run() - workflow_events_module = _mock_repo_for_run(monkeypatch, workflow_run=workflow_run, sqlite_engine=sqlite_engine) + workflow_events_module = _workflow_events_mock_repo_for_run( + monkeypatch, workflow_run=workflow_run, sqlite_engine=sqlite_engine + ) msg_generator = Mock() msg_generator.retrieve_events.return_value = [{"event": "streamed"}] monkeypatch.setattr(workflow_events_module, "MessageGenerator", lambda: msg_generator) api = WorkflowEventsApi() handler = unwrap(api.get) - app_model = _app() - end_user = _end_user() + app_model = _workflow_events_app() + end_user = _workflow_events_end_user() with app.test_request_context("/workflow/run-1/events?user=u1", method="GET"): response = handler(api, app_model=app_model, end_user=end_user, workflow_run_id="run-1") @@ -174,7 +953,9 @@ class TestWorkflowEventsApi: self, app: Flask, monkeypatch: pytest.MonkeyPatch, sqlite_engine: Engine ) -> None: workflow_run = _workflow_run() - workflow_events_module = _mock_repo_for_run(monkeypatch, workflow_run=workflow_run, sqlite_engine=sqlite_engine) + workflow_events_module = _workflow_events_mock_repo_for_run( + monkeypatch, workflow_run=workflow_run, sqlite_engine=sqlite_engine + ) msg_generator = Mock() snapshot_calls: list[dict[str, object]] = [] @@ -187,8 +968,8 @@ class TestWorkflowEventsApi: api = WorkflowEventsApi() handler = unwrap(api.get) - app_model = _app() - end_user = _end_user() + app_model = _workflow_events_app() + end_user = _workflow_events_end_user() with app.test_request_context("/workflow/run-1/events?user=u1&include_state_snapshot=true", method="GET"): response = handler(api, app_model=app_model, end_user=end_user, workflow_run_id="run-1") diff --git a/api/tests/unit_tests/controllers/service_api/app/test_workflow_fields.py b/api/tests/unit_tests/controllers/service_api/app/test_workflow_fields.py deleted file mode 100644 index caa6ddc0d93..00000000000 --- a/api/tests/unit_tests/controllers/service_api/app/test_workflow_fields.py +++ /dev/null @@ -1,36 +0,0 @@ -from controllers.service_api.app.workflow import WorkflowRunResponse -from graphon.enums import WorkflowExecutionStatus -from libs.helper import dump_response -from models.workflow import WorkflowRun - - -def _workflow_run(status: WorkflowExecutionStatus, outputs: str | None = '{"foo": "bar"}') -> WorkflowRun: - return WorkflowRun( - id="run-id", - workflow_id="workflow-id", - status=status, - inputs="{}", - outputs=outputs, - error=None, - total_steps=1, - total_tokens=2, - elapsed_time=3.5, - ) - - -def test_workflow_run_serializer_normalizes_status_enum() -> None: - response = dump_response(WorkflowRunResponse, _workflow_run(WorkflowExecutionStatus.PAUSED)) - - assert response["status"] == "paused" - - -def test_workflow_run_serializer_paused_returns_empty_outputs() -> None: - response = dump_response(WorkflowRunResponse, _workflow_run(WorkflowExecutionStatus.PAUSED)) - - assert response["outputs"] == {} - - -def test_workflow_run_serializer_running_returns_outputs() -> None: - response = dump_response(WorkflowRunResponse, _workflow_run(WorkflowExecutionStatus.RUNNING)) - - assert response["outputs"] == {"foo": "bar"} diff --git a/api/tests/unit_tests/controllers/service_api/dataset/rag_pipeline/test_rag_pipeline_workflow.py b/api/tests/unit_tests/controllers/service_api/dataset/rag_pipeline/test_rag_pipeline_workflow.py index b1b10993968..0112481074c 100644 --- a/api/tests/unit_tests/controllers/service_api/dataset/rag_pipeline/test_rag_pipeline_workflow.py +++ b/api/tests/unit_tests/controllers/service_api/dataset/rag_pipeline/test_rag_pipeline_workflow.py @@ -1,25 +1,8 @@ -""" -Unit tests for Service API RAG Pipeline Workflow controllers. - -Tests coverage for: -- DatasourceNodeRunPayload Pydantic model -- PipelineRunApiEntity / DatasourceNodeRunApiEntity model validation -- RAG pipeline service interfaces -- File upload validation for pipelines -- Endpoint tests for DatasourcePluginsApi, DatasourceNodeRunApi, - PipelineRunApi, and KnowledgebasePipelineFileUploadApi - -Strategy: -- Endpoint methods on these resources have no billing decorators on the method - itself. ``method_decorators = [validate_dataset_token]`` is only invoked by - Flask-RESTx dispatch, not by direct calls, so we call methods directly. -- Dataset ownership checks run against the shared SQLite fixture; pipeline and - upload responses use real mapped instances while external services stay mocked. -""" - +import ast import io import uuid from datetime import UTC, datetime +from pathlib import Path from types import SimpleNamespace from unittest.mock import Mock, patch @@ -48,10 +31,12 @@ from controllers.service_api.dataset.rag_pipeline.rag_pipeline_workflow import ( DatasourcePluginsApi, KnowledgebasePipelineFileUploadApi, PipelineRunApi, + PipelineUploadFileResponse, ) from core.app.apps.pipeline.pipeline_generator import PipelineGenerator from core.app.entities.app_invoke_entities import InvokeFrom from extensions.storage.storage_type import StorageType +from libs.helper import dump_response from models.account import Account from models.dataset import Dataset, Pipeline from models.enums import CreatorUserRole @@ -66,6 +51,26 @@ from services.rag_pipeline.entity.pipeline_service_api_entities import ( PipelineRunApiEntity, ) from services.rag_pipeline.rag_pipeline import RagPipelineService +from tests.helpers.unit.model_factories import make_upload_file + +""" +Unit tests for Service API RAG Pipeline Workflow controllers. + +Tests coverage for: +- DatasourceNodeRunPayload Pydantic model +- PipelineRunApiEntity / DatasourceNodeRunApiEntity model validation +- RAG pipeline service interfaces +- File upload validation for pipelines +- Endpoint tests for DatasourcePluginsApi, DatasourceNodeRunApi, + PipelineRunApi, and KnowledgebasePipelineFileUploadApi + +Strategy: +- Endpoint methods on these resources have no billing decorators on the method + itself. ``method_decorators = [validate_dataset_token]`` is only invoked by + Flask-RESTx dispatch, not by direct calls, so we call methods directly. +- Dataset ownership checks run against the shared SQLite fixture; pipeline and + upload responses use real mapped instances while external services stay mocked. +""" def _persist_dataset(session: Session, *, tenant_id: str, dataset_id: str) -> Dataset: @@ -89,756 +94,835 @@ def _persist_pipeline(session: Session, *, tenant_id: str) -> Pipeline: return pipeline -@pytest.fixture(autouse=True) -def _bind_database(sqlite_session_factory: sessionmaker[Session], monkeypatch: pytest.MonkeyPatch): - session_proxy = scoped_session(sqlite_session_factory) - monkeypatch.setattr(workflow_module.db, "session", session_proxy) - query = CredentialQueryRepository(session_factory=sqlite_session_factory) - providers = DatasourceProviderService( - credentials=SQLAlchemyDatasourceCredentialRepository(session_factory=sqlite_session_factory) - ) - monkeypatch.setattr( - workflow_module, - "application_services", - lambda: SimpleNamespace(credential_queries=query, data_sources=SimpleNamespace(providers=providers)), - ) - yield - session_proxy.remove() - - -class TestDatasourceNodeRunPayload: - """Test suite for DatasourceNodeRunPayload Pydantic model.""" - - def test_payload_with_required_fields(self): - """Test payload with required fields.""" - payload = DatasourceNodeRunPayload( - inputs={"key": "value"}, datasource_type="online_document", is_published=True +class TestRagPipelineWorkflowCases: + @staticmethod + @pytest.fixture(autouse=True) + def _bind_database(sqlite_session_factory: sessionmaker[Session], monkeypatch: pytest.MonkeyPatch): + session_proxy = scoped_session(sqlite_session_factory) + monkeypatch.setattr(workflow_module.db, "session", session_proxy) + query = CredentialQueryRepository(session_factory=sqlite_session_factory) + providers = DatasourceProviderService( + credentials=SQLAlchemyDatasourceCredentialRepository(session_factory=sqlite_session_factory) ) - assert payload.inputs == {"key": "value"} - assert payload.datasource_type == "online_document" - assert payload.is_published is True - assert payload.credential_id is None - - def test_payload_with_credential_id(self): - """Test payload with optional credential_id.""" - payload = DatasourceNodeRunPayload( - inputs={"url": "https://example.com"}, - datasource_type="online_document", - credential_id="cred_123", - is_published=False, + monkeypatch.setattr( + workflow_module, + "application_services", + lambda: SimpleNamespace(credential_queries=query, data_sources=SimpleNamespace(providers=providers)), ) - assert payload.credential_id == "cred_123" - assert payload.is_published is False + yield + session_proxy.remove() - def test_payload_with_complex_inputs(self): - """Test payload with complex nested inputs.""" - complex_inputs = { - "config": {"url": "https://api.example.com", "headers": {"Authorization": "Bearer token"}}, - "parameters": {"limit": 100, "offset": 0}, - "options": ["opt1", "opt2"], - } - payload = DatasourceNodeRunPayload(inputs=complex_inputs, datasource_type="api", is_published=True) - assert payload.inputs == complex_inputs + class TestDatasourceNodeRunPayload: + """Test suite for DatasourceNodeRunPayload Pydantic model.""" - def test_payload_with_empty_inputs(self): - """Test payload with empty inputs dict.""" - payload = DatasourceNodeRunPayload(inputs={}, datasource_type="local_file", is_published=True) - assert payload.inputs == {} + def test_payload_with_required_fields(self): + """Test payload with required fields.""" + payload = DatasourceNodeRunPayload( + inputs={"key": "value"}, datasource_type="online_document", is_published=True + ) + assert payload.inputs == {"key": "value"} + assert payload.datasource_type == "online_document" + assert payload.is_published is True + assert payload.credential_id is None - @pytest.mark.parametrize("datasource_type", ["online_document", "local_file", "api", "database", "website"]) - def test_payload_common_datasource_types(self, datasource_type): - """Test payload with common datasource types.""" - payload = DatasourceNodeRunPayload(inputs={}, datasource_type=datasource_type, is_published=True) - assert payload.datasource_type == datasource_type + def test_payload_with_credential_id(self): + """Test payload with optional credential_id.""" + payload = DatasourceNodeRunPayload( + inputs={"url": "https://example.com"}, + datasource_type="online_document", + credential_id="cred_123", + is_published=False, + ) + assert payload.credential_id == "cred_123" + assert payload.is_published is False + def test_payload_with_complex_inputs(self): + """Test payload with complex nested inputs.""" + complex_inputs = { + "config": {"url": "https://api.example.com", "headers": {"Authorization": "Bearer token"}}, + "parameters": {"limit": 100, "offset": 0}, + "options": ["opt1", "opt2"], + } + payload = DatasourceNodeRunPayload(inputs=complex_inputs, datasource_type="api", is_published=True) + assert payload.inputs == complex_inputs -class TestPipelineErrors: - """Test pipeline-related error types.""" + def test_payload_with_empty_inputs(self): + """Test payload with empty inputs dict.""" + payload = DatasourceNodeRunPayload(inputs={}, datasource_type="local_file", is_published=True) + assert payload.inputs == {} - def test_pipeline_run_error_can_be_raised(self): - """Test PipelineRunError can be raised.""" - error = PipelineRunError(description="Pipeline execution failed") - assert error is not None + @pytest.mark.parametrize("datasource_type", ["online_document", "local_file", "api", "database", "website"]) + def test_payload_common_datasource_types(self, datasource_type): + """Test payload with common datasource types.""" + payload = DatasourceNodeRunPayload(inputs={}, datasource_type=datasource_type, is_published=True) + assert payload.datasource_type == datasource_type - def test_pipeline_run_error_with_description(self): - """Test PipelineRunError captures description.""" - error = PipelineRunError(description="Timeout during node execution") - # The error should have the description attribute - assert hasattr(error, "description") + class TestPipelineErrors: + """Test pipeline-related error types.""" + def test_pipeline_run_error_can_be_raised(self): + """Test PipelineRunError can be raised.""" + error = PipelineRunError(description="Pipeline execution failed") + assert error is not None -class TestFileUploadErrors: - """Test file upload error types for pipelines.""" + def test_pipeline_run_error_with_description(self): + """Test PipelineRunError captures description.""" + error = PipelineRunError(description="Timeout during node execution") + # The error should have the description attribute + assert hasattr(error, "description") - def test_no_file_uploaded_error(self): - """Test NoFileUploadedError can be raised.""" - error = NoFileUploadedError() - assert error is not None + class TestFileUploadErrors: + """Test file upload error types for pipelines.""" - def test_too_many_files_error(self): - """Test TooManyFilesError can be raised.""" - error = TooManyFilesError() - assert error is not None + def test_no_file_uploaded_error(self): + """Test NoFileUploadedError can be raised.""" + error = NoFileUploadedError() + assert error is not None - def test_filename_not_exists_error(self): - """Test FilenameNotExistsError can be raised.""" - error = FilenameNotExistsError() - assert error is not None + def test_too_many_files_error(self): + """Test TooManyFilesError can be raised.""" + error = TooManyFilesError() + assert error is not None - def test_file_too_large_error(self): - """Test FileTooLargeError can be raised.""" - error = FileTooLargeServiceError("File exceeds size limit") - assert error is not None + def test_filename_not_exists_error(self): + """Test FilenameNotExistsError can be raised.""" + error = FilenameNotExistsError() + assert error is not None - def test_unsupported_file_type_error(self): - """Test UnsupportedFileTypeError can be raised.""" - error = UnsupportedFileTypeError() - assert error is not None + def test_file_too_large_error(self): + """Test FileTooLargeError can be raised.""" + error = FileTooLargeServiceError("File exceeds size limit") + assert error is not None + def test_unsupported_file_type_error(self): + """Test UnsupportedFileTypeError can be raised.""" + error = UnsupportedFileTypeError() + assert error is not None -class TestRagPipelineService: - """Test RagPipelineService interface.""" + class TestRagPipelineService: + """Test RagPipelineService interface.""" - def test_get_datasource_plugins_method_exists(self): - """Test RagPipelineService.get_datasource_plugins exists.""" - assert hasattr(RagPipelineService, "get_datasource_plugins") + def test_get_datasource_plugins_method_exists(self): + """Test RagPipelineService.get_datasource_plugins exists.""" + assert hasattr(RagPipelineService, "get_datasource_plugins") - def test_get_pipeline_method_exists(self): - """Test RagPipelineService.get_pipeline exists.""" - assert hasattr(RagPipelineService, "get_pipeline") + def test_get_pipeline_method_exists(self): + """Test RagPipelineService.get_pipeline exists.""" + assert hasattr(RagPipelineService, "get_pipeline") - def test_run_datasource_workflow_node_method_exists(self): - """Test RagPipelineService.run_datasource_workflow_node exists.""" - assert hasattr(RagPipelineService, "run_datasource_workflow_node") + def test_run_datasource_workflow_node_method_exists(self): + """Test RagPipelineService.run_datasource_workflow_node exists.""" + assert hasattr(RagPipelineService, "run_datasource_workflow_node") - def test_get_pipeline_templates_method_exists(self): - """Test RagPipelineService.get_pipeline_templates exists.""" - assert hasattr(RagPipelineService, "get_pipeline_templates") + def test_get_pipeline_templates_method_exists(self): + """Test RagPipelineService.get_pipeline_templates exists.""" + assert hasattr(RagPipelineService, "get_pipeline_templates") - def test_get_pipeline_template_detail_method_exists(self): - """Test RagPipelineService.get_pipeline_template_detail exists.""" - assert hasattr(RagPipelineService, "get_pipeline_template_detail") + def test_get_pipeline_template_detail_method_exists(self): + """Test RagPipelineService.get_pipeline_template_detail exists.""" + assert hasattr(RagPipelineService, "get_pipeline_template_detail") + class TestInvokeFrom: + """Test InvokeFrom enum for pipeline invocation.""" -class TestInvokeFrom: - """Test InvokeFrom enum for pipeline invocation.""" + def test_published_pipeline_invoke_from(self): + """Test PUBLISHED_PIPELINE InvokeFrom value exists.""" + assert hasattr(InvokeFrom, "PUBLISHED_PIPELINE") - def test_published_pipeline_invoke_from(self): - """Test PUBLISHED_PIPELINE InvokeFrom value exists.""" - assert hasattr(InvokeFrom, "PUBLISHED_PIPELINE") + def test_debugger_invoke_from(self): + """Test DEBUGGER InvokeFrom value exists.""" + assert hasattr(InvokeFrom, "DEBUGGER") - def test_debugger_invoke_from(self): - """Test DEBUGGER InvokeFrom value exists.""" - assert hasattr(InvokeFrom, "DEBUGGER") + class TestPipelineResponseModes: + """Test pipeline response mode patterns.""" + def test_streaming_mode(self): + """Test streaming response mode.""" + mode = "streaming" + valid_modes = ["streaming", "blocking"] + assert mode in valid_modes -class TestPipelineResponseModes: - """Test pipeline response mode patterns.""" + def test_blocking_mode(self): + """Test blocking response mode.""" + mode = "blocking" + valid_modes = ["streaming", "blocking"] + assert mode in valid_modes - def test_streaming_mode(self): - """Test streaming response mode.""" - mode = "streaming" - valid_modes = ["streaming", "blocking"] - assert mode in valid_modes + class TestDatasourceTypes: + """Test common datasource types for pipelines.""" - def test_blocking_mode(self): - """Test blocking response mode.""" - mode = "blocking" - valid_modes = ["streaming", "blocking"] - assert mode in valid_modes + @pytest.mark.parametrize("ds_type", ["online_document", "local_file", "website", "api", "database"]) + def test_datasource_type_valid(self, ds_type): + """Test common datasource types are strings.""" + assert isinstance(ds_type, str) + assert len(ds_type) > 0 + class TestPipelineFileUploadResponse: + """Test file upload response structure for pipelines.""" -class TestDatasourceTypes: - """Test common datasource types for pipelines.""" + def test_upload_response_fields(self): + """Test expected fields in upload response.""" + expected_fields = ["id", "name", "size", "extension", "mime_type", "created_by", "created_at"] - @pytest.mark.parametrize("ds_type", ["online_document", "local_file", "website", "api", "database"]) - def test_datasource_type_valid(self, ds_type): - """Test common datasource types are strings.""" - assert isinstance(ds_type, str) - assert len(ds_type) > 0 + # Create mock response + mock_response = { + "id": str(uuid.uuid4()), + "name": "document.pdf", + "size": 1024, + "extension": "pdf", + "mime_type": "application/pdf", + "created_by": str(uuid.uuid4()), + "created_at": "2024-01-01T00:00:00Z", + } + for field in expected_fields: + assert field in mock_response -class TestPipelineFileUploadResponse: - """Test file upload response structure for pipelines.""" + class TestPipelineNodeExecution: + """Test pipeline node execution patterns.""" - def test_upload_response_fields(self): - """Test expected fields in upload response.""" - expected_fields = ["id", "name", "size", "extension", "mime_type", "created_by", "created_at"] + def test_node_id_is_string(self): + """Test node_id is a string identifier.""" + node_id = "node_abc123" + assert isinstance(node_id, str) + assert len(node_id) > 0 - # Create mock response - mock_response = { - "id": str(uuid.uuid4()), - "name": "document.pdf", - "size": 1024, - "extension": "pdf", - "mime_type": "application/pdf", - "created_by": str(uuid.uuid4()), - "created_at": "2024-01-01T00:00:00Z", - } + def test_pipeline_id_is_uuid(self): + """Test pipeline_id is a valid UUID string.""" + pipeline_id = str(uuid.uuid4()) + assert len(pipeline_id) == 36 + assert "-" in pipeline_id - for field in expected_fields: - assert field in mock_response + class TestCredentialHandling: + """Test credential handling patterns.""" + def test_credential_id_is_optional(self): + """Test credential_id can be None.""" + payload = DatasourceNodeRunPayload( + inputs={}, datasource_type="local_file", is_published=True, credential_id=None + ) + assert payload.credential_id is None -class TestPipelineNodeExecution: - """Test pipeline node execution patterns.""" + def test_credential_id_can_be_provided(self): + """Test credential_id can be set.""" + payload = DatasourceNodeRunPayload( + inputs={}, datasource_type="api", is_published=True, credential_id="cred_oauth_123" + ) + assert payload.credential_id == "cred_oauth_123" - def test_node_id_is_string(self): - """Test node_id is a string identifier.""" - node_id = "node_abc123" - assert isinstance(node_id, str) - assert len(node_id) > 0 + class TestPublishedVsDraft: + """Test published vs draft pipeline patterns.""" - def test_pipeline_id_is_uuid(self): - """Test pipeline_id is a valid UUID string.""" - pipeline_id = str(uuid.uuid4()) - assert len(pipeline_id) == 36 - assert "-" in pipeline_id + def test_is_published_true(self): + """Test is_published=True for published pipelines.""" + payload = DatasourceNodeRunPayload(inputs={}, datasource_type="online_document", is_published=True) + assert payload.is_published is True + def test_is_published_false_for_draft(self): + """Test is_published=False for draft pipelines.""" + payload = DatasourceNodeRunPayload(inputs={}, datasource_type="online_document", is_published=False) + assert payload.is_published is False -class TestCredentialHandling: - """Test credential handling patterns.""" + class TestPipelineInputVariables: + """Test pipeline input variable patterns.""" - def test_credential_id_is_optional(self): - """Test credential_id can be None.""" - payload = DatasourceNodeRunPayload( - inputs={}, datasource_type="local_file", is_published=True, credential_id=None - ) - assert payload.credential_id is None + def test_inputs_as_dict(self): + """Test inputs are passed as dictionary.""" + inputs = {"url": "https://example.com/doc.pdf", "timeout": 30, "retry": True} + payload = DatasourceNodeRunPayload(inputs=inputs, datasource_type="online_document", is_published=True) + assert payload.inputs["url"] == "https://example.com/doc.pdf" + assert payload.inputs["timeout"] == 30 + assert payload.inputs["retry"] is True - def test_credential_id_can_be_provided(self): - """Test credential_id can be set.""" - payload = DatasourceNodeRunPayload( - inputs={}, datasource_type="api", is_published=True, credential_id="cred_oauth_123" - ) - assert payload.credential_id == "cred_oauth_123" + def test_inputs_with_list_values(self): + """Test inputs with list values.""" + inputs = {"urls": ["https://example.com/1", "https://example.com/2"], "tags": ["tag1", "tag2", "tag3"]} + payload = DatasourceNodeRunPayload(inputs=inputs, datasource_type="online_document", is_published=True) + assert len(payload.inputs["urls"]) == 2 + assert len(payload.inputs["tags"]) == 3 + # --------------------------------------------------------------------------- + # PipelineRunApiEntity / DatasourceNodeRunApiEntity Model Tests + # --------------------------------------------------------------------------- -class TestPublishedVsDraft: - """Test published vs draft pipeline patterns.""" + class TestPipelineRunApiEntity: + """Test PipelineRunApiEntity Pydantic model.""" - def test_is_published_true(self): - """Test is_published=True for published pipelines.""" - payload = DatasourceNodeRunPayload(inputs={}, datasource_type="online_document", is_published=True) - assert payload.is_published is True + def test_entity_with_all_fields(self): + """Test entity with all required fields.""" + entity = PipelineRunApiEntity( + inputs={"key": "value"}, + datasource_type="online_document", + datasource_info_list=[{"url": "https://example.com"}], + start_node_id="node_1", + is_published=True, + response_mode="streaming", + ) + assert entity.datasource_type == "online_document" + assert entity.response_mode == "streaming" + assert entity.is_published is True - def test_is_published_false_for_draft(self): - """Test is_published=False for draft pipelines.""" - payload = DatasourceNodeRunPayload(inputs={}, datasource_type="online_document", is_published=False) - assert payload.is_published is False + def test_entity_blocking_response_mode(self): + """Test entity with blocking response mode.""" + entity = PipelineRunApiEntity( + inputs={}, + datasource_type="local_file", + datasource_info_list=[], + start_node_id="node_start", + is_published=False, + response_mode="blocking", + ) + assert entity.response_mode == "blocking" + assert entity.is_published is False + def test_entity_missing_required_field(self): + """Test entity raises on missing required field.""" + with pytest.raises( + ValueError, match="4 validation errors for PipelineRunApiEntity\\\ndatasource_info_list" + ): + PipelineRunApiEntity.model_validate( + { + "inputs": {}, + "datasource_type": "online_document", + # missing datasource_info_list, start_node_id, etc. + } + ) -class TestPipelineInputVariables: - """Test pipeline input variable patterns.""" + class TestDatasourceNodeRunApiEntity: + """Test DatasourceNodeRunApiEntity Pydantic model.""" - def test_inputs_as_dict(self): - """Test inputs are passed as dictionary.""" - inputs = {"url": "https://example.com/doc.pdf", "timeout": 30, "retry": True} - payload = DatasourceNodeRunPayload(inputs=inputs, datasource_type="online_document", is_published=True) - assert payload.inputs["url"] == "https://example.com/doc.pdf" - assert payload.inputs["timeout"] == 30 - assert payload.inputs["retry"] is True + def test_entity_with_all_fields(self): + """Test entity with all fields.""" + entity = DatasourceNodeRunApiEntity( + pipeline_id=str(uuid.uuid4()), + node_id="node_abc", + inputs={"url": "https://example.com"}, + datasource_type="website", + is_published=True, + ) + assert entity.node_id == "node_abc" + assert entity.credential_id is None - def test_inputs_with_list_values(self): - """Test inputs with list values.""" - inputs = {"urls": ["https://example.com/1", "https://example.com/2"], "tags": ["tag1", "tag2", "tag3"]} - payload = DatasourceNodeRunPayload(inputs=inputs, datasource_type="online_document", is_published=True) - assert len(payload.inputs["urls"]) == 2 - assert len(payload.inputs["tags"]) == 3 + def test_entity_with_credential(self): + """Test entity with credential_id.""" + entity = DatasourceNodeRunApiEntity( + pipeline_id=str(uuid.uuid4()), + node_id="node_xyz", + inputs={}, + datasource_type="api", + credential_id="cred_123", + is_published=False, + ) + assert entity.credential_id == "cred_123" + # --------------------------------------------------------------------------- + # Endpoint Tests + # --------------------------------------------------------------------------- -# --------------------------------------------------------------------------- -# PipelineRunApiEntity / DatasourceNodeRunApiEntity Model Tests -# --------------------------------------------------------------------------- + class TestDatasourcePluginsApiGet: + """Tests for DatasourcePluginsApi.get(). + The original source delegates directly to ``RagPipelineService`` without + an inline dataset query, so no ``db`` patching is needed. + """ -class TestPipelineRunApiEntity: - """Test PipelineRunApiEntity Pydantic model.""" + @patch("controllers.service_api.dataset.rag_pipeline.rag_pipeline_workflow.RagPipelineService") + def test_get_plugins_success(self, mock_svc_cls, app: Flask, sqlite_session: Session): + """Test successful retrieval of datasource plugins.""" + tenant_id = str(uuid.uuid4()) + dataset_id = str(uuid.uuid4()) - def test_entity_with_all_fields(self): - """Test entity with all required fields.""" - entity = PipelineRunApiEntity( - inputs={"key": "value"}, - datasource_type="online_document", - datasource_info_list=[{"url": "https://example.com"}], - start_node_id="node_1", - is_published=True, - response_mode="streaming", - ) - assert entity.datasource_type == "online_document" - assert entity.response_mode == "streaming" - assert entity.is_published is True + _persist_dataset(sqlite_session, tenant_id=tenant_id, dataset_id=dataset_id) - def test_entity_blocking_response_mode(self): - """Test entity with blocking response mode.""" - entity = PipelineRunApiEntity( - inputs={}, - datasource_type="local_file", - datasource_info_list=[], - start_node_id="node_start", - is_published=False, - response_mode="blocking", - ) - assert entity.response_mode == "blocking" - assert entity.is_published is False - - def test_entity_missing_required_field(self): - """Test entity raises on missing required field.""" - with pytest.raises(ValueError, match="4 validation errors for PipelineRunApiEntity\\\ndatasource_info_list"): - PipelineRunApiEntity.model_validate( + datasource_plugins = [ { - "inputs": {}, + "node_id": "node-datasource-1", + "plugin_id": "plugin-a", + "provider_name": "provider-a", "datasource_type": "online_document", - # missing datasource_info_list, start_node_id, etc. + "title": "Online Docs", + "user_input_variables": [ + {"variable": "url", "label": "URL", "type": "text-input", "required": True} + ], + "credentials": [ + {"id": "cred-1", "name": "Default credential", "type": "oauth2", "is_default": True} + ], } + ] + mock_svc_instance = Mock() + mock_svc_instance.get_datasource_plugins.return_value = datasource_plugins + mock_svc_cls.return_value = mock_svc_instance + + with app.test_request_context("/datasets/test/pipeline/datasource-plugins?is_published=true"): + api = DatasourcePluginsApi() + response, status = api.get(tenant_id=tenant_id, dataset_id=dataset_id) + + assert status == 200 + assert response == datasource_plugins + mock_svc_instance.get_datasource_plugins.assert_called_once_with( + tenant_id=tenant_id, + dataset_id=dataset_id, + is_published=True, + credential_query=workflow_module.application_services().credential_queries, + datasource_providers=workflow_module.application_services().data_sources.providers, ) + @patch("controllers.service_api.dataset.rag_pipeline.rag_pipeline_workflow.RagPipelineService") + def test_get_plugins_parses_false_is_published_query(self, mock_svc_cls, app: Flask, sqlite_session: Session): + """Test false query string is parsed as boolean False.""" + tenant_id = str(uuid.uuid4()) + dataset_id = str(uuid.uuid4()) -class TestDatasourceNodeRunApiEntity: - """Test DatasourceNodeRunApiEntity Pydantic model.""" + _persist_dataset(sqlite_session, tenant_id=tenant_id, dataset_id=dataset_id) + mock_svc_instance = Mock() + mock_svc_instance.get_datasource_plugins.return_value = [] + mock_svc_cls.return_value = mock_svc_instance - def test_entity_with_all_fields(self): - """Test entity with all fields.""" - entity = DatasourceNodeRunApiEntity( - pipeline_id=str(uuid.uuid4()), - node_id="node_abc", - inputs={"url": "https://example.com"}, - datasource_type="website", - is_published=True, + with app.test_request_context("/datasets/test/pipeline/datasource-plugins?is_published=false"): + api = DatasourcePluginsApi() + response, status = api.get(tenant_id=tenant_id, dataset_id=dataset_id) + + assert status == 200 + assert response == [] + mock_svc_instance.get_datasource_plugins.assert_called_once_with( + tenant_id=tenant_id, + dataset_id=dataset_id, + is_published=False, + credential_query=workflow_module.application_services().credential_queries, + datasource_providers=workflow_module.application_services().data_sources.providers, + ) + + def test_get_plugins_not_found(self, app: Flask, sqlite_session: Session): + """Test NotFound when dataset check fails.""" + tenant_id = str(uuid.uuid4()) + dataset_id = str(uuid.uuid4()) + _persist_dataset(sqlite_session, tenant_id="other-tenant", dataset_id=dataset_id) + + with app.test_request_context("/datasets/test/pipeline/datasource-plugins"): + api = DatasourcePluginsApi() + with pytest.raises(NotFound): + api.get(tenant_id=tenant_id, dataset_id=dataset_id) + + @patch("controllers.service_api.dataset.rag_pipeline.rag_pipeline_workflow.RagPipelineService") + def test_get_plugins_empty_list(self, mock_svc_cls, app: Flask, sqlite_session: Session): + """Test empty plugin list.""" + tenant_id = str(uuid.uuid4()) + dataset_id = str(uuid.uuid4()) + _persist_dataset(sqlite_session, tenant_id=tenant_id, dataset_id=dataset_id) + mock_svc_instance = Mock() + mock_svc_instance.get_datasource_plugins.return_value = [] + mock_svc_cls.return_value = mock_svc_instance + + with app.test_request_context("/datasets/test/pipeline/datasource-plugins"): + api = DatasourcePluginsApi() + response, status = api.get(tenant_id=tenant_id, dataset_id=dataset_id) + + assert status == 200 + assert response == [] + + class TestDatasourceNodeRunApiPost: + """Tests for DatasourceNodeRunApi.post(). + + The source asserts ``isinstance(current_user, Account)`` and delegates to + ``RagPipelineService`` and ``PipelineGenerator``, so we patch those plus + ``current_user``. ``post`` is wrapped in ``@model_validate``, which parses + the JSON request body live, so payloads are supplied via + ``test_request_context(json=...)`` and validation runs before the dataset + ownership guard. + """ + + @patch("controllers.service_api.dataset.rag_pipeline.rag_pipeline_workflow.helper") + @patch("controllers.service_api.dataset.rag_pipeline.rag_pipeline_workflow.PipelineGenerator") + @patch( + "controllers.service_api.dataset.rag_pipeline.rag_pipeline_workflow.current_user", + new_callable=lambda: Account(name="Test Account", email="test@example.com"), ) - assert entity.node_id == "node_abc" - assert entity.credential_id is None + @patch("controllers.service_api.dataset.rag_pipeline.rag_pipeline_workflow.RagPipelineService") + def test_post_success( + self, + mock_svc_cls, + current_account, + mock_gen, + mock_helper, + app: Flask, + sqlite_session: Session, + ): + """Test successful datasource node run.""" + tenant_id = str(uuid.uuid4()) + dataset_id = str(uuid.uuid4()) + node_id = "node_abc" - def test_entity_with_credential(self): - """Test entity with credential_id.""" - entity = DatasourceNodeRunApiEntity( - pipeline_id=str(uuid.uuid4()), - node_id="node_xyz", - inputs={}, - datasource_type="api", - credential_id="cred_123", - is_published=False, + _persist_dataset(sqlite_session, tenant_id=tenant_id, dataset_id=dataset_id) + + pipeline = _persist_pipeline(sqlite_session, tenant_id=tenant_id) + mock_svc_instance = Mock() + mock_svc_instance.get_pipeline.return_value = pipeline + mock_svc_instance.run_datasource_workflow_node.return_value = iter(["event1"]) + mock_svc_cls.return_value = mock_svc_instance + + mock_gen.convert_to_event_stream.return_value = iter(["stream_event"]) + mock_helper.compact_generate_response.return_value = {"result": "ok"} + + with app.test_request_context( + "/datasets/test/pipeline/datasource/nodes/node_abc/run", + method="POST", + json={ + "inputs": {"url": "https://example.com"}, + "datasource_type": "online_document", + "is_published": True, + }, + ): + api = DatasourceNodeRunApi() + response = api.post(tenant_id=tenant_id, dataset_id=dataset_id, node_id=node_id) + + assert response == {"result": "ok"} + mock_svc_instance.get_pipeline.assert_called_once_with(tenant_id=tenant_id, dataset_id=dataset_id) + mock_svc_instance.get_pipeline.assert_called_once_with(tenant_id=tenant_id, dataset_id=dataset_id) + mock_svc_instance.run_datasource_workflow_node.assert_called_once() + + def test_post_not_found(self, app: Flask, sqlite_session: Session): + """Test NotFound when dataset check fails.""" + + # `@model_validate` parses the body before the ownership guard, so a + # valid payload is required to reach the NotFound branch. + with app.test_request_context( + "/datasets/test/pipeline/datasource/nodes/n1/run", + method="POST", + json={"inputs": {}, "datasource_type": "online_document", "is_published": True}, + ): + api = DatasourceNodeRunApi() + with pytest.raises(NotFound): + api.post(tenant_id=str(uuid.uuid4()), dataset_id=str(uuid.uuid4()), node_id="n1") + + @patch( + "controllers.service_api.dataset.rag_pipeline.rag_pipeline_workflow.current_user", + new="not_account", ) - assert entity.credential_id == "cred_123" + def test_post_fails_when_current_user_not_account(self, app: Flask, sqlite_session: Session): + """Test AssertionError when current_user is not an Account instance.""" + tenant_id = str(uuid.uuid4()) + dataset_id = str(uuid.uuid4()) + _persist_dataset(sqlite_session, tenant_id=tenant_id, dataset_id=dataset_id) + with app.test_request_context( + "/datasets/test/pipeline/datasource/nodes/n1/run", + method="POST", + json={"inputs": {}, "datasource_type": "local_file", "is_published": True}, + ): + api = DatasourceNodeRunApi() + with pytest.raises(AssertionError): + api.post(tenant_id=tenant_id, dataset_id=dataset_id, node_id="n1") -# --------------------------------------------------------------------------- -# Endpoint Tests -# --------------------------------------------------------------------------- + class TestPipelineRunApiPost: + """Tests for PipelineRunApi.post().""" + @patch("controllers.service_api.dataset.rag_pipeline.rag_pipeline_workflow.helper") + @patch("controllers.service_api.dataset.rag_pipeline.rag_pipeline_workflow.PipelineGenerateService") + @patch( + "controllers.service_api.dataset.rag_pipeline.rag_pipeline_workflow.current_user", + new_callable=lambda: Account(name="Test Account", email="test@example.com"), + ) + @patch("controllers.service_api.dataset.rag_pipeline.rag_pipeline_workflow.RagPipelineService") + @patch("controllers.service_api.dataset.rag_pipeline.rag_pipeline_workflow.service_api_ns") + def test_post_success_streaming( + self, + mock_ns, + mock_svc_cls, + mock_current_user, + mock_gen_svc, + mock_helper, + app: Flask, + sqlite_session: Session, + pipeline_application: PipelineGenerator, + ): + """Test successful pipeline run with streaming response.""" + tenant_id = str(uuid.uuid4()) + dataset_id = str(uuid.uuid4()) -class TestDatasourcePluginsApiGet: - """Tests for DatasourcePluginsApi.get(). + _persist_dataset(sqlite_session, tenant_id=tenant_id, dataset_id=dataset_id) - The original source delegates directly to ``RagPipelineService`` without - an inline dataset query, so no ``db`` patching is needed. - """ - - @patch("controllers.service_api.dataset.rag_pipeline.rag_pipeline_workflow.RagPipelineService") - def test_get_plugins_success(self, mock_svc_cls, app: Flask, sqlite_session: Session): - """Test successful retrieval of datasource plugins.""" - tenant_id = str(uuid.uuid4()) - dataset_id = str(uuid.uuid4()) - - _persist_dataset(sqlite_session, tenant_id=tenant_id, dataset_id=dataset_id) - - datasource_plugins = [ - { - "node_id": "node-datasource-1", - "plugin_id": "plugin-a", - "provider_name": "provider-a", - "datasource_type": "online_document", - "title": "Online Docs", - "user_input_variables": [{"variable": "url", "label": "URL", "type": "text-input", "required": True}], - "credentials": [{"id": "cred-1", "name": "Default credential", "type": "oauth2", "is_default": True}], - } - ] - mock_svc_instance = Mock() - mock_svc_instance.get_datasource_plugins.return_value = datasource_plugins - mock_svc_cls.return_value = mock_svc_instance - - with app.test_request_context("/datasets/test/pipeline/datasource-plugins?is_published=true"): - api = DatasourcePluginsApi() - response, status = api.get(tenant_id=tenant_id, dataset_id=dataset_id) - - assert status == 200 - assert response == datasource_plugins - mock_svc_instance.get_datasource_plugins.assert_called_once_with( - tenant_id=tenant_id, - dataset_id=dataset_id, - is_published=True, - credential_query=workflow_module.application_services().credential_queries, - datasource_providers=workflow_module.application_services().data_sources.providers, - ) - - @patch("controllers.service_api.dataset.rag_pipeline.rag_pipeline_workflow.RagPipelineService") - def test_get_plugins_parses_false_is_published_query(self, mock_svc_cls, app: Flask, sqlite_session: Session): - """Test false query string is parsed as boolean False.""" - tenant_id = str(uuid.uuid4()) - dataset_id = str(uuid.uuid4()) - - _persist_dataset(sqlite_session, tenant_id=tenant_id, dataset_id=dataset_id) - mock_svc_instance = Mock() - mock_svc_instance.get_datasource_plugins.return_value = [] - mock_svc_cls.return_value = mock_svc_instance - - with app.test_request_context("/datasets/test/pipeline/datasource-plugins?is_published=false"): - api = DatasourcePluginsApi() - response, status = api.get(tenant_id=tenant_id, dataset_id=dataset_id) - - assert status == 200 - assert response == [] - mock_svc_instance.get_datasource_plugins.assert_called_once_with( - tenant_id=tenant_id, - dataset_id=dataset_id, - is_published=False, - credential_query=workflow_module.application_services().credential_queries, - datasource_providers=workflow_module.application_services().data_sources.providers, - ) - - def test_get_plugins_not_found(self, app: Flask, sqlite_session: Session): - """Test NotFound when dataset check fails.""" - tenant_id = str(uuid.uuid4()) - dataset_id = str(uuid.uuid4()) - _persist_dataset(sqlite_session, tenant_id="other-tenant", dataset_id=dataset_id) - - with app.test_request_context("/datasets/test/pipeline/datasource-plugins"): - api = DatasourcePluginsApi() - with pytest.raises(NotFound): - api.get(tenant_id=tenant_id, dataset_id=dataset_id) - - @patch("controllers.service_api.dataset.rag_pipeline.rag_pipeline_workflow.RagPipelineService") - def test_get_plugins_empty_list(self, mock_svc_cls, app: Flask, sqlite_session: Session): - """Test empty plugin list.""" - tenant_id = str(uuid.uuid4()) - dataset_id = str(uuid.uuid4()) - _persist_dataset(sqlite_session, tenant_id=tenant_id, dataset_id=dataset_id) - mock_svc_instance = Mock() - mock_svc_instance.get_datasource_plugins.return_value = [] - mock_svc_cls.return_value = mock_svc_instance - - with app.test_request_context("/datasets/test/pipeline/datasource-plugins"): - api = DatasourcePluginsApi() - response, status = api.get(tenant_id=tenant_id, dataset_id=dataset_id) - - assert status == 200 - assert response == [] - - -class TestDatasourceNodeRunApiPost: - """Tests for DatasourceNodeRunApi.post(). - - The source asserts ``isinstance(current_user, Account)`` and delegates to - ``RagPipelineService`` and ``PipelineGenerator``, so we patch those plus - ``current_user``. ``post`` is wrapped in ``@model_validate``, which parses - the JSON request body live, so payloads are supplied via - ``test_request_context(json=...)`` and validation runs before the dataset - ownership guard. - """ - - @patch("controllers.service_api.dataset.rag_pipeline.rag_pipeline_workflow.helper") - @patch("controllers.service_api.dataset.rag_pipeline.rag_pipeline_workflow.PipelineGenerator") - @patch( - "controllers.service_api.dataset.rag_pipeline.rag_pipeline_workflow.current_user", - new_callable=lambda: Account(name="Test Account", email="test@example.com"), - ) - @patch("controllers.service_api.dataset.rag_pipeline.rag_pipeline_workflow.RagPipelineService") - def test_post_success( - self, - mock_svc_cls, - current_account, - mock_gen, - mock_helper, - app: Flask, - sqlite_session: Session, - ): - """Test successful datasource node run.""" - tenant_id = str(uuid.uuid4()) - dataset_id = str(uuid.uuid4()) - node_id = "node_abc" - - _persist_dataset(sqlite_session, tenant_id=tenant_id, dataset_id=dataset_id) - - pipeline = _persist_pipeline(sqlite_session, tenant_id=tenant_id) - mock_svc_instance = Mock() - mock_svc_instance.get_pipeline.return_value = pipeline - mock_svc_instance.run_datasource_workflow_node.return_value = iter(["event1"]) - mock_svc_cls.return_value = mock_svc_instance - - mock_gen.convert_to_event_stream.return_value = iter(["stream_event"]) - mock_helper.compact_generate_response.return_value = {"result": "ok"} - - with app.test_request_context( - "/datasets/test/pipeline/datasource/nodes/node_abc/run", - method="POST", - json={ - "inputs": {"url": "https://example.com"}, + mock_ns.payload = { + "inputs": {"key": "val"}, "datasource_type": "online_document", + "datasource_info_list": [], + "start_node_id": "node_1", "is_published": True, - }, - ): - api = DatasourceNodeRunApi() - response = api.post(tenant_id=tenant_id, dataset_id=dataset_id, node_id=node_id) + "response_mode": "streaming", + } - assert response == {"result": "ok"} - mock_svc_instance.get_pipeline.assert_called_once_with(tenant_id=tenant_id, dataset_id=dataset_id) - mock_svc_instance.get_pipeline.assert_called_once_with(tenant_id=tenant_id, dataset_id=dataset_id) - mock_svc_instance.run_datasource_workflow_node.assert_called_once() + pipeline = _persist_pipeline(sqlite_session, tenant_id=tenant_id) + mock_svc_instance = Mock() + mock_svc_instance.get_pipeline.return_value = pipeline + mock_svc_cls.return_value = mock_svc_instance - def test_post_not_found(self, app: Flask, sqlite_session: Session): - """Test NotFound when dataset check fails.""" + mock_gen_svc.generate.return_value = {"result": "ok"} + mock_helper.compact_generate_response.return_value = {"result": "ok"} - # `@model_validate` parses the body before the ownership guard, so a - # valid payload is required to reach the NotFound branch. - with app.test_request_context( - "/datasets/test/pipeline/datasource/nodes/n1/run", - method="POST", - json={"inputs": {}, "datasource_type": "online_document", "is_published": True}, - ): - api = DatasourceNodeRunApi() - with pytest.raises(NotFound): - api.post(tenant_id=str(uuid.uuid4()), dataset_id=str(uuid.uuid4()), node_id="n1") + with app.test_request_context("/datasets/test/pipeline/run", method="POST"): + api = PipelineRunApi() + response = api.post.__wrapped__(api, sqlite_session, tenant_id=tenant_id, dataset_id=dataset_id) - @patch( - "controllers.service_api.dataset.rag_pipeline.rag_pipeline_workflow.current_user", - new="not_account", - ) - def test_post_fails_when_current_user_not_account(self, app: Flask, sqlite_session: Session): - """Test AssertionError when current_user is not an Account instance.""" - tenant_id = str(uuid.uuid4()) - dataset_id = str(uuid.uuid4()) - _persist_dataset(sqlite_session, tenant_id=tenant_id, dataset_id=dataset_id) + assert response == {"result": "ok"} + mock_svc_cls.assert_called_once_with(sqlite_session) + mock_gen_svc.generate.assert_called_once() + assert mock_gen_svc.generate.call_args.kwargs["generator"] is pipeline_application + assert mock_gen_svc.generate.call_args.kwargs["session"] is sqlite_session - with app.test_request_context( - "/datasets/test/pipeline/datasource/nodes/n1/run", - method="POST", - json={"inputs": {}, "datasource_type": "local_file", "is_published": True}, - ): - api = DatasourceNodeRunApi() - with pytest.raises(AssertionError): - api.post(tenant_id=tenant_id, dataset_id=dataset_id, node_id="n1") + def test_post_not_found(self, app: Flask, sqlite_session: Session): + """Test NotFound when dataset check fails.""" + with app.test_request_context("/datasets/test/pipeline/run", method="POST"): + api = PipelineRunApi() + with pytest.raises(NotFound): + api.post.__wrapped__( + api, + sqlite_session, + tenant_id=str(uuid.uuid4()), + dataset_id=str(uuid.uuid4()), + ) + @patch("controllers.service_api.dataset.rag_pipeline.rag_pipeline_workflow.current_user", new="not_account") + @patch("controllers.service_api.dataset.rag_pipeline.rag_pipeline_workflow.service_api_ns") + def test_post_forbidden_non_account_user(self, mock_ns, app: Flask, sqlite_session: Session): + """Test Forbidden when current_user is not an Account.""" + tenant_id = str(uuid.uuid4()) + dataset_id = str(uuid.uuid4()) + _persist_dataset(sqlite_session, tenant_id=tenant_id, dataset_id=dataset_id) + mock_ns.payload = { + "inputs": {}, + "datasource_type": "online_document", + "datasource_info_list": [], + "start_node_id": "node_1", + "is_published": True, + "response_mode": "blocking", + } -class TestPipelineRunApiPost: - """Tests for PipelineRunApi.post().""" + with app.test_request_context("/datasets/test/pipeline/run", method="POST"): + api = PipelineRunApi() + with pytest.raises(Forbidden): + api.post.__wrapped__( + api, + sqlite_session, + tenant_id=tenant_id, + dataset_id=dataset_id, + ) - @patch("controllers.service_api.dataset.rag_pipeline.rag_pipeline_workflow.helper") - @patch("controllers.service_api.dataset.rag_pipeline.rag_pipeline_workflow.PipelineGenerateService") - @patch( - "controllers.service_api.dataset.rag_pipeline.rag_pipeline_workflow.current_user", - new_callable=lambda: Account(name="Test Account", email="test@example.com"), - ) - @patch("controllers.service_api.dataset.rag_pipeline.rag_pipeline_workflow.RagPipelineService") - @patch("controllers.service_api.dataset.rag_pipeline.rag_pipeline_workflow.service_api_ns") - def test_post_success_streaming( - self, - mock_ns, - mock_svc_cls, - mock_current_user, - mock_gen_svc, - mock_helper, - app: Flask, - sqlite_session: Session, - pipeline_application: PipelineGenerator, - ): - """Test successful pipeline run with streaming response.""" - tenant_id = str(uuid.uuid4()) - dataset_id = str(uuid.uuid4()) + class TestFileUploadApiPost: + """Tests for KnowledgebasePipelineFileUploadApi.post().""" - _persist_dataset(sqlite_session, tenant_id=tenant_id, dataset_id=dataset_id) - - mock_ns.payload = { - "inputs": {"key": "val"}, - "datasource_type": "online_document", - "datasource_info_list": [], - "start_node_id": "node_1", - "is_published": True, - "response_mode": "streaming", - } - - pipeline = _persist_pipeline(sqlite_session, tenant_id=tenant_id) - mock_svc_instance = Mock() - mock_svc_instance.get_pipeline.return_value = pipeline - mock_svc_cls.return_value = mock_svc_instance - - mock_gen_svc.generate.return_value = {"result": "ok"} - mock_helper.compact_generate_response.return_value = {"result": "ok"} - - with app.test_request_context("/datasets/test/pipeline/run", method="POST"): - api = PipelineRunApi() - response = api.post.__wrapped__(api, sqlite_session, tenant_id=tenant_id, dataset_id=dataset_id) - - assert response == {"result": "ok"} - mock_svc_cls.assert_called_once_with(sqlite_session) - mock_gen_svc.generate.assert_called_once() - assert mock_gen_svc.generate.call_args.kwargs["generator"] is pipeline_application - assert mock_gen_svc.generate.call_args.kwargs["session"] is sqlite_session - - def test_post_not_found(self, app: Flask, sqlite_session: Session): - """Test NotFound when dataset check fails.""" - with app.test_request_context("/datasets/test/pipeline/run", method="POST"): - api = PipelineRunApi() - with pytest.raises(NotFound): - api.post.__wrapped__( - api, - sqlite_session, - tenant_id=str(uuid.uuid4()), - dataset_id=str(uuid.uuid4()), - ) - - @patch("controllers.service_api.dataset.rag_pipeline.rag_pipeline_workflow.current_user", new="not_account") - @patch("controllers.service_api.dataset.rag_pipeline.rag_pipeline_workflow.service_api_ns") - def test_post_forbidden_non_account_user(self, mock_ns, app: Flask, sqlite_session: Session): - """Test Forbidden when current_user is not an Account.""" - tenant_id = str(uuid.uuid4()) - dataset_id = str(uuid.uuid4()) - _persist_dataset(sqlite_session, tenant_id=tenant_id, dataset_id=dataset_id) - mock_ns.payload = { - "inputs": {}, - "datasource_type": "online_document", - "datasource_info_list": [], - "start_node_id": "node_1", - "is_published": True, - "response_mode": "blocking", - } - - with app.test_request_context("/datasets/test/pipeline/run", method="POST"): - api = PipelineRunApi() - with pytest.raises(Forbidden): - api.post.__wrapped__( - api, - sqlite_session, - tenant_id=tenant_id, - dataset_id=dataset_id, - ) - - -class TestFileUploadApiPost: - """Tests for KnowledgebasePipelineFileUploadApi.post().""" - - @patch("controllers.service_api.dataset.rag_pipeline.rag_pipeline_workflow.FileService") - @patch( - "controllers.service_api.dataset.rag_pipeline.rag_pipeline_workflow.current_user", - new_callable=lambda: Account(name="Upload User", email="upload@example.com"), - ) - def test_upload_success(self, current_account, mock_file_svc_cls, app: Flask, sqlite_engine): - """Test successful file upload.""" - current_account.id = str(uuid.uuid4()) - upload = UploadFile( - tenant_id=str(uuid.uuid4()), - storage_type=StorageType.LOCAL, - key="pipeline/doc.pdf", - name="doc.pdf", - size=1024, - extension="pdf", - mime_type="application/pdf", - created_by_role=CreatorUserRole.ACCOUNT, - created_by=current_account.id, - created_at=datetime(2024, 1, 1, tzinfo=UTC), - used=False, + @patch("controllers.service_api.dataset.rag_pipeline.rag_pipeline_workflow.FileService") + @patch( + "controllers.service_api.dataset.rag_pipeline.rag_pipeline_workflow.current_user", + new_callable=lambda: Account(name="Upload User", email="upload@example.com"), ) + def test_upload_success(self, current_account, mock_file_svc_cls, app: Flask, sqlite_engine): + """Test successful file upload.""" + current_account.id = str(uuid.uuid4()) + upload = UploadFile( + tenant_id=str(uuid.uuid4()), + storage_type=StorageType.LOCAL, + key="pipeline/doc.pdf", + name="doc.pdf", + size=1024, + extension="pdf", + mime_type="application/pdf", + created_by_role=CreatorUserRole.ACCOUNT, + created_by=current_account.id, + created_at=datetime(2024, 1, 1, tzinfo=UTC), + used=False, + ) - mock_file_svc_instance = Mock() - mock_file_svc_instance.upload_file.return_value = upload - mock_file_svc_cls.return_value = mock_file_svc_instance + mock_file_svc_instance = Mock() + mock_file_svc_instance.upload_file.return_value = upload + mock_file_svc_cls.return_value = mock_file_svc_instance - file_data = FileStorage( - stream=io.BytesIO(b"fake pdf content"), - filename="doc.pdf", - content_type="application/pdf", + file_data = FileStorage( + stream=io.BytesIO(b"fake pdf content"), + filename="doc.pdf", + content_type="application/pdf", + ) + + with ( + patch.object(workflow_module, "db", SimpleNamespace(engine=sqlite_engine)), + app.test_request_context( + "/datasets/pipeline/file-upload", + method="POST", + content_type="multipart/form-data", + data={"file": file_data}, + ), + ): + response, status = KnowledgebasePipelineFileUploadApi().post(tenant_id=str(uuid.uuid4())) + + assert status == 201 + assert response["name"] == "doc.pdf" + assert response["extension"] == "pdf" + + @patch( + "controllers.service_api.dataset.rag_pipeline.rag_pipeline_workflow.FeatureService" + ".get_knowledge_file_size_limit", + return_value=15, ) + @patch("controllers.service_api.dataset.rag_pipeline.rag_pipeline_workflow.FileService") + @patch( + "controllers.service_api.dataset.rag_pipeline.rag_pipeline_workflow.current_user", + new_callable=lambda: Account(name="Upload User", email="upload@example.com"), + ) + def test_upload_file_too_large_returns_http_413( + self, current_account, mock_file_svc_cls, mock_get_limit, app: Flask, sqlite_engine + ): + mock_file_svc_cls.return_value.upload_file.side_effect = FileTooLargeServiceError() + file_data = FileStorage( + stream=io.BytesIO(b"oversized content"), + filename="doc.pdf", + content_type="application/pdf", + ) - with ( - patch.object(workflow_module, "db", SimpleNamespace(engine=sqlite_engine)), - app.test_request_context( + with ( + patch.object(workflow_module, "db", SimpleNamespace(engine=sqlite_engine)), + app.test_request_context( + "/datasets/pipeline/file-upload", + method="POST", + content_type="multipart/form-data", + data={"file": file_data}, + ), + ): + with pytest.raises(FileTooLargeHTTPError) as exc_info: + KnowledgebasePipelineFileUploadApi().post(tenant_id="tenant-1") + + assert exc_info.value.code == 413 + assert exc_info.value.error_code == "file_too_large" + mock_get_limit.assert_called_once_with("tenant-1") + + @patch("controllers.service_api.dataset.rag_pipeline.rag_pipeline_workflow.FileService") + @patch( + "controllers.service_api.dataset.rag_pipeline.rag_pipeline_workflow.current_user", + new_callable=lambda: Account(name="Upload User", email="upload@example.com"), + ) + def test_upload_unsupported_file_type_returns_http_415( + self, current_account, mock_file_svc_cls, app: Flask, sqlite_engine + ): + # Regression: the service error once shadowed the HTTP error of the same name, + # so this surfaced as a 400 `invalid_param` instead of the documented 415. + mock_file_svc_cls.return_value.upload_file.side_effect = UnsupportedFileTypeError() + file_data = FileStorage( + stream=io.BytesIO(b"binary"), + filename="doc.exe", + content_type="application/octet-stream", + ) + + with ( + patch.object(workflow_module, "db", SimpleNamespace(engine=sqlite_engine)), + patch.object(workflow_module.FeatureService, "get_knowledge_file_size_limit", return_value=15), + app.test_request_context( + "/datasets/pipeline/file-upload", + method="POST", + content_type="multipart/form-data", + data={"file": file_data}, + ), + ): + with pytest.raises(UnsupportedFileTypeHTTPError) as exc_info: + KnowledgebasePipelineFileUploadApi().post(tenant_id="tenant-1") + + assert exc_info.value.code == 415 + assert exc_info.value.error_code == "unsupported_file_type" + + def test_upload_no_file(self, app: Flask): + """Test error when no file is uploaded.""" + with app.test_request_context( "/datasets/pipeline/file-upload", method="POST", content_type="multipart/form-data", - data={"file": file_data}, - ), - ): - response, status = KnowledgebasePipelineFileUploadApi().post(tenant_id=str(uuid.uuid4())) + ): + api = KnowledgebasePipelineFileUploadApi() + with pytest.raises(NoFileUploadedError): + api.post(tenant_id=str(uuid.uuid4())) - assert status == 201 - assert response["name"] == "doc.pdf" - assert response["extension"] == "pdf" - @patch( - "controllers.service_api.dataset.rag_pipeline.rag_pipeline_workflow.FeatureService" - ".get_knowledge_file_size_limit", - return_value=15, +""" +Unit tests for Service API knowledge pipeline file-upload serialization. +""" + + +def _upload_file(*, created_at: datetime) -> UploadFile: + return make_upload_file( + file_id="file-1", + name="test.pdf", + size=123, + extension="pdf", + mime_type="application/pdf", + created_at=created_at, ) - @patch("controllers.service_api.dataset.rag_pipeline.rag_pipeline_workflow.FileService") - @patch( - "controllers.service_api.dataset.rag_pipeline.rag_pipeline_workflow.current_user", - new_callable=lambda: Account(name="Upload User", email="upload@example.com"), - ) - def test_upload_file_too_large_returns_http_413( - self, current_account, mock_file_svc_cls, mock_get_limit, app: Flask, sqlite_engine - ): - mock_file_svc_cls.return_value.upload_file.side_effect = FileTooLargeServiceError() - file_data = FileStorage( - stream=io.BytesIO(b"oversized content"), - filename="doc.pdf", - content_type="application/pdf", + + +class TestRagPipelineFileUploadSerializationCases: + @staticmethod + def test_file_upload_created_at_is_isoformat_string(): + created_at = datetime(2026, 2, 8, 12, 0, 0, tzinfo=UTC) + upload_file = _upload_file(created_at=created_at) + + result = dump_response(PipelineUploadFileResponse, upload_file) + assert result["created_at"] == created_at.isoformat() + + @staticmethod + def test_file_upload_created_at_none_serializes_to_null(): + upload_file = _upload_file(created_at=datetime(2026, 2, 8, 12, 0, 0, tzinfo=UTC)) + upload_file.created_at = None + + result = dump_response(PipelineUploadFileResponse, upload_file) + assert result["created_at"] is None + + +""" +Unit tests for Service API knowledge pipeline route registration. +""" + + +class TestRagPipelineRouteRegistrationCases: + @staticmethod + def test_rag_pipeline_routes_registered(): + api_dir = Path(__file__).resolve().parents[6] + + service_api_init = api_dir / "controllers" / "service_api" / "__init__.py" + rag_pipeline_workflow = ( + api_dir / "controllers" / "service_api" / "dataset" / "rag_pipeline" / "rag_pipeline_workflow.py" ) - with ( - patch.object(workflow_module, "db", SimpleNamespace(engine=sqlite_engine)), - app.test_request_context( - "/datasets/pipeline/file-upload", - method="POST", - content_type="multipart/form-data", - data={"file": file_data}, - ), - ): - with pytest.raises(FileTooLargeHTTPError) as exc_info: - KnowledgebasePipelineFileUploadApi().post(tenant_id="tenant-1") + assert service_api_init.exists() + assert rag_pipeline_workflow.exists() - assert exc_info.value.code == 413 - assert exc_info.value.error_code == "file_too_large" - mock_get_limit.assert_called_once_with("tenant-1") - - @patch("controllers.service_api.dataset.rag_pipeline.rag_pipeline_workflow.FileService") - @patch( - "controllers.service_api.dataset.rag_pipeline.rag_pipeline_workflow.current_user", - new_callable=lambda: Account(name="Upload User", email="upload@example.com"), - ) - def test_upload_unsupported_file_type_returns_http_415( - self, current_account, mock_file_svc_cls, app: Flask, sqlite_engine - ): - # Regression: the service error once shadowed the HTTP error of the same name, - # so this surfaced as a 400 `invalid_param` instead of the documented 415. - mock_file_svc_cls.return_value.upload_file.side_effect = UnsupportedFileTypeError() - file_data = FileStorage( - stream=io.BytesIO(b"binary"), - filename="doc.exe", - content_type="application/octet-stream", + init_tree = ast.parse(service_api_init.read_text(encoding="utf-8")) + import_found = False + for node in ast.walk(init_tree): + if not isinstance(node, ast.ImportFrom): + continue + if node.module != "dataset.rag_pipeline" or node.level != 1: + continue + if any(alias.name == "rag_pipeline_workflow" for alias in node.names): + import_found = True + break + assert import_found, ( + "from .dataset.rag_pipeline import rag_pipeline_workflow not found in service_api/__init__.py" ) - with ( - patch.object(workflow_module, "db", SimpleNamespace(engine=sqlite_engine)), - patch.object(workflow_module.FeatureService, "get_knowledge_file_size_limit", return_value=15), - app.test_request_context( - "/datasets/pipeline/file-upload", - method="POST", - content_type="multipart/form-data", - data={"file": file_data}, - ), - ): - with pytest.raises(UnsupportedFileTypeHTTPError) as exc_info: - KnowledgebasePipelineFileUploadApi().post(tenant_id="tenant-1") + workflow_tree = ast.parse(rag_pipeline_workflow.read_text(encoding="utf-8")) + route_paths: set[str] = set() - assert exc_info.value.code == 415 - assert exc_info.value.error_code == "unsupported_file_type" + for node in ast.walk(workflow_tree): + if not isinstance(node, ast.ClassDef): + continue + for decorator in node.decorator_list: + if not isinstance(decorator, ast.Call): + continue + if not isinstance(decorator.func, ast.Attribute): + continue + if decorator.func.attr != "route": + continue + if not decorator.args: + continue + first_arg = decorator.args[0] + if isinstance(first_arg, ast.Constant) and isinstance(first_arg.value, str): + route_paths.add(first_arg.value) - def test_upload_no_file(self, app: Flask): - """Test error when no file is uploaded.""" - with app.test_request_context( - "/datasets/pipeline/file-upload", - method="POST", - content_type="multipart/form-data", - ): - api = KnowledgebasePipelineFileUploadApi() - with pytest.raises(NoFileUploadedError): - api.post(tenant_id=str(uuid.uuid4())) + assert "/datasets//pipeline/datasource-plugins" in route_paths + assert "/datasets//pipeline/datasource/nodes//run" in route_paths + assert "/datasets//pipeline/run" in route_paths + assert "/datasets/pipeline/file-upload" in route_paths diff --git a/api/tests/unit_tests/controllers/service_api/dataset/test_dataset.py b/api/tests/unit_tests/controllers/service_api/dataset/test_dataset.py new file mode 100644 index 00000000000..f94799f04d6 --- /dev/null +++ b/api/tests/unit_tests/controllers/service_api/dataset/test_dataset.py @@ -0,0 +1,1495 @@ +import uuid +from datetime import UTC, datetime +from inspect import unwrap +from types import SimpleNamespace +from typing import Literal, cast +from unittest.mock import MagicMock, patch + +import pytest +from flask import Flask, request +from sqlalchemy.orm import Session +from werkzeug.exceptions import Forbidden, NotFound, UnprocessableEntity + +import services.errors.base +from controllers.service_api.dataset.dataset import ( + DatasetCreatePayload, + DatasetListQuery, + DatasetUpdatePayload, + TagBindingPayload, + TagCreatePayload, + TagDeletePayload, + TagUnbindingPayload, + TagUpdatePayload, +) +from controllers.service_api.dataset.error import DatasetInUseError, DatasetNameDuplicateError, InvalidActionError +from models.account import Account, Tenant, TenantAccountRole +from models.dataset import AppDatasetJoin, Dataset, DatasetMetadata, DatasetPermissionEnum, Document +from models.enums import PermissionEnum, TagType +from models.model import App, Tag, TagBinding + +"""Unit tests for Service API dataset controller behavior. + +Service boundaries stay mocked, while ORM collaborators are concrete model instances +persisted through the shared SQLite session fixture. +""" + + +# --------------------------------------------------------------------------- +# Helpers +# --------------------------------------------------------------------------- +DATASET_MODEL_TABLES = ( + Account, + Tenant, + Dataset, + Document, + App, + AppDatasetJoin, + DatasetMetadata, + Tag, + TagBinding, +) + + +def make_dataset( + session: Session, + tenant: Tenant, + account: Account, + **overrides: object, +) -> Dataset: + """Create and flush a real dataset so its database-backed properties can be serialized.""" + + base: dict[str, object] = { + "id": str(uuid.uuid4()), + "tenant_id": tenant.id, + "name": "Dataset", + "description": "desc", + "provider": "vendor", + "permission": PermissionEnum.ONLY_ME, + "data_source_type": None, + "indexing_technique": "economy", + "created_by": account.id, + "created_at": datetime(2024, 1, 1, 12, 0, 0, tzinfo=UTC), + "updated_by": None, + "updated_at": datetime(2024, 1, 1, 12, 0, 0, tzinfo=UTC), + "embedding_model": None, + "embedding_model_provider": None, + "retrieval_model": None, + "summary_index_setting": None, + "built_in_field_enabled": False, + "pipeline_id": None, + "runtime_mode": "general", + "chunk_structure": None, + "icon_info": None, + "enable_api": False, + "is_multimodal": False, + } + base.update(overrides) + created_at = cast(datetime, base.pop("created_at")) + updated_at = cast(datetime, base.pop("updated_at")) + dataset = Dataset(**base) # type: ignore[bad-argument-type] + dataset.created_at = created_at + dataset.updated_at = updated_at + session.add(dataset) + session.flush() + return dataset + + +DATASET_DETAIL_KEYS = { + "id", + "name", + "description", + "provider", + "permission", + "data_source_type", + "indexing_technique", + "app_count", + "document_count", + "word_count", + "created_by", + "author_name", + "created_at", + "updated_by", + "updated_at", + "embedding_model", + "embedding_model_provider", + "embedding_available", + "retrieval_model_dict", + "summary_index_setting", + "tags", + "doc_form", + "external_knowledge_info", + "external_retrieval_model", + "doc_metadata", + "built_in_field_enabled", + "pipeline_id", + "runtime_mode", + "chunk_structure", + "icon_info", + "is_published", + "total_documents", + "total_available_documents", + "enable_api", + "is_multimodal", + "maintainer", +} + + +def assert_dataset_detail_shape(response: dict[str, object], *, with_partial_members: bool = False) -> None: + expected_keys = set(DATASET_DETAIL_KEYS) + if with_partial_members: + expected_keys.add("partial_member_list") + assert set(response) == expected_keys + assert isinstance(response["created_at"], int) + assert isinstance(response["updated_at"], int) + retrieval_model = response["retrieval_model_dict"] + assert isinstance(retrieval_model, dict) + assert set(retrieval_model) == { + "search_method", + "reranking_enable", + "reranking_mode", + "reranking_model", + "weights", + "top_k", + "score_threshold_enabled", + "score_threshold", + } + external_retrieval_model = response["external_retrieval_model"] + if external_retrieval_model is not None: + assert isinstance(external_retrieval_model, dict) + assert set(external_retrieval_model) == { + "top_k", + "score_threshold", + "score_threshold_enabled", + } + if not with_partial_members: + assert "partial_member_list" not in response + + +class TestDatasetApisCases: + pytestmark = pytest.mark.parametrize("sqlite_session", [DATASET_MODEL_TABLES], indirect=True) + + @staticmethod + @pytest.fixture + def controller_session(sqlite_session: Session) -> Session: + """Expose the shared SQLite session under the controller-focused fixture name.""" + return sqlite_session + + @staticmethod + @pytest.fixture + def tenant(controller_session: Session) -> Tenant: + tenant = Tenant(name="Dataset API Tenant") + controller_session.add(tenant) + controller_session.flush() + return tenant + + @staticmethod + @pytest.fixture + def account(controller_session: Session, tenant: Tenant, monkeypatch: pytest.MonkeyPatch) -> Account: + account = Account(name="Dataset API User", email=f"dataset-api-{uuid.uuid4()}@example.com") + account.role = TenantAccountRole.OWNER + account._current_tenant = tenant + controller_session.add(account) + controller_session.flush() + + # Inject the concrete account at the controller boundary without relying on Flask-Login globals. + from controllers.service_api.dataset import dataset as dataset_module + + monkeypatch.setattr(dataset_module, "current_user", account) + return account + + @staticmethod + @pytest.fixture + def dataset(controller_session: Session, tenant: Tenant, account: Account) -> Dataset: + return make_dataset(controller_session, tenant, account) + + # --------------------------------------------------------------------------- + # API endpoint tests — DatasetListApi + # --------------------------------------------------------------------------- + + @pytest.mark.usefixtures("app_query_services") + class TestDatasetListApiGet: + """Test suite for DatasetListApi.get() endpoint.""" + + @patch("controllers.service_api.dataset.dataset.create_plugin_provider_manager") + @patch("controllers.service_api.dataset.dataset.DatasetService") + def test_list_datasets_success( + self, + mock_dataset_svc: MagicMock, + mock_provider_mgr: MagicMock, + app: Flask, + account: Account, + tenant: Tenant, + controller_session: Session, + ) -> None: + from controllers.service_api.dataset.dataset import DatasetListApi + + mock_dataset_svc.get_datasets.return_value = ([make_dataset(controller_session, tenant, account)], 1) + mock_provider_mgr.return_value.get_configurations.return_value.get_models.return_value = list[object]() + + with app.test_request_context("/datasets?page=1&limit=20", method="GET"): + api = DatasetListApi() + response, status = unwrap(api.get)(api, controller_session, tenant_id=tenant.id) + + assert status == 200 + assert set(response) == {"data", "has_more", "limit", "total", "page"} + assert response["has_more"] is False + assert response["limit"] == 20 + assert response["total"] == 1 + assert response["page"] == 1 + assert len(response["data"]) == 1 + assert_dataset_detail_shape(response["data"][0]) + + @patch("controllers.service_api.dataset.dataset.create_plugin_provider_manager") + @patch("controllers.service_api.dataset.dataset.DatasetService") + def test_list_datasets_preserves_repeated_tag_ids( + self, + mock_dataset_svc: MagicMock, + mock_provider_mgr: MagicMock, + app: Flask, + account: Account, + tenant: Tenant, + controller_session: Session, + ) -> None: + from controllers.service_api.dataset.dataset import DatasetListApi + + mock_dataset_svc.get_datasets.return_value = ([make_dataset(controller_session, tenant, account)], 1) + mock_provider_mgr.return_value.get_configurations.return_value.get_models.return_value = list[object]() + + with app.test_request_context("/datasets?tag_ids=tag-a&tag_ids=tag-b", method="GET"): + api = DatasetListApi() + response, status = unwrap(api.get)(api, controller_session, tenant_id=tenant.id) + page, limit, session, tenant_id, user, keyword, tag_ids, include_all = ( + mock_dataset_svc.get_datasets.call_args.args + ) + assert user is account + + assert status == 200 + assert response["total"] == 1 + assert (page, limit, session, tenant_id, keyword, tag_ids, include_all) == ( + 1, + 20, + controller_session, + tenant.id, + None, + ["tag-a", "tag-b"], + False, + ) + + @patch("controllers.service_api.dataset.dataset.create_plugin_provider_manager") + @patch("controllers.service_api.dataset.dataset.DatasetService") + def test_list_datasets_has_more_false_on_last_page_exact_limit( + self, + mock_dataset_svc: MagicMock, + mock_provider_mgr: MagicMock, + app: Flask, + account: Account, + tenant: Tenant, + controller_session: Session, + ) -> None: + """A full last page must set has_more false instead of forcing another fetch.""" + from controllers.service_api.dataset.dataset import DatasetListApi + + page_size = 20 + dataset = make_dataset(controller_session, tenant, account) + mock_dataset_svc.get_datasets.return_value = ([dataset] * page_size, page_size) + mock_provider_mgr.return_value.get_configurations.return_value.get_models.return_value = list[object]() + + with app.test_request_context(f"/datasets?page=1&limit={page_size}", method="GET"): + api = DatasetListApi() + response, status = unwrap(api.get)(api, controller_session, tenant_id=tenant.id) + + assert status == 200 + assert response["has_more"] is False + assert response["limit"] == page_size + assert response["total"] == page_size + assert response["page"] == 1 + + @patch("controllers.service_api.dataset.dataset.create_plugin_provider_manager") + @patch("controllers.service_api.dataset.dataset.DatasetService") + def test_list_datasets_has_more_true_when_limit_exceeds_cap( + self, + mock_dataset_svc: MagicMock, + mock_provider_mgr: MagicMock, + app: Flask, + account: Account, + tenant: Tenant, + controller_session: Session, + ) -> None: + """limit>100 still reports remaining rows after the server cap of 100.""" + from controllers.service_api.dataset.dataset import DatasetListApi + + returned_count = 100 + total = 150 + dataset = make_dataset(controller_session, tenant, account) + mock_dataset_svc.get_datasets.return_value = ([dataset] * returned_count, total) + mock_provider_mgr.return_value.get_configurations.return_value.get_models.return_value = list[object]() + + with app.test_request_context("/datasets?page=1&limit=200", method="GET"): + api = DatasetListApi() + response, status = unwrap(api.get)(api, controller_session, tenant_id=tenant.id) + + assert status == 200 + assert response["has_more"] is True + assert response["limit"] == 100 + assert response["total"] == total + assert response["page"] == 1 + assert mock_dataset_svc.get_datasets.call_args.args[1] == 100 + + @patch("controllers.service_api.dataset.dataset.create_plugin_provider_manager") + @patch("controllers.service_api.dataset.dataset.DatasetService") + def test_list_datasets_reports_the_page_size_the_query_used_for_limit_zero( + self, + mock_dataset_svc: MagicMock, + mock_provider_mgr: MagicMock, + app: Flask, + account: Account, + tenant: Tenant, + controller_session: Session, + ) -> None: + """limit=0 is floored to one row by the query, so the response must report one.""" + from controllers.service_api.dataset.dataset import DatasetListApi + + total = 7 + dataset = make_dataset(controller_session, tenant, account) + mock_dataset_svc.get_datasets.return_value = ([dataset], total) + mock_provider_mgr.return_value.get_configurations.return_value.get_models.return_value = list[object]() + + with app.test_request_context("/datasets?page=1&limit=0", method="GET"): + api = DatasetListApi() + response, status = unwrap(api.get)(api, controller_session, tenant_id=tenant.id) + + assert status == 200 + assert response["limit"] == 1 + assert response["page"] == 1 + assert response["has_more"] is True + assert mock_dataset_svc.get_datasets.call_args.args[:2] == (1, 1) + + @pytest.mark.usefixtures("account") + @patch("controllers.service_api.dataset.dataset.create_plugin_provider_manager") + @patch("controllers.service_api.dataset.dataset.DatasetService") + def test_list_datasets_has_more_false_past_the_last_page_for_limit_zero( + self, + mock_dataset_svc: MagicMock, + mock_provider_mgr: MagicMock, + app: Flask, + tenant: Tenant, + controller_session: Session, + ) -> None: + """A page past the end must end the walk instead of reporting more rows forever.""" + from controllers.service_api.dataset.dataset import DatasetListApi + + total = 7 + mock_dataset_svc.get_datasets.return_value = (list[Dataset](), total) + mock_provider_mgr.return_value.get_configurations.return_value.get_models.return_value = list[object]() + + with app.test_request_context("/datasets?page=8&limit=0", method="GET"): + api = DatasetListApi() + response, status = unwrap(api.get)(api, controller_session, tenant_id=tenant.id) + + assert status == 200 + assert response["has_more"] is False + assert response["limit"] == 1 + assert response["page"] == 8 + assert response["data"] == [] + + @patch("controllers.service_api.dataset.dataset.create_plugin_provider_manager") + @patch("controllers.service_api.dataset.dataset.DatasetService") + def test_list_datasets_reports_the_page_size_the_query_used_for_negative_limit( + self, + mock_dataset_svc: MagicMock, + mock_provider_mgr: MagicMock, + app: Flask, + account: Account, + tenant: Tenant, + controller_session: Session, + ) -> None: + """A negative limit is floored the same way, and is never echoed back.""" + from controllers.service_api.dataset.dataset import DatasetListApi + + total = 7 + dataset = make_dataset(controller_session, tenant, account) + mock_dataset_svc.get_datasets.return_value = ([dataset], total) + mock_provider_mgr.return_value.get_configurations.return_value.get_models.return_value = list[object]() + + with app.test_request_context("/datasets?page=1&limit=-1", method="GET"): + api = DatasetListApi() + response, status = unwrap(api.get)(api, controller_session, tenant_id=tenant.id) + + assert status == 200 + assert response["limit"] == 1 + assert response["page"] == 1 + assert mock_dataset_svc.get_datasets.call_args.args[:2] == (1, 1) + + @patch("controllers.service_api.dataset.dataset.create_plugin_provider_manager") + @patch("controllers.service_api.dataset.dataset.DatasetService") + def test_list_datasets_reports_the_page_the_query_used_for_page_zero( + self, + mock_dataset_svc: MagicMock, + mock_provider_mgr: MagicMock, + app: Flask, + account: Account, + tenant: Tenant, + controller_session: Session, + ) -> None: + """page=0 is served as page 1, so the response must not offer page 1 as the next one.""" + from controllers.service_api.dataset.dataset import DatasetListApi + + total = 7 + dataset = make_dataset(controller_session, tenant, account) + mock_dataset_svc.get_datasets.return_value = ([dataset] * total, total) + mock_provider_mgr.return_value.get_configurations.return_value.get_models.return_value = list[object]() + + with app.test_request_context("/datasets?page=0&limit=20", method="GET"): + api = DatasetListApi() + response, status = unwrap(api.get)(api, controller_session, tenant_id=tenant.id) + + assert status == 200 + assert response["page"] == 1 + assert response["has_more"] is False + assert mock_dataset_svc.get_datasets.call_args.args[:2] == (1, 20) + + class TestDatasetListApiPost: + """Test suite for DatasetListApi.post() endpoint.""" + + @patch("controllers.service_api.dataset.dataset.DatasetService") + def test_create_dataset_success( + self, + mock_dataset_svc: MagicMock, + app: Flask, + account: Account, + tenant: Tenant, + controller_session: Session, + ) -> None: + from controllers.service_api.dataset.dataset import DatasetCreatePayload, DatasetListApi + + mock_dataset_svc.create_empty_dataset.return_value = make_dataset( + controller_session, tenant, account, name="New Dataset" + ) + + with app.test_request_context( + "/datasets", + method="POST", + json={"name": "New Dataset"}, + ): + api = DatasetListApi() + # `post` is wrapped in @model_validate, so the unwrapped view expects + # the validated model where the decorator would have injected it. + validated_payload = DatasetCreatePayload.model_validate(request.get_json() or {}) + response, status = unwrap(api.post)(api, validated_payload, controller_session, tenant_id=tenant.id) + + assert status == 200 + assert_dataset_detail_shape(response) + assert response["name"] == "New Dataset" + mock_dataset_svc.create_empty_dataset.assert_called_once() + + @pytest.mark.usefixtures("account") + @patch("controllers.service_api.dataset.dataset.DatasetService") + def test_create_dataset_duplicate_name( + self, + mock_dataset_svc: MagicMock, + app: Flask, + tenant: Tenant, + controller_session: Session, + ) -> None: + from controllers.service_api.dataset.dataset import DatasetCreatePayload, DatasetListApi + + mock_dataset_svc.create_empty_dataset.side_effect = services.errors.dataset.DatasetNameDuplicateError() + + with app.test_request_context( + "/datasets", + method="POST", + json={"name": "Existing Dataset"}, + ): + api = DatasetListApi() + validated_payload = DatasetCreatePayload.model_validate(request.get_json() or {}) + with pytest.raises(DatasetNameDuplicateError): + unwrap(api.post)(api, validated_payload, controller_session, tenant_id=tenant.id) + + @pytest.mark.usefixtures("account") + @patch("controllers.service_api.wraps.FeatureService") + @patch("controllers.service_api.wraps.validate_and_get_api_token") + def test_invalid_body_is_rejected_before_the_handler_runs( + self, + mock_validate_token: MagicMock, + mock_feature_svc: MagicMock, + app: Flask, + tenant: Tenant, + ) -> None: + """The tests above unwrap the view, so this is what covers the decorator.""" + from controllers.service_api.dataset.dataset import DatasetListApi + + mock_validate_token.return_value = SimpleNamespace(tenant_id=tenant.id) + mock_feature_svc.get_knowledge_rate_limit.return_value = SimpleNamespace(enabled=False) + + with app.test_request_context("/datasets", method="POST", json={}): + with pytest.raises(UnprocessableEntity): + DatasetListApi().post(tenant_id=tenant.id) + + # --------------------------------------------------------------------------- + # API endpoint tests — DatasetApi + # --------------------------------------------------------------------------- + + class TestDatasetApiGet: + """Test suite for DatasetApi.get() endpoint.""" + + @patch("controllers.service_api.dataset.dataset.create_plugin_provider_manager") + @patch("controllers.service_api.dataset.dataset.DatasetService") + def test_get_dataset_success( + self, + mock_dataset_svc: MagicMock, + mock_provider_mgr: MagicMock, + app: Flask, + dataset: Dataset, + controller_session: Session, + ) -> None: + from controllers.service_api.dataset.dataset import DatasetApi + + mock_dataset_svc.get_dataset.return_value = dataset + mock_dataset_svc.check_dataset_permission.return_value = None + mock_provider_mgr.return_value.get_configurations.return_value.get_models.return_value = list[object]() + + with app.test_request_context( + f"/datasets/{dataset.id}", + method="GET", + ): + api = DatasetApi() + response, status = unwrap(api.get)(api, controller_session, _=dataset.tenant_id, dataset_id=dataset.id) + + assert status == 200 + assert_dataset_detail_shape(response) + assert response["embedding_available"] is True + assert response["retrieval_model_dict"]["search_method"] == "keyword_search" + + @patch("controllers.service_api.dataset.dataset.DatasetPermissionService") + @patch("controllers.service_api.dataset.dataset.create_plugin_provider_manager") + @patch("controllers.service_api.dataset.dataset.DatasetService") + def test_get_dataset_partial_members_shape( + self, + mock_dataset_svc: MagicMock, + mock_provider_mgr: MagicMock, + mock_perm_svc: MagicMock, + app: Flask, + dataset: Dataset, + controller_session: Session, + ) -> None: + from controllers.service_api.dataset.dataset import DatasetApi + + dataset.permission = PermissionEnum.PARTIAL_TEAM + mock_dataset_svc.get_dataset.return_value = dataset + mock_dataset_svc.check_dataset_permission.return_value = None + mock_perm_svc.get_dataset_partial_member_list.return_value = ["user-1", "user-2"] + mock_provider_mgr.return_value.get_configurations.return_value.get_models.return_value = list[object]() + + with app.test_request_context( + f"/datasets/{dataset.id}", + method="GET", + ): + api = DatasetApi() + response, status = unwrap(api.get)(api, controller_session, _=dataset.tenant_id, dataset_id=dataset.id) + + assert status == 200 + assert_dataset_detail_shape(response, with_partial_members=True) + assert response["partial_member_list"] == ["user-1", "user-2"] + + @patch("controllers.service_api.dataset.dataset.create_plugin_provider_manager") + @patch("controllers.service_api.dataset.dataset.DatasetService") + def test_get_dataset_uses_default_external_retrieval_model( + self, + mock_dataset_svc: MagicMock, + mock_provider_mgr: MagicMock, + app: Flask, + dataset: Dataset, + controller_session: Session, + ) -> None: + from controllers.service_api.dataset.dataset import DatasetApi + + dataset.retrieval_model = None + mock_dataset_svc.get_dataset.return_value = dataset + mock_dataset_svc.check_dataset_permission.return_value = None + mock_provider_mgr.return_value.get_configurations.return_value.get_models.return_value = list[object]() + + with app.test_request_context(f"/datasets/{dataset.id}", method="GET"): + api = DatasetApi() + response, status = unwrap(api.get)(api, controller_session, _=dataset.tenant_id, dataset_id=dataset.id) + + assert status == 200 + assert_dataset_detail_shape(response) + assert response["external_retrieval_model"] == { + "top_k": 2, + "score_threshold": 0.0, + "score_threshold_enabled": None, + } + + @patch("controllers.service_api.dataset.dataset.DatasetService") + def test_get_dataset_not_found( + self, + mock_dataset_svc: MagicMock, + app: Flask, + dataset: Dataset, + controller_session: Session, + ) -> None: + from controllers.service_api.dataset.dataset import DatasetApi + + mock_dataset_svc.get_dataset.return_value = None + + with app.test_request_context( + f"/datasets/{dataset.id}", + method="GET", + ): + api = DatasetApi() + with pytest.raises(NotFound): + unwrap(api.get)(api, controller_session, _=dataset.tenant_id, dataset_id=dataset.id) + + @patch("controllers.service_api.dataset.dataset.DatasetService") + def test_get_dataset_no_permission( + self, + mock_dataset_svc: MagicMock, + app: Flask, + dataset: Dataset, + controller_session: Session, + ) -> None: + from controllers.service_api.dataset.dataset import DatasetApi + + mock_dataset_svc.get_dataset.return_value = dataset + mock_dataset_svc.check_dataset_permission.side_effect = services.errors.base.NoPermissionError() + + with app.test_request_context( + f"/datasets/{dataset.id}", + method="GET", + ): + api = DatasetApi() + with pytest.raises(Forbidden): + unwrap(api.get)(api, controller_session, _=dataset.tenant_id, dataset_id=dataset.id) + + class TestDatasetApiPatch: + """Test suite for DatasetApi.patch() endpoint.""" + + @patch("controllers.service_api.dataset.dataset.DatasetPermissionService") + @patch("controllers.service_api.dataset.dataset.DatasetService") + def test_patch_dataset_success_shape( + self, + mock_dataset_svc: MagicMock, + mock_perm_svc: MagicMock, + app: Flask, + dataset: Dataset, + controller_session: Session, + ) -> None: + from controllers.service_api.dataset.dataset import DatasetApi, DatasetUpdatePayload + + dataset.name = "Updated Dataset" + mock_dataset_svc.get_dataset.return_value = dataset + mock_dataset_svc.update_dataset.return_value = dataset + mock_perm_svc.check_permission.return_value = None + mock_perm_svc.get_dataset_partial_member_list.return_value = ["user-1"] + + payload = { + "name": "Updated Dataset", + "permission": "partial_members", + "partial_member_list": [{"user_id": "user-1", "role": "editor"}], + } + with app.test_request_context( + f"/datasets/{dataset.id}", + method="PATCH", + json=payload, + ): + api = DatasetApi() + # `patch` is wrapped in @model_validate, so the unwrapped view expects + # the validated model where the decorator would have injected it. + validated_payload = DatasetUpdatePayload.model_validate(request.get_json() or {}) + response, status = unwrap(api.patch)( + api, + validated_payload, + controller_session, + _=dataset.tenant_id, + dataset_id=dataset.id, + ) + + assert status == 200 + assert_dataset_detail_shape(response, with_partial_members=True) + assert response["name"] == "Updated Dataset" + assert response["partial_member_list"] == ["user-1"] + mock_dataset_svc.update_dataset.assert_called_once() + _, update_data, _ = mock_dataset_svc.update_dataset.call_args.args + session = mock_dataset_svc.update_dataset.call_args.kwargs["session"] + assert session is controller_session + assert update_data["name"] == "Updated Dataset" + assert update_data["permission"] == "partial_members" + mock_perm_svc.update_partial_member_list.assert_called_once_with( + dataset.tenant_id, + dataset.id, + [{"user_id": "user-1", "role": "editor"}], + controller_session, + ) + + class TestDatasetApiDelete: + """Test suite for DatasetApi.delete() endpoint.""" + + @patch("controllers.service_api.dataset.dataset.DatasetPermissionService") + @patch("controllers.service_api.dataset.dataset.DatasetService") + def test_delete_dataset_success( + self, + mock_dataset_svc: MagicMock, + mock_perm_svc: MagicMock, + app: Flask, + dataset: Dataset, + controller_session: Session, + ) -> None: + from controllers.service_api.dataset.dataset import DatasetApi + + mock_dataset_svc.delete_dataset.return_value = True + + with app.test_request_context( + f"/datasets/{dataset.id}", + method="DELETE", + ): + api = DatasetApi() + result = unwrap(api.delete)(api, controller_session, _=dataset.tenant_id, dataset_id=dataset.id) + + assert result == ("", 204) + mock_perm_svc.clear_partial_member_list.assert_called_once_with(dataset.id, controller_session) + + @patch("controllers.service_api.dataset.dataset.DatasetService") + def test_delete_dataset_not_found( + self, + mock_dataset_svc: MagicMock, + app: Flask, + dataset: Dataset, + controller_session: Session, + ) -> None: + from controllers.service_api.dataset.dataset import DatasetApi + + mock_dataset_svc.delete_dataset.return_value = False + + with app.test_request_context( + f"/datasets/{dataset.id}", + method="DELETE", + ): + api = DatasetApi() + with pytest.raises(NotFound): + unwrap(api.delete)(api, controller_session, _=dataset.tenant_id, dataset_id=dataset.id) + + @patch("controllers.service_api.dataset.dataset.DatasetService") + def test_delete_dataset_in_use( + self, + mock_dataset_svc: MagicMock, + app: Flask, + dataset: Dataset, + controller_session: Session, + ) -> None: + from controllers.service_api.dataset.dataset import DatasetApi + + mock_dataset_svc.delete_dataset.side_effect = services.errors.dataset.DatasetInUseError() + + with app.test_request_context( + f"/datasets/{dataset.id}", + method="DELETE", + ): + api = DatasetApi() + with pytest.raises(DatasetInUseError): + unwrap(api.delete)(api, controller_session, _=dataset.tenant_id, dataset_id=dataset.id) + + # --------------------------------------------------------------------------- + # API endpoint tests — DocumentStatusApi + # --------------------------------------------------------------------------- + + class TestDocumentStatusApiPatch: + """Test suite for DocumentStatusApi.patch() endpoint.""" + + @patch("controllers.service_api.dataset.dataset.DocumentService") + @patch("controllers.service_api.dataset.dataset.DatasetService") + def test_batch_update_status_success( + self, + mock_dataset_svc: MagicMock, + mock_doc_svc: MagicMock, + app: Flask, + tenant: Tenant, + dataset: Dataset, + ) -> None: + from controllers.service_api.dataset.dataset import DocumentStatusApi + + mock_dataset_svc.get_dataset.return_value = dataset + mock_dataset_svc.check_dataset_permission.return_value = None + mock_dataset_svc.check_dataset_model_setting.return_value = None + mock_doc_svc.batch_update_document_status.return_value = None + + with app.test_request_context( + f"/datasets/{dataset.id}/documents/status/enable", + method="PATCH", + json={"document_ids": ["doc-1", "doc-2"]}, + ): + api = DocumentStatusApi() + response, status = api.patch( + tenant_id=tenant.id, + dataset_id=dataset.id, + action="enable", + ) + + assert status == 200 + assert response["result"] == "success" + + @patch("controllers.service_api.dataset.dataset.DatasetService") + def test_batch_update_status_dataset_not_found( + self, + mock_dataset_svc: MagicMock, + app: Flask, + tenant: Tenant, + dataset: Dataset, + ) -> None: + from controllers.service_api.dataset.dataset import DocumentStatusApi + + mock_dataset_svc.get_dataset.return_value = None + + with app.test_request_context( + f"/datasets/{dataset.id}/documents/status/enable", + method="PATCH", + json={"document_ids": ["doc-1"]}, + ): + api = DocumentStatusApi() + with pytest.raises(NotFound): + api.patch( + tenant_id=tenant.id, + dataset_id=dataset.id, + action="enable", + ) + + @patch("controllers.service_api.dataset.dataset.DatasetService") + def test_batch_update_status_permission_error( + self, + mock_dataset_svc: MagicMock, + app: Flask, + tenant: Tenant, + dataset: Dataset, + ) -> None: + from controllers.service_api.dataset.dataset import DocumentStatusApi + + mock_dataset_svc.get_dataset.return_value = dataset + mock_dataset_svc.check_dataset_permission.side_effect = services.errors.base.NoPermissionError( + "No permission" + ) + + with app.test_request_context( + f"/datasets/{dataset.id}/documents/status/enable", + method="PATCH", + json={"document_ids": ["doc-1"]}, + ): + api = DocumentStatusApi() + with pytest.raises(Forbidden): + api.patch( + tenant_id=tenant.id, + dataset_id=dataset.id, + action="enable", + ) + + @patch("controllers.service_api.dataset.dataset.DocumentService") + @patch("controllers.service_api.dataset.dataset.DatasetService") + def test_batch_update_status_indexing_error( + self, + mock_dataset_svc: MagicMock, + mock_doc_svc: MagicMock, + app: Flask, + tenant: Tenant, + dataset: Dataset, + ) -> None: + from controllers.service_api.dataset.dataset import DocumentStatusApi + + mock_dataset_svc.get_dataset.return_value = dataset + mock_dataset_svc.check_dataset_permission.return_value = None + mock_dataset_svc.check_dataset_model_setting.return_value = None + mock_doc_svc.batch_update_document_status.side_effect = services.errors.document.DocumentIndexingError() + + with app.test_request_context( + f"/datasets/{dataset.id}/documents/status/enable", + method="PATCH", + json={"document_ids": ["doc-1"]}, + ): + api = DocumentStatusApi() + with pytest.raises(InvalidActionError): + api.patch( + tenant_id=tenant.id, + dataset_id=dataset.id, + action="enable", + ) + + @patch("controllers.service_api.dataset.dataset.DocumentService") + @patch("controllers.service_api.dataset.dataset.DatasetService") + def test_batch_update_status_value_error( + self, + mock_dataset_svc: MagicMock, + mock_doc_svc: MagicMock, + app: Flask, + tenant: Tenant, + dataset: Dataset, + ) -> None: + from controllers.service_api.dataset.dataset import DocumentStatusApi + + mock_dataset_svc.get_dataset.return_value = dataset + mock_dataset_svc.check_dataset_permission.return_value = None + mock_dataset_svc.check_dataset_model_setting.return_value = None + mock_doc_svc.batch_update_document_status.side_effect = ValueError("Invalid action") + + with app.test_request_context( + f"/datasets/{dataset.id}/documents/status/enable", + method="PATCH", + json={"document_ids": ["doc-1"]}, + ): + api = DocumentStatusApi() + with pytest.raises(InvalidActionError): + api.patch( + tenant_id=tenant.id, + dataset_id=dataset.id, + action="enable", + ) + + +"""Unit tests for Service API dataset request payloads.""" + + +class TestDatasetPayloadsCases: + class TestDatasetCreatePayload: + """Test suite for DatasetCreatePayload Pydantic model.""" + + def test_payload_with_required_name(self) -> None: + payload = DatasetCreatePayload(name="Test Dataset") + assert payload.name == "Test Dataset" + assert payload.description == "" + assert payload.permission == DatasetPermissionEnum.ONLY_ME + + def test_payload_with_all_fields(self) -> None: + payload = DatasetCreatePayload( + name="Full Dataset", + description="A comprehensive dataset description", + indexing_technique="high_quality", + permission=DatasetPermissionEnum.ALL_TEAM, + provider="vendor", + embedding_model="text-embedding-ada-002", + embedding_model_provider="openai", + ) + assert payload.name == "Full Dataset" + assert payload.description == "A comprehensive dataset description" + assert payload.indexing_technique == "high_quality" + assert payload.permission == DatasetPermissionEnum.ALL_TEAM + assert payload.provider == "vendor" + assert payload.embedding_model == "text-embedding-ada-002" + assert payload.embedding_model_provider == "openai" + + def test_payload_name_length_validation_min(self) -> None: + with pytest.raises(ValueError, match="1 validation error for DatasetCreatePayload\\\nname"): + DatasetCreatePayload(name="") + + def test_payload_name_length_validation_max(self) -> None: + with pytest.raises(ValueError, match="1 validation error for DatasetCreatePayload\\\nname"): + DatasetCreatePayload(name="A" * 41) + + def test_payload_description_max_length(self) -> None: + with pytest.raises(ValueError, match="1 validation error for DatasetCreatePayload\\\ndescription"): + DatasetCreatePayload(name="Dataset", description="A" * 401) + + @pytest.mark.parametrize("technique", ["high_quality", "economy"]) + def test_payload_valid_indexing_techniques(self, technique: Literal["high_quality", "economy"]) -> None: + payload = DatasetCreatePayload(name="Dataset", indexing_technique=technique) + assert payload.indexing_technique == technique + + def test_payload_with_external_knowledge_settings(self) -> None: + payload = DatasetCreatePayload( + name="External Dataset", external_knowledge_api_id="api_123", external_knowledge_id="knowledge_456" + ) + assert payload.external_knowledge_api_id == "api_123" + assert payload.external_knowledge_id == "knowledge_456" + + class TestDatasetUpdatePayload: + """Test suite for DatasetUpdatePayload Pydantic model.""" + + def test_payload_all_optional(self) -> None: + payload = DatasetUpdatePayload() + assert payload.name is None + assert payload.description is None + assert payload.permission is None + + def test_payload_with_partial_update(self) -> None: + payload = DatasetUpdatePayload(name="Updated Name", description="Updated description") + assert payload.name == "Updated Name" + assert payload.description == "Updated description" + + def test_payload_with_permission_change(self) -> None: + payload = DatasetUpdatePayload( + permission=DatasetPermissionEnum.PARTIAL_TEAM, + partial_member_list=[{"user_id": "user_123", "role": "editor"}], + ) + assert payload.permission == DatasetPermissionEnum.PARTIAL_TEAM + assert payload.partial_member_list is not None + assert len(payload.partial_member_list) == 1 + + def test_payload_name_length_validation(self) -> None: + with pytest.raises(ValueError, match="1 validation error for DatasetUpdatePayload\\\nname"): + DatasetUpdatePayload(name="") + with pytest.raises(ValueError, match="1 validation error for DatasetUpdatePayload\\\nname"): + DatasetUpdatePayload(name="A" * 41) + + class TestDatasetListQuery: + """Test suite for DatasetListQuery Pydantic model.""" + + def test_query_with_defaults(self) -> None: + query = DatasetListQuery() + assert query.page == 1 + assert query.limit == 20 + assert query.keyword is None + assert query.include_all is False + assert query.tag_ids == [] + + def test_query_with_all_filters(self) -> None: + query = DatasetListQuery( + page=3, limit=50, keyword="machine learning", include_all=True, tag_ids=["tag1", "tag2", "tag3"] + ) + assert query.page == 3 + assert query.limit == 50 + assert query.keyword == "machine learning" + assert query.include_all is True + assert len(query.tag_ids) == 3 + + def test_query_with_tag_filter(self) -> None: + query = DatasetListQuery(tag_ids=["tag_abc", "tag_def"]) + assert query.tag_ids == ["tag_abc", "tag_def"] + + class TestTagCreatePayload: + """Test suite for TagCreatePayload Pydantic model.""" + + def test_payload_with_name(self) -> None: + payload = TagCreatePayload(name="New Tag") + assert payload.name == "New Tag" + + def test_payload_name_length_min(self) -> None: + with pytest.raises(ValueError, match="1 validation error for TagCreatePayload\\\nname"): + TagCreatePayload(name="") + + def test_payload_name_length_max(self) -> None: + with pytest.raises(ValueError, match="1 validation error for TagCreatePayload\\\nname"): + TagCreatePayload(name="A" * 51) + + def test_payload_with_unicode_name(self) -> None: + payload = TagCreatePayload(name="标签 🏷️ Тег") + assert payload.name == "标签 🏷️ Тег" + + class TestTagUpdatePayload: + """Test suite for TagUpdatePayload Pydantic model.""" + + def test_payload_with_name_and_id(self) -> None: + payload = TagUpdatePayload(name="Updated Tag", tag_id="tag_123") + assert payload.name == "Updated Tag" + assert payload.tag_id == "tag_123" + + def test_payload_requires_tag_id(self) -> None: + with pytest.raises(ValueError, match="1 validation error for TagUpdatePayload\\\ntag_id"): + TagUpdatePayload.model_validate({"name": "Updated Tag"}) + + class TestTagDeletePayload: + """Test suite for TagDeletePayload Pydantic model.""" + + def test_payload_with_tag_id(self) -> None: + payload = TagDeletePayload(tag_id="tag_to_delete") + assert payload.tag_id == "tag_to_delete" + + def test_payload_requires_tag_id(self) -> None: + with pytest.raises(ValueError, match="1 validation error for TagDeletePayload\\\ntag_id"): + TagDeletePayload.model_validate({}) + + class TestTagBindingPayload: + """Test suite for TagBindingPayload Pydantic model.""" + + def test_payload_with_valid_data(self) -> None: + payload = TagBindingPayload(tag_ids=["tag1", "tag2"], target_id="dataset_123") + assert len(payload.tag_ids) == 2 + assert payload.target_id == "dataset_123" + + def test_payload_rejects_empty_tag_ids(self) -> None: + with pytest.raises(ValueError, match="1 validation error for TagBindingPayload\\\ntag_ids") as exc_info: + TagBindingPayload(tag_ids=[], target_id="dataset_123") + assert "Tag IDs is required" in str(exc_info.value) + + def test_payload_single_tag_id(self) -> None: + payload = TagBindingPayload(tag_ids=["single_tag"], target_id="dataset_456") + assert payload.tag_ids == ["single_tag"] + + class TestTagUnbindingPayload: + """Test suite for TagUnbindingPayload Pydantic model.""" + + def test_payload_with_valid_data(self) -> None: + payload = TagUnbindingPayload(tag_ids=["tag_123"], target_id="dataset_456") + assert payload.tag_ids == ["tag_123"] + assert payload.target_id == "dataset_456" + + def test_payload_normalizes_legacy_tag_id(self) -> None: + payload = TagUnbindingPayload(tag_id="tag_123", target_id="dataset_456") + assert payload.tag_ids == ["tag_123"] + assert payload.target_id == "dataset_456" + + def test_payload_rejects_empty_tag_ids(self) -> None: + with pytest.raises( + ValueError, + match=("1 validation error for TagUnbindingPayload\\\n Value error, Tag IDs is required\\."), + ) as exc_info: + TagUnbindingPayload(tag_ids=[], target_id="dataset_456") + assert "Tag IDs is required" in str(exc_info.value) + + +"""Unit tests for Service API dataset tag controller behavior. + +Service boundaries stay mocked, while users, tenants, and tags are real ORM objects +persisted in SQLite. Controller database calls share that SQLite session so assertions +cover the concrete objects and session passed across the controller boundary. +""" + + +TAG_MODEL_TABLES = (Account, Tenant, Tag) + + +def make_tag( + session: Session, + tenant: Tenant, + account: Account, + *, + id: str, + name: str, + binding_count: int | None = None, +) -> Tag: + """Create and flush a real tag, optionally adding the aggregate count returned by TagService.""" + + tag = Tag(tenant_id=tenant.id, type=TagType.KNOWLEDGE, name=name, created_by=account.id) + tag.id = id + session.add(tag) + session.flush() + if binding_count is not None: + tag.__dict__["binding_count"] = binding_count + return tag + + +class TestDatasetTagApisCases: + pytestmark = pytest.mark.parametrize("sqlite_session", [TAG_MODEL_TABLES], indirect=True) + + @staticmethod + @pytest.fixture + def controller_session(sqlite_session: Session) -> Session: + """Expose the shared SQLite session under the controller-focused fixture name.""" + return sqlite_session + + @staticmethod + @pytest.fixture + def tenant(controller_session: Session) -> Tenant: + tenant = Tenant(name="Dataset Tag API Tenant") + controller_session.add(tenant) + controller_session.flush() + return tenant + + @staticmethod + @pytest.fixture + def account(controller_session: Session, tenant: Tenant, monkeypatch: pytest.MonkeyPatch) -> Account: + account = Account(name="Dataset Tag API User", email=f"dataset-tag-api-{uuid.uuid4()}@example.com") + account.role = TenantAccountRole.OWNER + account._current_tenant = tenant + controller_session.add(account) + controller_session.flush() + + # Inject the concrete account at the controller boundary without relying on Flask-Login globals. + from controllers.service_api.dataset import dataset as dataset_module + + monkeypatch.setattr(dataset_module, "current_user", account) + return account + + class TestDatasetTagsApiGet: + """Test suite for DatasetTagsApi.get() endpoint.""" + + @patch("controllers.service_api.dataset.dataset.TagService") + def test_list_tags_success( + self, + mock_tag_svc: MagicMock, + app: Flask, + account: Account, + tenant: Tenant, + controller_session: Session, + ) -> None: + from controllers.service_api.dataset.dataset import DatasetTagsApi + + tag = make_tag(controller_session, tenant, account, id="tag-1", name="Test Tag", binding_count=0) + mock_tag_svc.get_tags.return_value = [tag] + + with app.test_request_context("/datasets/tags", method="GET"): + api = DatasetTagsApi() + response, status = unwrap(api.get)(api, controller_session, _=None) + + assert status == 200 + assert response == [{"id": "tag-1", "name": "Test Tag", "type": "knowledge", "binding_count": "0"}] + mock_tag_svc.get_tags.assert_called_once_with("knowledge", tenant.id, session=controller_session) + + class TestDatasetTagsApiPost: + """Test suite for DatasetTagsApi.post() endpoint.""" + + @patch("controllers.service_api.dataset.dataset.TagService") + def test_create_tag_success( + self, + mock_tag_svc: MagicMock, + app: Flask, + account: Account, + tenant: Tenant, + controller_session: Session, + ) -> None: + from controllers.service_api.dataset.dataset import DatasetTagsApi, TagCreatePayload + + tag = make_tag(controller_session, tenant, account, id="tag-new", name="New Tag") + mock_tag_svc.save_tags.return_value = tag + + with app.test_request_context( + "/datasets/tags", + method="POST", + json={"name": "New Tag"}, + ): + api = DatasetTagsApi() + payload = TagCreatePayload.model_validate(request.get_json() or {}) + response, status = unwrap(api.post)(api, payload, controller_session, _=None) + + assert status == 200 + assert response == {"id": "tag-new", "name": "New Tag", "type": "knowledge", "binding_count": "0"} + mock_tag_svc.save_tags.assert_called_once() + + def test_create_tag_forbidden(self, app: Flask, account: Account) -> None: + from controllers.service_api.dataset.dataset import DatasetTagsApi + + account.role = TenantAccountRole.NORMAL + + with app.test_request_context( + "/datasets/tags", + method="POST", + json={"name": "New Tag"}, + ): + api = DatasetTagsApi() + with pytest.raises(Forbidden): + api.post(_=None) + + class TestDatasetTagsApiPatch: + """Test suite for DatasetTagsApi.patch() endpoint.""" + + @patch("controllers.service_api.dataset.dataset.TagService") + @patch("controllers.service_api.dataset.dataset.service_api_ns") + def test_update_tag_success( + self, + mock_service_api_ns: MagicMock, + mock_tag_svc: MagicMock, + app: Flask, + account: Account, + tenant: Tenant, + controller_session: Session, + ) -> None: + from controllers.service_api.dataset.dataset import DatasetTagsApi, TagUpdatePayload + + tag = make_tag(controller_session, tenant, account, id="tag-1", name="Updated Tag") + mock_tag_svc.update_tags.return_value = tag + mock_tag_svc.get_tag_binding_count.return_value = 5 + mock_service_api_ns.payload = {"name": "Updated Tag", "tag_id": "tag-1"} + + with app.test_request_context( + "/datasets/tags", + method="PATCH", + json={"name": "Updated Tag", "tag_id": "tag-1"}, + ): + api = DatasetTagsApi() + payload = TagUpdatePayload.model_validate(request.get_json() or {}) + response, status = unwrap(api.patch)(api, payload, controller_session, _=None) + + assert status == 200 + assert response == {"id": "tag-1", "name": "Updated Tag", "type": "knowledge", "binding_count": "5"} + mock_tag_svc.update_tags.assert_called_once() + update_payload, tag_id, session = mock_tag_svc.update_tags.call_args.args + assert update_payload.name == "Updated Tag" + assert tag_id == "tag-1" + assert session is controller_session + + def test_update_tag_forbidden(self, app: Flask, account: Account) -> None: + from controllers.service_api.dataset.dataset import DatasetTagsApi + + account.role = TenantAccountRole.NORMAL + + with app.test_request_context( + "/datasets/tags", + method="PATCH", + json={"name": "Updated Tag", "tag_id": "tag-1"}, + ): + api = DatasetTagsApi() + with pytest.raises(Forbidden): + api.patch(_=None) + + class TestDatasetTagsApiDelete: + """Test suite for DatasetTagsApi.delete() endpoint.""" + + @pytest.mark.usefixtures("account") + @patch("controllers.service_api.dataset.dataset.TagService") + @patch("controllers.service_api.dataset.dataset.service_api_ns") + def test_delete_tag_success( + self, + mock_service_api_ns: MagicMock, + mock_tag_svc: MagicMock, + app: Flask, + controller_session: Session, + ) -> None: + from controllers.service_api.dataset.dataset import DatasetTagsApi, TagDeletePayload + + mock_tag_svc.delete_tag.return_value = None + mock_service_api_ns.payload = {"tag_id": "tag-1"} + + with app.test_request_context( + "/datasets/tags", + method="DELETE", + json={"tag_id": "tag-1"}, + ): + api = DatasetTagsApi() + payload = TagDeletePayload.model_validate(request.get_json() or {}) + result = unwrap(api.delete)(api, payload, controller_session, _=None) + + assert result == ("", 204) + mock_tag_svc.delete_tag.assert_called_once_with("tag-1", controller_session, tag_type=TagType.KNOWLEDGE) + + class TestDatasetTagsBindingStatusApi: + """Test suite for DatasetTagsBindingStatusApi endpoints.""" + + @patch("controllers.service_api.dataset.dataset.TagService") + def test_get_dataset_tags_binding_status( + self, + mock_tag_svc: MagicMock, + app: Flask, + account: Account, + tenant: Tenant, + controller_session: Session, + ) -> None: + from controllers.service_api.dataset.dataset import DatasetTagsBindingStatusApi + + tag = make_tag(controller_session, tenant, account, id="tag_1", name="Test Tag") + mock_tag_svc.get_tags_by_target_id.return_value = [tag] + + with app.test_request_context("/", method="GET"): + api = DatasetTagsBindingStatusApi() + response, status_code = unwrap(api.get)(api, controller_session, tenant.id, dataset_id="dataset_123") + + assert status_code == 200 + assert response["data"] == [{"id": "tag_1", "name": "Test Tag"}] + assert response["total"] == 1 + mock_tag_svc.get_tags_by_target_id.assert_called_once_with( + "knowledge", tenant.id, "dataset_123", controller_session + ) + + class TestDatasetTagBindingApiPost: + """Test suite for DatasetTagBindingApi.post() endpoint.""" + + @pytest.mark.usefixtures("account") + @patch("controllers.service_api.dataset.dataset.TagService") + def test_bind_tags_success( + self, + mock_tag_svc: MagicMock, + app: Flask, + controller_session: Session, + ) -> None: + from controllers.service_api.dataset.dataset import DatasetTagBindingApi, TagBindingPayload + + mock_tag_svc.save_tag_binding.return_value = None + + with app.test_request_context( + "/datasets/tags/binding", + method="POST", + json={"tag_ids": ["tag-1"], "target_id": "ds-1"}, + ): + api = DatasetTagBindingApi() + payload = TagBindingPayload.model_validate(request.get_json() or {}) + result = unwrap(api.post)(api, payload, controller_session, _=None) + + assert result == ("", 204) + from services.tag_service import TagBindingCreatePayload + + mock_tag_svc.save_tag_binding.assert_called_once_with( + TagBindingCreatePayload(tag_ids=["tag-1"], target_id="ds-1", type=TagType.KNOWLEDGE), + controller_session, + ) + + def test_bind_tags_forbidden(self, app: Flask, account: Account) -> None: + from controllers.service_api.dataset.dataset import DatasetTagBindingApi + + account.role = TenantAccountRole.NORMAL + + with app.test_request_context( + "/datasets/tags/binding", + method="POST", + json={"tag_ids": ["tag-1"], "target_id": "ds-1"}, + ): + api = DatasetTagBindingApi() + with pytest.raises(Forbidden): + api.post(_=None) + + class TestDatasetTagUnbindingApiPost: + """Test suite for DatasetTagUnbindingApi.post() endpoint.""" + + @pytest.mark.usefixtures("account") + @patch("controllers.service_api.dataset.dataset.TagService") + def test_unbind_tag_success( + self, + mock_tag_svc: MagicMock, + app: Flask, + controller_session: Session, + ) -> None: + from controllers.service_api.dataset.dataset import DatasetTagUnbindingApi, TagUnbindingPayload + + mock_tag_svc.delete_tag_binding.return_value = None + + with app.test_request_context( + "/datasets/tags/unbinding", + method="POST", + json={"tag_ids": ["tag-1"], "target_id": "ds-1"}, + ): + api = DatasetTagUnbindingApi() + payload = TagUnbindingPayload.model_validate(request.get_json() or {}) + result = unwrap(api.post)(api, payload, controller_session, _=None) + + assert result == ("", 204) + from services.tag_service import TagBindingDeletePayload + + mock_tag_svc.delete_tag_binding.assert_called_once_with( + TagBindingDeletePayload(tag_ids=["tag-1"], target_id="ds-1", type=TagType.KNOWLEDGE), + controller_session, + ) + + @pytest.mark.usefixtures("account") + @patch("controllers.service_api.dataset.dataset.TagService") + def test_unbind_legacy_tag_id_success( + self, + mock_tag_svc: MagicMock, + app: Flask, + controller_session: Session, + ) -> None: + from controllers.service_api.dataset.dataset import DatasetTagUnbindingApi, TagUnbindingPayload + + mock_tag_svc.delete_tag_binding.return_value = None + + with app.test_request_context( + "/datasets/tags/unbinding", + method="POST", + json={"tag_id": "tag-1", "target_id": "ds-1"}, + ): + api = DatasetTagUnbindingApi() + payload = TagUnbindingPayload.model_validate(request.get_json() or {}) + result = unwrap(api.post)(api, payload, controller_session, _=None) + + assert result == ("", 204) + from services.tag_service import TagBindingDeletePayload + + mock_tag_svc.delete_tag_binding.assert_called_once_with( + TagBindingDeletePayload(tag_ids=["tag-1"], target_id="ds-1", type=TagType.KNOWLEDGE), + controller_session, + ) + + def test_unbind_tag_forbidden(self, app: Flask, account: Account) -> None: + from controllers.service_api.dataset.dataset import DatasetTagUnbindingApi + + account.role = TenantAccountRole.NORMAL + + with app.test_request_context( + "/datasets/tags/unbinding", + method="POST", + json={"tag_ids": ["tag-1"], "target_id": "ds-1"}, + ): + api = DatasetTagUnbindingApi() + with pytest.raises(Forbidden): + api.post(_=None) diff --git a/api/tests/unit_tests/controllers/service_api/dataset/test_dataset_apis.py b/api/tests/unit_tests/controllers/service_api/dataset/test_dataset_apis.py deleted file mode 100644 index 5ec542c7a56..00000000000 --- a/api/tests/unit_tests/controllers/service_api/dataset/test_dataset_apis.py +++ /dev/null @@ -1,934 +0,0 @@ -"""Unit tests for Service API dataset controller behavior. - -Service boundaries stay mocked, while ORM collaborators are concrete model instances -persisted through the shared SQLite session fixture. -""" - -import uuid -from datetime import UTC, datetime -from inspect import unwrap -from types import SimpleNamespace -from typing import cast -from unittest.mock import MagicMock, patch - -import pytest -from flask import Flask, request -from sqlalchemy.orm import Session -from werkzeug.exceptions import Forbidden, NotFound, UnprocessableEntity - -import services.errors.base -from controllers.service_api.dataset.error import DatasetInUseError, DatasetNameDuplicateError, InvalidActionError -from models.account import Account, Tenant, TenantAccountRole -from models.dataset import AppDatasetJoin, Dataset, DatasetMetadata, Document -from models.enums import PermissionEnum -from models.model import App, Tag, TagBinding - -# --------------------------------------------------------------------------- -# Helpers -# --------------------------------------------------------------------------- -DATASET_MODEL_TABLES = ( - Account, - Tenant, - Dataset, - Document, - App, - AppDatasetJoin, - DatasetMetadata, - Tag, - TagBinding, -) -pytestmark = pytest.mark.parametrize("sqlite_session", [DATASET_MODEL_TABLES], indirect=True) - - -@pytest.fixture -def controller_session(sqlite_session: Session) -> Session: - """Expose the shared SQLite session under the controller-focused fixture name.""" - return sqlite_session - - -@pytest.fixture -def tenant(controller_session: Session) -> Tenant: - tenant = Tenant(name="Dataset API Tenant") - controller_session.add(tenant) - controller_session.flush() - return tenant - - -@pytest.fixture -def account(controller_session: Session, tenant: Tenant, monkeypatch: pytest.MonkeyPatch) -> Account: - account = Account(name="Dataset API User", email=f"dataset-api-{uuid.uuid4()}@example.com") - account.role = TenantAccountRole.OWNER - account._current_tenant = tenant - controller_session.add(account) - controller_session.flush() - - # Inject the concrete account at the controller boundary without relying on Flask-Login globals. - from controllers.service_api.dataset import dataset as dataset_module - - monkeypatch.setattr(dataset_module, "current_user", account) - return account - - -def make_dataset( - session: Session, - tenant: Tenant, - account: Account, - **overrides: object, -) -> Dataset: - """Create and flush a real dataset so its database-backed properties can be serialized.""" - - base: dict[str, object] = { - "id": str(uuid.uuid4()), - "tenant_id": tenant.id, - "name": "Dataset", - "description": "desc", - "provider": "vendor", - "permission": PermissionEnum.ONLY_ME, - "data_source_type": None, - "indexing_technique": "economy", - "created_by": account.id, - "created_at": datetime(2024, 1, 1, 12, 0, 0, tzinfo=UTC), - "updated_by": None, - "updated_at": datetime(2024, 1, 1, 12, 0, 0, tzinfo=UTC), - "embedding_model": None, - "embedding_model_provider": None, - "retrieval_model": None, - "summary_index_setting": None, - "built_in_field_enabled": False, - "pipeline_id": None, - "runtime_mode": "general", - "chunk_structure": None, - "icon_info": None, - "enable_api": False, - "is_multimodal": False, - } - base.update(overrides) - created_at = cast(datetime, base.pop("created_at")) - updated_at = cast(datetime, base.pop("updated_at")) - dataset = Dataset(**base) # type: ignore[bad-argument-type] - dataset.created_at = created_at - dataset.updated_at = updated_at - session.add(dataset) - session.flush() - return dataset - - -@pytest.fixture -def dataset(controller_session: Session, tenant: Tenant, account: Account) -> Dataset: - return make_dataset(controller_session, tenant, account) - - -DATASET_DETAIL_KEYS = { - "id", - "name", - "description", - "provider", - "permission", - "data_source_type", - "indexing_technique", - "app_count", - "document_count", - "word_count", - "created_by", - "author_name", - "created_at", - "updated_by", - "updated_at", - "embedding_model", - "embedding_model_provider", - "embedding_available", - "retrieval_model_dict", - "summary_index_setting", - "tags", - "doc_form", - "external_knowledge_info", - "external_retrieval_model", - "doc_metadata", - "built_in_field_enabled", - "pipeline_id", - "runtime_mode", - "chunk_structure", - "icon_info", - "is_published", - "total_documents", - "total_available_documents", - "enable_api", - "is_multimodal", - "maintainer", -} - - -def assert_dataset_detail_shape(response: dict[str, object], *, with_partial_members: bool = False) -> None: - expected_keys = set(DATASET_DETAIL_KEYS) - if with_partial_members: - expected_keys.add("partial_member_list") - assert set(response) == expected_keys - assert isinstance(response["created_at"], int) - assert isinstance(response["updated_at"], int) - retrieval_model = response["retrieval_model_dict"] - assert isinstance(retrieval_model, dict) - assert set(retrieval_model) == { - "search_method", - "reranking_enable", - "reranking_mode", - "reranking_model", - "weights", - "top_k", - "score_threshold_enabled", - "score_threshold", - } - external_retrieval_model = response["external_retrieval_model"] - if external_retrieval_model is not None: - assert isinstance(external_retrieval_model, dict) - assert set(external_retrieval_model) == { - "top_k", - "score_threshold", - "score_threshold_enabled", - } - if not with_partial_members: - assert "partial_member_list" not in response - - -# --------------------------------------------------------------------------- -# API endpoint tests — DatasetListApi -# --------------------------------------------------------------------------- - - -@pytest.mark.usefixtures("app_query_services") -class TestDatasetListApiGet: - """Test suite for DatasetListApi.get() endpoint.""" - - @patch("controllers.service_api.dataset.dataset.create_plugin_provider_manager") - @patch("controllers.service_api.dataset.dataset.DatasetService") - def test_list_datasets_success( - self, - mock_dataset_svc: MagicMock, - mock_provider_mgr: MagicMock, - app: Flask, - account: Account, - tenant: Tenant, - controller_session: Session, - ) -> None: - from controllers.service_api.dataset.dataset import DatasetListApi - - mock_dataset_svc.get_datasets.return_value = ([make_dataset(controller_session, tenant, account)], 1) - mock_provider_mgr.return_value.get_configurations.return_value.get_models.return_value = list[object]() - - with app.test_request_context("/datasets?page=1&limit=20", method="GET"): - api = DatasetListApi() - response, status = unwrap(api.get)(api, controller_session, tenant_id=tenant.id) - - assert status == 200 - assert set(response) == {"data", "has_more", "limit", "total", "page"} - assert response["has_more"] is False - assert response["limit"] == 20 - assert response["total"] == 1 - assert response["page"] == 1 - assert len(response["data"]) == 1 - assert_dataset_detail_shape(response["data"][0]) - - @patch("controllers.service_api.dataset.dataset.create_plugin_provider_manager") - @patch("controllers.service_api.dataset.dataset.DatasetService") - def test_list_datasets_preserves_repeated_tag_ids( - self, - mock_dataset_svc: MagicMock, - mock_provider_mgr: MagicMock, - app: Flask, - account: Account, - tenant: Tenant, - controller_session: Session, - ) -> None: - from controllers.service_api.dataset.dataset import DatasetListApi - - mock_dataset_svc.get_datasets.return_value = ([make_dataset(controller_session, tenant, account)], 1) - mock_provider_mgr.return_value.get_configurations.return_value.get_models.return_value = list[object]() - - with app.test_request_context("/datasets?tag_ids=tag-a&tag_ids=tag-b", method="GET"): - api = DatasetListApi() - response, status = unwrap(api.get)(api, controller_session, tenant_id=tenant.id) - page, limit, session, tenant_id, user, keyword, tag_ids, include_all = ( - mock_dataset_svc.get_datasets.call_args.args - ) - assert user is account - - assert status == 200 - assert response["total"] == 1 - assert (page, limit, session, tenant_id, keyword, tag_ids, include_all) == ( - 1, - 20, - controller_session, - tenant.id, - None, - ["tag-a", "tag-b"], - False, - ) - - @patch("controllers.service_api.dataset.dataset.create_plugin_provider_manager") - @patch("controllers.service_api.dataset.dataset.DatasetService") - def test_list_datasets_has_more_false_on_last_page_exact_limit( - self, - mock_dataset_svc: MagicMock, - mock_provider_mgr: MagicMock, - app: Flask, - account: Account, - tenant: Tenant, - controller_session: Session, - ) -> None: - """A full last page must set has_more false instead of forcing another fetch.""" - from controllers.service_api.dataset.dataset import DatasetListApi - - page_size = 20 - dataset = make_dataset(controller_session, tenant, account) - mock_dataset_svc.get_datasets.return_value = ([dataset] * page_size, page_size) - mock_provider_mgr.return_value.get_configurations.return_value.get_models.return_value = list[object]() - - with app.test_request_context(f"/datasets?page=1&limit={page_size}", method="GET"): - api = DatasetListApi() - response, status = unwrap(api.get)(api, controller_session, tenant_id=tenant.id) - - assert status == 200 - assert response["has_more"] is False - assert response["limit"] == page_size - assert response["total"] == page_size - assert response["page"] == 1 - - @patch("controllers.service_api.dataset.dataset.create_plugin_provider_manager") - @patch("controllers.service_api.dataset.dataset.DatasetService") - def test_list_datasets_has_more_true_when_limit_exceeds_cap( - self, - mock_dataset_svc: MagicMock, - mock_provider_mgr: MagicMock, - app: Flask, - account: Account, - tenant: Tenant, - controller_session: Session, - ) -> None: - """limit>100 still reports remaining rows after the server cap of 100.""" - from controllers.service_api.dataset.dataset import DatasetListApi - - returned_count = 100 - total = 150 - dataset = make_dataset(controller_session, tenant, account) - mock_dataset_svc.get_datasets.return_value = ([dataset] * returned_count, total) - mock_provider_mgr.return_value.get_configurations.return_value.get_models.return_value = list[object]() - - with app.test_request_context("/datasets?page=1&limit=200", method="GET"): - api = DatasetListApi() - response, status = unwrap(api.get)(api, controller_session, tenant_id=tenant.id) - - assert status == 200 - assert response["has_more"] is True - assert response["limit"] == 100 - assert response["total"] == total - assert response["page"] == 1 - assert mock_dataset_svc.get_datasets.call_args.args[1] == 100 - - @patch("controllers.service_api.dataset.dataset.create_plugin_provider_manager") - @patch("controllers.service_api.dataset.dataset.DatasetService") - def test_list_datasets_reports_the_page_size_the_query_used_for_limit_zero( - self, - mock_dataset_svc: MagicMock, - mock_provider_mgr: MagicMock, - app: Flask, - account: Account, - tenant: Tenant, - controller_session: Session, - ) -> None: - """limit=0 is floored to one row by the query, so the response must report one.""" - from controllers.service_api.dataset.dataset import DatasetListApi - - total = 7 - dataset = make_dataset(controller_session, tenant, account) - mock_dataset_svc.get_datasets.return_value = ([dataset], total) - mock_provider_mgr.return_value.get_configurations.return_value.get_models.return_value = list[object]() - - with app.test_request_context("/datasets?page=1&limit=0", method="GET"): - api = DatasetListApi() - response, status = unwrap(api.get)(api, controller_session, tenant_id=tenant.id) - - assert status == 200 - assert response["limit"] == 1 - assert response["page"] == 1 - assert response["has_more"] is True - assert mock_dataset_svc.get_datasets.call_args.args[:2] == (1, 1) - - @pytest.mark.usefixtures("account") - @patch("controllers.service_api.dataset.dataset.create_plugin_provider_manager") - @patch("controllers.service_api.dataset.dataset.DatasetService") - def test_list_datasets_has_more_false_past_the_last_page_for_limit_zero( - self, - mock_dataset_svc: MagicMock, - mock_provider_mgr: MagicMock, - app: Flask, - tenant: Tenant, - controller_session: Session, - ) -> None: - """A page past the end must end the walk instead of reporting more rows forever.""" - from controllers.service_api.dataset.dataset import DatasetListApi - - total = 7 - mock_dataset_svc.get_datasets.return_value = (list[Dataset](), total) - mock_provider_mgr.return_value.get_configurations.return_value.get_models.return_value = list[object]() - - with app.test_request_context("/datasets?page=8&limit=0", method="GET"): - api = DatasetListApi() - response, status = unwrap(api.get)(api, controller_session, tenant_id=tenant.id) - - assert status == 200 - assert response["has_more"] is False - assert response["limit"] == 1 - assert response["page"] == 8 - assert response["data"] == [] - - @patch("controllers.service_api.dataset.dataset.create_plugin_provider_manager") - @patch("controllers.service_api.dataset.dataset.DatasetService") - def test_list_datasets_reports_the_page_size_the_query_used_for_negative_limit( - self, - mock_dataset_svc: MagicMock, - mock_provider_mgr: MagicMock, - app: Flask, - account: Account, - tenant: Tenant, - controller_session: Session, - ) -> None: - """A negative limit is floored the same way, and is never echoed back.""" - from controllers.service_api.dataset.dataset import DatasetListApi - - total = 7 - dataset = make_dataset(controller_session, tenant, account) - mock_dataset_svc.get_datasets.return_value = ([dataset], total) - mock_provider_mgr.return_value.get_configurations.return_value.get_models.return_value = list[object]() - - with app.test_request_context("/datasets?page=1&limit=-1", method="GET"): - api = DatasetListApi() - response, status = unwrap(api.get)(api, controller_session, tenant_id=tenant.id) - - assert status == 200 - assert response["limit"] == 1 - assert response["page"] == 1 - assert mock_dataset_svc.get_datasets.call_args.args[:2] == (1, 1) - - @patch("controllers.service_api.dataset.dataset.create_plugin_provider_manager") - @patch("controllers.service_api.dataset.dataset.DatasetService") - def test_list_datasets_reports_the_page_the_query_used_for_page_zero( - self, - mock_dataset_svc: MagicMock, - mock_provider_mgr: MagicMock, - app: Flask, - account: Account, - tenant: Tenant, - controller_session: Session, - ) -> None: - """page=0 is served as page 1, so the response must not offer page 1 as the next one.""" - from controllers.service_api.dataset.dataset import DatasetListApi - - total = 7 - dataset = make_dataset(controller_session, tenant, account) - mock_dataset_svc.get_datasets.return_value = ([dataset] * total, total) - mock_provider_mgr.return_value.get_configurations.return_value.get_models.return_value = list[object]() - - with app.test_request_context("/datasets?page=0&limit=20", method="GET"): - api = DatasetListApi() - response, status = unwrap(api.get)(api, controller_session, tenant_id=tenant.id) - - assert status == 200 - assert response["page"] == 1 - assert response["has_more"] is False - assert mock_dataset_svc.get_datasets.call_args.args[:2] == (1, 20) - - -class TestDatasetListApiPost: - """Test suite for DatasetListApi.post() endpoint.""" - - @patch("controllers.service_api.dataset.dataset.DatasetService") - def test_create_dataset_success( - self, - mock_dataset_svc: MagicMock, - app: Flask, - account: Account, - tenant: Tenant, - controller_session: Session, - ) -> None: - from controllers.service_api.dataset.dataset import DatasetCreatePayload, DatasetListApi - - mock_dataset_svc.create_empty_dataset.return_value = make_dataset( - controller_session, tenant, account, name="New Dataset" - ) - - with app.test_request_context( - "/datasets", - method="POST", - json={"name": "New Dataset"}, - ): - api = DatasetListApi() - # `post` is wrapped in @model_validate, so the unwrapped view expects - # the validated model where the decorator would have injected it. - validated_payload = DatasetCreatePayload.model_validate(request.get_json() or {}) - response, status = unwrap(api.post)(api, validated_payload, controller_session, tenant_id=tenant.id) - - assert status == 200 - assert_dataset_detail_shape(response) - assert response["name"] == "New Dataset" - mock_dataset_svc.create_empty_dataset.assert_called_once() - - @pytest.mark.usefixtures("account") - @patch("controllers.service_api.dataset.dataset.DatasetService") - def test_create_dataset_duplicate_name( - self, - mock_dataset_svc: MagicMock, - app: Flask, - tenant: Tenant, - controller_session: Session, - ) -> None: - from controllers.service_api.dataset.dataset import DatasetCreatePayload, DatasetListApi - - mock_dataset_svc.create_empty_dataset.side_effect = services.errors.dataset.DatasetNameDuplicateError() - - with app.test_request_context( - "/datasets", - method="POST", - json={"name": "Existing Dataset"}, - ): - api = DatasetListApi() - validated_payload = DatasetCreatePayload.model_validate(request.get_json() or {}) - with pytest.raises(DatasetNameDuplicateError): - unwrap(api.post)(api, validated_payload, controller_session, tenant_id=tenant.id) - - @pytest.mark.usefixtures("account") - @patch("controllers.service_api.wraps.FeatureService") - @patch("controllers.service_api.wraps.validate_and_get_api_token") - def test_invalid_body_is_rejected_before_the_handler_runs( - self, - mock_validate_token: MagicMock, - mock_feature_svc: MagicMock, - app: Flask, - tenant: Tenant, - ) -> None: - """The tests above unwrap the view, so this is what covers the decorator.""" - from controllers.service_api.dataset.dataset import DatasetListApi - - mock_validate_token.return_value = SimpleNamespace(tenant_id=tenant.id) - mock_feature_svc.get_knowledge_rate_limit.return_value = SimpleNamespace(enabled=False) - - with app.test_request_context("/datasets", method="POST", json={}): - with pytest.raises(UnprocessableEntity): - DatasetListApi().post(tenant_id=tenant.id) - - -# --------------------------------------------------------------------------- -# API endpoint tests — DatasetApi -# --------------------------------------------------------------------------- - - -class TestDatasetApiGet: - """Test suite for DatasetApi.get() endpoint.""" - - @patch("controllers.service_api.dataset.dataset.create_plugin_provider_manager") - @patch("controllers.service_api.dataset.dataset.DatasetService") - def test_get_dataset_success( - self, - mock_dataset_svc: MagicMock, - mock_provider_mgr: MagicMock, - app: Flask, - dataset: Dataset, - controller_session: Session, - ) -> None: - from controllers.service_api.dataset.dataset import DatasetApi - - mock_dataset_svc.get_dataset.return_value = dataset - mock_dataset_svc.check_dataset_permission.return_value = None - mock_provider_mgr.return_value.get_configurations.return_value.get_models.return_value = list[object]() - - with app.test_request_context( - f"/datasets/{dataset.id}", - method="GET", - ): - api = DatasetApi() - response, status = unwrap(api.get)(api, controller_session, _=dataset.tenant_id, dataset_id=dataset.id) - - assert status == 200 - assert_dataset_detail_shape(response) - assert response["embedding_available"] is True - assert response["retrieval_model_dict"]["search_method"] == "keyword_search" - - @patch("controllers.service_api.dataset.dataset.DatasetPermissionService") - @patch("controllers.service_api.dataset.dataset.create_plugin_provider_manager") - @patch("controllers.service_api.dataset.dataset.DatasetService") - def test_get_dataset_partial_members_shape( - self, - mock_dataset_svc: MagicMock, - mock_provider_mgr: MagicMock, - mock_perm_svc: MagicMock, - app: Flask, - dataset: Dataset, - controller_session: Session, - ) -> None: - from controllers.service_api.dataset.dataset import DatasetApi - - dataset.permission = PermissionEnum.PARTIAL_TEAM - mock_dataset_svc.get_dataset.return_value = dataset - mock_dataset_svc.check_dataset_permission.return_value = None - mock_perm_svc.get_dataset_partial_member_list.return_value = ["user-1", "user-2"] - mock_provider_mgr.return_value.get_configurations.return_value.get_models.return_value = list[object]() - - with app.test_request_context( - f"/datasets/{dataset.id}", - method="GET", - ): - api = DatasetApi() - response, status = unwrap(api.get)(api, controller_session, _=dataset.tenant_id, dataset_id=dataset.id) - - assert status == 200 - assert_dataset_detail_shape(response, with_partial_members=True) - assert response["partial_member_list"] == ["user-1", "user-2"] - - @patch("controllers.service_api.dataset.dataset.create_plugin_provider_manager") - @patch("controllers.service_api.dataset.dataset.DatasetService") - def test_get_dataset_uses_default_external_retrieval_model( - self, - mock_dataset_svc: MagicMock, - mock_provider_mgr: MagicMock, - app: Flask, - dataset: Dataset, - controller_session: Session, - ) -> None: - from controllers.service_api.dataset.dataset import DatasetApi - - dataset.retrieval_model = None - mock_dataset_svc.get_dataset.return_value = dataset - mock_dataset_svc.check_dataset_permission.return_value = None - mock_provider_mgr.return_value.get_configurations.return_value.get_models.return_value = list[object]() - - with app.test_request_context(f"/datasets/{dataset.id}", method="GET"): - api = DatasetApi() - response, status = unwrap(api.get)(api, controller_session, _=dataset.tenant_id, dataset_id=dataset.id) - - assert status == 200 - assert_dataset_detail_shape(response) - assert response["external_retrieval_model"] == { - "top_k": 2, - "score_threshold": 0.0, - "score_threshold_enabled": None, - } - - @patch("controllers.service_api.dataset.dataset.DatasetService") - def test_get_dataset_not_found( - self, - mock_dataset_svc: MagicMock, - app: Flask, - dataset: Dataset, - controller_session: Session, - ) -> None: - from controllers.service_api.dataset.dataset import DatasetApi - - mock_dataset_svc.get_dataset.return_value = None - - with app.test_request_context( - f"/datasets/{dataset.id}", - method="GET", - ): - api = DatasetApi() - with pytest.raises(NotFound): - unwrap(api.get)(api, controller_session, _=dataset.tenant_id, dataset_id=dataset.id) - - @patch("controllers.service_api.dataset.dataset.DatasetService") - def test_get_dataset_no_permission( - self, - mock_dataset_svc: MagicMock, - app: Flask, - dataset: Dataset, - controller_session: Session, - ) -> None: - from controllers.service_api.dataset.dataset import DatasetApi - - mock_dataset_svc.get_dataset.return_value = dataset - mock_dataset_svc.check_dataset_permission.side_effect = services.errors.base.NoPermissionError() - - with app.test_request_context( - f"/datasets/{dataset.id}", - method="GET", - ): - api = DatasetApi() - with pytest.raises(Forbidden): - unwrap(api.get)(api, controller_session, _=dataset.tenant_id, dataset_id=dataset.id) - - -class TestDatasetApiPatch: - """Test suite for DatasetApi.patch() endpoint.""" - - @patch("controllers.service_api.dataset.dataset.DatasetPermissionService") - @patch("controllers.service_api.dataset.dataset.DatasetService") - def test_patch_dataset_success_shape( - self, - mock_dataset_svc: MagicMock, - mock_perm_svc: MagicMock, - app: Flask, - dataset: Dataset, - controller_session: Session, - ) -> None: - from controllers.service_api.dataset.dataset import DatasetApi, DatasetUpdatePayload - - dataset.name = "Updated Dataset" - mock_dataset_svc.get_dataset.return_value = dataset - mock_dataset_svc.update_dataset.return_value = dataset - mock_perm_svc.check_permission.return_value = None - mock_perm_svc.get_dataset_partial_member_list.return_value = ["user-1"] - - payload = { - "name": "Updated Dataset", - "permission": "partial_members", - "partial_member_list": [{"user_id": "user-1", "role": "editor"}], - } - with app.test_request_context( - f"/datasets/{dataset.id}", - method="PATCH", - json=payload, - ): - api = DatasetApi() - # `patch` is wrapped in @model_validate, so the unwrapped view expects - # the validated model where the decorator would have injected it. - validated_payload = DatasetUpdatePayload.model_validate(request.get_json() or {}) - response, status = unwrap(api.patch)( - api, - validated_payload, - controller_session, - _=dataset.tenant_id, - dataset_id=dataset.id, - ) - - assert status == 200 - assert_dataset_detail_shape(response, with_partial_members=True) - assert response["name"] == "Updated Dataset" - assert response["partial_member_list"] == ["user-1"] - mock_dataset_svc.update_dataset.assert_called_once() - _, update_data, _ = mock_dataset_svc.update_dataset.call_args.args - session = mock_dataset_svc.update_dataset.call_args.kwargs["session"] - assert session is controller_session - assert update_data["name"] == "Updated Dataset" - assert update_data["permission"] == "partial_members" - mock_perm_svc.update_partial_member_list.assert_called_once_with( - dataset.tenant_id, - dataset.id, - [{"user_id": "user-1", "role": "editor"}], - controller_session, - ) - - -class TestDatasetApiDelete: - """Test suite for DatasetApi.delete() endpoint.""" - - @patch("controllers.service_api.dataset.dataset.DatasetPermissionService") - @patch("controllers.service_api.dataset.dataset.DatasetService") - def test_delete_dataset_success( - self, - mock_dataset_svc: MagicMock, - mock_perm_svc: MagicMock, - app: Flask, - dataset: Dataset, - controller_session: Session, - ) -> None: - from controllers.service_api.dataset.dataset import DatasetApi - - mock_dataset_svc.delete_dataset.return_value = True - - with app.test_request_context( - f"/datasets/{dataset.id}", - method="DELETE", - ): - api = DatasetApi() - result = unwrap(api.delete)(api, controller_session, _=dataset.tenant_id, dataset_id=dataset.id) - - assert result == ("", 204) - mock_perm_svc.clear_partial_member_list.assert_called_once_with(dataset.id, controller_session) - - @patch("controllers.service_api.dataset.dataset.DatasetService") - def test_delete_dataset_not_found( - self, - mock_dataset_svc: MagicMock, - app: Flask, - dataset: Dataset, - controller_session: Session, - ) -> None: - from controllers.service_api.dataset.dataset import DatasetApi - - mock_dataset_svc.delete_dataset.return_value = False - - with app.test_request_context( - f"/datasets/{dataset.id}", - method="DELETE", - ): - api = DatasetApi() - with pytest.raises(NotFound): - unwrap(api.delete)(api, controller_session, _=dataset.tenant_id, dataset_id=dataset.id) - - @patch("controllers.service_api.dataset.dataset.DatasetService") - def test_delete_dataset_in_use( - self, - mock_dataset_svc: MagicMock, - app: Flask, - dataset: Dataset, - controller_session: Session, - ) -> None: - from controllers.service_api.dataset.dataset import DatasetApi - - mock_dataset_svc.delete_dataset.side_effect = services.errors.dataset.DatasetInUseError() - - with app.test_request_context( - f"/datasets/{dataset.id}", - method="DELETE", - ): - api = DatasetApi() - with pytest.raises(DatasetInUseError): - unwrap(api.delete)(api, controller_session, _=dataset.tenant_id, dataset_id=dataset.id) - - -# --------------------------------------------------------------------------- -# API endpoint tests — DocumentStatusApi -# --------------------------------------------------------------------------- - - -class TestDocumentStatusApiPatch: - """Test suite for DocumentStatusApi.patch() endpoint.""" - - @patch("controllers.service_api.dataset.dataset.DocumentService") - @patch("controllers.service_api.dataset.dataset.DatasetService") - def test_batch_update_status_success( - self, - mock_dataset_svc: MagicMock, - mock_doc_svc: MagicMock, - app: Flask, - tenant: Tenant, - dataset: Dataset, - ) -> None: - from controllers.service_api.dataset.dataset import DocumentStatusApi - - mock_dataset_svc.get_dataset.return_value = dataset - mock_dataset_svc.check_dataset_permission.return_value = None - mock_dataset_svc.check_dataset_model_setting.return_value = None - mock_doc_svc.batch_update_document_status.return_value = None - - with app.test_request_context( - f"/datasets/{dataset.id}/documents/status/enable", - method="PATCH", - json={"document_ids": ["doc-1", "doc-2"]}, - ): - api = DocumentStatusApi() - response, status = api.patch( - tenant_id=tenant.id, - dataset_id=dataset.id, - action="enable", - ) - - assert status == 200 - assert response["result"] == "success" - - @patch("controllers.service_api.dataset.dataset.DatasetService") - def test_batch_update_status_dataset_not_found( - self, - mock_dataset_svc: MagicMock, - app: Flask, - tenant: Tenant, - dataset: Dataset, - ) -> None: - from controllers.service_api.dataset.dataset import DocumentStatusApi - - mock_dataset_svc.get_dataset.return_value = None - - with app.test_request_context( - f"/datasets/{dataset.id}/documents/status/enable", - method="PATCH", - json={"document_ids": ["doc-1"]}, - ): - api = DocumentStatusApi() - with pytest.raises(NotFound): - api.patch( - tenant_id=tenant.id, - dataset_id=dataset.id, - action="enable", - ) - - @patch("controllers.service_api.dataset.dataset.DatasetService") - def test_batch_update_status_permission_error( - self, - mock_dataset_svc: MagicMock, - app: Flask, - tenant: Tenant, - dataset: Dataset, - ) -> None: - from controllers.service_api.dataset.dataset import DocumentStatusApi - - mock_dataset_svc.get_dataset.return_value = dataset - mock_dataset_svc.check_dataset_permission.side_effect = services.errors.base.NoPermissionError("No permission") - - with app.test_request_context( - f"/datasets/{dataset.id}/documents/status/enable", - method="PATCH", - json={"document_ids": ["doc-1"]}, - ): - api = DocumentStatusApi() - with pytest.raises(Forbidden): - api.patch( - tenant_id=tenant.id, - dataset_id=dataset.id, - action="enable", - ) - - @patch("controllers.service_api.dataset.dataset.DocumentService") - @patch("controllers.service_api.dataset.dataset.DatasetService") - def test_batch_update_status_indexing_error( - self, - mock_dataset_svc: MagicMock, - mock_doc_svc: MagicMock, - app: Flask, - tenant: Tenant, - dataset: Dataset, - ) -> None: - from controllers.service_api.dataset.dataset import DocumentStatusApi - - mock_dataset_svc.get_dataset.return_value = dataset - mock_dataset_svc.check_dataset_permission.return_value = None - mock_dataset_svc.check_dataset_model_setting.return_value = None - mock_doc_svc.batch_update_document_status.side_effect = services.errors.document.DocumentIndexingError() - - with app.test_request_context( - f"/datasets/{dataset.id}/documents/status/enable", - method="PATCH", - json={"document_ids": ["doc-1"]}, - ): - api = DocumentStatusApi() - with pytest.raises(InvalidActionError): - api.patch( - tenant_id=tenant.id, - dataset_id=dataset.id, - action="enable", - ) - - @patch("controllers.service_api.dataset.dataset.DocumentService") - @patch("controllers.service_api.dataset.dataset.DatasetService") - def test_batch_update_status_value_error( - self, - mock_dataset_svc: MagicMock, - mock_doc_svc: MagicMock, - app: Flask, - tenant: Tenant, - dataset: Dataset, - ) -> None: - from controllers.service_api.dataset.dataset import DocumentStatusApi - - mock_dataset_svc.get_dataset.return_value = dataset - mock_dataset_svc.check_dataset_permission.return_value = None - mock_dataset_svc.check_dataset_model_setting.return_value = None - mock_doc_svc.batch_update_document_status.side_effect = ValueError("Invalid action") - - with app.test_request_context( - f"/datasets/{dataset.id}/documents/status/enable", - method="PATCH", - json={"document_ids": ["doc-1"]}, - ): - api = DocumentStatusApi() - with pytest.raises(InvalidActionError): - api.patch( - tenant_id=tenant.id, - dataset_id=dataset.id, - action="enable", - ) diff --git a/api/tests/unit_tests/controllers/service_api/dataset/test_dataset_payloads.py b/api/tests/unit_tests/controllers/service_api/dataset/test_dataset_payloads.py deleted file mode 100644 index dc88606fe94..00000000000 --- a/api/tests/unit_tests/controllers/service_api/dataset/test_dataset_payloads.py +++ /dev/null @@ -1,210 +0,0 @@ -"""Unit tests for Service API dataset request payloads.""" - -from typing import Literal - -import pytest - -from controllers.service_api.dataset.dataset import ( - DatasetCreatePayload, - DatasetListQuery, - DatasetUpdatePayload, - TagBindingPayload, - TagCreatePayload, - TagDeletePayload, - TagUnbindingPayload, - TagUpdatePayload, -) -from models.dataset import DatasetPermissionEnum - - -class TestDatasetCreatePayload: - """Test suite for DatasetCreatePayload Pydantic model.""" - - def test_payload_with_required_name(self) -> None: - payload = DatasetCreatePayload(name="Test Dataset") - assert payload.name == "Test Dataset" - assert payload.description == "" - assert payload.permission == DatasetPermissionEnum.ONLY_ME - - def test_payload_with_all_fields(self) -> None: - payload = DatasetCreatePayload( - name="Full Dataset", - description="A comprehensive dataset description", - indexing_technique="high_quality", - permission=DatasetPermissionEnum.ALL_TEAM, - provider="vendor", - embedding_model="text-embedding-ada-002", - embedding_model_provider="openai", - ) - assert payload.name == "Full Dataset" - assert payload.description == "A comprehensive dataset description" - assert payload.indexing_technique == "high_quality" - assert payload.permission == DatasetPermissionEnum.ALL_TEAM - assert payload.provider == "vendor" - assert payload.embedding_model == "text-embedding-ada-002" - assert payload.embedding_model_provider == "openai" - - def test_payload_name_length_validation_min(self) -> None: - with pytest.raises(ValueError, match="1 validation error for DatasetCreatePayload\\\nname"): - DatasetCreatePayload(name="") - - def test_payload_name_length_validation_max(self) -> None: - with pytest.raises(ValueError, match="1 validation error for DatasetCreatePayload\\\nname"): - DatasetCreatePayload(name="A" * 41) - - def test_payload_description_max_length(self) -> None: - with pytest.raises(ValueError, match="1 validation error for DatasetCreatePayload\\\ndescription"): - DatasetCreatePayload(name="Dataset", description="A" * 401) - - @pytest.mark.parametrize("technique", ["high_quality", "economy"]) - def test_payload_valid_indexing_techniques(self, technique: Literal["high_quality", "economy"]) -> None: - payload = DatasetCreatePayload(name="Dataset", indexing_technique=technique) - assert payload.indexing_technique == technique - - def test_payload_with_external_knowledge_settings(self) -> None: - payload = DatasetCreatePayload( - name="External Dataset", external_knowledge_api_id="api_123", external_knowledge_id="knowledge_456" - ) - assert payload.external_knowledge_api_id == "api_123" - assert payload.external_knowledge_id == "knowledge_456" - - -class TestDatasetUpdatePayload: - """Test suite for DatasetUpdatePayload Pydantic model.""" - - def test_payload_all_optional(self) -> None: - payload = DatasetUpdatePayload() - assert payload.name is None - assert payload.description is None - assert payload.permission is None - - def test_payload_with_partial_update(self) -> None: - payload = DatasetUpdatePayload(name="Updated Name", description="Updated description") - assert payload.name == "Updated Name" - assert payload.description == "Updated description" - - def test_payload_with_permission_change(self) -> None: - payload = DatasetUpdatePayload( - permission=DatasetPermissionEnum.PARTIAL_TEAM, - partial_member_list=[{"user_id": "user_123", "role": "editor"}], - ) - assert payload.permission == DatasetPermissionEnum.PARTIAL_TEAM - assert payload.partial_member_list is not None - assert len(payload.partial_member_list) == 1 - - def test_payload_name_length_validation(self) -> None: - with pytest.raises(ValueError, match="1 validation error for DatasetUpdatePayload\\\nname"): - DatasetUpdatePayload(name="") - with pytest.raises(ValueError, match="1 validation error for DatasetUpdatePayload\\\nname"): - DatasetUpdatePayload(name="A" * 41) - - -class TestDatasetListQuery: - """Test suite for DatasetListQuery Pydantic model.""" - - def test_query_with_defaults(self) -> None: - query = DatasetListQuery() - assert query.page == 1 - assert query.limit == 20 - assert query.keyword is None - assert query.include_all is False - assert query.tag_ids == [] - - def test_query_with_all_filters(self) -> None: - query = DatasetListQuery( - page=3, limit=50, keyword="machine learning", include_all=True, tag_ids=["tag1", "tag2", "tag3"] - ) - assert query.page == 3 - assert query.limit == 50 - assert query.keyword == "machine learning" - assert query.include_all is True - assert len(query.tag_ids) == 3 - - def test_query_with_tag_filter(self) -> None: - query = DatasetListQuery(tag_ids=["tag_abc", "tag_def"]) - assert query.tag_ids == ["tag_abc", "tag_def"] - - -class TestTagCreatePayload: - """Test suite for TagCreatePayload Pydantic model.""" - - def test_payload_with_name(self) -> None: - payload = TagCreatePayload(name="New Tag") - assert payload.name == "New Tag" - - def test_payload_name_length_min(self) -> None: - with pytest.raises(ValueError, match="1 validation error for TagCreatePayload\\\nname"): - TagCreatePayload(name="") - - def test_payload_name_length_max(self) -> None: - with pytest.raises(ValueError, match="1 validation error for TagCreatePayload\\\nname"): - TagCreatePayload(name="A" * 51) - - def test_payload_with_unicode_name(self) -> None: - payload = TagCreatePayload(name="标签 🏷️ Тег") - assert payload.name == "标签 🏷️ Тег" - - -class TestTagUpdatePayload: - """Test suite for TagUpdatePayload Pydantic model.""" - - def test_payload_with_name_and_id(self) -> None: - payload = TagUpdatePayload(name="Updated Tag", tag_id="tag_123") - assert payload.name == "Updated Tag" - assert payload.tag_id == "tag_123" - - def test_payload_requires_tag_id(self) -> None: - with pytest.raises(ValueError, match="1 validation error for TagUpdatePayload\\\ntag_id"): - TagUpdatePayload.model_validate({"name": "Updated Tag"}) - - -class TestTagDeletePayload: - """Test suite for TagDeletePayload Pydantic model.""" - - def test_payload_with_tag_id(self) -> None: - payload = TagDeletePayload(tag_id="tag_to_delete") - assert payload.tag_id == "tag_to_delete" - - def test_payload_requires_tag_id(self) -> None: - with pytest.raises(ValueError, match="1 validation error for TagDeletePayload\\\ntag_id"): - TagDeletePayload.model_validate({}) - - -class TestTagBindingPayload: - """Test suite for TagBindingPayload Pydantic model.""" - - def test_payload_with_valid_data(self) -> None: - payload = TagBindingPayload(tag_ids=["tag1", "tag2"], target_id="dataset_123") - assert len(payload.tag_ids) == 2 - assert payload.target_id == "dataset_123" - - def test_payload_rejects_empty_tag_ids(self) -> None: - with pytest.raises(ValueError, match="1 validation error for TagBindingPayload\\\ntag_ids") as exc_info: - TagBindingPayload(tag_ids=[], target_id="dataset_123") - assert "Tag IDs is required" in str(exc_info.value) - - def test_payload_single_tag_id(self) -> None: - payload = TagBindingPayload(tag_ids=["single_tag"], target_id="dataset_456") - assert payload.tag_ids == ["single_tag"] - - -class TestTagUnbindingPayload: - """Test suite for TagUnbindingPayload Pydantic model.""" - - def test_payload_with_valid_data(self) -> None: - payload = TagUnbindingPayload(tag_ids=["tag_123"], target_id="dataset_456") - assert payload.tag_ids == ["tag_123"] - assert payload.target_id == "dataset_456" - - def test_payload_normalizes_legacy_tag_id(self) -> None: - payload = TagUnbindingPayload(tag_id="tag_123", target_id="dataset_456") - assert payload.tag_ids == ["tag_123"] - assert payload.target_id == "dataset_456" - - def test_payload_rejects_empty_tag_ids(self) -> None: - with pytest.raises( - ValueError, - match=("1 validation error for TagUnbindingPayload\\\n Value error, Tag IDs is required\\."), - ) as exc_info: - TagUnbindingPayload(tag_ids=[], target_id="dataset_456") - assert "Tag IDs is required" in str(exc_info.value) diff --git a/api/tests/unit_tests/controllers/service_api/dataset/test_dataset_tag_apis.py b/api/tests/unit_tests/controllers/service_api/dataset/test_dataset_tag_apis.py deleted file mode 100644 index 55aae9be652..00000000000 --- a/api/tests/unit_tests/controllers/service_api/dataset/test_dataset_tag_apis.py +++ /dev/null @@ -1,378 +0,0 @@ -"""Unit tests for Service API dataset tag controller behavior. - -Service boundaries stay mocked, while users, tenants, and tags are real ORM objects -persisted in SQLite. Controller database calls share that SQLite session so assertions -cover the concrete objects and session passed across the controller boundary. -""" - -import uuid -from inspect import unwrap -from unittest.mock import MagicMock, patch - -import pytest -from flask import Flask, request -from sqlalchemy.orm import Session -from werkzeug.exceptions import Forbidden - -from models.account import Account, Tenant, TenantAccountRole -from models.enums import TagType -from models.model import Tag - -TAG_MODEL_TABLES = (Account, Tenant, Tag) -pytestmark = pytest.mark.parametrize("sqlite_session", [TAG_MODEL_TABLES], indirect=True) - - -@pytest.fixture -def controller_session(sqlite_session: Session) -> Session: - """Expose the shared SQLite session under the controller-focused fixture name.""" - return sqlite_session - - -@pytest.fixture -def tenant(controller_session: Session) -> Tenant: - tenant = Tenant(name="Dataset Tag API Tenant") - controller_session.add(tenant) - controller_session.flush() - return tenant - - -@pytest.fixture -def account(controller_session: Session, tenant: Tenant, monkeypatch: pytest.MonkeyPatch) -> Account: - account = Account(name="Dataset Tag API User", email=f"dataset-tag-api-{uuid.uuid4()}@example.com") - account.role = TenantAccountRole.OWNER - account._current_tenant = tenant - controller_session.add(account) - controller_session.flush() - - # Inject the concrete account at the controller boundary without relying on Flask-Login globals. - from controllers.service_api.dataset import dataset as dataset_module - - monkeypatch.setattr(dataset_module, "current_user", account) - return account - - -def make_tag( - session: Session, - tenant: Tenant, - account: Account, - *, - id: str, - name: str, - binding_count: int | None = None, -) -> Tag: - """Create and flush a real tag, optionally adding the aggregate count returned by TagService.""" - - tag = Tag(tenant_id=tenant.id, type=TagType.KNOWLEDGE, name=name, created_by=account.id) - tag.id = id - session.add(tag) - session.flush() - if binding_count is not None: - tag.__dict__["binding_count"] = binding_count - return tag - - -class TestDatasetTagsApiGet: - """Test suite for DatasetTagsApi.get() endpoint.""" - - @patch("controllers.service_api.dataset.dataset.TagService") - def test_list_tags_success( - self, - mock_tag_svc: MagicMock, - app: Flask, - account: Account, - tenant: Tenant, - controller_session: Session, - ) -> None: - from controllers.service_api.dataset.dataset import DatasetTagsApi - - tag = make_tag(controller_session, tenant, account, id="tag-1", name="Test Tag", binding_count=0) - mock_tag_svc.get_tags.return_value = [tag] - - with app.test_request_context("/datasets/tags", method="GET"): - api = DatasetTagsApi() - response, status = unwrap(api.get)(api, controller_session, _=None) - - assert status == 200 - assert response == [{"id": "tag-1", "name": "Test Tag", "type": "knowledge", "binding_count": "0"}] - mock_tag_svc.get_tags.assert_called_once_with("knowledge", tenant.id, session=controller_session) - - -class TestDatasetTagsApiPost: - """Test suite for DatasetTagsApi.post() endpoint.""" - - @patch("controllers.service_api.dataset.dataset.TagService") - def test_create_tag_success( - self, - mock_tag_svc: MagicMock, - app: Flask, - account: Account, - tenant: Tenant, - controller_session: Session, - ) -> None: - from controllers.service_api.dataset.dataset import DatasetTagsApi, TagCreatePayload - - tag = make_tag(controller_session, tenant, account, id="tag-new", name="New Tag") - mock_tag_svc.save_tags.return_value = tag - - with app.test_request_context( - "/datasets/tags", - method="POST", - json={"name": "New Tag"}, - ): - api = DatasetTagsApi() - payload = TagCreatePayload.model_validate(request.get_json() or {}) - response, status = unwrap(api.post)(api, payload, controller_session, _=None) - - assert status == 200 - assert response == {"id": "tag-new", "name": "New Tag", "type": "knowledge", "binding_count": "0"} - mock_tag_svc.save_tags.assert_called_once() - - def test_create_tag_forbidden(self, app: Flask, account: Account) -> None: - from controllers.service_api.dataset.dataset import DatasetTagsApi - - account.role = TenantAccountRole.NORMAL - - with app.test_request_context( - "/datasets/tags", - method="POST", - json={"name": "New Tag"}, - ): - api = DatasetTagsApi() - with pytest.raises(Forbidden): - api.post(_=None) - - -class TestDatasetTagsApiPatch: - """Test suite for DatasetTagsApi.patch() endpoint.""" - - @patch("controllers.service_api.dataset.dataset.TagService") - @patch("controllers.service_api.dataset.dataset.service_api_ns") - def test_update_tag_success( - self, - mock_service_api_ns: MagicMock, - mock_tag_svc: MagicMock, - app: Flask, - account: Account, - tenant: Tenant, - controller_session: Session, - ) -> None: - from controllers.service_api.dataset.dataset import DatasetTagsApi, TagUpdatePayload - - tag = make_tag(controller_session, tenant, account, id="tag-1", name="Updated Tag") - mock_tag_svc.update_tags.return_value = tag - mock_tag_svc.get_tag_binding_count.return_value = 5 - mock_service_api_ns.payload = {"name": "Updated Tag", "tag_id": "tag-1"} - - with app.test_request_context( - "/datasets/tags", - method="PATCH", - json={"name": "Updated Tag", "tag_id": "tag-1"}, - ): - api = DatasetTagsApi() - payload = TagUpdatePayload.model_validate(request.get_json() or {}) - response, status = unwrap(api.patch)(api, payload, controller_session, _=None) - - assert status == 200 - assert response == {"id": "tag-1", "name": "Updated Tag", "type": "knowledge", "binding_count": "5"} - mock_tag_svc.update_tags.assert_called_once() - update_payload, tag_id, session = mock_tag_svc.update_tags.call_args.args - assert update_payload.name == "Updated Tag" - assert tag_id == "tag-1" - assert session is controller_session - - def test_update_tag_forbidden(self, app: Flask, account: Account) -> None: - from controllers.service_api.dataset.dataset import DatasetTagsApi - - account.role = TenantAccountRole.NORMAL - - with app.test_request_context( - "/datasets/tags", - method="PATCH", - json={"name": "Updated Tag", "tag_id": "tag-1"}, - ): - api = DatasetTagsApi() - with pytest.raises(Forbidden): - api.patch(_=None) - - -class TestDatasetTagsApiDelete: - """Test suite for DatasetTagsApi.delete() endpoint.""" - - @pytest.mark.usefixtures("account") - @patch("controllers.service_api.dataset.dataset.TagService") - @patch("controllers.service_api.dataset.dataset.service_api_ns") - def test_delete_tag_success( - self, - mock_service_api_ns: MagicMock, - mock_tag_svc: MagicMock, - app: Flask, - controller_session: Session, - ) -> None: - from controllers.service_api.dataset.dataset import DatasetTagsApi, TagDeletePayload - - mock_tag_svc.delete_tag.return_value = None - mock_service_api_ns.payload = {"tag_id": "tag-1"} - - with app.test_request_context( - "/datasets/tags", - method="DELETE", - json={"tag_id": "tag-1"}, - ): - api = DatasetTagsApi() - payload = TagDeletePayload.model_validate(request.get_json() or {}) - result = unwrap(api.delete)(api, payload, controller_session, _=None) - - assert result == ("", 204) - mock_tag_svc.delete_tag.assert_called_once_with("tag-1", controller_session, tag_type=TagType.KNOWLEDGE) - - -class TestDatasetTagsBindingStatusApi: - """Test suite for DatasetTagsBindingStatusApi endpoints.""" - - @patch("controllers.service_api.dataset.dataset.TagService") - def test_get_dataset_tags_binding_status( - self, - mock_tag_svc: MagicMock, - app: Flask, - account: Account, - tenant: Tenant, - controller_session: Session, - ) -> None: - from controllers.service_api.dataset.dataset import DatasetTagsBindingStatusApi - - tag = make_tag(controller_session, tenant, account, id="tag_1", name="Test Tag") - mock_tag_svc.get_tags_by_target_id.return_value = [tag] - - with app.test_request_context("/", method="GET"): - api = DatasetTagsBindingStatusApi() - response, status_code = unwrap(api.get)(api, controller_session, tenant.id, dataset_id="dataset_123") - - assert status_code == 200 - assert response["data"] == [{"id": "tag_1", "name": "Test Tag"}] - assert response["total"] == 1 - mock_tag_svc.get_tags_by_target_id.assert_called_once_with( - "knowledge", tenant.id, "dataset_123", controller_session - ) - - -class TestDatasetTagBindingApiPost: - """Test suite for DatasetTagBindingApi.post() endpoint.""" - - @pytest.mark.usefixtures("account") - @patch("controllers.service_api.dataset.dataset.TagService") - def test_bind_tags_success( - self, - mock_tag_svc: MagicMock, - app: Flask, - controller_session: Session, - ) -> None: - from controllers.service_api.dataset.dataset import DatasetTagBindingApi, TagBindingPayload - - mock_tag_svc.save_tag_binding.return_value = None - - with app.test_request_context( - "/datasets/tags/binding", - method="POST", - json={"tag_ids": ["tag-1"], "target_id": "ds-1"}, - ): - api = DatasetTagBindingApi() - payload = TagBindingPayload.model_validate(request.get_json() or {}) - result = unwrap(api.post)(api, payload, controller_session, _=None) - - assert result == ("", 204) - from services.tag_service import TagBindingCreatePayload - - mock_tag_svc.save_tag_binding.assert_called_once_with( - TagBindingCreatePayload(tag_ids=["tag-1"], target_id="ds-1", type=TagType.KNOWLEDGE), - controller_session, - ) - - def test_bind_tags_forbidden(self, app: Flask, account: Account) -> None: - from controllers.service_api.dataset.dataset import DatasetTagBindingApi - - account.role = TenantAccountRole.NORMAL - - with app.test_request_context( - "/datasets/tags/binding", - method="POST", - json={"tag_ids": ["tag-1"], "target_id": "ds-1"}, - ): - api = DatasetTagBindingApi() - with pytest.raises(Forbidden): - api.post(_=None) - - -class TestDatasetTagUnbindingApiPost: - """Test suite for DatasetTagUnbindingApi.post() endpoint.""" - - @pytest.mark.usefixtures("account") - @patch("controllers.service_api.dataset.dataset.TagService") - def test_unbind_tag_success( - self, - mock_tag_svc: MagicMock, - app: Flask, - controller_session: Session, - ) -> None: - from controllers.service_api.dataset.dataset import DatasetTagUnbindingApi, TagUnbindingPayload - - mock_tag_svc.delete_tag_binding.return_value = None - - with app.test_request_context( - "/datasets/tags/unbinding", - method="POST", - json={"tag_ids": ["tag-1"], "target_id": "ds-1"}, - ): - api = DatasetTagUnbindingApi() - payload = TagUnbindingPayload.model_validate(request.get_json() or {}) - result = unwrap(api.post)(api, payload, controller_session, _=None) - - assert result == ("", 204) - from services.tag_service import TagBindingDeletePayload - - mock_tag_svc.delete_tag_binding.assert_called_once_with( - TagBindingDeletePayload(tag_ids=["tag-1"], target_id="ds-1", type=TagType.KNOWLEDGE), - controller_session, - ) - - @pytest.mark.usefixtures("account") - @patch("controllers.service_api.dataset.dataset.TagService") - def test_unbind_legacy_tag_id_success( - self, - mock_tag_svc: MagicMock, - app: Flask, - controller_session: Session, - ) -> None: - from controllers.service_api.dataset.dataset import DatasetTagUnbindingApi, TagUnbindingPayload - - mock_tag_svc.delete_tag_binding.return_value = None - - with app.test_request_context( - "/datasets/tags/unbinding", - method="POST", - json={"tag_id": "tag-1", "target_id": "ds-1"}, - ): - api = DatasetTagUnbindingApi() - payload = TagUnbindingPayload.model_validate(request.get_json() or {}) - result = unwrap(api.post)(api, payload, controller_session, _=None) - - assert result == ("", 204) - from services.tag_service import TagBindingDeletePayload - - mock_tag_svc.delete_tag_binding.assert_called_once_with( - TagBindingDeletePayload(tag_ids=["tag-1"], target_id="ds-1", type=TagType.KNOWLEDGE), - controller_session, - ) - - def test_unbind_tag_forbidden(self, app: Flask, account: Account) -> None: - from controllers.service_api.dataset.dataset import DatasetTagUnbindingApi - - account.role = TenantAccountRole.NORMAL - - with app.test_request_context( - "/datasets/tags/unbinding", - method="POST", - json={"tag_ids": ["tag-1"], "target_id": "ds-1"}, - ): - api = DatasetTagUnbindingApi() - with pytest.raises(Forbidden): - api.post(_=None) diff --git a/api/tests/unit_tests/controllers/service_api/dataset/test_document.py b/api/tests/unit_tests/controllers/service_api/dataset/test_document.py index a89b15bd2eb..e59853690a3 100644 --- a/api/tests/unit_tests/controllers/service_api/dataset/test_document.py +++ b/api/tests/unit_tests/controllers/service_api/dataset/test_document.py @@ -65,8 +65,8 @@ from services.errors.file import FileTooLargeError as FileTooLargeServiceError from services.knowledge.dataset_service import DocumentService from services.knowledge.entities.knowledge_entities import ProcessRule, RetrievalModel from services.knowledge.resource_scope import DatasetRef -from tests.unit_tests.config_override import config_overrides_context -from tests.unit_tests.model_factories import make_account, make_upload_file +from tests.helpers.unit.config_override import config_overrides_context +from tests.helpers.unit.model_factories import make_account, make_upload_file def _document_data_source_info() -> dict[str, str]: diff --git a/api/tests/unit_tests/controllers/service_api/dataset/test_rag_pipeline_file_upload_serialization.py b/api/tests/unit_tests/controllers/service_api/dataset/test_rag_pipeline_file_upload_serialization.py deleted file mode 100644 index aa2b1f644af..00000000000 --- a/api/tests/unit_tests/controllers/service_api/dataset/test_rag_pipeline_file_upload_serialization.py +++ /dev/null @@ -1,37 +0,0 @@ -""" -Unit tests for Service API knowledge pipeline file-upload serialization. -""" - -from datetime import UTC, datetime - -from controllers.service_api.dataset.rag_pipeline.rag_pipeline_workflow import PipelineUploadFileResponse -from libs.helper import dump_response -from models.model import UploadFile -from tests.unit_tests.model_factories import make_upload_file - - -def _upload_file(*, created_at: datetime) -> UploadFile: - return make_upload_file( - file_id="file-1", - name="test.pdf", - size=123, - extension="pdf", - mime_type="application/pdf", - created_at=created_at, - ) - - -def test_file_upload_created_at_is_isoformat_string(): - created_at = datetime(2026, 2, 8, 12, 0, 0, tzinfo=UTC) - upload_file = _upload_file(created_at=created_at) - - result = dump_response(PipelineUploadFileResponse, upload_file) - assert result["created_at"] == created_at.isoformat() - - -def test_file_upload_created_at_none_serializes_to_null(): - upload_file = _upload_file(created_at=datetime(2026, 2, 8, 12, 0, 0, tzinfo=UTC)) - upload_file.created_at = None - - result = dump_response(PipelineUploadFileResponse, upload_file) - assert result["created_at"] is None diff --git a/api/tests/unit_tests/controllers/service_api/dataset/test_rag_pipeline_route_registration.py b/api/tests/unit_tests/controllers/service_api/dataset/test_rag_pipeline_route_registration.py deleted file mode 100644 index 184e37014b6..00000000000 --- a/api/tests/unit_tests/controllers/service_api/dataset/test_rag_pipeline_route_registration.py +++ /dev/null @@ -1,54 +0,0 @@ -""" -Unit tests for Service API knowledge pipeline route registration. -""" - -import ast -from pathlib import Path - - -def test_rag_pipeline_routes_registered(): - api_dir = Path(__file__).resolve().parents[5] - - service_api_init = api_dir / "controllers" / "service_api" / "__init__.py" - rag_pipeline_workflow = ( - api_dir / "controllers" / "service_api" / "dataset" / "rag_pipeline" / "rag_pipeline_workflow.py" - ) - - assert service_api_init.exists() - assert rag_pipeline_workflow.exists() - - init_tree = ast.parse(service_api_init.read_text(encoding="utf-8")) - import_found = False - for node in ast.walk(init_tree): - if not isinstance(node, ast.ImportFrom): - continue - if node.module != "dataset.rag_pipeline" or node.level != 1: - continue - if any(alias.name == "rag_pipeline_workflow" for alias in node.names): - import_found = True - break - assert import_found, "from .dataset.rag_pipeline import rag_pipeline_workflow not found in service_api/__init__.py" - - workflow_tree = ast.parse(rag_pipeline_workflow.read_text(encoding="utf-8")) - route_paths: set[str] = set() - - for node in ast.walk(workflow_tree): - if not isinstance(node, ast.ClassDef): - continue - for decorator in node.decorator_list: - if not isinstance(decorator, ast.Call): - continue - if not isinstance(decorator.func, ast.Attribute): - continue - if decorator.func.attr != "route": - continue - if not decorator.args: - continue - first_arg = decorator.args[0] - if isinstance(first_arg, ast.Constant) and isinstance(first_arg.value, str): - route_paths.add(first_arg.value) - - assert "/datasets//pipeline/datasource-plugins" in route_paths - assert "/datasets//pipeline/datasource/nodes//run" in route_paths - assert "/datasets//pipeline/run" in route_paths - assert "/datasets/pipeline/file-upload" in route_paths diff --git a/api/tests/unit_tests/controllers/service_api/dataset/test_dataset_segment.py b/api/tests/unit_tests/controllers/service_api/dataset/test_segment.py similarity index 99% rename from api/tests/unit_tests/controllers/service_api/dataset/test_dataset_segment.py rename to api/tests/unit_tests/controllers/service_api/dataset/test_segment.py index be28636eae7..9685ed00bf5 100644 --- a/api/tests/unit_tests/controllers/service_api/dataset/test_dataset_segment.py +++ b/api/tests/unit_tests/controllers/service_api/dataset/test_segment.py @@ -45,7 +45,7 @@ from models.enums import DataSourceType, DocumentCreatedFrom, IndexingStatus, Se from services.api_token_service import CachedApiToken from services.knowledge.dataset_service import DocumentService, SegmentService from services.knowledge.segments.application import SegmentMutationService -from tests.unit_tests.model_factories import make_account +from tests.helpers.unit.model_factories import make_account @pytest.fixture(autouse=True) diff --git a/api/tests/unit_tests/controllers/test_knowledge_fs_service_contracts.py b/api/tests/unit_tests/controllers/service_api/knowledge_fs/test_resources.py similarity index 99% rename from api/tests/unit_tests/controllers/test_knowledge_fs_service_contracts.py rename to api/tests/unit_tests/controllers/service_api/knowledge_fs/test_resources.py index 4295596d119..22eb54ecce2 100644 --- a/api/tests/unit_tests/controllers/test_knowledge_fs_service_contracts.py +++ b/api/tests/unit_tests/controllers/service_api/knowledge_fs/test_resources.py @@ -30,14 +30,15 @@ from services.knowledge_fs.product_dto import ( KnowledgeFSServiceSourceUpdatePayload, ) from services.knowledge_fs.service_api_authorization import KnowledgeFSServiceApiProfile -from tests.unit_tests.config_override import apply_config_overrides -from tests.unit_tests.controllers.knowledge_fs_controller_fakes import ( +from tests.helpers.unit.config_override import apply_config_overrides +from tests.helpers.unit.controllers.knowledge_fs_controller_fakes import ( FakeComponent, FakeRuntime, Recorder, invocation, narrowed, ) +from tests.unit_tests.conftest import app as app # noqa: PLC0414 _NOW = "2026-09-17T00:00:00Z" _FAILURE = { diff --git a/api/tests/unit_tests/controllers/test_swagger.py b/api/tests/unit_tests/controllers/service_api/test___init__.py similarity index 99% rename from api/tests/unit_tests/controllers/test_swagger.py rename to api/tests/unit_tests/controllers/service_api/test___init__.py index 3c49bd90f89..c779271eb2c 100644 --- a/api/tests/unit_tests/controllers/test_swagger.py +++ b/api/tests/unit_tests/controllers/service_api/test___init__.py @@ -7,6 +7,8 @@ from pathlib import Path import pytest from flask import Blueprint, Flask +from tests.unit_tests.conftest import app as app # noqa: PLC0414 + @pytest.fixture(autouse=True) def _swagger_config(config_overrides) -> None: diff --git a/api/tests/unit_tests/controllers/service_api/test_trace_session_id_parsing.py b/api/tests/unit_tests/controllers/service_api/test_trace_session_id_parsing.py deleted file mode 100644 index 05fcc5e400d..00000000000 --- a/api/tests/unit_tests/controllers/service_api/test_trace_session_id_parsing.py +++ /dev/null @@ -1,181 +0,0 @@ -from unittest.mock import MagicMock, patch - -import pytest -from flask import Flask -from werkzeug.exceptions import BadRequest - -from controllers.service_api.app import completion as completion_module -from controllers.service_api.app import workflow as workflow_module -from core.helper.trace_id_helper import get_trace_session_id -from models.enums import EndUserType -from models.model import App, AppMode, EndUser -from tests.unit_tests.model_factories import make_app, make_end_user - - -class _Request: - def __init__(self, *, headers=None, args=None, json=None, is_json=True): - self.headers = headers or {} - self.args = args or {} - self.json = json - self.is_json = is_json - - -def test_trace_session_id_header_query_body_priority_matches_service_api_contract(): - req = _Request( - headers={"X-Trace-Session-Id": "header"}, - args={"trace_session_id": "query"}, - json={"trace_session_id": "body"}, - ) - - assert get_trace_session_id(req) == "header" - - -def test_trace_session_id_invalid_highest_priority_raises_bad_request(): - req = _Request( - headers={"X-Trace-Session-Id": " "}, - args={"trace_session_id": "query"}, - json={"trace_session_id": "body"}, - ) - - with pytest.raises(BadRequest): - get_trace_session_id(req) - - -def _app(mode: AppMode) -> App: - return make_app(mode=mode, icon_type=None) - - -def _end_user() -> EndUser: - return make_end_user(end_user_id="user-1", end_user_type=EndUserType.SERVICE_API) - - -def _assert_generate_trace_session_id(mock_generate_service: MagicMock, expected: str) -> None: - _, kwargs = mock_generate_service.generate.call_args - assert kwargs["args"]["trace_session_id"] == expected - - -@patch("controllers.service_api.app.completion.AppGenerateService") -@patch("controllers.service_api.app.completion.service_api_ns") -def test_chat_api_rejects_invalid_highest_priority_query_trace_session_id_without_generating( - mock_service_api_ns: MagicMock, - mock_generate_service: MagicMock, - app: Flask, -): - payload = {"inputs": {}, "query": "hello", "trace_session_id": "body-session"} - mock_service_api_ns.payload = payload - - with app.test_request_context( - "/chat-messages?trace_session_id=%20%20%20", - method="POST", - json=payload, - ): - with pytest.raises(BadRequest): - completion_module.ChatApi().post.__wrapped__( - completion_module.ChatApi(), - _app(AppMode.CHAT), - _end_user(), - ) - - mock_generate_service.generate.assert_not_called() - - -@patch("controllers.service_api.app.workflow.AppGenerateService") -@patch("controllers.service_api.app.workflow.service_api_ns") -def test_workflow_run_api_rejects_invalid_highest_priority_body_trace_session_id_without_generating( - mock_service_api_ns: MagicMock, - mock_generate_service: MagicMock, - app: Flask, -): - payload = {"inputs": {}, "trace_session_id": 123} - mock_service_api_ns.payload = payload - - with app.test_request_context("/workflows/run", method="POST", json=payload): - with pytest.raises(BadRequest): - workflow_module.WorkflowRunApi().post.__wrapped__( - workflow_module.WorkflowRunApi(), - _app(AppMode.WORKFLOW), - _end_user(), - ) - - mock_generate_service.generate.assert_not_called() - - -@patch("controllers.service_api.app.completion.helper.compact_generate_response", return_value={"answer": "ok"}) -@patch("controllers.service_api.app.completion.AppGenerateService") -@patch("controllers.service_api.app.completion.service_api_ns") -def test_completion_api_passes_header_trace_session_id_when_body_value_is_invalid_lower_priority( - mock_service_api_ns: MagicMock, - mock_generate_service: MagicMock, - mock_compact: MagicMock, - app: Flask, -): - payload = {"inputs": {}, "trace_session_id": 123} - mock_service_api_ns.payload = payload - mock_generate_service.generate.return_value = "response" - - with app.test_request_context( - "/completion-messages", - method="POST", - json=payload, - headers={"X-Trace-Session-Id": " header-session "}, - ): - response = completion_module.CompletionApi().post.__wrapped__( - completion_module.CompletionApi(), - _app(AppMode.COMPLETION), - _end_user(), - ) - - assert response == {"answer": "ok"} - _assert_generate_trace_session_id(mock_generate_service, "header-session") - - -@patch("controllers.service_api.app.completion.helper.compact_generate_response", return_value={"answer": "ok"}) -@patch("controllers.service_api.app.completion.AppGenerateService") -@patch("controllers.service_api.app.completion.service_api_ns") -def test_chat_api_passes_query_trace_session_id_when_body_value_is_invalid_lower_priority( - mock_service_api_ns: MagicMock, - mock_generate_service: MagicMock, - mock_compact: MagicMock, - app: Flask, -): - payload = {"inputs": {}, "query": "hello", "trace_session_id": 123} - mock_service_api_ns.payload = payload - mock_generate_service.generate.return_value = "response" - - with app.test_request_context( - "/chat-messages?trace_session_id=query-session", - method="POST", - json=payload, - ): - response = completion_module.ChatApi().post.__wrapped__( - completion_module.ChatApi(), - _app(AppMode.CHAT), - _end_user(), - ) - - assert response == {"answer": "ok"} - _assert_generate_trace_session_id(mock_generate_service, "query-session") - - -@patch("controllers.service_api.app.workflow.helper.compact_generate_response", return_value={"result": "ok"}) -@patch("controllers.service_api.app.workflow.AppGenerateService") -@patch("controllers.service_api.app.workflow.service_api_ns") -def test_workflow_run_api_passes_body_trace_session_id( - mock_service_api_ns: MagicMock, - mock_generate_service: MagicMock, - mock_compact: MagicMock, - app: Flask, -): - payload = {"inputs": {}, "trace_session_id": " body-session "} - mock_service_api_ns.payload = payload - mock_generate_service.generate.return_value = "response" - - with app.test_request_context("/workflows/run", method="POST", json=payload): - response = workflow_module.WorkflowRunApi().post.__wrapped__( - workflow_module.WorkflowRunApi(), - _app(AppMode.WORKFLOW), - _end_user(), - ) - - assert response == {"result": "ok"} - _assert_generate_trace_session_id(mock_generate_service, "body-session") diff --git a/api/tests/unit_tests/controllers/service_api/test_wraps.py b/api/tests/unit_tests/controllers/service_api/test_wraps.py index a49b1af8df1..c16cab461f3 100644 --- a/api/tests/unit_tests/controllers/service_api/test_wraps.py +++ b/api/tests/unit_tests/controllers/service_api/test_wraps.py @@ -41,7 +41,7 @@ from models.resource_access_token import ( ) from services.api_token_service import ApiTokenCache from services.knowledge_fs.control_space_commands import KnowledgeFSControlSpaceCommandService -from tests.unit_tests.config_override import config_overrides_context +from tests.helpers.unit.config_override import config_overrides_context class _RecordingEndUserCommands: diff --git a/api/tests/unit_tests/controllers/test_app_import_quota.py b/api/tests/unit_tests/controllers/test_app_import_quota.py deleted file mode 100644 index d01fa8f3df4..00000000000 --- a/api/tests/unit_tests/controllers/test_app_import_quota.py +++ /dev/null @@ -1,134 +0,0 @@ -"""Quota failures retain their HTTP semantics across both app import surfaces.""" - -from collections.abc import Callable -from inspect import unwrap -from typing import Protocol, cast -from uuid import uuid4 - -import httpx -import pytest -from flask import Flask -from flask_restx import Resource -from sqlalchemy import select -from sqlalchemy.orm import Session - -from constants.dsl_version import CURRENT_APP_DSL_VERSION -from controllers.console.app import app_import as console -from controllers.inner_api.app import dsl as inner -from controllers.openapi import app_dsl as openapi -from controllers.openapi._errors import OpenApiErrorFormatter -from controllers.openapi._models import AppDslImportPayload -from enums import DeploymentEdition -from extensions.ext_redis import redis_client -from libs.external_api import ExternalApi -from machinery.context import RequestContext -from models.account import Account, Tenant, TenantAccountJoin, TenantAccountRole -from models.model import App -from services.app_dsl_service import AppDslService, PendingData -from services.billing_service import BillingService, _BillingHTTPStatusError -from services.entities.feature_entities import FeatureModel, LimitationModel -from services.feature_service import FeatureService -from tests.unit_tests.config_override import apply_config_overrides -from tests.unit_tests.controllers.conftest import ControllerTestServices - - -class _EndpointView(Protocol): - __handler__: Callable[..., tuple[object, int]] - - -@pytest.mark.parametrize( - ("surface", "confirm"), - [("console", False), ("console", True), ("openapi", False), ("openapi", True), ("inner", False)], -) -@pytest.mark.parametrize("failure", ["timeout", "upstream_503", "invalid_response", "quota_exhausted"]) -def test_import_quota_failure_http_status( - app_query_services: ControllerTestServices, - sqlite_session: Session, - monkeypatch: pytest.MonkeyPatch, - surface: str, - confirm: bool, - failure: str, -) -> None: - apply_config_overrides( - monkeypatch, DEPLOYMENT_EDITION=DeploymentEdition.CLOUD, INNER_API=True, INNER_API_KEY="test-inner-api-key" - ) - monkeypatch.setattr("controllers.console.wraps._is_setup_completed", lambda: True) - account = Account(name="Importer", email="importer@example.com") - tenant = Tenant(name="Workspace") - account.id, tenant.id = str(uuid4()), str(uuid4()) - sqlite_session.add_all( - [ - account, - tenant, - TenantAccountJoin(tenant_id=tenant.id, account_id=account.id, role=TenantAccountRole.OWNER), - ] - ) - sqlite_session.commit() - context = RequestContext("request-1", None, account.id, tenant.id) - content = ( - f"version: {CURRENT_APP_DSL_VERSION}\napp: {{name: Imported, mode: workflow}}\n" - "workflow: {graph: {nodes: [], edges: []}, features: {}}\n" - ) - pending = PendingData(tenant_id=tenant.id, account_id=account.id, import_mode="yaml-content", yaml_content=content) - monkeypatch.setattr(redis_client, "get", lambda _: pending.model_dump_json()) - - def unexpected_side_effect(*_args: object, **_kwargs: object) -> None: - pytest.fail("A failed quota check must preserve pending data and avoid database writes") - - monkeypatch.setattr(redis_client, "delete", unexpected_side_effect) - monkeypatch.setattr(AppDslService, "_create_or_update_app", unexpected_side_effect) - if failure == "quota_exhausted": - monkeypatch.setattr( - FeatureService, - "get_features", - lambda *_args, **_kwargs: FeatureModel(apps=LimitationModel(size=5, limit=5)), - ) - else: - - def failed_billing(*_args: object, **_kwargs: object) -> dict[str, str]: - if failure == "timeout": - raise httpx.ReadTimeout("Private billing timeout details") - if failure == "upstream_503": - raise _BillingHTTPStatusError("Private billing failure details", 503) - return {"subscription": "Private invalid billing response"} - - monkeypatch.setattr(BillingService, "_send_request", failed_billing) - - app = Flask(__name__) - app.config["TESTING"] = True - app.extensions["application_services"] = app_query_services - api = ExternalApi(app, error_body_formatter=OpenApiErrorFormatter() if surface == "openapi" else None) - - class AdmittedImport(Resource): - def post(self) -> tuple[object, int]: - if surface == "inner": - return inner.EnterpriseAppDSLImport().post(workspace_id=tenant.id) - if surface == "console": - if confirm: - controller = console.AppImportConfirmApi() - return unwrap(controller.post)(controller, context, import_id="pending-1") - controller = console.AppImportApi() - return unwrap(controller.post)(controller, context) - if confirm: - controller = openapi.AppDslImportConfirmApi() - return cast(_EndpointView, controller.post).__handler__( - controller, context, workspace_id=tenant.id, import_id="pending-1" - ) - controller = openapi.AppDslImportApi() - return cast(_EndpointView, controller.post).__handler__( - controller, - context, - workspace_id=tenant.id, - body=AppDslImportPayload(mode="yaml-content", yaml_content=content), - ) - - api.add_resource(AdmittedImport, "/imports") - response = app.test_client().post( - "/imports", - json={"mode": "yaml-content", "yaml_content": content, "creator_email": account.email}, - headers={"X-Inner-Api-Key": "test-inner-api-key"}, - ) - - assert response.status_code == (403 if failure == "quota_exhausted" else 500) - assert "Private" not in response.get_data(as_text=True) - assert sqlite_session.scalar(select(App.id).where(App.tenant_id == tenant.id)) is None diff --git a/api/tests/unit_tests/controllers/test_compare_versions.py b/api/tests/unit_tests/controllers/test_compare_versions.py deleted file mode 100644 index f4c296898aa..00000000000 --- a/api/tests/unit_tests/controllers/test_compare_versions.py +++ /dev/null @@ -1,24 +0,0 @@ -import pytest - -from controllers.console.system import _has_new_version - - -@pytest.mark.parametrize( - ("latest_version", "current_version", "expected"), - [ - ("1.0.1", "1.0.0", True), - ("1.1.0", "1.0.0", True), - ("2.0.0", "1.9.9", True), - ("1.0.0", "1.0.0", False), - ("1.0.0", "1.0.1", False), - ("1.0.0", "2.0.0", False), - ("1.0.1", "1.0.0-beta", True), - ("1.0.0", "1.0.0-alpha", True), - ("1.0.0-beta", "1.0.0-alpha", True), - ("1.0.0", "1.0.0-rc1", True), - ("1.0.0", "0.9.9", True), - ("1.0.0", "1.0.0-dev", True), - ], -) -def test_has_new_version(latest_version: str, current_version: str, expected: bool) -> None: - assert _has_new_version(latest_version=latest_version, current_version=current_version) == expected diff --git a/api/tests/unit_tests/controllers/test_conversation_rename_payload.py b/api/tests/unit_tests/controllers/test_conversation_rename_payload.py deleted file mode 100644 index 494176cbd9e..00000000000 --- a/api/tests/unit_tests/controllers/test_conversation_rename_payload.py +++ /dev/null @@ -1,20 +0,0 @@ -import pytest -from pydantic import ValidationError - -from controllers.console.explore.conversation import ConversationRenamePayload as ConsolePayload -from controllers.service_api.app.conversation import ConversationRenamePayload as ServicePayload - - -@pytest.mark.parametrize("payload_cls", [ConsolePayload, ServicePayload]) -def test_payload_allows_auto_generate_without_name(payload_cls): - payload = payload_cls.model_validate({"auto_generate": True}) - - assert payload.auto_generate is True - assert payload.name is None - - -@pytest.mark.parametrize("payload_cls", [ConsolePayload, ServicePayload]) -@pytest.mark.parametrize("value", [None, "", " "]) -def test_payload_requires_name_when_not_auto_generate(payload_cls, value): - with pytest.raises(ValidationError): - payload_cls.model_validate({"name": value, "auto_generate": False}) diff --git a/api/tests/unit_tests/controllers/test_error_response_contract_check.py b/api/tests/unit_tests/controllers/test_error_response_contract_check.py deleted file mode 100644 index 085379e97ea..00000000000 --- a/api/tests/unit_tests/controllers/test_error_response_contract_check.py +++ /dev/null @@ -1,195 +0,0 @@ -import pytest -from flask import Blueprint, Flask, Response, json, request_finished -from flask_restx import Namespace, Resource -from pydantic import BaseModel - -from controllers.common.schema import register_response_schema_models -from libs.exception import BaseHTTPException, RelayedHTTPError -from libs.external_api import ExternalApi, error_responses -from tests.unit_tests.controllers import error_contract -from tests.unit_tests.controllers.error_contract import check_response, violation - - -class _ItemGoneError(BaseHTTPException): - error_code = "item_gone" - code = 404 - - -class _ItemLockedDetails(BaseModel): - locked_by: str - - -class _ItemLockedError(BaseHTTPException): - error_code = "item_locked" - code = 404 - details_model = _ItemLockedDetails - - -class _LegacyConflict(BaseModel): - error: str - - -def _api() -> tuple[Flask, ExternalApi]: - app = Flask(__name__) - blueprint = Blueprint("contract_check", __name__, url_prefix="/check") - api = ExternalApi(blueprint) - ns = Namespace("items", path="/") - register_response_schema_models(ns, _LegacyConflict) - - @ns.route("/items") - class Items(Resource): - @ns.response(409, "Legacy conflict", ns.models[_LegacyConflict.__name__]) - @error_responses(_ItemGoneError, _ItemLockedError) - def get(self) -> dict[str, str]: - return {} - - @ns.route("/relayed") - class Relayed(Resource): - @error_responses(_ItemGoneError, RelayedHTTPError) - def get(self) -> dict[str, str]: - return {} - - api.add_namespace(ns) - app.register_blueprint(blueprint) - return app, api - - -def _check(body: object, status: int, path: str = "/check/items") -> str | None: - app, api = _api() - with app.test_request_context(path): - return violation(api, Response(json.dumps(body), status=status, mimetype="application/json")) - - -def test_rejects_a_body_without_the_standard_fields() -> None: - found = _check({"error": "gone"}, 404) - assert found is not None - assert "lacks" in found - - -def test_rejects_a_code_the_operation_does_not_document_for_the_status() -> None: - found = _check({"code": "surprise", "message": "m", "status": 404}, 404) - assert found is not None - assert "'surprise' is not documented" in found - - -def test_accepts_documented_and_implied_codes() -> None: - assert _check({"code": "item_gone", "message": "m", "status": 404}, 404) is None - assert _check({"code": "not_found", "message": "m", "status": 404}, 404) is None - - -def test_generic_codes_satisfy_the_full_schema_of_a_status_that_documents_business_errors() -> None: - # A generic 404 answers an operation that documents its own 404 codes: the status's complete - # schema (not only ``default``) must admit it, so the open error body stays in the ``anyOf``. - app, api = _api() - with app.test_request_context("/check/items"): - schema = error_contract._document(api)["paths"]["/items"]["get"]["responses"]["404"]["content"][ # type: ignore[index] - "application/json" - ]["schema"] - assert schema["anyOf"][-1] == {"$ref": "#/components/schemas/ErrorResponse"} - assert _check({"code": "not_found", "message": "m", "status": 404}, 404) is None - - -def test_validates_a_documented_code_against_its_component() -> None: - assert _check({"code": "item_locked", "message": "m", "status": 404, "details": {"locked_by": "a"}}, 404) is None - - found = _check({"code": "item_locked", "message": "m", "status": 404, "details": {}}, 404) - assert found is not None - assert "'item_locked' component" in found - assert "locked_by" in found - - -def test_accepts_a_body_matching_a_legacy_model_the_status_declares() -> None: - assert _check({"error": "legacy"}, 409) is None - assert _check({"code": "conflict", "message": "m", "status": 409}, 409) is None - - -def test_rejects_a_documented_code_served_with_another_status(monkeypatch: pytest.MonkeyPatch) -> None: - # On the surfaces the lint covers every served code must be documented or generic. - monkeypatch.setattr(error_contract, "_documents_business_errors", lambda _api: True) - found = _check({"code": "item_gone", "message": "m", "status": 409}, 409) - assert found is not None - assert "'item_gone' is not documented" in found - - -def test_accepts_any_code_from_an_operation_that_relays_errors(monkeypatch: pytest.MonkeyPatch) -> None: - monkeypatch.setattr(error_contract, "_documents_business_errors", lambda _api: True) - assert _check({"code": "downstream_busy", "message": "m", "status": 503}, 503, "/check/relayed") is None - assert _check({"code": "item_gone", "message": "m", "status": 404}, 404, "/check/relayed") is None - # Relaying never excuses a body that is not the error body. - found = _check({"error": "x"}, 503, "/check/relayed") - assert found is not None - assert "lacks" in found - - -def test_only_checks_documented_statuses_on_apis_without_business_errors() -> None: - # A test-built API (like /openapi/v1) documents nothing, so an undocumented status stays open. - assert _check({"code": "item_gone", "message": "m", "status": 409}, 409) is None - - -def test_rejects_a_body_matching_neither_the_legacy_nor_the_standard_model() -> None: - found = _check({"code": "conflict", "message": "m"}, 409) - assert found is not None - assert "lacks ['status']" in found - - -@pytest.mark.parametrize( - ("body", "field"), - [ - ({"code": "item_gone", "message": {"nested": True}, "status": 404}, "message"), - ({"code": ["item_gone"], "message": "m", "status": 404}, "code"), - ({"code": "item_gone", "message": "m", "status": "404"}, "status"), - ], -) -def test_rejects_standard_fields_of_the_wrong_type(body: dict[str, object], field: str) -> None: - found = _check(body, 404) - assert found is not None - assert "does not match ErrorResponse" in found - assert f"{field}:" in found - - -@pytest.mark.allow_error_contract_violations -def test_checks_responses_served_by_an_api_mounted_on_the_app() -> None: - app = Flask(__name__) - api = ExternalApi(app) - - @api.route("/items") - class Items(Resource): - def get(self) -> tuple[dict[str, str], int]: - return {"wrong": "shape"}, 400 - - violations: list[str] = [] - - def receiver(sender: Flask, response: Response, **_: object) -> None: - check_response(sender, response, violations) - - with request_finished.connected_to(receiver, app): - app.test_client().get("/items") - - assert len(violations) == 1 - assert "lacks" in violations[0] - - -@pytest.mark.allow_error_contract_violations -def test_reports_responses_it_cannot_check() -> None: - app = Flask(__name__) - api = ExternalApi(app) - - class _Unregistered(BaseModel): - value: str - - @api.route("/items") - class Items(Resource): - @api.response(200, "OK", _Unregistered.__name__) - def get(self) -> tuple[dict[str, str], int]: - return {"wrong": "shape"}, 400 - - violations: list[str] = [] - - def receiver(sender: Flask, response: Response, **_: object) -> None: - check_response(sender, response, violations) - - with request_finished.connected_to(receiver, app): - app.test_client().get("/items") - - assert len(violations) == 1 - assert "cannot check the error contract" in violations[0] diff --git a/api/tests/unit_tests/controllers/test_knowledge_fs_product_controllers.py b/api/tests/unit_tests/controllers/test_knowledge_fs_product_controllers.py deleted file mode 100644 index ccb12b1dce4..00000000000 --- a/api/tests/unit_tests/controllers/test_knowledge_fs_product_controllers.py +++ /dev/null @@ -1,1889 +0,0 @@ -from __future__ import annotations - -import ast -import inspect -from collections.abc import Callable, Generator -from contextlib import contextmanager -from dataclasses import dataclass -from datetime import UTC, datetime -from io import BytesIO -from pathlib import Path -from types import FunctionType -from typing import Never - -import pytest -from flask import Flask -from werkzeug.exceptions import Forbidden, NotFound, ServiceUnavailable - -from controllers.common import wraps as common_wraps -from controllers.console import console_ns -from controllers.console import wraps as console_wraps -from controllers.console.knowledge_fs import resources as console_resources -from controllers.console.knowledge_fs.error import ( - KnowledgeFSConflictHTTPError, - KnowledgeFSRequestTooLargeHTTPError, - KnowledgeFSResourceNotFoundHTTPError, -) -from controllers.console.wraps import RBACPermission -from controllers.service_api import service_api_ns -from controllers.service_api.knowledge_fs import resources as service_resources -from enums import DeploymentEdition -from models import Account -from models.enums import ApiTokenType -from models.model import ApiToken -from services.knowledge_fs.capability_broker import KnowledgeFSIssuedProductCapability -from services.knowledge_fs.control_space_commands import KnowledgeFSControlSpaceIntentConflictError -from services.knowledge_fs.download_service import KnowledgeFSDownloadUnavailableError -from services.knowledge_fs.object_storage import KnowledgeFSObjectMetadata -from services.knowledge_fs.product_dto import ( - KnowledgeFSAdmittedQueryRequest, - KnowledgeFSDocumentDownloadDescriptor, - KnowledgeFSDocumentMultimodalManifest, - KnowledgeFSDocumentStagedUploadAcceptedResponse, - KnowledgeFSDocumentStagedUploadPayload, - KnowledgeFSDocumentUploadAcceptedResponse, - KnowledgeFSDurableDeletionAcceptedResponse, - KnowledgeFSLogicalDocumentDeletePayload, - KnowledgeFSNamespacePreviewJobResponse, - KnowledgeFSNamespacePreviewPageResponse, - KnowledgeFSSmallFileUploadResponse, - KnowledgeFSSpaceCreatePayload, - KnowledgeFSStagedUploadResponse, -) -from services.knowledge_fs.product_remote import ( - KnowledgeFSProductRequestRejectedError, - KnowledgeFSProductResourceNotFoundError, - KnowledgeFSRemoteMultipartFile, - KnowledgeFSRemoteSSEResponse, -) -from services.knowledge_fs.service_api_authorization import ( - KnowledgeFSServiceApiAuthorizationError, - KnowledgeFSServiceApiProfile, - KnowledgeFSServiceApiScopeError, -) -from tests.unit_tests.config_override import apply_config_overrides -from tests.unit_tests.controllers.knowledge_fs_controller_fakes import ( - FakeComponent, - FakeRuntime, - Recorder, - invocation, - narrowed, -) - -_API_ROOT = Path(__file__).resolve().parents[3] - - -@pytest.fixture(autouse=True) -def _enabled_knowledge_fs(config_overrides: Callable[..., None]) -> None: - config_overrides( - DEPLOYMENT_EDITION=DeploymentEdition.COMMUNITY, - KNOWLEDGE_FS_ENABLED=True, - KNOWLEDGE_FS_COMMUNITY_DEV_ENABLED=True, - ) - - -def test_disabled_knowledge_fs_rejects_console_work_before_preview_upload_or_query_admission( - config_overrides: Callable[..., None], -) -> None: - config_overrides(KNOWLEDGE_FS_ENABLED=False) - routes = ( - ( - console_resources.KnowledgeFSInitialSourcePreviewApi, - console_resources.KnowledgeFSInitialSourcePreviewApi.post, - (), - ), - ( - console_resources.KnowledgeFSInitialSourcePreviewJobsApi, - console_resources.KnowledgeFSInitialSourcePreviewJobsApi.post, - (), - ), - (console_resources.KnowledgeFSStagedUploadsApi, console_resources.KnowledgeFSStagedUploadsApi.post, ()), - ( - console_resources.KnowledgeFSSpaceQueryAdmissionApi, - console_resources.KnowledgeFSSpaceQueryAdmissionApi.post, - ("space-1",), - ), - ) - for resource_type, view, args in routes: - handler = inspect.unwrap( - view, - stop=lambda view: "_knowledge_fs_errors..decorated" in view.__code__.co_qualname, - ) - with pytest.raises(NotFound): - handler(resource_type(), *args) - - -def test_disabled_knowledge_fs_hides_jwks_even_when_capability_v2_is_configured( - config_overrides: Callable[..., None], - monkeypatch: pytest.MonkeyPatch, -) -> None: - config_overrides(KNOWLEDGE_FS_ENABLED=False, KNOWLEDGE_FS_CAPABILITY_V2_ENABLED=True) - monkeypatch.setattr( - console_resources.session_factory, - "get_session_maker", - lambda: pytest.fail("disabled JWKS must not access the database"), - ) - with pytest.raises(NotFound): - console_resources.KnowledgeFSJWKSApi().get() - - -def test_disabled_knowledge_fs_rejects_service_api_before_token_validation( - config_overrides: Callable[..., None], -) -> None: - config_overrides(KNOWLEDGE_FS_ENABLED=False) - handler = inspect.unwrap( - service_resources.KnowledgeFSServiceSourceProvidersApi.get, - stop=lambda view: "_service_api_errors..decorated" in view.__code__.co_qualname, - ) - with pytest.raises(NotFound): - handler(service_resources.KnowledgeFSServiceSourceProvidersApi(), "space-1") - - -@dataclass(frozen=True) -class _FacadeRuntime[T]: - facade: T - - -@dataclass(frozen=True) -class _BrokerRuntime[T]: - broker: T - - -def _account(account_id: str = "account-1") -> Account: - account = Account(name="Test User", email="user@example.com") - account.id = account_id - return account - - -def _dataset_api_token(tenant_id: str | None = "tenant-1") -> ApiToken: - api_token = ApiToken() - api_token.id = "token-1" - api_token.tenant_id = tenant_id - api_token.type = ApiTokenType.DATASET - api_token.token = "dataset-test-key" - return api_token - - -def _issued_capability( - *, token: str, operation_id: str, knowledge_space_id: str = "space-1" -) -> KnowledgeFSIssuedProductCapability: - return KnowledgeFSIssuedProductCapability( - token=token, - expires_at=datetime(2030, 1, 1, tzinfo=UTC), - operation_id=operation_id, - knowledge_space_id=knowledge_space_id, - knowledge_space_revision=1, - trace_id="trace-1", - ) - - -def _service_profile(control_space_id: str) -> KnowledgeFSServiceApiProfile: - return KnowledgeFSServiceApiProfile( - tenant_id="tenant-1", - control_space_id=control_space_id, - api_token_id="token-1", - principal_id="token-1", - knowledge_space_id="space-1", - knowledge_space_revision=1, - membership_epoch=0, - space_acl_epoch=0, - external_access_epoch=0, - content_policy_revision=0, - ) - - -def _rbac_wrapper(view: FunctionType) -> FunctionType: - current = inspect.unwrap( - view, - stop=lambda candidate: ( - isinstance(candidate, FunctionType) and "rbac_permission_required" in candidate.__code__.co_qualname - ), - ) - assert isinstance(current, FunctionType) - assert "rbac_permission_required" in current.__code__.co_qualname, "RBAC permission wrapper is missing" - return current - - -def test_console_and_service_api_routes_are_registered() -> None: - console_urls = {url for route in console_ns.resources for url in route.urls if url.startswith("/knowledge-fs/")} - service_urls = {url for route in service_api_ns.resources for url in route.urls if url.startswith("/knowledge-fs/")} - - assert { - "/knowledge-fs/spaces", - "/knowledge-fs/uploads", - "/knowledge-fs/uploads/", - "/knowledge-fs/spaces/", - "/knowledge-fs/spaces//tags", - "/knowledge-fs/spaces//permissions", - "/knowledge-fs/spaces//members", - "/knowledge-fs/spaces//app-bindings", - ("/knowledge-fs/spaces//app-bindings//"), - "/knowledge-fs/spaces//external-access", - "/knowledge-fs/spaces//settings", - "/knowledge-fs/spaces//documents", - "/knowledge-fs/spaces//documents//multimodal", - ( - "/knowledge-fs/spaces//documents//multimodal/" - "/asset" - ), - "/knowledge-fs/spaces//logical-documents", - "/knowledge-fs/spaces//logical-documents/download-zip", - "/knowledge-fs/spaces//logical-documents/", - "/knowledge-fs/spaces//logical-documents//download", - "/knowledge-fs/spaces//sources", - "/knowledge-fs/spaces//source-connections", - "/knowledge-fs/spaces//source-providers", - ("/knowledge-fs/spaces//source-connections//refresh"), - "/knowledge-fs/spaces//sources//sync", - "/knowledge-fs/spaces//sources//crawl-preview", - "/knowledge-fs/spaces//sources//workflow-imports", - "/knowledge-fs/spaces//sources//sync-policy", - "/knowledge-fs/spaces//source-workflows/", - "/knowledge-fs/spaces//source-workflows//cancel", - "/knowledge-fs/spaces//source-workflows//retry", - "/knowledge-fs/spaces//source-workflows//pages", - "/knowledge-fs/spaces//source-workflows//selection", - "/knowledge-fs/spaces//sources//async-import", - "/knowledge-fs/spaces//queries", - "/knowledge-fs/spaces//research-tasks", - "/knowledge-fs/spaces//traces", - "/knowledge-fs/spaces//golden-questions/bulk-import", - "/knowledge-fs/spaces//golden-questions/evidence-matches", - "/knowledge-fs/spaces//upload-sessions", - ( - "/knowledge-fs/spaces//upload-sessions/" - "/parts//presign" - ), - ("/knowledge-fs/spaces//upload-sessions//complete"), - ("/knowledge-fs/spaces//upload-sessions//abort"), - ("/knowledge-fs/spaces//upload-sessions//small-file"), - "/knowledge-fs/spaces//query-stream-capability", - "/knowledge-fs/tasks//stream-capability", - "/knowledge-fs/query-stream", - "/knowledge-fs/research-tasks//events", - "/knowledge-fs/.well-known/jwks.json", - }.issubset(console_urls) - assert service_urls == { - "/knowledge-fs/spaces//logical-documents", - "/knowledge-fs/spaces//logical-documents/", - "/knowledge-fs/spaces//document-references/resolve", - "/knowledge-fs/spaces//deletion-jobs/", - "/knowledge-fs/spaces//deletion-jobs//retry", - "/knowledge-fs/spaces//deletion-batches/", - "/knowledge-fs/spaces//research-tasks//resume", - "/knowledge-fs/spaces//background-tasks", - "/knowledge-fs/spaces//settings/migrations/", - "/knowledge-fs/spaces//sources//sync", - "/knowledge-fs/spaces//sources//crawl-preview", - "/knowledge-fs/spaces//sources//crawl-import", - "/knowledge-fs/spaces//sources//workflow-imports", - "/knowledge-fs/spaces//sources//sync-policy", - "/knowledge-fs/spaces//source-connections", - "/knowledge-fs/spaces//source-providers", - "/knowledge-fs/spaces//query-images", - "/knowledge-fs/spaces//source-connections//refresh", - "/knowledge-fs/spaces//source-workflows", - "/knowledge-fs/spaces//source-workflows/", - "/knowledge-fs/spaces//source-workflows//cancel", - "/knowledge-fs/spaces//source-workflows//retry", - "/knowledge-fs/spaces//source-workflows//pages", - "/knowledge-fs/spaces//source-workflows//selection", - "/knowledge-fs/spaces//logical-documents//processing-tasks", - "/knowledge-fs/spaces//logical-documents//processing-tasks/", - "/knowledge-fs/spaces//bulk-jobs/", - "/knowledge-fs/spaces//documents", - "/knowledge-fs/spaces//documents/", - "/knowledge-fs/spaces//documents//outline", - "/knowledge-fs/spaces//documents//revisions", - "/knowledge-fs/spaces//documents//revisions//chunks", - "/knowledge-fs/spaces//documents//revisions//chunks/", - "/knowledge-fs/spaces//documents/bulk", - "/knowledge-fs/spaces//documents/reindex", - "/knowledge-fs/spaces//jobs/", - "/knowledge-fs/spaces//jobs//retry", - "/knowledge-fs/spaces//queries/admission", - "/knowledge-fs/query-stream", - "/knowledge-fs/spaces//research-tasks", - "/knowledge-fs/spaces//research-tasks/", - "/knowledge-fs/spaces//research-tasks//partials", - "/knowledge-fs/spaces//research-tasks/plan", - "/knowledge-fs/spaces//settings", - "/knowledge-fs/spaces//sources", - "/knowledge-fs/spaces//sources/", - "/knowledge-fs/spaces//sources//crawl", - "/knowledge-fs/spaces//sources//files", - "/knowledge-fs/spaces//sources//import", - "/knowledge-fs/spaces//sources//import-files", - "/knowledge-fs/spaces//sources//pages", - "/knowledge-fs/spaces//sources//test", - "/knowledge-fs/spaces//traces", - "/knowledge-fs/spaces//traces/", - "/knowledge-fs/spaces//traces//conflicts", - "/knowledge-fs/spaces//traces//evidence", - "/knowledge-fs/spaces//traces//missing", - } - - -def test_initial_website_preview_job_routes_delegate_to_namespace_capabilities( - monkeypatch: pytest.MonkeyPatch, -) -> None: - account = _account() - preview_service = FakeComponent(require_visible_credential=Recorder()) - facade = FakeComponent( - create_namespace_source_preview=Recorder( - KnowledgeFSNamespacePreviewJobResponse( - job_id="job-1", - status="queued", - configuration_fingerprint="f" * 64, - expires_at="2030-01-01T00:00:00Z", - ) - ), - get_namespace_source_preview=Recorder( - KnowledgeFSNamespacePreviewJobResponse( - job_id="job-1", - status="completed", - configuration_fingerprint="f" * 64, - expires_at="2030-01-01T00:00:00Z", - pages=[ - KnowledgeFSNamespacePreviewPageResponse( - page_id="page-1", - source_url="https://example.com/page", - title="Page", - ) - ], - ) - ), - cancel_namespace_source_preview=Recorder( - KnowledgeFSNamespacePreviewJobResponse( - job_id="job-1", - status="consumed", - configuration_fingerprint="f" * 64, - expires_at="2030-01-01T00:00:00Z", - ) - ), - ) - monkeypatch.setattr(console_resources, "current_account_with_tenant", lambda: (account, "tenant-1")) - monkeypatch.setattr(console_resources, "KnowledgeFSInitialSourcePreviewService", lambda _: preview_service) - monkeypatch.setattr(console_resources, "_console_services", lambda: FakeRuntime(facade=facade)) - app = Flask(__name__) - - with app.test_request_context( - method="POST", - json={ - "kind": "website_crawl", - "credentialId": "credential-1", - "datasource": "website", - "pluginId": "plugin-1", - "provider": "firecrawl", - "parameters": {"url": "https://example.com"}, - }, - ): - response, status = inspect.unwrap(console_resources.KnowledgeFSInitialSourcePreviewJobsApi.post)( - console_resources.KnowledgeFSInitialSourcePreviewJobsApi() - ) - assert status == 202 - assert response == {"job_id": "job-1", "status": "pending"} - assert len(preview_service.require_visible_credential.calls) == 1 - assert len(facade.create_namespace_source_preview.calls) == 1 - - with app.test_request_context(): - response = inspect.unwrap(console_resources.KnowledgeFSInitialSourcePreviewJobApi.get)( - console_resources.KnowledgeFSInitialSourcePreviewJobApi(), "job-1" - ) - assert response["status"] == "completed" - assert response["result"]["pages"] == [ - { - "description": None, - "page_id": "page-1", - "source_url": "https://example.com/page", - "title": "Page", - } - ] - - facade.get_namespace_source_preview.return_value = KnowledgeFSNamespacePreviewJobResponse( - job_id="job-1", - status="queued", - configuration_fingerprint="f" * 64, - expires_at="2030-01-01T00:00:00Z", - ) - with app.test_request_context(): - response = inspect.unwrap(console_resources.KnowledgeFSInitialSourcePreviewJobApi.get)( - console_resources.KnowledgeFSInitialSourcePreviewJobApi(), "job-1" - ) - assert response["status"] == "pending" - assert response["result"] is None - - with app.test_request_context(method="DELETE"): - response = inspect.unwrap(console_resources.KnowledgeFSInitialSourcePreviewJobApi.delete)( - console_resources.KnowledgeFSInitialSourcePreviewJobApi(), "job-1" - ) - assert response == {"job_id": "job-1", "result": None, "status": "completed"} - - -@pytest.mark.parametrize( - ("api_class", "method_name", "path_args"), - [ - ( - console_resources.KnowledgeFSSpaceLogicalDocumentDownloadApi, - console_resources.KnowledgeFSSpaceLogicalDocumentDownloadApi.get, - {"control_space_id": "control-1", "document_id": "document-1"}, - ), - ( - console_resources.KnowledgeFSSpaceLogicalDocumentsDownloadApi, - console_resources.KnowledgeFSSpaceLogicalDocumentsDownloadApi.post, - {"control_space_id": "control-1"}, - ), - ], -) -def test_document_download_routes_deny_callers_without_dataset_download_permission( - monkeypatch: pytest.MonkeyPatch, - api_class: type, - method_name: FunctionType, - path_args: dict[str, str], -) -> None: - permission_gate = Recorder(raises=Forbidden()) - apply_config_overrides(monkeypatch, RBAC_ENABLED=True) - monkeypatch.setattr( - common_wraps, - "current_account_with_tenant", - lambda: (_account(), "tenant-1"), - ) - monkeypatch.setattr("controllers.common.rbac.checks.RBACService.CheckAccess.check", permission_gate) - # The download routes carry the knowledge rate limit too; it needs a logged-in account. - monkeypatch.setattr(console_wraps, "check_knowledge_rate_limit", lambda: None) - monkeypatch.setattr( - console_resources, - "_console_services", - lambda: pytest.fail("permission denial must happen before KnowledgeFS access"), - ) - - permission_wrapper = _rbac_wrapper(method_name) - with pytest.raises(Forbidden): - permission_wrapper(api_class(), **path_args) - - assert permission_gate.calls == [ - invocation( - "tenant-1", - "account-1", - scene=RBACPermission.DATASET_DOCUMENT_DOWNLOAD, - resource_type=None, - resource_id=None, - ) - ] - - -def test_single_document_download_allows_dataset_download_permission_and_keeps_tenant_scope( - monkeypatch: pytest.MonkeyPatch, -) -> None: - descriptor = KnowledgeFSDocumentDownloadDescriptor( - document_id="document-1", - filename="guide.md", - mime_type="text/markdown", - object_key="namespaces/tenant-1/spaces/space-1/documents/guide.md", - sha256="a" * 64, - size_bytes=4, - ) - facade = FakeComponent(prepare_logical_document_download=Recorder(descriptor)) - download_service = FakeComponent(load_stream=Recorder(iter((b"body",)))) - permission_gate = Recorder(True) - apply_config_overrides(monkeypatch, RBAC_ENABLED=True) - monkeypatch.setattr( - common_wraps, - "current_account_with_tenant", - lambda: (_account(), "tenant-1"), - ) - monkeypatch.setattr("controllers.common.rbac.checks.RBACService.CheckAccess.check", permission_gate) - # The download routes carry the knowledge rate limit too; it needs a logged-in account. - monkeypatch.setattr(console_wraps, "check_knowledge_rate_limit", lambda: None) - monkeypatch.setattr(console_resources, "_actor", lambda: ("account-1", "tenant-1")) - monkeypatch.setattr(console_resources, "_console_services", lambda: FakeRuntime(facade=facade)) - monkeypatch.setattr(console_resources, "KnowledgeFSDownloadService", lambda: download_service) - permission_wrapper = _rbac_wrapper(console_resources.KnowledgeFSSpaceLogicalDocumentDownloadApi.get) - app = Flask(__name__) - - with app.test_request_context(): - response = permission_wrapper( - console_resources.KnowledgeFSSpaceLogicalDocumentDownloadApi(), - control_space_id="control-1", - document_id="document-1", - ) - - assert response.get_data() == b"body" - assert response.headers["Content-Disposition"] == "attachment; filename*=UTF-8''guide.md" - assert permission_gate.calls == [ - invocation( - "tenant-1", - "account-1", - scene=RBACPermission.DATASET_DOCUMENT_DOWNLOAD, - resource_type=None, - resource_id=None, - ) - ] - assert facade.prepare_logical_document_download.calls == [ - invocation( - tenant_id="tenant-1", - account_id="account-1", - control_space_id="control-1", - document_id="document-1", - ) - ] - - -def test_single_document_download_preserves_resource_not_found_after_permission_allows( - monkeypatch: pytest.MonkeyPatch, -) -> None: - facade = FakeComponent( - prepare_logical_document_download=Recorder(raises=KnowledgeFSProductResourceNotFoundError("document not found")) - ) - permission_gate = Recorder(True) - download_service_factory = Recorder() - apply_config_overrides(monkeypatch, RBAC_ENABLED=True) - monkeypatch.setattr( - common_wraps, - "current_account_with_tenant", - lambda: (_account(), "tenant-1"), - ) - monkeypatch.setattr("controllers.common.rbac.checks.RBACService.CheckAccess.check", permission_gate) - # The download routes carry the knowledge rate limit too; it needs a logged-in account. - monkeypatch.setattr(console_wraps, "check_knowledge_rate_limit", lambda: None) - monkeypatch.setattr(console_resources, "_actor", lambda: ("account-1", "tenant-1")) - monkeypatch.setattr(console_resources, "_console_services", lambda: FakeRuntime(facade=facade)) - monkeypatch.setattr(console_resources, "KnowledgeFSDownloadService", download_service_factory) - permission_wrapper = _rbac_wrapper(console_resources.KnowledgeFSSpaceLogicalDocumentDownloadApi.get) - app = Flask(__name__) - - with app.test_request_context(), pytest.raises(KnowledgeFSResourceNotFoundHTTPError): - permission_wrapper( - console_resources.KnowledgeFSSpaceLogicalDocumentDownloadApi(), - control_space_id="foreign-control", - document_id="document-1", - ) - - assert len(permission_gate.calls) == 1 - assert facade.prepare_logical_document_download.calls == [ - invocation( - tenant_id="tenant-1", - account_id="account-1", - control_space_id="foreign-control", - document_id="document-1", - ) - ] - assert download_service_factory.calls == [] - - -def test_single_document_download_maps_storage_unavailable_after_permission_allows( - monkeypatch: pytest.MonkeyPatch, -) -> None: - descriptor = KnowledgeFSDocumentDownloadDescriptor( - document_id="document-1", - filename="guide.md", - mime_type="text/markdown", - object_key="namespaces/tenant-1/spaces/space-1/documents/guide.md", - sha256="a" * 64, - size_bytes=4, - ) - facade = FakeComponent(prepare_logical_document_download=Recorder(descriptor)) - download_service = FakeComponent( - load_stream=Recorder(raises=KnowledgeFSDownloadUnavailableError("storage unavailable")) - ) - permission_gate = Recorder(True) - apply_config_overrides(monkeypatch, RBAC_ENABLED=True) - monkeypatch.setattr( - common_wraps, - "current_account_with_tenant", - lambda: (_account(), "tenant-1"), - ) - monkeypatch.setattr("controllers.common.rbac.checks.RBACService.CheckAccess.check", permission_gate) - # The download routes carry the knowledge rate limit too; it needs a logged-in account. - monkeypatch.setattr(console_wraps, "check_knowledge_rate_limit", lambda: None) - monkeypatch.setattr(console_resources, "_actor", lambda: ("account-1", "tenant-1")) - monkeypatch.setattr(console_resources, "_console_services", lambda: FakeRuntime(facade=facade)) - monkeypatch.setattr(console_resources, "KnowledgeFSDownloadService", lambda: download_service) - permission_wrapper = _rbac_wrapper(console_resources.KnowledgeFSSpaceLogicalDocumentDownloadApi.get) - app = Flask(__name__) - - with app.test_request_context(), pytest.raises(ServiceUnavailable): - permission_wrapper( - console_resources.KnowledgeFSSpaceLogicalDocumentDownloadApi(), - control_space_id="control-1", - document_id="document-1", - ) - - assert len(permission_gate.calls) == 1 - assert download_service.load_stream.calls == [invocation(descriptor)] - - -def test_batch_document_download_maps_storage_unavailable_after_permission_allows( - monkeypatch: pytest.MonkeyPatch, -) -> None: - descriptor = KnowledgeFSDocumentDownloadDescriptor( - document_id="document-1", - filename="guide.md", - mime_type="text/markdown", - object_key="namespaces/tenant-1/spaces/space-1/documents/guide.md", - sha256="a" * 64, - size_bytes=4, - ) - facade = FakeComponent(prepare_logical_document_download=Recorder(descriptor)) - - @contextmanager - def unavailable_zip() -> Generator[Never]: - raise KnowledgeFSDownloadUnavailableError("storage unavailable") - yield - - download_service = FakeComponent(build_zip_tempfile=Recorder(unavailable_zip())) - permission_gate = Recorder(True) - apply_config_overrides(monkeypatch, RBAC_ENABLED=True) - monkeypatch.setattr( - common_wraps, - "current_account_with_tenant", - lambda: (_account(), "tenant-1"), - ) - monkeypatch.setattr("controllers.common.rbac.checks.RBACService.CheckAccess.check", permission_gate) - # The download routes carry the knowledge rate limit too; it needs a logged-in account. - monkeypatch.setattr(console_wraps, "check_knowledge_rate_limit", lambda: None) - monkeypatch.setattr(console_resources, "_actor", lambda: ("account-1", "tenant-1")) - monkeypatch.setattr(console_resources, "_console_services", lambda: FakeRuntime(facade=facade)) - monkeypatch.setattr(console_resources, "KnowledgeFSDownloadService", lambda: download_service) - permission_wrapper = _rbac_wrapper(console_resources.KnowledgeFSSpaceLogicalDocumentsDownloadApi.post) - app = Flask(__name__) - - with ( - app.test_request_context(json={"document_ids": ["document-1"]}), - pytest.raises(ServiceUnavailable), - ): - permission_wrapper( - console_resources.KnowledgeFSSpaceLogicalDocumentsDownloadApi(), - control_space_id="control-1", - ) - - assert len(permission_gate.calls) == 1 - assert download_service.build_zip_tempfile.calls == [invocation([descriptor])] - - -def test_knowledge_fs_request_and_response_schemas_are_registered() -> None: - assert { - "KnowledgeFSSpaceCreatePayload", - "KnowledgeFSAppBindingPayload", - "KnowledgeFSAppBindingListResponse", - "KnowledgeFSSpaceUpdatePayload", - "KnowledgeFSMembersReplacePayload", - "KnowledgeFSExternalAccessPayload", - "KnowledgeFSSettingsPayload", - "KnowledgeFSSourceCreatePayload", - "KnowledgeFSQueryCreatePayload", - "KnowledgeFSQueryStreamCapabilityResponse", - "KnowledgeFSResearchTaskCreatePayload", - "KnowledgeFSStreamCapabilityPayload", - "KnowledgeFSSpaceListResponse", - "KnowledgeFSSpaceDetailResponse", - "KnowledgeFSDocumentStagedUploadPayload", - "KnowledgeFSDocumentStagedUploadAcceptedResponse", - "KnowledgeFSDocumentMultimodalAssetQuery", - "KnowledgeFSDocumentMultimodalManifestResponse", - "KnowledgeFSStreamCapabilityResponse", - "KnowledgeFSJWKSResponse", - "KnowledgeFSSmallFileUploadResponse", - "KnowledgeFSCrawlPreviewPageListQuery", - "KnowledgeFSCrawlPreviewPageListResponse", - "KnowledgeFSCrawlPreviewSelectionPayload", - "KnowledgeFSSourceConnectionCreatePayload", - "KnowledgeFSSourceConnectionListQuery", - "KnowledgeFSSourceConnectionListResponse", - "KnowledgeFSSourceConnectionRefreshPayload", - "KnowledgeFSSourceProviderListResponse", - "KnowledgeFSSourceSyncPolicyPayload", - "KnowledgeFSSourceSyncPolicyResponse", - "KnowledgeFSSourceWorkflowImportPayload", - "KnowledgeFSSourceWorkflowCancelPayload", - "KnowledgeFSSourceWorkflowResponse", - "KnowledgeFSUploadPartPresignPayload", - "KnowledgeFSUploadSessionAbortPayload", - "KnowledgeFSUploadSessionCompletePayload", - "KnowledgeFSUploadSessionCreatePayload", - "KnowledgeFSPresignedUploadResponse", - "KnowledgeFSUploadSessionCreateResponse", - "KnowledgeFSUploadSessionMutationResponse", - "KnowledgeFSGoldenQuestionBulkImportPayload", - "KnowledgeFSGoldenQuestionBulkImportResponse", - "KnowledgeFSGoldenQuestionEvidenceMatchPayload", - "KnowledgeFSGoldenQuestionEvidenceMatchResponse", - }.issubset(console_ns.models) - assert { - "KnowledgeFSQueryCreatePayload", - "KnowledgeFSDocumentListResponse", - "KnowledgeFSResearchTaskListResponse", - "KnowledgeFSSettingsPayload", - "KnowledgeFSSettingsResponse", - "KnowledgeFSSourceCreatePayload", - "KnowledgeFSSourceListResponse", - "KnowledgeFSSourceResponse", - "KnowledgeFSTraceListResponse", - }.issubset(service_api_ns.models) - assert console_ns.models["KnowledgeFSSpaceCreatePayload"].__schema__["additionalProperties"] is False - query_capability_schema = console_ns.models["KnowledgeFSQueryStreamCapabilityResponse"].__schema__ - assert set(query_capability_schema["required"]) == {"expires_at", "operation_id", "token", "url"} - assert query_capability_schema["properties"]["operation_id"]["const"] == "createQuery" - - -def _document_multimodal_manifest(*, object_key: str | None = None) -> KnowledgeFSDocumentMultimodalManifest: - asset_ref = ( - { - "contentType": "image/png", - "objectKey": object_key, - "sha256": "a" * 64, - "variants": { - "thumbnail": { - "contentType": "image/webp", - "objectKey": object_key.replace("full.png", "thumbnail.webp"), - "sha256": "b" * 64, - } - }, - } - if object_key is not None - else None - ) - return KnowledgeFSDocumentMultimodalManifest.model_validate( - { - "artifactHash": "c" * 64, - "createdAt": "2026-08-14T12:00:00.000Z", - "documentAssetId": "018f0d60-7a49-7cc2-9c1b-5b36f18f2c44", - "id": "018f0d60-7a49-7cc2-9c1b-5b36f18f2c45", - "items": [ - { - "assetRef": asset_ref, - "caption": "Architecture", - "endOffset": 42, - "id": "figure:1", - "modality": "image", - "ocrText": "Service A to Service B", - "pageNumber": 2, - "parseElementId": "parse-element-1", - "sectionPath": ["Architecture"], - "startOffset": 42, - "title": "System diagram", - } - ], - "knowledgeSpaceId": "018f0d60-7a49-7cc2-9c1b-5b36f18f2c43", - "manifestVersion": "document-multimodal@1", - "updatedAt": "2026-08-14T12:01:00.000Z", - "version": 3, - } - ) - - -def test_document_multimodal_manifest_console_bff_exposes_only_authorized_asset_urls( - monkeypatch: pytest.MonkeyPatch, -) -> None: - document_id = "018f0d60-7a49-7cc2-9c1b-5b36f18f2c44" - manifest = _document_multimodal_manifest( - object_key=( - f"tenant-1/spaces/018f0d60-7a49-7cc2-9c1b-5b36f18f2c43/documents/{document_id}/assets/figure-1-full.png" - ) - ) - facade = FakeComponent(get_document_multimodal_manifest=Recorder(manifest)) - monkeypatch.setattr(console_resources, "_actor", lambda: ("account-1", "tenant-1")) - monkeypatch.setattr(console_resources, "_console_services", lambda: FakeRuntime(facade=facade)) - apply_config_overrides(monkeypatch, CONSOLE_API_URL="") - app = Flask(__name__) - - with app.test_request_context(): - get = inspect.unwrap(console_resources.KnowledgeFSSpaceDocumentMultimodalManifestApi.get) - response = get( - console_resources.KnowledgeFSSpaceDocumentMultimodalManifestApi(), - control_space_id="control-1", - document_id=document_id, - ) - - item = response["items"][0] - assert item["asset_url"] == ( - f"/console/api/knowledge-fs/spaces/control-1/documents/{document_id}/multimodal/figure%3A1/asset" - ) - assert item["thumbnail_url"] == f"{item['asset_url']}?variant=thumbnail" - assert item["parse_element_id"] == "parse-element-1" - assert "objectKey" not in str(response) - assert "sha256" not in str(response) - assert facade.get_document_multimodal_manifest.calls == [ - invocation( - tenant_id="tenant-1", - account_id="account-1", - control_space_id="control-1", - document_id=document_id, - ) - ] - - -def test_document_multimodal_asset_console_bff_reauthorizes_and_streams_scoped_images( - monkeypatch: pytest.MonkeyPatch, -) -> None: - document_id = "018f0d60-7a49-7cc2-9c1b-5b36f18f2c44" - object_key = ( - f"tenant-1/spaces/018f0d60-7a49-7cc2-9c1b-5b36f18f2c43/documents/{document_id}/assets/figure-1-full.png" - ) - manifest = _document_multimodal_manifest(object_key=object_key) - facade = FakeComponent(get_document_multimodal_manifest=Recorder(manifest)) - object_storage = FakeComponent( - head_object=Recorder( - KnowledgeFSObjectMetadata( - checksum_sha256_base64="checksum", - content_type="image/png", - key=object_key, - metadata={}, - size_bytes=4, - ) - ) - ) - download_service = FakeComponent(load_stream=Recorder(iter((b"im", b"age")))) - download_service_factory = Recorder(download_service) - monkeypatch.setattr(console_resources, "_actor", lambda: ("account-1", "tenant-1")) - monkeypatch.setattr(console_resources, "_console_services", lambda: FakeRuntime(facade=facade)) - monkeypatch.setattr(console_resources, "KnowledgeFSObjectStorageService", lambda: object_storage) - monkeypatch.setattr(console_resources, "KnowledgeFSDownloadService", download_service_factory) - app = Flask(__name__) - - with app.test_request_context(): - get = inspect.unwrap(console_resources.KnowledgeFSSpaceDocumentMultimodalAssetApi.get) - response = get( - console_resources.KnowledgeFSSpaceDocumentMultimodalAssetApi(), - control_space_id="control-1", - document_id=document_id, - item_id="figure:1", - ) - - assert b"".join(response.response) == b"image" - assert response.content_type == "image/png" - assert response.headers["Content-Disposition"] == "inline" - assert response.headers["Content-Security-Policy"] == "default-src 'none'; sandbox" - assert response.headers["X-Content-Type-Options"] == "nosniff" - assert object_storage.head_object.calls == [invocation(key=object_key)] - assert download_service_factory.calls == [invocation(object_storage=object_storage)] - assert len(download_service.load_stream.calls) == 1 - - -def test_document_multimodal_asset_console_bff_rejects_cross_document_object_keys( - monkeypatch: pytest.MonkeyPatch, -) -> None: - document_id = "018f0d60-7a49-7cc2-9c1b-5b36f18f2c44" - manifest = _document_multimodal_manifest( - object_key=( - "tenant-1/spaces/018f0d60-7a49-7cc2-9c1b-5b36f18f2c43/" - "documents/018f0d60-7a49-7cc2-9c1b-5b36f18f2cff/assets/foreign.png" - ) - ) - facade = FakeComponent(get_document_multimodal_manifest=Recorder(manifest)) - object_storage_factory = Recorder() - monkeypatch.setattr(console_resources, "_actor", lambda: ("account-1", "tenant-1")) - monkeypatch.setattr(console_resources, "_console_services", lambda: FakeRuntime(facade=facade)) - monkeypatch.setattr(console_resources, "KnowledgeFSObjectStorageService", object_storage_factory) - app = Flask(__name__) - get = inspect.unwrap(console_resources.KnowledgeFSSpaceDocumentMultimodalAssetApi.get) - - with app.test_request_context(), pytest.raises(NotFound): - get( - console_resources.KnowledgeFSSpaceDocumentMultimodalAssetApi(), - control_space_id="control-1", - document_id=document_id, - item_id="figure:1", - ) - - assert object_storage_factory.calls == [] - - -def test_source_workflow_import_contract_is_discriminated_idempotent_and_accepted() -> None: - import_schema = console_ns.models["KnowledgeFSSourceWorkflowImportPayload"].__schema__ - assert import_schema["discriminator"] == { - "propertyName": "kind", - "mapping": { - "online-document-import": "#/components/schemas/KnowledgeFSOnlineDocumentWorkflowImportPayload", - "online-drive-import": "#/components/schemas/KnowledgeFSOnlineDriveWorkflowImportPayload", - }, - } - assert import_schema["oneOf"] == [ - {"$ref": "#/components/schemas/KnowledgeFSOnlineDocumentWorkflowImportPayload"}, - {"$ref": "#/components/schemas/KnowledgeFSOnlineDriveWorkflowImportPayload"}, - ] - - api_doc = console_resources.KnowledgeFSSpaceSourceWorkflowImportApi.post.__apidoc__ - assert api_doc["params"]["Idempotency-Key"]["required"] is True - assert api_doc["params"]["Idempotency-Key"]["minLength"] == 8 - assert api_doc["params"]["Idempotency-Key"]["maxLength"] == 255 - assert set(api_doc["responses"]) == {"202"} - - -def test_small_file_console_bff_reads_only_through_facade_and_returns_no_capability_material( - monkeypatch: pytest.MonkeyPatch, -) -> None: - calls: list[dict[str, object]] = [] - - class Facade: - def upload_small_file(self, **kwargs: object) -> KnowledgeFSSmallFileUploadResponse: - calls.append({name: value for name, value in kwargs.items() if name != "body_reader"}) - body_reader = kwargs["body_reader"] - assert callable(body_reader) - assert body_reader(8 * 1024 * 1024) == b"tiny" - return KnowledgeFSSmallFileUploadResponse.model_validate( - { - "session": { - "compilationJobId": "compilation-1", - "completedAt": 2_000_000, - "documentAssetId": "asset-1", - "expectedSizeBytes": 4, - "expiresAt": 2_060_000, - "id": "session-1", - "mode": "small_fallback", - "status": "completed", - } - } - ) - - monkeypatch.setattr(console_resources, "_actor", lambda: ("account-1", "tenant-1")) - monkeypatch.setattr(console_resources, "_console_services", lambda: _FacadeRuntime(facade=Facade())) - app = Flask(__name__) - - with app.test_request_context( - method="POST", - data={"file": (BytesIO(b"tiny"), "small.txt")}, - content_type="multipart/form-data", - ): - post = inspect.unwrap(console_resources.KnowledgeFSSpaceSmallFileUploadApi.post) - response = post(console_resources.KnowledgeFSSpaceSmallFileUploadApi(), "control-1", "session-1") - - assert response["session"]["status"] == "completed" - assert "token" not in str(response).lower() - assert "url" not in str(response).lower() - assert calls == [ - { - "tenant_id": "tenant-1", - "account_id": "account-1", - "control_space_id": "control-1", - "upload_session_id": "session-1", - } - ] - - -def test_document_upload_console_bff_reads_only_through_facade_and_returns_accepted_response( - monkeypatch: pytest.MonkeyPatch, -) -> None: - calls: list[dict[str, object]] = [] - - class Facade: - def create_document(self, **kwargs: object) -> KnowledgeFSDocumentUploadAcceptedResponse: - calls.append({name: value for name, value in kwargs.items() if name != "body_reader"}) - body_reader = kwargs["body_reader"] - assert callable(body_reader) - assert body_reader(15 * 1024 * 1024) == KnowledgeFSRemoteMultipartFile( - filename="notes.md", - content_type="text/markdown", - body=b"# Notes", - ) - return KnowledgeFSDocumentUploadAcceptedResponse.model_validate( - { - "asset": { - "createdAt": "2030-01-01T00:00:00Z", - "filename": "notes.md", - "id": "asset-1", - "knowledgeSpaceId": "space-1", - "metadata": {}, - "mimeType": "text/markdown", - "objectKey": "documents/asset-1/notes.md", - "parserStatus": "pending", - "sha256": "sha256", - "sizeBytes": 7, - "sourceId": None, - "updatedAt": None, - "version": 1, - }, - "assetStatusUrl": "/knowledge-spaces/space-1/documents/asset-1", - "compilationJob": {"id": "job-1", "stage": "queued"}, - "documentRevision": 1, - "logicalDocument": {"id": "document-1", "revision": 1}, - "logicalDocumentId": "document-1", - "statusUrl": "/knowledge-spaces/space-1/logical-documents/document-1/tasks/job-1", - } - ) - - monkeypatch.setattr(console_resources, "_actor", lambda: ("account-1", "tenant-1")) - monkeypatch.setattr(console_resources, "_console_services", lambda: _FacadeRuntime(facade=Facade())) - app = Flask(__name__) - - with app.test_request_context( - method="POST", - data={"file": (BytesIO(b"# Notes"), "notes.md", "text/markdown")}, - content_type="multipart/form-data", - ): - post = inspect.unwrap(console_resources.KnowledgeFSSpaceDocumentsApi.post) - response, status = post(console_resources.KnowledgeFSSpaceDocumentsApi(), "control-1") - - assert status == 202 - assert response["logical_document_id"] == "document-1" - assert calls == [ - { - "tenant_id": "tenant-1", - "account_id": "account-1", - "control_space_id": "control-1", - } - ] - - -def test_staged_document_claim_passes_only_upload_id_to_the_admission_service( - monkeypatch: pytest.MonkeyPatch, -) -> None: - staged_uploads = FakeComponent( - claim=Recorder( - KnowledgeFSDocumentStagedUploadAcceptedResponse( - upload_id="staged-1", - document_asset_id="asset-1", - compilation_job_id="job-1", - ) - ) - ) - monkeypatch.setattr(console_resources, "_actor", lambda: ("account-1", "tenant-1")) - monkeypatch.setattr(console_resources, "_staged_uploads", lambda: staged_uploads) - app = Flask(__name__) - - with app.test_request_context(method="POST", json={"upload_id": "staged-1"}): - post = inspect.unwrap(console_resources.KnowledgeFSSpaceDocumentsApi.post) - response, status = post(console_resources.KnowledgeFSSpaceDocumentsApi(), "control-1") - - assert status == 202 - assert response == { - "status": "accepted", - "upload_id": "staged-1", - "document_asset_id": "asset-1", - "compilation_job_id": "job-1", - } - assert len(staged_uploads.claim.calls) == 1 - claimed = narrowed(staged_uploads.claim.last_call.kwargs["payload"], KnowledgeFSDocumentStagedUploadPayload) - assert claimed.upload_id == "staged-1" - - -@pytest.mark.parametrize("file_size_limit_mb", [15, 50]) -def test_workspace_staging_upload_persists_before_space_creation( - monkeypatch: pytest.MonkeyPatch, - file_size_limit_mb: int, -) -> None: - admission_events: list[str] = [] - - class RecordingAdmission: - active = False - - @contextmanager - def admit(self, *, reserved_bytes: int) -> Generator[None]: - admission_events.append(f"admit:{reserved_bytes}") - self.active = True - try: - yield - finally: - self.active = False - admission_events.append("release") - - admission = RecordingAdmission() - staged = KnowledgeFSStagedUploadResponse( - id="staged-1", - file_name="notes.md", - content_type="text/markdown", - size_bytes=7, - status="uploaded", - expires_at=datetime(2030, 1, 1, tzinfo=UTC), - ) - - def stage(**_: object) -> KnowledgeFSStagedUploadResponse: - assert admission.active - admission_events.append("stage") - return staged - - staged_uploads = FakeComponent(stage=Recorder(effect=stage)) - account = _account() - monkeypatch.setattr(console_resources, "current_account_with_tenant", lambda: (account, "tenant-1")) - monkeypatch.setattr(console_resources, "_staged_uploads", lambda: staged_uploads) - monkeypatch.setattr(console_resources, "DEFAULT_KNOWLEDGE_FS_BUFFERED_UPLOAD_ADMISSION", admission) - monkeypatch.setattr( - console_resources.FeatureService, - "get_knowledge_file_size_limit", - lambda _tenant_id: file_size_limit_mb, - ) - app = Flask(__name__) - - with app.test_request_context( - method="POST", - data={"file": (BytesIO(b"# Notes"), "notes.md", "text/markdown")}, - content_type="multipart/form-data", - ): - post = inspect.unwrap(console_resources.KnowledgeFSStagedUploadsApi.post) - response, status = post(console_resources.KnowledgeFSStagedUploadsApi()) - - assert status == 201 - assert response["id"] == "staged-1" - assert admission_events == [f"admit:{file_size_limit_mb * 1024 * 1024}", "stage", "release"] - assert staged_uploads.stage.calls == [ - invocation( - tenant_id="tenant-1", - account=account, - file_name="notes.md", - content_type="text/markdown", - body=b"# Notes", - file_size_limit_mb=file_size_limit_mb, - ) - ] - - -def test_workspace_staging_zero_byte_plan_limit_returns_request_too_large_before_admission( - monkeypatch: pytest.MonkeyPatch, -) -> None: - class FailingAdmission: - def admit(self, *, reserved_bytes: int) -> Never: - _ = reserved_bytes - raise AssertionError("a disabled upload limit must fail before admission") - - staged_uploads = FakeComponent(stage=Recorder()) - account = _account() - monkeypatch.setattr(console_resources, "current_account_with_tenant", lambda: (account, "tenant-1")) - monkeypatch.setattr(console_resources, "_staged_uploads", lambda: staged_uploads) - monkeypatch.setattr(console_resources, "DEFAULT_KNOWLEDGE_FS_BUFFERED_UPLOAD_ADMISSION", FailingAdmission()) - monkeypatch.setattr( - console_resources.FeatureService, - "get_knowledge_file_size_limit", - lambda _tenant_id: 0, - ) - app = Flask(__name__) - - with app.test_request_context( - method="POST", - data={"file": (BytesIO(b"# Notes"), "notes.md", "text/markdown")}, - content_type="multipart/form-data", - ): - post = inspect.unwrap(console_resources.KnowledgeFSStagedUploadsApi.post) - with pytest.raises(KnowledgeFSProductRequestRejectedError) as rejected: - post(console_resources.KnowledgeFSStagedUploadsApi()) - - assert rejected.value.status_code == 413 - assert staged_uploads.stage.calls == [] - - -def test_logical_document_delete_accepts_initial_row_version( - monkeypatch: pytest.MonkeyPatch, -) -> None: - calls: list[dict[str, object]] = [] - - class Facade: - def delete_logical_document(self, **kwargs: object) -> KnowledgeFSDurableDeletionAcceptedResponse: - calls.append(kwargs) - payload = kwargs["payload"] - assert isinstance(payload, KnowledgeFSLogicalDocumentDeletePayload) - assert payload.expected_revision == 0 - return KnowledgeFSDurableDeletionAcceptedResponse.model_validate( - { - "job": { - "checkpoint": "requested", - "createdAt": "2030-01-01T00:00:00Z", - "id": "00000000-0000-4000-8000-000000000001", - "knowledgeSpaceId": "space-1", - "runState": "queued", - "targetId": "00000000-0000-4000-8000-000000000002", - "targetType": "logical_document", - "updatedAt": "2030-01-01T00:00:00Z", - }, - "statusUrl": "/deletion-jobs/job-1", - } - ) - - monkeypatch.setattr(console_resources, "_actor", lambda: ("account-1", "tenant-1")) - monkeypatch.setattr(console_resources, "_console_services", lambda: _FacadeRuntime(facade=Facade())) - app = Flask(__name__) - - with app.test_request_context( - method="DELETE", - json={"expectedRevision": 0}, - headers={"Idempotency-Key": "delete-logical-document-once"}, - ): - delete = inspect.unwrap(console_resources.KnowledgeFSSpaceLogicalDocumentApi.delete) - response, status = delete( - console_resources.KnowledgeFSSpaceLogicalDocumentApi(), - "control-1", - "document-1", - ) - - assert status == 202 - assert response["job"]["target_type"] == "logical_document" - assert len(calls) == 1 - assert {name: value for name, value in calls[0].items() if name != "payload"} == { - "tenant_id": "tenant-1", - "account_id": "account-1", - "control_space_id": "control-1", - "document_id": "document-1", - "idempotency_key": "delete-logical-document-once", - } - - -def test_small_file_console_bff_maps_oversize_to_413() -> None: - app = Flask(__name__) - with app.test_request_context( - method="POST", - data={"file": (BytesIO(b"large"), "large.bin")}, - content_type="multipart/form-data", - ): - with pytest.raises(KnowledgeFSProductRequestRejectedError) as oversized: - console_resources._read_small_file_body(4) - assert oversized.value.status_code == 413 - - @console_resources._knowledge_fs_errors - def reject() -> Never: - raise KnowledgeFSProductRequestRejectedError(status_code=413) - - with pytest.raises(KnowledgeFSRequestTooLargeHTTPError): - reject() - - -def test_console_knowledge_rate_limit_covers_every_write_download_and_retrieval_entrypoint() -> None: - """Mirror the legacy dataset console: every mutating, download, or retrieval route is rate limited. - - Pure reads stay unlimited. The exemptions are workspace-level source previews (like legacy - data source browsing), a staged-upload discard, the deprecated buffered query route that - fails closed anyway, and the internal capability-token stream transport. - """ - space = "/knowledge-fs/spaces/" - exempt_writes = { - ("post", "/knowledge-fs/source-provider-preview"), - ("post", "/knowledge-fs/source-provider-preview/jobs"), - ("delete", "/knowledge-fs/source-provider-preview/jobs/"), - ("delete", "/knowledge-fs/uploads/"), - ("post", f"{space}/queries"), - ("post", "_QUERY_STREAM_PROXY_PATH"), - } - rate_limited_reads = {("get", f"{space}/logical-documents//download")} - tree = ast.parse( - Path(console_resources.__file__).read_text(encoding="utf-8"), - filename=console_resources.__file__, - ) - seen: set[tuple[str, str]] = set() - - for class_node in (node for node in tree.body if isinstance(node, ast.ClassDef)): - routes = [ - ast.unparse(decorator.args[0]).strip("'\"") - for decorator in class_node.decorator_list - if isinstance(decorator, ast.Call) - and isinstance(decorator.func, ast.Attribute) - and decorator.func.attr == "route" - ] - if not routes: - continue - for method_node in ( - node - for node in class_node.body - if isinstance(node, ast.FunctionDef) and node.name in {"delete", "get", "patch", "post", "put"} - ): - key = (method_node.name, routes[0]) - seen.add(key) - rate_limited = any( - isinstance(decorator, ast.Call) - and isinstance(decorator.func, ast.Name) - and decorator.func.id == "cloud_edition_billing_rate_limit_check" - and len(decorator.args) == 1 - and isinstance(decorator.args[0], ast.Constant) - and decorator.args[0].value == "knowledge" - for decorator in method_node.decorator_list - ) - expected = key in rate_limited_reads if method_node.name == "get" else key not in exempt_writes - assert rate_limited == expected, f"{key} rate limited={rate_limited}, expected {expected}" - - assert exempt_writes <= seen - assert rate_limited_reads <= seen - - -def test_space_update_and_delete_publish_their_actual_http_status_contracts() -> None: - patch_responses = console_resources.KnowledgeFSSpaceApi.patch.__apidoc__["responses"] - delete_responses = console_resources.KnowledgeFSSpaceApi.delete.__apidoc__["responses"] - - assert set(patch_responses) == {"200"} - assert set(delete_responses) == {"204"} - - -def test_space_create_profile_intent_matches_the_exact_kfs_pending_configuration_shape() -> None: - payload = KnowledgeFSSpaceCreatePayload.model_validate( - { - "name": "Technical space", - "slug": "technical-space", - "embedding": { - "pluginId": "langgenius/openai", - "provider": "openai", - "model": "text-embedding-3-small", - }, - "retrieval": { - "defaultMode": "deep", - "reasoningModel": { - "pluginId": "langgenius/openai", - "provider": "openai", - "model": "gpt-4.1-mini", - }, - "rerank": { - "enabled": True, - "model": { - "pluginId": "langgenius/cohere", - "provider": "cohere", - "model": "rerank-v3.5", - }, - }, - "scoreThreshold": {"enabled": True, "stage": "mode-final", "value": 0.6}, - "topK": 20, - }, - } - ) - - assert payload.embedding is not None - assert payload.embedding is not None - assert payload.embedding.model_dump(mode="json", by_alias=True) == { - "pluginId": "langgenius/openai", - "provider": "openai", - "model": "text-embedding-3-small", - } - assert payload.retrieval is not None - assert payload.retrieval is not None - assert payload.retrieval.model_dump(mode="json", by_alias=True) == { - "defaultMode": "deep", - "reasoningModel": { - "pluginId": "langgenius/openai", - "provider": "openai", - "model": "gpt-4.1-mini", - }, - "rerank": { - "enabled": True, - "model": { - "pluginId": "langgenius/cohere", - "provider": "cohere", - "model": "rerank-v3.5", - }, - }, - "scoreThreshold": {"enabled": True, "stage": "mode-final", "value": 0.6}, - "topK": 20, - } - - with pytest.raises(ValueError, match="retrieval.defaultMode"): - KnowledgeFSSpaceCreatePayload.model_validate( - { - "name": "Invalid preset", - "slug": "invalid-preset", - "embedding": {"pluginId": "plugin", "provider": "provider", "model": "embedding"}, - "retrieval": { - "defaultMode": "auto", - "reasoningModel": {"pluginId": "plugin", "provider": "provider", "model": "reasoning"}, - "rerank": { - "enabled": True, - "model": {"pluginId": "plugin", "provider": "provider", "model": "rerank"}, - }, - "scoreThreshold": {"enabled": False, "stage": "mode-final"}, - "topK": 10, - }, - } - ) - - -def test_space_create_allows_deferred_model_setup_but_rejects_incomplete_fast_profile() -> None: - payload = KnowledgeFSSpaceCreatePayload.model_validate( - { - "name": "Setup later", - "slug": "setup-later", - } - ) - - assert payload.embedding is None - assert payload.retrieval is None - - with pytest.raises(ValueError, match="requires an embedding model"): - KnowledgeFSSpaceCreatePayload.model_validate( - { - "name": "Incomplete fast profile", - "slug": "incomplete-fast-profile", - "retrieval": { - "defaultMode": "fast", - "reasoningModel": { - "pluginId": "plugin", - "provider": "provider", - "model": "reasoning", - }, - "rerank": { - "enabled": True, - "model": {"pluginId": "plugin", "provider": "provider", "model": "rerank"}, - }, - "scoreThreshold": {"enabled": False, "stage": "mode-final"}, - "topK": 10, - }, - } - ) - - -def test_jwks_http_resource_returns_only_public_keys(monkeypatch: pytest.MonkeyPatch) -> None: - public_jwks = { - "keys": [ - {"alg": "RS256", "e": "AQAB", "kid": "current", "kty": "RSA", "n": "modulus-1", "use": "sig"}, - {"alg": "RS256", "e": "AQAB", "kid": "previous", "kty": "RSA", "n": "modulus-2", "use": "sig"}, - ] - } - issuer = FakeComponent(public_jwks=Recorder(public_jwks)) - apply_config_overrides( - monkeypatch, - KNOWLEDGE_FS_CAPABILITY_V2_ENABLED=True, - KNOWLEDGE_FS_JWKS_CACHE_MAX_AGE_SECONDS=123, - ) - monkeypatch.setattr(console_resources.session_factory, "get_session_maker", lambda: object()) - monkeypatch.setattr(console_resources, "create_configured_knowledge_fs_capability_issuer", lambda **_: issuer) - app = Flask(__name__) - - with app.app_context(): - response = console_resources.KnowledgeFSJWKSApi().get() - - payload = response.get_json() - assert payload == public_jwks - assert response.headers["Cache-Control"] == "public, max-age=123, must-revalidate" - assert not {"d", "p", "q", "dp", "dq", "qi"} & set().union(*(key.keys() for key in payload["keys"])) - - -def test_service_profile_authorizes_the_dataset_key_for_the_route_space(monkeypatch: pytest.MonkeyPatch) -> None: - profile = KnowledgeFSServiceApiProfile( - tenant_id="tenant-1", - control_space_id="control-2", - api_token_id="token-1", - principal_id="token-1", - knowledge_space_id="space-1", - knowledge_space_revision=1, - membership_epoch=0, - space_acl_epoch=0, - external_access_epoch=0, - content_policy_revision=0, - ) - authorization = FakeComponent(authorize=Recorder(profile)) - runtime = FakeRuntime(service_api_authorization=authorization) - monkeypatch.setattr(service_resources, "validate_and_get_api_token", Recorder(_dataset_api_token())) - app = Flask(__name__) - with app.test_request_context(headers={"Authorization": "Bearer dataset-test-key"}): - result = service_resources._profile( - runtime, # type: ignore[arg-type] - operation_id="listDocuments", - control_space_id="control-2", - ) - - assert result is profile - assert authorization.authorize.calls == [ - invocation( - api_token_id="token-1", - tenant_id="tenant-1", - control_space_id="control-2", - ) - ] - - -def test_service_profile_applies_the_knowledge_rate_limit_before_authorizing( - monkeypatch: pytest.MonkeyPatch, -) -> None: - """Every KnowledgeFS service route shares the legacy dataset service API rate limit.""" - authorization = FakeComponent(authorize=Recorder()) - runtime = FakeRuntime(service_api_authorization=authorization) - api_token = _dataset_api_token() - monkeypatch.setattr(service_resources, "validate_and_get_api_token", Recorder(api_token)) - rate_limit = Recorder(raises=Forbidden("rate limited")) - monkeypatch.setattr(service_resources, "check_knowledge_rate_limit", rate_limit) - - with pytest.raises(Forbidden, match="rate limited"): - service_resources._profile( - runtime, # type: ignore[arg-type] - operation_id="listDocuments", - control_space_id="control-2", - ) - - assert rate_limit.calls == [invocation(api_token)] - assert authorization.authorize.calls == [] - - -def test_service_profile_rejects_a_dataset_key_without_a_workspace(monkeypatch: pytest.MonkeyPatch) -> None: - authorization = FakeComponent(authorize=Recorder()) - runtime = FakeRuntime(service_api_authorization=authorization) - monkeypatch.setattr(service_resources, "validate_and_get_api_token", Recorder(_dataset_api_token(tenant_id=None))) - - with pytest.raises(KnowledgeFSServiceApiAuthorizationError, match="not workspace-scoped"): - service_resources._profile( - runtime, # type: ignore[arg-type] - operation_id="listDocuments", - control_space_id="control-2", - ) - - assert authorization.authorize.calls == [] - - -def test_product_modules_do_not_import_dify_dataset_or_document_services() -> None: - paths = [ - *(path for path in _API_ROOT.glob("services/knowledge_fs/*.py") if not path.name.startswith("upgrade_")), - *_API_ROOT.glob("controllers/console/knowledge_fs/*.py"), - *_API_ROOT.glob("controllers/service_api/knowledge_fs/*.py"), - ] - forbidden_modules = { - "models.dataset", - "services.dataset_service", - "services.document_service", - "controllers.console.datasets", - "controllers.service_api.dataset", - } - - for path in paths: - tree = ast.parse(path.read_text(encoding="utf-8"), filename=str(path)) - imported_modules = set() - for node in ast.walk(tree): - if not isinstance(node, ast.ImportFrom) or node.module is None: - continue - # Native vector credentials share the legacy models.dataset module. - # This bridge may read the tenant binding, never Dataset/Document data. - if ( - path == _API_ROOT / "services/knowledge_fs/vector_store.py" - and node.module == "models.dataset" - and {alias.name for alias in node.names} == {"TidbAuthBinding"} - ): - continue - imported_modules.add(node.module) - assert not any( - module == forbidden or module.startswith(f"{forbidden}.") - for module in imported_modules - for forbidden in forbidden_modules - ), path - - -def test_research_task_stream_url_binds_task_and_parent_space_without_token( - monkeypatch: pytest.MonkeyPatch, -) -> None: - apply_config_overrides(monkeypatch, CONSOLE_API_URL="https://dify.test") - url = console_resources._research_task_events_url( - task_id="task/one", - knowledge_space_id="space one", - ) - - assert ( - url == "https://dify.test/console/api/knowledge-fs/research-tasks/task%2Fone/events?knowledgeSpaceId=space+one" - ) - assert "token" not in url.lower() - - -def test_query_stream_capability_issues_exact_space_grant_without_token_in_url( - monkeypatch: pytest.MonkeyPatch, -) -> None: - calls: list[dict[str, object]] = [] - - class Broker: - def issue_interactive(self, **kwargs: object) -> KnowledgeFSIssuedProductCapability: - calls.append(kwargs) - return _issued_capability(token="query-capability", operation_id="createQuery") - - apply_config_overrides(monkeypatch, CONSOLE_API_URL="https://dify.test") - monkeypatch.setattr(console_resources, "_actor", lambda: ("account-1", "tenant-1")) - monkeypatch.setattr( - console_resources, - "_console_services", - lambda: _BrokerRuntime(broker=Broker()), - ) - app = Flask(__name__) - - with app.test_request_context(): - post = inspect.unwrap(console_resources.KnowledgeFSSpaceQueryStreamCapabilityApi.post) - response = post(console_resources.KnowledgeFSSpaceQueryStreamCapabilityApi(), "control-1") - - assert response == { - "expires_at": "2030-01-01T00:00:00Z", - "operation_id": "createQuery", - "token": "query-capability", - "url": "https://dify.test/console/api/knowledge-fs/query-stream", - } - assert calls == [ - { - "account_id": "account-1", - "control_space_id": "control-1", - "operation_id": "createQuery", - "tenant_id": "tenant-1", - } - ] - assert "token" not in response["url"].lower() - assert console_resources.KnowledgeFSSpaceQueryStreamCapabilityApi.post.__apidoc__["deprecated"] is True - - -def test_query_admission_binds_validated_mode_to_resolved_kfs_space(monkeypatch: pytest.MonkeyPatch) -> None: - calls: list[dict[str, object]] = [] - - class Broker: - def issue_interactive(self, **kwargs: object) -> KnowledgeFSIssuedProductCapability: - calls.append(kwargs) - return _issued_capability(token="query-capability", operation_id="createQuery") - - apply_config_overrides(monkeypatch, CONSOLE_API_URL="https://dify.test") - monkeypatch.setattr(console_resources, "_actor", lambda: ("account-1", "tenant-1")) - monkeypatch.setattr( - console_resources, - "_console_services", - lambda: _BrokerRuntime(broker=Broker()), - ) - app = Flask(__name__) - - with app.test_request_context(json={"query": "What changed?", "mode": "auto"}): - post = inspect.unwrap(console_resources.KnowledgeFSSpaceQueryAdmissionApi.post) - response = post(console_resources.KnowledgeFSSpaceQueryAdmissionApi(), "control-1") - - assert response["operation_id"] == "createQuery" - assert response["trace_id"] == "trace-1" - assert response["request"] == { - "activeDocumentIds": [], - "activeEntityIds": [], - "knowledgeSpaceId": "space-1", - "mode": "auto", - "query": "What changed?", - } - assert response["url"] == "https://dify.test/console/api/knowledge-fs/query-stream" - assert calls == [ - { - "account_id": "account-1", - "control_space_id": "control-1", - "operation_id": "createQuery", - "tenant_id": "tenant-1", - } - ] - - -def test_console_query_stream_proxy_forwards_only_the_admitted_capability( - monkeypatch: pytest.MonkeyPatch, -) -> None: - closed: list[str] = [] - - def close() -> None: - closed.append("closed") - - facade = FakeComponent( - stream_query=Recorder( - KnowledgeFSRemoteSSEResponse( - status_code=200, - headers=( - ("content-type", "text/event-stream"), - ("x-query-run-id", "query-1"), - ), - chunks=iter((b"event: answer\n", b'data: {"answer":"ok"}\n\n')), - close=close, - ) - ) - ) - monkeypatch.setattr(console_resources, "_console_services", lambda: FakeRuntime(facade=facade)) - app = Flask(__name__) - - with app.test_request_context( - json={ - "activeDocumentIds": [], - "activeEntityIds": [], - "knowledgeSpaceId": "space-1", - "mode": "fast", - "query": "hello", - }, - headers={ - "Authorization": "Bearer capability-token", - "Cookie": "session=must-not-forward", - "X-Trace-ID": "trace-1", - }, - ): - post = inspect.unwrap(console_resources.KnowledgeFSQueryStreamProxyApi.post) - response = post(console_resources.KnowledgeFSQueryStreamProxyApi()) - - assert response.status_code == 200 - assert response.headers["Content-Type"] == "text/event-stream" - assert response.headers["X-Accel-Buffering"] == "no" - assert b"".join(response.response) == b'event: answer\ndata: {"answer":"ok"}\n\n' - assert closed == ["closed"] - call = facade.stream_query.last_call.kwargs - assert call["capability_token"] == "capability-token" - assert call["trace_id"] == "trace-1" - assert narrowed(call["payload"], KnowledgeFSAdmittedQueryRequest).knowledge_space_id == "space-1" - assert set(call) == {"capability_token", "trace_id", "payload"} - - -def test_service_query_stream_proxy_forwards_only_the_admitted_capability( - monkeypatch: pytest.MonkeyPatch, -) -> None: - closed: list[str] = [] - - def close() -> None: - closed.append("closed") - - facade = FakeComponent( - stream_query=Recorder( - KnowledgeFSRemoteSSEResponse( - status_code=200, - headers=( - ("content-type", "text/event-stream"), - ("x-query-run-id", "query-1"), - ), - chunks=iter((b"event: answer\n", b'data: {"answer":"ok"}\n\n')), - close=close, - ) - ) - ) - monkeypatch.setattr(service_resources, "_runtime", lambda: FakeRuntime(facade=facade)) - app = Flask(__name__) - - with app.test_request_context( - json={ - "activeDocumentIds": [], - "activeEntityIds": [], - "knowledgeSpaceId": "space-1", - "mode": "fast", - "query": "hello", - }, - headers={ - "Authorization": "Bearer capability-token", - "Cookie": "session=must-not-forward", - "X-Trace-ID": "trace-1", - }, - ): - post = inspect.unwrap(service_resources.KnowledgeFSServiceQueryStreamProxyApi.post) - response = post(service_resources.KnowledgeFSServiceQueryStreamProxyApi()) - - assert response.status_code == 200 - assert response.headers["Content-Type"] == "text/event-stream" - assert response.headers["X-Accel-Buffering"] == "no" - assert b"".join(response.response) == b'event: answer\ndata: {"answer":"ok"}\n\n' - assert closed == ["closed"] - call = facade.stream_query.last_call.kwargs - assert call["capability_token"] == "capability-token" - assert call["trace_id"] == "trace-1" - assert narrowed(call["payload"], KnowledgeFSAdmittedQueryRequest).knowledge_space_id == "space-1" - assert set(call) == {"capability_token", "trace_id", "payload"} - - -@pytest.mark.parametrize( - ("headers", "error_type"), - [ - ({"X-Trace-ID": "trace-1"}, service_resources.KnowledgeFSInvalidCredentialHTTPError), - ( - {"Authorization": "Bearer capability-token"}, - service_resources.KnowledgeFSServiceInvalidRequestHTTPError, - ), - ], -) -def test_service_query_stream_proxy_rejects_invalid_internal_transport_headers( - headers: dict[str, str], - error_type: type[Exception], -) -> None: - app = Flask(__name__) - - with app.test_request_context(headers=headers), pytest.raises(error_type): - service_resources._stream_capability() - - -def test_task_stream_capability_uses_broker( - monkeypatch: pytest.MonkeyPatch, -) -> None: - calls: list[dict[str, object]] = [] - - class Broker: - def issue_interactive(self, **kwargs: object) -> KnowledgeFSIssuedProductCapability: - calls.append(kwargs) - return _issued_capability(token="direct-capability", operation_id="streamResearchTask") - - runtime = _BrokerRuntime(broker=Broker()) - apply_config_overrides(monkeypatch, CONSOLE_API_URL="https://dify.test") - monkeypatch.setattr(console_resources, "_actor", lambda: ("account-1", "tenant-1")) - monkeypatch.setattr(console_resources, "_console_services", lambda: runtime) - app = Flask(__name__) - - with app.test_request_context(json={"control_space_id": "control-1"}): - stream_post = inspect.unwrap(console_resources.KnowledgeFSTaskStreamCapabilityApi.post) - stream_response = stream_post(console_resources.KnowledgeFSTaskStreamCapabilityApi(), "task-1") - - assert stream_response["operation_id"] == "streamResearchTask" - assert calls == [ - { - "tenant_id": "tenant-1", - "account_id": "account-1", - "control_space_id": "control-1", - "operation_id": "streamResearchTask", - "resource_id": "task-1", - }, - ] - - -def test_service_query_admission_uses_broker(monkeypatch: pytest.MonkeyPatch) -> None: - calls: list[dict[str, object]] = [] - profile = _service_profile("control-1") - authorization = FakeComponent(authorize=Recorder(profile)) - - class Broker: - def issue_service(self, **kwargs: object) -> KnowledgeFSIssuedProductCapability: - calls.append(kwargs) - return _issued_capability(token="service-capability", operation_id="createQuery") - - @dataclass(frozen=True) - class ServiceRuntime: - service_api_authorization: FakeComponent - broker: Broker - - runtime = ServiceRuntime(service_api_authorization=authorization, broker=Broker()) - monkeypatch.setattr(service_resources, "validate_and_get_api_token", Recorder(_dataset_api_token())) - apply_config_overrides(monkeypatch, SERVICE_API_URL="https://api.dify.test") - monkeypatch.setattr(service_resources, "_runtime", lambda: runtime) - app = Flask(__name__) - - with app.test_request_context( - json={"query": "What changed?", "mode": "fast"}, - headers={"Authorization": "Bearer dataset-test-key"}, - ): - post = inspect.unwrap(service_resources.KnowledgeFSServiceQueryAdmissionApi.post) - response = post(service_resources.KnowledgeFSServiceQueryAdmissionApi(), "control-1") - - assert response["operation_id"] == "createQuery" - assert response["trace_id"] == "trace-1" - assert response["request"]["knowledgeSpaceId"] == "space-1" - assert response["url"] == "https://api.dify.test/v1/knowledge-fs/query-stream" - assert authorization.authorize.calls == [ - invocation( - api_token_id="token-1", - tenant_id="tenant-1", - control_space_id="control-1", - ) - ] - assert calls == [{"profile": profile, "operation_id": "createQuery"}] - - -@pytest.mark.parametrize( - ("raised", "expected"), - [ - # An unknown or wrong-kind key is an invalid credential (401). - ( - KnowledgeFSServiceApiAuthorizationError("bad key"), - service_resources.KnowledgeFSInvalidCredentialHTTPError, - ), - # A real key bound to other knowledge bases is out of scope for this space (403). - ( - KnowledgeFSServiceApiScopeError("out of scope"), - service_resources.KnowledgeFSServiceAccessDeniedHTTPError, - ), - ], -) -def test_service_api_errors_distinguish_unknown_key_from_out_of_scope_key( - raised: Exception, - expected: type[Exception], -) -> None: - @service_resources._service_api_errors - def view() -> None: - raise raised - - with pytest.raises(expected): - view() - - -def test_space_create_member_schema_requires_viewer_role() -> None: - payload_schema = console_ns.models["KnowledgeFSSpaceCreatePayload"].__schema__ - member_schema = console_ns.models["KnowledgeFSSpaceCreateMemberPayload"].__schema__ - - assert payload_schema["properties"]["members"]["maxItems"] == 1_000 - assert member_schema["properties"]["role"]["const"] == "viewer" - assert set(member_schema["required"]) == {"account_id", "role"} - assert member_schema["additionalProperties"] is False - - -def test_space_create_initial_access_idempotency_conflict_returns_409() -> None: - @console_resources._knowledge_fs_errors - def conflicting_creation() -> None: - raise KnowledgeFSControlSpaceIntentConflictError("different initial access") - - with pytest.raises(KnowledgeFSConflictHTTPError) as exc: - conflicting_creation() - - assert exc.value.code == 409 diff --git a/api/tests/unit_tests/controllers/test_knowledge_fs_resource_delegation.py b/api/tests/unit_tests/controllers/test_knowledge_fs_resource_delegation.py deleted file mode 100644 index 7c26d73393b..00000000000 --- a/api/tests/unit_tests/controllers/test_knowledge_fs_resource_delegation.py +++ /dev/null @@ -1,2201 +0,0 @@ -from __future__ import annotations - -from collections.abc import Callable -from dataclasses import dataclass, field -from datetime import UTC, datetime -from hashlib import sha256 -from http import HTTPStatus -from inspect import unwrap -from typing import Literal, cast - -import pytest -from flask import Flask -from flask_restx import Resource -from pydantic import BaseModel -from werkzeug.exceptions import NotFound, ServiceUnavailable - -from controllers.console.knowledge_fs import resources as console_resources -from controllers.console.knowledge_fs.error import ( - KnowledgeFSAccessDeniedHTTPError, - KnowledgeFSConflictHTTPError, - KnowledgeFSInvalidRequestHTTPError, - KnowledgeFSOperationUnavailableHTTPError, - KnowledgeFSRateLimitHTTPError, - KnowledgeFSRequestRejectedHTTPError, - KnowledgeFSRequestTooLargeHTTPError, - KnowledgeFSResourceNotFoundHTTPError, - KnowledgeFSSpaceNotFoundHTTPError, - KnowledgeFSTimeoutHTTPError, - KnowledgeFSUpstreamUnavailableHTTPError, -) -from controllers.service_api.knowledge_fs import resources as service_resources -from controllers.service_api.knowledge_fs.error import ( - KnowledgeFSInvalidCredentialHTTPError, - KnowledgeFSServiceAccessDeniedHTTPError, - KnowledgeFSServiceConflictHTTPError, - KnowledgeFSServiceInvalidRequestHTTPError, - KnowledgeFSServiceOperationUnavailableHTTPError, - KnowledgeFSServiceRateLimitHTTPError, - KnowledgeFSServiceRequestRejectedHTTPError, - KnowledgeFSServiceRequestTooLargeHTTPError, - KnowledgeFSServiceResourceNotFoundHTTPError, - KnowledgeFSServiceUpstreamUnavailableHTTPError, -) -from enums import DeploymentEdition -from libs.exception import BaseHTTPException -from machinery.context import RequestContext -from models import Account -from models.knowledge_fs import KnowledgeFSAppSpaceJoinType -from services.knowledge_fs import source_edit_service -from services.knowledge_fs.capability_broker import KnowledgeFSIssuedProductCapability -from services.knowledge_fs.product_dto import ( - KnowledgeFSAtomicSourceUpdatePayload, - KnowledgeFSOverviewBaseStatsResponse, - KnowledgeFSOverviewQueryOutcomesResponse, - KnowledgeFSPublicFailureResponse, - KnowledgeFSQueryCreatePayload, - KnowledgeFSSourceImportConfigurationPayload, - KnowledgeFSSourceResponse, - KnowledgeFSSourceSyncPolicyResponse, - KnowledgeFSSourceUpdatePayload, - KnowledgeFSSourceWorkflowResponse, - KnowledgeFSStreamCapabilityPayload, -) -from services.knowledge_fs.product_remote import KnowledgeFSOperationUnavailableError -from services.knowledge_fs.runtime import KnowledgeFSRuntime -from tests.unit_tests.config_override import apply_config_overrides -from tests.unit_tests.controllers.knowledge_fs_controller_fakes import ( - FakeComponent, - FakeRuntime, - Recorder, - invocation, - narrowed, - raising, -) - -_RAW_RESULT = object() - - -@pytest.fixture(autouse=True) -def _enabled_knowledge_fs(config_overrides: Callable[..., None]) -> None: - config_overrides(DEPLOYMENT_EDITION=DeploymentEdition.CLOUD, KNOWLEDGE_FS_ENABLED=True) - - -def source_response(**overrides: object) -> KnowledgeFSSourceResponse: - return KnowledgeFSSourceResponse.model_validate( - { - "id": "source-1", - "created_at": datetime(2026, 1, 1, tzinfo=UTC), - "updated_at": datetime(2026, 1, 1, tzinfo=UTC), - "knowledge_space_id": "space-1", - "name": "Source", - "permission_scope": [], - "status": "active", - "metadata": {}, - "type": "web", - "uri": "https://example.com", - "version": 1, - **overrides, - } - ) - - -def sync_policy_response(revision: int) -> KnowledgeFSSourceSyncPolicyResponse: - return KnowledgeFSSourceSyncPolicyResponse.model_validate( - { - "created_at": datetime(2026, 1, 1, tzinfo=UTC), - "enabled": True, - "expected_source_version": 1, - "id": "policy-1", - "knowledge_space_id": "space-1", - "mode": "interval", - "revision": revision, - "source_id": "source-1", - "updated_at": datetime(2026, 1, 1, tzinfo=UTC), - } - ) - - -def issued_capability() -> KnowledgeFSIssuedProductCapability: - return KnowledgeFSIssuedProductCapability( - token="capability-token", - expires_at=datetime(2026, 7, 21, tzinfo=UTC), - operation_id="createQuery", - knowledge_space_id="knowledge-space-1", - knowledge_space_revision=1, - trace_id="trace-1", - ) - - -def _account(account_id: str) -> Account: - account = Account(name="Test User", email="user@example.com") - account.id = account_id - return account - - -@dataclass(frozen=True) -class _SettingsUpdateResult: - """The only field the settings route reads before dumping the facade result.""" - - migration: None = None - - -@dataclass(frozen=True) -class _MembersPayload: - members: tuple[str, ...] - query_images: list[str] = field(default_factory=list) - - -def _invoke[R](resource_type: type[Resource], handler: Callable[..., R], *args: object) -> R: - return unwrap(handler)(resource_type(), *args) - - -# Each case names the public resource, HTTP method, route arguments, and the one -# application boundary that owns the operation. Route-specific identifiers are -# asserted below so accidental delegation to a sibling resource is observable. -_CONSOLE_DELEGATION_CASES = ( - ( - console_resources.KnowledgeFSSpacesApi, - console_resources.KnowledgeFSSpacesApi.get, - (), - "application", - "list_spaces", - dict[str, object](), - ), - ( - console_resources.KnowledgeFSSpacesApi, - console_resources.KnowledgeFSSpacesApi.post, - ( - RequestContext( - request_id="request-1", trace_id=None, account_id="account-1", active_workspace_id="tenant-1" - ), - ), - "application", - "create_space", - dict[str, object](), - ), - ( - console_resources.KnowledgeFSSpaceApi, - console_resources.KnowledgeFSSpaceApi.get, - ("space-1",), - "application", - "get_space", - {"control_space_id": "space-1"}, - ), - ( - console_resources.KnowledgeFSSpaceApi, - console_resources.KnowledgeFSSpaceApi.patch, - ("space-1",), - "application", - "update_space", - {"control_space_id": "space-1"}, - ), - ( - console_resources.KnowledgeFSSpaceApi, - console_resources.KnowledgeFSSpaceApi.delete, - ("space-1",), - "application", - "delete_space", - {"control_space_id": "space-1"}, - ), - ( - console_resources.KnowledgeFSSpacePermissionsApi, - console_resources.KnowledgeFSSpacePermissionsApi.get, - ("space-1",), - "control_plane", - "list_permissions", - {"control_space_id": "space-1"}, - ), - ( - console_resources.KnowledgeFSSpaceMembersApi, - console_resources.KnowledgeFSSpaceMembersApi.put, - ("space-1",), - "control_plane", - "replace_members", - {"control_space_id": "space-1"}, - ), - ( - console_resources.KnowledgeFSSpaceExternalAccessApi, - console_resources.KnowledgeFSSpaceExternalAccessApi.get, - ("space-1",), - "control_plane", - "get_external_access", - {"control_space_id": "space-1"}, - ), - ( - console_resources.KnowledgeFSSpaceExternalAccessApi, - console_resources.KnowledgeFSSpaceExternalAccessApi.put, - ("space-1",), - "control_plane", - "update_external_access", - {"control_space_id": "space-1"}, - ), - ( - console_resources.KnowledgeFSSpaceAppBindingsApi, - console_resources.KnowledgeFSSpaceAppBindingsApi.get, - ("space-1",), - "app_bindings", - "list", - {"control_space_id": "space-1"}, - ), - ( - console_resources.KnowledgeFSSpaceAppBindingsApi, - console_resources.KnowledgeFSSpaceAppBindingsApi.put, - ("space-1",), - "app_bindings", - "upsert", - {"control_space_id": "space-1"}, - ), - ( - console_resources.KnowledgeFSSpaceAppBindingApi, - console_resources.KnowledgeFSSpaceAppBindingApi.delete, - ("space-1", KnowledgeFSAppSpaceJoinType.AGENT.value, "app-1"), - "app_bindings", - "revoke", - { - "control_space_id": "space-1", - "app_id": "app-1", - "caller_kind": KnowledgeFSAppSpaceJoinType.AGENT, - }, - ), - ( - console_resources.KnowledgeFSSpaceSettingsApi, - console_resources.KnowledgeFSSpaceSettingsApi.get, - ("space-1",), - "facade", - "get_settings", - {"control_space_id": "space-1"}, - ), - ( - console_resources.KnowledgeFSSpaceSettingsApi, - console_resources.KnowledgeFSSpaceSettingsApi.patch, - ("space-1",), - "facade", - "update_settings", - {"control_space_id": "space-1"}, - ), - ( - console_resources.KnowledgeFSSpaceOverviewQueryOutcomesApi, - console_resources.KnowledgeFSSpaceOverviewQueryOutcomesApi.get, - ("space-1",), - "facade", - "get_overview_query_outcomes", - {"control_space_id": "space-1", "window": "24h"}, - ), - ( - console_resources.KnowledgeFSSpaceOverviewInventoryApi, - console_resources.KnowledgeFSSpaceOverviewInventoryApi.get, - ("space-1",), - "facade", - "get_overview_inventory", - {"control_space_id": "space-1"}, - ), - ( - console_resources.KnowledgeFSSpaceOverviewAttentionApi, - console_resources.KnowledgeFSSpaceOverviewAttentionApi.get, - ("space-1",), - "facade", - "list_overview_attention", - {"control_space_id": "space-1", "include_dismissed": False, "limit": 50}, - ), - ( - console_resources.KnowledgeFSSpaceOverviewActivityApi, - console_resources.KnowledgeFSSpaceOverviewActivityApi.get, - ("space-1",), - "facade", - "list_overview_activity", - { - "control_space_id": "space-1", - "action": None, - "actor_id": None, - "actor_type": None, - "cursor": None, - "from_at": None, - "limit": 50, - "resource_type": None, - "result": None, - "to_at": None, - }, - ), - ( - console_resources.KnowledgeFSSpaceOverviewHealthApi, - console_resources.KnowledgeFSSpaceOverviewHealthApi.get, - ("space-1",), - "facade", - "get_overview_health", - {"control_space_id": "space-1"}, - ), - ( - console_resources.KnowledgeFSSpaceDocumentsApi, - console_resources.KnowledgeFSSpaceDocumentsApi.get, - ("space-1",), - "facade", - "list_documents", - {"control_space_id": "space-1"}, - ), - ( - console_resources.KnowledgeFSSpaceBulkDocumentsApi, - console_resources.KnowledgeFSSpaceBulkDocumentsApi.delete, - ("space-1",), - "facade", - "bulk_delete_documents", - {"control_space_id": "space-1"}, - ), - ( - console_resources.KnowledgeFSSpaceDocumentReindexApi, - console_resources.KnowledgeFSSpaceDocumentReindexApi.post, - ("space-1",), - "facade", - "reindex_documents", - {"control_space_id": "space-1"}, - ), - ( - console_resources.KnowledgeFSSpaceDocumentApi, - console_resources.KnowledgeFSSpaceDocumentApi.get, - ("space-1", "document-1"), - "facade", - "get_document", - {"control_space_id": "space-1", "document_id": "document-1"}, - ), - ( - console_resources.KnowledgeFSSpaceDocumentApi, - console_resources.KnowledgeFSSpaceDocumentApi.patch, - ("space-1", "document-1"), - "facade", - "update_document_metadata", - {"control_space_id": "space-1", "document_id": "document-1"}, - ), - ( - console_resources.KnowledgeFSSpaceMetadataApi, - console_resources.KnowledgeFSSpaceMetadataApi.get, - ("space-1",), - "facade", - "list_metadata_fields", - {"control_space_id": "space-1", "cursor": None, "limit": 100}, - ), - ( - console_resources.KnowledgeFSSpaceMetadataApi, - console_resources.KnowledgeFSSpaceMetadataApi.post, - ("space-1",), - "facade", - "create_metadata_field", - {"control_space_id": "space-1"}, - ), - ( - console_resources.KnowledgeFSSpaceMetadataFieldApi, - console_resources.KnowledgeFSSpaceMetadataFieldApi.patch, - ("space-1", "field-1"), - "facade", - "update_metadata_field", - {"control_space_id": "space-1", "field_id": "field-1"}, - ), - ( - console_resources.KnowledgeFSSpaceDocumentApi, - console_resources.KnowledgeFSSpaceDocumentApi.delete, - ("space-1", "document-1"), - "facade", - "delete_document", - {"control_space_id": "space-1", "document_id": "document-1"}, - ), - ( - console_resources.KnowledgeFSSpaceLogicalDocumentApi, - console_resources.KnowledgeFSSpaceLogicalDocumentApi.delete, - ("space-1", "document-1"), - "facade", - "delete_logical_document", - {"control_space_id": "space-1", "document_id": "document-1"}, - ), - ( - console_resources.KnowledgeFSSpaceDocumentOutlineApi, - console_resources.KnowledgeFSSpaceDocumentOutlineApi.get, - ("space-1", "document-1"), - "facade", - "get_document_outline", - {"control_space_id": "space-1", "document_id": "document-1"}, - ), - ( - console_resources.KnowledgeFSSpaceDocumentRevisionsApi, - console_resources.KnowledgeFSSpaceDocumentRevisionsApi.get, - ("space-1", "document-1"), - "facade", - "list_document_revisions", - {"control_space_id": "space-1", "document_id": "document-1"}, - ), - ( - console_resources.KnowledgeFSSpaceDocumentChunksApi, - console_resources.KnowledgeFSSpaceDocumentChunksApi.get, - ("space-1", "document-1", 3), - "facade", - "list_document_chunks", - {"control_space_id": "space-1", "document_id": "document-1", "revision": 3}, - ), - ( - console_resources.KnowledgeFSSpaceDocumentChunkApi, - console_resources.KnowledgeFSSpaceDocumentChunkApi.get, - ("space-1", "document-1", 3, "chunk-1"), - "facade", - "get_document_chunk", - { - "control_space_id": "space-1", - "document_id": "document-1", - "revision": 3, - "chunk_id": "chunk-1", - }, - ), - ( - console_resources.KnowledgeFSSpaceCompilationJobApi, - console_resources.KnowledgeFSSpaceCompilationJobApi.get, - ("space-1", "job-1"), - "facade", - "get_compilation_job", - {"control_space_id": "space-1", "job_id": "job-1"}, - ), - ( - console_resources.KnowledgeFSSpaceCompilationJobApi, - console_resources.KnowledgeFSSpaceCompilationJobApi.delete, - ("space-1", "job-1"), - "facade", - "cancel_compilation_job", - {"control_space_id": "space-1", "job_id": "job-1"}, - ), - ( - console_resources.KnowledgeFSSpaceCompilationJobRetryApi, - console_resources.KnowledgeFSSpaceCompilationJobRetryApi.post, - ("space-1", "job-1"), - "facade", - "retry_compilation_job", - {"control_space_id": "space-1", "job_id": "job-1"}, - ), - ( - console_resources.KnowledgeFSSpaceBulkJobApi, - console_resources.KnowledgeFSSpaceBulkJobApi.get, - ("space-1", "job-1"), - "facade", - "get_bulk_job", - {"control_space_id": "space-1", "job_id": "job-1"}, - ), - ( - console_resources.KnowledgeFSSpaceBackgroundTasksApi, - console_resources.KnowledgeFSSpaceBackgroundTasksApi.get, - ("space-1",), - "facade", - "list_background_tasks", - {"control_space_id": "space-1", "cursor": None, "limit": 50}, - ), - ( - console_resources.KnowledgeFSSpaceBackgroundTaskCancelApi, - console_resources.KnowledgeFSSpaceBackgroundTaskCancelApi.post, - ("space-1", "source", "task-1"), - "facade", - "cancel_background_task", - {"control_space_id": "space-1", "task_kind": "source", "task_id": "task-1"}, - ), - ( - console_resources.KnowledgeFSSpaceBackgroundTaskRetryApi, - console_resources.KnowledgeFSSpaceBackgroundTaskRetryApi.post, - ("space-1", "document_bulk", "task-1"), - "facade", - "retry_background_task", - { - "control_space_id": "space-1", - "task_kind": "document_bulk", - "task_id": "task-1", - }, - ), - ( - console_resources.KnowledgeFSSpaceSourcesApi, - console_resources.KnowledgeFSSpaceSourcesApi.get, - ("space-1",), - "facade", - "list_sources", - {"control_space_id": "space-1", "limit": 50}, - ), - ( - console_resources.KnowledgeFSSpaceSourcesApi, - console_resources.KnowledgeFSSpaceSourcesApi.post, - ("space-1",), - "facade", - "create_source", - {"control_space_id": "space-1"}, - ), - ( - console_resources.KnowledgeFSSpaceSourceApi, - console_resources.KnowledgeFSSpaceSourceApi.get, - ("space-1", "source-1"), - "facade", - "get_source", - {"control_space_id": "space-1", "source_id": "source-1"}, - ), - ( - console_resources.KnowledgeFSSpaceSourceApi, - console_resources.KnowledgeFSSpaceSourceApi.patch, - ("space-1", "source-1"), - "facade", - "update_source", - {"control_space_id": "space-1", "source_id": "source-1"}, - ), - ( - console_resources.KnowledgeFSSpaceSourceApi, - console_resources.KnowledgeFSSpaceSourceApi.delete, - ("space-1", "source-1"), - "facade", - "delete_source", - {"control_space_id": "space-1", "source_id": "source-1"}, - ), - ( - console_resources.KnowledgeFSSpaceSourceTestApi, - console_resources.KnowledgeFSSpaceSourceTestApi.post, - ("space-1", "source-1"), - "facade", - "test_source", - {"control_space_id": "space-1", "source_id": "source-1"}, - ), - ( - console_resources.KnowledgeFSSpaceSourceSyncApi, - console_resources.KnowledgeFSSpaceSourceSyncApi.post, - ("space-1", "source-1"), - "facade", - "sync_source", - {"control_space_id": "space-1", "source_id": "source-1"}, - ), - ( - console_resources.KnowledgeFSSpaceSourceWorkflowImportApi, - console_resources.KnowledgeFSSpaceSourceWorkflowImportApi.post, - ("space-1", "source-1"), - "facade", - "import_source_workflow", - {"control_space_id": "space-1", "source_id": "source-1"}, - ), - ( - console_resources.KnowledgeFSSpaceSourcePagesApi, - console_resources.KnowledgeFSSpaceSourcePagesApi.get, - ("space-1", "source-1"), - "facade", - "list_source_pages", - {"control_space_id": "space-1", "source_id": "source-1"}, - ), - ( - console_resources.KnowledgeFSSpaceSourcePageImportApi, - console_resources.KnowledgeFSSpaceSourcePageImportApi.post, - ("space-1", "source-1"), - "facade", - "import_source_pages", - {"control_space_id": "space-1", "source_id": "source-1"}, - ), - ( - console_resources.KnowledgeFSSpaceSourceFilesApi, - console_resources.KnowledgeFSSpaceSourceFilesApi.get, - ("space-1", "source-1"), - "facade", - "list_source_files", - {"control_space_id": "space-1", "source_id": "source-1"}, - ), - ( - console_resources.KnowledgeFSSpaceSourceFileImportApi, - console_resources.KnowledgeFSSpaceSourceFileImportApi.post, - ("space-1", "source-1"), - "facade", - "import_source_files", - {"control_space_id": "space-1", "source_id": "source-1"}, - ), - ( - console_resources.KnowledgeFSSpaceResearchTasksApi, - console_resources.KnowledgeFSSpaceResearchTasksApi.get, - ("space-1",), - "facade", - "list_research_tasks", - {"control_space_id": "space-1"}, - ), - ( - console_resources.KnowledgeFSSpaceResearchTasksApi, - console_resources.KnowledgeFSSpaceResearchTasksApi.post, - ("space-1",), - "facade", - "create_research_task", - {"control_space_id": "space-1"}, - ), - ( - console_resources.KnowledgeFSSpaceResearchTaskPlanApi, - console_resources.KnowledgeFSSpaceResearchTaskPlanApi.post, - ("space-1",), - "facade", - "plan_research_task", - {"control_space_id": "space-1"}, - ), - ( - console_resources.KnowledgeFSSpaceResearchTaskApi, - console_resources.KnowledgeFSSpaceResearchTaskApi.get, - ("space-1", "task-1"), - "facade", - "get_research_task", - {"control_space_id": "space-1", "task_id": "task-1"}, - ), - ( - console_resources.KnowledgeFSSpaceResearchTaskApi, - console_resources.KnowledgeFSSpaceResearchTaskApi.delete, - ("space-1", "task-1"), - "facade", - "cancel_research_task", - {"control_space_id": "space-1", "task_id": "task-1"}, - ), - ( - console_resources.KnowledgeFSSpaceResearchTaskPartialsApi, - console_resources.KnowledgeFSSpaceResearchTaskPartialsApi.get, - ("space-1", "task-1"), - "facade", - "list_research_task_partials", - {"control_space_id": "space-1", "task_id": "task-1"}, - ), - ( - console_resources.KnowledgeFSSpaceTracesApi, - console_resources.KnowledgeFSSpaceTracesApi.get, - ("space-1",), - "facade", - "list_traces", - {"control_space_id": "space-1"}, - ), - ( - console_resources.KnowledgeFSSpaceTraceApi, - console_resources.KnowledgeFSSpaceTraceApi.get, - ("space-1", "trace-1"), - "facade", - "get_trace", - {"control_space_id": "space-1", "trace_id": "trace-1"}, - ), - ( - console_resources.KnowledgeFSSpaceTraceEvidenceApi, - console_resources.KnowledgeFSSpaceTraceEvidenceApi.get, - ("space-1", "trace-1"), - "facade", - "list_trace_entries", - {"control_space_id": "space-1", "trace_id": "trace-1", "kind": "evidence"}, - ), - ( - console_resources.KnowledgeFSSpaceTraceConflictsApi, - console_resources.KnowledgeFSSpaceTraceConflictsApi.get, - ("space-1", "trace-1"), - "facade", - "list_trace_entries", - {"control_space_id": "space-1", "trace_id": "trace-1", "kind": "conflicts"}, - ), - ( - console_resources.KnowledgeFSSpaceTraceMissingApi, - console_resources.KnowledgeFSSpaceTraceMissingApi.get, - ("space-1", "trace-1"), - "facade", - "list_trace_entries", - {"control_space_id": "space-1", "trace_id": "trace-1", "kind": "missing"}, - ), - ( - console_resources.KnowledgeFSSpaceUploadSessionsApi, - console_resources.KnowledgeFSSpaceUploadSessionsApi.post, - ("space-1",), - "facade", - "create_upload_session", - {"control_space_id": "space-1"}, - ), - ( - console_resources.KnowledgeFSSpaceUploadSessionPartPresignApi, - console_resources.KnowledgeFSSpaceUploadSessionPartPresignApi.post, - ("space-1", "session-1", 3), - "facade", - "presign_upload_session_part", - {"control_space_id": "space-1", "upload_session_id": "session-1", "part_number": 3}, - ), - ( - console_resources.KnowledgeFSSpaceUploadSessionCompleteApi, - console_resources.KnowledgeFSSpaceUploadSessionCompleteApi.post, - ("space-1", "session-1"), - "facade", - "complete_upload_session", - {"control_space_id": "space-1", "upload_session_id": "session-1"}, - ), - ( - console_resources.KnowledgeFSSpaceUploadSessionAbortApi, - console_resources.KnowledgeFSSpaceUploadSessionAbortApi.post, - ("space-1", "session-1"), - "facade", - "abort_upload_session", - {"control_space_id": "space-1", "upload_session_id": "session-1"}, - ), - ( - console_resources.KnowledgeFSSpaceSmallFileUploadApi, - console_resources.KnowledgeFSSpaceSmallFileUploadApi.post, - ("space-1", "session-1"), - "facade", - "upload_small_file", - {"control_space_id": "space-1", "upload_session_id": "session-1"}, - ), -) - - -@pytest.mark.parametrize( - ("class_name", "method_name", "route_args", "component_name", "delegate_name", "expected_fields"), - _CONSOLE_DELEGATION_CASES, -) -def test_console_resources_delegate_one_tenant_scoped_product_operation( - monkeypatch: pytest.MonkeyPatch, - class_name: type[Resource], - method_name: Callable[..., object], - route_args: tuple[object, ...], - component_name: str, - delegate_name: str, - expected_fields: dict[str, object], -) -> None: - payload = ( - KnowledgeFSSourceUpdatePayload(name="Renamed") - if class_name.__name__ == "KnowledgeFSSpaceSourceApi" and method_name.__name__ == "patch" - else _MembersPayload(members=("member-1",)) - ) - delegate = Recorder(_RAW_RESULT) - components = {name: FakeComponent() for name in ("application", "control_plane", "app_bindings", "facade")} - components[component_name] = FakeComponent(**{delegate_name: delegate}) - runtime = FakeRuntime( - application=components["application"], - control_plane=components["control_plane"], - app_bindings=components["app_bindings"], - facade=components["facade"], - ) - dump_response = Recorder(effect=lambda schema, raw: (schema.__name__, raw)) - monkeypatch.setattr(console_resources, "_actor", lambda: ("account-1", "tenant-1")) - monkeypatch.setattr(console_resources, "_console_services", lambda: runtime) - monkeypatch.setattr(console_resources, "_payload", lambda _: payload) - monkeypatch.setattr(console_resources, "_idempotency_key", lambda: "idempotency-1") - monkeypatch.setattr(console_resources, "_query_pairs", lambda _: (("normalized", "true"),)) - monkeypatch.setattr(console_resources, "dump_response", dump_response) - app = Flask(__name__) - - with app.test_request_context("/", method="POST"): - result = _invoke(class_name, method_name, *route_args) - - assert len(delegate.calls) == 1 - call_fields = delegate.last_call.kwargs - assert call_fields["tenant_id"] == "tenant-1" - assert call_fields.get("account_id", call_fields.get("actor_account_id")) == "account-1" - for field_name, expected in expected_fields.items(): - assert call_fields[field_name] == expected - if "payload" in call_fields: - if class_name.__name__ == "KnowledgeFSSpaceSourceApi" and method_name.__name__ == "patch": - forwarded = narrowed(call_fields["payload"], BaseModel) - source_update = narrowed(payload, KnowledgeFSSourceUpdatePayload) - assert forwarded.model_dump(exclude_none=True) == source_update.model_dump(exclude_none=True) - else: - assert call_fields["payload"] is payload - if "members" in call_fields: - assert call_fields["members"] == ("member-1",) - if "idempotency_key" in call_fields: - assert call_fields["idempotency_key"] == "idempotency-1" - if dump_response.calls: - assert dump_response.last_call.args[1] is _RAW_RESULT - else: - assert result == ("", HTTPStatus.NO_CONTENT) - - -def test_console_document_reference_delegates_exact_asset_identity(monkeypatch: pytest.MonkeyPatch) -> None: - facade = FakeComponent(resolve_document_reference=Recorder(_RAW_RESULT)) - runtime = FakeRuntime(facade=facade) - monkeypatch.setattr(console_resources, "_actor", lambda: ("account-1", "tenant-1")) - monkeypatch.setattr(console_resources, "_console_services", lambda: runtime) - monkeypatch.setattr(console_resources, "dump_response", lambda _schema, raw: raw) - app = Flask(__name__) - - with app.test_request_context("/?document_asset_id=asset-1&document_asset_version=7"): - result = _invoke( - console_resources.KnowledgeFSSpaceDocumentReferenceApi, - console_resources.KnowledgeFSSpaceDocumentReferenceApi.get, - "space-1", - ) - - assert result is _RAW_RESULT - assert facade.resolve_document_reference.calls == [ - invocation( - tenant_id="tenant-1", - account_id="account-1", - control_space_id="space-1", - document_asset_id="asset-1", - document_asset_version=7, - ) - ] - - -def test_console_space_list_preserves_repeated_creator_filters(monkeypatch: pytest.MonkeyPatch) -> None: - application = FakeComponent(list_spaces=Recorder(_RAW_RESULT)) - runtime = FakeRuntime(application=application) - monkeypatch.setattr(console_resources, "_actor", lambda: ("account-1", "tenant-1")) - monkeypatch.setattr(console_resources, "_console_services", lambda: runtime) - monkeypatch.setattr(console_resources, "dump_response", lambda _schema, raw: raw) - app = Flask(__name__) - - with app.test_request_context( - "/?page=2&limit=10&creator_ids=creator-1&creator_ids=creator-2&tag_ids=tag-1&tag_ids=tag-2&query=%20Support%20" - ): - result = _invoke(console_resources.KnowledgeFSSpacesApi, console_resources.KnowledgeFSSpacesApi.get) - - assert application.list_spaces.calls == [ - invocation( - tenant_id="tenant-1", - account_id="account-1", - page=2, - limit=10, - creator_ids=["creator-1", "creator-2"], - tag_ids=["tag-1", "tag-2"], - query="Support", - ) - ] - assert result is _RAW_RESULT - - -def test_console_source_update_does_not_start_sync_before_resource_selection( - monkeypatch: pytest.MonkeyPatch, -) -> None: - facade = FakeComponent(update_source=Recorder(_RAW_RESULT)) - runtime = FakeRuntime(facade=facade) - payload = KnowledgeFSSourceUpdatePayload.model_validate( - { - "expectedVersion": 3, - "status": "disabled", - "uri": "https://new.example.com", - } - ) - monkeypatch.setattr(console_resources, "_actor", lambda: ("account-1", "tenant-1")) - monkeypatch.setattr(console_resources, "_console_services", lambda: runtime) - monkeypatch.setattr(console_resources, "_payload", lambda _: payload) - monkeypatch.setattr(console_resources, "dump_response", lambda _schema, raw: raw) - app = Flask(__name__) - - with app.test_request_context("/", method="PATCH"): - result = _invoke( - console_resources.KnowledgeFSSpaceSourceApi, - console_resources.KnowledgeFSSpaceSourceApi.patch, - "space-1", - "source-1", - ) - - forwarded = narrowed(facade.update_source.last_call.kwargs["payload"], KnowledgeFSAtomicSourceUpdatePayload) - assert forwarded is not payload - assert forwarded.status == "disabled" - assert forwarded.sync_after_update is None - assert result is _RAW_RESULT - - -@pytest.mark.parametrize( - ("selection", "expected_import_kind"), - [ - ( - {"kind": "website_crawl", "sourceUrls": ["https://new.example.com/a"]}, - "website-crawl-import", - ), - ( - { - "kind": "online_document", - "items": [ - { - "pageId": "page-1", - "providerItemId": "provider-page-1", - "type": "page", - "workspaceId": "workspace-1", - } - ], - }, - "online-document-import", - ), - ( - { - "kind": "online_drive", - "items": [{"id": "file-1", "name": "Plan.pdf", "providerItemId": "provider-file-1"}], - }, - "online-drive-import", - ), - ], -) -def test_console_source_update_commits_complete_edit_selection( - monkeypatch: pytest.MonkeyPatch, - selection: dict[str, object], - expected_import_kind: str, -) -> None: - selection_kind = selection["kind"] - assert isinstance(selection_kind, str) - provider_kind = {"online_document": "online-document", "online_drive": "online-drive"}.get(selection_kind) - original = source_response( - metadata={"providerKind": provider_kind} if provider_kind is not None else {}, - type="web" if selection_kind == "website_crawl" else "connector", - uri="https://old.example.com", - ) - updated = source_response(sync_workflow=None) - facade = FakeComponent( - get_source=Recorder(results=[original, updated]), - update_source=Recorder(source_response(version=4)), - ) - workflow = KnowledgeFSSourceWorkflowResponse.model_validate( - { - "checkpoint": "queued", - "created_at": datetime(2026, 1, 1, tzinfo=UTC), - "execution_attempts": 0, - "id": "workflow-1", - "knowledge_space_id": "space-1", - "kind": "source-import", - "max_execution_attempts": 3, - "progress_completed": 0, - "progress_failed": 0, - "progress_skipped": 0, - "state": "queued", - "updated_at": datetime(2026, 1, 1, tzinfo=UTC), - } - ) - runtime = FakeRuntime(facade=facade) - payload = KnowledgeFSSourceUpdatePayload.model_validate( - { - "expectedVersion": 3, - "name": "Edited source", - "selection": selection, - "syncPolicy": {"enabled": True, "mode": "interval"}, - } - ) - commit = Recorder(workflow) - monkeypatch.setattr(console_resources, "_actor", lambda: ("account-1", "tenant-1")) - monkeypatch.setattr(console_resources, "_console_services", lambda: runtime) - monkeypatch.setattr(console_resources, "_payload", lambda _: payload) - monkeypatch.setattr(source_edit_service, "commit_source_import", commit) - monkeypatch.setattr(console_resources, "dump_response", lambda _schema, raw: raw) - app = Flask(__name__) - - with app.test_request_context("/", method="PATCH"): - result = _invoke( - console_resources.KnowledgeFSSpaceSourceApi, - console_resources.KnowledgeFSSpaceSourceApi.patch, - "space-1", - "source-1", - ) - - assert facade.update_source.calls == [] - source_payload = narrowed(commit.last_call.kwargs["source_update"], KnowledgeFSSourceImportConfigurationPayload) - assert source_payload.status == "disabled" - import_payload = narrowed(commit.last_call.kwargs["payload"], BaseModel) - assert import_payload.model_dump()["kind"] == expected_import_kind - assert commit.last_call.kwargs["idempotency_key"] == "source-edit:source-1:3" - assert updated.sync_workflow is workflow - assert result is updated - - -def test_console_source_update_does_not_import_an_unchanged_complete_selection( - monkeypatch: pytest.MonkeyPatch, -) -> None: - source = source_response( - metadata={"crawled": {"https://example.com/a": {}}}, - sync_policy=None, - type="web", - uri="https://example.com", - version=3, - ) - facade = FakeComponent( - get_source=Recorder(source), - get_source_sync_policy=Recorder(sync_policy_response(revision=2)), - update_source=Recorder(source), - update_source_sync_policy=Recorder(), - ) - runtime = FakeRuntime(facade=facade) - payload = KnowledgeFSSourceUpdatePayload.model_validate( - { - "selection": {"kind": "website_crawl", "sourceUrls": ["https://example.com/a"]}, - "syncPolicy": {"enabled": True, "mode": "interval"}, - } - ) - commit = Recorder() - monkeypatch.setattr(console_resources, "_actor", lambda: ("account-1", "tenant-1")) - monkeypatch.setattr(console_resources, "_console_services", lambda: runtime) - monkeypatch.setattr(console_resources, "_payload", lambda _: payload) - monkeypatch.setattr(source_edit_service, "commit_source_import", commit) - monkeypatch.setattr(console_resources, "dump_response", lambda _schema, raw: raw) - app = Flask(__name__) - - with app.test_request_context("/", method="PATCH"): - result = _invoke( - console_resources.KnowledgeFSSpaceSourceApi, - console_resources.KnowledgeFSSpaceSourceApi.patch, - "space-1", - "source-1", - ) - - assert len(facade.update_source.calls) == 1 - assert commit.calls == [] - assert facade.update_source_sync_policy.calls == [] - command = narrowed(facade.update_source.last_call.kwargs["payload"], KnowledgeFSAtomicSourceUpdatePayload) - assert command.expected_version == 3 - assert command.expected_policy_revision == 2 - assert command.sync_policy is not None - assert command.sync_policy.mode == "interval" - assert result is source - - -def test_source_edit_reimports_when_selection_adds_a_canonical_duplicate() -> None: - source = source_response( - metadata={"crawled": {"https://docs.dify.ai/a": {}}}, - status="active", - type="web", - uri="https://docs.dify.ai", - ) - payload = KnowledgeFSSourceUpdatePayload.model_validate( - { - "selection": { - "kind": "website_crawl", - "sourceUrls": [ - "https://docs.dify.ai/a/", - "https://DOCS.dify.ai:443/a#section", - ], - }, - "syncPolicy": {"enabled": True, "mode": "interval"}, - } - ) - - assert source_edit_service.source_edit_requires_import(source, payload) is True - - -def test_source_edit_ignores_order_when_canonical_url_multiplicity_is_unchanged() -> None: - source = source_response( - metadata={ - "crawled": { - "https://docs.dify.ai/a/": {}, - "https://DOCS.dify.ai:443/a#stored": {}, - } - }, - status="active", - type="web", - uri="https://docs.dify.ai", - ) - payload = KnowledgeFSSourceUpdatePayload.model_validate( - { - "selection": { - "kind": "website_crawl", - "sourceUrls": [ - "https://docs.dify.ai/a#selected", - "https://docs.dify.ai/a/", - ], - }, - "syncPolicy": {"enabled": True, "mode": "interval"}, - } - ) - - assert source_edit_service.source_edit_requires_import(source, payload) is False - - -def test_source_edit_uses_latest_effective_website_selection_instead_of_stale_crawl_metadata() -> None: - source = source_response( - metadata={ - "crawled": {"https://docs.dify.ai/old": {}}, - "initialPreview": {"canonicalSourceUrls": ["https://docs.dify.ai/old"]}, - "__knowledgeFsWebsiteSelection": { - "sourceUrls": ["https://docs.dify.ai/current/"], - "version": 1, - }, - }, - status="active", - type="web", - uri="https://docs.dify.ai", - ) - unchanged = KnowledgeFSSourceUpdatePayload.model_validate( - { - "selection": { - "kind": "website_crawl", - "sourceUrls": ["https://DOCS.dify.ai:443/current#selected"], - }, - "syncPolicy": {"enabled": True, "mode": "interval"}, - } - ) - changed = KnowledgeFSSourceUpdatePayload.model_validate( - { - "selection": { - "kind": "website_crawl", - "sourceUrls": ["https://docs.dify.ai/old"], - }, - "syncPolicy": {"enabled": True, "mode": "interval"}, - } - ) - - assert source_edit_service.source_edit_requires_import(source, unchanged) is False - assert source_edit_service.source_edit_requires_import(source, changed) is True - - -def test_source_edit_reimports_an_unchanged_selection_when_source_is_in_error() -> None: - provider_item_id = '["workspace-a","page-a"]' - identity_hash = sha256(f"online-document\0{provider_item_id}".encode()).hexdigest() - source = source_response( - metadata={ - "providerKind": "online-document", - "__knowledgeFsProviderSelection": {"identityHashes": [identity_hash]}, - }, - status="error", - type="connector", - uri="notion://connection-a", - ) - payload = KnowledgeFSSourceUpdatePayload.model_validate( - { - "selection": { - "kind": "online_document", - "items": [ - { - "pageId": "page-a", - "providerItemId": provider_item_id, - "type": "page", - "workspaceId": "workspace-a", - } - ], - }, - "syncPolicy": {"enabled": True, "mode": "interval"}, - } - ) - - assert source_edit_service.source_edit_requires_import(source, payload) is True - - -def test_console_source_update_applies_policy_without_creating_import_work( - monkeypatch: pytest.MonkeyPatch, -) -> None: - source = source_response(version=5, sync_policy=None) - policy = sync_policy_response(revision=3) - source.sync_policy = policy - facade = FakeComponent( - get_source=Recorder(source), - get_source_sync_policy=Recorder(sync_policy_response(revision=2)), - update_source=Recorder(source), - update_source_sync_policy=Recorder(), - ) - runtime = FakeRuntime(facade=facade) - payload = KnowledgeFSSourceUpdatePayload.model_validate( - {"syncPolicy": {"customIntervalSeconds": 7200, "enabled": True, "mode": "custom"}} - ) - commit = Recorder() - monkeypatch.setattr(console_resources, "_actor", lambda: ("account-1", "tenant-1")) - monkeypatch.setattr(console_resources, "_console_services", lambda: runtime) - monkeypatch.setattr(console_resources, "_payload", lambda _: payload) - monkeypatch.setattr(source_edit_service, "commit_source_import", commit) - monkeypatch.setattr(console_resources, "dump_response", lambda _schema, raw: raw) - app = Flask(__name__) - - with app.test_request_context("/", method="PATCH"): - result = _invoke( - console_resources.KnowledgeFSSpaceSourceApi, - console_resources.KnowledgeFSSpaceSourceApi.patch, - "space-1", - "source-1", - ) - - assert len(facade.update_source.calls) == 1 - assert commit.calls == [] - assert facade.update_source_sync_policy.calls == [] - policy_payload = narrowed(facade.update_source.last_call.kwargs["payload"], KnowledgeFSAtomicSourceUpdatePayload) - assert policy_payload.expected_policy_revision == 2 - assert policy_payload.expected_version == 5 - assert policy_payload.sync_policy is not None - assert policy_payload.sync_policy.custom_interval_seconds == 7200 - assert source.sync_policy is policy - assert result is source - - -def test_console_metadata_delete_forwards_row_version_cas(monkeypatch: pytest.MonkeyPatch) -> None: - facade = FakeComponent(delete_metadata_field=Recorder(_RAW_RESULT)) - runtime = FakeRuntime(facade=facade) - monkeypatch.setattr(console_resources, "_actor", lambda: ("account-1", "tenant-1")) - monkeypatch.setattr(console_resources, "_console_services", lambda: runtime) - monkeypatch.setattr(console_resources, "dump_response", lambda _schema, raw: raw) - app = Flask(__name__) - - with app.test_request_context("/?expectedRowVersion=7", method="DELETE"): - result = _invoke( - console_resources.KnowledgeFSSpaceMetadataFieldApi, - console_resources.KnowledgeFSSpaceMetadataFieldApi.delete, - "space-1", - "field-1", - ) - - assert facade.delete_metadata_field.calls == [ - invocation( - tenant_id="tenant-1", - account_id="account-1", - control_space_id="space-1", - field_id="field-1", - expected_row_version=7, - ) - ] - assert result is _RAW_RESULT - - -def test_console_overview_stats_composes_kfs_metrics_with_dify_app_bindings( - monkeypatch: pytest.MonkeyPatch, -) -> None: - stats = KnowledgeFSOverviewBaseStatsResponse.model_validate( - { - "current": { - "freshSourceCount": 2, - "knowledgeCount": 13, - "latestSourceSyncAt": "2026-07-23T11:00:00.000Z", - "linkedAppCount": 0, - "sourceCount": 3, - "staleSourceCount": 1, - }, - "generatedAt": "2026-07-23T11:59:59.000Z", - "knowledgeSpaceId": "knowledge-space-1", - "windows": { - "24h": { - "answerRate": 0.8, - "answeredQueryCount": 8, - "queryCount": 10, - "since": "2026-07-22T12:00:00.000Z", - }, - "7d": { - "answerRate": 0.75, - "answeredQueryCount": 75, - "queryCount": 100, - "since": "2026-07-16T12:00:00.000Z", - }, - "30d": { - "answerRate": 0.7, - "answeredQueryCount": 210, - "queryCount": 300, - "since": "2026-06-23T12:00:00.000Z", - }, - }, - } - ) - outcomes = KnowledgeFSOverviewQueryOutcomesResponse.model_validate( - { - "buckets": [], - "current": { - "answerRate": 0.8, - "answered": 96, - "lowConfidence": 12, - "noEvidence": 12, - "queryCount": 120, - }, - "generatedAt": "2026-07-23T12:00:00.000Z", - "knowledgeSpaceId": "knowledge-space-1", - "previous": { - "answerRate": 0.85, - "answered": 85, - "lowConfidence": 8, - "noEvidence": 7, - "queryCount": 100, - }, - "previousSince": "2026-07-09T12:00:00.000Z", - "since": "2026-07-16T12:00:00.000Z", - "window": "7d", - } - ) - runtime = FakeRuntime( - facade=FakeComponent( - get_overview_query_outcomes=Recorder(outcomes), - get_overview_stats=Recorder(stats), - ), - app_bindings=FakeComponent(count_active=Recorder(7)), - ) - dump_response = Recorder(effect=lambda _schema, raw: raw) - monkeypatch.setattr(console_resources, "_actor", lambda: ("account-1", "tenant-1")) - monkeypatch.setattr(console_resources, "_console_services", lambda: runtime) - monkeypatch.setattr(console_resources, "dump_response", dump_response) - app = Flask(__name__) - - with app.test_request_context("/?window=7d"): - result = _invoke( - console_resources.KnowledgeFSSpaceOverviewStatsApi, - console_resources.KnowledgeFSSpaceOverviewStatsApi.get, - "control-space-1", - ) - - assert runtime.facade.get_overview_stats.calls == [ - invocation( - tenant_id="tenant-1", - account_id="account-1", - control_space_id="control-space-1", - ) - ] - assert runtime.facade.get_overview_query_outcomes.calls == [ - invocation( - tenant_id="tenant-1", - account_id="account-1", - control_space_id="control-space-1", - window="7d", - ) - ] - assert runtime.app_bindings.count_active.calls == [ - invocation( - tenant_id="tenant-1", - actor_account_id="account-1", - control_space_id="control-space-1", - ) - ] - assert result.documents == 13 - assert result.linked_apps == 7 - assert result.queries.value == 120 - assert result.queries.previous_value == 100 - assert result.queries.change_rate == pytest.approx(0.2) - assert result.answer_rate.value == pytest.approx(0.8) - assert result.answer_rate.change_percentage_points == pytest.approx(-5) - assert result.freshness_seconds == 3600 - assert result.generated_at == outcomes.generated_at - assert dump_response.last_call.args[1] is result - - -_SERVICE_DELEGATION_CASES = ( - ( - service_resources.KnowledgeFSServiceBulkDocumentsApi, - service_resources.KnowledgeFSServiceBulkDocumentsApi.delete, - ("space-1",), - "bulkDeleteDocuments", - dict[str, object](), - ), - ( - service_resources.KnowledgeFSServiceDocumentReindexApi, - service_resources.KnowledgeFSServiceDocumentReindexApi.post, - ("space-1",), - "reindexDocuments", - dict[str, object](), - ), - ( - service_resources.KnowledgeFSServiceDocumentApi, - service_resources.KnowledgeFSServiceDocumentApi.get, - ("space-1", "document-1"), - "getDocument", - {"resource_id": "document-1", "path_parameters": (("documentId", "document-1"),)}, - ), - ( - service_resources.KnowledgeFSServiceDocumentApi, - service_resources.KnowledgeFSServiceDocumentApi.patch, - ("space-1", "document-1"), - "updateDocumentMetadata", - {"resource_id": "document-1", "path_parameters": (("documentId", "document-1"),)}, - ), - ( - service_resources.KnowledgeFSServiceDocumentApi, - service_resources.KnowledgeFSServiceDocumentApi.delete, - ("space-1", "document-1"), - "deleteDocument", - {"resource_id": "document-1", "path_parameters": (("documentId", "document-1"),)}, - ), - ( - service_resources.KnowledgeFSServiceDocumentOutlineApi, - service_resources.KnowledgeFSServiceDocumentOutlineApi.get, - ("space-1", "document-1"), - "getDocumentOutline", - {"resource_id": "document-1", "path_parameters": (("documentId", "document-1"),)}, - ), - ( - service_resources.KnowledgeFSServiceDocumentRevisionsApi, - service_resources.KnowledgeFSServiceDocumentRevisionsApi.get, - ("space-1", "document-1"), - "listDocumentRevisions", - {"resource_id": "document-1", "path_parameters": (("documentId", "document-1"),)}, - ), - ( - service_resources.KnowledgeFSServiceDocumentChunksApi, - service_resources.KnowledgeFSServiceDocumentChunksApi.get, - ("space-1", "document-1", 3), - "listDocumentChunks", - { - "resource_id": "document-1", - "path_parameters": (("documentId", "document-1"), ("revision", "3")), - }, - ), - ( - service_resources.KnowledgeFSServiceDocumentChunkApi, - service_resources.KnowledgeFSServiceDocumentChunkApi.get, - ("space-1", "document-1", 3, "chunk-1"), - "getDocumentChunk", - { - "resource_id": "document-1", - "path_parameters": (("documentId", "document-1"), ("revision", "3"), ("chunkId", "chunk-1")), - }, - ), - ( - service_resources.KnowledgeFSServiceCompilationJobApi, - service_resources.KnowledgeFSServiceCompilationJobApi.get, - ("space-1", "job-1"), - "getCompilationJob", - {"resource_id": "job-1"}, - ), - ( - service_resources.KnowledgeFSServiceCompilationJobApi, - service_resources.KnowledgeFSServiceCompilationJobApi.delete, - ("space-1", "job-1"), - "cancelCompilationJob", - {"resource_id": "job-1"}, - ), - ( - service_resources.KnowledgeFSServiceCompilationJobRetryApi, - service_resources.KnowledgeFSServiceCompilationJobRetryApi.post, - ("space-1", "job-1"), - "retryCompilationJob", - {"resource_id": "job-1"}, - ), - ( - service_resources.KnowledgeFSServiceBulkJobApi, - service_resources.KnowledgeFSServiceBulkJobApi.get, - ("space-1", "job-1"), - "getBulkJob", - {"resource_id": "job-1"}, - ), - ( - service_resources.KnowledgeFSServiceSourceApi, - service_resources.KnowledgeFSServiceSourceApi.get, - ("space-1", "source-1"), - "getSource", - {"resource_id": "source-1", "path_parameters": (("sourceId", "source-1"),)}, - ), - ( - service_resources.KnowledgeFSServiceSourceApi, - service_resources.KnowledgeFSServiceSourceApi.patch, - ("space-1", "source-1"), - "updateSource", - {"resource_id": "source-1", "path_parameters": (("sourceId", "source-1"),)}, - ), - ( - service_resources.KnowledgeFSServiceSourceApi, - service_resources.KnowledgeFSServiceSourceApi.delete, - ("space-1", "source-1"), - "deleteSource", - {"resource_id": "source-1", "path_parameters": (("sourceId", "source-1"),)}, - ), - ( - service_resources.KnowledgeFSServiceSourceTestApi, - service_resources.KnowledgeFSServiceSourceTestApi.post, - ("space-1", "source-1"), - "testSource", - {"resource_id": "source-1", "path_parameters": (("sourceId", "source-1"),)}, - ), - ( - service_resources.KnowledgeFSServiceSourceCrawlApi, - service_resources.KnowledgeFSServiceSourceCrawlApi.post, - ("space-1", "source-1"), - "previewSourceCrawl", - {"resource_id": "source-1", "path_parameters": (("sourceId", "source-1"),)}, - ), - ( - service_resources.KnowledgeFSServiceSourcePagesApi, - service_resources.KnowledgeFSServiceSourcePagesApi.get, - ("space-1", "source-1"), - "listSourcePages", - {"resource_id": "source-1", "path_parameters": (("sourceId", "source-1"),)}, - ), - ( - service_resources.KnowledgeFSServiceSourcePageImportApi, - service_resources.KnowledgeFSServiceSourcePageImportApi.post, - ("space-1", "source-1"), - "importSourceWorkflow", - {"resource_id": "source-1", "path_parameters": (("sourceId", "source-1"),)}, - ), - ( - service_resources.KnowledgeFSServiceSourceFilesApi, - service_resources.KnowledgeFSServiceSourceFilesApi.get, - ("space-1", "source-1"), - "listSourceFiles", - {"resource_id": "source-1", "path_parameters": (("sourceId", "source-1"),)}, - ), - ( - service_resources.KnowledgeFSServiceSourceFileImportApi, - service_resources.KnowledgeFSServiceSourceFileImportApi.post, - ("space-1", "source-1"), - "importSourceWorkflow", - {"resource_id": "source-1", "path_parameters": (("sourceId", "source-1"),)}, - ), - ( - service_resources.KnowledgeFSServiceResearchTasksApi, - service_resources.KnowledgeFSServiceResearchTasksApi.post, - ("space-1",), - "createResearchTask", - {"bind_space_in_body": True}, - ), - ( - service_resources.KnowledgeFSServiceResearchTaskPlanApi, - service_resources.KnowledgeFSServiceResearchTaskPlanApi.post, - ("space-1",), - "planResearchTask", - {"bind_space_in_body": True}, - ), - ( - service_resources.KnowledgeFSServiceResearchTaskApi, - service_resources.KnowledgeFSServiceResearchTaskApi.get, - ("space-1", "task-1"), - "getResearchTask", - {"resource_id": "task-1"}, - ), - ( - service_resources.KnowledgeFSServiceResearchTaskApi, - service_resources.KnowledgeFSServiceResearchTaskApi.delete, - ("space-1", "task-1"), - "cancelResearchTask", - {"resource_id": "task-1"}, - ), - ( - service_resources.KnowledgeFSServiceResearchTaskPartialsApi, - service_resources.KnowledgeFSServiceResearchTaskPartialsApi.get, - ("space-1", "task-1"), - "listResearchTaskPartials", - {"resource_id": "task-1"}, - ), - ( - service_resources.KnowledgeFSServiceTraceApi, - service_resources.KnowledgeFSServiceTraceApi.get, - ("space-1", "trace-1"), - "getTrace", - {"resource_id": "trace-1", "path_parameters": (("traceId", "trace-1"),)}, - ), - ( - service_resources.KnowledgeFSServiceTraceEvidenceApi, - service_resources.KnowledgeFSServiceTraceEvidenceApi.get, - ("space-1", "trace-1"), - "listTraceEvidence", - {"resource_id": "trace-1", "path_parameters": (("traceId", "trace-1"),)}, - ), - ( - service_resources.KnowledgeFSServiceTraceConflictsApi, - service_resources.KnowledgeFSServiceTraceConflictsApi.get, - ("space-1", "trace-1"), - "listTraceConflicts", - {"resource_id": "trace-1", "path_parameters": (("traceId", "trace-1"),)}, - ), - ( - service_resources.KnowledgeFSServiceTraceMissingApi, - service_resources.KnowledgeFSServiceTraceMissingApi.get, - ("space-1", "trace-1"), - "listTraceMissing", - {"resource_id": "trace-1", "path_parameters": (("traceId", "trace-1"),)}, - ), -) - - -@pytest.mark.parametrize( - ("class_name", "method_name", "route_args", "operation_id", "expected_fields"), - _SERVICE_DELEGATION_CASES, -) -def test_service_resources_bind_route_identifiers_to_one_declared_operation( - monkeypatch: pytest.MonkeyPatch, - class_name: type[Resource], - method_name: Callable[..., object], - route_args: tuple[object, ...], - operation_id: str, - expected_fields: dict[str, object], -) -> None: - payload = object() - execute = Recorder(_RAW_RESULT) - dump_response = Recorder(effect=lambda schema, raw: (schema.__name__, raw)) - monkeypatch.setattr(service_resources, "_execute_service_operation", execute) - monkeypatch.setattr(service_resources, "_payload", lambda _: payload) - monkeypatch.setattr(service_resources, "_idempotency_key", lambda: "idempotency-1") - monkeypatch.setattr(service_resources, "_query_pairs", lambda _: (("normalized", "true"),)) - monkeypatch.setattr(service_resources, "dump_response", dump_response) - monkeypatch.setattr(service_resources, "durable_page_import", lambda value: value) - monkeypatch.setattr(service_resources, "durable_file_import", lambda value: value) - monkeypatch.setattr( - service_resources, - "_deletion_accepted", - lambda _space, raw: dump_response(service_resources.KnowledgeFSDurableDeletionAcceptedResponse, raw), - ) - monkeypatch.setattr( - service_resources, - "_bulk_deletion_accepted", - lambda _space, raw: dump_response(service_resources.KnowledgeFSDurableDeletionAcceptedResponse, raw), - ) - app = Flask(__name__) - - with app.test_request_context("/", method="POST", headers={"Idempotency-Key": "idempotency-1"}): - _invoke(class_name, method_name, *route_args) - - assert len(execute.calls) == 1 - call_fields = execute.last_call.kwargs - assert call_fields["control_space_id"] == "space-1" - assert call_fields["operation_id"] == operation_id - for field_name, expected in expected_fields.items(): - assert call_fields[field_name] == expected - if "payload" in call_fields: - assert call_fields["payload"] is payload - if "headers" in call_fields: - assert call_fields["headers"] == (("Idempotency-Key", "idempotency-1"),) - if "query" in call_fields: - assert call_fields["query"] == (("normalized", "true"),) - assert dump_response.last_call.args[1] is _RAW_RESULT - - -@pytest.mark.parametrize( - ("class_name", "method_name", "operation_id"), - [ - ( - service_resources.KnowledgeFSServiceDocumentsApi, - service_resources.KnowledgeFSServiceDocumentsApi.get, - "listDocuments", - ), - ( - service_resources.KnowledgeFSServiceSettingsApi, - service_resources.KnowledgeFSServiceSettingsApi.get, - "getSettings", - ), - ( - service_resources.KnowledgeFSServiceSettingsApi, - service_resources.KnowledgeFSServiceSettingsApi.patch, - "updateSettings", - ), - ( - service_resources.KnowledgeFSServiceSourcesApi, - service_resources.KnowledgeFSServiceSourcesApi.get, - "listSources", - ), - ( - service_resources.KnowledgeFSServiceSourcesApi, - service_resources.KnowledgeFSServiceSourcesApi.post, - "createSource", - ), - ( - service_resources.KnowledgeFSServiceResearchTasksApi, - service_resources.KnowledgeFSServiceResearchTasksApi.get, - "listResearchTasks", - ), - ( - service_resources.KnowledgeFSServiceTracesApi, - service_resources.KnowledgeFSServiceTracesApi.get, - "listTraces", - ), - ], -) -@pytest.mark.parametrize("cursor", [None, "cursor-2"]) -def test_service_credential_routes_validate_profile_before_facade_delegation( - monkeypatch: pytest.MonkeyPatch, - class_name: type[Resource], - method_name: Callable[..., object], - operation_id: str, - cursor: str | None, -) -> None: - facade = FakeComponent( - execute_service=Recorder(_RAW_RESULT), - update_service_settings=Recorder(_SettingsUpdateResult()), - ) - runtime = FakeRuntime(facade=facade) - profile = object() - validate_profile = Recorder(profile) - monkeypatch.setattr(service_resources, "_runtime", lambda: runtime) - monkeypatch.setattr(service_resources, "_profile", validate_profile) - monkeypatch.setattr(service_resources, "_payload", lambda _: object()) - monkeypatch.setattr(service_resources, "dump_response", lambda schema, raw: (schema.__name__, raw)) - app = Flask(__name__) - query_string = {"cursor": cursor} if cursor is not None else None - - with app.test_request_context("/", method="POST", query_string=query_string): - _invoke(class_name, method_name, "space-1") - - assert validate_profile.calls == [invocation(runtime, operation_id=operation_id, control_space_id="space-1")] - if operation_id == "updateSettings": - assert facade.update_service_settings.last_call.kwargs["profile"] is profile - assert facade.execute_service.calls == [] - return - assert facade.execute_service.last_call.kwargs["profile"] is profile - assert facade.execute_service.last_call.kwargs["operation_id"] == operation_id - if "query" in facade.execute_service.last_call.kwargs: - expected_query = (("limit", "50"),) + ((("cursor", cursor),) if cursor else ()) - assert facade.execute_service.last_call.kwargs["query"] == expected_query - - -@pytest.mark.parametrize( - "class_name", - [ - console_resources.KnowledgeFSSpaceQueriesApi, - ], -) -def test_deprecated_buffered_routes_fail_closed(class_name: type[console_resources.KnowledgeFSSpaceQueriesApi]) -> None: - with pytest.raises(KnowledgeFSOperationUnavailableError, match="deprecated"): - _invoke(class_name, class_name.post, "space-1") - - -def test_console_stream_capabilities_bind_the_authorized_resource(monkeypatch: pytest.MonkeyPatch) -> None: - broker = FakeComponent(issue_interactive=Recorder(issued_capability())) - payloads = iter( - [ - KnowledgeFSQueryCreatePayload(query="question", mode="deep"), - KnowledgeFSStreamCapabilityPayload(control_space_id="space-2"), - ] - ) - apply_config_overrides(monkeypatch, CONSOLE_API_URL="https://dify.example") - monkeypatch.setattr(console_resources, "_actor", lambda: ("account-1", "tenant-1")) - monkeypatch.setattr( - console_resources, - "_console_services", - lambda: FakeRuntime(broker=broker), - ) - monkeypatch.setattr(console_resources, "_payload", lambda _: next(payloads)) - monkeypatch.setattr(console_resources, "dump_response", lambda _, response: response) - app = Flask(__name__) - - with app.test_request_context("/", method="POST"): - query = _invoke( - console_resources.KnowledgeFSSpaceQueryAdmissionApi, - console_resources.KnowledgeFSSpaceQueryAdmissionApi.post, - "space-1", - ) - task = _invoke( - console_resources.KnowledgeFSTaskStreamCapabilityApi, - console_resources.KnowledgeFSTaskStreamCapabilityApi.post, - "task/1", - ) - legacy_query = _invoke( - console_resources.KnowledgeFSSpaceQueryStreamCapabilityApi, - console_resources.KnowledgeFSSpaceQueryStreamCapabilityApi.post, - "space-1", - ) - - assert query.request.knowledge_space_id == "knowledge-space-1" - assert query.url == "https://dify.example/console/api/knowledge-fs/query-stream" - assert ( - task.url == "https://dify.example/console/api/knowledge-fs/research-tasks/task%2F1/events" - "?knowledgeSpaceId=knowledge-space-1" - ) - assert legacy_query.url == "https://dify.example/console/api/knowledge-fs/query-stream" - assert broker.issue_interactive.calls[0].kwargs["operation_id"] == "createQuery" - assert broker.issue_interactive.calls[1].kwargs == { - "tenant_id": "tenant-1", - "account_id": "account-1", - "control_space_id": "space-2", - "operation_id": "streamResearchTask", - "resource_id": "task/1", - } - assert broker.issue_interactive.calls[2].kwargs["operation_id"] == "createQuery" - - -def test_service_query_admission_binds_profile_space_and_payload(monkeypatch: pytest.MonkeyPatch) -> None: - broker = FakeComponent(issue_service=Recorder(issued_capability())) - runtime = FakeRuntime(broker=broker) - profile = object() - apply_config_overrides(monkeypatch, SERVICE_API_URL="https://api.dify.example") - monkeypatch.setattr(service_resources, "_runtime", lambda: runtime) - monkeypatch.setattr(service_resources, "_profile", lambda *_args, **_kwargs: profile) - monkeypatch.setattr(service_resources, "_payload", lambda _: KnowledgeFSQueryCreatePayload(query="question")) - monkeypatch.setattr(service_resources, "dump_response", lambda _, response: response) - app = Flask(__name__) - - with app.test_request_context("/", method="POST"): - response = _invoke( - service_resources.KnowledgeFSServiceQueryAdmissionApi, - service_resources.KnowledgeFSServiceQueryAdmissionApi.post, - "space-1", - ) - - assert broker.issue_service.calls == [invocation(profile=profile, operation_id="createQuery")] - assert response.request.knowledge_space_id == "knowledge-space-1" - assert response.request.query == "question" - assert response.url == "https://api.dify.example/v1/knowledge-fs/query-stream" - - -def test_console_resource_helpers_validate_feature_payload_headers_and_query_pairs( - monkeypatch: pytest.MonkeyPatch, -) -> None: - runtime = FakeRuntime() - session_maker = object() - get_runtime = Recorder(runtime) - apply_config_overrides(monkeypatch, DEPLOYMENT_EDITION=DeploymentEdition.CLOUD, KNOWLEDGE_FS_ENABLED=True) - monkeypatch.setattr(console_resources.session_factory, "get_session_maker", lambda: session_maker) - monkeypatch.setattr(console_resources, "get_knowledge_fs_runtime", get_runtime) - monkeypatch.setattr( - console_resources, - "current_account_with_tenant", - lambda: (_account("account-1"), "tenant-1"), - ) - app = Flask(__name__) - - assert console_resources._console_services() is runtime - assert get_runtime.calls == [invocation(session_maker)] - assert console_resources._actor() == ("account-1", "tenant-1") - query = KnowledgeFSQueryCreatePayload(query="question", mode="fast", active_document_ids=["doc-1"]) - assert console_resources._query_pairs(query) == ( - ("query", "question"), - ("mode", "fast"), - ("activeDocumentIds", "['doc-1']"), - ("activeEntityIds", "[]"), - ) - with app.test_request_context( - "/", - method="POST", - json={"query": "question"}, - headers={"Idempotency-Key": "request-key"}, - ): - assert console_resources._payload(KnowledgeFSQueryCreatePayload).query == "question" - assert console_resources._idempotency_key() == "request-key" - - apply_config_overrides(monkeypatch, DEPLOYMENT_EDITION=DeploymentEdition.COMMUNITY) - with pytest.raises(NotFound): - console_resources._console_services() - - apply_config_overrides(monkeypatch, KNOWLEDGE_FS_COMMUNITY_DEV_ENABLED=True) - assert console_resources._console_services() is runtime - - apply_config_overrides(monkeypatch, KNOWLEDGE_FS_ENABLED=False) - with pytest.raises(NotFound): - console_resources._console_services() - - -def test_service_resource_helpers_validate_feature_bearer_headers_and_boolean_queries( - monkeypatch: pytest.MonkeyPatch, -) -> None: - runtime = FakeRuntime() - session_maker = object() - get_runtime = Recorder(runtime) - apply_config_overrides(monkeypatch, DEPLOYMENT_EDITION=DeploymentEdition.CLOUD, KNOWLEDGE_FS_ENABLED=True) - monkeypatch.setattr(service_resources.session_factory, "get_session_maker", lambda: session_maker) - monkeypatch.setattr(service_resources, "get_knowledge_fs_runtime", get_runtime) - app = Flask(__name__) - - assert service_resources._runtime() is runtime - assert get_runtime.calls == [invocation(session_maker)] - - class Query(BaseModel): - enabled: bool = True - disabled: bool = False - count: int = 2 - - query = Query() - assert service_resources._query_pairs(query) == ( - ("enabled", "true"), - ("disabled", "false"), - ("count", "2"), - ) - with app.test_request_context( - "/", - method="POST", - json={"query": "question"}, - headers={"Authorization": "bEaReR credential-value ", "Idempotency-Key": "request-key"}, - ): - assert service_resources._payload(KnowledgeFSQueryCreatePayload).query == "question" - assert service_resources._idempotency_key() == "request-key" - - apply_config_overrides(monkeypatch, DEPLOYMENT_EDITION=DeploymentEdition.COMMUNITY) - with pytest.raises(NotFound): - service_resources._runtime() - - apply_config_overrides(monkeypatch, KNOWLEDGE_FS_COMMUNITY_DEV_ENABLED=True) - assert service_resources._runtime() is runtime - - apply_config_overrides(monkeypatch, KNOWLEDGE_FS_ENABLED=False) - with pytest.raises(NotFound): - service_resources._runtime() - - -def test_console_request_rejections_preserve_http_status_contracts() -> None: - from services.knowledge_fs.product_remote import KnowledgeFSProductRequestRejectedError - - expected: dict[Literal[400, 403, 409, 413, 422, 429, 503, 504], type[BaseHTTPException]] = { - 400: KnowledgeFSInvalidRequestHTTPError, - 403: KnowledgeFSAccessDeniedHTTPError, - 409: KnowledgeFSConflictHTTPError, - 413: KnowledgeFSRequestTooLargeHTTPError, - 422: KnowledgeFSRequestRejectedHTTPError, - 429: KnowledgeFSRateLimitHTTPError, - 503: KnowledgeFSOperationUnavailableHTTPError, - 504: KnowledgeFSTimeoutHTTPError, - } - for status, http_error in expected.items(): - guarded = console_resources._knowledge_fs_errors( - raising(KnowledgeFSProductRequestRejectedError(status_code=status)) - ) - - with pytest.raises(http_error) as raised: - guarded() - - assert raised.value.code == status - assert raised.value.data == { - "code": http_error.error_code, - "message": http_error.description, - "status": status, - } - - -@pytest.mark.parametrize( - ("status", "failure_code", "category", "safe_message"), - [ - ( - 503, - "DOCUMENT_COUNT_QUOTA_UNAVAILABLE", - "dependency", - "The workspace document quota could not be verified. Try again later.", - ), - ( - 503, - "VECTOR_SPACE_QUOTA_UNAVAILABLE", - "dependency", - "The workspace vector storage quota could not be verified. Try again later.", - ), - (504, "KNOWLEDGE_FS_TIMEOUT", "timeout", "The Agent Knowledge Base operation timed out. Try again later."), - ], -) -def test_console_unavailable_and_timeout_rejections_preserve_safe_failure_metadata( - status: Literal[503, 504], failure_code: str, category: str, safe_message: str -) -> None: - from services.knowledge_fs.product_remote import KnowledgeFSProductRequestRejectedError - - expected_failure = { - "action": "retry", - "category": category, - "code": failure_code, - "message": safe_message, - "retryPolicy": "manual", - "traceId": "trace-quota", - "parameters": {"retryAfterSeconds": 30}, - } - failure = KnowledgeFSPublicFailureResponse.model_validate( - {**expected_failure, "message": "Private upstream diagnostic that must not reach the client"} - ) - guarded = console_resources._knowledge_fs_errors( - raising(KnowledgeFSProductRequestRejectedError(status_code=status, failure=failure)) - ) - http_error = KnowledgeFSOperationUnavailableHTTPError if status == 503 else KnowledgeFSTimeoutHTTPError - - with pytest.raises(http_error) as raised: - guarded() - - assert raised.value.code == status - assert raised.value.data == { - "code": http_error.error_code, - "message": safe_message, - "status": status, - "failure": expected_failure, - } - - -def test_console_request_rejection_preserves_safe_failure_metadata() -> None: - from services.knowledge_fs.product_remote import KnowledgeFSProductRequestRejectedError - - failure = KnowledgeFSPublicFailureResponse.model_validate( - { - "action": "configure_model", - "category": "configuration", - "code": "MODEL_SELECTION_NOT_FOUND", - "message": "Select another model before retrying.", - "retryPolicy": "after_configuration", - "traceId": "trace-model", - } - ) - guarded = console_resources._knowledge_fs_errors( - raising(KnowledgeFSProductRequestRejectedError(status_code=422, failure=failure)) - ) - - with pytest.raises(KnowledgeFSRequestRejectedHTTPError) as raised: - guarded() - - assert raised.value.data == { - "code": "knowledge_fs_request_rejected", - "message": "The Agent Knowledge Base operation requires a configuration change before it can continue.", - "status": 422, - "failure": { - "action": "configure_model", - "category": "configuration", - "code": "MODEL_SELECTION_NOT_FOUND", - "message": "The Agent Knowledge Base operation requires a configuration change before it can continue.", - "retryPolicy": "after_configuration", - "traceId": "trace-model", - }, - } - - -def test_jwks_resource_fails_closed_for_disabled_missing_and_misconfigured_issuers( - monkeypatch: pytest.MonkeyPatch, -) -> None: - from services.knowledge_fs_capability import KnowledgeFSCapabilityConfigurationError - - app = Flask(__name__) - apply_config_overrides(monkeypatch, KNOWLEDGE_FS_CAPABILITY_V2_ENABLED=False) - with app.app_context(), pytest.raises(NotFound): - console_resources.KnowledgeFSJWKSApi().get() - - apply_config_overrides(monkeypatch, KNOWLEDGE_FS_CAPABILITY_V2_ENABLED=True) - monkeypatch.setattr(console_resources.session_factory, "get_session_maker", lambda: object()) - monkeypatch.setattr(console_resources, "create_configured_knowledge_fs_capability_issuer", lambda **_: None) - with app.app_context(), pytest.raises(NotFound): - console_resources.KnowledgeFSJWKSApi().get() - - def misconfigured(**_: object) -> None: - raise KnowledgeFSCapabilityConfigurationError("missing signing key") - - monkeypatch.setattr(console_resources, "create_configured_knowledge_fs_capability_issuer", misconfigured) - with app.app_context(), pytest.raises(ServiceUnavailable, match="not configured"): - console_resources.KnowledgeFSJWKSApi().get() - - -def test_console_error_adapter_maps_every_domain_boundary_to_the_stable_http_contract() -> None: - from pydantic import ValidationError - - from services.knowledge_fs.app_binding_management import KnowledgeFSAppBindingManagementError - from services.knowledge_fs.control_plane_service import KnowledgeFSControlPlaneInvariantError - from services.knowledge_fs.product_authorization import KnowledgeFSProductNotFoundError - from services.knowledge_fs.product_remote import ( - KnowledgeFSProductRemoteError, - KnowledgeFSProductResourceNotFoundError, - ) - from services.knowledge_fs_capability import KnowledgeFSCapabilityConfigurationError - - with pytest.raises(ValidationError) as raised_validation: - KnowledgeFSQueryCreatePayload.model_validate({"query": ""}) - validation_error = raised_validation.value - mappings = ( - (KnowledgeFSProductNotFoundError("hidden"), KnowledgeFSSpaceNotFoundHTTPError), - (KnowledgeFSProductResourceNotFoundError("missing child"), KnowledgeFSResourceNotFoundHTTPError), - (KnowledgeFSOperationUnavailableError("manifest mismatch"), KnowledgeFSOperationUnavailableHTTPError), - (KnowledgeFSProductRemoteError("upstream unavailable"), KnowledgeFSUpstreamUnavailableHTTPError), - (KnowledgeFSAppBindingManagementError("invalid binding"), KnowledgeFSInvalidRequestHTTPError), - (KnowledgeFSControlPlaneInvariantError("missing revision"), KnowledgeFSInvalidRequestHTTPError), - (validation_error, KnowledgeFSInvalidRequestHTTPError), - (PermissionError("forbidden"), KnowledgeFSAccessDeniedHTTPError), - (KnowledgeFSCapabilityConfigurationError("missing issuer"), KnowledgeFSOperationUnavailableHTTPError), - ) - - for domain_error, http_error in mappings: - guarded = console_resources._knowledge_fs_errors(raising(domain_error)) - - with pytest.raises(http_error): - guarded() - - assert KnowledgeFSUpstreamUnavailableHTTPError.code == HTTPStatus.SERVICE_UNAVAILABLE - - -def test_service_error_adapter_maps_every_domain_boundary_to_the_stable_http_contract() -> None: - from pydantic import ValidationError - - from services.knowledge_fs.product_remote import ( - KnowledgeFSProductRemoteError, - KnowledgeFSProductRequestRejectedError, - KnowledgeFSProductResourceNotFoundError, - ) - from services.knowledge_fs.service_api_authorization import KnowledgeFSServiceApiAuthorizationError - - with pytest.raises(ValidationError) as raised_validation: - KnowledgeFSQueryCreatePayload.model_validate({"query": ""}) - validation_error = raised_validation.value - mappings = ( - (KnowledgeFSServiceApiAuthorizationError("revoked"), KnowledgeFSInvalidCredentialHTTPError), - (KnowledgeFSOperationUnavailableError("manifest mismatch"), KnowledgeFSServiceOperationUnavailableHTTPError), - ( - KnowledgeFSProductResourceNotFoundError("missing child"), - KnowledgeFSServiceResourceNotFoundHTTPError, - ), - (KnowledgeFSProductRemoteError("upstream unavailable"), KnowledgeFSServiceUpstreamUnavailableHTTPError), - (validation_error, KnowledgeFSServiceInvalidRequestHTTPError), - ) - - for domain_error, http_error in mappings: - guarded = service_resources._service_api_errors(raising(domain_error)) - - with pytest.raises(http_error): - guarded() - - rejected_mappings = ( - (400, KnowledgeFSServiceInvalidRequestHTTPError), - (403, KnowledgeFSServiceAccessDeniedHTTPError), - (409, KnowledgeFSServiceConflictHTTPError), - (413, KnowledgeFSServiceRequestTooLargeHTTPError), - (422, KnowledgeFSServiceRequestRejectedHTTPError), - (429, KnowledgeFSServiceRateLimitHTTPError), - ) - for status, http_error in rejected_mappings: - guarded = service_resources._service_api_errors( - raising(KnowledgeFSProductRequestRejectedError(status_code=status)) - ) - with pytest.raises(http_error): - guarded() - - -def test_console_small_file_reader_rejects_malformed_and_empty_multipart_bodies() -> None: - from io import BytesIO - - from services.knowledge_fs.product_remote import KnowledgeFSProductRequestRejectedError - - app = Flask(__name__) - with app.test_request_context( - "/", - method="POST", - data=b"x", - environ_overrides={"CONTENT_LENGTH": str(64 * 1024 + 2)}, - ): - with pytest.raises(KnowledgeFSProductRequestRejectedError) as too_large: - console_resources._read_small_file_body(1) - assert too_large.value.status_code == HTTPStatus.REQUEST_ENTITY_TOO_LARGE - - malformed_payloads = ( - dict[str, object](), - {"file": (BytesIO(b"content"), "")}, - {"file": (BytesIO(b""), "empty.txt")}, - ) - for payload in malformed_payloads: - with app.test_request_context("/", method="POST", data=payload, content_type="multipart/form-data"): - with pytest.raises(KnowledgeFSProductRequestRejectedError) as rejected: - console_resources._read_small_file_body(10) - assert rejected.value.status_code == HTTPStatus.UNPROCESSABLE_ENTITY - - -def test_console_app_binding_route_rejects_unknown_caller_kind_before_revoke( - monkeypatch: pytest.MonkeyPatch, -) -> None: - from services.knowledge_fs.app_binding_management import KnowledgeFSAppBindingManagementError - - revoke = Recorder() - monkeypatch.setattr(console_resources, "_actor", lambda: ("account-1", "tenant-1")) - monkeypatch.setattr( - console_resources, - "_console_services", - lambda: FakeRuntime(app_bindings=FakeComponent(revoke=revoke)), - ) - - with pytest.raises(KnowledgeFSAppBindingManagementError, match="caller kind"): - _invoke( - console_resources.KnowledgeFSSpaceAppBindingApi, - console_resources.KnowledgeFSSpaceAppBindingApi.delete, - "space-1", - "unknown-caller", - "app-1", - ) - - assert revoke.calls == [] - - -def test_service_profile_hides_unknown_operations_before_dataset_key_authorization( - monkeypatch: pytest.MonkeyPatch, -) -> None: - authorization = FakeComponent(authorize=Recorder()) - # The profile check must fail before it touches any runtime component but authorization. - runtime = cast(KnowledgeFSRuntime, FakeRuntime(service_api_authorization=authorization)) - monkeypatch.setattr(service_resources, "product_operation_action", raising(KeyError("unknown"))) - - with pytest.raises(KnowledgeFSOperationUnavailableError, match="unknownOperation"): - service_resources._profile( - runtime, - operation_id="unknownOperation", - control_space_id="space-1", - ) - - assert authorization.authorize.calls == [] - - -def test_service_operation_helper_forwards_the_complete_validated_request( - monkeypatch: pytest.MonkeyPatch, -) -> None: - profile = object() - payload = KnowledgeFSQueryCreatePayload(query="question") - facade = FakeComponent(execute_service=Recorder({"id": "result-1"})) - runtime = FakeRuntime(facade=facade) - monkeypatch.setattr(service_resources, "_runtime", lambda: runtime) - monkeypatch.setattr(service_resources, "_profile", Recorder(profile)) - - result = service_resources._execute_service_operation( - control_space_id="space-1", - operation_id="createResearchTask", - payload=payload, - query=(("cursor", "next"),), - bind_space_in_body=True, - resource_id="task-1", - path_parameters=(("taskId", "task-1"),), - headers=(("Idempotency-Key", "request-key"),), - ) - - assert result == {"id": "result-1"} - assert facade.execute_service.calls == [ - invocation( - profile=profile, - operation_id="createResearchTask", - payload=payload, - query=(("cursor", "next"),), - bind_space_in_body=True, - resource_id="task-1", - path_parameters=(("taskId", "task-1"),), - headers=(("Idempotency-Key", "request-key"),), - ) - ] diff --git a/api/tests/unit_tests/controllers/test_mcp_app_api.py b/api/tests/unit_tests/controllers/test_mcp_app_api.py deleted file mode 100644 index 62937cb26b2..00000000000 --- a/api/tests/unit_tests/controllers/test_mcp_app_api.py +++ /dev/null @@ -1,44 +0,0 @@ -"""SQLite-backed unit tests for MCP app controller database helpers.""" - -import pytest -from sqlalchemy.orm import Session - -from controllers.mcp.mcp import MCPAppApi -from models.model import App, AppMCPServer, AppMode, IconType -from services.app.mcp_server_service import AppMCPServerStatus - -TENANT_ID = "11111111-1111-1111-1111-111111111111" -APP_ID = "22222222-2222-2222-2222-222222222222" - - -@pytest.mark.parametrize("sqlite_session", [(App, AppMCPServer)], indirect=True) -def test_get_mcp_server_and_app_returns_persisted_models(sqlite_session: Session) -> None: - app = App( - tenant_id=TENANT_ID, - name="MCP App", - description="", - mode=AppMode.ADVANCED_CHAT, - icon_type=IconType.EMOJI, - icon="🤖", - icon_background="#FFFFFF", - enable_site=False, - enable_api=False, - max_active_requests=0, - ) - app.id = APP_ID - server = AppMCPServer( - tenant_id=TENANT_ID, - app_id=APP_ID, - name="Test MCP Server", - description="Test server", - server_code="server-1", - status=AppMCPServerStatus.ACTIVE, - parameters="{}", - ) - sqlite_session.add_all([app, server]) - sqlite_session.commit() - - result_server, result_app = MCPAppApi()._get_mcp_server_and_app("server-1", sqlite_session) - - assert result_server is server - assert result_app is app diff --git a/api/tests/unit_tests/controllers/trigger/test_webhook.py b/api/tests/unit_tests/controllers/trigger/test_webhook.py index 0e80e896b94..8537780e1e2 100644 --- a/api/tests/unit_tests/controllers/trigger/test_webhook.py +++ b/api/tests/unit_tests/controllers/trigger/test_webhook.py @@ -8,7 +8,7 @@ import controllers.trigger.webhook as module from models.trigger import WorkflowWebhookTrigger from models.workflow import Workflow from services.errors.app import QuotaExceededError -from tests.unit_tests.model_factories import make_workflow +from tests.helpers.unit.model_factories import make_workflow class _RequestStub: diff --git a/api/tests/unit_tests/controllers/web/test_app.py b/api/tests/unit_tests/controllers/web/test_app.py index 80f591f218e..bb661be0c2c 100644 --- a/api/tests/unit_tests/controllers/web/test_app.py +++ b/api/tests/unit_tests/controllers/web/test_app.py @@ -29,7 +29,7 @@ from services.webapp_access_query_service import ( WebAppAccessReferenceRequiredError, WebAppAccessUnavailableError, ) -from tests.unit_tests.config_override import config_overrides_context +from tests.helpers.unit.config_override import config_overrides_context PASSPORT_SECRET = "test-secret-that-is-at-least-32-bytes" OTHER_PASSPORT_SECRET = "other-secret-that-is-at-least-32-bytes" diff --git a/api/tests/unit_tests/controllers/web/test_audio.py b/api/tests/unit_tests/controllers/web/test_audio.py index 3163c106836..b86a439f3c4 100644 --- a/api/tests/unit_tests/controllers/web/test_audio.py +++ b/api/tests/unit_tests/controllers/web/test_audio.py @@ -29,8 +29,8 @@ from models.model import App, AppMode, AppModelConfig, EndUser from services.app_ref_service import AppRef, MessageRef from services.audio_service import FILE_SIZE_LIMIT, AudioService from services.errors.audio import ProviderNotSupportSpeechToTextServiceError -from tests.unit_tests.audio_runtime_fixtures import AudioRuntimeObservations, record_calls -from tests.unit_tests.model_factories import make_app, make_end_user, make_message +from tests.helpers.unit.audio_runtime_fixtures import AudioRuntimeObservations, record_calls +from tests.helpers.unit.model_factories import make_app, make_end_user, make_message @pytest.fixture diff --git a/api/tests/unit_tests/controllers/web/test_completion.py b/api/tests/unit_tests/controllers/web/test_completion.py index a9e8c595ea3..c42856c193c 100644 --- a/api/tests/unit_tests/controllers/web/test_completion.py +++ b/api/tests/unit_tests/controllers/web/test_completion.py @@ -25,7 +25,7 @@ from core.app.apps.agent_app.errors import AgentAppNotPublishedError from core.errors.error import ModelCurrentlyNotSupportError, ProviderTokenNotInitError, QuotaExceededError from graphon.model_runtime.errors.invoke import InvokeError from models.model import App, AppMode, EndUser -from tests.unit_tests.model_factories import make_app, make_end_user +from tests.helpers.unit.model_factories import make_app, make_end_user def _app(mode: AppMode) -> App: diff --git a/api/tests/unit_tests/controllers/web/test_pydantic_models.py b/api/tests/unit_tests/controllers/web/test_conversation.py similarity index 100% rename from api/tests/unit_tests/controllers/web/test_pydantic_models.py rename to api/tests/unit_tests/controllers/web/test_conversation.py diff --git a/api/tests/unit_tests/controllers/web/test_files.py b/api/tests/unit_tests/controllers/web/test_files.py index 30c36a8af24..52c3ec1dc37 100644 --- a/api/tests/unit_tests/controllers/web/test_files.py +++ b/api/tests/unit_tests/controllers/web/test_files.py @@ -21,7 +21,7 @@ from libs.exception import BaseHTTPException from models.enums import CreatorUserRole from models.model import App, AppMode, EndUser, UploadFile from services.errors import file as file_errors -from tests.unit_tests.model_factories import make_end_user, make_upload_file +from tests.helpers.unit.model_factories import make_end_user, make_upload_file def _app_model() -> App: diff --git a/api/tests/unit_tests/controllers/web/test_web_forgot_password.py b/api/tests/unit_tests/controllers/web/test_forgot_password.py similarity index 100% rename from api/tests/unit_tests/controllers/web/test_web_forgot_password.py rename to api/tests/unit_tests/controllers/web/test_forgot_password.py diff --git a/api/tests/unit_tests/controllers/web/test_human_input_file_upload.py b/api/tests/unit_tests/controllers/web/test_human_input_file_upload.py index 42117e7fbe5..74633eee425 100644 --- a/api/tests/unit_tests/controllers/web/test_human_input_file_upload.py +++ b/api/tests/unit_tests/controllers/web/test_human_input_file_upload.py @@ -53,8 +53,8 @@ from services.remote_file_service import RemoteFileNotFoundError as RemoteFileNo from services.remote_file_service import RemoteFileUnavailableError as RemoteFileUnavailableServiceError from services.remote_file_service import RemoteFileUploadResult from services.remote_file_service import RemoteFileUrlBlockedError as RemoteFileUrlBlockedServiceError -from tests.unit_tests.controllers.error_body import raised_error -from tests.unit_tests.model_factories import make_account, make_upload_file +from tests.helpers.unit.controllers.error_body import raised_error +from tests.helpers.unit.model_factories import make_account, make_upload_file @pytest.fixture diff --git a/api/tests/unit_tests/controllers/web/test_web_login.py b/api/tests/unit_tests/controllers/web/test_login.py similarity index 99% rename from api/tests/unit_tests/controllers/web/test_web_login.py rename to api/tests/unit_tests/controllers/web/test_login.py index 62c1d392498..d2a05c06697 100644 --- a/api/tests/unit_tests/controllers/web/test_web_login.py +++ b/api/tests/unit_tests/controllers/web/test_login.py @@ -40,7 +40,7 @@ from services.webapp_access_query_service import ( WebAppAccessAppNotFoundError, WebAppAccessUnavailableError, ) -from tests.unit_tests.config_override import apply_config_overrides +from tests.helpers.unit.config_override import apply_config_overrides @pytest.fixture diff --git a/api/tests/unit_tests/controllers/web/test_message.py b/api/tests/unit_tests/controllers/web/test_message.py new file mode 100644 index 00000000000..148ac9ecc3d --- /dev/null +++ b/api/tests/unit_tests/controllers/web/test_message.py @@ -0,0 +1,419 @@ +from __future__ import annotations + +import builtins +import inspect +import json +import uuid +from datetime import datetime +from pathlib import Path +from types import ModuleType, SimpleNamespace +from unittest.mock import ANY, MagicMock, patch +from uuid import uuid4 + +import pytest +from flask import Flask +from flask.views import MethodView +from werkzeug.exceptions import NotFound + +from controllers.common.controller_schemas import MessageFeedbackPayload, MessageListQuery +from controllers.web.error import ( + AppMoreLikeThisDisabledError, + NotChatAppError, + NotCompletionAppError, +) +from controllers.web.message import ( + MessageFeedbackApi, + MessageMoreLikeThisApi, + MessageMoreLikeThisQuery, + MessageSuggestedQuestionApi, +) +from core.app.entities.app_invoke_entities import InvokeFrom +from core.entities.execution_extra_content import HumanInputContent +from models.enums import ConversationFromSource, EndUserType, FeedbackFromSource, FeedbackRating +from models.model import ( + App, + AppMode, + Conversation, + EndUser, + Message, + MessageAgentThought, + MessageFeedback, +) +from services.errors.app import MoreLikeThisDisabledError +from services.errors.message import MessageNotExistsError +from tests.helpers.unit.model_factories import make_end_user + +"""Unit tests for controllers.web.message — feedback, more-like-this, suggested questions.""" + + +def _chat_app() -> App: + return App(id="app-1", tenant_id="tenant-1", mode=AppMode.CHAT) + + +def _completion_app() -> App: + return App(id="app-1", tenant_id="tenant-1", mode=AppMode.COMPLETION) + + +def _end_user() -> EndUser: + return make_end_user(end_user_id="eu-1") + + +# The @model_validate and @with_session decorators wrap the handlers; tests +# call the undecorated function so they can pass a pydantic payload / a fake +# session directly. +_feedback_post = inspect.unwrap(MessageFeedbackApi.post) +_more_like_this_get = inspect.unwrap(MessageMoreLikeThisApi.get) + + +class TestMessageEndpointsCases: + # --------------------------------------------------------------------------- + # MessageFeedbackApi + # --------------------------------------------------------------------------- + class TestMessageFeedbackApi: + @patch("controllers.web.message.MessageService.create_feedback") + def test_feedback_success(self, mock_create: MagicMock, app: Flask) -> None: + payload = MessageFeedbackPayload.model_validate({"rating": "like", "content": "great"}) + msg_id = uuid4() + + with app.test_request_context(f"/messages/{msg_id}/feedbacks", method="POST"): + result = _feedback_post(MessageFeedbackApi(), payload, _chat_app(), _end_user(), msg_id) + + assert result == {"result": "success"} + mock_create.assert_called_once() + + @patch("controllers.web.message.MessageService.create_feedback") + def test_feedback_null_rating(self, mock_create: MagicMock, app: Flask) -> None: + payload = MessageFeedbackPayload.model_validate({"rating": None}) + msg_id = uuid4() + + with app.test_request_context(f"/messages/{msg_id}/feedbacks", method="POST"): + result = _feedback_post(MessageFeedbackApi(), payload, _chat_app(), _end_user(), msg_id) + + assert result == {"result": "success"} + + @patch( + "controllers.web.message.MessageService.create_feedback", + side_effect=MessageNotExistsError(), + ) + def test_feedback_message_not_found(self, mock_create: MagicMock, app: Flask) -> None: + payload = MessageFeedbackPayload.model_validate({"rating": "dislike"}) + msg_id = uuid4() + + with app.test_request_context(f"/messages/{msg_id}/feedbacks", method="POST"): + with pytest.raises(NotFound, match="Message Not Exists"): + _feedback_post(MessageFeedbackApi(), payload, _chat_app(), _end_user(), msg_id) + + # --------------------------------------------------------------------------- + # MessageMoreLikeThisApi + # --------------------------------------------------------------------------- + class TestMessageMoreLikeThisApi: + def test_wrong_mode_raises(self, app: Flask) -> None: + msg_id = uuid4() + query = MessageMoreLikeThisQuery.model_validate({"response_mode": "blocking"}) + session = MagicMock() + with app.test_request_context(f"/messages/{msg_id}/more-like-this?response_mode=blocking"): + with pytest.raises(NotCompletionAppError): + _more_like_this_get(MessageMoreLikeThisApi(), query, session, _chat_app(), _end_user(), msg_id) + + @patch("controllers.web.message.helper.compact_generate_response", return_value={"answer": "similar"}) + @patch("controllers.web.message.AppGenerateService.generate_more_like_this") + def test_happy_path(self, mock_gen: MagicMock, mock_compact: MagicMock, app: Flask) -> None: + msg_id = uuid4() + mock_gen.return_value = "response" + query = MessageMoreLikeThisQuery.model_validate({"response_mode": "blocking"}) + session = MagicMock() + + with app.test_request_context(f"/messages/{msg_id}/more-like-this?response_mode=blocking"): + result = _more_like_this_get( + MessageMoreLikeThisApi(), query, session, _completion_app(), _end_user(), msg_id + ) + + assert result == {"answer": "similar"} + + @patch( + "controllers.web.message.AppGenerateService.generate_more_like_this", + side_effect=MessageNotExistsError(), + ) + def test_message_not_found(self, mock_gen: MagicMock, app: Flask) -> None: + msg_id = uuid4() + query = MessageMoreLikeThisQuery.model_validate({"response_mode": "blocking"}) + session = MagicMock() + with app.test_request_context(f"/messages/{msg_id}/more-like-this?response_mode=blocking"): + with pytest.raises(NotFound, match="Message Not Exists"): + _more_like_this_get( + MessageMoreLikeThisApi(), query, session, _completion_app(), _end_user(), msg_id + ) + + @patch( + "controllers.web.message.AppGenerateService.generate_more_like_this", + side_effect=MoreLikeThisDisabledError(), + ) + def test_feature_disabled(self, mock_gen: MagicMock, app: Flask) -> None: + msg_id = uuid4() + query = MessageMoreLikeThisQuery.model_validate({"response_mode": "blocking"}) + session = MagicMock() + with app.test_request_context(f"/messages/{msg_id}/more-like-this?response_mode=blocking"): + with pytest.raises(AppMoreLikeThisDisabledError): + _more_like_this_get( + MessageMoreLikeThisApi(), query, session, _completion_app(), _end_user(), msg_id + ) + + # --------------------------------------------------------------------------- + # MessageSuggestedQuestionApi + # --------------------------------------------------------------------------- + class TestMessageSuggestedQuestionApi: + def test_wrong_mode_raises(self, app: Flask) -> None: + msg_id = uuid4() + with app.test_request_context(f"/messages/{msg_id}/suggested-questions"): + with pytest.raises(NotChatAppError): + MessageSuggestedQuestionApi().get(_completion_app(), _end_user(), msg_id) + + @patch("controllers.web.message.MessageService.get_suggested_questions_after_answer") + def test_happy_path(self, mock_suggest: MagicMock, app: Flask) -> None: + msg_id = uuid4() + mock_suggest.return_value = ["What about X?", "Tell me more about Y."] + + with app.test_request_context(f"/messages/{msg_id}/suggested-questions"): + result = MessageSuggestedQuestionApi().get(_chat_app(), _end_user(), msg_id) + + assert result["data"] == ["What about X?", "Tell me more about Y."] + + @patch( + "controllers.web.message.MessageService.get_suggested_questions_after_answer", + side_effect=MessageNotExistsError(), + ) + def test_message_not_found(self, mock_suggest: MagicMock, app: Flask) -> None: + msg_id = uuid4() + with app.test_request_context(f"/messages/{msg_id}/suggested-questions"): + with pytest.raises(NotFound, match="Message not found"): + MessageSuggestedQuestionApi().get(_chat_app(), _end_user(), msg_id) + + +"""Unit tests for controllers.web.message message list mapping.""" + + +# Ensure flask_restx.api finds MethodView during import. +if not hasattr(builtins, "MethodView"): + builtins.MethodView = MethodView # type: ignore[attr-defined] + + +def _load_controller_module(): + """Import controllers.web.message using a stub package.""" + + import importlib + import importlib.util + import sys + + parent_module_name = "controllers.web" + module_name = f"{parent_module_name}.message" + + if parent_module_name not in sys.modules: + from flask_restx import Namespace + + stub = ModuleType(parent_module_name) + web_controller_dir = Path(__file__).resolve().parents[4] / "controllers" / "web" + stub.__file__ = str(web_controller_dir / "__init__.py") + stub.__path__ = [str(web_controller_dir)] + stub.__package__ = "controllers" + stub.__spec__ = importlib.util.spec_from_loader(parent_module_name, loader=None, is_package=True) + stub.web_ns = Namespace("web", description="Web API", path="/") + sys.modules[parent_module_name] = stub + + wraps_module_name = f"{parent_module_name}.wraps" + if wraps_module_name not in sys.modules: + wraps_stub = ModuleType(wraps_module_name) + + class WebApiResource: + pass + + wraps_stub.WebApiResource = WebApiResource + sys.modules[wraps_module_name] = wraps_stub + + return importlib.import_module(module_name) + + +message_module = _load_controller_module() +MessageListApi = message_module.MessageListApi + + +class TestMessageListCases: + @staticmethod + @pytest.fixture + def app() -> Flask: + app = Flask(__name__) + app.config["TESTING"] = True + return app + + @staticmethod + @pytest.mark.parametrize( + "sqlite_session", + [(App, Conversation, EndUser, Message, MessageAgentThought, MessageFeedback)], + indirect=True, + ) + def test_message_list_mapping(app: Flask, monkeypatch: pytest.MonkeyPatch, sqlite_session) -> None: + conversation_id = str(uuid4()) + message_id = str(uuid4()) + + created_at = datetime(2024, 1, 1, 12, 0, 0) + resource_created_at = datetime(2024, 1, 1, 13, 0, 0) + thought_created_at = datetime(2024, 1, 1, 14, 0, 0) + + retriever_resource = { + "id": "res-obj", + "message_id": message_id, + "position": 2, + "dataset_id": "ds-1", + "dataset_name": "dataset", + "document_id": "doc-1", + "document_name": "document", + "data_source_type": "file", + "segment_id": "seg-1", + "score": 0.9, + "hit_count": 1, + "word_count": 10, + "segment_position": 0, + "index_node_hash": "hash", + "content": "content", + "created_at": int(resource_created_at.timestamp()), + } + + agent_thought = MessageAgentThought( + message_chain_id="chain-1", + message_id=message_id, + position=1, + created_by_role="end_user", + created_by="end-user-1", + thought="thinking", + tool="tool", + tool_labels_str=json.dumps({"label": "value"}), + tool_input="{}", + observation="observed", + message_files=json.dumps(["file-a"]), + ) + agent_thought.id = "thought-1" + agent_thought.created_at = thought_created_at + + message_files = [ + {"id": "file-dict", "filename": "a.txt", "type": "file", "transfer_method": "local"}, + {"id": "file-obj", "filename": "b.txt", "type": "file", "transfer_method": "local"}, + ] + + app_model = App( + id="app-1", + tenant_id="tenant-1", + name="Chat App", + mode=AppMode.CHAT, + enable_site=False, + enable_api=False, + ) + end_user = EndUser( + id="end-user-1", + tenant_id="tenant-1", + app_id=app_model.id, + type=EndUserType.BROWSER, + name="Web User", + session_id="session-1", + ) + conversation = Conversation( + id=conversation_id, + app_id=app_model.id, + mode=AppMode.CHAT, + name="Conversation", + _inputs={}, + status="normal", + from_source=ConversationFromSource.API, + from_end_user_id=end_user.id, + ) + message = Message( + id=message_id, + app_id=app_model.id, + conversation_id=conversation_id, + parent_message_id=None, + _inputs={"foo": "bar"}, + query="hello", + message={}, + answer="answer", + message_unit_price=0, + message_price_unit=0, + answer_unit_price=0, + answer_price_unit=0, + provider_response_latency=0, + total_price=0, + currency="USD", + invoke_from=InvokeFrom.SERVICE_API, + from_source=ConversationFromSource.API, + from_end_user_id=end_user.id, + app_mode=AppMode.CHAT, + message_metadata=json.dumps( + { + "meta": "value", + "retriever_resources": [ + {"id": "res-dict", "message_id": message_id, "position": 1}, + retriever_resource, + ], + } + ), + status="normal", + error=None, + ) + message.created_at = created_at + message.set_extra_contents( + [ + HumanInputContent( + workflow_run_id=str(uuid.uuid4()), + submitted=True, + ).model_dump(mode="json") + ] + ) + feedback = MessageFeedback( + app_id=app_model.id, + conversation_id=conversation_id, + message_id=message_id, + rating=FeedbackRating.LIKE, + from_source=FeedbackFromSource.USER, + from_end_user_id=end_user.id, + ) + sqlite_session.add_all([app_model, end_user, conversation, message, feedback, agent_thought]) + sqlite_session.commit() + + pagination = SimpleNamespace(limit=20, has_more=False, data=[message]) + + def message_files_with_session(_self, *, session): + del session + return message_files + + monkeypatch.setattr(Message, "message_files_with_session", message_files_with_session) + + with ( + patch.object(message_module.MessageService, "pagination_by_first_id", return_value=pagination) as mock_page, + patch.object(message_module.db, "session", return_value=sqlite_session), + app.test_request_context(f"/messages?conversation_id={conversation_id}&limit=20"), + ): + query = MessageListQuery.model_validate({"conversation_id": conversation_id, "limit": 20}) + response = inspect.unwrap(MessageListApi.get)(MessageListApi(), query, app_model, end_user) + + mock_page.assert_called_once_with(app_model, end_user, conversation_id, None, 20, session=ANY) + assert response["limit"] == 20 + assert response["has_more"] is False + assert len(response["data"]) == 1 + + item = response["data"][0] + assert item["id"] == message_id + assert item["conversation_id"] == conversation_id + assert item["inputs"] == {"foo": "bar"} + assert item["answer"] == "answer" + assert item["feedback"]["rating"] == "like" + assert item["metadata"]["meta"] == "value" + assert item["created_at"] == int(created_at.timestamp()) + + assert item["retriever_resources"][0]["id"] == "res-dict" + assert item["retriever_resources"][1]["id"] == "res-obj" + assert item["retriever_resources"][1]["created_at"] == int(resource_created_at.timestamp()) + + assert item["agent_thoughts"][0]["chain_id"] == "chain-1" + assert item["agent_thoughts"][0]["created_at"] == int(thought_created_at.timestamp()) + assert item["extra_contents"][0]["workflow_run_id"] == message.extra_contents[0]["workflow_run_id"] + assert item["extra_contents"][0]["submitted"] == message.extra_contents[0]["submitted"] + + assert item["message_files"][0]["id"] == "file-dict" + assert item["message_files"][1]["id"] == "file-obj" diff --git a/api/tests/unit_tests/controllers/web/test_message_endpoints.py b/api/tests/unit_tests/controllers/web/test_message_endpoints.py deleted file mode 100644 index 9fc9ad2fedc..00000000000 --- a/api/tests/unit_tests/controllers/web/test_message_endpoints.py +++ /dev/null @@ -1,168 +0,0 @@ -"""Unit tests for controllers.web.message — feedback, more-like-this, suggested questions.""" - -from __future__ import annotations - -import inspect -from unittest.mock import MagicMock, patch -from uuid import uuid4 - -import pytest -from flask import Flask -from werkzeug.exceptions import NotFound - -from controllers.common.controller_schemas import MessageFeedbackPayload -from controllers.web.error import ( - AppMoreLikeThisDisabledError, - NotChatAppError, - NotCompletionAppError, -) -from controllers.web.message import ( - MessageFeedbackApi, - MessageMoreLikeThisApi, - MessageMoreLikeThisQuery, - MessageSuggestedQuestionApi, -) -from models.model import App, AppMode, EndUser -from services.errors.app import MoreLikeThisDisabledError -from services.errors.message import MessageNotExistsError -from tests.unit_tests.model_factories import make_end_user - - -def _chat_app() -> App: - return App(id="app-1", tenant_id="tenant-1", mode=AppMode.CHAT) - - -def _completion_app() -> App: - return App(id="app-1", tenant_id="tenant-1", mode=AppMode.COMPLETION) - - -def _end_user() -> EndUser: - return make_end_user(end_user_id="eu-1") - - -# The @model_validate and @with_session decorators wrap the handlers; tests -# call the undecorated function so they can pass a pydantic payload / a fake -# session directly. -_feedback_post = inspect.unwrap(MessageFeedbackApi.post) -_more_like_this_get = inspect.unwrap(MessageMoreLikeThisApi.get) - - -# --------------------------------------------------------------------------- -# MessageFeedbackApi -# --------------------------------------------------------------------------- -class TestMessageFeedbackApi: - @patch("controllers.web.message.MessageService.create_feedback") - def test_feedback_success(self, mock_create: MagicMock, app: Flask) -> None: - payload = MessageFeedbackPayload.model_validate({"rating": "like", "content": "great"}) - msg_id = uuid4() - - with app.test_request_context(f"/messages/{msg_id}/feedbacks", method="POST"): - result = _feedback_post(MessageFeedbackApi(), payload, _chat_app(), _end_user(), msg_id) - - assert result == {"result": "success"} - mock_create.assert_called_once() - - @patch("controllers.web.message.MessageService.create_feedback") - def test_feedback_null_rating(self, mock_create: MagicMock, app: Flask) -> None: - payload = MessageFeedbackPayload.model_validate({"rating": None}) - msg_id = uuid4() - - with app.test_request_context(f"/messages/{msg_id}/feedbacks", method="POST"): - result = _feedback_post(MessageFeedbackApi(), payload, _chat_app(), _end_user(), msg_id) - - assert result == {"result": "success"} - - @patch( - "controllers.web.message.MessageService.create_feedback", - side_effect=MessageNotExistsError(), - ) - def test_feedback_message_not_found(self, mock_create: MagicMock, app: Flask) -> None: - payload = MessageFeedbackPayload.model_validate({"rating": "dislike"}) - msg_id = uuid4() - - with app.test_request_context(f"/messages/{msg_id}/feedbacks", method="POST"): - with pytest.raises(NotFound, match="Message Not Exists"): - _feedback_post(MessageFeedbackApi(), payload, _chat_app(), _end_user(), msg_id) - - -# --------------------------------------------------------------------------- -# MessageMoreLikeThisApi -# --------------------------------------------------------------------------- -class TestMessageMoreLikeThisApi: - def test_wrong_mode_raises(self, app: Flask) -> None: - msg_id = uuid4() - query = MessageMoreLikeThisQuery.model_validate({"response_mode": "blocking"}) - session = MagicMock() - with app.test_request_context(f"/messages/{msg_id}/more-like-this?response_mode=blocking"): - with pytest.raises(NotCompletionAppError): - _more_like_this_get(MessageMoreLikeThisApi(), query, session, _chat_app(), _end_user(), msg_id) - - @patch("controllers.web.message.helper.compact_generate_response", return_value={"answer": "similar"}) - @patch("controllers.web.message.AppGenerateService.generate_more_like_this") - def test_happy_path(self, mock_gen: MagicMock, mock_compact: MagicMock, app: Flask) -> None: - msg_id = uuid4() - mock_gen.return_value = "response" - query = MessageMoreLikeThisQuery.model_validate({"response_mode": "blocking"}) - session = MagicMock() - - with app.test_request_context(f"/messages/{msg_id}/more-like-this?response_mode=blocking"): - result = _more_like_this_get( - MessageMoreLikeThisApi(), query, session, _completion_app(), _end_user(), msg_id - ) - - assert result == {"answer": "similar"} - - @patch( - "controllers.web.message.AppGenerateService.generate_more_like_this", - side_effect=MessageNotExistsError(), - ) - def test_message_not_found(self, mock_gen: MagicMock, app: Flask) -> None: - msg_id = uuid4() - query = MessageMoreLikeThisQuery.model_validate({"response_mode": "blocking"}) - session = MagicMock() - with app.test_request_context(f"/messages/{msg_id}/more-like-this?response_mode=blocking"): - with pytest.raises(NotFound, match="Message Not Exists"): - _more_like_this_get(MessageMoreLikeThisApi(), query, session, _completion_app(), _end_user(), msg_id) - - @patch( - "controllers.web.message.AppGenerateService.generate_more_like_this", - side_effect=MoreLikeThisDisabledError(), - ) - def test_feature_disabled(self, mock_gen: MagicMock, app: Flask) -> None: - msg_id = uuid4() - query = MessageMoreLikeThisQuery.model_validate({"response_mode": "blocking"}) - session = MagicMock() - with app.test_request_context(f"/messages/{msg_id}/more-like-this?response_mode=blocking"): - with pytest.raises(AppMoreLikeThisDisabledError): - _more_like_this_get(MessageMoreLikeThisApi(), query, session, _completion_app(), _end_user(), msg_id) - - -# --------------------------------------------------------------------------- -# MessageSuggestedQuestionApi -# --------------------------------------------------------------------------- -class TestMessageSuggestedQuestionApi: - def test_wrong_mode_raises(self, app: Flask) -> None: - msg_id = uuid4() - with app.test_request_context(f"/messages/{msg_id}/suggested-questions"): - with pytest.raises(NotChatAppError): - MessageSuggestedQuestionApi().get(_completion_app(), _end_user(), msg_id) - - @patch("controllers.web.message.MessageService.get_suggested_questions_after_answer") - def test_happy_path(self, mock_suggest: MagicMock, app: Flask) -> None: - msg_id = uuid4() - mock_suggest.return_value = ["What about X?", "Tell me more about Y."] - - with app.test_request_context(f"/messages/{msg_id}/suggested-questions"): - result = MessageSuggestedQuestionApi().get(_chat_app(), _end_user(), msg_id) - - assert result["data"] == ["What about X?", "Tell me more about Y."] - - @patch( - "controllers.web.message.MessageService.get_suggested_questions_after_answer", - side_effect=MessageNotExistsError(), - ) - def test_message_not_found(self, mock_suggest: MagicMock, app: Flask) -> None: - msg_id = uuid4() - with app.test_request_context(f"/messages/{msg_id}/suggested-questions"): - with pytest.raises(NotFound, match="Message not found"): - MessageSuggestedQuestionApi().get(_chat_app(), _end_user(), msg_id) diff --git a/api/tests/unit_tests/controllers/web/test_message_list.py b/api/tests/unit_tests/controllers/web/test_message_list.py deleted file mode 100644 index 1b5d79ef324..00000000000 --- a/api/tests/unit_tests/controllers/web/test_message_list.py +++ /dev/null @@ -1,255 +0,0 @@ -"""Unit tests for controllers.web.message message list mapping.""" - -from __future__ import annotations - -import builtins -import inspect -import json -import uuid -from datetime import datetime -from pathlib import Path -from types import ModuleType, SimpleNamespace -from unittest.mock import ANY, patch -from uuid import uuid4 - -import pytest -from flask import Flask -from flask.views import MethodView - -from controllers.common.controller_schemas import MessageListQuery -from core.app.entities.app_invoke_entities import InvokeFrom -from core.entities.execution_extra_content import HumanInputContent -from models.enums import ConversationFromSource, EndUserType, FeedbackFromSource, FeedbackRating -from models.model import ( - App, - AppMode, - Conversation, - EndUser, - Message, - MessageAgentThought, - MessageFeedback, -) - -# Ensure flask_restx.api finds MethodView during import. -if not hasattr(builtins, "MethodView"): - builtins.MethodView = MethodView # type: ignore[attr-defined] - - -def _load_controller_module(): - """Import controllers.web.message using a stub package.""" - - import importlib - import importlib.util - import sys - - parent_module_name = "controllers.web" - module_name = f"{parent_module_name}.message" - - if parent_module_name not in sys.modules: - from flask_restx import Namespace - - stub = ModuleType(parent_module_name) - web_controller_dir = Path(__file__).resolve().parents[4] / "controllers" / "web" - stub.__file__ = str(web_controller_dir / "__init__.py") - stub.__path__ = [str(web_controller_dir)] - stub.__package__ = "controllers" - stub.__spec__ = importlib.util.spec_from_loader(parent_module_name, loader=None, is_package=True) - stub.web_ns = Namespace("web", description="Web API", path="/") - sys.modules[parent_module_name] = stub - - wraps_module_name = f"{parent_module_name}.wraps" - if wraps_module_name not in sys.modules: - wraps_stub = ModuleType(wraps_module_name) - - class WebApiResource: - pass - - wraps_stub.WebApiResource = WebApiResource - sys.modules[wraps_module_name] = wraps_stub - - return importlib.import_module(module_name) - - -message_module = _load_controller_module() -MessageListApi = message_module.MessageListApi - - -@pytest.fixture -def app() -> Flask: - app = Flask(__name__) - app.config["TESTING"] = True - return app - - -@pytest.mark.parametrize( - "sqlite_session", - [(App, Conversation, EndUser, Message, MessageAgentThought, MessageFeedback)], - indirect=True, -) -def test_message_list_mapping(app: Flask, monkeypatch: pytest.MonkeyPatch, sqlite_session) -> None: - conversation_id = str(uuid4()) - message_id = str(uuid4()) - - created_at = datetime(2024, 1, 1, 12, 0, 0) - resource_created_at = datetime(2024, 1, 1, 13, 0, 0) - thought_created_at = datetime(2024, 1, 1, 14, 0, 0) - - retriever_resource = { - "id": "res-obj", - "message_id": message_id, - "position": 2, - "dataset_id": "ds-1", - "dataset_name": "dataset", - "document_id": "doc-1", - "document_name": "document", - "data_source_type": "file", - "segment_id": "seg-1", - "score": 0.9, - "hit_count": 1, - "word_count": 10, - "segment_position": 0, - "index_node_hash": "hash", - "content": "content", - "created_at": int(resource_created_at.timestamp()), - } - - agent_thought = MessageAgentThought( - message_chain_id="chain-1", - message_id=message_id, - position=1, - created_by_role="end_user", - created_by="end-user-1", - thought="thinking", - tool="tool", - tool_labels_str=json.dumps({"label": "value"}), - tool_input="{}", - observation="observed", - message_files=json.dumps(["file-a"]), - ) - agent_thought.id = "thought-1" - agent_thought.created_at = thought_created_at - - message_files = [ - {"id": "file-dict", "filename": "a.txt", "type": "file", "transfer_method": "local"}, - {"id": "file-obj", "filename": "b.txt", "type": "file", "transfer_method": "local"}, - ] - - app_model = App( - id="app-1", - tenant_id="tenant-1", - name="Chat App", - mode=AppMode.CHAT, - enable_site=False, - enable_api=False, - ) - end_user = EndUser( - id="end-user-1", - tenant_id="tenant-1", - app_id=app_model.id, - type=EndUserType.BROWSER, - name="Web User", - session_id="session-1", - ) - conversation = Conversation( - id=conversation_id, - app_id=app_model.id, - mode=AppMode.CHAT, - name="Conversation", - _inputs={}, - status="normal", - from_source=ConversationFromSource.API, - from_end_user_id=end_user.id, - ) - message = Message( - id=message_id, - app_id=app_model.id, - conversation_id=conversation_id, - parent_message_id=None, - _inputs={"foo": "bar"}, - query="hello", - message={}, - answer="answer", - message_unit_price=0, - message_price_unit=0, - answer_unit_price=0, - answer_price_unit=0, - provider_response_latency=0, - total_price=0, - currency="USD", - invoke_from=InvokeFrom.SERVICE_API, - from_source=ConversationFromSource.API, - from_end_user_id=end_user.id, - app_mode=AppMode.CHAT, - message_metadata=json.dumps( - { - "meta": "value", - "retriever_resources": [ - {"id": "res-dict", "message_id": message_id, "position": 1}, - retriever_resource, - ], - } - ), - status="normal", - error=None, - ) - message.created_at = created_at - message.set_extra_contents( - [ - HumanInputContent( - workflow_run_id=str(uuid.uuid4()), - submitted=True, - ).model_dump(mode="json") - ] - ) - feedback = MessageFeedback( - app_id=app_model.id, - conversation_id=conversation_id, - message_id=message_id, - rating=FeedbackRating.LIKE, - from_source=FeedbackFromSource.USER, - from_end_user_id=end_user.id, - ) - sqlite_session.add_all([app_model, end_user, conversation, message, feedback, agent_thought]) - sqlite_session.commit() - - pagination = SimpleNamespace(limit=20, has_more=False, data=[message]) - - def message_files_with_session(_self, *, session): - del session - return message_files - - monkeypatch.setattr(Message, "message_files_with_session", message_files_with_session) - - with ( - patch.object(message_module.MessageService, "pagination_by_first_id", return_value=pagination) as mock_page, - patch.object(message_module.db, "session", return_value=sqlite_session), - app.test_request_context(f"/messages?conversation_id={conversation_id}&limit=20"), - ): - query = MessageListQuery.model_validate({"conversation_id": conversation_id, "limit": 20}) - response = inspect.unwrap(MessageListApi.get)(MessageListApi(), query, app_model, end_user) - - mock_page.assert_called_once_with(app_model, end_user, conversation_id, None, 20, session=ANY) - assert response["limit"] == 20 - assert response["has_more"] is False - assert len(response["data"]) == 1 - - item = response["data"][0] - assert item["id"] == message_id - assert item["conversation_id"] == conversation_id - assert item["inputs"] == {"foo": "bar"} - assert item["answer"] == "answer" - assert item["feedback"]["rating"] == "like" - assert item["metadata"]["meta"] == "value" - assert item["created_at"] == int(created_at.timestamp()) - - assert item["retriever_resources"][0]["id"] == "res-dict" - assert item["retriever_resources"][1]["id"] == "res-obj" - assert item["retriever_resources"][1]["created_at"] == int(resource_created_at.timestamp()) - - assert item["agent_thoughts"][0]["chain_id"] == "chain-1" - assert item["agent_thoughts"][0]["created_at"] == int(thought_created_at.timestamp()) - assert item["extra_contents"][0]["workflow_run_id"] == message.extra_contents[0]["workflow_run_id"] - assert item["extra_contents"][0]["submitted"] == message.extra_contents[0]["submitted"] - - assert item["message_files"][0]["id"] == "file-dict" - assert item["message_files"][1]["id"] == "file-obj" diff --git a/api/tests/unit_tests/controllers/web/test_web_passport.py b/api/tests/unit_tests/controllers/web/test_passport.py similarity index 100% rename from api/tests/unit_tests/controllers/web/test_web_passport.py rename to api/tests/unit_tests/controllers/web/test_passport.py diff --git a/api/tests/unit_tests/controllers/web/test_public_app_admission.py b/api/tests/unit_tests/controllers/web/test_public_app_admission.py deleted file mode 100644 index cf2aa7e3fbd..00000000000 --- a/api/tests/unit_tests/controllers/web/test_public_app_admission.py +++ /dev/null @@ -1,358 +0,0 @@ -"""Public bootstrap admission with real SQL repositories and the canonical hook.""" - -from dataclasses import dataclass, replace -from datetime import UTC, datetime -from unittest.mock import MagicMock -from uuid import uuid4 - -import pytest -from flask import Flask -from flask.testing import FlaskClient -from pytest_mock import MockerFixture -from sqlalchemy import func, select -from sqlalchemy.exc import OperationalError -from sqlalchemy.orm import Session, sessionmaker - -from controllers.web import app as app_controller -from controllers.web import bp, login, passport, site, wraps -from enums import DeploymentEdition -from extensions.ext_application_services import ApplicationServices -from libs.passport import PassportService -from models.account import Tenant -from models.enums import EndUserType -from models.model import App, AppMode, CustomizeTokenStrategy, EndUser, Message, Site -from models.workflow import Workflow, WorkflowRun, WorkflowType -from repositories.app_definition_query_repository import AppDefinitionQueryRepository -from repositories.web_passport_repository import WebPassportRepository -from repositories.webapp_access_query_repository import WebAppAccessQueryRepository -from services.app_definition_query_service import AppDefinitionQueryService -from services.entities.feature_entities import FeatureModel -from services.file_service import FileService -from services.web_app_runtime_query_service import WebAppRuntimeQueryService -from services.web_authentication_adapters import PassportWebAppSessionGateway -from services.web_passport_gateways import PassportTokenGateway -from services.web_passport_service import WebPassportService -from tests.unit_tests.config_override import apply_config_overrides - - -@dataclass -class AdmissionHarness: - app: App - old_user: EndUser - client: FlaskClient - services: ApplicationServices - issue: MagicMock - auth_calls: tuple[MagicMock, ...] - mode_lookup: MagicMock - login_verify: MagicMock - session_verify: MagicMock - metadata_calls: tuple[MagicMock, ...] - access_repo: WebAppAccessQueryRepository - - -def _harness( - session: Session, - factory: sessionmaker[Session], - monkeypatch: pytest.MonkeyPatch, - mocker: MockerFixture, - base_services: ApplicationServices, - *, - published: bool, - missing: bool = False, -) -> AdmissionHarness: - app = App( - id=str(uuid4()), - tenant_id=str(uuid4()), - mode=AppMode.ADVANCED_CHAT, - name="Must not leak", - enable_site=True, - enable_api=True, - ) - old_user = EndUser( - id=str(uuid4()), - tenant_id=app.tenant_id, - app_id=app.id, - session_id="old", - _is_anonymous=True, - type=EndUserType.BROWSER, - ) - if not missing: - tenant = Tenant(name="Private workspace") - tenant.id = app.tenant_id - session.add(tenant) - workflow = Workflow( - id=str(uuid4()), - tenant_id=app.tenant_id, - app_id=app.id, - type=WorkflowType.CHAT, - version="published" if published else "draft", - graph='{"nodes":[]}', - features="{}", - created_by=str(uuid4()), - environment_variables=[], - conversation_variables=[], - rag_pipeline_variables=[], - ) - if published: - app.workflow_id = workflow.id - session.add_all( - [ - app, - old_user, - workflow, - Site( - app_id=app.id, - code="fixture-code", - title="Must not leak", - default_language="en-US", - customize_token_strategy=CustomizeTokenStrategy.UUID, - ), - ] - ) - session.commit() - - access_repo = WebAppAccessQueryRepository(session_factory=factory) - passport_repo = WebPassportRepository(session_factory=factory, generate_session_id=lambda: "new-session") - access = base_services.webapp_access - monkeypatch.setattr(access, "_access", access_repo) - monkeypatch.setattr(access, "_webapp_auth_enabled", False) - mode_lookup = mocker.spy(access._policy, "get_access_mode") - auth = base_services.web_passport._auth - auth_calls = ( - mocker.spy(auth, "is_webapp_auth_enabled"), - mocker.spy(auth, "get_app_auth_type"), - ) - tokens = PassportTokenGateway(passport=PassportService()) - issue = mocker.patch.object(PassportService, "issue", return_value="new-passport") - mocker.patch.object( - PassportService, - "verify", - return_value={"app_id": app.id, "app_code": "fixture-code", "end_user_id": old_user.id}, - ) - app_sessions = PassportWebAppSessionGateway(sessions=access_repo, app_access=access) - session_verify = mocker.spy(app_sessions, "verify") - authentication = base_services.web_authentication - monkeypatch.setattr(authentication, "_app_access", access) - monkeypatch.setattr(authentication, "_app_sessions", app_sessions) - monkeypatch.setattr(authentication, "_private_app_access_enabled", False) - login_verify = mocker.spy(authentication._tokens, "verify_access_token") - definitions = AppDefinitionQueryRepository(session_factory=factory) - services = replace( - base_services, - webapp_access=access, - web_authentication=authentication, - web_passport=WebPassportService( - passports=passport_repo, - auth=auth, - tokens=tokens, - now=lambda: datetime.now(UTC), - access_token_expire_minutes=60, - ), - app_definitions=AppDefinitionQueryService(definitions=definitions, builtin_icon_url_prefix=""), - web_app_runtime=WebAppRuntimeQueryService( - runtime=definitions, - file_service=FileService(session_factory=factory), - workspace_features=lambda _: FeatureModel(), - files_url="https://files.example", - deployment_edition=DeploymentEdition.CLOUD, - ), - ) - metadata_calls = tuple( - mocker.spy(services.app_definitions, method) - for method in ( - "get_mode", - "get_parameters", - "get_public_parameters", - "get_tool_icons", - "get_summary", - "get_site_configuration", - ) - ) + (mocker.spy(services.web_app_runtime, "get_bootstrap"),) - for module in (app_controller, login, passport, site, wraps): - monkeypatch.setattr(module, "application_services", lambda: services) - monkeypatch.setattr(wraps, "extract_webapp_passport", lambda *_: "old-passport") - monkeypatch.setattr(login, "extract_webapp_passport", lambda *_: "old-passport") - monkeypatch.setattr(wraps.SystemFeatureService, "is_webapp_auth_enabled", lambda: False) - apply_config_overrides( - monkeypatch, - DEPLOYMENT_EDITION=DeploymentEdition.CLOUD, - NETWORK_ACCESS_TRUSTED_PROXY_CIDRS="172.18.0.0/16", - ) - flask = Flask(__name__) - flask.config.update(TESTING=True, RESTX_ERROR_404_HELP=False) - flask.register_blueprint(bp) - return AdmissionHarness( - app=app, - old_user=old_user, - client=flask.test_client(), - services=services, - issue=issue, - auth_calls=auth_calls, - mode_lookup=mode_lookup, - login_verify=login_verify, - session_verify=session_verify, - metadata_calls=metadata_calls, - access_repo=access_repo, - ) - - -ROUTES = ( - "/webapp/access-mode?appCode=fixture-code", - "/login/status?app_code=fixture-code", - "/passport?user_id=new-session", - "/site", - "/meta", - "/parameters", -) - - -@pytest.mark.parametrize("route", ROUTES) -@pytest.mark.parametrize("missing", [False, True], ids=["unpublished", "missing"]) -def test_all_bootstrap_routes_share_404_before_user_or_metadata_side_effects( - sqlite_session: Session, - sqlite_session_factory: sessionmaker[Session], - monkeypatch: pytest.MonkeyPatch, - mocker: MockerFixture, - account_application_services: ApplicationServices, - route: str, - missing: bool, -) -> None: - h = _harness( - sqlite_session, - sqlite_session_factory, - monkeypatch, - mocker, - account_application_services, - published=False, - missing=missing, - ) - response = h.client.get( - "/api" + route, - headers={"X-App-Code": "fixture-code"}, - environ_overrides={"REMOTE_ADDR": "172.18.0.2", "HTTP_X_FORWARDED_FOR": "203.0.113.42"}, - ) - assert response.status_code == 404 - assert ( - response.data == b'{"client_ip":"203.0.113.42","code":"app_not_found","message":"App not found.","status":404}' - ) - assert response.headers["Content-Type"] == "application/json" - assert response.headers["Cache-Control"] == "no-store" - assert all(call.call_count == 0 for call in h.metadata_calls) - assert all(call.call_count == 0 for call in h.auth_calls) - h.mode_lookup.assert_not_called() - h.issue.assert_not_called() - h.login_verify.assert_not_called() - h.session_verify.assert_not_called() - assert sqlite_session.scalar(select(func.count()).select_from(EndUser)) == (0 if missing else 1) - assert sqlite_session.scalar(select(func.count()).select_from(WorkflowRun)) == 0 - assert sqlite_session.scalar(select(func.count()).select_from(Message)) == 0 - - -@pytest.mark.parametrize("route", ROUTES) -def test_published_bootstrap_remains_successful_with_normal_passport_behavior( - sqlite_session: Session, - sqlite_session_factory: sessionmaker[Session], - monkeypatch: pytest.MonkeyPatch, - route: str, - mocker: MockerFixture, - account_application_services: ApplicationServices, -) -> None: - h = _harness( - sqlite_session, sqlite_session_factory, monkeypatch, mocker, account_application_services, published=True - ) - response = h.client.get("/api" + route, headers={"X-App-Code": "fixture-code"}) - assert response.status_code == 200 - if route.startswith("/passport"): - assert response.get_json() == {"access_token": "new-passport"} - h.issue.assert_called_once() - assert sqlite_session.scalar(select(func.count()).select_from(EndUser)) == 2 - else: - h.issue.assert_not_called() - assert sqlite_session.scalar(select(func.count()).select_from(EndUser)) == 1 - if route.startswith("/webapp"): - assert response.get_json() == {"accessMode": "public"} - elif route.startswith("/login"): - assert response.get_json() == {"logged_in": True, "app_logged_in": True} - h.session_verify.assert_called_once_with(token="old-passport", app_code="fixture-code", user_id=None) - h.login_verify.assert_not_called() - elif route == "/site": - assert response.get_json()["site"]["title"] == "Must not leak" - elif route == "/meta": - assert response.get_json() == {"tool_icons": {}} - else: - assert "user_input_form" in response.get_json() - - -@pytest.mark.parametrize("route", ["/site", "/meta", "/parameters"]) -def test_old_passport_cannot_read_metadata_after_published_version_is_removed( - sqlite_session: Session, - sqlite_session_factory: sessionmaker[Session], - monkeypatch: pytest.MonkeyPatch, - route: str, - mocker: MockerFixture, - account_application_services: ApplicationServices, -) -> None: - h = _harness( - sqlite_session, sqlite_session_factory, monkeypatch, mocker, account_application_services, published=True - ) - assert h.access_repo.is_app_available(h.app.id) - h.app.workflow_id = None - sqlite_session.commit() - response = h.client.get( - "/api" + route, headers={"X-App-Code": "fixture-code"}, environ_overrides={"REMOTE_ADDR": "203.0.113.42"} - ) - assert response.status_code == 404 - assert response.get_json()["code"] == "app_not_found" - assert response.get_json()["client_ip"] == "203.0.113.42" - assert all(call.call_count == 0 for call in h.metadata_calls) - - -@pytest.mark.parametrize("changed_owner", ["site-app", "end-user-app", "end-user-tenant"]) -def test_old_passport_cannot_cross_persisted_app_or_tenant_owner_chain( - sqlite_session: Session, - sqlite_session_factory: sessionmaker[Session], - monkeypatch: pytest.MonkeyPatch, - mocker: MockerFixture, - account_application_services: ApplicationServices, - changed_owner: str, -) -> None: - h = _harness( - sqlite_session, sqlite_session_factory, monkeypatch, mocker, account_application_services, published=True - ) - if changed_owner == "site-app": - persisted_site = sqlite_session.scalar(select(Site).where(Site.app_id == h.app.id)) - assert persisted_site is not None - persisted_site.app_id = str(uuid4()) - elif changed_owner == "end-user-app": - h.old_user.app_id = str(uuid4()) - else: - h.old_user.tenant_id = str(uuid4()) - sqlite_session.commit() - response = h.client.get("/api/site", headers={"X-App-Code": "fixture-code"}) - assert response.status_code == 404 - assert response.get_json()["code"] == "app_not_found" - assert all(call.call_count == 0 for call in h.metadata_calls) - - -@pytest.mark.parametrize("route", ROUTES) -def test_real_database_failure_is_not_reclassified_as_unpublished( - sqlite_session: Session, - sqlite_session_factory: sessionmaker[Session], - monkeypatch: pytest.MonkeyPatch, - route: str, - mocker: MockerFixture, - account_application_services: ApplicationServices, -) -> None: - h = _harness( - sqlite_session, sqlite_session_factory, monkeypatch, mocker, account_application_services, published=False - ) - failure = OperationalError("select", {}, RuntimeError("database unavailable")) - mocker.patch.object(Session, "scalar", side_effect=failure) - mocker.patch.object(Session, "execute", side_effect=failure) - response = h.client.get("/api" + route, headers={"X-App-Code": "fixture-code"}) - assert response.status_code == (503 if route.startswith(("/webapp", "/login")) else 500) - assert response.get_json()["code"] != "app_not_found" - assert "client_ip" not in response.get_json() - h.issue.assert_not_called() - h.login_verify.assert_not_called() - h.session_verify.assert_not_called() diff --git a/api/tests/unit_tests/controllers/web/test_remote_files.py b/api/tests/unit_tests/controllers/web/test_remote_files.py index 9454984a923..9503834b3c7 100644 --- a/api/tests/unit_tests/controllers/web/test_remote_files.py +++ b/api/tests/unit_tests/controllers/web/test_remote_files.py @@ -40,7 +40,7 @@ from services.remote_file_service import RemoteFileInvalidUrlError as RemoteFile from services.remote_file_service import RemoteFileNotFoundError as RemoteFileNotFoundServiceError from services.remote_file_service import RemoteFileUnavailableError as RemoteFileUnavailableServiceError from services.remote_file_service import RemoteFileUrlBlockedError as RemoteFileUrlBlockedServiceError -from tests.unit_tests.model_factories import make_end_user +from tests.helpers.unit.model_factories import make_end_user REMOTE_FILE_ERROR_CASES: tuple[tuple[type[RemoteFileError], type[BaseHTTPException], str, int], ...] = ( (RemoteFileInvalidUrlServiceError, RemoteFileInvalidUrlError, "remote_file_invalid_url", 400), diff --git a/api/tests/unit_tests/controllers/web/test_saved_message.py b/api/tests/unit_tests/controllers/web/test_saved_message.py index 6e2f3def794..57dc6c4a876 100644 --- a/api/tests/unit_tests/controllers/web/test_saved_message.py +++ b/api/tests/unit_tests/controllers/web/test_saved_message.py @@ -22,7 +22,7 @@ from services.saved_message_service import ( SavedMessagePage, SavedMessageRecord, ) -from tests.unit_tests.model_factories import make_end_user +from tests.helpers.unit.model_factories import make_end_user _INPUT_FILE_URL = "https://example.com/input.pdf" _CREATED_AT = datetime(2026, 1, 2, 3, 4, 5, tzinfo=UTC) diff --git a/api/tests/unit_tests/controllers/web/test_site.py b/api/tests/unit_tests/controllers/web/test_site.py index 6e2800e9452..37629f1ee92 100644 --- a/api/tests/unit_tests/controllers/web/test_site.py +++ b/api/tests/unit_tests/controllers/web/test_site.py @@ -18,7 +18,7 @@ from services.web_app_runtime_query_service import ( WebAppBootstrap, WebAppRuntimeUnavailableError, ) -from tests.unit_tests.config_override import apply_config_overrides +from tests.helpers.unit.config_override import apply_config_overrides def _bootstrap() -> WebAppBootstrap: diff --git a/api/tests/unit_tests/controllers/web/test_workflow.py b/api/tests/unit_tests/controllers/web/test_workflow.py index 35459762f5c..79597949d33 100644 --- a/api/tests/unit_tests/controllers/web/test_workflow.py +++ b/api/tests/unit_tests/controllers/web/test_workflow.py @@ -19,7 +19,7 @@ from models.model import App, AppMode, EndUser from services.errors.app import ( TriggerWorkflowServiceModeUnavailableError as TriggerWorkflowServiceModeUnavailableServiceError, ) -from tests.unit_tests.model_factories import make_end_user +from tests.helpers.unit.model_factories import make_end_user def _workflow_app() -> App: diff --git a/api/tests/unit_tests/controllers/web/test_workflow_events.py b/api/tests/unit_tests/controllers/web/test_workflow_events.py index 343f11552a2..acfaf804f7e 100644 --- a/api/tests/unit_tests/controllers/web/test_workflow_events.py +++ b/api/tests/unit_tests/controllers/web/test_workflow_events.py @@ -14,7 +14,7 @@ from graphon.enums import WorkflowExecutionStatus from models.enums import CreatorUserRole, WorkflowRunTriggeredFrom from models.model import App, AppMode, EndUser from models.workflow import WorkflowRun, WorkflowType -from tests.unit_tests.model_factories import make_end_user +from tests.helpers.unit.model_factories import make_end_user def _workflow_app() -> App: diff --git a/api/tests/unit_tests/controllers/web/test_wraps.py b/api/tests/unit_tests/controllers/web/test_wraps.py index 353b9514d42..c93753a2b67 100644 --- a/api/tests/unit_tests/controllers/web/test_wraps.py +++ b/api/tests/unit_tests/controllers/web/test_wraps.py @@ -1,236 +1,596 @@ +from dataclasses import dataclass, replace +from datetime import UTC, datetime from typing import Never from unittest import mock +from unittest.mock import MagicMock from uuid import uuid4 import pytest from flask import Blueprint, Flask -from sqlalchemy.orm import Session +from flask.testing import FlaskClient +from pytest_mock import MockerFixture +from sqlalchemy import func, select +from sqlalchemy.exc import OperationalError +from sqlalchemy.orm import Session, sessionmaker from werkzeug.exceptions import NotFound, Unauthorized +from controllers.web import app as app_controller +from controllers.web import bp, login, passport, site, wraps from core.logging.context import clear_request_context, get_identity_context -from models import App, EndUser -from models.model import AppModelConfig -from tests.unit_tests.config_override import apply_config_overrides +from enums import DeploymentEdition +from extensions.ext_application_services import ApplicationServices +from libs.passport import PassportService +from models import App as _wraps_App +from models import EndUser as _wraps_EndUser +from models.account import Tenant +from models.enums import EndUserType +from models.model import App as _public_app_admission_App +from models.model import AppMode, AppModelConfig, CustomizeTokenStrategy, Message, Site +from models.model import EndUser as _public_app_admission_EndUser +from models.workflow import Workflow, WorkflowRun, WorkflowType +from repositories.app_definition_query_repository import AppDefinitionQueryRepository +from repositories.web_passport_repository import WebPassportRepository +from repositories.webapp_access_query_repository import WebAppAccessQueryRepository +from services.app_definition_query_service import AppDefinitionQueryService +from services.entities.feature_entities import FeatureModel +from services.file_service import FileService +from services.web_app_runtime_query_service import WebAppRuntimeQueryService +from services.web_authentication_adapters import PassportWebAppSessionGateway +from services.web_passport_gateways import PassportTokenGateway +from services.web_passport_service import WebPassportService +from tests.helpers.unit.config_override import apply_config_overrides + +"""Public bootstrap admission with real SQL repositories and the canonical hook.""" -@pytest.fixture(autouse=True) -def _reset_logging_context(): - clear_request_context() - yield - clear_request_context() +@dataclass +class AdmissionHarness: + app: _public_app_admission_App + old_user: _public_app_admission_EndUser + client: FlaskClient + services: ApplicationServices + issue: MagicMock + auth_calls: tuple[MagicMock, ...] + mode_lookup: MagicMock + login_verify: MagicMock + session_verify: MagicMock + metadata_calls: tuple[MagicMock, ...] + access_repo: WebAppAccessQueryRepository -def test_validate_jwt_token_sets_logging_identity_before_view() -> None: - from controllers.web import wraps - - app_model = App(id="app-id", tenant_id="tenant-id") - end_user = EndUser(id="end-user-id", tenant_id="tenant-id", type=None) - clear_request_context() - - @wraps.validate_jwt_token - def protected_view(received_app, received_user): - assert get_identity_context() == ("tenant-id", "end-user-id", "end_user") - return received_app, received_user - - with mock.patch.object(wraps, "decode_jwt_token", return_value=(app_model, end_user)): - result = protected_view() - - assert result == (app_model, end_user) - - -def test_validate_jwt_token_does_not_set_identity_when_authentication_fails() -> None: - from controllers.web import wraps - - clear_request_context() - - @wraps.validate_jwt_token - def protected_view(_app, _user): - raise AssertionError("view must not be called") - - with ( - mock.patch.object(wraps, "decode_jwt_token", side_effect=Unauthorized()), - pytest.raises(Unauthorized), - ): - protected_view() - - assert get_identity_context() == ("", "", "") - - -def test_decode_jwt_token_uses_shared_session_factory(sqlite_session: Session) -> None: - from controllers.web import wraps - from models.enums import EndUserType - from models.model import AppMode, CustomizeTokenStrategy, Site - - tenant_id = str(uuid4()) - app_model = App( - tenant_id=tenant_id, - mode=AppMode.CHAT.value, - name="test-app", +def _harness( + session: Session, + factory: sessionmaker[Session], + monkeypatch: pytest.MonkeyPatch, + mocker: MockerFixture, + base_services: ApplicationServices, + *, + published: bool, + missing: bool = False, +) -> AdmissionHarness: + app = _public_app_admission_App( + id=str(uuid4()), + tenant_id=str(uuid4()), + mode=AppMode.ADVANCED_CHAT, + name="Must not leak", enable_site=True, enable_api=True, ) - sqlite_session.add(app_model) - sqlite_session.flush() - config = AppModelConfig(app_id=app_model.id) - config.id = str(uuid4()) - app_model.app_model_config_id = config.id - sqlite_session.add(config) - sqlite_session.commit() - - site = Site( - app_id=app_model.id, - title="test-site", - default_language="en-US", - customize_token_strategy=CustomizeTokenStrategy.NOT_ALLOW, - code="app-code", - ) - end_user = EndUser( - tenant_id=tenant_id, - app_id=app_model.id, + old_user = _public_app_admission_EndUser( + id=str(uuid4()), + tenant_id=app.tenant_id, + app_id=app.id, + session_id="old", + _is_anonymous=True, type=EndUserType.BROWSER, - session_id="session-id", ) - sqlite_session.add_all((site, end_user)) - sqlite_session.commit() + if not missing: + tenant = Tenant(name="Private workspace") + tenant.id = app.tenant_id + session.add(tenant) + workflow = Workflow( + id=str(uuid4()), + tenant_id=app.tenant_id, + app_id=app.id, + type=WorkflowType.CHAT, + version="published" if published else "draft", + graph='{"nodes":[]}', + features="{}", + created_by=str(uuid4()), + environment_variables=[], + conversation_variables=[], + rag_pipeline_variables=[], + ) + if published: + app.workflow_id = workflow.id + session.add_all( + [ + app, + old_user, + workflow, + Site( + app_id=app.id, + code="fixture-code", + title="Must not leak", + default_language="en-US", + customize_token_strategy=CustomizeTokenStrategy.UUID, + ), + ] + ) + session.commit() - with ( - mock.patch.object(wraps, "extract_webapp_passport", return_value="jwt-token"), - mock.patch.object(wraps, "PassportService") as mock_passport_service, - mock.patch.object( - wraps.SystemFeatureService, - "is_webapp_auth_enabled", - return_value=False, + access_repo = WebAppAccessQueryRepository(session_factory=factory) + passport_repo = WebPassportRepository(session_factory=factory, generate_session_id=lambda: "new-session") + access = base_services.webapp_access + monkeypatch.setattr(access, "_access", access_repo) + monkeypatch.setattr(access, "_webapp_auth_enabled", False) + mode_lookup = mocker.spy(access._policy, "get_access_mode") + auth = base_services.web_passport._auth + auth_calls = ( + mocker.spy(auth, "is_webapp_auth_enabled"), + mocker.spy(auth, "get_app_auth_type"), + ) + tokens = PassportTokenGateway(passport=PassportService()) + issue = mocker.patch.object(PassportService, "issue", return_value="new-passport") + mocker.patch.object( + PassportService, + "verify", + return_value={"app_id": app.id, "app_code": "fixture-code", "end_user_id": old_user.id}, + ) + app_sessions = PassportWebAppSessionGateway(sessions=access_repo, app_access=access) + session_verify = mocker.spy(app_sessions, "verify") + authentication = base_services.web_authentication + monkeypatch.setattr(authentication, "_app_access", access) + monkeypatch.setattr(authentication, "_app_sessions", app_sessions) + monkeypatch.setattr(authentication, "_private_app_access_enabled", False) + login_verify = mocker.spy(authentication._tokens, "verify_access_token") + definitions = AppDefinitionQueryRepository(session_factory=factory) + services = replace( + base_services, + webapp_access=access, + web_authentication=authentication, + web_passport=WebPassportService( + passports=passport_repo, + auth=auth, + tokens=tokens, + now=lambda: datetime.now(UTC), + access_token_expire_minutes=60, ), - ): - mock_passport_service.return_value.verify.return_value = { - "app_code": "app-code", - "app_id": app_model.id, - "end_user_id": end_user.id, - } - - with Flask(__name__).test_request_context("/", headers={"X-App-Code": "app-code"}): - result_app, result_end_user = wraps.decode_jwt_token() - - assert result_app.id == app_model.id - assert result_end_user.id == end_user.id - - -@pytest.mark.parametrize("missing", ["app", "site", "code", "disabled", "unpublished", "end_user", "resource"]) -def test_post_missing_app_is_canonical_but_other_not_found_keeps_its_owner( - sqlite_session: Session, monkeypatch: pytest.MonkeyPatch, missing: str -) -> None: - from controllers.common.app_access_error import register_app_access_error_metadata - from controllers.web import wraps - from libs.external_api import ExternalApi - from models.enums import EndUserType - from models.model import AppMode, CustomizeTokenStrategy, Site - - app_id, tenant_id, end_user_id = str(uuid4()), str(uuid4()), str(uuid4()) - config_id = str(uuid4()) - if missing != "app": - sqlite_session.add( - App( - id=app_id, - tenant_id=tenant_id, - name="fixture", - mode=AppMode.CHAT, - enable_site=missing != "disabled", - enable_api=True, - app_model_config_id=config_id if missing != "unpublished" else None, - ) + app_definitions=AppDefinitionQueryService(definitions=definitions, builtin_icon_url_prefix=""), + web_app_runtime=WebAppRuntimeQueryService( + runtime=definitions, + file_service=FileService(session_factory=factory), + workspace_features=lambda _: FeatureModel(), + files_url="https://files.example", + deployment_edition=DeploymentEdition.CLOUD, + ), + ) + metadata_calls = tuple( + mocker.spy(services.app_definitions, method) + for method in ( + "get_mode", + "get_parameters", + "get_public_parameters", + "get_tool_icons", + "get_summary", + "get_site_configuration", ) - if missing != "unpublished": - config = AppModelConfig(app_id=app_id) - config.id = config_id - sqlite_session.add(config) - if missing != "site": - sqlite_session.add( - Site( - app_id=app_id, - code="fixture-code", - title="fixture", - default_language="en-US", - customize_token_strategy=CustomizeTokenStrategy.NOT_ALLOW, - ) - ) - if missing == "resource": - sqlite_session.add( - EndUser( - id=end_user_id, - tenant_id=tenant_id, - app_id=app_id, - type=EndUserType.BROWSER, - session_id="fixture-session", - ) - ) - sqlite_session.commit() - apply_config_overrides(monkeypatch, NETWORK_ACCESS_TRUSTED_PROXY_CIDRS="172.18.0.0/16") - app = Flask(__name__) - app.config.update(TESTING=True, RESTX_ERROR_404_HELP=False) - bp = Blueprint("web_post_fixture", __name__, url_prefix="/api") - register_app_access_error_metadata(bp, surface="web") - api = ExternalApi(bp) - calls: list[str] = [] + ) + (mocker.spy(services.web_app_runtime, "get_bootstrap"),) + for module in (app_controller, login, passport, site, wraps): + monkeypatch.setattr(module, "application_services", lambda: services) + monkeypatch.setattr(wraps, "extract_webapp_passport", lambda *_: "old-passport") + monkeypatch.setattr(login, "extract_webapp_passport", lambda *_: "old-passport") + monkeypatch.setattr(wraps.SystemFeatureService, "is_webapp_auth_enabled", lambda: False) + apply_config_overrides( + monkeypatch, + DEPLOYMENT_EDITION=DeploymentEdition.CLOUD, + NETWORK_ACCESS_TRUSTED_PROXY_CIDRS="172.18.0.0/16", + ) + flask = Flask(__name__) + flask.config.update(TESTING=True, RESTX_ERROR_404_HELP=False) + flask.register_blueprint(bp) + return AdmissionHarness( + app=app, + old_user=old_user, + client=flask.test_client(), + services=services, + issue=issue, + auth_calls=auth_calls, + mode_lookup=mode_lookup, + login_verify=login_verify, + session_verify=session_verify, + metadata_calls=metadata_calls, + access_repo=access_repo, + ) - class ProtectedResource(wraps.WebApiResource): - def post(self, _app: App, _end_user: EndUser) -> Never: - calls.append("view") - raise NotFound("Conversation not found.") - api.add_resource(ProtectedResource, "/chat-messages") - app.register_blueprint(bp) - with ( - mock.patch.object(wraps, "extract_webapp_passport", return_value="fixture-passport"), - mock.patch.object(wraps, "PassportService") as passport, - mock.patch.object(wraps.SystemFeatureService, "is_webapp_auth_enabled", return_value=False), - ): - passport.return_value.verify.return_value = { - "app_code": None if missing == "code" else "fixture-code", - "app_id": app_id, - "end_user_id": end_user_id, - } - response = app.test_client().post( - "/api/chat-messages", - json={}, +ROUTES = ( + "/webapp/access-mode?appCode=fixture-code", + "/login/status?app_code=fixture-code", + "/passport?user_id=new-session", + "/site", + "/meta", + "/parameters", +) + + +class TestPublicAppAdmissionCases: + @staticmethod + @pytest.mark.parametrize("route", ROUTES) + @pytest.mark.parametrize("missing", [False, True], ids=["unpublished", "missing"]) + def test_all_bootstrap_routes_share_404_before_user_or_metadata_side_effects( + sqlite_session: Session, + sqlite_session_factory: sessionmaker[Session], + monkeypatch: pytest.MonkeyPatch, + mocker: MockerFixture, + account_application_services: ApplicationServices, + route: str, + missing: bool, + ) -> None: + h = _harness( + sqlite_session, + sqlite_session_factory, + monkeypatch, + mocker, + account_application_services, + published=False, + missing=missing, + ) + response = h.client.get( + "/api" + route, headers={"X-App-Code": "fixture-code"}, - environ_overrides={"REMOTE_ADDR": "203.0.113.42"}, + environ_overrides={"REMOTE_ADDR": "172.18.0.2", "HTTP_X_FORWARDED_FOR": "203.0.113.42"}, ) - - assert response.status_code == 404 - if missing in {"app", "site", "code", "disabled", "unpublished"}: + assert response.status_code == 404 assert ( response.data == b'{"client_ip":"203.0.113.42","code":"app_not_found","message":"App not found.","status":404}' ) assert response.headers["Content-Type"] == "application/json" assert response.headers["Cache-Control"] == "no-store" - else: - assert response.get_json()["code"] == "not_found" + assert all(call.call_count == 0 for call in h.metadata_calls) + assert all(call.call_count == 0 for call in h.auth_calls) + h.mode_lookup.assert_not_called() + h.issue.assert_not_called() + h.login_verify.assert_not_called() + h.session_verify.assert_not_called() + assert sqlite_session.scalar(select(func.count()).select_from(_public_app_admission_EndUser)) == ( + 0 if missing else 1 + ) + assert sqlite_session.scalar(select(func.count()).select_from(WorkflowRun)) == 0 + assert sqlite_session.scalar(select(func.count()).select_from(Message)) == 0 + + @staticmethod + @pytest.mark.parametrize("route", ROUTES) + def test_published_bootstrap_remains_successful_with_normal_passport_behavior( + sqlite_session: Session, + sqlite_session_factory: sessionmaker[Session], + monkeypatch: pytest.MonkeyPatch, + route: str, + mocker: MockerFixture, + account_application_services: ApplicationServices, + ) -> None: + h = _harness( + sqlite_session, sqlite_session_factory, monkeypatch, mocker, account_application_services, published=True + ) + response = h.client.get("/api" + route, headers={"X-App-Code": "fixture-code"}) + assert response.status_code == 200 + if route.startswith("/passport"): + assert response.get_json() == {"access_token": "new-passport"} + h.issue.assert_called_once() + assert sqlite_session.scalar(select(func.count()).select_from(_public_app_admission_EndUser)) == 2 + else: + h.issue.assert_not_called() + assert sqlite_session.scalar(select(func.count()).select_from(_public_app_admission_EndUser)) == 1 + if route.startswith("/webapp"): + assert response.get_json() == {"accessMode": "public"} + elif route.startswith("/login"): + assert response.get_json() == {"logged_in": True, "app_logged_in": True} + h.session_verify.assert_called_once_with(token="old-passport", app_code="fixture-code", user_id=None) + h.login_verify.assert_not_called() + elif route == "/site": + assert response.get_json()["site"]["title"] == "Must not leak" + elif route == "/meta": + assert response.get_json() == {"tool_icons": {}} + else: + assert "user_input_form" in response.get_json() + + @staticmethod + @pytest.mark.parametrize("route", ["/site", "/meta", "/parameters"]) + def test_old_passport_cannot_read_metadata_after_published_version_is_removed( + sqlite_session: Session, + sqlite_session_factory: sessionmaker[Session], + monkeypatch: pytest.MonkeyPatch, + route: str, + mocker: MockerFixture, + account_application_services: ApplicationServices, + ) -> None: + h = _harness( + sqlite_session, sqlite_session_factory, monkeypatch, mocker, account_application_services, published=True + ) + assert h.access_repo.is_app_available(h.app.id) + h.app.workflow_id = None + sqlite_session.commit() + response = h.client.get( + "/api" + route, headers={"X-App-Code": "fixture-code"}, environ_overrides={"REMOTE_ADDR": "203.0.113.42"} + ) + assert response.status_code == 404 + assert response.get_json()["code"] == "app_not_found" + assert response.get_json()["client_ip"] == "203.0.113.42" + assert all(call.call_count == 0 for call in h.metadata_calls) + + @staticmethod + @pytest.mark.parametrize("changed_owner", ["site-app", "end-user-app", "end-user-tenant"]) + def test_old_passport_cannot_cross_persisted_app_or_tenant_owner_chain( + sqlite_session: Session, + sqlite_session_factory: sessionmaker[Session], + monkeypatch: pytest.MonkeyPatch, + mocker: MockerFixture, + account_application_services: ApplicationServices, + changed_owner: str, + ) -> None: + h = _harness( + sqlite_session, sqlite_session_factory, monkeypatch, mocker, account_application_services, published=True + ) + if changed_owner == "site-app": + persisted_site = sqlite_session.scalar(select(Site).where(Site.app_id == h.app.id)) + assert persisted_site is not None + persisted_site.app_id = str(uuid4()) + elif changed_owner == "end-user-app": + h.old_user.app_id = str(uuid4()) + else: + h.old_user.tenant_id = str(uuid4()) + sqlite_session.commit() + response = h.client.get("/api/site", headers={"X-App-Code": "fixture-code"}) + assert response.status_code == 404 + assert response.get_json()["code"] == "app_not_found" + assert all(call.call_count == 0 for call in h.metadata_calls) + + @staticmethod + @pytest.mark.parametrize("route", ROUTES) + def test_real_database_failure_is_not_reclassified_as_unpublished( + sqlite_session: Session, + sqlite_session_factory: sessionmaker[Session], + monkeypatch: pytest.MonkeyPatch, + route: str, + mocker: MockerFixture, + account_application_services: ApplicationServices, + ) -> None: + h = _harness( + sqlite_session, sqlite_session_factory, monkeypatch, mocker, account_application_services, published=False + ) + failure = OperationalError("select", {}, RuntimeError("database unavailable")) + mocker.patch.object(Session, "scalar", side_effect=failure) + mocker.patch.object(Session, "execute", side_effect=failure) + response = h.client.get("/api" + route, headers={"X-App-Code": "fixture-code"}) + assert response.status_code == (503 if route.startswith(("/webapp", "/login")) else 500) + assert response.get_json()["code"] != "app_not_found" assert "client_ip" not in response.get_json() - assert "Cache-Control" not in response.headers + h.issue.assert_not_called() + h.login_verify.assert_not_called() + h.session_verify.assert_not_called() + + +class TestWrapsCases: + @staticmethod + @pytest.fixture(autouse=True) + def _reset_logging_context(): + clear_request_context() + yield + clear_request_context() + + @staticmethod + def test_validate_jwt_token_sets_logging_identity_before_view() -> None: + from controllers.web import wraps + + app_model = _wraps_App(id="app-id", tenant_id="tenant-id") + end_user = _wraps_EndUser(id="end-user-id", tenant_id="tenant-id", type=None) + clear_request_context() + + @wraps.validate_jwt_token + def protected_view(received_app, received_user): + assert get_identity_context() == ("tenant-id", "end-user-id", "end_user") + return received_app, received_user + + with mock.patch.object(wraps, "decode_jwt_token", return_value=(app_model, end_user)): + result = protected_view() + + assert result == (app_model, end_user) + + @staticmethod + def test_validate_jwt_token_does_not_set_identity_when_authentication_fails() -> None: + from controllers.web import wraps + + clear_request_context() + + @wraps.validate_jwt_token + def protected_view(_app, _user): + raise AssertionError("view must not be called") + + with ( + mock.patch.object(wraps, "decode_jwt_token", side_effect=Unauthorized()), + pytest.raises(Unauthorized), + ): + protected_view() + + assert get_identity_context() == ("", "", "") + + @staticmethod + def test_decode_jwt_token_uses_shared_session_factory(sqlite_session: Session) -> None: + from controllers.web import wraps + from models.enums import EndUserType + from models.model import AppMode, CustomizeTokenStrategy, Site + + tenant_id = str(uuid4()) + app_model = _wraps_App( + tenant_id=tenant_id, + mode=AppMode.CHAT.value, + name="test-app", + enable_site=True, + enable_api=True, + ) + sqlite_session.add(app_model) + sqlite_session.flush() + config = AppModelConfig(app_id=app_model.id) + config.id = str(uuid4()) + app_model.app_model_config_id = config.id + sqlite_session.add(config) + sqlite_session.commit() + + site = Site( + app_id=app_model.id, + title="test-site", + default_language="en-US", + customize_token_strategy=CustomizeTokenStrategy.NOT_ALLOW, + code="app-code", + ) + end_user = _wraps_EndUser( + tenant_id=tenant_id, + app_id=app_model.id, + type=EndUserType.BROWSER, + session_id="session-id", + ) + sqlite_session.add_all((site, end_user)) + sqlite_session.commit() + + with ( + mock.patch.object(wraps, "extract_webapp_passport", return_value="jwt-token"), + mock.patch.object(wraps, "PassportService") as mock_passport_service, + mock.patch.object( + wraps.SystemFeatureService, + "is_webapp_auth_enabled", + return_value=False, + ), + ): + mock_passport_service.return_value.verify.return_value = { + "app_code": "app-code", + "app_id": app_model.id, + "end_user_id": end_user.id, + } + + with Flask(__name__).test_request_context("/", headers={"X-App-Code": "app-code"}): + result_app, result_end_user = wraps.decode_jwt_token() + + assert result_app.id == app_model.id + assert result_end_user.id == end_user.id + + @staticmethod + @pytest.mark.parametrize("missing", ["app", "site", "code", "disabled", "unpublished", "end_user", "resource"]) + def test_post_missing_app_is_canonical_but_other_not_found_keeps_its_owner( + sqlite_session: Session, monkeypatch: pytest.MonkeyPatch, missing: str + ) -> None: + from controllers.common.app_access_error import register_app_access_error_metadata + from controllers.web import wraps + from libs.external_api import ExternalApi + from models.enums import EndUserType + from models.model import AppMode, CustomizeTokenStrategy, Site + + app_id, tenant_id, end_user_id = str(uuid4()), str(uuid4()), str(uuid4()) + config_id = str(uuid4()) + if missing != "app": + sqlite_session.add( + _wraps_App( + id=app_id, + tenant_id=tenant_id, + name="fixture", + mode=AppMode.CHAT, + enable_site=missing != "disabled", + enable_api=True, + app_model_config_id=config_id if missing != "unpublished" else None, + ) + ) + if missing != "unpublished": + config = AppModelConfig(app_id=app_id) + config.id = config_id + sqlite_session.add(config) + if missing != "site": + sqlite_session.add( + Site( + app_id=app_id, + code="fixture-code", + title="fixture", + default_language="en-US", + customize_token_strategy=CustomizeTokenStrategy.NOT_ALLOW, + ) + ) if missing == "resource": - assert response.get_json()["message"] == "Conversation not found." - assert calls == (["view"] if missing == "resource" else []) + sqlite_session.add( + _wraps_EndUser( + id=end_user_id, + tenant_id=tenant_id, + app_id=app_id, + type=EndUserType.BROWSER, + session_id="fixture-session", + ) + ) + sqlite_session.commit() + apply_config_overrides(monkeypatch, NETWORK_ACCESS_TRUSTED_PROXY_CIDRS="172.18.0.0/16") + app = Flask(__name__) + app.config.update(TESTING=True, RESTX_ERROR_404_HELP=False) + bp = Blueprint("web_post_fixture", __name__, url_prefix="/api") + register_app_access_error_metadata(bp, surface="web") + api = ExternalApi(bp) + calls: list[str] = [] + class ProtectedResource(wraps.WebApiResource): + def post(self, _app: _wraps_App, _end_user: _wraps_EndUser) -> Never: + calls.append("view") + raise NotFound("Conversation not found.") -@pytest.mark.usefixtures("app_query_services") -@pytest.mark.parametrize("token", [None, "expired-token"]) -def test_web_auth_rejects_unknown_app_code_with_http_not_found(app, monkeypatch, token): - from controllers.web import wraps - from libs.external_api import ExternalApi + api.add_resource(ProtectedResource, "/chat-messages") + app.register_blueprint(bp) + with ( + mock.patch.object(wraps, "extract_webapp_passport", return_value="fixture-passport"), + mock.patch.object(wraps, "PassportService") as passport, + mock.patch.object(wraps.SystemFeatureService, "is_webapp_auth_enabled", return_value=False), + ): + passport.return_value.verify.return_value = { + "app_code": None if missing == "code" else "fixture-code", + "app_id": app_id, + "end_user_id": end_user_id, + } + response = app.test_client().post( + "/api/chat-messages", + json={}, + headers={"X-App-Code": "fixture-code"}, + environ_overrides={"REMOTE_ADDR": "203.0.113.42"}, + ) - class ProtectedApi(wraps.WebApiResource): - def get(self, _app, _user): - pytest.fail("An unknown WebApp must not reach the protected resource") + assert response.status_code == 404 + if missing in {"app", "site", "code", "disabled", "unpublished"}: + assert ( + response.data + == b'{"client_ip":"203.0.113.42","code":"app_not_found","message":"App not found.","status":404}' + ) + assert response.headers["Content-Type"] == "application/json" + assert response.headers["Cache-Control"] == "no-store" + else: + assert response.get_json()["code"] == "not_found" + assert "client_ip" not in response.get_json() + assert "Cache-Control" not in response.headers + if missing == "resource": + assert response.get_json()["message"] == "Conversation not found." + assert calls == (["view"] if missing == "resource" else []) - def reject_token(_self, _token): - raise Unauthorized("Authentication has expired.") + @staticmethod + @pytest.mark.usefixtures("app_query_services") + @pytest.mark.parametrize("token", [None, "expired-token"]) + def test_web_auth_rejects_unknown_app_code_with_http_not_found(app, monkeypatch, token): + from controllers.web import wraps + from libs.external_api import ExternalApi - monkeypatch.setattr(wraps.SystemFeatureService, "is_webapp_auth_enabled", staticmethod(lambda: True)) - monkeypatch.setattr(wraps, "extract_webapp_passport", lambda _code, _request: token) - monkeypatch.setattr(wraps.PassportService, "verify", reject_token) - api = ExternalApi(app) - api.add_resource(ProtectedApi, "/protected") - response = app.test_client().get("/protected", headers={"X-App-Code": "does-not-exist"}) - assert response.status_code == 404 - assert response.get_json()["code"] == "app_not_found" + class ProtectedApi(wraps.WebApiResource): + def get(self, _app, _user): + pytest.fail("An unknown WebApp must not reach the protected resource") + + def reject_token(_self, _token): + raise Unauthorized("Authentication has expired.") + + monkeypatch.setattr(wraps.SystemFeatureService, "is_webapp_auth_enabled", staticmethod(lambda: True)) + monkeypatch.setattr(wraps, "extract_webapp_passport", lambda _code, _request: token) + monkeypatch.setattr(wraps.PassportService, "verify", reject_token) + api = ExternalApi(app) + api.add_resource(ProtectedApi, "/protected") + response = app.test_client().get("/protected", headers={"X-App-Code": "does-not-exist"}) + assert response.status_code == 404 + assert response.get_json()["code"] == "app_not_found" diff --git a/api/tests/unit_tests/core/agent/test_base_agent_runner.py b/api/tests/unit_tests/core/agent/test_base_agent_runner.py index 476d7432148..6f3c23e9d32 100644 --- a/api/tests/unit_tests/core/agent/test_base_agent_runner.py +++ b/api/tests/unit_tests/core/agent/test_base_agent_runner.py @@ -43,8 +43,8 @@ from graphon.model_runtime.entities import LLMUsage, PromptMessageTool from models.enums import ConversationFromSource, CreatorUserRole, MessageFileBelongsTo from models.model import AppMode, AppModelConfig, Conversation, Message, MessageAgentThought, MessageFile, UploadFile from services.workflow_run_agg import WorkflowRunAgg +from tests.helpers.unit.core.model_fixtures import make_model_config from tests.tool_fixtures import make_runtime_tool -from tests.unit_tests.core.model_fixtures import make_model_config def _message( diff --git a/api/tests/unit_tests/core/app/app_config/common/test_parameters_mapping.py b/api/tests/unit_tests/core/app/app_config/common/parameters_mapping/test___init__.py similarity index 99% rename from api/tests/unit_tests/core/app/app_config/common/test_parameters_mapping.py rename to api/tests/unit_tests/core/app/app_config/common/parameters_mapping/test___init__.py index fdf2cb3d85f..9e0a550c856 100644 --- a/api/tests/unit_tests/core/app/app_config/common/test_parameters_mapping.py +++ b/api/tests/unit_tests/core/app/app_config/common/parameters_mapping/test___init__.py @@ -4,7 +4,7 @@ import pytest # Module under test from core.app.app_config.common import parameters_mapping -from tests.unit_tests.config_override import apply_config_overrides +from tests.helpers.unit.config_override import apply_config_overrides type _DefaultedFeatureKey = Literal[ "suggested_questions", diff --git a/api/tests/unit_tests/core/app/app_config/common/test_sensitive_word_avoidance_manager.py b/api/tests/unit_tests/core/app/app_config/common/sensitive_word_avoidance/test_manager.py similarity index 100% rename from api/tests/unit_tests/core/app/app_config/common/test_sensitive_word_avoidance_manager.py rename to api/tests/unit_tests/core/app/app_config/common/sensitive_word_avoidance/test_manager.py diff --git a/api/tests/unit_tests/core/app/app_config/easy_ui_based_app/test_agent_manager.py b/api/tests/unit_tests/core/app/app_config/easy_ui_based_app/agent/test_manager.py similarity index 100% rename from api/tests/unit_tests/core/app/app_config/easy_ui_based_app/test_agent_manager.py rename to api/tests/unit_tests/core/app/app_config/easy_ui_based_app/agent/test_manager.py diff --git a/api/tests/unit_tests/core/app/app_config/easy_ui_based_app/test_dataset_manager.py b/api/tests/unit_tests/core/app/app_config/easy_ui_based_app/dataset/test_manager.py similarity index 100% rename from api/tests/unit_tests/core/app/app_config/easy_ui_based_app/test_dataset_manager.py rename to api/tests/unit_tests/core/app/app_config/easy_ui_based_app/dataset/test_manager.py diff --git a/api/tests/unit_tests/core/app/app_config/easy_ui_based_app/test_model_config_converter.py b/api/tests/unit_tests/core/app/app_config/easy_ui_based_app/model_config/test_converter.py similarity index 100% rename from api/tests/unit_tests/core/app/app_config/easy_ui_based_app/test_model_config_converter.py rename to api/tests/unit_tests/core/app/app_config/easy_ui_based_app/model_config/test_converter.py diff --git a/api/tests/unit_tests/core/app/app_config/easy_ui_based_app/model_config/test_manager.py b/api/tests/unit_tests/core/app/app_config/easy_ui_based_app/model_config/test_manager.py new file mode 100644 index 00000000000..0894c73771a --- /dev/null +++ b/api/tests/unit_tests/core/app/app_config/easy_ui_based_app/model_config/test_manager.py @@ -0,0 +1,376 @@ +from types import SimpleNamespace +from typing import cast +from unittest.mock import patch + +import pytest +from pytest_mock import MockerFixture + +from core.app.app_config.easy_ui_based_app.model_config.manager import ModelConfigManager +from core.app.app_config.entities import ModelConfigEntity +from core.entities.model_entities import ModelStatus, ModelWithProviderEntity, SimpleModelProviderEntity +from core.entities.provider_configuration import ProviderConfigurations +from core.plugin.impl.model_runtime_factory import PluginModelAssembly +from graphon.model_runtime.entities.common_entities import I18nObject +from graphon.model_runtime.entities.model_entities import FetchFrom, ModelPropertyKey, ModelType +from graphon.model_runtime.entities.provider_entities import ProviderEntity +from models.model import AppModelConfigDict +from models.provider_ids import ModelProviderID + +# Target + + +class TestModelConfigManagerCases: + # ----------------------------- + # Fixtures + # ----------------------------- + + @staticmethod + @pytest.fixture + def valid_completion_params() -> dict[str, object]: + return {"temperature": 0.7, "stop": ["\n"]} + + @staticmethod + @pytest.fixture + def valid_model_list(provider_entities: list[ProviderEntity]) -> list[ModelWithProviderEntity]: + return [ + ModelWithProviderEntity( + model="gpt-4", + label=I18nObject(en_US="GPT-4"), + model_type=ModelType.LLM, + fetch_from=FetchFrom.PREDEFINED_MODEL, + model_properties={ModelPropertyKey.MODE: "chat"}, + status=ModelStatus.ACTIVE, + provider=SimpleModelProviderEntity(provider_entities[0]), + ) + ] + + @staticmethod + @pytest.fixture + def provider_entities() -> list[ProviderEntity]: + return [ + ProviderEntity( + provider="langgenius/openai/openai", + label=I18nObject(en_US="OpenAI"), + supported_model_types=[ModelType.LLM], + configurate_methods=[], + ) + ] + + @staticmethod + @pytest.fixture + def valid_config() -> AppModelConfigDict: + return cast( + AppModelConfigDict, + { + "model": { + "provider": "langgenius/openai/openai", + "name": "gpt-4", + "completion_params": {"temperature": 0.5, "stop": ["END"]}, + } + }, + ) + + # ----------------------------- + # Test Class + # ----------------------------- + + class TestModelConfigManager: + @staticmethod + def _patch_model_assembly( + mocker: MockerFixture, *, provider_entities: list[ProviderEntity], model_list: list[ModelWithProviderEntity] + ) -> PluginModelAssembly: + assembly = PluginModelAssembly(tenant_id="tenant1") + configurations = ProviderConfigurations(tenant_id="tenant1") + + def get_providers() -> list[ProviderEntity]: + return provider_entities + + def get_configurations(tenant_id: str) -> ProviderConfigurations: + assert tenant_id == assembly.tenant_id + return configurations + + def get_models( + self: ProviderConfigurations, *, provider: str, model_type: ModelType + ) -> list[ModelWithProviderEntity]: + assert self is configurations + assert provider in [entity.provider for entity in provider_entities] + assert model_type == ModelType.LLM + return model_list + + def create_assembly(*, tenant_id: str) -> PluginModelAssembly: + assert tenant_id == assembly.tenant_id + return assembly + + mocker.patch.object(assembly.model_provider_factory, "get_providers", new=get_providers) + mocker.patch.object(assembly.provider_manager, "get_configurations", new=get_configurations) + mocker.patch.object(ProviderConfigurations, "get_models", new=get_models) + mocker.patch( + "core.app.app_config.easy_ui_based_app.model_config.manager.create_plugin_model_assembly", + new=create_assembly, + ) + return assembly + + # ========================================================== + # convert + # ========================================================== + + def test_convert_success(self, valid_config: AppModelConfigDict) -> None: + result = ModelConfigManager.convert(valid_config) + + assert result.provider == "langgenius/openai/openai" + assert result.model == "gpt-4" + assert result.parameters == {"temperature": 0.5} + assert result.stop == ["END"] + + def test_convert_missing_model(self) -> None: + with pytest.raises(ValueError, match="model is required"): + ModelConfigManager.convert(cast(AppModelConfigDict, {})) + + def test_convert_without_stop(self) -> None: + config = { + "model": { + "provider": "langgenius/openai/openai", + "name": "gpt-4", + "completion_params": {"temperature": 0.9}, + } + } + result = ModelConfigManager.convert(cast(AppModelConfigDict, config)) + assert result.stop == [] + assert result.parameters == {"temperature": 0.9} + + # ========================================================== + # validate_model_completion_params + # ========================================================== + + @pytest.mark.parametrize( + "invalid_cp", + [None, "string", 123, []], + ) + def test_validate_model_completion_params_invalid_type(self, invalid_cp: object) -> None: + with pytest.raises(ValueError, match="must be of object type"): + ModelConfigManager.validate_model_completion_params(invalid_cp) # pyrefly: ignore[bad-argument-type] + + def test_validate_model_completion_params_default_stop(self) -> None: + cp = {"temperature": 0.2} + result = ModelConfigManager.validate_model_completion_params(cp) + assert result["stop"] == [] + + def test_validate_model_completion_params_invalid_stop_type(self) -> None: + cp = {"stop": "invalid"} + with pytest.raises(ValueError, match="must be of list type"): + ModelConfigManager.validate_model_completion_params(cp) + + def test_validate_model_completion_params_stop_length_exceeded(self) -> None: + cp = {"stop": [1, 2, 3, 4, 5]} + with pytest.raises(ValueError, match="less than 4"): + ModelConfigManager.validate_model_completion_params(cp) + + # ========================================================== + # validate_and_set_defaults + # ========================================================== + + def test_validate_and_set_defaults_success( + self, + mocker: MockerFixture, + valid_config: AppModelConfigDict, + provider_entities: list[ProviderEntity], + valid_model_list: list[ModelWithProviderEntity], + ) -> None: + self._patch_model_assembly( + mocker, + provider_entities=provider_entities, + model_list=valid_model_list, + ) + + updated_config, keys = ModelConfigManager.validate_and_set_defaults("tenant1", valid_config) + + assert updated_config["model"]["mode"] == "chat" + assert keys == ["model"] + + def test_validate_and_set_defaults_missing_model(self) -> None: + with pytest.raises(ValueError, match="model is required"): + ModelConfigManager.validate_and_set_defaults("tenant1", {}) + + def test_validate_and_set_defaults_model_not_dict(self) -> None: + with pytest.raises(ValueError, match="object type"): + ModelConfigManager.validate_and_set_defaults("tenant1", {"model": "invalid"}) + + def test_validate_and_set_defaults_missing_provider( + self, mocker: MockerFixture, provider_entities: list[ProviderEntity] + ) -> None: + config = {"model": {"name": "gpt-4", "completion_params": dict[str, object]()}} + self._patch_model_assembly(mocker, provider_entities=provider_entities, model_list=[]) + + with pytest.raises(ValueError, match="model.provider is required"): + ModelConfigManager.validate_and_set_defaults("tenant1", config) + + def test_validate_and_set_defaults_invalid_provider( + self, mocker: MockerFixture, provider_entities: list[ProviderEntity] + ) -> None: + config = { + "model": {"provider": "invalid/provider", "name": "gpt-4", "completion_params": dict[str, object]()} + } + self._patch_model_assembly(mocker, provider_entities=provider_entities, model_list=[]) + + with pytest.raises(ValueError, match="model.provider is required"): + ModelConfigManager.validate_and_set_defaults("tenant1", config) + + def test_validate_and_set_defaults_missing_name( + self, mocker: MockerFixture, provider_entities: list[ProviderEntity] + ) -> None: + config = {"model": {"provider": "langgenius/openai/openai", "completion_params": dict[str, object]()}} + self._patch_model_assembly(mocker, provider_entities=provider_entities, model_list=[]) + + with pytest.raises(ValueError, match="model.name is required"): + ModelConfigManager.validate_and_set_defaults("tenant1", config) + + def test_validate_and_set_defaults_empty_models( + self, mocker: MockerFixture, provider_entities: list[ProviderEntity] + ) -> None: + config = { + "model": { + "provider": "langgenius/openai/openai", + "name": "gpt-4", + "completion_params": dict[str, object](), + } + } + self._patch_model_assembly(mocker, provider_entities=provider_entities, model_list=[]) + + with pytest.raises(ValueError, match="must be in the specified model list"): + ModelConfigManager.validate_and_set_defaults("tenant1", config) + + def test_validate_and_set_defaults_invalid_model_name( + self, + mocker: MockerFixture, + provider_entities: list[ProviderEntity], + valid_model_list: list[ModelWithProviderEntity], + ) -> None: + config = { + "model": { + "provider": "langgenius/openai/openai", + "name": "invalid", + "completion_params": dict[str, object](), + } + } + self._patch_model_assembly( + mocker, + provider_entities=provider_entities, + model_list=valid_model_list, + ) + + with pytest.raises(ValueError, match="must be in the specified model list"): + ModelConfigManager.validate_and_set_defaults("tenant1", config) + + def test_validate_and_set_defaults_default_mode_when_missing( + self, + mocker: MockerFixture, + provider_entities: list[ProviderEntity], + valid_model_list: list[ModelWithProviderEntity], + ) -> None: + model = valid_model_list[0] + model.model_properties = dict[ModelPropertyKey, object]() + + config = { + "model": { + "provider": "langgenius/openai/openai", + "name": "gpt-4", + "completion_params": dict[str, object](), + } + } + self._patch_model_assembly(mocker, provider_entities=provider_entities, model_list=[model]) + + updated_config, _ = ModelConfigManager.validate_and_set_defaults("tenant1", config) + + assert updated_config["model"]["mode"] == "completion" + + def test_validate_and_set_defaults_missing_completion_params( + self, + mocker: MockerFixture, + provider_entities: list[ProviderEntity], + valid_model_list: list[ModelWithProviderEntity], + ) -> None: + config = {"model": {"provider": "langgenius/openai/openai", "name": "gpt-4"}} + self._patch_model_assembly( + mocker, + provider_entities=provider_entities, + model_list=valid_model_list, + ) + + with pytest.raises(ValueError, match="completion_params is required"): + ModelConfigManager.validate_and_set_defaults("tenant1", config) + + def test_validate_and_set_defaults_provider_without_slash_converted( + self, + mocker: MockerFixture, + provider_entities: list[ProviderEntity], + valid_model_list: list[ModelWithProviderEntity], + ) -> None: + """ + Covers branch where provider does not contain '/' and + ModelProviderID conversion is triggered (line 64). + """ + config = { + "model": { + "provider": "openai", # no slash -> triggers conversion + "name": "gpt-4", + "completion_params": dict[str, object](), + } + } + + self._patch_model_assembly(mocker, provider_entities=provider_entities, model_list=valid_model_list) + + updated_config, _ = ModelConfigManager.validate_and_set_defaults("tenant1", config) + + assert updated_config["model"]["provider"] == "langgenius/openai/openai" + + +class TestEasyUiModelConfigManagerCases: + @staticmethod + def test_validate_and_set_defaults_reuses_single_model_assembly(): + provider_name = str(ModelProviderID("openai")) + provider_entity = SimpleNamespace(provider=provider_name) + model = SimpleNamespace(model="gpt-4o-mini", model_properties={ModelPropertyKey.MODE: "chat"}) + provider_configurations = SimpleNamespace(get_models=lambda **kwargs: [model]) + assembly = SimpleNamespace( + model_provider_factory=SimpleNamespace(get_providers=lambda: [provider_entity]), + provider_manager=SimpleNamespace(get_configurations=lambda tenant_id: provider_configurations), + ) + config = { + "model": { + "provider": "openai", + "name": "gpt-4o-mini", + "completion_params": {"stop": []}, + } + } + + with patch( + "core.app.app_config.easy_ui_based_app.model_config.manager.create_plugin_model_assembly", + return_value=assembly, + ) as mock_assembly: + result, keys = ModelConfigManager.validate_and_set_defaults("tenant-1", config) + + assert result["model"]["provider"] == provider_name + assert result["model"]["mode"] == "chat" + assert keys == ["model"] + mock_assembly.assert_called_once_with(tenant_id="tenant-1") + + @staticmethod + def test_convert_keeps_model_config_shape(): + config = { + "model": { + "provider": "openai", + "name": "gpt-4o-mini", + "mode": "chat", + "completion_params": {"temperature": 0.3, "stop": ["END"]}, + } + } + + result = ModelConfigManager.convert(config) + + assert result == ModelConfigEntity( + provider="openai", + model="gpt-4o-mini", + mode="chat", + parameters={"temperature": 0.3}, + stop=["END"], + ) diff --git a/api/tests/unit_tests/core/app/app_config/easy_ui_based_app/test_prompt_template_manager.py b/api/tests/unit_tests/core/app/app_config/easy_ui_based_app/prompt_template/test_manager.py similarity index 100% rename from api/tests/unit_tests/core/app/app_config/easy_ui_based_app/test_prompt_template_manager.py rename to api/tests/unit_tests/core/app/app_config/easy_ui_based_app/prompt_template/test_manager.py diff --git a/api/tests/unit_tests/core/app/app_config/easy_ui_based_app/test_model_config_manager.py b/api/tests/unit_tests/core/app/app_config/easy_ui_based_app/test_model_config_manager.py deleted file mode 100644 index 26b0ead50e9..00000000000 --- a/api/tests/unit_tests/core/app/app_config/easy_ui_based_app/test_model_config_manager.py +++ /dev/null @@ -1,309 +0,0 @@ -from typing import cast - -import pytest -from pytest_mock import MockerFixture - -# Target -from core.app.app_config.easy_ui_based_app.model_config.manager import ModelConfigManager -from core.entities.model_entities import ModelStatus, ModelWithProviderEntity, SimpleModelProviderEntity -from core.entities.provider_configuration import ProviderConfigurations -from core.plugin.impl.model_runtime_factory import PluginModelAssembly -from graphon.model_runtime.entities.common_entities import I18nObject -from graphon.model_runtime.entities.model_entities import FetchFrom, ModelPropertyKey, ModelType -from graphon.model_runtime.entities.provider_entities import ProviderEntity -from models.model import AppModelConfigDict - -# ----------------------------- -# Fixtures -# ----------------------------- - - -@pytest.fixture -def valid_completion_params() -> dict[str, object]: - return {"temperature": 0.7, "stop": ["\n"]} - - -@pytest.fixture -def valid_model_list(provider_entities: list[ProviderEntity]) -> list[ModelWithProviderEntity]: - return [ - ModelWithProviderEntity( - model="gpt-4", - label=I18nObject(en_US="GPT-4"), - model_type=ModelType.LLM, - fetch_from=FetchFrom.PREDEFINED_MODEL, - model_properties={ModelPropertyKey.MODE: "chat"}, - status=ModelStatus.ACTIVE, - provider=SimpleModelProviderEntity(provider_entities[0]), - ) - ] - - -@pytest.fixture -def provider_entities() -> list[ProviderEntity]: - return [ - ProviderEntity( - provider="langgenius/openai/openai", - label=I18nObject(en_US="OpenAI"), - supported_model_types=[ModelType.LLM], - configurate_methods=[], - ) - ] - - -@pytest.fixture -def valid_config() -> AppModelConfigDict: - return cast( - AppModelConfigDict, - { - "model": { - "provider": "langgenius/openai/openai", - "name": "gpt-4", - "completion_params": {"temperature": 0.5, "stop": ["END"]}, - } - }, - ) - - -# ----------------------------- -# Test Class -# ----------------------------- - - -class TestModelConfigManager: - @staticmethod - def _patch_model_assembly( - mocker: MockerFixture, *, provider_entities: list[ProviderEntity], model_list: list[ModelWithProviderEntity] - ) -> PluginModelAssembly: - assembly = PluginModelAssembly(tenant_id="tenant1") - configurations = ProviderConfigurations(tenant_id="tenant1") - - def get_providers() -> list[ProviderEntity]: - return provider_entities - - def get_configurations(tenant_id: str) -> ProviderConfigurations: - assert tenant_id == assembly.tenant_id - return configurations - - def get_models( - self: ProviderConfigurations, *, provider: str, model_type: ModelType - ) -> list[ModelWithProviderEntity]: - assert self is configurations - assert provider in [entity.provider for entity in provider_entities] - assert model_type == ModelType.LLM - return model_list - - def create_assembly(*, tenant_id: str) -> PluginModelAssembly: - assert tenant_id == assembly.tenant_id - return assembly - - mocker.patch.object(assembly.model_provider_factory, "get_providers", new=get_providers) - mocker.patch.object(assembly.provider_manager, "get_configurations", new=get_configurations) - mocker.patch.object(ProviderConfigurations, "get_models", new=get_models) - mocker.patch( - "core.app.app_config.easy_ui_based_app.model_config.manager.create_plugin_model_assembly", - new=create_assembly, - ) - return assembly - - # ========================================================== - # convert - # ========================================================== - - def test_convert_success(self, valid_config: AppModelConfigDict) -> None: - result = ModelConfigManager.convert(valid_config) - - assert result.provider == "langgenius/openai/openai" - assert result.model == "gpt-4" - assert result.parameters == {"temperature": 0.5} - assert result.stop == ["END"] - - def test_convert_missing_model(self) -> None: - with pytest.raises(ValueError, match="model is required"): - ModelConfigManager.convert(cast(AppModelConfigDict, {})) - - def test_convert_without_stop(self) -> None: - config = { - "model": { - "provider": "langgenius/openai/openai", - "name": "gpt-4", - "completion_params": {"temperature": 0.9}, - } - } - result = ModelConfigManager.convert(cast(AppModelConfigDict, config)) - assert result.stop == [] - assert result.parameters == {"temperature": 0.9} - - # ========================================================== - # validate_model_completion_params - # ========================================================== - - @pytest.mark.parametrize( - "invalid_cp", - [None, "string", 123, []], - ) - def test_validate_model_completion_params_invalid_type(self, invalid_cp: object) -> None: - with pytest.raises(ValueError, match="must be of object type"): - ModelConfigManager.validate_model_completion_params(invalid_cp) # pyrefly: ignore[bad-argument-type] - - def test_validate_model_completion_params_default_stop(self) -> None: - cp = {"temperature": 0.2} - result = ModelConfigManager.validate_model_completion_params(cp) - assert result["stop"] == [] - - def test_validate_model_completion_params_invalid_stop_type(self) -> None: - cp = {"stop": "invalid"} - with pytest.raises(ValueError, match="must be of list type"): - ModelConfigManager.validate_model_completion_params(cp) - - def test_validate_model_completion_params_stop_length_exceeded(self) -> None: - cp = {"stop": [1, 2, 3, 4, 5]} - with pytest.raises(ValueError, match="less than 4"): - ModelConfigManager.validate_model_completion_params(cp) - - # ========================================================== - # validate_and_set_defaults - # ========================================================== - - def test_validate_and_set_defaults_success( - self, - mocker: MockerFixture, - valid_config: AppModelConfigDict, - provider_entities: list[ProviderEntity], - valid_model_list: list[ModelWithProviderEntity], - ) -> None: - self._patch_model_assembly( - mocker, - provider_entities=provider_entities, - model_list=valid_model_list, - ) - - updated_config, keys = ModelConfigManager.validate_and_set_defaults("tenant1", valid_config) - - assert updated_config["model"]["mode"] == "chat" - assert keys == ["model"] - - def test_validate_and_set_defaults_missing_model(self) -> None: - with pytest.raises(ValueError, match="model is required"): - ModelConfigManager.validate_and_set_defaults("tenant1", {}) - - def test_validate_and_set_defaults_model_not_dict(self) -> None: - with pytest.raises(ValueError, match="object type"): - ModelConfigManager.validate_and_set_defaults("tenant1", {"model": "invalid"}) - - def test_validate_and_set_defaults_missing_provider( - self, mocker: MockerFixture, provider_entities: list[ProviderEntity] - ) -> None: - config = {"model": {"name": "gpt-4", "completion_params": dict[str, object]()}} - self._patch_model_assembly(mocker, provider_entities=provider_entities, model_list=[]) - - with pytest.raises(ValueError, match="model.provider is required"): - ModelConfigManager.validate_and_set_defaults("tenant1", config) - - def test_validate_and_set_defaults_invalid_provider( - self, mocker: MockerFixture, provider_entities: list[ProviderEntity] - ) -> None: - config = {"model": {"provider": "invalid/provider", "name": "gpt-4", "completion_params": dict[str, object]()}} - self._patch_model_assembly(mocker, provider_entities=provider_entities, model_list=[]) - - with pytest.raises(ValueError, match="model.provider is required"): - ModelConfigManager.validate_and_set_defaults("tenant1", config) - - def test_validate_and_set_defaults_missing_name( - self, mocker: MockerFixture, provider_entities: list[ProviderEntity] - ) -> None: - config = {"model": {"provider": "langgenius/openai/openai", "completion_params": dict[str, object]()}} - self._patch_model_assembly(mocker, provider_entities=provider_entities, model_list=[]) - - with pytest.raises(ValueError, match="model.name is required"): - ModelConfigManager.validate_and_set_defaults("tenant1", config) - - def test_validate_and_set_defaults_empty_models( - self, mocker: MockerFixture, provider_entities: list[ProviderEntity] - ) -> None: - config = { - "model": {"provider": "langgenius/openai/openai", "name": "gpt-4", "completion_params": dict[str, object]()} - } - self._patch_model_assembly(mocker, provider_entities=provider_entities, model_list=[]) - - with pytest.raises(ValueError, match="must be in the specified model list"): - ModelConfigManager.validate_and_set_defaults("tenant1", config) - - def test_validate_and_set_defaults_invalid_model_name( - self, - mocker: MockerFixture, - provider_entities: list[ProviderEntity], - valid_model_list: list[ModelWithProviderEntity], - ) -> None: - config = { - "model": { - "provider": "langgenius/openai/openai", - "name": "invalid", - "completion_params": dict[str, object](), - } - } - self._patch_model_assembly( - mocker, - provider_entities=provider_entities, - model_list=valid_model_list, - ) - - with pytest.raises(ValueError, match="must be in the specified model list"): - ModelConfigManager.validate_and_set_defaults("tenant1", config) - - def test_validate_and_set_defaults_default_mode_when_missing( - self, - mocker: MockerFixture, - provider_entities: list[ProviderEntity], - valid_model_list: list[ModelWithProviderEntity], - ) -> None: - model = valid_model_list[0] - model.model_properties = dict[ModelPropertyKey, object]() - - config = { - "model": {"provider": "langgenius/openai/openai", "name": "gpt-4", "completion_params": dict[str, object]()} - } - self._patch_model_assembly(mocker, provider_entities=provider_entities, model_list=[model]) - - updated_config, _ = ModelConfigManager.validate_and_set_defaults("tenant1", config) - - assert updated_config["model"]["mode"] == "completion" - - def test_validate_and_set_defaults_missing_completion_params( - self, - mocker: MockerFixture, - provider_entities: list[ProviderEntity], - valid_model_list: list[ModelWithProviderEntity], - ) -> None: - config = {"model": {"provider": "langgenius/openai/openai", "name": "gpt-4"}} - self._patch_model_assembly( - mocker, - provider_entities=provider_entities, - model_list=valid_model_list, - ) - - with pytest.raises(ValueError, match="completion_params is required"): - ModelConfigManager.validate_and_set_defaults("tenant1", config) - - def test_validate_and_set_defaults_provider_without_slash_converted( - self, - mocker: MockerFixture, - provider_entities: list[ProviderEntity], - valid_model_list: list[ModelWithProviderEntity], - ) -> None: - """ - Covers branch where provider does not contain '/' and - ModelProviderID conversion is triggered (line 64). - """ - config = { - "model": { - "provider": "openai", # no slash -> triggers conversion - "name": "gpt-4", - "completion_params": dict[str, object](), - } - } - - self._patch_model_assembly(mocker, provider_entities=provider_entities, model_list=valid_model_list) - - updated_config, _ = ModelConfigManager.validate_and_set_defaults("tenant1", config) - - assert updated_config["model"]["provider"] == "langgenius/openai/openai" diff --git a/api/tests/unit_tests/core/app/app_config/easy_ui_based_app/test_variables_manager.py b/api/tests/unit_tests/core/app/app_config/easy_ui_based_app/variables/test_manager.py similarity index 100% rename from api/tests/unit_tests/core/app/app_config/easy_ui_based_app/test_variables_manager.py rename to api/tests/unit_tests/core/app/app_config/easy_ui_based_app/variables/test_manager.py diff --git a/api/tests/unit_tests/core/app/app_config/features/test_additional_feature_managers.py b/api/tests/unit_tests/core/app/app_config/features/suggested_questions_after_answer/test_manager.py similarity index 100% rename from api/tests/unit_tests/core/app/app_config/features/test_additional_feature_managers.py rename to api/tests/unit_tests/core/app/app_config/features/suggested_questions_after_answer/test_manager.py diff --git a/api/tests/unit_tests/core/app/app_config/workflow_ui_based_app/test_workflow_ui_based_app_manager.py b/api/tests/unit_tests/core/app/app_config/workflow_ui_based_app/variables/test_manager.py similarity index 100% rename from api/tests/unit_tests/core/app/app_config/workflow_ui_based_app/test_workflow_ui_based_app_manager.py rename to api/tests/unit_tests/core/app/app_config/workflow_ui_based_app/variables/test_manager.py diff --git a/api/tests/unit_tests/core/app/apps/advanced_chat/test_app_config_manager.py b/api/tests/unit_tests/core/app/apps/advanced_chat/test_app_config_manager.py index 8dcd1a995b3..16cb9a949d7 100644 --- a/api/tests/unit_tests/core/app/apps/advanced_chat/test_app_config_manager.py +++ b/api/tests/unit_tests/core/app/apps/advanced_chat/test_app_config_manager.py @@ -3,7 +3,7 @@ from unittest.mock import patch from core.app.apps.advanced_chat.app_config_manager import AdvancedChatAppConfigManager from models.model import App, AppMode from models.workflow import Workflow, WorkflowType -from tests.unit_tests.model_factories import make_app, make_workflow +from tests.helpers.unit.model_factories import make_app, make_workflow def _app() -> App: diff --git a/api/tests/unit_tests/core/app/apps/advanced_chat/test_app_generator.py b/api/tests/unit_tests/core/app/apps/advanced_chat/test_app_generator.py index a1c4bdfb5d2..9663f197699 100644 --- a/api/tests/unit_tests/core/app/apps/advanced_chat/test_app_generator.py +++ b/api/tests/unit_tests/core/app/apps/advanced_chat/test_app_generator.py @@ -2,18 +2,28 @@ from __future__ import annotations import json import logging +import sys +import time +from collections.abc import Mapping from contextlib import contextmanager from decimal import Decimal -from types import SimpleNamespace +from types import ModuleType, SimpleNamespace +from typing import Any from unittest.mock import MagicMock import pytest from pydantic import BaseModel, ValidationError -from sqlalchemy import Engine, event +from pytest_mock import MockerFixture +from sqlalchemy import Engine as _app_generator_Engine +from sqlalchemy import event as sqlalchemy_event +from sqlalchemy.engine import Engine as _advanced_chat_app_generator_Engine from sqlalchemy.orm import Session +import core.app.apps.advanced_chat.app_generator as app_generator_module +import core.workflow.nodes.human_input.entities # noqa: F401 from constants import UUID_NIL from core.app.app_config.entities import AppAdditionalFeatures, WorkflowUIBasedAppConfig +from core.app.apps.advanced_chat import app_generator as adv_app_gen_module from core.app.apps.advanced_chat.app_generator import AdvancedChatAppGenerator from core.app.apps.advanced_chat.generate_task_pipeline import ( ConversationSnapshot, @@ -21,18 +31,51 @@ from core.app.apps.advanced_chat.generate_task_pipeline import ( WorkflowSnapshot, ) from core.app.apps.exc import GenerateTaskStoppedError +from core.app.apps.workflow import app_generator as wf_app_gen_module from core.app.entities.app_invoke_entities import AdvancedChatAppGenerateEntity, InvokeFrom +from core.app.task_pipeline import message_cycle_manager +from core.app.task_pipeline.message_cycle_manager import MessageCycleManager from core.ops.ops_trace_manager import TraceQueueManager +from core.workflow import node_factory as node_factory_module +from core.workflow.node_factory import DifyNodeFactory +from core.workflow.system_variables import build_system_variables +from graphon.engine import Engine as _pause_resume_Engine +from graphon.engine.command import InMemoryChannel +from graphon.engine_events import ( + EngineEvent, + GraphRunPausedEvent, + GraphRunStartedEvent, + GraphRunSucceededEvent, + NodeRunSucceededEvent, +) +from graphon.entities import WorkflowStartReason +from graphon.entities.base_node_data import BaseNodeData, RetryConfig +from graphon.entities.pause_reason import SchedulingPause +from graphon.enums import BuiltinNodeTypes, NodeType, WorkflowNodeExecutionStatus +from graphon.graph import Graph +from graphon.node_events import NodeRunResult, PauseRequestedEvent +from graphon.nodes.base.entities import OutputVariableEntity +from graphon.nodes.base.node import Node +from graphon.nodes.end.entities import EndNodeData +from graphon.nodes.start.entities import StartNodeData +from graphon.runtime import RuntimeState, VariablePool from libs.datetime_utils import naive_utc_now -from models.account import Account +from models import Account as _advanced_chat_app_generator_Account +from models import Workflow as _advanced_chat_app_generator_Workflow +from models.account import Account as _app_generator_Account +from models.account import Account as _pause_resume_Account from models.enums import ConversationFromSource, EndUserType, MessageStatus from models.model import App, AppMode, Conversation, EndUser, Message -from models.workflow import Workflow, WorkflowType +from models.workflow import Workflow as _app_generator_Workflow +from models.workflow import Workflow as _pause_resume_Workflow +from models.workflow import WorkflowType +from services.errors.conversation import ConversationNotExistsError from services.workflow_run_agg import WorkflowRunAgg -from tests.unit_tests.config_override import apply_config_overrides +from tests.helpers.unit.config_override import apply_config_overrides +from tests.workflow_test_utils import build_test_graph_init_params -def _make_app(*, app_id: str = "app", tenant_id: str = "tenant") -> App: +def _app_generator_make_app(*, app_id: str = "app", tenant_id: str = "tenant") -> App: return App( id=app_id, tenant_id=tenant_id, @@ -43,27 +86,27 @@ def _make_app(*, app_id: str = "app", tenant_id: str = "tenant") -> App: ) -def _make_workflow( +def _app_generator_make_workflow( *, workflow_id: str = "workflow-id", tenant_id: str = "tenant", app_id: str = "app", features: dict[str, object] | None = None, -) -> Workflow: - return Workflow( +) -> _app_generator_Workflow: + return _app_generator_Workflow( id=workflow_id, tenant_id=tenant_id, app_id=app_id, type=WorkflowType.CHAT, - version=Workflow.VERSION_DRAFT, + version=_app_generator_Workflow.VERSION_DRAFT, graph="{}", features=json.dumps(features or {}), created_by="user", ) -def _make_account(*, account_id: str = "user-id") -> Account: - account = Account(name="Advanced Chat User", email=f"{account_id}@example.com") +def _app_generator_make_account(*, account_id: str = "user-id") -> _app_generator_Account: + account = _app_generator_Account(name="Advanced Chat User", email=f"{account_id}@example.com") account.id = account_id return account @@ -116,9 +159,9 @@ class TestAdvancedChatAppGeneratorValidation: with pytest.raises(ValueError, match="query is required"): generator.generate( - app_model=_make_app(), - workflow=_make_workflow(), - user=_make_account(), + app_model=_app_generator_make_app(), + workflow=_app_generator_make_workflow(), + user=_app_generator_make_account(), args={"inputs": {}}, invoke_from=InvokeFrom.WEB_APP, workflow_run_id="run-id", @@ -131,9 +174,9 @@ class TestAdvancedChatAppGeneratorValidation: with pytest.raises(ValueError, match="query must be a string"): generator.generate( - app_model=_make_app(), - workflow=_make_workflow(), - user=_make_account(), + app_model=_app_generator_make_app(), + workflow=_app_generator_make_workflow(), + user=_app_generator_make_account(), args={"inputs": {}, "query": 123}, invoke_from=InvokeFrom.WEB_APP, workflow_run_id="run-id", @@ -146,10 +189,10 @@ class TestAdvancedChatAppGeneratorValidation: with pytest.raises(ValueError, match="node_id is required"): generator.single_iteration_generate( - app_model=_make_app(), - workflow=_make_workflow(), + app_model=_app_generator_make_app(), + workflow=_app_generator_make_workflow(), node_id="", - user=_make_account(), + user=_app_generator_make_account(), args={"inputs": {}}, streaming=False, session=unbound_session, @@ -157,10 +200,10 @@ class TestAdvancedChatAppGeneratorValidation: with pytest.raises(ValueError, match="inputs is required"): generator.single_iteration_generate( - app_model=_make_app(), - workflow=_make_workflow(), + app_model=_app_generator_make_app(), + workflow=_app_generator_make_workflow(), node_id="node", - user=_make_account(), + user=_app_generator_make_account(), args={}, streaming=False, session=unbound_session, @@ -171,10 +214,10 @@ class TestAdvancedChatAppGeneratorValidation: with pytest.raises(ValueError, match="node_id is required"): generator.single_loop_generate( - app_model=_make_app(), - workflow=_make_workflow(), + app_model=_app_generator_make_app(), + workflow=_app_generator_make_workflow(), node_id="", - user=_make_account(), + user=_app_generator_make_account(), args=SimpleNamespace(inputs={}), streaming=False, session=unbound_session, @@ -182,10 +225,10 @@ class TestAdvancedChatAppGeneratorValidation: with pytest.raises(ValueError, match="inputs is required"): generator.single_loop_generate( - app_model=_make_app(), - workflow=_make_workflow(), + app_model=_app_generator_make_app(), + workflow=_app_generator_make_workflow(), node_id="node", - user=_make_account(), + user=_app_generator_make_account(), args=SimpleNamespace(inputs=None), streaming=False, session=unbound_session, @@ -205,7 +248,7 @@ class TestAdvancedChatAppGeneratorInternals: ) def test_generate_loads_conversation_and_files( - self, monkeypatch: pytest.MonkeyPatch, unbound_session: Session, sqlite_engine: Engine + self, monkeypatch: pytest.MonkeyPatch, unbound_session: Session, sqlite_engine: _app_generator_Engine ): generator = AdvancedChatAppGenerator(execution_driver=WorkflowRunAgg.run) app_config = self._build_app_config() @@ -270,8 +313,8 @@ class TestAdvancedChatAppGeneratorInternals: user.id = "user-id" result = generator.generate( - app_model=_make_app(), - workflow=_make_workflow(), + app_model=_app_generator_make_app(), + workflow=_app_generator_make_workflow(), user=user, args={ "query": "hello", @@ -321,8 +364,8 @@ class TestAdvancedChatAppGeneratorInternals: monkeypatch.setattr(generator, "_generate", _fake_generate) result = generator.resume( - app_model=_make_app(app_id="app-id"), - workflow=_make_workflow(), + app_model=_app_generator_make_app(app_id="app-id"), + workflow=_app_generator_make_workflow(), user=_make_end_user(), conversation=_make_conversation(), message=_make_message(), @@ -340,7 +383,7 @@ class TestAdvancedChatAppGeneratorInternals: assert captured_graph_runtime_state is not None def test_single_iteration_generate_builds_debug_task( - self, monkeypatch: pytest.MonkeyPatch, unbound_session: Session, sqlite_engine: Engine + self, monkeypatch: pytest.MonkeyPatch, unbound_session: Session, sqlite_engine: _app_generator_Engine ): generator = AdvancedChatAppGenerator(execution_driver=WorkflowRunAgg.run) app_config = self._build_app_config() @@ -348,7 +391,7 @@ class TestAdvancedChatAppGeneratorInternals: prefill_calls: list[object] = [] draft_sessions: list[object] = [] var_loader = SimpleNamespace(loader="draft") - workflow = _make_workflow() + workflow = _app_generator_make_workflow() session = unbound_session monkeypatch.setattr( @@ -385,10 +428,10 @@ class TestAdvancedChatAppGeneratorInternals: monkeypatch.setattr(generator, "_generate", _fake_generate) result = generator.single_iteration_generate( - app_model=_make_app(), + app_model=_app_generator_make_app(), workflow=workflow, node_id="node-1", - user=_make_account(), + user=_app_generator_make_account(), args={"inputs": {"foo": "bar"}, "trace_session_id": "session-1"}, streaming=False, session=session, @@ -403,7 +446,7 @@ class TestAdvancedChatAppGeneratorInternals: assert captured["application_generate_entity"].extras["trace_session_id"] == "session-1" def test_single_loop_generate_builds_debug_task( - self, monkeypatch: pytest.MonkeyPatch, unbound_session: Session, sqlite_engine: Engine + self, monkeypatch: pytest.MonkeyPatch, unbound_session: Session, sqlite_engine: _app_generator_Engine ): generator = AdvancedChatAppGenerator(execution_driver=WorkflowRunAgg.run) app_config = self._build_app_config() @@ -411,7 +454,7 @@ class TestAdvancedChatAppGeneratorInternals: prefill_calls: list[object] = [] draft_sessions: list[object] = [] var_loader = SimpleNamespace(loader="draft") - workflow = _make_workflow() + workflow = _app_generator_make_workflow() session = unbound_session monkeypatch.setattr( @@ -448,10 +491,10 @@ class TestAdvancedChatAppGeneratorInternals: monkeypatch.setattr(generator, "_generate", _fake_generate) result = generator.single_loop_generate( - app_model=_make_app(), + app_model=_app_generator_make_app(), workflow=workflow, node_id="node-2", - user=_make_account(), + user=_app_generator_make_account(), args=SimpleNamespace(inputs={"foo": "bar"}, trace_session_id="session-1"), streaming=False, session=session, @@ -466,7 +509,7 @@ class TestAdvancedChatAppGeneratorInternals: assert captured["application_generate_entity"].extras["trace_session_id"] == "session-1" def test_generate_internal_flow_initial_conversation_with_pause_layer( - self, monkeypatch: pytest.MonkeyPatch, sqlite_session: Session, sqlite_engine: Engine + self, monkeypatch: pytest.MonkeyPatch, sqlite_session: Session, sqlite_engine: _app_generator_Engine ): generator = AdvancedChatAppGenerator(execution_driver=WorkflowRunAgg.run) generator._dialogue_count = 0 @@ -486,8 +529,8 @@ class TestAdvancedChatAppGeneratorInternals: workflow_run_id="run-id", ) - app = _make_app() - workflow = _make_workflow(workflow_id="wf-1", features={"feature": True}) + app = _app_generator_make_app() + workflow = _app_generator_make_workflow(workflow_id="wf-1", features={"feature": True}) conversation = _make_conversation(conversation_id="conv-1") message = _make_message(message_id="msg-1", conversation_id=conversation.id) sqlite_session.add_all([app, workflow, conversation, message]) @@ -498,7 +541,7 @@ class TestAdvancedChatAppGeneratorInternals: nonlocal commit_count commit_count += 1 - event.listen(sqlite_session, "after_commit", _record_commit) + sqlalchemy_event.listen(sqlite_session, "after_commit", _record_commit) captured: dict[str, object] = {} thread_data: dict[str, object] = {} init_records = MagicMock(return_value=(conversation, message)) @@ -553,7 +596,7 @@ class TestAdvancedChatAppGeneratorInternals: response = generator._generate( workflow=workflow, - user=_make_account(account_id="user"), + user=_app_generator_make_account(account_id="user"), invoke_from=InvokeFrom.WEB_APP, application_generate_entity=application_generate_entity, session=sqlite_session, @@ -581,7 +624,7 @@ class TestAdvancedChatAppGeneratorInternals: assert isinstance(captured["message"], MessageSnapshot) def test_generate_internal_flow_with_existing_records_skips_init( - self, monkeypatch: pytest.MonkeyPatch, sqlite_session: Session, sqlite_engine: Engine + self, monkeypatch: pytest.MonkeyPatch, sqlite_session: Session, sqlite_engine: _app_generator_Engine ): generator = AdvancedChatAppGenerator(execution_driver=WorkflowRunAgg.run) generator._dialogue_count = 0 @@ -601,8 +644,8 @@ class TestAdvancedChatAppGeneratorInternals: workflow_run_id="run-id", ) - app = _make_app() - workflow = _make_workflow(workflow_id="wf-2") + app = _app_generator_make_app() + workflow = _app_generator_make_workflow(workflow_id="wf-2") conversation = _make_conversation(conversation_id="conv-2") message = _make_message(message_id="msg-2", conversation_id=conversation.id) sqlite_session.add_all([app, workflow, conversation, message]) @@ -613,7 +656,7 @@ class TestAdvancedChatAppGeneratorInternals: nonlocal commit_count commit_count += 1 - event.listen(sqlite_session, "after_commit", _record_commit) + sqlalchemy_event.listen(sqlite_session, "after_commit", _record_commit) init_records = MagicMock() get_thread_messages_length = MagicMock(return_value=0) thread_data: dict[str, object] = {} @@ -665,7 +708,7 @@ class TestAdvancedChatAppGeneratorInternals: response = generator._generate( workflow=workflow, - user=_make_account(account_id="user"), + user=_app_generator_make_account(account_id="user"), invoke_from=InvokeFrom.WEB_APP, application_generate_entity=application_generate_entity, session=sqlite_session, @@ -685,7 +728,7 @@ class TestAdvancedChatAppGeneratorInternals: assert commit_count == 0 def test_generate_worker_raises_when_workflow_not_found( - self, monkeypatch: pytest.MonkeyPatch, sqlite_session: Session, sqlite_engine: Engine + self, monkeypatch: pytest.MonkeyPatch, sqlite_session: Session, sqlite_engine: _app_generator_Engine ): generator = AdvancedChatAppGenerator(execution_driver=WorkflowRunAgg.run) generator._dialogue_count = 1 @@ -734,7 +777,7 @@ class TestAdvancedChatAppGeneratorInternals: ) def test_generate_worker_raises_when_app_not_found_for_internal_call( - self, monkeypatch: pytest.MonkeyPatch, sqlite_session: Session, sqlite_engine: Engine + self, monkeypatch: pytest.MonkeyPatch, sqlite_session: Session, sqlite_engine: _app_generator_Engine ): generator = AdvancedChatAppGenerator(execution_driver=WorkflowRunAgg.run) generator._dialogue_count = 1 @@ -756,7 +799,7 @@ class TestAdvancedChatAppGeneratorInternals: generator._get_conversation = MagicMock(return_value=_make_conversation(conversation_id="conv")) generator._get_message = MagicMock(return_value=_make_message(message_id="msg", conversation_id="conv")) - sqlite_session.add(_make_workflow()) + sqlite_session.add(_app_generator_make_workflow()) sqlite_session.commit() @contextmanager @@ -785,7 +828,7 @@ class TestAdvancedChatAppGeneratorInternals: ) def test_generate_worker_handles_stopped_error( - self, monkeypatch: pytest.MonkeyPatch, sqlite_session: Session, sqlite_engine: Engine + self, monkeypatch: pytest.MonkeyPatch, sqlite_session: Session, sqlite_engine: _app_generator_Engine ): generator = AdvancedChatAppGenerator(execution_driver=WorkflowRunAgg.run) generator._dialogue_count = 1 @@ -815,7 +858,7 @@ class TestAdvancedChatAppGeneratorInternals: monkeypatch.setattr("core.app.apps.advanced_chat.app_generator.preserve_flask_contexts", _fake_context) - sqlite_session.add_all([_make_app(), _make_workflow()]) + sqlite_session.add_all([_app_generator_make_app(), _app_generator_make_workflow()]) sqlite_session.commit() class _Runner: @@ -851,7 +894,7 @@ class TestAdvancedChatAppGeneratorInternals: assert restore_workflow_run_graph.call_args.kwargs["workflow_run_id"] == "run-id" def test_generate_worker_handles_validation_error( - self, monkeypatch: pytest.MonkeyPatch, sqlite_session: Session, sqlite_engine: Engine + self, monkeypatch: pytest.MonkeyPatch, sqlite_session: Session, sqlite_engine: _app_generator_Engine ): generator = AdvancedChatAppGenerator(execution_driver=WorkflowRunAgg.run) generator._dialogue_count = 1 @@ -884,7 +927,7 @@ class TestAdvancedChatAppGeneratorInternals: queue_manager = MagicMock() generator._get_conversation = MagicMock(return_value=_make_conversation(conversation_id="conv")) generator._get_message = MagicMock(return_value=_make_message(message_id="msg", conversation_id="conv")) - sqlite_session.add_all([_make_app(), _make_workflow()]) + sqlite_session.add_all([_app_generator_make_app(), _app_generator_make_workflow()]) sqlite_session.commit() @contextmanager @@ -922,10 +965,10 @@ class TestAdvancedChatAppGeneratorInternals: queue_manager.publish_error.assert_called_once() def test_generate_worker_handles_value_and_unknown_errors( - self, monkeypatch: pytest.MonkeyPatch, sqlite_session: Session, sqlite_engine: Engine + self, monkeypatch: pytest.MonkeyPatch, sqlite_session: Session, sqlite_engine: _app_generator_Engine ): app_config = self._build_app_config() - sqlite_session.add_all([_make_app(), _make_workflow()]) + sqlite_session.add_all([_app_generator_make_app(), _app_generator_make_workflow()]) sqlite_session.commit() @contextmanager @@ -1040,7 +1083,7 @@ class TestAdvancedChatAppGeneratorInternals: status=MessageStatus.NORMAL, answer="", ), - user=_make_account(), + user=_app_generator_make_account(), draft_var_saver_factory=lambda **kwargs: None, stream=False, ) @@ -1088,7 +1131,7 @@ class TestAdvancedChatAppGeneratorInternals: status=MessageStatus.NORMAL, answer="", ), - user=_make_account(), + user=_app_generator_make_account(), draft_var_saver_factory=lambda **kwargs: None, stream=False, ) @@ -1096,7 +1139,7 @@ class TestAdvancedChatAppGeneratorInternals: assert "Failed to process generate task pipeline, conversation_id: conv" in caplog.messages def test_generate_worker_handles_invoke_auth_error( - self, monkeypatch: pytest.MonkeyPatch, sqlite_session: Session, sqlite_engine: Engine + self, monkeypatch: pytest.MonkeyPatch, sqlite_session: Session, sqlite_engine: _app_generator_Engine ): generator = AdvancedChatAppGenerator(execution_driver=WorkflowRunAgg.run) generator._dialogue_count = 1 @@ -1127,7 +1170,7 @@ class TestAdvancedChatAppGeneratorInternals: generator._get_conversation = MagicMock(return_value=_make_conversation(conversation_id="conv")) generator._get_message = MagicMock(return_value=_make_message(message_id="msg", conversation_id="conv")) - sqlite_session.add_all([_make_app(), _make_workflow(), _make_end_user()]) + sqlite_session.add_all([_app_generator_make_app(), _app_generator_make_workflow(), _make_end_user()]) sqlite_session.commit() class _Runner: @@ -1168,7 +1211,7 @@ class TestAdvancedChatAppGeneratorInternals: assert queue_manager.publish_error.called def test_generate_debugger_enables_retrieve_source( - self, monkeypatch: pytest.MonkeyPatch, unbound_session: Session, sqlite_engine: Engine + self, monkeypatch: pytest.MonkeyPatch, unbound_session: Session, sqlite_engine: _app_generator_Engine ): generator = AdvancedChatAppGenerator(execution_driver=WorkflowRunAgg.run) @@ -1223,9 +1266,9 @@ class TestAdvancedChatAppGeneratorInternals: monkeypatch.setattr(generator, "_generate", _fake_generate) - app_model = _make_app() - workflow = _make_workflow() - user = _make_account(account_id="user") + app_model = _app_generator_make_app() + workflow = _app_generator_make_workflow() + user = _app_generator_make_account(account_id="user") result = generator.generate( app_model=app_model, @@ -1243,7 +1286,7 @@ class TestAdvancedChatAppGeneratorInternals: assert captured["application_generate_entity"].query == "hello" def test_generate_service_api_sets_parent_message_id( - self, monkeypatch: pytest.MonkeyPatch, unbound_session: Session, sqlite_engine: Engine + self, monkeypatch: pytest.MonkeyPatch, unbound_session: Session, sqlite_engine: _app_generator_Engine ): generator = AdvancedChatAppGenerator(execution_driver=WorkflowRunAgg.run) @@ -1298,8 +1341,8 @@ class TestAdvancedChatAppGeneratorInternals: monkeypatch.setattr(generator, "_generate", _fake_generate) - app_model = _make_app() - workflow = _make_workflow() + app_model = _app_generator_make_app() + workflow = _app_generator_make_workflow() user = _make_end_user(end_user_id="end-user", session_id="session") generator.generate( @@ -1371,8 +1414,8 @@ class TestAdvancedChatAppGeneratorResume: monkeypatch.setattr(generator, "_generate", _fake_generate) result = generator.resume( - app_model=_make_app(app_id="app-id"), - workflow=_make_workflow(), + app_model=_app_generator_make_app(app_id="app-id"), + workflow=_app_generator_make_workflow(), user=_make_end_user(), conversation=_make_conversation(), message=_make_message(), @@ -1419,8 +1462,8 @@ class TestAdvancedChatAppGeneratorResume: monkeypatch.setattr(generator, "_generate", _fake_generate) result = generator.resume( - app_model=_make_app(app_id="app-id"), - workflow=_make_workflow(), + app_model=_app_generator_make_app(app_id="app-id"), + workflow=_app_generator_make_workflow(), user=_make_end_user(), conversation=_make_conversation(), message=_make_message(), @@ -1434,3 +1477,555 @@ class TestAdvancedChatAppGeneratorResume: assert result.ok is True assert captured_entity is not None assert captured_entity.trace_manager is existing_trace_manager + + +def _make_app_config() -> WorkflowUIBasedAppConfig: + return WorkflowUIBasedAppConfig( + tenant_id="tenant-id", + app_id="app-id", + app_mode=AppMode.ADVANCED_CHAT, + workflow_id="workflow-id", + additional_features=AppAdditionalFeatures(), + variables=[], + ) + + +def _make_generate_entity(app_config: WorkflowUIBasedAppConfig) -> AdvancedChatAppGenerateEntity: + return AdvancedChatAppGenerateEntity( + task_id="task-id", + app_config=app_config, + file_upload_config=None, + conversation_id=None, + inputs={}, + query="hello", + files=[], + parent_message_id=None, + user_id="user-id", + stream=True, + invoke_from=InvokeFrom.WEB_APP, + extras={}, + workflow_run_id="workflow-run-id", + ) + + +def _app() -> App: + return App( + id="app-id", + tenant_id="tenant-id", + name="Advanced chat app", + description="", + mode=AppMode.ADVANCED_CHAT, + enable_site=False, + enable_api=False, + ) + + +def _workflow() -> _advanced_chat_app_generator_Workflow: + return _advanced_chat_app_generator_Workflow.new( + tenant_id="tenant-id", + app_id="app-id", + type="chat", + version=_advanced_chat_app_generator_Workflow.VERSION_DRAFT, + graph='{"nodes": [], "edges": []}', + features="{}", + created_by="user-id", + environment_variables=[], + conversation_variables=[], + rag_pipeline_variables=[], + ) + + +def _account() -> _advanced_chat_app_generator_Account: + account = _advanced_chat_app_generator_Account(name="User", email="user@example.com") + account.id = "user-id" + return account + + +def test_init_generate_records_sets_conversation_metadata(sqlite_session: Session): + app_config = _make_app_config() + entity = _make_generate_entity(app_config) + + generator = AdvancedChatAppGenerator(execution_driver=WorkflowRunAgg.run) + + conversation, _ = generator._init_generate_records( + entity, + conversation=None, + session=sqlite_session, + ) + + assert entity.conversation_id == conversation.id + assert conversation.id is not None + assert entity.is_new_conversation is True + + +def test_init_generate_records_marks_existing_conversation(sqlite_session: Session): + app_config = _make_app_config() + entity = _make_generate_entity(app_config) + + existing_conversation = Conversation( + app_id=app_config.app_id, + app_model_config_id=None, + model_provider=None, + override_model_configs=None, + model_id=None, + mode=app_config.app_mode.value, + name="existing", + _inputs={}, + introduction="", + system_instruction="", + system_instruction_tokens=0, + status="normal", + invoke_from=InvokeFrom.WEB_APP.value, + from_source=ConversationFromSource.API, + from_end_user_id="user-id", + from_account_id=None, + ) + existing_conversation.id = "existing-conversation-id" + sqlite_session.add(existing_conversation) + sqlite_session.flush() + + generator = AdvancedChatAppGenerator(execution_driver=WorkflowRunAgg.run) + + conversation, _ = generator._init_generate_records( + entity, + conversation=existing_conversation, + session=sqlite_session, + ) + + assert entity.conversation_id == "existing-conversation-id" + assert conversation is existing_conversation + assert entity.is_new_conversation is False + + +def test_generate_falls_back_to_new_conversation_when_conversation_missing( + monkeypatch: pytest.MonkeyPatch, + sqlite_engine: _advanced_chat_app_generator_Engine, + sqlite_session: Session, +): + app_config = _make_app_config() + workflow = _workflow() + app_model = _app() + user = _account() + + def raise_conversation_not_exists(**_kwargs): + raise ConversationNotExistsError() + + monkeypatch.setattr( + "core.app.apps.advanced_chat.app_generator.ConversationService.get_conversation", + raise_conversation_not_exists, + ) + monkeypatch.setattr( + "core.app.apps.advanced_chat.app_generator.FileUploadConfigManager.convert", + lambda *_args, **_kwargs: None, + ) + monkeypatch.setattr( + "core.app.apps.advanced_chat.app_generator.AdvancedChatAppConfigManager.get_app_config", + lambda **_kwargs: app_config, + ) + monkeypatch.setattr(app_generator_module, "db", SimpleNamespace(engine=sqlite_engine)) + trace_manager = object.__new__(TraceQueueManager) + monkeypatch.setattr( + "core.app.apps.advanced_chat.app_generator.TraceQueueManager", + lambda **_kwargs: trace_manager, + ) + monkeypatch.setattr( + "core.app.apps.advanced_chat.app_generator.DifyCoreRepositoryFactory.create_workflow_execution_repository", + lambda **_kwargs: SimpleNamespace(), + ) + monkeypatch.setattr( + "core.app.apps.advanced_chat.app_generator.DifyCoreRepositoryFactory.create_workflow_node_execution_repository", + lambda **_kwargs: SimpleNamespace(), + ) + + captured: dict[str, object] = {} + session = sqlite_session + + def fake_generate(self, **kwargs): + captured.update(kwargs) + return {"status": "ok"} + + monkeypatch.setattr(AdvancedChatAppGenerator, "_generate", fake_generate) + + result = AdvancedChatAppGenerator(execution_driver=WorkflowRunAgg.run).generate( + app_model=app_model, + workflow=workflow, + user=user, + args={"inputs": {}, "query": "hello", "conversation_id": "missing-conversation-id"}, + invoke_from=InvokeFrom.SERVICE_API, + workflow_run_id="workflow-run-id", + streaming=False, + session=session, + ) + + assert result == {"status": "ok"} + assert captured["conversation"] is None + assert captured["session"] is session + application_generate_entity = captured["application_generate_entity"] + assert isinstance(application_generate_entity, AdvancedChatAppGenerateEntity) + assert application_generate_entity.conversation_id is None + + +def test_message_cycle_manager_uses_new_conversation_flag(monkeypatch: pytest.MonkeyPatch): + app_config = _make_app_config() + entity = _make_generate_entity(app_config) + entity.conversation_id = "existing-conversation-id" + entity.is_new_conversation = True + entity.extras = {"auto_generate_conversation_name": True} + + captured = {} + + class DummyThread: + def __init__(self, **kwargs): + self.kwargs = kwargs + self.started = False + + def start(self): + self.started = True + + def fake_thread[**P](*args: P.args, **kwargs: P.kwargs): + thread = DummyThread(**kwargs) + captured["thread"] = thread + return thread + + monkeypatch.setattr(message_cycle_manager, "Timer", fake_thread) + + manager = MessageCycleManager(application_generate_entity=entity, task_state=MagicMock()) + thread = manager.generate_conversation_name(conversation_id="existing-conversation-id", query="hello") + + assert thread is captured["thread"] + assert thread.started is True + assert entity.is_new_conversation is False + + +if "core.ops.ops_trace_manager" not in sys.modules: + ops_stub = ModuleType("core.ops.ops_trace_manager") + + class _StubTraceQueueManager: + def __init__(self, *_, **__): + pass + + ops_stub.TraceQueueManager = _StubTraceQueueManager + sys.modules["core.ops.ops_trace_manager"] = ops_stub + + +class _StubToolNodeData(BaseNodeData): + type: NodeType = BuiltinNodeTypes.TOOL + pause_on: bool = False + + +class _StubToolNode(Node[_StubToolNodeData]): + node_type = BuiltinNodeTypes.TOOL + + @classmethod + def version(cls) -> str: + return "1" + + def __init__( + self, + node_id: str, + data: _StubToolNodeData, + *, + init_params, + runtime_state, + **_kwargs: Any, + ) -> None: + super().__init__( + node_id=node_id, + data=data, + init_params=init_params, + runtime_state=runtime_state, + ) + + def _get_error_strategy(self): + return self._node_data.error_strategy + + def _get_retry_config(self) -> RetryConfig: + return self._node_data.retry_config + + def _get_title(self) -> str: + return self._node_data.title + + def _get_description(self): + return self._node_data.desc + + def _get_default_value_dict(self) -> dict[str, Any]: + return self._node_data.default_value_dict + + def get_base_node_data(self) -> BaseNodeData: + return self._node_data + + def _run(self): + if self.node_data.pause_on: + yield PauseRequestedEvent(reason=SchedulingPause(message="test pause")) + return + + result = NodeRunResult( + status=WorkflowNodeExecutionStatus.SUCCEEDED, + outputs={"value": f"{self.id}-done"}, + ) + yield self._convert_node_run_result_to_node_event(result) + + +def _patch_tool_node(mocker: MockerFixture): + original_resolve_node_class = node_factory_module.resolve_workflow_node_class + + def _patched_resolve_node_class( + *, + node_type: NodeType, + node_version: str, + node_data: Mapping[str, Any] | BaseNodeData | None = None, + ) -> type[Node]: + if node_type == BuiltinNodeTypes.TOOL: + return _StubToolNode + return original_resolve_node_class( + node_type=node_type, + node_version=node_version, + node_data=node_data, + ) + + mocker.patch.object(node_factory_module, "resolve_workflow_node_class", side_effect=_patched_resolve_node_class) + + +def _node_data(node_type: NodeType, data: BaseNodeData) -> dict[str, object]: + node_data = data.model_dump() + node_data["type"] = str(node_type) + return node_data + + +def _build_graph_config(*, pause_on: str | None) -> dict[str, object]: + start_data = StartNodeData(title="start", variables=[]) + tool_data_a = _StubToolNodeData(title="tool", pause_on=pause_on == "tool_a") + tool_data_b = _StubToolNodeData(title="tool", pause_on=pause_on == "tool_b") + tool_data_c = _StubToolNodeData(title="tool", pause_on=pause_on == "tool_c") + end_data = EndNodeData( + title="end", + outputs=[OutputVariableEntity(variable="result", value_selector=["tool_c", "value"])], + desc=None, + ) + + nodes = [ + {"id": "start", "data": _node_data(BuiltinNodeTypes.START, start_data)}, + {"id": "tool_a", "data": _node_data(BuiltinNodeTypes.TOOL, tool_data_a)}, + {"id": "tool_b", "data": _node_data(BuiltinNodeTypes.TOOL, tool_data_b)}, + {"id": "tool_c", "data": _node_data(BuiltinNodeTypes.TOOL, tool_data_c)}, + {"id": "end", "data": _node_data(BuiltinNodeTypes.END, end_data)}, + ] + edges = [ + {"source": "start", "target": "tool_a"}, + {"source": "tool_a", "target": "tool_b"}, + {"source": "tool_b", "target": "tool_c"}, + {"source": "tool_c", "target": "end"}, + ] + return {"nodes": nodes, "edges": edges} + + +def _build_graph(runtime_state: RuntimeState, *, pause_on: str | None) -> Graph: + graph_config = _build_graph_config(pause_on=pause_on) + params = build_test_graph_init_params( + workflow_id="workflow", + graph_config=graph_config, + tenant_id="tenant", + app_id="app", + user_id="user", + user_from="account", + invoke_from="service-api", + call_depth=0, + ) + + node_factory = DifyNodeFactory( + init_params=params, + runtime_state=runtime_state, + ) + + return Graph.init(graph_config=graph_config, node_factory=node_factory, root_node_id="start") + + +def _build_runtime_state(run_id: str) -> RuntimeState: + variable_pool = VariablePool.from_bootstrap( + system_variables=build_system_variables(user_id="user", app_id="app", workflow_id="workflow"), + user_inputs={}, + conversation_variables=[], + ) + variable_pool.add(("sys", "workflow_run_id"), run_id) + return RuntimeState(workflow_id="test-workflow", variable_pool=variable_pool, start_at=time.perf_counter()) + + +def _pause_resume_make_app() -> App: + return App( + id="app", + tenant_id="tenant", + name="Pause Resume App", + mode=AppMode.WORKFLOW, + enable_site=False, + enable_api=False, + ) + + +def _pause_resume_make_workflow() -> _pause_resume_Workflow: + return _pause_resume_Workflow( + id="workflow", + tenant_id="tenant", + app_id="app", + type=WorkflowType.CHAT, + version=_pause_resume_Workflow.VERSION_DRAFT, + graph="{}", + features="{}", + created_by="user", + ) + + +def _pause_resume_make_account() -> _pause_resume_Account: + account = _pause_resume_Account(name="Pause Resume User", email="pause@example.com") + account.id = "user" + return account + + +def _pause_resume_make_conversation() -> Conversation: + return Conversation( + id="conv", + app_id="app", + mode=AppMode.ADVANCED_CHAT, + name="Pause Resume Conversation", + _inputs={}, + from_source=ConversationFromSource.API, + ) + + +def _pause_resume_make_message() -> Message: + return Message( + id="msg", + app_id="app", + conversation_id="conv", + _inputs={}, + query="query", + message={}, + answer="answer", + message_unit_price=Decimal(0), + answer_unit_price=Decimal(0), + currency="USD", + from_source=ConversationFromSource.API, + ) + + +def _run_with_optional_pause(runtime_state: RuntimeState, *, pause_on: str | None) -> list[EngineEvent]: + command_channel = InMemoryChannel() + graph = _build_graph(runtime_state, pause_on=pause_on) + engine = _pause_resume_Engine( + graph=graph, + runtime_state=runtime_state, + command_channel=command_channel, + ) + + events: list[EngineEvent] = [] + for event in engine.run(): + events.append(event) + return events + + +def _node_successes(events: list[EngineEvent]) -> list[str]: + return [evt.node_id for evt in events if isinstance(evt, NodeRunSucceededEvent)] + + +def test_workflow_app_pause_resume_matches_baseline(mocker: MockerFixture): + _patch_tool_node(mocker) + + baseline_state = _build_runtime_state("baseline") + baseline_events = _run_with_optional_pause(baseline_state, pause_on=None) + assert isinstance(baseline_events[-1], GraphRunSucceededEvent) + baseline_nodes = _node_successes(baseline_events) + baseline_outputs = baseline_state.outputs + + paused_state = _build_runtime_state("paused-run") + paused_events = _run_with_optional_pause(paused_state, pause_on="tool_a") + assert isinstance(paused_events[-1], GraphRunPausedEvent) + paused_nodes = _node_successes(paused_events) + snapshot = paused_state.dumps() + + resumed_state = RuntimeState.from_snapshot(snapshot) + + generator = wf_app_gen_module.WorkflowAppGenerator(execution_driver=WorkflowRunAgg.run) + + def _fake_generate(**kwargs): + state: RuntimeState = kwargs["graph_runtime_state"] + events = _run_with_optional_pause(state, pause_on=None) + return _node_successes(events) + + mocker.patch.object(generator, "_generate", side_effect=_fake_generate) + + resumed_nodes = generator.resume( + app_model=_pause_resume_make_app(), + workflow=_pause_resume_make_workflow(), + user=_pause_resume_make_account(), + application_generate_entity=SimpleNamespace( + stream=False, + invoke_from=InvokeFrom.SERVICE_API, + trace_manager=SimpleNamespace(), + ), + graph_runtime_state=resumed_state, + workflow_execution_repository=SimpleNamespace(), + workflow_node_execution_repository=SimpleNamespace(), + ) + + assert paused_nodes + resumed_nodes == baseline_nodes + assert resumed_state.outputs == baseline_outputs + + +def test_advanced_chat_pause_resume_matches_baseline(mocker: MockerFixture, unbound_session: Session): + _patch_tool_node(mocker) + + baseline_state = _build_runtime_state("adv-baseline") + baseline_events = _run_with_optional_pause(baseline_state, pause_on=None) + assert isinstance(baseline_events[-1], GraphRunSucceededEvent) + baseline_nodes = _node_successes(baseline_events) + baseline_outputs = baseline_state.outputs + + paused_state = _build_runtime_state("adv-paused") + paused_events = _run_with_optional_pause(paused_state, pause_on="tool_a") + assert isinstance(paused_events[-1], GraphRunPausedEvent) + paused_nodes = _node_successes(paused_events) + snapshot = paused_state.dumps() + + resumed_state = RuntimeState.from_snapshot(snapshot) + + generator = adv_app_gen_module.AdvancedChatAppGenerator(execution_driver=WorkflowRunAgg.run) + + def _fake_generate(**kwargs): + state: RuntimeState = kwargs["graph_runtime_state"] + events = _run_with_optional_pause(state, pause_on=None) + return _node_successes(events) + + mocker.patch.object(generator, "_generate", side_effect=_fake_generate) + + resumed_nodes = generator.resume( + app_model=_pause_resume_make_app(), + workflow=_pause_resume_make_workflow(), + user=_pause_resume_make_account(), + conversation=_pause_resume_make_conversation(), + message=_pause_resume_make_message(), + session=unbound_session, + application_generate_entity=SimpleNamespace( + stream=False, + invoke_from=InvokeFrom.SERVICE_API, + trace_manager=SimpleNamespace(), + ), + workflow_execution_repository=SimpleNamespace(), + workflow_node_execution_repository=SimpleNamespace(), + graph_runtime_state=resumed_state, + ) + + assert paused_nodes + resumed_nodes == baseline_nodes + assert resumed_state.outputs == baseline_outputs + + +def test_resume_emits_resumption_start_reason(mocker: MockerFixture) -> None: + _patch_tool_node(mocker) + + paused_state = _build_runtime_state("resume-reason") + paused_events = _run_with_optional_pause(paused_state, pause_on="tool_a") + initial_start = next(event for event in paused_events if isinstance(event, GraphRunStartedEvent)) + assert initial_start.reason == WorkflowStartReason.INITIAL + + resumed_state = RuntimeState.from_snapshot(paused_state.dumps()) + resumed_events = _run_with_optional_pause(resumed_state, pause_on=None) + resume_start = next(event for event in resumed_events if isinstance(event, GraphRunStartedEvent)) + assert resume_start.reason == WorkflowStartReason.RESUMPTION diff --git a/api/tests/unit_tests/core/app/apps/advanced_chat/test_app_runner.py b/api/tests/unit_tests/core/app/apps/advanced_chat/test_app_runner.py new file mode 100644 index 00000000000..8e90f19e04a --- /dev/null +++ b/api/tests/unit_tests/core/app/apps/advanced_chat/test_app_runner.py @@ -0,0 +1,435 @@ +import json +from contextlib import contextmanager +from unittest.mock import MagicMock, patch +from uuid import uuid4 + +import pytest +from sqlalchemy import event, select +from sqlalchemy.engine import Engine +from sqlalchemy.orm import Session, sessionmaker + +import core.app.apps.advanced_chat.app_runner as module +from core.app.app_config.entities import WorkflowUIBasedAppConfig +from core.app.apps.advanced_chat import app_runner as app_runner_module +from core.app.apps.advanced_chat.app_runner import AdvancedChatAppRunner +from core.app.entities.app_invoke_entities import AdvancedChatAppGenerateEntity, InvokeFrom +from core.app.entities.queue_entities import QueueAnnotationReplyEvent, QueueStopEvent +from core.moderation.base import ModerationError +from factories import variable_factory +from graphon.variables import SegmentType +from models import App as _app_runner_conversation_variables_App +from models import Conversation as _app_runner_conversation_variables_Conversation +from models import ConversationVariable +from models import Message as _app_runner_conversation_variables_Message +from models.model import App as _app_runner_input_moderation_App +from models.model import AppMode, IconType, MessageAnnotation +from models.model import Conversation as _app_runner_input_moderation_Conversation +from models.model import Message as _app_runner_input_moderation_Message +from models.workflow import Workflow, WorkflowType + +"""SQLite-backed conversation-variable synchronization tests for AdvancedChatAppRunner.""" + + +APP_ID = "11111111-1111-1111-1111-111111111111" +CONVERSATION_ID = "22222222-2222-2222-2222-222222222222" +OTHER_CONVERSATION_ID = "22222222-2222-2222-2222-222222222223" +VAR_1_ID = "33333333-3333-3333-3333-333333333333" +VAR_2_ID = "33333333-3333-3333-3333-333333333334" + + +def _variable(variable_id: str, name: str, value: str): + return variable_factory.build_conversation_variable_from_mapping( + { + "id": variable_id, + "name": name, + "value_type": SegmentType.STRING, + "value": value, + } + ) + + +def _runner(workflow_variables: list[object]) -> AdvancedChatAppRunner: + workflow = Workflow( + id="55555555-5555-5555-5555-555555555555", + tenant_id="66666666-6666-6666-6666-666666666666", + app_id=APP_ID, + type=WorkflowType.CHAT, + version=Workflow.VERSION_DRAFT, + graph="{}", + features="{}", + created_by="44444444-4444-4444-4444-444444444444", + ) + workflow.conversation_variables = workflow_variables + conversation = _app_runner_conversation_variables_Conversation(id=CONVERSATION_ID, app_id=APP_ID) + return AdvancedChatAppRunner( + application_generate_entity=MagicMock(), + queue_manager=MagicMock(), + conversation=conversation, + message=_app_runner_conversation_variables_Message(id="77777777-7777-7777-7777-777777777777"), + dialogue_count=1, + variable_loader=MagicMock(), + workflow=workflow, + system_user_id="44444444-4444-4444-4444-444444444444", + app=_app_runner_conversation_variables_App(id=APP_ID, tenant_id=workflow.tenant_id), + ) + + +def _bind_runner_sessions(monkeypatch: pytest.MonkeyPatch, sqlite_session: Session) -> None: + engine = sqlite_session.get_bind() + monkeypatch.setattr( + app_runner_module, + "create_session", + lambda: Session(engine, expire_on_commit=False), + ) + + +def _persist_variable(session: Session, *, variable: object, conversation_id: str = CONVERSATION_ID) -> None: + session.add( + ConversationVariable.from_variable( + app_id=APP_ID, + conversation_id=conversation_id, + variable=variable, + ) + ) + session.commit() + + +class TestAppRunnerConversationVariablesCases: + @staticmethod + @pytest.mark.parametrize("sqlite_session", [(ConversationVariable,)], indirect=True) + def test_missing_conversation_variables_are_added(monkeypatch: pytest.MonkeyPatch, sqlite_session: Session) -> None: + existing_variable = _variable(VAR_1_ID, "existing_var", "default1") + new_variable = _variable(VAR_2_ID, "new_var", "default2") + _persist_variable(sqlite_session, variable=existing_variable) + _persist_variable(sqlite_session, variable=new_variable, conversation_id=OTHER_CONVERSATION_ID) + _bind_runner_sessions(monkeypatch, sqlite_session) + + variables = _runner([existing_variable, new_variable])._initialize_conversation_variables() + + assert [variable.id for variable in variables] == [VAR_1_ID, VAR_2_ID] + persisted = sqlite_session.scalars( + select(ConversationVariable) + .where(ConversationVariable.conversation_id == CONVERSATION_ID) + .order_by(ConversationVariable.id) + ).all() + assert [variable.id for variable in persisted] == [VAR_1_ID, VAR_2_ID] + + @staticmethod + @pytest.mark.parametrize("sqlite_session", [(ConversationVariable,)], indirect=True) + def test_no_variables_creates_all(monkeypatch: pytest.MonkeyPatch, sqlite_session: Session) -> None: + workflow_variables = [ + _variable(VAR_1_ID, "var1", "default1"), + _variable(VAR_2_ID, "var2", "default2"), + ] + _bind_runner_sessions(monkeypatch, sqlite_session) + + variables = _runner(workflow_variables)._initialize_conversation_variables() + + assert [variable.id for variable in variables] == [VAR_1_ID, VAR_2_ID] + persisted = sqlite_session.scalars(select(ConversationVariable).order_by(ConversationVariable.id)).all() + assert [variable.id for variable in persisted] == [VAR_1_ID, VAR_2_ID] + + @staticmethod + @pytest.mark.parametrize("sqlite_session", [(ConversationVariable,)], indirect=True) + def test_all_variables_exist_no_changes(monkeypatch: pytest.MonkeyPatch, sqlite_session: Session) -> None: + workflow_variables = [ + _variable(VAR_1_ID, "var1", "default1"), + _variable(VAR_2_ID, "var2", "default2"), + ] + for variable in workflow_variables: + _persist_variable(sqlite_session, variable=variable) + _bind_runner_sessions(monkeypatch, sqlite_session) + + variables = _runner(workflow_variables)._initialize_conversation_variables() + + assert [variable.id for variable in variables] == [VAR_1_ID, VAR_2_ID] + persisted = sqlite_session.scalars(select(ConversationVariable)).all() + assert len(persisted) == 2 + + @staticmethod + @pytest.mark.parametrize("sqlite_session", [(ConversationVariable,)], indirect=True) + def test_non_uuid_conversation_variable_is_created_once( + monkeypatch: pytest.MonkeyPatch, sqlite_session: Session + ) -> None: + author_id = "opt-comp-prompt-var" + variable = _variable(author_id, "optimization_comparison_prompt", "-") + _bind_runner_sessions(monkeypatch, sqlite_session) + runner = _runner([variable]) + + first = runner._initialize_conversation_variables() + second = runner._initialize_conversation_variables() + + expected_row_id = ConversationVariable.storage_id(variable) + assert variable.id == author_id + assert [item.id for item in first] == [author_id] + assert [item.id for item in second] == [author_id] + persisted = sqlite_session.scalars(select(ConversationVariable)).all() + assert [row.id for row in persisted] == [expected_row_id] + assert persisted[0].to_variable().id == author_id + + +MINIMAL_GRAPH = { + "nodes": [ + { + "id": "start", + "data": { + "type": "start", + "title": "Start", + }, + } + ], + "edges": [], +} + + +def _patch_common_run_deps(runner: AdvancedChatAppRunner): + """Context manager that patches common heavy deps used by run().""" + return patch.multiple( + "core.app.apps.advanced_chat.app_runner", + RedisChannel=MagicMock(), + redis_client=MagicMock(), + WorkflowEntry=MagicMock(**{"return_value.run.return_value": iter([])}), + RuntimeState=MagicMock(), + ) + + +class TestAppRunnerInputModerationCases: + @staticmethod + @pytest.fixture + def build_runner(sqlite_session: Session): + """Construct a minimal AdvancedChatAppRunner with heavy dependencies mocked.""" + app_id = str(uuid4()) + workflow_id = str(uuid4()) + tenant_id = str(uuid4()) + + # Mocks for constructor args + mock_queue_manager = MagicMock() + + conversation = _app_runner_input_moderation_Conversation(id=str(uuid4()), app_id=app_id) + message = _app_runner_input_moderation_Message(id=str(uuid4()), app_id=app_id, conversation_id=conversation.id) + workflow = Workflow( + id=workflow_id, + tenant_id=tenant_id, + app_id=app_id, + type=WorkflowType.CHAT, + version=Workflow.VERSION_DRAFT, + graph=json.dumps(MINIMAL_GRAPH), + features="{}", + created_by=str(uuid4()), + ) + + app_config = WorkflowUIBasedAppConfig( + app_id=app_id, workflow_id=workflow_id, tenant_id=tenant_id, app_mode=AppMode.ADVANCED_CHAT + ) + + app = _app_runner_input_moderation_App( + id=app_id, + tenant_id=tenant_id, + name="Advanced chat app", + mode=AppMode.ADVANCED_CHAT, + icon_type=IconType.EMOJI, + icon="chat", + icon_background="#ffffff", + enable_site=False, + enable_api=False, + ) + sqlite_session.add(app) + sqlite_session.commit() + + gen = AdvancedChatAppGenerateEntity( + app_config=app_config, + inputs={"q": "raw"}, + query="raw-query", + files=[], + user_id=str(uuid4()), + invoke_from=InvokeFrom.SERVICE_API, + workflow_run_id=str(uuid4()), + task_id=str(uuid4()), + stream=True, + ) + + runner = AdvancedChatAppRunner( + application_generate_entity=gen, + queue_manager=mock_queue_manager, + conversation=conversation, + message=message, + dialogue_count=1, + variable_loader=MagicMock(), + workflow=workflow, + system_user_id=str(uuid4()), + app=app, + ) + + return runner + + @staticmethod + def test_handle_input_moderation_stops_on_moderation_error(build_runner): + runner = build_runner + + # moderation_for_inputs raises ModerationError -> should stop and emit stop event + with ( + patch.object(runner, "moderation_for_inputs", side_effect=ModerationError("blocked")), + patch.object(runner, "_complete_with_stream_output") as mock_complete, + ): + stop, new_inputs, new_query = runner.handle_input_moderation( + app_record=runner._app, + app_generate_entity=runner.application_generate_entity, + inputs={"k": "v"}, + query="hello", + message_id="mid", + ) + + assert stop is True + # inputs/query should be unchanged on error path + assert new_inputs == {"k": "v"} + assert new_query == "hello" + # ensure stopped_by reason is INPUT_MODERATION + assert mock_complete.called + args, kwargs = mock_complete.call_args + assert kwargs.get("stopped_by") == QueueStopEvent.StopBy.INPUT_MODERATION + + @staticmethod + def test_run_applies_overridden_inputs_and_query_from_moderation(build_runner): + runner = build_runner + runner.application_generate_entity.call_depth = 5 + + overridden_inputs = {"q": "sanitized"} + overridden_query = "sanitized-query" + + with ( + _patch_common_run_deps(runner), + patch.object( + runner, + "moderation_for_inputs", + return_value=(True, overridden_inputs, overridden_query), + ) as mock_moderate, + patch.object(runner, "handle_annotation_reply", return_value=False) as mock_anno, + patch.object(runner, "_init_graph", return_value=MagicMock()) as mock_init_graph, + ): + runner.prepare() + + # moderation called with original values + mock_moderate.assert_called_once() + + # application_generate_entity should be updated to overridden values + assert runner.application_generate_entity.inputs == overridden_inputs + assert runner.application_generate_entity.query == overridden_query + + # annotation reply should use the new query + mock_anno.assert_called() + assert mock_anno.call_args.kwargs.get("query") == overridden_query + + # since not stopped, graph initialization should proceed + assert mock_init_graph.called + assert mock_init_graph.call_args.kwargs["call_depth"] == 5 + + @staticmethod + def test_run_returns_early_when_direct_output_via_handle_input_moderation(build_runner): + runner = build_runner + + with ( + _patch_common_run_deps(runner), + # Simulate handle_input_moderation signalling to stop + patch.object( + runner, + "handle_input_moderation", + return_value=( + True, + runner.application_generate_entity.inputs, + runner.application_generate_entity.query, + ), + ) as mock_handle, + patch.object(runner, "_init_graph") as mock_init_graph, + patch.object(runner, "handle_annotation_reply") as mock_anno, + ): + runner.prepare() + + mock_handle.assert_called_once() + # Ensure no further steps executed + mock_anno.assert_not_called() + mock_init_graph.assert_not_called() + + @staticmethod + def test_run_publishes_annotation_after_commit(build_runner, sqlite_engine: Engine): + runner = build_runner + events: list[str] = [] + + def record_commit(session: Session) -> None: + if session.get_bind() is sqlite_engine: + events.append("commit") + + event.listen(Session, "after_commit", record_commit) + annotation_reply = MessageAnnotation( + app_id=runner._app.id, + question="question", + content="annotated answer", + account_id=str(uuid4()), + ) + + def publish(event): + if isinstance(event, QueueAnnotationReplyEvent): + events.append("publish") + + with ( + _patch_common_run_deps(runner), + patch.object( + runner, + "handle_input_moderation", + return_value=( + False, + runner.application_generate_entity.inputs, + runner.application_generate_entity.query, + ), + ), + patch.object(runner, "handle_annotation_reply", return_value=annotation_reply), + patch.object(runner, "_publish_event", side_effect=publish), + patch.object(runner, "_complete_with_stream_output"), + ): + runner.prepare() + event.remove(Session, "after_commit", record_commit) + + assert events == ["commit", "publish"] + + @staticmethod + def test_run_closes_scoped_session_before_workflow_run(build_runner, sqlite_session_factory: sessionmaker[Session]): + runner = build_runner + events = [] + + @contextmanager + def observed_session(): + with sqlite_session_factory() as session: + yield session + events.append("close") + + workflow_entry = MagicMock() + + def run_workflow(**kwargs): + events.append("run") + return iter([]) + + workflow_entry.run.side_effect = run_workflow + + with ( + patch.object(module, "create_session", observed_session), + patch.object(module, "RedisChannel"), + patch.object(module, "redis_client"), + patch.object(module, "WorkflowEntry", return_value=workflow_entry), + patch.object( + runner, + "handle_input_moderation", + return_value=( + False, + runner.application_generate_entity.inputs, + runner.application_generate_entity.query, + ), + ), + patch.object(runner, "handle_annotation_reply", return_value=False), + patch.object(runner, "_initialize_conversation_variables", return_value=[]), + patch.object(runner, "_init_graph", return_value=MagicMock()), + ): + prepared = runner.prepare() + assert prepared is not None + assert events[-1] == "close" + list(prepared.entry.run()) + + assert events[-2:] == ["close", "run"] diff --git a/api/tests/unit_tests/core/app/apps/advanced_chat/test_app_runner_conversation_variables.py b/api/tests/unit_tests/core/app/apps/advanced_chat/test_app_runner_conversation_variables.py deleted file mode 100644 index 4fc147d7355..00000000000 --- a/api/tests/unit_tests/core/app/apps/advanced_chat/test_app_runner_conversation_variables.py +++ /dev/null @@ -1,149 +0,0 @@ -"""SQLite-backed conversation-variable synchronization tests for AdvancedChatAppRunner.""" - -from unittest.mock import MagicMock - -import pytest -from sqlalchemy import select -from sqlalchemy.orm import Session - -from core.app.apps.advanced_chat import app_runner as app_runner_module -from core.app.apps.advanced_chat.app_runner import AdvancedChatAppRunner -from factories import variable_factory -from graphon.variables import SegmentType -from models import App, Conversation, ConversationVariable, Message -from models.workflow import Workflow, WorkflowType - -APP_ID = "11111111-1111-1111-1111-111111111111" -CONVERSATION_ID = "22222222-2222-2222-2222-222222222222" -OTHER_CONVERSATION_ID = "22222222-2222-2222-2222-222222222223" -VAR_1_ID = "33333333-3333-3333-3333-333333333333" -VAR_2_ID = "33333333-3333-3333-3333-333333333334" - - -def _variable(variable_id: str, name: str, value: str): - return variable_factory.build_conversation_variable_from_mapping( - { - "id": variable_id, - "name": name, - "value_type": SegmentType.STRING, - "value": value, - } - ) - - -def _runner(workflow_variables: list[object]) -> AdvancedChatAppRunner: - workflow = Workflow( - id="55555555-5555-5555-5555-555555555555", - tenant_id="66666666-6666-6666-6666-666666666666", - app_id=APP_ID, - type=WorkflowType.CHAT, - version=Workflow.VERSION_DRAFT, - graph="{}", - features="{}", - created_by="44444444-4444-4444-4444-444444444444", - ) - workflow.conversation_variables = workflow_variables - conversation = Conversation(id=CONVERSATION_ID, app_id=APP_ID) - return AdvancedChatAppRunner( - application_generate_entity=MagicMock(), - queue_manager=MagicMock(), - conversation=conversation, - message=Message(id="77777777-7777-7777-7777-777777777777"), - dialogue_count=1, - variable_loader=MagicMock(), - workflow=workflow, - system_user_id="44444444-4444-4444-4444-444444444444", - app=App(id=APP_ID, tenant_id=workflow.tenant_id), - ) - - -def _bind_runner_sessions(monkeypatch: pytest.MonkeyPatch, sqlite_session: Session) -> None: - engine = sqlite_session.get_bind() - monkeypatch.setattr( - app_runner_module, - "create_session", - lambda: Session(engine, expire_on_commit=False), - ) - - -def _persist_variable(session: Session, *, variable: object, conversation_id: str = CONVERSATION_ID) -> None: - session.add( - ConversationVariable.from_variable( - app_id=APP_ID, - conversation_id=conversation_id, - variable=variable, - ) - ) - session.commit() - - -@pytest.mark.parametrize("sqlite_session", [(ConversationVariable,)], indirect=True) -def test_missing_conversation_variables_are_added(monkeypatch: pytest.MonkeyPatch, sqlite_session: Session) -> None: - existing_variable = _variable(VAR_1_ID, "existing_var", "default1") - new_variable = _variable(VAR_2_ID, "new_var", "default2") - _persist_variable(sqlite_session, variable=existing_variable) - _persist_variable(sqlite_session, variable=new_variable, conversation_id=OTHER_CONVERSATION_ID) - _bind_runner_sessions(monkeypatch, sqlite_session) - - variables = _runner([existing_variable, new_variable])._initialize_conversation_variables() - - assert [variable.id for variable in variables] == [VAR_1_ID, VAR_2_ID] - persisted = sqlite_session.scalars( - select(ConversationVariable) - .where(ConversationVariable.conversation_id == CONVERSATION_ID) - .order_by(ConversationVariable.id) - ).all() - assert [variable.id for variable in persisted] == [VAR_1_ID, VAR_2_ID] - - -@pytest.mark.parametrize("sqlite_session", [(ConversationVariable,)], indirect=True) -def test_no_variables_creates_all(monkeypatch: pytest.MonkeyPatch, sqlite_session: Session) -> None: - workflow_variables = [ - _variable(VAR_1_ID, "var1", "default1"), - _variable(VAR_2_ID, "var2", "default2"), - ] - _bind_runner_sessions(monkeypatch, sqlite_session) - - variables = _runner(workflow_variables)._initialize_conversation_variables() - - assert [variable.id for variable in variables] == [VAR_1_ID, VAR_2_ID] - persisted = sqlite_session.scalars(select(ConversationVariable).order_by(ConversationVariable.id)).all() - assert [variable.id for variable in persisted] == [VAR_1_ID, VAR_2_ID] - - -@pytest.mark.parametrize("sqlite_session", [(ConversationVariable,)], indirect=True) -def test_all_variables_exist_no_changes(monkeypatch: pytest.MonkeyPatch, sqlite_session: Session) -> None: - workflow_variables = [ - _variable(VAR_1_ID, "var1", "default1"), - _variable(VAR_2_ID, "var2", "default2"), - ] - for variable in workflow_variables: - _persist_variable(sqlite_session, variable=variable) - _bind_runner_sessions(monkeypatch, sqlite_session) - - variables = _runner(workflow_variables)._initialize_conversation_variables() - - assert [variable.id for variable in variables] == [VAR_1_ID, VAR_2_ID] - persisted = sqlite_session.scalars(select(ConversationVariable)).all() - assert len(persisted) == 2 - - -@pytest.mark.parametrize("sqlite_session", [(ConversationVariable,)], indirect=True) -def test_non_uuid_conversation_variable_is_created_once( - monkeypatch: pytest.MonkeyPatch, sqlite_session: Session -) -> None: - author_id = "opt-comp-prompt-var" - variable = _variable(author_id, "optimization_comparison_prompt", "-") - _bind_runner_sessions(monkeypatch, sqlite_session) - runner = _runner([variable]) - - first = runner._initialize_conversation_variables() - second = runner._initialize_conversation_variables() - - expected_row_id = ConversationVariable.storage_id(variable) - assert variable.id == author_id - assert [item.id for item in first] == [author_id] - assert [item.id for item in second] == [author_id] - persisted = sqlite_session.scalars(select(ConversationVariable)).all() - assert [row.id for row in persisted] == [expected_row_id] - assert persisted[0].to_variable().id == author_id diff --git a/api/tests/unit_tests/core/app/apps/advanced_chat/test_app_runner_input_moderation.py b/api/tests/unit_tests/core/app/apps/advanced_chat/test_app_runner_input_moderation.py deleted file mode 100644 index 76309013bf9..00000000000 --- a/api/tests/unit_tests/core/app/apps/advanced_chat/test_app_runner_input_moderation.py +++ /dev/null @@ -1,270 +0,0 @@ -import json -from contextlib import contextmanager -from unittest.mock import MagicMock, patch -from uuid import uuid4 - -import pytest -from sqlalchemy import event -from sqlalchemy.engine import Engine -from sqlalchemy.orm import Session, sessionmaker - -import core.app.apps.advanced_chat.app_runner as module -from core.app.app_config.entities import WorkflowUIBasedAppConfig -from core.app.apps.advanced_chat.app_runner import AdvancedChatAppRunner -from core.app.entities.app_invoke_entities import AdvancedChatAppGenerateEntity, InvokeFrom -from core.app.entities.queue_entities import QueueAnnotationReplyEvent, QueueStopEvent -from core.moderation.base import ModerationError -from models.model import App, AppMode, Conversation, IconType, Message, MessageAnnotation -from models.workflow import Workflow, WorkflowType - -MINIMAL_GRAPH = { - "nodes": [ - { - "id": "start", - "data": { - "type": "start", - "title": "Start", - }, - } - ], - "edges": [], -} - - -@pytest.fixture -def build_runner(sqlite_session: Session): - """Construct a minimal AdvancedChatAppRunner with heavy dependencies mocked.""" - app_id = str(uuid4()) - workflow_id = str(uuid4()) - tenant_id = str(uuid4()) - - # Mocks for constructor args - mock_queue_manager = MagicMock() - - conversation = Conversation(id=str(uuid4()), app_id=app_id) - message = Message(id=str(uuid4()), app_id=app_id, conversation_id=conversation.id) - workflow = Workflow( - id=workflow_id, - tenant_id=tenant_id, - app_id=app_id, - type=WorkflowType.CHAT, - version=Workflow.VERSION_DRAFT, - graph=json.dumps(MINIMAL_GRAPH), - features="{}", - created_by=str(uuid4()), - ) - - app_config = WorkflowUIBasedAppConfig( - app_id=app_id, workflow_id=workflow_id, tenant_id=tenant_id, app_mode=AppMode.ADVANCED_CHAT - ) - - app = App( - id=app_id, - tenant_id=tenant_id, - name="Advanced chat app", - mode=AppMode.ADVANCED_CHAT, - icon_type=IconType.EMOJI, - icon="chat", - icon_background="#ffffff", - enable_site=False, - enable_api=False, - ) - sqlite_session.add(app) - sqlite_session.commit() - - gen = AdvancedChatAppGenerateEntity( - app_config=app_config, - inputs={"q": "raw"}, - query="raw-query", - files=[], - user_id=str(uuid4()), - invoke_from=InvokeFrom.SERVICE_API, - workflow_run_id=str(uuid4()), - task_id=str(uuid4()), - stream=True, - ) - - runner = AdvancedChatAppRunner( - application_generate_entity=gen, - queue_manager=mock_queue_manager, - conversation=conversation, - message=message, - dialogue_count=1, - variable_loader=MagicMock(), - workflow=workflow, - system_user_id=str(uuid4()), - app=app, - ) - - return runner - - -def _patch_common_run_deps(runner: AdvancedChatAppRunner): - """Context manager that patches common heavy deps used by run().""" - return patch.multiple( - "core.app.apps.advanced_chat.app_runner", - RedisChannel=MagicMock(), - redis_client=MagicMock(), - WorkflowEntry=MagicMock(**{"return_value.run.return_value": iter([])}), - RuntimeState=MagicMock(), - ) - - -def test_handle_input_moderation_stops_on_moderation_error(build_runner): - runner = build_runner - - # moderation_for_inputs raises ModerationError -> should stop and emit stop event - with ( - patch.object(runner, "moderation_for_inputs", side_effect=ModerationError("blocked")), - patch.object(runner, "_complete_with_stream_output") as mock_complete, - ): - stop, new_inputs, new_query = runner.handle_input_moderation( - app_record=runner._app, - app_generate_entity=runner.application_generate_entity, - inputs={"k": "v"}, - query="hello", - message_id="mid", - ) - - assert stop is True - # inputs/query should be unchanged on error path - assert new_inputs == {"k": "v"} - assert new_query == "hello" - # ensure stopped_by reason is INPUT_MODERATION - assert mock_complete.called - args, kwargs = mock_complete.call_args - assert kwargs.get("stopped_by") == QueueStopEvent.StopBy.INPUT_MODERATION - - -def test_run_applies_overridden_inputs_and_query_from_moderation(build_runner): - runner = build_runner - runner.application_generate_entity.call_depth = 5 - - overridden_inputs = {"q": "sanitized"} - overridden_query = "sanitized-query" - - with ( - _patch_common_run_deps(runner), - patch.object( - runner, - "moderation_for_inputs", - return_value=(True, overridden_inputs, overridden_query), - ) as mock_moderate, - patch.object(runner, "handle_annotation_reply", return_value=False) as mock_anno, - patch.object(runner, "_init_graph", return_value=MagicMock()) as mock_init_graph, - ): - runner.prepare() - - # moderation called with original values - mock_moderate.assert_called_once() - - # application_generate_entity should be updated to overridden values - assert runner.application_generate_entity.inputs == overridden_inputs - assert runner.application_generate_entity.query == overridden_query - - # annotation reply should use the new query - mock_anno.assert_called() - assert mock_anno.call_args.kwargs.get("query") == overridden_query - - # since not stopped, graph initialization should proceed - assert mock_init_graph.called - assert mock_init_graph.call_args.kwargs["call_depth"] == 5 - - -def test_run_returns_early_when_direct_output_via_handle_input_moderation(build_runner): - runner = build_runner - - with ( - _patch_common_run_deps(runner), - # Simulate handle_input_moderation signalling to stop - patch.object( - runner, - "handle_input_moderation", - return_value=(True, runner.application_generate_entity.inputs, runner.application_generate_entity.query), - ) as mock_handle, - patch.object(runner, "_init_graph") as mock_init_graph, - patch.object(runner, "handle_annotation_reply") as mock_anno, - ): - runner.prepare() - - mock_handle.assert_called_once() - # Ensure no further steps executed - mock_anno.assert_not_called() - mock_init_graph.assert_not_called() - - -def test_run_publishes_annotation_after_commit(build_runner, sqlite_engine: Engine): - runner = build_runner - events: list[str] = [] - - def record_commit(session: Session) -> None: - if session.get_bind() is sqlite_engine: - events.append("commit") - - event.listen(Session, "after_commit", record_commit) - annotation_reply = MessageAnnotation( - app_id=runner._app.id, - question="question", - content="annotated answer", - account_id=str(uuid4()), - ) - - def publish(event): - if isinstance(event, QueueAnnotationReplyEvent): - events.append("publish") - - with ( - _patch_common_run_deps(runner), - patch.object( - runner, - "handle_input_moderation", - return_value=(False, runner.application_generate_entity.inputs, runner.application_generate_entity.query), - ), - patch.object(runner, "handle_annotation_reply", return_value=annotation_reply), - patch.object(runner, "_publish_event", side_effect=publish), - patch.object(runner, "_complete_with_stream_output"), - ): - runner.prepare() - event.remove(Session, "after_commit", record_commit) - - assert events == ["commit", "publish"] - - -def test_run_closes_scoped_session_before_workflow_run(build_runner, sqlite_session_factory: sessionmaker[Session]): - runner = build_runner - events = [] - - @contextmanager - def observed_session(): - with sqlite_session_factory() as session: - yield session - events.append("close") - - workflow_entry = MagicMock() - - def run_workflow(**kwargs): - events.append("run") - return iter([]) - - workflow_entry.run.side_effect = run_workflow - - with ( - patch.object(module, "create_session", observed_session), - patch.object(module, "RedisChannel"), - patch.object(module, "redis_client"), - patch.object(module, "WorkflowEntry", return_value=workflow_entry), - patch.object( - runner, - "handle_input_moderation", - return_value=(False, runner.application_generate_entity.inputs, runner.application_generate_entity.query), - ), - patch.object(runner, "handle_annotation_reply", return_value=False), - patch.object(runner, "_initialize_conversation_variables", return_value=[]), - patch.object(runner, "_init_graph", return_value=MagicMock()), - ): - prepared = runner.prepare() - assert prepared is not None - assert events[-1] == "close" - list(prepared.entry.run()) - - assert events[-2:] == ["close", "run"] diff --git a/api/tests/unit_tests/core/app/apps/advanced_chat/test_generate_task_pipeline.py b/api/tests/unit_tests/core/app/apps/advanced_chat/test_generate_task_pipeline.py index 4314aa8b55a..b7f94f6acf9 100644 --- a/api/tests/unit_tests/core/app/apps/advanced_chat/test_generate_task_pipeline.py +++ b/api/tests/unit_tests/core/app/apps/advanced_chat/test_generate_task_pipeline.py @@ -1,5 +1,6 @@ from __future__ import annotations +from contextlib import nullcontext from datetime import datetime from types import SimpleNamespace from unittest import mock @@ -9,24 +10,67 @@ from pytest_mock import MockerFixture from sqlalchemy import select from sqlalchemy.orm import Session, sessionmaker -from core.app.app_config.entities import WorkflowUIBasedAppConfig +from core.app.app_config.entities import AppAdditionalFeatures, WorkflowUIBasedAppConfig from core.app.apps.advanced_chat import generate_task_pipeline as pipeline_module +from core.app.apps.advanced_chat.generate_task_pipeline import ( + AdvancedChatAppGenerateTaskPipeline, + ConversationSnapshot, + MessageSnapshot, + WorkflowSnapshot, +) from core.app.entities.app_invoke_entities import AdvancedChatAppGenerateEntity, InvokeFrom from core.app.entities.queue_entities import ( + QueueAdvancedChatMessageEndEvent, + QueueAnnotationReplyEvent, + QueueErrorEvent, + QueueHumanInputFormFilledEvent, + QueueHumanInputFormTimeoutEvent, + QueueIterationCompletedEvent, + QueueIterationNextEvent, + QueueIterationStartEvent, + QueueLoopCompletedEvent, + QueueLoopNextEvent, + QueueLoopStartEvent, + QueueMessageReplaceEvent, + QueueNodeExceptionEvent, + QueueNodeFailedEvent, QueuePingEvent, + QueueReasoningChunkEvent, + QueueRetrieverResourcesEvent, + QueueStopEvent, QueueTextChunkEvent, + QueueWorkflowFailedEvent, QueueWorkflowPartialSuccessEvent, QueueWorkflowPausedEvent, + QueueWorkflowStartedEvent, QueueWorkflowSucceededEvent, ) -from core.app.entities.task_entities import StreamEvent -from core.workflow.nodes.human_input.pause_reason import HumanInputRequired -from graphon.enums import WorkflowExecutionStatus, WorkflowType +from core.app.entities.task_entities import ( + AdvancedChatPausedBlockingResponse, + AnnotationReply, + AnnotationReplyAccount, + HumanInputRequiredResponse, + MessageAudioStreamResponse, + MessageEndStreamResponse, + PingStreamResponse, + ReasoningChunkStreamResponse, + StreamEvent, +) +from core.base.tts.app_generator_tts_publisher import AudioTrunk +from core.workflow.nodes.human_input.entities import UserActionConfig +from core.workflow.nodes.human_input.pause_reason import DifyHITLEventType, HumanInputRequired +from core.workflow.system_variables import build_system_variables +from graphon.enums import BuiltinNodeTypes, WorkflowExecutionStatus, WorkflowType +from graphon.file import FileTransferMethod, FileType +from graphon.model_runtime.entities.llm_entities import LLMUsage +from graphon.runtime import RuntimeState, VariablePool +from libs.datetime_utils import naive_utc_now from models.enums import CreatorUserRole, MessageStatus, WorkflowRunTriggeredFrom from models.execution_extra_content import HumanInputContent -from models.model import AppMode, EndUser, Message +from models.model import AppMode, EndUser, Message, MessageFile from models.workflow import WorkflowRun from repositories.sqlalchemy_api_workflow_run_repository import DifyAPISQLAlchemyWorkflowRunRepository +from tests.workflow_test_utils import build_test_variable_pool def _build_generate_entity() -> AdvancedChatAppGenerateEntity: @@ -524,3 +568,677 @@ def test_process_stream_response_breaks_after_workflow_partial_success() -> None assert [resp.event for resp in responses] == [StreamEvent.WORKFLOW_FINISHED] pipeline._handle_workflow_partial_success_event.assert_called_once_with(partial_event, trace_manager=None) pipeline._base_task_pipeline.ping_stream_response.assert_not_called() + + +def _make_pipeline(): + app_config = WorkflowUIBasedAppConfig( + tenant_id="tenant", + app_id="app", + app_mode=AppMode.ADVANCED_CHAT, + additional_features=AppAdditionalFeatures(), + variables=[], + workflow_id="workflow-id", + ) + application_generate_entity = AdvancedChatAppGenerateEntity.model_construct( + task_id="task", + app_config=app_config, + inputs={}, + query="hello", + files=[], + user_id="user", + stream=False, + invoke_from=InvokeFrom.WEB_APP, + extras={}, + trace_manager=None, + workflow_run_id="run-id", + ) + + message = MessageSnapshot( + id="message-id", + query="hello", + created_at=naive_utc_now(), + status=MessageStatus.NORMAL, + answer="", + ) + conversation = ConversationSnapshot(id="conv-id", mode=AppMode.ADVANCED_CHAT) + workflow = WorkflowSnapshot(id="workflow-id", tenant_id="tenant", features_dict={}) + user = EndUser(tenant_id="tenant", type="session", name="tester", session_id="session") + + pipeline = AdvancedChatAppGenerateTaskPipeline( + application_generate_entity=application_generate_entity, + workflow=workflow, + queue_manager=SimpleNamespace(invoke_from=InvokeFrom.WEB_APP, graph_runtime_state=None), + conversation=conversation, + message=message, + user=user, + stream=False, + dialogue_count=1, + draft_var_saver_factory=lambda **kwargs: None, + ) + + return pipeline + + +def _persist_message(session: Session, *, message_id: str = "message-id") -> Message: + message = Message( + app_id="app", + model_provider="provider", + model_id="model", + override_model_configs=None, + conversation_id="conv-id", + _inputs={}, + query="hello", + message="", + message_tokens=0, + message_unit_price=0, + message_price_unit=0, + answer="", + answer_tokens=0, + answer_unit_price=0, + answer_price_unit=0, + parent_message_id=None, + provider_response_latency=0, + total_price=0, + currency="USD", + invoke_from=InvokeFrom.WEB_APP, + from_source="api", + from_end_user_id="end-user", + from_account_id=None, + app_mode=AppMode.ADVANCED_CHAT, + status=MessageStatus.PAUSED, + ) + message.id = message_id + session.add(message) + session.commit() + return message + + +class TestAdvancedChatGenerateTaskPipeline: + def test_ensure_workflow_initialized_raises(self): + pipeline = _make_pipeline() + + with pytest.raises(ValueError, match="workflow run not initialized"): + pipeline._ensure_workflow_initialized() + + def test_to_blocking_response_returns_message_end(self): + pipeline = _make_pipeline() + pipeline._task_state.answer = "done" + + def _gen(): + yield MessageEndStreamResponse(task_id="task", id="message-id", metadata={"k": "v"}) + + response = pipeline._to_blocking_response(_gen()) + + assert response.data.answer == "done" + assert response.data.metadata == {"k": "v"} + + def test_to_blocking_response_falls_back_to_human_input_required_when_pause_event_missing(self): + pipeline = _make_pipeline() + pipeline._task_state.answer = "partial answer" + pipeline._workflow_run_id = "run-id" + pipeline._graph_runtime_state = RuntimeState( + workflow_id="test-workflow", + variable_pool=build_test_variable_pool( + variables=build_system_variables(workflow_execution_id="run-id"), + ), + start_at=0.0, + llm_usage=LLMUsage.empty_usage().model_copy(update={"total_tokens": 7}), + node_run_steps=3, + ) + + def _gen(): + yield HumanInputRequiredResponse( + task_id="task", + workflow_run_id="run-id", + data=HumanInputRequiredResponse.Data( + form_id="form-1", + node_id="node-1", + node_title="Approval", + form_content="Need approval", + inputs=[], + actions=[UserActionConfig(id="approve", title="Approve")], + display_in_ui=True, + form_token="token-1", + resolved_default_values={}, + expiration_time=123, + ), + ) + + response = pipeline._to_blocking_response(_gen()) + + assert isinstance(response, AdvancedChatPausedBlockingResponse) + assert response.data.workflow_run_id == "run-id" + assert response.data.status == "paused" + assert response.data.paused_nodes == ["node-1"] + assert response.data.reasons == [ + { + "TYPE": DifyHITLEventType.HUMAN_INPUT_REQUIRED, + "form_id": "form-1", + "node_id": "node-1", + "node_title": "Approval", + "form_content": "Need approval", + "inputs": [], + "actions": [{"id": "approve", "title": "Approve", "button_style": "default"}], + "display_in_ui": True, + "form_token": "token-1", + "approval_channels": [], + "resolved_default_values": {}, + "expiration_time": 123, + } + ] + + def test_handle_text_chunk_event_updates_state(self): + pipeline = _make_pipeline() + pipeline._message_cycle_manager = SimpleNamespace( + message_to_stream_response=lambda **kwargs: MessageEndStreamResponse( + task_id="task", id="message-id", metadata={} + ) + ) + + event = SimpleNamespace(text="hi", from_variable_selector=None) + + responses = list(pipeline._handle_text_chunk_event(event)) + + assert pipeline._task_state.answer == "hi" + assert responses + + def test_handle_reasoning_chunk_event_emits_on_nonempty(self): + pipeline = _make_pipeline() + event = QueueReasoningChunkEvent(reasoning="pondering", from_node_id="llm-1", is_final=False) + + responses = list(pipeline._handle_reasoning_chunk_event(event)) + + assert len(responses) == 1 + response = responses[0] + assert isinstance(response, ReasoningChunkStreamResponse) + assert response.data.message_id == pipeline._message_id + assert response.data.reasoning == "pondering" + assert response.data.node_id == "llm-1" + assert response.data.is_final is False + # reasoning never touches the answer stream + assert pipeline._task_state.answer == "" + + def test_handle_reasoning_chunk_event_drops_empty_nonfinal(self): + pipeline = _make_pipeline() + event = QueueReasoningChunkEvent(reasoning="", from_node_id="llm-1", is_final=False) + + responses = list(pipeline._handle_reasoning_chunk_event(event)) + + assert responses == [] + + def test_handle_reasoning_chunk_event_emits_empty_final_marker(self): + pipeline = _make_pipeline() + event = QueueReasoningChunkEvent(reasoning="", from_node_id="llm-1", is_final=True) + + responses = list(pipeline._handle_reasoning_chunk_event(event)) + + assert len(responses) == 1 + response = responses[0] + assert isinstance(response, ReasoningChunkStreamResponse) + assert response.data.reasoning == "" + assert response.data.is_final is True + + def test_listen_audio_msg_returns_audio_stream(self): + pipeline = _make_pipeline() + publisher = SimpleNamespace(check_and_get_audio=lambda: AudioTrunk(status="responding", audio="data")) + + response = pipeline._listen_audio_msg(publisher=publisher, task_id="task") + + assert isinstance(response, MessageAudioStreamResponse) + + def test_handle_ping_event(self): + pipeline = _make_pipeline() + pipeline._base_task_pipeline.ping_stream_response = lambda: PingStreamResponse(task_id="task") + + responses = list(pipeline._handle_ping_event(QueuePingEvent())) + + assert isinstance(responses[0], PingStreamResponse) + + def test_handle_error_event(self): + pipeline = _make_pipeline() + pipeline._base_task_pipeline.handle_error = lambda **kwargs: ValueError("boom") + pipeline._base_task_pipeline.error_to_stream_response = lambda err: err + + responses = list(pipeline._handle_error_event(QueueErrorEvent(error=ValueError("boom")))) + + assert isinstance(responses[0], ValueError) + + def test_handle_workflow_started_event_sets_run_id(self, sqlite_session: Session): + pipeline = _make_pipeline() + message = _persist_message(sqlite_session) + other_message = _persist_message(sqlite_session, message_id="other-message-id") + pipeline._graph_runtime_state = RuntimeState( + workflow_id="test-workflow", + variable_pool=build_test_variable_pool(variables=build_system_variables(workflow_execution_id="run-id")), + start_at=0.0, + ) + pipeline._workflow_response_converter.workflow_start_to_stream_response = lambda **kwargs: "started" + + responses = list(pipeline._handle_workflow_started_event(QueueWorkflowStartedEvent())) + + assert pipeline._workflow_run_id == "run-id" + assert responses == ["started"] + + sqlite_session.refresh(message) + sqlite_session.refresh(other_message) + assert message.workflow_run_id == "run-id" + assert other_message.workflow_run_id is None + + def test_message_end_to_stream_response_strips_annotation_reply(self): + pipeline = _make_pipeline() + pipeline._task_state.metadata.annotation_reply = AnnotationReply( + id="ann", + account=AnnotationReplyAccount(id="acc", name="acc"), + ) + + response = pipeline._message_end_to_stream_response() + + assert "annotation_reply" not in response.metadata + + def test_handle_output_moderation_chunk_publishes_stop(self): + pipeline = _make_pipeline() + events: list[object] = [] + + class _Moderation: + def should_direct_output(self): + return True + + def get_final_output(self): + return "final" + + pipeline._base_task_pipeline.output_moderation_handler = _Moderation() + pipeline._base_task_pipeline.queue_manager = SimpleNamespace( + publish=lambda event, pub_from: events.append(event) + ) + + result = pipeline._handle_output_moderation_chunk("ignored") + + assert result is True + assert pipeline._task_state.answer == "final" + assert any(isinstance(event, QueueTextChunkEvent) for event in events) + assert any(isinstance(event, QueueStopEvent) for event in events) + + def test_handle_node_succeeded_event_records_files(self): + pipeline = _make_pipeline() + pipeline._workflow_response_converter.fetch_files_from_node_outputs = lambda outputs: [ + {"type": "file", "transfer_method": "local"} + ] + pipeline._workflow_response_converter.workflow_node_finish_to_stream_response = lambda **kwargs: "done" + pipeline._save_output_for_event = lambda event, node_execution_id: None + + event = SimpleNamespace( + node_type=BuiltinNodeTypes.ANSWER, + outputs={"k": "v"}, + node_execution_id="exec", + node_id="node", + ) + + responses = list(pipeline._handle_node_succeeded_event(event)) + + assert responses == ["done"] + assert pipeline._recorded_files + + def test_handle_node_succeeded_event_records_llm_reasoning(self): + pipeline = _make_pipeline() + pipeline._workflow_response_converter.fetch_files_from_node_outputs = lambda outputs: [] + pipeline._workflow_response_converter.workflow_node_finish_to_stream_response = lambda **kwargs: "done" + pipeline._save_output_for_event = lambda event, node_execution_id: None + + event = SimpleNamespace( + node_type=BuiltinNodeTypes.LLM, + outputs={"reasoning_content": "first pass "}, + node_execution_id="exec", + node_id="llm-1", + ) + + list(pipeline._handle_node_succeeded_event(event)) + + assert pipeline._task_state.metadata.reasoning == {"llm-1": "first pass "} + + def test_handle_node_succeeded_event_accumulates_reasoning_across_passes(self): + pipeline = _make_pipeline() + pipeline._workflow_response_converter.fetch_files_from_node_outputs = lambda outputs: [] + pipeline._workflow_response_converter.workflow_node_finish_to_stream_response = lambda **kwargs: "done" + pipeline._save_output_for_event = lambda event, node_execution_id: None + + def _llm_event(reasoning: str): + return SimpleNamespace( + node_type=BuiltinNodeTypes.LLM, + outputs={"reasoning_content": reasoning}, + node_execution_id="exec", + node_id="llm-1", + ) + + # Same node id across iteration/loop passes must accumulate, not overwrite. + list(pipeline._handle_node_succeeded_event(_llm_event("pass one "))) + list(pipeline._handle_node_succeeded_event(_llm_event("pass two"))) + + assert pipeline._task_state.metadata.reasoning == {"llm-1": "pass one pass two"} + + def test_iteration_and_loop_handlers(self): + pipeline = _make_pipeline() + pipeline._workflow_run_id = "run-id" + pipeline._workflow_response_converter.workflow_iteration_start_to_stream_response = lambda **kwargs: ( + "iter_start" + ) + pipeline._workflow_response_converter.workflow_iteration_next_to_stream_response = lambda **kwargs: "iter_next" + pipeline._workflow_response_converter.workflow_iteration_completed_to_stream_response = lambda **kwargs: ( + "iter_done" + ) + pipeline._workflow_response_converter.workflow_loop_start_to_stream_response = lambda **kwargs: "loop_start" + pipeline._workflow_response_converter.workflow_loop_next_to_stream_response = lambda **kwargs: "loop_next" + pipeline._workflow_response_converter.workflow_loop_completed_to_stream_response = lambda **kwargs: "loop_done" + + iter_start = QueueIterationStartEvent( + node_execution_id="exec", + node_id="node", + node_type=BuiltinNodeTypes.LLM, + node_title="LLM", + start_at=naive_utc_now(), + ) + iter_next = QueueIterationNextEvent( + index=1, + node_execution_id="exec", + node_id="node", + node_type=BuiltinNodeTypes.LLM, + node_title="LLM", + ) + iter_done = QueueIterationCompletedEvent( + node_execution_id="exec", + node_id="node", + node_type=BuiltinNodeTypes.LLM, + node_title="LLM", + start_at=naive_utc_now(), + ) + loop_start = QueueLoopStartEvent( + node_execution_id="exec", + node_id="node", + node_type=BuiltinNodeTypes.LLM, + node_title="LLM", + start_at=naive_utc_now(), + ) + loop_next = QueueLoopNextEvent( + index=1, + node_execution_id="exec", + node_id="node", + node_type=BuiltinNodeTypes.LLM, + node_title="LLM", + ) + loop_done = QueueLoopCompletedEvent( + node_execution_id="exec", + node_id="node", + node_type=BuiltinNodeTypes.LLM, + node_title="LLM", + start_at=naive_utc_now(), + ) + + assert list(pipeline._handle_iteration_start_event(iter_start)) == ["iter_start"] + assert list(pipeline._handle_iteration_next_event(iter_next)) == ["iter_next"] + assert list(pipeline._handle_iteration_completed_event(iter_done)) == ["iter_done"] + assert list(pipeline._handle_loop_start_event(loop_start)) == ["loop_start"] + assert list(pipeline._handle_loop_next_event(loop_next)) == ["loop_next"] + assert list(pipeline._handle_loop_completed_event(loop_done)) == ["loop_done"] + + def test_workflow_finish_handlers(self): + pipeline = _make_pipeline() + pipeline._workflow_run_id = "run-id" + pipeline._graph_runtime_state = RuntimeState( + workflow_id="test-workflow", + variable_pool=VariablePool.from_bootstrap( + system_variables=build_system_variables(workflow_execution_id="run-id") + ), + start_at=0.0, + ) + pipeline._workflow_response_converter.workflow_finish_to_stream_response = lambda **kwargs: "finish" + pipeline._workflow_response_converter.workflow_pause_to_stream_response = lambda **kwargs: ["pause"] + pipeline._persist_human_input_extra_content = lambda **kwargs: None + pipeline._save_message = lambda **kwargs: None + pipeline._base_task_pipeline.queue_manager.publish = lambda *args, **kwargs: None + pipeline._base_task_pipeline.handle_error = lambda **kwargs: ValueError("boom") + pipeline._base_task_pipeline.error_to_stream_response = lambda err: err + pipeline._get_message = lambda **kwargs: Message(id="message-id") + + succeeded_responses = list(pipeline._handle_workflow_succeeded_event(QueueWorkflowSucceededEvent(outputs={}))) + assert len(succeeded_responses) == 2 + assert isinstance(succeeded_responses[0], MessageEndStreamResponse) + assert succeeded_responses[1] == "finish" + + partial_success_responses = list( + pipeline._handle_workflow_partial_success_event( + QueueWorkflowPartialSuccessEvent(exceptions_count=1, outputs={}) + ) + ) + assert len(partial_success_responses) == 2 + assert isinstance(partial_success_responses[0], MessageEndStreamResponse) + assert partial_success_responses[1] == "finish" + assert ( + list(pipeline._handle_workflow_failed_event(QueueWorkflowFailedEvent(error="err", exceptions_count=1)))[0] + == "finish" + ) + assert list(pipeline._handle_workflow_paused_event(QueueWorkflowPausedEvent(reasons=[], outputs={}))) == [ + "pause" + ] + + def test_node_failure_handlers(self): + pipeline = _make_pipeline() + pipeline._workflow_response_converter.workflow_node_finish_to_stream_response = lambda **kwargs: "node_finish" + pipeline._save_output_for_event = lambda event, node_execution_id: None + + failed_event = QueueNodeFailedEvent( + node_execution_id="exec", + node_id="node", + node_type=BuiltinNodeTypes.LLM, + start_at=naive_utc_now(), + inputs={}, + outputs={}, + process_data={}, + error="err", + ) + exc_event = QueueNodeExceptionEvent( + node_execution_id="exec", + node_id="node", + node_type=BuiltinNodeTypes.LLM, + start_at=naive_utc_now(), + inputs={}, + outputs={}, + process_data={}, + error="err", + ) + + assert list(pipeline._handle_node_failed_events(failed_event)) == ["node_finish"] + assert list(pipeline._handle_node_failed_events(exc_event)) == ["node_finish"] + + def test_handle_text_chunk_event_tracks_streaming_metrics(self): + pipeline = _make_pipeline() + published: list[object] = [] + + class _Publisher: + def publish(self, message): + published.append(message) + + pipeline._message_cycle_manager = SimpleNamespace(message_to_stream_response=lambda **kwargs: "chunk") + + event = SimpleNamespace(text="hi", from_variable_selector=["a"]) + queue_message = SimpleNamespace(event=event) + + responses = list( + pipeline._handle_text_chunk_event(event, tts_publisher=_Publisher(), queue_message=queue_message) + ) + + assert responses == ["chunk"] + assert pipeline._task_state.is_streaming_response is True + assert pipeline._task_state.first_token_time is not None + assert pipeline._task_state.last_token_time is not None + assert pipeline._task_state.answer == "hi" + assert published == [queue_message] + + def test_handle_output_moderation_chunk_appends_token(self): + pipeline = _make_pipeline() + seen: list[str] = [] + + class _Moderation: + def should_direct_output(self): + return False + + def append_new_token(self, text): + seen.append(text) + + pipeline._base_task_pipeline.output_moderation_handler = _Moderation() + + result = pipeline._handle_output_moderation_chunk("token") + + assert result is False + assert seen == ["token"] + + def test_handle_retriever_and_annotation_events(self, monkeypatch: pytest.MonkeyPatch, unbound_session: Session): + pipeline = _make_pipeline() + calls = {"retriever": 0, "annotation": 0} + + def _hit_retriever(event): + calls["retriever"] += 1 + + def _hit_annotation(_manager, event, session): + calls["annotation"] += 1 + + pipeline._message_cycle_manager.handle_retriever_resources = _hit_retriever + monkeypatch.setattr(type(pipeline._message_cycle_manager), "handle_annotation_reply", _hit_annotation) + + retriever_event = QueueRetrieverResourcesEvent(retriever_resources=[]) + annotation_event = QueueAnnotationReplyEvent(message_annotation_id="ann") + + monkeypatch.setattr(pipeline, "_database_session", lambda: nullcontext(unbound_session)) + + assert list(pipeline._handle_retriever_resources_event(retriever_event)) == [] + assert list(pipeline._handle_annotation_reply_event(annotation_event)) == [] + assert calls == {"retriever": 1, "annotation": 1} + + def test_handle_message_replace_event(self): + pipeline = _make_pipeline() + pipeline._message_cycle_manager.message_replace_to_stream_response = lambda **kwargs: "replace" + + event = QueueMessageReplaceEvent( + text="new", + reason=QueueMessageReplaceEvent.MessageReplaceReason.OUTPUT_MODERATION, + ) + + assert list(pipeline._handle_message_replace_event(event)) == ["replace"] + + def test_handle_human_input_events(self): + pipeline = _make_pipeline() + persisted: list[str] = [] + pipeline._persist_human_input_extra_content = lambda **kwargs: persisted.append("saved") + pipeline._workflow_response_converter.human_input_form_filled_to_stream_response = lambda **kwargs: "filled" + pipeline._workflow_response_converter.human_input_form_timeout_to_stream_response = lambda **kwargs: "timeout" + + filled_event = QueueHumanInputFormFilledEvent( + form_id="form", + node_id="node", + node_type=BuiltinNodeTypes.LLM, + node_title="title", + rendered_content="content", + action_id="action", + action_text="action", + ) + timeout_event = QueueHumanInputFormTimeoutEvent( + form_id="form", + node_id="node", + node_type=BuiltinNodeTypes.LLM, + node_title="title", + expiration_time=naive_utc_now(), + ) + + assert list(pipeline._handle_human_input_form_filled_event(filled_event)) == ["filled"] + assert list(pipeline._handle_human_input_form_timeout_event(timeout_event)) == ["timeout"] + assert persisted == ["saved"] + + def test_save_message_preserves_full_answer_and_sets_usage(self, sqlite_session: Session): + pipeline = _make_pipeline() + pipeline._recorded_files = [ + { + "type": FileType.IMAGE, + "transfer_method": FileTransferMethod.REMOTE_URL, + "remote_url": "http://example.com/file.png", + "related_id": "file-id", + } + ] + # The answer is stored verbatim; markdown image links are never stripped. + pipeline._task_state.answer = "![img](http://example.com/file.png) hello ![inline](http://llm.com/img.jpg)" + pipeline._task_state.is_streaming_response = True + pipeline._task_state.first_token_time = pipeline._base_task_pipeline.start_at + 0.1 + pipeline._task_state.last_token_time = pipeline._base_task_pipeline.start_at + 0.2 + + message = _persist_message(sqlite_session) + + graph_runtime_state = RuntimeState( + workflow_id="test-workflow", + variable_pool=VariablePool.from_bootstrap( + system_variables=build_system_variables(workflow_execution_id="run-id") + ), + start_at=0.0, + ) + + pipeline._save_message(session=sqlite_session, graph_runtime_state=graph_runtime_state) + sqlite_session.commit() + + assert message.status == MessageStatus.NORMAL + assert message.answer == "![img](http://example.com/file.png) hello ![inline](http://llm.com/img.jpg)" + assert message.message_metadata + assert sqlite_session.query(MessageFile).filter_by(message_id=message.id).count() == 1 + + def test_handle_stop_event_saves_message_for_moderation(self): + pipeline = _make_pipeline() + pipeline._message_end_to_stream_response = lambda: "end" + saved: list[str] = [] + + def _save_message(**kwargs): + saved.append("saved") + + pipeline._save_message = _save_message + + responses = list(pipeline._handle_stop_event(QueueStopEvent(stopped_by=QueueStopEvent.StopBy.INPUT_MODERATION))) + + assert responses == ["end"] + assert saved == ["saved"] + + def test_handle_message_end_event_applies_output_moderation(self): + pipeline = _make_pipeline() + pipeline._graph_runtime_state = RuntimeState( + workflow_id="test-workflow", + variable_pool=VariablePool.from_bootstrap( + system_variables=build_system_variables(workflow_execution_id="run-id") + ), + start_at=0.0, + ) + pipeline._base_task_pipeline.handle_output_moderation_when_task_finished = lambda answer: "safe" + pipeline._message_cycle_manager.message_replace_to_stream_response = lambda **kwargs: "replace" + pipeline._message_end_to_stream_response = lambda: "end" + + saved: list[str] = [] + + def _save_message(**kwargs): + saved.append("saved") + + pipeline._save_message = _save_message + + responses = list(pipeline._handle_advanced_chat_message_end_event(QueueAdvancedChatMessageEndEvent())) + + assert responses == ["replace", "end"] + assert saved == ["saved"] + + def test_dispatch_event_handles_node_exception(self): + pipeline = _make_pipeline() + pipeline._workflow_response_converter.workflow_node_finish_to_stream_response = lambda **kwargs: "failed" + pipeline._save_output_for_event = lambda *args, **kwargs: None + + event = QueueNodeExceptionEvent( + node_execution_id="exec", + node_id="node", + node_type=BuiltinNodeTypes.LLM, + start_at=naive_utc_now(), + inputs={}, + outputs={}, + process_data={}, + error="err", + ) + + assert list(pipeline._dispatch_event(event)) == ["failed"] diff --git a/api/tests/unit_tests/core/app/apps/advanced_chat/test_generate_task_pipeline_core.py b/api/tests/unit_tests/core/app/apps/advanced_chat/test_generate_task_pipeline_core.py deleted file mode 100644 index 5be3635fba4..00000000000 --- a/api/tests/unit_tests/core/app/apps/advanced_chat/test_generate_task_pipeline_core.py +++ /dev/null @@ -1,738 +0,0 @@ -from __future__ import annotations - -from contextlib import nullcontext -from types import SimpleNamespace - -import pytest -from sqlalchemy.orm import Session - -from core.app.app_config.entities import AppAdditionalFeatures, WorkflowUIBasedAppConfig -from core.app.apps.advanced_chat.generate_task_pipeline import ( - AdvancedChatAppGenerateTaskPipeline, - ConversationSnapshot, - MessageSnapshot, - WorkflowSnapshot, -) -from core.app.entities.app_invoke_entities import AdvancedChatAppGenerateEntity, InvokeFrom -from core.app.entities.queue_entities import ( - QueueAdvancedChatMessageEndEvent, - QueueAnnotationReplyEvent, - QueueErrorEvent, - QueueHumanInputFormFilledEvent, - QueueHumanInputFormTimeoutEvent, - QueueIterationCompletedEvent, - QueueIterationNextEvent, - QueueIterationStartEvent, - QueueLoopCompletedEvent, - QueueLoopNextEvent, - QueueLoopStartEvent, - QueueMessageReplaceEvent, - QueueNodeExceptionEvent, - QueueNodeFailedEvent, - QueuePingEvent, - QueueReasoningChunkEvent, - QueueRetrieverResourcesEvent, - QueueStopEvent, - QueueTextChunkEvent, - QueueWorkflowFailedEvent, - QueueWorkflowPartialSuccessEvent, - QueueWorkflowPausedEvent, - QueueWorkflowStartedEvent, - QueueWorkflowSucceededEvent, -) -from core.app.entities.task_entities import ( - AdvancedChatPausedBlockingResponse, - AnnotationReply, - AnnotationReplyAccount, - HumanInputRequiredResponse, - MessageAudioStreamResponse, - MessageEndStreamResponse, - PingStreamResponse, - ReasoningChunkStreamResponse, -) -from core.base.tts.app_generator_tts_publisher import AudioTrunk -from core.workflow.nodes.human_input.entities import UserActionConfig -from core.workflow.nodes.human_input.pause_reason import DifyHITLEventType -from core.workflow.system_variables import build_system_variables -from graphon.enums import BuiltinNodeTypes -from graphon.file import FileTransferMethod, FileType -from graphon.model_runtime.entities.llm_entities import LLMUsage -from graphon.runtime import RuntimeState, VariablePool -from libs.datetime_utils import naive_utc_now -from models.enums import MessageStatus -from models.model import AppMode, EndUser, Message, MessageFile -from tests.workflow_test_utils import build_test_variable_pool - - -def _make_pipeline(): - app_config = WorkflowUIBasedAppConfig( - tenant_id="tenant", - app_id="app", - app_mode=AppMode.ADVANCED_CHAT, - additional_features=AppAdditionalFeatures(), - variables=[], - workflow_id="workflow-id", - ) - application_generate_entity = AdvancedChatAppGenerateEntity.model_construct( - task_id="task", - app_config=app_config, - inputs={}, - query="hello", - files=[], - user_id="user", - stream=False, - invoke_from=InvokeFrom.WEB_APP, - extras={}, - trace_manager=None, - workflow_run_id="run-id", - ) - - message = MessageSnapshot( - id="message-id", - query="hello", - created_at=naive_utc_now(), - status=MessageStatus.NORMAL, - answer="", - ) - conversation = ConversationSnapshot(id="conv-id", mode=AppMode.ADVANCED_CHAT) - workflow = WorkflowSnapshot(id="workflow-id", tenant_id="tenant", features_dict={}) - user = EndUser(tenant_id="tenant", type="session", name="tester", session_id="session") - - pipeline = AdvancedChatAppGenerateTaskPipeline( - application_generate_entity=application_generate_entity, - workflow=workflow, - queue_manager=SimpleNamespace(invoke_from=InvokeFrom.WEB_APP, graph_runtime_state=None), - conversation=conversation, - message=message, - user=user, - stream=False, - dialogue_count=1, - draft_var_saver_factory=lambda **kwargs: None, - ) - - return pipeline - - -def _persist_message(session: Session, *, message_id: str = "message-id") -> Message: - message = Message( - app_id="app", - model_provider="provider", - model_id="model", - override_model_configs=None, - conversation_id="conv-id", - _inputs={}, - query="hello", - message="", - message_tokens=0, - message_unit_price=0, - message_price_unit=0, - answer="", - answer_tokens=0, - answer_unit_price=0, - answer_price_unit=0, - parent_message_id=None, - provider_response_latency=0, - total_price=0, - currency="USD", - invoke_from=InvokeFrom.WEB_APP, - from_source="api", - from_end_user_id="end-user", - from_account_id=None, - app_mode=AppMode.ADVANCED_CHAT, - status=MessageStatus.PAUSED, - ) - message.id = message_id - session.add(message) - session.commit() - return message - - -class TestAdvancedChatGenerateTaskPipeline: - def test_ensure_workflow_initialized_raises(self): - pipeline = _make_pipeline() - - with pytest.raises(ValueError, match="workflow run not initialized"): - pipeline._ensure_workflow_initialized() - - def test_to_blocking_response_returns_message_end(self): - pipeline = _make_pipeline() - pipeline._task_state.answer = "done" - - def _gen(): - yield MessageEndStreamResponse(task_id="task", id="message-id", metadata={"k": "v"}) - - response = pipeline._to_blocking_response(_gen()) - - assert response.data.answer == "done" - assert response.data.metadata == {"k": "v"} - - def test_to_blocking_response_falls_back_to_human_input_required_when_pause_event_missing(self): - pipeline = _make_pipeline() - pipeline._task_state.answer = "partial answer" - pipeline._workflow_run_id = "run-id" - pipeline._graph_runtime_state = RuntimeState( - workflow_id="test-workflow", - variable_pool=build_test_variable_pool( - variables=build_system_variables(workflow_execution_id="run-id"), - ), - start_at=0.0, - llm_usage=LLMUsage.empty_usage().model_copy(update={"total_tokens": 7}), - node_run_steps=3, - ) - - def _gen(): - yield HumanInputRequiredResponse( - task_id="task", - workflow_run_id="run-id", - data=HumanInputRequiredResponse.Data( - form_id="form-1", - node_id="node-1", - node_title="Approval", - form_content="Need approval", - inputs=[], - actions=[UserActionConfig(id="approve", title="Approve")], - display_in_ui=True, - form_token="token-1", - resolved_default_values={}, - expiration_time=123, - ), - ) - - response = pipeline._to_blocking_response(_gen()) - - assert isinstance(response, AdvancedChatPausedBlockingResponse) - assert response.data.workflow_run_id == "run-id" - assert response.data.status == "paused" - assert response.data.paused_nodes == ["node-1"] - assert response.data.reasons == [ - { - "TYPE": DifyHITLEventType.HUMAN_INPUT_REQUIRED, - "form_id": "form-1", - "node_id": "node-1", - "node_title": "Approval", - "form_content": "Need approval", - "inputs": [], - "actions": [{"id": "approve", "title": "Approve", "button_style": "default"}], - "display_in_ui": True, - "form_token": "token-1", - "approval_channels": [], - "resolved_default_values": {}, - "expiration_time": 123, - } - ] - - def test_handle_text_chunk_event_updates_state(self): - pipeline = _make_pipeline() - pipeline._message_cycle_manager = SimpleNamespace( - message_to_stream_response=lambda **kwargs: MessageEndStreamResponse( - task_id="task", id="message-id", metadata={} - ) - ) - - event = SimpleNamespace(text="hi", from_variable_selector=None) - - responses = list(pipeline._handle_text_chunk_event(event)) - - assert pipeline._task_state.answer == "hi" - assert responses - - def test_handle_reasoning_chunk_event_emits_on_nonempty(self): - pipeline = _make_pipeline() - event = QueueReasoningChunkEvent(reasoning="pondering", from_node_id="llm-1", is_final=False) - - responses = list(pipeline._handle_reasoning_chunk_event(event)) - - assert len(responses) == 1 - response = responses[0] - assert isinstance(response, ReasoningChunkStreamResponse) - assert response.data.message_id == pipeline._message_id - assert response.data.reasoning == "pondering" - assert response.data.node_id == "llm-1" - assert response.data.is_final is False - # reasoning never touches the answer stream - assert pipeline._task_state.answer == "" - - def test_handle_reasoning_chunk_event_drops_empty_nonfinal(self): - pipeline = _make_pipeline() - event = QueueReasoningChunkEvent(reasoning="", from_node_id="llm-1", is_final=False) - - responses = list(pipeline._handle_reasoning_chunk_event(event)) - - assert responses == [] - - def test_handle_reasoning_chunk_event_emits_empty_final_marker(self): - pipeline = _make_pipeline() - event = QueueReasoningChunkEvent(reasoning="", from_node_id="llm-1", is_final=True) - - responses = list(pipeline._handle_reasoning_chunk_event(event)) - - assert len(responses) == 1 - response = responses[0] - assert isinstance(response, ReasoningChunkStreamResponse) - assert response.data.reasoning == "" - assert response.data.is_final is True - - def test_listen_audio_msg_returns_audio_stream(self): - pipeline = _make_pipeline() - publisher = SimpleNamespace(check_and_get_audio=lambda: AudioTrunk(status="responding", audio="data")) - - response = pipeline._listen_audio_msg(publisher=publisher, task_id="task") - - assert isinstance(response, MessageAudioStreamResponse) - - def test_handle_ping_event(self): - pipeline = _make_pipeline() - pipeline._base_task_pipeline.ping_stream_response = lambda: PingStreamResponse(task_id="task") - - responses = list(pipeline._handle_ping_event(QueuePingEvent())) - - assert isinstance(responses[0], PingStreamResponse) - - def test_handle_error_event(self): - pipeline = _make_pipeline() - pipeline._base_task_pipeline.handle_error = lambda **kwargs: ValueError("boom") - pipeline._base_task_pipeline.error_to_stream_response = lambda err: err - - responses = list(pipeline._handle_error_event(QueueErrorEvent(error=ValueError("boom")))) - - assert isinstance(responses[0], ValueError) - - def test_handle_workflow_started_event_sets_run_id(self, sqlite_session: Session): - pipeline = _make_pipeline() - message = _persist_message(sqlite_session) - other_message = _persist_message(sqlite_session, message_id="other-message-id") - pipeline._graph_runtime_state = RuntimeState( - workflow_id="test-workflow", - variable_pool=build_test_variable_pool(variables=build_system_variables(workflow_execution_id="run-id")), - start_at=0.0, - ) - pipeline._workflow_response_converter.workflow_start_to_stream_response = lambda **kwargs: "started" - - responses = list(pipeline._handle_workflow_started_event(QueueWorkflowStartedEvent())) - - assert pipeline._workflow_run_id == "run-id" - assert responses == ["started"] - - sqlite_session.refresh(message) - sqlite_session.refresh(other_message) - assert message.workflow_run_id == "run-id" - assert other_message.workflow_run_id is None - - def test_message_end_to_stream_response_strips_annotation_reply(self): - pipeline = _make_pipeline() - pipeline._task_state.metadata.annotation_reply = AnnotationReply( - id="ann", - account=AnnotationReplyAccount(id="acc", name="acc"), - ) - - response = pipeline._message_end_to_stream_response() - - assert "annotation_reply" not in response.metadata - - def test_handle_output_moderation_chunk_publishes_stop(self): - pipeline = _make_pipeline() - events: list[object] = [] - - class _Moderation: - def should_direct_output(self): - return True - - def get_final_output(self): - return "final" - - pipeline._base_task_pipeline.output_moderation_handler = _Moderation() - pipeline._base_task_pipeline.queue_manager = SimpleNamespace( - publish=lambda event, pub_from: events.append(event) - ) - - result = pipeline._handle_output_moderation_chunk("ignored") - - assert result is True - assert pipeline._task_state.answer == "final" - assert any(isinstance(event, QueueTextChunkEvent) for event in events) - assert any(isinstance(event, QueueStopEvent) for event in events) - - def test_handle_node_succeeded_event_records_files(self): - pipeline = _make_pipeline() - pipeline._workflow_response_converter.fetch_files_from_node_outputs = lambda outputs: [ - {"type": "file", "transfer_method": "local"} - ] - pipeline._workflow_response_converter.workflow_node_finish_to_stream_response = lambda **kwargs: "done" - pipeline._save_output_for_event = lambda event, node_execution_id: None - - event = SimpleNamespace( - node_type=BuiltinNodeTypes.ANSWER, - outputs={"k": "v"}, - node_execution_id="exec", - node_id="node", - ) - - responses = list(pipeline._handle_node_succeeded_event(event)) - - assert responses == ["done"] - assert pipeline._recorded_files - - def test_handle_node_succeeded_event_records_llm_reasoning(self): - pipeline = _make_pipeline() - pipeline._workflow_response_converter.fetch_files_from_node_outputs = lambda outputs: [] - pipeline._workflow_response_converter.workflow_node_finish_to_stream_response = lambda **kwargs: "done" - pipeline._save_output_for_event = lambda event, node_execution_id: None - - event = SimpleNamespace( - node_type=BuiltinNodeTypes.LLM, - outputs={"reasoning_content": "first pass "}, - node_execution_id="exec", - node_id="llm-1", - ) - - list(pipeline._handle_node_succeeded_event(event)) - - assert pipeline._task_state.metadata.reasoning == {"llm-1": "first pass "} - - def test_handle_node_succeeded_event_accumulates_reasoning_across_passes(self): - pipeline = _make_pipeline() - pipeline._workflow_response_converter.fetch_files_from_node_outputs = lambda outputs: [] - pipeline._workflow_response_converter.workflow_node_finish_to_stream_response = lambda **kwargs: "done" - pipeline._save_output_for_event = lambda event, node_execution_id: None - - def _llm_event(reasoning: str): - return SimpleNamespace( - node_type=BuiltinNodeTypes.LLM, - outputs={"reasoning_content": reasoning}, - node_execution_id="exec", - node_id="llm-1", - ) - - # Same node id across iteration/loop passes must accumulate, not overwrite. - list(pipeline._handle_node_succeeded_event(_llm_event("pass one "))) - list(pipeline._handle_node_succeeded_event(_llm_event("pass two"))) - - assert pipeline._task_state.metadata.reasoning == {"llm-1": "pass one pass two"} - - def test_iteration_and_loop_handlers(self): - pipeline = _make_pipeline() - pipeline._workflow_run_id = "run-id" - pipeline._workflow_response_converter.workflow_iteration_start_to_stream_response = lambda **kwargs: ( - "iter_start" - ) - pipeline._workflow_response_converter.workflow_iteration_next_to_stream_response = lambda **kwargs: "iter_next" - pipeline._workflow_response_converter.workflow_iteration_completed_to_stream_response = lambda **kwargs: ( - "iter_done" - ) - pipeline._workflow_response_converter.workflow_loop_start_to_stream_response = lambda **kwargs: "loop_start" - pipeline._workflow_response_converter.workflow_loop_next_to_stream_response = lambda **kwargs: "loop_next" - pipeline._workflow_response_converter.workflow_loop_completed_to_stream_response = lambda **kwargs: "loop_done" - - iter_start = QueueIterationStartEvent( - node_execution_id="exec", - node_id="node", - node_type=BuiltinNodeTypes.LLM, - node_title="LLM", - start_at=naive_utc_now(), - ) - iter_next = QueueIterationNextEvent( - index=1, - node_execution_id="exec", - node_id="node", - node_type=BuiltinNodeTypes.LLM, - node_title="LLM", - ) - iter_done = QueueIterationCompletedEvent( - node_execution_id="exec", - node_id="node", - node_type=BuiltinNodeTypes.LLM, - node_title="LLM", - start_at=naive_utc_now(), - ) - loop_start = QueueLoopStartEvent( - node_execution_id="exec", - node_id="node", - node_type=BuiltinNodeTypes.LLM, - node_title="LLM", - start_at=naive_utc_now(), - ) - loop_next = QueueLoopNextEvent( - index=1, - node_execution_id="exec", - node_id="node", - node_type=BuiltinNodeTypes.LLM, - node_title="LLM", - ) - loop_done = QueueLoopCompletedEvent( - node_execution_id="exec", - node_id="node", - node_type=BuiltinNodeTypes.LLM, - node_title="LLM", - start_at=naive_utc_now(), - ) - - assert list(pipeline._handle_iteration_start_event(iter_start)) == ["iter_start"] - assert list(pipeline._handle_iteration_next_event(iter_next)) == ["iter_next"] - assert list(pipeline._handle_iteration_completed_event(iter_done)) == ["iter_done"] - assert list(pipeline._handle_loop_start_event(loop_start)) == ["loop_start"] - assert list(pipeline._handle_loop_next_event(loop_next)) == ["loop_next"] - assert list(pipeline._handle_loop_completed_event(loop_done)) == ["loop_done"] - - def test_workflow_finish_handlers(self): - pipeline = _make_pipeline() - pipeline._workflow_run_id = "run-id" - pipeline._graph_runtime_state = RuntimeState( - workflow_id="test-workflow", - variable_pool=VariablePool.from_bootstrap( - system_variables=build_system_variables(workflow_execution_id="run-id") - ), - start_at=0.0, - ) - pipeline._workflow_response_converter.workflow_finish_to_stream_response = lambda **kwargs: "finish" - pipeline._workflow_response_converter.workflow_pause_to_stream_response = lambda **kwargs: ["pause"] - pipeline._persist_human_input_extra_content = lambda **kwargs: None - pipeline._save_message = lambda **kwargs: None - pipeline._base_task_pipeline.queue_manager.publish = lambda *args, **kwargs: None - pipeline._base_task_pipeline.handle_error = lambda **kwargs: ValueError("boom") - pipeline._base_task_pipeline.error_to_stream_response = lambda err: err - pipeline._get_message = lambda **kwargs: Message(id="message-id") - - succeeded_responses = list(pipeline._handle_workflow_succeeded_event(QueueWorkflowSucceededEvent(outputs={}))) - assert len(succeeded_responses) == 2 - assert isinstance(succeeded_responses[0], MessageEndStreamResponse) - assert succeeded_responses[1] == "finish" - - partial_success_responses = list( - pipeline._handle_workflow_partial_success_event( - QueueWorkflowPartialSuccessEvent(exceptions_count=1, outputs={}) - ) - ) - assert len(partial_success_responses) == 2 - assert isinstance(partial_success_responses[0], MessageEndStreamResponse) - assert partial_success_responses[1] == "finish" - assert ( - list(pipeline._handle_workflow_failed_event(QueueWorkflowFailedEvent(error="err", exceptions_count=1)))[0] - == "finish" - ) - assert list(pipeline._handle_workflow_paused_event(QueueWorkflowPausedEvent(reasons=[], outputs={}))) == [ - "pause" - ] - - def test_node_failure_handlers(self): - pipeline = _make_pipeline() - pipeline._workflow_response_converter.workflow_node_finish_to_stream_response = lambda **kwargs: "node_finish" - pipeline._save_output_for_event = lambda event, node_execution_id: None - - failed_event = QueueNodeFailedEvent( - node_execution_id="exec", - node_id="node", - node_type=BuiltinNodeTypes.LLM, - start_at=naive_utc_now(), - inputs={}, - outputs={}, - process_data={}, - error="err", - ) - exc_event = QueueNodeExceptionEvent( - node_execution_id="exec", - node_id="node", - node_type=BuiltinNodeTypes.LLM, - start_at=naive_utc_now(), - inputs={}, - outputs={}, - process_data={}, - error="err", - ) - - assert list(pipeline._handle_node_failed_events(failed_event)) == ["node_finish"] - assert list(pipeline._handle_node_failed_events(exc_event)) == ["node_finish"] - - def test_handle_text_chunk_event_tracks_streaming_metrics(self): - pipeline = _make_pipeline() - published: list[object] = [] - - class _Publisher: - def publish(self, message): - published.append(message) - - pipeline._message_cycle_manager = SimpleNamespace(message_to_stream_response=lambda **kwargs: "chunk") - - event = SimpleNamespace(text="hi", from_variable_selector=["a"]) - queue_message = SimpleNamespace(event=event) - - responses = list( - pipeline._handle_text_chunk_event(event, tts_publisher=_Publisher(), queue_message=queue_message) - ) - - assert responses == ["chunk"] - assert pipeline._task_state.is_streaming_response is True - assert pipeline._task_state.first_token_time is not None - assert pipeline._task_state.last_token_time is not None - assert pipeline._task_state.answer == "hi" - assert published == [queue_message] - - def test_handle_output_moderation_chunk_appends_token(self): - pipeline = _make_pipeline() - seen: list[str] = [] - - class _Moderation: - def should_direct_output(self): - return False - - def append_new_token(self, text): - seen.append(text) - - pipeline._base_task_pipeline.output_moderation_handler = _Moderation() - - result = pipeline._handle_output_moderation_chunk("token") - - assert result is False - assert seen == ["token"] - - def test_handle_retriever_and_annotation_events(self, monkeypatch: pytest.MonkeyPatch, unbound_session: Session): - pipeline = _make_pipeline() - calls = {"retriever": 0, "annotation": 0} - - def _hit_retriever(event): - calls["retriever"] += 1 - - def _hit_annotation(_manager, event, session): - calls["annotation"] += 1 - - pipeline._message_cycle_manager.handle_retriever_resources = _hit_retriever - monkeypatch.setattr(type(pipeline._message_cycle_manager), "handle_annotation_reply", _hit_annotation) - - retriever_event = QueueRetrieverResourcesEvent(retriever_resources=[]) - annotation_event = QueueAnnotationReplyEvent(message_annotation_id="ann") - - monkeypatch.setattr(pipeline, "_database_session", lambda: nullcontext(unbound_session)) - - assert list(pipeline._handle_retriever_resources_event(retriever_event)) == [] - assert list(pipeline._handle_annotation_reply_event(annotation_event)) == [] - assert calls == {"retriever": 1, "annotation": 1} - - def test_handle_message_replace_event(self): - pipeline = _make_pipeline() - pipeline._message_cycle_manager.message_replace_to_stream_response = lambda **kwargs: "replace" - - event = QueueMessageReplaceEvent( - text="new", - reason=QueueMessageReplaceEvent.MessageReplaceReason.OUTPUT_MODERATION, - ) - - assert list(pipeline._handle_message_replace_event(event)) == ["replace"] - - def test_handle_human_input_events(self): - pipeline = _make_pipeline() - persisted: list[str] = [] - pipeline._persist_human_input_extra_content = lambda **kwargs: persisted.append("saved") - pipeline._workflow_response_converter.human_input_form_filled_to_stream_response = lambda **kwargs: "filled" - pipeline._workflow_response_converter.human_input_form_timeout_to_stream_response = lambda **kwargs: "timeout" - - filled_event = QueueHumanInputFormFilledEvent( - form_id="form", - node_id="node", - node_type=BuiltinNodeTypes.LLM, - node_title="title", - rendered_content="content", - action_id="action", - action_text="action", - ) - timeout_event = QueueHumanInputFormTimeoutEvent( - form_id="form", - node_id="node", - node_type=BuiltinNodeTypes.LLM, - node_title="title", - expiration_time=naive_utc_now(), - ) - - assert list(pipeline._handle_human_input_form_filled_event(filled_event)) == ["filled"] - assert list(pipeline._handle_human_input_form_timeout_event(timeout_event)) == ["timeout"] - assert persisted == ["saved"] - - def test_save_message_preserves_full_answer_and_sets_usage(self, sqlite_session: Session): - pipeline = _make_pipeline() - pipeline._recorded_files = [ - { - "type": FileType.IMAGE, - "transfer_method": FileTransferMethod.REMOTE_URL, - "remote_url": "http://example.com/file.png", - "related_id": "file-id", - } - ] - # The answer is stored verbatim; markdown image links are never stripped. - pipeline._task_state.answer = "![img](http://example.com/file.png) hello ![inline](http://llm.com/img.jpg)" - pipeline._task_state.is_streaming_response = True - pipeline._task_state.first_token_time = pipeline._base_task_pipeline.start_at + 0.1 - pipeline._task_state.last_token_time = pipeline._base_task_pipeline.start_at + 0.2 - - message = _persist_message(sqlite_session) - - graph_runtime_state = RuntimeState( - workflow_id="test-workflow", - variable_pool=VariablePool.from_bootstrap( - system_variables=build_system_variables(workflow_execution_id="run-id") - ), - start_at=0.0, - ) - - pipeline._save_message(session=sqlite_session, graph_runtime_state=graph_runtime_state) - sqlite_session.commit() - - assert message.status == MessageStatus.NORMAL - assert message.answer == "![img](http://example.com/file.png) hello ![inline](http://llm.com/img.jpg)" - assert message.message_metadata - assert sqlite_session.query(MessageFile).filter_by(message_id=message.id).count() == 1 - - def test_handle_stop_event_saves_message_for_moderation(self): - pipeline = _make_pipeline() - pipeline._message_end_to_stream_response = lambda: "end" - saved: list[str] = [] - - def _save_message(**kwargs): - saved.append("saved") - - pipeline._save_message = _save_message - - responses = list(pipeline._handle_stop_event(QueueStopEvent(stopped_by=QueueStopEvent.StopBy.INPUT_MODERATION))) - - assert responses == ["end"] - assert saved == ["saved"] - - def test_handle_message_end_event_applies_output_moderation(self): - pipeline = _make_pipeline() - pipeline._graph_runtime_state = RuntimeState( - workflow_id="test-workflow", - variable_pool=VariablePool.from_bootstrap( - system_variables=build_system_variables(workflow_execution_id="run-id") - ), - start_at=0.0, - ) - pipeline._base_task_pipeline.handle_output_moderation_when_task_finished = lambda answer: "safe" - pipeline._message_cycle_manager.message_replace_to_stream_response = lambda **kwargs: "replace" - pipeline._message_end_to_stream_response = lambda: "end" - - saved: list[str] = [] - - def _save_message(**kwargs): - saved.append("saved") - - pipeline._save_message = _save_message - - responses = list(pipeline._handle_advanced_chat_message_end_event(QueueAdvancedChatMessageEndEvent())) - - assert responses == ["replace", "end"] - assert saved == ["saved"] - - def test_dispatch_event_handles_node_exception(self): - pipeline = _make_pipeline() - pipeline._workflow_response_converter.workflow_node_finish_to_stream_response = lambda **kwargs: "failed" - pipeline._save_output_for_event = lambda *args, **kwargs: None - - event = QueueNodeExceptionEvent( - node_execution_id="exec", - node_id="node", - node_type=BuiltinNodeTypes.LLM, - start_at=naive_utc_now(), - inputs={}, - outputs={}, - process_data={}, - error="err", - ) - - assert list(pipeline._dispatch_event(event)) == ["failed"] diff --git a/api/tests/unit_tests/core/app/apps/agent_app/test_app_generator.py b/api/tests/unit_tests/core/app/apps/agent_app/test_app_generator.py index 7cec02c4fe3..d7bd1f2d32f 100644 --- a/api/tests/unit_tests/core/app/apps/agent_app/test_app_generator.py +++ b/api/tests/unit_tests/core/app/apps/agent_app/test_app_generator.py @@ -1,3 +1,58 @@ +from __future__ import annotations + +import contextlib +import inspect +from decimal import Decimal +from types import SimpleNamespace +from typing import Any +from unittest.mock import MagicMock + +import pytest +from pytest_mock import MockerFixture +from sqlalchemy import event, select +from sqlalchemy.orm import Session + +import core.app.apps.agent_app.app_generator as module +import core.app.features.annotation_reply.annotation_reply as annotation_mod +import core.moderation.input_moderation as input_moderation_mod +from core.app.apps.agent_app import app_generator +from core.app.apps.agent_app.app_generator import ( + AgentAppGenerator, + AgentAppGeneratorError, + AgentAppNotPublishedError, +) +from core.app.apps.exc import GenerateTaskStoppedError +from core.app.entities.app_invoke_entities import InvokeFrom, UserFrom +from core.app.entities.queue_entities import ( + QueueAnnotationReplyEvent, + QueueLLMChunkEvent, + QueueMessageEndEvent, +) +from core.moderation.base import ModerationError +from models import Account as _app_generator_Account +from models import AppModelConfig +from models.account import Account as _resolve_agent_Account +from models.agent import ( + Agent, + AgentConfigDraft, + AgentConfigDraftType, + AgentConfigSnapshot, + AgentConfigVersionKind, + AgentScope, + AgentSource, + AgentStatus, + AgentWorkingResourceStatus, + AgentWorkspaceBinding, +) +from models.agent_config_entities import AgentSoulConfig +from models.enums import AppStatus, ConversationFromSource +from models.model import App, AppMode, Conversation, Message, MessageAnnotation +from services.agent.workspace_service import ( + AgentWorkspaceBindingGenerationMismatchError, + AgentWorkspaceService, +) +from tests.helpers.unit.model_factories import make_account, make_app, make_conversation, make_message + """Unit tests for AgentAppGenerator.generate() and its worker thread. Mirrors the agent_chat generator tests: every collaborator (config manager, @@ -7,42 +62,17 @@ patched so no real pydantic entity is built, and the worker's flask-context manager is replaced with a no-op so the thread body can run inline. """ -from __future__ import annotations - -import contextlib -import inspect -from decimal import Decimal -from types import SimpleNamespace - -import pytest -from pytest_mock import MockerFixture -from sqlalchemy import event -from sqlalchemy.orm import Session - -import core.app.apps.agent_app.app_generator as module -from core.app.apps.agent_app.app_generator import ( - AgentAppGenerator, - AgentAppGeneratorError, -) -from core.app.apps.exc import GenerateTaskStoppedError -from core.app.entities.app_invoke_entities import InvokeFrom, UserFrom -from core.app.entities.queue_entities import QueueAnnotationReplyEvent -from models import Account, AppModelConfig -from models.agent import Agent, AgentConfigSnapshot, AgentScope, AgentSource, AgentStatus -from models.agent_config_entities import AgentSoulConfig -from models.enums import AppStatus, ConversationFromSource -from models.model import App, AppMode, Conversation, Message, MessageAnnotation MODULE = "core.app.apps.agent_app.app_generator" -def _account(user_id: str = "user") -> Account: - account = Account(name="User", email=f"{user_id}@example.com") +def _app_generator_account(user_id: str = "user") -> _app_generator_Account: + account = _app_generator_Account(name="User", email=f"{user_id}@example.com") account.id = user_id return account -def _app(*, app_model_config_id: str | None = None) -> App: +def _app_generator_app(*, app_model_config_id: str | None = None) -> App: return App( id="app1", tenant_id="tenant", @@ -70,7 +100,7 @@ def _agent(*, agent_id: str = "agent1") -> Agent: ) -def _snapshot(*, snapshot_id: str = "snap1", agent_id: str = "agent1") -> AgentConfigSnapshot: +def _app_generator_snapshot(*, snapshot_id: str = "snap1", agent_id: str = "agent1") -> AgentConfigSnapshot: return AgentConfigSnapshot( id=snapshot_id, tenant_id="tenant", @@ -117,380 +147,1036 @@ def _message(*, query: str = "query") -> Message: _CURRENT_SESSION: Session | None = None -@pytest.fixture(autouse=True) -def _bind_real_session(sqlite_session: Session, monkeypatch: pytest.MonkeyPatch): - global _CURRENT_SESSION - _CURRENT_SESSION = sqlite_session - monkeypatch.setattr(module.db, "session", sqlite_session) - yield - _CURRENT_SESSION = None - - def _session() -> Session: assert _CURRENT_SESSION is not None return _CURRENT_SESSION -@pytest.fixture -def generator(mocker: MockerFixture) -> AgentAppGenerator: - gen = AgentAppGenerator() - mocker.patch(f"{MODULE}.current_app", new=mocker.MagicMock(_get_current_object=mocker.MagicMock())) - mocker.patch(f"{MODULE}.contextvars.copy_context", return_value="ctx") - return gen - - -class TestGenerateGuards: - def test_rejects_blocking_mode(self, generator: AgentAppGenerator): - with pytest.raises(AgentAppGeneratorError, match="only supports streaming"): - generator.generate( - app_model=_app(), - user=_account("u"), - args={}, - invoke_from=InvokeFrom.WEB_APP, - session=_session(), - streaming=False, - ) - - def test_requires_query(self, generator: AgentAppGenerator): - with pytest.raises(AgentAppGeneratorError, match="query is required"): - generator.generate( - app_model=_app(), - user=_account("u"), - args={"inputs": {}}, - invoke_from=InvokeFrom.WEB_APP, - session=_session(), - ) - - def test_rejects_blank_query(self, generator: AgentAppGenerator): - with pytest.raises(AgentAppGeneratorError, match="query is required"): - generator.generate( - app_model=_app(), - user=_account("u"), - args={"query": " ", "inputs": {}}, - invoke_from=InvokeFrom.WEB_APP, - session=_session(), - ) - - -class TestGenerateSuccess: - def test_session_scope_config_version_id_preserves_draft_or_snapshot_id(self): - assert ( - AgentAppGenerator._session_scope_config_version_id( - invoke_from=InvokeFrom.DEBUGGER, config_version_id="draft-1" - ) - == "draft-1" - ) - assert ( - AgentAppGenerator._session_scope_config_version_id( - invoke_from=InvokeFrom.WEB_APP, config_version_id="snapshot-1" - ) - == "snapshot-1" - ) - - def test_generate_orchestrates_and_starts_worker(self, generator, mocker: MockerFixture): - config = AppModelConfig(app_id="app1") - config.id = "config-1" - session = _session() - session.add(config) - session.commit() - app_model = _app(app_model_config_id=config.id) - user = _account() - - generator._resolve_agent = mocker.MagicMock(return_value=(_agent(), "snap1", "snapshot", AgentSoulConfig())) - generator._prepare_user_inputs = mocker.MagicMock(return_value={"x": 1}) - generator._init_generate_records = mocker.MagicMock(return_value=(_conversation(), _message())) - generator._handle_response = mocker.MagicMock(return_value="raw-response") - - mocker.patch( - f"{MODULE}.AgentAppConfigManager.get_app_config", - return_value=mocker.MagicMock(variables=[], tenant_id="tenant", app_id="app1"), - ) - mocker.patch(f"{MODULE}.ModelConfigConverter.convert", return_value=mocker.MagicMock(model="gpt-4o-mini")) - file_upload_config = mocker.MagicMock() - mocker.patch(f"{MODULE}.FileUploadConfigManager.convert", return_value=file_upload_config) - parsed_file = mocker.MagicMock() - build_files = mocker.patch(f"{MODULE}.file_factory.build_from_mappings", return_value=[parsed_file]) - mocker.patch(f"{MODULE}.TraceQueueManager", return_value=mocker.MagicMock()) - generate_entity = mocker.patch( - f"{MODULE}.AgentAppGenerateEntity", return_value=mocker.MagicMock(task_id="t", user_id="user") - ) - mocker.patch(f"{MODULE}.MessageBasedAppQueueManager", return_value=mocker.MagicMock()) - thread_obj = mocker.MagicMock() - thread_constructor = mocker.patch(f"{MODULE}.threading.Thread", return_value=thread_obj) - mocker.patch(f"{MODULE}.AgentAppGenerateResponseConverter.convert", return_value={"result": "ok"}) - file_mappings = [ - { - "type": "image", - "transfer_method": "local_file", - "url": "", - "upload_file_id": "upload-file-1", - } - ] - - result = generator.generate( - app_model=app_model, - user=user, - args={"query": "hello", "inputs": {"name": "world"}, "files": file_mappings}, - invoke_from=InvokeFrom.WEB_APP, - session=session, - streaming=True, - ) - - assert result == {"result": "ok"} - thread_obj.start.assert_called_once() - worker_call = thread_constructor.call_args - inspect.signature(worker_call.kwargs["target"]).bind(**worker_call.kwargs["kwargs"]) - generator._resolve_agent.assert_called_once_with( - app_model, - invoke_from=InvokeFrom.WEB_APP, - draft_type=None, - user=user, - session=session, - conversation=None, - ) - assert session.get(AppModelConfig, "config-1") is config - build_files.assert_called_once() - assert build_files.call_args.kwargs["mappings"] == file_mappings - assert generate_entity.call_args.kwargs["files"] == [parsed_file] - assert generate_entity.call_args.kwargs["file_upload_config"] is file_upload_config - assert "agent_runtime_exit_intent" not in generate_entity.call_args.kwargs - - def test_generate_loads_existing_conversation(self, generator: AgentAppGenerator, mocker: MockerFixture): - app_model = _app() - generator._resolve_agent = mocker.MagicMock( - return_value=(_agent(agent_id="a"), "snap1", "snapshot", AgentSoulConfig()) - ) - generator._prepare_user_inputs = mocker.MagicMock(return_value={}) - generator._init_generate_records = mocker.MagicMock(return_value=(_conversation(), _message())) - generator._handle_response = mocker.MagicMock(return_value="raw") - get_conv = mocker.patch(f"{MODULE}.ConversationService.get_conversation", return_value=_conversation()) - mocker.patch(f"{MODULE}.AgentAppConfigManager.get_app_config", return_value=mocker.MagicMock(variables=[])) - mocker.patch(f"{MODULE}.load_annotation_reply_config", return_value={"enabled": False}) - mocker.patch(f"{MODULE}.ModelConfigConverter.convert", return_value=mocker.MagicMock()) - mocker.patch(f"{MODULE}.TraceQueueManager", return_value=mocker.MagicMock()) - mocker.patch(f"{MODULE}.AgentAppGenerateEntity", return_value=mocker.MagicMock()) - mocker.patch(f"{MODULE}.MessageBasedAppQueueManager", return_value=mocker.MagicMock()) - mocker.patch(f"{MODULE}.threading.Thread", return_value=mocker.MagicMock()) - mocker.patch(f"{MODULE}.AgentAppGenerateResponseConverter.convert", return_value={"result": "ok"}) - session = _session() - user = _account() - - generator.generate( - app_model=app_model, - user=user, - args={"query": "hi", "inputs": {}, "conversation_id": "conv"}, - invoke_from=InvokeFrom.WEB_APP, - session=session, - streaming=True, - ) - - get_conv.assert_called_once_with( - app_model=app_model, - conversation_id="conv", - user=user, - session=session, - ) - assert generator._resolve_agent.call_args.kwargs["conversation"].id == "conv" - assert generator._init_generate_records.call_args.kwargs["session"] is session - - def test_generate_does_not_include_trace_session_id_in_extras( - self, generator: AgentAppGenerator, mocker: MockerFixture - ): - app_model = _app() - user = _account() - - generator._resolve_agent = mocker.MagicMock(return_value=(_agent(), "snap1", "snapshot", AgentSoulConfig())) - generator._prepare_user_inputs = mocker.MagicMock(return_value={}) - generator._init_generate_records = mocker.MagicMock(return_value=(_conversation(), _message())) - generator._handle_response = mocker.MagicMock(return_value="raw-response") - - mocker.patch( - f"{MODULE}.AgentAppConfigManager.get_app_config", - return_value=mocker.MagicMock(variables=[], tenant_id="tenant", app_id="app1"), - ) - mocker.patch(f"{MODULE}.ModelConfigConverter.convert", return_value=mocker.MagicMock(model="gpt-4o-mini")) - mocker.patch(f"{MODULE}.TraceQueueManager", return_value=mocker.MagicMock()) - generate_entity = mocker.patch( - f"{MODULE}.AgentAppGenerateEntity", return_value=mocker.MagicMock(task_id="t", user_id="user") - ) - mocker.patch(f"{MODULE}.MessageBasedAppQueueManager", return_value=mocker.MagicMock()) - mocker.patch(f"{MODULE}.threading.Thread", return_value=mocker.MagicMock()) - mocker.patch(f"{MODULE}.AgentAppGenerateResponseConverter.convert", return_value={"result": "ok"}) - - generator.generate( - app_model=app_model, - user=user, - args={"query": "hello", "inputs": {}, "trace_session_id": "session-1"}, - invoke_from=InvokeFrom.WEB_APP, - session=_session(), - streaming=True, - ) - - assert generate_entity.call_args.kwargs["extras"] == {"auto_generate_conversation_name": True} - - -class TestGenerateWorker: +class TestAppGeneratorCases: + @staticmethod @pytest.fixture(autouse=True) - def patch_context(self, mocker: MockerFixture): - @contextlib.contextmanager - def ctx_manager[**P](*args: P.args, **kwargs: P.kwargs): - yield + def _bind_real_session(sqlite_session: Session, monkeypatch: pytest.MonkeyPatch): + global _CURRENT_SESSION + _CURRENT_SESSION = sqlite_session + monkeypatch.setattr(module.db, "session", sqlite_session) + yield + _CURRENT_SESSION = None - mocker.patch("libs.flask_utils.preserve_flask_contexts", ctx_manager) + @staticmethod + @pytest.fixture + def generator(mocker: MockerFixture) -> AgentAppGenerator: + gen = AgentAppGenerator() + mocker.patch(f"{MODULE}.current_app", new=mocker.MagicMock(_get_current_object=mocker.MagicMock())) + mocker.patch(f"{MODULE}.contextvars.copy_context", return_value="ctx") + return gen - def _wire( - self, - generator: AgentAppGenerator, - mocker: MockerFixture, - *, - run_side_effect=None, - handled=False, - guard_query="query", - ): - generator._get_conversation = mocker.MagicMock(return_value=_conversation()) - generator._get_message = mocker.MagicMock(return_value=_message()) - generator._run_input_guards = mocker.MagicMock(return_value=(handled, guard_query, None)) - resolved_agent = _agent(agent_id="a") - resolved_config = _snapshot(snapshot_id="s", agent_id="a") - resolver_sessions: list[Session] = [] + class TestGenerateGuards: + def test_rejects_blocking_mode(self, generator: AgentAppGenerator): + with pytest.raises(AgentAppGeneratorError, match="only supports streaming"): + generator.generate( + app_model=_app_generator_app(), + user=_app_generator_account("u"), + args={}, + invoke_from=InvokeFrom.WEB_APP, + session=_session(), + streaming=False, + ) - def resolve_agent_by_id(**kwargs): - resolver_sessions.append(kwargs["session"]) - return resolved_agent, resolved_config, AgentSoulConfig() + def test_requires_query(self, generator: AgentAppGenerator): + with pytest.raises(AgentAppGeneratorError, match="query is required"): + generator.generate( + app_model=_app_generator_app(), + user=_app_generator_account("u"), + args={"inputs": {}}, + invoke_from=InvokeFrom.WEB_APP, + session=_session(), + ) - generator._resolve_agent_by_id = mocker.MagicMock(side_effect=resolve_agent_by_id) - session = _session() - if session.get(App, "app1") is None: - session.add(_app()) + def test_rejects_blank_query(self, generator: AgentAppGenerator): + with pytest.raises(AgentAppGeneratorError, match="query is required"): + generator.generate( + app_model=_app_generator_app(), + user=_app_generator_account("u"), + args={"query": " ", "inputs": {}}, + invoke_from=InvokeFrom.WEB_APP, + session=_session(), + ) + + class TestGenerateSuccess: + def test_session_scope_config_version_id_preserves_draft_or_snapshot_id(self): + assert ( + AgentAppGenerator._session_scope_config_version_id( + invoke_from=InvokeFrom.DEBUGGER, config_version_id="draft-1" + ) + == "draft-1" + ) + assert ( + AgentAppGenerator._session_scope_config_version_id( + invoke_from=InvokeFrom.WEB_APP, config_version_id="snapshot-1" + ) + == "snapshot-1" + ) + + def test_generate_orchestrates_and_starts_worker(self, generator, mocker: MockerFixture): + config = AppModelConfig(app_id="app1") + config.id = "config-1" + session = _session() + session.add(config) session.commit() - mocker.patch(f"{MODULE}.DifyRunContext", return_value=mocker.MagicMock()) - mocker.patch(f"{MODULE}.AgentAppRuntimeRequestBuilder", return_value=mocker.MagicMock()) - mocker.patch(f"{MODULE}.create_agent_backend_run_client", return_value=mocker.MagicMock()) - mocker.patch(f"{MODULE}.AgentBackendRunEventAdapter", return_value=mocker.MagicMock()) - mocker.patch(f"{MODULE}.AgentAppWorkspaceStore", return_value=mocker.MagicMock()) - runner = mocker.MagicMock() - if run_side_effect is not None: - runner.run.side_effect = run_side_effect - mocker.patch(f"{MODULE}.AgentAppRunner", return_value=runner) - return runner, resolver_sessions + app_model = _app_generator_app(app_model_config_id=config.id) + user = _app_generator_account() - def _call( - self, - generator, - mocker: MockerFixture, - queue_manager, - *, - query="query", - session_scope_config_version_id="s", - files=(), - file_upload_config=None, - ): - generator._generate_worker( - flask_app=mocker.MagicMock(), - context=mocker.MagicMock(), - application_generate_entity=mocker.MagicMock( - app_config=SimpleNamespace(app_id="app1", tenant_id="tenant"), - agent_id="a", - agent_config_snapshot_id="s", - agent_session_scope_config_version_id=session_scope_config_version_id, - model_conf=mocker.MagicMock(model="m"), - query=query, - files=files, - file_upload_config=file_upload_config, - ), - queue_manager=queue_manager, - conversation_id="conv", - message_id="msg", - user_from=UserFrom.END_USER, - ) + generator._resolve_agent = mocker.MagicMock(return_value=(_agent(), "snap1", "snapshot", AgentSoulConfig())) + generator._prepare_user_inputs = mocker.MagicMock(return_value={"x": 1}) + generator._init_generate_records = mocker.MagicMock(return_value=(_conversation(), _message())) + generator._handle_response = mocker.MagicMock(return_value="raw-response") - def test_happy_path_runs_backend(self, generator: AgentAppGenerator, mocker: MockerFixture): - runner, resolver_sessions = self._wire(generator, mocker) - queue_manager = mocker.MagicMock() - self._call(generator, mocker, queue_manager) - runner.run.assert_called_once() - assert resolver_sessions == [generator._resolve_agent_by_id.call_args.kwargs["session"]] - assert resolver_sessions[0].get_bind() is not None - assert runner.run.call_args.kwargs["home_snapshot_id"] == "home-1" - assert "home_snapshot_ref" not in runner.run.call_args.kwargs - queue_manager.publish_error.assert_not_called() + mocker.patch( + f"{MODULE}.AgentAppConfigManager.get_app_config", + return_value=mocker.MagicMock(variables=[], tenant_id="tenant", app_id="app1"), + ) + mocker.patch(f"{MODULE}.ModelConfigConverter.convert", return_value=mocker.MagicMock(model="gpt-4o-mini")) + file_upload_config = mocker.MagicMock() + mocker.patch(f"{MODULE}.FileUploadConfigManager.convert", return_value=file_upload_config) + parsed_file = mocker.MagicMock() + build_files = mocker.patch(f"{MODULE}.file_factory.build_from_mappings", return_value=[parsed_file]) + mocker.patch(f"{MODULE}.TraceQueueManager", return_value=mocker.MagicMock()) + generate_entity = mocker.patch( + f"{MODULE}.AgentAppGenerateEntity", return_value=mocker.MagicMock(task_id="t", user_id="user") + ) + mocker.patch(f"{MODULE}.MessageBasedAppQueueManager", return_value=mocker.MagicMock()) + thread_obj = mocker.MagicMock() + thread_constructor = mocker.patch(f"{MODULE}.threading.Thread", return_value=thread_obj) + mocker.patch(f"{MODULE}.AgentAppGenerateResponseConverter.convert", return_value={"result": "ok"}) + file_mappings = [ + { + "type": "image", + "transfer_method": "local_file", + "url": "", + "upload_file_id": "upload-file-1", + } + ] - def test_worker_passes_session_scope_config_version_to_runner(self, generator, mocker: MockerFixture): - runner, _ = self._wire(generator, mocker) - queue_manager = mocker.MagicMock() + result = generator.generate( + app_model=app_model, + user=user, + args={"query": "hello", "inputs": {"name": "world"}, "files": file_mappings}, + invoke_from=InvokeFrom.WEB_APP, + session=session, + streaming=True, + ) - self._call(generator, mocker, queue_manager, session_scope_config_version_id=None) + assert result == {"result": "ok"} + thread_obj.start.assert_called_once() + worker_call = thread_constructor.call_args + inspect.signature(worker_call.kwargs["target"]).bind(**worker_call.kwargs["kwargs"]) + generator._resolve_agent.assert_called_once_with( + app_model, + invoke_from=InvokeFrom.WEB_APP, + draft_type=None, + user=user, + session=session, + conversation=None, + ) + assert session.get(AppModelConfig, "config-1") is config + build_files.assert_called_once() + assert build_files.call_args.kwargs["mappings"] == file_mappings + assert generate_entity.call_args.kwargs["files"] == [parsed_file] + assert generate_entity.call_args.kwargs["file_upload_config"] is file_upload_config + assert "agent_runtime_exit_intent" not in generate_entity.call_args.kwargs - assert runner.run.call_args.kwargs["agent_config_snapshot_id"] == "s" - assert runner.run.call_args.kwargs["session_scope_snapshot_id"] is None + def test_generate_loads_existing_conversation(self, generator: AgentAppGenerator, mocker: MockerFixture): + app_model = _app_generator_app() + generator._resolve_agent = mocker.MagicMock( + return_value=(_agent(agent_id="a"), "snap1", "snapshot", AgentSoulConfig()) + ) + generator._prepare_user_inputs = mocker.MagicMock(return_value={}) + generator._init_generate_records = mocker.MagicMock(return_value=(_conversation(), _message())) + generator._handle_response = mocker.MagicMock(return_value="raw") + get_conv = mocker.patch(f"{MODULE}.ConversationService.get_conversation", return_value=_conversation()) + mocker.patch(f"{MODULE}.AgentAppConfigManager.get_app_config", return_value=mocker.MagicMock(variables=[])) + mocker.patch(f"{MODULE}.load_annotation_reply_config", return_value={"enabled": False}) + mocker.patch(f"{MODULE}.ModelConfigConverter.convert", return_value=mocker.MagicMock()) + mocker.patch(f"{MODULE}.TraceQueueManager", return_value=mocker.MagicMock()) + mocker.patch(f"{MODULE}.AgentAppGenerateEntity", return_value=mocker.MagicMock()) + mocker.patch(f"{MODULE}.MessageBasedAppQueueManager", return_value=mocker.MagicMock()) + mocker.patch(f"{MODULE}.threading.Thread", return_value=mocker.MagicMock()) + mocker.patch(f"{MODULE}.AgentAppGenerateResponseConverter.convert", return_value={"result": "ok"}) + session = _session() + user = _app_generator_account() - def test_worker_passes_files_to_backend_runner_without_rewriting_query(self, generator, mocker: MockerFixture): - runner, _ = self._wire(generator, mocker, guard_query="你看得见这张图片吗") - queue_manager = mocker.MagicMock() - files = [mocker.MagicMock(), mocker.MagicMock()] + generator.generate( + app_model=app_model, + user=user, + args={"query": "hi", "inputs": {}, "conversation_id": "conv"}, + invoke_from=InvokeFrom.WEB_APP, + session=session, + streaming=True, + ) - self._call( + get_conv.assert_called_once_with( + app_model=app_model, + conversation_id="conv", + user=user, + session=session, + ) + assert generator._resolve_agent.call_args.kwargs["conversation"].id == "conv" + assert generator._init_generate_records.call_args.kwargs["session"] is session + + def test_generate_does_not_include_trace_session_id_in_extras( + self, generator: AgentAppGenerator, mocker: MockerFixture + ): + app_model = _app_generator_app() + user = _app_generator_account() + + generator._resolve_agent = mocker.MagicMock(return_value=(_agent(), "snap1", "snapshot", AgentSoulConfig())) + generator._prepare_user_inputs = mocker.MagicMock(return_value={}) + generator._init_generate_records = mocker.MagicMock(return_value=(_conversation(), _message())) + generator._handle_response = mocker.MagicMock(return_value="raw-response") + + mocker.patch( + f"{MODULE}.AgentAppConfigManager.get_app_config", + return_value=mocker.MagicMock(variables=[], tenant_id="tenant", app_id="app1"), + ) + mocker.patch(f"{MODULE}.ModelConfigConverter.convert", return_value=mocker.MagicMock(model="gpt-4o-mini")) + mocker.patch(f"{MODULE}.TraceQueueManager", return_value=mocker.MagicMock()) + generate_entity = mocker.patch( + f"{MODULE}.AgentAppGenerateEntity", return_value=mocker.MagicMock(task_id="t", user_id="user") + ) + mocker.patch(f"{MODULE}.MessageBasedAppQueueManager", return_value=mocker.MagicMock()) + mocker.patch(f"{MODULE}.threading.Thread", return_value=mocker.MagicMock()) + mocker.patch(f"{MODULE}.AgentAppGenerateResponseConverter.convert", return_value={"result": "ok"}) + + generator.generate( + app_model=app_model, + user=user, + args={"query": "hello", "inputs": {}, "trace_session_id": "session-1"}, + invoke_from=InvokeFrom.WEB_APP, + session=_session(), + streaming=True, + ) + + assert generate_entity.call_args.kwargs["extras"] == {"auto_generate_conversation_name": True} + + class TestGenerateWorker: + @pytest.fixture(autouse=True) + def patch_context(self, mocker: MockerFixture): + @contextlib.contextmanager + def ctx_manager[**P](*args: P.args, **kwargs: P.kwargs): + yield + + mocker.patch("libs.flask_utils.preserve_flask_contexts", ctx_manager) + + def _wire( + self, + generator: AgentAppGenerator, + mocker: MockerFixture, + *, + run_side_effect=None, + handled=False, + guard_query="query", + ): + generator._get_conversation = mocker.MagicMock(return_value=_conversation()) + generator._get_message = mocker.MagicMock(return_value=_message()) + generator._run_input_guards = mocker.MagicMock(return_value=(handled, guard_query, None)) + resolved_agent = _agent(agent_id="a") + resolved_config = _app_generator_snapshot(snapshot_id="s", agent_id="a") + resolver_sessions: list[Session] = [] + + def resolve_agent_by_id(**kwargs): + resolver_sessions.append(kwargs["session"]) + return resolved_agent, resolved_config, AgentSoulConfig() + + generator._resolve_agent_by_id = mocker.MagicMock(side_effect=resolve_agent_by_id) + session = _session() + if session.get(App, "app1") is None: + session.add(_app_generator_app()) + session.commit() + mocker.patch(f"{MODULE}.DifyRunContext", return_value=mocker.MagicMock()) + mocker.patch(f"{MODULE}.AgentAppRuntimeRequestBuilder", return_value=mocker.MagicMock()) + mocker.patch(f"{MODULE}.create_agent_backend_run_client", return_value=mocker.MagicMock()) + mocker.patch(f"{MODULE}.AgentBackendRunEventAdapter", return_value=mocker.MagicMock()) + mocker.patch(f"{MODULE}.AgentAppWorkspaceStore", return_value=mocker.MagicMock()) + runner = mocker.MagicMock() + if run_side_effect is not None: + runner.run.side_effect = run_side_effect + mocker.patch(f"{MODULE}.AgentAppRunner", return_value=runner) + return runner, resolver_sessions + + def _call( + self, generator, - mocker, + mocker: MockerFixture, queue_manager, - query="你看得见这张图片吗", - files=files, - ) + *, + query="query", + session_scope_config_version_id="s", + files=(), + file_upload_config=None, + ): + generator._generate_worker( + flask_app=mocker.MagicMock(), + context=mocker.MagicMock(), + application_generate_entity=mocker.MagicMock( + app_config=SimpleNamespace(app_id="app1", tenant_id="tenant"), + agent_id="a", + agent_config_snapshot_id="s", + agent_session_scope_config_version_id=session_scope_config_version_id, + model_conf=mocker.MagicMock(model="m"), + query=query, + files=files, + file_upload_config=file_upload_config, + ), + queue_manager=queue_manager, + conversation_id="conv", + message_id="msg", + user_from=UserFrom.END_USER, + ) - assert runner.run.call_args.kwargs["query"] == "你看得见这张图片吗" - assert runner.run.call_args.kwargs["files"] == files - assert runner.run.call_args.kwargs["image_detail_config"] is None - - def test_input_guard_short_circuit_skips_backend(self, generator, mocker: MockerFixture): - runner, _ = self._wire(generator, mocker, handled=True) - queue_manager = mocker.MagicMock() - self._call(generator, mocker, queue_manager) - runner.run.assert_not_called() - - def test_annotation_reply_publishes_after_guard_transaction_commits(self, generator, mocker: MockerFixture): - runner, _ = self._wire(generator, mocker, handled=True) - annotation_reply = MessageAnnotation( - app_id="app1", - question="query", - content="annotated answer", - account_id="user", - ) - generator._run_input_guards.return_value = (True, "query", annotation_reply) - events: list[str] = [] - queue_manager = mocker.MagicMock() - - def record_commit(_session: Session) -> None: - events.append("commit") - - def publish(event, *_args): - if isinstance(event, QueueAnnotationReplyEvent): - events.append("publish") - - queue_manager.publish.side_effect = publish - - event.listen(type(_session()), "after_commit", record_commit) - try: + def test_happy_path_runs_backend(self, generator: AgentAppGenerator, mocker: MockerFixture): + runner, resolver_sessions = self._wire(generator, mocker) + queue_manager = mocker.MagicMock() self._call(generator, mocker, queue_manager) - finally: - event.remove(type(_session()), "after_commit", record_commit) + runner.run.assert_called_once() + assert resolver_sessions == [generator._resolve_agent_by_id.call_args.kwargs["session"]] + assert resolver_sessions[0].get_bind() is not None + assert runner.run.call_args.kwargs["home_snapshot_id"] == "home-1" + assert "home_snapshot_ref" not in runner.run.call_args.kwargs + queue_manager.publish_error.assert_not_called() - assert events == ["commit", "publish"] - runner.run.assert_not_called() + def test_worker_passes_session_scope_config_version_to_runner(self, generator, mocker: MockerFixture): + runner, _ = self._wire(generator, mocker) + queue_manager = mocker.MagicMock() - def test_generate_task_stopped_is_swallowed(self, generator, mocker: MockerFixture): - self._wire(generator, mocker, run_side_effect=GenerateTaskStoppedError()) - # A locally stopped consumer does not prove runtime cancellation. The - # runner must hand off the typed terminal it actually observed. - mocker.patch( - "core.ops.base_trace_instance.application_tracing_hooks", - side_effect=AssertionError("stop intent must not invent a tracing terminal"), + self._call(generator, mocker, queue_manager, session_scope_config_version_id=None) + + assert runner.run.call_args.kwargs["agent_config_snapshot_id"] == "s" + assert runner.run.call_args.kwargs["session_scope_snapshot_id"] is None + + def test_worker_passes_files_to_backend_runner_without_rewriting_query(self, generator, mocker: MockerFixture): + runner, _ = self._wire(generator, mocker, guard_query="你看得见这张图片吗") + queue_manager = mocker.MagicMock() + files = [mocker.MagicMock(), mocker.MagicMock()] + + self._call( + generator, + mocker, + queue_manager, + query="你看得见这张图片吗", + files=files, + ) + + assert runner.run.call_args.kwargs["query"] == "你看得见这张图片吗" + assert runner.run.call_args.kwargs["files"] == files + assert runner.run.call_args.kwargs["image_detail_config"] is None + + def test_input_guard_short_circuit_skips_backend(self, generator, mocker: MockerFixture): + runner, _ = self._wire(generator, mocker, handled=True) + queue_manager = mocker.MagicMock() + self._call(generator, mocker, queue_manager) + runner.run.assert_not_called() + + def test_annotation_reply_publishes_after_guard_transaction_commits(self, generator, mocker: MockerFixture): + runner, _ = self._wire(generator, mocker, handled=True) + annotation_reply = MessageAnnotation( + app_id="app1", + question="query", + content="annotated answer", + account_id="user", + ) + generator._run_input_guards.return_value = (True, "query", annotation_reply) + events: list[str] = [] + queue_manager = mocker.MagicMock() + + def record_commit(_session: Session) -> None: + events.append("commit") + + def publish(event, *_args): + if isinstance(event, QueueAnnotationReplyEvent): + events.append("publish") + + queue_manager.publish.side_effect = publish + + event.listen(type(_session()), "after_commit", record_commit) + try: + self._call(generator, mocker, queue_manager) + finally: + event.remove(type(_session()), "after_commit", record_commit) + + assert events == ["commit", "publish"] + runner.run.assert_not_called() + + def test_generate_task_stopped_is_swallowed(self, generator, mocker: MockerFixture): + self._wire(generator, mocker, run_side_effect=GenerateTaskStoppedError()) + # A locally stopped consumer does not prove runtime cancellation. The + # runner must hand off the typed terminal it actually observed. + mocker.patch( + "core.ops.base_trace_instance.application_tracing_hooks", + side_effect=AssertionError("stop intent must not invent a tracing terminal"), + ) + queue_manager = mocker.MagicMock() + self._call(generator, mocker, queue_manager) + queue_manager.publish_error.assert_not_called() + + def test_unexpected_error_is_published(self, generator: AgentAppGenerator, mocker: MockerFixture): + self._wire(generator, mocker, run_side_effect=ValueError("boom")) + queue_manager = mocker.MagicMock() + self._call(generator, mocker, queue_manager) + assert queue_manager.publish_error.called + + +"""Unit tests for AgentAppGenerator._run_input_guards. + +The guards apply content moderation + annotation reply before the Agent backend +call, mirroring the EasyUI chat runners. They can short-circuit the turn (a +blocked/preset moderation answer or a matched annotation) and publish a direct +answer, or pass through a possibly moderation-sanitized query. +""" + + +class _FakeQueueManager: + def __init__(self) -> None: + self.events: list[Any] = [] + + def publish(self, event: Any, _from: Any) -> None: + self.events.append(event) + + +def _make_entity(query: str = "hello") -> SimpleNamespace: + return SimpleNamespace( + app_config=SimpleNamespace(app_id="app-1", tenant_id="tenant-1"), + model_conf=SimpleNamespace(model="gpt-4o-mini"), + inputs={}, + query=query, + trace_manager=None, + user_id="user-1", + invoke_from=SimpleNamespace(), + ) + + +def _input_guards_app() -> App: + return make_app(name="Agent App", mode=AppMode.AGENT_CHAT, icon_type=None) + + +def _input_guards_message() -> Message: + return make_message(message_id="msg-1") + + +def _patch_moderation(monkeypatch: pytest.MonkeyPatch, *, returns=None, raises: Exception | None = None) -> None: + class _FakeModeration: + def check(self, **kwargs: Any): + if raises is not None: + raise raises + return returns + + monkeypatch.setattr(input_moderation_mod, "InputModeration", _FakeModeration) + + +def _patch_annotation(monkeypatch: pytest.MonkeyPatch, *, reply=None) -> None: + class _FakeAnnotation: + def query(self, **kwargs: Any): + return reply + + monkeypatch.setattr(annotation_mod, "AnnotationReplyFeature", _FakeAnnotation) + + +def _answer_text(events: list[Any]) -> str: + end = next(e for e in events if isinstance(e, QueueMessageEndEvent)) + return end.llm_result.message.content + + +def _saved_user_query(events: list[Any]) -> str: + end = next(e for e in events if isinstance(e, QueueMessageEndEvent)) + prompt_messages = end.llm_result.prompt_messages + assert len(prompt_messages) == 1 + return prompt_messages[0].content + + +class TestInputGuardsCases: + class TestRunInputGuards: + def test_no_guards_passes_through(self, monkeypatch: pytest.MonkeyPatch, sqlite_session: Session): + _patch_moderation(monkeypatch, returns=(False, {}, "hello")) + _patch_annotation(monkeypatch, reply=None) + qm = _FakeQueueManager() + + handled, query, annotation_reply = AgentAppGenerator()._run_input_guards( + session=sqlite_session, + application_generate_entity=_make_entity("hello"), + app_model=_input_guards_app(), + message=_input_guards_message(), + queue_manager=qm, + ) + + assert handled is False + assert query == "hello" + assert annotation_reply is None + assert qm.events == [] + + def test_moderation_override_sanitizes_query(self, monkeypatch: pytest.MonkeyPatch, sqlite_session: Session): + _patch_moderation(monkeypatch, returns=(True, {}, "[redacted]")) + _patch_annotation(monkeypatch, reply=None) + qm = _FakeQueueManager() + + handled, query, annotation_reply = AgentAppGenerator()._run_input_guards( + session=sqlite_session, + application_generate_entity=_make_entity("leak my secret"), + app_model=_input_guards_app(), + message=_input_guards_message(), + queue_manager=qm, + ) + + assert handled is False + assert query == "[redacted]" + assert annotation_reply is None + assert qm.events == [] + + def test_moderation_block_short_circuits(self, monkeypatch: pytest.MonkeyPatch, sqlite_session: Session): + _patch_moderation(monkeypatch, raises=ModerationError("blocked preset answer")) + _patch_annotation(monkeypatch, reply=None) + qm = _FakeQueueManager() + + handled, _, annotation_reply = AgentAppGenerator()._run_input_guards( + session=sqlite_session, + application_generate_entity=_make_entity("forbidden"), + app_model=_input_guards_app(), + message=_input_guards_message(), + queue_manager=qm, + ) + + assert handled is True + assert annotation_reply is None + assert any(isinstance(e, QueueLLMChunkEvent) for e in qm.events) + assert _answer_text(qm.events) == "blocked preset answer" + assert _saved_user_query(qm.events) == "forbidden" + + def test_annotation_hit_short_circuits(self, monkeypatch: pytest.MonkeyPatch, sqlite_session: Session): + _patch_moderation(monkeypatch, returns=(False, {}, "what is your name")) + annotation = MessageAnnotation( + app_id="app-1", + question="what is your name", + content="I am the annotated Iris.", + account_id="account-1", + ) + annotation.id = "anno-1" + _patch_annotation(monkeypatch, reply=annotation) + qm = _FakeQueueManager() + + handled, _, annotation_reply = AgentAppGenerator()._run_input_guards( + session=sqlite_session, + application_generate_entity=_make_entity("what is your name"), + app_model=_input_guards_app(), + message=_input_guards_message(), + queue_manager=qm, + ) + + assert handled is True + assert annotation_reply.id == "anno-1" + assert qm.events == [] + + +"""Unit tests for AgentAppGenerator agent/snapshot resolution. + +Covers the DB-backed resolution helpers (the bound roster Agent + its published +Agent Soul snapshot) including every not-found error path, using persisted rows +in the shared SQLite test database. +""" + + +_SOUL_DICT = { + "model": { + "plugin_id": "langgenius/openai", + "model_provider": "langgenius/openai/openai", + "model": "gpt-4o-mini", + }, + "prompt": {"system_prompt": "You are Iris."}, +} + + +def _resolve_agent_agent( + *, + agent_id: str = "agent-1", + scope: AgentScope = AgentScope.ROSTER, + source: AgentSource = AgentSource.AGENT_APP, + active_snapshot_id: str | None = "snap-1", + published: bool = True, +) -> Agent: + return Agent( + id=agent_id, + tenant_id="t1", + name=f"Agent {agent_id}", + scope=scope, + source=source, + app_id="app-1" if scope == AgentScope.ROSTER else None, + backing_app_id="app-1" if scope == AgentScope.WORKFLOW_ONLY else None, + active_config_snapshot_id=active_snapshot_id, + active_config_has_model=True, + active_config_is_published=published, + created_by="creator-1", + updated_by="updater-1", + ) + + +def _resolve_agent_snapshot( + *, snapshot_id: str = "snap-1", agent_id: str = "agent-1", home_snapshot_id: str = "home-1" +) -> AgentConfigSnapshot: + return AgentConfigSnapshot( + id=snapshot_id, + tenant_id="t1", + agent_id=agent_id, + version=1, + config_snapshot=AgentSoulConfig.model_validate(_SOUL_DICT), + home_snapshot_id=home_snapshot_id, + created_by="creator-1", + ) + + +def _resolve_agent_app() -> App: + return make_app( + tenant_id="t1", + name="Agent App", + mode=AppMode.AGENT_CHAT, + icon_type=None, + enable_site=False, + enable_api=False, + ) + + +def _resolve_agent_account() -> _resolve_agent_Account: + return make_account(account_id="user-1", name="Agent User", email="agent-user@example.com") + + +def _resolve_agent_conversation(*, binding_id: str) -> Conversation: + return make_conversation(agent_workspace_binding_id=binding_id) + + +def _binding( + *, + home_snapshot_id: str = "home-1", + version_id: str = "snap-1", + version_kind: AgentConfigVersionKind = AgentConfigVersionKind.SNAPSHOT, +) -> AgentWorkspaceBinding: + return AgentWorkspaceBinding( + id="binding-1", + tenant_id="t1", + app_id="app-1", + workspace_id="workspace-1", + agent_id="agent-1", + base_home_snapshot_id=home_snapshot_id, + agent_config_version_id=version_id, + agent_config_version_kind=version_kind, + backend_binding_ref="backend-binding-1", + status=AgentWorkingResourceStatus.ACTIVE, + ) + + +def _persist(session: Session, *models: object) -> None: + session.add_all(models) + session.commit() + + +class TestResolveAgentCases: + class TestResolveAgentById: + def test_success_returns_agent_snapshot_soul(self, sqlite_session: Session): + agent = _resolve_agent_agent() + snapshot = _resolve_agent_snapshot() + sqlite_session.add_all([agent, snapshot]) + sqlite_session.commit() + + resolved_agent, resolved_snapshot, soul = AgentAppGenerator._resolve_agent_by_id( + tenant_id="t1", agent_id="agent-1", snapshot_id="snap-1", session=sqlite_session + ) + + assert resolved_agent is agent + assert resolved_snapshot is snapshot + assert soul.prompt.system_prompt == "You are Iris." + assert soul.model is not None + assert soul.model.model == "gpt-4o-mini" + + def test_agent_missing_raises(self, sqlite_session: Session): + with pytest.raises(AgentAppGeneratorError, match="Agent not found"): + AgentAppGenerator._resolve_agent_by_id( + tenant_id="t1", agent_id="x", snapshot_id="snap-1", session=sqlite_session + ) + + def test_no_published_version_raises(self, sqlite_session: Session): + sqlite_session.add(_resolve_agent_agent(active_snapshot_id=None)) + sqlite_session.commit() + with pytest.raises(AgentAppGeneratorError, match="no published version"): + AgentAppGenerator._resolve_agent_by_id( + tenant_id="t1", agent_id="agent-1", snapshot_id=None, session=sqlite_session + ) + + def test_snapshot_missing_raises(self, sqlite_session: Session): + sqlite_session.add(_resolve_agent_agent()) + sqlite_session.commit() + with pytest.raises(AgentAppGeneratorError, match="published version not found"): + AgentAppGenerator._resolve_agent_by_id( + tenant_id="t1", + agent_id="agent-1", + snapshot_id="snap-1", + session=sqlite_session, + ) + + class TestResolveDebugDraft: + def test_missing_shared_draft_is_created_with_supplied_session(self, sqlite_session: Session): + agent = _resolve_agent_agent() + sqlite_session.add_all([agent, _resolve_agent_snapshot()]) + sqlite_session.commit() + flush_count = 0 + + def _record_flush(_session: Session, _context: object) -> None: + nonlocal flush_count + flush_count += 1 + + event.listen(sqlite_session, "after_flush", _record_flush) + + draft = AgentAppGenerator._resolve_debug_draft( + tenant_id="t1", + agent=agent, + draft_type=None, + draft_id=None, + account_id=None, + session=sqlite_session, + ) + + assert draft.draft_type == AgentConfigDraftType.DRAFT + assert draft.base_snapshot_id == "snap-1" + assert sqlite_session.scalar(select(AgentConfigDraft).where(AgentConfigDraft.id == draft.id)) is draft + assert flush_count == 1 + + def test_stale_workflow_only_shared_draft_is_rebased_to_active_snapshot(self, sqlite_session: Session): + agent = _resolve_agent_agent(scope=AgentScope.WORKFLOW_ONLY, active_snapshot_id="snap-2") + draft = AgentConfigDraft( + id="draft-1", + tenant_id="t1", + agent_id="agent-1", + draft_type=AgentConfigDraftType.DRAFT, + account_id=None, + draft_owner_key="", + base_snapshot_id="snap-1", + home_snapshot_id="home-1", + config_snapshot=AgentSoulConfig.model_validate({"prompt": {"system_prompt": "old"}}), + ) + active_snapshot = AgentConfigSnapshot( + id="snap-2", + tenant_id="t1", + agent_id="agent-1", + version=2, + config_snapshot=AgentSoulConfig.model_validate({"prompt": {"system_prompt": "new"}}), + home_snapshot_id="home-2", + created_by="creator-1", + ) + sqlite_session.add_all([agent, draft, active_snapshot]) + sqlite_session.commit() + flush_count = 0 + + def _record_flush(_session: Session, _context: object) -> None: + nonlocal flush_count + flush_count += 1 + + event.listen(sqlite_session, "after_flush", _record_flush) + + resolved = AgentAppGenerator._resolve_debug_draft( + tenant_id="t1", + agent=agent, + draft_type=None, + account_id=None, + session=sqlite_session, + ) + + assert resolved is draft + assert resolved.id == "draft-1" + assert resolved.base_snapshot_id == "snap-2" + assert resolved.home_snapshot_id == "home-2" + assert resolved.config_snapshot_dict["prompt"]["system_prompt"] == "new" + assert flush_count == 1 + + def test_build_draft_is_not_rebased_to_active_snapshot(self, sqlite_session: Session): + agent = _resolve_agent_agent(scope=AgentScope.WORKFLOW_ONLY, active_snapshot_id="snap-2") + draft = AgentConfigDraft( + id="build-draft-1", + tenant_id="t1", + agent_id="agent-1", + draft_type=AgentConfigDraftType.DEBUG_BUILD, + account_id="account-1", + draft_owner_key="account-1", + base_snapshot_id="snap-1", + home_snapshot_id="home-build", + config_snapshot=AgentSoulConfig.model_validate({"prompt": {"system_prompt": "build edit"}}), + ) + sqlite_session.add_all([agent, draft]) + sqlite_session.commit() + + resolved = AgentAppGenerator._resolve_debug_draft( + tenant_id="t1", + agent=agent, + draft_type=AgentConfigDraftType.DEBUG_BUILD.value, + account_id="account-1", + session=sqlite_session, + ) + + assert resolved is draft + assert resolved.base_snapshot_id == "snap-1" + assert resolved.config_snapshot_dict["prompt"]["system_prompt"] == "build edit" + assert not sqlite_session.dirty + + def test_build_draft_uses_exact_draft_id(self, sqlite_session: Session): + agent = _resolve_agent_agent(scope=AgentScope.WORKFLOW_ONLY, active_snapshot_id="snap-2") + draft = AgentConfigDraft( + id="exact-build-draft", + tenant_id="t1", + agent_id="agent-1", + draft_type=AgentConfigDraftType.DEBUG_BUILD, + account_id="account-1", + draft_owner_key="account-1", + base_snapshot_id="snap-1", + home_snapshot_id="home-build", + config_snapshot=AgentSoulConfig(), + ) + newer_decoy = AgentConfigDraft( + id="newer-build-draft", + tenant_id="t1", + agent_id="agent-1", + draft_type=AgentConfigDraftType.DEBUG_BUILD, + account_id="account-1", + draft_owner_key="account-1-decoy", + base_snapshot_id="snap-2", + home_snapshot_id="home-decoy", + config_snapshot=AgentSoulConfig(prompt={"system_prompt": "decoy"}), + ) + sqlite_session.add_all([agent, draft, newer_decoy]) + sqlite_session.commit() + + resolved = AgentAppGenerator._resolve_debug_draft( + tenant_id="t1", + agent=agent, + draft_type=AgentConfigDraftType.DEBUG_BUILD.value, + draft_id="exact-build-draft", + account_id="account-1", + session=sqlite_session, + ) + + assert resolved is draft + assert resolved.id == "exact-build-draft" + + class TestResolveAgent: + @pytest.fixture(autouse=True) + def _publish_visibility(self, monkeypatch: pytest.MonkeyPatch) -> None: + def is_publish_visible(*, agent: Agent, **_kwargs: object) -> bool: + return agent.active_config_is_published + + monkeypatch.setattr( + app_generator, + "agent_has_workflow_callable_active_snapshot", + is_publish_visible, + ) + + def test_success_chains_to_resolve_by_id(self, sqlite_session: Session): + bound_agent = _resolve_agent_agent() + snapshot = _resolve_agent_snapshot() + _persist(sqlite_session, bound_agent, snapshot) + app_model = _resolve_agent_app() + + agent, config_id, config_version_kind, soul = AgentAppGenerator()._resolve_agent( + app_model, + invoke_from=InvokeFrom.WEB_APP, + draft_type=None, + user=_resolve_agent_account(), + session=sqlite_session, + ) + + assert agent is bound_agent + assert config_id == snapshot.id + assert config_version_kind == "snapshot" + assert soul.model is not None + + def test_unpublished_draft_still_resolves_active_snapshot( + self, monkeypatch: pytest.MonkeyPatch, sqlite_session: Session + ): + bound_agent = _resolve_agent_agent(published=False) + snapshot = _resolve_agent_snapshot() + _persist(sqlite_session, bound_agent, snapshot) + app_model = _resolve_agent_app() + monkeypatch.setattr(app_generator, "agent_has_workflow_callable_active_snapshot", lambda **_kwargs: True) + + agent, config_id, config_version_kind, soul = AgentAppGenerator()._resolve_agent( + app_model, + invoke_from=InvokeFrom.WEB_APP, + draft_type=None, + user=_resolve_agent_account(), + session=sqlite_session, + ) + + assert agent is bound_agent + assert config_id == snapshot.id + assert config_version_kind == "snapshot" + assert soul.prompt.system_prompt == "You are Iris." + + def test_existing_conversation_resolves_binding_snapshot_instead_of_latest_active_snapshot( + self, monkeypatch: pytest.MonkeyPatch, sqlite_session: Session + ): + bound_agent = _resolve_agent_agent(active_snapshot_id="snap-2") + pinned_snapshot = _resolve_agent_snapshot() + _persist(sqlite_session, bound_agent, pinned_snapshot) + conversation = _resolve_agent_conversation(binding_id="binding-1") + binding = _binding() + get_active_binding = MagicMock(return_value=binding) + validate_generation = MagicMock() + monkeypatch.setattr(AgentWorkspaceService, "get_active_binding", get_active_binding) + monkeypatch.setattr(AgentWorkspaceService, "validate_binding_generation", validate_generation) + app_model = _resolve_agent_app() + + _, config_id, config_version_kind, soul = AgentAppGenerator()._resolve_agent( + app_model, + invoke_from=InvokeFrom.WEB_APP, + draft_type=None, + user=_resolve_agent_account(), + session=sqlite_session, + conversation=conversation, + ) + + assert config_id == "snap-1" + assert config_version_kind == "snapshot" + assert soul.prompt.system_prompt == "You are Iris." + assert get_active_binding.call_args.kwargs["binding_id"] == "binding-1" + validate_generation.assert_called_once_with( + binding, + base_home_snapshot_id="home-1", + agent_config_version_id="snap-1", + agent_config_version_kind=AgentConfigVersionKind.SNAPSHOT, + ) + + def test_existing_conversation_rejects_unavailable_binding( + self, monkeypatch: pytest.MonkeyPatch, sqlite_session: Session + ): + bound_agent = _resolve_agent_agent(active_snapshot_id="snap-active") + _persist(sqlite_session, bound_agent) + conversation = _resolve_agent_conversation(binding_id="binding-missing") + monkeypatch.setattr(AgentWorkspaceService, "get_active_binding", MagicMock(return_value=None)) + + with pytest.raises(AgentAppGeneratorError, match="Conversation participant Binding is unavailable"): + AgentAppGenerator()._resolve_agent( + _resolve_agent_app(), + invoke_from=InvokeFrom.WEB_APP, + draft_type=None, + user=_resolve_agent_account(), + session=sqlite_session, + conversation=conversation, + ) + + @pytest.mark.parametrize( + ("binding_home_id", "binding_version_kind", "snapshot_home_id"), + [ + ("home-binding", AgentConfigVersionKind.SNAPSHOT, "home-other"), + ("home-pinned", AgentConfigVersionKind.DRAFT, "home-pinned"), + ], ) - queue_manager = mocker.MagicMock() - self._call(generator, mocker, queue_manager) - queue_manager.publish_error.assert_not_called() + def test_existing_conversation_generation_mismatch_does_not_fallback_to_active_snapshot( + self, + monkeypatch: pytest.MonkeyPatch, + sqlite_session: Session, + binding_home_id: str, + binding_version_kind: AgentConfigVersionKind, + snapshot_home_id: str, + ): + bound_agent = _resolve_agent_agent(active_snapshot_id="snap-active") + conversation = _resolve_agent_conversation(binding_id="binding-1") + binding = _binding( + home_snapshot_id=binding_home_id, + version_id="snap-pinned", + version_kind=binding_version_kind, + ) + pinned_snapshot = _resolve_agent_snapshot(snapshot_id="snap-pinned", home_snapshot_id=snapshot_home_id) + _persist(sqlite_session, bound_agent, pinned_snapshot) + monkeypatch.setattr(AgentWorkspaceService, "get_active_binding", MagicMock(return_value=binding)) - def test_unexpected_error_is_published(self, generator: AgentAppGenerator, mocker: MockerFixture): - self._wire(generator, mocker, run_side_effect=ValueError("boom")) - queue_manager = mocker.MagicMock() - self._call(generator, mocker, queue_manager) - assert queue_manager.publish_error.called + with pytest.raises(AgentWorkspaceBindingGenerationMismatchError): + AgentAppGenerator()._resolve_agent( + _resolve_agent_app(), + invoke_from=InvokeFrom.WEB_APP, + draft_type=None, + user=_resolve_agent_account(), + session=sqlite_session, + conversation=conversation, + ) + + def test_unpublished_imported_agent_is_not_available_to_public_runtime(self, sqlite_session: Session): + bound_agent = _resolve_agent_agent(source=AgentSource.IMPORTED, published=False) + _persist(sqlite_session, bound_agent) + app_model = _resolve_agent_app() + + with pytest.raises(AgentAppNotPublishedError, match="not been published"): + AgentAppGenerator()._resolve_agent( + app_model, + invoke_from=InvokeFrom.WEB_APP, + draft_type=None, + user=_resolve_agent_account(), + session=sqlite_session, + ) + + def test_never_published_agent_app_is_not_available_to_public_runtime(self, sqlite_session: Session): + bound_agent = _resolve_agent_agent(published=False) + _persist(sqlite_session, bound_agent) + + with pytest.raises(AgentAppNotPublishedError, match="not been published"): + AgentAppGenerator()._resolve_agent( + _resolve_agent_app(), + invoke_from=InvokeFrom.WEB_APP, + draft_type=None, + user=_resolve_agent_account(), + session=sqlite_session, + ) + + def test_unpublished_imported_agent_remains_available_to_debugger(self, sqlite_session: Session): + bound_agent = _resolve_agent_agent(source=AgentSource.IMPORTED, published=False) + draft = AgentConfigDraft( + id="draft-1", + tenant_id="t1", + agent_id="agent-1", + draft_type=AgentConfigDraftType.DRAFT, + account_id=None, + draft_owner_key="", + base_snapshot_id="snap-1", + home_snapshot_id="home-1", + config_snapshot=AgentSoulConfig.model_validate(_SOUL_DICT), + ) + _persist(sqlite_session, bound_agent, draft) + app_model = _resolve_agent_app() + + agent, config_id, config_version_kind, soul = AgentAppGenerator()._resolve_agent( + app_model, + invoke_from=InvokeFrom.DEBUGGER, + draft_type=None, + user=_resolve_agent_account(), + session=sqlite_session, + ) + + assert agent is bound_agent + assert config_id == draft.id + assert config_version_kind == "draft" + assert soul.prompt.system_prompt == "You are Iris." + + def test_agent_without_active_snapshot_raises_before_model_resolution(self, sqlite_session: Session): + bound_agent = _resolve_agent_agent(active_snapshot_id=None, published=False) + _persist(sqlite_session, bound_agent) + app_model = _resolve_agent_app() + + with pytest.raises(AgentAppNotPublishedError, match="not been published"): + AgentAppGenerator()._resolve_agent( + app_model, + invoke_from=InvokeFrom.WEB_APP, + draft_type=None, + user=_resolve_agent_account(), + session=sqlite_session, + ) + + def test_unbound_app_raises(self, sqlite_session: Session): + app_model = _resolve_agent_app() + with pytest.raises(AgentAppGeneratorError, match="has no bound Agent"): + AgentAppGenerator()._resolve_agent( + app_model, + invoke_from=InvokeFrom.WEB_APP, + draft_type=None, + user=_resolve_agent_account(), + session=sqlite_session, + ) diff --git a/api/tests/unit_tests/core/app/apps/agent_app/test_input_guards.py b/api/tests/unit_tests/core/app/apps/agent_app/test_input_guards.py deleted file mode 100644 index eaf50049612..00000000000 --- a/api/tests/unit_tests/core/app/apps/agent_app/test_input_guards.py +++ /dev/null @@ -1,165 +0,0 @@ -"""Unit tests for AgentAppGenerator._run_input_guards. - -The guards apply content moderation + annotation reply before the Agent backend -call, mirroring the EasyUI chat runners. They can short-circuit the turn (a -blocked/preset moderation answer or a matched annotation) and publish a direct -answer, or pass through a possibly moderation-sanitized query. -""" - -from __future__ import annotations - -from types import SimpleNamespace -from typing import Any - -import pytest -from sqlalchemy.orm import Session - -import core.app.features.annotation_reply.annotation_reply as annotation_mod -import core.moderation.input_moderation as input_moderation_mod -from core.app.apps.agent_app.app_generator import AgentAppGenerator -from core.app.entities.queue_entities import ( - QueueLLMChunkEvent, - QueueMessageEndEvent, -) -from core.moderation.base import ModerationError -from models.model import App, AppMode, Message, MessageAnnotation -from tests.unit_tests.model_factories import make_app, make_message - - -class _FakeQueueManager: - def __init__(self) -> None: - self.events: list[Any] = [] - - def publish(self, event: Any, _from: Any) -> None: - self.events.append(event) - - -def _make_entity(query: str = "hello") -> SimpleNamespace: - return SimpleNamespace( - app_config=SimpleNamespace(app_id="app-1", tenant_id="tenant-1"), - model_conf=SimpleNamespace(model="gpt-4o-mini"), - inputs={}, - query=query, - trace_manager=None, - user_id="user-1", - invoke_from=SimpleNamespace(), - ) - - -def _app() -> App: - return make_app(name="Agent App", mode=AppMode.AGENT_CHAT, icon_type=None) - - -def _message() -> Message: - return make_message(message_id="msg-1") - - -def _patch_moderation(monkeypatch: pytest.MonkeyPatch, *, returns=None, raises: Exception | None = None) -> None: - class _FakeModeration: - def check(self, **kwargs: Any): - if raises is not None: - raise raises - return returns - - monkeypatch.setattr(input_moderation_mod, "InputModeration", _FakeModeration) - - -def _patch_annotation(monkeypatch: pytest.MonkeyPatch, *, reply=None) -> None: - class _FakeAnnotation: - def query(self, **kwargs: Any): - return reply - - monkeypatch.setattr(annotation_mod, "AnnotationReplyFeature", _FakeAnnotation) - - -def _answer_text(events: list[Any]) -> str: - end = next(e for e in events if isinstance(e, QueueMessageEndEvent)) - return end.llm_result.message.content - - -def _saved_user_query(events: list[Any]) -> str: - end = next(e for e in events if isinstance(e, QueueMessageEndEvent)) - prompt_messages = end.llm_result.prompt_messages - assert len(prompt_messages) == 1 - return prompt_messages[0].content - - -class TestRunInputGuards: - def test_no_guards_passes_through(self, monkeypatch: pytest.MonkeyPatch, sqlite_session: Session): - _patch_moderation(monkeypatch, returns=(False, {}, "hello")) - _patch_annotation(monkeypatch, reply=None) - qm = _FakeQueueManager() - - handled, query, annotation_reply = AgentAppGenerator()._run_input_guards( - session=sqlite_session, - application_generate_entity=_make_entity("hello"), - app_model=_app(), - message=_message(), - queue_manager=qm, - ) - - assert handled is False - assert query == "hello" - assert annotation_reply is None - assert qm.events == [] - - def test_moderation_override_sanitizes_query(self, monkeypatch: pytest.MonkeyPatch, sqlite_session: Session): - _patch_moderation(monkeypatch, returns=(True, {}, "[redacted]")) - _patch_annotation(monkeypatch, reply=None) - qm = _FakeQueueManager() - - handled, query, annotation_reply = AgentAppGenerator()._run_input_guards( - session=sqlite_session, - application_generate_entity=_make_entity("leak my secret"), - app_model=_app(), - message=_message(), - queue_manager=qm, - ) - - assert handled is False - assert query == "[redacted]" - assert annotation_reply is None - assert qm.events == [] - - def test_moderation_block_short_circuits(self, monkeypatch: pytest.MonkeyPatch, sqlite_session: Session): - _patch_moderation(monkeypatch, raises=ModerationError("blocked preset answer")) - _patch_annotation(monkeypatch, reply=None) - qm = _FakeQueueManager() - - handled, _, annotation_reply = AgentAppGenerator()._run_input_guards( - session=sqlite_session, - application_generate_entity=_make_entity("forbidden"), - app_model=_app(), - message=_message(), - queue_manager=qm, - ) - - assert handled is True - assert annotation_reply is None - assert any(isinstance(e, QueueLLMChunkEvent) for e in qm.events) - assert _answer_text(qm.events) == "blocked preset answer" - assert _saved_user_query(qm.events) == "forbidden" - - def test_annotation_hit_short_circuits(self, monkeypatch: pytest.MonkeyPatch, sqlite_session: Session): - _patch_moderation(monkeypatch, returns=(False, {}, "what is your name")) - annotation = MessageAnnotation( - app_id="app-1", - question="what is your name", - content="I am the annotated Iris.", - account_id="account-1", - ) - annotation.id = "anno-1" - _patch_annotation(monkeypatch, reply=annotation) - qm = _FakeQueueManager() - - handled, _, annotation_reply = AgentAppGenerator()._run_input_guards( - session=sqlite_session, - application_generate_entity=_make_entity("what is your name"), - app_model=_app(), - message=_message(), - queue_manager=qm, - ) - - assert handled is True - assert annotation_reply.id == "anno-1" - assert qm.events == [] diff --git a/api/tests/unit_tests/core/app/apps/agent_app/test_resolve_agent.py b/api/tests/unit_tests/core/app/apps/agent_app/test_resolve_agent.py deleted file mode 100644 index 9ed6f108680..00000000000 --- a/api/tests/unit_tests/core/app/apps/agent_app/test_resolve_agent.py +++ /dev/null @@ -1,537 +0,0 @@ -"""Unit tests for AgentAppGenerator agent/snapshot resolution. - -Covers the DB-backed resolution helpers (the bound roster Agent + its published -Agent Soul snapshot) including every not-found error path, using persisted rows -in the shared SQLite test database. -""" - -from __future__ import annotations - -from unittest.mock import MagicMock - -import pytest -from sqlalchemy import event, select -from sqlalchemy.orm import Session - -from core.app.apps.agent_app import app_generator -from core.app.apps.agent_app.app_generator import AgentAppGenerator, AgentAppGeneratorError, AgentAppNotPublishedError -from core.app.entities.app_invoke_entities import InvokeFrom -from models.account import Account -from models.agent import ( - Agent, - AgentConfigDraft, - AgentConfigDraftType, - AgentConfigSnapshot, - AgentConfigVersionKind, - AgentScope, - AgentSource, - AgentWorkingResourceStatus, - AgentWorkspaceBinding, -) -from models.agent_config_entities import AgentSoulConfig -from models.model import App, AppMode, Conversation -from services.agent.workspace_service import ( - AgentWorkspaceBindingGenerationMismatchError, - AgentWorkspaceService, -) -from tests.unit_tests.model_factories import make_account, make_app, make_conversation - -_SOUL_DICT = { - "model": { - "plugin_id": "langgenius/openai", - "model_provider": "langgenius/openai/openai", - "model": "gpt-4o-mini", - }, - "prompt": {"system_prompt": "You are Iris."}, -} - - -def _agent( - *, - agent_id: str = "agent-1", - scope: AgentScope = AgentScope.ROSTER, - source: AgentSource = AgentSource.AGENT_APP, - active_snapshot_id: str | None = "snap-1", - published: bool = True, -) -> Agent: - return Agent( - id=agent_id, - tenant_id="t1", - name=f"Agent {agent_id}", - scope=scope, - source=source, - app_id="app-1" if scope == AgentScope.ROSTER else None, - backing_app_id="app-1" if scope == AgentScope.WORKFLOW_ONLY else None, - active_config_snapshot_id=active_snapshot_id, - active_config_has_model=True, - active_config_is_published=published, - created_by="creator-1", - updated_by="updater-1", - ) - - -def _snapshot( - *, snapshot_id: str = "snap-1", agent_id: str = "agent-1", home_snapshot_id: str = "home-1" -) -> AgentConfigSnapshot: - return AgentConfigSnapshot( - id=snapshot_id, - tenant_id="t1", - agent_id=agent_id, - version=1, - config_snapshot=AgentSoulConfig.model_validate(_SOUL_DICT), - home_snapshot_id=home_snapshot_id, - created_by="creator-1", - ) - - -def _app() -> App: - return make_app( - tenant_id="t1", - name="Agent App", - mode=AppMode.AGENT_CHAT, - icon_type=None, - enable_site=False, - enable_api=False, - ) - - -def _account() -> Account: - return make_account(account_id="user-1", name="Agent User", email="agent-user@example.com") - - -def _conversation(*, binding_id: str) -> Conversation: - return make_conversation(agent_workspace_binding_id=binding_id) - - -def _binding( - *, - home_snapshot_id: str = "home-1", - version_id: str = "snap-1", - version_kind: AgentConfigVersionKind = AgentConfigVersionKind.SNAPSHOT, -) -> AgentWorkspaceBinding: - return AgentWorkspaceBinding( - id="binding-1", - tenant_id="t1", - app_id="app-1", - workspace_id="workspace-1", - agent_id="agent-1", - base_home_snapshot_id=home_snapshot_id, - agent_config_version_id=version_id, - agent_config_version_kind=version_kind, - backend_binding_ref="backend-binding-1", - status=AgentWorkingResourceStatus.ACTIVE, - ) - - -def _persist(session: Session, *models: object) -> None: - session.add_all(models) - session.commit() - - -class TestResolveAgentById: - def test_success_returns_agent_snapshot_soul(self, sqlite_session: Session): - agent = _agent() - snapshot = _snapshot() - sqlite_session.add_all([agent, snapshot]) - sqlite_session.commit() - - resolved_agent, resolved_snapshot, soul = AgentAppGenerator._resolve_agent_by_id( - tenant_id="t1", agent_id="agent-1", snapshot_id="snap-1", session=sqlite_session - ) - - assert resolved_agent is agent - assert resolved_snapshot is snapshot - assert soul.prompt.system_prompt == "You are Iris." - assert soul.model is not None - assert soul.model.model == "gpt-4o-mini" - - def test_agent_missing_raises(self, sqlite_session: Session): - with pytest.raises(AgentAppGeneratorError, match="Agent not found"): - AgentAppGenerator._resolve_agent_by_id( - tenant_id="t1", agent_id="x", snapshot_id="snap-1", session=sqlite_session - ) - - def test_no_published_version_raises(self, sqlite_session: Session): - sqlite_session.add(_agent(active_snapshot_id=None)) - sqlite_session.commit() - with pytest.raises(AgentAppGeneratorError, match="no published version"): - AgentAppGenerator._resolve_agent_by_id( - tenant_id="t1", agent_id="agent-1", snapshot_id=None, session=sqlite_session - ) - - def test_snapshot_missing_raises(self, sqlite_session: Session): - sqlite_session.add(_agent()) - sqlite_session.commit() - with pytest.raises(AgentAppGeneratorError, match="published version not found"): - AgentAppGenerator._resolve_agent_by_id( - tenant_id="t1", - agent_id="agent-1", - snapshot_id="snap-1", - session=sqlite_session, - ) - - -class TestResolveDebugDraft: - def test_missing_shared_draft_is_created_with_supplied_session(self, sqlite_session: Session): - agent = _agent() - sqlite_session.add_all([agent, _snapshot()]) - sqlite_session.commit() - flush_count = 0 - - def _record_flush(_session: Session, _context: object) -> None: - nonlocal flush_count - flush_count += 1 - - event.listen(sqlite_session, "after_flush", _record_flush) - - draft = AgentAppGenerator._resolve_debug_draft( - tenant_id="t1", - agent=agent, - draft_type=None, - draft_id=None, - account_id=None, - session=sqlite_session, - ) - - assert draft.draft_type == AgentConfigDraftType.DRAFT - assert draft.base_snapshot_id == "snap-1" - assert sqlite_session.scalar(select(AgentConfigDraft).where(AgentConfigDraft.id == draft.id)) is draft - assert flush_count == 1 - - def test_stale_workflow_only_shared_draft_is_rebased_to_active_snapshot(self, sqlite_session: Session): - agent = _agent(scope=AgentScope.WORKFLOW_ONLY, active_snapshot_id="snap-2") - draft = AgentConfigDraft( - id="draft-1", - tenant_id="t1", - agent_id="agent-1", - draft_type=AgentConfigDraftType.DRAFT, - account_id=None, - draft_owner_key="", - base_snapshot_id="snap-1", - home_snapshot_id="home-1", - config_snapshot=AgentSoulConfig.model_validate({"prompt": {"system_prompt": "old"}}), - ) - active_snapshot = AgentConfigSnapshot( - id="snap-2", - tenant_id="t1", - agent_id="agent-1", - version=2, - config_snapshot=AgentSoulConfig.model_validate({"prompt": {"system_prompt": "new"}}), - home_snapshot_id="home-2", - created_by="creator-1", - ) - sqlite_session.add_all([agent, draft, active_snapshot]) - sqlite_session.commit() - flush_count = 0 - - def _record_flush(_session: Session, _context: object) -> None: - nonlocal flush_count - flush_count += 1 - - event.listen(sqlite_session, "after_flush", _record_flush) - - resolved = AgentAppGenerator._resolve_debug_draft( - tenant_id="t1", - agent=agent, - draft_type=None, - account_id=None, - session=sqlite_session, - ) - - assert resolved is draft - assert resolved.id == "draft-1" - assert resolved.base_snapshot_id == "snap-2" - assert resolved.home_snapshot_id == "home-2" - assert resolved.config_snapshot_dict["prompt"]["system_prompt"] == "new" - assert flush_count == 1 - - def test_build_draft_is_not_rebased_to_active_snapshot(self, sqlite_session: Session): - agent = _agent(scope=AgentScope.WORKFLOW_ONLY, active_snapshot_id="snap-2") - draft = AgentConfigDraft( - id="build-draft-1", - tenant_id="t1", - agent_id="agent-1", - draft_type=AgentConfigDraftType.DEBUG_BUILD, - account_id="account-1", - draft_owner_key="account-1", - base_snapshot_id="snap-1", - home_snapshot_id="home-build", - config_snapshot=AgentSoulConfig.model_validate({"prompt": {"system_prompt": "build edit"}}), - ) - sqlite_session.add_all([agent, draft]) - sqlite_session.commit() - - resolved = AgentAppGenerator._resolve_debug_draft( - tenant_id="t1", - agent=agent, - draft_type=AgentConfigDraftType.DEBUG_BUILD.value, - account_id="account-1", - session=sqlite_session, - ) - - assert resolved is draft - assert resolved.base_snapshot_id == "snap-1" - assert resolved.config_snapshot_dict["prompt"]["system_prompt"] == "build edit" - assert not sqlite_session.dirty - - def test_build_draft_uses_exact_draft_id(self, sqlite_session: Session): - agent = _agent(scope=AgentScope.WORKFLOW_ONLY, active_snapshot_id="snap-2") - draft = AgentConfigDraft( - id="exact-build-draft", - tenant_id="t1", - agent_id="agent-1", - draft_type=AgentConfigDraftType.DEBUG_BUILD, - account_id="account-1", - draft_owner_key="account-1", - base_snapshot_id="snap-1", - home_snapshot_id="home-build", - config_snapshot=AgentSoulConfig(), - ) - newer_decoy = AgentConfigDraft( - id="newer-build-draft", - tenant_id="t1", - agent_id="agent-1", - draft_type=AgentConfigDraftType.DEBUG_BUILD, - account_id="account-1", - draft_owner_key="account-1-decoy", - base_snapshot_id="snap-2", - home_snapshot_id="home-decoy", - config_snapshot=AgentSoulConfig(prompt={"system_prompt": "decoy"}), - ) - sqlite_session.add_all([agent, draft, newer_decoy]) - sqlite_session.commit() - - resolved = AgentAppGenerator._resolve_debug_draft( - tenant_id="t1", - agent=agent, - draft_type=AgentConfigDraftType.DEBUG_BUILD.value, - draft_id="exact-build-draft", - account_id="account-1", - session=sqlite_session, - ) - - assert resolved is draft - assert resolved.id == "exact-build-draft" - - -class TestResolveAgent: - @pytest.fixture(autouse=True) - def _publish_visibility(self, monkeypatch: pytest.MonkeyPatch) -> None: - def is_publish_visible(*, agent: Agent, **_kwargs: object) -> bool: - return agent.active_config_is_published - - monkeypatch.setattr( - app_generator, - "agent_has_workflow_callable_active_snapshot", - is_publish_visible, - ) - - def test_success_chains_to_resolve_by_id(self, sqlite_session: Session): - bound_agent = _agent() - snapshot = _snapshot() - _persist(sqlite_session, bound_agent, snapshot) - app_model = _app() - - agent, config_id, config_version_kind, soul = AgentAppGenerator()._resolve_agent( - app_model, - invoke_from=InvokeFrom.WEB_APP, - draft_type=None, - user=_account(), - session=sqlite_session, - ) - - assert agent is bound_agent - assert config_id == snapshot.id - assert config_version_kind == "snapshot" - assert soul.model is not None - - def test_unpublished_draft_still_resolves_active_snapshot( - self, monkeypatch: pytest.MonkeyPatch, sqlite_session: Session - ): - bound_agent = _agent(published=False) - snapshot = _snapshot() - _persist(sqlite_session, bound_agent, snapshot) - app_model = _app() - monkeypatch.setattr(app_generator, "agent_has_workflow_callable_active_snapshot", lambda **_kwargs: True) - - agent, config_id, config_version_kind, soul = AgentAppGenerator()._resolve_agent( - app_model, - invoke_from=InvokeFrom.WEB_APP, - draft_type=None, - user=_account(), - session=sqlite_session, - ) - - assert agent is bound_agent - assert config_id == snapshot.id - assert config_version_kind == "snapshot" - assert soul.prompt.system_prompt == "You are Iris." - - def test_existing_conversation_resolves_binding_snapshot_instead_of_latest_active_snapshot( - self, monkeypatch: pytest.MonkeyPatch, sqlite_session: Session - ): - bound_agent = _agent(active_snapshot_id="snap-2") - pinned_snapshot = _snapshot() - _persist(sqlite_session, bound_agent, pinned_snapshot) - conversation = _conversation(binding_id="binding-1") - binding = _binding() - get_active_binding = MagicMock(return_value=binding) - validate_generation = MagicMock() - monkeypatch.setattr(AgentWorkspaceService, "get_active_binding", get_active_binding) - monkeypatch.setattr(AgentWorkspaceService, "validate_binding_generation", validate_generation) - app_model = _app() - - _, config_id, config_version_kind, soul = AgentAppGenerator()._resolve_agent( - app_model, - invoke_from=InvokeFrom.WEB_APP, - draft_type=None, - user=_account(), - session=sqlite_session, - conversation=conversation, - ) - - assert config_id == "snap-1" - assert config_version_kind == "snapshot" - assert soul.prompt.system_prompt == "You are Iris." - assert get_active_binding.call_args.kwargs["binding_id"] == "binding-1" - validate_generation.assert_called_once_with( - binding, - base_home_snapshot_id="home-1", - agent_config_version_id="snap-1", - agent_config_version_kind=AgentConfigVersionKind.SNAPSHOT, - ) - - def test_existing_conversation_rejects_unavailable_binding( - self, monkeypatch: pytest.MonkeyPatch, sqlite_session: Session - ): - bound_agent = _agent(active_snapshot_id="snap-active") - _persist(sqlite_session, bound_agent) - conversation = _conversation(binding_id="binding-missing") - monkeypatch.setattr(AgentWorkspaceService, "get_active_binding", MagicMock(return_value=None)) - - with pytest.raises(AgentAppGeneratorError, match="Conversation participant Binding is unavailable"): - AgentAppGenerator()._resolve_agent( - _app(), - invoke_from=InvokeFrom.WEB_APP, - draft_type=None, - user=_account(), - session=sqlite_session, - conversation=conversation, - ) - - @pytest.mark.parametrize( - ("binding_home_id", "binding_version_kind", "snapshot_home_id"), - [ - ("home-binding", AgentConfigVersionKind.SNAPSHOT, "home-other"), - ("home-pinned", AgentConfigVersionKind.DRAFT, "home-pinned"), - ], - ) - def test_existing_conversation_generation_mismatch_does_not_fallback_to_active_snapshot( - self, - monkeypatch: pytest.MonkeyPatch, - sqlite_session: Session, - binding_home_id: str, - binding_version_kind: AgentConfigVersionKind, - snapshot_home_id: str, - ): - bound_agent = _agent(active_snapshot_id="snap-active") - conversation = _conversation(binding_id="binding-1") - binding = _binding( - home_snapshot_id=binding_home_id, - version_id="snap-pinned", - version_kind=binding_version_kind, - ) - pinned_snapshot = _snapshot(snapshot_id="snap-pinned", home_snapshot_id=snapshot_home_id) - _persist(sqlite_session, bound_agent, pinned_snapshot) - monkeypatch.setattr(AgentWorkspaceService, "get_active_binding", MagicMock(return_value=binding)) - - with pytest.raises(AgentWorkspaceBindingGenerationMismatchError): - AgentAppGenerator()._resolve_agent( - _app(), - invoke_from=InvokeFrom.WEB_APP, - draft_type=None, - user=_account(), - session=sqlite_session, - conversation=conversation, - ) - - def test_unpublished_imported_agent_is_not_available_to_public_runtime(self, sqlite_session: Session): - bound_agent = _agent(source=AgentSource.IMPORTED, published=False) - _persist(sqlite_session, bound_agent) - app_model = _app() - - with pytest.raises(AgentAppNotPublishedError, match="not been published"): - AgentAppGenerator()._resolve_agent( - app_model, - invoke_from=InvokeFrom.WEB_APP, - draft_type=None, - user=_account(), - session=sqlite_session, - ) - - def test_never_published_agent_app_is_not_available_to_public_runtime(self, sqlite_session: Session): - bound_agent = _agent(published=False) - _persist(sqlite_session, bound_agent) - - with pytest.raises(AgentAppNotPublishedError, match="not been published"): - AgentAppGenerator()._resolve_agent( - _app(), - invoke_from=InvokeFrom.WEB_APP, - draft_type=None, - user=_account(), - session=sqlite_session, - ) - - def test_unpublished_imported_agent_remains_available_to_debugger(self, sqlite_session: Session): - bound_agent = _agent(source=AgentSource.IMPORTED, published=False) - draft = AgentConfigDraft( - id="draft-1", - tenant_id="t1", - agent_id="agent-1", - draft_type=AgentConfigDraftType.DRAFT, - account_id=None, - draft_owner_key="", - base_snapshot_id="snap-1", - home_snapshot_id="home-1", - config_snapshot=AgentSoulConfig.model_validate(_SOUL_DICT), - ) - _persist(sqlite_session, bound_agent, draft) - app_model = _app() - - agent, config_id, config_version_kind, soul = AgentAppGenerator()._resolve_agent( - app_model, - invoke_from=InvokeFrom.DEBUGGER, - draft_type=None, - user=_account(), - session=sqlite_session, - ) - - assert agent is bound_agent - assert config_id == draft.id - assert config_version_kind == "draft" - assert soul.prompt.system_prompt == "You are Iris." - - def test_agent_without_active_snapshot_raises_before_model_resolution(self, sqlite_session: Session): - bound_agent = _agent(active_snapshot_id=None, published=False) - _persist(sqlite_session, bound_agent) - app_model = _app() - - with pytest.raises(AgentAppNotPublishedError, match="not been published"): - AgentAppGenerator()._resolve_agent( - app_model, - invoke_from=InvokeFrom.WEB_APP, - draft_type=None, - user=_account(), - session=sqlite_session, - ) - - def test_unbound_app_raises(self, sqlite_session: Session): - app_model = _app() - with pytest.raises(AgentAppGeneratorError, match="has no bound Agent"): - AgentAppGenerator()._resolve_agent( - app_model, - invoke_from=InvokeFrom.WEB_APP, - draft_type=None, - user=_account(), - session=sqlite_session, - ) diff --git a/api/tests/unit_tests/core/app/apps/agent_app/test_runtime_request_builder.py b/api/tests/unit_tests/core/app/apps/agent_app/test_runtime_request_builder.py index e2572315149..d0116bb44a8 100644 --- a/api/tests/unit_tests/core/app/apps/agent_app/test_runtime_request_builder.py +++ b/api/tests/unit_tests/core/app/apps/agent_app/test_runtime_request_builder.py @@ -42,8 +42,8 @@ from graphon.model_runtime.entities.message_entities import ImagePromptMessageCo from graphon.model_runtime.entities.model_entities import ModelFeature from models.agent_config_entities import AgentSoulConfig from services.tracing.runtime_bridge import RuntimeApplicationTracingHooks -from tests.unit_tests.config_override import apply_config_overrides -from tests.unit_tests.knowledge_fs_core_fakes import RecordingModelFactory, llm_schema +from tests.helpers.unit.config_override import apply_config_overrides +from tests.helpers.unit.knowledge_fs_core_fakes import RecordingModelFactory, llm_schema @dataclass(frozen=True) diff --git a/api/tests/unit_tests/core/app/apps/agent_chat/test_agent_chat_app_config_manager.py b/api/tests/unit_tests/core/app/apps/agent_chat/test_app_config_manager.py similarity index 100% rename from api/tests/unit_tests/core/app/apps/agent_chat/test_agent_chat_app_config_manager.py rename to api/tests/unit_tests/core/app/apps/agent_chat/test_app_config_manager.py diff --git a/api/tests/unit_tests/core/app/apps/agent_chat/test_agent_chat_app_generator.py b/api/tests/unit_tests/core/app/apps/agent_chat/test_app_generator.py similarity index 99% rename from api/tests/unit_tests/core/app/apps/agent_chat/test_agent_chat_app_generator.py rename to api/tests/unit_tests/core/app/apps/agent_chat/test_app_generator.py index 1da9f8f77c4..e8d66310c26 100644 --- a/api/tests/unit_tests/core/app/apps/agent_chat/test_agent_chat_app_generator.py +++ b/api/tests/unit_tests/core/app/apps/agent_chat/test_app_generator.py @@ -24,7 +24,7 @@ from graphon.model_runtime.errors.invoke import InvokeAuthorizationError from models import Account from models.model import App, AppMode, AppModelConfig, Conversation, Message from services.workflow_run_agg import WorkflowRunAgg -from tests.unit_tests.config_override import apply_config_overrides +from tests.helpers.unit.config_override import apply_config_overrides @dataclass(frozen=True) diff --git a/api/tests/unit_tests/core/app/apps/agent_chat/test_agent_chat_app_runner.py b/api/tests/unit_tests/core/app/apps/agent_chat/test_app_runner.py similarity index 100% rename from api/tests/unit_tests/core/app/apps/agent_chat/test_agent_chat_app_runner.py rename to api/tests/unit_tests/core/app/apps/agent_chat/test_app_runner.py diff --git a/api/tests/unit_tests/core/app/apps/agent_chat/test_agent_chat_generate_response_converter.py b/api/tests/unit_tests/core/app/apps/agent_chat/test_generate_response_converter.py similarity index 100% rename from api/tests/unit_tests/core/app/apps/agent_chat/test_agent_chat_generate_response_converter.py rename to api/tests/unit_tests/core/app/apps/agent_chat/test_generate_response_converter.py diff --git a/api/tests/unit_tests/core/app/apps/chat/test_app_config_manager.py b/api/tests/unit_tests/core/app/apps/chat/test_app_config_manager.py index 2acd2137278..30fb3b34197 100644 --- a/api/tests/unit_tests/core/app/apps/chat/test_app_config_manager.py +++ b/api/tests/unit_tests/core/app/apps/chat/test_app_config_manager.py @@ -7,7 +7,7 @@ from sqlalchemy.orm import Session from core.app.app_config.entities import EasyUIBasedAppModelConfigFrom, ModelConfigEntity, PromptTemplateEntity from core.app.apps.chat.app_config_manager import ChatAppConfigManager from models.model import App, AppMode, AppModelConfig -from tests.unit_tests.model_factories import make_app +from tests.helpers.unit.model_factories import make_app def _app() -> App: diff --git a/api/tests/unit_tests/core/app/apps/chat/test_app_generator_and_runner.py b/api/tests/unit_tests/core/app/apps/chat/test_app_generator.py similarity index 100% rename from api/tests/unit_tests/core/app/apps/chat/test_app_generator_and_runner.py rename to api/tests/unit_tests/core/app/apps/chat/test_app_generator.py diff --git a/api/tests/unit_tests/core/app/apps/chat/test_base_app_runner_multimodal.py b/api/tests/unit_tests/core/app/apps/chat/test_base_app_runner_multimodal.py deleted file mode 100644 index c41c33487eb..00000000000 --- a/api/tests/unit_tests/core/app/apps/chat/test_base_app_runner_multimodal.py +++ /dev/null @@ -1,337 +0,0 @@ -"""Test multimodal image output handling in BaseAppRunner.""" - -from unittest.mock import MagicMock, patch -from uuid import uuid4 - -import pytest -from sqlalchemy import func, select -from sqlalchemy.orm import Session - -from core.app.apps.base_app_runner import AppRunner -from core.app.entities.app_invoke_entities import InvokeFrom -from graphon.file import FileTransferMethod, FileType -from graphon.model_runtime.entities.message_entities import ImagePromptMessageContent -from models.enums import CreatorUserRole -from models.model import MessageFile -from models.tools import ToolFile - - -class TestBaseAppRunnerMultimodal: - """Test that BaseAppRunner correctly handles multimodal image content.""" - - @pytest.fixture - def mock_user_id(self): - """Mock user ID.""" - return str(uuid4()) - - @pytest.fixture - def mock_tenant_id(self): - """Mock tenant ID.""" - return str(uuid4()) - - @pytest.fixture - def mock_message_id(self): - """Mock message ID.""" - return str(uuid4()) - - @pytest.fixture - def mock_queue_manager(self): - """Create a mock queue manager.""" - manager = MagicMock() - manager.invoke_from = InvokeFrom.SERVICE_API - return manager - - @pytest.fixture - def tool_file(self, mock_user_id: str, mock_tenant_id: str) -> ToolFile: - """Create a real transient tool-file model returned by the external file manager.""" - return ToolFile( - user_id=mock_user_id, - tenant_id=mock_tenant_id, - conversation_id=None, - file_key="generated/image.png", - mimetype="image/png", - original_url="http://example.com/image.png", - name="image.png", - size=68, - ) - - def test_handle_multimodal_image_content_with_url( - self, - mock_user_id, - mock_tenant_id, - mock_message_id, - mock_queue_manager, - tool_file, - sqlite_session: Session, - ): - """Test handling image from URL.""" - # Arrange - image_url = "http://example.com/image.png" - content = ImagePromptMessageContent( - url=image_url, - format="png", - mime_type="image/png", - ) - - with patch("core.app.apps.base_app_runner.ToolFileManager", autospec=True) as mock_mgr_class: - # Setup mock tool file manager - mock_mgr = MagicMock() - mock_mgr.create_file_by_url.return_value = tool_file - mock_mgr_class.return_value = mock_mgr - - message_file_id = AppRunner()._handle_multimodal_image_content( - session=sqlite_session, - content=content, - message_id=mock_message_id, - user_id=mock_user_id, - tenant_id=mock_tenant_id, - queue_manager=mock_queue_manager, - ) - - mock_mgr.create_file_by_url.assert_called_once_with( - user_id=mock_user_id, - tenant_id=mock_tenant_id, - file_url=image_url, - conversation_id=None, - ) - message_file = sqlite_session.get(MessageFile, message_file_id) - assert message_file is not None - assert message_file.message_id == mock_message_id - assert message_file.type == FileType.IMAGE - assert message_file.transfer_method == FileTransferMethod.TOOL_FILE - assert message_file.belongs_to == "assistant" - assert message_file.created_by == mock_user_id - assert message_file.upload_file_id == tool_file.id - mock_queue_manager.publish.assert_not_called() - - def test_handle_multimodal_image_content_with_base64( - self, - mock_user_id, - mock_tenant_id, - mock_message_id, - mock_queue_manager, - tool_file, - sqlite_session: Session, - ): - """Test handling image from base64 data.""" - # Arrange - import base64 - - # Create a small test image (1x1 PNG) - test_image_data = base64.b64encode( - b"\x89PNG\r\n\x1a\n\x00\x00\x00\rIHDR\x00\x00\x00\x01\x00\x00\x00\x01\x08\x02\x00\x00\x00\x90wS\xde" - ).decode() - content = ImagePromptMessageContent( - base64_data=test_image_data, - format="png", - mime_type="image/png", - ) - - with patch("core.app.apps.base_app_runner.ToolFileManager", autospec=True) as mock_mgr_class: - # Setup mock tool file manager - mock_mgr = MagicMock() - mock_mgr.create_file_by_raw.return_value = tool_file - mock_mgr_class.return_value = mock_mgr - - message_file_id = AppRunner()._handle_multimodal_image_content( - session=sqlite_session, - content=content, - message_id=mock_message_id, - user_id=mock_user_id, - tenant_id=mock_tenant_id, - queue_manager=mock_queue_manager, - ) - - mock_mgr.create_file_by_raw.assert_called_once() - call_kwargs = mock_mgr.create_file_by_raw.call_args[1] - assert call_kwargs["user_id"] == mock_user_id - assert call_kwargs["tenant_id"] == mock_tenant_id - assert call_kwargs["conversation_id"] is None - assert "file_binary" in call_kwargs - assert call_kwargs["mimetype"] == "image/png" - assert call_kwargs["filename"].startswith("generated_image") - assert call_kwargs["filename"].endswith(".png") - assert sqlite_session.get(MessageFile, message_file_id) is not None - mock_queue_manager.publish.assert_not_called() - - def test_handle_multimodal_image_content_with_base64_data_uri( - self, - mock_user_id, - mock_tenant_id, - mock_message_id, - mock_queue_manager, - tool_file, - sqlite_session: Session, - ): - """Test handling image from base64 data with URI prefix.""" - # Arrange - # Data URI format: data:image/png;base64, - test_image_data = ( - "iVBORw0KGgoAAAANSUhEUgAAAAEAAAABCAYAAAAfFcSJAAAADUlEQVR42mNk+M9QDwADhgGAWjR9awAAAABJRU5ErkJggg==" - ) - content = ImagePromptMessageContent( - base64_data=f"data:image/png;base64,{test_image_data}", - format="png", - mime_type="image/png", - ) - - with patch("core.app.apps.base_app_runner.ToolFileManager", autospec=True) as mock_mgr_class: - # Setup mock tool file manager - mock_mgr = MagicMock() - mock_mgr.create_file_by_raw.return_value = tool_file - mock_mgr_class.return_value = mock_mgr - - message_file_id = AppRunner()._handle_multimodal_image_content( - session=sqlite_session, - content=content, - message_id=mock_message_id, - user_id=mock_user_id, - tenant_id=mock_tenant_id, - queue_manager=mock_queue_manager, - ) - - mock_mgr.create_file_by_raw.assert_called_once() - call_kwargs = mock_mgr.create_file_by_raw.call_args[1] - assert "file_binary" in call_kwargs - assert sqlite_session.get(MessageFile, message_file_id) is not None - - def test_handle_multimodal_image_content_without_url_or_base64( - self, - mock_user_id, - mock_tenant_id, - mock_message_id, - mock_queue_manager, - sqlite_session: Session, - ): - """Test handling image content without URL or base64 data.""" - # Arrange - content = ImagePromptMessageContent( - url="", - base64_data="", - format="png", - mime_type="image/png", - ) - - with patch("core.app.apps.base_app_runner.ToolFileManager", autospec=True) as mock_mgr_class: - result = AppRunner()._handle_multimodal_image_content( - session=sqlite_session, - content=content, - message_id=mock_message_id, - user_id=mock_user_id, - tenant_id=mock_tenant_id, - queue_manager=mock_queue_manager, - ) - - assert result is None - assert sqlite_session.scalar(select(func.count()).select_from(MessageFile)) == 0 - mock_mgr_class.assert_not_called() - mock_queue_manager.publish.assert_not_called() - - def test_handle_multimodal_image_content_with_error( - self, - mock_user_id, - mock_tenant_id, - mock_message_id, - mock_queue_manager, - sqlite_session: Session, - ): - """Test handling image content when an error occurs.""" - # Arrange - image_url = "http://example.com/image.png" - content = ImagePromptMessageContent( - url=image_url, - format="png", - mime_type="image/png", - ) - - with patch("core.app.apps.base_app_runner.ToolFileManager", autospec=True) as mock_mgr_class: - mock_mgr = MagicMock() - mock_mgr.create_file_by_url.side_effect = Exception("Network error") - mock_mgr_class.return_value = mock_mgr - - result = AppRunner()._handle_multimodal_image_content( - session=sqlite_session, - content=content, - message_id=mock_message_id, - user_id=mock_user_id, - tenant_id=mock_tenant_id, - queue_manager=mock_queue_manager, - ) - - assert result is None - assert sqlite_session.scalar(select(func.count()).select_from(MessageFile)) == 0 - mock_queue_manager.publish.assert_not_called() - - def test_handle_multimodal_image_content_debugger_mode( - self, - mock_user_id, - mock_tenant_id, - mock_message_id, - mock_queue_manager, - tool_file, - sqlite_session: Session, - ): - """Test that debugger mode sets correct created_by_role.""" - # Arrange - image_url = "http://example.com/image.png" - content = ImagePromptMessageContent( - url=image_url, - format="png", - mime_type="image/png", - ) - mock_queue_manager.invoke_from = InvokeFrom.DEBUGGER - - with patch("core.app.apps.base_app_runner.ToolFileManager", autospec=True) as mock_mgr_class: - mock_mgr = MagicMock() - mock_mgr.create_file_by_url.return_value = tool_file - mock_mgr_class.return_value = mock_mgr - - message_file_id = AppRunner()._handle_multimodal_image_content( - session=sqlite_session, - content=content, - message_id=mock_message_id, - user_id=mock_user_id, - tenant_id=mock_tenant_id, - queue_manager=mock_queue_manager, - ) - - message_file = sqlite_session.get(MessageFile, message_file_id) - assert message_file is not None - assert message_file.created_by_role == CreatorUserRole.ACCOUNT - - def test_handle_multimodal_image_content_service_api_mode( - self, - mock_user_id, - mock_tenant_id, - mock_message_id, - mock_queue_manager, - tool_file, - sqlite_session: Session, - ): - """Test that service API mode sets correct created_by_role.""" - # Arrange - image_url = "http://example.com/image.png" - content = ImagePromptMessageContent( - url=image_url, - format="png", - mime_type="image/png", - ) - mock_queue_manager.invoke_from = InvokeFrom.SERVICE_API - - with patch("core.app.apps.base_app_runner.ToolFileManager", autospec=True) as mock_mgr_class: - mock_mgr = MagicMock() - mock_mgr.create_file_by_url.return_value = tool_file - mock_mgr_class.return_value = mock_mgr - - message_file_id = AppRunner()._handle_multimodal_image_content( - session=sqlite_session, - content=content, - message_id=mock_message_id, - user_id=mock_user_id, - tenant_id=mock_tenant_id, - queue_manager=mock_queue_manager, - ) - - message_file = sqlite_session.get(MessageFile, message_file_id) - assert message_file is not None - assert message_file.created_by_role == CreatorUserRole.END_USER diff --git a/api/tests/unit_tests/core/app/apps/common/test_workflow_response_converter.py b/api/tests/unit_tests/core/app/apps/common/test_workflow_response_converter.py index dd6cd0e919d..c80352323d3 100644 --- a/api/tests/unit_tests/core/app/apps/common/test_workflow_response_converter.py +++ b/api/tests/unit_tests/core/app/apps/common/test_workflow_response_converter.py @@ -1,258 +1,1821 @@ +import json +import uuid from collections.abc import Mapping, Sequence +from dataclasses import dataclass +from datetime import UTC, datetime, timedelta +from types import SimpleNamespace +from typing import Any +from unittest.mock import MagicMock +import pytest +from sqlalchemy.orm import Session + +from core.app.app_config.entities import WorkflowUIBasedAppConfig +from core.app.apps.common import workflow_response_converter from core.app.apps.common.workflow_response_converter import WorkflowResponseConverter +from core.app.apps.workflow.app_runner import WorkflowAppRunner +from core.app.entities.app_invoke_entities import InvokeFrom, WorkflowAppGenerateEntity +from core.app.entities.queue_entities import ( + NodeExecutionSnapshot, + QueueEvent, + QueueHumanInputFormFilledEvent, + QueueHumanInputFormTimeoutEvent, + QueueIterationStartEvent, + QueueLoopStartEvent, + QueueNodeExceptionEvent, + QueueNodeFailedEvent, + QueueNodeRetryEvent, + QueueNodeStartedEvent, + QueueNodeSucceededEvent, + QueueWorkflowPausedEvent, + QueueWorkflowStartedEvent, +) +from core.app.entities.task_entities import HumanInputRequiredResponse, WorkflowPauseStreamResponse +from core.workflow.nodes.human_input.entities import ( + ParagraphInputConfig, + SelectInputConfig, + StringListSource, + UserActionConfig, +) +from core.workflow.nodes.human_input.enums import ValueSourceType +from core.workflow.nodes.human_input.pause_reason import HumanInputRequired +from core.workflow.system_variables import build_system_variables +from graphon.engine_events import GraphRunPausedEvent +from graphon.entities import WorkflowStartReason +from graphon.entities.pause_reason import HitlRequired +from graphon.enums import BuiltinNodeTypes from graphon.file import FILE_MODEL_IDENTITY, File, FileTransferMethod, FileType -from graphon.variables.segments import ArrayFileSegment, FileSegment +from graphon.runtime import RuntimeState, VariablePool +from graphon.variables.segments import ArrayFileSegment, FileSegment, StringSegment +from libs.datetime_utils import naive_utc_now +from models import Account as _workflow_response_converter_truncation_Account +from models.account import Account as _workflow_pause_events_Account +from models.account import Account as _workflow_response_converter_human_input_Account +from models.account import Account as _workflow_response_converter_resumption_Account +from models.human_input import HumanInputForm, HumanInputFormRecipient, RecipientType +from models.model import AppMode +from models.workflow import Workflow, WorkflowNodeExecutionModel, WorkflowType -class TestWorkflowResponseConverterFetchFilesFromVariableValue: - """Test class for WorkflowResponseConverter._fetch_files_from_variable_value method""" +class TestWorkflowResponseConverterCases: + class TestWorkflowResponseConverterFetchFilesFromVariableValue: + """Test class for WorkflowResponseConverter._fetch_files_from_variable_value method""" - def create_test_file(self, file_id: str = "test_file_1") -> File: - """Create a test File object""" - return File( - file_id=file_id, - file_type=FileType.DOCUMENT, - transfer_method=FileTransferMethod.LOCAL_FILE, - related_id="related_123", - filename=f"{file_id}.txt", - extension=".txt", - mime_type="text/plain", - size=1024, - storage_key="storage_key_123", + def create_test_file(self, file_id: str = "test_file_1") -> File: + """Create a test File object""" + return File( + file_id=file_id, + file_type=FileType.DOCUMENT, + transfer_method=FileTransferMethod.LOCAL_FILE, + related_id="related_123", + filename=f"{file_id}.txt", + extension=".txt", + mime_type="text/plain", + size=1024, + storage_key="storage_key_123", + ) + + def create_file_dict(self, file_id: str = "test_file_dict"): + """Create a file dictionary with correct dify_model_identity""" + return { + "dify_model_identity": FILE_MODEL_IDENTITY, + "id": file_id, + "tenant_id": "test_tenant", + "type": "document", + "transfer_method": "local_file", + "related_id": "related_456", + "filename": f"{file_id}.txt", + "extension": ".txt", + "mime_type": "text/plain", + "size": 2048, + "url": "http://example.com/file.txt", + } + + def test_fetch_files_from_variable_value_with_none(self): + """Test with None input""" + # The method signature expects Union[dict, list, Segment], but implementation handles None + # We'll test the actual behavior by passing an empty dict instead + result = WorkflowResponseConverter._fetch_files_from_variable_value(None) + assert result == [] + + def test_fetch_files_from_variable_value_with_empty_dict(self): + """Test with empty dictionary""" + result = WorkflowResponseConverter._fetch_files_from_variable_value({}) + assert result == [] + + def test_fetch_files_from_variable_value_with_empty_list(self): + """Test with empty list""" + result = WorkflowResponseConverter._fetch_files_from_variable_value([]) + assert result == [] + + def test_fetch_files_from_variable_value_with_file_segment(self): + """Test with valid FileSegment""" + test_file = self.create_test_file("segment_file") + file_segment = FileSegment(value=test_file) + + result = WorkflowResponseConverter._fetch_files_from_variable_value(file_segment) + + assert len(result) == 1 + assert isinstance(result[0], dict) + assert result[0]["id"] == "segment_file" + assert result[0]["dify_model_identity"] == FILE_MODEL_IDENTITY + + def test_fetch_files_from_variable_value_with_array_file_segment_single(self): + """Test with ArrayFileSegment containing single file""" + test_file = self.create_test_file("array_file_1") + array_segment = ArrayFileSegment(value=[test_file]) + + result = WorkflowResponseConverter._fetch_files_from_variable_value(array_segment) + + assert len(result) == 1 + assert isinstance(result[0], dict) + assert result[0]["id"] == "array_file_1" + + def test_fetch_files_from_variable_value_with_array_file_segment_multiple(self): + """Test with ArrayFileSegment containing multiple files""" + test_file_1 = self.create_test_file("array_file_1") + test_file_2 = self.create_test_file("array_file_2") + array_segment = ArrayFileSegment(value=[test_file_1, test_file_2]) + + result = WorkflowResponseConverter._fetch_files_from_variable_value(array_segment) + + assert len(result) == 2 + assert result[0]["id"] == "array_file_1" + assert result[1]["id"] == "array_file_2" + + def test_fetch_files_from_variable_value_with_array_file_segment_empty(self): + """Test with ArrayFileSegment containing empty array""" + array_segment = ArrayFileSegment(value=[]) + + result = WorkflowResponseConverter._fetch_files_from_variable_value(array_segment) + + assert result == [] + + def test_fetch_files_from_variable_value_with_list_of_file_dicts(self): + """Test with list containing file dictionaries""" + file_dict_1 = self.create_file_dict("list_file_1") + file_dict_2 = self.create_file_dict("list_file_2") + test_list = [file_dict_1, file_dict_2] + + result = WorkflowResponseConverter._fetch_files_from_variable_value(test_list) + + assert len(result) == 2 + assert result[0]["id"] == "list_file_1" + assert result[1]["id"] == "list_file_2" + + def test_fetch_files_from_variable_value_with_list_of_file_objects(self): + """Test with list containing File objects""" + file_obj_1 = self.create_test_file("list_obj_1") + file_obj_2 = self.create_test_file("list_obj_2") + test_list = [file_obj_1, file_obj_2] + + result = WorkflowResponseConverter._fetch_files_from_variable_value(test_list) + + assert len(result) == 2 + assert result[0]["id"] == "list_obj_1" + assert result[1]["id"] == "list_obj_2" + + def test_fetch_files_from_variable_value_with_list_mixed_valid_invalid(self): + """Test with list containing mix of valid files and invalid items""" + file_dict = self.create_file_dict("mixed_file") + invalid_dict = {"not_a_file": "value"} + test_list = [file_dict, invalid_dict, "string_item", 123] + + result = WorkflowResponseConverter._fetch_files_from_variable_value(test_list) + + assert len(result) == 1 + assert result[0]["id"] == "mixed_file" + + def test_fetch_files_from_variable_value_with_list_nested_structures(self): + """Test with list containing nested structures""" + file_dict = self.create_file_dict("nested_file") + nested_list = [file_dict, ["inner_list"]] + test_list = [nested_list, {"nested": "dict"}] + + result = WorkflowResponseConverter._fetch_files_from_variable_value(test_list) + + # Should not process nested structures in list items + assert result == [] + + def test_fetch_files_from_variable_value_with_dict_incorrect_identity(self): + """Test with dictionary having incorrect dify_model_identity""" + invalid_dict = {"dify_model_identity": "wrong_identity", "id": "invalid_file", "filename": "test.txt"} + + result = WorkflowResponseConverter._fetch_files_from_variable_value(invalid_dict) + + assert result == [] + + def test_fetch_files_from_variable_value_with_dict_missing_identity(self): + """Test with dictionary missing dify_model_identity""" + invalid_dict = {"id": "no_identity_file", "filename": "test.txt"} + + result = WorkflowResponseConverter._fetch_files_from_variable_value(invalid_dict) + + assert result == [] + + def test_fetch_files_from_variable_value_with_dict_file_object(self): + """Test with dictionary containing File object""" + file_obj = self.create_test_file("dict_obj_file") + test_dict = {"file_key": file_obj} + + result = WorkflowResponseConverter._fetch_files_from_variable_value(test_dict) + + # Should not extract File objects from dict values + assert result == [] + + def test_fetch_files_from_variable_value_with_mixed_data_types(self): + """Test with various mixed data types""" + mixed_data = {"string": "text", "number": 42, "boolean": True, "null": None, "dify_model_identity": "wrong"} + + result = WorkflowResponseConverter._fetch_files_from_variable_value(mixed_data) + + assert result == [] + + def test_fetch_files_from_variable_value_with_invalid_objects(self): + """Test with invalid objects that are not supported types""" + # Test with an invalid dict that doesn't match expected patterns + invalid_dict = {"custom_key": "custom_value"} + + result = WorkflowResponseConverter._fetch_files_from_variable_value(invalid_dict) + + assert result == [] + + def test_fetch_files_from_variable_value_with_string_input(self): + """Test with string input (unsupported type)""" + # Since method expects Union[dict, list, Segment], test with empty list instead + result = WorkflowResponseConverter._fetch_files_from_variable_value([]) + + assert result == [] + + def test_fetch_files_from_variable_value_with_number_input(self): + """Test with number input (unsupported type)""" + # Test with list containing numbers (should be ignored) + result = WorkflowResponseConverter._fetch_files_from_variable_value([42, "string", None]) + + assert result == [] + + def test_fetch_files_from_variable_value_return_type_is_sequence(self): + """Test that return type is Sequence[Mapping[str, Any]]""" + file_dict = self.create_file_dict("type_test_file") + + result = WorkflowResponseConverter._fetch_files_from_variable_value(file_dict) + + assert isinstance(result, Sequence) + assert len(result) == 1 + assert isinstance(result[0], Mapping) + assert all(isinstance(key, str) for key in result[0]) + + def test_fetch_files_from_variable_value_preserves_file_properties(self): + """Test that all file properties are preserved in the result""" + original_file = self.create_test_file("property_test") + file_segment = FileSegment(value=original_file) + + result = WorkflowResponseConverter._fetch_files_from_variable_value(file_segment) + + assert len(result) == 1 + file_dict = result[0] + assert file_dict["id"] == "property_test" + assert "tenant_id" not in file_dict + assert file_dict["type"] == "document" + assert file_dict["transfer_method"] == "local_file" + assert file_dict["filename"] == "property_test.txt" + assert file_dict["extension"] == ".txt" + assert file_dict["mime_type"] == "text/plain" + assert file_dict["size"] == 1024 + + def test_fetch_files_from_variable_value_with_complex_nested_scenario(self): + """Test complex scenario with nested valid and invalid data""" + file_dict = self.create_file_dict("complex_file") + file_obj = self.create_test_file("complex_obj") + + # Complex nested structure + complex_data = [ + file_dict, # Valid file dict + file_obj, # Valid file object + { # Invalid dict + "not_file": "data", + "nested": {"deep": "value"}, + }, + [ # Nested list (should be ignored) + self.create_file_dict("nested_file") + ], + "string", # Invalid string + None, # None value + 42, # Invalid number + ] + + result = WorkflowResponseConverter._fetch_files_from_variable_value(complex_data) + + assert len(result) == 2 + assert result[0]["id"] == "complex_file" + assert result[1]["id"] == "complex_obj" + + +def _workflow_response_converter_human_input_build_converter() -> WorkflowResponseConverter: + system_variables = build_system_variables( + files=[], + user_id="user-1", + app_id="app-1", + workflow_id="wf-1", + workflow_execution_id="run-1", + ) + app_entity = WorkflowAppGenerateEntity( + task_id="task-1", + app_config=WorkflowUIBasedAppConfig( + app_id="app-1", tenant_id="tenant-1", workflow_id="wf-1", app_mode=AppMode.WORKFLOW + ), + invoke_from=InvokeFrom.EXPLORE, + files=[], + inputs={}, + workflow_execution_id="run-1", + call_depth=0, + user_id="user-1", + stream=True, + ) + account = _workflow_response_converter_human_input_Account(name="tester", email="tester@example.com") + account.id = "acc-1" + return WorkflowResponseConverter( + application_generate_entity=app_entity, + user=account, + system_variables=system_variables, + ) + + +class TestWorkflowResponseConverterHumanInputCases: + @staticmethod + def test_human_input_form_filled_stream_response_contains_rendered_content() -> None: + converter = _workflow_response_converter_human_input_build_converter() + converter.workflow_start_to_stream_response( + task_id="task-1", + workflow_run_id="run-1", + workflow_id="wf-1", + reason=WorkflowStartReason.INITIAL, ) - def create_file_dict(self, file_id: str = "test_file_dict"): - """Create a file dictionary with correct dify_model_identity""" - return { - "dify_model_identity": FILE_MODEL_IDENTITY, - "id": file_id, - "tenant_id": "test_tenant", - "type": "document", - "transfer_method": "local_file", - "related_id": "related_456", - "filename": f"{file_id}.txt", - "extension": ".txt", - "mime_type": "text/plain", - "size": 2048, - "url": "http://example.com/file.txt", + queue_event = QueueHumanInputFormFilledEvent( + form_id="form-1", + node_id="node-1", + node_type="human-input", + node_title="Human Input", + rendered_content="# Title\nvalue", + action_id="Approve", + action_text="Approve", + ) + + resp = converter.human_input_form_filled_to_stream_response(event=queue_event, task_id="task-1") + + assert resp.workflow_run_id == "run-1" + assert resp.data.form_id == "form-1" + assert resp.data.node_id == "node-1" + assert resp.data.node_title == "Human Input" + assert resp.data.rendered_content.startswith("# Title") + assert resp.data.action_id == "Approve" + + @staticmethod + def test_human_input_form_filled_stream_response_serializes_submitted_data() -> None: + converter = _workflow_response_converter_human_input_build_converter() + converter.workflow_start_to_stream_response( + task_id="task-1", + workflow_run_id="run-1", + workflow_id="wf-1", + reason=WorkflowStartReason.INITIAL, + ) + + queue_event = QueueHumanInputFormFilledEvent( + form_id="form-1", + node_id="node-1", + node_type="human-input", + node_title="Human Input", + rendered_content="# Title\nvalue", + action_id="Approve", + action_text="Approve", + submitted_data={ + "decision": StringSegment(value="approve"), + "comment": StringSegment(value="looks good"), + }, + ) + + resp = converter.human_input_form_filled_to_stream_response(event=queue_event, task_id="task-1") + + assert resp.data.submitted_data == { + "decision": "approve", + "comment": "looks good", } - def test_fetch_files_from_variable_value_with_none(self): - """Test with None input""" - # The method signature expects Union[dict, list, Segment], but implementation handles None - # We'll test the actual behavior by passing an empty dict instead - result = WorkflowResponseConverter._fetch_files_from_variable_value(None) - assert result == [] + @staticmethod + def test_human_input_form_timeout_stream_response_contains_timeout_metadata() -> None: + converter = _workflow_response_converter_human_input_build_converter() + converter.workflow_start_to_stream_response( + task_id="task-1", + workflow_run_id="run-1", + workflow_id="wf-1", + reason=WorkflowStartReason.INITIAL, + ) - def test_fetch_files_from_variable_value_with_empty_dict(self): - """Test with empty dictionary""" - result = WorkflowResponseConverter._fetch_files_from_variable_value({}) - assert result == [] + queue_event = QueueHumanInputFormTimeoutEvent( + form_id="form-1", + node_id="node-1", + node_type="human-input", + node_title="Human Input", + expiration_time=datetime(2025, 1, 1, tzinfo=UTC), + ) - def test_fetch_files_from_variable_value_with_empty_list(self): - """Test with empty list""" - result = WorkflowResponseConverter._fetch_files_from_variable_value([]) - assert result == [] + resp = converter.human_input_form_timeout_to_stream_response(event=queue_event, task_id="task-1") - def test_fetch_files_from_variable_value_with_file_segment(self): - """Test with valid FileSegment""" - test_file = self.create_test_file("segment_file") - file_segment = FileSegment(value=test_file) + assert resp.workflow_run_id == "run-1" + assert resp.data.form_id == "form-1" + assert resp.data.node_id == "node-1" + assert resp.data.node_title == "Human Input" + assert resp.data.expiration_time == 1735689600 - result = WorkflowResponseConverter._fetch_files_from_variable_value(file_segment) - assert len(result) == 1 - assert isinstance(result[0], dict) - assert result[0]["id"] == "segment_file" - assert result[0]["dify_model_identity"] == FILE_MODEL_IDENTITY +def _workflow_response_converter_resumption_build_converter() -> WorkflowResponseConverter: + """Construct a minimal WorkflowResponseConverter for testing.""" + system_variables = build_system_variables( + files=[], + user_id="user-1", + app_id="app-1", + workflow_id="wf-1", + workflow_execution_id="run-1", + ) + app_entity = WorkflowAppGenerateEntity( + task_id="task-1", + app_config=WorkflowUIBasedAppConfig( + app_id="app-1", tenant_id="tenant-1", workflow_id="wf-1", app_mode=AppMode.WORKFLOW + ), + invoke_from=InvokeFrom.EXPLORE, + files=[], + inputs={}, + workflow_execution_id="run-1", + call_depth=0, + user_id="user-1", + stream=True, + ) + account = _workflow_response_converter_resumption_Account(name="tester", email="tester@example.com") + account.id = "acc-1" + return WorkflowResponseConverter( + application_generate_entity=app_entity, + user=account, + system_variables=system_variables, + ) - def test_fetch_files_from_variable_value_with_array_file_segment_single(self): - """Test with ArrayFileSegment containing single file""" - test_file = self.create_test_file("array_file_1") - array_segment = ArrayFileSegment(value=[test_file]) - result = WorkflowResponseConverter._fetch_files_from_variable_value(array_segment) +class TestWorkflowResponseConverterResumptionCases: + @staticmethod + def test_workflow_start_stream_response_carries_resumption_reason() -> None: + converter = _workflow_response_converter_resumption_build_converter() + resp = converter.workflow_start_to_stream_response( + task_id="task-1", + workflow_run_id="run-1", + workflow_id="wf-1", + reason=WorkflowStartReason.RESUMPTION, + ) + assert resp.data.reason is WorkflowStartReason.RESUMPTION - assert len(result) == 1 - assert isinstance(result[0], dict) - assert result[0]["id"] == "array_file_1" + @staticmethod + def test_workflow_start_stream_response_carries_initial_reason() -> None: + converter = _workflow_response_converter_resumption_build_converter() + resp = converter.workflow_start_to_stream_response( + task_id="task-1", + workflow_run_id="run-1", + workflow_id="wf-1", + reason=WorkflowStartReason.INITIAL, + ) + assert resp.data.reason is WorkflowStartReason.INITIAL - def test_fetch_files_from_variable_value_with_array_file_segment_multiple(self): - """Test with ArrayFileSegment containing multiple files""" - test_file_1 = self.create_test_file("array_file_1") - test_file_2 = self.create_test_file("array_file_2") - array_segment = ArrayFileSegment(value=[test_file_1, test_file_2]) + @staticmethod + def test_node_start_preserves_the_allocated_index_when_events_arrive_out_of_order(): + converter = _workflow_response_converter_resumption_build_converter() + converter.workflow_start_to_stream_response( + task_id="task-1", workflow_run_id="run-1", workflow_id="wf-1", reason=WorkflowStartReason.INITIAL + ) + for index in (7, 3): + response = converter.workflow_node_start_to_stream_response( + task_id="task-1", + event=QueueNodeStartedEvent( + node_execution_id=f"execution-{index}", + node_id=f"node-{index}", + node_title="Start", + node_type=BuiltinNodeTypes.START, + node_run_index=index, + start_at=datetime(2026, 9, 7), + provider_type="", + provider_id="", + ), + ) + assert response is not None + assert response.data.index == index - result = WorkflowResponseConverter._fetch_files_from_variable_value(array_segment) + @staticmethod + @pytest.mark.parametrize("provider_type", ["workflow", "builtin"]) + def test_workflow_tool_detail_hint_in_live_and_saved_traces(monkeypatch, provider_type): + monkeypatch.setattr(workflow_response_converter.ToolManager, "get_tool_icon", lambda **_kwargs: "icon") + converter = _workflow_response_converter_resumption_build_converter() + converter.workflow_start_to_stream_response( + task_id="task-1", workflow_run_id="run-1", workflow_id="wf-1", reason=WorkflowStartReason.INITIAL + ) + response = converter.workflow_node_start_to_stream_response( + task_id="task-1", + event=QueueNodeStartedEvent( + node_execution_id="tool-execution", + node_id="tool", + node_title="Tool", + node_type=BuiltinNodeTypes.TOOL, + start_at=datetime(2026, 9, 7), + provider_type=provider_type, + provider_id="provider-1", + ), + ) + saved = WorkflowNodeExecutionModel( + node_type=BuiltinNodeTypes.TOOL, + execution_metadata=json.dumps({"tool_info": {"provider_type": provider_type, "provider_id": "provider-1"}}), + ) + assert response is not None + assert response.data.extras == saved.extras + assert response.data.extras.get("workflow_tool", False) == (provider_type == "workflow") - assert len(result) == 2 - assert result[0]["id"] == "array_file_1" - assert result[1]["id"] == "array_file_2" + @staticmethod + @pytest.mark.parametrize("next_node_already_persisted", [False, True]) + def test_resumed_node_metadata_and_sequence(next_node_already_persisted): + started_at = datetime(2026, 9, 7) + converter = _workflow_response_converter_resumption_build_converter() + converter.workflow_start_to_stream_response( + task_id="task-1", workflow_run_id="run-1", workflow_id="wf-1", reason=WorkflowStartReason.INITIAL + ) - def test_fetch_files_from_variable_value_with_array_file_segment_empty(self): - """Test with ArrayFileSegment containing empty array""" - array_segment = ArrayFileSegment(value=[]) + def node_start(execution_id, title, node_type, index): + return QueueNodeStartedEvent( + node_execution_id=execution_id, + node_id=execution_id, + node_run_index=index, + node_title=title, + node_type=node_type, + start_at=started_at, + provider_type="workflow", + provider_id="tool-1", + ) - result = WorkflowResponseConverter._fetch_files_from_variable_value(array_segment) + request = node_start("request", "Request", BuiltinNodeTypes.START, 1) + tool = node_start("tool", "Approval tool", BuiltinNodeTypes.TOOL, 2) + for index, event in enumerate((request, tool), start=1): + response = converter.workflow_node_start_to_stream_response(event=event, task_id="task-1") + assert response is not None + assert response.data.index == index - assert result == [] - - def test_fetch_files_from_variable_value_with_list_of_file_dicts(self): - """Test with list containing file dictionaries""" - file_dict_1 = self.create_file_dict("list_file_1") - file_dict_2 = self.create_file_dict("list_file_2") - test_list = [file_dict_1, file_dict_2] - - result = WorkflowResponseConverter._fetch_files_from_variable_value(test_list) - - assert len(result) == 2 - assert result[0]["id"] == "list_file_1" - assert result[1]["id"] == "list_file_2" - - def test_fetch_files_from_variable_value_with_list_of_file_objects(self): - """Test with list containing File objects""" - file_obj_1 = self.create_test_file("list_obj_1") - file_obj_2 = self.create_test_file("list_obj_2") - test_list = [file_obj_1, file_obj_2] - - result = WorkflowResponseConverter._fetch_files_from_variable_value(test_list) - - assert len(result) == 2 - assert result[0]["id"] == "list_obj_1" - assert result[1]["id"] == "list_obj_2" - - def test_fetch_files_from_variable_value_with_list_mixed_valid_invalid(self): - """Test with list containing mix of valid files and invalid items""" - file_dict = self.create_file_dict("mixed_file") - invalid_dict = {"not_a_file": "value"} - test_list = [file_dict, invalid_dict, "string_item", 123] - - result = WorkflowResponseConverter._fetch_files_from_variable_value(test_list) - - assert len(result) == 1 - assert result[0]["id"] == "mixed_file" - - def test_fetch_files_from_variable_value_with_list_nested_structures(self): - """Test with list containing nested structures""" - file_dict = self.create_file_dict("nested_file") - nested_list = [file_dict, ["inner_list"]] - test_list = [nested_list, {"nested": "dict"}] - - result = WorkflowResponseConverter._fetch_files_from_variable_value(test_list) - - # Should not process nested structures in list items - assert result == [] - - def test_fetch_files_from_variable_value_with_dict_incorrect_identity(self): - """Test with dictionary having incorrect dify_model_identity""" - invalid_dict = {"dify_model_identity": "wrong_identity", "id": "invalid_file", "filename": "test.txt"} - - result = WorkflowResponseConverter._fetch_files_from_variable_value(invalid_dict) - - assert result == [] - - def test_fetch_files_from_variable_value_with_dict_missing_identity(self): - """Test with dictionary missing dify_model_identity""" - invalid_dict = {"id": "no_identity_file", "filename": "test.txt"} - - result = WorkflowResponseConverter._fetch_files_from_variable_value(invalid_dict) - - assert result == [] - - def test_fetch_files_from_variable_value_with_dict_file_object(self): - """Test with dictionary containing File object""" - file_obj = self.create_test_file("dict_obj_file") - test_dict = {"file_key": file_obj} - - result = WorkflowResponseConverter._fetch_files_from_variable_value(test_dict) - - # Should not extract File objects from dict values - assert result == [] - - def test_fetch_files_from_variable_value_with_mixed_data_types(self): - """Test with various mixed data types""" - mixed_data = {"string": "text", "number": 42, "boolean": True, "null": None, "dify_model_identity": "wrong"} - - result = WorkflowResponseConverter._fetch_files_from_variable_value(mixed_data) - - assert result == [] - - def test_fetch_files_from_variable_value_with_invalid_objects(self): - """Test with invalid objects that are not supported types""" - # Test with an invalid dict that doesn't match expected patterns - invalid_dict = {"custom_key": "custom_value"} - - result = WorkflowResponseConverter._fetch_files_from_variable_value(invalid_dict) - - assert result == [] - - def test_fetch_files_from_variable_value_with_string_input(self): - """Test with string input (unsupported type)""" - # Since method expects Union[dict, list, Segment], test with empty list instead - result = WorkflowResponseConverter._fetch_files_from_variable_value([]) - - assert result == [] - - def test_fetch_files_from_variable_value_with_number_input(self): - """Test with number input (unsupported type)""" - # Test with list containing numbers (should be ignored) - result = WorkflowResponseConverter._fetch_files_from_variable_value([42, "string", None]) - - assert result == [] - - def test_fetch_files_from_variable_value_return_type_is_sequence(self): - """Test that return type is Sequence[Mapping[str, Any]]""" - file_dict = self.create_file_dict("type_test_file") - - result = WorkflowResponseConverter._fetch_files_from_variable_value(file_dict) - - assert isinstance(result, Sequence) - assert len(result) == 1 - assert isinstance(result[0], Mapping) - assert all(isinstance(key, str) for key in result[0]) - - def test_fetch_files_from_variable_value_preserves_file_properties(self): - """Test that all file properties are preserved in the result""" - original_file = self.create_test_file("property_test") - file_segment = FileSegment(value=original_file) - - result = WorkflowResponseConverter._fetch_files_from_variable_value(file_segment) - - assert len(result) == 1 - file_dict = result[0] - assert file_dict["id"] == "property_test" - assert "tenant_id" not in file_dict - assert file_dict["type"] == "document" - assert file_dict["transfer_method"] == "local_file" - assert file_dict["filename"] == "property_test.txt" - assert file_dict["extension"] == ".txt" - assert file_dict["mime_type"] == "text/plain" - assert file_dict["size"] == 1024 - - def test_fetch_files_from_variable_value_with_complex_nested_scenario(self): - """Test complex scenario with nested valid and invalid data""" - file_dict = self.create_file_dict("complex_file") - file_obj = self.create_test_file("complex_obj") - - # Complex nested structure - complex_data = [ - file_dict, # Valid file dict - file_obj, # Valid file object - { # Invalid dict - "not_file": "data", - "nested": {"deep": "value"}, - }, - [ # Nested list (should be ignored) - self.create_file_dict("nested_file") - ], - "string", # Invalid string - None, # None value - 42, # Invalid number + snapshots = [ + NodeExecutionSnapshot(execution_id="request", title="Request", index=1, start_at=started_at), + NodeExecutionSnapshot(execution_id="tool", title="Approval tool", index=2, start_at=started_at), ] + if next_node_already_persisted: + snapshots.append(NodeExecutionSnapshot(execution_id="end", title="Result", index=3, start_at=started_at)) + # Resume history survives the actual queued event's serialization boundary. + queued = QueueWorkflowStartedEvent.model_validate_json( + QueueWorkflowStartedEvent(node_execution_snapshots=tuple(snapshots)).model_dump_json() + ) - result = WorkflowResponseConverter._fetch_files_from_variable_value(complex_data) + converter = _workflow_response_converter_resumption_build_converter() + converter.workflow_start_to_stream_response( + task_id="task-1", + workflow_run_id="run-1", + workflow_id="wf-1", + reason=WorkflowStartReason.RESUMPTION, + node_execution_snapshots=queued.node_execution_snapshots, + ) + response = converter.workflow_node_finish_to_stream_response( + event=QueueNodeSucceededEvent( + node_execution_id="tool", + node_id="tool", + node_type=BuiltinNodeTypes.TOOL, + start_at=started_at, + finished_at=started_at + timedelta(seconds=4), + outputs={"decision": "approve"}, + ), + task_id="task-1", + ) + assert response is not None + assert (response.data.title, response.data.index, response.data.elapsed_time) == ("Approval tool", 2, 4) + assert response.data.outputs == {"decision": "approve"} - assert len(result) == 2 - assert result[0]["id"] == "complex_file" - assert result[1]["id"] == "complex_obj" + for index, execution_id in enumerate(("end", "next"), start=3): + response = converter.workflow_node_start_to_stream_response( + event=node_start(execution_id, "Result", BuiltinNodeTypes.END, index), task_id="task-1" + ) + assert response is not None + assert response.data.index == index + + +""" +Unit tests for WorkflowResponseConverter focusing on process_data truncation functionality. +""" + + +@dataclass +class TestCase: + """Test case data for table-driven tests.""" + + name: str + invoke_from: InvokeFrom + expected_truncation_enabled: bool + description: str + + +class TestConverterTruncationCases: + class TestWorkflowResponseConverter: + """Test truncation in WorkflowResponseConverter.""" + + def create_generate_entity(self) -> WorkflowAppGenerateEntity: + """Create a validated workflow invocation for response conversion.""" + return WorkflowAppGenerateEntity( + app_config=WorkflowUIBasedAppConfig( + tenant_id="test-tenant-id", app_id="app", workflow_id="wf-id", app_mode=AppMode.WORKFLOW + ), + task_id="task", + workflow_execution_id="initial-run-id", + user_id="test-user-id", + invoke_from=InvokeFrom.WEB_APP, + inputs={}, + files=[], + stream=True, + ) + + def create_workflow_response_converter(self) -> WorkflowResponseConverter: + """Create a WorkflowResponseConverter for testing.""" + + entity = self.create_generate_entity() + mock_user = _workflow_response_converter_truncation_Account( + name="Test User", + email="test@example.com", + ) + mock_user.id = "test-user-id" + + system_variables = build_system_variables(workflow_id="wf-id", workflow_execution_id="initial-run-id") + return WorkflowResponseConverter( + application_generate_entity=entity, + user=mock_user, + system_variables=system_variables, + ) + + def create_node_started_event(self, *, node_execution_id: str | None = None) -> QueueNodeStartedEvent: + """Create a QueueNodeStartedEvent for testing.""" + return QueueNodeStartedEvent( + node_execution_id=node_execution_id or str(uuid.uuid4()), + node_id="test-node-id", + node_title="Test Node", + node_type=BuiltinNodeTypes.CODE, + start_at=naive_utc_now(), + in_iteration_id=None, + in_loop_id=None, + provider_type="built-in", + provider_id="code", + ) + + def create_node_succeeded_event( + self, + *, + node_execution_id: str, + process_data: Mapping[str, Any] | None = None, + ) -> QueueNodeSucceededEvent: + """Create a QueueNodeSucceededEvent for testing.""" + return QueueNodeSucceededEvent( + node_id="test-node-id", + node_type=BuiltinNodeTypes.CODE, + node_execution_id=node_execution_id, + start_at=naive_utc_now(), + in_iteration_id=None, + in_loop_id=None, + inputs={}, + process_data=process_data or {}, + outputs={}, + execution_metadata={}, + ) + + def create_node_retry_event( + self, + *, + node_execution_id: str, + process_data: Mapping[str, Any] | None = None, + ) -> QueueNodeRetryEvent: + """Create a QueueNodeRetryEvent for testing.""" + return QueueNodeRetryEvent( + inputs={"data": "inputs"}, + outputs={"data": "outputs"}, + process_data=process_data or {}, + error="oops", + retry_index=1, + node_id="test-node-id", + node_type=BuiltinNodeTypes.CODE, + node_title="test code", + provider_type="built-in", + provider_id="code", + node_execution_id=node_execution_id, + start_at=naive_utc_now(), + in_iteration_id=None, + in_loop_id=None, + ) + + def test_workflow_node_finish_response_uses_truncated_process_data(self): + """Test that node finish response uses get_response_process_data().""" + converter = self.create_workflow_response_converter() + + original_data = {"large_field": "x" * 10000, "metadata": "info"} + truncated_data = {"large_field": "[TRUNCATED]", "metadata": "info"} + + converter.workflow_start_to_stream_response( + task_id="bootstrap", + workflow_run_id="run-id", + workflow_id="wf-id", + reason=WorkflowStartReason.INITIAL, + ) + start_event = self.create_node_started_event() + converter.workflow_node_start_to_stream_response( + event=start_event, + task_id="test-task-id", + ) + + event = self.create_node_succeeded_event( + node_execution_id=start_event.node_execution_id, + process_data=original_data, + ) + + def fake_truncate(mapping): + if mapping == dict(original_data): + return truncated_data, True + return mapping, False + + converter._truncator.truncate_variable_mapping = fake_truncate # type: ignore[assignment] + + response = converter.workflow_node_finish_to_stream_response( + event=event, + task_id="test-task-id", + ) + + # Response should use truncated data, not original + assert response is not None + assert response.data.process_data == truncated_data + assert response.data.process_data != original_data + assert response.data.process_data_truncated is True + + def test_workflow_node_finish_response_without_truncation(self): + """Test node finish response when no truncation is applied.""" + converter = self.create_workflow_response_converter() + + original_data = {"small": "data"} + + converter.workflow_start_to_stream_response( + task_id="bootstrap", + workflow_run_id="run-id", + workflow_id="wf-id", + reason=WorkflowStartReason.INITIAL, + ) + start_event = self.create_node_started_event() + converter.workflow_node_start_to_stream_response( + event=start_event, + task_id="test-task-id", + ) + + event = self.create_node_succeeded_event( + node_execution_id=start_event.node_execution_id, + process_data=original_data, + ) + + def fake_truncate(mapping): + return mapping, False + + converter._truncator.truncate_variable_mapping = fake_truncate # type: ignore[assignment] + + response = converter.workflow_node_finish_to_stream_response( + event=event, + task_id="test-task-id", + ) + + # Response should use original data + assert response is not None + assert response.data.process_data == original_data + assert response.data.process_data_truncated is False + + def test_workflow_node_finish_response_with_none_process_data(self): + """Test node finish response when process_data is None.""" + converter = self.create_workflow_response_converter() + + converter.workflow_start_to_stream_response( + task_id="bootstrap", + workflow_run_id="run-id", + workflow_id="wf-id", + reason=WorkflowStartReason.INITIAL, + ) + start_event = self.create_node_started_event() + converter.workflow_node_start_to_stream_response( + event=start_event, + task_id="test-task-id", + ) + + event = self.create_node_succeeded_event( + node_execution_id=start_event.node_execution_id, + process_data=None, + ) + + def fake_truncate(mapping): + return mapping, False + + converter._truncator.truncate_variable_mapping = fake_truncate # type: ignore[assignment] + + response = converter.workflow_node_finish_to_stream_response( + event=event, + task_id="test-task-id", + ) + + # Response should normalize missing process_data to an empty mapping + assert response is not None + assert response.data.process_data == {} + assert response.data.process_data_truncated is False + + def test_workflow_node_finish_response_prefers_event_finished_at( + self, + monkeypatch: pytest.MonkeyPatch, + ) -> None: + """Finished timestamps should come from the event, not delayed queue processing time.""" + converter = self.create_workflow_response_converter() + start_at = datetime(2024, 1, 1, 0, 0, 0, tzinfo=UTC).replace(tzinfo=None) + finished_at = datetime(2024, 1, 1, 0, 0, 2, tzinfo=UTC).replace(tzinfo=None) + delayed_processing_time = datetime(2024, 1, 1, 0, 0, 10, tzinfo=UTC).replace(tzinfo=None) + + monkeypatch.setattr( + "core.app.apps.common.workflow_response_converter.naive_utc_now", + lambda: delayed_processing_time, + ) + converter.workflow_start_to_stream_response( + task_id="bootstrap", + workflow_run_id="run-id", + workflow_id="wf-id", + reason=WorkflowStartReason.INITIAL, + ) + + event = QueueNodeSucceededEvent( + node_id="test-node-id", + node_type=BuiltinNodeTypes.CODE, + node_execution_id="node-exec-1", + start_at=start_at, + finished_at=finished_at, + in_iteration_id=None, + in_loop_id=None, + inputs={}, + process_data={}, + outputs={}, + execution_metadata={}, + ) + + response = converter.workflow_node_finish_to_stream_response( + event=event, + task_id="test-task-id", + ) + + assert response is not None + assert response.data.elapsed_time == 2.0 + assert response.data.finished_at == int(finished_at.timestamp()) + + def test_workflow_node_retry_response_uses_truncated_process_data(self): + """Test that node retry response uses get_response_process_data().""" + converter = self.create_workflow_response_converter() + + original_data = {"large_field": "x" * 10000, "metadata": "info"} + truncated_data = {"large_field": "[TRUNCATED]", "metadata": "info"} + + converter.workflow_start_to_stream_response( + task_id="bootstrap", + workflow_run_id="run-id", + workflow_id="wf-id", + reason=WorkflowStartReason.INITIAL, + ) + start_event = self.create_node_started_event() + converter.workflow_node_start_to_stream_response( + event=start_event, + task_id="test-task-id", + ) + + event = self.create_node_retry_event( + node_execution_id=start_event.node_execution_id, + process_data=original_data, + ) + + def fake_truncate(mapping): + if mapping == dict(original_data): + return truncated_data, True + return mapping, False + + converter._truncator.truncate_variable_mapping = fake_truncate # type: ignore[assignment] + + response = converter.workflow_node_retry_to_stream_response( + event=event, + task_id="test-task-id", + ) + + # Response should use truncated data, not original + assert response is not None + assert response.data.process_data == truncated_data + assert response.data.process_data != original_data + assert response.data.process_data_truncated is True + + def test_workflow_node_retry_response_without_truncation(self): + """Test node retry response when no truncation is applied.""" + converter = self.create_workflow_response_converter() + + original_data = {"small": "data"} + + converter.workflow_start_to_stream_response( + task_id="bootstrap", + workflow_run_id="run-id", + workflow_id="wf-id", + reason=WorkflowStartReason.INITIAL, + ) + start_event = self.create_node_started_event() + converter.workflow_node_start_to_stream_response( + event=start_event, + task_id="test-task-id", + ) + + event = self.create_node_retry_event( + node_execution_id=start_event.node_execution_id, + process_data=original_data, + ) + + def fake_truncate(mapping): + return mapping, False + + converter._truncator.truncate_variable_mapping = fake_truncate # type: ignore[assignment] + + response = converter.workflow_node_retry_to_stream_response( + event=event, + task_id="test-task-id", + ) + + assert response is not None + assert response.data.process_data == original_data + assert response.data.process_data_truncated is False + + def test_iteration_and_loop_nodes_return_none(self): + """Test that iteration and loop nodes return None (no streaming events).""" + converter = self.create_workflow_response_converter() + + iteration_event = QueueNodeSucceededEvent( + node_id="iteration-node", + node_type=BuiltinNodeTypes.ITERATION, + node_execution_id=str(uuid.uuid4()), + start_at=naive_utc_now(), + in_iteration_id=None, + in_loop_id=None, + inputs={}, + process_data={}, + outputs={}, + execution_metadata={}, + ) + + response = converter.workflow_node_finish_to_stream_response( + event=iteration_event, + task_id="test-task-id", + ) + assert response is None + + loop_event = iteration_event.model_copy(update={"node_type": BuiltinNodeTypes.LOOP}) + response = converter.workflow_node_finish_to_stream_response( + event=loop_event, + task_id="test-task-id", + ) + assert response is None + + def test_finish_without_start_raises(self): + """Ensure finish responses require a prior workflow start.""" + converter = self.create_workflow_response_converter() + event = self.create_node_succeeded_event( + node_execution_id=str(uuid.uuid4()), + process_data={}, + ) + + with pytest.raises(ValueError, match="workflow_run_id missing before streaming workflow events"): + converter.workflow_node_finish_to_stream_response( + event=event, + task_id="test-task-id", + ) + + class TestWorkflowResponseConverterServiceApiTruncation: + """Test class for Service API truncation functionality in WorkflowResponseConverter.""" + + def create_test_app_generate_entity(self, invoke_from: InvokeFrom) -> WorkflowAppGenerateEntity: + """Create a test WorkflowAppGenerateEntity with specified invoke_from.""" + # Create a minimal WorkflowUIBasedAppConfig for testing + app_config = WorkflowUIBasedAppConfig( + tenant_id="test_tenant", + app_id="test_app", + app_mode=AppMode.WORKFLOW, + workflow_id="test_workflow_id", + ) + + entity = WorkflowAppGenerateEntity( + task_id="test_task_id", + app_id="test_app_id", + app_config=app_config, + tenant_id="test_tenant", + app_mode=AppMode.WORKFLOW, + invoke_from=invoke_from, + inputs={"test_input": "test_value"}, + user_id="test_user_id", + stream=True, + files=[], + workflow_execution_id="test_workflow_exec_id", + ) + return entity + + def create_test_user(self) -> _workflow_response_converter_truncation_Account: + """Create a test user account.""" + account = _workflow_response_converter_truncation_Account( + name="Test User", + email="test@example.com", + ) + # Manually set the ID for testing purposes + account.id = "test_user_id" + return account + + def create_test_system_variables(self): + """Create test system variables.""" + return build_system_variables() + + def create_test_converter(self, invoke_from: InvokeFrom) -> WorkflowResponseConverter: + """Create WorkflowResponseConverter with specified invoke_from.""" + entity = self.create_test_app_generate_entity(invoke_from) + user = self.create_test_user() + system_variables = self.create_test_system_variables() + + converter = WorkflowResponseConverter( + application_generate_entity=entity, + user=user, + system_variables=system_variables, + ) + # ensure `workflow_run_id` is set. + converter.workflow_start_to_stream_response( + task_id="test-task-id", + workflow_run_id="test-workflow-run-id", + workflow_id="test-workflow-id", + reason=WorkflowStartReason.INITIAL, + ) + return converter + + @pytest.mark.parametrize( + "test_case", + [ + TestCase( + name="service_api_truncation_disabled", + invoke_from=InvokeFrom.SERVICE_API, + expected_truncation_enabled=False, + description="Service API calls should have truncation disabled", + ), + TestCase( + name="web_app_truncation_enabled", + invoke_from=InvokeFrom.WEB_APP, + expected_truncation_enabled=True, + description="Web app calls should have truncation enabled", + ), + TestCase( + name="debugger_truncation_enabled", + invoke_from=InvokeFrom.DEBUGGER, + expected_truncation_enabled=True, + description="Debugger calls should have truncation enabled", + ), + TestCase( + name="explore_truncation_enabled", + invoke_from=InvokeFrom.EXPLORE, + expected_truncation_enabled=True, + description="Explore calls should have truncation enabled", + ), + TestCase( + name="published_pipeline_truncation_enabled", + invoke_from=InvokeFrom.PUBLISHED_PIPELINE, + expected_truncation_enabled=True, + description="Published pipeline calls should have truncation enabled", + ), + ], + ids=lambda x: x.name, + ) + def test_truncator_selection_based_on_invoke_from(self, test_case: TestCase): + """Test that the correct truncator is selected based on invoke_from.""" + converter = self.create_test_converter(test_case.invoke_from) + + # Test truncation behavior instead of checking private attribute + + # Create a test event with large data + large_value = {"key": ["x"] * 2000} # Large data that would be truncated + + event = QueueNodeSucceededEvent( + node_execution_id="test_node_exec_id", + node_id="test_node", + node_type=BuiltinNodeTypes.LLM, + start_at=naive_utc_now(), + inputs=large_value, + process_data=large_value, + outputs=large_value, + error=None, + execution_metadata=None, + in_iteration_id=None, + in_loop_id=None, + ) + + response = converter.workflow_node_finish_to_stream_response( + event=event, + task_id="test_task", + ) + + # Verify response is not None + assert response is not None + + # Verify truncation behavior matches expectations + if test_case.expected_truncation_enabled: + # Truncation should be enabled for non-service-api calls + assert response.data.inputs_truncated + assert response.data.process_data_truncated + assert response.data.outputs_truncated + else: + # SERVICE_API should not truncate + assert not response.data.inputs_truncated + assert not response.data.process_data_truncated + assert not response.data.outputs_truncated + + def test_service_api_truncator_no_op_mapping(self): + """Test that Service API truncator doesn't truncate variable mappings.""" + converter = self.create_test_converter(InvokeFrom.SERVICE_API) + + # Create a test event with large data + large_value: dict[str, Any] = { + "large_string": "x" * 10000, # Large string + "large_list": list(range(2000)), # Large array + "nested_data": {"deep_nested": {"very_deep": {"value": "x" * 5000}}}, + } + + event = QueueNodeSucceededEvent( + node_execution_id="test_node_exec_id", + node_id="test_node", + node_type=BuiltinNodeTypes.LLM, + start_at=naive_utc_now(), + inputs=large_value, + process_data=large_value, + outputs=large_value, + error=None, + execution_metadata=None, + in_iteration_id=None, + in_loop_id=None, + ) + + response = converter.workflow_node_finish_to_stream_response( + event=event, + task_id="test_task", + ) + + # Verify response is not None + data = response.data + assert data.inputs == large_value + assert data.process_data == large_value + assert data.outputs == large_value + # Service API should not truncate + assert data.inputs_truncated is False + assert data.process_data_truncated is False + assert data.outputs_truncated is False + + def test_web_app_truncator_works_normally(self): + """Test that web app truncator still works normally.""" + converter = self.create_test_converter(InvokeFrom.WEB_APP) + + # Create a test event with large data + large_value = { + "large_string": "x" * 10000, # Large string + "large_list": list(range(2000)), # Large array + } + + event = QueueNodeSucceededEvent( + node_execution_id="test_node_exec_id", + node_id="test_node", + node_type=BuiltinNodeTypes.LLM, + start_at=naive_utc_now(), + inputs=large_value, + process_data=large_value, + outputs=large_value, + error=None, + execution_metadata=None, + in_iteration_id=None, + in_loop_id=None, + ) + + response = converter.workflow_node_finish_to_stream_response( + event=event, + task_id="test_task", + ) + + # Verify response is not None + assert response is not None + + # Web app should truncate + data = response.data + assert data.inputs != large_value + assert data.process_data != large_value + assert data.outputs != large_value + # The exact behavior depends on VariableTruncator implementation + # Just verify that truncation flags are present + assert data.inputs_truncated is True + assert data.process_data_truncated is True + assert data.outputs_truncated is True + + @staticmethod + def _create_event_by_type( + type_: QueueEvent, inputs: Mapping[str, Any], process_data: Mapping[str, Any], outputs: Mapping[str, Any] + ) -> QueueNodeSucceededEvent | QueueNodeFailedEvent | QueueNodeExceptionEvent: + if type_ == QueueEvent.NODE_SUCCEEDED: + return QueueNodeSucceededEvent( + node_execution_id="test_node_exec_id", + node_id="test_node", + node_type=BuiltinNodeTypes.LLM, + start_at=naive_utc_now(), + inputs=inputs, + process_data=process_data, + outputs=outputs, + error=None, + execution_metadata=None, + in_iteration_id=None, + in_loop_id=None, + ) + elif type_ == QueueEvent.NODE_FAILED: + return QueueNodeFailedEvent( + node_execution_id="test_node_exec_id", + node_id="test_node", + node_type=BuiltinNodeTypes.LLM, + start_at=naive_utc_now(), + inputs=inputs, + process_data=process_data, + outputs=outputs, + error="oops", + execution_metadata=None, + in_iteration_id=None, + in_loop_id=None, + ) + elif type_ == QueueEvent.NODE_EXCEPTION: + return QueueNodeExceptionEvent( + node_execution_id="test_node_exec_id", + node_id="test_node", + node_type=BuiltinNodeTypes.LLM, + start_at=naive_utc_now(), + inputs=inputs, + process_data=process_data, + outputs=outputs, + error="oops", + execution_metadata=None, + in_iteration_id=None, + in_loop_id=None, + ) + else: + raise Exception("unknown type.") + + @pytest.mark.parametrize( + "event_type", + [ + QueueEvent.NODE_SUCCEEDED, + QueueEvent.NODE_FAILED, + QueueEvent.NODE_EXCEPTION, + ], + ) + def test_service_api_node_finish_event_no_truncation(self, event_type: QueueEvent): + """Test that Service API doesn't truncate node finish events.""" + converter = self.create_test_converter(InvokeFrom.SERVICE_API) + # Create test event with large data + large_inputs = {"input1": "x" * 5000, "input2": list(range(2000))} + large_process_data = {"process1": "y" * 5000, "process2": {"nested": ["z"] * 2000}} + large_outputs = {"output1": "result" * 1000, "output2": list(range(2000))} + + event = ( + TestConverterTruncationCases.TestWorkflowResponseConverterServiceApiTruncation._create_event_by_type( + event_type, large_inputs, large_process_data, large_outputs + ) + ) + + response = converter.workflow_node_finish_to_stream_response( + event=event, + task_id="test_task", + ) + + # Verify response is not None + assert response is not None + + # Verify response contains full data (not truncated) + assert response.data.inputs == large_inputs + assert response.data.process_data == large_process_data + assert response.data.outputs == large_outputs + assert not response.data.inputs_truncated + assert not response.data.process_data_truncated + assert not response.data.outputs_truncated + + def test_service_api_node_retry_event_no_truncation(self): + """Test that Service API doesn't truncate node retry events.""" + converter = self.create_test_converter(InvokeFrom.SERVICE_API) + + # Create test event with large data + large_inputs = {"retry_input": "x" * 5000} + large_process_data = {"retry_process": "y" * 5000} + large_outputs = {"retry_output": "z" * 5000} + + # First, we need to store a snapshot by simulating a start event + start_event = QueueNodeStartedEvent( + node_execution_id="test_node_exec_id", + node_id="test_node", + node_type=BuiltinNodeTypes.LLM, + node_title="Test Node", + node_run_index=1, + start_at=naive_utc_now(), + in_iteration_id=None, + in_loop_id=None, + agent_strategy=None, + provider_type="plugin", + provider_id="test/test_plugin", + ) + converter.workflow_node_start_to_stream_response(event=start_event, task_id="test_task") + + # Now create retry event + event = QueueNodeRetryEvent( + node_execution_id="test_node_exec_id", + node_id="test_node", + node_type=BuiltinNodeTypes.LLM, + node_title="Test Node", + node_run_index=1, + start_at=naive_utc_now(), + inputs=large_inputs, + process_data=large_process_data, + outputs=large_outputs, + error="Retry error", + execution_metadata=None, + in_iteration_id=None, + in_loop_id=None, + retry_index=1, + provider_type="plugin", + provider_id="test/test_plugin", + ) + + response = converter.workflow_node_retry_to_stream_response( + event=event, + task_id="test_task", + ) + + # Verify response is not None + assert response is not None + + # Verify response contains full data (not truncated) + assert response.data.inputs == large_inputs + assert response.data.process_data == large_process_data + assert response.data.outputs == large_outputs + assert not response.data.inputs_truncated + assert not response.data.process_data_truncated + assert not response.data.outputs_truncated + + def test_service_api_iteration_events_no_truncation(self): + """Test that Service API doesn't truncate iteration events.""" + converter = self.create_test_converter(InvokeFrom.SERVICE_API) + + # Test iteration start event + large_value = {"iteration_input": ["x"] * 2000} + + start_event = QueueIterationStartEvent( + node_execution_id="test_iter_exec_id", + node_id="test_iteration", + node_type=BuiltinNodeTypes.ITERATION, + node_title="Test Iteration", + start_at=naive_utc_now(), + inputs=large_value, + metadata={}, + ) + + response = converter.workflow_iteration_start_to_stream_response( + task_id="test_task", + workflow_execution_id="test_workflow_exec_id", + event=start_event, + ) + + assert response is not None + assert response.data.inputs == large_value + assert not response.data.inputs_truncated + + def test_service_api_loop_events_no_truncation(self): + """Test that Service API doesn't truncate loop events.""" + converter = self.create_test_converter(InvokeFrom.SERVICE_API) + + # Test loop start event + large_inputs = {"loop_input": ["x"] * 2000} + + start_event = QueueLoopStartEvent( + node_execution_id="test_loop_exec_id", + node_id="test_loop", + node_type=BuiltinNodeTypes.LOOP, + node_title="Test Loop", + start_at=naive_utc_now(), + inputs=large_inputs, + metadata={}, + ) + + response = converter.workflow_loop_start_to_stream_response( + task_id="test_task", + workflow_execution_id="test_workflow_exec_id", + event=start_event, + ) + + assert response is not None + assert response.data.inputs == large_inputs + assert not response.data.inputs_truncated + + def test_web_app_node_finish_event_truncation_works(self): + """Test that web app still truncates node finish events.""" + converter = self.create_test_converter(InvokeFrom.WEB_APP) + + # Create test event with large data that should be truncated + large_inputs = {"input1": ["x"] * 2000} + large_process_data = {"process1": ["y"] * 2000} + large_outputs = {"output1": ["z"] * 2000} + + event = QueueNodeSucceededEvent( + node_execution_id="test_node_exec_id", + node_id="test_node", + node_type=BuiltinNodeTypes.LLM, + start_at=naive_utc_now(), + inputs=large_inputs, + process_data=large_process_data, + outputs=large_outputs, + error=None, + execution_metadata=None, + in_iteration_id=None, + in_loop_id=None, + ) + + response = converter.workflow_node_finish_to_stream_response( + event=event, + task_id="test_task", + ) + + # Verify response is not None + assert response is not None + + # Verify response contains truncated data + # The exact behavior depends on VariableTruncator implementation + # Just verify truncation flags are set correctly (may or may not be truncated depending on size) + # At minimum, the truncation mechanism should work + assert isinstance(response.data.inputs, dict) + assert response.data.inputs_truncated + assert isinstance(response.data.process_data, dict) + assert response.data.process_data_truncated + assert isinstance(response.data.outputs, dict) + assert response.data.outputs_truncated + + +class _RecordingWorkflowAppRunner(WorkflowAppRunner): + def __init__(self, **kwargs): + super().__init__(**kwargs) + self.published_events = [] + + def _publish_event(self, event): + self.published_events.append(event) + + +def _persist_human_input_form( + session: Session, + *, + recipients: list[tuple[RecipientType, str]] | None = None, +) -> datetime: + expiration_time = datetime(2024, 1, 1, tzinfo=UTC) + form = HumanInputForm( + id="form-1", + tenant_id="tenant-id", + app_id="app-id", + workflow_run_id="run-id", + node_id="node-id", + form_definition='{"display_in_ui": true}', + rendered_content="Rendered", + expiration_time=expiration_time, + ) + recipient_models = [ + HumanInputFormRecipient( + id=f"recipient-{index}", + form_id=form.id, + delivery_id=f"delivery-{index}", + recipient_type=recipient_type, + recipient_payload="{}", + access_token=access_token, + ) + for index, (recipient_type, access_token) in enumerate(recipients or ()) + ] + session.add(form) + session.add_all(recipient_models) + session.commit() + return expiration_time + + +def _build_runner(): + app_entity = SimpleNamespace( + app_config=SimpleNamespace(app_id="app-id"), + inputs={}, + files=[], + invoke_from=InvokeFrom.SERVICE_API, + single_iteration_run=None, + single_loop_run=None, + workflow_execution_id="run-id", + user_id="user-id", + ) + workflow = Workflow.new( + tenant_id="tenant-id", + app_id="app-id", + type=WorkflowType.WORKFLOW, + version=Workflow.VERSION_DRAFT, + graph=json.dumps({}), + features="{}", + created_by="account-id", + environment_variables=[], + conversation_variables=[], + rag_pipeline_variables=[], + ) + workflow.id = "workflow-id" + queue_manager = SimpleNamespace(publish=lambda event, pub_from: None) + return _RecordingWorkflowAppRunner( + application_generate_entity=app_entity, + queue_manager=queue_manager, + variable_loader=MagicMock(), + workflow=workflow, + system_user_id="sys-user", + root_node_id=None, + graph_engine_layers=(), + graph_runtime_state=None, + ) + + +def _workflow_pause_events_build_converter( + *, invoke_from: InvokeFrom = InvokeFrom.SERVICE_API +) -> WorkflowResponseConverter: + application_generate_entity = WorkflowAppGenerateEntity( + task_id="task", + workflow_execution_id="run-id", + user_id="user", + stream=True, + inputs={}, + files=[], + invoke_from=invoke_from, + app_config=WorkflowUIBasedAppConfig( + app_id="app-id", tenant_id="tenant-id", workflow_id="workflow-id", app_mode=AppMode.WORKFLOW + ), + ) + system_variables = build_system_variables( + user_id="user", + app_id="app-id", + workflow_id="workflow-id", + workflow_execution_id="run-id", + ) + user = _workflow_pause_events_Account(name="Tester", email="tester@example.com") + user.id = "account-id" + return WorkflowResponseConverter( + application_generate_entity=application_generate_entity, + user=user, + system_variables=system_variables, + ) + + +def _build_paused_human_input_response( + session: Session, + recipients: list[tuple[RecipientType, str]], +): + """Drive the live OPENAPI pause path with persisted forms and recipients.""" + converter = _workflow_pause_events_build_converter(invoke_from=InvokeFrom.OPENAPI) + converter.workflow_start_to_stream_response( + task_id="task", + workflow_run_id="run-id", + workflow_id="workflow-id", + reason=WorkflowStartReason.INITIAL, + ) + + _persist_human_input_form(session, recipients=recipients) + + reason = HumanInputRequired( + form_id="form-1", + form_content="Rendered", + inputs=[ParagraphInputConfig(output_variable_name="field")], + actions=[UserActionConfig(id="approve", title="Approve")], + node_id="node-id", + node_title="Human Step", + ) + queue_event = QueueWorkflowPausedEvent( + reasons=[reason], + outputs={}, + paused_nodes=["node-id"], + ) + + runtime_state = RuntimeState(workflow_id="test-workflow", variable_pool=VariablePool(), start_at=0.0) + responses = converter.workflow_pause_to_stream_response( + event=queue_event, + task_id="task", + graph_runtime_state=runtime_state, + ) + assert isinstance(responses[0], HumanInputRequiredResponse) + return responses + + +class TestWorkflowPauseEventsCases: + @staticmethod + @pytest.fixture + def sqlite_pause_session(sqlite_session: Session, monkeypatch: pytest.MonkeyPatch) -> Session: + """Bind pause-response queries to the shared SQLite session's database.""" + monkeypatch.setattr(workflow_response_converter, "db", SimpleNamespace(engine=sqlite_session.get_bind())) + return sqlite_session + + @staticmethod + @pytest.mark.parametrize("source_node_id", ["node-human", "nested-human"]) + def test_graph_run_paused_event_uses_projected_nodes_for_highlighting( + monkeypatch: pytest.MonkeyPatch, source_node_id: str + ): + from tests.unit_tests.core.app.apps.test_workflow_app_runner import _workflow_entry + + runner = _build_runner() + graph_reason = HitlRequired( + session_id="form-1", + node_id=source_node_id, + node_title="Human Step", + ) + event = GraphRunPausedEvent(reasons=[graph_reason], outputs={"foo": "bar"}) + workflow_entry = _workflow_entry(runner) + + enriched_reason = HumanInputRequired( + form_id="form-1", + form_content="content", + inputs=[], + actions=[], + node_id="node-human", + node_title="Human Step", + ) + monkeypatch.setattr("core.app.apps.workflow_app_runner.dispatch_human_input_email_task", MagicMock()) + + runner.handle_event(workflow_entry, event, pause_reasons=[enriched_reason]) + + assert len(runner.published_events) == 1 + queue_event = runner.published_events[0] + assert isinstance(queue_event, QueueWorkflowPausedEvent) + assert queue_event.reasons == [enriched_reason] + assert queue_event.outputs == {"foo": "bar"} + assert queue_event.paused_nodes == ["node-human"] + assert graph_reason.node_id == source_node_id + + @staticmethod + @pytest.mark.parametrize( + "sqlite_session", + [(HumanInputForm, HumanInputFormRecipient)], + indirect=True, + ) + def test_queue_workflow_paused_event_to_stream_responses(sqlite_pause_session: Session): + converter = _workflow_pause_events_build_converter() + converter.workflow_start_to_stream_response( + task_id="task", + workflow_run_id="run-id", + workflow_id="workflow-id", + reason=WorkflowStartReason.INITIAL, + ) + + expiration_time = _persist_human_input_form( + sqlite_pause_session, + recipients=[ + (RecipientType.CONSOLE, "console-token"), + (RecipientType.BACKSTAGE, "backstage-token"), + ], + ) + + reason = HumanInputRequired( + form_id="form-1", + form_content="Rendered", + inputs=[ParagraphInputConfig(output_variable_name="field")], + actions=[UserActionConfig(id="approve", title="Approve")], + node_id="node-id", + node_title="Human Step", + ) + queue_event = QueueWorkflowPausedEvent( + reasons=[reason], + outputs={"answer": "value"}, + paused_nodes=["node-id"], + ) + + runtime_state = RuntimeState(workflow_id="test-workflow", variable_pool=VariablePool(), start_at=0.0) + responses = converter.workflow_pause_to_stream_response( + event=queue_event, + task_id="task", + graph_runtime_state=runtime_state, + ) + + assert isinstance(responses[-1], WorkflowPauseStreamResponse) + pause_resp = responses[-1] + assert pause_resp.workflow_run_id == "run-id" + assert pause_resp.data.paused_nodes == ["node-id"] + assert pause_resp.data.outputs == {} + assert pause_resp.data.reasons[0]["form_id"] == "form-1" + + assert isinstance(responses[0], HumanInputRequiredResponse) + hi_resp = responses[0] + assert hi_resp.data.form_id == "form-1" + assert hi_resp.data.node_id == "node-id" + assert hi_resp.data.node_title == "Human Step" + assert hi_resp.data.inputs[0].output_variable_name == "field" + assert hi_resp.data.actions[0].id == "approve" + assert hi_resp.data.display_in_ui is True + assert hi_resp.data.form_token is None + assert hi_resp.data.approval_channels == ["console"] + assert hi_resp.data.expiration_time == int(expiration_time.timestamp()) + + @staticmethod + @pytest.mark.parametrize( + "sqlite_session", + [(HumanInputForm, HumanInputFormRecipient)], + indirect=True, + ) + def test_openapi_pause_without_web_app_recipient_emits_approval_channels(sqlite_pause_session: Session): + responses = _build_paused_human_input_response( + sqlite_pause_session, + recipients=[ + (RecipientType.EMAIL_MEMBER, "email-token"), + (RecipientType.BACKSTAGE, "backstage-token"), + ], + ) + + hi_resp = responses[0] + assert hi_resp.data.form_token is None + assert hi_resp.data.approval_channels == ["console", "email"] + + pause_resp = responses[-1] + assert pause_resp.data.reasons[0]["approval_channels"] == ["console", "email"] + + @staticmethod + @pytest.mark.parametrize( + "sqlite_session", + [(HumanInputForm, HumanInputFormRecipient)], + indirect=True, + ) + def test_openapi_pause_with_web_app_recipient_sets_token_and_channels(sqlite_pause_session: Session): + responses = _build_paused_human_input_response( + sqlite_pause_session, + recipients=[ + (RecipientType.STANDALONE_WEB_APP, "web-app-token"), + (RecipientType.BACKSTAGE, "backstage-token"), + ], + ) + + hi_resp = responses[0] + assert hi_resp.data.form_token == "web-app-token" + assert hi_resp.data.approval_channels == ["console"] + + pause_resp = responses[-1] + assert pause_resp.data.reasons[0]["approval_channels"] == ["console"] + + @staticmethod + @pytest.mark.parametrize( + "sqlite_session", + [(HumanInputForm, HumanInputFormRecipient)], + indirect=True, + ) + def test_queue_workflow_paused_event_resolves_variable_select_options(sqlite_pause_session: Session): + converter = _workflow_pause_events_build_converter() + converter.workflow_start_to_stream_response( + task_id="task", + workflow_run_id="run-id", + workflow_id="workflow-id", + reason=WorkflowStartReason.INITIAL, + ) + + _persist_human_input_form(sqlite_pause_session) + + reason = HumanInputRequired( + form_id="form-1", + form_content="Rendered", + inputs=[ + SelectInputConfig( + output_variable_name="decision", + option_source=StringListSource( + type=ValueSourceType.VARIABLE, + selector=["start", "options"], + value=[], + ), + ) + ], + actions=[UserActionConfig(id="approve", title="Approve")], + node_id="node-id", + node_title="Human Step", + ) + queue_event = QueueWorkflowPausedEvent( + reasons=[reason], + outputs={}, + paused_nodes=["node-id"], + ) + + runtime_state = RuntimeState(workflow_id="test-workflow", variable_pool=VariablePool(), start_at=0.0) + runtime_state.variable_pool.add(("start", "options"), ["approve", "reject"]) + responses = converter.workflow_pause_to_stream_response( + event=queue_event, + task_id="task", + graph_runtime_state=runtime_state, + ) + + assert isinstance(responses[0], HumanInputRequiredResponse) + hi_resp = responses[0] + assert hi_resp.data.inputs[0].option_source.value == ["approve", "reject"] + + assert isinstance(responses[-1], WorkflowPauseStreamResponse) + pause_resp = responses[-1] + assert pause_resp.data.reasons[0]["inputs"][0]["option_source"]["value"] == ["approve", "reject"] diff --git a/api/tests/unit_tests/core/app/apps/common/test_workflow_response_converter_human_input.py b/api/tests/unit_tests/core/app/apps/common/test_workflow_response_converter_human_input.py deleted file mode 100644 index 6f5ae204f83..00000000000 --- a/api/tests/unit_tests/core/app/apps/common/test_workflow_response_converter_human_input.py +++ /dev/null @@ -1,127 +0,0 @@ -from datetime import UTC, datetime - -from core.app.app_config.entities import WorkflowUIBasedAppConfig -from core.app.apps.common.workflow_response_converter import WorkflowResponseConverter -from core.app.entities.app_invoke_entities import InvokeFrom, WorkflowAppGenerateEntity -from core.app.entities.queue_entities import QueueHumanInputFormFilledEvent, QueueHumanInputFormTimeoutEvent -from core.workflow.system_variables import build_system_variables -from graphon.entities import WorkflowStartReason -from graphon.variables.segments import StringSegment -from models.account import Account -from models.model import AppMode - - -def _build_converter() -> WorkflowResponseConverter: - system_variables = build_system_variables( - files=[], - user_id="user-1", - app_id="app-1", - workflow_id="wf-1", - workflow_execution_id="run-1", - ) - app_entity = WorkflowAppGenerateEntity( - task_id="task-1", - app_config=WorkflowUIBasedAppConfig( - app_id="app-1", tenant_id="tenant-1", workflow_id="wf-1", app_mode=AppMode.WORKFLOW - ), - invoke_from=InvokeFrom.EXPLORE, - files=[], - inputs={}, - workflow_execution_id="run-1", - call_depth=0, - user_id="user-1", - stream=True, - ) - account = Account(name="tester", email="tester@example.com") - account.id = "acc-1" - return WorkflowResponseConverter( - application_generate_entity=app_entity, - user=account, - system_variables=system_variables, - ) - - -def test_human_input_form_filled_stream_response_contains_rendered_content() -> None: - converter = _build_converter() - converter.workflow_start_to_stream_response( - task_id="task-1", - workflow_run_id="run-1", - workflow_id="wf-1", - reason=WorkflowStartReason.INITIAL, - ) - - queue_event = QueueHumanInputFormFilledEvent( - form_id="form-1", - node_id="node-1", - node_type="human-input", - node_title="Human Input", - rendered_content="# Title\nvalue", - action_id="Approve", - action_text="Approve", - ) - - resp = converter.human_input_form_filled_to_stream_response(event=queue_event, task_id="task-1") - - assert resp.workflow_run_id == "run-1" - assert resp.data.form_id == "form-1" - assert resp.data.node_id == "node-1" - assert resp.data.node_title == "Human Input" - assert resp.data.rendered_content.startswith("# Title") - assert resp.data.action_id == "Approve" - - -def test_human_input_form_filled_stream_response_serializes_submitted_data() -> None: - converter = _build_converter() - converter.workflow_start_to_stream_response( - task_id="task-1", - workflow_run_id="run-1", - workflow_id="wf-1", - reason=WorkflowStartReason.INITIAL, - ) - - queue_event = QueueHumanInputFormFilledEvent( - form_id="form-1", - node_id="node-1", - node_type="human-input", - node_title="Human Input", - rendered_content="# Title\nvalue", - action_id="Approve", - action_text="Approve", - submitted_data={ - "decision": StringSegment(value="approve"), - "comment": StringSegment(value="looks good"), - }, - ) - - resp = converter.human_input_form_filled_to_stream_response(event=queue_event, task_id="task-1") - - assert resp.data.submitted_data == { - "decision": "approve", - "comment": "looks good", - } - - -def test_human_input_form_timeout_stream_response_contains_timeout_metadata() -> None: - converter = _build_converter() - converter.workflow_start_to_stream_response( - task_id="task-1", - workflow_run_id="run-1", - workflow_id="wf-1", - reason=WorkflowStartReason.INITIAL, - ) - - queue_event = QueueHumanInputFormTimeoutEvent( - form_id="form-1", - node_id="node-1", - node_type="human-input", - node_title="Human Input", - expiration_time=datetime(2025, 1, 1, tzinfo=UTC), - ) - - resp = converter.human_input_form_timeout_to_stream_response(event=queue_event, task_id="task-1") - - assert resp.workflow_run_id == "run-1" - assert resp.data.form_id == "form-1" - assert resp.data.node_id == "node-1" - assert resp.data.node_title == "Human Input" - assert resp.data.expiration_time == 1735689600 diff --git a/api/tests/unit_tests/core/app/apps/common/test_workflow_response_converter_resumption.py b/api/tests/unit_tests/core/app/apps/common/test_workflow_response_converter_resumption.py deleted file mode 100644 index e56ba99e15d..00000000000 --- a/api/tests/unit_tests/core/app/apps/common/test_workflow_response_converter_resumption.py +++ /dev/null @@ -1,194 +0,0 @@ -import json -from datetime import datetime, timedelta - -import pytest - -from core.app.app_config.entities import WorkflowUIBasedAppConfig -from core.app.apps.common import workflow_response_converter -from core.app.apps.common.workflow_response_converter import WorkflowResponseConverter -from core.app.entities.app_invoke_entities import InvokeFrom, WorkflowAppGenerateEntity -from core.app.entities.queue_entities import ( - NodeExecutionSnapshot, - QueueNodeStartedEvent, - QueueNodeSucceededEvent, - QueueWorkflowStartedEvent, -) -from core.workflow.system_variables import build_system_variables -from graphon.entities import WorkflowStartReason -from graphon.enums import BuiltinNodeTypes -from models.account import Account -from models.model import AppMode -from models.workflow import WorkflowNodeExecutionModel - - -def _build_converter() -> WorkflowResponseConverter: - """Construct a minimal WorkflowResponseConverter for testing.""" - system_variables = build_system_variables( - files=[], - user_id="user-1", - app_id="app-1", - workflow_id="wf-1", - workflow_execution_id="run-1", - ) - app_entity = WorkflowAppGenerateEntity( - task_id="task-1", - app_config=WorkflowUIBasedAppConfig( - app_id="app-1", tenant_id="tenant-1", workflow_id="wf-1", app_mode=AppMode.WORKFLOW - ), - invoke_from=InvokeFrom.EXPLORE, - files=[], - inputs={}, - workflow_execution_id="run-1", - call_depth=0, - user_id="user-1", - stream=True, - ) - account = Account(name="tester", email="tester@example.com") - account.id = "acc-1" - return WorkflowResponseConverter( - application_generate_entity=app_entity, - user=account, - system_variables=system_variables, - ) - - -def test_workflow_start_stream_response_carries_resumption_reason() -> None: - converter = _build_converter() - resp = converter.workflow_start_to_stream_response( - task_id="task-1", - workflow_run_id="run-1", - workflow_id="wf-1", - reason=WorkflowStartReason.RESUMPTION, - ) - assert resp.data.reason is WorkflowStartReason.RESUMPTION - - -def test_workflow_start_stream_response_carries_initial_reason() -> None: - converter = _build_converter() - resp = converter.workflow_start_to_stream_response( - task_id="task-1", - workflow_run_id="run-1", - workflow_id="wf-1", - reason=WorkflowStartReason.INITIAL, - ) - assert resp.data.reason is WorkflowStartReason.INITIAL - - -def test_node_start_preserves_the_allocated_index_when_events_arrive_out_of_order(): - converter = _build_converter() - converter.workflow_start_to_stream_response( - task_id="task-1", workflow_run_id="run-1", workflow_id="wf-1", reason=WorkflowStartReason.INITIAL - ) - for index in (7, 3): - response = converter.workflow_node_start_to_stream_response( - task_id="task-1", - event=QueueNodeStartedEvent( - node_execution_id=f"execution-{index}", - node_id=f"node-{index}", - node_title="Start", - node_type=BuiltinNodeTypes.START, - node_run_index=index, - start_at=datetime(2026, 9, 7), - provider_type="", - provider_id="", - ), - ) - assert response is not None - assert response.data.index == index - - -@pytest.mark.parametrize("provider_type", ["workflow", "builtin"]) -def test_workflow_tool_detail_hint_in_live_and_saved_traces(monkeypatch, provider_type): - monkeypatch.setattr(workflow_response_converter.ToolManager, "get_tool_icon", lambda **_kwargs: "icon") - converter = _build_converter() - converter.workflow_start_to_stream_response( - task_id="task-1", workflow_run_id="run-1", workflow_id="wf-1", reason=WorkflowStartReason.INITIAL - ) - response = converter.workflow_node_start_to_stream_response( - task_id="task-1", - event=QueueNodeStartedEvent( - node_execution_id="tool-execution", - node_id="tool", - node_title="Tool", - node_type=BuiltinNodeTypes.TOOL, - start_at=datetime(2026, 9, 7), - provider_type=provider_type, - provider_id="provider-1", - ), - ) - saved = WorkflowNodeExecutionModel( - node_type=BuiltinNodeTypes.TOOL, - execution_metadata=json.dumps({"tool_info": {"provider_type": provider_type, "provider_id": "provider-1"}}), - ) - assert response is not None - assert response.data.extras == saved.extras - assert response.data.extras.get("workflow_tool", False) == (provider_type == "workflow") - - -@pytest.mark.parametrize("next_node_already_persisted", [False, True]) -def test_resumed_node_metadata_and_sequence(next_node_already_persisted): - started_at = datetime(2026, 9, 7) - converter = _build_converter() - converter.workflow_start_to_stream_response( - task_id="task-1", workflow_run_id="run-1", workflow_id="wf-1", reason=WorkflowStartReason.INITIAL - ) - - def node_start(execution_id, title, node_type, index): - return QueueNodeStartedEvent( - node_execution_id=execution_id, - node_id=execution_id, - node_run_index=index, - node_title=title, - node_type=node_type, - start_at=started_at, - provider_type="workflow", - provider_id="tool-1", - ) - - request = node_start("request", "Request", BuiltinNodeTypes.START, 1) - tool = node_start("tool", "Approval tool", BuiltinNodeTypes.TOOL, 2) - for index, event in enumerate((request, tool), start=1): - response = converter.workflow_node_start_to_stream_response(event=event, task_id="task-1") - assert response is not None - assert response.data.index == index - - snapshots = [ - NodeExecutionSnapshot(execution_id="request", title="Request", index=1, start_at=started_at), - NodeExecutionSnapshot(execution_id="tool", title="Approval tool", index=2, start_at=started_at), - ] - if next_node_already_persisted: - snapshots.append(NodeExecutionSnapshot(execution_id="end", title="Result", index=3, start_at=started_at)) - # Resume history survives the actual queued event's serialization boundary. - queued = QueueWorkflowStartedEvent.model_validate_json( - QueueWorkflowStartedEvent(node_execution_snapshots=tuple(snapshots)).model_dump_json() - ) - - converter = _build_converter() - converter.workflow_start_to_stream_response( - task_id="task-1", - workflow_run_id="run-1", - workflow_id="wf-1", - reason=WorkflowStartReason.RESUMPTION, - node_execution_snapshots=queued.node_execution_snapshots, - ) - response = converter.workflow_node_finish_to_stream_response( - event=QueueNodeSucceededEvent( - node_execution_id="tool", - node_id="tool", - node_type=BuiltinNodeTypes.TOOL, - start_at=started_at, - finished_at=started_at + timedelta(seconds=4), - outputs={"decision": "approve"}, - ), - task_id="task-1", - ) - assert response is not None - assert (response.data.title, response.data.index, response.data.elapsed_time) == ("Approval tool", 2, 4) - assert response.data.outputs == {"decision": "approve"} - - for index, execution_id in enumerate(("end", "next"), start=3): - response = converter.workflow_node_start_to_stream_response( - event=node_start(execution_id, "Result", BuiltinNodeTypes.END, index), task_id="task-1" - ) - assert response is not None - assert response.data.index == index diff --git a/api/tests/unit_tests/core/app/apps/common/test_workflow_response_converter_truncation.py b/api/tests/unit_tests/core/app/apps/common/test_workflow_response_converter_truncation.py deleted file mode 100644 index 93149691cbb..00000000000 --- a/api/tests/unit_tests/core/app/apps/common/test_workflow_response_converter_truncation.py +++ /dev/null @@ -1,885 +0,0 @@ -""" -Unit tests for WorkflowResponseConverter focusing on process_data truncation functionality. -""" - -import uuid -from collections.abc import Mapping -from dataclasses import dataclass -from datetime import UTC, datetime -from typing import Any - -import pytest - -from core.app.app_config.entities import WorkflowUIBasedAppConfig -from core.app.apps.common.workflow_response_converter import WorkflowResponseConverter -from core.app.entities.app_invoke_entities import InvokeFrom, WorkflowAppGenerateEntity -from core.app.entities.queue_entities import ( - QueueEvent, - QueueIterationStartEvent, - QueueLoopStartEvent, - QueueNodeExceptionEvent, - QueueNodeFailedEvent, - QueueNodeRetryEvent, - QueueNodeStartedEvent, - QueueNodeSucceededEvent, -) -from core.workflow.system_variables import build_system_variables -from graphon.entities import WorkflowStartReason -from graphon.enums import BuiltinNodeTypes -from libs.datetime_utils import naive_utc_now -from models import Account -from models.model import AppMode - - -class TestWorkflowResponseConverter: - """Test truncation in WorkflowResponseConverter.""" - - def create_generate_entity(self) -> WorkflowAppGenerateEntity: - """Create a validated workflow invocation for response conversion.""" - return WorkflowAppGenerateEntity( - app_config=WorkflowUIBasedAppConfig( - tenant_id="test-tenant-id", app_id="app", workflow_id="wf-id", app_mode=AppMode.WORKFLOW - ), - task_id="task", - workflow_execution_id="initial-run-id", - user_id="test-user-id", - invoke_from=InvokeFrom.WEB_APP, - inputs={}, - files=[], - stream=True, - ) - - def create_workflow_response_converter(self) -> WorkflowResponseConverter: - """Create a WorkflowResponseConverter for testing.""" - - entity = self.create_generate_entity() - mock_user = Account( - name="Test User", - email="test@example.com", - ) - mock_user.id = "test-user-id" - - system_variables = build_system_variables(workflow_id="wf-id", workflow_execution_id="initial-run-id") - return WorkflowResponseConverter( - application_generate_entity=entity, - user=mock_user, - system_variables=system_variables, - ) - - def create_node_started_event(self, *, node_execution_id: str | None = None) -> QueueNodeStartedEvent: - """Create a QueueNodeStartedEvent for testing.""" - return QueueNodeStartedEvent( - node_execution_id=node_execution_id or str(uuid.uuid4()), - node_id="test-node-id", - node_title="Test Node", - node_type=BuiltinNodeTypes.CODE, - start_at=naive_utc_now(), - in_iteration_id=None, - in_loop_id=None, - provider_type="built-in", - provider_id="code", - ) - - def create_node_succeeded_event( - self, - *, - node_execution_id: str, - process_data: Mapping[str, Any] | None = None, - ) -> QueueNodeSucceededEvent: - """Create a QueueNodeSucceededEvent for testing.""" - return QueueNodeSucceededEvent( - node_id="test-node-id", - node_type=BuiltinNodeTypes.CODE, - node_execution_id=node_execution_id, - start_at=naive_utc_now(), - in_iteration_id=None, - in_loop_id=None, - inputs={}, - process_data=process_data or {}, - outputs={}, - execution_metadata={}, - ) - - def create_node_retry_event( - self, - *, - node_execution_id: str, - process_data: Mapping[str, Any] | None = None, - ) -> QueueNodeRetryEvent: - """Create a QueueNodeRetryEvent for testing.""" - return QueueNodeRetryEvent( - inputs={"data": "inputs"}, - outputs={"data": "outputs"}, - process_data=process_data or {}, - error="oops", - retry_index=1, - node_id="test-node-id", - node_type=BuiltinNodeTypes.CODE, - node_title="test code", - provider_type="built-in", - provider_id="code", - node_execution_id=node_execution_id, - start_at=naive_utc_now(), - in_iteration_id=None, - in_loop_id=None, - ) - - def test_workflow_node_finish_response_uses_truncated_process_data(self): - """Test that node finish response uses get_response_process_data().""" - converter = self.create_workflow_response_converter() - - original_data = {"large_field": "x" * 10000, "metadata": "info"} - truncated_data = {"large_field": "[TRUNCATED]", "metadata": "info"} - - converter.workflow_start_to_stream_response( - task_id="bootstrap", - workflow_run_id="run-id", - workflow_id="wf-id", - reason=WorkflowStartReason.INITIAL, - ) - start_event = self.create_node_started_event() - converter.workflow_node_start_to_stream_response( - event=start_event, - task_id="test-task-id", - ) - - event = self.create_node_succeeded_event( - node_execution_id=start_event.node_execution_id, - process_data=original_data, - ) - - def fake_truncate(mapping): - if mapping == dict(original_data): - return truncated_data, True - return mapping, False - - converter._truncator.truncate_variable_mapping = fake_truncate # type: ignore[assignment] - - response = converter.workflow_node_finish_to_stream_response( - event=event, - task_id="test-task-id", - ) - - # Response should use truncated data, not original - assert response is not None - assert response.data.process_data == truncated_data - assert response.data.process_data != original_data - assert response.data.process_data_truncated is True - - def test_workflow_node_finish_response_without_truncation(self): - """Test node finish response when no truncation is applied.""" - converter = self.create_workflow_response_converter() - - original_data = {"small": "data"} - - converter.workflow_start_to_stream_response( - task_id="bootstrap", - workflow_run_id="run-id", - workflow_id="wf-id", - reason=WorkflowStartReason.INITIAL, - ) - start_event = self.create_node_started_event() - converter.workflow_node_start_to_stream_response( - event=start_event, - task_id="test-task-id", - ) - - event = self.create_node_succeeded_event( - node_execution_id=start_event.node_execution_id, - process_data=original_data, - ) - - def fake_truncate(mapping): - return mapping, False - - converter._truncator.truncate_variable_mapping = fake_truncate # type: ignore[assignment] - - response = converter.workflow_node_finish_to_stream_response( - event=event, - task_id="test-task-id", - ) - - # Response should use original data - assert response is not None - assert response.data.process_data == original_data - assert response.data.process_data_truncated is False - - def test_workflow_node_finish_response_with_none_process_data(self): - """Test node finish response when process_data is None.""" - converter = self.create_workflow_response_converter() - - converter.workflow_start_to_stream_response( - task_id="bootstrap", - workflow_run_id="run-id", - workflow_id="wf-id", - reason=WorkflowStartReason.INITIAL, - ) - start_event = self.create_node_started_event() - converter.workflow_node_start_to_stream_response( - event=start_event, - task_id="test-task-id", - ) - - event = self.create_node_succeeded_event( - node_execution_id=start_event.node_execution_id, - process_data=None, - ) - - def fake_truncate(mapping): - return mapping, False - - converter._truncator.truncate_variable_mapping = fake_truncate # type: ignore[assignment] - - response = converter.workflow_node_finish_to_stream_response( - event=event, - task_id="test-task-id", - ) - - # Response should normalize missing process_data to an empty mapping - assert response is not None - assert response.data.process_data == {} - assert response.data.process_data_truncated is False - - def test_workflow_node_finish_response_prefers_event_finished_at( - self, - monkeypatch: pytest.MonkeyPatch, - ) -> None: - """Finished timestamps should come from the event, not delayed queue processing time.""" - converter = self.create_workflow_response_converter() - start_at = datetime(2024, 1, 1, 0, 0, 0, tzinfo=UTC).replace(tzinfo=None) - finished_at = datetime(2024, 1, 1, 0, 0, 2, tzinfo=UTC).replace(tzinfo=None) - delayed_processing_time = datetime(2024, 1, 1, 0, 0, 10, tzinfo=UTC).replace(tzinfo=None) - - monkeypatch.setattr( - "core.app.apps.common.workflow_response_converter.naive_utc_now", - lambda: delayed_processing_time, - ) - converter.workflow_start_to_stream_response( - task_id="bootstrap", - workflow_run_id="run-id", - workflow_id="wf-id", - reason=WorkflowStartReason.INITIAL, - ) - - event = QueueNodeSucceededEvent( - node_id="test-node-id", - node_type=BuiltinNodeTypes.CODE, - node_execution_id="node-exec-1", - start_at=start_at, - finished_at=finished_at, - in_iteration_id=None, - in_loop_id=None, - inputs={}, - process_data={}, - outputs={}, - execution_metadata={}, - ) - - response = converter.workflow_node_finish_to_stream_response( - event=event, - task_id="test-task-id", - ) - - assert response is not None - assert response.data.elapsed_time == 2.0 - assert response.data.finished_at == int(finished_at.timestamp()) - - def test_workflow_node_retry_response_uses_truncated_process_data(self): - """Test that node retry response uses get_response_process_data().""" - converter = self.create_workflow_response_converter() - - original_data = {"large_field": "x" * 10000, "metadata": "info"} - truncated_data = {"large_field": "[TRUNCATED]", "metadata": "info"} - - converter.workflow_start_to_stream_response( - task_id="bootstrap", - workflow_run_id="run-id", - workflow_id="wf-id", - reason=WorkflowStartReason.INITIAL, - ) - start_event = self.create_node_started_event() - converter.workflow_node_start_to_stream_response( - event=start_event, - task_id="test-task-id", - ) - - event = self.create_node_retry_event( - node_execution_id=start_event.node_execution_id, - process_data=original_data, - ) - - def fake_truncate(mapping): - if mapping == dict(original_data): - return truncated_data, True - return mapping, False - - converter._truncator.truncate_variable_mapping = fake_truncate # type: ignore[assignment] - - response = converter.workflow_node_retry_to_stream_response( - event=event, - task_id="test-task-id", - ) - - # Response should use truncated data, not original - assert response is not None - assert response.data.process_data == truncated_data - assert response.data.process_data != original_data - assert response.data.process_data_truncated is True - - def test_workflow_node_retry_response_without_truncation(self): - """Test node retry response when no truncation is applied.""" - converter = self.create_workflow_response_converter() - - original_data = {"small": "data"} - - converter.workflow_start_to_stream_response( - task_id="bootstrap", - workflow_run_id="run-id", - workflow_id="wf-id", - reason=WorkflowStartReason.INITIAL, - ) - start_event = self.create_node_started_event() - converter.workflow_node_start_to_stream_response( - event=start_event, - task_id="test-task-id", - ) - - event = self.create_node_retry_event( - node_execution_id=start_event.node_execution_id, - process_data=original_data, - ) - - def fake_truncate(mapping): - return mapping, False - - converter._truncator.truncate_variable_mapping = fake_truncate # type: ignore[assignment] - - response = converter.workflow_node_retry_to_stream_response( - event=event, - task_id="test-task-id", - ) - - assert response is not None - assert response.data.process_data == original_data - assert response.data.process_data_truncated is False - - def test_iteration_and_loop_nodes_return_none(self): - """Test that iteration and loop nodes return None (no streaming events).""" - converter = self.create_workflow_response_converter() - - iteration_event = QueueNodeSucceededEvent( - node_id="iteration-node", - node_type=BuiltinNodeTypes.ITERATION, - node_execution_id=str(uuid.uuid4()), - start_at=naive_utc_now(), - in_iteration_id=None, - in_loop_id=None, - inputs={}, - process_data={}, - outputs={}, - execution_metadata={}, - ) - - response = converter.workflow_node_finish_to_stream_response( - event=iteration_event, - task_id="test-task-id", - ) - assert response is None - - loop_event = iteration_event.model_copy(update={"node_type": BuiltinNodeTypes.LOOP}) - response = converter.workflow_node_finish_to_stream_response( - event=loop_event, - task_id="test-task-id", - ) - assert response is None - - def test_finish_without_start_raises(self): - """Ensure finish responses require a prior workflow start.""" - converter = self.create_workflow_response_converter() - event = self.create_node_succeeded_event( - node_execution_id=str(uuid.uuid4()), - process_data={}, - ) - - with pytest.raises(ValueError, match="workflow_run_id missing before streaming workflow events"): - converter.workflow_node_finish_to_stream_response( - event=event, - task_id="test-task-id", - ) - - -@dataclass -class TestCase: - """Test case data for table-driven tests.""" - - name: str - invoke_from: InvokeFrom - expected_truncation_enabled: bool - description: str - - -class TestWorkflowResponseConverterServiceApiTruncation: - """Test class for Service API truncation functionality in WorkflowResponseConverter.""" - - def create_test_app_generate_entity(self, invoke_from: InvokeFrom) -> WorkflowAppGenerateEntity: - """Create a test WorkflowAppGenerateEntity with specified invoke_from.""" - # Create a minimal WorkflowUIBasedAppConfig for testing - app_config = WorkflowUIBasedAppConfig( - tenant_id="test_tenant", - app_id="test_app", - app_mode=AppMode.WORKFLOW, - workflow_id="test_workflow_id", - ) - - entity = WorkflowAppGenerateEntity( - task_id="test_task_id", - app_id="test_app_id", - app_config=app_config, - tenant_id="test_tenant", - app_mode=AppMode.WORKFLOW, - invoke_from=invoke_from, - inputs={"test_input": "test_value"}, - user_id="test_user_id", - stream=True, - files=[], - workflow_execution_id="test_workflow_exec_id", - ) - return entity - - def create_test_user(self) -> Account: - """Create a test user account.""" - account = Account( - name="Test User", - email="test@example.com", - ) - # Manually set the ID for testing purposes - account.id = "test_user_id" - return account - - def create_test_system_variables(self): - """Create test system variables.""" - return build_system_variables() - - def create_test_converter(self, invoke_from: InvokeFrom) -> WorkflowResponseConverter: - """Create WorkflowResponseConverter with specified invoke_from.""" - entity = self.create_test_app_generate_entity(invoke_from) - user = self.create_test_user() - system_variables = self.create_test_system_variables() - - converter = WorkflowResponseConverter( - application_generate_entity=entity, - user=user, - system_variables=system_variables, - ) - # ensure `workflow_run_id` is set. - converter.workflow_start_to_stream_response( - task_id="test-task-id", - workflow_run_id="test-workflow-run-id", - workflow_id="test-workflow-id", - reason=WorkflowStartReason.INITIAL, - ) - return converter - - @pytest.mark.parametrize( - "test_case", - [ - TestCase( - name="service_api_truncation_disabled", - invoke_from=InvokeFrom.SERVICE_API, - expected_truncation_enabled=False, - description="Service API calls should have truncation disabled", - ), - TestCase( - name="web_app_truncation_enabled", - invoke_from=InvokeFrom.WEB_APP, - expected_truncation_enabled=True, - description="Web app calls should have truncation enabled", - ), - TestCase( - name="debugger_truncation_enabled", - invoke_from=InvokeFrom.DEBUGGER, - expected_truncation_enabled=True, - description="Debugger calls should have truncation enabled", - ), - TestCase( - name="explore_truncation_enabled", - invoke_from=InvokeFrom.EXPLORE, - expected_truncation_enabled=True, - description="Explore calls should have truncation enabled", - ), - TestCase( - name="published_pipeline_truncation_enabled", - invoke_from=InvokeFrom.PUBLISHED_PIPELINE, - expected_truncation_enabled=True, - description="Published pipeline calls should have truncation enabled", - ), - ], - ids=lambda x: x.name, - ) - def test_truncator_selection_based_on_invoke_from(self, test_case: TestCase): - """Test that the correct truncator is selected based on invoke_from.""" - converter = self.create_test_converter(test_case.invoke_from) - - # Test truncation behavior instead of checking private attribute - - # Create a test event with large data - large_value = {"key": ["x"] * 2000} # Large data that would be truncated - - event = QueueNodeSucceededEvent( - node_execution_id="test_node_exec_id", - node_id="test_node", - node_type=BuiltinNodeTypes.LLM, - start_at=naive_utc_now(), - inputs=large_value, - process_data=large_value, - outputs=large_value, - error=None, - execution_metadata=None, - in_iteration_id=None, - in_loop_id=None, - ) - - response = converter.workflow_node_finish_to_stream_response( - event=event, - task_id="test_task", - ) - - # Verify response is not None - assert response is not None - - # Verify truncation behavior matches expectations - if test_case.expected_truncation_enabled: - # Truncation should be enabled for non-service-api calls - assert response.data.inputs_truncated - assert response.data.process_data_truncated - assert response.data.outputs_truncated - else: - # SERVICE_API should not truncate - assert not response.data.inputs_truncated - assert not response.data.process_data_truncated - assert not response.data.outputs_truncated - - def test_service_api_truncator_no_op_mapping(self): - """Test that Service API truncator doesn't truncate variable mappings.""" - converter = self.create_test_converter(InvokeFrom.SERVICE_API) - - # Create a test event with large data - large_value: dict[str, Any] = { - "large_string": "x" * 10000, # Large string - "large_list": list(range(2000)), # Large array - "nested_data": {"deep_nested": {"very_deep": {"value": "x" * 5000}}}, - } - - event = QueueNodeSucceededEvent( - node_execution_id="test_node_exec_id", - node_id="test_node", - node_type=BuiltinNodeTypes.LLM, - start_at=naive_utc_now(), - inputs=large_value, - process_data=large_value, - outputs=large_value, - error=None, - execution_metadata=None, - in_iteration_id=None, - in_loop_id=None, - ) - - response = converter.workflow_node_finish_to_stream_response( - event=event, - task_id="test_task", - ) - - # Verify response is not None - data = response.data - assert data.inputs == large_value - assert data.process_data == large_value - assert data.outputs == large_value - # Service API should not truncate - assert data.inputs_truncated is False - assert data.process_data_truncated is False - assert data.outputs_truncated is False - - def test_web_app_truncator_works_normally(self): - """Test that web app truncator still works normally.""" - converter = self.create_test_converter(InvokeFrom.WEB_APP) - - # Create a test event with large data - large_value = { - "large_string": "x" * 10000, # Large string - "large_list": list(range(2000)), # Large array - } - - event = QueueNodeSucceededEvent( - node_execution_id="test_node_exec_id", - node_id="test_node", - node_type=BuiltinNodeTypes.LLM, - start_at=naive_utc_now(), - inputs=large_value, - process_data=large_value, - outputs=large_value, - error=None, - execution_metadata=None, - in_iteration_id=None, - in_loop_id=None, - ) - - response = converter.workflow_node_finish_to_stream_response( - event=event, - task_id="test_task", - ) - - # Verify response is not None - assert response is not None - - # Web app should truncate - data = response.data - assert data.inputs != large_value - assert data.process_data != large_value - assert data.outputs != large_value - # The exact behavior depends on VariableTruncator implementation - # Just verify that truncation flags are present - assert data.inputs_truncated is True - assert data.process_data_truncated is True - assert data.outputs_truncated is True - - @staticmethod - def _create_event_by_type( - type_: QueueEvent, inputs: Mapping[str, Any], process_data: Mapping[str, Any], outputs: Mapping[str, Any] - ) -> QueueNodeSucceededEvent | QueueNodeFailedEvent | QueueNodeExceptionEvent: - if type_ == QueueEvent.NODE_SUCCEEDED: - return QueueNodeSucceededEvent( - node_execution_id="test_node_exec_id", - node_id="test_node", - node_type=BuiltinNodeTypes.LLM, - start_at=naive_utc_now(), - inputs=inputs, - process_data=process_data, - outputs=outputs, - error=None, - execution_metadata=None, - in_iteration_id=None, - in_loop_id=None, - ) - elif type_ == QueueEvent.NODE_FAILED: - return QueueNodeFailedEvent( - node_execution_id="test_node_exec_id", - node_id="test_node", - node_type=BuiltinNodeTypes.LLM, - start_at=naive_utc_now(), - inputs=inputs, - process_data=process_data, - outputs=outputs, - error="oops", - execution_metadata=None, - in_iteration_id=None, - in_loop_id=None, - ) - elif type_ == QueueEvent.NODE_EXCEPTION: - return QueueNodeExceptionEvent( - node_execution_id="test_node_exec_id", - node_id="test_node", - node_type=BuiltinNodeTypes.LLM, - start_at=naive_utc_now(), - inputs=inputs, - process_data=process_data, - outputs=outputs, - error="oops", - execution_metadata=None, - in_iteration_id=None, - in_loop_id=None, - ) - else: - raise Exception("unknown type.") - - @pytest.mark.parametrize( - "event_type", - [ - QueueEvent.NODE_SUCCEEDED, - QueueEvent.NODE_FAILED, - QueueEvent.NODE_EXCEPTION, - ], - ) - def test_service_api_node_finish_event_no_truncation(self, event_type: QueueEvent): - """Test that Service API doesn't truncate node finish events.""" - converter = self.create_test_converter(InvokeFrom.SERVICE_API) - # Create test event with large data - large_inputs = {"input1": "x" * 5000, "input2": list(range(2000))} - large_process_data = {"process1": "y" * 5000, "process2": {"nested": ["z"] * 2000}} - large_outputs = {"output1": "result" * 1000, "output2": list(range(2000))} - - event = TestWorkflowResponseConverterServiceApiTruncation._create_event_by_type( - event_type, large_inputs, large_process_data, large_outputs - ) - - response = converter.workflow_node_finish_to_stream_response( - event=event, - task_id="test_task", - ) - - # Verify response is not None - assert response is not None - - # Verify response contains full data (not truncated) - assert response.data.inputs == large_inputs - assert response.data.process_data == large_process_data - assert response.data.outputs == large_outputs - assert not response.data.inputs_truncated - assert not response.data.process_data_truncated - assert not response.data.outputs_truncated - - def test_service_api_node_retry_event_no_truncation(self): - """Test that Service API doesn't truncate node retry events.""" - converter = self.create_test_converter(InvokeFrom.SERVICE_API) - - # Create test event with large data - large_inputs = {"retry_input": "x" * 5000} - large_process_data = {"retry_process": "y" * 5000} - large_outputs = {"retry_output": "z" * 5000} - - # First, we need to store a snapshot by simulating a start event - start_event = QueueNodeStartedEvent( - node_execution_id="test_node_exec_id", - node_id="test_node", - node_type=BuiltinNodeTypes.LLM, - node_title="Test Node", - node_run_index=1, - start_at=naive_utc_now(), - in_iteration_id=None, - in_loop_id=None, - agent_strategy=None, - provider_type="plugin", - provider_id="test/test_plugin", - ) - converter.workflow_node_start_to_stream_response(event=start_event, task_id="test_task") - - # Now create retry event - event = QueueNodeRetryEvent( - node_execution_id="test_node_exec_id", - node_id="test_node", - node_type=BuiltinNodeTypes.LLM, - node_title="Test Node", - node_run_index=1, - start_at=naive_utc_now(), - inputs=large_inputs, - process_data=large_process_data, - outputs=large_outputs, - error="Retry error", - execution_metadata=None, - in_iteration_id=None, - in_loop_id=None, - retry_index=1, - provider_type="plugin", - provider_id="test/test_plugin", - ) - - response = converter.workflow_node_retry_to_stream_response( - event=event, - task_id="test_task", - ) - - # Verify response is not None - assert response is not None - - # Verify response contains full data (not truncated) - assert response.data.inputs == large_inputs - assert response.data.process_data == large_process_data - assert response.data.outputs == large_outputs - assert not response.data.inputs_truncated - assert not response.data.process_data_truncated - assert not response.data.outputs_truncated - - def test_service_api_iteration_events_no_truncation(self): - """Test that Service API doesn't truncate iteration events.""" - converter = self.create_test_converter(InvokeFrom.SERVICE_API) - - # Test iteration start event - large_value = {"iteration_input": ["x"] * 2000} - - start_event = QueueIterationStartEvent( - node_execution_id="test_iter_exec_id", - node_id="test_iteration", - node_type=BuiltinNodeTypes.ITERATION, - node_title="Test Iteration", - start_at=naive_utc_now(), - inputs=large_value, - metadata={}, - ) - - response = converter.workflow_iteration_start_to_stream_response( - task_id="test_task", - workflow_execution_id="test_workflow_exec_id", - event=start_event, - ) - - assert response is not None - assert response.data.inputs == large_value - assert not response.data.inputs_truncated - - def test_service_api_loop_events_no_truncation(self): - """Test that Service API doesn't truncate loop events.""" - converter = self.create_test_converter(InvokeFrom.SERVICE_API) - - # Test loop start event - large_inputs = {"loop_input": ["x"] * 2000} - - start_event = QueueLoopStartEvent( - node_execution_id="test_loop_exec_id", - node_id="test_loop", - node_type=BuiltinNodeTypes.LOOP, - node_title="Test Loop", - start_at=naive_utc_now(), - inputs=large_inputs, - metadata={}, - ) - - response = converter.workflow_loop_start_to_stream_response( - task_id="test_task", - workflow_execution_id="test_workflow_exec_id", - event=start_event, - ) - - assert response is not None - assert response.data.inputs == large_inputs - assert not response.data.inputs_truncated - - def test_web_app_node_finish_event_truncation_works(self): - """Test that web app still truncates node finish events.""" - converter = self.create_test_converter(InvokeFrom.WEB_APP) - - # Create test event with large data that should be truncated - large_inputs = {"input1": ["x"] * 2000} - large_process_data = {"process1": ["y"] * 2000} - large_outputs = {"output1": ["z"] * 2000} - - event = QueueNodeSucceededEvent( - node_execution_id="test_node_exec_id", - node_id="test_node", - node_type=BuiltinNodeTypes.LLM, - start_at=naive_utc_now(), - inputs=large_inputs, - process_data=large_process_data, - outputs=large_outputs, - error=None, - execution_metadata=None, - in_iteration_id=None, - in_loop_id=None, - ) - - response = converter.workflow_node_finish_to_stream_response( - event=event, - task_id="test_task", - ) - - # Verify response is not None - assert response is not None - - # Verify response contains truncated data - # The exact behavior depends on VariableTruncator implementation - # Just verify truncation flags are set correctly (may or may not be truncated depending on size) - # At minimum, the truncation mechanism should work - assert isinstance(response.data.inputs, dict) - assert response.data.inputs_truncated - assert isinstance(response.data.process_data, dict) - assert response.data.process_data_truncated - assert isinstance(response.data.outputs, dict) - assert response.data.outputs_truncated diff --git a/api/tests/unit_tests/core/app/apps/completion/test_completion_app_config_manager.py b/api/tests/unit_tests/core/app/apps/completion/test_app_config_manager.py similarity index 100% rename from api/tests/unit_tests/core/app/apps/completion/test_completion_app_config_manager.py rename to api/tests/unit_tests/core/app/apps/completion/test_app_config_manager.py diff --git a/api/tests/unit_tests/core/app/apps/completion/test_completion_completion_app_generator.py b/api/tests/unit_tests/core/app/apps/completion/test_app_generator.py similarity index 100% rename from api/tests/unit_tests/core/app/apps/completion/test_completion_completion_app_generator.py rename to api/tests/unit_tests/core/app/apps/completion/test_app_generator.py diff --git a/api/tests/unit_tests/core/app/apps/completion/test_app_runner.py b/api/tests/unit_tests/core/app/apps/completion/test_app_runner.py index ecd8499c103..6a463c05e23 100644 --- a/api/tests/unit_tests/core/app/apps/completion/test_app_runner.py +++ b/api/tests/unit_tests/core/app/apps/completion/test_app_runner.py @@ -47,7 +47,7 @@ from models.enums import ConversationFromSource from models.enums import InvokeFrom as ModelInvokeFrom from models.model import App, AppMode, Conversation, IconType, Message from services.knowledge.external.service import ExternalDatasetService -from tests.unit_tests.core.model_fixtures import make_model_config +from tests.helpers.unit.core.model_fixtures import make_model_config APP_ID = "00000000-0000-0000-0000-000000000001" TENANT_ID = "00000000-0000-0000-0000-000000000002" diff --git a/api/tests/unit_tests/core/app/apps/completion/test_completion_generate_response_converter.py b/api/tests/unit_tests/core/app/apps/completion/test_generate_response_converter.py similarity index 100% rename from api/tests/unit_tests/core/app/apps/completion/test_completion_generate_response_converter.py rename to api/tests/unit_tests/core/app/apps/completion/test_generate_response_converter.py diff --git a/api/tests/unit_tests/core/app/apps/pipeline/test_pipeline_generate_response_converter.py b/api/tests/unit_tests/core/app/apps/pipeline/test_generate_response_converter.py similarity index 100% rename from api/tests/unit_tests/core/app/apps/pipeline/test_pipeline_generate_response_converter.py rename to api/tests/unit_tests/core/app/apps/pipeline/test_generate_response_converter.py diff --git a/api/tests/unit_tests/core/app/apps/pipeline/test_pipeline_runner.py b/api/tests/unit_tests/core/app/apps/pipeline/test_pipeline_runner.py index 9632de39766..35637ce277a 100644 --- a/api/tests/unit_tests/core/app/apps/pipeline/test_pipeline_runner.py +++ b/api/tests/unit_tests/core/app/apps/pipeline/test_pipeline_runner.py @@ -35,7 +35,7 @@ from models.enums import DocumentCreatedFrom from models.model import AppMode, EndUser from models.workflow import Workflow, WorkflowType from repositories.knowledge.document_repository import SQLAlchemyDocumentRepository -from tests.unit_tests.model_factories import make_dataset, make_document, make_end_user, make_workflow +from tests.helpers.unit.model_factories import make_dataset, make_document, make_end_user, make_workflow def _pipeline(*, tenant_id: str = "tenant", pipeline_id: str = "pipe") -> Pipeline: diff --git a/api/tests/unit_tests/core/app/apps/test_advanced_chat_app_generator.py b/api/tests/unit_tests/core/app/apps/test_advanced_chat_app_generator.py deleted file mode 100644 index f8b8599162b..00000000000 --- a/api/tests/unit_tests/core/app/apps/test_advanced_chat_app_generator.py +++ /dev/null @@ -1,239 +0,0 @@ -from __future__ import annotations - -from types import SimpleNamespace -from unittest.mock import MagicMock - -import pytest -from sqlalchemy.engine import Engine -from sqlalchemy.orm import Session - -import core.app.apps.advanced_chat.app_generator as app_generator_module -from core.app.app_config.entities import AppAdditionalFeatures, WorkflowUIBasedAppConfig -from core.app.apps.advanced_chat.app_generator import AdvancedChatAppGenerator -from core.app.entities.app_invoke_entities import AdvancedChatAppGenerateEntity, InvokeFrom -from core.app.task_pipeline import message_cycle_manager -from core.app.task_pipeline.message_cycle_manager import MessageCycleManager -from core.ops.ops_trace_manager import TraceQueueManager -from models import Account, Workflow -from models.enums import ConversationFromSource -from models.model import App, AppMode, Conversation -from services.errors.conversation import ConversationNotExistsError -from services.workflow_run_agg import WorkflowRunAgg - - -def _make_app_config() -> WorkflowUIBasedAppConfig: - return WorkflowUIBasedAppConfig( - tenant_id="tenant-id", - app_id="app-id", - app_mode=AppMode.ADVANCED_CHAT, - workflow_id="workflow-id", - additional_features=AppAdditionalFeatures(), - variables=[], - ) - - -def _make_generate_entity(app_config: WorkflowUIBasedAppConfig) -> AdvancedChatAppGenerateEntity: - return AdvancedChatAppGenerateEntity( - task_id="task-id", - app_config=app_config, - file_upload_config=None, - conversation_id=None, - inputs={}, - query="hello", - files=[], - parent_message_id=None, - user_id="user-id", - stream=True, - invoke_from=InvokeFrom.WEB_APP, - extras={}, - workflow_run_id="workflow-run-id", - ) - - -def _app() -> App: - return App( - id="app-id", - tenant_id="tenant-id", - name="Advanced chat app", - description="", - mode=AppMode.ADVANCED_CHAT, - enable_site=False, - enable_api=False, - ) - - -def _workflow() -> Workflow: - return Workflow.new( - tenant_id="tenant-id", - app_id="app-id", - type="chat", - version=Workflow.VERSION_DRAFT, - graph='{"nodes": [], "edges": []}', - features="{}", - created_by="user-id", - environment_variables=[], - conversation_variables=[], - rag_pipeline_variables=[], - ) - - -def _account() -> Account: - account = Account(name="User", email="user@example.com") - account.id = "user-id" - return account - - -def test_init_generate_records_sets_conversation_metadata(sqlite_session: Session): - app_config = _make_app_config() - entity = _make_generate_entity(app_config) - - generator = AdvancedChatAppGenerator(execution_driver=WorkflowRunAgg.run) - - conversation, _ = generator._init_generate_records( - entity, - conversation=None, - session=sqlite_session, - ) - - assert entity.conversation_id == conversation.id - assert conversation.id is not None - assert entity.is_new_conversation is True - - -def test_init_generate_records_marks_existing_conversation(sqlite_session: Session): - app_config = _make_app_config() - entity = _make_generate_entity(app_config) - - existing_conversation = Conversation( - app_id=app_config.app_id, - app_model_config_id=None, - model_provider=None, - override_model_configs=None, - model_id=None, - mode=app_config.app_mode.value, - name="existing", - _inputs={}, - introduction="", - system_instruction="", - system_instruction_tokens=0, - status="normal", - invoke_from=InvokeFrom.WEB_APP.value, - from_source=ConversationFromSource.API, - from_end_user_id="user-id", - from_account_id=None, - ) - existing_conversation.id = "existing-conversation-id" - sqlite_session.add(existing_conversation) - sqlite_session.flush() - - generator = AdvancedChatAppGenerator(execution_driver=WorkflowRunAgg.run) - - conversation, _ = generator._init_generate_records( - entity, - conversation=existing_conversation, - session=sqlite_session, - ) - - assert entity.conversation_id == "existing-conversation-id" - assert conversation is existing_conversation - assert entity.is_new_conversation is False - - -def test_generate_falls_back_to_new_conversation_when_conversation_missing( - monkeypatch: pytest.MonkeyPatch, - sqlite_engine: Engine, - sqlite_session: Session, -): - app_config = _make_app_config() - workflow = _workflow() - app_model = _app() - user = _account() - - def raise_conversation_not_exists(**_kwargs): - raise ConversationNotExistsError() - - monkeypatch.setattr( - "core.app.apps.advanced_chat.app_generator.ConversationService.get_conversation", - raise_conversation_not_exists, - ) - monkeypatch.setattr( - "core.app.apps.advanced_chat.app_generator.FileUploadConfigManager.convert", - lambda *_args, **_kwargs: None, - ) - monkeypatch.setattr( - "core.app.apps.advanced_chat.app_generator.AdvancedChatAppConfigManager.get_app_config", - lambda **_kwargs: app_config, - ) - monkeypatch.setattr(app_generator_module, "db", SimpleNamespace(engine=sqlite_engine)) - trace_manager = object.__new__(TraceQueueManager) - monkeypatch.setattr( - "core.app.apps.advanced_chat.app_generator.TraceQueueManager", - lambda **_kwargs: trace_manager, - ) - monkeypatch.setattr( - "core.app.apps.advanced_chat.app_generator.DifyCoreRepositoryFactory.create_workflow_execution_repository", - lambda **_kwargs: SimpleNamespace(), - ) - monkeypatch.setattr( - "core.app.apps.advanced_chat.app_generator.DifyCoreRepositoryFactory.create_workflow_node_execution_repository", - lambda **_kwargs: SimpleNamespace(), - ) - - captured: dict[str, object] = {} - session = sqlite_session - - def fake_generate(self, **kwargs): - captured.update(kwargs) - return {"status": "ok"} - - monkeypatch.setattr(AdvancedChatAppGenerator, "_generate", fake_generate) - - result = AdvancedChatAppGenerator(execution_driver=WorkflowRunAgg.run).generate( - app_model=app_model, - workflow=workflow, - user=user, - args={"inputs": {}, "query": "hello", "conversation_id": "missing-conversation-id"}, - invoke_from=InvokeFrom.SERVICE_API, - workflow_run_id="workflow-run-id", - streaming=False, - session=session, - ) - - assert result == {"status": "ok"} - assert captured["conversation"] is None - assert captured["session"] is session - application_generate_entity = captured["application_generate_entity"] - assert isinstance(application_generate_entity, AdvancedChatAppGenerateEntity) - assert application_generate_entity.conversation_id is None - - -def test_message_cycle_manager_uses_new_conversation_flag(monkeypatch: pytest.MonkeyPatch): - app_config = _make_app_config() - entity = _make_generate_entity(app_config) - entity.conversation_id = "existing-conversation-id" - entity.is_new_conversation = True - entity.extras = {"auto_generate_conversation_name": True} - - captured = {} - - class DummyThread: - def __init__(self, **kwargs): - self.kwargs = kwargs - self.started = False - - def start(self): - self.started = True - - def fake_thread[**P](*args: P.args, **kwargs: P.kwargs): - thread = DummyThread(**kwargs) - captured["thread"] = thread - return thread - - monkeypatch.setattr(message_cycle_manager, "Timer", fake_thread) - - manager = MessageCycleManager(application_generate_entity=entity, task_state=MagicMock()) - thread = manager.generate_conversation_name(conversation_id="existing-conversation-id", query="hello") - - assert thread is captured["thread"] - assert thread.started is True - assert entity.is_new_conversation is False diff --git a/api/tests/unit_tests/core/app/apps/test_base_app_queue_manager.py b/api/tests/unit_tests/core/app/apps/test_base_app_queue_manager.py index 36a1da95b72..9c3186d79aa 100644 --- a/api/tests/unit_tests/core/app/apps/test_base_app_queue_manager.py +++ b/api/tests/unit_tests/core/app/apps/test_base_app_queue_manager.py @@ -7,11 +7,15 @@ import pytest from redis.exceptions import ConnectionError as RedisConnectionError from core.app.apps.base_app_queue_manager import AppQueueManager, PublishFrom +from core.app.apps.exc import GenerateTaskStoppedError from core.app.apps.execution_coordinator import AppExecutionState +from core.app.apps.message_based_app_queue_manager import MessageBasedAppQueueManager +from core.app.apps.workflow.app_queue_manager import WorkflowAppQueueManager from core.app.entities.app_invoke_entities import InvokeFrom -from core.app.entities.queue_entities import QueueErrorEvent +from core.app.entities.queue_entities import QueueErrorEvent, QueueTextChunkEvent from graphon.engine.command import AbortCommand from models import Tenant +from models.model import AppMode class DummyQueueManager(AppQueueManager): @@ -140,3 +144,51 @@ class TestBaseAppQueueManager: assert isinstance(command, AbortCommand) assert command.reason == "stream closed" assert "Failed to send Engine abort command for task t1" in caplog.text + + +def _message_queue_manager(app_mode: str) -> MessageBasedAppQueueManager: + with patch("core.app.apps.base_app_queue_manager.redis_client") as mock_redis: + mock_redis.setex.return_value = True + return MessageBasedAppQueueManager( + task_id="task-1", + user_id="user-1", + invoke_from=InvokeFrom.DEBUGGER, + conversation_id="conversation-1", + app_mode=app_mode, + message_id="message-1", + ) + + +def _workflow_queue_manager(app_mode: str) -> WorkflowAppQueueManager: + with patch("core.app.apps.base_app_queue_manager.redis_client") as mock_redis: + mock_redis.setex.return_value = True + return WorkflowAppQueueManager( + task_id="task-1", + user_id="user-1", + invoke_from=InvokeFrom.DEBUGGER, + app_mode=app_mode, + ) + + +def test_message_queue_does_not_raise_legacy_stop_for_advanced_chat() -> None: + manager = _message_queue_manager(AppMode.ADVANCED_CHAT.value) + + with patch.object(manager, "_is_stopped", return_value=True): + manager.publish(QueueTextChunkEvent(text="chunk"), PublishFrom.APPLICATION_MANAGER) + + +def test_workflow_queue_does_not_read_legacy_stop_flag() -> None: + manager = _workflow_queue_manager(AppMode.WORKFLOW.value) + + with patch.object(manager, "_is_stopped", return_value=True) as is_stopped: + manager.publish(QueueTextChunkEvent(text="chunk"), PublishFrom.APPLICATION_MANAGER) + + is_stopped.assert_not_called() + + +def test_message_queue_keeps_legacy_stop_for_non_graphengine_chat() -> None: + manager = _message_queue_manager(AppMode.CHAT.value) + + with patch.object(manager, "_is_stopped", return_value=True): + with pytest.raises(GenerateTaskStoppedError): + manager.publish(QueueTextChunkEvent(text="chunk"), PublishFrom.APPLICATION_MANAGER) diff --git a/api/tests/unit_tests/core/app/apps/test_base_app_runner.py b/api/tests/unit_tests/core/app/apps/test_base_app_runner.py index b6d387338da..563f8c188b6 100644 --- a/api/tests/unit_tests/core/app/apps/test_base_app_runner.py +++ b/api/tests/unit_tests/core/app/apps/test_base_app_runner.py @@ -2,6 +2,8 @@ from __future__ import annotations import logging from dataclasses import dataclass +from unittest.mock import MagicMock, patch +from uuid import uuid4 import pytest from sqlalchemy import event, func, select @@ -28,6 +30,7 @@ from core.app.entities.queue_entities import ( QueueMessageEndEvent, QueueMessageFileEvent, ) +from graphon.file import FileTransferMethod, FileType from graphon.model_runtime.entities.llm_entities import LLMResult, LLMResultChunk, LLMResultChunkDelta, LLMUsage from graphon.model_runtime.entities.message_entities import ( AssistantPromptMessage, @@ -37,7 +40,332 @@ from graphon.model_runtime.entities.message_entities import ( ) from graphon.model_runtime.entities.model_entities import AIModelEntity, ModelPropertyKey from graphon.model_runtime.errors.invoke import InvokeBadRequestError +from models.enums import CreatorUserRole from models.model import App, AppMode, Message, MessageFile +from models.tools import ToolFile + +"""Test multimodal image output handling in BaseAppRunner.""" + + +class TestBaseAppRunnerMultimodal: + """Test that BaseAppRunner correctly handles multimodal image content.""" + + @pytest.fixture + def mock_user_id(self): + """Mock user ID.""" + return str(uuid4()) + + @pytest.fixture + def mock_tenant_id(self): + """Mock tenant ID.""" + return str(uuid4()) + + @pytest.fixture + def mock_message_id(self): + """Mock message ID.""" + return str(uuid4()) + + @pytest.fixture + def mock_queue_manager(self): + """Create a mock queue manager.""" + manager = MagicMock() + manager.invoke_from = InvokeFrom.SERVICE_API + return manager + + @pytest.fixture + def tool_file(self, mock_user_id: str, mock_tenant_id: str) -> ToolFile: + """Create a real transient tool-file model returned by the external file manager.""" + return ToolFile( + user_id=mock_user_id, + tenant_id=mock_tenant_id, + conversation_id=None, + file_key="generated/image.png", + mimetype="image/png", + original_url="http://example.com/image.png", + name="image.png", + size=68, + ) + + def test_handle_multimodal_image_content_with_url( + self, + mock_user_id, + mock_tenant_id, + mock_message_id, + mock_queue_manager, + tool_file, + sqlite_session: Session, + ): + """Test handling image from URL.""" + # Arrange + image_url = "http://example.com/image.png" + content = ImagePromptMessageContent( + url=image_url, + format="png", + mime_type="image/png", + ) + + with patch("core.app.apps.base_app_runner.ToolFileManager", autospec=True) as mock_mgr_class: + # Setup mock tool file manager + mock_mgr = MagicMock() + mock_mgr.create_file_by_url.return_value = tool_file + mock_mgr_class.return_value = mock_mgr + + message_file_id = AppRunner()._handle_multimodal_image_content( + session=sqlite_session, + content=content, + message_id=mock_message_id, + user_id=mock_user_id, + tenant_id=mock_tenant_id, + queue_manager=mock_queue_manager, + ) + + mock_mgr.create_file_by_url.assert_called_once_with( + user_id=mock_user_id, + tenant_id=mock_tenant_id, + file_url=image_url, + conversation_id=None, + ) + message_file = sqlite_session.get(MessageFile, message_file_id) + assert message_file is not None + assert message_file.message_id == mock_message_id + assert message_file.type == FileType.IMAGE + assert message_file.transfer_method == FileTransferMethod.TOOL_FILE + assert message_file.belongs_to == "assistant" + assert message_file.created_by == mock_user_id + assert message_file.upload_file_id == tool_file.id + mock_queue_manager.publish.assert_not_called() + + def test_handle_multimodal_image_content_with_base64( + self, + mock_user_id, + mock_tenant_id, + mock_message_id, + mock_queue_manager, + tool_file, + sqlite_session: Session, + ): + """Test handling image from base64 data.""" + # Arrange + import base64 + + # Create a small test image (1x1 PNG) + test_image_data = base64.b64encode( + b"\x89PNG\r\n\x1a\n\x00\x00\x00\rIHDR\x00\x00\x00\x01\x00\x00\x00\x01\x08\x02\x00\x00\x00\x90wS\xde" + ).decode() + content = ImagePromptMessageContent( + base64_data=test_image_data, + format="png", + mime_type="image/png", + ) + + with patch("core.app.apps.base_app_runner.ToolFileManager", autospec=True) as mock_mgr_class: + # Setup mock tool file manager + mock_mgr = MagicMock() + mock_mgr.create_file_by_raw.return_value = tool_file + mock_mgr_class.return_value = mock_mgr + + message_file_id = AppRunner()._handle_multimodal_image_content( + session=sqlite_session, + content=content, + message_id=mock_message_id, + user_id=mock_user_id, + tenant_id=mock_tenant_id, + queue_manager=mock_queue_manager, + ) + + mock_mgr.create_file_by_raw.assert_called_once() + call_kwargs = mock_mgr.create_file_by_raw.call_args[1] + assert call_kwargs["user_id"] == mock_user_id + assert call_kwargs["tenant_id"] == mock_tenant_id + assert call_kwargs["conversation_id"] is None + assert "file_binary" in call_kwargs + assert call_kwargs["mimetype"] == "image/png" + assert call_kwargs["filename"].startswith("generated_image") + assert call_kwargs["filename"].endswith(".png") + assert sqlite_session.get(MessageFile, message_file_id) is not None + mock_queue_manager.publish.assert_not_called() + + def test_handle_multimodal_image_content_with_base64_data_uri( + self, + mock_user_id, + mock_tenant_id, + mock_message_id, + mock_queue_manager, + tool_file, + sqlite_session: Session, + ): + """Test handling image from base64 data with URI prefix.""" + # Arrange + # Data URI format: data:image/png;base64, + test_image_data = ( + "iVBORw0KGgoAAAANSUhEUgAAAAEAAAABCAYAAAAfFcSJAAAADUlEQVR42mNk+M9QDwADhgGAWjR9awAAAABJRU5ErkJggg==" + ) + content = ImagePromptMessageContent( + base64_data=f"data:image/png;base64,{test_image_data}", + format="png", + mime_type="image/png", + ) + + with patch("core.app.apps.base_app_runner.ToolFileManager", autospec=True) as mock_mgr_class: + # Setup mock tool file manager + mock_mgr = MagicMock() + mock_mgr.create_file_by_raw.return_value = tool_file + mock_mgr_class.return_value = mock_mgr + + message_file_id = AppRunner()._handle_multimodal_image_content( + session=sqlite_session, + content=content, + message_id=mock_message_id, + user_id=mock_user_id, + tenant_id=mock_tenant_id, + queue_manager=mock_queue_manager, + ) + + mock_mgr.create_file_by_raw.assert_called_once() + call_kwargs = mock_mgr.create_file_by_raw.call_args[1] + assert "file_binary" in call_kwargs + assert sqlite_session.get(MessageFile, message_file_id) is not None + + def test_handle_multimodal_image_content_without_url_or_base64( + self, + mock_user_id, + mock_tenant_id, + mock_message_id, + mock_queue_manager, + sqlite_session: Session, + ): + """Test handling image content without URL or base64 data.""" + # Arrange + content = ImagePromptMessageContent( + url="", + base64_data="", + format="png", + mime_type="image/png", + ) + + with patch("core.app.apps.base_app_runner.ToolFileManager", autospec=True) as mock_mgr_class: + result = AppRunner()._handle_multimodal_image_content( + session=sqlite_session, + content=content, + message_id=mock_message_id, + user_id=mock_user_id, + tenant_id=mock_tenant_id, + queue_manager=mock_queue_manager, + ) + + assert result is None + assert sqlite_session.scalar(select(func.count()).select_from(MessageFile)) == 0 + mock_mgr_class.assert_not_called() + mock_queue_manager.publish.assert_not_called() + + def test_handle_multimodal_image_content_with_error( + self, + mock_user_id, + mock_tenant_id, + mock_message_id, + mock_queue_manager, + sqlite_session: Session, + ): + """Test handling image content when an error occurs.""" + # Arrange + image_url = "http://example.com/image.png" + content = ImagePromptMessageContent( + url=image_url, + format="png", + mime_type="image/png", + ) + + with patch("core.app.apps.base_app_runner.ToolFileManager", autospec=True) as mock_mgr_class: + mock_mgr = MagicMock() + mock_mgr.create_file_by_url.side_effect = Exception("Network error") + mock_mgr_class.return_value = mock_mgr + + result = AppRunner()._handle_multimodal_image_content( + session=sqlite_session, + content=content, + message_id=mock_message_id, + user_id=mock_user_id, + tenant_id=mock_tenant_id, + queue_manager=mock_queue_manager, + ) + + assert result is None + assert sqlite_session.scalar(select(func.count()).select_from(MessageFile)) == 0 + mock_queue_manager.publish.assert_not_called() + + def test_handle_multimodal_image_content_debugger_mode( + self, + mock_user_id, + mock_tenant_id, + mock_message_id, + mock_queue_manager, + tool_file, + sqlite_session: Session, + ): + """Test that debugger mode sets correct created_by_role.""" + # Arrange + image_url = "http://example.com/image.png" + content = ImagePromptMessageContent( + url=image_url, + format="png", + mime_type="image/png", + ) + mock_queue_manager.invoke_from = InvokeFrom.DEBUGGER + + with patch("core.app.apps.base_app_runner.ToolFileManager", autospec=True) as mock_mgr_class: + mock_mgr = MagicMock() + mock_mgr.create_file_by_url.return_value = tool_file + mock_mgr_class.return_value = mock_mgr + + message_file_id = AppRunner()._handle_multimodal_image_content( + session=sqlite_session, + content=content, + message_id=mock_message_id, + user_id=mock_user_id, + tenant_id=mock_tenant_id, + queue_manager=mock_queue_manager, + ) + + message_file = sqlite_session.get(MessageFile, message_file_id) + assert message_file is not None + assert message_file.created_by_role == CreatorUserRole.ACCOUNT + + def test_handle_multimodal_image_content_service_api_mode( + self, + mock_user_id, + mock_tenant_id, + mock_message_id, + mock_queue_manager, + tool_file, + sqlite_session: Session, + ): + """Test that service API mode sets correct created_by_role.""" + # Arrange + image_url = "http://example.com/image.png" + content = ImagePromptMessageContent( + url=image_url, + format="png", + mime_type="image/png", + ) + mock_queue_manager.invoke_from = InvokeFrom.SERVICE_API + + with patch("core.app.apps.base_app_runner.ToolFileManager", autospec=True) as mock_mgr_class: + mock_mgr = MagicMock() + mock_mgr.create_file_by_url.return_value = tool_file + mock_mgr_class.return_value = mock_mgr + + message_file_id = AppRunner()._handle_multimodal_image_content( + session=sqlite_session, + content=content, + message_id=mock_message_id, + user_id=mock_user_id, + tenant_id=mock_tenant_id, + queue_manager=mock_queue_manager, + ) + + message_file = sqlite_session.get(MessageFile, message_file_id) + assert message_file is not None + assert message_file.created_by_role == CreatorUserRole.END_USER @dataclass(eq=False, repr=False) diff --git a/api/tests/unit_tests/core/app/apps/test_legacy_stop_graphengine_lifecycle.py b/api/tests/unit_tests/core/app/apps/test_legacy_stop_graphengine_lifecycle.py deleted file mode 100644 index a9aafd5949d..00000000000 --- a/api/tests/unit_tests/core/app/apps/test_legacy_stop_graphengine_lifecycle.py +++ /dev/null @@ -1,59 +0,0 @@ -from unittest.mock import patch - -import pytest - -from core.app.apps.base_app_queue_manager import PublishFrom -from core.app.apps.exc import GenerateTaskStoppedError -from core.app.apps.message_based_app_queue_manager import MessageBasedAppQueueManager -from core.app.apps.workflow.app_queue_manager import WorkflowAppQueueManager -from core.app.entities.app_invoke_entities import InvokeFrom -from core.app.entities.queue_entities import QueueTextChunkEvent -from models.model import AppMode - - -def _message_queue_manager(app_mode: str) -> MessageBasedAppQueueManager: - with patch("core.app.apps.base_app_queue_manager.redis_client") as mock_redis: - mock_redis.setex.return_value = True - return MessageBasedAppQueueManager( - task_id="task-1", - user_id="user-1", - invoke_from=InvokeFrom.DEBUGGER, - conversation_id="conversation-1", - app_mode=app_mode, - message_id="message-1", - ) - - -def _workflow_queue_manager(app_mode: str) -> WorkflowAppQueueManager: - with patch("core.app.apps.base_app_queue_manager.redis_client") as mock_redis: - mock_redis.setex.return_value = True - return WorkflowAppQueueManager( - task_id="task-1", - user_id="user-1", - invoke_from=InvokeFrom.DEBUGGER, - app_mode=app_mode, - ) - - -def test_message_queue_does_not_raise_legacy_stop_for_advanced_chat() -> None: - manager = _message_queue_manager(AppMode.ADVANCED_CHAT.value) - - with patch.object(manager, "_is_stopped", return_value=True): - manager.publish(QueueTextChunkEvent(text="chunk"), PublishFrom.APPLICATION_MANAGER) - - -def test_workflow_queue_does_not_read_legacy_stop_flag() -> None: - manager = _workflow_queue_manager(AppMode.WORKFLOW.value) - - with patch.object(manager, "_is_stopped", return_value=True) as is_stopped: - manager.publish(QueueTextChunkEvent(text="chunk"), PublishFrom.APPLICATION_MANAGER) - - is_stopped.assert_not_called() - - -def test_message_queue_keeps_legacy_stop_for_non_graphengine_chat() -> None: - manager = _message_queue_manager(AppMode.CHAT.value) - - with patch.object(manager, "_is_stopped", return_value=True): - with pytest.raises(GenerateTaskStoppedError): - manager.publish(QueueTextChunkEvent(text="chunk"), PublishFrom.APPLICATION_MANAGER) diff --git a/api/tests/unit_tests/core/app/apps/test_message_based_app_generator.py b/api/tests/unit_tests/core/app/apps/test_message_based_app_generator.py index c130b350b04..c2c9d7443d9 100644 --- a/api/tests/unit_tests/core/app/apps/test_message_based_app_generator.py +++ b/api/tests/unit_tests/core/app/apps/test_message_based_app_generator.py @@ -19,7 +19,7 @@ from core.app.entities.app_invoke_entities import ChatAppGenerateEntity, InvokeF from models.account import Account from models.model import App, AppMode, Conversation, Message from services.errors.app_model_config import AppModelConfigBrokenError -from tests.unit_tests.model_factories import make_account, make_app +from tests.helpers.unit.model_factories import make_account, make_app @dataclass(eq=False, repr=False) diff --git a/api/tests/unit_tests/core/app/apps/test_pause_resume.py b/api/tests/unit_tests/core/app/apps/test_pause_resume.py deleted file mode 100644 index 660fac5dce1..00000000000 --- a/api/tests/unit_tests/core/app/apps/test_pause_resume.py +++ /dev/null @@ -1,376 +0,0 @@ -import sys -import time -from collections.abc import Mapping -from decimal import Decimal -from types import ModuleType, SimpleNamespace -from typing import Any - -from pytest_mock import MockerFixture -from sqlalchemy.orm import Session - -import core.workflow.nodes.human_input.entities # noqa: F401 -from core.app.apps.advanced_chat import app_generator as adv_app_gen_module -from core.app.apps.workflow import app_generator as wf_app_gen_module -from core.app.entities.app_invoke_entities import InvokeFrom -from core.workflow import node_factory as node_factory_module -from core.workflow.node_factory import DifyNodeFactory -from core.workflow.system_variables import build_system_variables -from graphon.engine import Engine -from graphon.engine.command import InMemoryChannel -from graphon.engine_events import ( - EngineEvent, - GraphRunPausedEvent, - GraphRunStartedEvent, - GraphRunSucceededEvent, - NodeRunSucceededEvent, -) -from graphon.entities import WorkflowStartReason -from graphon.entities.base_node_data import BaseNodeData, RetryConfig -from graphon.entities.pause_reason import SchedulingPause -from graphon.enums import BuiltinNodeTypes, NodeType, WorkflowNodeExecutionStatus -from graphon.graph import Graph -from graphon.node_events import NodeRunResult, PauseRequestedEvent -from graphon.nodes.base.entities import OutputVariableEntity -from graphon.nodes.base.node import Node -from graphon.nodes.end.entities import EndNodeData -from graphon.nodes.start.entities import StartNodeData -from graphon.runtime import RuntimeState, VariablePool -from models.account import Account -from models.enums import ConversationFromSource -from models.model import App, AppMode, Conversation, Message -from models.workflow import Workflow, WorkflowType -from services.workflow_run_agg import WorkflowRunAgg -from tests.workflow_test_utils import build_test_graph_init_params - -if "core.ops.ops_trace_manager" not in sys.modules: - ops_stub = ModuleType("core.ops.ops_trace_manager") - - class _StubTraceQueueManager: - def __init__(self, *_, **__): - pass - - ops_stub.TraceQueueManager = _StubTraceQueueManager - sys.modules["core.ops.ops_trace_manager"] = ops_stub - - -class _StubToolNodeData(BaseNodeData): - type: NodeType = BuiltinNodeTypes.TOOL - pause_on: bool = False - - -class _StubToolNode(Node[_StubToolNodeData]): - node_type = BuiltinNodeTypes.TOOL - - @classmethod - def version(cls) -> str: - return "1" - - def __init__( - self, - node_id: str, - data: _StubToolNodeData, - *, - init_params, - runtime_state, - **_kwargs: Any, - ) -> None: - super().__init__( - node_id=node_id, - data=data, - init_params=init_params, - runtime_state=runtime_state, - ) - - def _get_error_strategy(self): - return self._node_data.error_strategy - - def _get_retry_config(self) -> RetryConfig: - return self._node_data.retry_config - - def _get_title(self) -> str: - return self._node_data.title - - def _get_description(self): - return self._node_data.desc - - def _get_default_value_dict(self) -> dict[str, Any]: - return self._node_data.default_value_dict - - def get_base_node_data(self) -> BaseNodeData: - return self._node_data - - def _run(self): - if self.node_data.pause_on: - yield PauseRequestedEvent(reason=SchedulingPause(message="test pause")) - return - - result = NodeRunResult( - status=WorkflowNodeExecutionStatus.SUCCEEDED, - outputs={"value": f"{self.id}-done"}, - ) - yield self._convert_node_run_result_to_node_event(result) - - -def _patch_tool_node(mocker: MockerFixture): - original_resolve_node_class = node_factory_module.resolve_workflow_node_class - - def _patched_resolve_node_class( - *, - node_type: NodeType, - node_version: str, - node_data: Mapping[str, Any] | BaseNodeData | None = None, - ) -> type[Node]: - if node_type == BuiltinNodeTypes.TOOL: - return _StubToolNode - return original_resolve_node_class( - node_type=node_type, - node_version=node_version, - node_data=node_data, - ) - - mocker.patch.object(node_factory_module, "resolve_workflow_node_class", side_effect=_patched_resolve_node_class) - - -def _node_data(node_type: NodeType, data: BaseNodeData) -> dict[str, object]: - node_data = data.model_dump() - node_data["type"] = str(node_type) - return node_data - - -def _build_graph_config(*, pause_on: str | None) -> dict[str, object]: - start_data = StartNodeData(title="start", variables=[]) - tool_data_a = _StubToolNodeData(title="tool", pause_on=pause_on == "tool_a") - tool_data_b = _StubToolNodeData(title="tool", pause_on=pause_on == "tool_b") - tool_data_c = _StubToolNodeData(title="tool", pause_on=pause_on == "tool_c") - end_data = EndNodeData( - title="end", - outputs=[OutputVariableEntity(variable="result", value_selector=["tool_c", "value"])], - desc=None, - ) - - nodes = [ - {"id": "start", "data": _node_data(BuiltinNodeTypes.START, start_data)}, - {"id": "tool_a", "data": _node_data(BuiltinNodeTypes.TOOL, tool_data_a)}, - {"id": "tool_b", "data": _node_data(BuiltinNodeTypes.TOOL, tool_data_b)}, - {"id": "tool_c", "data": _node_data(BuiltinNodeTypes.TOOL, tool_data_c)}, - {"id": "end", "data": _node_data(BuiltinNodeTypes.END, end_data)}, - ] - edges = [ - {"source": "start", "target": "tool_a"}, - {"source": "tool_a", "target": "tool_b"}, - {"source": "tool_b", "target": "tool_c"}, - {"source": "tool_c", "target": "end"}, - ] - return {"nodes": nodes, "edges": edges} - - -def _build_graph(runtime_state: RuntimeState, *, pause_on: str | None) -> Graph: - graph_config = _build_graph_config(pause_on=pause_on) - params = build_test_graph_init_params( - workflow_id="workflow", - graph_config=graph_config, - tenant_id="tenant", - app_id="app", - user_id="user", - user_from="account", - invoke_from="service-api", - call_depth=0, - ) - - node_factory = DifyNodeFactory( - init_params=params, - runtime_state=runtime_state, - ) - - return Graph.init(graph_config=graph_config, node_factory=node_factory, root_node_id="start") - - -def _build_runtime_state(run_id: str) -> RuntimeState: - variable_pool = VariablePool.from_bootstrap( - system_variables=build_system_variables(user_id="user", app_id="app", workflow_id="workflow"), - user_inputs={}, - conversation_variables=[], - ) - variable_pool.add(("sys", "workflow_run_id"), run_id) - return RuntimeState(workflow_id="test-workflow", variable_pool=variable_pool, start_at=time.perf_counter()) - - -def _make_app() -> App: - return App( - id="app", - tenant_id="tenant", - name="Pause Resume App", - mode=AppMode.WORKFLOW, - enable_site=False, - enable_api=False, - ) - - -def _make_workflow() -> Workflow: - return Workflow( - id="workflow", - tenant_id="tenant", - app_id="app", - type=WorkflowType.CHAT, - version=Workflow.VERSION_DRAFT, - graph="{}", - features="{}", - created_by="user", - ) - - -def _make_account() -> Account: - account = Account(name="Pause Resume User", email="pause@example.com") - account.id = "user" - return account - - -def _make_conversation() -> Conversation: - return Conversation( - id="conv", - app_id="app", - mode=AppMode.ADVANCED_CHAT, - name="Pause Resume Conversation", - _inputs={}, - from_source=ConversationFromSource.API, - ) - - -def _make_message() -> Message: - return Message( - id="msg", - app_id="app", - conversation_id="conv", - _inputs={}, - query="query", - message={}, - answer="answer", - message_unit_price=Decimal(0), - answer_unit_price=Decimal(0), - currency="USD", - from_source=ConversationFromSource.API, - ) - - -def _run_with_optional_pause(runtime_state: RuntimeState, *, pause_on: str | None) -> list[EngineEvent]: - command_channel = InMemoryChannel() - graph = _build_graph(runtime_state, pause_on=pause_on) - engine = Engine( - graph=graph, - runtime_state=runtime_state, - command_channel=command_channel, - ) - - events: list[EngineEvent] = [] - for event in engine.run(): - events.append(event) - return events - - -def _node_successes(events: list[EngineEvent]) -> list[str]: - return [evt.node_id for evt in events if isinstance(evt, NodeRunSucceededEvent)] - - -def test_workflow_app_pause_resume_matches_baseline(mocker: MockerFixture): - _patch_tool_node(mocker) - - baseline_state = _build_runtime_state("baseline") - baseline_events = _run_with_optional_pause(baseline_state, pause_on=None) - assert isinstance(baseline_events[-1], GraphRunSucceededEvent) - baseline_nodes = _node_successes(baseline_events) - baseline_outputs = baseline_state.outputs - - paused_state = _build_runtime_state("paused-run") - paused_events = _run_with_optional_pause(paused_state, pause_on="tool_a") - assert isinstance(paused_events[-1], GraphRunPausedEvent) - paused_nodes = _node_successes(paused_events) - snapshot = paused_state.dumps() - - resumed_state = RuntimeState.from_snapshot(snapshot) - - generator = wf_app_gen_module.WorkflowAppGenerator(execution_driver=WorkflowRunAgg.run) - - def _fake_generate(**kwargs): - state: RuntimeState = kwargs["graph_runtime_state"] - events = _run_with_optional_pause(state, pause_on=None) - return _node_successes(events) - - mocker.patch.object(generator, "_generate", side_effect=_fake_generate) - - resumed_nodes = generator.resume( - app_model=_make_app(), - workflow=_make_workflow(), - user=_make_account(), - application_generate_entity=SimpleNamespace( - stream=False, - invoke_from=InvokeFrom.SERVICE_API, - trace_manager=SimpleNamespace(), - ), - graph_runtime_state=resumed_state, - workflow_execution_repository=SimpleNamespace(), - workflow_node_execution_repository=SimpleNamespace(), - ) - - assert paused_nodes + resumed_nodes == baseline_nodes - assert resumed_state.outputs == baseline_outputs - - -def test_advanced_chat_pause_resume_matches_baseline(mocker: MockerFixture, unbound_session: Session): - _patch_tool_node(mocker) - - baseline_state = _build_runtime_state("adv-baseline") - baseline_events = _run_with_optional_pause(baseline_state, pause_on=None) - assert isinstance(baseline_events[-1], GraphRunSucceededEvent) - baseline_nodes = _node_successes(baseline_events) - baseline_outputs = baseline_state.outputs - - paused_state = _build_runtime_state("adv-paused") - paused_events = _run_with_optional_pause(paused_state, pause_on="tool_a") - assert isinstance(paused_events[-1], GraphRunPausedEvent) - paused_nodes = _node_successes(paused_events) - snapshot = paused_state.dumps() - - resumed_state = RuntimeState.from_snapshot(snapshot) - - generator = adv_app_gen_module.AdvancedChatAppGenerator(execution_driver=WorkflowRunAgg.run) - - def _fake_generate(**kwargs): - state: RuntimeState = kwargs["graph_runtime_state"] - events = _run_with_optional_pause(state, pause_on=None) - return _node_successes(events) - - mocker.patch.object(generator, "_generate", side_effect=_fake_generate) - - resumed_nodes = generator.resume( - app_model=_make_app(), - workflow=_make_workflow(), - user=_make_account(), - conversation=_make_conversation(), - message=_make_message(), - session=unbound_session, - application_generate_entity=SimpleNamespace( - stream=False, - invoke_from=InvokeFrom.SERVICE_API, - trace_manager=SimpleNamespace(), - ), - workflow_execution_repository=SimpleNamespace(), - workflow_node_execution_repository=SimpleNamespace(), - graph_runtime_state=resumed_state, - ) - - assert paused_nodes + resumed_nodes == baseline_nodes - assert resumed_state.outputs == baseline_outputs - - -def test_resume_emits_resumption_start_reason(mocker: MockerFixture) -> None: - _patch_tool_node(mocker) - - paused_state = _build_runtime_state("resume-reason") - paused_events = _run_with_optional_pause(paused_state, pause_on="tool_a") - initial_start = next(event for event in paused_events if isinstance(event, GraphRunStartedEvent)) - assert initial_start.reason == WorkflowStartReason.INITIAL - - resumed_state = RuntimeState.from_snapshot(paused_state.dumps()) - resumed_events = _run_with_optional_pause(resumed_state, pause_on=None) - resume_start = next(event for event in resumed_events if isinstance(event, GraphRunStartedEvent)) - assert resume_start.reason == WorkflowStartReason.RESUMPTION diff --git a/api/tests/unit_tests/core/app/apps/test_trace_session_id_generate_extras.py b/api/tests/unit_tests/core/app/apps/test_trace_session_id_generate_extras.py deleted file mode 100644 index ef261689752..00000000000 --- a/api/tests/unit_tests/core/app/apps/test_trace_session_id_generate_extras.py +++ /dev/null @@ -1,11 +0,0 @@ -from core.helper.trace_id_helper import extract_trace_session_id_from_args - - -def test_extract_trace_session_id_from_args_for_generator_extras(): - assert extract_trace_session_id_from_args({"trace_session_id": "session-1"}) == { - "trace_session_id": "session-1", - } - - -def test_extract_trace_session_id_from_args_missing_value_keeps_extras_clean(): - assert extract_trace_session_id_from_args({"inputs": {}}) == {} diff --git a/api/tests/unit_tests/core/app/apps/test_workflow_app_generator.py b/api/tests/unit_tests/core/app/apps/test_workflow_app_generator.py deleted file mode 100644 index af7b0ba1c4f..00000000000 --- a/api/tests/unit_tests/core/app/apps/test_workflow_app_generator.py +++ /dev/null @@ -1,570 +0,0 @@ -"""SQLite-backed tests for workflow app generation and worker reload behavior.""" - -import contextlib -import contextvars -import json -from types import SimpleNamespace -from unittest.mock import MagicMock - -import pytest -from flask import Flask -from sqlalchemy.orm import Session, sessionmaker - -import core.app.apps.workflow.app_generator as app_generator_module -from core.app.app_config.entities import WorkflowUIBasedAppConfig -from core.app.apps.draft_variable_saver import DraftVariableSaverFactory -from core.app.apps.workflow.app_generator import WorkflowAppGenerator -from core.app.apps.workflow.app_queue_manager import WorkflowAppQueueManager -from core.app.entities.app_invoke_entities import InvokeFrom, WorkflowAppGenerateEntity -from core.app.entities.queue_entities import QueueWorkflowSucceededEvent -from core.app.layers.pause_state_persist_layer import PauseStateLayerConfig -from core.ops.ops_trace_manager import OpsTraceManager, TraceQueueManager -from core.repositories import SQLAlchemyWorkflowExecutionRepository, SQLAlchemyWorkflowNodeExecutionRepository -from graphon.enums import WorkflowExecutionStatus -from graphon.runtime import RuntimeState, VariablePool -from graphon.variable_loader import DUMMY_VARIABLE_LOADER -from models.enums import CreatorUserRole, EndUserType, WorkflowRunTriggeredFrom -from models.model import App, AppMode, EndUser -from models.snippet import CustomizedSnippet -from models.workflow import Workflow, WorkflowKind, WorkflowNodeExecutionTriggeredFrom, WorkflowRun -from services.workflow_run_agg import WorkflowRunAgg -from tests.unit_tests.model_factories import make_workflow - - -@pytest.fixture(autouse=True) -def disable_trace_delivery(monkeypatch: pytest.MonkeyPatch) -> None: - """Construct real trace managers without provider lookup or background delivery.""" - monkeypatch.setattr(OpsTraceManager, "get_ops_trace_instance", lambda app_id: None) - monkeypatch.setattr(TraceQueueManager, "start_timer", lambda self: None) - monkeypatch.setattr("core.telemetry.gateway.is_enterprise_telemetry_enabled", lambda: False) - - -def _workflow( - *, - workflow_id: str = "workflow", - app_id: str = "app", - tenant_id: str = "tenant", - kind: WorkflowKind = WorkflowKind.STANDARD, -) -> Workflow: - return make_workflow( - workflow_id=workflow_id, - tenant_id=tenant_id, - app_id=app_id, - kind=kind, - version="1", - created_by="creator", - graph={"nodes": [{"id": "start", "data": {"type": "start", "title": "Start", "variables": []}}], "edges": []}, - ) - - -def _persist_generator_rows(sqlite_session: Session) -> tuple[App, Workflow, EndUser]: - workflow = _workflow() - app = App( - id="app", - tenant_id="tenant", - name="Workflow app", - description="", - mode=AppMode.WORKFLOW, - workflow_id=workflow.id, - enable_site=True, - enable_api=True, - max_active_requests=0, - ) - end_user = EndUser( - id="user", - tenant_id="tenant", - app_id=app.id, - type=EndUserType.SERVICE_API, - name="End user", - session_id="end-user-session", - ) - sqlite_session.add_all([workflow, app, end_user]) - sqlite_session.commit() - return app, workflow, end_user - - -def _app_config(app: App, workflow: Workflow) -> WorkflowUIBasedAppConfig: - return WorkflowUIBasedAppConfig( - tenant_id=app.tenant_id, - app_id=app.id, - app_mode=app.mode, - workflow_id=workflow.id, - ) - - -def _generate_entity( - app: App, - workflow: Workflow, - end_user: EndUser, - *, - invoke_from: InvokeFrom = InvokeFrom.SERVICE_API, - stream: bool = True, -) -> WorkflowAppGenerateEntity: - return WorkflowAppGenerateEntity( - task_id="task", - app_config=_app_config(app, workflow), - inputs={}, - files=[], - user_id=end_user.id, - stream=stream, - invoke_from=invoke_from, - trace_manager=TraceQueueManager(app_id=app.id, user_id=end_user.id), - workflow_execution_id="run", - ) - - -def _runtime_state() -> RuntimeState: - return RuntimeState(workflow_id="test-workflow", variable_pool=VariablePool(), start_at=0.0) - - -def _repositories( - sqlite_session_factory: sessionmaker[Session], app: App, end_user: EndUser -) -> tuple[SQLAlchemyWorkflowExecutionRepository, SQLAlchemyWorkflowNodeExecutionRepository]: - return ( - SQLAlchemyWorkflowExecutionRepository( - session_factory=sqlite_session_factory, - tenant_id=app.tenant_id, - user=end_user, - app_id=app.id, - triggered_from=WorkflowRunTriggeredFrom.APP_RUN, - ), - SQLAlchemyWorkflowNodeExecutionRepository( - session_factory=sqlite_session_factory, - tenant_id=app.tenant_id, - user=end_user, - app_id=app.id, - triggered_from=WorkflowNodeExecutionTriggeredFrom.WORKFLOW_RUN, - ), - ) - - -def test_worker_executes_and_persists_a_workflow_through_the_injected_driver( - app: Flask, sqlite_session: Session, sqlite_session_factory: sessionmaker[Session] -) -> None: - app_model, workflow, end_user = _persist_generator_rows(sqlite_session) - workflow.graph = json.dumps( - { - "nodes": [ - {"id": "start", "data": {"type": "start", "title": "Start", "variables": []}}, - { - "id": "end", - "data": { - "type": "end", - "title": "End", - "outputs": [{"variable": "caller", "value_selector": ["sys", "user_id"]}], - }, - }, - ], - "edges": [{"id": "start-end", "source": "start", "target": "end"}], - } - ) - sqlite_session.commit() - entity = WorkflowAppGenerateEntity( - task_id="worker-task", - app_config=_app_config(app_model, workflow), - inputs={}, - files=[], - user_id=end_user.id, - stream=True, - invoke_from=InvokeFrom.SERVICE_API, - workflow_execution_id="worker-run", - ) - queue = WorkflowAppQueueManager(entity.task_id, end_user.id, entity.invoke_from, AppMode.WORKFLOW) - execution_repository, node_repository = _repositories(sqlite_session_factory, app_model, end_user) - - WorkflowAppGenerator(execution_driver=WorkflowRunAgg.run)._generate_worker( - flask_app=app, - application_generate_entity=entity, - queue_manager=queue, - context=contextvars.copy_context(), - variable_loader=DUMMY_VARIABLE_LOADER, - workflow_execution_repository=execution_repository, - workflow_node_execution_repository=node_repository, - ) - - events = [message.event for message in queue.listen()] - assert isinstance(events[-1], QueueWorkflowSucceededEvent), events - assert events[-1].outputs == {"caller": end_user.session_id} - run = sqlite_session.get(WorkflowRun, entity.workflow_execution_id) - assert run is not None - assert run.status == "succeeded" - assert run.outputs_dict == {"caller": end_user.session_id} - - -def test_ensure_snippet_start_node_in_worker_returns_standard_workflow_without_lookup( - sqlite_session: Session, -) -> None: - workflow = _workflow() - sqlite_session.add(workflow) - sqlite_session.commit() - - result = WorkflowAppGenerator._ensure_snippet_start_node_in_worker(session=sqlite_session, workflow=workflow) - - assert result is workflow - - -def test_ensure_snippet_start_node_in_worker_returns_snippet_workflow_when_snippet_missing( - sqlite_session: Session, -) -> None: - workflow = _workflow(app_id="snippet-1", tenant_id="tenant-1", kind=WorkflowKind.SNIPPET) - other_tenant_snippet = CustomizedSnippet( - id="snippet-1", - tenant_id="tenant-2", - name="Other tenant snippet", - description="", - type="node", - ) - sqlite_session.add_all([workflow, other_tenant_snippet]) - sqlite_session.commit() - - result = WorkflowAppGenerator._ensure_snippet_start_node_in_worker(session=sqlite_session, workflow=workflow) - - assert result is workflow - - -def test_ensure_snippet_start_node_in_worker_applies_snippet_start_injection( - monkeypatch: pytest.MonkeyPatch, sqlite_session: Session -) -> None: - workflow = _workflow(app_id="snippet-1", tenant_id="tenant-1", kind=WorkflowKind.SNIPPET) - snippet = CustomizedSnippet( - id="snippet-1", - tenant_id="tenant-1", - name="Matching snippet", - description="", - type="node", - ) - sqlite_session.add_all([workflow, snippet]) - sqlite_session.commit() - ensure_start_node = MagicMock(return_value=workflow) - monkeypatch.setattr( - "services.snippet_generate_service.SnippetGenerateService.ensure_start_node_for_worker", - ensure_start_node, - ) - - result = WorkflowAppGenerator._ensure_snippet_start_node_in_worker(session=sqlite_session, workflow=workflow) - - assert result is workflow - ensure_start_node.assert_called_once_with(workflow, snippet) - - -def test_generate_includes_parent_trace_context_in_extras( - monkeypatch: pytest.MonkeyPatch, sqlite_session: Session -) -> None: - generator = WorkflowAppGenerator(execution_driver=WorkflowRunAgg.run) - app, workflow, end_user = _persist_generator_rows(sqlite_session) - - monkeypatch.setattr( - "core.app.apps.workflow.app_generator.WorkflowAppGenerator._bind_file_access_scope", - lambda *args, **kwargs: contextlib.nullcontext(), - ) - monkeypatch.setattr( - "core.app.apps.workflow.app_generator.WorkflowAppConfigManager.get_app_config", - lambda *args, **kwargs: _app_config(app, workflow), - ) - monkeypatch.setattr( - "core.app.apps.workflow.app_generator.file_factory.build_from_mappings", lambda *args, **kwargs: [] - ) - repository_tenant_ids: dict[str, str] = {} - workflow_execution_factory = app_generator_module.DifyCoreRepositoryFactory.create_workflow_execution_repository - workflow_node_execution_factory = ( - app_generator_module.DifyCoreRepositoryFactory.create_workflow_node_execution_repository - ) - - def create_workflow_execution_repository(**kwargs): - repository_tenant_ids["workflow"] = kwargs["tenant_id"] - return workflow_execution_factory(**kwargs) - - def create_workflow_node_execution_repository(**kwargs): - repository_tenant_ids["node"] = kwargs["tenant_id"] - return workflow_node_execution_factory(**kwargs) - - monkeypatch.setattr( - "core.app.apps.workflow.app_generator.DifyCoreRepositoryFactory.create_workflow_execution_repository", - create_workflow_execution_repository, - ) - monkeypatch.setattr( - "core.app.apps.workflow.app_generator.DifyCoreRepositoryFactory.create_workflow_node_execution_repository", - create_workflow_node_execution_repository, - ) - monkeypatch.setattr("core.app.apps.workflow.app_generator.db", SimpleNamespace(engine=sqlite_session.get_bind())) - monkeypatch.setattr(generator, "_prepare_user_inputs", lambda *, user_inputs, **kwargs: user_inputs) - - captured = {} - - def fake_generate(**kwargs): - captured.update(kwargs) - return {"data": {}} - - monkeypatch.setattr(generator, "_generate", fake_generate) - - result = generator.generate( - app_model=app, - workflow=workflow, - user=end_user, - args={ - "inputs": {"query": "hello"}, - "files": [], - "external_trace_id": "trace-1", - "parent_trace_context": { - "parent_workflow_run_id": "outer-workflow-run-1", - "parent_node_execution_id": "outer-node-execution-1", - }, - "trace_session_id": "session-1", - }, - invoke_from=InvokeFrom.SERVICE_API, - streaming=False, - call_depth=0, - ) - - assert result == {"data": {}} - application_generate_entity = captured["application_generate_entity"] - assert isinstance(application_generate_entity, WorkflowAppGenerateEntity) - extras = application_generate_entity.extras - assert extras["external_trace_id"] == "trace-1" - assert extras["parent_trace_context"].model_dump() == { - "parent_workflow_run_id": "outer-workflow-run-1", - "parent_node_execution_id": "outer-node-execution-1", - } - assert extras["trace_session_id"] == "session-1" - assert isinstance(captured["workflow_execution_repository"], SQLAlchemyWorkflowExecutionRepository) - assert isinstance(captured["workflow_node_execution_repository"], SQLAlchemyWorkflowNodeExecutionRepository) - assert repository_tenant_ids == {"workflow": app.tenant_id, "node": app.tenant_id} - - -def test_resume_delegates_to_generate( - monkeypatch: pytest.MonkeyPatch, - sqlite_session: Session, - sqlite_session_factory: sessionmaker[Session], -) -> None: - generator = WorkflowAppGenerator(execution_driver=WorkflowRunAgg.run) - app, workflow, end_user = _persist_generator_rows(sqlite_session) - mock_generate = MagicMock(return_value="ok") - monkeypatch.setattr(generator, "_generate", mock_generate) - - application_generate_entity = _generate_entity( - app, - workflow, - end_user, - invoke_from=InvokeFrom.DEBUGGER, - stream=False, - ) - runtime_state = _runtime_state() - workflow_execution_repository, workflow_node_execution_repository = _repositories( - sqlite_session_factory, app, end_user - ) - pause_config = PauseStateLayerConfig( - session_factory=sqlite_session_factory, - state_owner_user_id="owner", - ) - - result = generator.resume( - app_model=app, - workflow=workflow, - user=end_user, - application_generate_entity=application_generate_entity, - graph_runtime_state=runtime_state, - workflow_execution_repository=workflow_execution_repository, - workflow_node_execution_repository=workflow_node_execution_repository, - graph_engine_layers=("layer",), - pause_state_config=pause_config, - ) - - assert result == "ok" - mock_generate.assert_called_once() - kwargs = mock_generate.call_args.kwargs - assert kwargs["graph_runtime_state"] is runtime_state - assert kwargs["pause_state_config"] is pause_config - assert kwargs["streaming"] is False - assert kwargs["invoke_from"] == InvokeFrom.DEBUGGER - - -def test_generate_forwards_pause_configuration_and_state( - monkeypatch: pytest.MonkeyPatch, - sqlite_session: Session, - sqlite_session_factory: sessionmaker[Session], -) -> None: - generator = WorkflowAppGenerator(execution_driver=WorkflowRunAgg.run) - app, workflow, end_user = _persist_generator_rows(sqlite_session) - - queue_manager = MagicMock() - monkeypatch.setattr( - "core.app.apps.workflow.app_generator.WorkflowAppQueueManager", - MagicMock(return_value=queue_manager), - ) - - monkeypatch.setattr( - "core.app.apps.workflow.app_generator.WorkflowAppGenerateResponseConverter.convert", - MagicMock(return_value="converted"), - ) - monkeypatch.setattr(WorkflowAppGenerator, "_handle_response", MagicMock(return_value="response")) - - draft_factory_tenant_ids: list[str] = [] - get_draft_var_saver_factory = WorkflowAppGenerator._get_draft_var_saver_factory - - def get_recording_draft_var_saver_factory( - invoke_from: InvokeFrom, account: EndUser, *, tenant_id: str - ) -> DraftVariableSaverFactory: - draft_factory_tenant_ids.append(tenant_id) - return get_draft_var_saver_factory(invoke_from, account, tenant_id=tenant_id) - - monkeypatch.setattr( - WorkflowAppGenerator, - "_get_draft_var_saver_factory", - staticmethod(get_recording_draft_var_saver_factory), - ) - - db_session = sqlite_session_factory() - monkeypatch.setattr(app_generator_module.db, "session", db_session) - - worker_kwargs: dict[str, object] = {} - - class DummyThread: - def __init__(self, target, kwargs): - worker_kwargs["target"] = target - worker_kwargs["kwargs"] = kwargs - - def start(self): - return None - - def join(self, timeout): - worker_kwargs["joined"] = True - worker_kwargs["join_timeout"] = timeout - - def is_alive(self): - return False - - monkeypatch.setattr("core.app.apps.workflow.app_generator.threading.Thread", DummyThread) - - application_generate_entity = _generate_entity(app, workflow, end_user) - graph_runtime_state = _runtime_state() - workflow_execution_repository, workflow_node_execution_repository = _repositories( - sqlite_session_factory, app, end_user - ) - - flask_app = Flask(__name__) - with flask_app.app_context(): - result = generator._generate( - app_model=app, - workflow=workflow, - user=end_user, - application_generate_entity=application_generate_entity, - invoke_from=InvokeFrom.SERVICE_API, - workflow_execution_repository=workflow_execution_repository, - workflow_node_execution_repository=workflow_node_execution_repository, - streaming=True, - graph_engine_layers=("base-layer",), - graph_runtime_state=graph_runtime_state, - pause_state_config=PauseStateLayerConfig( - session_factory=sqlite_session_factory, - state_owner_user_id="owner", - ), - ) - - assert result == "converted" - graph_engine_layers = worker_kwargs["kwargs"]["graph_engine_layers"] - assert graph_engine_layers == ("base-layer",) - assert worker_kwargs["kwargs"]["pause_state_config"].state_owner_user_id == "owner" - assert worker_kwargs["kwargs"]["graph_runtime_state"] is graph_runtime_state - assert worker_kwargs["joined"] is True - assert worker_kwargs["join_timeout"] == 300 - assert draft_factory_tenant_ids == [app.tenant_id] - - -def test_resume_path_runs_worker_with_runtime_state( - monkeypatch: pytest.MonkeyPatch, - sqlite_session: Session, - sqlite_session_factory: sessionmaker[Session], -) -> None: - generator = WorkflowAppGenerator(execution_driver=WorkflowRunAgg.run) - app, workflow, end_user = _persist_generator_rows(sqlite_session) - workflow_run = WorkflowRun( - id="run", - tenant_id=workflow.tenant_id, - app_id=app.id, - workflow_id=workflow.id, - type=workflow.type, - triggered_from=WorkflowRunTriggeredFrom.APP_RUN, - version=workflow.version, - graph=workflow.graph, - inputs="{}", - status=WorkflowExecutionStatus.RUNNING, - created_by_role=CreatorUserRole.END_USER, - created_by=end_user.id, - ) - sqlite_session.add(workflow_run) - sqlite_session.commit() - runtime_state = _runtime_state() - - queue_manager = MagicMock() - monkeypatch.setattr( - "core.app.apps.workflow.app_generator.WorkflowAppQueueManager", - MagicMock(return_value=queue_manager), - ) - - monkeypatch.setattr(generator, "_handle_response", MagicMock(return_value="raw-response")) - monkeypatch.setattr( - "core.app.apps.workflow.app_generator.WorkflowAppGenerateResponseConverter.convert", - MagicMock(side_effect=lambda response, invoke_from: response), - ) - - db_session = sqlite_session_factory() - monkeypatch.setattr(app_generator_module.db, "session", db_session) - - runner_instance = MagicMock() - runner_instance.prepare.return_value = None - - def runner_ctor(**kwargs): - assert kwargs["graph_runtime_state"] is runtime_state - assert kwargs["workflow"].id == workflow.id - assert kwargs["system_user_id"] == end_user.session_id - assert kwargs["queue_manager"] is queue_manager - return runner_instance - - monkeypatch.setattr( - "core.app.apps.workflow.app_generator.WorkflowAppRunner", - MagicMock(side_effect=runner_ctor), - ) - - worker_lifecycle: dict[str, object] = {} - - class ImmediateThread: - def __init__(self, target, kwargs): - target(**kwargs) - - def start(self): - return None - - def join(self, timeout): - worker_lifecycle["joined"] = True - worker_lifecycle["join_timeout"] = timeout - - def is_alive(self): - return False - - monkeypatch.setattr("core.app.apps.workflow.app_generator.threading.Thread", ImmediateThread) - - pause_config = PauseStateLayerConfig( - session_factory=sqlite_session_factory, - state_owner_user_id="owner", - ) - - application_generate_entity = _generate_entity(app, workflow, end_user) - workflow_execution_repository, workflow_node_execution_repository = _repositories( - sqlite_session_factory, app, end_user - ) - - result = generator.resume( - app_model=app, - workflow=workflow, - user=end_user, - application_generate_entity=application_generate_entity, - graph_runtime_state=runtime_state, - workflow_execution_repository=workflow_execution_repository, - workflow_node_execution_repository=workflow_node_execution_repository, - pause_state_config=pause_config, - ) - - assert result == "raw-response" - assert worker_lifecycle["joined"] is True - assert worker_lifecycle["join_timeout"] == 300 - runner_instance.prepare.assert_called_once() diff --git a/api/tests/unit_tests/core/app/apps/test_workflow_app_runner.py b/api/tests/unit_tests/core/app/apps/test_workflow_app_runner.py new file mode 100644 index 00000000000..9a565925dc0 --- /dev/null +++ b/api/tests/unit_tests/core/app/apps/test_workflow_app_runner.py @@ -0,0 +1,1737 @@ +from __future__ import annotations + +import json +from collections.abc import Generator, Mapping +from copy import deepcopy +from datetime import UTC, datetime, timedelta +from time import perf_counter +from types import SimpleNamespace +from typing import override +from unittest.mock import MagicMock, PropertyMock, patch +from urllib.parse import parse_qs, urlparse + +import pytest +from pydantic import ValidationError +from pytest_mock import MockerFixture +from sqlalchemy import select +from sqlalchemy.orm import Session + +from configs import dify_config +from core.app.app_config.entities import WorkflowUIBasedAppConfig +from core.app.apps.base_app_queue_manager import PublishFrom +from core.app.apps.common import workflow_response_converter +from core.app.apps.execution_coordinator import AppExecutionState +from core.app.apps.workflow.app_queue_manager import WorkflowAppQueueManager +from core.app.apps.workflow.app_runner import WorkflowAppRunner +from core.app.apps.workflow_app_runner import PreparedWorkflowRun, WorkflowBasedAppRunner +from core.app.entities.app_invoke_entities import DIFY_RUN_CONTEXT_KEY, InvokeFrom, UserFrom, WorkflowAppGenerateEntity +from core.app.entities.queue_entities import ( + AppQueueEvent, + QueueAgentLogEvent, + QueueHumanInputFormFilledEvent, + QueueHumanInputFormTimeoutEvent, + QueueIterationCompletedEvent, + QueueIterationStartEvent, + QueueLoopCompletedEvent, + QueueLoopStartEvent, + QueueNodeExceptionEvent, + QueueNodeFailedEvent, + QueueNodeRetryEvent, + QueueNodeStartedEvent, + QueueNodeSucceededEvent, + QueueReasoningChunkEvent, + QueueStopEvent, + QueueTextChunkEvent, + QueueWorkflowPausedEvent, + QueueWorkflowStartedEvent, + QueueWorkflowSucceededEvent, +) +from core.credit_usage import CreditUsageAppType +from core.db.session_factory import session_factory +from core.repositories.human_input_repository import HumanInputFormRepositoryImpl, HumanInputFormSubmissionRepository +from core.repositories.sqlalchemy_workflow_execution_repository import SQLAlchemyWorkflowExecutionRepository +from core.repositories.sqlalchemy_workflow_node_execution_repository import SQLAlchemyWorkflowNodeExecutionRepository +from core.workflow.node_runtime import DifyFileReferenceFactory +from core.workflow.nodes.agent.events import NodeRunAgentLogEvent +from core.workflow.nodes.human_input.boundary import human_input_container_selector +from core.workflow.nodes.human_input.callback import DifyHITLCallback +from core.workflow.nodes.human_input.entities import ( + FileInputConfig, + FileListInputConfig, + FormDefinition, + HumanInputNodeData, + ParagraphInputConfig, + UserActionConfig, +) +from core.workflow.nodes.human_input.enums import HumanInputFormKind, HumanInputFormStatus +from core.workflow.nodes.human_input.pause_reason import HumanInputRequired +from core.workflow.system_variables import default_system_variables +from core.workflow.workflow_entry import WorkflowEntry, iter_dify_graph_engine_events +from extensions.storage.storage_type import StorageType +from graphon.engine.layer import Layer +from graphon.engine.ready_queue import StartTask +from graphon.engine_events import ( + EngineEvent, + GraphRunAbortedEvent, + GraphRunFailedEvent, + GraphRunPausedEvent, + GraphRunStartedEvent, + GraphRunSucceededEvent, + NodeRunExceptionEvent, + NodeRunFailedEvent, + NodeRunIterationStartedEvent, + NodeRunIterationSucceededEvent, + NodeRunLoopFailedEvent, + NodeRunLoopStartedEvent, + NodeRunLoopSucceededEvent, + NodeRunReasoningChunkEvent, + NodeRunRetryEvent, + NodeRunStartedEvent, + NodeRunStreamChunkEvent, + NodeRunSucceededEvent, +) +from graphon.entities import WorkflowStartReason +from graphon.entities.pause_reason import HitlRequired +from graphon.enums import BuiltinNodeTypes, WorkflowNodeExecutionMetadataKey, WorkflowType +from graphon.file.helpers import verify_file_signature +from graphon.graph import Graph +from graphon.node_events import NodeRunResult +from graphon.nodes.human_input.entities import HumanInputNodeData as GraphonHumanInputNodeData +from graphon.nodes.human_input.human_input_node import HumanInputNode +from graphon.runtime import RuntimeState, VariablePool +from graphon.runtime.execution import ROOT_FRAME_ID +from graphon.variables.variables import StringVariable +from models import Account, UploadFile +from models.enums import CreatorUserRole, WorkflowRunTriggeredFrom +from models.execution_extra_content import HumanInputContent +from models.human_input import HumanInputForm +from models.model import AppMode +from models.workflow import Workflow, WorkflowKind, WorkflowNodeExecutionTriggeredFrom +from services.workflow_run_agg import WorkflowRunAgg +from services.workflow_run_index import WorkflowRunIndex +from tests.helpers.unit.model_factories import make_workflow +from tests.unit_tests.core.app.apps.advanced_chat.test_generate_task_pipeline import ( + _build_generate_entity, + _build_pipeline, +) +from tests.unit_tests.core.app.apps.common.test_workflow_response_converter import ( + _workflow_response_converter_human_input_build_converter as _build_converter, +) +from tests.workflow_test_utils import build_test_graph_init_params, build_test_run_context + + +def _workflow(graph: dict[str, object] | None = None) -> Workflow: + return make_workflow( + tenant_id="tenant", + app_id="app", + graph=json.dumps(graph or {}, default=str), + created_by="account", + ) + + +def _workflow_entry(runner: WorkflowBasedAppRunner, state: RuntimeState | None = None) -> WorkflowEntry: + state = state or RuntimeState(workflow_id="workflow", variable_pool=VariablePool(), start_at=0.0) + graph = runner._init_graph( + graph_config={ + "nodes": [{"id": "start", "data": {"type": "start", "title": "Start", "variables": []}}], + "edges": [], + }, + graph_runtime_state=state, + user_from=UserFrom.ACCOUNT, + invoke_from=InvokeFrom.DEBUGGER, + call_depth=0, + ) + return WorkflowEntry(graph_config=graph.graph_config, graph=graph, call_depth=0, graph_runtime_state=state) + + +def _nested_container_graph(ownership: str) -> dict: + nodes = [ + {"id": "start", "data": {"type": "start", "title": "Start", "variables": []}}, + { + "id": "outer", + "data": { + "type": "loop", + "title": "Outer loop", + "start_node_id": "outer-start", + "loop_count": 1, + "break_conditions": [], + "logical_operator": "and", + }, + }, + {"id": "outer-start", "data": {"type": "loop-start", "title": "Loop start"}}, + { + "id": "inner", + "data": { + "type": "iteration", + "title": "Inner iteration", + "start_node_id": "inner-start", + "iterator_selector": ["start", "items"], + "output_selector": ["inner", "item"], + }, + }, + {"id": "inner-start", "data": {"type": "iteration-start", "title": "Iteration start"}}, + { + "id": "deep", + "data": { + "type": "loop", + "title": "Nested loop", + "start_node_id": "deep-start", + "loop_count": 1, + "break_conditions": [], + "logical_operator": "and", + }, + }, + {"id": "deep-start", "data": {"type": "loop-start", "title": "Nested loop start"}}, + {"id": "end", "data": {"type": "end", "title": "Output", "outputs": []}}, + ] + owners = {"outer-start": "outer", "inner": "outer", "inner-start": "inner", "deep": "inner", "deep-start": "deep"} + for node in nodes: + owner = owners.get(node["id"], "") + if ownership == "canonical": + node["data"]["container_id"] = owner + # Canonical ownership, including the empty root owner, takes precedence. + node["parentId"] = "stale-owner" + node["data"]["loop_id"] = "stale-owner" + elif ownership == "parent": + if owner: + node["parentId"] = owner + elif owner: + node["data"]["loop_id"] = "deep" if owner == "deep" else "outer" + if owner in {"inner", "deep"}: + node["data"]["iteration_id"] = "inner" + return { + "nodes": nodes, + "edges": [ + {"id": "start-outer", "source": "start", "target": "outer"}, + {"id": "outer-end", "source": "outer", "target": "end"}, + {"id": "outer-start-inner", "source": "outer-start", "target": "inner"}, + {"id": "inner-start-deep", "source": "inner-start", "target": "deep"}, + ], + } + + +class TestWorkflowAppRunnerCoreCases: + class TestWorkflowBasedAppRunner: + @pytest.mark.parametrize( + ("event_class", "node_type"), + [ + (NodeRunIterationStartedEvent, BuiltinNodeTypes.ITERATION), + (NodeRunIterationSucceededEvent, BuiltinNodeTypes.ITERATION), + (NodeRunLoopStartedEvent, BuiltinNodeTypes.LOOP), + (NodeRunLoopSucceededEvent, BuiltinNodeTypes.LOOP), + ], + ) + def test_container_queue_events_preserve_lifecycle_and_ancestor_metadata( + self, + event_class: type[ + NodeRunIterationStartedEvent + | NodeRunIterationSucceededEvent + | NodeRunLoopStartedEvent + | NodeRunLoopSucceededEvent + ], + node_type: str, + ) -> None: + queue = WorkflowAppQueueManager( + task_id="metadata-contract", user_id="account", invoke_from=InvokeFrom.DEBUGGER, app_mode="workflow" + ) + runner = WorkflowBasedAppRunner(queue_manager=queue, app_id="app") + entry = _workflow_entry(runner) + event = event_class( + id="container-execution", + node_id="container", + node_type=node_type, + node_title="Container", + start_at=datetime.now(UTC), + metadata={"iteration_length": 2, "total_tokens": 1}, + node_run_result=NodeRunResult( + metadata={ + WorkflowNodeExecutionMetadataKey.LOOP_ID: "outer", + WorkflowNodeExecutionMetadataKey.LOOP_INDEX: 3, + WorkflowNodeExecutionMetadataKey.TOTAL_TOKENS: 2, + } + ), + ) + original_event = event.model_dump() + + runner.handle_event(entry, event) + queue.stop_listen(execution_state=AppExecutionState.TERMINAL) + messages = list(queue.listen()) + + assert len(messages) == 1 + queued_event = messages[0].event + assert isinstance( + queued_event, + QueueIterationStartEvent | QueueIterationCompletedEvent | QueueLoopStartEvent | QueueLoopCompletedEvent, + ) + assert queued_event.metadata == { + "iteration_length": 2, + "total_tokens": 2, + "loop_id": "outer", + "loop_index": 3, + } + assert event.model_dump() == original_event + + def test_resolve_user_from(self): + runner = WorkflowBasedAppRunner(queue_manager=SimpleNamespace(), app_id="app") + + assert runner._resolve_user_from(InvokeFrom.EXPLORE) == UserFrom.ACCOUNT + assert runner._resolve_user_from(InvokeFrom.DEBUGGER) == UserFrom.ACCOUNT + assert runner._resolve_user_from(InvokeFrom.WEB_APP) == UserFrom.END_USER + + def test_init_graph_validates_graph_structure(self): + runner = WorkflowBasedAppRunner(queue_manager=SimpleNamespace(), app_id="app") + + runtime_state = RuntimeState( + workflow_id="test-workflow", + variable_pool=VariablePool.from_bootstrap(system_variables=default_system_variables()), + start_at=0.0, + ) + + with pytest.raises(ValidationError, match="nodes"): + runner._init_graph( + graph_config={}, + graph_runtime_state=runtime_state, + user_from=UserFrom.ACCOUNT, + invoke_from=InvokeFrom.DEBUGGER, + call_depth=0, + ) + + with pytest.raises(ValidationError, match="nodes"): + runner._init_graph( + graph_config={"nodes": {}, "edges": []}, + graph_runtime_state=runtime_state, + user_from=UserFrom.ACCOUNT, + invoke_from=InvokeFrom.DEBUGGER, + call_depth=0, + ) + + with pytest.raises(ValidationError, match="edges"): + runner._init_graph( + graph_config={"nodes": [], "edges": {}}, + graph_runtime_state=runtime_state, + user_from=UserFrom.ACCOUNT, + invoke_from=InvokeFrom.DEBUGGER, + call_depth=0, + ) + + def test_init_graph_includes_trace_session_id_in_run_context(self, monkeypatch: pytest.MonkeyPatch): + runner = WorkflowBasedAppRunner(queue_manager=SimpleNamespace(), app_id="app") + runtime_state = RuntimeState( + workflow_id="test-workflow", + variable_pool=VariablePool.from_bootstrap(system_variables=default_system_variables()), + start_at=0.0, + ) + captured = {} + + def fake_from_graph_init_context(**kwargs): + graph_init_context = kwargs["graph_init_context"] + captured["run_context"] = graph_init_context.run_context + captured["call_depth"] = graph_init_context.call_depth + captured["runtime_state"] = kwargs["runtime_state"] + return SimpleNamespace() + + monkeypatch.setattr( + "core.app.apps.workflow_app_runner.DifyNodeFactory.from_graph_init_context", + fake_from_graph_init_context, + ) + monkeypatch.setattr("core.app.apps.workflow_app_runner.Graph.init", lambda **_kwargs: SimpleNamespace()) + + runner._init_graph( + graph_config={"nodes": [], "edges": []}, + graph_runtime_state=runtime_state, + user_from=UserFrom.ACCOUNT, + invoke_from=InvokeFrom.DEBUGGER, + root_node_id="root", + trace_session_id="session-1", + call_depth=3, + ) + + assert captured["run_context"][DIFY_RUN_CONTEXT_KEY].trace_session_id == "session-1" + assert captured["call_depth"] == 3 + assert captured["runtime_state"] is runtime_state + + @pytest.mark.parametrize("ownership", ["canonical", "parent", "legacy"]) + def test_init_graph_preserves_container_execution_scopes(self, ownership: str) -> None: + runner = WorkflowBasedAppRunner(queue_manager=SimpleNamespace(), app_id="app") + runtime_state = RuntimeState( + workflow_id="test-workflow", + variable_pool=VariablePool.from_bootstrap(system_variables=default_system_variables()), + start_at=0.0, + ) + graph_config = _nested_container_graph(ownership) + original_config = deepcopy(graph_config) + graph = runner._init_graph( + graph_config=graph_config, + graph_runtime_state=runtime_state, + user_from=UserFrom.ACCOUNT, + invoke_from=InvokeFrom.DEBUGGER, + root_node_id="start", + call_depth=0, + ) + + assert set(graph.nodes) == {"start", "outer", "end"} + outer_graph = Graph.init( + graph_config=graph.graph_config, + node_factory=graph.node_factory, + root_node_id="outer-start", + container_id="outer", + ) + assert set(outer_graph.nodes) == {"outer-start", "inner"} + inner_graph = Graph.init( + graph_config=outer_graph.graph_config, + node_factory=outer_graph.node_factory, + root_node_id="inner-start", + container_id="inner", + ) + assert set(inner_graph.nodes) == {"inner-start", "deep"} + deep_graph = Graph.init( + graph_config=inner_graph.graph_config, + node_factory=inner_graph.node_factory, + root_node_id="deep-start", + container_id="deep", + ) + assert set(deep_graph.nodes) == {"deep-start"} + assert graph_config == original_config + + def test_prepare_single_node_execution_requires_run(self): + runner = WorkflowBasedAppRunner(queue_manager=SimpleNamespace(), app_id="app") + + workflow = _workflow() + + with pytest.raises(ValueError, match="Neither single_iteration_run nor single_loop_run"): + runner._prepare_single_node_execution( + workflow, None, None, user_id="00000000-0000-0000-0000-000000000001" + ) + + @pytest.mark.parametrize("ownership", ["canonical", "parent", "legacy"]) + @pytest.mark.parametrize("node_id", ["outer", "inner", "deep"]) + def test_single_container_debugging_preserves_descendants(self, ownership: str, node_id: str): + runner = WorkflowBasedAppRunner(queue_manager=SimpleNamespace(), app_id="app") + graph_runtime_state = RuntimeState( + workflow_id="test-workflow", + variable_pool=VariablePool.from_bootstrap(system_variables=default_system_variables()), + start_at=0.0, + ) + + graph_config = _nested_container_graph(ownership) + workflow = _workflow(graph_config) + workflow.id = "workflow" + graph_runtime_state.variable_pool.add(["start", "items"], ["first", "second"]) + + graph, variable_pool = runner._get_graph_and_variable_pool_for_single_node_run( + workflow=workflow, + node_id=node_id, + user_inputs={}, + graph_runtime_state=graph_runtime_state, + node_type_label="iteration", + user_id="00000000-0000-0000-0000-000000000001", + ) + + assert set(graph.nodes) == {node_id} + assert variable_pool is graph_runtime_state.variable_pool + deep_graph = Graph.init( + graph_config=graph.graph_config, + node_factory=graph.node_factory, + root_node_id="deep-start", + container_id="deep", + ) + assert set(deep_graph.nodes) == {"deep-start"} + + def test_get_graph_and_variable_pool_for_single_node_run_does_not_mutate_inputs( + self, monkeypatch: pytest.MonkeyPatch + ): + runner = WorkflowBasedAppRunner(queue_manager=SimpleNamespace(), app_id="app") + graph_runtime_state = RuntimeState( + workflow_id="workflow", + variable_pool=VariablePool.from_bootstrap(system_variables=default_system_variables()), + start_at=0.0, + ) + graph_config = { + "nodes": [ + {"id": "node-1", "data": {"type": "start", "version": "1"}}, + {"id": "node-child", "data": {"type": "start", "version": "1", "iteration_id": "node-1"}}, + {"id": "outside", "data": {"type": "start", "version": "1"}}, + ], + "edges": [ + {"id": "edge-1", "source": "node-1", "target": "node-child"}, + {"id": "edge-2", "source": "outside", "target": "node-1"}, + ], + } + expected_graph_config = deepcopy(graph_config) + workflow = SimpleNamespace(tenant_id="tenant", id="workflow", graph_dict=graph_config) + user_inputs = {"question": "hello"} + expected_user_inputs = dict(user_inputs) + captured_graph_config = {} + + def fake_graph_init(**kwargs): + captured_graph_config["value"] = kwargs["graph_config"] + return SimpleNamespace() + + class _NodeCls: + @staticmethod + def extract_variable_selector_to_variable_mapping(graph_config, config): + return {} + + from core.app.apps import workflow_app_runner + + monkeypatch.setattr("core.app.apps.workflow_app_runner.Graph.init", fake_graph_init) + monkeypatch.setattr(workflow_app_runner, "resolve_workflow_node_class", lambda **_kwargs: _NodeCls) + monkeypatch.setattr( + "core.app.apps.workflow_app_runner.load_into_variable_pool", + lambda **kwargs: kwargs["user_inputs"].update({"from_loader": "local-only"}), + ) + monkeypatch.setattr( + "core.app.apps.workflow_app_runner.WorkflowEntry.mapping_user_inputs_to_variable_pool", + lambda **kwargs: kwargs["user_inputs"].update({"from_mapping": "local-only"}), + ) + + graph, variable_pool = runner._get_graph_and_variable_pool_for_single_node_run( + workflow=workflow, + node_id="node-1", + user_inputs=user_inputs, + graph_runtime_state=graph_runtime_state, + node_type_label="iteration", + user_id="00000000-0000-0000-0000-000000000001", + ) + + assert graph is not None + assert variable_pool is graph_runtime_state.variable_pool + assert workflow.graph_dict == expected_graph_config + assert user_inputs == expected_user_inputs + assert captured_graph_config["value"] is not workflow.graph_dict + assert [node["id"] for node in captured_graph_config["value"]["nodes"]] == ["node-1", "node-child"] + assert captured_graph_config["value"]["edges"] == [ + {"id": "edge-1", "source": "node-1", "target": "node-child"} + ] + + def test_get_graph_and_variable_pool_for_single_node_run_includes_trace_session_id( + self, monkeypatch: pytest.MonkeyPatch + ): + runner = WorkflowBasedAppRunner(queue_manager=SimpleNamespace(), app_id="app") + graph_runtime_state = RuntimeState( + workflow_id="test-workflow", + variable_pool=VariablePool.from_bootstrap(system_variables=default_system_variables()), + start_at=0.0, + ) + graph_config = { + "nodes": [{"id": "node-1", "data": {"type": "start", "version": "1"}}], + "edges": [], + } + workflow = _workflow(graph_config) + workflow.id = "workflow" + captured = {} + + def fake_from_graph_init_context(**kwargs): + captured["run_context"] = kwargs["graph_init_context"].run_context + captured["runtime_state"] = kwargs["runtime_state"] + return SimpleNamespace() + + class _NodeCls: + @staticmethod + def extract_variable_selector_to_variable_mapping(graph_config, config): + return {} + + from core.app.apps import workflow_app_runner + + monkeypatch.setattr( + "core.app.apps.workflow_app_runner.DifyNodeFactory.from_graph_init_context", + fake_from_graph_init_context, + ) + monkeypatch.setattr("core.app.apps.workflow_app_runner.Graph.init", lambda **kwargs: SimpleNamespace()) + monkeypatch.setattr(workflow_app_runner, "resolve_workflow_node_class", lambda **_kwargs: _NodeCls) + monkeypatch.setattr("core.app.apps.workflow_app_runner.load_into_variable_pool", lambda **kwargs: None) + monkeypatch.setattr( + "core.app.apps.workflow_app_runner.WorkflowEntry.mapping_user_inputs_to_variable_pool", + lambda **kwargs: None, + ) + + runner._get_graph_and_variable_pool_for_single_node_run( + workflow=workflow, + node_id="node-1", + user_inputs={}, + graph_runtime_state=graph_runtime_state, + node_type_label="iteration", + user_id="00000000-0000-0000-0000-000000000001", + trace_session_id="session-1", + ) + + assert captured["run_context"][DIFY_RUN_CONTEXT_KEY].trace_session_id == "session-1" + assert captured["runtime_state"] is graph_runtime_state + + def test_get_graph_and_variable_pool_preloads_constructor_variables_before_graph_init( + self, monkeypatch: pytest.MonkeyPatch + ): + variable_loader = SimpleNamespace( + load_variables=lambda selectors: ( + [ + StringVariable( + name="conversation_id", + value="conv-1", + selector=["sys", "conversation_id"], + ) + ] + if selectors + else [] + ) + ) + runner = WorkflowBasedAppRunner( + queue_manager=SimpleNamespace(), + variable_loader=variable_loader, + app_id="app", + ) + graph_runtime_state = RuntimeState( + workflow_id="test-workflow", + variable_pool=VariablePool.from_bootstrap(system_variables=default_system_variables()), + start_at=0.0, + ) + + workflow = _workflow( + { + "nodes": [ + {"id": "loop-node", "data": {"type": "loop", "version": "1", "title": "Loop"}}, + { + "id": "llm-child", + "data": { + "type": "llm", + "version": "1", + "loop_id": "loop-node", + "memory": object(), + }, + }, + ], + "edges": [], + } + ) + workflow.id = "workflow" + + class _LoopNodeCls: + @staticmethod + def extract_variable_selector_to_variable_mapping(graph_config, config): + return {} + + def _graph_init(**kwargs): + variable_pool = graph_runtime_state.variable_pool + assert variable_pool.get(["sys", "conversation_id"]) is not None + return SimpleNamespace() + + monkeypatch.setattr( + "core.app.apps.workflow_app_runner.Graph.init", + _graph_init, + ) + monkeypatch.setattr( + "core.app.apps.workflow_app_runner.resolve_workflow_node_class", + lambda **_kwargs: _LoopNodeCls, + ) + monkeypatch.setattr( + "core.app.apps.workflow_app_runner.load_into_variable_pool", + lambda **kwargs: None, + ) + monkeypatch.setattr( + "core.app.apps.workflow_app_runner.WorkflowEntry.mapping_user_inputs_to_variable_pool", + lambda **kwargs: None, + ) + + graph, variable_pool = runner._get_graph_and_variable_pool_for_single_node_run( + workflow=workflow, + node_id="loop-node", + user_inputs={}, + graph_runtime_state=graph_runtime_state, + node_type_label="loop", + ) + + assert graph is not None + assert variable_pool.get(["sys", "conversation_id"]).value == "conv-1" + + def test_handle_graph_run_events_and_pause_notifications(self, monkeypatch: pytest.MonkeyPatch): + published: list[object] = [] + + class _QueueManager: + def publish(self, event, publish_from): + published.append((event, publish_from)) + + runner = WorkflowBasedAppRunner(queue_manager=_QueueManager(), app_id="app") + graph_runtime_state = RuntimeState( + workflow_id="test-workflow", + variable_pool=VariablePool.from_bootstrap(system_variables=default_system_variables()), + start_at=0.0, + ) + workflow_entry = _workflow_entry(runner, graph_runtime_state) + + emails: list[dict] = [] + + class _Dispatch: + def apply_async(self, *, kwargs, queue): + emails.append({"kwargs": kwargs, "queue": queue}) + + monkeypatch.setattr( + "core.app.apps.workflow_app_runner.dispatch_human_input_email_task", + _Dispatch(), + ) + enriched_reason = HumanInputRequired( + form_id="form", + form_content="content", + node_id="node-1", + node_title="Node", + ) + + reason = HitlRequired( + session_id="form", + node_id="node-1", + node_title="Node", + ) + + runner.handle_event(workflow_entry, GraphRunStartedEvent()) + runner.handle_event(workflow_entry, GraphRunSucceededEvent(outputs={"ok": True})) + runner.handle_event( + workflow_entry, GraphRunPausedEvent(reasons=[reason], outputs={}), pause_reasons=[enriched_reason] + ) + + assert any(isinstance(event, QueueWorkflowStartedEvent) for event, _ in published) + assert any(isinstance(event, QueueWorkflowSucceededEvent) for event, _ in published) + paused_event = next(event for event, _ in published if isinstance(event, QueueWorkflowPausedEvent)) + assert paused_event.paused_nodes == ["node-1"] + assert emails + + def test_handle_graph_aborted_publishes_stopped_terminal(self): + published: list[object] = [] + + class _QueueManager: + def publish(self, event, publish_from): + del publish_from + published.append(event) + + runner = WorkflowBasedAppRunner(queue_manager=_QueueManager(), app_id="app") + workflow_entry = _workflow_entry(runner) + + runner.handle_event(workflow_entry, GraphRunAbortedEvent(reason="User requested stop", outputs={})) + + event = published[-1] + assert isinstance(event, QueueStopEvent) + assert event.get_stop_reason() == "User requested stop" + + def test_handle_node_events_publishes_queue_events(self): + published: list[object] = [] + + class _QueueManager: + def publish(self, event, publish_from): + published.append(event) + + runner = WorkflowBasedAppRunner(queue_manager=_QueueManager(), app_id="app") + graph_runtime_state = RuntimeState( + workflow_id="test-workflow", + variable_pool=VariablePool.from_bootstrap(system_variables=default_system_variables()), + start_at=0.0, + ) + workflow_entry = _workflow_entry(runner, graph_runtime_state) + + runner.handle_event( + workflow_entry, + NodeRunStartedEvent( + id="exec", + node_id="node", + node_type=BuiltinNodeTypes.START, + node_title="Start", + start_at=datetime.now(UTC), + ), + ) + runner.handle_event( + workflow_entry, + NodeRunStreamChunkEvent( + id="exec", + node_id="node", + node_type=BuiltinNodeTypes.START, + selector=["node", "text"], + chunk="hi", + is_final=False, + ), + ) + runner.handle_event( + workflow_entry, + NodeRunReasoningChunkEvent( + id="exec", + node_id="node", + node_type=BuiltinNodeTypes.LLM, + selector=["node", "reasoning_content"], + chunk="thinking", + is_final=False, + ), + ) + runner.handle_event( + workflow_entry, + NodeRunAgentLogEvent( + id="exec", + node_id="node", + node_type=BuiltinNodeTypes.START, + message_id="msg", + label="label", + node_execution_id="exec", + parent_id=None, + error=None, + status="done", + data={}, + metadata={}, + ), + ) + runner.handle_event( + workflow_entry, + NodeRunIterationSucceededEvent( + id="exec", + node_id="node", + node_type=BuiltinNodeTypes.LLM, + node_title="Iter", + start_at=datetime.now(UTC), + inputs={}, + outputs={"ok": True}, + metadata={}, + steps=1, + ), + ) + runner.handle_event( + workflow_entry, + NodeRunLoopFailedEvent( + id="exec", + node_id="node", + node_type=BuiltinNodeTypes.LLM, + node_title="Loop", + start_at=datetime.now(UTC), + inputs={}, + outputs={}, + metadata={}, + steps=1, + error="boom", + ), + ) + + assert any(isinstance(event, QueueTextChunkEvent) for event in published) + assert any(isinstance(event, QueueReasoningChunkEvent) for event in published) + assert any(isinstance(event, QueueAgentLogEvent) for event in published) + assert any(isinstance(event, QueueIterationCompletedEvent) for event in published) + assert any(isinstance(event, QueueLoopCompletedEvent) for event in published) + + @pytest.mark.parametrize("ownership", ["canonical", "parent", "legacy"]) + def test_nested_engine_events_retain_both_container_ancestors(self, ownership: str): + from tests.unit_tests.services.test_workflow_persistence import _make_layer + + published: list[object] = [] + + class _QueueManager: + def publish(self, event, publish_from): + del publish_from + published.append(event) + + graph_config = _nested_container_graph(ownership) + graph_config["nodes"].append( + { + "id": "deep-end", + "data": {"type": "end", "title": "Nested End", "outputs": [], "container_id": "deep"}, + } + ) + graph_config["edges"].append({"id": "deep-output", "source": "deep-start", "target": "deep-end"}) + layer, _, node_repo, state = _make_layer(graph_data=graph_config) + persisted_starts: dict[str, dict] = {} + + class _StartMetadataLayer(Layer): + @override + def on_event(self, event): + if isinstance(event, NodeRunStartedEvent): + persisted_starts[event.node_id] = dict(node_repo.saved[-1].metadata) + + state.variable_pool.add(["start", "items"], ["item"]) + runner = WorkflowBasedAppRunner(queue_manager=_QueueManager(), app_id="app") + graph = runner._init_graph( + graph_config=graph_config, + graph_runtime_state=state, + user_from=UserFrom.ACCOUNT, + invoke_from=InvokeFrom.DEBUGGER, + call_depth=0, + ) + entry = WorkflowEntry(graph_config=graph.graph_config, graph=graph, call_depth=0, graph_runtime_state=state) + engine = entry.graph_engine + index = WorkflowRunIndex(()) + layer.set_node_run_indices(index.indices) + engine.add_layer(index) + engine.add_layer(layer) + engine.add_layer(_StartMetadataLayer()) + for event in iter_dify_graph_engine_events(engine): + runner.handle_event(entry, event) + + expected = { + "outer": (None, None), + "inner": (None, "outer"), + "deep": ("inner", "outer"), + "deep-end": ("inner", "deep"), + } + started = {event.node_id: event for event in published if isinstance(event, QueueNodeStartedEvent)} + persisted = {event.node_id: event for event in node_repo.saved} + for node_id, (iteration_id, loop_id) in expected.items(): + assert (started[node_id].in_iteration_id, started[node_id].in_loop_id) == (iteration_id, loop_id) + for metadata in (persisted_starts[node_id], persisted[node_id].metadata): + assert (metadata.get("iteration_id"), metadata.get("loop_id")) == (iteration_id, loop_id) + for metadata in (persisted_starts["deep-end"], persisted["deep-end"].metadata): + assert metadata["iteration_index"] == 0 + assert metadata["loop_index"] == 0 + for event in published: + if isinstance( + event, + QueueIterationStartEvent + | QueueIterationCompletedEvent + | QueueLoopStartEvent + | QueueLoopCompletedEvent, + ): + assert (event.metadata.get("iteration_id"), event.metadata.get("loop_id")) == expected[ + event.node_id + ] + + def test_handle_nested_answer_success_publishes_text_chunk(self): + published: list[object] = [] + + class _QueueManager: + def publish(self, event, publish_from): + published.append(event) + + runner = WorkflowBasedAppRunner(queue_manager=_QueueManager(), app_id="app") + workflow_entry = _workflow_entry(runner) + + runner.handle_event( + workflow_entry, + NodeRunSucceededEvent( + id="exec", + node_id="answer-in-iteration", + node_type=BuiltinNodeTypes.ANSWER, + start_at=datetime.now(UTC), + finished_at=datetime.now(UTC), + node_run_result=NodeRunResult( + outputs={"answer": "inside iteration"}, metadata={"iteration_id": "iteration"} + ), + container_id="iteration", + ), + ) + + text_chunk = published[0] + assert isinstance(text_chunk, QueueTextChunkEvent) + assert text_chunk.text == "inside iteration" + assert text_chunk.from_variable_selector == ["answer-in-iteration", "answer"] + assert text_chunk.in_iteration_id == "iteration" + assert any(isinstance(event, QueueNodeSucceededEvent) for event in published) + + def test_handle_root_answer_success_does_not_publish_extra_text_chunk(self): + published: list[object] = [] + + class _QueueManager: + def publish(self, event, publish_from): + published.append(event) + + runner = WorkflowBasedAppRunner(queue_manager=_QueueManager(), app_id="app") + workflow_entry = _workflow_entry(runner) + + runner.handle_event( + workflow_entry, + NodeRunSucceededEvent( + id="exec", + node_id="answer", + node_type=BuiltinNodeTypes.ANSWER, + start_at=datetime.now(UTC), + finished_at=datetime.now(UTC), + node_run_result=NodeRunResult(outputs={"answer": "root answer"}), + ), + ) + + assert not any(isinstance(event, QueueTextChunkEvent) for event in published) + assert len(published) == 1 + assert isinstance(published[0], QueueNodeSucceededEvent) + + @pytest.mark.parametrize( + ("event_factory", "queue_event_cls"), + [ + ( + lambda result, start_at, finished_at: NodeRunSucceededEvent( + id="exec", + node_id="node", + node_type=BuiltinNodeTypes.START, + start_at=start_at, + finished_at=finished_at, + node_run_result=result, + ), + QueueNodeSucceededEvent, + ), + ( + lambda result, start_at, finished_at: NodeRunFailedEvent( + id="exec", + node_id="node", + node_type=BuiltinNodeTypes.START, + start_at=start_at, + finished_at=finished_at, + error="boom", + node_run_result=result, + ), + QueueNodeFailedEvent, + ), + ( + lambda result, start_at, finished_at: NodeRunExceptionEvent( + id="exec", + node_id="node", + node_type=BuiltinNodeTypes.START, + start_at=start_at, + finished_at=finished_at, + error="boom", + node_run_result=result, + ), + QueueNodeExceptionEvent, + ), + ( + lambda result, start_at, _finished_at: NodeRunRetryEvent( + id="exec", + node_id="node", + node_type=BuiltinNodeTypes.START, + node_title="Start", + start_at=start_at, + error="boom", + retry_index=1, + node_run_result=result, + ), + QueueNodeRetryEvent, + ), + ], + ) + def test_handle_start_node_result_events_project_outputs(self, event_factory, queue_event_cls): + published: list[object] = [] + + class _QueueManager: + def publish(self, event, publish_from): + published.append(event) + + runner = WorkflowBasedAppRunner(queue_manager=_QueueManager(), app_id="app") + graph_runtime_state = RuntimeState( + workflow_id="test-workflow", + variable_pool=VariablePool.from_bootstrap(system_variables=default_system_variables()), + start_at=0.0, + ) + workflow_entry = _workflow_entry(runner, graph_runtime_state) + started_at = datetime.now(UTC) + finished_at = datetime.now(UTC) + result = NodeRunResult( + inputs={"question": "hello"}, + outputs={ + "question": "hello", + "sys.query": "hello", + "env.API_KEY": "secret", + "conversation.session_id": "session-1", + }, + ) + + runner.handle_event(workflow_entry, event_factory(result, started_at, finished_at)) + + queue_event = published[-1] + assert isinstance(queue_event, queue_event_cls) + assert queue_event.outputs == {"question": "hello"} + + +class _DummyQueueManager: + def __init__(self): + self.published = [] + + def publish(self, event, _from): + self.published.append(event) + + +class _DummyRuntimeState: + variable_pool = object() + + +class _DummyGraphEngine: + def __init__(self): + self.runtime_state = _DummyRuntimeState() + + +class _DummyWorkflowEntry: + def __init__(self): + self.graph_engine = _DummyGraphEngine() + + +class TestWorkflowAppRunnerNotificationsCases: + @staticmethod + def test_handle_pause_event_enqueues_email_task(monkeypatch: pytest.MonkeyPatch): + queue_manager = _DummyQueueManager() + runner = WorkflowBasedAppRunner(queue_manager=queue_manager, app_id="app-id") + workflow_entry = _DummyWorkflowEntry() + + graph_reason = HitlRequired(session_id="form-123", node_id="node-1", node_title="Review") + event = GraphRunPausedEvent(reasons=[graph_reason], outputs={}) + + email_task = MagicMock() + enriched_reason = HumanInputRequired( + form_id="form-123", + form_content="content", + inputs=[], + actions=[], + node_id="node-1", + node_title="Review", + ) + monkeypatch.setattr("core.app.apps.workflow_app_runner.dispatch_human_input_email_task", email_task) + + runner.handle_event(workflow_entry, event, pause_reasons=[enriched_reason]) + + email_task.apply_async.assert_called_once() + kwargs = email_task.apply_async.call_args.kwargs["kwargs"] + assert kwargs["form_id"] == "form-123" + assert kwargs["node_title"] == "Review" + + assert any(isinstance(evt, QueueWorkflowPausedEvent) for evt in queue_manager.published) + + +def _make_graph_state(): + variable_pool = VariablePool.from_bootstrap( + system_variables=default_system_variables(), + user_inputs={}, + environment_variables=[], + conversation_variables=[], + ) + return ( + MagicMock(), + variable_pool, + RuntimeState(workflow_id="test-workflow", variable_pool=variable_pool, start_at=0.0), + ) + + +class TestWorkflowAppRunnerSingleNodeCases: + @staticmethod + @pytest.fixture + def queue_manager() -> WorkflowAppQueueManager: + return WorkflowAppQueueManager( + task_id="task-id", user_id="user", invoke_from=InvokeFrom.SERVICE_API, app_mode=AppMode.WORKFLOW + ) + + @staticmethod + @pytest.mark.parametrize( + ("single_iteration_run", "single_loop_run"), + [ + (WorkflowAppGenerateEntity.SingleIterationRunEntity(node_id="iter", inputs={}), None), + (None, WorkflowAppGenerateEntity.SingleLoopRunEntity(node_id="loop", inputs={})), + ], + ) + def test_run_uses_single_node_execution_branch( + single_iteration_run: WorkflowAppGenerateEntity.SingleIterationRunEntity | None, + single_loop_run: WorkflowAppGenerateEntity.SingleLoopRunEntity | None, + queue_manager: WorkflowAppQueueManager, + ) -> None: + app_generate_entity = WorkflowAppGenerateEntity( + app_config=WorkflowUIBasedAppConfig( + app_id="app", tenant_id="tenant", workflow_id="workflow", app_mode=AppMode.WORKFLOW + ), + inputs={}, + files=[], + user_id="user", + invoke_from=InvokeFrom.SERVICE_API, + workflow_execution_id="execution-id", + task_id="task-id", + stream=True, + extras={"trace_session_id": "session-1"}, + single_iteration_run=single_iteration_run, + single_loop_run=single_loop_run, + ) + + workflow = Workflow( + tenant_id="tenant", + app_id="app", + id="workflow", + type="workflow", + version="v1", + graph=json.dumps({"nodes": [], "edges": []}), + ) + workflow.environment_variables = [] + + runner = WorkflowAppRunner( + application_generate_entity=app_generate_entity, + queue_manager=queue_manager, + variable_loader=MagicMock(), + workflow=workflow, + system_user_id="system-user", + ) + + graph, variable_pool, graph_runtime_state = _make_graph_state() + mock_workflow_entry = MagicMock() + mock_workflow_entry.graph_engine = MagicMock() + mock_workflow_entry.graph_engine.layer = MagicMock() + mock_workflow_entry.run.return_value = iter([]) + + with ( + patch("core.app.apps.workflow.app_runner.RedisChannel"), + patch("core.app.apps.workflow.app_runner.redis_client"), + patch("core.app.apps.workflow.app_runner.WorkflowEntry", return_value=mock_workflow_entry) as entry_class, + patch.object( + runner, + "_prepare_single_node_execution", + return_value=( + graph, + variable_pool, + graph_runtime_state, + ), + ) as prepare_single, + patch.object(runner, "_init_graph") as init_graph, + ): + runner.prepare() + + prepare_single.assert_called_once_with( + workflow=workflow, + single_iteration_run=single_iteration_run, + single_loop_run=single_loop_run, + user_id="user", + app_type=CreditUsageAppType.WORKFLOW, + trace_session_id="session-1", + ) + init_graph.assert_not_called() + + assert entry_class.call_args.kwargs["graph_runtime_state"] is graph_runtime_state + + @staticmethod + def test_single_node_run_validates_target_node_config(queue_manager: WorkflowAppQueueManager) -> None: + runner = WorkflowBasedAppRunner( + queue_manager=queue_manager, + variable_loader=MagicMock(), + app_id="app", + ) + + workflow = Workflow( + id="workflow", + tenant_id="tenant", + graph=json.dumps( + { + "nodes": [ + { + "id": "loop-node", + "data": {}, + } + ], + "edges": [], + } + ), + ) + + _, _, graph_runtime_state = _make_graph_state() + + with pytest.raises(ValidationError, match="nodes.0.data.type"): + runner._get_graph_and_variable_pool_for_single_node_run( + workflow=workflow, + node_id="loop-node", + user_inputs={}, + graph_runtime_state=graph_runtime_state, + node_type_label="loop", + user_id="00000000-0000-0000-0000-000000000001", + ) + + @staticmethod + def test_single_node_run_rejects_invalid_loop_count(queue_manager: WorkflowAppQueueManager) -> None: + runner = WorkflowBasedAppRunner( + queue_manager=queue_manager, + app_id="app", + ) + + workflow = Workflow( + id="workflow", + tenant_id="tenant", + graph=json.dumps( + { + "nodes": [ + { + "id": "loop-node", + "data": { + "type": "loop", + "title": "Loop", + "loop_count": 0, + "start_node_id": "loop-start", + "break_conditions": [], + "logical_operator": "and", + }, + }, + { + "id": "loop-start", + "data": {"type": "loop-start", "title": "Loop start", "container_id": "loop-node"}, + }, + ], + "edges": [], + } + ), + ) + + _, _, graph_runtime_state = _make_graph_state() + with pytest.raises(ValidationError, match="loop_count"): + runner._get_graph_and_variable_pool_for_single_node_run( + workflow=workflow, + node_id="loop-node", + user_inputs={}, + graph_runtime_state=graph_runtime_state, + node_type_label="loop", + user_id="00000000-0000-0000-0000-000000000001", + ) + + @staticmethod + def test_run_adds_inputs_with_snippet_compatible_start_aliases(queue_manager: WorkflowAppQueueManager) -> None: + app_generate_entity = WorkflowAppGenerateEntity( + app_config=WorkflowUIBasedAppConfig( + app_id="app", tenant_id="tenant", workflow_id="workflow", app_mode=AppMode.WORKFLOW + ), + inputs={"question": "hello"}, + files=[], + user_id="user", + invoke_from=InvokeFrom.SERVICE_API, + workflow_execution_id="execution-id", + task_id="task-id", + stream=True, + call_depth=4, + ) + + workflow = Workflow( + tenant_id="tenant", + app_id="app", + id="workflow", + type="workflow", + version="v1", + graph=json.dumps({"nodes": [], "edges": []}), + kind=WorkflowKind.SNIPPET, + ) + workflow.environment_variables = [] + + runner = WorkflowAppRunner( + application_generate_entity=app_generate_entity, + queue_manager=queue_manager, + variable_loader=MagicMock(), + workflow=workflow, + system_user_id="system-user", + ) + + mock_workflow_entry = MagicMock() + mock_workflow_entry.graph_engine = MagicMock() + mock_workflow_entry.graph_engine.layer = MagicMock() + mock_workflow_entry.run.return_value = iter([]) + + with ( + patch("core.app.apps.workflow.app_runner.RedisChannel"), + patch("core.app.apps.workflow.app_runner.redis_client"), + patch("core.app.apps.workflow.app_runner.WorkflowEntry", return_value=mock_workflow_entry), + patch("core.app.apps.workflow.app_runner.build_system_variables", return_value={}), + patch("core.app.apps.workflow.app_runner.build_bootstrap_variables", return_value=[]), + patch("core.app.apps.workflow.app_runner.add_variables_to_pool"), + patch("core.app.apps.workflow.app_runner.get_default_root_node_id", return_value="root-node"), + patch( + "core.app.apps.workflow.app_runner.get_compatible_start_aliases", return_value=("legacy-start",) + ) as aliases, + patch("core.app.apps.workflow.app_runner.add_node_inputs_to_pool") as add_inputs, + patch.object(runner, "_init_graph", return_value=MagicMock()) as init_graph, + ): + runner.prepare() + + aliases.assert_called_once_with(workflow_kind="snippet", root_node_id="root-node") + add_inputs.assert_called_once() + assert add_inputs.call_args.kwargs["node_id"] == "root-node" + assert add_inputs.call_args.kwargs["inputs"] == {"question": "hello"} + assert add_inputs.call_args.kwargs["aliases"] == ("legacy-start",) + assert init_graph.call_args.kwargs["call_depth"] == 4 + + +def _save_form(session: Session, *, status: HumanInputFormStatus) -> HumanInputForm: + now = datetime.now(UTC) + definition = FormDefinition( + form_content="Decision: {{#$output.answer#}}", + rendered_content="Decision: {{#$output.answer#}}", + inputs=[ParagraphInputConfig(output_variable_name="answer")], + user_actions=[UserActionConfig(id="approve", title="Approve")], + expiration_time=now + timedelta(hours=1), + node_title="Approval", + ) + form = HumanInputForm( + tenant_id="tenant", + app_id="app", + workflow_run_id="run-1", + node_id="same-human-node", + form_kind=HumanInputFormKind.RUNTIME, + form_definition=definition.model_dump_json(), + rendered_content=definition.rendered_content, + expiration_time=definition.expiration_time, + status=status, + selected_action_id="approve" if status == HumanInputFormStatus.SUBMITTED else None, + submitted_data=json.dumps({"answer": "approved"}) if status == HumanInputFormStatus.SUBMITTED else None, + submitted_at=now if status == HumanInputFormStatus.SUBMITTED else None, + created_at=now, + ) + session.add(form) + session.commit() + return form + + +def _make_paused_workflow( + runner: WorkflowBasedAppRunner, forms: list[HumanInputForm], *, human_input_form: HumanInputForm | None = None +) -> WorkflowEntry: + state = RuntimeState(workflow_id="workflow", variable_pool=VariablePool(), start_at=perf_counter()) + state.variable_pool.add(("sys", "workflow_run_id"), "run-1") + state.variable_pool.add(("sys", "app_id"), "app") + graph = runner._init_graph( + graph_config={ + "nodes": [{"id": "start", "data": {"type": "start", "title": "Start", "variables": []}}], + "edges": [], + }, + graph_runtime_state=state, + tenant_id="tenant", + user_id="user", + user_from=UserFrom.END_USER, + invoke_from=InvokeFrom.WEB_APP, + call_depth=0, + ) + if human_input_form is not None: + node_data = HumanInputNodeData( + title="Approval", + form_content="Decision: {{#$output.answer#}}", + inputs=[ParagraphInputConfig(output_variable_name="answer")], + user_actions=[UserActionConfig(id="approve", title="Approve")], + ) + callback = DifyHITLCallback( + form_repository=HumanInputFormRepositoryImpl( + tenant_id="tenant", app_id="app", workflow_execution_id="run-1" + ), + node_data=node_data, + execution_id_getter=lambda: human_input_form.id, + ) + graph = ( + Graph.new() + .add_root(graph.root_node) + .add_node( + HumanInputNode( + node_id=human_input_form.node_id, + data=GraphonHumanInputNodeData.model_validate(node_data.model_dump()), + init_params=build_test_graph_init_params(tenant_id="tenant", app_id="app"), + runtime_state=state, + hitl_callback=callback, + ), + from_node_id=graph.root_node.id, + ) + .build() + ) + # A resumed Human Input node keeps the execution identity used to create its form. + state.graph_execution.get_or_create_node_execution( + frame_id=ROOT_FRAME_ID, node_id=human_input_form.node_id + ).execution_id = human_input_form.id + state.graph_execution.start() + for form in forms: + state.graph_execution.pause(HitlRequired(session_id=form.id, node_id=form.node_id, node_title="Approval")) + state.variable_pool.add(human_input_container_selector(form.id), "outer-tool") + state.defer_ready_task(StartTask(frame_id=ROOT_FRAME_ID, node_id=graph.root_node.id)) + return WorkflowEntry( + graph=graph, + graph_config=graph.graph_config or {}, + call_depth=0, + graph_runtime_state=state, + ) + + +def _resume_events(runner: WorkflowBasedAppRunner, entry: WorkflowEntry) -> Generator[EngineEvent]: + """Resume through orchestration with actual repositories and the real Engine.""" + entity = WorkflowAppGenerateEntity( + task_id="task", + app_config=WorkflowUIBasedAppConfig( + tenant_id="tenant", + app_id="app", + app_mode=AppMode.WORKFLOW, + workflow_id="workflow", + ), + inputs={}, + files=[], + user_id="user", + stream=True, + invoke_from=InvokeFrom.DEBUGGER, + workflow_execution_id="run-1", + ) + user = Account(name="Tester", email="tester@example.com") + user.id = "user" + factory = session_factory.get_session_maker() + runs = SQLAlchemyWorkflowExecutionRepository( + factory, + tenant_id="tenant", + app_id="app", + user=user, + triggered_from=WorkflowRunTriggeredFrom.DEBUGGING, + ) + nodes = SQLAlchemyWorkflowNodeExecutionRepository( + factory, + tenant_id="tenant", + app_id="app", + user=user, + triggered_from=WorkflowNodeExecutionTriggeredFrom.WORKFLOW_RUN, + ) + prepared = PreparedWorkflowRun(entry, WorkflowType.WORKFLOW, "draft", entity, ()) + return WorkflowRunAgg( + prepared, + runner, + None, + workflow_execution_repository=runs, + workflow_node_execution_repository=nodes, + ).iter_events() + + +class TestWorkflowHumanInputCompletionCases: + @staticmethod + @pytest.fixture + def workflow_queue(mocker: MockerFixture) -> tuple[WorkflowAppQueueManager, MagicMock]: + queue = WorkflowAppQueueManager("task", "user", InvokeFrom.DEBUGGER, AppMode.WORKFLOW) + return queue, mocker.spy(queue, "publish") + + @staticmethod + @pytest.mark.parametrize("status", [HumanInputFormStatus.SUBMITTED, HumanInputFormStatus.TIMEOUT]) + @pytest.mark.parametrize("runs_human_input_node", [False, True], ids=["hidden-form", "direct-completion"]) + def test_resume_publishes_only_the_completed_form_among_repeated_node_invocations( + sqlite_session: Session, + workflow_queue: tuple[WorkflowAppQueueManager, MagicMock], + status: HumanInputFormStatus, + runs_human_input_node: bool, + ) -> None: + selected = _save_form(sqlite_session, status=status) + waiting = _save_form(sqlite_session, status=HumanInputFormStatus.WAITING) + expired = _save_form(sqlite_session, status=HumanInputFormStatus.EXPIRED) + _save_form(sqlite_session, status=HumanInputFormStatus.SUBMITTED) + queue, publish_spy = workflow_queue + runner = WorkflowBasedAppRunner(queue_manager=queue, app_id="app") + entry = _make_paused_workflow( + runner, [selected, waiting, selected, expired], human_input_form=selected if runs_human_input_node else None + ) + original_reasons = tuple(entry.graph_engine.runtime_state.graph_execution.pause_reasons) + + for event in _resume_events(runner, entry): + runner.handle_event(entry, event) + + published = [call.args[0] for call in publish_spy.call_args_list] + assert isinstance(published[0], QueueWorkflowStartedEvent) + completions = [ + item + for item in published + if isinstance(item, QueueHumanInputFormFilledEvent | QueueHumanInputFormTimeoutEvent) + ] + assert len(completions) == 1 + completion = completions[0] + assert completion.form_id == selected.id + assert completion.node_id == "outer-tool" + if status == HumanInputFormStatus.SUBMITTED: + assert isinstance(completion, QueueHumanInputFormFilledEvent) + assert (completion.action_id, completion.action_text, completion.rendered_content) == ( + "approve", + "Approve", + "Decision: approved", + ) + assert completion.submitted_data is not None + assert completion.submitted_data["answer"].to_object() == "approved" + else: + assert isinstance(completion, QueueHumanInputFormTimeoutEvent) + assert len(original_reasons) == 4 + assert all( + isinstance(reason, HitlRequired) and reason.node_id == "same-human-node" for reason in original_reasons + ) + assert entry.graph_engine.runtime_state.graph_execution.pause_reasons == [] + + @staticmethod + @pytest.mark.parametrize("owner_field", ["tenant_id", "app_id", "workflow_run_id"]) + def test_resume_rejects_form_outside_the_trusted_execution_owner( + sqlite_session: Session, workflow_queue: tuple[WorkflowAppQueueManager, MagicMock], owner_field: str + ) -> None: + form = _save_form(sqlite_session, status=HumanInputFormStatus.SUBMITTED) + setattr(form, owner_field, "another-owner") + sqlite_session.commit() + queue, publish_spy = workflow_queue + runner = WorkflowBasedAppRunner(queue_manager=queue, app_id="app") + entry = _make_paused_workflow(runner, [form]) + + events = list(_resume_events(runner, entry)) + assert isinstance(events[-1], GraphRunFailedEvent) + assert "does not belong" in events[-1].error + + assert not any(isinstance(call.args[0], QueueHumanInputFormFilledEvent) for call in publish_spy.call_args_list) + + @staticmethod + def test_resume_refreshes_expired_file_and_file_list_urls_before_form_completion( + sqlite_session: Session, + workflow_queue: tuple[WorkflowAppQueueManager, MagicMock], + monkeypatch: pytest.MonkeyPatch, + mocker: MockerFixture, + ) -> None: + mocker.patch.object( + type(workflow_response_converter.db), + "engine", + new_callable=PropertyMock, + return_value=sqlite_session.get_bind(), + ) + submitted_at = 1_700_000_000 + monkeypatch.setattr("core.app.workflow.file_runtime.time.time", lambda: submitted_at) + factory = DifyFileReferenceFactory(build_test_run_context()) + saved_files = [] + for name in ("first.txt", "second.txt"): + upload = UploadFile( + tenant_id="tenant", + storage_type=StorageType.LOCAL, + key=name, + name=name, + size=12, + extension="txt", + mime_type="text/plain", + created_by="user", + created_by_role=CreatorUserRole.END_USER, + created_at=datetime.now(UTC), + used=False, + ) + sqlite_session.add(upload) + sqlite_session.commit() + saved_files.append( + factory.build_from_mapping( + mapping={"transfer_method": "local_file", "upload_file_id": upload.id} + ).to_dict() + ) + + def url_is_valid(file: Mapping[str, object]) -> bool: + url, file_id = file["url"], file["related_id"] + assert isinstance(url, str) + assert isinstance(file_id, str) + query = parse_qs(urlparse(url).query) + return verify_file_signature( + upload_file_id=file_id, + timestamp=query["timestamp"][0], + nonce=query["nonce"][0], + sign=query["sign"][0], + ) + + assert all(url_is_valid(file) for file in saved_files) + form = _save_form(sqlite_session, status=HumanInputFormStatus.SUBMITTED) + definition = FormDefinition.model_validate_json(form.form_definition) + definition.inputs.extend( + [ + FileInputConfig(output_variable_name="attachment"), + FileListInputConfig(output_variable_name="attachments"), + ] + ) + form.form_definition = definition.model_dump_json() + form.submitted_data = json.dumps( + {"answer": "approved", "attachment": saved_files[0], "attachments": saved_files} + ) + sqlite_session.commit() + + monkeypatch.setattr( + "core.app.workflow.file_runtime.time.time", lambda: submitted_at + dify_config.FILES_ACCESS_TIMEOUT + 1 + ) + assert not any(url_is_valid(file) for file in saved_files) + queue, publish_spy = workflow_queue + runner = WorkflowBasedAppRunner(queue_manager=queue, app_id="app") + entry = _make_paused_workflow(runner, [form]) + for event in _resume_events(runner, entry): + runner.handle_event(entry, event) + + completions = [ + call.args[0] + for call in publish_spy.call_args_list + if isinstance(call.args[0], QueueHumanInputFormFilledEvent) + ] + assert len(completions) == 1 + converter = _build_converter() + converter.workflow_start_to_stream_response( + task_id="task", workflow_run_id="run-1", workflow_id="workflow", reason=WorkflowStartReason.RESUMPTION + ) + response = converter.human_input_form_filled_to_stream_response(event=completions[0], task_id="task") + data = response.data.submitted_data + assert data is not None + assert data["answer"] == "approved" + assert response.data.rendered_content == "Decision: approved" + restored_files = [data["attachment"], *data["attachments"]] + for restored, saved in zip(restored_files, [saved_files[0], *saved_files], strict=True): + assert (restored["related_id"], restored["filename"]) == (saved["related_id"], saved["filename"]) + assert url_is_valid(restored) + + @staticmethod + @pytest.mark.parametrize("late_status", [HumanInputFormStatus.SUBMITTED, HumanInputFormStatus.TIMEOUT]) + def test_form_completed_during_resume_is_published_once_before_the_terminal_event( + sqlite_session: Session, + workflow_queue: tuple[WorkflowAppQueueManager, MagicMock], + late_status: HumanInputFormStatus, + ) -> None: + first = _save_form(sqlite_session, status=HumanInputFormStatus.SUBMITTED) + second = _save_form(sqlite_session, status=HumanInputFormStatus.WAITING) + queue, publish_spy = workflow_queue + runner = WorkflowBasedAppRunner(queue_manager=queue, app_id="app") + entry = _make_paused_workflow(runner, [first, second]) + + def complete_second_form_during_execution(event: AppQueueEvent, publish_from: PublishFrom) -> None: + WorkflowAppQueueManager.publish(queue, event, publish_from) + # The second submission arrives after the initial completion query, + # while this real Engine attempt is already processing node events. + if not isinstance(event, QueueNodeStartedEvent): + return + repository = HumanInputFormSubmissionRepository() + if late_status == HumanInputFormStatus.SUBMITTED: + repository.mark_submitted( + form_id=second.id, + recipient_id=None, + selected_action_id="approve", + form_data={"answer": "approved"}, + submission_user_id=None, + submission_end_user_id="user", + ) + else: + repository.mark_timeout(form_id=second.id, timeout_status=late_status) + + publish_spy.side_effect = complete_second_form_during_execution + for event in _resume_events(runner, entry): + runner.handle_event(entry, event) + + published = [call.args[0] for call in publish_spy.call_args_list] + completions = [ + event + for event in published + if isinstance(event, QueueHumanInputFormFilledEvent | QueueHumanInputFormTimeoutEvent) + ] + assert [event.form_id for event in completions] == [first.id, second.id] + assert isinstance(published[0], QueueWorkflowStartedEvent) + assert isinstance(published[-1], QueueWorkflowSucceededEvent) + assert published.index(completions[-1]) < len(published) - 1 + + @staticmethod + def test_waiting_form_expiring_during_resume_publishes_timeout_before_success( + sqlite_session: Session, workflow_queue: tuple[WorkflowAppQueueManager, MagicMock] + ) -> None: + form = _save_form(sqlite_session, status=HumanInputFormStatus.WAITING) + form.expiration_time = datetime.now(UTC) - timedelta(seconds=1) + sqlite_session.commit() + queue, publish_spy = workflow_queue + runner = WorkflowBasedAppRunner(queue_manager=queue, app_id="app") + entry = _make_paused_workflow(runner, [form], human_input_form=form) + + for event in _resume_events(runner, entry): + runner.handle_event(entry, event) + + published = [call.args[0] for call in publish_spy.call_args_list] + assert isinstance(published[-1], QueueWorkflowSucceededEvent) + timeouts = [event for event in published if isinstance(event, QueueHumanInputFormTimeoutEvent)] + assert [event.form_id for event in timeouts] == [form.id] + assert timeouts[0].node_id == "outer-tool" + assert published.index(timeouts[0]) < len(published) - 1 + sqlite_session.refresh(form) + assert form.status == HumanInputFormStatus.TIMEOUT + + @staticmethod + def test_chat_history_keeps_repeated_forms_by_form_id(sqlite_session: Session) -> None: + pipeline = _build_pipeline() + pipeline._workflow_tenant_id = "tenant" + pipeline._application_generate_entity = _build_generate_entity() + pipeline._workflow_response_converter = _build_converter() + pipeline._workflow_response_converter.workflow_start_to_stream_response( + task_id="task-1", workflow_run_id="run-1", workflow_id="workflow-1", reason=WorkflowStartReason.INITIAL + ) + for form_id in ("first-form", "second-form", "first-form"): + event = QueueHumanInputFormFilledEvent( + form_id=form_id, + node_id="same-human-node", + node_type=BuiltinNodeTypes.HUMAN_INPUT, + node_title="Approval", + rendered_content="Decision: approved", + action_id="approve", + action_text="Approve", + ) + list(pipeline._handle_human_input_form_filled_event(event)) + + contents = sqlite_session.scalars(select(HumanInputContent)).all() + assert {(content.form_id, content.message_id) for content in contents} == { + ("first-form", "message-1"), + ("second-form", "message-1"), + } + assert len(contents) == 2 diff --git a/api/tests/unit_tests/core/app/apps/test_workflow_app_runner_core.py b/api/tests/unit_tests/core/app/apps/test_workflow_app_runner_core.py deleted file mode 100644 index 4c98b62b71a..00000000000 --- a/api/tests/unit_tests/core/app/apps/test_workflow_app_runner_core.py +++ /dev/null @@ -1,958 +0,0 @@ -from __future__ import annotations - -import json -from copy import deepcopy -from datetime import UTC, datetime -from types import SimpleNamespace -from typing import override - -import pytest -from pydantic import ValidationError - -from core.app.apps.execution_coordinator import AppExecutionState -from core.app.apps.workflow.app_queue_manager import WorkflowAppQueueManager -from core.app.apps.workflow_app_runner import WorkflowBasedAppRunner -from core.app.entities.app_invoke_entities import DIFY_RUN_CONTEXT_KEY, InvokeFrom, UserFrom -from core.app.entities.queue_entities import ( - QueueAgentLogEvent, - QueueIterationCompletedEvent, - QueueIterationStartEvent, - QueueLoopCompletedEvent, - QueueLoopStartEvent, - QueueNodeExceptionEvent, - QueueNodeFailedEvent, - QueueNodeRetryEvent, - QueueNodeStartedEvent, - QueueNodeSucceededEvent, - QueueReasoningChunkEvent, - QueueStopEvent, - QueueTextChunkEvent, - QueueWorkflowPausedEvent, - QueueWorkflowStartedEvent, - QueueWorkflowSucceededEvent, -) -from core.workflow.nodes.agent.events import NodeRunAgentLogEvent -from core.workflow.nodes.human_input.pause_reason import HumanInputRequired -from core.workflow.system_variables import default_system_variables -from core.workflow.workflow_entry import WorkflowEntry, iter_dify_graph_engine_events -from graphon.engine.layer import Layer -from graphon.engine_events import ( - GraphRunAbortedEvent, - GraphRunPausedEvent, - GraphRunStartedEvent, - GraphRunSucceededEvent, - NodeRunExceptionEvent, - NodeRunFailedEvent, - NodeRunIterationStartedEvent, - NodeRunIterationSucceededEvent, - NodeRunLoopFailedEvent, - NodeRunLoopStartedEvent, - NodeRunLoopSucceededEvent, - NodeRunReasoningChunkEvent, - NodeRunRetryEvent, - NodeRunStartedEvent, - NodeRunStreamChunkEvent, - NodeRunSucceededEvent, -) -from graphon.entities.pause_reason import HitlRequired -from graphon.enums import BuiltinNodeTypes, WorkflowNodeExecutionMetadataKey -from graphon.graph import Graph -from graphon.node_events import NodeRunResult -from graphon.runtime import RuntimeState, VariablePool -from graphon.variables.variables import StringVariable -from models.workflow import Workflow -from services.workflow_run_index import WorkflowRunIndex -from tests.unit_tests.model_factories import make_workflow - - -def _workflow(graph: dict[str, object] | None = None) -> Workflow: - return make_workflow( - tenant_id="tenant", - app_id="app", - graph=json.dumps(graph or {}, default=str), - created_by="account", - ) - - -def _workflow_entry(runner: WorkflowBasedAppRunner, state: RuntimeState | None = None) -> WorkflowEntry: - state = state or RuntimeState(workflow_id="workflow", variable_pool=VariablePool(), start_at=0.0) - graph = runner._init_graph( - graph_config={ - "nodes": [{"id": "start", "data": {"type": "start", "title": "Start", "variables": []}}], - "edges": [], - }, - graph_runtime_state=state, - user_from=UserFrom.ACCOUNT, - invoke_from=InvokeFrom.DEBUGGER, - call_depth=0, - ) - return WorkflowEntry(graph_config=graph.graph_config, graph=graph, call_depth=0, graph_runtime_state=state) - - -def _nested_container_graph(ownership: str) -> dict: - nodes = [ - {"id": "start", "data": {"type": "start", "title": "Start", "variables": []}}, - { - "id": "outer", - "data": { - "type": "loop", - "title": "Outer loop", - "start_node_id": "outer-start", - "loop_count": 1, - "break_conditions": [], - "logical_operator": "and", - }, - }, - {"id": "outer-start", "data": {"type": "loop-start", "title": "Loop start"}}, - { - "id": "inner", - "data": { - "type": "iteration", - "title": "Inner iteration", - "start_node_id": "inner-start", - "iterator_selector": ["start", "items"], - "output_selector": ["inner", "item"], - }, - }, - {"id": "inner-start", "data": {"type": "iteration-start", "title": "Iteration start"}}, - { - "id": "deep", - "data": { - "type": "loop", - "title": "Nested loop", - "start_node_id": "deep-start", - "loop_count": 1, - "break_conditions": [], - "logical_operator": "and", - }, - }, - {"id": "deep-start", "data": {"type": "loop-start", "title": "Nested loop start"}}, - {"id": "end", "data": {"type": "end", "title": "Output", "outputs": []}}, - ] - owners = {"outer-start": "outer", "inner": "outer", "inner-start": "inner", "deep": "inner", "deep-start": "deep"} - for node in nodes: - owner = owners.get(node["id"], "") - if ownership == "canonical": - node["data"]["container_id"] = owner - # Canonical ownership, including the empty root owner, takes precedence. - node["parentId"] = "stale-owner" - node["data"]["loop_id"] = "stale-owner" - elif ownership == "parent": - if owner: - node["parentId"] = owner - elif owner: - node["data"]["loop_id"] = "deep" if owner == "deep" else "outer" - if owner in {"inner", "deep"}: - node["data"]["iteration_id"] = "inner" - return { - "nodes": nodes, - "edges": [ - {"id": "start-outer", "source": "start", "target": "outer"}, - {"id": "outer-end", "source": "outer", "target": "end"}, - {"id": "outer-start-inner", "source": "outer-start", "target": "inner"}, - {"id": "inner-start-deep", "source": "inner-start", "target": "deep"}, - ], - } - - -class TestWorkflowBasedAppRunner: - @pytest.mark.parametrize( - ("event_class", "node_type"), - [ - (NodeRunIterationStartedEvent, BuiltinNodeTypes.ITERATION), - (NodeRunIterationSucceededEvent, BuiltinNodeTypes.ITERATION), - (NodeRunLoopStartedEvent, BuiltinNodeTypes.LOOP), - (NodeRunLoopSucceededEvent, BuiltinNodeTypes.LOOP), - ], - ) - def test_container_queue_events_preserve_lifecycle_and_ancestor_metadata( - self, - event_class: type[ - NodeRunIterationStartedEvent - | NodeRunIterationSucceededEvent - | NodeRunLoopStartedEvent - | NodeRunLoopSucceededEvent - ], - node_type: str, - ) -> None: - queue = WorkflowAppQueueManager( - task_id="metadata-contract", user_id="account", invoke_from=InvokeFrom.DEBUGGER, app_mode="workflow" - ) - runner = WorkflowBasedAppRunner(queue_manager=queue, app_id="app") - entry = _workflow_entry(runner) - event = event_class( - id="container-execution", - node_id="container", - node_type=node_type, - node_title="Container", - start_at=datetime.now(UTC), - metadata={"iteration_length": 2, "total_tokens": 1}, - node_run_result=NodeRunResult( - metadata={ - WorkflowNodeExecutionMetadataKey.LOOP_ID: "outer", - WorkflowNodeExecutionMetadataKey.LOOP_INDEX: 3, - WorkflowNodeExecutionMetadataKey.TOTAL_TOKENS: 2, - } - ), - ) - original_event = event.model_dump() - - runner.handle_event(entry, event) - queue.stop_listen(execution_state=AppExecutionState.TERMINAL) - messages = list(queue.listen()) - - assert len(messages) == 1 - queued_event = messages[0].event - assert isinstance( - queued_event, - QueueIterationStartEvent | QueueIterationCompletedEvent | QueueLoopStartEvent | QueueLoopCompletedEvent, - ) - assert queued_event.metadata == { - "iteration_length": 2, - "total_tokens": 2, - "loop_id": "outer", - "loop_index": 3, - } - assert event.model_dump() == original_event - - def test_resolve_user_from(self): - runner = WorkflowBasedAppRunner(queue_manager=SimpleNamespace(), app_id="app") - - assert runner._resolve_user_from(InvokeFrom.EXPLORE) == UserFrom.ACCOUNT - assert runner._resolve_user_from(InvokeFrom.DEBUGGER) == UserFrom.ACCOUNT - assert runner._resolve_user_from(InvokeFrom.WEB_APP) == UserFrom.END_USER - - def test_init_graph_validates_graph_structure(self): - runner = WorkflowBasedAppRunner(queue_manager=SimpleNamespace(), app_id="app") - - runtime_state = RuntimeState( - workflow_id="test-workflow", - variable_pool=VariablePool.from_bootstrap(system_variables=default_system_variables()), - start_at=0.0, - ) - - with pytest.raises(ValidationError, match="nodes"): - runner._init_graph( - graph_config={}, - graph_runtime_state=runtime_state, - user_from=UserFrom.ACCOUNT, - invoke_from=InvokeFrom.DEBUGGER, - call_depth=0, - ) - - with pytest.raises(ValidationError, match="nodes"): - runner._init_graph( - graph_config={"nodes": {}, "edges": []}, - graph_runtime_state=runtime_state, - user_from=UserFrom.ACCOUNT, - invoke_from=InvokeFrom.DEBUGGER, - call_depth=0, - ) - - with pytest.raises(ValidationError, match="edges"): - runner._init_graph( - graph_config={"nodes": [], "edges": {}}, - graph_runtime_state=runtime_state, - user_from=UserFrom.ACCOUNT, - invoke_from=InvokeFrom.DEBUGGER, - call_depth=0, - ) - - def test_init_graph_includes_trace_session_id_in_run_context(self, monkeypatch: pytest.MonkeyPatch): - runner = WorkflowBasedAppRunner(queue_manager=SimpleNamespace(), app_id="app") - runtime_state = RuntimeState( - workflow_id="test-workflow", - variable_pool=VariablePool.from_bootstrap(system_variables=default_system_variables()), - start_at=0.0, - ) - captured = {} - - def fake_from_graph_init_context(**kwargs): - graph_init_context = kwargs["graph_init_context"] - captured["run_context"] = graph_init_context.run_context - captured["call_depth"] = graph_init_context.call_depth - captured["runtime_state"] = kwargs["runtime_state"] - return SimpleNamespace() - - monkeypatch.setattr( - "core.app.apps.workflow_app_runner.DifyNodeFactory.from_graph_init_context", - fake_from_graph_init_context, - ) - monkeypatch.setattr("core.app.apps.workflow_app_runner.Graph.init", lambda **_kwargs: SimpleNamespace()) - - runner._init_graph( - graph_config={"nodes": [], "edges": []}, - graph_runtime_state=runtime_state, - user_from=UserFrom.ACCOUNT, - invoke_from=InvokeFrom.DEBUGGER, - root_node_id="root", - trace_session_id="session-1", - call_depth=3, - ) - - assert captured["run_context"][DIFY_RUN_CONTEXT_KEY].trace_session_id == "session-1" - assert captured["call_depth"] == 3 - assert captured["runtime_state"] is runtime_state - - @pytest.mark.parametrize("ownership", ["canonical", "parent", "legacy"]) - def test_init_graph_preserves_container_execution_scopes(self, ownership: str) -> None: - runner = WorkflowBasedAppRunner(queue_manager=SimpleNamespace(), app_id="app") - runtime_state = RuntimeState( - workflow_id="test-workflow", - variable_pool=VariablePool.from_bootstrap(system_variables=default_system_variables()), - start_at=0.0, - ) - graph_config = _nested_container_graph(ownership) - original_config = deepcopy(graph_config) - graph = runner._init_graph( - graph_config=graph_config, - graph_runtime_state=runtime_state, - user_from=UserFrom.ACCOUNT, - invoke_from=InvokeFrom.DEBUGGER, - root_node_id="start", - call_depth=0, - ) - - assert set(graph.nodes) == {"start", "outer", "end"} - outer_graph = Graph.init( - graph_config=graph.graph_config, - node_factory=graph.node_factory, - root_node_id="outer-start", - container_id="outer", - ) - assert set(outer_graph.nodes) == {"outer-start", "inner"} - inner_graph = Graph.init( - graph_config=outer_graph.graph_config, - node_factory=outer_graph.node_factory, - root_node_id="inner-start", - container_id="inner", - ) - assert set(inner_graph.nodes) == {"inner-start", "deep"} - deep_graph = Graph.init( - graph_config=inner_graph.graph_config, - node_factory=inner_graph.node_factory, - root_node_id="deep-start", - container_id="deep", - ) - assert set(deep_graph.nodes) == {"deep-start"} - assert graph_config == original_config - - def test_prepare_single_node_execution_requires_run(self): - runner = WorkflowBasedAppRunner(queue_manager=SimpleNamespace(), app_id="app") - - workflow = _workflow() - - with pytest.raises(ValueError, match="Neither single_iteration_run nor single_loop_run"): - runner._prepare_single_node_execution(workflow, None, None, user_id="00000000-0000-0000-0000-000000000001") - - @pytest.mark.parametrize("ownership", ["canonical", "parent", "legacy"]) - @pytest.mark.parametrize("node_id", ["outer", "inner", "deep"]) - def test_single_container_debugging_preserves_descendants(self, ownership: str, node_id: str): - runner = WorkflowBasedAppRunner(queue_manager=SimpleNamespace(), app_id="app") - graph_runtime_state = RuntimeState( - workflow_id="test-workflow", - variable_pool=VariablePool.from_bootstrap(system_variables=default_system_variables()), - start_at=0.0, - ) - - graph_config = _nested_container_graph(ownership) - workflow = _workflow(graph_config) - workflow.id = "workflow" - graph_runtime_state.variable_pool.add(["start", "items"], ["first", "second"]) - - graph, variable_pool = runner._get_graph_and_variable_pool_for_single_node_run( - workflow=workflow, - node_id=node_id, - user_inputs={}, - graph_runtime_state=graph_runtime_state, - node_type_label="iteration", - user_id="00000000-0000-0000-0000-000000000001", - ) - - assert set(graph.nodes) == {node_id} - assert variable_pool is graph_runtime_state.variable_pool - deep_graph = Graph.init( - graph_config=graph.graph_config, - node_factory=graph.node_factory, - root_node_id="deep-start", - container_id="deep", - ) - assert set(deep_graph.nodes) == {"deep-start"} - - def test_get_graph_and_variable_pool_for_single_node_run_does_not_mutate_inputs( - self, monkeypatch: pytest.MonkeyPatch - ): - runner = WorkflowBasedAppRunner(queue_manager=SimpleNamespace(), app_id="app") - graph_runtime_state = RuntimeState( - workflow_id="workflow", - variable_pool=VariablePool.from_bootstrap(system_variables=default_system_variables()), - start_at=0.0, - ) - graph_config = { - "nodes": [ - {"id": "node-1", "data": {"type": "start", "version": "1"}}, - {"id": "node-child", "data": {"type": "start", "version": "1", "iteration_id": "node-1"}}, - {"id": "outside", "data": {"type": "start", "version": "1"}}, - ], - "edges": [ - {"id": "edge-1", "source": "node-1", "target": "node-child"}, - {"id": "edge-2", "source": "outside", "target": "node-1"}, - ], - } - expected_graph_config = deepcopy(graph_config) - workflow = SimpleNamespace(tenant_id="tenant", id="workflow", graph_dict=graph_config) - user_inputs = {"question": "hello"} - expected_user_inputs = dict(user_inputs) - captured_graph_config = {} - - def fake_graph_init(**kwargs): - captured_graph_config["value"] = kwargs["graph_config"] - return SimpleNamespace() - - class _NodeCls: - @staticmethod - def extract_variable_selector_to_variable_mapping(graph_config, config): - return {} - - from core.app.apps import workflow_app_runner - - monkeypatch.setattr("core.app.apps.workflow_app_runner.Graph.init", fake_graph_init) - monkeypatch.setattr(workflow_app_runner, "resolve_workflow_node_class", lambda **_kwargs: _NodeCls) - monkeypatch.setattr( - "core.app.apps.workflow_app_runner.load_into_variable_pool", - lambda **kwargs: kwargs["user_inputs"].update({"from_loader": "local-only"}), - ) - monkeypatch.setattr( - "core.app.apps.workflow_app_runner.WorkflowEntry.mapping_user_inputs_to_variable_pool", - lambda **kwargs: kwargs["user_inputs"].update({"from_mapping": "local-only"}), - ) - - graph, variable_pool = runner._get_graph_and_variable_pool_for_single_node_run( - workflow=workflow, - node_id="node-1", - user_inputs=user_inputs, - graph_runtime_state=graph_runtime_state, - node_type_label="iteration", - user_id="00000000-0000-0000-0000-000000000001", - ) - - assert graph is not None - assert variable_pool is graph_runtime_state.variable_pool - assert workflow.graph_dict == expected_graph_config - assert user_inputs == expected_user_inputs - assert captured_graph_config["value"] is not workflow.graph_dict - assert [node["id"] for node in captured_graph_config["value"]["nodes"]] == ["node-1", "node-child"] - assert captured_graph_config["value"]["edges"] == [{"id": "edge-1", "source": "node-1", "target": "node-child"}] - - def test_get_graph_and_variable_pool_for_single_node_run_includes_trace_session_id( - self, monkeypatch: pytest.MonkeyPatch - ): - runner = WorkflowBasedAppRunner(queue_manager=SimpleNamespace(), app_id="app") - graph_runtime_state = RuntimeState( - workflow_id="test-workflow", - variable_pool=VariablePool.from_bootstrap(system_variables=default_system_variables()), - start_at=0.0, - ) - graph_config = { - "nodes": [{"id": "node-1", "data": {"type": "start", "version": "1"}}], - "edges": [], - } - workflow = _workflow(graph_config) - workflow.id = "workflow" - captured = {} - - def fake_from_graph_init_context(**kwargs): - captured["run_context"] = kwargs["graph_init_context"].run_context - captured["runtime_state"] = kwargs["runtime_state"] - return SimpleNamespace() - - class _NodeCls: - @staticmethod - def extract_variable_selector_to_variable_mapping(graph_config, config): - return {} - - from core.app.apps import workflow_app_runner - - monkeypatch.setattr( - "core.app.apps.workflow_app_runner.DifyNodeFactory.from_graph_init_context", - fake_from_graph_init_context, - ) - monkeypatch.setattr("core.app.apps.workflow_app_runner.Graph.init", lambda **kwargs: SimpleNamespace()) - monkeypatch.setattr(workflow_app_runner, "resolve_workflow_node_class", lambda **_kwargs: _NodeCls) - monkeypatch.setattr("core.app.apps.workflow_app_runner.load_into_variable_pool", lambda **kwargs: None) - monkeypatch.setattr( - "core.app.apps.workflow_app_runner.WorkflowEntry.mapping_user_inputs_to_variable_pool", - lambda **kwargs: None, - ) - - runner._get_graph_and_variable_pool_for_single_node_run( - workflow=workflow, - node_id="node-1", - user_inputs={}, - graph_runtime_state=graph_runtime_state, - node_type_label="iteration", - user_id="00000000-0000-0000-0000-000000000001", - trace_session_id="session-1", - ) - - assert captured["run_context"][DIFY_RUN_CONTEXT_KEY].trace_session_id == "session-1" - assert captured["runtime_state"] is graph_runtime_state - - def test_get_graph_and_variable_pool_preloads_constructor_variables_before_graph_init( - self, monkeypatch: pytest.MonkeyPatch - ): - variable_loader = SimpleNamespace( - load_variables=lambda selectors: ( - [ - StringVariable( - name="conversation_id", - value="conv-1", - selector=["sys", "conversation_id"], - ) - ] - if selectors - else [] - ) - ) - runner = WorkflowBasedAppRunner( - queue_manager=SimpleNamespace(), - variable_loader=variable_loader, - app_id="app", - ) - graph_runtime_state = RuntimeState( - workflow_id="test-workflow", - variable_pool=VariablePool.from_bootstrap(system_variables=default_system_variables()), - start_at=0.0, - ) - - workflow = _workflow( - { - "nodes": [ - {"id": "loop-node", "data": {"type": "loop", "version": "1", "title": "Loop"}}, - { - "id": "llm-child", - "data": { - "type": "llm", - "version": "1", - "loop_id": "loop-node", - "memory": object(), - }, - }, - ], - "edges": [], - } - ) - workflow.id = "workflow" - - class _LoopNodeCls: - @staticmethod - def extract_variable_selector_to_variable_mapping(graph_config, config): - return {} - - def _graph_init(**kwargs): - variable_pool = graph_runtime_state.variable_pool - assert variable_pool.get(["sys", "conversation_id"]) is not None - return SimpleNamespace() - - monkeypatch.setattr( - "core.app.apps.workflow_app_runner.Graph.init", - _graph_init, - ) - monkeypatch.setattr( - "core.app.apps.workflow_app_runner.resolve_workflow_node_class", - lambda **_kwargs: _LoopNodeCls, - ) - monkeypatch.setattr( - "core.app.apps.workflow_app_runner.load_into_variable_pool", - lambda **kwargs: None, - ) - monkeypatch.setattr( - "core.app.apps.workflow_app_runner.WorkflowEntry.mapping_user_inputs_to_variable_pool", - lambda **kwargs: None, - ) - - graph, variable_pool = runner._get_graph_and_variable_pool_for_single_node_run( - workflow=workflow, - node_id="loop-node", - user_inputs={}, - graph_runtime_state=graph_runtime_state, - node_type_label="loop", - ) - - assert graph is not None - assert variable_pool.get(["sys", "conversation_id"]).value == "conv-1" - - def test_handle_graph_run_events_and_pause_notifications(self, monkeypatch: pytest.MonkeyPatch): - published: list[object] = [] - - class _QueueManager: - def publish(self, event, publish_from): - published.append((event, publish_from)) - - runner = WorkflowBasedAppRunner(queue_manager=_QueueManager(), app_id="app") - graph_runtime_state = RuntimeState( - workflow_id="test-workflow", - variable_pool=VariablePool.from_bootstrap(system_variables=default_system_variables()), - start_at=0.0, - ) - workflow_entry = _workflow_entry(runner, graph_runtime_state) - - emails: list[dict] = [] - - class _Dispatch: - def apply_async(self, *, kwargs, queue): - emails.append({"kwargs": kwargs, "queue": queue}) - - monkeypatch.setattr( - "core.app.apps.workflow_app_runner.dispatch_human_input_email_task", - _Dispatch(), - ) - enriched_reason = HumanInputRequired( - form_id="form", - form_content="content", - node_id="node-1", - node_title="Node", - ) - - reason = HitlRequired( - session_id="form", - node_id="node-1", - node_title="Node", - ) - - runner.handle_event(workflow_entry, GraphRunStartedEvent()) - runner.handle_event(workflow_entry, GraphRunSucceededEvent(outputs={"ok": True})) - runner.handle_event( - workflow_entry, GraphRunPausedEvent(reasons=[reason], outputs={}), pause_reasons=[enriched_reason] - ) - - assert any(isinstance(event, QueueWorkflowStartedEvent) for event, _ in published) - assert any(isinstance(event, QueueWorkflowSucceededEvent) for event, _ in published) - paused_event = next(event for event, _ in published if isinstance(event, QueueWorkflowPausedEvent)) - assert paused_event.paused_nodes == ["node-1"] - assert emails - - def test_handle_graph_aborted_publishes_stopped_terminal(self): - published: list[object] = [] - - class _QueueManager: - def publish(self, event, publish_from): - del publish_from - published.append(event) - - runner = WorkflowBasedAppRunner(queue_manager=_QueueManager(), app_id="app") - workflow_entry = _workflow_entry(runner) - - runner.handle_event(workflow_entry, GraphRunAbortedEvent(reason="User requested stop", outputs={})) - - event = published[-1] - assert isinstance(event, QueueStopEvent) - assert event.get_stop_reason() == "User requested stop" - - def test_handle_node_events_publishes_queue_events(self): - published: list[object] = [] - - class _QueueManager: - def publish(self, event, publish_from): - published.append(event) - - runner = WorkflowBasedAppRunner(queue_manager=_QueueManager(), app_id="app") - graph_runtime_state = RuntimeState( - workflow_id="test-workflow", - variable_pool=VariablePool.from_bootstrap(system_variables=default_system_variables()), - start_at=0.0, - ) - workflow_entry = _workflow_entry(runner, graph_runtime_state) - - runner.handle_event( - workflow_entry, - NodeRunStartedEvent( - id="exec", - node_id="node", - node_type=BuiltinNodeTypes.START, - node_title="Start", - start_at=datetime.now(UTC), - ), - ) - runner.handle_event( - workflow_entry, - NodeRunStreamChunkEvent( - id="exec", - node_id="node", - node_type=BuiltinNodeTypes.START, - selector=["node", "text"], - chunk="hi", - is_final=False, - ), - ) - runner.handle_event( - workflow_entry, - NodeRunReasoningChunkEvent( - id="exec", - node_id="node", - node_type=BuiltinNodeTypes.LLM, - selector=["node", "reasoning_content"], - chunk="thinking", - is_final=False, - ), - ) - runner.handle_event( - workflow_entry, - NodeRunAgentLogEvent( - id="exec", - node_id="node", - node_type=BuiltinNodeTypes.START, - message_id="msg", - label="label", - node_execution_id="exec", - parent_id=None, - error=None, - status="done", - data={}, - metadata={}, - ), - ) - runner.handle_event( - workflow_entry, - NodeRunIterationSucceededEvent( - id="exec", - node_id="node", - node_type=BuiltinNodeTypes.LLM, - node_title="Iter", - start_at=datetime.now(UTC), - inputs={}, - outputs={"ok": True}, - metadata={}, - steps=1, - ), - ) - runner.handle_event( - workflow_entry, - NodeRunLoopFailedEvent( - id="exec", - node_id="node", - node_type=BuiltinNodeTypes.LLM, - node_title="Loop", - start_at=datetime.now(UTC), - inputs={}, - outputs={}, - metadata={}, - steps=1, - error="boom", - ), - ) - - assert any(isinstance(event, QueueTextChunkEvent) for event in published) - assert any(isinstance(event, QueueReasoningChunkEvent) for event in published) - assert any(isinstance(event, QueueAgentLogEvent) for event in published) - assert any(isinstance(event, QueueIterationCompletedEvent) for event in published) - assert any(isinstance(event, QueueLoopCompletedEvent) for event in published) - - @pytest.mark.parametrize("ownership", ["canonical", "parent", "legacy"]) - def test_nested_engine_events_retain_both_container_ancestors(self, ownership: str): - from tests.unit_tests.core.app.workflow.test_persistence_layer import _make_layer - - published: list[object] = [] - - class _QueueManager: - def publish(self, event, publish_from): - del publish_from - published.append(event) - - graph_config = _nested_container_graph(ownership) - graph_config["nodes"].append( - {"id": "deep-end", "data": {"type": "end", "title": "Nested End", "outputs": [], "container_id": "deep"}} - ) - graph_config["edges"].append({"id": "deep-output", "source": "deep-start", "target": "deep-end"}) - layer, _, node_repo, state = _make_layer(graph_data=graph_config) - persisted_starts: dict[str, dict] = {} - - class _StartMetadataLayer(Layer): - @override - def on_event(self, event): - if isinstance(event, NodeRunStartedEvent): - persisted_starts[event.node_id] = dict(node_repo.saved[-1].metadata) - - state.variable_pool.add(["start", "items"], ["item"]) - runner = WorkflowBasedAppRunner(queue_manager=_QueueManager(), app_id="app") - graph = runner._init_graph( - graph_config=graph_config, - graph_runtime_state=state, - user_from=UserFrom.ACCOUNT, - invoke_from=InvokeFrom.DEBUGGER, - call_depth=0, - ) - entry = WorkflowEntry(graph_config=graph.graph_config, graph=graph, call_depth=0, graph_runtime_state=state) - engine = entry.graph_engine - index = WorkflowRunIndex(()) - layer.set_node_run_indices(index.indices) - engine.add_layer(index) - engine.add_layer(layer) - engine.add_layer(_StartMetadataLayer()) - for event in iter_dify_graph_engine_events(engine): - runner.handle_event(entry, event) - - expected = { - "outer": (None, None), - "inner": (None, "outer"), - "deep": ("inner", "outer"), - "deep-end": ("inner", "deep"), - } - started = {event.node_id: event for event in published if isinstance(event, QueueNodeStartedEvent)} - persisted = {event.node_id: event for event in node_repo.saved} - for node_id, (iteration_id, loop_id) in expected.items(): - assert (started[node_id].in_iteration_id, started[node_id].in_loop_id) == (iteration_id, loop_id) - for metadata in (persisted_starts[node_id], persisted[node_id].metadata): - assert (metadata.get("iteration_id"), metadata.get("loop_id")) == (iteration_id, loop_id) - for metadata in (persisted_starts["deep-end"], persisted["deep-end"].metadata): - assert metadata["iteration_index"] == 0 - assert metadata["loop_index"] == 0 - for event in published: - if isinstance( - event, - QueueIterationStartEvent | QueueIterationCompletedEvent | QueueLoopStartEvent | QueueLoopCompletedEvent, - ): - assert (event.metadata.get("iteration_id"), event.metadata.get("loop_id")) == expected[event.node_id] - - def test_handle_nested_answer_success_publishes_text_chunk(self): - published: list[object] = [] - - class _QueueManager: - def publish(self, event, publish_from): - published.append(event) - - runner = WorkflowBasedAppRunner(queue_manager=_QueueManager(), app_id="app") - workflow_entry = _workflow_entry(runner) - - runner.handle_event( - workflow_entry, - NodeRunSucceededEvent( - id="exec", - node_id="answer-in-iteration", - node_type=BuiltinNodeTypes.ANSWER, - start_at=datetime.now(UTC), - finished_at=datetime.now(UTC), - node_run_result=NodeRunResult( - outputs={"answer": "inside iteration"}, metadata={"iteration_id": "iteration"} - ), - container_id="iteration", - ), - ) - - text_chunk = published[0] - assert isinstance(text_chunk, QueueTextChunkEvent) - assert text_chunk.text == "inside iteration" - assert text_chunk.from_variable_selector == ["answer-in-iteration", "answer"] - assert text_chunk.in_iteration_id == "iteration" - assert any(isinstance(event, QueueNodeSucceededEvent) for event in published) - - def test_handle_root_answer_success_does_not_publish_extra_text_chunk(self): - published: list[object] = [] - - class _QueueManager: - def publish(self, event, publish_from): - published.append(event) - - runner = WorkflowBasedAppRunner(queue_manager=_QueueManager(), app_id="app") - workflow_entry = _workflow_entry(runner) - - runner.handle_event( - workflow_entry, - NodeRunSucceededEvent( - id="exec", - node_id="answer", - node_type=BuiltinNodeTypes.ANSWER, - start_at=datetime.now(UTC), - finished_at=datetime.now(UTC), - node_run_result=NodeRunResult(outputs={"answer": "root answer"}), - ), - ) - - assert not any(isinstance(event, QueueTextChunkEvent) for event in published) - assert len(published) == 1 - assert isinstance(published[0], QueueNodeSucceededEvent) - - @pytest.mark.parametrize( - ("event_factory", "queue_event_cls"), - [ - ( - lambda result, start_at, finished_at: NodeRunSucceededEvent( - id="exec", - node_id="node", - node_type=BuiltinNodeTypes.START, - start_at=start_at, - finished_at=finished_at, - node_run_result=result, - ), - QueueNodeSucceededEvent, - ), - ( - lambda result, start_at, finished_at: NodeRunFailedEvent( - id="exec", - node_id="node", - node_type=BuiltinNodeTypes.START, - start_at=start_at, - finished_at=finished_at, - error="boom", - node_run_result=result, - ), - QueueNodeFailedEvent, - ), - ( - lambda result, start_at, finished_at: NodeRunExceptionEvent( - id="exec", - node_id="node", - node_type=BuiltinNodeTypes.START, - start_at=start_at, - finished_at=finished_at, - error="boom", - node_run_result=result, - ), - QueueNodeExceptionEvent, - ), - ( - lambda result, start_at, _finished_at: NodeRunRetryEvent( - id="exec", - node_id="node", - node_type=BuiltinNodeTypes.START, - node_title="Start", - start_at=start_at, - error="boom", - retry_index=1, - node_run_result=result, - ), - QueueNodeRetryEvent, - ), - ], - ) - def test_handle_start_node_result_events_project_outputs(self, event_factory, queue_event_cls): - published: list[object] = [] - - class _QueueManager: - def publish(self, event, publish_from): - published.append(event) - - runner = WorkflowBasedAppRunner(queue_manager=_QueueManager(), app_id="app") - graph_runtime_state = RuntimeState( - workflow_id="test-workflow", - variable_pool=VariablePool.from_bootstrap(system_variables=default_system_variables()), - start_at=0.0, - ) - workflow_entry = _workflow_entry(runner, graph_runtime_state) - started_at = datetime.now(UTC) - finished_at = datetime.now(UTC) - result = NodeRunResult( - inputs={"question": "hello"}, - outputs={ - "question": "hello", - "sys.query": "hello", - "env.API_KEY": "secret", - "conversation.session_id": "session-1", - }, - ) - - runner.handle_event(workflow_entry, event_factory(result, started_at, finished_at)) - - queue_event = published[-1] - assert isinstance(queue_event, queue_event_cls) - assert queue_event.outputs == {"question": "hello"} diff --git a/api/tests/unit_tests/core/app/apps/test_workflow_app_runner_notifications.py b/api/tests/unit_tests/core/app/apps/test_workflow_app_runner_notifications.py deleted file mode 100644 index d0437d0fa22..00000000000 --- a/api/tests/unit_tests/core/app/apps/test_workflow_app_runner_notifications.py +++ /dev/null @@ -1,60 +0,0 @@ -from unittest.mock import MagicMock - -import pytest - -from core.app.apps.workflow_app_runner import WorkflowBasedAppRunner -from core.app.entities.queue_entities import QueueWorkflowPausedEvent -from core.workflow.nodes.human_input.pause_reason import HumanInputRequired -from graphon.engine_events import GraphRunPausedEvent -from graphon.entities.pause_reason import HitlRequired - - -class _DummyQueueManager: - def __init__(self): - self.published = [] - - def publish(self, event, _from): - self.published.append(event) - - -class _DummyRuntimeState: - variable_pool = object() - - -class _DummyGraphEngine: - def __init__(self): - self.runtime_state = _DummyRuntimeState() - - -class _DummyWorkflowEntry: - def __init__(self): - self.graph_engine = _DummyGraphEngine() - - -def test_handle_pause_event_enqueues_email_task(monkeypatch: pytest.MonkeyPatch): - queue_manager = _DummyQueueManager() - runner = WorkflowBasedAppRunner(queue_manager=queue_manager, app_id="app-id") - workflow_entry = _DummyWorkflowEntry() - - graph_reason = HitlRequired(session_id="form-123", node_id="node-1", node_title="Review") - event = GraphRunPausedEvent(reasons=[graph_reason], outputs={}) - - email_task = MagicMock() - enriched_reason = HumanInputRequired( - form_id="form-123", - form_content="content", - inputs=[], - actions=[], - node_id="node-1", - node_title="Review", - ) - monkeypatch.setattr("core.app.apps.workflow_app_runner.dispatch_human_input_email_task", email_task) - - runner.handle_event(workflow_entry, event, pause_reasons=[enriched_reason]) - - email_task.apply_async.assert_called_once() - kwargs = email_task.apply_async.call_args.kwargs["kwargs"] - assert kwargs["form_id"] == "form-123" - assert kwargs["node_title"] == "Review" - - assert any(isinstance(evt, QueueWorkflowPausedEvent) for evt in queue_manager.published) diff --git a/api/tests/unit_tests/core/app/apps/test_workflow_app_runner_single_node.py b/api/tests/unit_tests/core/app/apps/test_workflow_app_runner_single_node.py deleted file mode 100644 index 5da553ba7ed..00000000000 --- a/api/tests/unit_tests/core/app/apps/test_workflow_app_runner_single_node.py +++ /dev/null @@ -1,265 +0,0 @@ -from __future__ import annotations - -import json -from unittest.mock import MagicMock, patch - -import pytest -from pydantic import ValidationError - -from core.app.app_config.entities import WorkflowUIBasedAppConfig -from core.app.apps.workflow.app_queue_manager import WorkflowAppQueueManager -from core.app.apps.workflow.app_runner import WorkflowAppRunner -from core.app.apps.workflow_app_runner import WorkflowBasedAppRunner -from core.app.entities.app_invoke_entities import InvokeFrom, WorkflowAppGenerateEntity -from core.credit_usage import CreditUsageAppType -from core.workflow.system_variables import default_system_variables -from graphon.runtime import RuntimeState, VariablePool -from models.model import AppMode -from models.workflow import Workflow, WorkflowKind - - -def _make_graph_state(): - variable_pool = VariablePool.from_bootstrap( - system_variables=default_system_variables(), - user_inputs={}, - environment_variables=[], - conversation_variables=[], - ) - return ( - MagicMock(), - variable_pool, - RuntimeState(workflow_id="test-workflow", variable_pool=variable_pool, start_at=0.0), - ) - - -@pytest.fixture -def queue_manager() -> WorkflowAppQueueManager: - return WorkflowAppQueueManager( - task_id="task-id", user_id="user", invoke_from=InvokeFrom.SERVICE_API, app_mode=AppMode.WORKFLOW - ) - - -@pytest.mark.parametrize( - ("single_iteration_run", "single_loop_run"), - [ - (WorkflowAppGenerateEntity.SingleIterationRunEntity(node_id="iter", inputs={}), None), - (None, WorkflowAppGenerateEntity.SingleLoopRunEntity(node_id="loop", inputs={})), - ], -) -def test_run_uses_single_node_execution_branch( - single_iteration_run: WorkflowAppGenerateEntity.SingleIterationRunEntity | None, - single_loop_run: WorkflowAppGenerateEntity.SingleLoopRunEntity | None, - queue_manager: WorkflowAppQueueManager, -) -> None: - app_generate_entity = WorkflowAppGenerateEntity( - app_config=WorkflowUIBasedAppConfig( - app_id="app", tenant_id="tenant", workflow_id="workflow", app_mode=AppMode.WORKFLOW - ), - inputs={}, - files=[], - user_id="user", - invoke_from=InvokeFrom.SERVICE_API, - workflow_execution_id="execution-id", - task_id="task-id", - stream=True, - extras={"trace_session_id": "session-1"}, - single_iteration_run=single_iteration_run, - single_loop_run=single_loop_run, - ) - - workflow = Workflow( - tenant_id="tenant", - app_id="app", - id="workflow", - type="workflow", - version="v1", - graph=json.dumps({"nodes": [], "edges": []}), - ) - workflow.environment_variables = [] - - runner = WorkflowAppRunner( - application_generate_entity=app_generate_entity, - queue_manager=queue_manager, - variable_loader=MagicMock(), - workflow=workflow, - system_user_id="system-user", - ) - - graph, variable_pool, graph_runtime_state = _make_graph_state() - mock_workflow_entry = MagicMock() - mock_workflow_entry.graph_engine = MagicMock() - mock_workflow_entry.graph_engine.layer = MagicMock() - mock_workflow_entry.run.return_value = iter([]) - - with ( - patch("core.app.apps.workflow.app_runner.RedisChannel"), - patch("core.app.apps.workflow.app_runner.redis_client"), - patch("core.app.apps.workflow.app_runner.WorkflowEntry", return_value=mock_workflow_entry) as entry_class, - patch.object( - runner, - "_prepare_single_node_execution", - return_value=( - graph, - variable_pool, - graph_runtime_state, - ), - ) as prepare_single, - patch.object(runner, "_init_graph") as init_graph, - ): - runner.prepare() - - prepare_single.assert_called_once_with( - workflow=workflow, - single_iteration_run=single_iteration_run, - single_loop_run=single_loop_run, - user_id="user", - app_type=CreditUsageAppType.WORKFLOW, - trace_session_id="session-1", - ) - init_graph.assert_not_called() - - assert entry_class.call_args.kwargs["graph_runtime_state"] is graph_runtime_state - - -def test_single_node_run_validates_target_node_config(queue_manager: WorkflowAppQueueManager) -> None: - runner = WorkflowBasedAppRunner( - queue_manager=queue_manager, - variable_loader=MagicMock(), - app_id="app", - ) - - workflow = Workflow( - id="workflow", - tenant_id="tenant", - graph=json.dumps( - { - "nodes": [ - { - "id": "loop-node", - "data": {}, - } - ], - "edges": [], - } - ), - ) - - _, _, graph_runtime_state = _make_graph_state() - - with pytest.raises(ValidationError, match="nodes.0.data.type"): - runner._get_graph_and_variable_pool_for_single_node_run( - workflow=workflow, - node_id="loop-node", - user_inputs={}, - graph_runtime_state=graph_runtime_state, - node_type_label="loop", - user_id="00000000-0000-0000-0000-000000000001", - ) - - -def test_single_node_run_rejects_invalid_loop_count(queue_manager: WorkflowAppQueueManager) -> None: - runner = WorkflowBasedAppRunner( - queue_manager=queue_manager, - app_id="app", - ) - - workflow = Workflow( - id="workflow", - tenant_id="tenant", - graph=json.dumps( - { - "nodes": [ - { - "id": "loop-node", - "data": { - "type": "loop", - "title": "Loop", - "loop_count": 0, - "start_node_id": "loop-start", - "break_conditions": [], - "logical_operator": "and", - }, - }, - { - "id": "loop-start", - "data": {"type": "loop-start", "title": "Loop start", "container_id": "loop-node"}, - }, - ], - "edges": [], - } - ), - ) - - _, _, graph_runtime_state = _make_graph_state() - with pytest.raises(ValidationError, match="loop_count"): - runner._get_graph_and_variable_pool_for_single_node_run( - workflow=workflow, - node_id="loop-node", - user_inputs={}, - graph_runtime_state=graph_runtime_state, - node_type_label="loop", - user_id="00000000-0000-0000-0000-000000000001", - ) - - -def test_run_adds_inputs_with_snippet_compatible_start_aliases(queue_manager: WorkflowAppQueueManager) -> None: - app_generate_entity = WorkflowAppGenerateEntity( - app_config=WorkflowUIBasedAppConfig( - app_id="app", tenant_id="tenant", workflow_id="workflow", app_mode=AppMode.WORKFLOW - ), - inputs={"question": "hello"}, - files=[], - user_id="user", - invoke_from=InvokeFrom.SERVICE_API, - workflow_execution_id="execution-id", - task_id="task-id", - stream=True, - call_depth=4, - ) - - workflow = Workflow( - tenant_id="tenant", - app_id="app", - id="workflow", - type="workflow", - version="v1", - graph=json.dumps({"nodes": [], "edges": []}), - kind=WorkflowKind.SNIPPET, - ) - workflow.environment_variables = [] - - runner = WorkflowAppRunner( - application_generate_entity=app_generate_entity, - queue_manager=queue_manager, - variable_loader=MagicMock(), - workflow=workflow, - system_user_id="system-user", - ) - - mock_workflow_entry = MagicMock() - mock_workflow_entry.graph_engine = MagicMock() - mock_workflow_entry.graph_engine.layer = MagicMock() - mock_workflow_entry.run.return_value = iter([]) - - with ( - patch("core.app.apps.workflow.app_runner.RedisChannel"), - patch("core.app.apps.workflow.app_runner.redis_client"), - patch("core.app.apps.workflow.app_runner.WorkflowEntry", return_value=mock_workflow_entry), - patch("core.app.apps.workflow.app_runner.build_system_variables", return_value={}), - patch("core.app.apps.workflow.app_runner.build_bootstrap_variables", return_value=[]), - patch("core.app.apps.workflow.app_runner.add_variables_to_pool"), - patch("core.app.apps.workflow.app_runner.get_default_root_node_id", return_value="root-node"), - patch( - "core.app.apps.workflow.app_runner.get_compatible_start_aliases", return_value=("legacy-start",) - ) as aliases, - patch("core.app.apps.workflow.app_runner.add_node_inputs_to_pool") as add_inputs, - patch.object(runner, "_init_graph", return_value=MagicMock()) as init_graph, - ): - runner.prepare() - - aliases.assert_called_once_with(workflow_kind="snippet", root_node_id="root-node") - add_inputs.assert_called_once() - assert add_inputs.call_args.kwargs["node_id"] == "root-node" - assert add_inputs.call_args.kwargs["inputs"] == {"question": "hello"} - assert add_inputs.call_args.kwargs["aliases"] == ("legacy-start",) - assert init_graph.call_args.kwargs["call_depth"] == 4 diff --git a/api/tests/unit_tests/core/app/apps/test_workflow_human_input_completion.py b/api/tests/unit_tests/core/app/apps/test_workflow_human_input_completion.py deleted file mode 100644 index 654a8338b7a..00000000000 --- a/api/tests/unit_tests/core/app/apps/test_workflow_human_input_completion.py +++ /dev/null @@ -1,463 +0,0 @@ -import json -from collections.abc import Generator, Mapping -from datetime import UTC, datetime, timedelta -from time import perf_counter -from unittest.mock import MagicMock, PropertyMock -from urllib.parse import parse_qs, urlparse - -import pytest -from pytest_mock import MockerFixture -from sqlalchemy import select -from sqlalchemy.orm import Session - -from configs import dify_config -from core.app.app_config.entities import WorkflowUIBasedAppConfig -from core.app.apps.base_app_queue_manager import PublishFrom -from core.app.apps.common import workflow_response_converter -from core.app.apps.workflow.app_queue_manager import WorkflowAppQueueManager -from core.app.apps.workflow_app_runner import PreparedWorkflowRun, WorkflowBasedAppRunner -from core.app.entities.app_invoke_entities import InvokeFrom, UserFrom, WorkflowAppGenerateEntity -from core.app.entities.queue_entities import ( - AppQueueEvent, - QueueHumanInputFormFilledEvent, - QueueHumanInputFormTimeoutEvent, - QueueNodeStartedEvent, - QueueWorkflowStartedEvent, - QueueWorkflowSucceededEvent, -) -from core.db.session_factory import session_factory -from core.repositories.human_input_repository import HumanInputFormRepositoryImpl, HumanInputFormSubmissionRepository -from core.repositories.sqlalchemy_workflow_execution_repository import SQLAlchemyWorkflowExecutionRepository -from core.repositories.sqlalchemy_workflow_node_execution_repository import SQLAlchemyWorkflowNodeExecutionRepository -from core.workflow.node_runtime import DifyFileReferenceFactory -from core.workflow.nodes.human_input.boundary import human_input_container_selector -from core.workflow.nodes.human_input.callback import DifyHITLCallback -from core.workflow.nodes.human_input.entities import ( - FileInputConfig, - FileListInputConfig, - FormDefinition, - HumanInputNodeData, - ParagraphInputConfig, - UserActionConfig, -) -from core.workflow.nodes.human_input.enums import HumanInputFormKind, HumanInputFormStatus -from core.workflow.workflow_entry import WorkflowEntry -from extensions.storage.storage_type import StorageType -from graphon.engine.ready_queue import StartTask -from graphon.engine_events import EngineEvent, GraphRunFailedEvent -from graphon.entities import WorkflowStartReason -from graphon.entities.pause_reason import HitlRequired -from graphon.enums import BuiltinNodeTypes, WorkflowType -from graphon.file.helpers import verify_file_signature -from graphon.graph import Graph -from graphon.nodes.human_input.entities import HumanInputNodeData as GraphonHumanInputNodeData -from graphon.nodes.human_input.human_input_node import HumanInputNode -from graphon.runtime import RuntimeState, VariablePool -from graphon.runtime.execution import ROOT_FRAME_ID -from models import Account, UploadFile -from models.enums import CreatorUserRole, WorkflowRunTriggeredFrom -from models.execution_extra_content import HumanInputContent -from models.human_input import HumanInputForm -from models.model import AppMode -from models.workflow import WorkflowNodeExecutionTriggeredFrom -from services.workflow_run_agg import WorkflowRunAgg -from tests.unit_tests.core.app.apps.advanced_chat.test_generate_task_pipeline import ( - _build_generate_entity, - _build_pipeline, -) -from tests.unit_tests.core.app.apps.common.test_workflow_response_converter_human_input import _build_converter -from tests.workflow_test_utils import build_test_graph_init_params, build_test_run_context - - -@pytest.fixture -def workflow_queue(mocker: MockerFixture) -> tuple[WorkflowAppQueueManager, MagicMock]: - queue = WorkflowAppQueueManager("task", "user", InvokeFrom.DEBUGGER, AppMode.WORKFLOW) - return queue, mocker.spy(queue, "publish") - - -def _save_form(session: Session, *, status: HumanInputFormStatus) -> HumanInputForm: - now = datetime.now(UTC) - definition = FormDefinition( - form_content="Decision: {{#$output.answer#}}", - rendered_content="Decision: {{#$output.answer#}}", - inputs=[ParagraphInputConfig(output_variable_name="answer")], - user_actions=[UserActionConfig(id="approve", title="Approve")], - expiration_time=now + timedelta(hours=1), - node_title="Approval", - ) - form = HumanInputForm( - tenant_id="tenant", - app_id="app", - workflow_run_id="run-1", - node_id="same-human-node", - form_kind=HumanInputFormKind.RUNTIME, - form_definition=definition.model_dump_json(), - rendered_content=definition.rendered_content, - expiration_time=definition.expiration_time, - status=status, - selected_action_id="approve" if status == HumanInputFormStatus.SUBMITTED else None, - submitted_data=json.dumps({"answer": "approved"}) if status == HumanInputFormStatus.SUBMITTED else None, - submitted_at=now if status == HumanInputFormStatus.SUBMITTED else None, - created_at=now, - ) - session.add(form) - session.commit() - return form - - -def _make_paused_workflow( - runner: WorkflowBasedAppRunner, forms: list[HumanInputForm], *, human_input_form: HumanInputForm | None = None -) -> WorkflowEntry: - state = RuntimeState(workflow_id="workflow", variable_pool=VariablePool(), start_at=perf_counter()) - state.variable_pool.add(("sys", "workflow_run_id"), "run-1") - state.variable_pool.add(("sys", "app_id"), "app") - graph = runner._init_graph( - graph_config={ - "nodes": [{"id": "start", "data": {"type": "start", "title": "Start", "variables": []}}], - "edges": [], - }, - graph_runtime_state=state, - tenant_id="tenant", - user_id="user", - user_from=UserFrom.END_USER, - invoke_from=InvokeFrom.WEB_APP, - call_depth=0, - ) - if human_input_form is not None: - node_data = HumanInputNodeData( - title="Approval", - form_content="Decision: {{#$output.answer#}}", - inputs=[ParagraphInputConfig(output_variable_name="answer")], - user_actions=[UserActionConfig(id="approve", title="Approve")], - ) - callback = DifyHITLCallback( - form_repository=HumanInputFormRepositoryImpl( - tenant_id="tenant", app_id="app", workflow_execution_id="run-1" - ), - node_data=node_data, - execution_id_getter=lambda: human_input_form.id, - ) - graph = ( - Graph.new() - .add_root(graph.root_node) - .add_node( - HumanInputNode( - node_id=human_input_form.node_id, - data=GraphonHumanInputNodeData.model_validate(node_data.model_dump()), - init_params=build_test_graph_init_params(tenant_id="tenant", app_id="app"), - runtime_state=state, - hitl_callback=callback, - ), - from_node_id=graph.root_node.id, - ) - .build() - ) - # A resumed Human Input node keeps the execution identity used to create its form. - state.graph_execution.get_or_create_node_execution( - frame_id=ROOT_FRAME_ID, node_id=human_input_form.node_id - ).execution_id = human_input_form.id - state.graph_execution.start() - for form in forms: - state.graph_execution.pause(HitlRequired(session_id=form.id, node_id=form.node_id, node_title="Approval")) - state.variable_pool.add(human_input_container_selector(form.id), "outer-tool") - state.defer_ready_task(StartTask(frame_id=ROOT_FRAME_ID, node_id=graph.root_node.id)) - return WorkflowEntry( - graph=graph, - graph_config=graph.graph_config or {}, - call_depth=0, - graph_runtime_state=state, - ) - - -def _resume_events(runner: WorkflowBasedAppRunner, entry: WorkflowEntry) -> Generator[EngineEvent]: - """Resume through orchestration with actual repositories and the real Engine.""" - entity = WorkflowAppGenerateEntity( - task_id="task", - app_config=WorkflowUIBasedAppConfig( - tenant_id="tenant", - app_id="app", - app_mode=AppMode.WORKFLOW, - workflow_id="workflow", - ), - inputs={}, - files=[], - user_id="user", - stream=True, - invoke_from=InvokeFrom.DEBUGGER, - workflow_execution_id="run-1", - ) - user = Account(name="Tester", email="tester@example.com") - user.id = "user" - factory = session_factory.get_session_maker() - runs = SQLAlchemyWorkflowExecutionRepository( - factory, - tenant_id="tenant", - app_id="app", - user=user, - triggered_from=WorkflowRunTriggeredFrom.DEBUGGING, - ) - nodes = SQLAlchemyWorkflowNodeExecutionRepository( - factory, - tenant_id="tenant", - app_id="app", - user=user, - triggered_from=WorkflowNodeExecutionTriggeredFrom.WORKFLOW_RUN, - ) - prepared = PreparedWorkflowRun(entry, WorkflowType.WORKFLOW, "draft", entity, ()) - return WorkflowRunAgg( - prepared, - runner, - None, - workflow_execution_repository=runs, - workflow_node_execution_repository=nodes, - ).iter_events() - - -@pytest.mark.parametrize("status", [HumanInputFormStatus.SUBMITTED, HumanInputFormStatus.TIMEOUT]) -@pytest.mark.parametrize("runs_human_input_node", [False, True], ids=["hidden-form", "direct-completion"]) -def test_resume_publishes_only_the_completed_form_among_repeated_node_invocations( - sqlite_session: Session, - workflow_queue: tuple[WorkflowAppQueueManager, MagicMock], - status: HumanInputFormStatus, - runs_human_input_node: bool, -) -> None: - selected = _save_form(sqlite_session, status=status) - waiting = _save_form(sqlite_session, status=HumanInputFormStatus.WAITING) - expired = _save_form(sqlite_session, status=HumanInputFormStatus.EXPIRED) - _save_form(sqlite_session, status=HumanInputFormStatus.SUBMITTED) - queue, publish_spy = workflow_queue - runner = WorkflowBasedAppRunner(queue_manager=queue, app_id="app") - entry = _make_paused_workflow( - runner, [selected, waiting, selected, expired], human_input_form=selected if runs_human_input_node else None - ) - original_reasons = tuple(entry.graph_engine.runtime_state.graph_execution.pause_reasons) - - for event in _resume_events(runner, entry): - runner.handle_event(entry, event) - - published = [call.args[0] for call in publish_spy.call_args_list] - assert isinstance(published[0], QueueWorkflowStartedEvent) - completions = [ - item for item in published if isinstance(item, QueueHumanInputFormFilledEvent | QueueHumanInputFormTimeoutEvent) - ] - assert len(completions) == 1 - completion = completions[0] - assert completion.form_id == selected.id - assert completion.node_id == "outer-tool" - if status == HumanInputFormStatus.SUBMITTED: - assert isinstance(completion, QueueHumanInputFormFilledEvent) - assert (completion.action_id, completion.action_text, completion.rendered_content) == ( - "approve", - "Approve", - "Decision: approved", - ) - assert completion.submitted_data is not None - assert completion.submitted_data["answer"].to_object() == "approved" - else: - assert isinstance(completion, QueueHumanInputFormTimeoutEvent) - assert len(original_reasons) == 4 - assert all(isinstance(reason, HitlRequired) and reason.node_id == "same-human-node" for reason in original_reasons) - assert entry.graph_engine.runtime_state.graph_execution.pause_reasons == [] - - -@pytest.mark.parametrize("owner_field", ["tenant_id", "app_id", "workflow_run_id"]) -def test_resume_rejects_form_outside_the_trusted_execution_owner( - sqlite_session: Session, workflow_queue: tuple[WorkflowAppQueueManager, MagicMock], owner_field: str -) -> None: - form = _save_form(sqlite_session, status=HumanInputFormStatus.SUBMITTED) - setattr(form, owner_field, "another-owner") - sqlite_session.commit() - queue, publish_spy = workflow_queue - runner = WorkflowBasedAppRunner(queue_manager=queue, app_id="app") - entry = _make_paused_workflow(runner, [form]) - - events = list(_resume_events(runner, entry)) - assert isinstance(events[-1], GraphRunFailedEvent) - assert "does not belong" in events[-1].error - - assert not any(isinstance(call.args[0], QueueHumanInputFormFilledEvent) for call in publish_spy.call_args_list) - - -def test_resume_refreshes_expired_file_and_file_list_urls_before_form_completion( - sqlite_session: Session, - workflow_queue: tuple[WorkflowAppQueueManager, MagicMock], - monkeypatch: pytest.MonkeyPatch, - mocker: MockerFixture, -) -> None: - mocker.patch.object( - type(workflow_response_converter.db), - "engine", - new_callable=PropertyMock, - return_value=sqlite_session.get_bind(), - ) - submitted_at = 1_700_000_000 - monkeypatch.setattr("core.app.workflow.file_runtime.time.time", lambda: submitted_at) - factory = DifyFileReferenceFactory(build_test_run_context()) - saved_files = [] - for name in ("first.txt", "second.txt"): - upload = UploadFile( - tenant_id="tenant", - storage_type=StorageType.LOCAL, - key=name, - name=name, - size=12, - extension="txt", - mime_type="text/plain", - created_by="user", - created_by_role=CreatorUserRole.END_USER, - created_at=datetime.now(UTC), - used=False, - ) - sqlite_session.add(upload) - sqlite_session.commit() - saved_files.append( - factory.build_from_mapping(mapping={"transfer_method": "local_file", "upload_file_id": upload.id}).to_dict() - ) - - def url_is_valid(file: Mapping[str, object]) -> bool: - url, file_id = file["url"], file["related_id"] - assert isinstance(url, str) - assert isinstance(file_id, str) - query = parse_qs(urlparse(url).query) - return verify_file_signature( - upload_file_id=file_id, - timestamp=query["timestamp"][0], - nonce=query["nonce"][0], - sign=query["sign"][0], - ) - - assert all(url_is_valid(file) for file in saved_files) - form = _save_form(sqlite_session, status=HumanInputFormStatus.SUBMITTED) - definition = FormDefinition.model_validate_json(form.form_definition) - definition.inputs.extend( - [FileInputConfig(output_variable_name="attachment"), FileListInputConfig(output_variable_name="attachments")] - ) - form.form_definition = definition.model_dump_json() - form.submitted_data = json.dumps({"answer": "approved", "attachment": saved_files[0], "attachments": saved_files}) - sqlite_session.commit() - - monkeypatch.setattr( - "core.app.workflow.file_runtime.time.time", lambda: submitted_at + dify_config.FILES_ACCESS_TIMEOUT + 1 - ) - assert not any(url_is_valid(file) for file in saved_files) - queue, publish_spy = workflow_queue - runner = WorkflowBasedAppRunner(queue_manager=queue, app_id="app") - entry = _make_paused_workflow(runner, [form]) - for event in _resume_events(runner, entry): - runner.handle_event(entry, event) - - completions = [ - call.args[0] for call in publish_spy.call_args_list if isinstance(call.args[0], QueueHumanInputFormFilledEvent) - ] - assert len(completions) == 1 - converter = _build_converter() - converter.workflow_start_to_stream_response( - task_id="task", workflow_run_id="run-1", workflow_id="workflow", reason=WorkflowStartReason.RESUMPTION - ) - response = converter.human_input_form_filled_to_stream_response(event=completions[0], task_id="task") - data = response.data.submitted_data - assert data is not None - assert data["answer"] == "approved" - assert response.data.rendered_content == "Decision: approved" - restored_files = [data["attachment"], *data["attachments"]] - for restored, saved in zip(restored_files, [saved_files[0], *saved_files], strict=True): - assert (restored["related_id"], restored["filename"]) == (saved["related_id"], saved["filename"]) - assert url_is_valid(restored) - - -@pytest.mark.parametrize("late_status", [HumanInputFormStatus.SUBMITTED, HumanInputFormStatus.TIMEOUT]) -def test_form_completed_during_resume_is_published_once_before_the_terminal_event( - sqlite_session: Session, - workflow_queue: tuple[WorkflowAppQueueManager, MagicMock], - late_status: HumanInputFormStatus, -) -> None: - first = _save_form(sqlite_session, status=HumanInputFormStatus.SUBMITTED) - second = _save_form(sqlite_session, status=HumanInputFormStatus.WAITING) - queue, publish_spy = workflow_queue - runner = WorkflowBasedAppRunner(queue_manager=queue, app_id="app") - entry = _make_paused_workflow(runner, [first, second]) - - def complete_second_form_during_execution(event: AppQueueEvent, publish_from: PublishFrom) -> None: - WorkflowAppQueueManager.publish(queue, event, publish_from) - # The second submission arrives after the initial completion query, - # while this real Engine attempt is already processing node events. - if not isinstance(event, QueueNodeStartedEvent): - return - repository = HumanInputFormSubmissionRepository() - if late_status == HumanInputFormStatus.SUBMITTED: - repository.mark_submitted( - form_id=second.id, - recipient_id=None, - selected_action_id="approve", - form_data={"answer": "approved"}, - submission_user_id=None, - submission_end_user_id="user", - ) - else: - repository.mark_timeout(form_id=second.id, timeout_status=late_status) - - publish_spy.side_effect = complete_second_form_during_execution - for event in _resume_events(runner, entry): - runner.handle_event(entry, event) - - published = [call.args[0] for call in publish_spy.call_args_list] - completions = [ - event - for event in published - if isinstance(event, QueueHumanInputFormFilledEvent | QueueHumanInputFormTimeoutEvent) - ] - assert [event.form_id for event in completions] == [first.id, second.id] - assert isinstance(published[0], QueueWorkflowStartedEvent) - assert isinstance(published[-1], QueueWorkflowSucceededEvent) - assert published.index(completions[-1]) < len(published) - 1 - - -def test_waiting_form_expiring_during_resume_publishes_timeout_before_success( - sqlite_session: Session, workflow_queue: tuple[WorkflowAppQueueManager, MagicMock] -) -> None: - form = _save_form(sqlite_session, status=HumanInputFormStatus.WAITING) - form.expiration_time = datetime.now(UTC) - timedelta(seconds=1) - sqlite_session.commit() - queue, publish_spy = workflow_queue - runner = WorkflowBasedAppRunner(queue_manager=queue, app_id="app") - entry = _make_paused_workflow(runner, [form], human_input_form=form) - - for event in _resume_events(runner, entry): - runner.handle_event(entry, event) - - published = [call.args[0] for call in publish_spy.call_args_list] - assert isinstance(published[-1], QueueWorkflowSucceededEvent) - timeouts = [event for event in published if isinstance(event, QueueHumanInputFormTimeoutEvent)] - assert [event.form_id for event in timeouts] == [form.id] - assert timeouts[0].node_id == "outer-tool" - assert published.index(timeouts[0]) < len(published) - 1 - sqlite_session.refresh(form) - assert form.status == HumanInputFormStatus.TIMEOUT - - -def test_chat_history_keeps_repeated_forms_by_form_id(sqlite_session: Session) -> None: - pipeline = _build_pipeline() - pipeline._workflow_tenant_id = "tenant" - pipeline._application_generate_entity = _build_generate_entity() - pipeline._workflow_response_converter = _build_converter() - pipeline._workflow_response_converter.workflow_start_to_stream_response( - task_id="task-1", workflow_run_id="run-1", workflow_id="workflow-1", reason=WorkflowStartReason.INITIAL - ) - for form_id in ("first-form", "second-form", "first-form"): - event = QueueHumanInputFormFilledEvent( - form_id=form_id, - node_id="same-human-node", - node_type=BuiltinNodeTypes.HUMAN_INPUT, - node_title="Approval", - rendered_content="Decision: approved", - action_id="approve", - action_text="Approve", - ) - list(pipeline._handle_human_input_form_filled_event(event)) - - contents = sqlite_session.scalars(select(HumanInputContent)).all() - assert {(content.form_id, content.message_id) for content in contents} == { - ("first-form", "message-1"), - ("second-form", "message-1"), - } - assert len(contents) == 2 diff --git a/api/tests/unit_tests/core/app/apps/test_workflow_pause_events.py b/api/tests/unit_tests/core/app/apps/test_workflow_pause_events.py deleted file mode 100644 index 16394e5e7ba..00000000000 --- a/api/tests/unit_tests/core/app/apps/test_workflow_pause_events.py +++ /dev/null @@ -1,383 +0,0 @@ -import json -from datetime import UTC, datetime -from types import SimpleNamespace -from unittest.mock import MagicMock - -import pytest -from sqlalchemy.orm import Session - -from core.app.app_config.entities import WorkflowUIBasedAppConfig -from core.app.apps.common import workflow_response_converter -from core.app.apps.common.workflow_response_converter import WorkflowResponseConverter -from core.app.apps.workflow.app_runner import WorkflowAppRunner -from core.app.entities.app_invoke_entities import InvokeFrom, WorkflowAppGenerateEntity -from core.app.entities.queue_entities import QueueWorkflowPausedEvent -from core.app.entities.task_entities import HumanInputRequiredResponse, WorkflowPauseStreamResponse -from core.workflow.nodes.human_input.entities import ( - ParagraphInputConfig, - SelectInputConfig, - StringListSource, - UserActionConfig, -) -from core.workflow.nodes.human_input.enums import ValueSourceType -from core.workflow.nodes.human_input.pause_reason import HumanInputRequired -from core.workflow.system_variables import build_system_variables -from graphon.engine_events import GraphRunPausedEvent -from graphon.entities import WorkflowStartReason -from graphon.entities.pause_reason import HitlRequired -from graphon.runtime import RuntimeState, VariablePool -from models.account import Account -from models.human_input import HumanInputForm, HumanInputFormRecipient, RecipientType -from models.model import AppMode -from models.workflow import Workflow, WorkflowType -from tests.unit_tests.core.app.apps.test_workflow_app_runner_core import _workflow_entry - - -class _RecordingWorkflowAppRunner(WorkflowAppRunner): - def __init__(self, **kwargs): - super().__init__(**kwargs) - self.published_events = [] - - def _publish_event(self, event): - self.published_events.append(event) - - -@pytest.fixture -def sqlite_pause_session(sqlite_session: Session, monkeypatch: pytest.MonkeyPatch) -> Session: - """Bind pause-response queries to the shared SQLite session's database.""" - monkeypatch.setattr(workflow_response_converter, "db", SimpleNamespace(engine=sqlite_session.get_bind())) - return sqlite_session - - -def _persist_human_input_form( - session: Session, - *, - recipients: list[tuple[RecipientType, str]] | None = None, -) -> datetime: - expiration_time = datetime(2024, 1, 1, tzinfo=UTC) - form = HumanInputForm( - id="form-1", - tenant_id="tenant-id", - app_id="app-id", - workflow_run_id="run-id", - node_id="node-id", - form_definition='{"display_in_ui": true}', - rendered_content="Rendered", - expiration_time=expiration_time, - ) - recipient_models = [ - HumanInputFormRecipient( - id=f"recipient-{index}", - form_id=form.id, - delivery_id=f"delivery-{index}", - recipient_type=recipient_type, - recipient_payload="{}", - access_token=access_token, - ) - for index, (recipient_type, access_token) in enumerate(recipients or ()) - ] - session.add(form) - session.add_all(recipient_models) - session.commit() - return expiration_time - - -def _build_runner(): - app_entity = SimpleNamespace( - app_config=SimpleNamespace(app_id="app-id"), - inputs={}, - files=[], - invoke_from=InvokeFrom.SERVICE_API, - single_iteration_run=None, - single_loop_run=None, - workflow_execution_id="run-id", - user_id="user-id", - ) - workflow = Workflow.new( - tenant_id="tenant-id", - app_id="app-id", - type=WorkflowType.WORKFLOW, - version=Workflow.VERSION_DRAFT, - graph=json.dumps({}), - features="{}", - created_by="account-id", - environment_variables=[], - conversation_variables=[], - rag_pipeline_variables=[], - ) - workflow.id = "workflow-id" - queue_manager = SimpleNamespace(publish=lambda event, pub_from: None) - return _RecordingWorkflowAppRunner( - application_generate_entity=app_entity, - queue_manager=queue_manager, - variable_loader=MagicMock(), - workflow=workflow, - system_user_id="sys-user", - root_node_id=None, - graph_engine_layers=(), - graph_runtime_state=None, - ) - - -@pytest.mark.parametrize("source_node_id", ["node-human", "nested-human"]) -def test_graph_run_paused_event_uses_projected_nodes_for_highlighting( - monkeypatch: pytest.MonkeyPatch, source_node_id: str -): - runner = _build_runner() - graph_reason = HitlRequired( - session_id="form-1", - node_id=source_node_id, - node_title="Human Step", - ) - event = GraphRunPausedEvent(reasons=[graph_reason], outputs={"foo": "bar"}) - workflow_entry = _workflow_entry(runner) - - enriched_reason = HumanInputRequired( - form_id="form-1", - form_content="content", - inputs=[], - actions=[], - node_id="node-human", - node_title="Human Step", - ) - monkeypatch.setattr("core.app.apps.workflow_app_runner.dispatch_human_input_email_task", MagicMock()) - - runner.handle_event(workflow_entry, event, pause_reasons=[enriched_reason]) - - assert len(runner.published_events) == 1 - queue_event = runner.published_events[0] - assert isinstance(queue_event, QueueWorkflowPausedEvent) - assert queue_event.reasons == [enriched_reason] - assert queue_event.outputs == {"foo": "bar"} - assert queue_event.paused_nodes == ["node-human"] - assert graph_reason.node_id == source_node_id - - -def _build_converter(*, invoke_from: InvokeFrom = InvokeFrom.SERVICE_API) -> WorkflowResponseConverter: - application_generate_entity = WorkflowAppGenerateEntity( - task_id="task", - workflow_execution_id="run-id", - user_id="user", - stream=True, - inputs={}, - files=[], - invoke_from=invoke_from, - app_config=WorkflowUIBasedAppConfig( - app_id="app-id", tenant_id="tenant-id", workflow_id="workflow-id", app_mode=AppMode.WORKFLOW - ), - ) - system_variables = build_system_variables( - user_id="user", - app_id="app-id", - workflow_id="workflow-id", - workflow_execution_id="run-id", - ) - user = Account(name="Tester", email="tester@example.com") - user.id = "account-id" - return WorkflowResponseConverter( - application_generate_entity=application_generate_entity, - user=user, - system_variables=system_variables, - ) - - -@pytest.mark.parametrize( - "sqlite_session", - [(HumanInputForm, HumanInputFormRecipient)], - indirect=True, -) -def test_queue_workflow_paused_event_to_stream_responses(sqlite_pause_session: Session): - converter = _build_converter() - converter.workflow_start_to_stream_response( - task_id="task", - workflow_run_id="run-id", - workflow_id="workflow-id", - reason=WorkflowStartReason.INITIAL, - ) - - expiration_time = _persist_human_input_form( - sqlite_pause_session, - recipients=[ - (RecipientType.CONSOLE, "console-token"), - (RecipientType.BACKSTAGE, "backstage-token"), - ], - ) - - reason = HumanInputRequired( - form_id="form-1", - form_content="Rendered", - inputs=[ParagraphInputConfig(output_variable_name="field")], - actions=[UserActionConfig(id="approve", title="Approve")], - node_id="node-id", - node_title="Human Step", - ) - queue_event = QueueWorkflowPausedEvent( - reasons=[reason], - outputs={"answer": "value"}, - paused_nodes=["node-id"], - ) - - runtime_state = RuntimeState(workflow_id="test-workflow", variable_pool=VariablePool(), start_at=0.0) - responses = converter.workflow_pause_to_stream_response( - event=queue_event, - task_id="task", - graph_runtime_state=runtime_state, - ) - - assert isinstance(responses[-1], WorkflowPauseStreamResponse) - pause_resp = responses[-1] - assert pause_resp.workflow_run_id == "run-id" - assert pause_resp.data.paused_nodes == ["node-id"] - assert pause_resp.data.outputs == {} - assert pause_resp.data.reasons[0]["form_id"] == "form-1" - - assert isinstance(responses[0], HumanInputRequiredResponse) - hi_resp = responses[0] - assert hi_resp.data.form_id == "form-1" - assert hi_resp.data.node_id == "node-id" - assert hi_resp.data.node_title == "Human Step" - assert hi_resp.data.inputs[0].output_variable_name == "field" - assert hi_resp.data.actions[0].id == "approve" - assert hi_resp.data.display_in_ui is True - assert hi_resp.data.form_token is None - assert hi_resp.data.approval_channels == ["console"] - assert hi_resp.data.expiration_time == int(expiration_time.timestamp()) - - -def _build_paused_human_input_response( - session: Session, - recipients: list[tuple[RecipientType, str]], -): - """Drive the live OPENAPI pause path with persisted forms and recipients.""" - converter = _build_converter(invoke_from=InvokeFrom.OPENAPI) - converter.workflow_start_to_stream_response( - task_id="task", - workflow_run_id="run-id", - workflow_id="workflow-id", - reason=WorkflowStartReason.INITIAL, - ) - - _persist_human_input_form(session, recipients=recipients) - - reason = HumanInputRequired( - form_id="form-1", - form_content="Rendered", - inputs=[ParagraphInputConfig(output_variable_name="field")], - actions=[UserActionConfig(id="approve", title="Approve")], - node_id="node-id", - node_title="Human Step", - ) - queue_event = QueueWorkflowPausedEvent( - reasons=[reason], - outputs={}, - paused_nodes=["node-id"], - ) - - runtime_state = RuntimeState(workflow_id="test-workflow", variable_pool=VariablePool(), start_at=0.0) - responses = converter.workflow_pause_to_stream_response( - event=queue_event, - task_id="task", - graph_runtime_state=runtime_state, - ) - assert isinstance(responses[0], HumanInputRequiredResponse) - return responses - - -@pytest.mark.parametrize( - "sqlite_session", - [(HumanInputForm, HumanInputFormRecipient)], - indirect=True, -) -def test_openapi_pause_without_web_app_recipient_emits_approval_channels(sqlite_pause_session: Session): - responses = _build_paused_human_input_response( - sqlite_pause_session, - recipients=[ - (RecipientType.EMAIL_MEMBER, "email-token"), - (RecipientType.BACKSTAGE, "backstage-token"), - ], - ) - - hi_resp = responses[0] - assert hi_resp.data.form_token is None - assert hi_resp.data.approval_channels == ["console", "email"] - - pause_resp = responses[-1] - assert pause_resp.data.reasons[0]["approval_channels"] == ["console", "email"] - - -@pytest.mark.parametrize( - "sqlite_session", - [(HumanInputForm, HumanInputFormRecipient)], - indirect=True, -) -def test_openapi_pause_with_web_app_recipient_sets_token_and_channels(sqlite_pause_session: Session): - responses = _build_paused_human_input_response( - sqlite_pause_session, - recipients=[ - (RecipientType.STANDALONE_WEB_APP, "web-app-token"), - (RecipientType.BACKSTAGE, "backstage-token"), - ], - ) - - hi_resp = responses[0] - assert hi_resp.data.form_token == "web-app-token" - assert hi_resp.data.approval_channels == ["console"] - - pause_resp = responses[-1] - assert pause_resp.data.reasons[0]["approval_channels"] == ["console"] - - -@pytest.mark.parametrize( - "sqlite_session", - [(HumanInputForm, HumanInputFormRecipient)], - indirect=True, -) -def test_queue_workflow_paused_event_resolves_variable_select_options(sqlite_pause_session: Session): - converter = _build_converter() - converter.workflow_start_to_stream_response( - task_id="task", - workflow_run_id="run-id", - workflow_id="workflow-id", - reason=WorkflowStartReason.INITIAL, - ) - - _persist_human_input_form(sqlite_pause_session) - - reason = HumanInputRequired( - form_id="form-1", - form_content="Rendered", - inputs=[ - SelectInputConfig( - output_variable_name="decision", - option_source=StringListSource( - type=ValueSourceType.VARIABLE, - selector=["start", "options"], - value=[], - ), - ) - ], - actions=[UserActionConfig(id="approve", title="Approve")], - node_id="node-id", - node_title="Human Step", - ) - queue_event = QueueWorkflowPausedEvent( - reasons=[reason], - outputs={}, - paused_nodes=["node-id"], - ) - - runtime_state = RuntimeState(workflow_id="test-workflow", variable_pool=VariablePool(), start_at=0.0) - runtime_state.variable_pool.add(("start", "options"), ["approve", "reject"]) - responses = converter.workflow_pause_to_stream_response( - event=queue_event, - task_id="task", - graph_runtime_state=runtime_state, - ) - - assert isinstance(responses[0], HumanInputRequiredResponse) - hi_resp = responses[0] - assert hi_resp.data.inputs[0].option_source.value == ["approve", "reject"] - - assert isinstance(responses[-1], WorkflowPauseStreamResponse) - pause_resp = responses[-1] - assert pause_resp.data.reasons[0]["inputs"][0]["option_source"]["value"] == ["approve", "reject"] diff --git a/api/tests/unit_tests/core/app/apps/workflow/test_app_config_manager.py b/api/tests/unit_tests/core/app/apps/workflow/test_app_config_manager.py index 3606bcb72d3..ace1a934d6b 100644 --- a/api/tests/unit_tests/core/app/apps/workflow/test_app_config_manager.py +++ b/api/tests/unit_tests/core/app/apps/workflow/test_app_config_manager.py @@ -3,7 +3,7 @@ from unittest.mock import patch from core.app.apps.workflow.app_config_manager import WorkflowAppConfigManager from models.model import App, AppMode from models.workflow import Workflow -from tests.unit_tests.model_factories import make_app, make_workflow +from tests.helpers.unit.model_factories import make_app, make_workflow def _app() -> App: diff --git a/api/tests/unit_tests/core/app/apps/workflow/test_app_generator.py b/api/tests/unit_tests/core/app/apps/workflow/test_app_generator.py new file mode 100644 index 00000000000..98b8d0411ae --- /dev/null +++ b/api/tests/unit_tests/core/app/apps/workflow/test_app_generator.py @@ -0,0 +1,1331 @@ +from __future__ import annotations + +import contextlib +import contextvars +import json +from types import SimpleNamespace +from unittest.mock import MagicMock, Mock + +import pytest +from flask import Flask +from sqlalchemy import inspect +from sqlalchemy.orm import Session, sessionmaker + +import core.app.apps.workflow.app_generator as _workflow_app_generator +from core.app.app_config.entities import AppAdditionalFeatures, WorkflowUIBasedAppConfig +from core.app.apps.draft_variable_saver import DraftVariableSaverFactory +from core.app.apps.exc import GenerateTaskStoppedError +from core.app.apps.workflow import app_generator as _app_generator_extra_app_generator_module +from core.app.apps.workflow.app_generator import WorkflowAppGenerator +from core.app.apps.workflow.app_queue_manager import WorkflowAppQueueManager +from core.app.entities.app_invoke_entities import InvokeFrom, WorkflowAppGenerateEntity +from core.app.entities.queue_entities import QueueWorkflowSucceededEvent +from core.app.layers.pause_state_persist_layer import PauseStateLayerConfig +from core.ops.ops_trace_manager import OpsTraceManager, TraceQueueManager +from core.repositories import SQLAlchemyWorkflowExecutionRepository, SQLAlchemyWorkflowNodeExecutionRepository +from graphon.enums import WorkflowExecutionStatus +from graphon.runtime import RuntimeState, VariablePool +from graphon.variable_loader import DUMMY_VARIABLE_LOADER +from models.enums import CreatorUserRole, EndUserType, WorkflowRunTriggeredFrom +from models.model import App, AppMode, EndUser +from models.snippet import CustomizedSnippet +from models.workflow import Workflow, WorkflowKind, WorkflowNodeExecutionTriggeredFrom, WorkflowRun, WorkflowType +from services.workflow_run_agg import WorkflowRunAgg +from tests.helpers.unit.model_factories import make_workflow + +"""SQLite-backed tests for workflow app generation and worker reload behavior.""" + + +def _workflow_app_generator_workflow( + *, + workflow_id: str = "workflow", + app_id: str = "app", + tenant_id: str = "tenant", + kind: WorkflowKind = WorkflowKind.STANDARD, +) -> Workflow: + return make_workflow( + workflow_id=workflow_id, + tenant_id=tenant_id, + app_id=app_id, + kind=kind, + version="1", + created_by="creator", + graph={"nodes": [{"id": "start", "data": {"type": "start", "title": "Start", "variables": []}}], "edges": []}, + ) + + +def _persist_generator_rows(sqlite_session: Session) -> tuple[App, Workflow, EndUser]: + workflow = _workflow_app_generator_workflow() + app = App( + id="app", + tenant_id="tenant", + name="Workflow app", + description="", + mode=AppMode.WORKFLOW, + workflow_id=workflow.id, + enable_site=True, + enable_api=True, + max_active_requests=0, + ) + end_user = EndUser( + id="user", + tenant_id="tenant", + app_id=app.id, + type=EndUserType.SERVICE_API, + name="End user", + session_id="end-user-session", + ) + sqlite_session.add_all([workflow, app, end_user]) + sqlite_session.commit() + return app, workflow, end_user + + +def _app_config(app: App, workflow: Workflow) -> WorkflowUIBasedAppConfig: + return WorkflowUIBasedAppConfig( + tenant_id=app.tenant_id, + app_id=app.id, + app_mode=app.mode, + workflow_id=workflow.id, + ) + + +def _generate_entity( + app: App, + workflow: Workflow, + end_user: EndUser, + *, + invoke_from: InvokeFrom = InvokeFrom.SERVICE_API, + stream: bool = True, +) -> WorkflowAppGenerateEntity: + return WorkflowAppGenerateEntity( + task_id="task", + app_config=_app_config(app, workflow), + inputs={}, + files=[], + user_id=end_user.id, + stream=stream, + invoke_from=invoke_from, + trace_manager=TraceQueueManager(app_id=app.id, user_id=end_user.id), + workflow_execution_id="run", + ) + + +def _runtime_state() -> RuntimeState: + return RuntimeState(workflow_id="test-workflow", variable_pool=VariablePool(), start_at=0.0) + + +def _repositories( + sqlite_session_factory: sessionmaker[Session], app: App, end_user: EndUser +) -> tuple[SQLAlchemyWorkflowExecutionRepository, SQLAlchemyWorkflowNodeExecutionRepository]: + return ( + SQLAlchemyWorkflowExecutionRepository( + session_factory=sqlite_session_factory, + tenant_id=app.tenant_id, + user=end_user, + app_id=app.id, + triggered_from=WorkflowRunTriggeredFrom.APP_RUN, + ), + SQLAlchemyWorkflowNodeExecutionRepository( + session_factory=sqlite_session_factory, + tenant_id=app.tenant_id, + user=end_user, + app_id=app.id, + triggered_from=WorkflowNodeExecutionTriggeredFrom.WORKFLOW_RUN, + ), + ) + + +class TestWorkflowAppGeneratorCases: + @staticmethod + @pytest.fixture(autouse=True) + def disable_trace_delivery(monkeypatch: pytest.MonkeyPatch) -> None: + """Construct real trace managers without provider lookup or background delivery.""" + monkeypatch.setattr(OpsTraceManager, "get_ops_trace_instance", lambda app_id: None) + monkeypatch.setattr(TraceQueueManager, "start_timer", lambda self: None) + monkeypatch.setattr("core.telemetry.gateway.is_enterprise_telemetry_enabled", lambda: False) + + @staticmethod + def test_worker_executes_and_persists_a_workflow_through_the_injected_driver( + app: Flask, sqlite_session: Session, sqlite_session_factory: sessionmaker[Session] + ) -> None: + app_model, workflow, end_user = _persist_generator_rows(sqlite_session) + workflow.graph = json.dumps( + { + "nodes": [ + {"id": "start", "data": {"type": "start", "title": "Start", "variables": []}}, + { + "id": "end", + "data": { + "type": "end", + "title": "End", + "outputs": [{"variable": "caller", "value_selector": ["sys", "user_id"]}], + }, + }, + ], + "edges": [{"id": "start-end", "source": "start", "target": "end"}], + } + ) + sqlite_session.commit() + entity = WorkflowAppGenerateEntity( + task_id="worker-task", + app_config=_app_config(app_model, workflow), + inputs={}, + files=[], + user_id=end_user.id, + stream=True, + invoke_from=InvokeFrom.SERVICE_API, + workflow_execution_id="worker-run", + ) + queue = WorkflowAppQueueManager(entity.task_id, end_user.id, entity.invoke_from, AppMode.WORKFLOW) + execution_repository, node_repository = _repositories(sqlite_session_factory, app_model, end_user) + + WorkflowAppGenerator(execution_driver=WorkflowRunAgg.run)._generate_worker( + flask_app=app, + application_generate_entity=entity, + queue_manager=queue, + context=contextvars.copy_context(), + variable_loader=DUMMY_VARIABLE_LOADER, + workflow_execution_repository=execution_repository, + workflow_node_execution_repository=node_repository, + ) + + events = [message.event for message in queue.listen()] + assert isinstance(events[-1], QueueWorkflowSucceededEvent), events + assert events[-1].outputs == {"caller": end_user.session_id} + run = sqlite_session.get(WorkflowRun, entity.workflow_execution_id) + assert run is not None + assert run.status == "succeeded" + assert run.outputs_dict == {"caller": end_user.session_id} + + @staticmethod + def test_ensure_snippet_start_node_in_worker_returns_standard_workflow_without_lookup( + sqlite_session: Session, + ) -> None: + workflow = _workflow_app_generator_workflow() + sqlite_session.add(workflow) + sqlite_session.commit() + + result = WorkflowAppGenerator._ensure_snippet_start_node_in_worker(session=sqlite_session, workflow=workflow) + + assert result is workflow + + @staticmethod + def test_ensure_snippet_start_node_in_worker_returns_snippet_workflow_when_snippet_missing( + sqlite_session: Session, + ) -> None: + workflow = _workflow_app_generator_workflow(app_id="snippet-1", tenant_id="tenant-1", kind=WorkflowKind.SNIPPET) + other_tenant_snippet = CustomizedSnippet( + id="snippet-1", + tenant_id="tenant-2", + name="Other tenant snippet", + description="", + type="node", + ) + sqlite_session.add_all([workflow, other_tenant_snippet]) + sqlite_session.commit() + + result = WorkflowAppGenerator._ensure_snippet_start_node_in_worker(session=sqlite_session, workflow=workflow) + + assert result is workflow + + @staticmethod + def test_ensure_snippet_start_node_in_worker_applies_snippet_start_injection( + monkeypatch: pytest.MonkeyPatch, sqlite_session: Session + ) -> None: + workflow = _workflow_app_generator_workflow(app_id="snippet-1", tenant_id="tenant-1", kind=WorkflowKind.SNIPPET) + snippet = CustomizedSnippet( + id="snippet-1", + tenant_id="tenant-1", + name="Matching snippet", + description="", + type="node", + ) + sqlite_session.add_all([workflow, snippet]) + sqlite_session.commit() + ensure_start_node = MagicMock(return_value=workflow) + monkeypatch.setattr( + "services.snippet_generate_service.SnippetGenerateService.ensure_start_node_for_worker", + ensure_start_node, + ) + + result = WorkflowAppGenerator._ensure_snippet_start_node_in_worker(session=sqlite_session, workflow=workflow) + + assert result is workflow + ensure_start_node.assert_called_once_with(workflow, snippet) + + @staticmethod + def test_generate_includes_parent_trace_context_in_extras( + monkeypatch: pytest.MonkeyPatch, sqlite_session: Session + ) -> None: + generator = WorkflowAppGenerator(execution_driver=WorkflowRunAgg.run) + app, workflow, end_user = _persist_generator_rows(sqlite_session) + + monkeypatch.setattr( + "core.app.apps.workflow.app_generator.WorkflowAppGenerator._bind_file_access_scope", + lambda *args, **kwargs: contextlib.nullcontext(), + ) + monkeypatch.setattr( + "core.app.apps.workflow.app_generator.WorkflowAppConfigManager.get_app_config", + lambda *args, **kwargs: _app_config(app, workflow), + ) + monkeypatch.setattr( + "core.app.apps.workflow.app_generator.file_factory.build_from_mappings", lambda *args, **kwargs: [] + ) + repository_tenant_ids: dict[str, str] = {} + workflow_execution_factory = ( + _workflow_app_generator.DifyCoreRepositoryFactory.create_workflow_execution_repository + ) + workflow_node_execution_factory = ( + _workflow_app_generator.DifyCoreRepositoryFactory.create_workflow_node_execution_repository + ) + + def create_workflow_execution_repository(**kwargs): + repository_tenant_ids["workflow"] = kwargs["tenant_id"] + return workflow_execution_factory(**kwargs) + + def create_workflow_node_execution_repository(**kwargs): + repository_tenant_ids["node"] = kwargs["tenant_id"] + return workflow_node_execution_factory(**kwargs) + + monkeypatch.setattr( + "core.app.apps.workflow.app_generator.DifyCoreRepositoryFactory.create_workflow_execution_repository", + create_workflow_execution_repository, + ) + monkeypatch.setattr( + "core.app.apps.workflow.app_generator.DifyCoreRepositoryFactory.create_workflow_node_execution_repository", + create_workflow_node_execution_repository, + ) + monkeypatch.setattr( + "core.app.apps.workflow.app_generator.db", SimpleNamespace(engine=sqlite_session.get_bind()) + ) + monkeypatch.setattr(generator, "_prepare_user_inputs", lambda *, user_inputs, **kwargs: user_inputs) + + captured = {} + + def fake_generate(**kwargs): + captured.update(kwargs) + return {"data": {}} + + monkeypatch.setattr(generator, "_generate", fake_generate) + + result = generator.generate( + app_model=app, + workflow=workflow, + user=end_user, + args={ + "inputs": {"query": "hello"}, + "files": [], + "external_trace_id": "trace-1", + "parent_trace_context": { + "parent_workflow_run_id": "outer-workflow-run-1", + "parent_node_execution_id": "outer-node-execution-1", + }, + "trace_session_id": "session-1", + }, + invoke_from=InvokeFrom.SERVICE_API, + streaming=False, + call_depth=0, + ) + + assert result == {"data": {}} + application_generate_entity = captured["application_generate_entity"] + assert isinstance(application_generate_entity, WorkflowAppGenerateEntity) + extras = application_generate_entity.extras + assert extras["external_trace_id"] == "trace-1" + assert extras["parent_trace_context"].model_dump() == { + "parent_workflow_run_id": "outer-workflow-run-1", + "parent_node_execution_id": "outer-node-execution-1", + } + assert extras["trace_session_id"] == "session-1" + assert isinstance(captured["workflow_execution_repository"], SQLAlchemyWorkflowExecutionRepository) + assert isinstance(captured["workflow_node_execution_repository"], SQLAlchemyWorkflowNodeExecutionRepository) + assert repository_tenant_ids == {"workflow": app.tenant_id, "node": app.tenant_id} + + @staticmethod + def test_resume_delegates_to_generate( + monkeypatch: pytest.MonkeyPatch, + sqlite_session: Session, + sqlite_session_factory: sessionmaker[Session], + ) -> None: + generator = WorkflowAppGenerator(execution_driver=WorkflowRunAgg.run) + app, workflow, end_user = _persist_generator_rows(sqlite_session) + mock_generate = MagicMock(return_value="ok") + monkeypatch.setattr(generator, "_generate", mock_generate) + + application_generate_entity = _generate_entity( + app, + workflow, + end_user, + invoke_from=InvokeFrom.DEBUGGER, + stream=False, + ) + runtime_state = _runtime_state() + workflow_execution_repository, workflow_node_execution_repository = _repositories( + sqlite_session_factory, app, end_user + ) + pause_config = PauseStateLayerConfig( + session_factory=sqlite_session_factory, + state_owner_user_id="owner", + ) + + result = generator.resume( + app_model=app, + workflow=workflow, + user=end_user, + application_generate_entity=application_generate_entity, + graph_runtime_state=runtime_state, + workflow_execution_repository=workflow_execution_repository, + workflow_node_execution_repository=workflow_node_execution_repository, + graph_engine_layers=("layer",), + pause_state_config=pause_config, + ) + + assert result == "ok" + mock_generate.assert_called_once() + kwargs = mock_generate.call_args.kwargs + assert kwargs["graph_runtime_state"] is runtime_state + assert kwargs["pause_state_config"] is pause_config + assert kwargs["streaming"] is False + assert kwargs["invoke_from"] == InvokeFrom.DEBUGGER + + @staticmethod + def test_generate_forwards_pause_configuration_and_state( + monkeypatch: pytest.MonkeyPatch, + sqlite_session: Session, + sqlite_session_factory: sessionmaker[Session], + ) -> None: + generator = WorkflowAppGenerator(execution_driver=WorkflowRunAgg.run) + app, workflow, end_user = _persist_generator_rows(sqlite_session) + + queue_manager = MagicMock() + monkeypatch.setattr( + "core.app.apps.workflow.app_generator.WorkflowAppQueueManager", + MagicMock(return_value=queue_manager), + ) + + monkeypatch.setattr( + "core.app.apps.workflow.app_generator.WorkflowAppGenerateResponseConverter.convert", + MagicMock(return_value="converted"), + ) + monkeypatch.setattr(WorkflowAppGenerator, "_handle_response", MagicMock(return_value="response")) + + draft_factory_tenant_ids: list[str] = [] + get_draft_var_saver_factory = WorkflowAppGenerator._get_draft_var_saver_factory + + def get_recording_draft_var_saver_factory( + invoke_from: InvokeFrom, account: EndUser, *, tenant_id: str + ) -> DraftVariableSaverFactory: + draft_factory_tenant_ids.append(tenant_id) + return get_draft_var_saver_factory(invoke_from, account, tenant_id=tenant_id) + + monkeypatch.setattr( + WorkflowAppGenerator, + "_get_draft_var_saver_factory", + staticmethod(get_recording_draft_var_saver_factory), + ) + + db_session = sqlite_session_factory() + monkeypatch.setattr(_workflow_app_generator.db, "session", db_session) + + worker_kwargs: dict[str, object] = {} + + class DummyThread: + def __init__(self, target, kwargs): + worker_kwargs["target"] = target + worker_kwargs["kwargs"] = kwargs + + def start(self): + return None + + def join(self, timeout): + worker_kwargs["joined"] = True + worker_kwargs["join_timeout"] = timeout + + def is_alive(self): + return False + + monkeypatch.setattr("core.app.apps.workflow.app_generator.threading.Thread", DummyThread) + + application_generate_entity = _generate_entity(app, workflow, end_user) + graph_runtime_state = _runtime_state() + workflow_execution_repository, workflow_node_execution_repository = _repositories( + sqlite_session_factory, app, end_user + ) + + flask_app = Flask(__name__) + with flask_app.app_context(): + result = generator._generate( + app_model=app, + workflow=workflow, + user=end_user, + application_generate_entity=application_generate_entity, + invoke_from=InvokeFrom.SERVICE_API, + workflow_execution_repository=workflow_execution_repository, + workflow_node_execution_repository=workflow_node_execution_repository, + streaming=True, + graph_engine_layers=("base-layer",), + graph_runtime_state=graph_runtime_state, + pause_state_config=PauseStateLayerConfig( + session_factory=sqlite_session_factory, + state_owner_user_id="owner", + ), + ) + + assert result == "converted" + graph_engine_layers = worker_kwargs["kwargs"]["graph_engine_layers"] + assert graph_engine_layers == ("base-layer",) + assert worker_kwargs["kwargs"]["pause_state_config"].state_owner_user_id == "owner" + assert worker_kwargs["kwargs"]["graph_runtime_state"] is graph_runtime_state + assert worker_kwargs["joined"] is True + assert worker_kwargs["join_timeout"] == 300 + assert draft_factory_tenant_ids == [app.tenant_id] + + @staticmethod + def test_resume_path_runs_worker_with_runtime_state( + monkeypatch: pytest.MonkeyPatch, + sqlite_session: Session, + sqlite_session_factory: sessionmaker[Session], + ) -> None: + generator = WorkflowAppGenerator(execution_driver=WorkflowRunAgg.run) + app, workflow, end_user = _persist_generator_rows(sqlite_session) + workflow_run = WorkflowRun( + id="run", + tenant_id=workflow.tenant_id, + app_id=app.id, + workflow_id=workflow.id, + type=workflow.type, + triggered_from=WorkflowRunTriggeredFrom.APP_RUN, + version=workflow.version, + graph=workflow.graph, + inputs="{}", + status=WorkflowExecutionStatus.RUNNING, + created_by_role=CreatorUserRole.END_USER, + created_by=end_user.id, + ) + sqlite_session.add(workflow_run) + sqlite_session.commit() + runtime_state = _runtime_state() + + queue_manager = MagicMock() + monkeypatch.setattr( + "core.app.apps.workflow.app_generator.WorkflowAppQueueManager", + MagicMock(return_value=queue_manager), + ) + + monkeypatch.setattr(generator, "_handle_response", MagicMock(return_value="raw-response")) + monkeypatch.setattr( + "core.app.apps.workflow.app_generator.WorkflowAppGenerateResponseConverter.convert", + MagicMock(side_effect=lambda response, invoke_from: response), + ) + + db_session = sqlite_session_factory() + monkeypatch.setattr(_workflow_app_generator.db, "session", db_session) + + runner_instance = MagicMock() + runner_instance.prepare.return_value = None + + def runner_ctor(**kwargs): + assert kwargs["graph_runtime_state"] is runtime_state + assert kwargs["workflow"].id == workflow.id + assert kwargs["system_user_id"] == end_user.session_id + assert kwargs["queue_manager"] is queue_manager + return runner_instance + + monkeypatch.setattr( + "core.app.apps.workflow.app_generator.WorkflowAppRunner", + MagicMock(side_effect=runner_ctor), + ) + + worker_lifecycle: dict[str, object] = {} + + class ImmediateThread: + def __init__(self, target, kwargs): + target(**kwargs) + + def start(self): + return None + + def join(self, timeout): + worker_lifecycle["joined"] = True + worker_lifecycle["join_timeout"] = timeout + + def is_alive(self): + return False + + monkeypatch.setattr("core.app.apps.workflow.app_generator.threading.Thread", ImmediateThread) + + pause_config = PauseStateLayerConfig( + session_factory=sqlite_session_factory, + state_owner_user_id="owner", + ) + + application_generate_entity = _generate_entity(app, workflow, end_user) + workflow_execution_repository, workflow_node_execution_repository = _repositories( + sqlite_session_factory, app, end_user + ) + + result = generator.resume( + app_model=app, + workflow=workflow, + user=end_user, + application_generate_entity=application_generate_entity, + graph_runtime_state=runtime_state, + workflow_execution_repository=workflow_execution_repository, + workflow_node_execution_repository=workflow_node_execution_repository, + pause_state_config=pause_config, + ) + + assert result == "raw-response" + assert worker_lifecycle["joined"] is True + assert worker_lifecycle["join_timeout"] == 300 + runner_instance.prepare.assert_called_once() + + +TENANT_ID = "00000000-0000-0000-0000-000000000001" +OTHER_TENANT_ID = "00000000-0000-0000-0000-000000000002" +APP_ID = "00000000-0000-0000-0000-000000000003" +WORKFLOW_ID = "00000000-0000-0000-0000-000000000004" +END_USER_ID = "00000000-0000-0000-0000-000000000005" +CREATOR_ID = "00000000-0000-0000-0000-000000000006" + + +def _app() -> App: + return App( + id=APP_ID, + tenant_id=TENANT_ID, + name="Workflow app", + description="", + mode=AppMode.WORKFLOW, + icon_type=None, + icon="", + icon_background=None, + app_model_config_id=None, + workflow_id=WORKFLOW_ID, + enable_site=False, + enable_api=True, + max_active_requests=None, + created_by=CREATOR_ID, + ) + + +def _persist_app(session: Session) -> App: + app = _app() + session.add(app) + session.commit() + return app + + +def _app_generator_extra_workflow( + *, + workflow_id: str = WORKFLOW_ID, + app_id: str = APP_ID, + tenant_id: str = TENANT_ID, + kind: WorkflowKind = WorkflowKind.STANDARD, +) -> Workflow: + workflow = Workflow.new( + tenant_id=tenant_id, + app_id=app_id, + type=WorkflowType.WORKFLOW.value, + version="1", + graph=json.dumps({"nodes": [], "edges": []}), + features="{}", + created_by=CREATOR_ID, + environment_variables=[], + conversation_variables=[], + rag_pipeline_variables=[], + kind=kind.value, + ) + workflow.id = workflow_id + return workflow + + +def _persist_workflow( + session: Session, + *, + workflow_id: str = WORKFLOW_ID, + app_id: str = APP_ID, + tenant_id: str = TENANT_ID, + kind: WorkflowKind = WorkflowKind.STANDARD, +) -> Workflow: + workflow = _app_generator_extra_workflow( + workflow_id=workflow_id, + app_id=app_id, + tenant_id=tenant_id, + kind=kind, + ) + session.add(workflow) + session.commit() + return workflow + + +def _end_user() -> EndUser: + return EndUser( + id=END_USER_ID, + tenant_id=TENANT_ID, + app_id=APP_ID, + type=EndUserType.BROWSER, + name="End user", + session_id="session-id", + ) + + +def _persist_end_user(session: Session) -> EndUser: + end_user = _end_user() + session.add(end_user) + session.commit() + return end_user + + +def _persist_snippet( + session: Session, + *, + snippet_id: str, + tenant_id: str = TENANT_ID, +) -> CustomizedSnippet: + snippet = CustomizedSnippet( + id=snippet_id, + tenant_id=tenant_id, + name="Snippet", + description=None, + type="node", + ) + session.add(snippet) + session.commit() + return snippet + + +class TestAppGeneratorExtraCases: + @staticmethod + @pytest.fixture(autouse=True) + def disable_trace_delivery(monkeypatch: pytest.MonkeyPatch) -> None: + """Construct real trace managers without provider lookup or background delivery.""" + monkeypatch.setattr(OpsTraceManager, "get_ops_trace_instance", lambda app_id: None) + monkeypatch.setattr(TraceQueueManager, "start_timer", lambda self: None) + monkeypatch.setattr("core.telemetry.gateway.is_enterprise_telemetry_enabled", lambda: False) + + @staticmethod + @pytest.fixture + def sqlite_generator_session( + monkeypatch: pytest.MonkeyPatch, + sqlite_engine, + sqlite_session: Session, + ) -> Session: + """Expose the shared SQLite engine and real Session through the Flask extension boundary.""" + monkeypatch.setattr( + _app_generator_extra_app_generator_module, + "db", + SimpleNamespace(engine=sqlite_engine, session=sqlite_session), + ) + return sqlite_session + + class TestWorkflowAppGeneratorValidation: + @pytest.mark.usefixtures("sqlite_generator_session") + def test_generate_stream_joins_worker_after_response_exhaustion( + self, monkeypatch: pytest.MonkeyPatch, sqlite_session: Session + ): + generator = WorkflowAppGenerator(execution_driver=WorkflowRunAgg.run) + app = _persist_app(sqlite_session) + workflow = _persist_workflow(sqlite_session) + user = _persist_end_user(sqlite_session) + worker_thread = Mock() + worker_thread.is_alive.return_value = False + app_config = WorkflowUIBasedAppConfig( + tenant_id="tenant", + app_id="app", + app_mode=AppMode.WORKFLOW, + additional_features=AppAdditionalFeatures(), + variables=[], + workflow_id="workflow-id", + ) + application_generate_entity = WorkflowAppGenerateEntity.model_construct( + task_id="task", + app_config=app_config, + inputs={}, + files=[], + user_id="user", + stream=True, + invoke_from=InvokeFrom.WEB_APP, + extras={}, + ) + + def response_stream(): + yield {"event": "workflow_finished"} + + monkeypatch.setattr(generator, "_bind_file_access_scope", lambda **kwargs: contextlib.nullcontext()) + monkeypatch.setattr( + "core.app.apps.workflow.app_generator.WorkflowAppQueueManager", + lambda **kwargs: SimpleNamespace(**kwargs), + ) + monkeypatch.setattr( + "core.app.apps.workflow.app_generator.current_app", + SimpleNamespace(_get_current_object=lambda: SimpleNamespace(name="flask")), + ) + monkeypatch.setattr("core.app.apps.workflow.app_generator.contextvars.copy_context", lambda: "ctx") + monkeypatch.setattr("core.app.apps.workflow.app_generator.threading.Thread", lambda **kwargs: worker_thread) + monkeypatch.setattr(generator, "_get_draft_var_saver_factory", lambda *args, **kwargs: "draft-factory") + monkeypatch.setattr(generator, "_handle_response", lambda **kwargs: response_stream()) + monkeypatch.setattr( + "core.app.apps.workflow.app_generator.WorkflowAppGenerateResponseConverter.convert", + lambda response, invoke_from: response, + ) + + managed_stream = generator._generate( + app_model=app, + workflow=workflow, + user=user, + application_generate_entity=application_generate_entity, + invoke_from=InvokeFrom.WEB_APP, + workflow_execution_repository=SimpleNamespace(), + workflow_node_execution_repository=SimpleNamespace(), + streaming=True, + ) + + worker_thread.start.assert_called_once_with() + worker_thread.join.assert_not_called() + assert list(managed_stream) == [{"event": "workflow_finished"}] + worker_thread.join.assert_called_once_with(timeout=300) + + def test_ensure_snippet_start_node_returns_original_for_non_snippet_workflow( + self, + unbound_session: Session, + ): + workflow = _app_generator_extra_workflow() + + result = WorkflowAppGenerator._ensure_snippet_start_node_in_worker( + session=unbound_session, + workflow=workflow, + ) + + assert result is workflow + + def test_ensure_snippet_start_node_returns_original_when_snippet_is_from_another_tenant( + self, + sqlite_session: Session, + ): + workflow = _persist_workflow(sqlite_session, kind=WorkflowKind.SNIPPET) + _persist_snippet(sqlite_session, snippet_id=APP_ID, tenant_id=OTHER_TENANT_ID) + + result = WorkflowAppGenerator._ensure_snippet_start_node_in_worker( + session=sqlite_session, + workflow=workflow, + ) + + assert result is workflow + + def test_ensure_snippet_start_node_delegates_when_snippet_exists( + self, + monkeypatch: pytest.MonkeyPatch, + sqlite_session: Session, + ): + workflow = _persist_workflow(sqlite_session, kind=WorkflowKind.SNIPPET) + snippet = _persist_snippet(sqlite_session, snippet_id=APP_ID) + injected_workflow = _app_generator_extra_workflow(workflow_id="workflow-injected") + ensure_start_node = Mock(return_value=injected_workflow) + monkeypatch.setattr( + "services.snippet_generate_service.SnippetGenerateService.ensure_start_node_for_worker", + ensure_start_node, + ) + + result = WorkflowAppGenerator._ensure_snippet_start_node_in_worker( + session=sqlite_session, + workflow=workflow, + ) + + assert result is injected_workflow + ensure_start_node.assert_called_once_with(workflow, snippet) + + def test_single_iteration_generate_validates_args(self, sqlite_session: Session): + generator = WorkflowAppGenerator(execution_driver=WorkflowRunAgg.run) + + with pytest.raises(ValueError, match="node_id is required"): + generator.single_iteration_generate( + app_model=_app(), + workflow=_app_generator_extra_workflow(), + node_id="", + user=_end_user(), + args={"inputs": {}}, + streaming=False, + session=sqlite_session, + ) + + with pytest.raises(ValueError, match="inputs is required"): + generator.single_iteration_generate( + app_model=_app(), + workflow=_app_generator_extra_workflow(), + node_id="node", + user=_end_user(), + args={}, + streaming=False, + session=sqlite_session, + ) + + def test_single_loop_generate_validates_args(self, sqlite_session: Session): + generator = WorkflowAppGenerator(execution_driver=WorkflowRunAgg.run) + + with pytest.raises(ValueError, match="node_id is required"): + generator.single_loop_generate( + app_model=_app(), + workflow=_app_generator_extra_workflow(), + node_id="", + user=_end_user(), + args=SimpleNamespace(inputs={}), + streaming=False, + session=sqlite_session, + ) + + @pytest.mark.usefixtures("sqlite_generator_session") + def test_single_iteration_generate_includes_trace_session_id_in_extras( + self, + monkeypatch: pytest.MonkeyPatch, + sqlite_session: Session, + ): + generator = WorkflowAppGenerator(execution_driver=WorkflowRunAgg.run) + app = _persist_app(sqlite_session) + workflow = _persist_workflow(sqlite_session) + user = _persist_end_user(sqlite_session) + app_config = WorkflowUIBasedAppConfig( + tenant_id=TENANT_ID, + app_id=APP_ID, + app_mode=AppMode.WORKFLOW, + additional_features=AppAdditionalFeatures(), + variables=[], + workflow_id=WORKFLOW_ID, + ) + captured: dict[str, object] = {} + repository_session_makers: list[sessionmaker[Session]] = [] + draft_sessions: list[Session] = [] + + monkeypatch.setattr( + "core.app.apps.workflow.app_generator.WorkflowAppConfigManager.get_app_config", + lambda **kwargs: app_config, + ) + monkeypatch.setattr( + "core.app.apps.workflow.app_generator.DifyCoreRepositoryFactory.create_workflow_execution_repository", + lambda **kwargs: repository_session_makers.append(kwargs["session_factory"]) or SimpleNamespace(), + ) + monkeypatch.setattr( + "core.app.apps.workflow.app_generator.DifyCoreRepositoryFactory.create_workflow_node_execution_repository", + lambda **kwargs: repository_session_makers.append(kwargs["session_factory"]) or SimpleNamespace(), + ) + monkeypatch.setattr( + "core.app.apps.workflow.app_generator.DraftVarLoader", lambda **kwargs: SimpleNamespace() + ) + monkeypatch.setattr( + "core.app.apps.workflow.app_generator.WorkflowDraftVariableService", + lambda session: ( + draft_sessions.append(session) + or SimpleNamespace(prefill_conversation_variable_default_values=lambda *args, **kwargs: None) + ), + ) + monkeypatch.setattr(generator, "_generate", lambda **kwargs: captured.update(kwargs) or {"ok": True}) + + generator.single_iteration_generate( + app_model=app, + workflow=workflow, + node_id="node-1", + user=user, + args={"inputs": {"foo": "bar"}, "trace_session_id": "session-1"}, + streaming=False, + session=sqlite_session, + ) + + assert captured["application_generate_entity"].extras["trace_session_id"] == "session-1" + assert len(repository_session_makers) == 2 + assert all(factory.kw["bind"] is sqlite_session.get_bind() for factory in repository_session_makers) + assert len(draft_sessions) == 1 + assert draft_sessions[0] is sqlite_session + + @pytest.mark.usefixtures("sqlite_generator_session") + def test_single_loop_generate_includes_trace_session_id_in_extras( + self, + monkeypatch: pytest.MonkeyPatch, + sqlite_session: Session, + ): + generator = WorkflowAppGenerator(execution_driver=WorkflowRunAgg.run) + app = _persist_app(sqlite_session) + workflow = _persist_workflow(sqlite_session) + user = _persist_end_user(sqlite_session) + app_config = WorkflowUIBasedAppConfig( + tenant_id=TENANT_ID, + app_id=APP_ID, + app_mode=AppMode.WORKFLOW, + additional_features=AppAdditionalFeatures(), + variables=[], + workflow_id=WORKFLOW_ID, + ) + captured: dict[str, object] = {} + repository_session_makers: list[sessionmaker[Session]] = [] + draft_sessions: list[Session] = [] + + monkeypatch.setattr( + "core.app.apps.workflow.app_generator.WorkflowAppConfigManager.get_app_config", + lambda **kwargs: app_config, + ) + monkeypatch.setattr( + "core.app.apps.workflow.app_generator.DifyCoreRepositoryFactory.create_workflow_execution_repository", + lambda **kwargs: repository_session_makers.append(kwargs["session_factory"]) or SimpleNamespace(), + ) + monkeypatch.setattr( + "core.app.apps.workflow.app_generator.DifyCoreRepositoryFactory.create_workflow_node_execution_repository", + lambda **kwargs: repository_session_makers.append(kwargs["session_factory"]) or SimpleNamespace(), + ) + monkeypatch.setattr( + "core.app.apps.workflow.app_generator.DraftVarLoader", lambda **kwargs: SimpleNamespace() + ) + monkeypatch.setattr( + "core.app.apps.workflow.app_generator.WorkflowDraftVariableService", + lambda session: ( + draft_sessions.append(session) + or SimpleNamespace(prefill_conversation_variable_default_values=lambda *args, **kwargs: None) + ), + ) + monkeypatch.setattr(generator, "_generate", lambda **kwargs: captured.update(kwargs) or {"ok": True}) + + generator.single_loop_generate( + app_model=app, + workflow=workflow, + node_id="node-2", + user=user, + args=SimpleNamespace(inputs={"foo": "bar"}, trace_session_id="session-1"), + streaming=False, + session=sqlite_session, + ) + + assert captured["application_generate_entity"].extras["trace_session_id"] == "session-1" + assert len(repository_session_makers) == 2 + assert all(factory.kw["bind"] is sqlite_session.get_bind() for factory in repository_session_makers) + assert len(draft_sessions) == 1 + assert draft_sessions[0] is sqlite_session + + with pytest.raises(ValueError, match="inputs is required"): + generator.single_loop_generate( + app_model=_app(), + workflow=_app_generator_extra_workflow(), + node_id="node", + user=_end_user(), + args=SimpleNamespace(inputs=None), + streaming=False, + session=sqlite_session, + ) + + class TestWorkflowAppGeneratorHandleResponse: + def test_handle_response_closed_file_raises_stopped(self, monkeypatch: pytest.MonkeyPatch): + generator = WorkflowAppGenerator(execution_driver=WorkflowRunAgg.run) + + app_config = WorkflowUIBasedAppConfig( + tenant_id="tenant", + app_id="app", + app_mode=AppMode.WORKFLOW, + additional_features=AppAdditionalFeatures(), + variables=[], + workflow_id="workflow-id", + ) + application_generate_entity = WorkflowAppGenerateEntity.model_construct( + task_id="task", + app_config=app_config, + inputs={}, + files=[], + user_id="user", + stream=False, + invoke_from=InvokeFrom.WEB_APP, + extras={}, + trace_manager=None, + workflow_execution_id="run-id", + call_depth=0, + ) + + class _Pipeline: + def __init__(self, **kwargs) -> None: + _ = kwargs + + def process(self): + raise ValueError("I/O operation on closed file.") + + monkeypatch.setattr( + "core.app.apps.workflow.app_generator.WorkflowAppGenerateTaskPipeline", + _Pipeline, + ) + + with pytest.raises(GenerateTaskStoppedError): + generator._handle_response( + application_generate_entity=application_generate_entity, + workflow=_app_generator_extra_workflow(), + queue_manager=SimpleNamespace(), + user=_end_user(), + draft_var_saver_factory=lambda **kwargs: None, + stream=False, + ) + + class TestWorkflowAppGeneratorGenerate: + @pytest.fixture + def generation(self, monkeypatch: pytest.MonkeyPatch, sqlite_generator_session: Session): + """Keep input preparation real and stop at the execution boundary.""" + generator = WorkflowAppGenerator(execution_driver=WorkflowRunAgg.run) + app = _persist_app(sqlite_generator_session) + workflow = _persist_workflow(sqlite_generator_session) + user = _persist_end_user(sqlite_generator_session) + execute = Mock(return_value={"ok": True}) + monkeypatch.setattr(generator, "_generate", execute) + return generator, app, workflow, user, execute + + @pytest.mark.parametrize("node_type", ["trigger-webhook", "trigger-schedule", "trigger-plugin"]) + @pytest.mark.parametrize("root_node_id", [None, "selected-trigger"]) + @pytest.mark.parametrize("invoke_from", [InvokeFrom.SERVICE_API, InvokeFrom.DEBUGGER]) + def test_generate_preserves_trigger_inputs(self, generation, node_type, root_node_id, invoke_from): + generator, app, workflow, user, execute = generation + workflow.graph = json.dumps( + { + "nodes": [ + {"id": "first-trigger", "data": {"type": node_type, "title": "First trigger"}}, + {"id": "selected-trigger", "data": {"type": node_type, "title": "Selected trigger"}}, + ], + "edges": [], + } + ) + inputs = {"event": {"message": "hello", "items": [1, 2]}, "enabled": False} + + result = generator.generate( + app_model=app, + workflow=workflow, + user=user, + args={"inputs": inputs}, + invoke_from=invoke_from, + streaming=False, + root_node_id=root_node_id, + ) + + assert result == {"ok": True} + assert execute.call_args.kwargs["application_generate_entity"].inputs == inputs + assert execute.call_args.kwargs["root_node_id"] == (root_node_id or "first-trigger") + + @pytest.mark.parametrize("root_node_id", [None, "start"]) + def test_generate_prepares_start_inputs(self, generation, root_node_id): + generator, app, workflow, user, execute = generation + workflow.graph = json.dumps( + { + "nodes": [ + { + "id": "start", + "data": { + "type": "start", + "title": "Start", + "variables": [ + { + "variable": "question", + "label": "Question", + "type": "text-input", + "required": True, + }, + {"variable": "count", "label": "Count", "type": "number", "default": "3"}, + ], + }, + } + ], + "edges": [], + } + ) + + generator.generate( + app_model=app, + workflow=workflow, + user=user, + args={"inputs": {"question": "hello", "unknown": "filtered"}}, + invoke_from=InvokeFrom.SERVICE_API, + streaming=False, + root_node_id=root_node_id, + ) + + assert execute.call_args.kwargs["application_generate_entity"].inputs == {"question": "hello", "count": 3} + assert execute.call_args.kwargs["root_node_id"] == "start" + + execute.reset_mock() + with pytest.raises(ValueError, match="question is required in input form"): + generator.generate( + app_model=app, + workflow=workflow, + user=user, + args={"inputs": {}}, + invoke_from=InvokeFrom.SERVICE_API, + streaming=False, + root_node_id=root_node_id, + ) + execute.assert_not_called() + + class TestWorkflowAppGeneratorResume: + def test_resume_restores_trace_manager_when_missing(self, monkeypatch: pytest.MonkeyPatch): + generator = WorkflowAppGenerator(execution_driver=WorkflowRunAgg.run) + app_config = WorkflowUIBasedAppConfig( + tenant_id="tenant", + app_id="app", + app_mode=AppMode.WORKFLOW, + additional_features=AppAdditionalFeatures(), + variables=[], + workflow_id="workflow-id", + ) + application_generate_entity = WorkflowAppGenerateEntity.model_construct( + task_id="task", + app_config=app_config, + inputs={}, + files=[], + user_id="user", + stream=False, + invoke_from=InvokeFrom.WEB_APP, + extras={}, + trace_manager=None, + workflow_execution_id="run-id", + call_depth=0, + ) + dummy_trace_queue_manager = type( + "_DummyTraceQueueManager", + (TraceQueueManager,), + { + "__init__": lambda self, app_id=None, user_id=None: ( + setattr(self, "app_id", app_id) or setattr(self, "user_id", user_id) + ) + }, + ) + monkeypatch.setattr( + "core.app.apps.workflow.app_generator.TraceQueueManager", + dummy_trace_queue_manager, + ) + captured_entity: WorkflowAppGenerateEntity | None = None + + def _fake_generate(**kwargs): + nonlocal captured_entity + captured_entity = kwargs["application_generate_entity"] + return SimpleNamespace(ok=True) + + monkeypatch.setattr(generator, "_generate", _fake_generate) + + result = generator.resume( + app_model=_app(), + workflow=_app_generator_extra_workflow(), + user=_end_user(), + application_generate_entity=application_generate_entity, + graph_runtime_state=SimpleNamespace(), + workflow_execution_repository=SimpleNamespace(), + workflow_node_execution_repository=SimpleNamespace(), + ) + + assert result.ok is True + assert captured_entity is not None + trace_manager = captured_entity.trace_manager + assert isinstance(trace_manager, dummy_trace_queue_manager) + assert trace_manager.app_id == APP_ID + assert trace_manager.user_id == "session-id" + + def test_resume_preserves_existing_trace_manager(self, monkeypatch: pytest.MonkeyPatch): + generator = WorkflowAppGenerator(execution_driver=WorkflowRunAgg.run) + app_config = WorkflowUIBasedAppConfig( + tenant_id="tenant", + app_id="app", + app_mode=AppMode.WORKFLOW, + additional_features=AppAdditionalFeatures(), + variables=[], + workflow_id="workflow-id", + ) + existing_trace_manager = SimpleNamespace(app_id="existing-app", user_id="existing-user") + application_generate_entity = WorkflowAppGenerateEntity.model_construct( + task_id="task", + app_config=app_config, + inputs={}, + files=[], + user_id="user", + stream=False, + invoke_from=InvokeFrom.WEB_APP, + extras={}, + trace_manager=existing_trace_manager, + workflow_execution_id="run-id", + call_depth=0, + ) + captured_entity: WorkflowAppGenerateEntity | None = None + + def _fake_generate(**kwargs): + nonlocal captured_entity + captured_entity = kwargs["application_generate_entity"] + return SimpleNamespace(ok=True) + + monkeypatch.setattr(generator, "_generate", _fake_generate) + + result = generator.resume( + app_model=_app(), + workflow=_app_generator_extra_workflow(), + user=_end_user(), + application_generate_entity=application_generate_entity, + graph_runtime_state=SimpleNamespace(), + workflow_execution_repository=SimpleNamespace(), + workflow_node_execution_repository=SimpleNamespace(), + ) + + assert result.ok is True + assert captured_entity is not None + assert captured_entity.trace_manager is existing_trace_manager + + class TestWorkflowAppGeneratorWorker: + def test_generate_worker_uses_end_user_session_for_external_invocation( + self, + monkeypatch: pytest.MonkeyPatch, + sqlite_session: Session, + ): + generator = WorkflowAppGenerator(execution_driver=WorkflowRunAgg.run) + _persist_app(sqlite_session) + _persist_workflow(sqlite_session) + _persist_end_user(sqlite_session) + + runner_kwargs = {} + + class _Runner: + def __init__(self, **kwargs): + runner_kwargs.update(kwargs) + + def prepare(self): + return None + + monkeypatch.setattr( + "core.app.apps.workflow.app_generator.preserve_flask_contexts", + lambda flask_app, context_vars: contextlib.nullcontext(), + ) + monkeypatch.setattr("core.app.apps.workflow.app_generator.WorkflowAppRunner", _Runner) + restore_workflow_run_graph = Mock() + monkeypatch.setattr(generator, "_restore_workflow_run_graph", restore_workflow_run_graph) + + app_config = WorkflowUIBasedAppConfig( + tenant_id=TENANT_ID, + app_id=APP_ID, + app_mode=AppMode.WORKFLOW, + additional_features=AppAdditionalFeatures(), + variables=[], + workflow_id=WORKFLOW_ID, + ) + application_generate_entity = WorkflowAppGenerateEntity.model_construct( + task_id="task", + app_config=app_config, + inputs={}, + files=[], + user_id=END_USER_ID, + stream=False, + invoke_from=InvokeFrom.WEB_APP, + extras={}, + trace_manager=None, + workflow_execution_id="run-id", + call_depth=0, + ) + + generator._generate_worker( + flask_app=SimpleNamespace(), + application_generate_entity=application_generate_entity, + queue_manager=SimpleNamespace(), + context=SimpleNamespace(), + variable_loader=SimpleNamespace(), + workflow_execution_repository=SimpleNamespace(), + workflow_node_execution_repository=SimpleNamespace(), + graph_runtime_state=SimpleNamespace(), + ) + + assert runner_kwargs["system_user_id"] == "session-id" + restore_workflow_run_graph.assert_called_once() + restore_kwargs = restore_workflow_run_graph.call_args.kwargs + assert isinstance(restore_kwargs["session"], Session) + assert restore_kwargs["workflow"] is runner_kwargs["workflow"] + assert restore_kwargs["workflow_run_id"] == "run-id" + assert inspect(runner_kwargs["workflow"]).detached is True diff --git a/api/tests/unit_tests/core/app/apps/workflow/test_app_generator_extra.py b/api/tests/unit_tests/core/app/apps/workflow/test_app_generator_extra.py deleted file mode 100644 index 98333d15881..00000000000 --- a/api/tests/unit_tests/core/app/apps/workflow/test_app_generator_extra.py +++ /dev/null @@ -1,766 +0,0 @@ -from __future__ import annotations - -import contextlib -import json -from types import SimpleNamespace -from unittest.mock import Mock - -import pytest -from sqlalchemy import inspect -from sqlalchemy.orm import Session, sessionmaker - -from core.app.app_config.entities import AppAdditionalFeatures, WorkflowUIBasedAppConfig -from core.app.apps.exc import GenerateTaskStoppedError -from core.app.apps.workflow import app_generator as app_generator_module -from core.app.apps.workflow.app_generator import WorkflowAppGenerator -from core.app.entities.app_invoke_entities import InvokeFrom, WorkflowAppGenerateEntity -from core.ops.ops_trace_manager import OpsTraceManager, TraceQueueManager -from models.enums import EndUserType -from models.model import App, AppMode, EndUser -from models.snippet import CustomizedSnippet -from models.workflow import Workflow, WorkflowKind, WorkflowType -from services.workflow_run_agg import WorkflowRunAgg - - -@pytest.fixture(autouse=True) -def disable_trace_delivery(monkeypatch: pytest.MonkeyPatch) -> None: - """Construct real trace managers without provider lookup or background delivery.""" - monkeypatch.setattr(OpsTraceManager, "get_ops_trace_instance", lambda app_id: None) - monkeypatch.setattr(TraceQueueManager, "start_timer", lambda self: None) - monkeypatch.setattr("core.telemetry.gateway.is_enterprise_telemetry_enabled", lambda: False) - - -TENANT_ID = "00000000-0000-0000-0000-000000000001" -OTHER_TENANT_ID = "00000000-0000-0000-0000-000000000002" -APP_ID = "00000000-0000-0000-0000-000000000003" -WORKFLOW_ID = "00000000-0000-0000-0000-000000000004" -END_USER_ID = "00000000-0000-0000-0000-000000000005" -CREATOR_ID = "00000000-0000-0000-0000-000000000006" - - -@pytest.fixture -def sqlite_generator_session( - monkeypatch: pytest.MonkeyPatch, - sqlite_engine, - sqlite_session: Session, -) -> Session: - """Expose the shared SQLite engine and real Session through the Flask extension boundary.""" - monkeypatch.setattr( - app_generator_module, - "db", - SimpleNamespace(engine=sqlite_engine, session=sqlite_session), - ) - return sqlite_session - - -def _app() -> App: - return App( - id=APP_ID, - tenant_id=TENANT_ID, - name="Workflow app", - description="", - mode=AppMode.WORKFLOW, - icon_type=None, - icon="", - icon_background=None, - app_model_config_id=None, - workflow_id=WORKFLOW_ID, - enable_site=False, - enable_api=True, - max_active_requests=None, - created_by=CREATOR_ID, - ) - - -def _persist_app(session: Session) -> App: - app = _app() - session.add(app) - session.commit() - return app - - -def _workflow( - *, - workflow_id: str = WORKFLOW_ID, - app_id: str = APP_ID, - tenant_id: str = TENANT_ID, - kind: WorkflowKind = WorkflowKind.STANDARD, -) -> Workflow: - workflow = Workflow.new( - tenant_id=tenant_id, - app_id=app_id, - type=WorkflowType.WORKFLOW.value, - version="1", - graph=json.dumps({"nodes": [], "edges": []}), - features="{}", - created_by=CREATOR_ID, - environment_variables=[], - conversation_variables=[], - rag_pipeline_variables=[], - kind=kind.value, - ) - workflow.id = workflow_id - return workflow - - -def _persist_workflow( - session: Session, - *, - workflow_id: str = WORKFLOW_ID, - app_id: str = APP_ID, - tenant_id: str = TENANT_ID, - kind: WorkflowKind = WorkflowKind.STANDARD, -) -> Workflow: - workflow = _workflow( - workflow_id=workflow_id, - app_id=app_id, - tenant_id=tenant_id, - kind=kind, - ) - session.add(workflow) - session.commit() - return workflow - - -def _end_user() -> EndUser: - return EndUser( - id=END_USER_ID, - tenant_id=TENANT_ID, - app_id=APP_ID, - type=EndUserType.BROWSER, - name="End user", - session_id="session-id", - ) - - -def _persist_end_user(session: Session) -> EndUser: - end_user = _end_user() - session.add(end_user) - session.commit() - return end_user - - -def _persist_snippet( - session: Session, - *, - snippet_id: str, - tenant_id: str = TENANT_ID, -) -> CustomizedSnippet: - snippet = CustomizedSnippet( - id=snippet_id, - tenant_id=tenant_id, - name="Snippet", - description=None, - type="node", - ) - session.add(snippet) - session.commit() - return snippet - - -class TestWorkflowAppGeneratorValidation: - @pytest.mark.usefixtures("sqlite_generator_session") - def test_generate_stream_joins_worker_after_response_exhaustion( - self, monkeypatch: pytest.MonkeyPatch, sqlite_session: Session - ): - generator = WorkflowAppGenerator(execution_driver=WorkflowRunAgg.run) - app = _persist_app(sqlite_session) - workflow = _persist_workflow(sqlite_session) - user = _persist_end_user(sqlite_session) - worker_thread = Mock() - worker_thread.is_alive.return_value = False - app_config = WorkflowUIBasedAppConfig( - tenant_id="tenant", - app_id="app", - app_mode=AppMode.WORKFLOW, - additional_features=AppAdditionalFeatures(), - variables=[], - workflow_id="workflow-id", - ) - application_generate_entity = WorkflowAppGenerateEntity.model_construct( - task_id="task", - app_config=app_config, - inputs={}, - files=[], - user_id="user", - stream=True, - invoke_from=InvokeFrom.WEB_APP, - extras={}, - ) - - def response_stream(): - yield {"event": "workflow_finished"} - - monkeypatch.setattr(generator, "_bind_file_access_scope", lambda **kwargs: contextlib.nullcontext()) - monkeypatch.setattr( - "core.app.apps.workflow.app_generator.WorkflowAppQueueManager", - lambda **kwargs: SimpleNamespace(**kwargs), - ) - monkeypatch.setattr( - "core.app.apps.workflow.app_generator.current_app", - SimpleNamespace(_get_current_object=lambda: SimpleNamespace(name="flask")), - ) - monkeypatch.setattr("core.app.apps.workflow.app_generator.contextvars.copy_context", lambda: "ctx") - monkeypatch.setattr("core.app.apps.workflow.app_generator.threading.Thread", lambda **kwargs: worker_thread) - monkeypatch.setattr(generator, "_get_draft_var_saver_factory", lambda *args, **kwargs: "draft-factory") - monkeypatch.setattr(generator, "_handle_response", lambda **kwargs: response_stream()) - monkeypatch.setattr( - "core.app.apps.workflow.app_generator.WorkflowAppGenerateResponseConverter.convert", - lambda response, invoke_from: response, - ) - - managed_stream = generator._generate( - app_model=app, - workflow=workflow, - user=user, - application_generate_entity=application_generate_entity, - invoke_from=InvokeFrom.WEB_APP, - workflow_execution_repository=SimpleNamespace(), - workflow_node_execution_repository=SimpleNamespace(), - streaming=True, - ) - - worker_thread.start.assert_called_once_with() - worker_thread.join.assert_not_called() - assert list(managed_stream) == [{"event": "workflow_finished"}] - worker_thread.join.assert_called_once_with(timeout=300) - - def test_ensure_snippet_start_node_returns_original_for_non_snippet_workflow( - self, - unbound_session: Session, - ): - workflow = _workflow() - - result = WorkflowAppGenerator._ensure_snippet_start_node_in_worker( - session=unbound_session, - workflow=workflow, - ) - - assert result is workflow - - def test_ensure_snippet_start_node_returns_original_when_snippet_is_from_another_tenant( - self, - sqlite_session: Session, - ): - workflow = _persist_workflow(sqlite_session, kind=WorkflowKind.SNIPPET) - _persist_snippet(sqlite_session, snippet_id=APP_ID, tenant_id=OTHER_TENANT_ID) - - result = WorkflowAppGenerator._ensure_snippet_start_node_in_worker( - session=sqlite_session, - workflow=workflow, - ) - - assert result is workflow - - def test_ensure_snippet_start_node_delegates_when_snippet_exists( - self, - monkeypatch: pytest.MonkeyPatch, - sqlite_session: Session, - ): - workflow = _persist_workflow(sqlite_session, kind=WorkflowKind.SNIPPET) - snippet = _persist_snippet(sqlite_session, snippet_id=APP_ID) - injected_workflow = _workflow(workflow_id="workflow-injected") - ensure_start_node = Mock(return_value=injected_workflow) - monkeypatch.setattr( - "services.snippet_generate_service.SnippetGenerateService.ensure_start_node_for_worker", - ensure_start_node, - ) - - result = WorkflowAppGenerator._ensure_snippet_start_node_in_worker( - session=sqlite_session, - workflow=workflow, - ) - - assert result is injected_workflow - ensure_start_node.assert_called_once_with(workflow, snippet) - - def test_single_iteration_generate_validates_args(self, sqlite_session: Session): - generator = WorkflowAppGenerator(execution_driver=WorkflowRunAgg.run) - - with pytest.raises(ValueError, match="node_id is required"): - generator.single_iteration_generate( - app_model=_app(), - workflow=_workflow(), - node_id="", - user=_end_user(), - args={"inputs": {}}, - streaming=False, - session=sqlite_session, - ) - - with pytest.raises(ValueError, match="inputs is required"): - generator.single_iteration_generate( - app_model=_app(), - workflow=_workflow(), - node_id="node", - user=_end_user(), - args={}, - streaming=False, - session=sqlite_session, - ) - - def test_single_loop_generate_validates_args(self, sqlite_session: Session): - generator = WorkflowAppGenerator(execution_driver=WorkflowRunAgg.run) - - with pytest.raises(ValueError, match="node_id is required"): - generator.single_loop_generate( - app_model=_app(), - workflow=_workflow(), - node_id="", - user=_end_user(), - args=SimpleNamespace(inputs={}), - streaming=False, - session=sqlite_session, - ) - - @pytest.mark.usefixtures("sqlite_generator_session") - def test_single_iteration_generate_includes_trace_session_id_in_extras( - self, - monkeypatch: pytest.MonkeyPatch, - sqlite_session: Session, - ): - generator = WorkflowAppGenerator(execution_driver=WorkflowRunAgg.run) - app = _persist_app(sqlite_session) - workflow = _persist_workflow(sqlite_session) - user = _persist_end_user(sqlite_session) - app_config = WorkflowUIBasedAppConfig( - tenant_id=TENANT_ID, - app_id=APP_ID, - app_mode=AppMode.WORKFLOW, - additional_features=AppAdditionalFeatures(), - variables=[], - workflow_id=WORKFLOW_ID, - ) - captured: dict[str, object] = {} - repository_session_makers: list[sessionmaker[Session]] = [] - draft_sessions: list[Session] = [] - - monkeypatch.setattr( - "core.app.apps.workflow.app_generator.WorkflowAppConfigManager.get_app_config", - lambda **kwargs: app_config, - ) - monkeypatch.setattr( - "core.app.apps.workflow.app_generator.DifyCoreRepositoryFactory.create_workflow_execution_repository", - lambda **kwargs: repository_session_makers.append(kwargs["session_factory"]) or SimpleNamespace(), - ) - monkeypatch.setattr( - "core.app.apps.workflow.app_generator.DifyCoreRepositoryFactory.create_workflow_node_execution_repository", - lambda **kwargs: repository_session_makers.append(kwargs["session_factory"]) or SimpleNamespace(), - ) - monkeypatch.setattr("core.app.apps.workflow.app_generator.DraftVarLoader", lambda **kwargs: SimpleNamespace()) - monkeypatch.setattr( - "core.app.apps.workflow.app_generator.WorkflowDraftVariableService", - lambda session: ( - draft_sessions.append(session) - or SimpleNamespace(prefill_conversation_variable_default_values=lambda *args, **kwargs: None) - ), - ) - monkeypatch.setattr(generator, "_generate", lambda **kwargs: captured.update(kwargs) or {"ok": True}) - - generator.single_iteration_generate( - app_model=app, - workflow=workflow, - node_id="node-1", - user=user, - args={"inputs": {"foo": "bar"}, "trace_session_id": "session-1"}, - streaming=False, - session=sqlite_session, - ) - - assert captured["application_generate_entity"].extras["trace_session_id"] == "session-1" - assert len(repository_session_makers) == 2 - assert all(factory.kw["bind"] is sqlite_session.get_bind() for factory in repository_session_makers) - assert len(draft_sessions) == 1 - assert draft_sessions[0] is sqlite_session - - @pytest.mark.usefixtures("sqlite_generator_session") - def test_single_loop_generate_includes_trace_session_id_in_extras( - self, - monkeypatch: pytest.MonkeyPatch, - sqlite_session: Session, - ): - generator = WorkflowAppGenerator(execution_driver=WorkflowRunAgg.run) - app = _persist_app(sqlite_session) - workflow = _persist_workflow(sqlite_session) - user = _persist_end_user(sqlite_session) - app_config = WorkflowUIBasedAppConfig( - tenant_id=TENANT_ID, - app_id=APP_ID, - app_mode=AppMode.WORKFLOW, - additional_features=AppAdditionalFeatures(), - variables=[], - workflow_id=WORKFLOW_ID, - ) - captured: dict[str, object] = {} - repository_session_makers: list[sessionmaker[Session]] = [] - draft_sessions: list[Session] = [] - - monkeypatch.setattr( - "core.app.apps.workflow.app_generator.WorkflowAppConfigManager.get_app_config", - lambda **kwargs: app_config, - ) - monkeypatch.setattr( - "core.app.apps.workflow.app_generator.DifyCoreRepositoryFactory.create_workflow_execution_repository", - lambda **kwargs: repository_session_makers.append(kwargs["session_factory"]) or SimpleNamespace(), - ) - monkeypatch.setattr( - "core.app.apps.workflow.app_generator.DifyCoreRepositoryFactory.create_workflow_node_execution_repository", - lambda **kwargs: repository_session_makers.append(kwargs["session_factory"]) or SimpleNamespace(), - ) - monkeypatch.setattr("core.app.apps.workflow.app_generator.DraftVarLoader", lambda **kwargs: SimpleNamespace()) - monkeypatch.setattr( - "core.app.apps.workflow.app_generator.WorkflowDraftVariableService", - lambda session: ( - draft_sessions.append(session) - or SimpleNamespace(prefill_conversation_variable_default_values=lambda *args, **kwargs: None) - ), - ) - monkeypatch.setattr(generator, "_generate", lambda **kwargs: captured.update(kwargs) or {"ok": True}) - - generator.single_loop_generate( - app_model=app, - workflow=workflow, - node_id="node-2", - user=user, - args=SimpleNamespace(inputs={"foo": "bar"}, trace_session_id="session-1"), - streaming=False, - session=sqlite_session, - ) - - assert captured["application_generate_entity"].extras["trace_session_id"] == "session-1" - assert len(repository_session_makers) == 2 - assert all(factory.kw["bind"] is sqlite_session.get_bind() for factory in repository_session_makers) - assert len(draft_sessions) == 1 - assert draft_sessions[0] is sqlite_session - - with pytest.raises(ValueError, match="inputs is required"): - generator.single_loop_generate( - app_model=_app(), - workflow=_workflow(), - node_id="node", - user=_end_user(), - args=SimpleNamespace(inputs=None), - streaming=False, - session=sqlite_session, - ) - - -class TestWorkflowAppGeneratorHandleResponse: - def test_handle_response_closed_file_raises_stopped(self, monkeypatch: pytest.MonkeyPatch): - generator = WorkflowAppGenerator(execution_driver=WorkflowRunAgg.run) - - app_config = WorkflowUIBasedAppConfig( - tenant_id="tenant", - app_id="app", - app_mode=AppMode.WORKFLOW, - additional_features=AppAdditionalFeatures(), - variables=[], - workflow_id="workflow-id", - ) - application_generate_entity = WorkflowAppGenerateEntity.model_construct( - task_id="task", - app_config=app_config, - inputs={}, - files=[], - user_id="user", - stream=False, - invoke_from=InvokeFrom.WEB_APP, - extras={}, - trace_manager=None, - workflow_execution_id="run-id", - call_depth=0, - ) - - class _Pipeline: - def __init__(self, **kwargs) -> None: - _ = kwargs - - def process(self): - raise ValueError("I/O operation on closed file.") - - monkeypatch.setattr( - "core.app.apps.workflow.app_generator.WorkflowAppGenerateTaskPipeline", - _Pipeline, - ) - - with pytest.raises(GenerateTaskStoppedError): - generator._handle_response( - application_generate_entity=application_generate_entity, - workflow=_workflow(), - queue_manager=SimpleNamespace(), - user=_end_user(), - draft_var_saver_factory=lambda **kwargs: None, - stream=False, - ) - - -class TestWorkflowAppGeneratorGenerate: - @pytest.fixture - def generation(self, monkeypatch: pytest.MonkeyPatch, sqlite_generator_session: Session): - """Keep input preparation real and stop at the execution boundary.""" - generator = WorkflowAppGenerator(execution_driver=WorkflowRunAgg.run) - app = _persist_app(sqlite_generator_session) - workflow = _persist_workflow(sqlite_generator_session) - user = _persist_end_user(sqlite_generator_session) - execute = Mock(return_value={"ok": True}) - monkeypatch.setattr(generator, "_generate", execute) - return generator, app, workflow, user, execute - - @pytest.mark.parametrize("node_type", ["trigger-webhook", "trigger-schedule", "trigger-plugin"]) - @pytest.mark.parametrize("root_node_id", [None, "selected-trigger"]) - @pytest.mark.parametrize("invoke_from", [InvokeFrom.SERVICE_API, InvokeFrom.DEBUGGER]) - def test_generate_preserves_trigger_inputs(self, generation, node_type, root_node_id, invoke_from): - generator, app, workflow, user, execute = generation - workflow.graph = json.dumps( - { - "nodes": [ - {"id": "first-trigger", "data": {"type": node_type, "title": "First trigger"}}, - {"id": "selected-trigger", "data": {"type": node_type, "title": "Selected trigger"}}, - ], - "edges": [], - } - ) - inputs = {"event": {"message": "hello", "items": [1, 2]}, "enabled": False} - - result = generator.generate( - app_model=app, - workflow=workflow, - user=user, - args={"inputs": inputs}, - invoke_from=invoke_from, - streaming=False, - root_node_id=root_node_id, - ) - - assert result == {"ok": True} - assert execute.call_args.kwargs["application_generate_entity"].inputs == inputs - assert execute.call_args.kwargs["root_node_id"] == (root_node_id or "first-trigger") - - @pytest.mark.parametrize("root_node_id", [None, "start"]) - def test_generate_prepares_start_inputs(self, generation, root_node_id): - generator, app, workflow, user, execute = generation - workflow.graph = json.dumps( - { - "nodes": [ - { - "id": "start", - "data": { - "type": "start", - "title": "Start", - "variables": [ - {"variable": "question", "label": "Question", "type": "text-input", "required": True}, - {"variable": "count", "label": "Count", "type": "number", "default": "3"}, - ], - }, - } - ], - "edges": [], - } - ) - - generator.generate( - app_model=app, - workflow=workflow, - user=user, - args={"inputs": {"question": "hello", "unknown": "filtered"}}, - invoke_from=InvokeFrom.SERVICE_API, - streaming=False, - root_node_id=root_node_id, - ) - - assert execute.call_args.kwargs["application_generate_entity"].inputs == {"question": "hello", "count": 3} - assert execute.call_args.kwargs["root_node_id"] == "start" - - execute.reset_mock() - with pytest.raises(ValueError, match="question is required in input form"): - generator.generate( - app_model=app, - workflow=workflow, - user=user, - args={"inputs": {}}, - invoke_from=InvokeFrom.SERVICE_API, - streaming=False, - root_node_id=root_node_id, - ) - execute.assert_not_called() - - -class TestWorkflowAppGeneratorResume: - def test_resume_restores_trace_manager_when_missing(self, monkeypatch: pytest.MonkeyPatch): - generator = WorkflowAppGenerator(execution_driver=WorkflowRunAgg.run) - app_config = WorkflowUIBasedAppConfig( - tenant_id="tenant", - app_id="app", - app_mode=AppMode.WORKFLOW, - additional_features=AppAdditionalFeatures(), - variables=[], - workflow_id="workflow-id", - ) - application_generate_entity = WorkflowAppGenerateEntity.model_construct( - task_id="task", - app_config=app_config, - inputs={}, - files=[], - user_id="user", - stream=False, - invoke_from=InvokeFrom.WEB_APP, - extras={}, - trace_manager=None, - workflow_execution_id="run-id", - call_depth=0, - ) - dummy_trace_queue_manager = type( - "_DummyTraceQueueManager", - (TraceQueueManager,), - { - "__init__": lambda self, app_id=None, user_id=None: ( - setattr(self, "app_id", app_id) or setattr(self, "user_id", user_id) - ) - }, - ) - monkeypatch.setattr( - "core.app.apps.workflow.app_generator.TraceQueueManager", - dummy_trace_queue_manager, - ) - captured_entity: WorkflowAppGenerateEntity | None = None - - def _fake_generate(**kwargs): - nonlocal captured_entity - captured_entity = kwargs["application_generate_entity"] - return SimpleNamespace(ok=True) - - monkeypatch.setattr(generator, "_generate", _fake_generate) - - result = generator.resume( - app_model=_app(), - workflow=_workflow(), - user=_end_user(), - application_generate_entity=application_generate_entity, - graph_runtime_state=SimpleNamespace(), - workflow_execution_repository=SimpleNamespace(), - workflow_node_execution_repository=SimpleNamespace(), - ) - - assert result.ok is True - assert captured_entity is not None - trace_manager = captured_entity.trace_manager - assert isinstance(trace_manager, dummy_trace_queue_manager) - assert trace_manager.app_id == APP_ID - assert trace_manager.user_id == "session-id" - - def test_resume_preserves_existing_trace_manager(self, monkeypatch: pytest.MonkeyPatch): - generator = WorkflowAppGenerator(execution_driver=WorkflowRunAgg.run) - app_config = WorkflowUIBasedAppConfig( - tenant_id="tenant", - app_id="app", - app_mode=AppMode.WORKFLOW, - additional_features=AppAdditionalFeatures(), - variables=[], - workflow_id="workflow-id", - ) - existing_trace_manager = SimpleNamespace(app_id="existing-app", user_id="existing-user") - application_generate_entity = WorkflowAppGenerateEntity.model_construct( - task_id="task", - app_config=app_config, - inputs={}, - files=[], - user_id="user", - stream=False, - invoke_from=InvokeFrom.WEB_APP, - extras={}, - trace_manager=existing_trace_manager, - workflow_execution_id="run-id", - call_depth=0, - ) - captured_entity: WorkflowAppGenerateEntity | None = None - - def _fake_generate(**kwargs): - nonlocal captured_entity - captured_entity = kwargs["application_generate_entity"] - return SimpleNamespace(ok=True) - - monkeypatch.setattr(generator, "_generate", _fake_generate) - - result = generator.resume( - app_model=_app(), - workflow=_workflow(), - user=_end_user(), - application_generate_entity=application_generate_entity, - graph_runtime_state=SimpleNamespace(), - workflow_execution_repository=SimpleNamespace(), - workflow_node_execution_repository=SimpleNamespace(), - ) - - assert result.ok is True - assert captured_entity is not None - assert captured_entity.trace_manager is existing_trace_manager - - -class TestWorkflowAppGeneratorWorker: - def test_generate_worker_uses_end_user_session_for_external_invocation( - self, - monkeypatch: pytest.MonkeyPatch, - sqlite_session: Session, - ): - generator = WorkflowAppGenerator(execution_driver=WorkflowRunAgg.run) - _persist_app(sqlite_session) - _persist_workflow(sqlite_session) - _persist_end_user(sqlite_session) - - runner_kwargs = {} - - class _Runner: - def __init__(self, **kwargs): - runner_kwargs.update(kwargs) - - def prepare(self): - return None - - monkeypatch.setattr( - "core.app.apps.workflow.app_generator.preserve_flask_contexts", - lambda flask_app, context_vars: contextlib.nullcontext(), - ) - monkeypatch.setattr("core.app.apps.workflow.app_generator.WorkflowAppRunner", _Runner) - restore_workflow_run_graph = Mock() - monkeypatch.setattr(generator, "_restore_workflow_run_graph", restore_workflow_run_graph) - - app_config = WorkflowUIBasedAppConfig( - tenant_id=TENANT_ID, - app_id=APP_ID, - app_mode=AppMode.WORKFLOW, - additional_features=AppAdditionalFeatures(), - variables=[], - workflow_id=WORKFLOW_ID, - ) - application_generate_entity = WorkflowAppGenerateEntity.model_construct( - task_id="task", - app_config=app_config, - inputs={}, - files=[], - user_id=END_USER_ID, - stream=False, - invoke_from=InvokeFrom.WEB_APP, - extras={}, - trace_manager=None, - workflow_execution_id="run-id", - call_depth=0, - ) - - generator._generate_worker( - flask_app=SimpleNamespace(), - application_generate_entity=application_generate_entity, - queue_manager=SimpleNamespace(), - context=SimpleNamespace(), - variable_loader=SimpleNamespace(), - workflow_execution_repository=SimpleNamespace(), - workflow_node_execution_repository=SimpleNamespace(), - graph_runtime_state=SimpleNamespace(), - ) - - assert runner_kwargs["system_user_id"] == "session-id" - restore_workflow_run_graph.assert_called_once() - restore_kwargs = restore_workflow_run_graph.call_args.kwargs - assert isinstance(restore_kwargs["session"], Session) - assert restore_kwargs["workflow"] is runner_kwargs["workflow"] - assert restore_kwargs["workflow_run_id"] == "run-id" - assert inspect(runner_kwargs["workflow"]).detached is True diff --git a/api/tests/unit_tests/core/app/apps/workflow/test_generate_task_pipeline.py b/api/tests/unit_tests/core/app/apps/workflow/test_generate_task_pipeline.py index 06f8bd462a6..a01c791fb64 100644 --- a/api/tests/unit_tests/core/app/apps/workflow/test_generate_task_pipeline.py +++ b/api/tests/unit_tests/core/app/apps/workflow/test_generate_task_pipeline.py @@ -1,30 +1,71 @@ +from __future__ import annotations + import json import time +from types import SimpleNamespace from unittest.mock import MagicMock import pytest +from sqlalchemy import select from sqlalchemy.orm import Session, sessionmaker -from core.app.app_config.entities import WorkflowUIBasedAppConfig +from core.app.app_config.entities import AppAdditionalFeatures, WorkflowUIBasedAppConfig from core.app.apps.workflow.app_queue_manager import WorkflowAppQueueManager from core.app.apps.workflow.generate_task_pipeline import WorkflowAppGenerateTaskPipeline from core.app.entities.app_invoke_entities import InvokeFrom, WorkflowAppGenerateEntity -from core.app.entities.queue_entities import QueueWorkflowStartedEvent -from core.workflow.system_variables import build_system_variables +from core.app.entities.queue_entities import ( + QueueAgentLogEvent, + QueueErrorEvent, + QueueHumanInputFormFilledEvent, + QueueHumanInputFormTimeoutEvent, + QueueIterationCompletedEvent, + QueueIterationNextEvent, + QueueIterationStartEvent, + QueueLoopCompletedEvent, + QueueLoopNextEvent, + QueueLoopStartEvent, + QueueNodeExceptionEvent, + QueueNodeFailedEvent, + QueueNodeRetryEvent, + QueueNodeStartedEvent, + QueueNodeSucceededEvent, + QueuePingEvent, + QueueReasoningChunkEvent, + QueueStopEvent, + QueueTextChunkEvent, + QueueWorkflowFailedEvent, + QueueWorkflowPartialSuccessEvent, + QueueWorkflowPausedEvent, + QueueWorkflowStartedEvent, + QueueWorkflowSucceededEvent, +) +from core.app.entities.task_entities import ( + ErrorStreamResponse, + HumanInputRequiredResponse, + MessageAudioEndStreamResponse, + MessageAudioStreamResponse, + PingStreamResponse, + ReasoningChunkStreamResponse, + WorkflowAppPausedBlockingResponse, + WorkflowFinishStreamResponse, + WorkflowStartStreamResponse, +) +from core.base.tts.app_generator_tts_publisher import AudioTrunk +from core.workflow.system_variables import build_system_variables, system_variables_to_mapping from extensions.ext_redis import RedisClientWrapper from graphon.entities import WorkflowStartReason -from graphon.runtime import RuntimeState +from graphon.enums import BuiltinNodeTypes, WorkflowExecutionStatus +from graphon.model_runtime.entities.llm_entities import LLMUsage +from graphon.runtime import RuntimeState, VariablePool +from libs.datetime_utils import naive_utc_now from models.account import Account -from models.model import AppMode -from models.workflow import Workflow, WorkflowType +from models.enums import CreatorUserRole +from models.model import AppMode, EndUser +from models.workflow import Workflow, WorkflowAppLog, WorkflowType +from tests.helpers.unit.model_factories import make_end_user, make_workflow from tests.workflow_test_utils import build_test_variable_pool -@pytest.fixture(autouse=True) -def bind_queue_redis(monkeypatch: pytest.MonkeyPatch, redis_transport: tuple[RedisClientWrapper, MagicMock]) -> None: - monkeypatch.setattr("core.app.apps.base_app_queue_manager.redis_client", redis_transport[0]) - - def _build_workflow_app_config() -> WorkflowUIBasedAppConfig: return WorkflowUIBasedAppConfig( tenant_id="tenant-id", @@ -80,29 +121,917 @@ def _build_pipeline(run_id: str, unbound_session_factory: sessionmaker[Session]) return pipeline -def test_workflow_app_log_saved_only_on_initial_start(unbound_session_factory: sessionmaker[Session]) -> None: - run_id = "run-initial" - pipeline = _build_pipeline(run_id, unbound_session_factory) - pipeline._save_workflow_app_log = MagicMock() +class TestGenerateTaskPipelineCases: + @staticmethod + @pytest.fixture(autouse=True) + def bind_queue_redis( + monkeypatch: pytest.MonkeyPatch, redis_transport: tuple[RedisClientWrapper, MagicMock] + ) -> None: + monkeypatch.setattr("core.app.apps.base_app_queue_manager.redis_client", redis_transport[0]) - event = QueueWorkflowStartedEvent(reason=WorkflowStartReason.INITIAL) - list(pipeline._handle_workflow_started_event(event)) + @staticmethod + def test_workflow_app_log_saved_only_on_initial_start(unbound_session_factory: sessionmaker[Session]) -> None: + run_id = "run-initial" + pipeline = _build_pipeline(run_id, unbound_session_factory) + pipeline._save_workflow_app_log = MagicMock() - pipeline._save_workflow_app_log.assert_called_once() - _, kwargs = pipeline._save_workflow_app_log.call_args - assert kwargs["workflow_run_id"] == run_id - assert pipeline._workflow_execution_id == run_id + event = QueueWorkflowStartedEvent(reason=WorkflowStartReason.INITIAL) + list(pipeline._handle_workflow_started_event(event)) + + pipeline._save_workflow_app_log.assert_called_once() + _, kwargs = pipeline._save_workflow_app_log.call_args + assert kwargs["workflow_run_id"] == run_id + assert pipeline._workflow_execution_id == run_id + + @staticmethod + def test_workflow_app_log_skipped_on_resumption_start( + unbound_session_factory: sessionmaker[Session], + ) -> None: + run_id = "run-resume" + pipeline = _build_pipeline(run_id, unbound_session_factory) + pipeline._save_workflow_app_log = MagicMock() + + event = QueueWorkflowStartedEvent(reason=WorkflowStartReason.RESUMPTION) + list(pipeline._handle_workflow_started_event(event)) + + pipeline._save_workflow_app_log.assert_not_called() + assert pipeline._workflow_execution_id == run_id -def test_workflow_app_log_skipped_on_resumption_start( - unbound_session_factory: sessionmaker[Session], -) -> None: - run_id = "run-resume" - pipeline = _build_pipeline(run_id, unbound_session_factory) - pipeline._save_workflow_app_log = MagicMock() +def _workflow() -> Workflow: + return make_workflow(workflow_id="workflow-id", tenant_id="tenant", app_id="app", graph="{}", created_by="user") - event = QueueWorkflowStartedEvent(reason=WorkflowStartReason.RESUMPTION) - list(pipeline._handle_workflow_started_event(event)) - pipeline._save_workflow_app_log.assert_not_called() - assert pipeline._workflow_execution_id == run_id +def _end_user(*, end_user_id: str = "user", session_id: str = "session") -> EndUser: + return make_end_user(end_user_id=end_user_id, tenant_id="tenant", app_id="app", session_id=session_id) + + +def _make_pipeline(): + app_config = WorkflowUIBasedAppConfig( + tenant_id="tenant", + app_id="app", + app_mode=AppMode.WORKFLOW, + additional_features=AppAdditionalFeatures(), + variables=[], + workflow_id="workflow-id", + ) + application_generate_entity = WorkflowAppGenerateEntity.model_construct( + task_id="task", + app_config=app_config, + inputs={}, + files=[], + user_id="user", + stream=False, + invoke_from=InvokeFrom.WEB_APP, + trace_manager=None, + workflow_execution_id="run-id", + extras={}, + call_depth=0, + ) + workflow = _workflow() + user = _end_user() + + pipeline = WorkflowAppGenerateTaskPipeline( + application_generate_entity=application_generate_entity, + workflow=workflow, + queue_manager=SimpleNamespace(invoke_from=InvokeFrom.WEB_APP, graph_runtime_state=None), + user=user, + stream=False, + draft_var_saver_factory=lambda **kwargs: None, + ) + + return pipeline + + +class TestGenerateTaskPipelineCoreCases: + class TestWorkflowGenerateTaskPipeline: + def test_to_blocking_response_falls_back_to_human_input_required_when_pause_event_missing(self): + pipeline = _make_pipeline() + pipeline._graph_runtime_state = RuntimeState( + workflow_id="test-workflow", + variable_pool=build_test_variable_pool( + variables=build_system_variables(workflow_execution_id="run-id"), + ), + start_at=0.0, + llm_usage=LLMUsage.empty_usage().model_copy(update={"total_tokens": 5}), + node_run_steps=2, + ) + + def _gen(): + yield HumanInputRequiredResponse( + task_id="task", + workflow_run_id="run-id", + data=HumanInputRequiredResponse.Data( + form_id="form-1", + node_id="node-1", + node_title="Human Input", + form_content="content", + expiration_time=1, + ), + ) + + response = pipeline._to_blocking_response(_gen()) + + assert isinstance(response, WorkflowAppPausedBlockingResponse) + assert response.workflow_run_id == "run-id" + assert response.data.status == WorkflowExecutionStatus.PAUSED + assert response.data.created_at == 0 + assert response.data.paused_nodes == ["node-1"] + assert response.data.reasons == [ + { + "TYPE": "human_input_required", + "form_id": "form-1", + "node_id": "node-1", + "node_title": "Human Input", + "form_content": "content", + "inputs": [], + "actions": [], + "display_in_ui": False, + "form_token": None, + "approval_channels": [], + "resolved_default_values": {}, + "expiration_time": 1, + } + ] + + def test_to_blocking_response_handles_finish(self): + pipeline = _make_pipeline() + + def _gen(): + yield WorkflowFinishStreamResponse( + task_id="task", + workflow_run_id="run", + data=WorkflowFinishStreamResponse.Data( + id="run", + workflow_id="workflow-id", + status=WorkflowExecutionStatus.SUCCEEDED, + outputs={"ok": True}, + error=None, + elapsed_time=1.0, + total_tokens=5, + total_steps=2, + created_at=1, + finished_at=2, + ), + ) + + response = pipeline._to_blocking_response(_gen()) + + assert response.data.outputs == {"ok": True} + + def test_listen_audio_msg_returns_audio_stream(self): + pipeline = _make_pipeline() + publisher = SimpleNamespace(check_and_get_audio=lambda: AudioTrunk(status="responding", audio="data")) + + response = pipeline._listen_audio_msg(publisher=publisher, task_id="task") + + assert isinstance(response, MessageAudioStreamResponse) + + def test_handle_ping_event(self): + pipeline = _make_pipeline() + pipeline._base_task_pipeline.ping_stream_response = lambda: PingStreamResponse(task_id="task") + + responses = list(pipeline._handle_ping_event(QueuePingEvent())) + + assert isinstance(responses[0], PingStreamResponse) + + def test_handle_error_event(self): + pipeline = _make_pipeline() + pipeline._base_task_pipeline.handle_error = lambda **kwargs: ValueError("boom") + pipeline._base_task_pipeline.error_to_stream_response = lambda err: err + + responses = list(pipeline._handle_error_event(QueueErrorEvent(error=ValueError("boom")))) + + assert isinstance(responses[0], ValueError) + + def test_handle_workflow_started_event_sets_run_id(self, monkeypatch: pytest.MonkeyPatch, sqlite_engine): + pipeline = _make_pipeline() + pipeline._graph_runtime_state = RuntimeState( + workflow_id="test-workflow", + variable_pool=build_test_variable_pool( + variables=build_system_variables(workflow_execution_id="run-id") + ), + start_at=0.0, + ) + pipeline._workflow_response_converter.workflow_start_to_stream_response = lambda **kwargs: "started" + + monkeypatch.setattr( + "core.app.apps.workflow.generate_task_pipeline.db", + SimpleNamespace(engine=sqlite_engine), + ) + monkeypatch.setattr(pipeline, "_save_workflow_app_log", lambda **kwargs: None) + + responses = list(pipeline._handle_workflow_started_event(QueueWorkflowStartedEvent())) + + assert pipeline._workflow_execution_id == "run-id" + assert responses == ["started"] + + def test_handle_node_succeeded_event_saves_output(self): + pipeline = _make_pipeline() + pipeline._workflow_response_converter.workflow_node_finish_to_stream_response = lambda **kwargs: "done" + pipeline._save_output_for_event = lambda event, node_execution_id: None + pipeline._workflow_execution_id = "run-id" + + event = QueueNodeSucceededEvent( + node_execution_id="exec", + node_id="node", + node_type=BuiltinNodeTypes.START, + start_at=naive_utc_now(), + inputs={}, + outputs={}, + process_data={}, + ) + + responses = list(pipeline._handle_node_succeeded_event(event)) + + assert responses == ["done"] + + def test_handle_workflow_failed_event_yields_error(self): + pipeline = _make_pipeline() + pipeline._workflow_execution_id = "run-id" + pipeline._graph_runtime_state = RuntimeState( + workflow_id="test-workflow", + variable_pool=build_test_variable_pool( + variables=build_system_variables(workflow_execution_id="run-id") + ), + start_at=0.0, + ) + pipeline._workflow_response_converter.workflow_finish_to_stream_response = lambda **kwargs: "finish" + pipeline._base_task_pipeline.handle_error = lambda **kwargs: ValueError("boom") + pipeline._base_task_pipeline.error_to_stream_response = lambda err: err + + responses = list( + pipeline._handle_workflow_failed_and_stop_events( + QueueWorkflowFailedEvent(error="fail", exceptions_count=1) + ) + ) + + assert responses[0] == "finish" + + def test_handle_text_chunk_event_publishes_tts(self): + pipeline = _make_pipeline() + published: list[object] = [] + + class _Publisher: + def publish(self, message): + published.append(message) + + event = QueueTextChunkEvent(text="hi", from_variable_selector=["x"]) + queue_message = SimpleNamespace(event=event) + + responses = list( + pipeline._handle_text_chunk_event(event, tts_publisher=_Publisher(), queue_message=queue_message) + ) + + assert responses[0].data.text == "hi" + assert published == [queue_message] + + def test_handle_reasoning_chunk_event_emits_on_nonempty(self): + pipeline = _make_pipeline() + event = QueueReasoningChunkEvent(reasoning="pondering", from_node_id="llm-1", is_final=False) + + responses = list(pipeline._handle_reasoning_chunk_event(event)) + + assert len(responses) == 1 + response = responses[0] + assert isinstance(response, ReasoningChunkStreamResponse) + # workflow runs have no message, so the id is omitted + assert response.data.message_id is None + assert response.data.reasoning == "pondering" + assert response.data.node_id == "llm-1" + assert response.data.is_final is False + + def test_handle_reasoning_chunk_event_drops_empty_nonfinal(self): + pipeline = _make_pipeline() + event = QueueReasoningChunkEvent(reasoning="", from_node_id="llm-1", is_final=False) + + responses = list(pipeline._handle_reasoning_chunk_event(event)) + + assert responses == [] + + def test_handle_reasoning_chunk_event_emits_empty_final_marker(self): + pipeline = _make_pipeline() + event = QueueReasoningChunkEvent(reasoning="", from_node_id="llm-1", is_final=True) + + responses = list(pipeline._handle_reasoning_chunk_event(event)) + + assert len(responses) == 1 + response = responses[0] + assert isinstance(response, ReasoningChunkStreamResponse) + assert response.data.reasoning == "" + assert response.data.is_final is True + + def test_dispatch_event_handles_node_failed(self): + pipeline = _make_pipeline() + pipeline._workflow_response_converter.workflow_node_finish_to_stream_response = lambda **kwargs: "done" + + event = QueueNodeFailedEvent( + node_execution_id="exec", + node_id="node", + node_type=BuiltinNodeTypes.START, + start_at=naive_utc_now(), + inputs={}, + outputs={}, + process_data={}, + error="err", + ) + + assert list(pipeline._dispatch_event(event)) == ["done"] + + def test_handle_stop_event_yields_finish(self): + pipeline = _make_pipeline() + pipeline._workflow_execution_id = "run-id" + pipeline._graph_runtime_state = RuntimeState( + workflow_id="test-workflow", + variable_pool=VariablePool.from_bootstrap( + system_variables=build_system_variables(workflow_execution_id="run-id") + ), + start_at=0.0, + ) + pipeline._workflow_response_converter.workflow_finish_to_stream_response = lambda **kwargs: "finish" + + responses = list( + pipeline._handle_workflow_failed_and_stop_events( + QueueStopEvent(stopped_by=QueueStopEvent.StopBy.USER_MANUAL) + ) + ) + + assert responses == ["finish"] + + @pytest.mark.parametrize("sqlite_session", [(WorkflowAppLog,)], indirect=True) + def test_save_workflow_app_log_created_from(self, sqlite_session: Session): + pipeline = _make_pipeline() + pipeline._application_generate_entity.invoke_from = InvokeFrom.SERVICE_API + pipeline._user_id = "user" + pipeline._save_workflow_app_log(session=sqlite_session, workflow_run_id="run-id") + sqlite_session.flush() + + saved_log = sqlite_session.scalar(select(WorkflowAppLog)) + assert saved_log is not None + assert saved_log.workflow_run_id == "run-id" + assert saved_log.created_from == "service-api" + + def test_iteration_loop_and_human_input_handlers(self): + pipeline = _make_pipeline() + pipeline._workflow_execution_id = "run-id" + pipeline._workflow_response_converter.workflow_iteration_start_to_stream_response = lambda **kwargs: "iter" + pipeline._workflow_response_converter.workflow_iteration_next_to_stream_response = lambda **kwargs: "next" + pipeline._workflow_response_converter.workflow_iteration_completed_to_stream_response = lambda **kwargs: ( + "done" + ) + pipeline._workflow_response_converter.workflow_loop_start_to_stream_response = lambda **kwargs: "loop" + pipeline._workflow_response_converter.workflow_loop_next_to_stream_response = lambda **kwargs: "loop_next" + pipeline._workflow_response_converter.workflow_loop_completed_to_stream_response = lambda **kwargs: ( + "loop_done" + ) + pipeline._workflow_response_converter.human_input_form_filled_to_stream_response = lambda **kwargs: "filled" + pipeline._workflow_response_converter.human_input_form_timeout_to_stream_response = lambda **kwargs: ( + "timeout" + ) + pipeline._workflow_response_converter.handle_agent_log = lambda **kwargs: "log" + + iter_start = QueueIterationStartEvent( + node_execution_id="exec", + node_id="node", + node_type=BuiltinNodeTypes.LLM, + node_title="LLM", + start_at=naive_utc_now(), + ) + iter_next = QueueIterationNextEvent( + index=1, + node_execution_id="exec", + node_id="node", + node_type=BuiltinNodeTypes.LLM, + node_title="LLM", + ) + iter_done = QueueIterationCompletedEvent( + node_execution_id="exec", + node_id="node", + node_type=BuiltinNodeTypes.LLM, + node_title="LLM", + start_at=naive_utc_now(), + ) + loop_start = QueueLoopStartEvent( + node_execution_id="exec", + node_id="node", + node_type=BuiltinNodeTypes.LLM, + node_title="LLM", + start_at=naive_utc_now(), + ) + loop_next = QueueLoopNextEvent( + index=1, + node_execution_id="exec", + node_id="node", + node_type=BuiltinNodeTypes.LLM, + node_title="LLM", + ) + loop_done = QueueLoopCompletedEvent( + node_execution_id="exec", + node_id="node", + node_type=BuiltinNodeTypes.LLM, + node_title="LLM", + start_at=naive_utc_now(), + ) + filled_event = QueueHumanInputFormFilledEvent( + form_id="form", + node_id="node", + node_type=BuiltinNodeTypes.LLM, + node_title="title", + rendered_content="content", + action_id="action", + action_text="action", + ) + timeout_event = QueueHumanInputFormTimeoutEvent( + form_id="form", + node_id="node", + node_type=BuiltinNodeTypes.LLM, + node_title="title", + expiration_time=naive_utc_now(), + ) + agent_event = QueueAgentLogEvent( + id="log", + label="label", + node_execution_id="exec", + parent_id=None, + error=None, + status="done", + data={}, + metadata={}, + node_id="node", + ) + + assert list(pipeline._handle_iteration_start_event(iter_start)) == ["iter"] + assert list(pipeline._handle_iteration_next_event(iter_next)) == ["next"] + assert list(pipeline._handle_iteration_completed_event(iter_done)) == ["done"] + assert list(pipeline._handle_loop_start_event(loop_start)) == ["loop"] + assert list(pipeline._handle_loop_next_event(loop_next)) == ["loop_next"] + assert list(pipeline._handle_loop_completed_event(loop_done)) == ["loop_done"] + assert list(pipeline._handle_human_input_form_filled_event(filled_event)) == ["filled"] + assert list(pipeline._handle_human_input_form_timeout_event(timeout_event)) == ["timeout"] + assert list(pipeline._handle_agent_log_event(agent_event)) == ["log"] + + def test_wrapper_process_stream_response_emits_audio_end(self, monkeypatch: pytest.MonkeyPatch): + pipeline = _make_pipeline() + pipeline._base_task_pipeline.stream = True + pipeline._workflow_features_dict = { + "text_to_speech": {"enabled": True, "autoPlay": "enabled", "voice": "v", "language": "en"} + } + pipeline._process_stream_response = lambda **kwargs: iter([PingStreamResponse(task_id="task")]) + + class _Publisher: + def __init__[**P](self, *args: P.args, **kwargs: P.kwargs): + self.calls = 0 + + def check_and_get_audio(self, *, block=False): + self.calls += 1 + if self.calls == 1: + assert not block + return AudioTrunk(status="responding", audio="data") + if self.calls == 2: + assert not block + return None + assert block + return AudioTrunk(status="finish", audio="") + + def publish(self, message): + return None + + def cancel(self): + return None + + monkeypatch.setattr( + "core.app.apps.workflow.generate_task_pipeline.AppGeneratorTTSPublisher", + _Publisher, + ) + + responses = list(pipeline._wrapper_process_stream_response()) + + assert any(isinstance(item, MessageAudioStreamResponse) for item in responses) + assert any(isinstance(item, MessageAudioEndStreamResponse) for item in responses) + + def test_init_with_end_user_sets_role_and_system_user(self): + app_config = WorkflowUIBasedAppConfig( + tenant_id="tenant", + app_id="app", + app_mode=AppMode.WORKFLOW, + additional_features=AppAdditionalFeatures(), + variables=[], + workflow_id="workflow-id", + ) + application_generate_entity = WorkflowAppGenerateEntity.model_construct( + task_id="task", + app_config=app_config, + inputs={}, + files=[], + user_id="end-user-id", + stream=False, + invoke_from=InvokeFrom.WEB_APP, + trace_manager=None, + workflow_execution_id="run-id", + extras={}, + call_depth=0, + ) + workflow = _workflow() + queue_manager = SimpleNamespace(invoke_from=InvokeFrom.WEB_APP, graph_runtime_state=None) + end_user = _end_user(end_user_id="end-user-id", session_id="session-id") + + pipeline = WorkflowAppGenerateTaskPipeline( + application_generate_entity=application_generate_entity, + workflow=workflow, + queue_manager=queue_manager, + user=end_user, + stream=False, + draft_var_saver_factory=lambda **kwargs: None, + ) + + assert pipeline._created_by_role == CreatorUserRole.END_USER + assert system_variables_to_mapping(pipeline._workflow_system_variables)["user_id"] == "session-id" + + def test_process_returns_stream_and_blocking_variants(self): + pipeline = _make_pipeline() + pipeline._base_task_pipeline.stream = True + pipeline._wrapper_process_stream_response = lambda **kwargs: iter([PingStreamResponse(task_id="task")]) + + stream_response = list(pipeline.process()) + assert len(stream_response) == 1 + assert stream_response[0].workflow_run_id is None + + pipeline._base_task_pipeline.stream = False + pipeline._wrapper_process_stream_response = lambda **kwargs: iter( + [ + WorkflowFinishStreamResponse( + task_id="task", + workflow_run_id="run-id", + data=WorkflowFinishStreamResponse.Data( + id="run-id", + workflow_id="workflow-id", + status=WorkflowExecutionStatus.SUCCEEDED, + outputs={}, + error=None, + elapsed_time=0.1, + total_tokens=0, + total_steps=0, + created_at=1, + finished_at=2, + ), + ) + ] + ) + + blocking_response = pipeline.process() + assert blocking_response.workflow_run_id == "run-id" + + def test_to_blocking_response_handles_error_and_unexpected_end(self): + pipeline = _make_pipeline() + + def _error_gen(): + yield ErrorStreamResponse(task_id="task", err=ValueError("boom")) + + with pytest.raises(ValueError, match="boom"): + pipeline._to_blocking_response(_error_gen()) + + def _unexpected_gen(): + yield PingStreamResponse(task_id="task") + + with pytest.raises(ValueError, match="queue listening stopped unexpectedly"): + pipeline._to_blocking_response(_unexpected_gen()) + + def test_to_stream_response_tracks_workflow_run_id(self): + pipeline = _make_pipeline() + + def _gen(): + yield WorkflowStartStreamResponse( + task_id="task", + workflow_run_id="run-id", + data=WorkflowStartStreamResponse.Data( + id="run-id", + workflow_id="workflow-id", + inputs={}, + created_at=1, + ), + ) + yield PingStreamResponse(task_id="task") + + stream_responses = list(pipeline._to_stream_response(_gen())) + assert stream_responses[0].workflow_run_id == "run-id" + assert stream_responses[1].workflow_run_id == "run-id" + + def test_listen_audio_msg_returns_none_without_publisher(self): + pipeline = _make_pipeline() + assert pipeline._listen_audio_msg(publisher=None, task_id="task") is None + + def test_wrapper_process_stream_response_without_tts(self): + pipeline = _make_pipeline() + pipeline._workflow_features_dict = {} + pipeline._process_stream_response = lambda **kwargs: iter([PingStreamResponse(task_id="task")]) + + responses = list(pipeline._wrapper_process_stream_response()) + assert responses == [PingStreamResponse(task_id="task")] + + def test_wrapper_process_stream_response_uses_a_blocking_terminal_read(self, monkeypatch: pytest.MonkeyPatch): + pipeline = _make_pipeline() + pipeline._base_task_pipeline.stream = True + pipeline._workflow_features_dict = { + "text_to_speech": {"enabled": True, "autoPlay": "enabled", "voice": "v", "language": "en"} + } + pipeline._process_stream_response = lambda **kwargs: iter([]) + + blocking_reads = [] + + class _Publisher: + def __init__[**P](self, *args: P.args, **kwargs: P.kwargs): + pass + + def check_and_get_audio(self, *, block=False): + blocking_reads.append(block) + return AudioTrunk(status="finish", audio="") + + def publish(self, message): + _ = message + + def cancel(self): + return None + + monkeypatch.setattr( + "core.app.apps.workflow.generate_task_pipeline.AppGeneratorTTSPublisher", + _Publisher, + ) + + responses = list(pipeline._wrapper_process_stream_response()) + + assert blocking_reads == [True] + assert any(isinstance(item, MessageAudioEndStreamResponse) for item in responses) + + def test_wrapper_process_stream_response_does_not_swallow_audio_queue_exceptions( + self, monkeypatch: pytest.MonkeyPatch + ): + pipeline = _make_pipeline() + pipeline._base_task_pipeline.stream = True + pipeline._workflow_features_dict = { + "text_to_speech": {"enabled": True, "autoPlay": "enabled", "voice": "v", "language": "en"} + } + pipeline._process_stream_response = lambda **kwargs: iter([]) + + class _Publisher: + def __init__[**P](self, *args: P.args, **kwargs: P.kwargs): + pass + + def check_and_get_audio(self, *, block=False): + assert block + raise RuntimeError("tts failure") + + def publish(self, message): + _ = message + + def cancel(self): + return None + + monkeypatch.setattr( + "core.app.apps.workflow.generate_task_pipeline.AppGeneratorTTSPublisher", + _Publisher, + ) + + with pytest.raises(RuntimeError, match="tts failure"): + list(pipeline._wrapper_process_stream_response()) + + @pytest.mark.parametrize("sqlite_session", [(WorkflowAppLog,)], indirect=True) + def test_database_session_rolls_back_on_error( + self, monkeypatch: pytest.MonkeyPatch, sqlite_engine, sqlite_session: Session + ): + pipeline = _make_pipeline() + pipeline._application_generate_entity.invoke_from = InvokeFrom.SERVICE_API + pipeline._user_id = "user" + monkeypatch.setattr( + "core.app.apps.workflow.generate_task_pipeline.db", + SimpleNamespace(engine=sqlite_engine), + ) + + def persist_then_fail() -> None: + with pipeline._database_session() as session: + pipeline._save_workflow_app_log(session=session, workflow_run_id="run-id") + session.flush() + raise RuntimeError("db error") + + with pytest.raises(RuntimeError, match="db error"): + persist_then_fail() + + sqlite_session.expire_all() + assert sqlite_session.scalar(select(WorkflowAppLog)) is None + + def test_node_retry_and_started_handlers_cover_none_and_value(self): + pipeline = _make_pipeline() + pipeline._workflow_execution_id = "run-id" + + retry_event = QueueNodeRetryEvent( + node_execution_id="exec", + node_id="node", + node_title="title", + node_type=BuiltinNodeTypes.LLM, + node_run_index=1, + start_at=naive_utc_now(), + provider_type="provider", + provider_id="provider-id", + error="error", + retry_index=1, + ) + started_event = QueueNodeStartedEvent( + node_execution_id="exec", + node_id="node", + node_title="title", + node_type=BuiltinNodeTypes.LLM, + node_run_index=1, + start_at=naive_utc_now(), + provider_type="provider", + provider_id="provider-id", + ) + + pipeline._workflow_response_converter.workflow_node_retry_to_stream_response = lambda **kwargs: None + assert list(pipeline._handle_node_retry_event(retry_event)) == [] + pipeline._workflow_response_converter.workflow_node_retry_to_stream_response = lambda **kwargs: "retry" + assert list(pipeline._handle_node_retry_event(retry_event)) == ["retry"] + + pipeline._workflow_response_converter.workflow_node_start_to_stream_response = lambda **kwargs: None + assert list(pipeline._handle_node_started_event(started_event)) == [] + pipeline._workflow_response_converter.workflow_node_start_to_stream_response = lambda **kwargs: "started" + assert list(pipeline._handle_node_started_event(started_event)) == ["started"] + + def test_handle_node_exception_event_saves_output(self): + pipeline = _make_pipeline() + saved_ids: list[str] = [] + pipeline._workflow_response_converter.workflow_node_finish_to_stream_response = lambda **kwargs: "failed" + pipeline._save_output_for_event = lambda event, node_execution_id: saved_ids.append(node_execution_id) + + event = QueueNodeExceptionEvent( + node_execution_id="exec-id", + node_id="node", + node_type=BuiltinNodeTypes.START, + start_at=naive_utc_now(), + inputs={}, + outputs={}, + process_data={}, + error="boom", + ) + + responses = list(pipeline._handle_node_failed_events(event)) + assert responses == ["failed"] + assert saved_ids == ["exec-id"] + + def test_success_partial_and_pause_handlers(self): + pipeline = _make_pipeline() + pipeline._workflow_execution_id = "run-id" + pipeline._graph_runtime_state = RuntimeState( + workflow_id="test-workflow", + variable_pool=VariablePool.from_bootstrap( + system_variables=build_system_variables(workflow_execution_id="run-id") + ), + start_at=0.0, + ) + + pipeline._workflow_response_converter.workflow_finish_to_stream_response = lambda **kwargs: "finish" + assert list(pipeline._handle_workflow_succeeded_event(QueueWorkflowSucceededEvent(outputs={}))) == [ + "finish" + ] + assert list( + pipeline._handle_workflow_partial_success_event( + QueueWorkflowPartialSuccessEvent(exceptions_count=2, outputs={}) + ) + ) == ["finish"] + + pipeline._workflow_response_converter.workflow_pause_to_stream_response = lambda **kwargs: [ + "pause-a", + "pause-b", + ] + pause_event = QueueWorkflowPausedEvent(reasons=[], outputs={}, paused_nodes=["node"]) + assert list(pipeline._handle_workflow_paused_event(pause_event)) == ["pause-a", "pause-b"] + + def test_text_chunk_handler_returns_empty_when_text_missing(self): + pipeline = _make_pipeline() + event = QueueTextChunkEvent.model_construct(text=None, from_variable_selector=None) + assert list(pipeline._handle_text_chunk_event(event)) == [] + + def test_dispatch_event_direct_failed_and_unhandled_paths(self): + pipeline = _make_pipeline() + pipeline._workflow_execution_id = "run-id" + pipeline._graph_runtime_state = RuntimeState( + workflow_id="test-workflow", + variable_pool=VariablePool.from_bootstrap( + system_variables=build_system_variables(workflow_execution_id="run-id") + ), + start_at=0.0, + ) + pipeline._handle_ping_event = lambda event, **kwargs: iter(["ping"]) + assert list(pipeline._dispatch_event(QueuePingEvent())) == ["ping"] + + pipeline._handle_workflow_failed_and_stop_events = lambda event, **kwargs: iter(["workflow-failed"]) + assert list(pipeline._dispatch_event(QueueWorkflowFailedEvent(error="failed", exceptions_count=1))) == [ + "workflow-failed" + ] + + assert list(pipeline._dispatch_event(SimpleNamespace())) == [] + + def test_process_stream_response_main_match_paths_and_cleanup(self): + pipeline = _make_pipeline() + pipeline._graph_runtime_state = RuntimeState( + workflow_id="test-workflow", + variable_pool=VariablePool.from_bootstrap( + system_variables=build_system_variables(workflow_execution_id="run-id") + ), + start_at=0.0, + ) + pipeline._base_task_pipeline.queue_manager.listen = lambda: iter( + [ + SimpleNamespace(event=QueueWorkflowStartedEvent()), + SimpleNamespace(event=QueueTextChunkEvent(text="hello")), + SimpleNamespace(event=QueuePingEvent()), + SimpleNamespace(event=QueueErrorEvent(error="e")), + ] + ) + pipeline._handle_workflow_started_event = lambda event, **kwargs: iter(["started"]) + pipeline._handle_text_chunk_event = lambda event, **kwargs: iter(["text"]) + pipeline._dispatch_event = lambda event, **kwargs: iter(["dispatched"]) + pipeline._handle_error_event = lambda event, **kwargs: iter(["error"]) + publisher_calls: list[object] = [] + + class _Publisher: + def publish(self, message): + publisher_calls.append(message) + + responses = list(pipeline._process_stream_response(tts_publisher=_Publisher())) + assert responses == ["started", "text", "dispatched", "error"] + assert publisher_calls == [] + + def test_process_stream_response_break_paths(self): + pipeline = _make_pipeline() + + pipeline._base_task_pipeline.queue_manager.listen = lambda: iter( + [SimpleNamespace(event=QueueWorkflowFailedEvent(error="fail", exceptions_count=1))] + ) + pipeline._handle_workflow_failed_and_stop_events = lambda event, **kwargs: iter(["failed"]) + assert list(pipeline._process_stream_response()) == ["failed"] + + pipeline._base_task_pipeline.queue_manager.listen = lambda: iter( + [SimpleNamespace(event=QueueWorkflowPausedEvent(reasons=[], outputs={}, paused_nodes=[]))] + ) + pipeline._handle_workflow_paused_event = lambda event, **kwargs: iter(["paused"]) + assert list(pipeline._process_stream_response()) == ["paused"] + + pipeline._base_task_pipeline.queue_manager.listen = lambda: iter( + [SimpleNamespace(event=QueueStopEvent(stopped_by=QueueStopEvent.StopBy.USER_MANUAL))] + ) + pipeline._handle_workflow_failed_and_stop_events = lambda event, **kwargs: iter(["stopped"]) + assert list(pipeline._process_stream_response()) == ["stopped"] + + @pytest.mark.parametrize("sqlite_session", [(WorkflowAppLog,)], indirect=True) + def test_save_workflow_app_log_covers_invoke_from_variants(self, sqlite_session: Session): + pipeline = _make_pipeline() + pipeline._user_id = "user-id" + + pipeline._application_generate_entity.invoke_from = InvokeFrom.EXPLORE + pipeline._save_workflow_app_log(session=sqlite_session, workflow_run_id="run-id") + + pipeline._application_generate_entity.invoke_from = InvokeFrom.WEB_APP + pipeline._save_workflow_app_log(session=sqlite_session, workflow_run_id="run-id-2") + sqlite_session.flush() + saved_logs = sqlite_session.scalars(select(WorkflowAppLog).order_by(WorkflowAppLog.workflow_run_id)).all() + assert [log.created_from for log in saved_logs] == ["installed-app", "web-app"] + + count_before = len(saved_logs) + pipeline._application_generate_entity.invoke_from = InvokeFrom.DEBUGGER + pipeline._save_workflow_app_log(session=sqlite_session, workflow_run_id="run-id-3") + sqlite_session.flush() + assert len(sqlite_session.scalars(select(WorkflowAppLog)).all()) == count_before + + pipeline._application_generate_entity.invoke_from = InvokeFrom.WEB_APP + pipeline._save_workflow_app_log(session=sqlite_session, workflow_run_id=None) + sqlite_session.flush() + assert len(sqlite_session.scalars(select(WorkflowAppLog)).all()) == count_before + + def test_save_output_for_event_writes_draft_variables(self): + pipeline = _make_pipeline() + saver_calls: list[tuple[object, object]] = [] + captured_factory_args: dict[str, object] = {} + + class _Saver: + def save(self, process_data, outputs): + saver_calls.append((process_data, outputs)) + + def _factory(**kwargs): + captured_factory_args.update(kwargs) + return _Saver() + + pipeline._draft_var_saver_factory = _factory + + event = QueueNodeSucceededEvent( + node_execution_id="exec-id", + node_id="node-id", + node_type=BuiltinNodeTypes.START, + in_loop_id="loop-id", + start_at=naive_utc_now(), + process_data={"k": "v"}, + outputs={"out": 1}, + ) + pipeline._save_output_for_event(event=event, node_execution_id="exec-id") + + assert captured_factory_args["node_execution_id"] == "exec-id" + assert captured_factory_args["enclosing_node_id"] == "loop-id" + assert saver_calls == [({"k": "v"}, {"out": 1})] diff --git a/api/tests/unit_tests/core/app/apps/workflow/test_generate_task_pipeline_core.py b/api/tests/unit_tests/core/app/apps/workflow/test_generate_task_pipeline_core.py deleted file mode 100644 index a8da5b417ab..00000000000 --- a/api/tests/unit_tests/core/app/apps/workflow/test_generate_task_pipeline_core.py +++ /dev/null @@ -1,923 +0,0 @@ -from __future__ import annotations - -from types import SimpleNamespace - -import pytest -from sqlalchemy import select -from sqlalchemy.orm import Session - -from core.app.app_config.entities import AppAdditionalFeatures, WorkflowUIBasedAppConfig -from core.app.apps.workflow.generate_task_pipeline import WorkflowAppGenerateTaskPipeline -from core.app.entities.app_invoke_entities import InvokeFrom, WorkflowAppGenerateEntity -from core.app.entities.queue_entities import ( - QueueAgentLogEvent, - QueueErrorEvent, - QueueHumanInputFormFilledEvent, - QueueHumanInputFormTimeoutEvent, - QueueIterationCompletedEvent, - QueueIterationNextEvent, - QueueIterationStartEvent, - QueueLoopCompletedEvent, - QueueLoopNextEvent, - QueueLoopStartEvent, - QueueNodeExceptionEvent, - QueueNodeFailedEvent, - QueueNodeRetryEvent, - QueueNodeStartedEvent, - QueueNodeSucceededEvent, - QueuePingEvent, - QueueReasoningChunkEvent, - QueueStopEvent, - QueueTextChunkEvent, - QueueWorkflowFailedEvent, - QueueWorkflowPartialSuccessEvent, - QueueWorkflowPausedEvent, - QueueWorkflowStartedEvent, - QueueWorkflowSucceededEvent, -) -from core.app.entities.task_entities import ( - ErrorStreamResponse, - HumanInputRequiredResponse, - MessageAudioEndStreamResponse, - MessageAudioStreamResponse, - PingStreamResponse, - ReasoningChunkStreamResponse, - WorkflowAppPausedBlockingResponse, - WorkflowFinishStreamResponse, - WorkflowStartStreamResponse, -) -from core.base.tts.app_generator_tts_publisher import AudioTrunk -from core.workflow.system_variables import build_system_variables, system_variables_to_mapping -from graphon.enums import BuiltinNodeTypes, WorkflowExecutionStatus -from graphon.model_runtime.entities.llm_entities import LLMUsage -from graphon.runtime import RuntimeState, VariablePool -from libs.datetime_utils import naive_utc_now -from models.enums import CreatorUserRole -from models.model import AppMode, EndUser -from models.workflow import Workflow, WorkflowAppLog -from tests.unit_tests.model_factories import make_end_user, make_workflow -from tests.workflow_test_utils import build_test_variable_pool - - -def _workflow() -> Workflow: - return make_workflow(workflow_id="workflow-id", tenant_id="tenant", app_id="app", graph="{}", created_by="user") - - -def _end_user(*, end_user_id: str = "user", session_id: str = "session") -> EndUser: - return make_end_user(end_user_id=end_user_id, tenant_id="tenant", app_id="app", session_id=session_id) - - -def _make_pipeline(): - app_config = WorkflowUIBasedAppConfig( - tenant_id="tenant", - app_id="app", - app_mode=AppMode.WORKFLOW, - additional_features=AppAdditionalFeatures(), - variables=[], - workflow_id="workflow-id", - ) - application_generate_entity = WorkflowAppGenerateEntity.model_construct( - task_id="task", - app_config=app_config, - inputs={}, - files=[], - user_id="user", - stream=False, - invoke_from=InvokeFrom.WEB_APP, - trace_manager=None, - workflow_execution_id="run-id", - extras={}, - call_depth=0, - ) - workflow = _workflow() - user = _end_user() - - pipeline = WorkflowAppGenerateTaskPipeline( - application_generate_entity=application_generate_entity, - workflow=workflow, - queue_manager=SimpleNamespace(invoke_from=InvokeFrom.WEB_APP, graph_runtime_state=None), - user=user, - stream=False, - draft_var_saver_factory=lambda **kwargs: None, - ) - - return pipeline - - -class TestWorkflowGenerateTaskPipeline: - def test_to_blocking_response_falls_back_to_human_input_required_when_pause_event_missing(self): - pipeline = _make_pipeline() - pipeline._graph_runtime_state = RuntimeState( - workflow_id="test-workflow", - variable_pool=build_test_variable_pool( - variables=build_system_variables(workflow_execution_id="run-id"), - ), - start_at=0.0, - llm_usage=LLMUsage.empty_usage().model_copy(update={"total_tokens": 5}), - node_run_steps=2, - ) - - def _gen(): - yield HumanInputRequiredResponse( - task_id="task", - workflow_run_id="run-id", - data=HumanInputRequiredResponse.Data( - form_id="form-1", - node_id="node-1", - node_title="Human Input", - form_content="content", - expiration_time=1, - ), - ) - - response = pipeline._to_blocking_response(_gen()) - - assert isinstance(response, WorkflowAppPausedBlockingResponse) - assert response.workflow_run_id == "run-id" - assert response.data.status == WorkflowExecutionStatus.PAUSED - assert response.data.created_at == 0 - assert response.data.paused_nodes == ["node-1"] - assert response.data.reasons == [ - { - "TYPE": "human_input_required", - "form_id": "form-1", - "node_id": "node-1", - "node_title": "Human Input", - "form_content": "content", - "inputs": [], - "actions": [], - "display_in_ui": False, - "form_token": None, - "approval_channels": [], - "resolved_default_values": {}, - "expiration_time": 1, - } - ] - - def test_to_blocking_response_handles_finish(self): - pipeline = _make_pipeline() - - def _gen(): - yield WorkflowFinishStreamResponse( - task_id="task", - workflow_run_id="run", - data=WorkflowFinishStreamResponse.Data( - id="run", - workflow_id="workflow-id", - status=WorkflowExecutionStatus.SUCCEEDED, - outputs={"ok": True}, - error=None, - elapsed_time=1.0, - total_tokens=5, - total_steps=2, - created_at=1, - finished_at=2, - ), - ) - - response = pipeline._to_blocking_response(_gen()) - - assert response.data.outputs == {"ok": True} - - def test_listen_audio_msg_returns_audio_stream(self): - pipeline = _make_pipeline() - publisher = SimpleNamespace(check_and_get_audio=lambda: AudioTrunk(status="responding", audio="data")) - - response = pipeline._listen_audio_msg(publisher=publisher, task_id="task") - - assert isinstance(response, MessageAudioStreamResponse) - - def test_handle_ping_event(self): - pipeline = _make_pipeline() - pipeline._base_task_pipeline.ping_stream_response = lambda: PingStreamResponse(task_id="task") - - responses = list(pipeline._handle_ping_event(QueuePingEvent())) - - assert isinstance(responses[0], PingStreamResponse) - - def test_handle_error_event(self): - pipeline = _make_pipeline() - pipeline._base_task_pipeline.handle_error = lambda **kwargs: ValueError("boom") - pipeline._base_task_pipeline.error_to_stream_response = lambda err: err - - responses = list(pipeline._handle_error_event(QueueErrorEvent(error=ValueError("boom")))) - - assert isinstance(responses[0], ValueError) - - def test_handle_workflow_started_event_sets_run_id(self, monkeypatch: pytest.MonkeyPatch, sqlite_engine): - pipeline = _make_pipeline() - pipeline._graph_runtime_state = RuntimeState( - workflow_id="test-workflow", - variable_pool=build_test_variable_pool(variables=build_system_variables(workflow_execution_id="run-id")), - start_at=0.0, - ) - pipeline._workflow_response_converter.workflow_start_to_stream_response = lambda **kwargs: "started" - - monkeypatch.setattr( - "core.app.apps.workflow.generate_task_pipeline.db", - SimpleNamespace(engine=sqlite_engine), - ) - monkeypatch.setattr(pipeline, "_save_workflow_app_log", lambda **kwargs: None) - - responses = list(pipeline._handle_workflow_started_event(QueueWorkflowStartedEvent())) - - assert pipeline._workflow_execution_id == "run-id" - assert responses == ["started"] - - def test_handle_node_succeeded_event_saves_output(self): - pipeline = _make_pipeline() - pipeline._workflow_response_converter.workflow_node_finish_to_stream_response = lambda **kwargs: "done" - pipeline._save_output_for_event = lambda event, node_execution_id: None - pipeline._workflow_execution_id = "run-id" - - event = QueueNodeSucceededEvent( - node_execution_id="exec", - node_id="node", - node_type=BuiltinNodeTypes.START, - start_at=naive_utc_now(), - inputs={}, - outputs={}, - process_data={}, - ) - - responses = list(pipeline._handle_node_succeeded_event(event)) - - assert responses == ["done"] - - def test_handle_workflow_failed_event_yields_error(self): - pipeline = _make_pipeline() - pipeline._workflow_execution_id = "run-id" - pipeline._graph_runtime_state = RuntimeState( - workflow_id="test-workflow", - variable_pool=build_test_variable_pool(variables=build_system_variables(workflow_execution_id="run-id")), - start_at=0.0, - ) - pipeline._workflow_response_converter.workflow_finish_to_stream_response = lambda **kwargs: "finish" - pipeline._base_task_pipeline.handle_error = lambda **kwargs: ValueError("boom") - pipeline._base_task_pipeline.error_to_stream_response = lambda err: err - - responses = list( - pipeline._handle_workflow_failed_and_stop_events(QueueWorkflowFailedEvent(error="fail", exceptions_count=1)) - ) - - assert responses[0] == "finish" - - def test_handle_text_chunk_event_publishes_tts(self): - pipeline = _make_pipeline() - published: list[object] = [] - - class _Publisher: - def publish(self, message): - published.append(message) - - event = QueueTextChunkEvent(text="hi", from_variable_selector=["x"]) - queue_message = SimpleNamespace(event=event) - - responses = list( - pipeline._handle_text_chunk_event(event, tts_publisher=_Publisher(), queue_message=queue_message) - ) - - assert responses[0].data.text == "hi" - assert published == [queue_message] - - def test_handle_reasoning_chunk_event_emits_on_nonempty(self): - pipeline = _make_pipeline() - event = QueueReasoningChunkEvent(reasoning="pondering", from_node_id="llm-1", is_final=False) - - responses = list(pipeline._handle_reasoning_chunk_event(event)) - - assert len(responses) == 1 - response = responses[0] - assert isinstance(response, ReasoningChunkStreamResponse) - # workflow runs have no message, so the id is omitted - assert response.data.message_id is None - assert response.data.reasoning == "pondering" - assert response.data.node_id == "llm-1" - assert response.data.is_final is False - - def test_handle_reasoning_chunk_event_drops_empty_nonfinal(self): - pipeline = _make_pipeline() - event = QueueReasoningChunkEvent(reasoning="", from_node_id="llm-1", is_final=False) - - responses = list(pipeline._handle_reasoning_chunk_event(event)) - - assert responses == [] - - def test_handle_reasoning_chunk_event_emits_empty_final_marker(self): - pipeline = _make_pipeline() - event = QueueReasoningChunkEvent(reasoning="", from_node_id="llm-1", is_final=True) - - responses = list(pipeline._handle_reasoning_chunk_event(event)) - - assert len(responses) == 1 - response = responses[0] - assert isinstance(response, ReasoningChunkStreamResponse) - assert response.data.reasoning == "" - assert response.data.is_final is True - - def test_dispatch_event_handles_node_failed(self): - pipeline = _make_pipeline() - pipeline._workflow_response_converter.workflow_node_finish_to_stream_response = lambda **kwargs: "done" - - event = QueueNodeFailedEvent( - node_execution_id="exec", - node_id="node", - node_type=BuiltinNodeTypes.START, - start_at=naive_utc_now(), - inputs={}, - outputs={}, - process_data={}, - error="err", - ) - - assert list(pipeline._dispatch_event(event)) == ["done"] - - def test_handle_stop_event_yields_finish(self): - pipeline = _make_pipeline() - pipeline._workflow_execution_id = "run-id" - pipeline._graph_runtime_state = RuntimeState( - workflow_id="test-workflow", - variable_pool=VariablePool.from_bootstrap( - system_variables=build_system_variables(workflow_execution_id="run-id") - ), - start_at=0.0, - ) - pipeline._workflow_response_converter.workflow_finish_to_stream_response = lambda **kwargs: "finish" - - responses = list( - pipeline._handle_workflow_failed_and_stop_events( - QueueStopEvent(stopped_by=QueueStopEvent.StopBy.USER_MANUAL) - ) - ) - - assert responses == ["finish"] - - @pytest.mark.parametrize("sqlite_session", [(WorkflowAppLog,)], indirect=True) - def test_save_workflow_app_log_created_from(self, sqlite_session: Session): - pipeline = _make_pipeline() - pipeline._application_generate_entity.invoke_from = InvokeFrom.SERVICE_API - pipeline._user_id = "user" - pipeline._save_workflow_app_log(session=sqlite_session, workflow_run_id="run-id") - sqlite_session.flush() - - saved_log = sqlite_session.scalar(select(WorkflowAppLog)) - assert saved_log is not None - assert saved_log.workflow_run_id == "run-id" - assert saved_log.created_from == "service-api" - - def test_iteration_loop_and_human_input_handlers(self): - pipeline = _make_pipeline() - pipeline._workflow_execution_id = "run-id" - pipeline._workflow_response_converter.workflow_iteration_start_to_stream_response = lambda **kwargs: "iter" - pipeline._workflow_response_converter.workflow_iteration_next_to_stream_response = lambda **kwargs: "next" - pipeline._workflow_response_converter.workflow_iteration_completed_to_stream_response = lambda **kwargs: "done" - pipeline._workflow_response_converter.workflow_loop_start_to_stream_response = lambda **kwargs: "loop" - pipeline._workflow_response_converter.workflow_loop_next_to_stream_response = lambda **kwargs: "loop_next" - pipeline._workflow_response_converter.workflow_loop_completed_to_stream_response = lambda **kwargs: "loop_done" - pipeline._workflow_response_converter.human_input_form_filled_to_stream_response = lambda **kwargs: "filled" - pipeline._workflow_response_converter.human_input_form_timeout_to_stream_response = lambda **kwargs: "timeout" - pipeline._workflow_response_converter.handle_agent_log = lambda **kwargs: "log" - - iter_start = QueueIterationStartEvent( - node_execution_id="exec", - node_id="node", - node_type=BuiltinNodeTypes.LLM, - node_title="LLM", - start_at=naive_utc_now(), - ) - iter_next = QueueIterationNextEvent( - index=1, - node_execution_id="exec", - node_id="node", - node_type=BuiltinNodeTypes.LLM, - node_title="LLM", - ) - iter_done = QueueIterationCompletedEvent( - node_execution_id="exec", - node_id="node", - node_type=BuiltinNodeTypes.LLM, - node_title="LLM", - start_at=naive_utc_now(), - ) - loop_start = QueueLoopStartEvent( - node_execution_id="exec", - node_id="node", - node_type=BuiltinNodeTypes.LLM, - node_title="LLM", - start_at=naive_utc_now(), - ) - loop_next = QueueLoopNextEvent( - index=1, - node_execution_id="exec", - node_id="node", - node_type=BuiltinNodeTypes.LLM, - node_title="LLM", - ) - loop_done = QueueLoopCompletedEvent( - node_execution_id="exec", - node_id="node", - node_type=BuiltinNodeTypes.LLM, - node_title="LLM", - start_at=naive_utc_now(), - ) - filled_event = QueueHumanInputFormFilledEvent( - form_id="form", - node_id="node", - node_type=BuiltinNodeTypes.LLM, - node_title="title", - rendered_content="content", - action_id="action", - action_text="action", - ) - timeout_event = QueueHumanInputFormTimeoutEvent( - form_id="form", - node_id="node", - node_type=BuiltinNodeTypes.LLM, - node_title="title", - expiration_time=naive_utc_now(), - ) - agent_event = QueueAgentLogEvent( - id="log", - label="label", - node_execution_id="exec", - parent_id=None, - error=None, - status="done", - data={}, - metadata={}, - node_id="node", - ) - - assert list(pipeline._handle_iteration_start_event(iter_start)) == ["iter"] - assert list(pipeline._handle_iteration_next_event(iter_next)) == ["next"] - assert list(pipeline._handle_iteration_completed_event(iter_done)) == ["done"] - assert list(pipeline._handle_loop_start_event(loop_start)) == ["loop"] - assert list(pipeline._handle_loop_next_event(loop_next)) == ["loop_next"] - assert list(pipeline._handle_loop_completed_event(loop_done)) == ["loop_done"] - assert list(pipeline._handle_human_input_form_filled_event(filled_event)) == ["filled"] - assert list(pipeline._handle_human_input_form_timeout_event(timeout_event)) == ["timeout"] - assert list(pipeline._handle_agent_log_event(agent_event)) == ["log"] - - def test_wrapper_process_stream_response_emits_audio_end(self, monkeypatch: pytest.MonkeyPatch): - pipeline = _make_pipeline() - pipeline._base_task_pipeline.stream = True - pipeline._workflow_features_dict = { - "text_to_speech": {"enabled": True, "autoPlay": "enabled", "voice": "v", "language": "en"} - } - pipeline._process_stream_response = lambda **kwargs: iter([PingStreamResponse(task_id="task")]) - - class _Publisher: - def __init__[**P](self, *args: P.args, **kwargs: P.kwargs): - self.calls = 0 - - def check_and_get_audio(self, *, block=False): - self.calls += 1 - if self.calls == 1: - assert not block - return AudioTrunk(status="responding", audio="data") - if self.calls == 2: - assert not block - return None - assert block - return AudioTrunk(status="finish", audio="") - - def publish(self, message): - return None - - def cancel(self): - return None - - monkeypatch.setattr( - "core.app.apps.workflow.generate_task_pipeline.AppGeneratorTTSPublisher", - _Publisher, - ) - - responses = list(pipeline._wrapper_process_stream_response()) - - assert any(isinstance(item, MessageAudioStreamResponse) for item in responses) - assert any(isinstance(item, MessageAudioEndStreamResponse) for item in responses) - - def test_init_with_end_user_sets_role_and_system_user(self): - app_config = WorkflowUIBasedAppConfig( - tenant_id="tenant", - app_id="app", - app_mode=AppMode.WORKFLOW, - additional_features=AppAdditionalFeatures(), - variables=[], - workflow_id="workflow-id", - ) - application_generate_entity = WorkflowAppGenerateEntity.model_construct( - task_id="task", - app_config=app_config, - inputs={}, - files=[], - user_id="end-user-id", - stream=False, - invoke_from=InvokeFrom.WEB_APP, - trace_manager=None, - workflow_execution_id="run-id", - extras={}, - call_depth=0, - ) - workflow = _workflow() - queue_manager = SimpleNamespace(invoke_from=InvokeFrom.WEB_APP, graph_runtime_state=None) - end_user = _end_user(end_user_id="end-user-id", session_id="session-id") - - pipeline = WorkflowAppGenerateTaskPipeline( - application_generate_entity=application_generate_entity, - workflow=workflow, - queue_manager=queue_manager, - user=end_user, - stream=False, - draft_var_saver_factory=lambda **kwargs: None, - ) - - assert pipeline._created_by_role == CreatorUserRole.END_USER - assert system_variables_to_mapping(pipeline._workflow_system_variables)["user_id"] == "session-id" - - def test_process_returns_stream_and_blocking_variants(self): - pipeline = _make_pipeline() - pipeline._base_task_pipeline.stream = True - pipeline._wrapper_process_stream_response = lambda **kwargs: iter([PingStreamResponse(task_id="task")]) - - stream_response = list(pipeline.process()) - assert len(stream_response) == 1 - assert stream_response[0].workflow_run_id is None - - pipeline._base_task_pipeline.stream = False - pipeline._wrapper_process_stream_response = lambda **kwargs: iter( - [ - WorkflowFinishStreamResponse( - task_id="task", - workflow_run_id="run-id", - data=WorkflowFinishStreamResponse.Data( - id="run-id", - workflow_id="workflow-id", - status=WorkflowExecutionStatus.SUCCEEDED, - outputs={}, - error=None, - elapsed_time=0.1, - total_tokens=0, - total_steps=0, - created_at=1, - finished_at=2, - ), - ) - ] - ) - - blocking_response = pipeline.process() - assert blocking_response.workflow_run_id == "run-id" - - def test_to_blocking_response_handles_error_and_unexpected_end(self): - pipeline = _make_pipeline() - - def _error_gen(): - yield ErrorStreamResponse(task_id="task", err=ValueError("boom")) - - with pytest.raises(ValueError, match="boom"): - pipeline._to_blocking_response(_error_gen()) - - def _unexpected_gen(): - yield PingStreamResponse(task_id="task") - - with pytest.raises(ValueError, match="queue listening stopped unexpectedly"): - pipeline._to_blocking_response(_unexpected_gen()) - - def test_to_stream_response_tracks_workflow_run_id(self): - pipeline = _make_pipeline() - - def _gen(): - yield WorkflowStartStreamResponse( - task_id="task", - workflow_run_id="run-id", - data=WorkflowStartStreamResponse.Data( - id="run-id", - workflow_id="workflow-id", - inputs={}, - created_at=1, - ), - ) - yield PingStreamResponse(task_id="task") - - stream_responses = list(pipeline._to_stream_response(_gen())) - assert stream_responses[0].workflow_run_id == "run-id" - assert stream_responses[1].workflow_run_id == "run-id" - - def test_listen_audio_msg_returns_none_without_publisher(self): - pipeline = _make_pipeline() - assert pipeline._listen_audio_msg(publisher=None, task_id="task") is None - - def test_wrapper_process_stream_response_without_tts(self): - pipeline = _make_pipeline() - pipeline._workflow_features_dict = {} - pipeline._process_stream_response = lambda **kwargs: iter([PingStreamResponse(task_id="task")]) - - responses = list(pipeline._wrapper_process_stream_response()) - assert responses == [PingStreamResponse(task_id="task")] - - def test_wrapper_process_stream_response_uses_a_blocking_terminal_read(self, monkeypatch: pytest.MonkeyPatch): - pipeline = _make_pipeline() - pipeline._base_task_pipeline.stream = True - pipeline._workflow_features_dict = { - "text_to_speech": {"enabled": True, "autoPlay": "enabled", "voice": "v", "language": "en"} - } - pipeline._process_stream_response = lambda **kwargs: iter([]) - - blocking_reads = [] - - class _Publisher: - def __init__[**P](self, *args: P.args, **kwargs: P.kwargs): - pass - - def check_and_get_audio(self, *, block=False): - blocking_reads.append(block) - return AudioTrunk(status="finish", audio="") - - def publish(self, message): - _ = message - - def cancel(self): - return None - - monkeypatch.setattr( - "core.app.apps.workflow.generate_task_pipeline.AppGeneratorTTSPublisher", - _Publisher, - ) - - responses = list(pipeline._wrapper_process_stream_response()) - - assert blocking_reads == [True] - assert any(isinstance(item, MessageAudioEndStreamResponse) for item in responses) - - def test_wrapper_process_stream_response_does_not_swallow_audio_queue_exceptions( - self, monkeypatch: pytest.MonkeyPatch - ): - pipeline = _make_pipeline() - pipeline._base_task_pipeline.stream = True - pipeline._workflow_features_dict = { - "text_to_speech": {"enabled": True, "autoPlay": "enabled", "voice": "v", "language": "en"} - } - pipeline._process_stream_response = lambda **kwargs: iter([]) - - class _Publisher: - def __init__[**P](self, *args: P.args, **kwargs: P.kwargs): - pass - - def check_and_get_audio(self, *, block=False): - assert block - raise RuntimeError("tts failure") - - def publish(self, message): - _ = message - - def cancel(self): - return None - - monkeypatch.setattr( - "core.app.apps.workflow.generate_task_pipeline.AppGeneratorTTSPublisher", - _Publisher, - ) - - with pytest.raises(RuntimeError, match="tts failure"): - list(pipeline._wrapper_process_stream_response()) - - @pytest.mark.parametrize("sqlite_session", [(WorkflowAppLog,)], indirect=True) - def test_database_session_rolls_back_on_error( - self, monkeypatch: pytest.MonkeyPatch, sqlite_engine, sqlite_session: Session - ): - pipeline = _make_pipeline() - pipeline._application_generate_entity.invoke_from = InvokeFrom.SERVICE_API - pipeline._user_id = "user" - monkeypatch.setattr( - "core.app.apps.workflow.generate_task_pipeline.db", - SimpleNamespace(engine=sqlite_engine), - ) - - def persist_then_fail() -> None: - with pipeline._database_session() as session: - pipeline._save_workflow_app_log(session=session, workflow_run_id="run-id") - session.flush() - raise RuntimeError("db error") - - with pytest.raises(RuntimeError, match="db error"): - persist_then_fail() - - sqlite_session.expire_all() - assert sqlite_session.scalar(select(WorkflowAppLog)) is None - - def test_node_retry_and_started_handlers_cover_none_and_value(self): - pipeline = _make_pipeline() - pipeline._workflow_execution_id = "run-id" - - retry_event = QueueNodeRetryEvent( - node_execution_id="exec", - node_id="node", - node_title="title", - node_type=BuiltinNodeTypes.LLM, - node_run_index=1, - start_at=naive_utc_now(), - provider_type="provider", - provider_id="provider-id", - error="error", - retry_index=1, - ) - started_event = QueueNodeStartedEvent( - node_execution_id="exec", - node_id="node", - node_title="title", - node_type=BuiltinNodeTypes.LLM, - node_run_index=1, - start_at=naive_utc_now(), - provider_type="provider", - provider_id="provider-id", - ) - - pipeline._workflow_response_converter.workflow_node_retry_to_stream_response = lambda **kwargs: None - assert list(pipeline._handle_node_retry_event(retry_event)) == [] - pipeline._workflow_response_converter.workflow_node_retry_to_stream_response = lambda **kwargs: "retry" - assert list(pipeline._handle_node_retry_event(retry_event)) == ["retry"] - - pipeline._workflow_response_converter.workflow_node_start_to_stream_response = lambda **kwargs: None - assert list(pipeline._handle_node_started_event(started_event)) == [] - pipeline._workflow_response_converter.workflow_node_start_to_stream_response = lambda **kwargs: "started" - assert list(pipeline._handle_node_started_event(started_event)) == ["started"] - - def test_handle_node_exception_event_saves_output(self): - pipeline = _make_pipeline() - saved_ids: list[str] = [] - pipeline._workflow_response_converter.workflow_node_finish_to_stream_response = lambda **kwargs: "failed" - pipeline._save_output_for_event = lambda event, node_execution_id: saved_ids.append(node_execution_id) - - event = QueueNodeExceptionEvent( - node_execution_id="exec-id", - node_id="node", - node_type=BuiltinNodeTypes.START, - start_at=naive_utc_now(), - inputs={}, - outputs={}, - process_data={}, - error="boom", - ) - - responses = list(pipeline._handle_node_failed_events(event)) - assert responses == ["failed"] - assert saved_ids == ["exec-id"] - - def test_success_partial_and_pause_handlers(self): - pipeline = _make_pipeline() - pipeline._workflow_execution_id = "run-id" - pipeline._graph_runtime_state = RuntimeState( - workflow_id="test-workflow", - variable_pool=VariablePool.from_bootstrap( - system_variables=build_system_variables(workflow_execution_id="run-id") - ), - start_at=0.0, - ) - - pipeline._workflow_response_converter.workflow_finish_to_stream_response = lambda **kwargs: "finish" - assert list(pipeline._handle_workflow_succeeded_event(QueueWorkflowSucceededEvent(outputs={}))) == ["finish"] - assert list( - pipeline._handle_workflow_partial_success_event( - QueueWorkflowPartialSuccessEvent(exceptions_count=2, outputs={}) - ) - ) == ["finish"] - - pipeline._workflow_response_converter.workflow_pause_to_stream_response = lambda **kwargs: [ - "pause-a", - "pause-b", - ] - pause_event = QueueWorkflowPausedEvent(reasons=[], outputs={}, paused_nodes=["node"]) - assert list(pipeline._handle_workflow_paused_event(pause_event)) == ["pause-a", "pause-b"] - - def test_text_chunk_handler_returns_empty_when_text_missing(self): - pipeline = _make_pipeline() - event = QueueTextChunkEvent.model_construct(text=None, from_variable_selector=None) - assert list(pipeline._handle_text_chunk_event(event)) == [] - - def test_dispatch_event_direct_failed_and_unhandled_paths(self): - pipeline = _make_pipeline() - pipeline._workflow_execution_id = "run-id" - pipeline._graph_runtime_state = RuntimeState( - workflow_id="test-workflow", - variable_pool=VariablePool.from_bootstrap( - system_variables=build_system_variables(workflow_execution_id="run-id") - ), - start_at=0.0, - ) - pipeline._handle_ping_event = lambda event, **kwargs: iter(["ping"]) - assert list(pipeline._dispatch_event(QueuePingEvent())) == ["ping"] - - pipeline._handle_workflow_failed_and_stop_events = lambda event, **kwargs: iter(["workflow-failed"]) - assert list(pipeline._dispatch_event(QueueWorkflowFailedEvent(error="failed", exceptions_count=1))) == [ - "workflow-failed" - ] - - assert list(pipeline._dispatch_event(SimpleNamespace())) == [] - - def test_process_stream_response_main_match_paths_and_cleanup(self): - pipeline = _make_pipeline() - pipeline._graph_runtime_state = RuntimeState( - workflow_id="test-workflow", - variable_pool=VariablePool.from_bootstrap( - system_variables=build_system_variables(workflow_execution_id="run-id") - ), - start_at=0.0, - ) - pipeline._base_task_pipeline.queue_manager.listen = lambda: iter( - [ - SimpleNamespace(event=QueueWorkflowStartedEvent()), - SimpleNamespace(event=QueueTextChunkEvent(text="hello")), - SimpleNamespace(event=QueuePingEvent()), - SimpleNamespace(event=QueueErrorEvent(error="e")), - ] - ) - pipeline._handle_workflow_started_event = lambda event, **kwargs: iter(["started"]) - pipeline._handle_text_chunk_event = lambda event, **kwargs: iter(["text"]) - pipeline._dispatch_event = lambda event, **kwargs: iter(["dispatched"]) - pipeline._handle_error_event = lambda event, **kwargs: iter(["error"]) - publisher_calls: list[object] = [] - - class _Publisher: - def publish(self, message): - publisher_calls.append(message) - - responses = list(pipeline._process_stream_response(tts_publisher=_Publisher())) - assert responses == ["started", "text", "dispatched", "error"] - assert publisher_calls == [] - - def test_process_stream_response_break_paths(self): - pipeline = _make_pipeline() - - pipeline._base_task_pipeline.queue_manager.listen = lambda: iter( - [SimpleNamespace(event=QueueWorkflowFailedEvent(error="fail", exceptions_count=1))] - ) - pipeline._handle_workflow_failed_and_stop_events = lambda event, **kwargs: iter(["failed"]) - assert list(pipeline._process_stream_response()) == ["failed"] - - pipeline._base_task_pipeline.queue_manager.listen = lambda: iter( - [SimpleNamespace(event=QueueWorkflowPausedEvent(reasons=[], outputs={}, paused_nodes=[]))] - ) - pipeline._handle_workflow_paused_event = lambda event, **kwargs: iter(["paused"]) - assert list(pipeline._process_stream_response()) == ["paused"] - - pipeline._base_task_pipeline.queue_manager.listen = lambda: iter( - [SimpleNamespace(event=QueueStopEvent(stopped_by=QueueStopEvent.StopBy.USER_MANUAL))] - ) - pipeline._handle_workflow_failed_and_stop_events = lambda event, **kwargs: iter(["stopped"]) - assert list(pipeline._process_stream_response()) == ["stopped"] - - @pytest.mark.parametrize("sqlite_session", [(WorkflowAppLog,)], indirect=True) - def test_save_workflow_app_log_covers_invoke_from_variants(self, sqlite_session: Session): - pipeline = _make_pipeline() - pipeline._user_id = "user-id" - - pipeline._application_generate_entity.invoke_from = InvokeFrom.EXPLORE - pipeline._save_workflow_app_log(session=sqlite_session, workflow_run_id="run-id") - - pipeline._application_generate_entity.invoke_from = InvokeFrom.WEB_APP - pipeline._save_workflow_app_log(session=sqlite_session, workflow_run_id="run-id-2") - sqlite_session.flush() - saved_logs = sqlite_session.scalars(select(WorkflowAppLog).order_by(WorkflowAppLog.workflow_run_id)).all() - assert [log.created_from for log in saved_logs] == ["installed-app", "web-app"] - - count_before = len(saved_logs) - pipeline._application_generate_entity.invoke_from = InvokeFrom.DEBUGGER - pipeline._save_workflow_app_log(session=sqlite_session, workflow_run_id="run-id-3") - sqlite_session.flush() - assert len(sqlite_session.scalars(select(WorkflowAppLog)).all()) == count_before - - pipeline._application_generate_entity.invoke_from = InvokeFrom.WEB_APP - pipeline._save_workflow_app_log(session=sqlite_session, workflow_run_id=None) - sqlite_session.flush() - assert len(sqlite_session.scalars(select(WorkflowAppLog)).all()) == count_before - - def test_save_output_for_event_writes_draft_variables(self): - pipeline = _make_pipeline() - saver_calls: list[tuple[object, object]] = [] - captured_factory_args: dict[str, object] = {} - - class _Saver: - def save(self, process_data, outputs): - saver_calls.append((process_data, outputs)) - - def _factory(**kwargs): - captured_factory_args.update(kwargs) - return _Saver() - - pipeline._draft_var_saver_factory = _factory - - event = QueueNodeSucceededEvent( - node_execution_id="exec-id", - node_id="node-id", - node_type=BuiltinNodeTypes.START, - in_loop_id="loop-id", - start_at=naive_utc_now(), - process_data={"k": "v"}, - outputs={"out": 1}, - ) - pipeline._save_output_for_event(event=event, node_execution_id="exec-id") - - assert captured_factory_args["node_execution_id"] == "exec-id" - assert captured_factory_args["enclosing_node_id"] == "loop-id" - assert saver_calls == [({"k": "v"}, {"out": 1})] diff --git a/api/tests/unit_tests/core/app/entities/test_app_invoke_entities.py b/api/tests/unit_tests/core/app/entities/test_app_invoke_entities.py index 60d856dd416..87c6dac64d5 100644 --- a/api/tests/unit_tests/core/app/entities/test_app_invoke_entities.py +++ b/api/tests/unit_tests/core/app/entities/test_app_invoke_entities.py @@ -200,3 +200,27 @@ def test_workflow_resumption_context_roundtrip(case: ResumptionContextCase): assert isinstance(entity, expected_type) assert entity.model_dump() == context.get_generate_entity().model_dump() assert entity.trace_manager is None + + +def test_openapi_variant_present(): + assert InvokeFrom.OPENAPI.value == "openapi" + + +def test_openapi_distinct_from_service_api(): + assert InvokeFrom.OPENAPI != InvokeFrom.SERVICE_API + + +def test_runs_as_account_only_for_console_contexts(): + # Console contexts (studio debugger / explore) run as the signed-in Account. + assert InvokeFrom.DEBUGGER.runs_as_account() is True + assert InvokeFrom.EXPLORE.runs_as_account() is True + # Everything else is attributed to an end user. + for invoke_from in ( + InvokeFrom.WEB_APP, + InvokeFrom.SERVICE_API, + InvokeFrom.OPENAPI, + InvokeFrom.TRIGGER, + InvokeFrom.PUBLISHED_PIPELINE, + InvokeFrom.VALIDATION, + ): + assert invoke_from.runs_as_account() is False diff --git a/api/tests/unit_tests/core/app/features/test_annotation_reply.py b/api/tests/unit_tests/core/app/features/annotation_reply/test_annotation_reply.py similarity index 98% rename from api/tests/unit_tests/core/app/features/test_annotation_reply.py rename to api/tests/unit_tests/core/app/features/annotation_reply/test_annotation_reply.py index a9001c96c61..ccb8c7c5a75 100644 --- a/api/tests/unit_tests/core/app/features/test_annotation_reply.py +++ b/api/tests/unit_tests/core/app/features/annotation_reply/test_annotation_reply.py @@ -11,7 +11,7 @@ from core.app.features.annotation_reply.annotation_reply import AnnotationReplyF from models.dataset import DatasetCollectionBinding from models.enums import CollectionBindingType, ConversationFromSource from models.model import App, AppAnnotationHitHistory, AppAnnotationSetting, Message, MessageAnnotation -from tests.unit_tests.model_factories import make_app, make_message +from tests.helpers.unit.model_factories import make_app, make_message TABLES = (AppAnnotationSetting, DatasetCollectionBinding, MessageAnnotation, AppAnnotationHitHistory) diff --git a/api/tests/unit_tests/core/app/features/test_hosting_moderation.py b/api/tests/unit_tests/core/app/features/hosting_moderation/test_hosting_moderation.py similarity index 100% rename from api/tests/unit_tests/core/app/features/test_hosting_moderation.py rename to api/tests/unit_tests/core/app/features/hosting_moderation/test_hosting_moderation.py diff --git a/api/tests/unit_tests/core/app/layers/test_execution_context_layer.py b/api/tests/unit_tests/core/app/layers/test_execution_context_layer.py index 3a4ab5c5290..f2f69dea4a8 100644 --- a/api/tests/unit_tests/core/app/layers/test_execution_context_layer.py +++ b/api/tests/unit_tests/core/app/layers/test_execution_context_layer.py @@ -18,10 +18,14 @@ from graphon.engine_events import GraphRunSucceededEvent from graphon.entities.base_node_data import BaseNodeData from graphon.file.runtime import peek_workflow_file_runtime, use_workflow_file_runtime from graphon.nodes.base.node import Node -from tests.unit_tests.core.workflow.test_workflow_tool_container import ( - _outer_graph, - _workflow_tool_node, - _workflow_tool_source, +from tests.unit_tests.core.workflow.test_workflow_tool_container_handler import ( + _workflow_tool_container_outer_graph as _outer_graph, +) +from tests.unit_tests.core.workflow.test_workflow_tool_container_handler import ( + _workflow_tool_container_workflow_tool_node as _workflow_tool_node, +) +from tests.unit_tests.core.workflow.test_workflow_tool_container_handler import ( + _workflow_tool_container_workflow_tool_source as _workflow_tool_source, ) diff --git a/api/tests/unit_tests/core/app/test_llm_quota.py b/api/tests/unit_tests/core/app/llm/test_quota.py similarity index 99% rename from api/tests/unit_tests/core/app/test_llm_quota.py rename to api/tests/unit_tests/core/app/llm/test_quota.py index e915ee88f17..2a0c413a103 100644 --- a/api/tests/unit_tests/core/app/test_llm_quota.py +++ b/api/tests/unit_tests/core/app/llm/test_quota.py @@ -43,7 +43,7 @@ from models.enums import ProviderQuotaType as ModelProviderQuotaType from models.provider import Provider, ProviderType from models.provider_ids import ModelProviderID from services.credit_pool_service import CreditPoolReservation, CreditPoolReservationState -from tests.unit_tests.core.model_fixtures import make_model_config +from tests.helpers.unit.core.model_fixtures import make_model_config def _model_instance(model_type: ModelType) -> ModelInstance: diff --git a/api/tests/unit_tests/core/app/task_pipeline/test_easy_ui_based_generate_task_pipeline.py b/api/tests/unit_tests/core/app/task_pipeline/test_easy_ui_based_generate_task_pipeline.py index 82bc7710f8d..f6ecac635fe 100644 --- a/api/tests/unit_tests/core/app/task_pipeline/test_easy_ui_based_generate_task_pipeline.py +++ b/api/tests/unit_tests/core/app/task_pipeline/test_easy_ui_based_generate_task_pipeline.py @@ -1,33 +1,56 @@ -from collections.abc import Generator, Iterator +from __future__ import annotations + +import json +import uuid +from collections.abc import Generator, Iterator, Sequence +from dataclasses import dataclass from datetime import UTC, datetime -from unittest.mock import Mock, patch +from decimal import Decimal +from threading import Thread +from typing import cast +from unittest.mock import MagicMock, Mock, patch import pytest -from sqlalchemy import Engine, event -from sqlalchemy.orm import Session +from sqlalchemy import Engine as _easy_ui_based_generate_task_pipeline_Engine +from sqlalchemy import event +from sqlalchemy.engine import Engine as _easy_ui_message_end_files_Engine +from sqlalchemy.orm import Session, sessionmaker from core.app.app_config.entities import ( + AppAdditionalFeatures, EasyUIBasedAppConfig, EasyUIBasedAppModelConfigFrom, ModelConfigEntity, PromptTemplateEntity, ) from core.app.apps.base_app_queue_manager import AppQueueManager, PublishFrom -from core.app.entities.app_invoke_entities import ChatAppGenerateEntity, InvokeFrom +from core.app.apps.message_based_app_queue_manager import MessageBasedAppQueueManager +from core.app.entities.app_invoke_entities import ChatAppGenerateEntity, CompletionAppGenerateEntity, InvokeFrom from core.app.entities.queue_entities import ( AppQueueEvent, MessageQueueMessage, QueueAgentMessageEvent, + QueueAgentThoughtEvent, + QueueAnnotationReplyEvent, QueueErrorEvent, QueueLLMChunkEvent, QueueMessageEndEvent, QueueMessageFileEvent, + QueueMessageReplaceEvent, QueuePingEvent, + QueueRetrieverResourcesEvent, + QueueStopEvent, WorkflowQueueMessage, ) from core.app.entities.task_entities import ( + AgentMessageStreamResponse, + AgentThoughtStreamResponse, + ChatbotAppStreamResponse, + CompletionAppStreamResponse, EasyUITaskState, ErrorStreamResponse, + MessageAudioEndStreamResponse, + MessageAudioStreamResponse, MessageEndStreamResponse, MessageFileStreamResponse, MessageReplaceStreamResponse, @@ -36,14 +59,19 @@ from core.app.entities.task_entities import ( StreamEvent, ) from core.app.task_pipeline.easy_ui_based_generate_task_pipeline import EasyUIBasedGenerateTaskPipeline -from core.base.tts import AppGeneratorTTSPublisher +from core.base.tts import AppGeneratorTTSPublisher, AudioTrunk +from core.ops.entities.trace_entity import TraceTaskName from core.ops.ops_trace_manager import TraceQueueManager +from extensions.ext_redis import RedisClientWrapper +from extensions.storage.storage_type import StorageType +from graphon.file import FileTransferMethod, FileType +from graphon.model_runtime.entities.llm_entities import LLMResult, LLMResultChunk, LLMResultChunkDelta, LLMUsage from graphon.model_runtime.entities.llm_entities import LLMResult as RuntimeLLMResult -from graphon.model_runtime.entities.llm_entities import LLMResultChunk, LLMResultChunkDelta, LLMUsage from graphon.model_runtime.entities.message_entities import AssistantPromptMessage, TextPromptMessageContent -from models.enums import ConversationFromSource -from models.model import AppMode, Conversation, Message -from tests.unit_tests.core.model_fixtures import make_model_config +from models.enums import ConversationFromSource, CreatorUserRole +from models.model import AppMode, Conversation, Message, MessageAgentThought, MessageFile, UploadFile +from tests.helpers.unit.core.model_fixtures import make_model_config +from tests.helpers.unit.model_factories import make_upload_file class _QueueManager(AppQueueManager): @@ -92,371 +120,2379 @@ def llm_result(content: str) -> RuntimeLLMResult: ) -@pytest.fixture -def committed_sessions(sqlite_engine: Engine) -> Iterator[list[Session]]: - sessions: list[Session] = [] - - def record_commit(session: Session) -> None: - if session.get_bind() is sqlite_engine: - sessions.append(session) - - event.listen(Session, "after_commit", record_commit) - yield sessions - event.remove(Session, "after_commit", record_commit) - - -class TestEasyUIBasedGenerateTaskPipelineProcessStreamResponse: - """Test cases for EasyUIBasedGenerateTaskPipeline._process_stream_response method.""" - +class TestEasyUiBasedGenerateTaskPipelineCases: + @staticmethod @pytest.fixture - def application_generate_entity(self): - """Create a validated application request with real model configuration.""" - return ChatAppGenerateEntity( - task_id="test-task-id", - app_config=EasyUIBasedAppConfig( - tenant_id="test-tenant-id", + def committed_sessions(sqlite_engine: _easy_ui_based_generate_task_pipeline_Engine) -> Iterator[list[Session]]: + sessions: list[Session] = [] + + def record_commit(session: Session) -> None: + if session.get_bind() is sqlite_engine: + sessions.append(session) + + event.listen(Session, "after_commit", record_commit) + yield sessions + event.remove(Session, "after_commit", record_commit) + + class TestEasyUIBasedGenerateTaskPipelineProcessStreamResponse: + """Test cases for EasyUIBasedGenerateTaskPipeline._process_stream_response method.""" + + @pytest.fixture + def application_generate_entity(self): + """Create a validated application request with real model configuration.""" + return ChatAppGenerateEntity( + task_id="test-task-id", + app_config=EasyUIBasedAppConfig( + tenant_id="test-tenant-id", + app_id="test-app-id", + app_mode=AppMode.CHAT, + app_model_config_from=EasyUIBasedAppModelConfigFrom.APP_LATEST_CONFIG, + app_model_config_dict={}, + model=ModelConfigEntity(provider="test-provider", model="test-model"), + prompt_template=PromptTemplateEntity(prompt_type=PromptTemplateEntity.PromptType.SIMPLE), + ), + model_conf=make_model_config(provider="test-provider", model="test-model", mode="chat"), + inputs={}, + files=[], + user_id="test-user-id", + stream=True, + invoke_from=InvokeFrom.WEB_APP, + ) + + @pytest.fixture + def queue_manager(self): + return _QueueManager() + + @pytest.fixture + def mock_message_cycle_manager(self): + """Create a mock message cycle manager.""" + manager = Mock() + manager.get_message_event_type.return_value = StreamEvent.MESSAGE + manager.message_to_stream_response.return_value = MessageStreamResponse( + task_id="test-task-id", id="test-message-id", answer="response" + ) + manager.message_file_to_stream_response.return_value = MessageFileStreamResponse( + task_id="test-task-id", + id="file-id", + type="image", + belongs_to="assistant", + url="https://example.com/file.png", + ) + manager.message_replace_to_stream_response.return_value = MessageReplaceStreamResponse( + task_id="test-task-id", answer="replacement", reason="moderation" + ) + manager.handle_retriever_resources = Mock() + manager.handle_annotation_reply.return_value = None + return manager + + @pytest.fixture + def conversation(self): + """Create a transient mapped conversation.""" + return Conversation( + id="test-conversation-id", app_id="test-app-id", - app_mode=AppMode.CHAT, - app_model_config_from=EasyUIBasedAppModelConfigFrom.APP_LATEST_CONFIG, - app_model_config_dict={}, - model=ModelConfigEntity(provider="test-provider", model="test-model"), - prompt_template=PromptTemplateEntity(prompt_type=PromptTemplateEntity.PromptType.SIMPLE), - ), - model_conf=make_model_config(provider="test-provider", model="test-model", mode="chat"), - inputs={}, - files=[], - user_id="test-user-id", - stream=True, - invoke_from=InvokeFrom.WEB_APP, - ) + mode=AppMode.CHAT, + name="Test Conversation", + status="normal", + from_source=ConversationFromSource.API, + _inputs={}, + ) - @pytest.fixture - def queue_manager(self): - return _QueueManager() + @pytest.fixture + def message(self): + """Create a transient mapped message.""" + message = Message(id="test-message-id") + message.created_at = datetime.fromtimestamp(1234567890, tz=UTC) + return message - @pytest.fixture - def mock_message_cycle_manager(self): - """Create a mock message cycle manager.""" - manager = Mock() - manager.get_message_event_type.return_value = StreamEvent.MESSAGE - manager.message_to_stream_response.return_value = MessageStreamResponse( - task_id="test-task-id", id="test-message-id", answer="response" - ) - manager.message_file_to_stream_response.return_value = MessageFileStreamResponse( - task_id="test-task-id", - id="file-id", - type="image", - belongs_to="assistant", - url="https://example.com/file.png", - ) - manager.message_replace_to_stream_response.return_value = MessageReplaceStreamResponse( - task_id="test-task-id", answer="replacement", reason="moderation" - ) - manager.handle_retriever_resources = Mock() - manager.handle_annotation_reply.return_value = None - return manager + @pytest.fixture + def task_state(self): + """Create real mutable state for the stream response pipeline.""" + return EasyUITaskState(llm_result=llm_result("")) - @pytest.fixture - def conversation(self): - """Create a transient mapped conversation.""" - return Conversation( - id="test-conversation-id", - app_id="test-app-id", - mode=AppMode.CHAT, - name="Test Conversation", - status="normal", - from_source=ConversationFromSource.API, - _inputs={}, - ) - - @pytest.fixture - def message(self): - """Create a transient mapped message.""" - message = Message(id="test-message-id") - message.created_at = datetime.fromtimestamp(1234567890, tz=UTC) - return message - - @pytest.fixture - def task_state(self): - """Create real mutable state for the stream response pipeline.""" - return EasyUITaskState(llm_result=llm_result("")) - - @pytest.fixture - def pipeline( - self, - application_generate_entity, - queue_manager, - conversation, - message, - mock_message_cycle_manager, - task_state, - ): - """Create an EasyUIBasedGenerateTaskPipeline instance with mocked dependencies.""" - pipeline = EasyUIBasedGenerateTaskPipeline( - application_generate_entity=application_generate_entity, - queue_manager=queue_manager, - conversation=conversation, - message=message, - stream=True, - ) - pipeline._message_cycle_manager = mock_message_cycle_manager - pipeline._task_state = task_state - return pipeline - - def test_get_message_event_type_called_once_when_first_llm_chunk_arrives( - self, pipeline, mock_message_cycle_manager - ): - """Expect get_message_event_type to be called when processing the first LLM chunk event.""" - # Setup a minimal LLM chunk event - chunk = llm_chunk("hi") - llm_chunk_event = QueueLLMChunkEvent(chunk=chunk) - message = queue_message(llm_chunk_event) - pipeline.queue_manager.messages = [message] - - # Execute - list(pipeline._process_stream_response(publisher=None, trace_manager=None)) - - # Assert - mock_message_cycle_manager.get_message_event_type.assert_called_once_with(message_id="test-message-id") - - def test_llm_chunk_event_with_text_content(self, pipeline, mock_message_cycle_manager, task_state): - """Test handling of LLM chunk events with text content.""" - # Setup - chunk = llm_chunk("Hello, world!") - - llm_chunk_event = QueueLLMChunkEvent(chunk=chunk) - - message = queue_message(llm_chunk_event) - pipeline.queue_manager.messages = [message] - - mock_message_cycle_manager.get_message_event_type.return_value = StreamEvent.MESSAGE - - # Execute - responses = list(pipeline._process_stream_response(publisher=None, trace_manager=None)) - - # Assert - assert len(responses) == 1 - mock_message_cycle_manager.message_to_stream_response.assert_called_once_with( - answer="Hello, world!", message_id="test-message-id", event_type=StreamEvent.MESSAGE - ) - assert task_state.llm_result.message.content == "Hello, world!" - - def test_llm_chunk_event_with_list_content(self, pipeline, mock_message_cycle_manager, task_state): - """Test handling of LLM chunk events with list content.""" - # Setup - text_content = TextPromptMessageContent(data="Hello") - chunk = llm_chunk("") - # Preserve the legacy mixed-list payload handled by _chunk_delta_text. - chunk.delta.message = AssistantPromptMessage.model_construct(content=[text_content, " world!"]) - - llm_chunk_event = QueueLLMChunkEvent(chunk=chunk) - - message = queue_message(llm_chunk_event) - pipeline.queue_manager.messages = [message] - - mock_message_cycle_manager.get_message_event_type.return_value = StreamEvent.MESSAGE - - # Execute - responses = list(pipeline._process_stream_response(publisher=None, trace_manager=None)) - - # Assert - assert len(responses) == 1 - mock_message_cycle_manager.message_to_stream_response.assert_called_once_with( - answer="Hello world!", message_id="test-message-id", event_type=StreamEvent.MESSAGE - ) - assert task_state.llm_result.message.content == "Hello world!" - - def test_agent_message_event(self, pipeline, mock_message_cycle_manager, task_state): - """Test handling of agent message events.""" - # Setup - chunk = llm_chunk("Agent response") - - agent_message_event = QueueAgentMessageEvent(chunk=chunk) - - message = queue_message(agent_message_event) - pipeline.queue_manager.messages = [message] - - # Ensure method under assertion is a mock to track calls - pipeline._agent_message_to_stream_response = Mock(return_value=Mock()) - - # Execute - responses = list(pipeline._process_stream_response(publisher=None, trace_manager=None)) - - # Assert - assert len(responses) == 1 - # Agent messages should use _agent_message_to_stream_response - pipeline._agent_message_to_stream_response.assert_called_once_with( - answer="Agent response", message_id="test-message-id" - ) - - def test_message_end_event(self, pipeline, mock_message_cycle_manager, task_state, committed_sessions): - """Test handling of message end events.""" - # Setup - result = llm_result("Final response") - message_end_event = QueueMessageEndEvent(llm_result=result) - - message = queue_message(message_end_event) - pipeline.queue_manager.messages = [message] - - pipeline._save_message = Mock() - pipeline._message_end_to_stream_response = Mock( - return_value=MessageEndStreamResponse(task_id="test-task-id", id="test-message-id") - ) - - responses = list(pipeline._process_stream_response(publisher=None, trace_manager=None)) - - # Assert - assert len(responses) == 1 - assert task_state.llm_result == result - session = pipeline._save_message.call_args.kwargs["session"] - assert isinstance(session, Session) - assert committed_sessions == [session] - pipeline._message_end_to_stream_response.assert_called_once() - - def test_error_event(self, pipeline, committed_sessions): - """Test handling of error events.""" - # Setup - error_event = QueueErrorEvent(error=Exception("Test error")) - - message = queue_message(error_event) - pipeline.queue_manager.messages = [message] - - pipeline.handle_error = Mock(return_value=Exception("Test error")) - pipeline.error_to_stream_response = Mock( - return_value=ErrorStreamResponse(task_id="test-task-id", err=Exception("Test error")) - ) - - responses = list(pipeline._process_stream_response(publisher=None, trace_manager=None)) - - # Assert - assert len(responses) == 1 - session = pipeline.handle_error.call_args.kwargs["session"] - assert isinstance(session, Session) - assert committed_sessions == [session] - pipeline.handle_error.assert_called_once_with(event=error_event, session=session, message_id="test-message-id") - pipeline.error_to_stream_response.assert_called_once() - - def test_ping_event(self, pipeline): - """Test handling of ping events.""" - # Setup - ping_event = QueuePingEvent() - - message = queue_message(ping_event) - pipeline.queue_manager.messages = [message] - - pipeline.ping_stream_response = Mock(return_value=PingStreamResponse(task_id="test-task-id")) - - # Execute - responses = list(pipeline._process_stream_response(publisher=None, trace_manager=None)) - - # Assert - assert len(responses) == 1 - pipeline.ping_stream_response.assert_called_once() - - def test_file_event(self, pipeline, mock_message_cycle_manager): - """Test handling of file events.""" - # Setup - file_event = QueueMessageFileEvent(message_file_id="file-id") - - message = queue_message(file_event) - pipeline.queue_manager.messages = [message] - - file_response = MessageFileStreamResponse( - task_id="test-task-id", - id="file-id", - type="image", - belongs_to="assistant", - url="https://example.com/file.png", - ) - mock_message_cycle_manager.message_file_to_stream_response.return_value = file_response - - # Execute - responses = list(pipeline._process_stream_response(publisher=None, trace_manager=None)) - - # Assert - assert len(responses) == 1 - assert responses[0] == file_response - mock_message_cycle_manager.message_file_to_stream_response.assert_called_once_with(file_event) - - def test_publisher_is_called_with_messages(self, pipeline): - """Test that publisher publishes messages when provided.""" - # Setup - publisher = _Publisher() - - ping_event = QueuePingEvent() - message = queue_message(ping_event) - pipeline.queue_manager.messages = [message] - - pipeline.ping_stream_response = Mock(return_value=PingStreamResponse(task_id="test-task-id")) - - # Execute - list(pipeline._process_stream_response(publisher=publisher, trace_manager=None)) - - # Assert - assert publisher.messages == [message] - - def test_trace_manager_passed_to_save_message(self, pipeline, committed_sessions): - """Test that trace manager is passed to _save_message.""" - # Setup - with ( - patch("core.ops.ops_trace_manager.OpsTraceManager.get_ops_trace_instance", return_value=None), - patch.object(TraceQueueManager, "start_timer"), + @pytest.fixture + def pipeline( + self, + application_generate_entity, + queue_manager, + conversation, + message, + mock_message_cycle_manager, + task_state, ): - trace_manager = TraceQueueManager(app_id="test-app-id") + """Create an EasyUIBasedGenerateTaskPipeline instance with mocked dependencies.""" + pipeline = EasyUIBasedGenerateTaskPipeline( + application_generate_entity=application_generate_entity, + queue_manager=queue_manager, + conversation=conversation, + message=message, + stream=True, + ) + pipeline._message_cycle_manager = mock_message_cycle_manager + pipeline._task_state = task_state + return pipeline - message_end_event = QueueMessageEndEvent(llm_result=None) + def test_get_message_event_type_called_once_when_first_llm_chunk_arrives( + self, pipeline, mock_message_cycle_manager + ): + """Expect get_message_event_type to be called when processing the first LLM chunk event.""" + # Setup a minimal LLM chunk event + chunk = llm_chunk("hi") + llm_chunk_event = QueueLLMChunkEvent(chunk=chunk) + message = queue_message(llm_chunk_event) + pipeline.queue_manager.messages = [message] - message = queue_message(message_end_event) - pipeline.queue_manager.messages = [message] + # Execute + list(pipeline._process_stream_response(publisher=None, trace_manager=None)) - pipeline._save_message = Mock() - pipeline._message_end_to_stream_response = Mock( - return_value=MessageEndStreamResponse(task_id="test-task-id", id="test-message-id") - ) + # Assert + mock_message_cycle_manager.get_message_event_type.assert_called_once_with(message_id="test-message-id") - list(pipeline._process_stream_response(publisher=None, trace_manager=trace_manager)) + def test_llm_chunk_event_with_text_content(self, pipeline, mock_message_cycle_manager, task_state): + """Test handling of LLM chunk events with text content.""" + # Setup + chunk = llm_chunk("Hello, world!") - # Assert - session = pipeline._save_message.call_args.kwargs["session"] - assert isinstance(session, Session) - assert committed_sessions == [session] - pipeline._save_message.assert_called_once_with(session=session, trace_manager=trace_manager) + llm_chunk_event = QueueLLMChunkEvent(chunk=chunk) - def test_multiple_events_sequence(self, pipeline, mock_message_cycle_manager, task_state): - """Test handling multiple events in sequence.""" - # Setup - chunk1 = llm_chunk("Hello") + message = queue_message(llm_chunk_event) + pipeline.queue_manager.messages = [message] - chunk2 = llm_chunk(" world!") + mock_message_cycle_manager.get_message_event_type.return_value = StreamEvent.MESSAGE - llm_chunk_event1 = QueueLLMChunkEvent(chunk=chunk1) + # Execute + responses = list(pipeline._process_stream_response(publisher=None, trace_manager=None)) - ping_event = QueuePingEvent() + # Assert + assert len(responses) == 1 + mock_message_cycle_manager.message_to_stream_response.assert_called_once_with( + answer="Hello, world!", message_id="test-message-id", event_type=StreamEvent.MESSAGE + ) + assert task_state.llm_result.message.content == "Hello, world!" - llm_chunk_event2 = QueueLLMChunkEvent(chunk=chunk2) + def test_llm_chunk_event_with_list_content(self, pipeline, mock_message_cycle_manager, task_state): + """Test handling of LLM chunk events with list content.""" + # Setup + text_content = TextPromptMessageContent(data="Hello") + chunk = llm_chunk("") + # Preserve the legacy mixed-list payload handled by _chunk_delta_text. + chunk.delta.message = AssistantPromptMessage.model_construct(content=[text_content, " world!"]) - mock_queue_messages = [ - queue_message(llm_chunk_event1), - queue_message(ping_event), - queue_message(llm_chunk_event2), - ] - pipeline.queue_manager.messages = mock_queue_messages + llm_chunk_event = QueueLLMChunkEvent(chunk=chunk) - mock_message_cycle_manager.get_message_event_type.return_value = StreamEvent.MESSAGE - pipeline.ping_stream_response = Mock(return_value=PingStreamResponse(task_id="test-task-id")) + message = queue_message(llm_chunk_event) + pipeline.queue_manager.messages = [message] - # Execute - responses = list(pipeline._process_stream_response(publisher=None, trace_manager=None)) + mock_message_cycle_manager.get_message_event_type.return_value = StreamEvent.MESSAGE - # Assert - assert len(responses) == 3 - assert task_state.llm_result.message.content == "Hello world!" + # Execute + responses = list(pipeline._process_stream_response(publisher=None, trace_manager=None)) - # Verify calls to message_to_stream_response - assert mock_message_cycle_manager.message_to_stream_response.call_count == 2 - mock_message_cycle_manager.message_to_stream_response.assert_any_call( - answer="Hello", message_id="test-message-id", event_type=StreamEvent.MESSAGE - ) - mock_message_cycle_manager.message_to_stream_response.assert_any_call( - answer=" world!", message_id="test-message-id", event_type=StreamEvent.MESSAGE - ) + # Assert + assert len(responses) == 1 + mock_message_cycle_manager.message_to_stream_response.assert_called_once_with( + answer="Hello world!", message_id="test-message-id", event_type=StreamEvent.MESSAGE + ) + assert task_state.llm_result.message.content == "Hello world!" + + def test_agent_message_event(self, pipeline, mock_message_cycle_manager, task_state): + """Test handling of agent message events.""" + # Setup + chunk = llm_chunk("Agent response") + + agent_message_event = QueueAgentMessageEvent(chunk=chunk) + + message = queue_message(agent_message_event) + pipeline.queue_manager.messages = [message] + + # Ensure method under assertion is a mock to track calls + pipeline._agent_message_to_stream_response = Mock(return_value=Mock()) + + # Execute + responses = list(pipeline._process_stream_response(publisher=None, trace_manager=None)) + + # Assert + assert len(responses) == 1 + # Agent messages should use _agent_message_to_stream_response + pipeline._agent_message_to_stream_response.assert_called_once_with( + answer="Agent response", message_id="test-message-id" + ) + + def test_message_end_event(self, pipeline, mock_message_cycle_manager, task_state, committed_sessions): + """Test handling of message end events.""" + # Setup + result = llm_result("Final response") + message_end_event = QueueMessageEndEvent(llm_result=result) + + message = queue_message(message_end_event) + pipeline.queue_manager.messages = [message] + + pipeline._save_message = Mock() + pipeline._message_end_to_stream_response = Mock( + return_value=MessageEndStreamResponse(task_id="test-task-id", id="test-message-id") + ) + + responses = list(pipeline._process_stream_response(publisher=None, trace_manager=None)) + + # Assert + assert len(responses) == 1 + assert task_state.llm_result == result + session = pipeline._save_message.call_args.kwargs["session"] + assert isinstance(session, Session) + assert committed_sessions == [session] + pipeline._message_end_to_stream_response.assert_called_once() + + def test_error_event(self, pipeline, committed_sessions): + """Test handling of error events.""" + # Setup + error_event = QueueErrorEvent(error=Exception("Test error")) + + message = queue_message(error_event) + pipeline.queue_manager.messages = [message] + + pipeline.handle_error = Mock(return_value=Exception("Test error")) + pipeline.error_to_stream_response = Mock( + return_value=ErrorStreamResponse(task_id="test-task-id", err=Exception("Test error")) + ) + + responses = list(pipeline._process_stream_response(publisher=None, trace_manager=None)) + + # Assert + assert len(responses) == 1 + session = pipeline.handle_error.call_args.kwargs["session"] + assert isinstance(session, Session) + assert committed_sessions == [session] + pipeline.handle_error.assert_called_once_with( + event=error_event, session=session, message_id="test-message-id" + ) + pipeline.error_to_stream_response.assert_called_once() + + def test_ping_event(self, pipeline): + """Test handling of ping events.""" + # Setup + ping_event = QueuePingEvent() + + message = queue_message(ping_event) + pipeline.queue_manager.messages = [message] + + pipeline.ping_stream_response = Mock(return_value=PingStreamResponse(task_id="test-task-id")) + + # Execute + responses = list(pipeline._process_stream_response(publisher=None, trace_manager=None)) + + # Assert + assert len(responses) == 1 + pipeline.ping_stream_response.assert_called_once() + + def test_file_event(self, pipeline, mock_message_cycle_manager): + """Test handling of file events.""" + # Setup + file_event = QueueMessageFileEvent(message_file_id="file-id") + + message = queue_message(file_event) + pipeline.queue_manager.messages = [message] + + file_response = MessageFileStreamResponse( + task_id="test-task-id", + id="file-id", + type="image", + belongs_to="assistant", + url="https://example.com/file.png", + ) + mock_message_cycle_manager.message_file_to_stream_response.return_value = file_response + + # Execute + responses = list(pipeline._process_stream_response(publisher=None, trace_manager=None)) + + # Assert + assert len(responses) == 1 + assert responses[0] == file_response + mock_message_cycle_manager.message_file_to_stream_response.assert_called_once_with(file_event) + + def test_publisher_is_called_with_messages(self, pipeline): + """Test that publisher publishes messages when provided.""" + # Setup + publisher = _Publisher() + + ping_event = QueuePingEvent() + message = queue_message(ping_event) + pipeline.queue_manager.messages = [message] + + pipeline.ping_stream_response = Mock(return_value=PingStreamResponse(task_id="test-task-id")) + + # Execute + list(pipeline._process_stream_response(publisher=publisher, trace_manager=None)) + + # Assert + assert publisher.messages == [message] + + def test_trace_manager_passed_to_save_message(self, pipeline, committed_sessions): + """Test that trace manager is passed to _save_message.""" + # Setup + with ( + patch("core.ops.ops_trace_manager.OpsTraceManager.get_ops_trace_instance", return_value=None), + patch.object(TraceQueueManager, "start_timer"), + ): + trace_manager = TraceQueueManager(app_id="test-app-id") + + message_end_event = QueueMessageEndEvent(llm_result=None) + + message = queue_message(message_end_event) + pipeline.queue_manager.messages = [message] + + pipeline._save_message = Mock() + pipeline._message_end_to_stream_response = Mock( + return_value=MessageEndStreamResponse(task_id="test-task-id", id="test-message-id") + ) + + list(pipeline._process_stream_response(publisher=None, trace_manager=trace_manager)) + + # Assert + session = pipeline._save_message.call_args.kwargs["session"] + assert isinstance(session, Session) + assert committed_sessions == [session] + pipeline._save_message.assert_called_once_with(session=session, trace_manager=trace_manager) + + def test_multiple_events_sequence(self, pipeline, mock_message_cycle_manager, task_state): + """Test handling multiple events in sequence.""" + # Setup + chunk1 = llm_chunk("Hello") + + chunk2 = llm_chunk(" world!") + + llm_chunk_event1 = QueueLLMChunkEvent(chunk=chunk1) + + ping_event = QueuePingEvent() + + llm_chunk_event2 = QueueLLMChunkEvent(chunk=chunk2) + + mock_queue_messages = [ + queue_message(llm_chunk_event1), + queue_message(ping_event), + queue_message(llm_chunk_event2), + ] + pipeline.queue_manager.messages = mock_queue_messages + + mock_message_cycle_manager.get_message_event_type.return_value = StreamEvent.MESSAGE + pipeline.ping_stream_response = Mock(return_value=PingStreamResponse(task_id="test-task-id")) + + # Execute + responses = list(pipeline._process_stream_response(publisher=None, trace_manager=None)) + + # Assert + assert len(responses) == 3 + assert task_state.llm_result.message.content == "Hello world!" + + # Verify calls to message_to_stream_response + assert mock_message_cycle_manager.message_to_stream_response.call_count == 2 + mock_message_cycle_manager.message_to_stream_response.assert_any_call( + answer="Hello", message_id="test-message-id", event_type=StreamEvent.MESSAGE + ) + mock_message_cycle_manager.message_to_stream_response.assert_any_call( + answer=" world!", message_id="test-message-id", event_type=StreamEvent.MESSAGE + ) + + +class _DummyModelConf: + def __init__(self) -> None: + self.model = "mock" + + +class _UnknownQueueEvent: + pass + + +@dataclass(eq=False, repr=False) +class _AnnotationReply: + content: str + + +@dataclass(eq=False, repr=False) +class _ModelConfigMode: + mode: str + + +@dataclass(eq=False, repr=False) +class _ProviderModelBundle: + model_type_instance: object + + +class _AudioPublisher: + def __init__(self, status: str, audio: str) -> None: + self._audio = AudioTrunk(status, audio) + + def check_and_get_audio(self) -> AudioTrunk: + return self._audio + + +class _TraceManagerDouble: + def __init__(self) -> None: + self.add_trace_task = Mock() + + +class _FakeQueueManager(AppQueueManager): + def __init__(self) -> None: + self._events: list[MessageQueueMessage | WorkflowQueueMessage] = [] + self.published_events: list[AppQueueEvent] = [] + + def set_events(self, events: list[MessageQueueMessage | WorkflowQueueMessage]) -> None: + self._events = events + + def listen(self) -> Generator[MessageQueueMessage | WorkflowQueueMessage]: + yield from self._events + + def publish(self, event: AppQueueEvent, pub_from: PublishFrom) -> None: + self.published_events.append(event) + + def _publish(self, event: AppQueueEvent, pub_from: PublishFrom) -> None: + self.published_events.append(event) + + +def _queue_message(event: AppQueueEvent) -> MessageQueueMessage: + return MessageQueueMessage( + task_id="task", + app_mode=AppMode.CHAT.value, + message_id="msg", + conversation_id="conv", + event=event, + ) + + +def _unknown_queue_message() -> MessageQueueMessage: + return MessageQueueMessage.model_construct( + task_id="task", + app_mode=AppMode.CHAT.value, + message_id="msg", + conversation_id="conv", + event=cast(AppQueueEvent, _UnknownQueueEvent()), + ) + + +def _make_conversation(app_mode: AppMode) -> Conversation: + conversation = Conversation( + app_id="app", + app_model_config_id=None, + model_provider=None, + override_model_configs=None, + model_id=None, + mode=app_mode, + name="conversation", + _inputs={}, + introduction="", + system_instruction="", + system_instruction_tokens=0, + status="normal", + invoke_from=InvokeFrom.WEB_APP, + from_source=ConversationFromSource.API, + from_end_user_id="user", + from_account_id=None, + ) + conversation.id = "conv" + conversation.mode = app_mode + return conversation + + +def _make_message() -> Message: + message = Message( + app_id="app", + conversation_id="conv", + _inputs={}, + query="query", + message="", + message_tokens=0, + message_unit_price=0, + message_price_unit=0, + answer="", + answer_tokens=0, + answer_unit_price=0, + answer_price_unit=0, + provider_response_latency=0, + total_price=0, + currency="USD", + invoke_from=InvokeFrom.WEB_APP, + from_source=ConversationFromSource.API, + from_end_user_id="user", + from_account_id=None, + app_mode=AppMode.CHAT, + ) + message.id = "msg" + message.created_at = datetime.now(UTC) + return message + + +def _message_file( + *, + file_id: str, + transfer_method: FileTransferMethod, + url: str | None, + upload_file_id: str | None, + file_type: FileType = FileType.IMAGE, +) -> MessageFile: + message_file = MessageFile( + message_id="msg", + type=file_type, + transfer_method=transfer_method, + created_by_role=CreatorUserRole.ACCOUNT, + created_by="user", + url=url, + upload_file_id=upload_file_id, + ) + message_file.id = file_id + return message_file + + +def _upload_file(*, file_id: str, name: str, mime_type: str, size: int, extension: str) -> UploadFile: + return make_upload_file( + file_id=file_id, + tenant_id="tenant", + key=f"uploads/{file_id}", + name=name, + size=size, + extension=extension, + mime_type=mime_type, + created_by="user", + created_at=datetime.now(UTC), + ) + + +def _agent_thought() -> MessageAgentThought: + thought = MessageAgentThought( + message_id="msg", + position=1, + created_by_role=CreatorUserRole.ACCOUNT, + created_by="user", + thought="t", + observation="o", + tool="tool", + tool_labels_str="{}", + tool_input="input", + answer="answer", + message_files="[]", + ) + thought.id = "thought" + return thought + + +def _make_app_config(app_mode: AppMode) -> EasyUIBasedAppConfig: + return EasyUIBasedAppConfig( + tenant_id="tenant", + app_id="app", + app_mode=app_mode, + app_model_config_from=EasyUIBasedAppModelConfigFrom.APP_LATEST_CONFIG, + app_model_config_id="model-config", + app_model_config_dict={}, + model=ModelConfigEntity(provider="mock", model="mock"), + prompt_template=PromptTemplateEntity( + prompt_type=PromptTemplateEntity.PromptType.SIMPLE, + simple_prompt_template="hi", + ), + additional_features=AppAdditionalFeatures(), + variables=[], + ) + + +def _make_entity(entity_cls, app_mode: AppMode): + app_config = _make_app_config(app_mode) + return entity_cls.model_construct( + task_id="task", + app_config=app_config, + model_conf=_DummyModelConf(), + file_upload_config=None, + conversation_id=None, + inputs={}, + query="hello", + files=[], + parent_message_id=None, + user_id="user", + stream=False, + invoke_from=InvokeFrom.WEB_APP, + extras={}, + call_depth=0, + trace_manager=None, + ) + + +def _make_pipeline( + entity_cls: type[ChatAppGenerateEntity] | type[CompletionAppGenerateEntity] = ChatAppGenerateEntity, + app_mode: AppMode = AppMode.CHAT, + *, + stream: bool = True, +) -> tuple[EasyUIBasedGenerateTaskPipeline, _FakeQueueManager]: + queue_manager = _FakeQueueManager() + pipeline = EasyUIBasedGenerateTaskPipeline( + application_generate_entity=_make_entity(entity_cls, app_mode), + queue_manager=queue_manager, + conversation=_make_conversation(app_mode), + message=_make_message(), + stream=stream, + ) + return pipeline, queue_manager + + +def _set_queue_events( + pipeline: EasyUIBasedGenerateTaskPipeline, + events: Sequence[MessageQueueMessage | WorkflowQueueMessage], +) -> None: + cast(_FakeQueueManager, pipeline.queue_manager).set_events(list(events)) + + +def _set_method(obj: object, name: str, value: object) -> None: + object.__setattr__(obj, name, value) + + +class TestEasyUiBasedGenerateTaskPipelineCoreCases: + class TestEasyUiBasedGenerateTaskPipeline: + def test_to_blocking_response_chat(self): + conversation = _make_conversation(AppMode.CHAT) + message = _make_message() + + pipeline = EasyUIBasedGenerateTaskPipeline( + application_generate_entity=_make_entity(ChatAppGenerateEntity, AppMode.CHAT), + queue_manager=_FakeQueueManager(), + conversation=conversation, + message=message, + stream=False, + ) + pipeline._task_state.llm_result.message.content = "answer" + + def _gen(): + yield MessageEndStreamResponse(task_id="task", id="msg") + + response = pipeline._to_blocking_response(_gen()) + + assert response.data.answer == "answer" + + def test_to_blocking_response_completion(self): + conversation = _make_conversation(AppMode.COMPLETION) + message = _make_message() + + pipeline = EasyUIBasedGenerateTaskPipeline( + application_generate_entity=_make_entity(CompletionAppGenerateEntity, AppMode.COMPLETION), + queue_manager=_FakeQueueManager(), + conversation=conversation, + message=message, + stream=False, + ) + pipeline._task_state.llm_result.message.content = "answer" + + def _gen(): + yield MessageEndStreamResponse(task_id="task", id="msg") + + response = pipeline._to_blocking_response(_gen()) + + assert response.data.answer == "answer" + + def test_listen_audio_msg_returns_none_when_no_publisher(self): + conversation = _make_conversation(AppMode.CHAT) + message = _make_message() + + pipeline = EasyUIBasedGenerateTaskPipeline( + application_generate_entity=_make_entity(ChatAppGenerateEntity, AppMode.CHAT), + queue_manager=_FakeQueueManager(), + conversation=conversation, + message=message, + stream=False, + ) + + assert pipeline._listen_audio_msg(publisher=None, task_id="task") is None + + def test_process_stream_response_handles_chunks_and_end(self, monkeypatch: pytest.MonkeyPatch): + conversation = _make_conversation(AppMode.CHAT) + message = _make_message() + + pipeline = EasyUIBasedGenerateTaskPipeline( + application_generate_entity=_make_entity(ChatAppGenerateEntity, AppMode.CHAT), + queue_manager=_FakeQueueManager(), + conversation=conversation, + message=message, + stream=True, + ) + + chunk = LLMResultChunk( + model="mock", + prompt_messages=[], + delta=LLMResultChunkDelta( + index=0, + message=AssistantPromptMessage( + content=[TextPromptMessageContent(data="hi"), TextPromptMessageContent(data="yo")] + ), + ), + ) + llm_result = LLMResult( + model="mock", + prompt_messages=[], + message=AssistantPromptMessage(content="done"), + usage=LLMUsage.empty_usage(), + ) + + events = [ + _queue_message(QueueLLMChunkEvent(chunk=chunk)), + _queue_message(QueueMessageReplaceEvent(text="replace", reason="output_moderation")), + _queue_message(QueuePingEvent()), + _queue_message(QueueMessageEndEvent(llm_result=llm_result)), + ] + + _set_queue_events(pipeline, events) + + def _message_event_type(message_id: str) -> StreamEvent: + return StreamEvent.MESSAGE + + def _message_end() -> MessageEndStreamResponse: + return MessageEndStreamResponse(task_id="task", id="msg") + + _set_method(pipeline._message_cycle_manager, "get_message_event_type", _message_event_type) + _set_method(pipeline, "handle_output_moderation_when_task_finished", lambda completion: None) + _set_method(pipeline, "_message_end_to_stream_response", _message_end) + _set_method(pipeline, "_save_message", lambda **kwargs: None) + + class _Session: + def __init__[**P](self, *args: P.args, **kwargs: P.kwargs): + pass + + def __enter__(self): + return self + + def __exit__(self, exc_type, exc, tb): + return False + + def commit(self): + return None + + monkeypatch.setattr( + "core.app.task_pipeline.easy_ui_based_generate_task_pipeline.session_factory.create_session", + lambda: _Session(), + ) + + responses = list(pipeline._process_stream_response(publisher=None)) + + message_response = next(item for item in responses if isinstance(item, MessageStreamResponse)) + assert message_response.answer == "hiyo" + replace_response = next(item for item in responses if isinstance(item, MessageReplaceStreamResponse)) + assert replace_response.answer == "replace" + assert any(isinstance(item, PingStreamResponse) for item in responses) + assert isinstance(responses[-1], MessageEndStreamResponse) + assert pipeline._task_state.llm_result.message.content == "done" + + def test_handle_output_moderation_chunk_directs_output(self): + conversation = _make_conversation(AppMode.CHAT) + message = _make_message() + + pipeline = EasyUIBasedGenerateTaskPipeline( + application_generate_entity=_make_entity(ChatAppGenerateEntity, AppMode.CHAT), + queue_manager=_FakeQueueManager(), + conversation=conversation, + message=message, + stream=True, + ) + + class _Moderation: + def should_direct_output(self): + return True + + def get_final_output(self): + return "final" + + _set_method(pipeline, "output_moderation_handler", _Moderation()) + + result = pipeline._handle_output_moderation_chunk("token") + + assert result is True + events = cast(_FakeQueueManager, pipeline.queue_manager).published_events + assert any(isinstance(event, QueueLLMChunkEvent) for event in events) + assert any(isinstance(event, QueueStopEvent) for event in events) + + def test_handle_stop_updates_usage(self, monkeypatch: pytest.MonkeyPatch): + conversation = _make_conversation(AppMode.CHAT) + message = _make_message() + + class _ModelType: + def calc_response_usage(self, model, credentials, prompt_tokens, completion_tokens): + return LLMUsage.from_metadata( + { + "prompt_tokens": prompt_tokens, + "completion_tokens": completion_tokens, + } + ) + + class _ModelConf: + def __init__(self) -> None: + self.model = "mock" + self.credentials = {} + self.provider_model_bundle = _ProviderModelBundle(model_type_instance=_ModelType()) + + app_config = _make_app_config(AppMode.CHAT) + application_generate_entity = ChatAppGenerateEntity.model_construct( + task_id="task", + app_config=app_config, + model_conf=_ModelConf(), + file_upload_config=None, + conversation_id=None, + inputs={}, + query="hello", + files=[], + parent_message_id=None, + user_id="user", + stream=False, + invoke_from=InvokeFrom.WEB_APP, + extras={}, + call_depth=0, + trace_manager=None, + ) + + pipeline = EasyUIBasedGenerateTaskPipeline( + application_generate_entity=application_generate_entity, + queue_manager=_FakeQueueManager(), + conversation=conversation, + message=message, + stream=False, + ) + pipeline._task_state.llm_result.prompt_messages = [AssistantPromptMessage(content="prompt")] + pipeline._task_state.llm_result.message = AssistantPromptMessage(content="answer") + + calls: list[int] = [] + + class _FakeModelInstance: + def __init__(self, provider_model_bundle, model): + pass + + def get_llm_num_tokens(self, messages): + calls.append(1) + return 10 if len(calls) == 1 else 5 + + monkeypatch.setattr( + "core.app.task_pipeline.easy_ui_based_generate_task_pipeline.ModelInstance", + _FakeModelInstance, + ) + + pipeline._handle_stop(QueueStopEvent(stopped_by=QueueStopEvent.StopBy.USER_MANUAL)) + + assert pipeline._task_state.llm_result.usage.prompt_tokens == 10 + assert pipeline._task_state.llm_result.usage.completion_tokens == 5 + + def test_record_files_builds_file_payloads(self, monkeypatch: pytest.MonkeyPatch): + conversation = _make_conversation(AppMode.CHAT) + message = _make_message() + + pipeline = EasyUIBasedGenerateTaskPipeline( + application_generate_entity=_make_entity(ChatAppGenerateEntity, AppMode.CHAT), + queue_manager=_FakeQueueManager(), + conversation=conversation, + message=message, + stream=False, + ) + + message_files = [ + _message_file( + file_id="mf-1", + transfer_method=FileTransferMethod.REMOTE_URL, + url="http://example.com/a.png", + upload_file_id=None, + ), + _message_file( + file_id="mf-2", + transfer_method=FileTransferMethod.LOCAL_FILE, + url="", + upload_file_id="upload-1", + ), + _message_file( + file_id="mf-3", + transfer_method=FileTransferMethod.TOOL_FILE, + url="tool/file.bin", + upload_file_id=None, + file_type=FileType.CUSTOM, + ), + ] + upload_files = [ + _upload_file( + file_id="upload-1", + name="local.png", + mime_type="image/png", + size=123, + extension="png", + ) + ] + + class _Result: + def __init__(self, items): + self._items = items + + def all(self): + return self._items + + class _Session: + def __init__[**P](self, *args: P.args, **kwargs: P.kwargs): + self.calls = 0 + + def __enter__(self): + return self + + def __exit__(self, exc_type, exc, tb): + return False + + def scalars(self, *args, **kwargs): + self.calls += 1 + return _Result(message_files if self.calls == 1 else upload_files) + + monkeypatch.setattr( + "core.app.task_pipeline.easy_ui_based_generate_task_pipeline.session_factory.create_session", + lambda: _Session(), + ) + monkeypatch.setattr( + "core.app.task_pipeline.message_file_utils.file_helpers.get_signed_file_url", + lambda **kwargs: "signed-url", + ) + monkeypatch.setattr( + "core.app.task_pipeline.message_file_utils.sign_tool_file", + lambda **kwargs: "signed-tool", + ) + + response = pipeline._message_end_to_stream_response() + files = response.files + + assert files + assert len(files) == 3 + + def test_process_stream_response_handles_annotation_and_error(self, monkeypatch: pytest.MonkeyPatch): + conversation = _make_conversation(AppMode.CHAT) + message = _make_message() + + pipeline = EasyUIBasedGenerateTaskPipeline( + application_generate_entity=_make_entity(ChatAppGenerateEntity, AppMode.CHAT), + queue_manager=_FakeQueueManager(), + conversation=conversation, + message=message, + stream=True, + ) + + agent_chunk = LLMResultChunk( + model="mock", + prompt_messages=[], + delta=LLMResultChunkDelta( + index=0, + message=AssistantPromptMessage(content="agent"), + ), + ) + + events = [ + _queue_message(QueueAnnotationReplyEvent(message_annotation_id="ann")), + _queue_message(QueueAgentThoughtEvent(agent_thought_id="thought")), + _queue_message(QueueMessageFileEvent(message_file_id="file")), + _queue_message(QueueAgentMessageEvent(chunk=agent_chunk)), + _queue_message(QueueErrorEvent(error=ValueError("boom"))), + ] + + _set_queue_events(pipeline, events) + + def _agent_thought_response(event: QueueAgentThoughtEvent) -> AgentThoughtStreamResponse: + return AgentThoughtStreamResponse(task_id="task", id=event.agent_thought_id, position=1) + + def _file_response(event: QueueMessageFileEvent) -> MessageFileStreamResponse: + return MessageFileStreamResponse( + task_id="task", id=event.message_file_id, type="image", belongs_to="user", url="file" + ) + + def _agent_message_response(answer: str, message_id: str) -> AgentMessageStreamResponse: + return AgentMessageStreamResponse(task_id="task", id=message_id, answer=answer) + + _set_method( + pipeline._message_cycle_manager, + "handle_annotation_reply", + lambda event, session: _AnnotationReply(content="annotated"), + ) + _set_method(pipeline, "_agent_thought_to_stream_response", _agent_thought_response) + _set_method(pipeline._message_cycle_manager, "message_file_to_stream_response", _file_response) + _set_method(pipeline, "_agent_message_to_stream_response", _agent_message_response) + _set_method(pipeline, "handle_error", lambda **kwargs: ValueError("boom")) + _set_method(pipeline, "error_to_stream_response", lambda err: ErrorStreamResponse(task_id="task", err=err)) + + class _Session: + def __init__[**P](self, *args: P.args, **kwargs: P.kwargs): + pass + + def __enter__(self): + return self + + def __exit__(self, exc_type, exc, tb): + return False + + def commit(self): + return None + + monkeypatch.setattr( + "core.app.task_pipeline.easy_ui_based_generate_task_pipeline.session_factory.create_session", + lambda: _Session(), + ) + + responses = list(pipeline._process_stream_response(publisher=None)) + + assert any(isinstance(response, AgentThoughtStreamResponse) for response in responses) + assert any(isinstance(response, MessageFileStreamResponse) for response in responses) + agent_response = next( + response for response in responses if isinstance(response, AgentMessageStreamResponse) + ) + assert agent_response.answer == "agent" + assert isinstance(responses[-1], ErrorStreamResponse) + assert isinstance(responses[-1].err, ValueError) + assert pipeline._task_state.llm_result.message.content == "annotated" + + def test_process_stream_response_error_event_adds_trace_task(self, monkeypatch: pytest.MonkeyPatch): + conversation = _make_conversation(AppMode.CHAT) + message = _make_message() + application_generate_entity = _make_entity(ChatAppGenerateEntity, AppMode.CHAT) + application_generate_entity.extras = {"trace_session_id": "session-1"} + + pipeline = EasyUIBasedGenerateTaskPipeline( + application_generate_entity=application_generate_entity, + queue_manager=_FakeQueueManager(), + conversation=conversation, + message=message, + stream=True, + ) + _set_queue_events(pipeline, [_queue_message(QueueErrorEvent(error=ValueError("boom")))]) + _set_method(pipeline, "handle_error", lambda **kwargs: ValueError("boom")) + _set_method(pipeline, "error_to_stream_response", lambda err: ErrorStreamResponse(task_id="task", err=err)) + + trace_manager_double = _TraceManagerDouble() + trace_manager = cast(TraceQueueManager, trace_manager_double) + + class _Session: + def __enter__(self): + return self + + def __exit__(self, exc_type, exc, tb): + return False + + def commit(self): + return None + + monkeypatch.setattr( + "core.app.task_pipeline.easy_ui_based_generate_task_pipeline.session_factory.create_session", + lambda: _Session(), + ) + + responses = list(pipeline._process_stream_response(publisher=None, trace_manager=trace_manager)) + + assert len(responses) == 1 + assert isinstance(responses[0], ErrorStreamResponse) + trace_manager_double.add_trace_task.assert_called_once() + trace_task = trace_manager_double.add_trace_task.call_args.args[0] + assert trace_task.trace_type == TraceTaskName.MESSAGE_TRACE + assert trace_task.conversation_id == "conv" + assert trace_task.message_id == "msg" + assert trace_task.kwargs["trace_session_id"] == "session-1" + + def test_agent_thought_to_stream_response_returns_payload(self, monkeypatch: pytest.MonkeyPatch): + conversation = _make_conversation(AppMode.CHAT) + message = _make_message() + + pipeline = EasyUIBasedGenerateTaskPipeline( + application_generate_entity=_make_entity(ChatAppGenerateEntity, AppMode.CHAT), + queue_manager=_FakeQueueManager(), + conversation=conversation, + message=message, + stream=True, + ) + + agent_thought = _agent_thought() + + class _Session: + def __init__[**P](self, *args: P.args, **kwargs: P.kwargs): + pass + + def __enter__(self): + return self + + def __exit__(self, exc_type, exc, tb): + return False + + def scalar(self, *args, **kwargs): + return agent_thought + + monkeypatch.setattr( + "core.app.task_pipeline.easy_ui_based_generate_task_pipeline.session_factory.create_session", + lambda: _Session(), + ) + + response = pipeline._agent_thought_to_stream_response(QueueAgentThoughtEvent(agent_thought_id="thought")) + + assert response is not None + assert response.id == "thought" + assert response.thought == "t" + + def test_agent_thought_to_stream_response_normalizes_null_display_fields(self, monkeypatch: pytest.MonkeyPatch): + conversation = _make_conversation(AppMode.CHAT) + message = _make_message() + + pipeline = EasyUIBasedGenerateTaskPipeline( + application_generate_entity=_make_entity(ChatAppGenerateEntity, AppMode.CHAT), + queue_manager=_FakeQueueManager(), + conversation=conversation, + message=message, + stream=True, + ) + + agent_thought = _agent_thought() + agent_thought.thought = None + agent_thought.answer = None + agent_thought.observation = None + agent_thought.tool = None + agent_thought.tool_input = None + agent_thought.message_files = None + + class _Session: + def __init__(self, *args, **kwargs): + pass + + def __enter__(self): + return self + + def __exit__(self, exc_type, exc, tb): + return False + + def scalar(self, *args, **kwargs): + return agent_thought + + monkeypatch.setattr( + "core.app.task_pipeline.easy_ui_based_generate_task_pipeline.session_factory.create_session", + lambda: _Session(), + ) + response = pipeline._agent_thought_to_stream_response(QueueAgentThoughtEvent(agent_thought_id="thought")) + + assert response is not None + assert response.thought == "" + assert response.observation == "" + assert response.tool == "" + assert response.tool_input == "" + assert response.model_dump(mode="json")["message_files"] == [] + + def test_process_routes_to_stream_and_starts_conversation_name_generation(self): + conversation = _make_conversation(AppMode.CHAT) + message = _make_message() + pipeline = EasyUIBasedGenerateTaskPipeline( + application_generate_entity=_make_entity(ChatAppGenerateEntity, AppMode.CHAT), + queue_manager=_FakeQueueManager(), + conversation=conversation, + message=message, + stream=True, + ) + pipeline._message_cycle_manager.generate_conversation_name = Mock(return_value=object()) + _set_method( + pipeline, + "_wrapper_process_stream_response", + lambda trace_manager: iter([PingStreamResponse(task_id="task")]), + ) + _set_method(pipeline, "_to_stream_response", lambda generator: "streamed") + + result = pipeline.process() + + assert result == "streamed" + pipeline._message_cycle_manager.generate_conversation_name.assert_called_once_with( + conversation_id="conv", query="hello", message_id="msg" + ) + + def test_process_routes_to_blocking_for_completion_mode(self): + conversation = _make_conversation(AppMode.COMPLETION) + message = _make_message() + pipeline = EasyUIBasedGenerateTaskPipeline( + application_generate_entity=_make_entity(CompletionAppGenerateEntity, AppMode.COMPLETION), + queue_manager=_FakeQueueManager(), + conversation=conversation, + message=message, + stream=False, + ) + pipeline._message_cycle_manager.generate_conversation_name = Mock() + _set_method( + pipeline, + "_wrapper_process_stream_response", + lambda trace_manager: iter([PingStreamResponse(task_id="task")]), + ) + _set_method(pipeline, "_to_blocking_response", lambda generator: "blocking") + + result = pipeline.process() + + assert result == "blocking" + pipeline._message_cycle_manager.generate_conversation_name.assert_not_called() + + def test_to_blocking_response_raises_error_stream_exception(self): + conversation = _make_conversation(AppMode.CHAT) + message = _make_message() + pipeline = EasyUIBasedGenerateTaskPipeline( + application_generate_entity=_make_entity(ChatAppGenerateEntity, AppMode.CHAT), + queue_manager=_FakeQueueManager(), + conversation=conversation, + message=message, + stream=False, + ) + + def _gen(): + yield ErrorStreamResponse(task_id="task", err=ValueError("stream error")) + + with pytest.raises(ValueError, match="stream error"): + pipeline._to_blocking_response(_gen()) + + def test_to_blocking_response_raises_when_generator_ends_without_message_end(self): + conversation = _make_conversation(AppMode.CHAT) + message = _make_message() + pipeline = EasyUIBasedGenerateTaskPipeline( + application_generate_entity=_make_entity(ChatAppGenerateEntity, AppMode.CHAT), + queue_manager=_FakeQueueManager(), + conversation=conversation, + message=message, + stream=False, + ) + + def _gen(): + yield PingStreamResponse(task_id="task") + + with pytest.raises(RuntimeError, match="queue listening stopped unexpectedly"): + pipeline._to_blocking_response(_gen()) + + def test_to_stream_response_wraps_completion_stream_events(self): + conversation = _make_conversation(AppMode.COMPLETION) + message = _make_message() + pipeline = EasyUIBasedGenerateTaskPipeline( + application_generate_entity=_make_entity(CompletionAppGenerateEntity, AppMode.COMPLETION), + queue_manager=_FakeQueueManager(), + conversation=conversation, + message=message, + stream=True, + ) + + def _gen(): + yield PingStreamResponse(task_id="task") + + response = list(pipeline._to_stream_response(_gen()))[0] + + assert isinstance(response, CompletionAppStreamResponse) + assert response.message_id == "msg" + + def test_to_stream_response_wraps_chat_stream_events(self): + conversation = _make_conversation(AppMode.CHAT) + message = _make_message() + pipeline = EasyUIBasedGenerateTaskPipeline( + application_generate_entity=_make_entity(ChatAppGenerateEntity, AppMode.CHAT), + queue_manager=_FakeQueueManager(), + conversation=conversation, + message=message, + stream=True, + ) + + def _gen(): + yield PingStreamResponse(task_id="task") + + response = list(pipeline._to_stream_response(_gen()))[0] + + assert isinstance(response, ChatbotAppStreamResponse) + assert response.conversation_id == "conv" + + def test_listen_audio_msg_returns_audio_response_for_non_finish_audio(self): + conversation = _make_conversation(AppMode.CHAT) + message = _make_message() + pipeline = EasyUIBasedGenerateTaskPipeline( + application_generate_entity=_make_entity(ChatAppGenerateEntity, AppMode.CHAT), + queue_manager=_FakeQueueManager(), + conversation=conversation, + message=message, + stream=True, + ) + publisher = cast( + AppGeneratorTTSPublisher, + _AudioPublisher(status="responding", audio="abc"), + ) + + response = pipeline._listen_audio_msg(publisher=publisher, task_id="task") + + assert isinstance(response, MessageAudioStreamResponse) + assert response.audio == "abc" + + def test_listen_audio_msg_returns_none_for_finish_audio(self): + conversation = _make_conversation(AppMode.CHAT) + message = _make_message() + pipeline = EasyUIBasedGenerateTaskPipeline( + application_generate_entity=_make_entity(ChatAppGenerateEntity, AppMode.CHAT), + queue_manager=_FakeQueueManager(), + conversation=conversation, + message=message, + stream=True, + ) + publisher = cast( + AppGeneratorTTSPublisher, + _AudioPublisher(status="finish", audio="abc"), + ) + + assert pipeline._listen_audio_msg(publisher=publisher, task_id="task") is None + + def test_wrapper_process_stream_response_without_tts_publisher(self): + conversation = _make_conversation(AppMode.CHAT) + message = _make_message() + pipeline = EasyUIBasedGenerateTaskPipeline( + application_generate_entity=_make_entity(ChatAppGenerateEntity, AppMode.CHAT), + queue_manager=_FakeQueueManager(), + conversation=conversation, + message=message, + stream=True, + ) + payload = PingStreamResponse(task_id="task") + _set_method(pipeline, "_process_stream_response", lambda publisher, trace_manager: iter([payload])) + + responses = list(pipeline._wrapper_process_stream_response()) + + assert responses == [payload] + + def test_wrapper_process_stream_response_with_tts_publisher(self, monkeypatch: pytest.MonkeyPatch): + conversation = _make_conversation(AppMode.CHAT) + message = _make_message() + entity = _make_entity(ChatAppGenerateEntity, AppMode.CHAT) + entity.app_config.app_model_config_dict = { + "text_to_speech": {"autoPlay": "enabled", "enabled": True, "voice": "v", "language": "en"} + } + pipeline = EasyUIBasedGenerateTaskPipeline( + application_generate_entity=entity, + queue_manager=_FakeQueueManager(), + conversation=conversation, + message=message, + stream=True, + ) + + class _Publisher: + def check_and_get_audio(self, *, block=False): + assert block + return AudioTrunk("finish", "") + + def publish(self, message): + assert message is None + + def cancel(self): + return None + + inline_audio = MessageAudioStreamResponse(task_id="task", audio="inline") + audio_calls = iter([inline_audio, None]) + payload = PingStreamResponse(task_id="task") + _set_method(pipeline, "_listen_audio_msg", lambda publisher, task_id: next(audio_calls)) + _set_method(pipeline, "_process_stream_response", lambda publisher, trace_manager: iter([payload])) + monkeypatch.setattr( + "core.app.task_pipeline.easy_ui_based_generate_task_pipeline.AppGeneratorTTSPublisher", + lambda tenant_id, voice, language, app_type: _Publisher(), + ) + + responses = list(pipeline._wrapper_process_stream_response()) + + assert responses[0] == inline_audio + assert responses[1] == payload + assert isinstance(responses[-1], MessageAudioEndStreamResponse) + assert "audio_type" not in responses[-1].model_dump() + + def test_wrapper_process_stream_response_waits_for_the_audio_terminal(self, monkeypatch: pytest.MonkeyPatch): + conversation = _make_conversation(AppMode.CHAT) + message = _make_message() + entity = _make_entity(ChatAppGenerateEntity, AppMode.CHAT) + entity.app_config.app_model_config_dict = { + "text_to_speech": {"autoPlay": "enabled", "enabled": True, "voice": "v", "language": "en"} + } + pipeline = EasyUIBasedGenerateTaskPipeline( + application_generate_entity=entity, + queue_manager=_FakeQueueManager(), + conversation=conversation, + message=message, + stream=True, + ) + + class _Publisher: + def __init__(self): + self._events = iter([AudioTrunk("responding", "later"), AudioTrunk("finish", "")]) + + def check_and_get_audio(self, *, block=False): + assert block + return next(self._events) + + def publish(self, message): + assert message is None + + def cancel(self): + return None + + _set_method(pipeline, "_process_stream_response", lambda publisher, trace_manager: iter([])) + monkeypatch.setattr( + "core.app.task_pipeline.easy_ui_based_generate_task_pipeline.AppGeneratorTTSPublisher", + lambda tenant_id, voice, language, app_type: _Publisher(), + ) + responses = list(pipeline._wrapper_process_stream_response()) + + assert any(isinstance(item, MessageAudioStreamResponse) for item in responses) + assert isinstance(responses[-1], MessageAudioEndStreamResponse) + + def test_wrapper_process_stream_response_ends_audio_before_reporting_tts_error( + self, monkeypatch: pytest.MonkeyPatch + ): + conversation = _make_conversation(AppMode.CHAT) + message = _make_message() + entity = _make_entity(ChatAppGenerateEntity, AppMode.CHAT) + entity.app_config.app_model_config_dict = { + "text_to_speech": {"autoPlay": "enabled", "enabled": True, "voice": "v", "language": "en"} + } + pipeline = EasyUIBasedGenerateTaskPipeline( + application_generate_entity=entity, + queue_manager=_FakeQueueManager(), + conversation=conversation, + message=message, + stream=True, + ) + error = RuntimeError("tts failed") + + class _Publisher: + def check_and_get_audio(self, *, block=False): + assert block + return AudioTrunk("error", b"", error=error) + + def publish(self, message): + assert message is None + + def cancel(self): + return None + + _set_method(pipeline, "_process_stream_response", lambda publisher, trace_manager: iter([])) + monkeypatch.setattr( + "core.app.task_pipeline.easy_ui_based_generate_task_pipeline.AppGeneratorTTSPublisher", + lambda tenant_id, voice, language, app_type: _Publisher(), + ) + + responses = list(pipeline._wrapper_process_stream_response()) + + assert isinstance(responses[-2], MessageAudioEndStreamResponse) + assert isinstance(responses[-1], ErrorStreamResponse) + assert responses[-1].err is error + + def test_wrapper_process_stream_response_ends_audio_before_a_main_error(self, monkeypatch: pytest.MonkeyPatch): + conversation = _make_conversation(AppMode.CHAT) + message = _make_message() + entity = _make_entity(ChatAppGenerateEntity, AppMode.CHAT) + entity.app_config.app_model_config_dict = { + "text_to_speech": {"autoPlay": "enabled", "enabled": True, "voice": "v", "language": "en"} + } + pipeline = EasyUIBasedGenerateTaskPipeline( + application_generate_entity=entity, + queue_manager=_FakeQueueManager(), + conversation=conversation, + message=message, + stream=True, + ) + publisher = Mock() + publisher.check_and_get_audio.return_value = None + error = ErrorStreamResponse(task_id="task", err=RuntimeError("generation failed")) + _set_method(pipeline, "_process_stream_response", lambda publisher, trace_manager: iter([error])) + monkeypatch.setattr( + "core.app.task_pipeline.easy_ui_based_generate_task_pipeline.AppGeneratorTTSPublisher", + lambda tenant_id, voice, language, app_type: publisher, + ) + + responses = list(pipeline._wrapper_process_stream_response()) + + assert isinstance(responses[-2], MessageAudioEndStreamResponse) + assert responses[-1] is error + assert publisher.cancel.called + publisher.publish.assert_not_called() + + def test_process_stream_response_handles_stop_event_and_output_replacement( + self, monkeypatch: pytest.MonkeyPatch + ): + conversation = _make_conversation(AppMode.CHAT) + message = _make_message() + pipeline = EasyUIBasedGenerateTaskPipeline( + application_generate_entity=_make_entity(ChatAppGenerateEntity, AppMode.CHAT), + queue_manager=_FakeQueueManager(), + conversation=conversation, + message=message, + stream=True, + ) + pipeline._task_state.llm_result.message.content = "raw answer" + _set_queue_events(pipeline, [_queue_message(QueueStopEvent(stopped_by=QueueStopEvent.StopBy.USER_MANUAL))]) + pipeline._handle_stop = Mock() + _set_method(pipeline, "handle_output_moderation_when_task_finished", lambda answer: "moderated answer") + _set_method( + pipeline._message_cycle_manager, + "message_replace_to_stream_response", + lambda answer: MessageReplaceStreamResponse(task_id="task", answer=answer, reason=""), + ) + _set_method(pipeline, "_save_message", lambda **kwargs: None) + _set_method( + pipeline, + "_message_end_to_stream_response", + lambda: MessageEndStreamResponse(task_id="task", id="msg"), + ) + + class _Session: + def __init__[**P](self, *args: P.args, **kwargs: P.kwargs): + pass + + def __enter__(self): + return self + + def __exit__(self, exc_type, exc, tb): + return False + + def commit(self): + return None + + monkeypatch.setattr( + "core.app.task_pipeline.easy_ui_based_generate_task_pipeline.session_factory.create_session", + lambda: _Session(), + ) + + responses = list(pipeline._process_stream_response(publisher=None)) + + assert isinstance(responses[0], MessageReplaceStreamResponse) + assert responses[0].answer == "moderated answer" + assert isinstance(responses[1], MessageEndStreamResponse) + pipeline._handle_stop.assert_called_once() + + def test_process_stream_response_handles_retriever_unknown_and_empty_chunk(self): + conversation = _make_conversation(AppMode.CHAT) + message = _make_message() + pipeline = EasyUIBasedGenerateTaskPipeline( + application_generate_entity=_make_entity(ChatAppGenerateEntity, AppMode.CHAT), + queue_manager=_FakeQueueManager(), + conversation=conversation, + message=message, + stream=True, + ) + retriever_event = QueueRetrieverResourcesEvent(retriever_resources=[]) + chunk = LLMResultChunk( + model="mock", + prompt_messages=[], + delta=LLMResultChunkDelta(index=0, message=AssistantPromptMessage(content=None)), + ) + handled = {"retriever": 0} + + def _handle_retriever_resources(event): + handled["retriever"] += 1 + + pipeline._message_cycle_manager.handle_retriever_resources = _handle_retriever_resources + _set_queue_events( + pipeline, + [ + _queue_message(retriever_event), + _unknown_queue_message(), + _queue_message(QueueLLMChunkEvent(chunk=chunk)), + ], + ) + + responses = list(pipeline._process_stream_response(publisher=None)) + + assert responses == [] + assert handled["retriever"] == 1 + + def test_process_stream_response_skips_when_output_moderation_directs_chunk(self): + conversation = _make_conversation(AppMode.CHAT) + message = _make_message() + pipeline = EasyUIBasedGenerateTaskPipeline( + application_generate_entity=_make_entity(ChatAppGenerateEntity, AppMode.CHAT), + queue_manager=_FakeQueueManager(), + conversation=conversation, + message=message, + stream=True, + ) + chunk = LLMResultChunk( + model="mock", + prompt_messages=[], + delta=LLMResultChunkDelta(index=0, message=AssistantPromptMessage(content="x")), + ) + pipeline._handle_output_moderation_chunk = lambda text: True + _set_queue_events(pipeline, [_queue_message(QueueLLMChunkEvent(chunk=chunk))]) + + responses = list(pipeline._process_stream_response(publisher=None)) + + assert responses == [] + + def test_process_stream_response_ignores_unsupported_chunk_content_types(self): + conversation = _make_conversation(AppMode.CHAT) + message = _make_message() + pipeline = EasyUIBasedGenerateTaskPipeline( + application_generate_entity=_make_entity(ChatAppGenerateEntity, AppMode.CHAT), + queue_manager=_FakeQueueManager(), + conversation=conversation, + message=message, + stream=True, + ) + chunk = LLMResultChunk.model_construct( + prompt_messages=[], + delta=LLMResultChunkDelta.model_construct( + message=AssistantPromptMessage.model_construct(content=[object(), "ok"]) + ), + ) + _set_method( + pipeline._message_cycle_manager, "get_message_event_type", lambda message_id: StreamEvent.MESSAGE + ) + _set_queue_events(pipeline, [_queue_message(QueueLLMChunkEvent.model_construct(chunk=chunk))]) + + responses = list(pipeline._process_stream_response(publisher=None)) + + assert len(responses) == 1 + assert isinstance(responses[0], MessageStreamResponse) + assert responses[0].answer == "ok" + assert pipeline._task_state.llm_result.message.content == "ok" + + def test_process_stream_response_skips_none_chunk_content(self): + conversation = _make_conversation(AppMode.CHAT) + message = _make_message() + pipeline = EasyUIBasedGenerateTaskPipeline( + application_generate_entity=_make_entity(ChatAppGenerateEntity, AppMode.CHAT), + queue_manager=_FakeQueueManager(), + conversation=conversation, + message=message, + stream=True, + ) + chunk = LLMResultChunk( + model="mock", + prompt_messages=[], + delta=LLMResultChunkDelta(index=0, message=AssistantPromptMessage(content=None)), + ) + pipeline._message_cycle_manager.message_to_stream_response = Mock() + _set_queue_events(pipeline, [_queue_message(QueueLLMChunkEvent(chunk=chunk))]) + + responses = list(pipeline._process_stream_response(publisher=None)) + + assert responses == [] + pipeline._message_cycle_manager.message_to_stream_response.assert_not_called() + assert pipeline._task_state.llm_result.message.content == "" + + def test_process_stream_response_reaches_post_loop_branch_with_thread_reference(self): + conversation = _make_conversation(AppMode.CHAT) + message = _make_message() + pipeline = EasyUIBasedGenerateTaskPipeline( + application_generate_entity=_make_entity(ChatAppGenerateEntity, AppMode.CHAT), + queue_manager=_FakeQueueManager(), + conversation=conversation, + message=message, + stream=True, + ) + pipeline._conversation_name_generate_thread = Thread() + _set_queue_events(pipeline, []) + + assert list(pipeline._process_stream_response(publisher=None)) == [] + + def test_save_message_persists_fields_and_emits_trace( + self, monkeypatch: pytest.MonkeyPatch, sqlite_session: Session + ): + conversation = _make_conversation(AppMode.CHAT) + message = _make_message() + application_generate_entity = _make_entity(ChatAppGenerateEntity, AppMode.CHAT) + application_generate_entity.extras = {"trace_session_id": "session-1"} + pipeline = EasyUIBasedGenerateTaskPipeline( + application_generate_entity=application_generate_entity, + queue_manager=_FakeQueueManager(), + conversation=conversation, + message=message, + stream=False, + ) + pipeline.start_at = 10.0 + _set_method(pipeline, "_model_config", _ModelConfigMode(mode="chat")) + pipeline._task_state.llm_result.prompt_messages = [AssistantPromptMessage(content="prompt")] + pipeline._task_state.llm_result.message = AssistantPromptMessage(content=" {{name}} hello ") + pipeline._task_state.llm_result.usage = LLMUsage.from_metadata( + {"prompt_tokens": 3, "completion_tokens": 5, "total_price": "1.23"} + ) + + message_obj = _make_message() + conversation_obj = _make_conversation(AppMode.CHAT) + session = sqlite_session + session.add_all([conversation_obj, message_obj]) + session.flush() + trace_manager_double = _TraceManagerDouble() + trace_manager = cast(TraceQueueManager, trace_manager_double) + sent_payloads: list[tuple[tuple[object, ...], dict[str, object]]] = [] + + monkeypatch.setattr( + "core.app.task_pipeline.easy_ui_based_generate_task_pipeline.PromptMessageUtil.prompt_messages_to_prompt_for_saving", + lambda mode, prompt_messages: "serialized-prompt", + ) + monkeypatch.setattr( + "core.app.task_pipeline.easy_ui_based_generate_task_pipeline.PromptTemplateParser.remove_template_variables", + lambda text: text.replace("{{name}}", "").strip(), + ) + monkeypatch.setattr( + "core.app.task_pipeline.easy_ui_based_generate_task_pipeline.naive_utc_now", + lambda: datetime(2024, 1, 1, tzinfo=UTC), + ) + monkeypatch.setattr( + "core.app.task_pipeline.easy_ui_based_generate_task_pipeline.time.perf_counter", lambda: 15.0 + ) + monkeypatch.setattr( + "core.app.task_pipeline.easy_ui_based_generate_task_pipeline.message_was_created.send", + lambda *args, **kwargs: sent_payloads.append((args, kwargs)), + ) + + pipeline._save_message(session=session, trace_manager=trace_manager) + + assert message_obj.message == "serialized-prompt" + assert message_obj.answer == "hello" + assert message_obj.provider_response_latency == 5.0 + trace_manager_double.add_trace_task.assert_called_once() + trace_task = trace_manager_double.add_trace_task.call_args.args[0] + assert trace_task.trace_type == TraceTaskName.MESSAGE_TRACE + assert trace_task.conversation_id == "conv" + assert trace_task.message_id == "msg" + assert trace_task.kwargs["trace_session_id"] == "session-1" + assert len(sent_payloads) == 1 + + def test_save_stopped_message_preserves_backend_reported_usage( + self, monkeypatch: pytest.MonkeyPatch, sqlite_session: Session + ) -> None: + pipeline = EasyUIBasedGenerateTaskPipeline( + application_generate_entity=_make_entity(ChatAppGenerateEntity, AppMode.CHAT), + queue_manager=_FakeQueueManager(), + conversation=_make_conversation(AppMode.CHAT), + message=_make_message(), + stream=True, + ) + _set_method(pipeline, "_model_config", _ModelConfigMode(mode="chat")) + pipeline._task_state.llm_result.usage = LLMUsage.empty_usage() + + message = _make_message() + message.message_tokens = 13 + message.answer_tokens = 8 + message.total_price = Decimal("0.21") + message.message_metadata = json.dumps({"usage": {"total_tokens": 21}}) + sqlite_session.add_all([_make_conversation(AppMode.CHAT), message]) + sqlite_session.flush() + monkeypatch.setattr( + "core.app.task_pipeline.easy_ui_based_generate_task_pipeline.PromptMessageUtil.prompt_messages_to_prompt_for_saving", + lambda _mode, _prompt_messages: "", + ) + monkeypatch.setattr( + "core.app.task_pipeline.easy_ui_based_generate_task_pipeline.message_was_created.send", + lambda *_args, **_kwargs: None, + ) + + pipeline._save_message(session=sqlite_session, preserve_existing_usage=True) + + assert message.message_tokens == 13 + assert message.answer_tokens == 8 + assert message.total_price == Decimal("0.21") + assert message.message_metadata is not None + assert json.loads(message.message_metadata)["usage"]["total_tokens"] == 21 + + @pytest.mark.parametrize("metadata", ["{", "[]"]) + def test_save_stopped_message_recovers_invalid_existing_metadata( + self, + metadata: str, + monkeypatch: pytest.MonkeyPatch, + sqlite_session: Session, + ) -> None: + pipeline = EasyUIBasedGenerateTaskPipeline( + application_generate_entity=_make_entity(ChatAppGenerateEntity, AppMode.CHAT), + queue_manager=_FakeQueueManager(), + conversation=_make_conversation(AppMode.CHAT), + message=_make_message(), + stream=True, + ) + _set_method(pipeline, "_model_config", _ModelConfigMode(mode="chat")) + pipeline._task_state.llm_result.usage = LLMUsage.empty_usage() + + message = _make_message() + message.message_tokens = 13 + message.answer_tokens = 8 + message.total_price = Decimal("0.21") + message.message_metadata = metadata + sqlite_session.add_all([_make_conversation(AppMode.CHAT), message]) + sqlite_session.flush() + monkeypatch.setattr( + "core.app.task_pipeline.easy_ui_based_generate_task_pipeline.PromptMessageUtil.prompt_messages_to_prompt_for_saving", + lambda _mode, _prompt_messages: "", + ) + monkeypatch.setattr( + "core.app.task_pipeline.easy_ui_based_generate_task_pipeline.message_was_created.send", + lambda *_args, **_kwargs: None, + ) + + pipeline._save_message(session=sqlite_session, preserve_existing_usage=True) + + assert message.message_tokens == 13 + assert message.answer_tokens == 8 + assert message.total_price == Decimal("0.21") + + def test_save_message_raises_when_message_not_found(self, sqlite_session: Session): + conversation = _make_conversation(AppMode.CHAT) + message = _make_message() + pipeline = EasyUIBasedGenerateTaskPipeline( + application_generate_entity=_make_entity(ChatAppGenerateEntity, AppMode.CHAT), + queue_manager=_FakeQueueManager(), + conversation=conversation, + message=message, + stream=False, + ) + session = sqlite_session + + with pytest.raises(ValueError, match="message msg not found"): + pipeline._save_message(session=session) + + def test_save_message_raises_when_conversation_not_found(self, sqlite_session: Session): + conversation = _make_conversation(AppMode.CHAT) + message = _make_message() + pipeline = EasyUIBasedGenerateTaskPipeline( + application_generate_entity=_make_entity(ChatAppGenerateEntity, AppMode.CHAT), + queue_manager=_FakeQueueManager(), + conversation=conversation, + message=message, + stream=False, + ) + session = sqlite_session + session.add(_make_message()) + session.flush() + + with pytest.raises(ValueError, match="Conversation conv not found"): + pipeline._save_message(session=session) + + def test_message_end_to_stream_response_includes_usage_metadata(self, monkeypatch: pytest.MonkeyPatch): + conversation = _make_conversation(AppMode.CHAT) + message = _make_message() + pipeline = EasyUIBasedGenerateTaskPipeline( + application_generate_entity=_make_entity(ChatAppGenerateEntity, AppMode.CHAT), + queue_manager=_FakeQueueManager(), + conversation=conversation, + message=message, + stream=False, + ) + pipeline._task_state.llm_result.usage = LLMUsage.from_metadata({"prompt_tokens": 1, "completion_tokens": 2}) + + class _Result: + def all(self): + return [] + + class _Session: + def __init__[**P](self, *args: P.args, **kwargs: P.kwargs): + pass + + def __enter__(self): + return self + + def __exit__(self, exc_type, exc, tb): + return False + + def scalars(self, *args, **kwargs): + return _Result() + + monkeypatch.setattr( + "core.app.task_pipeline.easy_ui_based_generate_task_pipeline.session_factory.create_session", + lambda: _Session(), + ) + + response = pipeline._message_end_to_stream_response() + + assert response.id == "msg" + usage_metadata = cast(dict[str, object], response.metadata["usage"]) + assert usage_metadata["prompt_tokens"] == 1 + + def test_record_files_returns_empty_list_when_message_has_no_files(self, monkeypatch: pytest.MonkeyPatch): + conversation = _make_conversation(AppMode.CHAT) + message = _make_message() + pipeline = EasyUIBasedGenerateTaskPipeline( + application_generate_entity=_make_entity(ChatAppGenerateEntity, AppMode.CHAT), + queue_manager=_FakeQueueManager(), + conversation=conversation, + message=message, + stream=False, + ) + + class _Result: + def all(self): + return [] + + class _Session: + def __init__[**P](self, *args: P.args, **kwargs: P.kwargs): + pass + + def __enter__(self): + return self + + def __exit__(self, exc_type, exc, tb): + return False + + def scalars(self, *args, **kwargs): + return _Result() + + monkeypatch.setattr( + "core.app.task_pipeline.easy_ui_based_generate_task_pipeline.session_factory.create_session", + lambda: _Session(), + ) + + response = pipeline._message_end_to_stream_response() + + assert response.files == [] + + def test_record_files_handles_local_fallback_and_tool_url_variants(self, monkeypatch: pytest.MonkeyPatch): + conversation = _make_conversation(AppMode.CHAT) + message = _make_message() + pipeline = EasyUIBasedGenerateTaskPipeline( + application_generate_entity=_make_entity(ChatAppGenerateEntity, AppMode.CHAT), + queue_manager=_FakeQueueManager(), + conversation=conversation, + message=message, + stream=False, + ) + message_files = [ + _message_file( + file_id="mf-local-fallback", + transfer_method=FileTransferMethod.LOCAL_FILE, + url="", + upload_file_id="upload-missing", + file_type=FileType.CUSTOM, + ), + _message_file( + file_id="mf-tool-http", + transfer_method=FileTransferMethod.TOOL_FILE, + url="http://cdn.example.com/file.txt?x=1", + upload_file_id=None, + file_type=FileType.CUSTOM, + ), + _message_file( + file_id="mf-tool-noext", + transfer_method=FileTransferMethod.TOOL_FILE, + url="tool/path/toolid", + upload_file_id=None, + file_type=FileType.CUSTOM, + ), + ] + + class _Result: + def __init__(self, items): + self._items = items + + def all(self): + return self._items + + class _Session: + def __init__[**P](self, *args: P.args, **kwargs: P.kwargs): + self.calls = 0 + + def __enter__(self): + return self + + def __exit__(self, exc_type, exc, tb): + return False + + def scalars(self, *args, **kwargs): + self.calls += 1 + return _Result(message_files if self.calls == 1 else []) + + monkeypatch.setattr( + "core.app.task_pipeline.easy_ui_based_generate_task_pipeline.session_factory.create_session", + lambda: _Session(), + ) + monkeypatch.setattr( + "core.app.task_pipeline.message_file_utils.file_helpers.get_signed_file_url", + lambda **kwargs: "local-fallback-signed", + ) + monkeypatch.setattr( + "core.app.task_pipeline.message_file_utils.sign_tool_file", + lambda **kwargs: "tool-signed", + ) + + response = pipeline._message_end_to_stream_response() + files = response.files + + assert files is not None + assert files[0]["url"] == "local-fallback-signed" + assert files[1]["filename"] == "file.txt" + assert files[2]["extension"] == ".bin" + + def test_agent_message_to_stream_response_builds_payload(self): + conversation = _make_conversation(AppMode.CHAT) + message = _make_message() + pipeline = EasyUIBasedGenerateTaskPipeline( + application_generate_entity=_make_entity(ChatAppGenerateEntity, AppMode.CHAT), + queue_manager=_FakeQueueManager(), + conversation=conversation, + message=message, + stream=True, + ) + + response = pipeline._agent_message_to_stream_response(answer="hello", message_id="msg") + + assert response.id == "msg" + assert response.answer == "hello" + + def test_agent_thought_to_stream_response_returns_none_when_not_found(self, monkeypatch: pytest.MonkeyPatch): + conversation = _make_conversation(AppMode.CHAT) + message = _make_message() + pipeline = EasyUIBasedGenerateTaskPipeline( + application_generate_entity=_make_entity(ChatAppGenerateEntity, AppMode.CHAT), + queue_manager=_FakeQueueManager(), + conversation=conversation, + message=message, + stream=True, + ) + + class _Session: + def __init__[**P](self, *args: P.args, **kwargs: P.kwargs): + pass + + def __enter__(self): + return self + + def __exit__(self, exc_type, exc, tb): + return False + + def scalar(self, *args, **kwargs): + return None + + monkeypatch.setattr( + "core.app.task_pipeline.easy_ui_based_generate_task_pipeline.session_factory.create_session", + lambda: _Session(), + ) + + response = pipeline._agent_thought_to_stream_response(QueueAgentThoughtEvent(agent_thought_id="missing")) + + assert response is None + + def test_handle_output_moderation_chunk_appends_token_when_not_directing(self): + conversation = _make_conversation(AppMode.CHAT) + message = _make_message() + pipeline = EasyUIBasedGenerateTaskPipeline( + application_generate_entity=_make_entity(ChatAppGenerateEntity, AppMode.CHAT), + queue_manager=_FakeQueueManager(), + conversation=conversation, + message=message, + stream=True, + ) + appended_tokens: list[str] = [] + + class _Moderation: + def should_direct_output(self): + return False + + def append_new_token(self, text): + appended_tokens.append(text) + + _set_method(pipeline, "output_moderation_handler", _Moderation()) + + result = pipeline._handle_output_moderation_chunk("next-token") + + assert result is False + assert appended_tokens == ["next-token"] + + +"""Exercise message-end file serialization against persisted SQLite rows. + +The suite covers empty results, all transfer methods, upload metadata batching, +and the fallback used when a local message file references a missing upload. +""" + + +SQLITE_MODELS = (MessageFile, UploadFile) + + +class TestEasyUiMessageEndFilesCases: + pytestmark = [ + pytest.mark.usefixtures("sqlite_session"), + pytest.mark.parametrize("sqlite_session", [SQLITE_MODELS], indirect=True), + ] + + class TestMessageEndStreamResponseFiles: + """Verify message-end file payloads from actual ORM query results.""" + + @pytest.fixture(autouse=True) + def bind_sqlite_engine( + self, sqlite_engine: _easy_ui_message_end_files_Engine, monkeypatch: pytest.MonkeyPatch + ) -> None: + """Bind sessions opened by the pipeline to the per-test SQLite engine.""" + + sqlite_session_maker = sessionmaker(bind=sqlite_engine, expire_on_commit=False) + monkeypatch.setattr("core.db.session_factory._session_maker", sqlite_session_maker) + + @pytest.fixture + def pipeline( + self, monkeypatch: pytest.MonkeyPatch, redis_transport: tuple[RedisClientWrapper, MagicMock] + ) -> EasyUIBasedGenerateTaskPipeline: + """Construct the pipeline with validated request, model, and task state.""" + monkeypatch.setattr("core.app.apps.base_app_queue_manager.redis_client", redis_transport[0]) + request = ChatAppGenerateEntity( + task_id=str(uuid.uuid4()), + app_config=EasyUIBasedAppConfig( + tenant_id="tenant-id", + app_id="app-id", + app_mode=AppMode.CHAT, + app_model_config_from=EasyUIBasedAppModelConfigFrom.APP_LATEST_CONFIG, + app_model_config_dict={}, + model=ModelConfigEntity(provider="test-provider", model="test-model"), + prompt_template=PromptTemplateEntity(prompt_type=PromptTemplateEntity.PromptType.SIMPLE), + ), + model_conf=make_model_config(provider="test-provider", model="test-model", mode="chat"), + inputs={}, + files=[], + user_id="user-id", + stream=True, + invoke_from=InvokeFrom.WEB_APP, + ) + conversation = Conversation( + id=str(uuid.uuid4()), + app_id="app-id", + mode=AppMode.CHAT, + name="Conversation", + from_source=ConversationFromSource.API, + _inputs={}, + ) + message = Message(id=str(uuid.uuid4())) + message.created_at = datetime.now() + queue = MessageBasedAppQueueManager( + task_id=request.task_id, + user_id=request.user_id, + invoke_from=request.invoke_from, + conversation_id=conversation.id, + app_mode=AppMode.CHAT, + message_id=message.id, + ) + return EasyUIBasedGenerateTaskPipeline( + application_generate_entity=request, + queue_manager=queue, + conversation=conversation, + message=message, + stream=True, + ) + + @staticmethod + def _message_file( + *, + transfer_method: FileTransferMethod, + url: str | None = None, + upload_file_id: str | None = None, + ) -> MessageFile: + return MessageFile( + message_id=str(uuid.uuid4()), + type=FileType.IMAGE, + transfer_method=transfer_method, + created_by_role=CreatorUserRole.ACCOUNT, + created_by=str(uuid.uuid4()), + url=url, + upload_file_id=upload_file_id, + ) + + @pytest.fixture + def message_file_local(self) -> MessageFile: + """Create an unpersisted local-file row.""" + + return self._message_file( + transfer_method=FileTransferMethod.LOCAL_FILE, + upload_file_id=str(uuid.uuid4()), + ) + + @pytest.fixture + def message_file_remote(self) -> MessageFile: + """Create an unpersisted remote-file row.""" + + return self._message_file( + transfer_method=FileTransferMethod.REMOTE_URL, + url="https://example.com/image.jpg", + ) + + @pytest.fixture + def message_file_tool(self) -> MessageFile: + """Create an unpersisted tool-file row.""" + + return self._message_file( + transfer_method=FileTransferMethod.TOOL_FILE, + url="tool_file_123.png", + ) + + @pytest.fixture + def upload_file(self, message_file_local: MessageFile) -> UploadFile: + """Create upload metadata matching the local message-file reference.""" + + upload = UploadFile( + tenant_id=str(uuid.uuid4()), + storage_type=StorageType.LOCAL, + key="uploads/test_image.png", + name="test_image.png", + size=1024, + extension="png", + mime_type="image/png", + created_by_role=CreatorUserRole.ACCOUNT, + created_by=str(uuid.uuid4()), + created_at=datetime.now(), + used=True, + ) + upload.id = message_file_local.upload_file_id or str(uuid.uuid4()) + return upload + + @staticmethod + def _persist(session: Session, *rows: MessageFile | UploadFile) -> None: + session.add_all(rows) + session.commit() + + def test_message_end_with_no_files( + self, sqlite_session: Session, pipeline: EasyUIBasedGenerateTaskPipeline + ) -> None: + """Rows for another message do not leak into an empty files array.""" + + unrelated_file = self._message_file( + transfer_method=FileTransferMethod.REMOTE_URL, + url="https://example.com/unrelated.png", + ) + self._persist(sqlite_session, unrelated_file) + + result = pipeline._message_end_to_stream_response() + + assert sqlite_session.get(MessageFile, unrelated_file.id) is unrelated_file + assert isinstance(result, MessageEndStreamResponse) + assert result.files == [] + assert result.id == pipeline._message_id + assert result.metadata == { + "usage": pipeline._task_state.llm_result.usage.model_dump(exclude_none=True), + "retriever_resources": [], + "reasoning": {}, + } + + def test_message_end_with_local_file( + self, + sqlite_session: Session, + pipeline: EasyUIBasedGenerateTaskPipeline, + message_file_local: MessageFile, + upload_file: UploadFile, + ) -> None: + """Local files include persisted upload metadata and a signed URL.""" + + message_file_local.message_id = pipeline._message_id + self._persist(sqlite_session, message_file_local, upload_file) + + with patch( + "core.app.task_pipeline.message_file_utils.file_helpers.get_signed_file_url", + return_value="https://example.com/signed-url?signature=abc123", + ) as get_signed_url: + result = pipeline._message_end_to_stream_response() + + assert result.files is not None + assert len(result.files) == 1 + file_dict = result.files[0] + assert file_dict["related_id"] == message_file_local.id + assert file_dict["filename"] == "test_image.png" + assert file_dict["mime_type"] == "image/png" + assert file_dict["size"] == 1024 + assert file_dict["extension"] == ".png" + assert file_dict["type"] == "image" + assert file_dict["transfer_method"] == FileTransferMethod.LOCAL_FILE.value + assert file_dict["url"].startswith("https://example.com/signed-url") + assert file_dict["upload_file_id"] == message_file_local.upload_file_id + assert file_dict["remote_url"] == "" + get_signed_url.assert_called_once_with(upload_file_id=upload_file.id) + + def test_message_end_with_remote_url( + self, sqlite_session: Session, pipeline: EasyUIBasedGenerateTaskPipeline, message_file_remote: MessageFile + ) -> None: + """Remote files retain their source URL and derived filename.""" + + message_file_remote.message_id = pipeline._message_id + self._persist(sqlite_session, message_file_remote) + + result = pipeline._message_end_to_stream_response() + + assert result.files is not None + assert len(result.files) == 1 + file_dict = result.files[0] + assert file_dict["related_id"] == message_file_remote.id + assert file_dict["filename"] == "image.jpg" + assert file_dict["url"] == "https://example.com/image.jpg" + assert file_dict["extension"] == ".jpg" + assert file_dict["type"] == "image" + assert file_dict["transfer_method"] == FileTransferMethod.REMOTE_URL.value + assert file_dict["remote_url"] == "https://example.com/image.jpg" + assert file_dict["upload_file_id"] == message_file_remote.id + + def test_message_end_with_tool_file_http( + self, sqlite_session: Session, pipeline: EasyUIBasedGenerateTaskPipeline, message_file_tool: MessageFile + ) -> None: + """HTTP tool-file URLs pass through unchanged.""" + + message_file_tool.message_id = pipeline._message_id + message_file_tool.url = "https://example.com/tool_file.png" + self._persist(sqlite_session, message_file_tool) + + result = pipeline._message_end_to_stream_response() + + assert result.files is not None + file_dict = result.files[0] + assert file_dict["url"] == "https://example.com/tool_file.png" + assert file_dict["filename"] == "tool_file.png" + assert file_dict["extension"] == ".png" + assert file_dict["transfer_method"] == FileTransferMethod.TOOL_FILE.value + + def test_message_end_with_tool_file_local( + self, sqlite_session: Session, pipeline: EasyUIBasedGenerateTaskPipeline, message_file_tool: MessageFile + ) -> None: + """Local tool-file identifiers are signed at the external boundary.""" + + message_file_tool.message_id = pipeline._message_id + self._persist(sqlite_session, message_file_tool) + + with patch( + "core.app.task_pipeline.message_file_utils.sign_tool_file", + return_value="https://example.com/signed-tool-file.png?signature=xyz", + ) as sign_tool: + result = pipeline._message_end_to_stream_response() + + assert result.files is not None + file_dict = result.files[0] + assert file_dict["url"].startswith("https://example.com/signed-tool-file.png") + assert file_dict["filename"] == "tool_file_123.png" + assert file_dict["extension"] == ".png" + assert file_dict["transfer_method"] == FileTransferMethod.TOOL_FILE.value + sign_tool.assert_called_once_with(tool_file_id="tool_file_123", extension=".png") + + def test_message_end_with_tool_file_long_extension( + self, sqlite_session: Session, pipeline: EasyUIBasedGenerateTaskPipeline, message_file_tool: MessageFile + ) -> None: + """Overlong tool-file extensions use the safe binary fallback.""" + + message_file_tool.message_id = pipeline._message_id + message_file_tool.url = "tool_file_abc.verylongextension" + self._persist(sqlite_session, message_file_tool) + + with patch( + "core.app.task_pipeline.message_file_utils.sign_tool_file", + return_value="https://example.com/signed.bin", + ) as sign_tool: + result = pipeline._message_end_to_stream_response() + + assert result.files is not None + assert result.files[0]["extension"] == ".bin" + sign_tool.assert_called_once_with(tool_file_id="tool_file_abc", extension=".bin") + + def test_message_end_with_multiple_files( + self, + sqlite_session: Session, + pipeline: EasyUIBasedGenerateTaskPipeline, + message_file_local: MessageFile, + message_file_remote: MessageFile, + upload_file: UploadFile, + ) -> None: + """The response contains every persisted file associated with the message.""" + + message_file_local.message_id = pipeline._message_id + message_file_remote.message_id = pipeline._message_id + self._persist(sqlite_session, message_file_local, message_file_remote, upload_file) + + with patch( + "core.app.task_pipeline.message_file_utils.file_helpers.get_signed_file_url", + return_value="https://example.com/signed-url?signature=abc123", + ): + result = pipeline._message_end_to_stream_response() + + assert result.files is not None + assert {file["related_id"] for file in result.files} == {message_file_local.id, message_file_remote.id} + + def test_message_end_with_local_file_no_upload_file( + self, sqlite_session: Session, pipeline: EasyUIBasedGenerateTaskPipeline, message_file_local: MessageFile + ) -> None: + """A missing upload row still signs the stored upload identifier.""" + + message_file_local.message_id = pipeline._message_id + self._persist(sqlite_session, message_file_local) + + with patch( + "core.app.task_pipeline.message_file_utils.file_helpers.get_signed_file_url", + return_value="https://example.com/fallback-url?signature=def456", + ) as get_signed_url: + result = pipeline._message_end_to_stream_response() + + assert result.files is not None + assert len(result.files) == 1 + assert result.files[0]["url"].startswith("https://example.com/fallback-url") + get_signed_url.assert_called_once_with(upload_file_id=str(message_file_local.upload_file_id)) diff --git a/api/tests/unit_tests/core/app/task_pipeline/test_easy_ui_based_generate_task_pipeline_core.py b/api/tests/unit_tests/core/app/task_pipeline/test_easy_ui_based_generate_task_pipeline_core.py deleted file mode 100644 index afe1c85e94d..00000000000 --- a/api/tests/unit_tests/core/app/task_pipeline/test_easy_ui_based_generate_task_pipeline_core.py +++ /dev/null @@ -1,1735 +0,0 @@ -from __future__ import annotations - -import json -from collections.abc import Generator, Sequence -from dataclasses import dataclass -from datetime import UTC, datetime -from decimal import Decimal -from threading import Thread -from typing import cast -from unittest.mock import Mock - -import pytest -from sqlalchemy.orm import Session - -from core.app.app_config.entities import ( - AppAdditionalFeatures, - EasyUIBasedAppConfig, - EasyUIBasedAppModelConfigFrom, - ModelConfigEntity, - PromptTemplateEntity, -) -from core.app.apps.base_app_queue_manager import AppQueueManager, PublishFrom -from core.app.entities.app_invoke_entities import ChatAppGenerateEntity, CompletionAppGenerateEntity, InvokeFrom -from core.app.entities.queue_entities import ( - AppQueueEvent, - MessageQueueMessage, - QueueAgentMessageEvent, - QueueAgentThoughtEvent, - QueueAnnotationReplyEvent, - QueueErrorEvent, - QueueLLMChunkEvent, - QueueMessageEndEvent, - QueueMessageFileEvent, - QueueMessageReplaceEvent, - QueuePingEvent, - QueueRetrieverResourcesEvent, - QueueStopEvent, - WorkflowQueueMessage, -) -from core.app.entities.task_entities import ( - AgentMessageStreamResponse, - AgentThoughtStreamResponse, - ChatbotAppStreamResponse, - CompletionAppStreamResponse, - ErrorStreamResponse, - MessageAudioEndStreamResponse, - MessageAudioStreamResponse, - MessageEndStreamResponse, - MessageFileStreamResponse, - MessageReplaceStreamResponse, - MessageStreamResponse, - PingStreamResponse, - StreamEvent, -) -from core.app.task_pipeline.easy_ui_based_generate_task_pipeline import EasyUIBasedGenerateTaskPipeline -from core.base.tts import AppGeneratorTTSPublisher, AudioTrunk -from core.ops.entities.trace_entity import TraceTaskName -from core.ops.ops_trace_manager import TraceQueueManager -from graphon.file import FileTransferMethod, FileType -from graphon.model_runtime.entities.llm_entities import LLMResult, LLMResultChunk, LLMResultChunkDelta, LLMUsage -from graphon.model_runtime.entities.message_entities import AssistantPromptMessage, TextPromptMessageContent -from models.enums import ConversationFromSource, CreatorUserRole -from models.model import AppMode, Conversation, Message, MessageAgentThought, MessageFile, UploadFile -from tests.unit_tests.model_factories import make_upload_file - - -class _DummyModelConf: - def __init__(self) -> None: - self.model = "mock" - - -class _UnknownQueueEvent: - pass - - -@dataclass(eq=False, repr=False) -class _AnnotationReply: - content: str - - -@dataclass(eq=False, repr=False) -class _ModelConfigMode: - mode: str - - -@dataclass(eq=False, repr=False) -class _ProviderModelBundle: - model_type_instance: object - - -class _AudioPublisher: - def __init__(self, status: str, audio: str) -> None: - self._audio = AudioTrunk(status, audio) - - def check_and_get_audio(self) -> AudioTrunk: - return self._audio - - -class _TraceManagerDouble: - def __init__(self) -> None: - self.add_trace_task = Mock() - - -class _FakeQueueManager(AppQueueManager): - def __init__(self) -> None: - self._events: list[MessageQueueMessage | WorkflowQueueMessage] = [] - self.published_events: list[AppQueueEvent] = [] - - def set_events(self, events: list[MessageQueueMessage | WorkflowQueueMessage]) -> None: - self._events = events - - def listen(self) -> Generator[MessageQueueMessage | WorkflowQueueMessage]: - yield from self._events - - def publish(self, event: AppQueueEvent, pub_from: PublishFrom) -> None: - self.published_events.append(event) - - def _publish(self, event: AppQueueEvent, pub_from: PublishFrom) -> None: - self.published_events.append(event) - - -def _queue_message(event: AppQueueEvent) -> MessageQueueMessage: - return MessageQueueMessage( - task_id="task", - app_mode=AppMode.CHAT.value, - message_id="msg", - conversation_id="conv", - event=event, - ) - - -def _unknown_queue_message() -> MessageQueueMessage: - return MessageQueueMessage.model_construct( - task_id="task", - app_mode=AppMode.CHAT.value, - message_id="msg", - conversation_id="conv", - event=cast(AppQueueEvent, _UnknownQueueEvent()), - ) - - -def _make_conversation(app_mode: AppMode) -> Conversation: - conversation = Conversation( - app_id="app", - app_model_config_id=None, - model_provider=None, - override_model_configs=None, - model_id=None, - mode=app_mode, - name="conversation", - _inputs={}, - introduction="", - system_instruction="", - system_instruction_tokens=0, - status="normal", - invoke_from=InvokeFrom.WEB_APP, - from_source=ConversationFromSource.API, - from_end_user_id="user", - from_account_id=None, - ) - conversation.id = "conv" - conversation.mode = app_mode - return conversation - - -def _make_message() -> Message: - message = Message( - app_id="app", - conversation_id="conv", - _inputs={}, - query="query", - message="", - message_tokens=0, - message_unit_price=0, - message_price_unit=0, - answer="", - answer_tokens=0, - answer_unit_price=0, - answer_price_unit=0, - provider_response_latency=0, - total_price=0, - currency="USD", - invoke_from=InvokeFrom.WEB_APP, - from_source=ConversationFromSource.API, - from_end_user_id="user", - from_account_id=None, - app_mode=AppMode.CHAT, - ) - message.id = "msg" - message.created_at = datetime.now(UTC) - return message - - -def _message_file( - *, - file_id: str, - transfer_method: FileTransferMethod, - url: str | None, - upload_file_id: str | None, - file_type: FileType = FileType.IMAGE, -) -> MessageFile: - message_file = MessageFile( - message_id="msg", - type=file_type, - transfer_method=transfer_method, - created_by_role=CreatorUserRole.ACCOUNT, - created_by="user", - url=url, - upload_file_id=upload_file_id, - ) - message_file.id = file_id - return message_file - - -def _upload_file(*, file_id: str, name: str, mime_type: str, size: int, extension: str) -> UploadFile: - return make_upload_file( - file_id=file_id, - tenant_id="tenant", - key=f"uploads/{file_id}", - name=name, - size=size, - extension=extension, - mime_type=mime_type, - created_by="user", - created_at=datetime.now(UTC), - ) - - -def _agent_thought() -> MessageAgentThought: - thought = MessageAgentThought( - message_id="msg", - position=1, - created_by_role=CreatorUserRole.ACCOUNT, - created_by="user", - thought="t", - observation="o", - tool="tool", - tool_labels_str="{}", - tool_input="input", - answer="answer", - message_files="[]", - ) - thought.id = "thought" - return thought - - -def _make_app_config(app_mode: AppMode) -> EasyUIBasedAppConfig: - return EasyUIBasedAppConfig( - tenant_id="tenant", - app_id="app", - app_mode=app_mode, - app_model_config_from=EasyUIBasedAppModelConfigFrom.APP_LATEST_CONFIG, - app_model_config_id="model-config", - app_model_config_dict={}, - model=ModelConfigEntity(provider="mock", model="mock"), - prompt_template=PromptTemplateEntity( - prompt_type=PromptTemplateEntity.PromptType.SIMPLE, - simple_prompt_template="hi", - ), - additional_features=AppAdditionalFeatures(), - variables=[], - ) - - -def _make_entity(entity_cls, app_mode: AppMode): - app_config = _make_app_config(app_mode) - return entity_cls.model_construct( - task_id="task", - app_config=app_config, - model_conf=_DummyModelConf(), - file_upload_config=None, - conversation_id=None, - inputs={}, - query="hello", - files=[], - parent_message_id=None, - user_id="user", - stream=False, - invoke_from=InvokeFrom.WEB_APP, - extras={}, - call_depth=0, - trace_manager=None, - ) - - -def _make_pipeline( - entity_cls: type[ChatAppGenerateEntity] | type[CompletionAppGenerateEntity] = ChatAppGenerateEntity, - app_mode: AppMode = AppMode.CHAT, - *, - stream: bool = True, -) -> tuple[EasyUIBasedGenerateTaskPipeline, _FakeQueueManager]: - queue_manager = _FakeQueueManager() - pipeline = EasyUIBasedGenerateTaskPipeline( - application_generate_entity=_make_entity(entity_cls, app_mode), - queue_manager=queue_manager, - conversation=_make_conversation(app_mode), - message=_make_message(), - stream=stream, - ) - return pipeline, queue_manager - - -def _set_queue_events( - pipeline: EasyUIBasedGenerateTaskPipeline, - events: Sequence[MessageQueueMessage | WorkflowQueueMessage], -) -> None: - cast(_FakeQueueManager, pipeline.queue_manager).set_events(list(events)) - - -def _set_method(obj: object, name: str, value: object) -> None: - object.__setattr__(obj, name, value) - - -class TestEasyUiBasedGenerateTaskPipeline: - def test_to_blocking_response_chat(self): - conversation = _make_conversation(AppMode.CHAT) - message = _make_message() - - pipeline = EasyUIBasedGenerateTaskPipeline( - application_generate_entity=_make_entity(ChatAppGenerateEntity, AppMode.CHAT), - queue_manager=_FakeQueueManager(), - conversation=conversation, - message=message, - stream=False, - ) - pipeline._task_state.llm_result.message.content = "answer" - - def _gen(): - yield MessageEndStreamResponse(task_id="task", id="msg") - - response = pipeline._to_blocking_response(_gen()) - - assert response.data.answer == "answer" - - def test_to_blocking_response_completion(self): - conversation = _make_conversation(AppMode.COMPLETION) - message = _make_message() - - pipeline = EasyUIBasedGenerateTaskPipeline( - application_generate_entity=_make_entity(CompletionAppGenerateEntity, AppMode.COMPLETION), - queue_manager=_FakeQueueManager(), - conversation=conversation, - message=message, - stream=False, - ) - pipeline._task_state.llm_result.message.content = "answer" - - def _gen(): - yield MessageEndStreamResponse(task_id="task", id="msg") - - response = pipeline._to_blocking_response(_gen()) - - assert response.data.answer == "answer" - - def test_listen_audio_msg_returns_none_when_no_publisher(self): - conversation = _make_conversation(AppMode.CHAT) - message = _make_message() - - pipeline = EasyUIBasedGenerateTaskPipeline( - application_generate_entity=_make_entity(ChatAppGenerateEntity, AppMode.CHAT), - queue_manager=_FakeQueueManager(), - conversation=conversation, - message=message, - stream=False, - ) - - assert pipeline._listen_audio_msg(publisher=None, task_id="task") is None - - def test_process_stream_response_handles_chunks_and_end(self, monkeypatch: pytest.MonkeyPatch): - conversation = _make_conversation(AppMode.CHAT) - message = _make_message() - - pipeline = EasyUIBasedGenerateTaskPipeline( - application_generate_entity=_make_entity(ChatAppGenerateEntity, AppMode.CHAT), - queue_manager=_FakeQueueManager(), - conversation=conversation, - message=message, - stream=True, - ) - - chunk = LLMResultChunk( - model="mock", - prompt_messages=[], - delta=LLMResultChunkDelta( - index=0, - message=AssistantPromptMessage( - content=[TextPromptMessageContent(data="hi"), TextPromptMessageContent(data="yo")] - ), - ), - ) - llm_result = LLMResult( - model="mock", - prompt_messages=[], - message=AssistantPromptMessage(content="done"), - usage=LLMUsage.empty_usage(), - ) - - events = [ - _queue_message(QueueLLMChunkEvent(chunk=chunk)), - _queue_message(QueueMessageReplaceEvent(text="replace", reason="output_moderation")), - _queue_message(QueuePingEvent()), - _queue_message(QueueMessageEndEvent(llm_result=llm_result)), - ] - - _set_queue_events(pipeline, events) - - def _message_event_type(message_id: str) -> StreamEvent: - return StreamEvent.MESSAGE - - def _message_end() -> MessageEndStreamResponse: - return MessageEndStreamResponse(task_id="task", id="msg") - - _set_method(pipeline._message_cycle_manager, "get_message_event_type", _message_event_type) - _set_method(pipeline, "handle_output_moderation_when_task_finished", lambda completion: None) - _set_method(pipeline, "_message_end_to_stream_response", _message_end) - _set_method(pipeline, "_save_message", lambda **kwargs: None) - - class _Session: - def __init__[**P](self, *args: P.args, **kwargs: P.kwargs): - pass - - def __enter__(self): - return self - - def __exit__(self, exc_type, exc, tb): - return False - - def commit(self): - return None - - monkeypatch.setattr( - "core.app.task_pipeline.easy_ui_based_generate_task_pipeline.session_factory.create_session", - lambda: _Session(), - ) - - responses = list(pipeline._process_stream_response(publisher=None)) - - message_response = next(item for item in responses if isinstance(item, MessageStreamResponse)) - assert message_response.answer == "hiyo" - replace_response = next(item for item in responses if isinstance(item, MessageReplaceStreamResponse)) - assert replace_response.answer == "replace" - assert any(isinstance(item, PingStreamResponse) for item in responses) - assert isinstance(responses[-1], MessageEndStreamResponse) - assert pipeline._task_state.llm_result.message.content == "done" - - def test_handle_output_moderation_chunk_directs_output(self): - conversation = _make_conversation(AppMode.CHAT) - message = _make_message() - - pipeline = EasyUIBasedGenerateTaskPipeline( - application_generate_entity=_make_entity(ChatAppGenerateEntity, AppMode.CHAT), - queue_manager=_FakeQueueManager(), - conversation=conversation, - message=message, - stream=True, - ) - - class _Moderation: - def should_direct_output(self): - return True - - def get_final_output(self): - return "final" - - _set_method(pipeline, "output_moderation_handler", _Moderation()) - - result = pipeline._handle_output_moderation_chunk("token") - - assert result is True - events = cast(_FakeQueueManager, pipeline.queue_manager).published_events - assert any(isinstance(event, QueueLLMChunkEvent) for event in events) - assert any(isinstance(event, QueueStopEvent) for event in events) - - def test_handle_stop_updates_usage(self, monkeypatch: pytest.MonkeyPatch): - conversation = _make_conversation(AppMode.CHAT) - message = _make_message() - - class _ModelType: - def calc_response_usage(self, model, credentials, prompt_tokens, completion_tokens): - return LLMUsage.from_metadata( - { - "prompt_tokens": prompt_tokens, - "completion_tokens": completion_tokens, - } - ) - - class _ModelConf: - def __init__(self) -> None: - self.model = "mock" - self.credentials = {} - self.provider_model_bundle = _ProviderModelBundle(model_type_instance=_ModelType()) - - app_config = _make_app_config(AppMode.CHAT) - application_generate_entity = ChatAppGenerateEntity.model_construct( - task_id="task", - app_config=app_config, - model_conf=_ModelConf(), - file_upload_config=None, - conversation_id=None, - inputs={}, - query="hello", - files=[], - parent_message_id=None, - user_id="user", - stream=False, - invoke_from=InvokeFrom.WEB_APP, - extras={}, - call_depth=0, - trace_manager=None, - ) - - pipeline = EasyUIBasedGenerateTaskPipeline( - application_generate_entity=application_generate_entity, - queue_manager=_FakeQueueManager(), - conversation=conversation, - message=message, - stream=False, - ) - pipeline._task_state.llm_result.prompt_messages = [AssistantPromptMessage(content="prompt")] - pipeline._task_state.llm_result.message = AssistantPromptMessage(content="answer") - - calls: list[int] = [] - - class _FakeModelInstance: - def __init__(self, provider_model_bundle, model): - pass - - def get_llm_num_tokens(self, messages): - calls.append(1) - return 10 if len(calls) == 1 else 5 - - monkeypatch.setattr( - "core.app.task_pipeline.easy_ui_based_generate_task_pipeline.ModelInstance", - _FakeModelInstance, - ) - - pipeline._handle_stop(QueueStopEvent(stopped_by=QueueStopEvent.StopBy.USER_MANUAL)) - - assert pipeline._task_state.llm_result.usage.prompt_tokens == 10 - assert pipeline._task_state.llm_result.usage.completion_tokens == 5 - - def test_record_files_builds_file_payloads(self, monkeypatch: pytest.MonkeyPatch): - conversation = _make_conversation(AppMode.CHAT) - message = _make_message() - - pipeline = EasyUIBasedGenerateTaskPipeline( - application_generate_entity=_make_entity(ChatAppGenerateEntity, AppMode.CHAT), - queue_manager=_FakeQueueManager(), - conversation=conversation, - message=message, - stream=False, - ) - - message_files = [ - _message_file( - file_id="mf-1", - transfer_method=FileTransferMethod.REMOTE_URL, - url="http://example.com/a.png", - upload_file_id=None, - ), - _message_file( - file_id="mf-2", - transfer_method=FileTransferMethod.LOCAL_FILE, - url="", - upload_file_id="upload-1", - ), - _message_file( - file_id="mf-3", - transfer_method=FileTransferMethod.TOOL_FILE, - url="tool/file.bin", - upload_file_id=None, - file_type=FileType.CUSTOM, - ), - ] - upload_files = [ - _upload_file( - file_id="upload-1", - name="local.png", - mime_type="image/png", - size=123, - extension="png", - ) - ] - - class _Result: - def __init__(self, items): - self._items = items - - def all(self): - return self._items - - class _Session: - def __init__[**P](self, *args: P.args, **kwargs: P.kwargs): - self.calls = 0 - - def __enter__(self): - return self - - def __exit__(self, exc_type, exc, tb): - return False - - def scalars(self, *args, **kwargs): - self.calls += 1 - return _Result(message_files if self.calls == 1 else upload_files) - - monkeypatch.setattr( - "core.app.task_pipeline.easy_ui_based_generate_task_pipeline.session_factory.create_session", - lambda: _Session(), - ) - monkeypatch.setattr( - "core.app.task_pipeline.message_file_utils.file_helpers.get_signed_file_url", - lambda **kwargs: "signed-url", - ) - monkeypatch.setattr( - "core.app.task_pipeline.message_file_utils.sign_tool_file", - lambda **kwargs: "signed-tool", - ) - - response = pipeline._message_end_to_stream_response() - files = response.files - - assert files - assert len(files) == 3 - - def test_process_stream_response_handles_annotation_and_error(self, monkeypatch: pytest.MonkeyPatch): - conversation = _make_conversation(AppMode.CHAT) - message = _make_message() - - pipeline = EasyUIBasedGenerateTaskPipeline( - application_generate_entity=_make_entity(ChatAppGenerateEntity, AppMode.CHAT), - queue_manager=_FakeQueueManager(), - conversation=conversation, - message=message, - stream=True, - ) - - agent_chunk = LLMResultChunk( - model="mock", - prompt_messages=[], - delta=LLMResultChunkDelta( - index=0, - message=AssistantPromptMessage(content="agent"), - ), - ) - - events = [ - _queue_message(QueueAnnotationReplyEvent(message_annotation_id="ann")), - _queue_message(QueueAgentThoughtEvent(agent_thought_id="thought")), - _queue_message(QueueMessageFileEvent(message_file_id="file")), - _queue_message(QueueAgentMessageEvent(chunk=agent_chunk)), - _queue_message(QueueErrorEvent(error=ValueError("boom"))), - ] - - _set_queue_events(pipeline, events) - - def _agent_thought_response(event: QueueAgentThoughtEvent) -> AgentThoughtStreamResponse: - return AgentThoughtStreamResponse(task_id="task", id=event.agent_thought_id, position=1) - - def _file_response(event: QueueMessageFileEvent) -> MessageFileStreamResponse: - return MessageFileStreamResponse( - task_id="task", id=event.message_file_id, type="image", belongs_to="user", url="file" - ) - - def _agent_message_response(answer: str, message_id: str) -> AgentMessageStreamResponse: - return AgentMessageStreamResponse(task_id="task", id=message_id, answer=answer) - - _set_method( - pipeline._message_cycle_manager, - "handle_annotation_reply", - lambda event, session: _AnnotationReply(content="annotated"), - ) - _set_method(pipeline, "_agent_thought_to_stream_response", _agent_thought_response) - _set_method(pipeline._message_cycle_manager, "message_file_to_stream_response", _file_response) - _set_method(pipeline, "_agent_message_to_stream_response", _agent_message_response) - _set_method(pipeline, "handle_error", lambda **kwargs: ValueError("boom")) - _set_method(pipeline, "error_to_stream_response", lambda err: ErrorStreamResponse(task_id="task", err=err)) - - class _Session: - def __init__[**P](self, *args: P.args, **kwargs: P.kwargs): - pass - - def __enter__(self): - return self - - def __exit__(self, exc_type, exc, tb): - return False - - def commit(self): - return None - - monkeypatch.setattr( - "core.app.task_pipeline.easy_ui_based_generate_task_pipeline.session_factory.create_session", - lambda: _Session(), - ) - - responses = list(pipeline._process_stream_response(publisher=None)) - - assert any(isinstance(response, AgentThoughtStreamResponse) for response in responses) - assert any(isinstance(response, MessageFileStreamResponse) for response in responses) - agent_response = next(response for response in responses if isinstance(response, AgentMessageStreamResponse)) - assert agent_response.answer == "agent" - assert isinstance(responses[-1], ErrorStreamResponse) - assert isinstance(responses[-1].err, ValueError) - assert pipeline._task_state.llm_result.message.content == "annotated" - - def test_process_stream_response_error_event_adds_trace_task(self, monkeypatch: pytest.MonkeyPatch): - conversation = _make_conversation(AppMode.CHAT) - message = _make_message() - application_generate_entity = _make_entity(ChatAppGenerateEntity, AppMode.CHAT) - application_generate_entity.extras = {"trace_session_id": "session-1"} - - pipeline = EasyUIBasedGenerateTaskPipeline( - application_generate_entity=application_generate_entity, - queue_manager=_FakeQueueManager(), - conversation=conversation, - message=message, - stream=True, - ) - _set_queue_events(pipeline, [_queue_message(QueueErrorEvent(error=ValueError("boom")))]) - _set_method(pipeline, "handle_error", lambda **kwargs: ValueError("boom")) - _set_method(pipeline, "error_to_stream_response", lambda err: ErrorStreamResponse(task_id="task", err=err)) - - trace_manager_double = _TraceManagerDouble() - trace_manager = cast(TraceQueueManager, trace_manager_double) - - class _Session: - def __enter__(self): - return self - - def __exit__(self, exc_type, exc, tb): - return False - - def commit(self): - return None - - monkeypatch.setattr( - "core.app.task_pipeline.easy_ui_based_generate_task_pipeline.session_factory.create_session", - lambda: _Session(), - ) - - responses = list(pipeline._process_stream_response(publisher=None, trace_manager=trace_manager)) - - assert len(responses) == 1 - assert isinstance(responses[0], ErrorStreamResponse) - trace_manager_double.add_trace_task.assert_called_once() - trace_task = trace_manager_double.add_trace_task.call_args.args[0] - assert trace_task.trace_type == TraceTaskName.MESSAGE_TRACE - assert trace_task.conversation_id == "conv" - assert trace_task.message_id == "msg" - assert trace_task.kwargs["trace_session_id"] == "session-1" - - def test_agent_thought_to_stream_response_returns_payload(self, monkeypatch: pytest.MonkeyPatch): - conversation = _make_conversation(AppMode.CHAT) - message = _make_message() - - pipeline = EasyUIBasedGenerateTaskPipeline( - application_generate_entity=_make_entity(ChatAppGenerateEntity, AppMode.CHAT), - queue_manager=_FakeQueueManager(), - conversation=conversation, - message=message, - stream=True, - ) - - agent_thought = _agent_thought() - - class _Session: - def __init__[**P](self, *args: P.args, **kwargs: P.kwargs): - pass - - def __enter__(self): - return self - - def __exit__(self, exc_type, exc, tb): - return False - - def scalar(self, *args, **kwargs): - return agent_thought - - monkeypatch.setattr( - "core.app.task_pipeline.easy_ui_based_generate_task_pipeline.session_factory.create_session", - lambda: _Session(), - ) - - response = pipeline._agent_thought_to_stream_response(QueueAgentThoughtEvent(agent_thought_id="thought")) - - assert response is not None - assert response.id == "thought" - assert response.thought == "t" - - def test_agent_thought_to_stream_response_normalizes_null_display_fields(self, monkeypatch: pytest.MonkeyPatch): - conversation = _make_conversation(AppMode.CHAT) - message = _make_message() - - pipeline = EasyUIBasedGenerateTaskPipeline( - application_generate_entity=_make_entity(ChatAppGenerateEntity, AppMode.CHAT), - queue_manager=_FakeQueueManager(), - conversation=conversation, - message=message, - stream=True, - ) - - agent_thought = _agent_thought() - agent_thought.thought = None - agent_thought.answer = None - agent_thought.observation = None - agent_thought.tool = None - agent_thought.tool_input = None - agent_thought.message_files = None - - class _Session: - def __init__(self, *args, **kwargs): - pass - - def __enter__(self): - return self - - def __exit__(self, exc_type, exc, tb): - return False - - def scalar(self, *args, **kwargs): - return agent_thought - - monkeypatch.setattr( - "core.app.task_pipeline.easy_ui_based_generate_task_pipeline.session_factory.create_session", - lambda: _Session(), - ) - response = pipeline._agent_thought_to_stream_response(QueueAgentThoughtEvent(agent_thought_id="thought")) - - assert response is not None - assert response.thought == "" - assert response.observation == "" - assert response.tool == "" - assert response.tool_input == "" - assert response.model_dump(mode="json")["message_files"] == [] - - def test_process_routes_to_stream_and_starts_conversation_name_generation(self): - conversation = _make_conversation(AppMode.CHAT) - message = _make_message() - pipeline = EasyUIBasedGenerateTaskPipeline( - application_generate_entity=_make_entity(ChatAppGenerateEntity, AppMode.CHAT), - queue_manager=_FakeQueueManager(), - conversation=conversation, - message=message, - stream=True, - ) - pipeline._message_cycle_manager.generate_conversation_name = Mock(return_value=object()) - _set_method( - pipeline, - "_wrapper_process_stream_response", - lambda trace_manager: iter([PingStreamResponse(task_id="task")]), - ) - _set_method(pipeline, "_to_stream_response", lambda generator: "streamed") - - result = pipeline.process() - - assert result == "streamed" - pipeline._message_cycle_manager.generate_conversation_name.assert_called_once_with( - conversation_id="conv", query="hello", message_id="msg" - ) - - def test_process_routes_to_blocking_for_completion_mode(self): - conversation = _make_conversation(AppMode.COMPLETION) - message = _make_message() - pipeline = EasyUIBasedGenerateTaskPipeline( - application_generate_entity=_make_entity(CompletionAppGenerateEntity, AppMode.COMPLETION), - queue_manager=_FakeQueueManager(), - conversation=conversation, - message=message, - stream=False, - ) - pipeline._message_cycle_manager.generate_conversation_name = Mock() - _set_method( - pipeline, - "_wrapper_process_stream_response", - lambda trace_manager: iter([PingStreamResponse(task_id="task")]), - ) - _set_method(pipeline, "_to_blocking_response", lambda generator: "blocking") - - result = pipeline.process() - - assert result == "blocking" - pipeline._message_cycle_manager.generate_conversation_name.assert_not_called() - - def test_to_blocking_response_raises_error_stream_exception(self): - conversation = _make_conversation(AppMode.CHAT) - message = _make_message() - pipeline = EasyUIBasedGenerateTaskPipeline( - application_generate_entity=_make_entity(ChatAppGenerateEntity, AppMode.CHAT), - queue_manager=_FakeQueueManager(), - conversation=conversation, - message=message, - stream=False, - ) - - def _gen(): - yield ErrorStreamResponse(task_id="task", err=ValueError("stream error")) - - with pytest.raises(ValueError, match="stream error"): - pipeline._to_blocking_response(_gen()) - - def test_to_blocking_response_raises_when_generator_ends_without_message_end(self): - conversation = _make_conversation(AppMode.CHAT) - message = _make_message() - pipeline = EasyUIBasedGenerateTaskPipeline( - application_generate_entity=_make_entity(ChatAppGenerateEntity, AppMode.CHAT), - queue_manager=_FakeQueueManager(), - conversation=conversation, - message=message, - stream=False, - ) - - def _gen(): - yield PingStreamResponse(task_id="task") - - with pytest.raises(RuntimeError, match="queue listening stopped unexpectedly"): - pipeline._to_blocking_response(_gen()) - - def test_to_stream_response_wraps_completion_stream_events(self): - conversation = _make_conversation(AppMode.COMPLETION) - message = _make_message() - pipeline = EasyUIBasedGenerateTaskPipeline( - application_generate_entity=_make_entity(CompletionAppGenerateEntity, AppMode.COMPLETION), - queue_manager=_FakeQueueManager(), - conversation=conversation, - message=message, - stream=True, - ) - - def _gen(): - yield PingStreamResponse(task_id="task") - - response = list(pipeline._to_stream_response(_gen()))[0] - - assert isinstance(response, CompletionAppStreamResponse) - assert response.message_id == "msg" - - def test_to_stream_response_wraps_chat_stream_events(self): - conversation = _make_conversation(AppMode.CHAT) - message = _make_message() - pipeline = EasyUIBasedGenerateTaskPipeline( - application_generate_entity=_make_entity(ChatAppGenerateEntity, AppMode.CHAT), - queue_manager=_FakeQueueManager(), - conversation=conversation, - message=message, - stream=True, - ) - - def _gen(): - yield PingStreamResponse(task_id="task") - - response = list(pipeline._to_stream_response(_gen()))[0] - - assert isinstance(response, ChatbotAppStreamResponse) - assert response.conversation_id == "conv" - - def test_listen_audio_msg_returns_audio_response_for_non_finish_audio(self): - conversation = _make_conversation(AppMode.CHAT) - message = _make_message() - pipeline = EasyUIBasedGenerateTaskPipeline( - application_generate_entity=_make_entity(ChatAppGenerateEntity, AppMode.CHAT), - queue_manager=_FakeQueueManager(), - conversation=conversation, - message=message, - stream=True, - ) - publisher = cast( - AppGeneratorTTSPublisher, - _AudioPublisher(status="responding", audio="abc"), - ) - - response = pipeline._listen_audio_msg(publisher=publisher, task_id="task") - - assert isinstance(response, MessageAudioStreamResponse) - assert response.audio == "abc" - - def test_listen_audio_msg_returns_none_for_finish_audio(self): - conversation = _make_conversation(AppMode.CHAT) - message = _make_message() - pipeline = EasyUIBasedGenerateTaskPipeline( - application_generate_entity=_make_entity(ChatAppGenerateEntity, AppMode.CHAT), - queue_manager=_FakeQueueManager(), - conversation=conversation, - message=message, - stream=True, - ) - publisher = cast( - AppGeneratorTTSPublisher, - _AudioPublisher(status="finish", audio="abc"), - ) - - assert pipeline._listen_audio_msg(publisher=publisher, task_id="task") is None - - def test_wrapper_process_stream_response_without_tts_publisher(self): - conversation = _make_conversation(AppMode.CHAT) - message = _make_message() - pipeline = EasyUIBasedGenerateTaskPipeline( - application_generate_entity=_make_entity(ChatAppGenerateEntity, AppMode.CHAT), - queue_manager=_FakeQueueManager(), - conversation=conversation, - message=message, - stream=True, - ) - payload = PingStreamResponse(task_id="task") - _set_method(pipeline, "_process_stream_response", lambda publisher, trace_manager: iter([payload])) - - responses = list(pipeline._wrapper_process_stream_response()) - - assert responses == [payload] - - def test_wrapper_process_stream_response_with_tts_publisher(self, monkeypatch: pytest.MonkeyPatch): - conversation = _make_conversation(AppMode.CHAT) - message = _make_message() - entity = _make_entity(ChatAppGenerateEntity, AppMode.CHAT) - entity.app_config.app_model_config_dict = { - "text_to_speech": {"autoPlay": "enabled", "enabled": True, "voice": "v", "language": "en"} - } - pipeline = EasyUIBasedGenerateTaskPipeline( - application_generate_entity=entity, - queue_manager=_FakeQueueManager(), - conversation=conversation, - message=message, - stream=True, - ) - - class _Publisher: - def check_and_get_audio(self, *, block=False): - assert block - return AudioTrunk("finish", "") - - def publish(self, message): - assert message is None - - def cancel(self): - return None - - inline_audio = MessageAudioStreamResponse(task_id="task", audio="inline") - audio_calls = iter([inline_audio, None]) - payload = PingStreamResponse(task_id="task") - _set_method(pipeline, "_listen_audio_msg", lambda publisher, task_id: next(audio_calls)) - _set_method(pipeline, "_process_stream_response", lambda publisher, trace_manager: iter([payload])) - monkeypatch.setattr( - "core.app.task_pipeline.easy_ui_based_generate_task_pipeline.AppGeneratorTTSPublisher", - lambda tenant_id, voice, language, app_type: _Publisher(), - ) - - responses = list(pipeline._wrapper_process_stream_response()) - - assert responses[0] == inline_audio - assert responses[1] == payload - assert isinstance(responses[-1], MessageAudioEndStreamResponse) - assert "audio_type" not in responses[-1].model_dump() - - def test_wrapper_process_stream_response_waits_for_the_audio_terminal(self, monkeypatch: pytest.MonkeyPatch): - conversation = _make_conversation(AppMode.CHAT) - message = _make_message() - entity = _make_entity(ChatAppGenerateEntity, AppMode.CHAT) - entity.app_config.app_model_config_dict = { - "text_to_speech": {"autoPlay": "enabled", "enabled": True, "voice": "v", "language": "en"} - } - pipeline = EasyUIBasedGenerateTaskPipeline( - application_generate_entity=entity, - queue_manager=_FakeQueueManager(), - conversation=conversation, - message=message, - stream=True, - ) - - class _Publisher: - def __init__(self): - self._events = iter([AudioTrunk("responding", "later"), AudioTrunk("finish", "")]) - - def check_and_get_audio(self, *, block=False): - assert block - return next(self._events) - - def publish(self, message): - assert message is None - - def cancel(self): - return None - - _set_method(pipeline, "_process_stream_response", lambda publisher, trace_manager: iter([])) - monkeypatch.setattr( - "core.app.task_pipeline.easy_ui_based_generate_task_pipeline.AppGeneratorTTSPublisher", - lambda tenant_id, voice, language, app_type: _Publisher(), - ) - responses = list(pipeline._wrapper_process_stream_response()) - - assert any(isinstance(item, MessageAudioStreamResponse) for item in responses) - assert isinstance(responses[-1], MessageAudioEndStreamResponse) - - def test_wrapper_process_stream_response_ends_audio_before_reporting_tts_error( - self, monkeypatch: pytest.MonkeyPatch - ): - conversation = _make_conversation(AppMode.CHAT) - message = _make_message() - entity = _make_entity(ChatAppGenerateEntity, AppMode.CHAT) - entity.app_config.app_model_config_dict = { - "text_to_speech": {"autoPlay": "enabled", "enabled": True, "voice": "v", "language": "en"} - } - pipeline = EasyUIBasedGenerateTaskPipeline( - application_generate_entity=entity, - queue_manager=_FakeQueueManager(), - conversation=conversation, - message=message, - stream=True, - ) - error = RuntimeError("tts failed") - - class _Publisher: - def check_and_get_audio(self, *, block=False): - assert block - return AudioTrunk("error", b"", error=error) - - def publish(self, message): - assert message is None - - def cancel(self): - return None - - _set_method(pipeline, "_process_stream_response", lambda publisher, trace_manager: iter([])) - monkeypatch.setattr( - "core.app.task_pipeline.easy_ui_based_generate_task_pipeline.AppGeneratorTTSPublisher", - lambda tenant_id, voice, language, app_type: _Publisher(), - ) - - responses = list(pipeline._wrapper_process_stream_response()) - - assert isinstance(responses[-2], MessageAudioEndStreamResponse) - assert isinstance(responses[-1], ErrorStreamResponse) - assert responses[-1].err is error - - def test_wrapper_process_stream_response_ends_audio_before_a_main_error(self, monkeypatch: pytest.MonkeyPatch): - conversation = _make_conversation(AppMode.CHAT) - message = _make_message() - entity = _make_entity(ChatAppGenerateEntity, AppMode.CHAT) - entity.app_config.app_model_config_dict = { - "text_to_speech": {"autoPlay": "enabled", "enabled": True, "voice": "v", "language": "en"} - } - pipeline = EasyUIBasedGenerateTaskPipeline( - application_generate_entity=entity, - queue_manager=_FakeQueueManager(), - conversation=conversation, - message=message, - stream=True, - ) - publisher = Mock() - publisher.check_and_get_audio.return_value = None - error = ErrorStreamResponse(task_id="task", err=RuntimeError("generation failed")) - _set_method(pipeline, "_process_stream_response", lambda publisher, trace_manager: iter([error])) - monkeypatch.setattr( - "core.app.task_pipeline.easy_ui_based_generate_task_pipeline.AppGeneratorTTSPublisher", - lambda tenant_id, voice, language, app_type: publisher, - ) - - responses = list(pipeline._wrapper_process_stream_response()) - - assert isinstance(responses[-2], MessageAudioEndStreamResponse) - assert responses[-1] is error - assert publisher.cancel.called - publisher.publish.assert_not_called() - - def test_process_stream_response_handles_stop_event_and_output_replacement(self, monkeypatch: pytest.MonkeyPatch): - conversation = _make_conversation(AppMode.CHAT) - message = _make_message() - pipeline = EasyUIBasedGenerateTaskPipeline( - application_generate_entity=_make_entity(ChatAppGenerateEntity, AppMode.CHAT), - queue_manager=_FakeQueueManager(), - conversation=conversation, - message=message, - stream=True, - ) - pipeline._task_state.llm_result.message.content = "raw answer" - _set_queue_events(pipeline, [_queue_message(QueueStopEvent(stopped_by=QueueStopEvent.StopBy.USER_MANUAL))]) - pipeline._handle_stop = Mock() - _set_method(pipeline, "handle_output_moderation_when_task_finished", lambda answer: "moderated answer") - _set_method( - pipeline._message_cycle_manager, - "message_replace_to_stream_response", - lambda answer: MessageReplaceStreamResponse(task_id="task", answer=answer, reason=""), - ) - _set_method(pipeline, "_save_message", lambda **kwargs: None) - _set_method( - pipeline, - "_message_end_to_stream_response", - lambda: MessageEndStreamResponse(task_id="task", id="msg"), - ) - - class _Session: - def __init__[**P](self, *args: P.args, **kwargs: P.kwargs): - pass - - def __enter__(self): - return self - - def __exit__(self, exc_type, exc, tb): - return False - - def commit(self): - return None - - monkeypatch.setattr( - "core.app.task_pipeline.easy_ui_based_generate_task_pipeline.session_factory.create_session", - lambda: _Session(), - ) - - responses = list(pipeline._process_stream_response(publisher=None)) - - assert isinstance(responses[0], MessageReplaceStreamResponse) - assert responses[0].answer == "moderated answer" - assert isinstance(responses[1], MessageEndStreamResponse) - pipeline._handle_stop.assert_called_once() - - def test_process_stream_response_handles_retriever_unknown_and_empty_chunk(self): - conversation = _make_conversation(AppMode.CHAT) - message = _make_message() - pipeline = EasyUIBasedGenerateTaskPipeline( - application_generate_entity=_make_entity(ChatAppGenerateEntity, AppMode.CHAT), - queue_manager=_FakeQueueManager(), - conversation=conversation, - message=message, - stream=True, - ) - retriever_event = QueueRetrieverResourcesEvent(retriever_resources=[]) - chunk = LLMResultChunk( - model="mock", - prompt_messages=[], - delta=LLMResultChunkDelta(index=0, message=AssistantPromptMessage(content=None)), - ) - handled = {"retriever": 0} - - def _handle_retriever_resources(event): - handled["retriever"] += 1 - - pipeline._message_cycle_manager.handle_retriever_resources = _handle_retriever_resources - _set_queue_events( - pipeline, - [ - _queue_message(retriever_event), - _unknown_queue_message(), - _queue_message(QueueLLMChunkEvent(chunk=chunk)), - ], - ) - - responses = list(pipeline._process_stream_response(publisher=None)) - - assert responses == [] - assert handled["retriever"] == 1 - - def test_process_stream_response_skips_when_output_moderation_directs_chunk(self): - conversation = _make_conversation(AppMode.CHAT) - message = _make_message() - pipeline = EasyUIBasedGenerateTaskPipeline( - application_generate_entity=_make_entity(ChatAppGenerateEntity, AppMode.CHAT), - queue_manager=_FakeQueueManager(), - conversation=conversation, - message=message, - stream=True, - ) - chunk = LLMResultChunk( - model="mock", - prompt_messages=[], - delta=LLMResultChunkDelta(index=0, message=AssistantPromptMessage(content="x")), - ) - pipeline._handle_output_moderation_chunk = lambda text: True - _set_queue_events(pipeline, [_queue_message(QueueLLMChunkEvent(chunk=chunk))]) - - responses = list(pipeline._process_stream_response(publisher=None)) - - assert responses == [] - - def test_process_stream_response_ignores_unsupported_chunk_content_types(self): - conversation = _make_conversation(AppMode.CHAT) - message = _make_message() - pipeline = EasyUIBasedGenerateTaskPipeline( - application_generate_entity=_make_entity(ChatAppGenerateEntity, AppMode.CHAT), - queue_manager=_FakeQueueManager(), - conversation=conversation, - message=message, - stream=True, - ) - chunk = LLMResultChunk.model_construct( - prompt_messages=[], - delta=LLMResultChunkDelta.model_construct( - message=AssistantPromptMessage.model_construct(content=[object(), "ok"]) - ), - ) - _set_method(pipeline._message_cycle_manager, "get_message_event_type", lambda message_id: StreamEvent.MESSAGE) - _set_queue_events(pipeline, [_queue_message(QueueLLMChunkEvent.model_construct(chunk=chunk))]) - - responses = list(pipeline._process_stream_response(publisher=None)) - - assert len(responses) == 1 - assert isinstance(responses[0], MessageStreamResponse) - assert responses[0].answer == "ok" - assert pipeline._task_state.llm_result.message.content == "ok" - - def test_process_stream_response_skips_none_chunk_content(self): - conversation = _make_conversation(AppMode.CHAT) - message = _make_message() - pipeline = EasyUIBasedGenerateTaskPipeline( - application_generate_entity=_make_entity(ChatAppGenerateEntity, AppMode.CHAT), - queue_manager=_FakeQueueManager(), - conversation=conversation, - message=message, - stream=True, - ) - chunk = LLMResultChunk( - model="mock", - prompt_messages=[], - delta=LLMResultChunkDelta(index=0, message=AssistantPromptMessage(content=None)), - ) - pipeline._message_cycle_manager.message_to_stream_response = Mock() - _set_queue_events(pipeline, [_queue_message(QueueLLMChunkEvent(chunk=chunk))]) - - responses = list(pipeline._process_stream_response(publisher=None)) - - assert responses == [] - pipeline._message_cycle_manager.message_to_stream_response.assert_not_called() - assert pipeline._task_state.llm_result.message.content == "" - - def test_process_stream_response_reaches_post_loop_branch_with_thread_reference(self): - conversation = _make_conversation(AppMode.CHAT) - message = _make_message() - pipeline = EasyUIBasedGenerateTaskPipeline( - application_generate_entity=_make_entity(ChatAppGenerateEntity, AppMode.CHAT), - queue_manager=_FakeQueueManager(), - conversation=conversation, - message=message, - stream=True, - ) - pipeline._conversation_name_generate_thread = Thread() - _set_queue_events(pipeline, []) - - assert list(pipeline._process_stream_response(publisher=None)) == [] - - def test_save_message_persists_fields_and_emits_trace( - self, monkeypatch: pytest.MonkeyPatch, sqlite_session: Session - ): - conversation = _make_conversation(AppMode.CHAT) - message = _make_message() - application_generate_entity = _make_entity(ChatAppGenerateEntity, AppMode.CHAT) - application_generate_entity.extras = {"trace_session_id": "session-1"} - pipeline = EasyUIBasedGenerateTaskPipeline( - application_generate_entity=application_generate_entity, - queue_manager=_FakeQueueManager(), - conversation=conversation, - message=message, - stream=False, - ) - pipeline.start_at = 10.0 - _set_method(pipeline, "_model_config", _ModelConfigMode(mode="chat")) - pipeline._task_state.llm_result.prompt_messages = [AssistantPromptMessage(content="prompt")] - pipeline._task_state.llm_result.message = AssistantPromptMessage(content=" {{name}} hello ") - pipeline._task_state.llm_result.usage = LLMUsage.from_metadata( - {"prompt_tokens": 3, "completion_tokens": 5, "total_price": "1.23"} - ) - - message_obj = _make_message() - conversation_obj = _make_conversation(AppMode.CHAT) - session = sqlite_session - session.add_all([conversation_obj, message_obj]) - session.flush() - trace_manager_double = _TraceManagerDouble() - trace_manager = cast(TraceQueueManager, trace_manager_double) - sent_payloads: list[tuple[tuple[object, ...], dict[str, object]]] = [] - - monkeypatch.setattr( - "core.app.task_pipeline.easy_ui_based_generate_task_pipeline.PromptMessageUtil.prompt_messages_to_prompt_for_saving", - lambda mode, prompt_messages: "serialized-prompt", - ) - monkeypatch.setattr( - "core.app.task_pipeline.easy_ui_based_generate_task_pipeline.PromptTemplateParser.remove_template_variables", - lambda text: text.replace("{{name}}", "").strip(), - ) - monkeypatch.setattr( - "core.app.task_pipeline.easy_ui_based_generate_task_pipeline.naive_utc_now", - lambda: datetime(2024, 1, 1, tzinfo=UTC), - ) - monkeypatch.setattr( - "core.app.task_pipeline.easy_ui_based_generate_task_pipeline.time.perf_counter", lambda: 15.0 - ) - monkeypatch.setattr( - "core.app.task_pipeline.easy_ui_based_generate_task_pipeline.message_was_created.send", - lambda *args, **kwargs: sent_payloads.append((args, kwargs)), - ) - - pipeline._save_message(session=session, trace_manager=trace_manager) - - assert message_obj.message == "serialized-prompt" - assert message_obj.answer == "hello" - assert message_obj.provider_response_latency == 5.0 - trace_manager_double.add_trace_task.assert_called_once() - trace_task = trace_manager_double.add_trace_task.call_args.args[0] - assert trace_task.trace_type == TraceTaskName.MESSAGE_TRACE - assert trace_task.conversation_id == "conv" - assert trace_task.message_id == "msg" - assert trace_task.kwargs["trace_session_id"] == "session-1" - assert len(sent_payloads) == 1 - - def test_save_stopped_message_preserves_backend_reported_usage( - self, monkeypatch: pytest.MonkeyPatch, sqlite_session: Session - ) -> None: - pipeline = EasyUIBasedGenerateTaskPipeline( - application_generate_entity=_make_entity(ChatAppGenerateEntity, AppMode.CHAT), - queue_manager=_FakeQueueManager(), - conversation=_make_conversation(AppMode.CHAT), - message=_make_message(), - stream=True, - ) - _set_method(pipeline, "_model_config", _ModelConfigMode(mode="chat")) - pipeline._task_state.llm_result.usage = LLMUsage.empty_usage() - - message = _make_message() - message.message_tokens = 13 - message.answer_tokens = 8 - message.total_price = Decimal("0.21") - message.message_metadata = json.dumps({"usage": {"total_tokens": 21}}) - sqlite_session.add_all([_make_conversation(AppMode.CHAT), message]) - sqlite_session.flush() - monkeypatch.setattr( - "core.app.task_pipeline.easy_ui_based_generate_task_pipeline.PromptMessageUtil.prompt_messages_to_prompt_for_saving", - lambda _mode, _prompt_messages: "", - ) - monkeypatch.setattr( - "core.app.task_pipeline.easy_ui_based_generate_task_pipeline.message_was_created.send", - lambda *_args, **_kwargs: None, - ) - - pipeline._save_message(session=sqlite_session, preserve_existing_usage=True) - - assert message.message_tokens == 13 - assert message.answer_tokens == 8 - assert message.total_price == Decimal("0.21") - assert message.message_metadata is not None - assert json.loads(message.message_metadata)["usage"]["total_tokens"] == 21 - - @pytest.mark.parametrize("metadata", ["{", "[]"]) - def test_save_stopped_message_recovers_invalid_existing_metadata( - self, - metadata: str, - monkeypatch: pytest.MonkeyPatch, - sqlite_session: Session, - ) -> None: - pipeline = EasyUIBasedGenerateTaskPipeline( - application_generate_entity=_make_entity(ChatAppGenerateEntity, AppMode.CHAT), - queue_manager=_FakeQueueManager(), - conversation=_make_conversation(AppMode.CHAT), - message=_make_message(), - stream=True, - ) - _set_method(pipeline, "_model_config", _ModelConfigMode(mode="chat")) - pipeline._task_state.llm_result.usage = LLMUsage.empty_usage() - - message = _make_message() - message.message_tokens = 13 - message.answer_tokens = 8 - message.total_price = Decimal("0.21") - message.message_metadata = metadata - sqlite_session.add_all([_make_conversation(AppMode.CHAT), message]) - sqlite_session.flush() - monkeypatch.setattr( - "core.app.task_pipeline.easy_ui_based_generate_task_pipeline.PromptMessageUtil.prompt_messages_to_prompt_for_saving", - lambda _mode, _prompt_messages: "", - ) - monkeypatch.setattr( - "core.app.task_pipeline.easy_ui_based_generate_task_pipeline.message_was_created.send", - lambda *_args, **_kwargs: None, - ) - - pipeline._save_message(session=sqlite_session, preserve_existing_usage=True) - - assert message.message_tokens == 13 - assert message.answer_tokens == 8 - assert message.total_price == Decimal("0.21") - - def test_save_message_raises_when_message_not_found(self, sqlite_session: Session): - conversation = _make_conversation(AppMode.CHAT) - message = _make_message() - pipeline = EasyUIBasedGenerateTaskPipeline( - application_generate_entity=_make_entity(ChatAppGenerateEntity, AppMode.CHAT), - queue_manager=_FakeQueueManager(), - conversation=conversation, - message=message, - stream=False, - ) - session = sqlite_session - - with pytest.raises(ValueError, match="message msg not found"): - pipeline._save_message(session=session) - - def test_save_message_raises_when_conversation_not_found(self, sqlite_session: Session): - conversation = _make_conversation(AppMode.CHAT) - message = _make_message() - pipeline = EasyUIBasedGenerateTaskPipeline( - application_generate_entity=_make_entity(ChatAppGenerateEntity, AppMode.CHAT), - queue_manager=_FakeQueueManager(), - conversation=conversation, - message=message, - stream=False, - ) - session = sqlite_session - session.add(_make_message()) - session.flush() - - with pytest.raises(ValueError, match="Conversation conv not found"): - pipeline._save_message(session=session) - - def test_message_end_to_stream_response_includes_usage_metadata(self, monkeypatch: pytest.MonkeyPatch): - conversation = _make_conversation(AppMode.CHAT) - message = _make_message() - pipeline = EasyUIBasedGenerateTaskPipeline( - application_generate_entity=_make_entity(ChatAppGenerateEntity, AppMode.CHAT), - queue_manager=_FakeQueueManager(), - conversation=conversation, - message=message, - stream=False, - ) - pipeline._task_state.llm_result.usage = LLMUsage.from_metadata({"prompt_tokens": 1, "completion_tokens": 2}) - - class _Result: - def all(self): - return [] - - class _Session: - def __init__[**P](self, *args: P.args, **kwargs: P.kwargs): - pass - - def __enter__(self): - return self - - def __exit__(self, exc_type, exc, tb): - return False - - def scalars(self, *args, **kwargs): - return _Result() - - monkeypatch.setattr( - "core.app.task_pipeline.easy_ui_based_generate_task_pipeline.session_factory.create_session", - lambda: _Session(), - ) - - response = pipeline._message_end_to_stream_response() - - assert response.id == "msg" - usage_metadata = cast(dict[str, object], response.metadata["usage"]) - assert usage_metadata["prompt_tokens"] == 1 - - def test_record_files_returns_empty_list_when_message_has_no_files(self, monkeypatch: pytest.MonkeyPatch): - conversation = _make_conversation(AppMode.CHAT) - message = _make_message() - pipeline = EasyUIBasedGenerateTaskPipeline( - application_generate_entity=_make_entity(ChatAppGenerateEntity, AppMode.CHAT), - queue_manager=_FakeQueueManager(), - conversation=conversation, - message=message, - stream=False, - ) - - class _Result: - def all(self): - return [] - - class _Session: - def __init__[**P](self, *args: P.args, **kwargs: P.kwargs): - pass - - def __enter__(self): - return self - - def __exit__(self, exc_type, exc, tb): - return False - - def scalars(self, *args, **kwargs): - return _Result() - - monkeypatch.setattr( - "core.app.task_pipeline.easy_ui_based_generate_task_pipeline.session_factory.create_session", - lambda: _Session(), - ) - - response = pipeline._message_end_to_stream_response() - - assert response.files == [] - - def test_record_files_handles_local_fallback_and_tool_url_variants(self, monkeypatch: pytest.MonkeyPatch): - conversation = _make_conversation(AppMode.CHAT) - message = _make_message() - pipeline = EasyUIBasedGenerateTaskPipeline( - application_generate_entity=_make_entity(ChatAppGenerateEntity, AppMode.CHAT), - queue_manager=_FakeQueueManager(), - conversation=conversation, - message=message, - stream=False, - ) - message_files = [ - _message_file( - file_id="mf-local-fallback", - transfer_method=FileTransferMethod.LOCAL_FILE, - url="", - upload_file_id="upload-missing", - file_type=FileType.CUSTOM, - ), - _message_file( - file_id="mf-tool-http", - transfer_method=FileTransferMethod.TOOL_FILE, - url="http://cdn.example.com/file.txt?x=1", - upload_file_id=None, - file_type=FileType.CUSTOM, - ), - _message_file( - file_id="mf-tool-noext", - transfer_method=FileTransferMethod.TOOL_FILE, - url="tool/path/toolid", - upload_file_id=None, - file_type=FileType.CUSTOM, - ), - ] - - class _Result: - def __init__(self, items): - self._items = items - - def all(self): - return self._items - - class _Session: - def __init__[**P](self, *args: P.args, **kwargs: P.kwargs): - self.calls = 0 - - def __enter__(self): - return self - - def __exit__(self, exc_type, exc, tb): - return False - - def scalars(self, *args, **kwargs): - self.calls += 1 - return _Result(message_files if self.calls == 1 else []) - - monkeypatch.setattr( - "core.app.task_pipeline.easy_ui_based_generate_task_pipeline.session_factory.create_session", - lambda: _Session(), - ) - monkeypatch.setattr( - "core.app.task_pipeline.message_file_utils.file_helpers.get_signed_file_url", - lambda **kwargs: "local-fallback-signed", - ) - monkeypatch.setattr( - "core.app.task_pipeline.message_file_utils.sign_tool_file", - lambda **kwargs: "tool-signed", - ) - - response = pipeline._message_end_to_stream_response() - files = response.files - - assert files is not None - assert files[0]["url"] == "local-fallback-signed" - assert files[1]["filename"] == "file.txt" - assert files[2]["extension"] == ".bin" - - def test_agent_message_to_stream_response_builds_payload(self): - conversation = _make_conversation(AppMode.CHAT) - message = _make_message() - pipeline = EasyUIBasedGenerateTaskPipeline( - application_generate_entity=_make_entity(ChatAppGenerateEntity, AppMode.CHAT), - queue_manager=_FakeQueueManager(), - conversation=conversation, - message=message, - stream=True, - ) - - response = pipeline._agent_message_to_stream_response(answer="hello", message_id="msg") - - assert response.id == "msg" - assert response.answer == "hello" - - def test_agent_thought_to_stream_response_returns_none_when_not_found(self, monkeypatch: pytest.MonkeyPatch): - conversation = _make_conversation(AppMode.CHAT) - message = _make_message() - pipeline = EasyUIBasedGenerateTaskPipeline( - application_generate_entity=_make_entity(ChatAppGenerateEntity, AppMode.CHAT), - queue_manager=_FakeQueueManager(), - conversation=conversation, - message=message, - stream=True, - ) - - class _Session: - def __init__[**P](self, *args: P.args, **kwargs: P.kwargs): - pass - - def __enter__(self): - return self - - def __exit__(self, exc_type, exc, tb): - return False - - def scalar(self, *args, **kwargs): - return None - - monkeypatch.setattr( - "core.app.task_pipeline.easy_ui_based_generate_task_pipeline.session_factory.create_session", - lambda: _Session(), - ) - - response = pipeline._agent_thought_to_stream_response(QueueAgentThoughtEvent(agent_thought_id="missing")) - - assert response is None - - def test_handle_output_moderation_chunk_appends_token_when_not_directing(self): - conversation = _make_conversation(AppMode.CHAT) - message = _make_message() - pipeline = EasyUIBasedGenerateTaskPipeline( - application_generate_entity=_make_entity(ChatAppGenerateEntity, AppMode.CHAT), - queue_manager=_FakeQueueManager(), - conversation=conversation, - message=message, - stream=True, - ) - appended_tokens: list[str] = [] - - class _Moderation: - def should_direct_output(self): - return False - - def append_new_token(self, text): - appended_tokens.append(text) - - _set_method(pipeline, "output_moderation_handler", _Moderation()) - - result = pipeline._handle_output_moderation_chunk("next-token") - - assert result is False - assert appended_tokens == ["next-token"] diff --git a/api/tests/unit_tests/core/app/task_pipeline/test_easy_ui_message_end_files.py b/api/tests/unit_tests/core/app/task_pipeline/test_easy_ui_message_end_files.py deleted file mode 100644 index 91fd4d77bd7..00000000000 --- a/api/tests/unit_tests/core/app/task_pipeline/test_easy_ui_message_end_files.py +++ /dev/null @@ -1,345 +0,0 @@ -"""Exercise message-end file serialization against persisted SQLite rows. - -The suite covers empty results, all transfer methods, upload metadata batching, -and the fallback used when a local message file references a missing upload. -""" - -import uuid -from datetime import datetime -from unittest.mock import MagicMock, patch - -import pytest -from sqlalchemy.engine import Engine -from sqlalchemy.orm import Session, sessionmaker - -from core.app.app_config.entities import ( - EasyUIBasedAppConfig, - EasyUIBasedAppModelConfigFrom, - ModelConfigEntity, - PromptTemplateEntity, -) -from core.app.apps.message_based_app_queue_manager import MessageBasedAppQueueManager -from core.app.entities.app_invoke_entities import ChatAppGenerateEntity, InvokeFrom -from core.app.entities.task_entities import MessageEndStreamResponse -from core.app.task_pipeline.easy_ui_based_generate_task_pipeline import EasyUIBasedGenerateTaskPipeline -from extensions.ext_redis import RedisClientWrapper -from extensions.storage.storage_type import StorageType -from graphon.file import FileTransferMethod, FileType -from models.enums import ConversationFromSource, CreatorUserRole -from models.model import AppMode, Conversation, Message, MessageFile, UploadFile -from tests.unit_tests.core.model_fixtures import make_model_config - -SQLITE_MODELS = (MessageFile, UploadFile) -pytestmark = [ - pytest.mark.usefixtures("sqlite_session"), - pytest.mark.parametrize("sqlite_session", [SQLITE_MODELS], indirect=True), -] - - -class TestMessageEndStreamResponseFiles: - """Verify message-end file payloads from actual ORM query results.""" - - @pytest.fixture(autouse=True) - def bind_sqlite_engine(self, sqlite_engine: Engine, monkeypatch: pytest.MonkeyPatch) -> None: - """Bind sessions opened by the pipeline to the per-test SQLite engine.""" - - sqlite_session_maker = sessionmaker(bind=sqlite_engine, expire_on_commit=False) - monkeypatch.setattr("core.db.session_factory._session_maker", sqlite_session_maker) - - @pytest.fixture - def pipeline( - self, monkeypatch: pytest.MonkeyPatch, redis_transport: tuple[RedisClientWrapper, MagicMock] - ) -> EasyUIBasedGenerateTaskPipeline: - """Construct the pipeline with validated request, model, and task state.""" - monkeypatch.setattr("core.app.apps.base_app_queue_manager.redis_client", redis_transport[0]) - request = ChatAppGenerateEntity( - task_id=str(uuid.uuid4()), - app_config=EasyUIBasedAppConfig( - tenant_id="tenant-id", - app_id="app-id", - app_mode=AppMode.CHAT, - app_model_config_from=EasyUIBasedAppModelConfigFrom.APP_LATEST_CONFIG, - app_model_config_dict={}, - model=ModelConfigEntity(provider="test-provider", model="test-model"), - prompt_template=PromptTemplateEntity(prompt_type=PromptTemplateEntity.PromptType.SIMPLE), - ), - model_conf=make_model_config(provider="test-provider", model="test-model", mode="chat"), - inputs={}, - files=[], - user_id="user-id", - stream=True, - invoke_from=InvokeFrom.WEB_APP, - ) - conversation = Conversation( - id=str(uuid.uuid4()), - app_id="app-id", - mode=AppMode.CHAT, - name="Conversation", - from_source=ConversationFromSource.API, - _inputs={}, - ) - message = Message(id=str(uuid.uuid4())) - message.created_at = datetime.now() - queue = MessageBasedAppQueueManager( - task_id=request.task_id, - user_id=request.user_id, - invoke_from=request.invoke_from, - conversation_id=conversation.id, - app_mode=AppMode.CHAT, - message_id=message.id, - ) - return EasyUIBasedGenerateTaskPipeline( - application_generate_entity=request, - queue_manager=queue, - conversation=conversation, - message=message, - stream=True, - ) - - @staticmethod - def _message_file( - *, - transfer_method: FileTransferMethod, - url: str | None = None, - upload_file_id: str | None = None, - ) -> MessageFile: - return MessageFile( - message_id=str(uuid.uuid4()), - type=FileType.IMAGE, - transfer_method=transfer_method, - created_by_role=CreatorUserRole.ACCOUNT, - created_by=str(uuid.uuid4()), - url=url, - upload_file_id=upload_file_id, - ) - - @pytest.fixture - def message_file_local(self) -> MessageFile: - """Create an unpersisted local-file row.""" - - return self._message_file( - transfer_method=FileTransferMethod.LOCAL_FILE, - upload_file_id=str(uuid.uuid4()), - ) - - @pytest.fixture - def message_file_remote(self) -> MessageFile: - """Create an unpersisted remote-file row.""" - - return self._message_file( - transfer_method=FileTransferMethod.REMOTE_URL, - url="https://example.com/image.jpg", - ) - - @pytest.fixture - def message_file_tool(self) -> MessageFile: - """Create an unpersisted tool-file row.""" - - return self._message_file( - transfer_method=FileTransferMethod.TOOL_FILE, - url="tool_file_123.png", - ) - - @pytest.fixture - def upload_file(self, message_file_local: MessageFile) -> UploadFile: - """Create upload metadata matching the local message-file reference.""" - - upload = UploadFile( - tenant_id=str(uuid.uuid4()), - storage_type=StorageType.LOCAL, - key="uploads/test_image.png", - name="test_image.png", - size=1024, - extension="png", - mime_type="image/png", - created_by_role=CreatorUserRole.ACCOUNT, - created_by=str(uuid.uuid4()), - created_at=datetime.now(), - used=True, - ) - upload.id = message_file_local.upload_file_id or str(uuid.uuid4()) - return upload - - @staticmethod - def _persist(session: Session, *rows: MessageFile | UploadFile) -> None: - session.add_all(rows) - session.commit() - - def test_message_end_with_no_files( - self, sqlite_session: Session, pipeline: EasyUIBasedGenerateTaskPipeline - ) -> None: - """Rows for another message do not leak into an empty files array.""" - - unrelated_file = self._message_file( - transfer_method=FileTransferMethod.REMOTE_URL, - url="https://example.com/unrelated.png", - ) - self._persist(sqlite_session, unrelated_file) - - result = pipeline._message_end_to_stream_response() - - assert sqlite_session.get(MessageFile, unrelated_file.id) is unrelated_file - assert isinstance(result, MessageEndStreamResponse) - assert result.files == [] - assert result.id == pipeline._message_id - assert result.metadata == { - "usage": pipeline._task_state.llm_result.usage.model_dump(exclude_none=True), - "retriever_resources": [], - "reasoning": {}, - } - - def test_message_end_with_local_file( - self, - sqlite_session: Session, - pipeline: EasyUIBasedGenerateTaskPipeline, - message_file_local: MessageFile, - upload_file: UploadFile, - ) -> None: - """Local files include persisted upload metadata and a signed URL.""" - - message_file_local.message_id = pipeline._message_id - self._persist(sqlite_session, message_file_local, upload_file) - - with patch( - "core.app.task_pipeline.message_file_utils.file_helpers.get_signed_file_url", - return_value="https://example.com/signed-url?signature=abc123", - ) as get_signed_url: - result = pipeline._message_end_to_stream_response() - - assert result.files is not None - assert len(result.files) == 1 - file_dict = result.files[0] - assert file_dict["related_id"] == message_file_local.id - assert file_dict["filename"] == "test_image.png" - assert file_dict["mime_type"] == "image/png" - assert file_dict["size"] == 1024 - assert file_dict["extension"] == ".png" - assert file_dict["type"] == "image" - assert file_dict["transfer_method"] == FileTransferMethod.LOCAL_FILE.value - assert file_dict["url"].startswith("https://example.com/signed-url") - assert file_dict["upload_file_id"] == message_file_local.upload_file_id - assert file_dict["remote_url"] == "" - get_signed_url.assert_called_once_with(upload_file_id=upload_file.id) - - def test_message_end_with_remote_url( - self, sqlite_session: Session, pipeline: EasyUIBasedGenerateTaskPipeline, message_file_remote: MessageFile - ) -> None: - """Remote files retain their source URL and derived filename.""" - - message_file_remote.message_id = pipeline._message_id - self._persist(sqlite_session, message_file_remote) - - result = pipeline._message_end_to_stream_response() - - assert result.files is not None - assert len(result.files) == 1 - file_dict = result.files[0] - assert file_dict["related_id"] == message_file_remote.id - assert file_dict["filename"] == "image.jpg" - assert file_dict["url"] == "https://example.com/image.jpg" - assert file_dict["extension"] == ".jpg" - assert file_dict["type"] == "image" - assert file_dict["transfer_method"] == FileTransferMethod.REMOTE_URL.value - assert file_dict["remote_url"] == "https://example.com/image.jpg" - assert file_dict["upload_file_id"] == message_file_remote.id - - def test_message_end_with_tool_file_http( - self, sqlite_session: Session, pipeline: EasyUIBasedGenerateTaskPipeline, message_file_tool: MessageFile - ) -> None: - """HTTP tool-file URLs pass through unchanged.""" - - message_file_tool.message_id = pipeline._message_id - message_file_tool.url = "https://example.com/tool_file.png" - self._persist(sqlite_session, message_file_tool) - - result = pipeline._message_end_to_stream_response() - - assert result.files is not None - file_dict = result.files[0] - assert file_dict["url"] == "https://example.com/tool_file.png" - assert file_dict["filename"] == "tool_file.png" - assert file_dict["extension"] == ".png" - assert file_dict["transfer_method"] == FileTransferMethod.TOOL_FILE.value - - def test_message_end_with_tool_file_local( - self, sqlite_session: Session, pipeline: EasyUIBasedGenerateTaskPipeline, message_file_tool: MessageFile - ) -> None: - """Local tool-file identifiers are signed at the external boundary.""" - - message_file_tool.message_id = pipeline._message_id - self._persist(sqlite_session, message_file_tool) - - with patch( - "core.app.task_pipeline.message_file_utils.sign_tool_file", - return_value="https://example.com/signed-tool-file.png?signature=xyz", - ) as sign_tool: - result = pipeline._message_end_to_stream_response() - - assert result.files is not None - file_dict = result.files[0] - assert file_dict["url"].startswith("https://example.com/signed-tool-file.png") - assert file_dict["filename"] == "tool_file_123.png" - assert file_dict["extension"] == ".png" - assert file_dict["transfer_method"] == FileTransferMethod.TOOL_FILE.value - sign_tool.assert_called_once_with(tool_file_id="tool_file_123", extension=".png") - - def test_message_end_with_tool_file_long_extension( - self, sqlite_session: Session, pipeline: EasyUIBasedGenerateTaskPipeline, message_file_tool: MessageFile - ) -> None: - """Overlong tool-file extensions use the safe binary fallback.""" - - message_file_tool.message_id = pipeline._message_id - message_file_tool.url = "tool_file_abc.verylongextension" - self._persist(sqlite_session, message_file_tool) - - with patch( - "core.app.task_pipeline.message_file_utils.sign_tool_file", - return_value="https://example.com/signed.bin", - ) as sign_tool: - result = pipeline._message_end_to_stream_response() - - assert result.files is not None - assert result.files[0]["extension"] == ".bin" - sign_tool.assert_called_once_with(tool_file_id="tool_file_abc", extension=".bin") - - def test_message_end_with_multiple_files( - self, - sqlite_session: Session, - pipeline: EasyUIBasedGenerateTaskPipeline, - message_file_local: MessageFile, - message_file_remote: MessageFile, - upload_file: UploadFile, - ) -> None: - """The response contains every persisted file associated with the message.""" - - message_file_local.message_id = pipeline._message_id - message_file_remote.message_id = pipeline._message_id - self._persist(sqlite_session, message_file_local, message_file_remote, upload_file) - - with patch( - "core.app.task_pipeline.message_file_utils.file_helpers.get_signed_file_url", - return_value="https://example.com/signed-url?signature=abc123", - ): - result = pipeline._message_end_to_stream_response() - - assert result.files is not None - assert {file["related_id"] for file in result.files} == {message_file_local.id, message_file_remote.id} - - def test_message_end_with_local_file_no_upload_file( - self, sqlite_session: Session, pipeline: EasyUIBasedGenerateTaskPipeline, message_file_local: MessageFile - ) -> None: - """A missing upload row still signs the stored upload identifier.""" - - message_file_local.message_id = pipeline._message_id - self._persist(sqlite_session, message_file_local) - - with patch( - "core.app.task_pipeline.message_file_utils.file_helpers.get_signed_file_url", - return_value="https://example.com/fallback-url?signature=def456", - ) as get_signed_url: - result = pipeline._message_end_to_stream_response() - - assert result.files is not None - assert len(result.files) == 1 - assert result.files[0]["url"].startswith("https://example.com/fallback-url") - get_signed_url.assert_called_once_with(upload_file_id=str(message_file_local.upload_file_id)) diff --git a/api/tests/unit_tests/core/app/task_pipeline/test_message_cycle_manager_optimization.py b/api/tests/unit_tests/core/app/task_pipeline/test_message_cycle_manager.py similarity index 99% rename from api/tests/unit_tests/core/app/task_pipeline/test_message_cycle_manager_optimization.py rename to api/tests/unit_tests/core/app/task_pipeline/test_message_cycle_manager.py index b2c268d99db..f50355ace26 100644 --- a/api/tests/unit_tests/core/app/task_pipeline/test_message_cycle_manager_optimization.py +++ b/api/tests/unit_tests/core/app/task_pipeline/test_message_cycle_manager.py @@ -23,7 +23,7 @@ from models.account import Account from models.base import TypeBase from models.enums import ConversationFromSource, CreatorUserRole, MessageFileBelongsTo from models.model import App, Conversation, MessageAnnotation, MessageFile -from tests.unit_tests.model_factories import make_app, make_conversation +from tests.helpers.unit.model_factories import make_app, make_conversation @dataclass(frozen=True) diff --git a/api/tests/unit_tests/core/app/test_easy_ui_model_config_manager.py b/api/tests/unit_tests/core/app/test_easy_ui_model_config_manager.py deleted file mode 100644 index 21c761c579b..00000000000 --- a/api/tests/unit_tests/core/app/test_easy_ui_model_config_manager.py +++ /dev/null @@ -1,57 +0,0 @@ -from types import SimpleNamespace -from unittest.mock import patch - -from core.app.app_config.easy_ui_based_app.model_config.manager import ModelConfigManager -from core.app.app_config.entities import ModelConfigEntity -from graphon.model_runtime.entities.model_entities import ModelPropertyKey -from models.provider_ids import ModelProviderID - - -def test_validate_and_set_defaults_reuses_single_model_assembly(): - provider_name = str(ModelProviderID("openai")) - provider_entity = SimpleNamespace(provider=provider_name) - model = SimpleNamespace(model="gpt-4o-mini", model_properties={ModelPropertyKey.MODE: "chat"}) - provider_configurations = SimpleNamespace(get_models=lambda **kwargs: [model]) - assembly = SimpleNamespace( - model_provider_factory=SimpleNamespace(get_providers=lambda: [provider_entity]), - provider_manager=SimpleNamespace(get_configurations=lambda tenant_id: provider_configurations), - ) - config = { - "model": { - "provider": "openai", - "name": "gpt-4o-mini", - "completion_params": {"stop": []}, - } - } - - with patch( - "core.app.app_config.easy_ui_based_app.model_config.manager.create_plugin_model_assembly", - return_value=assembly, - ) as mock_assembly: - result, keys = ModelConfigManager.validate_and_set_defaults("tenant-1", config) - - assert result["model"]["provider"] == provider_name - assert result["model"]["mode"] == "chat" - assert keys == ["model"] - mock_assembly.assert_called_once_with(tenant_id="tenant-1") - - -def test_convert_keeps_model_config_shape(): - config = { - "model": { - "provider": "openai", - "name": "gpt-4o-mini", - "mode": "chat", - "completion_params": {"temperature": 0.3, "stop": ["END"]}, - } - } - - result = ModelConfigManager.convert(config) - - assert result == ModelConfigEntity( - provider="openai", - model="gpt-4o-mini", - mode="chat", - parameters={"temperature": 0.3}, - stop=["END"], - ) diff --git a/api/tests/unit_tests/core/app/test_invoke_from.py b/api/tests/unit_tests/core/app/test_invoke_from.py deleted file mode 100644 index 33a4d2edf11..00000000000 --- a/api/tests/unit_tests/core/app/test_invoke_from.py +++ /dev/null @@ -1,25 +0,0 @@ -from core.app.entities.app_invoke_entities import InvokeFrom - - -def test_openapi_variant_present(): - assert InvokeFrom.OPENAPI.value == "openapi" - - -def test_openapi_distinct_from_service_api(): - assert InvokeFrom.OPENAPI != InvokeFrom.SERVICE_API - - -def test_runs_as_account_only_for_console_contexts(): - # Console contexts (studio debugger / explore) run as the signed-in Account. - assert InvokeFrom.DEBUGGER.runs_as_account() is True - assert InvokeFrom.EXPLORE.runs_as_account() is True - # Everything else is attributed to an end user. - for invoke_from in ( - InvokeFrom.WEB_APP, - InvokeFrom.SERVICE_API, - InvokeFrom.OPENAPI, - InvokeFrom.TRIGGER, - InvokeFrom.PUBLISHED_PIPELINE, - InvokeFrom.VALIDATION, - ): - assert invoke_from.runs_as_account() is False diff --git a/api/tests/unit_tests/core/workflow/graph_engine/layers/conftest.py b/api/tests/unit_tests/core/app/workflow/layers/conftest.py similarity index 100% rename from api/tests/unit_tests/core/workflow/graph_engine/layers/conftest.py rename to api/tests/unit_tests/core/app/workflow/layers/conftest.py diff --git a/api/tests/unit_tests/core/app/workflow/layers/test_observability.py b/api/tests/unit_tests/core/app/workflow/layers/test_observability.py new file mode 100644 index 00000000000..55d75486f2c --- /dev/null +++ b/api/tests/unit_tests/core/app/workflow/layers/test_observability.py @@ -0,0 +1,499 @@ +from __future__ import annotations + +from concurrent.futures import ThreadPoolExecutor +from threading import Event +from types import SimpleNamespace + +import pytest +from opentelemetry.sdk.trace import ReadableSpan, SpanProcessor +from opentelemetry.trace import StatusCode, get_current_span + +from core.app.workflow.layers.observability import ObservabilityLayer +from core.workflow.workflow_entry import WorkflowEntry +from extensions.otel.semconv import DifySpanAttributes +from graphon.engine_events import GraphRunAbortedEvent, GraphRunSucceededEvent +from graphon.graph import Graph +from graphon.nodes.start import StartNode +from graphon.nodes.start.entities import StartNodeData +from graphon.runtime import RuntimeState, VariablePool +from tests.helpers.unit.config_override import apply_config_overrides +from tests.workflow_test_utils import build_test_graph_init_params + + +class TestObservabilityLayerExtraCases: + class TestObservabilityLayerExtras: + def test_init_tracer_disables_when_get_tracer_fails( + self, monkeypatch: pytest.MonkeyPatch, caplog: pytest.LogCaptureFixture + ): + apply_config_overrides(monkeypatch, ENABLE_OTEL=True) + monkeypatch.setattr("core.app.workflow.layers.observability.is_instrument_flag_enabled", lambda: False) + + def _raise(*_args, **_kwargs): + raise RuntimeError("tracer init failed") + + monkeypatch.setattr("core.app.workflow.layers.observability.get_tracer", _raise) + + layer = ObservabilityLayer() + + assert layer._is_disabled is True + assert layer._tracer is None + assert "Failed to get OpenTelemetry tracer" in caplog.text + + def test_node_run_context_skips_without_execution_id(self): + layer = ObservabilityLayer() + layer._is_disabled = False + calls: list[str] = [] + layer._tracer = SimpleNamespace(start_span=lambda *_args, **_kwargs: calls.append("called")) + + with layer.node_run_context(SimpleNamespace(execution_id=None, title="node", id="node")): + assert calls == [] + + assert calls == [] + + def test_node_runs_when_span_creation_fails(self, caplog: pytest.LogCaptureFixture): + layer = ObservabilityLayer() + layer._is_disabled = False + + def _raise(*_args, **_kwargs): + raise RuntimeError("start failed") + + layer._tracer = SimpleNamespace(start_span=_raise) + + node_executed = False + with layer.node_run_context(SimpleNamespace(execution_id="exec", title="node", id="node")): + node_executed = True + + assert node_executed + assert "start failed" in caplog.text + + +""" +Tests for ObservabilityLayer. + +Test coverage: +- Initialization and enable/disable logic +- Node span lifecycle (start, end, error handling) +- Parser integration (default, tool, LLM, and retrieval parsers) +- Result event parameter extraction (inputs/outputs) +- Graph lifecycle management +- Disabled mode behavior +""" + + +class TestObservabilityCases: + @staticmethod + @pytest.fixture(autouse=True) + def _otel_config(config_overrides) -> None: + config_overrides(ENABLE_OTEL=True) + + @staticmethod + @pytest.mark.parametrize("tracing_enabled", [False, True]) + def test_unbound_node_still_runs_and_preserves_body_errors(config_overrides, tracing_enabled): + config_overrides(ENABLE_OTEL=tracing_enabled) + node = StartNode( + node_id="start", + data=StartNodeData(title="Start", variables=[]), + init_params=build_test_graph_init_params(), + runtime_state=RuntimeState(workflow_id="workflow", variable_pool=VariablePool(), start_at=0), + ) + layer = ObservabilityLayer() + parent = get_current_span() + executed = [] + + with layer.node_run_context(node): + executed.append("completed") + + with pytest.raises(ValueError, match="node failed"), layer.node_run_context(node): + raise ValueError("node failed") + + assert executed == ["completed"] + assert get_current_span() is parent + + @staticmethod + @pytest.mark.usefixtures("mock_is_instrument_flag_enabled_true", "tracer_provider_with_memory_exporter") + def test_instrument_flag_enables_workflow_node_spans(memory_span_exporter, config_overrides, monkeypatch): + config_overrides(ENABLE_OTEL=False, GRAPH_ENGINE_MAX_WORKERS=1) + monkeypatch.setattr("core.workflow.workflow_entry.is_instrument_flag_enabled", lambda: True) + state = RuntimeState(workflow_id="workflow", variable_pool=VariablePool(), start_at=0) + node = StartNode( + node_id="start", + data=StartNodeData(title="Start", variables=[]), + init_params=build_test_graph_init_params(), + runtime_state=state, + ) + entry = WorkflowEntry( + graph_config={}, + graph=Graph.new().add_root(node).build(), + call_depth=0, + graph_runtime_state=state, + ) + + events = list(entry.run()) + + assert isinstance(events[-1], GraphRunSucceededEvent) + spans = memory_span_exporter.get_finished_spans() + assert [span.name for span in spans] == [node.title] + + class TestObservabilityLayerNodeSpanLifecycle: + """Test node span creation and lifecycle management.""" + + @pytest.mark.usefixtures("mock_is_instrument_flag_enabled_false") + def test_node_span_created_and_ended( + self, tracer_provider_with_memory_exporter, memory_span_exporter, mock_llm_node + ): + """Test that span is created on node start and ended on node end.""" + layer = ObservabilityLayer() + layer.on_graph_start() + + with layer.node_run_context(mock_llm_node): + layer.on_node_run_end(mock_llm_node, None) + + spans = memory_span_exporter.get_finished_spans() + assert len(spans) == 1 + assert spans[0].name == mock_llm_node.title + assert spans[0].status.status_code == StatusCode.OK + + @pytest.mark.usefixtures("mock_is_instrument_flag_enabled_false") + def test_node_error_recorded_in_span( + self, tracer_provider_with_memory_exporter, memory_span_exporter, mock_llm_node + ): + """Test that node execution errors are recorded in span.""" + layer = ObservabilityLayer() + layer.on_graph_start() + + error = ValueError("Test error") + with layer.node_run_context(mock_llm_node): + layer.on_node_run_end(mock_llm_node, error) + + spans = memory_span_exporter.get_finished_spans() + assert len(spans) == 1 + assert spans[0].status.status_code == StatusCode.ERROR + assert len(spans[0].events) > 0 + assert any("exception" in event.name.lower() for event in spans[0].events) + + @pytest.mark.usefixtures("mock_is_instrument_flag_enabled_false") + def test_node_end_without_start_handled_gracefully( + self, tracer_provider_with_memory_exporter, memory_span_exporter, mock_llm_node + ): + """Test that ending a node without start doesn't crash.""" + layer = ObservabilityLayer() + layer.on_graph_start() + + layer.on_node_run_end(mock_llm_node, None) + + spans = memory_span_exporter.get_finished_spans() + assert len(spans) == 0 + + @pytest.mark.usefixtures("mock_is_instrument_flag_enabled_false") + def test_suspension_deactivates_span_until_resume( + self, tracer_provider_with_memory_exporter, memory_span_exporter, mock_tool_node + ): + layer = ObservabilityLayer() + parent = get_current_span() + + with layer.node_run_context(mock_tool_node): + tool_span = get_current_span() + assert tool_span.is_recording() + + assert get_current_span() is parent + assert memory_span_exporter.get_finished_spans() == () + + with layer.node_run_context(mock_tool_node): + assert get_current_span() is tool_span + layer.on_node_run_end(mock_tool_node, None) + + assert get_current_span() is parent + spans = memory_span_exporter.get_finished_spans() + assert len(spans) == 1 + assert spans[0].context == tool_span.get_span_context() + + @pytest.mark.usefixtures("mock_is_instrument_flag_enabled_false") + def test_activation_does_not_swallow_node_errors( + self, tracer_provider_with_memory_exporter, memory_span_exporter, mock_llm_node + ): + layer = ObservabilityLayer() + parent = get_current_span() + error = ValueError("node failed") + + with pytest.raises(ValueError, match="node failed"): + with layer.node_run_context(mock_llm_node): + raise error + + assert get_current_span() is parent + layer.on_node_run_end(mock_llm_node, error) + spans = memory_span_exporter.get_finished_spans() + assert len(spans) == 1 + assert spans[0].status.status_code == StatusCode.ERROR + + @pytest.mark.usefixtures("mock_is_instrument_flag_enabled_false") + def test_parser_failure_still_ends_span( + self, tracer_provider_with_memory_exporter, memory_span_exporter, mock_start_node, monkeypatch, caplog + ): + layer = ObservabilityLayer() + + def fail_parse(**_kwargs): + raise ValueError("invalid node attributes") + + monkeypatch.setattr(layer._default_parser, "parse", fail_parse) + + with layer.node_run_context(mock_start_node): + layer.on_node_run_end(mock_start_node, None) + + assert len(memory_span_exporter.get_finished_spans()) == 1 + assert "invalid node attributes" in caplog.text + + class TestObservabilityLayerParserIntegration: + """Test parser integration for different node types.""" + + @pytest.mark.usefixtures("mock_is_instrument_flag_enabled_false") + def test_default_parser_used_for_regular_node( + self, tracer_provider_with_memory_exporter, memory_span_exporter, mock_start_node + ): + """Test that default parser is used for non-tool nodes.""" + layer = ObservabilityLayer() + layer.on_graph_start() + + with layer.node_run_context(mock_start_node): + layer.on_node_run_end(mock_start_node, None) + + spans = memory_span_exporter.get_finished_spans() + assert len(spans) == 1 + attrs = spans[0].attributes + assert attrs["node.id"] == mock_start_node.id + assert attrs["node.execution_id"] == mock_start_node.execution_id + assert attrs["node.type"] == mock_start_node.node_type + + @pytest.mark.usefixtures("mock_is_instrument_flag_enabled_false") + def test_tool_parser_used_for_tool_node( + self, tracer_provider_with_memory_exporter, memory_span_exporter, mock_tool_node + ): + """Test that tool parser is used for tool nodes.""" + layer = ObservabilityLayer() + layer.on_graph_start() + + with layer.node_run_context(mock_tool_node): + layer.on_node_run_end(mock_tool_node, None) + + spans = memory_span_exporter.get_finished_spans() + assert len(spans) == 1 + attrs = spans[0].attributes + assert attrs["node.id"] == mock_tool_node.id + assert attrs["gen_ai.tool.name"] == mock_tool_node.title + assert attrs["gen_ai.tool.type"] == mock_tool_node._node_data.provider_type.value + + @pytest.mark.usefixtures("mock_is_instrument_flag_enabled_false") + def test_llm_parser_used_for_llm_node( + self, tracer_provider_with_memory_exporter, memory_span_exporter, mock_llm_node, mock_result_event + ): + """Test that LLM parser is used for LLM nodes and extracts LLM-specific attributes.""" + from graphon.node_events import NodeRunResult + + mock_result_event.node_run_result = NodeRunResult( + inputs={}, + outputs={"text": "test completion", "finish_reason": "stop"}, + process_data={ + "model_name": "gpt-4", + "model_provider": "openai", + "usage": {"prompt_tokens": 10, "completion_tokens": 20, "total_tokens": 30}, + "prompts": [{"role": "user", "text": "test prompt"}], + }, + metadata={}, + ) + + layer = ObservabilityLayer() + layer.on_graph_start() + + with layer.node_run_context(mock_llm_node): + layer.on_node_run_end(mock_llm_node, None, mock_result_event) + + spans = memory_span_exporter.get_finished_spans() + assert len(spans) == 1 + attrs = spans[0].attributes + assert attrs["node.id"] == mock_llm_node.id + assert attrs["gen_ai.request.model"] == "gpt-4" + assert attrs["gen_ai.provider.name"] == "openai" + assert attrs["gen_ai.usage.input_tokens"] == 10 + assert attrs["gen_ai.usage.output_tokens"] == 20 + assert attrs["gen_ai.usage.total_tokens"] == 30 + assert attrs["gen_ai.completion"] == "test completion" + assert attrs["gen_ai.response.finish_reason"] == "stop" + + @pytest.mark.usefixtures("mock_is_instrument_flag_enabled_false") + def test_retrieval_parser_used_for_retrieval_node( + self, tracer_provider_with_memory_exporter, memory_span_exporter, mock_retrieval_node, mock_result_event + ): + """Test that retrieval parser is used for retrieval nodes and extracts retrieval-specific attributes.""" + from graphon.node_events import NodeRunResult + + mock_result_event.node_run_result = NodeRunResult( + inputs={"query": "test query"}, + outputs={"result": [{"content": "test content", "metadata": {"score": 0.9, "document_id": "doc1"}}]}, + process_data={}, + metadata={}, + ) + + layer = ObservabilityLayer() + layer.on_graph_start() + + with layer.node_run_context(mock_retrieval_node): + layer.on_node_run_end(mock_retrieval_node, None, mock_result_event) + + spans = memory_span_exporter.get_finished_spans() + assert len(spans) == 1 + attrs = spans[0].attributes + assert attrs["node.id"] == mock_retrieval_node.id + assert attrs["retrieval.query"] == "test query" + assert "retrieval.document" in attrs + + @pytest.mark.usefixtures("mock_is_instrument_flag_enabled_false") + def test_result_event_extracts_inputs_and_outputs( + self, tracer_provider_with_memory_exporter, memory_span_exporter, mock_start_node, mock_result_event + ): + """Test that result_event parameter allows parsers to extract inputs and outputs.""" + from graphon.node_events import NodeRunResult + + mock_result_event.node_run_result = NodeRunResult( + inputs={"input_key": "input_value"}, + outputs={"output_key": "output_value"}, + process_data={}, + metadata={}, + ) + + layer = ObservabilityLayer() + layer.on_graph_start() + + with layer.node_run_context(mock_start_node): + layer.on_node_run_end(mock_start_node, None, mock_result_event) + + spans = memory_span_exporter.get_finished_spans() + assert len(spans) == 1 + attrs = spans[0].attributes + assert "input.value" in attrs + assert "output.value" in attrs + + class TestObservabilityLayerGraphLifecycle: + """Test graph lifecycle management.""" + + @pytest.mark.usefixtures("mock_is_instrument_flag_enabled_false") + def test_on_graph_end_with_no_unfinished_spans( + self, tracer_provider_with_memory_exporter, memory_span_exporter, mock_llm_node + ): + """Test that on_graph_end handles normal completion.""" + layer = ObservabilityLayer() + layer.on_graph_start() + + with layer.node_run_context(mock_llm_node): + layer.on_node_run_end(mock_llm_node, None) + layer.on_graph_end(None) + + spans = memory_span_exporter.get_finished_spans() + assert len(spans) == 1 + + @pytest.mark.usefixtures("mock_is_instrument_flag_enabled_false") + def test_on_graph_end_ends_suspended_spans( + self, tracer_provider_with_memory_exporter, memory_span_exporter, mock_llm_node + ): + """Pause and abort cleanup must export spans that never reached node end.""" + layer = ObservabilityLayer() + layer.on_graph_start() + + with layer.node_run_context(mock_llm_node): + span = get_current_span() + + layer.on_graph_end(None) + layer.on_graph_end(None) + + assert not span.is_recording() + spans = memory_span_exporter.get_finished_spans() + assert len(spans) == 1 + assert spans[0].name == mock_llm_node.title + + @pytest.mark.usefixtures("mock_is_instrument_flag_enabled_false") + def test_graph_end_and_late_worker_each_end_spans_once( + self, tracer_provider_with_memory_exporter, memory_span_exporter, mock_start_node, mock_llm_node, caplog + ): + """An aborted graph can finish while a worker still owns another span.""" + graph_ending = Event() + worker_finished = Event() + + class EndBarrier(SpanProcessor): + def on_end(self, span: ReadableSpan) -> None: + if span.name == mock_start_node.title: + graph_ending.set() + assert worker_finished.wait(5) + + tracer_provider_with_memory_exporter.add_span_processor(EndBarrier()) + layer = ObservabilityLayer() + for node in (mock_start_node, mock_llm_node): + with layer.node_run_context(node): + pass + + def finish_worker() -> None: + assert graph_ending.wait(5) + try: + layer.on_node_run_end(mock_llm_node, None) + finally: + worker_finished.set() + + with ThreadPoolExecutor(max_workers=1) as executor: + worker = executor.submit(finish_worker) + try: + layer.on_graph_end(None) + finally: + worker.result(timeout=5) + + layer.on_graph_end(None) + spans = memory_span_exporter.get_finished_spans() + assert len(spans) == 2 + assert {span.name for span in spans} == {mock_start_node.title, mock_llm_node.title} + assert "Calling end() on an ended span" not in caplog.text + + @pytest.mark.usefixtures("mock_is_instrument_flag_enabled_false") + def test_graph_aborted_event_records_reason_on_current_span( + self, tracer_provider_with_memory_exporter, memory_span_exporter, mock_start_node + ): + layer = ObservabilityLayer() + layer.on_graph_start() + with layer.node_run_context(mock_start_node): + layer.on_event(GraphRunAbortedEvent(reason="worker shutdown", outputs={})) + layer.on_node_run_end(mock_start_node, None) + + spans = memory_span_exporter.get_finished_spans() + assert len(spans) == 1 + assert spans[0].attributes[DifySpanAttributes.WORKFLOW_ABORT_REASON] == "worker shutdown" + assert any( + event.name == "dify.workflow.aborted" + and event.attributes[DifySpanAttributes.WORKFLOW_ABORT_REASON] == "worker shutdown" + for event in spans[0].events + ) + + class TestObservabilityLayerDisabledMode: + """Test behavior when layer is disabled.""" + + @pytest.mark.usefixtures("mock_is_instrument_flag_enabled_false") + def test_disabled_mode_skips_node_start( + self, tracer_provider_with_memory_exporter, memory_span_exporter, mock_start_node, config_overrides + ): + """Test that disabled layer doesn't create spans on node start.""" + config_overrides(ENABLE_OTEL=False) + layer = ObservabilityLayer() + + layer.on_graph_start() + with layer.node_run_context(mock_start_node): + layer.on_node_run_end(mock_start_node, None) + + spans = memory_span_exporter.get_finished_spans() + assert len(spans) == 0 + + @pytest.mark.usefixtures("mock_is_instrument_flag_enabled_false") + def test_disabled_mode_skips_node_end( + self, tracer_provider_with_memory_exporter, memory_span_exporter, mock_llm_node, config_overrides + ): + """Test that disabled layer doesn't process node end.""" + config_overrides(ENABLE_OTEL=False) + layer = ObservabilityLayer() + + layer.on_node_run_end(mock_llm_node, None) + + spans = memory_span_exporter.get_finished_spans() + assert len(spans) == 0 diff --git a/api/tests/unit_tests/core/app/workflow/layers/test_persistence_inspector_publish.py b/api/tests/unit_tests/core/app/workflow/layers/test_persistence_inspector_publish.py deleted file mode 100644 index 31ff06ff71d..00000000000 --- a/api/tests/unit_tests/core/app/workflow/layers/test_persistence_inspector_publish.py +++ /dev/null @@ -1,186 +0,0 @@ -"""Verify the workflow persistence layer fans Inspector deltas to redis pub/sub. - -The hook lives in ``services/workflow_persistence.py``: -every ``_handle_node_*`` and the terminal ``_handle_graph_run_*`` handlers -call into ``services.workflow.inspector_events.publish_node_changed`` / -``publish_workflow_completed`` after the DB write succeeds. Those calls are -the only thing the Inspector SSE stream listens to, so any future refactor of -the persistence layer must keep them in place. - -We don't reconstruct a full workflow engine here — the handlers are tested -in isolation by patching just the moving parts they touch -(``_workflow_execution`` + ``_node_execution_cache``) and asserting against -the publisher module's call sites. This keeps the test compact and tied to -the contract, not the implementation. -""" - -from __future__ import annotations - -from datetime import datetime -from typing import Any -from unittest.mock import MagicMock - -import pytest - -from graphon.engine_events import GraphRunPausedEvent, NodeRunStartedEvent -from services import workflow_persistence as persistence_mod -from services.workflow_persistence import WorkflowPersistenceLayer - - -@pytest.fixture -def layer() -> WorkflowPersistenceLayer: - """Use real layer initialization with repository and runtime test doubles.""" - instance = WorkflowPersistenceLayer( - application_generate_entity=MagicMock(), - workflow_info=MagicMock(workflow_id="wf-1"), - workflow_execution_repository=MagicMock(), - workflow_node_execution_repository=MagicMock(), - ) - # Use a SimpleNamespace-like spec so Pydantic-validated callsites (e.g. - # ``WorkflowNodeExecution.new`` requires real strings) get the right types. - workflow_execution = MagicMock() - workflow_execution.id_ = "run-1" - workflow_execution.workflow_id = "wf-1" - workflow_execution.status = MagicMock(value="succeeded") - workflow_execution.outputs = {} - workflow_execution.error_message = None - workflow_execution.exceptions_count = 0 - workflow_execution.finished_at = None - instance.set_node_run_indices({"exec-1": 1}) - instance._workflow_execution = workflow_execution - # `runtime_state` is a layer-base property; stub it. - instance._runtime_state = MagicMock(total_tokens=0, node_run_steps=0, outputs={}, exceptions_count=0) - return instance - - -@pytest.fixture -def capture_publishes(monkeypatch: pytest.MonkeyPatch) -> dict[str, list]: - """Replace the two publishers with capture lists so each test can assert - on the exact arguments.""" - calls: dict[str, list] = {"node": [], "workflow": []} - - def fake_node(*, workflow_run_id: str, node_id: str, status: str) -> None: - calls["node"].append({"workflow_run_id": workflow_run_id, "node_id": node_id, "status": status}) - - def fake_workflow(*, workflow_run_id: str, status: str) -> None: - calls["workflow"].append({"workflow_run_id": workflow_run_id, "status": status}) - - monkeypatch.setattr(persistence_mod, "_inspector_publish_node_changed", fake_node) - monkeypatch.setattr(persistence_mod, "_inspector_publish_workflow_completed", fake_workflow) - return calls - - -# ────────────────────────────────────────────────────────────────────────────── -# Graph-level publish hooks -# ────────────────────────────────────────────────────────────────────────────── - - -def _graph_event(**kwargs: Any) -> MagicMock: - return MagicMock(**kwargs) - - -def test_graph_run_succeeded_publishes_workflow_completed(layer, capture_publishes): - layer._workflow_execution.status = MagicMock(value="succeeded") - layer._handle_graph_run_succeeded(_graph_event(outputs={"text": "hi"})) - assert capture_publishes["workflow"] == [{"workflow_run_id": "run-1", "status": "succeeded"}] - assert capture_publishes["node"] == [] - - -def test_graph_run_partial_succeeded_publishes_workflow_completed(layer, capture_publishes): - layer._workflow_execution.status = MagicMock(value="partial-succeeded") - layer._handle_graph_run_partial_succeeded(_graph_event(outputs={}, exceptions_count=1)) - assert capture_publishes["workflow"] == [{"workflow_run_id": "run-1", "status": "partial-succeeded"}] - - -def test_graph_run_failed_publishes_workflow_completed(layer, capture_publishes): - layer._workflow_execution.status = MagicMock(value="failed") - layer._handle_graph_run_failed(_graph_event(error="boom", exceptions_count=0)) - assert capture_publishes["workflow"] == [{"workflow_run_id": "run-1", "status": "failed"}] - - -def test_graph_run_aborted_publishes_workflow_completed(layer, capture_publishes): - layer._workflow_execution.status = MagicMock(value="stopped") - layer._handle_graph_run_aborted(_graph_event(reason="user stop")) - assert capture_publishes["workflow"] == [{"workflow_run_id": "run-1", "status": "stopped"}] - - -def test_graph_run_paused_does_not_publish_completion(layer, capture_publishes): - """Pause is not a terminal state — the Inspector keeps waiting for either - resume or a real terminal event.""" - layer.on_event(GraphRunPausedEvent(outputs={})) - assert capture_publishes["workflow"] == [] - assert capture_publishes["node"] == [] - - -# ────────────────────────────────────────────────────────────────────────────── -# Node-level publish hooks -# ────────────────────────────────────────────────────────────────────────────── - - -def _node_started_event(node_id: str = "agent-1", exec_id: str = "exec-1") -> NodeRunStartedEvent: - return NodeRunStartedEvent( - id=exec_id, - node_id=node_id, - node_type="agent", - node_title="Greeter", - predecessor_node_id=None, - container_id="", - node_version="1", - start_at=datetime(2026, 5, 26, 0, 0, 0), - ) - - -def _seed_node_execution(layer: WorkflowPersistenceLayer, exec_id: str, node_id: str) -> None: - """Inject a domain execution into the cache so the success / fail / etc - handlers (which look it up by id) can run without going through started.""" - layer._node_execution_cache[exec_id] = MagicMock( - id=exec_id, node_id=node_id, status=MagicMock(value="running"), outputs={}, error=None - ) - - -def test_node_started_publishes_running(layer, capture_publishes): - layer._handle_node_started(_node_started_event()) - assert capture_publishes["node"] == [{"workflow_run_id": "run-1", "node_id": "agent-1", "status": "running"}] - - -def test_node_retry_publishes_retry(layer, capture_publishes, monkeypatch: pytest.MonkeyPatch): - _seed_node_execution(layer, exec_id="exec-1", node_id="agent-1") - monkeypatch.setattr(layer, "_append_retry_history", lambda *args: None) - event = MagicMock(id="exec-1", error="rate limit") - layer._handle_node_retry(event) - assert capture_publishes["node"] == [{"workflow_run_id": "run-1", "node_id": "agent-1", "status": "retry"}] - - -def test_node_succeeded_publishes_succeeded(layer, capture_publishes, monkeypatch: pytest.MonkeyPatch): - _seed_node_execution(layer, exec_id="exec-1", node_id="agent-1") - # Stub the inner _update_node_execution so we don't have to construct a - # full NodeRunResult — we only want to confirm the publish happens after. - monkeypatch.setattr(layer, "_update_node_execution", lambda *a, **kw: None) - event = MagicMock(id="exec-1", node_run_result=MagicMock(), finished_at=datetime.now()) - layer._handle_node_succeeded(event) - assert capture_publishes["node"] == [{"workflow_run_id": "run-1", "node_id": "agent-1", "status": "succeeded"}] - - -def test_node_failed_publishes_failed(layer, capture_publishes, monkeypatch: pytest.MonkeyPatch): - _seed_node_execution(layer, exec_id="exec-1", node_id="agent-1") - monkeypatch.setattr(layer, "_update_node_execution", lambda *a, **kw: None) - event = MagicMock(id="exec-1", node_run_result=MagicMock(), error="bad", finished_at=datetime.now()) - layer._handle_node_failed(event) - assert capture_publishes["node"] == [{"workflow_run_id": "run-1", "node_id": "agent-1", "status": "failed"}] - - -def test_node_exception_publishes_exception(layer, capture_publishes, monkeypatch: pytest.MonkeyPatch): - _seed_node_execution(layer, exec_id="exec-1", node_id="agent-1") - monkeypatch.setattr(layer, "_update_node_execution", lambda *a, **kw: None) - event = MagicMock(id="exec-1", node_run_result=MagicMock(), error="oom", finished_at=datetime.now()) - layer._handle_node_exception(event) - assert capture_publishes["node"] == [{"workflow_run_id": "run-1", "node_id": "agent-1", "status": "exception"}] - - -def test_node_pause_requested_does_not_publish(layer, capture_publishes, monkeypatch: pytest.MonkeyPatch): - """Node pause is not an Inspector-visible state — no publish.""" - _seed_node_execution(layer, exec_id="exec-1", node_id="agent-1") - monkeypatch.setattr(layer, "_update_node_execution", lambda *a, **kw: None) - event = MagicMock(id="exec-1", node_run_result=MagicMock()) - layer._handle_node_pause_requested(event) - assert capture_publishes["node"] == [] diff --git a/api/tests/unit_tests/core/app/workflow/test_node_factory.py b/api/tests/unit_tests/core/app/workflow/test_node_factory.py deleted file mode 100644 index 8e4ed261e0f..00000000000 --- a/api/tests/unit_tests/core/app/workflow/test_node_factory.py +++ /dev/null @@ -1,166 +0,0 @@ -from collections.abc import Callable -from types import SimpleNamespace - -import pytest - -from core.app.entities.app_invoke_entities import InvokeFrom, UserFrom, build_dify_run_context -from core.workflow.node_factory import DifyNodeFactory -from graphon.enums import BuiltinNodeTypes - - -class DummyNode: - def __init__(self, *, node_id, data, init_params, runtime_state, **kwargs): - self.id = node_id - self.data = data - self.init_params = init_params - self.runtime_state = runtime_state - self.kwargs = kwargs - - -class DummyCodeNode(DummyNode): - @classmethod - def default_code_providers(cls): - return () - - -class DummyTemplateTransformNode(DummyNode): - pass - - -class DummyHttpRequestNode(DummyNode): - pass - - -class DummyKnowledgeRetrievalNode(DummyNode): - pass - - -class DummyDocumentExtractorNode(DummyNode): - pass - - -class TestDifyNodeFactory: - @pytest.fixture(autouse=True) - def _node_config(self, config_overrides: Callable[..., None]) -> None: - config_overrides( - CODE_MAX_STRING_LENGTH=10, - CODE_MAX_NUMBER=10, - CODE_MIN_NUMBER=-10, - CODE_MAX_PRECISION=4, - CODE_MAX_DEPTH=2, - CODE_MAX_NUMBER_ARRAY_LENGTH=2, - CODE_MAX_STRING_ARRAY_LENGTH=2, - CODE_MAX_OBJECT_ARRAY_LENGTH=2, - TEMPLATE_TRANSFORM_MAX_LENGTH=100, - UNSTRUCTURED_API_URL="http://u", - UNSTRUCTURED_API_KEY="key", - ) - - @staticmethod - def _stub_node_resolution(monkeypatch: pytest.MonkeyPatch, node_class): - monkeypatch.setattr( - "core.workflow.node_factory.resolve_workflow_node_class", - lambda **_kwargs: node_class, - ) - - def _factory(self): - run_context = build_dify_run_context( - tenant_id="tenant", - app_id="app", - user_id="user", - user_from=UserFrom.END_USER, - invoke_from=InvokeFrom.WEB_APP, - ) - - return DifyNodeFactory( - init_params=SimpleNamespace(run_context=run_context), - runtime_state=SimpleNamespace(), - ) - - def test_create_node_unknown_type(self): - factory = self._factory() - - with pytest.raises(ValueError, match="No class mapping found for node type: unknown"): - factory.create_node({"id": "node-1", "data": {"type": "unknown"}}) - - def test_create_node_missing_mapping(self, monkeypatch: pytest.MonkeyPatch): - factory = self._factory() - monkeypatch.setattr("core.workflow.node_factory.get_node_type_classes_mapping", lambda: {}) - - with pytest.raises(ValueError, match="No class mapping found for node type: start"): - factory.create_node({"id": "node-1", "data": {"type": BuiltinNodeTypes.START}}) - - def test_create_node_missing_latest_class(self, monkeypatch: pytest.MonkeyPatch): - factory = self._factory() - monkeypatch.setattr( - "core.workflow.node_factory.get_node_type_classes_mapping", - lambda: {BuiltinNodeTypes.START: {"1": None}}, - ) - monkeypatch.setattr("core.workflow.node_factory.LATEST_VERSION", "latest") - - with pytest.raises(ValueError, match="No latest version class found for node type: start"): - factory.create_node({"id": "node-1", "data": {"type": BuiltinNodeTypes.START}}) - - def test_create_node_selects_versioned_class(self, monkeypatch: pytest.MonkeyPatch): - factory = self._factory() - selected_versions: list[tuple[str, str]] = [] - - class DummyNodeV2(DummyNode): - pass - - def _get_mapping(): - selected_versions.append(("snapshot", "called")) - return {BuiltinNodeTypes.START: {"1": DummyNode, "2": DummyNodeV2}} - - monkeypatch.setattr("core.workflow.node_factory.get_node_type_classes_mapping", _get_mapping) - - node = factory.create_node({"id": "node-1", "data": {"type": BuiltinNodeTypes.START, "version": "2"}}) - - assert isinstance(node, DummyNodeV2) - assert node.id == "node-1" - assert selected_versions == [("snapshot", "called")] - - def test_create_node_code_branch(self, monkeypatch: pytest.MonkeyPatch): - factory = self._factory() - self._stub_node_resolution(monkeypatch, DummyCodeNode) - - node = factory.create_node({"id": "node-1", "data": {"type": BuiltinNodeTypes.CODE}}) - - assert isinstance(node, DummyCodeNode) - assert node.id == "node-1" - - def test_create_node_template_transform_branch(self, monkeypatch: pytest.MonkeyPatch): - factory = self._factory() - self._stub_node_resolution(monkeypatch, DummyTemplateTransformNode) - - node = factory.create_node({"id": "node-1", "data": {"type": BuiltinNodeTypes.TEMPLATE_TRANSFORM}}) - - assert isinstance(node, DummyTemplateTransformNode) - assert "jinja2_template_renderer" in node.kwargs - - def test_create_node_http_request_branch(self, monkeypatch: pytest.MonkeyPatch): - factory = self._factory() - self._stub_node_resolution(monkeypatch, DummyHttpRequestNode) - - node = factory.create_node({"id": "node-1", "data": {"type": BuiltinNodeTypes.HTTP_REQUEST}}) - - assert isinstance(node, DummyHttpRequestNode) - assert "http_request_config" in node.kwargs - - def test_create_node_knowledge_retrieval_branch(self, monkeypatch: pytest.MonkeyPatch): - factory = self._factory() - self._stub_node_resolution(monkeypatch, DummyKnowledgeRetrievalNode) - - node = factory.create_node({"id": "node-1", "data": {"type": BuiltinNodeTypes.KNOWLEDGE_RETRIEVAL}}) - - assert isinstance(node, DummyKnowledgeRetrievalNode) - assert node.kwargs == {} - - def test_create_node_document_extractor_branch(self, monkeypatch: pytest.MonkeyPatch): - factory = self._factory() - self._stub_node_resolution(monkeypatch, DummyDocumentExtractorNode) - - node = factory.create_node({"id": "node-1", "data": {"type": BuiltinNodeTypes.DOCUMENT_EXTRACTOR}}) - - assert isinstance(node, DummyDocumentExtractorNode) - assert "unstructured_api_config" in node.kwargs diff --git a/api/tests/unit_tests/core/app/workflow/test_observability_layer_extra.py b/api/tests/unit_tests/core/app/workflow/test_observability_layer_extra.py deleted file mode 100644 index 27a6f4a70cf..00000000000 --- a/api/tests/unit_tests/core/app/workflow/test_observability_layer_extra.py +++ /dev/null @@ -1,54 +0,0 @@ -from __future__ import annotations - -from types import SimpleNamespace - -import pytest - -from core.app.workflow.layers.observability import ObservabilityLayer -from tests.unit_tests.config_override import apply_config_overrides - - -class TestObservabilityLayerExtras: - def test_init_tracer_disables_when_get_tracer_fails( - self, monkeypatch: pytest.MonkeyPatch, caplog: pytest.LogCaptureFixture - ): - apply_config_overrides(monkeypatch, ENABLE_OTEL=True) - monkeypatch.setattr("core.app.workflow.layers.observability.is_instrument_flag_enabled", lambda: False) - - def _raise(*_args, **_kwargs): - raise RuntimeError("tracer init failed") - - monkeypatch.setattr("core.app.workflow.layers.observability.get_tracer", _raise) - - layer = ObservabilityLayer() - - assert layer._is_disabled is True - assert layer._tracer is None - assert "Failed to get OpenTelemetry tracer" in caplog.text - - def test_node_run_context_skips_without_execution_id(self): - layer = ObservabilityLayer() - layer._is_disabled = False - calls: list[str] = [] - layer._tracer = SimpleNamespace(start_span=lambda *_args, **_kwargs: calls.append("called")) - - with layer.node_run_context(SimpleNamespace(execution_id=None, title="node", id="node")): - assert calls == [] - - assert calls == [] - - def test_node_runs_when_span_creation_fails(self, caplog: pytest.LogCaptureFixture): - layer = ObservabilityLayer() - layer._is_disabled = False - - def _raise(*_args, **_kwargs): - raise RuntimeError("start failed") - - layer._tracer = SimpleNamespace(start_span=_raise) - - node_executed = False - with layer.node_run_context(SimpleNamespace(execution_id="exec", title="node", id="node")): - node_executed = True - - assert node_executed - assert "start failed" in caplog.text diff --git a/api/tests/unit_tests/core/app/workflow/test_persistence_layer.py b/api/tests/unit_tests/core/app/workflow/test_persistence_layer.py deleted file mode 100644 index cdf3f984378..00000000000 --- a/api/tests/unit_tests/core/app/workflow/test_persistence_layer.py +++ /dev/null @@ -1,1216 +0,0 @@ -from __future__ import annotations - -from dataclasses import replace -from datetime import UTC, datetime, timedelta -from types import SimpleNamespace - -import pytest -from flask import Flask -from sqlalchemy import event as sqlalchemy_event -from sqlalchemy import select -from sqlalchemy.orm import Session, sessionmaker - -from core.app.entities.app_invoke_entities import AdvancedChatAppGenerateEntity, WorkflowAppGenerateEntity -from core.app.workflow.retry_history import RETRY_HISTORY_PROCESS_DATA_KEY -from core.ops.base_trace_instance import bind_application_tracing_hooks -from core.ops.ops_trace_manager import TraceTask, TraceTaskName -from core.repositories.sqlalchemy_workflow_execution_repository import SQLAlchemyWorkflowExecutionRepository -from core.repositories.sqlalchemy_workflow_node_execution_repository import SQLAlchemyWorkflowNodeExecutionRepository -from core.tools.workflow_as_tool.repository import WorkflowToolSource -from core.workflow.node_execution_process_data import ( - WORKFLOW_TOOL_INVOCATION_ID_KEY, - WORKFLOW_TOOL_ROOT_APP_ID_KEY, -) -from core.workflow.system_variables import SystemVariableKey, build_system_variables -from enums import WorkflowKind -from graphon.engine.event.node_failure import NodeFailureHandler -from graphon.engine.event.processor import NodeEventProcessor -from graphon.engine.event.stream import EventStream -from graphon.engine.frame import ExecutionFrame, FrameRegistry -from graphon.engine.scheduler import Scheduler -from graphon.engine.worker import NodeEventTask -from graphon.engine_events import ( - EngineEvent, - GraphRunAbortedEvent, - GraphRunFailedEvent, - GraphRunPartialSucceededEvent, - GraphRunPausedEvent, - GraphRunStartedEvent, - GraphRunSucceededEvent, - NodeEvent, - NodeRunExceptionEvent, - NodeRunFailedEvent, - NodeRunPauseRequestedEvent, - NodeRunRetryEvent, - NodeRunStartedEvent, - NodeRunSucceededEvent, -) -from graphon.entities import WorkflowNodeExecution, WorkflowStartReason -from graphon.entities.pause_reason import SchedulingPause -from graphon.enums import ( - BuiltinNodeTypes, - WorkflowExecutionStatus, - WorkflowNodeExecutionMetadataKey, - WorkflowNodeExecutionStatus, - WorkflowType, -) -from graphon.graph import Graph -from graphon.model_runtime.entities.llm_entities import LLMUsage -from graphon.node_events import NodeRunResult -from graphon.nodes.start.entities import StartNodeData -from graphon.nodes.start.start_node import StartNode -from graphon.runtime import ReadOnlyRuntimeStateWrapper, RuntimeState, VariablePool -from graphon.runtime.execution import ROOT_FRAME_ID -from models import Account, WorkflowRun -from models.enums import WorkflowRunTriggeredFrom -from models.workflow import WorkflowNodeExecutionModel, WorkflowNodeExecutionTriggeredFrom -from repositories.entities.workflow_execution import WorkflowExecution -from repositories.sqlalchemy_api_workflow_node_execution_repository import ( - DifyAPISQLAlchemyWorkflowNodeExecutionRepository, -) -from services.tracing.runtime_bridge import RuntimeApplicationTracingHooks -from services.workflow_persistence import PersistenceWorkflowInfo, WorkflowPersistenceLayer -from tests.unit_tests.config_override import apply_config_overrides -from tests.workflow_test_utils import build_test_graph_init_params - - -@pytest.fixture(autouse=True) -def _application_tracing_disabled(monkeypatch: pytest.MonkeyPatch) -> None: - apply_config_overrides(monkeypatch, OPS_AGENT_TRACE_ENABLED=False) - - -class _RepoRecorder: - def __init__(self) -> None: - self.saved: list[object] = [] - self.synchronously_saved: list[object] = [] - self.saved_exec_data: list[object] = [] - - def save(self, entity): - self.saved.append(entity) - - def save_many(self, entities): - self.saved.extend(entities) - - def save_synchronously(self, entity): - self.synchronously_saved.append(entity) - - def save_execution_data(self, entity): - self.saved_exec_data.append(entity) - - -def _naive_utc_now() -> datetime: - return datetime.now(UTC).replace(tzinfo=None) - - -def _make_layer( - system_variables: list | None = None, - *, - extras: dict | None = None, - graph_data: dict | None = None, - trace_manager: object | None = None, - workflow_execution_repo=None, - workflow_node_execution_repo=None, - node_run_indices=None, -): - system_variables = system_variables or build_system_variables( - workflow_execution_id="run-id", - conversation_id="conv-id", - ) - runtime_state = RuntimeState( - workflow_id="test-workflow", - variable_pool=VariablePool.from_bootstrap(system_variables=system_variables), - start_at=0.0, - ) - read_only_state = ReadOnlyRuntimeStateWrapper(runtime_state) - - application_generate_entity = WorkflowAppGenerateEntity.model_construct( - task_id="task", - app_config=SimpleNamespace(app_id="app", tenant_id="tenant"), - inputs={"foo": "bar"}, - files=[], - user_id="user", - stream=False, - invoke_from=None, - trace_manager=None, - workflow_execution_id="run-id", - extras=extras or {}, - call_depth=0, - ) - - workflow_info = PersistenceWorkflowInfo( - workflow_id="workflow-id", - workflow_type=WorkflowType.WORKFLOW, - version="1", - graph_data=graph_data if graph_data is not None else {"nodes": [], "edges": []}, - ) - - workflow_execution_repo = workflow_execution_repo if workflow_execution_repo is not None else _RepoRecorder() - workflow_node_execution_repo = ( - workflow_node_execution_repo if workflow_node_execution_repo is not None else _RepoRecorder() - ) - - layer = WorkflowPersistenceLayer( - application_generate_entity=application_generate_entity, - workflow_info=workflow_info, - workflow_execution_repository=workflow_execution_repo, - workflow_node_execution_repository=workflow_node_execution_repo, - trace_manager=trace_manager, - ) - layer.initialize(read_only_state, command_channel=None) - layer.set_node_run_indices(node_run_indices or {}) - - return layer, workflow_execution_repo, workflow_node_execution_repo, runtime_state - - -def _make_sql_layer(session_factory, *, node_run_indices=None): - user = Account(name="Test", email="test@example.com") - user.id = "user" - layer, _, _, _ = _make_layer( - node_run_indices=node_run_indices, - workflow_execution_repo=SQLAlchemyWorkflowExecutionRepository( - session_factory=session_factory, - tenant_id="tenant", - user=user, - app_id="caller-app", - triggered_from=WorkflowRunTriggeredFrom.APP_RUN, - ), - workflow_node_execution_repo=SQLAlchemyWorkflowNodeExecutionRepository( - session_factory=session_factory, - tenant_id="tenant", - user=user, - app_id="caller-app", - triggered_from=WorkflowNodeExecutionTriggeredFrom.WORKFLOW_RUN, - ), - ) - layer._application_generate_entity.app_config.app_id = "caller-app" - return layer - - -def test_tool_trace_identity_survives_start_pause_and_resume(): - layer, _, node_repo, _ = _make_layer(node_run_indices={"tool-exec": 1}) - layer.on_event(GraphRunStartedEvent()) - started_at = _naive_utc_now() - layer.on_event( - NodeRunStartedEvent( - id="tool-exec", - node_id="tool", - node_type=BuiltinNodeTypes.TOOL, - node_title="Nested tool", - provider_type="workflow", - provider_id="provider", - start_at=started_at, - node_run_result=NodeRunResult(process_data={WORKFLOW_TOOL_INVOCATION_ID_KEY: "invocation"}), - ) - ) - for event in ( - NodeRunPauseRequestedEvent( - id="tool-exec", node_id="tool", node_type=BuiltinNodeTypes.TOOL, reason=SchedulingPause(message="pause") - ), - NodeRunSucceededEvent(id="tool-exec", node_id="tool", node_type=BuiltinNodeTypes.TOOL, start_at=started_at), - ): - layer.on_event(event) - execution = node_repo.saved_exec_data[-1] - assert execution.process_data[WORKFLOW_TOOL_INVOCATION_ID_KEY] == "invocation" - assert execution.metadata[WorkflowNodeExecutionMetadataKey.TOOL_INFO] == { - "provider_type": "workflow", - "provider_id": "provider", - } - - -def _tool_source(): - return WorkflowToolSource( - app_id="source-app", - workflow_id="source-workflow", - graph_config={"nodes": [{"id": "source-loop", "data": {"type": "loop"}}], "edges": []}, - features_dict={}, - environment_variables=[], - workflow_kind=WorkflowKind.STANDARD, - ) - - -def test_workflow_tool_agent_nodes_persist_in_source_app_under_only_root_run(sqlite_session_factory): - layer = _make_sql_layer(sqlite_session_factory, node_run_indices={"caller-exec": 1, "source-agent-exec": 1}) - listen = layer.create_workflow_tool_event_listener(_tool_source(), "workflow-id", "caller-exec") - layer.on_graph_start() - layer.on_event(GraphRunStartedEvent()) - started_at = _naive_utc_now() - layer.on_event( - NodeRunStartedEvent( - id="caller-exec", - node_id="caller", - node_type=BuiltinNodeTypes.TOOL, - node_title="Workflow Tool", - start_at=started_at, - ) - ) - listen( - NodeRunStartedEvent( - id="source-agent-exec", - node_id="agent", - node_type=BuiltinNodeTypes.AGENT, - node_version="2", - node_title="Agent", - start_at=started_at, - container_id="source-loop", - node_run_result=NodeRunResult( - process_data={WORKFLOW_TOOL_ROOT_APP_ID_KEY: "untrusted-app"}, metadata={"loop_id": "source-loop"} - ), - ) - ) - - with sqlite_session_factory() as session: - rows = {row.id: row for row in session.scalars(select(WorkflowNodeExecutionModel))} - assert rows["caller-exec"].app_id == "caller-app" - source_row = rows["source-agent-exec"] - assert (source_row.tenant_id, source_row.app_id, source_row.workflow_id, source_row.workflow_run_id) == ( - "tenant", - "source-app", - "source-workflow", - "run-id", - ) - assert source_row.status == "running" - assert source_row.process_data_dict[WORKFLOW_TOOL_ROOT_APP_ID_KEY] == "caller-app" - assert source_row.execution_metadata_dict["loop_id"] == "source-loop" - assert session.scalars(select(WorkflowRun)).one().workflow_id == "workflow-id" - - listen( - NodeRunSucceededEvent( - id="source-agent-exec", - node_id="agent", - node_type=BuiltinNodeTypes.AGENT, - node_title="Agent", - start_at=started_at, - node_run_result=NodeRunResult(status=WorkflowNodeExecutionStatus.SUCCEEDED, outputs={"answer": "ok"}), - ) - ) - layer.on_event(GraphRunSucceededEvent(outputs={"result": "caller-result"})) - - with sqlite_session_factory() as session: - run = session.scalars(select(WorkflowRun)).one() - assert (run.app_id, run.workflow_id, run.status, run.outputs_dict) == ( - "caller-app", - "workflow-id", - "succeeded", - {"result": "caller-result"}, - ) - row = session.get(WorkflowNodeExecutionModel, "source-agent-exec") - assert row.status == "succeeded" - assert row.outputs_dict == {"answer": "ok"} - assert row.process_data_dict[WORKFLOW_TOOL_ROOT_APP_ID_KEY] == "caller-app" - - -def test_nested_tool_callers_keep_distinct_invocations_in_shared_source_layers( - sqlite_session_factory: sessionmaker[Session], -) -> None: - ids = ["a-call-1", "a-call-2", "b-call-1", "b-call-2", "c-child"] - layer = _make_sql_layer(sqlite_session_factory, node_run_indices=dict.fromkeys(ids, 1)) - source = _tool_source() - nested = replace(source, app_id="nested-app", workflow_id="nested-workflow") - first_call = layer.create_workflow_tool_event_listener(source, "workflow-id", "a-call-1") - second_call = layer.create_workflow_tool_event_listener(source, "workflow-id", "a-call-2") - nested_call = layer.create_workflow_tool_event_listener(nested, source.workflow_id, "b-call-1") - layer.on_event(GraphRunStartedEvent()) - for execution_id, listener in zip( - ids, [layer.on_event, layer.on_event, first_call, second_call, nested_call], strict=True - ): - listener( - NodeRunStartedEvent( - id=execution_id, - node_id="same-canvas-node", - node_type=BuiltinNodeTypes.TOOL, - node_title="Tool", - start_at=_naive_utc_now(), - ) - ) - repository = DifyAPISQLAlchemyWorkflowNodeExecutionRepository(sqlite_session_factory) - for caller, expected, source_app, source_workflow in ( - ("a-call-1", "b-call-1", "source-app", "source-workflow"), - ("a-call-2", "b-call-2", "source-app", "source-workflow"), - ("b-call-1", "c-child", "nested-app", "nested-workflow"), - ): - children = repository.get_workflow_tool_executions("tenant", "run-id", caller) - assert [child.id for child in children] == [expected] - child = children[0] - assert (child.app_id, child.workflow_id, child.workflow_run_id) == (source_app, source_workflow, "run-id") - assert child.triggered_from_workflow_id == ("source-workflow" if caller == "b-call-1" else "workflow-id") - assert child.triggered_from_node_execution_id == caller - - -def test_workflow_tool_resume_updates_existing_running_container(sqlite_session_factory): - layer = _make_sql_layer(sqlite_session_factory, node_run_indices={"loop-exec": 1}) - listen = layer.create_workflow_tool_event_listener(_tool_source(), "workflow-id", "caller-exec") - layer.on_event(GraphRunStartedEvent()) - started_at = _naive_utc_now() - listen( - NodeRunStartedEvent( - id="loop-exec", - node_id="source-loop", - node_type=BuiltinNodeTypes.LOOP, - node_title="Approval Loop", - start_at=started_at, - ) - ) - layer.on_event(GraphRunPausedEvent(outputs={}, reasons=[])) - - resumed = _make_sql_layer(sqlite_session_factory, node_run_indices={"loop-exec": 1}) - resume_listener = resumed.create_workflow_tool_event_listener( - _tool_source(), - "workflow-id", - "caller-exec", - node_executions=resumed._workflow_node_execution_repository.for_workflow_tool( - "source-app" - ).get_by_workflow_execution("run-id"), - ) - resumed.on_graph_start() - resumed.on_event(GraphRunStartedEvent(reason=WorkflowStartReason.RESUMPTION)) - resume_listener( - NodeRunSucceededEvent( - id="loop-exec", - node_id="source-loop", - node_type=BuiltinNodeTypes.LOOP, - node_title="Approval Loop", - start_at=started_at, - node_run_result=NodeRunResult(status=WorkflowNodeExecutionStatus.SUCCEEDED, outputs={"approved": True}), - ) - ) - resumed.on_event(GraphRunSucceededEvent(outputs={"answer": "approved"})) - - with sqlite_session_factory() as session: - node = session.scalars(select(WorkflowNodeExecutionModel)).one() - assert (node.app_id, node.workflow_id, node.workflow_run_id, node.status) == ( - "source-app", - "source-workflow", - "run-id", - "succeeded", - ) - assert node.outputs_dict == {"approved": True} - run = session.scalars(select(WorkflowRun)).one() - assert (run.app_id, run.workflow_id, run.status) == ("caller-app", "workflow-id", "succeeded") - - -@pytest.mark.parametrize("source_workflow_id", ["older-workflow", "workflow-id"]) -def test_root_resume_isolates_same_app_workflow_tool_origin(sqlite_session_factory, source_workflow_id): - source = replace(_tool_source(), app_id="caller-app", workflow_id=source_workflow_id) - layer = _make_sql_layer(sqlite_session_factory, node_run_indices={"caller-exec": 1, "older-exec": 1}) - listen = layer.create_workflow_tool_event_listener(source, "workflow-id", "caller-exec") - layer.on_event(GraphRunStartedEvent()) - started_at = _naive_utc_now() - layer.on_event( - NodeRunStartedEvent( - id="caller-exec", - node_id="caller", - node_type=BuiltinNodeTypes.TOOL, - node_title="Caller", - start_at=started_at, - ) - ) - listen( - NodeRunStartedEvent( - id="older-exec", - node_id="older-node", - node_type=BuiltinNodeTypes.LLM, - node_title="Older version", - start_at=started_at, - ) - ) - layer.on_event(GraphRunPausedEvent()) - - resumed = _make_sql_layer(sqlite_session_factory, node_run_indices={"caller-exec": 1, "older-exec": 1}) - resumed.set_node_execution_history(resumed._workflow_node_execution_repository.get_by_workflow_execution("run-id")) - resume_listener = resumed.create_workflow_tool_event_listener( - source, - "workflow-id", - "caller-exec", - node_executions=resumed._workflow_node_execution_repository.for_workflow_tool( - source.app_id - ).get_by_workflow_execution("run-id"), - ) - resumed.on_event(GraphRunStartedEvent(reason=WorkflowStartReason.RESUMPTION)) - resume_listener( - NodeRunSucceededEvent( - id="older-exec", - node_id="older-node", - node_type=BuiltinNodeTypes.LLM, - start_at=started_at, - node_run_result=NodeRunResult(status=WorkflowNodeExecutionStatus.SUCCEEDED), - ) - ) - resumed.on_event(GraphRunFailedEvent(error="caller failed later")) - - with sqlite_session_factory() as session: - rows = {row.id: row for row in session.scalars(select(WorkflowNodeExecutionModel))} - assert rows["caller-exec"].status == "failed" - assert rows["older-exec"].status == "succeeded" - assert rows["older-exec"].workflow_id == source_workflow_id - - -@pytest.mark.parametrize("source", [False, True]) -@pytest.mark.parametrize("outcome", ["succeeded", "failed", "aborted", "paused"]) -def test_retry_starts_preserve_original_execution_and_attempt_history( - sqlite_session_factory: sessionmaker[Session], - monkeypatch: pytest.MonkeyPatch, - source: bool, - outcome: str, -) -> None: - layer = _make_sql_layer(sqlite_session_factory, node_run_indices={"retry-exec": 1}) - listen = ( - layer.create_workflow_tool_event_listener(_tool_source(), "workflow-id", "caller-exec") - if source - else layer.on_event - ) - layer.on_event(GraphRunStartedEvent()) - workflow_id = "source-workflow" if source else "workflow-id" - state = RuntimeState(workflow_id=workflow_id, variable_pool=VariablePool(), start_at=0) - frames = FrameRegistry() - node = StartNode( - node_id="child", - data=StartNodeData(title="Child", variables=[]), - init_params=build_test_graph_init_params(workflow_id=workflow_id), - runtime_state=state, - ) - graph = Graph.new().add_root(node).build() - frames.register( - ExecutionFrame( - frame_id=ROOT_FRAME_ID, - graph=graph, - state=state, - scheduler=Scheduler(graph, state, ROOT_FRAME_ID), - failure_handler=NodeFailureHandler(graph, state.graph_execution), - ) - ) - - event_stream = EventStream(layers=[]) - collect = event_stream.collect - - def collect_source_event(event: EngineEvent) -> None: - assert isinstance(event, NodeEvent) - listen(event) - collect(event) - - monkeypatch.setattr(event_stream, "collect", collect_source_event) - processor = NodeEventProcessor(state.graph_execution, event_stream, frames, {}) - - def dispatch(event: NodeEvent) -> None: - processor.dispatch(NodeEventTask(frame_id=ROOT_FRAME_ID, event=event)) - - started_at = _naive_utc_now() - started = NodeRunStartedEvent( - id="retry-exec", - node_id="child", - node_type=BuiltinNodeTypes.LLM, - node_title="Child", - start_at=started_at, - ) - dispatch(started) - for attempt in (1, 2): - dispatch( - NodeRunRetryEvent( - id="retry-exec", - node_id="child", - node_type=BuiltinNodeTypes.LLM, - node_title="Child", - start_at=started_at, - retry_index=attempt, - error=f"attempt-{attempt}", - node_run_result=NodeRunResult(outputs={"attempt": attempt}), - ) - ) - dispatch(started) - if outcome == "succeeded": - dispatch( - NodeRunSucceededEvent( - id="retry-exec", - node_id="child", - node_type=BuiltinNodeTypes.LLM, - start_at=started_at, - node_run_result=NodeRunResult(status=WorkflowNodeExecutionStatus.SUCCEEDED, outputs={"ok": True}), - ) - ) - elif outcome == "failed": - layer.on_event(GraphRunFailedEvent(error="terminal failure")) - elif outcome == "aborted": - layer.on_event(GraphRunAbortedEvent(reason="terminal failure")) - else: - commits = [] - - def record_commit(_session: Session) -> None: - commits.append(True) - - sqlalchemy_event.listen(Session, "after_commit", record_commit) - try: - layer.persist_unfinished_nodes_synchronously() - finally: - sqlalchemy_event.remove(Session, "after_commit", record_commit) - assert len(commits) == 1 - - with sqlite_session_factory() as session: - row = session.scalars(select(WorkflowNodeExecutionModel)).one() - assert row.index == 1 - assert row.created_at == started_at - assert row.workflow_id == workflow_id - assert row.app_id == ("source-app" if source else "caller-app") - if outcome == "succeeded": - assert row.status == "succeeded" - assert row.outputs_dict == {"ok": True} - elif outcome == "paused": - assert row.status == "retry" - assert row.finished_at is None - else: - assert row.status == "failed" - assert row.error == "terminal failure" - assert row.finished_at is not None - history = row.process_data_dict[RETRY_HISTORY_PROCESS_DATA_KEY] - assert [(attempt["retry_index"], attempt["error"]) for attempt in history] == [ - (1, "attempt-1"), - (2, "attempt-2"), - ] - - -@pytest.mark.parametrize( - "event", [GraphRunFailedEvent(error="failed", exceptions_count=1), GraphRunAbortedEvent(reason="aborted")] -) -def test_root_terminal_failure_finishes_running_workflow_tool_nodes(sqlite_session_factory, event): - layer = _make_sql_layer(sqlite_session_factory, node_run_indices={"source-exec": 1}) - listen = layer.create_workflow_tool_event_listener(_tool_source(), "workflow-id", "caller-exec") - layer.on_event(GraphRunStartedEvent()) - listen( - NodeRunStartedEvent( - id="source-exec", - node_id="child", - node_type=BuiltinNodeTypes.LLM, - node_title="Child", - start_at=_naive_utc_now(), - ) - ) - - layer.on_event(event) - - with sqlite_session_factory() as session: - node = session.scalars(select(WorkflowNodeExecutionModel)).one() - assert node.status == "failed" - assert node.error == (event.error if isinstance(event, GraphRunFailedEvent) else event.reason) - assert node.finished_at is not None - assert session.scalars(select(WorkflowRun)).one().workflow_id == "workflow-id" - - -def test_terminal_failure_batches_only_unfinished_nodes_per_source( - sqlite_session_factory: sessionmaker[Session], -) -> None: - ids = ( - "root-a", - "root-b", - "root-retry", - "root-paused", - "source-a", - "source-b", - "source-retry", - "source-paused", - "source-done", - ) - layer = _make_sql_layer(sqlite_session_factory, node_run_indices=dict.fromkeys(ids, 1)) - listen = layer.create_workflow_tool_event_listener(_tool_source(), "workflow-id", "root-a") - layer.on_event(GraphRunStartedEvent()) - started_at = _naive_utc_now() - for execution_id in ids: - listener = listen if execution_id.startswith("source") else layer.on_event - listener( - NodeRunStartedEvent( - id=execution_id, - node_id=execution_id, - node_type=BuiltinNodeTypes.CODE, - node_title="Code", - start_at=started_at, - ) - ) - if execution_id.endswith("paused"): - listener( - NodeRunPauseRequestedEvent( - id=execution_id, - node_id=execution_id, - node_type=BuiltinNodeTypes.CODE, - reason=SchedulingPause(message="pause"), - ) - ) - elif execution_id.endswith("retry"): - listener( - NodeRunRetryEvent( - id=execution_id, - node_id=execution_id, - node_type=BuiltinNodeTypes.CODE, - node_title="Code", - start_at=started_at, - retry_index=1, - error="retry failed", - ) - ) - listen( - NodeRunSucceededEvent( - id="source-done", - node_id="source-done", - node_type=BuiltinNodeTypes.CODE, - start_at=started_at, - node_run_result=NodeRunResult(status=WorkflowNodeExecutionStatus.SUCCEEDED, outputs={"answer": "kept"}), - ) - ) - commits = [] - - def record_commit(_session: Session) -> None: - commits.append(True) - - sqlalchemy_event.listen(Session, "after_commit", record_commit) - try: - layer.on_event(GraphRunFailedEvent(error="terminal failure", exceptions_count=0)) - finally: - sqlalchemy_event.remove(Session, "after_commit", record_commit) - - with sqlite_session_factory() as session: - rows = {row.id: row for row in session.scalars(select(WorkflowNodeExecutionModel))} - assert {key: row.status for key, row in rows.items()} == { - "root-a": "failed", - "root-b": "failed", - "root-retry": "failed", - "root-paused": "paused", - "source-a": "failed", - "source-b": "failed", - "source-retry": "failed", - "source-paused": "paused", - "source-done": "succeeded", - } - assert rows["source-done"].outputs_dict == {"answer": "kept"} - for execution_id in ("root-a", "root-b", "root-retry", "source-a", "source-b", "source-retry"): - assert rows[execution_id].error == "terminal failure" - assert rows[execution_id].finished_at is not None - for execution_id in ("root-retry", "source-retry"): - assert rows[execution_id].process_data_dict[RETRY_HISTORY_PROCESS_DATA_KEY][0]["error"] == "retry failed" - assert rows["source-a"].triggered_from_node_execution_id == "root-a" - assert rows["source-a"].app_id == "source-app" - # One commit per source owner, plus the workflow summary. - assert len(commits) == 3 - - -class TestWorkflowPersistenceLayer: - def test_start_freezes_tracing_only_after_workflow_save(self, monkeypatch: pytest.MonkeyPatch) -> None: - apply_config_overrides(monkeypatch, OPS_AGENT_TRACE_ENABLED=True) - layer, workflow_repo, _, _ = _make_layer(extras={"trace_session_id": "external-session"}) - calls: list[tuple[str, str, str, str, object]] = [] - - def freeze( - entity: AdvancedChatAppGenerateEntity | WorkflowAppGenerateEntity, - *, - workflow_id: str, - workflow_run_id: str, - conversation_id: object, - graph_data: object, - ) -> None: - assert workflow_repo.synchronously_saved - assert graph_data == {"nodes": [], "edges": []} - saved_execution = workflow_repo.synchronously_saved[-1] - assert isinstance(saved_execution, WorkflowExecution) - assert saved_execution.id_ == workflow_run_id - calls.append( - (entity.app_config.app_id, entity.app_config.tenant_id, workflow_id, workflow_run_id, conversation_id) - ) - - monkeypatch.setattr(RuntimeApplicationTracingHooks, "freeze_workflow_run", staticmethod(freeze)) - hooks = RuntimeApplicationTracingHooks() - app = Flask(__name__) - bind_application_tracing_hooks(app, hooks) - - with app.app_context(): - layer._handle_graph_run_started() - assert calls == [("app", "tenant", "workflow-id", "run-id", "conv-id")] - - def test_get_execution_id_requires_system_variable(self): - layer, _, _, _ = _make_layer(build_system_variables()) - - with pytest.raises(ValueError, match="workflow_execution_id must be provided"): - layer._get_execution_id() - - def test_prepare_workflow_inputs_excludes_conversation_id(self, monkeypatch: pytest.MonkeyPatch): - layer, _, _, _ = _make_layer() - - monkeypatch.setattr( - "core.workflow.workflow_entry.WorkflowEntry.handle_special_values", - lambda inputs: inputs, - ) - - inputs = layer._prepare_workflow_inputs() - - assert "sys.conversation_id" not in inputs - assert inputs[f"sys.{SystemVariableKey.WORKFLOW_EXECUTION_ID.value}"] == "run-id" - - @pytest.mark.parametrize( - "status", - [ - WorkflowNodeExecutionStatus.RUNNING, - WorkflowNodeExecutionStatus.RETRY, - WorkflowNodeExecutionStatus.PAUSED, - WorkflowNodeExecutionStatus.SUCCEEDED, - ], - ) - def test_fail_running_node_executions_marks_only_live_attempts_failed(self, status): - layer, _, node_repo, _ = _make_layer() - - execution = WorkflowNodeExecution( - id="exec-id", - workflow_id="workflow-id", - workflow_execution_id="run-id", - index=1, - node_id="node", - node_type=BuiltinNodeTypes.START, - title="Start", - status=status, - created_at=_naive_utc_now(), - ) - layer._node_execution_cache[execution.id] = execution - - layer._fail_running_node_executions(error_message="boom") - - if status in (WorkflowNodeExecutionStatus.RUNNING, WorkflowNodeExecutionStatus.RETRY): - assert execution.status == WorkflowNodeExecutionStatus.FAILED - assert execution.error == "boom" - assert execution.finished_at is not None - assert node_repo.saved == [execution] - else: - assert execution.status == status - assert execution.finished_at is None - assert node_repo.saved == [] - - def test_resumption_restores_container_execution_before_terminal_event(self): - layer, _, node_repo, _ = _make_layer() - started_at = _naive_utc_now() - execution = WorkflowNodeExecution( - id="loop-exec", - workflow_id="workflow-id", - workflow_execution_id="run-id", - index=4, - node_id="loop", - node_type=BuiltinNodeTypes.LOOP, - title="Loop", - status=WorkflowNodeExecutionStatus.RUNNING, - created_at=started_at, - ) - layer.set_node_execution_history([execution]) - - layer.on_event(GraphRunStartedEvent(reason=WorkflowStartReason.RESUMPTION)) - layer.on_event( - NodeRunSucceededEvent( - id=execution.id, - node_id=execution.node_id, - node_type=execution.node_type, - start_at=started_at, - finished_at=started_at + timedelta(seconds=2), - node_run_result=NodeRunResult(status=WorkflowNodeExecutionStatus.SUCCEEDED), - ) - ) - - assert execution.status == WorkflowNodeExecutionStatus.SUCCEEDED - assert execution.elapsed_time == 2.0 - assert execution.index == 4 - - def test_handle_graph_run_succeeded_updates_execution(self): - layer, exec_repo, _, runtime_state = _make_layer() - layer.on_event(GraphRunStartedEvent()) - usage = LLMUsage.empty_usage() - usage.total_tokens = 3 - runtime_state.add_llm_usage(usage) - for _ in range(2): - runtime_state.increment_node_run_steps() - runtime_state.set_output("out", "v") - - layer.on_event(GraphRunSucceededEvent(outputs={"ok": True})) - - saved = exec_repo.synchronously_saved[-1] - assert saved.status == WorkflowExecutionStatus.SUCCEEDED - assert saved.total_tokens == 3 - assert saved.total_steps == 2 - - def test_handle_graph_run_partial_succeeded_updates_execution(self): - layer, exec_repo, _, runtime_state = _make_layer() - layer.on_event(GraphRunStartedEvent()) - usage = LLMUsage.empty_usage() - usage.total_tokens = 5 - runtime_state.add_llm_usage(usage) - for _ in range(4): - runtime_state.increment_node_run_steps() - runtime_state._graph_execution = SimpleNamespace(exceptions_count=2) - - layer.on_event(GraphRunPartialSucceededEvent(outputs={"ok": True}, exceptions_count=2)) - - saved = exec_repo.synchronously_saved[-1] - assert saved.status == WorkflowExecutionStatus.PARTIAL_SUCCEEDED - assert saved.exceptions_count == 2 - assert saved.total_tokens == 5 - - def test_handle_graph_run_failed_marks_nodes_and_enqueues_trace(self): - trace_tasks: list[object] = [] - trace_manager = SimpleNamespace(user_id="user", add_trace_task=lambda task: trace_tasks.append(task)) - layer, exec_repo, node_repo, _ = _make_layer(extras={"external_trace_id": "trace"}, trace_manager=trace_manager) - layer.on_event(GraphRunStartedEvent()) - - running = WorkflowNodeExecution( - id="node-exec", - workflow_id="workflow-id", - workflow_execution_id="run-id", - index=1, - node_id="node", - node_type=BuiltinNodeTypes.START, - title="Start", - created_at=_naive_utc_now(), - ) - layer._node_execution_cache[running.id] = running - - layer.on_event(GraphRunFailedEvent(error="boom", exceptions_count=1)) - - assert node_repo.saved - assert exec_repo.synchronously_saved[-1].status == WorkflowExecutionStatus.FAILED - assert trace_tasks - - def test_handle_graph_run_succeeded_enqueues_parent_trace_context(self, monkeypatch: pytest.MonkeyPatch): - trace_tasks: list[TraceTask] = [] - trace_manager = SimpleNamespace(user_id="user", add_trace_task=lambda task: trace_tasks.append(task)) - layer, _, _, _ = _make_layer( - extras={ - "external_trace_id": "trace", - "trace_session_id": "session-1", - "parent_trace_context": { - "parent_workflow_run_id": "outer-workflow-run-1", - "parent_node_execution_id": "outer-node-execution-1", - }, - }, - trace_manager=trace_manager, - ) - layer.on_event(GraphRunStartedEvent()) - - captured: dict[str, object] = {} - - def fake_workflow_trace( - self: TraceTask, - *, - workflow_run_id: str | None, - conversation_id: str | None, - user_id: str | None, - total_tokens_override: int | None = None, - ): - captured["trace_type"] = self.trace_type - captured["external_trace_id"] = self.kwargs.get("external_trace_id") - captured["trace_session_id"] = self.kwargs.get("trace_session_id") - captured["parent_trace_context"] = self.kwargs.get("parent_trace_context") - captured["workflow_run_id"] = workflow_run_id - return {"ok": True} - - monkeypatch.setattr(TraceTask, "workflow_trace", fake_workflow_trace) - - layer.on_event(GraphRunSucceededEvent(outputs={"ok": True})) - - assert trace_tasks - trace_task = trace_tasks[0] - assert trace_task.trace_type == TraceTaskName.WORKFLOW_TRACE - assert trace_task.kwargs["external_trace_id"] == "trace" - assert trace_task.kwargs["trace_session_id"] == "session-1" - assert trace_task.kwargs["parent_trace_context"] == { - "parent_workflow_run_id": "outer-workflow-run-1", - "parent_node_execution_id": "outer-node-execution-1", - } - - trace_task.execute() - - assert captured["trace_type"] == TraceTaskName.WORKFLOW_TRACE - assert captured["external_trace_id"] == "trace" - assert captured["trace_session_id"] == "session-1" - assert captured["parent_trace_context"] == { - "parent_workflow_run_id": "outer-workflow-run-1", - "parent_node_execution_id": "outer-node-execution-1", - } - - def test_handle_graph_run_aborted_sets_status(self): - layer, exec_repo, _, _ = _make_layer() - layer.on_event(GraphRunStartedEvent()) - - layer.on_event(GraphRunAbortedEvent(reason=None, outputs={})) - - saved = exec_repo.synchronously_saved[-1] - assert saved.status == WorkflowExecutionStatus.STOPPED - assert saved.error_message - - def test_handle_node_started_and_retry(self): - layer, _, node_repo, _ = _make_layer(node_run_indices={"exec": 1}) - layer.on_event(GraphRunStartedEvent()) - - start_event = NodeRunStartedEvent( - id="exec", - node_id="node", - node_type=BuiltinNodeTypes.START, - node_title="Start", - start_at=_naive_utc_now(), - predecessor_node_id="prev", - ) - layer.on_event(start_event) - - assert node_repo.saved - assert "exec" in layer._node_execution_cache - assert node_repo.saved[-1].created_at == start_event.start_at - - retry_event = NodeRunRetryEvent( - id="exec", - node_id="node", - node_type=BuiltinNodeTypes.START, - node_title="Start", - start_at=_naive_utc_now(), - error="retry", - retry_index=1, - ) - layer.on_event(retry_event) - assert node_repo.saved_exec_data - - def test_agent_v2_caller_row_is_saved_synchronously_before_node_run(self): - layer, _, node_repo, _ = _make_layer(node_run_indices={"agent-exec": 1}) - layer.on_event(GraphRunStartedEvent()) - - layer.on_event( - NodeRunStartedEvent( - id="agent-exec", - node_id="agent-node", - node_type=BuiltinNodeTypes.AGENT, - node_version="2", - node_title="Agent", - start_at=_naive_utc_now(), - ) - ) - - assert [execution.id for execution in node_repo.synchronously_saved] == ["agent-exec"] - assert node_repo.saved == [] - - def test_retry_history_is_preserved_after_node_succeeds(self): - layer, _, node_repo, _ = _make_layer(node_run_indices={"exec": 1}) - layer.on_event(GraphRunStartedEvent()) - started_at = _naive_utc_now() - layer.on_event( - NodeRunStartedEvent( - id="exec", - node_id="node", - node_type=BuiltinNodeTypes.HTTP_REQUEST, - node_title="HTTP", - start_at=started_at, - ) - ) - - for retry_index in (1, 2): - layer.on_event( - NodeRunRetryEvent( - id="exec", - node_id="node", - node_type=BuiltinNodeTypes.HTTP_REQUEST, - node_title="HTTP", - start_at=started_at, - error=f"attempt {retry_index} failed", - retry_index=retry_index, - node_run_result=NodeRunResult( - inputs={"attempt": retry_index}, - process_data={"request": f"attempt-{retry_index}"}, - outputs={"status_code": 500, "body": f"failure-{retry_index}"}, - metadata={WorkflowNodeExecutionMetadataKey.ITERATION_ID: "iteration-1"}, - ), - ) - ) - - layer.on_event( - NodeRunSucceededEvent( - id="exec", - node_id="node", - node_type=BuiltinNodeTypes.HTTP_REQUEST, - start_at=started_at, - node_run_result=NodeRunResult( - inputs={"attempt": 3}, - process_data={"request": "successful-attempt"}, - outputs={"status_code": 200, "body": "ok"}, - metadata={}, - ), - ) - ) - - saved_execution = node_repo.saved_exec_data[-1] - assert saved_execution.status == WorkflowNodeExecutionStatus.SUCCEEDED - assert saved_execution.process_data["request"] == "successful-attempt" - retry_history = saved_execution.process_data["__dify_retry_history"] - assert [attempt["retry_index"] for attempt in retry_history] == [1, 2] - assert retry_history[0]["inputs"] == {"attempt": 1} - assert retry_history[0]["process_data"] == {"request": "attempt-1"} - assert retry_history[0]["outputs"] == {"status_code": 500, "body": "failure-1"} - assert retry_history[0]["error"] == "attempt 1 failed" - assert retry_history[0]["execution_metadata"] == {"iteration_id": "iteration-1"} - - @pytest.mark.parametrize( - ("event_type", "expected_status"), - [ - (NodeRunFailedEvent, WorkflowNodeExecutionStatus.FAILED), - (NodeRunExceptionEvent, WorkflowNodeExecutionStatus.EXCEPTION), - (GraphRunFailedEvent, WorkflowNodeExecutionStatus.FAILED), - (GraphRunAbortedEvent, WorkflowNodeExecutionStatus.FAILED), - ], - ) - def test_retry_history_is_preserved_after_terminal_error(self, event_type, expected_status): - layer, _, node_repo, _ = _make_layer(node_run_indices={"exec": 1}) - layer.on_event(GraphRunStartedEvent()) - started_at = _naive_utc_now() - layer.on_event( - NodeRunStartedEvent( - id="exec", - node_id="node", - node_type=BuiltinNodeTypes.LLM, - node_title="LLM", - start_at=started_at, - ) - ) - layer.on_event( - NodeRunRetryEvent( - id="exec", - node_id="node", - node_type=BuiltinNodeTypes.LLM, - node_title="LLM", - start_at=started_at, - error="retry failed", - retry_index=1, - node_run_result=NodeRunResult(outputs={"attempt": 1}), - ) - ) - - if event_type is GraphRunFailedEvent: - terminal_event = event_type(error="terminal failure") - elif event_type is GraphRunAbortedEvent: - terminal_event = event_type(reason="terminal failure") - else: - terminal_event = event_type( - id="exec", - node_id="node", - node_type=BuiltinNodeTypes.LLM, - start_at=started_at, - error="terminal failure", - node_run_result=NodeRunResult(process_data={"terminal": True}), - ) - node_repo.saved.clear() - layer.on_event(terminal_event) - - saved_execution = node_repo.saved[-1] - assert saved_execution.status == expected_status - assert saved_execution.error == "terminal failure" - assert saved_execution.finished_at is not None - if event_type in (NodeRunFailedEvent, NodeRunExceptionEvent): - assert saved_execution.process_data["terminal"] is True - assert saved_execution.process_data["__dify_retry_history"][0]["retry_index"] == 1 - - def test_pause_flush_persists_retry_history_synchronously(self): - layer, _, node_repo, _ = _make_layer(node_run_indices={"exec": 1}) - layer.on_event(GraphRunStartedEvent()) - started_at = _naive_utc_now() - layer.on_event( - NodeRunStartedEvent( - id="exec", node_id="node", node_type=BuiltinNodeTypes.LLM, node_title="LLM", start_at=started_at - ) - ) - layer.on_event( - NodeRunRetryEvent( - id="exec", - node_id="node", - node_type=BuiltinNodeTypes.LLM, - node_title="LLM", - start_at=started_at, - error="retry failed", - retry_index=1, - ) - ) - assert node_repo.synchronously_saved == [] - - layer.persist_unfinished_nodes_synchronously() - - saved_execution = node_repo.synchronously_saved[-1] - assert saved_execution.status == WorkflowNodeExecutionStatus.RETRY - assert saved_execution.finished_at is None - assert saved_execution.process_data["__dify_retry_history"][0]["error"] == "retry failed" - - def test_handle_node_pause_requested_skips_outputs(self): - layer, _, _, _ = _make_layer(node_run_indices={"exec": 1}) - layer.on_event(GraphRunStartedEvent()) - start_event = NodeRunStartedEvent( - id="exec", - node_id="node", - node_type=BuiltinNodeTypes.LLM, - node_title="LLM", - start_at=_naive_utc_now(), - ) - layer.on_event(start_event) - - domain_execution = layer._node_execution_cache["exec"] - domain_execution.inputs = {"old": True} - - result = NodeRunResult( - inputs={"new": True}, - outputs={"out": 1}, - process_data={ - "p": 1, - "workflow_agent_binding_id": "workflow-binding-1", - "workflow_tool_invocation_id": "tool-call-1", - }, - metadata={}, - ) - pause_event = NodeRunPauseRequestedEvent( - id="exec", - node_id="node", - node_type=BuiltinNodeTypes.LLM, - reason=SchedulingPause(message="pause"), - node_run_result=result, - ) - layer.on_event(pause_event) - - assert domain_execution.status == WorkflowNodeExecutionStatus.PAUSED - assert domain_execution.inputs == {"old": True} - assert domain_execution.process_data == { - "workflow_agent_binding_id": "workflow-binding-1", - "workflow_tool_invocation_id": "tool-call-1", - } - - def test_handle_node_retry_preserves_workflow_agent_binding_identity(self): - layer, _, _, _ = _make_layer(node_run_indices={"exec": 1}) - layer.on_event(GraphRunStartedEvent()) - started_at = _naive_utc_now() - layer.on_event( - NodeRunStartedEvent( - id="exec", - node_id="node", - node_type=BuiltinNodeTypes.AGENT, - node_title="Agent", - start_at=started_at, - ) - ) - - layer.on_event( - NodeRunRetryEvent( - id="exec", - node_id="node", - node_type=BuiltinNodeTypes.AGENT, - node_title="Agent", - start_at=started_at, - error="retry", - retry_index=1, - node_run_result=NodeRunResult( - process_data={ - "workflow_agent_binding_id": "workflow-binding-1", - "workflow_tool_invocation_id": "tool-call-1", - }, - ), - ) - ) - - assert layer._node_execution_cache["exec"].process_data["workflow_agent_binding_id"] == "workflow-binding-1" - assert layer._node_execution_cache["exec"].process_data["workflow_tool_invocation_id"] == "tool-call-1" - - def test_get_node_execution_raises_for_missing(self): - layer, _, _, _ = _make_layer() - with pytest.raises(ValueError, match="Node execution not found"): - layer._get_node_execution("missing") - - def test_get_workflow_execution_raises_when_uninitialized(self): - layer, _, _, _ = _make_layer() - - with pytest.raises(ValueError, match="workflow execution not initialized"): - layer._get_workflow_execution() diff --git a/api/tests/unit_tests/core/base/test_app_generator_tts_publisher.py b/api/tests/unit_tests/core/base/tts/test_app_generator_tts_publisher.py similarity index 99% rename from api/tests/unit_tests/core/base/test_app_generator_tts_publisher.py rename to api/tests/unit_tests/core/base/tts/test_app_generator_tts_publisher.py index fc6418b843a..91fa11e31f4 100644 --- a/api/tests/unit_tests/core/base/test_app_generator_tts_publisher.py +++ b/api/tests/unit_tests/core/base/tts/test_app_generator_tts_publisher.py @@ -19,7 +19,7 @@ from graphon.model_runtime.entities.model_entities import ModelPropertyKey from graphon.model_runtime.errors.invoke import InvokeBadRequestError from libs.datetime_utils import naive_utc_now from models.model import AppMode -from tests.unit_tests.audio_runtime_fixtures import AudioRuntimeObservations +from tests.helpers.unit.audio_runtime_fixtures import AudioRuntimeObservations @pytest.fixture(autouse=True) diff --git a/api/tests/unit_tests/core/base/tts/test_audio_mime.py b/api/tests/unit_tests/core/base/tts/test_audio_mime.py index 071be49ab76..88d05cbeaf9 100644 --- a/api/tests/unit_tests/core/base/tts/test_audio_mime.py +++ b/api/tests/unit_tests/core/base/tts/test_audio_mime.py @@ -13,7 +13,7 @@ from core.model_manager import ModelManager from core.plugin.entities.plugin_daemon import TTSAudioChunk from graphon.model_runtime.entities.model_entities import ModelPropertyKey, ModelType from graphon.model_runtime.errors.invoke import InvokeBadRequestError -from tests.unit_tests.audio_runtime_fixtures import AudioRuntimeObservations +from tests.helpers.unit.audio_runtime_fixtures import AudioRuntimeObservations def test_inspect_audio_stream_preserves_the_prefix_with_matching_chunk_mime_type() -> None: diff --git a/api/tests/unit_tests/core/callback_handler/test_agent_tool_callback_handler.py b/api/tests/unit_tests/core/callback_handler/test_agent_tool_callback_handler.py index 9d22e2f633f..6245b1ace23 100644 --- a/api/tests/unit_tests/core/callback_handler/test_agent_tool_callback_handler.py +++ b/api/tests/unit_tests/core/callback_handler/test_agent_tool_callback_handler.py @@ -5,7 +5,7 @@ from pytest_mock import MockerFixture import core.callback_handler.agent_tool_callback_handler as module from core.callback_handler.agent_tool_callback_handler import DifyAgentCallbackHandler -from tests.unit_tests.config_override import apply_config_overrides +from tests.helpers.unit.config_override import apply_config_overrides # ----------------------------- # Fixtures diff --git a/api/tests/unit_tests/core/datasource/__base/test_datasource_plugin.py b/api/tests/unit_tests/core/datasource/__base/test_datasource_plugin.py index 0cf48bb30d2..97391de88b7 100644 --- a/api/tests/unit_tests/core/datasource/__base/test_datasource_plugin.py +++ b/api/tests/unit_tests/core/datasource/__base/test_datasource_plugin.py @@ -3,7 +3,7 @@ from collections.abc import Callable from core.datasource.__base.datasource_plugin import DatasourcePlugin from core.datasource.__base.datasource_runtime import DatasourceRuntime from core.datasource.entities.datasource_entities import DatasourceProviderType -from tests.unit_tests.core.datasource.factories import datasource_entity +from tests.helpers.unit.core.datasource.factories import datasource_entity class ConcreteDatasourcePlugin(DatasourcePlugin): diff --git a/api/tests/unit_tests/core/datasource/__base/test_datasource_provider.py b/api/tests/unit_tests/core/datasource/__base/test_datasource_provider.py index 1b26decb945..babe0e6597b 100644 --- a/api/tests/unit_tests/core/datasource/__base/test_datasource_provider.py +++ b/api/tests/unit_tests/core/datasource/__base/test_datasource_provider.py @@ -10,7 +10,7 @@ from core.datasource.entities.datasource_entities import ( from core.entities.provider_entities import ProviderConfig, ProviderConfigType from core.tools.entities.common_entities import I18nObject from core.tools.errors import ToolProviderCredentialValidationError -from tests.unit_tests.core.datasource.factories import provider_entity +from tests.helpers.unit.core.datasource.factories import provider_entity class ConcreteDatasourcePluginProviderController(DatasourcePluginProviderController): diff --git a/api/tests/unit_tests/core/datasource/entities/test_datasource_entities.py b/api/tests/unit_tests/core/datasource/entities/test_datasource_entities.py index 74c3147bf28..988183bc4e7 100644 --- a/api/tests/unit_tests/core/datasource/entities/test_datasource_entities.py +++ b/api/tests/unit_tests/core/datasource/entities/test_datasource_entities.py @@ -1,5 +1,9 @@ -import pytest +from collections.abc import Generator +import pytest +from pytest_mock import MockerFixture + +from core.datasource.datasource_manager import DatasourceManager from core.datasource.entities.datasource_entities import ( DatasourceEntity, DatasourceIdentity, @@ -30,6 +34,8 @@ from core.datasource.entities.datasource_entities import ( ) from core.tools.entities.common_entities import I18nObject from core.tools.entities.tool_entities import ToolLabelEnum +from graphon.node_events import StreamCompletedEvent +from graphon.runtime import VariablePool def test_online_drive_remote_metadata_survives_runtime_serialization(): @@ -393,3 +399,42 @@ def test_datasource_message(): msg_json = DatasourceMessage(type="json", message={"json_object": {"k": "v"}}) assert msg_json.message.json_object == {"k": "v"} + + +"""Unit coverage for datasource event-stream accumulation.""" + + +def _gen_var_stream() -> Generator[DatasourceMessage]: + # produce a streamed variable "a"="xy" + yield DatasourceMessage( + type=DatasourceMessage.MessageType.VARIABLE, + message=DatasourceMessage.VariableMessage(variable_name="a", variable_value="x", stream=True), + meta=None, + ) + yield DatasourceMessage( + type=DatasourceMessage.MessageType.VARIABLE, + message=DatasourceMessage.VariableMessage(variable_name="a", variable_value="y", stream=True), + meta=None, + ) + + +def test_stream_node_events_accumulates_variables(mocker: MockerFixture) -> None: + mocker.patch.object(DatasourceManager, "stream_online_results", return_value=_gen_var_stream()) + events = list( + DatasourceManager.stream_node_events( + node_id="A", + user_id="u", + datasource_name="ds", + datasource_type="online_document", + provider_id="p/x", + tenant_id="t", + parameters_for_log={}, + datasource_info={"user_id": "u"}, + variable_pool=VariablePool(), + datasource_param=None, + online_drive_request=None, + credentials={}, + ) + ) + assert isinstance(events[-1], StreamCompletedEvent) + assert events[-1].node_run_result.outputs["a"] == "xy" diff --git a/api/tests/unit_tests/core/datasource/local_file/test_local_file_plugin.py b/api/tests/unit_tests/core/datasource/local_file/test_local_file_plugin.py index 1ce21c6b5f2..23ce89ffd98 100644 --- a/api/tests/unit_tests/core/datasource/local_file/test_local_file_plugin.py +++ b/api/tests/unit_tests/core/datasource/local_file/test_local_file_plugin.py @@ -3,7 +3,7 @@ from core.datasource.entities.datasource_entities import ( DatasourceProviderType, ) from core.datasource.local_file.local_file_plugin import LocalFileDatasourcePlugin -from tests.unit_tests.core.datasource.factories import datasource_entity +from tests.helpers.unit.core.datasource.factories import datasource_entity class TestLocalFileDatasourcePlugin: diff --git a/api/tests/unit_tests/core/datasource/local_file/test_local_file_provider.py b/api/tests/unit_tests/core/datasource/local_file/test_local_file_provider.py index aa6ca328f18..c7ff966f731 100644 --- a/api/tests/unit_tests/core/datasource/local_file/test_local_file_provider.py +++ b/api/tests/unit_tests/core/datasource/local_file/test_local_file_provider.py @@ -5,7 +5,7 @@ from core.datasource.entities.datasource_entities import ( ) from core.datasource.local_file.local_file_plugin import LocalFileDatasourcePlugin from core.datasource.local_file.local_file_provider import LocalFileDatasourcePluginProviderController -from tests.unit_tests.core.datasource.factories import datasource_entity, provider_entity +from tests.helpers.unit.core.datasource.factories import datasource_entity, provider_entity class TestLocalFileDatasourcePluginProviderController: diff --git a/api/tests/unit_tests/core/datasource/online_document/test_online_document_plugin.py b/api/tests/unit_tests/core/datasource/online_document/test_online_document_plugin.py index b8f4b1c91cb..7b1f0c4130a 100644 --- a/api/tests/unit_tests/core/datasource/online_document/test_online_document_plugin.py +++ b/api/tests/unit_tests/core/datasource/online_document/test_online_document_plugin.py @@ -8,7 +8,7 @@ from core.datasource.entities.datasource_entities import ( OnlineDocumentPagesMessage, ) from core.datasource.online_document.online_document_plugin import OnlineDocumentDatasourcePlugin -from tests.unit_tests.core.datasource.factories import datasource_entity +from tests.helpers.unit.core.datasource.factories import datasource_entity class TestOnlineDocumentDatasourcePlugin: diff --git a/api/tests/unit_tests/core/datasource/online_document/test_online_document_provider.py b/api/tests/unit_tests/core/datasource/online_document/test_online_document_provider.py index f09e10adcc3..1365e0be1ae 100644 --- a/api/tests/unit_tests/core/datasource/online_document/test_online_document_provider.py +++ b/api/tests/unit_tests/core/datasource/online_document/test_online_document_provider.py @@ -1,88 +1,1886 @@ +import json +from collections.abc import Callable, Iterator, Mapping +from unittest.mock import Mock, patch +from uuid import uuid4 + +import httpx import pytest +from sqlalchemy.engine import Engine +from sqlalchemy.orm import Session, sessionmaker from core.datasource.entities.datasource_entities import ( + DatasourceEntity, + DatasourceIdentity, + DatasourceProviderEntityWithPlugin, + DatasourceProviderIdentity, DatasourceProviderType, ) from core.datasource.online_document.online_document_plugin import OnlineDocumentDatasourcePlugin -from core.datasource.online_document.online_document_provider import OnlineDocumentDatasourcePluginProviderController -from tests.unit_tests.core.datasource.factories import datasource_entity, provider_entity +from core.datasource.online_document.online_document_provider import ( + OnlineDocumentDatasourcePluginProviderController, +) +from core.rag.extractor import notion_extractor as notion_extractor_module +from core.rag.extractor.entity.extract_setting import StoredDocumentExtractionInput +from core.rag.extractor.notion_extractor import NotionExtractor +from core.rag.models.document import Document +from core.tools.entities.common_entities import I18nObject +from extensions.application_services.data_sources import build_data_source_credentials +from models.base import TypeBase +from models.dataset import Document as DocumentModel +from models.enums import DataSourceType, DocumentCreatedFrom +from tests.helpers.unit.core.datasource.factories import datasource_entity, provider_entity -class TestOnlineDocumentDatasourcePluginProviderController: - def test_init(self) -> None: - # Arrange - entity = provider_entity(DatasourceProviderType.ONLINE_DOCUMENT) - plugin_id = "test_plugin_id" - plugin_unique_identifier = "test_plugin_uid" - tenant_id = "test_tenant_id" +class TestOnlineDocumentProviderCases: + class TestOnlineDocumentDatasourcePluginProviderController: + def test_init(self) -> None: + # Arrange + entity = provider_entity(DatasourceProviderType.ONLINE_DOCUMENT) + plugin_id = "test_plugin_id" + plugin_unique_identifier = "test_plugin_uid" + tenant_id = "test_tenant_id" - # Act - controller = OnlineDocumentDatasourcePluginProviderController( - entity=entity, - plugin_id=plugin_id, - plugin_unique_identifier=plugin_unique_identifier, - tenant_id=tenant_id, + # Act + controller = OnlineDocumentDatasourcePluginProviderController( + entity=entity, + plugin_id=plugin_id, + plugin_unique_identifier=plugin_unique_identifier, + tenant_id=tenant_id, + ) + + # Assert + assert controller.entity == entity + assert controller.plugin_id == plugin_id + assert controller.plugin_unique_identifier == plugin_unique_identifier + assert controller.tenant_id == tenant_id + + def test_provider_type(self) -> None: + # Arrange + entity = provider_entity(DatasourceProviderType.ONLINE_DOCUMENT) + controller = OnlineDocumentDatasourcePluginProviderController( + entity=entity, plugin_id="test", plugin_unique_identifier="test", tenant_id="test" + ) + + # Assert + assert controller.provider_type == DatasourceProviderType.ONLINE_DOCUMENT + + def test_get_datasource_success(self) -> None: + # Arrange + source = datasource_entity("target_datasource") + + entity = provider_entity(DatasourceProviderType.ONLINE_DOCUMENT) + entity.datasources = [source] + entity.identity.icon = "test_icon" + + plugin_unique_identifier = "test_plugin_uid" + tenant_id = "test_tenant_id" + + controller = OnlineDocumentDatasourcePluginProviderController( + entity=entity, + plugin_id="test_plugin_id", + plugin_unique_identifier=plugin_unique_identifier, + tenant_id=tenant_id, + ) + + # Act + result = controller.get_datasource("target_datasource") + + # Assert + assert isinstance(result, OnlineDocumentDatasourcePlugin) + assert result.entity == source + assert result.tenant_id == tenant_id + assert result.icon == "test_icon" + assert result.plugin_unique_identifier == plugin_unique_identifier + assert result.runtime.tenant_id == tenant_id + + def test_get_datasource_not_found(self) -> None: + # Arrange + source = datasource_entity("other_datasource") + + entity = provider_entity(DatasourceProviderType.ONLINE_DOCUMENT) + entity.datasources = [source] + + controller = OnlineDocumentDatasourcePluginProviderController( + entity=entity, + plugin_id="test_plugin_id", + plugin_unique_identifier="test_plugin_uid", + tenant_id="test_tenant_id", + ) + + # Act & Assert + with pytest.raises(ValueError, match="Datasource with name missing_datasource not found"): + controller.get_datasource("missing_datasource") + + +"""Comprehensive unit tests for Notion datasource provider. + +This test module covers all aspects of the Notion provider including: +- Notion API integration with proper authentication +- Page retrieval (single pages and databases) +- Block content parsing (headings, paragraphs, tables, nested blocks) +- Authentication handling (OAuth tokens, integration tokens, credential management) +- Error handling for API failures +- Pagination handling for large datasets +- Last edited time tracking + +All tests use mocking to avoid external dependencies and ensure fast, reliable execution. +Tests follow the Arrange-Act-Assert pattern for clarity. +""" + + +def _extraction_input(document: DocumentModel) -> StoredDocumentExtractionInput: + """Detach the values the extractor needs, as the indexing source adapter does.""" + return StoredDocumentExtractionInput( + id=document.id, + tenant_id=document.tenant_id, + dataset_id=document.dataset_id, + data_source_info_dict=document.data_source_info_dict, + ) + + +class TestNotionProviderCases: + @staticmethod + @pytest.fixture + def database(sqlite_engine: Engine, monkeypatch: pytest.MonkeyPatch) -> Iterator[Session]: + """Bind a real session for Notion document metadata persistence.""" + + TypeBase.metadata.create_all(sqlite_engine, tables=[TypeBase.metadata.tables[DocumentModel.__tablename__]]) + with Session(sqlite_engine, expire_on_commit=False) as session: + monkeypatch.setattr(notion_extractor_module.db, "session", session) + yield session + + @staticmethod + @pytest.fixture + def persisted_document(database: Session) -> DocumentModel: + document = DocumentModel( + id=str(uuid4()), + tenant_id=str(uuid4()), + dataset_id=str(uuid4()), + position=1, + data_source_type=DataSourceType.NOTION_IMPORT, + data_source_info=json.dumps({"last_edited_time": "2024-01-01T00:00:00.000Z"}), + batch="batch", + name="Notion page", + created_from=DocumentCreatedFrom.WEB, + created_by=str(uuid4()), ) + database.add(document) + database.commit() + return document - # Assert - assert controller.entity == entity - assert controller.plugin_id == plugin_id - assert controller.plugin_unique_identifier == plugin_unique_identifier - assert controller.tenant_id == tenant_id + class TestNotionExtractorAuthentication: + """Tests for Notion authentication handling. - def test_provider_type(self) -> None: - # Arrange - entity = provider_entity(DatasourceProviderType.ONLINE_DOCUMENT) - controller = OnlineDocumentDatasourcePluginProviderController( - entity=entity, plugin_id="test", plugin_unique_identifier="test", tenant_id="test" + Covers: + - OAuth token authentication + - Integration token fallback + - Credential retrieval from database + - Missing credential error handling + """ + + @pytest.fixture + def mock_document_model(self) -> StoredDocumentExtractionInput: + """Create detached document values with persisted-format Notion metadata.""" + return StoredDocumentExtractionInput( + id="test-doc-id", + tenant_id="tenant-789", + dataset_id="dataset-id", + data_source_info_dict={"last_edited_time": "2024-01-01T00:00:00.000Z"}, + ) + + def test_init_with_explicit_token(self, mock_document_model: StoredDocumentExtractionInput) -> None: + """Test NotionExtractor initialization with explicit access token.""" + # Arrange & Act + extractor = NotionExtractor( + notion_workspace_id="workspace-123", + notion_obj_id="page-456", + notion_page_type="page", + tenant_id="tenant-789", + notion_access_token="explicit-token-abc", + document_model=mock_document_model, + notion_token_loader=lambda: "token", + ) + + # Assert + assert extractor._notion_access_token == "explicit-token-abc" + assert extractor._notion_workspace_id == "workspace-123" + assert extractor._notion_obj_id == "page-456" + assert extractor._notion_page_type == "page" + + @patch("services.data_source.provider_service.DatasourceProviderService.get_datasource_credentials") + def test_init_with_credential_id( + self, + get_credentials: Mock, + mock_document_model: StoredDocumentExtractionInput, + sqlite_session_factory: sessionmaker[Session], + ) -> None: + """Test NotionExtractor initialization with credential ID retrieval.""" + # Arrange + providers = build_data_source_credentials(database_client=sqlite_session_factory).providers + get_credentials.return_value = {"integration_secret": "credential-token-xyz"} + + # Act + extractor = NotionExtractor( + notion_workspace_id="workspace-123", + notion_obj_id="page-456", + notion_page_type="page", + tenant_id="tenant-789", + credential_id="cred-123", + document_model=mock_document_model, + notion_token_loader=lambda: providers.get_stored_notion_access_token( + tenant_id="tenant-789", credential_id="cred-123" + ), + ) + + # Assert + assert extractor._notion_access_token == "credential-token-xyz" + get_credentials.assert_called_once_with( + tenant_id="tenant-789", + credential_id="cred-123", + provider="notion_datasource", + plugin_id="langgenius/notion_datasource", + ) + + @patch("services.data_source.provider_service.DatasourceProviderService.get_datasource_credentials") + def test_init_with_integration_token_fallback( + self, + mock_get_token: Mock, + mock_document_model: StoredDocumentExtractionInput, + config_overrides: Callable[..., None], + sqlite_session_factory: sessionmaker[Session], + ) -> None: + """Test NotionExtractor falls back to integration token when credential not found.""" + # Arrange + mock_get_token.side_effect = Exception("No credential id found") + config_overrides(NOTION_INTEGRATION_TOKEN="integration-token-fallback") + providers = build_data_source_credentials(database_client=sqlite_session_factory).providers + + # Act + extractor = NotionExtractor( + notion_workspace_id="workspace-123", + notion_obj_id="page-456", + notion_page_type="page", + tenant_id="tenant-789", + credential_id=None, + document_model=mock_document_model, + notion_token_loader=lambda: providers.get_stored_notion_access_token( + tenant_id="tenant-789", credential_id=None + ), + ) + + # Assert + assert extractor._notion_access_token == "integration-token-fallback" + + @patch("services.data_source.provider_service.DatasourceProviderService.get_datasource_credentials") + def test_init_missing_credentials_raises_error( + self, + mock_get_token: Mock, + mock_document_model: StoredDocumentExtractionInput, + config_overrides: Callable[..., None], + sqlite_session_factory: sessionmaker[Session], + ) -> None: + """Test NotionExtractor raises error when no credentials available.""" + # Arrange + mock_get_token.side_effect = Exception("No credential id found") + config_overrides(NOTION_INTEGRATION_TOKEN=None) + providers = build_data_source_credentials(database_client=sqlite_session_factory).providers + + # Act & Assert + with pytest.raises( + ValueError, + match=("Must specify `integration_token` or set environment variable `NOTION_INTEGRATION_TOKEN`\\."), + ) as exc_info: + NotionExtractor( + notion_workspace_id="workspace-123", + notion_obj_id="page-456", + notion_page_type="page", + tenant_id="tenant-789", + credential_id=None, + document_model=mock_document_model, + notion_token_loader=lambda: providers.get_stored_notion_access_token( + tenant_id="tenant-789", credential_id=None + ), + ) + assert "Must specify `integration_token`" in str(exc_info.value) + + class TestNotionExtractorPageRetrieval: + """Tests for Notion page retrieval functionality. + + Covers: + - Single page retrieval + - Database page retrieval with pagination + - Block content extraction + - Nested block handling + """ + + @pytest.fixture + def extractor(self) -> NotionExtractor: + """Create a NotionExtractor instance for testing.""" + return NotionExtractor( + notion_workspace_id="workspace-123", + notion_obj_id="page-456", + notion_page_type="page", + tenant_id="tenant-789", + notion_access_token="test-token", + notion_token_loader=lambda: "token", + ) + + def _create_response(self, data: Mapping[str, object], status_code: int = 200) -> httpx.Response: + """Build an HTTP response with serialized JSON for the extractor to parse.""" + return httpx.Response(status_code, json=data) + + def _create_block( + self, block_id: str, block_type: str, text_content: str, has_children: bool = False + ) -> dict[str, object]: + """Helper to create a Notion block structure.""" + return { + "object": "block", + "id": block_id, + "type": block_type, + "has_children": has_children, + block_type: { + "rich_text": [ + { + "type": "text", + "text": {"content": text_content}, + "plain_text": text_content, + } + ] + }, + } + + @patch("httpx.request") + def test_get_notion_block_data_simple_page(self, mock_request: Mock, extractor: NotionExtractor) -> None: + """Test retrieving simple page with basic blocks.""" + # Arrange + mock_data = { + "object": "list", + "results": [ + self._create_block("block-1", "paragraph", "First paragraph"), + self._create_block("block-2", "paragraph", "Second paragraph"), + ], + "next_cursor": None, + "has_more": False, + } + mock_request.return_value = self._create_response(mock_data) + + # Act + result = extractor._get_notion_block_data("page-456") + + # Assert + assert len(result) == 2 + assert "First paragraph" in result[0] + assert "Second paragraph" in result[1] + mock_request.assert_called_once() + + @patch("httpx.request") + def test_get_notion_block_data_with_headings(self, mock_request: Mock, extractor: NotionExtractor) -> None: + """Test retrieving page with heading blocks.""" + # Arrange + mock_data = { + "object": "list", + "results": [ + self._create_block("block-1", "heading_1", "Main Title"), + self._create_block("block-2", "heading_2", "Subtitle"), + self._create_block("block-3", "paragraph", "Content text"), + self._create_block("block-4", "heading_3", "Sub-subtitle"), + ], + "next_cursor": None, + "has_more": False, + } + mock_request.return_value = self._create_response(mock_data) + + # Act + result = extractor._get_notion_block_data("page-456") + + # Assert + assert len(result) == 4 + assert "# Main Title" in result[0] + assert "## Subtitle" in result[1] + assert "Content text" in result[2] + assert "### Sub-subtitle" in result[3] + + @patch("httpx.request") + def test_get_notion_block_data_with_pagination(self, mock_request: Mock, extractor: NotionExtractor) -> None: + """Test retrieving page with paginated results.""" + # Arrange + first_page = { + "object": "list", + "results": [self._create_block("block-1", "paragraph", "First page content")], + "next_cursor": "cursor-abc", + "has_more": True, + } + second_page = { + "object": "list", + "results": [self._create_block("block-2", "paragraph", "Second page content")], + "next_cursor": None, + "has_more": False, + } + mock_request.side_effect = [ + self._create_response(first_page), + self._create_response(second_page), + ] + + # Act + result = extractor._get_notion_block_data("page-456") + + # Assert + assert len(result) == 2 + assert "First page content" in result[0] + assert "Second page content" in result[1] + assert mock_request.call_count == 2 + + @patch("httpx.request") + def test_get_notion_block_data_with_nested_blocks(self, mock_request: Mock, extractor: NotionExtractor) -> None: + """Test retrieving page with nested block structure.""" + # Arrange + # First call returns parent blocks + parent_data = { + "object": "list", + "results": [ + self._create_block("block-1", "paragraph", "Parent block", has_children=True), + ], + "next_cursor": None, + "has_more": False, + } + # Second call returns child blocks + child_data = { + "object": "list", + "results": [ + self._create_block("block-child-1", "paragraph", "Child block"), + ], + "next_cursor": None, + "has_more": False, + } + mock_request.side_effect = [ + self._create_response(parent_data), + self._create_response(child_data), + ] + + # Act + result = extractor._get_notion_block_data("page-456") + + # Assert + assert len(result) == 1 + assert "Parent block" in result[0] + assert "Child block" in result[0] + assert mock_request.call_count == 2 + + @patch("httpx.request") + def test_get_notion_block_data_error_handling(self, mock_request: Mock, extractor: NotionExtractor) -> None: + """Test error handling for failed API requests.""" + # Arrange + mock_request.return_value = self._create_response({}, status_code=404) + + # Act & Assert + with pytest.raises(ValueError, match="Error fetching Notion block data: \\{\\}") as exc_info: + extractor._get_notion_block_data("page-456") + assert "Error fetching Notion block data" in str(exc_info.value) + + @patch("httpx.request") + def test_get_notion_block_data_invalid_response(self, mock_request: Mock, extractor: NotionExtractor) -> None: + """Test handling of invalid API response structure.""" + # Arrange + mock_request.return_value = self._create_response({"invalid": "structure"}) + + # Act & Assert + with pytest.raises(ValueError, match="Error fetching Notion block data") as exc_info: + extractor._get_notion_block_data("page-456") + assert "Error fetching Notion block data" in str(exc_info.value) + + @patch("httpx.request") + def test_get_notion_block_data_http_error(self, mock_request: Mock, extractor: NotionExtractor) -> None: + """Test handling of HTTP errors during request.""" + # Arrange + mock_request.side_effect = httpx.HTTPError("Network error") + + # Act & Assert + with pytest.raises(ValueError, match="Error fetching Notion block data") as exc_info: + extractor._get_notion_block_data("page-456") + assert "Error fetching Notion block data" in str(exc_info.value) + + class TestNotionExtractorDatabaseRetrieval: + """Tests for Notion database retrieval functionality. + + Covers: + - Database query with pagination + - Property extraction (title, rich_text, select, multi_select, etc.) + - Row formatting + - Empty database handling + """ + + @pytest.fixture + def extractor(self) -> NotionExtractor: + """Create a NotionExtractor instance for testing.""" + return NotionExtractor( + notion_workspace_id="workspace-123", + notion_obj_id="database-789", + notion_page_type="database", + tenant_id="tenant-789", + notion_access_token="test-token", + notion_token_loader=lambda: "token", + ) + + def _create_database_page( + self, page_id: str, properties: Mapping[str, Mapping[str, object]] + ) -> dict[str, object]: + """Helper to create a database page structure.""" + formatted_properties = {} + for prop_name, prop_data in properties.items(): + prop_type = prop_data["type"] + assert isinstance(prop_type, str) + formatted_properties[prop_name] = {"type": prop_type, prop_type: prop_data["value"]} + return { + "object": "page", + "id": page_id, + "properties": formatted_properties, + "url": f"https://notion.so/{page_id}", + } + + @patch("httpx.post") + def test_get_notion_database_data_simple(self, mock_post: Mock, extractor: NotionExtractor) -> None: + """Test retrieving simple database with basic properties.""" + # Arrange + mock_response_data = { + "object": "list", + "results": [ + self._create_database_page( + "page-1", + { + "Title": {"type": "title", "value": [{"plain_text": "Task 1"}]}, + "Status": {"type": "select", "value": {"name": "In Progress"}}, + }, + ), + self._create_database_page( + "page-2", + { + "Title": {"type": "title", "value": [{"plain_text": "Task 2"}]}, + "Status": {"type": "select", "value": {"name": "Done"}}, + }, + ), + ], + "has_more": False, + "next_cursor": None, + } + mock_response = httpx.Response(200, json=mock_response_data) + mock_post.return_value = mock_response + + # Act + result = extractor._get_notion_database_data("database-789") + + # Assert + assert len(result) == 1 + content = result[0].page_content + assert "Title:Task 1" in content + assert "Status:In Progress" in content + assert "Title:Task 2" in content + assert "Status:Done" in content + + @patch("httpx.post") + def test_get_notion_database_data_with_pagination(self, mock_post: Mock, extractor: NotionExtractor) -> None: + """Test retrieving database with paginated results.""" + # Arrange + first_response_data = { + "object": "list", + "results": [ + self._create_database_page( + "page-1", {"Title": {"type": "title", "value": [{"plain_text": "Page 1"}]}} + ), + ], + "has_more": True, + "next_cursor": "cursor-xyz", + } + first_response = httpx.Response(200, json=first_response_data) + second_response_data = { + "object": "list", + "results": [ + self._create_database_page( + "page-2", {"Title": {"type": "title", "value": [{"plain_text": "Page 2"}]}} + ), + ], + "has_more": False, + "next_cursor": None, + } + second_response = httpx.Response(200, json=second_response_data) + mock_post.side_effect = [first_response, second_response] + + # Act + result = extractor._get_notion_database_data("database-789") + + # Assert + assert len(result) == 1 + content = result[0].page_content + assert "Title:Page 1" in content + assert "Title:Page 2" in content + assert mock_post.call_count == 2 + + @patch("httpx.post") + def test_get_notion_database_data_multi_select(self, mock_post: Mock, extractor: NotionExtractor) -> None: + """Test database with multi_select property type.""" + # Arrange + mock_response_data = { + "object": "list", + "results": [ + self._create_database_page( + "page-1", + { + "Title": {"type": "title", "value": [{"plain_text": "Project"}]}, + "Tags": { + "type": "multi_select", + "value": [{"name": "urgent"}, {"name": "frontend"}], + }, + }, + ), + ], + "has_more": False, + "next_cursor": None, + } + mock_response = httpx.Response(200, json=mock_response_data) + mock_post.return_value = mock_response + + # Act + result = extractor._get_notion_database_data("database-789") + + # Assert + assert len(result) == 1 + content = result[0].page_content + assert "Title:Project" in content + assert "Tags:" in content + + @patch("httpx.post") + def test_get_notion_database_data_empty_properties(self, mock_post: Mock, extractor: NotionExtractor) -> None: + """Test database with empty property values.""" + # Arrange + mock_response_data = { + "object": "list", + "results": [ + self._create_database_page( + "page-1", + { + "Title": {"type": "title", "value": []}, + "Status": {"type": "select", "value": None}, + }, + ), + ], + "has_more": False, + "next_cursor": None, + } + mock_response = httpx.Response(200, json=mock_response_data) + mock_post.return_value = mock_response + + # Act + result = extractor._get_notion_database_data("database-789") + + # Assert + assert len(result) == 1 + # Empty properties should be filtered out + content = result[0].page_content + assert "Row Page URL:" in content + + @patch("httpx.post") + def test_get_notion_database_data_empty_results(self, mock_post: Mock, extractor: NotionExtractor) -> None: + """Test handling of empty database.""" + # Arrange + mock_response_data = { + "object": "list", + "results": list[object](), + "has_more": False, + "next_cursor": None, + } + mock_response = httpx.Response(200, json=mock_response_data) + mock_post.return_value = mock_response + + # Act + result = extractor._get_notion_database_data("database-789") + + # Assert + assert len(result) == 0 + + @patch("httpx.post") + def test_get_notion_database_data_missing_results(self, mock_post: Mock, extractor: NotionExtractor) -> None: + """Test handling of malformed API response.""" + # Arrange + mock_response_data = {"object": "list"} + mock_response = httpx.Response(200, json=mock_response_data) + mock_post.return_value = mock_response + + # Act + result = extractor._get_notion_database_data("database-789") + + # Assert + assert len(result) == 0 + + class TestNotionExtractorTableParsing: + """Tests for Notion table block parsing. + + Covers: + - Table header extraction + - Table row parsing + - Markdown table formatting + - Empty cell handling + """ + + @pytest.fixture + def extractor(self) -> NotionExtractor: + """Create a NotionExtractor instance for testing.""" + return NotionExtractor( + notion_workspace_id="workspace-123", + notion_obj_id="page-456", + notion_page_type="page", + tenant_id="tenant-789", + notion_access_token="test-token", + notion_token_loader=lambda: "token", + ) + + @patch("httpx.request") + def test_read_table_rows_simple(self, mock_request: Mock, extractor: NotionExtractor) -> None: + """Test reading simple table with headers and rows.""" + # Arrange + mock_data = { + "object": "list", + "results": [ + { + "object": "block", + "type": "table_row", + "table_row": { + "cells": [ + [{"text": {"content": "Name"}}], + [{"text": {"content": "Age"}}], + ] + }, + }, + { + "object": "block", + "type": "table_row", + "table_row": { + "cells": [ + [{"text": {"content": "Alice"}}], + [{"text": {"content": "30"}}], + ] + }, + }, + { + "object": "block", + "type": "table_row", + "table_row": { + "cells": [ + [{"text": {"content": "Bob"}}], + [{"text": {"content": "25"}}], + ] + }, + }, + ], + "next_cursor": None, + "has_more": False, + } + mock_request.return_value = httpx.Response(200, json=mock_data) + + # Act + result = extractor._read_table_rows("table-block-123") + + # Assert + assert "| Name | Age |" in result + assert "| --- | --- |" in result + assert "| Alice | 30 |" in result + assert "| Bob | 25 |" in result + + @patch("httpx.request") + def test_read_table_rows_with_empty_cells(self, mock_request: Mock, extractor: NotionExtractor) -> None: + """Test reading table with empty cells.""" + # Arrange + mock_data = { + "object": "list", + "results": [ + { + "object": "block", + "type": "table_row", + "table_row": {"cells": [[{"text": {"content": "Col1"}}], [{"text": {"content": "Col2"}}]]}, + }, + { + "object": "block", + "type": "table_row", + "table_row": {"cells": [[{"text": {"content": "Value1"}}], list[object]()]}, + }, + ], + "next_cursor": None, + "has_more": False, + } + mock_request.return_value = httpx.Response(200, json=mock_data) + + # Act + result = extractor._read_table_rows("table-block-123") + + # Assert + assert "| Col1 | Col2 |" in result + assert "| --- | --- |" in result + # Empty cells are handled by the table parsing logic + assert "Value1" in result + + @patch("httpx.request") + def test_read_table_rows_with_pagination(self, mock_request: Mock, extractor: NotionExtractor) -> None: + """Test reading table with paginated results.""" + # Arrange + first_page = { + "object": "list", + "results": [ + { + "object": "block", + "type": "table_row", + "table_row": {"cells": [[{"text": {"content": "Header"}}]]}, + }, + ], + "next_cursor": "cursor-abc", + "has_more": True, + } + second_page = { + "object": "list", + "results": [ + { + "object": "block", + "type": "table_row", + "table_row": {"cells": [[{"text": {"content": "Row1"}}]]}, + }, + ], + "next_cursor": None, + "has_more": False, + } + mock_request.side_effect = [httpx.Response(200, json=first_page), httpx.Response(200, json=second_page)] + + # Act + result = extractor._read_table_rows("table-block-123") + + # Assert + assert "| Header |" in result + assert mock_request.call_count == 2 + + class TestNotionExtractorLastEditedTime: + """Tests for last edited time tracking. + + Covers: + - Page last edited time retrieval + - Database last edited time retrieval + - Document model update + """ + + @pytest.fixture + def mock_document_model(self) -> StoredDocumentExtractionInput: + """Create detached document values with persisted-format Notion metadata.""" + return StoredDocumentExtractionInput( + id="test-doc-id", + tenant_id="tenant-789", + dataset_id="dataset-id", + data_source_info_dict={"last_edited_time": "2024-01-01T00:00:00.000Z"}, + ) + + @pytest.fixture + def extractor_page(self, mock_document_model: StoredDocumentExtractionInput) -> NotionExtractor: + """Create a NotionExtractor instance for page testing.""" + return NotionExtractor( + notion_workspace_id="workspace-123", + notion_obj_id="page-456", + notion_page_type="page", + tenant_id="tenant-789", + notion_access_token="test-token", + document_model=mock_document_model, + notion_token_loader=lambda: "token", + ) + + @pytest.fixture + def extractor_database(self, mock_document_model: StoredDocumentExtractionInput) -> NotionExtractor: + """Create a NotionExtractor instance for database testing.""" + return NotionExtractor( + notion_workspace_id="workspace-123", + notion_obj_id="database-789", + notion_page_type="database", + tenant_id="tenant-789", + notion_access_token="test-token", + document_model=mock_document_model, + notion_token_loader=lambda: "token", + ) + + @patch("httpx.request") + def test_get_notion_last_edited_time_page(self, mock_request: Mock, extractor_page: NotionExtractor) -> None: + """Test retrieving last edited time for a page.""" + # Arrange + mock_response_data = { + "object": "page", + "id": "page-456", + "last_edited_time": "2024-11-27T12:00:00.000Z", + } + mock_response = httpx.Response(200, json=mock_response_data) + mock_request.return_value = mock_response + + # Act + result = extractor_page.get_notion_last_edited_time() + + # Assert + assert result == "2024-11-27T12:00:00.000Z" + mock_request.assert_called_once() + call_args = mock_request.call_args + assert "pages/page-456" in call_args[0][1] + + @patch("httpx.request") + def test_get_notion_last_edited_time_database( + self, mock_request: Mock, extractor_database: NotionExtractor + ) -> None: + """Test retrieving last edited time for a database.""" + # Arrange + mock_response_data = { + "object": "database", + "id": "database-789", + "last_edited_time": "2024-11-27T15:30:00.000Z", + } + mock_response = httpx.Response(200, json=mock_response_data) + mock_request.return_value = mock_response + + # Act + result = extractor_database.get_notion_last_edited_time() + + # Assert + assert result == "2024-11-27T15:30:00.000Z" + mock_request.assert_called_once() + call_args = mock_request.call_args + assert "databases/database-789" in call_args[0][1] + + @patch("httpx.request") + def test_update_last_edited_time( + self, + mock_request: Mock, + extractor_page: NotionExtractor, + database: Session, + persisted_document: DocumentModel, + ) -> None: + """Test updating document model with last edited time.""" + # Arrange + mock_response_data = { + "object": "page", + "id": "page-456", + "last_edited_time": "2024-11-27T18:00:00.000Z", + } + mock_response = httpx.Response(200, json=mock_response_data) + mock_request.return_value = mock_response + + # Act + extractor_page.update_last_edited_time(_extraction_input(persisted_document)) + + # Assert + database.expire(persisted_document) + assert persisted_document.data_source_info_dict["last_edited_time"] == "2024-11-27T18:00:00.000Z" + + def test_update_last_edited_time_no_document(self, extractor_page: NotionExtractor) -> None: + """Test update_last_edited_time with None document model.""" + # Act & Assert - should not raise error + extractor_page.update_last_edited_time(None) + + class TestNotionExtractorIntegration: + """Integration tests for complete extraction workflow. + + Covers: + - Full page extraction workflow + - Full database extraction workflow + - Document creation + - Error handling in extract method + """ + + @pytest.fixture + def mock_document_model(self) -> StoredDocumentExtractionInput: + """Create detached document values with persisted-format Notion metadata.""" + return StoredDocumentExtractionInput( + id="test-doc-id", + tenant_id="tenant-789", + dataset_id="dataset-id", + data_source_info_dict={"last_edited_time": "2024-01-01T00:00:00.000Z"}, + ) + + @patch("httpx.request") + def test_extract_page_complete_workflow( + self, mock_request: Mock, database: Session, persisted_document: DocumentModel + ) -> None: + """Test complete page extraction workflow.""" + # Arrange + extractor = NotionExtractor( + notion_workspace_id="workspace-123", + notion_obj_id="page-456", + notion_page_type="page", + tenant_id="tenant-789", + notion_access_token="test-token", + document_model=_extraction_input(persisted_document), + notion_token_loader=lambda: "token", + ) + + # Mock last edited time request + last_edited_response_data = { + "object": "page", + "last_edited_time": "2024-11-27T20:00:00.000Z", + } + last_edited_response = httpx.Response(200, json=last_edited_response_data) + + # Mock block data request + block_response_data = { + "object": "list", + "results": [ + { + "object": "block", + "id": "block-1", + "type": "heading_1", + "has_children": False, + "heading_1": { + "rich_text": [{"type": "text", "text": {"content": "Test Page"}, "plain_text": "Test Page"}] + }, + }, + { + "object": "block", + "id": "block-2", + "type": "paragraph", + "has_children": False, + "paragraph": { + "rich_text": [ + {"type": "text", "text": {"content": "Test content"}, "plain_text": "Test content"} + ] + }, + }, + ], + "next_cursor": None, + "has_more": False, + } + block_response = httpx.Response(200, json=block_response_data) + + mock_request.side_effect = [last_edited_response, block_response] + + # Act + documents = extractor.extract() + + # Assert + assert len(documents) == 1 + assert isinstance(documents[0], Document) + assert "# Test Page" in documents[0].page_content + assert "Test content" in documents[0].page_content + database.expire(persisted_document) + assert persisted_document.data_source_info_dict["last_edited_time"] == "2024-11-27T20:00:00.000Z" + + @patch("httpx.post") + @patch("httpx.request") + def test_extract_database_complete_workflow( + self, + mock_request: Mock, + mock_post: Mock, + database: Session, + persisted_document: DocumentModel, + ) -> None: + """Test complete database extraction workflow.""" + # Arrange + extractor = NotionExtractor( + notion_workspace_id="workspace-123", + notion_obj_id="database-789", + notion_page_type="database", + tenant_id="tenant-789", + notion_access_token="test-token", + document_model=_extraction_input(persisted_document), + notion_token_loader=lambda: "token", + ) + + # Mock last edited time request + last_edited_response_data = { + "object": "database", + "last_edited_time": "2024-11-27T20:00:00.000Z", + } + last_edited_response = httpx.Response(200, json=last_edited_response_data) + mock_request.return_value = last_edited_response + + # Mock database query request + database_response_data = { + "object": "list", + "results": [ + { + "object": "page", + "id": "page-1", + "properties": { + "Name": {"type": "title", "title": [{"plain_text": "Item 1"}]}, + "Status": {"type": "select", "select": {"name": "Active"}}, + }, + "url": "https://notion.so/page-1", + } + ], + "has_more": False, + "next_cursor": None, + } + database_response = httpx.Response(200, json=database_response_data) + mock_post.return_value = database_response + + # Act + documents = extractor.extract() + + # Assert + assert len(documents) == 1 + assert isinstance(documents[0], Document) + assert "Name:Item 1" in documents[0].page_content + assert "Status:Active" in documents[0].page_content + database.expire(persisted_document) + assert persisted_document.data_source_info_dict["last_edited_time"] == "2024-11-27T20:00:00.000Z" + + def test_extract_invalid_page_type(self) -> None: + """Test extract with invalid page type.""" + # Arrange + extractor = NotionExtractor( + notion_workspace_id="workspace-123", + notion_obj_id="invalid-456", + notion_page_type="invalid_type", + tenant_id="tenant-789", + notion_access_token="test-token", + notion_token_loader=lambda: "token", + ) + + # Act & Assert + with pytest.raises(ValueError, match="notion page type not supported") as exc_info: + extractor.extract() + assert "notion page type not supported" in str(exc_info.value) + + class TestNotionExtractorReadBlock: + """Tests for nested block reading functionality. + + Covers: + - Recursive block reading + - Indentation handling + - Child page handling + """ + + @pytest.fixture + def extractor(self) -> NotionExtractor: + """Create a NotionExtractor instance for testing.""" + return NotionExtractor( + notion_workspace_id="workspace-123", + notion_obj_id="page-456", + notion_page_type="page", + tenant_id="tenant-789", + notion_access_token="test-token", + notion_token_loader=lambda: "token", + ) + + @patch("httpx.request") + def test_read_block_with_indentation(self, mock_request: Mock, extractor: NotionExtractor) -> None: + """Test reading nested blocks with proper indentation.""" + # Arrange + mock_data = { + "object": "list", + "results": [ + { + "object": "block", + "id": "block-1", + "type": "paragraph", + "has_children": False, + "paragraph": { + "rich_text": [ + {"type": "text", "text": {"content": "Nested content"}, "plain_text": "Nested content"} + ] + }, + } + ], + "next_cursor": None, + "has_more": False, + } + mock_request.return_value = httpx.Response(200, json=mock_data) + + # Act + result = extractor._read_block("block-parent", num_tabs=2) + + # Assert + assert "\t\tNested content" in result + + @patch("httpx.request") + def test_read_block_skip_child_page(self, mock_request: Mock, extractor: NotionExtractor) -> None: + """Test that child_page blocks don't recurse.""" + # Arrange + mock_data = { + "object": "list", + "results": [ + { + "object": "block", + "id": "block-1", + "type": "child_page", + "has_children": True, + "child_page": {"title": "Child Page"}, + } + ], + "next_cursor": None, + "has_more": False, + } + mock_request.return_value = httpx.Response(200, json=mock_data) + + # Act + extractor._read_block("block-parent") + + # Assert + # Should only be called once (no recursion for child_page) + assert mock_request.call_count == 1 + + class TestNotionProviderController: + """Tests for Notion datasource provider controller integration. + + Covers: + - Provider initialization + - Datasource retrieval + - Provider type verification + """ + + @pytest.fixture + def mock_entity(self) -> DatasourceProviderEntityWithPlugin: + """Create a validated Notion provider entity.""" + return DatasourceProviderEntityWithPlugin( + identity=DatasourceProviderIdentity( + author="langgenius", + name="notion_datasource", + icon="notion-icon.png", + description=I18nObject(en_US="Notion documents"), + label=I18nObject(en_US="Notion"), + ), + provider_type=DatasourceProviderType.ONLINE_DOCUMENT, + ) + + def test_provider_controller_initialization(self, mock_entity: DatasourceProviderEntityWithPlugin) -> None: + """Test OnlineDocumentDatasourcePluginProviderController initialization.""" + # Act + controller = OnlineDocumentDatasourcePluginProviderController( + entity=mock_entity, + plugin_id="langgenius/notion_datasource", + plugin_unique_identifier="notion-unique-id", + tenant_id="tenant-123", + ) + + # Assert + assert controller.plugin_id == "langgenius/notion_datasource" + assert controller.plugin_unique_identifier == "notion-unique-id" + assert controller.tenant_id == "tenant-123" + assert controller.provider_type == DatasourceProviderType.ONLINE_DOCUMENT + + def test_provider_controller_get_datasource(self, mock_entity: DatasourceProviderEntityWithPlugin) -> None: + """Test retrieving datasource from controller.""" + # Arrange + mock_datasource_entity = DatasourceEntity( + identity=DatasourceIdentity( + author="langgenius", + name="notion_datasource", + label=I18nObject(en_US="Notion"), + provider="notion_datasource", + ), + description=I18nObject(en_US="Notion documents"), + ) + mock_entity.datasources = [mock_datasource_entity] + + controller = OnlineDocumentDatasourcePluginProviderController( + entity=mock_entity, + plugin_id="langgenius/notion_datasource", + plugin_unique_identifier="notion-unique-id", + tenant_id="tenant-123", + ) + + # Act + datasource = controller.get_datasource("notion_datasource") + + # Assert + assert datasource is not None + assert datasource.tenant_id == "tenant-123" + + def test_provider_controller_datasource_not_found( + self, mock_entity: DatasourceProviderEntityWithPlugin + ) -> None: + """Test error when datasource not found.""" + # Arrange + mock_entity.datasources = [] + controller = OnlineDocumentDatasourcePluginProviderController( + entity=mock_entity, + plugin_id="langgenius/notion_datasource", + plugin_unique_identifier="notion-unique-id", + tenant_id="tenant-123", + ) + + # Act & Assert + with pytest.raises(ValueError, match="Datasource with name nonexistent_datasource not found") as exc_info: + controller.get_datasource("nonexistent_datasource") + assert "not found" in str(exc_info.value) + + class TestNotionExtractorAdvancedBlockTypes: + """Tests for advanced Notion block types and edge cases. + + Covers: + - Various block types (code, quote, lists, toggle, callout) + - Empty blocks + - Multiple rich text elements + - Mixed block types in realistic scenarios + """ + + @pytest.fixture + def extractor(self) -> NotionExtractor: + """Create a NotionExtractor instance for testing. + + Returns: + NotionExtractor: Configured extractor with test credentials + """ + return NotionExtractor( + notion_workspace_id="workspace-123", + notion_obj_id="page-456", + notion_page_type="page", + tenant_id="tenant-789", + notion_access_token="test-token", + notion_token_loader=lambda: "token", + ) + + def _create_block_with_rich_text( + self, block_id: str, block_type: str, rich_text_items: list[str], has_children: bool = False + ) -> dict[str, object]: + """Helper to create a Notion block with multiple rich text elements. + + Args: + block_id: Unique identifier for the block + block_type: Type of block (paragraph, heading_1, etc.) + rich_text_items: List of text content strings + has_children: Whether the block has child blocks + + Returns: + dict: Notion block structure with rich text elements + """ + rich_text_array = [ + {"type": "text", "text": {"content": text}, "plain_text": text} for text in rich_text_items + ] + return { + "object": "block", + "id": block_id, + "type": block_type, + "has_children": has_children, + block_type: {"rich_text": rich_text_array}, + } + + @patch("httpx.request") + def test_get_notion_block_data_with_list_blocks(self, mock_request: Mock, extractor: NotionExtractor) -> None: + """Test retrieving page with bulleted and numbered list items. + + Both list types should be extracted with their content. + """ + # Arrange + mock_data = { + "object": "list", + "results": [ + self._create_block_with_rich_text("block-1", "bulleted_list_item", ["Bullet item"]), + self._create_block_with_rich_text("block-2", "numbered_list_item", ["Numbered item"]), + ], + "next_cursor": None, + "has_more": False, + } + mock_request.return_value = httpx.Response(200, json=mock_data) + + # Act + result = extractor._get_notion_block_data("page-456") + + # Assert + assert len(result) == 2 + assert "Bullet item" in result[0] + assert "Numbered item" in result[1] + + @patch("httpx.request") + def test_get_notion_block_data_with_special_blocks( + self, mock_request: Mock, extractor: NotionExtractor + ) -> None: + """Test retrieving page with code, quote, and callout blocks. + + Special block types should preserve their content correctly. + """ + # Arrange + mock_data = { + "object": "list", + "results": [ + self._create_block_with_rich_text("block-1", "code", ["print('code')"]), + self._create_block_with_rich_text("block-2", "quote", ["Quoted text"]), + self._create_block_with_rich_text("block-3", "callout", ["Important note"]), + ], + "next_cursor": None, + "has_more": False, + } + mock_request.return_value = httpx.Response(200, json=mock_data) + + # Act + result = extractor._get_notion_block_data("page-456") + + # Assert + assert len(result) == 3 + assert "print('code')" in result[0] + assert "Quoted text" in result[1] + assert "Important note" in result[2] + + @patch("httpx.request") + def test_get_notion_block_data_with_toggle_block(self, mock_request: Mock, extractor: NotionExtractor) -> None: + """Test retrieving page with toggle block containing children. + + Toggle blocks can have nested content that should be extracted. + """ + # Arrange + parent_data = { + "object": "list", + "results": [ + self._create_block_with_rich_text("block-1", "toggle", ["Toggle header"], has_children=True), + ], + "next_cursor": None, + "has_more": False, + } + child_data = { + "object": "list", + "results": [ + self._create_block_with_rich_text("block-child-1", "paragraph", ["Hidden content"]), + ], + "next_cursor": None, + "has_more": False, + } + mock_request.side_effect = [ + httpx.Response(200, json=parent_data), + httpx.Response(200, json=child_data), + ] + + # Act + result = extractor._get_notion_block_data("page-456") + + # Assert + assert len(result) == 1 + assert "Toggle header" in result[0] + assert "Hidden content" in result[0] + + @patch("httpx.request") + def test_get_notion_block_data_mixed_block_types(self, mock_request: Mock, extractor: NotionExtractor) -> None: + """Test retrieving page with mixed block types. + + Real Notion pages contain various block types mixed together. + This tests a realistic scenario with multiple block types. + """ + # Arrange + mock_data = { + "object": "list", + "results": [ + self._create_block_with_rich_text("block-1", "heading_1", ["Project Documentation"]), + self._create_block_with_rich_text("block-2", "paragraph", ["This is an introduction."]), + self._create_block_with_rich_text("block-3", "heading_2", ["Features"]), + self._create_block_with_rich_text("block-4", "bulleted_list_item", ["Feature A"]), + self._create_block_with_rich_text("block-5", "code", ["npm install package"]), + ], + "next_cursor": None, + "has_more": False, + } + mock_request.return_value = httpx.Response(200, json=mock_data) + + # Act + result = extractor._get_notion_block_data("page-456") + + # Assert + assert len(result) == 5 + assert "# Project Documentation" in result[0] + assert "This is an introduction" in result[1] + assert "## Features" in result[2] + assert "Feature A" in result[3] + assert "npm install package" in result[4] + + class TestNotionExtractorDatabaseAdvanced: + """Tests for advanced database scenarios and property types. + + Covers: + - Various property types (date, number, checkbox, url, email, phone, status) + - Rich text properties + - Large database pagination + """ + + @pytest.fixture + def extractor(self) -> NotionExtractor: + """Create a NotionExtractor instance for database testing. + + Returns: + NotionExtractor: Configured extractor for database operations + """ + return NotionExtractor( + notion_workspace_id="workspace-123", + notion_obj_id="database-789", + notion_page_type="database", + tenant_id="tenant-789", + notion_access_token="test-token", + notion_token_loader=lambda: "token", + ) + + def _create_database_page_with_properties( + self, page_id: str, properties: Mapping[str, Mapping[str, object]] + ) -> dict[str, object]: + """Helper to create a database page with various property types. + + Args: + page_id: Unique identifier for the page + properties: Dictionary of property names to property configurations + + Returns: + dict: Notion database page structure + """ + formatted_properties = {} + for prop_name, prop_data in properties.items(): + prop_type = prop_data["type"] + assert isinstance(prop_type, str) + formatted_properties[prop_name] = {"type": prop_type, prop_type: prop_data["value"]} + return { + "object": "page", + "id": page_id, + "properties": formatted_properties, + "url": f"https://notion.so/{page_id}", + } + + @patch("httpx.post") + def test_get_notion_database_data_with_various_property_types( + self, mock_post: Mock, extractor: NotionExtractor + ) -> None: + """Test database with multiple property types. + + Tests date, number, checkbox, URL, email, phone, and status properties. + All property types should be extracted correctly. + """ + # Arrange + mock_response_data = { + "object": "list", + "results": [ + self._create_database_page_with_properties( + "page-1", + { + "Title": {"type": "title", "value": [{"plain_text": "Test Entry"}]}, + "Date": {"type": "date", "value": {"start": "2024-11-27", "end": None}}, + "Price": {"type": "number", "value": 99.99}, + "Completed": {"type": "checkbox", "value": True}, + "Link": {"type": "url", "value": "https://example.com"}, + "Email": {"type": "email", "value": "test@example.com"}, + "Phone": {"type": "phone_number", "value": "+1-555-0123"}, + "Status": {"type": "status", "value": {"name": "Active"}}, + }, + ), + ], + "has_more": False, + "next_cursor": None, + } + mock_response = httpx.Response(200, json=mock_response_data) + mock_post.return_value = mock_response + + # Act + result = extractor._get_notion_database_data("database-789") + + # Assert + assert len(result) == 1 + content = result[0].page_content + assert "Title:Test Entry" in content + assert "Date:" in content + assert "Price:99.99" in content + assert "Completed:True" in content + assert "Link:https://example.com" in content + assert "Email:test@example.com" in content + assert "Phone:+1-555-0123" in content + assert "Status:Active" in content + + @patch("httpx.post") + def test_get_notion_database_data_large_pagination(self, mock_post: Mock, extractor: NotionExtractor) -> None: + """Test database with multiple pages of results. + + Large databases require multiple API calls with cursor-based pagination. + This tests that all pages are retrieved correctly. + """ + # Arrange - Create 3 pages of results + page1_response_data = { + "object": "list", + "results": [ + self._create_database_page_with_properties( + f"page-{i}", {"Title": {"type": "title", "value": [{"plain_text": f"Item {i}"}]}} + ) + for i in range(1, 4) + ], + "has_more": True, + "next_cursor": "cursor-1", + } + page1_response = httpx.Response(200, json=page1_response_data) + + page2_response_data = { + "object": "list", + "results": [ + self._create_database_page_with_properties( + f"page-{i}", {"Title": {"type": "title", "value": [{"plain_text": f"Item {i}"}]}} + ) + for i in range(4, 7) + ], + "has_more": True, + "next_cursor": "cursor-2", + } + page2_response = httpx.Response(200, json=page2_response_data) + + page3_response_data = { + "object": "list", + "results": [ + self._create_database_page_with_properties( + f"page-{i}", {"Title": {"type": "title", "value": [{"plain_text": f"Item {i}"}]}} + ) + for i in range(7, 10) + ], + "has_more": False, + "next_cursor": None, + } + page3_response = httpx.Response(200, json=page3_response_data) + + mock_post.side_effect = [page1_response, page2_response, page3_response] + + # Act + result = extractor._get_notion_database_data("database-789") + + # Assert + assert len(result) == 1 + content = result[0].page_content + # Verify all items from all pages are present + for i in range(1, 10): + assert f"Title:Item {i}" in content + # Verify pagination was called correctly + assert mock_post.call_count == 3 + + @patch("httpx.post") + def test_get_notion_database_data_with_rich_text_property( + self, mock_post: Mock, extractor: NotionExtractor + ) -> None: + """Test database with rich_text property type. + + Rich text properties can contain formatted text and should be extracted. + """ + # Arrange + mock_response_data = { + "object": "list", + "results": [ + self._create_database_page_with_properties( + "page-1", + { + "Title": {"type": "title", "value": [{"plain_text": "Note"}]}, + "Description": { + "type": "rich_text", + "value": [{"plain_text": "This is a detailed description"}], + }, + }, + ), + ], + "has_more": False, + "next_cursor": None, + } + mock_response = httpx.Response(200, json=mock_response_data) + mock_post.return_value = mock_response + + # Act + result = extractor._get_notion_database_data("database-789") + + # Assert + assert len(result) == 1 + content = result[0].page_content + assert "Title:Note" in content + assert "Description:This is a detailed description" in content + + class TestNotionExtractorErrorScenarios: + """Tests for error handling and edge cases. + + Covers: + - Network timeouts + - Rate limiting + - Invalid tokens + - Malformed responses + - Missing required fields + - API version mismatches + """ + + @pytest.fixture + def extractor(self) -> NotionExtractor: + """Create a NotionExtractor instance for error testing. + + Returns: + NotionExtractor: Configured extractor for error scenarios + """ + return NotionExtractor( + notion_workspace_id="workspace-123", + notion_obj_id="page-456", + notion_page_type="page", + tenant_id="tenant-789", + notion_access_token="test-token", + notion_token_loader=lambda: "token", + ) + + @pytest.mark.parametrize( + "error_value", + [httpx.TimeoutException("Request timed out"), httpx.ConnectError("Connection failed")], + ids=["timeout", "connection"], ) + @patch("httpx.request") + def test_get_notion_block_data_network_errors( + self, mock_request: Mock, extractor: NotionExtractor, error_value: httpx.TransportError + ) -> None: + """Test handling of various network errors. - # Assert - assert controller.provider_type == DatasourceProviderType.ONLINE_DOCUMENT + Network issues (timeouts, connection failures) should raise appropriate errors. + """ + # Arrange + mock_request.side_effect = error_value - def test_get_datasource_success(self) -> None: - # Arrange - source = datasource_entity("target_datasource") + # Act & Assert + with pytest.raises(ValueError, match="Error fetching Notion block data") as exc_info: + extractor._get_notion_block_data("page-456") + assert "Error fetching Notion block data" in str(exc_info.value) - entity = provider_entity(DatasourceProviderType.ONLINE_DOCUMENT) - entity.datasources = [source] - entity.identity.icon = "test_icon" - - plugin_unique_identifier = "test_plugin_uid" - tenant_id = "test_tenant_id" - - controller = OnlineDocumentDatasourcePluginProviderController( - entity=entity, - plugin_id="test_plugin_id", - plugin_unique_identifier=plugin_unique_identifier, - tenant_id=tenant_id, + @pytest.mark.parametrize( + ("status_code", "description"), + [ + (401, "Unauthorized"), + (403, "Forbidden"), + (404, "Not Found"), + (429, "Rate limit exceeded"), + ], ) + @patch("httpx.request") + def test_get_notion_block_data_http_status_errors( + self, mock_request: Mock, extractor: NotionExtractor, status_code: int, description: str + ) -> None: + """Test handling of various HTTP status errors. - # Act - result = controller.get_datasource("target_datasource") + Different HTTP error codes (401, 403, 404, 429) should be handled appropriately. + """ + # Arrange + mock_response = httpx.Response(status_code, text=description) + mock_request.return_value = mock_response - # Assert - assert isinstance(result, OnlineDocumentDatasourcePlugin) - assert result.entity == source - assert result.tenant_id == tenant_id - assert result.icon == "test_icon" - assert result.plugin_unique_identifier == plugin_unique_identifier - assert result.runtime.tenant_id == tenant_id + # Act & Assert + with pytest.raises( + ValueError, + match=( + "(?:Error fetching Notion block data: Forbidden|Error fetching Notion " + "block data: Not Found|Error fetching Notion block data: Rate limit " + "exceeded|Error fetching Notion block data: Unauthorized)" + ), + ) as exc_info: + extractor._get_notion_block_data("page-456") + assert "Error fetching Notion block data" in str(exc_info.value) - def test_get_datasource_not_found(self) -> None: - # Arrange - source = datasource_entity("other_datasource") - - entity = provider_entity(DatasourceProviderType.ONLINE_DOCUMENT) - entity.datasources = [source] - - controller = OnlineDocumentDatasourcePluginProviderController( - entity=entity, - plugin_id="test_plugin_id", - plugin_unique_identifier="test_plugin_uid", - tenant_id="test_tenant_id", + @pytest.mark.parametrize( + "response_data", + [ + {"object": "list"}, + {"object": "list", "results": "not a list"}, + {"object": "list", "results": None}, + ], + ids=["missing results field", "results not a list", "results is None"], ) + @patch("httpx.request") + def test_get_notion_block_data_malformed_responses( + self, mock_request: Mock, extractor: NotionExtractor, response_data: dict[str, object] + ) -> None: + """Test handling of malformed API responses. - # Act & Assert - with pytest.raises(ValueError, match="Datasource with name missing_datasource not found"): - controller.get_datasource("missing_datasource") + Various malformed responses should be handled gracefully. + """ + # Arrange + mock_response = httpx.Response(200, json=response_data) + mock_request.return_value = mock_response + + # Act & Assert + with pytest.raises(ValueError, match="Error fetching Notion block data") as exc_info: + extractor._get_notion_block_data("page-456") + assert "Error fetching Notion block data" in str(exc_info.value) + + @patch("httpx.post") + def test_get_notion_database_data_with_query_filter(self, mock_post: Mock, extractor: NotionExtractor) -> None: + """Test database query with custom filter. + + Databases can be queried with filters to retrieve specific rows. + """ + # Arrange + mock_response_data = { + "object": "list", + "results": [ + { + "object": "page", + "id": "page-1", + "properties": { + "Title": {"type": "title", "title": [{"plain_text": "Filtered Item"}]}, + "Status": {"type": "select", "select": {"name": "Active"}}, + }, + "url": "https://notion.so/page-1", + } + ], + "has_more": False, + "next_cursor": None, + } + mock_response = httpx.Response(200, json=mock_response_data) + mock_post.return_value = mock_response + + # Create a custom query filter + query_filter = {"filter": {"property": "Status", "select": {"equals": "Active"}}} + + # Act + result = extractor._get_notion_database_data("database-789", query_dict=query_filter) + + # Assert + assert len(result) == 1 + content = result[0].page_content + assert "Title:Filtered Item" in content + assert "Status:Active" in content + # Verify the filter was passed to the API + mock_post.assert_called_once() + call_args = mock_post.call_args + assert "filter" in call_args[1]["json"] + + class TestNotionExtractorTableAdvanced: + """Tests for advanced table scenarios. + + Covers: + - Tables with many columns + - Tables with complex cell content + - Empty tables + """ + + @pytest.fixture + def extractor(self) -> NotionExtractor: + """Create a NotionExtractor instance for table testing. + + Returns: + NotionExtractor: Configured extractor for table operations + """ + return NotionExtractor( + notion_workspace_id="workspace-123", + notion_obj_id="page-456", + notion_page_type="page", + tenant_id="tenant-789", + notion_access_token="test-token", + notion_token_loader=lambda: "token", + ) + + @patch("httpx.request") + def test_read_table_rows_with_many_columns(self, mock_request: Mock, extractor: NotionExtractor) -> None: + """Test reading table with many columns. + + Tables can have numerous columns; all should be extracted correctly. + """ + # Arrange - Create a table with 10 columns + headers = [f"Col{i}" for i in range(1, 11)] + values = [f"Val{i}" for i in range(1, 11)] + + mock_data = { + "object": "list", + "results": [ + { + "object": "block", + "type": "table_row", + "table_row": {"cells": [[{"text": {"content": h}}] for h in headers]}, + }, + { + "object": "block", + "type": "table_row", + "table_row": {"cells": [[{"text": {"content": v}}] for v in values]}, + }, + ], + "next_cursor": None, + "has_more": False, + } + mock_request.return_value = httpx.Response(200, json=mock_data) + + # Act + result = extractor._read_table_rows("table-block-123") + + # Assert + for header in headers: + assert header in result + for value in values: + assert value in result + # Verify markdown table structure + assert "| --- |" in result diff --git a/api/tests/unit_tests/core/datasource/online_drive/test_online_drive_plugin.py b/api/tests/unit_tests/core/datasource/online_drive/test_online_drive_plugin.py index 0ddcffabbc8..48d9c4b42df 100644 --- a/api/tests/unit_tests/core/datasource/online_drive/test_online_drive_plugin.py +++ b/api/tests/unit_tests/core/datasource/online_drive/test_online_drive_plugin.py @@ -9,7 +9,7 @@ from core.datasource.entities.datasource_entities import ( OnlineDriveDownloadFileRequest, ) from core.datasource.online_drive.online_drive_plugin import OnlineDriveDatasourcePlugin -from tests.unit_tests.core.datasource.factories import datasource_entity +from tests.helpers.unit.core.datasource.factories import datasource_entity class TestOnlineDriveDatasourcePlugin: diff --git a/api/tests/unit_tests/core/datasource/online_drive/test_online_drive_provider.py b/api/tests/unit_tests/core/datasource/online_drive/test_online_drive_provider.py index 557b89fe3a7..701dcfcf9d4 100644 --- a/api/tests/unit_tests/core/datasource/online_drive/test_online_drive_provider.py +++ b/api/tests/unit_tests/core/datasource/online_drive/test_online_drive_provider.py @@ -3,7 +3,7 @@ import pytest from core.datasource.entities.datasource_entities import DatasourceProviderType from core.datasource.online_drive.online_drive_plugin import OnlineDriveDatasourcePlugin from core.datasource.online_drive.online_drive_provider import OnlineDriveDatasourcePluginProviderController -from tests.unit_tests.core.datasource.factories import datasource_entity, provider_entity +from tests.helpers.unit.core.datasource.factories import datasource_entity, provider_entity class TestOnlineDriveDatasourcePluginProviderController: diff --git a/api/tests/unit_tests/core/datasource/test_datasource_file_manager.py b/api/tests/unit_tests/core/datasource/test_datasource_file_manager.py index 5aeac8de816..8b729e59d1a 100644 --- a/api/tests/unit_tests/core/datasource/test_datasource_file_manager.py +++ b/api/tests/unit_tests/core/datasource/test_datasource_file_manager.py @@ -12,7 +12,7 @@ from graphon.file import FileTransferMethod, FileType from models.enums import CreatorUserRole from models.model import MessageFile, UploadFile from models.tools import ToolFile -from tests.unit_tests.model_factories import make_upload_file +from tests.helpers.unit.model_factories import make_upload_file def _upload_file(id: str, *, key: str, mime_type: str) -> UploadFile: diff --git a/api/tests/unit_tests/core/datasource/test_notion_provider.py b/api/tests/unit_tests/core/datasource/test_notion_provider.py deleted file mode 100644 index acb4a61300a..00000000000 --- a/api/tests/unit_tests/core/datasource/test_notion_provider.py +++ /dev/null @@ -1,1800 +0,0 @@ -"""Comprehensive unit tests for Notion datasource provider. - -This test module covers all aspects of the Notion provider including: -- Notion API integration with proper authentication -- Page retrieval (single pages and databases) -- Block content parsing (headings, paragraphs, tables, nested blocks) -- Authentication handling (OAuth tokens, integration tokens, credential management) -- Error handling for API failures -- Pagination handling for large datasets -- Last edited time tracking - -All tests use mocking to avoid external dependencies and ensure fast, reliable execution. -Tests follow the Arrange-Act-Assert pattern for clarity. -""" - -import json -from collections.abc import Callable, Iterator, Mapping -from unittest.mock import Mock, patch -from uuid import uuid4 - -import httpx -import pytest -from sqlalchemy.engine import Engine -from sqlalchemy.orm import Session, sessionmaker - -from core.datasource.entities.datasource_entities import ( - DatasourceEntity, - DatasourceIdentity, - DatasourceProviderEntityWithPlugin, - DatasourceProviderIdentity, - DatasourceProviderType, -) -from core.datasource.online_document.online_document_provider import ( - OnlineDocumentDatasourcePluginProviderController, -) -from core.rag.extractor import notion_extractor as notion_extractor_module -from core.rag.extractor.entity.extract_setting import StoredDocumentExtractionInput -from core.rag.extractor.notion_extractor import NotionExtractor -from core.rag.models.document import Document -from core.tools.entities.common_entities import I18nObject -from extensions.application_services.data_sources import build_data_source_credentials -from models.base import TypeBase -from models.dataset import Document as DocumentModel -from models.enums import DataSourceType, DocumentCreatedFrom - - -@pytest.fixture -def database(sqlite_engine: Engine, monkeypatch: pytest.MonkeyPatch) -> Iterator[Session]: - """Bind a real session for Notion document metadata persistence.""" - - TypeBase.metadata.create_all(sqlite_engine, tables=[TypeBase.metadata.tables[DocumentModel.__tablename__]]) - with Session(sqlite_engine, expire_on_commit=False) as session: - monkeypatch.setattr(notion_extractor_module.db, "session", session) - yield session - - -@pytest.fixture -def persisted_document(database: Session) -> DocumentModel: - document = DocumentModel( - id=str(uuid4()), - tenant_id=str(uuid4()), - dataset_id=str(uuid4()), - position=1, - data_source_type=DataSourceType.NOTION_IMPORT, - data_source_info=json.dumps({"last_edited_time": "2024-01-01T00:00:00.000Z"}), - batch="batch", - name="Notion page", - created_from=DocumentCreatedFrom.WEB, - created_by=str(uuid4()), - ) - database.add(document) - database.commit() - return document - - -def _extraction_input(document: DocumentModel) -> StoredDocumentExtractionInput: - """Detach the values the extractor needs, as the indexing source adapter does.""" - return StoredDocumentExtractionInput( - id=document.id, - tenant_id=document.tenant_id, - dataset_id=document.dataset_id, - data_source_info_dict=document.data_source_info_dict, - ) - - -class TestNotionExtractorAuthentication: - """Tests for Notion authentication handling. - - Covers: - - OAuth token authentication - - Integration token fallback - - Credential retrieval from database - - Missing credential error handling - """ - - @pytest.fixture - def mock_document_model(self) -> StoredDocumentExtractionInput: - """Create detached document values with persisted-format Notion metadata.""" - return StoredDocumentExtractionInput( - id="test-doc-id", - tenant_id="tenant-789", - dataset_id="dataset-id", - data_source_info_dict={"last_edited_time": "2024-01-01T00:00:00.000Z"}, - ) - - def test_init_with_explicit_token(self, mock_document_model: StoredDocumentExtractionInput) -> None: - """Test NotionExtractor initialization with explicit access token.""" - # Arrange & Act - extractor = NotionExtractor( - notion_workspace_id="workspace-123", - notion_obj_id="page-456", - notion_page_type="page", - tenant_id="tenant-789", - notion_access_token="explicit-token-abc", - document_model=mock_document_model, - notion_token_loader=lambda: "token", - ) - - # Assert - assert extractor._notion_access_token == "explicit-token-abc" - assert extractor._notion_workspace_id == "workspace-123" - assert extractor._notion_obj_id == "page-456" - assert extractor._notion_page_type == "page" - - @patch("services.data_source.provider_service.DatasourceProviderService.get_datasource_credentials") - def test_init_with_credential_id( - self, - get_credentials: Mock, - mock_document_model: StoredDocumentExtractionInput, - sqlite_session_factory: sessionmaker[Session], - ) -> None: - """Test NotionExtractor initialization with credential ID retrieval.""" - # Arrange - providers = build_data_source_credentials(database_client=sqlite_session_factory).providers - get_credentials.return_value = {"integration_secret": "credential-token-xyz"} - - # Act - extractor = NotionExtractor( - notion_workspace_id="workspace-123", - notion_obj_id="page-456", - notion_page_type="page", - tenant_id="tenant-789", - credential_id="cred-123", - document_model=mock_document_model, - notion_token_loader=lambda: providers.get_stored_notion_access_token( - tenant_id="tenant-789", credential_id="cred-123" - ), - ) - - # Assert - assert extractor._notion_access_token == "credential-token-xyz" - get_credentials.assert_called_once_with( - tenant_id="tenant-789", - credential_id="cred-123", - provider="notion_datasource", - plugin_id="langgenius/notion_datasource", - ) - - @patch("services.data_source.provider_service.DatasourceProviderService.get_datasource_credentials") - def test_init_with_integration_token_fallback( - self, - mock_get_token: Mock, - mock_document_model: StoredDocumentExtractionInput, - config_overrides: Callable[..., None], - sqlite_session_factory: sessionmaker[Session], - ) -> None: - """Test NotionExtractor falls back to integration token when credential not found.""" - # Arrange - mock_get_token.side_effect = Exception("No credential id found") - config_overrides(NOTION_INTEGRATION_TOKEN="integration-token-fallback") - providers = build_data_source_credentials(database_client=sqlite_session_factory).providers - - # Act - extractor = NotionExtractor( - notion_workspace_id="workspace-123", - notion_obj_id="page-456", - notion_page_type="page", - tenant_id="tenant-789", - credential_id=None, - document_model=mock_document_model, - notion_token_loader=lambda: providers.get_stored_notion_access_token( - tenant_id="tenant-789", credential_id=None - ), - ) - - # Assert - assert extractor._notion_access_token == "integration-token-fallback" - - @patch("services.data_source.provider_service.DatasourceProviderService.get_datasource_credentials") - def test_init_missing_credentials_raises_error( - self, - mock_get_token: Mock, - mock_document_model: StoredDocumentExtractionInput, - config_overrides: Callable[..., None], - sqlite_session_factory: sessionmaker[Session], - ) -> None: - """Test NotionExtractor raises error when no credentials available.""" - # Arrange - mock_get_token.side_effect = Exception("No credential id found") - config_overrides(NOTION_INTEGRATION_TOKEN=None) - providers = build_data_source_credentials(database_client=sqlite_session_factory).providers - - # Act & Assert - with pytest.raises( - ValueError, - match=("Must specify `integration_token` or set environment variable `NOTION_INTEGRATION_TOKEN`\\."), - ) as exc_info: - NotionExtractor( - notion_workspace_id="workspace-123", - notion_obj_id="page-456", - notion_page_type="page", - tenant_id="tenant-789", - credential_id=None, - document_model=mock_document_model, - notion_token_loader=lambda: providers.get_stored_notion_access_token( - tenant_id="tenant-789", credential_id=None - ), - ) - assert "Must specify `integration_token`" in str(exc_info.value) - - -class TestNotionExtractorPageRetrieval: - """Tests for Notion page retrieval functionality. - - Covers: - - Single page retrieval - - Database page retrieval with pagination - - Block content extraction - - Nested block handling - """ - - @pytest.fixture - def extractor(self) -> NotionExtractor: - """Create a NotionExtractor instance for testing.""" - return NotionExtractor( - notion_workspace_id="workspace-123", - notion_obj_id="page-456", - notion_page_type="page", - tenant_id="tenant-789", - notion_access_token="test-token", - notion_token_loader=lambda: "token", - ) - - def _create_response(self, data: Mapping[str, object], status_code: int = 200) -> httpx.Response: - """Build an HTTP response with serialized JSON for the extractor to parse.""" - return httpx.Response(status_code, json=data) - - def _create_block( - self, block_id: str, block_type: str, text_content: str, has_children: bool = False - ) -> dict[str, object]: - """Helper to create a Notion block structure.""" - return { - "object": "block", - "id": block_id, - "type": block_type, - "has_children": has_children, - block_type: { - "rich_text": [ - { - "type": "text", - "text": {"content": text_content}, - "plain_text": text_content, - } - ] - }, - } - - @patch("httpx.request") - def test_get_notion_block_data_simple_page(self, mock_request: Mock, extractor: NotionExtractor) -> None: - """Test retrieving simple page with basic blocks.""" - # Arrange - mock_data = { - "object": "list", - "results": [ - self._create_block("block-1", "paragraph", "First paragraph"), - self._create_block("block-2", "paragraph", "Second paragraph"), - ], - "next_cursor": None, - "has_more": False, - } - mock_request.return_value = self._create_response(mock_data) - - # Act - result = extractor._get_notion_block_data("page-456") - - # Assert - assert len(result) == 2 - assert "First paragraph" in result[0] - assert "Second paragraph" in result[1] - mock_request.assert_called_once() - - @patch("httpx.request") - def test_get_notion_block_data_with_headings(self, mock_request: Mock, extractor: NotionExtractor) -> None: - """Test retrieving page with heading blocks.""" - # Arrange - mock_data = { - "object": "list", - "results": [ - self._create_block("block-1", "heading_1", "Main Title"), - self._create_block("block-2", "heading_2", "Subtitle"), - self._create_block("block-3", "paragraph", "Content text"), - self._create_block("block-4", "heading_3", "Sub-subtitle"), - ], - "next_cursor": None, - "has_more": False, - } - mock_request.return_value = self._create_response(mock_data) - - # Act - result = extractor._get_notion_block_data("page-456") - - # Assert - assert len(result) == 4 - assert "# Main Title" in result[0] - assert "## Subtitle" in result[1] - assert "Content text" in result[2] - assert "### Sub-subtitle" in result[3] - - @patch("httpx.request") - def test_get_notion_block_data_with_pagination(self, mock_request: Mock, extractor: NotionExtractor) -> None: - """Test retrieving page with paginated results.""" - # Arrange - first_page = { - "object": "list", - "results": [self._create_block("block-1", "paragraph", "First page content")], - "next_cursor": "cursor-abc", - "has_more": True, - } - second_page = { - "object": "list", - "results": [self._create_block("block-2", "paragraph", "Second page content")], - "next_cursor": None, - "has_more": False, - } - mock_request.side_effect = [ - self._create_response(first_page), - self._create_response(second_page), - ] - - # Act - result = extractor._get_notion_block_data("page-456") - - # Assert - assert len(result) == 2 - assert "First page content" in result[0] - assert "Second page content" in result[1] - assert mock_request.call_count == 2 - - @patch("httpx.request") - def test_get_notion_block_data_with_nested_blocks(self, mock_request: Mock, extractor: NotionExtractor) -> None: - """Test retrieving page with nested block structure.""" - # Arrange - # First call returns parent blocks - parent_data = { - "object": "list", - "results": [ - self._create_block("block-1", "paragraph", "Parent block", has_children=True), - ], - "next_cursor": None, - "has_more": False, - } - # Second call returns child blocks - child_data = { - "object": "list", - "results": [ - self._create_block("block-child-1", "paragraph", "Child block"), - ], - "next_cursor": None, - "has_more": False, - } - mock_request.side_effect = [ - self._create_response(parent_data), - self._create_response(child_data), - ] - - # Act - result = extractor._get_notion_block_data("page-456") - - # Assert - assert len(result) == 1 - assert "Parent block" in result[0] - assert "Child block" in result[0] - assert mock_request.call_count == 2 - - @patch("httpx.request") - def test_get_notion_block_data_error_handling(self, mock_request: Mock, extractor: NotionExtractor) -> None: - """Test error handling for failed API requests.""" - # Arrange - mock_request.return_value = self._create_response({}, status_code=404) - - # Act & Assert - with pytest.raises(ValueError, match="Error fetching Notion block data: \\{\\}") as exc_info: - extractor._get_notion_block_data("page-456") - assert "Error fetching Notion block data" in str(exc_info.value) - - @patch("httpx.request") - def test_get_notion_block_data_invalid_response(self, mock_request: Mock, extractor: NotionExtractor) -> None: - """Test handling of invalid API response structure.""" - # Arrange - mock_request.return_value = self._create_response({"invalid": "structure"}) - - # Act & Assert - with pytest.raises(ValueError, match="Error fetching Notion block data") as exc_info: - extractor._get_notion_block_data("page-456") - assert "Error fetching Notion block data" in str(exc_info.value) - - @patch("httpx.request") - def test_get_notion_block_data_http_error(self, mock_request: Mock, extractor: NotionExtractor) -> None: - """Test handling of HTTP errors during request.""" - # Arrange - mock_request.side_effect = httpx.HTTPError("Network error") - - # Act & Assert - with pytest.raises(ValueError, match="Error fetching Notion block data") as exc_info: - extractor._get_notion_block_data("page-456") - assert "Error fetching Notion block data" in str(exc_info.value) - - -class TestNotionExtractorDatabaseRetrieval: - """Tests for Notion database retrieval functionality. - - Covers: - - Database query with pagination - - Property extraction (title, rich_text, select, multi_select, etc.) - - Row formatting - - Empty database handling - """ - - @pytest.fixture - def extractor(self) -> NotionExtractor: - """Create a NotionExtractor instance for testing.""" - return NotionExtractor( - notion_workspace_id="workspace-123", - notion_obj_id="database-789", - notion_page_type="database", - tenant_id="tenant-789", - notion_access_token="test-token", - notion_token_loader=lambda: "token", - ) - - def _create_database_page(self, page_id: str, properties: Mapping[str, Mapping[str, object]]) -> dict[str, object]: - """Helper to create a database page structure.""" - formatted_properties = {} - for prop_name, prop_data in properties.items(): - prop_type = prop_data["type"] - assert isinstance(prop_type, str) - formatted_properties[prop_name] = {"type": prop_type, prop_type: prop_data["value"]} - return { - "object": "page", - "id": page_id, - "properties": formatted_properties, - "url": f"https://notion.so/{page_id}", - } - - @patch("httpx.post") - def test_get_notion_database_data_simple(self, mock_post: Mock, extractor: NotionExtractor) -> None: - """Test retrieving simple database with basic properties.""" - # Arrange - mock_response_data = { - "object": "list", - "results": [ - self._create_database_page( - "page-1", - { - "Title": {"type": "title", "value": [{"plain_text": "Task 1"}]}, - "Status": {"type": "select", "value": {"name": "In Progress"}}, - }, - ), - self._create_database_page( - "page-2", - { - "Title": {"type": "title", "value": [{"plain_text": "Task 2"}]}, - "Status": {"type": "select", "value": {"name": "Done"}}, - }, - ), - ], - "has_more": False, - "next_cursor": None, - } - mock_response = httpx.Response(200, json=mock_response_data) - mock_post.return_value = mock_response - - # Act - result = extractor._get_notion_database_data("database-789") - - # Assert - assert len(result) == 1 - content = result[0].page_content - assert "Title:Task 1" in content - assert "Status:In Progress" in content - assert "Title:Task 2" in content - assert "Status:Done" in content - - @patch("httpx.post") - def test_get_notion_database_data_with_pagination(self, mock_post: Mock, extractor: NotionExtractor) -> None: - """Test retrieving database with paginated results.""" - # Arrange - first_response_data = { - "object": "list", - "results": [ - self._create_database_page("page-1", {"Title": {"type": "title", "value": [{"plain_text": "Page 1"}]}}), - ], - "has_more": True, - "next_cursor": "cursor-xyz", - } - first_response = httpx.Response(200, json=first_response_data) - second_response_data = { - "object": "list", - "results": [ - self._create_database_page("page-2", {"Title": {"type": "title", "value": [{"plain_text": "Page 2"}]}}), - ], - "has_more": False, - "next_cursor": None, - } - second_response = httpx.Response(200, json=second_response_data) - mock_post.side_effect = [first_response, second_response] - - # Act - result = extractor._get_notion_database_data("database-789") - - # Assert - assert len(result) == 1 - content = result[0].page_content - assert "Title:Page 1" in content - assert "Title:Page 2" in content - assert mock_post.call_count == 2 - - @patch("httpx.post") - def test_get_notion_database_data_multi_select(self, mock_post: Mock, extractor: NotionExtractor) -> None: - """Test database with multi_select property type.""" - # Arrange - mock_response_data = { - "object": "list", - "results": [ - self._create_database_page( - "page-1", - { - "Title": {"type": "title", "value": [{"plain_text": "Project"}]}, - "Tags": { - "type": "multi_select", - "value": [{"name": "urgent"}, {"name": "frontend"}], - }, - }, - ), - ], - "has_more": False, - "next_cursor": None, - } - mock_response = httpx.Response(200, json=mock_response_data) - mock_post.return_value = mock_response - - # Act - result = extractor._get_notion_database_data("database-789") - - # Assert - assert len(result) == 1 - content = result[0].page_content - assert "Title:Project" in content - assert "Tags:" in content - - @patch("httpx.post") - def test_get_notion_database_data_empty_properties(self, mock_post: Mock, extractor: NotionExtractor) -> None: - """Test database with empty property values.""" - # Arrange - mock_response_data = { - "object": "list", - "results": [ - self._create_database_page( - "page-1", - { - "Title": {"type": "title", "value": []}, - "Status": {"type": "select", "value": None}, - }, - ), - ], - "has_more": False, - "next_cursor": None, - } - mock_response = httpx.Response(200, json=mock_response_data) - mock_post.return_value = mock_response - - # Act - result = extractor._get_notion_database_data("database-789") - - # Assert - assert len(result) == 1 - # Empty properties should be filtered out - content = result[0].page_content - assert "Row Page URL:" in content - - @patch("httpx.post") - def test_get_notion_database_data_empty_results(self, mock_post: Mock, extractor: NotionExtractor) -> None: - """Test handling of empty database.""" - # Arrange - mock_response_data = { - "object": "list", - "results": list[object](), - "has_more": False, - "next_cursor": None, - } - mock_response = httpx.Response(200, json=mock_response_data) - mock_post.return_value = mock_response - - # Act - result = extractor._get_notion_database_data("database-789") - - # Assert - assert len(result) == 0 - - @patch("httpx.post") - def test_get_notion_database_data_missing_results(self, mock_post: Mock, extractor: NotionExtractor) -> None: - """Test handling of malformed API response.""" - # Arrange - mock_response_data = {"object": "list"} - mock_response = httpx.Response(200, json=mock_response_data) - mock_post.return_value = mock_response - - # Act - result = extractor._get_notion_database_data("database-789") - - # Assert - assert len(result) == 0 - - -class TestNotionExtractorTableParsing: - """Tests for Notion table block parsing. - - Covers: - - Table header extraction - - Table row parsing - - Markdown table formatting - - Empty cell handling - """ - - @pytest.fixture - def extractor(self) -> NotionExtractor: - """Create a NotionExtractor instance for testing.""" - return NotionExtractor( - notion_workspace_id="workspace-123", - notion_obj_id="page-456", - notion_page_type="page", - tenant_id="tenant-789", - notion_access_token="test-token", - notion_token_loader=lambda: "token", - ) - - @patch("httpx.request") - def test_read_table_rows_simple(self, mock_request: Mock, extractor: NotionExtractor) -> None: - """Test reading simple table with headers and rows.""" - # Arrange - mock_data = { - "object": "list", - "results": [ - { - "object": "block", - "type": "table_row", - "table_row": { - "cells": [ - [{"text": {"content": "Name"}}], - [{"text": {"content": "Age"}}], - ] - }, - }, - { - "object": "block", - "type": "table_row", - "table_row": { - "cells": [ - [{"text": {"content": "Alice"}}], - [{"text": {"content": "30"}}], - ] - }, - }, - { - "object": "block", - "type": "table_row", - "table_row": { - "cells": [ - [{"text": {"content": "Bob"}}], - [{"text": {"content": "25"}}], - ] - }, - }, - ], - "next_cursor": None, - "has_more": False, - } - mock_request.return_value = httpx.Response(200, json=mock_data) - - # Act - result = extractor._read_table_rows("table-block-123") - - # Assert - assert "| Name | Age |" in result - assert "| --- | --- |" in result - assert "| Alice | 30 |" in result - assert "| Bob | 25 |" in result - - @patch("httpx.request") - def test_read_table_rows_with_empty_cells(self, mock_request: Mock, extractor: NotionExtractor) -> None: - """Test reading table with empty cells.""" - # Arrange - mock_data = { - "object": "list", - "results": [ - { - "object": "block", - "type": "table_row", - "table_row": {"cells": [[{"text": {"content": "Col1"}}], [{"text": {"content": "Col2"}}]]}, - }, - { - "object": "block", - "type": "table_row", - "table_row": {"cells": [[{"text": {"content": "Value1"}}], list[object]()]}, - }, - ], - "next_cursor": None, - "has_more": False, - } - mock_request.return_value = httpx.Response(200, json=mock_data) - - # Act - result = extractor._read_table_rows("table-block-123") - - # Assert - assert "| Col1 | Col2 |" in result - assert "| --- | --- |" in result - # Empty cells are handled by the table parsing logic - assert "Value1" in result - - @patch("httpx.request") - def test_read_table_rows_with_pagination(self, mock_request: Mock, extractor: NotionExtractor) -> None: - """Test reading table with paginated results.""" - # Arrange - first_page = { - "object": "list", - "results": [ - { - "object": "block", - "type": "table_row", - "table_row": {"cells": [[{"text": {"content": "Header"}}]]}, - }, - ], - "next_cursor": "cursor-abc", - "has_more": True, - } - second_page = { - "object": "list", - "results": [ - { - "object": "block", - "type": "table_row", - "table_row": {"cells": [[{"text": {"content": "Row1"}}]]}, - }, - ], - "next_cursor": None, - "has_more": False, - } - mock_request.side_effect = [httpx.Response(200, json=first_page), httpx.Response(200, json=second_page)] - - # Act - result = extractor._read_table_rows("table-block-123") - - # Assert - assert "| Header |" in result - assert mock_request.call_count == 2 - - -class TestNotionExtractorLastEditedTime: - """Tests for last edited time tracking. - - Covers: - - Page last edited time retrieval - - Database last edited time retrieval - - Document model update - """ - - @pytest.fixture - def mock_document_model(self) -> StoredDocumentExtractionInput: - """Create detached document values with persisted-format Notion metadata.""" - return StoredDocumentExtractionInput( - id="test-doc-id", - tenant_id="tenant-789", - dataset_id="dataset-id", - data_source_info_dict={"last_edited_time": "2024-01-01T00:00:00.000Z"}, - ) - - @pytest.fixture - def extractor_page(self, mock_document_model: StoredDocumentExtractionInput) -> NotionExtractor: - """Create a NotionExtractor instance for page testing.""" - return NotionExtractor( - notion_workspace_id="workspace-123", - notion_obj_id="page-456", - notion_page_type="page", - tenant_id="tenant-789", - notion_access_token="test-token", - document_model=mock_document_model, - notion_token_loader=lambda: "token", - ) - - @pytest.fixture - def extractor_database(self, mock_document_model: StoredDocumentExtractionInput) -> NotionExtractor: - """Create a NotionExtractor instance for database testing.""" - return NotionExtractor( - notion_workspace_id="workspace-123", - notion_obj_id="database-789", - notion_page_type="database", - tenant_id="tenant-789", - notion_access_token="test-token", - document_model=mock_document_model, - notion_token_loader=lambda: "token", - ) - - @patch("httpx.request") - def test_get_notion_last_edited_time_page(self, mock_request: Mock, extractor_page: NotionExtractor) -> None: - """Test retrieving last edited time for a page.""" - # Arrange - mock_response_data = { - "object": "page", - "id": "page-456", - "last_edited_time": "2024-11-27T12:00:00.000Z", - } - mock_response = httpx.Response(200, json=mock_response_data) - mock_request.return_value = mock_response - - # Act - result = extractor_page.get_notion_last_edited_time() - - # Assert - assert result == "2024-11-27T12:00:00.000Z" - mock_request.assert_called_once() - call_args = mock_request.call_args - assert "pages/page-456" in call_args[0][1] - - @patch("httpx.request") - def test_get_notion_last_edited_time_database( - self, mock_request: Mock, extractor_database: NotionExtractor - ) -> None: - """Test retrieving last edited time for a database.""" - # Arrange - mock_response_data = { - "object": "database", - "id": "database-789", - "last_edited_time": "2024-11-27T15:30:00.000Z", - } - mock_response = httpx.Response(200, json=mock_response_data) - mock_request.return_value = mock_response - - # Act - result = extractor_database.get_notion_last_edited_time() - - # Assert - assert result == "2024-11-27T15:30:00.000Z" - mock_request.assert_called_once() - call_args = mock_request.call_args - assert "databases/database-789" in call_args[0][1] - - @patch("httpx.request") - def test_update_last_edited_time( - self, - mock_request: Mock, - extractor_page: NotionExtractor, - database: Session, - persisted_document: DocumentModel, - ) -> None: - """Test updating document model with last edited time.""" - # Arrange - mock_response_data = { - "object": "page", - "id": "page-456", - "last_edited_time": "2024-11-27T18:00:00.000Z", - } - mock_response = httpx.Response(200, json=mock_response_data) - mock_request.return_value = mock_response - - # Act - extractor_page.update_last_edited_time(_extraction_input(persisted_document)) - - # Assert - database.expire(persisted_document) - assert persisted_document.data_source_info_dict["last_edited_time"] == "2024-11-27T18:00:00.000Z" - - def test_update_last_edited_time_no_document(self, extractor_page: NotionExtractor) -> None: - """Test update_last_edited_time with None document model.""" - # Act & Assert - should not raise error - extractor_page.update_last_edited_time(None) - - -class TestNotionExtractorIntegration: - """Integration tests for complete extraction workflow. - - Covers: - - Full page extraction workflow - - Full database extraction workflow - - Document creation - - Error handling in extract method - """ - - @pytest.fixture - def mock_document_model(self) -> StoredDocumentExtractionInput: - """Create detached document values with persisted-format Notion metadata.""" - return StoredDocumentExtractionInput( - id="test-doc-id", - tenant_id="tenant-789", - dataset_id="dataset-id", - data_source_info_dict={"last_edited_time": "2024-01-01T00:00:00.000Z"}, - ) - - @patch("httpx.request") - def test_extract_page_complete_workflow( - self, mock_request: Mock, database: Session, persisted_document: DocumentModel - ) -> None: - """Test complete page extraction workflow.""" - # Arrange - extractor = NotionExtractor( - notion_workspace_id="workspace-123", - notion_obj_id="page-456", - notion_page_type="page", - tenant_id="tenant-789", - notion_access_token="test-token", - document_model=_extraction_input(persisted_document), - notion_token_loader=lambda: "token", - ) - - # Mock last edited time request - last_edited_response_data = { - "object": "page", - "last_edited_time": "2024-11-27T20:00:00.000Z", - } - last_edited_response = httpx.Response(200, json=last_edited_response_data) - - # Mock block data request - block_response_data = { - "object": "list", - "results": [ - { - "object": "block", - "id": "block-1", - "type": "heading_1", - "has_children": False, - "heading_1": { - "rich_text": [{"type": "text", "text": {"content": "Test Page"}, "plain_text": "Test Page"}] - }, - }, - { - "object": "block", - "id": "block-2", - "type": "paragraph", - "has_children": False, - "paragraph": { - "rich_text": [ - {"type": "text", "text": {"content": "Test content"}, "plain_text": "Test content"} - ] - }, - }, - ], - "next_cursor": None, - "has_more": False, - } - block_response = httpx.Response(200, json=block_response_data) - - mock_request.side_effect = [last_edited_response, block_response] - - # Act - documents = extractor.extract() - - # Assert - assert len(documents) == 1 - assert isinstance(documents[0], Document) - assert "# Test Page" in documents[0].page_content - assert "Test content" in documents[0].page_content - database.expire(persisted_document) - assert persisted_document.data_source_info_dict["last_edited_time"] == "2024-11-27T20:00:00.000Z" - - @patch("httpx.post") - @patch("httpx.request") - def test_extract_database_complete_workflow( - self, - mock_request: Mock, - mock_post: Mock, - database: Session, - persisted_document: DocumentModel, - ) -> None: - """Test complete database extraction workflow.""" - # Arrange - extractor = NotionExtractor( - notion_workspace_id="workspace-123", - notion_obj_id="database-789", - notion_page_type="database", - tenant_id="tenant-789", - notion_access_token="test-token", - document_model=_extraction_input(persisted_document), - notion_token_loader=lambda: "token", - ) - - # Mock last edited time request - last_edited_response_data = { - "object": "database", - "last_edited_time": "2024-11-27T20:00:00.000Z", - } - last_edited_response = httpx.Response(200, json=last_edited_response_data) - mock_request.return_value = last_edited_response - - # Mock database query request - database_response_data = { - "object": "list", - "results": [ - { - "object": "page", - "id": "page-1", - "properties": { - "Name": {"type": "title", "title": [{"plain_text": "Item 1"}]}, - "Status": {"type": "select", "select": {"name": "Active"}}, - }, - "url": "https://notion.so/page-1", - } - ], - "has_more": False, - "next_cursor": None, - } - database_response = httpx.Response(200, json=database_response_data) - mock_post.return_value = database_response - - # Act - documents = extractor.extract() - - # Assert - assert len(documents) == 1 - assert isinstance(documents[0], Document) - assert "Name:Item 1" in documents[0].page_content - assert "Status:Active" in documents[0].page_content - database.expire(persisted_document) - assert persisted_document.data_source_info_dict["last_edited_time"] == "2024-11-27T20:00:00.000Z" - - def test_extract_invalid_page_type(self) -> None: - """Test extract with invalid page type.""" - # Arrange - extractor = NotionExtractor( - notion_workspace_id="workspace-123", - notion_obj_id="invalid-456", - notion_page_type="invalid_type", - tenant_id="tenant-789", - notion_access_token="test-token", - notion_token_loader=lambda: "token", - ) - - # Act & Assert - with pytest.raises(ValueError, match="notion page type not supported") as exc_info: - extractor.extract() - assert "notion page type not supported" in str(exc_info.value) - - -class TestNotionExtractorReadBlock: - """Tests for nested block reading functionality. - - Covers: - - Recursive block reading - - Indentation handling - - Child page handling - """ - - @pytest.fixture - def extractor(self) -> NotionExtractor: - """Create a NotionExtractor instance for testing.""" - return NotionExtractor( - notion_workspace_id="workspace-123", - notion_obj_id="page-456", - notion_page_type="page", - tenant_id="tenant-789", - notion_access_token="test-token", - notion_token_loader=lambda: "token", - ) - - @patch("httpx.request") - def test_read_block_with_indentation(self, mock_request: Mock, extractor: NotionExtractor) -> None: - """Test reading nested blocks with proper indentation.""" - # Arrange - mock_data = { - "object": "list", - "results": [ - { - "object": "block", - "id": "block-1", - "type": "paragraph", - "has_children": False, - "paragraph": { - "rich_text": [ - {"type": "text", "text": {"content": "Nested content"}, "plain_text": "Nested content"} - ] - }, - } - ], - "next_cursor": None, - "has_more": False, - } - mock_request.return_value = httpx.Response(200, json=mock_data) - - # Act - result = extractor._read_block("block-parent", num_tabs=2) - - # Assert - assert "\t\tNested content" in result - - @patch("httpx.request") - def test_read_block_skip_child_page(self, mock_request: Mock, extractor: NotionExtractor) -> None: - """Test that child_page blocks don't recurse.""" - # Arrange - mock_data = { - "object": "list", - "results": [ - { - "object": "block", - "id": "block-1", - "type": "child_page", - "has_children": True, - "child_page": {"title": "Child Page"}, - } - ], - "next_cursor": None, - "has_more": False, - } - mock_request.return_value = httpx.Response(200, json=mock_data) - - # Act - extractor._read_block("block-parent") - - # Assert - # Should only be called once (no recursion for child_page) - assert mock_request.call_count == 1 - - -class TestNotionProviderController: - """Tests for Notion datasource provider controller integration. - - Covers: - - Provider initialization - - Datasource retrieval - - Provider type verification - """ - - @pytest.fixture - def mock_entity(self) -> DatasourceProviderEntityWithPlugin: - """Create a validated Notion provider entity.""" - return DatasourceProviderEntityWithPlugin( - identity=DatasourceProviderIdentity( - author="langgenius", - name="notion_datasource", - icon="notion-icon.png", - description=I18nObject(en_US="Notion documents"), - label=I18nObject(en_US="Notion"), - ), - provider_type=DatasourceProviderType.ONLINE_DOCUMENT, - ) - - def test_provider_controller_initialization(self, mock_entity: DatasourceProviderEntityWithPlugin) -> None: - """Test OnlineDocumentDatasourcePluginProviderController initialization.""" - # Act - controller = OnlineDocumentDatasourcePluginProviderController( - entity=mock_entity, - plugin_id="langgenius/notion_datasource", - plugin_unique_identifier="notion-unique-id", - tenant_id="tenant-123", - ) - - # Assert - assert controller.plugin_id == "langgenius/notion_datasource" - assert controller.plugin_unique_identifier == "notion-unique-id" - assert controller.tenant_id == "tenant-123" - assert controller.provider_type == DatasourceProviderType.ONLINE_DOCUMENT - - def test_provider_controller_get_datasource(self, mock_entity: DatasourceProviderEntityWithPlugin) -> None: - """Test retrieving datasource from controller.""" - # Arrange - mock_datasource_entity = DatasourceEntity( - identity=DatasourceIdentity( - author="langgenius", - name="notion_datasource", - label=I18nObject(en_US="Notion"), - provider="notion_datasource", - ), - description=I18nObject(en_US="Notion documents"), - ) - mock_entity.datasources = [mock_datasource_entity] - - controller = OnlineDocumentDatasourcePluginProviderController( - entity=mock_entity, - plugin_id="langgenius/notion_datasource", - plugin_unique_identifier="notion-unique-id", - tenant_id="tenant-123", - ) - - # Act - datasource = controller.get_datasource("notion_datasource") - - # Assert - assert datasource is not None - assert datasource.tenant_id == "tenant-123" - - def test_provider_controller_datasource_not_found(self, mock_entity: DatasourceProviderEntityWithPlugin) -> None: - """Test error when datasource not found.""" - # Arrange - mock_entity.datasources = [] - controller = OnlineDocumentDatasourcePluginProviderController( - entity=mock_entity, - plugin_id="langgenius/notion_datasource", - plugin_unique_identifier="notion-unique-id", - tenant_id="tenant-123", - ) - - # Act & Assert - with pytest.raises(ValueError, match="Datasource with name nonexistent_datasource not found") as exc_info: - controller.get_datasource("nonexistent_datasource") - assert "not found" in str(exc_info.value) - - -class TestNotionExtractorAdvancedBlockTypes: - """Tests for advanced Notion block types and edge cases. - - Covers: - - Various block types (code, quote, lists, toggle, callout) - - Empty blocks - - Multiple rich text elements - - Mixed block types in realistic scenarios - """ - - @pytest.fixture - def extractor(self) -> NotionExtractor: - """Create a NotionExtractor instance for testing. - - Returns: - NotionExtractor: Configured extractor with test credentials - """ - return NotionExtractor( - notion_workspace_id="workspace-123", - notion_obj_id="page-456", - notion_page_type="page", - tenant_id="tenant-789", - notion_access_token="test-token", - notion_token_loader=lambda: "token", - ) - - def _create_block_with_rich_text( - self, block_id: str, block_type: str, rich_text_items: list[str], has_children: bool = False - ) -> dict[str, object]: - """Helper to create a Notion block with multiple rich text elements. - - Args: - block_id: Unique identifier for the block - block_type: Type of block (paragraph, heading_1, etc.) - rich_text_items: List of text content strings - has_children: Whether the block has child blocks - - Returns: - dict: Notion block structure with rich text elements - """ - rich_text_array = [{"type": "text", "text": {"content": text}, "plain_text": text} for text in rich_text_items] - return { - "object": "block", - "id": block_id, - "type": block_type, - "has_children": has_children, - block_type: {"rich_text": rich_text_array}, - } - - @patch("httpx.request") - def test_get_notion_block_data_with_list_blocks(self, mock_request: Mock, extractor: NotionExtractor) -> None: - """Test retrieving page with bulleted and numbered list items. - - Both list types should be extracted with their content. - """ - # Arrange - mock_data = { - "object": "list", - "results": [ - self._create_block_with_rich_text("block-1", "bulleted_list_item", ["Bullet item"]), - self._create_block_with_rich_text("block-2", "numbered_list_item", ["Numbered item"]), - ], - "next_cursor": None, - "has_more": False, - } - mock_request.return_value = httpx.Response(200, json=mock_data) - - # Act - result = extractor._get_notion_block_data("page-456") - - # Assert - assert len(result) == 2 - assert "Bullet item" in result[0] - assert "Numbered item" in result[1] - - @patch("httpx.request") - def test_get_notion_block_data_with_special_blocks(self, mock_request: Mock, extractor: NotionExtractor) -> None: - """Test retrieving page with code, quote, and callout blocks. - - Special block types should preserve their content correctly. - """ - # Arrange - mock_data = { - "object": "list", - "results": [ - self._create_block_with_rich_text("block-1", "code", ["print('code')"]), - self._create_block_with_rich_text("block-2", "quote", ["Quoted text"]), - self._create_block_with_rich_text("block-3", "callout", ["Important note"]), - ], - "next_cursor": None, - "has_more": False, - } - mock_request.return_value = httpx.Response(200, json=mock_data) - - # Act - result = extractor._get_notion_block_data("page-456") - - # Assert - assert len(result) == 3 - assert "print('code')" in result[0] - assert "Quoted text" in result[1] - assert "Important note" in result[2] - - @patch("httpx.request") - def test_get_notion_block_data_with_toggle_block(self, mock_request: Mock, extractor: NotionExtractor) -> None: - """Test retrieving page with toggle block containing children. - - Toggle blocks can have nested content that should be extracted. - """ - # Arrange - parent_data = { - "object": "list", - "results": [ - self._create_block_with_rich_text("block-1", "toggle", ["Toggle header"], has_children=True), - ], - "next_cursor": None, - "has_more": False, - } - child_data = { - "object": "list", - "results": [ - self._create_block_with_rich_text("block-child-1", "paragraph", ["Hidden content"]), - ], - "next_cursor": None, - "has_more": False, - } - mock_request.side_effect = [ - httpx.Response(200, json=parent_data), - httpx.Response(200, json=child_data), - ] - - # Act - result = extractor._get_notion_block_data("page-456") - - # Assert - assert len(result) == 1 - assert "Toggle header" in result[0] - assert "Hidden content" in result[0] - - @patch("httpx.request") - def test_get_notion_block_data_mixed_block_types(self, mock_request: Mock, extractor: NotionExtractor) -> None: - """Test retrieving page with mixed block types. - - Real Notion pages contain various block types mixed together. - This tests a realistic scenario with multiple block types. - """ - # Arrange - mock_data = { - "object": "list", - "results": [ - self._create_block_with_rich_text("block-1", "heading_1", ["Project Documentation"]), - self._create_block_with_rich_text("block-2", "paragraph", ["This is an introduction."]), - self._create_block_with_rich_text("block-3", "heading_2", ["Features"]), - self._create_block_with_rich_text("block-4", "bulleted_list_item", ["Feature A"]), - self._create_block_with_rich_text("block-5", "code", ["npm install package"]), - ], - "next_cursor": None, - "has_more": False, - } - mock_request.return_value = httpx.Response(200, json=mock_data) - - # Act - result = extractor._get_notion_block_data("page-456") - - # Assert - assert len(result) == 5 - assert "# Project Documentation" in result[0] - assert "This is an introduction" in result[1] - assert "## Features" in result[2] - assert "Feature A" in result[3] - assert "npm install package" in result[4] - - -class TestNotionExtractorDatabaseAdvanced: - """Tests for advanced database scenarios and property types. - - Covers: - - Various property types (date, number, checkbox, url, email, phone, status) - - Rich text properties - - Large database pagination - """ - - @pytest.fixture - def extractor(self) -> NotionExtractor: - """Create a NotionExtractor instance for database testing. - - Returns: - NotionExtractor: Configured extractor for database operations - """ - return NotionExtractor( - notion_workspace_id="workspace-123", - notion_obj_id="database-789", - notion_page_type="database", - tenant_id="tenant-789", - notion_access_token="test-token", - notion_token_loader=lambda: "token", - ) - - def _create_database_page_with_properties( - self, page_id: str, properties: Mapping[str, Mapping[str, object]] - ) -> dict[str, object]: - """Helper to create a database page with various property types. - - Args: - page_id: Unique identifier for the page - properties: Dictionary of property names to property configurations - - Returns: - dict: Notion database page structure - """ - formatted_properties = {} - for prop_name, prop_data in properties.items(): - prop_type = prop_data["type"] - assert isinstance(prop_type, str) - formatted_properties[prop_name] = {"type": prop_type, prop_type: prop_data["value"]} - return { - "object": "page", - "id": page_id, - "properties": formatted_properties, - "url": f"https://notion.so/{page_id}", - } - - @patch("httpx.post") - def test_get_notion_database_data_with_various_property_types( - self, mock_post: Mock, extractor: NotionExtractor - ) -> None: - """Test database with multiple property types. - - Tests date, number, checkbox, URL, email, phone, and status properties. - All property types should be extracted correctly. - """ - # Arrange - mock_response_data = { - "object": "list", - "results": [ - self._create_database_page_with_properties( - "page-1", - { - "Title": {"type": "title", "value": [{"plain_text": "Test Entry"}]}, - "Date": {"type": "date", "value": {"start": "2024-11-27", "end": None}}, - "Price": {"type": "number", "value": 99.99}, - "Completed": {"type": "checkbox", "value": True}, - "Link": {"type": "url", "value": "https://example.com"}, - "Email": {"type": "email", "value": "test@example.com"}, - "Phone": {"type": "phone_number", "value": "+1-555-0123"}, - "Status": {"type": "status", "value": {"name": "Active"}}, - }, - ), - ], - "has_more": False, - "next_cursor": None, - } - mock_response = httpx.Response(200, json=mock_response_data) - mock_post.return_value = mock_response - - # Act - result = extractor._get_notion_database_data("database-789") - - # Assert - assert len(result) == 1 - content = result[0].page_content - assert "Title:Test Entry" in content - assert "Date:" in content - assert "Price:99.99" in content - assert "Completed:True" in content - assert "Link:https://example.com" in content - assert "Email:test@example.com" in content - assert "Phone:+1-555-0123" in content - assert "Status:Active" in content - - @patch("httpx.post") - def test_get_notion_database_data_large_pagination(self, mock_post: Mock, extractor: NotionExtractor) -> None: - """Test database with multiple pages of results. - - Large databases require multiple API calls with cursor-based pagination. - This tests that all pages are retrieved correctly. - """ - # Arrange - Create 3 pages of results - page1_response_data = { - "object": "list", - "results": [ - self._create_database_page_with_properties( - f"page-{i}", {"Title": {"type": "title", "value": [{"plain_text": f"Item {i}"}]}} - ) - for i in range(1, 4) - ], - "has_more": True, - "next_cursor": "cursor-1", - } - page1_response = httpx.Response(200, json=page1_response_data) - - page2_response_data = { - "object": "list", - "results": [ - self._create_database_page_with_properties( - f"page-{i}", {"Title": {"type": "title", "value": [{"plain_text": f"Item {i}"}]}} - ) - for i in range(4, 7) - ], - "has_more": True, - "next_cursor": "cursor-2", - } - page2_response = httpx.Response(200, json=page2_response_data) - - page3_response_data = { - "object": "list", - "results": [ - self._create_database_page_with_properties( - f"page-{i}", {"Title": {"type": "title", "value": [{"plain_text": f"Item {i}"}]}} - ) - for i in range(7, 10) - ], - "has_more": False, - "next_cursor": None, - } - page3_response = httpx.Response(200, json=page3_response_data) - - mock_post.side_effect = [page1_response, page2_response, page3_response] - - # Act - result = extractor._get_notion_database_data("database-789") - - # Assert - assert len(result) == 1 - content = result[0].page_content - # Verify all items from all pages are present - for i in range(1, 10): - assert f"Title:Item {i}" in content - # Verify pagination was called correctly - assert mock_post.call_count == 3 - - @patch("httpx.post") - def test_get_notion_database_data_with_rich_text_property( - self, mock_post: Mock, extractor: NotionExtractor - ) -> None: - """Test database with rich_text property type. - - Rich text properties can contain formatted text and should be extracted. - """ - # Arrange - mock_response_data = { - "object": "list", - "results": [ - self._create_database_page_with_properties( - "page-1", - { - "Title": {"type": "title", "value": [{"plain_text": "Note"}]}, - "Description": { - "type": "rich_text", - "value": [{"plain_text": "This is a detailed description"}], - }, - }, - ), - ], - "has_more": False, - "next_cursor": None, - } - mock_response = httpx.Response(200, json=mock_response_data) - mock_post.return_value = mock_response - - # Act - result = extractor._get_notion_database_data("database-789") - - # Assert - assert len(result) == 1 - content = result[0].page_content - assert "Title:Note" in content - assert "Description:This is a detailed description" in content - - -class TestNotionExtractorErrorScenarios: - """Tests for error handling and edge cases. - - Covers: - - Network timeouts - - Rate limiting - - Invalid tokens - - Malformed responses - - Missing required fields - - API version mismatches - """ - - @pytest.fixture - def extractor(self) -> NotionExtractor: - """Create a NotionExtractor instance for error testing. - - Returns: - NotionExtractor: Configured extractor for error scenarios - """ - return NotionExtractor( - notion_workspace_id="workspace-123", - notion_obj_id="page-456", - notion_page_type="page", - tenant_id="tenant-789", - notion_access_token="test-token", - notion_token_loader=lambda: "token", - ) - - @pytest.mark.parametrize( - "error_value", - [httpx.TimeoutException("Request timed out"), httpx.ConnectError("Connection failed")], - ids=["timeout", "connection"], - ) - @patch("httpx.request") - def test_get_notion_block_data_network_errors( - self, mock_request: Mock, extractor: NotionExtractor, error_value: httpx.TransportError - ) -> None: - """Test handling of various network errors. - - Network issues (timeouts, connection failures) should raise appropriate errors. - """ - # Arrange - mock_request.side_effect = error_value - - # Act & Assert - with pytest.raises(ValueError, match="Error fetching Notion block data") as exc_info: - extractor._get_notion_block_data("page-456") - assert "Error fetching Notion block data" in str(exc_info.value) - - @pytest.mark.parametrize( - ("status_code", "description"), - [ - (401, "Unauthorized"), - (403, "Forbidden"), - (404, "Not Found"), - (429, "Rate limit exceeded"), - ], - ) - @patch("httpx.request") - def test_get_notion_block_data_http_status_errors( - self, mock_request: Mock, extractor: NotionExtractor, status_code: int, description: str - ) -> None: - """Test handling of various HTTP status errors. - - Different HTTP error codes (401, 403, 404, 429) should be handled appropriately. - """ - # Arrange - mock_response = httpx.Response(status_code, text=description) - mock_request.return_value = mock_response - - # Act & Assert - with pytest.raises( - ValueError, - match=( - "(?:Error fetching Notion block data: Forbidden|Error fetching Notion " - "block data: Not Found|Error fetching Notion block data: Rate limit " - "exceeded|Error fetching Notion block data: Unauthorized)" - ), - ) as exc_info: - extractor._get_notion_block_data("page-456") - assert "Error fetching Notion block data" in str(exc_info.value) - - @pytest.mark.parametrize( - "response_data", - [ - {"object": "list"}, - {"object": "list", "results": "not a list"}, - {"object": "list", "results": None}, - ], - ids=["missing results field", "results not a list", "results is None"], - ) - @patch("httpx.request") - def test_get_notion_block_data_malformed_responses( - self, mock_request: Mock, extractor: NotionExtractor, response_data: dict[str, object] - ) -> None: - """Test handling of malformed API responses. - - Various malformed responses should be handled gracefully. - """ - # Arrange - mock_response = httpx.Response(200, json=response_data) - mock_request.return_value = mock_response - - # Act & Assert - with pytest.raises(ValueError, match="Error fetching Notion block data") as exc_info: - extractor._get_notion_block_data("page-456") - assert "Error fetching Notion block data" in str(exc_info.value) - - @patch("httpx.post") - def test_get_notion_database_data_with_query_filter(self, mock_post: Mock, extractor: NotionExtractor) -> None: - """Test database query with custom filter. - - Databases can be queried with filters to retrieve specific rows. - """ - # Arrange - mock_response_data = { - "object": "list", - "results": [ - { - "object": "page", - "id": "page-1", - "properties": { - "Title": {"type": "title", "title": [{"plain_text": "Filtered Item"}]}, - "Status": {"type": "select", "select": {"name": "Active"}}, - }, - "url": "https://notion.so/page-1", - } - ], - "has_more": False, - "next_cursor": None, - } - mock_response = httpx.Response(200, json=mock_response_data) - mock_post.return_value = mock_response - - # Create a custom query filter - query_filter = {"filter": {"property": "Status", "select": {"equals": "Active"}}} - - # Act - result = extractor._get_notion_database_data("database-789", query_dict=query_filter) - - # Assert - assert len(result) == 1 - content = result[0].page_content - assert "Title:Filtered Item" in content - assert "Status:Active" in content - # Verify the filter was passed to the API - mock_post.assert_called_once() - call_args = mock_post.call_args - assert "filter" in call_args[1]["json"] - - -class TestNotionExtractorTableAdvanced: - """Tests for advanced table scenarios. - - Covers: - - Tables with many columns - - Tables with complex cell content - - Empty tables - """ - - @pytest.fixture - def extractor(self) -> NotionExtractor: - """Create a NotionExtractor instance for table testing. - - Returns: - NotionExtractor: Configured extractor for table operations - """ - return NotionExtractor( - notion_workspace_id="workspace-123", - notion_obj_id="page-456", - notion_page_type="page", - tenant_id="tenant-789", - notion_access_token="test-token", - notion_token_loader=lambda: "token", - ) - - @patch("httpx.request") - def test_read_table_rows_with_many_columns(self, mock_request: Mock, extractor: NotionExtractor) -> None: - """Test reading table with many columns. - - Tables can have numerous columns; all should be extracted correctly. - """ - # Arrange - Create a table with 10 columns - headers = [f"Col{i}" for i in range(1, 11)] - values = [f"Val{i}" for i in range(1, 11)] - - mock_data = { - "object": "list", - "results": [ - { - "object": "block", - "type": "table_row", - "table_row": {"cells": [[{"text": {"content": h}}] for h in headers]}, - }, - { - "object": "block", - "type": "table_row", - "table_row": {"cells": [[{"text": {"content": v}}] for v in values]}, - }, - ], - "next_cursor": None, - "has_more": False, - } - mock_request.return_value = httpx.Response(200, json=mock_data) - - # Act - result = extractor._read_table_rows("table-block-123") - - # Assert - for header in headers: - assert header in result - for value in values: - assert value in result - # Verify markdown table structure - assert "| --- |" in result diff --git a/api/tests/unit_tests/core/datasource/test_stream_node_events.py b/api/tests/unit_tests/core/datasource/test_stream_node_events.py deleted file mode 100644 index 1b3005af1ac..00000000000 --- a/api/tests/unit_tests/core/datasource/test_stream_node_events.py +++ /dev/null @@ -1,46 +0,0 @@ -"""Unit coverage for datasource event-stream accumulation.""" - -from collections.abc import Generator - -from pytest_mock import MockerFixture - -from core.datasource.datasource_manager import DatasourceManager -from core.datasource.entities.datasource_entities import DatasourceMessage -from graphon.node_events import StreamCompletedEvent -from graphon.runtime import VariablePool - - -def _gen_var_stream() -> Generator[DatasourceMessage]: - # produce a streamed variable "a"="xy" - yield DatasourceMessage( - type=DatasourceMessage.MessageType.VARIABLE, - message=DatasourceMessage.VariableMessage(variable_name="a", variable_value="x", stream=True), - meta=None, - ) - yield DatasourceMessage( - type=DatasourceMessage.MessageType.VARIABLE, - message=DatasourceMessage.VariableMessage(variable_name="a", variable_value="y", stream=True), - meta=None, - ) - - -def test_stream_node_events_accumulates_variables(mocker: MockerFixture) -> None: - mocker.patch.object(DatasourceManager, "stream_online_results", return_value=_gen_var_stream()) - events = list( - DatasourceManager.stream_node_events( - node_id="A", - user_id="u", - datasource_name="ds", - datasource_type="online_document", - provider_id="p/x", - tenant_id="t", - parameters_for_log={}, - datasource_info={"user_id": "u"}, - variable_pool=VariablePool(), - datasource_param=None, - online_drive_request=None, - credentials={}, - ) - ) - assert isinstance(events[-1], StreamCompletedEvent) - assert events[-1].node_run_result.outputs["a"] == "xy" diff --git a/api/tests/unit_tests/core/datasource/test_website_crawl.py b/api/tests/unit_tests/core/datasource/test_website_crawl.py deleted file mode 100644 index 20b8fd32089..00000000000 --- a/api/tests/unit_tests/core/datasource/test_website_crawl.py +++ /dev/null @@ -1,1786 +0,0 @@ -""" -Unit tests for website crawling functionality. - -This module tests the core website crawling features including: -- URL crawling logic with different providers -- Robots.txt respect and compliance -- Max depth limiting for crawl operations -- Content extraction from web pages -- Link following logic and navigation - -The tests cover multiple crawl providers (Firecrawl, WaterCrawl, JinaReader) -and ensure proper handling of crawl options, status checking, and data retrieval. -""" - -from unittest.mock import Mock, create_autospec, patch - -import httpx -import pytest -from pytest_mock import MockerFixture - -import services.data_source.website_service as website_service_module -from core.datasource.entities.datasource_entities import ( - DatasourceEntity, - DatasourceIdentity, - DatasourceProviderEntityWithPlugin, - DatasourceProviderIdentity, - DatasourceProviderType, -) -from core.datasource.website_crawl.website_crawl_plugin import WebsiteCrawlDatasourcePlugin -from core.datasource.website_crawl.website_crawl_provider import WebsiteCrawlDatasourcePluginProviderController -from core.rag.extractor.watercrawl.provider import WaterCrawlProvider -from services.data_source.credential_gateway import DatasourceProviderCredentialStore -from services.data_source.website_service import CrawlOptions, CrawlRequest, WebsiteService - -# ============================================================================ -# Fixtures -# ============================================================================ - - -@pytest.fixture -def mock_datasource_entity() -> DatasourceEntity: - """Create a mock datasource entity for testing.""" - return DatasourceEntity( - identity=DatasourceIdentity( - author="test_author", - name="test_datasource", - label={"en_US": "Test Datasource", "zh_Hans": "测试数据源"}, - provider="test_provider", - icon="test_icon.svg", - ), - parameters=[], - description={"en_US": "Test datasource description", "zh_Hans": "测试数据源描述"}, - ) - - -@pytest.fixture -def mock_provider_entity(mock_datasource_entity: DatasourceEntity) -> DatasourceProviderEntityWithPlugin: - """Create a mock provider entity with plugin for testing.""" - return DatasourceProviderEntityWithPlugin( - identity=DatasourceProviderIdentity( - author="test_author", - name="test_provider", - description={"en_US": "Test Provider", "zh_Hans": "测试提供者"}, - icon="test_icon.svg", - label={"en_US": "Test Provider", "zh_Hans": "测试提供者"}, - ), - credentials_schema=[], - provider_type=DatasourceProviderType.WEBSITE_CRAWL, - datasources=[mock_datasource_entity], - ) - - -@pytest.fixture -def crawl_options() -> CrawlOptions: - """Create default crawl options for testing.""" - return CrawlOptions( - limit=10, - crawl_sub_pages=True, - only_main_content=True, - includes="/blog/*,/docs/*", - excludes="/admin/*,/private/*", - max_depth=3, - use_sitemap=True, - ) - - -@pytest.fixture -def crawl_request(crawl_options: CrawlOptions) -> CrawlRequest: - """Create a crawl request for testing.""" - return CrawlRequest(url="https://example.com", provider="watercrawl", options=crawl_options) - - -# ============================================================================ -# Test CrawlOptions -# ============================================================================ - - -class TestCrawlOptions: - """Test suite for CrawlOptions data class.""" - - def test_crawl_options_defaults(self) -> None: - """Test that CrawlOptions has correct default values.""" - options = CrawlOptions() - - assert options.limit == 1 - assert options.crawl_sub_pages is False - assert options.only_main_content is False - assert options.includes is None - assert options.excludes is None - assert options.prompt is None - assert options.max_depth is None - assert options.use_sitemap is True - - def test_get_include_paths_with_values(self, crawl_options: CrawlOptions) -> None: - """Test parsing include paths from comma-separated string.""" - paths = crawl_options.get_include_paths() - - assert len(paths) == 2 - assert "/blog/*" in paths - assert "/docs/*" in paths - - def test_get_include_paths_empty(self) -> None: - """Test that empty includes returns empty list.""" - options = CrawlOptions(includes=None) - paths = options.get_include_paths() - - assert paths == [] - - def test_get_exclude_paths_with_values(self, crawl_options: CrawlOptions) -> None: - """Test parsing exclude paths from comma-separated string.""" - paths = crawl_options.get_exclude_paths() - - assert len(paths) == 2 - assert "/admin/*" in paths - assert "/private/*" in paths - - def test_get_exclude_paths_empty(self) -> None: - """Test that empty excludes returns empty list.""" - options = CrawlOptions(excludes=None) - paths = options.get_exclude_paths() - - assert paths == [] - - def test_max_depth_limiting(self) -> None: - """Test that max_depth can be set to limit crawl depth.""" - options = CrawlOptions(max_depth=5, crawl_sub_pages=True) - - assert options.max_depth == 5 - assert options.crawl_sub_pages is True - - -# ============================================================================ -# Test WebsiteCrawlDatasourcePlugin -# ============================================================================ - - -class TestWebsiteCrawlDatasourcePlugin: - """Test suite for WebsiteCrawlDatasourcePlugin.""" - - def test_plugin_initialization(self, mock_datasource_entity: DatasourceEntity) -> None: - """Test that plugin initializes correctly with required parameters.""" - from core.datasource.__base.datasource_runtime import DatasourceRuntime - - runtime = DatasourceRuntime(tenant_id="test_tenant", credentials={}) - plugin = WebsiteCrawlDatasourcePlugin( - entity=mock_datasource_entity, - runtime=runtime, - tenant_id="test_tenant", - icon="test_icon.svg", - plugin_unique_identifier="test_plugin_id", - ) - - assert plugin.tenant_id == "test_tenant" - assert plugin.plugin_unique_identifier == "test_plugin_id" - assert plugin.entity == mock_datasource_entity - assert plugin.datasource_provider_type() == DatasourceProviderType.WEBSITE_CRAWL - - def test_get_website_crawl(self, mock_datasource_entity: DatasourceEntity, mocker: MockerFixture) -> None: - """Test that get_website_crawl calls PluginDatasourceManager correctly.""" - from core.datasource.__base.datasource_runtime import DatasourceRuntime - - runtime = DatasourceRuntime(tenant_id="test_tenant", credentials={"api_key": "test_key"}) - plugin = WebsiteCrawlDatasourcePlugin( - entity=mock_datasource_entity, - runtime=runtime, - tenant_id="test_tenant", - icon="test_icon.svg", - plugin_unique_identifier="test_plugin_id", - ) - - # Mock the PluginDatasourceManager - mock_manager = mocker.patch("core.datasource.website_crawl.website_crawl_plugin.PluginDatasourceManager") - mock_instance = mock_manager.return_value - mock_instance.get_website_crawl.return_value = iter([]) - - datasource_params = {"url": "https://example.com", "max_depth": 2} - - plugin.get_website_crawl( - user_id="test_user", datasource_parameters=datasource_params, provider_type="watercrawl" - ) - - # Verify the manager was called with correct parameters - mock_instance.get_website_crawl.assert_called_once_with( - tenant_id="test_tenant", - user_id="test_user", - datasource_provider=mock_datasource_entity.identity.provider, - datasource_name=mock_datasource_entity.identity.name, - credentials={"api_key": "test_key"}, - datasource_parameters=datasource_params, - provider_type="watercrawl", - ) - - -# ============================================================================ -# Test WebsiteCrawlDatasourcePluginProviderController -# ============================================================================ - - -class TestWebsiteCrawlDatasourcePluginProviderController: - """Test suite for WebsiteCrawlDatasourcePluginProviderController.""" - - def test_provider_controller_initialization(self, mock_provider_entity: DatasourceProviderEntityWithPlugin) -> None: - """Test provider controller initialization.""" - controller = WebsiteCrawlDatasourcePluginProviderController( - entity=mock_provider_entity, - plugin_id="test_plugin_id", - plugin_unique_identifier="test_unique_id", - tenant_id="test_tenant", - ) - - assert controller.plugin_id == "test_plugin_id" - assert controller.plugin_unique_identifier == "test_unique_id" - assert controller.provider_type == DatasourceProviderType.WEBSITE_CRAWL - - def test_get_datasource_success(self, mock_provider_entity: DatasourceProviderEntityWithPlugin) -> None: - """Test retrieving a datasource by name.""" - controller = WebsiteCrawlDatasourcePluginProviderController( - entity=mock_provider_entity, - plugin_id="test_plugin_id", - plugin_unique_identifier="test_unique_id", - tenant_id="test_tenant", - ) - - datasource = controller.get_datasource("test_datasource") - - assert isinstance(datasource, WebsiteCrawlDatasourcePlugin) - assert datasource.tenant_id == "test_tenant" - assert datasource.plugin_unique_identifier == "test_unique_id" - - def test_get_datasource_not_found(self, mock_provider_entity: DatasourceProviderEntityWithPlugin) -> None: - """Test that ValueError is raised when datasource is not found.""" - controller = WebsiteCrawlDatasourcePluginProviderController( - entity=mock_provider_entity, - plugin_id="test_plugin_id", - plugin_unique_identifier="test_unique_id", - tenant_id="test_tenant", - ) - - with pytest.raises(ValueError, match="Datasource with name nonexistent not found"): - controller.get_datasource("nonexistent") - - -# ============================================================================ -# Test WaterCrawl Provider - URL Crawling Logic -# ============================================================================ - - -class TestWaterCrawlProvider: - """Test suite for WaterCrawl provider crawling functionality.""" - - def test_crawl_url_basic(self, mocker: MockerFixture) -> None: - """Test basic URL crawling without sub-pages.""" - mock_client = mocker.patch("core.rag.extractor.watercrawl.provider.WaterCrawlAPIClient") - mock_instance = mock_client.return_value - mock_instance.create_crawl_request.return_value = {"uuid": "test-job-123"} - - provider = WaterCrawlProvider(api_key="test_key") - result = provider.crawl_url("https://example.com", options={"crawl_sub_pages": False}) - - assert result["status"] == "active" - assert result["job_id"] == "test-job-123" - - # Verify spider options for single page crawl - call_args = mock_instance.create_crawl_request.call_args - spider_options = call_args.kwargs["spider_options"] - assert spider_options["max_depth"] == 1 - assert spider_options["page_limit"] == 1 - - def test_crawl_url_with_sub_pages(self, mocker: MockerFixture) -> None: - """Test URL crawling with sub-pages enabled.""" - mock_client = mocker.patch("core.rag.extractor.watercrawl.provider.WaterCrawlAPIClient") - mock_instance = mock_client.return_value - mock_instance.create_crawl_request.return_value = {"uuid": "test-job-456"} - - provider = WaterCrawlProvider(api_key="test_key") - options = {"crawl_sub_pages": True, "limit": 50, "max_depth": 3} - result = provider.crawl_url("https://example.com", options=options) - - assert result["status"] == "active" - assert result["job_id"] == "test-job-456" - - # Verify spider options for multi-page crawl - call_args = mock_instance.create_crawl_request.call_args - spider_options = call_args.kwargs["spider_options"] - assert spider_options["max_depth"] == 3 - assert spider_options["page_limit"] == 50 - - def test_crawl_url_max_depth_limiting(self, mocker: MockerFixture) -> None: - """Test that max_depth properly limits crawl depth.""" - mock_client = mocker.patch("core.rag.extractor.watercrawl.provider.WaterCrawlAPIClient") - mock_instance = mock_client.return_value - mock_instance.create_crawl_request.return_value = {"uuid": "test-job-789"} - - provider = WaterCrawlProvider(api_key="test_key") - - # Test with max_depth of 2 - options = {"crawl_sub_pages": True, "max_depth": 2, "limit": 100} - provider.crawl_url("https://example.com", options=options) - - call_args = mock_instance.create_crawl_request.call_args - spider_options = call_args.kwargs["spider_options"] - assert spider_options["max_depth"] == 2 - - def test_crawl_url_with_include_exclude_paths(self, mocker: MockerFixture) -> None: - """Test URL crawling with include and exclude path filters.""" - mock_client = mocker.patch("core.rag.extractor.watercrawl.provider.WaterCrawlAPIClient") - mock_instance = mock_client.return_value - mock_instance.create_crawl_request.return_value = {"uuid": "test-job-101"} - - provider = WaterCrawlProvider(api_key="test_key") - options = { - "crawl_sub_pages": True, - "includes": "/blog/*,/docs/*", - "excludes": "/admin/*,/private/*", - "limit": 20, - } - provider.crawl_url("https://example.com", options=options) - - call_args = mock_instance.create_crawl_request.call_args - spider_options = call_args.kwargs["spider_options"] - - # Verify include paths - assert len(spider_options["include_paths"]) == 2 - assert "/blog/*" in spider_options["include_paths"] - assert "/docs/*" in spider_options["include_paths"] - - # Verify exclude paths - assert len(spider_options["exclude_paths"]) == 2 - assert "/admin/*" in spider_options["exclude_paths"] - assert "/private/*" in spider_options["exclude_paths"] - - def test_crawl_url_content_extraction_options(self, mocker: MockerFixture) -> None: - """Test that content extraction options are properly configured.""" - mock_client = mocker.patch("core.rag.extractor.watercrawl.provider.WaterCrawlAPIClient") - mock_instance = mock_client.return_value - mock_instance.create_crawl_request.return_value = {"uuid": "test-job-202"} - - provider = WaterCrawlProvider(api_key="test_key") - options = {"only_main_content": True, "wait_time": 2000} - provider.crawl_url("https://example.com", options=options) - - call_args = mock_instance.create_crawl_request.call_args - page_options = call_args.kwargs["page_options"] - - # Verify content extraction settings - assert page_options["only_main_content"] is True - assert page_options["wait_time"] == 2000 - assert page_options["include_html"] is False - - def test_crawl_url_minimum_wait_time(self, mocker: MockerFixture) -> None: - """Test that wait_time has a minimum value of 1000ms.""" - mock_client = mocker.patch("core.rag.extractor.watercrawl.provider.WaterCrawlAPIClient") - mock_instance = mock_client.return_value - mock_instance.create_crawl_request.return_value = {"uuid": "test-job-303"} - - provider = WaterCrawlProvider(api_key="test_key") - options = {"wait_time": 500} # Below minimum - provider.crawl_url("https://example.com", options=options) - - call_args = mock_instance.create_crawl_request.call_args - page_options = call_args.kwargs["page_options"] - - # Should be clamped to minimum of 1000 - assert page_options["wait_time"] == 1000 - - -# ============================================================================ -# Test Crawl Status and Results -# ============================================================================ - - -class TestCrawlStatus: - """Test suite for crawl status checking and result retrieval.""" - - def test_get_crawl_status_active(self, mocker: MockerFixture) -> None: - """Test getting status of an active crawl job.""" - mock_client = mocker.patch("core.rag.extractor.watercrawl.provider.WaterCrawlAPIClient") - mock_instance = mock_client.return_value - mock_instance.get_crawl_request.return_value = { - "uuid": "test-job-123", - "status": "running", - "number_of_documents": 5, - "options": {"spider_options": {"page_limit": 10}}, - "duration": None, - } - - provider = WaterCrawlProvider(api_key="test_key") - status = provider.get_crawl_status("test-job-123") - - assert status["status"] == "active" - assert status["job_id"] == "test-job-123" - assert status["total"] == 10 - assert status["current"] == 5 - assert status["data"] == [] - - def test_get_crawl_status_completed(self, mocker: MockerFixture) -> None: - """Test getting status of a completed crawl job with results.""" - mock_client = mocker.patch("core.rag.extractor.watercrawl.provider.WaterCrawlAPIClient") - mock_instance = mock_client.return_value - mock_instance.get_crawl_request.return_value = { - "uuid": "test-job-456", - "status": "completed", - "number_of_documents": 10, - "options": {"spider_options": {"page_limit": 10}}, - "duration": "00:00:15.500000", - } - mock_instance.get_crawl_request_results.return_value = { - "results": [ - { - "url": "https://example.com/page1", - "result": { - "markdown": "# Page 1 Content", - "metadata": {"title": "Page 1", "description": "First page"}, - }, - } - ], - "next": None, - } - - provider = WaterCrawlProvider(api_key="test_key") - status = provider.get_crawl_status("test-job-456") - - assert status["status"] == "completed" - assert status["job_id"] == "test-job-456" - assert status["total"] == 10 - assert status["current"] == 10 - assert len(status["data"]) == 1 - assert status["time_consuming"] == 15.5 - - def test_get_crawl_url_data(self, mocker: MockerFixture) -> None: - """Test retrieving specific URL data from crawl results.""" - mock_client = mocker.patch("core.rag.extractor.watercrawl.provider.WaterCrawlAPIClient") - mock_instance = mock_client.return_value - mock_instance.get_crawl_request_results.return_value = { - "results": [ - { - "url": "https://example.com/target", - "result": { - "markdown": "# Target Page", - "metadata": {"title": "Target", "description": "Target page description"}, - }, - } - ], - "next": None, - } - - provider = WaterCrawlProvider(api_key="test_key") - data = provider.get_crawl_url_data("test-job-789", "https://example.com/target") - - assert data is not None - assert data["source_url"] == "https://example.com/target" - assert data["title"] == "Target" - assert data["markdown"] == "# Target Page" - - def test_get_crawl_url_data_not_found(self, mocker: MockerFixture) -> None: - """Test that None is returned when URL is not in results.""" - mock_client = mocker.patch("core.rag.extractor.watercrawl.provider.WaterCrawlAPIClient") - mock_instance = mock_client.return_value - mock_instance.get_crawl_request_results.return_value = {"results": list[object](), "next": None} - - provider = WaterCrawlProvider(api_key="test_key") - data = provider.get_crawl_url_data("test-job-789", "https://example.com/nonexistent") - - assert data is None - - -# ============================================================================ -# Test WebsiteService - Multi-Provider Support -# ============================================================================ - - -class TestWebsiteService: - """Test suite for WebsiteService with multiple providers.""" - - @patch("services.data_source.website_service.current_user") - @patch("services.data_source.website_service.DatasourceProviderService") - def test_crawl_url_firecrawl( - self, mock_provider_service: Mock, mock_current_user: Mock, mocker: MockerFixture - ) -> None: - """Test crawling with Firecrawl provider.""" - # Setup mocks - mock_current_user.current_tenant_id = "test_tenant" - mock_provider_service.return_value.get_datasource_credentials.return_value = { - "firecrawl_api_key": "test_key", - "base_url": "https://api.firecrawl.dev", - } - - mock_firecrawl = mocker.patch("services.data_source.website_service.FirecrawlApp") - mock_firecrawl_instance = mock_firecrawl.return_value - mock_firecrawl_instance.crawl_url.return_value = "job-123" - - # Mock redis - mocker.patch("services.data_source.website_service.redis_client") - - from services.data_source.website_service import WebsiteCrawlApiRequest - - api_request = WebsiteCrawlApiRequest( - provider="firecrawl", - url="https://example.com", - options={"limit": 10, "crawl_sub_pages": True, "only_main_content": True}, - ) - - result = WebsiteService( - providers=website_service_module.DatasourceProviderService( - credentials=create_autospec(DatasourceProviderCredentialStore, instance=True) - ) - ).crawl_url(api_request) - - assert result["status"] == "active" - assert result["job_id"] == "job-123" - - @patch("services.data_source.website_service.current_user") - @patch("services.data_source.website_service.DatasourceProviderService") - def test_crawl_url_watercrawl( - self, mock_provider_service: Mock, mock_current_user: Mock, mocker: MockerFixture - ) -> None: - """Test crawling with WaterCrawl provider.""" - # Setup mocks - mock_current_user.current_tenant_id = "test_tenant" - mock_provider_service.return_value.get_datasource_credentials.return_value = { - "api_key": "test_key", - "base_url": "https://app.watercrawl.dev", - } - - mock_watercrawl = mocker.patch("services.data_source.website_service.WaterCrawlProvider") - mock_watercrawl_instance = mock_watercrawl.return_value - mock_watercrawl_instance.crawl_url.return_value = {"status": "active", "job_id": "job-456"} - - from services.data_source.website_service import WebsiteCrawlApiRequest - - api_request = WebsiteCrawlApiRequest( - provider="watercrawl", - url="https://example.com", - options={"limit": 20, "crawl_sub_pages": True, "max_depth": 2}, - ) - - result = WebsiteService( - providers=website_service_module.DatasourceProviderService( - credentials=create_autospec(DatasourceProviderCredentialStore, instance=True) - ) - ).crawl_url(api_request) - - assert result["status"] == "active" - assert result["job_id"] == "job-456" - - @patch("services.data_source.website_service.current_user") - @patch("services.data_source.website_service.DatasourceProviderService") - def test_crawl_url_jinareader( - self, mock_provider_service: Mock, mock_current_user: Mock, mocker: MockerFixture - ) -> None: - """Test crawling with JinaReader provider.""" - # Setup mocks - mock_current_user.current_tenant_id = "test_tenant" - mock_provider_service.return_value.get_datasource_credentials.return_value = { - "api_key": "test_key", - } - - mock_response = httpx.Response(200, json={"code": 200, "data": {"taskId": "task-789"}}) - mocker.patch( - "services.data_source.website_service._adaptive_http_client.post", - return_value=mock_response, - ) - - from services.data_source.website_service import WebsiteCrawlApiRequest - - api_request = WebsiteCrawlApiRequest( - provider="jinareader", - url="https://example.com", - options={"limit": 15, "crawl_sub_pages": True, "use_sitemap": True}, - ) - - result = WebsiteService( - providers=website_service_module.DatasourceProviderService( - credentials=create_autospec(DatasourceProviderCredentialStore, instance=True) - ) - ).crawl_url(api_request) - - assert result["status"] == "active" - assert result["job_id"] == "task-789" - - def test_document_create_args_validate_success(self) -> None: - """Test validation of valid document creation arguments.""" - args = {"provider": "watercrawl", "url": "https://example.com", "options": {"limit": 10}} - - # Should not raise any exception - WebsiteService.document_create_args_validate(args) - - def test_document_create_args_validate_missing_provider(self) -> None: - """Test validation fails when provider is missing.""" - args = {"url": "https://example.com", "options": {"limit": 10}} - - with pytest.raises(ValueError, match="Provider is required"): - WebsiteService.document_create_args_validate(args) - - def test_document_create_args_validate_missing_url(self) -> None: - """Test validation fails when URL is missing.""" - args = {"provider": "watercrawl", "options": {"limit": 10}} - - with pytest.raises(ValueError, match="URL is required"): - WebsiteService.document_create_args_validate(args) - - def test_document_create_args_validate_missing_options(self) -> None: - """Test validation fails when options are missing.""" - args = {"provider": "watercrawl", "url": "https://example.com"} - - with pytest.raises(ValueError, match="Options are required"): - WebsiteService.document_create_args_validate(args) - - -# ============================================================================ -# Test Link Following Logic -# ============================================================================ - - -class TestLinkFollowingLogic: - """Test suite for link following and navigation logic.""" - - def test_link_following_with_includes(self, mocker: MockerFixture) -> None: - """Test that only links matching include patterns are followed.""" - mock_client = mocker.patch("core.rag.extractor.watercrawl.provider.WaterCrawlAPIClient") - mock_instance = mock_client.return_value - mock_instance.create_crawl_request.return_value = {"uuid": "test-job"} - - provider = WaterCrawlProvider(api_key="test_key") - options = {"crawl_sub_pages": True, "includes": "/blog/*,/news/*", "limit": 50} - provider.crawl_url("https://example.com", options=options) - - call_args = mock_instance.create_crawl_request.call_args - spider_options = call_args.kwargs["spider_options"] - - # Verify include paths are set for link filtering - assert "/blog/*" in spider_options["include_paths"] - assert "/news/*" in spider_options["include_paths"] - - def test_link_following_with_excludes(self, mocker: MockerFixture) -> None: - """Test that links matching exclude patterns are not followed.""" - mock_client = mocker.patch("core.rag.extractor.watercrawl.provider.WaterCrawlAPIClient") - mock_instance = mock_client.return_value - mock_instance.create_crawl_request.return_value = {"uuid": "test-job"} - - provider = WaterCrawlProvider(api_key="test_key") - options = {"crawl_sub_pages": True, "excludes": "/login/*,/logout/*", "limit": 50} - provider.crawl_url("https://example.com", options=options) - - call_args = mock_instance.create_crawl_request.call_args - spider_options = call_args.kwargs["spider_options"] - - # Verify exclude paths are set to prevent following certain links - assert "/login/*" in spider_options["exclude_paths"] - assert "/logout/*" in spider_options["exclude_paths"] - - def test_link_following_respects_max_depth(self, mocker: MockerFixture) -> None: - """Test that link following stops at specified max depth.""" - mock_client = mocker.patch("core.rag.extractor.watercrawl.provider.WaterCrawlAPIClient") - mock_instance = mock_client.return_value - mock_instance.create_crawl_request.return_value = {"uuid": "test-job"} - - provider = WaterCrawlProvider(api_key="test_key") - - # Test depth of 1 (only start page) - options = {"crawl_sub_pages": True, "max_depth": 1, "limit": 100} - provider.crawl_url("https://example.com", options=options) - - call_args = mock_instance.create_crawl_request.call_args - spider_options = call_args.kwargs["spider_options"] - assert spider_options["max_depth"] == 1 - - def test_link_following_page_limit(self, mocker: MockerFixture) -> None: - """Test that link following respects page limit.""" - mock_client = mocker.patch("core.rag.extractor.watercrawl.provider.WaterCrawlAPIClient") - mock_instance = mock_client.return_value - mock_instance.create_crawl_request.return_value = {"uuid": "test-job"} - - provider = WaterCrawlProvider(api_key="test_key") - options = {"crawl_sub_pages": True, "limit": 25, "max_depth": 5} - provider.crawl_url("https://example.com", options=options) - - call_args = mock_instance.create_crawl_request.call_args - spider_options = call_args.kwargs["spider_options"] - - # Verify page limit is set correctly - assert spider_options["page_limit"] == 25 - - -# ============================================================================ -# Test Robots.txt Respect (Implicit in Provider Implementation) -# ============================================================================ - - -class TestRobotsTxtRespect: - """ - Test suite for robots.txt compliance. - - Note: Robots.txt respect is typically handled by the underlying crawl - providers (Firecrawl, WaterCrawl, JinaReader). These tests verify that - the service layer properly configures providers to respect robots.txt. - """ - - def test_watercrawl_provider_respects_robots_txt(self, mocker: MockerFixture) -> None: - """ - Test that WaterCrawl provider is configured to respect robots.txt. - - WaterCrawl respects robots.txt by default in its implementation. - This test verifies the provider is initialized correctly. - """ - mock_client = mocker.patch("core.rag.extractor.watercrawl.provider.WaterCrawlAPIClient") - - provider = WaterCrawlProvider(api_key="test_key", base_url="https://app.watercrawl.dev/") - - # Verify provider is initialized with proper client - assert provider.client is not None - mock_client.assert_called_once_with("test_key", "https://app.watercrawl.dev/") - - def test_firecrawl_provider_respects_robots_txt(self) -> None: - """ - Test that Firecrawl provider respects robots.txt. - - Firecrawl respects robots.txt by default. This test ensures - the provider is configured correctly. - """ - from core.rag.extractor.firecrawl.firecrawl_app import FirecrawlApp - - # FirecrawlApp respects robots.txt in its implementation - app = FirecrawlApp(api_key="test_key", base_url="https://api.firecrawl.dev") - - assert app.api_key == "test_key" - assert app.base_url == "https://api.firecrawl.dev" - - def test_crawl_respects_domain_restrictions(self, mocker: MockerFixture) -> None: - """ - Test that crawl operations respect domain restrictions. - - This ensures that crawlers don't follow links to external domains - unless explicitly configured to do so. - """ - mock_client = mocker.patch("core.rag.extractor.watercrawl.provider.WaterCrawlAPIClient") - mock_instance = mock_client.return_value - mock_instance.create_crawl_request.return_value = {"uuid": "test-job"} - - provider = WaterCrawlProvider(api_key="test_key") - provider.crawl_url("https://example.com", options={"crawl_sub_pages": True}) - - call_args = mock_instance.create_crawl_request.call_args - spider_options = call_args.kwargs["spider_options"] - - # Verify allowed_domains is initialized (empty means same domain only) - assert "allowed_domains" in spider_options - assert isinstance(spider_options["allowed_domains"], list) - - -# ============================================================================ -# Test Content Extraction -# ============================================================================ - - -class TestContentExtraction: - """Test suite for content extraction from crawled pages.""" - - def test_structure_data_with_metadata(self, mocker: MockerFixture) -> None: - """Test that content is properly structured with metadata.""" - mocker.patch("core.rag.extractor.watercrawl.provider.WaterCrawlAPIClient") - - provider = WaterCrawlProvider(api_key="test_key") - - result_object = { - "url": "https://example.com/page", - "result": { - "markdown": "# Page Title\n\nPage content here.", - "metadata": { - "og:title": "Page Title", - "title": "Fallback Title", - "description": "Page description", - }, - }, - } - - structured = provider._structure_data(result_object) - - assert structured["title"] == "Page Title" - assert structured["description"] == "Page description" - assert structured["source_url"] == "https://example.com/page" - assert structured["markdown"] == "# Page Title\n\nPage content here." - - def test_structure_data_fallback_title(self, mocker: MockerFixture) -> None: - """Test that fallback title is used when og:title is not available.""" - mocker.patch("core.rag.extractor.watercrawl.provider.WaterCrawlAPIClient") - - provider = WaterCrawlProvider(api_key="test_key") - - result_object = { - "url": "https://example.com/page", - "result": {"markdown": "Content", "metadata": {"title": "Fallback Title"}}, - } - - structured = provider._structure_data(result_object) - - assert structured["title"] == "Fallback Title" - - def test_structure_data_invalid_result(self, mocker: MockerFixture) -> None: - """Test that ValueError is raised for invalid result objects.""" - mocker.patch("core.rag.extractor.watercrawl.provider.WaterCrawlAPIClient") - - provider = WaterCrawlProvider(api_key="test_key") - - # Result is a string instead of dict - result_object = {"url": "https://example.com/page", "result": "invalid string result"} - - with pytest.raises(ValueError, match="Invalid result object"): - provider._structure_data(result_object) - - def test_scrape_url_content_extraction(self, mocker: MockerFixture) -> None: - """Test content extraction from single URL scraping.""" - mock_client = mocker.patch("core.rag.extractor.watercrawl.provider.WaterCrawlAPIClient") - mock_instance = mock_client.return_value - mock_instance.scrape_url.return_value = { - "url": "https://example.com", - "result": { - "markdown": "# Main Content", - "metadata": {"og:title": "Example Page", "description": "Example description"}, - }, - } - - provider = WaterCrawlProvider(api_key="test_key") - result = provider.scrape_url("https://example.com") - - assert result["title"] == "Example Page" - assert result["description"] == "Example description" - assert result["markdown"] == "# Main Content" - assert result["source_url"] == "https://example.com" - - def test_only_main_content_extraction(self, mocker: MockerFixture) -> None: - """Test that only_main_content option filters out non-content elements.""" - mock_client = mocker.patch("core.rag.extractor.watercrawl.provider.WaterCrawlAPIClient") - mock_instance = mock_client.return_value - mock_instance.create_crawl_request.return_value = {"uuid": "test-job"} - - provider = WaterCrawlProvider(api_key="test_key") - options = {"only_main_content": True, "crawl_sub_pages": False} - provider.crawl_url("https://example.com", options=options) - - call_args = mock_instance.create_crawl_request.call_args - page_options = call_args.kwargs["page_options"] - - # Verify main content extraction is enabled - assert page_options["only_main_content"] is True - assert page_options["include_html"] is False - - -# ============================================================================ -# Test Error Handling -# ============================================================================ - - -class TestErrorHandling: - """Test suite for error handling in crawl operations.""" - - @patch("services.data_source.website_service.current_user") - @patch("services.data_source.website_service.DatasourceProviderService") - def test_invalid_provider_error(self, mock_provider_service: Mock, mock_current_user: Mock) -> None: - """Test that invalid provider raises ValueError.""" - from services.data_source.website_service import WebsiteCrawlApiRequest - - # Setup mocks - mock_current_user.current_tenant_id = "test_tenant" - mock_provider_service.return_value.get_datasource_credentials.return_value = { - "api_key": "test_key", - } - - api_request = WebsiteCrawlApiRequest( - provider="invalid_provider", url="https://example.com", options={"limit": 10} - ) - - # The error should be raised when trying to crawl with invalid provider - with pytest.raises(ValueError, match="Invalid provider"): - WebsiteService( - providers=website_service_module.DatasourceProviderService( - credentials=create_autospec(DatasourceProviderCredentialStore, instance=True) - ) - ).crawl_url(api_request) - - def test_missing_api_key_error(self, mocker: MockerFixture) -> None: - """Test that missing API key is handled properly at the httpx client level.""" - # Mock the client to avoid actual httpx initialization - mock_client = mocker.patch("core.rag.extractor.watercrawl.provider.WaterCrawlAPIClient") - - # Create provider with mocked client - should work with mock - WaterCrawlProvider(api_key="test_key") - - # Verify the client was initialized with the API key - mock_client.assert_called_once_with("test_key", None) - - def test_crawl_status_for_nonexistent_job(self, mocker: MockerFixture) -> None: - """Test handling of status check for non-existent job.""" - mock_client = mocker.patch("core.rag.extractor.watercrawl.provider.WaterCrawlAPIClient") - mock_instance = mock_client.return_value - - # Simulate API error for non-existent job - from core.rag.extractor.watercrawl.exceptions import WaterCrawlBadRequestError - - mock_response = httpx.Response(404, json={"message": "Crawl job not found", "errors": {}}) - mock_instance.get_crawl_request.side_effect = WaterCrawlBadRequestError(mock_response) - - provider = WaterCrawlProvider(api_key="test_key") - - with pytest.raises(WaterCrawlBadRequestError): - provider.get_crawl_status("nonexistent-job-id") - - -# ============================================================================ -# Integration-style Tests -# ============================================================================ - - -class TestCrawlWorkflow: - """Integration-style tests for complete crawl workflows.""" - - def test_complete_crawl_workflow(self, mocker: MockerFixture) -> None: - """Test a complete crawl workflow from start to finish.""" - mock_client = mocker.patch("core.rag.extractor.watercrawl.provider.WaterCrawlAPIClient") - mock_instance = mock_client.return_value - - # Step 1: Start crawl - mock_instance.create_crawl_request.return_value = {"uuid": "workflow-job-123"} - - provider = WaterCrawlProvider(api_key="test_key") - crawl_result = provider.crawl_url( - "https://example.com", options={"crawl_sub_pages": True, "limit": 5, "max_depth": 2} - ) - - assert crawl_result["job_id"] == "workflow-job-123" - - # Step 2: Check status (running) - mock_instance.get_crawl_request.return_value = { - "uuid": "workflow-job-123", - "status": "running", - "number_of_documents": 3, - "options": {"spider_options": {"page_limit": 5}}, - } - - status = provider.get_crawl_status("workflow-job-123") - assert status["status"] == "active" - assert status["current"] == 3 - - # Step 3: Check status (completed) - mock_instance.get_crawl_request.return_value = { - "uuid": "workflow-job-123", - "status": "completed", - "number_of_documents": 5, - "options": {"spider_options": {"page_limit": 5}}, - "duration": "00:00:10.000000", - } - mock_instance.get_crawl_request_results.return_value = { - "results": [ - { - "url": "https://example.com/page1", - "result": {"markdown": "Content 1", "metadata": {"title": "Page 1"}}, - }, - { - "url": "https://example.com/page2", - "result": {"markdown": "Content 2", "metadata": {"title": "Page 2"}}, - }, - ], - "next": None, - } - - status = provider.get_crawl_status("workflow-job-123") - assert status["status"] == "completed" - assert status["current"] == 5 - assert len(status["data"]) == 2 - - # Step 4: Get specific URL data - data = provider.get_crawl_url_data("workflow-job-123", "https://example.com/page1") - assert data is not None - assert data["title"] == "Page 1" - - def test_single_page_scrape_workflow(self, mocker: MockerFixture) -> None: - """Test workflow for scraping a single page without crawling.""" - mock_client = mocker.patch("core.rag.extractor.watercrawl.provider.WaterCrawlAPIClient") - mock_instance = mock_client.return_value - mock_instance.scrape_url.return_value = { - "url": "https://example.com/single-page", - "result": { - "markdown": "# Single Page\n\nThis is a single page scrape.", - "metadata": {"og:title": "Single Page", "description": "A single page"}, - }, - } - - provider = WaterCrawlProvider(api_key="test_key") - result = provider.scrape_url("https://example.com/single-page") - - assert result["title"] == "Single Page" - assert result["description"] == "A single page" - markdown = result["markdown"] - assert markdown is not None - assert "Single Page" in markdown - assert result["source_url"] == "https://example.com/single-page" - - -# ============================================================================ -# Test Advanced Crawl Scenarios -# ============================================================================ - - -class TestAdvancedCrawlScenarios: - """ - Test suite for advanced and edge-case crawling scenarios. - - This class tests complex crawling situations including: - - Pagination handling - - Large-scale crawls - - Concurrent crawl management - - Retry mechanisms - - Timeout handling - """ - - def test_pagination_in_crawl_results(self, mocker: MockerFixture) -> None: - """ - Test that pagination is properly handled when retrieving crawl results. - - When a crawl produces many results, they are paginated. This test - ensures that the provider correctly iterates through all pages. - """ - mock_client = mocker.patch("core.rag.extractor.watercrawl.provider.WaterCrawlAPIClient") - mock_instance = mock_client.return_value - - # Mock paginated responses - first page has 'next', second page doesn't - mock_instance.get_crawl_request_results.side_effect = [ - { - "results": [ - { - "url": f"https://example.com/page{i}", - "result": {"markdown": f"Content {i}", "metadata": {"title": f"Page {i}"}}, - } - for i in range(1, 101) - ], - "next": "page2", - }, - { - "results": [ - { - "url": f"https://example.com/page{i}", - "result": {"markdown": f"Content {i}", "metadata": {"title": f"Page {i}"}}, - } - for i in range(101, 151) - ], - "next": None, - }, - ] - - provider = WaterCrawlProvider(api_key="test_key") - - # Collect all results from paginated response - results = list(provider._get_results("test-job-id")) - - # Verify all pages were retrieved - assert len(results) == 150 - assert results[0]["title"] == "Page 1" - assert results[149]["title"] == "Page 150" - - def test_large_scale_crawl_configuration(self, mocker: MockerFixture) -> None: - """ - Test configuration for large-scale crawls with high page limits. - - Large-scale crawls require specific configuration to handle - hundreds or thousands of pages efficiently. - """ - mock_client = mocker.patch("core.rag.extractor.watercrawl.provider.WaterCrawlAPIClient") - mock_instance = mock_client.return_value - mock_instance.create_crawl_request.return_value = {"uuid": "large-crawl-job"} - - provider = WaterCrawlProvider(api_key="test_key") - - # Configure for large-scale crawl: 1000 pages, depth 5 - options = { - "crawl_sub_pages": True, - "limit": 1000, - "max_depth": 5, - "only_main_content": True, - "wait_time": 1500, - } - result = provider.crawl_url("https://example.com", options=options) - - # Verify crawl was initiated - assert result["status"] == "active" - assert result["job_id"] == "large-crawl-job" - - # Verify spider options for large crawl - call_args = mock_instance.create_crawl_request.call_args - spider_options = call_args.kwargs["spider_options"] - assert spider_options["page_limit"] == 1000 - assert spider_options["max_depth"] == 5 - - def test_crawl_with_custom_wait_time(self, mocker: MockerFixture) -> None: - """ - Test that custom wait times are properly applied to page loads. - - Wait times are crucial for dynamic content that loads via JavaScript. - This ensures pages have time to fully render before extraction. - """ - mock_client = mocker.patch("core.rag.extractor.watercrawl.provider.WaterCrawlAPIClient") - mock_instance = mock_client.return_value - mock_instance.create_crawl_request.return_value = {"uuid": "wait-test-job"} - - provider = WaterCrawlProvider(api_key="test_key") - - # Test with 3-second wait time for JavaScript-heavy pages - options = {"wait_time": 3000, "only_main_content": True} - provider.crawl_url("https://example.com/dynamic-page", options=options) - - call_args = mock_instance.create_crawl_request.call_args - page_options = call_args.kwargs["page_options"] - - # Verify wait time is set correctly - assert page_options["wait_time"] == 3000 - - def test_crawl_status_progress_tracking(self, mocker: MockerFixture) -> None: - """ - Test that crawl progress is accurately tracked and reported. - - Progress tracking allows users to monitor long-running crawls - and estimate completion time. - """ - mock_client = mocker.patch("core.rag.extractor.watercrawl.provider.WaterCrawlAPIClient") - mock_instance = mock_client.return_value - - # Simulate crawl at 60% completion - mock_instance.get_crawl_request.return_value = { - "uuid": "progress-job", - "status": "running", - "number_of_documents": 60, - "options": {"spider_options": {"page_limit": 100}}, - "duration": "00:01:30.000000", - } - - provider = WaterCrawlProvider(api_key="test_key") - status = provider.get_crawl_status("progress-job") - - # Verify progress metrics - assert status["status"] == "active" - assert status["current"] == 60 - assert status["total"] == 100 - # Calculate progress percentage - progress_percentage = (status["current"] / status["total"]) * 100 - assert progress_percentage == 60.0 - - def test_crawl_with_sitemap_usage(self, mocker: MockerFixture) -> None: - """ - Test that sitemap.xml is utilized when use_sitemap is enabled. - - Sitemaps provide a structured list of URLs, making crawls more - efficient and comprehensive. - """ - mock_client = mocker.patch("core.rag.extractor.watercrawl.provider.WaterCrawlAPIClient") - mock_instance = mock_client.return_value - mock_instance.create_crawl_request.return_value = {"uuid": "sitemap-job"} - - provider = WaterCrawlProvider(api_key="test_key") - - # Enable sitemap usage - options = {"crawl_sub_pages": True, "use_sitemap": True, "limit": 50} - provider.crawl_url("https://example.com", options=options) - - # Note: use_sitemap is passed to the service layer but not directly - # to WaterCrawl spider_options. This test verifies the option is accepted. - call_args = mock_instance.create_crawl_request.call_args - assert call_args is not None - - def test_empty_crawl_results(self, mocker: MockerFixture) -> None: - """ - Test handling of crawls that return no results. - - This can occur when all pages are excluded or no content matches - the extraction criteria. - """ - mock_client = mocker.patch("core.rag.extractor.watercrawl.provider.WaterCrawlAPIClient") - mock_instance = mock_client.return_value - mock_instance.get_crawl_request.return_value = { - "uuid": "empty-job", - "status": "completed", - "number_of_documents": 0, - "options": {"spider_options": {"page_limit": 10}}, - "duration": "00:00:05.000000", - } - mock_instance.get_crawl_request_results.return_value = {"results": list[object](), "next": None} - - provider = WaterCrawlProvider(api_key="test_key") - status = provider.get_crawl_status("empty-job") - - # Verify empty results are handled correctly - assert status["status"] == "completed" - assert status["current"] == 0 - assert status["total"] == 10 - assert len(status["data"]) == 0 - - def test_crawl_with_multiple_include_patterns(self, mocker: MockerFixture) -> None: - """ - Test crawling with multiple include patterns for fine-grained control. - - Multiple patterns allow targeting specific sections of a website - while excluding others. - """ - mock_client = mocker.patch("core.rag.extractor.watercrawl.provider.WaterCrawlAPIClient") - mock_instance = mock_client.return_value - mock_instance.create_crawl_request.return_value = {"uuid": "multi-pattern-job"} - - provider = WaterCrawlProvider(api_key="test_key") - - # Multiple include patterns for different content types - options = { - "crawl_sub_pages": True, - "includes": "/blog/*,/news/*,/articles/*,/docs/*", - "limit": 100, - } - provider.crawl_url("https://example.com", options=options) - - call_args = mock_instance.create_crawl_request.call_args - spider_options = call_args.kwargs["spider_options"] - - # Verify all include patterns are set - assert len(spider_options["include_paths"]) == 4 - assert "/blog/*" in spider_options["include_paths"] - assert "/news/*" in spider_options["include_paths"] - assert "/articles/*" in spider_options["include_paths"] - assert "/docs/*" in spider_options["include_paths"] - - def test_crawl_duration_calculation(self, mocker: MockerFixture) -> None: - """ - Test accurate calculation of crawl duration from time strings. - - Duration tracking helps analyze crawl performance and optimize - configuration for future crawls. - """ - mock_client = mocker.patch("core.rag.extractor.watercrawl.provider.WaterCrawlAPIClient") - mock_instance = mock_client.return_value - - # Test various duration formats - test_cases = [ - ("00:00:10.500000", 10.5), # 10.5 seconds - ("00:01:30.250000", 90.25), # 1 minute 30.25 seconds - ("01:15:45.750000", 4545.75), # 1 hour 15 minutes 45.75 seconds - ] - - for duration_str, expected_seconds in test_cases: - mock_instance.get_crawl_request.return_value = { - "uuid": "duration-test", - "status": "completed", - "number_of_documents": 10, - "options": {"spider_options": {"page_limit": 10}}, - "duration": duration_str, - } - mock_instance.get_crawl_request_results.return_value = {"results": list[object](), "next": None} - - provider = WaterCrawlProvider(api_key="test_key") - status = provider.get_crawl_status("duration-test") - - # Verify duration is calculated correctly - assert abs(status["time_consuming"] - expected_seconds) < 0.01 - - -# ============================================================================ -# Test Provider-Specific Features -# ============================================================================ - - -class TestProviderSpecificFeatures: - """ - Test suite for provider-specific features and behaviors. - - Different crawl providers (Firecrawl, WaterCrawl, JinaReader) have - unique features and API behaviors that require specific testing. - """ - - @patch("services.data_source.website_service.current_user") - @patch("services.data_source.website_service.DatasourceProviderService") - def test_firecrawl_with_prompt_parameter( - self, mock_provider_service: Mock, mock_current_user: Mock, mocker: MockerFixture - ) -> None: - """ - Test Firecrawl's prompt parameter for AI-guided extraction. - - Firecrawl v2 supports prompts to guide content extraction using AI, - allowing for semantic filtering of crawled content. - """ - # Setup mocks - mock_current_user.current_tenant_id = "test_tenant" - mock_provider_service.return_value.get_datasource_credentials.return_value = { - "firecrawl_api_key": "test_key", - "base_url": "https://api.firecrawl.dev", - } - - mock_firecrawl = mocker.patch("services.data_source.website_service.FirecrawlApp") - mock_firecrawl_instance = mock_firecrawl.return_value - mock_firecrawl_instance.crawl_url.return_value = "prompt-job-123" - - # Mock redis - mocker.patch("services.data_source.website_service.redis_client") - - from services.data_source.website_service import WebsiteCrawlApiRequest - - # Include a prompt for AI-guided extraction - api_request = WebsiteCrawlApiRequest( - provider="firecrawl", - url="https://example.com", - options={ - "limit": 20, - "crawl_sub_pages": True, - "only_main_content": True, - "prompt": "Extract only technical documentation and API references", - }, - ) - - result = WebsiteService( - providers=website_service_module.DatasourceProviderService( - credentials=create_autospec(DatasourceProviderCredentialStore, instance=True) - ) - ).crawl_url(api_request) - - assert result["status"] == "active" - assert result["job_id"] == "prompt-job-123" - - # Verify prompt was passed to Firecrawl - call_args = mock_firecrawl_instance.crawl_url.call_args - params = call_args[0][1] # Second argument is params - assert "prompt" in params - assert params["prompt"] == "Extract only technical documentation and API references" - - @patch("services.data_source.website_service.current_user") - @patch("services.data_source.website_service.DatasourceProviderService") - def test_jinareader_single_page_mode( - self, mock_provider_service: Mock, mock_current_user: Mock, mocker: MockerFixture - ) -> None: - """ - Test JinaReader's single-page scraping mode. - - JinaReader can scrape individual pages without crawling, - useful for quick content extraction. - """ - # Setup mocks - mock_current_user.current_tenant_id = "test_tenant" - mock_provider_service.return_value.get_datasource_credentials.return_value = { - "api_key": "test_key", - } - - response_data = { - "code": 200, - "data": { - "title": "Single Page Title", - "content": "Page content here", - "url": "https://example.com/page", - }, - } - mock_response = httpx.Response(200, json=response_data) - mocker.patch("services.data_source.website_service._jina_http_client.get", return_value=mock_response) - - from services.data_source.website_service import WebsiteCrawlApiRequest - - # Single page mode (crawl_sub_pages = False) - api_request = WebsiteCrawlApiRequest( - provider="jinareader", url="https://example.com/page", options={"crawl_sub_pages": False, "limit": 1} - ) - - result = WebsiteService( - providers=website_service_module.DatasourceProviderService( - credentials=create_autospec(DatasourceProviderCredentialStore, instance=True) - ) - ).crawl_url(api_request) - - # In single-page mode, JinaReader returns data immediately - assert result["status"] == "active" - assert "data" in result - - @patch("services.data_source.website_service.current_user") - @patch("services.data_source.website_service.DatasourceProviderService") - def test_watercrawl_with_tag_filtering( - self, mock_provider_service: Mock, mock_current_user: Mock, mocker: MockerFixture - ) -> None: - """ - Test WaterCrawl's HTML tag filtering capabilities. - - WaterCrawl allows including or excluding specific HTML tags - during content extraction for precise control. - """ - # Setup mocks - mock_current_user.current_tenant_id = "test_tenant" - mock_provider_service.return_value.get_datasource_credentials.return_value = { - "api_key": "test_key", - "base_url": "https://app.watercrawl.dev", - } - - mock_watercrawl = mocker.patch("services.data_source.website_service.WaterCrawlProvider") - mock_watercrawl_instance = mock_watercrawl.return_value - mock_watercrawl_instance.crawl_url.return_value = {"status": "active", "job_id": "tag-filter-job"} - - from services.data_source.website_service import WebsiteCrawlApiRequest - - # Configure with tag filtering - api_request = WebsiteCrawlApiRequest( - provider="watercrawl", - url="https://example.com", - options={ - "limit": 10, - "crawl_sub_pages": True, - "exclude_tags": "nav,footer,aside", - "include_tags": "article,main", - }, - ) - - result = WebsiteService( - providers=website_service_module.DatasourceProviderService( - credentials=create_autospec(DatasourceProviderCredentialStore, instance=True) - ) - ).crawl_url(api_request) - - assert result["status"] == "active" - assert result["job_id"] == "tag-filter-job" - - def test_firecrawl_base_url_configuration(self) -> None: - """ - Test that Firecrawl can be configured with custom base URLs. - - This is important for self-hosted Firecrawl instances or - different API endpoints. - """ - from core.rag.extractor.firecrawl.firecrawl_app import FirecrawlApp - - # Test with custom base URL - custom_base_url = "https://custom-firecrawl.example.com" - app = FirecrawlApp(api_key="test_key", base_url=custom_base_url) - - assert app.base_url == custom_base_url - assert app.api_key == "test_key" - - def test_watercrawl_base_url_default(self, mocker: MockerFixture) -> None: - """ - Test WaterCrawl's default base URL configuration. - - Verifies that the provider uses the correct default URL when - none is specified. - """ - mock_client = mocker.patch("core.rag.extractor.watercrawl.provider.WaterCrawlAPIClient") - - # Create provider without specifying base_url - WaterCrawlProvider(api_key="test_key") - - # Verify default base URL is used - mock_client.assert_called_once_with("test_key", None) - - -# ============================================================================ -# Test Data Structure and Validation -# ============================================================================ - - -class TestDataStructureValidation: - """ - Test suite for data structure validation and transformation. - - Ensures that crawled data is properly structured, validated, - and transformed into the expected format. - """ - - def test_crawl_request_to_api_request_conversion(self) -> None: - """ - Test conversion from API request to internal CrawlRequest format. - - This conversion ensures that external API parameters are properly - mapped to internal data structures. - """ - from services.data_source.website_service import WebsiteCrawlApiRequest - - # Create API request with all options - api_request = WebsiteCrawlApiRequest( - provider="watercrawl", - url="https://example.com", - options={ - "limit": 50, - "crawl_sub_pages": True, - "only_main_content": True, - "includes": "/blog/*", - "excludes": "/admin/*", - "prompt": "Extract main content", - "max_depth": 3, - "use_sitemap": True, - }, - ) - - # Convert to internal format - crawl_request = api_request.to_crawl_request() - - # Verify all fields are properly converted - assert crawl_request.url == "https://example.com" - assert crawl_request.provider == "watercrawl" - assert crawl_request.options.limit == 50 - assert crawl_request.options.crawl_sub_pages is True - assert crawl_request.options.only_main_content is True - assert crawl_request.options.includes == "/blog/*" - assert crawl_request.options.excludes == "/admin/*" - assert crawl_request.options.prompt == "Extract main content" - assert crawl_request.options.max_depth == 3 - assert crawl_request.options.use_sitemap is True - - def test_crawl_options_path_parsing(self) -> None: - """ - Test that include/exclude paths are correctly parsed from strings. - - Paths can be provided as comma-separated strings and must be - split into individual patterns. - """ - # Test with multiple paths - options = CrawlOptions(includes="/blog/*,/news/*,/docs/*", excludes="/admin/*,/private/*,/test/*") - - include_paths = options.get_include_paths() - exclude_paths = options.get_exclude_paths() - - # Verify parsing - assert len(include_paths) == 3 - assert "/blog/*" in include_paths - assert "/news/*" in include_paths - assert "/docs/*" in include_paths - - assert len(exclude_paths) == 3 - assert "/admin/*" in exclude_paths - assert "/private/*" in exclude_paths - assert "/test/*" in exclude_paths - - def test_crawl_options_with_whitespace(self) -> None: - """ - Test that whitespace in path strings is handled correctly. - - Users might include spaces around commas, which should be - handled gracefully. - """ - # Test with spaces around commas - options = CrawlOptions(includes=" /blog/* , /news/* , /docs/* ", excludes=" /admin/* , /private/* ") - - include_paths = options.get_include_paths() - exclude_paths = options.get_exclude_paths() - - # Verify paths are trimmed (note: current implementation doesn't trim, - # so paths will include spaces - this documents current behavior) - assert len(include_paths) == 3 - assert len(exclude_paths) == 2 - - def test_website_crawl_message_structure(self) -> None: - """ - Test the structure of WebsiteCrawlMessage entity. - - This entity wraps crawl results and must have the correct structure - for downstream processing. - """ - from core.datasource.entities.datasource_entities import WebsiteCrawlMessage, WebSiteInfo - - # Create a crawl message with results - web_info = WebSiteInfo(status="completed", web_info_list=[], total=10, completed=10) - - message = WebsiteCrawlMessage(result=web_info) - - # Verify structure - assert message.result.status == "completed" - assert message.result.total == 10 - assert message.result.completed == 10 - assert isinstance(message.result.web_info_list, list) - - def test_datasource_identity_structure(self) -> None: - """ - Test that DatasourceIdentity contains all required fields. - - Identity information is crucial for tracking and managing - datasource instances. - """ - identity = DatasourceIdentity( - author="test_author", - name="test_datasource", - label={"en_US": "Test Datasource", "zh_Hans": "测试数据源"}, - provider="test_provider", - icon="test_icon.svg", - ) - - # Verify all fields are present - assert identity.author == "test_author" - assert identity.name == "test_datasource" - assert identity.provider == "test_provider" - assert identity.icon == "test_icon.svg" - # I18nObject has attributes, not dict keys - assert identity.label.en_US == "Test Datasource" - assert identity.label.zh_Hans == "测试数据源" - - -# ============================================================================ -# Test Edge Cases and Boundary Conditions -# ============================================================================ - - -class TestEdgeCasesAndBoundaries: - """ - Test suite for edge cases and boundary conditions. - - These tests ensure robust handling of unusual inputs, limits, - and exceptional scenarios. - """ - - def test_crawl_with_zero_limit(self, mocker: MockerFixture) -> None: - """ - Test behavior when limit is set to zero. - - A zero limit should be handled gracefully, potentially defaulting - to a minimum value or raising an error. - """ - mock_client = mocker.patch("core.rag.extractor.watercrawl.provider.WaterCrawlAPIClient") - mock_instance = mock_client.return_value - mock_instance.create_crawl_request.return_value = {"uuid": "zero-limit-job"} - - provider = WaterCrawlProvider(api_key="test_key") - - # Attempt crawl with zero limit - options = {"crawl_sub_pages": True, "limit": 0} - result = provider.crawl_url("https://example.com", options=options) - - # Verify crawl was created (implementation may handle this differently) - assert result["status"] == "active" - - def test_crawl_with_very_large_limit(self, mocker: MockerFixture) -> None: - """ - Test crawl configuration with extremely large page limits. - - Very large limits should be accepted but may be subject to - provider-specific constraints. - """ - mock_client = mocker.patch("core.rag.extractor.watercrawl.provider.WaterCrawlAPIClient") - mock_instance = mock_client.return_value - mock_instance.create_crawl_request.return_value = {"uuid": "large-limit-job"} - - provider = WaterCrawlProvider(api_key="test_key") - - # Test with very large limit (10,000 pages) - options = {"crawl_sub_pages": True, "limit": 10000, "max_depth": 10} - result = provider.crawl_url("https://example.com", options=options) - - assert result["status"] == "active" - - call_args = mock_instance.create_crawl_request.call_args - spider_options = call_args.kwargs["spider_options"] - assert spider_options["page_limit"] == 10000 - - def test_crawl_with_empty_url(self) -> None: - """ - Test that empty URLs are rejected with appropriate error. - - Empty or invalid URLs should fail validation before attempting - to crawl. - """ - from services.data_source.website_service import WebsiteCrawlApiRequest - - # Empty URL should raise ValueError during validation - with pytest.raises(ValueError, match="URL is required"): - WebsiteCrawlApiRequest.from_args({"provider": "watercrawl", "url": "", "options": {"limit": 10}}) - - def test_crawl_with_special_characters_in_paths(self, mocker: MockerFixture) -> None: - """ - Test handling of special characters in include/exclude paths. - - Paths may contain special regex characters that need proper escaping - or handling. - """ - mock_client = mocker.patch("core.rag.extractor.watercrawl.provider.WaterCrawlAPIClient") - mock_instance = mock_client.return_value - mock_instance.create_crawl_request.return_value = {"uuid": "special-chars-job"} - - provider = WaterCrawlProvider(api_key="test_key") - - # Include paths with special characters - options = { - "crawl_sub_pages": True, - "includes": "/blog/[0-9]+/*,/category/(tech|science)/*", - "limit": 20, - } - provider.crawl_url("https://example.com", options=options) - - call_args = mock_instance.create_crawl_request.call_args - spider_options = call_args.kwargs["spider_options"] - - # Verify special characters are preserved - assert "/blog/[0-9]+/*" in spider_options["include_paths"] - assert "/category/(tech|science)/*" in spider_options["include_paths"] - - def test_crawl_status_with_null_duration(self, mocker: MockerFixture) -> None: - """ - Test handling of null/missing duration in crawl status. - - Duration may be null for active crawls or if timing data is unavailable. - """ - mock_client = mocker.patch("core.rag.extractor.watercrawl.provider.WaterCrawlAPIClient") - mock_instance = mock_client.return_value - mock_instance.get_crawl_request.return_value = { - "uuid": "null-duration-job", - "status": "running", - "number_of_documents": 5, - "options": {"spider_options": {"page_limit": 10}}, - "duration": None, # Null duration - } - - provider = WaterCrawlProvider(api_key="test_key") - status = provider.get_crawl_status("null-duration-job") - - # Verify null duration is handled (should default to 0) - assert status["time_consuming"] == 0 - - def test_structure_data_with_missing_metadata_fields(self, mocker: MockerFixture) -> None: - """ - Test content extraction when metadata fields are missing. - - Not all pages have complete metadata, so extraction should - handle missing fields gracefully. - """ - mocker.patch("core.rag.extractor.watercrawl.provider.WaterCrawlAPIClient") - - provider = WaterCrawlProvider(api_key="test_key") - - # Result with minimal metadata - result_object = { - "url": "https://example.com/minimal", - "result": { - "markdown": "# Minimal Content", - "metadata": dict[str, object](), # Empty metadata - }, - } - - structured = provider._structure_data(result_object) - - # Verify graceful handling of missing metadata - assert structured["title"] is None - assert structured["description"] is None - assert structured["source_url"] == "https://example.com/minimal" - assert structured["markdown"] == "# Minimal Content" - - def test_get_results_with_empty_pages(self, mocker: MockerFixture) -> None: - """ - Test pagination handling when some pages return empty results. - - Empty pages in pagination cause the loop to break early in the - current implementation, as per the code logic in _get_results. - """ - mock_client = mocker.patch("core.rag.extractor.watercrawl.provider.WaterCrawlAPIClient") - mock_instance = mock_client.return_value - - # First page has results, second page is empty (breaks loop) - mock_instance.get_crawl_request_results.side_effect = [ - { - "results": [ - { - "url": "https://example.com/page1", - "result": {"markdown": "Content 1", "metadata": {"title": "Page 1"}}, - } - ], - "next": "page2", - }, - {"results": list[object](), "next": None}, # Empty page breaks the loop - ] - - provider = WaterCrawlProvider(api_key="test_key") - results = list(provider._get_results("test-job")) - - # Current implementation breaks on empty results - # This documents the actual behavior - assert len(results) == 1 - assert results[0]["title"] == "Page 1" diff --git a/api/tests/unit_tests/core/datasource/website_crawl/test_website_crawl_plugin.py b/api/tests/unit_tests/core/datasource/website_crawl/test_website_crawl_plugin.py index c388b1ed46a..acf2f4a8446 100644 --- a/api/tests/unit_tests/core/datasource/website_crawl/test_website_crawl_plugin.py +++ b/api/tests/unit_tests/core/datasource/website_crawl/test_website_crawl_plugin.py @@ -1,101 +1,1855 @@ from collections.abc import Generator -from unittest.mock import patch +from unittest.mock import Mock, create_autospec, patch +import httpx import pytest +from pytest_mock import MockerFixture +import services.data_source.website_service as website_service_module from core.datasource.__base.datasource_runtime import DatasourceRuntime from core.datasource.entities.datasource_entities import ( DatasourceEntity, + DatasourceIdentity, + DatasourceProviderEntityWithPlugin, + DatasourceProviderIdentity, DatasourceProviderType, WebsiteCrawlMessage, ) from core.datasource.website_crawl.website_crawl_plugin import WebsiteCrawlDatasourcePlugin -from tests.unit_tests.core.datasource.factories import datasource_entity +from core.datasource.website_crawl.website_crawl_provider import WebsiteCrawlDatasourcePluginProviderController +from core.rag.extractor.watercrawl.provider import WaterCrawlProvider +from services.data_source.credential_gateway import DatasourceProviderCredentialStore +from services.data_source.website_service import CrawlOptions, CrawlRequest, WebsiteService +from tests.helpers.unit.core.datasource.factories import datasource_entity + +""" +Unit tests for website crawling functionality. + +This module tests the core website crawling features including: +- URL crawling logic with different providers +- Robots.txt respect and compliance +- Max depth limiting for crawl operations +- Content extraction from web pages +- Link following logic and navigation + +The tests cover multiple crawl providers (Firecrawl, WaterCrawl, JinaReader) +and ensure proper handling of crawl options, status checking, and data retrieval. +""" -class TestWebsiteCrawlDatasourcePlugin: +class TestWebsiteCrawlCases: + # ============================================================================ + # Fixtures + # ============================================================================ + + @staticmethod @pytest.fixture - def entity(self) -> DatasourceEntity: - entity = datasource_entity("test_name") - entity.identity.provider = "test-provider" - entity.identity.name = "test-name" - return entity + def mock_datasource_entity() -> DatasourceEntity: + """Create a mock datasource entity for testing.""" + return DatasourceEntity( + identity=DatasourceIdentity( + author="test_author", + name="test_datasource", + label={"en_US": "Test Datasource", "zh_Hans": "测试数据源"}, + provider="test_provider", + icon="test_icon.svg", + ), + parameters=[], + description={"en_US": "Test datasource description", "zh_Hans": "测试数据源描述"}, + ) + @staticmethod @pytest.fixture - def runtime(self) -> DatasourceRuntime: - runtime = DatasourceRuntime(tenant_id="test_tenant") - runtime.credentials = {"api_key": "test-key"} - return runtime - - def test_init(self, entity: DatasourceEntity, runtime: DatasourceRuntime) -> None: - # Arrange - tenant_id = "test-tenant-id" - icon = "test-icon" - plugin_unique_identifier = "test-plugin-id" - - # Act - plugin = WebsiteCrawlDatasourcePlugin( - entity=entity, - runtime=runtime, - tenant_id=tenant_id, - icon=icon, - plugin_unique_identifier=plugin_unique_identifier, + def mock_provider_entity(mock_datasource_entity: DatasourceEntity) -> DatasourceProviderEntityWithPlugin: + """Create a mock provider entity with plugin for testing.""" + return DatasourceProviderEntityWithPlugin( + identity=DatasourceProviderIdentity( + author="test_author", + name="test_provider", + description={"en_US": "Test Provider", "zh_Hans": "测试提供者"}, + icon="test_icon.svg", + label={"en_US": "Test Provider", "zh_Hans": "测试提供者"}, + ), + credentials_schema=[], + provider_type=DatasourceProviderType.WEBSITE_CRAWL, + datasources=[mock_datasource_entity], ) - # Assert - assert plugin.tenant_id == tenant_id - assert plugin.plugin_unique_identifier == plugin_unique_identifier - assert plugin.entity == entity - assert plugin.runtime == runtime - assert plugin.icon == icon - - def test_datasource_provider_type(self, entity: DatasourceEntity, runtime: DatasourceRuntime) -> None: - # Arrange - plugin = WebsiteCrawlDatasourcePlugin( - entity=entity, runtime=runtime, tenant_id="test", icon="test", plugin_unique_identifier="test" + @staticmethod + @pytest.fixture + def crawl_options() -> CrawlOptions: + """Create default crawl options for testing.""" + return CrawlOptions( + limit=10, + crawl_sub_pages=True, + only_main_content=True, + includes="/blog/*,/docs/*", + excludes="/admin/*,/private/*", + max_depth=3, + use_sitemap=True, ) - # Act & Assert - assert plugin.datasource_provider_type() == DatasourceProviderType.WEBSITE_CRAWL + @staticmethod + @pytest.fixture + def crawl_request(crawl_options: CrawlOptions) -> CrawlRequest: + """Create a crawl request for testing.""" + return CrawlRequest(url="https://example.com", provider="watercrawl", options=crawl_options) - def test_get_website_crawl(self, entity: DatasourceEntity, runtime: DatasourceRuntime) -> None: - # Arrange - plugin = WebsiteCrawlDatasourcePlugin( - entity=entity, - runtime=runtime, - tenant_id="test-tenant-id", - icon="test-icon", - plugin_unique_identifier="test-plugin-id", - ) + # ============================================================================ + # Test CrawlOptions + # ============================================================================ - user_id = "test-user-id" - datasource_parameters = {"url": "https://example.com"} - provider_type = "firecrawl" + class TestCrawlOptions: + """Test suite for CrawlOptions data class.""" - message = WebsiteCrawlMessage() + def test_crawl_options_defaults(self) -> None: + """Test that CrawlOptions has correct default values.""" + options = CrawlOptions() - # Mock PluginDatasourceManager - with patch("core.datasource.website_crawl.website_crawl_plugin.PluginDatasourceManager") as mock_manager_class: - mock_manager = mock_manager_class.return_value - mock_manager.get_website_crawl.return_value = (msg for msg in [message]) + assert options.limit == 1 + assert options.crawl_sub_pages is False + assert options.only_main_content is False + assert options.includes is None + assert options.excludes is None + assert options.prompt is None + assert options.max_depth is None + assert options.use_sitemap is True + + def test_get_include_paths_with_values(self, crawl_options: CrawlOptions) -> None: + """Test parsing include paths from comma-separated string.""" + paths = crawl_options.get_include_paths() + + assert len(paths) == 2 + assert "/blog/*" in paths + assert "/docs/*" in paths + + def test_get_include_paths_empty(self) -> None: + """Test that empty includes returns empty list.""" + options = CrawlOptions(includes=None) + paths = options.get_include_paths() + + assert paths == [] + + def test_get_exclude_paths_with_values(self, crawl_options: CrawlOptions) -> None: + """Test parsing exclude paths from comma-separated string.""" + paths = crawl_options.get_exclude_paths() + + assert len(paths) == 2 + assert "/admin/*" in paths + assert "/private/*" in paths + + def test_get_exclude_paths_empty(self) -> None: + """Test that empty excludes returns empty list.""" + options = CrawlOptions(excludes=None) + paths = options.get_exclude_paths() + + assert paths == [] + + def test_max_depth_limiting(self) -> None: + """Test that max_depth can be set to limit crawl depth.""" + options = CrawlOptions(max_depth=5, crawl_sub_pages=True) + + assert options.max_depth == 5 + assert options.crawl_sub_pages is True + + # ============================================================================ + # Test WebsiteCrawlDatasourcePlugin + # ============================================================================ + + class TestWebsiteCrawlDatasourcePlugin: + """Test suite for WebsiteCrawlDatasourcePlugin.""" + + def test_plugin_initialization(self, mock_datasource_entity: DatasourceEntity) -> None: + """Test that plugin initializes correctly with required parameters.""" + from core.datasource.__base.datasource_runtime import DatasourceRuntime + + runtime = DatasourceRuntime(tenant_id="test_tenant", credentials={}) + plugin = WebsiteCrawlDatasourcePlugin( + entity=mock_datasource_entity, + runtime=runtime, + tenant_id="test_tenant", + icon="test_icon.svg", + plugin_unique_identifier="test_plugin_id", + ) + + assert plugin.tenant_id == "test_tenant" + assert plugin.plugin_unique_identifier == "test_plugin_id" + assert plugin.entity == mock_datasource_entity + assert plugin.datasource_provider_type() == DatasourceProviderType.WEBSITE_CRAWL + + def test_get_website_crawl(self, mock_datasource_entity: DatasourceEntity, mocker: MockerFixture) -> None: + """Test that get_website_crawl calls PluginDatasourceManager correctly.""" + from core.datasource.__base.datasource_runtime import DatasourceRuntime + + runtime = DatasourceRuntime(tenant_id="test_tenant", credentials={"api_key": "test_key"}) + plugin = WebsiteCrawlDatasourcePlugin( + entity=mock_datasource_entity, + runtime=runtime, + tenant_id="test_tenant", + icon="test_icon.svg", + plugin_unique_identifier="test_plugin_id", + ) + + # Mock the PluginDatasourceManager + mock_manager = mocker.patch("core.datasource.website_crawl.website_crawl_plugin.PluginDatasourceManager") + mock_instance = mock_manager.return_value + mock_instance.get_website_crawl.return_value = iter([]) + + datasource_params = {"url": "https://example.com", "max_depth": 2} + + plugin.get_website_crawl( + user_id="test_user", datasource_parameters=datasource_params, provider_type="watercrawl" + ) + + # Verify the manager was called with correct parameters + mock_instance.get_website_crawl.assert_called_once_with( + tenant_id="test_tenant", + user_id="test_user", + datasource_provider=mock_datasource_entity.identity.provider, + datasource_name=mock_datasource_entity.identity.name, + credentials={"api_key": "test_key"}, + datasource_parameters=datasource_params, + provider_type="watercrawl", + ) + + # ============================================================================ + # Test WebsiteCrawlDatasourcePluginProviderController + # ============================================================================ + + class TestWebsiteCrawlDatasourcePluginProviderController: + """Test suite for WebsiteCrawlDatasourcePluginProviderController.""" + + def test_provider_controller_initialization( + self, mock_provider_entity: DatasourceProviderEntityWithPlugin + ) -> None: + """Test provider controller initialization.""" + controller = WebsiteCrawlDatasourcePluginProviderController( + entity=mock_provider_entity, + plugin_id="test_plugin_id", + plugin_unique_identifier="test_unique_id", + tenant_id="test_tenant", + ) + + assert controller.plugin_id == "test_plugin_id" + assert controller.plugin_unique_identifier == "test_unique_id" + assert controller.provider_type == DatasourceProviderType.WEBSITE_CRAWL + + def test_get_datasource_success(self, mock_provider_entity: DatasourceProviderEntityWithPlugin) -> None: + """Test retrieving a datasource by name.""" + controller = WebsiteCrawlDatasourcePluginProviderController( + entity=mock_provider_entity, + plugin_id="test_plugin_id", + plugin_unique_identifier="test_unique_id", + tenant_id="test_tenant", + ) + + datasource = controller.get_datasource("test_datasource") + + assert isinstance(datasource, WebsiteCrawlDatasourcePlugin) + assert datasource.tenant_id == "test_tenant" + assert datasource.plugin_unique_identifier == "test_unique_id" + + def test_get_datasource_not_found(self, mock_provider_entity: DatasourceProviderEntityWithPlugin) -> None: + """Test that ValueError is raised when datasource is not found.""" + controller = WebsiteCrawlDatasourcePluginProviderController( + entity=mock_provider_entity, + plugin_id="test_plugin_id", + plugin_unique_identifier="test_unique_id", + tenant_id="test_tenant", + ) + + with pytest.raises(ValueError, match="Datasource with name nonexistent not found"): + controller.get_datasource("nonexistent") + + # ============================================================================ + # Test WaterCrawl Provider - URL Crawling Logic + # ============================================================================ + + class TestWaterCrawlProvider: + """Test suite for WaterCrawl provider crawling functionality.""" + + def test_crawl_url_basic(self, mocker: MockerFixture) -> None: + """Test basic URL crawling without sub-pages.""" + mock_client = mocker.patch("core.rag.extractor.watercrawl.provider.WaterCrawlAPIClient") + mock_instance = mock_client.return_value + mock_instance.create_crawl_request.return_value = {"uuid": "test-job-123"} + + provider = WaterCrawlProvider(api_key="test_key") + result = provider.crawl_url("https://example.com", options={"crawl_sub_pages": False}) + + assert result["status"] == "active" + assert result["job_id"] == "test-job-123" + + # Verify spider options for single page crawl + call_args = mock_instance.create_crawl_request.call_args + spider_options = call_args.kwargs["spider_options"] + assert spider_options["max_depth"] == 1 + assert spider_options["page_limit"] == 1 + + def test_crawl_url_with_sub_pages(self, mocker: MockerFixture) -> None: + """Test URL crawling with sub-pages enabled.""" + mock_client = mocker.patch("core.rag.extractor.watercrawl.provider.WaterCrawlAPIClient") + mock_instance = mock_client.return_value + mock_instance.create_crawl_request.return_value = {"uuid": "test-job-456"} + + provider = WaterCrawlProvider(api_key="test_key") + options = {"crawl_sub_pages": True, "limit": 50, "max_depth": 3} + result = provider.crawl_url("https://example.com", options=options) + + assert result["status"] == "active" + assert result["job_id"] == "test-job-456" + + # Verify spider options for multi-page crawl + call_args = mock_instance.create_crawl_request.call_args + spider_options = call_args.kwargs["spider_options"] + assert spider_options["max_depth"] == 3 + assert spider_options["page_limit"] == 50 + + def test_crawl_url_max_depth_limiting(self, mocker: MockerFixture) -> None: + """Test that max_depth properly limits crawl depth.""" + mock_client = mocker.patch("core.rag.extractor.watercrawl.provider.WaterCrawlAPIClient") + mock_instance = mock_client.return_value + mock_instance.create_crawl_request.return_value = {"uuid": "test-job-789"} + + provider = WaterCrawlProvider(api_key="test_key") + + # Test with max_depth of 2 + options = {"crawl_sub_pages": True, "max_depth": 2, "limit": 100} + provider.crawl_url("https://example.com", options=options) + + call_args = mock_instance.create_crawl_request.call_args + spider_options = call_args.kwargs["spider_options"] + assert spider_options["max_depth"] == 2 + + def test_crawl_url_with_include_exclude_paths(self, mocker: MockerFixture) -> None: + """Test URL crawling with include and exclude path filters.""" + mock_client = mocker.patch("core.rag.extractor.watercrawl.provider.WaterCrawlAPIClient") + mock_instance = mock_client.return_value + mock_instance.create_crawl_request.return_value = {"uuid": "test-job-101"} + + provider = WaterCrawlProvider(api_key="test_key") + options = { + "crawl_sub_pages": True, + "includes": "/blog/*,/docs/*", + "excludes": "/admin/*,/private/*", + "limit": 20, + } + provider.crawl_url("https://example.com", options=options) + + call_args = mock_instance.create_crawl_request.call_args + spider_options = call_args.kwargs["spider_options"] + + # Verify include paths + assert len(spider_options["include_paths"]) == 2 + assert "/blog/*" in spider_options["include_paths"] + assert "/docs/*" in spider_options["include_paths"] + + # Verify exclude paths + assert len(spider_options["exclude_paths"]) == 2 + assert "/admin/*" in spider_options["exclude_paths"] + assert "/private/*" in spider_options["exclude_paths"] + + def test_crawl_url_content_extraction_options(self, mocker: MockerFixture) -> None: + """Test that content extraction options are properly configured.""" + mock_client = mocker.patch("core.rag.extractor.watercrawl.provider.WaterCrawlAPIClient") + mock_instance = mock_client.return_value + mock_instance.create_crawl_request.return_value = {"uuid": "test-job-202"} + + provider = WaterCrawlProvider(api_key="test_key") + options = {"only_main_content": True, "wait_time": 2000} + provider.crawl_url("https://example.com", options=options) + + call_args = mock_instance.create_crawl_request.call_args + page_options = call_args.kwargs["page_options"] + + # Verify content extraction settings + assert page_options["only_main_content"] is True + assert page_options["wait_time"] == 2000 + assert page_options["include_html"] is False + + def test_crawl_url_minimum_wait_time(self, mocker: MockerFixture) -> None: + """Test that wait_time has a minimum value of 1000ms.""" + mock_client = mocker.patch("core.rag.extractor.watercrawl.provider.WaterCrawlAPIClient") + mock_instance = mock_client.return_value + mock_instance.create_crawl_request.return_value = {"uuid": "test-job-303"} + + provider = WaterCrawlProvider(api_key="test_key") + options = {"wait_time": 500} # Below minimum + provider.crawl_url("https://example.com", options=options) + + call_args = mock_instance.create_crawl_request.call_args + page_options = call_args.kwargs["page_options"] + + # Should be clamped to minimum of 1000 + assert page_options["wait_time"] == 1000 + + # ============================================================================ + # Test Crawl Status and Results + # ============================================================================ + + class TestCrawlStatus: + """Test suite for crawl status checking and result retrieval.""" + + def test_get_crawl_status_active(self, mocker: MockerFixture) -> None: + """Test getting status of an active crawl job.""" + mock_client = mocker.patch("core.rag.extractor.watercrawl.provider.WaterCrawlAPIClient") + mock_instance = mock_client.return_value + mock_instance.get_crawl_request.return_value = { + "uuid": "test-job-123", + "status": "running", + "number_of_documents": 5, + "options": {"spider_options": {"page_limit": 10}}, + "duration": None, + } + + provider = WaterCrawlProvider(api_key="test_key") + status = provider.get_crawl_status("test-job-123") + + assert status["status"] == "active" + assert status["job_id"] == "test-job-123" + assert status["total"] == 10 + assert status["current"] == 5 + assert status["data"] == [] + + def test_get_crawl_status_completed(self, mocker: MockerFixture) -> None: + """Test getting status of a completed crawl job with results.""" + mock_client = mocker.patch("core.rag.extractor.watercrawl.provider.WaterCrawlAPIClient") + mock_instance = mock_client.return_value + mock_instance.get_crawl_request.return_value = { + "uuid": "test-job-456", + "status": "completed", + "number_of_documents": 10, + "options": {"spider_options": {"page_limit": 10}}, + "duration": "00:00:15.500000", + } + mock_instance.get_crawl_request_results.return_value = { + "results": [ + { + "url": "https://example.com/page1", + "result": { + "markdown": "# Page 1 Content", + "metadata": {"title": "Page 1", "description": "First page"}, + }, + } + ], + "next": None, + } + + provider = WaterCrawlProvider(api_key="test_key") + status = provider.get_crawl_status("test-job-456") + + assert status["status"] == "completed" + assert status["job_id"] == "test-job-456" + assert status["total"] == 10 + assert status["current"] == 10 + assert len(status["data"]) == 1 + assert status["time_consuming"] == 15.5 + + def test_get_crawl_url_data(self, mocker: MockerFixture) -> None: + """Test retrieving specific URL data from crawl results.""" + mock_client = mocker.patch("core.rag.extractor.watercrawl.provider.WaterCrawlAPIClient") + mock_instance = mock_client.return_value + mock_instance.get_crawl_request_results.return_value = { + "results": [ + { + "url": "https://example.com/target", + "result": { + "markdown": "# Target Page", + "metadata": {"title": "Target", "description": "Target page description"}, + }, + } + ], + "next": None, + } + + provider = WaterCrawlProvider(api_key="test_key") + data = provider.get_crawl_url_data("test-job-789", "https://example.com/target") + + assert data is not None + assert data["source_url"] == "https://example.com/target" + assert data["title"] == "Target" + assert data["markdown"] == "# Target Page" + + def test_get_crawl_url_data_not_found(self, mocker: MockerFixture) -> None: + """Test that None is returned when URL is not in results.""" + mock_client = mocker.patch("core.rag.extractor.watercrawl.provider.WaterCrawlAPIClient") + mock_instance = mock_client.return_value + mock_instance.get_crawl_request_results.return_value = {"results": list[object](), "next": None} + + provider = WaterCrawlProvider(api_key="test_key") + data = provider.get_crawl_url_data("test-job-789", "https://example.com/nonexistent") + + assert data is None + + # ============================================================================ + # Test WebsiteService - Multi-Provider Support + # ============================================================================ + + class TestWebsiteService: + """Test suite for WebsiteService with multiple providers.""" + + @patch("services.data_source.website_service.current_user") + @patch("services.data_source.website_service.DatasourceProviderService") + def test_crawl_url_firecrawl( + self, mock_provider_service: Mock, mock_current_user: Mock, mocker: MockerFixture + ) -> None: + """Test crawling with Firecrawl provider.""" + # Setup mocks + mock_current_user.current_tenant_id = "test_tenant" + mock_provider_service.return_value.get_datasource_credentials.return_value = { + "firecrawl_api_key": "test_key", + "base_url": "https://api.firecrawl.dev", + } + + mock_firecrawl = mocker.patch("services.data_source.website_service.FirecrawlApp") + mock_firecrawl_instance = mock_firecrawl.return_value + mock_firecrawl_instance.crawl_url.return_value = "job-123" + + # Mock redis + mocker.patch("services.data_source.website_service.redis_client") + + from services.data_source.website_service import WebsiteCrawlApiRequest + + api_request = WebsiteCrawlApiRequest( + provider="firecrawl", + url="https://example.com", + options={"limit": 10, "crawl_sub_pages": True, "only_main_content": True}, + ) + + result = WebsiteService( + providers=website_service_module.DatasourceProviderService( + credentials=create_autospec(DatasourceProviderCredentialStore, instance=True) + ) + ).crawl_url(api_request) + + assert result["status"] == "active" + assert result["job_id"] == "job-123" + + @patch("services.data_source.website_service.current_user") + @patch("services.data_source.website_service.DatasourceProviderService") + def test_crawl_url_watercrawl( + self, mock_provider_service: Mock, mock_current_user: Mock, mocker: MockerFixture + ) -> None: + """Test crawling with WaterCrawl provider.""" + # Setup mocks + mock_current_user.current_tenant_id = "test_tenant" + mock_provider_service.return_value.get_datasource_credentials.return_value = { + "api_key": "test_key", + "base_url": "https://app.watercrawl.dev", + } + + mock_watercrawl = mocker.patch("services.data_source.website_service.WaterCrawlProvider") + mock_watercrawl_instance = mock_watercrawl.return_value + mock_watercrawl_instance.crawl_url.return_value = {"status": "active", "job_id": "job-456"} + + from services.data_source.website_service import WebsiteCrawlApiRequest + + api_request = WebsiteCrawlApiRequest( + provider="watercrawl", + url="https://example.com", + options={"limit": 20, "crawl_sub_pages": True, "max_depth": 2}, + ) + + result = WebsiteService( + providers=website_service_module.DatasourceProviderService( + credentials=create_autospec(DatasourceProviderCredentialStore, instance=True) + ) + ).crawl_url(api_request) + + assert result["status"] == "active" + assert result["job_id"] == "job-456" + + @patch("services.data_source.website_service.current_user") + @patch("services.data_source.website_service.DatasourceProviderService") + def test_crawl_url_jinareader( + self, mock_provider_service: Mock, mock_current_user: Mock, mocker: MockerFixture + ) -> None: + """Test crawling with JinaReader provider.""" + # Setup mocks + mock_current_user.current_tenant_id = "test_tenant" + mock_provider_service.return_value.get_datasource_credentials.return_value = { + "api_key": "test_key", + } + + mock_response = httpx.Response(200, json={"code": 200, "data": {"taskId": "task-789"}}) + mocker.patch( + "services.data_source.website_service._adaptive_http_client.post", + return_value=mock_response, + ) + + from services.data_source.website_service import WebsiteCrawlApiRequest + + api_request = WebsiteCrawlApiRequest( + provider="jinareader", + url="https://example.com", + options={"limit": 15, "crawl_sub_pages": True, "use_sitemap": True}, + ) + + result = WebsiteService( + providers=website_service_module.DatasourceProviderService( + credentials=create_autospec(DatasourceProviderCredentialStore, instance=True) + ) + ).crawl_url(api_request) + + assert result["status"] == "active" + assert result["job_id"] == "task-789" + + def test_document_create_args_validate_success(self) -> None: + """Test validation of valid document creation arguments.""" + args = {"provider": "watercrawl", "url": "https://example.com", "options": {"limit": 10}} + + # Should not raise any exception + WebsiteService.document_create_args_validate(args) + + def test_document_create_args_validate_missing_provider(self) -> None: + """Test validation fails when provider is missing.""" + args = {"url": "https://example.com", "options": {"limit": 10}} + + with pytest.raises(ValueError, match="Provider is required"): + WebsiteService.document_create_args_validate(args) + + def test_document_create_args_validate_missing_url(self) -> None: + """Test validation fails when URL is missing.""" + args = {"provider": "watercrawl", "options": {"limit": 10}} + + with pytest.raises(ValueError, match="URL is required"): + WebsiteService.document_create_args_validate(args) + + def test_document_create_args_validate_missing_options(self) -> None: + """Test validation fails when options are missing.""" + args = {"provider": "watercrawl", "url": "https://example.com"} + + with pytest.raises(ValueError, match="Options are required"): + WebsiteService.document_create_args_validate(args) + + # ============================================================================ + # Test Link Following Logic + # ============================================================================ + + class TestLinkFollowingLogic: + """Test suite for link following and navigation logic.""" + + def test_link_following_with_includes(self, mocker: MockerFixture) -> None: + """Test that only links matching include patterns are followed.""" + mock_client = mocker.patch("core.rag.extractor.watercrawl.provider.WaterCrawlAPIClient") + mock_instance = mock_client.return_value + mock_instance.create_crawl_request.return_value = {"uuid": "test-job"} + + provider = WaterCrawlProvider(api_key="test_key") + options = {"crawl_sub_pages": True, "includes": "/blog/*,/news/*", "limit": 50} + provider.crawl_url("https://example.com", options=options) + + call_args = mock_instance.create_crawl_request.call_args + spider_options = call_args.kwargs["spider_options"] + + # Verify include paths are set for link filtering + assert "/blog/*" in spider_options["include_paths"] + assert "/news/*" in spider_options["include_paths"] + + def test_link_following_with_excludes(self, mocker: MockerFixture) -> None: + """Test that links matching exclude patterns are not followed.""" + mock_client = mocker.patch("core.rag.extractor.watercrawl.provider.WaterCrawlAPIClient") + mock_instance = mock_client.return_value + mock_instance.create_crawl_request.return_value = {"uuid": "test-job"} + + provider = WaterCrawlProvider(api_key="test_key") + options = {"crawl_sub_pages": True, "excludes": "/login/*,/logout/*", "limit": 50} + provider.crawl_url("https://example.com", options=options) + + call_args = mock_instance.create_crawl_request.call_args + spider_options = call_args.kwargs["spider_options"] + + # Verify exclude paths are set to prevent following certain links + assert "/login/*" in spider_options["exclude_paths"] + assert "/logout/*" in spider_options["exclude_paths"] + + def test_link_following_respects_max_depth(self, mocker: MockerFixture) -> None: + """Test that link following stops at specified max depth.""" + mock_client = mocker.patch("core.rag.extractor.watercrawl.provider.WaterCrawlAPIClient") + mock_instance = mock_client.return_value + mock_instance.create_crawl_request.return_value = {"uuid": "test-job"} + + provider = WaterCrawlProvider(api_key="test_key") + + # Test depth of 1 (only start page) + options = {"crawl_sub_pages": True, "max_depth": 1, "limit": 100} + provider.crawl_url("https://example.com", options=options) + + call_args = mock_instance.create_crawl_request.call_args + spider_options = call_args.kwargs["spider_options"] + assert spider_options["max_depth"] == 1 + + def test_link_following_page_limit(self, mocker: MockerFixture) -> None: + """Test that link following respects page limit.""" + mock_client = mocker.patch("core.rag.extractor.watercrawl.provider.WaterCrawlAPIClient") + mock_instance = mock_client.return_value + mock_instance.create_crawl_request.return_value = {"uuid": "test-job"} + + provider = WaterCrawlProvider(api_key="test_key") + options = {"crawl_sub_pages": True, "limit": 25, "max_depth": 5} + provider.crawl_url("https://example.com", options=options) + + call_args = mock_instance.create_crawl_request.call_args + spider_options = call_args.kwargs["spider_options"] + + # Verify page limit is set correctly + assert spider_options["page_limit"] == 25 + + # ============================================================================ + # Test Robots.txt Respect (Implicit in Provider Implementation) + # ============================================================================ + + class TestRobotsTxtRespect: + """ + Test suite for robots.txt compliance. + + Note: Robots.txt respect is typically handled by the underlying crawl + providers (Firecrawl, WaterCrawl, JinaReader). These tests verify that + the service layer properly configures providers to respect robots.txt. + """ + + def test_watercrawl_provider_respects_robots_txt(self, mocker: MockerFixture) -> None: + """ + Test that WaterCrawl provider is configured to respect robots.txt. + + WaterCrawl respects robots.txt by default in its implementation. + This test verifies the provider is initialized correctly. + """ + mock_client = mocker.patch("core.rag.extractor.watercrawl.provider.WaterCrawlAPIClient") + + provider = WaterCrawlProvider(api_key="test_key", base_url="https://app.watercrawl.dev/") + + # Verify provider is initialized with proper client + assert provider.client is not None + mock_client.assert_called_once_with("test_key", "https://app.watercrawl.dev/") + + def test_firecrawl_provider_respects_robots_txt(self) -> None: + """ + Test that Firecrawl provider respects robots.txt. + + Firecrawl respects robots.txt by default. This test ensures + the provider is configured correctly. + """ + from core.rag.extractor.firecrawl.firecrawl_app import FirecrawlApp + + # FirecrawlApp respects robots.txt in its implementation + app = FirecrawlApp(api_key="test_key", base_url="https://api.firecrawl.dev") + + assert app.api_key == "test_key" + assert app.base_url == "https://api.firecrawl.dev" + + def test_crawl_respects_domain_restrictions(self, mocker: MockerFixture) -> None: + """ + Test that crawl operations respect domain restrictions. + + This ensures that crawlers don't follow links to external domains + unless explicitly configured to do so. + """ + mock_client = mocker.patch("core.rag.extractor.watercrawl.provider.WaterCrawlAPIClient") + mock_instance = mock_client.return_value + mock_instance.create_crawl_request.return_value = {"uuid": "test-job"} + + provider = WaterCrawlProvider(api_key="test_key") + provider.crawl_url("https://example.com", options={"crawl_sub_pages": True}) + + call_args = mock_instance.create_crawl_request.call_args + spider_options = call_args.kwargs["spider_options"] + + # Verify allowed_domains is initialized (empty means same domain only) + assert "allowed_domains" in spider_options + assert isinstance(spider_options["allowed_domains"], list) + + # ============================================================================ + # Test Content Extraction + # ============================================================================ + + class TestContentExtraction: + """Test suite for content extraction from crawled pages.""" + + def test_structure_data_with_metadata(self, mocker: MockerFixture) -> None: + """Test that content is properly structured with metadata.""" + mocker.patch("core.rag.extractor.watercrawl.provider.WaterCrawlAPIClient") + + provider = WaterCrawlProvider(api_key="test_key") + + result_object = { + "url": "https://example.com/page", + "result": { + "markdown": "# Page Title\n\nPage content here.", + "metadata": { + "og:title": "Page Title", + "title": "Fallback Title", + "description": "Page description", + }, + }, + } + + structured = provider._structure_data(result_object) + + assert structured["title"] == "Page Title" + assert structured["description"] == "Page description" + assert structured["source_url"] == "https://example.com/page" + assert structured["markdown"] == "# Page Title\n\nPage content here." + + def test_structure_data_fallback_title(self, mocker: MockerFixture) -> None: + """Test that fallback title is used when og:title is not available.""" + mocker.patch("core.rag.extractor.watercrawl.provider.WaterCrawlAPIClient") + + provider = WaterCrawlProvider(api_key="test_key") + + result_object = { + "url": "https://example.com/page", + "result": {"markdown": "Content", "metadata": {"title": "Fallback Title"}}, + } + + structured = provider._structure_data(result_object) + + assert structured["title"] == "Fallback Title" + + def test_structure_data_invalid_result(self, mocker: MockerFixture) -> None: + """Test that ValueError is raised for invalid result objects.""" + mocker.patch("core.rag.extractor.watercrawl.provider.WaterCrawlAPIClient") + + provider = WaterCrawlProvider(api_key="test_key") + + # Result is a string instead of dict + result_object = {"url": "https://example.com/page", "result": "invalid string result"} + + with pytest.raises(ValueError, match="Invalid result object"): + provider._structure_data(result_object) + + def test_scrape_url_content_extraction(self, mocker: MockerFixture) -> None: + """Test content extraction from single URL scraping.""" + mock_client = mocker.patch("core.rag.extractor.watercrawl.provider.WaterCrawlAPIClient") + mock_instance = mock_client.return_value + mock_instance.scrape_url.return_value = { + "url": "https://example.com", + "result": { + "markdown": "# Main Content", + "metadata": {"og:title": "Example Page", "description": "Example description"}, + }, + } + + provider = WaterCrawlProvider(api_key="test_key") + result = provider.scrape_url("https://example.com") + + assert result["title"] == "Example Page" + assert result["description"] == "Example description" + assert result["markdown"] == "# Main Content" + assert result["source_url"] == "https://example.com" + + def test_only_main_content_extraction(self, mocker: MockerFixture) -> None: + """Test that only_main_content option filters out non-content elements.""" + mock_client = mocker.patch("core.rag.extractor.watercrawl.provider.WaterCrawlAPIClient") + mock_instance = mock_client.return_value + mock_instance.create_crawl_request.return_value = {"uuid": "test-job"} + + provider = WaterCrawlProvider(api_key="test_key") + options = {"only_main_content": True, "crawl_sub_pages": False} + provider.crawl_url("https://example.com", options=options) + + call_args = mock_instance.create_crawl_request.call_args + page_options = call_args.kwargs["page_options"] + + # Verify main content extraction is enabled + assert page_options["only_main_content"] is True + assert page_options["include_html"] is False + + # ============================================================================ + # Test Error Handling + # ============================================================================ + + class TestErrorHandling: + """Test suite for error handling in crawl operations.""" + + @patch("services.data_source.website_service.current_user") + @patch("services.data_source.website_service.DatasourceProviderService") + def test_invalid_provider_error(self, mock_provider_service: Mock, mock_current_user: Mock) -> None: + """Test that invalid provider raises ValueError.""" + from services.data_source.website_service import WebsiteCrawlApiRequest + + # Setup mocks + mock_current_user.current_tenant_id = "test_tenant" + mock_provider_service.return_value.get_datasource_credentials.return_value = { + "api_key": "test_key", + } + + api_request = WebsiteCrawlApiRequest( + provider="invalid_provider", url="https://example.com", options={"limit": 10} + ) + + # The error should be raised when trying to crawl with invalid provider + with pytest.raises(ValueError, match="Invalid provider"): + WebsiteService( + providers=website_service_module.DatasourceProviderService( + credentials=create_autospec(DatasourceProviderCredentialStore, instance=True) + ) + ).crawl_url(api_request) + + def test_missing_api_key_error(self, mocker: MockerFixture) -> None: + """Test that missing API key is handled properly at the httpx client level.""" + # Mock the client to avoid actual httpx initialization + mock_client = mocker.patch("core.rag.extractor.watercrawl.provider.WaterCrawlAPIClient") + + # Create provider with mocked client - should work with mock + WaterCrawlProvider(api_key="test_key") + + # Verify the client was initialized with the API key + mock_client.assert_called_once_with("test_key", None) + + def test_crawl_status_for_nonexistent_job(self, mocker: MockerFixture) -> None: + """Test handling of status check for non-existent job.""" + mock_client = mocker.patch("core.rag.extractor.watercrawl.provider.WaterCrawlAPIClient") + mock_instance = mock_client.return_value + + # Simulate API error for non-existent job + from core.rag.extractor.watercrawl.exceptions import WaterCrawlBadRequestError + + mock_response = httpx.Response(404, json={"message": "Crawl job not found", "errors": {}}) + mock_instance.get_crawl_request.side_effect = WaterCrawlBadRequestError(mock_response) + + provider = WaterCrawlProvider(api_key="test_key") + + with pytest.raises(WaterCrawlBadRequestError): + provider.get_crawl_status("nonexistent-job-id") + + # ============================================================================ + # Integration-style Tests + # ============================================================================ + + class TestCrawlWorkflow: + """Integration-style tests for complete crawl workflows.""" + + def test_complete_crawl_workflow(self, mocker: MockerFixture) -> None: + """Test a complete crawl workflow from start to finish.""" + mock_client = mocker.patch("core.rag.extractor.watercrawl.provider.WaterCrawlAPIClient") + mock_instance = mock_client.return_value + + # Step 1: Start crawl + mock_instance.create_crawl_request.return_value = {"uuid": "workflow-job-123"} + + provider = WaterCrawlProvider(api_key="test_key") + crawl_result = provider.crawl_url( + "https://example.com", options={"crawl_sub_pages": True, "limit": 5, "max_depth": 2} + ) + + assert crawl_result["job_id"] == "workflow-job-123" + + # Step 2: Check status (running) + mock_instance.get_crawl_request.return_value = { + "uuid": "workflow-job-123", + "status": "running", + "number_of_documents": 3, + "options": {"spider_options": {"page_limit": 5}}, + } + + status = provider.get_crawl_status("workflow-job-123") + assert status["status"] == "active" + assert status["current"] == 3 + + # Step 3: Check status (completed) + mock_instance.get_crawl_request.return_value = { + "uuid": "workflow-job-123", + "status": "completed", + "number_of_documents": 5, + "options": {"spider_options": {"page_limit": 5}}, + "duration": "00:00:10.000000", + } + mock_instance.get_crawl_request_results.return_value = { + "results": [ + { + "url": "https://example.com/page1", + "result": {"markdown": "Content 1", "metadata": {"title": "Page 1"}}, + }, + { + "url": "https://example.com/page2", + "result": {"markdown": "Content 2", "metadata": {"title": "Page 2"}}, + }, + ], + "next": None, + } + + status = provider.get_crawl_status("workflow-job-123") + assert status["status"] == "completed" + assert status["current"] == 5 + assert len(status["data"]) == 2 + + # Step 4: Get specific URL data + data = provider.get_crawl_url_data("workflow-job-123", "https://example.com/page1") + assert data is not None + assert data["title"] == "Page 1" + + def test_single_page_scrape_workflow(self, mocker: MockerFixture) -> None: + """Test workflow for scraping a single page without crawling.""" + mock_client = mocker.patch("core.rag.extractor.watercrawl.provider.WaterCrawlAPIClient") + mock_instance = mock_client.return_value + mock_instance.scrape_url.return_value = { + "url": "https://example.com/single-page", + "result": { + "markdown": "# Single Page\n\nThis is a single page scrape.", + "metadata": {"og:title": "Single Page", "description": "A single page"}, + }, + } + + provider = WaterCrawlProvider(api_key="test_key") + result = provider.scrape_url("https://example.com/single-page") + + assert result["title"] == "Single Page" + assert result["description"] == "A single page" + markdown = result["markdown"] + assert markdown is not None + assert "Single Page" in markdown + assert result["source_url"] == "https://example.com/single-page" + + # ============================================================================ + # Test Advanced Crawl Scenarios + # ============================================================================ + + class TestAdvancedCrawlScenarios: + """ + Test suite for advanced and edge-case crawling scenarios. + + This class tests complex crawling situations including: + - Pagination handling + - Large-scale crawls + - Concurrent crawl management + - Retry mechanisms + - Timeout handling + """ + + def test_pagination_in_crawl_results(self, mocker: MockerFixture) -> None: + """ + Test that pagination is properly handled when retrieving crawl results. + + When a crawl produces many results, they are paginated. This test + ensures that the provider correctly iterates through all pages. + """ + mock_client = mocker.patch("core.rag.extractor.watercrawl.provider.WaterCrawlAPIClient") + mock_instance = mock_client.return_value + + # Mock paginated responses - first page has 'next', second page doesn't + mock_instance.get_crawl_request_results.side_effect = [ + { + "results": [ + { + "url": f"https://example.com/page{i}", + "result": {"markdown": f"Content {i}", "metadata": {"title": f"Page {i}"}}, + } + for i in range(1, 101) + ], + "next": "page2", + }, + { + "results": [ + { + "url": f"https://example.com/page{i}", + "result": {"markdown": f"Content {i}", "metadata": {"title": f"Page {i}"}}, + } + for i in range(101, 151) + ], + "next": None, + }, + ] + + provider = WaterCrawlProvider(api_key="test_key") + + # Collect all results from paginated response + results = list(provider._get_results("test-job-id")) + + # Verify all pages were retrieved + assert len(results) == 150 + assert results[0]["title"] == "Page 1" + assert results[149]["title"] == "Page 150" + + def test_large_scale_crawl_configuration(self, mocker: MockerFixture) -> None: + """ + Test configuration for large-scale crawls with high page limits. + + Large-scale crawls require specific configuration to handle + hundreds or thousands of pages efficiently. + """ + mock_client = mocker.patch("core.rag.extractor.watercrawl.provider.WaterCrawlAPIClient") + mock_instance = mock_client.return_value + mock_instance.create_crawl_request.return_value = {"uuid": "large-crawl-job"} + + provider = WaterCrawlProvider(api_key="test_key") + + # Configure for large-scale crawl: 1000 pages, depth 5 + options = { + "crawl_sub_pages": True, + "limit": 1000, + "max_depth": 5, + "only_main_content": True, + "wait_time": 1500, + } + result = provider.crawl_url("https://example.com", options=options) + + # Verify crawl was initiated + assert result["status"] == "active" + assert result["job_id"] == "large-crawl-job" + + # Verify spider options for large crawl + call_args = mock_instance.create_crawl_request.call_args + spider_options = call_args.kwargs["spider_options"] + assert spider_options["page_limit"] == 1000 + assert spider_options["max_depth"] == 5 + + def test_crawl_with_custom_wait_time(self, mocker: MockerFixture) -> None: + """ + Test that custom wait times are properly applied to page loads. + + Wait times are crucial for dynamic content that loads via JavaScript. + This ensures pages have time to fully render before extraction. + """ + mock_client = mocker.patch("core.rag.extractor.watercrawl.provider.WaterCrawlAPIClient") + mock_instance = mock_client.return_value + mock_instance.create_crawl_request.return_value = {"uuid": "wait-test-job"} + + provider = WaterCrawlProvider(api_key="test_key") + + # Test with 3-second wait time for JavaScript-heavy pages + options = {"wait_time": 3000, "only_main_content": True} + provider.crawl_url("https://example.com/dynamic-page", options=options) + + call_args = mock_instance.create_crawl_request.call_args + page_options = call_args.kwargs["page_options"] + + # Verify wait time is set correctly + assert page_options["wait_time"] == 3000 + + def test_crawl_status_progress_tracking(self, mocker: MockerFixture) -> None: + """ + Test that crawl progress is accurately tracked and reported. + + Progress tracking allows users to monitor long-running crawls + and estimate completion time. + """ + mock_client = mocker.patch("core.rag.extractor.watercrawl.provider.WaterCrawlAPIClient") + mock_instance = mock_client.return_value + + # Simulate crawl at 60% completion + mock_instance.get_crawl_request.return_value = { + "uuid": "progress-job", + "status": "running", + "number_of_documents": 60, + "options": {"spider_options": {"page_limit": 100}}, + "duration": "00:01:30.000000", + } + + provider = WaterCrawlProvider(api_key="test_key") + status = provider.get_crawl_status("progress-job") + + # Verify progress metrics + assert status["status"] == "active" + assert status["current"] == 60 + assert status["total"] == 100 + # Calculate progress percentage + progress_percentage = (status["current"] / status["total"]) * 100 + assert progress_percentage == 60.0 + + def test_crawl_with_sitemap_usage(self, mocker: MockerFixture) -> None: + """ + Test that sitemap.xml is utilized when use_sitemap is enabled. + + Sitemaps provide a structured list of URLs, making crawls more + efficient and comprehensive. + """ + mock_client = mocker.patch("core.rag.extractor.watercrawl.provider.WaterCrawlAPIClient") + mock_instance = mock_client.return_value + mock_instance.create_crawl_request.return_value = {"uuid": "sitemap-job"} + + provider = WaterCrawlProvider(api_key="test_key") + + # Enable sitemap usage + options = {"crawl_sub_pages": True, "use_sitemap": True, "limit": 50} + provider.crawl_url("https://example.com", options=options) + + # Note: use_sitemap is passed to the service layer but not directly + # to WaterCrawl spider_options. This test verifies the option is accepted. + call_args = mock_instance.create_crawl_request.call_args + assert call_args is not None + + def test_empty_crawl_results(self, mocker: MockerFixture) -> None: + """ + Test handling of crawls that return no results. + + This can occur when all pages are excluded or no content matches + the extraction criteria. + """ + mock_client = mocker.patch("core.rag.extractor.watercrawl.provider.WaterCrawlAPIClient") + mock_instance = mock_client.return_value + mock_instance.get_crawl_request.return_value = { + "uuid": "empty-job", + "status": "completed", + "number_of_documents": 0, + "options": {"spider_options": {"page_limit": 10}}, + "duration": "00:00:05.000000", + } + mock_instance.get_crawl_request_results.return_value = {"results": list[object](), "next": None} + + provider = WaterCrawlProvider(api_key="test_key") + status = provider.get_crawl_status("empty-job") + + # Verify empty results are handled correctly + assert status["status"] == "completed" + assert status["current"] == 0 + assert status["total"] == 10 + assert len(status["data"]) == 0 + + def test_crawl_with_multiple_include_patterns(self, mocker: MockerFixture) -> None: + """ + Test crawling with multiple include patterns for fine-grained control. + + Multiple patterns allow targeting specific sections of a website + while excluding others. + """ + mock_client = mocker.patch("core.rag.extractor.watercrawl.provider.WaterCrawlAPIClient") + mock_instance = mock_client.return_value + mock_instance.create_crawl_request.return_value = {"uuid": "multi-pattern-job"} + + provider = WaterCrawlProvider(api_key="test_key") + + # Multiple include patterns for different content types + options = { + "crawl_sub_pages": True, + "includes": "/blog/*,/news/*,/articles/*,/docs/*", + "limit": 100, + } + provider.crawl_url("https://example.com", options=options) + + call_args = mock_instance.create_crawl_request.call_args + spider_options = call_args.kwargs["spider_options"] + + # Verify all include patterns are set + assert len(spider_options["include_paths"]) == 4 + assert "/blog/*" in spider_options["include_paths"] + assert "/news/*" in spider_options["include_paths"] + assert "/articles/*" in spider_options["include_paths"] + assert "/docs/*" in spider_options["include_paths"] + + def test_crawl_duration_calculation(self, mocker: MockerFixture) -> None: + """ + Test accurate calculation of crawl duration from time strings. + + Duration tracking helps analyze crawl performance and optimize + configuration for future crawls. + """ + mock_client = mocker.patch("core.rag.extractor.watercrawl.provider.WaterCrawlAPIClient") + mock_instance = mock_client.return_value + + # Test various duration formats + test_cases = [ + ("00:00:10.500000", 10.5), # 10.5 seconds + ("00:01:30.250000", 90.25), # 1 minute 30.25 seconds + ("01:15:45.750000", 4545.75), # 1 hour 15 minutes 45.75 seconds + ] + + for duration_str, expected_seconds in test_cases: + mock_instance.get_crawl_request.return_value = { + "uuid": "duration-test", + "status": "completed", + "number_of_documents": 10, + "options": {"spider_options": {"page_limit": 10}}, + "duration": duration_str, + } + mock_instance.get_crawl_request_results.return_value = {"results": list[object](), "next": None} + + provider = WaterCrawlProvider(api_key="test_key") + status = provider.get_crawl_status("duration-test") + + # Verify duration is calculated correctly + assert abs(status["time_consuming"] - expected_seconds) < 0.01 + + # ============================================================================ + # Test Provider-Specific Features + # ============================================================================ + + class TestProviderSpecificFeatures: + """ + Test suite for provider-specific features and behaviors. + + Different crawl providers (Firecrawl, WaterCrawl, JinaReader) have + unique features and API behaviors that require specific testing. + """ + + @patch("services.data_source.website_service.current_user") + @patch("services.data_source.website_service.DatasourceProviderService") + def test_firecrawl_with_prompt_parameter( + self, mock_provider_service: Mock, mock_current_user: Mock, mocker: MockerFixture + ) -> None: + """ + Test Firecrawl's prompt parameter for AI-guided extraction. + + Firecrawl v2 supports prompts to guide content extraction using AI, + allowing for semantic filtering of crawled content. + """ + # Setup mocks + mock_current_user.current_tenant_id = "test_tenant" + mock_provider_service.return_value.get_datasource_credentials.return_value = { + "firecrawl_api_key": "test_key", + "base_url": "https://api.firecrawl.dev", + } + + mock_firecrawl = mocker.patch("services.data_source.website_service.FirecrawlApp") + mock_firecrawl_instance = mock_firecrawl.return_value + mock_firecrawl_instance.crawl_url.return_value = "prompt-job-123" + + # Mock redis + mocker.patch("services.data_source.website_service.redis_client") + + from services.data_source.website_service import WebsiteCrawlApiRequest + + # Include a prompt for AI-guided extraction + api_request = WebsiteCrawlApiRequest( + provider="firecrawl", + url="https://example.com", + options={ + "limit": 20, + "crawl_sub_pages": True, + "only_main_content": True, + "prompt": "Extract only technical documentation and API references", + }, + ) + + result = WebsiteService( + providers=website_service_module.DatasourceProviderService( + credentials=create_autospec(DatasourceProviderCredentialStore, instance=True) + ) + ).crawl_url(api_request) + + assert result["status"] == "active" + assert result["job_id"] == "prompt-job-123" + + # Verify prompt was passed to Firecrawl + call_args = mock_firecrawl_instance.crawl_url.call_args + params = call_args[0][1] # Second argument is params + assert "prompt" in params + assert params["prompt"] == "Extract only technical documentation and API references" + + @patch("services.data_source.website_service.current_user") + @patch("services.data_source.website_service.DatasourceProviderService") + def test_jinareader_single_page_mode( + self, mock_provider_service: Mock, mock_current_user: Mock, mocker: MockerFixture + ) -> None: + """ + Test JinaReader's single-page scraping mode. + + JinaReader can scrape individual pages without crawling, + useful for quick content extraction. + """ + # Setup mocks + mock_current_user.current_tenant_id = "test_tenant" + mock_provider_service.return_value.get_datasource_credentials.return_value = { + "api_key": "test_key", + } + + response_data = { + "code": 200, + "data": { + "title": "Single Page Title", + "content": "Page content here", + "url": "https://example.com/page", + }, + } + mock_response = httpx.Response(200, json=response_data) + mocker.patch("services.data_source.website_service._jina_http_client.get", return_value=mock_response) + + from services.data_source.website_service import WebsiteCrawlApiRequest + + # Single page mode (crawl_sub_pages = False) + api_request = WebsiteCrawlApiRequest( + provider="jinareader", url="https://example.com/page", options={"crawl_sub_pages": False, "limit": 1} + ) + + result = WebsiteService( + providers=website_service_module.DatasourceProviderService( + credentials=create_autospec(DatasourceProviderCredentialStore, instance=True) + ) + ).crawl_url(api_request) + + # In single-page mode, JinaReader returns data immediately + assert result["status"] == "active" + assert "data" in result + + @patch("services.data_source.website_service.current_user") + @patch("services.data_source.website_service.DatasourceProviderService") + def test_watercrawl_with_tag_filtering( + self, mock_provider_service: Mock, mock_current_user: Mock, mocker: MockerFixture + ) -> None: + """ + Test WaterCrawl's HTML tag filtering capabilities. + + WaterCrawl allows including or excluding specific HTML tags + during content extraction for precise control. + """ + # Setup mocks + mock_current_user.current_tenant_id = "test_tenant" + mock_provider_service.return_value.get_datasource_credentials.return_value = { + "api_key": "test_key", + "base_url": "https://app.watercrawl.dev", + } + + mock_watercrawl = mocker.patch("services.data_source.website_service.WaterCrawlProvider") + mock_watercrawl_instance = mock_watercrawl.return_value + mock_watercrawl_instance.crawl_url.return_value = {"status": "active", "job_id": "tag-filter-job"} + + from services.data_source.website_service import WebsiteCrawlApiRequest + + # Configure with tag filtering + api_request = WebsiteCrawlApiRequest( + provider="watercrawl", + url="https://example.com", + options={ + "limit": 10, + "crawl_sub_pages": True, + "exclude_tags": "nav,footer,aside", + "include_tags": "article,main", + }, + ) + + result = WebsiteService( + providers=website_service_module.DatasourceProviderService( + credentials=create_autospec(DatasourceProviderCredentialStore, instance=True) + ) + ).crawl_url(api_request) + + assert result["status"] == "active" + assert result["job_id"] == "tag-filter-job" + + def test_firecrawl_base_url_configuration(self) -> None: + """ + Test that Firecrawl can be configured with custom base URLs. + + This is important for self-hosted Firecrawl instances or + different API endpoints. + """ + from core.rag.extractor.firecrawl.firecrawl_app import FirecrawlApp + + # Test with custom base URL + custom_base_url = "https://custom-firecrawl.example.com" + app = FirecrawlApp(api_key="test_key", base_url=custom_base_url) + + assert app.base_url == custom_base_url + assert app.api_key == "test_key" + + def test_watercrawl_base_url_default(self, mocker: MockerFixture) -> None: + """ + Test WaterCrawl's default base URL configuration. + + Verifies that the provider uses the correct default URL when + none is specified. + """ + mock_client = mocker.patch("core.rag.extractor.watercrawl.provider.WaterCrawlAPIClient") + + # Create provider without specifying base_url + WaterCrawlProvider(api_key="test_key") + + # Verify default base URL is used + mock_client.assert_called_once_with("test_key", None) + + # ============================================================================ + # Test Data Structure and Validation + # ============================================================================ + + class TestDataStructureValidation: + """ + Test suite for data structure validation and transformation. + + Ensures that crawled data is properly structured, validated, + and transformed into the expected format. + """ + + def test_crawl_request_to_api_request_conversion(self) -> None: + """ + Test conversion from API request to internal CrawlRequest format. + + This conversion ensures that external API parameters are properly + mapped to internal data structures. + """ + from services.data_source.website_service import WebsiteCrawlApiRequest + + # Create API request with all options + api_request = WebsiteCrawlApiRequest( + provider="watercrawl", + url="https://example.com", + options={ + "limit": 50, + "crawl_sub_pages": True, + "only_main_content": True, + "includes": "/blog/*", + "excludes": "/admin/*", + "prompt": "Extract main content", + "max_depth": 3, + "use_sitemap": True, + }, + ) + + # Convert to internal format + crawl_request = api_request.to_crawl_request() + + # Verify all fields are properly converted + assert crawl_request.url == "https://example.com" + assert crawl_request.provider == "watercrawl" + assert crawl_request.options.limit == 50 + assert crawl_request.options.crawl_sub_pages is True + assert crawl_request.options.only_main_content is True + assert crawl_request.options.includes == "/blog/*" + assert crawl_request.options.excludes == "/admin/*" + assert crawl_request.options.prompt == "Extract main content" + assert crawl_request.options.max_depth == 3 + assert crawl_request.options.use_sitemap is True + + def test_crawl_options_path_parsing(self) -> None: + """ + Test that include/exclude paths are correctly parsed from strings. + + Paths can be provided as comma-separated strings and must be + split into individual patterns. + """ + # Test with multiple paths + options = CrawlOptions(includes="/blog/*,/news/*,/docs/*", excludes="/admin/*,/private/*,/test/*") + + include_paths = options.get_include_paths() + exclude_paths = options.get_exclude_paths() + + # Verify parsing + assert len(include_paths) == 3 + assert "/blog/*" in include_paths + assert "/news/*" in include_paths + assert "/docs/*" in include_paths + + assert len(exclude_paths) == 3 + assert "/admin/*" in exclude_paths + assert "/private/*" in exclude_paths + assert "/test/*" in exclude_paths + + def test_crawl_options_with_whitespace(self) -> None: + """ + Test that whitespace in path strings is handled correctly. + + Users might include spaces around commas, which should be + handled gracefully. + """ + # Test with spaces around commas + options = CrawlOptions(includes=" /blog/* , /news/* , /docs/* ", excludes=" /admin/* , /private/* ") + + include_paths = options.get_include_paths() + exclude_paths = options.get_exclude_paths() + + # Verify paths are trimmed (note: current implementation doesn't trim, + # so paths will include spaces - this documents current behavior) + assert len(include_paths) == 3 + assert len(exclude_paths) == 2 + + def test_website_crawl_message_structure(self) -> None: + """ + Test the structure of WebsiteCrawlMessage entity. + + This entity wraps crawl results and must have the correct structure + for downstream processing. + """ + from core.datasource.entities.datasource_entities import WebsiteCrawlMessage, WebSiteInfo + + # Create a crawl message with results + web_info = WebSiteInfo(status="completed", web_info_list=[], total=10, completed=10) + + message = WebsiteCrawlMessage(result=web_info) + + # Verify structure + assert message.result.status == "completed" + assert message.result.total == 10 + assert message.result.completed == 10 + assert isinstance(message.result.web_info_list, list) + + def test_datasource_identity_structure(self) -> None: + """ + Test that DatasourceIdentity contains all required fields. + + Identity information is crucial for tracking and managing + datasource instances. + """ + identity = DatasourceIdentity( + author="test_author", + name="test_datasource", + label={"en_US": "Test Datasource", "zh_Hans": "测试数据源"}, + provider="test_provider", + icon="test_icon.svg", + ) + + # Verify all fields are present + assert identity.author == "test_author" + assert identity.name == "test_datasource" + assert identity.provider == "test_provider" + assert identity.icon == "test_icon.svg" + # I18nObject has attributes, not dict keys + assert identity.label.en_US == "Test Datasource" + assert identity.label.zh_Hans == "测试数据源" + + # ============================================================================ + # Test Edge Cases and Boundary Conditions + # ============================================================================ + + class TestEdgeCasesAndBoundaries: + """ + Test suite for edge cases and boundary conditions. + + These tests ensure robust handling of unusual inputs, limits, + and exceptional scenarios. + """ + + def test_crawl_with_zero_limit(self, mocker: MockerFixture) -> None: + """ + Test behavior when limit is set to zero. + + A zero limit should be handled gracefully, potentially defaulting + to a minimum value or raising an error. + """ + mock_client = mocker.patch("core.rag.extractor.watercrawl.provider.WaterCrawlAPIClient") + mock_instance = mock_client.return_value + mock_instance.create_crawl_request.return_value = {"uuid": "zero-limit-job"} + + provider = WaterCrawlProvider(api_key="test_key") + + # Attempt crawl with zero limit + options = {"crawl_sub_pages": True, "limit": 0} + result = provider.crawl_url("https://example.com", options=options) + + # Verify crawl was created (implementation may handle this differently) + assert result["status"] == "active" + + def test_crawl_with_very_large_limit(self, mocker: MockerFixture) -> None: + """ + Test crawl configuration with extremely large page limits. + + Very large limits should be accepted but may be subject to + provider-specific constraints. + """ + mock_client = mocker.patch("core.rag.extractor.watercrawl.provider.WaterCrawlAPIClient") + mock_instance = mock_client.return_value + mock_instance.create_crawl_request.return_value = {"uuid": "large-limit-job"} + + provider = WaterCrawlProvider(api_key="test_key") + + # Test with very large limit (10,000 pages) + options = {"crawl_sub_pages": True, "limit": 10000, "max_depth": 10} + result = provider.crawl_url("https://example.com", options=options) + + assert result["status"] == "active" + + call_args = mock_instance.create_crawl_request.call_args + spider_options = call_args.kwargs["spider_options"] + assert spider_options["page_limit"] == 10000 + + def test_crawl_with_empty_url(self) -> None: + """ + Test that empty URLs are rejected with appropriate error. + + Empty or invalid URLs should fail validation before attempting + to crawl. + """ + from services.data_source.website_service import WebsiteCrawlApiRequest + + # Empty URL should raise ValueError during validation + with pytest.raises(ValueError, match="URL is required"): + WebsiteCrawlApiRequest.from_args({"provider": "watercrawl", "url": "", "options": {"limit": 10}}) + + def test_crawl_with_special_characters_in_paths(self, mocker: MockerFixture) -> None: + """ + Test handling of special characters in include/exclude paths. + + Paths may contain special regex characters that need proper escaping + or handling. + """ + mock_client = mocker.patch("core.rag.extractor.watercrawl.provider.WaterCrawlAPIClient") + mock_instance = mock_client.return_value + mock_instance.create_crawl_request.return_value = {"uuid": "special-chars-job"} + + provider = WaterCrawlProvider(api_key="test_key") + + # Include paths with special characters + options = { + "crawl_sub_pages": True, + "includes": "/blog/[0-9]+/*,/category/(tech|science)/*", + "limit": 20, + } + provider.crawl_url("https://example.com", options=options) + + call_args = mock_instance.create_crawl_request.call_args + spider_options = call_args.kwargs["spider_options"] + + # Verify special characters are preserved + assert "/blog/[0-9]+/*" in spider_options["include_paths"] + assert "/category/(tech|science)/*" in spider_options["include_paths"] + + def test_crawl_status_with_null_duration(self, mocker: MockerFixture) -> None: + """ + Test handling of null/missing duration in crawl status. + + Duration may be null for active crawls or if timing data is unavailable. + """ + mock_client = mocker.patch("core.rag.extractor.watercrawl.provider.WaterCrawlAPIClient") + mock_instance = mock_client.return_value + mock_instance.get_crawl_request.return_value = { + "uuid": "null-duration-job", + "status": "running", + "number_of_documents": 5, + "options": {"spider_options": {"page_limit": 10}}, + "duration": None, # Null duration + } + + provider = WaterCrawlProvider(api_key="test_key") + status = provider.get_crawl_status("null-duration-job") + + # Verify null duration is handled (should default to 0) + assert status["time_consuming"] == 0 + + def test_structure_data_with_missing_metadata_fields(self, mocker: MockerFixture) -> None: + """ + Test content extraction when metadata fields are missing. + + Not all pages have complete metadata, so extraction should + handle missing fields gracefully. + """ + mocker.patch("core.rag.extractor.watercrawl.provider.WaterCrawlAPIClient") + + provider = WaterCrawlProvider(api_key="test_key") + + # Result with minimal metadata + result_object = { + "url": "https://example.com/minimal", + "result": { + "markdown": "# Minimal Content", + "metadata": dict[str, object](), # Empty metadata + }, + } + + structured = provider._structure_data(result_object) + + # Verify graceful handling of missing metadata + assert structured["title"] is None + assert structured["description"] is None + assert structured["source_url"] == "https://example.com/minimal" + assert structured["markdown"] == "# Minimal Content" + + def test_get_results_with_empty_pages(self, mocker: MockerFixture) -> None: + """ + Test pagination handling when some pages return empty results. + + Empty pages in pagination cause the loop to break early in the + current implementation, as per the code logic in _get_results. + """ + mock_client = mocker.patch("core.rag.extractor.watercrawl.provider.WaterCrawlAPIClient") + mock_instance = mock_client.return_value + + # First page has results, second page is empty (breaks loop) + mock_instance.get_crawl_request_results.side_effect = [ + { + "results": [ + { + "url": "https://example.com/page1", + "result": {"markdown": "Content 1", "metadata": {"title": "Page 1"}}, + } + ], + "next": "page2", + }, + {"results": list[object](), "next": None}, # Empty page breaks the loop + ] + + provider = WaterCrawlProvider(api_key="test_key") + results = list(provider._get_results("test-job")) + + # Current implementation breaks on empty results + # This documents the actual behavior + assert len(results) == 1 + assert results[0]["title"] == "Page 1" + + +class TestWebsiteCrawlPluginCases: + class TestWebsiteCrawlDatasourcePlugin: + @pytest.fixture + def entity(self) -> DatasourceEntity: + entity = datasource_entity("test_name") + entity.identity.provider = "test-provider" + entity.identity.name = "test-name" + return entity + + @pytest.fixture + def runtime(self) -> DatasourceRuntime: + runtime = DatasourceRuntime(tenant_id="test_tenant") + runtime.credentials = {"api_key": "test-key"} + return runtime + + def test_init(self, entity: DatasourceEntity, runtime: DatasourceRuntime) -> None: + # Arrange + tenant_id = "test-tenant-id" + icon = "test-icon" + plugin_unique_identifier = "test-plugin-id" # Act - result = plugin.get_website_crawl( - user_id=user_id, datasource_parameters=datasource_parameters, provider_type=provider_type + plugin = WebsiteCrawlDatasourcePlugin( + entity=entity, + runtime=runtime, + tenant_id=tenant_id, + icon=icon, + plugin_unique_identifier=plugin_unique_identifier, ) # Assert - assert isinstance(result, Generator) - messages = list(result) - assert len(messages) == 1 - assert messages[0] == message + assert plugin.tenant_id == tenant_id + assert plugin.plugin_unique_identifier == plugin_unique_identifier + assert plugin.entity == entity + assert plugin.runtime == runtime + assert plugin.icon == icon - mock_manager.get_website_crawl.assert_called_once_with( - tenant_id="test-tenant-id", - user_id=user_id, - datasource_provider="test-provider", - datasource_name="test-name", - credentials={"api_key": "test-key"}, - datasource_parameters=datasource_parameters, - provider_type=provider_type, + def test_datasource_provider_type(self, entity: DatasourceEntity, runtime: DatasourceRuntime) -> None: + # Arrange + plugin = WebsiteCrawlDatasourcePlugin( + entity=entity, runtime=runtime, tenant_id="test", icon="test", plugin_unique_identifier="test" ) + + # Act & Assert + assert plugin.datasource_provider_type() == DatasourceProviderType.WEBSITE_CRAWL + + def test_get_website_crawl(self, entity: DatasourceEntity, runtime: DatasourceRuntime) -> None: + # Arrange + plugin = WebsiteCrawlDatasourcePlugin( + entity=entity, + runtime=runtime, + tenant_id="test-tenant-id", + icon="test-icon", + plugin_unique_identifier="test-plugin-id", + ) + + user_id = "test-user-id" + datasource_parameters = {"url": "https://example.com"} + provider_type = "firecrawl" + + message = WebsiteCrawlMessage() + + # Mock PluginDatasourceManager + with patch( + "core.datasource.website_crawl.website_crawl_plugin.PluginDatasourceManager" + ) as mock_manager_class: + mock_manager = mock_manager_class.return_value + mock_manager.get_website_crawl.return_value = (msg for msg in [message]) + + # Act + result = plugin.get_website_crawl( + user_id=user_id, datasource_parameters=datasource_parameters, provider_type=provider_type + ) + + # Assert + assert isinstance(result, Generator) + messages = list(result) + assert len(messages) == 1 + assert messages[0] == message + + mock_manager.get_website_crawl.assert_called_once_with( + tenant_id="test-tenant-id", + user_id=user_id, + datasource_provider="test-provider", + datasource_name="test-name", + credentials={"api_key": "test-key"}, + datasource_parameters=datasource_parameters, + provider_type=provider_type, + ) diff --git a/api/tests/unit_tests/core/datasource/website_crawl/test_website_crawl_provider.py b/api/tests/unit_tests/core/datasource/website_crawl/test_website_crawl_provider.py index 3c988f224b4..74810260efc 100644 --- a/api/tests/unit_tests/core/datasource/website_crawl/test_website_crawl_provider.py +++ b/api/tests/unit_tests/core/datasource/website_crawl/test_website_crawl_provider.py @@ -7,7 +7,7 @@ from core.datasource.entities.datasource_entities import ( ) from core.datasource.website_crawl.website_crawl_plugin import WebsiteCrawlDatasourcePlugin from core.datasource.website_crawl.website_crawl_provider import WebsiteCrawlDatasourcePluginProviderController -from tests.unit_tests.core.datasource.factories import datasource_entity, provider_entity +from tests.helpers.unit.core.datasource.factories import datasource_entity, provider_entity class TestWebsiteCrawlDatasourcePluginProviderController: diff --git a/api/tests/unit_tests/core/entities/test_entities_agent_entities.py b/api/tests/unit_tests/core/entities/test_agent_entities.py similarity index 100% rename from api/tests/unit_tests/core/entities/test_entities_agent_entities.py rename to api/tests/unit_tests/core/entities/test_agent_entities.py diff --git a/api/tests/unit_tests/core/entities/test_entities_document_task.py b/api/tests/unit_tests/core/entities/test_document_task.py similarity index 100% rename from api/tests/unit_tests/core/entities/test_entities_document_task.py rename to api/tests/unit_tests/core/entities/test_document_task.py diff --git a/api/tests/unit_tests/core/entities/test_entities_embedding_type.py b/api/tests/unit_tests/core/entities/test_embedding_type.py similarity index 100% rename from api/tests/unit_tests/core/entities/test_entities_embedding_type.py rename to api/tests/unit_tests/core/entities/test_embedding_type.py diff --git a/api/tests/unit_tests/core/entities/test_entities_provider_configuration.py b/api/tests/unit_tests/core/entities/test_entities_provider_configuration.py deleted file mode 100644 index 6ebcd317360..00000000000 --- a/api/tests/unit_tests/core/entities/test_entities_provider_configuration.py +++ /dev/null @@ -1,1474 +0,0 @@ -from __future__ import annotations - -import json -import logging -from collections.abc import Iterator -from contextlib import contextmanager -from unittest.mock import Mock, PropertyMock, patch - -import pytest -from sqlalchemy import Engine, event, select -from sqlalchemy.orm import Session - -from constants import HIDDEN_VALUE -from core.entities.model_entities import ModelStatus, ProviderModelWithStatusEntity -from core.entities.provider_configuration import ProviderConfiguration, ProviderConfigurations -from core.entities.provider_entities import ( - CredentialConfiguration, - CustomConfiguration, - CustomModelConfiguration, - CustomProviderConfiguration, - ModelLoadBalancingConfiguration, - ModelSettings, - ProviderQuotaType, - QuotaConfiguration, - QuotaUnit, - RestrictModel, - SystemConfiguration, - SystemConfigurationStatus, -) -from core.helper.model_provider_cache import ProviderCredentialsCacheType -from core.plugin.impl.model_runtime_factory import PluginModelAssembly, create_model_type_instance -from extensions.ext_database import db -from graphon.model_runtime.entities.common_entities import I18nObject -from graphon.model_runtime.entities.model_entities import AIModelEntity, FetchFrom, ModelType -from graphon.model_runtime.entities.provider_entities import ( - ConfigurateMethod, - CredentialFormSchema, - FieldModelSchema, - FormType, - ModelCredentialSchema, - ProviderCredentialSchema, - ProviderEntity, -) -from graphon.model_runtime.model_providers.base.large_language_model import LargeLanguageModel -from models.enums import CredentialSourceType -from models.provider import ( - LoadBalancingModelConfig, - Provider, - ProviderCredential, - ProviderModel, - ProviderModelCredential, - ProviderModelSetting, - ProviderType, - TenantPreferredModelProvider, -) -from models.provider_ids import ModelProviderID - - -def _build_provider_configuration(*, provider_name: str = "openai") -> ProviderConfiguration: - provider_entity = ProviderEntity( - provider=provider_name, - label=I18nObject(en_US="OpenAI"), - supported_model_types=[ModelType.LLM], - configurate_methods=[ConfigurateMethod.PREDEFINED_MODEL], - ) - system_configuration = SystemConfiguration( - enabled=True, - credentials={"api_key": "test-key"}, - current_quota_type=ProviderQuotaType.TRIAL, - quota_configurations=[ - QuotaConfiguration( - quota_type=ProviderQuotaType.TRIAL, - quota_unit=QuotaUnit.TOKENS, - quota_limit=1_000, - quota_used=0, - is_valid=True, - restrict_models=[], - ) - ], - ) - with patch("core.entities.provider_configuration.original_provider_configurate_methods", {}): - return ProviderConfiguration( - tenant_id="tenant-1", - provider=provider_entity, - preferred_provider_type=ProviderType.SYSTEM, - using_provider_type=ProviderType.SYSTEM, - system_configuration=system_configuration, - custom_configuration=CustomConfiguration(provider=None, models=[]), - model_settings=[], - ) - - -def _build_ai_model(name: str, *, model_type: ModelType = ModelType.LLM) -> AIModelEntity: - return AIModelEntity( - model=name, - label=I18nObject(en_US=name), - model_type=model_type, - fetch_from=FetchFrom.PREDEFINED_MODEL, - model_properties={}, - ) - - -def _build_provider_model(name: str) -> ProviderModelWithStatusEntity: - return ProviderModelWithStatusEntity(**_build_ai_model(name).model_dump(), status=ModelStatus.ACTIVE) - - -def _build_secret_provider_schema() -> ProviderCredentialSchema: - return ProviderCredentialSchema( - credential_form_schemas=[ - CredentialFormSchema( - variable="openai_api_key", - label=I18nObject(en_US="API Key"), - type=FormType.SECRET_INPUT, - ) - ] - ) - - -def _build_secret_model_schema() -> ModelCredentialSchema: - return ModelCredentialSchema( - model=FieldModelSchema(label=I18nObject(en_US="Model")), - credential_form_schemas=[ - CredentialFormSchema( - variable="openai_api_key", - label=I18nObject(en_US="API Key"), - type=FormType.SECRET_INPUT, - ) - ], - ) - - -def test_extract_secret_variables_returns_only_secret_inputs() -> None: - configuration = _build_provider_configuration() - credential_form_schemas = [ - CredentialFormSchema( - variable="api_key", - label=I18nObject(en_US="API Key"), - type=FormType.SECRET_INPUT, - ), - CredentialFormSchema( - variable="endpoint", - label=I18nObject(en_US="Endpoint"), - type=FormType.TEXT_INPUT, - ), - ] - - secret_variables = configuration.extract_secret_variables(credential_form_schemas) - assert secret_variables == ["api_key"] - - -def test_obfuscated_credentials_masks_only_secret_fields() -> None: - configuration = _build_provider_configuration() - credential_form_schemas = [ - CredentialFormSchema( - variable="api_key", - label=I18nObject(en_US="API Key"), - type=FormType.SECRET_INPUT, - ), - CredentialFormSchema( - variable="endpoint", - label=I18nObject(en_US="Endpoint"), - type=FormType.TEXT_INPUT, - ), - ] - - with patch( - "core.entities.provider_configuration.encrypter.obfuscated_token", - side_effect=lambda value: f"masked-{value[-2:]}", - ): - obfuscated = configuration.obfuscated_credentials( - credentials={"api_key": "sk-test-1234", "endpoint": "https://api.example.com"}, - credential_form_schemas=credential_form_schemas, - ) - - assert obfuscated["api_key"] == "masked-34" - assert obfuscated["endpoint"] == "https://api.example.com" - - -def test_provider_configurations_behave_like_keyed_container() -> None: - configuration = _build_provider_configuration() - provider_key = str(ModelProviderID("openai")) - configurations = ProviderConfigurations(tenant_id="tenant-1") - - configurations[provider_key] = configuration - - assert "openai" in configurations - assert configurations["openai"] is configuration - assert configurations.get("openai") is configuration - assert configurations.to_list() == [configuration] - assert list(configurations) == [(provider_key, configuration)] - - -def test_provider_configurations_get_models_forwards_filters() -> None: - configuration = _build_provider_configuration() - provider_key = str(ModelProviderID("openai")) - configurations = ProviderConfigurations(tenant_id="tenant-1") - configurations[provider_key] = configuration - expected_model = _build_provider_model("gpt-4o") - - with patch.object(ProviderConfiguration, "get_provider_models", return_value=[expected_model]) as mock_get: - models = configurations.get_models(provider="openai", model_type=ModelType.LLM, only_active=True) - - mock_get.assert_called_once_with(ModelType.LLM, True) - assert models == [expected_model] - - -def test_provider_configurations_get_models_skips_non_matching_provider_filter() -> None: - configuration = _build_provider_configuration() - provider_key = str(ModelProviderID("openai")) - configurations = ProviderConfigurations(tenant_id="tenant-1") - configurations[provider_key] = configuration - - with patch.object( - ProviderConfiguration, "get_provider_models", return_value=[_build_provider_model("gpt-4o")] - ) as mock_get: - models = configurations.get_models(provider="anthropic", model_type=ModelType.LLM, only_active=True) - - assert models == [] - mock_get.assert_not_called() - - -def test_get_current_credentials_custom_provider_checks_current_credential() -> None: - configuration = _build_provider_configuration() - configuration.using_provider_type = ProviderType.CUSTOM - configuration.custom_configuration.provider = CustomProviderConfiguration( - credentials={"api_key": "provider-key"}, - current_credential_id="credential-1", - current_credential_name="Primary", - available_credentials=[], - ) - - with patch("core.helper.credential_utils.check_credential_policy_compliance") as mock_check: - credentials = configuration.get_current_credentials(ModelType.LLM, "gpt-4o") - - assert credentials == {"api_key": "provider-key"} - assert mock_check.call_count == 1 - assert mock_check.call_args.kwargs["credential_id"] == "credential-1" - assert mock_check.call_args.kwargs["provider"] == "openai" - - -def test_get_current_credentials_custom_provider_checks_all_available_credentials() -> None: - configuration = _build_provider_configuration() - configuration.using_provider_type = ProviderType.CUSTOM - configuration.custom_configuration.provider = CustomProviderConfiguration( - credentials={"api_key": "provider-key"}, - available_credentials=[ - CredentialConfiguration(credential_id="cred-1", credential_name="First"), - CredentialConfiguration(credential_id="cred-2", credential_name="Second"), - ], - ) - - with patch("core.helper.credential_utils.check_credential_policy_compliance") as mock_check: - credentials = configuration.get_current_credentials(ModelType.LLM, "gpt-4o") - - assert credentials == {"api_key": "provider-key"} - assert [c.kwargs["credential_id"] for c in mock_check.call_args_list] == ["cred-1", "cred-2"] - assert all(c.kwargs["provider"] == "openai" for c in mock_check.call_args_list) - - -def test_get_system_configuration_status_returns_none_when_current_quota_missing() -> None: - configuration = _build_provider_configuration() - configuration.system_configuration.current_quota_type = ProviderQuotaType.FREE - - status = configuration.get_system_configuration_status() - assert status is None - - -def test_get_provider_names_supports_legacy_and_full_plugin_id() -> None: - configuration = _build_provider_configuration() - configuration.provider.provider = "langgenius/openai/openai" - - provider_names = configuration._get_provider_names() - assert provider_names == ["langgenius/openai/openai", "openai"] - - -def test_get_model_type_instance_and_schema_delegate_to_factory() -> None: - configuration = _build_provider_configuration() - assembly = PluginModelAssembly(tenant_id="tenant-1") - schema = _build_ai_model("gpt-4o") - - with ( - patch("core.entities.provider_configuration.create_plugin_model_assembly", return_value=assembly) as builder, - patch( - "core.entities.provider_configuration.create_model_type_instance", wraps=create_model_type_instance - ) as model_builder, - patch.object(assembly.model_runtime, "get_model_schema", return_value=schema) as get_schema, - ): - model_type_instance = configuration.get_model_type_instance(ModelType.LLM) - model_schema = configuration.get_model_schema(ModelType.LLM, "gpt-4o", {"api_key": "x"}) - - assert isinstance(model_type_instance, LargeLanguageModel) - assert model_schema is schema - assert builder.call_count == 2 - model_builder.assert_called_once_with( - runtime=assembly.model_runtime, - provider_schema=configuration.provider, - model_type=ModelType.LLM, - ) - get_schema.assert_called_once_with( - provider="openai", model_type=ModelType.LLM, model="gpt-4o", credentials={"api_key": "x"} - ) - - -def test_get_model_type_instance_and_schema_reuse_bound_runtime_factory() -> None: - configuration = _build_provider_configuration() - bound_runtime = PluginModelAssembly(tenant_id="tenant-1").model_runtime - schema = _build_ai_model("gpt-4o") - configuration.bind_model_runtime(bound_runtime) - - with ( - patch("core.entities.provider_configuration.ModelProviderFactory") as factory_cls, - patch("core.entities.provider_configuration.create_plugin_model_assembly") as assembly_builder, - patch( - "core.entities.provider_configuration.create_model_type_instance", wraps=create_model_type_instance - ) as model_builder, - patch.object(bound_runtime, "get_model_schema", return_value=schema) as get_schema, - ): - model_type_instance = configuration.get_model_type_instance(ModelType.LLM) - model_schema = configuration.get_model_schema(ModelType.LLM, "gpt-4o", {"api_key": "x"}) - - assert isinstance(model_type_instance, LargeLanguageModel) - assert model_schema is schema - factory_cls.assert_not_called() - assembly_builder.assert_not_called() - model_builder.assert_called_once_with( - runtime=bound_runtime, - provider_schema=configuration.provider, - model_type=ModelType.LLM, - ) - get_schema.assert_called_once_with( - provider="openai", model_type=ModelType.LLM, model="gpt-4o", credentials={"api_key": "x"} - ) - - -def test_get_provider_model_returns_none_when_model_not_found() -> None: - configuration = _build_provider_configuration() - other_model = _build_provider_model("other-model") - - with patch.object(ProviderConfiguration, "get_provider_models", return_value=[other_model]): - selected = configuration.get_provider_model(ModelType.LLM, "gpt-4o") - - assert selected is None - - -def test_get_provider_models_system_deduplicates_sorts_and_filters_active() -> None: - configuration = _build_provider_configuration() - configuration.provider.position = {"llm": ["b-model", "a-model"]} - configuration.model_settings = [ - ModelSettings(model="a-model", model_type=ModelType.LLM, enabled=False, load_balancing_configs=[]) - ] - provider_schema = ProviderEntity( - provider="openai", - label=I18nObject(en_US="OpenAI"), - supported_model_types=[ModelType.LLM], - configurate_methods=[ConfigurateMethod.PREDEFINED_MODEL], - models=[_build_ai_model("a-model"), _build_ai_model("b-model"), _build_ai_model("a-model")], - ) - assembly = PluginModelAssembly(tenant_id="tenant-1") - - with ( - patch("core.entities.provider_configuration.create_plugin_model_assembly", return_value=assembly), - patch.object(assembly.model_provider_factory, "get_model_provider", return_value=provider_schema), - ): - all_models = configuration.get_provider_models(model_type=ModelType.LLM, only_active=False) - active_models = configuration.get_provider_models(model_type=ModelType.LLM, only_active=True) - - assert [model.model for model in all_models] == ["b-model", "a-model"] - assert [model.status for model in all_models] == [ModelStatus.ACTIVE, ModelStatus.DISABLED] - assert [model.model for model in active_models] == ["b-model"] - - -def test_get_provider_models_system_filters_requested_model() -> None: - configuration = _build_provider_configuration() - provider_schema = ProviderEntity( - provider="openai", - label=I18nObject(en_US="OpenAI"), - supported_model_types=[ModelType.LLM], - configurate_methods=[ConfigurateMethod.PREDEFINED_MODEL], - models=[_build_ai_model("a-model"), _build_ai_model("target-model"), _build_ai_model("b-model")], - ) - assembly = PluginModelAssembly(tenant_id="tenant-1") - - with ( - patch("core.entities.provider_configuration.create_plugin_model_assembly", return_value=assembly), - patch.object(assembly.model_provider_factory, "get_model_provider", return_value=provider_schema), - ): - models = configuration.get_provider_models( - model_type=ModelType.LLM, - only_active=False, - model="target-model", - ) - - assert [model.model for model in models] == ["target-model"] - - -def test_get_provider_models_system_customizable_filters_requested_restricted_model() -> None: - provider = ProviderEntity( - provider="openai", - label=I18nObject(en_US="OpenAI"), - supported_model_types=[ModelType.LLM], - configurate_methods=[ConfigurateMethod.CUSTOMIZABLE_MODEL], - ) - system_configuration = SystemConfiguration( - enabled=True, - credentials={"api_key": "test-key"}, - current_quota_type=ProviderQuotaType.TRIAL, - quota_configurations=[ - QuotaConfiguration( - quota_type=ProviderQuotaType.TRIAL, - quota_unit=QuotaUnit.TOKENS, - quota_limit=1_000, - quota_used=0, - is_valid=True, - restrict_models=[ - RestrictModel(model="target-model", base_model_name="base-model", model_type=ModelType.LLM), - RestrictModel(model="other-model", base_model_name="base-model", model_type=ModelType.LLM), - ], - ) - ], - ) - provider_schema = ProviderEntity( - provider="openai", - label=I18nObject(en_US="OpenAI"), - supported_model_types=[ModelType.LLM], - configurate_methods=[ConfigurateMethod.PREDEFINED_MODEL], - models=[], - ) - assembly = PluginModelAssembly(tenant_id="tenant-1") - - with patch("core.entities.provider_configuration.original_provider_configurate_methods", {}): - configuration = ProviderConfiguration( - tenant_id="tenant-1", - provider=provider, - preferred_provider_type=ProviderType.SYSTEM, - using_provider_type=ProviderType.SYSTEM, - system_configuration=system_configuration, - custom_configuration=CustomConfiguration(provider=None, models=[]), - model_settings=[], - ) - - with ( - patch("core.entities.provider_configuration.create_plugin_model_assembly", return_value=assembly), - patch.object(assembly.model_provider_factory, "get_model_provider", return_value=provider_schema), - patch.object( - ProviderConfiguration, - "get_model_schema", - side_effect=lambda *args, **kwargs: _build_ai_model(kwargs["model"]), - ) as mock_get_model_schema, - ): - models = configuration.get_provider_models( - model_type=ModelType.LLM, - only_active=False, - model="target-model", - ) - - assert [model.model for model in models] == ["target-model"] - mock_get_model_schema.assert_called_once() - assert mock_get_model_schema.call_args.kwargs["model"] == "target-model" - - -def test_get_custom_provider_models_sets_status_for_removed_credentials_and_invalid_lb_configs() -> None: - configuration = _build_provider_configuration() - configuration.using_provider_type = ProviderType.CUSTOM - configuration.custom_configuration.provider = CustomProviderConfiguration(credentials={"api_key": "provider-key"}) - configuration.custom_configuration.models = [ - CustomModelConfiguration( - model="custom-model", - model_type=ModelType.LLM, - credentials=None, - available_model_credentials=[CredentialConfiguration(credential_id="c-1", credential_name="first")], - ) - ] - provider_schema = ProviderEntity( - provider="openai", - label=I18nObject(en_US="OpenAI"), - supported_model_types=[ModelType.LLM], - configurate_methods=[ConfigurateMethod.PREDEFINED_MODEL], - models=[_build_ai_model("base-model")], - ) - model_setting_map = { - ModelType.LLM: { - "base-model": ModelSettings( - model="base-model", - model_type=ModelType.LLM, - enabled=True, - load_balancing_enabled=True, - load_balancing_configs=[ - ModelLoadBalancingConfiguration( - id="lb-base", - name="LB Base", - credentials={}, - credential_source_type=CredentialSourceType.PROVIDER, - ) - ], - ), - "custom-model": ModelSettings( - model="custom-model", - model_type=ModelType.LLM, - enabled=True, - load_balancing_enabled=True, - load_balancing_configs=[ - ModelLoadBalancingConfiguration( - id="lb-custom", - name="LB Custom", - credentials={}, - credential_source_type=CredentialSourceType.CUSTOM_MODEL, - ) - ], - ), - } - } - - with patch.object(ProviderConfiguration, "get_model_schema", return_value=_build_ai_model("custom-model")): - models = configuration._get_custom_provider_models( - model_types=[ModelType.LLM], - provider_schema=provider_schema, - model_setting_map=model_setting_map, - ) - - status_map = {model.model: model.status for model in models} - invalid_lb_map = {model.model: model.has_invalid_load_balancing_configs for model in models} - assert status_map["base-model"] == ModelStatus.ACTIVE - assert status_map["custom-model"] == ModelStatus.CREDENTIAL_REMOVED - assert invalid_lb_map["base-model"] is True - assert invalid_lb_map["custom-model"] is True - - -def test_get_custom_provider_models_filters_requested_base_model() -> None: - configuration = _build_provider_configuration() - configuration.using_provider_type = ProviderType.CUSTOM - configuration.custom_configuration.provider = CustomProviderConfiguration(credentials={"api_key": "provider-key"}) - provider_schema = ProviderEntity( - provider="openai", - label=I18nObject(en_US="OpenAI"), - supported_model_types=[ModelType.LLM], - configurate_methods=[ConfigurateMethod.PREDEFINED_MODEL], - models=[_build_ai_model("base-model"), _build_ai_model("target-model")], - ) - - models = configuration._get_custom_provider_models( - model_types=[ModelType.LLM], - provider_schema=provider_schema, - model_setting_map={}, - model="target-model", - ) - - assert [model.model for model in models] == ["target-model"] - - -def test_get_provider_models_reuses_cached_provider_schema() -> None: - configuration = _build_provider_configuration() - provider_schema = ProviderEntity( - provider="openai", - label=I18nObject(en_US="OpenAI"), - supported_model_types=[ModelType.LLM], - configurate_methods=[ConfigurateMethod.PREDEFINED_MODEL], - models=[_build_ai_model("a-model"), _build_ai_model("b-model")], - ) - configuration.provider = provider_schema - - with patch( - "core.entities.provider_configuration.create_plugin_model_assembly", - ) as mock_assembly_builder: - configuration.get_provider_models(model_type=ModelType.LLM, model="a-model") - configuration.get_provider_models(model_type=ModelType.LLM, model="b-model") - - mock_assembly_builder.assert_not_called() - - -def test_validator_adds_predefined_model_for_customizable_provider_with_restrictions() -> None: - provider = ProviderEntity( - provider="openai", - label=I18nObject(en_US="OpenAI"), - supported_model_types=[ModelType.LLM], - configurate_methods=[ConfigurateMethod.CUSTOMIZABLE_MODEL], - ) - system_configuration = SystemConfiguration( - enabled=True, - credentials={"api_key": "test-key"}, - current_quota_type=ProviderQuotaType.TRIAL, - quota_configurations=[ - QuotaConfiguration( - quota_type=ProviderQuotaType.TRIAL, - quota_unit=QuotaUnit.TOKENS, - quota_limit=100, - quota_used=0, - is_valid=True, - restrict_models=[ - RestrictModel(model="restricted", base_model_name="base-model", model_type=ModelType.LLM) - ], - ) - ], - ) - with patch("core.entities.provider_configuration.original_provider_configurate_methods", {}): - configuration = ProviderConfiguration( - tenant_id="tenant-1", - provider=provider, - preferred_provider_type=ProviderType.SYSTEM, - using_provider_type=ProviderType.SYSTEM, - system_configuration=system_configuration, - custom_configuration=CustomConfiguration(provider=None, models=[]), - model_settings=[], - ) - - assert ConfigurateMethod.PREDEFINED_MODEL in configuration.provider.configurate_methods - - -def test_get_current_credentials_system_handles_disable_and_restricted_base_model() -> None: - configuration = _build_provider_configuration() - configuration.model_settings = [ - ModelSettings(model="gpt-4o", model_type=ModelType.LLM, enabled=False, load_balancing_configs=[]) - ] - - with pytest.raises(ValueError, match="Model gpt-4o is disabled"): - configuration.get_current_credentials(ModelType.LLM, "gpt-4o") - - configuration.model_settings = [] - configuration.system_configuration.quota_configurations[0].restrict_models = [ - RestrictModel(model="gpt-4o", base_model_name="base-model", model_type=ModelType.LLM) - ] - credentials = configuration.get_current_credentials(ModelType.LLM, "gpt-4o") - assert credentials["base_model_name"] == "base-model" - - -def test_get_current_credentials_prefers_model_specific_custom_credentials() -> None: - configuration = _build_provider_configuration() - configuration.using_provider_type = ProviderType.CUSTOM - configuration.custom_configuration.models = [ - CustomModelConfiguration( - model="gpt-4o", - model_type=ModelType.LLM, - credentials={"api_key": "model-key"}, - ) - ] - configuration.custom_configuration.provider = CustomProviderConfiguration(credentials={"api_key": "provider-key"}) - - credentials = configuration.get_current_credentials(ModelType.LLM, "gpt-4o") - assert credentials == {"api_key": "model-key"} - - -def test_get_system_configuration_status_falsey_quota_returns_unsupported() -> None: - class _FalseyQuota: - quota_type = ProviderQuotaType.TRIAL - is_valid = True - - def __bool__(self) -> bool: - return False - - configuration = _build_provider_configuration() - configuration.system_configuration.quota_configurations = [_FalseyQuota()] # type: ignore[list-item] - assert configuration.get_system_configuration_status() == SystemConfigurationStatus.UNSUPPORTED - - -def test_get_custom_provider_models_skips_schema_models_with_mismatched_type() -> None: - configuration = _build_provider_configuration() - provider_schema = ProviderEntity( - provider="openai", - label=I18nObject(en_US="OpenAI"), - supported_model_types=[ModelType.LLM, ModelType.TEXT_EMBEDDING], - configurate_methods=[ConfigurateMethod.PREDEFINED_MODEL], - models=[ - _build_ai_model("llm-model", model_type=ModelType.LLM), - _build_ai_model("embed-model", model_type=ModelType.TEXT_EMBEDDING), - ], - ) - - models = configuration._get_custom_provider_models( - model_types=[ModelType.LLM], - provider_schema=provider_schema, - model_setting_map={}, - ) - - assert any(model.model == "llm-model" for model in models) - assert all(model.model != "embed-model" for model in models) - - -def test_get_custom_provider_models_skips_custom_models_on_schema_error_or_none( - caplog: pytest.LogCaptureFixture, -) -> None: - configuration = _build_provider_configuration() - configuration.custom_configuration.models = [ - CustomModelConfiguration(model="error-custom", model_type=ModelType.LLM, credentials={"k": "v"}), - CustomModelConfiguration(model="none-custom", model_type=ModelType.LLM, credentials={"k": "v"}), - CustomModelConfiguration(model="ok-custom", model_type=ModelType.LLM, credentials={"k": "v"}), - ] - provider_schema = ProviderEntity( - provider="openai", - label=I18nObject(en_US="OpenAI"), - supported_model_types=[ModelType.LLM], - configurate_methods=[ConfigurateMethod.PREDEFINED_MODEL], - models=[], - ) - - def _schema(*, model_type: ModelType, model: str, credentials: dict | None): - if model == "error-custom": - raise RuntimeError("boom") - if model == "none-custom": - return None - return _build_ai_model(model) - - with caplog.at_level(logging.WARNING, logger="core.entities.provider_configuration"): - with patch.object(ProviderConfiguration, "get_model_schema", side_effect=_schema): - models = configuration._get_custom_provider_models( - model_types=[ModelType.LLM], - provider_schema=provider_schema, - model_setting_map={}, - ) - - assert "get custom model schema failed, boom" in caplog.messages - assert any(model.model == "ok-custom" for model in models) - assert all(model.model != "none-custom" for model in models) - - -@pytest.fixture -def sqlite_provider_session( - sqlite_session: Session, - sqlite_engine: Engine, -) -> Iterator[Session]: - """Bind provider-owned sessions to the same isolated SQLite database as the test.""" - with patch.object(type(db), "engine", new_callable=PropertyMock, return_value=sqlite_engine): - yield sqlite_session - - -def _provider_credential( - session: Session, - *, - name: str = "API KEY 1", - tenant_id: str = "tenant-1", - provider_name: str = "openai", - encrypted_config: str = "{}", -) -> ProviderCredential: - record = ProviderCredential( - tenant_id=tenant_id, - provider_name=provider_name, - credential_name=name, - encrypted_config=encrypted_config, - ) - session.add(record) - session.commit() - return record - - -def _provider_record( - session: Session, - *, - credential_id: str | None = None, - tenant_id: str = "tenant-1", - provider_name: str = "openai", -) -> Provider: - record = Provider( - tenant_id=tenant_id, - provider_name=provider_name, - provider_type=ProviderType.CUSTOM, - credential_id=credential_id, - is_valid=True, - ) - session.add(record) - session.commit() - return record - - -def _model_credential( - session: Session, - *, - name: str = "API KEY 1", - tenant_id: str = "tenant-1", - provider_name: str = "openai", - model: str = "gpt-4o", - encrypted_config: str = "{}", -) -> ProviderModelCredential: - record = ProviderModelCredential( - tenant_id=tenant_id, - provider_name=provider_name, - model_name=model, - model_type=ModelType.LLM, - credential_name=name, - encrypted_config=encrypted_config, - ) - session.add(record) - session.commit() - return record - - -def _provider_model_record( - session: Session, - *, - credential_id: str | None = None, - tenant_id: str = "tenant-1", - provider_name: str = "openai", - model: str = "gpt-4o", -) -> ProviderModel: - record = ProviderModel( - tenant_id=tenant_id, - provider_name=provider_name, - model_name=model, - model_type=ModelType.LLM, - credential_id=credential_id, - is_valid=True, - ) - session.add(record) - session.commit() - return record - - -def _load_balancing_config( - session: Session, - *, - credential_id: str, - source: CredentialSourceType, - name: str = "Old", -) -> LoadBalancingModelConfig: - record = LoadBalancingModelConfig( - tenant_id="tenant-1", - provider_name="openai", - model_name="gpt-4o", - model_type=ModelType.LLM, - name=name, - encrypted_config="{}", - credential_id=credential_id, - credential_source_type=source, - ) - session.add(record) - session.commit() - return record - - -@contextmanager -def _raise_on_sql(engine: Engine, table_name: str, operation: str) -> Iterator[None]: - """Fail one table operation while production still owns a real transaction.""" - - def fail_target(_conn, _cursor, statement, _parameters, _context, _executemany): - if statement.lstrip().upper().startswith(operation) and table_name in statement: - raise RuntimeError(f"forced {operation} failure for {table_name}") - - event.listen(engine, "before_cursor_execute", fail_target) - try: - yield - finally: - event.remove(engine, "before_cursor_execute", fail_target) - - -@contextmanager -def _mock_cache_boundaries() -> Iterator[tuple[Mock, Mock]]: - with ( - patch("core.entities.provider_configuration.ProviderCredentialsCache") as credentials_cache, - patch.object(ProviderConfiguration, "_invalidate_provider_configuration_cache") as configuration_cache, - ): - yield credentials_cache, configuration_cache - - -def test_generate_credential_names_from_real_rows_and_tenant_isolation( - sqlite_provider_session: Session, -) -> None: - configuration = _build_provider_configuration() - _provider_credential(sqlite_provider_session, name="API KEY 9") - _provider_credential(sqlite_provider_session, name="legacy") - _provider_credential(sqlite_provider_session, name="API KEY 50", tenant_id="other-tenant") - _model_credential(sqlite_provider_session, name="API KEY 4") - assert configuration._generate_provider_credential_name(sqlite_provider_session) == "API KEY 10" - assert ( - configuration._generate_custom_model_credential_name("gpt-4o", ModelType.LLM, sqlite_provider_session) - == "API KEY 5" - ) - - -def test_validate_provider_credentials_reuses_hidden_secret(sqlite_provider_session: Session) -> None: - configuration = _build_provider_configuration() - configuration.provider.provider_credential_schema = _build_secret_provider_schema() - credential = _provider_credential(sqlite_provider_session, encrypted_config='{"openai_api_key":"enc-old"}') - assembly = PluginModelAssembly(tenant_id="tenant-1") - with ( - patch("core.entities.provider_configuration.create_plugin_model_assembly", return_value=assembly), - patch.object(assembly.model_provider_factory, "get_model_provider", return_value=configuration.provider), - patch.object(assembly.model_runtime, "validate_provider_credentials", new=lambda **_kwargs: None), - patch("core.entities.provider_configuration.encrypter.decrypt_token", return_value="raw"), - patch("core.entities.provider_configuration.encrypter.encrypt_token", return_value="enc-new"), - ): - result = configuration.validate_provider_credentials( - {"openai_api_key": HIDDEN_VALUE}, credential_id=credential.id - ) - assert result == {"openai_api_key": "enc-new"} - - -def test_preferred_provider_state_updates_and_is_tenant_scoped(sqlite_provider_session: Session) -> None: - configuration = _build_provider_configuration() - configuration.preferred_provider_type = ProviderType.CUSTOM - other = TenantPreferredModelProvider( - tenant_id="other-tenant", provider_name="openai", preferred_provider_type=ProviderType.CUSTOM - ) - current = TenantPreferredModelProvider( - tenant_id="tenant-1", provider_name="openai", preferred_provider_type=ProviderType.CUSTOM - ) - sqlite_provider_session.add_all([other, current]) - sqlite_provider_session.commit() - assert configuration.switch_preferred_provider_type(ProviderType.SYSTEM, session=sqlite_provider_session) - sqlite_provider_session.refresh(current) - sqlite_provider_session.refresh(other) - assert current.preferred_provider_type == ProviderType.SYSTEM - assert other.preferred_provider_type == ProviderType.CUSTOM - - -def test_provider_record_duplicate_and_setting_helpers_use_real_session(sqlite_provider_session: Session) -> None: - configuration = _build_provider_configuration() - provider = _provider_record(sqlite_provider_session) - _provider_record(sqlite_provider_session, tenant_id="other-tenant") - credential = _provider_credential(sqlite_provider_session, name="Main") - _provider_credential(sqlite_provider_session, name="Main", tenant_id="other-tenant") - setting = ProviderModelSetting( - tenant_id="tenant-1", - provider_name="openai", - model_name="gpt-4o", - model_type=ModelType.LLM, - ) - sqlite_provider_session.add(setting) - sqlite_provider_session.commit() - assert configuration._get_provider_record(sqlite_provider_session).id == provider.id - assert configuration._check_provider_credential_name_exists("Main", sqlite_provider_session) - assert not configuration._check_provider_credential_name_exists( - "Main", sqlite_provider_session, exclude_id=credential.id - ) - assert configuration._get_provider_model_setting(ModelType.LLM, "gpt-4o", sqlite_provider_session).id == setting.id - - -def test_create_provider_credential_persists_provider_and_rejects_duplicate( - sqlite_provider_session: Session, -) -> None: - configuration = _build_provider_configuration() - with ( - patch.object(ProviderConfiguration, "validate_provider_credentials", return_value={"api_key": "enc"}), - _mock_cache_boundaries() as (credentials_cache, configuration_cache), - ): - created_id = configuration.create_provider_credential({"api_key": "raw"}, "Main") - credential = sqlite_provider_session.scalar( - select(ProviderCredential).where(ProviderCredential.credential_name == "Main") - ) - provider = sqlite_provider_session.scalar(select(Provider).where(Provider.tenant_id == "tenant-1")) - assert credential is not None - assert provider is not None - assert created_id == credential.id - assert provider.credential_id == credential.id - credentials_cache.assert_called_once_with( - tenant_id="tenant-1", - identity_id=provider.id, - cache_type=ProviderCredentialsCacheType.PROVIDER, - ) - credentials_cache.return_value.delete.assert_called_once_with() - configuration_cache.assert_called_once_with( - preferred_model_providers=True, - provider_credentials=True, - ) - with pytest.raises(ValueError, match="already exists"): - configuration.create_provider_credential({"api_key": "raw"}, "Main") - - -def test_update_provider_credential_propagates_to_load_balancing(sqlite_provider_session: Session) -> None: - configuration = _build_provider_configuration() - credential = _provider_credential(sqlite_provider_session, name="Old") - provider = _provider_record(sqlite_provider_session, credential_id=credential.id) - lb_config = _load_balancing_config( - sqlite_provider_session, credential_id=credential.id, source=CredentialSourceType.PROVIDER - ) - with ( - patch.object(ProviderConfiguration, "validate_provider_credentials", return_value={"api_key": "enc-new"}), - _mock_cache_boundaries() as (credentials_cache, configuration_cache), - ): - configuration.update_provider_credential({"api_key": "raw"}, credential.id, "New") - sqlite_provider_session.expire_all() - persisted_credential = sqlite_provider_session.get(ProviderCredential, credential.id) - persisted_lb = sqlite_provider_session.get(LoadBalancingModelConfig, lb_config.id) - assert persisted_credential is not None - assert persisted_credential.credential_name == "New" - assert persisted_lb is not None - assert persisted_lb.name == "New" - assert json.loads(persisted_lb.encrypted_config) == {"api_key": "enc-new"} - assert {cache_call.kwargs["identity_id"] for cache_call in credentials_cache.call_args_list} == { - provider.id, - lb_config.id, - } - assert {cache_call.kwargs["cache_type"] for cache_call in credentials_cache.call_args_list} == { - ProviderCredentialsCacheType.PROVIDER, - ProviderCredentialsCacheType.LOAD_BALANCING_MODEL, - } - assert credentials_cache.return_value.delete.call_count == 2 - configuration_cache.assert_called_once_with( - provider_credentials=True, - provider_load_balancing_configs=True, - ) - - -def test_switch_active_provider_credential_updates_persisted_state_and_cache( - sqlite_provider_session: Session, -) -> None: - configuration = _build_provider_configuration() - configuration.preferred_provider_type = ProviderType.CUSTOM - first = _provider_credential(sqlite_provider_session, name="First") - second = _provider_credential(sqlite_provider_session, name="Second") - provider = _provider_record(sqlite_provider_session, credential_id=first.id) - provider_id = provider.id - - with _mock_cache_boundaries() as (credentials_cache, configuration_cache): - configuration.switch_active_provider_credential(second.id) - - sqlite_provider_session.expire_all() - persisted_provider = sqlite_provider_session.get(Provider, provider_id) - assert persisted_provider is not None - assert persisted_provider.credential_id == second.id - credentials_cache.assert_called_once_with( - tenant_id="tenant-1", - identity_id=provider_id, - cache_type=ProviderCredentialsCacheType.PROVIDER, - ) - credentials_cache.return_value.delete.assert_called_once_with() - configuration_cache.assert_not_called() - - -def test_deleting_active_provider_credential_switches_preference_to_system( - sqlite_provider_session: Session, -) -> None: - configuration = _build_provider_configuration() - configuration.preferred_provider_type = ProviderType.CUSTOM - first = _provider_credential(sqlite_provider_session, name="First") - active = _provider_credential(sqlite_provider_session, name="Active") - provider = _provider_record(sqlite_provider_session, credential_id=active.id) - preferred_provider = TenantPreferredModelProvider( - tenant_id="tenant-1", - provider_name="openai", - preferred_provider_type=ProviderType.CUSTOM, - ) - sqlite_provider_session.add(preferred_provider) - sqlite_provider_session.commit() - first_id = first.id - active_id = active.id - provider_id = provider.id - preferred_provider_id = preferred_provider.id - - with _mock_cache_boundaries() as (credentials_cache, configuration_cache): - configuration.delete_provider_credential(active_id) - - sqlite_provider_session.expire_all() - assert sqlite_provider_session.get(ProviderCredential, active_id) is None - assert sqlite_provider_session.get(ProviderCredential, first_id) is not None - persisted_provider = sqlite_provider_session.get(Provider, provider_id) - persisted_preference = sqlite_provider_session.get(TenantPreferredModelProvider, preferred_provider_id) - assert persisted_provider is not None - assert persisted_provider.credential_id is None - assert persisted_preference is not None - assert persisted_preference.preferred_provider_type == ProviderType.SYSTEM - credentials_cache.assert_called_once_with( - tenant_id="tenant-1", - identity_id=provider_id, - cache_type=ProviderCredentialsCacheType.PROVIDER, - ) - credentials_cache.return_value.delete.assert_called_once_with() - configuration_cache.assert_called_once_with( - preferred_model_providers=True, - provider_credentials=True, - provider_load_balancing_configs=False, - ) - - -def test_specific_provider_credential_decrypts_and_obfuscates(sqlite_provider_session: Session) -> None: - configuration = _build_provider_configuration() - configuration.provider.provider_credential_schema = _build_secret_provider_schema() - credential = _provider_credential(sqlite_provider_session, encrypted_config='{"openai_api_key":"enc"}') - with ( - patch("core.entities.provider_configuration.encrypter.decrypt_token", return_value="raw"), - patch("core.entities.provider_configuration.encrypter.obfuscated_token", return_value="masked"), - ): - result = configuration._get_specific_provider_credential(credential.id) - assert result == {"openai_api_key": "masked"} - with pytest.raises(ValueError, match="not found"): - configuration._get_specific_provider_credential("missing") - - -def test_validate_custom_model_credentials_reuses_hidden_secret(sqlite_provider_session: Session) -> None: - configuration = _build_provider_configuration() - configuration.provider.model_credential_schema = _build_secret_model_schema() - credential = _model_credential(sqlite_provider_session, encrypted_config='{"openai_api_key":"enc-old"}') - assembly = PluginModelAssembly(tenant_id="tenant-1") - with ( - patch("core.entities.provider_configuration.create_plugin_model_assembly", return_value=assembly), - patch.object(assembly.model_provider_factory, "get_model_provider", return_value=configuration.provider), - patch.object(assembly.model_runtime, "validate_model_credentials", new=lambda **_kwargs: None), - patch("core.entities.provider_configuration.encrypter.decrypt_token", return_value="raw"), - patch("core.entities.provider_configuration.encrypter.encrypt_token", return_value="enc-new"), - ): - result = configuration.validate_custom_model_credentials( - ModelType.LLM, - "gpt-4o", - {"openai_api_key": HIDDEN_VALUE}, - credential_id=credential.id, - ) - assert result == {"openai_api_key": "enc-new"} - - -def test_specific_custom_model_credential_preserves_secret_when_decryption_fails( - sqlite_provider_session: Session, - caplog: pytest.LogCaptureFixture, -) -> None: - configuration = _build_provider_configuration() - configuration.provider.model_credential_schema = _build_secret_model_schema() - credential = _model_credential( - sqlite_provider_session, - name="Main", - encrypted_config='{"openai_api_key":"enc-secret"}', - ) - - with ( - caplog.at_level(logging.ERROR, logger="core.entities.provider_configuration"), - patch("core.entities.provider_configuration.encrypter.decrypt_token", side_effect=RuntimeError("boom")), - patch.object( - ProviderConfiguration, - "obfuscated_credentials", - side_effect=lambda credentials, credential_form_schemas: credentials, - ), - ): - result = configuration._get_specific_custom_model_credential(ModelType.LLM, "gpt-4o", credential.id) - - assert result == { - "current_credential_id": credential.id, - "current_credential_name": "Main", - "credentials": {"openai_api_key": "enc-secret"}, - } - assert caplog.messages.count("Failed to decrypt model credential secret variable openai_api_key") == 1 - - -def test_create_update_and_delete_custom_model_credential(sqlite_provider_session: Session) -> None: - configuration = _build_provider_configuration() - with ( - patch.object(ProviderConfiguration, "validate_custom_model_credentials", return_value={"api_key": "enc"}), - _mock_cache_boundaries() as (credentials_cache, configuration_cache), - ): - created_id = configuration.create_custom_model_credential(ModelType.LLM, "gpt-4o", {"api_key": "raw"}, "Main") - credential = sqlite_provider_session.scalar(select(ProviderModelCredential)) - model = sqlite_provider_session.scalar(select(ProviderModel)) - assert credential is not None - assert model is not None - assert created_id == credential.id - assert model.credential_id == credential.id - credential_id = credential.id - model_id = model.id - lb_config = _load_balancing_config( - sqlite_provider_session, - credential_id=credential_id, - source=CredentialSourceType.CUSTOM_MODEL, - ) - lb_config_id = lb_config.id - credentials_cache.assert_called_once_with( - tenant_id="tenant-1", - identity_id=model_id, - cache_type=ProviderCredentialsCacheType.MODEL, - ) - credentials_cache.return_value.delete.assert_called_once_with() - configuration_cache.assert_called_once_with( - provider_models=True, - provider_model_credentials=True, - ) - - with ( - patch.object(ProviderConfiguration, "validate_custom_model_credentials", return_value={"api_key": "enc-2"}), - _mock_cache_boundaries() as (credentials_cache, configuration_cache), - ): - configuration.update_custom_model_credential( - ModelType.LLM, "gpt-4o", {"api_key": "raw"}, "Renamed", credential_id - ) - sqlite_provider_session.expire_all() - persisted_credential = sqlite_provider_session.get(ProviderModelCredential, credential_id) - persisted_lb = sqlite_provider_session.get(LoadBalancingModelConfig, lb_config_id) - assert persisted_credential is not None - assert persisted_credential.credential_name == "Renamed" - assert json.loads(persisted_credential.encrypted_config) == {"api_key": "enc-2"} - assert persisted_lb is not None - assert persisted_lb.name == "Renamed" - assert json.loads(persisted_lb.encrypted_config) == {"api_key": "enc-2"} - assert {cache_call.kwargs["identity_id"] for cache_call in credentials_cache.call_args_list} == { - model_id, - lb_config_id, - } - assert {cache_call.kwargs["cache_type"] for cache_call in credentials_cache.call_args_list} == { - ProviderCredentialsCacheType.MODEL, - ProviderCredentialsCacheType.LOAD_BALANCING_MODEL, - } - assert credentials_cache.return_value.delete.call_count == 2 - configuration_cache.assert_called_once_with( - provider_models=True, - provider_model_credentials=True, - provider_load_balancing_configs=True, - ) - - with _mock_cache_boundaries() as (credentials_cache, configuration_cache): - configuration.delete_custom_model_credential(ModelType.LLM, "gpt-4o", credential_id) - sqlite_provider_session.expire_all() - assert sqlite_provider_session.get(ProviderModelCredential, credential_id) is None - assert sqlite_provider_session.get(ProviderModel, model_id) is None - assert sqlite_provider_session.get(LoadBalancingModelConfig, lb_config_id) is None - assert {cache_call.kwargs["identity_id"] for cache_call in credentials_cache.call_args_list} == { - model_id, - lb_config_id, - } - assert {cache_call.kwargs["cache_type"] for cache_call in credentials_cache.call_args_list} == { - ProviderCredentialsCacheType.MODEL, - ProviderCredentialsCacheType.LOAD_BALANCING_MODEL, - } - assert credentials_cache.return_value.delete.call_count == 2 - configuration_cache.assert_called_once_with( - provider_models=True, - provider_model_credentials=True, - provider_load_balancing_configs=True, - ) - - -def test_add_and_switch_custom_model_credential(sqlite_provider_session: Session) -> None: - configuration = _build_provider_configuration() - first = _model_credential(sqlite_provider_session, name="First") - second = _model_credential(sqlite_provider_session, name="Second") - with _mock_cache_boundaries() as (credentials_cache, configuration_cache): - configuration.add_model_credential_to_model(ModelType.LLM, "gpt-4o", first.id) - configuration.switch_custom_model_credential(ModelType.LLM, "gpt-4o", second.id) - model = sqlite_provider_session.scalar(select(ProviderModel)) - assert model is not None - assert model.credential_id == second.id - credentials_cache.assert_called_once_with( - tenant_id="tenant-1", - identity_id=model.id, - cache_type=ProviderCredentialsCacheType.MODEL, - ) - credentials_cache.return_value.delete.assert_called_once_with() - assert configuration_cache.call_count == 2 - for cache_call in configuration_cache.call_args_list: - assert cache_call.kwargs == {"provider_models": True} - with pytest.raises(ValueError, match="Can't add same credential"): - configuration.add_model_credential_to_model(ModelType.LLM, "gpt-4o", second.id) - - -def test_model_settings_and_load_balancing_persist(sqlite_provider_session: Session) -> None: - configuration = _build_provider_configuration() - with patch.object(configuration, "_invalidate_provider_configuration_cache") as configuration_cache: - configuration.disable_model(ModelType.LLM, "gpt-4o") - configuration_cache.assert_called_once_with(provider_model_settings=True) - persisted_setting = sqlite_provider_session.scalar(select(ProviderModelSetting)) - assert persisted_setting is not None - assert persisted_setting.enabled is False - with patch.object(configuration, "_invalidate_provider_configuration_cache") as configuration_cache: - configuration.enable_model(ModelType.LLM, "gpt-4o") - configuration_cache.assert_called_once_with(provider_model_settings=True) - sqlite_provider_session.expire_all() - refreshed_setting = sqlite_provider_session.get(ProviderModelSetting, persisted_setting.id) - assert refreshed_setting is not None - assert refreshed_setting.enabled is True - - first = _provider_credential(sqlite_provider_session, name="First") - second = _provider_credential(sqlite_provider_session, name="Second") - _load_balancing_config(sqlite_provider_session, credential_id=first.id, source=CredentialSourceType.PROVIDER) - _load_balancing_config(sqlite_provider_session, credential_id=second.id, source=CredentialSourceType.PROVIDER) - with patch.object(configuration, "_invalidate_provider_configuration_cache") as configuration_cache: - configuration.enable_model_load_balancing(ModelType.LLM, "gpt-4o") - configuration_cache.assert_called_once_with(provider_model_settings=True) - sqlite_provider_session.expire_all() - refreshed_setting = sqlite_provider_session.get(ProviderModelSetting, persisted_setting.id) - assert refreshed_setting is not None - assert refreshed_setting.load_balancing_enabled is True - with patch.object(configuration, "_invalidate_provider_configuration_cache") as configuration_cache: - configuration.disable_model_load_balancing(ModelType.LLM, "gpt-4o") - configuration_cache.assert_called_once_with(provider_model_settings=True) - sqlite_provider_session.expire_all() - refreshed_setting = sqlite_provider_session.get(ProviderModelSetting, persisted_setting.id) - assert refreshed_setting is not None - assert refreshed_setting.load_balancing_enabled is False - - -def test_provider_create_rolls_back_on_insert_failure( - sqlite_provider_session: Session, - sqlite_engine: Engine, -) -> None: - configuration = _build_provider_configuration() - with ( - patch.object(ProviderConfiguration, "validate_provider_credentials", return_value={"api_key": "enc"}), - _mock_cache_boundaries(), - _raise_on_sql(sqlite_engine, "provider_credentials", "INSERT"), - pytest.raises(RuntimeError, match="forced INSERT"), - ): - configuration.create_provider_credential({"api_key": "raw"}, "Main") - assert sqlite_provider_session.scalar(select(ProviderCredential)) is None - assert sqlite_provider_session.scalar(select(Provider)) is None - - -def test_custom_model_create_rolls_back_on_insert_failure( - sqlite_provider_session: Session, - sqlite_engine: Engine, -) -> None: - configuration = _build_provider_configuration() - with ( - patch.object(ProviderConfiguration, "validate_custom_model_credentials", return_value={"api_key": "enc"}), - _mock_cache_boundaries(), - _raise_on_sql(sqlite_engine, "provider_model_credentials", "INSERT"), - pytest.raises(RuntimeError, match="forced INSERT"), - ): - configuration.create_custom_model_credential(ModelType.LLM, "gpt-4o", {"api_key": "raw"}, "Main") - assert sqlite_provider_session.scalar(select(ProviderModelCredential)) is None - assert sqlite_provider_session.scalar(select(ProviderModel)) is None - - -def test_provider_update_rolls_back_on_update_failure( - sqlite_provider_session: Session, - sqlite_engine: Engine, -) -> None: - configuration = _build_provider_configuration() - credential = _provider_credential( - sqlite_provider_session, - name="Old", - encrypted_config='{"api_key":"enc-old"}', - ) - _provider_record(sqlite_provider_session, credential_id=credential.id) - - with ( - patch.object(ProviderConfiguration, "validate_provider_credentials", return_value={"api_key": "enc-new"}), - _mock_cache_boundaries() as (credentials_cache, configuration_cache), - _raise_on_sql(sqlite_engine, "provider_credentials", "UPDATE"), - pytest.raises(RuntimeError, match="forced UPDATE"), - ): - configuration.update_provider_credential({"api_key": "raw"}, credential.id, "New") - - sqlite_provider_session.expire_all() - persisted = sqlite_provider_session.get(ProviderCredential, credential.id) - assert persisted is not None - assert persisted.credential_name == "Old" - assert json.loads(persisted.encrypted_config) == {"api_key": "enc-old"} - credentials_cache.assert_not_called() - configuration_cache.assert_not_called() - - -def test_provider_delete_rolls_back_on_delete_failure( - sqlite_provider_session: Session, - sqlite_engine: Engine, -) -> None: - configuration = _build_provider_configuration() - credential = _provider_credential(sqlite_provider_session) - provider = _provider_record(sqlite_provider_session, credential_id=credential.id) - - with ( - _mock_cache_boundaries() as (credentials_cache, configuration_cache), - _raise_on_sql(sqlite_engine, "provider_credentials", "DELETE"), - pytest.raises(RuntimeError, match="forced DELETE"), - ): - configuration.delete_provider_credential(credential.id) - - sqlite_provider_session.expire_all() - assert sqlite_provider_session.get(ProviderCredential, credential.id) is not None - persisted_provider = sqlite_provider_session.get(Provider, provider.id) - assert persisted_provider is not None - assert persisted_provider.credential_id == credential.id - credentials_cache.assert_called_once_with( - tenant_id="tenant-1", - identity_id=provider.id, - cache_type=ProviderCredentialsCacheType.PROVIDER, - ) - credentials_cache.return_value.delete.assert_called_once_with() - configuration_cache.assert_not_called() - - -def test_provider_switch_rolls_back_on_update_failure( - sqlite_provider_session: Session, - sqlite_engine: Engine, -) -> None: - configuration = _build_provider_configuration() - first = _provider_credential(sqlite_provider_session, name="First") - second = _provider_credential(sqlite_provider_session, name="Second") - provider = _provider_record(sqlite_provider_session, credential_id=first.id) - - with ( - _mock_cache_boundaries() as (credentials_cache, configuration_cache), - _raise_on_sql(sqlite_engine, "providers", "UPDATE"), - pytest.raises(RuntimeError, match="forced UPDATE"), - ): - configuration.switch_active_provider_credential(second.id) - - sqlite_provider_session.expire_all() - persisted_provider = sqlite_provider_session.get(Provider, provider.id) - assert persisted_provider is not None - assert persisted_provider.credential_id == first.id - credentials_cache.assert_not_called() - configuration_cache.assert_not_called() - - -def test_custom_model_update_rolls_back_on_update_failure( - sqlite_provider_session: Session, - sqlite_engine: Engine, -) -> None: - configuration = _build_provider_configuration() - credential = _model_credential( - sqlite_provider_session, - name="Old", - encrypted_config='{"api_key":"enc-old"}', - ) - _provider_model_record(sqlite_provider_session, credential_id=credential.id) - - with ( - patch.object( - ProviderConfiguration, - "validate_custom_model_credentials", - return_value={"api_key": "enc-new"}, - ), - _mock_cache_boundaries() as (credentials_cache, configuration_cache), - _raise_on_sql(sqlite_engine, "provider_model_credentials", "UPDATE"), - pytest.raises(RuntimeError, match="forced UPDATE"), - ): - configuration.update_custom_model_credential( - ModelType.LLM, - "gpt-4o", - {"api_key": "raw"}, - "New", - credential.id, - ) - - sqlite_provider_session.expire_all() - persisted = sqlite_provider_session.get(ProviderModelCredential, credential.id) - assert persisted is not None - assert persisted.credential_name == "Old" - assert json.loads(persisted.encrypted_config) == {"api_key": "enc-old"} - credentials_cache.assert_not_called() - configuration_cache.assert_not_called() - - -def test_custom_model_delete_rolls_back_on_delete_failure( - sqlite_provider_session: Session, - sqlite_engine: Engine, -) -> None: - configuration = _build_provider_configuration() - credential = _model_credential(sqlite_provider_session) - model = _provider_model_record(sqlite_provider_session, credential_id=credential.id) - - with ( - _mock_cache_boundaries() as (credentials_cache, configuration_cache), - _raise_on_sql(sqlite_engine, "provider_model_credentials", "DELETE"), - pytest.raises(RuntimeError, match="forced DELETE"), - ): - configuration.delete_custom_model_credential(ModelType.LLM, "gpt-4o", credential.id) - - sqlite_provider_session.expire_all() - assert sqlite_provider_session.get(ProviderModelCredential, credential.id) is not None - persisted_model = sqlite_provider_session.get(ProviderModel, model.id) - assert persisted_model is not None - assert persisted_model.credential_id == credential.id - credentials_cache.assert_called_once_with( - tenant_id="tenant-1", - identity_id=model.id, - cache_type=ProviderCredentialsCacheType.MODEL, - ) - credentials_cache.return_value.delete.assert_called_once_with() - configuration_cache.assert_not_called() - - -def test_custom_model_switch_rolls_back_on_update_failure( - sqlite_provider_session: Session, - sqlite_engine: Engine, -) -> None: - configuration = _build_provider_configuration() - first = _model_credential(sqlite_provider_session, name="First") - second = _model_credential(sqlite_provider_session, name="Second") - model = _provider_model_record(sqlite_provider_session, credential_id=first.id) - - with ( - _mock_cache_boundaries() as (credentials_cache, configuration_cache), - _raise_on_sql(sqlite_engine, "provider_models", "UPDATE"), - pytest.raises(RuntimeError, match="forced UPDATE"), - ): - configuration.switch_custom_model_credential(ModelType.LLM, "gpt-4o", second.id) - - sqlite_provider_session.expire_all() - persisted_model = sqlite_provider_session.get(ProviderModel, model.id) - assert persisted_model is not None - assert persisted_model.credential_id == first.id - credentials_cache.assert_not_called() - configuration_cache.assert_not_called() diff --git a/api/tests/unit_tests/core/entities/test_entities_execution_extra_content.py b/api/tests/unit_tests/core/entities/test_execution_extra_content.py similarity index 100% rename from api/tests/unit_tests/core/entities/test_entities_execution_extra_content.py rename to api/tests/unit_tests/core/entities/test_execution_extra_content.py diff --git a/api/tests/unit_tests/core/entities/test_entities_knowledge_entities.py b/api/tests/unit_tests/core/entities/test_knowledge_entities.py similarity index 100% rename from api/tests/unit_tests/core/entities/test_entities_knowledge_entities.py rename to api/tests/unit_tests/core/entities/test_knowledge_entities.py diff --git a/api/tests/unit_tests/core/entities/test_entities_mcp_provider.py b/api/tests/unit_tests/core/entities/test_mcp_provider.py similarity index 99% rename from api/tests/unit_tests/core/entities/test_entities_mcp_provider.py rename to api/tests/unit_tests/core/entities/test_mcp_provider.py index ae6db615395..16db07c2f7f 100644 --- a/api/tests/unit_tests/core/entities/test_entities_mcp_provider.py +++ b/api/tests/unit_tests/core/entities/test_mcp_provider.py @@ -68,7 +68,7 @@ def test_from_db_model_maps_fields() -> None: def test_redirect_url_uses_console_api_url(monkeypatch: pytest.MonkeyPatch) -> None: # Arrange entity = _build_mcp_provider_entity() - from tests.unit_tests.config_override import apply_config_overrides + from tests.helpers.unit.config_override import apply_config_overrides apply_config_overrides(monkeypatch, CONSOLE_API_URL="https://console.example.com") diff --git a/api/tests/unit_tests/core/entities/test_entities_model_entities.py b/api/tests/unit_tests/core/entities/test_model_entities.py similarity index 100% rename from api/tests/unit_tests/core/entities/test_entities_model_entities.py rename to api/tests/unit_tests/core/entities/test_model_entities.py diff --git a/api/tests/unit_tests/core/entities/test_entities_parameter_entities.py b/api/tests/unit_tests/core/entities/test_parameter_entities.py similarity index 100% rename from api/tests/unit_tests/core/entities/test_entities_parameter_entities.py rename to api/tests/unit_tests/core/entities/test_parameter_entities.py diff --git a/api/tests/unit_tests/core/entities/test_provider_configuration.py b/api/tests/unit_tests/core/entities/test_provider_configuration.py new file mode 100644 index 00000000000..b3bc29f8d77 --- /dev/null +++ b/api/tests/unit_tests/core/entities/test_provider_configuration.py @@ -0,0 +1,2026 @@ +from __future__ import annotations + +import json +import logging +from collections.abc import Iterator +from contextlib import contextmanager +from types import SimpleNamespace +from unittest.mock import MagicMock, Mock, PropertyMock, patch + +import pytest +from sqlalchemy import Engine, event, select +from sqlalchemy.orm import Session + +from constants import HIDDEN_VALUE +from core.entities import provider_configuration as provider_configuration_module +from core.entities.model_entities import ModelStatus, ProviderModelWithStatusEntity +from core.entities.provider_configuration import ProviderConfiguration, ProviderConfigurations +from core.entities.provider_configuration import ( + SystemConfigurationStatus as _provider_configuration_SystemConfigurationStatus, +) +from core.entities.provider_entities import ( + CredentialConfiguration, + CustomConfiguration, + CustomModelConfiguration, + CustomProviderConfiguration, + ModelLoadBalancingConfiguration, + ModelSettings, + ProviderQuotaType, + QuotaConfiguration, + QuotaUnit, + RestrictModel, + SystemConfiguration, +) +from core.entities.provider_entities import ( + SystemConfigurationStatus as _entities_provider_configuration_SystemConfigurationStatus, +) +from core.helper.model_provider_cache import ProviderCredentialsCacheType +from core.plugin.impl.model_runtime_factory import PluginModelAssembly, create_model_type_instance +from extensions.ext_database import db +from graphon.model_runtime.entities.common_entities import I18nObject +from graphon.model_runtime.entities.model_entities import AIModelEntity, FetchFrom, ModelType +from graphon.model_runtime.entities.provider_entities import ( + ConfigurateMethod, + CredentialFormSchema, + FieldModelSchema, + FormOption, + FormType, + ModelCredentialSchema, + ProviderCredentialSchema, + ProviderEntity, +) +from graphon.model_runtime.model_providers.base.large_language_model import LargeLanguageModel +from models.enums import CredentialSourceType +from models.provider import ( + LoadBalancingModelConfig, + Provider, + ProviderCredential, + ProviderModel, + ProviderModelCredential, + ProviderModelSetting, + ProviderType, + TenantPreferredModelProvider, +) +from models.provider_ids import ModelProviderID + + +def _build_provider_configuration(*, provider_name: str = "openai") -> ProviderConfiguration: + provider_entity = ProviderEntity( + provider=provider_name, + label=I18nObject(en_US="OpenAI"), + supported_model_types=[ModelType.LLM], + configurate_methods=[ConfigurateMethod.PREDEFINED_MODEL], + ) + system_configuration = SystemConfiguration( + enabled=True, + credentials={"api_key": "test-key"}, + current_quota_type=ProviderQuotaType.TRIAL, + quota_configurations=[ + QuotaConfiguration( + quota_type=ProviderQuotaType.TRIAL, + quota_unit=QuotaUnit.TOKENS, + quota_limit=1_000, + quota_used=0, + is_valid=True, + restrict_models=[], + ) + ], + ) + with patch("core.entities.provider_configuration.original_provider_configurate_methods", {}): + return ProviderConfiguration( + tenant_id="tenant-1", + provider=provider_entity, + preferred_provider_type=ProviderType.SYSTEM, + using_provider_type=ProviderType.SYSTEM, + system_configuration=system_configuration, + custom_configuration=CustomConfiguration(provider=None, models=[]), + model_settings=[], + ) + + +def _build_ai_model(name: str, *, model_type: ModelType = ModelType.LLM) -> AIModelEntity: + return AIModelEntity( + model=name, + label=I18nObject(en_US=name), + model_type=model_type, + fetch_from=FetchFrom.PREDEFINED_MODEL, + model_properties={}, + ) + + +def _build_provider_model(name: str) -> ProviderModelWithStatusEntity: + return ProviderModelWithStatusEntity(**_build_ai_model(name).model_dump(), status=ModelStatus.ACTIVE) + + +def _build_secret_provider_schema() -> ProviderCredentialSchema: + return ProviderCredentialSchema( + credential_form_schemas=[ + CredentialFormSchema( + variable="openai_api_key", + label=I18nObject(en_US="API Key"), + type=FormType.SECRET_INPUT, + ) + ] + ) + + +def _build_secret_model_schema() -> ModelCredentialSchema: + return ModelCredentialSchema( + model=FieldModelSchema(label=I18nObject(en_US="Model")), + credential_form_schemas=[ + CredentialFormSchema( + variable="openai_api_key", + label=I18nObject(en_US="API Key"), + type=FormType.SECRET_INPUT, + ) + ], + ) + + +def _provider_credential( + session: Session, + *, + name: str = "API KEY 1", + tenant_id: str = "tenant-1", + provider_name: str = "openai", + encrypted_config: str = "{}", +) -> ProviderCredential: + record = ProviderCredential( + tenant_id=tenant_id, + provider_name=provider_name, + credential_name=name, + encrypted_config=encrypted_config, + ) + session.add(record) + session.commit() + return record + + +def _provider_record( + session: Session, + *, + credential_id: str | None = None, + tenant_id: str = "tenant-1", + provider_name: str = "openai", +) -> Provider: + record = Provider( + tenant_id=tenant_id, + provider_name=provider_name, + provider_type=ProviderType.CUSTOM, + credential_id=credential_id, + is_valid=True, + ) + session.add(record) + session.commit() + return record + + +def _model_credential( + session: Session, + *, + name: str = "API KEY 1", + tenant_id: str = "tenant-1", + provider_name: str = "openai", + model: str = "gpt-4o", + encrypted_config: str = "{}", +) -> ProviderModelCredential: + record = ProviderModelCredential( + tenant_id=tenant_id, + provider_name=provider_name, + model_name=model, + model_type=ModelType.LLM, + credential_name=name, + encrypted_config=encrypted_config, + ) + session.add(record) + session.commit() + return record + + +def _provider_model_record( + session: Session, + *, + credential_id: str | None = None, + tenant_id: str = "tenant-1", + provider_name: str = "openai", + model: str = "gpt-4o", +) -> ProviderModel: + record = ProviderModel( + tenant_id=tenant_id, + provider_name=provider_name, + model_name=model, + model_type=ModelType.LLM, + credential_id=credential_id, + is_valid=True, + ) + session.add(record) + session.commit() + return record + + +def _load_balancing_config( + session: Session, + *, + credential_id: str, + source: CredentialSourceType, + name: str = "Old", +) -> LoadBalancingModelConfig: + record = LoadBalancingModelConfig( + tenant_id="tenant-1", + provider_name="openai", + model_name="gpt-4o", + model_type=ModelType.LLM, + name=name, + encrypted_config="{}", + credential_id=credential_id, + credential_source_type=source, + ) + session.add(record) + session.commit() + return record + + +@contextmanager +def _raise_on_sql(engine: Engine, table_name: str, operation: str) -> Iterator[None]: + """Fail one table operation while production still owns a real transaction.""" + + def fail_target(_conn, _cursor, statement, _parameters, _context, _executemany): + if statement.lstrip().upper().startswith(operation) and table_name in statement: + raise RuntimeError(f"forced {operation} failure for {table_name}") + + event.listen(engine, "before_cursor_execute", fail_target) + try: + yield + finally: + event.remove(engine, "before_cursor_execute", fail_target) + + +@contextmanager +def _mock_cache_boundaries() -> Iterator[tuple[Mock, Mock]]: + with ( + patch("core.entities.provider_configuration.ProviderCredentialsCache") as credentials_cache, + patch.object(ProviderConfiguration, "_invalidate_provider_configuration_cache") as configuration_cache, + ): + yield credentials_cache, configuration_cache + + +class TestEntitiesProviderConfigurationCases: + @staticmethod + def test_extract_secret_variables_returns_only_secret_inputs() -> None: + configuration = _build_provider_configuration() + credential_form_schemas = [ + CredentialFormSchema( + variable="api_key", + label=I18nObject(en_US="API Key"), + type=FormType.SECRET_INPUT, + ), + CredentialFormSchema( + variable="endpoint", + label=I18nObject(en_US="Endpoint"), + type=FormType.TEXT_INPUT, + ), + ] + + secret_variables = configuration.extract_secret_variables(credential_form_schemas) + assert secret_variables == ["api_key"] + + @staticmethod + def test_obfuscated_credentials_masks_only_secret_fields() -> None: + configuration = _build_provider_configuration() + credential_form_schemas = [ + CredentialFormSchema( + variable="api_key", + label=I18nObject(en_US="API Key"), + type=FormType.SECRET_INPUT, + ), + CredentialFormSchema( + variable="endpoint", + label=I18nObject(en_US="Endpoint"), + type=FormType.TEXT_INPUT, + ), + ] + + with patch( + "core.entities.provider_configuration.encrypter.obfuscated_token", + side_effect=lambda value: f"masked-{value[-2:]}", + ): + obfuscated = configuration.obfuscated_credentials( + credentials={"api_key": "sk-test-1234", "endpoint": "https://api.example.com"}, + credential_form_schemas=credential_form_schemas, + ) + + assert obfuscated["api_key"] == "masked-34" + assert obfuscated["endpoint"] == "https://api.example.com" + + @staticmethod + def test_provider_configurations_behave_like_keyed_container() -> None: + configuration = _build_provider_configuration() + provider_key = str(ModelProviderID("openai")) + configurations = ProviderConfigurations(tenant_id="tenant-1") + + configurations[provider_key] = configuration + + assert "openai" in configurations + assert configurations["openai"] is configuration + assert configurations.get("openai") is configuration + assert configurations.to_list() == [configuration] + assert list(configurations) == [(provider_key, configuration)] + + @staticmethod + def test_provider_configurations_get_models_forwards_filters() -> None: + configuration = _build_provider_configuration() + provider_key = str(ModelProviderID("openai")) + configurations = ProviderConfigurations(tenant_id="tenant-1") + configurations[provider_key] = configuration + expected_model = _build_provider_model("gpt-4o") + + with patch.object(ProviderConfiguration, "get_provider_models", return_value=[expected_model]) as mock_get: + models = configurations.get_models(provider="openai", model_type=ModelType.LLM, only_active=True) + + mock_get.assert_called_once_with(ModelType.LLM, True) + assert models == [expected_model] + + @staticmethod + def test_provider_configurations_get_models_skips_non_matching_provider_filter() -> None: + configuration = _build_provider_configuration() + provider_key = str(ModelProviderID("openai")) + configurations = ProviderConfigurations(tenant_id="tenant-1") + configurations[provider_key] = configuration + + with patch.object( + ProviderConfiguration, "get_provider_models", return_value=[_build_provider_model("gpt-4o")] + ) as mock_get: + models = configurations.get_models(provider="anthropic", model_type=ModelType.LLM, only_active=True) + + assert models == [] + mock_get.assert_not_called() + + @staticmethod + def test_get_current_credentials_custom_provider_checks_current_credential() -> None: + configuration = _build_provider_configuration() + configuration.using_provider_type = ProviderType.CUSTOM + configuration.custom_configuration.provider = CustomProviderConfiguration( + credentials={"api_key": "provider-key"}, + current_credential_id="credential-1", + current_credential_name="Primary", + available_credentials=[], + ) + + with patch("core.helper.credential_utils.check_credential_policy_compliance") as mock_check: + credentials = configuration.get_current_credentials(ModelType.LLM, "gpt-4o") + + assert credentials == {"api_key": "provider-key"} + assert mock_check.call_count == 1 + assert mock_check.call_args.kwargs["credential_id"] == "credential-1" + assert mock_check.call_args.kwargs["provider"] == "openai" + + @staticmethod + def test_get_current_credentials_custom_provider_checks_all_available_credentials() -> None: + configuration = _build_provider_configuration() + configuration.using_provider_type = ProviderType.CUSTOM + configuration.custom_configuration.provider = CustomProviderConfiguration( + credentials={"api_key": "provider-key"}, + available_credentials=[ + CredentialConfiguration(credential_id="cred-1", credential_name="First"), + CredentialConfiguration(credential_id="cred-2", credential_name="Second"), + ], + ) + + with patch("core.helper.credential_utils.check_credential_policy_compliance") as mock_check: + credentials = configuration.get_current_credentials(ModelType.LLM, "gpt-4o") + + assert credentials == {"api_key": "provider-key"} + assert [c.kwargs["credential_id"] for c in mock_check.call_args_list] == ["cred-1", "cred-2"] + assert all(c.kwargs["provider"] == "openai" for c in mock_check.call_args_list) + + @staticmethod + def test_get_system_configuration_status_returns_none_when_current_quota_missing() -> None: + configuration = _build_provider_configuration() + configuration.system_configuration.current_quota_type = ProviderQuotaType.FREE + + status = configuration.get_system_configuration_status() + assert status is None + + @staticmethod + def test_get_provider_names_supports_legacy_and_full_plugin_id() -> None: + configuration = _build_provider_configuration() + configuration.provider.provider = "langgenius/openai/openai" + + provider_names = configuration._get_provider_names() + assert provider_names == ["langgenius/openai/openai", "openai"] + + @staticmethod + def test_get_model_type_instance_and_schema_delegate_to_factory() -> None: + configuration = _build_provider_configuration() + assembly = PluginModelAssembly(tenant_id="tenant-1") + schema = _build_ai_model("gpt-4o") + + with ( + patch( + "core.entities.provider_configuration.create_plugin_model_assembly", return_value=assembly + ) as builder, + patch( + "core.entities.provider_configuration.create_model_type_instance", wraps=create_model_type_instance + ) as model_builder, + patch.object(assembly.model_runtime, "get_model_schema", return_value=schema) as get_schema, + ): + model_type_instance = configuration.get_model_type_instance(ModelType.LLM) + model_schema = configuration.get_model_schema(ModelType.LLM, "gpt-4o", {"api_key": "x"}) + + assert isinstance(model_type_instance, LargeLanguageModel) + assert model_schema is schema + assert builder.call_count == 2 + model_builder.assert_called_once_with( + runtime=assembly.model_runtime, + provider_schema=configuration.provider, + model_type=ModelType.LLM, + ) + get_schema.assert_called_once_with( + provider="openai", model_type=ModelType.LLM, model="gpt-4o", credentials={"api_key": "x"} + ) + + @staticmethod + def test_get_model_type_instance_and_schema_reuse_bound_runtime_factory() -> None: + configuration = _build_provider_configuration() + bound_runtime = PluginModelAssembly(tenant_id="tenant-1").model_runtime + schema = _build_ai_model("gpt-4o") + configuration.bind_model_runtime(bound_runtime) + + with ( + patch("core.entities.provider_configuration.ModelProviderFactory") as factory_cls, + patch("core.entities.provider_configuration.create_plugin_model_assembly") as assembly_builder, + patch( + "core.entities.provider_configuration.create_model_type_instance", wraps=create_model_type_instance + ) as model_builder, + patch.object(bound_runtime, "get_model_schema", return_value=schema) as get_schema, + ): + model_type_instance = configuration.get_model_type_instance(ModelType.LLM) + model_schema = configuration.get_model_schema(ModelType.LLM, "gpt-4o", {"api_key": "x"}) + + assert isinstance(model_type_instance, LargeLanguageModel) + assert model_schema is schema + factory_cls.assert_not_called() + assembly_builder.assert_not_called() + model_builder.assert_called_once_with( + runtime=bound_runtime, + provider_schema=configuration.provider, + model_type=ModelType.LLM, + ) + get_schema.assert_called_once_with( + provider="openai", model_type=ModelType.LLM, model="gpt-4o", credentials={"api_key": "x"} + ) + + @staticmethod + def test_get_provider_model_returns_none_when_model_not_found() -> None: + configuration = _build_provider_configuration() + other_model = _build_provider_model("other-model") + + with patch.object(ProviderConfiguration, "get_provider_models", return_value=[other_model]): + selected = configuration.get_provider_model(ModelType.LLM, "gpt-4o") + + assert selected is None + + @staticmethod + def test_get_provider_models_system_deduplicates_sorts_and_filters_active() -> None: + configuration = _build_provider_configuration() + configuration.provider.position = {"llm": ["b-model", "a-model"]} + configuration.model_settings = [ + ModelSettings(model="a-model", model_type=ModelType.LLM, enabled=False, load_balancing_configs=[]) + ] + provider_schema = ProviderEntity( + provider="openai", + label=I18nObject(en_US="OpenAI"), + supported_model_types=[ModelType.LLM], + configurate_methods=[ConfigurateMethod.PREDEFINED_MODEL], + models=[_build_ai_model("a-model"), _build_ai_model("b-model"), _build_ai_model("a-model")], + ) + assembly = PluginModelAssembly(tenant_id="tenant-1") + + with ( + patch("core.entities.provider_configuration.create_plugin_model_assembly", return_value=assembly), + patch.object(assembly.model_provider_factory, "get_model_provider", return_value=provider_schema), + ): + all_models = configuration.get_provider_models(model_type=ModelType.LLM, only_active=False) + active_models = configuration.get_provider_models(model_type=ModelType.LLM, only_active=True) + + assert [model.model for model in all_models] == ["b-model", "a-model"] + assert [model.status for model in all_models] == [ModelStatus.ACTIVE, ModelStatus.DISABLED] + assert [model.model for model in active_models] == ["b-model"] + + @staticmethod + def test_get_provider_models_system_filters_requested_model() -> None: + configuration = _build_provider_configuration() + provider_schema = ProviderEntity( + provider="openai", + label=I18nObject(en_US="OpenAI"), + supported_model_types=[ModelType.LLM], + configurate_methods=[ConfigurateMethod.PREDEFINED_MODEL], + models=[_build_ai_model("a-model"), _build_ai_model("target-model"), _build_ai_model("b-model")], + ) + assembly = PluginModelAssembly(tenant_id="tenant-1") + + with ( + patch("core.entities.provider_configuration.create_plugin_model_assembly", return_value=assembly), + patch.object(assembly.model_provider_factory, "get_model_provider", return_value=provider_schema), + ): + models = configuration.get_provider_models( + model_type=ModelType.LLM, + only_active=False, + model="target-model", + ) + + assert [model.model for model in models] == ["target-model"] + + @staticmethod + def test_get_provider_models_system_customizable_filters_requested_restricted_model() -> None: + provider = ProviderEntity( + provider="openai", + label=I18nObject(en_US="OpenAI"), + supported_model_types=[ModelType.LLM], + configurate_methods=[ConfigurateMethod.CUSTOMIZABLE_MODEL], + ) + system_configuration = SystemConfiguration( + enabled=True, + credentials={"api_key": "test-key"}, + current_quota_type=ProviderQuotaType.TRIAL, + quota_configurations=[ + QuotaConfiguration( + quota_type=ProviderQuotaType.TRIAL, + quota_unit=QuotaUnit.TOKENS, + quota_limit=1_000, + quota_used=0, + is_valid=True, + restrict_models=[ + RestrictModel(model="target-model", base_model_name="base-model", model_type=ModelType.LLM), + RestrictModel(model="other-model", base_model_name="base-model", model_type=ModelType.LLM), + ], + ) + ], + ) + provider_schema = ProviderEntity( + provider="openai", + label=I18nObject(en_US="OpenAI"), + supported_model_types=[ModelType.LLM], + configurate_methods=[ConfigurateMethod.PREDEFINED_MODEL], + models=[], + ) + assembly = PluginModelAssembly(tenant_id="tenant-1") + + with patch("core.entities.provider_configuration.original_provider_configurate_methods", {}): + configuration = ProviderConfiguration( + tenant_id="tenant-1", + provider=provider, + preferred_provider_type=ProviderType.SYSTEM, + using_provider_type=ProviderType.SYSTEM, + system_configuration=system_configuration, + custom_configuration=CustomConfiguration(provider=None, models=[]), + model_settings=[], + ) + + with ( + patch("core.entities.provider_configuration.create_plugin_model_assembly", return_value=assembly), + patch.object(assembly.model_provider_factory, "get_model_provider", return_value=provider_schema), + patch.object( + ProviderConfiguration, + "get_model_schema", + side_effect=lambda *args, **kwargs: _build_ai_model(kwargs["model"]), + ) as mock_get_model_schema, + ): + models = configuration.get_provider_models( + model_type=ModelType.LLM, + only_active=False, + model="target-model", + ) + + assert [model.model for model in models] == ["target-model"] + mock_get_model_schema.assert_called_once() + assert mock_get_model_schema.call_args.kwargs["model"] == "target-model" + + @staticmethod + def test_get_custom_provider_models_sets_status_for_removed_credentials_and_invalid_lb_configs() -> None: + configuration = _build_provider_configuration() + configuration.using_provider_type = ProviderType.CUSTOM + configuration.custom_configuration.provider = CustomProviderConfiguration( + credentials={"api_key": "provider-key"} + ) + configuration.custom_configuration.models = [ + CustomModelConfiguration( + model="custom-model", + model_type=ModelType.LLM, + credentials=None, + available_model_credentials=[CredentialConfiguration(credential_id="c-1", credential_name="first")], + ) + ] + provider_schema = ProviderEntity( + provider="openai", + label=I18nObject(en_US="OpenAI"), + supported_model_types=[ModelType.LLM], + configurate_methods=[ConfigurateMethod.PREDEFINED_MODEL], + models=[_build_ai_model("base-model")], + ) + model_setting_map = { + ModelType.LLM: { + "base-model": ModelSettings( + model="base-model", + model_type=ModelType.LLM, + enabled=True, + load_balancing_enabled=True, + load_balancing_configs=[ + ModelLoadBalancingConfiguration( + id="lb-base", + name="LB Base", + credentials={}, + credential_source_type=CredentialSourceType.PROVIDER, + ) + ], + ), + "custom-model": ModelSettings( + model="custom-model", + model_type=ModelType.LLM, + enabled=True, + load_balancing_enabled=True, + load_balancing_configs=[ + ModelLoadBalancingConfiguration( + id="lb-custom", + name="LB Custom", + credentials={}, + credential_source_type=CredentialSourceType.CUSTOM_MODEL, + ) + ], + ), + } + } + + with patch.object(ProviderConfiguration, "get_model_schema", return_value=_build_ai_model("custom-model")): + models = configuration._get_custom_provider_models( + model_types=[ModelType.LLM], + provider_schema=provider_schema, + model_setting_map=model_setting_map, + ) + + status_map = {model.model: model.status for model in models} + invalid_lb_map = {model.model: model.has_invalid_load_balancing_configs for model in models} + assert status_map["base-model"] == ModelStatus.ACTIVE + assert status_map["custom-model"] == ModelStatus.CREDENTIAL_REMOVED + assert invalid_lb_map["base-model"] is True + assert invalid_lb_map["custom-model"] is True + + @staticmethod + def test_get_custom_provider_models_filters_requested_base_model() -> None: + configuration = _build_provider_configuration() + configuration.using_provider_type = ProviderType.CUSTOM + configuration.custom_configuration.provider = CustomProviderConfiguration( + credentials={"api_key": "provider-key"} + ) + provider_schema = ProviderEntity( + provider="openai", + label=I18nObject(en_US="OpenAI"), + supported_model_types=[ModelType.LLM], + configurate_methods=[ConfigurateMethod.PREDEFINED_MODEL], + models=[_build_ai_model("base-model"), _build_ai_model("target-model")], + ) + + models = configuration._get_custom_provider_models( + model_types=[ModelType.LLM], + provider_schema=provider_schema, + model_setting_map={}, + model="target-model", + ) + + assert [model.model for model in models] == ["target-model"] + + @staticmethod + def test_get_provider_models_reuses_cached_provider_schema() -> None: + configuration = _build_provider_configuration() + provider_schema = ProviderEntity( + provider="openai", + label=I18nObject(en_US="OpenAI"), + supported_model_types=[ModelType.LLM], + configurate_methods=[ConfigurateMethod.PREDEFINED_MODEL], + models=[_build_ai_model("a-model"), _build_ai_model("b-model")], + ) + configuration.provider = provider_schema + + with patch( + "core.entities.provider_configuration.create_plugin_model_assembly", + ) as mock_assembly_builder: + configuration.get_provider_models(model_type=ModelType.LLM, model="a-model") + configuration.get_provider_models(model_type=ModelType.LLM, model="b-model") + + mock_assembly_builder.assert_not_called() + + @staticmethod + def test_validator_adds_predefined_model_for_customizable_provider_with_restrictions() -> None: + provider = ProviderEntity( + provider="openai", + label=I18nObject(en_US="OpenAI"), + supported_model_types=[ModelType.LLM], + configurate_methods=[ConfigurateMethod.CUSTOMIZABLE_MODEL], + ) + system_configuration = SystemConfiguration( + enabled=True, + credentials={"api_key": "test-key"}, + current_quota_type=ProviderQuotaType.TRIAL, + quota_configurations=[ + QuotaConfiguration( + quota_type=ProviderQuotaType.TRIAL, + quota_unit=QuotaUnit.TOKENS, + quota_limit=100, + quota_used=0, + is_valid=True, + restrict_models=[ + RestrictModel(model="restricted", base_model_name="base-model", model_type=ModelType.LLM) + ], + ) + ], + ) + with patch("core.entities.provider_configuration.original_provider_configurate_methods", {}): + configuration = ProviderConfiguration( + tenant_id="tenant-1", + provider=provider, + preferred_provider_type=ProviderType.SYSTEM, + using_provider_type=ProviderType.SYSTEM, + system_configuration=system_configuration, + custom_configuration=CustomConfiguration(provider=None, models=[]), + model_settings=[], + ) + + assert ConfigurateMethod.PREDEFINED_MODEL in configuration.provider.configurate_methods + + @staticmethod + def test_get_current_credentials_system_handles_disable_and_restricted_base_model() -> None: + configuration = _build_provider_configuration() + configuration.model_settings = [ + ModelSettings(model="gpt-4o", model_type=ModelType.LLM, enabled=False, load_balancing_configs=[]) + ] + + with pytest.raises(ValueError, match="Model gpt-4o is disabled"): + configuration.get_current_credentials(ModelType.LLM, "gpt-4o") + + configuration.model_settings = [] + configuration.system_configuration.quota_configurations[0].restrict_models = [ + RestrictModel(model="gpt-4o", base_model_name="base-model", model_type=ModelType.LLM) + ] + credentials = configuration.get_current_credentials(ModelType.LLM, "gpt-4o") + assert credentials["base_model_name"] == "base-model" + + @staticmethod + def test_get_current_credentials_prefers_model_specific_custom_credentials() -> None: + configuration = _build_provider_configuration() + configuration.using_provider_type = ProviderType.CUSTOM + configuration.custom_configuration.models = [ + CustomModelConfiguration( + model="gpt-4o", + model_type=ModelType.LLM, + credentials={"api_key": "model-key"}, + ) + ] + configuration.custom_configuration.provider = CustomProviderConfiguration( + credentials={"api_key": "provider-key"} + ) + + credentials = configuration.get_current_credentials(ModelType.LLM, "gpt-4o") + assert credentials == {"api_key": "model-key"} + + @staticmethod + def test_get_system_configuration_status_falsey_quota_returns_unsupported() -> None: + class _FalseyQuota: + quota_type = ProviderQuotaType.TRIAL + is_valid = True + + def __bool__(self) -> bool: + return False + + configuration = _build_provider_configuration() + configuration.system_configuration.quota_configurations = [_FalseyQuota()] # type: ignore[list-item] + assert ( + configuration.get_system_configuration_status() + == _entities_provider_configuration_SystemConfigurationStatus.UNSUPPORTED + ) + + @staticmethod + def test_get_custom_provider_models_skips_schema_models_with_mismatched_type() -> None: + configuration = _build_provider_configuration() + provider_schema = ProviderEntity( + provider="openai", + label=I18nObject(en_US="OpenAI"), + supported_model_types=[ModelType.LLM, ModelType.TEXT_EMBEDDING], + configurate_methods=[ConfigurateMethod.PREDEFINED_MODEL], + models=[ + _build_ai_model("llm-model", model_type=ModelType.LLM), + _build_ai_model("embed-model", model_type=ModelType.TEXT_EMBEDDING), + ], + ) + + models = configuration._get_custom_provider_models( + model_types=[ModelType.LLM], + provider_schema=provider_schema, + model_setting_map={}, + ) + + assert any(model.model == "llm-model" for model in models) + assert all(model.model != "embed-model" for model in models) + + @staticmethod + def test_get_custom_provider_models_skips_custom_models_on_schema_error_or_none( + caplog: pytest.LogCaptureFixture, + ) -> None: + configuration = _build_provider_configuration() + configuration.custom_configuration.models = [ + CustomModelConfiguration(model="error-custom", model_type=ModelType.LLM, credentials={"k": "v"}), + CustomModelConfiguration(model="none-custom", model_type=ModelType.LLM, credentials={"k": "v"}), + CustomModelConfiguration(model="ok-custom", model_type=ModelType.LLM, credentials={"k": "v"}), + ] + provider_schema = ProviderEntity( + provider="openai", + label=I18nObject(en_US="OpenAI"), + supported_model_types=[ModelType.LLM], + configurate_methods=[ConfigurateMethod.PREDEFINED_MODEL], + models=[], + ) + + def _schema(*, model_type: ModelType, model: str, credentials: dict | None): + if model == "error-custom": + raise RuntimeError("boom") + if model == "none-custom": + return None + return _build_ai_model(model) + + with caplog.at_level(logging.WARNING, logger="core.entities.provider_configuration"): + with patch.object(ProviderConfiguration, "get_model_schema", side_effect=_schema): + models = configuration._get_custom_provider_models( + model_types=[ModelType.LLM], + provider_schema=provider_schema, + model_setting_map={}, + ) + + assert "get custom model schema failed, boom" in caplog.messages + assert any(model.model == "ok-custom" for model in models) + assert all(model.model != "none-custom" for model in models) + + @staticmethod + @pytest.fixture + def sqlite_provider_session( + sqlite_session: Session, + sqlite_engine: Engine, + ) -> Iterator[Session]: + """Bind provider-owned sessions to the same isolated SQLite database as the test.""" + with patch.object(type(db), "engine", new_callable=PropertyMock, return_value=sqlite_engine): + yield sqlite_session + + @staticmethod + def test_generate_credential_names_from_real_rows_and_tenant_isolation( + sqlite_provider_session: Session, + ) -> None: + configuration = _build_provider_configuration() + _provider_credential(sqlite_provider_session, name="API KEY 9") + _provider_credential(sqlite_provider_session, name="legacy") + _provider_credential(sqlite_provider_session, name="API KEY 50", tenant_id="other-tenant") + _model_credential(sqlite_provider_session, name="API KEY 4") + assert configuration._generate_provider_credential_name(sqlite_provider_session) == "API KEY 10" + assert ( + configuration._generate_custom_model_credential_name("gpt-4o", ModelType.LLM, sqlite_provider_session) + == "API KEY 5" + ) + + @staticmethod + def test_validate_provider_credentials_reuses_hidden_secret(sqlite_provider_session: Session) -> None: + configuration = _build_provider_configuration() + configuration.provider.provider_credential_schema = _build_secret_provider_schema() + credential = _provider_credential(sqlite_provider_session, encrypted_config='{"openai_api_key":"enc-old"}') + assembly = PluginModelAssembly(tenant_id="tenant-1") + with ( + patch("core.entities.provider_configuration.create_plugin_model_assembly", return_value=assembly), + patch.object(assembly.model_provider_factory, "get_model_provider", return_value=configuration.provider), + patch.object(assembly.model_runtime, "validate_provider_credentials", new=lambda **_kwargs: None), + patch("core.entities.provider_configuration.encrypter.decrypt_token", return_value="raw"), + patch("core.entities.provider_configuration.encrypter.encrypt_token", return_value="enc-new"), + ): + result = configuration.validate_provider_credentials( + {"openai_api_key": HIDDEN_VALUE}, credential_id=credential.id + ) + assert result == {"openai_api_key": "enc-new"} + + @staticmethod + def test_preferred_provider_state_updates_and_is_tenant_scoped(sqlite_provider_session: Session) -> None: + configuration = _build_provider_configuration() + configuration.preferred_provider_type = ProviderType.CUSTOM + other = TenantPreferredModelProvider( + tenant_id="other-tenant", provider_name="openai", preferred_provider_type=ProviderType.CUSTOM + ) + current = TenantPreferredModelProvider( + tenant_id="tenant-1", provider_name="openai", preferred_provider_type=ProviderType.CUSTOM + ) + sqlite_provider_session.add_all([other, current]) + sqlite_provider_session.commit() + assert configuration.switch_preferred_provider_type(ProviderType.SYSTEM, session=sqlite_provider_session) + sqlite_provider_session.refresh(current) + sqlite_provider_session.refresh(other) + assert current.preferred_provider_type == ProviderType.SYSTEM + assert other.preferred_provider_type == ProviderType.CUSTOM + + @staticmethod + def test_provider_record_duplicate_and_setting_helpers_use_real_session(sqlite_provider_session: Session) -> None: + configuration = _build_provider_configuration() + provider = _provider_record(sqlite_provider_session) + _provider_record(sqlite_provider_session, tenant_id="other-tenant") + credential = _provider_credential(sqlite_provider_session, name="Main") + _provider_credential(sqlite_provider_session, name="Main", tenant_id="other-tenant") + setting = ProviderModelSetting( + tenant_id="tenant-1", + provider_name="openai", + model_name="gpt-4o", + model_type=ModelType.LLM, + ) + sqlite_provider_session.add(setting) + sqlite_provider_session.commit() + assert configuration._get_provider_record(sqlite_provider_session).id == provider.id + assert configuration._check_provider_credential_name_exists("Main", sqlite_provider_session) + assert not configuration._check_provider_credential_name_exists( + "Main", sqlite_provider_session, exclude_id=credential.id + ) + assert ( + configuration._get_provider_model_setting(ModelType.LLM, "gpt-4o", sqlite_provider_session).id == setting.id + ) + + @staticmethod + def test_create_provider_credential_persists_provider_and_rejects_duplicate( + sqlite_provider_session: Session, + ) -> None: + configuration = _build_provider_configuration() + with ( + patch.object(ProviderConfiguration, "validate_provider_credentials", return_value={"api_key": "enc"}), + _mock_cache_boundaries() as (credentials_cache, configuration_cache), + ): + created_id = configuration.create_provider_credential({"api_key": "raw"}, "Main") + credential = sqlite_provider_session.scalar( + select(ProviderCredential).where(ProviderCredential.credential_name == "Main") + ) + provider = sqlite_provider_session.scalar(select(Provider).where(Provider.tenant_id == "tenant-1")) + assert credential is not None + assert provider is not None + assert created_id == credential.id + assert provider.credential_id == credential.id + credentials_cache.assert_called_once_with( + tenant_id="tenant-1", + identity_id=provider.id, + cache_type=ProviderCredentialsCacheType.PROVIDER, + ) + credentials_cache.return_value.delete.assert_called_once_with() + configuration_cache.assert_called_once_with( + preferred_model_providers=True, + provider_credentials=True, + ) + with pytest.raises(ValueError, match="already exists"): + configuration.create_provider_credential({"api_key": "raw"}, "Main") + + @staticmethod + def test_update_provider_credential_propagates_to_load_balancing(sqlite_provider_session: Session) -> None: + configuration = _build_provider_configuration() + credential = _provider_credential(sqlite_provider_session, name="Old") + provider = _provider_record(sqlite_provider_session, credential_id=credential.id) + lb_config = _load_balancing_config( + sqlite_provider_session, credential_id=credential.id, source=CredentialSourceType.PROVIDER + ) + with ( + patch.object(ProviderConfiguration, "validate_provider_credentials", return_value={"api_key": "enc-new"}), + _mock_cache_boundaries() as (credentials_cache, configuration_cache), + ): + configuration.update_provider_credential({"api_key": "raw"}, credential.id, "New") + sqlite_provider_session.expire_all() + persisted_credential = sqlite_provider_session.get(ProviderCredential, credential.id) + persisted_lb = sqlite_provider_session.get(LoadBalancingModelConfig, lb_config.id) + assert persisted_credential is not None + assert persisted_credential.credential_name == "New" + assert persisted_lb is not None + assert persisted_lb.name == "New" + assert json.loads(persisted_lb.encrypted_config) == {"api_key": "enc-new"} + assert {cache_call.kwargs["identity_id"] for cache_call in credentials_cache.call_args_list} == { + provider.id, + lb_config.id, + } + assert {cache_call.kwargs["cache_type"] for cache_call in credentials_cache.call_args_list} == { + ProviderCredentialsCacheType.PROVIDER, + ProviderCredentialsCacheType.LOAD_BALANCING_MODEL, + } + assert credentials_cache.return_value.delete.call_count == 2 + configuration_cache.assert_called_once_with( + provider_credentials=True, + provider_load_balancing_configs=True, + ) + + @staticmethod + def test_switch_active_provider_credential_updates_persisted_state_and_cache( + sqlite_provider_session: Session, + ) -> None: + configuration = _build_provider_configuration() + configuration.preferred_provider_type = ProviderType.CUSTOM + first = _provider_credential(sqlite_provider_session, name="First") + second = _provider_credential(sqlite_provider_session, name="Second") + provider = _provider_record(sqlite_provider_session, credential_id=first.id) + provider_id = provider.id + + with _mock_cache_boundaries() as (credentials_cache, configuration_cache): + configuration.switch_active_provider_credential(second.id) + + sqlite_provider_session.expire_all() + persisted_provider = sqlite_provider_session.get(Provider, provider_id) + assert persisted_provider is not None + assert persisted_provider.credential_id == second.id + credentials_cache.assert_called_once_with( + tenant_id="tenant-1", + identity_id=provider_id, + cache_type=ProviderCredentialsCacheType.PROVIDER, + ) + credentials_cache.return_value.delete.assert_called_once_with() + configuration_cache.assert_not_called() + + @staticmethod + def test_deleting_active_provider_credential_switches_preference_to_system( + sqlite_provider_session: Session, + ) -> None: + configuration = _build_provider_configuration() + configuration.preferred_provider_type = ProviderType.CUSTOM + first = _provider_credential(sqlite_provider_session, name="First") + active = _provider_credential(sqlite_provider_session, name="Active") + provider = _provider_record(sqlite_provider_session, credential_id=active.id) + preferred_provider = TenantPreferredModelProvider( + tenant_id="tenant-1", + provider_name="openai", + preferred_provider_type=ProviderType.CUSTOM, + ) + sqlite_provider_session.add(preferred_provider) + sqlite_provider_session.commit() + first_id = first.id + active_id = active.id + provider_id = provider.id + preferred_provider_id = preferred_provider.id + + with _mock_cache_boundaries() as (credentials_cache, configuration_cache): + configuration.delete_provider_credential(active_id) + + sqlite_provider_session.expire_all() + assert sqlite_provider_session.get(ProviderCredential, active_id) is None + assert sqlite_provider_session.get(ProviderCredential, first_id) is not None + persisted_provider = sqlite_provider_session.get(Provider, provider_id) + persisted_preference = sqlite_provider_session.get(TenantPreferredModelProvider, preferred_provider_id) + assert persisted_provider is not None + assert persisted_provider.credential_id is None + assert persisted_preference is not None + assert persisted_preference.preferred_provider_type == ProviderType.SYSTEM + credentials_cache.assert_called_once_with( + tenant_id="tenant-1", + identity_id=provider_id, + cache_type=ProviderCredentialsCacheType.PROVIDER, + ) + credentials_cache.return_value.delete.assert_called_once_with() + configuration_cache.assert_called_once_with( + preferred_model_providers=True, + provider_credentials=True, + provider_load_balancing_configs=False, + ) + + @staticmethod + def test_specific_provider_credential_decrypts_and_obfuscates(sqlite_provider_session: Session) -> None: + configuration = _build_provider_configuration() + configuration.provider.provider_credential_schema = _build_secret_provider_schema() + credential = _provider_credential(sqlite_provider_session, encrypted_config='{"openai_api_key":"enc"}') + with ( + patch("core.entities.provider_configuration.encrypter.decrypt_token", return_value="raw"), + patch("core.entities.provider_configuration.encrypter.obfuscated_token", return_value="masked"), + ): + result = configuration._get_specific_provider_credential(credential.id) + assert result == {"openai_api_key": "masked"} + with pytest.raises(ValueError, match="not found"): + configuration._get_specific_provider_credential("missing") + + @staticmethod + def test_validate_custom_model_credentials_reuses_hidden_secret(sqlite_provider_session: Session) -> None: + configuration = _build_provider_configuration() + configuration.provider.model_credential_schema = _build_secret_model_schema() + credential = _model_credential(sqlite_provider_session, encrypted_config='{"openai_api_key":"enc-old"}') + assembly = PluginModelAssembly(tenant_id="tenant-1") + with ( + patch("core.entities.provider_configuration.create_plugin_model_assembly", return_value=assembly), + patch.object(assembly.model_provider_factory, "get_model_provider", return_value=configuration.provider), + patch.object(assembly.model_runtime, "validate_model_credentials", new=lambda **_kwargs: None), + patch("core.entities.provider_configuration.encrypter.decrypt_token", return_value="raw"), + patch("core.entities.provider_configuration.encrypter.encrypt_token", return_value="enc-new"), + ): + result = configuration.validate_custom_model_credentials( + ModelType.LLM, + "gpt-4o", + {"openai_api_key": HIDDEN_VALUE}, + credential_id=credential.id, + ) + assert result == {"openai_api_key": "enc-new"} + + @staticmethod + def test_specific_custom_model_credential_preserves_secret_when_decryption_fails( + sqlite_provider_session: Session, + caplog: pytest.LogCaptureFixture, + ) -> None: + configuration = _build_provider_configuration() + configuration.provider.model_credential_schema = _build_secret_model_schema() + credential = _model_credential( + sqlite_provider_session, + name="Main", + encrypted_config='{"openai_api_key":"enc-secret"}', + ) + + with ( + caplog.at_level(logging.ERROR, logger="core.entities.provider_configuration"), + patch("core.entities.provider_configuration.encrypter.decrypt_token", side_effect=RuntimeError("boom")), + patch.object( + ProviderConfiguration, + "obfuscated_credentials", + side_effect=lambda credentials, credential_form_schemas: credentials, + ), + ): + result = configuration._get_specific_custom_model_credential(ModelType.LLM, "gpt-4o", credential.id) + + assert result == { + "current_credential_id": credential.id, + "current_credential_name": "Main", + "credentials": {"openai_api_key": "enc-secret"}, + } + assert caplog.messages.count("Failed to decrypt model credential secret variable openai_api_key") == 1 + + @staticmethod + def test_create_update_and_delete_custom_model_credential(sqlite_provider_session: Session) -> None: + configuration = _build_provider_configuration() + with ( + patch.object(ProviderConfiguration, "validate_custom_model_credentials", return_value={"api_key": "enc"}), + _mock_cache_boundaries() as (credentials_cache, configuration_cache), + ): + created_id = configuration.create_custom_model_credential( + ModelType.LLM, "gpt-4o", {"api_key": "raw"}, "Main" + ) + credential = sqlite_provider_session.scalar(select(ProviderModelCredential)) + model = sqlite_provider_session.scalar(select(ProviderModel)) + assert credential is not None + assert model is not None + assert created_id == credential.id + assert model.credential_id == credential.id + credential_id = credential.id + model_id = model.id + lb_config = _load_balancing_config( + sqlite_provider_session, + credential_id=credential_id, + source=CredentialSourceType.CUSTOM_MODEL, + ) + lb_config_id = lb_config.id + credentials_cache.assert_called_once_with( + tenant_id="tenant-1", + identity_id=model_id, + cache_type=ProviderCredentialsCacheType.MODEL, + ) + credentials_cache.return_value.delete.assert_called_once_with() + configuration_cache.assert_called_once_with( + provider_models=True, + provider_model_credentials=True, + ) + + with ( + patch.object(ProviderConfiguration, "validate_custom_model_credentials", return_value={"api_key": "enc-2"}), + _mock_cache_boundaries() as (credentials_cache, configuration_cache), + ): + configuration.update_custom_model_credential( + ModelType.LLM, "gpt-4o", {"api_key": "raw"}, "Renamed", credential_id + ) + sqlite_provider_session.expire_all() + persisted_credential = sqlite_provider_session.get(ProviderModelCredential, credential_id) + persisted_lb = sqlite_provider_session.get(LoadBalancingModelConfig, lb_config_id) + assert persisted_credential is not None + assert persisted_credential.credential_name == "Renamed" + assert json.loads(persisted_credential.encrypted_config) == {"api_key": "enc-2"} + assert persisted_lb is not None + assert persisted_lb.name == "Renamed" + assert json.loads(persisted_lb.encrypted_config) == {"api_key": "enc-2"} + assert {cache_call.kwargs["identity_id"] for cache_call in credentials_cache.call_args_list} == { + model_id, + lb_config_id, + } + assert {cache_call.kwargs["cache_type"] for cache_call in credentials_cache.call_args_list} == { + ProviderCredentialsCacheType.MODEL, + ProviderCredentialsCacheType.LOAD_BALANCING_MODEL, + } + assert credentials_cache.return_value.delete.call_count == 2 + configuration_cache.assert_called_once_with( + provider_models=True, + provider_model_credentials=True, + provider_load_balancing_configs=True, + ) + + with _mock_cache_boundaries() as (credentials_cache, configuration_cache): + configuration.delete_custom_model_credential(ModelType.LLM, "gpt-4o", credential_id) + sqlite_provider_session.expire_all() + assert sqlite_provider_session.get(ProviderModelCredential, credential_id) is None + assert sqlite_provider_session.get(ProviderModel, model_id) is None + assert sqlite_provider_session.get(LoadBalancingModelConfig, lb_config_id) is None + assert {cache_call.kwargs["identity_id"] for cache_call in credentials_cache.call_args_list} == { + model_id, + lb_config_id, + } + assert {cache_call.kwargs["cache_type"] for cache_call in credentials_cache.call_args_list} == { + ProviderCredentialsCacheType.MODEL, + ProviderCredentialsCacheType.LOAD_BALANCING_MODEL, + } + assert credentials_cache.return_value.delete.call_count == 2 + configuration_cache.assert_called_once_with( + provider_models=True, + provider_model_credentials=True, + provider_load_balancing_configs=True, + ) + + @staticmethod + def test_add_and_switch_custom_model_credential(sqlite_provider_session: Session) -> None: + configuration = _build_provider_configuration() + first = _model_credential(sqlite_provider_session, name="First") + second = _model_credential(sqlite_provider_session, name="Second") + with _mock_cache_boundaries() as (credentials_cache, configuration_cache): + configuration.add_model_credential_to_model(ModelType.LLM, "gpt-4o", first.id) + configuration.switch_custom_model_credential(ModelType.LLM, "gpt-4o", second.id) + model = sqlite_provider_session.scalar(select(ProviderModel)) + assert model is not None + assert model.credential_id == second.id + credentials_cache.assert_called_once_with( + tenant_id="tenant-1", + identity_id=model.id, + cache_type=ProviderCredentialsCacheType.MODEL, + ) + credentials_cache.return_value.delete.assert_called_once_with() + assert configuration_cache.call_count == 2 + for cache_call in configuration_cache.call_args_list: + assert cache_call.kwargs == {"provider_models": True} + with pytest.raises(ValueError, match="Can't add same credential"): + configuration.add_model_credential_to_model(ModelType.LLM, "gpt-4o", second.id) + + @staticmethod + def test_model_settings_and_load_balancing_persist(sqlite_provider_session: Session) -> None: + configuration = _build_provider_configuration() + with patch.object(configuration, "_invalidate_provider_configuration_cache") as configuration_cache: + configuration.disable_model(ModelType.LLM, "gpt-4o") + configuration_cache.assert_called_once_with(provider_model_settings=True) + persisted_setting = sqlite_provider_session.scalar(select(ProviderModelSetting)) + assert persisted_setting is not None + assert persisted_setting.enabled is False + with patch.object(configuration, "_invalidate_provider_configuration_cache") as configuration_cache: + configuration.enable_model(ModelType.LLM, "gpt-4o") + configuration_cache.assert_called_once_with(provider_model_settings=True) + sqlite_provider_session.expire_all() + refreshed_setting = sqlite_provider_session.get(ProviderModelSetting, persisted_setting.id) + assert refreshed_setting is not None + assert refreshed_setting.enabled is True + + first = _provider_credential(sqlite_provider_session, name="First") + second = _provider_credential(sqlite_provider_session, name="Second") + _load_balancing_config(sqlite_provider_session, credential_id=first.id, source=CredentialSourceType.PROVIDER) + _load_balancing_config(sqlite_provider_session, credential_id=second.id, source=CredentialSourceType.PROVIDER) + with patch.object(configuration, "_invalidate_provider_configuration_cache") as configuration_cache: + configuration.enable_model_load_balancing(ModelType.LLM, "gpt-4o") + configuration_cache.assert_called_once_with(provider_model_settings=True) + sqlite_provider_session.expire_all() + refreshed_setting = sqlite_provider_session.get(ProviderModelSetting, persisted_setting.id) + assert refreshed_setting is not None + assert refreshed_setting.load_balancing_enabled is True + with patch.object(configuration, "_invalidate_provider_configuration_cache") as configuration_cache: + configuration.disable_model_load_balancing(ModelType.LLM, "gpt-4o") + configuration_cache.assert_called_once_with(provider_model_settings=True) + sqlite_provider_session.expire_all() + refreshed_setting = sqlite_provider_session.get(ProviderModelSetting, persisted_setting.id) + assert refreshed_setting is not None + assert refreshed_setting.load_balancing_enabled is False + + @staticmethod + def test_provider_create_rolls_back_on_insert_failure( + sqlite_provider_session: Session, + sqlite_engine: Engine, + ) -> None: + configuration = _build_provider_configuration() + with ( + patch.object(ProviderConfiguration, "validate_provider_credentials", return_value={"api_key": "enc"}), + _mock_cache_boundaries(), + _raise_on_sql(sqlite_engine, "provider_credentials", "INSERT"), + pytest.raises(RuntimeError, match="forced INSERT"), + ): + configuration.create_provider_credential({"api_key": "raw"}, "Main") + assert sqlite_provider_session.scalar(select(ProviderCredential)) is None + assert sqlite_provider_session.scalar(select(Provider)) is None + + @staticmethod + def test_custom_model_create_rolls_back_on_insert_failure( + sqlite_provider_session: Session, + sqlite_engine: Engine, + ) -> None: + configuration = _build_provider_configuration() + with ( + patch.object(ProviderConfiguration, "validate_custom_model_credentials", return_value={"api_key": "enc"}), + _mock_cache_boundaries(), + _raise_on_sql(sqlite_engine, "provider_model_credentials", "INSERT"), + pytest.raises(RuntimeError, match="forced INSERT"), + ): + configuration.create_custom_model_credential(ModelType.LLM, "gpt-4o", {"api_key": "raw"}, "Main") + assert sqlite_provider_session.scalar(select(ProviderModelCredential)) is None + assert sqlite_provider_session.scalar(select(ProviderModel)) is None + + @staticmethod + def test_provider_update_rolls_back_on_update_failure( + sqlite_provider_session: Session, + sqlite_engine: Engine, + ) -> None: + configuration = _build_provider_configuration() + credential = _provider_credential( + sqlite_provider_session, + name="Old", + encrypted_config='{"api_key":"enc-old"}', + ) + _provider_record(sqlite_provider_session, credential_id=credential.id) + + with ( + patch.object(ProviderConfiguration, "validate_provider_credentials", return_value={"api_key": "enc-new"}), + _mock_cache_boundaries() as (credentials_cache, configuration_cache), + _raise_on_sql(sqlite_engine, "provider_credentials", "UPDATE"), + pytest.raises(RuntimeError, match="forced UPDATE"), + ): + configuration.update_provider_credential({"api_key": "raw"}, credential.id, "New") + + sqlite_provider_session.expire_all() + persisted = sqlite_provider_session.get(ProviderCredential, credential.id) + assert persisted is not None + assert persisted.credential_name == "Old" + assert json.loads(persisted.encrypted_config) == {"api_key": "enc-old"} + credentials_cache.assert_not_called() + configuration_cache.assert_not_called() + + @staticmethod + def test_provider_delete_rolls_back_on_delete_failure( + sqlite_provider_session: Session, + sqlite_engine: Engine, + ) -> None: + configuration = _build_provider_configuration() + credential = _provider_credential(sqlite_provider_session) + provider = _provider_record(sqlite_provider_session, credential_id=credential.id) + + with ( + _mock_cache_boundaries() as (credentials_cache, configuration_cache), + _raise_on_sql(sqlite_engine, "provider_credentials", "DELETE"), + pytest.raises(RuntimeError, match="forced DELETE"), + ): + configuration.delete_provider_credential(credential.id) + + sqlite_provider_session.expire_all() + assert sqlite_provider_session.get(ProviderCredential, credential.id) is not None + persisted_provider = sqlite_provider_session.get(Provider, provider.id) + assert persisted_provider is not None + assert persisted_provider.credential_id == credential.id + credentials_cache.assert_called_once_with( + tenant_id="tenant-1", + identity_id=provider.id, + cache_type=ProviderCredentialsCacheType.PROVIDER, + ) + credentials_cache.return_value.delete.assert_called_once_with() + configuration_cache.assert_not_called() + + @staticmethod + def test_provider_switch_rolls_back_on_update_failure( + sqlite_provider_session: Session, + sqlite_engine: Engine, + ) -> None: + configuration = _build_provider_configuration() + first = _provider_credential(sqlite_provider_session, name="First") + second = _provider_credential(sqlite_provider_session, name="Second") + provider = _provider_record(sqlite_provider_session, credential_id=first.id) + + with ( + _mock_cache_boundaries() as (credentials_cache, configuration_cache), + _raise_on_sql(sqlite_engine, "providers", "UPDATE"), + pytest.raises(RuntimeError, match="forced UPDATE"), + ): + configuration.switch_active_provider_credential(second.id) + + sqlite_provider_session.expire_all() + persisted_provider = sqlite_provider_session.get(Provider, provider.id) + assert persisted_provider is not None + assert persisted_provider.credential_id == first.id + credentials_cache.assert_not_called() + configuration_cache.assert_not_called() + + @staticmethod + def test_custom_model_update_rolls_back_on_update_failure( + sqlite_provider_session: Session, + sqlite_engine: Engine, + ) -> None: + configuration = _build_provider_configuration() + credential = _model_credential( + sqlite_provider_session, + name="Old", + encrypted_config='{"api_key":"enc-old"}', + ) + _provider_model_record(sqlite_provider_session, credential_id=credential.id) + + with ( + patch.object( + ProviderConfiguration, + "validate_custom_model_credentials", + return_value={"api_key": "enc-new"}, + ), + _mock_cache_boundaries() as (credentials_cache, configuration_cache), + _raise_on_sql(sqlite_engine, "provider_model_credentials", "UPDATE"), + pytest.raises(RuntimeError, match="forced UPDATE"), + ): + configuration.update_custom_model_credential( + ModelType.LLM, + "gpt-4o", + {"api_key": "raw"}, + "New", + credential.id, + ) + + sqlite_provider_session.expire_all() + persisted = sqlite_provider_session.get(ProviderModelCredential, credential.id) + assert persisted is not None + assert persisted.credential_name == "Old" + assert json.loads(persisted.encrypted_config) == {"api_key": "enc-old"} + credentials_cache.assert_not_called() + configuration_cache.assert_not_called() + + @staticmethod + def test_custom_model_delete_rolls_back_on_delete_failure( + sqlite_provider_session: Session, + sqlite_engine: Engine, + ) -> None: + configuration = _build_provider_configuration() + credential = _model_credential(sqlite_provider_session) + model = _provider_model_record(sqlite_provider_session, credential_id=credential.id) + + with ( + _mock_cache_boundaries() as (credentials_cache, configuration_cache), + _raise_on_sql(sqlite_engine, "provider_model_credentials", "DELETE"), + pytest.raises(RuntimeError, match="forced DELETE"), + ): + configuration.delete_custom_model_credential(ModelType.LLM, "gpt-4o", credential.id) + + sqlite_provider_session.expire_all() + assert sqlite_provider_session.get(ProviderModelCredential, credential.id) is not None + persisted_model = sqlite_provider_session.get(ProviderModel, model.id) + assert persisted_model is not None + assert persisted_model.credential_id == credential.id + credentials_cache.assert_called_once_with( + tenant_id="tenant-1", + identity_id=model.id, + cache_type=ProviderCredentialsCacheType.MODEL, + ) + credentials_cache.return_value.delete.assert_called_once_with() + configuration_cache.assert_not_called() + + @staticmethod + def test_custom_model_switch_rolls_back_on_update_failure( + sqlite_provider_session: Session, + sqlite_engine: Engine, + ) -> None: + configuration = _build_provider_configuration() + first = _model_credential(sqlite_provider_session, name="First") + second = _model_credential(sqlite_provider_session, name="Second") + model = _provider_model_record(sqlite_provider_session, credential_id=first.id) + + with ( + _mock_cache_boundaries() as (credentials_cache, configuration_cache), + _raise_on_sql(sqlite_engine, "provider_models", "UPDATE"), + pytest.raises(RuntimeError, match="forced UPDATE"), + ): + configuration.switch_custom_model_credential(ModelType.LLM, "gpt-4o", second.id) + + sqlite_provider_session.expire_all() + persisted_model = sqlite_provider_session.get(ProviderModel, model.id) + assert persisted_model is not None + assert persisted_model.credential_id == first.id + credentials_cache.assert_not_called() + configuration_cache.assert_not_called() + + +"""Provider-configuration behavior with persisted SQLite credential records.""" + + +class TestProviderConfigurationCases: + @staticmethod + @pytest.fixture + def mock_provider_entity() -> ProviderEntity: + """Mock provider entity with basic configuration""" + provider_entity = ProviderEntity( + provider="openai", + label=I18nObject(en_US="OpenAI", zh_Hans="OpenAI"), + description=I18nObject(en_US="OpenAI provider", zh_Hans="OpenAI 提供商"), + icon_small=I18nObject(en_US="icon.png", zh_Hans="icon.png"), + background="background.png", + help=None, + supported_model_types=[ModelType.LLM], + configurate_methods=[ConfigurateMethod.PREDEFINED_MODEL], + provider_credential_schema=None, + model_credential_schema=None, + ) + + return provider_entity + + @staticmethod + @pytest.fixture + def mock_system_configuration() -> SystemConfiguration: + """Mock system configuration""" + quota_config = QuotaConfiguration( + quota_type=ProviderQuotaType.TRIAL, + quota_unit=QuotaUnit.TOKENS, + quota_limit=1000, + quota_used=0, + is_valid=True, + restrict_models=[RestrictModel(model="gpt-4", reason="Experimental", model_type=ModelType.LLM)], + ) + + system_config = SystemConfiguration( + enabled=True, + credentials={"openai_api_key": "test_key"}, + quota_configurations=[quota_config], + current_quota_type=ProviderQuotaType.TRIAL, + ) + + return system_config + + @staticmethod + @pytest.fixture + def mock_custom_configuration() -> CustomConfiguration: + """Mock custom configuration""" + custom_config = CustomConfiguration(provider=None, models=[]) + return custom_config + + @staticmethod + @pytest.fixture + def provider_configuration( + mock_provider_entity: ProviderEntity, + mock_system_configuration: SystemConfiguration, + mock_custom_configuration: CustomConfiguration, + ) -> ProviderConfiguration: + """Create a test provider configuration instance""" + with patch( + "core.entities.provider_configuration.original_provider_configurate_methods", + dict[str, list[ConfigurateMethod]](), + ): + return ProviderConfiguration( + tenant_id="test_tenant", + provider=mock_provider_entity, + preferred_provider_type=ProviderType.SYSTEM, + using_provider_type=ProviderType.SYSTEM, + system_configuration=mock_system_configuration, + custom_configuration=mock_custom_configuration, + model_settings=[], + ) + + class TestProviderConfiguration: + """Test cases for ProviderConfiguration class""" + + def test_get_current_credentials_system_provider_success( + self, provider_configuration: ProviderConfiguration + ) -> None: + """Test successfully getting credentials from system provider""" + # Arrange + provider_configuration.using_provider_type = ProviderType.SYSTEM + + # Act + credentials = provider_configuration.get_current_credentials(ModelType.LLM, "gpt-4") + + # Assert + assert credentials == {"openai_api_key": "test_key"} + + def test_get_current_credentials_model_disabled(self, provider_configuration: ProviderConfiguration) -> None: + """Test getting credentials when model is disabled""" + # Arrange + model_setting = ModelSettings( + model="gpt-4", + model_type=ModelType.LLM, + enabled=False, + load_balancing_configs=[], + has_invalid_load_balancing_configs=False, + ) + provider_configuration.model_settings = [model_setting] + + # Act & Assert + with pytest.raises(ValueError, match="Model gpt-4 is disabled"): + provider_configuration.get_current_credentials(ModelType.LLM, "gpt-4") + + def test_get_current_credentials_custom_provider_with_models( + self, provider_configuration: ProviderConfiguration + ) -> None: + """Test getting credentials from custom provider with model configurations""" + # Arrange + provider_configuration.using_provider_type = ProviderType.CUSTOM + + mock_model_config = Mock() + mock_model_config.model_type = ModelType.LLM + mock_model_config.model = "gpt-4" + mock_model_config.credentials = {"openai_api_key": "custom_key"} + provider_configuration.custom_configuration.models = [mock_model_config] + + # Act + credentials = provider_configuration.get_current_credentials(ModelType.LLM, "gpt-4") + + # Assert + assert credentials == {"openai_api_key": "custom_key"} + + def test_get_system_configuration_status_active(self, provider_configuration: ProviderConfiguration) -> None: + """Test getting active system configuration status""" + # Arrange + provider_configuration.system_configuration.enabled = True + + # Act + status = provider_configuration.get_system_configuration_status() + + # Assert + assert status == _provider_configuration_SystemConfigurationStatus.ACTIVE + + def test_get_system_configuration_status_unsupported( + self, provider_configuration: ProviderConfiguration + ) -> None: + """Test getting unsupported system configuration status""" + # Arrange + provider_configuration.system_configuration.enabled = False + + # Act + status = provider_configuration.get_system_configuration_status() + + # Assert + assert status == _provider_configuration_SystemConfigurationStatus.UNSUPPORTED + + def test_get_system_configuration_status_quota_exceeded( + self, provider_configuration: ProviderConfiguration + ) -> None: + """Test getting quota exceeded system configuration status""" + # Arrange + provider_configuration.system_configuration.enabled = True + quota_config = provider_configuration.system_configuration.quota_configurations[0] + quota_config.is_valid = False + + # Act + status = provider_configuration.get_system_configuration_status() + + # Assert + assert status == _provider_configuration_SystemConfigurationStatus.QUOTA_EXCEEDED + + def test_is_custom_configuration_available_with_provider( + self, provider_configuration: ProviderConfiguration + ) -> None: + """Test custom configuration availability with provider credentials""" + # Arrange + mock_provider = Mock() + mock_provider.available_credentials = ["openai_api_key"] + provider_configuration.custom_configuration.provider = mock_provider + provider_configuration.custom_configuration.models = [] + + # Act + result = provider_configuration.is_custom_configuration_available() + + # Assert + assert result is True + + def test_is_custom_configuration_available_with_models( + self, provider_configuration: ProviderConfiguration + ) -> None: + """Test custom configuration availability with model configurations""" + # Arrange + provider_configuration.custom_configuration.provider = None + provider_configuration.custom_configuration.models = [Mock()] + + # Act + result = provider_configuration.is_custom_configuration_available() + + # Assert + assert result is True + + def test_is_custom_configuration_available_false(self, provider_configuration: ProviderConfiguration) -> None: + """Test custom configuration not available""" + # Arrange + provider_configuration.custom_configuration.provider = None + provider_configuration.custom_configuration.models = [] + + # Act + result = provider_configuration.is_custom_configuration_available() + + # Assert + assert result is False + + @pytest.mark.parametrize("sqlite_session", [(Provider,)], indirect=True) + def test_get_provider_record_found( + self, provider_configuration: ProviderConfiguration, sqlite_session: Session + ) -> None: + """Test getting provider record successfully""" + provider = Provider(tenant_id="test_tenant", provider_name="openai") + sqlite_session.add(provider) + sqlite_session.commit() + + result = provider_configuration._get_provider_record(sqlite_session) + + assert result is provider + + @pytest.mark.parametrize("sqlite_session", [(Provider,)], indirect=True) + def test_get_provider_record_not_found( + self, provider_configuration: ProviderConfiguration, sqlite_session: Session + ) -> None: + """Test getting provider record when not found""" + sqlite_session.add(Provider(tenant_id="other_tenant", provider_name="openai")) + sqlite_session.commit() + + result = provider_configuration._get_provider_record(sqlite_session) + + assert result is None + + def test_init_with_customizable_model_only( + self, + mock_provider_entity: ProviderEntity, + mock_system_configuration: SystemConfiguration, + mock_custom_configuration: CustomConfiguration, + ) -> None: + """Test initialization with customizable model only configuration""" + # Arrange + mock_provider_entity.configurate_methods = [ConfigurateMethod.CUSTOMIZABLE_MODEL] + + # Act + with patch( + "core.entities.provider_configuration.original_provider_configurate_methods", + dict[str, list[ConfigurateMethod]](), + ): + config = ProviderConfiguration( + tenant_id="test_tenant", + provider=mock_provider_entity, + preferred_provider_type=ProviderType.SYSTEM, + using_provider_type=ProviderType.SYSTEM, + system_configuration=mock_system_configuration, + custom_configuration=mock_custom_configuration, + model_settings=[], + ) + + # Assert + assert ConfigurateMethod.PREDEFINED_MODEL in config.provider.configurate_methods + assert config.provider is not mock_provider_entity + assert mock_provider_entity.configurate_methods == [ConfigurateMethod.CUSTOMIZABLE_MODEL] + + def test_get_current_credentials_with_restricted_models( + self, provider_configuration: ProviderConfiguration + ) -> None: + """Test getting credentials with model restrictions""" + # Arrange + provider_configuration.using_provider_type = ProviderType.SYSTEM + + # Act + credentials = provider_configuration.get_current_credentials(ModelType.LLM, "gpt-3.5-turbo") + + # Assert + assert credentials is not None + assert "openai_api_key" in credentials + + @pytest.mark.parametrize("sqlite_session", [(Provider, ProviderCredential)], indirect=True) + def test_get_specific_provider_credential_success( + self, + monkeypatch: pytest.MonkeyPatch, + provider_configuration: ProviderConfiguration, + sqlite_session: Session, + ) -> None: + """Test getting specific provider credential successfully""" + credential_id = "test_credential_id" + credential = ProviderCredential( + tenant_id="test_tenant", + provider_name="openai", + credential_name="primary", + encrypted_config='{"openai_api_key": "encrypted_key"}', + ) + credential.id = credential_id + sqlite_session.add_all((Provider(tenant_id="test_tenant", provider_name="openai"), credential)) + sqlite_session.commit() + monkeypatch.setattr( + provider_configuration_module, + "db", + SimpleNamespace(engine=sqlite_session.get_bind()), + ) + + result = provider_configuration._get_specific_provider_credential(credential_id) + + assert result == {"openai_api_key": "encrypted_key"} + + @pytest.mark.parametrize("sqlite_session", [(Provider, ProviderCredential)], indirect=True) + def test_get_specific_provider_credential_with_mixed_provider_names( + self, + monkeypatch: pytest.MonkeyPatch, + provider_configuration: ProviderConfiguration, + sqlite_session: Session, + ) -> None: + provider_configuration.provider.provider = "langgenius/deepseek/deepseek" + credential = ProviderCredential( + tenant_id="test_tenant", + provider_name="langgenius/deepseek/deepseek", + credential_name="primary", + encrypted_config='{"api_key": "encrypted_key"}', + ) + sqlite_session.add_all((Provider(tenant_id="test_tenant", provider_name="deepseek"), credential)) + sqlite_session.commit() + monkeypatch.setattr( + provider_configuration_module, + "db", + SimpleNamespace(engine=sqlite_session.get_bind()), + ) + + result = provider_configuration._get_specific_provider_credential(credential.id) + + assert result == {"api_key": "encrypted_key"} + + @pytest.mark.parametrize("sqlite_session", [(Provider, ProviderCredential)], indirect=True) + def test_get_specific_provider_credential_not_found( + self, + monkeypatch: pytest.MonkeyPatch, + provider_configuration: ProviderConfiguration, + sqlite_session: Session, + ) -> None: + """Test getting specific provider credential when not found""" + credential_id = "nonexistent_credential_id" + sqlite_session.add(Provider(tenant_id="other_tenant", provider_name="openai")) + sqlite_session.commit() + monkeypatch.setattr( + provider_configuration_module, + "db", + SimpleNamespace(engine=sqlite_session.get_bind()), + ) + + with pytest.raises(ValueError, match=credential_id): + provider_configuration._get_specific_provider_credential(credential_id) + + def test_extract_secret_variables_with_secret_input( + self, provider_configuration: ProviderConfiguration + ) -> None: + """Test extracting secret variables from credential form schemas""" + # Arrange + credential_form_schemas = [ + CredentialFormSchema( + variable="api_key", + label=I18nObject(en_US="API Key", zh_Hans="API 密钥"), + type=FormType.SECRET_INPUT, + required=True, + ), + CredentialFormSchema( + variable="model_name", + label=I18nObject(en_US="Model Name", zh_Hans="模型名称"), + type=FormType.TEXT_INPUT, + required=True, + ), + CredentialFormSchema( + variable="secret_token", + label=I18nObject(en_US="Secret Token", zh_Hans="密钥令牌"), + type=FormType.SECRET_INPUT, + required=False, + ), + ] + + # Act + secret_variables = provider_configuration.extract_secret_variables(credential_form_schemas) + + # Assert + assert len(secret_variables) == 2 + assert "api_key" in secret_variables + assert "secret_token" in secret_variables + assert "model_name" not in secret_variables + + def test_extract_secret_variables_no_secret_input(self, provider_configuration: ProviderConfiguration) -> None: + """Test extracting secret variables when no secret input fields exist""" + # Arrange + credential_form_schemas = [ + CredentialFormSchema( + variable="model_name", + label=I18nObject(en_US="Model Name", zh_Hans="模型名称"), + type=FormType.TEXT_INPUT, + required=True, + ), + CredentialFormSchema( + variable="temperature", + label=I18nObject(en_US="Temperature", zh_Hans="温度"), + type=FormType.SELECT, + required=True, + options=[FormOption(label=I18nObject(en_US="0.1", zh_Hans="0.1"), value="0.1")], + ), + ] + + # Act + secret_variables = provider_configuration.extract_secret_variables(credential_form_schemas) + + # Assert + assert len(secret_variables) == 0 + + def test_extract_secret_variables_empty_list(self, provider_configuration: ProviderConfiguration) -> None: + """Test extracting secret variables from empty credential form schemas""" + # Arrange + credential_form_schemas = [] + + # Act + secret_variables = provider_configuration.extract_secret_variables(credential_form_schemas) + + # Assert + assert len(secret_variables) == 0 + + @patch("core.entities.provider_configuration.encrypter") + def test_obfuscated_credentials_with_secret_variables( + self, mock_encrypter: MagicMock, provider_configuration: ProviderConfiguration + ) -> None: + """Test obfuscating credentials with secret variables""" + # Arrange + credentials = { + "api_key": "sk-1234567890abcdef", + "model_name": "gpt-4", + "secret_token": "secret_value_123", + "temperature": "0.7", + } + + credential_form_schemas = [ + CredentialFormSchema( + variable="api_key", + label=I18nObject(en_US="API Key", zh_Hans="API 密钥"), + type=FormType.SECRET_INPUT, + required=True, + ), + CredentialFormSchema( + variable="model_name", + label=I18nObject(en_US="Model Name", zh_Hans="模型名称"), + type=FormType.TEXT_INPUT, + required=True, + ), + CredentialFormSchema( + variable="secret_token", + label=I18nObject(en_US="Secret Token", zh_Hans="密钥令牌"), + type=FormType.SECRET_INPUT, + required=False, + ), + CredentialFormSchema( + variable="temperature", + label=I18nObject(en_US="Temperature", zh_Hans="温度"), + type=FormType.TEXT_INPUT, + required=True, + ), + ] + + mock_encrypter.obfuscated_token.side_effect = lambda x: f"***{x[-4:]}" + + # Act + obfuscated = provider_configuration.obfuscated_credentials(credentials, credential_form_schemas) + + # Assert + assert obfuscated["api_key"] == "***cdef" + assert obfuscated["model_name"] == "gpt-4" # Not obfuscated + assert obfuscated["secret_token"] == "***_123" + assert obfuscated["temperature"] == "0.7" # Not obfuscated + + # Verify encrypter was called for secret fields only + assert mock_encrypter.obfuscated_token.call_count == 2 + mock_encrypter.obfuscated_token.assert_any_call("sk-1234567890abcdef") + mock_encrypter.obfuscated_token.assert_any_call("secret_value_123") + + def test_obfuscated_credentials_no_secret_variables( + self, provider_configuration: ProviderConfiguration + ) -> None: + """Test obfuscating credentials when no secret variables exist""" + # Arrange + credentials = { + "model_name": "gpt-4", + "temperature": "0.7", + "max_tokens": "1000", + } + + credential_form_schemas = [ + CredentialFormSchema( + variable="model_name", + label=I18nObject(en_US="Model Name", zh_Hans="模型名称"), + type=FormType.TEXT_INPUT, + required=True, + ), + CredentialFormSchema( + variable="temperature", + label=I18nObject(en_US="Temperature", zh_Hans="温度"), + type=FormType.TEXT_INPUT, + required=True, + ), + CredentialFormSchema( + variable="max_tokens", + label=I18nObject(en_US="Max Tokens", zh_Hans="最大令牌数"), + type=FormType.TEXT_INPUT, + required=True, + ), + ] + + # Act + obfuscated = provider_configuration.obfuscated_credentials(credentials, credential_form_schemas) + + # Assert + assert obfuscated == credentials # No changes expected + + def test_obfuscated_credentials_empty_credentials(self, provider_configuration: ProviderConfiguration) -> None: + """Test obfuscating empty credentials""" + # Arrange + credentials = {} + credential_form_schemas = [] + + # Act + obfuscated = provider_configuration.obfuscated_credentials(credentials, credential_form_schemas) + + # Assert + assert obfuscated == {} diff --git a/api/tests/unit_tests/core/entities/test_entities_provider_entities.py b/api/tests/unit_tests/core/entities/test_provider_entities.py similarity index 100% rename from api/tests/unit_tests/core/entities/test_entities_provider_entities.py rename to api/tests/unit_tests/core/entities/test_provider_entities.py diff --git a/api/tests/unit_tests/core/helper/code_executor/jinja2/test_jinja2_sandbox.py b/api/tests/unit_tests/core/helper/code_executor/jinja2/test_jinja2_transformer.py similarity index 100% rename from api/tests/unit_tests/core/helper/code_executor/jinja2/test_jinja2_sandbox.py rename to api/tests/unit_tests/core/helper/code_executor/jinja2/test_jinja2_transformer.py diff --git a/api/tests/unit_tests/core/helper/test_moderation.py b/api/tests/unit_tests/core/helper/test_moderation.py index 063e8d5120b..474d6607ff5 100644 --- a/api/tests/unit_tests/core/helper/test_moderation.py +++ b/api/tests/unit_tests/core/helper/test_moderation.py @@ -16,7 +16,7 @@ from graphon.model_runtime.entities.provider_entities import ConfigurateMethod, from graphon.model_runtime.errors.invoke import InvokeAuthorizationError, InvokeBadRequestError from graphon.model_runtime.model_providers.model_provider_factory import ModelProviderFactory from models.provider import ProviderType -from tests.unit_tests.core.model_fixtures import make_model_config +from tests.helpers.unit.core.model_fixtures import make_model_config OPENAI_PROVIDER = "langgenius/openai/openai" ANTHROPIC_PROVIDER = "langgenius/anthropic/anthropic" diff --git a/api/tests/unit_tests/utils/test_module_import_helper_behavior.py b/api/tests/unit_tests/core/helper/test_module_import_helper.py similarity index 84% rename from api/tests/unit_tests/utils/test_module_import_helper_behavior.py rename to api/tests/unit_tests/core/helper/test_module_import_helper.py index 5165cfd1149..6dfd72854d9 100644 --- a/api/tests/unit_tests/utils/test_module_import_helper_behavior.py +++ b/api/tests/unit_tests/core/helper/test_module_import_helper.py @@ -1,9 +1,9 @@ from pathlib import Path from core.helper.module_import_helper import import_module_from_source, load_single_subclass_from_source -from tests.unit_tests.utils.module_import_fixtures.parent_class import ParentClass +from tests.helpers.unit.utils.module_import_fixtures.parent_class import ParentClass -FIXTURE_DIR = Path(__file__).parent / "module_import_fixtures" +FIXTURE_DIR = Path(__file__).parents[3] / "helpers" / "unit" / "utils" / "module_import_fixtures" def test_loading_subclass_from_source() -> None: diff --git a/api/tests/unit_tests/utils/position_helper/test_position_helper.py b/api/tests/unit_tests/core/helper/test_position_helper.py similarity index 100% rename from api/tests/unit_tests/utils/position_helper/test_position_helper.py rename to api/tests/unit_tests/core/helper/test_position_helper.py diff --git a/api/tests/unit_tests/core/helper/test_trace_id_helper.py b/api/tests/unit_tests/core/helper/test_trace_id_helper.py index f7ae926a709..b5f6d6a11aa 100644 --- a/api/tests/unit_tests/core/helper/test_trace_id_helper.py +++ b/api/tests/unit_tests/core/helper/test_trace_id_helper.py @@ -1,4 +1,8 @@ +import re +from unittest import mock + import pytest +from opentelemetry.trace import NonRecordingSpan, SpanContext, use_span from werkzeug.exceptions import BadRequest from core.helper.trace_id_helper import ( @@ -12,6 +16,16 @@ from core.helper.trace_id_helper import ( ) +def test_extract_trace_session_id_from_args_for_generator_extras(): + assert extract_trace_session_id_from_args({"trace_session_id": "session-1"}) == { + "trace_session_id": "session-1", + } + + +def test_extract_trace_session_id_from_args_missing_value_keeps_extras_clean(): + assert extract_trace_session_id_from_args({"inputs": {}}) == {} + + class DummyRequest: def __init__( self, @@ -279,3 +293,95 @@ class TestTraceIdHelper: ) assert result == {} + + +"""Tests for trace helper functions.""" + + +class TestGetSpanIdFromOtelContext: + def test_returns_none_without_span(self): + from core.helper.trace_id_helper import get_span_id_from_otel_context + + with mock.patch("opentelemetry.trace.get_current_span", return_value=None, autospec=True): + result = get_span_id_from_otel_context() + assert result is None + + def test_returns_span_id_when_available(self): + from core.helper.trace_id_helper import get_span_id_from_otel_context + + span = NonRecordingSpan(SpanContext(trace_id=1, span_id=0x051581BF3BB55C45, is_remote=False)) + with use_span(span): + result = get_span_id_from_otel_context() + assert result == "051581bf3bb55c45" + + def test_returns_none_on_exception(self): + from core.helper.trace_id_helper import get_span_id_from_otel_context + + with mock.patch("opentelemetry.trace.get_current_span", side_effect=Exception("Test error"), autospec=True): + result = get_span_id_from_otel_context() + assert result is None + + +class TestGenerateTraceparentHeader: + def test_generates_valid_format(self): + from core.helper.trace_id_helper import generate_traceparent_header + + with mock.patch("opentelemetry.trace.get_current_span", return_value=None, autospec=True): + result = generate_traceparent_header() + + assert result is not None + # Format: 00-{trace_id}-{span_id}-01 + parts = result.split("-") + assert len(parts) == 4 + assert parts[0] == "00" # version + assert len(parts[1]) == 32 # trace_id (32 hex chars) + assert len(parts[2]) == 16 # span_id (16 hex chars) + assert parts[3] == "01" # flags + + def test_uses_otel_context_when_available(self): + from core.helper.trace_id_helper import generate_traceparent_header + + span = NonRecordingSpan( + SpanContext( + trace_id=0x5B8AA5A2D2C872E8321CF37308D69DF2, + span_id=0x051581BF3BB55C45, + is_remote=False, + ) + ) + with use_span(span): + result = generate_traceparent_header() + assert result == "00-5b8aa5a2d2c872e8321cf37308d69df2-051581bf3bb55c45-01" + + def test_generates_hex_only_values(self): + from core.helper.trace_id_helper import generate_traceparent_header + + with mock.patch("opentelemetry.trace.get_current_span", return_value=None, autospec=True): + result = generate_traceparent_header() + + parts = result.split("-") + # All parts should be valid hex + assert re.match(r"^[0-9a-f]+$", parts[1]) + assert re.match(r"^[0-9a-f]+$", parts[2]) + + +class TestParseTraceparentHeader: + def test_parses_valid_traceparent(self): + from core.helper.trace_id_helper import parse_traceparent_header + + traceparent = "00-5b8aa5a2d2c872e8321cf37308d69df2-051581bf3bb55c45-01" + result = parse_traceparent_header(traceparent) + + assert result == "5b8aa5a2d2c872e8321cf37308d69df2" + + def test_returns_none_for_invalid_format(self): + from core.helper.trace_id_helper import parse_traceparent_header + + # Wrong number of parts + assert parse_traceparent_header("00-abc-def") is None + # Wrong trace_id length + assert parse_traceparent_header("00-abc-def-01") is None + + def test_returns_none_for_empty_string(self): + from core.helper.trace_id_helper import parse_traceparent_header + + assert parse_traceparent_header("") is None diff --git a/api/tests/unit_tests/core/llm_generator/output_parser/test_structured_output.py b/api/tests/unit_tests/core/llm_generator/output_parser/test_structured_output.py index 548cc63000d..fd22a8295b7 100644 --- a/api/tests/unit_tests/core/llm_generator/output_parser/test_structured_output.py +++ b/api/tests/unit_tests/core/llm_generator/output_parser/test_structured_output.py @@ -1,10 +1,13 @@ import inspect import json from collections.abc import Iterator +from decimal import Decimal from types import SimpleNamespace -from unittest.mock import patch +from typing import Literal, NotRequired +from unittest.mock import MagicMock, patch import pytest +from typing_extensions import TypedDict from core.llm_generator.output_parser.errors import OutputParserError from core.llm_generator.output_parser.structured_output import ( @@ -24,6 +27,7 @@ from graphon.model_runtime.entities.llm_entities import ( LLMResult, LLMResultChunk, LLMResultChunkDelta, + LLMResultChunkWithStructuredOutput, LLMResultWithStructuredOutput, LLMUsage, ) @@ -519,3 +523,475 @@ class TestStructuredOutput: def test_parse_structured_output_empty_string(self): with pytest.raises(OutputParserError): _parse_structured_output("") + + +class _StructuredOutputCase(TypedDict): + name: str + provider: str + model_name: str + support_structure_output: bool + stream: bool + json_schema: dict[str, object] + expected_llm_response: LLMResult | list[LLMResultChunk] | list[MagicMock] + expected_result_type: type[LLMResultWithStructuredOutput] | Literal["generator"] | None + should_raise: bool + parameter_rules: NotRequired[list[ParameterRule]] + expected_response_format: NotRequired[str] + expected_error: NotRequired[type[Exception]] + + +def _response_format_rule(*options: str) -> ParameterRule: + return ParameterRule( + name="response_format", + label=I18nObject(en_US="Response Format"), + type=ParameterType.STRING, + options=list(options), + ) + + +def create_mock_usage(prompt_tokens: int = 10, completion_tokens: int = 5) -> LLMUsage: + """Create a mock LLMUsage with all required fields""" + return LLMUsage( + prompt_tokens=prompt_tokens, + prompt_unit_price=Decimal("0.001"), + prompt_price_unit=Decimal(1), + prompt_price=Decimal(str(prompt_tokens)) * Decimal("0.001"), + completion_tokens=completion_tokens, + completion_unit_price=Decimal("0.002"), + completion_price_unit=Decimal(1), + completion_price=Decimal(str(completion_tokens)) * Decimal("0.002"), + total_tokens=prompt_tokens + completion_tokens, + total_price=Decimal(str(prompt_tokens)) * Decimal("0.001") + Decimal(str(completion_tokens)) * Decimal("0.002"), + currency="USD", + latency=1.5, + ) + + +def get_model_entity(provider: str, model_name: str, support_structure_output: bool = False) -> AIModelEntity: + """Create a mock AIModelEntity for testing""" + model_schema = MagicMock() + model_schema.model = model_name + model_schema.provider = provider + model_schema.model_type = ModelType.LLM + model_schema.model_provider = provider + model_schema.model_name = model_name + model_schema.support_structure_output = support_structure_output + model_schema.parameter_rules = list[object]() + + return model_schema + + +def get_model_instance() -> MagicMock: + """Create a mock ModelInstance for testing""" + mock_instance = MagicMock() + mock_instance.provider = "openai" + mock_instance.credentials = dict[str, object]() + return mock_instance + + +def test_structured_output_parser() -> None: + """Test cases for invoke_llm_with_structured_output function""" + + testcases: list[_StructuredOutputCase] = [ + # Test case 1: Model with native structured output support, non-streaming + { + "name": "native_structured_output_non_streaming", + "provider": "openai", + "model_name": "gpt-4o", + "support_structure_output": True, + "stream": False, + "json_schema": {"type": "object", "properties": {"name": {"type": "string"}}}, + "expected_llm_response": LLMResult( + model="gpt-4o", + message=AssistantPromptMessage(content='{"name": "test"}'), + usage=create_mock_usage(prompt_tokens=10, completion_tokens=5), + ), + "expected_result_type": LLMResultWithStructuredOutput, + "should_raise": False, + }, + # Test case 2: Model with native structured output support, streaming + { + "name": "native_structured_output_streaming", + "provider": "openai", + "model_name": "gpt-4o", + "support_structure_output": True, + "stream": True, + "json_schema": {"type": "object", "properties": {"name": {"type": "string"}}}, + "expected_llm_response": [ + LLMResultChunk( + model="gpt-4o", + prompt_messages=[UserPromptMessage(content="test")], + system_fingerprint="test", + delta=LLMResultChunkDelta( + index=0, + message=AssistantPromptMessage(content='{"name":'), + usage=create_mock_usage(prompt_tokens=10, completion_tokens=2), + ), + ), + LLMResultChunk( + model="gpt-4o", + prompt_messages=[UserPromptMessage(content="test")], + system_fingerprint="test", + delta=LLMResultChunkDelta( + index=0, + message=AssistantPromptMessage(content=' "test"}'), + usage=create_mock_usage(prompt_tokens=10, completion_tokens=3), + ), + ), + ], + "expected_result_type": "generator", + "should_raise": False, + }, + # Test case 3: Model without native structured output support, non-streaming + { + "name": "prompt_based_structured_output_non_streaming", + "provider": "anthropic", + "model_name": "claude-3-sonnet", + "support_structure_output": False, + "stream": False, + "json_schema": {"type": "object", "properties": {"answer": {"type": "string"}}}, + "expected_llm_response": LLMResult( + model="claude-3-sonnet", + message=AssistantPromptMessage(content='{"answer": "test response"}'), + usage=create_mock_usage(prompt_tokens=15, completion_tokens=8), + ), + "expected_result_type": LLMResultWithStructuredOutput, + "should_raise": False, + }, + # Test case 4: Model without native structured output support, streaming + { + "name": "prompt_based_structured_output_streaming", + "provider": "anthropic", + "model_name": "claude-3-sonnet", + "support_structure_output": False, + "stream": True, + "json_schema": {"type": "object", "properties": {"answer": {"type": "string"}}}, + "expected_llm_response": [ + LLMResultChunk( + model="claude-3-sonnet", + prompt_messages=[UserPromptMessage(content="test")], + system_fingerprint="test", + delta=LLMResultChunkDelta( + index=0, + message=AssistantPromptMessage(content='{"answer": "test'), + usage=create_mock_usage(prompt_tokens=15, completion_tokens=3), + ), + ), + LLMResultChunk( + model="claude-3-sonnet", + prompt_messages=[UserPromptMessage(content="test")], + system_fingerprint="test", + delta=LLMResultChunkDelta( + index=0, + message=AssistantPromptMessage(content=' response"}'), + usage=create_mock_usage(prompt_tokens=15, completion_tokens=5), + ), + ), + ], + "expected_result_type": "generator", + "should_raise": False, + }, + # Test case 5: Streaming with list content + { + "name": "streaming_with_list_content", + "provider": "openai", + "model_name": "gpt-4o", + "support_structure_output": True, + "stream": True, + "json_schema": {"type": "object", "properties": {"data": {"type": "string"}}}, + "expected_llm_response": [ + LLMResultChunk( + model="gpt-4o", + prompt_messages=[UserPromptMessage(content="test")], + system_fingerprint="test", + delta=LLMResultChunkDelta( + index=0, + message=AssistantPromptMessage( + content=[ + TextPromptMessageContent(data='{"data":'), + ] + ), + usage=create_mock_usage(prompt_tokens=10, completion_tokens=2), + ), + ), + LLMResultChunk( + model="gpt-4o", + prompt_messages=[UserPromptMessage(content="test")], + system_fingerprint="test", + delta=LLMResultChunkDelta( + index=0, + message=AssistantPromptMessage( + content=[ + TextPromptMessageContent(data=' "value"}'), + ] + ), + usage=create_mock_usage(prompt_tokens=10, completion_tokens=3), + ), + ), + ], + "expected_result_type": "generator", + "should_raise": False, + }, + # Test case 6: Error case - non-string LLM response content (non-streaming) + { + "name": "error_non_string_content_non_streaming", + "provider": "openai", + "model_name": "gpt-4o", + "support_structure_output": True, + "stream": False, + "json_schema": {"type": "object", "properties": {"name": {"type": "string"}}}, + "expected_llm_response": LLMResult( + model="gpt-4o", + message=AssistantPromptMessage(content=None), # Non-string content + usage=create_mock_usage(prompt_tokens=10, completion_tokens=5), + ), + "expected_result_type": None, + "should_raise": True, + "expected_error": OutputParserError, + }, + # Test case 7: JSON repair scenario + { + "name": "json_repair_scenario", + "provider": "openai", + "model_name": "gpt-4o", + "support_structure_output": True, + "stream": False, + "json_schema": {"type": "object", "properties": {"name": {"type": "string"}}}, + "expected_llm_response": LLMResult( + model="gpt-4o", + message=AssistantPromptMessage(content='{"name": "test"'), # Invalid JSON - missing closing brace + usage=create_mock_usage(prompt_tokens=10, completion_tokens=5), + ), + "expected_result_type": LLMResultWithStructuredOutput, + "should_raise": False, + }, + # Test case 8: Model with parameter rules for response format + { + "name": "model_with_parameter_rules", + "provider": "openai", + "model_name": "gpt-4o", + "support_structure_output": True, + "stream": False, + "json_schema": {"type": "object", "properties": {"result": {"type": "string"}}}, + "parameter_rules": [_response_format_rule("json_schema")], + "expected_response_format": "json_schema", + "expected_llm_response": LLMResult( + model="gpt-4o", + message=AssistantPromptMessage(content='{"result": "success"}'), + usage=create_mock_usage(prompt_tokens=10, completion_tokens=5), + ), + "expected_result_type": LLMResultWithStructuredOutput, + "should_raise": False, + }, + # Test case 9: Model without native support but with JSON response format rules + { + "name": "non_native_with_json_rules", + "provider": "anthropic", + "model_name": "claude-3-sonnet", + "support_structure_output": False, + "stream": False, + "json_schema": {"type": "object", "properties": {"output": {"type": "string"}}}, + "parameter_rules": [_response_format_rule("JSON")], + "expected_response_format": "JSON", + "expected_llm_response": LLMResult( + model="claude-3-sonnet", + message=AssistantPromptMessage(content='{"output": "result"}'), + usage=create_mock_usage(prompt_tokens=15, completion_tokens=8), + ), + "expected_result_type": LLMResultWithStructuredOutput, + "should_raise": False, + }, + ] + + for case in testcases: + # Setup model entity + model_schema = get_model_entity(case["provider"], case["model_name"], case["support_structure_output"]) + + # Add parameter rules if specified + if "parameter_rules" in case: + model_schema.parameter_rules = case["parameter_rules"] + + # Setup model instance + model_instance = get_model_instance() + model_instance.invoke_llm.return_value = case["expected_llm_response"] + + # Setup prompt messages + prompt_messages = [ + SystemPromptMessage(content="You are a helpful assistant."), + UserPromptMessage(content="Generate a response according to the schema."), + ] + + if case["should_raise"]: + # Test error cases + expected_error = case.get("expected_error") + assert expected_error is not None + with pytest.raises(expected_error): # noqa: PT012 + if case["stream"]: + result_generator = invoke_llm_with_structured_output( + provider=case["provider"], + model_schema=model_schema, + model_instance=model_instance, + prompt_messages=prompt_messages, + json_schema=case["json_schema"], + stream=case["stream"], + ) + # Consume the generator to trigger the error + list(result_generator) + else: + invoke_llm_with_structured_output( + provider=case["provider"], + model_schema=model_schema, + model_instance=model_instance, + prompt_messages=prompt_messages, + json_schema=case["json_schema"], + stream=case["stream"], + ) + else: + # Test successful cases + with patch( + "core.llm_generator.output_parser.structured_output.json_repair.loads", autospec=True + ) as mock_json_repair: + # Configure json_repair mock for cases that need it + if case["name"] == "json_repair_scenario": + mock_json_repair.return_value = {"name": "test"} + + result = invoke_llm_with_structured_output( + provider=case["provider"], + model_schema=model_schema, + model_instance=model_instance, + prompt_messages=prompt_messages, + json_schema=case["json_schema"], + stream=case["stream"], + model_parameters={"temperature": 0.7, "max_tokens": 100}, + ) + + if case["expected_result_type"] == "generator": + # Test streaming results + assert hasattr(result, "__iter__") + chunks = list(result) + assert len(chunks) > 0 + + # Verify all chunks are LLMResultChunkWithStructuredOutput + for chunk in chunks[:-1]: # All except last + assert isinstance(chunk, LLMResultChunkWithStructuredOutput) + assert chunk.model == case["model_name"] + + # Last chunk should have structured output + last_chunk = chunks[-1] + assert isinstance(last_chunk, LLMResultChunkWithStructuredOutput) + assert last_chunk.structured_output is not None + assert isinstance(last_chunk.structured_output, dict) + else: + # Test non-streaming results + expected_type = case["expected_result_type"] + assert isinstance(expected_type, type) + assert isinstance(result, expected_type) + assert result.model == case["model_name"] + assert result.structured_output is not None + assert isinstance(result.structured_output, dict) + + # Verify model_instance.invoke_llm was called with correct parameters + model_instance.invoke_llm.assert_called_once() + call_args = model_instance.invoke_llm.call_args + + assert call_args.kwargs["stream"] == case["stream"] + assert "user" not in call_args.kwargs + assert "temperature" in call_args.kwargs["model_parameters"] + if "expected_response_format" in case: + assert call_args.kwargs["model_parameters"]["response_format"] == case["expected_response_format"] + assert "max_tokens" in call_args.kwargs["model_parameters"] + + +def test_parse_structured_output_edge_cases() -> None: + """Test edge cases for structured output parsing""" + + # Test case with list that contains dict (reasoning model scenario) + testcase_list_with_dict = { + "name": "list_with_dict_parsing", + "provider": "deepseek", + "model_name": "deepseek-r1", + "support_structure_output": False, + "stream": False, + "json_schema": {"type": "object", "properties": {"thought": {"type": "string"}}}, + "expected_llm_response": LLMResult( + model="deepseek-r1", + message=AssistantPromptMessage(content='[{"thought": "reasoning process"}, "other content"]'), + usage=create_mock_usage(prompt_tokens=10, completion_tokens=5), + ), + "expected_result_type": LLMResultWithStructuredOutput, + "should_raise": False, + } + + # Setup for list parsing test + model_schema = get_model_entity( + testcase_list_with_dict["provider"], + testcase_list_with_dict["model_name"], + testcase_list_with_dict["support_structure_output"], + ) + + model_instance = get_model_instance() + model_instance.invoke_llm.return_value = testcase_list_with_dict["expected_llm_response"] + + prompt_messages = [UserPromptMessage(content="Test reasoning")] + + with patch( + "core.llm_generator.output_parser.structured_output.json_repair.loads", autospec=True + ) as mock_json_repair: + # Mock json_repair to return a list with dict + mock_json_repair.return_value = [{"thought": "reasoning process"}, "other content"] + + result = invoke_llm_with_structured_output( + provider=testcase_list_with_dict["provider"], + model_schema=model_schema, + model_instance=model_instance, + prompt_messages=prompt_messages, + json_schema=testcase_list_with_dict["json_schema"], + stream=testcase_list_with_dict["stream"], + ) + + assert isinstance(result, LLMResultWithStructuredOutput) + assert result.structured_output == {"thought": "reasoning process"} + + +def test_model_specific_schema_preparation() -> None: + """Test schema preparation for different model types""" + + # Test Gemini model + gemini_case = { + "provider": "google", + "model_name": "gemini-pro", + "support_structure_output": True, + "stream": False, + "json_schema": {"type": "object", "properties": {"result": {"type": "boolean"}}, "additionalProperties": False}, + } + + model_schema = get_model_entity( + gemini_case["provider"], gemini_case["model_name"], gemini_case["support_structure_output"] + ) + + model_instance = get_model_instance() + model_instance.invoke_llm.return_value = LLMResult( + model="gemini-pro", + message=AssistantPromptMessage(content='{"result": "true"}'), + usage=create_mock_usage(prompt_tokens=10, completion_tokens=5), + ) + + prompt_messages = [UserPromptMessage(content="Test")] + + result = invoke_llm_with_structured_output( + provider=gemini_case["provider"], + model_schema=model_schema, + model_instance=model_instance, + prompt_messages=prompt_messages, + json_schema=gemini_case["json_schema"], + stream=gemini_case["stream"], + ) + + assert isinstance(result, LLMResultWithStructuredOutput) + + # Verify model_instance.invoke_llm was called and check the schema preparation + model_instance.invoke_llm.assert_called_once() + call_args = model_instance.invoke_llm.call_args + + # For Gemini, the schema should not have additionalProperties and boolean should be converted to string + assert "json_schema" in call_args.kwargs["model_parameters"] diff --git a/api/tests/unit_tests/core/llm_generator/test_llm_generator.py b/api/tests/unit_tests/core/llm_generator/test_llm_generator.py index eb469ddb6bb..71ca06b7010 100644 --- a/api/tests/unit_tests/core/llm_generator/test_llm_generator.py +++ b/api/tests/unit_tests/core/llm_generator/test_llm_generator.py @@ -1,30 +1,34 @@ -"""Tests for LLM generation and database-backed instruction modification.""" - import json from collections.abc import Iterator -from datetime import datetime +from datetime import datetime, timedelta from decimal import Decimal from types import SimpleNamespace from unittest.mock import MagicMock, Mock, patch from uuid import uuid4 import pytest +from sqlalchemy import event +from sqlalchemy.exc import SQLAlchemyError from sqlalchemy.orm import Session, scoped_session, sessionmaker +import core.llm_generator.llm_generator as generator_module from core.app.app_config.entities import ModelConfig from core.app.entities.app_invoke_entities import CreditUsageCreatedBy from core.llm_generator import llm_generator as llm_generator_module from core.llm_generator.entities import RuleCodeGeneratePayload, RuleGeneratePayload, RuleStructuredOutputPayload -from core.llm_generator.llm_generator import LLMGenerator +from core.llm_generator.llm_generator import LLMGenerator, _parse_string_list from core.model_context import get_credit_usage_metadata from core.model_manager import ModelManager from core.plugin.impl.base import _get_plugin_daemon_request_timeout from core.plugin.impl.model_runtime_factory import create_plugin_model_manager +from core.workflow.generator import tool_catalogue as tool_catalogue_module +from core.workflow.generator.tool_catalogue import ToolCatalogueEntry from graphon.enums import WorkflowNodeExecutionStatus from graphon.model_runtime.entities.llm_entities import LLMMode, LLMResult, LLMUsage from graphon.model_runtime.entities.message_entities import AssistantPromptMessage from graphon.model_runtime.entities.model_entities import ModelType, ParameterType from graphon.model_runtime.errors.invoke import InvokeAuthorizationError, InvokeError +from models.dataset import Dataset from models.enums import ConversationFromSource, CreatorUserRole from models.model import App, AppMode, Message from models.workflow import ( @@ -33,30 +37,10 @@ from models.workflow import ( WorkflowNodeExecutionTriggeredFrom, ) from services.workflow_service import WorkflowService -from tests.unit_tests.core.model_fixtures import make_model_config, make_model_instance +from tests.helpers.unit.core.model_fixtures import make_model_config, make_model_instance +from tests.helpers.unit.model_factories import make_dataset - -@pytest.fixture -def database(sqlite_session: Session, monkeypatch: pytest.MonkeyPatch) -> Iterator[Session]: - """Bind the shared SQLite session to the generator's scoped-session interface.""" - - registry = scoped_session(lambda: sqlite_session) - monkeypatch.setattr(llm_generator_module.db, "session", registry) - try: - yield sqlite_session - finally: - registry.remove() - - -@pytest.fixture -def llm_invocation(monkeypatch: pytest.MonkeyPatch) -> Mock: - model_instance = make_model_instance(provider="openai", model="gpt-4") - invocation = Mock(return_value=_llm_result('{"modified": "workflow"}')) - monkeypatch.setattr(model_instance, "invoke_llm", invocation) - model_manager = create_plugin_model_manager(tenant_id="tenant") - monkeypatch.setattr(model_manager, "get_model_instance", Mock(return_value=model_instance)) - monkeypatch.setattr(llm_generator_module.ModelManager, "for_tenant", Mock(return_value=model_manager)) - return invocation +"""Tests for LLM generation and database-backed instruction modification.""" def _llm_result(content: str) -> LLMResult: @@ -171,676 +155,740 @@ def _persist_node_execution( return execution -class TestLLMGenerator: +class TestLlmGeneratorCases: + @staticmethod @pytest.fixture - def mock_model_instance(self): - with patch("core.llm_generator.llm_generator.ModelManager.for_tenant") as mock_manager: - instance = MagicMock() - mock_manager.return_value.get_default_model_instance.return_value = instance - mock_manager.return_value.get_model_instance.return_value = instance - yield instance + def database(sqlite_session: Session, monkeypatch: pytest.MonkeyPatch) -> Iterator[Session]: + """Bind the shared SQLite session to the generator's scoped-session interface.""" + registry = scoped_session(lambda: sqlite_session) + monkeypatch.setattr(llm_generator_module.db, "session", registry) + try: + yield sqlite_session + finally: + registry.remove() + + @staticmethod @pytest.fixture - def model_config_entity(self): - return ModelConfig(provider="openai", name="gpt-4", mode=LLMMode.CHAT, completion_params={"temperature": 0.7}) + def llm_invocation(monkeypatch: pytest.MonkeyPatch) -> Mock: + model_instance = make_model_instance(provider="openai", model="gpt-4") + invocation = Mock(return_value=_llm_result('{"modified": "workflow"}')) + monkeypatch.setattr(model_instance, "invoke_llm", invocation) + model_manager = create_plugin_model_manager(tenant_id="tenant") + monkeypatch.setattr(model_manager, "get_model_instance", Mock(return_value=model_instance)) + monkeypatch.setattr(llm_generator_module.ModelManager, "for_tenant", Mock(return_value=model_manager)) + return invocation - def test_generate_conversation_name_success(self, mock_model_instance): - mock_response = MagicMock() - mock_response.message.get_text_content.return_value = json.dumps({"Your Output": "Test Conversation Name"}) - mock_model_instance.invoke_llm.return_value = mock_response + class TestLLMGenerator: + @pytest.fixture + def mock_model_instance(self): + with patch("core.llm_generator.llm_generator.ModelManager.for_tenant") as mock_manager: + instance = MagicMock() + mock_manager.return_value.get_default_model_instance.return_value = instance + mock_manager.return_value.get_model_instance.return_value = instance + yield instance - with patch("core.llm_generator.llm_generator.TraceQueueManager") as mock_trace: - name = LLMGenerator.generate_conversation_name( - "tenant_id", "test query", "conversation-1", "app-1", message_id="message-1" + @pytest.fixture + def model_config_entity(self): + return ModelConfig( + provider="openai", name="gpt-4", mode=LLMMode.CHAT, completion_params={"temperature": 0.7} ) - assert name == "Test Conversation Name" - mock_trace.assert_called_once_with(app_id="app-1") - trace_task = mock_trace.return_value.add_trace_task.call_args.args[0] - assert trace_task.message_id == "message-1" - def test_generate_conversation_name_truncated(self, mock_model_instance): - long_query = "a" * 2100 - mock_response = MagicMock() - mock_response.message.get_text_content.return_value = json.dumps({"Your Output": "Short Name"}) - mock_model_instance.invoke_llm.return_value = mock_response + def test_generate_conversation_name_success(self, mock_model_instance): + mock_response = MagicMock() + mock_response.message.get_text_content.return_value = json.dumps({"Your Output": "Test Conversation Name"}) + mock_model_instance.invoke_llm.return_value = mock_response - with patch("core.llm_generator.llm_generator.TraceQueueManager"): - name = LLMGenerator.generate_conversation_name("tenant_id", long_query) - assert name == "Short Name" + with patch("core.llm_generator.llm_generator.TraceQueueManager") as mock_trace: + name = LLMGenerator.generate_conversation_name( + "tenant_id", "test query", "conversation-1", "app-1", message_id="message-1" + ) + assert name == "Test Conversation Name" + mock_trace.assert_called_once_with(app_id="app-1") + trace_task = mock_trace.return_value.add_trace_task.call_args.args[0] + assert trace_task.message_id == "message-1" - def test_generate_conversation_name_empty_answer(self, mock_model_instance): - mock_response = MagicMock() - mock_response.message.get_text_content.return_value = "" - mock_model_instance.invoke_llm.return_value = mock_response + def test_generate_conversation_name_truncated(self, mock_model_instance): + long_query = "a" * 2100 + mock_response = MagicMock() + mock_response.message.get_text_content.return_value = json.dumps({"Your Output": "Short Name"}) + mock_model_instance.invoke_llm.return_value = mock_response - name = LLMGenerator.generate_conversation_name("tenant_id", "test query") - assert name == "test query" + with patch("core.llm_generator.llm_generator.TraceQueueManager"): + name = LLMGenerator.generate_conversation_name("tenant_id", long_query) + assert name == "Short Name" - def test_generate_conversation_name_json_repair(self, mock_model_instance): - mock_response = MagicMock() - # Invalid JSON that json_repair can fix - mock_response.message.get_text_content.return_value = "{'Your Output': 'Repaired Name'}" - mock_model_instance.invoke_llm.return_value = mock_response + def test_generate_conversation_name_empty_answer(self, mock_model_instance): + mock_response = MagicMock() + mock_response.message.get_text_content.return_value = "" + mock_model_instance.invoke_llm.return_value = mock_response - with patch("core.llm_generator.llm_generator.TraceQueueManager"): - name = LLMGenerator.generate_conversation_name("tenant_id", "test query") - assert name == "Repaired Name" - - def test_generate_conversation_name_not_dict_result(self, mock_model_instance): - mock_response = MagicMock() - mock_response.message.get_text_content.return_value = '["not a dict"]' - mock_model_instance.invoke_llm.return_value = mock_response - with patch("core.llm_generator.llm_generator.TraceQueueManager"): name = LLMGenerator.generate_conversation_name("tenant_id", "test query") assert name == "test query" - def test_generate_conversation_name_no_output_in_dict(self, mock_model_instance): - mock_response = MagicMock() - mock_response.message.get_text_content.return_value = '{"something": "else"}' - mock_model_instance.invoke_llm.return_value = mock_response - with patch("core.llm_generator.llm_generator.TraceQueueManager"): - name = LLMGenerator.generate_conversation_name("tenant_id", "test query") - assert name == "test query" + def test_generate_conversation_name_json_repair(self, mock_model_instance): + mock_response = MagicMock() + # Invalid JSON that json_repair can fix + mock_response.message.get_text_content.return_value = "{'Your Output': 'Repaired Name'}" + mock_model_instance.invoke_llm.return_value = mock_response - def test_generate_conversation_name_long_output(self, mock_model_instance): - long_output = "a" * 100 - mock_response = MagicMock() - mock_response.message.get_text_content.return_value = json.dumps({"Your Output": long_output}) - mock_model_instance.invoke_llm.return_value = mock_response + with patch("core.llm_generator.llm_generator.TraceQueueManager"): + name = LLMGenerator.generate_conversation_name("tenant_id", "test query") + assert name == "Repaired Name" - with patch("core.llm_generator.llm_generator.TraceQueueManager"): - name = LLMGenerator.generate_conversation_name("tenant_id", "test query") - assert len(name) == 78 # 75 + "..." - assert name.endswith("...") + def test_generate_conversation_name_not_dict_result(self, mock_model_instance): + mock_response = MagicMock() + mock_response.message.get_text_content.return_value = '["not a dict"]' + mock_model_instance.invoke_llm.return_value = mock_response + with patch("core.llm_generator.llm_generator.TraceQueueManager"): + name = LLMGenerator.generate_conversation_name("tenant_id", "test query") + assert name == "test query" - def test_generate_suggested_questions_after_answer_success(self, mock_model_instance): - mock_response = MagicMock() - mock_response.message.get_text_content.return_value = '["Question 1?", "Question 2?"]' - mock_model_instance.invoke_llm.return_value = mock_response + def test_generate_conversation_name_no_output_in_dict(self, mock_model_instance): + mock_response = MagicMock() + mock_response.message.get_text_content.return_value = '{"something": "else"}' + mock_model_instance.invoke_llm.return_value = mock_response + with patch("core.llm_generator.llm_generator.TraceQueueManager"): + name = LLMGenerator.generate_conversation_name("tenant_id", "test query") + assert name == "test query" - questions = LLMGenerator.generate_suggested_questions_after_answer("tenant_id", "histories") - assert len(questions) == 2 - assert questions[0] == "Question 1?" - assert mock_model_instance.invoke_llm.call_args.kwargs["model_parameters"] == { - "max_tokens": 256, - "temperature": 0.0, - } + def test_generate_conversation_name_long_output(self, mock_model_instance): + long_output = "a" * 100 + mock_response = MagicMock() + mock_response.message.get_text_content.return_value = json.dumps({"Your Output": long_output}) + mock_model_instance.invoke_llm.return_value = mock_response - def test_generate_suggested_questions_after_answer_uses_lowest_reasoning_effort(self, mock_model_instance): - mock_response = MagicMock() - mock_response.message.get_text_content.return_value = '["Question 1?"]' - mock_model_instance.invoke_llm.return_value = mock_response - mock_model_instance.get_model_schema.return_value.parameter_rules = [ - SimpleNamespace( - name="reasoning_effort", - type=ParameterType.STRING, - options=["minimal", "low", "medium", "high"], - ) - ] + with patch("core.llm_generator.llm_generator.TraceQueueManager"): + name = LLMGenerator.generate_conversation_name("tenant_id", "test query") + assert len(name) == 78 # 75 + "..." + assert name.endswith("...") - questions = LLMGenerator.generate_suggested_questions_after_answer("tenant_id", "histories") + def test_generate_suggested_questions_after_answer_success(self, mock_model_instance): + mock_response = MagicMock() + mock_response.message.get_text_content.return_value = '["Question 1?", "Question 2?"]' + mock_model_instance.invoke_llm.return_value = mock_response - assert questions == ["Question 1?"] - assert mock_model_instance.invoke_llm.call_args.kwargs["model_parameters"] == { - "max_tokens": 256, - "temperature": 0.0, - "reasoning_effort": "minimal", - } + questions = LLMGenerator.generate_suggested_questions_after_answer("tenant_id", "histories") + assert len(questions) == 2 + assert questions[0] == "Question 1?" + assert mock_model_instance.invoke_llm.call_args.kwargs["model_parameters"] == { + "max_tokens": 256, + "temperature": 0.0, + } - def test_generate_suggested_questions_after_answer_uses_defaults_when_schema_lookup_fails( - self, mock_model_instance - ): - mock_response = MagicMock() - mock_response.message.get_text_content.return_value = '["Question 1?"]' - mock_model_instance.invoke_llm.return_value = mock_response - mock_model_instance.get_model_schema.side_effect = ValueError("schema unavailable") + def test_generate_suggested_questions_after_answer_uses_lowest_reasoning_effort(self, mock_model_instance): + mock_response = MagicMock() + mock_response.message.get_text_content.return_value = '["Question 1?"]' + mock_model_instance.invoke_llm.return_value = mock_response + mock_model_instance.get_model_schema.return_value.parameter_rules = [ + SimpleNamespace( + name="reasoning_effort", + type=ParameterType.STRING, + options=["minimal", "low", "medium", "high"], + ) + ] - questions = LLMGenerator.generate_suggested_questions_after_answer("tenant_id", "histories") + questions = LLMGenerator.generate_suggested_questions_after_answer("tenant_id", "histories") - assert questions == ["Question 1?"] - assert mock_model_instance.invoke_llm.call_args.kwargs["model_parameters"] == { - "max_tokens": 256, - "temperature": 0.0, - } + assert questions == ["Question 1?"] + assert mock_model_instance.invoke_llm.call_args.kwargs["model_parameters"] == { + "max_tokens": 256, + "temperature": 0.0, + "reasoning_effort": "minimal", + } - @pytest.mark.parametrize( - ("parameter_type", "options", "expected_value"), - [ - (ParameterType.BOOLEAN, [], False), - (ParameterType.STRING, ["enabled", "disabled"], "disabled"), - ], - ) - def test_generate_suggested_questions_after_answer_disables_thinking( - self, mock_model_instance, parameter_type, options, expected_value - ): - mock_response = MagicMock() - mock_response.message.get_text_content.return_value = '["Question 1?"]' - mock_model_instance.invoke_llm.return_value = mock_response - mock_model_instance.get_model_schema.return_value.parameter_rules = [ - SimpleNamespace(name="thinking", type=parameter_type, options=options), - SimpleNamespace(name="reasoning_effort", type=ParameterType.STRING, options=["low", "high"]), - ] + def test_generate_suggested_questions_after_answer_uses_defaults_when_schema_lookup_fails( + self, mock_model_instance + ): + mock_response = MagicMock() + mock_response.message.get_text_content.return_value = '["Question 1?"]' + mock_model_instance.invoke_llm.return_value = mock_response + mock_model_instance.get_model_schema.side_effect = ValueError("schema unavailable") - questions = LLMGenerator.generate_suggested_questions_after_answer("tenant_id", "histories") + questions = LLMGenerator.generate_suggested_questions_after_answer("tenant_id", "histories") - assert questions == ["Question 1?"] - assert mock_model_instance.invoke_llm.call_args.kwargs["model_parameters"] == { - "max_tokens": 256, - "temperature": 0.0, - "thinking": expected_value, - } + assert questions == ["Question 1?"] + assert mock_model_instance.invoke_llm.call_args.kwargs["model_parameters"] == { + "max_tokens": 256, + "temperature": 0.0, + } - def test_generate_suggested_questions_after_answer_auth_error(self, mock_model_instance): - with patch("core.llm_generator.llm_generator.ModelManager.for_tenant") as mock_manager: - mock_manager.return_value.get_default_model_instance.side_effect = InvokeAuthorizationError("Auth failed") + @pytest.mark.parametrize( + ("parameter_type", "options", "expected_value"), + [ + (ParameterType.BOOLEAN, [], False), + (ParameterType.STRING, ["enabled", "disabled"], "disabled"), + ], + ) + def test_generate_suggested_questions_after_answer_disables_thinking( + self, mock_model_instance, parameter_type, options, expected_value + ): + mock_response = MagicMock() + mock_response.message.get_text_content.return_value = '["Question 1?"]' + mock_model_instance.invoke_llm.return_value = mock_response + mock_model_instance.get_model_schema.return_value.parameter_rules = [ + SimpleNamespace(name="thinking", type=parameter_type, options=options), + SimpleNamespace(name="reasoning_effort", type=ParameterType.STRING, options=["low", "high"]), + ] + + questions = LLMGenerator.generate_suggested_questions_after_answer("tenant_id", "histories") + + assert questions == ["Question 1?"] + assert mock_model_instance.invoke_llm.call_args.kwargs["model_parameters"] == { + "max_tokens": 256, + "temperature": 0.0, + "thinking": expected_value, + } + + def test_generate_suggested_questions_after_answer_auth_error(self, mock_model_instance): + with patch("core.llm_generator.llm_generator.ModelManager.for_tenant") as mock_manager: + mock_manager.return_value.get_default_model_instance.side_effect = InvokeAuthorizationError( + "Auth failed" + ) + questions = LLMGenerator.generate_suggested_questions_after_answer("tenant_id", "histories") + assert questions == [] + + def test_generate_suggested_questions_after_answer_model_resolution_error(self, mock_model_instance): + with patch("core.llm_generator.llm_generator.ModelManager.for_tenant") as mock_manager: + mock_manager.return_value.get_default_model_instance.side_effect = ValueError( + "unsupported default model" + ) + + questions = LLMGenerator.generate_suggested_questions_after_answer("tenant_id", "histories") + + assert questions == [] + + def test_generate_suggested_questions_after_answer_invoke_error(self, mock_model_instance): + mock_model_instance.invoke_llm.side_effect = InvokeError("Invoke failed") questions = LLMGenerator.generate_suggested_questions_after_answer("tenant_id", "histories") assert questions == [] - def test_generate_suggested_questions_after_answer_model_resolution_error(self, mock_model_instance): - with patch("core.llm_generator.llm_generator.ModelManager.for_tenant") as mock_manager: - mock_manager.return_value.get_default_model_instance.side_effect = ValueError("unsupported default model") - + def test_generate_suggested_questions_after_answer_exception(self, mock_model_instance): + mock_model_instance.invoke_llm.side_effect = Exception("Random error") questions = LLMGenerator.generate_suggested_questions_after_answer("tenant_id", "histories") + assert questions == [] - assert questions == [] + @patch("core.llm_generator.llm_generator.ModelManager.for_tenant") + def test_generate_suggested_questions_after_answer_with_custom_model_and_prompt(self, mock_for_tenant): + custom_model_instance = MagicMock() + custom_response = MagicMock() + custom_response.message.get_text_content.return_value = '["Question 1?"]' + custom_model_instance.invoke_llm.return_value = custom_response - def test_generate_suggested_questions_after_answer_invoke_error(self, mock_model_instance): - mock_model_instance.invoke_llm.side_effect = InvokeError("Invoke failed") - questions = LLMGenerator.generate_suggested_questions_after_answer("tenant_id", "histories") - assert questions == [] + mock_for_tenant.return_value.get_model_instance.return_value = custom_model_instance - def test_generate_suggested_questions_after_answer_exception(self, mock_model_instance): - mock_model_instance.invoke_llm.side_effect = Exception("Random error") - questions = LLMGenerator.generate_suggested_questions_after_answer("tenant_id", "histories") - assert questions == [] - - @patch("core.llm_generator.llm_generator.ModelManager.for_tenant") - def test_generate_suggested_questions_after_answer_with_custom_model_and_prompt(self, mock_for_tenant): - custom_model_instance = MagicMock() - custom_response = MagicMock() - custom_response.message.get_text_content.return_value = '["Question 1?"]' - custom_model_instance.invoke_llm.return_value = custom_response - - mock_for_tenant.return_value.get_model_instance.return_value = custom_model_instance - - questions = LLMGenerator.generate_suggested_questions_after_answer( - "tenant_id", - "histories", - instruction_prompt="custom prompt", - model_config={ - "provider": "openai", - "name": "gpt-4o", - "completion_params": {"temperature": 0.2}, - }, - ) - - assert questions == ["Question 1?"] - mock_for_tenant.return_value.get_model_instance.assert_called_once_with( - tenant_id="tenant_id", - model_type=ModelType.LLM, - provider="openai", - model="gpt-4o", - ) - - invoke_kwargs = custom_model_instance.invoke_llm.call_args.kwargs - assert invoke_kwargs["model_parameters"] == {"temperature": 0.2} - assert invoke_kwargs["stop"] == [] - assert "custom prompt" in invoke_kwargs["prompt_messages"][0].content - - @patch("core.llm_generator.llm_generator.ModelManager.for_tenant") - def test_generate_suggested_questions_after_answer_with_custom_model_without_completion_params( - self, mock_for_tenant - ): - custom_model_instance = MagicMock() - custom_response = MagicMock() - custom_response.message.get_text_content.return_value = '["Question 1?"]' - custom_model_instance.invoke_llm.return_value = custom_response - mock_for_tenant.return_value.get_model_instance.return_value = custom_model_instance - - questions = LLMGenerator.generate_suggested_questions_after_answer( - "tenant_id", - "histories", - model_config={"provider": "openai", "name": "gpt-4o"}, - ) - - assert questions == ["Question 1?"] - invoke_kwargs = custom_model_instance.invoke_llm.call_args.kwargs - assert invoke_kwargs["model_parameters"] == {} - assert invoke_kwargs["stop"] == [] - - @patch("core.llm_generator.llm_generator.ModelManager.for_tenant") - def test_generate_suggested_questions_after_answer_fallback_to_default_model(self, mock_for_tenant): - default_model_instance = MagicMock() - default_response = MagicMock() - default_response.message.get_text_content.return_value = '["Question 1?"]' - default_model_instance.invoke_llm.return_value = default_response - - mock_for_tenant.return_value.get_model_instance.side_effect = ValueError("invalid configured model") - mock_for_tenant.return_value.get_default_model_instance.return_value = default_model_instance - - questions = LLMGenerator.generate_suggested_questions_after_answer( - "tenant_id", - "histories", - model_config={ - "provider": "openai", - "name": "not-found-model", - "completion_params": {"temperature": 0.2}, - }, - ) - - assert questions == ["Question 1?"] - mock_for_tenant.return_value.get_default_model_instance.assert_called_once_with( - tenant_id="tenant_id", - model_type=ModelType.LLM, - ) - assert default_model_instance.invoke_llm.call_args.kwargs["model_parameters"] == { - "max_tokens": 256, - "temperature": 0.0, - } - assert default_model_instance.invoke_llm.call_args.kwargs["stop"] == [] - - @patch("core.llm_generator.llm_generator.ModelManager.for_tenant") - def test_generate_suggested_questions_after_answer_drops_non_positive_max_tokens(self, mock_for_tenant): - custom_model_instance = MagicMock() - custom_response = MagicMock() - custom_response.message.get_text_content.return_value = '["Question 1?"]' - custom_model_instance.invoke_llm.return_value = custom_response - mock_for_tenant.return_value.get_model_instance.return_value = custom_model_instance - - questions = LLMGenerator.generate_suggested_questions_after_answer( - "tenant_id", - "histories", - model_config={ - "provider": "openai", - "name": "gpt-4o", - "completion_params": { - "temperature": 0.2, - "max_tokens": 0, - "stop": ["END"], + questions = LLMGenerator.generate_suggested_questions_after_answer( + "tenant_id", + "histories", + instruction_prompt="custom prompt", + model_config={ + "provider": "openai", + "name": "gpt-4o", + "completion_params": {"temperature": 0.2}, }, - }, - ) + ) - assert questions == ["Question 1?"] - invoke_kwargs = custom_model_instance.invoke_llm.call_args.kwargs - assert invoke_kwargs["model_parameters"] == {"temperature": 0.2} - assert invoke_kwargs["stop"] == ["END"] + assert questions == ["Question 1?"] + mock_for_tenant.return_value.get_model_instance.assert_called_once_with( + tenant_id="tenant_id", + model_type=ModelType.LLM, + provider="openai", + model="gpt-4o", + ) - @pytest.mark.parametrize("use_configured_model", [False, True]) - def test_prepared_suggested_questions_defer_provider_calls_without_resolving_again( - self, monkeypatch: pytest.MonkeyPatch, use_configured_model: bool - ) -> None: - model_instance = make_model_instance(provider="openai", model="custom-model") - schema = make_model_config(provider="openai", model="custom-model", mode="chat").model_schema - get_schema = Mock(return_value=schema) - invocation = Mock() - monkeypatch.setattr(model_instance, "get_model_schema", get_schema) - monkeypatch.setattr(model_instance, "invoke_llm", invocation) - manager = create_plugin_model_manager(tenant_id="tenant_id") - monkeypatch.setattr(manager, "get_model_instance", Mock(return_value=model_instance)) - monkeypatch.setattr(manager, "get_default_model_instance", Mock(return_value=model_instance)) - resolve_manager = Mock(return_value=manager) - monkeypatch.setattr(ModelManager, "for_tenant", resolve_manager) - original_metadata = get_credit_usage_metadata() - original_timeout = _get_plugin_daemon_request_timeout() + invoke_kwargs = custom_model_instance.invoke_llm.call_args.kwargs + assert invoke_kwargs["model_parameters"] == {"temperature": 0.2} + assert invoke_kwargs["stop"] == [] + assert "custom prompt" in invoke_kwargs["prompt_messages"][0].content - def invoke(**_kwargs: object) -> LLMResult: - metadata = get_credit_usage_metadata() - assert metadata is not None - assert metadata["created_by"] == CreditUsageCreatedBy.SUGGESTED_QUESTIONS - timeout = _get_plugin_daemon_request_timeout() - assert timeout is not None - assert timeout.read == 30.0 - return _llm_result('["Next question?"]') + @patch("core.llm_generator.llm_generator.ModelManager.for_tenant") + def test_generate_suggested_questions_after_answer_with_custom_model_without_completion_params( + self, mock_for_tenant + ): + custom_model_instance = MagicMock() + custom_response = MagicMock() + custom_response.message.get_text_content.return_value = '["Question 1?"]' + custom_model_instance.invoke_llm.return_value = custom_response + mock_for_tenant.return_value.get_model_instance.return_value = custom_model_instance - invocation.side_effect = invoke - model_config = ( - { - "provider": "openai", - "name": "custom-model", - "completion_params": {"temperature": 0.2, "stop": ["END"], "max_tokens": 0}, + questions = LLMGenerator.generate_suggested_questions_after_answer( + "tenant_id", + "histories", + model_config={"provider": "openai", "name": "gpt-4o"}, + ) + + assert questions == ["Question 1?"] + invoke_kwargs = custom_model_instance.invoke_llm.call_args.kwargs + assert invoke_kwargs["model_parameters"] == {} + assert invoke_kwargs["stop"] == [] + + @patch("core.llm_generator.llm_generator.ModelManager.for_tenant") + def test_generate_suggested_questions_after_answer_fallback_to_default_model(self, mock_for_tenant): + default_model_instance = MagicMock() + default_response = MagicMock() + default_response.message.get_text_content.return_value = '["Question 1?"]' + default_model_instance.invoke_llm.return_value = default_response + + mock_for_tenant.return_value.get_model_instance.side_effect = ValueError("invalid configured model") + mock_for_tenant.return_value.get_default_model_instance.return_value = default_model_instance + + questions = LLMGenerator.generate_suggested_questions_after_answer( + "tenant_id", + "histories", + model_config={ + "provider": "openai", + "name": "not-found-model", + "completion_params": {"temperature": 0.2}, + }, + ) + + assert questions == ["Question 1?"] + mock_for_tenant.return_value.get_default_model_instance.assert_called_once_with( + tenant_id="tenant_id", + model_type=ModelType.LLM, + ) + assert default_model_instance.invoke_llm.call_args.kwargs["model_parameters"] == { + "max_tokens": 256, + "temperature": 0.0, } - if use_configured_model - else None - ) - prepared_model = LLMGenerator.prepare_suggested_questions_model("tenant_id", model_config=model_config) + assert default_model_instance.invoke_llm.call_args.kwargs["stop"] == [] - assert prepared_model is not None - get_schema.assert_not_called() - invocation.assert_not_called() - resolve_manager.side_effect = AssertionError("Model lookup must finish in the preparation phase") + @patch("core.llm_generator.llm_generator.ModelManager.for_tenant") + def test_generate_suggested_questions_after_answer_drops_non_positive_max_tokens(self, mock_for_tenant): + custom_model_instance = MagicMock() + custom_response = MagicMock() + custom_response.message.get_text_content.return_value = '["Question 1?"]' + custom_model_instance.invoke_llm.return_value = custom_response + mock_for_tenant.return_value.get_model_instance.return_value = custom_model_instance - result = LLMGenerator.invoke_suggested_questions_after_answer( - prepared_model, "Human: hello\nAssistant: world", instruction_prompt="Ask a follow-up" - ) - - assert result == ["Next question?"] - parameters = invocation.call_args.kwargs - assert parameters["model_parameters"] == ( - {"temperature": 0.2} if use_configured_model else {"max_tokens": 256, "temperature": 0.0} - ) - assert parameters["stop"] == (["END"] if use_configured_model else []) - assert parameters["stream"] is False - assert "Human: hello\nAssistant: world" in parameters["prompt_messages"][0].content - assert "Ask a follow-up" in parameters["prompt_messages"][0].content - assert get_credit_usage_metadata() == original_metadata - assert _get_plugin_daemon_request_timeout() == original_timeout - - def test_generate_rule_config_no_variable_success(self, mock_model_instance, model_config_entity): - payload = RuleGeneratePayload( - instruction="test instruction", model_config=model_config_entity, no_variable=True - ) - mock_response = MagicMock() - mock_response.message.get_text_content.return_value = "Generated Prompt" - mock_model_instance.invoke_llm.return_value = mock_response - - result = LLMGenerator.generate_rule_config("tenant_id", payload) - assert result["prompt"] == "Generated Prompt" - assert result["error"] == "" - - def test_generate_rule_config_no_variable_invoke_error(self, mock_model_instance, model_config_entity): - payload = RuleGeneratePayload( - instruction="test instruction", model_config=model_config_entity, no_variable=True - ) - mock_model_instance.invoke_llm.side_effect = InvokeError("Invoke failed") - - result = LLMGenerator.generate_rule_config("tenant_id", payload) - assert "Failed to generate rule config" in result["error"] - - def test_generate_rule_config_no_variable_exception(self, mock_model_instance, model_config_entity): - payload = RuleGeneratePayload( - instruction="test instruction", model_config=model_config_entity, no_variable=True - ) - mock_model_instance.invoke_llm.side_effect = Exception("Random error") - - result = LLMGenerator.generate_rule_config("tenant_id", payload) - assert "Failed to generate rule config" in result["error"] - assert "Random error" in result["error"] - - def test_generate_rule_config_with_variable_success(self, mock_model_instance, model_config_entity): - payload = RuleGeneratePayload( - instruction="test instruction", model_config=model_config_entity, no_variable=False - ) - # Mocking 3 calls for invoke_llm - mock_res1 = MagicMock() - mock_res1.message.get_text_content.return_value = "Step 1 Prompt" - - mock_res2 = MagicMock() - mock_res2.message.get_text_content.return_value = '"var1", "var2"' - - mock_res3 = MagicMock() - mock_res3.message.get_text_content.return_value = "Opening Statement" - - mock_model_instance.invoke_llm.side_effect = [mock_res1, mock_res2, mock_res3] - - result = LLMGenerator.generate_rule_config("tenant_id", payload) - assert result["prompt"] == "Step 1 Prompt" - assert result["variables"] == ["var1", "var2"] - assert result["opening_statement"] == "Opening Statement" - assert result["error"] == "" - - def test_generate_rule_config_with_variable_step1_error(self, mock_model_instance, model_config_entity): - payload = RuleGeneratePayload( - instruction="test instruction", model_config=model_config_entity, no_variable=False - ) - mock_model_instance.invoke_llm.side_effect = InvokeError("Step 1 Failed") - - result = LLMGenerator.generate_rule_config("tenant_id", payload) - assert "Failed to generate prefix prompt" in result["error"] - - def test_generate_rule_config_with_variable_step2_error(self, mock_model_instance, model_config_entity): - payload = RuleGeneratePayload( - instruction="test instruction", model_config=model_config_entity, no_variable=False - ) - mock_res1 = MagicMock() - mock_res1.message.get_text_content.return_value = "Step 1 Prompt" - - # Step 2 fails - mock_model_instance.invoke_llm.side_effect = [mock_res1, InvokeError("Step 2 Failed"), MagicMock()] - - result = LLMGenerator.generate_rule_config("tenant_id", payload) - assert "Failed to generate variables" in result["error"] - - def test_generate_rule_config_with_variable_step3_error(self, mock_model_instance, model_config_entity): - payload = RuleGeneratePayload( - instruction="test instruction", model_config=model_config_entity, no_variable=False - ) - mock_res1 = MagicMock() - mock_res1.message.get_text_content.return_value = "Step 1 Prompt" - - mock_res2 = MagicMock() - mock_res2.message.get_text_content.return_value = '"var1"' - - # Step 3 fails - mock_model_instance.invoke_llm.side_effect = [mock_res1, mock_res2, InvokeError("Step 3 Failed")] - - result = LLMGenerator.generate_rule_config("tenant_id", payload) - assert "Failed to generate conversation opener" in result["error"] - - def test_generate_rule_config_with_variable_exception(self, mock_model_instance, model_config_entity): - payload = RuleGeneratePayload( - instruction="test instruction", model_config=model_config_entity, no_variable=False - ) - # Mock any step to throw Exception - mock_model_instance.invoke_llm.side_effect = Exception("Unexpected multi-step error") - - result = LLMGenerator.generate_rule_config("tenant_id", payload) - assert "Unexpected multi-step error" in result["error"] - - def test_generate_code_python_success(self, mock_model_instance, model_config_entity): - payload = RuleCodeGeneratePayload( - instruction="print hello", code_language="python", model_config=model_config_entity - ) - mock_response = MagicMock() - mock_response.message.get_text_content.return_value = "print('hello')" - mock_model_instance.invoke_llm.return_value = mock_response - - result = LLMGenerator.generate_code("tenant_id", payload) - assert result["code"] == "print('hello')" - assert result["language"] == "python" - - def test_generate_code_javascript_success(self, mock_model_instance, model_config_entity): - payload = RuleCodeGeneratePayload( - instruction="console log hello", code_language="javascript", model_config=model_config_entity - ) - mock_response = MagicMock() - mock_response.message.get_text_content.return_value = "console.log('hello')" - mock_model_instance.invoke_llm.return_value = mock_response - - result = LLMGenerator.generate_code("tenant_id", payload) - assert result["code"] == "console.log('hello')" - assert result["language"] == "javascript" - - def test_generate_code_invoke_error(self, mock_model_instance, model_config_entity): - payload = RuleCodeGeneratePayload(instruction="error", code_language="python", model_config=model_config_entity) - mock_model_instance.invoke_llm.side_effect = InvokeError("Invoke failed") - - result = LLMGenerator.generate_code("tenant_id", payload) - assert "Failed to generate code" in result["error"] - - def test_generate_code_exception(self, mock_model_instance, model_config_entity): - payload = RuleCodeGeneratePayload(instruction="error", code_language="python", model_config=model_config_entity) - mock_model_instance.invoke_llm.side_effect = Exception("Random error") - - result = LLMGenerator.generate_code("tenant_id", payload) - assert "Random error" in result["error"] - - def test_generate_qa_document_success(self, mock_model_instance): - mock_response = _llm_result("QA Document Content") - mock_model_instance.invoke_llm.return_value = mock_response - - result = LLMGenerator.generate_qa_document("tenant_id", "query", "English") - assert result == "QA Document Content" - - def test_generate_qa_document_type_error(self, mock_model_instance): - mock_model_instance.invoke_llm.return_value = "Not an LLMResult" - - with pytest.raises(TypeError, match="Expected LLMResult when stream=False"): - LLMGenerator.generate_qa_document("tenant_id", "query", "English") - - def test_generate_structured_output_success(self, mock_model_instance, model_config_entity): - payload = RuleStructuredOutputPayload(instruction="generate schema", model_config=model_config_entity) - mock_response = MagicMock() - mock_response.message.get_text_content.return_value = '{"type": "object", "properties": {}}' - mock_model_instance.invoke_llm.return_value = mock_response - - result = LLMGenerator.generate_structured_output("tenant_id", payload) - parsed_output = json.loads(result["output"]) - assert parsed_output["type"] == "object" - assert result["error"] == "" - - def test_generate_structured_output_json_repair(self, mock_model_instance, model_config_entity): - payload = RuleStructuredOutputPayload(instruction="generate schema", model_config=model_config_entity) - mock_response = MagicMock() - mock_response.message.get_text_content.return_value = "{'type': 'object'}" - mock_model_instance.invoke_llm.return_value = mock_response - - result = LLMGenerator.generate_structured_output("tenant_id", payload) - parsed_output = json.loads(result["output"]) - assert parsed_output["type"] == "object" - - def test_generate_structured_output_not_dict_or_list(self, mock_model_instance, model_config_entity): - payload = RuleStructuredOutputPayload(instruction="generate schema", model_config=model_config_entity) - mock_response = MagicMock() - mock_response.message.get_text_content.return_value = "true" # parsed as bool - mock_model_instance.invoke_llm.return_value = mock_response - - result = LLMGenerator.generate_structured_output("tenant_id", payload) - assert "Failed to parse structured output" in result["error"] - - def test_generate_structured_output_invoke_error(self, mock_model_instance, model_config_entity): - payload = RuleStructuredOutputPayload(instruction="error", model_config=model_config_entity) - mock_model_instance.invoke_llm.side_effect = InvokeError("Invoke failed") - - result = LLMGenerator.generate_structured_output("tenant_id", payload) - assert "Invoke failed" in result["error"] - - def test_generate_structured_output_exception(self, mock_model_instance, model_config_entity): - payload = RuleStructuredOutputPayload(instruction="error", model_config=model_config_entity) - mock_model_instance.invoke_llm.side_effect = Exception("Random error") - - result = LLMGenerator.generate_structured_output("tenant_id", payload) - assert "Random error" in result["error"] - - def test_instruction_modify_legacy_without_last_run_uses_real_empty_query( - self, - database: Session, - llm_invocation: Mock, - model_config_entity: ModelConfig, - ): - app = _persist_app(database) - llm_invocation.return_value = _llm_result('{"modified": "prompt"}') - - result = LLMGenerator.instruction_modify_legacy( - app.tenant_id, - app.id, - "current_val", - "Test {{#last_run#}} and {{#current#}} and {{#error_message#}}", - model_config_entity, - "ideal", - session=database, - ) - - assert result == {"modified": "prompt"} - user_payload = json.loads(llm_invocation.call_args.kwargs["prompt_messages"][1].content) - assert "null" in user_payload["instruction"] - assert "current_val" in user_payload["instruction"] - - def test_instruction_modify_legacy_reads_latest_tenant_scoped_message( - self, - database: Session, - llm_invocation: Mock, - model_config_entity: ModelConfig, - ): - app = _persist_app(database) - _persist_message( - database, - app, - query="older question", - answer="older answer", - created_at=datetime(2026, 1, 1), - ) - _persist_message( - database, - app, - query="latest question", - answer="latest answer", - created_at=datetime(2026, 1, 2), - ) - other_app = _persist_app(database) - _persist_message( - database, - other_app, - query="other tenant question", - created_at=datetime(2026, 1, 3), - ) - llm_invocation.return_value = _llm_result('{"modified": "prompt"}') - - result = LLMGenerator.instruction_modify_legacy( - app.tenant_id, app.id, "current", "instruction", model_config_entity, "ideal", session=database - ) - - assert result == {"modified": "prompt"} - user_payload = json.loads(llm_invocation.call_args.kwargs["prompt_messages"][1].content) - assert user_payload["last_run"]["query"] == "latest question" - assert user_payload["last_run"]["answer"] == "latest answer" - assert "older question" not in json.dumps(user_payload) - assert "other tenant question" not in json.dumps(user_payload) - - def test_instruction_modify_workflow_app_not_found( - self, - database: Session, - sqlite_session_factory: sessionmaker[Session], - model_config_entity: ModelConfig, - ): - workflow_service = WorkflowService(sqlite_session_factory) - - with pytest.raises(ValueError, match="App not found"): - LLMGenerator.instruction_modify_workflow( - str(uuid4()), - str(uuid4()), - "node", - "current", - "instruction", - model_config_entity, - "ideal", - workflow_service, - session=database, + questions = LLMGenerator.generate_suggested_questions_after_answer( + "tenant_id", + "histories", + model_config={ + "provider": "openai", + "name": "gpt-4o", + "completion_params": { + "temperature": 0.2, + "max_tokens": 0, + "stop": ["END"], + }, + }, ) - def test_instruction_modify_workflow_rejects_app_from_another_tenant( - self, - database: Session, - sqlite_session_factory: sessionmaker[Session], - model_config_entity: ModelConfig, - ): - app = _persist_app(database) - workflow_service = WorkflowService(sqlite_session_factory) + assert questions == ["Question 1?"] + invoke_kwargs = custom_model_instance.invoke_llm.call_args.kwargs + assert invoke_kwargs["model_parameters"] == {"temperature": 0.2} + assert invoke_kwargs["stop"] == ["END"] - with pytest.raises(ValueError, match="App not found"): - LLMGenerator.instruction_modify_workflow( - str(uuid4()), + @pytest.mark.parametrize("use_configured_model", [False, True]) + def test_prepared_suggested_questions_defer_provider_calls_without_resolving_again( + self, monkeypatch: pytest.MonkeyPatch, use_configured_model: bool + ) -> None: + model_instance = make_model_instance(provider="openai", model="custom-model") + schema = make_model_config(provider="openai", model="custom-model", mode="chat").model_schema + get_schema = Mock(return_value=schema) + invocation = Mock() + monkeypatch.setattr(model_instance, "get_model_schema", get_schema) + monkeypatch.setattr(model_instance, "invoke_llm", invocation) + manager = create_plugin_model_manager(tenant_id="tenant_id") + monkeypatch.setattr(manager, "get_model_instance", Mock(return_value=model_instance)) + monkeypatch.setattr(manager, "get_default_model_instance", Mock(return_value=model_instance)) + resolve_manager = Mock(return_value=manager) + monkeypatch.setattr(ModelManager, "for_tenant", resolve_manager) + original_metadata = get_credit_usage_metadata() + original_timeout = _get_plugin_daemon_request_timeout() + + def invoke(**_kwargs: object) -> LLMResult: + metadata = get_credit_usage_metadata() + assert metadata is not None + assert metadata["created_by"] == CreditUsageCreatedBy.SUGGESTED_QUESTIONS + timeout = _get_plugin_daemon_request_timeout() + assert timeout is not None + assert timeout.read == 30.0 + return _llm_result('["Next question?"]') + + invocation.side_effect = invoke + model_config = ( + { + "provider": "openai", + "name": "custom-model", + "completion_params": {"temperature": 0.2, "stop": ["END"], "max_tokens": 0}, + } + if use_configured_model + else None + ) + prepared_model = LLMGenerator.prepare_suggested_questions_model("tenant_id", model_config=model_config) + + assert prepared_model is not None + get_schema.assert_not_called() + invocation.assert_not_called() + resolve_manager.side_effect = AssertionError("Model lookup must finish in the preparation phase") + + result = LLMGenerator.invoke_suggested_questions_after_answer( + prepared_model, "Human: hello\nAssistant: world", instruction_prompt="Ask a follow-up" + ) + + assert result == ["Next question?"] + parameters = invocation.call_args.kwargs + assert parameters["model_parameters"] == ( + {"temperature": 0.2} if use_configured_model else {"max_tokens": 256, "temperature": 0.0} + ) + assert parameters["stop"] == (["END"] if use_configured_model else []) + assert parameters["stream"] is False + assert "Human: hello\nAssistant: world" in parameters["prompt_messages"][0].content + assert "Ask a follow-up" in parameters["prompt_messages"][0].content + assert get_credit_usage_metadata() == original_metadata + assert _get_plugin_daemon_request_timeout() == original_timeout + + def test_generate_rule_config_no_variable_success(self, mock_model_instance, model_config_entity): + payload = RuleGeneratePayload( + instruction="test instruction", model_config=model_config_entity, no_variable=True + ) + mock_response = MagicMock() + mock_response.message.get_text_content.return_value = "Generated Prompt" + mock_model_instance.invoke_llm.return_value = mock_response + + result = LLMGenerator.generate_rule_config("tenant_id", payload) + assert result["prompt"] == "Generated Prompt" + assert result["error"] == "" + + def test_generate_rule_config_no_variable_invoke_error(self, mock_model_instance, model_config_entity): + payload = RuleGeneratePayload( + instruction="test instruction", model_config=model_config_entity, no_variable=True + ) + mock_model_instance.invoke_llm.side_effect = InvokeError("Invoke failed") + + result = LLMGenerator.generate_rule_config("tenant_id", payload) + assert "Failed to generate rule config" in result["error"] + + def test_generate_rule_config_no_variable_exception(self, mock_model_instance, model_config_entity): + payload = RuleGeneratePayload( + instruction="test instruction", model_config=model_config_entity, no_variable=True + ) + mock_model_instance.invoke_llm.side_effect = Exception("Random error") + + result = LLMGenerator.generate_rule_config("tenant_id", payload) + assert "Failed to generate rule config" in result["error"] + assert "Random error" in result["error"] + + def test_generate_rule_config_with_variable_success(self, mock_model_instance, model_config_entity): + payload = RuleGeneratePayload( + instruction="test instruction", model_config=model_config_entity, no_variable=False + ) + # Mocking 3 calls for invoke_llm + mock_res1 = MagicMock() + mock_res1.message.get_text_content.return_value = "Step 1 Prompt" + + mock_res2 = MagicMock() + mock_res2.message.get_text_content.return_value = '"var1", "var2"' + + mock_res3 = MagicMock() + mock_res3.message.get_text_content.return_value = "Opening Statement" + + mock_model_instance.invoke_llm.side_effect = [mock_res1, mock_res2, mock_res3] + + result = LLMGenerator.generate_rule_config("tenant_id", payload) + assert result["prompt"] == "Step 1 Prompt" + assert result["variables"] == ["var1", "var2"] + assert result["opening_statement"] == "Opening Statement" + assert result["error"] == "" + + def test_generate_rule_config_with_variable_step1_error(self, mock_model_instance, model_config_entity): + payload = RuleGeneratePayload( + instruction="test instruction", model_config=model_config_entity, no_variable=False + ) + mock_model_instance.invoke_llm.side_effect = InvokeError("Step 1 Failed") + + result = LLMGenerator.generate_rule_config("tenant_id", payload) + assert "Failed to generate prefix prompt" in result["error"] + + def test_generate_rule_config_with_variable_step2_error(self, mock_model_instance, model_config_entity): + payload = RuleGeneratePayload( + instruction="test instruction", model_config=model_config_entity, no_variable=False + ) + mock_res1 = MagicMock() + mock_res1.message.get_text_content.return_value = "Step 1 Prompt" + + # Step 2 fails + mock_model_instance.invoke_llm.side_effect = [mock_res1, InvokeError("Step 2 Failed"), MagicMock()] + + result = LLMGenerator.generate_rule_config("tenant_id", payload) + assert "Failed to generate variables" in result["error"] + + def test_generate_rule_config_with_variable_step3_error(self, mock_model_instance, model_config_entity): + payload = RuleGeneratePayload( + instruction="test instruction", model_config=model_config_entity, no_variable=False + ) + mock_res1 = MagicMock() + mock_res1.message.get_text_content.return_value = "Step 1 Prompt" + + mock_res2 = MagicMock() + mock_res2.message.get_text_content.return_value = '"var1"' + + # Step 3 fails + mock_model_instance.invoke_llm.side_effect = [mock_res1, mock_res2, InvokeError("Step 3 Failed")] + + result = LLMGenerator.generate_rule_config("tenant_id", payload) + assert "Failed to generate conversation opener" in result["error"] + + def test_generate_rule_config_with_variable_exception(self, mock_model_instance, model_config_entity): + payload = RuleGeneratePayload( + instruction="test instruction", model_config=model_config_entity, no_variable=False + ) + # Mock any step to throw Exception + mock_model_instance.invoke_llm.side_effect = Exception("Unexpected multi-step error") + + result = LLMGenerator.generate_rule_config("tenant_id", payload) + assert "Unexpected multi-step error" in result["error"] + + def test_generate_code_python_success(self, mock_model_instance, model_config_entity): + payload = RuleCodeGeneratePayload( + instruction="print hello", code_language="python", model_config=model_config_entity + ) + mock_response = MagicMock() + mock_response.message.get_text_content.return_value = "print('hello')" + mock_model_instance.invoke_llm.return_value = mock_response + + result = LLMGenerator.generate_code("tenant_id", payload) + assert result["code"] == "print('hello')" + assert result["language"] == "python" + + def test_generate_code_javascript_success(self, mock_model_instance, model_config_entity): + payload = RuleCodeGeneratePayload( + instruction="console log hello", code_language="javascript", model_config=model_config_entity + ) + mock_response = MagicMock() + mock_response.message.get_text_content.return_value = "console.log('hello')" + mock_model_instance.invoke_llm.return_value = mock_response + + result = LLMGenerator.generate_code("tenant_id", payload) + assert result["code"] == "console.log('hello')" + assert result["language"] == "javascript" + + def test_generate_code_invoke_error(self, mock_model_instance, model_config_entity): + payload = RuleCodeGeneratePayload( + instruction="error", code_language="python", model_config=model_config_entity + ) + mock_model_instance.invoke_llm.side_effect = InvokeError("Invoke failed") + + result = LLMGenerator.generate_code("tenant_id", payload) + assert "Failed to generate code" in result["error"] + + def test_generate_code_exception(self, mock_model_instance, model_config_entity): + payload = RuleCodeGeneratePayload( + instruction="error", code_language="python", model_config=model_config_entity + ) + mock_model_instance.invoke_llm.side_effect = Exception("Random error") + + result = LLMGenerator.generate_code("tenant_id", payload) + assert "Random error" in result["error"] + + def test_generate_qa_document_success(self, mock_model_instance): + mock_response = _llm_result("QA Document Content") + mock_model_instance.invoke_llm.return_value = mock_response + + result = LLMGenerator.generate_qa_document("tenant_id", "query", "English") + assert result == "QA Document Content" + + def test_generate_qa_document_type_error(self, mock_model_instance): + mock_model_instance.invoke_llm.return_value = "Not an LLMResult" + + with pytest.raises(TypeError, match="Expected LLMResult when stream=False"): + LLMGenerator.generate_qa_document("tenant_id", "query", "English") + + def test_generate_structured_output_success(self, mock_model_instance, model_config_entity): + payload = RuleStructuredOutputPayload(instruction="generate schema", model_config=model_config_entity) + mock_response = MagicMock() + mock_response.message.get_text_content.return_value = '{"type": "object", "properties": {}}' + mock_model_instance.invoke_llm.return_value = mock_response + + result = LLMGenerator.generate_structured_output("tenant_id", payload) + parsed_output = json.loads(result["output"]) + assert parsed_output["type"] == "object" + assert result["error"] == "" + + def test_generate_structured_output_json_repair(self, mock_model_instance, model_config_entity): + payload = RuleStructuredOutputPayload(instruction="generate schema", model_config=model_config_entity) + mock_response = MagicMock() + mock_response.message.get_text_content.return_value = "{'type': 'object'}" + mock_model_instance.invoke_llm.return_value = mock_response + + result = LLMGenerator.generate_structured_output("tenant_id", payload) + parsed_output = json.loads(result["output"]) + assert parsed_output["type"] == "object" + + def test_generate_structured_output_not_dict_or_list(self, mock_model_instance, model_config_entity): + payload = RuleStructuredOutputPayload(instruction="generate schema", model_config=model_config_entity) + mock_response = MagicMock() + mock_response.message.get_text_content.return_value = "true" # parsed as bool + mock_model_instance.invoke_llm.return_value = mock_response + + result = LLMGenerator.generate_structured_output("tenant_id", payload) + assert "Failed to parse structured output" in result["error"] + + def test_generate_structured_output_invoke_error(self, mock_model_instance, model_config_entity): + payload = RuleStructuredOutputPayload(instruction="error", model_config=model_config_entity) + mock_model_instance.invoke_llm.side_effect = InvokeError("Invoke failed") + + result = LLMGenerator.generate_structured_output("tenant_id", payload) + assert "Invoke failed" in result["error"] + + def test_generate_structured_output_exception(self, mock_model_instance, model_config_entity): + payload = RuleStructuredOutputPayload(instruction="error", model_config=model_config_entity) + mock_model_instance.invoke_llm.side_effect = Exception("Random error") + + result = LLMGenerator.generate_structured_output("tenant_id", payload) + assert "Random error" in result["error"] + + def test_instruction_modify_legacy_without_last_run_uses_real_empty_query( + self, + database: Session, + llm_invocation: Mock, + model_config_entity: ModelConfig, + ): + app = _persist_app(database) + llm_invocation.return_value = _llm_result('{"modified": "prompt"}') + + result = LLMGenerator.instruction_modify_legacy( + app.tenant_id, app.id, - "node", - "current", - "instruction", + "current_val", + "Test {{#last_run#}} and {{#current#}} and {{#error_message#}}", model_config_entity, "ideal", - workflow_service, session=database, ) - def test_instruction_modify_workflow_requires_draft_workflow( - self, - database: Session, - sqlite_session_factory: sessionmaker[Session], - model_config_entity: ModelConfig, - ): - app = _persist_app(database) - workflow_service = WorkflowService(sqlite_session_factory) + assert result == {"modified": "prompt"} + user_payload = json.loads(llm_invocation.call_args.kwargs["prompt_messages"][1].content) + assert "null" in user_payload["instruction"] + assert "current_val" in user_payload["instruction"] - with pytest.raises(ValueError, match="Workflow not found"): - LLMGenerator.instruction_modify_workflow( + def test_instruction_modify_legacy_reads_latest_tenant_scoped_message( + self, + database: Session, + llm_invocation: Mock, + model_config_entity: ModelConfig, + ): + app = _persist_app(database) + _persist_message( + database, + app, + query="older question", + answer="older answer", + created_at=datetime(2026, 1, 1), + ) + _persist_message( + database, + app, + query="latest question", + answer="latest answer", + created_at=datetime(2026, 1, 2), + ) + other_app = _persist_app(database) + _persist_message( + database, + other_app, + query="other tenant question", + created_at=datetime(2026, 1, 3), + ) + llm_invocation.return_value = _llm_result('{"modified": "prompt"}') + + result = LLMGenerator.instruction_modify_legacy( + app.tenant_id, app.id, "current", "instruction", model_config_entity, "ideal", session=database + ) + + assert result == {"modified": "prompt"} + user_payload = json.loads(llm_invocation.call_args.kwargs["prompt_messages"][1].content) + assert user_payload["last_run"]["query"] == "latest question" + assert user_payload["last_run"]["answer"] == "latest answer" + assert "older question" not in json.dumps(user_payload) + assert "other tenant question" not in json.dumps(user_payload) + + def test_instruction_modify_workflow_app_not_found( + self, + database: Session, + sqlite_session_factory: sessionmaker[Session], + model_config_entity: ModelConfig, + ): + workflow_service = WorkflowService(sqlite_session_factory) + + with pytest.raises(ValueError, match="App not found"): + LLMGenerator.instruction_modify_workflow( + str(uuid4()), + str(uuid4()), + "node", + "current", + "instruction", + model_config_entity, + "ideal", + workflow_service, + session=database, + ) + + def test_instruction_modify_workflow_rejects_app_from_another_tenant( + self, + database: Session, + sqlite_session_factory: sessionmaker[Session], + model_config_entity: ModelConfig, + ): + app = _persist_app(database) + workflow_service = WorkflowService(sqlite_session_factory) + + with pytest.raises(ValueError, match="App not found"): + LLMGenerator.instruction_modify_workflow( + str(uuid4()), + app.id, + "node", + "current", + "instruction", + model_config_entity, + "ideal", + workflow_service, + session=database, + ) + + def test_instruction_modify_workflow_requires_draft_workflow( + self, + database: Session, + sqlite_session_factory: sessionmaker[Session], + model_config_entity: ModelConfig, + ): + app = _persist_app(database) + workflow_service = WorkflowService(sqlite_session_factory) + + with pytest.raises(ValueError, match="Workflow not found"): + LLMGenerator.instruction_modify_workflow( + app.tenant_id, + app.id, + "node", + "current", + "instruction", + model_config_entity, + "ideal", + workflow_service, + session=database, + ) + + def test_instruction_modify_workflow_uses_last_run( + self, + database: Session, + sqlite_session_factory: sessionmaker[Session], + llm_invocation: Mock, + model_config_entity: ModelConfig, + ): + app = _persist_app(database) + workflow = _persist_workflow(database, app, node_type="llm") + _persist_node_execution( + database, + app, + workflow, + inputs={"input": "value"}, + agent_log=[{"status": "s", "error": "", "data": {"step": 1}}], + ) + workflow_service = WorkflowService(sqlite_session_factory) + + result = LLMGenerator.instruction_modify_workflow( app.tenant_id, app.id, "node", @@ -852,245 +900,401 @@ class TestLLMGenerator: session=database, ) - def test_instruction_modify_workflow_uses_last_run( - self, - database: Session, - sqlite_session_factory: sessionmaker[Session], - llm_invocation: Mock, - model_config_entity: ModelConfig, - ): - app = _persist_app(database) - workflow = _persist_workflow(database, app, node_type="llm") - _persist_node_execution( - database, - app, - workflow, - inputs={"input": "value"}, - agent_log=[{"status": "s", "error": "", "data": {"step": 1}}], + assert result == {"modified": "workflow"} + user_payload = json.loads(llm_invocation.call_args.kwargs["prompt_messages"][1].content) + assert user_payload["last_run"]["inputs"] == {"input": "value"} + assert user_payload["last_run"]["agent_log"][0]["data"] == {"step": 1} + + def test_instruction_modify_workflow_accepts_empty_agent_log( + self, + database: Session, + sqlite_session_factory: sessionmaker[Session], + llm_invocation: Mock, + model_config_entity: ModelConfig, + ): + app = _persist_app(database) + workflow = _persist_workflow(database, app, node_type="llm") + _persist_node_execution(database, app, workflow, agent_log=[]) + workflow_service = WorkflowService(sqlite_session_factory) + + result = LLMGenerator.instruction_modify_workflow( + app.tenant_id, + app.id, + "node", + "current", + "instruction", + model_config_entity, + "ideal", + workflow_service, + session=database, + ) + + assert result == {"modified": "workflow"} + user_payload = json.loads(llm_invocation.call_args.kwargs["prompt_messages"][1].content) + assert user_payload["last_run"]["agent_log"] == [] + + @pytest.mark.parametrize( + "node_type", + [ + "code", + None, + ], ) - workflow_service = WorkflowService(sqlite_session_factory) + def test_instruction_modify_workflow_falls_back_without_last_run( + self, + database: Session, + sqlite_session_factory: sessionmaker[Session], + llm_invocation: Mock, + model_config_entity: ModelConfig, + node_type: str | None, + ): + app = _persist_app(database) + _persist_workflow(database, app, node_type=node_type) + workflow_service = WorkflowService(sqlite_session_factory) + llm_invocation.return_value = _llm_result('{"modified": "fallback"}') - result = LLMGenerator.instruction_modify_workflow( - app.tenant_id, - app.id, - "node", - "current", - "instruction", - model_config_entity, - "ideal", - workflow_service, - session=database, + result = LLMGenerator.instruction_modify_workflow( + app.tenant_id, + app.id, + "node", + "current", + "instruction", + model_config_entity, + "ideal", + workflow_service, + session=database, + ) + + assert result == {"modified": "fallback"} + + def test_instruction_modify_workflow_falls_back_for_unknown_node_type( + self, + database: Session, + sqlite_session_factory: sessionmaker[Session], + llm_invocation: Mock, + model_config_entity: ModelConfig, + ): + app = _persist_app(database) + _persist_workflow(database, app, node_type="unknown") + workflow_service = WorkflowService(sqlite_session_factory) + + result = LLMGenerator.instruction_modify_workflow( + app.tenant_id, + app.id, + "node", + "current", + "instruction", + model_config_entity, + "ideal", + workflow_service, + session=database, + ) + + assert result == {"modified": "workflow"} + system_prompt = llm_invocation.call_args.kwargs["prompt_messages"][0].content + assert system_prompt == llm_generator_module.LLM_MODIFY_PROMPT_SYSTEM + + @pytest.mark.parametrize( + ("raw_output", "error_fragment"), + [ + ("No braces here", "Could not find a valid JSON object"), + ("[1, 2, 3]", "Could not find a valid JSON object"), + ], ) + def test_instruction_modify_rejects_invalid_model_output( + self, + database: Session, + mock_model_instance: MagicMock, + model_config_entity: ModelConfig, + raw_output: str, + error_fragment: str, + ): + app = _persist_app(database) + response = MagicMock() + response.message.get_text_content.return_value = raw_output + mock_model_instance.invoke_llm.return_value = response - assert result == {"modified": "workflow"} - user_payload = json.loads(llm_invocation.call_args.kwargs["prompt_messages"][1].content) - assert user_payload["last_run"]["inputs"] == {"input": "value"} - assert user_payload["last_run"]["agent_log"][0]["data"] == {"step": 1} + result = LLMGenerator.instruction_modify_legacy( + app.tenant_id, app.id, "current", "instruction", model_config_entity, "ideal", session=database + ) - def test_instruction_modify_workflow_accepts_empty_agent_log( - self, - database: Session, - sqlite_session_factory: sessionmaker[Session], - llm_invocation: Mock, - model_config_entity: ModelConfig, - ): - app = _persist_app(database) - workflow = _persist_workflow(database, app, node_type="llm") - _persist_node_execution(database, app, workflow, agent_log=[]) - workflow_service = WorkflowService(sqlite_session_factory) + assert error_fragment in result["error"] - result = LLMGenerator.instruction_modify_workflow( - app.tenant_id, - app.id, - "node", - "current", - "instruction", - model_config_entity, - "ideal", - workflow_service, - session=database, + @pytest.mark.parametrize( + ("model_error", "error_fragment"), + [(InvokeError("invoke failed"), "invoke failed"), (RuntimeError("boom"), "boom")], ) + def test_instruction_modify_handles_model_errors( + self, + database: Session, + mock_model_instance: MagicMock, + model_config_entity: ModelConfig, + model_error: Exception, + error_fragment: str, + ): + app = _persist_app(database) + mock_model_instance.invoke_llm.side_effect = model_error - assert result == {"modified": "workflow"} - user_payload = json.loads(llm_invocation.call_args.kwargs["prompt_messages"][1].content) - assert user_payload["last_run"]["agent_log"] == [] + result = LLMGenerator.instruction_modify_legacy( + app.tenant_id, app.id, "current", "instruction", model_config_entity, "ideal", session=database + ) - @pytest.mark.parametrize( - "node_type", - [ - "code", - None, - ], + assert error_fragment.lower() in result["error"].lower() + + # --------------------------------------------------------------------------- + # Prompt-generation telemetry emit-site tests + # --------------------------------------------------------------------------- + + class TestPromptGenerationTelemetryEmit: + """Verify that LLM generator methods call telemetry_emit with a + PromptGenerationEvent so the metric pipeline is not accidentally broken.""" + + @pytest.fixture(autouse=True) + def _setup(self, llm_invocation: Mock) -> None: + self.llm_invocation = llm_invocation + + @patch("core.llm_generator.llm_generator.telemetry_emit") + def test_generate_code_emits_on_success(self, mock_emit: MagicMock) -> None: + self.llm_invocation.return_value = _llm_result("print('hello')") + model_cfg = ModelConfig(provider="openai", name="gpt-4", mode="chat", completion_params={}) + args = RuleCodeGeneratePayload( + instruction="write hello world", model_config=model_cfg, code_language="python" + ) + + LLMGenerator.generate_code(tenant_id="t-1", args=args, app_id="app-1") + + mock_emit.assert_called_once() + event = mock_emit.call_args[0][0] + assert type(event).__name__ == "PromptGenerationEvent" + assert event.payload["operation_type"] == "code_generate" + assert event.payload["model_provider"] == "openai" + assert event.payload["model_name"] == "gpt-4" + assert event.context.tenant_id == "t-1" + assert event.context.app_id == "app-1" + assert event.payload["error"] is None + + @patch("core.llm_generator.llm_generator.telemetry_emit") + def test_generate_code_emits_on_failure(self, mock_emit: MagicMock) -> None: + self.llm_invocation.side_effect = InvokeError("model down") + model_cfg = ModelConfig(provider="openai", name="gpt-4", mode="chat", completion_params={}) + args = RuleCodeGeneratePayload( + instruction="write hello world", model_config=model_cfg, code_language="python" + ) + + result = LLMGenerator.generate_code(tenant_id="t-1", args=args) + + mock_emit.assert_called_once() + event = mock_emit.call_args[0][0] + assert event.payload["error"] == "model down" + assert "Failed to generate code" in result["error"] + + @patch("core.llm_generator.llm_generator.telemetry_emit") + def test_generate_rule_config_no_variable_emits(self, mock_emit: MagicMock) -> None: + self.llm_invocation.return_value = _llm_result("generated prompt") + model_cfg = ModelConfig(provider="anthropic", name="claude-3", mode="chat", completion_params={}) + args = RuleGeneratePayload(instruction="be helpful", model_config=model_cfg, no_variable=True) + + LLMGenerator.generate_rule_config(tenant_id="t-1", args=args, app_id="app-2") + + mock_emit.assert_called_once() + event = mock_emit.call_args[0][0] + assert event.payload["operation_type"] == "rule_generate" + assert event.payload["model_provider"] == "anthropic" + assert event.context.app_id == "app-2" + + @patch("core.llm_generator.llm_generator.telemetry_emit") + def test_generate_rule_config_no_variable_emits_on_failure(self, mock_emit: MagicMock) -> None: + self.llm_invocation.side_effect = InvokeError("auth fail") + model_cfg = ModelConfig(provider="anthropic", name="claude-3", mode="chat", completion_params={}) + args = RuleGeneratePayload(instruction="be helpful", model_config=model_cfg, no_variable=True) + + result = LLMGenerator.generate_rule_config(tenant_id="t-1", args=args) + + mock_emit.assert_called_once() + event = mock_emit.call_args[0][0] + assert event.payload["error"] == "auth fail" + assert "error" in result + + +def _llm_generator_missing_llm_result(content: str) -> LLMResult: + """Build a real non-streaming LLM response around deterministic test content.""" + + return LLMResult( + model="test-model", + message=AssistantPromptMessage(content=content), + usage=LLMUsage.empty_usage(), ) - def test_instruction_modify_workflow_falls_back_without_last_run( - self, - database: Session, - sqlite_session_factory: sessionmaker[Session], - llm_invocation: Mock, - model_config_entity: ModelConfig, - node_type: str | None, - ): - app = _persist_app(database) - _persist_workflow(database, app, node_type=node_type) - workflow_service = WorkflowService(sqlite_session_factory) - llm_invocation.return_value = _llm_result('{"modified": "fallback"}') - result = LLMGenerator.instruction_modify_workflow( - app.tenant_id, - app.id, - "node", - "current", - "instruction", - model_config_entity, - "ideal", - workflow_service, - session=database, - ) - assert result == {"modified": "fallback"} +def _model_manager(monkeypatch: pytest.MonkeyPatch) -> tuple[ModelManager, MagicMock, MagicMock]: + """Compose real model objects with controllable resolution and invocation boundaries.""" + model_manager = create_plugin_model_manager(tenant_id="tenant") + model_instance = make_model_instance(provider="openai", model="gpt-4") + invocation = MagicMock() + resolve_default = MagicMock(return_value=model_instance) + monkeypatch.setattr(model_instance, "invoke_llm", invocation) + monkeypatch.setattr(model_manager, "get_default_model_instance", resolve_default) + return model_manager, invocation, resolve_default - def test_instruction_modify_workflow_falls_back_for_unknown_node_type( - self, - database: Session, - sqlite_session_factory: sessionmaker[Session], - llm_invocation: Mock, - model_config_entity: ModelConfig, - ): - app = _persist_app(database) - _persist_workflow(database, app, node_type="unknown") - workflow_service = WorkflowService(sqlite_session_factory) - result = LLMGenerator.instruction_modify_workflow( - app.tenant_id, - app.id, - "node", - "current", - "instruction", - model_config_entity, - "ideal", - workflow_service, - session=database, - ) - - assert result == {"modified": "workflow"} - system_prompt = llm_invocation.call_args.kwargs["prompt_messages"][0].content - assert system_prompt == llm_generator_module.LLM_MODIFY_PROMPT_SYSTEM - - @pytest.mark.parametrize( - ("raw_output", "error_fragment"), - [ - ("No braces here", "Could not find a valid JSON object"), - ("[1, 2, 3]", "Could not find a valid JSON object"), - ], +def _dataset(*, dataset_id: str, tenant_id: str, name: str, created_at: datetime) -> Dataset: + return make_dataset( + dataset_id=dataset_id, + tenant_id=tenant_id, + name=name, + created_by="account-id", + created_at=created_at, ) - def test_instruction_modify_rejects_invalid_model_output( - self, - database: Session, - mock_model_instance: MagicMock, - model_config_entity: ModelConfig, - raw_output: str, - error_fragment: str, - ): - app = _persist_app(database) - response = MagicMock() - response.message.get_text_content.return_value = raw_output - mock_model_instance.invoke_llm.return_value = response - - result = LLMGenerator.instruction_modify_legacy( - app.tenant_id, app.id, "current", "instruction", model_config_entity, "ideal", session=database - ) - - assert error_fragment in result["error"] - - @pytest.mark.parametrize( - ("model_error", "error_fragment"), - [(InvokeError("invoke failed"), "invoke failed"), (RuntimeError("boom"), "boom")], - ) - def test_instruction_modify_handles_model_errors( - self, - database: Session, - mock_model_instance: MagicMock, - model_config_entity: ModelConfig, - model_error: Exception, - error_fragment: str, - ): - app = _persist_app(database) - mock_model_instance.invoke_llm.side_effect = model_error - - result = LLMGenerator.instruction_modify_legacy( - app.tenant_id, app.id, "current", "instruction", model_config_entity, "ideal", session=database - ) - - assert error_fragment.lower() in result["error"].lower() -# --------------------------------------------------------------------------- -# Prompt-generation telemetry emit-site tests -# --------------------------------------------------------------------------- +class TestLlmGeneratorMissingCases: + @staticmethod + @pytest.fixture + def dataset_session(sqlite_session: Session, monkeypatch: pytest.MonkeyPatch) -> Session: + """Bind the real SQLite session to the production database extension.""" + monkeypatch.setattr(generator_module.db, "session", sqlite_session) + return sqlite_session -class TestPromptGenerationTelemetryEmit: - """Verify that LLM generator methods call telemetry_emit with a - PromptGenerationEvent so the metric pipeline is not accidentally broken.""" + class TestParseStringList: + def test_empty(self): + assert _parse_string_list("") == [] - @pytest.fixture(autouse=True) - def _setup(self, llm_invocation: Mock) -> None: - self.llm_invocation = llm_invocation + def test_no_match(self): + assert _parse_string_list("no list here") == [] - @patch("core.llm_generator.llm_generator.telemetry_emit") - def test_generate_code_emits_on_success(self, mock_emit: MagicMock) -> None: - self.llm_invocation.return_value = _llm_result("print('hello')") - model_cfg = ModelConfig(provider="openai", name="gpt-4", mode="chat", completion_params={}) - args = RuleCodeGeneratePayload(instruction="write hello world", model_config=model_cfg, code_language="python") + def test_valid_json(self): + assert _parse_string_list('["item1", "item2"]') == ["item1", "item2"] - LLMGenerator.generate_code(tenant_id="t-1", args=args, app_id="app-1") + def test_with_surrounding_text(self): + assert _parse_string_list('Here is the list: ["a", "b"] enjoy!') == ["a", "b"] - mock_emit.assert_called_once() - event = mock_emit.call_args[0][0] - assert type(event).__name__ == "PromptGenerationEvent" - assert event.payload["operation_type"] == "code_generate" - assert event.payload["model_provider"] == "openai" - assert event.payload["model_name"] == "gpt-4" - assert event.context.tenant_id == "t-1" - assert event.context.app_id == "app-1" - assert event.payload["error"] is None + def test_invalid_json_fallback(self): + # json_repair can fix missing quotes + assert _parse_string_list("[item1, item2]") == ["item1", "item2"] - @patch("core.llm_generator.llm_generator.telemetry_emit") - def test_generate_code_emits_on_failure(self, mock_emit: MagicMock) -> None: - self.llm_invocation.side_effect = InvokeError("model down") - model_cfg = ModelConfig(provider="openai", name="gpt-4", mode="chat", completion_params={}) - args = RuleCodeGeneratePayload(instruction="write hello world", model_config=model_cfg, code_language="python") + def test_completely_invalid_json(self): + assert _parse_string_list("[{}}]") == [] - result = LLMGenerator.generate_code(tenant_id="t-1", args=args) + def test_not_a_list(self): + assert _parse_string_list('{"a": "b"}') == [] - mock_emit.assert_called_once() - event = mock_emit.call_args[0][0] - assert event.payload["error"] == "model down" - assert "Failed to generate code" in result["error"] + def test_filter_non_strings(self): + assert _parse_string_list('["a", 1, "b", {"foo": "bar"}]') == ["a", "b"] - @patch("core.llm_generator.llm_generator.telemetry_emit") - def test_generate_rule_config_no_variable_emits(self, mock_emit: MagicMock) -> None: - self.llm_invocation.return_value = _llm_result("generated prompt") - model_cfg = ModelConfig(provider="anthropic", name="claude-3", mode="chat", completion_params={}) - args = RuleGeneratePayload(instruction="be helpful", model_config=model_cfg, no_variable=True) + class TestGenerateWorkflowInstructionSuggestions: + @patch("core.llm_generator.llm_generator.ModelManager.for_tenant") + def test_no_default_model(self, mock_for_tenant, monkeypatch: pytest.MonkeyPatch): + model_manager, _, resolve_default = _model_manager(monkeypatch) + resolve_default.side_effect = RuntimeError("no default model") + mock_for_tenant.return_value = model_manager - LLMGenerator.generate_rule_config(tenant_id="t-1", args=args, app_id="app-2") + assert LLMGenerator.generate_workflow_instruction_suggestions("tenant", mode="workflow") == [] - mock_emit.assert_called_once() - event = mock_emit.call_args[0][0] - assert event.payload["operation_type"] == "rule_generate" - assert event.payload["model_provider"] == "anthropic" - assert event.context.app_id == "app-2" + @patch("core.llm_generator.llm_generator.ModelManager.for_tenant") + @patch("core.llm_generator.llm_generator.LLMGenerator._build_suggestion_context") + def test_llm_success(self, mock_build_context, mock_for_tenant, monkeypatch: pytest.MonkeyPatch): + mock_build_context.return_value = "context" + model_manager, invocation, _ = _model_manager(monkeypatch) + invocation.return_value = _llm_generator_missing_llm_result('["idea 1", "idea 2"]') + mock_for_tenant.return_value = model_manager - @patch("core.llm_generator.llm_generator.telemetry_emit") - def test_generate_rule_config_no_variable_emits_on_failure(self, mock_emit: MagicMock) -> None: - self.llm_invocation.side_effect = InvokeError("auth fail") - model_cfg = ModelConfig(provider="anthropic", name="claude-3", mode="chat", completion_params={}) - args = RuleGeneratePayload(instruction="be helpful", model_config=model_cfg, no_variable=True) + result = LLMGenerator.generate_workflow_instruction_suggestions("tenant", mode="workflow") + assert result == ["idea 1", "idea 2"] + invocation.assert_called_once() - result = LLMGenerator.generate_rule_config(tenant_id="t-1", args=args) + @patch("core.llm_generator.llm_generator.ModelManager.for_tenant") + @patch("core.llm_generator.llm_generator.LLMGenerator._build_suggestion_context") + def test_llm_error(self, mock_build_context, mock_for_tenant, monkeypatch: pytest.MonkeyPatch): + mock_build_context.return_value = "context" + model_manager, invocation, _ = _model_manager(monkeypatch) + invocation.side_effect = RuntimeError("API error") + mock_for_tenant.return_value = model_manager - mock_emit.assert_called_once() - event = mock_emit.call_args[0][0] - assert event.payload["error"] == "auth fail" - assert "error" in result + result = LLMGenerator.generate_workflow_instruction_suggestions("tenant", mode="workflow") + assert result == [] + invocation.assert_called_once() + + @patch("core.llm_generator.llm_generator.ModelManager.for_tenant") + @patch("core.llm_generator.llm_generator.LLMGenerator._build_suggestion_context") + def test_llm_bad_output(self, mock_build_context, mock_for_tenant, monkeypatch: pytest.MonkeyPatch): + mock_build_context.return_value = "context" + model_manager, invocation, _ = _model_manager(monkeypatch) + invocation.return_value = _llm_generator_missing_llm_result("Not a list") + mock_for_tenant.return_value = model_manager + + result = LLMGenerator.generate_workflow_instruction_suggestions("tenant", mode="workflow") + assert result == [] + invocation.assert_called_once() + + @pytest.mark.parametrize("sqlite_session", [(Dataset,)], indirect=True) + class TestBuildSuggestionContext: + def test_both_success(self, dataset_session: Session, monkeypatch: pytest.MonkeyPatch): + now = datetime.now() + dataset_session.add_all( + ( + _dataset(dataset_id="kb-1", tenant_id="tenant", name="kb1", created_at=now), + _dataset( + dataset_id="kb-2", + tenant_id="tenant", + name="kb2", + created_at=now - timedelta(seconds=1), + ), + _dataset(dataset_id="other-kb", tenant_id="other", name="private", created_at=now), + ) + ) + dataset_session.commit() + + def build_tool_catalogue(_tenant_id: str) -> list[ToolCatalogueEntry]: + return [ + ToolCatalogueEntry( + provider_name="provider", + provider_type="builtin", + plugin_id="", + tool_name="tool1", + tool_label="tool1", + description="First tool", + ), + ToolCatalogueEntry( + provider_name="provider", + provider_type="builtin", + plugin_id="", + tool_name="tool2", + tool_label="tool2", + description="Second tool", + ), + ] + + # Keep the real module and formatter; only isolate provider/plugin discovery. + monkeypatch.setattr(tool_catalogue_module, "build_tool_catalogue", build_tool_catalogue) + + result = LLMGenerator._build_suggestion_context("tenant") + assert "Knowledge bases:\n- kb1\n- kb2" in result + assert ( + 'Installed tools:\n- provider/tool1 [provider_id="provider"; tool_name="tool1"] — First tool\n' + '- provider/tool2 [provider_id="provider"; tool_name="tool2"] — Second tool' + ) in result + + def test_both_fail(self, dataset_session: Session, monkeypatch: pytest.MonkeyPatch): + def fail_query(_orm_execute_state: object) -> None: + raise SQLAlchemyError("DB error") + + def fail_tool_catalogue(_tenant_id: str) -> list[ToolCatalogueEntry]: + raise RuntimeError("Tool error") + + event.listen(dataset_session, "do_orm_execute", fail_query) + monkeypatch.setattr(tool_catalogue_module, "build_tool_catalogue", fail_tool_catalogue) + + try: + assert LLMGenerator._build_suggestion_context("tenant") == "" + finally: + event.remove(dataset_session, "do_orm_execute", fail_query) + + class TestWorkflowServiceInterface: + def test_real_workflow_service_exposes_protocol_methods(self): + from core.llm_generator.llm_generator import WorkflowServiceInterface + + service: WorkflowServiceInterface = WorkflowService(sessionmaker()) + + assert callable(service.get_draft_workflow) + assert callable(service.get_node_last_run) diff --git a/api/tests/unit_tests/core/llm_generator/test_llm_generator_missing.py b/api/tests/unit_tests/core/llm_generator/test_llm_generator_missing.py deleted file mode 100644 index 8631c7fb752..00000000000 --- a/api/tests/unit_tests/core/llm_generator/test_llm_generator_missing.py +++ /dev/null @@ -1,206 +0,0 @@ -from datetime import datetime, timedelta -from unittest.mock import MagicMock, patch - -import pytest -from sqlalchemy import event -from sqlalchemy.exc import SQLAlchemyError -from sqlalchemy.orm import Session, sessionmaker - -import core.llm_generator.llm_generator as generator_module -from core.llm_generator.llm_generator import LLMGenerator, _parse_string_list -from core.model_manager import ModelManager -from core.plugin.impl.model_runtime_factory import create_plugin_model_manager -from core.workflow.generator import tool_catalogue as tool_catalogue_module -from core.workflow.generator.tool_catalogue import ToolCatalogueEntry -from graphon.model_runtime.entities.llm_entities import LLMResult, LLMUsage -from graphon.model_runtime.entities.message_entities import AssistantPromptMessage -from models.dataset import Dataset -from services.workflow_service import WorkflowService -from tests.unit_tests.core.model_fixtures import make_model_instance -from tests.unit_tests.model_factories import make_dataset - - -@pytest.fixture -def dataset_session(sqlite_session: Session, monkeypatch: pytest.MonkeyPatch) -> Session: - """Bind the real SQLite session to the production database extension.""" - - monkeypatch.setattr(generator_module.db, "session", sqlite_session) - return sqlite_session - - -def _llm_result(content: str) -> LLMResult: - """Build a real non-streaming LLM response around deterministic test content.""" - - return LLMResult( - model="test-model", - message=AssistantPromptMessage(content=content), - usage=LLMUsage.empty_usage(), - ) - - -def _model_manager(monkeypatch: pytest.MonkeyPatch) -> tuple[ModelManager, MagicMock, MagicMock]: - """Compose real model objects with controllable resolution and invocation boundaries.""" - model_manager = create_plugin_model_manager(tenant_id="tenant") - model_instance = make_model_instance(provider="openai", model="gpt-4") - invocation = MagicMock() - resolve_default = MagicMock(return_value=model_instance) - monkeypatch.setattr(model_instance, "invoke_llm", invocation) - monkeypatch.setattr(model_manager, "get_default_model_instance", resolve_default) - return model_manager, invocation, resolve_default - - -def _dataset(*, dataset_id: str, tenant_id: str, name: str, created_at: datetime) -> Dataset: - return make_dataset( - dataset_id=dataset_id, - tenant_id=tenant_id, - name=name, - created_by="account-id", - created_at=created_at, - ) - - -class TestParseStringList: - def test_empty(self): - assert _parse_string_list("") == [] - - def test_no_match(self): - assert _parse_string_list("no list here") == [] - - def test_valid_json(self): - assert _parse_string_list('["item1", "item2"]') == ["item1", "item2"] - - def test_with_surrounding_text(self): - assert _parse_string_list('Here is the list: ["a", "b"] enjoy!') == ["a", "b"] - - def test_invalid_json_fallback(self): - # json_repair can fix missing quotes - assert _parse_string_list("[item1, item2]") == ["item1", "item2"] - - def test_completely_invalid_json(self): - assert _parse_string_list("[{}}]") == [] - - def test_not_a_list(self): - assert _parse_string_list('{"a": "b"}') == [] - - def test_filter_non_strings(self): - assert _parse_string_list('["a", 1, "b", {"foo": "bar"}]') == ["a", "b"] - - -class TestGenerateWorkflowInstructionSuggestions: - @patch("core.llm_generator.llm_generator.ModelManager.for_tenant") - def test_no_default_model(self, mock_for_tenant, monkeypatch: pytest.MonkeyPatch): - model_manager, _, resolve_default = _model_manager(monkeypatch) - resolve_default.side_effect = RuntimeError("no default model") - mock_for_tenant.return_value = model_manager - - assert LLMGenerator.generate_workflow_instruction_suggestions("tenant", mode="workflow") == [] - - @patch("core.llm_generator.llm_generator.ModelManager.for_tenant") - @patch("core.llm_generator.llm_generator.LLMGenerator._build_suggestion_context") - def test_llm_success(self, mock_build_context, mock_for_tenant, monkeypatch: pytest.MonkeyPatch): - mock_build_context.return_value = "context" - model_manager, invocation, _ = _model_manager(monkeypatch) - invocation.return_value = _llm_result('["idea 1", "idea 2"]') - mock_for_tenant.return_value = model_manager - - result = LLMGenerator.generate_workflow_instruction_suggestions("tenant", mode="workflow") - assert result == ["idea 1", "idea 2"] - invocation.assert_called_once() - - @patch("core.llm_generator.llm_generator.ModelManager.for_tenant") - @patch("core.llm_generator.llm_generator.LLMGenerator._build_suggestion_context") - def test_llm_error(self, mock_build_context, mock_for_tenant, monkeypatch: pytest.MonkeyPatch): - mock_build_context.return_value = "context" - model_manager, invocation, _ = _model_manager(monkeypatch) - invocation.side_effect = RuntimeError("API error") - mock_for_tenant.return_value = model_manager - - result = LLMGenerator.generate_workflow_instruction_suggestions("tenant", mode="workflow") - assert result == [] - invocation.assert_called_once() - - @patch("core.llm_generator.llm_generator.ModelManager.for_tenant") - @patch("core.llm_generator.llm_generator.LLMGenerator._build_suggestion_context") - def test_llm_bad_output(self, mock_build_context, mock_for_tenant, monkeypatch: pytest.MonkeyPatch): - mock_build_context.return_value = "context" - model_manager, invocation, _ = _model_manager(monkeypatch) - invocation.return_value = _llm_result("Not a list") - mock_for_tenant.return_value = model_manager - - result = LLMGenerator.generate_workflow_instruction_suggestions("tenant", mode="workflow") - assert result == [] - invocation.assert_called_once() - - -@pytest.mark.parametrize("sqlite_session", [(Dataset,)], indirect=True) -class TestBuildSuggestionContext: - def test_both_success(self, dataset_session: Session, monkeypatch: pytest.MonkeyPatch): - now = datetime.now() - dataset_session.add_all( - ( - _dataset(dataset_id="kb-1", tenant_id="tenant", name="kb1", created_at=now), - _dataset( - dataset_id="kb-2", - tenant_id="tenant", - name="kb2", - created_at=now - timedelta(seconds=1), - ), - _dataset(dataset_id="other-kb", tenant_id="other", name="private", created_at=now), - ) - ) - dataset_session.commit() - - def build_tool_catalogue(_tenant_id: str) -> list[ToolCatalogueEntry]: - return [ - ToolCatalogueEntry( - provider_name="provider", - provider_type="builtin", - plugin_id="", - tool_name="tool1", - tool_label="tool1", - description="First tool", - ), - ToolCatalogueEntry( - provider_name="provider", - provider_type="builtin", - plugin_id="", - tool_name="tool2", - tool_label="tool2", - description="Second tool", - ), - ] - - # Keep the real module and formatter; only isolate provider/plugin discovery. - monkeypatch.setattr(tool_catalogue_module, "build_tool_catalogue", build_tool_catalogue) - - result = LLMGenerator._build_suggestion_context("tenant") - assert "Knowledge bases:\n- kb1\n- kb2" in result - assert ( - 'Installed tools:\n- provider/tool1 [provider_id="provider"; tool_name="tool1"] — First tool\n' - '- provider/tool2 [provider_id="provider"; tool_name="tool2"] — Second tool' - ) in result - - def test_both_fail(self, dataset_session: Session, monkeypatch: pytest.MonkeyPatch): - def fail_query(_orm_execute_state: object) -> None: - raise SQLAlchemyError("DB error") - - def fail_tool_catalogue(_tenant_id: str) -> list[ToolCatalogueEntry]: - raise RuntimeError("Tool error") - - event.listen(dataset_session, "do_orm_execute", fail_query) - monkeypatch.setattr(tool_catalogue_module, "build_tool_catalogue", fail_tool_catalogue) - - try: - assert LLMGenerator._build_suggestion_context("tenant") == "" - finally: - event.remove(dataset_session, "do_orm_execute", fail_query) - - -class TestWorkflowServiceInterface: - def test_real_workflow_service_exposes_protocol_methods(self): - from core.llm_generator.llm_generator import WorkflowServiceInterface - - service: WorkflowServiceInterface = WorkflowService(sessionmaker()) - - assert callable(service.get_draft_workflow) - assert callable(service.get_node_last_run) diff --git a/api/tests/unit_tests/core/logging/test_trace_helpers.py b/api/tests/unit_tests/core/logging/test_trace_helpers.py deleted file mode 100644 index 707a0049982..00000000000 --- a/api/tests/unit_tests/core/logging/test_trace_helpers.py +++ /dev/null @@ -1,95 +0,0 @@ -"""Tests for trace helper functions.""" - -import re -from unittest import mock - -from opentelemetry.trace import NonRecordingSpan, SpanContext, use_span - - -class TestGetSpanIdFromOtelContext: - def test_returns_none_without_span(self): - from core.helper.trace_id_helper import get_span_id_from_otel_context - - with mock.patch("opentelemetry.trace.get_current_span", return_value=None, autospec=True): - result = get_span_id_from_otel_context() - assert result is None - - def test_returns_span_id_when_available(self): - from core.helper.trace_id_helper import get_span_id_from_otel_context - - span = NonRecordingSpan(SpanContext(trace_id=1, span_id=0x051581BF3BB55C45, is_remote=False)) - with use_span(span): - result = get_span_id_from_otel_context() - assert result == "051581bf3bb55c45" - - def test_returns_none_on_exception(self): - from core.helper.trace_id_helper import get_span_id_from_otel_context - - with mock.patch("opentelemetry.trace.get_current_span", side_effect=Exception("Test error"), autospec=True): - result = get_span_id_from_otel_context() - assert result is None - - -class TestGenerateTraceparentHeader: - def test_generates_valid_format(self): - from core.helper.trace_id_helper import generate_traceparent_header - - with mock.patch("opentelemetry.trace.get_current_span", return_value=None, autospec=True): - result = generate_traceparent_header() - - assert result is not None - # Format: 00-{trace_id}-{span_id}-01 - parts = result.split("-") - assert len(parts) == 4 - assert parts[0] == "00" # version - assert len(parts[1]) == 32 # trace_id (32 hex chars) - assert len(parts[2]) == 16 # span_id (16 hex chars) - assert parts[3] == "01" # flags - - def test_uses_otel_context_when_available(self): - from core.helper.trace_id_helper import generate_traceparent_header - - span = NonRecordingSpan( - SpanContext( - trace_id=0x5B8AA5A2D2C872E8321CF37308D69DF2, - span_id=0x051581BF3BB55C45, - is_remote=False, - ) - ) - with use_span(span): - result = generate_traceparent_header() - assert result == "00-5b8aa5a2d2c872e8321cf37308d69df2-051581bf3bb55c45-01" - - def test_generates_hex_only_values(self): - from core.helper.trace_id_helper import generate_traceparent_header - - with mock.patch("opentelemetry.trace.get_current_span", return_value=None, autospec=True): - result = generate_traceparent_header() - - parts = result.split("-") - # All parts should be valid hex - assert re.match(r"^[0-9a-f]+$", parts[1]) - assert re.match(r"^[0-9a-f]+$", parts[2]) - - -class TestParseTraceparentHeader: - def test_parses_valid_traceparent(self): - from core.helper.trace_id_helper import parse_traceparent_header - - traceparent = "00-5b8aa5a2d2c872e8321cf37308d69df2-051581bf3bb55c45-01" - result = parse_traceparent_header(traceparent) - - assert result == "5b8aa5a2d2c872e8321cf37308d69df2" - - def test_returns_none_for_invalid_format(self): - from core.helper.trace_id_helper import parse_traceparent_header - - # Wrong number of parts - assert parse_traceparent_header("00-abc-def") is None - # Wrong trace_id length - assert parse_traceparent_header("00-abc-def-01") is None - - def test_returns_none_for_empty_string(self): - from core.helper.trace_id_helper import parse_traceparent_header - - assert parse_traceparent_header("") is None diff --git a/api/tests/unit_tests/core/mcp/auth/test_auth_flow.py b/api/tests/unit_tests/core/mcp/auth/test_auth_flow.py index 3ff42fa7422..c10d4cb1ce1 100644 --- a/api/tests/unit_tests/core/mcp/auth/test_auth_flow.py +++ b/api/tests/unit_tests/core/mcp/auth/test_auth_flow.py @@ -41,7 +41,7 @@ from core.mcp.types import ( OAuthTokens, ProtectedResourceMetadata, ) -from tests.unit_tests.core.mcp.fixtures import make_provider +from tests.helpers.unit.core.mcp.fixtures import make_provider @pytest.fixture diff --git a/api/tests/unit_tests/core/mcp/client/test_session.py b/api/tests/unit_tests/core/mcp/client/test_session.py deleted file mode 100644 index 65f6a8a6f52..00000000000 --- a/api/tests/unit_tests/core/mcp/client/test_session.py +++ /dev/null @@ -1,468 +0,0 @@ -import queue -import threading -from typing import Any - -from core.mcp import types -from core.mcp.entities import RequestContext -from core.mcp.session.base_session import RequestResponder -from core.mcp.session.client_session import DEFAULT_CLIENT_INFO, ClientSession -from core.mcp.types import ( - LATEST_PROTOCOL_VERSION, - ClientNotification, - ClientRequest, - Implementation, - InitializedNotification, - InitializeRequest, - InitializeResult, - JSONRPCMessage, - JSONRPCNotification, - JSONRPCRequest, - JSONRPCResponse, - ServerCapabilities, - ServerResult, - SessionMessage, -) - - -def test_client_session_initialize() -> None: - # Create synchronous queues to replace async streams - client_to_server: queue.Queue[SessionMessage] = queue.Queue() - server_to_client: queue.Queue[SessionMessage] = queue.Queue() - - initialized_notification = None - - def mock_server() -> None: - nonlocal initialized_notification - - # Receive initialization request - session_message = client_to_server.get(timeout=5.0) - jsonrpc_request = session_message.message - assert isinstance(jsonrpc_request.root, JSONRPCRequest) - request = ClientRequest.model_validate( - jsonrpc_request.root.model_dump(by_alias=True, mode="json", exclude_none=True) - ) - assert isinstance(request.root, InitializeRequest) - - # Create response - result = ServerResult( - InitializeResult( - protocolVersion=LATEST_PROTOCOL_VERSION, - capabilities=ServerCapabilities( - logging=None, - resources=None, - tools=None, - experimental=None, - prompts=None, - ), - serverInfo=Implementation(name="mock-server", version="0.1.0"), - instructions="The server instructions.", - ) - ) - - # Send response - server_to_client.put( - SessionMessage( - message=JSONRPCMessage( - JSONRPCResponse( - jsonrpc="2.0", - id=jsonrpc_request.root.id, - result=result.model_dump(by_alias=True, mode="json", exclude_none=True), - ) - ) - ) - ) - - # Receive initialized notification - session_notification = client_to_server.get(timeout=5.0) - jsonrpc_notification = session_notification.message - assert isinstance(jsonrpc_notification.root, JSONRPCNotification) - initialized_notification = ClientNotification.model_validate( - jsonrpc_notification.root.model_dump(by_alias=True, mode="json", exclude_none=True) - ) - - # Create message handler - def message_handler( - message: RequestResponder[types.ServerRequest, types.ClientResult] | types.ServerNotification | Exception, - ) -> None: - if isinstance(message, Exception): - raise message - - # Start mock server thread - server_thread = threading.Thread(target=mock_server, daemon=True) - server_thread.start() - - # Create and use client session - with ClientSession( - server_to_client, - client_to_server, - message_handler=message_handler, - ) as session: - result = session.initialize() - - # Wait for server thread to complete - server_thread.join(timeout=10.0) - - # Assert results - assert isinstance(result, InitializeResult) - assert result.protocolVersion == LATEST_PROTOCOL_VERSION - assert isinstance(result.capabilities, ServerCapabilities) - assert result.serverInfo == Implementation(name="mock-server", version="0.1.0") - assert result.instructions == "The server instructions." - - # Check that client sent initialized notification - assert initialized_notification - assert isinstance(initialized_notification.root, InitializedNotification) - - -def test_client_session_custom_client_info() -> None: - # Create synchronous queues to replace async streams - client_to_server: queue.Queue[SessionMessage] = queue.Queue() - server_to_client: queue.Queue[SessionMessage] = queue.Queue() - - custom_client_info = Implementation(name="test-client", version="1.2.3") - received_client_info = None - - def mock_server() -> None: - nonlocal received_client_info - - session_message = client_to_server.get(timeout=5.0) - jsonrpc_request = session_message.message - assert isinstance(jsonrpc_request.root, JSONRPCRequest) - request = ClientRequest.model_validate( - jsonrpc_request.root.model_dump(by_alias=True, mode="json", exclude_none=True) - ) - assert isinstance(request.root, InitializeRequest) - received_client_info = request.root.params.clientInfo - - result = ServerResult( - InitializeResult( - protocolVersion=LATEST_PROTOCOL_VERSION, - capabilities=ServerCapabilities(), - serverInfo=Implementation(name="mock-server", version="0.1.0"), - ) - ) - - server_to_client.put( - SessionMessage( - message=JSONRPCMessage( - JSONRPCResponse( - jsonrpc="2.0", - id=jsonrpc_request.root.id, - result=result.model_dump(by_alias=True, mode="json", exclude_none=True), - ) - ) - ) - ) - # Receive initialized notification - client_to_server.get(timeout=5.0) - - # Start mock server thread - server_thread = threading.Thread(target=mock_server, daemon=True) - server_thread.start() - - with ClientSession( - server_to_client, - client_to_server, - client_info=custom_client_info, - ) as session: - session.initialize() - - # Wait for server thread to complete - server_thread.join(timeout=10.0) - - # Assert that custom client info was sent - assert received_client_info == custom_client_info - - -def test_client_session_default_client_info() -> None: - # Create synchronous queues to replace async streams - client_to_server: queue.Queue[SessionMessage] = queue.Queue() - server_to_client: queue.Queue[SessionMessage] = queue.Queue() - - received_client_info = None - - def mock_server() -> None: - nonlocal received_client_info - - session_message = client_to_server.get(timeout=5.0) - jsonrpc_request = session_message.message - assert isinstance(jsonrpc_request.root, JSONRPCRequest) - request = ClientRequest.model_validate( - jsonrpc_request.root.model_dump(by_alias=True, mode="json", exclude_none=True) - ) - assert isinstance(request.root, InitializeRequest) - received_client_info = request.root.params.clientInfo - - result = ServerResult( - InitializeResult( - protocolVersion=LATEST_PROTOCOL_VERSION, - capabilities=ServerCapabilities(), - serverInfo=Implementation(name="mock-server", version="0.1.0"), - ) - ) - - server_to_client.put( - SessionMessage( - message=JSONRPCMessage( - JSONRPCResponse( - jsonrpc="2.0", - id=jsonrpc_request.root.id, - result=result.model_dump(by_alias=True, mode="json", exclude_none=True), - ) - ) - ) - ) - # Receive initialized notification - client_to_server.get(timeout=5.0) - - # Start mock server thread - server_thread = threading.Thread(target=mock_server, daemon=True) - server_thread.start() - - with ClientSession( - server_to_client, - client_to_server, - ) as session: - session.initialize() - - # Wait for server thread to complete - server_thread.join(timeout=10.0) - - # Assert that default client info was used - assert received_client_info == DEFAULT_CLIENT_INFO - - -def test_client_session_version_negotiation_success() -> None: - # Create synchronous queues to replace async streams - client_to_server: queue.Queue[SessionMessage] = queue.Queue() - server_to_client: queue.Queue[SessionMessage] = queue.Queue() - - def mock_server() -> None: - session_message = client_to_server.get(timeout=5.0) - jsonrpc_request = session_message.message - assert isinstance(jsonrpc_request.root, JSONRPCRequest) - request = ClientRequest.model_validate( - jsonrpc_request.root.model_dump(by_alias=True, mode="json", exclude_none=True) - ) - assert isinstance(request.root, InitializeRequest) - - # Send supported protocol version - result = ServerResult( - InitializeResult( - protocolVersion=LATEST_PROTOCOL_VERSION, - capabilities=ServerCapabilities(), - serverInfo=Implementation(name="mock-server", version="0.1.0"), - ) - ) - - server_to_client.put( - SessionMessage( - message=JSONRPCMessage( - JSONRPCResponse( - jsonrpc="2.0", - id=jsonrpc_request.root.id, - result=result.model_dump(by_alias=True, mode="json", exclude_none=True), - ) - ) - ) - ) - # Receive initialized notification - client_to_server.get(timeout=5.0) - - # Start mock server thread - server_thread = threading.Thread(target=mock_server, daemon=True) - server_thread.start() - - with ClientSession( - server_to_client, - client_to_server, - ) as session: - result = session.initialize() - - # Wait for server thread to complete - server_thread.join(timeout=10.0) - - # Should successfully initialize - assert isinstance(result, InitializeResult) - assert result.protocolVersion == LATEST_PROTOCOL_VERSION - - -def test_client_session_version_negotiation_failure() -> None: - # Create synchronous queues to replace async streams - client_to_server: queue.Queue[SessionMessage] = queue.Queue() - server_to_client: queue.Queue[SessionMessage] = queue.Queue() - - def mock_server() -> None: - session_message = client_to_server.get(timeout=5.0) - jsonrpc_request = session_message.message - assert isinstance(jsonrpc_request.root, JSONRPCRequest) - request = ClientRequest.model_validate( - jsonrpc_request.root.model_dump(by_alias=True, mode="json", exclude_none=True) - ) - assert isinstance(request.root, InitializeRequest) - - # Send unsupported protocol version - result = ServerResult( - InitializeResult( - protocolVersion="99.99.99", # Unsupported version - capabilities=ServerCapabilities(), - serverInfo=Implementation(name="mock-server", version="0.1.0"), - ) - ) - - server_to_client.put( - SessionMessage( - message=JSONRPCMessage( - JSONRPCResponse( - jsonrpc="2.0", - id=jsonrpc_request.root.id, - result=result.model_dump(by_alias=True, mode="json", exclude_none=True), - ) - ) - ) - ) - - # Start mock server thread - server_thread = threading.Thread(target=mock_server, daemon=True) - server_thread.start() - - with ClientSession( - server_to_client, - client_to_server, - ) as session: - import pytest - - with pytest.raises(RuntimeError, match="Unsupported protocol version"): - session.initialize() - - # Wait for server thread to complete - server_thread.join(timeout=10.0) - - -def test_client_capabilities_default() -> None: - # Create synchronous queues to replace async streams - client_to_server: queue.Queue[SessionMessage] = queue.Queue() - server_to_client: queue.Queue[SessionMessage] = queue.Queue() - - received_capabilities = None - - def mock_server() -> None: - nonlocal received_capabilities - - session_message = client_to_server.get(timeout=5.0) - jsonrpc_request = session_message.message - assert isinstance(jsonrpc_request.root, JSONRPCRequest) - request = ClientRequest.model_validate( - jsonrpc_request.root.model_dump(by_alias=True, mode="json", exclude_none=True) - ) - assert isinstance(request.root, InitializeRequest) - received_capabilities = request.root.params.capabilities - - result = ServerResult( - InitializeResult( - protocolVersion=LATEST_PROTOCOL_VERSION, - capabilities=ServerCapabilities(), - serverInfo=Implementation(name="mock-server", version="0.1.0"), - ) - ) - - server_to_client.put( - SessionMessage( - message=JSONRPCMessage( - JSONRPCResponse( - jsonrpc="2.0", - id=jsonrpc_request.root.id, - result=result.model_dump(by_alias=True, mode="json", exclude_none=True), - ) - ) - ) - ) - # Receive initialized notification - client_to_server.get(timeout=5.0) - - # Start mock server thread - server_thread = threading.Thread(target=mock_server, daemon=True) - server_thread.start() - - with ClientSession( - server_to_client, - client_to_server, - ) as session: - session.initialize() - - # Wait for server thread to complete - server_thread.join(timeout=10.0) - - # Assert default capabilities - assert received_capabilities is not None - - -def test_client_capabilities_with_custom_callbacks() -> None: - # Create synchronous queues to replace async streams - client_to_server: queue.Queue[SessionMessage] = queue.Queue() - server_to_client: queue.Queue[SessionMessage] = queue.Queue() - - def custom_sampling_callback( - context: RequestContext["ClientSession", Any], - params: types.CreateMessageRequestParams, - ) -> types.CreateMessageResult | types.ErrorData: - return types.CreateMessageResult( - model="test-model", - role="assistant", - content=types.TextContent(type="text", text="Custom response"), - ) - - def custom_list_roots_callback( - context: RequestContext["ClientSession", Any], - ) -> types.ListRootsResult | types.ErrorData: - return types.ListRootsResult(roots=[]) - - def mock_server() -> None: - session_message = client_to_server.get(timeout=5.0) - jsonrpc_request = session_message.message - assert isinstance(jsonrpc_request.root, JSONRPCRequest) - request = ClientRequest.model_validate( - jsonrpc_request.root.model_dump(by_alias=True, mode="json", exclude_none=True) - ) - assert isinstance(request.root, InitializeRequest) - - result = ServerResult( - InitializeResult( - protocolVersion=LATEST_PROTOCOL_VERSION, - capabilities=ServerCapabilities(), - serverInfo=Implementation(name="mock-server", version="0.1.0"), - ) - ) - - server_to_client.put( - SessionMessage( - message=JSONRPCMessage( - JSONRPCResponse( - jsonrpc="2.0", - id=jsonrpc_request.root.id, - result=result.model_dump(by_alias=True, mode="json", exclude_none=True), - ) - ) - ) - ) - # Receive initialized notification - client_to_server.get(timeout=5.0) - - # Start mock server thread - server_thread = threading.Thread(target=mock_server, daemon=True) - server_thread.start() - - with ClientSession( - server_to_client, - client_to_server, - sampling_callback=custom_sampling_callback, - list_roots_callback=custom_list_roots_callback, - ) as session: - result = session.initialize() - - # Wait for server thread to complete - server_thread.join(timeout=10.0) - - # Verify initialization succeeded - assert isinstance(result, InitializeResult) - assert result.protocolVersion == LATEST_PROTOCOL_VERSION diff --git a/api/tests/unit_tests/core/mcp/client/test_sse.py b/api/tests/unit_tests/core/mcp/client/test_sse_client.py similarity index 100% rename from api/tests/unit_tests/core/mcp/client/test_sse.py rename to api/tests/unit_tests/core/mcp/client/test_sse_client.py diff --git a/api/tests/unit_tests/core/mcp/client/test_streamable_http.py b/api/tests/unit_tests/core/mcp/client/test_streamable_client.py similarity index 100% rename from api/tests/unit_tests/core/mcp/client/test_streamable_http.py rename to api/tests/unit_tests/core/mcp/client/test_streamable_client.py diff --git a/api/tests/unit_tests/core/mcp/session/test_client_session.py b/api/tests/unit_tests/core/mcp/session/test_client_session.py index 0dcccdcf539..ab990a98e42 100644 --- a/api/tests/unit_tests/core/mcp/session/test_client_session.py +++ b/api/tests/unit_tests/core/mcp/session/test_client_session.py @@ -1,586 +1,1066 @@ import queue +import threading +from typing import Any from unittest.mock import MagicMock import pytest from pydantic import AnyUrl from core.mcp import types -from core.mcp.session.base_session import RequestResponder, SessionMessage +from core.mcp.entities import RequestContext +from core.mcp.session.base_session import RequestResponder +from core.mcp.session.base_session import SessionMessage as _client_session_SessionMessage from core.mcp.session.client_session import ( + DEFAULT_CLIENT_INFO, ClientSession, _default_list_roots_callback, _default_logging_callback, _default_message_handler, _default_sampling_callback, ) +from core.mcp.types import ( + LATEST_PROTOCOL_VERSION, + ClientNotification, + ClientRequest, + Implementation, + InitializedNotification, + InitializeRequest, + InitializeResult, + JSONRPCMessage, + JSONRPCNotification, + JSONRPCRequest, + JSONRPCResponse, + ServerCapabilities, + ServerResult, +) +from core.mcp.types import ( + SessionMessage as _session_SessionMessage, +) -@pytest.fixture -def streams(): - return queue.Queue(), queue.Queue() +class TestSessionCases: + @staticmethod + def test_client_session_initialize() -> None: + # Create synchronous queues to replace async streams + client_to_server: queue.Queue[_session_SessionMessage] = queue.Queue() + server_to_client: queue.Queue[_session_SessionMessage] = queue.Queue() + initialized_notification = None -def test_client_session_init(streams): - read_stream, write_stream = streams - session = ClientSession(read_stream, write_stream) + def mock_server() -> None: + nonlocal initialized_notification - assert session._client_info.name == "Dify" - assert session._sampling_callback == _default_sampling_callback - assert session._list_roots_callback == _default_list_roots_callback - assert session._logging_callback == _default_logging_callback - assert session._message_handler == _default_message_handler + # Receive initialization request + session_message = client_to_server.get(timeout=5.0) + jsonrpc_request = session_message.message + assert isinstance(jsonrpc_request.root, JSONRPCRequest) + request = ClientRequest.model_validate( + jsonrpc_request.root.model_dump(by_alias=True, mode="json", exclude_none=True) + ) + assert isinstance(request.root, InitializeRequest) + # Create response + result = ServerResult( + InitializeResult( + protocolVersion=LATEST_PROTOCOL_VERSION, + capabilities=ServerCapabilities( + logging=None, + resources=None, + tools=None, + experimental=None, + prompts=None, + ), + serverInfo=Implementation(name="mock-server", version="0.1.0"), + instructions="The server instructions.", + ) + ) -def test_client_session_init_custom(streams): - read_stream, write_stream = streams - sampling_cb = MagicMock() - list_roots_cb = MagicMock() - logging_cb = MagicMock() - msg_handler = MagicMock() - client_info = types.Implementation(name="Custom", version="1.0") + # Send response + server_to_client.put( + _session_SessionMessage( + message=JSONRPCMessage( + JSONRPCResponse( + jsonrpc="2.0", + id=jsonrpc_request.root.id, + result=result.model_dump(by_alias=True, mode="json", exclude_none=True), + ) + ) + ) + ) - session = ClientSession( - read_stream, - write_stream, - sampling_callback=sampling_cb, - list_roots_callback=list_roots_cb, - logging_callback=logging_cb, - message_handler=msg_handler, - client_info=client_info, - ) + # Receive initialized notification + session_notification = client_to_server.get(timeout=5.0) + jsonrpc_notification = session_notification.message + assert isinstance(jsonrpc_notification.root, JSONRPCNotification) + initialized_notification = ClientNotification.model_validate( + jsonrpc_notification.root.model_dump(by_alias=True, mode="json", exclude_none=True) + ) - assert session._client_info == client_info - assert session._sampling_callback == sampling_cb - assert session._list_roots_callback == list_roots_cb - assert session._logging_callback == logging_cb - assert session._message_handler == msg_handler + # Create message handler + def message_handler( + message: RequestResponder[types.ServerRequest, types.ClientResult] | types.ServerNotification | Exception, + ) -> None: + if isinstance(message, Exception): + raise message + # Start mock server thread + server_thread = threading.Thread(target=mock_server, daemon=True) + server_thread.start() -def test_initialize_success(streams): - read_stream, write_stream = streams - session = ClientSession(read_stream, write_stream) + # Create and use client session + with ClientSession( + server_to_client, + client_to_server, + message_handler=message_handler, + ) as session: + result = session.initialize() - expected_result = types.InitializeResult( - protocolVersion=types.LATEST_PROTOCOL_VERSION, - capabilities=types.ServerCapabilities(), - serverInfo=types.Implementation(name="test-server", version="1.0"), - ) + # Wait for server thread to complete + server_thread.join(timeout=10.0) - def mock_server(): - # Handle initialize request - msg = write_stream.get(timeout=2) - req_id = msg.message.root.id + # Assert results + assert isinstance(result, InitializeResult) + assert result.protocolVersion == LATEST_PROTOCOL_VERSION + assert isinstance(result.capabilities, ServerCapabilities) + assert result.serverInfo == Implementation(name="mock-server", version="0.1.0") + assert result.instructions == "The server instructions." - resp = types.JSONRPCResponse(jsonrpc="2.0", id=req_id, result=expected_result.model_dump()) - read_stream.put(SessionMessage(message=types.JSONRPCMessage(resp))) + # Check that client sent initialized notification + assert initialized_notification + assert isinstance(initialized_notification.root, InitializedNotification) - # Expect initialized notification - notif = write_stream.get(timeout=2) - assert notif.message.root.method == "notifications/initialized" + @staticmethod + def test_client_session_custom_client_info() -> None: + # Create synchronous queues to replace async streams + client_to_server: queue.Queue[_session_SessionMessage] = queue.Queue() + server_to_client: queue.Queue[_session_SessionMessage] = queue.Queue() - import threading + custom_client_info = Implementation(name="test-client", version="1.2.3") + received_client_info = None - t = threading.Thread(target=mock_server, daemon=True) - t.start() + def mock_server() -> None: + nonlocal received_client_info - with session: - result = session.initialize() - assert result.protocolVersion == types.LATEST_PROTOCOL_VERSION - assert result.serverInfo.name == "test-server" + session_message = client_to_server.get(timeout=5.0) + jsonrpc_request = session_message.message + assert isinstance(jsonrpc_request.root, JSONRPCRequest) + request = ClientRequest.model_validate( + jsonrpc_request.root.model_dump(by_alias=True, mode="json", exclude_none=True) + ) + assert isinstance(request.root, InitializeRequest) + received_client_info = request.root.params.clientInfo - t.join(timeout=1) + result = ServerResult( + InitializeResult( + protocolVersion=LATEST_PROTOCOL_VERSION, + capabilities=ServerCapabilities(), + serverInfo=Implementation(name="mock-server", version="0.1.0"), + ) + ) + server_to_client.put( + _session_SessionMessage( + message=JSONRPCMessage( + JSONRPCResponse( + jsonrpc="2.0", + id=jsonrpc_request.root.id, + result=result.model_dump(by_alias=True, mode="json", exclude_none=True), + ) + ) + ) + ) + # Receive initialized notification + client_to_server.get(timeout=5.0) -def test_initialize_custom_capabilities(streams): - read_stream, write_stream = streams - session = ClientSession( - read_stream, write_stream, sampling_callback=lambda c, p: None, list_roots_callback=lambda c: None - ) + # Start mock server thread + server_thread = threading.Thread(target=mock_server, daemon=True) + server_thread.start() - def mock_server(): - msg = write_stream.get(timeout=2) - params = msg.message.root.params - # Check that capabilities are set because we provided custom callbacks - assert params["capabilities"]["sampling"] is not None - assert params["capabilities"]["roots"]["listChanged"] is True - - req_id = msg.message.root.id - resp = types.JSONRPCResponse( - jsonrpc="2.0", - id=req_id, - result={ - "protocolVersion": types.LATEST_PROTOCOL_VERSION, - "capabilities": {}, - "serverInfo": {"name": "test", "version": "1.0"}, - }, - ) - read_stream.put(SessionMessage(message=types.JSONRPCMessage(resp))) - write_stream.get(timeout=2) # initialized notif - - import threading - - t = threading.Thread(target=mock_server, daemon=True) - t.start() - - with session: - session.initialize() - t.join(timeout=1) - - -def test_initialize_unsupported_version(streams): - read_stream, write_stream = streams - session = ClientSession(read_stream, write_stream) - - def mock_server(): - msg = write_stream.get(timeout=2) - req_id = msg.message.root.id - resp = types.JSONRPCResponse( - jsonrpc="2.0", - id=req_id, - result={ - "protocolVersion": "0.0.1", # Unsupported - "capabilities": {}, - "serverInfo": {"name": "test", "version": "1.0"}, - }, - ) - read_stream.put(SessionMessage(message=types.JSONRPCMessage(resp))) - - import threading - - t = threading.Thread(target=mock_server, daemon=True) - t.start() - - with session: - with pytest.raises(RuntimeError, match="Unsupported protocol version"): + with ClientSession( + server_to_client, + client_to_server, + client_info=custom_client_info, + ) as session: session.initialize() - t.join(timeout=1) + # Wait for server thread to complete + server_thread.join(timeout=10.0) + + # Assert that custom client info was sent + assert received_client_info == custom_client_info + + @staticmethod + def test_client_session_default_client_info() -> None: + # Create synchronous queues to replace async streams + client_to_server: queue.Queue[_session_SessionMessage] = queue.Queue() + server_to_client: queue.Queue[_session_SessionMessage] = queue.Queue() + + received_client_info = None + + def mock_server() -> None: + nonlocal received_client_info + + session_message = client_to_server.get(timeout=5.0) + jsonrpc_request = session_message.message + assert isinstance(jsonrpc_request.root, JSONRPCRequest) + request = ClientRequest.model_validate( + jsonrpc_request.root.model_dump(by_alias=True, mode="json", exclude_none=True) + ) + assert isinstance(request.root, InitializeRequest) + received_client_info = request.root.params.clientInfo + + result = ServerResult( + InitializeResult( + protocolVersion=LATEST_PROTOCOL_VERSION, + capabilities=ServerCapabilities(), + serverInfo=Implementation(name="mock-server", version="0.1.0"), + ) + ) + + server_to_client.put( + _session_SessionMessage( + message=JSONRPCMessage( + JSONRPCResponse( + jsonrpc="2.0", + id=jsonrpc_request.root.id, + result=result.model_dump(by_alias=True, mode="json", exclude_none=True), + ) + ) + ) + ) + # Receive initialized notification + client_to_server.get(timeout=5.0) + + # Start mock server thread + server_thread = threading.Thread(target=mock_server, daemon=True) + server_thread.start() + + with ClientSession( + server_to_client, + client_to_server, + ) as session: + session.initialize() + + # Wait for server thread to complete + server_thread.join(timeout=10.0) + + # Assert that default client info was used + assert received_client_info == DEFAULT_CLIENT_INFO + + @staticmethod + def test_client_session_version_negotiation_success() -> None: + # Create synchronous queues to replace async streams + client_to_server: queue.Queue[_session_SessionMessage] = queue.Queue() + server_to_client: queue.Queue[_session_SessionMessage] = queue.Queue() + + def mock_server() -> None: + session_message = client_to_server.get(timeout=5.0) + jsonrpc_request = session_message.message + assert isinstance(jsonrpc_request.root, JSONRPCRequest) + request = ClientRequest.model_validate( + jsonrpc_request.root.model_dump(by_alias=True, mode="json", exclude_none=True) + ) + assert isinstance(request.root, InitializeRequest) + + # Send supported protocol version + result = ServerResult( + InitializeResult( + protocolVersion=LATEST_PROTOCOL_VERSION, + capabilities=ServerCapabilities(), + serverInfo=Implementation(name="mock-server", version="0.1.0"), + ) + ) + + server_to_client.put( + _session_SessionMessage( + message=JSONRPCMessage( + JSONRPCResponse( + jsonrpc="2.0", + id=jsonrpc_request.root.id, + result=result.model_dump(by_alias=True, mode="json", exclude_none=True), + ) + ) + ) + ) + # Receive initialized notification + client_to_server.get(timeout=5.0) + + # Start mock server thread + server_thread = threading.Thread(target=mock_server, daemon=True) + server_thread.start() + + with ClientSession( + server_to_client, + client_to_server, + ) as session: + result = session.initialize() + + # Wait for server thread to complete + server_thread.join(timeout=10.0) + + # Should successfully initialize + assert isinstance(result, InitializeResult) + assert result.protocolVersion == LATEST_PROTOCOL_VERSION + + @staticmethod + def test_client_session_version_negotiation_failure() -> None: + # Create synchronous queues to replace async streams + client_to_server: queue.Queue[_session_SessionMessage] = queue.Queue() + server_to_client: queue.Queue[_session_SessionMessage] = queue.Queue() + + def mock_server() -> None: + session_message = client_to_server.get(timeout=5.0) + jsonrpc_request = session_message.message + assert isinstance(jsonrpc_request.root, JSONRPCRequest) + request = ClientRequest.model_validate( + jsonrpc_request.root.model_dump(by_alias=True, mode="json", exclude_none=True) + ) + assert isinstance(request.root, InitializeRequest) + + # Send unsupported protocol version + result = ServerResult( + InitializeResult( + protocolVersion="99.99.99", # Unsupported version + capabilities=ServerCapabilities(), + serverInfo=Implementation(name="mock-server", version="0.1.0"), + ) + ) + + server_to_client.put( + _session_SessionMessage( + message=JSONRPCMessage( + JSONRPCResponse( + jsonrpc="2.0", + id=jsonrpc_request.root.id, + result=result.model_dump(by_alias=True, mode="json", exclude_none=True), + ) + ) + ) + ) + + # Start mock server thread + server_thread = threading.Thread(target=mock_server, daemon=True) + server_thread.start() + + with ClientSession( + server_to_client, + client_to_server, + ) as session: + import pytest + + with pytest.raises(RuntimeError, match="Unsupported protocol version"): + session.initialize() + + # Wait for server thread to complete + server_thread.join(timeout=10.0) + + @staticmethod + def test_client_capabilities_default() -> None: + # Create synchronous queues to replace async streams + client_to_server: queue.Queue[_session_SessionMessage] = queue.Queue() + server_to_client: queue.Queue[_session_SessionMessage] = queue.Queue() + + received_capabilities = None + + def mock_server() -> None: + nonlocal received_capabilities + + session_message = client_to_server.get(timeout=5.0) + jsonrpc_request = session_message.message + assert isinstance(jsonrpc_request.root, JSONRPCRequest) + request = ClientRequest.model_validate( + jsonrpc_request.root.model_dump(by_alias=True, mode="json", exclude_none=True) + ) + assert isinstance(request.root, InitializeRequest) + received_capabilities = request.root.params.capabilities + + result = ServerResult( + InitializeResult( + protocolVersion=LATEST_PROTOCOL_VERSION, + capabilities=ServerCapabilities(), + serverInfo=Implementation(name="mock-server", version="0.1.0"), + ) + ) + + server_to_client.put( + _session_SessionMessage( + message=JSONRPCMessage( + JSONRPCResponse( + jsonrpc="2.0", + id=jsonrpc_request.root.id, + result=result.model_dump(by_alias=True, mode="json", exclude_none=True), + ) + ) + ) + ) + # Receive initialized notification + client_to_server.get(timeout=5.0) + + # Start mock server thread + server_thread = threading.Thread(target=mock_server, daemon=True) + server_thread.start() + + with ClientSession( + server_to_client, + client_to_server, + ) as session: + session.initialize() + + # Wait for server thread to complete + server_thread.join(timeout=10.0) + + # Assert default capabilities + assert received_capabilities is not None + + @staticmethod + def test_client_capabilities_with_custom_callbacks() -> None: + # Create synchronous queues to replace async streams + client_to_server: queue.Queue[_session_SessionMessage] = queue.Queue() + server_to_client: queue.Queue[_session_SessionMessage] = queue.Queue() + + def custom_sampling_callback( + context: RequestContext["ClientSession", Any], + params: types.CreateMessageRequestParams, + ) -> types.CreateMessageResult | types.ErrorData: + return types.CreateMessageResult( + model="test-model", + role="assistant", + content=types.TextContent(type="text", text="Custom response"), + ) + + def custom_list_roots_callback( + context: RequestContext["ClientSession", Any], + ) -> types.ListRootsResult | types.ErrorData: + return types.ListRootsResult(roots=[]) + + def mock_server() -> None: + session_message = client_to_server.get(timeout=5.0) + jsonrpc_request = session_message.message + assert isinstance(jsonrpc_request.root, JSONRPCRequest) + request = ClientRequest.model_validate( + jsonrpc_request.root.model_dump(by_alias=True, mode="json", exclude_none=True) + ) + assert isinstance(request.root, InitializeRequest) + + result = ServerResult( + InitializeResult( + protocolVersion=LATEST_PROTOCOL_VERSION, + capabilities=ServerCapabilities(), + serverInfo=Implementation(name="mock-server", version="0.1.0"), + ) + ) + + server_to_client.put( + _session_SessionMessage( + message=JSONRPCMessage( + JSONRPCResponse( + jsonrpc="2.0", + id=jsonrpc_request.root.id, + result=result.model_dump(by_alias=True, mode="json", exclude_none=True), + ) + ) + ) + ) + # Receive initialized notification + client_to_server.get(timeout=5.0) + + # Start mock server thread + server_thread = threading.Thread(target=mock_server, daemon=True) + server_thread.start() + + with ClientSession( + server_to_client, + client_to_server, + sampling_callback=custom_sampling_callback, + list_roots_callback=custom_list_roots_callback, + ) as session: + result = session.initialize() + + # Wait for server thread to complete + server_thread.join(timeout=10.0) + + # Verify initialization succeeded + assert isinstance(result, InitializeResult) + assert result.protocolVersion == LATEST_PROTOCOL_VERSION -def test_send_ping(streams): - read_stream, write_stream = streams - session = ClientSession(read_stream, write_stream) - - def mock_server(): - msg = write_stream.get(timeout=2) - assert msg.message.root.method == "ping" - req_id = msg.message.root.id - resp = types.JSONRPCResponse(jsonrpc="2.0", id=req_id, result={}) - read_stream.put(SessionMessage(message=types.JSONRPCMessage(resp))) - - import threading - - t = threading.Thread(target=mock_server, daemon=True) - t.start() - - with session: - session.send_ping() - t.join(timeout=1) - - -def test_send_progress_notification(streams): - read_stream, write_stream = streams - session = ClientSession(read_stream, write_stream) - - session.send_progress_notification(progress_token="token", progress=50.0, total=100.0) - - msg = write_stream.get_nowait() - assert msg.message.root.method == "notifications/progress" - assert msg.message.root.params["progressToken"] == "token" - assert msg.message.root.params["progress"] == 50.0 - assert msg.message.root.params["total"] == 100.0 - - -def test_set_logging_level(streams): - read_stream, write_stream = streams - session = ClientSession(read_stream, write_stream) - - def mock_server(): - msg = write_stream.get(timeout=2) - assert msg.message.root.method == "logging/setLevel" - assert msg.message.root.params["level"] == "debug" - req_id = msg.message.root.id - resp = types.JSONRPCResponse(jsonrpc="2.0", id=req_id, result={}) - read_stream.put(SessionMessage(message=types.JSONRPCMessage(resp))) - - import threading - - t = threading.Thread(target=mock_server, daemon=True) - t.start() - - with session: - session.set_logging_level("debug") - t.join(timeout=1) - - -def test_list_resources(streams): - read_stream, write_stream = streams - session = ClientSession(read_stream, write_stream) - - def mock_server(): - msg = write_stream.get(timeout=2) - assert msg.message.root.method == "resources/list" - req_id = msg.message.root.id - resp = types.JSONRPCResponse(jsonrpc="2.0", id=req_id, result={"resources": []}) - read_stream.put(SessionMessage(message=types.JSONRPCMessage(resp))) - - import threading - - t = threading.Thread(target=mock_server, daemon=True) - t.start() - - with session: - result = session.list_resources() - assert result.resources == [] - t.join(timeout=1) - - -def test_list_resource_templates(streams): - read_stream, write_stream = streams - session = ClientSession(read_stream, write_stream) - - def mock_server(): - msg = write_stream.get(timeout=2) - assert msg.message.root.method == "resources/templates/list" - req_id = msg.message.root.id - resp = types.JSONRPCResponse(jsonrpc="2.0", id=req_id, result={"resourceTemplates": []}) - read_stream.put(SessionMessage(message=types.JSONRPCMessage(resp))) - - import threading - - t = threading.Thread(target=mock_server, daemon=True) - t.start() - - with session: - result = session.list_resource_templates() - assert result.resourceTemplates == [] - t.join(timeout=1) - - -def test_read_resource(streams): - read_stream, write_stream = streams - session = ClientSession(read_stream, write_stream) - uri = AnyUrl("file:///test") - - def mock_server(): - msg = write_stream.get(timeout=2) - assert msg.message.root.method == "resources/read" - assert msg.message.root.params["uri"] == str(uri) - req_id = msg.message.root.id - resp = types.JSONRPCResponse(jsonrpc="2.0", id=req_id, result={"contents": []}) - read_stream.put(SessionMessage(message=types.JSONRPCMessage(resp))) - - import threading - - t = threading.Thread(target=mock_server, daemon=True) - t.start() - - with session: - result = session.read_resource(uri) - assert result.contents == [] - t.join(timeout=1) - - -def test_subscribe_resource(streams): - read_stream, write_stream = streams - session = ClientSession(read_stream, write_stream) - uri = AnyUrl("file:///test") - - def mock_server(): - msg = write_stream.get(timeout=2) - assert msg.message.root.method == "resources/subscribe" - assert msg.message.root.params["uri"] == str(uri) - req_id = msg.message.root.id - resp = types.JSONRPCResponse(jsonrpc="2.0", id=req_id, result={}) - read_stream.put(SessionMessage(message=types.JSONRPCMessage(resp))) - - import threading - - t = threading.Thread(target=mock_server, daemon=True) - t.start() - - with session: - session.subscribe_resource(uri) - t.join(timeout=1) - - -def test_unsubscribe_resource(streams): - read_stream, write_stream = streams - session = ClientSession(read_stream, write_stream) - uri = AnyUrl("file:///test") - - def mock_server(): - msg = write_stream.get(timeout=2) - assert msg.message.root.method == "resources/unsubscribe" - assert msg.message.root.params["uri"] == str(uri) - req_id = msg.message.root.id - resp = types.JSONRPCResponse(jsonrpc="2.0", id=req_id, result={}) - read_stream.put(SessionMessage(message=types.JSONRPCMessage(resp))) - - import threading - - t = threading.Thread(target=mock_server, daemon=True) - t.start() - - with session: - session.unsubscribe_resource(uri) - t.join(timeout=1) - - -def test_call_tool(streams): - read_stream, write_stream = streams - session = ClientSession(read_stream, write_stream) - - def mock_server(): - msg = write_stream.get(timeout=2) - assert msg.message.root.method == "tools/call" - assert msg.message.root.params["name"] == "test-tool" - assert msg.message.root.params["arguments"] == {"arg": 1} - req_id = msg.message.root.id - resp = types.JSONRPCResponse(jsonrpc="2.0", id=req_id, result={"content": [], "isError": False}) - read_stream.put(SessionMessage(message=types.JSONRPCMessage(resp))) - - import threading - - t = threading.Thread(target=mock_server, daemon=True) - t.start() - - with session: - result = session.call_tool("test-tool", arguments={"arg": 1}) - assert result.isError is False - t.join(timeout=1) - - -def test_list_prompts(streams): - read_stream, write_stream = streams - session = ClientSession(read_stream, write_stream) - - def mock_server(): - msg = write_stream.get(timeout=2) - assert msg.message.root.method == "prompts/list" - req_id = msg.message.root.id - resp = types.JSONRPCResponse(jsonrpc="2.0", id=req_id, result={"prompts": []}) - read_stream.put(SessionMessage(message=types.JSONRPCMessage(resp))) - - import threading - - t = threading.Thread(target=mock_server, daemon=True) - t.start() - - with session: - result = session.list_prompts() - assert result.prompts == [] - t.join(timeout=1) - - -def test_get_prompt(streams): - read_stream, write_stream = streams - session = ClientSession(read_stream, write_stream) - - def mock_server(): - msg = write_stream.get(timeout=2) - assert msg.message.root.method == "prompts/get" - assert msg.message.root.params["name"] == "test-prompt" - req_id = msg.message.root.id - resp = types.JSONRPCResponse(jsonrpc="2.0", id=req_id, result={"messages": []}) - read_stream.put(SessionMessage(message=types.JSONRPCMessage(resp))) - - import threading - - t = threading.Thread(target=mock_server, daemon=True) - t.start() - - with session: - result = session.get_prompt("test-prompt") - assert result.messages == [] - t.join(timeout=1) - - -def test_complete(streams): - read_stream, write_stream = streams - session = ClientSession(read_stream, write_stream) - ref = types.PromptReference(type="ref/prompt", name="test") - - def mock_server(): - msg = write_stream.get(timeout=2) - assert msg.message.root.method == "completion/complete" - req_id = msg.message.root.id - resp = types.JSONRPCResponse(jsonrpc="2.0", id=req_id, result={"completion": {"values": [], "hasMore": False}}) - read_stream.put(SessionMessage(message=types.JSONRPCMessage(resp))) - - import threading - - t = threading.Thread(target=mock_server, daemon=True) - t.start() - - with session: - result = session.complete(ref, argument={"name": "val", "value": "x"}) - assert result.completion.hasMore is False - t.join(timeout=1) - - -def test_list_tools(streams): - read_stream, write_stream = streams - session = ClientSession(read_stream, write_stream) - - def mock_server(): - msg = write_stream.get(timeout=2) - assert msg.message.root.method == "tools/list" - req_id = msg.message.root.id - resp = types.JSONRPCResponse(jsonrpc="2.0", id=req_id, result={"tools": []}) - read_stream.put(SessionMessage(message=types.JSONRPCMessage(resp))) - - import threading - - t = threading.Thread(target=mock_server, daemon=True) - t.start() - - with session: - result = session.list_tools() - assert result.tools == [] - t.join(timeout=1) - - -def test_send_roots_list_changed(streams): - read_stream, write_stream = streams - session = ClientSession(read_stream, write_stream) - - session.send_roots_list_changed() - - msg = write_stream.get_nowait() - assert msg.message.root.method == "notifications/roots/list_changed" - - -def test_received_request_sampling(streams): - read_stream, write_stream = streams - sampling_cb = MagicMock( - return_value=types.CreateMessageResult( - role="assistant", content=types.TextContent(type="text", text="hello"), model="gpt-4" - ) - ) - session = ClientSession(read_stream, write_stream, sampling_callback=sampling_cb) - - req = types.ServerRequest( - root=types.CreateMessageRequest( - method="sampling/createMessage", params=types.CreateMessageRequestParams(messages=[], maxTokens=100) - ) - ) - - responder = RequestResponder(request_id=1, request_meta=None, request=req, session=session, on_complete=MagicMock()) - - session._received_request(responder) - - msg = write_stream.get_nowait() - assert msg.message.root.result["model"] == "gpt-4" - sampling_cb.assert_called_once() - - -def test_received_request_list_roots(streams): - read_stream, write_stream = streams - list_roots_cb = MagicMock(return_value=types.ListRootsResult(roots=[])) - session = ClientSession(read_stream, write_stream, list_roots_callback=list_roots_cb) - - req = types.ServerRequest(root=types.ListRootsRequest(method="roots/list")) - - responder = RequestResponder(request_id=1, request_meta=None, request=req, session=session, on_complete=MagicMock()) - - session._received_request(responder) - - msg = write_stream.get_nowait() - assert msg.message.root.result["roots"] == [] - list_roots_cb.assert_called_once() - - -def test_received_request_ping(streams): - read_stream, write_stream = streams - session = ClientSession(read_stream, write_stream) - - req = types.ServerRequest(root=types.PingRequest(method="ping")) - - responder = RequestResponder(request_id=1, request_meta=None, request=req, session=session, on_complete=MagicMock()) - - session._received_request(responder) - - msg = write_stream.get_nowait() - assert msg.message.root.result == {} - - -def test_handle_incoming(streams): - read_stream, write_stream = streams - msg_handler = MagicMock() - session = ClientSession(read_stream, write_stream, message_handler=msg_handler) - - item = MagicMock() - session._handle_incoming(item) - msg_handler.assert_called_once_with(item) - - -def test_received_notification_logging(streams): - read_stream, write_stream = streams - logging_cb = MagicMock() - session = ClientSession(read_stream, write_stream, logging_callback=logging_cb) - - notif = types.ServerNotification( - root=types.LoggingMessageNotification( - method="notifications/message", - params=types.LoggingMessageNotificationParams(level="info", data={"msg": "test"}), - ) - ) - - session._received_notification(notif) - logging_cb.assert_called_once() - assert logging_cb.call_args[0][0].level == "info" - - -def test_default_message_handler(streams): - # Exception case - with pytest.raises(ValueError, match="test error"): - _default_message_handler(Exception("test error")) - - # Notification case - should do nothing - _default_message_handler(types.ServerNotification(root=types.ToolListChangedNotification())) - - # RequestResponder case - should do nothing - read_stream, write_stream = streams - session = ClientSession(read_stream, write_stream) - with RequestResponder( - request_id=1, - request_meta=None, - request=types.ServerRequest(root=types.PingRequest()), - session=session, - on_complete=lambda _responder: None, - ) as responder: - _default_message_handler(responder) - assert not responder.completed - - -def test_default_sampling_callback(): - ctx = MagicMock() - params = MagicMock() - res = _default_sampling_callback(ctx, params) - assert res.code == types.INVALID_REQUEST - assert "not supported" in res.message - - -def test_default_list_roots_callback(): - ctx = MagicMock() - res = _default_list_roots_callback(ctx) - assert res.code == types.INVALID_REQUEST - assert "not supported" in res.message - - -def test_default_logging_callback(): - params = MagicMock() - _default_logging_callback(params) # Should do nothing - - -def test_received_notification_unknown(streams): - read_stream, write_stream = streams - session = ClientSession(read_stream, write_stream) - - # Use a notification type that is NOT LoggingMessageNotification - notif = types.ServerNotification( - root=types.ResourceListChangedNotification(method="notifications/resources/list_changed") - ) - - session._received_notification(notif) # Should just pass (case _:) + + +class TestClientSessionCases: + @staticmethod + @pytest.fixture + def streams(): + return queue.Queue(), queue.Queue() + + @staticmethod + def test_client_session_init(streams): + read_stream, write_stream = streams + session = ClientSession(read_stream, write_stream) + + assert session._client_info.name == "Dify" + assert session._sampling_callback == _default_sampling_callback + assert session._list_roots_callback == _default_list_roots_callback + assert session._logging_callback == _default_logging_callback + assert session._message_handler == _default_message_handler + + @staticmethod + def test_client_session_init_custom(streams): + read_stream, write_stream = streams + sampling_cb = MagicMock() + list_roots_cb = MagicMock() + logging_cb = MagicMock() + msg_handler = MagicMock() + client_info = types.Implementation(name="Custom", version="1.0") + + session = ClientSession( + read_stream, + write_stream, + sampling_callback=sampling_cb, + list_roots_callback=list_roots_cb, + logging_callback=logging_cb, + message_handler=msg_handler, + client_info=client_info, + ) + + assert session._client_info == client_info + assert session._sampling_callback == sampling_cb + assert session._list_roots_callback == list_roots_cb + assert session._logging_callback == logging_cb + assert session._message_handler == msg_handler + + @staticmethod + def test_initialize_success(streams): + read_stream, write_stream = streams + session = ClientSession(read_stream, write_stream) + + expected_result = types.InitializeResult( + protocolVersion=types.LATEST_PROTOCOL_VERSION, + capabilities=types.ServerCapabilities(), + serverInfo=types.Implementation(name="test-server", version="1.0"), + ) + + def mock_server(): + # Handle initialize request + msg = write_stream.get(timeout=2) + req_id = msg.message.root.id + + resp = types.JSONRPCResponse(jsonrpc="2.0", id=req_id, result=expected_result.model_dump()) + read_stream.put(_client_session_SessionMessage(message=types.JSONRPCMessage(resp))) + + # Expect initialized notification + notif = write_stream.get(timeout=2) + assert notif.message.root.method == "notifications/initialized" + + import threading + + t = threading.Thread(target=mock_server, daemon=True) + t.start() + + with session: + result = session.initialize() + assert result.protocolVersion == types.LATEST_PROTOCOL_VERSION + assert result.serverInfo.name == "test-server" + + t.join(timeout=1) + + @staticmethod + def test_initialize_custom_capabilities(streams): + read_stream, write_stream = streams + session = ClientSession( + read_stream, write_stream, sampling_callback=lambda c, p: None, list_roots_callback=lambda c: None + ) + + def mock_server(): + msg = write_stream.get(timeout=2) + params = msg.message.root.params + # Check that capabilities are set because we provided custom callbacks + assert params["capabilities"]["sampling"] is not None + assert params["capabilities"]["roots"]["listChanged"] is True + + req_id = msg.message.root.id + resp = types.JSONRPCResponse( + jsonrpc="2.0", + id=req_id, + result={ + "protocolVersion": types.LATEST_PROTOCOL_VERSION, + "capabilities": {}, + "serverInfo": {"name": "test", "version": "1.0"}, + }, + ) + read_stream.put(_client_session_SessionMessage(message=types.JSONRPCMessage(resp))) + write_stream.get(timeout=2) # initialized notif + + import threading + + t = threading.Thread(target=mock_server, daemon=True) + t.start() + + with session: + session.initialize() + t.join(timeout=1) + + @staticmethod + def test_initialize_unsupported_version(streams): + read_stream, write_stream = streams + session = ClientSession(read_stream, write_stream) + + def mock_server(): + msg = write_stream.get(timeout=2) + req_id = msg.message.root.id + resp = types.JSONRPCResponse( + jsonrpc="2.0", + id=req_id, + result={ + "protocolVersion": "0.0.1", # Unsupported + "capabilities": {}, + "serverInfo": {"name": "test", "version": "1.0"}, + }, + ) + read_stream.put(_client_session_SessionMessage(message=types.JSONRPCMessage(resp))) + + import threading + + t = threading.Thread(target=mock_server, daemon=True) + t.start() + + with session: + with pytest.raises(RuntimeError, match="Unsupported protocol version"): + session.initialize() + t.join(timeout=1) + + @staticmethod + def test_send_ping(streams): + read_stream, write_stream = streams + session = ClientSession(read_stream, write_stream) + + def mock_server(): + msg = write_stream.get(timeout=2) + assert msg.message.root.method == "ping" + req_id = msg.message.root.id + resp = types.JSONRPCResponse(jsonrpc="2.0", id=req_id, result={}) + read_stream.put(_client_session_SessionMessage(message=types.JSONRPCMessage(resp))) + + import threading + + t = threading.Thread(target=mock_server, daemon=True) + t.start() + + with session: + session.send_ping() + t.join(timeout=1) + + @staticmethod + def test_send_progress_notification(streams): + read_stream, write_stream = streams + session = ClientSession(read_stream, write_stream) + + session.send_progress_notification(progress_token="token", progress=50.0, total=100.0) + + msg = write_stream.get_nowait() + assert msg.message.root.method == "notifications/progress" + assert msg.message.root.params["progressToken"] == "token" + assert msg.message.root.params["progress"] == 50.0 + assert msg.message.root.params["total"] == 100.0 + + @staticmethod + def test_set_logging_level(streams): + read_stream, write_stream = streams + session = ClientSession(read_stream, write_stream) + + def mock_server(): + msg = write_stream.get(timeout=2) + assert msg.message.root.method == "logging/setLevel" + assert msg.message.root.params["level"] == "debug" + req_id = msg.message.root.id + resp = types.JSONRPCResponse(jsonrpc="2.0", id=req_id, result={}) + read_stream.put(_client_session_SessionMessage(message=types.JSONRPCMessage(resp))) + + import threading + + t = threading.Thread(target=mock_server, daemon=True) + t.start() + + with session: + session.set_logging_level("debug") + t.join(timeout=1) + + @staticmethod + def test_list_resources(streams): + read_stream, write_stream = streams + session = ClientSession(read_stream, write_stream) + + def mock_server(): + msg = write_stream.get(timeout=2) + assert msg.message.root.method == "resources/list" + req_id = msg.message.root.id + resp = types.JSONRPCResponse(jsonrpc="2.0", id=req_id, result={"resources": []}) + read_stream.put(_client_session_SessionMessage(message=types.JSONRPCMessage(resp))) + + import threading + + t = threading.Thread(target=mock_server, daemon=True) + t.start() + + with session: + result = session.list_resources() + assert result.resources == [] + t.join(timeout=1) + + @staticmethod + def test_list_resource_templates(streams): + read_stream, write_stream = streams + session = ClientSession(read_stream, write_stream) + + def mock_server(): + msg = write_stream.get(timeout=2) + assert msg.message.root.method == "resources/templates/list" + req_id = msg.message.root.id + resp = types.JSONRPCResponse(jsonrpc="2.0", id=req_id, result={"resourceTemplates": []}) + read_stream.put(_client_session_SessionMessage(message=types.JSONRPCMessage(resp))) + + import threading + + t = threading.Thread(target=mock_server, daemon=True) + t.start() + + with session: + result = session.list_resource_templates() + assert result.resourceTemplates == [] + t.join(timeout=1) + + @staticmethod + def test_read_resource(streams): + read_stream, write_stream = streams + session = ClientSession(read_stream, write_stream) + uri = AnyUrl("file:///test") + + def mock_server(): + msg = write_stream.get(timeout=2) + assert msg.message.root.method == "resources/read" + assert msg.message.root.params["uri"] == str(uri) + req_id = msg.message.root.id + resp = types.JSONRPCResponse(jsonrpc="2.0", id=req_id, result={"contents": []}) + read_stream.put(_client_session_SessionMessage(message=types.JSONRPCMessage(resp))) + + import threading + + t = threading.Thread(target=mock_server, daemon=True) + t.start() + + with session: + result = session.read_resource(uri) + assert result.contents == [] + t.join(timeout=1) + + @staticmethod + def test_subscribe_resource(streams): + read_stream, write_stream = streams + session = ClientSession(read_stream, write_stream) + uri = AnyUrl("file:///test") + + def mock_server(): + msg = write_stream.get(timeout=2) + assert msg.message.root.method == "resources/subscribe" + assert msg.message.root.params["uri"] == str(uri) + req_id = msg.message.root.id + resp = types.JSONRPCResponse(jsonrpc="2.0", id=req_id, result={}) + read_stream.put(_client_session_SessionMessage(message=types.JSONRPCMessage(resp))) + + import threading + + t = threading.Thread(target=mock_server, daemon=True) + t.start() + + with session: + session.subscribe_resource(uri) + t.join(timeout=1) + + @staticmethod + def test_unsubscribe_resource(streams): + read_stream, write_stream = streams + session = ClientSession(read_stream, write_stream) + uri = AnyUrl("file:///test") + + def mock_server(): + msg = write_stream.get(timeout=2) + assert msg.message.root.method == "resources/unsubscribe" + assert msg.message.root.params["uri"] == str(uri) + req_id = msg.message.root.id + resp = types.JSONRPCResponse(jsonrpc="2.0", id=req_id, result={}) + read_stream.put(_client_session_SessionMessage(message=types.JSONRPCMessage(resp))) + + import threading + + t = threading.Thread(target=mock_server, daemon=True) + t.start() + + with session: + session.unsubscribe_resource(uri) + t.join(timeout=1) + + @staticmethod + def test_call_tool(streams): + read_stream, write_stream = streams + session = ClientSession(read_stream, write_stream) + + def mock_server(): + msg = write_stream.get(timeout=2) + assert msg.message.root.method == "tools/call" + assert msg.message.root.params["name"] == "test-tool" + assert msg.message.root.params["arguments"] == {"arg": 1} + req_id = msg.message.root.id + resp = types.JSONRPCResponse(jsonrpc="2.0", id=req_id, result={"content": [], "isError": False}) + read_stream.put(_client_session_SessionMessage(message=types.JSONRPCMessage(resp))) + + import threading + + t = threading.Thread(target=mock_server, daemon=True) + t.start() + + with session: + result = session.call_tool("test-tool", arguments={"arg": 1}) + assert result.isError is False + t.join(timeout=1) + + @staticmethod + def test_list_prompts(streams): + read_stream, write_stream = streams + session = ClientSession(read_stream, write_stream) + + def mock_server(): + msg = write_stream.get(timeout=2) + assert msg.message.root.method == "prompts/list" + req_id = msg.message.root.id + resp = types.JSONRPCResponse(jsonrpc="2.0", id=req_id, result={"prompts": []}) + read_stream.put(_client_session_SessionMessage(message=types.JSONRPCMessage(resp))) + + import threading + + t = threading.Thread(target=mock_server, daemon=True) + t.start() + + with session: + result = session.list_prompts() + assert result.prompts == [] + t.join(timeout=1) + + @staticmethod + def test_get_prompt(streams): + read_stream, write_stream = streams + session = ClientSession(read_stream, write_stream) + + def mock_server(): + msg = write_stream.get(timeout=2) + assert msg.message.root.method == "prompts/get" + assert msg.message.root.params["name"] == "test-prompt" + req_id = msg.message.root.id + resp = types.JSONRPCResponse(jsonrpc="2.0", id=req_id, result={"messages": []}) + read_stream.put(_client_session_SessionMessage(message=types.JSONRPCMessage(resp))) + + import threading + + t = threading.Thread(target=mock_server, daemon=True) + t.start() + + with session: + result = session.get_prompt("test-prompt") + assert result.messages == [] + t.join(timeout=1) + + @staticmethod + def test_complete(streams): + read_stream, write_stream = streams + session = ClientSession(read_stream, write_stream) + ref = types.PromptReference(type="ref/prompt", name="test") + + def mock_server(): + msg = write_stream.get(timeout=2) + assert msg.message.root.method == "completion/complete" + req_id = msg.message.root.id + resp = types.JSONRPCResponse( + jsonrpc="2.0", id=req_id, result={"completion": {"values": [], "hasMore": False}} + ) + read_stream.put(_client_session_SessionMessage(message=types.JSONRPCMessage(resp))) + + import threading + + t = threading.Thread(target=mock_server, daemon=True) + t.start() + + with session: + result = session.complete(ref, argument={"name": "val", "value": "x"}) + assert result.completion.hasMore is False + t.join(timeout=1) + + @staticmethod + def test_list_tools(streams): + read_stream, write_stream = streams + session = ClientSession(read_stream, write_stream) + + def mock_server(): + msg = write_stream.get(timeout=2) + assert msg.message.root.method == "tools/list" + req_id = msg.message.root.id + resp = types.JSONRPCResponse(jsonrpc="2.0", id=req_id, result={"tools": []}) + read_stream.put(_client_session_SessionMessage(message=types.JSONRPCMessage(resp))) + + import threading + + t = threading.Thread(target=mock_server, daemon=True) + t.start() + + with session: + result = session.list_tools() + assert result.tools == [] + t.join(timeout=1) + + @staticmethod + def test_send_roots_list_changed(streams): + read_stream, write_stream = streams + session = ClientSession(read_stream, write_stream) + + session.send_roots_list_changed() + + msg = write_stream.get_nowait() + assert msg.message.root.method == "notifications/roots/list_changed" + + @staticmethod + def test_received_request_sampling(streams): + read_stream, write_stream = streams + sampling_cb = MagicMock( + return_value=types.CreateMessageResult( + role="assistant", content=types.TextContent(type="text", text="hello"), model="gpt-4" + ) + ) + session = ClientSession(read_stream, write_stream, sampling_callback=sampling_cb) + + req = types.ServerRequest( + root=types.CreateMessageRequest( + method="sampling/createMessage", params=types.CreateMessageRequestParams(messages=[], maxTokens=100) + ) + ) + + responder = RequestResponder( + request_id=1, request_meta=None, request=req, session=session, on_complete=MagicMock() + ) + + session._received_request(responder) + + msg = write_stream.get_nowait() + assert msg.message.root.result["model"] == "gpt-4" + sampling_cb.assert_called_once() + + @staticmethod + def test_received_request_list_roots(streams): + read_stream, write_stream = streams + list_roots_cb = MagicMock(return_value=types.ListRootsResult(roots=[])) + session = ClientSession(read_stream, write_stream, list_roots_callback=list_roots_cb) + + req = types.ServerRequest(root=types.ListRootsRequest(method="roots/list")) + + responder = RequestResponder( + request_id=1, request_meta=None, request=req, session=session, on_complete=MagicMock() + ) + + session._received_request(responder) + + msg = write_stream.get_nowait() + assert msg.message.root.result["roots"] == [] + list_roots_cb.assert_called_once() + + @staticmethod + def test_received_request_ping(streams): + read_stream, write_stream = streams + session = ClientSession(read_stream, write_stream) + + req = types.ServerRequest(root=types.PingRequest(method="ping")) + + responder = RequestResponder( + request_id=1, request_meta=None, request=req, session=session, on_complete=MagicMock() + ) + + session._received_request(responder) + + msg = write_stream.get_nowait() + assert msg.message.root.result == {} + + @staticmethod + def test_handle_incoming(streams): + read_stream, write_stream = streams + msg_handler = MagicMock() + session = ClientSession(read_stream, write_stream, message_handler=msg_handler) + + item = MagicMock() + session._handle_incoming(item) + msg_handler.assert_called_once_with(item) + + @staticmethod + def test_received_notification_logging(streams): + read_stream, write_stream = streams + logging_cb = MagicMock() + session = ClientSession(read_stream, write_stream, logging_callback=logging_cb) + + notif = types.ServerNotification( + root=types.LoggingMessageNotification( + method="notifications/message", + params=types.LoggingMessageNotificationParams(level="info", data={"msg": "test"}), + ) + ) + + session._received_notification(notif) + logging_cb.assert_called_once() + assert logging_cb.call_args[0][0].level == "info" + + @staticmethod + def test_default_message_handler(streams): + # Exception case + with pytest.raises(ValueError, match="test error"): + _default_message_handler(Exception("test error")) + + # Notification case - should do nothing + _default_message_handler(types.ServerNotification(root=types.ToolListChangedNotification())) + + # RequestResponder case - should do nothing + read_stream, write_stream = streams + session = ClientSession(read_stream, write_stream) + with RequestResponder( + request_id=1, + request_meta=None, + request=types.ServerRequest(root=types.PingRequest()), + session=session, + on_complete=lambda _responder: None, + ) as responder: + _default_message_handler(responder) + assert not responder.completed + + @staticmethod + def test_default_sampling_callback(): + ctx = MagicMock() + params = MagicMock() + res = _default_sampling_callback(ctx, params) + assert res.code == types.INVALID_REQUEST + assert "not supported" in res.message + + @staticmethod + def test_default_list_roots_callback(): + ctx = MagicMock() + res = _default_list_roots_callback(ctx) + assert res.code == types.INVALID_REQUEST + assert "not supported" in res.message + + @staticmethod + def test_default_logging_callback(): + params = MagicMock() + _default_logging_callback(params) # Should do nothing + + @staticmethod + def test_received_notification_unknown(streams): + read_stream, write_stream = streams + session = ClientSession(read_stream, write_stream) + + # Use a notification type that is NOT LoggingMessageNotification + notif = types.ServerNotification( + root=types.ResourceListChangedNotification(method="notifications/resources/list_changed") + ) + + session._received_notification(notif) diff --git a/api/tests/unit_tests/core/mcp/test_auth_client_inheritance.py b/api/tests/unit_tests/core/mcp/test_auth_client.py similarity index 100% rename from api/tests/unit_tests/core/mcp/test_auth_client_inheritance.py rename to api/tests/unit_tests/core/mcp/test_auth_client.py diff --git a/api/tests/unit_tests/core/mcp/test_mcp_client.py b/api/tests/unit_tests/core/mcp/test_mcp_client.py index 3a62d571047..6176feb1393 100644 --- a/api/tests/unit_tests/core/mcp/test_mcp_client.py +++ b/api/tests/unit_tests/core/mcp/test_mcp_client.py @@ -1,6 +1,8 @@ -"""Unit tests for MCP client.""" - -from contextlib import ExitStack +from collections.abc import Generator, Mapping +from contextlib import ExitStack, contextmanager +from datetime import timedelta +from queue import Queue +from threading import Thread from types import SimpleNamespace, TracebackType from unittest.mock import MagicMock, Mock, patch @@ -14,8 +16,24 @@ from core.entities.mcp_provider import MCPProviderEntity from core.mcp.auth_client import MCPClientWithAuthRetry from core.mcp.error import MCPAuthError, MCPConnectionError from core.mcp.mcp_client import MCPClient -from core.mcp.types import CallToolResult, ListToolsResult, TextContent, Tool, ToolAnnotations -from tests.unit_tests.core.mcp.fixtures import make_provider +from core.mcp.session.client_session import ClientSession +from core.mcp.types import ( + CallToolResult, + ErrorData, + JSONRPCError, + JSONRPCMessage, + JSONRPCRequest, + JSONRPCResponse, + ListToolsRequest, + ListToolsResult, + SessionMessage, + TextContent, + Tool, + ToolAnnotations, +) +from tests.helpers.unit.core.mcp.fixtures import make_provider + +"""Unit tests for MCP client.""" class TestMCPClient: @@ -805,3 +823,261 @@ class TestMCPClientWithAuthRetry: assert mock_execute_retry.call_args[0][0].__name__ == "invoke_tool" assert mock_execute_retry.call_args[0][1] == "test-tool" assert mock_execute_retry.call_args[0][2] == {"arg": "val"} + + +"""Offline pagination regressions using the real client and in-memory JSON-RPC.""" + + +def tool(name: str) -> Tool: + return Tool(name=name, description=name, inputSchema={"type": "object", "properties": {}}) + + +@contextmanager +def client_for_pages( + pages: Mapping[str | None, ListToolsResult | ErrorData], client: MCPClient | None = None +) -> Generator[tuple[MCPClient, list[ListToolsRequest]]]: + incoming: Queue[SessionMessage | Exception | None] = Queue() + outgoing: Queue[SessionMessage | None] = Queue() + requests: list[ListToolsRequest] = [] + errors: list[Exception] = [] + + def serve_pages() -> None: + while (message := outgoing.get()) is not None: + request = message.message.root + assert isinstance(request, JSONRPCRequest) + try: + payload = ListToolsRequest.model_validate({"method": request.method, "params": request.params}) + requests.append(payload) + cursor = payload.params.cursor if payload.params is not None else None + page = pages[cursor] + response: JSONRPCResponse | JSONRPCError + if isinstance(page, ErrorData): + response = JSONRPCError(jsonrpc="2.0", id=request.id, error=page) + else: + response = JSONRPCResponse( + jsonrpc="2.0", + id=request.id, + result=page.model_dump(mode="json", by_alias=True, exclude_none=True), + ) + except Exception as error: + errors.append(error) + response = JSONRPCError( + jsonrpc="2.0", id=request.id, error=ErrorData(code=500, message="Unexpected pagination request") + ) + incoming.put(SessionMessage(message=JSONRPCMessage(response))) + + thread = Thread(target=serve_pages, daemon=True) + thread.start() + client = client if client is not None else MCPClient("http://127.0.0.1:1/mcp") + try: + with client._exit_stack: + client._session = client._exit_stack.enter_context( + ClientSession(incoming, outgoing, read_timeout_seconds=timedelta(seconds=2)) + ) + yield client, requests + finally: + outgoing.put(None) + thread.join(timeout=2) + assert not thread.is_alive() + assert not errors, errors + + +def test_single_page_control() -> None: + with client_for_pages({None: ListToolsResult(tools=[tool("alpha")])}) as (client, requests): + assert [x.name for x in client.list_tools()] == ["alpha"] + assert len(requests) == 1 + + +@pytest.mark.parametrize("first_page", [[tool("alpha")], []], ids=["populated-first-page", "empty-first-page"]) +def test_all_pages_are_available(first_page: list[Tool]) -> None: + cursor = "opaque:+/=?page2" + with client_for_pages( + { + None: ListToolsResult(tools=first_page, nextCursor=cursor), + cursor: ListToolsResult(tools=[tool("beta")]), + } + ) as (client, requests): + expected = [t.name for t in first_page] + ["beta"] + assert [x.name for x in client.list_tools()] == expected + assert requests[1].params is not None + assert requests[1].params.cursor == cursor + + +def test_no_session_control() -> None: + with pytest.raises(ValueError, match="Session not initialized"): + MCPClient("http://127.0.0.1:1/mcp").list_tools() + + +@pytest.mark.parametrize("cursors", [("opaque-A", "opaque-A"), ("opaque-A", "opaque-B", "opaque-A")]) +def test_repeated_cursor_is_a_bounded_error(cursors: tuple[str, ...]) -> None: + pages = { + None if index == 0 else cursors[index - 1]: ListToolsResult(tools=[tool(f"page-{index}")], nextCursor=cursor) + for index, cursor in enumerate(cursors) + } + with client_for_pages(pages) as (client, requests): + with pytest.raises(MCPConnectionError, match="repeated pagination cursor"): + client.list_tools() + assert len(requests) == len(cursors) + + +def test_later_page_auth_error_restarts_clean_inventory(monkeypatch: pytest.MonkeyPatch) -> None: + client = MCPClientWithAuthRetry("http://127.0.0.1:1/mcp") + old_pages = { + None: ListToolsResult(tools=[tool("stale-alpha")], nextCursor="opaque-old"), + "opaque-old": ErrorData(code=401, message="controlled later-page authentication failure"), + } + new_pages = { + None: ListToolsResult(tools=[tool("fresh-alpha")], nextCursor="opaque-new"), + "opaque-new": ListToolsResult(tools=[tool("beta")]), + } + auth_errors: list[MCPAuthError] = [] + cleanup: list[str] = [] + reconnects: list[str] = [] + with ( + client_for_pages(old_pages, client) as (_, old_requests), + client_for_pages(new_pages) as ( + reconnected_client, + new_requests, + ), + ): + old_session = client._session + assert old_session is not None + client._initialized = True + client._exit_stack.callback(lambda: cleanup.append("closed")) + + def reconnect() -> None: + reconnects.append("reconnected") + client._session = reconnected_client._session + + # Patch only the OAuth exchange and transport reconnect, not session requests. + monkeypatch.setattr(client, "_handle_auth_error", auth_errors.append) + monkeypatch.setattr(client, "_initialize", reconnect) + assert [t.name for t in client.list_tools()] == ["fresh-alpha", "beta"] + assert [p.params.cursor if p.params is not None else None for p in old_requests] == [None, "opaque-old"] + assert [p.params.cursor if p.params is not None else None for p in new_requests] == [None, "opaque-new"] + assert len(auth_errors) == 1 + assert str(auth_errors[0]) == "controlled later-page authentication failure" + assert cleanup == ["closed"] + assert reconnects == ["reconnected"] + assert old_session._receiver_future is not None + assert old_session._receiver_future.done() + assert client._has_retried is False + + +@pytest.mark.parametrize("page_sizes", [(10_000,), (9_999, 1)], ids=["single-page", "two-pages"]) +def test_exact_tool_limit_on_terminal_page(page_sizes: tuple[int, ...]) -> None: + pages: dict[str | None, ListToolsResult] = {} + expected: list[str] = [] + for index, size in enumerate(page_sizes): + names = [f"page-{index}-tool-{number}" for number in range(size)] + expected.extend(names) + cursor = None if index == 0 else f"page-{index}" + next_cursor = f"page-{index + 1}" if index + 1 < len(page_sizes) else None + pages[cursor] = ListToolsResult(tools=[tool(name) for name in names], nextCursor=next_cursor) + with client_for_pages(pages) as (client, requests): + assert [result.name for result in client.list_tools()] == expected + assert len(requests) == len(page_sizes) + + +@pytest.mark.parametrize("next_cursor", [None, "more"]) +@pytest.mark.parametrize("page_sizes", [(10_001,), (9_999, 2)], ids=["oversized-page", "accumulated-overflow"]) +def test_tool_limit_overflow_never_returns_partial_inventory( + page_sizes: tuple[int, ...], next_cursor: str | None +) -> None: + pages: dict[str | None, ListToolsResult] = {} + for index, size in enumerate(page_sizes): + cursor = None if index == 0 else f"page-{index}" + continuation = f"page-{index + 1}" if index + 1 < len(page_sizes) else next_cursor + pages[cursor] = ListToolsResult( + tools=[tool(f"page-{index}-tool-{number}") for number in range(size)], nextCursor=continuation + ) + with client_for_pages(pages) as (client, requests): + with pytest.raises(MCPConnectionError, match="10,000-tool discovery limit"): + client.list_tools() + assert len(requests) == len(page_sizes) + + +@pytest.mark.parametrize("next_cursor", ["more", ""]) +@pytest.mark.parametrize("first_page_size", [10_000, 9_999], ids=["single-page", "two-pages"]) +def test_exact_tool_limit_with_continuation_fails_without_fetching_again( + first_page_size: int, next_cursor: str +) -> None: + pages: dict[str | None, ListToolsResult] = { + None: ListToolsResult( + tools=[tool(f"tool-{number}") for number in range(first_page_size)], + nextCursor=next_cursor if first_page_size == 10_000 else "last-tool", + ) + } + if first_page_size < 10_000: + pages["last-tool"] = ListToolsResult(tools=[tool("tool-9999")], nextCursor=next_cursor) + with client_for_pages(pages) as (client, requests): + with pytest.raises(MCPConnectionError, match="10,000-tool discovery limit"): + client.list_tools() + assert len(requests) == (1 if first_page_size == 10_000 else 2) + + +@pytest.mark.parametrize("auth_retry", [False, True], ids=["base-client", "auth-retry-client"]) +def test_empty_pages_exhaust_request_budget_without_extra_request_or_auth_retry(auth_retry: bool) -> None: + client = MCPClientWithAuthRetry("http://127.0.0.1:1/mcp") if auth_retry else MCPClient("http://127.0.0.1:1/mcp") + pages = { + None if index == 0 else f"fresh-cursor-{index - 1}": ListToolsResult( + tools=[], nextCursor=f"fresh-cursor-{index}" + ) + for index in range(100) + } + with client_for_pages(pages, client) as (_, requests): + with pytest.raises(MCPConnectionError, match="100-request discovery limit") as exc_info: + client.list_tools() + assert not isinstance(exc_info.value, MCPAuthError) + assert len(requests) == 100 + + +@pytest.mark.parametrize( + "last_page_tools", [[], [tool("last-page-tool")]], ids=["empty-terminal", "populated-terminal"] +) +def test_terminal_page_on_last_allowed_request_succeeds(last_page_tools: list[Tool]) -> None: + pages = { + None if index == 0 else f"fresh-cursor-{index - 1}": ListToolsResult( + tools=[], nextCursor=f"fresh-cursor-{index}" + ) + for index in range(99) + } + pages["fresh-cursor-98"] = ListToolsResult(tools=last_page_tools) + with client_for_pages(pages) as (client, requests): + assert client.list_tools() == last_page_tools + assert len(requests) == 100 + + +@pytest.mark.parametrize("error_on_first_page", [True, False], ids=["first-page", "later-page"]) +def test_session_error_propagates_unchanged(error_on_first_page: bool, monkeypatch: pytest.MonkeyPatch) -> None: + error = RuntimeError("controlled session error") + calls: list[str | None] = [] + with client_for_pages({None: ListToolsResult(tools=[tool("alpha")], nextCursor="more")}) as (client, _): + session = client._session + assert session is not None + list_tools = session.list_tools + + def fail_on_page(cursor: str | None = None) -> ListToolsResult: + calls.append(cursor) + if error_on_first_page or cursor is not None: + raise error + return list_tools(cursor=cursor) + + monkeypatch.setattr(session, "list_tools", fail_on_page) + with pytest.raises(RuntimeError) as exc_info: + client.list_tools() + assert exc_info.value is error + assert calls == ([None] if error_on_first_page else [None, "more"]) + + +def test_empty_string_is_an_opaque_continuation_cursor() -> None: + with client_for_pages( + { + None: ListToolsResult(tools=[tool("alpha")], nextCursor=""), + "": ListToolsResult(tools=[tool("beta")]), + } + ) as (client, requests): + assert [result.name for result in client.list_tools()] == ["alpha", "beta"] + assert len(requests) == 2 + assert requests[1].params is not None + assert requests[1].params.cursor == "" diff --git a/api/tests/unit_tests/core/mcp/test_mcp_tool_pagination.py b/api/tests/unit_tests/core/mcp/test_mcp_tool_pagination.py deleted file mode 100644 index ca80d80bed4..00000000000 --- a/api/tests/unit_tests/core/mcp/test_mcp_tool_pagination.py +++ /dev/null @@ -1,280 +0,0 @@ -"""Offline pagination regressions using the real client and in-memory JSON-RPC.""" - -from collections.abc import Generator, Mapping -from contextlib import contextmanager -from datetime import timedelta -from queue import Queue -from threading import Thread - -import pytest - -from core.mcp.auth_client import MCPClientWithAuthRetry -from core.mcp.error import MCPAuthError, MCPConnectionError -from core.mcp.mcp_client import MCPClient -from core.mcp.session.client_session import ClientSession -from core.mcp.types import ( - ErrorData, - JSONRPCError, - JSONRPCMessage, - JSONRPCRequest, - JSONRPCResponse, - ListToolsRequest, - ListToolsResult, - SessionMessage, - Tool, -) - - -def tool(name: str) -> Tool: - return Tool(name=name, description=name, inputSchema={"type": "object", "properties": {}}) - - -@contextmanager -def client_for_pages( - pages: Mapping[str | None, ListToolsResult | ErrorData], client: MCPClient | None = None -) -> Generator[tuple[MCPClient, list[ListToolsRequest]]]: - incoming: Queue[SessionMessage | Exception | None] = Queue() - outgoing: Queue[SessionMessage | None] = Queue() - requests: list[ListToolsRequest] = [] - errors: list[Exception] = [] - - def serve_pages() -> None: - while (message := outgoing.get()) is not None: - request = message.message.root - assert isinstance(request, JSONRPCRequest) - try: - payload = ListToolsRequest.model_validate({"method": request.method, "params": request.params}) - requests.append(payload) - cursor = payload.params.cursor if payload.params is not None else None - page = pages[cursor] - response: JSONRPCResponse | JSONRPCError - if isinstance(page, ErrorData): - response = JSONRPCError(jsonrpc="2.0", id=request.id, error=page) - else: - response = JSONRPCResponse( - jsonrpc="2.0", - id=request.id, - result=page.model_dump(mode="json", by_alias=True, exclude_none=True), - ) - except Exception as error: - errors.append(error) - response = JSONRPCError( - jsonrpc="2.0", id=request.id, error=ErrorData(code=500, message="Unexpected pagination request") - ) - incoming.put(SessionMessage(message=JSONRPCMessage(response))) - - thread = Thread(target=serve_pages, daemon=True) - thread.start() - client = client if client is not None else MCPClient("http://127.0.0.1:1/mcp") - try: - with client._exit_stack: - client._session = client._exit_stack.enter_context( - ClientSession(incoming, outgoing, read_timeout_seconds=timedelta(seconds=2)) - ) - yield client, requests - finally: - outgoing.put(None) - thread.join(timeout=2) - assert not thread.is_alive() - assert not errors, errors - - -def test_single_page_control() -> None: - with client_for_pages({None: ListToolsResult(tools=[tool("alpha")])}) as (client, requests): - assert [x.name for x in client.list_tools()] == ["alpha"] - assert len(requests) == 1 - - -@pytest.mark.parametrize("first_page", [[tool("alpha")], []], ids=["populated-first-page", "empty-first-page"]) -def test_all_pages_are_available(first_page: list[Tool]) -> None: - cursor = "opaque:+/=?page2" - with client_for_pages( - { - None: ListToolsResult(tools=first_page, nextCursor=cursor), - cursor: ListToolsResult(tools=[tool("beta")]), - } - ) as (client, requests): - expected = [t.name for t in first_page] + ["beta"] - assert [x.name for x in client.list_tools()] == expected - assert requests[1].params is not None - assert requests[1].params.cursor == cursor - - -def test_no_session_control() -> None: - with pytest.raises(ValueError, match="Session not initialized"): - MCPClient("http://127.0.0.1:1/mcp").list_tools() - - -@pytest.mark.parametrize("cursors", [("opaque-A", "opaque-A"), ("opaque-A", "opaque-B", "opaque-A")]) -def test_repeated_cursor_is_a_bounded_error(cursors: tuple[str, ...]) -> None: - pages = { - None if index == 0 else cursors[index - 1]: ListToolsResult(tools=[tool(f"page-{index}")], nextCursor=cursor) - for index, cursor in enumerate(cursors) - } - with client_for_pages(pages) as (client, requests): - with pytest.raises(MCPConnectionError, match="repeated pagination cursor"): - client.list_tools() - assert len(requests) == len(cursors) - - -def test_later_page_auth_error_restarts_clean_inventory(monkeypatch: pytest.MonkeyPatch) -> None: - client = MCPClientWithAuthRetry("http://127.0.0.1:1/mcp") - old_pages = { - None: ListToolsResult(tools=[tool("stale-alpha")], nextCursor="opaque-old"), - "opaque-old": ErrorData(code=401, message="controlled later-page authentication failure"), - } - new_pages = { - None: ListToolsResult(tools=[tool("fresh-alpha")], nextCursor="opaque-new"), - "opaque-new": ListToolsResult(tools=[tool("beta")]), - } - auth_errors: list[MCPAuthError] = [] - cleanup: list[str] = [] - reconnects: list[str] = [] - with ( - client_for_pages(old_pages, client) as (_, old_requests), - client_for_pages(new_pages) as ( - reconnected_client, - new_requests, - ), - ): - old_session = client._session - assert old_session is not None - client._initialized = True - client._exit_stack.callback(lambda: cleanup.append("closed")) - - def reconnect() -> None: - reconnects.append("reconnected") - client._session = reconnected_client._session - - # Patch only the OAuth exchange and transport reconnect, not session requests. - monkeypatch.setattr(client, "_handle_auth_error", auth_errors.append) - monkeypatch.setattr(client, "_initialize", reconnect) - assert [t.name for t in client.list_tools()] == ["fresh-alpha", "beta"] - assert [p.params.cursor if p.params is not None else None for p in old_requests] == [None, "opaque-old"] - assert [p.params.cursor if p.params is not None else None for p in new_requests] == [None, "opaque-new"] - assert len(auth_errors) == 1 - assert str(auth_errors[0]) == "controlled later-page authentication failure" - assert cleanup == ["closed"] - assert reconnects == ["reconnected"] - assert old_session._receiver_future is not None - assert old_session._receiver_future.done() - assert client._has_retried is False - - -@pytest.mark.parametrize("page_sizes", [(10_000,), (9_999, 1)], ids=["single-page", "two-pages"]) -def test_exact_tool_limit_on_terminal_page(page_sizes: tuple[int, ...]) -> None: - pages: dict[str | None, ListToolsResult] = {} - expected: list[str] = [] - for index, size in enumerate(page_sizes): - names = [f"page-{index}-tool-{number}" for number in range(size)] - expected.extend(names) - cursor = None if index == 0 else f"page-{index}" - next_cursor = f"page-{index + 1}" if index + 1 < len(page_sizes) else None - pages[cursor] = ListToolsResult(tools=[tool(name) for name in names], nextCursor=next_cursor) - with client_for_pages(pages) as (client, requests): - assert [result.name for result in client.list_tools()] == expected - assert len(requests) == len(page_sizes) - - -@pytest.mark.parametrize("next_cursor", [None, "more"]) -@pytest.mark.parametrize("page_sizes", [(10_001,), (9_999, 2)], ids=["oversized-page", "accumulated-overflow"]) -def test_tool_limit_overflow_never_returns_partial_inventory( - page_sizes: tuple[int, ...], next_cursor: str | None -) -> None: - pages: dict[str | None, ListToolsResult] = {} - for index, size in enumerate(page_sizes): - cursor = None if index == 0 else f"page-{index}" - continuation = f"page-{index + 1}" if index + 1 < len(page_sizes) else next_cursor - pages[cursor] = ListToolsResult( - tools=[tool(f"page-{index}-tool-{number}") for number in range(size)], nextCursor=continuation - ) - with client_for_pages(pages) as (client, requests): - with pytest.raises(MCPConnectionError, match="10,000-tool discovery limit"): - client.list_tools() - assert len(requests) == len(page_sizes) - - -@pytest.mark.parametrize("next_cursor", ["more", ""]) -@pytest.mark.parametrize("first_page_size", [10_000, 9_999], ids=["single-page", "two-pages"]) -def test_exact_tool_limit_with_continuation_fails_without_fetching_again( - first_page_size: int, next_cursor: str -) -> None: - pages: dict[str | None, ListToolsResult] = { - None: ListToolsResult( - tools=[tool(f"tool-{number}") for number in range(first_page_size)], - nextCursor=next_cursor if first_page_size == 10_000 else "last-tool", - ) - } - if first_page_size < 10_000: - pages["last-tool"] = ListToolsResult(tools=[tool("tool-9999")], nextCursor=next_cursor) - with client_for_pages(pages) as (client, requests): - with pytest.raises(MCPConnectionError, match="10,000-tool discovery limit"): - client.list_tools() - assert len(requests) == (1 if first_page_size == 10_000 else 2) - - -@pytest.mark.parametrize("auth_retry", [False, True], ids=["base-client", "auth-retry-client"]) -def test_empty_pages_exhaust_request_budget_without_extra_request_or_auth_retry(auth_retry: bool) -> None: - client = MCPClientWithAuthRetry("http://127.0.0.1:1/mcp") if auth_retry else MCPClient("http://127.0.0.1:1/mcp") - pages = { - None if index == 0 else f"fresh-cursor-{index - 1}": ListToolsResult( - tools=[], nextCursor=f"fresh-cursor-{index}" - ) - for index in range(100) - } - with client_for_pages(pages, client) as (_, requests): - with pytest.raises(MCPConnectionError, match="100-request discovery limit") as exc_info: - client.list_tools() - assert not isinstance(exc_info.value, MCPAuthError) - assert len(requests) == 100 - - -@pytest.mark.parametrize( - "last_page_tools", [[], [tool("last-page-tool")]], ids=["empty-terminal", "populated-terminal"] -) -def test_terminal_page_on_last_allowed_request_succeeds(last_page_tools: list[Tool]) -> None: - pages = { - None if index == 0 else f"fresh-cursor-{index - 1}": ListToolsResult( - tools=[], nextCursor=f"fresh-cursor-{index}" - ) - for index in range(99) - } - pages["fresh-cursor-98"] = ListToolsResult(tools=last_page_tools) - with client_for_pages(pages) as (client, requests): - assert client.list_tools() == last_page_tools - assert len(requests) == 100 - - -@pytest.mark.parametrize("error_on_first_page", [True, False], ids=["first-page", "later-page"]) -def test_session_error_propagates_unchanged(error_on_first_page: bool, monkeypatch: pytest.MonkeyPatch) -> None: - error = RuntimeError("controlled session error") - calls: list[str | None] = [] - with client_for_pages({None: ListToolsResult(tools=[tool("alpha")], nextCursor="more")}) as (client, _): - session = client._session - assert session is not None - list_tools = session.list_tools - - def fail_on_page(cursor: str | None = None) -> ListToolsResult: - calls.append(cursor) - if error_on_first_page or cursor is not None: - raise error - return list_tools(cursor=cursor) - - monkeypatch.setattr(session, "list_tools", fail_on_page) - with pytest.raises(RuntimeError) as exc_info: - client.list_tools() - assert exc_info.value is error - assert calls == ([None] if error_on_first_page else [None, "more"]) - - -def test_empty_string_is_an_opaque_continuation_cursor() -> None: - with client_for_pages( - { - None: ListToolsResult(tools=[tool("alpha")], nextCursor=""), - "": ListToolsResult(tools=[tool("beta")]), - } - ) as (client, requests): - assert [result.name for result in client.list_tools()] == ["alpha", "beta"] - assert len(requests) == 2 - assert requests[1].params is not None - assert requests[1].params.cursor == "" diff --git a/api/tests/unit_tests/core/model_runtime/test_model_provider_factory.py b/api/tests/unit_tests/core/model_runtime/test_model_provider_factory.py deleted file mode 100644 index 2b51dc8182c..00000000000 --- a/api/tests/unit_tests/core/model_runtime/test_model_provider_factory.py +++ /dev/null @@ -1,425 +0,0 @@ -from unittest.mock import Mock - -import pytest - -from core.plugin.impl.model_runtime_factory import create_model_type_instance -from graphon.model_runtime.entities.common_entities import I18nObject -from graphon.model_runtime.entities.model_entities import AIModelEntity, FetchFrom, ModelType -from graphon.model_runtime.entities.provider_entities import ( - ConfigurateMethod, - CredentialFormSchema, - FieldModelSchema, - FormType, - ModelCredentialSchema, - ProviderCredentialSchema, - ProviderEntity, -) -from graphon.model_runtime.model_providers.base.large_language_model import LargeLanguageModel -from graphon.model_runtime.model_providers.base.moderation_model import ModerationModel -from graphon.model_runtime.model_providers.base.rerank_model import RerankModel -from graphon.model_runtime.model_providers.base.speech2text_model import Speech2TextModel -from graphon.model_runtime.model_providers.base.text_embedding_model import TextEmbeddingModel -from graphon.model_runtime.model_providers.base.tts_model import TTSModel -from graphon.model_runtime.model_providers.model_provider_factory import ModelProviderFactory - - -def _build_model(model: str, model_type: ModelType) -> AIModelEntity: - return AIModelEntity( - model=model, - label=I18nObject(en_US=model), - model_type=model_type, - fetch_from=FetchFrom.PREDEFINED_MODEL, - model_properties={}, - ) - - -def _build_provider( - *, - provider: str, - provider_name: str, - supported_model_types: list[ModelType], - models: list[AIModelEntity] | None = None, - provider_credential_schema: ProviderCredentialSchema | None = None, - model_credential_schema: ModelCredentialSchema | None = None, -) -> ProviderEntity: - return ProviderEntity( - provider=provider, - provider_name=provider_name, - label=I18nObject(en_US=provider_name or provider), - supported_model_types=supported_model_types, - configurate_methods=[ConfigurateMethod.PREDEFINED_MODEL], - models=models or [], - provider_credential_schema=provider_credential_schema, - model_credential_schema=model_credential_schema, - ) - - -class _FakeModelRuntime: - def __init__(self, providers: list[ProviderEntity]) -> None: - self._providers = providers - self.validate_provider_credentials = Mock() - self.validate_model_credentials = Mock() - self.get_model_schema = Mock() - self.get_provider_icon = Mock() - - def fetch_model_providers(self) -> list[ProviderEntity]: - return self._providers - - -def test_model_provider_factory_resolves_runtime_provider_name() -> None: - provider = ProviderEntity( - provider="langgenius/openai/openai", - provider_name="openai", - label=I18nObject(en_US="OpenAI"), - supported_model_types=[ModelType.LLM], - configurate_methods=[ConfigurateMethod.PREDEFINED_MODEL], - ) - factory = ModelProviderFactory(runtime=_FakeModelRuntime([provider])) - - provider_schema = factory.get_model_provider("openai") - - assert provider_schema.provider == "langgenius/openai/openai" - assert provider_schema.provider_name == "openai" - - -def test_model_provider_factory_resolves_canonical_short_name_independent_of_provider_order() -> None: - providers = [ - ProviderEntity( - provider="acme/openai/openai", - provider_name="", - label=I18nObject(en_US="Acme OpenAI"), - supported_model_types=[ModelType.LLM], - configurate_methods=[ConfigurateMethod.PREDEFINED_MODEL], - ), - ProviderEntity( - provider="langgenius/openai/openai", - provider_name="openai", - label=I18nObject(en_US="OpenAI"), - supported_model_types=[ModelType.LLM], - configurate_methods=[ConfigurateMethod.PREDEFINED_MODEL], - ), - ] - factory = ModelProviderFactory(runtime=_FakeModelRuntime(providers)) - - provider_schema = factory.get_model_provider("openai") - - assert provider_schema.provider == "langgenius/openai/openai" - assert provider_schema.provider_name == "openai" - - -def test_model_provider_factory_requires_runtime() -> None: - with pytest.raises(ValueError, match="runtime is required"): - ModelProviderFactory(runtime=None) # type: ignore[arg-type] - - -def test_model_provider_factory_get_providers_returns_runtime_providers() -> None: - providers = [ - _build_provider( - provider="langgenius/openai/openai", - provider_name="openai", - supported_model_types=[ModelType.LLM], - ) - ] - factory = ModelProviderFactory(runtime=_FakeModelRuntime(providers)) - - result = factory.get_providers() - - assert list(result) == providers - assert result is not providers - - -def test_model_provider_factory_get_provider_schema_delegates_to_provider_lookup() -> None: - provider = _build_provider( - provider="langgenius/openai/openai", - provider_name="openai", - supported_model_types=[ModelType.LLM], - ) - factory = ModelProviderFactory(runtime=_FakeModelRuntime([provider])) - - result = factory.get_provider_schema("openai") - - assert result is provider - - -def test_model_provider_factory_raises_for_unknown_provider() -> None: - factory = ModelProviderFactory( - runtime=_FakeModelRuntime( - [ - _build_provider( - provider="langgenius/openai/openai", - provider_name="openai", - supported_model_types=[ModelType.LLM], - ) - ] - ) - ) - - with pytest.raises(ValueError, match="Invalid provider: anthropic"): - factory.get_model_provider("anthropic") - - -def test_model_provider_factory_get_models_filters_provider_and_model_type() -> None: - providers = [ - _build_provider( - provider="langgenius/openai/openai", - provider_name="openai", - supported_model_types=[ModelType.LLM, ModelType.TTS], - models=[_build_model("gpt-4o-mini", ModelType.LLM), _build_model("tts-1", ModelType.TTS)], - ), - _build_provider( - provider="langgenius/cohere/cohere", - provider_name="cohere", - supported_model_types=[ModelType.RERANK], - models=[_build_model("rerank-v3", ModelType.RERANK)], - ), - ] - factory = ModelProviderFactory(runtime=_FakeModelRuntime(providers)) - - results = factory.get_models(provider="openai", model_type=ModelType.LLM) - - assert len(results) == 1 - assert results[0].provider == "langgenius/openai/openai" - assert [model.model for model in results[0].models] == ["gpt-4o-mini"] - - -def test_model_provider_factory_get_models_skips_providers_without_requested_model_type() -> None: - providers = [ - _build_provider( - provider="langgenius/openai/openai", - provider_name="openai", - supported_model_types=[ModelType.LLM], - models=[_build_model("gpt-4o-mini", ModelType.LLM)], - ), - _build_provider( - provider="langgenius/elevenlabs/elevenlabs", - provider_name="elevenlabs", - supported_model_types=[ModelType.TTS], - models=[_build_model("eleven_multilingual_v2", ModelType.TTS)], - ), - ] - factory = ModelProviderFactory(runtime=_FakeModelRuntime(providers)) - - results = factory.get_models(model_type=ModelType.TTS) - - assert len(results) == 1 - assert results[0].provider == "langgenius/elevenlabs/elevenlabs" - assert [model.model for model in results[0].models] == ["eleven_multilingual_v2"] - - -def test_model_provider_factory_get_models_without_model_type_keeps_all_provider_models() -> None: - providers = [ - _build_provider( - provider="langgenius/openai/openai", - provider_name="openai", - supported_model_types=[ModelType.LLM, ModelType.TTS], - models=[_build_model("gpt-4o-mini", ModelType.LLM), _build_model("tts-1", ModelType.TTS)], - ) - ] - factory = ModelProviderFactory(runtime=_FakeModelRuntime(providers)) - - results = factory.get_models(provider="openai") - - assert len(results) == 1 - assert [model.model for model in results[0].models] == ["gpt-4o-mini", "tts-1"] - - -def test_model_provider_factory_validates_provider_credentials() -> None: - runtime = _FakeModelRuntime( - [ - _build_provider( - provider="langgenius/openai/openai", - provider_name="openai", - supported_model_types=[ModelType.LLM], - provider_credential_schema=ProviderCredentialSchema( - credential_form_schemas=[ - CredentialFormSchema( - variable="api_key", - label=I18nObject(en_US="API key"), - type=FormType.SECRET_INPUT, - required=True, - ) - ] - ), - ) - ] - ) - factory = ModelProviderFactory(runtime=runtime) - - filtered = factory.provider_credentials_validate( - provider="openai", - credentials={"api_key": "secret", "ignored": "value"}, - ) - - assert filtered == {"api_key": "secret"} - runtime.validate_provider_credentials.assert_called_once_with( - provider="langgenius/openai/openai", - credentials={"api_key": "secret"}, - ) - - -def test_model_provider_factory_provider_credentials_validate_requires_schema() -> None: - factory = ModelProviderFactory( - runtime=_FakeModelRuntime( - [ - _build_provider( - provider="langgenius/openai/openai", - provider_name="openai", - supported_model_types=[ModelType.LLM], - ) - ] - ) - ) - - with pytest.raises(ValueError, match="Provider openai does not have provider_credential_schema"): - factory.provider_credentials_validate(provider="openai", credentials={"api_key": "secret"}) - - -def test_model_provider_factory_validates_model_credentials() -> None: - runtime = _FakeModelRuntime( - [ - _build_provider( - provider="langgenius/openai/openai", - provider_name="openai", - supported_model_types=[ModelType.LLM], - model_credential_schema=ModelCredentialSchema( - model=FieldModelSchema(label=I18nObject(en_US="Model")), - credential_form_schemas=[ - CredentialFormSchema( - variable="api_key", - label=I18nObject(en_US="API key"), - type=FormType.SECRET_INPUT, - required=True, - ) - ], - ), - ) - ] - ) - factory = ModelProviderFactory(runtime=runtime) - - filtered = factory.model_credentials_validate( - provider="openai", - model_type=ModelType.LLM, - model="gpt-4o-mini", - credentials={"api_key": "secret", "ignored": "value"}, - ) - - assert filtered == {"api_key": "secret"} - runtime.validate_model_credentials.assert_called_once_with( - provider="langgenius/openai/openai", - model_type=ModelType.LLM, - model="gpt-4o-mini", - credentials={"api_key": "secret"}, - ) - - -def test_model_provider_factory_model_credentials_validate_requires_schema() -> None: - factory = ModelProviderFactory( - runtime=_FakeModelRuntime( - [ - _build_provider( - provider="langgenius/openai/openai", - provider_name="openai", - supported_model_types=[ModelType.LLM], - ) - ] - ) - ) - - with pytest.raises(ValueError, match="Provider openai does not have model_credential_schema"): - factory.model_credentials_validate( - provider="openai", - model_type=ModelType.LLM, - model="gpt-4o-mini", - credentials={"api_key": "secret"}, - ) - - -def test_model_provider_factory_get_model_schema_and_icon_use_canonical_provider() -> None: - runtime = _FakeModelRuntime( - [ - _build_provider( - provider="langgenius/openai/openai", - provider_name="openai", - supported_model_types=[ModelType.LLM], - ) - ] - ) - runtime.get_model_schema.return_value = "schema" - runtime.get_provider_icon.return_value = (b"icon", "image/png") - factory = ModelProviderFactory(runtime=runtime) - - assert ( - factory.get_model_schema( - provider="openai", - model_type=ModelType.LLM, - model="gpt-4o-mini", - credentials=None, - ) - == "schema" - ) - assert factory.get_provider_icon("openai", "icon_small", "en_US") == (b"icon", "image/png") - runtime.get_model_schema.assert_called_once_with( - provider="langgenius/openai/openai", - model_type=ModelType.LLM, - model="gpt-4o-mini", - credentials={}, - ) - runtime.get_provider_icon.assert_called_once_with( - provider="langgenius/openai/openai", - icon_type="icon_small", - lang="en_US", - ) - - -@pytest.mark.parametrize( - ("model_type", "expected_type"), - [ - (ModelType.LLM, LargeLanguageModel), - (ModelType.TEXT_EMBEDDING, TextEmbeddingModel), - (ModelType.RERANK, RerankModel), - (ModelType.SPEECH2TEXT, Speech2TextModel), - (ModelType.MODERATION, ModerationModel), - (ModelType.TTS, TTSModel), - ], -) -def test_create_model_type_instance_builds_model_wrappers( - model_type: ModelType, - expected_type: type[object], -) -> None: - runtime = _FakeModelRuntime( - [ - _build_provider( - provider="langgenius/openai/openai", - provider_name="openai", - supported_model_types=[model_type], - ) - ] - ) - - instance = create_model_type_instance( - runtime=runtime, - provider_schema=runtime.fetch_model_providers()[0], - model_type=model_type, - ) - - assert isinstance(instance, expected_type) - - -def test_create_model_type_instance_rejects_unsupported_model_type() -> None: - runtime = _FakeModelRuntime( - [ - _build_provider( - provider="langgenius/openai/openai", - provider_name="openai", - supported_model_types=[ModelType.LLM], - ) - ] - ) - - with pytest.raises(ValueError, match="Unsupported model type: unsupported"): - create_model_type_instance( - runtime=runtime, - provider_schema=runtime.fetch_model_providers()[0], - model_type="unsupported", # type: ignore[arg-type] - ) diff --git a/api/tests/unit_tests/core/moderation/test_sensitive_word_filter.py b/api/tests/unit_tests/core/moderation/keywords/test_keywords.py similarity index 99% rename from api/tests/unit_tests/core/moderation/test_sensitive_word_filter.py rename to api/tests/unit_tests/core/moderation/keywords/test_keywords.py index e8232645fc4..4f316354839 100644 --- a/api/tests/unit_tests/core/moderation/test_sensitive_word_filter.py +++ b/api/tests/unit_tests/core/moderation/keywords/test_keywords.py @@ -1,3 +1,21 @@ +import time + +import pytest + +from core.moderation.base import ModerationAction, ModerationInputsResult, ModerationOutputsResult +from core.moderation.keywords.keywords import KeywordsModeration + + +def test_validate_config_ignores_trailing_blank_keyword_row() -> None: + config = { + "inputs_config": {"enabled": True, "preset_response": "Blocked"}, + "outputs_config": {"enabled": False}, + "keywords": "\n".join(f"word{i}" for i in range(100)) + "\n", + } + + KeywordsModeration.validate_config("tenant", config) + + """ Unit tests for sensitive word filter (KeywordsModeration). @@ -9,13 +27,6 @@ This module tests the sensitive word filtering functionality including: - Input and output moderation scenarios """ -import time - -import pytest - -from core.moderation.base import ModerationAction, ModerationInputsResult, ModerationOutputsResult -from core.moderation.keywords.keywords import KeywordsModeration - class TestConfigValidation: """Test configuration validation for KeywordsModeration.""" diff --git a/api/tests/unit_tests/core/moderation/test_content_moderation.py b/api/tests/unit_tests/core/moderation/openai_moderation/test_openai_moderation.py similarity index 100% rename from api/tests/unit_tests/core/moderation/test_content_moderation.py rename to api/tests/unit_tests/core/moderation/openai_moderation/test_openai_moderation.py diff --git a/api/tests/unit_tests/core/moderation/test_keyword_row_validation.py b/api/tests/unit_tests/core/moderation/test_keyword_row_validation.py deleted file mode 100644 index 0ba4d99cdf6..00000000000 --- a/api/tests/unit_tests/core/moderation/test_keyword_row_validation.py +++ /dev/null @@ -1,11 +0,0 @@ -from core.moderation.keywords.keywords import KeywordsModeration - - -def test_validate_config_ignores_trailing_blank_keyword_row() -> None: - config = { - "inputs_config": {"enabled": True, "preset_response": "Blocked"}, - "outputs_config": {"enabled": False}, - "keywords": "\n".join(f"word{i}" for i in range(100)) + "\n", - } - - KeywordsModeration.validate_config("tenant", config) diff --git a/api/tests/unit_tests/core/ops/test_config_entity.py b/api/tests/unit_tests/core/ops/entities/test_config_entity.py similarity index 100% rename from api/tests/unit_tests/core/ops/test_config_entity.py rename to api/tests/unit_tests/core/ops/entities/test_config_entity.py diff --git a/api/tests/unit_tests/core/ops/test_trace_session_metadata.py b/api/tests/unit_tests/core/ops/entities/test_trace_entity.py similarity index 100% rename from api/tests/unit_tests/core/ops/test_trace_session_metadata.py rename to api/tests/unit_tests/core/ops/entities/test_trace_entity.py diff --git a/api/tests/unit_tests/core/ops/test_application_tracing_hooks.py b/api/tests/unit_tests/core/ops/test_application_tracing_hooks.py deleted file mode 100644 index 994d0cad57b..00000000000 --- a/api/tests/unit_tests/core/ops/test_application_tracing_hooks.py +++ /dev/null @@ -1,52 +0,0 @@ -"""The execution port is gated before Flask lookup and bound per application.""" - -from contextvars import Context - -import pytest -from flask import Flask, has_app_context - -from core.ops.base_trace_instance import ( - ApplicationTracingHooksNotRegisteredError, - application_tracing_hooks, - bind_application_tracing_hooks, -) -from services.tracing.runtime_bridge import RuntimeApplicationTracingHooks -from tests.unit_tests.config_override import apply_config_overrides - - -def test_disabled_port_needs_no_application_context(monkeypatch: pytest.MonkeyPatch) -> None: - apply_config_overrides(monkeypatch, OPS_AGENT_TRACE_ENABLED=False) - isolated_context = Context() - assert not isolated_context.run(has_app_context) - assert isolated_context.run(application_tracing_hooks) is None - - -def test_enabled_port_reports_missing_registration_without_context(monkeypatch: pytest.MonkeyPatch) -> None: - apply_config_overrides(monkeypatch, OPS_AGENT_TRACE_ENABLED=True) - isolated_context = Context() - assert not isolated_context.run(has_app_context) - with pytest.raises(ApplicationTracingHooksNotRegisteredError, match="application_tracing_hooks_not_registered"): - isolated_context.run(application_tracing_hooks) - - -def test_enabled_port_reports_missing_registration_in_app(monkeypatch: pytest.MonkeyPatch) -> None: - apply_config_overrides(monkeypatch, OPS_AGENT_TRACE_ENABLED=True) - with Flask(__name__).app_context(): - with pytest.raises(ApplicationTracingHooksNotRegisteredError, match="application_tracing_hooks_not_registered"): - application_tracing_hooks() - - -def test_port_does_not_leak_between_flask_applications(monkeypatch: pytest.MonkeyPatch) -> None: - apply_config_overrides(monkeypatch, OPS_AGENT_TRACE_ENABLED=True) - first, second = Flask("first"), Flask("second") - first_hooks = RuntimeApplicationTracingHooks() - second_hooks = RuntimeApplicationTracingHooks() - bind_application_tracing_hooks(first, first_hooks) - bind_application_tracing_hooks(second, second_hooks) - with first.app_context(): - assert application_tracing_hooks() is first_hooks - with second.app_context(): - assert application_tracing_hooks() is second_hooks - assert application_tracing_hooks() is first_hooks - apply_config_overrides(monkeypatch, OPS_AGENT_TRACE_ENABLED=False) - assert application_tracing_hooks() is None diff --git a/api/tests/unit_tests/core/ops/test_base_trace_instance.py b/api/tests/unit_tests/core/ops/test_base_trace_instance.py index c09883aded0..fdf8a4ff54f 100644 --- a/api/tests/unit_tests/core/ops/test_base_trace_instance.py +++ b/api/tests/unit_tests/core/ops/test_base_trace_instance.py @@ -1,15 +1,69 @@ +from contextvars import Context from types import SimpleNamespace import pytest +from flask import Flask, has_app_context from sqlalchemy import Engine, select from sqlalchemy.orm import Session -from core.ops.base_trace_instance import BaseTraceInstance +from core.ops.base_trace_instance import ( + ApplicationTracingHooksNotRegisteredError, + BaseTraceInstance, + application_tracing_hooks, + bind_application_tracing_hooks, +) from core.ops.entities.config_entity import BaseTracingConfig from core.ops.entities.trace_entity import BaseTraceInfo from models import Account, App, Tenant, TenantAccountJoin from models.account import TenantAccountRole from models.model import AppMode +from services.tracing.runtime_bridge import RuntimeApplicationTracingHooks +from tests.helpers.unit.config_override import apply_config_overrides + +"""The execution port is gated before Flask lookup and bound per application.""" + + +class TestApplicationTracingHooksCases: + @staticmethod + def test_disabled_port_needs_no_application_context(monkeypatch: pytest.MonkeyPatch) -> None: + apply_config_overrides(monkeypatch, OPS_AGENT_TRACE_ENABLED=False) + isolated_context = Context() + assert not isolated_context.run(has_app_context) + assert isolated_context.run(application_tracing_hooks) is None + + @staticmethod + def test_enabled_port_reports_missing_registration_without_context(monkeypatch: pytest.MonkeyPatch) -> None: + apply_config_overrides(monkeypatch, OPS_AGENT_TRACE_ENABLED=True) + isolated_context = Context() + assert not isolated_context.run(has_app_context) + with pytest.raises(ApplicationTracingHooksNotRegisteredError, match="application_tracing_hooks_not_registered"): + isolated_context.run(application_tracing_hooks) + + @staticmethod + def test_enabled_port_reports_missing_registration_in_app(monkeypatch: pytest.MonkeyPatch) -> None: + apply_config_overrides(monkeypatch, OPS_AGENT_TRACE_ENABLED=True) + with Flask(__name__).app_context(): + with pytest.raises( + ApplicationTracingHooksNotRegisteredError, match="application_tracing_hooks_not_registered" + ): + application_tracing_hooks() + + @staticmethod + def test_port_does_not_leak_between_flask_applications(monkeypatch: pytest.MonkeyPatch) -> None: + apply_config_overrides(monkeypatch, OPS_AGENT_TRACE_ENABLED=True) + first, second = Flask("first"), Flask("second") + first_hooks = RuntimeApplicationTracingHooks() + second_hooks = RuntimeApplicationTracingHooks() + bind_application_tracing_hooks(first, first_hooks) + bind_application_tracing_hooks(second, second_hooks) + with first.app_context(): + assert application_tracing_hooks() is first_hooks + with second.app_context(): + assert application_tracing_hooks() is second_hooks + assert application_tracing_hooks() is first_hooks + apply_config_overrides(monkeypatch, OPS_AGENT_TRACE_ENABLED=False) + assert application_tracing_hooks() is None + TABLES = (App, Account, Tenant, TenantAccountJoin) @@ -22,13 +76,6 @@ class ConcreteTraceInstance(BaseTraceInstance): super().trace(trace_info) -@pytest.fixture(autouse=True) -def _bind_production_sessions(monkeypatch: pytest.MonkeyPatch, sqlite_engine: Engine) -> None: - """Bind the service-owned ORM session to the test's SQLite engine.""" - database = SimpleNamespace(engine=sqlite_engine) - monkeypatch.setattr("core.ops.base_trace_instance.db", database) - - def _persist_app(session: Session, *, created_by: str | None) -> App: app = App( id="some_app_id", @@ -63,146 +110,155 @@ def _trace_instance() -> ConcreteTraceInstance: return ConcreteTraceInstance(BaseTracingConfig()) -@pytest.mark.parametrize("sqlite_session", [TABLES], indirect=True) -def test_get_service_account_with_tenant_app_not_found(sqlite_session: Session): - with pytest.raises(ValueError, match="App with id some_app_id not found"): - _trace_instance().get_service_account_with_tenant("some_app_id") +class TestBaseTraceInstanceCases: + @staticmethod + @pytest.fixture(autouse=True) + def _bind_production_sessions(monkeypatch: pytest.MonkeyPatch, sqlite_engine: Engine) -> None: + """Bind the service-owned ORM session to the test's SQLite engine.""" + database = SimpleNamespace(engine=sqlite_engine) + monkeypatch.setattr("core.ops.base_trace_instance.db", database) + @staticmethod + @pytest.mark.parametrize("sqlite_session", [TABLES], indirect=True) + def test_get_service_account_with_tenant_app_not_found(sqlite_session: Session): + with pytest.raises(ValueError, match="App with id some_app_id not found"): + _trace_instance().get_service_account_with_tenant("some_app_id") -@pytest.mark.parametrize("sqlite_session", [TABLES], indirect=True) -def test_get_service_account_with_tenant_no_creator(sqlite_session: Session): - _persist_app(sqlite_session, created_by=None) + @staticmethod + @pytest.mark.parametrize("sqlite_session", [TABLES], indirect=True) + def test_get_service_account_with_tenant_no_creator(sqlite_session: Session): + _persist_app(sqlite_session, created_by=None) - with pytest.raises(ValueError, match="App with id some_app_id has no creator"): - _trace_instance().get_service_account_with_tenant("some_app_id") + with pytest.raises(ValueError, match="App with id some_app_id has no creator"): + _trace_instance().get_service_account_with_tenant("some_app_id") + @staticmethod + @pytest.mark.parametrize("sqlite_session", [TABLES], indirect=True) + def test_get_service_account_with_tenant_creator_not_found(sqlite_session: Session): + _persist_app(sqlite_session, created_by="creator_id") -@pytest.mark.parametrize("sqlite_session", [TABLES], indirect=True) -def test_get_service_account_with_tenant_creator_not_found(sqlite_session: Session): - _persist_app(sqlite_session, created_by="creator_id") + with pytest.raises(ValueError, match="Creator account with id creator_id not found for app some_app_id"): + _trace_instance().get_service_account_with_tenant("some_app_id") - with pytest.raises(ValueError, match="Creator account with id creator_id not found for app some_app_id"): - _trace_instance().get_service_account_with_tenant("some_app_id") + @staticmethod + @pytest.mark.parametrize("sqlite_session", [TABLES], indirect=True) + def test_get_service_account_with_tenant_tenant_not_found(sqlite_session: Session): + _persist_app(sqlite_session, created_by="creator_id") + _persist_account(sqlite_session) + with pytest.raises(ValueError, match="App tenant tenant_id not found for account creator_id"): + _trace_instance().get_service_account_with_tenant("some_app_id") -@pytest.mark.parametrize("sqlite_session", [TABLES], indirect=True) -def test_get_service_account_with_tenant_tenant_not_found(sqlite_session: Session): - _persist_app(sqlite_session, created_by="creator_id") - _persist_account(sqlite_session) + @staticmethod + @pytest.mark.parametrize("current", [True, False]) + @pytest.mark.parametrize("sqlite_session", [TABLES], indirect=True) + def test_get_service_account_with_tenant_success(sqlite_session: Session, current: bool): + _persist_app(sqlite_session, created_by="creator_id") + _persist_account(sqlite_session) + tenant = Tenant(name="Workspace") + tenant.id = "tenant_id" + sqlite_session.add_all( + [ + tenant, + TenantAccountJoin( + tenant_id=tenant.id, + account_id="creator_id", + current=current, + role=TenantAccountRole.OWNER, + ), + ] + ) + sqlite_session.commit() - with pytest.raises(ValueError, match="App tenant tenant_id not found for account creator_id"): - _trace_instance().get_service_account_with_tenant("some_app_id") + result = _trace_instance().get_service_account_with_tenant("some_app_id") + assert result.id == "creator_id" + assert result.current_tenant_id == "tenant_id" + assert result.current_role == TenantAccountRole.OWNER -@pytest.mark.parametrize("current", [True, False]) -@pytest.mark.parametrize("sqlite_session", [TABLES], indirect=True) -def test_get_service_account_with_tenant_success(sqlite_session: Session, current: bool): - _persist_app(sqlite_session, created_by="creator_id") - _persist_account(sqlite_session) - tenant = Tenant(name="Workspace") - tenant.id = "tenant_id" - sqlite_session.add_all( - [ - tenant, - TenantAccountJoin( - tenant_id=tenant.id, - account_id="creator_id", - current=current, - role=TenantAccountRole.OWNER, - ), - ] - ) - sqlite_session.commit() + @staticmethod + @pytest.mark.parametrize("sqlite_session", [TABLES], indirect=True) + def test_get_service_account_with_tenant_uses_app_workspace(sqlite_session: Session): + _persist_app(sqlite_session, created_by="creator_id") + _persist_account(sqlite_session) + app_tenant = Tenant(name="App workspace") + app_tenant.id = "tenant_id" + current_tenant = Tenant(name="Current workspace") + current_tenant.id = "other_tenant_id" + sqlite_session.add_all( + [ + app_tenant, + current_tenant, + TenantAccountJoin( + tenant_id="tenant_id", + account_id="creator_id", + current=False, + role=TenantAccountRole.EDITOR, + ), + TenantAccountJoin( + tenant_id="other_tenant_id", + account_id="creator_id", + current=True, + role=TenantAccountRole.OWNER, + ), + ] + ) + sqlite_session.commit() - result = _trace_instance().get_service_account_with_tenant("some_app_id") + result = _trace_instance().get_service_account_with_tenant("some_app_id") - assert result.id == "creator_id" - assert result.current_tenant_id == "tenant_id" - assert result.current_role == TenantAccountRole.OWNER + assert result.id == "creator_id" + assert result.current_tenant_id == "tenant_id" + assert result.current_role == TenantAccountRole.EDITOR + sqlite_session.expire_all() + memberships = sqlite_session.scalars( + select(TenantAccountJoin).where(TenantAccountJoin.account_id == "creator_id") + ).all() + assert {membership.tenant_id: membership.current for membership in memberships} == { + "tenant_id": False, + "other_tenant_id": True, + } + @staticmethod + @pytest.mark.parametrize("sqlite_session", [TABLES], indirect=True) + def test_get_service_account_with_tenant_requires_app_membership(sqlite_session: Session): + _persist_app(sqlite_session, created_by="creator_id") + _persist_account(sqlite_session) + app_tenant = Tenant(name="App workspace") + app_tenant.id = "tenant_id" + current_tenant = Tenant(name="Current workspace") + current_tenant.id = "other_tenant_id" + sqlite_session.add_all( + [ + app_tenant, + current_tenant, + TenantAccountJoin( + tenant_id="other_tenant_id", + account_id="creator_id", + current=True, + role=TenantAccountRole.OWNER, + ), + ] + ) + sqlite_session.commit() -@pytest.mark.parametrize("sqlite_session", [TABLES], indirect=True) -def test_get_service_account_with_tenant_uses_app_workspace(sqlite_session: Session): - _persist_app(sqlite_session, created_by="creator_id") - _persist_account(sqlite_session) - app_tenant = Tenant(name="App workspace") - app_tenant.id = "tenant_id" - current_tenant = Tenant(name="Current workspace") - current_tenant.id = "other_tenant_id" - sqlite_session.add_all( - [ - app_tenant, - current_tenant, + with pytest.raises(ValueError, match="App tenant tenant_id not found for account creator_id"): + _trace_instance().get_service_account_with_tenant("some_app_id") + + @staticmethod + @pytest.mark.parametrize("sqlite_session", [TABLES], indirect=True) + def test_get_service_account_with_tenant_requires_existing_app_tenant(sqlite_session: Session): + _persist_app(sqlite_session, created_by="creator_id") + _persist_account(sqlite_session) + sqlite_session.add( TenantAccountJoin( tenant_id="tenant_id", account_id="creator_id", - current=False, - role=TenantAccountRole.EDITOR, - ), - TenantAccountJoin( - tenant_id="other_tenant_id", - account_id="creator_id", current=True, role=TenantAccountRole.OWNER, - ), - ] - ) - sqlite_session.commit() - - result = _trace_instance().get_service_account_with_tenant("some_app_id") - - assert result.id == "creator_id" - assert result.current_tenant_id == "tenant_id" - assert result.current_role == TenantAccountRole.EDITOR - sqlite_session.expire_all() - memberships = sqlite_session.scalars( - select(TenantAccountJoin).where(TenantAccountJoin.account_id == "creator_id") - ).all() - assert {membership.tenant_id: membership.current for membership in memberships} == { - "tenant_id": False, - "other_tenant_id": True, - } - - -@pytest.mark.parametrize("sqlite_session", [TABLES], indirect=True) -def test_get_service_account_with_tenant_requires_app_membership(sqlite_session: Session): - _persist_app(sqlite_session, created_by="creator_id") - _persist_account(sqlite_session) - app_tenant = Tenant(name="App workspace") - app_tenant.id = "tenant_id" - current_tenant = Tenant(name="Current workspace") - current_tenant.id = "other_tenant_id" - sqlite_session.add_all( - [ - app_tenant, - current_tenant, - TenantAccountJoin( - tenant_id="other_tenant_id", - account_id="creator_id", - current=True, - role=TenantAccountRole.OWNER, - ), - ] - ) - sqlite_session.commit() - - with pytest.raises(ValueError, match="App tenant tenant_id not found for account creator_id"): - _trace_instance().get_service_account_with_tenant("some_app_id") - - -@pytest.mark.parametrize("sqlite_session", [TABLES], indirect=True) -def test_get_service_account_with_tenant_requires_existing_app_tenant(sqlite_session: Session): - _persist_app(sqlite_session, created_by="creator_id") - _persist_account(sqlite_session) - sqlite_session.add( - TenantAccountJoin( - tenant_id="tenant_id", - account_id="creator_id", - current=True, - role=TenantAccountRole.OWNER, + ) ) - ) - sqlite_session.commit() + sqlite_session.commit() - with pytest.raises(ValueError, match="App tenant tenant_id not found for account creator_id"): - _trace_instance().get_service_account_with_tenant("some_app_id") + with pytest.raises(ValueError, match="App tenant tenant_id not found for account creator_id"): + _trace_instance().get_service_account_with_tenant("some_app_id") diff --git a/api/tests/unit_tests/core/ops/test_lookup_helpers.py b/api/tests/unit_tests/core/ops/test_lookup_helpers.py deleted file mode 100644 index 90c2127b6ad..00000000000 --- a/api/tests/unit_tests/core/ops/test_lookup_helpers.py +++ /dev/null @@ -1,595 +0,0 @@ -"""Unit tests for lookup helper functions in core.ops.ops_trace_manager. - -Covers: -- _lookup_app_and_workspace_names -- _lookup_credential_name -- _lookup_llm_credential_info -- TraceTask._get_user_id_from_metadata -""" - -import uuid -from collections.abc import Iterator -from contextlib import contextmanager -from unittest.mock import PropertyMock, patch - -import pytest -from sqlalchemy import Engine, event -from sqlalchemy.orm import Session - -from core.tools.entities.tool_entities import ApiProviderSchemaType -from extensions.ext_database import db -from graphon.model_runtime.entities.model_entities import ModelType -from models.account import Tenant -from models.base import TypeBase -from models.model import App, AppMode, IconType -from models.provider import Provider, ProviderCredential, ProviderModel, ProviderModelCredential, ProviderType -from models.tools import ApiToolProvider, BuiltinToolProvider, MCPToolProvider, WorkflowToolProvider - -# --------------------------------------------------------------------------- -# Helpers -# --------------------------------------------------------------------------- - - -@pytest.fixture -def orm_session(sqlite_engine: Engine) -> Iterator[Session]: - models = ( - App, - Tenant, - Provider, - ProviderCredential, - ProviderModel, - ProviderModelCredential, - BuiltinToolProvider, - ApiToolProvider, - WorkflowToolProvider, - MCPToolProvider, - ) - tables = [model.metadata.tables[model.__tablename__] for model in models] - TypeBase.metadata.create_all(sqlite_engine, tables=tables) - - with patch.object(type(db), "engine", new_callable=PropertyMock, return_value=sqlite_engine): - with Session(sqlite_engine, expire_on_commit=False) as session: - yield session - - -def _persist_app(session: Session, *, tenant_id: str, name: str = "MyApp") -> App: - app = App( - id=str(uuid.uuid4()), - tenant_id=tenant_id, - name=name, - mode=AppMode.WORKFLOW, - icon_type=IconType.EMOJI, - icon="workflow", - icon_background="#FFFFFF", - enable_site=True, - enable_api=False, - ) - session.add(app) - session.commit() - return app - - -def _persist_tenant(session: Session, *, name: str = "MyWorkspace") -> Tenant: - tenant = Tenant(name=name) - session.add(tenant) - session.commit() - return tenant - - -def _persist_tool_provider( - session: Session, provider_type: str -) -> BuiltinToolProvider | ApiToolProvider | WorkflowToolProvider | MCPToolProvider: - tenant_id = str(uuid.uuid4()) - user_id = str(uuid.uuid4()) - if provider_type in {"builtin", "plugin"}: - provider = BuiltinToolProvider( - name="CredentialA", - tenant_id=tenant_id, - user_id=user_id, - provider="test/provider", - ) - elif provider_type == "api": - provider = ApiToolProvider( - name="CredentialA", - icon="icon.svg", - schema="{}", - schema_type_str=ApiProviderSchemaType.OPENAPI, - user_id=user_id, - tenant_id=tenant_id, - description="API provider", - tools_str="[]", - credentials_str="{}", - ) - elif provider_type == "workflow": - provider = WorkflowToolProvider( - name="CredentialA", - label="CredentialA", - icon="icon.svg", - app_id=str(uuid.uuid4()), - version="1", - user_id=user_id, - tenant_id=tenant_id, - description="Workflow provider", - ) - elif provider_type == "mcp": - provider = MCPToolProvider( - name="CredentialA", - server_identifier="credential-a", - server_url="https://example.com/mcp", - server_url_hash="credential-a-hash", - icon="icon.svg", - tenant_id=tenant_id, - user_id=user_id, - ) - else: - raise ValueError(f"unsupported provider type: {provider_type}") - - session.add(provider) - session.commit() - return provider - - -def _persist_provider_credential( - session: Session, - *, - tenant_id: str, - credential_name: str = "ProvCredName", -) -> ProviderCredential: - credential = ProviderCredential( - tenant_id=tenant_id, - provider_name="openai", - credential_name=credential_name, - encrypted_config="{}", - ) - session.add(credential) - session.commit() - return credential - - -def _persist_model_credential( - session: Session, - *, - tenant_id: str, - credential_name: str = "ModelCredName", -) -> ProviderModelCredential: - credential = ProviderModelCredential( - tenant_id=tenant_id, - provider_name="openai", - model_name="gpt-4", - model_type=ModelType.LLM, - credential_name=credential_name, - encrypted_config="{}", - ) - session.add(credential) - session.commit() - return credential - - -def _persist_provider( - session: Session, - *, - tenant_id: str, - credential_id: str | None, -) -> Provider: - provider = Provider( - tenant_id=tenant_id, - provider_name="openai", - provider_type=ProviderType.CUSTOM, - credential_id=credential_id, - ) - session.add(provider) - session.commit() - return provider - - -def _persist_provider_model( - session: Session, - *, - tenant_id: str, - credential_id: str | None, -) -> ProviderModel: - model = ProviderModel( - tenant_id=tenant_id, - provider_name="openai", - model_name="gpt-4", - model_type=ModelType.LLM, - credential_id=credential_id, - ) - session.add(model) - session.commit() - return model - - -@contextmanager -def _raise_on_table(engine: Engine, table_name: str) -> Iterator[None]: - """Raise only when SQL targets the named table, leaving other real lookups intact.""" - - def fail_target_query(_conn, _cursor, statement, _parameters, _context, _executemany): - if f"FROM {table_name}" in statement: - raise RuntimeError(f"forced failure for {table_name}") - - event.listen(engine, "before_cursor_execute", fail_target_query) - try: - yield - finally: - event.remove(engine, "before_cursor_execute", fail_target_query) - - -# --------------------------------------------------------------------------- -# _lookup_app_and_workspace_names -# --------------------------------------------------------------------------- - - -class TestLookupAppAndWorkspaceNames: - """Tests for _lookup_app_and_workspace_names(app_id, tenant_id).""" - - def test_both_found(self, orm_session: Session): - """Returns (app_name, workspace_name) when both records exist.""" - from core.ops.ops_trace_manager import _lookup_app_and_workspace_names - - tenant = _persist_tenant(orm_session) - app = _persist_app(orm_session, tenant_id=tenant.id) - app_name, workspace_name = _lookup_app_and_workspace_names(app.id, tenant.id) - - assert app_name == "MyApp" - assert workspace_name == "MyWorkspace" - - def test_app_only_found(self, orm_session: Session): - """Returns (app_name, '') when tenant record is absent.""" - from core.ops.ops_trace_manager import _lookup_app_and_workspace_names - - app = _persist_app(orm_session, tenant_id=str(uuid.uuid4())) - app_name, workspace_name = _lookup_app_and_workspace_names(app.id, str(uuid.uuid4())) - - assert app_name == "MyApp" - assert workspace_name == "" - - def test_tenant_only_found(self, orm_session: Session): - """Returns ('', workspace_name) when app record is absent.""" - from core.ops.ops_trace_manager import _lookup_app_and_workspace_names - - tenant = _persist_tenant(orm_session) - app_name, workspace_name = _lookup_app_and_workspace_names(str(uuid.uuid4()), tenant.id) - - assert app_name == "" - assert workspace_name == "MyWorkspace" - - def test_neither_found(self, orm_session: Session): - """Returns ('', '') when both DB lookups return None.""" - from core.ops.ops_trace_manager import _lookup_app_and_workspace_names - - app_name, workspace_name = _lookup_app_and_workspace_names(str(uuid.uuid4()), str(uuid.uuid4())) - - assert app_name == "" - assert workspace_name == "" - - def test_none_inputs_skips_db(self): - """Returns ('', '') immediately when both IDs are None — no DB access.""" - from core.ops.ops_trace_manager import _lookup_app_and_workspace_names - - app_name, workspace_name = _lookup_app_and_workspace_names(None, None) - - assert app_name == "" - assert workspace_name == "" - - def test_app_id_none_only_queries_tenant(self, orm_session: Session): - """When app_id is None, only the tenant query is issued.""" - from core.ops.ops_trace_manager import _lookup_app_and_workspace_names - - tenant = _persist_tenant(orm_session, name="OnlyWorkspace") - app_name, workspace_name = _lookup_app_and_workspace_names(None, tenant.id) - - assert app_name == "" - assert workspace_name == "OnlyWorkspace" - - def test_tenant_id_none_only_queries_app(self, orm_session: Session): - """When tenant_id is None, only the app query is issued.""" - from core.ops.ops_trace_manager import _lookup_app_and_workspace_names - - app = _persist_app(orm_session, tenant_id=str(uuid.uuid4()), name="OnlyApp") - app_name, workspace_name = _lookup_app_and_workspace_names(app.id, None) - - assert app_name == "OnlyApp" - assert workspace_name == "" - - -# --------------------------------------------------------------------------- -# _lookup_credential_name -# --------------------------------------------------------------------------- - - -class TestLookupCredentialName: - """Tests for _lookup_credential_name(credential_id, provider_type).""" - - @pytest.mark.parametrize("provider_type", ["builtin", "plugin", "api", "workflow", "mcp"]) - def test_known_provider_types_return_name(self, provider_type: str, orm_session: Session): - """Each valid provider_type results in a DB query and returns the credential name.""" - from core.ops.ops_trace_manager import _lookup_credential_name - - provider = _persist_tool_provider(orm_session, provider_type) - result = _lookup_credential_name(provider.id, provider_type) - - assert result == "CredentialA" - - def test_credential_not_found_returns_empty_string(self, orm_session: Session): - """Returns '' when DB yields None for the given credential_id.""" - from core.ops.ops_trace_manager import _lookup_credential_name - - result = _lookup_credential_name(str(uuid.uuid4()), "api") - - assert result == "" - - def test_invalid_provider_type_returns_empty_string_without_db(self): - """Returns '' immediately for an unrecognised provider_type — no DB access.""" - from core.ops.ops_trace_manager import _lookup_credential_name - - result = _lookup_credential_name(str(uuid.uuid4()), "unknown_type") - - assert result == "" - - def test_none_credential_id_returns_empty_string_without_db(self): - """Returns '' immediately when credential_id is None — no DB access.""" - from core.ops.ops_trace_manager import _lookup_credential_name - - result = _lookup_credential_name(None, "api") - - assert result == "" - - def test_none_provider_type_returns_empty_string_without_db(self): - """Returns '' immediately when provider_type is None — no DB access.""" - from core.ops.ops_trace_manager import _lookup_credential_name - - result = _lookup_credential_name(str(uuid.uuid4()), None) - - assert result == "" - - def test_builtin_and_plugin_map_to_same_model(self): - """Both 'builtin' and 'plugin' provider_types query BuiltinToolProvider.""" - from core.ops.ops_trace_manager import _PROVIDER_TYPE_TO_MODEL - from models.tools import BuiltinToolProvider - - assert _PROVIDER_TYPE_TO_MODEL["builtin"] is BuiltinToolProvider - assert _PROVIDER_TYPE_TO_MODEL["plugin"] is BuiltinToolProvider - - def test_api_maps_to_api_tool_provider(self): - """'api' maps to ApiToolProvider.""" - from core.ops.ops_trace_manager import _PROVIDER_TYPE_TO_MODEL - from models.tools import ApiToolProvider - - assert _PROVIDER_TYPE_TO_MODEL["api"] is ApiToolProvider - - def test_workflow_maps_to_workflow_tool_provider(self): - """'workflow' maps to WorkflowToolProvider.""" - from core.ops.ops_trace_manager import _PROVIDER_TYPE_TO_MODEL - from models.tools import WorkflowToolProvider - - assert _PROVIDER_TYPE_TO_MODEL["workflow"] is WorkflowToolProvider - - def test_mcp_maps_to_mcp_tool_provider(self): - """'mcp' maps to MCPToolProvider.""" - from core.ops.ops_trace_manager import _PROVIDER_TYPE_TO_MODEL - from models.tools import MCPToolProvider - - assert _PROVIDER_TYPE_TO_MODEL["mcp"] is MCPToolProvider - - -# --------------------------------------------------------------------------- -# _lookup_llm_credential_info -# --------------------------------------------------------------------------- - - -class TestLookupLlmCredentialInfo: - """Tests for _lookup_llm_credential_info(tenant_id, provider, model, model_type).""" - - def test_model_level_credential_found(self, orm_session: Session): - """Returns model-level credential_id and name when ProviderModel has a credential.""" - from core.ops.ops_trace_manager import _lookup_llm_credential_info - - tenant_id = str(uuid.uuid4()) - model_credential = _persist_model_credential(orm_session, tenant_id=tenant_id) - _persist_provider(orm_session, tenant_id=tenant_id, credential_id=None) - _persist_provider_model(orm_session, tenant_id=tenant_id, credential_id=model_credential.id) - - decoy_tenant_id = str(uuid.uuid4()) - decoy_credential = _persist_model_credential( - orm_session, - tenant_id=decoy_tenant_id, - credential_name="WrongTenantCredential", - ) - _persist_provider(orm_session, tenant_id=decoy_tenant_id, credential_id=None) - _persist_provider_model(orm_session, tenant_id=decoy_tenant_id, credential_id=decoy_credential.id) - - cred_id, cred_name = _lookup_llm_credential_info(tenant_id, "openai", "gpt-4") - - assert cred_id == model_credential.id - assert cred_name == "ModelCredName" - - def test_provider_level_fallback_when_no_model_credential(self, orm_session: Session): - """Falls back to provider-level credential when ProviderModel has no credential_id.""" - from core.ops.ops_trace_manager import _lookup_llm_credential_info - - tenant_id = str(uuid.uuid4()) - provider_credential = _persist_provider_credential(orm_session, tenant_id=tenant_id) - _persist_provider(orm_session, tenant_id=tenant_id, credential_id=provider_credential.id) - _persist_provider_model(orm_session, tenant_id=tenant_id, credential_id=None) - - cred_id, cred_name = _lookup_llm_credential_info(tenant_id, "openai", "gpt-4") - - assert cred_id == provider_credential.id - assert cred_name == "ProvCredName" - - def test_provider_level_fallback_when_no_model_record(self, orm_session: Session): - """Falls back to provider-level credential when no ProviderModel row exists.""" - from core.ops.ops_trace_manager import _lookup_llm_credential_info - - tenant_id = str(uuid.uuid4()) - provider_credential = _persist_provider_credential(orm_session, tenant_id=tenant_id) - _persist_provider(orm_session, tenant_id=tenant_id, credential_id=provider_credential.id) - - cred_id, cred_name = _lookup_llm_credential_info(tenant_id, "openai", "gpt-4") - - assert cred_id == provider_credential.id - assert cred_name == "ProvCredName" - - def test_no_model_arg_uses_provider_level_only(self, orm_session: Session): - """When model is None, skips ProviderModel query and uses provider credential.""" - from core.ops.ops_trace_manager import _lookup_llm_credential_info - - tenant_id = str(uuid.uuid4()) - provider_credential = _persist_provider_credential(orm_session, tenant_id=tenant_id) - _persist_provider(orm_session, tenant_id=tenant_id, credential_id=provider_credential.id) - - cred_id, cred_name = _lookup_llm_credential_info(tenant_id, "openai", None) - - assert cred_id == provider_credential.id - assert cred_name == "ProvCredName" - - def test_provider_not_found_returns_none_and_empty(self, orm_session: Session): - """Returns (None, '') when Provider record does not exist.""" - from core.ops.ops_trace_manager import _lookup_llm_credential_info - - other_tenant_id = str(uuid.uuid4()) - _persist_provider(orm_session, tenant_id=other_tenant_id, credential_id=None) - tenant_id = str(uuid.uuid4()) - - cred_id, cred_name = _lookup_llm_credential_info(tenant_id, "openai", "gpt-4") - - assert cred_id is None - assert cred_name == "" - - def test_none_tenant_id_returns_none_and_empty_without_db(self): - """Returns (None, '') immediately when tenant_id is None — no DB access.""" - from core.ops.ops_trace_manager import _lookup_llm_credential_info - - cred_id, cred_name = _lookup_llm_credential_info(None, "openai", "gpt-4") - - assert cred_id is None - assert cred_name == "" - - def test_none_provider_returns_none_and_empty_without_db(self): - """Returns (None, '') immediately when provider is None — no DB access.""" - from core.ops.ops_trace_manager import _lookup_llm_credential_info - - cred_id, cred_name = _lookup_llm_credential_info(str(uuid.uuid4()), None, "gpt-4") - - assert cred_id is None - assert cred_name == "" - - def test_db_error_on_outer_query_returns_none_and_empty(self, orm_session: Session, sqlite_engine: Engine): - """Returns (None, '') and logs a warning when the outer DB query raises.""" - from core.ops.ops_trace_manager import _lookup_llm_credential_info - - with _raise_on_table(sqlite_engine, "providers"): - cred_id, cred_name = _lookup_llm_credential_info(str(uuid.uuid4()), "openai", "gpt-4") - - assert cred_id is None - assert cred_name == "" - - def test_credential_name_lookup_failure_returns_id_with_empty_name( - self, orm_session: Session, sqlite_engine: Engine - ): - """When credential name sub-query fails, returns cred_id but '' for name.""" - from core.ops.ops_trace_manager import _lookup_llm_credential_info - - tenant_id = str(uuid.uuid4()) - provider_credential = _persist_provider_credential(orm_session, tenant_id=tenant_id) - _persist_provider(orm_session, tenant_id=tenant_id, credential_id=provider_credential.id) - - with _raise_on_table(sqlite_engine, "provider_credentials"): - cred_id, cred_name = _lookup_llm_credential_info(tenant_id, "openai", "gpt-4") - - assert cred_id == provider_credential.id - assert cred_name == "" - - def test_no_credential_on_provider_or_model_returns_none_id(self, orm_session: Session): - """Returns (None, '') when neither provider nor model has a credential_id.""" - from core.ops.ops_trace_manager import _lookup_llm_credential_info - - tenant_id = str(uuid.uuid4()) - _persist_provider(orm_session, tenant_id=tenant_id, credential_id=None) - _persist_provider_model(orm_session, tenant_id=tenant_id, credential_id=None) - - cred_id, cred_name = _lookup_llm_credential_info(tenant_id, "openai", "gpt-4") - - assert cred_id is None - assert cred_name == "" - - -# --------------------------------------------------------------------------- -# TraceTask._get_user_id_from_metadata -# --------------------------------------------------------------------------- - - -class TestGetUserIdFromMetadata: - """Tests for TraceTask._get_user_id_from_metadata(metadata). - - Pure dict logic — no DB access required. - """ - - @pytest.fixture - def get_user_id(self): - """Return the classmethod under test.""" - from core.ops.ops_trace_manager import TraceTask - - return TraceTask._get_user_id_from_metadata - - def test_from_end_user_id_has_highest_priority(self, get_user_id): - """from_end_user_id takes precedence over all other keys.""" - metadata = { - "from_end_user_id": "eu-abc", - "from_account_id": "acc-xyz", - "user_id": "u-123", - } - assert get_user_id(metadata) == "end_user:eu-abc" - - def test_from_account_id_used_when_no_end_user(self, get_user_id): - """from_account_id is used when from_end_user_id is absent.""" - metadata = { - "from_account_id": "acc-xyz", - "user_id": "u-123", - } - assert get_user_id(metadata) == "account:acc-xyz" - - def test_user_id_used_when_no_end_user_or_account(self, get_user_id): - """user_id is used when both higher-priority keys are absent.""" - metadata = {"user_id": "u-123"} - assert get_user_id(metadata) == "user:u-123" - - def test_returns_anonymous_when_all_keys_absent(self, get_user_id): - """Returns 'anonymous' when metadata has none of the expected keys.""" - assert get_user_id({}) == "anonymous" - - def test_empty_string_end_user_id_is_skipped(self, get_user_id): - """Empty string for from_end_user_id is falsy and falls through to next key.""" - metadata = { - "from_end_user_id": "", - "from_account_id": "acc-xyz", - } - assert get_user_id(metadata) == "account:acc-xyz" - - def test_empty_string_account_id_is_skipped(self, get_user_id): - """Empty string for from_account_id is falsy and falls through to user_id.""" - metadata = { - "from_end_user_id": "", - "from_account_id": "", - "user_id": "u-123", - } - assert get_user_id(metadata) == "user:u-123" - - def test_empty_string_user_id_falls_through_to_anonymous(self, get_user_id): - """Empty string for user_id is falsy, so 'anonymous' is returned.""" - metadata = { - "from_end_user_id": "", - "from_account_id": "", - "user_id": "", - } - assert get_user_id(metadata) == "anonymous" - - def test_only_from_end_user_id_present(self, get_user_id): - """Minimal case: only from_end_user_id present.""" - assert get_user_id({"from_end_user_id": "eu-only"}) == "end_user:eu-only" - - def test_irrelevant_keys_do_not_interfere(self, get_user_id): - """Extra metadata keys have no effect on the result.""" - metadata = {"invoke_from": "web", "app_id": "a1"} - assert get_user_id(metadata) == "anonymous" diff --git a/api/tests/unit_tests/core/ops/test_ops_trace_manager.py b/api/tests/unit_tests/core/ops/test_ops_trace_manager.py index a7cc2ba847a..db5a0e38ee6 100644 --- a/api/tests/unit_tests/core/ops/test_ops_trace_manager.py +++ b/api/tests/unit_tests/core/ops/test_ops_trace_manager.py @@ -1,45 +1,636 @@ -"""SQLite-backed tests for :mod:`core.ops.ops_trace_manager`.""" - from __future__ import annotations import builtins import json import queue import sys +import uuid from collections.abc import Iterator +from contextlib import contextmanager from datetime import datetime, timedelta from decimal import Decimal from types import SimpleNamespace -from unittest.mock import PropertyMock, patch -from uuid import UUID +from unittest.mock import MagicMock, PropertyMock, patch +from uuid import UUID, uuid4 import pytest +from Crypto.PublicKey import RSA from flask import Flask -from sqlalchemy import Engine +from pydantic import ValidationError +from sqlalchemy import Engine, event from sqlalchemy.orm import Session, scoped_session, sessionmaker import core.ops.ops_trace_manager as module +import models.engine as model_engine from core.moderation.base import ModerationAction, ModerationInputsResult +from core.ops.base_trace_instance import bind_application_tracing_hooks from core.ops.exceptions import TraceProviderNotInstalledError from core.ops.ops_trace_manager import OpsTraceManager, TraceQueueManager, TraceTask, TraceTaskName from core.rag.models.document import Document as RetrievalDocument +from core.tools.entities.tool_entities import ApiProviderSchemaType +from extensions.ext_database import db +from extensions.ext_key_provider import key_provider_manager from graphon.enums import WorkflowExecutionStatus from graphon.file import FileTransferMethod, FileType +from graphon.model_runtime.entities.model_entities import ModelType +from libs import gmpy2_pkcs10aep_cipher, rsa +from libs.key_providers.rsa_key_provider import RSAKeyProvider +from models.account import Tenant +from models.base import TypeBase from models.enums import ConversationFromSource, CreatorUserRole, MessageStatus, WorkflowRunTriggeredFrom from models.model import ( App, AppMode, AppModelConfig, Conversation, + IconType, Message, MessageAgentThought, MessageFile, TraceAppConfig, ) +from models.provider import Provider, ProviderCredential, ProviderModel, ProviderModelCredential, ProviderType +from models.tools import ApiToolProvider, BuiltinToolProvider, MCPToolProvider, WorkflowToolProvider from models.workflow import WorkflowAppLog, WorkflowAppLogCreatedFrom, WorkflowRun, WorkflowType from repositories.sqlalchemy_api_workflow_run_repository import DifyAPISQLAlchemyWorkflowRunRepository -from tests.unit_tests.config_override import apply_config_overrides -from tests.unit_tests.model_factories import make_app +from services.tracing.runtime_bridge import RuntimeApplicationTracingHooks +from tests.helpers.unit.config_override import apply_config_overrides +from tests.helpers.unit.model_factories import make_app, make_tenant + +"""Unit tests for lookup helper functions in core.ops.ops_trace_manager. + +Covers: +- _lookup_app_and_workspace_names +- _lookup_credential_name +- _lookup_llm_credential_info +- TraceTask._get_user_id_from_metadata +""" + + +def _persist_app(session: Session, *, tenant_id: str, name: str = "MyApp") -> App: + app = App( + id=str(uuid.uuid4()), + tenant_id=tenant_id, + name=name, + mode=AppMode.WORKFLOW, + icon_type=IconType.EMOJI, + icon="workflow", + icon_background="#FFFFFF", + enable_site=True, + enable_api=False, + ) + session.add(app) + session.commit() + return app + + +def _persist_tenant(session: Session, *, name: str = "MyWorkspace") -> Tenant: + tenant = Tenant(name=name) + session.add(tenant) + session.commit() + return tenant + + +def _persist_tool_provider( + session: Session, provider_type: str +) -> BuiltinToolProvider | ApiToolProvider | WorkflowToolProvider | MCPToolProvider: + tenant_id = str(uuid.uuid4()) + user_id = str(uuid.uuid4()) + if provider_type in {"builtin", "plugin"}: + provider = BuiltinToolProvider( + name="CredentialA", + tenant_id=tenant_id, + user_id=user_id, + provider="test/provider", + ) + elif provider_type == "api": + provider = ApiToolProvider( + name="CredentialA", + icon="icon.svg", + schema="{}", + schema_type_str=ApiProviderSchemaType.OPENAPI, + user_id=user_id, + tenant_id=tenant_id, + description="API provider", + tools_str="[]", + credentials_str="{}", + ) + elif provider_type == "workflow": + provider = WorkflowToolProvider( + name="CredentialA", + label="CredentialA", + icon="icon.svg", + app_id=str(uuid.uuid4()), + version="1", + user_id=user_id, + tenant_id=tenant_id, + description="Workflow provider", + ) + elif provider_type == "mcp": + provider = MCPToolProvider( + name="CredentialA", + server_identifier="credential-a", + server_url="https://example.com/mcp", + server_url_hash="credential-a-hash", + icon="icon.svg", + tenant_id=tenant_id, + user_id=user_id, + ) + else: + raise ValueError(f"unsupported provider type: {provider_type}") + + session.add(provider) + session.commit() + return provider + + +def _persist_provider_credential( + session: Session, + *, + tenant_id: str, + credential_name: str = "ProvCredName", +) -> ProviderCredential: + credential = ProviderCredential( + tenant_id=tenant_id, + provider_name="openai", + credential_name=credential_name, + encrypted_config="{}", + ) + session.add(credential) + session.commit() + return credential + + +def _persist_model_credential( + session: Session, + *, + tenant_id: str, + credential_name: str = "ModelCredName", +) -> ProviderModelCredential: + credential = ProviderModelCredential( + tenant_id=tenant_id, + provider_name="openai", + model_name="gpt-4", + model_type=ModelType.LLM, + credential_name=credential_name, + encrypted_config="{}", + ) + session.add(credential) + session.commit() + return credential + + +def _persist_provider( + session: Session, + *, + tenant_id: str, + credential_id: str | None, +) -> Provider: + provider = Provider( + tenant_id=tenant_id, + provider_name="openai", + provider_type=ProviderType.CUSTOM, + credential_id=credential_id, + ) + session.add(provider) + session.commit() + return provider + + +def _persist_provider_model( + session: Session, + *, + tenant_id: str, + credential_id: str | None, +) -> ProviderModel: + model = ProviderModel( + tenant_id=tenant_id, + provider_name="openai", + model_name="gpt-4", + model_type=ModelType.LLM, + credential_id=credential_id, + ) + session.add(model) + session.commit() + return model + + +@contextmanager +def _raise_on_table(engine: Engine, table_name: str) -> Iterator[None]: + """Raise only when SQL targets the named table, leaving other real lookups intact.""" + + def fail_target_query(_conn, _cursor, statement, _parameters, _context, _executemany): + if f"FROM {table_name}" in statement: + raise RuntimeError(f"forced failure for {table_name}") + + event.listen(engine, "before_cursor_execute", fail_target_query) + try: + yield + finally: + event.remove(engine, "before_cursor_execute", fail_target_query) + + +class TestLookupHelpersCases: + # --------------------------------------------------------------------------- + # Helpers + # --------------------------------------------------------------------------- + + @staticmethod + @pytest.fixture + def orm_session(sqlite_engine: Engine) -> Iterator[Session]: + models = ( + App, + Tenant, + Provider, + ProviderCredential, + ProviderModel, + ProviderModelCredential, + BuiltinToolProvider, + ApiToolProvider, + WorkflowToolProvider, + MCPToolProvider, + ) + tables = [model.metadata.tables[model.__tablename__] for model in models] + TypeBase.metadata.create_all(sqlite_engine, tables=tables) + + with patch.object(type(db), "engine", new_callable=PropertyMock, return_value=sqlite_engine): + with Session(sqlite_engine, expire_on_commit=False) as session: + yield session + + # --------------------------------------------------------------------------- + # _lookup_app_and_workspace_names + # --------------------------------------------------------------------------- + + class TestLookupAppAndWorkspaceNames: + """Tests for _lookup_app_and_workspace_names(app_id, tenant_id).""" + + def test_both_found(self, orm_session: Session): + """Returns (app_name, workspace_name) when both records exist.""" + from core.ops.ops_trace_manager import _lookup_app_and_workspace_names + + tenant = _persist_tenant(orm_session) + app = _persist_app(orm_session, tenant_id=tenant.id) + app_name, workspace_name = _lookup_app_and_workspace_names(app.id, tenant.id) + + assert app_name == "MyApp" + assert workspace_name == "MyWorkspace" + + def test_app_only_found(self, orm_session: Session): + """Returns (app_name, '') when tenant record is absent.""" + from core.ops.ops_trace_manager import _lookup_app_and_workspace_names + + app = _persist_app(orm_session, tenant_id=str(uuid.uuid4())) + app_name, workspace_name = _lookup_app_and_workspace_names(app.id, str(uuid.uuid4())) + + assert app_name == "MyApp" + assert workspace_name == "" + + def test_tenant_only_found(self, orm_session: Session): + """Returns ('', workspace_name) when app record is absent.""" + from core.ops.ops_trace_manager import _lookup_app_and_workspace_names + + tenant = _persist_tenant(orm_session) + app_name, workspace_name = _lookup_app_and_workspace_names(str(uuid.uuid4()), tenant.id) + + assert app_name == "" + assert workspace_name == "MyWorkspace" + + def test_neither_found(self, orm_session: Session): + """Returns ('', '') when both DB lookups return None.""" + from core.ops.ops_trace_manager import _lookup_app_and_workspace_names + + app_name, workspace_name = _lookup_app_and_workspace_names(str(uuid.uuid4()), str(uuid.uuid4())) + + assert app_name == "" + assert workspace_name == "" + + def test_none_inputs_skips_db(self): + """Returns ('', '') immediately when both IDs are None — no DB access.""" + from core.ops.ops_trace_manager import _lookup_app_and_workspace_names + + app_name, workspace_name = _lookup_app_and_workspace_names(None, None) + + assert app_name == "" + assert workspace_name == "" + + def test_app_id_none_only_queries_tenant(self, orm_session: Session): + """When app_id is None, only the tenant query is issued.""" + from core.ops.ops_trace_manager import _lookup_app_and_workspace_names + + tenant = _persist_tenant(orm_session, name="OnlyWorkspace") + app_name, workspace_name = _lookup_app_and_workspace_names(None, tenant.id) + + assert app_name == "" + assert workspace_name == "OnlyWorkspace" + + def test_tenant_id_none_only_queries_app(self, orm_session: Session): + """When tenant_id is None, only the app query is issued.""" + from core.ops.ops_trace_manager import _lookup_app_and_workspace_names + + app = _persist_app(orm_session, tenant_id=str(uuid.uuid4()), name="OnlyApp") + app_name, workspace_name = _lookup_app_and_workspace_names(app.id, None) + + assert app_name == "OnlyApp" + assert workspace_name == "" + + # --------------------------------------------------------------------------- + # _lookup_credential_name + # --------------------------------------------------------------------------- + + class TestLookupCredentialName: + """Tests for _lookup_credential_name(credential_id, provider_type).""" + + @pytest.mark.parametrize("provider_type", ["builtin", "plugin", "api", "workflow", "mcp"]) + def test_known_provider_types_return_name(self, provider_type: str, orm_session: Session): + """Each valid provider_type results in a DB query and returns the credential name.""" + from core.ops.ops_trace_manager import _lookup_credential_name + + provider = _persist_tool_provider(orm_session, provider_type) + result = _lookup_credential_name(provider.id, provider_type) + + assert result == "CredentialA" + + def test_credential_not_found_returns_empty_string(self, orm_session: Session): + """Returns '' when DB yields None for the given credential_id.""" + from core.ops.ops_trace_manager import _lookup_credential_name + + result = _lookup_credential_name(str(uuid.uuid4()), "api") + + assert result == "" + + def test_invalid_provider_type_returns_empty_string_without_db(self): + """Returns '' immediately for an unrecognised provider_type — no DB access.""" + from core.ops.ops_trace_manager import _lookup_credential_name + + result = _lookup_credential_name(str(uuid.uuid4()), "unknown_type") + + assert result == "" + + def test_none_credential_id_returns_empty_string_without_db(self): + """Returns '' immediately when credential_id is None — no DB access.""" + from core.ops.ops_trace_manager import _lookup_credential_name + + result = _lookup_credential_name(None, "api") + + assert result == "" + + def test_none_provider_type_returns_empty_string_without_db(self): + """Returns '' immediately when provider_type is None — no DB access.""" + from core.ops.ops_trace_manager import _lookup_credential_name + + result = _lookup_credential_name(str(uuid.uuid4()), None) + + assert result == "" + + def test_builtin_and_plugin_map_to_same_model(self): + """Both 'builtin' and 'plugin' provider_types query BuiltinToolProvider.""" + from core.ops.ops_trace_manager import _PROVIDER_TYPE_TO_MODEL + from models.tools import BuiltinToolProvider + + assert _PROVIDER_TYPE_TO_MODEL["builtin"] is BuiltinToolProvider + assert _PROVIDER_TYPE_TO_MODEL["plugin"] is BuiltinToolProvider + + def test_api_maps_to_api_tool_provider(self): + """'api' maps to ApiToolProvider.""" + from core.ops.ops_trace_manager import _PROVIDER_TYPE_TO_MODEL + from models.tools import ApiToolProvider + + assert _PROVIDER_TYPE_TO_MODEL["api"] is ApiToolProvider + + def test_workflow_maps_to_workflow_tool_provider(self): + """'workflow' maps to WorkflowToolProvider.""" + from core.ops.ops_trace_manager import _PROVIDER_TYPE_TO_MODEL + from models.tools import WorkflowToolProvider + + assert _PROVIDER_TYPE_TO_MODEL["workflow"] is WorkflowToolProvider + + def test_mcp_maps_to_mcp_tool_provider(self): + """'mcp' maps to MCPToolProvider.""" + from core.ops.ops_trace_manager import _PROVIDER_TYPE_TO_MODEL + from models.tools import MCPToolProvider + + assert _PROVIDER_TYPE_TO_MODEL["mcp"] is MCPToolProvider + + # --------------------------------------------------------------------------- + # _lookup_llm_credential_info + # --------------------------------------------------------------------------- + + class TestLookupLlmCredentialInfo: + """Tests for _lookup_llm_credential_info(tenant_id, provider, model, model_type).""" + + def test_model_level_credential_found(self, orm_session: Session): + """Returns model-level credential_id and name when ProviderModel has a credential.""" + from core.ops.ops_trace_manager import _lookup_llm_credential_info + + tenant_id = str(uuid.uuid4()) + model_credential = _persist_model_credential(orm_session, tenant_id=tenant_id) + _persist_provider(orm_session, tenant_id=tenant_id, credential_id=None) + _persist_provider_model(orm_session, tenant_id=tenant_id, credential_id=model_credential.id) + + decoy_tenant_id = str(uuid.uuid4()) + decoy_credential = _persist_model_credential( + orm_session, + tenant_id=decoy_tenant_id, + credential_name="WrongTenantCredential", + ) + _persist_provider(orm_session, tenant_id=decoy_tenant_id, credential_id=None) + _persist_provider_model(orm_session, tenant_id=decoy_tenant_id, credential_id=decoy_credential.id) + + cred_id, cred_name = _lookup_llm_credential_info(tenant_id, "openai", "gpt-4") + + assert cred_id == model_credential.id + assert cred_name == "ModelCredName" + + def test_provider_level_fallback_when_no_model_credential(self, orm_session: Session): + """Falls back to provider-level credential when ProviderModel has no credential_id.""" + from core.ops.ops_trace_manager import _lookup_llm_credential_info + + tenant_id = str(uuid.uuid4()) + provider_credential = _persist_provider_credential(orm_session, tenant_id=tenant_id) + _persist_provider(orm_session, tenant_id=tenant_id, credential_id=provider_credential.id) + _persist_provider_model(orm_session, tenant_id=tenant_id, credential_id=None) + + cred_id, cred_name = _lookup_llm_credential_info(tenant_id, "openai", "gpt-4") + + assert cred_id == provider_credential.id + assert cred_name == "ProvCredName" + + def test_provider_level_fallback_when_no_model_record(self, orm_session: Session): + """Falls back to provider-level credential when no ProviderModel row exists.""" + from core.ops.ops_trace_manager import _lookup_llm_credential_info + + tenant_id = str(uuid.uuid4()) + provider_credential = _persist_provider_credential(orm_session, tenant_id=tenant_id) + _persist_provider(orm_session, tenant_id=tenant_id, credential_id=provider_credential.id) + + cred_id, cred_name = _lookup_llm_credential_info(tenant_id, "openai", "gpt-4") + + assert cred_id == provider_credential.id + assert cred_name == "ProvCredName" + + def test_no_model_arg_uses_provider_level_only(self, orm_session: Session): + """When model is None, skips ProviderModel query and uses provider credential.""" + from core.ops.ops_trace_manager import _lookup_llm_credential_info + + tenant_id = str(uuid.uuid4()) + provider_credential = _persist_provider_credential(orm_session, tenant_id=tenant_id) + _persist_provider(orm_session, tenant_id=tenant_id, credential_id=provider_credential.id) + + cred_id, cred_name = _lookup_llm_credential_info(tenant_id, "openai", None) + + assert cred_id == provider_credential.id + assert cred_name == "ProvCredName" + + def test_provider_not_found_returns_none_and_empty(self, orm_session: Session): + """Returns (None, '') when Provider record does not exist.""" + from core.ops.ops_trace_manager import _lookup_llm_credential_info + + other_tenant_id = str(uuid.uuid4()) + _persist_provider(orm_session, tenant_id=other_tenant_id, credential_id=None) + tenant_id = str(uuid.uuid4()) + + cred_id, cred_name = _lookup_llm_credential_info(tenant_id, "openai", "gpt-4") + + assert cred_id is None + assert cred_name == "" + + def test_none_tenant_id_returns_none_and_empty_without_db(self): + """Returns (None, '') immediately when tenant_id is None — no DB access.""" + from core.ops.ops_trace_manager import _lookup_llm_credential_info + + cred_id, cred_name = _lookup_llm_credential_info(None, "openai", "gpt-4") + + assert cred_id is None + assert cred_name == "" + + def test_none_provider_returns_none_and_empty_without_db(self): + """Returns (None, '') immediately when provider is None — no DB access.""" + from core.ops.ops_trace_manager import _lookup_llm_credential_info + + cred_id, cred_name = _lookup_llm_credential_info(str(uuid.uuid4()), None, "gpt-4") + + assert cred_id is None + assert cred_name == "" + + def test_db_error_on_outer_query_returns_none_and_empty(self, orm_session: Session, sqlite_engine: Engine): + """Returns (None, '') and logs a warning when the outer DB query raises.""" + from core.ops.ops_trace_manager import _lookup_llm_credential_info + + with _raise_on_table(sqlite_engine, "providers"): + cred_id, cred_name = _lookup_llm_credential_info(str(uuid.uuid4()), "openai", "gpt-4") + + assert cred_id is None + assert cred_name == "" + + def test_credential_name_lookup_failure_returns_id_with_empty_name( + self, orm_session: Session, sqlite_engine: Engine + ): + """When credential name sub-query fails, returns cred_id but '' for name.""" + from core.ops.ops_trace_manager import _lookup_llm_credential_info + + tenant_id = str(uuid.uuid4()) + provider_credential = _persist_provider_credential(orm_session, tenant_id=tenant_id) + _persist_provider(orm_session, tenant_id=tenant_id, credential_id=provider_credential.id) + + with _raise_on_table(sqlite_engine, "provider_credentials"): + cred_id, cred_name = _lookup_llm_credential_info(tenant_id, "openai", "gpt-4") + + assert cred_id == provider_credential.id + assert cred_name == "" + + def test_no_credential_on_provider_or_model_returns_none_id(self, orm_session: Session): + """Returns (None, '') when neither provider nor model has a credential_id.""" + from core.ops.ops_trace_manager import _lookup_llm_credential_info + + tenant_id = str(uuid.uuid4()) + _persist_provider(orm_session, tenant_id=tenant_id, credential_id=None) + _persist_provider_model(orm_session, tenant_id=tenant_id, credential_id=None) + + cred_id, cred_name = _lookup_llm_credential_info(tenant_id, "openai", "gpt-4") + + assert cred_id is None + assert cred_name == "" + + # --------------------------------------------------------------------------- + # TraceTask._get_user_id_from_metadata + # --------------------------------------------------------------------------- + + class TestGetUserIdFromMetadata: + """Tests for TraceTask._get_user_id_from_metadata(metadata). + + Pure dict logic — no DB access required. + """ + + @pytest.fixture + def get_user_id(self): + """Return the classmethod under test.""" + from core.ops.ops_trace_manager import TraceTask + + return TraceTask._get_user_id_from_metadata + + def test_from_end_user_id_has_highest_priority(self, get_user_id): + """from_end_user_id takes precedence over all other keys.""" + metadata = { + "from_end_user_id": "eu-abc", + "from_account_id": "acc-xyz", + "user_id": "u-123", + } + assert get_user_id(metadata) == "end_user:eu-abc" + + def test_from_account_id_used_when_no_end_user(self, get_user_id): + """from_account_id is used when from_end_user_id is absent.""" + metadata = { + "from_account_id": "acc-xyz", + "user_id": "u-123", + } + assert get_user_id(metadata) == "account:acc-xyz" + + def test_user_id_used_when_no_end_user_or_account(self, get_user_id): + """user_id is used when both higher-priority keys are absent.""" + metadata = {"user_id": "u-123"} + assert get_user_id(metadata) == "user:u-123" + + def test_returns_anonymous_when_all_keys_absent(self, get_user_id): + """Returns 'anonymous' when metadata has none of the expected keys.""" + assert get_user_id({}) == "anonymous" + + def test_empty_string_end_user_id_is_skipped(self, get_user_id): + """Empty string for from_end_user_id is falsy and falls through to next key.""" + metadata = { + "from_end_user_id": "", + "from_account_id": "acc-xyz", + } + assert get_user_id(metadata) == "account:acc-xyz" + + def test_empty_string_account_id_is_skipped(self, get_user_id): + """Empty string for from_account_id is falsy and falls through to user_id.""" + metadata = { + "from_end_user_id": "", + "from_account_id": "", + "user_id": "u-123", + } + assert get_user_id(metadata) == "user:u-123" + + def test_empty_string_user_id_falls_through_to_anonymous(self, get_user_id): + """Empty string for user_id is falsy, so 'anonymous' is returned.""" + metadata = { + "from_end_user_id": "", + "from_account_id": "", + "user_id": "", + } + assert get_user_id(metadata) == "anonymous" + + def test_only_from_end_user_id_present(self, get_user_id): + """Minimal case: only from_end_user_id present.""" + assert get_user_id({"from_end_user_id": "eu-only"}) == "end_user:eu-only" + + def test_irrelevant_keys_do_not_interfere(self, get_user_id): + """Extra metadata keys have no effect on the result.""" + metadata = {"invoke_from": "web", "app_id": "a1"} + assert get_user_id(metadata) == "anonymous" + + +"""SQLite-backed tests for :mod:`core.ops.ops_trace_manager`.""" class DummyConfig: @@ -156,46 +747,6 @@ class RecordingDispatcher: self.options.append(kwargs) -@pytest.fixture -def database(sqlite_engine: Engine, sqlite_session: Session) -> Iterator[Session]: - session_proxy = scoped_session(lambda: sqlite_session) - with ( - patch.object(module.db, "session", session_proxy), - patch.object(type(module.db), "engine", new_callable=PropertyMock, return_value=sqlite_engine), - ): - yield sqlite_session - - -@pytest.fixture -def trace_environment(monkeypatch: pytest.MonkeyPatch) -> Iterator[None]: - monkeypatch.setattr(module, "provider_config_map", FakeProviderMap({"dummy": PROVIDER_ENTRY})) - monkeypatch.setattr(module, "unified_provider_config_map", FakeProviderMap({})) - apply_config_overrides(monkeypatch, OPS_TRACE_UNIFIED_ENABLED=False) - OpsTraceManager.ops_trace_instances_cache.clear() - OpsTraceManager.decrypted_configs_cache.clear() - monkeypatch.setattr(module.threading, "Timer", DummyTimer) - monkeypatch.setattr(module, "trace_manager_queue", queue.Queue()) - monkeypatch.setattr(module, "trace_manager_timer", None) - monkeypatch.setattr("core.telemetry.gateway.is_enterprise_telemetry_enabled", lambda: False) - - app = Flask(__name__) - with app.app_context(): - yield - - -@pytest.fixture -def encryption_functions( - monkeypatch: pytest.MonkeyPatch, -) -> tuple[EncryptTokenRecorder, BatchDecryptTokenRecorder, ObfuscatedTokenRecorder]: - encrypt = EncryptTokenRecorder() - decrypt = BatchDecryptTokenRecorder() - obfuscate = ObfuscatedTokenRecorder() - monkeypatch.setattr(module, "encrypt_token", encrypt) - monkeypatch.setattr(module, "batch_decrypt_token", decrypt) - monkeypatch.setattr(module, "obfuscated_token", obfuscate) - return encrypt, decrypt, obfuscate - - def _app(session: Session, *, app_id: str = "app-id", tracing: str | None = None) -> App: app = make_app(app_id=app_id, name="App", description="description", icon_type=None, tracing=tracing) session.add(app) @@ -298,510 +849,837 @@ def _message_data(**overrides): ) -def test_encrypt_decrypt_obfuscate_and_cache( - trace_environment: None, - encryption_functions: tuple[EncryptTokenRecorder, BatchDecryptTokenRecorder, ObfuscatedTokenRecorder], -) -> None: - encrypted = OpsTraceManager.encrypt_tracing_config( - "tenant-1", "dummy", {"secret_value": "value", "other_value": "info"} +class TestOpsTraceManagerCases: + @staticmethod + @pytest.fixture + def database(sqlite_engine: Engine, sqlite_session: Session) -> Iterator[Session]: + session_proxy = scoped_session(lambda: sqlite_session) + with ( + patch.object(module.db, "session", session_proxy), + patch.object(type(module.db), "engine", new_callable=PropertyMock, return_value=sqlite_engine), + ): + yield sqlite_session + + @staticmethod + @pytest.fixture + def trace_environment(monkeypatch: pytest.MonkeyPatch) -> Iterator[None]: + monkeypatch.setattr(module, "provider_config_map", FakeProviderMap({"dummy": PROVIDER_ENTRY})) + monkeypatch.setattr(module, "unified_provider_config_map", FakeProviderMap({})) + apply_config_overrides(monkeypatch, OPS_TRACE_UNIFIED_ENABLED=False) + OpsTraceManager.ops_trace_instances_cache.clear() + OpsTraceManager.decrypted_configs_cache.clear() + monkeypatch.setattr(module.threading, "Timer", DummyTimer) + monkeypatch.setattr(module, "trace_manager_queue", queue.Queue()) + monkeypatch.setattr(module, "trace_manager_timer", None) + monkeypatch.setattr("core.telemetry.gateway.is_enterprise_telemetry_enabled", lambda: False) + + app = Flask(__name__) + with app.app_context(): + yield + + @staticmethod + @pytest.fixture + def encryption_functions( + monkeypatch: pytest.MonkeyPatch, + ) -> tuple[EncryptTokenRecorder, BatchDecryptTokenRecorder, ObfuscatedTokenRecorder]: + encrypt = EncryptTokenRecorder() + decrypt = BatchDecryptTokenRecorder() + obfuscate = ObfuscatedTokenRecorder() + monkeypatch.setattr(module, "encrypt_token", encrypt) + monkeypatch.setattr(module, "batch_decrypt_token", decrypt) + monkeypatch.setattr(module, "obfuscated_token", obfuscate) + return encrypt, decrypt, obfuscate + + @staticmethod + def test_encrypt_decrypt_obfuscate_and_cache( + trace_environment: None, + encryption_functions: tuple[EncryptTokenRecorder, BatchDecryptTokenRecorder, ObfuscatedTokenRecorder], + ) -> None: + encrypted = OpsTraceManager.encrypt_tracing_config( + "tenant-1", "dummy", {"secret_value": "value", "other_value": "info"} + ) + assert encrypted == {"secret_value": "enc-value", "other_value": "info"} + preserved = OpsTraceManager.encrypt_tracing_config( + "tenant-1", "dummy", {"secret_value": "*"}, current_trace_config={"secret_value": "keep"} + ) + assert preserved["secret_value"] == "keep" + first = OpsTraceManager.decrypt_tracing_config("tenant-1", "dummy", encrypted) + second = OpsTraceManager.decrypt_tracing_config("tenant-1", "dummy", encrypted) + assert first == second + assert len(encryption_functions[1].calls) == 1 + obfuscated = OpsTraceManager.obfuscated_decrypt_token("dummy", first) + assert obfuscated["secret_value"] == "ob-dec-enc-value" + assert encryption_functions[2].calls == ["dec-enc-value"] + + @staticmethod + def test_decrypted_config_reads_real_trace_and_app_rows( + trace_environment: None, + encryption_functions, + database: Session, + ) -> None: + app = _app(database) + trace = TraceAppConfig( + app_id=app.id, + tracing_provider="dummy", + tracing_config={"secret_value": "encrypted", "other_value": "info"}, + ) + database.add(trace) + database.commit() + result = OpsTraceManager.get_decrypted_tracing_config(app.id, "dummy") + assert result == {"secret_value": "dec-encrypted", "other_value": "info"} + assert OpsTraceManager.get_decrypted_tracing_config(app.id, "missing") is None + + null_config_app = _app(database, app_id="app-null-config") + database.add(TraceAppConfig(app_id=null_config_app.id, tracing_provider="dummy", tracing_config=None)) + database.commit() + with pytest.raises(ValueError, match="Tracing config cannot be None"): + OpsTraceManager.get_decrypted_tracing_config(null_config_app.id, "dummy") + + database.delete(app) + database.commit() + with pytest.raises(ValueError, match="App not found"): + OpsTraceManager.get_decrypted_tracing_config("app-id", "dummy") + + @staticmethod + def test_ops_trace_instance_uses_persisted_enabled_state_and_cache( + trace_environment: None, + encryption_functions, + database: Session, + ) -> None: + app = _app(database, tracing=json.dumps({"enabled": False, "tracing_provider": "dummy"})) + assert OpsTraceManager.get_ops_trace_instance(app.id) is None + app.tracing = json.dumps({"enabled": True, "tracing_provider": "dummy"}) + database.add( + TraceAppConfig(app_id=app.id, tracing_provider="dummy", tracing_config={"secret_value": "encrypted"}) + ) + database.commit() + instance = OpsTraceManager.get_ops_trace_instance(app.id) + assert isinstance(instance, DummyTraceInstance) + assert OpsTraceManager.get_ops_trace_instance(app.id) is instance + + app.tracing = json.dumps({"enabled": True, "tracing_provider": "missing"}) + database.commit() + assert OpsTraceManager.get_ops_trace_instance(app.id) is None + + assert OpsTraceManager.get_ops_trace_instance(None) is None + assert OpsTraceManager.get_ops_trace_instance("tenant-storage-id") is None + assert OpsTraceManager.get_ops_trace_instance("missing") is None + + @staticmethod + @pytest.mark.parametrize( + ("enabled", "registered", "expected_type"), + [ + (False, False, DummyTraceInstance), + (False, True, DummyTraceInstance), + (True, False, DummyTraceInstance), + (True, True, DummyUnifiedTraceInstance), + ], ) - assert encrypted == {"secret_value": "enc-value", "other_value": "info"} - preserved = OpsTraceManager.encrypt_tracing_config( - "tenant-1", "dummy", {"secret_value": "*"}, current_trace_config={"secret_value": "keep"} - ) - assert preserved["secret_value"] == "keep" - first = OpsTraceManager.decrypt_tracing_config("tenant-1", "dummy", encrypted) - second = OpsTraceManager.decrypt_tracing_config("tenant-1", "dummy", encrypted) - assert first == second - assert len(encryption_functions[1].calls) == 1 - obfuscated = OpsTraceManager.obfuscated_decrypt_token("dummy", first) - assert obfuscated["secret_value"] == "ob-dec-enc-value" - assert encryption_functions[2].calls == ["dec-enc-value"] + def test_ops_trace_instance_routes_by_unified_switch( + enabled: bool, + registered: bool, + expected_type: type[DummyTraceInstance], + monkeypatch: pytest.MonkeyPatch, + trace_environment: None, + encryption_functions, + database: Session, + ) -> None: + app = _app(database, tracing=json.dumps({"enabled": True, "tracing_provider": "dummy"})) + database.add(TraceAppConfig(app_id=app.id, tracing_provider="dummy", tracing_config={})) + database.commit() + apply_config_overrides(monkeypatch, OPS_TRACE_UNIFIED_ENABLED=enabled) + entries = {"dummy": UNIFIED_PROVIDER_ENTRY} if registered else {} + monkeypatch.setattr(module, "unified_provider_config_map", FakeProviderMap(entries)) + instance = OpsTraceManager.get_ops_trace_instance(app.id) -def test_decrypted_config_reads_real_trace_and_app_rows( - trace_environment: None, - encryption_functions, - database: Session, -) -> None: - app = _app(database) - trace = TraceAppConfig( - app_id=app.id, - tracing_provider="dummy", - tracing_config={"secret_value": "encrypted", "other_value": "info"}, - ) - database.add(trace) - database.commit() - result = OpsTraceManager.get_decrypted_tracing_config(app.id, "dummy") - assert result == {"secret_value": "dec-encrypted", "other_value": "info"} - assert OpsTraceManager.get_decrypted_tracing_config(app.id, "missing") is None + assert type(instance) is expected_type - null_config_app = _app(database, app_id="app-null-config") - database.add(TraceAppConfig(app_id=null_config_app.id, tracing_provider="dummy", tracing_config=None)) - database.commit() - with pytest.raises(ValueError, match="Tracing config cannot be None"): - OpsTraceManager.get_decrypted_tracing_config(null_config_app.id, "dummy") - - database.delete(app) - database.commit() - with pytest.raises(ValueError, match="App not found"): - OpsTraceManager.get_decrypted_tracing_config("app-id", "dummy") - - -def test_ops_trace_instance_uses_persisted_enabled_state_and_cache( - trace_environment: None, - encryption_functions, - database: Session, -) -> None: - app = _app(database, tracing=json.dumps({"enabled": False, "tracing_provider": "dummy"})) - assert OpsTraceManager.get_ops_trace_instance(app.id) is None - app.tracing = json.dumps({"enabled": True, "tracing_provider": "dummy"}) - database.add(TraceAppConfig(app_id=app.id, tracing_provider="dummy", tracing_config={"secret_value": "encrypted"})) - database.commit() - instance = OpsTraceManager.get_ops_trace_instance(app.id) - assert isinstance(instance, DummyTraceInstance) - assert OpsTraceManager.get_ops_trace_instance(app.id) is instance - - app.tracing = json.dumps({"enabled": True, "tracing_provider": "missing"}) - database.commit() - assert OpsTraceManager.get_ops_trace_instance(app.id) is None - - assert OpsTraceManager.get_ops_trace_instance(None) is None - assert OpsTraceManager.get_ops_trace_instance("tenant-storage-id") is None - assert OpsTraceManager.get_ops_trace_instance("missing") is None - - -@pytest.mark.parametrize( - ("enabled", "registered", "expected_type"), - [ - (False, False, DummyTraceInstance), - (False, True, DummyTraceInstance), - (True, False, DummyTraceInstance), - (True, True, DummyUnifiedTraceInstance), - ], -) -def test_ops_trace_instance_routes_by_unified_switch( - enabled: bool, - registered: bool, - expected_type: type[DummyTraceInstance], - monkeypatch: pytest.MonkeyPatch, - trace_environment: None, - encryption_functions, - database: Session, -) -> None: - app = _app(database, tracing=json.dumps({"enabled": True, "tracing_provider": "dummy"})) - database.add(TraceAppConfig(app_id=app.id, tracing_provider="dummy", tracing_config={})) - database.commit() - apply_config_overrides(monkeypatch, OPS_TRACE_UNIFIED_ENABLED=enabled) - entries = {"dummy": UNIFIED_PROVIDER_ENTRY} if registered else {} - monkeypatch.setattr(module, "unified_provider_config_map", FakeProviderMap(entries)) - - instance = OpsTraceManager.get_ops_trace_instance(app.id) - - assert type(instance) is expected_type - - -def test_registered_unified_provider_does_not_fallback_when_construction_fails( - monkeypatch: pytest.MonkeyPatch, - trace_environment: None, - encryption_functions, - database: Session, -) -> None: - app = _app(database, tracing=json.dumps({"enabled": True, "tracing_provider": "dummy"})) - database.add(TraceAppConfig(app_id=app.id, tracing_provider="dummy", tracing_config={})) - database.commit() - apply_config_overrides(monkeypatch, OPS_TRACE_UNIFIED_ENABLED=True) - monkeypatch.setattr( - module, - "unified_provider_config_map", - FakeProviderMap( - { - "dummy": { - "config_class": DummyConfig, - "trace_instance": FailingUnifiedTraceInstance, + @staticmethod + def test_registered_unified_provider_does_not_fallback_when_construction_fails( + monkeypatch: pytest.MonkeyPatch, + trace_environment: None, + encryption_functions, + database: Session, + ) -> None: + app = _app(database, tracing=json.dumps({"enabled": True, "tracing_provider": "dummy"})) + database.add(TraceAppConfig(app_id=app.id, tracing_provider="dummy", tracing_config={})) + database.commit() + apply_config_overrides(monkeypatch, OPS_TRACE_UNIFIED_ENABLED=True) + monkeypatch.setattr( + module, + "unified_provider_config_map", + FakeProviderMap( + { + "dummy": { + "config_class": DummyConfig, + "trace_instance": FailingUnifiedTraceInstance, + } } - } - ), - ) + ), + ) - with pytest.raises(RuntimeError, match="unified constructor failed"): - OpsTraceManager.get_ops_trace_instance(app.id) + with pytest.raises(RuntimeError, match="unified constructor failed"): + OpsTraceManager.get_ops_trace_instance(app.id) + @staticmethod + def test_unified_and_legacy_instances_have_separate_cache_entries( + monkeypatch: pytest.MonkeyPatch, + trace_environment: None, + encryption_functions, + database: Session, + ) -> None: + app = _app(database, tracing=json.dumps({"enabled": True, "tracing_provider": "dummy"})) + database.add(TraceAppConfig(app_id=app.id, tracing_provider="dummy", tracing_config={})) + database.commit() + monkeypatch.setattr(module, "unified_provider_config_map", FakeProviderMap({"dummy": UNIFIED_PROVIDER_ENTRY})) -def test_unified_and_legacy_instances_have_separate_cache_entries( - monkeypatch: pytest.MonkeyPatch, - trace_environment: None, - encryption_functions, - database: Session, -) -> None: - app = _app(database, tracing=json.dumps({"enabled": True, "tracing_provider": "dummy"})) - database.add(TraceAppConfig(app_id=app.id, tracing_provider="dummy", tracing_config={})) - database.commit() - monkeypatch.setattr(module, "unified_provider_config_map", FakeProviderMap({"dummy": UNIFIED_PROVIDER_ENTRY})) + apply_config_overrides(monkeypatch, OPS_TRACE_UNIFIED_ENABLED=False) + legacy = OpsTraceManager.get_ops_trace_instance(app.id) + apply_config_overrides(monkeypatch, OPS_TRACE_UNIFIED_ENABLED=True) + unified = OpsTraceManager.get_ops_trace_instance(app.id) - apply_config_overrides(monkeypatch, OPS_TRACE_UNIFIED_ENABLED=False) - legacy = OpsTraceManager.get_ops_trace_instance(app.id) - apply_config_overrides(monkeypatch, OPS_TRACE_UNIFIED_ENABLED=True) - unified = OpsTraceManager.get_ops_trace_instance(app.id) + assert type(legacy) is DummyTraceInstance + assert type(unified) is DummyUnifiedTraceInstance + assert legacy is not unified - assert type(legacy) is DummyTraceInstance - assert type(unified) is DummyUnifiedTraceInstance - assert legacy is not unified + @staticmethod + def test_message_config_lookup_uses_real_conversation_and_model_config(database: Session) -> None: + app = _app(database) + config = AppModelConfig(app_id=app.id, model='{"provider":"openai"}') + database.add(config) + database.commit() + conversation, message = _conversation_message(database, app, config=config) + result = OpsTraceManager.get_app_config_through_message_id(message.id) + assert result.id == config.id + conversation.app_model_config_id = None + conversation.override_model_configs = json.dumps({"provider": "override"}) + database.commit() + override = OpsTraceManager.get_app_config_through_message_id(message.id) + assert json.loads(override) == {"provider": "override"} -def test_message_config_lookup_uses_real_conversation_and_model_config(database: Session) -> None: - app = _app(database) - config = AppModelConfig(app_id=app.id, model='{"provider":"openai"}') - database.add(config) - database.commit() - conversation, message = _conversation_message(database, app, config=config) - result = OpsTraceManager.get_app_config_through_message_id(message.id) - assert result.id == config.id + assert OpsTraceManager.get_app_config_through_message_id("missing") is None - conversation.app_model_config_id = None - conversation.override_model_configs = json.dumps({"provider": "override"}) - database.commit() - override = OpsTraceManager.get_app_config_through_message_id(message.id) - assert json.loads(override) == {"provider": "override"} - - assert OpsTraceManager.get_app_config_through_message_id("missing") is None - - -@pytest.mark.parametrize("imports_cached", [False, True], ids=["cold", "warm"]) -@pytest.mark.parametrize("module_name", ["dify_trace_weave", "weave"]) -def test_provider_loader_identifies_missing_optional_dependency(module_name: str, imports_cached: bool) -> None: - with patch.dict(sys.modules): - if imports_cached: - module.OpsTraceProviderConfigMap()["weave"] - # Cached typed SDK submodules can bypass a missing top-level package. - # Simulate its absence consistently and restore the complete cache after. - for cached_name in list(sys.modules): - if cached_name.partition(".")[0] in {"dify_trace_weave", "weave"}: - del sys.modules[cached_name] - with patch.dict(sys.modules, {module_name: None}): - with pytest.raises(TraceProviderNotInstalledError, match=f"weave.*{module_name}") as caught: + @staticmethod + @pytest.mark.parametrize("imports_cached", [False, True], ids=["cold", "warm"]) + @pytest.mark.parametrize("module_name", ["dify_trace_weave", "weave"]) + def test_provider_loader_identifies_missing_optional_dependency(module_name: str, imports_cached: bool) -> None: + with patch.dict(sys.modules): + if imports_cached: module.OpsTraceProviderConfigMap()["weave"] + # Cached typed SDK submodules can bypass a missing top-level package. + # Simulate its absence consistently and restore the complete cache after. + for cached_name in list(sys.modules): + if cached_name.partition(".")[0] in {"dify_trace_weave", "weave"}: + del sys.modules[cached_name] + with patch.dict(sys.modules, {module_name: None}): + with pytest.raises(TraceProviderNotInstalledError, match=f"weave.*{module_name}") as caught: + module.OpsTraceProviderConfigMap()["weave"] - cause = caught.value.__cause__ - assert isinstance(cause, ModuleNotFoundError) - assert cause.name is not None - assert cause.name.partition(".")[0] == module_name + cause = caught.value.__cause__ + assert isinstance(cause, ModuleNotFoundError) + assert cause.name is not None + assert cause.name.partition(".")[0] == module_name - -@pytest.mark.parametrize( - "failure", - [ - ImportError("cannot import name 'WeaveDataTrace'"), - ModuleNotFoundError("SDK import failed without identifying a missing module"), - ModuleNotFoundError("No module named 'json.missing_module'", name="json.missing_module"), - ], -) -def test_provider_loader_preserves_import_errors_in_installed_packages(failure: ImportError) -> None: - original_import = builtins.__import__ - - def broken_import(name, *args, **kwargs): - if name.startswith("dify_trace_weave."): - raise failure - return original_import(name, *args, **kwargs) - - with patch("builtins.__import__", side_effect=broken_import), pytest.raises(ImportError) as caught: - module.OpsTraceProviderConfigMap()["weave"] - - assert caught.value is failure - - -def test_message_trace_reads_real_conversation_app_and_message_file( - trace_environment: None, - database: Session, -) -> None: - app = _app(database) - _, message = _conversation_message(database, app) - file = MessageFile( - message_id=message.id, - type=FileType.DOCUMENT, - transfer_method=FileTransferMethod.REMOTE_URL, - created_by_role=CreatorUserRole.ACCOUNT, - created_by="user-1", - url="path/to/file", + @staticmethod + @pytest.mark.parametrize( + "failure", + [ + ImportError("cannot import name 'WeaveDataTrace'"), + ModuleNotFoundError("SDK import failed without identifying a missing module"), + ModuleNotFoundError("No module named 'json.missing_module'", name="json.missing_module"), + ], ) - database.add(file) - database.commit() - result = TraceTask( - trace_type=TraceTaskName.MESSAGE_TRACE, - message_id=message.id, - ).preprocess() - assert result.message_id == message.id - assert result.conversation_mode == AppMode.CHAT - assert result.file_list[0].endswith("path/to/file") - assert result.metadata["tenant_id"] == "tenant-1" + def test_provider_loader_preserves_import_errors_in_installed_packages(failure: ImportError) -> None: + original_import = builtins.__import__ + + def broken_import(name, *args, **kwargs): + if name.startswith("dify_trace_weave."): + raise failure + return original_import(name, *args, **kwargs) + + with patch("builtins.__import__", side_effect=broken_import), pytest.raises(ImportError) as caught: + module.OpsTraceProviderConfigMap()["weave"] + + assert caught.value is failure + + @staticmethod + def test_message_trace_reads_real_conversation_app_and_message_file( + trace_environment: None, + database: Session, + ) -> None: + app = _app(database) + _, message = _conversation_message(database, app) + file = MessageFile( + message_id=message.id, + type=FileType.DOCUMENT, + transfer_method=FileTransferMethod.REMOTE_URL, + created_by_role=CreatorUserRole.ACCOUNT, + created_by="user-1", + url="path/to/file", + ) + database.add(file) + database.commit() + result = TraceTask( + trace_type=TraceTaskName.MESSAGE_TRACE, + message_id=message.id, + ).preprocess() + assert result.message_id == message.id + assert result.conversation_mode == AppMode.CHAT + assert result.file_list[0].endswith("path/to/file") + assert result.metadata["tenant_id"] == "tenant-1" + + @staticmethod + def test_workflow_log_enriches_moderation_and_suggested_question_traces( + database: Session, + ) -> None: + app = _app(database) + _, message = _conversation_message(database, app) + log = WorkflowAppLog( + tenant_id="tenant-1", + app_id="app-id", + workflow_id="workflow-1", + workflow_run_id="run-1", + created_from=WorkflowAppLogCreatedFrom.WEB_APP, + created_by_role=CreatorUserRole.ACCOUNT, + created_by="user-1", + ) + database.add(log) + database.commit() + task = TraceTask(trace_type=TraceTaskName.MODERATION_TRACE, message_id=message.id) + moderation = ModerationInputsResult( + action=ModerationAction.DIRECT_OUTPUT, + preset_response="no", + query="q", + flagged=True, + ) + result = task.moderation_trace( + message.id, + {"start": 1, "end": 2}, + moderation_result=moderation, + inputs={"source": "payload"}, + ) + assert result.message_id == log.id + assert result.flagged is True + assert result.inputs == {"source": "payload"} + suggested = task.suggested_question_trace(message.id, {"start": 1, "end": 2}, suggested_question=["q1"]) + assert suggested.message_id == log.id + assert suggested.suggested_question == ["q1"] + + @staticmethod + def test_dataset_retrieval_trace_serializes_documents( + trace_environment: None, + database: Session, + ) -> None: + app = _app(database) + _, message = _conversation_message(database, app) + document = RetrievalDocument(page_content="value") + + result = TraceTask(trace_type=TraceTaskName.DATASET_RETRIEVAL_TRACE).dataset_retrieval_trace( + message.id, + {"start": 1, "end": 2}, + documents=[document], + ) + + assert result.documents == [document.model_dump()] + assert result.documents[0]["page_content"] == "value" + assert result.metadata["tenant_id"] == "tenant-1" + + @staticmethod + def test_workflow_trace_reads_real_workflow_log_from_owned_session( + monkeypatch: pytest.MonkeyPatch, + trace_environment: None, + database: Session, + sqlite_session_factory: sessionmaker[Session], + ) -> None: + app = _app(database) + log = WorkflowAppLog( + tenant_id=app.tenant_id, + app_id=app.id, + workflow_id="workflow-1", + workflow_run_id="run-1", + created_from=WorkflowAppLogCreatedFrom.WEB_APP, + created_by_role=CreatorUserRole.ACCOUNT, + created_by="user-1", + ) + workflow_run = WorkflowRun( + id="run-1", + tenant_id=app.tenant_id, + app_id=app.id, + workflow_id="workflow-1", + type=WorkflowType.WORKFLOW, + triggered_from=WorkflowRunTriggeredFrom.APP_RUN, + version="3", + graph="{}", + inputs=json.dumps({"query": "search"}), + status=WorkflowExecutionStatus.SUCCEEDED, + outputs=json.dumps({"out": "value"}), + error=None, + elapsed_time=10, + total_tokens=12, + total_steps=1, + created_by_role=CreatorUserRole.ACCOUNT, + created_by="user-1", + created_at=datetime(2025, 2, 20, 10, 0, 0), + finished_at=datetime(2025, 2, 20, 10, 0, 5), + ) + database.add_all([log, workflow_run]) + database.commit() + repo = DifyAPISQLAlchemyWorkflowRunRepository(sqlite_session_factory) + monkeypatch.setattr(TraceTask, "_get_workflow_run_repo", classmethod(lambda cls: repo)) + monkeypatch.setattr(TraceTask, "_calculate_workflow_token_split", classmethod(lambda cls, *_a, **_k: (5, 7))) + result = TraceTask(trace_type=TraceTaskName.WORKFLOW_TRACE).workflow_trace( + workflow_run_id="run-1", conversation_id=None, user_id="user-1" + ) + assert result.workflow_run_id == "run-1" + assert result.workflow_id == "workflow-1" + assert result.workflow_app_log_id == log.id + assert result.prompt_tokens == 5 + assert result.completion_tokens == 7 + + @staticmethod + def test_tool_trace_reads_real_message_file(database: Session) -> None: + app = _app(database) + _, message = _conversation_message(database, app) + file = MessageFile( + message_id=message.id, + type=FileType.DOCUMENT, + transfer_method=FileTransferMethod.REMOTE_URL, + created_by_role=CreatorUserRole.ACCOUNT, + created_by="user-1", + url="tool/file", + ) + thought = MessageAgentThought( + message_id=message.id, + position=1, + created_by_role=CreatorUserRole.ACCOUNT, + created_by="user-1", + tool="tool-a", + tool_meta_str=json.dumps( + {"tool-a": {"tool_config": {}, "time_cost": 5, "error": "", "tool_parameters": {}}} + ), + ) + database.add_all([file, thought]) + database.commit() + result = TraceTask(trace_type=TraceTaskName.TOOL_TRACE).tool_trace( + message.id, {"start": 1, "end": 2}, tool_name="tool-a", tool_inputs={}, tool_outputs="result" + ) + assert result.tool_name == "tool-a" + assert result.time_cost == 5 + assert result.message_file_data.id == file.id + + @staticmethod + def test_node_execution_trace_resolves_real_message_by_conversation_and_run( + trace_environment: None, database: Session + ) -> None: + app = _app(database) + conversation, message = _conversation_message(database, app) + result = TraceTask(trace_type=TraceTaskName.NODE_EXECUTION_TRACE).node_execution_trace( + node_execution_data={ + "tenant_id": app.tenant_id, + "app_id": app.id, + "conversation_id": conversation.id, + "workflow_execution_id": message.workflow_run_id, + "workflow_id": "workflow-1", + "node_execution_id": "node-execution-1", + "node_id": "node-1", + "node_type": "llm", + "title": "Node", + "status": "succeeded", + } + ) + assert result.message_id == message.id + assert result.metadata["conversation_id"] == conversation.id + + @staticmethod + def test_trace_helpers_and_streaming_metrics(trace_environment: None) -> None: + assert OpsTraceManager.check_trace_config_is_effective({}, "dummy") + assert OpsTraceManager.get_trace_config_project_key({}, "dummy") == "fake-key" + assert OpsTraceManager.get_trace_config_project_url({}, "dummy") == "https://project.fake" + task = TraceTask(trace_type=TraceTaskName.MESSAGE_TRACE, message_id="message-1") + assert task.conversation_trace(foo="bar") == {"foo": "bar"} + assert task._extract_streaming_metrics(Message(message_metadata="invalid")) == {} + assert task.generate_name_trace("conversation", {"start": 1, "end": 2}, tenant_id=None) == {} + generated = task.generate_name_trace( + "conversation", + {"start": 1, "end": 2}, + tenant_id="tenant-1", + generate_conversation_name="name", + inputs="query", + ) + assert generated.outputs == "name" + assert generated.tenant_id == "tenant-1" + assert generated.message_id == "message-1" + + @staticmethod + def test_trace_queue_preserves_source_app_when_another_manager_drains( + monkeypatch: pytest.MonkeyPatch, trace_environment: None + ) -> None: + monkeypatch.setattr( + OpsTraceManager, + "get_ops_trace_instance", + classmethod(lambda cls, app_id: app_id == "source-app-id"), + ) + draining_manager = TraceQueueManager(app_id="draining-app-id", user_id="user-1") + source_manager = TraceQueueManager(app_id="source-app-id", user_id="user-2") + task = TraceTask( + trace_type=TraceTaskName.GENERATE_NAME_TRACE, + conversation_id="conversation-1", + timer={"start": 1, "end": 2}, + tenant_id="tenant-1", + generate_conversation_name="name", + inputs="query", + ) + source_manager.add_trace_task(task) + assert draining_manager.collect_tasks() == [task] + + recording_storage = RecordingStorage() + dispatcher = RecordingDispatcher() + monkeypatch.setattr(module.storage, "save", recording_storage.save) + monkeypatch.setattr(module.process_trace_tasks, "apply_async", dispatcher.apply_async) + file_id = UUID("00000000-0000-0000-0000-000000000123") + monkeypatch.setattr(module, "uuid4", lambda: file_id) + source_manager.add_trace_task(task) + draining_manager.run() + + assert len(recording_storage.writes) == 1 + path, data = recording_storage.writes[0] + assert path.endswith(f"source-app-id/{file_id.hex}.json") + assert json.loads(data)["app_id"] == "source-app-id" + assert dispatcher.payloads == [{"file_id": file_id.hex, "app_id": "source-app-id"}] + + @staticmethod + def test_trace_queue_persists_with_caller_supplied_file_id( + monkeypatch: pytest.MonkeyPatch, trace_environment: None + ) -> None: + monkeypatch.setattr(OpsTraceManager, "get_ops_trace_instance", classmethod(lambda cls, _app_id: True)) + manager = TraceQueueManager(app_id="app-id", user_id="user-1") + task = TraceTask( + trace_type=TraceTaskName.GENERATE_NAME_TRACE, + conversation_id="conversation-1", + timer={"start": 1, "end": 2}, + tenant_id="tenant-1", + generate_conversation_name="name", + inputs="query", + ) + task.app_id = "app-id" + recording_storage = RecordingStorage() + monkeypatch.setattr(module.storage, "save", recording_storage.save) + + file_info = manager.persist_trace_task(task, file_id="workflow-final-run-1") + + assert file_info == {"file_id": "workflow-final-run-1", "app_id": "app-id"} + path, data = recording_storage.writes[0] + payload = json.loads(data) + assert path == "ops_trace/app-id/workflow-final-run-1.json" + assert UUID(payload["trace_info"]["operation_id"]) + + @staticmethod + def test_trace_queue_persistence_error_propagates(monkeypatch: pytest.MonkeyPatch, trace_environment: None) -> None: + monkeypatch.setattr(OpsTraceManager, "get_ops_trace_instance", classmethod(lambda cls, _app_id: True)) + manager = TraceQueueManager(app_id="app-id", user_id="user-1") + task = TraceTask(trace_type=TraceTaskName.GENERATE_NAME_TRACE) + task.app_id = "app-id" + + def fail_save(_path: str, _data: bytes) -> None: + raise OSError("storage unavailable") + + monkeypatch.setattr(module.storage, "save", fail_save) + + with pytest.raises(OSError, match="storage unavailable"): + manager.persist_trace_task(task, file_id="workflow-final-run-1") + + @staticmethod + def test_trace_queue_enqueue_error_propagates(monkeypatch: pytest.MonkeyPatch, trace_environment: None) -> None: + monkeypatch.setattr(OpsTraceManager, "get_ops_trace_instance", classmethod(lambda cls, _app_id: True)) + manager = TraceQueueManager(app_id="app-id", user_id="user-1") + + def fail_enqueue(*_args: object, **_kwargs: object) -> None: + raise ConnectionError("broker unavailable") + + monkeypatch.setattr(module.process_trace_tasks, "apply_async", fail_enqueue) + + with pytest.raises(ConnectionError, match="broker unavailable"): + manager.enqueue_persisted_trace({"file_id": "workflow-final-run-1", "app_id": "app-id"}) -def test_workflow_log_enriches_moderation_and_suggested_question_traces( - database: Session, -) -> None: - app = _app(database) - _, message = _conversation_message(database, app) - log = WorkflowAppLog( - tenant_id="tenant-1", - app_id="app-id", - workflow_id="workflow-1", - workflow_run_id="run-1", - created_from=WorkflowAppLogCreatedFrom.WEB_APP, - created_by_role=CreatorUserRole.ACCOUNT, - created_by="user-1", - ) - database.add(log) - database.commit() - task = TraceTask(trace_type=TraceTaskName.MODERATION_TRACE, message_id=message.id) - moderation = ModerationInputsResult( - action=ModerationAction.DIRECT_OUTPUT, - preset_response="no", - query="q", - flagged=True, - ) - result = task.moderation_trace( - message.id, - {"start": 1, "end": 2}, - moderation_result=moderation, - inputs={"source": "payload"}, - ) - assert result.message_id == log.id - assert result.flagged is True - assert result.inputs == {"source": "payload"} - suggested = task.suggested_question_trace(message.id, {"start": 1, "end": 2}, suggested_question=["q1"]) - assert suggested.message_id == log.id - assert suggested.suggested_question == ["q1"] +"""Unit tests for TraceQueueManager telemetry guard. + +Verifies that TraceQueueManager.add_trace_task() only enqueues tasks when at +least one consumer is active: +- Enterprise telemetry is enabled (_enterprise_telemetry_enabled=True), OR +- A third-party trace instance (Langfuse, etc.) is configured + +When neither is active, tasks are silently dropped to avoid unnecessary work. + +When BOTH are false, tasks are silently dropped (correct behavior). +""" -def test_dataset_retrieval_trace_serializes_documents( - trace_environment: None, - database: Session, -) -> None: - app = _app(database) - _, message = _conversation_message(database, app) - document = RetrievalDocument(page_content="value") +class TestTraceQueueManagerCases: + @staticmethod + @pytest.fixture(autouse=True) + def _application_tracing_disabled(monkeypatch: pytest.MonkeyPatch) -> None: + apply_config_overrides(monkeypatch, OPS_AGENT_TRACE_ENABLED=False) - result = TraceTask(trace_type=TraceTaskName.DATASET_RETRIEVAL_TRACE).dataset_retrieval_trace( - message.id, - {"start": 1, "end": 2}, - documents=[document], - ) + @staticmethod + @pytest.fixture + def trace_queue_manager_and_task(): + """Exercise the production guard without starting background dispatch.""" + from core.ops.entities.trace_entity import TraceTaskName + from core.ops.ops_trace_manager import TraceQueueManager, TraceTask - assert result.documents == [document.model_dump()] - assert result.documents[0]["page_content"] == "value" - assert result.metadata["tenant_id"] == "tenant-1" + with patch.object(TraceQueueManager, "start_timer"): + yield TraceQueueManager, TraceTask, TraceTaskName + + class TestTraceQueueManagerTelemetryGuard: + """Test TraceQueueManager's telemetry guard in add_trace_task().""" + + def test_task_not_enqueued_when_telemetry_disabled_and_no_trace_instance(self, trace_queue_manager_and_task): + """Verify task is NOT enqueued when telemetry disabled and no trace instance. + + This is the core guard: when _enterprise_telemetry_enabled=False AND + trace_instance=None, the task should be silently dropped. + """ + trace_queue_manager, trace_task_type, trace_task_name = trace_queue_manager_and_task + + trace_queue = queue.Queue() + + trace_task = trace_task_type(trace_type=trace_task_name.WORKFLOW_TRACE) + + with ( + patch("core.telemetry.gateway.is_enterprise_telemetry_enabled", return_value=False), + patch("core.ops.ops_trace_manager.OpsTraceManager.get_ops_trace_instance", return_value=None), + patch("core.ops.ops_trace_manager.trace_manager_queue", trace_queue), + ): + manager = trace_queue_manager(app_id="test-app-id") + manager.add_trace_task(trace_task) + + assert trace_queue.empty() + + def test_task_enqueued_when_telemetry_enabled(self, trace_queue_manager_and_task): + """Verify task IS enqueued when enterprise telemetry is enabled. + + When _enterprise_telemetry_enabled=True, the task should be enqueued + regardless of trace_instance state. + """ + trace_queue_manager, trace_task_type, trace_task_name = trace_queue_manager_and_task + + trace_queue = queue.Queue() + + trace_task = trace_task_type(trace_type=trace_task_name.WORKFLOW_TRACE) + + with ( + patch("core.telemetry.gateway.is_enterprise_telemetry_enabled", return_value=True), + patch("core.ops.ops_trace_manager.OpsTraceManager.get_ops_trace_instance", return_value=None), + patch("core.ops.ops_trace_manager.trace_manager_queue", trace_queue), + ): + manager = trace_queue_manager(app_id="test-app-id") + manager.add_trace_task(trace_task) + + called_task = trace_queue.get_nowait() + assert called_task is trace_task + assert trace_queue.empty() + assert called_task.app_id == "test-app-id" + + def test_task_enqueued_when_trace_instance_configured(self, trace_queue_manager_and_task): + """Verify task IS enqueued when third-party trace instance is configured. + + When trace_instance is not None (e.g., Langfuse configured), the task + should be enqueued even if enterprise telemetry is disabled. + """ + trace_queue_manager, trace_task_type, trace_task_name = trace_queue_manager_and_task + + trace_queue = queue.Queue() + + mock_trace_instance = MagicMock() + + trace_task = trace_task_type(trace_type=trace_task_name.WORKFLOW_TRACE) + + with ( + patch("core.telemetry.gateway.is_enterprise_telemetry_enabled", return_value=False), + patch( + "core.ops.ops_trace_manager.OpsTraceManager.get_ops_trace_instance", + return_value=mock_trace_instance, + ), + patch("core.ops.ops_trace_manager.trace_manager_queue", trace_queue), + ): + manager = trace_queue_manager(app_id="test-app-id") + manager.add_trace_task(trace_task) + + called_task = trace_queue.get_nowait() + assert called_task is trace_task + assert trace_queue.empty() + assert called_task.app_id == "test-app-id" + + def test_task_enqueued_when_both_telemetry_and_trace_instance_enabled(self, trace_queue_manager_and_task): + """Verify task IS enqueued when both telemetry and trace instance are enabled. + + When both _enterprise_telemetry_enabled=True AND trace_instance is set, + the task should definitely be enqueued. + """ + trace_queue_manager, trace_task_type, trace_task_name = trace_queue_manager_and_task + + trace_queue = queue.Queue() + + mock_trace_instance = MagicMock() + + trace_task = trace_task_type(trace_type=trace_task_name.WORKFLOW_TRACE) + + with ( + patch("core.telemetry.gateway.is_enterprise_telemetry_enabled", return_value=True), + patch( + "core.ops.ops_trace_manager.OpsTraceManager.get_ops_trace_instance", + return_value=mock_trace_instance, + ), + patch("core.ops.ops_trace_manager.trace_manager_queue", trace_queue), + ): + manager = trace_queue_manager(app_id="test-app-id") + manager.add_trace_task(trace_task) + + called_task = trace_queue.get_nowait() + assert called_task is trace_task + assert trace_queue.empty() + assert called_task.app_id == "test-app-id" + + def test_app_id_set_before_enqueue(self, trace_queue_manager_and_task): + """Verify app_id is set on the task before enqueuing. + + The guard logic sets trace_task.app_id = self.app_id before calling + trace_manager_queue.put(trace_task). This test verifies that behavior. + """ + trace_queue_manager, trace_task_type, trace_task_name = trace_queue_manager_and_task + + trace_queue = queue.Queue() + + trace_task = trace_task_type(trace_type=trace_task_name.WORKFLOW_TRACE) + + with ( + patch("core.telemetry.gateway.is_enterprise_telemetry_enabled", return_value=True), + patch("core.ops.ops_trace_manager.OpsTraceManager.get_ops_trace_instance", return_value=None), + patch("core.ops.ops_trace_manager.trace_manager_queue", trace_queue), + ): + manager = trace_queue_manager(app_id="expected-app-id") + manager.add_trace_task(trace_task) + + called_task = trace_queue.get_nowait() + assert called_task is trace_task + assert trace_queue.empty() + assert called_task.app_id == "expected-app-id" + + @staticmethod + @pytest.mark.parametrize("enabled_at_start", [False, True]) + def test_application_pipeline_snapshots_switch_without_constructing_sdk( + trace_queue_manager_and_task, monkeypatch: pytest.MonkeyPatch, enabled_at_start: bool + ): + apply_config_overrides(monkeypatch, OPS_AGENT_TRACE_ENABLED=True) + trace_queue_manager, trace_task_class, trace_task_name = trace_queue_manager_and_task + trace_queue = queue.Queue() + calls: list[object] = [] + current_enabled = enabled_at_start + + def queue_enabled(app_id: object) -> bool: + calls.append(app_id) + return current_enabled + + monkeypatch.setattr(RuntimeApplicationTracingHooks, "queue_enabled_at_start", staticmethod(queue_enabled)) + hooks = RuntimeApplicationTracingHooks() + app = Flask(__name__) + bind_application_tracing_hooks(app, hooks) + with ( + app.app_context(), + patch("core.telemetry.gateway.is_enterprise_telemetry_enabled", return_value=False), + patch("core.ops.ops_trace_manager.OpsTraceManager.get_ops_trace_instance") as legacy, + patch("core.ops.ops_trace_manager.trace_manager_queue", trace_queue), + ): + manager = trace_queue_manager(app_id="app-id") + current_enabled = not enabled_at_start + task = trace_task_class(trace_type=trace_task_name.WORKFLOW_TRACE) + manager.add_trace_task(task) + legacy.assert_not_called() + assert calls == ["app-id"] + assert manager.trace_instance is None + assert trace_queue.qsize() == int(enabled_at_start) -def test_workflow_trace_reads_real_workflow_log_from_owned_session( - monkeypatch: pytest.MonkeyPatch, - trace_environment: None, - database: Session, - sqlite_session_factory: sessionmaker[Session], -) -> None: - app = _app(database) - log = WorkflowAppLog( - tenant_id=app.tenant_id, - app_id=app.id, - workflow_id="workflow-1", - workflow_run_id="run-1", - created_from=WorkflowAppLogCreatedFrom.WEB_APP, - created_by_role=CreatorUserRole.ACCOUNT, - created_by="user-1", - ) - workflow_run = WorkflowRun( - id="run-1", - tenant_id=app.tenant_id, - app_id=app.id, - workflow_id="workflow-1", - type=WorkflowType.WORKFLOW, - triggered_from=WorkflowRunTriggeredFrom.APP_RUN, - version="3", - graph="{}", - inputs=json.dumps({"query": "search"}), - status=WorkflowExecutionStatus.SUCCEEDED, - outputs=json.dumps({"out": "value"}), - error=None, - elapsed_time=10, - total_tokens=12, - total_steps=1, - created_by_role=CreatorUserRole.ACCOUNT, - created_by="user-1", - created_at=datetime(2025, 2, 20, 10, 0, 0), - finished_at=datetime(2025, 2, 20, 10, 0, 5), - ) - database.add_all([log, workflow_run]) - database.commit() - repo = DifyAPISQLAlchemyWorkflowRunRepository(sqlite_session_factory) - monkeypatch.setattr(TraceTask, "_get_workflow_run_repo", classmethod(lambda cls: repo)) - monkeypatch.setattr(TraceTask, "_calculate_workflow_token_split", classmethod(lambda cls, *_a, **_k: (5, 7))) - result = TraceTask(trace_type=TraceTaskName.WORKFLOW_TRACE).workflow_trace( - workflow_run_id="run-1", conversation_id=None, user_id="user-1" - ) - assert result.workflow_run_id == "run-1" - assert result.workflow_id == "workflow-1" - assert result.workflow_app_log_id == log.id - assert result.prompt_tokens == 5 - assert result.completion_tokens == 7 +"""Nullable provider credentials never reach RSA/base64 as non-string tokens.""" -def test_tool_trace_reads_real_message_file(database: Session) -> None: - app = _app(database) - _, message = _conversation_message(database, app) - file = MessageFile( - message_id=message.id, - type=FileType.DOCUMENT, - transfer_method=FileTransferMethod.REMOTE_URL, - created_by_role=CreatorUserRole.ACCOUNT, - created_by="user-1", - url="tool/file", - ) - thought = MessageAgentThought( - message_id=message.id, - position=1, - created_by_role=CreatorUserRole.ACCOUNT, - created_by="user-1", - tool="tool-a", - tool_meta_str=json.dumps({"tool-a": {"tool_config": {}, "time_cost": 5, "error": "", "tool_parameters": {}}}), - ) - database.add_all([file, thought]) - database.commit() - result = TraceTask(trace_type=TraceTaskName.TOOL_TRACE).tool_trace( - message.id, {"start": 1, "end": 2}, tool_name="tool-a", tool_inputs={}, tool_outputs="result" - ) - assert result.tool_name == "tool-a" - assert result.time_cost == 5 - assert result.message_file_data.id == file.id +@pytest.fixture +def workspace(sqlite_session: Session, monkeypatch: pytest.MonkeyPatch) -> str: + key = RSA.generate(2048) + tenant_id = str(uuid4()) + sqlite_session.add(make_tenant(tenant_id=tenant_id, encrypt_public_key=key.publickey().export_key().decode())) + sqlite_session.commit() + monkeypatch.setattr(model_engine.db, "session", sqlite_session) + monkeypatch.setattr(key_provider_manager, "_provider", RSAKeyProvider()) + monkeypatch.setattr(rsa, "get_decrypt_decoding", lambda _tenant: (key, gmpy2_pkcs10aep_cipher.new(key))) + return tenant_id -def test_node_execution_trace_resolves_real_message_by_conversation_and_run( - trace_environment: None, database: Session -) -> None: - app = _app(database) - conversation, message = _conversation_message(database, app) - result = TraceTask(trace_type=TraceTaskName.NODE_EXECUTION_TRACE).node_execution_trace( - node_execution_data={ - "tenant_id": app.tenant_id, - "app_id": app.id, - "conversation_id": conversation.id, - "workflow_execution_id": message.workflow_run_id, - "workflow_id": "workflow-1", - "node_execution_id": "node-execution-1", - "node_id": "node-1", - "node_type": "llm", - "title": "Node", - "status": "succeeded", - } - ) - assert result.message_id == message.id - assert result.metadata["conversation_id"] == conversation.id +class TestNullableTracingCredentialsCases: + @staticmethod + def test_databricks_switches_auth_with_real_encryption_and_explicit_null(workspace: str) -> None: + pat = OpsTraceManager.encrypt_tracing_config( + workspace, + "databricks", + { + "host": "https://workspace.test", + "experiment_id": "12", + "personal_access_token": "synthetic-pat", + "client_id": None, + "client_secret": None, + }, + ) + assert pat["personal_access_token"] != "synthetic-pat" + decrypted_pat = OpsTraceManager.decrypt_tracing_config(workspace, "databricks", pat) + assert decrypted_pat["personal_access_token"] == "synthetic-pat" + assert decrypted_pat["client_secret"] is None + oauth = OpsTraceManager.encrypt_tracing_config( + workspace, + "databricks", + { + "host": "https://workspace.test", + "experiment_id": "12", + "personal_access_token": None, + "client_id": "synthetic-client", + "client_secret": "synthetic-secret", + }, + pat, + ) + assert oauth["personal_access_token"] is None + assert oauth["client_secret"] != "synthetic-secret" + decrypted_oauth = OpsTraceManager.decrypt_tracing_config(workspace, "databricks", oauth) + assert decrypted_oauth["personal_access_token"] is None + assert decrypted_oauth["client_secret"] == "synthetic-secret" -def test_trace_helpers_and_streaming_metrics(trace_environment: None) -> None: - assert OpsTraceManager.check_trace_config_is_effective({}, "dummy") - assert OpsTraceManager.get_trace_config_project_key({}, "dummy") == "fake-key" - assert OpsTraceManager.get_trace_config_project_url({}, "dummy") == "https://project.fake" - task = TraceTask(trace_type=TraceTaskName.MESSAGE_TRACE, message_id="message-1") - assert task.conversation_trace(foo="bar") == {"foo": "bar"} - assert task._extract_streaming_metrics(Message(message_metadata="invalid")) == {} - assert task.generate_name_trace("conversation", {"start": 1, "end": 2}, tenant_id=None) == {} - generated = task.generate_name_trace( - "conversation", - {"start": 1, "end": 2}, - tenant_id="tenant-1", - generate_conversation_name="name", - inputs="query", - ) - assert generated.outputs == "name" - assert generated.tenant_id == "tenant-1" - assert generated.message_id == "message-1" + @staticmethod + def test_other_optional_secret_model_can_be_cleared_without_losing_validation(workspace: str) -> None: + password = OpsTraceManager.encrypt_tracing_config( + workspace, + "mlflow", + { + "tracking_uri": "https://mlflow.test", + "experiment_id": "0", + "username": "user", + "password": "synthetic-password", + }, + ) + assert OpsTraceManager.decrypt_tracing_config(workspace, "mlflow", password)["password"] == "synthetic-password" + cleared = OpsTraceManager.encrypt_tracing_config(workspace, "mlflow", {**password, "password": None}, password) + assert cleared["password"] is None + assert OpsTraceManager.decrypt_tracing_config(workspace, "mlflow", cleared)["password"] is None - -def test_trace_queue_preserves_source_app_when_another_manager_drains( - monkeypatch: pytest.MonkeyPatch, trace_environment: None -) -> None: - monkeypatch.setattr( - OpsTraceManager, - "get_ops_trace_instance", - classmethod(lambda cls, app_id: app_id == "source-app-id"), - ) - draining_manager = TraceQueueManager(app_id="draining-app-id", user_id="user-1") - source_manager = TraceQueueManager(app_id="source-app-id", user_id="user-2") - task = TraceTask( - trace_type=TraceTaskName.GENERATE_NAME_TRACE, - conversation_id="conversation-1", - timer={"start": 1, "end": 2}, - tenant_id="tenant-1", - generate_conversation_name="name", - inputs="query", - ) - source_manager.add_trace_task(task) - assert draining_manager.collect_tasks() == [task] - - recording_storage = RecordingStorage() - dispatcher = RecordingDispatcher() - monkeypatch.setattr(module.storage, "save", recording_storage.save) - monkeypatch.setattr(module.process_trace_tasks, "apply_async", dispatcher.apply_async) - file_id = UUID("00000000-0000-0000-0000-000000000123") - monkeypatch.setattr(module, "uuid4", lambda: file_id) - source_manager.add_trace_task(task) - draining_manager.run() - - assert len(recording_storage.writes) == 1 - path, data = recording_storage.writes[0] - assert path.endswith(f"source-app-id/{file_id.hex}.json") - assert json.loads(data)["app_id"] == "source-app-id" - assert dispatcher.payloads == [{"file_id": file_id.hex, "app_id": "source-app-id"}] - - -def test_trace_queue_persists_with_caller_supplied_file_id( - monkeypatch: pytest.MonkeyPatch, trace_environment: None -) -> None: - monkeypatch.setattr(OpsTraceManager, "get_ops_trace_instance", classmethod(lambda cls, _app_id: True)) - manager = TraceQueueManager(app_id="app-id", user_id="user-1") - task = TraceTask( - trace_type=TraceTaskName.GENERATE_NAME_TRACE, - conversation_id="conversation-1", - timer={"start": 1, "end": 2}, - tenant_id="tenant-1", - generate_conversation_name="name", - inputs="query", - ) - task.app_id = "app-id" - recording_storage = RecordingStorage() - monkeypatch.setattr(module.storage, "save", recording_storage.save) - - file_info = manager.persist_trace_task(task, file_id="workflow-final-run-1") - - assert file_info == {"file_id": "workflow-final-run-1", "app_id": "app-id"} - path, data = recording_storage.writes[0] - payload = json.loads(data) - assert path == "ops_trace/app-id/workflow-final-run-1.json" - assert UUID(payload["trace_info"]["operation_id"]) - - -def test_trace_queue_persistence_error_propagates(monkeypatch: pytest.MonkeyPatch, trace_environment: None) -> None: - monkeypatch.setattr(OpsTraceManager, "get_ops_trace_instance", classmethod(lambda cls, _app_id: True)) - manager = TraceQueueManager(app_id="app-id", user_id="user-1") - task = TraceTask(trace_type=TraceTaskName.GENERATE_NAME_TRACE) - task.app_id = "app-id" - - def fail_save(_path: str, _data: bytes) -> None: - raise OSError("storage unavailable") - - monkeypatch.setattr(module.storage, "save", fail_save) - - with pytest.raises(OSError, match="storage unavailable"): - manager.persist_trace_task(task, file_id="workflow-final-run-1") - - -def test_trace_queue_enqueue_error_propagates(monkeypatch: pytest.MonkeyPatch, trace_environment: None) -> None: - monkeypatch.setattr(OpsTraceManager, "get_ops_trace_instance", classmethod(lambda cls, _app_id: True)) - manager = TraceQueueManager(app_id="app-id", user_id="user-1") - - def fail_enqueue(*_args: object, **_kwargs: object) -> None: - raise ConnectionError("broker unavailable") - - monkeypatch.setattr(module.process_trace_tasks, "apply_async", fail_enqueue) - - with pytest.raises(ConnectionError, match="broker unavailable"): - manager.enqueue_persisted_trace({"file_id": "workflow-final-run-1", "app_id": "app-id"}) + @staticmethod + def test_non_nullable_secret_is_still_rejected_by_provider_model(workspace: str) -> None: + with pytest.raises(ValidationError, match="secret_key"): + OpsTraceManager.encrypt_tracing_config( + workspace, "langfuse", {"public_key": "synthetic-public", "secret_key": None} + ) + encrypted = OpsTraceManager.encrypt_tracing_config( + workspace, "langfuse", {"public_key": "synthetic-public", "secret_key": "synthetic-secret"} + ) + with pytest.raises(ValidationError, match="secret_key"): + OpsTraceManager.decrypt_tracing_config(workspace, "langfuse", {**encrypted, "secret_key": None}) diff --git a/api/tests/unit_tests/core/ops/test_trace_queue_manager.py b/api/tests/unit_tests/core/ops/test_trace_queue_manager.py deleted file mode 100644 index 51a924127ee..00000000000 --- a/api/tests/unit_tests/core/ops/test_trace_queue_manager.py +++ /dev/null @@ -1,204 +0,0 @@ -"""Unit tests for TraceQueueManager telemetry guard. - -Verifies that TraceQueueManager.add_trace_task() only enqueues tasks when at -least one consumer is active: -- Enterprise telemetry is enabled (_enterprise_telemetry_enabled=True), OR -- A third-party trace instance (Langfuse, etc.) is configured - -When neither is active, tasks are silently dropped to avoid unnecessary work. - -When BOTH are false, tasks are silently dropped (correct behavior). -""" - -import queue -from unittest.mock import MagicMock, patch - -import pytest -from flask import Flask - -from core.ops.base_trace_instance import bind_application_tracing_hooks -from services.tracing.runtime_bridge import RuntimeApplicationTracingHooks -from tests.unit_tests.config_override import apply_config_overrides - - -@pytest.fixture(autouse=True) -def _application_tracing_disabled(monkeypatch: pytest.MonkeyPatch) -> None: - apply_config_overrides(monkeypatch, OPS_AGENT_TRACE_ENABLED=False) - - -@pytest.fixture -def trace_queue_manager_and_task(): - """Exercise the production guard without starting background dispatch.""" - from core.ops.entities.trace_entity import TraceTaskName - from core.ops.ops_trace_manager import TraceQueueManager, TraceTask - - with patch.object(TraceQueueManager, "start_timer"): - yield TraceQueueManager, TraceTask, TraceTaskName - - -class TestTraceQueueManagerTelemetryGuard: - """Test TraceQueueManager's telemetry guard in add_trace_task().""" - - def test_task_not_enqueued_when_telemetry_disabled_and_no_trace_instance(self, trace_queue_manager_and_task): - """Verify task is NOT enqueued when telemetry disabled and no trace instance. - - This is the core guard: when _enterprise_telemetry_enabled=False AND - trace_instance=None, the task should be silently dropped. - """ - trace_queue_manager, trace_task_type, trace_task_name = trace_queue_manager_and_task - - trace_queue = queue.Queue() - - trace_task = trace_task_type(trace_type=trace_task_name.WORKFLOW_TRACE) - - with ( - patch("core.telemetry.gateway.is_enterprise_telemetry_enabled", return_value=False), - patch("core.ops.ops_trace_manager.OpsTraceManager.get_ops_trace_instance", return_value=None), - patch("core.ops.ops_trace_manager.trace_manager_queue", trace_queue), - ): - manager = trace_queue_manager(app_id="test-app-id") - manager.add_trace_task(trace_task) - - assert trace_queue.empty() - - def test_task_enqueued_when_telemetry_enabled(self, trace_queue_manager_and_task): - """Verify task IS enqueued when enterprise telemetry is enabled. - - When _enterprise_telemetry_enabled=True, the task should be enqueued - regardless of trace_instance state. - """ - trace_queue_manager, trace_task_type, trace_task_name = trace_queue_manager_and_task - - trace_queue = queue.Queue() - - trace_task = trace_task_type(trace_type=trace_task_name.WORKFLOW_TRACE) - - with ( - patch("core.telemetry.gateway.is_enterprise_telemetry_enabled", return_value=True), - patch("core.ops.ops_trace_manager.OpsTraceManager.get_ops_trace_instance", return_value=None), - patch("core.ops.ops_trace_manager.trace_manager_queue", trace_queue), - ): - manager = trace_queue_manager(app_id="test-app-id") - manager.add_trace_task(trace_task) - - called_task = trace_queue.get_nowait() - assert called_task is trace_task - assert trace_queue.empty() - assert called_task.app_id == "test-app-id" - - def test_task_enqueued_when_trace_instance_configured(self, trace_queue_manager_and_task): - """Verify task IS enqueued when third-party trace instance is configured. - - When trace_instance is not None (e.g., Langfuse configured), the task - should be enqueued even if enterprise telemetry is disabled. - """ - trace_queue_manager, trace_task_type, trace_task_name = trace_queue_manager_and_task - - trace_queue = queue.Queue() - - mock_trace_instance = MagicMock() - - trace_task = trace_task_type(trace_type=trace_task_name.WORKFLOW_TRACE) - - with ( - patch("core.telemetry.gateway.is_enterprise_telemetry_enabled", return_value=False), - patch( - "core.ops.ops_trace_manager.OpsTraceManager.get_ops_trace_instance", return_value=mock_trace_instance - ), - patch("core.ops.ops_trace_manager.trace_manager_queue", trace_queue), - ): - manager = trace_queue_manager(app_id="test-app-id") - manager.add_trace_task(trace_task) - - called_task = trace_queue.get_nowait() - assert called_task is trace_task - assert trace_queue.empty() - assert called_task.app_id == "test-app-id" - - def test_task_enqueued_when_both_telemetry_and_trace_instance_enabled(self, trace_queue_manager_and_task): - """Verify task IS enqueued when both telemetry and trace instance are enabled. - - When both _enterprise_telemetry_enabled=True AND trace_instance is set, - the task should definitely be enqueued. - """ - trace_queue_manager, trace_task_type, trace_task_name = trace_queue_manager_and_task - - trace_queue = queue.Queue() - - mock_trace_instance = MagicMock() - - trace_task = trace_task_type(trace_type=trace_task_name.WORKFLOW_TRACE) - - with ( - patch("core.telemetry.gateway.is_enterprise_telemetry_enabled", return_value=True), - patch( - "core.ops.ops_trace_manager.OpsTraceManager.get_ops_trace_instance", return_value=mock_trace_instance - ), - patch("core.ops.ops_trace_manager.trace_manager_queue", trace_queue), - ): - manager = trace_queue_manager(app_id="test-app-id") - manager.add_trace_task(trace_task) - - called_task = trace_queue.get_nowait() - assert called_task is trace_task - assert trace_queue.empty() - assert called_task.app_id == "test-app-id" - - def test_app_id_set_before_enqueue(self, trace_queue_manager_and_task): - """Verify app_id is set on the task before enqueuing. - - The guard logic sets trace_task.app_id = self.app_id before calling - trace_manager_queue.put(trace_task). This test verifies that behavior. - """ - trace_queue_manager, trace_task_type, trace_task_name = trace_queue_manager_and_task - - trace_queue = queue.Queue() - - trace_task = trace_task_type(trace_type=trace_task_name.WORKFLOW_TRACE) - - with ( - patch("core.telemetry.gateway.is_enterprise_telemetry_enabled", return_value=True), - patch("core.ops.ops_trace_manager.OpsTraceManager.get_ops_trace_instance", return_value=None), - patch("core.ops.ops_trace_manager.trace_manager_queue", trace_queue), - ): - manager = trace_queue_manager(app_id="expected-app-id") - manager.add_trace_task(trace_task) - - called_task = trace_queue.get_nowait() - assert called_task is trace_task - assert trace_queue.empty() - assert called_task.app_id == "expected-app-id" - - -@pytest.mark.parametrize("enabled_at_start", [False, True]) -def test_application_pipeline_snapshots_switch_without_constructing_sdk( - trace_queue_manager_and_task, monkeypatch: pytest.MonkeyPatch, enabled_at_start: bool -): - apply_config_overrides(monkeypatch, OPS_AGENT_TRACE_ENABLED=True) - trace_queue_manager, trace_task_class, trace_task_name = trace_queue_manager_and_task - trace_queue = queue.Queue() - calls: list[object] = [] - current_enabled = enabled_at_start - - def queue_enabled(app_id: object) -> bool: - calls.append(app_id) - return current_enabled - - monkeypatch.setattr(RuntimeApplicationTracingHooks, "queue_enabled_at_start", staticmethod(queue_enabled)) - hooks = RuntimeApplicationTracingHooks() - app = Flask(__name__) - bind_application_tracing_hooks(app, hooks) - with ( - app.app_context(), - patch("core.telemetry.gateway.is_enterprise_telemetry_enabled", return_value=False), - patch("core.ops.ops_trace_manager.OpsTraceManager.get_ops_trace_instance") as legacy, - patch("core.ops.ops_trace_manager.trace_manager_queue", trace_queue), - ): - manager = trace_queue_manager(app_id="app-id") - current_enabled = not enabled_at_start - task = trace_task_class(trace_type=trace_task_name.WORKFLOW_TRACE) - manager.add_trace_task(task) - legacy.assert_not_called() - assert calls == ["app-id"] - assert manager.trace_instance is None - assert trace_queue.qsize() == int(enabled_at_start) diff --git a/api/tests/unit_tests/core/ops/test_utils.py b/api/tests/unit_tests/core/ops/test_utils.py index 313f345b2d8..b1dbd9c0974 100644 --- a/api/tests/unit_tests/core/ops/test_utils.py +++ b/api/tests/unit_tests/core/ops/test_utils.py @@ -18,7 +18,7 @@ from core.ops.utils import ( ) from models.enums import ConversationFromSource from models.model import Message -from tests.unit_tests.model_factories import make_message +from tests.helpers.unit.model_factories import make_message @pytest.fixture diff --git a/api/tests/unit_tests/core/ops/test_workflow_tool_export_boundary.py b/api/tests/unit_tests/core/ops/test_workflow_tool_export_boundary.py deleted file mode 100644 index 81166c68e6a..00000000000 --- a/api/tests/unit_tests/core/ops/test_workflow_tool_export_boundary.py +++ /dev/null @@ -1,132 +0,0 @@ -"""External traces retain the caller's Tool contract without source workflow internals.""" - -import json -from datetime import datetime - -import pytest -import requests -from sqlalchemy.orm import Session - -from core.ops.entities.trace_entity import WorkflowTraceInfo -from core.ops.unified_trace.trace_builder import CanonicalTraceBuilder, RepositoryWorkflowExecutionLoader -from extensions.ext_database import db -from models import Account -from models.enums import CreatorUserRole -from models.workflow import WorkflowNodeExecutionModel, WorkflowNodeExecutionTriggeredFrom -from tests.unit_tests.services.tracing.test_mlflow_legacy_destination import ( - exported_spans, -) -from tests.unit_tests.services.tracing.test_mlflow_legacy_destination import ( - transport as transport, # noqa: PLC0414 - Reuse the real SDK's isolated requests boundary. -) - - -@pytest.mark.parametrize("exporter", ["canonical", "mlflow"]) -def test_external_workflow_trace_omits_source_internals( - exporter: str, - sqlite_session: Session, - monkeypatch: pytest.MonkeyPatch, - transport: list[requests.PreparedRequest], -) -> None: - started_at = datetime(2025, 1, 1) - public_inputs = {"query": "public tool input"} - public_outputs = {"answer": "public tool output"} - for index, (execution_id, app_id, origin) in enumerate( - [ - ("outer-tool", "caller-app", WorkflowNodeExecutionTriggeredFrom.WORKFLOW_RUN), - ("source-llm", "source-app", WorkflowNodeExecutionTriggeredFrom.WORKFLOW_TOOL), - ("recursive-llm", "caller-app", WorkflowNodeExecutionTriggeredFrom.WORKFLOW_TOOL), - ] - ): - is_outer = execution_id == "outer-tool" - sqlite_session.add( - WorkflowNodeExecutionModel( - id=execution_id, - tenant_id="tenant-1", - app_id=app_id, - workflow_id=f"{app_id}-workflow", - workflow_run_id="run-1", - triggered_from_workflow_id=None if is_outer else "caller-app-workflow", - triggered_from_node_execution_id=None if is_outer else "outer-tool", - triggered_from=origin, - index=index, - node_execution_id=execution_id, - node_id=execution_id, - node_type="tool" if is_outer else "llm", - title=execution_id, - inputs=json.dumps(public_inputs if is_outer else {"secret": "private source input"}), - outputs=json.dumps(public_outputs if is_outer else {"text": "private source output"}), - process_data=json.dumps( - {} - if is_outer - else { - "prompts": [{"role": "system", "text": "private source prompt"}], - "workflow_tool_root_app_id": "caller-app", - } - ), - status="succeeded", - elapsed_time=1, - created_at=started_at, - created_by_role=CreatorUserRole.ACCOUNT, - created_by="account-1", - ) - ) - sqlite_session.commit() - info = WorkflowTraceInfo( - tenant_id="tenant-1", - workflow_id="caller-workflow", - workflow_run_id="run-1", - workflow_run_elapsed_time=1, - workflow_run_status="succeeded", - workflow_run_inputs={}, - workflow_run_outputs={}, - workflow_run_version="1", - total_tokens=0, - file_list=[], - query="", - metadata={"app_id": "caller-app"}, - start_time=started_at, - end_time=started_at, - ) - monkeypatch.setattr(type(db), "engine", property(lambda _: sqlite_session.get_bind())) - monkeypatch.setattr(db, "session", sqlite_session) - if exporter == "canonical": - account = Account(name="Caller maintainer", email="caller@example.com") - account.id = "account-1" - trace = CanonicalTraceBuilder(RepositoryWorkflowExecutionLoader(lambda _: account)).build(info) - assert trace is not None - assert len(trace.spans) == 2 - payload = trace.model_dump_json() - tool = next(span for span in trace.spans if span.id == "outer-tool") - assert tool.inputs == public_inputs - assert tool.outputs == public_outputs - else: - from dify_trace_mlflow.config import MLflowConfig - from dify_trace_mlflow.mlflow_trace import MLflowDataTrace - from mlflow.entities import Span - from mlflow.tracing.trace_manager import InMemoryTraceManager - - monkeypatch.setattr(InMemoryTraceManager, "_instance", InMemoryTraceManager()) - # OSS owns a private synchronous SDK pipeline. Inspect its real wire spans - # instead of replacing the unrelated process-wide MLflow tracer. - instance = MLflowDataTrace(MLflowConfig(tracking_uri="https://workflow-boundary.example", experiment_id="0")) - try: - instance.trace(info) - spans = [Span.from_otel_proto(span) for span in exported_spans(transport)] - assert len(spans) == 2 - payload = json.dumps([span.to_dict() for span in spans]) - tool = next(span for span in spans if span.name == "outer-tool") - assert tool.inputs == public_inputs - assert tool.outputs == public_outputs - assert transport - assert {request.url for request in transport} <= { - "https://workflow-boundary.example/version", - "https://workflow-boundary.example/api/3.0/mlflow/traces", - "https://workflow-boundary.example/v1/traces", - } - finally: - instance._provider.shutdown() - assert "private source" not in payload - assert "source-llm" not in payload - assert "recursive-llm" not in payload - assert "workflow_tool_root_app_id" not in payload diff --git a/api/tests/unit_tests/core/ops/unified_trace/test_trace_builder.py b/api/tests/unit_tests/core/ops/unified_trace/test_trace_builder.py index 3d1d0ee6325..799fdb4be33 100644 --- a/api/tests/unit_tests/core/ops/unified_trace/test_trace_builder.py +++ b/api/tests/unit_tests/core/ops/unified_trace/test_trace_builder.py @@ -1,3 +1,5 @@ +import importlib as _test_importlib +import json from copy import deepcopy from datetime import datetime, timedelta from types import SimpleNamespace @@ -6,6 +8,8 @@ from unittest.mock import MagicMock from uuid import UUID import pytest +import requests +from sqlalchemy.orm import Session from core.ops.entities.trace_entity import ( BaseTraceInfo, @@ -24,6 +28,140 @@ from core.ops.unified_trace.trace_builder import ( WorkflowExecutionLike, resolve_session_id, ) +from extensions.ext_database import db +from models import Account +from models.enums import CreatorUserRole +from models.workflow import WorkflowNodeExecutionModel, WorkflowNodeExecutionTriggeredFrom + +"""External traces retain the caller's Tool contract without source workflow internals.""" + + +exported_spans = _test_importlib.import_module( + "tests.unit_tests.providers.trace.trace-mlflow.src.dify_trace_mlflow.test_mlflow_trace" +).exported_spans + + +class TestWorkflowToolExportBoundaryCases: + import importlib as _test_importlib + + transport = _test_importlib.import_module( + "tests.unit_tests.providers.trace.trace-mlflow.src.dify_trace_mlflow.test_mlflow_trace" + ).transport + + transport = staticmethod(transport) + + @staticmethod + @pytest.mark.parametrize("exporter", ["canonical", "mlflow"]) + def test_external_workflow_trace_omits_source_internals( + exporter: str, + sqlite_session: Session, + monkeypatch: pytest.MonkeyPatch, + transport: list[requests.PreparedRequest], + ) -> None: + started_at = datetime(2025, 1, 1) + public_inputs = {"query": "public tool input"} + public_outputs = {"answer": "public tool output"} + for index, (execution_id, app_id, origin) in enumerate( + [ + ("outer-tool", "caller-app", WorkflowNodeExecutionTriggeredFrom.WORKFLOW_RUN), + ("source-llm", "source-app", WorkflowNodeExecutionTriggeredFrom.WORKFLOW_TOOL), + ("recursive-llm", "caller-app", WorkflowNodeExecutionTriggeredFrom.WORKFLOW_TOOL), + ] + ): + is_outer = execution_id == "outer-tool" + sqlite_session.add( + WorkflowNodeExecutionModel( + id=execution_id, + tenant_id="tenant-1", + app_id=app_id, + workflow_id=f"{app_id}-workflow", + workflow_run_id="run-1", + triggered_from_workflow_id=None if is_outer else "caller-app-workflow", + triggered_from_node_execution_id=None if is_outer else "outer-tool", + triggered_from=origin, + index=index, + node_execution_id=execution_id, + node_id=execution_id, + node_type="tool" if is_outer else "llm", + title=execution_id, + inputs=json.dumps(public_inputs if is_outer else {"secret": "private source input"}), + outputs=json.dumps(public_outputs if is_outer else {"text": "private source output"}), + process_data=json.dumps( + {} + if is_outer + else { + "prompts": [{"role": "system", "text": "private source prompt"}], + "workflow_tool_root_app_id": "caller-app", + } + ), + status="succeeded", + elapsed_time=1, + created_at=started_at, + created_by_role=CreatorUserRole.ACCOUNT, + created_by="account-1", + ) + ) + sqlite_session.commit() + info = WorkflowTraceInfo( + tenant_id="tenant-1", + workflow_id="caller-workflow", + workflow_run_id="run-1", + workflow_run_elapsed_time=1, + workflow_run_status="succeeded", + workflow_run_inputs={}, + workflow_run_outputs={}, + workflow_run_version="1", + total_tokens=0, + file_list=[], + query="", + metadata={"app_id": "caller-app"}, + start_time=started_at, + end_time=started_at, + ) + monkeypatch.setattr(type(db), "engine", property(lambda _: sqlite_session.get_bind())) + monkeypatch.setattr(db, "session", sqlite_session) + if exporter == "canonical": + account = Account(name="Caller maintainer", email="caller@example.com") + account.id = "account-1" + trace = CanonicalTraceBuilder(RepositoryWorkflowExecutionLoader(lambda _: account)).build(info) + assert trace is not None + assert len(trace.spans) == 2 + payload = trace.model_dump_json() + tool = next(span for span in trace.spans if span.id == "outer-tool") + assert tool.inputs == public_inputs + assert tool.outputs == public_outputs + else: + from dify_trace_mlflow.config import MLflowConfig + from dify_trace_mlflow.mlflow_trace import MLflowDataTrace + from mlflow.entities import Span + from mlflow.tracing.trace_manager import InMemoryTraceManager + + monkeypatch.setattr(InMemoryTraceManager, "_instance", InMemoryTraceManager()) + # OSS owns a private synchronous SDK pipeline. Inspect its real wire spans + # instead of replacing the unrelated process-wide MLflow tracer. + instance = MLflowDataTrace( + MLflowConfig(tracking_uri="https://workflow-boundary.example", experiment_id="0") + ) + try: + instance.trace(info) + spans = [Span.from_otel_proto(span) for span in exported_spans(transport)] + assert len(spans) == 2 + payload = json.dumps([span.to_dict() for span in spans]) + tool = next(span for span in spans if span.name == "outer-tool") + assert tool.inputs == public_inputs + assert tool.outputs == public_outputs + assert transport + assert {request.url for request in transport} <= { + "https://workflow-boundary.example/version", + "https://workflow-boundary.example/api/3.0/mlflow/traces", + "https://workflow-boundary.example/v1/traces", + } + finally: + instance._provider.shutdown() + assert "private source" not in payload + assert "source-llm" not in payload + assert "recursive-llm" not in payload + assert "workflow_tool_root_app_id" not in payload def make_workflow_trace_info(**overrides: object) -> WorkflowTraceInfo: @@ -47,47 +185,6 @@ def make_workflow_trace_info(**overrides: object) -> WorkflowTraceInfo: return WorkflowTraceInfo.model_validate(values) -def test_custom_session_id_wins_over_conversation_id() -> None: - info = make_workflow_trace_info( - conversation_id="conversation-1", - metadata={"trace_session_id": "customer-session"}, - ) - - assert resolve_session_id(info) == "customer-session" - - -def test_workflow_session_falls_back_to_conversation_then_run() -> None: - assert resolve_session_id(make_workflow_trace_info(conversation_id="conversation-1")) == "conversation-1" - assert resolve_session_id(make_workflow_trace_info()) == "run-1" - - -def test_nested_workflow_session_falls_back_to_parent_workflow() -> None: - info = make_workflow_trace_info( - metadata={ - "parent_trace_context": { - "parent_workflow_run_id": "parent-run", - "parent_node_execution_id": "parent-node-execution", - } - } - ) - - assert resolve_session_id(info) == "parent-run" - - -def test_message_session_falls_back_to_message_conversation() -> None: - info = MessageTraceInfo( - conversation_model="chat", - message_tokens=0, - answer_tokens=0, - total_tokens=0, - conversation_mode="chat", - message_data=SimpleNamespace(conversation_id="conversation-1", created_at=datetime(2025, 1, 1)), - metadata={}, - ) - - assert resolve_session_id(info) == "conversation-1" - - def node_execution(**overrides: object) -> WorkflowExecutionLike: values: dict[str, object] = { "id": "node-exec", @@ -140,99 +237,6 @@ CHATFLOW_INPUTS: dict[str, object] = { } -def test_chatflow_message_uses_query_while_workflow_keeps_complete_inputs() -> None: - builder = CanonicalTraceBuilder(lambda _info: []) - - trace = builder.build( - workflow_info( - message_id="message-1", - query="hi", - workflow_run_inputs=CHATFLOW_INPUTS, - ) - ) - - assert trace is not None - spans = {span.id: span for span in trace.spans} - assert trace.root_span_id == "message-1" - assert [span.id for span in trace.spans[:2]] == ["message-1", "run-1"] - assert spans["message-1"].name == "chatflow_run-1" - assert spans["message-1"].inputs == "hi" - assert spans["message-1"].metadata["trace_entity_type"] == "message" - assert spans["message-1"].publishes_parent_context is True - assert spans["run-1"].name == "workflow_run-1" - assert spans["run-1"].parent_id == "message-1" - assert spans["run-1"].inputs == CHATFLOW_INPUTS - - -def test_chatflow_message_falls_back_to_complete_inputs_for_empty_query() -> None: - builder = CanonicalTraceBuilder(lambda _info: []) - - trace = builder.build( - workflow_info( - message_id="message-1", - query="", - workflow_run_inputs=CHATFLOW_INPUTS, - ) - ) - - assert trace is not None - assert trace.spans[0].inputs == CHATFLOW_INPUTS - - -def test_workflow_without_message_keeps_complete_inputs_on_root() -> None: - builder = CanonicalTraceBuilder(lambda _info: []) - - trace = builder.build(workflow_info(query="hi", workflow_run_inputs=CHATFLOW_INPUTS)) - - assert trace is not None - assert trace.root_span_id == "run-1" - assert trace.spans[0].inputs == CHATFLOW_INPUTS - - -def test_build_workflow_trace_is_parent_first_and_uses_wrappers() -> None: - container = node_execution(id="iteration-exec", node_id="iteration", node_type="iteration", title="Items") - child = node_execution( - id="llm-exec", - node_id="llm", - node_type="llm", - title="Summarize", - iteration_id="iteration", - iteration_index=0, - process_data={ - "prompts": [{"role": "user", "text": "hello"}], - "model_provider": "openai", - "model_name": "gpt-4", - "usage": {"prompt_tokens": 10, "completion_tokens": 20, "total_tokens": 30}, - }, - ) - loader = MagicMock(return_value=[child, container]) - builder = CanonicalTraceBuilder(loader) - - trace = builder.build(workflow_info()) - - assert trace is not None - assert [span.id for span in trace.spans] == [ - "run-1", - "iteration-exec", - "iteration:iteration-exec:0", - "llm-exec", - ] - spans = {span.id: span for span in trace.spans} - assert spans["llm-exec"].parent_id == "iteration:iteration-exec:0" - assert spans["llm-exec"].kind is CanonicalSpanKind.LLM - assert spans["llm-exec"].metadata["total_tokens"] == 30 - loader.assert_called_once() - - -def test_workflow_tool_is_marked_as_nested_workflow_parent() -> None: - builder = CanonicalTraceBuilder(lambda _info: [node_execution(id="tool-exec")]) - - trace = builder.build(workflow_info()) - - assert trace is not None - assert trace.spans[-1].can_parent_workflow is True - - def retry_attempt(retry_index: object, **overrides: object) -> dict[str, object]: values: dict[str, object] = { "retry_index": retry_index, @@ -249,225 +253,6 @@ def retry_attempt(retry_index: object, **overrides: object) -> dict[str, object] return values -def test_node_metadata_contains_compact_retry_summary_and_skips_malformed_entries() -> None: - history = [ - retry_attempt(1), - "malformed", - retry_attempt(True), - retry_attempt(2, error="attempt 2 timed out", elapsed_time=2.5), - retry_attempt(3), - ] - builder = CanonicalTraceBuilder(lambda _info: [node_execution(process_data={"__dify_retry_history": history})]) - - trace = builder.build(workflow_info()) - - assert trace is not None - metadata = trace.spans[-1].metadata - assert metadata["retry_count"] == 3 - assert metadata["retry_attempts"] == [ - { - "retry_index": 1, - "error": "attempt 1 failed", - "elapsed_time": 1.0, - "created_at": 1_700_000_000, - "finished_at": 1_700_000_001, - }, - { - "retry_index": 2, - "error": "attempt 2 timed out", - "elapsed_time": 2.5, - "created_at": 1_700_000_000, - "finished_at": 1_700_000_001, - }, - { - "retry_index": 3, - "error": "attempt 3 failed", - "elapsed_time": 3.0, - "created_at": 1_700_000_000, - "finished_at": 1_700_000_001, - }, - ] - assert all("inputs" not in attempt for attempt in metadata["retry_attempts"]) - assert all("process_data" not in attempt for attempt in metadata["retry_attempts"]) - assert all("outputs" not in attempt for attempt in metadata["retry_attempts"]) - assert all("execution_metadata" not in attempt for attempt in metadata["retry_attempts"]) - - -def test_node_without_retry_history_has_no_retry_metadata() -> None: - builder = CanonicalTraceBuilder(lambda _info: [node_execution()]) - - trace = builder.build(workflow_info()) - - assert trace is not None - assert "retry_count" not in trace.spans[-1].metadata - assert "retry_attempts" not in trace.spans[-1].metadata - - -def test_failed_node_preserves_error_without_mutating_trace_metadata() -> None: - metadata = {"app_id": "app-1", "custom": {"nested": True}} - original = deepcopy(metadata) - builder = CanonicalTraceBuilder( - lambda _info: [node_execution(status="failed", error="boom", metadata={"attempt": 1})] - ) - - trace = builder.build(workflow_info(metadata=metadata)) - - assert trace is not None - assert trace.spans[-1].status is CanonicalSpanStatus.ERROR - assert trace.spans[-1].error == "boom" - assert metadata == original - - -def test_nested_workflow_exposes_typed_external_parent() -> None: - builder = CanonicalTraceBuilder(lambda _info: []) - info = workflow_info( - metadata={ - "app_id": "app-1", - "parent_trace_context": { - "parent_workflow_run_id": "outer-run", - "parent_node_execution_id": "outer-tool", - }, - } - ) - - trace = builder.build(info) - - assert trace is not None - assert trace.external_parent is not None - assert trace.external_parent.parent_node_execution_id == "outer-tool" - - -def test_repository_loader_scopes_repository_to_trace(monkeypatch: pytest.MonkeyPatch) -> None: - repository = MagicMock() - repository.get_by_workflow_execution.return_value = [node_execution()] - factory = MagicMock() - factory.create_workflow_node_execution_repository.return_value = repository - monkeypatch.setattr("core.ops.unified_trace.trace_builder.DifyCoreRepositoryFactory", factory) - monkeypatch.setattr("core.ops.unified_trace.trace_builder.db", MagicMock(engine="engine")) - account = MagicMock() - get_account = MagicMock(return_value=account) - loader = RepositoryWorkflowExecutionLoader(get_account) - info = workflow_info() - - result = loader(info) - - assert len(result) == 1 - get_account.assert_called_once_with("app-1") - factory.create_workflow_node_execution_repository.assert_called_once() - call = factory.create_workflow_node_execution_repository.call_args.kwargs - assert call["tenant_id"] == "tenant-1" - assert call["app_id"] == "app-1" - repository.get_by_workflow_execution.assert_called_once_with(workflow_execution_id="run-1") - - -def test_message_trace_does_not_load_workflow_executions() -> None: - loader = MagicMock() - builder = CanonicalTraceBuilder(loader) - info = MessageTraceInfo( - conversation_model="chat", - message_tokens=2, - answer_tokens=3, - total_tokens=5, - conversation_mode="chat", - message_id="message-1", - message_data=SimpleNamespace( - id="message-1", - conversation_id="conversation-1", - answer="hello", - created_at=datetime(2025, 1, 1), - updated_at=datetime(2025, 1, 1, 0, 0, 1), - ), - inputs="hi", - metadata={}, - ) - - trace = builder.build(info) - - assert trace is not None - assert trace.root_span_id == "message-1" - assert [span.name for span in trace.spans] == ["message", "llm"] - assert trace.spans[0].outputs == "hello" - assert trace.spans[0].metadata["trace_entity_type"] == "message" - assert trace.spans[0].publishes_parent_context is True - assert trace.spans[1].parent_id == "message-1" - assert trace.spans[1].kind is CanonicalSpanKind.LLM - assert trace.spans[1].metadata["total_tokens"] == 5 - loader.assert_not_called() - - -def test_generate_name_uses_message_parent_and_conversation_session() -> None: - builder = CanonicalTraceBuilder(lambda _info: []) - info = GenerateNameTraceInfo( - tenant_id="tenant-1", - conversation_id="conversation-1", - message_id="message-1", - inputs="title prompt", - outputs="title", - metadata={}, - ) - - trace = builder.build(info) - - assert trace is not None - assert trace.trace_id == "message-1" - assert trace.session_id == "conversation-1" - assert trace.required_parent_context_id == "message-1" - assert trace.spans[0].parent_id is None - - -@pytest.mark.parametrize( - "info", - [ - ModerationTraceInfo( - message_id="message-1", - message_data=SimpleNamespace(id="message-1"), - flagged=False, - action="direct_output", - preset_response="", - query="hello", - metadata={}, - ), - SuggestedQuestionTraceInfo( - message_id="message-1", - message_data=SimpleNamespace(id="message-1"), - total_tokens=1, - suggested_question=["next"], - level="info", - metadata={}, - ), - DatasetRetrievalTraceInfo( - message_id="message-1", - message_data=SimpleNamespace(id="message-1"), - documents=[], - metadata={}, - ), - ToolTraceInfo( - message_id="message-1", - tool_name="search", - tool_inputs={}, - tool_outputs="done", - tool_config={}, - time_cost=0.1, - tool_parameters={}, - metadata={}, - ), - GenerateNameTraceInfo( - tenant_id="tenant-1", - conversation_id="conversation-1", - message_id="message-1", - metadata={}, - ), - ], -) -def test_standalone_message_child_uses_explicit_required_parent(info: BaseTraceInfo) -> None: - trace = CanonicalTraceBuilder(lambda _info: []).build(info) - - assert trace is not None - assert trace.spans[0].parent_id is None - assert trace.required_parent_context_id == "message-1" - assert_fragment_is_parent_first(trace) - - def assert_fragment_is_parent_first(trace: CanonicalTrace) -> None: seen: set[str] = set() ids = [span.id for span in trace.spans] @@ -481,56 +266,411 @@ def assert_fragment_is_parent_first(trace: CanonicalTrace) -> None: seen.add(span.id) -def test_generate_name_without_message_remains_root_in_conversation_session() -> None: - builder = CanonicalTraceBuilder(lambda _info: []) - info = GenerateNameTraceInfo( - tenant_id="tenant-1", - conversation_id="conversation-1", - inputs="title prompt", - outputs="title", - metadata={}, +class TestTraceBuilderCases: + @staticmethod + def test_custom_session_id_wins_over_conversation_id() -> None: + info = make_workflow_trace_info( + conversation_id="conversation-1", + metadata={"trace_session_id": "customer-session"}, + ) + + assert resolve_session_id(info) == "customer-session" + + @staticmethod + def test_workflow_session_falls_back_to_conversation_then_run() -> None: + assert resolve_session_id(make_workflow_trace_info(conversation_id="conversation-1")) == "conversation-1" + assert resolve_session_id(make_workflow_trace_info()) == "run-1" + + @staticmethod + def test_nested_workflow_session_falls_back_to_parent_workflow() -> None: + info = make_workflow_trace_info( + metadata={ + "parent_trace_context": { + "parent_workflow_run_id": "parent-run", + "parent_node_execution_id": "parent-node-execution", + } + } + ) + + assert resolve_session_id(info) == "parent-run" + + @staticmethod + def test_message_session_falls_back_to_message_conversation() -> None: + info = MessageTraceInfo( + conversation_model="chat", + message_tokens=0, + answer_tokens=0, + total_tokens=0, + conversation_mode="chat", + message_data=SimpleNamespace(conversation_id="conversation-1", created_at=datetime(2025, 1, 1)), + metadata={}, + ) + + assert resolve_session_id(info) == "conversation-1" + + @staticmethod + def test_chatflow_message_uses_query_while_workflow_keeps_complete_inputs() -> None: + builder = CanonicalTraceBuilder(lambda _info: []) + + trace = builder.build( + workflow_info( + message_id="message-1", + query="hi", + workflow_run_inputs=CHATFLOW_INPUTS, + ) + ) + + assert trace is not None + spans = {span.id: span for span in trace.spans} + assert trace.root_span_id == "message-1" + assert [span.id for span in trace.spans[:2]] == ["message-1", "run-1"] + assert spans["message-1"].name == "chatflow_run-1" + assert spans["message-1"].inputs == "hi" + assert spans["message-1"].metadata["trace_entity_type"] == "message" + assert spans["message-1"].publishes_parent_context is True + assert spans["run-1"].name == "workflow_run-1" + assert spans["run-1"].parent_id == "message-1" + assert spans["run-1"].inputs == CHATFLOW_INPUTS + + @staticmethod + def test_chatflow_message_falls_back_to_complete_inputs_for_empty_query() -> None: + builder = CanonicalTraceBuilder(lambda _info: []) + + trace = builder.build( + workflow_info( + message_id="message-1", + query="", + workflow_run_inputs=CHATFLOW_INPUTS, + ) + ) + + assert trace is not None + assert trace.spans[0].inputs == CHATFLOW_INPUTS + + @staticmethod + def test_workflow_without_message_keeps_complete_inputs_on_root() -> None: + builder = CanonicalTraceBuilder(lambda _info: []) + + trace = builder.build(workflow_info(query="hi", workflow_run_inputs=CHATFLOW_INPUTS)) + + assert trace is not None + assert trace.root_span_id == "run-1" + assert trace.spans[0].inputs == CHATFLOW_INPUTS + + @staticmethod + def test_build_workflow_trace_is_parent_first_and_uses_wrappers() -> None: + container = node_execution(id="iteration-exec", node_id="iteration", node_type="iteration", title="Items") + child = node_execution( + id="llm-exec", + node_id="llm", + node_type="llm", + title="Summarize", + iteration_id="iteration", + iteration_index=0, + process_data={ + "prompts": [{"role": "user", "text": "hello"}], + "model_provider": "openai", + "model_name": "gpt-4", + "usage": {"prompt_tokens": 10, "completion_tokens": 20, "total_tokens": 30}, + }, + ) + loader = MagicMock(return_value=[child, container]) + builder = CanonicalTraceBuilder(loader) + + trace = builder.build(workflow_info()) + + assert trace is not None + assert [span.id for span in trace.spans] == [ + "run-1", + "iteration-exec", + "iteration:iteration-exec:0", + "llm-exec", + ] + spans = {span.id: span for span in trace.spans} + assert spans["llm-exec"].parent_id == "iteration:iteration-exec:0" + assert spans["llm-exec"].kind is CanonicalSpanKind.LLM + assert spans["llm-exec"].metadata["total_tokens"] == 30 + loader.assert_called_once() + + @staticmethod + def test_workflow_tool_is_marked_as_nested_workflow_parent() -> None: + builder = CanonicalTraceBuilder(lambda _info: [node_execution(id="tool-exec")]) + + trace = builder.build(workflow_info()) + + assert trace is not None + assert trace.spans[-1].can_parent_workflow is True + + @staticmethod + def test_node_metadata_contains_compact_retry_summary_and_skips_malformed_entries() -> None: + history = [ + retry_attempt(1), + "malformed", + retry_attempt(True), + retry_attempt(2, error="attempt 2 timed out", elapsed_time=2.5), + retry_attempt(3), + ] + builder = CanonicalTraceBuilder(lambda _info: [node_execution(process_data={"__dify_retry_history": history})]) + + trace = builder.build(workflow_info()) + + assert trace is not None + metadata = trace.spans[-1].metadata + assert metadata["retry_count"] == 3 + assert metadata["retry_attempts"] == [ + { + "retry_index": 1, + "error": "attempt 1 failed", + "elapsed_time": 1.0, + "created_at": 1_700_000_000, + "finished_at": 1_700_000_001, + }, + { + "retry_index": 2, + "error": "attempt 2 timed out", + "elapsed_time": 2.5, + "created_at": 1_700_000_000, + "finished_at": 1_700_000_001, + }, + { + "retry_index": 3, + "error": "attempt 3 failed", + "elapsed_time": 3.0, + "created_at": 1_700_000_000, + "finished_at": 1_700_000_001, + }, + ] + assert all("inputs" not in attempt for attempt in metadata["retry_attempts"]) + assert all("process_data" not in attempt for attempt in metadata["retry_attempts"]) + assert all("outputs" not in attempt for attempt in metadata["retry_attempts"]) + assert all("execution_metadata" not in attempt for attempt in metadata["retry_attempts"]) + + @staticmethod + def test_node_without_retry_history_has_no_retry_metadata() -> None: + builder = CanonicalTraceBuilder(lambda _info: [node_execution()]) + + trace = builder.build(workflow_info()) + + assert trace is not None + assert "retry_count" not in trace.spans[-1].metadata + assert "retry_attempts" not in trace.spans[-1].metadata + + @staticmethod + def test_failed_node_preserves_error_without_mutating_trace_metadata() -> None: + metadata = {"app_id": "app-1", "custom": {"nested": True}} + original = deepcopy(metadata) + builder = CanonicalTraceBuilder( + lambda _info: [node_execution(status="failed", error="boom", metadata={"attempt": 1})] + ) + + trace = builder.build(workflow_info(metadata=metadata)) + + assert trace is not None + assert trace.spans[-1].status is CanonicalSpanStatus.ERROR + assert trace.spans[-1].error == "boom" + assert metadata == original + + @staticmethod + def test_nested_workflow_exposes_typed_external_parent() -> None: + builder = CanonicalTraceBuilder(lambda _info: []) + info = workflow_info( + metadata={ + "app_id": "app-1", + "parent_trace_context": { + "parent_workflow_run_id": "outer-run", + "parent_node_execution_id": "outer-tool", + }, + } + ) + + trace = builder.build(info) + + assert trace is not None + assert trace.external_parent is not None + assert trace.external_parent.parent_node_execution_id == "outer-tool" + + @staticmethod + def test_repository_loader_scopes_repository_to_trace(monkeypatch: pytest.MonkeyPatch) -> None: + repository = MagicMock() + repository.get_by_workflow_execution.return_value = [node_execution()] + factory = MagicMock() + factory.create_workflow_node_execution_repository.return_value = repository + monkeypatch.setattr("core.ops.unified_trace.trace_builder.DifyCoreRepositoryFactory", factory) + monkeypatch.setattr("core.ops.unified_trace.trace_builder.db", MagicMock(engine="engine")) + account = MagicMock() + get_account = MagicMock(return_value=account) + loader = RepositoryWorkflowExecutionLoader(get_account) + info = workflow_info() + + result = loader(info) + + assert len(result) == 1 + get_account.assert_called_once_with("app-1") + factory.create_workflow_node_execution_repository.assert_called_once() + call = factory.create_workflow_node_execution_repository.call_args.kwargs + assert call["tenant_id"] == "tenant-1" + assert call["app_id"] == "app-1" + repository.get_by_workflow_execution.assert_called_once_with(workflow_execution_id="run-1") + + @staticmethod + def test_message_trace_does_not_load_workflow_executions() -> None: + loader = MagicMock() + builder = CanonicalTraceBuilder(loader) + info = MessageTraceInfo( + conversation_model="chat", + message_tokens=2, + answer_tokens=3, + total_tokens=5, + conversation_mode="chat", + message_id="message-1", + message_data=SimpleNamespace( + id="message-1", + conversation_id="conversation-1", + answer="hello", + created_at=datetime(2025, 1, 1), + updated_at=datetime(2025, 1, 1, 0, 0, 1), + ), + inputs="hi", + metadata={}, + ) + + trace = builder.build(info) + + assert trace is not None + assert trace.root_span_id == "message-1" + assert [span.name for span in trace.spans] == ["message", "llm"] + assert trace.spans[0].outputs == "hello" + assert trace.spans[0].metadata["trace_entity_type"] == "message" + assert trace.spans[0].publishes_parent_context is True + assert trace.spans[1].parent_id == "message-1" + assert trace.spans[1].kind is CanonicalSpanKind.LLM + assert trace.spans[1].metadata["total_tokens"] == 5 + loader.assert_not_called() + + @staticmethod + def test_generate_name_uses_message_parent_and_conversation_session() -> None: + builder = CanonicalTraceBuilder(lambda _info: []) + info = GenerateNameTraceInfo( + tenant_id="tenant-1", + conversation_id="conversation-1", + message_id="message-1", + inputs="title prompt", + outputs="title", + metadata={}, + ) + + trace = builder.build(info) + + assert trace is not None + assert trace.trace_id == "message-1" + assert trace.session_id == "conversation-1" + assert trace.required_parent_context_id == "message-1" + assert trace.spans[0].parent_id is None + + @staticmethod + @pytest.mark.parametrize( + "info", + [ + ModerationTraceInfo( + message_id="message-1", + message_data=SimpleNamespace(id="message-1"), + flagged=False, + action="direct_output", + preset_response="", + query="hello", + metadata={}, + ), + SuggestedQuestionTraceInfo( + message_id="message-1", + message_data=SimpleNamespace(id="message-1"), + total_tokens=1, + suggested_question=["next"], + level="info", + metadata={}, + ), + DatasetRetrievalTraceInfo( + message_id="message-1", + message_data=SimpleNamespace(id="message-1"), + documents=[], + metadata={}, + ), + ToolTraceInfo( + message_id="message-1", + tool_name="search", + tool_inputs={}, + tool_outputs="done", + tool_config={}, + time_cost=0.1, + tool_parameters={}, + metadata={}, + ), + GenerateNameTraceInfo( + tenant_id="tenant-1", + conversation_id="conversation-1", + message_id="message-1", + metadata={}, + ), + ], ) + def test_standalone_message_child_uses_explicit_required_parent(info: BaseTraceInfo) -> None: + trace = CanonicalTraceBuilder(lambda _info: []).build(info) - trace = builder.build(info) + assert trace is not None + assert trace.spans[0].parent_id is None + assert trace.required_parent_context_id == "message-1" + assert_fragment_is_parent_first(trace) - assert trace is not None - assert trace.session_id == "conversation-1" - assert trace.required_parent_context_id is None - assert trace.spans[0].parent_id is None + @staticmethod + def test_generate_name_without_message_remains_root_in_conversation_session() -> None: + builder = CanonicalTraceBuilder(lambda _info: []) + info = GenerateNameTraceInfo( + tenant_id="tenant-1", + conversation_id="conversation-1", + inputs="title prompt", + outputs="title", + metadata={}, + ) + trace = builder.build(info) -def test_standalone_trace_reuses_persisted_operation_id() -> None: - builder = CanonicalTraceBuilder(lambda _info: []) - info = GenerateNameTraceInfo( - tenant_id="tenant-1", - conversation_id="conversation-1", - inputs="title prompt", - outputs="title", - operation_id="00000000-0000-0000-0000-000000000099", - metadata={}, - ) + assert trace is not None + assert trace.session_id == "conversation-1" + assert trace.required_parent_context_id is None + assert trace.spans[0].parent_id is None - first = builder.build(info) - second = builder.build(info) + @staticmethod + def test_standalone_trace_reuses_persisted_operation_id() -> None: + builder = CanonicalTraceBuilder(lambda _info: []) + info = GenerateNameTraceInfo( + tenant_id="tenant-1", + conversation_id="conversation-1", + inputs="title prompt", + outputs="title", + operation_id="00000000-0000-0000-0000-000000000099", + metadata={}, + ) - assert first is not None - assert second is not None - assert first.root_span_id == "00000000-0000-0000-0000-000000000099" - assert second.root_span_id == first.root_span_id + first = builder.build(info) + second = builder.build(info) + assert first is not None + assert second is not None + assert first.root_span_id == "00000000-0000-0000-0000-000000000099" + assert second.root_span_id == first.root_span_id -def test_standalone_trace_accepts_legacy_payload_without_operation_id() -> None: - info = GenerateNameTraceInfo.model_validate( - { - "tenant_id": "tenant-1", - "conversation_id": "conversation-1", - "inputs": "title prompt", - "outputs": "title", - "metadata": {}, - } - ) + @staticmethod + def test_standalone_trace_accepts_legacy_payload_without_operation_id() -> None: + info = GenerateNameTraceInfo.model_validate( + { + "tenant_id": "tenant-1", + "conversation_id": "conversation-1", + "inputs": "title prompt", + "outputs": "title", + "metadata": {}, + } + ) - trace = CanonicalTraceBuilder(lambda _info: []).build(info) + trace = CanonicalTraceBuilder(lambda _info: []).build(info) - assert trace is not None - assert UUID(trace.root_span_id) + assert trace is not None + assert UUID(trace.root_span_id) diff --git a/api/tests/unit_tests/core/plugin/test_backwards_invocation_app.py b/api/tests/unit_tests/core/plugin/backwards_invocation/test_app.py similarity index 99% rename from api/tests/unit_tests/core/plugin/test_backwards_invocation_app.py rename to api/tests/unit_tests/core/plugin/backwards_invocation/test_app.py index ef9cd5fc0ce..328b025690e 100644 --- a/api/tests/unit_tests/core/plugin/test_backwards_invocation_app.py +++ b/api/tests/unit_tests/core/plugin/backwards_invocation/test_app.py @@ -16,7 +16,7 @@ from models.enums import EndUserType from models.model import App, AppMode, AppModelConfig, EndUser from models.workflow import Workflow from services.workflow_run_agg import WorkflowRunAgg -from tests.unit_tests.model_factories import make_app, make_end_user, make_workflow +from tests.helpers.unit.model_factories import make_app, make_end_user, make_workflow class _Chunk(BaseModel): diff --git a/api/tests/unit_tests/core/plugin/test_backwards_invocation_model.py b/api/tests/unit_tests/core/plugin/backwards_invocation/test_model.py similarity index 68% rename from api/tests/unit_tests/core/plugin/test_backwards_invocation_model.py rename to api/tests/unit_tests/core/plugin/backwards_invocation/test_model.py index 1080f5bbe40..d675d9745ca 100644 --- a/api/tests/unit_tests/core/plugin/test_backwards_invocation_model.py +++ b/api/tests/unit_tests/core/plugin/backwards_invocation/test_model.py @@ -1,17 +1,28 @@ import json +from collections.abc import Callable, Iterator from dataclasses import dataclass, field from unittest.mock import patch import pytest +from pydantic import BaseModel +from configs import dify_config from core.entities.embedding_type import EmbeddingInputType from core.entities.provider_configuration import ProviderConfiguration, ProviderModelBundle -from core.entities.provider_entities import CustomConfiguration, SystemConfiguration +from core.entities.provider_entities import ( + CustomConfiguration, + ProviderQuotaType, + QuotaConfiguration, + QuotaUnit, + SystemConfiguration, +) +from core.model_context import get_credit_usage_metadata from core.model_manager import ModelInstance from core.plugin.backwards_invocation.model import PluginModelBackwardsInvocation from core.plugin.entities.plugin_daemon import TTSAudioChunk from core.plugin.entities.request import ( InvokableModelTokenLimits, + RequestInvokeLLM, RequestInvokeMultimodalEmbedding, RequestInvokeRerank, RequestInvokeSummary, @@ -20,8 +31,10 @@ from core.plugin.entities.request import ( RequestListModels, ) from core.plugin.impl.model_runtime_factory import create_plugin_model_runtime +from enums import DeploymentEdition from graphon.model_runtime.entities.common_entities import I18nObject -from graphon.model_runtime.entities.message_entities import UserPromptMessage +from graphon.model_runtime.entities.llm_entities import LLMResultChunk, LLMResultChunkDelta, LLMUsage +from graphon.model_runtime.entities.message_entities import AssistantPromptMessage, UserPromptMessage from graphon.model_runtime.entities.model_entities import ( AIModelEntity, FetchFrom, @@ -30,10 +43,15 @@ from graphon.model_runtime.entities.model_entities import ( ModelType, ) from graphon.model_runtime.entities.provider_entities import ProviderEntity +from graphon.model_runtime.model_providers.base.large_language_model import LargeLanguageModel +from graphon.model_runtime.model_providers.base.rerank_model import RerankModel +from graphon.model_runtime.model_providers.base.text_embedding_model import TextEmbeddingModel from graphon.model_runtime.model_providers.base.tts_model import TTSModel from graphon.model_runtime.protocols.tts_runtime import TTSChunk from models.account import Tenant from models.provider import ProviderType +from services.billing_service import BillingService +from tests.helpers.unit.knowledge_fs_core_fakes import Marker def _tts_model_instance() -> ModelInstance: @@ -130,7 +148,7 @@ def test_invoke_summary_uses_same_user_scope_for_token_helpers(): ) -def _tenant() -> Tenant: +def _backwards_invocation_model_tenant() -> Tenant: tenant = Tenant(name="Test Workspace") tenant.id = "tenant-1" return tenant @@ -185,7 +203,7 @@ class _ProviderConfigurations: @dataclass(frozen=True) -class _ProviderManager: +class _BackwardsInvocationModelProvidermanager: configurations: _ProviderConfigurations def get_configurations(self, _tenant_id: str) -> _ProviderConfigurations: @@ -221,7 +239,7 @@ class _SchemaModelInstance: def test_invoke_text_embedding_forwards_input_type_to_bound_model_instance(): - tenant = _tenant() + tenant = _backwards_invocation_model_tenant() payload = RequestInvokeTextEmbedding( provider="langgenius/openai/openai", model="text-embedding-3-small", @@ -252,7 +270,7 @@ def test_invoke_text_embedding_forwards_input_type_to_bound_model_instance(): def test_invoke_multimodal_embedding_uses_bound_model_instance(): - tenant = _tenant() + tenant = _backwards_invocation_model_tenant() documents = [{"content": "AQID", "content_type": "image", "file_id": "asset.png"}] payload = RequestInvokeMultimodalEmbedding( provider="langgenius/clip/clip", @@ -287,7 +305,7 @@ def test_invoke_multimodal_embedding_uses_bound_model_instance(): def test_invoke_rerank_accepts_omitted_threshold_and_top_n(): - tenant = _tenant() + tenant = _backwards_invocation_model_tenant() payload = RequestInvokeRerank( provider="langgenius/cohere/cohere", model="rerank-v3.5", @@ -332,7 +350,7 @@ def test_list_models_returns_only_active_tenant_models_with_installed_identity(m status="active", ) configurations = _ProviderConfigurations(models=[active_model]) - provider_manager = _ProviderManager(configurations=configurations) + provider_manager = _BackwardsInvocationModelProvidermanager(configurations=configurations) installed_plugin = _InstalledPlugin( plugin_id="langgenius/openai", plugin_unique_identifier="langgenius/openai:1.0.0@sha256:installed", @@ -519,3 +537,173 @@ def test_backwards_event_stream_serializes_a_deferred_mime_error(): chunks = list(PluginModelBackwardsInvocation.convert_to_event_stream(response)) assert json.loads(chunks[0])["error"].startswith("TTS provider output MIME does not match") + + +"""Exercise backwards invocation through the real hosted quota owner, with no external I/O.""" + + +_MODEL_TYPE_CLASSES: dict[ModelType, type[LargeLanguageModel] | type[TextEmbeddingModel] | type[RerankModel]] = { + ModelType.LLM: LargeLanguageModel, + ModelType.TEXT_EMBEDDING: TextEmbeddingModel, + ModelType.RERANK: RerankModel, +} + + +@dataclass(frozen=True) +class _ModelBillingAttributionProvidermanager: + bundle: ProviderModelBundle + + def get_provider_model_bundle(self, **_kwargs: object) -> ProviderModelBundle: + return self.bundle + + def get_configurations(self, _tenant_id: str) -> dict[str, ProviderConfiguration]: + return {self.bundle.configuration.provider.provider: self.bundle.configuration} + + +def _model_billing_attribution_tenant() -> Tenant: + tenant = Tenant(name="Test Workspace") + tenant.id = "tenant-1" + return tenant + + +def _hosted_bundle(model_type: ModelType) -> ProviderModelBundle: + provider = ProviderEntity( + provider="langgenius/openai/openai", + label=I18nObject(en_US="OpenAI"), + supported_model_types=[model_type], + configurate_methods=[], + ) + configuration = ProviderConfiguration( + tenant_id="tenant-1", + provider=provider, + preferred_provider_type=ProviderType.SYSTEM, + using_provider_type=ProviderType.SYSTEM, + system_configuration=SystemConfiguration( + enabled=True, + current_quota_type=ProviderQuotaType.TRIAL, + quota_configurations=[ + QuotaConfiguration( + quota_type=ProviderQuotaType.TRIAL, + quota_unit=QuotaUnit.CREDITS, + quota_limit=200, + quota_used=0, + is_valid=True, + restrict_models=[], + ) + ], + ), + custom_configuration=CustomConfiguration(provider=None), + model_settings=[], + ) + model_type_instance = _MODEL_TYPE_CLASSES[model_type]( + provider_schema=provider, + model_runtime=create_plugin_model_runtime(tenant_id="tenant-1", user_id="synthetic-end-user"), + ) + return ProviderModelBundle(configuration=configuration, model_type_instance=model_type_instance) + + +@pytest.mark.parametrize( + ("payload_type", "invoke_method", "provider_method", "model_type", "body"), + [ + ( + RequestInvokeLLM, + PluginModelBackwardsInvocation.invoke_llm, + "invoke_llm", + ModelType.LLM, + {"mode": "chat", "stream": True}, + ), + ( + RequestInvokeTextEmbedding, + PluginModelBackwardsInvocation.invoke_text_embedding, + "invoke_text_embedding", + ModelType.TEXT_EMBEDDING, + {"texts": ["query"], "input_type": "query"}, + ), + ( + RequestInvokeMultimodalEmbedding, + PluginModelBackwardsInvocation.invoke_multimodal_embedding, + "invoke_multimodal_embedding", + ModelType.TEXT_EMBEDDING, + {"documents": [{"content": "AQID", "content_type": "image"}]}, + ), + ( + RequestInvokeRerank, + PluginModelBackwardsInvocation.invoke_rerank, + "invoke_rerank", + ModelType.RERANK, + {"query": "query", "docs": ["doc"]}, + ), + ], +) +@pytest.mark.parametrize( + "created_by", + ["knowledge_indexing", "knowledge_retrieval", "app", None], + ids=["indexing", "retrieval", "answer", "legacy-plugin"], +) +def test_model_invocation_preserves_billing_origin( + payload_type: type[BaseModel], + invoke_method: Callable[..., object], + provider_method: str, + model_type: ModelType, + body: dict[str, object], + created_by: str | None, + config_overrides: Callable[..., None], +) -> None: + config_overrides(DEPLOYMENT_EDITION=DeploymentEdition.CLOUD) + provider_manager = _ModelBillingAttributionProvidermanager(bundle=_hosted_bundle(model_type)) + metadata = {"app_type": "new_rag", "created_by": created_by} if created_by else {} + payload = payload_type.model_validate( + {"provider": "langgenius/openai/openai", "model": "selected-model", **body, **metadata} + ) + requests: list[tuple[str, dict[str, object]]] = [] + + def billing_transport(method: str, endpoint: str, *, json: dict[str, object]) -> dict[str, object]: + assert method == "POST" + requests.append((endpoint, json)) + if endpoint == "/quota/reserve": + return {"reservation_id": "reservation-1", "available": 197, "reserved": 3} + assert endpoint == "/quota/commit" + return {"available": 197, "reserved": 0, "refunded": 0} + + chunk = LLMResultChunk( + model="selected-model", + prompt_messages=[], + delta=LLMResultChunkDelta( + index=0, message=AssistantPromptMessage(content="OK"), usage=LLMUsage.empty_usage(), finish_reason="stop" + ), + ) + provider_result = (item for item in [chunk]) if model_type == ModelType.LLM else Marker("provider-result") + with ( + patch("core.model_manager.create_plugin_provider_manager", return_value=provider_manager), + patch("core.app.llm.quota.create_plugin_provider_manager", return_value=provider_manager), + patch.object(type(dify_config), "get_model_credits", return_value=3), + patch.object(ProviderConfiguration, "get_current_credentials", return_value={"api_key": "local-test-only"}), + patch.object(ProviderConfiguration, "get_provider_model", return_value=None), + patch.object(ModelInstance, provider_method, return_value=provider_result), + patch.object(BillingService, "_send_quota_request", side_effect=billing_transport), + ): + response = invoke_method("synthetic-end-user", _model_billing_attribution_tenant(), payload) + assert get_credit_usage_metadata() is None + if model_type == ModelType.LLM: + assert isinstance(response, Iterator) + first_chunk = list(response)[0] + assert isinstance(first_chunk, LLMResultChunk) + assert first_chunk.delta.message.content == "OK" + + assert [endpoint for endpoint, _ in requests] == ["/quota/reserve", "/quota/commit"] + expected_metadata = { + "source": "llm.invoke" if model_type == ModelType.LLM else "model.invoke", + "provider": "langgenius/openai/openai", + "model": "selected-model", + "app_type": "new_rag" if created_by else "unknown", + "created_by": created_by or "plugin_api", + } + if model_type != ModelType.LLM: + expected_metadata["model_type"] = model_type.value + reserve, commit = requests[0][1], requests[1][1] + assert reserve["tenant_id"] == commit["tenant_id"] == "tenant-1" + assert reserve["bucket"] == commit["bucket"] == "trial" + assert reserve["feature_key"] == commit["feature_key"] == "credit_pool" + assert reserve["amount"] == commit["actual_amount"] == 3 + assert reserve["meta"] == expected_metadata + assert commit["meta"] == {**expected_metadata, "request_id": reserve["request_id"]} diff --git a/api/tests/unit_tests/core/plugin/test_plugin_entities.py b/api/tests/unit_tests/core/plugin/entities/test_plugin_daemon.py similarity index 100% rename from api/tests/unit_tests/core/plugin/test_plugin_entities.py rename to api/tests/unit_tests/core/plugin/entities/test_plugin_daemon.py diff --git a/api/tests/unit_tests/core/plugin/entities/test_request.py b/api/tests/unit_tests/core/plugin/entities/test_request.py index 57bc7c9aae8..d1e9e850a56 100644 --- a/api/tests/unit_tests/core/plugin/entities/test_request.py +++ b/api/tests/unit_tests/core/plugin/entities/test_request.py @@ -1,7 +1,10 @@ +import binascii +from collections.abc import Mapping + import pytest from pydantic import ValidationError -from core.plugin.entities.request import RequestRequestDownloadFile +from core.plugin.entities.request import RequestRequestDownloadFile, TriggerDispatchResponse from core.workflow.file_reference import build_file_reference @@ -144,3 +147,20 @@ def test_request_download_file_rejects_non_remote_with_url(transfer_method: str) }, } ) + + +def test_trigger_dispatch_response() -> None: + raw_http_response = b'HTTP/1.1 200 OK\r\nContent-Type: application/json\r\n\r\n{"message": "Hello, world!"}' + + data: Mapping[str, object] = { + "user_id": "123", + "events": ["event1", "event2"], + "response": binascii.hexlify(raw_http_response).decode(), + "payload": {"key": "value"}, + } + + response = TriggerDispatchResponse.model_validate(data) + + assert response.response.status_code == 200 + assert response.response.headers["Content-Type"] == "application/json" + assert response.response.get_data(as_text=True) == '{"message": "Hello, world!"}' diff --git a/api/tests/unit_tests/core/plugin/impl/test_agent_client.py b/api/tests/unit_tests/core/plugin/impl/test_agent.py similarity index 100% rename from api/tests/unit_tests/core/plugin/impl/test_agent_client.py rename to api/tests/unit_tests/core/plugin/impl/test_agent.py diff --git a/api/tests/unit_tests/core/plugin/impl/test_asset_manager.py b/api/tests/unit_tests/core/plugin/impl/test_asset.py similarity index 100% rename from api/tests/unit_tests/core/plugin/impl/test_asset_manager.py rename to api/tests/unit_tests/core/plugin/impl/test_asset.py diff --git a/api/tests/unit_tests/core/plugin/impl/test_base_client_impl.py b/api/tests/unit_tests/core/plugin/impl/test_base.py similarity index 100% rename from api/tests/unit_tests/core/plugin/impl/test_base_client_impl.py rename to api/tests/unit_tests/core/plugin/impl/test_base.py diff --git a/api/tests/unit_tests/core/plugin/impl/test_datasource_manager.py b/api/tests/unit_tests/core/plugin/impl/test_datasource.py similarity index 100% rename from api/tests/unit_tests/core/plugin/impl/test_datasource_manager.py rename to api/tests/unit_tests/core/plugin/impl/test_datasource.py diff --git a/api/tests/unit_tests/core/plugin/impl/test_debugging_client.py b/api/tests/unit_tests/core/plugin/impl/test_debugging.py similarity index 100% rename from api/tests/unit_tests/core/plugin/impl/test_debugging_client.py rename to api/tests/unit_tests/core/plugin/impl/test_debugging.py diff --git a/api/tests/unit_tests/core/plugin/impl/test_endpoint.py b/api/tests/unit_tests/core/plugin/impl/test_endpoint.py new file mode 100644 index 00000000000..543af9501a4 --- /dev/null +++ b/api/tests/unit_tests/core/plugin/impl/test_endpoint.py @@ -0,0 +1,417 @@ +from collections.abc import Callable +from unittest.mock import patch + +import httpx +import pytest +from pytest_mock import MockerFixture + +from core.plugin.impl.endpoint import PluginEndpointClient +from core.plugin.impl.exc import PluginDaemonInternalServerError + + +class TestEndpointClientImplCases: + class TestPluginEndpointClientImpl: + @pytest.mark.parametrize( + "mutate", + [ + pytest.param( + lambda client: client.create_endpoint("tenant-1", "user-1", "org/plugin:1", "endpoint-a", {}), + id="create", + ), + pytest.param( + lambda client: client.update_endpoint("tenant-1", "user-1", "endpoint-1", "renamed", {}), + id="update", + ), + pytest.param(lambda client: client.delete_endpoint("tenant-1", "user-1", "endpoint-1"), id="delete"), + pytest.param(lambda client: client.enable_endpoint("tenant-1", "user-1", "endpoint-1"), id="enable"), + pytest.param(lambda client: client.disable_endpoint("tenant-1", "user-1", "endpoint-1"), id="disable"), + ], + ) + def test_mutation_rejects_false_daemon_result( + self, mocker: MockerFixture, mutate: Callable[[PluginEndpointClient], bool] + ) -> None: + client = PluginEndpointClient() + mocker.patch.object( + client, + "_request", + return_value=httpx.Response( + 200, + request=httpx.Request("POST", "https://daemon.test/endpoint"), + json={"code": 0, "message": "success", "data": False}, + ), + ) + + with pytest.raises(PluginDaemonInternalServerError, match="Plugin daemon failed to"): + mutate(client) + + def test_create_endpoint(self, mocker: MockerFixture): + client = PluginEndpointClient() + request_mock = mocker.patch.object(client, "_request_with_plugin_daemon_response", return_value=True) + + result = client.create_endpoint("tenant-1", "user-1", "org/plugin:1", "endpoint-a", {"k": "v"}) + + assert result + assert request_mock.call_count == 1 + args = request_mock.call_args.args + kwargs = request_mock.call_args.kwargs + assert args[:3] == ("POST", "plugin/tenant-1/endpoint/setup", bool) + assert kwargs["data"]["plugin_unique_identifier"] == "org/plugin:1" + + def test_list_endpoints(self, mocker: MockerFixture): + client = PluginEndpointClient() + request_mock = mocker.patch.object( + client, "_request_with_plugin_daemon_response", return_value=["endpoint"] + ) + + result = client.list_endpoints("tenant-1", "user-1", 2, 20) + + assert result == ["endpoint"] + assert request_mock.call_args.args[1] == "plugin/tenant-1/endpoint/list" + assert request_mock.call_args.kwargs["params"] == {"page": 2, "page_size": 20} + + def test_list_endpoints_for_single_plugin(self, mocker: MockerFixture): + client = PluginEndpointClient() + request_mock = mocker.patch.object( + client, "_request_with_plugin_daemon_response", return_value=["endpoint"] + ) + + result = client.list_endpoints_for_single_plugin("tenant-1", "user-1", "org/plugin", 1, 10) + + assert result == ["endpoint"] + assert request_mock.call_args.args[1] == "plugin/tenant-1/endpoint/list/plugin" + assert request_mock.call_args.kwargs["params"] == { + "plugin_id": "org/plugin", + "page": 1, + "page_size": 10, + } + + def test_update_endpoint(self, mocker: MockerFixture): + client = PluginEndpointClient() + request_mock = mocker.patch.object(client, "_request_with_plugin_daemon_response", return_value=True) + + result = client.update_endpoint("tenant-1", "user-1", "endpoint-1", "renamed", {"x": 1}) + + assert result + assert request_mock.call_args.args[:3] == ("POST", "plugin/tenant-1/endpoint/update", bool) + + def test_enable_and_disable_endpoint(self, mocker: MockerFixture): + client = PluginEndpointClient() + request_mock = mocker.patch.object(client, "_request_with_plugin_daemon_response", return_value=True) + + assert client.enable_endpoint("tenant-1", "user-1", "endpoint-1") + assert client.disable_endpoint("tenant-1", "user-1", "endpoint-1") + + calls = request_mock.call_args_list + assert calls[0].args[1] == "plugin/tenant-1/endpoint/enable" + assert calls[1].args[1] == "plugin/tenant-1/endpoint/disable" + + def test_delete_endpoint_idempotent_and_re_raise(self, mocker: MockerFixture): + client = PluginEndpointClient() + request_mock = mocker.patch.object(client, "_request_with_plugin_daemon_response") + + request_mock.side_effect = PluginDaemonInternalServerError("record not found") + assert client.delete_endpoint("tenant-1", "user-1", "endpoint-1") + + request_mock.side_effect = PluginDaemonInternalServerError("permission denied") + with pytest.raises(PluginDaemonInternalServerError) as exc_info: + client.delete_endpoint("tenant-1", "user-1", "endpoint-1") + assert "permission denied" in exc_info.value.description + + +"""Unit tests for PluginEndpointClient functionality. + +This test module covers the endpoint client operations including: +- Successful endpoint deletion +- Idempotent delete behavior (record not found) +- Non-idempotent delete behavior (other errors) + +Tests follow the Arrange-Act-Assert pattern for clarity. +""" + + +class TestEndpointClientCases: + @staticmethod + @pytest.fixture(autouse=True) + def _patch_shared_httpx_client(): + """Patch module-level client methods to delegate to module httpx.request/stream.""" + + def request[**P](*args: P.args, **kwargs: P.kwargs) -> httpx.Response: + response = httpx.request(*args, **kwargs) + method, url = kwargs["method"], kwargs["url"] + assert isinstance(method, str) + assert isinstance(url, str) + response.request = httpx.Request(method, url) + return response + + with ( + patch("core.plugin.impl.base._httpx_client.request", side_effect=request), + patch("core.plugin.impl.base._httpx_client.stream", side_effect=lambda **kw: httpx.stream(**kw)), + ): + yield + + class TestPluginEndpointClientDelete: + """Unit tests for PluginEndpointClient delete_endpoint operation. + + Tests cover: + - Successful endpoint deletion + - Idempotent behavior when endpoint is already deleted (record not found) + - Non-idempotent behavior for other errors + """ + + @pytest.fixture + def endpoint_client(self): + """Create a PluginEndpointClient instance for testing.""" + return PluginEndpointClient() + + @pytest.fixture + def mock_config(self, config_overrides: Callable[..., None]): + """Mock plugin daemon configuration.""" + config_overrides(PLUGIN_DAEMON_URL="http://127.0.0.1:5002", PLUGIN_DAEMON_KEY="test-api-key") + + def test_delete_endpoint_success(self, endpoint_client, mock_config): + """Test successful endpoint deletion. + + Given: + - A valid tenant_id, user_id, and endpoint_id + - The plugin daemon returns success response + When: + - delete_endpoint is called + Then: + - The method should return True + - The request should be made with correct parameters + """ + # Arrange + tenant_id = "tenant-123" + user_id = "user-456" + endpoint_id = "endpoint-789" + + mock_response = httpx.Response( + 200, + json={ + "code": 0, + "message": "success", + "data": True, + }, + ) + + with patch("httpx.request", return_value=mock_response, autospec=True): + # Act + result = endpoint_client.delete_endpoint( + tenant_id=tenant_id, + user_id=user_id, + endpoint_id=endpoint_id, + ) + + # Assert + assert result is True + + def test_delete_endpoint_idempotent_record_not_found(self, endpoint_client, mock_config): + """Test idempotent delete behavior when endpoint is already deleted. + + Given: + - A valid tenant_id, user_id, and endpoint_id + - The plugin daemon returns "record not found" error + When: + - delete_endpoint is called + Then: + - The method should return True (idempotent behavior) + - No exception should be raised + """ + # Arrange + tenant_id = "tenant-123" + user_id = "user-456" + endpoint_id = "endpoint-789" + + mock_response = httpx.Response( + 200, + json={ + "code": -1, + "message": ( + '{"error_type": "PluginDaemonInternalServerError", ' + '"message": "failed to remove endpoint: record not found"}' + ), + }, + ) + + with patch("httpx.request", return_value=mock_response, autospec=True): + # Act + result = endpoint_client.delete_endpoint( + tenant_id=tenant_id, + user_id=user_id, + endpoint_id=endpoint_id, + ) + + # Assert - should return True instead of raising an error + assert result is True + + def test_delete_endpoint_non_idempotent_other_errors(self, endpoint_client, mock_config): + """Test non-idempotent delete behavior for other errors. + + Given: + - A valid tenant_id, user_id, and endpoint_id + - The plugin daemon returns a different error (not "record not found") + When: + - delete_endpoint is called + Then: + - The method should raise PluginDaemonInternalServerError + """ + # Arrange + tenant_id = "tenant-123" + user_id = "user-456" + endpoint_id = "endpoint-789" + + mock_response = httpx.Response( + 200, + json={ + "code": -1, + "message": ( + '{"error_type": "PluginDaemonInternalServerError", ' + '"message": "failed to remove endpoint: internal server error"}' + ), + }, + ) + + with patch("httpx.request", return_value=mock_response, autospec=True): + # Act & Assert + with pytest.raises(PluginDaemonInternalServerError) as exc_info: + endpoint_client.delete_endpoint( + tenant_id=tenant_id, + user_id=user_id, + endpoint_id=endpoint_id, + ) + + # Assert - the error message should not be "record not found" + assert "record not found" not in str(exc_info.value.description) + + def test_delete_endpoint_idempotent_case_insensitive(self, endpoint_client, mock_config): + """Test idempotent delete behavior with case-insensitive error message. + + Given: + - A valid tenant_id, user_id, and endpoint_id + - The plugin daemon returns "Record Not Found" error (different case) + When: + - delete_endpoint is called + Then: + - The method should return True (idempotent behavior) + """ + # Arrange + tenant_id = "tenant-123" + user_id = "user-456" + endpoint_id = "endpoint-789" + + mock_response = httpx.Response( + 200, + json={ + "code": -1, + "message": '{"error_type": "PluginDaemonInternalServerError", "message": "Record Not Found"}', + }, + ) + + with patch("httpx.request", return_value=mock_response, autospec=True): + # Act + result = endpoint_client.delete_endpoint( + tenant_id=tenant_id, + user_id=user_id, + endpoint_id=endpoint_id, + ) + + # Assert - should still return True + assert result is True + + def test_delete_endpoint_multiple_calls_idempotent(self, endpoint_client, mock_config): + """Test that multiple delete calls are idempotent. + + Given: + - A valid tenant_id, user_id, and endpoint_id + - The first call succeeds + - Subsequent calls return "record not found" + When: + - delete_endpoint is called multiple times + Then: + - All calls should return True + """ + # Arrange + tenant_id = "tenant-123" + user_id = "user-456" + endpoint_id = "endpoint-789" + + # First call - success + mock_response_success = httpx.Response( + 200, + json={ + "code": 0, + "message": "success", + "data": True, + }, + ) + + # Second call - record not found + mock_response_not_found = httpx.Response( + 200, + json={ + "code": -1, + "message": ( + '{"error_type": "PluginDaemonInternalServerError", ' + '"message": "failed to remove endpoint: record not found"}' + ), + }, + ) + + with patch("httpx.request", autospec=True) as mock_request: + # Act - first call + mock_request.return_value = mock_response_success + result1 = endpoint_client.delete_endpoint( + tenant_id=tenant_id, + user_id=user_id, + endpoint_id=endpoint_id, + ) + + # Act - second call (already deleted) + mock_request.return_value = mock_response_not_found + result2 = endpoint_client.delete_endpoint( + tenant_id=tenant_id, + user_id=user_id, + endpoint_id=endpoint_id, + ) + + # Assert - both should return True + assert result1 is True + assert result2 is True + + def test_delete_endpoint_non_idempotent_unauthorized_error(self, endpoint_client, mock_config): + """Test that authorization errors are not treated as idempotent. + + Given: + - A valid tenant_id, user_id, and endpoint_id + - The plugin daemon returns an unauthorized error + When: + - delete_endpoint is called + Then: + - The method should raise the appropriate error (not return True) + """ + # Arrange + tenant_id = "tenant-123" + user_id = "user-456" + endpoint_id = "endpoint-789" + + mock_response = httpx.Response( + 200, + json={ + "code": -1, + "message": '{"error_type": "PluginDaemonUnauthorizedError", "message": "unauthorized access"}', + }, + ) + + with patch("httpx.request", return_value=mock_response, autospec=True): + # Act & Assert + with pytest.raises( + Exception, match="req_id: PluginDaemonUnauthorizedError: unauthorized access" + ) as exc_info: + endpoint_client.delete_endpoint( + tenant_id=tenant_id, + user_id=user_id, + endpoint_id=endpoint_id, + ) + + # Assert - should not return True for unauthorized errors + assert exc_info.value.__class__.__name__ == "PluginDaemonUnauthorizedError" diff --git a/api/tests/unit_tests/core/plugin/impl/test_endpoint_client_impl.py b/api/tests/unit_tests/core/plugin/impl/test_endpoint_client_impl.py deleted file mode 100644 index e8d487c336d..00000000000 --- a/api/tests/unit_tests/core/plugin/impl/test_endpoint_client_impl.py +++ /dev/null @@ -1,112 +0,0 @@ -from collections.abc import Callable - -import httpx -import pytest -from pytest_mock import MockerFixture - -from core.plugin.impl.endpoint import PluginEndpointClient -from core.plugin.impl.exc import PluginDaemonInternalServerError - - -class TestPluginEndpointClientImpl: - @pytest.mark.parametrize( - "mutate", - [ - pytest.param( - lambda client: client.create_endpoint("tenant-1", "user-1", "org/plugin:1", "endpoint-a", {}), - id="create", - ), - pytest.param( - lambda client: client.update_endpoint("tenant-1", "user-1", "endpoint-1", "renamed", {}), - id="update", - ), - pytest.param(lambda client: client.delete_endpoint("tenant-1", "user-1", "endpoint-1"), id="delete"), - pytest.param(lambda client: client.enable_endpoint("tenant-1", "user-1", "endpoint-1"), id="enable"), - pytest.param(lambda client: client.disable_endpoint("tenant-1", "user-1", "endpoint-1"), id="disable"), - ], - ) - def test_mutation_rejects_false_daemon_result( - self, mocker: MockerFixture, mutate: Callable[[PluginEndpointClient], bool] - ) -> None: - client = PluginEndpointClient() - mocker.patch.object( - client, - "_request", - return_value=httpx.Response( - 200, - request=httpx.Request("POST", "https://daemon.test/endpoint"), - json={"code": 0, "message": "success", "data": False}, - ), - ) - - with pytest.raises(PluginDaemonInternalServerError, match="Plugin daemon failed to"): - mutate(client) - - def test_create_endpoint(self, mocker: MockerFixture): - client = PluginEndpointClient() - request_mock = mocker.patch.object(client, "_request_with_plugin_daemon_response", return_value=True) - - result = client.create_endpoint("tenant-1", "user-1", "org/plugin:1", "endpoint-a", {"k": "v"}) - - assert result - assert request_mock.call_count == 1 - args = request_mock.call_args.args - kwargs = request_mock.call_args.kwargs - assert args[:3] == ("POST", "plugin/tenant-1/endpoint/setup", bool) - assert kwargs["data"]["plugin_unique_identifier"] == "org/plugin:1" - - def test_list_endpoints(self, mocker: MockerFixture): - client = PluginEndpointClient() - request_mock = mocker.patch.object(client, "_request_with_plugin_daemon_response", return_value=["endpoint"]) - - result = client.list_endpoints("tenant-1", "user-1", 2, 20) - - assert result == ["endpoint"] - assert request_mock.call_args.args[1] == "plugin/tenant-1/endpoint/list" - assert request_mock.call_args.kwargs["params"] == {"page": 2, "page_size": 20} - - def test_list_endpoints_for_single_plugin(self, mocker: MockerFixture): - client = PluginEndpointClient() - request_mock = mocker.patch.object(client, "_request_with_plugin_daemon_response", return_value=["endpoint"]) - - result = client.list_endpoints_for_single_plugin("tenant-1", "user-1", "org/plugin", 1, 10) - - assert result == ["endpoint"] - assert request_mock.call_args.args[1] == "plugin/tenant-1/endpoint/list/plugin" - assert request_mock.call_args.kwargs["params"] == { - "plugin_id": "org/plugin", - "page": 1, - "page_size": 10, - } - - def test_update_endpoint(self, mocker: MockerFixture): - client = PluginEndpointClient() - request_mock = mocker.patch.object(client, "_request_with_plugin_daemon_response", return_value=True) - - result = client.update_endpoint("tenant-1", "user-1", "endpoint-1", "renamed", {"x": 1}) - - assert result - assert request_mock.call_args.args[:3] == ("POST", "plugin/tenant-1/endpoint/update", bool) - - def test_enable_and_disable_endpoint(self, mocker: MockerFixture): - client = PluginEndpointClient() - request_mock = mocker.patch.object(client, "_request_with_plugin_daemon_response", return_value=True) - - assert client.enable_endpoint("tenant-1", "user-1", "endpoint-1") - assert client.disable_endpoint("tenant-1", "user-1", "endpoint-1") - - calls = request_mock.call_args_list - assert calls[0].args[1] == "plugin/tenant-1/endpoint/enable" - assert calls[1].args[1] == "plugin/tenant-1/endpoint/disable" - - def test_delete_endpoint_idempotent_and_re_raise(self, mocker: MockerFixture): - client = PluginEndpointClient() - request_mock = mocker.patch.object(client, "_request_with_plugin_daemon_response") - - request_mock.side_effect = PluginDaemonInternalServerError("record not found") - assert client.delete_endpoint("tenant-1", "user-1", "endpoint-1") - - request_mock.side_effect = PluginDaemonInternalServerError("permission denied") - with pytest.raises(PluginDaemonInternalServerError) as exc_info: - client.delete_endpoint("tenant-1", "user-1", "endpoint-1") - assert "permission denied" in exc_info.value.description diff --git a/api/tests/unit_tests/core/plugin/impl/test_exc_impl.py b/api/tests/unit_tests/core/plugin/impl/test_exc.py similarity index 100% rename from api/tests/unit_tests/core/plugin/impl/test_exc_impl.py rename to api/tests/unit_tests/core/plugin/impl/test_exc.py diff --git a/api/tests/unit_tests/core/plugin/impl/test_model.py b/api/tests/unit_tests/core/plugin/impl/test_model.py new file mode 100644 index 00000000000..36e695096c6 --- /dev/null +++ b/api/tests/unit_tests/core/plugin/impl/test_model.py @@ -0,0 +1,916 @@ +from __future__ import annotations + +import io +import json +from types import SimpleNamespace + +import httpx +import pytest +from pytest_mock import MockerFixture + +from core.plugin.entities.plugin import PluginInstallationSource +from core.plugin.entities.plugin_daemon import ( + PluginDaemonInnerError, + PluginModelProviderDeclaration, + PluginStringResultResponse, +) +from core.plugin.impl import base as client_module +from core.plugin.impl.exc import PluginInvokeError, PluginLLMPollingUnsupportedError +from core.plugin.impl.model import PluginModelClient +from core.plugin.impl.model_runtime_factory import create_model_type_instance, create_plugin_model_runtime +from core.plugin.plugin_service import PluginService +from graphon.model_runtime.entities.common_entities import I18nObject +from graphon.model_runtime.entities.llm_entities import ( + LLMPollingResult, + LLMPollingStatus, + LLMResult, + LLMResultChunk, + LLMResultChunkDelta, + LLMUsage, +) +from graphon.model_runtime.entities.message_entities import AssistantPromptMessage +from graphon.model_runtime.entities.model_entities import ModelType +from graphon.model_runtime.entities.provider_entities import ConfigurateMethod, ProviderEntity + +PLUGIN_IDENTIFIER = "org/plugin:1.0.0@aaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaa" + + +class TestModelClientCases: + class TestPluginModelClient: + def test_fetch_model_providers(self, mocker: MockerFixture): + client = PluginModelClient() + request_mock = mocker.patch.object( + client, "_request_with_plugin_daemon_response", return_value=["provider-a"] + ) + + result = client.fetch_model_providers("tenant-1") + + assert result == ["provider-a"] + assert request_mock.call_args.args[:2] == ( + "GET", + "plugin/tenant-1/management/models", + ) + assert request_mock.call_args.kwargs["params"] == {"page": 1, "page_size": 256} + + def test_fetch_model_provider_bindings(self, mocker: MockerFixture): + client = PluginModelClient() + request_mock = mocker.patch.object( + client, "_request_with_plugin_daemon_response", return_value=["binding-a"] + ) + + result = client.fetch_model_provider_bindings("tenant-1") + + assert result == ["binding-a"] + assert request_mock.call_args.args[:2] == ( + "GET", + "plugin/tenant-1/management/models/bindings", + ) + assert "params" not in request_mock.call_args.kwargs + + def test_get_model_schema(self, mocker: MockerFixture): + client = PluginModelClient() + schema = SimpleNamespace(name="schema") + stream_mock = mocker.patch.object( + client, + "_request_with_plugin_daemon_response_stream", + return_value=iter([SimpleNamespace(model_schema=schema)]), + ) + + result = client.get_model_schema( + tenant_id="tenant-1", + user_id="user-1", + plugin_id="org/plugin", + provider="provider-a", + model_type="llm", + model="gpt-test", + credentials={"api_key": "key"}, + plugin_unique_identifier=PLUGIN_IDENTIFIER, + ) + + assert result is schema + assert stream_mock.call_args.args[:2] == ("POST", "plugin/tenant-1/dispatch/model/schema") + + def test_get_model_schema_empty_stream_returns_none(self, mocker: MockerFixture): + client = PluginModelClient() + mocker.patch.object(client, "_request_with_plugin_daemon_response_stream", return_value=iter([])) + + result = client.get_model_schema( + "tenant-1", + "user-1", + "org/plugin", + "provider-a", + "llm", + "gpt-test", + {}, + plugin_unique_identifier=PLUGIN_IDENTIFIER, + ) + + assert result is None + + def test_validate_provider_credentials(self, mocker: MockerFixture): + client = PluginModelClient() + stream_mock = mocker.patch.object( + client, + "_request_with_plugin_daemon_response_stream", + return_value=iter([SimpleNamespace(result=True, credentials={"api_key": "new"})]), + ) + credentials = {"api_key": "old"} + + result = client.validate_provider_credentials( + tenant_id="tenant-1", + user_id="user-1", + plugin_id="org/plugin", + provider="provider-a", + credentials=credentials, + plugin_unique_identifier=PLUGIN_IDENTIFIER, + ) + + assert result is True + assert credentials["api_key"] == "new" + assert stream_mock.call_args.args[:2] == ( + "POST", + "plugin/tenant-1/dispatch/model/validate_provider_credentials", + ) + + def test_validate_provider_credentials_without_dict_update(self, mocker: MockerFixture): + client = PluginModelClient() + mocker.patch.object( + client, + "_request_with_plugin_daemon_response_stream", + return_value=iter([SimpleNamespace(result=False, credentials="not-a-dict")]), + ) + credentials = {"api_key": "same"} + + result = client.validate_provider_credentials( + "tenant-1", + "user-1", + "org/plugin", + "provider-a", + credentials, + plugin_unique_identifier=PLUGIN_IDENTIFIER, + ) + + assert result is False + assert credentials == {"api_key": "same"} + + def test_validate_provider_credentials_empty_returns_false(self, mocker: MockerFixture): + client = PluginModelClient() + mocker.patch.object(client, "_request_with_plugin_daemon_response_stream", return_value=iter([])) + + assert ( + client.validate_provider_credentials( + "tenant-1", + "user-1", + "org/plugin", + "provider-a", + {}, + plugin_unique_identifier=PLUGIN_IDENTIFIER, + ) + is False + ) + + def test_validate_model_credentials(self, mocker: MockerFixture): + client = PluginModelClient() + stream_mock = mocker.patch.object( + client, + "_request_with_plugin_daemon_response_stream", + return_value=iter([SimpleNamespace(result=True, credentials={"token": "rotated"})]), + ) + credentials = {"token": "old"} + + result = client.validate_model_credentials( + tenant_id="tenant-1", + user_id="user-1", + plugin_id="org/plugin", + provider="provider-a", + model_type="llm", + model="gpt-test", + credentials=credentials, + plugin_unique_identifier=PLUGIN_IDENTIFIER, + ) + + assert result is True + assert credentials["token"] == "rotated" + assert stream_mock.call_args.args[:2] == ( + "POST", + "plugin/tenant-1/dispatch/model/validate_model_credentials", + ) + + def test_validate_model_credentials_empty_returns_false(self, mocker: MockerFixture): + client = PluginModelClient() + mocker.patch.object(client, "_request_with_plugin_daemon_response_stream", return_value=iter([])) + + assert ( + client.validate_model_credentials( + "tenant-1", + "user-1", + "org/plugin", + "provider-a", + "llm", + "gpt-test", + {}, + plugin_unique_identifier=PLUGIN_IDENTIFIER, + ) + is False + ) + + def test_invoke_llm(self, mocker: MockerFixture): + client = PluginModelClient() + stream_mock = mocker.patch.object( + client, "_request_with_plugin_daemon_response_stream", return_value=iter(["chunk-1"]) + ) + + result = list( + client.invoke_llm( + tenant_id="tenant-1", + user_id="user-1", + plugin_id="org/plugin", + provider="provider-a", + model="gpt-test", + credentials={"api_key": "key"}, + prompt_messages=[], + model_parameters={"temperature": 0.1}, + tools=[], + stop=["STOP"], + stream=False, + app_id="app-1", + plugin_unique_identifier=PLUGIN_IDENTIFIER, + ) + ) + + assert result == ["chunk-1"] + call_kwargs = stream_mock.call_args.kwargs + assert call_kwargs["path"] == "plugin/tenant-1/dispatch/llm/invoke" + assert call_kwargs["data"]["app_id"] == "app-1" + assert call_kwargs["data"]["data"]["stream"] is False + assert call_kwargs["data"]["data"]["model_parameters"] == {"temperature": 0.1} + + def test_invoke_llm_omits_app_id_when_missing(self, mocker: MockerFixture): + client = PluginModelClient() + stream_mock = mocker.patch.object( + client, "_request_with_plugin_daemon_response_stream", return_value=iter([]) + ) + + list( + client.invoke_llm( + tenant_id="tenant-1", + user_id="user-1", + plugin_id="org/plugin", + provider="provider-a", + model="gpt-test", + credentials={}, + prompt_messages=[], + plugin_unique_identifier=PLUGIN_IDENTIFIER, + ) + ) + + assert "app_id" not in stream_mock.call_args.kwargs["data"] + + def test_invoke_llm_wraps_plugin_daemon_inner_error(self, mocker: MockerFixture): + client = PluginModelClient() + + def _boom(): + raise PluginDaemonInnerError(code=-500, message="invoke failed") + yield # pragma: no cover + + mocker.patch.object(client, "_request_with_plugin_daemon_response_stream", return_value=_boom()) + + with pytest.raises(ValueError, match="invoke failed-500"): + list( + client.invoke_llm( + tenant_id="tenant-1", + user_id="user-1", + plugin_id="org/plugin", + provider="provider-a", + model="gpt-test", + credentials={}, + prompt_messages=[], + plugin_unique_identifier=PLUGIN_IDENTIFIER, + ) + ) + + def test_start_llm_polling(self, mocker: MockerFixture): + client = PluginModelClient() + polling_result = LLMPollingResult( + status=LLMPollingStatus.RUNNING, + plugin_state={"task_id": "poll-1"}, + next_check_after_seconds=3, + ) + request_mock = mocker.patch.object( + client, + "_request_with_plugin_daemon_response", + return_value=polling_result, + ) + + result = client.start_llm_polling( + tenant_id="tenant-1", + user_id="user-1", + plugin_id="org/plugin", + provider="provider-a", + model="gpt-test", + credentials={"api_key": "key"}, + prompt_messages=[], + model_parameters={"temperature": 0.1}, + tools=[], + stop=["STOP"], + json_schema={"type": "object"}, + plugin_unique_identifier=PLUGIN_IDENTIFIER, + ) + + assert result == polling_result + call_kwargs = request_mock.call_args.kwargs + assert call_kwargs["path"] == "plugin/tenant-1/dispatch/model/polling/start" + assert call_kwargs["data"]["data"] == { + "provider": "provider-a", + "model_type": "llm", + "model": "gpt-test", + "credentials": {"api_key": "key"}, + "prompt_messages": [], + "model_parameters": {"temperature": 0.1}, + "tools": [], + "stop": ["STOP"], + "stream": False, + "json_schema": {"type": "object"}, + } + + def test_check_llm_polling(self, mocker: MockerFixture): + client = PluginModelClient() + polling_result = LLMPollingResult( + status=LLMPollingStatus.SUCCEEDED, + result=LLMResult( + model="gpt-test", + prompt_messages=[], + message=AssistantPromptMessage(content="done"), + usage=LLMUsage.empty_usage(), + ), + ) + request_mock = mocker.patch.object( + client, + "_request_with_plugin_daemon_response", + return_value=polling_result, + ) + + result = client.check_llm_polling( + tenant_id="tenant-1", + user_id="user-1", + plugin_id="org/plugin", + provider="provider-a", + model="gpt-test", + credentials={"api_key": "key"}, + plugin_state={"task_id": "poll-1"}, + plugin_unique_identifier=PLUGIN_IDENTIFIER, + ) + + assert result == polling_result + call_kwargs = request_mock.call_args.kwargs + assert call_kwargs["path"] == "plugin/tenant-1/dispatch/model/polling/check" + assert call_kwargs["data"]["data"] == { + "provider": "provider-a", + "model_type": "llm", + "model": "gpt-test", + "credentials": {"api_key": "key"}, + "plugin_state": {"task_id": "poll-1"}, + } + + def test_start_llm_polling_maps_unsupported_polling_invoke_error(self, mocker: MockerFixture): + client = PluginModelClient() + mocker.patch.object( + client, + "_request_with_plugin_daemon_response", + side_effect=PluginInvokeError( + json.dumps( + { + "error_type": PluginLLMPollingUnsupportedError.__name__, + "message": "Model `gpt-test` does not support polling.", + } + ) + ), + ) + + with pytest.raises(PluginLLMPollingUnsupportedError): + client.start_llm_polling( + tenant_id="tenant-1", + user_id="user-1", + plugin_id="org/plugin", + provider="provider-a", + model="gpt-test", + credentials={"api_key": "key"}, + prompt_messages=[], + plugin_unique_identifier=PLUGIN_IDENTIFIER, + ) + + def test_get_llm_num_tokens(self, mocker: MockerFixture): + client = PluginModelClient() + mocker.patch.object( + client, + "_request_with_plugin_daemon_response_stream", + return_value=iter([SimpleNamespace(num_tokens=42)]), + ) + + result = client.get_llm_num_tokens( + tenant_id="tenant-1", + user_id="user-1", + plugin_id="org/plugin", + provider="provider-a", + model_type="llm", + model="gpt-test", + credentials={}, + prompt_messages=[], + tools=[], + plugin_unique_identifier=PLUGIN_IDENTIFIER, + ) + + assert result == 42 + + def test_get_llm_num_tokens_empty_returns_zero(self, mocker: MockerFixture): + client = PluginModelClient() + mocker.patch.object(client, "_request_with_plugin_daemon_response_stream", return_value=iter([])) + + assert ( + client.get_llm_num_tokens( + "tenant-1", + "user-1", + "org/plugin", + "provider-a", + "llm", + "gpt-test", + {}, + [], + plugin_unique_identifier=PLUGIN_IDENTIFIER, + ) + == 0 + ) + + def test_invoke_text_embedding(self, mocker: MockerFixture): + client = PluginModelClient() + embedding_result = SimpleNamespace(data=[[0.1, 0.2]]) + mocker.patch.object( + client, "_request_with_plugin_daemon_response_stream", return_value=iter([embedding_result]) + ) + + result = client.invoke_text_embedding( + tenant_id="tenant-1", + user_id="user-1", + plugin_id="org/plugin", + provider="provider-a", + model="embedding-a", + credentials={}, + texts=["hello"], + input_type="search_document", + plugin_unique_identifier=PLUGIN_IDENTIFIER, + ) + + assert result is embedding_result + + def test_invoke_text_embedding_empty_raises(self, mocker: MockerFixture): + client = PluginModelClient() + mocker.patch.object(client, "_request_with_plugin_daemon_response_stream", return_value=iter([])) + + with pytest.raises(ValueError, match="Failed to invoke text embedding"): + client.invoke_text_embedding( + "tenant-1", + "user-1", + "org/plugin", + "provider-a", + "embedding-a", + {}, + ["hello"], + "x", + plugin_unique_identifier=PLUGIN_IDENTIFIER, + ) + + def test_invoke_multimodal_embedding(self, mocker: MockerFixture): + client = PluginModelClient() + embedding_result = SimpleNamespace(data=[[0.3, 0.4]]) + mocker.patch.object( + client, "_request_with_plugin_daemon_response_stream", return_value=iter([embedding_result]) + ) + + result = client.invoke_multimodal_embedding( + tenant_id="tenant-1", + user_id="user-1", + plugin_id="org/plugin", + provider="provider-a", + model="embedding-a", + credentials={}, + documents=[{"type": "image", "value": "abc"}], + input_type="search_document", + plugin_unique_identifier=PLUGIN_IDENTIFIER, + ) + + assert result is embedding_result + + def test_invoke_multimodal_embedding_empty_raises(self, mocker: MockerFixture): + client = PluginModelClient() + mocker.patch.object(client, "_request_with_plugin_daemon_response_stream", return_value=iter([])) + + with pytest.raises(ValueError, match="Failed to invoke file embedding"): + client.invoke_multimodal_embedding( + "tenant-1", + "user-1", + "org/plugin", + "provider-a", + "embedding-a", + {}, + [{"type": "image"}], + "x", + plugin_unique_identifier=PLUGIN_IDENTIFIER, + ) + + def test_get_text_embedding_num_tokens(self, mocker: MockerFixture): + client = PluginModelClient() + mocker.patch.object( + client, + "_request_with_plugin_daemon_response_stream", + return_value=iter([SimpleNamespace(num_tokens=[1, 2, 3])]), + ) + + assert client.get_text_embedding_num_tokens( + "tenant-1", + "user-1", + "org/plugin", + "provider-a", + "embedding-a", + {}, + ["a"], + plugin_unique_identifier=PLUGIN_IDENTIFIER, + ) == [ + 1, + 2, + 3, + ] + + def test_get_text_embedding_num_tokens_empty_returns_list(self, mocker: MockerFixture): + client = PluginModelClient() + mocker.patch.object(client, "_request_with_plugin_daemon_response_stream", return_value=iter([])) + + assert ( + client.get_text_embedding_num_tokens( + "tenant-1", + "user-1", + "org/plugin", + "provider-a", + "embedding-a", + {}, + ["a"], + plugin_unique_identifier=PLUGIN_IDENTIFIER, + ) + == [] + ) + + def test_invoke_rerank(self, mocker: MockerFixture): + client = PluginModelClient() + rerank_result = SimpleNamespace(scores=[0.9]) + mocker.patch.object( + client, "_request_with_plugin_daemon_response_stream", return_value=iter([rerank_result]) + ) + + result = client.invoke_rerank( + tenant_id="tenant-1", + user_id="user-1", + plugin_id="org/plugin", + provider="provider-a", + model="rerank-a", + credentials={}, + query="q", + docs=["doc-1"], + score_threshold=0.2, + top_n=5, + plugin_unique_identifier=PLUGIN_IDENTIFIER, + ) + + assert result is rerank_result + + def test_invoke_rerank_empty_raises(self, mocker: MockerFixture): + client = PluginModelClient() + mocker.patch.object(client, "_request_with_plugin_daemon_response_stream", return_value=iter([])) + + with pytest.raises(ValueError, match="Failed to invoke rerank"): + client.invoke_rerank( + "tenant-1", + "user-1", + "org/plugin", + "provider-a", + "rerank-a", + {}, + "q", + ["doc-1"], + plugin_unique_identifier=PLUGIN_IDENTIFIER, + ) + + def test_invoke_multimodal_rerank(self, mocker: MockerFixture): + client = PluginModelClient() + rerank_result = SimpleNamespace(scores=[0.8]) + mocker.patch.object( + client, "_request_with_plugin_daemon_response_stream", return_value=iter([rerank_result]) + ) + + result = client.invoke_multimodal_rerank( + tenant_id="tenant-1", + user_id="user-1", + plugin_id="org/plugin", + provider="provider-a", + model="rerank-a", + credentials={}, + query={"type": "text", "value": "q"}, + docs=[{"type": "image", "value": "doc"}], + score_threshold=0.1, + top_n=3, + plugin_unique_identifier=PLUGIN_IDENTIFIER, + ) + + assert result is rerank_result + + def test_invoke_multimodal_rerank_empty_raises(self, mocker: MockerFixture): + client = PluginModelClient() + mocker.patch.object(client, "_request_with_plugin_daemon_response_stream", return_value=iter([])) + + with pytest.raises(ValueError, match="Failed to invoke multimodal rerank"): + client.invoke_multimodal_rerank( + "tenant-1", + "user-1", + "org/plugin", + "provider-a", + "rerank-a", + {}, + {"type": "text"}, + [{"type": "image"}], + plugin_unique_identifier=PLUGIN_IDENTIFIER, + ) + + def test_invoke_tts(self, mocker: MockerFixture): + client = PluginModelClient() + mocker.patch.object( + client, + "_request_with_plugin_daemon_response_stream", + return_value=iter( + [ + SimpleNamespace(result="68656c6c6f", mime_type="audio/wav"), + SimpleNamespace(result="21", mime_type="audio/wav"), + ] + ), + ) + + result = list( + client.invoke_tts( + tenant_id="tenant-1", + user_id="user-1", + plugin_id="org/plugin", + provider="provider-a", + model="tts-a", + credentials={}, + content_text="hello", + voice="alloy", + plugin_unique_identifier=PLUGIN_IDENTIFIER, + ) + ) + + assert result == [b"hello", b"!"] + assert [chunk.mime_type for chunk in result] == ["audio/wav", "audio/wav"] + + def test_invoke_tts_wraps_plugin_daemon_inner_error(self, mocker: MockerFixture): + client = PluginModelClient() + + def _boom(): + raise PluginDaemonInnerError(code=-400, message="tts error") + yield # pragma: no cover + + mocker.patch.object(client, "_request_with_plugin_daemon_response_stream", return_value=_boom()) + + with pytest.raises(ValueError, match="tts error-400"): + list( + client.invoke_tts( + "tenant-1", + "user-1", + "org/plugin", + "provider-a", + "tts-a", + {}, + "hello", + "alloy", + plugin_unique_identifier=PLUGIN_IDENTIFIER, + ) + ) + + def test_get_tts_model_voices(self, mocker: MockerFixture): + client = PluginModelClient() + mocker.patch.object( + client, + "_request_with_plugin_daemon_response_stream", + return_value=iter( + [ + SimpleNamespace( + voices=[ + SimpleNamespace(name="Alloy", value="alloy"), + SimpleNamespace(name="Echo", value="echo"), + ] + ) + ] + ), + ) + + result = client.get_tts_model_voices( + tenant_id="tenant-1", + user_id="user-1", + plugin_id="org/plugin", + provider="provider-a", + model="tts-a", + credentials={}, + language="en", + plugin_unique_identifier=PLUGIN_IDENTIFIER, + ) + + assert result == [{"name": "Alloy", "value": "alloy"}, {"name": "Echo", "value": "echo"}] + + def test_get_tts_model_voices_empty_returns_list(self, mocker: MockerFixture): + client = PluginModelClient() + mocker.patch.object(client, "_request_with_plugin_daemon_response_stream", return_value=iter([])) + + assert ( + client.get_tts_model_voices( + "tenant-1", + "user-1", + "org/plugin", + "provider-a", + "tts-a", + {}, + plugin_unique_identifier=PLUGIN_IDENTIFIER, + ) + == [] + ) + + def test_invoke_speech_to_text(self, mocker: MockerFixture): + client = PluginModelClient() + stream_mock = mocker.patch.object( + client, + "_request_with_plugin_daemon_response_stream", + return_value=iter([PluginStringResultResponse(result="transcribed text")]), + ) + + result = client.invoke_speech_to_text( + tenant_id="tenant-1", + user_id="user-1", + plugin_id="org/plugin", + provider="provider-a", + model="stt-a", + credentials={}, + file=io.BytesIO(b"abc"), + plugin_unique_identifier=PLUGIN_IDENTIFIER, + ) + + assert result == "transcribed text" + assert stream_mock.call_args.kwargs["data"]["data"]["file"] == "616263" + + def test_invoke_speech_to_text_empty_raises(self, mocker: MockerFixture): + client = PluginModelClient() + mocker.patch.object(client, "_request_with_plugin_daemon_response_stream", return_value=iter([])) + + with pytest.raises(ValueError, match="Failed to invoke speech to text"): + client.invoke_speech_to_text( + "tenant-1", + "user-1", + "org/plugin", + "provider-a", + "stt-a", + {}, + io.BytesIO(b"abc"), + plugin_unique_identifier=PLUGIN_IDENTIFIER, + ) + + def test_invoke_moderation(self, mocker: MockerFixture): + client = PluginModelClient() + stream_mock = mocker.patch.object( + client, + "_request_with_plugin_daemon_response_stream", + return_value=iter([SimpleNamespace(result=True)]), + ) + + result = client.invoke_moderation( + tenant_id="tenant-1", + user_id="user-1", + plugin_id="org/plugin", + provider="provider-a", + model="moderation-a", + credentials={}, + text="safe text", + plugin_unique_identifier=PLUGIN_IDENTIFIER, + ) + + assert result is True + assert stream_mock.call_args.kwargs["path"] == "plugin/tenant-1/dispatch/moderation/invoke" + + def test_invoke_moderation_empty_raises(self, mocker: MockerFixture): + client = PluginModelClient() + mocker.patch.object(client, "_request_with_plugin_daemon_response_stream", return_value=iter([])) + + with pytest.raises(ValueError, match="Failed to invoke moderation"): + client.invoke_moderation( + "tenant-1", + "user-1", + "org/plugin", + "provider-a", + "moderation-a", + {}, + "unsafe", + plugin_unique_identifier=PLUGIN_IDENTIFIER, + ) + + +"""Model dispatches reuse prepared package identities without discovery I/O.""" + + +IDENTIFIER = f"langgenius/openai:1.0.0@{'a' * 32}" + + +def _provider() -> PluginModelProviderDeclaration: + return PluginModelProviderDeclaration( + provider="langgenius/openai/openai", + provider_name="openai", + label=I18nObject(en_US="OpenAI"), + supported_model_types=[ModelType.LLM], + configurate_methods=[ConfigurateMethod.PREDEFINED_MODEL], + plugin_unique_identifier=IDENTIFIER, + installation_source=PluginInstallationSource.Marketplace, + ) + + +class TestModelDispatchBindingCases: + from tests.unit_tests.services.plugin.conftest import mock_installer + + mock_installer = staticmethod(mock_installer) + + from tests.unit_tests.services.plugin.conftest import mock_features + + mock_features = staticmethod(mock_features) + + @staticmethod + @pytest.mark.parametrize("binding", ["assembly", "discovery", "legacy"]) + def test_model_dispatch_does_not_discover_providers(monkeypatch: pytest.MonkeyPatch, binding: str) -> None: + provider = _provider() + runtime = create_plugin_model_runtime(tenant_id="tenant-a") + if binding == "discovery": + + def fetch( + *, tenant_id: str, client: PluginModelClient | None = None + ) -> list[PluginModelProviderDeclaration]: + assert tenant_id == "tenant-a" + assert client is runtime.client + return [provider] + + monkeypatch.setattr(PluginService, "fetch_plugin_model_providers", fetch) + assert runtime.fetch_model_providers() == [provider] + else: + schema = ProviderEntity.model_validate(provider.model_dump()) if binding == "legacy" else provider + create_model_type_instance(runtime=runtime, provider_schema=schema, model_type=ModelType.LLM) + + def unexpected_discovery(*_args: object, **_kwargs: object) -> None: + raise AssertionError("Invocation must reuse the prepared provider identity") + + monkeypatch.setattr(PluginService, "fetch_plugin_model_providers", unexpected_discovery) + requests: list[httpx.Request] = [] + chunk = LLMResultChunk( + model="model", delta=LLMResultChunkDelta(index=0, message=AssistantPromptMessage(content="ok")) + ) + + def handle(request: httpx.Request) -> httpx.Response: + assert request.method == "POST" + assert request.url.path == "/plugin/tenant-a/dispatch/llm/invoke" + requests.append(request) + return httpx.Response( + 200, text="data: " + json.dumps({"code": 0, "message": "", "data": chunk.model_dump()}) + ) + + with httpx.Client(transport=httpx.MockTransport(handle)) as client: + monkeypatch.setattr(client_module, "_httpx_client", client) + result = runtime.invoke_llm( + provider=provider.provider, + model="model", + credentials={"api_key": "secret"}, + model_parameters={}, + prompt_messages=[], + tools=None, + stop=None, + stream=True, + ) + assert list(result) == [chunk] + + assert len(requests) == 1 + assert requests[0].headers["X-Plugin-ID"] == "langgenius/openai" + expected_identifier = None if binding == "legacy" else IDENTIFIER + assert requests[0].headers.get("X-Plugin-Unique-Identifier") == expected_identifier + + @staticmethod + def test_model_schema_cache_is_scoped_to_prepared_package() -> None: + runtime = create_plugin_model_runtime(tenant_id="tenant-a") + provider = _provider() + runtime.bind_provider_schema(provider) + first = runtime._get_schema_cache_key( + provider=provider.provider, model_type=ModelType.LLM, model="model", credentials={} + ) + runtime.bind_provider_schema( + provider.model_copy(update={"plugin_unique_identifier": f"langgenius/openai:2.0.0@{'b' * 32}"}) + ) + second = runtime._get_schema_cache_key( + provider=provider.provider, model_type=ModelType.LLM, model="model", credentials={} + ) + assert first != second diff --git a/api/tests/unit_tests/core/plugin/impl/test_model_client.py b/api/tests/unit_tests/core/plugin/impl/test_model_client.py deleted file mode 100644 index ca5965eb21c..00000000000 --- a/api/tests/unit_tests/core/plugin/impl/test_model_client.py +++ /dev/null @@ -1,786 +0,0 @@ -from __future__ import annotations - -import io -import json -from types import SimpleNamespace - -import pytest -from pytest_mock import MockerFixture - -from core.plugin.entities.plugin_daemon import PluginDaemonInnerError, PluginStringResultResponse -from core.plugin.impl.exc import PluginInvokeError, PluginLLMPollingUnsupportedError -from core.plugin.impl.model import PluginModelClient -from graphon.model_runtime.entities.llm_entities import LLMPollingResult, LLMPollingStatus, LLMResult, LLMUsage -from graphon.model_runtime.entities.message_entities import AssistantPromptMessage - -PLUGIN_IDENTIFIER = "org/plugin:1.0.0@aaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaa" - - -class TestPluginModelClient: - def test_fetch_model_providers(self, mocker: MockerFixture): - client = PluginModelClient() - request_mock = mocker.patch.object(client, "_request_with_plugin_daemon_response", return_value=["provider-a"]) - - result = client.fetch_model_providers("tenant-1") - - assert result == ["provider-a"] - assert request_mock.call_args.args[:2] == ( - "GET", - "plugin/tenant-1/management/models", - ) - assert request_mock.call_args.kwargs["params"] == {"page": 1, "page_size": 256} - - def test_fetch_model_provider_bindings(self, mocker: MockerFixture): - client = PluginModelClient() - request_mock = mocker.patch.object(client, "_request_with_plugin_daemon_response", return_value=["binding-a"]) - - result = client.fetch_model_provider_bindings("tenant-1") - - assert result == ["binding-a"] - assert request_mock.call_args.args[:2] == ( - "GET", - "plugin/tenant-1/management/models/bindings", - ) - assert "params" not in request_mock.call_args.kwargs - - def test_get_model_schema(self, mocker: MockerFixture): - client = PluginModelClient() - schema = SimpleNamespace(name="schema") - stream_mock = mocker.patch.object( - client, - "_request_with_plugin_daemon_response_stream", - return_value=iter([SimpleNamespace(model_schema=schema)]), - ) - - result = client.get_model_schema( - tenant_id="tenant-1", - user_id="user-1", - plugin_id="org/plugin", - provider="provider-a", - model_type="llm", - model="gpt-test", - credentials={"api_key": "key"}, - plugin_unique_identifier=PLUGIN_IDENTIFIER, - ) - - assert result is schema - assert stream_mock.call_args.args[:2] == ("POST", "plugin/tenant-1/dispatch/model/schema") - - def test_get_model_schema_empty_stream_returns_none(self, mocker: MockerFixture): - client = PluginModelClient() - mocker.patch.object(client, "_request_with_plugin_daemon_response_stream", return_value=iter([])) - - result = client.get_model_schema( - "tenant-1", - "user-1", - "org/plugin", - "provider-a", - "llm", - "gpt-test", - {}, - plugin_unique_identifier=PLUGIN_IDENTIFIER, - ) - - assert result is None - - def test_validate_provider_credentials(self, mocker: MockerFixture): - client = PluginModelClient() - stream_mock = mocker.patch.object( - client, - "_request_with_plugin_daemon_response_stream", - return_value=iter([SimpleNamespace(result=True, credentials={"api_key": "new"})]), - ) - credentials = {"api_key": "old"} - - result = client.validate_provider_credentials( - tenant_id="tenant-1", - user_id="user-1", - plugin_id="org/plugin", - provider="provider-a", - credentials=credentials, - plugin_unique_identifier=PLUGIN_IDENTIFIER, - ) - - assert result is True - assert credentials["api_key"] == "new" - assert stream_mock.call_args.args[:2] == ( - "POST", - "plugin/tenant-1/dispatch/model/validate_provider_credentials", - ) - - def test_validate_provider_credentials_without_dict_update(self, mocker: MockerFixture): - client = PluginModelClient() - mocker.patch.object( - client, - "_request_with_plugin_daemon_response_stream", - return_value=iter([SimpleNamespace(result=False, credentials="not-a-dict")]), - ) - credentials = {"api_key": "same"} - - result = client.validate_provider_credentials( - "tenant-1", - "user-1", - "org/plugin", - "provider-a", - credentials, - plugin_unique_identifier=PLUGIN_IDENTIFIER, - ) - - assert result is False - assert credentials == {"api_key": "same"} - - def test_validate_provider_credentials_empty_returns_false(self, mocker: MockerFixture): - client = PluginModelClient() - mocker.patch.object(client, "_request_with_plugin_daemon_response_stream", return_value=iter([])) - - assert ( - client.validate_provider_credentials( - "tenant-1", - "user-1", - "org/plugin", - "provider-a", - {}, - plugin_unique_identifier=PLUGIN_IDENTIFIER, - ) - is False - ) - - def test_validate_model_credentials(self, mocker: MockerFixture): - client = PluginModelClient() - stream_mock = mocker.patch.object( - client, - "_request_with_plugin_daemon_response_stream", - return_value=iter([SimpleNamespace(result=True, credentials={"token": "rotated"})]), - ) - credentials = {"token": "old"} - - result = client.validate_model_credentials( - tenant_id="tenant-1", - user_id="user-1", - plugin_id="org/plugin", - provider="provider-a", - model_type="llm", - model="gpt-test", - credentials=credentials, - plugin_unique_identifier=PLUGIN_IDENTIFIER, - ) - - assert result is True - assert credentials["token"] == "rotated" - assert stream_mock.call_args.args[:2] == ( - "POST", - "plugin/tenant-1/dispatch/model/validate_model_credentials", - ) - - def test_validate_model_credentials_empty_returns_false(self, mocker: MockerFixture): - client = PluginModelClient() - mocker.patch.object(client, "_request_with_plugin_daemon_response_stream", return_value=iter([])) - - assert ( - client.validate_model_credentials( - "tenant-1", - "user-1", - "org/plugin", - "provider-a", - "llm", - "gpt-test", - {}, - plugin_unique_identifier=PLUGIN_IDENTIFIER, - ) - is False - ) - - def test_invoke_llm(self, mocker: MockerFixture): - client = PluginModelClient() - stream_mock = mocker.patch.object( - client, "_request_with_plugin_daemon_response_stream", return_value=iter(["chunk-1"]) - ) - - result = list( - client.invoke_llm( - tenant_id="tenant-1", - user_id="user-1", - plugin_id="org/plugin", - provider="provider-a", - model="gpt-test", - credentials={"api_key": "key"}, - prompt_messages=[], - model_parameters={"temperature": 0.1}, - tools=[], - stop=["STOP"], - stream=False, - app_id="app-1", - plugin_unique_identifier=PLUGIN_IDENTIFIER, - ) - ) - - assert result == ["chunk-1"] - call_kwargs = stream_mock.call_args.kwargs - assert call_kwargs["path"] == "plugin/tenant-1/dispatch/llm/invoke" - assert call_kwargs["data"]["app_id"] == "app-1" - assert call_kwargs["data"]["data"]["stream"] is False - assert call_kwargs["data"]["data"]["model_parameters"] == {"temperature": 0.1} - - def test_invoke_llm_omits_app_id_when_missing(self, mocker: MockerFixture): - client = PluginModelClient() - stream_mock = mocker.patch.object(client, "_request_with_plugin_daemon_response_stream", return_value=iter([])) - - list( - client.invoke_llm( - tenant_id="tenant-1", - user_id="user-1", - plugin_id="org/plugin", - provider="provider-a", - model="gpt-test", - credentials={}, - prompt_messages=[], - plugin_unique_identifier=PLUGIN_IDENTIFIER, - ) - ) - - assert "app_id" not in stream_mock.call_args.kwargs["data"] - - def test_invoke_llm_wraps_plugin_daemon_inner_error(self, mocker: MockerFixture): - client = PluginModelClient() - - def _boom(): - raise PluginDaemonInnerError(code=-500, message="invoke failed") - yield # pragma: no cover - - mocker.patch.object(client, "_request_with_plugin_daemon_response_stream", return_value=_boom()) - - with pytest.raises(ValueError, match="invoke failed-500"): - list( - client.invoke_llm( - tenant_id="tenant-1", - user_id="user-1", - plugin_id="org/plugin", - provider="provider-a", - model="gpt-test", - credentials={}, - prompt_messages=[], - plugin_unique_identifier=PLUGIN_IDENTIFIER, - ) - ) - - def test_start_llm_polling(self, mocker: MockerFixture): - client = PluginModelClient() - polling_result = LLMPollingResult( - status=LLMPollingStatus.RUNNING, - plugin_state={"task_id": "poll-1"}, - next_check_after_seconds=3, - ) - request_mock = mocker.patch.object( - client, - "_request_with_plugin_daemon_response", - return_value=polling_result, - ) - - result = client.start_llm_polling( - tenant_id="tenant-1", - user_id="user-1", - plugin_id="org/plugin", - provider="provider-a", - model="gpt-test", - credentials={"api_key": "key"}, - prompt_messages=[], - model_parameters={"temperature": 0.1}, - tools=[], - stop=["STOP"], - json_schema={"type": "object"}, - plugin_unique_identifier=PLUGIN_IDENTIFIER, - ) - - assert result == polling_result - call_kwargs = request_mock.call_args.kwargs - assert call_kwargs["path"] == "plugin/tenant-1/dispatch/model/polling/start" - assert call_kwargs["data"]["data"] == { - "provider": "provider-a", - "model_type": "llm", - "model": "gpt-test", - "credentials": {"api_key": "key"}, - "prompt_messages": [], - "model_parameters": {"temperature": 0.1}, - "tools": [], - "stop": ["STOP"], - "stream": False, - "json_schema": {"type": "object"}, - } - - def test_check_llm_polling(self, mocker: MockerFixture): - client = PluginModelClient() - polling_result = LLMPollingResult( - status=LLMPollingStatus.SUCCEEDED, - result=LLMResult( - model="gpt-test", - prompt_messages=[], - message=AssistantPromptMessage(content="done"), - usage=LLMUsage.empty_usage(), - ), - ) - request_mock = mocker.patch.object( - client, - "_request_with_plugin_daemon_response", - return_value=polling_result, - ) - - result = client.check_llm_polling( - tenant_id="tenant-1", - user_id="user-1", - plugin_id="org/plugin", - provider="provider-a", - model="gpt-test", - credentials={"api_key": "key"}, - plugin_state={"task_id": "poll-1"}, - plugin_unique_identifier=PLUGIN_IDENTIFIER, - ) - - assert result == polling_result - call_kwargs = request_mock.call_args.kwargs - assert call_kwargs["path"] == "plugin/tenant-1/dispatch/model/polling/check" - assert call_kwargs["data"]["data"] == { - "provider": "provider-a", - "model_type": "llm", - "model": "gpt-test", - "credentials": {"api_key": "key"}, - "plugin_state": {"task_id": "poll-1"}, - } - - def test_start_llm_polling_maps_unsupported_polling_invoke_error(self, mocker: MockerFixture): - client = PluginModelClient() - mocker.patch.object( - client, - "_request_with_plugin_daemon_response", - side_effect=PluginInvokeError( - json.dumps( - { - "error_type": PluginLLMPollingUnsupportedError.__name__, - "message": "Model `gpt-test` does not support polling.", - } - ) - ), - ) - - with pytest.raises(PluginLLMPollingUnsupportedError): - client.start_llm_polling( - tenant_id="tenant-1", - user_id="user-1", - plugin_id="org/plugin", - provider="provider-a", - model="gpt-test", - credentials={"api_key": "key"}, - prompt_messages=[], - plugin_unique_identifier=PLUGIN_IDENTIFIER, - ) - - def test_get_llm_num_tokens(self, mocker: MockerFixture): - client = PluginModelClient() - mocker.patch.object( - client, - "_request_with_plugin_daemon_response_stream", - return_value=iter([SimpleNamespace(num_tokens=42)]), - ) - - result = client.get_llm_num_tokens( - tenant_id="tenant-1", - user_id="user-1", - plugin_id="org/plugin", - provider="provider-a", - model_type="llm", - model="gpt-test", - credentials={}, - prompt_messages=[], - tools=[], - plugin_unique_identifier=PLUGIN_IDENTIFIER, - ) - - assert result == 42 - - def test_get_llm_num_tokens_empty_returns_zero(self, mocker: MockerFixture): - client = PluginModelClient() - mocker.patch.object(client, "_request_with_plugin_daemon_response_stream", return_value=iter([])) - - assert ( - client.get_llm_num_tokens( - "tenant-1", - "user-1", - "org/plugin", - "provider-a", - "llm", - "gpt-test", - {}, - [], - plugin_unique_identifier=PLUGIN_IDENTIFIER, - ) - == 0 - ) - - def test_invoke_text_embedding(self, mocker: MockerFixture): - client = PluginModelClient() - embedding_result = SimpleNamespace(data=[[0.1, 0.2]]) - mocker.patch.object( - client, "_request_with_plugin_daemon_response_stream", return_value=iter([embedding_result]) - ) - - result = client.invoke_text_embedding( - tenant_id="tenant-1", - user_id="user-1", - plugin_id="org/plugin", - provider="provider-a", - model="embedding-a", - credentials={}, - texts=["hello"], - input_type="search_document", - plugin_unique_identifier=PLUGIN_IDENTIFIER, - ) - - assert result is embedding_result - - def test_invoke_text_embedding_empty_raises(self, mocker: MockerFixture): - client = PluginModelClient() - mocker.patch.object(client, "_request_with_plugin_daemon_response_stream", return_value=iter([])) - - with pytest.raises(ValueError, match="Failed to invoke text embedding"): - client.invoke_text_embedding( - "tenant-1", - "user-1", - "org/plugin", - "provider-a", - "embedding-a", - {}, - ["hello"], - "x", - plugin_unique_identifier=PLUGIN_IDENTIFIER, - ) - - def test_invoke_multimodal_embedding(self, mocker: MockerFixture): - client = PluginModelClient() - embedding_result = SimpleNamespace(data=[[0.3, 0.4]]) - mocker.patch.object( - client, "_request_with_plugin_daemon_response_stream", return_value=iter([embedding_result]) - ) - - result = client.invoke_multimodal_embedding( - tenant_id="tenant-1", - user_id="user-1", - plugin_id="org/plugin", - provider="provider-a", - model="embedding-a", - credentials={}, - documents=[{"type": "image", "value": "abc"}], - input_type="search_document", - plugin_unique_identifier=PLUGIN_IDENTIFIER, - ) - - assert result is embedding_result - - def test_invoke_multimodal_embedding_empty_raises(self, mocker: MockerFixture): - client = PluginModelClient() - mocker.patch.object(client, "_request_with_plugin_daemon_response_stream", return_value=iter([])) - - with pytest.raises(ValueError, match="Failed to invoke file embedding"): - client.invoke_multimodal_embedding( - "tenant-1", - "user-1", - "org/plugin", - "provider-a", - "embedding-a", - {}, - [{"type": "image"}], - "x", - plugin_unique_identifier=PLUGIN_IDENTIFIER, - ) - - def test_get_text_embedding_num_tokens(self, mocker: MockerFixture): - client = PluginModelClient() - mocker.patch.object( - client, - "_request_with_plugin_daemon_response_stream", - return_value=iter([SimpleNamespace(num_tokens=[1, 2, 3])]), - ) - - assert client.get_text_embedding_num_tokens( - "tenant-1", - "user-1", - "org/plugin", - "provider-a", - "embedding-a", - {}, - ["a"], - plugin_unique_identifier=PLUGIN_IDENTIFIER, - ) == [ - 1, - 2, - 3, - ] - - def test_get_text_embedding_num_tokens_empty_returns_list(self, mocker: MockerFixture): - client = PluginModelClient() - mocker.patch.object(client, "_request_with_plugin_daemon_response_stream", return_value=iter([])) - - assert ( - client.get_text_embedding_num_tokens( - "tenant-1", - "user-1", - "org/plugin", - "provider-a", - "embedding-a", - {}, - ["a"], - plugin_unique_identifier=PLUGIN_IDENTIFIER, - ) - == [] - ) - - def test_invoke_rerank(self, mocker: MockerFixture): - client = PluginModelClient() - rerank_result = SimpleNamespace(scores=[0.9]) - mocker.patch.object(client, "_request_with_plugin_daemon_response_stream", return_value=iter([rerank_result])) - - result = client.invoke_rerank( - tenant_id="tenant-1", - user_id="user-1", - plugin_id="org/plugin", - provider="provider-a", - model="rerank-a", - credentials={}, - query="q", - docs=["doc-1"], - score_threshold=0.2, - top_n=5, - plugin_unique_identifier=PLUGIN_IDENTIFIER, - ) - - assert result is rerank_result - - def test_invoke_rerank_empty_raises(self, mocker: MockerFixture): - client = PluginModelClient() - mocker.patch.object(client, "_request_with_plugin_daemon_response_stream", return_value=iter([])) - - with pytest.raises(ValueError, match="Failed to invoke rerank"): - client.invoke_rerank( - "tenant-1", - "user-1", - "org/plugin", - "provider-a", - "rerank-a", - {}, - "q", - ["doc-1"], - plugin_unique_identifier=PLUGIN_IDENTIFIER, - ) - - def test_invoke_multimodal_rerank(self, mocker: MockerFixture): - client = PluginModelClient() - rerank_result = SimpleNamespace(scores=[0.8]) - mocker.patch.object(client, "_request_with_plugin_daemon_response_stream", return_value=iter([rerank_result])) - - result = client.invoke_multimodal_rerank( - tenant_id="tenant-1", - user_id="user-1", - plugin_id="org/plugin", - provider="provider-a", - model="rerank-a", - credentials={}, - query={"type": "text", "value": "q"}, - docs=[{"type": "image", "value": "doc"}], - score_threshold=0.1, - top_n=3, - plugin_unique_identifier=PLUGIN_IDENTIFIER, - ) - - assert result is rerank_result - - def test_invoke_multimodal_rerank_empty_raises(self, mocker: MockerFixture): - client = PluginModelClient() - mocker.patch.object(client, "_request_with_plugin_daemon_response_stream", return_value=iter([])) - - with pytest.raises(ValueError, match="Failed to invoke multimodal rerank"): - client.invoke_multimodal_rerank( - "tenant-1", - "user-1", - "org/plugin", - "provider-a", - "rerank-a", - {}, - {"type": "text"}, - [{"type": "image"}], - plugin_unique_identifier=PLUGIN_IDENTIFIER, - ) - - def test_invoke_tts(self, mocker: MockerFixture): - client = PluginModelClient() - mocker.patch.object( - client, - "_request_with_plugin_daemon_response_stream", - return_value=iter( - [ - SimpleNamespace(result="68656c6c6f", mime_type="audio/wav"), - SimpleNamespace(result="21", mime_type="audio/wav"), - ] - ), - ) - - result = list( - client.invoke_tts( - tenant_id="tenant-1", - user_id="user-1", - plugin_id="org/plugin", - provider="provider-a", - model="tts-a", - credentials={}, - content_text="hello", - voice="alloy", - plugin_unique_identifier=PLUGIN_IDENTIFIER, - ) - ) - - assert result == [b"hello", b"!"] - assert [chunk.mime_type for chunk in result] == ["audio/wav", "audio/wav"] - - def test_invoke_tts_wraps_plugin_daemon_inner_error(self, mocker: MockerFixture): - client = PluginModelClient() - - def _boom(): - raise PluginDaemonInnerError(code=-400, message="tts error") - yield # pragma: no cover - - mocker.patch.object(client, "_request_with_plugin_daemon_response_stream", return_value=_boom()) - - with pytest.raises(ValueError, match="tts error-400"): - list( - client.invoke_tts( - "tenant-1", - "user-1", - "org/plugin", - "provider-a", - "tts-a", - {}, - "hello", - "alloy", - plugin_unique_identifier=PLUGIN_IDENTIFIER, - ) - ) - - def test_get_tts_model_voices(self, mocker: MockerFixture): - client = PluginModelClient() - mocker.patch.object( - client, - "_request_with_plugin_daemon_response_stream", - return_value=iter( - [ - SimpleNamespace( - voices=[ - SimpleNamespace(name="Alloy", value="alloy"), - SimpleNamespace(name="Echo", value="echo"), - ] - ) - ] - ), - ) - - result = client.get_tts_model_voices( - tenant_id="tenant-1", - user_id="user-1", - plugin_id="org/plugin", - provider="provider-a", - model="tts-a", - credentials={}, - language="en", - plugin_unique_identifier=PLUGIN_IDENTIFIER, - ) - - assert result == [{"name": "Alloy", "value": "alloy"}, {"name": "Echo", "value": "echo"}] - - def test_get_tts_model_voices_empty_returns_list(self, mocker: MockerFixture): - client = PluginModelClient() - mocker.patch.object(client, "_request_with_plugin_daemon_response_stream", return_value=iter([])) - - assert ( - client.get_tts_model_voices( - "tenant-1", - "user-1", - "org/plugin", - "provider-a", - "tts-a", - {}, - plugin_unique_identifier=PLUGIN_IDENTIFIER, - ) - == [] - ) - - def test_invoke_speech_to_text(self, mocker: MockerFixture): - client = PluginModelClient() - stream_mock = mocker.patch.object( - client, - "_request_with_plugin_daemon_response_stream", - return_value=iter([PluginStringResultResponse(result="transcribed text")]), - ) - - result = client.invoke_speech_to_text( - tenant_id="tenant-1", - user_id="user-1", - plugin_id="org/plugin", - provider="provider-a", - model="stt-a", - credentials={}, - file=io.BytesIO(b"abc"), - plugin_unique_identifier=PLUGIN_IDENTIFIER, - ) - - assert result == "transcribed text" - assert stream_mock.call_args.kwargs["data"]["data"]["file"] == "616263" - - def test_invoke_speech_to_text_empty_raises(self, mocker: MockerFixture): - client = PluginModelClient() - mocker.patch.object(client, "_request_with_plugin_daemon_response_stream", return_value=iter([])) - - with pytest.raises(ValueError, match="Failed to invoke speech to text"): - client.invoke_speech_to_text( - "tenant-1", - "user-1", - "org/plugin", - "provider-a", - "stt-a", - {}, - io.BytesIO(b"abc"), - plugin_unique_identifier=PLUGIN_IDENTIFIER, - ) - - def test_invoke_moderation(self, mocker: MockerFixture): - client = PluginModelClient() - stream_mock = mocker.patch.object( - client, - "_request_with_plugin_daemon_response_stream", - return_value=iter([SimpleNamespace(result=True)]), - ) - - result = client.invoke_moderation( - tenant_id="tenant-1", - user_id="user-1", - plugin_id="org/plugin", - provider="provider-a", - model="moderation-a", - credentials={}, - text="safe text", - plugin_unique_identifier=PLUGIN_IDENTIFIER, - ) - - assert result is True - assert stream_mock.call_args.kwargs["path"] == "plugin/tenant-1/dispatch/moderation/invoke" - - def test_invoke_moderation_empty_raises(self, mocker: MockerFixture): - client = PluginModelClient() - mocker.patch.object(client, "_request_with_plugin_daemon_response_stream", return_value=iter([])) - - with pytest.raises(ValueError, match="Failed to invoke moderation"): - client.invoke_moderation( - "tenant-1", - "user-1", - "org/plugin", - "provider-a", - "moderation-a", - {}, - "unsafe", - plugin_unique_identifier=PLUGIN_IDENTIFIER, - ) diff --git a/api/tests/unit_tests/core/plugin/test_model_runtime_adapter.py b/api/tests/unit_tests/core/plugin/impl/test_model_runtime.py similarity index 100% rename from api/tests/unit_tests/core/plugin/test_model_runtime_adapter.py rename to api/tests/unit_tests/core/plugin/impl/test_model_runtime.py diff --git a/api/tests/unit_tests/core/plugin/impl/test_model_runtime_factory.py b/api/tests/unit_tests/core/plugin/impl/test_model_runtime_factory.py index 4111b1e9ae0..f96a21db858 100644 --- a/api/tests/unit_tests/core/plugin/impl/test_model_runtime_factory.py +++ b/api/tests/unit_tests/core/plugin/impl/test_model_runtime_factory.py @@ -1,12 +1,437 @@ -from unittest.mock import patch +from unittest.mock import Mock, patch + +import pytest from core.model_manager import ModelManager -from core.plugin.impl.model_runtime_factory import create_plugin_model_assembly, create_plugin_model_runtime +from core.plugin.impl.model_runtime_factory import ( + create_model_type_instance, + create_plugin_model_assembly, + create_plugin_model_runtime, +) from core.plugin.plugin_service import PluginService from core.provider_manager import ProviderManager +from graphon.model_runtime.entities.common_entities import I18nObject +from graphon.model_runtime.entities.model_entities import AIModelEntity, FetchFrom, ModelType +from graphon.model_runtime.entities.provider_entities import ( + ConfigurateMethod, + CredentialFormSchema, + FieldModelSchema, + FormType, + ModelCredentialSchema, + ProviderCredentialSchema, + ProviderEntity, +) +from graphon.model_runtime.model_providers.base.large_language_model import LargeLanguageModel +from graphon.model_runtime.model_providers.base.moderation_model import ModerationModel +from graphon.model_runtime.model_providers.base.rerank_model import RerankModel +from graphon.model_runtime.model_providers.base.speech2text_model import Speech2TextModel +from graphon.model_runtime.model_providers.base.text_embedding_model import TextEmbeddingModel +from graphon.model_runtime.model_providers.base.tts_model import TTSModel from graphon.model_runtime.model_providers.model_provider_factory import ModelProviderFactory +def _build_model(model: str, model_type: ModelType) -> AIModelEntity: + return AIModelEntity( + model=model, + label=I18nObject(en_US=model), + model_type=model_type, + fetch_from=FetchFrom.PREDEFINED_MODEL, + model_properties={}, + ) + + +def _build_provider( + *, + provider: str, + provider_name: str, + supported_model_types: list[ModelType], + models: list[AIModelEntity] | None = None, + provider_credential_schema: ProviderCredentialSchema | None = None, + model_credential_schema: ModelCredentialSchema | None = None, +) -> ProviderEntity: + return ProviderEntity( + provider=provider, + provider_name=provider_name, + label=I18nObject(en_US=provider_name or provider), + supported_model_types=supported_model_types, + configurate_methods=[ConfigurateMethod.PREDEFINED_MODEL], + models=models or [], + provider_credential_schema=provider_credential_schema, + model_credential_schema=model_credential_schema, + ) + + +class _FakeModelRuntime: + def __init__(self, providers: list[ProviderEntity]) -> None: + self._providers = providers + self.validate_provider_credentials = Mock() + self.validate_model_credentials = Mock() + self.get_model_schema = Mock() + self.get_provider_icon = Mock() + + def fetch_model_providers(self) -> list[ProviderEntity]: + return self._providers + + +def test_model_provider_factory_resolves_runtime_provider_name() -> None: + provider = ProviderEntity( + provider="langgenius/openai/openai", + provider_name="openai", + label=I18nObject(en_US="OpenAI"), + supported_model_types=[ModelType.LLM], + configurate_methods=[ConfigurateMethod.PREDEFINED_MODEL], + ) + factory = ModelProviderFactory(runtime=_FakeModelRuntime([provider])) + + provider_schema = factory.get_model_provider("openai") + + assert provider_schema.provider == "langgenius/openai/openai" + assert provider_schema.provider_name == "openai" + + +def test_model_provider_factory_resolves_canonical_short_name_independent_of_provider_order() -> None: + providers = [ + ProviderEntity( + provider="acme/openai/openai", + provider_name="", + label=I18nObject(en_US="Acme OpenAI"), + supported_model_types=[ModelType.LLM], + configurate_methods=[ConfigurateMethod.PREDEFINED_MODEL], + ), + ProviderEntity( + provider="langgenius/openai/openai", + provider_name="openai", + label=I18nObject(en_US="OpenAI"), + supported_model_types=[ModelType.LLM], + configurate_methods=[ConfigurateMethod.PREDEFINED_MODEL], + ), + ] + factory = ModelProviderFactory(runtime=_FakeModelRuntime(providers)) + + provider_schema = factory.get_model_provider("openai") + + assert provider_schema.provider == "langgenius/openai/openai" + assert provider_schema.provider_name == "openai" + + +def test_model_provider_factory_requires_runtime() -> None: + with pytest.raises(ValueError, match="runtime is required"): + ModelProviderFactory(runtime=None) # type: ignore[arg-type] + + +def test_model_provider_factory_get_providers_returns_runtime_providers() -> None: + providers = [ + _build_provider( + provider="langgenius/openai/openai", + provider_name="openai", + supported_model_types=[ModelType.LLM], + ) + ] + factory = ModelProviderFactory(runtime=_FakeModelRuntime(providers)) + + result = factory.get_providers() + + assert list(result) == providers + assert result is not providers + + +def test_model_provider_factory_get_provider_schema_delegates_to_provider_lookup() -> None: + provider = _build_provider( + provider="langgenius/openai/openai", + provider_name="openai", + supported_model_types=[ModelType.LLM], + ) + factory = ModelProviderFactory(runtime=_FakeModelRuntime([provider])) + + result = factory.get_provider_schema("openai") + + assert result is provider + + +def test_model_provider_factory_raises_for_unknown_provider() -> None: + factory = ModelProviderFactory( + runtime=_FakeModelRuntime( + [ + _build_provider( + provider="langgenius/openai/openai", + provider_name="openai", + supported_model_types=[ModelType.LLM], + ) + ] + ) + ) + + with pytest.raises(ValueError, match="Invalid provider: anthropic"): + factory.get_model_provider("anthropic") + + +def test_model_provider_factory_get_models_filters_provider_and_model_type() -> None: + providers = [ + _build_provider( + provider="langgenius/openai/openai", + provider_name="openai", + supported_model_types=[ModelType.LLM, ModelType.TTS], + models=[_build_model("gpt-4o-mini", ModelType.LLM), _build_model("tts-1", ModelType.TTS)], + ), + _build_provider( + provider="langgenius/cohere/cohere", + provider_name="cohere", + supported_model_types=[ModelType.RERANK], + models=[_build_model("rerank-v3", ModelType.RERANK)], + ), + ] + factory = ModelProviderFactory(runtime=_FakeModelRuntime(providers)) + + results = factory.get_models(provider="openai", model_type=ModelType.LLM) + + assert len(results) == 1 + assert results[0].provider == "langgenius/openai/openai" + assert [model.model for model in results[0].models] == ["gpt-4o-mini"] + + +def test_model_provider_factory_get_models_skips_providers_without_requested_model_type() -> None: + providers = [ + _build_provider( + provider="langgenius/openai/openai", + provider_name="openai", + supported_model_types=[ModelType.LLM], + models=[_build_model("gpt-4o-mini", ModelType.LLM)], + ), + _build_provider( + provider="langgenius/elevenlabs/elevenlabs", + provider_name="elevenlabs", + supported_model_types=[ModelType.TTS], + models=[_build_model("eleven_multilingual_v2", ModelType.TTS)], + ), + ] + factory = ModelProviderFactory(runtime=_FakeModelRuntime(providers)) + + results = factory.get_models(model_type=ModelType.TTS) + + assert len(results) == 1 + assert results[0].provider == "langgenius/elevenlabs/elevenlabs" + assert [model.model for model in results[0].models] == ["eleven_multilingual_v2"] + + +def test_model_provider_factory_get_models_without_model_type_keeps_all_provider_models() -> None: + providers = [ + _build_provider( + provider="langgenius/openai/openai", + provider_name="openai", + supported_model_types=[ModelType.LLM, ModelType.TTS], + models=[_build_model("gpt-4o-mini", ModelType.LLM), _build_model("tts-1", ModelType.TTS)], + ) + ] + factory = ModelProviderFactory(runtime=_FakeModelRuntime(providers)) + + results = factory.get_models(provider="openai") + + assert len(results) == 1 + assert [model.model for model in results[0].models] == ["gpt-4o-mini", "tts-1"] + + +def test_model_provider_factory_validates_provider_credentials() -> None: + runtime = _FakeModelRuntime( + [ + _build_provider( + provider="langgenius/openai/openai", + provider_name="openai", + supported_model_types=[ModelType.LLM], + provider_credential_schema=ProviderCredentialSchema( + credential_form_schemas=[ + CredentialFormSchema( + variable="api_key", + label=I18nObject(en_US="API key"), + type=FormType.SECRET_INPUT, + required=True, + ) + ] + ), + ) + ] + ) + factory = ModelProviderFactory(runtime=runtime) + + filtered = factory.provider_credentials_validate( + provider="openai", + credentials={"api_key": "secret", "ignored": "value"}, + ) + + assert filtered == {"api_key": "secret"} + runtime.validate_provider_credentials.assert_called_once_with( + provider="langgenius/openai/openai", + credentials={"api_key": "secret"}, + ) + + +def test_model_provider_factory_provider_credentials_validate_requires_schema() -> None: + factory = ModelProviderFactory( + runtime=_FakeModelRuntime( + [ + _build_provider( + provider="langgenius/openai/openai", + provider_name="openai", + supported_model_types=[ModelType.LLM], + ) + ] + ) + ) + + with pytest.raises(ValueError, match="Provider openai does not have provider_credential_schema"): + factory.provider_credentials_validate(provider="openai", credentials={"api_key": "secret"}) + + +def test_model_provider_factory_validates_model_credentials() -> None: + runtime = _FakeModelRuntime( + [ + _build_provider( + provider="langgenius/openai/openai", + provider_name="openai", + supported_model_types=[ModelType.LLM], + model_credential_schema=ModelCredentialSchema( + model=FieldModelSchema(label=I18nObject(en_US="Model")), + credential_form_schemas=[ + CredentialFormSchema( + variable="api_key", + label=I18nObject(en_US="API key"), + type=FormType.SECRET_INPUT, + required=True, + ) + ], + ), + ) + ] + ) + factory = ModelProviderFactory(runtime=runtime) + + filtered = factory.model_credentials_validate( + provider="openai", + model_type=ModelType.LLM, + model="gpt-4o-mini", + credentials={"api_key": "secret", "ignored": "value"}, + ) + + assert filtered == {"api_key": "secret"} + runtime.validate_model_credentials.assert_called_once_with( + provider="langgenius/openai/openai", + model_type=ModelType.LLM, + model="gpt-4o-mini", + credentials={"api_key": "secret"}, + ) + + +def test_model_provider_factory_model_credentials_validate_requires_schema() -> None: + factory = ModelProviderFactory( + runtime=_FakeModelRuntime( + [ + _build_provider( + provider="langgenius/openai/openai", + provider_name="openai", + supported_model_types=[ModelType.LLM], + ) + ] + ) + ) + + with pytest.raises(ValueError, match="Provider openai does not have model_credential_schema"): + factory.model_credentials_validate( + provider="openai", + model_type=ModelType.LLM, + model="gpt-4o-mini", + credentials={"api_key": "secret"}, + ) + + +def test_model_provider_factory_get_model_schema_and_icon_use_canonical_provider() -> None: + runtime = _FakeModelRuntime( + [ + _build_provider( + provider="langgenius/openai/openai", + provider_name="openai", + supported_model_types=[ModelType.LLM], + ) + ] + ) + runtime.get_model_schema.return_value = "schema" + runtime.get_provider_icon.return_value = (b"icon", "image/png") + factory = ModelProviderFactory(runtime=runtime) + + assert ( + factory.get_model_schema( + provider="openai", + model_type=ModelType.LLM, + model="gpt-4o-mini", + credentials=None, + ) + == "schema" + ) + assert factory.get_provider_icon("openai", "icon_small", "en_US") == (b"icon", "image/png") + runtime.get_model_schema.assert_called_once_with( + provider="langgenius/openai/openai", + model_type=ModelType.LLM, + model="gpt-4o-mini", + credentials={}, + ) + runtime.get_provider_icon.assert_called_once_with( + provider="langgenius/openai/openai", + icon_type="icon_small", + lang="en_US", + ) + + +@pytest.mark.parametrize( + ("model_type", "expected_type"), + [ + (ModelType.LLM, LargeLanguageModel), + (ModelType.TEXT_EMBEDDING, TextEmbeddingModel), + (ModelType.RERANK, RerankModel), + (ModelType.SPEECH2TEXT, Speech2TextModel), + (ModelType.MODERATION, ModerationModel), + (ModelType.TTS, TTSModel), + ], +) +def test_create_model_type_instance_builds_model_wrappers( + model_type: ModelType, + expected_type: type[object], +) -> None: + runtime = _FakeModelRuntime( + [ + _build_provider( + provider="langgenius/openai/openai", + provider_name="openai", + supported_model_types=[model_type], + ) + ] + ) + + instance = create_model_type_instance( + runtime=runtime, + provider_schema=runtime.fetch_model_providers()[0], + model_type=model_type, + ) + + assert isinstance(instance, expected_type) + + +def test_create_model_type_instance_rejects_unsupported_model_type() -> None: + runtime = _FakeModelRuntime( + [ + _build_provider( + provider="langgenius/openai/openai", + provider_name="openai", + supported_model_types=[ModelType.LLM], + ) + ] + ) + + with pytest.raises(ValueError, match="Unsupported model type: unsupported"): + create_model_type_instance( + runtime=runtime, + provider_schema=runtime.fetch_model_providers()[0], + model_type="unsupported", # type: ignore[arg-type] + ) + + def test_plugin_model_assembly_reuses_single_runtime_across_views(): runtime = create_plugin_model_runtime(tenant_id="tenant-1", user_id="user-1") provider_factory = ModelProviderFactory(runtime=runtime) diff --git a/api/tests/unit_tests/core/plugin/impl/test_oauth_handler.py b/api/tests/unit_tests/core/plugin/impl/test_oauth.py similarity index 66% rename from api/tests/unit_tests/core/plugin/impl/test_oauth_handler.py rename to api/tests/unit_tests/core/plugin/impl/test_oauth.py index 71c4d1b8a8e..17a05bad0d2 100644 --- a/api/tests/unit_tests/core/plugin/impl/test_oauth_handler.py +++ b/api/tests/unit_tests/core/plugin/impl/test_oauth.py @@ -1,9 +1,12 @@ +import json from io import BytesIO from types import SimpleNamespace import pytest from pytest_mock import MockerFixture from werkzeug import Request +from werkzeug.datastructures import Headers +from werkzeug.test import EnvironBuilder from core.plugin.impl.oauth import OAuthHandler @@ -146,3 +149,71 @@ class TestOAuthHandler: assert raw.startswith(b"POST /oauth/callback?code=123 HTTP/1.1\r\n") assert b"X-Test: yes\r\n" in raw assert raw.endswith(b"body-data") + + +def test_oauth_convert_request_to_raw_data() -> None: + oauth_handler = OAuthHandler() + builder = EnvironBuilder( + method="GET", + path="/test", + headers=Headers({"Content-Type": "application/json"}), + ) + request = Request(builder.get_environ()) + raw_request_bytes = oauth_handler._convert_request_to_raw_data(request) + + assert b"GET /test? HTTP/1.1" in raw_request_bytes + assert b"Content-Type: application/json" in raw_request_bytes + assert b"\r\n\r\n" in raw_request_bytes + + +def test_oauth_convert_request_to_raw_data_with_query_params() -> None: + oauth_handler = OAuthHandler() + builder = EnvironBuilder( + method="GET", + path="/test", + query_string="code=abc123&state=xyz789", + headers=Headers({"Content-Type": "application/json"}), + ) + request = Request(builder.get_environ()) + raw_request_bytes = oauth_handler._convert_request_to_raw_data(request) + + assert b"GET /test?code=abc123&state=xyz789 HTTP/1.1" in raw_request_bytes + assert b"Content-Type: application/json" in raw_request_bytes + assert b"\r\n\r\n" in raw_request_bytes + + +def test_oauth_convert_request_to_raw_data_with_post_body() -> None: + oauth_handler = OAuthHandler() + builder = EnvironBuilder( + method="POST", + path="/test", + data="param1=value1¶m2=value2", + headers=Headers({"Content-Type": "application/x-www-form-urlencoded"}), + ) + request = Request(builder.get_environ()) + raw_request_bytes = oauth_handler._convert_request_to_raw_data(request) + + assert b"POST /test? HTTP/1.1" in raw_request_bytes + assert b"Content-Type: application/x-www-form-urlencoded" in raw_request_bytes + assert b"\r\n\r\n" in raw_request_bytes + assert b"param1=value1¶m2=value2" in raw_request_bytes + + +def test_oauth_convert_request_to_raw_data_with_json_body() -> None: + oauth_handler = OAuthHandler() + json_data = {"code": "abc123", "state": "xyz789", "grant_type": "authorization_code"} + builder = EnvironBuilder( + method="POST", + path="/test", + data=json.dumps(json_data), + headers=Headers({"Content-Type": "application/json"}), + ) + request = Request(builder.get_environ()) + raw_request_bytes = oauth_handler._convert_request_to_raw_data(request) + + assert b"POST /test? HTTP/1.1" in raw_request_bytes + assert b"Content-Type: application/json" in raw_request_bytes + assert b"\r\n\r\n" in raw_request_bytes + assert b'"code": "abc123"' in raw_request_bytes + assert b'"state": "xyz789"' in raw_request_bytes + assert b'"grant_type": "authorization_code"' in raw_request_bytes diff --git a/api/tests/unit_tests/core/plugin/impl/test_plugin.py b/api/tests/unit_tests/core/plugin/impl/test_plugin.py new file mode 100644 index 00000000000..6718b370967 --- /dev/null +++ b/api/tests/unit_tests/core/plugin/impl/test_plugin.py @@ -0,0 +1,3274 @@ +import datetime +import json +from typing import Any +from unittest.mock import patch + +import httpx +import pytest +from packaging.version import Version +from pydantic import BaseModel +from requests import HTTPError + +from core.plugin.entities.bundle import PluginBundleDependency, PluginBundleDependencyType +from core.plugin.entities.plugin import ( + MissingPluginDependency, + PluginCategory, + PluginDeclaration, + PluginEntity, + PluginInstallation, + PluginInstallationSource, + PluginResourceRequirements, +) +from core.plugin.entities.plugin_daemon import ( + CredentialType, + PluginDaemonInnerError, + PluginDecodeResponse, + PluginInstalledIdsDaemonResponse, + PluginInstallTask, + PluginInstallTaskStartResponse, + PluginInstallTaskStatus, + PluginListResponse, + PluginListWithoutTotalResponse, + PluginReadmeResponse, + PluginVerification, +) +from core.plugin.impl.base import BasePluginClient, use_plugin_daemon_request_timeout +from core.plugin.impl.exc import ( + PluginDaemonBadRequestError, + PluginDaemonClientSideError, + PluginDaemonInternalServerError, + PluginDaemonNotFoundError, + PluginDaemonUnauthorizedError, + PluginInvokeError, + PluginNotFoundError, + PluginPermissionDeniedError, + PluginUniqueIdentifierError, +) +from core.plugin.impl.plugin import PluginInstaller +from core.plugin.impl.tool import PluginToolManager +from core.tools.entities.common_entities import I18nObject +from graphon.model_runtime.errors.invoke import ( + InvokeAuthorizationError, + InvokeBadRequestError, + InvokeConnectionError, + InvokeRateLimitError, + InvokeServerUnavailableError, +) +from graphon.model_runtime.errors.validate import CredentialsValidateFailedError +from models.provider_ids import GenericProviderID + +""" +Unit tests for Plugin Manager (PluginInstaller). + +This module tests the plugin management functionality including: +- Plugin discovery and listing +- Plugin loading and installation +- Plugin validation and manifest parsing +- Version compatibility checks +- Dependency resolution +""" + + +class TestPluginManagerCases: + class TestPluginDiscovery: + """Test plugin discovery functionality.""" + + @pytest.fixture + def plugin_installer(self): + """Create a PluginInstaller instance for testing.""" + return PluginInstaller() + + @pytest.fixture + def mock_plugin_entity(self): + """Create a mock PluginEntity for testing.""" + return PluginEntity( + id="entity-123", + created_at=datetime.datetime(2023, 1, 1, 0, 0, 0), + updated_at=datetime.datetime(2023, 1, 1, 0, 0, 0), + tenant_id="test-tenant", + endpoints_setups=0, + endpoints_active=0, + runtime_type="remote", + source=PluginInstallationSource.Marketplace, + meta={}, + plugin_id="plugin-123", + plugin_unique_identifier="test-org/test-plugin/1.0.0", + version="1.0.0", + checksum="abc123", + name="Test Plugin", + installation_id="install-123", + declaration=PluginDeclaration( + version="1.0.0", + author="test-author", + name="test-plugin", + description=I18nObject(en_US="Test plugin description", zh_Hans="测试插件描述"), + icon="icon.png", + label=I18nObject(en_US="Test Plugin", zh_Hans="测试插件"), + category=PluginCategory.Tool, + created_at=datetime.datetime.now(), + resource=PluginResourceRequirements(memory=512, permission=None), + plugins=PluginDeclaration.Plugins(), + meta=PluginDeclaration.Meta(version="1.0.0"), + ), + ) + + def test_list_plugins_success(self, plugin_installer, mock_plugin_entity): + """Test successful plugin listing.""" + # Arrange: Mock the HTTP response for listing plugins + mock_response = PluginListResponse(list=[mock_plugin_entity], total=1) + + with patch.object( + plugin_installer, "_request_with_plugin_daemon_response", return_value=mock_response + ) as mock_request: + # Act: List plugins for a tenant + result = plugin_installer.list_plugins("test-tenant") + + # Assert: Verify the request was made correctly + mock_request.assert_called_once() + assert len(result) == 1 + assert result[0].plugin_id == "plugin-123" + assert result[0].name == "Test Plugin" + + def test_list_plugins_with_pagination(self, plugin_installer, mock_plugin_entity): + """Test plugin listing with pagination support.""" + # Arrange: Mock paginated response + mock_response = PluginListResponse(list=[mock_plugin_entity], total=10) + + with patch.object( + plugin_installer, "_request_with_plugin_daemon_response", return_value=mock_response + ) as mock_request: + # Act: List plugins with pagination + result = plugin_installer.list_plugins_with_total("test-tenant", page=1, page_size=5) + + # Assert: Verify pagination parameters + mock_request.assert_called_once() + call_args = mock_request.call_args + assert call_args[1]["params"]["page"] == 1 + assert call_args[1]["params"]["page_size"] == 5 + assert result.total == 10 + + def test_list_plugins_by_category(self, plugin_installer, mock_plugin_entity): + """Test category plugin listing without total.""" + mock_response = PluginListWithoutTotalResponse(list=[mock_plugin_entity], has_more=True) + + with patch.object( + plugin_installer, "_request_with_plugin_daemon_response", return_value=mock_response + ) as mock_request: + result = plugin_installer.list_plugins_by_category( + "test-tenant", + category=PluginCategory.Tool, + page=2, + page_size=10, + query="weather", + tags=["search", "rag"], + language="zh_Hans", + ) + + mock_request.assert_called_once() + call_args = mock_request.call_args + assert call_args.args[1] == "plugin/test-tenant/management/tool/list" + assert call_args.args[2] is PluginListWithoutTotalResponse + assert call_args.kwargs["params"]["page"] == 2 + assert call_args.kwargs["params"]["page_size"] == 10 + assert call_args.kwargs["params"]["query"] == "weather" + assert call_args.kwargs["params"]["tags"] == ["search", "rag"] + assert call_args.kwargs["params"]["language"] == "zh_Hans" + assert result.list == [mock_plugin_entity] + assert result.has_more is True + + def test_list_plugins_empty_result(self, plugin_installer): + """Test plugin listing when no plugins are installed.""" + # Arrange: Mock empty response + mock_response = PluginListResponse(list=[], total=0) + + with patch.object(plugin_installer, "_request_with_plugin_daemon_response", return_value=mock_response): + # Act: List plugins + result = plugin_installer.list_plugins("test-tenant") + + # Assert: Verify empty list is returned + assert len(result) == 0 + + def test_list_installed_plugin_ids(self, plugin_installer): + """The lightweight ID endpoint is unpaginated and does not request plugin details.""" + mock_response = PluginInstalledIdsDaemonResponse(plugin_ids=["langgenius/openai", "langgenius/anthropic"]) + + with patch.object( + plugin_installer, "_request_with_plugin_daemon_response", return_value=mock_response + ) as mock_request: + result = plugin_installer.list_installed_plugin_ids("test-tenant", PluginCategory.Tool) + + mock_request.assert_called_once_with( + "GET", + "plugin/test-tenant/management/installation/ids", + PluginInstalledIdsDaemonResponse, + params={"category": "tool"}, + ) + assert result == ["langgenius/openai", "langgenius/anthropic"] + + def test_fetch_plugin_by_identifier_found(self, plugin_installer): + """Test fetching a plugin by its unique identifier when it exists.""" + # Arrange: Mock successful fetch + with patch.object( + plugin_installer, "_request_with_plugin_daemon_response", return_value=True + ) as mock_request: + # Act: Fetch plugin by identifier + result = plugin_installer.fetch_plugin_by_identifier("test-tenant", "test-org/test-plugin/1.0.0") + + # Assert: Verify the plugin was found + assert result is True + mock_request.assert_called_once() + + def test_fetch_plugin_by_identifier_not_found(self, plugin_installer): + """Test fetching a plugin by identifier when it doesn't exist.""" + # Arrange: Mock not found response + with patch.object(plugin_installer, "_request_with_plugin_daemon_response", return_value=False): + # Act: Fetch non-existent plugin + result = plugin_installer.fetch_plugin_by_identifier("test-tenant", "non-existent/plugin/1.0.0") + + # Assert: Verify the plugin was not found + assert result is False + + class TestPluginLoading: + """Test plugin loading and installation functionality.""" + + @pytest.fixture + def plugin_installer(self): + """Create a PluginInstaller instance for testing.""" + return PluginInstaller() + + @pytest.fixture + def mock_plugin_declaration(self): + """Create a mock PluginDeclaration for testing.""" + return PluginDeclaration( + version="1.0.0", + author="test-author", + name="test-plugin", + description=I18nObject(en_US="Test plugin", zh_Hans="测试插件"), + icon="icon.png", + label=I18nObject(en_US="Test Plugin", zh_Hans="测试插件"), + category=PluginCategory.Tool, + created_at=datetime.datetime.now(), + resource=PluginResourceRequirements(memory=512, permission=None), + plugins=PluginDeclaration.Plugins(), + meta=PluginDeclaration.Meta(version="1.0.0"), + ) + + def test_upload_pkg_success(self, plugin_installer, mock_plugin_declaration): + """Test successful plugin package upload.""" + # Arrange: Create mock package data and expected response + pkg_data = b"mock-plugin-package-data" + mock_response = PluginDecodeResponse( + unique_identifier="test-org/test-plugin/1.0.0", + manifest=mock_plugin_declaration, + verification=PluginVerification(authorized_category=PluginVerification.AuthorizedCategory.Community), + ) + + with patch.object( + plugin_installer, "_request_with_plugin_daemon_response", return_value=mock_response + ) as mock_request: + # Act: Upload plugin package + result = plugin_installer.upload_pkg("test-tenant", pkg_data, verify_signature=False) + + # Assert: Verify upload was successful + assert result.unique_identifier == "test-org/test-plugin/1.0.0" + assert result.manifest.name == "test-plugin" + mock_request.assert_called_once() + + def test_upload_pkg_with_signature_verification(self, plugin_installer, mock_plugin_declaration): + """Test plugin package upload with signature verification enabled.""" + # Arrange: Create mock package data + pkg_data = b"signed-plugin-package" + mock_response = PluginDecodeResponse( + unique_identifier="verified-org/verified-plugin/1.0.0", + manifest=mock_plugin_declaration, + verification=PluginVerification(authorized_category=PluginVerification.AuthorizedCategory.Partner), + ) + + with patch.object( + plugin_installer, "_request_with_plugin_daemon_response", return_value=mock_response + ) as mock_request: + # Act: Upload with signature verification + result = plugin_installer.upload_pkg("test-tenant", pkg_data, verify_signature=True) + + # Assert: Verify signature verification was requested + call_args = mock_request.call_args + assert call_args[1]["data"]["verify_signature"] == "true" + assert result.verification.authorized_category == PluginVerification.AuthorizedCategory.Partner + + def test_install_from_identifiers_success(self, plugin_installer): + """Test successful plugin installation from identifiers.""" + # Arrange: Mock installation response + mock_response = PluginInstallTaskStartResponse(all_installed=False, task_id="task-123") + + with patch.object( + plugin_installer, "_request_with_plugin_daemon_response", return_value=mock_response + ) as mock_request: + # Act: Install plugins from identifiers + result = plugin_installer.install_from_identifiers( + tenant_id="test-tenant", + identifiers=["plugin1/1.0.0", "plugin2/2.0.0"], + source=PluginInstallationSource.Marketplace, + metas=[{"key": "value1"}, {"key": "value2"}], + ) + + # Assert: Verify installation task was created + assert result.task_id == "task-123" + assert result.all_installed is False + mock_request.assert_called_once() + + def test_install_from_identifiers_all_installed(self, plugin_installer): + """Test installation when all plugins are already installed.""" + # Arrange: Mock response indicating all plugins are installed + mock_response = PluginInstallTaskStartResponse(all_installed=True, task_id="") + + with patch.object(plugin_installer, "_request_with_plugin_daemon_response", return_value=mock_response): + # Act: Attempt to install already-installed plugins + result = plugin_installer.install_from_identifiers( + tenant_id="test-tenant", + identifiers=["existing-plugin/1.0.0"], + source=PluginInstallationSource.Package, + metas=[{}], + ) + + # Assert: Verify all_installed flag is True + assert result.all_installed is True + + def test_fetch_plugin_installation_task(self, plugin_installer): + """Test fetching a specific plugin installation task.""" + # Arrange: Mock installation task + mock_task = PluginInstallTask( + id="task-123", + created_at=datetime.datetime.now(), + updated_at=datetime.datetime.now(), + status=PluginInstallTaskStatus.Running, + total_plugins=3, + completed_plugins=1, + plugins=[], + ) + + with patch.object( + plugin_installer, "_request_with_plugin_daemon_response", return_value=mock_task + ) as mock_request: + # Act: Fetch installation task + result = plugin_installer.fetch_plugin_installation_task("test-tenant", "task-123") + + # Assert: Verify task details + assert result.status == PluginInstallTaskStatus.Running + assert result.total_plugins == 3 + assert result.completed_plugins == 1 + mock_request.assert_called_once() + + def test_uninstall_plugin_success(self, plugin_installer): + """Test successful plugin uninstallation.""" + # Arrange: Mock successful uninstall + with patch.object( + plugin_installer, "_request_with_plugin_daemon_response", return_value=True + ) as mock_request: + # Act: Uninstall plugin + result = plugin_installer.uninstall("test-tenant", "install-123") + + # Assert: Verify uninstallation succeeded + assert result is True + mock_request.assert_called_once() + + def test_upgrade_plugin_success(self, plugin_installer): + """Test successful plugin upgrade.""" + # Arrange: Mock upgrade response + mock_response = PluginInstallTaskStartResponse(all_installed=False, task_id="upgrade-task-123") + + with patch.object( + plugin_installer, "_request_with_plugin_daemon_response", return_value=mock_response + ) as mock_request: + # Act: Upgrade plugin + result = plugin_installer.upgrade_plugin( + tenant_id="test-tenant", + original_plugin_unique_identifier="plugin/1.0.0", + new_plugin_unique_identifier="plugin/2.0.0", + source=PluginInstallationSource.Marketplace, + meta={"upgrade": "true"}, + ) + + # Assert: Verify upgrade task was created + assert result.task_id == "upgrade-task-123" + mock_request.assert_called_once() + + class TestPluginValidation: + """Test plugin validation and manifest parsing.""" + + @pytest.fixture + def plugin_installer(self): + """Create a PluginInstaller instance for testing.""" + return PluginInstaller() + + def test_fetch_plugin_manifest_success(self, plugin_installer): + """Test successful plugin manifest fetching.""" + # Arrange: Create a valid plugin declaration + mock_manifest = PluginDeclaration( + version="1.0.0", + author="test-author", + name="test-plugin", + description=I18nObject(en_US="Test plugin", zh_Hans="测试插件"), + icon="icon.png", + label=I18nObject(en_US="Test Plugin", zh_Hans="测试插件"), + category=PluginCategory.Tool, + created_at=datetime.datetime.now(), + resource=PluginResourceRequirements(memory=512, permission=None), + plugins=PluginDeclaration.Plugins(), + meta=PluginDeclaration.Meta(version="1.0.0", minimum_dify_version="0.6.0"), + ) + + with patch.object( + plugin_installer, "_request_with_plugin_daemon_response", return_value=mock_manifest + ) as mock_request: + # Act: Fetch plugin manifest + result = plugin_installer.fetch_plugin_manifest("test-tenant", "test-org/test-plugin/1.0.0") + + # Assert: Verify manifest was fetched correctly + assert result.name == "test-plugin" + assert result.version == "1.0.0" + assert result.author == "test-author" + assert result.meta.minimum_dify_version == "0.6.0" + mock_request.assert_called_once() + + def test_decode_plugin_from_identifier(self, plugin_installer): + """Test decoding plugin information from identifier.""" + # Arrange: Create mock decode response + mock_declaration = PluginDeclaration( + version="2.0.0", + author="decode-author", + name="decode-plugin", + description=I18nObject(en_US="Decoded plugin", zh_Hans="解码插件"), + icon="icon.png", + label=I18nObject(en_US="Decode Plugin", zh_Hans="解码插件"), + category=PluginCategory.Model, + created_at=datetime.datetime.now(), + resource=PluginResourceRequirements(memory=1024, permission=None), + plugins=PluginDeclaration.Plugins(), + meta=PluginDeclaration.Meta(version="2.0.0"), + ) + + mock_response = PluginDecodeResponse( + unique_identifier="org/decode-plugin/2.0.0", + manifest=mock_declaration, + verification=None, + ) + + with patch.object(plugin_installer, "_request_with_plugin_daemon_response", return_value=mock_response): + # Act: Decode plugin from identifier + result = plugin_installer.decode_plugin_from_identifier("test-tenant", "org/decode-plugin/2.0.0") + + # Assert: Verify decoded information + assert result.unique_identifier == "org/decode-plugin/2.0.0" + assert result.manifest.name == "decode-plugin" + # Category will be Extension unless a model provider entity is provided + assert result.manifest.category == PluginCategory.Extension + + def test_plugin_manifest_invalid_version_format(self): + """Test that invalid version format raises validation error.""" + # Arrange & Act & Assert: Creating a declaration with invalid version should fail + with pytest.raises(ValueError, match="Invalid version format"): + PluginDeclaration( + version="invalid-version", # Invalid version format + author="test-author", + name="test-plugin", + description=I18nObject(en_US="Test", zh_Hans="测试"), + icon="icon.png", + label=I18nObject(en_US="Test", zh_Hans="测试"), + category=PluginCategory.Tool, + created_at=datetime.datetime.now(), + resource=PluginResourceRequirements(memory=512, permission=None), + plugins=PluginDeclaration.Plugins(), + meta=PluginDeclaration.Meta(version="1.0.0"), + ) + + def test_plugin_manifest_invalid_author_format(self): + """Test that invalid author format raises validation error.""" + # Arrange & Act & Assert: Creating a declaration with invalid author should fail + with pytest.raises(ValueError, match="1 validation error for PluginDeclaration\\\nauthor"): + PluginDeclaration( + version="1.0.0", + author="invalid author with spaces!@#", # Invalid author format + name="test-plugin", + description=I18nObject(en_US="Test", zh_Hans="测试"), + icon="icon.png", + label=I18nObject(en_US="Test", zh_Hans="测试"), + category=PluginCategory.Tool, + created_at=datetime.datetime.now(), + resource=PluginResourceRequirements(memory=512, permission=None), + plugins=PluginDeclaration.Plugins(), + meta=PluginDeclaration.Meta(version="1.0.0"), + ) + + def test_plugin_manifest_invalid_name_format(self): + """Test that invalid plugin name format raises validation error.""" + # Arrange & Act & Assert: Creating a declaration with invalid name should fail + with pytest.raises(ValueError, match="1 validation error for PluginDeclaration\\\nname"): + PluginDeclaration( + version="1.0.0", + author="test-author", + name="Invalid_Plugin_Name_With_Uppercase", # Invalid name format + description=I18nObject(en_US="Test", zh_Hans="测试"), + icon="icon.png", + label=I18nObject(en_US="Test", zh_Hans="测试"), + category=PluginCategory.Tool, + created_at=datetime.datetime.now(), + resource=PluginResourceRequirements(memory=512, permission=None), + plugins=PluginDeclaration.Plugins(), + meta=PluginDeclaration.Meta(version="1.0.0"), + ) + + def test_fetch_plugin_readme_success(self, plugin_installer): + """Test successful plugin readme fetching.""" + # Arrange: Mock readme response + mock_response = PluginReadmeResponse(content="# Test Plugin\n\nThis is a test plugin.", language="en_US") + + with patch.object(plugin_installer, "_request_with_plugin_daemon_response", return_value=mock_response): + # Act: Fetch plugin readme + result = plugin_installer.fetch_plugin_readme("test-tenant", "test-org/test-plugin/1.0.0", "en_US") + + # Assert: Verify readme content + assert result == "# Test Plugin\n\nThis is a test plugin." + + def test_fetch_plugin_readme_not_found(self, plugin_installer): + """Test fetching readme when it doesn't exist (404 error).""" + # Arrange: Mock HTTP 404 error - the actual implementation catches HTTPError from requests library + mock_error = HTTPError("404 Not Found") + + with patch.object(plugin_installer, "_request_with_plugin_daemon_response", side_effect=mock_error): + # Act: Fetch non-existent readme + result = plugin_installer.fetch_plugin_readme("test-tenant", "test-org/test-plugin/1.0.0", "en_US") + + # Assert: Verify empty string is returned for 404 + assert result == "" + + class TestVersionCompatibility: + """Test version compatibility checks.""" + + def test_valid_version_format(self): + """Test that valid semantic versions are accepted.""" + # Arrange & Act: Create declarations with various valid version formats + valid_versions = ["1.0.0", "2.1.3", "0.0.1", "10.20.30"] + + for version in valid_versions: + # Assert: All valid versions should be accepted + declaration = PluginDeclaration( + version=version, + author="test-author", + name="test-plugin", + description=I18nObject(en_US="Test", zh_Hans="测试"), + icon="icon.png", + label=I18nObject(en_US="Test", zh_Hans="测试"), + category=PluginCategory.Tool, + created_at=datetime.datetime.now(), + resource=PluginResourceRequirements(memory=512, permission=None), + plugins=PluginDeclaration.Plugins(), + meta=PluginDeclaration.Meta(version=version), + ) + assert declaration.version == version + + def test_minimum_dify_version_validation(self): + """Test minimum Dify version validation.""" + # Arrange & Act: Create declaration with minimum Dify version + declaration = PluginDeclaration( + version="1.0.0", + author="test-author", + name="test-plugin", + description=I18nObject(en_US="Test", zh_Hans="测试"), + icon="icon.png", + label=I18nObject(en_US="Test", zh_Hans="测试"), + category=PluginCategory.Tool, + created_at=datetime.datetime.now(), + resource=PluginResourceRequirements(memory=512, permission=None), + plugins=PluginDeclaration.Plugins(), + meta=PluginDeclaration.Meta(version="1.0.0", minimum_dify_version="0.6.0"), + ) + + # Assert: Verify minimum version is set correctly + assert declaration.meta.minimum_dify_version == "0.6.0" + + def test_invalid_minimum_dify_version(self): + """Test that invalid minimum Dify version format raises error.""" + # Arrange & Act & Assert: Invalid minimum version should raise ValueError + with pytest.raises(ValueError, match="Invalid version format"): + PluginDeclaration.Meta(version="1.0.0", minimum_dify_version="invalid.version") + + def test_version_comparison_logic(self): + """Test version comparison using packaging.version.Version.""" + # Arrange: Create version objects for comparison + v1 = Version("1.0.0") + v2 = Version("2.0.0") + v3 = Version("1.5.0") + + # Act & Assert: Verify version comparison works correctly + assert v1 < v2 + assert v2 > v1 + assert v1 < v3 < v2 + assert v1 == Version("1.0.0") + + def test_plugin_upgrade_version_check(self): + """Test that plugin upgrade requires newer version.""" + # Arrange: Define old and new versions + old_version = Version("1.0.0") + new_version = Version("2.0.0") + same_version = Version("1.0.0") + + # Act & Assert: Verify version upgrade logic + assert new_version > old_version # Valid upgrade + assert not (same_version > old_version) # Invalid upgrade (same version) + + class TestDependencyResolution: + """Test plugin dependency resolution.""" + + @pytest.fixture + def plugin_installer(self): + """Create a PluginInstaller instance for testing.""" + return PluginInstaller() + + def test_upload_bundle_with_dependencies(self, plugin_installer): + """Test uploading a plugin bundle and extracting dependencies.""" + # Arrange: Create mock bundle data and dependencies + bundle_data = b"mock-bundle-data" + mock_dependencies = [ + PluginBundleDependency( + type=PluginBundleDependencyType.Marketplace, + value=PluginBundleDependency.Marketplace(organization="org1", plugin="plugin1", version="1.0.0"), + ), + PluginBundleDependency( + type=PluginBundleDependencyType.Github, + value=PluginBundleDependency.Github( + repo_address="https://github.com/org/repo", + repo="org/repo", + release="v1.0.0", + packages="plugin.zip", + ), + ), + ] + + with patch.object( + plugin_installer, "_request_with_plugin_daemon_response", return_value=mock_dependencies + ) as mock_request: + # Act: Upload bundle + result = plugin_installer.upload_bundle("test-tenant", bundle_data, verify_signature=False) + + # Assert: Verify dependencies were extracted + assert len(result) == 2 + assert result[0].type == PluginBundleDependencyType.Marketplace + assert result[1].type == PluginBundleDependencyType.Github + mock_request.assert_called_once() + + def test_fetch_missing_dependencies(self, plugin_installer): + """Test fetching missing dependencies for plugins.""" + # Arrange: Mock missing dependencies response + mock_missing = [ + MissingPluginDependency(plugin_unique_identifier="dep1/1.0.0", current_identifier=None), + MissingPluginDependency(plugin_unique_identifier="dep2/2.0.0", current_identifier="dep2/1.0.0"), + ] + + with patch.object( + plugin_installer, "_request_with_plugin_daemon_response", return_value=mock_missing + ) as mock_request: + # Act: Fetch missing dependencies + result = plugin_installer.fetch_missing_dependencies("test-tenant", ["plugin1/1.0.0", "plugin2/2.0.0"]) + + # Assert: Verify missing dependencies were identified + assert len(result) == 2 + assert result[0].plugin_unique_identifier == "dep1/1.0.0" + assert result[1].current_identifier == "dep2/1.0.0" + mock_request.assert_called_once() + + def test_fetch_missing_dependencies_none_missing(self, plugin_installer): + """Test fetching missing dependencies when all are satisfied.""" + # Arrange: Mock empty missing dependencies + with patch.object(plugin_installer, "_request_with_plugin_daemon_response", return_value=[]): + # Act: Fetch missing dependencies + result = plugin_installer.fetch_missing_dependencies("test-tenant", ["plugin1/1.0.0"]) + + # Assert: Verify no missing dependencies + assert len(result) == 0 + + def test_fetch_plugin_installation_by_ids(self, plugin_installer): + """Test fetching plugin installations by their IDs.""" + # Arrange: Create mock plugin installations + mock_installations = [ + PluginInstallation( + id="install-1", + created_at=datetime.datetime.now(), + updated_at=datetime.datetime.now(), + tenant_id="test-tenant", + endpoints_setups=0, + endpoints_active=0, + runtime_type="remote", + source=PluginInstallationSource.Marketplace, + meta={}, + plugin_id="plugin-1", + plugin_unique_identifier="org/plugin1/1.0.0", + version="1.0.0", + checksum="abc123", + declaration=PluginDeclaration( + version="1.0.0", + author="author1", + name="plugin1", + description=I18nObject(en_US="Plugin 1", zh_Hans="插件1"), + icon="icon.png", + label=I18nObject(en_US="Plugin 1", zh_Hans="插件1"), + category=PluginCategory.Tool, + created_at=datetime.datetime.now(), + resource=PluginResourceRequirements(memory=512, permission=None), + plugins=PluginDeclaration.Plugins(), + meta=PluginDeclaration.Meta(version="1.0.0"), + ), + ) + ] + + with patch.object( + plugin_installer, "_request_with_plugin_daemon_response", return_value=mock_installations + ) as mock_request: + # Act: Fetch installations by IDs + result = plugin_installer.fetch_plugin_installation_by_ids("test-tenant", ["plugin-1", "plugin-2"]) + + # Assert: Verify installations were fetched + assert len(result) == 1 + assert result[0].plugin_id == "plugin-1" + mock_request.assert_called_once() + + def test_dependency_chain_resolution(self, plugin_installer): + """Test resolving a chain of dependencies.""" + # Arrange: Create a dependency chain scenario + # Plugin A depends on Plugin B, Plugin B depends on Plugin C + mock_missing = [ + MissingPluginDependency(plugin_unique_identifier="plugin-b/1.0.0", current_identifier=None), + MissingPluginDependency(plugin_unique_identifier="plugin-c/1.0.0", current_identifier=None), + ] + + with patch.object(plugin_installer, "_request_with_plugin_daemon_response", return_value=mock_missing): + # Act: Fetch missing dependencies for Plugin A + result = plugin_installer.fetch_missing_dependencies("test-tenant", ["plugin-a/1.0.0"]) + + # Assert: Verify all dependencies in the chain are identified + assert len(result) == 2 + identifiers = [dep.plugin_unique_identifier for dep in result] + assert "plugin-b/1.0.0" in identifiers + assert "plugin-c/1.0.0" in identifiers + + def test_check_tools_existence(self, plugin_installer): + """Test checking if plugin tools exist.""" + # Arrange: Create provider IDs to check using the correct format + provider_ids = [ + GenericProviderID("org1/plugin1/provider1"), + GenericProviderID("org2/plugin2/provider2"), + ] + + # Mock response indicating first exists, second doesn't + mock_response = [True, False] + + with patch.object( + plugin_installer, "_request_with_plugin_daemon_response", return_value=mock_response + ) as mock_request: + # Act: Check tools existence + result = plugin_installer.check_tools_existence("test-tenant", provider_ids) + + # Assert: Verify existence check results + assert len(result) == 2 + assert result[0] is True + assert result[1] is False + mock_request.assert_called_once() + + class TestPluginTaskManagement: + """Test plugin installation task management.""" + + @pytest.fixture + def plugin_installer(self): + """Create a PluginInstaller instance for testing.""" + return PluginInstaller() + + def test_fetch_plugin_installation_tasks(self, plugin_installer): + """Test fetching multiple plugin installation tasks.""" + # Arrange: Create mock installation tasks + mock_tasks = [ + PluginInstallTask( + id="task-1", + created_at=datetime.datetime.now(), + updated_at=datetime.datetime.now(), + status=PluginInstallTaskStatus.Running, + total_plugins=2, + completed_plugins=1, + plugins=[], + ), + PluginInstallTask( + id="task-2", + created_at=datetime.datetime.now(), + updated_at=datetime.datetime.now(), + status=PluginInstallTaskStatus.Success, + total_plugins=1, + completed_plugins=1, + plugins=[], + ), + ] + + with patch.object( + plugin_installer, "_request_with_plugin_daemon_response", return_value=mock_tasks + ) as mock_request: + # Act: Fetch installation tasks + result = plugin_installer.fetch_plugin_installation_tasks("test-tenant", page=1, page_size=10) + + # Assert: Verify tasks were fetched + assert len(result) == 2 + assert result[0].status == PluginInstallTaskStatus.Running + assert result[1].status == PluginInstallTaskStatus.Success + mock_request.assert_called_once() + + def test_delete_plugin_installation_task(self, plugin_installer): + """Test deleting a specific plugin installation task.""" + # Arrange: Mock successful deletion + with patch.object( + plugin_installer, "_request_with_plugin_daemon_response", return_value=True + ) as mock_request: + # Act: Delete installation task + result = plugin_installer.delete_plugin_installation_task("test-tenant", "task-123") + + # Assert: Verify deletion succeeded + assert result is True + mock_request.assert_called_once() + + def test_delete_all_plugin_installation_task_items(self, plugin_installer): + """Test deleting all plugin installation task items.""" + # Arrange: Mock successful deletion of all items + with patch.object( + plugin_installer, "_request_with_plugin_daemon_response", return_value=True + ) as mock_request: + # Act: Delete all task items + result = plugin_installer.delete_all_plugin_installation_task_items("test-tenant") + + # Assert: Verify all items were deleted + assert result is True + mock_request.assert_called_once() + + def test_delete_plugin_installation_task_item(self, plugin_installer): + """Test deleting a specific item from an installation task.""" + # Arrange: Mock successful item deletion + with patch.object( + plugin_installer, "_request_with_plugin_daemon_response", return_value=True + ) as mock_request: + # Act: Delete specific task item + result = plugin_installer.delete_plugin_installation_task_item( + "test-tenant", "task-123", "plugin-identifier" + ) + + # Assert: Verify item was deleted + assert result is True + mock_request.assert_called_once() + + class TestErrorHandling: + """Test error handling in plugin manager.""" + + @pytest.fixture + def plugin_installer(self): + """Create a PluginInstaller instance for testing.""" + return PluginInstaller() + + def test_plugin_not_found_error(self, plugin_installer): + """Test handling of plugin not found error.""" + # Arrange: Mock plugin daemon not found error + with patch.object( + plugin_installer, + "_request_with_plugin_daemon_response", + side_effect=PluginDaemonNotFoundError("Plugin not found"), + ): + # Act & Assert: Verify error is raised + with pytest.raises(PluginDaemonNotFoundError): + plugin_installer.fetch_plugin_manifest("test-tenant", "non-existent/plugin/1.0.0") + + def test_plugin_bad_request_error(self, plugin_installer): + """Test handling of bad request error.""" + # Arrange: Mock bad request error + with patch.object( + plugin_installer, + "_request_with_plugin_daemon_response", + side_effect=PluginDaemonBadRequestError("Invalid request"), + ): + # Act & Assert: Verify error is raised + with pytest.raises(PluginDaemonBadRequestError): + plugin_installer.install_from_identifiers( + "test-tenant", [], PluginInstallationSource.Marketplace, [] + ) + + def test_plugin_internal_server_error(self, plugin_installer): + """Test handling of internal server error.""" + # Arrange: Mock internal server error + with patch.object( + plugin_installer, + "_request_with_plugin_daemon_response", + side_effect=PluginDaemonInternalServerError("Internal error"), + ): + # Act & Assert: Verify error is raised + with pytest.raises(PluginDaemonInternalServerError): + plugin_installer.list_plugins("test-tenant") + + def test_http_error_handling(self, plugin_installer): + """Test handling of HTTP errors during requests.""" + # Arrange: Mock HTTP error + with patch.object(plugin_installer, "_request", side_effect=httpx.RequestError("Connection failed")): + # Act & Assert: Verify appropriate error handling + with pytest.raises(httpx.RequestError): + plugin_installer._request("GET", "test/path") + + class TestPluginCategoryDetection: + """Test automatic plugin category detection.""" + + def test_category_defaults_to_extension_without_tool_provider(self): + """Test that plugins without tool providers default to Extension category.""" + # Arrange: Create declaration - category is auto-detected based on provider presence + # The model_validator in PluginDeclaration automatically sets category based on which provider is present + # Since we're not providing a tool provider entity, it defaults to Extension + # This test verifies that explicitly set categories are preserved + declaration = PluginDeclaration( + version="1.0.0", + author="test-author", + name="tool-plugin", + description=I18nObject(en_US="Tool plugin", zh_Hans="工具插件"), + icon="icon.png", + label=I18nObject(en_US="Tool Plugin", zh_Hans="工具插件"), + category=PluginCategory.Extension, # Will be Extension without a tool provider + created_at=datetime.datetime.now(), + resource=PluginResourceRequirements(memory=512, permission=None), + plugins=PluginDeclaration.Plugins(), + meta=PluginDeclaration.Meta(version="1.0.0"), + ) + + # Assert: Verify category defaults to Extension when no provider is specified + assert declaration.category == PluginCategory.Extension + + def test_category_defaults_to_extension_without_model_provider(self): + """Test that plugins without model providers default to Extension category.""" + # Arrange: Create declaration - without a model provider entity, defaults to Extension + # The category is auto-detected in the model_validator based on provider presence + declaration = PluginDeclaration( + version="1.0.0", + author="test-author", + name="model-plugin", + description=I18nObject(en_US="Model plugin", zh_Hans="模型插件"), + icon="icon.png", + label=I18nObject(en_US="Model Plugin", zh_Hans="模型插件"), + category=PluginCategory.Extension, # Will be Extension without a model provider + created_at=datetime.datetime.now(), + resource=PluginResourceRequirements(memory=1024, permission=None), + plugins=PluginDeclaration.Plugins(), + meta=PluginDeclaration.Meta(version="1.0.0"), + ) + + # Assert: Verify category defaults to Extension when no provider is specified + assert declaration.category == PluginCategory.Extension + + def test_extension_category_default(self): + """Test that plugins without specific providers default to Extension.""" + # Arrange: Create declaration without specific provider type + declaration = PluginDeclaration( + version="1.0.0", + author="test-author", + name="extension-plugin", + description=I18nObject(en_US="Extension plugin", zh_Hans="扩展插件"), + icon="icon.png", + label=I18nObject(en_US="Extension Plugin", zh_Hans="扩展插件"), + category=PluginCategory.Extension, + created_at=datetime.datetime.now(), + resource=PluginResourceRequirements(memory=512, permission=None), + plugins=PluginDeclaration.Plugins(), + meta=PluginDeclaration.Meta(version="1.0.0"), + ) + + # Assert: Verify category is Extension + assert declaration.category == PluginCategory.Extension + + class TestPluginResourceRequirements: + """Test plugin resource requirements and permissions.""" + + def test_default_resource_requirements(self): + """ + Test that plugin resource requirements can be created with default values. + + Resource requirements define the memory and permissions needed for a plugin to run. + This test verifies that a basic resource requirement with only memory can be created. + """ + # Arrange & Act: Create resource requirements with only memory specified + resources = PluginResourceRequirements(memory=512, permission=None) + + # Assert: Verify memory is set correctly and permissions are None + assert resources.memory == 512 + assert resources.permission is None + + def test_resource_requirements_with_tool_permission(self): + """ + Test plugin resource requirements with tool permissions enabled. + + Tool permissions allow a plugin to provide tool functionality. + This test verifies that tool permissions can be properly configured. + """ + # Arrange & Act: Create resource requirements with tool permissions + resources = PluginResourceRequirements( + memory=1024, + permission=PluginResourceRequirements.Permission( + tool=PluginResourceRequirements.Permission.Tool(enabled=True) + ), + ) + + # Assert: Verify tool permissions are enabled + assert resources.memory == 1024 + assert resources.permission is not None + assert resources.permission.tool is not None + assert resources.permission.tool.enabled is True + + def test_resource_requirements_with_model_permissions(self): + """ + Test plugin resource requirements with model permissions. + + Model permissions allow a plugin to provide various AI model capabilities + including LLM, text embedding, rerank, TTS, speech-to-text, and moderation. + """ + # Arrange & Act: Create resource requirements with comprehensive model permissions + resources = PluginResourceRequirements( + memory=2048, + permission=PluginResourceRequirements.Permission( + model=PluginResourceRequirements.Permission.Model( + enabled=True, + llm=True, + text_embedding=True, + rerank=True, + tts=False, + speech2text=False, + moderation=True, + ) + ), + ) + + # Assert: Verify all model permissions are set correctly + assert resources.memory == 2048 + assert resources.permission.model.enabled is True + assert resources.permission.model.llm is True + assert resources.permission.model.text_embedding is True + assert resources.permission.model.rerank is True + assert resources.permission.model.tts is False + assert resources.permission.model.speech2text is False + assert resources.permission.model.moderation is True + + def test_resource_requirements_with_storage_permission(self): + """ + Test plugin resource requirements with storage permissions. + + Storage permissions allow a plugin to persist data with size limits. + The size must be between 1KB (1024 bytes) and 1GB (1073741824 bytes). + """ + # Arrange & Act: Create resource requirements with storage permissions + resources = PluginResourceRequirements( + memory=512, + permission=PluginResourceRequirements.Permission( + storage=PluginResourceRequirements.Permission.Storage(enabled=True, size=10485760) # 10MB + ), + ) + + # Assert: Verify storage permissions and size limits + assert resources.permission.storage.enabled is True + assert resources.permission.storage.size == 10485760 + + def test_resource_requirements_with_endpoint_permission(self): + """ + Test plugin resource requirements with endpoint permissions. + + Endpoint permissions allow a plugin to expose HTTP endpoints. + """ + # Arrange & Act: Create resource requirements with endpoint permissions + resources = PluginResourceRequirements( + memory=1024, + permission=PluginResourceRequirements.Permission( + endpoint=PluginResourceRequirements.Permission.Endpoint(enabled=True) + ), + ) + + # Assert: Verify endpoint permissions are enabled + assert resources.permission.endpoint.enabled is True + + def test_resource_requirements_with_node_permission(self): + """ + Test plugin resource requirements with node permissions. + + Node permissions allow a plugin to provide custom workflow nodes. + """ + # Arrange & Act: Create resource requirements with node permissions + resources = PluginResourceRequirements( + memory=768, + permission=PluginResourceRequirements.Permission( + node=PluginResourceRequirements.Permission.Node(enabled=True) + ), + ) + + # Assert: Verify node permissions are enabled + assert resources.permission.node.enabled is True + + class TestPluginInstallationSources: + """Test different plugin installation sources.""" + + def test_marketplace_installation_source(self): + """ + Test plugin installation from marketplace source. + + Marketplace is the official plugin distribution channel where + verified and community plugins are available for installation. + """ + # Arrange & Act: Use marketplace as installation source + source = PluginInstallationSource.Marketplace + + # Assert: Verify source type + assert source == PluginInstallationSource.Marketplace + assert source.value == "marketplace" + + def test_github_installation_source(self): + """ + Test plugin installation from GitHub source. + + GitHub source allows installing plugins directly from GitHub repositories, + useful for development and testing unreleased versions. + """ + # Arrange & Act: Use GitHub as installation source + source = PluginInstallationSource.Github + + # Assert: Verify source type + assert source == PluginInstallationSource.Github + assert source.value == "github" + + def test_package_installation_source(self): + """ + Test plugin installation from package source. + + Package source allows installing plugins from local .difypkg files, + useful for private or custom plugins. + """ + # Arrange & Act: Use package as installation source + source = PluginInstallationSource.Package + + # Assert: Verify source type + assert source == PluginInstallationSource.Package + assert source.value == "package" + + def test_remote_installation_source(self): + """ + Test plugin installation from remote source. + + Remote source allows installing plugins from custom remote URLs. + """ + # Arrange & Act: Use remote as installation source + source = PluginInstallationSource.Remote + + # Assert: Verify source type + assert source == PluginInstallationSource.Remote + assert source.value == "remote" + + class TestPluginBundleOperations: + """Test plugin bundle operations and dependency extraction.""" + + @pytest.fixture + def plugin_installer(self): + """Create a PluginInstaller instance for testing.""" + return PluginInstaller() + + def test_upload_bundle_with_marketplace_dependencies(self, plugin_installer): + """ + Test uploading a bundle with marketplace dependencies. + + Marketplace dependencies reference plugins available in the official marketplace + by organization, plugin name, and version. + """ + # Arrange: Create mock bundle with marketplace dependencies + bundle_data = b"mock-marketplace-bundle" + mock_dependencies = [ + PluginBundleDependency( + type=PluginBundleDependencyType.Marketplace, + value=PluginBundleDependency.Marketplace( + organization="langgenius", plugin="search-tool", version="1.2.0" + ), + ) + ] + + with patch.object(plugin_installer, "_request_with_plugin_daemon_response", return_value=mock_dependencies): + # Act: Upload bundle + result = plugin_installer.upload_bundle("test-tenant", bundle_data) + + # Assert: Verify marketplace dependency was extracted + assert len(result) == 1 + assert result[0].type == PluginBundleDependencyType.Marketplace + assert isinstance(result[0].value, PluginBundleDependency.Marketplace) + assert result[0].value.organization == "langgenius" + assert result[0].value.plugin == "search-tool" + + def test_upload_bundle_with_github_dependencies(self, plugin_installer): + """ + Test uploading a bundle with GitHub dependencies. + + GitHub dependencies reference plugins hosted on GitHub repositories + with specific releases and package files. + """ + # Arrange: Create mock bundle with GitHub dependencies + bundle_data = b"mock-github-bundle" + mock_dependencies = [ + PluginBundleDependency( + type=PluginBundleDependencyType.Github, + value=PluginBundleDependency.Github( + repo_address="https://github.com/example/plugin", + repo="example/plugin", + release="v2.0.0", + packages="plugin-v2.0.0.zip", + ), + ) + ] + + with patch.object(plugin_installer, "_request_with_plugin_daemon_response", return_value=mock_dependencies): + # Act: Upload bundle + result = plugin_installer.upload_bundle("test-tenant", bundle_data) + + # Assert: Verify GitHub dependency was extracted + assert len(result) == 1 + assert result[0].type == PluginBundleDependencyType.Github + assert isinstance(result[0].value, PluginBundleDependency.Github) + assert result[0].value.repo == "example/plugin" + assert result[0].value.release == "v2.0.0" + + def test_upload_bundle_with_package_dependencies(self, plugin_installer): + """ + Test uploading a bundle with package dependencies. + + Package dependencies include the full plugin manifest and unique identifier, + allowing for self-contained plugin bundles. + """ + # Arrange: Create mock bundle with package dependencies + bundle_data = b"mock-package-bundle" + mock_manifest = PluginDeclaration( + version="1.5.0", + author="bundle-author", + name="bundled-plugin", + description=I18nObject(en_US="Bundled plugin", zh_Hans="捆绑插件"), + icon="icon.png", + label=I18nObject(en_US="Bundled Plugin", zh_Hans="捆绑插件"), + category=PluginCategory.Extension, + created_at=datetime.datetime.now(), + resource=PluginResourceRequirements(memory=512, permission=None), + plugins=PluginDeclaration.Plugins(), + meta=PluginDeclaration.Meta(version="1.5.0"), + ) + + mock_dependencies = [ + PluginBundleDependency( + type=PluginBundleDependencyType.Package, + value=PluginBundleDependency.Package( + unique_identifier="org/bundled-plugin/1.5.0", manifest=mock_manifest + ), + ) + ] + + with patch.object(plugin_installer, "_request_with_plugin_daemon_response", return_value=mock_dependencies): + # Act: Upload bundle + result = plugin_installer.upload_bundle("test-tenant", bundle_data) + + # Assert: Verify package dependency was extracted with manifest + assert len(result) == 1 + assert result[0].type == PluginBundleDependencyType.Package + assert isinstance(result[0].value, PluginBundleDependency.Package) + assert result[0].value.unique_identifier == "org/bundled-plugin/1.5.0" + assert result[0].value.manifest.name == "bundled-plugin" + + def test_upload_bundle_with_mixed_dependencies(self, plugin_installer): + """ + Test uploading a bundle with multiple dependency types. + + Real-world plugin bundles often have dependencies from various sources: + marketplace plugins, GitHub repositories, and packaged plugins. + """ + # Arrange: Create mock bundle with mixed dependencies + bundle_data = b"mock-mixed-bundle" + mock_dependencies = [ + PluginBundleDependency( + type=PluginBundleDependencyType.Marketplace, + value=PluginBundleDependency.Marketplace(organization="org1", plugin="plugin1", version="1.0.0"), + ), + PluginBundleDependency( + type=PluginBundleDependencyType.Github, + value=PluginBundleDependency.Github( + repo_address="https://github.com/org2/plugin2", + repo="org2/plugin2", + release="v1.0.0", + packages="plugin2.zip", + ), + ), + ] + + with patch.object(plugin_installer, "_request_with_plugin_daemon_response", return_value=mock_dependencies): + # Act: Upload bundle + result = plugin_installer.upload_bundle("test-tenant", bundle_data, verify_signature=True) + + # Assert: Verify all dependency types were extracted + assert len(result) == 2 + assert result[0].type == PluginBundleDependencyType.Marketplace + assert result[1].type == PluginBundleDependencyType.Github + + class TestPluginTaskStatusTransitions: + """Test plugin installation task status transitions and lifecycle.""" + + @pytest.fixture + def plugin_installer(self): + """Create a PluginInstaller instance for testing.""" + return PluginInstaller() + + def test_task_status_pending(self, plugin_installer): + """ + Test plugin installation task in pending status. + + Pending status indicates the task has been created but not yet started. + No plugins have been processed yet. + """ + # Arrange: Create mock task in pending status + mock_task = PluginInstallTask( + id="pending-task", + created_at=datetime.datetime.now(), + updated_at=datetime.datetime.now(), + status=PluginInstallTaskStatus.Pending, + total_plugins=3, + completed_plugins=0, # No plugins completed yet + plugins=[], + ) + + with patch.object(plugin_installer, "_request_with_plugin_daemon_response", return_value=mock_task): + # Act: Fetch task + result = plugin_installer.fetch_plugin_installation_task("test-tenant", "pending-task") + + # Assert: Verify pending status + assert result.status == PluginInstallTaskStatus.Pending + assert result.completed_plugins == 0 + assert result.total_plugins == 3 + + def test_task_status_running(self, plugin_installer): + """ + Test plugin installation task in running status. + + Running status indicates the task is actively installing plugins. + Some plugins may be completed while others are still in progress. + """ + # Arrange: Create mock task in running status + mock_task = PluginInstallTask( + id="running-task", + created_at=datetime.datetime.now(), + updated_at=datetime.datetime.now(), + status=PluginInstallTaskStatus.Running, + total_plugins=5, + completed_plugins=2, # 2 out of 5 completed + plugins=[], + ) + + with patch.object(plugin_installer, "_request_with_plugin_daemon_response", return_value=mock_task): + # Act: Fetch task + result = plugin_installer.fetch_plugin_installation_task("test-tenant", "running-task") + + # Assert: Verify running status and progress + assert result.status == PluginInstallTaskStatus.Running + assert result.completed_plugins == 2 + assert result.total_plugins == 5 + assert result.completed_plugins < result.total_plugins + + def test_task_status_success(self, plugin_installer): + """ + Test plugin installation task in success status. + + Success status indicates all plugins in the task have been + successfully installed without errors. + """ + # Arrange: Create mock task in success status + mock_task = PluginInstallTask( + id="success-task", + created_at=datetime.datetime.now(), + updated_at=datetime.datetime.now(), + status=PluginInstallTaskStatus.Success, + total_plugins=4, + completed_plugins=4, # All plugins completed + plugins=[], + ) + + with patch.object(plugin_installer, "_request_with_plugin_daemon_response", return_value=mock_task): + # Act: Fetch task + result = plugin_installer.fetch_plugin_installation_task("test-tenant", "success-task") + + # Assert: Verify success status and completion + assert result.status == PluginInstallTaskStatus.Success + assert result.completed_plugins == result.total_plugins + assert result.completed_plugins == 4 + + def test_task_status_failed(self, plugin_installer): + """ + Test plugin installation task in failed status. + + Failed status indicates the task encountered errors during installation. + Some plugins may have been installed before the failure occurred. + """ + # Arrange: Create mock task in failed status + mock_task = PluginInstallTask( + id="failed-task", + created_at=datetime.datetime.now(), + updated_at=datetime.datetime.now(), + status=PluginInstallTaskStatus.Failed, + total_plugins=3, + completed_plugins=1, # Only 1 completed before failure + plugins=[], + ) + + with patch.object(plugin_installer, "_request_with_plugin_daemon_response", return_value=mock_task): + # Act: Fetch task + result = plugin_installer.fetch_plugin_installation_task("test-tenant", "failed-task") + + # Assert: Verify failed status + assert result.status == PluginInstallTaskStatus.Failed + assert result.completed_plugins < result.total_plugins + + class TestPluginI18nSupport: + """Test plugin internationalization (i18n) support.""" + + def test_plugin_with_multiple_languages(self): + """ + Test plugin declaration with multiple language support. + + Plugins should support multiple languages for descriptions and labels + to provide localized experiences for users worldwide. + """ + # Arrange & Act: Create plugin with English and Chinese support + declaration = PluginDeclaration( + version="1.0.0", + author="i18n-author", + name="multilang-plugin", + description=I18nObject( + en_US="A plugin with multilingual support", + zh_Hans="支持多语言的插件", + ja_JP="多言語対応のプラグイン", + ), + icon="icon.png", + label=I18nObject(en_US="Multilingual Plugin", zh_Hans="多语言插件", ja_JP="多言語プラグイン"), + category=PluginCategory.Extension, + created_at=datetime.datetime.now(), + resource=PluginResourceRequirements(memory=512, permission=None), + plugins=PluginDeclaration.Plugins(), + meta=PluginDeclaration.Meta(version="1.0.0"), + ) + + # Assert: Verify all language variants are preserved + assert declaration.description.en_US == "A plugin with multilingual support" + assert declaration.description.zh_Hans == "支持多语言的插件" + assert declaration.label.en_US == "Multilingual Plugin" + assert declaration.label.zh_Hans == "多语言插件" + + def test_plugin_readme_language_variants(self): + """ + Test fetching plugin README in different languages. + + Plugins can provide README files in multiple languages to help + users understand the plugin in their preferred language. + """ + # Arrange: Create plugin installer instance + plugin_installer = PluginInstaller() + + # Mock README responses for different languages + english_readme = PluginReadmeResponse( + content="# English README\n\nThis is the English version.", language="en_US" + ) + + chinese_readme = PluginReadmeResponse(content="# 中文说明\n\n这是中文版本。", language="zh_Hans") + + # Test English README + with patch.object(plugin_installer, "_request_with_plugin_daemon_response", return_value=english_readme): + # Act: Fetch English README + result_en = plugin_installer.fetch_plugin_readme("test-tenant", "plugin/1.0.0", "en_US") + + # Assert: Verify English content + assert "English README" in result_en + + # Test Chinese README + with patch.object(plugin_installer, "_request_with_plugin_daemon_response", return_value=chinese_readme): + # Act: Fetch Chinese README + result_zh = plugin_installer.fetch_plugin_readme("test-tenant", "plugin/1.0.0", "zh_Hans") + + # Assert: Verify Chinese content + assert "中文说明" in result_zh + + +"""Comprehensive unit tests for Plugin Runtime functionality. + +This test module covers all aspects of plugin runtime including: +- Plugin execution through the plugin daemon +- Sandbox isolation via HTTP communication +- Resource limits (timeout, memory constraints) +- Error handling for various failure scenarios +- Plugin communication (request/response patterns, streaming) + +All tests use mocking to avoid external dependencies and ensure fast, reliable execution. +Tests follow the Arrange-Act-Assert pattern for clarity. +""" + + +class TestPluginRuntimeCases: + @staticmethod + @pytest.fixture(autouse=True) + def _patch_shared_httpx_client(): + """Make BasePluginClient's module-level httpx client delegate to patched httpx.request/stream. + + After refactor, code uses core.plugin.impl.base._httpx_client directly. + Patch its request/stream to route through module-level httpx so existing mocks still apply. + """ + + def request[**P](*args: P.args, **kwargs: P.kwargs) -> httpx.Response: + response = httpx.request(*args, **kwargs) + method, url = kwargs["method"], kwargs["url"] + assert isinstance(method, str) + assert isinstance(url, str) + response.request = httpx.Request(method, url) + return response + + with ( + patch("core.plugin.impl.base._httpx_client.request", side_effect=request), + patch("core.plugin.impl.base._httpx_client.stream", side_effect=lambda **kw: httpx.stream(**kw)), + ): + yield + + class TestPluginRuntimeExecution: + """Unit tests for plugin execution functionality. + + Tests cover: + - Successful plugin invocation + - Request preparation and headers + - Response parsing + - Streaming responses + """ + + @pytest.fixture + def plugin_client(self): + """Create a BasePluginClient instance for testing.""" + return BasePluginClient() + + @pytest.fixture + def mock_config(self, config_overrides): + """Mock plugin daemon configuration.""" + config_overrides( + PLUGIN_DAEMON_URL="http://127.0.0.1:5002", + PLUGIN_DAEMON_KEY="test-api-key", + ) + + def test_request_preparation(self, plugin_client, mock_config): + """Test that requests are properly prepared with correct headers and URL.""" + # Arrange + path = "plugin/test-tenant/management/list" + headers = {"Custom-Header": "value"} + data = {"key": "value"} + params = {"page": 1} + + # Act + url, prepared_headers, prepared_data, prepared_params, files = plugin_client._prepare_request( + path, headers, data, params, None + ) + + # Assert + assert url == "http://127.0.0.1:5002/plugin/test-tenant/management/list" + assert prepared_headers["X-Api-Key"] == "test-api-key" + assert prepared_headers["Custom-Header"] == "value" + assert prepared_headers["Accept-Encoding"] == "gzip, deflate, br" + assert prepared_data == data + assert prepared_params == params + + def test_request_with_json_content_type(self, plugin_client, mock_config): + """Test request preparation with JSON content type.""" + # Arrange + path = "plugin/test-tenant/management/install" + headers = {"Content-Type": "application/json"} + data = {"plugin_id": "test-plugin"} + + # Act + url, prepared_headers, prepared_data, prepared_params, files = plugin_client._prepare_request( + path, headers, data, None, None + ) + + # Assert + assert prepared_headers["Content-Type"] == "application/json" + assert prepared_data == json.dumps(data) + + def test_successful_request_execution(self, plugin_client, mock_config): + """Test successful HTTP request execution.""" + # Arrange + mock_response = httpx.Response(200, json={"result": "success"}) + + with patch("httpx.request", return_value=mock_response, autospec=True) as mock_request: + # Act + response = plugin_client._request("GET", "plugin/test-tenant/management/list") + + # Assert + assert response.status_code == 200 + mock_request.assert_called_once() + call_kwargs = mock_request.call_args[1] + assert call_kwargs["method"] == "GET" + assert "http://127.0.0.1:5002/plugin/test-tenant/management/list" in call_kwargs["url"] + assert call_kwargs["headers"]["X-Api-Key"] == "test-api-key" + + def test_request_with_timeout_configuration(self, plugin_client, mock_config): + """Test that timeout configuration is properly applied.""" + # Arrange + mock_response = httpx.Response(200) + + with patch("httpx.request", return_value=mock_response, autospec=True) as mock_request: + # Act + plugin_client._request("GET", "plugin/test-tenant/test") + + # Assert + call_kwargs = mock_request.call_args[1] + assert "timeout" in call_kwargs + + def test_request_timeout_can_be_scoped_to_current_context(self, plugin_client, mock_config): + mock_response = httpx.Response(200) + + with ( + patch("core.plugin.impl.base.plugin_daemon_request_timeout", httpx.Timeout(600.0)), + patch("httpx.request", return_value=mock_response, autospec=True) as mock_request, + ): + with use_plugin_daemon_request_timeout(30.0): + plugin_client._request("GET", "plugin/test-tenant/test") + plugin_client._request("GET", "plugin/test-tenant/test") + + scoped_timeout = mock_request.call_args_list[0].kwargs["timeout"] + default_timeout = mock_request.call_args_list[1].kwargs["timeout"] + assert isinstance(scoped_timeout, httpx.Timeout) + assert isinstance(default_timeout, httpx.Timeout) + assert scoped_timeout.read == 30.0 + assert default_timeout.read == 600.0 + + @pytest.mark.parametrize("timeout_seconds", [0.0, -1.0]) + def test_request_timeout_override_rejects_non_positive_values(self, timeout_seconds): + with pytest.raises(ValueError, match="greater than zero"): + with use_plugin_daemon_request_timeout(timeout_seconds): + pass + + def test_request_connection_error(self, plugin_client, mock_config): + """Test handling of connection errors during request.""" + # Arrange + with patch("httpx.request", side_effect=httpx.RequestError("Connection failed"), autospec=True): + # Act & Assert + with pytest.raises(PluginDaemonInnerError) as exc_info: + plugin_client._request("GET", "plugin/test-tenant/test") + assert exc_info.value.code == -500 + assert "Request to Plugin Daemon Service failed" in exc_info.value.message + + class TestPluginRuntimeSandboxIsolation: + """Unit tests for plugin sandbox isolation. + + Tests cover: + - Isolated execution environment via HTTP + - API key authentication + - Request/response boundaries + - Plugin daemon communication protocol + """ + + @pytest.fixture + def plugin_client(self): + """Create a BasePluginClient instance for testing.""" + return BasePluginClient() + + @pytest.fixture + def mock_config(self, config_overrides): + """Mock plugin daemon configuration.""" + config_overrides( + PLUGIN_DAEMON_URL="http://127.0.0.1:5002", + PLUGIN_DAEMON_KEY="secure-api-key", + ) + + def test_api_key_authentication(self, plugin_client, mock_config): + """Test that all requests include API key for authentication.""" + # Arrange + mock_response = httpx.Response(200, json={"code": 0, "message": "", "data": True}) + + with patch("httpx.request", return_value=mock_response, autospec=True) as mock_request: + # Act + plugin_client._request("GET", "plugin/test-tenant/test") + + # Assert + call_kwargs = mock_request.call_args[1] + assert call_kwargs["headers"]["X-Api-Key"] == "secure-api-key" + + def test_isolated_plugin_execution_via_http(self, plugin_client, mock_config): + """Test that plugin execution is isolated via HTTP communication.""" + + # Arrange + class TestResponse(BaseModel): + result: str + + mock_response = httpx.Response( + 200, json={"code": 0, "message": "", "data": {"result": "isolated_execution"}} + ) + + with patch("httpx.request", return_value=mock_response, autospec=True): + # Act + result = plugin_client._request_with_plugin_daemon_response( + "POST", "plugin/test-tenant/dispatch/tool/invoke", TestResponse, data={"tool": "test"} + ) + + # Assert + assert result.result == "isolated_execution" + + def test_plugin_daemon_unauthorized_error(self, plugin_client, mock_config): + """Test handling of unauthorized access to plugin daemon.""" + # Arrange + error_message = json.dumps( + {"error_type": "PluginDaemonUnauthorizedError", "message": "Unauthorized access"} + ) + mock_response = httpx.Response(200, json={"code": -1, "message": error_message, "data": None}) + + with patch("httpx.request", return_value=mock_response, autospec=True): + # Act & Assert + with pytest.raises(PluginDaemonUnauthorizedError) as exc_info: + plugin_client._request_with_plugin_daemon_response("GET", "plugin/test-tenant/test", bool) + assert "Unauthorized access" in exc_info.value.description + + def test_plugin_permission_denied(self, plugin_client, mock_config): + """Test handling of permission denied errors.""" + # Arrange + error_message = json.dumps( + {"error_type": "PluginPermissionDeniedError", "message": "Permission denied for this operation"} + ) + mock_response = httpx.Response(200, json={"code": -1, "message": error_message, "data": None}) + + with patch("httpx.request", return_value=mock_response, autospec=True): + # Act & Assert + with pytest.raises(PluginPermissionDeniedError) as exc_info: + plugin_client._request_with_plugin_daemon_response("POST", "plugin/test-tenant/test", bool) + assert "Permission denied" in exc_info.value.description + + class TestPluginRuntimeResourceLimits: + """Unit tests for plugin resource limits. + + Tests cover: + - Timeout enforcement + - Memory constraints + - Resource limit violations + - Graceful degradation + """ + + @pytest.fixture + def plugin_client(self): + """Create a BasePluginClient instance for testing.""" + return BasePluginClient() + + @pytest.fixture + def mock_config(self, config_overrides): + """Mock plugin daemon configuration with timeout.""" + config_overrides( + PLUGIN_DAEMON_URL="http://127.0.0.1:5002", + PLUGIN_DAEMON_KEY="test-key", + ) + with patch("core.plugin.impl.base.plugin_daemon_request_timeout", httpx.Timeout(30.0)): + yield + + def test_timeout_configuration_applied(self, plugin_client, mock_config): + """Test that timeout configuration is properly applied to requests.""" + # Arrange + mock_response = httpx.Response(200) + + with patch("httpx.request", return_value=mock_response, autospec=True) as mock_request: + # Act + plugin_client._request("GET", "plugin/test-tenant/test") + + # Assert + call_kwargs = mock_request.call_args[1] + assert call_kwargs["timeout"] is not None + + def test_timeout_error_handling(self, plugin_client, mock_config): + """Test handling of timeout errors.""" + # Arrange + with patch("httpx.request", side_effect=httpx.TimeoutException("Request timeout"), autospec=True): + # Act & Assert + with pytest.raises(PluginDaemonInnerError) as exc_info: + plugin_client._request("GET", "plugin/test-tenant/test") + assert exc_info.value.code == -500 + + def test_streaming_request_timeout(self, plugin_client, mock_config): + """Test timeout handling for streaming requests.""" + # Arrange + with patch("httpx.stream", side_effect=httpx.TimeoutException("Stream timeout"), autospec=True): + # Act & Assert + with pytest.raises(PluginDaemonInnerError) as exc_info: + list(plugin_client._stream_request("POST", "plugin/test-tenant/stream")) + assert exc_info.value.code == -500 + + def test_resource_limit_error_from_daemon(self, plugin_client, mock_config): + """Test handling of resource limit errors from plugin daemon.""" + # Arrange + error_message = json.dumps( + {"error_type": "PluginDaemonInternalServerError", "message": "Resource limit exceeded"} + ) + mock_response = httpx.Response(200, json={"code": -1, "message": error_message, "data": None}) + + with patch("httpx.request", return_value=mock_response, autospec=True): + # Act & Assert + with pytest.raises(PluginDaemonInternalServerError) as exc_info: + plugin_client._request_with_plugin_daemon_response("POST", "plugin/test-tenant/test", bool) + assert "Resource limit exceeded" in exc_info.value.description + + class TestPluginRuntimeErrorHandling: + """Unit tests for plugin runtime error handling. + + Tests cover: + - Various error types (invoke, validation, connection) + - Error propagation and transformation + - User-friendly error messages + - Error recovery mechanisms + """ + + @pytest.fixture + def plugin_client(self): + """Create a BasePluginClient instance for testing.""" + return BasePluginClient() + + @pytest.fixture + def mock_config(self, config_overrides): + """Mock plugin daemon configuration.""" + config_overrides( + PLUGIN_DAEMON_URL="http://127.0.0.1:5002", + PLUGIN_DAEMON_KEY="test-key", + ) + + def test_plugin_invoke_rate_limit_error(self, plugin_client, mock_config): + """Test handling of rate limit errors during plugin invocation.""" + # Arrange + invoke_error = { + "error_type": "InvokeRateLimitError", + "message": "Rate limit exceeded", + "args": {"description": "Rate limit exceeded"}, + } + error_message = json.dumps({"error_type": "PluginInvokeError", "message": json.dumps(invoke_error)}) + mock_response = httpx.Response(200, json={"code": -1, "message": error_message, "data": None}) + + with patch("httpx.request", return_value=mock_response, autospec=True): + # Act & Assert + with pytest.raises(InvokeRateLimitError) as exc_info: + plugin_client._request_with_plugin_daemon_response("POST", "plugin/test-tenant/invoke", bool) + assert "Rate limit exceeded" in exc_info.value.description + + def test_plugin_invoke_authorization_error(self, plugin_client, mock_config): + """Test handling of authorization errors during plugin invocation.""" + # Arrange + invoke_error = { + "error_type": "InvokeAuthorizationError", + "message": "Invalid credentials", + "args": {"description": "Invalid credentials"}, + } + error_message = json.dumps({"error_type": "PluginInvokeError", "message": json.dumps(invoke_error)}) + mock_response = httpx.Response(200, json={"code": -1, "message": error_message, "data": None}) + + with patch("httpx.request", return_value=mock_response, autospec=True): + # Act & Assert + with pytest.raises(InvokeAuthorizationError) as exc_info: + plugin_client._request_with_plugin_daemon_response("POST", "plugin/test-tenant/invoke", bool) + assert "Invalid credentials" in exc_info.value.description + + def test_plugin_invoke_bad_request_error(self, plugin_client, mock_config): + """Test handling of bad request errors during plugin invocation.""" + # Arrange + invoke_error = { + "error_type": "InvokeBadRequestError", + "message": "Invalid parameters", + "args": {"description": "Invalid parameters"}, + } + error_message = json.dumps({"error_type": "PluginInvokeError", "message": json.dumps(invoke_error)}) + mock_response = httpx.Response(200, json={"code": -1, "message": error_message, "data": None}) + + with patch("httpx.request", return_value=mock_response, autospec=True): + # Act & Assert + with pytest.raises(InvokeBadRequestError) as exc_info: + plugin_client._request_with_plugin_daemon_response("POST", "plugin/test-tenant/invoke", bool) + assert "Invalid parameters" in exc_info.value.description + + def test_plugin_invoke_connection_error(self, plugin_client, mock_config): + """Test handling of connection errors during plugin invocation.""" + # Arrange + invoke_error = { + "error_type": "InvokeConnectionError", + "message": "Connection to external service failed", + "args": {"description": "Connection to external service failed"}, + } + error_message = json.dumps({"error_type": "PluginInvokeError", "message": json.dumps(invoke_error)}) + mock_response = httpx.Response(200, json={"code": -1, "message": error_message, "data": None}) + + with patch("httpx.request", return_value=mock_response, autospec=True): + # Act & Assert + with pytest.raises(InvokeConnectionError) as exc_info: + plugin_client._request_with_plugin_daemon_response("POST", "plugin/test-tenant/invoke", bool) + assert "Connection to external service failed" in exc_info.value.description + + def test_plugin_invoke_server_unavailable_error(self, plugin_client, mock_config): + """Test handling of server unavailable errors during plugin invocation.""" + # Arrange + invoke_error = { + "error_type": "InvokeServerUnavailableError", + "message": "Service temporarily unavailable", + "args": {"description": "Service temporarily unavailable"}, + } + error_message = json.dumps({"error_type": "PluginInvokeError", "message": json.dumps(invoke_error)}) + mock_response = httpx.Response(200, json={"code": -1, "message": error_message, "data": None}) + + with patch("httpx.request", return_value=mock_response, autospec=True): + # Act & Assert + with pytest.raises(InvokeServerUnavailableError) as exc_info: + plugin_client._request_with_plugin_daemon_response("POST", "plugin/test-tenant/invoke", bool) + assert "Service temporarily unavailable" in exc_info.value.description + + def test_credentials_validation_error(self, plugin_client, mock_config): + """Test handling of credential validation errors.""" + # Arrange + invoke_error = { + "error_type": "CredentialsValidateFailedError", + "message": "Invalid API key format", + } + error_message = json.dumps({"error_type": "PluginInvokeError", "message": json.dumps(invoke_error)}) + mock_response = httpx.Response(200, json={"code": -1, "message": error_message, "data": None}) + + with patch("httpx.request", return_value=mock_response, autospec=True): + # Act & Assert + with pytest.raises(CredentialsValidateFailedError) as exc_info: + plugin_client._request_with_plugin_daemon_response("POST", "plugin/test-tenant/validate", bool) + assert "Invalid API key format" in str(exc_info.value) + + def test_plugin_not_found_error(self, plugin_client, mock_config): + """Test handling of plugin not found errors.""" + # Arrange + error_message = json.dumps( + {"error_type": "PluginNotFoundError", "message": "Plugin with ID 'test-plugin' not found"} + ) + mock_response = httpx.Response(200, json={"code": -1, "message": error_message, "data": None}) + + with patch("httpx.request", return_value=mock_response, autospec=True): + # Act & Assert + with pytest.raises(PluginNotFoundError) as exc_info: + plugin_client._request_with_plugin_daemon_response("GET", "plugin/test-tenant/get", bool) + assert "Plugin with ID 'test-plugin' not found" in exc_info.value.description + + def test_plugin_unique_identifier_error(self, plugin_client, mock_config): + """Test handling of unique identifier errors.""" + # Arrange + error_message = json.dumps( + {"error_type": "PluginUniqueIdentifierError", "message": "Invalid plugin identifier format"} + ) + mock_response = httpx.Response(200, json={"code": -1, "message": error_message, "data": None}) + + with patch("httpx.request", return_value=mock_response, autospec=True): + # Act & Assert + with pytest.raises(PluginUniqueIdentifierError) as exc_info: + plugin_client._request_with_plugin_daemon_response("POST", "plugin/test-tenant/install", bool) + assert "Invalid plugin identifier format" in exc_info.value.description + + def test_daemon_bad_request_error(self, plugin_client, mock_config): + """Test handling of daemon bad request errors.""" + # Arrange + error_message = json.dumps( + {"error_type": "PluginDaemonBadRequestError", "message": "Missing required parameter"} + ) + mock_response = httpx.Response(200, json={"code": -1, "message": error_message, "data": None}) + + with patch("httpx.request", return_value=mock_response, autospec=True): + # Act & Assert + with pytest.raises(PluginDaemonBadRequestError) as exc_info: + plugin_client._request_with_plugin_daemon_response("POST", "plugin/test-tenant/test", bool) + assert "Missing required parameter" in exc_info.value.description + + def test_daemon_not_found_error(self, plugin_client, mock_config): + """Test handling of daemon not found errors.""" + # Arrange + error_message = json.dumps({"error_type": "PluginDaemonNotFoundError", "message": "Resource not found"}) + mock_response = httpx.Response(200, json={"code": -1, "message": error_message, "data": None}) + + with patch("httpx.request", return_value=mock_response, autospec=True): + # Act & Assert + with pytest.raises(PluginDaemonNotFoundError) as exc_info: + plugin_client._request_with_plugin_daemon_response("GET", "plugin/test-tenant/resource", bool) + assert "Resource not found" in exc_info.value.description + + def test_generic_plugin_invoke_error(self, plugin_client, mock_config): + """Test handling of generic plugin invoke errors.""" + # Arrange + # Create a proper nested JSON structure for PluginInvokeError + invoke_error_message = json.dumps( + {"error_type": "UnknownInvokeError", "message": "Generic plugin execution error"} + ) + error_message = json.dumps({"error_type": "PluginInvokeError", "message": invoke_error_message}) + mock_response = httpx.Response(200, json={"code": -1, "message": error_message, "data": None}) + + with patch("httpx.request", return_value=mock_response, autospec=True): + # Act & Assert + with pytest.raises(PluginInvokeError) as exc_info: + plugin_client._request_with_plugin_daemon_response("POST", "plugin/test-tenant/invoke", bool) + assert exc_info.value.description is not None + + def test_unknown_error_type(self, plugin_client, mock_config): + """Test handling of unknown error types.""" + # Arrange + error_message = json.dumps({"error_type": "UnknownErrorType", "message": "Unknown error occurred"}) + mock_response = httpx.Response(200, json={"code": -1, "message": error_message, "data": None}) + + with patch("httpx.request", return_value=mock_response, autospec=True): + # Act & Assert + with pytest.raises( + Exception, + match=("got unknown error from plugin daemon: UnknownErrorType, message: Unknown error occurred"), + ) as exc_info: + plugin_client._request_with_plugin_daemon_response("POST", "plugin/test-tenant/test", bool) + assert "got unknown error from plugin daemon" in str(exc_info.value) + + def test_http_status_error_handling(self, plugin_client, mock_config): + """Test handling of HTTP status errors.""" + # Arrange + mock_response = httpx.Response( + 500, request=httpx.Request("GET", "http://127.0.0.1:5002/plugin/test-tenant/test") + ) + + with patch("httpx.request", return_value=mock_response, autospec=True): + # Act & Assert + with pytest.raises(PluginDaemonInternalServerError): + plugin_client._request_with_plugin_daemon_response("GET", "plugin/test-tenant/test", bool) + + def test_empty_data_response_error(self, plugin_client, mock_config): + """Test handling of empty data in successful response.""" + # Arrange + mock_response = httpx.Response(200, json={"code": 0, "message": "", "data": None}) + + with patch("httpx.request", return_value=mock_response, autospec=True): + # Act & Assert + with pytest.raises(ValueError, match="got empty data from plugin daemon") as exc_info: + plugin_client._request_with_plugin_daemon_response("GET", "plugin/test-tenant/test", bool) + assert "got empty data from plugin daemon" in str(exc_info.value) + + class TestPluginRuntimeCommunication: + """Unit tests for plugin communication patterns. + + Tests cover: + - Request/response communication + - Streaming responses + - Data serialization/deserialization + - Message formatting + """ + + @pytest.fixture + def plugin_client(self): + """Create a BasePluginClient instance for testing.""" + return BasePluginClient() + + @pytest.fixture + def mock_config(self, config_overrides): + """Mock plugin daemon configuration.""" + config_overrides( + PLUGIN_DAEMON_URL="http://127.0.0.1:5002", + PLUGIN_DAEMON_KEY="test-key", + ) + + def test_request_response_communication(self, plugin_client, mock_config): + """Test basic request/response communication pattern.""" + + # Arrange + class TestModel(BaseModel): + value: str + count: int + + mock_response = httpx.Response(200, json={"code": 0, "message": "", "data": {"value": "test", "count": 42}}) + + with patch("httpx.request", return_value=mock_response, autospec=True): + # Act + result = plugin_client._request_with_plugin_daemon_response( + "POST", "plugin/test-tenant/test", TestModel, data={"input": "data"} + ) + + # Assert + assert isinstance(result, TestModel) + assert result.value == "test" + assert result.count == 42 + + def test_streaming_response_communication(self, plugin_client, mock_config): + """Test streaming response communication pattern.""" + + # Arrange + class StreamModel(BaseModel): + chunk: str + + stream_data = [ + 'data: {"code": 0, "message": "", "data": {"chunk": "first"}}', + 'data: {"code": 0, "message": "", "data": {"chunk": "second"}}', + 'data: {"code": 0, "message": "", "data": {"chunk": "third"}}', + ] + + mock_response = httpx.Response(200, text="\n".join(stream_data)) + + with patch("httpx.stream", autospec=True) as mock_stream: + mock_stream.return_value.__enter__.return_value = mock_response + + # Act + results = list( + plugin_client._request_with_plugin_daemon_response_stream( + "POST", "plugin/test-tenant/stream", StreamModel + ) + ) + + # Assert + assert len(results) == 3 + assert all(isinstance(r, StreamModel) for r in results) + assert results[0].chunk == "first" + assert results[1].chunk == "second" + assert results[2].chunk == "third" + + def test_streaming_with_error_in_stream(self, plugin_client, mock_config): + """Test error handling in streaming responses.""" + # Arrange + # Create proper error structure for -500 code + error_obj = json.dumps({"error_type": "PluginDaemonInnerError", "message": "Stream error occurred"}) + stream_data = [ + 'data: {"code": 0, "message": "", "data": {"chunk": "first"}}', + f'data: {{"code": -500, "message": {json.dumps(error_obj)}, "data": null}}', + ] + + mock_response = httpx.Response(200, text="\n".join(stream_data)) + + with patch("httpx.stream", autospec=True) as mock_stream: + mock_stream.return_value.__enter__.return_value = mock_response + + # Act + class StreamModel(BaseModel): + chunk: str + + results = plugin_client._request_with_plugin_daemon_response_stream( + "POST", "plugin/test-tenant/stream", StreamModel + ) + + # Assert + first_result = next(results) + assert first_result.chunk == "first" + + with pytest.raises(PluginDaemonInnerError) as exc_info: + next(results) + assert exc_info.value.code == -500 + + def test_streaming_connection_error(self, plugin_client, mock_config): + """Test connection error during streaming.""" + # Arrange + with patch("httpx.stream", side_effect=httpx.RequestError("Stream connection failed"), autospec=True): + # Act & Assert + with pytest.raises(PluginDaemonInnerError) as exc_info: + list(plugin_client._stream_request("POST", "plugin/test-tenant/stream")) + assert exc_info.value.code == -500 + + def test_request_with_model_parsing(self, plugin_client, mock_config): + """Test request with direct model parsing (without daemon response wrapper).""" + + # Arrange + class DirectModel(BaseModel): + status: str + data: dict[str, Any] + + mock_response = httpx.Response(200, json={"status": "success", "data": {"key": "value"}}) + + with patch("httpx.request", return_value=mock_response, autospec=True): + # Act + result = plugin_client._request_with_model("GET", "plugin/test-tenant/direct", DirectModel) + + # Assert + assert isinstance(result, DirectModel) + assert result.status == "success" + assert result.data == {"key": "value"} + + def test_streaming_with_model_parsing(self, plugin_client, mock_config): + """Test streaming with direct model parsing.""" + + # Arrange + class StreamItem(BaseModel): + id: int + text: str + + stream_data = [ + '{"id": 1, "text": "first"}', + '{"id": 2, "text": "second"}', + ] + + mock_response = httpx.Response(200, text="\n".join(stream_data)) + + with patch("httpx.stream", autospec=True) as mock_stream: + mock_stream.return_value.__enter__.return_value = mock_response + + # Act + results = list( + plugin_client._stream_request_with_model("POST", "plugin/test-tenant/stream", StreamItem) + ) + + # Assert + assert len(results) == 2 + assert results[0].id == 1 + assert results[0].text == "first" + assert results[1].id == 2 + assert results[1].text == "second" + + def test_streaming_skips_empty_lines(self, plugin_client, mock_config): + """Test that streaming properly skips empty lines.""" + + # Arrange + class StreamModel(BaseModel): + value: str + + stream_data = [ + "", + '{"code": 0, "message": "", "data": {"value": "first"}}', + "", + "", + '{"code": 0, "message": "", "data": {"value": "second"}}', + "", + ] + + mock_response = httpx.Response(200, text="\n".join(stream_data)) + + with patch("httpx.stream", autospec=True) as mock_stream: + mock_stream.return_value.__enter__.return_value = mock_response + + # Act + results = list( + plugin_client._request_with_plugin_daemon_response_stream( + "POST", "plugin/test-tenant/stream", StreamModel + ) + ) + + # Assert + assert len(results) == 2 + assert results[0].value == "first" + assert results[1].value == "second" + + class TestPluginToolManagerIntegration: + """Integration tests for PluginToolManager. + + Tests cover: + - Tool invocation + - Credential validation + - Runtime parameter retrieval + - Tool provider management + """ + + @pytest.fixture + def tool_manager(self): + """Create a PluginToolManager instance for testing.""" + return PluginToolManager() + + @pytest.fixture + def mock_config(self, config_overrides): + """Mock plugin daemon configuration.""" + config_overrides( + PLUGIN_DAEMON_URL="http://127.0.0.1:5002", + PLUGIN_DAEMON_KEY="test-key", + ) + + def test_tool_invocation_success(self, tool_manager, mock_config): + """Test successful tool invocation.""" + # Arrange + stream_data = [ + 'data: {"code": 0, "message": "", "data": {"type": "text", "message": {"text": "Result"}}}', + ] + + mock_response = httpx.Response(200, text="\n".join(stream_data)) + + with patch("httpx.stream", autospec=True) as mock_stream: + mock_stream.return_value.__enter__.return_value = mock_response + + # Act + results = list( + tool_manager.invoke( + tenant_id="test-tenant", + user_id="test-user", + tool_provider="langgenius/test-plugin/test-provider", + tool_name="test-tool", + credentials={"api_key": "test-key"}, + credential_type=CredentialType.API_KEY, + tool_parameters={"param1": "value1"}, + ) + ) + + # Assert + assert len(results) > 0 + assert results[0].type == "text" + + def test_validate_provider_credentials_success(self, tool_manager, mock_config): + """Test successful provider credential validation.""" + # Arrange + stream_data = [ + 'data: {"code": 0, "message": "", "data": {"result": true}}', + ] + + mock_response = httpx.Response(200, text="\n".join(stream_data)) + + with patch("httpx.stream", autospec=True) as mock_stream: + mock_stream.return_value.__enter__.return_value = mock_response + + # Act + result = tool_manager.validate_provider_credentials( + tenant_id="test-tenant", + user_id="test-user", + provider="langgenius/test-plugin/test-provider", + credentials={"api_key": "valid-key"}, + ) + + # Assert + assert result is True + + def test_validate_provider_credentials_failure(self, tool_manager, mock_config): + """Test failed provider credential validation.""" + # Arrange + stream_data = [ + 'data: {"code": 0, "message": "", "data": {"result": false}}', + ] + + mock_response = httpx.Response(200, text="\n".join(stream_data)) + + with patch("httpx.stream", autospec=True) as mock_stream: + mock_stream.return_value.__enter__.return_value = mock_response + + # Act + result = tool_manager.validate_provider_credentials( + tenant_id="test-tenant", + user_id="test-user", + provider="langgenius/test-plugin/test-provider", + credentials={"api_key": "invalid-key"}, + ) + + # Assert + assert result is False + + def test_validate_datasource_credentials_success(self, tool_manager, mock_config): + """Test successful datasource credential validation.""" + # Arrange + stream_data = [ + 'data: {"code": 0, "message": "", "data": {"result": true}}', + ] + + mock_response = httpx.Response(200, text="\n".join(stream_data)) + + with patch("httpx.stream", autospec=True) as mock_stream: + mock_stream.return_value.__enter__.return_value = mock_response + + # Act + result = tool_manager.validate_datasource_credentials( + tenant_id="test-tenant", + user_id="test-user", + provider="langgenius/test-plugin/test-datasource", + credentials={"connection_string": "valid"}, + ) + + # Assert + assert result is True + + class TestPluginInstallerIntegration: + """Integration tests for PluginInstaller. + + Tests cover: + - Plugin installation + - Plugin listing + - Plugin uninstallation + - Package upload + """ + + @pytest.fixture + def installer(self): + """Create a PluginInstaller instance for testing.""" + return PluginInstaller() + + @pytest.fixture + def mock_config(self, config_overrides): + """Mock plugin daemon configuration.""" + config_overrides( + PLUGIN_DAEMON_URL="http://127.0.0.1:5002", + PLUGIN_DAEMON_KEY="test-key", + ) + + def test_list_plugins_success(self, installer, mock_config): + """Test successful plugin listing.""" + # Arrange + mock_response = httpx.Response( + 200, + json={ + "code": 0, + "message": "", + "data": { + "list": [], + "total": 0, + }, + }, + ) + + with patch("httpx.request", return_value=mock_response, autospec=True): + # Act + result = installer.list_plugins("test-tenant") + + # Assert + assert isinstance(result, list) + + def test_uninstall_plugin_success(self, installer, mock_config): + """Test successful plugin uninstallation.""" + # Arrange + mock_response = httpx.Response(200, json={"code": 0, "message": "", "data": True}) + + with patch("httpx.request", return_value=mock_response, autospec=True): + # Act + result = installer.uninstall("test-tenant", "plugin-installation-id") + + # Assert + assert result is True + + def test_fetch_plugin_by_identifier_success(self, installer, mock_config): + """Test successful plugin fetch by identifier.""" + # Arrange + mock_response = httpx.Response(200, json={"code": 0, "message": "", "data": True}) + + with patch("httpx.request", return_value=mock_response, autospec=True): + # Act + result = installer.fetch_plugin_by_identifier("test-tenant", "plugin-identifier") + + # Assert + assert result is True + + class TestPluginRuntimeEdgeCases: + """Tests for edge cases and corner scenarios in plugin runtime. + + Tests cover: + - Malformed responses + - Unexpected data types + - Concurrent requests + - Large payloads + """ + + @pytest.fixture + def plugin_client(self): + """Create a BasePluginClient instance for testing.""" + return BasePluginClient() + + @pytest.fixture + def mock_config(self, config_overrides): + """Mock plugin daemon configuration.""" + config_overrides( + PLUGIN_DAEMON_URL="http://127.0.0.1:5002", + PLUGIN_DAEMON_KEY="test-key", + ) + + def test_malformed_json_response(self, plugin_client, mock_config): + """Test handling of malformed JSON responses.""" + # Arrange + mock_response = httpx.Response(200, content=b"invalid JSON") + + with patch("httpx.request", return_value=mock_response, autospec=True): + # Act & Assert + with pytest.raises( + ValueError, + match=( + "Failed to parse response from plugin daemon to " + "PluginDaemonBasicResponse \\[bool\\], url: plugin/test\\-tenant/test" + ), + ): + plugin_client._request_with_plugin_daemon_response("GET", "plugin/test-tenant/test", bool) + + def test_invalid_response_structure(self, plugin_client, mock_config): + """Test handling of invalid response structure.""" + # Arrange + # Missing required fields in response + mock_response = httpx.Response(200, json={"invalid": "structure"}) + + with patch("httpx.request", return_value=mock_response, autospec=True): + # Act & Assert + with pytest.raises( + ValueError, + match=( + "Failed to parse response from plugin daemon to " + "PluginDaemonBasicResponse \\[bool\\], url: plugin/test\\-tenant/test" + ), + ): + plugin_client._request_with_plugin_daemon_response("GET", "plugin/test-tenant/test", bool) + + def test_streaming_with_invalid_json_line(self, plugin_client, mock_config): + """Test streaming with invalid JSON in one line.""" + # Arrange + stream_data = [ + 'data: {"code": 0, "message": "", "data": {"value": "valid"}}', + "data: {invalid json}", + ] + + mock_response = httpx.Response(200, text="\n".join(stream_data)) + + with patch("httpx.stream", autospec=True) as mock_stream: + mock_stream.return_value.__enter__.return_value = mock_response + + # Act + class StreamModel(BaseModel): + value: str + + results = plugin_client._request_with_plugin_daemon_response_stream( + "POST", "plugin/test-tenant/stream", StreamModel + ) + + # Assert + first_result = next(results) + assert first_result.value == "valid" + + with pytest.raises(ValueError, match="\\{invalid json\\}"): + next(results) + + def test_request_with_bytes_data(self, plugin_client, mock_config): + """Test request with bytes data.""" + # Arrange + mock_response = httpx.Response(200) + + with patch("httpx.request", return_value=mock_response, autospec=True) as mock_request: + # Act + plugin_client._request("POST", "plugin/test-tenant/upload", data=b"binary data") + + # Assert + call_kwargs = mock_request.call_args[1] + assert call_kwargs["content"] == b"binary data" + + def test_request_with_files(self, plugin_client, mock_config): + """Test request with file upload.""" + # Arrange + mock_response = httpx.Response(200) + + files = {"file": ("test.txt", b"file content", "text/plain")} + + with patch("httpx.request", return_value=mock_response, autospec=True) as mock_request: + # Act + plugin_client._request("POST", "plugin/test-tenant/upload", files=files) + + # Assert + call_kwargs = mock_request.call_args[1] + assert call_kwargs["files"] == files + + def test_streaming_empty_response(self, plugin_client, mock_config): + """Test streaming with empty response.""" + # Arrange + mock_response = httpx.Response(200, content=b"") + + with patch("httpx.stream", autospec=True) as mock_stream: + mock_stream.return_value.__enter__.return_value = mock_response + + # Act + results = list(plugin_client._stream_request("POST", "plugin/test-tenant/stream")) + + # Assert + assert len(results) == 0 + + def test_daemon_inner_error_with_code_500(self, plugin_client, mock_config): + """Test handling of daemon inner error with code -500 in stream.""" + # Arrange + error_obj = json.dumps({"error_type": "PluginDaemonInnerError", "message": "Internal error"}) + stream_data = [ + f'data: {{"code": -500, "message": {json.dumps(error_obj)}, "data": null}}', + ] + + mock_response = httpx.Response(200, text="\n".join(stream_data)) + + with patch("httpx.stream", autospec=True) as mock_stream: + mock_stream.return_value.__enter__.return_value = mock_response + + # Act & Assert + class StreamModel(BaseModel): + data: str + + results = plugin_client._request_with_plugin_daemon_response_stream( + "POST", "plugin/test-tenant/stream", StreamModel + ) + with pytest.raises(PluginDaemonInnerError) as exc_info: + next(results) + assert exc_info.value.code == -500 + + def test_non_json_error_message(self, plugin_client, mock_config): + """Test handling of non-JSON error message.""" + # Arrange + mock_response = httpx.Response(200, json={"code": -1, "message": "Plain text error message", "data": None}) + + with patch("httpx.request", return_value=mock_response, autospec=True): + # Act & Assert + with pytest.raises(ValueError, match="Plain text error message, code: \\-1") as exc_info: + plugin_client._request_with_plugin_daemon_response("GET", "plugin/test-tenant/test", bool) + assert "Plain text error message" in str(exc_info.value) + + class TestPluginRuntimeAdvancedScenarios: + """Advanced test scenarios for plugin runtime. + + Tests cover: + - Complex error recovery + - Concurrent request handling + - Plugin state management + - Advanced streaming patterns + """ + + @pytest.fixture + def plugin_client(self): + """Create a BasePluginClient instance for testing.""" + return BasePluginClient() + + @pytest.fixture + def mock_config(self, config_overrides): + """Mock plugin daemon configuration.""" + config_overrides( + PLUGIN_DAEMON_URL="http://127.0.0.1:5002", + PLUGIN_DAEMON_KEY="test-key", + ) + + def test_multiple_sequential_requests(self, plugin_client, mock_config): + """Test multiple sequential requests to the same endpoint.""" + # Arrange + mock_response = httpx.Response(200, json={"code": 0, "message": "", "data": True}) + + with patch("httpx.request", return_value=mock_response, autospec=True) as mock_request: + # Act + for i in range(5): + result = plugin_client._request_with_plugin_daemon_response( + "GET", f"plugin/test-tenant/test/{i}", bool + ) + assert result is True + + # Assert + assert mock_request.call_count == 5 + + def test_request_with_complex_nested_data(self, plugin_client, mock_config): + """Test request with complex nested data structures.""" + + # Arrange + class ComplexModel(BaseModel): + nested: dict[str, Any] + items: list[dict[str, Any]] + + complex_data = { + "nested": {"level1": {"level2": {"level3": "deep_value"}}}, + "items": [ + {"id": 1, "name": "item1"}, + {"id": 2, "name": "item2"}, + ], + } + + mock_response = httpx.Response(200, json={"code": 0, "message": "", "data": complex_data}) + + with patch("httpx.request", return_value=mock_response, autospec=True): + # Act + result = plugin_client._request_with_plugin_daemon_response( + "POST", "plugin/test-tenant/complex", ComplexModel + ) + + # Assert + assert result.nested["level1"]["level2"]["level3"] == "deep_value" + assert len(result.items) == 2 + assert result.items[0]["id"] == 1 + + def test_streaming_with_multiple_chunk_types(self, plugin_client, mock_config): + """Test streaming with different chunk types in sequence.""" + + # Arrange + class MultiTypeModel(BaseModel): + type: str + data: dict[str, Any] + + stream_data = [ + '{"code": 0, "message": "", "data": {"type": "start", "data": {"status": "initializing"}}}', + '{"code": 0, "message": "", "data": {"type": "progress", "data": {"percent": 50}}}', + '{"code": 0, "message": "", "data": {"type": "complete", "data": {"result": "success"}}}', + ] + + mock_response = httpx.Response(200, text="\n".join(stream_data)) + + with patch("httpx.stream", autospec=True) as mock_stream: + mock_stream.return_value.__enter__.return_value = mock_response + + # Act + results = list( + plugin_client._request_with_plugin_daemon_response_stream( + "POST", "plugin/test-tenant/multi-stream", MultiTypeModel + ) + ) + + # Assert + assert len(results) == 3 + assert results[0].type == "start" + assert results[1].type == "progress" + assert results[2].type == "complete" + assert results[1].data["percent"] == 50 + + def test_error_recovery_with_retry_pattern(self, plugin_client, mock_config): + """Test error recovery pattern (simulated retry logic).""" + # Arrange + call_count = 0 + + def side_effect[**P](*args: P.args, **kwargs: P.kwargs): + nonlocal call_count + call_count += 1 + if call_count < 3: + raise httpx.RequestError("Temporary failure") + mock_response = httpx.Response(200) + return mock_response + + with patch("httpx.request", side_effect=side_effect, autospec=True): + # Act & Assert - First two calls should fail + with pytest.raises(PluginDaemonInnerError): + plugin_client._request("GET", "plugin/test-tenant/test") + + with pytest.raises(PluginDaemonInnerError): + plugin_client._request("GET", "plugin/test-tenant/test") + + # Third call should succeed + response = plugin_client._request("GET", "plugin/test-tenant/test") + assert response.status_code == 200 + + def test_request_with_custom_headers_preservation(self, plugin_client, mock_config): + """Test that custom headers are preserved through request pipeline.""" + # Arrange + custom_headers = { + "X-Custom-Header": "custom-value", + "X-Request-ID": "req-123", + "X-Tenant-ID": "tenant-456", + } + + mock_response = httpx.Response(200) + + with patch("httpx.request", return_value=mock_response, autospec=True) as mock_request: + # Act + plugin_client._request("GET", "plugin/test-tenant/test", headers=custom_headers) + + # Assert + call_kwargs = mock_request.call_args[1] + for key, value in custom_headers.items(): + assert call_kwargs["headers"][key] == value + + def test_streaming_with_large_chunks(self, plugin_client, mock_config): + """Test streaming with large data chunks.""" + + # Arrange + class LargeChunkModel(BaseModel): + chunk_id: int + data: str + + # Create large chunks (simulating large data transfer) + large_data = "x" * 10000 # 10KB of data + stream_data = [ + f'{{"code": 0, "message": "", "data": {{"chunk_id": {i}, "data": "{large_data}"}}}}' for i in range(10) + ] + + mock_response = httpx.Response(200, text="\n".join(stream_data)) + + with patch("httpx.stream", autospec=True) as mock_stream: + mock_stream.return_value.__enter__.return_value = mock_response + + # Act + results = list( + plugin_client._request_with_plugin_daemon_response_stream( + "POST", "plugin/test-tenant/large-stream", LargeChunkModel + ) + ) + + # Assert + assert len(results) == 10 + for i, result in enumerate(results): + assert result.chunk_id == i + assert len(result.data) == 10000 + + class TestPluginRuntimeSecurityAndValidation: + """Tests for security and validation aspects of plugin runtime. + + Tests cover: + - Input validation + - Security headers + - Authentication failures + - Authorization checks + """ + + @pytest.fixture + def plugin_client(self): + """Create a BasePluginClient instance for testing.""" + return BasePluginClient() + + @pytest.fixture + def mock_config(self, config_overrides): + """Mock plugin daemon configuration.""" + config_overrides( + PLUGIN_DAEMON_URL="http://127.0.0.1:5002", + PLUGIN_DAEMON_KEY="secure-key-123", + ) + + def test_api_key_header_always_present(self, plugin_client, mock_config): + """Test that API key header is always included in requests.""" + # Arrange + mock_response = httpx.Response(200) + + with patch("httpx.request", return_value=mock_response, autospec=True) as mock_request: + # Act + plugin_client._request("GET", "plugin/test-tenant/test") + + # Assert + call_kwargs = mock_request.call_args[1] + assert "X-Api-Key" in call_kwargs["headers"] + assert call_kwargs["headers"]["X-Api-Key"] == "secure-key-123" + + def test_request_with_sensitive_data_in_body(self, plugin_client, mock_config): + """Test handling of sensitive data in request body.""" + # Arrange + sensitive_data = { + "api_key": "secret-api-key", + "password": "secret-password", + "credentials": {"token": "secret-token"}, + } + + mock_response = httpx.Response(200, json={"code": 0, "message": "", "data": True}) + + with patch("httpx.request", return_value=mock_response, autospec=True) as mock_request: + # Act + plugin_client._request_with_plugin_daemon_response( + "POST", + "plugin/test-tenant/validate", + bool, + data=sensitive_data, + headers={"Content-Type": "application/json"}, + ) + + # Assert - Verify data was sent + call_kwargs = mock_request.call_args[1] + assert "content" in call_kwargs or "data" in call_kwargs + + def test_unauthorized_access_with_invalid_key(self, plugin_client, mock_config): + """Test handling of unauthorized access with invalid API key.""" + # Arrange + error_message = json.dumps({"error_type": "PluginDaemonUnauthorizedError", "message": "Invalid API key"}) + mock_response = httpx.Response(200, json={"code": -1, "message": error_message, "data": None}) + + with patch("httpx.request", return_value=mock_response, autospec=True): + # Act & Assert + with pytest.raises(PluginDaemonUnauthorizedError) as exc_info: + plugin_client._request_with_plugin_daemon_response("GET", "plugin/test-tenant/test", bool) + assert "Invalid API key" in exc_info.value.description + + def test_request_parameter_validation(self, plugin_client, mock_config): + """Test validation of request parameters.""" + # Arrange + invalid_params = { + "page": -1, # Invalid negative page + "limit": 0, # Invalid zero limit + } + + error_message = json.dumps( + {"error_type": "PluginDaemonBadRequestError", "message": "Invalid parameters: page must be positive"} + ) + mock_response = httpx.Response(200, json={"code": -1, "message": error_message, "data": None}) + + with patch("httpx.request", return_value=mock_response, autospec=True): + # Act & Assert + with pytest.raises(PluginDaemonBadRequestError) as exc_info: + plugin_client._request_with_plugin_daemon_response( + "GET", "plugin/test-tenant/list", list, params=invalid_params + ) + assert "Invalid parameters" in exc_info.value.description + + def test_content_type_header_validation(self, plugin_client, mock_config): + """Test that Content-Type header is properly set for JSON requests.""" + # Arrange + mock_response = httpx.Response(200) + + with patch("httpx.request", return_value=mock_response, autospec=True) as mock_request: + # Act + plugin_client._request( + "POST", + "plugin/test-tenant/test", + headers={"Content-Type": "application/json"}, + data={"key": "value"}, + ) + + # Assert + call_kwargs = mock_request.call_args[1] + assert call_kwargs["headers"]["Content-Type"] == "application/json" + + class TestPluginRuntimePerformanceScenarios: + """Tests for performance-related scenarios in plugin runtime. + + Tests cover: + - High-volume streaming + - Concurrent operations simulation + - Memory-efficient processing + - Timeout handling under load + """ + + @pytest.fixture + def plugin_client(self): + """Create a BasePluginClient instance for testing.""" + return BasePluginClient() + + @pytest.fixture + def mock_config(self, config_overrides): + """Mock plugin daemon configuration.""" + config_overrides( + PLUGIN_DAEMON_URL="http://127.0.0.1:5002", + PLUGIN_DAEMON_KEY="test-key", + ) + + def test_high_volume_streaming(self, plugin_client, mock_config): + """Test streaming with high volume of chunks.""" + + # Arrange + class StreamChunk(BaseModel): + index: int + value: str + + # Generate 100 chunks + stream_data = [ + f'{{"code": 0, "message": "", "data": {{"index": {i}, "value": "chunk_{i}"}}}}' for i in range(100) + ] + + mock_response = httpx.Response(200, text="\n".join(stream_data)) + + with patch("httpx.stream", autospec=True) as mock_stream: + mock_stream.return_value.__enter__.return_value = mock_response + + # Act + results = list( + plugin_client._request_with_plugin_daemon_response_stream( + "POST", "plugin/test-tenant/high-volume", StreamChunk + ) + ) + + # Assert + assert len(results) == 100 + assert results[0].index == 0 + assert results[99].index == 99 + assert results[50].value == "chunk_50" + + def test_streaming_memory_efficiency(self, plugin_client, mock_config): + """Test that streaming processes chunks one at a time (memory efficient).""" + + # Arrange + class ChunkModel(BaseModel): + data: str + + processed_chunks = [] + + def process_chunk(chunk): + """Simulate processing each chunk individually.""" + processed_chunks.append(chunk.data) + return chunk + + stream_data = [f'{{"code": 0, "message": "", "data": {{"data": "chunk_{i}"}}}}' for i in range(10)] + + mock_response = httpx.Response(200, text="\n".join(stream_data)) + + with patch("httpx.stream", autospec=True) as mock_stream: + mock_stream.return_value.__enter__.return_value = mock_response + + # Act - Process chunks one by one + for chunk in plugin_client._request_with_plugin_daemon_response_stream( + "POST", "plugin/test-tenant/stream", ChunkModel + ): + process_chunk(chunk) + + # Assert + assert len(processed_chunks) == 10 + + def test_timeout_with_slow_response(self, plugin_client, mock_config): + """Test timeout handling with slow response simulation.""" + # Arrange + with patch( + "httpx.request", side_effect=httpx.TimeoutException("Request timed out after 30s"), autospec=True + ): + # Act & Assert + with pytest.raises(PluginDaemonInnerError) as exc_info: + plugin_client._request("GET", "plugin/test-tenant/slow-endpoint") + assert exc_info.value.code == -500 + + def test_concurrent_request_simulation(self, plugin_client, mock_config): + """Test simulation of concurrent requests (sequential execution in test).""" + # Arrange + mock_response = httpx.Response(200, json={"code": 0, "message": "", "data": True}) + + request_results = [] + + with patch("httpx.request", return_value=mock_response, autospec=True): + # Act - Simulate 10 concurrent requests + for i in range(10): + result = plugin_client._request_with_plugin_daemon_response( + "GET", f"plugin/test-tenant/concurrent/{i}", bool + ) + request_results.append(result) + + # Assert + assert len(request_results) == 10 + assert all(result is True for result in request_results) + + class TestPluginToolManagerAdvanced: + """Advanced tests for PluginToolManager functionality. + + Tests cover: + - Complex tool invocations + - Runtime parameter handling + - Tool provider discovery + - Advanced credential scenarios + """ + + @pytest.fixture + def tool_manager(self): + """Create a PluginToolManager instance for testing.""" + return PluginToolManager() + + @pytest.fixture + def mock_config(self, config_overrides): + """Mock plugin daemon configuration.""" + config_overrides( + PLUGIN_DAEMON_URL="http://127.0.0.1:5002", + PLUGIN_DAEMON_KEY="test-key", + ) + + def test_tool_invocation_with_complex_parameters(self, tool_manager, mock_config): + """Test tool invocation with complex parameter structures.""" + # Arrange + complex_params = { + "simple_string": "value", + "number": 42, + "boolean": True, + "nested_object": {"key1": "value1", "key2": ["item1", "item2"]}, + "array": [1, 2, 3, 4, 5], + } + + stream_data = [ + ( + 'data: {"code": 0, "message": "", "data": {"type": "text", ' + '"message": {"text": "Complex params processed"}}}' + ), + ] + + mock_response = httpx.Response(200, text="\n".join(stream_data)) + + with patch("httpx.stream", autospec=True) as mock_stream: + mock_stream.return_value.__enter__.return_value = mock_response + + # Act + results = list( + tool_manager.invoke( + tenant_id="test-tenant", + user_id="test-user", + tool_provider="langgenius/test-plugin/test-provider", + tool_name="complex-tool", + credentials={"api_key": "test-key"}, + credential_type=CredentialType.API_KEY, + tool_parameters=complex_params, + ) + ) + + # Assert + assert len(results) > 0 + + def test_tool_invocation_with_conversation_context(self, tool_manager, mock_config): + """Test tool invocation with conversation context.""" + # Arrange + stream_data = [ + 'data: {"code": 0, "message": "", "data": {"type": "text", ' + '"message": {"text": "Context-aware result"}}}', + ] + + mock_response = httpx.Response(200, text="\n".join(stream_data)) + + with patch("httpx.stream", autospec=True) as mock_stream: + mock_stream.return_value.__enter__.return_value = mock_response + + # Act + results = list( + tool_manager.invoke( + tenant_id="test-tenant", + user_id="test-user", + tool_provider="langgenius/test-plugin/test-provider", + tool_name="test-tool", + credentials={"api_key": "test-key"}, + credential_type=CredentialType.API_KEY, + tool_parameters={"query": "test"}, + conversation_id="conv-123", + app_id="app-456", + message_id="msg-789", + ) + ) + + # Assert + assert len(results) > 0 + + def test_get_runtime_parameters_success(self, tool_manager, mock_config): + """Test successful retrieval of runtime parameters.""" + # Arrange + stream_data = [ + 'data: {"code": 0, "message": "", "data": {"parameters": []}}', + ] + + mock_response = httpx.Response(200, text="\n".join(stream_data)) + + with patch("httpx.stream", autospec=True) as mock_stream: + mock_stream.return_value.__enter__.return_value = mock_response + + # Act + result = tool_manager.get_runtime_parameters( + tenant_id="test-tenant", + user_id="test-user", + provider="langgenius/test-plugin/test-provider", + credentials={"api_key": "test-key"}, + tool="test-tool", + ) + + # Assert + assert isinstance(result, list) + + def test_validate_credentials_with_oauth(self, tool_manager, mock_config): + """Test credential validation with OAuth credentials.""" + # Arrange + oauth_credentials = { + "access_token": "oauth-token-123", + "refresh_token": "refresh-token-456", + "expires_at": 1234567890, + } + + stream_data = [ + 'data: {"code": 0, "message": "", "data": {"result": true}}', + ] + + mock_response = httpx.Response(200, text="\n".join(stream_data)) + + with patch("httpx.stream", autospec=True) as mock_stream: + mock_stream.return_value.__enter__.return_value = mock_response + + # Act + result = tool_manager.validate_provider_credentials( + tenant_id="test-tenant", + user_id="test-user", + provider="langgenius/test-plugin/oauth-provider", + credentials=oauth_credentials, + ) + + # Assert + assert result is True + + class TestPluginInstallerAdvanced: + """Advanced tests for PluginInstaller functionality. + + Tests cover: + - Plugin package upload + - Bundle installation + - Plugin upgrade scenarios + - Dependency management + """ + + @pytest.fixture + def installer(self): + """Create a PluginInstaller instance for testing.""" + return PluginInstaller() + + @pytest.fixture + def mock_config(self, config_overrides): + """Mock plugin daemon configuration.""" + config_overrides( + PLUGIN_DAEMON_URL="http://127.0.0.1:5002", + PLUGIN_DAEMON_KEY="test-key", + ) + + def test_upload_plugin_package_success(self, installer, mock_config): + """Test successful plugin package upload.""" + # Arrange + plugin_package = b"fake-plugin-package-data" + mock_response = httpx.Response( + 200, + json={ + "code": 0, + "message": "", + "data": { + "unique_identifier": "test-org/test-plugin", + "manifest": { + "version": "1.0.0", + "author": "test-org", + "name": "test-plugin", + "description": {"en_US": "Test plugin"}, + "icon": "icon.png", + "label": {"en_US": "Test Plugin"}, + "created_at": "2024-01-01T00:00:00Z", + "resource": {"memory": 256}, + "plugins": {}, + "meta": {}, + }, + "verification": None, + }, + }, + ) + + with patch("httpx.request", return_value=mock_response, autospec=True): + # Act + result = installer.upload_pkg("test-tenant", plugin_package, verify_signature=False) + + # Assert + assert result.unique_identifier == "test-org/test-plugin" + + def test_fetch_plugin_readme_success(self, installer, mock_config): + """Test successful plugin readme fetch.""" + # Arrange + mock_response = httpx.Response( + 200, + json={ + "code": 0, + "message": "", + "data": {"content": "# Plugin README\n\nThis is a test plugin.", "language": "en"}, + }, + ) + + with patch("httpx.request", return_value=mock_response, autospec=True): + # Act + result = installer.fetch_plugin_readme("test-tenant", "test-org/test-plugin", "en") + + # Assert + assert "Plugin README" in result + assert "test plugin" in result + + def test_fetch_plugin_readme_not_found(self, installer, mock_config): + """Test plugin readme fetch when readme doesn't exist.""" + # Arrange + mock_response = httpx.Response(404) + + with patch("httpx.request", return_value=mock_response, autospec=True): + # Act & Assert - Should raise PluginDaemonClientSideError for 404 + with pytest.raises(PluginDaemonClientSideError): + installer.fetch_plugin_readme("test-tenant", "test-org/test-plugin", "en") + + def test_list_plugins_with_pagination(self, installer, mock_config): + """Test plugin listing with pagination.""" + # Arrange + mock_response = httpx.Response( + 200, + json={ + "code": 0, + "message": "", + "data": { + "list": [], + "total": 50, + }, + }, + ) + + with patch("httpx.request", return_value=mock_response, autospec=True): + # Act + result = installer.list_plugins_with_total("test-tenant", page=2, page_size=20) + + # Assert + assert result.total == 50 + assert isinstance(result.list, list) + + def test_check_tools_existence(self, installer, mock_config): + """Test checking existence of multiple tools.""" + # Arrange + from models.provider_ids import GenericProviderID + + provider_ids = [ + GenericProviderID("langgenius/plugin1/provider1"), + GenericProviderID("langgenius/plugin2/provider2"), + ] + + mock_response = httpx.Response(200, json={"code": 0, "message": "", "data": [True, False]}) + + with patch("httpx.request", return_value=mock_response, autospec=True): + # Act + result = installer.check_tools_existence("test-tenant", provider_ids) + + # Assert + assert len(result) == 2 + assert result[0] is True + assert result[1] is False diff --git a/api/tests/unit_tests/core/plugin/impl/test_tool_manager.py b/api/tests/unit_tests/core/plugin/impl/test_tool.py similarity index 100% rename from api/tests/unit_tests/core/plugin/impl/test_tool_manager.py rename to api/tests/unit_tests/core/plugin/impl/test_tool.py diff --git a/api/tests/unit_tests/core/plugin/impl/test_trigger_client.py b/api/tests/unit_tests/core/plugin/impl/test_trigger.py similarity index 100% rename from api/tests/unit_tests/core/plugin/impl/test_trigger_client.py rename to api/tests/unit_tests/core/plugin/impl/test_trigger.py diff --git a/api/tests/unit_tests/core/plugin/test_endpoint_client.py b/api/tests/unit_tests/core/plugin/test_endpoint_client.py deleted file mode 100644 index 082ba9b0c89..00000000000 --- a/api/tests/unit_tests/core/plugin/test_endpoint_client.py +++ /dev/null @@ -1,305 +0,0 @@ -"""Unit tests for PluginEndpointClient functionality. - -This test module covers the endpoint client operations including: -- Successful endpoint deletion -- Idempotent delete behavior (record not found) -- Non-idempotent delete behavior (other errors) - -Tests follow the Arrange-Act-Assert pattern for clarity. -""" - -from collections.abc import Callable -from unittest.mock import patch - -import httpx -import pytest - -from core.plugin.impl.endpoint import PluginEndpointClient -from core.plugin.impl.exc import PluginDaemonInternalServerError - - -@pytest.fixture(autouse=True) -def _patch_shared_httpx_client(): - """Patch module-level client methods to delegate to module httpx.request/stream.""" - - def request[**P](*args: P.args, **kwargs: P.kwargs) -> httpx.Response: - response = httpx.request(*args, **kwargs) - method, url = kwargs["method"], kwargs["url"] - assert isinstance(method, str) - assert isinstance(url, str) - response.request = httpx.Request(method, url) - return response - - with ( - patch("core.plugin.impl.base._httpx_client.request", side_effect=request), - patch("core.plugin.impl.base._httpx_client.stream", side_effect=lambda **kw: httpx.stream(**kw)), - ): - yield - - -class TestPluginEndpointClientDelete: - """Unit tests for PluginEndpointClient delete_endpoint operation. - - Tests cover: - - Successful endpoint deletion - - Idempotent behavior when endpoint is already deleted (record not found) - - Non-idempotent behavior for other errors - """ - - @pytest.fixture - def endpoint_client(self): - """Create a PluginEndpointClient instance for testing.""" - return PluginEndpointClient() - - @pytest.fixture - def mock_config(self, config_overrides: Callable[..., None]): - """Mock plugin daemon configuration.""" - config_overrides(PLUGIN_DAEMON_URL="http://127.0.0.1:5002", PLUGIN_DAEMON_KEY="test-api-key") - - def test_delete_endpoint_success(self, endpoint_client, mock_config): - """Test successful endpoint deletion. - - Given: - - A valid tenant_id, user_id, and endpoint_id - - The plugin daemon returns success response - When: - - delete_endpoint is called - Then: - - The method should return True - - The request should be made with correct parameters - """ - # Arrange - tenant_id = "tenant-123" - user_id = "user-456" - endpoint_id = "endpoint-789" - - mock_response = httpx.Response( - 200, - json={ - "code": 0, - "message": "success", - "data": True, - }, - ) - - with patch("httpx.request", return_value=mock_response, autospec=True): - # Act - result = endpoint_client.delete_endpoint( - tenant_id=tenant_id, - user_id=user_id, - endpoint_id=endpoint_id, - ) - - # Assert - assert result is True - - def test_delete_endpoint_idempotent_record_not_found(self, endpoint_client, mock_config): - """Test idempotent delete behavior when endpoint is already deleted. - - Given: - - A valid tenant_id, user_id, and endpoint_id - - The plugin daemon returns "record not found" error - When: - - delete_endpoint is called - Then: - - The method should return True (idempotent behavior) - - No exception should be raised - """ - # Arrange - tenant_id = "tenant-123" - user_id = "user-456" - endpoint_id = "endpoint-789" - - mock_response = httpx.Response( - 200, - json={ - "code": -1, - "message": ( - '{"error_type": "PluginDaemonInternalServerError", ' - '"message": "failed to remove endpoint: record not found"}' - ), - }, - ) - - with patch("httpx.request", return_value=mock_response, autospec=True): - # Act - result = endpoint_client.delete_endpoint( - tenant_id=tenant_id, - user_id=user_id, - endpoint_id=endpoint_id, - ) - - # Assert - should return True instead of raising an error - assert result is True - - def test_delete_endpoint_non_idempotent_other_errors(self, endpoint_client, mock_config): - """Test non-idempotent delete behavior for other errors. - - Given: - - A valid tenant_id, user_id, and endpoint_id - - The plugin daemon returns a different error (not "record not found") - When: - - delete_endpoint is called - Then: - - The method should raise PluginDaemonInternalServerError - """ - # Arrange - tenant_id = "tenant-123" - user_id = "user-456" - endpoint_id = "endpoint-789" - - mock_response = httpx.Response( - 200, - json={ - "code": -1, - "message": ( - '{"error_type": "PluginDaemonInternalServerError", ' - '"message": "failed to remove endpoint: internal server error"}' - ), - }, - ) - - with patch("httpx.request", return_value=mock_response, autospec=True): - # Act & Assert - with pytest.raises(PluginDaemonInternalServerError) as exc_info: - endpoint_client.delete_endpoint( - tenant_id=tenant_id, - user_id=user_id, - endpoint_id=endpoint_id, - ) - - # Assert - the error message should not be "record not found" - assert "record not found" not in str(exc_info.value.description) - - def test_delete_endpoint_idempotent_case_insensitive(self, endpoint_client, mock_config): - """Test idempotent delete behavior with case-insensitive error message. - - Given: - - A valid tenant_id, user_id, and endpoint_id - - The plugin daemon returns "Record Not Found" error (different case) - When: - - delete_endpoint is called - Then: - - The method should return True (idempotent behavior) - """ - # Arrange - tenant_id = "tenant-123" - user_id = "user-456" - endpoint_id = "endpoint-789" - - mock_response = httpx.Response( - 200, - json={ - "code": -1, - "message": '{"error_type": "PluginDaemonInternalServerError", "message": "Record Not Found"}', - }, - ) - - with patch("httpx.request", return_value=mock_response, autospec=True): - # Act - result = endpoint_client.delete_endpoint( - tenant_id=tenant_id, - user_id=user_id, - endpoint_id=endpoint_id, - ) - - # Assert - should still return True - assert result is True - - def test_delete_endpoint_multiple_calls_idempotent(self, endpoint_client, mock_config): - """Test that multiple delete calls are idempotent. - - Given: - - A valid tenant_id, user_id, and endpoint_id - - The first call succeeds - - Subsequent calls return "record not found" - When: - - delete_endpoint is called multiple times - Then: - - All calls should return True - """ - # Arrange - tenant_id = "tenant-123" - user_id = "user-456" - endpoint_id = "endpoint-789" - - # First call - success - mock_response_success = httpx.Response( - 200, - json={ - "code": 0, - "message": "success", - "data": True, - }, - ) - - # Second call - record not found - mock_response_not_found = httpx.Response( - 200, - json={ - "code": -1, - "message": ( - '{"error_type": "PluginDaemonInternalServerError", ' - '"message": "failed to remove endpoint: record not found"}' - ), - }, - ) - - with patch("httpx.request", autospec=True) as mock_request: - # Act - first call - mock_request.return_value = mock_response_success - result1 = endpoint_client.delete_endpoint( - tenant_id=tenant_id, - user_id=user_id, - endpoint_id=endpoint_id, - ) - - # Act - second call (already deleted) - mock_request.return_value = mock_response_not_found - result2 = endpoint_client.delete_endpoint( - tenant_id=tenant_id, - user_id=user_id, - endpoint_id=endpoint_id, - ) - - # Assert - both should return True - assert result1 is True - assert result2 is True - - def test_delete_endpoint_non_idempotent_unauthorized_error(self, endpoint_client, mock_config): - """Test that authorization errors are not treated as idempotent. - - Given: - - A valid tenant_id, user_id, and endpoint_id - - The plugin daemon returns an unauthorized error - When: - - delete_endpoint is called - Then: - - The method should raise the appropriate error (not return True) - """ - # Arrange - tenant_id = "tenant-123" - user_id = "user-456" - endpoint_id = "endpoint-789" - - mock_response = httpx.Response( - 200, - json={ - "code": -1, - "message": '{"error_type": "PluginDaemonUnauthorizedError", "message": "unauthorized access"}', - }, - ) - - with patch("httpx.request", return_value=mock_response, autospec=True): - # Act & Assert - with pytest.raises( - Exception, match="req_id: PluginDaemonUnauthorizedError: unauthorized access" - ) as exc_info: - endpoint_client.delete_endpoint( - tenant_id=tenant_id, - user_id=user_id, - endpoint_id=endpoint_id, - ) - - # Assert - should not return True for unauthorized errors - assert exc_info.value.__class__.__name__ == "PluginDaemonUnauthorizedError" diff --git a/api/tests/unit_tests/core/plugin/test_model_billing_attribution.py b/api/tests/unit_tests/core/plugin/test_model_billing_attribution.py deleted file mode 100644 index 431cb8e1e3d..00000000000 --- a/api/tests/unit_tests/core/plugin/test_model_billing_attribution.py +++ /dev/null @@ -1,207 +0,0 @@ -"""Exercise backwards invocation through the real hosted quota owner, with no external I/O.""" - -from collections.abc import Callable, Iterator -from dataclasses import dataclass -from unittest.mock import patch - -import pytest -from pydantic import BaseModel - -from configs import dify_config -from core.entities.provider_configuration import ProviderConfiguration, ProviderModelBundle -from core.entities.provider_entities import ( - CustomConfiguration, - ProviderQuotaType, - QuotaConfiguration, - QuotaUnit, - SystemConfiguration, -) -from core.model_context import get_credit_usage_metadata -from core.model_manager import ModelInstance -from core.plugin.backwards_invocation.model import PluginModelBackwardsInvocation -from core.plugin.entities.request import ( - RequestInvokeLLM, - RequestInvokeMultimodalEmbedding, - RequestInvokeRerank, - RequestInvokeTextEmbedding, -) -from core.plugin.impl.model_runtime_factory import create_plugin_model_runtime -from enums import DeploymentEdition -from graphon.model_runtime.entities.common_entities import I18nObject -from graphon.model_runtime.entities.llm_entities import LLMResultChunk, LLMResultChunkDelta, LLMUsage -from graphon.model_runtime.entities.message_entities import AssistantPromptMessage -from graphon.model_runtime.entities.model_entities import ModelType -from graphon.model_runtime.entities.provider_entities import ProviderEntity -from graphon.model_runtime.model_providers.base.large_language_model import LargeLanguageModel -from graphon.model_runtime.model_providers.base.rerank_model import RerankModel -from graphon.model_runtime.model_providers.base.text_embedding_model import TextEmbeddingModel -from models.account import Tenant -from models.provider import ProviderType -from services.billing_service import BillingService -from tests.unit_tests.knowledge_fs_core_fakes import Marker - -_MODEL_TYPE_CLASSES: dict[ModelType, type[LargeLanguageModel] | type[TextEmbeddingModel] | type[RerankModel]] = { - ModelType.LLM: LargeLanguageModel, - ModelType.TEXT_EMBEDDING: TextEmbeddingModel, - ModelType.RERANK: RerankModel, -} - - -@dataclass(frozen=True) -class _ProviderManager: - bundle: ProviderModelBundle - - def get_provider_model_bundle(self, **_kwargs: object) -> ProviderModelBundle: - return self.bundle - - def get_configurations(self, _tenant_id: str) -> dict[str, ProviderConfiguration]: - return {self.bundle.configuration.provider.provider: self.bundle.configuration} - - -def _tenant() -> Tenant: - tenant = Tenant(name="Test Workspace") - tenant.id = "tenant-1" - return tenant - - -def _hosted_bundle(model_type: ModelType) -> ProviderModelBundle: - provider = ProviderEntity( - provider="langgenius/openai/openai", - label=I18nObject(en_US="OpenAI"), - supported_model_types=[model_type], - configurate_methods=[], - ) - configuration = ProviderConfiguration( - tenant_id="tenant-1", - provider=provider, - preferred_provider_type=ProviderType.SYSTEM, - using_provider_type=ProviderType.SYSTEM, - system_configuration=SystemConfiguration( - enabled=True, - current_quota_type=ProviderQuotaType.TRIAL, - quota_configurations=[ - QuotaConfiguration( - quota_type=ProviderQuotaType.TRIAL, - quota_unit=QuotaUnit.CREDITS, - quota_limit=200, - quota_used=0, - is_valid=True, - restrict_models=[], - ) - ], - ), - custom_configuration=CustomConfiguration(provider=None), - model_settings=[], - ) - model_type_instance = _MODEL_TYPE_CLASSES[model_type]( - provider_schema=provider, - model_runtime=create_plugin_model_runtime(tenant_id="tenant-1", user_id="synthetic-end-user"), - ) - return ProviderModelBundle(configuration=configuration, model_type_instance=model_type_instance) - - -@pytest.mark.parametrize( - ("payload_type", "invoke_method", "provider_method", "model_type", "body"), - [ - ( - RequestInvokeLLM, - PluginModelBackwardsInvocation.invoke_llm, - "invoke_llm", - ModelType.LLM, - {"mode": "chat", "stream": True}, - ), - ( - RequestInvokeTextEmbedding, - PluginModelBackwardsInvocation.invoke_text_embedding, - "invoke_text_embedding", - ModelType.TEXT_EMBEDDING, - {"texts": ["query"], "input_type": "query"}, - ), - ( - RequestInvokeMultimodalEmbedding, - PluginModelBackwardsInvocation.invoke_multimodal_embedding, - "invoke_multimodal_embedding", - ModelType.TEXT_EMBEDDING, - {"documents": [{"content": "AQID", "content_type": "image"}]}, - ), - ( - RequestInvokeRerank, - PluginModelBackwardsInvocation.invoke_rerank, - "invoke_rerank", - ModelType.RERANK, - {"query": "query", "docs": ["doc"]}, - ), - ], -) -@pytest.mark.parametrize( - "created_by", - ["knowledge_indexing", "knowledge_retrieval", "app", None], - ids=["indexing", "retrieval", "answer", "legacy-plugin"], -) -def test_model_invocation_preserves_billing_origin( - payload_type: type[BaseModel], - invoke_method: Callable[..., object], - provider_method: str, - model_type: ModelType, - body: dict[str, object], - created_by: str | None, - config_overrides: Callable[..., None], -) -> None: - config_overrides(DEPLOYMENT_EDITION=DeploymentEdition.CLOUD) - provider_manager = _ProviderManager(bundle=_hosted_bundle(model_type)) - metadata = {"app_type": "new_rag", "created_by": created_by} if created_by else {} - payload = payload_type.model_validate( - {"provider": "langgenius/openai/openai", "model": "selected-model", **body, **metadata} - ) - requests: list[tuple[str, dict[str, object]]] = [] - - def billing_transport(method: str, endpoint: str, *, json: dict[str, object]) -> dict[str, object]: - assert method == "POST" - requests.append((endpoint, json)) - if endpoint == "/quota/reserve": - return {"reservation_id": "reservation-1", "available": 197, "reserved": 3} - assert endpoint == "/quota/commit" - return {"available": 197, "reserved": 0, "refunded": 0} - - chunk = LLMResultChunk( - model="selected-model", - prompt_messages=[], - delta=LLMResultChunkDelta( - index=0, message=AssistantPromptMessage(content="OK"), usage=LLMUsage.empty_usage(), finish_reason="stop" - ), - ) - provider_result = (item for item in [chunk]) if model_type == ModelType.LLM else Marker("provider-result") - with ( - patch("core.model_manager.create_plugin_provider_manager", return_value=provider_manager), - patch("core.app.llm.quota.create_plugin_provider_manager", return_value=provider_manager), - patch.object(type(dify_config), "get_model_credits", return_value=3), - patch.object(ProviderConfiguration, "get_current_credentials", return_value={"api_key": "local-test-only"}), - patch.object(ProviderConfiguration, "get_provider_model", return_value=None), - patch.object(ModelInstance, provider_method, return_value=provider_result), - patch.object(BillingService, "_send_quota_request", side_effect=billing_transport), - ): - response = invoke_method("synthetic-end-user", _tenant(), payload) - assert get_credit_usage_metadata() is None - if model_type == ModelType.LLM: - assert isinstance(response, Iterator) - first_chunk = list(response)[0] - assert isinstance(first_chunk, LLMResultChunk) - assert first_chunk.delta.message.content == "OK" - - assert [endpoint for endpoint, _ in requests] == ["/quota/reserve", "/quota/commit"] - expected_metadata = { - "source": "llm.invoke" if model_type == ModelType.LLM else "model.invoke", - "provider": "langgenius/openai/openai", - "model": "selected-model", - "app_type": "new_rag" if created_by else "unknown", - "created_by": created_by or "plugin_api", - } - if model_type != ModelType.LLM: - expected_metadata["model_type"] = model_type.value - reserve, commit = requests[0][1], requests[1][1] - assert reserve["tenant_id"] == commit["tenant_id"] == "tenant-1" - assert reserve["bucket"] == commit["bucket"] == "trial" - assert reserve["feature_key"] == commit["feature_key"] == "credit_pool" - assert reserve["amount"] == commit["actual_amount"] == 3 - assert reserve["meta"] == expected_metadata - assert commit["meta"] == {**expected_metadata, "request_id": reserve["request_id"]} diff --git a/api/tests/unit_tests/core/plugin/test_plugin_manager.py b/api/tests/unit_tests/core/plugin/test_plugin_manager.py deleted file mode 100644 index c08dda3d3b1..00000000000 --- a/api/tests/unit_tests/core/plugin/test_plugin_manager.py +++ /dev/null @@ -1,1470 +0,0 @@ -""" -Unit tests for Plugin Manager (PluginInstaller). - -This module tests the plugin management functionality including: -- Plugin discovery and listing -- Plugin loading and installation -- Plugin validation and manifest parsing -- Version compatibility checks -- Dependency resolution -""" - -import datetime -from unittest.mock import patch - -import httpx -import pytest -from packaging.version import Version -from requests import HTTPError - -from core.plugin.entities.bundle import PluginBundleDependency, PluginBundleDependencyType -from core.plugin.entities.plugin import ( - MissingPluginDependency, - PluginCategory, - PluginDeclaration, - PluginEntity, - PluginInstallation, - PluginInstallationSource, - PluginResourceRequirements, -) -from core.plugin.entities.plugin_daemon import ( - PluginDecodeResponse, - PluginInstalledIdsDaemonResponse, - PluginInstallTask, - PluginInstallTaskStartResponse, - PluginInstallTaskStatus, - PluginListResponse, - PluginListWithoutTotalResponse, - PluginReadmeResponse, - PluginVerification, -) -from core.plugin.impl.exc import ( - PluginDaemonBadRequestError, - PluginDaemonInternalServerError, - PluginDaemonNotFoundError, -) -from core.plugin.impl.plugin import PluginInstaller -from core.tools.entities.common_entities import I18nObject -from models.provider_ids import GenericProviderID - - -class TestPluginDiscovery: - """Test plugin discovery functionality.""" - - @pytest.fixture - def plugin_installer(self): - """Create a PluginInstaller instance for testing.""" - return PluginInstaller() - - @pytest.fixture - def mock_plugin_entity(self): - """Create a mock PluginEntity for testing.""" - return PluginEntity( - id="entity-123", - created_at=datetime.datetime(2023, 1, 1, 0, 0, 0), - updated_at=datetime.datetime(2023, 1, 1, 0, 0, 0), - tenant_id="test-tenant", - endpoints_setups=0, - endpoints_active=0, - runtime_type="remote", - source=PluginInstallationSource.Marketplace, - meta={}, - plugin_id="plugin-123", - plugin_unique_identifier="test-org/test-plugin/1.0.0", - version="1.0.0", - checksum="abc123", - name="Test Plugin", - installation_id="install-123", - declaration=PluginDeclaration( - version="1.0.0", - author="test-author", - name="test-plugin", - description=I18nObject(en_US="Test plugin description", zh_Hans="测试插件描述"), - icon="icon.png", - label=I18nObject(en_US="Test Plugin", zh_Hans="测试插件"), - category=PluginCategory.Tool, - created_at=datetime.datetime.now(), - resource=PluginResourceRequirements(memory=512, permission=None), - plugins=PluginDeclaration.Plugins(), - meta=PluginDeclaration.Meta(version="1.0.0"), - ), - ) - - def test_list_plugins_success(self, plugin_installer, mock_plugin_entity): - """Test successful plugin listing.""" - # Arrange: Mock the HTTP response for listing plugins - mock_response = PluginListResponse(list=[mock_plugin_entity], total=1) - - with patch.object( - plugin_installer, "_request_with_plugin_daemon_response", return_value=mock_response - ) as mock_request: - # Act: List plugins for a tenant - result = plugin_installer.list_plugins("test-tenant") - - # Assert: Verify the request was made correctly - mock_request.assert_called_once() - assert len(result) == 1 - assert result[0].plugin_id == "plugin-123" - assert result[0].name == "Test Plugin" - - def test_list_plugins_with_pagination(self, plugin_installer, mock_plugin_entity): - """Test plugin listing with pagination support.""" - # Arrange: Mock paginated response - mock_response = PluginListResponse(list=[mock_plugin_entity], total=10) - - with patch.object( - plugin_installer, "_request_with_plugin_daemon_response", return_value=mock_response - ) as mock_request: - # Act: List plugins with pagination - result = plugin_installer.list_plugins_with_total("test-tenant", page=1, page_size=5) - - # Assert: Verify pagination parameters - mock_request.assert_called_once() - call_args = mock_request.call_args - assert call_args[1]["params"]["page"] == 1 - assert call_args[1]["params"]["page_size"] == 5 - assert result.total == 10 - - def test_list_plugins_by_category(self, plugin_installer, mock_plugin_entity): - """Test category plugin listing without total.""" - mock_response = PluginListWithoutTotalResponse(list=[mock_plugin_entity], has_more=True) - - with patch.object( - plugin_installer, "_request_with_plugin_daemon_response", return_value=mock_response - ) as mock_request: - result = plugin_installer.list_plugins_by_category( - "test-tenant", - category=PluginCategory.Tool, - page=2, - page_size=10, - query="weather", - tags=["search", "rag"], - language="zh_Hans", - ) - - mock_request.assert_called_once() - call_args = mock_request.call_args - assert call_args.args[1] == "plugin/test-tenant/management/tool/list" - assert call_args.args[2] is PluginListWithoutTotalResponse - assert call_args.kwargs["params"]["page"] == 2 - assert call_args.kwargs["params"]["page_size"] == 10 - assert call_args.kwargs["params"]["query"] == "weather" - assert call_args.kwargs["params"]["tags"] == ["search", "rag"] - assert call_args.kwargs["params"]["language"] == "zh_Hans" - assert result.list == [mock_plugin_entity] - assert result.has_more is True - - def test_list_plugins_empty_result(self, plugin_installer): - """Test plugin listing when no plugins are installed.""" - # Arrange: Mock empty response - mock_response = PluginListResponse(list=[], total=0) - - with patch.object(plugin_installer, "_request_with_plugin_daemon_response", return_value=mock_response): - # Act: List plugins - result = plugin_installer.list_plugins("test-tenant") - - # Assert: Verify empty list is returned - assert len(result) == 0 - - def test_list_installed_plugin_ids(self, plugin_installer): - """The lightweight ID endpoint is unpaginated and does not request plugin details.""" - mock_response = PluginInstalledIdsDaemonResponse(plugin_ids=["langgenius/openai", "langgenius/anthropic"]) - - with patch.object( - plugin_installer, "_request_with_plugin_daemon_response", return_value=mock_response - ) as mock_request: - result = plugin_installer.list_installed_plugin_ids("test-tenant", PluginCategory.Tool) - - mock_request.assert_called_once_with( - "GET", - "plugin/test-tenant/management/installation/ids", - PluginInstalledIdsDaemonResponse, - params={"category": "tool"}, - ) - assert result == ["langgenius/openai", "langgenius/anthropic"] - - def test_fetch_plugin_by_identifier_found(self, plugin_installer): - """Test fetching a plugin by its unique identifier when it exists.""" - # Arrange: Mock successful fetch - with patch.object(plugin_installer, "_request_with_plugin_daemon_response", return_value=True) as mock_request: - # Act: Fetch plugin by identifier - result = plugin_installer.fetch_plugin_by_identifier("test-tenant", "test-org/test-plugin/1.0.0") - - # Assert: Verify the plugin was found - assert result is True - mock_request.assert_called_once() - - def test_fetch_plugin_by_identifier_not_found(self, plugin_installer): - """Test fetching a plugin by identifier when it doesn't exist.""" - # Arrange: Mock not found response - with patch.object(plugin_installer, "_request_with_plugin_daemon_response", return_value=False): - # Act: Fetch non-existent plugin - result = plugin_installer.fetch_plugin_by_identifier("test-tenant", "non-existent/plugin/1.0.0") - - # Assert: Verify the plugin was not found - assert result is False - - -class TestPluginLoading: - """Test plugin loading and installation functionality.""" - - @pytest.fixture - def plugin_installer(self): - """Create a PluginInstaller instance for testing.""" - return PluginInstaller() - - @pytest.fixture - def mock_plugin_declaration(self): - """Create a mock PluginDeclaration for testing.""" - return PluginDeclaration( - version="1.0.0", - author="test-author", - name="test-plugin", - description=I18nObject(en_US="Test plugin", zh_Hans="测试插件"), - icon="icon.png", - label=I18nObject(en_US="Test Plugin", zh_Hans="测试插件"), - category=PluginCategory.Tool, - created_at=datetime.datetime.now(), - resource=PluginResourceRequirements(memory=512, permission=None), - plugins=PluginDeclaration.Plugins(), - meta=PluginDeclaration.Meta(version="1.0.0"), - ) - - def test_upload_pkg_success(self, plugin_installer, mock_plugin_declaration): - """Test successful plugin package upload.""" - # Arrange: Create mock package data and expected response - pkg_data = b"mock-plugin-package-data" - mock_response = PluginDecodeResponse( - unique_identifier="test-org/test-plugin/1.0.0", - manifest=mock_plugin_declaration, - verification=PluginVerification(authorized_category=PluginVerification.AuthorizedCategory.Community), - ) - - with patch.object( - plugin_installer, "_request_with_plugin_daemon_response", return_value=mock_response - ) as mock_request: - # Act: Upload plugin package - result = plugin_installer.upload_pkg("test-tenant", pkg_data, verify_signature=False) - - # Assert: Verify upload was successful - assert result.unique_identifier == "test-org/test-plugin/1.0.0" - assert result.manifest.name == "test-plugin" - mock_request.assert_called_once() - - def test_upload_pkg_with_signature_verification(self, plugin_installer, mock_plugin_declaration): - """Test plugin package upload with signature verification enabled.""" - # Arrange: Create mock package data - pkg_data = b"signed-plugin-package" - mock_response = PluginDecodeResponse( - unique_identifier="verified-org/verified-plugin/1.0.0", - manifest=mock_plugin_declaration, - verification=PluginVerification(authorized_category=PluginVerification.AuthorizedCategory.Partner), - ) - - with patch.object( - plugin_installer, "_request_with_plugin_daemon_response", return_value=mock_response - ) as mock_request: - # Act: Upload with signature verification - result = plugin_installer.upload_pkg("test-tenant", pkg_data, verify_signature=True) - - # Assert: Verify signature verification was requested - call_args = mock_request.call_args - assert call_args[1]["data"]["verify_signature"] == "true" - assert result.verification.authorized_category == PluginVerification.AuthorizedCategory.Partner - - def test_install_from_identifiers_success(self, plugin_installer): - """Test successful plugin installation from identifiers.""" - # Arrange: Mock installation response - mock_response = PluginInstallTaskStartResponse(all_installed=False, task_id="task-123") - - with patch.object( - plugin_installer, "_request_with_plugin_daemon_response", return_value=mock_response - ) as mock_request: - # Act: Install plugins from identifiers - result = plugin_installer.install_from_identifiers( - tenant_id="test-tenant", - identifiers=["plugin1/1.0.0", "plugin2/2.0.0"], - source=PluginInstallationSource.Marketplace, - metas=[{"key": "value1"}, {"key": "value2"}], - ) - - # Assert: Verify installation task was created - assert result.task_id == "task-123" - assert result.all_installed is False - mock_request.assert_called_once() - - def test_install_from_identifiers_all_installed(self, plugin_installer): - """Test installation when all plugins are already installed.""" - # Arrange: Mock response indicating all plugins are installed - mock_response = PluginInstallTaskStartResponse(all_installed=True, task_id="") - - with patch.object(plugin_installer, "_request_with_plugin_daemon_response", return_value=mock_response): - # Act: Attempt to install already-installed plugins - result = plugin_installer.install_from_identifiers( - tenant_id="test-tenant", - identifiers=["existing-plugin/1.0.0"], - source=PluginInstallationSource.Package, - metas=[{}], - ) - - # Assert: Verify all_installed flag is True - assert result.all_installed is True - - def test_fetch_plugin_installation_task(self, plugin_installer): - """Test fetching a specific plugin installation task.""" - # Arrange: Mock installation task - mock_task = PluginInstallTask( - id="task-123", - created_at=datetime.datetime.now(), - updated_at=datetime.datetime.now(), - status=PluginInstallTaskStatus.Running, - total_plugins=3, - completed_plugins=1, - plugins=[], - ) - - with patch.object( - plugin_installer, "_request_with_plugin_daemon_response", return_value=mock_task - ) as mock_request: - # Act: Fetch installation task - result = plugin_installer.fetch_plugin_installation_task("test-tenant", "task-123") - - # Assert: Verify task details - assert result.status == PluginInstallTaskStatus.Running - assert result.total_plugins == 3 - assert result.completed_plugins == 1 - mock_request.assert_called_once() - - def test_uninstall_plugin_success(self, plugin_installer): - """Test successful plugin uninstallation.""" - # Arrange: Mock successful uninstall - with patch.object(plugin_installer, "_request_with_plugin_daemon_response", return_value=True) as mock_request: - # Act: Uninstall plugin - result = plugin_installer.uninstall("test-tenant", "install-123") - - # Assert: Verify uninstallation succeeded - assert result is True - mock_request.assert_called_once() - - def test_upgrade_plugin_success(self, plugin_installer): - """Test successful plugin upgrade.""" - # Arrange: Mock upgrade response - mock_response = PluginInstallTaskStartResponse(all_installed=False, task_id="upgrade-task-123") - - with patch.object( - plugin_installer, "_request_with_plugin_daemon_response", return_value=mock_response - ) as mock_request: - # Act: Upgrade plugin - result = plugin_installer.upgrade_plugin( - tenant_id="test-tenant", - original_plugin_unique_identifier="plugin/1.0.0", - new_plugin_unique_identifier="plugin/2.0.0", - source=PluginInstallationSource.Marketplace, - meta={"upgrade": "true"}, - ) - - # Assert: Verify upgrade task was created - assert result.task_id == "upgrade-task-123" - mock_request.assert_called_once() - - -class TestPluginValidation: - """Test plugin validation and manifest parsing.""" - - @pytest.fixture - def plugin_installer(self): - """Create a PluginInstaller instance for testing.""" - return PluginInstaller() - - def test_fetch_plugin_manifest_success(self, plugin_installer): - """Test successful plugin manifest fetching.""" - # Arrange: Create a valid plugin declaration - mock_manifest = PluginDeclaration( - version="1.0.0", - author="test-author", - name="test-plugin", - description=I18nObject(en_US="Test plugin", zh_Hans="测试插件"), - icon="icon.png", - label=I18nObject(en_US="Test Plugin", zh_Hans="测试插件"), - category=PluginCategory.Tool, - created_at=datetime.datetime.now(), - resource=PluginResourceRequirements(memory=512, permission=None), - plugins=PluginDeclaration.Plugins(), - meta=PluginDeclaration.Meta(version="1.0.0", minimum_dify_version="0.6.0"), - ) - - with patch.object( - plugin_installer, "_request_with_plugin_daemon_response", return_value=mock_manifest - ) as mock_request: - # Act: Fetch plugin manifest - result = plugin_installer.fetch_plugin_manifest("test-tenant", "test-org/test-plugin/1.0.0") - - # Assert: Verify manifest was fetched correctly - assert result.name == "test-plugin" - assert result.version == "1.0.0" - assert result.author == "test-author" - assert result.meta.minimum_dify_version == "0.6.0" - mock_request.assert_called_once() - - def test_decode_plugin_from_identifier(self, plugin_installer): - """Test decoding plugin information from identifier.""" - # Arrange: Create mock decode response - mock_declaration = PluginDeclaration( - version="2.0.0", - author="decode-author", - name="decode-plugin", - description=I18nObject(en_US="Decoded plugin", zh_Hans="解码插件"), - icon="icon.png", - label=I18nObject(en_US="Decode Plugin", zh_Hans="解码插件"), - category=PluginCategory.Model, - created_at=datetime.datetime.now(), - resource=PluginResourceRequirements(memory=1024, permission=None), - plugins=PluginDeclaration.Plugins(), - meta=PluginDeclaration.Meta(version="2.0.0"), - ) - - mock_response = PluginDecodeResponse( - unique_identifier="org/decode-plugin/2.0.0", - manifest=mock_declaration, - verification=None, - ) - - with patch.object(plugin_installer, "_request_with_plugin_daemon_response", return_value=mock_response): - # Act: Decode plugin from identifier - result = plugin_installer.decode_plugin_from_identifier("test-tenant", "org/decode-plugin/2.0.0") - - # Assert: Verify decoded information - assert result.unique_identifier == "org/decode-plugin/2.0.0" - assert result.manifest.name == "decode-plugin" - # Category will be Extension unless a model provider entity is provided - assert result.manifest.category == PluginCategory.Extension - - def test_plugin_manifest_invalid_version_format(self): - """Test that invalid version format raises validation error.""" - # Arrange & Act & Assert: Creating a declaration with invalid version should fail - with pytest.raises(ValueError, match="Invalid version format"): - PluginDeclaration( - version="invalid-version", # Invalid version format - author="test-author", - name="test-plugin", - description=I18nObject(en_US="Test", zh_Hans="测试"), - icon="icon.png", - label=I18nObject(en_US="Test", zh_Hans="测试"), - category=PluginCategory.Tool, - created_at=datetime.datetime.now(), - resource=PluginResourceRequirements(memory=512, permission=None), - plugins=PluginDeclaration.Plugins(), - meta=PluginDeclaration.Meta(version="1.0.0"), - ) - - def test_plugin_manifest_invalid_author_format(self): - """Test that invalid author format raises validation error.""" - # Arrange & Act & Assert: Creating a declaration with invalid author should fail - with pytest.raises(ValueError, match="1 validation error for PluginDeclaration\\\nauthor"): - PluginDeclaration( - version="1.0.0", - author="invalid author with spaces!@#", # Invalid author format - name="test-plugin", - description=I18nObject(en_US="Test", zh_Hans="测试"), - icon="icon.png", - label=I18nObject(en_US="Test", zh_Hans="测试"), - category=PluginCategory.Tool, - created_at=datetime.datetime.now(), - resource=PluginResourceRequirements(memory=512, permission=None), - plugins=PluginDeclaration.Plugins(), - meta=PluginDeclaration.Meta(version="1.0.0"), - ) - - def test_plugin_manifest_invalid_name_format(self): - """Test that invalid plugin name format raises validation error.""" - # Arrange & Act & Assert: Creating a declaration with invalid name should fail - with pytest.raises(ValueError, match="1 validation error for PluginDeclaration\\\nname"): - PluginDeclaration( - version="1.0.0", - author="test-author", - name="Invalid_Plugin_Name_With_Uppercase", # Invalid name format - description=I18nObject(en_US="Test", zh_Hans="测试"), - icon="icon.png", - label=I18nObject(en_US="Test", zh_Hans="测试"), - category=PluginCategory.Tool, - created_at=datetime.datetime.now(), - resource=PluginResourceRequirements(memory=512, permission=None), - plugins=PluginDeclaration.Plugins(), - meta=PluginDeclaration.Meta(version="1.0.0"), - ) - - def test_fetch_plugin_readme_success(self, plugin_installer): - """Test successful plugin readme fetching.""" - # Arrange: Mock readme response - mock_response = PluginReadmeResponse(content="# Test Plugin\n\nThis is a test plugin.", language="en_US") - - with patch.object(plugin_installer, "_request_with_plugin_daemon_response", return_value=mock_response): - # Act: Fetch plugin readme - result = plugin_installer.fetch_plugin_readme("test-tenant", "test-org/test-plugin/1.0.0", "en_US") - - # Assert: Verify readme content - assert result == "# Test Plugin\n\nThis is a test plugin." - - def test_fetch_plugin_readme_not_found(self, plugin_installer): - """Test fetching readme when it doesn't exist (404 error).""" - # Arrange: Mock HTTP 404 error - the actual implementation catches HTTPError from requests library - mock_error = HTTPError("404 Not Found") - - with patch.object(plugin_installer, "_request_with_plugin_daemon_response", side_effect=mock_error): - # Act: Fetch non-existent readme - result = plugin_installer.fetch_plugin_readme("test-tenant", "test-org/test-plugin/1.0.0", "en_US") - - # Assert: Verify empty string is returned for 404 - assert result == "" - - -class TestVersionCompatibility: - """Test version compatibility checks.""" - - def test_valid_version_format(self): - """Test that valid semantic versions are accepted.""" - # Arrange & Act: Create declarations with various valid version formats - valid_versions = ["1.0.0", "2.1.3", "0.0.1", "10.20.30"] - - for version in valid_versions: - # Assert: All valid versions should be accepted - declaration = PluginDeclaration( - version=version, - author="test-author", - name="test-plugin", - description=I18nObject(en_US="Test", zh_Hans="测试"), - icon="icon.png", - label=I18nObject(en_US="Test", zh_Hans="测试"), - category=PluginCategory.Tool, - created_at=datetime.datetime.now(), - resource=PluginResourceRequirements(memory=512, permission=None), - plugins=PluginDeclaration.Plugins(), - meta=PluginDeclaration.Meta(version=version), - ) - assert declaration.version == version - - def test_minimum_dify_version_validation(self): - """Test minimum Dify version validation.""" - # Arrange & Act: Create declaration with minimum Dify version - declaration = PluginDeclaration( - version="1.0.0", - author="test-author", - name="test-plugin", - description=I18nObject(en_US="Test", zh_Hans="测试"), - icon="icon.png", - label=I18nObject(en_US="Test", zh_Hans="测试"), - category=PluginCategory.Tool, - created_at=datetime.datetime.now(), - resource=PluginResourceRequirements(memory=512, permission=None), - plugins=PluginDeclaration.Plugins(), - meta=PluginDeclaration.Meta(version="1.0.0", minimum_dify_version="0.6.0"), - ) - - # Assert: Verify minimum version is set correctly - assert declaration.meta.minimum_dify_version == "0.6.0" - - def test_invalid_minimum_dify_version(self): - """Test that invalid minimum Dify version format raises error.""" - # Arrange & Act & Assert: Invalid minimum version should raise ValueError - with pytest.raises(ValueError, match="Invalid version format"): - PluginDeclaration.Meta(version="1.0.0", minimum_dify_version="invalid.version") - - def test_version_comparison_logic(self): - """Test version comparison using packaging.version.Version.""" - # Arrange: Create version objects for comparison - v1 = Version("1.0.0") - v2 = Version("2.0.0") - v3 = Version("1.5.0") - - # Act & Assert: Verify version comparison works correctly - assert v1 < v2 - assert v2 > v1 - assert v1 < v3 < v2 - assert v1 == Version("1.0.0") - - def test_plugin_upgrade_version_check(self): - """Test that plugin upgrade requires newer version.""" - # Arrange: Define old and new versions - old_version = Version("1.0.0") - new_version = Version("2.0.0") - same_version = Version("1.0.0") - - # Act & Assert: Verify version upgrade logic - assert new_version > old_version # Valid upgrade - assert not (same_version > old_version) # Invalid upgrade (same version) - - -class TestDependencyResolution: - """Test plugin dependency resolution.""" - - @pytest.fixture - def plugin_installer(self): - """Create a PluginInstaller instance for testing.""" - return PluginInstaller() - - def test_upload_bundle_with_dependencies(self, plugin_installer): - """Test uploading a plugin bundle and extracting dependencies.""" - # Arrange: Create mock bundle data and dependencies - bundle_data = b"mock-bundle-data" - mock_dependencies = [ - PluginBundleDependency( - type=PluginBundleDependencyType.Marketplace, - value=PluginBundleDependency.Marketplace(organization="org1", plugin="plugin1", version="1.0.0"), - ), - PluginBundleDependency( - type=PluginBundleDependencyType.Github, - value=PluginBundleDependency.Github( - repo_address="https://github.com/org/repo", - repo="org/repo", - release="v1.0.0", - packages="plugin.zip", - ), - ), - ] - - with patch.object( - plugin_installer, "_request_with_plugin_daemon_response", return_value=mock_dependencies - ) as mock_request: - # Act: Upload bundle - result = plugin_installer.upload_bundle("test-tenant", bundle_data, verify_signature=False) - - # Assert: Verify dependencies were extracted - assert len(result) == 2 - assert result[0].type == PluginBundleDependencyType.Marketplace - assert result[1].type == PluginBundleDependencyType.Github - mock_request.assert_called_once() - - def test_fetch_missing_dependencies(self, plugin_installer): - """Test fetching missing dependencies for plugins.""" - # Arrange: Mock missing dependencies response - mock_missing = [ - MissingPluginDependency(plugin_unique_identifier="dep1/1.0.0", current_identifier=None), - MissingPluginDependency(plugin_unique_identifier="dep2/2.0.0", current_identifier="dep2/1.0.0"), - ] - - with patch.object( - plugin_installer, "_request_with_plugin_daemon_response", return_value=mock_missing - ) as mock_request: - # Act: Fetch missing dependencies - result = plugin_installer.fetch_missing_dependencies("test-tenant", ["plugin1/1.0.0", "plugin2/2.0.0"]) - - # Assert: Verify missing dependencies were identified - assert len(result) == 2 - assert result[0].plugin_unique_identifier == "dep1/1.0.0" - assert result[1].current_identifier == "dep2/1.0.0" - mock_request.assert_called_once() - - def test_fetch_missing_dependencies_none_missing(self, plugin_installer): - """Test fetching missing dependencies when all are satisfied.""" - # Arrange: Mock empty missing dependencies - with patch.object(plugin_installer, "_request_with_plugin_daemon_response", return_value=[]): - # Act: Fetch missing dependencies - result = plugin_installer.fetch_missing_dependencies("test-tenant", ["plugin1/1.0.0"]) - - # Assert: Verify no missing dependencies - assert len(result) == 0 - - def test_fetch_plugin_installation_by_ids(self, plugin_installer): - """Test fetching plugin installations by their IDs.""" - # Arrange: Create mock plugin installations - mock_installations = [ - PluginInstallation( - id="install-1", - created_at=datetime.datetime.now(), - updated_at=datetime.datetime.now(), - tenant_id="test-tenant", - endpoints_setups=0, - endpoints_active=0, - runtime_type="remote", - source=PluginInstallationSource.Marketplace, - meta={}, - plugin_id="plugin-1", - plugin_unique_identifier="org/plugin1/1.0.0", - version="1.0.0", - checksum="abc123", - declaration=PluginDeclaration( - version="1.0.0", - author="author1", - name="plugin1", - description=I18nObject(en_US="Plugin 1", zh_Hans="插件1"), - icon="icon.png", - label=I18nObject(en_US="Plugin 1", zh_Hans="插件1"), - category=PluginCategory.Tool, - created_at=datetime.datetime.now(), - resource=PluginResourceRequirements(memory=512, permission=None), - plugins=PluginDeclaration.Plugins(), - meta=PluginDeclaration.Meta(version="1.0.0"), - ), - ) - ] - - with patch.object( - plugin_installer, "_request_with_plugin_daemon_response", return_value=mock_installations - ) as mock_request: - # Act: Fetch installations by IDs - result = plugin_installer.fetch_plugin_installation_by_ids("test-tenant", ["plugin-1", "plugin-2"]) - - # Assert: Verify installations were fetched - assert len(result) == 1 - assert result[0].plugin_id == "plugin-1" - mock_request.assert_called_once() - - def test_dependency_chain_resolution(self, plugin_installer): - """Test resolving a chain of dependencies.""" - # Arrange: Create a dependency chain scenario - # Plugin A depends on Plugin B, Plugin B depends on Plugin C - mock_missing = [ - MissingPluginDependency(plugin_unique_identifier="plugin-b/1.0.0", current_identifier=None), - MissingPluginDependency(plugin_unique_identifier="plugin-c/1.0.0", current_identifier=None), - ] - - with patch.object(plugin_installer, "_request_with_plugin_daemon_response", return_value=mock_missing): - # Act: Fetch missing dependencies for Plugin A - result = plugin_installer.fetch_missing_dependencies("test-tenant", ["plugin-a/1.0.0"]) - - # Assert: Verify all dependencies in the chain are identified - assert len(result) == 2 - identifiers = [dep.plugin_unique_identifier for dep in result] - assert "plugin-b/1.0.0" in identifiers - assert "plugin-c/1.0.0" in identifiers - - def test_check_tools_existence(self, plugin_installer): - """Test checking if plugin tools exist.""" - # Arrange: Create provider IDs to check using the correct format - provider_ids = [ - GenericProviderID("org1/plugin1/provider1"), - GenericProviderID("org2/plugin2/provider2"), - ] - - # Mock response indicating first exists, second doesn't - mock_response = [True, False] - - with patch.object( - plugin_installer, "_request_with_plugin_daemon_response", return_value=mock_response - ) as mock_request: - # Act: Check tools existence - result = plugin_installer.check_tools_existence("test-tenant", provider_ids) - - # Assert: Verify existence check results - assert len(result) == 2 - assert result[0] is True - assert result[1] is False - mock_request.assert_called_once() - - -class TestPluginTaskManagement: - """Test plugin installation task management.""" - - @pytest.fixture - def plugin_installer(self): - """Create a PluginInstaller instance for testing.""" - return PluginInstaller() - - def test_fetch_plugin_installation_tasks(self, plugin_installer): - """Test fetching multiple plugin installation tasks.""" - # Arrange: Create mock installation tasks - mock_tasks = [ - PluginInstallTask( - id="task-1", - created_at=datetime.datetime.now(), - updated_at=datetime.datetime.now(), - status=PluginInstallTaskStatus.Running, - total_plugins=2, - completed_plugins=1, - plugins=[], - ), - PluginInstallTask( - id="task-2", - created_at=datetime.datetime.now(), - updated_at=datetime.datetime.now(), - status=PluginInstallTaskStatus.Success, - total_plugins=1, - completed_plugins=1, - plugins=[], - ), - ] - - with patch.object( - plugin_installer, "_request_with_plugin_daemon_response", return_value=mock_tasks - ) as mock_request: - # Act: Fetch installation tasks - result = plugin_installer.fetch_plugin_installation_tasks("test-tenant", page=1, page_size=10) - - # Assert: Verify tasks were fetched - assert len(result) == 2 - assert result[0].status == PluginInstallTaskStatus.Running - assert result[1].status == PluginInstallTaskStatus.Success - mock_request.assert_called_once() - - def test_delete_plugin_installation_task(self, plugin_installer): - """Test deleting a specific plugin installation task.""" - # Arrange: Mock successful deletion - with patch.object(plugin_installer, "_request_with_plugin_daemon_response", return_value=True) as mock_request: - # Act: Delete installation task - result = plugin_installer.delete_plugin_installation_task("test-tenant", "task-123") - - # Assert: Verify deletion succeeded - assert result is True - mock_request.assert_called_once() - - def test_delete_all_plugin_installation_task_items(self, plugin_installer): - """Test deleting all plugin installation task items.""" - # Arrange: Mock successful deletion of all items - with patch.object(plugin_installer, "_request_with_plugin_daemon_response", return_value=True) as mock_request: - # Act: Delete all task items - result = plugin_installer.delete_all_plugin_installation_task_items("test-tenant") - - # Assert: Verify all items were deleted - assert result is True - mock_request.assert_called_once() - - def test_delete_plugin_installation_task_item(self, plugin_installer): - """Test deleting a specific item from an installation task.""" - # Arrange: Mock successful item deletion - with patch.object(plugin_installer, "_request_with_plugin_daemon_response", return_value=True) as mock_request: - # Act: Delete specific task item - result = plugin_installer.delete_plugin_installation_task_item( - "test-tenant", "task-123", "plugin-identifier" - ) - - # Assert: Verify item was deleted - assert result is True - mock_request.assert_called_once() - - -class TestErrorHandling: - """Test error handling in plugin manager.""" - - @pytest.fixture - def plugin_installer(self): - """Create a PluginInstaller instance for testing.""" - return PluginInstaller() - - def test_plugin_not_found_error(self, plugin_installer): - """Test handling of plugin not found error.""" - # Arrange: Mock plugin daemon not found error - with patch.object( - plugin_installer, - "_request_with_plugin_daemon_response", - side_effect=PluginDaemonNotFoundError("Plugin not found"), - ): - # Act & Assert: Verify error is raised - with pytest.raises(PluginDaemonNotFoundError): - plugin_installer.fetch_plugin_manifest("test-tenant", "non-existent/plugin/1.0.0") - - def test_plugin_bad_request_error(self, plugin_installer): - """Test handling of bad request error.""" - # Arrange: Mock bad request error - with patch.object( - plugin_installer, - "_request_with_plugin_daemon_response", - side_effect=PluginDaemonBadRequestError("Invalid request"), - ): - # Act & Assert: Verify error is raised - with pytest.raises(PluginDaemonBadRequestError): - plugin_installer.install_from_identifiers("test-tenant", [], PluginInstallationSource.Marketplace, []) - - def test_plugin_internal_server_error(self, plugin_installer): - """Test handling of internal server error.""" - # Arrange: Mock internal server error - with patch.object( - plugin_installer, - "_request_with_plugin_daemon_response", - side_effect=PluginDaemonInternalServerError("Internal error"), - ): - # Act & Assert: Verify error is raised - with pytest.raises(PluginDaemonInternalServerError): - plugin_installer.list_plugins("test-tenant") - - def test_http_error_handling(self, plugin_installer): - """Test handling of HTTP errors during requests.""" - # Arrange: Mock HTTP error - with patch.object(plugin_installer, "_request", side_effect=httpx.RequestError("Connection failed")): - # Act & Assert: Verify appropriate error handling - with pytest.raises(httpx.RequestError): - plugin_installer._request("GET", "test/path") - - -class TestPluginCategoryDetection: - """Test automatic plugin category detection.""" - - def test_category_defaults_to_extension_without_tool_provider(self): - """Test that plugins without tool providers default to Extension category.""" - # Arrange: Create declaration - category is auto-detected based on provider presence - # The model_validator in PluginDeclaration automatically sets category based on which provider is present - # Since we're not providing a tool provider entity, it defaults to Extension - # This test verifies that explicitly set categories are preserved - declaration = PluginDeclaration( - version="1.0.0", - author="test-author", - name="tool-plugin", - description=I18nObject(en_US="Tool plugin", zh_Hans="工具插件"), - icon="icon.png", - label=I18nObject(en_US="Tool Plugin", zh_Hans="工具插件"), - category=PluginCategory.Extension, # Will be Extension without a tool provider - created_at=datetime.datetime.now(), - resource=PluginResourceRequirements(memory=512, permission=None), - plugins=PluginDeclaration.Plugins(), - meta=PluginDeclaration.Meta(version="1.0.0"), - ) - - # Assert: Verify category defaults to Extension when no provider is specified - assert declaration.category == PluginCategory.Extension - - def test_category_defaults_to_extension_without_model_provider(self): - """Test that plugins without model providers default to Extension category.""" - # Arrange: Create declaration - without a model provider entity, defaults to Extension - # The category is auto-detected in the model_validator based on provider presence - declaration = PluginDeclaration( - version="1.0.0", - author="test-author", - name="model-plugin", - description=I18nObject(en_US="Model plugin", zh_Hans="模型插件"), - icon="icon.png", - label=I18nObject(en_US="Model Plugin", zh_Hans="模型插件"), - category=PluginCategory.Extension, # Will be Extension without a model provider - created_at=datetime.datetime.now(), - resource=PluginResourceRequirements(memory=1024, permission=None), - plugins=PluginDeclaration.Plugins(), - meta=PluginDeclaration.Meta(version="1.0.0"), - ) - - # Assert: Verify category defaults to Extension when no provider is specified - assert declaration.category == PluginCategory.Extension - - def test_extension_category_default(self): - """Test that plugins without specific providers default to Extension.""" - # Arrange: Create declaration without specific provider type - declaration = PluginDeclaration( - version="1.0.0", - author="test-author", - name="extension-plugin", - description=I18nObject(en_US="Extension plugin", zh_Hans="扩展插件"), - icon="icon.png", - label=I18nObject(en_US="Extension Plugin", zh_Hans="扩展插件"), - category=PluginCategory.Extension, - created_at=datetime.datetime.now(), - resource=PluginResourceRequirements(memory=512, permission=None), - plugins=PluginDeclaration.Plugins(), - meta=PluginDeclaration.Meta(version="1.0.0"), - ) - - # Assert: Verify category is Extension - assert declaration.category == PluginCategory.Extension - - -class TestPluginResourceRequirements: - """Test plugin resource requirements and permissions.""" - - def test_default_resource_requirements(self): - """ - Test that plugin resource requirements can be created with default values. - - Resource requirements define the memory and permissions needed for a plugin to run. - This test verifies that a basic resource requirement with only memory can be created. - """ - # Arrange & Act: Create resource requirements with only memory specified - resources = PluginResourceRequirements(memory=512, permission=None) - - # Assert: Verify memory is set correctly and permissions are None - assert resources.memory == 512 - assert resources.permission is None - - def test_resource_requirements_with_tool_permission(self): - """ - Test plugin resource requirements with tool permissions enabled. - - Tool permissions allow a plugin to provide tool functionality. - This test verifies that tool permissions can be properly configured. - """ - # Arrange & Act: Create resource requirements with tool permissions - resources = PluginResourceRequirements( - memory=1024, - permission=PluginResourceRequirements.Permission( - tool=PluginResourceRequirements.Permission.Tool(enabled=True) - ), - ) - - # Assert: Verify tool permissions are enabled - assert resources.memory == 1024 - assert resources.permission is not None - assert resources.permission.tool is not None - assert resources.permission.tool.enabled is True - - def test_resource_requirements_with_model_permissions(self): - """ - Test plugin resource requirements with model permissions. - - Model permissions allow a plugin to provide various AI model capabilities - including LLM, text embedding, rerank, TTS, speech-to-text, and moderation. - """ - # Arrange & Act: Create resource requirements with comprehensive model permissions - resources = PluginResourceRequirements( - memory=2048, - permission=PluginResourceRequirements.Permission( - model=PluginResourceRequirements.Permission.Model( - enabled=True, - llm=True, - text_embedding=True, - rerank=True, - tts=False, - speech2text=False, - moderation=True, - ) - ), - ) - - # Assert: Verify all model permissions are set correctly - assert resources.memory == 2048 - assert resources.permission.model.enabled is True - assert resources.permission.model.llm is True - assert resources.permission.model.text_embedding is True - assert resources.permission.model.rerank is True - assert resources.permission.model.tts is False - assert resources.permission.model.speech2text is False - assert resources.permission.model.moderation is True - - def test_resource_requirements_with_storage_permission(self): - """ - Test plugin resource requirements with storage permissions. - - Storage permissions allow a plugin to persist data with size limits. - The size must be between 1KB (1024 bytes) and 1GB (1073741824 bytes). - """ - # Arrange & Act: Create resource requirements with storage permissions - resources = PluginResourceRequirements( - memory=512, - permission=PluginResourceRequirements.Permission( - storage=PluginResourceRequirements.Permission.Storage(enabled=True, size=10485760) # 10MB - ), - ) - - # Assert: Verify storage permissions and size limits - assert resources.permission.storage.enabled is True - assert resources.permission.storage.size == 10485760 - - def test_resource_requirements_with_endpoint_permission(self): - """ - Test plugin resource requirements with endpoint permissions. - - Endpoint permissions allow a plugin to expose HTTP endpoints. - """ - # Arrange & Act: Create resource requirements with endpoint permissions - resources = PluginResourceRequirements( - memory=1024, - permission=PluginResourceRequirements.Permission( - endpoint=PluginResourceRequirements.Permission.Endpoint(enabled=True) - ), - ) - - # Assert: Verify endpoint permissions are enabled - assert resources.permission.endpoint.enabled is True - - def test_resource_requirements_with_node_permission(self): - """ - Test plugin resource requirements with node permissions. - - Node permissions allow a plugin to provide custom workflow nodes. - """ - # Arrange & Act: Create resource requirements with node permissions - resources = PluginResourceRequirements( - memory=768, - permission=PluginResourceRequirements.Permission( - node=PluginResourceRequirements.Permission.Node(enabled=True) - ), - ) - - # Assert: Verify node permissions are enabled - assert resources.permission.node.enabled is True - - -class TestPluginInstallationSources: - """Test different plugin installation sources.""" - - def test_marketplace_installation_source(self): - """ - Test plugin installation from marketplace source. - - Marketplace is the official plugin distribution channel where - verified and community plugins are available for installation. - """ - # Arrange & Act: Use marketplace as installation source - source = PluginInstallationSource.Marketplace - - # Assert: Verify source type - assert source == PluginInstallationSource.Marketplace - assert source.value == "marketplace" - - def test_github_installation_source(self): - """ - Test plugin installation from GitHub source. - - GitHub source allows installing plugins directly from GitHub repositories, - useful for development and testing unreleased versions. - """ - # Arrange & Act: Use GitHub as installation source - source = PluginInstallationSource.Github - - # Assert: Verify source type - assert source == PluginInstallationSource.Github - assert source.value == "github" - - def test_package_installation_source(self): - """ - Test plugin installation from package source. - - Package source allows installing plugins from local .difypkg files, - useful for private or custom plugins. - """ - # Arrange & Act: Use package as installation source - source = PluginInstallationSource.Package - - # Assert: Verify source type - assert source == PluginInstallationSource.Package - assert source.value == "package" - - def test_remote_installation_source(self): - """ - Test plugin installation from remote source. - - Remote source allows installing plugins from custom remote URLs. - """ - # Arrange & Act: Use remote as installation source - source = PluginInstallationSource.Remote - - # Assert: Verify source type - assert source == PluginInstallationSource.Remote - assert source.value == "remote" - - -class TestPluginBundleOperations: - """Test plugin bundle operations and dependency extraction.""" - - @pytest.fixture - def plugin_installer(self): - """Create a PluginInstaller instance for testing.""" - return PluginInstaller() - - def test_upload_bundle_with_marketplace_dependencies(self, plugin_installer): - """ - Test uploading a bundle with marketplace dependencies. - - Marketplace dependencies reference plugins available in the official marketplace - by organization, plugin name, and version. - """ - # Arrange: Create mock bundle with marketplace dependencies - bundle_data = b"mock-marketplace-bundle" - mock_dependencies = [ - PluginBundleDependency( - type=PluginBundleDependencyType.Marketplace, - value=PluginBundleDependency.Marketplace( - organization="langgenius", plugin="search-tool", version="1.2.0" - ), - ) - ] - - with patch.object(plugin_installer, "_request_with_plugin_daemon_response", return_value=mock_dependencies): - # Act: Upload bundle - result = plugin_installer.upload_bundle("test-tenant", bundle_data) - - # Assert: Verify marketplace dependency was extracted - assert len(result) == 1 - assert result[0].type == PluginBundleDependencyType.Marketplace - assert isinstance(result[0].value, PluginBundleDependency.Marketplace) - assert result[0].value.organization == "langgenius" - assert result[0].value.plugin == "search-tool" - - def test_upload_bundle_with_github_dependencies(self, plugin_installer): - """ - Test uploading a bundle with GitHub dependencies. - - GitHub dependencies reference plugins hosted on GitHub repositories - with specific releases and package files. - """ - # Arrange: Create mock bundle with GitHub dependencies - bundle_data = b"mock-github-bundle" - mock_dependencies = [ - PluginBundleDependency( - type=PluginBundleDependencyType.Github, - value=PluginBundleDependency.Github( - repo_address="https://github.com/example/plugin", - repo="example/plugin", - release="v2.0.0", - packages="plugin-v2.0.0.zip", - ), - ) - ] - - with patch.object(plugin_installer, "_request_with_plugin_daemon_response", return_value=mock_dependencies): - # Act: Upload bundle - result = plugin_installer.upload_bundle("test-tenant", bundle_data) - - # Assert: Verify GitHub dependency was extracted - assert len(result) == 1 - assert result[0].type == PluginBundleDependencyType.Github - assert isinstance(result[0].value, PluginBundleDependency.Github) - assert result[0].value.repo == "example/plugin" - assert result[0].value.release == "v2.0.0" - - def test_upload_bundle_with_package_dependencies(self, plugin_installer): - """ - Test uploading a bundle with package dependencies. - - Package dependencies include the full plugin manifest and unique identifier, - allowing for self-contained plugin bundles. - """ - # Arrange: Create mock bundle with package dependencies - bundle_data = b"mock-package-bundle" - mock_manifest = PluginDeclaration( - version="1.5.0", - author="bundle-author", - name="bundled-plugin", - description=I18nObject(en_US="Bundled plugin", zh_Hans="捆绑插件"), - icon="icon.png", - label=I18nObject(en_US="Bundled Plugin", zh_Hans="捆绑插件"), - category=PluginCategory.Extension, - created_at=datetime.datetime.now(), - resource=PluginResourceRequirements(memory=512, permission=None), - plugins=PluginDeclaration.Plugins(), - meta=PluginDeclaration.Meta(version="1.5.0"), - ) - - mock_dependencies = [ - PluginBundleDependency( - type=PluginBundleDependencyType.Package, - value=PluginBundleDependency.Package( - unique_identifier="org/bundled-plugin/1.5.0", manifest=mock_manifest - ), - ) - ] - - with patch.object(plugin_installer, "_request_with_plugin_daemon_response", return_value=mock_dependencies): - # Act: Upload bundle - result = plugin_installer.upload_bundle("test-tenant", bundle_data) - - # Assert: Verify package dependency was extracted with manifest - assert len(result) == 1 - assert result[0].type == PluginBundleDependencyType.Package - assert isinstance(result[0].value, PluginBundleDependency.Package) - assert result[0].value.unique_identifier == "org/bundled-plugin/1.5.0" - assert result[0].value.manifest.name == "bundled-plugin" - - def test_upload_bundle_with_mixed_dependencies(self, plugin_installer): - """ - Test uploading a bundle with multiple dependency types. - - Real-world plugin bundles often have dependencies from various sources: - marketplace plugins, GitHub repositories, and packaged plugins. - """ - # Arrange: Create mock bundle with mixed dependencies - bundle_data = b"mock-mixed-bundle" - mock_dependencies = [ - PluginBundleDependency( - type=PluginBundleDependencyType.Marketplace, - value=PluginBundleDependency.Marketplace(organization="org1", plugin="plugin1", version="1.0.0"), - ), - PluginBundleDependency( - type=PluginBundleDependencyType.Github, - value=PluginBundleDependency.Github( - repo_address="https://github.com/org2/plugin2", - repo="org2/plugin2", - release="v1.0.0", - packages="plugin2.zip", - ), - ), - ] - - with patch.object(plugin_installer, "_request_with_plugin_daemon_response", return_value=mock_dependencies): - # Act: Upload bundle - result = plugin_installer.upload_bundle("test-tenant", bundle_data, verify_signature=True) - - # Assert: Verify all dependency types were extracted - assert len(result) == 2 - assert result[0].type == PluginBundleDependencyType.Marketplace - assert result[1].type == PluginBundleDependencyType.Github - - -class TestPluginTaskStatusTransitions: - """Test plugin installation task status transitions and lifecycle.""" - - @pytest.fixture - def plugin_installer(self): - """Create a PluginInstaller instance for testing.""" - return PluginInstaller() - - def test_task_status_pending(self, plugin_installer): - """ - Test plugin installation task in pending status. - - Pending status indicates the task has been created but not yet started. - No plugins have been processed yet. - """ - # Arrange: Create mock task in pending status - mock_task = PluginInstallTask( - id="pending-task", - created_at=datetime.datetime.now(), - updated_at=datetime.datetime.now(), - status=PluginInstallTaskStatus.Pending, - total_plugins=3, - completed_plugins=0, # No plugins completed yet - plugins=[], - ) - - with patch.object(plugin_installer, "_request_with_plugin_daemon_response", return_value=mock_task): - # Act: Fetch task - result = plugin_installer.fetch_plugin_installation_task("test-tenant", "pending-task") - - # Assert: Verify pending status - assert result.status == PluginInstallTaskStatus.Pending - assert result.completed_plugins == 0 - assert result.total_plugins == 3 - - def test_task_status_running(self, plugin_installer): - """ - Test plugin installation task in running status. - - Running status indicates the task is actively installing plugins. - Some plugins may be completed while others are still in progress. - """ - # Arrange: Create mock task in running status - mock_task = PluginInstallTask( - id="running-task", - created_at=datetime.datetime.now(), - updated_at=datetime.datetime.now(), - status=PluginInstallTaskStatus.Running, - total_plugins=5, - completed_plugins=2, # 2 out of 5 completed - plugins=[], - ) - - with patch.object(plugin_installer, "_request_with_plugin_daemon_response", return_value=mock_task): - # Act: Fetch task - result = plugin_installer.fetch_plugin_installation_task("test-tenant", "running-task") - - # Assert: Verify running status and progress - assert result.status == PluginInstallTaskStatus.Running - assert result.completed_plugins == 2 - assert result.total_plugins == 5 - assert result.completed_plugins < result.total_plugins - - def test_task_status_success(self, plugin_installer): - """ - Test plugin installation task in success status. - - Success status indicates all plugins in the task have been - successfully installed without errors. - """ - # Arrange: Create mock task in success status - mock_task = PluginInstallTask( - id="success-task", - created_at=datetime.datetime.now(), - updated_at=datetime.datetime.now(), - status=PluginInstallTaskStatus.Success, - total_plugins=4, - completed_plugins=4, # All plugins completed - plugins=[], - ) - - with patch.object(plugin_installer, "_request_with_plugin_daemon_response", return_value=mock_task): - # Act: Fetch task - result = plugin_installer.fetch_plugin_installation_task("test-tenant", "success-task") - - # Assert: Verify success status and completion - assert result.status == PluginInstallTaskStatus.Success - assert result.completed_plugins == result.total_plugins - assert result.completed_plugins == 4 - - def test_task_status_failed(self, plugin_installer): - """ - Test plugin installation task in failed status. - - Failed status indicates the task encountered errors during installation. - Some plugins may have been installed before the failure occurred. - """ - # Arrange: Create mock task in failed status - mock_task = PluginInstallTask( - id="failed-task", - created_at=datetime.datetime.now(), - updated_at=datetime.datetime.now(), - status=PluginInstallTaskStatus.Failed, - total_plugins=3, - completed_plugins=1, # Only 1 completed before failure - plugins=[], - ) - - with patch.object(plugin_installer, "_request_with_plugin_daemon_response", return_value=mock_task): - # Act: Fetch task - result = plugin_installer.fetch_plugin_installation_task("test-tenant", "failed-task") - - # Assert: Verify failed status - assert result.status == PluginInstallTaskStatus.Failed - assert result.completed_plugins < result.total_plugins - - -class TestPluginI18nSupport: - """Test plugin internationalization (i18n) support.""" - - def test_plugin_with_multiple_languages(self): - """ - Test plugin declaration with multiple language support. - - Plugins should support multiple languages for descriptions and labels - to provide localized experiences for users worldwide. - """ - # Arrange & Act: Create plugin with English and Chinese support - declaration = PluginDeclaration( - version="1.0.0", - author="i18n-author", - name="multilang-plugin", - description=I18nObject( - en_US="A plugin with multilingual support", - zh_Hans="支持多语言的插件", - ja_JP="多言語対応のプラグイン", - ), - icon="icon.png", - label=I18nObject(en_US="Multilingual Plugin", zh_Hans="多语言插件", ja_JP="多言語プラグイン"), - category=PluginCategory.Extension, - created_at=datetime.datetime.now(), - resource=PluginResourceRequirements(memory=512, permission=None), - plugins=PluginDeclaration.Plugins(), - meta=PluginDeclaration.Meta(version="1.0.0"), - ) - - # Assert: Verify all language variants are preserved - assert declaration.description.en_US == "A plugin with multilingual support" - assert declaration.description.zh_Hans == "支持多语言的插件" - assert declaration.label.en_US == "Multilingual Plugin" - assert declaration.label.zh_Hans == "多语言插件" - - def test_plugin_readme_language_variants(self): - """ - Test fetching plugin README in different languages. - - Plugins can provide README files in multiple languages to help - users understand the plugin in their preferred language. - """ - # Arrange: Create plugin installer instance - plugin_installer = PluginInstaller() - - # Mock README responses for different languages - english_readme = PluginReadmeResponse( - content="# English README\n\nThis is the English version.", language="en_US" - ) - - chinese_readme = PluginReadmeResponse(content="# 中文说明\n\n这是中文版本。", language="zh_Hans") - - # Test English README - with patch.object(plugin_installer, "_request_with_plugin_daemon_response", return_value=english_readme): - # Act: Fetch English README - result_en = plugin_installer.fetch_plugin_readme("test-tenant", "plugin/1.0.0", "en_US") - - # Assert: Verify English content - assert "English README" in result_en - - # Test Chinese README - with patch.object(plugin_installer, "_request_with_plugin_daemon_response", return_value=chinese_readme): - # Act: Fetch Chinese README - result_zh = plugin_installer.fetch_plugin_readme("test-tenant", "plugin/1.0.0", "zh_Hans") - - # Assert: Verify Chinese content - assert "中文说明" in result_zh diff --git a/api/tests/unit_tests/core/plugin/test_plugin_runtime.py b/api/tests/unit_tests/core/plugin/test_plugin_runtime.py deleted file mode 100644 index 922aed4567f..00000000000 --- a/api/tests/unit_tests/core/plugin/test_plugin_runtime.py +++ /dev/null @@ -1,1812 +0,0 @@ -"""Comprehensive unit tests for Plugin Runtime functionality. - -This test module covers all aspects of plugin runtime including: -- Plugin execution through the plugin daemon -- Sandbox isolation via HTTP communication -- Resource limits (timeout, memory constraints) -- Error handling for various failure scenarios -- Plugin communication (request/response patterns, streaming) - -All tests use mocking to avoid external dependencies and ensure fast, reliable execution. -Tests follow the Arrange-Act-Assert pattern for clarity. -""" - -import json -from typing import Any -from unittest.mock import patch - -import httpx -import pytest -from pydantic import BaseModel - -from core.plugin.entities.plugin_daemon import ( - CredentialType, - PluginDaemonInnerError, -) -from core.plugin.impl.base import BasePluginClient, use_plugin_daemon_request_timeout -from core.plugin.impl.exc import ( - PluginDaemonBadRequestError, - PluginDaemonClientSideError, - PluginDaemonInternalServerError, - PluginDaemonNotFoundError, - PluginDaemonUnauthorizedError, - PluginInvokeError, - PluginNotFoundError, - PluginPermissionDeniedError, - PluginUniqueIdentifierError, -) -from core.plugin.impl.plugin import PluginInstaller -from core.plugin.impl.tool import PluginToolManager -from graphon.model_runtime.errors.invoke import ( - InvokeAuthorizationError, - InvokeBadRequestError, - InvokeConnectionError, - InvokeRateLimitError, - InvokeServerUnavailableError, -) -from graphon.model_runtime.errors.validate import CredentialsValidateFailedError - - -@pytest.fixture(autouse=True) -def _patch_shared_httpx_client(): - """Make BasePluginClient's module-level httpx client delegate to patched httpx.request/stream. - - After refactor, code uses core.plugin.impl.base._httpx_client directly. - Patch its request/stream to route through module-level httpx so existing mocks still apply. - """ - - def request[**P](*args: P.args, **kwargs: P.kwargs) -> httpx.Response: - response = httpx.request(*args, **kwargs) - method, url = kwargs["method"], kwargs["url"] - assert isinstance(method, str) - assert isinstance(url, str) - response.request = httpx.Request(method, url) - return response - - with ( - patch("core.plugin.impl.base._httpx_client.request", side_effect=request), - patch("core.plugin.impl.base._httpx_client.stream", side_effect=lambda **kw: httpx.stream(**kw)), - ): - yield - - -class TestPluginRuntimeExecution: - """Unit tests for plugin execution functionality. - - Tests cover: - - Successful plugin invocation - - Request preparation and headers - - Response parsing - - Streaming responses - """ - - @pytest.fixture - def plugin_client(self): - """Create a BasePluginClient instance for testing.""" - return BasePluginClient() - - @pytest.fixture - def mock_config(self, config_overrides): - """Mock plugin daemon configuration.""" - config_overrides( - PLUGIN_DAEMON_URL="http://127.0.0.1:5002", - PLUGIN_DAEMON_KEY="test-api-key", - ) - - def test_request_preparation(self, plugin_client, mock_config): - """Test that requests are properly prepared with correct headers and URL.""" - # Arrange - path = "plugin/test-tenant/management/list" - headers = {"Custom-Header": "value"} - data = {"key": "value"} - params = {"page": 1} - - # Act - url, prepared_headers, prepared_data, prepared_params, files = plugin_client._prepare_request( - path, headers, data, params, None - ) - - # Assert - assert url == "http://127.0.0.1:5002/plugin/test-tenant/management/list" - assert prepared_headers["X-Api-Key"] == "test-api-key" - assert prepared_headers["Custom-Header"] == "value" - assert prepared_headers["Accept-Encoding"] == "gzip, deflate, br" - assert prepared_data == data - assert prepared_params == params - - def test_request_with_json_content_type(self, plugin_client, mock_config): - """Test request preparation with JSON content type.""" - # Arrange - path = "plugin/test-tenant/management/install" - headers = {"Content-Type": "application/json"} - data = {"plugin_id": "test-plugin"} - - # Act - url, prepared_headers, prepared_data, prepared_params, files = plugin_client._prepare_request( - path, headers, data, None, None - ) - - # Assert - assert prepared_headers["Content-Type"] == "application/json" - assert prepared_data == json.dumps(data) - - def test_successful_request_execution(self, plugin_client, mock_config): - """Test successful HTTP request execution.""" - # Arrange - mock_response = httpx.Response(200, json={"result": "success"}) - - with patch("httpx.request", return_value=mock_response, autospec=True) as mock_request: - # Act - response = plugin_client._request("GET", "plugin/test-tenant/management/list") - - # Assert - assert response.status_code == 200 - mock_request.assert_called_once() - call_kwargs = mock_request.call_args[1] - assert call_kwargs["method"] == "GET" - assert "http://127.0.0.1:5002/plugin/test-tenant/management/list" in call_kwargs["url"] - assert call_kwargs["headers"]["X-Api-Key"] == "test-api-key" - - def test_request_with_timeout_configuration(self, plugin_client, mock_config): - """Test that timeout configuration is properly applied.""" - # Arrange - mock_response = httpx.Response(200) - - with patch("httpx.request", return_value=mock_response, autospec=True) as mock_request: - # Act - plugin_client._request("GET", "plugin/test-tenant/test") - - # Assert - call_kwargs = mock_request.call_args[1] - assert "timeout" in call_kwargs - - def test_request_timeout_can_be_scoped_to_current_context(self, plugin_client, mock_config): - mock_response = httpx.Response(200) - - with ( - patch("core.plugin.impl.base.plugin_daemon_request_timeout", httpx.Timeout(600.0)), - patch("httpx.request", return_value=mock_response, autospec=True) as mock_request, - ): - with use_plugin_daemon_request_timeout(30.0): - plugin_client._request("GET", "plugin/test-tenant/test") - plugin_client._request("GET", "plugin/test-tenant/test") - - scoped_timeout = mock_request.call_args_list[0].kwargs["timeout"] - default_timeout = mock_request.call_args_list[1].kwargs["timeout"] - assert isinstance(scoped_timeout, httpx.Timeout) - assert isinstance(default_timeout, httpx.Timeout) - assert scoped_timeout.read == 30.0 - assert default_timeout.read == 600.0 - - @pytest.mark.parametrize("timeout_seconds", [0.0, -1.0]) - def test_request_timeout_override_rejects_non_positive_values(self, timeout_seconds): - with pytest.raises(ValueError, match="greater than zero"): - with use_plugin_daemon_request_timeout(timeout_seconds): - pass - - def test_request_connection_error(self, plugin_client, mock_config): - """Test handling of connection errors during request.""" - # Arrange - with patch("httpx.request", side_effect=httpx.RequestError("Connection failed"), autospec=True): - # Act & Assert - with pytest.raises(PluginDaemonInnerError) as exc_info: - plugin_client._request("GET", "plugin/test-tenant/test") - assert exc_info.value.code == -500 - assert "Request to Plugin Daemon Service failed" in exc_info.value.message - - -class TestPluginRuntimeSandboxIsolation: - """Unit tests for plugin sandbox isolation. - - Tests cover: - - Isolated execution environment via HTTP - - API key authentication - - Request/response boundaries - - Plugin daemon communication protocol - """ - - @pytest.fixture - def plugin_client(self): - """Create a BasePluginClient instance for testing.""" - return BasePluginClient() - - @pytest.fixture - def mock_config(self, config_overrides): - """Mock plugin daemon configuration.""" - config_overrides( - PLUGIN_DAEMON_URL="http://127.0.0.1:5002", - PLUGIN_DAEMON_KEY="secure-api-key", - ) - - def test_api_key_authentication(self, plugin_client, mock_config): - """Test that all requests include API key for authentication.""" - # Arrange - mock_response = httpx.Response(200, json={"code": 0, "message": "", "data": True}) - - with patch("httpx.request", return_value=mock_response, autospec=True) as mock_request: - # Act - plugin_client._request("GET", "plugin/test-tenant/test") - - # Assert - call_kwargs = mock_request.call_args[1] - assert call_kwargs["headers"]["X-Api-Key"] == "secure-api-key" - - def test_isolated_plugin_execution_via_http(self, plugin_client, mock_config): - """Test that plugin execution is isolated via HTTP communication.""" - - # Arrange - class TestResponse(BaseModel): - result: str - - mock_response = httpx.Response(200, json={"code": 0, "message": "", "data": {"result": "isolated_execution"}}) - - with patch("httpx.request", return_value=mock_response, autospec=True): - # Act - result = plugin_client._request_with_plugin_daemon_response( - "POST", "plugin/test-tenant/dispatch/tool/invoke", TestResponse, data={"tool": "test"} - ) - - # Assert - assert result.result == "isolated_execution" - - def test_plugin_daemon_unauthorized_error(self, plugin_client, mock_config): - """Test handling of unauthorized access to plugin daemon.""" - # Arrange - error_message = json.dumps({"error_type": "PluginDaemonUnauthorizedError", "message": "Unauthorized access"}) - mock_response = httpx.Response(200, json={"code": -1, "message": error_message, "data": None}) - - with patch("httpx.request", return_value=mock_response, autospec=True): - # Act & Assert - with pytest.raises(PluginDaemonUnauthorizedError) as exc_info: - plugin_client._request_with_plugin_daemon_response("GET", "plugin/test-tenant/test", bool) - assert "Unauthorized access" in exc_info.value.description - - def test_plugin_permission_denied(self, plugin_client, mock_config): - """Test handling of permission denied errors.""" - # Arrange - error_message = json.dumps( - {"error_type": "PluginPermissionDeniedError", "message": "Permission denied for this operation"} - ) - mock_response = httpx.Response(200, json={"code": -1, "message": error_message, "data": None}) - - with patch("httpx.request", return_value=mock_response, autospec=True): - # Act & Assert - with pytest.raises(PluginPermissionDeniedError) as exc_info: - plugin_client._request_with_plugin_daemon_response("POST", "plugin/test-tenant/test", bool) - assert "Permission denied" in exc_info.value.description - - -class TestPluginRuntimeResourceLimits: - """Unit tests for plugin resource limits. - - Tests cover: - - Timeout enforcement - - Memory constraints - - Resource limit violations - - Graceful degradation - """ - - @pytest.fixture - def plugin_client(self): - """Create a BasePluginClient instance for testing.""" - return BasePluginClient() - - @pytest.fixture - def mock_config(self, config_overrides): - """Mock plugin daemon configuration with timeout.""" - config_overrides( - PLUGIN_DAEMON_URL="http://127.0.0.1:5002", - PLUGIN_DAEMON_KEY="test-key", - ) - with patch("core.plugin.impl.base.plugin_daemon_request_timeout", httpx.Timeout(30.0)): - yield - - def test_timeout_configuration_applied(self, plugin_client, mock_config): - """Test that timeout configuration is properly applied to requests.""" - # Arrange - mock_response = httpx.Response(200) - - with patch("httpx.request", return_value=mock_response, autospec=True) as mock_request: - # Act - plugin_client._request("GET", "plugin/test-tenant/test") - - # Assert - call_kwargs = mock_request.call_args[1] - assert call_kwargs["timeout"] is not None - - def test_timeout_error_handling(self, plugin_client, mock_config): - """Test handling of timeout errors.""" - # Arrange - with patch("httpx.request", side_effect=httpx.TimeoutException("Request timeout"), autospec=True): - # Act & Assert - with pytest.raises(PluginDaemonInnerError) as exc_info: - plugin_client._request("GET", "plugin/test-tenant/test") - assert exc_info.value.code == -500 - - def test_streaming_request_timeout(self, plugin_client, mock_config): - """Test timeout handling for streaming requests.""" - # Arrange - with patch("httpx.stream", side_effect=httpx.TimeoutException("Stream timeout"), autospec=True): - # Act & Assert - with pytest.raises(PluginDaemonInnerError) as exc_info: - list(plugin_client._stream_request("POST", "plugin/test-tenant/stream")) - assert exc_info.value.code == -500 - - def test_resource_limit_error_from_daemon(self, plugin_client, mock_config): - """Test handling of resource limit errors from plugin daemon.""" - # Arrange - error_message = json.dumps( - {"error_type": "PluginDaemonInternalServerError", "message": "Resource limit exceeded"} - ) - mock_response = httpx.Response(200, json={"code": -1, "message": error_message, "data": None}) - - with patch("httpx.request", return_value=mock_response, autospec=True): - # Act & Assert - with pytest.raises(PluginDaemonInternalServerError) as exc_info: - plugin_client._request_with_plugin_daemon_response("POST", "plugin/test-tenant/test", bool) - assert "Resource limit exceeded" in exc_info.value.description - - -class TestPluginRuntimeErrorHandling: - """Unit tests for plugin runtime error handling. - - Tests cover: - - Various error types (invoke, validation, connection) - - Error propagation and transformation - - User-friendly error messages - - Error recovery mechanisms - """ - - @pytest.fixture - def plugin_client(self): - """Create a BasePluginClient instance for testing.""" - return BasePluginClient() - - @pytest.fixture - def mock_config(self, config_overrides): - """Mock plugin daemon configuration.""" - config_overrides( - PLUGIN_DAEMON_URL="http://127.0.0.1:5002", - PLUGIN_DAEMON_KEY="test-key", - ) - - def test_plugin_invoke_rate_limit_error(self, plugin_client, mock_config): - """Test handling of rate limit errors during plugin invocation.""" - # Arrange - invoke_error = { - "error_type": "InvokeRateLimitError", - "message": "Rate limit exceeded", - "args": {"description": "Rate limit exceeded"}, - } - error_message = json.dumps({"error_type": "PluginInvokeError", "message": json.dumps(invoke_error)}) - mock_response = httpx.Response(200, json={"code": -1, "message": error_message, "data": None}) - - with patch("httpx.request", return_value=mock_response, autospec=True): - # Act & Assert - with pytest.raises(InvokeRateLimitError) as exc_info: - plugin_client._request_with_plugin_daemon_response("POST", "plugin/test-tenant/invoke", bool) - assert "Rate limit exceeded" in exc_info.value.description - - def test_plugin_invoke_authorization_error(self, plugin_client, mock_config): - """Test handling of authorization errors during plugin invocation.""" - # Arrange - invoke_error = { - "error_type": "InvokeAuthorizationError", - "message": "Invalid credentials", - "args": {"description": "Invalid credentials"}, - } - error_message = json.dumps({"error_type": "PluginInvokeError", "message": json.dumps(invoke_error)}) - mock_response = httpx.Response(200, json={"code": -1, "message": error_message, "data": None}) - - with patch("httpx.request", return_value=mock_response, autospec=True): - # Act & Assert - with pytest.raises(InvokeAuthorizationError) as exc_info: - plugin_client._request_with_plugin_daemon_response("POST", "plugin/test-tenant/invoke", bool) - assert "Invalid credentials" in exc_info.value.description - - def test_plugin_invoke_bad_request_error(self, plugin_client, mock_config): - """Test handling of bad request errors during plugin invocation.""" - # Arrange - invoke_error = { - "error_type": "InvokeBadRequestError", - "message": "Invalid parameters", - "args": {"description": "Invalid parameters"}, - } - error_message = json.dumps({"error_type": "PluginInvokeError", "message": json.dumps(invoke_error)}) - mock_response = httpx.Response(200, json={"code": -1, "message": error_message, "data": None}) - - with patch("httpx.request", return_value=mock_response, autospec=True): - # Act & Assert - with pytest.raises(InvokeBadRequestError) as exc_info: - plugin_client._request_with_plugin_daemon_response("POST", "plugin/test-tenant/invoke", bool) - assert "Invalid parameters" in exc_info.value.description - - def test_plugin_invoke_connection_error(self, plugin_client, mock_config): - """Test handling of connection errors during plugin invocation.""" - # Arrange - invoke_error = { - "error_type": "InvokeConnectionError", - "message": "Connection to external service failed", - "args": {"description": "Connection to external service failed"}, - } - error_message = json.dumps({"error_type": "PluginInvokeError", "message": json.dumps(invoke_error)}) - mock_response = httpx.Response(200, json={"code": -1, "message": error_message, "data": None}) - - with patch("httpx.request", return_value=mock_response, autospec=True): - # Act & Assert - with pytest.raises(InvokeConnectionError) as exc_info: - plugin_client._request_with_plugin_daemon_response("POST", "plugin/test-tenant/invoke", bool) - assert "Connection to external service failed" in exc_info.value.description - - def test_plugin_invoke_server_unavailable_error(self, plugin_client, mock_config): - """Test handling of server unavailable errors during plugin invocation.""" - # Arrange - invoke_error = { - "error_type": "InvokeServerUnavailableError", - "message": "Service temporarily unavailable", - "args": {"description": "Service temporarily unavailable"}, - } - error_message = json.dumps({"error_type": "PluginInvokeError", "message": json.dumps(invoke_error)}) - mock_response = httpx.Response(200, json={"code": -1, "message": error_message, "data": None}) - - with patch("httpx.request", return_value=mock_response, autospec=True): - # Act & Assert - with pytest.raises(InvokeServerUnavailableError) as exc_info: - plugin_client._request_with_plugin_daemon_response("POST", "plugin/test-tenant/invoke", bool) - assert "Service temporarily unavailable" in exc_info.value.description - - def test_credentials_validation_error(self, plugin_client, mock_config): - """Test handling of credential validation errors.""" - # Arrange - invoke_error = { - "error_type": "CredentialsValidateFailedError", - "message": "Invalid API key format", - } - error_message = json.dumps({"error_type": "PluginInvokeError", "message": json.dumps(invoke_error)}) - mock_response = httpx.Response(200, json={"code": -1, "message": error_message, "data": None}) - - with patch("httpx.request", return_value=mock_response, autospec=True): - # Act & Assert - with pytest.raises(CredentialsValidateFailedError) as exc_info: - plugin_client._request_with_plugin_daemon_response("POST", "plugin/test-tenant/validate", bool) - assert "Invalid API key format" in str(exc_info.value) - - def test_plugin_not_found_error(self, plugin_client, mock_config): - """Test handling of plugin not found errors.""" - # Arrange - error_message = json.dumps( - {"error_type": "PluginNotFoundError", "message": "Plugin with ID 'test-plugin' not found"} - ) - mock_response = httpx.Response(200, json={"code": -1, "message": error_message, "data": None}) - - with patch("httpx.request", return_value=mock_response, autospec=True): - # Act & Assert - with pytest.raises(PluginNotFoundError) as exc_info: - plugin_client._request_with_plugin_daemon_response("GET", "plugin/test-tenant/get", bool) - assert "Plugin with ID 'test-plugin' not found" in exc_info.value.description - - def test_plugin_unique_identifier_error(self, plugin_client, mock_config): - """Test handling of unique identifier errors.""" - # Arrange - error_message = json.dumps( - {"error_type": "PluginUniqueIdentifierError", "message": "Invalid plugin identifier format"} - ) - mock_response = httpx.Response(200, json={"code": -1, "message": error_message, "data": None}) - - with patch("httpx.request", return_value=mock_response, autospec=True): - # Act & Assert - with pytest.raises(PluginUniqueIdentifierError) as exc_info: - plugin_client._request_with_plugin_daemon_response("POST", "plugin/test-tenant/install", bool) - assert "Invalid plugin identifier format" in exc_info.value.description - - def test_daemon_bad_request_error(self, plugin_client, mock_config): - """Test handling of daemon bad request errors.""" - # Arrange - error_message = json.dumps( - {"error_type": "PluginDaemonBadRequestError", "message": "Missing required parameter"} - ) - mock_response = httpx.Response(200, json={"code": -1, "message": error_message, "data": None}) - - with patch("httpx.request", return_value=mock_response, autospec=True): - # Act & Assert - with pytest.raises(PluginDaemonBadRequestError) as exc_info: - plugin_client._request_with_plugin_daemon_response("POST", "plugin/test-tenant/test", bool) - assert "Missing required parameter" in exc_info.value.description - - def test_daemon_not_found_error(self, plugin_client, mock_config): - """Test handling of daemon not found errors.""" - # Arrange - error_message = json.dumps({"error_type": "PluginDaemonNotFoundError", "message": "Resource not found"}) - mock_response = httpx.Response(200, json={"code": -1, "message": error_message, "data": None}) - - with patch("httpx.request", return_value=mock_response, autospec=True): - # Act & Assert - with pytest.raises(PluginDaemonNotFoundError) as exc_info: - plugin_client._request_with_plugin_daemon_response("GET", "plugin/test-tenant/resource", bool) - assert "Resource not found" in exc_info.value.description - - def test_generic_plugin_invoke_error(self, plugin_client, mock_config): - """Test handling of generic plugin invoke errors.""" - # Arrange - # Create a proper nested JSON structure for PluginInvokeError - invoke_error_message = json.dumps( - {"error_type": "UnknownInvokeError", "message": "Generic plugin execution error"} - ) - error_message = json.dumps({"error_type": "PluginInvokeError", "message": invoke_error_message}) - mock_response = httpx.Response(200, json={"code": -1, "message": error_message, "data": None}) - - with patch("httpx.request", return_value=mock_response, autospec=True): - # Act & Assert - with pytest.raises(PluginInvokeError) as exc_info: - plugin_client._request_with_plugin_daemon_response("POST", "plugin/test-tenant/invoke", bool) - assert exc_info.value.description is not None - - def test_unknown_error_type(self, plugin_client, mock_config): - """Test handling of unknown error types.""" - # Arrange - error_message = json.dumps({"error_type": "UnknownErrorType", "message": "Unknown error occurred"}) - mock_response = httpx.Response(200, json={"code": -1, "message": error_message, "data": None}) - - with patch("httpx.request", return_value=mock_response, autospec=True): - # Act & Assert - with pytest.raises( - Exception, - match=("got unknown error from plugin daemon: UnknownErrorType, message: Unknown error occurred"), - ) as exc_info: - plugin_client._request_with_plugin_daemon_response("POST", "plugin/test-tenant/test", bool) - assert "got unknown error from plugin daemon" in str(exc_info.value) - - def test_http_status_error_handling(self, plugin_client, mock_config): - """Test handling of HTTP status errors.""" - # Arrange - mock_response = httpx.Response( - 500, request=httpx.Request("GET", "http://127.0.0.1:5002/plugin/test-tenant/test") - ) - - with patch("httpx.request", return_value=mock_response, autospec=True): - # Act & Assert - with pytest.raises(PluginDaemonInternalServerError): - plugin_client._request_with_plugin_daemon_response("GET", "plugin/test-tenant/test", bool) - - def test_empty_data_response_error(self, plugin_client, mock_config): - """Test handling of empty data in successful response.""" - # Arrange - mock_response = httpx.Response(200, json={"code": 0, "message": "", "data": None}) - - with patch("httpx.request", return_value=mock_response, autospec=True): - # Act & Assert - with pytest.raises(ValueError, match="got empty data from plugin daemon") as exc_info: - plugin_client._request_with_plugin_daemon_response("GET", "plugin/test-tenant/test", bool) - assert "got empty data from plugin daemon" in str(exc_info.value) - - -class TestPluginRuntimeCommunication: - """Unit tests for plugin communication patterns. - - Tests cover: - - Request/response communication - - Streaming responses - - Data serialization/deserialization - - Message formatting - """ - - @pytest.fixture - def plugin_client(self): - """Create a BasePluginClient instance for testing.""" - return BasePluginClient() - - @pytest.fixture - def mock_config(self, config_overrides): - """Mock plugin daemon configuration.""" - config_overrides( - PLUGIN_DAEMON_URL="http://127.0.0.1:5002", - PLUGIN_DAEMON_KEY="test-key", - ) - - def test_request_response_communication(self, plugin_client, mock_config): - """Test basic request/response communication pattern.""" - - # Arrange - class TestModel(BaseModel): - value: str - count: int - - mock_response = httpx.Response(200, json={"code": 0, "message": "", "data": {"value": "test", "count": 42}}) - - with patch("httpx.request", return_value=mock_response, autospec=True): - # Act - result = plugin_client._request_with_plugin_daemon_response( - "POST", "plugin/test-tenant/test", TestModel, data={"input": "data"} - ) - - # Assert - assert isinstance(result, TestModel) - assert result.value == "test" - assert result.count == 42 - - def test_streaming_response_communication(self, plugin_client, mock_config): - """Test streaming response communication pattern.""" - - # Arrange - class StreamModel(BaseModel): - chunk: str - - stream_data = [ - 'data: {"code": 0, "message": "", "data": {"chunk": "first"}}', - 'data: {"code": 0, "message": "", "data": {"chunk": "second"}}', - 'data: {"code": 0, "message": "", "data": {"chunk": "third"}}', - ] - - mock_response = httpx.Response(200, text="\n".join(stream_data)) - - with patch("httpx.stream", autospec=True) as mock_stream: - mock_stream.return_value.__enter__.return_value = mock_response - - # Act - results = list( - plugin_client._request_with_plugin_daemon_response_stream( - "POST", "plugin/test-tenant/stream", StreamModel - ) - ) - - # Assert - assert len(results) == 3 - assert all(isinstance(r, StreamModel) for r in results) - assert results[0].chunk == "first" - assert results[1].chunk == "second" - assert results[2].chunk == "third" - - def test_streaming_with_error_in_stream(self, plugin_client, mock_config): - """Test error handling in streaming responses.""" - # Arrange - # Create proper error structure for -500 code - error_obj = json.dumps({"error_type": "PluginDaemonInnerError", "message": "Stream error occurred"}) - stream_data = [ - 'data: {"code": 0, "message": "", "data": {"chunk": "first"}}', - f'data: {{"code": -500, "message": {json.dumps(error_obj)}, "data": null}}', - ] - - mock_response = httpx.Response(200, text="\n".join(stream_data)) - - with patch("httpx.stream", autospec=True) as mock_stream: - mock_stream.return_value.__enter__.return_value = mock_response - - # Act - class StreamModel(BaseModel): - chunk: str - - results = plugin_client._request_with_plugin_daemon_response_stream( - "POST", "plugin/test-tenant/stream", StreamModel - ) - - # Assert - first_result = next(results) - assert first_result.chunk == "first" - - with pytest.raises(PluginDaemonInnerError) as exc_info: - next(results) - assert exc_info.value.code == -500 - - def test_streaming_connection_error(self, plugin_client, mock_config): - """Test connection error during streaming.""" - # Arrange - with patch("httpx.stream", side_effect=httpx.RequestError("Stream connection failed"), autospec=True): - # Act & Assert - with pytest.raises(PluginDaemonInnerError) as exc_info: - list(plugin_client._stream_request("POST", "plugin/test-tenant/stream")) - assert exc_info.value.code == -500 - - def test_request_with_model_parsing(self, plugin_client, mock_config): - """Test request with direct model parsing (without daemon response wrapper).""" - - # Arrange - class DirectModel(BaseModel): - status: str - data: dict[str, Any] - - mock_response = httpx.Response(200, json={"status": "success", "data": {"key": "value"}}) - - with patch("httpx.request", return_value=mock_response, autospec=True): - # Act - result = plugin_client._request_with_model("GET", "plugin/test-tenant/direct", DirectModel) - - # Assert - assert isinstance(result, DirectModel) - assert result.status == "success" - assert result.data == {"key": "value"} - - def test_streaming_with_model_parsing(self, plugin_client, mock_config): - """Test streaming with direct model parsing.""" - - # Arrange - class StreamItem(BaseModel): - id: int - text: str - - stream_data = [ - '{"id": 1, "text": "first"}', - '{"id": 2, "text": "second"}', - ] - - mock_response = httpx.Response(200, text="\n".join(stream_data)) - - with patch("httpx.stream", autospec=True) as mock_stream: - mock_stream.return_value.__enter__.return_value = mock_response - - # Act - results = list(plugin_client._stream_request_with_model("POST", "plugin/test-tenant/stream", StreamItem)) - - # Assert - assert len(results) == 2 - assert results[0].id == 1 - assert results[0].text == "first" - assert results[1].id == 2 - assert results[1].text == "second" - - def test_streaming_skips_empty_lines(self, plugin_client, mock_config): - """Test that streaming properly skips empty lines.""" - - # Arrange - class StreamModel(BaseModel): - value: str - - stream_data = [ - "", - '{"code": 0, "message": "", "data": {"value": "first"}}', - "", - "", - '{"code": 0, "message": "", "data": {"value": "second"}}', - "", - ] - - mock_response = httpx.Response(200, text="\n".join(stream_data)) - - with patch("httpx.stream", autospec=True) as mock_stream: - mock_stream.return_value.__enter__.return_value = mock_response - - # Act - results = list( - plugin_client._request_with_plugin_daemon_response_stream( - "POST", "plugin/test-tenant/stream", StreamModel - ) - ) - - # Assert - assert len(results) == 2 - assert results[0].value == "first" - assert results[1].value == "second" - - -class TestPluginToolManagerIntegration: - """Integration tests for PluginToolManager. - - Tests cover: - - Tool invocation - - Credential validation - - Runtime parameter retrieval - - Tool provider management - """ - - @pytest.fixture - def tool_manager(self): - """Create a PluginToolManager instance for testing.""" - return PluginToolManager() - - @pytest.fixture - def mock_config(self, config_overrides): - """Mock plugin daemon configuration.""" - config_overrides( - PLUGIN_DAEMON_URL="http://127.0.0.1:5002", - PLUGIN_DAEMON_KEY="test-key", - ) - - def test_tool_invocation_success(self, tool_manager, mock_config): - """Test successful tool invocation.""" - # Arrange - stream_data = [ - 'data: {"code": 0, "message": "", "data": {"type": "text", "message": {"text": "Result"}}}', - ] - - mock_response = httpx.Response(200, text="\n".join(stream_data)) - - with patch("httpx.stream", autospec=True) as mock_stream: - mock_stream.return_value.__enter__.return_value = mock_response - - # Act - results = list( - tool_manager.invoke( - tenant_id="test-tenant", - user_id="test-user", - tool_provider="langgenius/test-plugin/test-provider", - tool_name="test-tool", - credentials={"api_key": "test-key"}, - credential_type=CredentialType.API_KEY, - tool_parameters={"param1": "value1"}, - ) - ) - - # Assert - assert len(results) > 0 - assert results[0].type == "text" - - def test_validate_provider_credentials_success(self, tool_manager, mock_config): - """Test successful provider credential validation.""" - # Arrange - stream_data = [ - 'data: {"code": 0, "message": "", "data": {"result": true}}', - ] - - mock_response = httpx.Response(200, text="\n".join(stream_data)) - - with patch("httpx.stream", autospec=True) as mock_stream: - mock_stream.return_value.__enter__.return_value = mock_response - - # Act - result = tool_manager.validate_provider_credentials( - tenant_id="test-tenant", - user_id="test-user", - provider="langgenius/test-plugin/test-provider", - credentials={"api_key": "valid-key"}, - ) - - # Assert - assert result is True - - def test_validate_provider_credentials_failure(self, tool_manager, mock_config): - """Test failed provider credential validation.""" - # Arrange - stream_data = [ - 'data: {"code": 0, "message": "", "data": {"result": false}}', - ] - - mock_response = httpx.Response(200, text="\n".join(stream_data)) - - with patch("httpx.stream", autospec=True) as mock_stream: - mock_stream.return_value.__enter__.return_value = mock_response - - # Act - result = tool_manager.validate_provider_credentials( - tenant_id="test-tenant", - user_id="test-user", - provider="langgenius/test-plugin/test-provider", - credentials={"api_key": "invalid-key"}, - ) - - # Assert - assert result is False - - def test_validate_datasource_credentials_success(self, tool_manager, mock_config): - """Test successful datasource credential validation.""" - # Arrange - stream_data = [ - 'data: {"code": 0, "message": "", "data": {"result": true}}', - ] - - mock_response = httpx.Response(200, text="\n".join(stream_data)) - - with patch("httpx.stream", autospec=True) as mock_stream: - mock_stream.return_value.__enter__.return_value = mock_response - - # Act - result = tool_manager.validate_datasource_credentials( - tenant_id="test-tenant", - user_id="test-user", - provider="langgenius/test-plugin/test-datasource", - credentials={"connection_string": "valid"}, - ) - - # Assert - assert result is True - - -class TestPluginInstallerIntegration: - """Integration tests for PluginInstaller. - - Tests cover: - - Plugin installation - - Plugin listing - - Plugin uninstallation - - Package upload - """ - - @pytest.fixture - def installer(self): - """Create a PluginInstaller instance for testing.""" - return PluginInstaller() - - @pytest.fixture - def mock_config(self, config_overrides): - """Mock plugin daemon configuration.""" - config_overrides( - PLUGIN_DAEMON_URL="http://127.0.0.1:5002", - PLUGIN_DAEMON_KEY="test-key", - ) - - def test_list_plugins_success(self, installer, mock_config): - """Test successful plugin listing.""" - # Arrange - mock_response = httpx.Response( - 200, - json={ - "code": 0, - "message": "", - "data": { - "list": [], - "total": 0, - }, - }, - ) - - with patch("httpx.request", return_value=mock_response, autospec=True): - # Act - result = installer.list_plugins("test-tenant") - - # Assert - assert isinstance(result, list) - - def test_uninstall_plugin_success(self, installer, mock_config): - """Test successful plugin uninstallation.""" - # Arrange - mock_response = httpx.Response(200, json={"code": 0, "message": "", "data": True}) - - with patch("httpx.request", return_value=mock_response, autospec=True): - # Act - result = installer.uninstall("test-tenant", "plugin-installation-id") - - # Assert - assert result is True - - def test_fetch_plugin_by_identifier_success(self, installer, mock_config): - """Test successful plugin fetch by identifier.""" - # Arrange - mock_response = httpx.Response(200, json={"code": 0, "message": "", "data": True}) - - with patch("httpx.request", return_value=mock_response, autospec=True): - # Act - result = installer.fetch_plugin_by_identifier("test-tenant", "plugin-identifier") - - # Assert - assert result is True - - -class TestPluginRuntimeEdgeCases: - """Tests for edge cases and corner scenarios in plugin runtime. - - Tests cover: - - Malformed responses - - Unexpected data types - - Concurrent requests - - Large payloads - """ - - @pytest.fixture - def plugin_client(self): - """Create a BasePluginClient instance for testing.""" - return BasePluginClient() - - @pytest.fixture - def mock_config(self, config_overrides): - """Mock plugin daemon configuration.""" - config_overrides( - PLUGIN_DAEMON_URL="http://127.0.0.1:5002", - PLUGIN_DAEMON_KEY="test-key", - ) - - def test_malformed_json_response(self, plugin_client, mock_config): - """Test handling of malformed JSON responses.""" - # Arrange - mock_response = httpx.Response(200, content=b"invalid JSON") - - with patch("httpx.request", return_value=mock_response, autospec=True): - # Act & Assert - with pytest.raises( - ValueError, - match=( - "Failed to parse response from plugin daemon to " - "PluginDaemonBasicResponse \\[bool\\], url: plugin/test\\-tenant/test" - ), - ): - plugin_client._request_with_plugin_daemon_response("GET", "plugin/test-tenant/test", bool) - - def test_invalid_response_structure(self, plugin_client, mock_config): - """Test handling of invalid response structure.""" - # Arrange - # Missing required fields in response - mock_response = httpx.Response(200, json={"invalid": "structure"}) - - with patch("httpx.request", return_value=mock_response, autospec=True): - # Act & Assert - with pytest.raises( - ValueError, - match=( - "Failed to parse response from plugin daemon to " - "PluginDaemonBasicResponse \\[bool\\], url: plugin/test\\-tenant/test" - ), - ): - plugin_client._request_with_plugin_daemon_response("GET", "plugin/test-tenant/test", bool) - - def test_streaming_with_invalid_json_line(self, plugin_client, mock_config): - """Test streaming with invalid JSON in one line.""" - # Arrange - stream_data = [ - 'data: {"code": 0, "message": "", "data": {"value": "valid"}}', - "data: {invalid json}", - ] - - mock_response = httpx.Response(200, text="\n".join(stream_data)) - - with patch("httpx.stream", autospec=True) as mock_stream: - mock_stream.return_value.__enter__.return_value = mock_response - - # Act - class StreamModel(BaseModel): - value: str - - results = plugin_client._request_with_plugin_daemon_response_stream( - "POST", "plugin/test-tenant/stream", StreamModel - ) - - # Assert - first_result = next(results) - assert first_result.value == "valid" - - with pytest.raises(ValueError, match="\\{invalid json\\}"): - next(results) - - def test_request_with_bytes_data(self, plugin_client, mock_config): - """Test request with bytes data.""" - # Arrange - mock_response = httpx.Response(200) - - with patch("httpx.request", return_value=mock_response, autospec=True) as mock_request: - # Act - plugin_client._request("POST", "plugin/test-tenant/upload", data=b"binary data") - - # Assert - call_kwargs = mock_request.call_args[1] - assert call_kwargs["content"] == b"binary data" - - def test_request_with_files(self, plugin_client, mock_config): - """Test request with file upload.""" - # Arrange - mock_response = httpx.Response(200) - - files = {"file": ("test.txt", b"file content", "text/plain")} - - with patch("httpx.request", return_value=mock_response, autospec=True) as mock_request: - # Act - plugin_client._request("POST", "plugin/test-tenant/upload", files=files) - - # Assert - call_kwargs = mock_request.call_args[1] - assert call_kwargs["files"] == files - - def test_streaming_empty_response(self, plugin_client, mock_config): - """Test streaming with empty response.""" - # Arrange - mock_response = httpx.Response(200, content=b"") - - with patch("httpx.stream", autospec=True) as mock_stream: - mock_stream.return_value.__enter__.return_value = mock_response - - # Act - results = list(plugin_client._stream_request("POST", "plugin/test-tenant/stream")) - - # Assert - assert len(results) == 0 - - def test_daemon_inner_error_with_code_500(self, plugin_client, mock_config): - """Test handling of daemon inner error with code -500 in stream.""" - # Arrange - error_obj = json.dumps({"error_type": "PluginDaemonInnerError", "message": "Internal error"}) - stream_data = [ - f'data: {{"code": -500, "message": {json.dumps(error_obj)}, "data": null}}', - ] - - mock_response = httpx.Response(200, text="\n".join(stream_data)) - - with patch("httpx.stream", autospec=True) as mock_stream: - mock_stream.return_value.__enter__.return_value = mock_response - - # Act & Assert - class StreamModel(BaseModel): - data: str - - results = plugin_client._request_with_plugin_daemon_response_stream( - "POST", "plugin/test-tenant/stream", StreamModel - ) - with pytest.raises(PluginDaemonInnerError) as exc_info: - next(results) - assert exc_info.value.code == -500 - - def test_non_json_error_message(self, plugin_client, mock_config): - """Test handling of non-JSON error message.""" - # Arrange - mock_response = httpx.Response(200, json={"code": -1, "message": "Plain text error message", "data": None}) - - with patch("httpx.request", return_value=mock_response, autospec=True): - # Act & Assert - with pytest.raises(ValueError, match="Plain text error message, code: \\-1") as exc_info: - plugin_client._request_with_plugin_daemon_response("GET", "plugin/test-tenant/test", bool) - assert "Plain text error message" in str(exc_info.value) - - -class TestPluginRuntimeAdvancedScenarios: - """Advanced test scenarios for plugin runtime. - - Tests cover: - - Complex error recovery - - Concurrent request handling - - Plugin state management - - Advanced streaming patterns - """ - - @pytest.fixture - def plugin_client(self): - """Create a BasePluginClient instance for testing.""" - return BasePluginClient() - - @pytest.fixture - def mock_config(self, config_overrides): - """Mock plugin daemon configuration.""" - config_overrides( - PLUGIN_DAEMON_URL="http://127.0.0.1:5002", - PLUGIN_DAEMON_KEY="test-key", - ) - - def test_multiple_sequential_requests(self, plugin_client, mock_config): - """Test multiple sequential requests to the same endpoint.""" - # Arrange - mock_response = httpx.Response(200, json={"code": 0, "message": "", "data": True}) - - with patch("httpx.request", return_value=mock_response, autospec=True) as mock_request: - # Act - for i in range(5): - result = plugin_client._request_with_plugin_daemon_response("GET", f"plugin/test-tenant/test/{i}", bool) - assert result is True - - # Assert - assert mock_request.call_count == 5 - - def test_request_with_complex_nested_data(self, plugin_client, mock_config): - """Test request with complex nested data structures.""" - - # Arrange - class ComplexModel(BaseModel): - nested: dict[str, Any] - items: list[dict[str, Any]] - - complex_data = { - "nested": {"level1": {"level2": {"level3": "deep_value"}}}, - "items": [ - {"id": 1, "name": "item1"}, - {"id": 2, "name": "item2"}, - ], - } - - mock_response = httpx.Response(200, json={"code": 0, "message": "", "data": complex_data}) - - with patch("httpx.request", return_value=mock_response, autospec=True): - # Act - result = plugin_client._request_with_plugin_daemon_response( - "POST", "plugin/test-tenant/complex", ComplexModel - ) - - # Assert - assert result.nested["level1"]["level2"]["level3"] == "deep_value" - assert len(result.items) == 2 - assert result.items[0]["id"] == 1 - - def test_streaming_with_multiple_chunk_types(self, plugin_client, mock_config): - """Test streaming with different chunk types in sequence.""" - - # Arrange - class MultiTypeModel(BaseModel): - type: str - data: dict[str, Any] - - stream_data = [ - '{"code": 0, "message": "", "data": {"type": "start", "data": {"status": "initializing"}}}', - '{"code": 0, "message": "", "data": {"type": "progress", "data": {"percent": 50}}}', - '{"code": 0, "message": "", "data": {"type": "complete", "data": {"result": "success"}}}', - ] - - mock_response = httpx.Response(200, text="\n".join(stream_data)) - - with patch("httpx.stream", autospec=True) as mock_stream: - mock_stream.return_value.__enter__.return_value = mock_response - - # Act - results = list( - plugin_client._request_with_plugin_daemon_response_stream( - "POST", "plugin/test-tenant/multi-stream", MultiTypeModel - ) - ) - - # Assert - assert len(results) == 3 - assert results[0].type == "start" - assert results[1].type == "progress" - assert results[2].type == "complete" - assert results[1].data["percent"] == 50 - - def test_error_recovery_with_retry_pattern(self, plugin_client, mock_config): - """Test error recovery pattern (simulated retry logic).""" - # Arrange - call_count = 0 - - def side_effect[**P](*args: P.args, **kwargs: P.kwargs): - nonlocal call_count - call_count += 1 - if call_count < 3: - raise httpx.RequestError("Temporary failure") - mock_response = httpx.Response(200) - return mock_response - - with patch("httpx.request", side_effect=side_effect, autospec=True): - # Act & Assert - First two calls should fail - with pytest.raises(PluginDaemonInnerError): - plugin_client._request("GET", "plugin/test-tenant/test") - - with pytest.raises(PluginDaemonInnerError): - plugin_client._request("GET", "plugin/test-tenant/test") - - # Third call should succeed - response = plugin_client._request("GET", "plugin/test-tenant/test") - assert response.status_code == 200 - - def test_request_with_custom_headers_preservation(self, plugin_client, mock_config): - """Test that custom headers are preserved through request pipeline.""" - # Arrange - custom_headers = { - "X-Custom-Header": "custom-value", - "X-Request-ID": "req-123", - "X-Tenant-ID": "tenant-456", - } - - mock_response = httpx.Response(200) - - with patch("httpx.request", return_value=mock_response, autospec=True) as mock_request: - # Act - plugin_client._request("GET", "plugin/test-tenant/test", headers=custom_headers) - - # Assert - call_kwargs = mock_request.call_args[1] - for key, value in custom_headers.items(): - assert call_kwargs["headers"][key] == value - - def test_streaming_with_large_chunks(self, plugin_client, mock_config): - """Test streaming with large data chunks.""" - - # Arrange - class LargeChunkModel(BaseModel): - chunk_id: int - data: str - - # Create large chunks (simulating large data transfer) - large_data = "x" * 10000 # 10KB of data - stream_data = [ - f'{{"code": 0, "message": "", "data": {{"chunk_id": {i}, "data": "{large_data}"}}}}' for i in range(10) - ] - - mock_response = httpx.Response(200, text="\n".join(stream_data)) - - with patch("httpx.stream", autospec=True) as mock_stream: - mock_stream.return_value.__enter__.return_value = mock_response - - # Act - results = list( - plugin_client._request_with_plugin_daemon_response_stream( - "POST", "plugin/test-tenant/large-stream", LargeChunkModel - ) - ) - - # Assert - assert len(results) == 10 - for i, result in enumerate(results): - assert result.chunk_id == i - assert len(result.data) == 10000 - - -class TestPluginRuntimeSecurityAndValidation: - """Tests for security and validation aspects of plugin runtime. - - Tests cover: - - Input validation - - Security headers - - Authentication failures - - Authorization checks - """ - - @pytest.fixture - def plugin_client(self): - """Create a BasePluginClient instance for testing.""" - return BasePluginClient() - - @pytest.fixture - def mock_config(self, config_overrides): - """Mock plugin daemon configuration.""" - config_overrides( - PLUGIN_DAEMON_URL="http://127.0.0.1:5002", - PLUGIN_DAEMON_KEY="secure-key-123", - ) - - def test_api_key_header_always_present(self, plugin_client, mock_config): - """Test that API key header is always included in requests.""" - # Arrange - mock_response = httpx.Response(200) - - with patch("httpx.request", return_value=mock_response, autospec=True) as mock_request: - # Act - plugin_client._request("GET", "plugin/test-tenant/test") - - # Assert - call_kwargs = mock_request.call_args[1] - assert "X-Api-Key" in call_kwargs["headers"] - assert call_kwargs["headers"]["X-Api-Key"] == "secure-key-123" - - def test_request_with_sensitive_data_in_body(self, plugin_client, mock_config): - """Test handling of sensitive data in request body.""" - # Arrange - sensitive_data = { - "api_key": "secret-api-key", - "password": "secret-password", - "credentials": {"token": "secret-token"}, - } - - mock_response = httpx.Response(200, json={"code": 0, "message": "", "data": True}) - - with patch("httpx.request", return_value=mock_response, autospec=True) as mock_request: - # Act - plugin_client._request_with_plugin_daemon_response( - "POST", - "plugin/test-tenant/validate", - bool, - data=sensitive_data, - headers={"Content-Type": "application/json"}, - ) - - # Assert - Verify data was sent - call_kwargs = mock_request.call_args[1] - assert "content" in call_kwargs or "data" in call_kwargs - - def test_unauthorized_access_with_invalid_key(self, plugin_client, mock_config): - """Test handling of unauthorized access with invalid API key.""" - # Arrange - error_message = json.dumps({"error_type": "PluginDaemonUnauthorizedError", "message": "Invalid API key"}) - mock_response = httpx.Response(200, json={"code": -1, "message": error_message, "data": None}) - - with patch("httpx.request", return_value=mock_response, autospec=True): - # Act & Assert - with pytest.raises(PluginDaemonUnauthorizedError) as exc_info: - plugin_client._request_with_plugin_daemon_response("GET", "plugin/test-tenant/test", bool) - assert "Invalid API key" in exc_info.value.description - - def test_request_parameter_validation(self, plugin_client, mock_config): - """Test validation of request parameters.""" - # Arrange - invalid_params = { - "page": -1, # Invalid negative page - "limit": 0, # Invalid zero limit - } - - error_message = json.dumps( - {"error_type": "PluginDaemonBadRequestError", "message": "Invalid parameters: page must be positive"} - ) - mock_response = httpx.Response(200, json={"code": -1, "message": error_message, "data": None}) - - with patch("httpx.request", return_value=mock_response, autospec=True): - # Act & Assert - with pytest.raises(PluginDaemonBadRequestError) as exc_info: - plugin_client._request_with_plugin_daemon_response( - "GET", "plugin/test-tenant/list", list, params=invalid_params - ) - assert "Invalid parameters" in exc_info.value.description - - def test_content_type_header_validation(self, plugin_client, mock_config): - """Test that Content-Type header is properly set for JSON requests.""" - # Arrange - mock_response = httpx.Response(200) - - with patch("httpx.request", return_value=mock_response, autospec=True) as mock_request: - # Act - plugin_client._request( - "POST", "plugin/test-tenant/test", headers={"Content-Type": "application/json"}, data={"key": "value"} - ) - - # Assert - call_kwargs = mock_request.call_args[1] - assert call_kwargs["headers"]["Content-Type"] == "application/json" - - -class TestPluginRuntimePerformanceScenarios: - """Tests for performance-related scenarios in plugin runtime. - - Tests cover: - - High-volume streaming - - Concurrent operations simulation - - Memory-efficient processing - - Timeout handling under load - """ - - @pytest.fixture - def plugin_client(self): - """Create a BasePluginClient instance for testing.""" - return BasePluginClient() - - @pytest.fixture - def mock_config(self, config_overrides): - """Mock plugin daemon configuration.""" - config_overrides( - PLUGIN_DAEMON_URL="http://127.0.0.1:5002", - PLUGIN_DAEMON_KEY="test-key", - ) - - def test_high_volume_streaming(self, plugin_client, mock_config): - """Test streaming with high volume of chunks.""" - - # Arrange - class StreamChunk(BaseModel): - index: int - value: str - - # Generate 100 chunks - stream_data = [ - f'{{"code": 0, "message": "", "data": {{"index": {i}, "value": "chunk_{i}"}}}}' for i in range(100) - ] - - mock_response = httpx.Response(200, text="\n".join(stream_data)) - - with patch("httpx.stream", autospec=True) as mock_stream: - mock_stream.return_value.__enter__.return_value = mock_response - - # Act - results = list( - plugin_client._request_with_plugin_daemon_response_stream( - "POST", "plugin/test-tenant/high-volume", StreamChunk - ) - ) - - # Assert - assert len(results) == 100 - assert results[0].index == 0 - assert results[99].index == 99 - assert results[50].value == "chunk_50" - - def test_streaming_memory_efficiency(self, plugin_client, mock_config): - """Test that streaming processes chunks one at a time (memory efficient).""" - - # Arrange - class ChunkModel(BaseModel): - data: str - - processed_chunks = [] - - def process_chunk(chunk): - """Simulate processing each chunk individually.""" - processed_chunks.append(chunk.data) - return chunk - - stream_data = [f'{{"code": 0, "message": "", "data": {{"data": "chunk_{i}"}}}}' for i in range(10)] - - mock_response = httpx.Response(200, text="\n".join(stream_data)) - - with patch("httpx.stream", autospec=True) as mock_stream: - mock_stream.return_value.__enter__.return_value = mock_response - - # Act - Process chunks one by one - for chunk in plugin_client._request_with_plugin_daemon_response_stream( - "POST", "plugin/test-tenant/stream", ChunkModel - ): - process_chunk(chunk) - - # Assert - assert len(processed_chunks) == 10 - - def test_timeout_with_slow_response(self, plugin_client, mock_config): - """Test timeout handling with slow response simulation.""" - # Arrange - with patch("httpx.request", side_effect=httpx.TimeoutException("Request timed out after 30s"), autospec=True): - # Act & Assert - with pytest.raises(PluginDaemonInnerError) as exc_info: - plugin_client._request("GET", "plugin/test-tenant/slow-endpoint") - assert exc_info.value.code == -500 - - def test_concurrent_request_simulation(self, plugin_client, mock_config): - """Test simulation of concurrent requests (sequential execution in test).""" - # Arrange - mock_response = httpx.Response(200, json={"code": 0, "message": "", "data": True}) - - request_results = [] - - with patch("httpx.request", return_value=mock_response, autospec=True): - # Act - Simulate 10 concurrent requests - for i in range(10): - result = plugin_client._request_with_plugin_daemon_response( - "GET", f"plugin/test-tenant/concurrent/{i}", bool - ) - request_results.append(result) - - # Assert - assert len(request_results) == 10 - assert all(result is True for result in request_results) - - -class TestPluginToolManagerAdvanced: - """Advanced tests for PluginToolManager functionality. - - Tests cover: - - Complex tool invocations - - Runtime parameter handling - - Tool provider discovery - - Advanced credential scenarios - """ - - @pytest.fixture - def tool_manager(self): - """Create a PluginToolManager instance for testing.""" - return PluginToolManager() - - @pytest.fixture - def mock_config(self, config_overrides): - """Mock plugin daemon configuration.""" - config_overrides( - PLUGIN_DAEMON_URL="http://127.0.0.1:5002", - PLUGIN_DAEMON_KEY="test-key", - ) - - def test_tool_invocation_with_complex_parameters(self, tool_manager, mock_config): - """Test tool invocation with complex parameter structures.""" - # Arrange - complex_params = { - "simple_string": "value", - "number": 42, - "boolean": True, - "nested_object": {"key1": "value1", "key2": ["item1", "item2"]}, - "array": [1, 2, 3, 4, 5], - } - - stream_data = [ - ( - 'data: {"code": 0, "message": "", "data": {"type": "text", ' - '"message": {"text": "Complex params processed"}}}' - ), - ] - - mock_response = httpx.Response(200, text="\n".join(stream_data)) - - with patch("httpx.stream", autospec=True) as mock_stream: - mock_stream.return_value.__enter__.return_value = mock_response - - # Act - results = list( - tool_manager.invoke( - tenant_id="test-tenant", - user_id="test-user", - tool_provider="langgenius/test-plugin/test-provider", - tool_name="complex-tool", - credentials={"api_key": "test-key"}, - credential_type=CredentialType.API_KEY, - tool_parameters=complex_params, - ) - ) - - # Assert - assert len(results) > 0 - - def test_tool_invocation_with_conversation_context(self, tool_manager, mock_config): - """Test tool invocation with conversation context.""" - # Arrange - stream_data = [ - 'data: {"code": 0, "message": "", "data": {"type": "text", "message": {"text": "Context-aware result"}}}', - ] - - mock_response = httpx.Response(200, text="\n".join(stream_data)) - - with patch("httpx.stream", autospec=True) as mock_stream: - mock_stream.return_value.__enter__.return_value = mock_response - - # Act - results = list( - tool_manager.invoke( - tenant_id="test-tenant", - user_id="test-user", - tool_provider="langgenius/test-plugin/test-provider", - tool_name="test-tool", - credentials={"api_key": "test-key"}, - credential_type=CredentialType.API_KEY, - tool_parameters={"query": "test"}, - conversation_id="conv-123", - app_id="app-456", - message_id="msg-789", - ) - ) - - # Assert - assert len(results) > 0 - - def test_get_runtime_parameters_success(self, tool_manager, mock_config): - """Test successful retrieval of runtime parameters.""" - # Arrange - stream_data = [ - 'data: {"code": 0, "message": "", "data": {"parameters": []}}', - ] - - mock_response = httpx.Response(200, text="\n".join(stream_data)) - - with patch("httpx.stream", autospec=True) as mock_stream: - mock_stream.return_value.__enter__.return_value = mock_response - - # Act - result = tool_manager.get_runtime_parameters( - tenant_id="test-tenant", - user_id="test-user", - provider="langgenius/test-plugin/test-provider", - credentials={"api_key": "test-key"}, - tool="test-tool", - ) - - # Assert - assert isinstance(result, list) - - def test_validate_credentials_with_oauth(self, tool_manager, mock_config): - """Test credential validation with OAuth credentials.""" - # Arrange - oauth_credentials = { - "access_token": "oauth-token-123", - "refresh_token": "refresh-token-456", - "expires_at": 1234567890, - } - - stream_data = [ - 'data: {"code": 0, "message": "", "data": {"result": true}}', - ] - - mock_response = httpx.Response(200, text="\n".join(stream_data)) - - with patch("httpx.stream", autospec=True) as mock_stream: - mock_stream.return_value.__enter__.return_value = mock_response - - # Act - result = tool_manager.validate_provider_credentials( - tenant_id="test-tenant", - user_id="test-user", - provider="langgenius/test-plugin/oauth-provider", - credentials=oauth_credentials, - ) - - # Assert - assert result is True - - -class TestPluginInstallerAdvanced: - """Advanced tests for PluginInstaller functionality. - - Tests cover: - - Plugin package upload - - Bundle installation - - Plugin upgrade scenarios - - Dependency management - """ - - @pytest.fixture - def installer(self): - """Create a PluginInstaller instance for testing.""" - return PluginInstaller() - - @pytest.fixture - def mock_config(self, config_overrides): - """Mock plugin daemon configuration.""" - config_overrides( - PLUGIN_DAEMON_URL="http://127.0.0.1:5002", - PLUGIN_DAEMON_KEY="test-key", - ) - - def test_upload_plugin_package_success(self, installer, mock_config): - """Test successful plugin package upload.""" - # Arrange - plugin_package = b"fake-plugin-package-data" - mock_response = httpx.Response( - 200, - json={ - "code": 0, - "message": "", - "data": { - "unique_identifier": "test-org/test-plugin", - "manifest": { - "version": "1.0.0", - "author": "test-org", - "name": "test-plugin", - "description": {"en_US": "Test plugin"}, - "icon": "icon.png", - "label": {"en_US": "Test Plugin"}, - "created_at": "2024-01-01T00:00:00Z", - "resource": {"memory": 256}, - "plugins": {}, - "meta": {}, - }, - "verification": None, - }, - }, - ) - - with patch("httpx.request", return_value=mock_response, autospec=True): - # Act - result = installer.upload_pkg("test-tenant", plugin_package, verify_signature=False) - - # Assert - assert result.unique_identifier == "test-org/test-plugin" - - def test_fetch_plugin_readme_success(self, installer, mock_config): - """Test successful plugin readme fetch.""" - # Arrange - mock_response = httpx.Response( - 200, - json={ - "code": 0, - "message": "", - "data": {"content": "# Plugin README\n\nThis is a test plugin.", "language": "en"}, - }, - ) - - with patch("httpx.request", return_value=mock_response, autospec=True): - # Act - result = installer.fetch_plugin_readme("test-tenant", "test-org/test-plugin", "en") - - # Assert - assert "Plugin README" in result - assert "test plugin" in result - - def test_fetch_plugin_readme_not_found(self, installer, mock_config): - """Test plugin readme fetch when readme doesn't exist.""" - # Arrange - mock_response = httpx.Response(404) - - with patch("httpx.request", return_value=mock_response, autospec=True): - # Act & Assert - Should raise PluginDaemonClientSideError for 404 - with pytest.raises(PluginDaemonClientSideError): - installer.fetch_plugin_readme("test-tenant", "test-org/test-plugin", "en") - - def test_list_plugins_with_pagination(self, installer, mock_config): - """Test plugin listing with pagination.""" - # Arrange - mock_response = httpx.Response( - 200, - json={ - "code": 0, - "message": "", - "data": { - "list": [], - "total": 50, - }, - }, - ) - - with patch("httpx.request", return_value=mock_response, autospec=True): - # Act - result = installer.list_plugins_with_total("test-tenant", page=2, page_size=20) - - # Assert - assert result.total == 50 - assert isinstance(result.list, list) - - def test_check_tools_existence(self, installer, mock_config): - """Test checking existence of multiple tools.""" - # Arrange - from models.provider_ids import GenericProviderID - - provider_ids = [ - GenericProviderID("langgenius/plugin1/provider1"), - GenericProviderID("langgenius/plugin2/provider2"), - ] - - mock_response = httpx.Response(200, json={"code": 0, "message": "", "data": [True, False]}) - - with patch("httpx.request", return_value=mock_response, autospec=True): - # Act - result = installer.check_tools_existence("test-tenant", provider_ids) - - # Assert - assert len(result) == 2 - assert result[0] is True - assert result[1] is False diff --git a/api/tests/unit_tests/services/plugin/test_plugin_service.py b/api/tests/unit_tests/core/plugin/test_plugin_service.py similarity index 99% rename from api/tests/unit_tests/services/plugin/test_plugin_service.py rename to api/tests/unit_tests/core/plugin/test_plugin_service.py index 11b38a39892..ce7aea701d4 100644 --- a/api/tests/unit_tests/services/plugin/test_plugin_service.py +++ b/api/tests/unit_tests/core/plugin/test_plugin_service.py @@ -24,6 +24,8 @@ from graphon.model_runtime.entities.common_entities import I18nObject from graphon.model_runtime.entities.provider_entities import ConfigurateMethod, ProviderEntity from models.provider import Provider, ProviderCredential, ProviderType, TenantPreferredModelProvider from services.entities.feature_entities import PluginInstallationPermissionModel, PluginInstallationScope +from tests.unit_tests.services.plugin.conftest import mock_features as mock_features # noqa: PLC0414 +from tests.unit_tests.services.plugin.conftest import mock_installer as mock_installer # noqa: PLC0414 MODULE = "core.plugin.plugin_service" TENANT_ID = "11111111-1111-1111-1111-111111111111" diff --git a/api/tests/unit_tests/core/prompt/test_advanced_prompt_transform.py b/api/tests/unit_tests/core/prompt/test_advanced_prompt_transform.py index 26554ea01f4..0873b109d1c 100644 --- a/api/tests/unit_tests/core/prompt/test_advanced_prompt_transform.py +++ b/api/tests/unit_tests/core/prompt/test_advanced_prompt_transform.py @@ -15,8 +15,8 @@ from graphon.model_runtime.entities.message_entities import ( TextPromptMessageContent, UserPromptMessage, ) -from tests.unit_tests.config_override import apply_config_overrides -from tests.unit_tests.core.model_fixtures import make_model_config, make_token_buffer_memory +from tests.helpers.unit.config_override import apply_config_overrides +from tests.helpers.unit.core.model_fixtures import make_model_config, make_token_buffer_memory def test__get_completion_model_prompt_messages(): diff --git a/api/tests/unit_tests/core/prompt/test_agent_history_prompt_transform.py b/api/tests/unit_tests/core/prompt/test_agent_history_prompt_transform.py index 34aa7cf0248..82c7d246158 100644 --- a/api/tests/unit_tests/core/prompt/test_agent_history_prompt_transform.py +++ b/api/tests/unit_tests/core/prompt/test_agent_history_prompt_transform.py @@ -10,7 +10,7 @@ from graphon.model_runtime.entities.message_entities import ( ToolPromptMessage, UserPromptMessage, ) -from tests.unit_tests.core.model_fixtures import make_model_config, make_token_buffer_memory +from tests.helpers.unit.core.model_fixtures import make_model_config, make_token_buffer_memory def test_get_prompt(mocker: MockerFixture) -> None: diff --git a/api/tests/unit_tests/core/prompt/test_prompt_transform.py b/api/tests/unit_tests/core/prompt/test_prompt_transform.py index 7e63daecc95..7f3f4b33c53 100644 --- a/api/tests/unit_tests/core/prompt/test_prompt_transform.py +++ b/api/tests/unit_tests/core/prompt/test_prompt_transform.py @@ -9,7 +9,7 @@ from core.prompt.prompt_transform import PromptTransform from graphon.model_runtime.entities.common_entities import I18nObject from graphon.model_runtime.entities.message_entities import AssistantPromptMessage, PromptMessage, UserPromptMessage from graphon.model_runtime.entities.model_entities import AIModelEntity, ModelPropertyKey, ParameterRule, ParameterType -from tests.unit_tests.core.model_fixtures import make_model_config, make_token_buffer_memory +from tests.helpers.unit.core.model_fixtures import make_model_config, make_token_buffer_memory class TestPromptTransform: diff --git a/api/tests/unit_tests/core/prompt/test_simple_prompt_transform.py b/api/tests/unit_tests/core/prompt/test_simple_prompt_transform.py index dbc5b617ec7..3894ae7f8b5 100644 --- a/api/tests/unit_tests/core/prompt/test_simple_prompt_transform.py +++ b/api/tests/unit_tests/core/prompt/test_simple_prompt_transform.py @@ -19,7 +19,7 @@ from graphon.model_runtime.entities.message_entities import ( UserPromptMessage, ) from models.model import AppMode -from tests.unit_tests.core.model_fixtures import make_model_config, make_token_buffer_memory +from tests.helpers.unit.core.model_fixtures import make_model_config, make_token_buffer_memory def test_get_common_chat_app_prompt_template_with_pcqm() -> None: diff --git a/api/tests/unit_tests/core/prompt/test_extract_thread_messages.py b/api/tests/unit_tests/core/prompt/utils/test_extract_thread_messages.py similarity index 100% rename from api/tests/unit_tests/core/prompt/test_extract_thread_messages.py rename to api/tests/unit_tests/core/prompt/utils/test_extract_thread_messages.py diff --git a/api/tests/unit_tests/core/prompt/test_prompt_message.py b/api/tests/unit_tests/core/prompt/utils/test_prompt_message_util.py similarity index 100% rename from api/tests/unit_tests/core/prompt/test_prompt_message.py rename to api/tests/unit_tests/core/prompt/utils/test_prompt_message_util.py diff --git a/api/tests/unit_tests/core/rag/datasource/keyword/jieba/test_jieba.py b/api/tests/unit_tests/core/rag/datasource/keyword/jieba/test_jieba.py index 41edd48db92..4f23dbc64e0 100644 --- a/api/tests/unit_tests/core/rag/datasource/keyword/jieba/test_jieba.py +++ b/api/tests/unit_tests/core/rag/datasource/keyword/jieba/test_jieba.py @@ -12,8 +12,8 @@ from core.rag.models.document import Document from models.dataset import Dataset, DatasetKeywordTable, DocumentSegment from repositories.knowledge.dataset_read_repository import get_dataset_keyword_table from repositories.knowledge.keyword_table_repository import load_keyword_table -from tests.unit_tests.config_override import apply_config_overrides -from tests.unit_tests.model_factories import make_dataset +from tests.helpers.unit.config_override import apply_config_overrides +from tests.helpers.unit.model_factories import make_dataset @dataclass diff --git a/api/tests/unit_tests/core/rag/datasource/keyword/test_keyword_base.py b/api/tests/unit_tests/core/rag/datasource/keyword/test_keyword_base.py index c3d2eba835c..9578b926840 100644 --- a/api/tests/unit_tests/core/rag/datasource/keyword/test_keyword_base.py +++ b/api/tests/unit_tests/core/rag/datasource/keyword/test_keyword_base.py @@ -6,7 +6,7 @@ from sqlalchemy.orm import Session from core.rag.datasource.keyword.keyword_base import BaseKeyword from core.rag.models.document import Document from models import Dataset -from tests.unit_tests.model_factories import make_dataset +from tests.helpers.unit.model_factories import make_dataset def _dataset() -> Dataset: diff --git a/api/tests/unit_tests/core/rag/datasource/keyword/test_keyword_factory.py b/api/tests/unit_tests/core/rag/datasource/keyword/test_keyword_factory.py index 4c5bae70ea8..755f7180e92 100644 --- a/api/tests/unit_tests/core/rag/datasource/keyword/test_keyword_factory.py +++ b/api/tests/unit_tests/core/rag/datasource/keyword/test_keyword_factory.py @@ -9,7 +9,7 @@ from core.rag.datasource.keyword.keyword_factory import Keyword from core.rag.datasource.keyword.keyword_type import KeyWordType from core.rag.models.document import Document from models.dataset import Dataset -from tests.unit_tests.config_override import apply_config_overrides +from tests.helpers.unit.config_override import apply_config_overrides def test_get_keyword_factory_returns_jieba_factory(monkeypatch: pytest.MonkeyPatch) -> None: diff --git a/api/tests/unit_tests/core/rag/datasource/test_retrieval_attachment_access.py b/api/tests/unit_tests/core/rag/datasource/test_retrieval_attachment_access.py deleted file mode 100644 index 8f89aca59d0..00000000000 --- a/api/tests/unit_tests/core/rag/datasource/test_retrieval_attachment_access.py +++ /dev/null @@ -1,238 +0,0 @@ -from __future__ import annotations - -from datetime import UTC, datetime -from uuid import uuid4 - -import pytest -from sqlalchemy import select -from sqlalchemy.orm import Session, sessionmaker - -from core.app.entities.app_invoke_entities import InvokeFrom, UserFrom -from core.app.file_access import ( - DatabaseFileAccessController, - FileAccessScope, - bind_file_access_scope, - get_current_file_access_scope, - grant_retriever_segment_access, - grant_upload_file_access, - is_retriever_segment_access_granted, -) -from core.rag.datasource.retrieval_service import RetrievalService -from core.rag.embedding.retrieval import RetrievalSegments -from core.rag.models.document import Document as RagDocument -from core.rag.retrieval import dataset_retrieval as dataset_retrieval_module -from core.rag.retrieval.dataset_retrieval import DatasetRetrieval -from core.workflow.nodes.knowledge_retrieval.retrieval import KnowledgeRetrievalRequest -from extensions.storage.storage_type import StorageType -from models import UploadFile -from models.dataset import Dataset, DocumentSegment, SegmentAttachmentBinding -from models.dataset import Document as DatasetDocument -from models.enums import CreatorUserRole, DataSourceType, DocumentCreatedFrom, SegmentStatus - - -def test_file_access_grants_ignore_empty_inputs_and_missing_scope() -> None: - grant_upload_file_access(["upload-file-id"]) - grant_retriever_segment_access(["segment-id"]) - assert is_retriever_segment_access_granted("segment-id") is True - - scope = FileAccessScope( - tenant_id=str(uuid4()), - user_id=str(uuid4()), - user_from=UserFrom.END_USER, - invoke_from=InvokeFrom.WEB_APP, - ) - - with bind_file_access_scope(scope): - grant_upload_file_access([""]) - grant_retriever_segment_access([""]) - current_scope = get_current_file_access_scope() - - assert current_scope is not None - assert current_scope.granted_upload_file_ids == frozenset() - assert current_scope.granted_retriever_segment_ids == frozenset() - - -@pytest.mark.parametrize("sqlite_session", [(UploadFile, SegmentAttachmentBinding)], indirect=True) -def test_segment_attachment_lookup_grants_returned_upload_files_to_current_scope(sqlite_session: Session) -> None: - tenant_id = str(uuid4()) - dataset_id = str(uuid4()) - upload_file_id = str(uuid4()) - segment_id = str(uuid4()) - user_id = str(uuid4()) - upload_file = UploadFile( - tenant_id=tenant_id, - storage_type=StorageType.LOCAL, - key="uploads/chart.png", - name="chart.png", - size=1024, - extension="png", - mime_type="image/png", - created_by_role=CreatorUserRole.END_USER, - created_by=user_id, - created_at=datetime.now(UTC), - used=False, - ) - upload_file.id = upload_file_id - binding = SegmentAttachmentBinding( - tenant_id=tenant_id, - dataset_id=dataset_id, - document_id=str(uuid4()), - segment_id=segment_id, - attachment_id=upload_file_id, - ) - sqlite_session.add_all([upload_file, binding]) - sqlite_session.commit() - scope = FileAccessScope( - tenant_id=tenant_id, - user_id=user_id, - user_from=UserFrom.END_USER, - invoke_from=InvokeFrom.WEB_APP, - ) - - with bind_file_access_scope(scope): - result = RetrievalService.get_segment_attachment_infos( - [upload_file_id], - dataset_refs=((tenant_id, dataset_id),), - session=sqlite_session, - ) - scoped_stmt = DatabaseFileAccessController().apply_upload_file_filters( - select(UploadFile).where(UploadFile.id == upload_file_id) - ) - - assert result[0]["attachment_id"] == upload_file_id - whereclause = str(scoped_stmt.whereclause) - assert "upload_files.created_by_role" in whereclause - assert "upload_files.id IN" in whereclause - - -@pytest.mark.parametrize("sqlite_session", [(UploadFile, SegmentAttachmentBinding)], indirect=True) -def test_segment_attachment_lookup_rejects_another_tenants_upload(sqlite_session: Session) -> None: - tenant_id = str(uuid4()) - foreign_tenant_id = str(uuid4()) - upload_file_id = str(uuid4()) - foreign_dataset_id = str(uuid4()) - upload_file = UploadFile( - tenant_id=foreign_tenant_id, - storage_type=StorageType.LOCAL, - key="uploads/foreign.png", - name="foreign.png", - size=1024, - extension="png", - mime_type="image/png", - created_by_role=CreatorUserRole.END_USER, - created_by=str(uuid4()), - created_at=datetime.now(UTC), - used=False, - ) - upload_file.id = upload_file_id - sqlite_session.add_all( - [ - upload_file, - SegmentAttachmentBinding( - tenant_id=foreign_tenant_id, - dataset_id=foreign_dataset_id, - document_id=str(uuid4()), - segment_id=str(uuid4()), - attachment_id=upload_file_id, - ), - ] - ) - sqlite_session.commit() - scope = FileAccessScope( - tenant_id=tenant_id, - user_id=str(uuid4()), - user_from=UserFrom.END_USER, - invoke_from=InvokeFrom.WEB_APP, - ) - - with bind_file_access_scope(scope): - result = RetrievalService.get_segment_attachment_infos( - [upload_file_id], - dataset_refs=((tenant_id, str(uuid4())),), - session=sqlite_session, - ) - current_scope = get_current_file_access_scope() - - assert result == [] - assert current_scope is not None - assert current_scope.granted_upload_file_ids == frozenset() - - -@pytest.mark.parametrize("sqlite_session", [(Dataset, DatasetDocument, DocumentSegment)], indirect=True) -def test_knowledge_retrieval_grants_returned_segments_to_current_scope( - monkeypatch: pytest.MonkeyPatch, - sqlite_session: Session, - sqlite_session_factory: sessionmaker[Session], -) -> None: - tenant_id = str(uuid4()) - dataset_id = str(uuid4()) - document_id = str(uuid4()) - segment_id = str(uuid4()) - created_by = str(uuid4()) - dataset = Dataset(tenant_id=tenant_id, name="Dataset", created_by=created_by) - dataset.id = dataset_id - document = DatasetDocument( - tenant_id=tenant_id, - dataset_id=dataset_id, - position=1, - data_source_type=DataSourceType.UPLOAD_FILE, - batch="batch-1", - name="Document", - created_from=DocumentCreatedFrom.API, - created_by=created_by, - doc_metadata={}, - ) - document.id = document_id - segment = DocumentSegment( - tenant_id=tenant_id, - dataset_id=dataset_id, - document_id=document_id, - position=1, - content="segment content", - word_count=10, - tokens=2, - created_by=created_by, - index_node_hash="hash", - hit_count=1, - ) - segment.id = segment_id - segment.status = SegmentStatus.COMPLETED - sqlite_session.add_all([dataset, document, segment]) - sqlite_session.commit() - - record = RetrievalSegments(segment=segment, score=0.8) - retrieval = DatasetRetrieval() - monkeypatch.setattr(retrieval, "_check_knowledge_rate_limit", lambda tenant_id: None) - monkeypatch.setattr(retrieval, "_get_available_datasets", lambda tenant_id, dataset_ids: [dataset]) - monkeypatch.setattr( - retrieval, - "multiple_retrieve", - lambda **kwargs: [RagDocument(page_content="segment content", provider="dify")], - ) - monkeypatch.setattr(RetrievalService, "format_retrieval_documents", lambda _session, documents: [record]) - monkeypatch.setattr(dataset_retrieval_module.session_factory, "create_session", sqlite_session_factory) - scope = FileAccessScope( - tenant_id=tenant_id, - user_id=str(uuid4()), - user_from=UserFrom.END_USER, - invoke_from=InvokeFrom.WEB_APP, - ) - - with bind_file_access_scope(scope): - results = retrieval.knowledge_retrieval( - sqlite_session, - KnowledgeRetrievalRequest( - tenant_id=tenant_id, - user_id=str(uuid4()), - app_id=str(uuid4()), - user_from=UserFrom.END_USER.value, - dataset_ids=[dataset_id], - query="desktop picture", - retrieval_mode="multiple", - ), - ) - current_scope = get_current_file_access_scope() - - assert results[0].metadata.segment_id == segment_id - assert current_scope is not None - assert segment_id in current_scope.granted_retriever_segment_ids diff --git a/api/tests/unit_tests/core/rag/datasource/test_datasource_retrieval.py b/api/tests/unit_tests/core/rag/datasource/test_retrieval_service.py similarity index 85% rename from api/tests/unit_tests/core/rag/datasource/test_datasource_retrieval.py rename to api/tests/unit_tests/core/rag/datasource/test_retrieval_service.py index df4d8b1477c..441f83fd1a9 100644 --- a/api/tests/unit_tests/core/rag/datasource/test_datasource_retrieval.py +++ b/api/tests/unit_tests/core/rag/datasource/test_retrieval_service.py @@ -1,3 +1,5 @@ +from __future__ import annotations + from dataclasses import dataclass from datetime import UTC, datetime from typing import Any @@ -6,24 +8,41 @@ from uuid import uuid4 import pytest from pytest_mock import MockerFixture +from sqlalchemy import select from sqlalchemy.engine import Engine from sqlalchemy.exc import SQLAlchemyError -from sqlalchemy.orm import Session +from sqlalchemy.orm import Session, sessionmaker +from core.app.entities.app_invoke_entities import InvokeFrom, UserFrom +from core.app.file_access import ( + DatabaseFileAccessController, + FileAccessScope, + bind_file_access_scope, + get_current_file_access_scope, + grant_retriever_segment_access, + grant_upload_file_access, + is_retriever_segment_access_granted, +) from core.model_context import get_credit_usage_metadata, use_credit_usage_metadata from core.rag.datasource import retrieval_service as retrieval_service_module from core.rag.datasource.retrieval_service import RetrievalService +from core.rag.embedding.retrieval import RetrievalSegments from core.rag.index_processor.constant.doc_type import DocType from core.rag.index_processor.constant.index_type import IndexStructureType from core.rag.index_processor.constant.query_type import QueryType from core.rag.models.document import Document +from core.rag.models.document import Document as RagDocument from core.rag.rerank.rerank_type import RerankMode +from core.rag.retrieval import dataset_retrieval as dataset_retrieval_module +from core.rag.retrieval.dataset_retrieval import DatasetRetrieval from core.rag.retrieval.retrieval_methods import RetrievalMethod +from core.workflow.nodes.knowledge_retrieval.retrieval import KnowledgeRetrievalRequest from extensions.storage.storage_type import StorageType +from models import UploadFile as _retrieval_attachment_access_UploadFile from models.dataset import ChildChunk, Dataset, DocumentSegment, DocumentSegmentSummary, SegmentAttachmentBinding from models.dataset import Document as DatasetDocument from models.enums import CreatorUserRole, DataSourceType, DocumentCreatedFrom, SegmentStatus, SummaryStatus -from models.model import UploadFile +from models.model import UploadFile as _datasource_retrieval_UploadFile def create_mock_document( @@ -155,8 +174,8 @@ def _upload_file( extension: str = "png", mime_type: str | None = None, size: int = 42, -) -> UploadFile: - upload = UploadFile( +) -> _datasource_retrieval_UploadFile: + upload = _datasource_retrieval_UploadFile( tenant_id=tenant_id, storage_type=StorageType.LOCAL, key=f"files/{file_id}", @@ -1248,3 +1267,217 @@ class TestRetrievalServiceInternals: ] ) assert mock_sign.call_count == 2 + + +def test_file_access_grants_ignore_empty_inputs_and_missing_scope() -> None: + grant_upload_file_access(["upload-file-id"]) + grant_retriever_segment_access(["segment-id"]) + assert is_retriever_segment_access_granted("segment-id") is True + + scope = FileAccessScope( + tenant_id=str(uuid4()), + user_id=str(uuid4()), + user_from=UserFrom.END_USER, + invoke_from=InvokeFrom.WEB_APP, + ) + + with bind_file_access_scope(scope): + grant_upload_file_access([""]) + grant_retriever_segment_access([""]) + current_scope = get_current_file_access_scope() + + assert current_scope is not None + assert current_scope.granted_upload_file_ids == frozenset() + assert current_scope.granted_retriever_segment_ids == frozenset() + + +@pytest.mark.parametrize( + "sqlite_session", [(_retrieval_attachment_access_UploadFile, SegmentAttachmentBinding)], indirect=True +) +def test_segment_attachment_lookup_grants_returned_upload_files_to_current_scope(sqlite_session: Session) -> None: + tenant_id = str(uuid4()) + dataset_id = str(uuid4()) + upload_file_id = str(uuid4()) + segment_id = str(uuid4()) + user_id = str(uuid4()) + upload_file = _retrieval_attachment_access_UploadFile( + tenant_id=tenant_id, + storage_type=StorageType.LOCAL, + key="uploads/chart.png", + name="chart.png", + size=1024, + extension="png", + mime_type="image/png", + created_by_role=CreatorUserRole.END_USER, + created_by=user_id, + created_at=datetime.now(UTC), + used=False, + ) + upload_file.id = upload_file_id + binding = SegmentAttachmentBinding( + tenant_id=tenant_id, + dataset_id=dataset_id, + document_id=str(uuid4()), + segment_id=segment_id, + attachment_id=upload_file_id, + ) + sqlite_session.add_all([upload_file, binding]) + sqlite_session.commit() + scope = FileAccessScope( + tenant_id=tenant_id, + user_id=user_id, + user_from=UserFrom.END_USER, + invoke_from=InvokeFrom.WEB_APP, + ) + + with bind_file_access_scope(scope): + result = RetrievalService.get_segment_attachment_infos( + [upload_file_id], + dataset_refs=((tenant_id, dataset_id),), + session=sqlite_session, + ) + scoped_stmt = DatabaseFileAccessController().apply_upload_file_filters( + select(_retrieval_attachment_access_UploadFile).where( + _retrieval_attachment_access_UploadFile.id == upload_file_id + ) + ) + + assert result[0]["attachment_id"] == upload_file_id + whereclause = str(scoped_stmt.whereclause) + assert "upload_files.created_by_role" in whereclause + assert "upload_files.id IN" in whereclause + + +@pytest.mark.parametrize( + "sqlite_session", [(_retrieval_attachment_access_UploadFile, SegmentAttachmentBinding)], indirect=True +) +def test_segment_attachment_lookup_rejects_another_tenants_upload(sqlite_session: Session) -> None: + tenant_id = str(uuid4()) + foreign_tenant_id = str(uuid4()) + upload_file_id = str(uuid4()) + foreign_dataset_id = str(uuid4()) + upload_file = _retrieval_attachment_access_UploadFile( + tenant_id=foreign_tenant_id, + storage_type=StorageType.LOCAL, + key="uploads/foreign.png", + name="foreign.png", + size=1024, + extension="png", + mime_type="image/png", + created_by_role=CreatorUserRole.END_USER, + created_by=str(uuid4()), + created_at=datetime.now(UTC), + used=False, + ) + upload_file.id = upload_file_id + sqlite_session.add_all( + [ + upload_file, + SegmentAttachmentBinding( + tenant_id=foreign_tenant_id, + dataset_id=foreign_dataset_id, + document_id=str(uuid4()), + segment_id=str(uuid4()), + attachment_id=upload_file_id, + ), + ] + ) + sqlite_session.commit() + scope = FileAccessScope( + tenant_id=tenant_id, + user_id=str(uuid4()), + user_from=UserFrom.END_USER, + invoke_from=InvokeFrom.WEB_APP, + ) + + with bind_file_access_scope(scope): + result = RetrievalService.get_segment_attachment_infos( + [upload_file_id], + dataset_refs=((tenant_id, str(uuid4())),), + session=sqlite_session, + ) + current_scope = get_current_file_access_scope() + + assert result == [] + assert current_scope is not None + assert current_scope.granted_upload_file_ids == frozenset() + + +@pytest.mark.parametrize("sqlite_session", [(Dataset, DatasetDocument, DocumentSegment)], indirect=True) +def test_knowledge_retrieval_grants_returned_segments_to_current_scope( + monkeypatch: pytest.MonkeyPatch, + sqlite_session: Session, + sqlite_session_factory: sessionmaker[Session], +) -> None: + tenant_id = str(uuid4()) + dataset_id = str(uuid4()) + document_id = str(uuid4()) + segment_id = str(uuid4()) + created_by = str(uuid4()) + dataset = Dataset(tenant_id=tenant_id, name="Dataset", created_by=created_by) + dataset.id = dataset_id + document = DatasetDocument( + tenant_id=tenant_id, + dataset_id=dataset_id, + position=1, + data_source_type=DataSourceType.UPLOAD_FILE, + batch="batch-1", + name="Document", + created_from=DocumentCreatedFrom.API, + created_by=created_by, + doc_metadata={}, + ) + document.id = document_id + segment = DocumentSegment( + tenant_id=tenant_id, + dataset_id=dataset_id, + document_id=document_id, + position=1, + content="segment content", + word_count=10, + tokens=2, + created_by=created_by, + index_node_hash="hash", + hit_count=1, + ) + segment.id = segment_id + segment.status = SegmentStatus.COMPLETED + sqlite_session.add_all([dataset, document, segment]) + sqlite_session.commit() + + record = RetrievalSegments(segment=segment, score=0.8) + retrieval = DatasetRetrieval() + monkeypatch.setattr(retrieval, "_check_knowledge_rate_limit", lambda tenant_id: None) + monkeypatch.setattr(retrieval, "_get_available_datasets", lambda tenant_id, dataset_ids: [dataset]) + monkeypatch.setattr( + retrieval, + "multiple_retrieve", + lambda **kwargs: [RagDocument(page_content="segment content", provider="dify")], + ) + monkeypatch.setattr(RetrievalService, "format_retrieval_documents", lambda _session, documents: [record]) + monkeypatch.setattr(dataset_retrieval_module.session_factory, "create_session", sqlite_session_factory) + scope = FileAccessScope( + tenant_id=tenant_id, + user_id=str(uuid4()), + user_from=UserFrom.END_USER, + invoke_from=InvokeFrom.WEB_APP, + ) + + with bind_file_access_scope(scope): + results = retrieval.knowledge_retrieval( + sqlite_session, + KnowledgeRetrievalRequest( + tenant_id=tenant_id, + user_id=str(uuid4()), + app_id=str(uuid4()), + user_from=UserFrom.END_USER.value, + dataset_ids=[dataset_id], + query="desktop picture", + retrieval_mode="multiple", + ), + ) + current_scope = get_current_file_access_scope() + + assert results[0].metadata.segment_id == segment_id + assert current_scope is not None + assert segment_id in current_scope.granted_retriever_segment_ids diff --git a/api/tests/unit_tests/core/rag/datasource/vdb/test_vector_factory.py b/api/tests/unit_tests/core/rag/datasource/vdb/test_vector_factory.py index b69512daef2..5b84bd5fb49 100644 --- a/api/tests/unit_tests/core/rag/datasource/vdb/test_vector_factory.py +++ b/api/tests/unit_tests/core/rag/datasource/vdb/test_vector_factory.py @@ -16,7 +16,7 @@ from extensions.storage.storage_type import StorageType from models.dataset import Dataset, Whitelist from models.enums import CreatorUserRole from models.model import UploadFile -from tests.unit_tests.config_override import apply_config_overrides +from tests.helpers.unit.config_override import apply_config_overrides def _dataset(**overrides) -> Dataset: diff --git a/api/tests/unit_tests/core/rag/docstore/test_dataset_docstore.py b/api/tests/unit_tests/core/rag/docstore/test_dataset_docstore.py index e9b48cdff8b..be5182a6bad 100644 --- a/api/tests/unit_tests/core/rag/docstore/test_dataset_docstore.py +++ b/api/tests/unit_tests/core/rag/docstore/test_dataset_docstore.py @@ -12,7 +12,7 @@ from sqlalchemy.orm import Session from core.rag.docstore.dataset_docstore import DatasetDocumentStore, DocumentSegment from core.rag.models.document import AttachmentDocument, ChildDocument, Document from models.dataset import ChildChunk, Dataset, SegmentAttachmentBinding -from tests.unit_tests.model_factories import make_dataset +from tests.helpers.unit.model_factories import make_dataset TENANT_ID = "00000000-0000-0000-0000-000000000001" DATASET_ID = "00000000-0000-0000-0000-000000000002" diff --git a/api/tests/unit_tests/core/rag/embedding/test_cached_embedding.py b/api/tests/unit_tests/core/rag/embedding/test_cached_embedding.py index 909fbbc85ac..8748bccf54d 100644 --- a/api/tests/unit_tests/core/rag/embedding/test_cached_embedding.py +++ b/api/tests/unit_tests/core/rag/embedding/test_cached_embedding.py @@ -1,13 +1,6 @@ -"""Unit tests for cached_embedding.py - CacheEmbedding class. - -This test file covers the methods not fully tested in test_embedding_service.py: -- embed_multimodal_documents -- embed_multimodal_query -- Error handling scenarios in embed_query (DEBUG mode) -""" - import base64 import logging +import pickle from collections.abc import Callable from dataclasses import dataclass from decimal import Decimal @@ -16,29 +9,38 @@ from unittest.mock import Mock, patch import numpy as np import pytest from sqlalchemy import event, func, select +from sqlalchemy.exc import IntegrityError from sqlalchemy.orm import Session +from core.entities.embedding_type import EmbeddingInputType +from core.rag.embedding import cached_embedding from core.rag.embedding import cached_embedding as cached_embedding_module from core.rag.embedding.cached_embedding import CacheEmbedding from graphon.model_runtime.entities.model_entities import ModelPropertyKey from graphon.model_runtime.entities.text_embedding_entities import EmbeddingResult, EmbeddingUsage +from graphon.model_runtime.errors.invoke import ( + InvokeAuthorizationError, + InvokeConnectionError, + InvokeRateLimitError, +) +from libs import helper from models.dataset import Embedding +"""Unit tests for cached_embedding.py - CacheEmbedding class. + +This test file covers the methods not fully tested in test_embedding_service.py: +- embed_multimodal_documents +- embed_multimodal_query +- Error handling scenarios in embed_query (DEBUG mode) +""" + @dataclass(frozen=True) class _DatabaseBinding: session: Session -@pytest.fixture -def embedding_session(sqlite_session: Session, monkeypatch: pytest.MonkeyPatch) -> Session: - """Bind CacheEmbedding to the shared SQLite session.""" - - monkeypatch.setattr(cached_embedding_module, "db", _DatabaseBinding(session=sqlite_session)) - return sqlite_session - - -def _persist_embedding( +def _cached_embedding_persist_embedding( session: Session, *, cache_key: str, @@ -58,541 +60,853 @@ def _persist_embedding( return embedding -class TestCacheEmbeddingMultimodalDocuments: - """Test suite for CacheEmbedding.embed_multimodal_documents method.""" - +class TestCachedEmbeddingCases: + @staticmethod @pytest.fixture - def mock_model_instance(self) -> Mock: - """Create a mock ModelInstance for testing.""" - model_instance = Mock() - model_instance.model = "vision-embedding-model" - model_instance.model_name = "vision-embedding-model" - model_instance.provider = "openai" - model_instance.credentials = {"api_key": "test-key"} + def embedding_session(sqlite_session: Session, monkeypatch: pytest.MonkeyPatch) -> Session: + """Bind CacheEmbedding to the shared SQLite session.""" - model_type_instance = Mock() - model_instance.model_type_instance = model_type_instance + monkeypatch.setattr(cached_embedding_module, "db", _DatabaseBinding(session=sqlite_session)) + return sqlite_session - model_schema = Mock() - model_schema.model_properties = {ModelPropertyKey.MAX_CHUNKS: 10} - model_type_instance.get_model_schema.return_value = model_schema + class TestCacheEmbeddingMultimodalDocuments: + """Test suite for CacheEmbedding.embed_multimodal_documents method.""" - return model_instance + @pytest.fixture + def mock_model_instance(self) -> Mock: + """Create a mock ModelInstance for testing.""" + model_instance = Mock() + model_instance.model = "vision-embedding-model" + model_instance.model_name = "vision-embedding-model" + model_instance.provider = "openai" + model_instance.credentials = {"api_key": "test-key"} - @pytest.fixture - def sample_multimodal_result(self) -> EmbeddingResult: - """Create a sample multimodal EmbeddingResult.""" - embedding_vector = np.random.randn(1536) - normalized_vector = (embedding_vector / np.linalg.norm(embedding_vector)).tolist() + model_type_instance = Mock() + model_instance.model_type_instance = model_type_instance - usage = EmbeddingUsage( - tokens=10, - total_tokens=10, - unit_price=Decimal("0.0001"), - price_unit=Decimal(1000), - total_price=Decimal("0.000001"), - currency="USD", - latency=0.5, - ) + model_schema = Mock() + model_schema.model_properties = {ModelPropertyKey.MAX_CHUNKS: 10} + model_type_instance.get_model_schema.return_value = model_schema - return EmbeddingResult( - model="vision-embedding-model", - embeddings=[normalized_vector], - usage=usage, - ) + return model_instance - def test_embed_single_multimodal_document_cache_miss( - self, - mock_model_instance: Mock, - sample_multimodal_result: EmbeddingResult, - embedding_session: Session, - ) -> None: - """Test embedding a single multimodal document when cache is empty.""" - cache_embedding = CacheEmbedding(mock_model_instance) - documents = [{"file_id": "file123", "content": "test content"}] - - mock_model_instance.invoke_multimodal_embedding.return_value = sample_multimodal_result - - result = cache_embedding.embed_multimodal_documents(documents) - - assert len(result) == 1 - assert isinstance(result[0], list) - assert len(result[0]) == 1536 - - mock_model_instance.invoke_multimodal_embedding.assert_called_once() - persisted = embedding_session.scalar(select(Embedding).where(Embedding.hash == "file123")) - assert persisted is not None - assert persisted.get_embedding() == result[0] - - def test_embed_multiple_multimodal_documents_cache_miss( - self, mock_model_instance: Mock, embedding_session: Session - ) -> None: - """Test embedding multiple multimodal documents when cache is empty.""" - cache_embedding = CacheEmbedding(mock_model_instance) - documents = [ - {"file_id": "file1", "content": "content 1"}, - {"file_id": "file2", "content": "content 2"}, - {"file_id": "file3", "content": "content 3"}, - ] - - embeddings = [] - for _ in range(3): - vector = np.random.randn(1536) - normalized = (vector / np.linalg.norm(vector)).tolist() - embeddings.append(normalized) - - usage = EmbeddingUsage( - tokens=30, - total_tokens=30, - unit_price=Decimal("0.0001"), - price_unit=Decimal(1000), - total_price=Decimal("0.000003"), - currency="USD", - latency=0.8, - ) - - embedding_result = EmbeddingResult( - model="vision-embedding-model", - embeddings=embeddings, - usage=usage, - ) - - mock_model_instance.invoke_multimodal_embedding.return_value = embedding_result - - result = cache_embedding.embed_multimodal_documents(documents) - - assert len(result) == 3 - assert all(len(emb) == 1536 for emb in result) - assert embedding_session.scalar(select(func.count()).select_from(Embedding)) == 3 - - def test_embed_multimodal_documents_cache_hit(self, mock_model_instance: Mock, embedding_session: Session) -> None: - """Test embedding multimodal documents when embeddings are cached.""" - cache_embedding = CacheEmbedding(mock_model_instance) - documents = [{"file_id": "file123"}] - - cached_vector = np.random.randn(1536) - normalized_cached = (cached_vector / np.linalg.norm(cached_vector)).tolist() - - _persist_embedding(embedding_session, cache_key="file123", vector=normalized_cached) - - result = cache_embedding.embed_multimodal_documents(documents) - - assert len(result) == 1 - assert result[0] == normalized_cached - mock_model_instance.invoke_multimodal_embedding.assert_not_called() - - def test_embed_multimodal_documents_partial_cache_hit( - self, mock_model_instance: Mock, embedding_session: Session - ) -> None: - """Test embedding multimodal documents with mixed cache hits and misses.""" - cache_embedding = CacheEmbedding(mock_model_instance) - documents = [ - {"file_id": "cached_file"}, - {"file_id": "new_file_1"}, - {"file_id": "new_file_2"}, - ] - - cached_vector = np.random.randn(1536) - normalized_cached = (cached_vector / np.linalg.norm(cached_vector)).tolist() - - _persist_embedding(embedding_session, cache_key="cached_file", vector=normalized_cached) - - new_embeddings = [] - for _ in range(2): - vector = np.random.randn(1536) - normalized = (vector / np.linalg.norm(vector)).tolist() - new_embeddings.append(normalized) - - usage = EmbeddingUsage( - tokens=20, - total_tokens=20, - unit_price=Decimal("0.0001"), - price_unit=Decimal(1000), - total_price=Decimal("0.000002"), - currency="USD", - latency=0.6, - ) - - embedding_result = EmbeddingResult( - model="vision-embedding-model", - embeddings=new_embeddings, - usage=usage, - ) - - mock_model_instance.invoke_multimodal_embedding.return_value = embedding_result - - result = cache_embedding.embed_multimodal_documents(documents) - - assert len(result) == 3 - assert result[0] == normalized_cached - assert embedding_session.scalar(select(func.count()).select_from(Embedding)) == 3 - - def test_embed_multimodal_documents_nan_handling( - self, - mock_model_instance: Mock, - embedding_session: Session, - caplog: pytest.LogCaptureFixture, - ) -> None: - """Test handling of NaN values in multimodal embeddings.""" - cache_embedding = CacheEmbedding(mock_model_instance) - documents = [{"file_id": "valid"}, {"file_id": "nan"}] - - valid_vector = np.random.randn(1536).tolist() - nan_vector = [float("nan")] * 1536 - - usage = EmbeddingUsage( - tokens=20, - total_tokens=20, - unit_price=Decimal("0.0001"), - price_unit=Decimal(1000), - total_price=Decimal("0.000002"), - currency="USD", - latency=0.5, - ) - - embedding_result = EmbeddingResult( - model="vision-embedding-model", - embeddings=[valid_vector, nan_vector], - usage=usage, - ) - - mock_model_instance.invoke_multimodal_embedding.return_value = embedding_result - - with caplog.at_level(logging.WARNING, logger="core.rag.embedding.cached_embedding"): - result = cache_embedding.embed_multimodal_documents(documents) - - assert len(result) == 2 - assert result[0] is not None - assert result[1] is None - assert embedding_session.scalar(select(func.count()).select_from(Embedding)) == 1 - - assert any(record.levelno == logging.WARNING for record in caplog.records) - - def test_embed_multimodal_documents_large_batch( - self, mock_model_instance: Mock, embedding_session: Session - ) -> None: - """Test embedding large batch of multimodal documents respecting MAX_CHUNKS.""" - cache_embedding = CacheEmbedding(mock_model_instance) - documents = [{"file_id": f"file{i}"} for i in range(25)] - - def create_batch_result(batch_size: int) -> EmbeddingResult: - embeddings = [] - for _ in range(batch_size): - vector = np.random.randn(1536) - normalized = (vector / np.linalg.norm(vector)).tolist() - embeddings.append(normalized) + @pytest.fixture + def sample_multimodal_result(self) -> EmbeddingResult: + """Create a sample multimodal EmbeddingResult.""" + embedding_vector = np.random.randn(1536) + normalized_vector = (embedding_vector / np.linalg.norm(embedding_vector)).tolist() usage = EmbeddingUsage( - tokens=batch_size * 10, - total_tokens=batch_size * 10, + tokens=10, + total_tokens=10, unit_price=Decimal("0.0001"), price_unit=Decimal(1000), - total_price=Decimal(str(batch_size * 0.000001)), + total_price=Decimal("0.000001"), currency="USD", latency=0.5, ) return EmbeddingResult( + model="vision-embedding-model", + embeddings=[normalized_vector], + usage=usage, + ) + + def test_embed_single_multimodal_document_cache_miss( + self, + mock_model_instance: Mock, + sample_multimodal_result: EmbeddingResult, + embedding_session: Session, + ) -> None: + """Test embedding a single multimodal document when cache is empty.""" + cache_embedding = CacheEmbedding(mock_model_instance) + documents = [{"file_id": "file123", "content": "test content"}] + + mock_model_instance.invoke_multimodal_embedding.return_value = sample_multimodal_result + + result = cache_embedding.embed_multimodal_documents(documents) + + assert len(result) == 1 + assert isinstance(result[0], list) + assert len(result[0]) == 1536 + + mock_model_instance.invoke_multimodal_embedding.assert_called_once() + persisted = embedding_session.scalar(select(Embedding).where(Embedding.hash == "file123")) + assert persisted is not None + assert persisted.get_embedding() == result[0] + + def test_embed_multiple_multimodal_documents_cache_miss( + self, mock_model_instance: Mock, embedding_session: Session + ) -> None: + """Test embedding multiple multimodal documents when cache is empty.""" + cache_embedding = CacheEmbedding(mock_model_instance) + documents = [ + {"file_id": "file1", "content": "content 1"}, + {"file_id": "file2", "content": "content 2"}, + {"file_id": "file3", "content": "content 3"}, + ] + + embeddings = [] + for _ in range(3): + vector = np.random.randn(1536) + normalized = (vector / np.linalg.norm(vector)).tolist() + embeddings.append(normalized) + + usage = EmbeddingUsage( + tokens=30, + total_tokens=30, + unit_price=Decimal("0.0001"), + price_unit=Decimal(1000), + total_price=Decimal("0.000003"), + currency="USD", + latency=0.8, + ) + + embedding_result = EmbeddingResult( model="vision-embedding-model", embeddings=embeddings, usage=usage, ) - batch_results = [create_batch_result(10), create_batch_result(10), create_batch_result(5)] - mock_model_instance.invoke_multimodal_embedding.side_effect = batch_results - - result = cache_embedding.embed_multimodal_documents(documents) - - assert len(result) == 25 - assert mock_model_instance.invoke_multimodal_embedding.call_count == 3 - assert embedding_session.scalar(select(func.count()).select_from(Embedding)) == 25 - - def test_embed_multimodal_documents_api_error(self, mock_model_instance: Mock, embedding_session: Session) -> None: - """Test handling of API errors during multimodal embedding.""" - cache_embedding = CacheEmbedding(mock_model_instance) - documents = [{"file_id": "file123"}] - - mock_model_instance.invoke_multimodal_embedding.side_effect = Exception("API Error") - - with pytest.raises(Exception, match="API Error"): - cache_embedding.embed_multimodal_documents(documents) - - assert not embedding_session.in_transaction() - assert embedding_session.scalar(select(func.count()).select_from(Embedding)) == 0 - - def test_embed_multimodal_documents_integrity_error_during_transform( - self, - mock_model_instance: Mock, - sample_multimodal_result: EmbeddingResult, - embedding_session: Session, - ) -> None: - """Test handling of IntegrityError during embedding transformation.""" - cache_embedding = CacheEmbedding(mock_model_instance) - documents = [{"file_id": "file123"}] - - mock_model_instance.invoke_multimodal_embedding.return_value = sample_multimodal_result - - injected = False - - def add_competing_row(session: Session, _flush_context: object, _instances: object) -> None: - nonlocal injected - if injected: - return - pending = next(item for item in session.new if isinstance(item, Embedding)) - session.add( - Embedding( - model_name=pending.model_name, - hash=pending.hash, - provider_name=pending.provider_name, - embedding=pending.embedding, - ) - ) - injected = True - - event.listen(embedding_session, "before_flush", add_competing_row) - try: - result = cache_embedding.embed_multimodal_documents(documents) - finally: - event.remove(embedding_session, "before_flush", add_competing_row) - - assert len(result) == 1 - assert not embedding_session.in_transaction() - assert embedding_session.scalar(select(func.count()).select_from(Embedding)) == 0 - - -class TestCacheEmbeddingMultimodalQuery: - """Test suite for CacheEmbedding.embed_multimodal_query method.""" - - @pytest.fixture - def mock_model_instance(self) -> Mock: - """Create a mock ModelInstance for testing.""" - model_instance = Mock() - model_instance.model = "vision-embedding-model" - model_instance.model_name = "vision-embedding-model" - model_instance.provider = "openai" - model_instance.credentials = {"api_key": "test-key"} - return model_instance - - def test_embed_multimodal_query_cache_miss(self, mock_model_instance: Mock) -> None: - """Test embedding multimodal query when Redis cache is empty.""" - cache_embedding = CacheEmbedding(mock_model_instance) - document = {"file_id": "file123"} - - vector = np.random.randn(1536) - normalized = (vector / np.linalg.norm(vector)).tolist() - - usage = EmbeddingUsage( - tokens=5, - total_tokens=5, - unit_price=Decimal("0.0001"), - price_unit=Decimal(1000), - total_price=Decimal("0.0000005"), - currency="USD", - latency=0.3, - ) - - embedding_result = EmbeddingResult( - model="vision-embedding-model", - embeddings=[normalized], - usage=usage, - ) - - with patch("core.rag.embedding.cached_embedding.redis_client") as mock_redis: - mock_redis.get.return_value = None mock_model_instance.invoke_multimodal_embedding.return_value = embedding_result - result = cache_embedding.embed_multimodal_query(document) + result = cache_embedding.embed_multimodal_documents(documents) - assert isinstance(result, list) - assert len(result) == 1536 - mock_redis.setex.assert_called_once() + assert len(result) == 3 + assert all(len(emb) == 1536 for emb in result) + assert embedding_session.scalar(select(func.count()).select_from(Embedding)) == 3 - def test_embed_multimodal_query_cache_hit(self, mock_model_instance: Mock) -> None: - """Test embedding multimodal query when Redis cache has the value.""" - cache_embedding = CacheEmbedding(mock_model_instance) - document = {"file_id": "file123"} + def test_embed_multimodal_documents_cache_hit( + self, mock_model_instance: Mock, embedding_session: Session + ) -> None: + """Test embedding multimodal documents when embeddings are cached.""" + cache_embedding = CacheEmbedding(mock_model_instance) + documents = [{"file_id": "file123"}] - embedding_vector = np.random.randn(1536) - vector_bytes = embedding_vector.tobytes() - encoded_vector = base64.b64encode(vector_bytes).decode("utf-8") + cached_vector = np.random.randn(1536) + normalized_cached = (cached_vector / np.linalg.norm(cached_vector)).tolist() - with patch("core.rag.embedding.cached_embedding.redis_client") as mock_redis: - mock_redis.get.return_value = encoded_vector.encode() + _cached_embedding_persist_embedding(embedding_session, cache_key="file123", vector=normalized_cached) - result = cache_embedding.embed_multimodal_query(document) + result = cache_embedding.embed_multimodal_documents(documents) - assert isinstance(result, list) - assert len(result) == 1536 - mock_redis.expire.assert_called_once() + assert len(result) == 1 + assert result[0] == normalized_cached mock_model_instance.invoke_multimodal_embedding.assert_not_called() - def test_embed_multimodal_query_nan_handling(self, mock_model_instance: Mock) -> None: - """Test handling of NaN values in multimodal query embeddings.""" - cache_embedding = CacheEmbedding(mock_model_instance) + def test_embed_multimodal_documents_partial_cache_hit( + self, mock_model_instance: Mock, embedding_session: Session + ) -> None: + """Test embedding multimodal documents with mixed cache hits and misses.""" + cache_embedding = CacheEmbedding(mock_model_instance) + documents = [ + {"file_id": "cached_file"}, + {"file_id": "new_file_1"}, + {"file_id": "new_file_2"}, + ] - nan_vector = [float("nan")] * 1536 + cached_vector = np.random.randn(1536) + normalized_cached = (cached_vector / np.linalg.norm(cached_vector)).tolist() - usage = EmbeddingUsage( - tokens=5, - total_tokens=5, - unit_price=Decimal("0.0001"), - price_unit=Decimal(1000), - total_price=Decimal("0.0000005"), - currency="USD", - latency=0.3, - ) + _cached_embedding_persist_embedding(embedding_session, cache_key="cached_file", vector=normalized_cached) - embedding_result = EmbeddingResult( - model="vision-embedding-model", - embeddings=[nan_vector], - usage=usage, - ) + new_embeddings = [] + for _ in range(2): + vector = np.random.randn(1536) + normalized = (vector / np.linalg.norm(vector)).tolist() + new_embeddings.append(normalized) - document = {"file_id": "file123"} + usage = EmbeddingUsage( + tokens=20, + total_tokens=20, + unit_price=Decimal("0.0001"), + price_unit=Decimal(1000), + total_price=Decimal("0.000002"), + currency="USD", + latency=0.6, + ) + + embedding_result = EmbeddingResult( + model="vision-embedding-model", + embeddings=new_embeddings, + usage=usage, + ) - with patch("core.rag.embedding.cached_embedding.redis_client") as mock_redis: - mock_redis.get.return_value = None mock_model_instance.invoke_multimodal_embedding.return_value = embedding_result - with pytest.raises(ValueError, match="Normalized embedding is nan please try again") as exc_info: - cache_embedding.embed_multimodal_query(document) + result = cache_embedding.embed_multimodal_documents(documents) - assert "Normalized embedding is nan" in str(exc_info.value) + assert len(result) == 3 + assert result[0] == normalized_cached + assert embedding_session.scalar(select(func.count()).select_from(Embedding)) == 3 - def test_embed_multimodal_query_api_error( - self, mock_model_instance: Mock, config_overrides: Callable[..., None] - ) -> None: - """Test handling of API errors during multimodal query embedding.""" - config_overrides(DEBUG=False) - cache_embedding = CacheEmbedding(mock_model_instance) - document = {"file_id": "file123"} + def test_embed_multimodal_documents_nan_handling( + self, + mock_model_instance: Mock, + embedding_session: Session, + caplog: pytest.LogCaptureFixture, + ) -> None: + """Test handling of NaN values in multimodal embeddings.""" + cache_embedding = CacheEmbedding(mock_model_instance) + documents = [{"file_id": "valid"}, {"file_id": "nan"}] + + valid_vector = np.random.randn(1536).tolist() + nan_vector = [float("nan")] * 1536 + + usage = EmbeddingUsage( + tokens=20, + total_tokens=20, + unit_price=Decimal("0.0001"), + price_unit=Decimal(1000), + total_price=Decimal("0.000002"), + currency="USD", + latency=0.5, + ) + + embedding_result = EmbeddingResult( + model="vision-embedding-model", + embeddings=[valid_vector, nan_vector], + usage=usage, + ) + + mock_model_instance.invoke_multimodal_embedding.return_value = embedding_result + + with caplog.at_level(logging.WARNING, logger="core.rag.embedding.cached_embedding"): + result = cache_embedding.embed_multimodal_documents(documents) + + assert len(result) == 2 + assert result[0] is not None + assert result[1] is None + assert embedding_session.scalar(select(func.count()).select_from(Embedding)) == 1 + + assert any(record.levelno == logging.WARNING for record in caplog.records) + + def test_embed_multimodal_documents_large_batch( + self, mock_model_instance: Mock, embedding_session: Session + ) -> None: + """Test embedding large batch of multimodal documents respecting MAX_CHUNKS.""" + cache_embedding = CacheEmbedding(mock_model_instance) + documents = [{"file_id": f"file{i}"} for i in range(25)] + + def create_batch_result(batch_size: int) -> EmbeddingResult: + embeddings = [] + for _ in range(batch_size): + vector = np.random.randn(1536) + normalized = (vector / np.linalg.norm(vector)).tolist() + embeddings.append(normalized) + + usage = EmbeddingUsage( + tokens=batch_size * 10, + total_tokens=batch_size * 10, + unit_price=Decimal("0.0001"), + price_unit=Decimal(1000), + total_price=Decimal(str(batch_size * 0.000001)), + currency="USD", + latency=0.5, + ) + + return EmbeddingResult( + model="vision-embedding-model", + embeddings=embeddings, + usage=usage, + ) + + batch_results = [create_batch_result(10), create_batch_result(10), create_batch_result(5)] + mock_model_instance.invoke_multimodal_embedding.side_effect = batch_results + + result = cache_embedding.embed_multimodal_documents(documents) + + assert len(result) == 25 + assert mock_model_instance.invoke_multimodal_embedding.call_count == 3 + assert embedding_session.scalar(select(func.count()).select_from(Embedding)) == 25 + + def test_embed_multimodal_documents_api_error( + self, mock_model_instance: Mock, embedding_session: Session + ) -> None: + """Test handling of API errors during multimodal embedding.""" + cache_embedding = CacheEmbedding(mock_model_instance) + documents = [{"file_id": "file123"}] - with patch("core.rag.embedding.cached_embedding.redis_client") as mock_redis: - mock_redis.get.return_value = None mock_model_instance.invoke_multimodal_embedding.side_effect = Exception("API Error") - with pytest.raises(Exception, match="API Error") as exc_info: - cache_embedding.embed_multimodal_query(document) + with pytest.raises(Exception, match="API Error"): + cache_embedding.embed_multimodal_documents(documents) - assert "API Error" in str(exc_info.value) + assert not embedding_session.in_transaction() + assert embedding_session.scalar(select(func.count()).select_from(Embedding)) == 0 - def test_embed_multimodal_query_redis_set_error( - self, mock_model_instance: Mock, config_overrides: Callable[..., None] - ) -> None: - """Test handling of Redis set errors during multimodal query embedding.""" - config_overrides(DEBUG=True) - cache_embedding = CacheEmbedding(mock_model_instance) - document = {"file_id": "file123"} + def test_embed_multimodal_documents_integrity_error_during_transform( + self, + mock_model_instance: Mock, + sample_multimodal_result: EmbeddingResult, + embedding_session: Session, + ) -> None: + """Test handling of IntegrityError during embedding transformation.""" + cache_embedding = CacheEmbedding(mock_model_instance) + documents = [{"file_id": "file123"}] - vector = np.random.randn(1536) - normalized = (vector / np.linalg.norm(vector)).tolist() + mock_model_instance.invoke_multimodal_embedding.return_value = sample_multimodal_result - usage = EmbeddingUsage( - tokens=5, - total_tokens=5, - unit_price=Decimal("0.0001"), - price_unit=Decimal(1000), - total_price=Decimal("0.0000005"), - currency="USD", - latency=0.3, - ) + injected = False - embedding_result = EmbeddingResult( - model="vision-embedding-model", - embeddings=[normalized], - usage=usage, - ) + def add_competing_row(session: Session, _flush_context: object, _instances: object) -> None: + nonlocal injected + if injected: + return + pending = next(item for item in session.new if isinstance(item, Embedding)) + session.add( + Embedding( + model_name=pending.model_name, + hash=pending.hash, + provider_name=pending.provider_name, + embedding=pending.embedding, + ) + ) + injected = True - with patch("core.rag.embedding.cached_embedding.redis_client") as mock_redis: - mock_redis.get.return_value = None - mock_model_instance.invoke_multimodal_embedding.return_value = embedding_result - mock_redis.setex.side_effect = RuntimeError("Redis Error") + event.listen(embedding_session, "before_flush", add_competing_row) + try: + result = cache_embedding.embed_multimodal_documents(documents) + finally: + event.remove(embedding_session, "before_flush", add_competing_row) - with pytest.raises(RuntimeError): - cache_embedding.embed_multimodal_query(document) + assert len(result) == 1 + assert not embedding_session.in_transaction() + assert embedding_session.scalar(select(func.count()).select_from(Embedding)) == 0 + class TestCacheEmbeddingMultimodalQuery: + """Test suite for CacheEmbedding.embed_multimodal_query method.""" -class TestCacheEmbeddingQueryErrors: - """Test suite for error handling in CacheEmbedding.embed_query method.""" + @pytest.fixture + def mock_model_instance(self) -> Mock: + """Create a mock ModelInstance for testing.""" + model_instance = Mock() + model_instance.model = "vision-embedding-model" + model_instance.model_name = "vision-embedding-model" + model_instance.provider = "openai" + model_instance.credentials = {"api_key": "test-key"} + return model_instance - @pytest.fixture - def mock_model_instance(self) -> Mock: - """Create a mock ModelInstance for testing.""" - model_instance = Mock() - model_instance.model = "text-embedding-ada-002" - model_instance.model_name = "text-embedding-ada-002" - model_instance.provider = "openai" - model_instance.credentials = {"api_key": "test-key"} - return model_instance + def test_embed_multimodal_query_cache_miss(self, mock_model_instance: Mock) -> None: + """Test embedding multimodal query when Redis cache is empty.""" + cache_embedding = CacheEmbedding(mock_model_instance) + document = {"file_id": "file123"} - def test_embed_query_api_error_debug_mode( - self, mock_model_instance: Mock, caplog: pytest.LogCaptureFixture, config_overrides: Callable[..., None] - ) -> None: - """Test handling of API errors in debug mode.""" - config_overrides(DEBUG=True) - cache_embedding = CacheEmbedding(mock_model_instance) - query = "test query" + vector = np.random.randn(1536) + normalized = (vector / np.linalg.norm(vector)).tolist() - with patch("core.rag.embedding.cached_embedding.redis_client") as mock_redis: - mock_redis.get.return_value = None - mock_model_instance.invoke_text_embedding.side_effect = RuntimeError("API Error") + usage = EmbeddingUsage( + tokens=5, + total_tokens=5, + unit_price=Decimal("0.0001"), + price_unit=Decimal(1000), + total_price=Decimal("0.0000005"), + currency="USD", + latency=0.3, + ) - with caplog.at_level(logging.ERROR, logger="core.rag.embedding.cached_embedding"): - with pytest.raises(RuntimeError) as exc_info: - cache_embedding.embed_query(query) + embedding_result = EmbeddingResult( + model="vision-embedding-model", + embeddings=[normalized], + usage=usage, + ) + + with patch("core.rag.embedding.cached_embedding.redis_client") as mock_redis: + mock_redis.get.return_value = None + mock_model_instance.invoke_multimodal_embedding.return_value = embedding_result + + result = cache_embedding.embed_multimodal_query(document) + + assert isinstance(result, list) + assert len(result) == 1536 + mock_redis.setex.assert_called_once() + + def test_embed_multimodal_query_cache_hit(self, mock_model_instance: Mock) -> None: + """Test embedding multimodal query when Redis cache has the value.""" + cache_embedding = CacheEmbedding(mock_model_instance) + document = {"file_id": "file123"} + + embedding_vector = np.random.randn(1536) + vector_bytes = embedding_vector.tobytes() + encoded_vector = base64.b64encode(vector_bytes).decode("utf-8") + + with patch("core.rag.embedding.cached_embedding.redis_client") as mock_redis: + mock_redis.get.return_value = encoded_vector.encode() + + result = cache_embedding.embed_multimodal_query(document) + + assert isinstance(result, list) + assert len(result) == 1536 + mock_redis.expire.assert_called_once() + mock_model_instance.invoke_multimodal_embedding.assert_not_called() + + def test_embed_multimodal_query_nan_handling(self, mock_model_instance: Mock) -> None: + """Test handling of NaN values in multimodal query embeddings.""" + cache_embedding = CacheEmbedding(mock_model_instance) + + nan_vector = [float("nan")] * 1536 + + usage = EmbeddingUsage( + tokens=5, + total_tokens=5, + unit_price=Decimal("0.0001"), + price_unit=Decimal(1000), + total_price=Decimal("0.0000005"), + currency="USD", + latency=0.3, + ) + + embedding_result = EmbeddingResult( + model="vision-embedding-model", + embeddings=[nan_vector], + usage=usage, + ) + + document = {"file_id": "file123"} + + with patch("core.rag.embedding.cached_embedding.redis_client") as mock_redis: + mock_redis.get.return_value = None + mock_model_instance.invoke_multimodal_embedding.return_value = embedding_result + + with pytest.raises(ValueError, match="Normalized embedding is nan please try again") as exc_info: + cache_embedding.embed_multimodal_query(document) + + assert "Normalized embedding is nan" in str(exc_info.value) + + def test_embed_multimodal_query_api_error( + self, mock_model_instance: Mock, config_overrides: Callable[..., None] + ) -> None: + """Test handling of API errors during multimodal query embedding.""" + config_overrides(DEBUG=False) + cache_embedding = CacheEmbedding(mock_model_instance) + document = {"file_id": "file123"} + + with patch("core.rag.embedding.cached_embedding.redis_client") as mock_redis: + mock_redis.get.return_value = None + mock_model_instance.invoke_multimodal_embedding.side_effect = Exception("API Error") + + with pytest.raises(Exception, match="API Error") as exc_info: + cache_embedding.embed_multimodal_query(document) assert "API Error" in str(exc_info.value) - assert any(record.levelno == logging.ERROR for record in caplog.records) - def test_embed_query_redis_set_error_debug_mode( - self, mock_model_instance: Mock, caplog: pytest.LogCaptureFixture, config_overrides: Callable[..., None] - ) -> None: - """Test handling of Redis set errors in debug mode.""" - config_overrides(DEBUG=True) - cache_embedding = CacheEmbedding(mock_model_instance) - query = "test query" + def test_embed_multimodal_query_redis_set_error( + self, mock_model_instance: Mock, config_overrides: Callable[..., None] + ) -> None: + """Test handling of Redis set errors during multimodal query embedding.""" + config_overrides(DEBUG=True) + cache_embedding = CacheEmbedding(mock_model_instance) + document = {"file_id": "file123"} - vector = np.random.randn(1536) - normalized = (vector / np.linalg.norm(vector)).tolist() + vector = np.random.randn(1536) + normalized = (vector / np.linalg.norm(vector)).tolist() - usage = EmbeddingUsage( - tokens=5, - total_tokens=5, + usage = EmbeddingUsage( + tokens=5, + total_tokens=5, + unit_price=Decimal("0.0001"), + price_unit=Decimal(1000), + total_price=Decimal("0.0000005"), + currency="USD", + latency=0.3, + ) + + embedding_result = EmbeddingResult( + model="vision-embedding-model", + embeddings=[normalized], + usage=usage, + ) + + with patch("core.rag.embedding.cached_embedding.redis_client") as mock_redis: + mock_redis.get.return_value = None + mock_model_instance.invoke_multimodal_embedding.return_value = embedding_result + mock_redis.setex.side_effect = RuntimeError("Redis Error") + + with pytest.raises(RuntimeError): + cache_embedding.embed_multimodal_query(document) + + class TestCacheEmbeddingQueryErrors: + """Test suite for error handling in CacheEmbedding.embed_query method.""" + + @pytest.fixture + def mock_model_instance(self) -> Mock: + """Create a mock ModelInstance for testing.""" + model_instance = Mock() + model_instance.model = "text-embedding-ada-002" + model_instance.model_name = "text-embedding-ada-002" + model_instance.provider = "openai" + model_instance.credentials = {"api_key": "test-key"} + return model_instance + + def test_embed_query_api_error_debug_mode( + self, mock_model_instance: Mock, caplog: pytest.LogCaptureFixture, config_overrides: Callable[..., None] + ) -> None: + """Test handling of API errors in debug mode.""" + config_overrides(DEBUG=True) + cache_embedding = CacheEmbedding(mock_model_instance) + query = "test query" + + with patch("core.rag.embedding.cached_embedding.redis_client") as mock_redis: + mock_redis.get.return_value = None + mock_model_instance.invoke_text_embedding.side_effect = RuntimeError("API Error") + + with caplog.at_level(logging.ERROR, logger="core.rag.embedding.cached_embedding"): + with pytest.raises(RuntimeError) as exc_info: + cache_embedding.embed_query(query) + + assert "API Error" in str(exc_info.value) + assert any(record.levelno == logging.ERROR for record in caplog.records) + + def test_embed_query_redis_set_error_debug_mode( + self, mock_model_instance: Mock, caplog: pytest.LogCaptureFixture, config_overrides: Callable[..., None] + ) -> None: + """Test handling of Redis set errors in debug mode.""" + config_overrides(DEBUG=True) + cache_embedding = CacheEmbedding(mock_model_instance) + query = "test query" + + vector = np.random.randn(1536) + normalized = (vector / np.linalg.norm(vector)).tolist() + + usage = EmbeddingUsage( + tokens=5, + total_tokens=5, + unit_price=Decimal("0.0001"), + price_unit=Decimal(1000), + total_price=Decimal("0.0000005"), + currency="USD", + latency=0.3, + ) + + embedding_result = EmbeddingResult( + model="text-embedding-ada-002", + embeddings=[normalized], + usage=usage, + ) + + with patch("core.rag.embedding.cached_embedding.redis_client") as mock_redis: + mock_redis.get.return_value = None + mock_model_instance.invoke_text_embedding.return_value = embedding_result + mock_redis.setex.side_effect = RuntimeError("Redis Error") + + with caplog.at_level(logging.ERROR, logger="core.rag.embedding.cached_embedding"): + with pytest.raises(RuntimeError): + cache_embedding.embed_query(query) + + assert any(record.levelno == logging.ERROR for record in caplog.records) + + class TestCacheEmbeddingInitialization: + """Test suite for CacheEmbedding initialization.""" + + def test_initialization_sets_model_instance(self) -> None: + """Test CacheEmbedding initialization stores the provided model instance.""" + model_instance = Mock() + model_instance.model = "test-model" + model_instance.model_name = "test-model" + model_instance.provider = "test-provider" + + cache_embedding = CacheEmbedding(model_instance) + + assert cache_embedding._model_instance == model_instance + + +"""Behavior-focused tests for :class:`CacheEmbedding`. + +Document and multimodal caches are persisted in SQLite so cache hits, misses, +provider isolation, uniqueness races, and rollback behavior exercise real +SQLAlchemy statements. Model providers and Redis remain mocked external I/O. +""" + + +@dataclass(frozen=True) +class _Database: + """Expose a real session through the interface used by the cache.""" + + session: Session + + +def _vector(dimension: int = 8, *, offset: float = 0) -> np.ndarray: + return np.arange(1 + offset, dimension + 1 + offset, dtype=float) + + +def _result(*vectors: np.ndarray) -> EmbeddingResult: + return EmbeddingResult( + model="text-embedding-ada-002", + embeddings=list(vectors), + usage=EmbeddingUsage( + tokens=len(vectors), + total_tokens=len(vectors), unit_price=Decimal("0.0001"), price_unit=Decimal(1000), - total_price=Decimal("0.0000005"), + total_price=Decimal("0.000001"), currency="USD", - latency=0.3, + latency=0.1, + ), + ) + + +def _embedding_service_persist_embedding( + session: Session, + *, + text: str, + vector: list[float], + model_name: str = "text-embedding-ada-002", + provider_name: str = "openai", +) -> Embedding: + row = Embedding( + model_name=model_name, + hash=helper.generate_text_hash(text), + provider_name=provider_name, + embedding=pickle.dumps(vector, protocol=pickle.HIGHEST_PROTOCOL), + ) + session.add(row) + session.commit() + return row + + +class TestEmbeddingServiceCases: + @staticmethod + @pytest.fixture + def sqlite_embedding_session(sqlite_session: Session, monkeypatch: pytest.MonkeyPatch) -> Session: + """Bind the shared SQLite session to the document-cache implementation.""" + + monkeypatch.setattr(cached_embedding, "db", _Database(sqlite_session)) + return sqlite_session + + @staticmethod + @pytest.fixture + def model_instance() -> Mock: + instance = Mock() + instance.model_name = "text-embedding-ada-002" + instance.provider = "openai" + instance.credentials = {"api_key": "test-key"} + schema = Mock(model_properties={ModelPropertyKey.MAX_CHUNKS: 10}) + instance.model_type_instance.get_model_schema.return_value = schema + return instance + + class TestDocumentCache: + def test_cache_miss_normalizes_and_persists( + self, sqlite_embedding_session: Session, model_instance: Mock + ) -> None: + model_instance.invoke_text_embedding.return_value = _result(_vector()) + + result = CacheEmbedding(model_instance).embed_documents(["new text"]) + + assert np.linalg.norm(result[0]) == pytest.approx(1.0) + model_instance.invoke_text_embedding.assert_called_once_with( + texts=["new text"], input_type=EmbeddingInputType.DOCUMENT + ) + persisted = sqlite_embedding_session.scalar(select(Embedding)) + assert persisted is not None + assert persisted.get_embedding() == result[0] + + def test_cache_hit_uses_persisted_vector(self, sqlite_embedding_session: Session, model_instance: Mock) -> None: + cached_vector = (_vector() / np.linalg.norm(_vector())).tolist() + row = _embedding_service_persist_embedding(sqlite_embedding_session, text="cached", vector=cached_vector) + + result = CacheEmbedding(model_instance).embed_documents(["cached"]) + + assert result == [cached_vector] + model_instance.invoke_text_embedding.assert_not_called() + assert sqlite_embedding_session.scalars(select(Embedding)).all() == [row] + + def test_partial_hit_preserves_input_order( + self, sqlite_embedding_session: Session, model_instance: Mock + ) -> None: + cached_vector = (_vector(offset=10) / np.linalg.norm(_vector(offset=10))).tolist() + _embedding_service_persist_embedding(sqlite_embedding_session, text="cached", vector=cached_vector) + model_instance.invoke_text_embedding.return_value = _result(_vector(), _vector(offset=20)) + + result = CacheEmbedding(model_instance).embed_documents(["cached", "new one", "new two"]) + + assert result[0] == cached_vector + assert all(np.linalg.norm(vector) == pytest.approx(1.0) for vector in result) + assert model_instance.invoke_text_embedding.call_args.kwargs["texts"] == ["new one", "new two"] + assert len(sqlite_embedding_session.scalars(select(Embedding)).all()) == 3 + + def test_large_batch_respects_model_chunk_limit( + self, sqlite_embedding_session: Session, model_instance: Mock + ) -> None: + texts = [f"text-{index}" for index in range(25)] + model_instance.invoke_text_embedding.side_effect = [ + _result(*(_vector(offset=index) for index in range(10))), + _result(*(_vector(offset=index) for index in range(10, 20))), + _result(*(_vector(offset=index) for index in range(20, 25))), + ] + + result = CacheEmbedding(model_instance).embed_documents(texts) + + assert len(result) == 25 + assert [len(call.kwargs["texts"]) for call in model_instance.invoke_text_embedding.call_args_list] == [ + 10, + 10, + 5, + ] + assert len(sqlite_embedding_session.scalars(select(Embedding)).all()) == 25 + + @pytest.mark.parametrize( + "provider_error", + [InvokeConnectionError("offline"), InvokeRateLimitError("limited"), InvokeAuthorizationError("denied")], ) + def test_provider_error_rolls_back_real_transaction( + self, + sqlite_embedding_session: Session, + model_instance: Mock, + provider_error: Exception, + ) -> None: + model_instance.invoke_text_embedding.side_effect = provider_error - embedding_result = EmbeddingResult( - model="text-embedding-ada-002", - embeddings=[normalized], - usage=usage, - ) + with pytest.raises(type(provider_error)): + CacheEmbedding(model_instance).embed_documents(["failure"]) - with patch("core.rag.embedding.cached_embedding.redis_client") as mock_redis: - mock_redis.get.return_value = None - mock_model_instance.invoke_text_embedding.return_value = embedding_result - mock_redis.setex.side_effect = RuntimeError("Redis Error") + assert not sqlite_embedding_session.in_transaction() + assert sqlite_embedding_session.scalar(select(Embedding)) is None - with caplog.at_level(logging.ERROR, logger="core.rag.embedding.cached_embedding"): - with pytest.raises(RuntimeError): - cache_embedding.embed_query(query) + def test_integrity_race_rolls_back_without_losing_result( + self, sqlite_embedding_session: Session, model_instance: Mock + ) -> None: + model_instance.invoke_text_embedding.return_value = _result(_vector()) - assert any(record.levelno == logging.ERROR for record in caplog.records) + def raise_integrity_error(_session: Session) -> None: + raise IntegrityError("duplicate embedding", {}, RuntimeError("unique constraint")) + event.listen(sqlite_embedding_session, "before_commit", raise_integrity_error) + try: + result = CacheEmbedding(model_instance).embed_documents(["racing insert"]) + finally: + event.remove(sqlite_embedding_session, "before_commit", raise_integrity_error) -class TestCacheEmbeddingInitialization: - """Test suite for CacheEmbedding initialization.""" + assert np.linalg.norm(result[0]) == pytest.approx(1.0) + assert not sqlite_embedding_session.in_transaction() + assert sqlite_embedding_session.scalar(select(Embedding)) is None - def test_initialization_sets_model_instance(self) -> None: - """Test CacheEmbedding initialization stores the provided model instance.""" - model_instance = Mock() - model_instance.model = "test-model" - model_instance.model_name = "test-model" - model_instance.provider = "test-provider" + def test_nan_vector_is_skipped_and_logged( + self, + sqlite_embedding_session: Session, + model_instance: Mock, + caplog: pytest.LogCaptureFixture, + ) -> None: + model_instance.invoke_text_embedding.return_value = _result(np.array([np.nan, 1.0])) - cache_embedding = CacheEmbedding(model_instance) + with caplog.at_level("WARNING", logger="core.rag.embedding.cached_embedding"): + result = CacheEmbedding(model_instance).embed_documents(["invalid"]) - assert cache_embedding._model_instance == model_instance + assert result == [None] + assert "Normalized embedding is nan" in caplog.text + assert sqlite_embedding_session.scalar(select(Embedding)) is None + + def test_duplicate_text_is_cached_once(self, sqlite_embedding_session: Session, model_instance: Mock) -> None: + model_instance.invoke_text_embedding.return_value = _result(_vector(), _vector()) + + result = CacheEmbedding(model_instance).embed_documents(["same", "same"]) + + assert result[0] == result[1] + assert len(sqlite_embedding_session.scalars(select(Embedding)).all()) == 1 + + def test_provider_and_model_are_part_of_cache_identity(self, sqlite_embedding_session: Session) -> None: + first = Mock(model_name="shared-model", provider="provider-a", credentials={}) + first.model_type_instance.get_model_schema.return_value = Mock( + model_properties={ModelPropertyKey.MAX_CHUNKS: 10} + ) + first.invoke_text_embedding.return_value = _result(_vector()) + second = Mock(model_name="shared-model", provider="provider-b", credentials={}) + second.model_type_instance.get_model_schema.return_value = Mock( + model_properties={ModelPropertyKey.MAX_CHUNKS: 10} + ) + second.invoke_text_embedding.return_value = _result(_vector(offset=10)) + + first_result = CacheEmbedding(first).embed_documents(["same text"]) + second_result = CacheEmbedding(second).embed_documents(["same text"]) + + assert first_result != second_result + assert len(sqlite_embedding_session.scalars(select(Embedding)).all()) == 2 + + def test_empty_input_does_not_touch_provider(self, model_instance: Mock) -> None: + assert CacheEmbedding(model_instance).embed_documents([]) == [] + model_instance.invoke_text_embedding.assert_not_called() + + class TestMultimodalDocumentCache: + def test_multimodal_miss_persists_and_hit_reuses( + self, sqlite_embedding_session: Session, model_instance: Mock + ) -> None: + document = {"file_id": "file-1"} + model_instance.invoke_multimodal_embedding.return_value = _result(_vector()) + cache = CacheEmbedding(model_instance) + + generated = cache.embed_multimodal_documents([document]) + cached = cache.embed_multimodal_documents([document]) + + assert cached == generated + model_instance.invoke_multimodal_embedding.assert_called_once_with( + multimodel_documents=[document], input_type=EmbeddingInputType.DOCUMENT + ) + assert sqlite_embedding_session.scalar(select(Embedding).where(Embedding.hash == "file-1")) is not None + + class TestQueryCache: + @patch("core.rag.embedding.cached_embedding.redis_client") + def test_cached_query_can_be_read_without_model_instance(self, redis: Mock) -> None: + vector = np.array([0.25, 0.75], dtype=float) + redis.get.return_value = base64.b64encode(vector.tobytes()) + + result = CacheEmbedding.get_cached_query_embedding("openai", "text-embedding-ada-002", "query") + + assert result == vector.tolist() + redis.get.assert_called_once_with(f"openai_text-embedding-ada-002_{helper.generate_text_hash('query')}") + redis.expire.assert_called_once_with( + f"openai_text-embedding-ada-002_{helper.generate_text_hash('query')}", 600 + ) + + @patch("core.rag.embedding.cached_embedding.redis_client") + def test_query_cache_miss_normalizes_and_stores(self, redis: Mock, model_instance: Mock) -> None: + redis.get.return_value = None + model_instance.invoke_text_embedding.return_value = _result(_vector()) + + result = CacheEmbedding(model_instance).embed_query("query") + + assert np.linalg.norm(result) == pytest.approx(1.0) + model_instance.invoke_text_embedding.assert_called_once_with( + texts=["query"], input_type=EmbeddingInputType.QUERY + ) + redis.setex.assert_called_once() + + @patch("core.rag.embedding.cached_embedding.redis_client") + def test_query_cache_hit_refreshes_ttl(self, redis: Mock, model_instance: Mock) -> None: + vector = np.array([0.25, 0.75], dtype=float) + redis.get.return_value = base64.b64encode(vector.tobytes()) + + result = CacheEmbedding(model_instance).embed_query("query") + + assert result == vector.tolist() + redis.expire.assert_called_once_with( + f"openai_text-embedding-ada-002_{helper.generate_text_hash('query')}", 600 + ) + model_instance.invoke_text_embedding.assert_not_called() + + @patch("core.rag.embedding.cached_embedding.redis_client") + def test_query_nan_is_rejected(self, redis: Mock, model_instance: Mock) -> None: + redis.get.return_value = None + model_instance.invoke_text_embedding.return_value = _result(np.array([np.nan, 1.0])) + + with pytest.raises(ValueError, match="Normalized embedding is nan"): + CacheEmbedding(model_instance).embed_query("query") + + redis.setex.assert_not_called() + + @patch("core.rag.embedding.cached_embedding.redis_client") + def test_redis_write_error_is_propagated(self, redis: Mock, model_instance: Mock) -> None: + redis.get.return_value = None + redis.setex.side_effect = ConnectionError("redis unavailable") + model_instance.invoke_text_embedding.return_value = _result(_vector()) + + with pytest.raises(ConnectionError, match="redis unavailable"): + CacheEmbedding(model_instance).embed_query("query") + + @patch("core.rag.embedding.cached_embedding.redis_client") + def test_multimodal_query_uses_file_cache_key(self, redis: Mock, model_instance: Mock) -> None: + redis.get.return_value = None + model_instance.invoke_multimodal_embedding.return_value = _result(_vector()) + + result = CacheEmbedding(model_instance).embed_multimodal_query({"file_id": "file-1"}) + + assert np.linalg.norm(result) == pytest.approx(1.0) + assert redis.setex.call_args.args[:2] == ("openai_text-embedding-ada-002_file-1", 600) diff --git a/api/tests/unit_tests/core/rag/embedding/test_embedding_service.py b/api/tests/unit_tests/core/rag/embedding/test_embedding_service.py deleted file mode 100644 index e116dc3454d..00000000000 --- a/api/tests/unit_tests/core/rag/embedding/test_embedding_service.py +++ /dev/null @@ -1,318 +0,0 @@ -"""Behavior-focused tests for :class:`CacheEmbedding`. - -Document and multimodal caches are persisted in SQLite so cache hits, misses, -provider isolation, uniqueness races, and rollback behavior exercise real -SQLAlchemy statements. Model providers and Redis remain mocked external I/O. -""" - -import base64 -import pickle -from dataclasses import dataclass -from decimal import Decimal -from unittest.mock import Mock, patch - -import numpy as np -import pytest -from sqlalchemy import event, select -from sqlalchemy.exc import IntegrityError -from sqlalchemy.orm import Session - -from core.entities.embedding_type import EmbeddingInputType -from core.rag.embedding import cached_embedding -from core.rag.embedding.cached_embedding import CacheEmbedding -from graphon.model_runtime.entities.model_entities import ModelPropertyKey -from graphon.model_runtime.entities.text_embedding_entities import EmbeddingResult, EmbeddingUsage -from graphon.model_runtime.errors.invoke import ( - InvokeAuthorizationError, - InvokeConnectionError, - InvokeRateLimitError, -) -from libs import helper -from models.dataset import Embedding - - -@dataclass(frozen=True) -class _Database: - """Expose a real session through the interface used by the cache.""" - - session: Session - - -@pytest.fixture -def sqlite_embedding_session(sqlite_session: Session, monkeypatch: pytest.MonkeyPatch) -> Session: - """Bind the shared SQLite session to the document-cache implementation.""" - - monkeypatch.setattr(cached_embedding, "db", _Database(sqlite_session)) - return sqlite_session - - -@pytest.fixture -def model_instance() -> Mock: - instance = Mock() - instance.model_name = "text-embedding-ada-002" - instance.provider = "openai" - instance.credentials = {"api_key": "test-key"} - schema = Mock(model_properties={ModelPropertyKey.MAX_CHUNKS: 10}) - instance.model_type_instance.get_model_schema.return_value = schema - return instance - - -def _vector(dimension: int = 8, *, offset: float = 0) -> np.ndarray: - return np.arange(1 + offset, dimension + 1 + offset, dtype=float) - - -def _result(*vectors: np.ndarray) -> EmbeddingResult: - return EmbeddingResult( - model="text-embedding-ada-002", - embeddings=list(vectors), - usage=EmbeddingUsage( - tokens=len(vectors), - total_tokens=len(vectors), - unit_price=Decimal("0.0001"), - price_unit=Decimal(1000), - total_price=Decimal("0.000001"), - currency="USD", - latency=0.1, - ), - ) - - -def _persist_embedding( - session: Session, - *, - text: str, - vector: list[float], - model_name: str = "text-embedding-ada-002", - provider_name: str = "openai", -) -> Embedding: - row = Embedding( - model_name=model_name, - hash=helper.generate_text_hash(text), - provider_name=provider_name, - embedding=pickle.dumps(vector, protocol=pickle.HIGHEST_PROTOCOL), - ) - session.add(row) - session.commit() - return row - - -class TestDocumentCache: - def test_cache_miss_normalizes_and_persists(self, sqlite_embedding_session: Session, model_instance: Mock) -> None: - model_instance.invoke_text_embedding.return_value = _result(_vector()) - - result = CacheEmbedding(model_instance).embed_documents(["new text"]) - - assert np.linalg.norm(result[0]) == pytest.approx(1.0) - model_instance.invoke_text_embedding.assert_called_once_with( - texts=["new text"], input_type=EmbeddingInputType.DOCUMENT - ) - persisted = sqlite_embedding_session.scalar(select(Embedding)) - assert persisted is not None - assert persisted.get_embedding() == result[0] - - def test_cache_hit_uses_persisted_vector(self, sqlite_embedding_session: Session, model_instance: Mock) -> None: - cached_vector = (_vector() / np.linalg.norm(_vector())).tolist() - row = _persist_embedding(sqlite_embedding_session, text="cached", vector=cached_vector) - - result = CacheEmbedding(model_instance).embed_documents(["cached"]) - - assert result == [cached_vector] - model_instance.invoke_text_embedding.assert_not_called() - assert sqlite_embedding_session.scalars(select(Embedding)).all() == [row] - - def test_partial_hit_preserves_input_order(self, sqlite_embedding_session: Session, model_instance: Mock) -> None: - cached_vector = (_vector(offset=10) / np.linalg.norm(_vector(offset=10))).tolist() - _persist_embedding(sqlite_embedding_session, text="cached", vector=cached_vector) - model_instance.invoke_text_embedding.return_value = _result(_vector(), _vector(offset=20)) - - result = CacheEmbedding(model_instance).embed_documents(["cached", "new one", "new two"]) - - assert result[0] == cached_vector - assert all(np.linalg.norm(vector) == pytest.approx(1.0) for vector in result) - assert model_instance.invoke_text_embedding.call_args.kwargs["texts"] == ["new one", "new two"] - assert len(sqlite_embedding_session.scalars(select(Embedding)).all()) == 3 - - def test_large_batch_respects_model_chunk_limit( - self, sqlite_embedding_session: Session, model_instance: Mock - ) -> None: - texts = [f"text-{index}" for index in range(25)] - model_instance.invoke_text_embedding.side_effect = [ - _result(*(_vector(offset=index) for index in range(10))), - _result(*(_vector(offset=index) for index in range(10, 20))), - _result(*(_vector(offset=index) for index in range(20, 25))), - ] - - result = CacheEmbedding(model_instance).embed_documents(texts) - - assert len(result) == 25 - assert [len(call.kwargs["texts"]) for call in model_instance.invoke_text_embedding.call_args_list] == [ - 10, - 10, - 5, - ] - assert len(sqlite_embedding_session.scalars(select(Embedding)).all()) == 25 - - @pytest.mark.parametrize( - "provider_error", - [InvokeConnectionError("offline"), InvokeRateLimitError("limited"), InvokeAuthorizationError("denied")], - ) - def test_provider_error_rolls_back_real_transaction( - self, - sqlite_embedding_session: Session, - model_instance: Mock, - provider_error: Exception, - ) -> None: - model_instance.invoke_text_embedding.side_effect = provider_error - - with pytest.raises(type(provider_error)): - CacheEmbedding(model_instance).embed_documents(["failure"]) - - assert not sqlite_embedding_session.in_transaction() - assert sqlite_embedding_session.scalar(select(Embedding)) is None - - def test_integrity_race_rolls_back_without_losing_result( - self, sqlite_embedding_session: Session, model_instance: Mock - ) -> None: - model_instance.invoke_text_embedding.return_value = _result(_vector()) - - def raise_integrity_error(_session: Session) -> None: - raise IntegrityError("duplicate embedding", {}, RuntimeError("unique constraint")) - - event.listen(sqlite_embedding_session, "before_commit", raise_integrity_error) - try: - result = CacheEmbedding(model_instance).embed_documents(["racing insert"]) - finally: - event.remove(sqlite_embedding_session, "before_commit", raise_integrity_error) - - assert np.linalg.norm(result[0]) == pytest.approx(1.0) - assert not sqlite_embedding_session.in_transaction() - assert sqlite_embedding_session.scalar(select(Embedding)) is None - - def test_nan_vector_is_skipped_and_logged( - self, - sqlite_embedding_session: Session, - model_instance: Mock, - caplog: pytest.LogCaptureFixture, - ) -> None: - model_instance.invoke_text_embedding.return_value = _result(np.array([np.nan, 1.0])) - - with caplog.at_level("WARNING", logger="core.rag.embedding.cached_embedding"): - result = CacheEmbedding(model_instance).embed_documents(["invalid"]) - - assert result == [None] - assert "Normalized embedding is nan" in caplog.text - assert sqlite_embedding_session.scalar(select(Embedding)) is None - - def test_duplicate_text_is_cached_once(self, sqlite_embedding_session: Session, model_instance: Mock) -> None: - model_instance.invoke_text_embedding.return_value = _result(_vector(), _vector()) - - result = CacheEmbedding(model_instance).embed_documents(["same", "same"]) - - assert result[0] == result[1] - assert len(sqlite_embedding_session.scalars(select(Embedding)).all()) == 1 - - def test_provider_and_model_are_part_of_cache_identity(self, sqlite_embedding_session: Session) -> None: - first = Mock(model_name="shared-model", provider="provider-a", credentials={}) - first.model_type_instance.get_model_schema.return_value = Mock( - model_properties={ModelPropertyKey.MAX_CHUNKS: 10} - ) - first.invoke_text_embedding.return_value = _result(_vector()) - second = Mock(model_name="shared-model", provider="provider-b", credentials={}) - second.model_type_instance.get_model_schema.return_value = Mock( - model_properties={ModelPropertyKey.MAX_CHUNKS: 10} - ) - second.invoke_text_embedding.return_value = _result(_vector(offset=10)) - - first_result = CacheEmbedding(first).embed_documents(["same text"]) - second_result = CacheEmbedding(second).embed_documents(["same text"]) - - assert first_result != second_result - assert len(sqlite_embedding_session.scalars(select(Embedding)).all()) == 2 - - def test_empty_input_does_not_touch_provider(self, model_instance: Mock) -> None: - assert CacheEmbedding(model_instance).embed_documents([]) == [] - model_instance.invoke_text_embedding.assert_not_called() - - -class TestMultimodalDocumentCache: - def test_multimodal_miss_persists_and_hit_reuses( - self, sqlite_embedding_session: Session, model_instance: Mock - ) -> None: - document = {"file_id": "file-1"} - model_instance.invoke_multimodal_embedding.return_value = _result(_vector()) - cache = CacheEmbedding(model_instance) - - generated = cache.embed_multimodal_documents([document]) - cached = cache.embed_multimodal_documents([document]) - - assert cached == generated - model_instance.invoke_multimodal_embedding.assert_called_once_with( - multimodel_documents=[document], input_type=EmbeddingInputType.DOCUMENT - ) - assert sqlite_embedding_session.scalar(select(Embedding).where(Embedding.hash == "file-1")) is not None - - -class TestQueryCache: - @patch("core.rag.embedding.cached_embedding.redis_client") - def test_cached_query_can_be_read_without_model_instance(self, redis: Mock) -> None: - vector = np.array([0.25, 0.75], dtype=float) - redis.get.return_value = base64.b64encode(vector.tobytes()) - - result = CacheEmbedding.get_cached_query_embedding("openai", "text-embedding-ada-002", "query") - - assert result == vector.tolist() - redis.get.assert_called_once_with(f"openai_text-embedding-ada-002_{helper.generate_text_hash('query')}") - redis.expire.assert_called_once_with(f"openai_text-embedding-ada-002_{helper.generate_text_hash('query')}", 600) - - @patch("core.rag.embedding.cached_embedding.redis_client") - def test_query_cache_miss_normalizes_and_stores(self, redis: Mock, model_instance: Mock) -> None: - redis.get.return_value = None - model_instance.invoke_text_embedding.return_value = _result(_vector()) - - result = CacheEmbedding(model_instance).embed_query("query") - - assert np.linalg.norm(result) == pytest.approx(1.0) - model_instance.invoke_text_embedding.assert_called_once_with( - texts=["query"], input_type=EmbeddingInputType.QUERY - ) - redis.setex.assert_called_once() - - @patch("core.rag.embedding.cached_embedding.redis_client") - def test_query_cache_hit_refreshes_ttl(self, redis: Mock, model_instance: Mock) -> None: - vector = np.array([0.25, 0.75], dtype=float) - redis.get.return_value = base64.b64encode(vector.tobytes()) - - result = CacheEmbedding(model_instance).embed_query("query") - - assert result == vector.tolist() - redis.expire.assert_called_once_with(f"openai_text-embedding-ada-002_{helper.generate_text_hash('query')}", 600) - model_instance.invoke_text_embedding.assert_not_called() - - @patch("core.rag.embedding.cached_embedding.redis_client") - def test_query_nan_is_rejected(self, redis: Mock, model_instance: Mock) -> None: - redis.get.return_value = None - model_instance.invoke_text_embedding.return_value = _result(np.array([np.nan, 1.0])) - - with pytest.raises(ValueError, match="Normalized embedding is nan"): - CacheEmbedding(model_instance).embed_query("query") - - redis.setex.assert_not_called() - - @patch("core.rag.embedding.cached_embedding.redis_client") - def test_redis_write_error_is_propagated(self, redis: Mock, model_instance: Mock) -> None: - redis.get.return_value = None - redis.setex.side_effect = ConnectionError("redis unavailable") - model_instance.invoke_text_embedding.return_value = _result(_vector()) - - with pytest.raises(ConnectionError, match="redis unavailable"): - CacheEmbedding(model_instance).embed_query("query") - - @patch("core.rag.embedding.cached_embedding.redis_client") - def test_multimodal_query_uses_file_cache_key(self, redis: Mock, model_instance: Mock) -> None: - redis.get.return_value = None - model_instance.invoke_multimodal_embedding.return_value = _result(_vector()) - - result = CacheEmbedding(model_instance).embed_multimodal_query({"file_id": "file-1"}) - - assert np.linalg.norm(result) == pytest.approx(1.0) - assert redis.setex.call_args.args[:2] == ("openai_text-embedding-ada-002_file-1", 600) diff --git a/api/tests/unit_tests/core/rag/extractor/firecrawl/test_firecrawl.py b/api/tests/unit_tests/core/rag/extractor/firecrawl/test_firecrawl.py deleted file mode 100644 index 9a33ff6effd..00000000000 --- a/api/tests/unit_tests/core/rag/extractor/firecrawl/test_firecrawl.py +++ /dev/null @@ -1,468 +0,0 @@ -"""Unit tests for Firecrawl app and extractor integration points.""" - -import json -from collections.abc import Mapping -from unittest.mock import MagicMock, create_autospec - -import pytest -from pytest_mock import MockerFixture - -import core.rag.extractor.firecrawl.firecrawl_app as firecrawl_module -from core.rag.extractor.firecrawl.firecrawl_app import FirecrawlApp -from core.rag.extractor.firecrawl.firecrawl_web_extractor import FirecrawlWebExtractor -from services.data_source.provider_service import DatasourceProviderService -from services.data_source.website_service import WebsiteService - - -def _response(status_code: int, json_data: Mapping[str, object] | None = None, text: str = "") -> MagicMock: - response = MagicMock() - response.status_code = status_code - response.text = text - response.json.return_value = json_data if json_data is not None else dict[str, object]() - return response - - -class TestFirecrawlApp: - def test_init_requires_api_key_for_default_base_url(self) -> None: - with pytest.raises(ValueError, match="No API key provided"): - FirecrawlApp(api_key=None, base_url="https://api.firecrawl.dev") - - def test_prepare_headers_and_build_url(self) -> None: - app = FirecrawlApp(api_key="fc-key", base_url="https://custom.firecrawl.dev/") - - assert app._prepare_headers() == { - "Content-Type": "application/json", - "Authorization": "Bearer fc-key", - } - assert app._build_url("/v2/crawl") == "https://custom.firecrawl.dev/v2/crawl" - - def test_requests_use_bounded_timeout(self, mocker: MockerFixture) -> None: - """Outbound requests must carry a bounded timeout so a hanging endpoint cannot block extraction.""" - app = FirecrawlApp(api_key="fc-key", base_url="https://custom.firecrawl.dev") - mock_post = mocker.patch("core.helper.ssrf_proxy.post", return_value=_response(200, {"id": "job-1"})) - mock_get = mocker.patch("core.helper.ssrf_proxy.get", return_value=_response(200, {"status": "completed"})) - - app._post_request("https://custom.firecrawl.dev/v1/crawl", {}, app._prepare_headers()) - app._get_request("https://custom.firecrawl.dev/v1/crawl/job-1", app._prepare_headers()) - - assert mock_post.call_args.kwargs["timeout"] == firecrawl_module._REQUEST_TIMEOUT - assert mock_get.call_args.kwargs["timeout"] == firecrawl_module._REQUEST_TIMEOUT - - def test_scrape_url_success(self, mocker: MockerFixture) -> None: - app = FirecrawlApp(api_key="fc-key", base_url="https://custom.firecrawl.dev") - mocker.patch( - "core.helper.ssrf_proxy.post", - return_value=_response( - 200, - { - "data": { - "metadata": { - "title": "t", - "description": "d", - "sourceURL": "https://example.com", - }, - "markdown": "body", - } - }, - ), - ) - - result = app.scrape_url("https://example.com", params={"onlyMainContent": False}) - - assert result == { - "title": "t", - "description": "d", - "source_url": "https://example.com", - "markdown": "body", - } - - def test_scrape_url_handles_known_error_status(self, mocker: MockerFixture) -> None: - app = FirecrawlApp(api_key="fc-key", base_url="https://custom.firecrawl.dev") - mock_handle = mocker.patch.object(app, "_handle_error", side_effect=Exception("boom")) - mocker.patch("core.helper.ssrf_proxy.post", return_value=_response(429, {"error": "limit"})) - - with pytest.raises(Exception, match="boom"): - app.scrape_url("https://example.com") - - mock_handle.assert_called_once() - - def test_scrape_url_unknown_status_raises(self, mocker: MockerFixture) -> None: - app = FirecrawlApp(api_key="fc-key", base_url="https://custom.firecrawl.dev") - mocker.patch("core.helper.ssrf_proxy.post", return_value=_response(404, text="Not Found")) - - with pytest.raises(Exception, match="Failed to scrape URL. Status code: 404"): - app.scrape_url("https://example.com") - - def test_crawl_url_success(self, mocker: MockerFixture) -> None: - app = FirecrawlApp(api_key="fc-key", base_url="https://custom.firecrawl.dev") - mocker.patch("core.helper.ssrf_proxy.post", return_value=_response(200, {"id": "job-1"})) - - assert app.crawl_url("https://example.com") == "job-1" - - def test_crawl_url_non_200_uses_error_handler(self, mocker: MockerFixture) -> None: - app = FirecrawlApp(api_key="fc-key", base_url="https://custom.firecrawl.dev") - mock_handle = mocker.patch.object(app, "_handle_error", side_effect=Exception("crawl failed")) - mocker.patch("core.helper.ssrf_proxy.post", return_value=_response(500, {"error": "server"})) - - with pytest.raises(Exception, match="crawl failed"): - app.crawl_url("https://example.com") - - mock_handle.assert_called_once() - - def test_map_success(self, mocker: MockerFixture) -> None: - app = FirecrawlApp(api_key="fc-key", base_url="https://custom.firecrawl.dev") - mocker.patch("core.helper.ssrf_proxy.post", return_value=_response(200, {"success": True, "links": ["a", "b"]})) - - assert app.map("https://example.com") == {"success": True, "links": ["a", "b"]} - - def test_map_known_error(self, mocker: MockerFixture) -> None: - app = FirecrawlApp(api_key="fc-key", base_url="https://custom.firecrawl.dev") - mock_handle = mocker.patch.object(app, "_handle_error", side_effect=Exception("map error")) - mocker.patch("core.helper.ssrf_proxy.post", return_value=_response(409, {"error": "conflict"})) - - with pytest.raises(Exception, match="map error"): - app.map("https://example.com") - mock_handle.assert_called_once() - - def test_map_unknown_error_raises(self, mocker: MockerFixture) -> None: - app = FirecrawlApp(api_key="fc-key", base_url="https://custom.firecrawl.dev") - mocker.patch("core.helper.ssrf_proxy.post", return_value=_response(418, text="teapot")) - - with pytest.raises(Exception, match="Failed to start map job. Status code: 418"): - app.map("https://example.com") - - def test_check_crawl_status_completed_with_data(self, mocker: MockerFixture) -> None: - app = FirecrawlApp(api_key="fc-key", base_url="https://custom.firecrawl.dev") - payload = { - "status": "completed", - "total": 2, - "completed": 2, - "data": [ - { - "metadata": {"title": "a", "description": "desc-a", "sourceURL": "https://a"}, - "markdown": "m-a", - }, - { - "metadata": {"title": "b", "description": "desc-b", "sourceURL": "https://b"}, - "markdown": "m-b", - }, - {"metadata": {"title": "skip"}}, - ], - } - mocker.patch("core.helper.ssrf_proxy.get", return_value=_response(200, payload)) - - save_calls: list[tuple[str, bytes]] = [] - delete_calls: list[str] = [] - - mock_storage = MagicMock() - mock_storage.exists.return_value = True - mock_storage.delete.side_effect = lambda key: delete_calls.append(key) - mock_storage.save.side_effect = lambda key, data: save_calls.append((key, data)) - mocker.patch.object(firecrawl_module, "storage", mock_storage) - - result = app.check_crawl_status("job-42") - - assert result["status"] == "completed" - assert result["total"] == 2 - assert result["current"] == 2 - assert len(result["data"]) == 2 - assert delete_calls == ["website_files/job-42.txt"] - assert len(save_calls) == 1 - assert save_calls[0][0] == "website_files/job-42.txt" - - def test_check_crawl_status_completed_with_zero_total_raises(self, mocker: MockerFixture) -> None: - app = FirecrawlApp(api_key="fc-key", base_url="https://custom.firecrawl.dev") - mocker.patch( - "core.helper.ssrf_proxy.get", return_value=_response(200, {"status": "completed", "total": 0, "data": []}) - ) - - with pytest.raises(Exception, match="No page found"): - app.check_crawl_status("job-1") - - def test_check_crawl_status_completed_with_null_total_raises(self, mocker: MockerFixture) -> None: - app = FirecrawlApp(api_key="fc-key", base_url="https://custom.firecrawl.dev") - mocker.patch( - "core.helper.ssrf_proxy.get", - return_value=_response(200, {"status": "completed", "total": None, "data": []}), - ) - - with pytest.raises(Exception, match="No page found"): - app.check_crawl_status("job-1") - - def test_check_crawl_status_non_completed(self, mocker: MockerFixture) -> None: - app = FirecrawlApp(api_key="fc-key", base_url="https://custom.firecrawl.dev") - payload = {"status": "processing", "total": 5, "completed": 1, "data": list[object]()} - mocker.patch("core.helper.ssrf_proxy.get", return_value=_response(200, payload)) - - assert app.check_crawl_status("job-1") == { - "status": "processing", - "total": 5, - "current": 1, - "data": [], - } - - def test_check_crawl_status_non_200_uses_error_handler(self, mocker: MockerFixture) -> None: - app = FirecrawlApp(api_key="fc-key", base_url="https://custom.firecrawl.dev") - mock_handle = mocker.patch.object(app, "_handle_error", side_effect=Exception("crawl error")) - mocker.patch("core.helper.ssrf_proxy.get", return_value=_response(500, {"error": "server"})) - - with pytest.raises(Exception, match="crawl error"): - app.check_crawl_status("job-1") - mock_handle.assert_called_once() - - def test_check_crawl_status_save_failure_raises(self, mocker: MockerFixture) -> None: - app = FirecrawlApp(api_key="fc-key", base_url="https://custom.firecrawl.dev") - payload = { - "status": "completed", - "total": 1, - "completed": 1, - "data": [{"metadata": {"title": "a", "sourceURL": "https://a"}, "markdown": "m-a"}], - } - mocker.patch("core.helper.ssrf_proxy.get", return_value=_response(200, payload)) - - mock_storage = MagicMock() - mock_storage.exists.return_value = False - mock_storage.save.side_effect = RuntimeError("save failed") - mocker.patch.object(firecrawl_module, "storage", mock_storage) - - with pytest.raises(Exception, match="Error saving crawl data"): - app.check_crawl_status("job-err") - - def test_check_crawl_status_follows_pagination(self, mocker: MockerFixture) -> None: - """When status is completed and next is present, follow pagination to collect all pages.""" - app = FirecrawlApp(api_key="fc-key", base_url="https://custom.firecrawl.dev") - page1 = { - "status": "completed", - "total": 3, - "completed": 3, - "next": "https://custom.firecrawl.dev/v2/crawl/job-42?skip=1", - "data": [{"metadata": {"title": "p1", "description": "", "sourceURL": "https://p1"}, "markdown": "m1"}], - } - page2 = { - "status": "completed", - "total": 3, - "completed": 3, - "next": "https://custom.firecrawl.dev/v2/crawl/job-42?skip=2", - "data": [{"metadata": {"title": "p2", "description": "", "sourceURL": "https://p2"}, "markdown": "m2"}], - } - page3 = { - "status": "completed", - "total": 3, - "completed": 3, - "data": [{"metadata": {"title": "p3", "description": "", "sourceURL": "https://p3"}, "markdown": "m3"}], - } - mocker.patch( - "core.helper.ssrf_proxy.get", - side_effect=[_response(200, page1), _response(200, page2), _response(200, page3)], - ) - mock_storage = MagicMock() - mock_storage.exists.return_value = False - mocker.patch.object(firecrawl_module, "storage", mock_storage) - - result = app.check_crawl_status("job-42") - - assert result["status"] == "completed" - assert result["total"] == 3 - assert len(result["data"]) == 3 - assert [d["title"] for d in result["data"]] == ["p1", "p2", "p3"] - - def test_check_crawl_status_pagination_error_raises(self, mocker: MockerFixture) -> None: - """An error while fetching a paginated page raises an exception; no partial data is returned.""" - app = FirecrawlApp(api_key="fc-key", base_url="https://custom.firecrawl.dev") - page1 = { - "status": "completed", - "total": 2, - "completed": 2, - "next": "https://custom.firecrawl.dev/v2/crawl/job-99?skip=1", - "data": [{"metadata": {"title": "p1", "description": "", "sourceURL": "https://p1"}, "markdown": "m1"}], - } - mocker.patch( - "core.helper.ssrf_proxy.get", side_effect=[_response(200, page1), _response(500, {"error": "server error"})] - ) - - with pytest.raises(Exception, match="fetch next crawl page"): - app.check_crawl_status("job-99") - - def test_check_crawl_status_pagination_capped_at_total(self, mocker: MockerFixture) -> None: - """Pagination stops once pages_processed reaches total, even if next is present.""" - app = FirecrawlApp(api_key="fc-key", base_url="https://custom.firecrawl.dev") - # total=1: only the first page should be processed; next must not be followed - page1 = { - "status": "completed", - "total": 1, - "completed": 1, - "next": "https://custom.firecrawl.dev/v2/crawl/job-cap?skip=1", - "data": [{"metadata": {"title": "p1", "description": "", "sourceURL": "https://p1"}, "markdown": "m1"}], - } - mock_get = mocker.patch("core.helper.ssrf_proxy.get", return_value=_response(200, page1)) - mock_storage = MagicMock() - mock_storage.exists.return_value = False - mocker.patch.object(firecrawl_module, "storage", mock_storage) - - result = app.check_crawl_status("job-cap") - - assert len(result["data"]) == 1 - mock_get.assert_called_once() # initial fetch only; next URL is not followed due to cap - - def test_extract_common_fields_and_status_formatter(self) -> None: - app = FirecrawlApp(api_key="fc-key", base_url="https://custom.firecrawl.dev") - - fields = app._extract_common_fields( - {"metadata": {"title": "t", "description": "d", "sourceURL": "u"}, "markdown": "m"} - ) - assert fields == {"title": "t", "description": "d", "source_url": "u", "markdown": "m"} - - status = app._format_crawl_status_response("completed", {"total": 1, "completed": 1}, [fields]) - assert status == {"status": "completed", "total": 1, "current": 1, "data": [fields]} - - def test_post_and_get_request_retry_logic(self, mocker: MockerFixture) -> None: - app = FirecrawlApp(api_key="fc-key", base_url="https://custom.firecrawl.dev") - sleep_mock = mocker.patch.object(firecrawl_module.time, "sleep") - - resp_502_a = _response(502) - resp_502_b = _response(502) - resp_200 = _response(200) - - mocker.patch("core.helper.ssrf_proxy.post", side_effect=[resp_502_a, resp_200]) - post_result = app._post_request("u", {"x": 1}, {"h": 1}, retries=3, backoff_factor=0.5) - assert post_result is resp_200 - - mocker.patch("core.helper.ssrf_proxy.get", side_effect=[resp_502_b, _response(200)]) - get_result = app._get_request("u", {"h": 1}, retries=3, backoff_factor=0.25) - assert get_result.status_code == 200 - - assert sleep_mock.call_count == 2 - - def test_post_and_get_request_return_last_502(self, mocker: MockerFixture) -> None: - app = FirecrawlApp(api_key="fc-key", base_url="https://custom.firecrawl.dev") - sleep_mock = mocker.patch.object(firecrawl_module.time, "sleep") - - last_post = _response(502) - mocker.patch("core.helper.ssrf_proxy.post", side_effect=[_response(502), last_post]) - assert app._post_request("u", {}, {}, retries=2).status_code == 502 - - last_get = _response(502) - mocker.patch("core.helper.ssrf_proxy.get", side_effect=[_response(502), last_get]) - assert app._get_request("u", {}, retries=2).status_code == 502 - - assert sleep_mock.call_count == 4 - - def test_handle_error_with_json_and_plain_text(self) -> None: - app = FirecrawlApp(api_key="fc-key", base_url="https://custom.firecrawl.dev") - - json_error = _response(400, {"message": "bad request"}) - with pytest.raises(Exception, match="bad request"): - app._handle_error(json_error, "run task") - - non_json = MagicMock() - non_json.status_code = 400 - non_json.text = "plain error" - non_json.json.side_effect = json.JSONDecodeError("bad", "x", 0) - - with pytest.raises(Exception, match="plain error"): - app._handle_error(non_json, "run task") - - def test_search_success(self, mocker: MockerFixture) -> None: - app = FirecrawlApp(api_key="fc-key", base_url="https://custom.firecrawl.dev") - mocker.patch( - "core.helper.ssrf_proxy.post", return_value=_response(200, {"success": True, "data": [{"url": "x"}]}) - ) - assert app.search("python")["success"] is True - - def test_search_warning_failure(self, mocker: MockerFixture) -> None: - app = FirecrawlApp(api_key="fc-key", base_url="https://custom.firecrawl.dev") - mocker.patch( - "core.helper.ssrf_proxy.post", return_value=_response(200, {"success": False, "warning": "bad search"}) - ) - with pytest.raises(Exception, match="bad search"): - app.search("python") - - def test_search_known_http_error(self, mocker: MockerFixture) -> None: - app = FirecrawlApp(api_key="fc-key", base_url="https://custom.firecrawl.dev") - mock_handle = mocker.patch.object(app, "_handle_error", side_effect=Exception("search error")) - mocker.patch("core.helper.ssrf_proxy.post", return_value=_response(408, {"error": "timeout"})) - with pytest.raises(Exception, match="search error"): - app.search("python") - mock_handle.assert_called_once() - - def test_search_unknown_http_error(self, mocker: MockerFixture) -> None: - app = FirecrawlApp(api_key="fc-key", base_url="https://custom.firecrawl.dev") - mocker.patch("core.helper.ssrf_proxy.post", return_value=_response(418, text="teapot")) - with pytest.raises(Exception, match="Failed to perform search. Status code: 418"): - app.search("python") - - -class TestFirecrawlWebExtractor: - def test_extract_crawl_mode_returns_document(self, mocker: MockerFixture) -> None: - mocker.patch( - "core.rag.extractor.firecrawl.firecrawl_web_extractor.WebsiteService.get_crawl_url_data", - return_value={ - "markdown": "crawl content", - "source_url": "https://example.com", - "description": "desc", - "title": "title", - }, - ) - - extractor = FirecrawlWebExtractor( - "https://example.com", - "job-1", - "tenant-1", - mode="crawl", - website_service=WebsiteService(providers=create_autospec(DatasourceProviderService, instance=True)), - ) - docs = extractor.extract() - - assert len(docs) == 1 - assert docs[0].page_content == "crawl content" - assert docs[0].metadata["source_url"] == "https://example.com" - - def test_extract_crawl_mode_with_missing_data_returns_empty(self, mocker: MockerFixture) -> None: - mocker.patch( - "core.rag.extractor.firecrawl.firecrawl_web_extractor.WebsiteService.get_crawl_url_data", - return_value=None, - ) - - extractor = FirecrawlWebExtractor( - "https://example.com", - "job-1", - "tenant-1", - mode="crawl", - website_service=WebsiteService(providers=create_autospec(DatasourceProviderService, instance=True)), - ) - assert extractor.extract() == [] - - def test_extract_scrape_mode_returns_document(self, mocker: MockerFixture) -> None: - mock_scrape = mocker.patch( - "core.rag.extractor.firecrawl.firecrawl_web_extractor.WebsiteService.get_scrape_url_data", - return_value={ - "markdown": "scrape content", - "source_url": "https://example.com", - "description": "desc", - "title": "title", - }, - ) - - extractor = FirecrawlWebExtractor( - "https://example.com", - "job-1", - "tenant-1", - mode="scrape", - only_main_content=False, - website_service=WebsiteService(providers=create_autospec(DatasourceProviderService, instance=True)), - ) - docs = extractor.extract() - - assert len(docs) == 1 - assert docs[0].page_content == "scrape content" - mock_scrape.assert_called_once_with("firecrawl", "https://example.com", "tenant-1", False) - - def test_extract_unknown_mode_returns_empty(self) -> None: - extractor = FirecrawlWebExtractor( - "https://example.com", - "job-1", - "tenant-1", - mode="unknown", - website_service=WebsiteService(providers=create_autospec(DatasourceProviderService, instance=True)), - ) - assert extractor.extract() == [] diff --git a/api/tests/unit_tests/core/rag/extractor/firecrawl/test_firecrawl_app.py b/api/tests/unit_tests/core/rag/extractor/firecrawl/test_firecrawl_app.py new file mode 100644 index 00000000000..9b22531cba2 --- /dev/null +++ b/api/tests/unit_tests/core/rag/extractor/firecrawl/test_firecrawl_app.py @@ -0,0 +1,627 @@ +import json +from collections.abc import Callable, Iterator, Mapping +from contextlib import ExitStack +from typing import Never +from unittest.mock import MagicMock, create_autospec + +import httpx +import pytest +from pytest_mock import MockerFixture + +import core.rag.extractor.firecrawl.firecrawl_app as firecrawl_module +from core.helper import ssrf_proxy +from core.rag.extractor.firecrawl.firecrawl_app import FirecrawlApp +from core.rag.extractor.firecrawl.firecrawl_web_extractor import FirecrawlWebExtractor +from core.rag.extractor.watercrawl.client import WaterCrawlAPIClient +from core.tools.errors import ToolSSRFError +from services.data_source.auth.firecrawl.firecrawl import FirecrawlAuth +from services.data_source.auth.watercrawl.watercrawl import WatercrawlAuth +from services.data_source.entities.api_key_auth import DataSourceApiKeyAuthCredentials +from services.data_source.provider_service import DatasourceProviderService +from services.data_source.website_service import WebsiteService + +"""Unit tests for Firecrawl app and extractor integration points.""" + + +def _response(status_code: int, json_data: Mapping[str, object] | None = None, text: str = "") -> MagicMock: + response = MagicMock() + response.status_code = status_code + response.text = text + response.json.return_value = json_data if json_data is not None else dict[str, object]() + return response + + +class TestFirecrawlCases: + class TestFirecrawlApp: + def test_init_requires_api_key_for_default_base_url(self) -> None: + with pytest.raises(ValueError, match="No API key provided"): + FirecrawlApp(api_key=None, base_url="https://api.firecrawl.dev") + + def test_prepare_headers_and_build_url(self) -> None: + app = FirecrawlApp(api_key="fc-key", base_url="https://custom.firecrawl.dev/") + + assert app._prepare_headers() == { + "Content-Type": "application/json", + "Authorization": "Bearer fc-key", + } + assert app._build_url("/v2/crawl") == "https://custom.firecrawl.dev/v2/crawl" + + def test_requests_use_bounded_timeout(self, mocker: MockerFixture) -> None: + """Outbound requests must carry a bounded timeout so a hanging endpoint cannot block extraction.""" + app = FirecrawlApp(api_key="fc-key", base_url="https://custom.firecrawl.dev") + mock_post = mocker.patch("core.helper.ssrf_proxy.post", return_value=_response(200, {"id": "job-1"})) + mock_get = mocker.patch("core.helper.ssrf_proxy.get", return_value=_response(200, {"status": "completed"})) + + app._post_request("https://custom.firecrawl.dev/v1/crawl", {}, app._prepare_headers()) + app._get_request("https://custom.firecrawl.dev/v1/crawl/job-1", app._prepare_headers()) + + assert mock_post.call_args.kwargs["timeout"] == firecrawl_module._REQUEST_TIMEOUT + assert mock_get.call_args.kwargs["timeout"] == firecrawl_module._REQUEST_TIMEOUT + + def test_scrape_url_success(self, mocker: MockerFixture) -> None: + app = FirecrawlApp(api_key="fc-key", base_url="https://custom.firecrawl.dev") + mocker.patch( + "core.helper.ssrf_proxy.post", + return_value=_response( + 200, + { + "data": { + "metadata": { + "title": "t", + "description": "d", + "sourceURL": "https://example.com", + }, + "markdown": "body", + } + }, + ), + ) + + result = app.scrape_url("https://example.com", params={"onlyMainContent": False}) + + assert result == { + "title": "t", + "description": "d", + "source_url": "https://example.com", + "markdown": "body", + } + + def test_scrape_url_handles_known_error_status(self, mocker: MockerFixture) -> None: + app = FirecrawlApp(api_key="fc-key", base_url="https://custom.firecrawl.dev") + mock_handle = mocker.patch.object(app, "_handle_error", side_effect=Exception("boom")) + mocker.patch("core.helper.ssrf_proxy.post", return_value=_response(429, {"error": "limit"})) + + with pytest.raises(Exception, match="boom"): + app.scrape_url("https://example.com") + + mock_handle.assert_called_once() + + def test_scrape_url_unknown_status_raises(self, mocker: MockerFixture) -> None: + app = FirecrawlApp(api_key="fc-key", base_url="https://custom.firecrawl.dev") + mocker.patch("core.helper.ssrf_proxy.post", return_value=_response(404, text="Not Found")) + + with pytest.raises(Exception, match="Failed to scrape URL. Status code: 404"): + app.scrape_url("https://example.com") + + def test_crawl_url_success(self, mocker: MockerFixture) -> None: + app = FirecrawlApp(api_key="fc-key", base_url="https://custom.firecrawl.dev") + mocker.patch("core.helper.ssrf_proxy.post", return_value=_response(200, {"id": "job-1"})) + + assert app.crawl_url("https://example.com") == "job-1" + + def test_crawl_url_non_200_uses_error_handler(self, mocker: MockerFixture) -> None: + app = FirecrawlApp(api_key="fc-key", base_url="https://custom.firecrawl.dev") + mock_handle = mocker.patch.object(app, "_handle_error", side_effect=Exception("crawl failed")) + mocker.patch("core.helper.ssrf_proxy.post", return_value=_response(500, {"error": "server"})) + + with pytest.raises(Exception, match="crawl failed"): + app.crawl_url("https://example.com") + + mock_handle.assert_called_once() + + def test_map_success(self, mocker: MockerFixture) -> None: + app = FirecrawlApp(api_key="fc-key", base_url="https://custom.firecrawl.dev") + mocker.patch( + "core.helper.ssrf_proxy.post", return_value=_response(200, {"success": True, "links": ["a", "b"]}) + ) + + assert app.map("https://example.com") == {"success": True, "links": ["a", "b"]} + + def test_map_known_error(self, mocker: MockerFixture) -> None: + app = FirecrawlApp(api_key="fc-key", base_url="https://custom.firecrawl.dev") + mock_handle = mocker.patch.object(app, "_handle_error", side_effect=Exception("map error")) + mocker.patch("core.helper.ssrf_proxy.post", return_value=_response(409, {"error": "conflict"})) + + with pytest.raises(Exception, match="map error"): + app.map("https://example.com") + mock_handle.assert_called_once() + + def test_map_unknown_error_raises(self, mocker: MockerFixture) -> None: + app = FirecrawlApp(api_key="fc-key", base_url="https://custom.firecrawl.dev") + mocker.patch("core.helper.ssrf_proxy.post", return_value=_response(418, text="teapot")) + + with pytest.raises(Exception, match="Failed to start map job. Status code: 418"): + app.map("https://example.com") + + def test_check_crawl_status_completed_with_data(self, mocker: MockerFixture) -> None: + app = FirecrawlApp(api_key="fc-key", base_url="https://custom.firecrawl.dev") + payload = { + "status": "completed", + "total": 2, + "completed": 2, + "data": [ + { + "metadata": {"title": "a", "description": "desc-a", "sourceURL": "https://a"}, + "markdown": "m-a", + }, + { + "metadata": {"title": "b", "description": "desc-b", "sourceURL": "https://b"}, + "markdown": "m-b", + }, + {"metadata": {"title": "skip"}}, + ], + } + mocker.patch("core.helper.ssrf_proxy.get", return_value=_response(200, payload)) + + save_calls: list[tuple[str, bytes]] = [] + delete_calls: list[str] = [] + + mock_storage = MagicMock() + mock_storage.exists.return_value = True + mock_storage.delete.side_effect = lambda key: delete_calls.append(key) + mock_storage.save.side_effect = lambda key, data: save_calls.append((key, data)) + mocker.patch.object(firecrawl_module, "storage", mock_storage) + + result = app.check_crawl_status("job-42") + + assert result["status"] == "completed" + assert result["total"] == 2 + assert result["current"] == 2 + assert len(result["data"]) == 2 + assert delete_calls == ["website_files/job-42.txt"] + assert len(save_calls) == 1 + assert save_calls[0][0] == "website_files/job-42.txt" + + def test_check_crawl_status_completed_with_zero_total_raises(self, mocker: MockerFixture) -> None: + app = FirecrawlApp(api_key="fc-key", base_url="https://custom.firecrawl.dev") + mocker.patch( + "core.helper.ssrf_proxy.get", + return_value=_response(200, {"status": "completed", "total": 0, "data": []}), + ) + + with pytest.raises(Exception, match="No page found"): + app.check_crawl_status("job-1") + + def test_check_crawl_status_completed_with_null_total_raises(self, mocker: MockerFixture) -> None: + app = FirecrawlApp(api_key="fc-key", base_url="https://custom.firecrawl.dev") + mocker.patch( + "core.helper.ssrf_proxy.get", + return_value=_response(200, {"status": "completed", "total": None, "data": []}), + ) + + with pytest.raises(Exception, match="No page found"): + app.check_crawl_status("job-1") + + def test_check_crawl_status_non_completed(self, mocker: MockerFixture) -> None: + app = FirecrawlApp(api_key="fc-key", base_url="https://custom.firecrawl.dev") + payload = {"status": "processing", "total": 5, "completed": 1, "data": list[object]()} + mocker.patch("core.helper.ssrf_proxy.get", return_value=_response(200, payload)) + + assert app.check_crawl_status("job-1") == { + "status": "processing", + "total": 5, + "current": 1, + "data": [], + } + + def test_check_crawl_status_non_200_uses_error_handler(self, mocker: MockerFixture) -> None: + app = FirecrawlApp(api_key="fc-key", base_url="https://custom.firecrawl.dev") + mock_handle = mocker.patch.object(app, "_handle_error", side_effect=Exception("crawl error")) + mocker.patch("core.helper.ssrf_proxy.get", return_value=_response(500, {"error": "server"})) + + with pytest.raises(Exception, match="crawl error"): + app.check_crawl_status("job-1") + mock_handle.assert_called_once() + + def test_check_crawl_status_save_failure_raises(self, mocker: MockerFixture) -> None: + app = FirecrawlApp(api_key="fc-key", base_url="https://custom.firecrawl.dev") + payload = { + "status": "completed", + "total": 1, + "completed": 1, + "data": [{"metadata": {"title": "a", "sourceURL": "https://a"}, "markdown": "m-a"}], + } + mocker.patch("core.helper.ssrf_proxy.get", return_value=_response(200, payload)) + + mock_storage = MagicMock() + mock_storage.exists.return_value = False + mock_storage.save.side_effect = RuntimeError("save failed") + mocker.patch.object(firecrawl_module, "storage", mock_storage) + + with pytest.raises(Exception, match="Error saving crawl data"): + app.check_crawl_status("job-err") + + def test_check_crawl_status_follows_pagination(self, mocker: MockerFixture) -> None: + """When status is completed and next is present, follow pagination to collect all pages.""" + app = FirecrawlApp(api_key="fc-key", base_url="https://custom.firecrawl.dev") + page1 = { + "status": "completed", + "total": 3, + "completed": 3, + "next": "https://custom.firecrawl.dev/v2/crawl/job-42?skip=1", + "data": [{"metadata": {"title": "p1", "description": "", "sourceURL": "https://p1"}, "markdown": "m1"}], + } + page2 = { + "status": "completed", + "total": 3, + "completed": 3, + "next": "https://custom.firecrawl.dev/v2/crawl/job-42?skip=2", + "data": [{"metadata": {"title": "p2", "description": "", "sourceURL": "https://p2"}, "markdown": "m2"}], + } + page3 = { + "status": "completed", + "total": 3, + "completed": 3, + "data": [{"metadata": {"title": "p3", "description": "", "sourceURL": "https://p3"}, "markdown": "m3"}], + } + mocker.patch( + "core.helper.ssrf_proxy.get", + side_effect=[_response(200, page1), _response(200, page2), _response(200, page3)], + ) + mock_storage = MagicMock() + mock_storage.exists.return_value = False + mocker.patch.object(firecrawl_module, "storage", mock_storage) + + result = app.check_crawl_status("job-42") + + assert result["status"] == "completed" + assert result["total"] == 3 + assert len(result["data"]) == 3 + assert [d["title"] for d in result["data"]] == ["p1", "p2", "p3"] + + def test_check_crawl_status_pagination_error_raises(self, mocker: MockerFixture) -> None: + """An error while fetching a paginated page raises an exception; no partial data is returned.""" + app = FirecrawlApp(api_key="fc-key", base_url="https://custom.firecrawl.dev") + page1 = { + "status": "completed", + "total": 2, + "completed": 2, + "next": "https://custom.firecrawl.dev/v2/crawl/job-99?skip=1", + "data": [{"metadata": {"title": "p1", "description": "", "sourceURL": "https://p1"}, "markdown": "m1"}], + } + mocker.patch( + "core.helper.ssrf_proxy.get", + side_effect=[_response(200, page1), _response(500, {"error": "server error"})], + ) + + with pytest.raises(Exception, match="fetch next crawl page"): + app.check_crawl_status("job-99") + + def test_check_crawl_status_pagination_capped_at_total(self, mocker: MockerFixture) -> None: + """Pagination stops once pages_processed reaches total, even if next is present.""" + app = FirecrawlApp(api_key="fc-key", base_url="https://custom.firecrawl.dev") + # total=1: only the first page should be processed; next must not be followed + page1 = { + "status": "completed", + "total": 1, + "completed": 1, + "next": "https://custom.firecrawl.dev/v2/crawl/job-cap?skip=1", + "data": [{"metadata": {"title": "p1", "description": "", "sourceURL": "https://p1"}, "markdown": "m1"}], + } + mock_get = mocker.patch("core.helper.ssrf_proxy.get", return_value=_response(200, page1)) + mock_storage = MagicMock() + mock_storage.exists.return_value = False + mocker.patch.object(firecrawl_module, "storage", mock_storage) + + result = app.check_crawl_status("job-cap") + + assert len(result["data"]) == 1 + mock_get.assert_called_once() # initial fetch only; next URL is not followed due to cap + + def test_extract_common_fields_and_status_formatter(self) -> None: + app = FirecrawlApp(api_key="fc-key", base_url="https://custom.firecrawl.dev") + + fields = app._extract_common_fields( + {"metadata": {"title": "t", "description": "d", "sourceURL": "u"}, "markdown": "m"} + ) + assert fields == {"title": "t", "description": "d", "source_url": "u", "markdown": "m"} + + status = app._format_crawl_status_response("completed", {"total": 1, "completed": 1}, [fields]) + assert status == {"status": "completed", "total": 1, "current": 1, "data": [fields]} + + def test_post_and_get_request_retry_logic(self, mocker: MockerFixture) -> None: + app = FirecrawlApp(api_key="fc-key", base_url="https://custom.firecrawl.dev") + sleep_mock = mocker.patch.object(firecrawl_module.time, "sleep") + + resp_502_a = _response(502) + resp_502_b = _response(502) + resp_200 = _response(200) + + mocker.patch("core.helper.ssrf_proxy.post", side_effect=[resp_502_a, resp_200]) + post_result = app._post_request("u", {"x": 1}, {"h": 1}, retries=3, backoff_factor=0.5) + assert post_result is resp_200 + + mocker.patch("core.helper.ssrf_proxy.get", side_effect=[resp_502_b, _response(200)]) + get_result = app._get_request("u", {"h": 1}, retries=3, backoff_factor=0.25) + assert get_result.status_code == 200 + + assert sleep_mock.call_count == 2 + + def test_post_and_get_request_return_last_502(self, mocker: MockerFixture) -> None: + app = FirecrawlApp(api_key="fc-key", base_url="https://custom.firecrawl.dev") + sleep_mock = mocker.patch.object(firecrawl_module.time, "sleep") + + last_post = _response(502) + mocker.patch("core.helper.ssrf_proxy.post", side_effect=[_response(502), last_post]) + assert app._post_request("u", {}, {}, retries=2).status_code == 502 + + last_get = _response(502) + mocker.patch("core.helper.ssrf_proxy.get", side_effect=[_response(502), last_get]) + assert app._get_request("u", {}, retries=2).status_code == 502 + + assert sleep_mock.call_count == 4 + + def test_handle_error_with_json_and_plain_text(self) -> None: + app = FirecrawlApp(api_key="fc-key", base_url="https://custom.firecrawl.dev") + + json_error = _response(400, {"message": "bad request"}) + with pytest.raises(Exception, match="bad request"): + app._handle_error(json_error, "run task") + + non_json = MagicMock() + non_json.status_code = 400 + non_json.text = "plain error" + non_json.json.side_effect = json.JSONDecodeError("bad", "x", 0) + + with pytest.raises(Exception, match="plain error"): + app._handle_error(non_json, "run task") + + def test_search_success(self, mocker: MockerFixture) -> None: + app = FirecrawlApp(api_key="fc-key", base_url="https://custom.firecrawl.dev") + mocker.patch( + "core.helper.ssrf_proxy.post", return_value=_response(200, {"success": True, "data": [{"url": "x"}]}) + ) + assert app.search("python")["success"] is True + + def test_search_warning_failure(self, mocker: MockerFixture) -> None: + app = FirecrawlApp(api_key="fc-key", base_url="https://custom.firecrawl.dev") + mocker.patch( + "core.helper.ssrf_proxy.post", return_value=_response(200, {"success": False, "warning": "bad search"}) + ) + with pytest.raises(Exception, match="bad search"): + app.search("python") + + def test_search_known_http_error(self, mocker: MockerFixture) -> None: + app = FirecrawlApp(api_key="fc-key", base_url="https://custom.firecrawl.dev") + mock_handle = mocker.patch.object(app, "_handle_error", side_effect=Exception("search error")) + mocker.patch("core.helper.ssrf_proxy.post", return_value=_response(408, {"error": "timeout"})) + with pytest.raises(Exception, match="search error"): + app.search("python") + mock_handle.assert_called_once() + + def test_search_unknown_http_error(self, mocker: MockerFixture) -> None: + app = FirecrawlApp(api_key="fc-key", base_url="https://custom.firecrawl.dev") + mocker.patch("core.helper.ssrf_proxy.post", return_value=_response(418, text="teapot")) + with pytest.raises(Exception, match="Failed to perform search. Status code: 418"): + app.search("python") + + class TestFirecrawlWebExtractor: + def test_extract_crawl_mode_returns_document(self, mocker: MockerFixture) -> None: + mocker.patch( + "core.rag.extractor.firecrawl.firecrawl_web_extractor.WebsiteService.get_crawl_url_data", + return_value={ + "markdown": "crawl content", + "source_url": "https://example.com", + "description": "desc", + "title": "title", + }, + ) + + extractor = FirecrawlWebExtractor( + "https://example.com", + "job-1", + "tenant-1", + mode="crawl", + website_service=WebsiteService(providers=create_autospec(DatasourceProviderService, instance=True)), + ) + docs = extractor.extract() + + assert len(docs) == 1 + assert docs[0].page_content == "crawl content" + assert docs[0].metadata["source_url"] == "https://example.com" + + def test_extract_crawl_mode_with_missing_data_returns_empty(self, mocker: MockerFixture) -> None: + mocker.patch( + "core.rag.extractor.firecrawl.firecrawl_web_extractor.WebsiteService.get_crawl_url_data", + return_value=None, + ) + + extractor = FirecrawlWebExtractor( + "https://example.com", + "job-1", + "tenant-1", + mode="crawl", + website_service=WebsiteService(providers=create_autospec(DatasourceProviderService, instance=True)), + ) + assert extractor.extract() == [] + + def test_extract_scrape_mode_returns_document(self, mocker: MockerFixture) -> None: + mock_scrape = mocker.patch( + "core.rag.extractor.firecrawl.firecrawl_web_extractor.WebsiteService.get_scrape_url_data", + return_value={ + "markdown": "scrape content", + "source_url": "https://example.com", + "description": "desc", + "title": "title", + }, + ) + + extractor = FirecrawlWebExtractor( + "https://example.com", + "job-1", + "tenant-1", + mode="scrape", + only_main_content=False, + website_service=WebsiteService(providers=create_autospec(DatasourceProviderService, instance=True)), + ) + docs = extractor.extract() + + assert len(docs) == 1 + assert docs[0].page_content == "scrape content" + mock_scrape.assert_called_once_with("firecrawl", "https://example.com", "tenant-1", False) + + def test_extract_unknown_mode_returns_empty(self) -> None: + extractor = FirecrawlWebExtractor( + "https://example.com", + "job-1", + "tenant-1", + mode="unknown", + website_service=WebsiteService(providers=create_autospec(DatasourceProviderService, instance=True)), + ) + assert extractor.extract() == [] + + +class TestCrawlerSsrfCases: + @staticmethod + @pytest.fixture + def proxy_transport( + monkeypatch: pytest.MonkeyPatch, + ) -> Iterator[Callable[[Callable[[httpx.Request], httpx.Response]], list[httpx.Request]]]: + with ExitStack() as stack: + + def install(handler: Callable[[httpx.Request], httpx.Response]) -> list[httpx.Request]: + requests: list[httpx.Request] = [] + + def respond(request: httpx.Request) -> httpx.Response: + requests.append(request) + return handler(request) + + client = stack.enter_context(httpx.Client(transport=httpx.MockTransport(respond))) + monkeypatch.setattr(ssrf_proxy, "_get_ssrf_client", lambda _verify: client) + + def reject_direct_request(*_args: object, **_kwargs: object) -> Never: + pytest.fail("Crawler request bypassed the shared SSRF client") + + monkeypatch.setattr(httpx, "get", reject_direct_request) + monkeypatch.setattr(httpx, "post", reject_direct_request) + monkeypatch.setattr(httpx, "Client", reject_direct_request) + return requests + + yield install + + @staticmethod + @pytest.mark.parametrize( + "operation", + [ + "firecrawl_post", + "firecrawl_get", + "watercrawl_post", + "watercrawl_get", + "watercrawl_stream", + "watercrawl_download", + ], + ) + def test_crawler_requests_honor_proxy_rejection( + operation: str, + proxy_transport: Callable[[Callable[[httpx.Request], httpx.Response]], list[httpx.Request]], + ) -> None: + rejected = httpx.Response(403, headers={"Server": "squid"}) + requests = proxy_transport(lambda _request: rejected) + firecrawl = FirecrawlApp(api_key="test-key", base_url="http://127.0.0.1") + watercrawl = WaterCrawlAPIClient(api_key="test-key", base_url="http://127.0.0.1") + operations: dict[str, Callable[[], object]] = { + "firecrawl_post": lambda: firecrawl.scrape_url("https://example.com"), + "firecrawl_get": lambda: firecrawl.check_crawl_status("job"), + "watercrawl_post": lambda: watercrawl.create_crawl_request(url="https://example.com"), + "watercrawl_get": lambda: watercrawl.get_crawl_request_results("job"), + "watercrawl_stream": lambda: list(watercrawl.monitor_crawl_request("job")), + "watercrawl_download": lambda: watercrawl.download_result({"result": "http://127.0.0.1/result.json"}), + } + + with pytest.raises(ToolSSRFError, match="blocked by SSRF protection"): + operations[operation]() + + assert len(requests) == 1 + assert requests[0].url.host == "127.0.0.1" + assert rejected.is_closed + + @staticmethod + def test_firecrawl_pagination_honors_proxy_rejection( + proxy_transport: Callable[[Callable[[httpx.Request], httpx.Response]], list[httpx.Request]], + ) -> None: + def respond(request: httpx.Request) -> httpx.Response: + if request.url.host == "crawler.example.com": + return httpx.Response( + 200, json={"status": "completed", "total": 2, "data": [], "next": "http://127.0.0.1/next"} + ) + return httpx.Response(403, headers={"Via": "squid"}) + + requests = proxy_transport(respond) + firecrawl = FirecrawlApp(api_key="test-key", base_url="https://crawler.example.com") + + with pytest.raises(ToolSSRFError, match="blocked by SSRF protection"): + firecrawl.check_crawl_status("job") + + assert [str(request.url) for request in requests] == [ + "https://crawler.example.com/v2/crawl/job", + "http://127.0.0.1/next", + ] + + @staticmethod + @pytest.mark.parametrize("provider", ["firecrawl", "watercrawl"]) + def test_credential_validation_honors_proxy_rejection( + provider: str, + proxy_transport: Callable[[Callable[[httpx.Request], httpx.Response]], list[httpx.Request]], + ) -> None: + requests = proxy_transport(lambda _request: httpx.Response(403, headers={"Server": "squid"})) + credentials = DataSourceApiKeyAuthCredentials( + "bearer" if provider == "firecrawl" else "x-api-key", "test-key", {"base_url": "http://127.0.0.1"} + ) + auth = FirecrawlAuth(credentials) if provider == "firecrawl" else WatercrawlAuth(credentials) + + with pytest.raises(ToolSSRFError, match="blocked by SSRF protection"): + auth.validate_credentials() + + assert len(requests) == 1 + assert requests[0].extensions["timeout"]["read"] == 10.0 + + @staticmethod + def test_watercrawl_stream_is_lazy_and_closed_after_events( + proxy_transport: Callable[[Callable[[httpx.Request], httpx.Response]], list[httpx.Request]], + ) -> None: + response = httpx.Response( + 200, + headers={"Content-Type": "text/event-stream"}, + stream=httpx.ByteStream(b'data: {"type": "result", "data": {"markdown": "page"}}\n\n'), + ) + requests = proxy_transport(lambda _request: response) + client = WaterCrawlAPIClient(api_key="test-key") + events = client.monitor_crawl_request("job") + + assert requests == [] + assert next(events) == {"type": "result", "data": {"markdown": "page"}} + events.close() + assert response.is_closed + assert requests[0].extensions["timeout"]["read"] is None + + @staticmethod + @pytest.mark.parametrize("status_code", [200, 500]) + def test_watercrawl_download_closes_response( + status_code: int, + proxy_transport: Callable[[Callable[[httpx.Request], httpx.Response]], list[httpx.Request]], + ) -> None: + response = httpx.Response(status_code, json={"markdown": "page"}) + requests = proxy_transport(lambda _request: response) + client = WaterCrawlAPIClient(api_key="test-key") + + if status_code == 200: + assert client.download_result({"result": "https://results.example.com/page.json"}) == { + "result": {"markdown": "page"} + } + else: + with pytest.raises(httpx.HTTPStatusError): + client.download_result({"result": "https://results.example.com/page.json"}) + + assert len(requests) == 1 + assert "X-API-Key" not in requests[0].headers + assert response.is_closed diff --git a/api/tests/unit_tests/core/rag/extractor/test_crawler_ssrf.py b/api/tests/unit_tests/core/rag/extractor/test_crawler_ssrf.py deleted file mode 100644 index 81e8ea05b2d..00000000000 --- a/api/tests/unit_tests/core/rag/extractor/test_crawler_ssrf.py +++ /dev/null @@ -1,158 +0,0 @@ -from collections.abc import Callable, Iterator -from contextlib import ExitStack -from typing import Never - -import httpx -import pytest - -from core.helper import ssrf_proxy -from core.rag.extractor.firecrawl.firecrawl_app import FirecrawlApp -from core.rag.extractor.watercrawl.client import WaterCrawlAPIClient -from core.tools.errors import ToolSSRFError -from services.data_source.auth.firecrawl.firecrawl import FirecrawlAuth -from services.data_source.auth.watercrawl.watercrawl import WatercrawlAuth -from services.data_source.entities.api_key_auth import DataSourceApiKeyAuthCredentials - - -@pytest.fixture -def proxy_transport( - monkeypatch: pytest.MonkeyPatch, -) -> Iterator[Callable[[Callable[[httpx.Request], httpx.Response]], list[httpx.Request]]]: - with ExitStack() as stack: - - def install(handler: Callable[[httpx.Request], httpx.Response]) -> list[httpx.Request]: - requests: list[httpx.Request] = [] - - def respond(request: httpx.Request) -> httpx.Response: - requests.append(request) - return handler(request) - - client = stack.enter_context(httpx.Client(transport=httpx.MockTransport(respond))) - monkeypatch.setattr(ssrf_proxy, "_get_ssrf_client", lambda _verify: client) - - def reject_direct_request(*_args: object, **_kwargs: object) -> Never: - pytest.fail("Crawler request bypassed the shared SSRF client") - - monkeypatch.setattr(httpx, "get", reject_direct_request) - monkeypatch.setattr(httpx, "post", reject_direct_request) - monkeypatch.setattr(httpx, "Client", reject_direct_request) - return requests - - yield install - - -@pytest.mark.parametrize( - "operation", - [ - "firecrawl_post", - "firecrawl_get", - "watercrawl_post", - "watercrawl_get", - "watercrawl_stream", - "watercrawl_download", - ], -) -def test_crawler_requests_honor_proxy_rejection( - operation: str, - proxy_transport: Callable[[Callable[[httpx.Request], httpx.Response]], list[httpx.Request]], -) -> None: - rejected = httpx.Response(403, headers={"Server": "squid"}) - requests = proxy_transport(lambda _request: rejected) - firecrawl = FirecrawlApp(api_key="test-key", base_url="http://127.0.0.1") - watercrawl = WaterCrawlAPIClient(api_key="test-key", base_url="http://127.0.0.1") - operations: dict[str, Callable[[], object]] = { - "firecrawl_post": lambda: firecrawl.scrape_url("https://example.com"), - "firecrawl_get": lambda: firecrawl.check_crawl_status("job"), - "watercrawl_post": lambda: watercrawl.create_crawl_request(url="https://example.com"), - "watercrawl_get": lambda: watercrawl.get_crawl_request_results("job"), - "watercrawl_stream": lambda: list(watercrawl.monitor_crawl_request("job")), - "watercrawl_download": lambda: watercrawl.download_result({"result": "http://127.0.0.1/result.json"}), - } - - with pytest.raises(ToolSSRFError, match="blocked by SSRF protection"): - operations[operation]() - - assert len(requests) == 1 - assert requests[0].url.host == "127.0.0.1" - assert rejected.is_closed - - -def test_firecrawl_pagination_honors_proxy_rejection( - proxy_transport: Callable[[Callable[[httpx.Request], httpx.Response]], list[httpx.Request]], -) -> None: - def respond(request: httpx.Request) -> httpx.Response: - if request.url.host == "crawler.example.com": - return httpx.Response( - 200, json={"status": "completed", "total": 2, "data": [], "next": "http://127.0.0.1/next"} - ) - return httpx.Response(403, headers={"Via": "squid"}) - - requests = proxy_transport(respond) - firecrawl = FirecrawlApp(api_key="test-key", base_url="https://crawler.example.com") - - with pytest.raises(ToolSSRFError, match="blocked by SSRF protection"): - firecrawl.check_crawl_status("job") - - assert [str(request.url) for request in requests] == [ - "https://crawler.example.com/v2/crawl/job", - "http://127.0.0.1/next", - ] - - -@pytest.mark.parametrize("provider", ["firecrawl", "watercrawl"]) -def test_credential_validation_honors_proxy_rejection( - provider: str, - proxy_transport: Callable[[Callable[[httpx.Request], httpx.Response]], list[httpx.Request]], -) -> None: - requests = proxy_transport(lambda _request: httpx.Response(403, headers={"Server": "squid"})) - credentials = DataSourceApiKeyAuthCredentials( - "bearer" if provider == "firecrawl" else "x-api-key", "test-key", {"base_url": "http://127.0.0.1"} - ) - auth = FirecrawlAuth(credentials) if provider == "firecrawl" else WatercrawlAuth(credentials) - - with pytest.raises(ToolSSRFError, match="blocked by SSRF protection"): - auth.validate_credentials() - - assert len(requests) == 1 - assert requests[0].extensions["timeout"]["read"] == 10.0 - - -def test_watercrawl_stream_is_lazy_and_closed_after_events( - proxy_transport: Callable[[Callable[[httpx.Request], httpx.Response]], list[httpx.Request]], -) -> None: - response = httpx.Response( - 200, - headers={"Content-Type": "text/event-stream"}, - stream=httpx.ByteStream(b'data: {"type": "result", "data": {"markdown": "page"}}\n\n'), - ) - requests = proxy_transport(lambda _request: response) - client = WaterCrawlAPIClient(api_key="test-key") - events = client.monitor_crawl_request("job") - - assert requests == [] - assert next(events) == {"type": "result", "data": {"markdown": "page"}} - events.close() - assert response.is_closed - assert requests[0].extensions["timeout"]["read"] is None - - -@pytest.mark.parametrize("status_code", [200, 500]) -def test_watercrawl_download_closes_response( - status_code: int, - proxy_transport: Callable[[Callable[[httpx.Request], httpx.Response]], list[httpx.Request]], -) -> None: - response = httpx.Response(status_code, json={"markdown": "page"}) - requests = proxy_transport(lambda _request: response) - client = WaterCrawlAPIClient(api_key="test-key") - - if status_code == 200: - assert client.download_result({"result": "https://results.example.com/page.json"}) == { - "result": {"markdown": "page"} - } - else: - with pytest.raises(httpx.HTTPStatusError): - client.download_result({"result": "https://results.example.com/page.json"}) - - assert len(requests) == 1 - assert "X-API-Key" not in requests[0].headers - assert response.is_closed diff --git a/api/tests/unit_tests/core/rag/extractor/test_excel_extractor.py b/api/tests/unit_tests/core/rag/extractor/test_excel_extractor.py index ddd7907727a..5b87121f12b 100644 --- a/api/tests/unit_tests/core/rag/extractor/test_excel_extractor.py +++ b/api/tests/unit_tests/core/rag/extractor/test_excel_extractor.py @@ -9,7 +9,7 @@ import core.rag.extractor.excel_extractor as excel_module from core.rag.extractor.excel_extractor import ExcelExtractor from models.base import TypeBase from models.model import UploadFile -from tests.unit_tests.config_override import apply_config_overrides +from tests.helpers.unit.config_override import apply_config_overrides @pytest.fixture diff --git a/api/tests/unit_tests/core/rag/extractor/test_extract_processor.py b/api/tests/unit_tests/core/rag/extractor/test_extract_processor.py index 344ee8d9d55..ee8681babbe 100644 --- a/api/tests/unit_tests/core/rag/extractor/test_extract_processor.py +++ b/api/tests/unit_tests/core/rag/extractor/test_extract_processor.py @@ -12,8 +12,8 @@ from core.rag.extractor.entity.extract_setting import ExtractSetting, NotionInfo from core.rag.extractor.extract_processor import ExtractProcessor from core.rag.models.document import Document from models.model import UploadFile -from tests.unit_tests.config_override import apply_config_overrides -from tests.unit_tests.model_factories import make_upload_file +from tests.helpers.unit.config_override import apply_config_overrides +from tests.helpers.unit.model_factories import make_upload_file def _upload_file(*, key: str, file_id: str = "upload-file-1") -> UploadFile: diff --git a/api/tests/unit_tests/core/rag/extractor/test_notion_extractor.py b/api/tests/unit_tests/core/rag/extractor/test_notion_extractor.py index 1bbeb00f260..ef4e0baf130 100644 --- a/api/tests/unit_tests/core/rag/extractor/test_notion_extractor.py +++ b/api/tests/unit_tests/core/rag/extractor/test_notion_extractor.py @@ -1,5 +1,3 @@ -"""Unit tests for Notion extraction, HTTP parsing, and persisted metadata updates.""" - import json import uuid from collections.abc import Iterator @@ -19,37 +17,10 @@ from models.base import TypeBase from models.dataset import Document as DocumentModel from models.enums import DataSourceType, DocumentCreatedFrom, IndexingStatus - -@pytest.fixture -def persisted_document( - monkeypatch: pytest.MonkeyPatch, - sqlite_engine: Engine, -) -> Iterator[tuple[sessionmaker[Session], DocumentModel]]: - """Persist a Notion document and bind the extractor's ``db.session`` to SQLite.""" - TypeBase.metadata.create_all(sqlite_engine, tables=[TypeBase.metadata.tables[DocumentModel.__tablename__]]) - session_maker = sessionmaker(bind=sqlite_engine, expire_on_commit=False) - document = DocumentModel( - id=str(uuid.uuid4()), - tenant_id=str(uuid.uuid4()), - dataset_id=str(uuid.uuid4()), - position=1, - data_source_type=DataSourceType.NOTION_IMPORT, - data_source_info=json.dumps({"source": "notion", "last_edited_time": "2025-01-01T00:00:00.000Z"}), - batch="batch", - name="Notion page", - created_from=DocumentCreatedFrom.API, - created_by=str(uuid.uuid4()), - indexing_status=IndexingStatus.COMPLETED, - doc_form=IndexStructureType.PARAGRAPH_INDEX, - ) - with session_maker() as sqlite_session: - sqlite_session.add(document) - sqlite_session.commit() - monkeypatch.setattr(notion_extractor, "db", SimpleNamespace(session=sqlite_session)) - yield session_maker, document +"""Unit tests for Notion extraction, HTTP parsing, and persisted metadata updates.""" -def _mock_response(data: object, status_code: int = 200, text: str = "") -> mock.Mock: +def _notion_extractor_mock_response(data: object, status_code: int = 200, text: str = "") -> mock.Mock: response = mock.Mock() response.status_code = status_code response.text = text @@ -57,37 +28,52 @@ def _mock_response(data: object, status_code: int = 200, text: str = "") -> mock return response -class TestNotionExtractorInitAndPublicMethods: - def test_init_with_explicit_token(self) -> None: - extractor = notion_extractor.NotionExtractor( - notion_workspace_id="ws", - notion_obj_id="obj", - notion_page_type="page", - tenant_id="tenant", - notion_access_token="token", - notion_token_loader=lambda: "token", +class TestNotionExtractorCases: + @staticmethod + @pytest.fixture + def persisted_document( + monkeypatch: pytest.MonkeyPatch, + sqlite_engine: Engine, + ) -> Iterator[tuple[sessionmaker[Session], DocumentModel]]: + """Persist a Notion document and bind the extractor's ``db.session`` to SQLite.""" + TypeBase.metadata.create_all(sqlite_engine, tables=[TypeBase.metadata.tables[DocumentModel.__tablename__]]) + session_maker = sessionmaker(bind=sqlite_engine, expire_on_commit=False) + document = DocumentModel( + id=str(uuid.uuid4()), + tenant_id=str(uuid.uuid4()), + dataset_id=str(uuid.uuid4()), + position=1, + data_source_type=DataSourceType.NOTION_IMPORT, + data_source_info=json.dumps({"source": "notion", "last_edited_time": "2025-01-01T00:00:00.000Z"}), + batch="batch", + name="Notion page", + created_from=DocumentCreatedFrom.API, + created_by=str(uuid.uuid4()), + indexing_status=IndexingStatus.COMPLETED, + doc_form=IndexStructureType.PARAGRAPH_INDEX, ) + with session_maker() as sqlite_session: + sqlite_session.add(document) + sqlite_session.commit() + monkeypatch.setattr(notion_extractor, "db", SimpleNamespace(session=sqlite_session)) + yield session_maker, document - assert extractor._notion_access_token == "token" + class TestNotionExtractorInitAndPublicMethods: + def test_init_with_explicit_token(self) -> None: + extractor = notion_extractor.NotionExtractor( + notion_workspace_id="ws", + notion_obj_id="obj", + notion_page_type="page", + tenant_id="tenant", + notion_access_token="token", + notion_token_loader=lambda: "token", + ) - def test_init_uses_injected_token_loader(self) -> None: - loader = mock.Mock(return_value="resolved-token") - extractor = notion_extractor.NotionExtractor( - notion_workspace_id="ws", - notion_obj_id="obj", - notion_page_type="page", - tenant_id="tenant", - credential_id="cred", - notion_token_loader=loader, - ) + assert extractor._notion_access_token == "token" - assert extractor._notion_access_token == "resolved-token" - loader.assert_called_once_with() - - def test_init_propagates_missing_token_from_loader(self) -> None: - loader = mock.Mock(side_effect=ValueError("Must specify `integration_token`")) - with pytest.raises(ValueError, match="Must specify `integration_token`"): - notion_extractor.NotionExtractor( + def test_init_uses_injected_token_loader(self) -> None: + loader = mock.Mock(return_value="resolved-token") + extractor = notion_extractor.NotionExtractor( notion_workspace_id="ws", notion_obj_id="obj", notion_page_type="page", @@ -96,607 +82,718 @@ class TestNotionExtractorInitAndPublicMethods: notion_token_loader=loader, ) - def test_extract_updates_last_edited_and_loads_documents(self, monkeypatch: pytest.MonkeyPatch) -> None: - extractor = notion_extractor.NotionExtractor( - notion_workspace_id="ws", - notion_obj_id="obj", - notion_page_type="page", - tenant_id="tenant", - notion_access_token="token", - notion_token_loader=lambda: "token", - ) + assert extractor._notion_access_token == "resolved-token" + loader.assert_called_once_with() - update_mock = mock.Mock() - load_mock = mock.Mock(return_value=[SimpleNamespace(page_content="doc")]) - monkeypatch.setattr(extractor, "update_last_edited_time", update_mock) - monkeypatch.setattr(extractor, "_load_data_as_documents", load_mock) + def test_init_propagates_missing_token_from_loader(self) -> None: + loader = mock.Mock(side_effect=ValueError("Must specify `integration_token`")) + with pytest.raises(ValueError, match="Must specify `integration_token`"): + notion_extractor.NotionExtractor( + notion_workspace_id="ws", + notion_obj_id="obj", + notion_page_type="page", + tenant_id="tenant", + credential_id="cred", + notion_token_loader=loader, + ) - docs = extractor.extract() + def test_extract_updates_last_edited_and_loads_documents(self, monkeypatch: pytest.MonkeyPatch) -> None: + extractor = notion_extractor.NotionExtractor( + notion_workspace_id="ws", + notion_obj_id="obj", + notion_page_type="page", + tenant_id="tenant", + notion_access_token="token", + notion_token_loader=lambda: "token", + ) - update_mock.assert_called_once_with(None) - load_mock.assert_called_once_with("obj", "page") - assert len(docs) == 1 + update_mock = mock.Mock() + load_mock = mock.Mock(return_value=[SimpleNamespace(page_content="doc")]) + monkeypatch.setattr(extractor, "update_last_edited_time", update_mock) + monkeypatch.setattr(extractor, "_load_data_as_documents", load_mock) - def test_load_data_as_documents_page_database_and_invalid(self, monkeypatch: pytest.MonkeyPatch) -> None: - extractor = notion_extractor.NotionExtractor( - notion_workspace_id="ws", - notion_obj_id="obj", - notion_page_type="page", - tenant_id="tenant", - notion_access_token="token", - notion_token_loader=lambda: "token", - ) + docs = extractor.extract() - monkeypatch.setattr(extractor, "_get_notion_block_data", lambda _: ["line1", "line2"]) - page_docs = extractor._load_data_as_documents("page-id", "page") - assert page_docs[0].page_content == "line1\nline2" + update_mock.assert_called_once_with(None) + load_mock.assert_called_once_with("obj", "page") + assert len(docs) == 1 - monkeypatch.setattr(extractor, "_get_notion_database_data", lambda _: [SimpleNamespace(page_content="db")]) - db_docs = extractor._load_data_as_documents("db-id", "database") - assert db_docs[0].page_content == "db" + def test_load_data_as_documents_page_database_and_invalid(self, monkeypatch: pytest.MonkeyPatch) -> None: + extractor = notion_extractor.NotionExtractor( + notion_workspace_id="ws", + notion_obj_id="obj", + notion_page_type="page", + tenant_id="tenant", + notion_access_token="token", + notion_token_loader=lambda: "token", + ) - with pytest.raises(ValueError, match="notion page type not supported"): - extractor._load_data_as_documents("obj", "unsupported") + monkeypatch.setattr(extractor, "_get_notion_block_data", lambda _: ["line1", "line2"]) + page_docs = extractor._load_data_as_documents("page-id", "page") + assert page_docs[0].page_content == "line1\nline2" + monkeypatch.setattr(extractor, "_get_notion_database_data", lambda _: [SimpleNamespace(page_content="db")]) + db_docs = extractor._load_data_as_documents("db-id", "database") + assert db_docs[0].page_content == "db" -class TestNotionDatabase: - def test_get_notion_database_data_parses_property_types_and_pagination(self, mocker: MockerFixture) -> None: - extractor = notion_extractor.NotionExtractor( - notion_workspace_id="ws", - notion_obj_id="obj", - notion_page_type="database", - tenant_id="tenant", - notion_access_token="token", - notion_token_loader=lambda: "token", - ) + with pytest.raises(ValueError, match="notion page type not supported"): + extractor._load_data_as_documents("obj", "unsupported") - first_page = { - "results": [ - { - "properties": { - "tags": { - "type": "multi_select", - "multi_select": [{"name": "A"}, {"name": "B"}], + class TestNotionDatabase: + def test_get_notion_database_data_parses_property_types_and_pagination(self, mocker: MockerFixture) -> None: + extractor = notion_extractor.NotionExtractor( + notion_workspace_id="ws", + notion_obj_id="obj", + notion_page_type="database", + tenant_id="tenant", + notion_access_token="token", + notion_token_loader=lambda: "token", + ) + + first_page = { + "results": [ + { + "properties": { + "tags": { + "type": "multi_select", + "multi_select": [{"name": "A"}, {"name": "B"}], + }, + "title_prop": {"type": "title", "title": [{"plain_text": "Title"}]}, + "empty_title": {"type": "title", "title": list[object]()}, + "rich": {"type": "rich_text", "rich_text": [{"plain_text": "RichText"}]}, + "empty_rich": {"type": "rich_text", "rich_text": list[object]()}, + "select_prop": {"type": "select", "select": {"name": "Selected"}}, + "empty_select": {"type": "select", "select": None}, + "status_prop": {"type": "status", "status": {"name": "Open"}}, + "empty_status": {"type": "status", "status": None}, + "number_prop": {"type": "number", "number": 10}, + "dict_prop": {"type": "date", "date": {"start": "2024-01-01", "end": None}}, }, - "title_prop": {"type": "title", "title": [{"plain_text": "Title"}]}, - "empty_title": {"type": "title", "title": list[object]()}, - "rich": {"type": "rich_text", "rich_text": [{"plain_text": "RichText"}]}, - "empty_rich": {"type": "rich_text", "rich_text": list[object]()}, - "select_prop": {"type": "select", "select": {"name": "Selected"}}, - "empty_select": {"type": "select", "select": None}, - "status_prop": {"type": "status", "status": {"name": "Open"}}, - "empty_status": {"type": "status", "status": None}, - "number_prop": {"type": "number", "number": 10}, - "dict_prop": {"type": "date", "date": {"start": "2024-01-01", "end": None}}, - }, - "url": "https://notion.so/page-1", - } - ], - "has_more": True, - "next_cursor": "cursor-2", - } - second_page = {"results": list[object](), "has_more": False, "next_cursor": None} + "url": "https://notion.so/page-1", + } + ], + "has_more": True, + "next_cursor": "cursor-2", + } + second_page = {"results": list[object](), "has_more": False, "next_cursor": None} - mock_post = mocker.patch("httpx.post", side_effect=[_mock_response(first_page), _mock_response(second_page)]) + mock_post = mocker.patch( + "httpx.post", + side_effect=[_notion_extractor_mock_response(first_page), _notion_extractor_mock_response(second_page)], + ) - docs = extractor._get_notion_database_data("db-1", query_dict={"filter": {"x": 1}}) + docs = extractor._get_notion_database_data("db-1", query_dict={"filter": {"x": 1}}) - assert len(docs) == 1 - content = docs[0].page_content - assert "tags:['A', 'B']" in content - assert "title_prop:Title" in content - assert "rich:RichText" in content - assert "number_prop:10" in content - assert "dict_prop:start:2024-01-01" in content - assert "Row Page URL:https://notion.so/page-1" in content - assert mock_post.call_count == 2 + assert len(docs) == 1 + content = docs[0].page_content + assert "tags:['A', 'B']" in content + assert "title_prop:Title" in content + assert "rich:RichText" in content + assert "number_prop:10" in content + assert "dict_prop:start:2024-01-01" in content + assert "Row Page URL:https://notion.so/page-1" in content + assert mock_post.call_count == 2 - def test_get_notion_database_data_joins_all_rich_text_segments(self, mocker: MockerFixture) -> None: - """Formatted text arrives as multiple segments; all of them must be kept.""" - extractor = notion_extractor.NotionExtractor( - notion_workspace_id="ws", - notion_obj_id="obj", - notion_page_type="database", - tenant_id="tenant", - notion_access_token="token", - notion_token_loader=lambda: "token", - ) + def test_get_notion_database_data_joins_all_rich_text_segments(self, mocker: MockerFixture) -> None: + """Formatted text arrives as multiple segments; all of them must be kept.""" + extractor = notion_extractor.NotionExtractor( + notion_workspace_id="ws", + notion_obj_id="obj", + notion_page_type="database", + tenant_id="tenant", + notion_access_token="token", + notion_token_loader=lambda: "token", + ) - page = { - "results": [ - { - "properties": { - "title_prop": { - "type": "title", - "title": [{"plain_text": "Hello "}, {"plain_text": "world"}], + page = { + "results": [ + { + "properties": { + "title_prop": { + "type": "title", + "title": [{"plain_text": "Hello "}, {"plain_text": "world"}], + }, + "rich": { + "type": "rich_text", + "rich_text": [ + {"plain_text": "first "}, + {"plain_text": "second"}, + {"plain_text": " third"}, + ], + }, }, - "rich": { - "type": "rich_text", - "rich_text": [ - {"plain_text": "first "}, - {"plain_text": "second"}, - {"plain_text": " third"}, - ], - }, - }, - "url": "https://notion.so/page-1", - } - ], - "has_more": False, - "next_cursor": None, - } - - mocker.patch("httpx.post", return_value=_mock_response(page)) - - docs = extractor._get_notion_database_data("db-1") - - assert len(docs) == 1 - content = docs[0].page_content - assert "rich:first second third" in content - assert "title_prop:Hello world" in content - - def test_get_notion_database_data_handles_missing_results_and_empty_content(self, mocker: MockerFixture) -> None: - extractor = notion_extractor.NotionExtractor( - notion_workspace_id="ws", - notion_obj_id="obj", - notion_page_type="database", - tenant_id="tenant", - notion_access_token="token", - notion_token_loader=lambda: "token", - ) - - mocker.patch("httpx.post", return_value=_mock_response({"results": None})) - assert extractor._get_notion_database_data("db-1") == [] - - def test_requests_use_bounded_timeout(self, mocker: MockerFixture) -> None: - """All outbound Notion API calls must carry a bounded timeout so a hanging endpoint cannot block extraction.""" - extractor = notion_extractor.NotionExtractor( - notion_workspace_id="ws", - notion_obj_id="obj", - notion_page_type="database", - tenant_id="tenant", - notion_access_token="token", - notion_token_loader=lambda: "token", - ) - - mock_post = mocker.patch("httpx.post", return_value=_mock_response({"results": None})) - extractor._get_notion_database_data("db-1") - assert mock_post.call_args.kwargs["timeout"] == notion_extractor._REQUEST_TIMEOUT - - mock_request = mocker.patch( - "httpx.request", return_value=_mock_response({"last_edited_time": "2024-01-01T00:00:00.000Z"}) - ) - extractor.get_notion_last_edited_time() - assert mock_request.call_args.kwargs["timeout"] == notion_extractor._REQUEST_TIMEOUT - - def test_get_notion_database_data_requires_access_token(self) -> None: - extractor = notion_extractor.NotionExtractor( - notion_workspace_id="ws", - notion_obj_id="obj", - notion_page_type="database", - tenant_id="tenant", - notion_access_token="token", - notion_token_loader=lambda: "token", - ) - extractor._notion_access_token = None - - with pytest.raises(AssertionError, match="Notion access token is required"): - extractor._get_notion_database_data("db-1") - - -class TestNotionBlocks: - def test_get_notion_block_data_success_with_table_headings_children_and_pagination( - self, mocker: MockerFixture - ) -> None: - extractor = notion_extractor.NotionExtractor( - notion_workspace_id="ws", - notion_obj_id="obj", - notion_page_type="page", - tenant_id="tenant", - notion_access_token="token", - notion_token_loader=lambda: "token", - ) - - first_response = { - "results": [ - {"type": "table", "id": "tbl-1", "has_children": False, "table": dict[str, object]()}, - { - "type": "heading_1", - "id": "h1", - "has_children": False, - "heading_1": {"rich_text": [{"text": {"content": "Heading"}}]}, - }, - { - "type": "paragraph", - "id": "p1", - "has_children": True, - "paragraph": {"rich_text": [{"text": {"content": "Paragraph"}}]}, - }, - { - "type": "child_page", - "id": "cp1", - "has_children": True, - "child_page": {"rich_text": list[object]()}, - }, - ], - "next_cursor": "cursor-2", - } - second_response = { - "results": [ - { - "type": "heading_2", - "id": "h2", - "has_children": False, - "heading_2": {"rich_text": [{"text": {"content": "SubHeading"}}]}, - } - ], - "next_cursor": None, - } - - mocker.patch("httpx.request", side_effect=[_mock_response(first_response), _mock_response(second_response)]) - mocker.patch.object(extractor, "_read_table_rows", return_value="TABLE") - mocker.patch.object(extractor, "_read_block", return_value="CHILD") - - lines = extractor._get_notion_block_data("page-1") - - assert lines[0] == "TABLE\n\n" - assert "# Heading" in lines[1] - assert "Paragraph\nCHILD\n\n" in lines[2] - assert "## SubHeading" in lines[-1] - - def test_get_notion_block_data_handles_http_error_and_invalid_payload(self, mocker: MockerFixture) -> None: - extractor = notion_extractor.NotionExtractor( - notion_workspace_id="ws", - notion_obj_id="obj", - notion_page_type="page", - tenant_id="tenant", - notion_access_token="token", - notion_token_loader=lambda: "token", - ) - - mocker.patch("httpx.request", side_effect=httpx.HTTPError("network")) - with pytest.raises(ValueError, match="Error fetching Notion block data"): - extractor._get_notion_block_data("page-1") - - mocker.patch("httpx.request", return_value=_mock_response({"bad": "payload"}, status_code=200)) - with pytest.raises(ValueError, match="Error fetching Notion block data"): - extractor._get_notion_block_data("page-1") - - mocker.patch("httpx.request", return_value=_mock_response({"results": []}, status_code=500, text="boom")) - with pytest.raises(ValueError, match="Error fetching Notion block data: boom"): - extractor._get_notion_block_data("page-1") - - def test_read_block_supports_heading_table_and_recursion(self, mocker: MockerFixture) -> None: - extractor = notion_extractor.NotionExtractor( - notion_workspace_id="ws", - notion_obj_id="obj", - notion_page_type="page", - tenant_id="tenant", - notion_access_token="token", - notion_token_loader=lambda: "token", - ) - - root_payload = { - "results": [ - { - "type": "heading_2", - "id": "h2", - "has_children": False, - "heading_2": {"rich_text": [{"text": {"content": "Root"}}]}, - }, - { - "type": "paragraph", - "id": "child-block", - "has_children": True, - "paragraph": {"rich_text": [{"text": {"content": "Parent"}}]}, - }, - {"type": "table", "id": "tbl-1", "has_children": False, "table": dict[str, object]()}, - ], - "next_cursor": None, - } - child_payload = { - "results": [ - { - "type": "paragraph", - "id": "leaf", - "has_children": False, - "paragraph": {"rich_text": [{"text": {"content": "Child"}}]}, - } - ], - "next_cursor": None, - } - - mocker.patch("httpx.request", side_effect=[_mock_response(root_payload), _mock_response(child_payload)]) - mocker.patch.object(extractor, "_read_table_rows", return_value="TABLE-MD") - - content = extractor._read_block("root") - - assert "## Root" in content - assert "Parent" in content - assert "Child" in content - assert "TABLE-MD" in content - - def test_read_block_breaks_on_missing_results(self, mocker: MockerFixture) -> None: - extractor = notion_extractor.NotionExtractor( - notion_workspace_id="ws", - notion_obj_id="obj", - notion_page_type="page", - tenant_id="tenant", - notion_access_token="token", - notion_token_loader=lambda: "token", - ) - mocker.patch("httpx.request", return_value=_mock_response({"results": None, "next_cursor": None})) - - assert extractor._read_block("root") == "" - - def test_read_table_rows_formats_markdown_with_pagination(self, mocker: MockerFixture) -> None: - extractor = notion_extractor.NotionExtractor( - notion_workspace_id="ws", - notion_obj_id="obj", - notion_page_type="page", - tenant_id="tenant", - notion_access_token="token", - notion_token_loader=lambda: "token", - ) - - page_one = { - "results": [ - { - "table_row": { - "cells": [ - [{"text": {"content": "H1"}}], - [{"text": {"content": "H2"}}], - ] + "url": "https://notion.so/page-1", } - }, - { - "table_row": { - "cells": [ - [{"text": {"content": "R1C1"}}], - [{"text": {"content": "R1C2"}}], - ] - } - }, - ], - "next_cursor": "next", - } - page_two = { - "results": [ - { - "table_row": { - "cells": [ - [{"text": {"content": "H1"}}], - list[object](), - ] - } - }, - { - "table_row": { - "cells": [ - [{"text": {"content": "R2C1"}}], - [{"text": {"content": "R2C2"}}], - ] - } - }, - ], - "next_cursor": None, - } - - mocker.patch("httpx.request", side_effect=[_mock_response(page_one), _mock_response(page_two)]) - - markdown = extractor._read_table_rows("tbl-1") - - assert "| H1 | H2 |" in markdown - assert "| R1C1 | R1C2 |" in markdown - assert "| H1 | |" in markdown - assert "| R2C1 | R2C2 |" in markdown - - def test_read_table_rows_joins_rich_text_segments_within_cell(self, mocker: MockerFixture) -> None: - extractor = notion_extractor.NotionExtractor( - notion_workspace_id="ws", - notion_obj_id="obj", - notion_page_type="page", - tenant_id="tenant", - notion_access_token="token", - notion_token_loader=lambda: "token", - ) - - # A cell with mixed formatting arrives as multiple rich text segments. - page = { - "results": [ - { - "table_row": { - "cells": [ - [{"text": {"content": "Name"}}], - [{"text": {"content": "Desc"}}], - ] - } - }, - { - "table_row": { - "cells": [ - [{"text": {"content": "item1"}}], - [{"text": {"content": "plain "}}, {"text": {"content": "bold"}}], - ] - } - }, - ], - "next_cursor": None, - } - - mocker.patch("httpx.request", side_effect=[_mock_response(page)]) - - markdown = extractor._read_table_rows("tbl-1") - - assert "| item1 | plain bold |" in markdown - # Every row must keep the same column count as the header. - for line in markdown.splitlines(): - if line.startswith("|"): - assert line.count("|") == 3 - - def test_read_table_rows_preserves_empty_data_cells_as_columns(self, mocker: MockerFixture) -> None: - extractor = notion_extractor.NotionExtractor( - notion_workspace_id="ws", - notion_obj_id="obj", - notion_page_type="page", - tenant_id="tenant", - notion_access_token="token", - notion_token_loader=lambda: "token", - ) - - page = { - "results": [ - { - "table_row": { - "cells": [ - [{"text": {"content": "A"}}], - [{"text": {"content": "B"}}], - [{"text": {"content": "C"}}], - ] - } - }, - { - "table_row": { - "cells": [ - [{"text": {"content": "a1"}}], - list[object](), - [{"text": {"content": "c1"}}], - ] - } - }, - ], - "next_cursor": None, - } - - mocker.patch("httpx.request", side_effect=[_mock_response(page)]) - - markdown = extractor._read_table_rows("tbl-1") - - # The empty middle cell must remain an empty column instead of - # collapsing and shifting the following cell left. - assert "| a1 | | c1 |" in markdown - - -class TestNotionMetadataAndCredentialMethods: - def test_update_last_edited_time_no_document_model(self) -> None: - extractor = notion_extractor.NotionExtractor( - notion_workspace_id="ws", - notion_obj_id="obj", - notion_page_type="page", - tenant_id="tenant", - notion_access_token="token", - notion_token_loader=lambda: "token", - ) - - assert extractor.update_last_edited_time(None) is None - - def test_update_last_edited_time_updates_document_and_commits( - self, - monkeypatch: pytest.MonkeyPatch, - persisted_document: tuple[sessionmaker[Session], DocumentModel], - ) -> None: - extractor = notion_extractor.NotionExtractor( - notion_workspace_id="ws", - notion_obj_id="obj", - notion_page_type="page", - tenant_id="tenant", - notion_access_token="token", - notion_token_loader=lambda: "token", - ) - monkeypatch.setattr(extractor, "get_notion_last_edited_time", lambda: "2026-01-01T00:00:00.000Z") - session_maker, document = persisted_document - - snapshot = StoredDocumentExtractionInput.model_validate(document) - extractor.update_last_edited_time(snapshot) - - # Closing the writer session rolls back an uncommitted update before the independent read below. - notion_extractor.db.session.close() - with session_maker() as verification_session: - stored_document = verification_session.get(DocumentModel, document.id) - assert stored_document is not None - assert stored_document.data_source_info_dict == { - "source": "notion", - "last_edited_time": "2026-01-01T00:00:00.000Z", + ], + "has_more": False, + "next_cursor": None, } - def test_last_edited_time_update_rejects_a_foreign_owner_snapshot( - self, - monkeypatch: pytest.MonkeyPatch, - persisted_document: tuple[sessionmaker[Session], DocumentModel], - ) -> None: - session_maker, document = persisted_document - extractor = notion_extractor.NotionExtractor( - notion_workspace_id="ws", - notion_obj_id="obj", - notion_page_type="page", - tenant_id=document.tenant_id, - notion_access_token="token", - notion_token_loader=lambda: "token", - ) - monkeypatch.setattr(extractor, "get_notion_last_edited_time", lambda: "2026-01-01T00:00:00.000Z") - snapshot = StoredDocumentExtractionInput.model_validate(document).model_copy(update={"tenant_id": "foreign"}) - extractor.update_last_edited_time(snapshot) - with session_maker() as session: - stored = session.get(DocumentModel, document.id) - assert stored is not None - assert stored.data_source_info_dict["last_edited_time"] == "2025-01-01T00:00:00.000Z" + mocker.patch("httpx.post", return_value=_notion_extractor_mock_response(page)) - def test_get_notion_last_edited_time_uses_page_and_database_urls(self, mocker: MockerFixture) -> None: - extractor_page = notion_extractor.NotionExtractor( - notion_workspace_id="ws", - notion_obj_id="page-id", - notion_page_type="page", - tenant_id="tenant", - notion_access_token="token", - notion_token_loader=lambda: "token", - ) - request_mock = mocker.patch( - "httpx.request", return_value=_mock_response({"last_edited_time": "2025-05-01T00:00:00.000Z"}) - ) + docs = extractor._get_notion_database_data("db-1") - assert extractor_page.get_notion_last_edited_time() == "2025-05-01T00:00:00.000Z" - assert "pages/page-id" in request_mock.call_args[0][1] + assert len(docs) == 1 + content = docs[0].page_content + assert "rich:first second third" in content + assert "title_prop:Hello world" in content - extractor_db = notion_extractor.NotionExtractor( - notion_workspace_id="ws", - notion_obj_id="db-id", - notion_page_type="database", - tenant_id="tenant", - notion_access_token="token", - notion_token_loader=lambda: "token", - ) - request_mock = mocker.patch( - "httpx.request", return_value=_mock_response({"last_edited_time": "2025-06-01T00:00:00.000Z"}) - ) + def test_get_notion_database_data_handles_missing_results_and_empty_content( + self, mocker: MockerFixture + ) -> None: + extractor = notion_extractor.NotionExtractor( + notion_workspace_id="ws", + notion_obj_id="obj", + notion_page_type="database", + tenant_id="tenant", + notion_access_token="token", + notion_token_loader=lambda: "token", + ) - assert extractor_db.get_notion_last_edited_time() == "2025-06-01T00:00:00.000Z" - assert "databases/db-id" in request_mock.call_args[0][1] + mocker.patch("httpx.post", return_value=_notion_extractor_mock_response({"results": None})) + assert extractor._get_notion_database_data("db-1") == [] - def test_get_notion_last_edited_time_requires_access_token(self) -> None: - extractor = notion_extractor.NotionExtractor( - notion_workspace_id="ws", - notion_obj_id="obj", - notion_page_type="page", - tenant_id="tenant", - notion_access_token="token", - notion_token_loader=lambda: "token", - ) - extractor._notion_access_token = None + def test_requests_use_bounded_timeout(self, mocker: MockerFixture) -> None: + """All outbound Notion API calls must carry a bounded timeout + so a hanging endpoint cannot block extraction. + """ + extractor = notion_extractor.NotionExtractor( + notion_workspace_id="ws", + notion_obj_id="obj", + notion_page_type="database", + tenant_id="tenant", + notion_access_token="token", + notion_token_loader=lambda: "token", + ) - with pytest.raises(AssertionError, match="Notion access token is required"): + mock_post = mocker.patch("httpx.post", return_value=_notion_extractor_mock_response({"results": None})) + extractor._get_notion_database_data("db-1") + assert mock_post.call_args.kwargs["timeout"] == notion_extractor._REQUEST_TIMEOUT + + mock_request = mocker.patch( + "httpx.request", + return_value=_notion_extractor_mock_response({"last_edited_time": "2024-01-01T00:00:00.000Z"}), + ) extractor.get_notion_last_edited_time() + assert mock_request.call_args.kwargs["timeout"] == notion_extractor._REQUEST_TIMEOUT + + def test_get_notion_database_data_requires_access_token(self) -> None: + extractor = notion_extractor.NotionExtractor( + notion_workspace_id="ws", + notion_obj_id="obj", + notion_page_type="database", + tenant_id="tenant", + notion_access_token="token", + notion_token_loader=lambda: "token", + ) + extractor._notion_access_token = None + + with pytest.raises(AssertionError, match="Notion access token is required"): + extractor._get_notion_database_data("db-1") + + class TestNotionBlocks: + def test_get_notion_block_data_success_with_table_headings_children_and_pagination( + self, mocker: MockerFixture + ) -> None: + extractor = notion_extractor.NotionExtractor( + notion_workspace_id="ws", + notion_obj_id="obj", + notion_page_type="page", + tenant_id="tenant", + notion_access_token="token", + notion_token_loader=lambda: "token", + ) + + first_response = { + "results": [ + {"type": "table", "id": "tbl-1", "has_children": False, "table": dict[str, object]()}, + { + "type": "heading_1", + "id": "h1", + "has_children": False, + "heading_1": {"rich_text": [{"text": {"content": "Heading"}}]}, + }, + { + "type": "paragraph", + "id": "p1", + "has_children": True, + "paragraph": {"rich_text": [{"text": {"content": "Paragraph"}}]}, + }, + { + "type": "child_page", + "id": "cp1", + "has_children": True, + "child_page": {"rich_text": list[object]()}, + }, + ], + "next_cursor": "cursor-2", + } + second_response = { + "results": [ + { + "type": "heading_2", + "id": "h2", + "has_children": False, + "heading_2": {"rich_text": [{"text": {"content": "SubHeading"}}]}, + } + ], + "next_cursor": None, + } + + mocker.patch( + "httpx.request", + side_effect=[ + _notion_extractor_mock_response(first_response), + _notion_extractor_mock_response(second_response), + ], + ) + mocker.patch.object(extractor, "_read_table_rows", return_value="TABLE") + mocker.patch.object(extractor, "_read_block", return_value="CHILD") + + lines = extractor._get_notion_block_data("page-1") + + assert lines[0] == "TABLE\n\n" + assert "# Heading" in lines[1] + assert "Paragraph\nCHILD\n\n" in lines[2] + assert "## SubHeading" in lines[-1] + + def test_get_notion_block_data_handles_http_error_and_invalid_payload(self, mocker: MockerFixture) -> None: + extractor = notion_extractor.NotionExtractor( + notion_workspace_id="ws", + notion_obj_id="obj", + notion_page_type="page", + tenant_id="tenant", + notion_access_token="token", + notion_token_loader=lambda: "token", + ) + + mocker.patch("httpx.request", side_effect=httpx.HTTPError("network")) + with pytest.raises(ValueError, match="Error fetching Notion block data"): + extractor._get_notion_block_data("page-1") + + mocker.patch( + "httpx.request", return_value=_notion_extractor_mock_response({"bad": "payload"}, status_code=200) + ) + with pytest.raises(ValueError, match="Error fetching Notion block data"): + extractor._get_notion_block_data("page-1") + + mocker.patch( + "httpx.request", + return_value=_notion_extractor_mock_response({"results": []}, status_code=500, text="boom"), + ) + with pytest.raises(ValueError, match="Error fetching Notion block data: boom"): + extractor._get_notion_block_data("page-1") + + def test_read_block_supports_heading_table_and_recursion(self, mocker: MockerFixture) -> None: + extractor = notion_extractor.NotionExtractor( + notion_workspace_id="ws", + notion_obj_id="obj", + notion_page_type="page", + tenant_id="tenant", + notion_access_token="token", + notion_token_loader=lambda: "token", + ) + + root_payload = { + "results": [ + { + "type": "heading_2", + "id": "h2", + "has_children": False, + "heading_2": {"rich_text": [{"text": {"content": "Root"}}]}, + }, + { + "type": "paragraph", + "id": "child-block", + "has_children": True, + "paragraph": {"rich_text": [{"text": {"content": "Parent"}}]}, + }, + {"type": "table", "id": "tbl-1", "has_children": False, "table": dict[str, object]()}, + ], + "next_cursor": None, + } + child_payload = { + "results": [ + { + "type": "paragraph", + "id": "leaf", + "has_children": False, + "paragraph": {"rich_text": [{"text": {"content": "Child"}}]}, + } + ], + "next_cursor": None, + } + + mocker.patch( + "httpx.request", + side_effect=[ + _notion_extractor_mock_response(root_payload), + _notion_extractor_mock_response(child_payload), + ], + ) + mocker.patch.object(extractor, "_read_table_rows", return_value="TABLE-MD") + + content = extractor._read_block("root") + + assert "## Root" in content + assert "Parent" in content + assert "Child" in content + assert "TABLE-MD" in content + + def test_read_block_breaks_on_missing_results(self, mocker: MockerFixture) -> None: + extractor = notion_extractor.NotionExtractor( + notion_workspace_id="ws", + notion_obj_id="obj", + notion_page_type="page", + tenant_id="tenant", + notion_access_token="token", + notion_token_loader=lambda: "token", + ) + mocker.patch( + "httpx.request", return_value=_notion_extractor_mock_response({"results": None, "next_cursor": None}) + ) + + assert extractor._read_block("root") == "" + + def test_read_table_rows_formats_markdown_with_pagination(self, mocker: MockerFixture) -> None: + extractor = notion_extractor.NotionExtractor( + notion_workspace_id="ws", + notion_obj_id="obj", + notion_page_type="page", + tenant_id="tenant", + notion_access_token="token", + notion_token_loader=lambda: "token", + ) + + page_one = { + "results": [ + { + "table_row": { + "cells": [ + [{"text": {"content": "H1"}}], + [{"text": {"content": "H2"}}], + ] + } + }, + { + "table_row": { + "cells": [ + [{"text": {"content": "R1C1"}}], + [{"text": {"content": "R1C2"}}], + ] + } + }, + ], + "next_cursor": "next", + } + page_two = { + "results": [ + { + "table_row": { + "cells": [ + [{"text": {"content": "H1"}}], + list[object](), + ] + } + }, + { + "table_row": { + "cells": [ + [{"text": {"content": "R2C1"}}], + [{"text": {"content": "R2C2"}}], + ] + } + }, + ], + "next_cursor": None, + } + + mocker.patch( + "httpx.request", + side_effect=[_notion_extractor_mock_response(page_one), _notion_extractor_mock_response(page_two)], + ) + + markdown = extractor._read_table_rows("tbl-1") + + assert "| H1 | H2 |" in markdown + assert "| R1C1 | R1C2 |" in markdown + assert "| H1 | |" in markdown + assert "| R2C1 | R2C2 |" in markdown + + def test_read_table_rows_joins_rich_text_segments_within_cell(self, mocker: MockerFixture) -> None: + extractor = notion_extractor.NotionExtractor( + notion_workspace_id="ws", + notion_obj_id="obj", + notion_page_type="page", + tenant_id="tenant", + notion_access_token="token", + notion_token_loader=lambda: "token", + ) + + # A cell with mixed formatting arrives as multiple rich text segments. + page = { + "results": [ + { + "table_row": { + "cells": [ + [{"text": {"content": "Name"}}], + [{"text": {"content": "Desc"}}], + ] + } + }, + { + "table_row": { + "cells": [ + [{"text": {"content": "item1"}}], + [{"text": {"content": "plain "}}, {"text": {"content": "bold"}}], + ] + } + }, + ], + "next_cursor": None, + } + + mocker.patch("httpx.request", side_effect=[_notion_extractor_mock_response(page)]) + + markdown = extractor._read_table_rows("tbl-1") + + assert "| item1 | plain bold |" in markdown + # Every row must keep the same column count as the header. + for line in markdown.splitlines(): + if line.startswith("|"): + assert line.count("|") == 3 + + def test_read_table_rows_preserves_empty_data_cells_as_columns(self, mocker: MockerFixture) -> None: + extractor = notion_extractor.NotionExtractor( + notion_workspace_id="ws", + notion_obj_id="obj", + notion_page_type="page", + tenant_id="tenant", + notion_access_token="token", + notion_token_loader=lambda: "token", + ) + + page = { + "results": [ + { + "table_row": { + "cells": [ + [{"text": {"content": "A"}}], + [{"text": {"content": "B"}}], + [{"text": {"content": "C"}}], + ] + } + }, + { + "table_row": { + "cells": [ + [{"text": {"content": "a1"}}], + list[object](), + [{"text": {"content": "c1"}}], + ] + } + }, + ], + "next_cursor": None, + } + + mocker.patch("httpx.request", side_effect=[_notion_extractor_mock_response(page)]) + + markdown = extractor._read_table_rows("tbl-1") + + # The empty middle cell must remain an empty column instead of + # collapsing and shifting the following cell left. + assert "| a1 | | c1 |" in markdown + + class TestNotionMetadataAndCredentialMethods: + def test_update_last_edited_time_no_document_model(self) -> None: + extractor = notion_extractor.NotionExtractor( + notion_workspace_id="ws", + notion_obj_id="obj", + notion_page_type="page", + tenant_id="tenant", + notion_access_token="token", + notion_token_loader=lambda: "token", + ) + + assert extractor.update_last_edited_time(None) is None + + def test_update_last_edited_time_updates_document_and_commits( + self, + monkeypatch: pytest.MonkeyPatch, + persisted_document: tuple[sessionmaker[Session], DocumentModel], + ) -> None: + extractor = notion_extractor.NotionExtractor( + notion_workspace_id="ws", + notion_obj_id="obj", + notion_page_type="page", + tenant_id="tenant", + notion_access_token="token", + notion_token_loader=lambda: "token", + ) + monkeypatch.setattr(extractor, "get_notion_last_edited_time", lambda: "2026-01-01T00:00:00.000Z") + session_maker, document = persisted_document + + snapshot = StoredDocumentExtractionInput.model_validate(document) + extractor.update_last_edited_time(snapshot) + + # Closing the writer session rolls back an uncommitted update before the independent read below. + notion_extractor.db.session.close() + with session_maker() as verification_session: + stored_document = verification_session.get(DocumentModel, document.id) + assert stored_document is not None + assert stored_document.data_source_info_dict == { + "source": "notion", + "last_edited_time": "2026-01-01T00:00:00.000Z", + } + + def test_last_edited_time_update_rejects_a_foreign_owner_snapshot( + self, + monkeypatch: pytest.MonkeyPatch, + persisted_document: tuple[sessionmaker[Session], DocumentModel], + ) -> None: + session_maker, document = persisted_document + extractor = notion_extractor.NotionExtractor( + notion_workspace_id="ws", + notion_obj_id="obj", + notion_page_type="page", + tenant_id=document.tenant_id, + notion_access_token="token", + notion_token_loader=lambda: "token", + ) + monkeypatch.setattr(extractor, "get_notion_last_edited_time", lambda: "2026-01-01T00:00:00.000Z") + snapshot = StoredDocumentExtractionInput.model_validate(document).model_copy( + update={"tenant_id": "foreign"} + ) + extractor.update_last_edited_time(snapshot) + with session_maker() as session: + stored = session.get(DocumentModel, document.id) + assert stored is not None + assert stored.data_source_info_dict["last_edited_time"] == "2025-01-01T00:00:00.000Z" + + def test_get_notion_last_edited_time_uses_page_and_database_urls(self, mocker: MockerFixture) -> None: + extractor_page = notion_extractor.NotionExtractor( + notion_workspace_id="ws", + notion_obj_id="page-id", + notion_page_type="page", + tenant_id="tenant", + notion_access_token="token", + notion_token_loader=lambda: "token", + ) + request_mock = mocker.patch( + "httpx.request", + return_value=_notion_extractor_mock_response({"last_edited_time": "2025-05-01T00:00:00.000Z"}), + ) + + assert extractor_page.get_notion_last_edited_time() == "2025-05-01T00:00:00.000Z" + assert "pages/page-id" in request_mock.call_args[0][1] + + extractor_db = notion_extractor.NotionExtractor( + notion_workspace_id="ws", + notion_obj_id="db-id", + notion_page_type="database", + tenant_id="tenant", + notion_access_token="token", + notion_token_loader=lambda: "token", + ) + request_mock = mocker.patch( + "httpx.request", + return_value=_notion_extractor_mock_response({"last_edited_time": "2025-06-01T00:00:00.000Z"}), + ) + + assert extractor_db.get_notion_last_edited_time() == "2025-06-01T00:00:00.000Z" + assert "databases/db-id" in request_mock.call_args[0][1] + + def test_get_notion_last_edited_time_requires_access_token(self) -> None: + extractor = notion_extractor.NotionExtractor( + notion_workspace_id="ws", + notion_obj_id="obj", + notion_page_type="page", + tenant_id="tenant", + notion_access_token="token", + notion_token_loader=lambda: "token", + ) + extractor._notion_access_token = None + + with pytest.raises(AssertionError, match="Notion access token is required"): + extractor.get_notion_last_edited_time() + + @staticmethod + def test_get_cell_text_uses_plain_text_for_mention_and_equation_segments() -> None: + # Notion rich text segments of type mention or equation carry plain_text + # but no "text" object; they must not be dropped from the cell content. + cell = [ + {"type": "text", "text": {"content": "see "}, "plain_text": "see "}, + {"type": "mention", "mention": {"type": "user", "user": dict[str, object]()}, "plain_text": "@alice"}, + {"type": "equation", "equation": {"expression": "e=mc^2"}, "plain_text": "e=mc^2"}, + ] + assert notion_extractor.NotionExtractor._get_cell_text(cell) == "see @alicee=mc^2" + # Empty cell stays an empty column. + assert notion_extractor.NotionExtractor._get_cell_text([]) == "" + # Fall back to text.content when plain_text is absent. + assert notion_extractor.NotionExtractor._get_cell_text([{"text": {"content": "x"}}]) == "x" -def test_get_cell_text_uses_plain_text_for_mention_and_equation_segments() -> None: - # Notion rich text segments of type mention or equation carry plain_text - # but no "text" object; they must not be dropped from the cell content. - cell = [ - {"type": "text", "text": {"content": "see "}, "plain_text": "see "}, - {"type": "mention", "mention": {"type": "user", "user": dict[str, object]()}, "plain_text": "@alice"}, - {"type": "equation", "equation": {"expression": "e=mc^2"}, "plain_text": "e=mc^2"}, +def _notion_page_rich_text_mock_response(data: dict[str, object]) -> mock.Mock: + response = mock.Mock() + response.status_code = 200 + response.text = "" + response.json.return_value = data + return response + + +def _extractor() -> notion_extractor.NotionExtractor: + return notion_extractor.NotionExtractor( + notion_workspace_id="ws", + notion_obj_id="obj", + notion_page_type="page", + tenant_id="tenant", + notion_access_token="token", + notion_token_loader=lambda: "token", + ) + + +def _mixed_rich_text() -> list[dict[str, object]]: + return [ + {"type": "text", "plain_text": "See ", "text": {"content": "See "}}, + { + "type": "mention", + "plain_text": "Project Alpha", + "mention": {"type": "page", "page": {"id": "page-id"}}, + }, + {"type": "equation", "plain_text": "E = mc^2", "equation": {"expression": "E = mc^2"}}, ] - assert notion_extractor.NotionExtractor._get_cell_text(cell) == "see @alicee=mc^2" - # Empty cell stays an empty column. - assert notion_extractor.NotionExtractor._get_cell_text([]) == "" - # Fall back to text.content when plain_text is absent. - assert notion_extractor.NotionExtractor._get_cell_text([{"text": {"content": "x"}}]) == "x" + + +def _paragraph_payload() -> dict[str, object]: + return { + "results": [ + { + "type": "paragraph", + "id": "paragraph-id", + "has_children": False, + "paragraph": {"rich_text": _mixed_rich_text()}, + } + ], + "next_cursor": None, + } + + +class TestNotionPageRichTextCases: + @staticmethod + def test_get_notion_block_data_preserves_mentions_and_equations() -> None: + extractor = _extractor() + + with mock.patch("httpx.request", return_value=_notion_page_rich_text_mock_response(_paragraph_payload())): + lines = extractor._get_notion_block_data("page-id") + + assert lines == ["See \nProject Alpha\nE = mc^2\n\n"] + + @staticmethod + def test_read_block_preserves_mentions_and_equations() -> None: + extractor = _extractor() + + with mock.patch("httpx.request", return_value=_notion_page_rich_text_mock_response(_paragraph_payload())): + content = extractor._read_block("child-block-id", num_tabs=1) + + assert content == "\tSee \n\tProject Alpha\n\tE = mc^2\n\n" diff --git a/api/tests/unit_tests/core/rag/extractor/test_notion_page_rich_text.py b/api/tests/unit_tests/core/rag/extractor/test_notion_page_rich_text.py deleted file mode 100644 index 15cbdc124f2..00000000000 --- a/api/tests/unit_tests/core/rag/extractor/test_notion_page_rich_text.py +++ /dev/null @@ -1,66 +0,0 @@ -from unittest import mock - -from core.rag.extractor import notion_extractor - - -def _mock_response(data: dict[str, object]) -> mock.Mock: - response = mock.Mock() - response.status_code = 200 - response.text = "" - response.json.return_value = data - return response - - -def _extractor() -> notion_extractor.NotionExtractor: - return notion_extractor.NotionExtractor( - notion_workspace_id="ws", - notion_obj_id="obj", - notion_page_type="page", - tenant_id="tenant", - notion_access_token="token", - notion_token_loader=lambda: "token", - ) - - -def _mixed_rich_text() -> list[dict[str, object]]: - return [ - {"type": "text", "plain_text": "See ", "text": {"content": "See "}}, - { - "type": "mention", - "plain_text": "Project Alpha", - "mention": {"type": "page", "page": {"id": "page-id"}}, - }, - {"type": "equation", "plain_text": "E = mc^2", "equation": {"expression": "E = mc^2"}}, - ] - - -def _paragraph_payload() -> dict[str, object]: - return { - "results": [ - { - "type": "paragraph", - "id": "paragraph-id", - "has_children": False, - "paragraph": {"rich_text": _mixed_rich_text()}, - } - ], - "next_cursor": None, - } - - -def test_get_notion_block_data_preserves_mentions_and_equations() -> None: - extractor = _extractor() - - with mock.patch("httpx.request", return_value=_mock_response(_paragraph_payload())): - lines = extractor._get_notion_block_data("page-id") - - assert lines == ["See \nProject Alpha\nE = mc^2\n\n"] - - -def test_read_block_preserves_mentions_and_equations() -> None: - extractor = _extractor() - - with mock.patch("httpx.request", return_value=_mock_response(_paragraph_payload())): - content = extractor._read_block("child-block-id", num_tabs=1) - - assert content == "\tSee \n\tProject Alpha\n\tE = mc^2\n\n" diff --git a/api/tests/unit_tests/core/rag/extractor/test_pdf_extractor.py b/api/tests/unit_tests/core/rag/extractor/test_pdf_extractor.py index ec940456784..0b9a211c5d5 100644 --- a/api/tests/unit_tests/core/rag/extractor/test_pdf_extractor.py +++ b/api/tests/unit_tests/core/rag/extractor/test_pdf_extractor.py @@ -9,7 +9,7 @@ from sqlalchemy.orm import Session import core.rag.extractor.pdf_extractor as pe from models.model import UploadFile -from tests.unit_tests.config_override import apply_config_overrides +from tests.helpers.unit.config_override import apply_config_overrides TENANT_ID = str(uuid4()) USER_ID = str(uuid4()) diff --git a/api/tests/unit_tests/core/rag/extractor/test_word_extractor.py b/api/tests/unit_tests/core/rag/extractor/test_word_extractor.py index a207418515d..a4d2f10a10d 100644 --- a/api/tests/unit_tests/core/rag/extractor/test_word_extractor.py +++ b/api/tests/unit_tests/core/rag/extractor/test_word_extractor.py @@ -23,7 +23,7 @@ from sqlalchemy.orm import Session import core.rag.extractor.word_extractor as we from core.rag.extractor.word_extractor import WordExtractor from models.model import UploadFile -from tests.unit_tests.config_override import apply_config_overrides +from tests.helpers.unit.config_override import apply_config_overrides class _TextOxmlElement(Protocol): diff --git a/api/tests/unit_tests/core/rag/extractor/unstructured/test_unstructured_extractors.py b/api/tests/unit_tests/core/rag/extractor/unstructured/test_unstructured_doc_extractor.py similarity index 100% rename from api/tests/unit_tests/core/rag/extractor/unstructured/test_unstructured_extractors.py rename to api/tests/unit_tests/core/rag/extractor/unstructured/test_unstructured_doc_extractor.py diff --git a/api/tests/unit_tests/core/rag/extractor/watercrawl/test_watercrawl.py b/api/tests/unit_tests/core/rag/extractor/watercrawl/test_client.py similarity index 100% rename from api/tests/unit_tests/core/rag/extractor/watercrawl/test_watercrawl.py rename to api/tests/unit_tests/core/rag/extractor/watercrawl/test_client.py diff --git a/api/tests/unit_tests/core/rag/indexing/processor/conftest.py b/api/tests/unit_tests/core/rag/index_processor/processor/conftest.py similarity index 100% rename from api/tests/unit_tests/core/rag/indexing/processor/conftest.py rename to api/tests/unit_tests/core/rag/index_processor/processor/conftest.py diff --git a/api/tests/unit_tests/core/rag/indexing/processor/test_paragraph_index_processor.py b/api/tests/unit_tests/core/rag/index_processor/processor/test_paragraph_index_processor.py similarity index 100% rename from api/tests/unit_tests/core/rag/indexing/processor/test_paragraph_index_processor.py rename to api/tests/unit_tests/core/rag/index_processor/processor/test_paragraph_index_processor.py diff --git a/api/tests/unit_tests/core/rag/indexing/processor/test_parent_child_index_processor.py b/api/tests/unit_tests/core/rag/index_processor/processor/test_parent_child_index_processor.py similarity index 99% rename from api/tests/unit_tests/core/rag/indexing/processor/test_parent_child_index_processor.py rename to api/tests/unit_tests/core/rag/index_processor/processor/test_parent_child_index_processor.py index db837981a58..85c1e79e232 100644 --- a/api/tests/unit_tests/core/rag/indexing/processor/test_parent_child_index_processor.py +++ b/api/tests/unit_tests/core/rag/index_processor/processor/test_parent_child_index_processor.py @@ -15,7 +15,7 @@ from models.account import Account from models.dataset import ChildChunk, Dataset, DatasetProcessRule, DocumentCreatedFrom, DocumentSegment from models.dataset import Document as DatasetDocument from models.enums import DataSourceType -from tests.unit_tests.config_override import config_overrides_context +from tests.helpers.unit.config_override import config_overrides_context class TestParentChildIndexProcessor: diff --git a/api/tests/unit_tests/core/rag/indexing/processor/test_qa_index_processor.py b/api/tests/unit_tests/core/rag/index_processor/processor/test_qa_index_processor.py similarity index 100% rename from api/tests/unit_tests/core/rag/indexing/processor/test_qa_index_processor.py rename to api/tests/unit_tests/core/rag/index_processor/processor/test_qa_index_processor.py diff --git a/api/tests/unit_tests/core/rag/indexing/test_index_processor.py b/api/tests/unit_tests/core/rag/index_processor/test_index_processor.py similarity index 89% rename from api/tests/unit_tests/core/rag/indexing/test_index_processor.py rename to api/tests/unit_tests/core/rag/index_processor/test_index_processor.py index a0ea675e96a..01a9564a9de 100644 --- a/api/tests/unit_tests/core/rag/indexing/test_index_processor.py +++ b/api/tests/unit_tests/core/rag/index_processor/test_index_processor.py @@ -3,14 +3,17 @@ from contextlib import nullcontext from types import SimpleNamespace from unittest.mock import MagicMock, Mock, patch +import pytest from flask import Flask from sqlalchemy import event from sqlalchemy.orm import Session, sessionmaker from core.rag.index_processor import index_processor as index_processor_module -from core.rag.index_processor.constant.index_type import IndexTechniqueType -from core.rag.index_processor.index_processor import IndexProcessor +from core.rag.index_processor.constant.index_type import IndexStructureType, IndexTechniqueType +from core.rag.index_processor.index_processor import IndexProcessor, IndexProcessorFactory from core.rag.index_processor.processor.paragraph_index_processor import ParagraphIndexProcessor +from core.rag.index_processor.processor.parent_child_index_processor import ParentChildIndexProcessor +from core.rag.index_processor.processor.qa_index_processor import QAIndexProcessor from core.rag.models.document import AttachmentDocument from core.workflow.nodes.knowledge_index.protocols import Preview, PreviewItem from extensions.ext_application_services import ApplicationServices @@ -312,3 +315,38 @@ class TestIndexProcessor: call_sessions = [call.kwargs["session"] for call in generate_summary.call_args_list] assert all(call_session is not sqlite_session for call_session in call_sessions) assert call_sessions == worker_sessions + + +class TestIndexProcessorFactory: + def test_requires_index_type(self) -> None: + factory = IndexProcessorFactory(index_type=None) + + with pytest.raises(ValueError, match="Index type must be specified"): + factory.init_index_processor() + + def test_builds_paragraph_processor(self) -> None: + factory = IndexProcessorFactory(index_type=IndexStructureType.PARAGRAPH_INDEX) + + processor = factory.init_index_processor() + + assert isinstance(processor, ParagraphIndexProcessor) + + def test_builds_qa_processor(self) -> None: + factory = IndexProcessorFactory(index_type=IndexStructureType.QA_INDEX) + + processor = factory.init_index_processor() + + assert isinstance(processor, QAIndexProcessor) + + def test_builds_parent_child_processor(self) -> None: + factory = IndexProcessorFactory(index_type=IndexStructureType.PARENT_CHILD_INDEX) + + processor = factory.init_index_processor() + + assert isinstance(processor, ParentChildIndexProcessor) + + def test_rejects_unsupported_index_type(self) -> None: + factory = IndexProcessorFactory(index_type="unsupported") + + with pytest.raises(ValueError, match="is not supported"): + factory.init_index_processor() diff --git a/api/tests/unit_tests/core/rag/indexing/test_index_processor_base.py b/api/tests/unit_tests/core/rag/index_processor/test_index_processor_base.py similarity index 99% rename from api/tests/unit_tests/core/rag/indexing/test_index_processor_base.py rename to api/tests/unit_tests/core/rag/index_processor/test_index_processor_base.py index aff0cfdeb40..b4103efdf23 100644 --- a/api/tests/unit_tests/core/rag/indexing/test_index_processor_base.py +++ b/api/tests/unit_tests/core/rag/index_processor/test_index_processor_base.py @@ -20,8 +20,8 @@ from models.dataset import Dataset from models.enums import CreatorUserRole from models.model import UploadFile from models.tools import ToolFile -from tests.unit_tests.config_override import config_overrides_context -from tests.unit_tests.model_factories import make_document +from tests.helpers.unit.config_override import config_overrides_context +from tests.helpers.unit.model_factories import make_document def _upload( diff --git a/api/tests/unit_tests/core/rag/indexing/test_index_processor_factory.py b/api/tests/unit_tests/core/rag/indexing/test_index_processor_factory.py deleted file mode 100644 index d6d5c31e815..00000000000 --- a/api/tests/unit_tests/core/rag/indexing/test_index_processor_factory.py +++ /dev/null @@ -1,42 +0,0 @@ -import pytest - -from core.rag.index_processor.constant.index_type import IndexStructureType -from core.rag.index_processor.index_processor import IndexProcessorFactory -from core.rag.index_processor.processor.paragraph_index_processor import ParagraphIndexProcessor -from core.rag.index_processor.processor.parent_child_index_processor import ParentChildIndexProcessor -from core.rag.index_processor.processor.qa_index_processor import QAIndexProcessor - - -class TestIndexProcessorFactory: - def test_requires_index_type(self) -> None: - factory = IndexProcessorFactory(index_type=None) - - with pytest.raises(ValueError, match="Index type must be specified"): - factory.init_index_processor() - - def test_builds_paragraph_processor(self) -> None: - factory = IndexProcessorFactory(index_type=IndexStructureType.PARAGRAPH_INDEX) - - processor = factory.init_index_processor() - - assert isinstance(processor, ParagraphIndexProcessor) - - def test_builds_qa_processor(self) -> None: - factory = IndexProcessorFactory(index_type=IndexStructureType.QA_INDEX) - - processor = factory.init_index_processor() - - assert isinstance(processor, QAIndexProcessor) - - def test_builds_parent_child_processor(self) -> None: - factory = IndexProcessorFactory(index_type=IndexStructureType.PARENT_CHILD_INDEX) - - processor = factory.init_index_processor() - - assert isinstance(processor, ParentChildIndexProcessor) - - def test_rejects_unsupported_index_type(self) -> None: - factory = IndexProcessorFactory(index_type="unsupported") - - with pytest.raises(ValueError, match="is not supported"): - factory.init_index_processor() diff --git a/api/tests/unit_tests/core/rag/rerank/test_reranker.py b/api/tests/unit_tests/core/rag/rerank/test_rerank_model.py similarity index 99% rename from api/tests/unit_tests/core/rag/rerank/test_reranker.py rename to api/tests/unit_tests/core/rag/rerank/test_rerank_model.py index 270ec9ab367..f4f26ed054d 100644 --- a/api/tests/unit_tests/core/rag/rerank/test_reranker.py +++ b/api/tests/unit_tests/core/rag/rerank/test_rerank_model.py @@ -37,7 +37,7 @@ from graphon.model_runtime.entities.rerank_entities import RerankDocument, Reran from graphon.model_runtime.model_providers.base.rerank_model import RerankModel from models.enums import CreatorUserRole from models.model import UploadFile -from tests.unit_tests.core.model_fixtures import make_model_config +from tests.helpers.unit.core.model_fixtures import make_model_config def create_model_instance() -> ModelInstance: diff --git a/api/tests/unit_tests/core/rag/retrieval/test_structured_chat_output_parser.py b/api/tests/unit_tests/core/rag/retrieval/output_parser/test_structured_chat.py similarity index 100% rename from api/tests/unit_tests/core/rag/retrieval/test_structured_chat_output_parser.py rename to api/tests/unit_tests/core/rag/retrieval/output_parser/test_structured_chat.py diff --git a/api/tests/unit_tests/core/rag/retrieval/test_multi_dataset_function_call_router.py b/api/tests/unit_tests/core/rag/retrieval/router/test_multi_dataset_function_call_router.py similarity index 100% rename from api/tests/unit_tests/core/rag/retrieval/test_multi_dataset_function_call_router.py rename to api/tests/unit_tests/core/rag/retrieval/router/test_multi_dataset_function_call_router.py diff --git a/api/tests/unit_tests/core/rag/retrieval/test_multi_dataset_react_route.py b/api/tests/unit_tests/core/rag/retrieval/router/test_multi_dataset_react_route.py similarity index 100% rename from api/tests/unit_tests/core/rag/retrieval/test_multi_dataset_react_route.py rename to api/tests/unit_tests/core/rag/retrieval/router/test_multi_dataset_react_route.py diff --git a/api/tests/unit_tests/core/rag/retrieval/test_dataset_retrieval.py b/api/tests/unit_tests/core/rag/retrieval/test_dataset_retrieval.py index d0a0fcd91a8..a2d419f7bac 100644 --- a/api/tests/unit_tests/core/rag/retrieval/test_dataset_retrieval.py +++ b/api/tests/unit_tests/core/rag/retrieval/test_dataset_retrieval.py @@ -1,6 +1,7 @@ import threading from collections.abc import Generator from contextlib import contextmanager, nullcontext +from dataclasses import dataclass from datetime import datetime from types import SimpleNamespace from typing import Any @@ -9,11 +10,11 @@ from uuid import uuid4 import pytest from flask import Flask, current_app -from sqlalchemy import event, select +from sqlalchemy import event, inspect, select from sqlalchemy.dialects import postgresql from sqlalchemy.engine import Engine -from sqlalchemy.exc import DBAPIError -from sqlalchemy.orm import Session +from sqlalchemy.exc import DBAPIError, InvalidRequestError +from sqlalchemy.orm import Session, sessionmaker import core.rag.retrieval.dataset_retrieval as dataset_retrieval_module from core.app.app_config.entities import ( @@ -27,16 +28,21 @@ from core.app.app_config.entities import ( ModelConfig as AppModelConfig, ) from core.app.app_config.entities import ModelConfig as WorkflowModelConfig -from core.app.entities.app_invoke_entities import InvokeFrom, ModelConfigWithCredentialsEntity +from core.app.entities.app_invoke_entities import InvokeFrom, ModelConfigWithCredentialsEntity, UserFrom +from core.app.file_access import FileAccessScope, bind_file_access_scope, get_current_file_access_scope from core.entities.agent_entities import PlanningStrategy from core.entities.model_entities import ModelStatus from core.rag.data_post_processor.data_post_processor import WeightsDict from core.rag.datasource.retrieval_service import RetrievalService +from core.rag.embedding.retrieval import RetrievalSegments from core.rag.entities import Condition as AppCondition from core.rag.index_processor.constant.doc_type import DocType -from core.rag.index_processor.constant.index_type import IndexStructureType -from core.rag.models.document import Document +from core.rag.index_processor.constant.index_type import IndexStructureType, IndexTechniqueType +from core.rag.models.document import Document as RetrievedDocument +from core.rag.models.document import Document as _dataset_retrieval_Document +from core.rag.models.document import Document as _dataset_retrieval_methods_Document from core.rag.rerank.rerank_type import RerankMode +from core.rag.retrieval import dataset_retrieval as retrieval_module from core.rag.retrieval.dataset_retrieval import DatasetRetrieval from core.rag.retrieval.retrieval_methods import RetrievalMethod from core.workflow.nodes.knowledge_retrieval import exc @@ -55,7 +61,9 @@ from models.dataset import ( SegmentAttachmentBinding, ) from models.dataset import Document as DatasetDocument -from models.enums import CreatorUserRole, DataSourceType, DocumentCreatedFrom +from models.dataset import Document as _retrieval_source_query_Document +from models.enums import CreatorUserRole, DataSourceType, DocumentCreatedFrom, IndexingStatus +from repositories.knowledge.retrieval_source_repository import get_retrieval_source_documents # ==================== Helper Functions ==================== @@ -66,7 +74,7 @@ def create_mock_document( score: float = 0.8, provider: str = "dify", additional_metadata: dict[str, Any] | None = None, -) -> Document: +) -> _dataset_retrieval_Document: """ Create a mock Document object for testing. @@ -98,7 +106,7 @@ def create_mock_document( if additional_metadata: metadata.update(additional_metadata) - return Document( + return _dataset_retrieval_Document( page_content=content, metadata=metadata, provider=provider, @@ -113,7 +121,7 @@ def _metadata_condition() -> AppMetadataFilteringCondition: return AppMetadataFilteringCondition(logical_operator="and", conditions=[]) -def create_side_effect_for_search(documents: list[Document]): +def create_side_effect_for_search(documents: list[_dataset_retrieval_Document]): """ Create a side effect function for mocking search methods. @@ -176,2031 +184,6 @@ def create_side_effect_with_exception(error_message: str): return side_effect -class TestRetrievalService: - """ - Comprehensive test suite for RetrievalService class. - - This test class validates all retrieval methods and their interactions, - including edge cases, error handling, and integration scenarios. - - Test Organization: - ================== - 1. Fixtures (lines ~190-240) - - mock_dataset: Standard dataset configuration - - sample_documents: Reusable test documents with varying scores - - mock_flask_app: Flask application context - - mock_thread_pool: Synchronous executor for deterministic testing - - 2. Vector Search Tests (lines ~240-350) - - Basic functionality - - Document filtering - - Empty results - - Metadata filtering - - Score thresholds - - 3. Keyword Search Tests (lines ~350-450) - - Basic keyword matching - - Special character handling - - Document filtering - - 4. Hybrid Search Tests (lines ~450-640) - - Vector + full-text combination - - Deduplication logic - - Weighted score merging - - 5. Full-Text Search Tests (lines ~640-680) - - BM25-based search - - 6. Score Merging Tests (lines ~680-790) - - Deduplication algorithms - - Score comparison - - Provider-specific handling - - 7. Error Handling Tests (lines ~790-920) - - Empty queries - - Non-existent datasets - - Exception propagation - - 8. Additional Tests (lines ~920-1080) - - Query escaping - - Reranking integration - - Top-K limiting - - Mocking Strategy: - ================= - Tests mock at the method level (embedding_search, keyword_search, etc.) - rather than the underlying Vector/Keyword classes. This approach: - - Avoids complexity of mocking ThreadPoolExecutor behavior - - Provides clearer test intent - - Makes tests more maintainable - - Properly simulates the in-place list modification pattern - - Common Patterns: - ================ - 1. **Arrange**: Set up mocks with side_effect functions - 2. **Act**: Call RetrievalService.retrieve() with specific parameters - 3. **Assert**: Verify results, mock calls, and side effects - - Example Test Structure: - ```python - def test_example(self, mock_get_dataset, mock_search, mock_dataset): - # Arrange: Set up test data and mocks - mock_get_dataset.return_value = mock_dataset - mock_search.side_effect = create_side_effect_for_search([doc1, doc2]) - - # Act: Execute the method under test - results = RetrievalService.retrieve(...) - - # Assert: Verify expectations - assert len(results) == 2 - mock_search.assert_called_once() - ``` - """ - - @pytest.fixture - def mock_dataset(self) -> Dataset: - """ - Create a Dataset object for testing. - - Returns: - Dataset with standard configuration. - """ - dataset = Dataset( - id=str(uuid4()), - tenant_id=str(uuid4()), - name="test_dataset", - indexing_technique="high_quality", - embedding_model="text-embedding-ada-002", - embedding_model_provider="openai", - retrieval_model={ - "search_method": RetrievalMethod.SEMANTIC_SEARCH, - "reranking_enable": False, - "top_k": 4, - "score_threshold_enabled": False, - }, - ) - return dataset - - @pytest.fixture - def sample_documents(self) -> list[Document]: - """ - Create sample documents for testing retrieval results. - - Returns: - list[Document]: List of mock documents with varying scores - """ - return [ - Document( - page_content="Python is a high-level programming language.", - metadata={ - "doc_id": "doc1", - "document_id": str(uuid4()), - "dataset_id": str(uuid4()), - "score": 0.95, - }, - provider="dify", - ), - Document( - page_content="JavaScript is widely used for web development.", - metadata={ - "doc_id": "doc2", - "document_id": str(uuid4()), - "dataset_id": str(uuid4()), - "score": 0.85, - }, - provider="dify", - ), - Document( - page_content="Machine learning is a subset of artificial intelligence.", - metadata={ - "doc_id": "doc3", - "document_id": str(uuid4()), - "dataset_id": str(uuid4()), - "score": 0.75, - }, - provider="dify", - ), - ] - - @pytest.fixture - def mock_flask_app(self): - """ - Create a mock Flask application context. - - Returns: - Mock: Flask app mock with app_context - """ - app = MagicMock() - app.app_context.return_value.__enter__ = Mock() - app.app_context.return_value.__exit__ = Mock() - return app - - @pytest.fixture - def retrieval_engine(self, sqlite_engine: Engine): - with patch("core.rag.datasource.retrieval_service.db", SimpleNamespace(engine=sqlite_engine)): - yield sqlite_engine - - @pytest.fixture(autouse=True) - def mock_thread_pool(self): - """ - Mock ThreadPoolExecutor to run tasks synchronously in tests. - - The RetrievalService uses ThreadPoolExecutor to run search operations - concurrently (embedding_search, keyword_search, full_text_index_search). - In tests, we want synchronous execution for: - - Deterministic behavior - - Easier debugging - - Avoiding race conditions - - Simpler assertions - - How it works: - ------------- - 1. Intercepts ThreadPoolExecutor creation - 2. Replaces submit() to execute functions immediately (synchronously) - 3. Functions modify shared all_documents list in-place - 4. Mocks concurrent.futures.wait() since tasks are already done - - Why this approach: - ------------------ - - RetrievalService.retrieve() creates a ThreadPoolExecutor context - - It submits search tasks that modify all_documents list - - concurrent.futures.wait() waits for all tasks to complete - - By executing synchronously, we avoid threading complexity in tests - - Returns: - Mock: Mocked ThreadPoolExecutor that executes tasks synchronously - """ - with patch("core.rag.datasource.retrieval_service.ThreadPoolExecutor") as mock_executor: - # Store futures to track submitted tasks (for debugging if needed) - futures_list = [] - - def sync_submit(fn, *args, **kwargs): - """ - Synchronous replacement for ThreadPoolExecutor.submit(). - - Instead of scheduling the function for async execution, - we execute it immediately in the current thread. - - Args: - fn: The function to execute (e.g., embedding_search) - *args, **kwargs: Arguments to pass to the function - - Returns: - Mock: A mock Future object - """ - future = Mock() - try: - # Execute immediately - this modifies all_documents in place - # The function signature is: fn(flask_app, dataset_id, query, - # top_k, all_documents, exceptions, ...) - fn(*args, **kwargs) - future.result.return_value = None - future.exception.return_value = None - except Exception as e: - # If function raises, store exception in future - future.result.return_value = None - future.exception.return_value = e - - futures_list.append(future) - return future - - # Set up the mock executor instance - mock_executor_instance = Mock() - mock_executor_instance.submit = sync_submit - - # Configure context manager behavior (__enter__ and __exit__) - mock_executor.return_value.__enter__.return_value = mock_executor_instance - mock_executor.return_value.__exit__.return_value = None - - # Mock concurrent.futures.wait to do nothing since tasks are already done - # In real code, this waits for all futures to complete - # In tests, futures complete immediately, so wait is a no-op - with patch("core.rag.datasource.retrieval_service.concurrent.futures.wait"): - # Mock concurrent.futures.as_completed for early error propagation - # In real code, this yields futures as they complete - # In tests, we yield all futures immediately since they're already done - def mock_as_completed(futures_list, timeout=None): - """Mock as_completed that yields futures immediately.""" - yield from futures_list - - with patch( - "core.rag.datasource.retrieval_service.concurrent.futures.as_completed", - side_effect=mock_as_completed, - ): - yield mock_executor - - # ==================== Vector Search Tests ==================== - - @patch("core.rag.datasource.retrieval_service.RetrievalService._retrieve") - @patch("core.rag.datasource.retrieval_service.RetrievalService._get_dataset") - def test_vector_search_basic(self, mock_get_dataset, mock_retrieve, mock_dataset, sample_documents): - """ - Test basic vector/semantic search functionality. - - This test validates the core vector search flow: - 1. Dataset is retrieved from database - 2. _retrieve is called via ThreadPoolExecutor - 3. Documents are added to shared all_documents list - 4. Results are returned to caller - - Verifies: - - Vector search is called with correct parameters - - Results are returned in expected format - - Score threshold is applied correctly - - Documents maintain their metadata and scores - """ - # ==================== ARRANGE ==================== - # Set up the mock dataset that will be "retrieved" from database - mock_get_dataset.return_value = mock_dataset - - # Create a side effect function that simulates _retrieve behavior - # _retrieve modifies the all_documents list in place - def side_effect_retrieve( - flask_app, - retrieval_method, - dataset, - query=None, - top_k=4, - score_threshold=None, - reranking_model=None, - reranking_mode="reranking_model", - weights=None, - document_ids_filter=None, - attachment_id=None, - all_documents=None, - exceptions=None, - ): - """Simulate _retrieve adding documents to the shared list.""" - if all_documents is not None: - all_documents.extend(sample_documents) - - mock_retrieve.side_effect = side_effect_retrieve - - # Define test parameters - query = "What is Python?" # Natural language query - top_k = 3 # Maximum number of results to return - score_threshold = 0.7 # Minimum relevance score (0.0 to 1.0) - - # ==================== ACT ==================== - # Call the retrieve method with SEMANTIC_SEARCH strategy - # This will: - # 1. Check if query is empty (early return if so) - # 2. Get the dataset using _get_dataset - # 3. Create ThreadPoolExecutor - # 4. Submit _retrieve task - # 5. Wait for completion - # 6. Return all_documents list - results = RetrievalService.retrieve( - retrieval_method=RetrievalMethod.SEMANTIC_SEARCH, - dataset_id=mock_dataset.id, - query=query, - top_k=top_k, - score_threshold=score_threshold, - ) - - # ==================== ASSERT ==================== - # Verify we got the expected number of documents - assert len(results) == 3, "Should return 3 documents from sample_documents" - - # Verify all results are Document objects (type safety) - assert all(isinstance(doc, Document) for doc in results), "All results should be Document instances" - - # Verify documents maintain their scores (highest score first in sample_documents) - assert results[0].metadata["score"] == 0.95, "First document should have highest score from sample_documents" - - # Verify _retrieve was called exactly once - # This confirms the search method was invoked by ThreadPoolExecutor - mock_retrieve.assert_called_once() - - @patch("core.rag.datasource.retrieval_service.RetrievalService._retrieve") - @patch("core.rag.datasource.retrieval_service.RetrievalService._get_dataset") - def test_vector_search_with_document_filter(self, mock_get_dataset, mock_retrieve, mock_dataset, sample_documents): - """ - Test vector search with document ID filtering. - - Verifies: - - Document ID filter is passed correctly to vector search - - Only specified documents are searched - """ - # Arrange - mock_get_dataset.return_value = mock_dataset - filtered_docs = [sample_documents[0]] - - def side_effect_retrieve( - flask_app, - retrieval_method, - dataset, - query=None, - top_k=4, - score_threshold=None, - reranking_model=None, - reranking_mode="reranking_model", - weights=None, - document_ids_filter=None, - attachment_id=None, - all_documents=None, - exceptions=None, - ): - if all_documents is not None: - all_documents.extend(filtered_docs) - - mock_retrieve.side_effect = side_effect_retrieve - document_ids_filter = [sample_documents[0].metadata["document_id"]] - - # Act - results = RetrievalService.retrieve( - retrieval_method=RetrievalMethod.SEMANTIC_SEARCH, - dataset_id=mock_dataset.id, - query="test query", - top_k=5, - document_ids_filter=document_ids_filter, - ) - - # Assert - assert len(results) == 1 - assert results[0].metadata["doc_id"] == "doc1" - # Verify document_ids_filter was passed - call_kwargs = mock_retrieve.call_args.kwargs - assert call_kwargs["document_ids_filter"] == document_ids_filter - - @patch("core.rag.datasource.retrieval_service.RetrievalService._retrieve") - @patch("core.rag.datasource.retrieval_service.RetrievalService._get_dataset") - def test_vector_search_empty_results(self, mock_get_dataset, mock_retrieve, mock_dataset): - """ - Test vector search when no results match the query. - - Verifies: - - Empty list is returned when no documents match - - No errors are raised - """ - # Arrange - mock_get_dataset.return_value = mock_dataset - # _retrieve doesn't add anything to all_documents - mock_retrieve.side_effect = lambda *args, **kwargs: None - - # Act - results = RetrievalService.retrieve( - retrieval_method=RetrievalMethod.SEMANTIC_SEARCH, - dataset_id=mock_dataset.id, - query="nonexistent query", - top_k=5, - ) - - # Assert - assert results == [] - - # ==================== Keyword Search Tests ==================== - - @patch("core.rag.datasource.retrieval_service.RetrievalService._retrieve") - @patch("core.rag.datasource.retrieval_service.RetrievalService._get_dataset") - def test_keyword_search_basic(self, mock_get_dataset, mock_retrieve, mock_dataset, sample_documents): - """ - Test basic keyword search functionality. - - Verifies: - - Keyword search is invoked correctly - - Query is escaped properly for search - - Results are returned in expected format - """ - # Arrange - mock_get_dataset.return_value = mock_dataset - - def side_effect_retrieve( - flask_app, - retrieval_method, - dataset, - query=None, - top_k=4, - score_threshold=None, - reranking_model=None, - reranking_mode="reranking_model", - weights=None, - document_ids_filter=None, - attachment_id=None, - all_documents=None, - exceptions=None, - ): - if all_documents is not None: - all_documents.extend(sample_documents) - - mock_retrieve.side_effect = side_effect_retrieve - - query = "Python programming" - top_k = 3 - - # Act - results = RetrievalService.retrieve( - retrieval_method=RetrievalMethod.KEYWORD_SEARCH, - dataset_id=mock_dataset.id, - query=query, - top_k=top_k, - ) - - # Assert - assert len(results) == 3 - assert all(isinstance(doc, Document) for doc in results) - mock_retrieve.assert_called_once() - - @patch("core.rag.datasource.retrieval_service.RetrievalService.keyword_search") - @patch("core.rag.datasource.retrieval_service.RetrievalService._get_dataset") - def test_keyword_search_with_special_characters(self, mock_get_dataset, mock_keyword_search, mock_dataset): - """ - Test keyword search with special characters in query. - - Verifies: - - Special characters are escaped correctly - - Search handles quotes and other special chars - """ - # Arrange - mock_get_dataset.return_value = mock_dataset - mock_keyword_search.side_effect = lambda *args, **kwargs: None - - query = 'Python "programming" language' - - # Act - RetrievalService.retrieve( - retrieval_method=RetrievalMethod.KEYWORD_SEARCH, - dataset_id=mock_dataset.id, - query=query, - top_k=5, - ) - - # Assert - # Verify that keyword_search was called - assert mock_keyword_search.called - # The query escaping happens inside keyword_search method - call_args = mock_keyword_search.call_args - assert call_args is not None - - @patch("core.rag.datasource.retrieval_service.RetrievalService.keyword_search") - @patch("core.rag.datasource.retrieval_service.RetrievalService._get_dataset") - def test_keyword_search_with_document_filter( - self, mock_get_dataset, mock_keyword_search, mock_dataset, sample_documents - ): - """ - Test keyword search with document ID filtering. - - Verifies: - - Document filter is applied to keyword search - - Only filtered documents are returned - """ - # Arrange - mock_get_dataset.return_value = mock_dataset - filtered_docs = [sample_documents[1]] - - def side_effect_keyword_search( - flask_app, dataset_id, query, top_k, all_documents, exceptions, document_ids_filter=None - ): - all_documents.extend(filtered_docs) - - mock_keyword_search.side_effect = side_effect_keyword_search - document_ids_filter = [sample_documents[1].metadata["document_id"]] - - # Act - results = RetrievalService.retrieve( - retrieval_method=RetrievalMethod.KEYWORD_SEARCH, - dataset_id=mock_dataset.id, - query="JavaScript", - top_k=5, - document_ids_filter=document_ids_filter, - ) - - # Assert - assert len(results) == 1 - assert results[0].metadata["doc_id"] == "doc2" - - # ==================== Hybrid Search Tests ==================== - - @patch("core.rag.datasource.retrieval_service.DataPostProcessor") - @patch("core.rag.datasource.retrieval_service.RetrievalService.full_text_index_search") - @patch("core.rag.datasource.retrieval_service.RetrievalService.embedding_search") - @patch("core.rag.datasource.retrieval_service.RetrievalService._get_dataset") - def test_hybrid_search_basic( - self, - mock_get_dataset, - mock_embedding_search, - mock_fulltext_search, - mock_data_processor_class, - mock_dataset, - sample_documents, - retrieval_engine, - ): - """ - Test basic hybrid search combining vector and full-text search. - - Verifies: - - Both vector and full-text search are executed - - Results are merged and deduplicated - - DataPostProcessor is invoked for score merging - """ - # Arrange - mock_get_dataset.return_value = mock_dataset - - # Vector search returns first 2 docs - def side_effect_embedding( - flask_app, - dataset_id, - query, - top_k, - score_threshold, - reranking_model, - all_documents, - retrieval_method, - exceptions, - document_ids_filter=None, - ): - all_documents.extend(sample_documents[:2]) - - mock_embedding_search.side_effect = side_effect_embedding - - # Full-text search returns last 2 docs (with overlap) - def side_effect_fulltext( - flask_app, - dataset_id, - query, - top_k, - score_threshold, - reranking_model, - all_documents, - retrieval_method, - exceptions, - document_ids_filter=None, - ): - all_documents.extend(sample_documents[1:]) - - mock_fulltext_search.side_effect = side_effect_fulltext - - # Mock DataPostProcessor - mock_processor_instance = Mock() - mock_processor_instance.invoke.return_value = sample_documents - mock_data_processor_class.return_value = mock_processor_instance - - # Act - results = RetrievalService.retrieve( - retrieval_method=RetrievalMethod.HYBRID_SEARCH, - dataset_id=mock_dataset.id, - query="Python programming", - top_k=3, - score_threshold=0.5, - ) - - # Assert - assert len(results) == 3 - mock_embedding_search.assert_called_once() - mock_fulltext_search.assert_called_once() - mock_processor_instance.invoke.assert_called_once() - processor_session = mock_data_processor_class.call_args.kwargs["session"] - assert isinstance(processor_session, Session) - assert processor_session.get_bind() is retrieval_engine - - @patch("core.rag.datasource.retrieval_service.DataPostProcessor") - @patch("core.rag.datasource.retrieval_service.RetrievalService.full_text_index_search") - @patch("core.rag.datasource.retrieval_service.RetrievalService.embedding_search") - @patch("core.rag.datasource.retrieval_service.RetrievalService._get_dataset") - def test_hybrid_search_deduplication( - self, - mock_get_dataset, - mock_embedding_search, - mock_fulltext_search, - mock_data_processor_class, - mock_dataset, - retrieval_engine, - ): - """ - Test that hybrid search properly deduplicates documents. - - Hybrid search combines results from multiple search methods (vector + full-text). - This can lead to duplicate documents when the same chunk is found by both methods. - - Scenario: - --------- - 1. Vector search finds document "duplicate_doc" with score 0.9 - 2. Full-text search also finds "duplicate_doc" but with score 0.6 - 3. Both searches find "unique_doc" - 4. Deduplication should keep only the higher-scoring version (0.9) - - Why deduplication matters: - -------------------------- - - Prevents showing the same content multiple times to users - - Ensures score consistency (keeps best match) - - Improves result quality and user experience - - Happens BEFORE reranking to avoid processing duplicates - - Verifies: - - Duplicate documents (same doc_id) are removed - - Higher scoring duplicate is retained - - Deduplication happens before post-processing - - Final result count is correct - """ - # ==================== ARRANGE ==================== - mock_get_dataset.return_value = mock_dataset - - # Create test documents with intentional duplication - # Same doc_id but different scores to test score comparison logic - doc1_high = Document( - page_content="Content 1", - metadata={ - "doc_id": "duplicate_doc", # Same doc_id as doc1_low - "score": 0.9, # Higher score - should be kept - "document_id": str(uuid4()), - }, - provider="dify", - ) - doc1_low = Document( - page_content="Content 1", - metadata={ - "doc_id": "duplicate_doc", # Same doc_id as doc1_high - "score": 0.6, # Lower score - should be discarded - "document_id": str(uuid4()), - }, - provider="dify", - ) - doc2 = Document( - page_content="Content 2", - metadata={ - "doc_id": "unique_doc", # Unique doc_id - "score": 0.8, - "document_id": str(uuid4()), - }, - provider="dify", - ) - - # Simulate vector search returning high-score duplicate + unique doc - def side_effect_embedding( - flask_app, - dataset_id, - query, - top_k, - score_threshold, - reranking_model, - all_documents, - retrieval_method, - exceptions, - document_ids_filter=None, - ): - """Vector search finds 2 documents including high-score duplicate.""" - all_documents.extend([doc1_high, doc2]) - - mock_embedding_search.side_effect = side_effect_embedding - - # Simulate full-text search returning low-score duplicate - def side_effect_fulltext( - flask_app, - dataset_id, - query, - top_k, - score_threshold, - reranking_model, - all_documents, - retrieval_method, - exceptions, - document_ids_filter=None, - ): - """Full-text search finds the same document but with lower score.""" - all_documents.extend([doc1_low]) - - mock_fulltext_search.side_effect = side_effect_fulltext - - # Mock DataPostProcessor to return deduplicated results - # In real implementation, _deduplicate_documents is called before this - mock_processor_instance = Mock() - mock_processor_instance.invoke.return_value = [doc1_high, doc2] - mock_data_processor_class.return_value = mock_processor_instance - - # ==================== ACT ==================== - # Execute hybrid search which should: - # 1. Run both embedding_search and full_text_index_search - # 2. Collect all results in all_documents (3 docs: 2 unique + 1 duplicate) - # 3. Call _deduplicate_documents to remove duplicate (keeps higher score) - # 4. Pass deduplicated results to DataPostProcessor - # 5. Return final results - results = RetrievalService.retrieve( - retrieval_method=RetrievalMethod.HYBRID_SEARCH, - dataset_id=mock_dataset.id, - query="test", - top_k=5, - ) - - # ==================== ASSERT ==================== - # Verify deduplication worked correctly - assert len(results) == 2, "Should have 2 unique documents after deduplication (not 3)" - - # Verify the correct documents are present - doc_ids = [doc.metadata["doc_id"] for doc in results] - assert "duplicate_doc" in doc_ids, "Duplicate doc should be present (higher score version)" - assert "unique_doc" in doc_ids, "Unique doc should be present" - processor_session = mock_data_processor_class.call_args.kwargs["session"] - assert isinstance(processor_session, Session) - assert processor_session.get_bind() is retrieval_engine - - # Implicitly verifies that doc1_low (score 0.6) was discarded - # in favor of doc1_high (score 0.9) - - @patch("core.rag.datasource.retrieval_service.DataPostProcessor") - @patch("core.rag.datasource.retrieval_service.RetrievalService.full_text_index_search") - @patch("core.rag.datasource.retrieval_service.RetrievalService.embedding_search") - @patch("core.rag.datasource.retrieval_service.RetrievalService._get_dataset") - def test_hybrid_search_with_weights( - self, - mock_get_dataset, - mock_embedding_search, - mock_fulltext_search, - mock_data_processor_class, - mock_dataset, - sample_documents, - retrieval_engine, - ): - """ - Test hybrid search with custom weights for score merging. - - Verifies: - - Weights are passed to DataPostProcessor - - Score merging respects weight configuration - """ - # Arrange - mock_get_dataset.return_value = mock_dataset - - def side_effect_embedding( - flask_app, - dataset_id, - query, - top_k, - score_threshold, - reranking_model, - all_documents, - retrieval_method, - exceptions, - document_ids_filter=None, - ): - all_documents.extend(sample_documents[:2]) - - mock_embedding_search.side_effect = side_effect_embedding - - def side_effect_fulltext( - flask_app, - dataset_id, - query, - top_k, - score_threshold, - reranking_model, - all_documents, - retrieval_method, - exceptions, - document_ids_filter=None, - ): - all_documents.extend(sample_documents[1:]) - - mock_fulltext_search.side_effect = side_effect_fulltext - - mock_processor_instance = Mock() - mock_processor_instance.invoke.return_value = sample_documents - mock_data_processor_class.return_value = mock_processor_instance - - weights = { - "vector_setting": { - "vector_weight": 0.7, - "embedding_provider_name": "openai", - "embedding_model_name": "text-embedding-ada-002", - }, - "keyword_setting": {"keyword_weight": 0.3}, - } - - # Act - results = RetrievalService.retrieve( - retrieval_method=RetrievalMethod.HYBRID_SEARCH, - dataset_id=mock_dataset.id, - query="test query", - top_k=3, - weights=weights, - reranking_mode="weighted_score", - ) - - # Assert - assert len(results) == 3 - # Verify DataPostProcessor was created with weights - mock_data_processor_class.assert_called_once() - call_args = mock_data_processor_class.call_args - assert call_args.args[3] == weights - processor_session = call_args.kwargs["session"] - assert isinstance(processor_session, Session) - assert processor_session.get_bind() is retrieval_engine - - @pytest.mark.parametrize("empty_query", ["", None]) - @patch("core.rag.datasource.retrieval_service.DataPostProcessor") - @patch("core.rag.datasource.retrieval_service.RetrievalService.embedding_search") - @patch("core.rag.datasource.retrieval_service.RetrievalService._get_dataset") - def test_hybrid_search_attachment_only_uses_image_query_type( - self, - mock_get_dataset, - mock_embedding_search, - mock_data_processor_class, - mock_dataset, - sample_documents, - empty_query, - retrieval_engine, - ): - """ - Regression test for GH #37116: attachment-only hybrid retrieval must use IMAGE_QUERY. - - When HYBRID_SEARCH is invoked with no text query and a non-None attachment_id, - DataPostProcessor.invoke must receive query_type=QueryType.IMAGE_QUERY. - """ - from core.rag.index_processor.constant.query_type import QueryType - - # Arrange - mock_get_dataset.return_value = mock_dataset - attachment_id = "upload-file-uuid-1234" - - def side_effect_embedding( - flask_app, - dataset_id, - query, - top_k, - score_threshold, - reranking_model, - all_documents, - retrieval_method, - exceptions, - document_ids_filter=None, - query_type=QueryType.TEXT_QUERY, - ): - all_documents.extend(sample_documents[:2]) - - mock_embedding_search.side_effect = side_effect_embedding - - mock_processor_instance = Mock() - mock_processor_instance.invoke.return_value = sample_documents[:2] - mock_data_processor_class.return_value = mock_processor_instance - - # Act: call retrieve with attachment_ids only, no text query - RetrievalService.retrieve( - retrieval_method=RetrievalMethod.HYBRID_SEARCH, - dataset_id=mock_dataset.id, - query=empty_query, - top_k=3, - score_threshold=0.5, - attachment_ids=[attachment_id], - ) - - # Assert: invoke must have been called with IMAGE_QUERY - mock_processor_instance.invoke.assert_called_once() - invoke_kwargs = mock_processor_instance.invoke.call_args.kwargs - assert invoke_kwargs["query_type"] == QueryType.IMAGE_QUERY, ( - "Attachment-only hybrid search must use IMAGE_QUERY for reranking, not TEXT_QUERY" - ) - assert invoke_kwargs["query"] == attachment_id, ( - "The rerank query must be the attachment_id, not the empty text query" - ) - processor_session = mock_data_processor_class.call_args.kwargs["session"] - assert isinstance(processor_session, Session) - assert processor_session.get_bind() is retrieval_engine - - # ==================== Full-Text Search Tests ==================== - - @patch("core.rag.datasource.retrieval_service.RetrievalService.full_text_index_search") - @patch("core.rag.datasource.retrieval_service.RetrievalService._get_dataset") - def test_fulltext_search_basic(self, mock_get_dataset, mock_fulltext_search, mock_dataset, sample_documents): - """ - Test basic full-text search functionality. - - Verifies: - - Full-text search is invoked correctly - - Results are returned in expected format - """ - # Arrange - mock_get_dataset.return_value = mock_dataset - - def side_effect_fulltext( - flask_app, - dataset_id, - query, - top_k, - score_threshold, - reranking_model, - all_documents, - retrieval_method, - exceptions, - document_ids_filter=None, - ): - all_documents.extend(sample_documents) - - mock_fulltext_search.side_effect = side_effect_fulltext - - # Act - results = RetrievalService.retrieve( - retrieval_method=RetrievalMethod.FULL_TEXT_SEARCH, - dataset_id=mock_dataset.id, - query="programming language", - top_k=3, - ) - - # Assert - assert len(results) == 3 - mock_fulltext_search.assert_called_once() - - # ==================== Score Merging Tests ==================== - - def test_deduplicate_documents_basic(self): - """ - Test basic document deduplication logic. - - Verifies: - - Documents with same doc_id are deduplicated - - First occurrence is kept by default - """ - # Arrange - doc1 = Document( - page_content="Content 1", - metadata={"doc_id": "doc1", "score": 0.8}, - provider="dify", - ) - doc2 = Document( - page_content="Content 2", - metadata={"doc_id": "doc2", "score": 0.7}, - provider="dify", - ) - doc1_duplicate = Document( - page_content="Content 1 duplicate", - metadata={"doc_id": "doc1", "score": 0.6}, - provider="dify", - ) - - documents = [doc1, doc2, doc1_duplicate] - - # Act - result = RetrievalService._deduplicate_documents(documents) - - # Assert - assert len(result) == 2 - doc_ids = [doc.metadata["doc_id"] for doc in result] - assert doc_ids == ["doc1", "doc2"] - - def test_deduplicate_documents_keeps_higher_score(self): - """ - Test that deduplication keeps document with higher score. - - Verifies: - - When duplicates exist, higher scoring version is retained - - Score comparison works correctly - """ - # Arrange - doc_low = Document( - page_content="Content", - metadata={"doc_id": "doc1", "score": 0.5}, - provider="dify", - ) - doc_high = Document( - page_content="Content", - metadata={"doc_id": "doc1", "score": 0.9}, - provider="dify", - ) - - # Low score first - documents = [doc_low, doc_high] - - # Act - result = RetrievalService._deduplicate_documents(documents) - - # Assert - assert len(result) == 1 - assert result[0].metadata["score"] == 0.9 - - def test_deduplicate_documents_empty_list(self): - """ - Test deduplication with empty document list. - - Verifies: - - Empty list returns empty list - - No errors are raised - """ - # Act - result = RetrievalService._deduplicate_documents([]) - - # Assert - assert result == [] - - def test_deduplicate_documents_non_dify_provider(self): - """ - Test deduplication with non-dify provider documents that have no doc_id. - - Verifies: - - External provider documents without doc_id use content-based deduplication - - Identical content from the same provider is collapsed to one result - """ - # Arrange - doc1 = Document( - page_content="External content", - metadata={"score": 0.8}, - provider="external", - ) - doc2 = Document( - page_content="External content", - metadata={"score": 0.7}, - provider="external", - ) - - documents = [doc1, doc2] - - # Act - result = RetrievalService._deduplicate_documents(documents) - - # Assert - # External documents without doc_id should use content-based dedup - assert len(result) == 1 - - def test_deduplicate_documents_non_dify_provider_with_doc_id_different_sources(self): - """ - Regression test for issue #35707. - - Two chunks from different source documents share identical text content but carry - different doc_ids. Before the fix, non-dify providers were forced into content-based - deduplication and the second chunk was silently dropped. After the fix, doc_id is used - as the dedup key for any provider that exposes it, so both chunks must be retained. - - Verifies: - - Non-dify provider documents with different doc_ids are NOT deduplicated even when - their page_content is identical. - """ - # Arrange — same content, different doc_ids, non-dify provider (e.g. Weaviate / Qdrant) - doc_a = Document( - page_content="Shared identical content", - metadata={"doc_id": "doc-from-file-a", "score": 0.85}, - provider="weaviate", - ) - doc_b = Document( - page_content="Shared identical content", - metadata={"doc_id": "doc-from-file-b", "score": 0.82}, - provider="weaviate", - ) - - # Act - result = RetrievalService._deduplicate_documents([doc_a, doc_b]) - - # Assert — both documents must be kept; losing either silently drops a source citation - assert len(result) == 2 - doc_ids = {doc.metadata["doc_id"] for doc in result} - assert doc_ids == {"doc-from-file-a", "doc-from-file-b"} - - def test_deduplicate_documents_non_dify_provider_with_same_doc_id(self): - """ - Test that non-dify provider documents sharing the same doc_id are deduplicated by - doc_id key (not by content), and the higher-scored duplicate is retained. - - Verifies: - - doc_id-based deduplication now applies to any provider, not only "dify" - - The document with the highest score wins when doc_ids collide - """ - # Arrange - doc_low = Document( - page_content="Content A", - metadata={"doc_id": "chunk-1", "score": 0.5}, - provider="qdrant", - ) - doc_high = Document( - page_content="Content A", - metadata={"doc_id": "chunk-1", "score": 0.9}, - provider="qdrant", - ) - - # Act - result = RetrievalService._deduplicate_documents([doc_low, doc_high]) - - # Assert - assert len(result) == 1 - assert result[0].metadata["score"] == 0.9 - - def test_deduplicate_documents_dify_provider_without_doc_id_falls_back_to_content(self): - """ - Test that a dify provider document without doc_id still falls back to content-based - deduplication (no regression from original behaviour). - - Verifies: - - Absence of doc_id triggers content-based dedup regardless of provider - - First occurrence is kept when content is identical - """ - # Arrange — dify docs with no doc_id, same content - doc1 = Document( - page_content="Same content", - metadata={"score": 0.8}, - provider="dify", - ) - doc2 = Document( - page_content="Same content", - metadata={"score": 0.9}, - provider="dify", - ) - - # Act - result = RetrievalService._deduplicate_documents([doc1, doc2]) - - # Assert — collapsed to one; first-seen wins (no score comparison in content branch) - assert len(result) == 1 - assert result[0].metadata["score"] == 0.8 - - # ==================== Metadata Filtering Tests ==================== - - @patch("core.rag.datasource.retrieval_service.RetrievalService._retrieve") - @patch("core.rag.datasource.retrieval_service.RetrievalService._get_dataset") - def test_vector_search_with_metadata_filter(self, mock_get_dataset, mock_retrieve, mock_dataset, sample_documents): - """ - Test vector search with metadata-based document filtering. - - Verifies: - - Metadata filters are applied correctly - - Only documents matching metadata criteria are returned - """ - # Arrange - mock_get_dataset.return_value = mock_dataset - - # Add metadata to documents - filtered_doc = sample_documents[0] - filtered_doc.metadata["category"] = "programming" - - def side_effect_retrieve( - flask_app, - retrieval_method, - dataset, - query=None, - top_k=4, - score_threshold=None, - reranking_model=None, - reranking_mode="reranking_model", - weights=None, - document_ids_filter=None, - attachment_id=None, - all_documents=None, - exceptions=None, - ): - if all_documents is not None: - all_documents.append(filtered_doc) - - mock_retrieve.side_effect = side_effect_retrieve - - # Act - results = RetrievalService.retrieve( - retrieval_method=RetrievalMethod.SEMANTIC_SEARCH, - dataset_id=mock_dataset.id, - query="Python", - top_k=5, - document_ids_filter=[filtered_doc.metadata["document_id"]], - ) - - # Assert - assert len(results) == 1 - assert results[0].metadata.get("category") == "programming" - - # ==================== Error Handling Tests ==================== - - @patch("core.rag.datasource.retrieval_service.RetrievalService._get_dataset") - def test_retrieve_with_empty_query(self, mock_get_dataset, mock_dataset): - """ - Test retrieval with empty query string. - - Verifies: - - Empty query returns empty results - - No search operations are performed - """ - # Arrange - mock_get_dataset.return_value = mock_dataset - - # Act - results = RetrievalService.retrieve( - retrieval_method=RetrievalMethod.SEMANTIC_SEARCH, - dataset_id=mock_dataset.id, - query="", - top_k=5, - ) - - # Assert - assert results == [] - - @patch("core.rag.datasource.retrieval_service.RetrievalService._get_dataset") - def test_retrieve_with_nonexistent_dataset(self, mock_get_dataset): - """ - Test retrieval with non-existent dataset ID. - - Verifies: - - Non-existent dataset returns empty results - - No errors are raised - """ - # Arrange - mock_get_dataset.return_value = None - - # Act - results = RetrievalService.retrieve( - retrieval_method=RetrievalMethod.SEMANTIC_SEARCH, - dataset_id="nonexistent_id", - query="test query", - top_k=5, - ) - - # Assert - assert results == [] - - @patch("core.rag.datasource.retrieval_service.RetrievalService._retrieve") - @patch("core.rag.datasource.retrieval_service.RetrievalService._get_dataset") - def test_retrieve_with_exception_handling(self, mock_get_dataset, mock_retrieve, mock_dataset): - """ - Test that exceptions during retrieval are properly handled. - - Verifies: - - Exceptions are caught and added to exceptions list - - ValueError is raised with exception messages - """ - # Arrange - mock_get_dataset.return_value = mock_dataset - - # Make _retrieve add an exception to the exceptions list - def side_effect_with_exception( - flask_app, - retrieval_method, - dataset, - query=None, - top_k=4, - score_threshold=None, - reranking_model=None, - reranking_mode="reranking_model", - weights=None, - document_ids_filter=None, - attachment_id=None, - all_documents=None, - exceptions=None, - ): - if exceptions is not None: - exceptions.append("Search failed") - - mock_retrieve.side_effect = side_effect_with_exception - - # Act & Assert - with pytest.raises(ValueError, match="Search failed") as exc_info: - RetrievalService.retrieve( - retrieval_method=RetrievalMethod.SEMANTIC_SEARCH, - dataset_id=mock_dataset.id, - query="test query", - top_k=5, - ) - - assert "Search failed" in str(exc_info.value) - - # ==================== Score Threshold Tests ==================== - - @patch("core.rag.datasource.retrieval_service.RetrievalService._retrieve") - @patch("core.rag.datasource.retrieval_service.RetrievalService._get_dataset") - def test_vector_search_with_score_threshold(self, mock_get_dataset, mock_retrieve, mock_dataset): - """ - Test vector search with score threshold filtering. - - Verifies: - - Score threshold is passed to search method - - Documents below threshold are filtered out - """ - # Arrange - mock_get_dataset.return_value = mock_dataset - - # Only return documents above threshold - high_score_doc = Document( - page_content="High relevance content", - metadata={"doc_id": "doc1", "score": 0.85}, - provider="dify", - ) - - def side_effect_retrieve( - flask_app, - retrieval_method, - dataset, - query=None, - top_k=4, - score_threshold=None, - reranking_model=None, - reranking_mode="reranking_model", - weights=None, - document_ids_filter=None, - attachment_id=None, - all_documents=None, - exceptions=None, - ): - if all_documents is not None: - all_documents.append(high_score_doc) - - mock_retrieve.side_effect = side_effect_retrieve - - score_threshold = 0.8 - - # Act - results = RetrievalService.retrieve( - retrieval_method=RetrievalMethod.SEMANTIC_SEARCH, - dataset_id=mock_dataset.id, - query="test query", - top_k=5, - score_threshold=score_threshold, - ) - - # Assert - assert len(results) == 1 - assert results[0].metadata["score"] >= score_threshold - - # ==================== Top-K Limiting Tests ==================== - - @patch("core.rag.datasource.retrieval_service.RetrievalService._retrieve") - @patch("core.rag.datasource.retrieval_service.RetrievalService._get_dataset") - def test_retrieve_respects_top_k_limit(self, mock_get_dataset, mock_retrieve, mock_dataset): - """ - Test that retrieval respects top_k parameter. - - Verifies: - - Only top_k documents are returned - - Limit is applied correctly - """ - # Arrange - mock_get_dataset.return_value = mock_dataset - - # Create more documents than top_k - many_docs = [ - Document( - page_content=f"Content {i}", - metadata={"doc_id": f"doc{i}", "score": 0.9 - i * 0.1}, - provider="dify", - ) - for i in range(10) - ] - - def side_effect_retrieve( - flask_app, - retrieval_method, - dataset, - query=None, - top_k=4, - score_threshold=None, - reranking_model=None, - reranking_mode="reranking_model", - weights=None, - document_ids_filter=None, - attachment_id=None, - all_documents=None, - exceptions=None, - ): - # Return only top_k documents - if all_documents is not None: - all_documents.extend(many_docs[:top_k]) - - mock_retrieve.side_effect = side_effect_retrieve - - top_k = 3 - - # Act - results = RetrievalService.retrieve( - retrieval_method=RetrievalMethod.SEMANTIC_SEARCH, - dataset_id=mock_dataset.id, - query="test query", - top_k=top_k, - ) - - # Assert - # Verify _retrieve was called - assert mock_retrieve.called - call_kwargs = mock_retrieve.call_args.kwargs - assert call_kwargs["top_k"] == top_k - # Verify we got the right number of results - assert len(results) == top_k - - # ==================== Query Escaping Tests ==================== - - def test_escape_query_for_search(self): - """ - Test query escaping for special characters. - - Verifies: - - Double quotes are properly escaped - - Other characters remain unchanged - """ - # Test cases with expected outputs - test_cases = [ - ("simple query", "simple query"), - ('query with "quotes"', 'query with \\"quotes\\"'), - ('"quoted phrase"', '\\"quoted phrase\\"'), - ("no special chars", "no special chars"), - ] - - for input_query, expected_output in test_cases: - result = RetrievalService.escape_query_for_search(input_query) - assert result == expected_output - - # ==================== Reranking Tests ==================== - - @patch("core.rag.datasource.retrieval_service.RetrievalService._retrieve") - @patch("core.rag.datasource.retrieval_service.RetrievalService._get_dataset") - def test_semantic_search_with_reranking(self, mock_get_dataset, mock_retrieve, mock_dataset, sample_documents): - """ - Test semantic search with reranking model. - - Verifies: - - Reranking is applied when configured - - DataPostProcessor is invoked with correct parameters - """ - # Arrange - mock_get_dataset.return_value = mock_dataset - - # Simulate reranking changing order - reranked_docs = list(reversed(sample_documents)) - - def side_effect_retrieve( - flask_app, - retrieval_method, - dataset, - query=None, - top_k=4, - score_threshold=None, - reranking_model=None, - reranking_mode="reranking_model", - weights=None, - document_ids_filter=None, - attachment_id=None, - all_documents=None, - exceptions=None, - ): - # _retrieve handles reranking internally - if all_documents is not None: - all_documents.extend(reranked_docs) - - mock_retrieve.side_effect = side_effect_retrieve - - reranking_model = { - "reranking_provider_name": "cohere", - "reranking_model_name": "rerank-english-v2.0", - } - - # Act - results = RetrievalService.retrieve( - retrieval_method=RetrievalMethod.SEMANTIC_SEARCH, - dataset_id=mock_dataset.id, - query="test query", - top_k=3, - reranking_model=reranking_model, - ) - - # Assert - # For semantic search with reranking, reranking_model should be passed - assert len(results) == 3 - call_kwargs = mock_retrieve.call_args.kwargs - assert call_kwargs["reranking_model"] == reranking_model - - # ==================== Multiple Retrieve Thread Tests ==================== - - @patch("core.rag.retrieval.dataset_retrieval.DataPostProcessor") - @patch("core.rag.retrieval.dataset_retrieval.DatasetRetrieval._retriever") - def test_multiple_retrieve_thread_skips_second_reranking_with_single_dataset( - self, mock_retriever, mock_data_processor_class, mock_flask_app, mock_dataset - ): - """ - Test that _multiple_retrieve_thread skips second reranking when dataset_count is 1. - - When there is only one dataset, the second reranking is unnecessary - because the documents are already ranked from the first retrieval. - This optimization avoids the overhead of reranking when it won't - provide any benefit. - - Verifies: - - DataPostProcessor is NOT called when dataset_count == 1 - - Documents are still added to all_documents - - Standard scoring logic is applied instead - """ - # Arrange - dataset_retrieval = DatasetRetrieval() - tenant_id = str(uuid4()) - - # Create test documents - doc1 = Document( - page_content="Test content 1", - metadata={"doc_id": "doc1", "score": 0.9, "document_id": str(uuid4()), "dataset_id": mock_dataset.id}, - provider="dify", - ) - doc2 = Document( - page_content="Test content 2", - metadata={"doc_id": "doc2", "score": 0.8, "document_id": str(uuid4()), "dataset_id": mock_dataset.id}, - provider="dify", - ) - - # Mock _retriever to return documents - def side_effect_retriever( - flask_app, - session, - dataset_id, - query, - top_k, - all_documents, - document_ids_filter, - metadata_condition, - attachment_ids, - ): - all_documents.extend([doc1, doc2]) - - mock_retriever.side_effect = side_effect_retriever - - # Set up dataset with high_quality indexing - mock_dataset.indexing_technique = "high_quality" - - all_documents = [] - - # Act - Call with dataset_count = 1 - dataset_retrieval._multiple_retrieve_thread( - flask_app=mock_flask_app, - available_datasets=[mock_dataset], - metadata_condition=None, - metadata_filter_document_ids=None, - all_documents=all_documents, - tenant_id=tenant_id, - reranking_enable=True, - reranking_mode="reranking_model", - reranking_model={"reranking_provider_name": "cohere", "reranking_model_name": "rerank-v2"}, - weights=None, - top_k=5, - score_threshold=0.5, - query="test query", - attachment_id=None, - dataset_count=1, # Single dataset - should skip second reranking - ) - - # Assert - # DataPostProcessor should NOT be called (second reranking skipped) - mock_data_processor_class.assert_not_called() - - # Documents should still be added to all_documents - assert len(all_documents) == 2 - assert all_documents[0].page_content == "Test content 1" - assert all_documents[1].page_content == "Test content 2" - - @patch("core.rag.retrieval.dataset_retrieval.DataPostProcessor") - @patch("core.rag.retrieval.dataset_retrieval.DatasetRetrieval._retriever") - @patch("core.rag.retrieval.dataset_retrieval.DatasetRetrieval.calculate_vector_score") - def test_multiple_retrieve_thread_performs_second_reranking_with_multiple_datasets( - self, mock_calculate_vector_score, mock_retriever, mock_data_processor_class, mock_flask_app, mock_dataset - ): - """ - Test that _multiple_retrieve_thread performs second reranking when dataset_count > 1. - - When there are multiple datasets, the second reranking is necessary - to merge and re-rank results from different datasets. This ensures - the most relevant documents across all datasets are returned. - - Verifies: - - DataPostProcessor IS called when dataset_count > 1 - - Reranking is applied with correct parameters - - Documents are processed correctly - """ - # Arrange - dataset_retrieval = DatasetRetrieval() - tenant_id = str(uuid4()) - - # Create test documents - doc1 = Document( - page_content="Test content 1", - metadata={"doc_id": "doc1", "score": 0.7, "document_id": str(uuid4()), "dataset_id": mock_dataset.id}, - provider="dify", - ) - doc2 = Document( - page_content="Test content 2", - metadata={"doc_id": "doc2", "score": 0.6, "document_id": str(uuid4()), "dataset_id": mock_dataset.id}, - provider="dify", - ) - - # Mock _retriever to return documents - def side_effect_retriever( - flask_app, - session, - dataset_id, - query, - top_k, - all_documents, - document_ids_filter, - metadata_condition, - attachment_ids, - ): - all_documents.extend([doc1, doc2]) - - mock_retriever.side_effect = side_effect_retriever - - # Set up dataset with high_quality indexing - mock_dataset.indexing_technique = "high_quality" - - # Mock DataPostProcessor instance and its invoke method - mock_processor_instance = Mock() - # Simulate reranking - return documents in reversed order with updated scores - reranked_docs = [ - Document( - page_content="Test content 2", - metadata={"doc_id": "doc2", "score": 0.95, "document_id": str(uuid4()), "dataset_id": mock_dataset.id}, - provider="dify", - ), - Document( - page_content="Test content 1", - metadata={"doc_id": "doc1", "score": 0.85, "document_id": str(uuid4()), "dataset_id": mock_dataset.id}, - provider="dify", - ), - ] - mock_processor_instance.invoke.return_value = reranked_docs - mock_data_processor_class.return_value = mock_processor_instance - - all_documents = [] - - # Create second dataset - mock_dataset2 = Dataset( - id=str(uuid4()), - indexing_technique="high_quality", - provider="dify", - ) - - # Act - Call with dataset_count = 2 - dataset_retrieval._multiple_retrieve_thread( - flask_app=mock_flask_app, - available_datasets=[mock_dataset, mock_dataset2], - metadata_condition=None, - metadata_filter_document_ids=None, - all_documents=all_documents, - tenant_id=tenant_id, - reranking_enable=True, - reranking_mode="reranking_model", - reranking_model={"reranking_provider_name": "cohere", "reranking_model_name": "rerank-v2"}, - weights=None, - top_k=5, - score_threshold=0.5, - query="test query", - attachment_id=None, - dataset_count=2, # Multiple datasets - should perform second reranking - ) - - # Assert - # DataPostProcessor SHOULD be called (second reranking performed) - mock_data_processor_class.assert_called_once() - assert isinstance(mock_data_processor_class.call_args.kwargs["session"], Session) - - # Verify invoke was called with correct parameters - mock_processor_instance.invoke.assert_called_once() - - # Documents should be added to all_documents after reranking - assert len(all_documents) == 2 - # The reranked order should be reflected - assert all_documents[0].page_content == "Test content 2" - assert all_documents[1].page_content == "Test content 1" - - @patch("core.rag.retrieval.dataset_retrieval.DataPostProcessor") - @patch("core.rag.retrieval.dataset_retrieval.DatasetRetrieval._retriever") - @patch("core.rag.retrieval.dataset_retrieval.DatasetRetrieval.calculate_vector_score") - def test_multiple_retrieve_thread_single_dataset_uses_standard_scoring( - self, mock_calculate_vector_score, mock_retriever, mock_data_processor_class, mock_flask_app, mock_dataset - ): - """ - Test that _multiple_retrieve_thread uses standard scoring when dataset_count is 1 - and reranking is enabled. - - When there's only one dataset, instead of using DataPostProcessor, - the method should fall through to the standard scoring logic - (calculate_vector_score for high_quality datasets). - - Verifies: - - DataPostProcessor is NOT called - - calculate_vector_score IS called for high_quality indexing - - Documents are scored correctly - """ - # Arrange - dataset_retrieval = DatasetRetrieval() - tenant_id = str(uuid4()) - - # Create test documents - doc1 = Document( - page_content="Test content 1", - metadata={"doc_id": "doc1", "score": 0.9, "document_id": str(uuid4()), "dataset_id": mock_dataset.id}, - provider="dify", - ) - doc2 = Document( - page_content="Test content 2", - metadata={"doc_id": "doc2", "score": 0.8, "document_id": str(uuid4()), "dataset_id": mock_dataset.id}, - provider="dify", - ) - - # Mock _retriever to return documents - def side_effect_retriever( - flask_app, - session, - dataset_id, - query, - top_k, - all_documents, - document_ids_filter, - metadata_condition, - attachment_ids, - ): - all_documents.extend([doc1, doc2]) - - mock_retriever.side_effect = side_effect_retriever - - # Set up dataset with high_quality indexing - mock_dataset.indexing_technique = "high_quality" - - # Mock calculate_vector_score to return scored documents - scored_docs = [ - Document( - page_content="Test content 1", - metadata={"doc_id": "doc1", "score": 0.95, "document_id": str(uuid4()), "dataset_id": mock_dataset.id}, - provider="dify", - ), - ] - mock_calculate_vector_score.return_value = scored_docs - - all_documents = [] - - # Act - Call with dataset_count = 1 - dataset_retrieval._multiple_retrieve_thread( - flask_app=mock_flask_app, - available_datasets=[mock_dataset], - metadata_condition=None, - metadata_filter_document_ids=None, - all_documents=all_documents, - tenant_id=tenant_id, - reranking_enable=True, # Reranking enabled but should be skipped for single dataset - reranking_mode="reranking_model", - reranking_model={"reranking_provider_name": "cohere", "reranking_model_name": "rerank-v2"}, - weights=None, - top_k=5, - score_threshold=0.5, - query="test query", - attachment_id=None, - dataset_count=1, - ) - - # Assert - # DataPostProcessor should NOT be called - mock_data_processor_class.assert_not_called() - - # calculate_vector_score SHOULD be called for high_quality datasets - mock_calculate_vector_score.assert_called_once() - call_args = mock_calculate_vector_score.call_args - assert call_args[0][1] == 5 # top_k - - # Documents should be added after standard scoring - assert len(all_documents) == 1 - assert all_documents[0].page_content == "Test content 1" - - -class TestRetrievalMethods: - """ - Test suite for RetrievalMethod enum and utility methods. - - The RetrievalMethod enum defines the available search strategies: - - 1. **SEMANTIC_SEARCH**: Vector-based similarity search using embeddings - - Best for: Natural language queries, conceptual similarity - - Uses: Embedding models (e.g., text-embedding-ada-002) - - Example: "What is machine learning?" matches "AI and ML concepts" - - 2. **FULL_TEXT_SEARCH**: BM25-based text matching - - Best for: Exact phrase matching, keyword presence - - Uses: BM25 algorithm with sparse vectors - - Example: "Python programming" matches documents with those exact terms - - 3. **HYBRID_SEARCH**: Combination of semantic + full-text - - Best for: Comprehensive search with both conceptual and exact matching - - Uses: Both embedding vectors and BM25, with score merging - - Example: Finds both semantically similar and keyword-matching documents - - 4. **KEYWORD_SEARCH**: Traditional keyword-based search (economy mode) - - Best for: Simple, fast searches without embeddings - - Uses: Jieba tokenization and keyword matching - - Example: Basic text search without vector database - - Utility Methods: - ================ - - is_support_semantic_search(): Check if method uses embeddings - - is_support_fulltext_search(): Check if method uses BM25 - - These utilities help determine which search operations to execute - in the RetrievalService.retrieve() method. - """ - - def test_retrieval_method_values(self): - """ - Test that all retrieval method constants are defined correctly. - - This ensures the enum values match the expected string constants - used throughout the codebase for configuration and API calls. - - Verifies: - - All expected retrieval methods exist - - Values are correct strings (not accidentally changed) - - String values match database/config expectations - """ - assert RetrievalMethod.SEMANTIC_SEARCH == "semantic_search" - assert RetrievalMethod.FULL_TEXT_SEARCH == "full_text_search" - assert RetrievalMethod.HYBRID_SEARCH == "hybrid_search" - assert RetrievalMethod.KEYWORD_SEARCH == "keyword_search" - - def test_is_support_semantic_search(self): - """ - Test semantic search support detection. - - Verifies: - - Semantic search method is detected - - Hybrid search method is detected (includes semantic) - - Other methods are not detected - """ - assert RetrievalMethod.is_support_semantic_search(RetrievalMethod.SEMANTIC_SEARCH) is True - assert RetrievalMethod.is_support_semantic_search(RetrievalMethod.HYBRID_SEARCH) is True - assert RetrievalMethod.is_support_semantic_search(RetrievalMethod.FULL_TEXT_SEARCH) is False - assert RetrievalMethod.is_support_semantic_search(RetrievalMethod.KEYWORD_SEARCH) is False - - def test_is_support_fulltext_search(self): - """ - Test full-text search support detection. - - Verifies: - - Full-text search method is detected - - Hybrid search method is detected (includes full-text) - - Other methods are not detected - """ - assert RetrievalMethod.is_support_fulltext_search(RetrievalMethod.FULL_TEXT_SEARCH) is True - assert RetrievalMethod.is_support_fulltext_search(RetrievalMethod.HYBRID_SEARCH) is True - assert RetrievalMethod.is_support_fulltext_search(RetrievalMethod.SEMANTIC_SEARCH) is False - assert RetrievalMethod.is_support_fulltext_search(RetrievalMethod.KEYWORD_SEARCH) is False - - -class TestDocumentModel: - """ - Test suite for Document model used in retrieval. - - The Document class is the core data structure for representing text chunks - in the retrieval system. It's based on Pydantic BaseModel for validation. - - Document Structure: - =================== - - **page_content** (str): The actual text content of the document chunk - - **metadata** (dict): Additional information about the document - - doc_id: Unique identifier for the chunk - - document_id: Parent document ID - - dataset_id: Dataset this document belongs to - - score: Relevance score from search (0.0 to 1.0) - - Custom fields: category, tags, timestamps, etc. - - **provider** (str): Source of the document ("dify" or "external") - - **vector** (list[float] | None): Embedding vector for semantic search - - **children** (list[ChildDocument] | None): Sub-chunks for hierarchical docs - - Document Lifecycle: - =================== - 1. **Creation**: Documents are created when text is indexed - - Content is chunked into manageable pieces - - Embeddings are generated for semantic search - - Metadata is attached for filtering and tracking - - 2. **Storage**: Documents are stored in vector databases - - Vector field stores embeddings - - Metadata enables filtering - - Provider tracks source (internal vs external) - - 3. **Retrieval**: Documents are returned from search operations - - Scores are added during search - - Multiple documents may be combined (hybrid search) - - Deduplication uses doc_id - - 4. **Post-processing**: Documents may be reranked or filtered - - Scores can be recalculated - - Content may be truncated or formatted - - Metadata is used for display - - Why Test the Document Model: - ============================ - - Ensures data structure integrity - - Validates Pydantic model behavior - - Confirms default values work correctly - - Tests equality comparison for deduplication - - Verifies metadata handling - - Related Classes: - ================ - - ChildDocument: For hierarchical document structures - - RetrievalSegments: Combines Document with database segment info - """ - - def test_document_creation_basic(self): - """ - Test basic Document object creation. - - Tests the minimal required fields and default values. - Only page_content is required; all other fields have defaults. - - Verifies: - - Document can be created with minimal fields - - Default values are set correctly - - Pydantic validation works - - No exceptions are raised - """ - doc = Document(page_content="Test content") - - assert doc.page_content == "Test content" - assert doc.metadata == {} # Empty dict by default - assert doc.provider == "dify" # Default provider - assert doc.vector is None # No embedding by default - assert doc.children is None # No child documents by default - - def test_document_creation_with_metadata(self): - """ - Test Document creation with metadata. - - Verifies: - - Metadata is stored correctly - - Metadata can contain various types - """ - metadata = { - "doc_id": "test_doc", - "score": 0.95, - "dataset_id": str(uuid4()), - "category": "test", - } - doc = Document(page_content="Test content", metadata=metadata) - - assert doc.metadata == metadata - assert doc.metadata["score"] == 0.95 - - def test_document_creation_with_vector(self): - """ - Test Document creation with embedding vector. - - Verifies: - - Vector embeddings can be stored - - Vector is optional - """ - vector = [0.1, 0.2, 0.3, 0.4, 0.5] - doc = Document(page_content="Test content", vector=vector) - - assert doc.vector == vector - assert doc.vector is not None - assert len(doc.vector) == 5 - - def test_document_with_external_provider(self): - """ - Test Document with external provider. - - Verifies: - - Provider can be set to external - - External documents are handled correctly - """ - doc = Document(page_content="External content", provider="external") - - assert doc.provider == "external" - - def test_document_equality(self): - """ - Test Document equality comparison. - - Verifies: - - Documents with same content are considered equal - - Metadata affects equality - """ - doc1 = Document(page_content="Content", metadata={"id": "1"}) - doc2 = Document(page_content="Content", metadata={"id": "1"}) - doc3 = Document(page_content="Different", metadata={"id": "1"}) - - assert doc1 == doc2 - assert doc1 != doc3 - - # ==================== Helper Functions ==================== @@ -2245,7 +228,7 @@ def create_mock_document_methods( score: float = 0.8, provider: str = "dify", additional_metadata: dict[str, Any] | None = None, -) -> Document: +) -> _dataset_retrieval_Document: """ Create a mock Document object for testing. @@ -2269,1614 +252,13 @@ def create_mock_document_methods( if additional_metadata: metadata.update(additional_metadata) - return Document( + return _dataset_retrieval_Document( page_content=content, metadata=metadata, provider=provider, ) -# ==================== Test _check_knowledge_rate_limit ==================== - - -class TestCheckKnowledgeRateLimit: - """ - Test suite for _check_knowledge_rate_limit method. - - The _check_knowledge_rate_limit method validates whether a tenant has - exceeded their knowledge retrieval rate limit. This is important for: - - Preventing abuse of the knowledge retrieval system - - Enforcing subscription plan limits - - Tracking usage for billing purposes - - Test Cases: - ============ - 1. Rate limit disabled - no exception raised - 2. Rate limit enabled but not exceeded - no exception raised - 3. Rate limit enabled and exceeded - RateLimitExceededError raised - 4. Redis operations are performed correctly - 5. RateLimitLog is created when limit is exceeded - """ - - @patch("core.rag.retrieval.dataset_retrieval.FeatureService") - @patch("core.rag.retrieval.dataset_retrieval.redis_client") - def test_rate_limit_disabled_no_exception(self, mock_redis, mock_feature_service): - """ - Test that when rate limit is disabled, no exception is raised. - - This test verifies the behavior when the tenant's subscription - does not have rate limiting enabled. - - Verifies: - - FeatureService.get_knowledge_rate_limit is called - - No Redis operations are performed - - No exception is raised - - Retrieval proceeds normally - """ - # Arrange - tenant_id = str(uuid4()) - dataset_retrieval = DatasetRetrieval() - - # Mock rate limit disabled - mock_limit = Mock() - mock_limit.enabled = False - mock_feature_service.get_knowledge_rate_limit.return_value = mock_limit - - # Act & Assert - should not raise any exception - dataset_retrieval._check_knowledge_rate_limit(tenant_id) - - # Verify FeatureService was called - mock_feature_service.get_knowledge_rate_limit.assert_called_once_with(tenant_id) - - # Verify no Redis operations were performed - assert not mock_redis.zadd.called - assert not mock_redis.zremrangebyscore.called - assert not mock_redis.zcard.called - - @patch("core.rag.retrieval.dataset_retrieval.FeatureService") - @patch("core.rag.retrieval.dataset_retrieval.redis_client") - @patch("core.rag.retrieval.dataset_retrieval.time") - def test_rate_limit_enabled_not_exceeded(self, mock_time, mock_redis, mock_feature_service): - """ - Test that when rate limit is enabled but not exceeded, no exception is raised. - - This test simulates a tenant making requests within their rate limit. - The Redis sorted set stores timestamps of recent requests, and old - requests (older than 60 seconds) are removed. - - Verifies: - - Redis zadd is called to track the request - - Redis zremrangebyscore removes old entries - - Redis zcard returns count within limit - - No exception is raised - """ - # Arrange - tenant_id = str(uuid4()) - dataset_retrieval = DatasetRetrieval() - - # Mock rate limit enabled with limit of 100 requests per minute - mock_limit = Mock() - mock_limit.enabled = True - mock_limit.limit = 100 - mock_limit.subscription_plan = "professional" - mock_feature_service.get_knowledge_rate_limit.return_value = mock_limit - - # Mock time - current_time = 1234567890000 # Current time in milliseconds - mock_time.time.return_value = current_time / 1000 # Return seconds - mock_time.time.__mul__ = lambda self, x: int(self * x) # Multiply to get milliseconds - - # Mock Redis operations - # zcard returns 50 (within limit of 100) - mock_redis.zcard.return_value = 50 - - # Act & Assert - should not raise any exception - dataset_retrieval._check_knowledge_rate_limit(tenant_id) - - # Verify Redis operations - expected_key = f"rate_limit_{tenant_id}" - mock_redis.zadd.assert_called_once_with(expected_key, {current_time: current_time}) - mock_redis.zremrangebyscore.assert_called_once_with(expected_key, 0, current_time - 60000) - mock_redis.zcard.assert_called_once_with(expected_key) - - @patch("core.rag.retrieval.dataset_retrieval.FeatureService") - @patch("core.rag.retrieval.dataset_retrieval.redis_client") - @patch("core.rag.retrieval.dataset_retrieval.time") - def test_rate_limit_enabled_exceeded_raises_exception( - self, mock_time, mock_redis, mock_feature_service, sqlite_session: Session - ): - """ - Test that when rate limit is enabled and exceeded, RateLimitExceededError is raised. - - This test simulates a tenant exceeding their rate limit. When the count - of recent requests exceeds the limit, an exception should be raised and - a RateLimitLog should be created. - - Verifies: - - Redis zcard returns count exceeding limit - - RateLimitExceededError is raised with correct message - - RateLimitLog is created in database - - Session operations are performed correctly - """ - # Arrange - tenant_id = str(uuid4()) - dataset_retrieval = DatasetRetrieval() - - # Mock rate limit enabled with limit of 100 requests per minute - mock_limit = Mock() - mock_limit.enabled = True - mock_limit.limit = 100 - mock_limit.subscription_plan = "professional" - mock_feature_service.get_knowledge_rate_limit.return_value = mock_limit - - # Mock time - current_time = 1234567890000 - mock_time.time.return_value = current_time / 1000 - - # Mock Redis operations - return count exceeding limit - mock_redis.zcard.return_value = 150 # Exceeds limit of 100 - - # Act & Assert - with pytest.raises(exc.RateLimitExceededError) as exc_info: - dataset_retrieval._check_knowledge_rate_limit(tenant_id) - - # Verify exception message - assert "knowledge base request rate limit" in str(exc_info.value) - - # Verify the independently committed audit row from an observer session. - added_log = sqlite_session.scalar(select(RateLimitLog).where(RateLimitLog.tenant_id == tenant_id)) - assert added_log is not None - assert added_log.tenant_id == tenant_id - assert added_log.subscription_plan == "professional" - assert added_log.operation == "knowledge" - - -# ==================== Test _get_available_datasets ==================== - - -class TestGetAvailableDatasets: - """ - Test suite for _get_available_datasets method. - - The _get_available_datasets method retrieves datasets that are available - for retrieval. A dataset is considered available if: - - It belongs to the specified tenant - - It's in the list of requested dataset_ids - - It has at least one completed, enabled, non-archived document OR - - It's an external provider dataset - - Note: Due to SQLAlchemy subquery complexity, full testing is done in - integration tests. Unit tests here verify basic behavior. - """ - - def test_method_exists_and_has_correct_signature(self): - """ - Test that the method exists and has the correct signature. - - Verifies: - - Method exists on DatasetRetrieval class - - Accepts tenant_id and dataset_ids parameters - """ - # Arrange - dataset_retrieval = DatasetRetrieval() - - # Assert - method exists - assert hasattr(dataset_retrieval, "_get_available_datasets") - # Assert - method is callable - assert callable(dataset_retrieval._get_available_datasets) - - -# ==================== Test knowledge_retrieval ==================== - - -class TestDatasetRetrievalKnowledgeRetrieval: - """ - Test suite for knowledge_retrieval method. - - The knowledge_retrieval method is the main entry point for retrieving - knowledge from datasets. It orchestrates the entire retrieval process: - 1. Checks rate limits - 2. Gets available datasets - 3. Applies metadata filtering if enabled - 4. Performs retrieval (single or multiple mode) - 5. Formats and returns results - - Test Cases: - ============ - 1. Single mode retrieval - 2. Multiple mode retrieval - 3. Metadata filtering disabled - 4. Metadata filtering automatic - 5. Metadata filtering manual - 6. External documents handling - 7. Dify documents handling - 8. Empty results handling - 9. Rate limit exceeded - 10. No available datasets - """ - - def test_knowledge_retrieval_single_mode_basic(self): - """ - Test knowledge_retrieval in single retrieval mode - basic check. - - Note: Full single mode testing requires complex model mocking and - is better suited for integration tests. This test verifies the - method accepts single mode requests. - - Verifies: - - Method can accept single mode request - - Request parameters are correctly structured - """ - # Arrange - tenant_id = str(uuid4()) - user_id = str(uuid4()) - app_id = str(uuid4()) - dataset_id = str(uuid4()) - - request = KnowledgeRetrievalRequest( - tenant_id=tenant_id, - user_id=user_id, - app_id=app_id, - user_from="web", - dataset_ids=[dataset_id], - query="What is Python?", - retrieval_mode="single", - model_provider="openai", - model_name="gpt-4", - model_mode="chat", - completion_params={"temperature": 0.7}, - ) - - # Assert - request is properly structured - assert request.retrieval_mode == "single" - assert request.model_provider == "openai" - assert request.model_name == "gpt-4" - assert request.model_mode == "chat" - - def test_knowledge_retrieval_multiple_mode(self, sqlite_session: Session): - """ - Test knowledge_retrieval in multiple retrieval mode. - - In multiple mode, retrieval is performed across all datasets and - results are combined and reranked. - - Verifies: - - Rate limit is checked - - Available datasets are retrieved - - Multiple retrieval is performed - - Results are combined and reranked - - Results are formatted correctly - """ - # Arrange - tenant_id = str(uuid4()) - user_id = str(uuid4()) - app_id = str(uuid4()) - dataset_id1 = str(uuid4()) - dataset_id2 = str(uuid4()) - - request = KnowledgeRetrievalRequest( - tenant_id=tenant_id, - user_id=user_id, - app_id=app_id, - user_from="web", - dataset_ids=[dataset_id1, dataset_id2], - query="What is Python?", - retrieval_mode="multiple", - top_k=5, - score_threshold=0.7, - reranking_enable=True, - reranking_mode="reranking_model", - reranking_model={"reranking_provider_name": "cohere", "reranking_model_name": "rerank-v2"}, - ) - - dataset_retrieval = DatasetRetrieval() - - # Mock _check_knowledge_rate_limit - with patch.object(dataset_retrieval, "_check_knowledge_rate_limit"): - # Mock _get_available_datasets - mock_dataset1 = create_mock_dataset_methods(dataset_id=dataset_id1, tenant_id=tenant_id) - mock_dataset2 = create_mock_dataset_methods(dataset_id=dataset_id2, tenant_id=tenant_id) - with patch.object( - dataset_retrieval, "_get_available_datasets", return_value=[mock_dataset1, mock_dataset2] - ): - # Mock get_metadata_filter_condition - with patch.object(dataset_retrieval, "get_metadata_filter_condition", return_value=(None, None)): - # Mock multiple_retrieve to return documents - doc1 = create_mock_document_methods( - "Python is great", - "doc1", - score=0.9, - provider="external", - additional_metadata={"dataset_name": "test dataset", "title": "Python"}, - ) - doc2 = create_mock_document_methods( - "Python is awesome", - "doc2", - score=0.8, - provider="external", - additional_metadata={"dataset_name": "test dataset", "title": "Python 2"}, - ) - with patch.object( - dataset_retrieval, "multiple_retrieve", return_value=[doc1, doc2] - ) as mock_multiple_retrieve: - result = dataset_retrieval.knowledge_retrieval(sqlite_session, request) - - assert len(result) == 2 - mock_multiple_retrieve.assert_called_once() - - def test_knowledge_retrieval_metadata_filtering_disabled(self, sqlite_session: Session): - """ - Test knowledge_retrieval with metadata filtering disabled. - - When metadata filtering is disabled, get_metadata_filter_condition is - NOT called (the method checks metadata_filtering_mode != "disabled"). - - Verifies: - - get_metadata_filter_condition is NOT called when mode is "disabled" - - Retrieval proceeds without metadata filters - """ - # Arrange - tenant_id = str(uuid4()) - user_id = str(uuid4()) - app_id = str(uuid4()) - dataset_id = str(uuid4()) - - request = KnowledgeRetrievalRequest( - tenant_id=tenant_id, - user_id=user_id, - app_id=app_id, - user_from="web", - dataset_ids=[dataset_id], - query="What is Python?", - retrieval_mode="multiple", - metadata_filtering_mode="disabled", - top_k=5, - ) - - dataset_retrieval = DatasetRetrieval() - - # Mock dependencies - with patch.object(dataset_retrieval, "_check_knowledge_rate_limit"): - mock_dataset = create_mock_dataset_methods(dataset_id=dataset_id, tenant_id=tenant_id) - with patch.object(dataset_retrieval, "_get_available_datasets", return_value=[mock_dataset]): - # Mock get_metadata_filter_condition - should NOT be called when disabled - with patch.object( - dataset_retrieval, - "get_metadata_filter_condition", - return_value=(None, None), - ) as mock_get_metadata: - with patch.object(dataset_retrieval, "multiple_retrieve", return_value=[]): - # Act - result = dataset_retrieval.knowledge_retrieval(sqlite_session, request) - - # Assert - assert isinstance(result, list) - # get_metadata_filter_condition should NOT be called when mode is "disabled" - mock_get_metadata.assert_not_called() - - def test_knowledge_retrieval_with_external_documents(self, sqlite_session: Session): - """ - Test knowledge_retrieval with external documents. - - External documents come from external knowledge bases and should - be formatted differently than Dify documents. - - Verifies: - - External documents are handled correctly - - Provider is set to "external" - - Metadata includes external-specific fields - """ - # Arrange - tenant_id = str(uuid4()) - user_id = str(uuid4()) - app_id = str(uuid4()) - dataset_id = str(uuid4()) - - request = KnowledgeRetrievalRequest( - tenant_id=tenant_id, - user_id=user_id, - app_id=app_id, - user_from="web", - dataset_ids=[dataset_id], - query="What is Python?", - retrieval_mode="multiple", - top_k=5, - ) - - dataset_retrieval = DatasetRetrieval() - - # Mock dependencies - with patch.object(dataset_retrieval, "_check_knowledge_rate_limit"): - mock_dataset = create_mock_dataset_methods(dataset_id=dataset_id, tenant_id=tenant_id, provider="external") - with patch.object(dataset_retrieval, "_get_available_datasets", return_value=[mock_dataset]): - with patch.object(dataset_retrieval, "get_metadata_filter_condition", return_value=(None, None)): - # Create external document - external_doc = create_mock_document_methods( - "External knowledge", - "doc1", - score=0.9, - provider="external", - additional_metadata={ - "dataset_id": dataset_id, - "dataset_name": "external_kb", - "document_id": "ext_doc1", - "title": "External Document", - }, - ) - with patch.object(dataset_retrieval, "multiple_retrieve", return_value=[external_doc]): - # Act - result = dataset_retrieval.knowledge_retrieval(sqlite_session, request) - - # Assert - assert isinstance(result, list) - if result: - assert result[0].metadata.data_source_type == "external" - - def test_knowledge_retrieval_empty_results(self, sqlite_session: Session): - """ - Test knowledge_retrieval when no documents are found. - - Verifies: - - Empty list is returned - - No errors are raised - - All dependencies are still called - """ - # Arrange - tenant_id = str(uuid4()) - user_id = str(uuid4()) - app_id = str(uuid4()) - dataset_id = str(uuid4()) - - request = KnowledgeRetrievalRequest( - tenant_id=tenant_id, - user_id=user_id, - app_id=app_id, - user_from="web", - dataset_ids=[dataset_id], - query="What is Python?", - retrieval_mode="multiple", - top_k=5, - ) - - dataset_retrieval = DatasetRetrieval() - - # Mock dependencies - with patch.object(dataset_retrieval, "_check_knowledge_rate_limit"): - mock_dataset = create_mock_dataset_methods(dataset_id=dataset_id, tenant_id=tenant_id) - with patch.object(dataset_retrieval, "_get_available_datasets", return_value=[mock_dataset]): - with patch.object(dataset_retrieval, "get_metadata_filter_condition", return_value=(None, None)): - # Mock multiple_retrieve to return empty list - with patch.object(dataset_retrieval, "multiple_retrieve", return_value=[]): - # Act - result = dataset_retrieval.knowledge_retrieval(sqlite_session, request) - - # Assert - assert result == [] - - def test_knowledge_retrieval_rate_limit_exceeded(self, sqlite_session: Session): - """ - Test knowledge_retrieval when rate limit is exceeded. - - Verifies: - - RateLimitExceededError is raised - - No further processing occurs - """ - # Arrange - tenant_id = str(uuid4()) - user_id = str(uuid4()) - app_id = str(uuid4()) - dataset_id = str(uuid4()) - - request = KnowledgeRetrievalRequest( - tenant_id=tenant_id, - user_id=user_id, - app_id=app_id, - user_from="web", - dataset_ids=[dataset_id], - query="What is Python?", - retrieval_mode="multiple", - top_k=5, - ) - - dataset_retrieval = DatasetRetrieval() - - # Mock _check_knowledge_rate_limit to raise exception - with patch.object( - dataset_retrieval, - "_check_knowledge_rate_limit", - side_effect=exc.RateLimitExceededError("Rate limit exceeded"), - ): - # Act & Assert - with pytest.raises(exc.RateLimitExceededError): - dataset_retrieval.knowledge_retrieval(sqlite_session, request) - - def test_knowledge_retrieval_no_available_datasets(self, sqlite_session: Session): - """ - Test knowledge_retrieval when no datasets are available. - - Verifies: - - Empty list is returned - - No retrieval is attempted - """ - # Arrange - tenant_id = str(uuid4()) - user_id = str(uuid4()) - app_id = str(uuid4()) - dataset_id = str(uuid4()) - - request = KnowledgeRetrievalRequest( - tenant_id=tenant_id, - user_id=user_id, - app_id=app_id, - user_from="web", - dataset_ids=[dataset_id], - query="What is Python?", - retrieval_mode="multiple", - top_k=5, - ) - - dataset_retrieval = DatasetRetrieval() - - # Mock dependencies - with patch.object(dataset_retrieval, "_check_knowledge_rate_limit"): - # Mock _get_available_datasets to return empty list - with patch.object(dataset_retrieval, "_get_available_datasets", return_value=[]): - # Act - result = dataset_retrieval.knowledge_retrieval(sqlite_session, request) - - # Assert - assert result == [] - - def test_knowledge_retrieval_handles_multiple_documents_with_different_scores(self): - """ - Test that knowledge_retrieval processes multiple documents with different scores. - - Note: Full sorting and position testing requires complex SQLAlchemy mocking - which is better suited for integration tests. This test verifies documents - with different scores can be created and have their metadata. - - Verifies: - - Documents can be created with different scores - - Score metadata is properly set - """ - # Create documents with different scores - doc1 = create_mock_document_methods("Low score", "doc1", score=0.6) - doc2 = create_mock_document_methods("High score", "doc2", score=0.95) - doc3 = create_mock_document_methods("Medium score", "doc3", score=0.8) - - # Assert - each document has the correct score - assert doc1.metadata["score"] == 0.6 - assert doc2.metadata["score"] == 0.95 - assert doc3.metadata["score"] == 0.8 - - # Assert - documents are correctly sorted (not the retrieval result, just the list) - unsorted = [doc1, doc2, doc3] - sorted_docs = sorted(unsorted, key=lambda d: d.metadata["score"], reverse=True) - assert [d.metadata["score"] for d in sorted_docs] == [0.95, 0.8, 0.6] - - -class TestProcessMetadataFilterFunc: - """ - Comprehensive test suite for process_metadata_filter_func method. - - This test class validates all metadata filtering conditions supported by - the DatasetRetrieval class, including string operations, numeric comparisons, - null checks, and list operations. - - Method Signature: - ================== - def process_metadata_filter_func( - self, sequence: int, condition: str, metadata_name: str, value: Any | None, filters: list - ) -> list: - - The method builds SQLAlchemy filter expressions by: - 1. Validating value is not None (except for empty/not empty conditions) - 2. Using DatasetDocument.doc_metadata JSON field operations - 3. Adding appropriate SQLAlchemy expressions to the filters list - 4. Returning the updated filters list - - Tests compile expressions against the real mapped JSON field for string, - numeric, null, and collection operators. - """ - - @pytest.fixture - def retrieval(self): - """ - Create a DatasetRetrieval instance for testing. - - Returns: - DatasetRetrieval: Instance to test process_metadata_filter_func - """ - return DatasetRetrieval() - - # ==================== String Condition Tests ==================== - - def test_contains_condition_string_value(self, retrieval): - """ - Test 'contains' condition with string value. - - Verifies: - - Filters list is populated with LIKE expression - - Pattern matching uses %value% syntax - """ - filters = [] - sequence = 0 - condition = "contains" - metadata_name = "author" - value = "John" - - result = retrieval.process_metadata_filter_func(sequence, condition, metadata_name, value, filters) - - assert result == filters - assert len(filters) == 1 - - def test_not_contains_condition(self, retrieval): - """ - Test 'not contains' condition. - - Verifies: - - Filters list is populated with NOT LIKE expression - - Pattern matching uses %value% syntax with negation - """ - filters = [] - sequence = 0 - condition = "not contains" - metadata_name = "title" - value = "banned" - - result = retrieval.process_metadata_filter_func(sequence, condition, metadata_name, value, filters) - - assert result == filters - assert len(filters) == 1 - - def test_start_with_condition(self, retrieval): - """ - Test 'start with' condition. - - Verifies: - - Filters list is populated with LIKE expression - - Pattern matching uses value% syntax - """ - filters = [] - sequence = 0 - condition = "start with" - metadata_name = "category" - value = "tech" - - result = retrieval.process_metadata_filter_func(sequence, condition, metadata_name, value, filters) - - assert result == filters - assert len(filters) == 1 - - def test_end_with_condition(self, retrieval): - """ - Test 'end with' condition. - - Verifies: - - Filters list is populated with LIKE expression - - Pattern matching uses %value syntax - """ - filters = [] - sequence = 0 - condition = "end with" - metadata_name = "filename" - value = ".pdf" - - result = retrieval.process_metadata_filter_func(sequence, condition, metadata_name, value, filters) - - assert result == filters - assert len(filters) == 1 - - # ==================== Equality Condition Tests ==================== - - def test_is_condition_with_string_value(self, retrieval): - """ - Test 'is' (=) condition with string value. - - Verifies: - - Filters list is populated with equality expression - - String comparison is used - """ - filters = [] - sequence = 0 - condition = "is" - metadata_name = "author" - value = "Jane Doe" - - result = retrieval.process_metadata_filter_func(sequence, condition, metadata_name, value, filters) - - assert result == filters - assert len(filters) == 1 - - def test_equals_condition_with_string_value(self, retrieval): - """ - Test '=' condition with string value. - - Verifies: - - Same behavior as 'is' condition - - String comparison is used - """ - filters = [] - sequence = 0 - condition = "=" - metadata_name = "category" - value = "technology" - - result = retrieval.process_metadata_filter_func(sequence, condition, metadata_name, value, filters) - - assert result == filters - assert len(filters) == 1 - - def test_is_condition_with_int_value(self, retrieval): - """ - Test 'is' condition with integer value. - - Verifies: - - Numeric comparison is used - - as_float() is called on the metadata field - """ - filters = [] - sequence = 0 - condition = "is" - metadata_name = "year" - value = 2023 - - result = retrieval.process_metadata_filter_func(sequence, condition, metadata_name, value, filters) - - assert result == filters - assert len(filters) == 1 - - def test_is_condition_with_float_value(self, retrieval): - """ - Test 'is' condition with float value. - - Verifies: - - Numeric comparison is used - - as_float() is called on the metadata field - """ - filters = [] - sequence = 0 - condition = "is" - metadata_name = "price" - value = 19.99 - - result = retrieval.process_metadata_filter_func(sequence, condition, metadata_name, value, filters) - - assert result == filters - assert len(filters) == 1 - - def test_is_not_condition_with_string_value(self, retrieval): - """ - Test 'is not' (≠) condition with string value. - - Verifies: - - Filters list is populated with inequality expression - - String comparison is used - """ - filters = [] - sequence = 0 - condition = "is not" - metadata_name = "author" - value = "Unknown" - - result = retrieval.process_metadata_filter_func(sequence, condition, metadata_name, value, filters) - - assert result == filters - assert len(filters) == 1 - - def test_not_equals_condition(self, retrieval): - """ - Test '≠' condition with string value. - - Verifies: - - Same behavior as 'is not' condition - - Inequality expression is used - """ - filters = [] - sequence = 0 - condition = "≠" - metadata_name = "category" - value = "archived" - - result = retrieval.process_metadata_filter_func(sequence, condition, metadata_name, value, filters) - - assert result == filters - assert len(filters) == 1 - - def test_is_not_condition_with_numeric_value(self, retrieval): - """ - Test 'is not' condition with numeric value. - - Verifies: - - Numeric inequality comparison is used - - as_float() is called on the metadata field - """ - filters = [] - sequence = 0 - condition = "is not" - metadata_name = "year" - value = 2000 - - result = retrieval.process_metadata_filter_func(sequence, condition, metadata_name, value, filters) - - assert result == filters - assert len(filters) == 1 - - # ==================== Null Condition Tests ==================== - - def test_empty_condition(self, retrieval): - """ - Test 'empty' condition (null check). - - Verifies: - - Filters list is populated with IS NULL expression - - Value can be None for this condition - """ - filters = [] - sequence = 0 - condition = "empty" - metadata_name = "author" - value = None - - result = retrieval.process_metadata_filter_func(sequence, condition, metadata_name, value, filters) - - assert result == filters - assert len(filters) == 1 - - def test_not_empty_condition(self, retrieval): - """ - Test 'not empty' condition (not null check). - - Verifies: - - Filters list is populated with IS NOT NULL expression - - Value can be None for this condition - """ - filters = [] - sequence = 0 - condition = "not empty" - metadata_name = "description" - value = None - - result = retrieval.process_metadata_filter_func(sequence, condition, metadata_name, value, filters) - - assert result == filters - assert len(filters) == 1 - - # ==================== Numeric Comparison Tests ==================== - - def test_before_condition(self, retrieval): - """ - Test 'before' (<) condition. - - Verifies: - - Filters list is populated with less than expression - - Numeric comparison is used - """ - filters = [] - sequence = 0 - condition = "before" - metadata_name = "year" - value = 2020 - - result = retrieval.process_metadata_filter_func(sequence, condition, metadata_name, value, filters) - - assert result == filters - assert len(filters) == 1 - - def test_less_than_condition(self, retrieval): - """ - Test '<' condition. - - Verifies: - - Same behavior as 'before' condition - - Less than expression is used - """ - filters = [] - sequence = 0 - condition = "<" - metadata_name = "price" - value = 100.0 - - result = retrieval.process_metadata_filter_func(sequence, condition, metadata_name, value, filters) - - assert result == filters - assert len(filters) == 1 - - def test_after_condition(self, retrieval): - """ - Test 'after' (>) condition. - - Verifies: - - Filters list is populated with greater than expression - - Numeric comparison is used - """ - filters = [] - sequence = 0 - condition = "after" - metadata_name = "year" - value = 2020 - - result = retrieval.process_metadata_filter_func(sequence, condition, metadata_name, value, filters) - - assert result == filters - assert len(filters) == 1 - - def test_greater_than_condition(self, retrieval): - """ - Test '>' condition. - - Verifies: - - Same behavior as 'after' condition - - Greater than expression is used - """ - filters = [] - sequence = 0 - condition = ">" - metadata_name = "rating" - value = 4.5 - - result = retrieval.process_metadata_filter_func(sequence, condition, metadata_name, value, filters) - - assert result == filters - assert len(filters) == 1 - - def test_less_than_or_equal_condition_unicode(self, retrieval): - """ - Test '≤' condition. - - Verifies: - - Filters list is populated with less than or equal expression - - Numeric comparison is used - """ - filters = [] - sequence = 0 - condition = "≤" - metadata_name = "price" - value = 50.0 - - result = retrieval.process_metadata_filter_func(sequence, condition, metadata_name, value, filters) - - assert result == filters - assert len(filters) == 1 - - def test_less_than_or_equal_condition_ascii(self, retrieval): - """ - Test '<=' condition. - - Verifies: - - Same behavior as '≤' condition - - Less than or equal expression is used - """ - filters = [] - sequence = 0 - condition = "<=" - metadata_name = "year" - value = 2023 - - result = retrieval.process_metadata_filter_func(sequence, condition, metadata_name, value, filters) - - assert result == filters - assert len(filters) == 1 - - def test_greater_than_or_equal_condition_unicode(self, retrieval): - """ - Test '≥' condition. - - Verifies: - - Filters list is populated with greater than or equal expression - - Numeric comparison is used - """ - filters = [] - sequence = 0 - condition = "≥" - metadata_name = "rating" - value = 3.5 - - result = retrieval.process_metadata_filter_func(sequence, condition, metadata_name, value, filters) - - assert result == filters - assert len(filters) == 1 - - def test_greater_than_or_equal_condition_ascii(self, retrieval): - """ - Test '>=' condition. - - Verifies: - - Same behavior as '≥' condition - - Greater than or equal expression is used - """ - filters = [] - sequence = 0 - condition = ">=" - metadata_name = "year" - value = 2000 - - result = retrieval.process_metadata_filter_func(sequence, condition, metadata_name, value, filters) - - assert result == filters - assert len(filters) == 1 - - # ==================== List/In Condition Tests ==================== - - def test_in_condition_with_comma_separated_string(self, retrieval): - """ - Test 'in' condition with comma-separated string value. - - Verifies: - - String is split into list - - Whitespace is trimmed from each value - - IN expression is created - """ - filters = [] - sequence = 0 - condition = "in" - metadata_name = "category" - value = "tech, science, AI " - - result = retrieval.process_metadata_filter_func(sequence, condition, metadata_name, value, filters) - - assert result == filters - assert len(filters) == 1 - - def test_in_condition_with_list_value(self, retrieval): - """ - Test 'in' condition with list value. - - Verifies: - - List is processed correctly - - None values are filtered out - - IN expression is created with valid values - """ - filters = [] - sequence = 0 - condition = "in" - metadata_name = "tags" - value = ["python", "javascript", None, "golang"] - - result = retrieval.process_metadata_filter_func(sequence, condition, metadata_name, value, filters) - - assert result == filters - assert len(filters) == 1 - - def test_in_condition_with_tuple_value(self, retrieval): - """ - Test 'in' condition with tuple value. - - Verifies: - - Tuple is processed like a list - - IN expression is created - """ - filters = [] - sequence = 0 - condition = "in" - metadata_name = "category" - value = ("tech", "science", "ai") - - result = retrieval.process_metadata_filter_func(sequence, condition, metadata_name, value, filters) - - assert result == filters - assert len(filters) == 1 - - def test_in_condition_with_empty_string(self, retrieval): - """ - Test 'in' condition with empty string value. - - Verifies: - - Empty string results in literal(False) filter - - No valid values to match - """ - filters = [] - sequence = 0 - condition = "in" - metadata_name = "category" - value = "" - - result = retrieval.process_metadata_filter_func(sequence, condition, metadata_name, value, filters) - - assert result == filters - assert len(filters) == 1 - # Verify it's a literal(False) expression - # This is a bit tricky to test without access to the actual expression - - def test_in_condition_with_only_whitespace(self, retrieval): - """ - Test 'in' condition with whitespace-only string value. - - Verifies: - - Whitespace-only string results in literal(False) filter - - All values are stripped and filtered out - """ - filters = [] - sequence = 0 - condition = "in" - metadata_name = "category" - value = " , , " - - result = retrieval.process_metadata_filter_func(sequence, condition, metadata_name, value, filters) - - assert result == filters - assert len(filters) == 1 - - def test_in_condition_with_single_string(self, retrieval): - """ - Test 'in' condition with single non-comma string. - - Verifies: - - Single string is treated as single-item list - - IN expression is created with one value - """ - filters = [] - sequence = 0 - condition = "in" - metadata_name = "category" - value = "technology" - - result = retrieval.process_metadata_filter_func(sequence, condition, metadata_name, value, filters) - - assert result == filters - assert len(filters) == 1 - - # ==================== Edge Case Tests ==================== - - def test_none_value_with_non_empty_condition(self, retrieval): - """ - Test None value with conditions that require value. - - Verifies: - - Original filters list is returned unchanged - - No filter is added for None values (except empty/not empty) - """ - filters = [] - sequence = 0 - condition = "contains" - metadata_name = "author" - value = None - - result = retrieval.process_metadata_filter_func(sequence, condition, metadata_name, value, filters) - - assert result == filters - assert len(filters) == 0 # No filter added - - def test_none_value_with_equals_condition(self, retrieval): - """ - Test None value with 'is' (=) condition. - - Verifies: - - Original filters list is returned unchanged - - No filter is added for None values - """ - filters = [] - sequence = 0 - condition = "is" - metadata_name = "author" - value = None - - result = retrieval.process_metadata_filter_func(sequence, condition, metadata_name, value, filters) - - assert result == filters - assert len(filters) == 0 - - def test_none_value_with_numeric_condition(self, retrieval): - """ - Test None value with numeric comparison condition. - - Verifies: - - Original filters list is returned unchanged - - No filter is added for None values - """ - filters = [] - sequence = 0 - condition = ">" - metadata_name = "year" - value = None - - result = retrieval.process_metadata_filter_func(sequence, condition, metadata_name, value, filters) - - assert result == filters - assert len(filters) == 0 - - def test_existing_filters_preserved(self, retrieval): - """ - Test that existing filters are preserved. - - Verifies: - - Existing filters in the list are not removed - - New filters are appended to the list - """ - existing_filter = MagicMock() - filters = [existing_filter] - sequence = 0 - condition = "contains" - metadata_name = "author" - value = "test" - - result = retrieval.process_metadata_filter_func(sequence, condition, metadata_name, value, filters) - - assert result == filters - assert len(filters) == 2 - assert filters[0] == existing_filter - - def test_multiple_filters_accumulated(self, retrieval): - """ - Test multiple calls to accumulate filters. - - Verifies: - - Each call adds a new filter to the list - - All filters are preserved across calls - """ - filters = [] - - # First filter - retrieval.process_metadata_filter_func(0, "contains", "author", "John", filters) - assert len(filters) == 1 - - # Second filter - retrieval.process_metadata_filter_func(1, ">", "year", 2020, filters) - assert len(filters) == 2 - - # Third filter - retrieval.process_metadata_filter_func(2, "is", "category", "tech", filters) - assert len(filters) == 3 - - def test_unknown_condition(self, retrieval): - """ - Test unknown/unsupported condition. - - Verifies: - - Original filters list is returned unchanged - - No filter is added for unknown conditions - """ - filters = [] - sequence = 0 - condition = "unknown_condition" - metadata_name = "author" - value = "test" - - result = retrieval.process_metadata_filter_func(sequence, condition, metadata_name, value, filters) - - assert result == filters - assert len(filters) == 0 - - def test_empty_string_value_with_contains(self, retrieval): - """ - Test empty string value with 'contains' condition. - - Verifies: - - Filter is added even with empty string - - LIKE expression is created - """ - filters = [] - sequence = 0 - condition = "contains" - metadata_name = "author" - value = "" - - result = retrieval.process_metadata_filter_func(sequence, condition, metadata_name, value, filters) - - assert result == filters - assert len(filters) == 1 - - def test_special_characters_in_value(self, retrieval): - """ - Test special characters in value string. - - Verifies: - - Special characters are handled in value - - LIKE expression is created correctly - """ - filters = [] - sequence = 0 - condition = "contains" - metadata_name = "title" - value = "C++ & Python's features" - - result = retrieval.process_metadata_filter_func(sequence, condition, metadata_name, value, filters) - - assert result == filters - assert len(filters) == 1 - - def test_zero_value_with_numeric_condition(self, retrieval): - """ - Test zero value with numeric comparison condition. - - Verifies: - - Zero is treated as valid value - - Numeric comparison is performed - """ - filters = [] - sequence = 0 - condition = ">" - metadata_name = "price" - value = 0 - - result = retrieval.process_metadata_filter_func(sequence, condition, metadata_name, value, filters) - - assert result == filters - assert len(filters) == 1 - - def test_negative_value_with_numeric_condition(self, retrieval): - """ - Test negative value with numeric comparison condition. - - Verifies: - - Negative numbers are handled correctly - - Numeric comparison is performed - """ - filters = [] - sequence = 0 - condition = "<" - metadata_name = "temperature" - value = -10.5 - - result = retrieval.process_metadata_filter_func(sequence, condition, metadata_name, value, filters) - - assert result == filters - assert len(filters) == 1 - - def test_float_value_with_integer_comparison(self, retrieval): - """ - Test float value with numeric comparison condition. - - Verifies: - - Float values work correctly - - Numeric comparison is performed - """ - filters = [] - sequence = 0 - condition = ">=" - metadata_name = "rating" - value = 4.5 - - result = retrieval.process_metadata_filter_func(sequence, condition, metadata_name, value, filters) - - assert result == filters - assert len(filters) == 1 - - -class TestKnowledgeRetrievalRegression: - @pytest.fixture - def mock_dataset(self) -> Dataset: - dataset = Dataset( - id=str(uuid4()), - tenant_id=str(uuid4()), - name="test_dataset", - indexing_technique="high_quality", - provider="dify", - ) - return dataset - - def test_multiple_retrieve_reranking_with_app_context(self, mock_dataset): - """ - Repro test for current bug: - reranking runs after `with flask_app.app_context():` exits. - `_multiple_retrieve_thread` catches exceptions and stores them into `thread_exceptions`, - so we must assert from that list (not from an outer try/except). - """ - dataset_retrieval = DatasetRetrieval() - flask_app = Flask(__name__) - tenant_id = str(uuid4()) - - # second dataset to ensure dataset_count > 1 reranking branch - secondary_dataset = Dataset( - id=str(uuid4()), - provider="dify", - indexing_technique="high_quality", - ) - - # retriever returns 1 doc into internal list (all_documents_item) - document = Document( - page_content="Context aware doc", - metadata={ - "doc_id": "doc1", - "score": 0.95, - "document_id": str(uuid4()), - "dataset_id": mock_dataset.id, - }, - provider="dify", - ) - - def fake_retriever( - flask_app, - session, - dataset_id, - query, - top_k, - all_documents, - document_ids_filter, - metadata_condition, - attachment_ids, - ): - all_documents.append(document) - - called = {"init": 0, "invoke": 0} - - class ContextRequiredPostProcessor: - def __init__(self, *args, **kwargs): - called["init"] += 1 - # will raise RuntimeError if no Flask app context exists - _ = current_app.name - - def invoke(self, *args, **kwargs): - called["invoke"] += 1 - _ = current_app.name - return kwargs.get("documents") or args[1] - - # output list from _multiple_retrieve_thread - all_documents: list[Document] = [] - - # The safe wrapper collects exceptions raised after tracing. - thread_exceptions: list[Exception] = [] - - def target(): - with ( - patch.object(dataset_retrieval, "_retriever", side_effect=fake_retriever), - patch( - "core.rag.retrieval.dataset_retrieval.DataPostProcessor", - ContextRequiredPostProcessor, - ), - ): - dataset_retrieval._multiple_retrieve_thread_safely( - flask_app=flask_app, - available_datasets=[mock_dataset, secondary_dataset], - metadata_condition=None, - metadata_filter_document_ids=None, - all_documents=all_documents, - tenant_id=tenant_id, - reranking_enable=True, - reranking_mode="reranking_model", - reranking_model={ - "reranking_provider_name": "cohere", - "reranking_model_name": "rerank-v2", - }, - weights=None, - top_k=3, - score_threshold=0.0, - query="test query", - attachment_id=None, - dataset_count=2, # force reranking branch - thread_exceptions=thread_exceptions, - ) - - t = threading.Thread(target=target) - t.start() - t.join() - - # Ensure reranking branch was actually executed - assert called["init"] >= 1, "DataPostProcessor was never constructed; reranking branch may not have run." - - # Current buggy code should record an exception (not raise it) - assert not thread_exceptions, thread_exceptions - - def test_run_retriever_thread_provides_session_to_retriever(self): - dataset_retrieval = DatasetRetrieval() - all_documents: list[Document] = [] - - with patch.object(dataset_retrieval, "_retriever") as mock_retriever: - dataset_retrieval._run_retriever_thread( - flask_app=_FakeFlaskApp(), - dataset_id="dataset-1", - query="test query", - top_k=3, - all_documents=all_documents, - document_ids_filter=None, - metadata_condition=None, - attachment_ids=None, - ) - - mock_retriever.assert_called_once() - assert isinstance(mock_retriever.call_args.kwargs["session"], Session) - - def test_run_retriever_thread_safely_records_retriever_exception(self): - dataset_retrieval = DatasetRetrieval() - all_documents: list[Document] = [] - cancel_event = threading.Event() - thread_exceptions: list[Exception] = [] - expected_error = RuntimeError("retrieval failed") - - with patch.object(dataset_retrieval, "_retriever", side_effect=expected_error): - dataset_retrieval._run_retriever_thread_safely( - flask_app=_FakeFlaskApp(), - dataset_id="dataset-1", - query="test query", - top_k=3, - all_documents=all_documents, - document_ids_filter=None, - metadata_condition=None, - attachment_ids=None, - cancel_event=cancel_event, - thread_exceptions=thread_exceptions, - ) - - assert cancel_event.is_set() - assert thread_exceptions == [expected_error] - - def test_run_retriever_thread_safely_skips_failed_dataset_when_requested(self, caplog): - dataset_retrieval = DatasetRetrieval() - all_documents: list[Document] = [] - cancel_event = threading.Event() - thread_exceptions: list[Exception] = [] - expected_error = RuntimeError("retrieval failed") - - with patch.object(dataset_retrieval, "_retriever", side_effect=expected_error): - dataset_retrieval._run_retriever_thread_safely( - flask_app=_FakeFlaskApp(), - dataset_id="dataset-1", - query="test query", - top_k=3, - all_documents=all_documents, - document_ids_filter=None, - metadata_condition=None, - attachment_ids=None, - cancel_event=cancel_event, - thread_exceptions=thread_exceptions, - skip_on_error=True, - ) - - assert not cancel_event.is_set() - assert thread_exceptions == [] - assert "dataset_id=dataset-1" in caplog.text - assert "Skipping dataset retrieval because retriever failed" in caplog.text - - def test_multiple_retrieve_thread_skips_failed_dataset(self, mock_dataset, caplog): - dataset_retrieval = DatasetRetrieval() - flask_app = Flask(__name__) - successful_dataset = Dataset( - id=str(uuid4()), - provider="dify", - indexing_technique="high_quality", - ) - document = Document( - page_content="successful doc", - metadata={ - "doc_id": "doc1", - "score": 0.95, - "document_id": str(uuid4()), - "dataset_id": successful_dataset.id, - }, - provider="dify", - ) - - def fake_retriever( - flask_app, - session, - dataset_id, - query, - top_k, - all_documents, - document_ids_filter, - metadata_condition, - attachment_ids, - ): - if dataset_id == mock_dataset.id: - raise RuntimeError("dataset unavailable") - all_documents.append(document) - - all_documents: list[Document] = [] - - with patch.object(dataset_retrieval, "_retriever", side_effect=fake_retriever): - dataset_retrieval._multiple_retrieve_thread( - flask_app=flask_app, - available_datasets=[mock_dataset, successful_dataset], - metadata_condition=None, - metadata_filter_document_ids=None, - all_documents=all_documents, - tenant_id=str(uuid4()), - reranking_enable=False, - reranking_mode="reranking_model", - reranking_model=None, - weights=None, - top_k=3, - score_threshold=0.0, - query="test query", - attachment_id=None, - dataset_count=2, - ) - - assert all_documents == [document] - assert f"dataset_id={mock_dataset.id}" in caplog.text - assert "Skipping dataset retrieval because retriever failed" in caplog.text - - class _FakeFlaskApp: def app_context(self): return nullcontext() @@ -3901,538 +283,6 @@ class _ImmediateThread: return self._alive -class TestDatasetRetrievalAdditionalHelpers: - @pytest.fixture - def retrieval(self) -> DatasetRetrieval: - return DatasetRetrieval() - - def test_llm_usage_and_record_usage(self, retrieval: DatasetRetrieval) -> None: - empty_usage = retrieval.llm_usage - assert empty_usage.total_tokens == 0 - - retrieval._record_usage(None) - assert retrieval.llm_usage.total_tokens == 0 - - usage_1 = LLMUsage.from_metadata({"prompt_tokens": 2, "completion_tokens": 3, "total_tokens": 5}) - usage_2 = LLMUsage.from_metadata({"prompt_tokens": 4, "completion_tokens": 1, "total_tokens": 5}) - retrieval._record_usage(usage_1) - retrieval._record_usage(usage_2) - assert retrieval.llm_usage.total_tokens == 10 - - def test_replace_metadata_filter_value(self, retrieval: DatasetRetrieval) -> None: - assert retrieval._replace_metadata_filter_value("plain", {}) == "plain" - replaced = retrieval._replace_metadata_filter_value( - "hello {{name}}\n\t{{missing}}", - {"name": "world"}, - ) - assert replaced == "hello world {{missing}}" - - def test_process_metadata_filter_in_with_scalar_fallback(self) -> None: - filters: list = [] - result = DatasetRetrieval.process_metadata_filter_func( - sequence=0, - condition="in", - metadata_name="category", - value=123, - filters=filters, - ) - assert result is filters - assert len(filters) == 1 - - def test_calculate_vector_score(self, retrieval: DatasetRetrieval) -> None: - doc_high = Document(page_content="a", metadata={"score": 0.9}, provider="dify") - doc_low = Document(page_content="b", metadata={"score": 0.2}, provider="dify") - doc_no_meta = Document(page_content="c", metadata={}, provider="dify") - - filtered = retrieval.calculate_vector_score([doc_low, doc_high, doc_no_meta], top_k=1, score_threshold=0.5) - assert len(filtered) == 1 - assert filtered[0].metadata["score"] == 0.9 - - assert retrieval.calculate_vector_score([doc_low], top_k=2, score_threshold=1.0) == [] - - def test_calculate_keyword_score(self, retrieval: DatasetRetrieval) -> None: - documents = [ - Document(page_content="python language", metadata={"doc_id": "1"}, provider="dify"), - Document(page_content="java language", metadata={"doc_id": "2"}, provider="dify"), - ] - keyword_handler = Mock() - keyword_handler.extract_keywords.side_effect = [ - ["python", "language"], - ["python", "language"], - ["java", "language"], - ] - - with patch("core.rag.retrieval.dataset_retrieval.JiebaKeywordTableHandler", return_value=keyword_handler): - ranked = retrieval.calculate_keyword_score("python language", documents, top_k=1) - - assert len(ranked) == 1 - assert "keywords" in ranked[0].metadata - assert ranked[0].metadata["doc_id"] == "1" - - def test_send_trace_task(self, retrieval: DatasetRetrieval) -> None: - trace_manager = Mock() - retrieval.application_generate_entity = SimpleNamespace(trace_manager=trace_manager) - docs = [Document(page_content="d", metadata={}, provider="dify")] - - retrieval._send_trace_task("m1", docs, {"cost": 1}) - trace_manager.add_trace_task.assert_called_once() - - retrieval.application_generate_entity = None - trace_manager.reset_mock() - retrieval._send_trace_task("m1", docs, {"cost": 1}) - trace_manager.add_trace_task.assert_not_called() - - def test_on_query(self, retrieval: DatasetRetrieval, sqlite_engine: Engine, sqlite_session: Session) -> None: - dataset_ids = [str(uuid4()), str(uuid4())] - app_id = str(uuid4()) - user_id = str(uuid4()) - - with patch("core.rag.retrieval.dataset_retrieval.db", SimpleNamespace(engine=sqlite_engine)): - retrieval._on_query( - query=None, - attachment_ids=None, - dataset_ids=[dataset_ids[0]], - app_id=app_id, - user_from="account", - user_id=user_id, - ) - assert sqlite_session.scalars(select(DatasetQuery)).all() == [] - - retrieval._on_query( - query="python", - attachment_ids=[str(uuid4())], - dataset_ids=dataset_ids, - app_id=app_id, - user_from="account", - user_id=user_id, - ) - added_queries = sqlite_session.scalars(select(DatasetQuery).order_by(DatasetQuery.dataset_id)).all() - assert [row.dataset_id for row in added_queries] == sorted(dataset_ids) - assert all(row.created_by == user_id for row in added_queries) - - def test_on_query_normalizes_workflow_end_user_role( - self, retrieval: DatasetRetrieval, sqlite_engine: Engine, sqlite_session: Session - ) -> None: - dataset_id = str(uuid4()) - with patch("core.rag.retrieval.dataset_retrieval.db", SimpleNamespace(engine=sqlite_engine)): - retrieval._on_query( - query="python", - attachment_ids=None, - dataset_ids=[dataset_id], - app_id=str(uuid4()), - user_from="end-user", - user_id=str(uuid4()), - ) - - added_query = sqlite_session.scalar(select(DatasetQuery).where(DatasetQuery.dataset_id == dataset_id)) - assert added_query is not None - assert added_query.created_by_role == CreatorUserRole.END_USER - - def test_handle_invoke_result(self, retrieval: DatasetRetrieval) -> None: - usage = LLMUsage.empty_usage() - chunk_1 = SimpleNamespace( - model="m1", - prompt_messages=[Mock()], - delta=SimpleNamespace(message=SimpleNamespace(content="hello "), usage=usage), - ) - chunk_2 = SimpleNamespace( - model="m1", - prompt_messages=[Mock()], - delta=SimpleNamespace( - message=SimpleNamespace(content=[SimpleNamespace(data="world")]), - usage=None, - ), - ) - - def _chunks() -> Generator[Any]: - yield chunk_1 - yield chunk_2 - - text, returned_usage = retrieval._handle_invoke_result(_chunks()) - assert text == "hello world" - assert returned_usage == usage - - def _empty_chunks() -> Generator[Any]: - yield from () - - text_empty, usage_empty = retrieval._handle_invoke_result(_empty_chunks()) - assert text_empty == "" - assert usage_empty == LLMUsage.empty_usage() - - def test_get_prompt_template(self, retrieval: DatasetRetrieval) -> None: - model_config_chat = ModelConfigWithCredentialsEntity.model_construct( - provider="openai", - model="gpt", - model_schema=Mock(), - mode="chat", - provider_model_bundle=Mock(), - credentials={}, - parameters={}, - stop=["x"], - ) - model_config_completion = ModelConfigWithCredentialsEntity.model_construct( - provider="openai", - model="gpt", - model_schema=Mock(), - mode="completion", - provider_model_bundle=Mock(), - credentials={}, - parameters={}, - stop=[], - ) - - with patch("core.rag.retrieval.dataset_retrieval.AdvancedPromptTransform") as mock_prompt_transform: - mock_prompt_transform.return_value.get_prompt.return_value = ["prompt"] - prompt_messages, stop = retrieval._get_prompt_template( - model_config=model_config_chat, - mode="chat", - metadata_fields=["author"], - query="python", - ) - assert prompt_messages == ["prompt"] - assert stop == ["x"] - - with patch( - "core.rag.retrieval.dataset_retrieval.METADATA_FILTER_COMPLETION_PROMPT", - "{input_text} {metadata_fields}", - ): - prompt_messages_completion, stop_completion = retrieval._get_prompt_template( - model_config=model_config_completion, - mode="completion", - metadata_fields=["author"], - query="python", - ) - assert prompt_messages_completion == ["prompt"] - assert stop_completion == [] - - with pytest.raises(ValueError, match="'unknown\\-mode' is not a valid ModelMode"): - retrieval._get_prompt_template( - model_config=model_config_chat, - mode="unknown-mode", - metadata_fields=[], - query="python", - ) - - def test_fetch_model_config_validation_and_success(self, retrieval: DatasetRetrieval) -> None: - with pytest.raises(ValueError, match="single_retrieval_config is required"): - retrieval._fetch_model_config("tenant-1", None) # type: ignore[arg-type] - - model_cfg = AppModelConfig(provider="openai", name="gpt", mode="chat", completion_params={"stop": ["END"]}) - model_instance = Mock() - model_instance.credentials = {"k": "v"} - model_instance.provider_model_bundle = Mock() - model_instance.model_type_instance = Mock() - model_instance.model_type_instance.get_model_schema.return_value = Mock() - - with ( - patch("core.rag.retrieval.dataset_retrieval.ModelManager.for_tenant") as mock_manager, - patch("core.rag.retrieval.dataset_retrieval.ModelConfigWithCredentialsEntity") as mock_cfg_entity, - ): - mock_manager.return_value.get_model_instance.return_value = model_instance - mock_cfg_entity.return_value = SimpleNamespace( - provider="openai", - model="gpt", - stop=["END"], - parameters={"temperature": 0.1}, - ) - - model_instance.provider_model_bundle.configuration.get_provider_model.return_value = None - with pytest.raises(ValueError, match="not exist"): - retrieval._fetch_model_config("tenant-1", model_cfg) - - provider_model = SimpleNamespace(status=ModelStatus.NO_CONFIGURE) - model_instance.provider_model_bundle.configuration.get_provider_model.return_value = provider_model - with pytest.raises(ValueError, match="credentials is not initialized"): - retrieval._fetch_model_config("tenant-1", model_cfg) - - provider_model.status = ModelStatus.NO_PERMISSION - with pytest.raises(ValueError, match="currently not support"): - retrieval._fetch_model_config("tenant-1", model_cfg) - - provider_model.status = ModelStatus.QUOTA_EXCEEDED - with pytest.raises(ValueError, match="quota exceeded"): - retrieval._fetch_model_config("tenant-1", model_cfg) - - provider_model.status = ModelStatus.ACTIVE - bad_mode_cfg = AppModelConfig(provider="openai", name="gpt", mode="chat") - bad_mode_cfg.mode = None # type: ignore[assignment] - with pytest.raises(ValueError, match="LLM mode is required"): - retrieval._fetch_model_config("tenant-1", bad_mode_cfg) - - model_instance.model_type_instance.get_model_schema.return_value = None - with pytest.raises(ValueError, match="not exist"): - retrieval._fetch_model_config("tenant-1", model_cfg) - - model_instance.model_type_instance.get_model_schema.return_value = Mock() - model_cfg_success = AppModelConfig( - provider="openai", - name="gpt", - mode="chat", - completion_params={"temperature": 0.1, "stop": ["END"]}, - ) - _, config = retrieval._fetch_model_config("tenant-1", model_cfg_success) - assert config.provider == "openai" - assert config.model == "gpt" - assert config.stop == ["END"] - assert "stop" not in config.parameters - - def test_automatic_metadata_filter_func(self, retrieval: DatasetRetrieval, sqlite_session: Session) -> None: - tenant_id = str(uuid4()) - dataset_id = str(uuid4()) - user_id = str(uuid4()) - sqlite_session.add( - DatasetMetadata( - tenant_id=tenant_id, - dataset_id=dataset_id, - type="string", - name="author", - created_by=user_id, - ) - ) - sqlite_session.commit() - model_instance = Mock() - model_instance.invoke_llm.return_value = iter([Mock()]) - model_config = ModelConfigWithCredentialsEntity.model_construct( - provider="openai", - model="gpt", - model_schema=Mock(), - mode="chat", - provider_model_bundle=Mock(), - credentials={}, - parameters={}, - stop=[], - ) - usage = LLMUsage.from_metadata({"prompt_tokens": 1, "completion_tokens": 1, "total_tokens": 2}) - with ( - patch.object(retrieval, "_fetch_model_config", return_value=(model_instance, model_config)), - patch.object(retrieval, "_get_prompt_template", return_value=(["prompt"], [])), - patch.object(retrieval, "_handle_invoke_result", return_value=('{"metadata_map":[]}', usage)), - patch("core.rag.retrieval.dataset_retrieval.parse_and_check_json_markdown") as mock_parse, - patch.object(retrieval, "_record_usage") as mock_record_usage, - ): - mock_parse.return_value = { - "metadata_map": [ - { - "metadata_field_name": "author", - "metadata_field_value": "Alice", - "comparison_operator": "contains", - }, - { - "metadata_field_name": "ignored", - "metadata_field_value": "value", - "comparison_operator": "contains", - }, - ] - } - result = retrieval._automatic_metadata_filter_func( - sqlite_session, - dataset_ids=[dataset_id], - query="python", - tenant_id=tenant_id, - user_id=user_id, - metadata_model_config=AppModelConfig(provider="openai", name="gpt", mode="chat"), - ) - - assert result == [{"metadata_name": "author", "value": "Alice", "condition": "contains"}] - mock_record_usage.assert_called_once_with(usage) - - with ( - patch.object(retrieval, "_fetch_model_config", side_effect=RuntimeError("boom")), - ): - with pytest.raises(RuntimeError, match="boom"): - retrieval._automatic_metadata_filter_func( - sqlite_session, - dataset_ids=[dataset_id], - query="python", - tenant_id=tenant_id, - user_id=user_id, - metadata_model_config=AppModelConfig(provider="openai", name="gpt", mode="chat"), - ) - - def test_get_metadata_filter_condition(self, retrieval: DatasetRetrieval, sqlite_session: Session) -> None: - tenant_id = str(uuid4()) - dataset_id = str(uuid4()) - document_id = str(uuid4()) - user_id = str(uuid4()) - sqlite_session.add( - DatasetDocument( - id=document_id, - tenant_id=tenant_id, - dataset_id=dataset_id, - position=1, - data_source_type=DataSourceType("upload_file"), - batch="batch-1", - name="document", - created_from=DocumentCreatedFrom("api"), - created_by=user_id, - indexing_status="completed", - enabled=True, - archived=False, - doc_metadata={"author": "Alice"}, - ) - ) - sqlite_session.commit() - - mapping, condition = retrieval.get_metadata_filter_condition( - sqlite_session, - dataset_ids=[dataset_id], - query="python", - tenant_id="tenant-1", - user_id="u1", - metadata_filtering_mode="disabled", - metadata_model_config=AppModelConfig(provider="openai", name="gpt", mode="chat"), - metadata_filtering_conditions=None, - inputs={}, - ) - assert mapping is None - assert condition is None - - automatic_filters = [{"condition": "contains", "metadata_name": "author", "value": "Alice"}] - with ( - patch.object(retrieval, "_automatic_metadata_filter_func", return_value=automatic_filters), - ): - mapping, condition = retrieval.get_metadata_filter_condition( - sqlite_session, - dataset_ids=[dataset_id], - query="python", - tenant_id="tenant-1", - user_id="u1", - metadata_filtering_mode="automatic", - metadata_model_config=AppModelConfig(provider="openai", name="gpt", mode="chat"), - metadata_filtering_conditions=AppMetadataFilteringCondition(logical_operator="or", conditions=[]), - inputs={}, - ) - assert mapping == {dataset_id: [document_id]} - assert condition is not None - assert condition.logical_operator == "or" - - manual_conditions = AppMetadataFilteringCondition( - logical_operator="and", - conditions=[AppCondition(name="author", comparison_operator="contains", value="{{name}}")], - ) - mapping, condition = retrieval.get_metadata_filter_condition( - sqlite_session, - dataset_ids=[dataset_id], - query="python", - tenant_id="tenant-1", - user_id="u1", - metadata_filtering_mode="manual", - metadata_model_config=AppModelConfig(provider="openai", name="gpt", mode="chat"), - metadata_filtering_conditions=manual_conditions, - inputs={"name": "Alice"}, - ) - assert mapping == {dataset_id: [document_id]} - assert condition is not None - assert condition.conditions - first_condition = condition.conditions[0] - assert first_condition.value == "Alice" - - with pytest.raises(ValueError, match="Invalid metadata filtering mode"): - retrieval.get_metadata_filter_condition( - sqlite_session, - dataset_ids=[dataset_id], - query="python", - tenant_id="tenant-1", - user_id="u1", - metadata_filtering_mode="unsupported", - metadata_model_config=AppModelConfig(provider="openai", name="gpt", mode="chat"), - metadata_filtering_conditions=None, - inputs={}, - ) - - def test_get_available_datasets(self, retrieval: DatasetRetrieval, sqlite_session: Session) -> None: - tenant_id = str(uuid4()) - other_tenant_id = str(uuid4()) - creator_id = str(uuid4()) - internal_id = str(uuid4()) - external_id = str(uuid4()) - unavailable_id = str(uuid4()) - decoy_id = str(uuid4()) - datasets = [ - Dataset( - id=internal_id, - tenant_id=tenant_id, - name="internal", - provider="dify", - indexing_technique="high_quality", - created_by=creator_id, - ), - Dataset( - id=external_id, - tenant_id=tenant_id, - name="external", - provider="external", - indexing_technique="high_quality", - created_by=creator_id, - ), - Dataset( - id=unavailable_id, - tenant_id=tenant_id, - name="unavailable", - provider="dify", - indexing_technique="high_quality", - created_by=creator_id, - ), - Dataset( - id=decoy_id, - tenant_id=other_tenant_id, - name="decoy", - provider="external", - indexing_technique="high_quality", - created_by=creator_id, - ), - ] - sqlite_session.add_all(datasets) - sqlite_session.add( - DatasetDocument( - tenant_id=tenant_id, - dataset_id=internal_id, - position=1, - data_source_type=DataSourceType("upload_file"), - batch="batch-1", - name="available document", - created_from=DocumentCreatedFrom("api"), - created_by=creator_id, - indexing_status="completed", - enabled=True, - archived=False, - ) - ) - sqlite_session.commit() - - available = retrieval._get_available_datasets(tenant_id, [internal_id, external_id, unavailable_id, decoy_id]) - - assert {dataset.id for dataset in available} == {internal_id, external_id} - - def test_check_knowledge_rate_limit(self, retrieval: DatasetRetrieval, sqlite_session: Session) -> None: - with ( - patch("core.rag.retrieval.dataset_retrieval.FeatureService.get_knowledge_rate_limit") as mock_limit, - patch("core.rag.retrieval.dataset_retrieval.redis_client") as mock_redis, - patch("core.rag.retrieval.dataset_retrieval.time.time", return_value=100.0), - ): - mock_limit.return_value = SimpleNamespace(enabled=True, limit=2, subscription_plan="pro") - mock_redis.zcard.return_value = 1 - retrieval._check_knowledge_rate_limit("tenant-1") - mock_redis.zadd.assert_called_once() - - tenant_id = str(uuid4()) - with ( - patch("core.rag.retrieval.dataset_retrieval.FeatureService.get_knowledge_rate_limit") as mock_limit, - patch("core.rag.retrieval.dataset_retrieval.redis_client") as mock_redis, - patch("core.rag.retrieval.dataset_retrieval.time.time", return_value=100.0), - ): - mock_limit.return_value = SimpleNamespace(enabled=True, limit=1, subscription_plan="pro") - mock_redis.zcard.return_value = 2 - with pytest.raises(exc.RateLimitExceededError): - retrieval._check_knowledge_rate_limit(tenant_id) - rate_limit_log = sqlite_session.scalar(select(RateLimitLog).where(RateLimitLog.tenant_id == tenant_id)) - assert rate_limit_log is not None - - with patch("core.rag.retrieval.dataset_retrieval.FeatureService.get_knowledge_rate_limit") as mock_limit: - mock_limit.return_value = SimpleNamespace(enabled=False) - retrieval._check_knowledge_rate_limit("tenant-1") - - def _doc( provider: str = "dify", content: str = "content", @@ -4441,7 +291,7 @@ def _doc( document_id: str = "document-1", doc_id: str = "node-1", extra: dict[str, Any] | None = None, -) -> Document: +) -> _dataset_retrieval_Document: metadata = { "score": score, "dataset_id": dataset_id, @@ -4450,7 +300,7 @@ def _doc( } if extra: metadata.update(extra) - return Document(page_content=content, metadata=metadata, provider=provider) + return _dataset_retrieval_Document(page_content=content, metadata=metadata, provider=provider) class _ImmediateThread: @@ -4497,170 +347,4313 @@ def _timer(): yield {"cost": 1} -class TestKnowledgeRetrievalCoverage: - @pytest.fixture(autouse=True) - def _orm_session(self, sqlite_session: Session) -> None: - self.orm_session = sqlite_session +class TestDatasetRetrievalCases: + class TestRetrievalService: + """ + Comprehensive test suite for RetrievalService class. - @pytest.fixture - def retrieval(self) -> DatasetRetrieval: - return DatasetRetrieval() + This test class validates all retrieval methods and their interactions, + including edge cases, error handling, and integration scenarios. - def test_returns_empty_when_query_missing(self, retrieval: DatasetRetrieval) -> None: - request = KnowledgeRetrievalRequest( - tenant_id="tenant-1", - user_id="user-1", - app_id="app-1", - user_from="workflow", - dataset_ids=["d1"], - query=None, - retrieval_mode="multiple", - ) - with ( - patch.object(retrieval, "_check_knowledge_rate_limit"), - patch.object(retrieval, "_get_available_datasets", return_value=[_dataset(id="d1")]), - ): - assert retrieval.knowledge_retrieval(self.orm_session, request) == [] + Test Organization: + ================== + 1. Fixtures (lines ~190-240) + - mock_dataset: Standard dataset configuration + - sample_documents: Reusable test documents with varying scores + - mock_flask_app: Flask application context + - mock_thread_pool: Synchronous executor for deterministic testing - def test_raises_when_metadata_model_config_missing(self, retrieval: DatasetRetrieval) -> None: - request = KnowledgeRetrievalRequest( - tenant_id="tenant-1", - user_id="user-1", - app_id="app-1", - user_from="workflow", - dataset_ids=["d1"], - query="query", - retrieval_mode="multiple", - metadata_filtering_mode="automatic", - metadata_model_config=None, - ) - with ( - patch.object(retrieval, "_check_knowledge_rate_limit"), - patch.object(retrieval, "_get_available_datasets", return_value=[_dataset(id="d1")]), - ): - with pytest.raises(ValueError, match="metadata_model_config is required"): - retrieval.knowledge_retrieval(self.orm_session, request) + 2. Vector Search Tests (lines ~240-350) + - Basic functionality + - Document filtering + - Empty results + - Metadata filtering + - Score thresholds - @pytest.mark.parametrize( - ("status", "error_cls"), - [ - (ModelStatus.NO_CONFIGURE, "ModelCredentialsNotInitializedError"), - (ModelStatus.NO_PERMISSION, "ModelNotSupportedError"), - (ModelStatus.QUOTA_EXCEEDED, "ModelQuotaExceededError"), - ], - ) - def test_single_mode_raises_for_model_status( - self, - retrieval: DatasetRetrieval, - status: ModelStatus, - error_cls: str, - ) -> None: - request = KnowledgeRetrievalRequest( - tenant_id="tenant-1", - user_id="user-1", - app_id="app-1", - user_from="workflow", - dataset_ids=["dataset-1"], - query="python", - retrieval_mode="single", - model_provider="openai", - model_name="gpt-4", - ) - provider_model_bundle = Mock() - provider_model_bundle.configuration.get_provider_model.return_value = SimpleNamespace(status=status) - model_type_instance = Mock() - model_type_instance.get_model_schema.return_value = Mock() - model_instance = SimpleNamespace( - provider_model_bundle=provider_model_bundle, - model_type_instance=model_type_instance, - credentials={}, - ) - with ( - patch.object(retrieval, "_check_knowledge_rate_limit"), - patch.object(retrieval, "_get_available_datasets", return_value=[_dataset(id="dataset-1")]), - patch("core.rag.retrieval.dataset_retrieval.ModelManager.for_tenant") as mock_model_manager, - ): - mock_model_manager.return_value.get_model_instance.return_value = model_instance - with pytest.raises( - Exception, - match=( - "(?:Dify Hosted OpenAI gpt\\-4 currently not support\\.|Model gpt\\-4 " - "credentials is not initialized\\.|Model provider openai quota " - "exceeded\\.)" + 3. Keyword Search Tests (lines ~350-450) + - Basic keyword matching + - Special character handling + - Document filtering + + 4. Hybrid Search Tests (lines ~450-640) + - Vector + full-text combination + - Deduplication logic + - Weighted score merging + + 5. Full-Text Search Tests (lines ~640-680) + - BM25-based search + + 6. Score Merging Tests (lines ~680-790) + - Deduplication algorithms + - Score comparison + - Provider-specific handling + + 7. Error Handling Tests (lines ~790-920) + - Empty queries + - Non-existent datasets + - Exception propagation + + 8. Additional Tests (lines ~920-1080) + - Query escaping + - Reranking integration + - Top-K limiting + + Mocking Strategy: + ================= + Tests mock at the method level (embedding_search, keyword_search, etc.) + rather than the underlying Vector/Keyword classes. This approach: + - Avoids complexity of mocking ThreadPoolExecutor behavior + - Provides clearer test intent + - Makes tests more maintainable + - Properly simulates the in-place list modification pattern + + Common Patterns: + ================ + 1. **Arrange**: Set up mocks with side_effect functions + 2. **Act**: Call RetrievalService.retrieve() with specific parameters + 3. **Assert**: Verify results, mock calls, and side effects + + Example Test Structure: + ```python + def test_example(self, mock_get_dataset, mock_search, mock_dataset): + # Arrange: Set up test data and mocks + mock_get_dataset.return_value = mock_dataset + mock_search.side_effect = create_side_effect_for_search([doc1, doc2]) + + # Act: Execute the method under test + results = RetrievalService.retrieve(...) + + # Assert: Verify expectations + assert len(results) == 2 + mock_search.assert_called_once() + ``` + """ + + @pytest.fixture + def mock_dataset(self) -> Dataset: + """ + Create a Dataset object for testing. + + Returns: + Dataset with standard configuration. + """ + dataset = Dataset( + id=str(uuid4()), + tenant_id=str(uuid4()), + name="test_dataset", + indexing_technique="high_quality", + embedding_model="text-embedding-ada-002", + embedding_model_provider="openai", + retrieval_model={ + "search_method": RetrievalMethod.SEMANTIC_SEARCH, + "reranking_enable": False, + "top_k": 4, + "score_threshold_enabled": False, + }, + ) + return dataset + + @pytest.fixture + def sample_documents(self) -> list[_dataset_retrieval_Document]: + """ + Create sample documents for testing retrieval results. + + Returns: + list[Document]: List of mock documents with varying scores + """ + return [ + _dataset_retrieval_Document( + page_content="Python is a high-level programming language.", + metadata={ + "doc_id": "doc1", + "document_id": str(uuid4()), + "dataset_id": str(uuid4()), + "score": 0.95, + }, + provider="dify", ), - ) as exc_info: - retrieval.knowledge_retrieval(self.orm_session, request) - mock_model_manager.assert_called_once_with(tenant_id="tenant-1", user_id="user-1") - assert error_cls in type(exc_info.value).__name__ + _dataset_retrieval_Document( + page_content="JavaScript is widely used for web development.", + metadata={ + "doc_id": "doc2", + "document_id": str(uuid4()), + "dataset_id": str(uuid4()), + "score": 0.85, + }, + provider="dify", + ), + _dataset_retrieval_Document( + page_content="Machine learning is a subset of artificial intelligence.", + metadata={ + "doc_id": "doc3", + "document_id": str(uuid4()), + "dataset_id": str(uuid4()), + "score": 0.75, + }, + provider="dify", + ), + ] + @pytest.fixture + def mock_flask_app(self): + """ + Create a mock Flask application context. -class TestRetrieveCoverage: - @pytest.fixture(autouse=True) - def _orm_session(self, sqlite_session: Session) -> None: - self.orm_session = sqlite_session + Returns: + Mock: Flask app mock with app_context + """ + app = MagicMock() + app.app_context.return_value.__enter__ = Mock() + app.app_context.return_value.__exit__ = Mock() + return app - @pytest.fixture - def retrieval(self) -> DatasetRetrieval: - return DatasetRetrieval() + @pytest.fixture + def retrieval_engine(self, sqlite_engine: Engine): + with patch("core.rag.datasource.retrieval_service.db", SimpleNamespace(engine=sqlite_engine)): + yield sqlite_engine - def _build_model_config(self, features: list[ModelFeature] | None = None): - model_type_instance = Mock() - model_type_instance.get_model_schema.return_value = SimpleNamespace(features=features or []) - provider_bundle = SimpleNamespace(model_type_instance=model_type_instance) - return ModelConfigWithCredentialsEntity.model_construct( - provider="openai", - model="gpt-4", - model_schema=Mock(), - mode="chat", - provider_model_bundle=provider_bundle, - credentials={}, - parameters={}, - stop=[], + @pytest.fixture(autouse=True) + def mock_thread_pool(self): + """ + Mock ThreadPoolExecutor to run tasks synchronously in tests. + + The RetrievalService uses ThreadPoolExecutor to run search operations + concurrently (embedding_search, keyword_search, full_text_index_search). + In tests, we want synchronous execution for: + - Deterministic behavior + - Easier debugging + - Avoiding race conditions + - Simpler assertions + + How it works: + ------------- + 1. Intercepts ThreadPoolExecutor creation + 2. Replaces submit() to execute functions immediately (synchronously) + 3. Functions modify shared all_documents list in-place + 4. Mocks concurrent.futures.wait() since tasks are already done + + Why this approach: + ------------------ + - RetrievalService.retrieve() creates a ThreadPoolExecutor context + - It submits search tasks that modify all_documents list + - concurrent.futures.wait() waits for all tasks to complete + - By executing synchronously, we avoid threading complexity in tests + + Returns: + Mock: Mocked ThreadPoolExecutor that executes tasks synchronously + """ + with patch("core.rag.datasource.retrieval_service.ThreadPoolExecutor") as mock_executor: + # Store futures to track submitted tasks (for debugging if needed) + futures_list = [] + + def sync_submit(fn, *args, **kwargs): + """ + Synchronous replacement for ThreadPoolExecutor.submit(). + + Instead of scheduling the function for async execution, + we execute it immediately in the current thread. + + Args: + fn: The function to execute (e.g., embedding_search) + *args, **kwargs: Arguments to pass to the function + + Returns: + Mock: A mock Future object + """ + future = Mock() + try: + # Execute immediately - this modifies all_documents in place + # The function signature is: fn(flask_app, dataset_id, query, + # top_k, all_documents, exceptions, ...) + fn(*args, **kwargs) + future.result.return_value = None + future.exception.return_value = None + except Exception as e: + # If function raises, store exception in future + future.result.return_value = None + future.exception.return_value = e + + futures_list.append(future) + return future + + # Set up the mock executor instance + mock_executor_instance = Mock() + mock_executor_instance.submit = sync_submit + + # Configure context manager behavior (__enter__ and __exit__) + mock_executor.return_value.__enter__.return_value = mock_executor_instance + mock_executor.return_value.__exit__.return_value = None + + # Mock concurrent.futures.wait to do nothing since tasks are already done + # In real code, this waits for all futures to complete + # In tests, futures complete immediately, so wait is a no-op + with patch("core.rag.datasource.retrieval_service.concurrent.futures.wait"): + # Mock concurrent.futures.as_completed for early error propagation + # In real code, this yields futures as they complete + # In tests, we yield all futures immediately since they're already done + def mock_as_completed(futures_list, timeout=None): + """Mock as_completed that yields futures immediately.""" + yield from futures_list + + with patch( + "core.rag.datasource.retrieval_service.concurrent.futures.as_completed", + side_effect=mock_as_completed, + ): + yield mock_executor + + # ==================== Vector Search Tests ==================== + + @patch("core.rag.datasource.retrieval_service.RetrievalService._retrieve") + @patch("core.rag.datasource.retrieval_service.RetrievalService._get_dataset") + def test_vector_search_basic(self, mock_get_dataset, mock_retrieve, mock_dataset, sample_documents): + """ + Test basic vector/semantic search functionality. + + This test validates the core vector search flow: + 1. Dataset is retrieved from database + 2. _retrieve is called via ThreadPoolExecutor + 3. Documents are added to shared all_documents list + 4. Results are returned to caller + + Verifies: + - Vector search is called with correct parameters + - Results are returned in expected format + - Score threshold is applied correctly + - Documents maintain their metadata and scores + """ + # ==================== ARRANGE ==================== + # Set up the mock dataset that will be "retrieved" from database + mock_get_dataset.return_value = mock_dataset + + # Create a side effect function that simulates _retrieve behavior + # _retrieve modifies the all_documents list in place + def side_effect_retrieve( + flask_app, + retrieval_method, + dataset, + query=None, + top_k=4, + score_threshold=None, + reranking_model=None, + reranking_mode="reranking_model", + weights=None, + document_ids_filter=None, + attachment_id=None, + all_documents=None, + exceptions=None, + ): + """Simulate _retrieve adding documents to the shared list.""" + if all_documents is not None: + all_documents.extend(sample_documents) + + mock_retrieve.side_effect = side_effect_retrieve + + # Define test parameters + query = "What is Python?" # Natural language query + top_k = 3 # Maximum number of results to return + score_threshold = 0.7 # Minimum relevance score (0.0 to 1.0) + + # ==================== ACT ==================== + # Call the retrieve method with SEMANTIC_SEARCH strategy + # This will: + # 1. Check if query is empty (early return if so) + # 2. Get the dataset using _get_dataset + # 3. Create ThreadPoolExecutor + # 4. Submit _retrieve task + # 5. Wait for completion + # 6. Return all_documents list + results = RetrievalService.retrieve( + retrieval_method=RetrievalMethod.SEMANTIC_SEARCH, + dataset_id=mock_dataset.id, + query=query, + top_k=top_k, + score_threshold=score_threshold, + ) + + # ==================== ASSERT ==================== + # Verify we got the expected number of documents + assert len(results) == 3, "Should return 3 documents from sample_documents" + + # Verify all results are Document objects (type safety) + assert all(isinstance(doc, _dataset_retrieval_Document) for doc in results), ( + "All results should be Document instances" + ) + + # Verify documents maintain their scores (highest score first in sample_documents) + assert results[0].metadata["score"] == 0.95, ( + "First document should have highest score from sample_documents" + ) + + # Verify _retrieve was called exactly once + # This confirms the search method was invoked by ThreadPoolExecutor + mock_retrieve.assert_called_once() + + @patch("core.rag.datasource.retrieval_service.RetrievalService._retrieve") + @patch("core.rag.datasource.retrieval_service.RetrievalService._get_dataset") + def test_vector_search_with_document_filter( + self, mock_get_dataset, mock_retrieve, mock_dataset, sample_documents + ): + """ + Test vector search with document ID filtering. + + Verifies: + - Document ID filter is passed correctly to vector search + - Only specified documents are searched + """ + # Arrange + mock_get_dataset.return_value = mock_dataset + filtered_docs = [sample_documents[0]] + + def side_effect_retrieve( + flask_app, + retrieval_method, + dataset, + query=None, + top_k=4, + score_threshold=None, + reranking_model=None, + reranking_mode="reranking_model", + weights=None, + document_ids_filter=None, + attachment_id=None, + all_documents=None, + exceptions=None, + ): + if all_documents is not None: + all_documents.extend(filtered_docs) + + mock_retrieve.side_effect = side_effect_retrieve + document_ids_filter = [sample_documents[0].metadata["document_id"]] + + # Act + results = RetrievalService.retrieve( + retrieval_method=RetrievalMethod.SEMANTIC_SEARCH, + dataset_id=mock_dataset.id, + query="test query", + top_k=5, + document_ids_filter=document_ids_filter, + ) + + # Assert + assert len(results) == 1 + assert results[0].metadata["doc_id"] == "doc1" + # Verify document_ids_filter was passed + call_kwargs = mock_retrieve.call_args.kwargs + assert call_kwargs["document_ids_filter"] == document_ids_filter + + @patch("core.rag.datasource.retrieval_service.RetrievalService._retrieve") + @patch("core.rag.datasource.retrieval_service.RetrievalService._get_dataset") + def test_vector_search_empty_results(self, mock_get_dataset, mock_retrieve, mock_dataset): + """ + Test vector search when no results match the query. + + Verifies: + - Empty list is returned when no documents match + - No errors are raised + """ + # Arrange + mock_get_dataset.return_value = mock_dataset + # _retrieve doesn't add anything to all_documents + mock_retrieve.side_effect = lambda *args, **kwargs: None + + # Act + results = RetrievalService.retrieve( + retrieval_method=RetrievalMethod.SEMANTIC_SEARCH, + dataset_id=mock_dataset.id, + query="nonexistent query", + top_k=5, + ) + + # Assert + assert results == [] + + # ==================== Keyword Search Tests ==================== + + @patch("core.rag.datasource.retrieval_service.RetrievalService._retrieve") + @patch("core.rag.datasource.retrieval_service.RetrievalService._get_dataset") + def test_keyword_search_basic(self, mock_get_dataset, mock_retrieve, mock_dataset, sample_documents): + """ + Test basic keyword search functionality. + + Verifies: + - Keyword search is invoked correctly + - Query is escaped properly for search + - Results are returned in expected format + """ + # Arrange + mock_get_dataset.return_value = mock_dataset + + def side_effect_retrieve( + flask_app, + retrieval_method, + dataset, + query=None, + top_k=4, + score_threshold=None, + reranking_model=None, + reranking_mode="reranking_model", + weights=None, + document_ids_filter=None, + attachment_id=None, + all_documents=None, + exceptions=None, + ): + if all_documents is not None: + all_documents.extend(sample_documents) + + mock_retrieve.side_effect = side_effect_retrieve + + query = "Python programming" + top_k = 3 + + # Act + results = RetrievalService.retrieve( + retrieval_method=RetrievalMethod.KEYWORD_SEARCH, + dataset_id=mock_dataset.id, + query=query, + top_k=top_k, + ) + + # Assert + assert len(results) == 3 + assert all(isinstance(doc, _dataset_retrieval_Document) for doc in results) + mock_retrieve.assert_called_once() + + @patch("core.rag.datasource.retrieval_service.RetrievalService.keyword_search") + @patch("core.rag.datasource.retrieval_service.RetrievalService._get_dataset") + def test_keyword_search_with_special_characters(self, mock_get_dataset, mock_keyword_search, mock_dataset): + """ + Test keyword search with special characters in query. + + Verifies: + - Special characters are escaped correctly + - Search handles quotes and other special chars + """ + # Arrange + mock_get_dataset.return_value = mock_dataset + mock_keyword_search.side_effect = lambda *args, **kwargs: None + + query = 'Python "programming" language' + + # Act + RetrievalService.retrieve( + retrieval_method=RetrievalMethod.KEYWORD_SEARCH, + dataset_id=mock_dataset.id, + query=query, + top_k=5, + ) + + # Assert + # Verify that keyword_search was called + assert mock_keyword_search.called + # The query escaping happens inside keyword_search method + call_args = mock_keyword_search.call_args + assert call_args is not None + + @patch("core.rag.datasource.retrieval_service.RetrievalService.keyword_search") + @patch("core.rag.datasource.retrieval_service.RetrievalService._get_dataset") + def test_keyword_search_with_document_filter( + self, mock_get_dataset, mock_keyword_search, mock_dataset, sample_documents + ): + """ + Test keyword search with document ID filtering. + + Verifies: + - Document filter is applied to keyword search + - Only filtered documents are returned + """ + # Arrange + mock_get_dataset.return_value = mock_dataset + filtered_docs = [sample_documents[1]] + + def side_effect_keyword_search( + flask_app, dataset_id, query, top_k, all_documents, exceptions, document_ids_filter=None + ): + all_documents.extend(filtered_docs) + + mock_keyword_search.side_effect = side_effect_keyword_search + document_ids_filter = [sample_documents[1].metadata["document_id"]] + + # Act + results = RetrievalService.retrieve( + retrieval_method=RetrievalMethod.KEYWORD_SEARCH, + dataset_id=mock_dataset.id, + query="JavaScript", + top_k=5, + document_ids_filter=document_ids_filter, + ) + + # Assert + assert len(results) == 1 + assert results[0].metadata["doc_id"] == "doc2" + + # ==================== Hybrid Search Tests ==================== + + @patch("core.rag.datasource.retrieval_service.DataPostProcessor") + @patch("core.rag.datasource.retrieval_service.RetrievalService.full_text_index_search") + @patch("core.rag.datasource.retrieval_service.RetrievalService.embedding_search") + @patch("core.rag.datasource.retrieval_service.RetrievalService._get_dataset") + def test_hybrid_search_basic( + self, + mock_get_dataset, + mock_embedding_search, + mock_fulltext_search, + mock_data_processor_class, + mock_dataset, + sample_documents, + retrieval_engine, + ): + """ + Test basic hybrid search combining vector and full-text search. + + Verifies: + - Both vector and full-text search are executed + - Results are merged and deduplicated + - DataPostProcessor is invoked for score merging + """ + # Arrange + mock_get_dataset.return_value = mock_dataset + + # Vector search returns first 2 docs + def side_effect_embedding( + flask_app, + dataset_id, + query, + top_k, + score_threshold, + reranking_model, + all_documents, + retrieval_method, + exceptions, + document_ids_filter=None, + ): + all_documents.extend(sample_documents[:2]) + + mock_embedding_search.side_effect = side_effect_embedding + + # Full-text search returns last 2 docs (with overlap) + def side_effect_fulltext( + flask_app, + dataset_id, + query, + top_k, + score_threshold, + reranking_model, + all_documents, + retrieval_method, + exceptions, + document_ids_filter=None, + ): + all_documents.extend(sample_documents[1:]) + + mock_fulltext_search.side_effect = side_effect_fulltext + + # Mock DataPostProcessor + mock_processor_instance = Mock() + mock_processor_instance.invoke.return_value = sample_documents + mock_data_processor_class.return_value = mock_processor_instance + + # Act + results = RetrievalService.retrieve( + retrieval_method=RetrievalMethod.HYBRID_SEARCH, + dataset_id=mock_dataset.id, + query="Python programming", + top_k=3, + score_threshold=0.5, + ) + + # Assert + assert len(results) == 3 + mock_embedding_search.assert_called_once() + mock_fulltext_search.assert_called_once() + mock_processor_instance.invoke.assert_called_once() + processor_session = mock_data_processor_class.call_args.kwargs["session"] + assert isinstance(processor_session, Session) + assert processor_session.get_bind() is retrieval_engine + + @patch("core.rag.datasource.retrieval_service.DataPostProcessor") + @patch("core.rag.datasource.retrieval_service.RetrievalService.full_text_index_search") + @patch("core.rag.datasource.retrieval_service.RetrievalService.embedding_search") + @patch("core.rag.datasource.retrieval_service.RetrievalService._get_dataset") + def test_hybrid_search_deduplication( + self, + mock_get_dataset, + mock_embedding_search, + mock_fulltext_search, + mock_data_processor_class, + mock_dataset, + retrieval_engine, + ): + """ + Test that hybrid search properly deduplicates documents. + + Hybrid search combines results from multiple search methods (vector + full-text). + This can lead to duplicate documents when the same chunk is found by both methods. + + Scenario: + --------- + 1. Vector search finds document "duplicate_doc" with score 0.9 + 2. Full-text search also finds "duplicate_doc" but with score 0.6 + 3. Both searches find "unique_doc" + 4. Deduplication should keep only the higher-scoring version (0.9) + + Why deduplication matters: + -------------------------- + - Prevents showing the same content multiple times to users + - Ensures score consistency (keeps best match) + - Improves result quality and user experience + - Happens BEFORE reranking to avoid processing duplicates + + Verifies: + - Duplicate documents (same doc_id) are removed + - Higher scoring duplicate is retained + - Deduplication happens before post-processing + - Final result count is correct + """ + # ==================== ARRANGE ==================== + mock_get_dataset.return_value = mock_dataset + + # Create test documents with intentional duplication + # Same doc_id but different scores to test score comparison logic + doc1_high = _dataset_retrieval_Document( + page_content="Content 1", + metadata={ + "doc_id": "duplicate_doc", # Same doc_id as doc1_low + "score": 0.9, # Higher score - should be kept + "document_id": str(uuid4()), + }, + provider="dify", + ) + doc1_low = _dataset_retrieval_Document( + page_content="Content 1", + metadata={ + "doc_id": "duplicate_doc", # Same doc_id as doc1_high + "score": 0.6, # Lower score - should be discarded + "document_id": str(uuid4()), + }, + provider="dify", + ) + doc2 = _dataset_retrieval_Document( + page_content="Content 2", + metadata={ + "doc_id": "unique_doc", # Unique doc_id + "score": 0.8, + "document_id": str(uuid4()), + }, + provider="dify", + ) + + # Simulate vector search returning high-score duplicate + unique doc + def side_effect_embedding( + flask_app, + dataset_id, + query, + top_k, + score_threshold, + reranking_model, + all_documents, + retrieval_method, + exceptions, + document_ids_filter=None, + ): + """Vector search finds 2 documents including high-score duplicate.""" + all_documents.extend([doc1_high, doc2]) + + mock_embedding_search.side_effect = side_effect_embedding + + # Simulate full-text search returning low-score duplicate + def side_effect_fulltext( + flask_app, + dataset_id, + query, + top_k, + score_threshold, + reranking_model, + all_documents, + retrieval_method, + exceptions, + document_ids_filter=None, + ): + """Full-text search finds the same document but with lower score.""" + all_documents.extend([doc1_low]) + + mock_fulltext_search.side_effect = side_effect_fulltext + + # Mock DataPostProcessor to return deduplicated results + # In real implementation, _deduplicate_documents is called before this + mock_processor_instance = Mock() + mock_processor_instance.invoke.return_value = [doc1_high, doc2] + mock_data_processor_class.return_value = mock_processor_instance + + # ==================== ACT ==================== + # Execute hybrid search which should: + # 1. Run both embedding_search and full_text_index_search + # 2. Collect all results in all_documents (3 docs: 2 unique + 1 duplicate) + # 3. Call _deduplicate_documents to remove duplicate (keeps higher score) + # 4. Pass deduplicated results to DataPostProcessor + # 5. Return final results + results = RetrievalService.retrieve( + retrieval_method=RetrievalMethod.HYBRID_SEARCH, + dataset_id=mock_dataset.id, + query="test", + top_k=5, + ) + + # ==================== ASSERT ==================== + # Verify deduplication worked correctly + assert len(results) == 2, "Should have 2 unique documents after deduplication (not 3)" + + # Verify the correct documents are present + doc_ids = [doc.metadata["doc_id"] for doc in results] + assert "duplicate_doc" in doc_ids, "Duplicate doc should be present (higher score version)" + assert "unique_doc" in doc_ids, "Unique doc should be present" + processor_session = mock_data_processor_class.call_args.kwargs["session"] + assert isinstance(processor_session, Session) + assert processor_session.get_bind() is retrieval_engine + + # Implicitly verifies that doc1_low (score 0.6) was discarded + # in favor of doc1_high (score 0.9) + + @patch("core.rag.datasource.retrieval_service.DataPostProcessor") + @patch("core.rag.datasource.retrieval_service.RetrievalService.full_text_index_search") + @patch("core.rag.datasource.retrieval_service.RetrievalService.embedding_search") + @patch("core.rag.datasource.retrieval_service.RetrievalService._get_dataset") + def test_hybrid_search_with_weights( + self, + mock_get_dataset, + mock_embedding_search, + mock_fulltext_search, + mock_data_processor_class, + mock_dataset, + sample_documents, + retrieval_engine, + ): + """ + Test hybrid search with custom weights for score merging. + + Verifies: + - Weights are passed to DataPostProcessor + - Score merging respects weight configuration + """ + # Arrange + mock_get_dataset.return_value = mock_dataset + + def side_effect_embedding( + flask_app, + dataset_id, + query, + top_k, + score_threshold, + reranking_model, + all_documents, + retrieval_method, + exceptions, + document_ids_filter=None, + ): + all_documents.extend(sample_documents[:2]) + + mock_embedding_search.side_effect = side_effect_embedding + + def side_effect_fulltext( + flask_app, + dataset_id, + query, + top_k, + score_threshold, + reranking_model, + all_documents, + retrieval_method, + exceptions, + document_ids_filter=None, + ): + all_documents.extend(sample_documents[1:]) + + mock_fulltext_search.side_effect = side_effect_fulltext + + mock_processor_instance = Mock() + mock_processor_instance.invoke.return_value = sample_documents + mock_data_processor_class.return_value = mock_processor_instance + + weights = { + "vector_setting": { + "vector_weight": 0.7, + "embedding_provider_name": "openai", + "embedding_model_name": "text-embedding-ada-002", + }, + "keyword_setting": {"keyword_weight": 0.3}, + } + + # Act + results = RetrievalService.retrieve( + retrieval_method=RetrievalMethod.HYBRID_SEARCH, + dataset_id=mock_dataset.id, + query="test query", + top_k=3, + weights=weights, + reranking_mode="weighted_score", + ) + + # Assert + assert len(results) == 3 + # Verify DataPostProcessor was created with weights + mock_data_processor_class.assert_called_once() + call_args = mock_data_processor_class.call_args + assert call_args.args[3] == weights + processor_session = call_args.kwargs["session"] + assert isinstance(processor_session, Session) + assert processor_session.get_bind() is retrieval_engine + + @pytest.mark.parametrize("empty_query", ["", None]) + @patch("core.rag.datasource.retrieval_service.DataPostProcessor") + @patch("core.rag.datasource.retrieval_service.RetrievalService.embedding_search") + @patch("core.rag.datasource.retrieval_service.RetrievalService._get_dataset") + def test_hybrid_search_attachment_only_uses_image_query_type( + self, + mock_get_dataset, + mock_embedding_search, + mock_data_processor_class, + mock_dataset, + sample_documents, + empty_query, + retrieval_engine, + ): + """ + Regression test for GH #37116: attachment-only hybrid retrieval must use IMAGE_QUERY. + + When HYBRID_SEARCH is invoked with no text query and a non-None attachment_id, + DataPostProcessor.invoke must receive query_type=QueryType.IMAGE_QUERY. + """ + from core.rag.index_processor.constant.query_type import QueryType + + # Arrange + mock_get_dataset.return_value = mock_dataset + attachment_id = "upload-file-uuid-1234" + + def side_effect_embedding( + flask_app, + dataset_id, + query, + top_k, + score_threshold, + reranking_model, + all_documents, + retrieval_method, + exceptions, + document_ids_filter=None, + query_type=QueryType.TEXT_QUERY, + ): + all_documents.extend(sample_documents[:2]) + + mock_embedding_search.side_effect = side_effect_embedding + + mock_processor_instance = Mock() + mock_processor_instance.invoke.return_value = sample_documents[:2] + mock_data_processor_class.return_value = mock_processor_instance + + # Act: call retrieve with attachment_ids only, no text query + RetrievalService.retrieve( + retrieval_method=RetrievalMethod.HYBRID_SEARCH, + dataset_id=mock_dataset.id, + query=empty_query, + top_k=3, + score_threshold=0.5, + attachment_ids=[attachment_id], + ) + + # Assert: invoke must have been called with IMAGE_QUERY + mock_processor_instance.invoke.assert_called_once() + invoke_kwargs = mock_processor_instance.invoke.call_args.kwargs + assert invoke_kwargs["query_type"] == QueryType.IMAGE_QUERY, ( + "Attachment-only hybrid search must use IMAGE_QUERY for reranking, not TEXT_QUERY" + ) + assert invoke_kwargs["query"] == attachment_id, ( + "The rerank query must be the attachment_id, not the empty text query" + ) + processor_session = mock_data_processor_class.call_args.kwargs["session"] + assert isinstance(processor_session, Session) + assert processor_session.get_bind() is retrieval_engine + + # ==================== Full-Text Search Tests ==================== + + @patch("core.rag.datasource.retrieval_service.RetrievalService.full_text_index_search") + @patch("core.rag.datasource.retrieval_service.RetrievalService._get_dataset") + def test_fulltext_search_basic(self, mock_get_dataset, mock_fulltext_search, mock_dataset, sample_documents): + """ + Test basic full-text search functionality. + + Verifies: + - Full-text search is invoked correctly + - Results are returned in expected format + """ + # Arrange + mock_get_dataset.return_value = mock_dataset + + def side_effect_fulltext( + flask_app, + dataset_id, + query, + top_k, + score_threshold, + reranking_model, + all_documents, + retrieval_method, + exceptions, + document_ids_filter=None, + ): + all_documents.extend(sample_documents) + + mock_fulltext_search.side_effect = side_effect_fulltext + + # Act + results = RetrievalService.retrieve( + retrieval_method=RetrievalMethod.FULL_TEXT_SEARCH, + dataset_id=mock_dataset.id, + query="programming language", + top_k=3, + ) + + # Assert + assert len(results) == 3 + mock_fulltext_search.assert_called_once() + + # ==================== Score Merging Tests ==================== + + def test_deduplicate_documents_basic(self): + """ + Test basic document deduplication logic. + + Verifies: + - Documents with same doc_id are deduplicated + - First occurrence is kept by default + """ + # Arrange + doc1 = _dataset_retrieval_Document( + page_content="Content 1", + metadata={"doc_id": "doc1", "score": 0.8}, + provider="dify", + ) + doc2 = _dataset_retrieval_Document( + page_content="Content 2", + metadata={"doc_id": "doc2", "score": 0.7}, + provider="dify", + ) + doc1_duplicate = _dataset_retrieval_Document( + page_content="Content 1 duplicate", + metadata={"doc_id": "doc1", "score": 0.6}, + provider="dify", + ) + + documents = [doc1, doc2, doc1_duplicate] + + # Act + result = RetrievalService._deduplicate_documents(documents) + + # Assert + assert len(result) == 2 + doc_ids = [doc.metadata["doc_id"] for doc in result] + assert doc_ids == ["doc1", "doc2"] + + def test_deduplicate_documents_keeps_higher_score(self): + """ + Test that deduplication keeps document with higher score. + + Verifies: + - When duplicates exist, higher scoring version is retained + - Score comparison works correctly + """ + # Arrange + doc_low = _dataset_retrieval_Document( + page_content="Content", + metadata={"doc_id": "doc1", "score": 0.5}, + provider="dify", + ) + doc_high = _dataset_retrieval_Document( + page_content="Content", + metadata={"doc_id": "doc1", "score": 0.9}, + provider="dify", + ) + + # Low score first + documents = [doc_low, doc_high] + + # Act + result = RetrievalService._deduplicate_documents(documents) + + # Assert + assert len(result) == 1 + assert result[0].metadata["score"] == 0.9 + + def test_deduplicate_documents_empty_list(self): + """ + Test deduplication with empty document list. + + Verifies: + - Empty list returns empty list + - No errors are raised + """ + # Act + result = RetrievalService._deduplicate_documents([]) + + # Assert + assert result == [] + + def test_deduplicate_documents_non_dify_provider(self): + """ + Test deduplication with non-dify provider documents that have no doc_id. + + Verifies: + - External provider documents without doc_id use content-based deduplication + - Identical content from the same provider is collapsed to one result + """ + # Arrange + doc1 = _dataset_retrieval_Document( + page_content="External content", + metadata={"score": 0.8}, + provider="external", + ) + doc2 = _dataset_retrieval_Document( + page_content="External content", + metadata={"score": 0.7}, + provider="external", + ) + + documents = [doc1, doc2] + + # Act + result = RetrievalService._deduplicate_documents(documents) + + # Assert + # External documents without doc_id should use content-based dedup + assert len(result) == 1 + + def test_deduplicate_documents_non_dify_provider_with_doc_id_different_sources(self): + """ + Regression test for issue #35707. + + Two chunks from different source documents share identical text content but carry + different doc_ids. Before the fix, non-dify providers were forced into content-based + deduplication and the second chunk was silently dropped. After the fix, doc_id is used + as the dedup key for any provider that exposes it, so both chunks must be retained. + + Verifies: + - Non-dify provider documents with different doc_ids are NOT deduplicated even when + their page_content is identical. + """ + # Arrange — same content, different doc_ids, non-dify provider (e.g. Weaviate / Qdrant) + doc_a = _dataset_retrieval_Document( + page_content="Shared identical content", + metadata={"doc_id": "doc-from-file-a", "score": 0.85}, + provider="weaviate", + ) + doc_b = _dataset_retrieval_Document( + page_content="Shared identical content", + metadata={"doc_id": "doc-from-file-b", "score": 0.82}, + provider="weaviate", + ) + + # Act + result = RetrievalService._deduplicate_documents([doc_a, doc_b]) + + # Assert — both documents must be kept; losing either silently drops a source citation + assert len(result) == 2 + doc_ids = {doc.metadata["doc_id"] for doc in result} + assert doc_ids == {"doc-from-file-a", "doc-from-file-b"} + + def test_deduplicate_documents_non_dify_provider_with_same_doc_id(self): + """ + Test that non-dify provider documents sharing the same doc_id are deduplicated by + doc_id key (not by content), and the higher-scored duplicate is retained. + + Verifies: + - doc_id-based deduplication now applies to any provider, not only "dify" + - The document with the highest score wins when doc_ids collide + """ + # Arrange + doc_low = _dataset_retrieval_Document( + page_content="Content A", + metadata={"doc_id": "chunk-1", "score": 0.5}, + provider="qdrant", + ) + doc_high = _dataset_retrieval_Document( + page_content="Content A", + metadata={"doc_id": "chunk-1", "score": 0.9}, + provider="qdrant", + ) + + # Act + result = RetrievalService._deduplicate_documents([doc_low, doc_high]) + + # Assert + assert len(result) == 1 + assert result[0].metadata["score"] == 0.9 + + def test_deduplicate_documents_dify_provider_without_doc_id_falls_back_to_content(self): + """ + Test that a dify provider document without doc_id still falls back to content-based + deduplication (no regression from original behaviour). + + Verifies: + - Absence of doc_id triggers content-based dedup regardless of provider + - First occurrence is kept when content is identical + """ + # Arrange — dify docs with no doc_id, same content + doc1 = _dataset_retrieval_Document( + page_content="Same content", + metadata={"score": 0.8}, + provider="dify", + ) + doc2 = _dataset_retrieval_Document( + page_content="Same content", + metadata={"score": 0.9}, + provider="dify", + ) + + # Act + result = RetrievalService._deduplicate_documents([doc1, doc2]) + + # Assert — collapsed to one; first-seen wins (no score comparison in content branch) + assert len(result) == 1 + assert result[0].metadata["score"] == 0.8 + + # ==================== Metadata Filtering Tests ==================== + + @patch("core.rag.datasource.retrieval_service.RetrievalService._retrieve") + @patch("core.rag.datasource.retrieval_service.RetrievalService._get_dataset") + def test_vector_search_with_metadata_filter( + self, mock_get_dataset, mock_retrieve, mock_dataset, sample_documents + ): + """ + Test vector search with metadata-based document filtering. + + Verifies: + - Metadata filters are applied correctly + - Only documents matching metadata criteria are returned + """ + # Arrange + mock_get_dataset.return_value = mock_dataset + + # Add metadata to documents + filtered_doc = sample_documents[0] + filtered_doc.metadata["category"] = "programming" + + def side_effect_retrieve( + flask_app, + retrieval_method, + dataset, + query=None, + top_k=4, + score_threshold=None, + reranking_model=None, + reranking_mode="reranking_model", + weights=None, + document_ids_filter=None, + attachment_id=None, + all_documents=None, + exceptions=None, + ): + if all_documents is not None: + all_documents.append(filtered_doc) + + mock_retrieve.side_effect = side_effect_retrieve + + # Act + results = RetrievalService.retrieve( + retrieval_method=RetrievalMethod.SEMANTIC_SEARCH, + dataset_id=mock_dataset.id, + query="Python", + top_k=5, + document_ids_filter=[filtered_doc.metadata["document_id"]], + ) + + # Assert + assert len(results) == 1 + assert results[0].metadata.get("category") == "programming" + + # ==================== Error Handling Tests ==================== + + @patch("core.rag.datasource.retrieval_service.RetrievalService._get_dataset") + def test_retrieve_with_empty_query(self, mock_get_dataset, mock_dataset): + """ + Test retrieval with empty query string. + + Verifies: + - Empty query returns empty results + - No search operations are performed + """ + # Arrange + mock_get_dataset.return_value = mock_dataset + + # Act + results = RetrievalService.retrieve( + retrieval_method=RetrievalMethod.SEMANTIC_SEARCH, + dataset_id=mock_dataset.id, + query="", + top_k=5, + ) + + # Assert + assert results == [] + + @patch("core.rag.datasource.retrieval_service.RetrievalService._get_dataset") + def test_retrieve_with_nonexistent_dataset(self, mock_get_dataset): + """ + Test retrieval with non-existent dataset ID. + + Verifies: + - Non-existent dataset returns empty results + - No errors are raised + """ + # Arrange + mock_get_dataset.return_value = None + + # Act + results = RetrievalService.retrieve( + retrieval_method=RetrievalMethod.SEMANTIC_SEARCH, + dataset_id="nonexistent_id", + query="test query", + top_k=5, + ) + + # Assert + assert results == [] + + @patch("core.rag.datasource.retrieval_service.RetrievalService._retrieve") + @patch("core.rag.datasource.retrieval_service.RetrievalService._get_dataset") + def test_retrieve_with_exception_handling(self, mock_get_dataset, mock_retrieve, mock_dataset): + """ + Test that exceptions during retrieval are properly handled. + + Verifies: + - Exceptions are caught and added to exceptions list + - ValueError is raised with exception messages + """ + # Arrange + mock_get_dataset.return_value = mock_dataset + + # Make _retrieve add an exception to the exceptions list + def side_effect_with_exception( + flask_app, + retrieval_method, + dataset, + query=None, + top_k=4, + score_threshold=None, + reranking_model=None, + reranking_mode="reranking_model", + weights=None, + document_ids_filter=None, + attachment_id=None, + all_documents=None, + exceptions=None, + ): + if exceptions is not None: + exceptions.append("Search failed") + + mock_retrieve.side_effect = side_effect_with_exception + + # Act & Assert + with pytest.raises(ValueError, match="Search failed") as exc_info: + RetrievalService.retrieve( + retrieval_method=RetrievalMethod.SEMANTIC_SEARCH, + dataset_id=mock_dataset.id, + query="test query", + top_k=5, + ) + + assert "Search failed" in str(exc_info.value) + + # ==================== Score Threshold Tests ==================== + + @patch("core.rag.datasource.retrieval_service.RetrievalService._retrieve") + @patch("core.rag.datasource.retrieval_service.RetrievalService._get_dataset") + def test_vector_search_with_score_threshold(self, mock_get_dataset, mock_retrieve, mock_dataset): + """ + Test vector search with score threshold filtering. + + Verifies: + - Score threshold is passed to search method + - Documents below threshold are filtered out + """ + # Arrange + mock_get_dataset.return_value = mock_dataset + + # Only return documents above threshold + high_score_doc = _dataset_retrieval_Document( + page_content="High relevance content", + metadata={"doc_id": "doc1", "score": 0.85}, + provider="dify", + ) + + def side_effect_retrieve( + flask_app, + retrieval_method, + dataset, + query=None, + top_k=4, + score_threshold=None, + reranking_model=None, + reranking_mode="reranking_model", + weights=None, + document_ids_filter=None, + attachment_id=None, + all_documents=None, + exceptions=None, + ): + if all_documents is not None: + all_documents.append(high_score_doc) + + mock_retrieve.side_effect = side_effect_retrieve + + score_threshold = 0.8 + + # Act + results = RetrievalService.retrieve( + retrieval_method=RetrievalMethod.SEMANTIC_SEARCH, + dataset_id=mock_dataset.id, + query="test query", + top_k=5, + score_threshold=score_threshold, + ) + + # Assert + assert len(results) == 1 + assert results[0].metadata["score"] >= score_threshold + + # ==================== Top-K Limiting Tests ==================== + + @patch("core.rag.datasource.retrieval_service.RetrievalService._retrieve") + @patch("core.rag.datasource.retrieval_service.RetrievalService._get_dataset") + def test_retrieve_respects_top_k_limit(self, mock_get_dataset, mock_retrieve, mock_dataset): + """ + Test that retrieval respects top_k parameter. + + Verifies: + - Only top_k documents are returned + - Limit is applied correctly + """ + # Arrange + mock_get_dataset.return_value = mock_dataset + + # Create more documents than top_k + many_docs = [ + _dataset_retrieval_Document( + page_content=f"Content {i}", + metadata={"doc_id": f"doc{i}", "score": 0.9 - i * 0.1}, + provider="dify", + ) + for i in range(10) + ] + + def side_effect_retrieve( + flask_app, + retrieval_method, + dataset, + query=None, + top_k=4, + score_threshold=None, + reranking_model=None, + reranking_mode="reranking_model", + weights=None, + document_ids_filter=None, + attachment_id=None, + all_documents=None, + exceptions=None, + ): + # Return only top_k documents + if all_documents is not None: + all_documents.extend(many_docs[:top_k]) + + mock_retrieve.side_effect = side_effect_retrieve + + top_k = 3 + + # Act + results = RetrievalService.retrieve( + retrieval_method=RetrievalMethod.SEMANTIC_SEARCH, + dataset_id=mock_dataset.id, + query="test query", + top_k=top_k, + ) + + # Assert + # Verify _retrieve was called + assert mock_retrieve.called + call_kwargs = mock_retrieve.call_args.kwargs + assert call_kwargs["top_k"] == top_k + # Verify we got the right number of results + assert len(results) == top_k + + # ==================== Query Escaping Tests ==================== + + def test_escape_query_for_search(self): + """ + Test query escaping for special characters. + + Verifies: + - Double quotes are properly escaped + - Other characters remain unchanged + """ + # Test cases with expected outputs + test_cases = [ + ("simple query", "simple query"), + ('query with "quotes"', 'query with \\"quotes\\"'), + ('"quoted phrase"', '\\"quoted phrase\\"'), + ("no special chars", "no special chars"), + ] + + for input_query, expected_output in test_cases: + result = RetrievalService.escape_query_for_search(input_query) + assert result == expected_output + + # ==================== Reranking Tests ==================== + + @patch("core.rag.datasource.retrieval_service.RetrievalService._retrieve") + @patch("core.rag.datasource.retrieval_service.RetrievalService._get_dataset") + def test_semantic_search_with_reranking(self, mock_get_dataset, mock_retrieve, mock_dataset, sample_documents): + """ + Test semantic search with reranking model. + + Verifies: + - Reranking is applied when configured + - DataPostProcessor is invoked with correct parameters + """ + # Arrange + mock_get_dataset.return_value = mock_dataset + + # Simulate reranking changing order + reranked_docs = list(reversed(sample_documents)) + + def side_effect_retrieve( + flask_app, + retrieval_method, + dataset, + query=None, + top_k=4, + score_threshold=None, + reranking_model=None, + reranking_mode="reranking_model", + weights=None, + document_ids_filter=None, + attachment_id=None, + all_documents=None, + exceptions=None, + ): + # _retrieve handles reranking internally + if all_documents is not None: + all_documents.extend(reranked_docs) + + mock_retrieve.side_effect = side_effect_retrieve + + reranking_model = { + "reranking_provider_name": "cohere", + "reranking_model_name": "rerank-english-v2.0", + } + + # Act + results = RetrievalService.retrieve( + retrieval_method=RetrievalMethod.SEMANTIC_SEARCH, + dataset_id=mock_dataset.id, + query="test query", + top_k=3, + reranking_model=reranking_model, + ) + + # Assert + # For semantic search with reranking, reranking_model should be passed + assert len(results) == 3 + call_kwargs = mock_retrieve.call_args.kwargs + assert call_kwargs["reranking_model"] == reranking_model + + # ==================== Multiple Retrieve Thread Tests ==================== + + @patch("core.rag.retrieval.dataset_retrieval.DataPostProcessor") + @patch("core.rag.retrieval.dataset_retrieval.DatasetRetrieval._retriever") + def test_multiple_retrieve_thread_skips_second_reranking_with_single_dataset( + self, mock_retriever, mock_data_processor_class, mock_flask_app, mock_dataset + ): + """ + Test that _multiple_retrieve_thread skips second reranking when dataset_count is 1. + + When there is only one dataset, the second reranking is unnecessary + because the documents are already ranked from the first retrieval. + This optimization avoids the overhead of reranking when it won't + provide any benefit. + + Verifies: + - DataPostProcessor is NOT called when dataset_count == 1 + - Documents are still added to all_documents + - Standard scoring logic is applied instead + """ + # Arrange + dataset_retrieval = DatasetRetrieval() + tenant_id = str(uuid4()) + + # Create test documents + doc1 = _dataset_retrieval_Document( + page_content="Test content 1", + metadata={"doc_id": "doc1", "score": 0.9, "document_id": str(uuid4()), "dataset_id": mock_dataset.id}, + provider="dify", + ) + doc2 = _dataset_retrieval_Document( + page_content="Test content 2", + metadata={"doc_id": "doc2", "score": 0.8, "document_id": str(uuid4()), "dataset_id": mock_dataset.id}, + provider="dify", + ) + + # Mock _retriever to return documents + def side_effect_retriever( + flask_app, + session, + dataset_id, + query, + top_k, + all_documents, + document_ids_filter, + metadata_condition, + attachment_ids, + ): + all_documents.extend([doc1, doc2]) + + mock_retriever.side_effect = side_effect_retriever + + # Set up dataset with high_quality indexing + mock_dataset.indexing_technique = "high_quality" + + all_documents = [] + + # Act - Call with dataset_count = 1 + dataset_retrieval._multiple_retrieve_thread( + flask_app=mock_flask_app, + available_datasets=[mock_dataset], + metadata_condition=None, + metadata_filter_document_ids=None, + all_documents=all_documents, + tenant_id=tenant_id, + reranking_enable=True, + reranking_mode="reranking_model", + reranking_model={"reranking_provider_name": "cohere", "reranking_model_name": "rerank-v2"}, + weights=None, + top_k=5, + score_threshold=0.5, + query="test query", + attachment_id=None, + dataset_count=1, # Single dataset - should skip second reranking + ) + + # Assert + # DataPostProcessor should NOT be called (second reranking skipped) + mock_data_processor_class.assert_not_called() + + # Documents should still be added to all_documents + assert len(all_documents) == 2 + assert all_documents[0].page_content == "Test content 1" + assert all_documents[1].page_content == "Test content 2" + + @patch("core.rag.retrieval.dataset_retrieval.DataPostProcessor") + @patch("core.rag.retrieval.dataset_retrieval.DatasetRetrieval._retriever") + @patch("core.rag.retrieval.dataset_retrieval.DatasetRetrieval.calculate_vector_score") + def test_multiple_retrieve_thread_performs_second_reranking_with_multiple_datasets( + self, mock_calculate_vector_score, mock_retriever, mock_data_processor_class, mock_flask_app, mock_dataset + ): + """ + Test that _multiple_retrieve_thread performs second reranking when dataset_count > 1. + + When there are multiple datasets, the second reranking is necessary + to merge and re-rank results from different datasets. This ensures + the most relevant documents across all datasets are returned. + + Verifies: + - DataPostProcessor IS called when dataset_count > 1 + - Reranking is applied with correct parameters + - Documents are processed correctly + """ + # Arrange + dataset_retrieval = DatasetRetrieval() + tenant_id = str(uuid4()) + + # Create test documents + doc1 = _dataset_retrieval_Document( + page_content="Test content 1", + metadata={"doc_id": "doc1", "score": 0.7, "document_id": str(uuid4()), "dataset_id": mock_dataset.id}, + provider="dify", + ) + doc2 = _dataset_retrieval_Document( + page_content="Test content 2", + metadata={"doc_id": "doc2", "score": 0.6, "document_id": str(uuid4()), "dataset_id": mock_dataset.id}, + provider="dify", + ) + + # Mock _retriever to return documents + def side_effect_retriever( + flask_app, + session, + dataset_id, + query, + top_k, + all_documents, + document_ids_filter, + metadata_condition, + attachment_ids, + ): + all_documents.extend([doc1, doc2]) + + mock_retriever.side_effect = side_effect_retriever + + # Set up dataset with high_quality indexing + mock_dataset.indexing_technique = "high_quality" + + # Mock DataPostProcessor instance and its invoke method + mock_processor_instance = Mock() + # Simulate reranking - return documents in reversed order with updated scores + reranked_docs = [ + _dataset_retrieval_Document( + page_content="Test content 2", + metadata={ + "doc_id": "doc2", + "score": 0.95, + "document_id": str(uuid4()), + "dataset_id": mock_dataset.id, + }, + provider="dify", + ), + _dataset_retrieval_Document( + page_content="Test content 1", + metadata={ + "doc_id": "doc1", + "score": 0.85, + "document_id": str(uuid4()), + "dataset_id": mock_dataset.id, + }, + provider="dify", + ), + ] + mock_processor_instance.invoke.return_value = reranked_docs + mock_data_processor_class.return_value = mock_processor_instance + + all_documents = [] + + # Create second dataset + mock_dataset2 = Dataset( + id=str(uuid4()), + indexing_technique="high_quality", + provider="dify", + ) + + # Act - Call with dataset_count = 2 + dataset_retrieval._multiple_retrieve_thread( + flask_app=mock_flask_app, + available_datasets=[mock_dataset, mock_dataset2], + metadata_condition=None, + metadata_filter_document_ids=None, + all_documents=all_documents, + tenant_id=tenant_id, + reranking_enable=True, + reranking_mode="reranking_model", + reranking_model={"reranking_provider_name": "cohere", "reranking_model_name": "rerank-v2"}, + weights=None, + top_k=5, + score_threshold=0.5, + query="test query", + attachment_id=None, + dataset_count=2, # Multiple datasets - should perform second reranking + ) + + # Assert + # DataPostProcessor SHOULD be called (second reranking performed) + mock_data_processor_class.assert_called_once() + assert isinstance(mock_data_processor_class.call_args.kwargs["session"], Session) + + # Verify invoke was called with correct parameters + mock_processor_instance.invoke.assert_called_once() + + # Documents should be added to all_documents after reranking + assert len(all_documents) == 2 + # The reranked order should be reflected + assert all_documents[0].page_content == "Test content 2" + assert all_documents[1].page_content == "Test content 1" + + @patch("core.rag.retrieval.dataset_retrieval.DataPostProcessor") + @patch("core.rag.retrieval.dataset_retrieval.DatasetRetrieval._retriever") + @patch("core.rag.retrieval.dataset_retrieval.DatasetRetrieval.calculate_vector_score") + def test_multiple_retrieve_thread_single_dataset_uses_standard_scoring( + self, mock_calculate_vector_score, mock_retriever, mock_data_processor_class, mock_flask_app, mock_dataset + ): + """ + Test that _multiple_retrieve_thread uses standard scoring when dataset_count is 1 + and reranking is enabled. + + When there's only one dataset, instead of using DataPostProcessor, + the method should fall through to the standard scoring logic + (calculate_vector_score for high_quality datasets). + + Verifies: + - DataPostProcessor is NOT called + - calculate_vector_score IS called for high_quality indexing + - Documents are scored correctly + """ + # Arrange + dataset_retrieval = DatasetRetrieval() + tenant_id = str(uuid4()) + + # Create test documents + doc1 = _dataset_retrieval_Document( + page_content="Test content 1", + metadata={"doc_id": "doc1", "score": 0.9, "document_id": str(uuid4()), "dataset_id": mock_dataset.id}, + provider="dify", + ) + doc2 = _dataset_retrieval_Document( + page_content="Test content 2", + metadata={"doc_id": "doc2", "score": 0.8, "document_id": str(uuid4()), "dataset_id": mock_dataset.id}, + provider="dify", + ) + + # Mock _retriever to return documents + def side_effect_retriever( + flask_app, + session, + dataset_id, + query, + top_k, + all_documents, + document_ids_filter, + metadata_condition, + attachment_ids, + ): + all_documents.extend([doc1, doc2]) + + mock_retriever.side_effect = side_effect_retriever + + # Set up dataset with high_quality indexing + mock_dataset.indexing_technique = "high_quality" + + # Mock calculate_vector_score to return scored documents + scored_docs = [ + _dataset_retrieval_Document( + page_content="Test content 1", + metadata={ + "doc_id": "doc1", + "score": 0.95, + "document_id": str(uuid4()), + "dataset_id": mock_dataset.id, + }, + provider="dify", + ), + ] + mock_calculate_vector_score.return_value = scored_docs + + all_documents = [] + + # Act - Call with dataset_count = 1 + dataset_retrieval._multiple_retrieve_thread( + flask_app=mock_flask_app, + available_datasets=[mock_dataset], + metadata_condition=None, + metadata_filter_document_ids=None, + all_documents=all_documents, + tenant_id=tenant_id, + reranking_enable=True, # Reranking enabled but should be skipped for single dataset + reranking_mode="reranking_model", + reranking_model={"reranking_provider_name": "cohere", "reranking_model_name": "rerank-v2"}, + weights=None, + top_k=5, + score_threshold=0.5, + query="test query", + attachment_id=None, + dataset_count=1, + ) + + # Assert + # DataPostProcessor should NOT be called + mock_data_processor_class.assert_not_called() + + # calculate_vector_score SHOULD be called for high_quality datasets + mock_calculate_vector_score.assert_called_once() + call_args = mock_calculate_vector_score.call_args + assert call_args[0][1] == 5 # top_k + + # Documents should be added after standard scoring + assert len(all_documents) == 1 + assert all_documents[0].page_content == "Test content 1" + + class TestRetrievalMethods: + """ + Test suite for RetrievalMethod enum and utility methods. + + The RetrievalMethod enum defines the available search strategies: + + 1. **SEMANTIC_SEARCH**: Vector-based similarity search using embeddings + - Best for: Natural language queries, conceptual similarity + - Uses: Embedding models (e.g., text-embedding-ada-002) + - Example: "What is machine learning?" matches "AI and ML concepts" + + 2. **FULL_TEXT_SEARCH**: BM25-based text matching + - Best for: Exact phrase matching, keyword presence + - Uses: BM25 algorithm with sparse vectors + - Example: "Python programming" matches documents with those exact terms + + 3. **HYBRID_SEARCH**: Combination of semantic + full-text + - Best for: Comprehensive search with both conceptual and exact matching + - Uses: Both embedding vectors and BM25, with score merging + - Example: Finds both semantically similar and keyword-matching documents + + 4. **KEYWORD_SEARCH**: Traditional keyword-based search (economy mode) + - Best for: Simple, fast searches without embeddings + - Uses: Jieba tokenization and keyword matching + - Example: Basic text search without vector database + + Utility Methods: + ================ + - is_support_semantic_search(): Check if method uses embeddings + - is_support_fulltext_search(): Check if method uses BM25 + + These utilities help determine which search operations to execute + in the RetrievalService.retrieve() method. + """ + + def test_retrieval_method_values(self): + """ + Test that all retrieval method constants are defined correctly. + + This ensures the enum values match the expected string constants + used throughout the codebase for configuration and API calls. + + Verifies: + - All expected retrieval methods exist + - Values are correct strings (not accidentally changed) + - String values match database/config expectations + """ + assert RetrievalMethod.SEMANTIC_SEARCH == "semantic_search" + assert RetrievalMethod.FULL_TEXT_SEARCH == "full_text_search" + assert RetrievalMethod.HYBRID_SEARCH == "hybrid_search" + assert RetrievalMethod.KEYWORD_SEARCH == "keyword_search" + + def test_is_support_semantic_search(self): + """ + Test semantic search support detection. + + Verifies: + - Semantic search method is detected + - Hybrid search method is detected (includes semantic) + - Other methods are not detected + """ + assert RetrievalMethod.is_support_semantic_search(RetrievalMethod.SEMANTIC_SEARCH) is True + assert RetrievalMethod.is_support_semantic_search(RetrievalMethod.HYBRID_SEARCH) is True + assert RetrievalMethod.is_support_semantic_search(RetrievalMethod.FULL_TEXT_SEARCH) is False + assert RetrievalMethod.is_support_semantic_search(RetrievalMethod.KEYWORD_SEARCH) is False + + def test_is_support_fulltext_search(self): + """ + Test full-text search support detection. + + Verifies: + - Full-text search method is detected + - Hybrid search method is detected (includes full-text) + - Other methods are not detected + """ + assert RetrievalMethod.is_support_fulltext_search(RetrievalMethod.FULL_TEXT_SEARCH) is True + assert RetrievalMethod.is_support_fulltext_search(RetrievalMethod.HYBRID_SEARCH) is True + assert RetrievalMethod.is_support_fulltext_search(RetrievalMethod.SEMANTIC_SEARCH) is False + assert RetrievalMethod.is_support_fulltext_search(RetrievalMethod.KEYWORD_SEARCH) is False + + class TestDocumentModel: + """ + Test suite for Document model used in retrieval. + + The Document class is the core data structure for representing text chunks + in the retrieval system. It's based on Pydantic BaseModel for validation. + + Document Structure: + =================== + - **page_content** (str): The actual text content of the document chunk + - **metadata** (dict): Additional information about the document + - doc_id: Unique identifier for the chunk + - document_id: Parent document ID + - dataset_id: Dataset this document belongs to + - score: Relevance score from search (0.0 to 1.0) + - Custom fields: category, tags, timestamps, etc. + - **provider** (str): Source of the document ("dify" or "external") + - **vector** (list[float] | None): Embedding vector for semantic search + - **children** (list[ChildDocument] | None): Sub-chunks for hierarchical docs + + Document Lifecycle: + =================== + 1. **Creation**: Documents are created when text is indexed + - Content is chunked into manageable pieces + - Embeddings are generated for semantic search + - Metadata is attached for filtering and tracking + + 2. **Storage**: Documents are stored in vector databases + - Vector field stores embeddings + - Metadata enables filtering + - Provider tracks source (internal vs external) + + 3. **Retrieval**: Documents are returned from search operations + - Scores are added during search + - Multiple documents may be combined (hybrid search) + - Deduplication uses doc_id + + 4. **Post-processing**: Documents may be reranked or filtered + - Scores can be recalculated + - Content may be truncated or formatted + - Metadata is used for display + + Why Test the Document Model: + ============================ + - Ensures data structure integrity + - Validates Pydantic model behavior + - Confirms default values work correctly + - Tests equality comparison for deduplication + - Verifies metadata handling + + Related Classes: + ================ + - ChildDocument: For hierarchical document structures + - RetrievalSegments: Combines Document with database segment info + """ + + def test_document_creation_basic(self): + """ + Test basic Document object creation. + + Tests the minimal required fields and default values. + Only page_content is required; all other fields have defaults. + + Verifies: + - Document can be created with minimal fields + - Default values are set correctly + - Pydantic validation works + - No exceptions are raised + """ + doc = _dataset_retrieval_Document(page_content="Test content") + + assert doc.page_content == "Test content" + assert doc.metadata == {} # Empty dict by default + assert doc.provider == "dify" # Default provider + assert doc.vector is None # No embedding by default + assert doc.children is None # No child documents by default + + def test_document_creation_with_metadata(self): + """ + Test Document creation with metadata. + + Verifies: + - Metadata is stored correctly + - Metadata can contain various types + """ + metadata = { + "doc_id": "test_doc", + "score": 0.95, + "dataset_id": str(uuid4()), + "category": "test", + } + doc = _dataset_retrieval_Document(page_content="Test content", metadata=metadata) + + assert doc.metadata == metadata + assert doc.metadata["score"] == 0.95 + + def test_document_creation_with_vector(self): + """ + Test Document creation with embedding vector. + + Verifies: + - Vector embeddings can be stored + - Vector is optional + """ + vector = [0.1, 0.2, 0.3, 0.4, 0.5] + doc = _dataset_retrieval_Document(page_content="Test content", vector=vector) + + assert doc.vector == vector + assert doc.vector is not None + assert len(doc.vector) == 5 + + def test_document_with_external_provider(self): + """ + Test Document with external provider. + + Verifies: + - Provider can be set to external + - External documents are handled correctly + """ + doc = _dataset_retrieval_Document(page_content="External content", provider="external") + + assert doc.provider == "external" + + def test_document_equality(self): + """ + Test Document equality comparison. + + Verifies: + - Documents with same content are considered equal + - Metadata affects equality + """ + doc1 = _dataset_retrieval_Document(page_content="Content", metadata={"id": "1"}) + doc2 = _dataset_retrieval_Document(page_content="Content", metadata={"id": "1"}) + doc3 = _dataset_retrieval_Document(page_content="Different", metadata={"id": "1"}) + + assert doc1 == doc2 + assert doc1 != doc3 + + # ==================== Test _check_knowledge_rate_limit ==================== + + class TestCheckKnowledgeRateLimit: + """ + Test suite for _check_knowledge_rate_limit method. + + The _check_knowledge_rate_limit method validates whether a tenant has + exceeded their knowledge retrieval rate limit. This is important for: + - Preventing abuse of the knowledge retrieval system + - Enforcing subscription plan limits + - Tracking usage for billing purposes + + Test Cases: + ============ + 1. Rate limit disabled - no exception raised + 2. Rate limit enabled but not exceeded - no exception raised + 3. Rate limit enabled and exceeded - RateLimitExceededError raised + 4. Redis operations are performed correctly + 5. RateLimitLog is created when limit is exceeded + """ + + @patch("core.rag.retrieval.dataset_retrieval.FeatureService") + @patch("core.rag.retrieval.dataset_retrieval.redis_client") + def test_rate_limit_disabled_no_exception(self, mock_redis, mock_feature_service): + """ + Test that when rate limit is disabled, no exception is raised. + + This test verifies the behavior when the tenant's subscription + does not have rate limiting enabled. + + Verifies: + - FeatureService.get_knowledge_rate_limit is called + - No Redis operations are performed + - No exception is raised + - Retrieval proceeds normally + """ + # Arrange + tenant_id = str(uuid4()) + dataset_retrieval = DatasetRetrieval() + + # Mock rate limit disabled + mock_limit = Mock() + mock_limit.enabled = False + mock_feature_service.get_knowledge_rate_limit.return_value = mock_limit + + # Act & Assert - should not raise any exception + dataset_retrieval._check_knowledge_rate_limit(tenant_id) + + # Verify FeatureService was called + mock_feature_service.get_knowledge_rate_limit.assert_called_once_with(tenant_id) + + # Verify no Redis operations were performed + assert not mock_redis.zadd.called + assert not mock_redis.zremrangebyscore.called + assert not mock_redis.zcard.called + + @patch("core.rag.retrieval.dataset_retrieval.FeatureService") + @patch("core.rag.retrieval.dataset_retrieval.redis_client") + @patch("core.rag.retrieval.dataset_retrieval.time") + def test_rate_limit_enabled_not_exceeded(self, mock_time, mock_redis, mock_feature_service): + """ + Test that when rate limit is enabled but not exceeded, no exception is raised. + + This test simulates a tenant making requests within their rate limit. + The Redis sorted set stores timestamps of recent requests, and old + requests (older than 60 seconds) are removed. + + Verifies: + - Redis zadd is called to track the request + - Redis zremrangebyscore removes old entries + - Redis zcard returns count within limit + - No exception is raised + """ + # Arrange + tenant_id = str(uuid4()) + dataset_retrieval = DatasetRetrieval() + + # Mock rate limit enabled with limit of 100 requests per minute + mock_limit = Mock() + mock_limit.enabled = True + mock_limit.limit = 100 + mock_limit.subscription_plan = "professional" + mock_feature_service.get_knowledge_rate_limit.return_value = mock_limit + + # Mock time + current_time = 1234567890000 # Current time in milliseconds + mock_time.time.return_value = current_time / 1000 # Return seconds + mock_time.time.__mul__ = lambda self, x: int(self * x) # Multiply to get milliseconds + + # Mock Redis operations + # zcard returns 50 (within limit of 100) + mock_redis.zcard.return_value = 50 + + # Act & Assert - should not raise any exception + dataset_retrieval._check_knowledge_rate_limit(tenant_id) + + # Verify Redis operations + expected_key = f"rate_limit_{tenant_id}" + mock_redis.zadd.assert_called_once_with(expected_key, {current_time: current_time}) + mock_redis.zremrangebyscore.assert_called_once_with(expected_key, 0, current_time - 60000) + mock_redis.zcard.assert_called_once_with(expected_key) + + @patch("core.rag.retrieval.dataset_retrieval.FeatureService") + @patch("core.rag.retrieval.dataset_retrieval.redis_client") + @patch("core.rag.retrieval.dataset_retrieval.time") + def test_rate_limit_enabled_exceeded_raises_exception( + self, mock_time, mock_redis, mock_feature_service, sqlite_session: Session + ): + """ + Test that when rate limit is enabled and exceeded, RateLimitExceededError is raised. + + This test simulates a tenant exceeding their rate limit. When the count + of recent requests exceeds the limit, an exception should be raised and + a RateLimitLog should be created. + + Verifies: + - Redis zcard returns count exceeding limit + - RateLimitExceededError is raised with correct message + - RateLimitLog is created in database + - Session operations are performed correctly + """ + # Arrange + tenant_id = str(uuid4()) + dataset_retrieval = DatasetRetrieval() + + # Mock rate limit enabled with limit of 100 requests per minute + mock_limit = Mock() + mock_limit.enabled = True + mock_limit.limit = 100 + mock_limit.subscription_plan = "professional" + mock_feature_service.get_knowledge_rate_limit.return_value = mock_limit + + # Mock time + current_time = 1234567890000 + mock_time.time.return_value = current_time / 1000 + + # Mock Redis operations - return count exceeding limit + mock_redis.zcard.return_value = 150 # Exceeds limit of 100 + + # Act & Assert + with pytest.raises(exc.RateLimitExceededError) as exc_info: + dataset_retrieval._check_knowledge_rate_limit(tenant_id) + + # Verify exception message + assert "knowledge base request rate limit" in str(exc_info.value) + + # Verify the independently committed audit row from an observer session. + added_log = sqlite_session.scalar(select(RateLimitLog).where(RateLimitLog.tenant_id == tenant_id)) + assert added_log is not None + assert added_log.tenant_id == tenant_id + assert added_log.subscription_plan == "professional" + assert added_log.operation == "knowledge" + + # ==================== Test _get_available_datasets ==================== + + class TestGetAvailableDatasets: + """ + Test suite for _get_available_datasets method. + + The _get_available_datasets method retrieves datasets that are available + for retrieval. A dataset is considered available if: + - It belongs to the specified tenant + - It's in the list of requested dataset_ids + - It has at least one completed, enabled, non-archived document OR + - It's an external provider dataset + + Note: Due to SQLAlchemy subquery complexity, full testing is done in + integration tests. Unit tests here verify basic behavior. + """ + + def test_method_exists_and_has_correct_signature(self): + """ + Test that the method exists and has the correct signature. + + Verifies: + - Method exists on DatasetRetrieval class + - Accepts tenant_id and dataset_ids parameters + """ + # Arrange + dataset_retrieval = DatasetRetrieval() + + # Assert - method exists + assert hasattr(dataset_retrieval, "_get_available_datasets") + # Assert - method is callable + assert callable(dataset_retrieval._get_available_datasets) + + # ==================== Test knowledge_retrieval ==================== + + class TestDatasetRetrievalKnowledgeRetrieval: + """ + Test suite for knowledge_retrieval method. + + The knowledge_retrieval method is the main entry point for retrieving + knowledge from datasets. It orchestrates the entire retrieval process: + 1. Checks rate limits + 2. Gets available datasets + 3. Applies metadata filtering if enabled + 4. Performs retrieval (single or multiple mode) + 5. Formats and returns results + + Test Cases: + ============ + 1. Single mode retrieval + 2. Multiple mode retrieval + 3. Metadata filtering disabled + 4. Metadata filtering automatic + 5. Metadata filtering manual + 6. External documents handling + 7. Dify documents handling + 8. Empty results handling + 9. Rate limit exceeded + 10. No available datasets + """ + + def test_knowledge_retrieval_single_mode_basic(self): + """ + Test knowledge_retrieval in single retrieval mode - basic check. + + Note: Full single mode testing requires complex model mocking and + is better suited for integration tests. This test verifies the + method accepts single mode requests. + + Verifies: + - Method can accept single mode request + - Request parameters are correctly structured + """ + # Arrange + tenant_id = str(uuid4()) + user_id = str(uuid4()) + app_id = str(uuid4()) + dataset_id = str(uuid4()) + + request = KnowledgeRetrievalRequest( + tenant_id=tenant_id, + user_id=user_id, + app_id=app_id, + user_from="web", + dataset_ids=[dataset_id], + query="What is Python?", + retrieval_mode="single", + model_provider="openai", + model_name="gpt-4", + model_mode="chat", + completion_params={"temperature": 0.7}, + ) + + # Assert - request is properly structured + assert request.retrieval_mode == "single" + assert request.model_provider == "openai" + assert request.model_name == "gpt-4" + assert request.model_mode == "chat" + + def test_knowledge_retrieval_multiple_mode(self, sqlite_session: Session): + """ + Test knowledge_retrieval in multiple retrieval mode. + + In multiple mode, retrieval is performed across all datasets and + results are combined and reranked. + + Verifies: + - Rate limit is checked + - Available datasets are retrieved + - Multiple retrieval is performed + - Results are combined and reranked + - Results are formatted correctly + """ + # Arrange + tenant_id = str(uuid4()) + user_id = str(uuid4()) + app_id = str(uuid4()) + dataset_id1 = str(uuid4()) + dataset_id2 = str(uuid4()) + + request = KnowledgeRetrievalRequest( + tenant_id=tenant_id, + user_id=user_id, + app_id=app_id, + user_from="web", + dataset_ids=[dataset_id1, dataset_id2], + query="What is Python?", + retrieval_mode="multiple", + top_k=5, + score_threshold=0.7, + reranking_enable=True, + reranking_mode="reranking_model", + reranking_model={"reranking_provider_name": "cohere", "reranking_model_name": "rerank-v2"}, + ) + + dataset_retrieval = DatasetRetrieval() + + # Mock _check_knowledge_rate_limit + with patch.object(dataset_retrieval, "_check_knowledge_rate_limit"): + # Mock _get_available_datasets + mock_dataset1 = create_mock_dataset_methods(dataset_id=dataset_id1, tenant_id=tenant_id) + mock_dataset2 = create_mock_dataset_methods(dataset_id=dataset_id2, tenant_id=tenant_id) + with patch.object( + dataset_retrieval, "_get_available_datasets", return_value=[mock_dataset1, mock_dataset2] + ): + # Mock get_metadata_filter_condition + with patch.object(dataset_retrieval, "get_metadata_filter_condition", return_value=(None, None)): + # Mock multiple_retrieve to return documents + doc1 = create_mock_document_methods( + "Python is great", + "doc1", + score=0.9, + provider="external", + additional_metadata={"dataset_name": "test dataset", "title": "Python"}, + ) + doc2 = create_mock_document_methods( + "Python is awesome", + "doc2", + score=0.8, + provider="external", + additional_metadata={"dataset_name": "test dataset", "title": "Python 2"}, + ) + with patch.object( + dataset_retrieval, "multiple_retrieve", return_value=[doc1, doc2] + ) as mock_multiple_retrieve: + result = dataset_retrieval.knowledge_retrieval(sqlite_session, request) + + assert len(result) == 2 + mock_multiple_retrieve.assert_called_once() + + def test_knowledge_retrieval_metadata_filtering_disabled(self, sqlite_session: Session): + """ + Test knowledge_retrieval with metadata filtering disabled. + + When metadata filtering is disabled, get_metadata_filter_condition is + NOT called (the method checks metadata_filtering_mode != "disabled"). + + Verifies: + - get_metadata_filter_condition is NOT called when mode is "disabled" + - Retrieval proceeds without metadata filters + """ + # Arrange + tenant_id = str(uuid4()) + user_id = str(uuid4()) + app_id = str(uuid4()) + dataset_id = str(uuid4()) + + request = KnowledgeRetrievalRequest( + tenant_id=tenant_id, + user_id=user_id, + app_id=app_id, + user_from="web", + dataset_ids=[dataset_id], + query="What is Python?", + retrieval_mode="multiple", + metadata_filtering_mode="disabled", + top_k=5, + ) + + dataset_retrieval = DatasetRetrieval() + + # Mock dependencies + with patch.object(dataset_retrieval, "_check_knowledge_rate_limit"): + mock_dataset = create_mock_dataset_methods(dataset_id=dataset_id, tenant_id=tenant_id) + with patch.object(dataset_retrieval, "_get_available_datasets", return_value=[mock_dataset]): + # Mock get_metadata_filter_condition - should NOT be called when disabled + with patch.object( + dataset_retrieval, + "get_metadata_filter_condition", + return_value=(None, None), + ) as mock_get_metadata: + with patch.object(dataset_retrieval, "multiple_retrieve", return_value=[]): + # Act + result = dataset_retrieval.knowledge_retrieval(sqlite_session, request) + + # Assert + assert isinstance(result, list) + # get_metadata_filter_condition should NOT be called when mode is "disabled" + mock_get_metadata.assert_not_called() + + def test_knowledge_retrieval_with_external_documents(self, sqlite_session: Session): + """ + Test knowledge_retrieval with external documents. + + External documents come from external knowledge bases and should + be formatted differently than Dify documents. + + Verifies: + - External documents are handled correctly + - Provider is set to "external" + - Metadata includes external-specific fields + """ + # Arrange + tenant_id = str(uuid4()) + user_id = str(uuid4()) + app_id = str(uuid4()) + dataset_id = str(uuid4()) + + request = KnowledgeRetrievalRequest( + tenant_id=tenant_id, + user_id=user_id, + app_id=app_id, + user_from="web", + dataset_ids=[dataset_id], + query="What is Python?", + retrieval_mode="multiple", + top_k=5, + ) + + dataset_retrieval = DatasetRetrieval() + + # Mock dependencies + with patch.object(dataset_retrieval, "_check_knowledge_rate_limit"): + mock_dataset = create_mock_dataset_methods( + dataset_id=dataset_id, tenant_id=tenant_id, provider="external" + ) + with patch.object(dataset_retrieval, "_get_available_datasets", return_value=[mock_dataset]): + with patch.object(dataset_retrieval, "get_metadata_filter_condition", return_value=(None, None)): + # Create external document + external_doc = create_mock_document_methods( + "External knowledge", + "doc1", + score=0.9, + provider="external", + additional_metadata={ + "dataset_id": dataset_id, + "dataset_name": "external_kb", + "document_id": "ext_doc1", + "title": "External Document", + }, + ) + with patch.object(dataset_retrieval, "multiple_retrieve", return_value=[external_doc]): + # Act + result = dataset_retrieval.knowledge_retrieval(sqlite_session, request) + + # Assert + assert isinstance(result, list) + if result: + assert result[0].metadata.data_source_type == "external" + + def test_knowledge_retrieval_empty_results(self, sqlite_session: Session): + """ + Test knowledge_retrieval when no documents are found. + + Verifies: + - Empty list is returned + - No errors are raised + - All dependencies are still called + """ + # Arrange + tenant_id = str(uuid4()) + user_id = str(uuid4()) + app_id = str(uuid4()) + dataset_id = str(uuid4()) + + request = KnowledgeRetrievalRequest( + tenant_id=tenant_id, + user_id=user_id, + app_id=app_id, + user_from="web", + dataset_ids=[dataset_id], + query="What is Python?", + retrieval_mode="multiple", + top_k=5, + ) + + dataset_retrieval = DatasetRetrieval() + + # Mock dependencies + with patch.object(dataset_retrieval, "_check_knowledge_rate_limit"): + mock_dataset = create_mock_dataset_methods(dataset_id=dataset_id, tenant_id=tenant_id) + with patch.object(dataset_retrieval, "_get_available_datasets", return_value=[mock_dataset]): + with patch.object(dataset_retrieval, "get_metadata_filter_condition", return_value=(None, None)): + # Mock multiple_retrieve to return empty list + with patch.object(dataset_retrieval, "multiple_retrieve", return_value=[]): + # Act + result = dataset_retrieval.knowledge_retrieval(sqlite_session, request) + + # Assert + assert result == [] + + def test_knowledge_retrieval_rate_limit_exceeded(self, sqlite_session: Session): + """ + Test knowledge_retrieval when rate limit is exceeded. + + Verifies: + - RateLimitExceededError is raised + - No further processing occurs + """ + # Arrange + tenant_id = str(uuid4()) + user_id = str(uuid4()) + app_id = str(uuid4()) + dataset_id = str(uuid4()) + + request = KnowledgeRetrievalRequest( + tenant_id=tenant_id, + user_id=user_id, + app_id=app_id, + user_from="web", + dataset_ids=[dataset_id], + query="What is Python?", + retrieval_mode="multiple", + top_k=5, + ) + + dataset_retrieval = DatasetRetrieval() + + # Mock _check_knowledge_rate_limit to raise exception + with patch.object( + dataset_retrieval, + "_check_knowledge_rate_limit", + side_effect=exc.RateLimitExceededError("Rate limit exceeded"), + ): + # Act & Assert + with pytest.raises(exc.RateLimitExceededError): + dataset_retrieval.knowledge_retrieval(sqlite_session, request) + + def test_knowledge_retrieval_no_available_datasets(self, sqlite_session: Session): + """ + Test knowledge_retrieval when no datasets are available. + + Verifies: + - Empty list is returned + - No retrieval is attempted + """ + # Arrange + tenant_id = str(uuid4()) + user_id = str(uuid4()) + app_id = str(uuid4()) + dataset_id = str(uuid4()) + + request = KnowledgeRetrievalRequest( + tenant_id=tenant_id, + user_id=user_id, + app_id=app_id, + user_from="web", + dataset_ids=[dataset_id], + query="What is Python?", + retrieval_mode="multiple", + top_k=5, + ) + + dataset_retrieval = DatasetRetrieval() + + # Mock dependencies + with patch.object(dataset_retrieval, "_check_knowledge_rate_limit"): + # Mock _get_available_datasets to return empty list + with patch.object(dataset_retrieval, "_get_available_datasets", return_value=[]): + # Act + result = dataset_retrieval.knowledge_retrieval(sqlite_session, request) + + # Assert + assert result == [] + + def test_knowledge_retrieval_handles_multiple_documents_with_different_scores(self): + """ + Test that knowledge_retrieval processes multiple documents with different scores. + + Note: Full sorting and position testing requires complex SQLAlchemy mocking + which is better suited for integration tests. This test verifies documents + with different scores can be created and have their metadata. + + Verifies: + - Documents can be created with different scores + - Score metadata is properly set + """ + # Create documents with different scores + doc1 = create_mock_document_methods("Low score", "doc1", score=0.6) + doc2 = create_mock_document_methods("High score", "doc2", score=0.95) + doc3 = create_mock_document_methods("Medium score", "doc3", score=0.8) + + # Assert - each document has the correct score + assert doc1.metadata["score"] == 0.6 + assert doc2.metadata["score"] == 0.95 + assert doc3.metadata["score"] == 0.8 + + # Assert - documents are correctly sorted (not the retrieval result, just the list) + unsorted = [doc1, doc2, doc3] + sorted_docs = sorted(unsorted, key=lambda d: d.metadata["score"], reverse=True) + assert [d.metadata["score"] for d in sorted_docs] == [0.95, 0.8, 0.6] + + class TestProcessMetadataFilterFunc: + """ + Comprehensive test suite for process_metadata_filter_func method. + + This test class validates all metadata filtering conditions supported by + the DatasetRetrieval class, including string operations, numeric comparisons, + null checks, and list operations. + + Method Signature: + ================== + def process_metadata_filter_func( + self, sequence: int, condition: str, metadata_name: str, value: Any | None, filters: list + ) -> list: + + The method builds SQLAlchemy filter expressions by: + 1. Validating value is not None (except for empty/not empty conditions) + 2. Using DatasetDocument.doc_metadata JSON field operations + 3. Adding appropriate SQLAlchemy expressions to the filters list + 4. Returning the updated filters list + + Tests compile expressions against the real mapped JSON field for string, + numeric, null, and collection operators. + """ + + @pytest.fixture + def retrieval(self): + """ + Create a DatasetRetrieval instance for testing. + + Returns: + DatasetRetrieval: Instance to test process_metadata_filter_func + """ + return DatasetRetrieval() + + # ==================== String Condition Tests ==================== + + def test_contains_condition_string_value(self, retrieval): + """ + Test 'contains' condition with string value. + + Verifies: + - Filters list is populated with LIKE expression + - Pattern matching uses %value% syntax + """ + filters = [] + sequence = 0 + condition = "contains" + metadata_name = "author" + value = "John" + + result = retrieval.process_metadata_filter_func(sequence, condition, metadata_name, value, filters) + + assert result == filters + assert len(filters) == 1 + + def test_not_contains_condition(self, retrieval): + """ + Test 'not contains' condition. + + Verifies: + - Filters list is populated with NOT LIKE expression + - Pattern matching uses %value% syntax with negation + """ + filters = [] + sequence = 0 + condition = "not contains" + metadata_name = "title" + value = "banned" + + result = retrieval.process_metadata_filter_func(sequence, condition, metadata_name, value, filters) + + assert result == filters + assert len(filters) == 1 + + def test_start_with_condition(self, retrieval): + """ + Test 'start with' condition. + + Verifies: + - Filters list is populated with LIKE expression + - Pattern matching uses value% syntax + """ + filters = [] + sequence = 0 + condition = "start with" + metadata_name = "category" + value = "tech" + + result = retrieval.process_metadata_filter_func(sequence, condition, metadata_name, value, filters) + + assert result == filters + assert len(filters) == 1 + + def test_end_with_condition(self, retrieval): + """ + Test 'end with' condition. + + Verifies: + - Filters list is populated with LIKE expression + - Pattern matching uses %value syntax + """ + filters = [] + sequence = 0 + condition = "end with" + metadata_name = "filename" + value = ".pdf" + + result = retrieval.process_metadata_filter_func(sequence, condition, metadata_name, value, filters) + + assert result == filters + assert len(filters) == 1 + + # ==================== Equality Condition Tests ==================== + + def test_is_condition_with_string_value(self, retrieval): + """ + Test 'is' (=) condition with string value. + + Verifies: + - Filters list is populated with equality expression + - String comparison is used + """ + filters = [] + sequence = 0 + condition = "is" + metadata_name = "author" + value = "Jane Doe" + + result = retrieval.process_metadata_filter_func(sequence, condition, metadata_name, value, filters) + + assert result == filters + assert len(filters) == 1 + + def test_equals_condition_with_string_value(self, retrieval): + """ + Test '=' condition with string value. + + Verifies: + - Same behavior as 'is' condition + - String comparison is used + """ + filters = [] + sequence = 0 + condition = "=" + metadata_name = "category" + value = "technology" + + result = retrieval.process_metadata_filter_func(sequence, condition, metadata_name, value, filters) + + assert result == filters + assert len(filters) == 1 + + def test_is_condition_with_int_value(self, retrieval): + """ + Test 'is' condition with integer value. + + Verifies: + - Numeric comparison is used + - as_float() is called on the metadata field + """ + filters = [] + sequence = 0 + condition = "is" + metadata_name = "year" + value = 2023 + + result = retrieval.process_metadata_filter_func(sequence, condition, metadata_name, value, filters) + + assert result == filters + assert len(filters) == 1 + + def test_is_condition_with_float_value(self, retrieval): + """ + Test 'is' condition with float value. + + Verifies: + - Numeric comparison is used + - as_float() is called on the metadata field + """ + filters = [] + sequence = 0 + condition = "is" + metadata_name = "price" + value = 19.99 + + result = retrieval.process_metadata_filter_func(sequence, condition, metadata_name, value, filters) + + assert result == filters + assert len(filters) == 1 + + def test_is_not_condition_with_string_value(self, retrieval): + """ + Test 'is not' (≠) condition with string value. + + Verifies: + - Filters list is populated with inequality expression + - String comparison is used + """ + filters = [] + sequence = 0 + condition = "is not" + metadata_name = "author" + value = "Unknown" + + result = retrieval.process_metadata_filter_func(sequence, condition, metadata_name, value, filters) + + assert result == filters + assert len(filters) == 1 + + def test_not_equals_condition(self, retrieval): + """ + Test '≠' condition with string value. + + Verifies: + - Same behavior as 'is not' condition + - Inequality expression is used + """ + filters = [] + sequence = 0 + condition = "≠" + metadata_name = "category" + value = "archived" + + result = retrieval.process_metadata_filter_func(sequence, condition, metadata_name, value, filters) + + assert result == filters + assert len(filters) == 1 + + def test_is_not_condition_with_numeric_value(self, retrieval): + """ + Test 'is not' condition with numeric value. + + Verifies: + - Numeric inequality comparison is used + - as_float() is called on the metadata field + """ + filters = [] + sequence = 0 + condition = "is not" + metadata_name = "year" + value = 2000 + + result = retrieval.process_metadata_filter_func(sequence, condition, metadata_name, value, filters) + + assert result == filters + assert len(filters) == 1 + + # ==================== Null Condition Tests ==================== + + def test_empty_condition(self, retrieval): + """ + Test 'empty' condition (null check). + + Verifies: + - Filters list is populated with IS NULL expression + - Value can be None for this condition + """ + filters = [] + sequence = 0 + condition = "empty" + metadata_name = "author" + value = None + + result = retrieval.process_metadata_filter_func(sequence, condition, metadata_name, value, filters) + + assert result == filters + assert len(filters) == 1 + + def test_not_empty_condition(self, retrieval): + """ + Test 'not empty' condition (not null check). + + Verifies: + - Filters list is populated with IS NOT NULL expression + - Value can be None for this condition + """ + filters = [] + sequence = 0 + condition = "not empty" + metadata_name = "description" + value = None + + result = retrieval.process_metadata_filter_func(sequence, condition, metadata_name, value, filters) + + assert result == filters + assert len(filters) == 1 + + # ==================== Numeric Comparison Tests ==================== + + def test_before_condition(self, retrieval): + """ + Test 'before' (<) condition. + + Verifies: + - Filters list is populated with less than expression + - Numeric comparison is used + """ + filters = [] + sequence = 0 + condition = "before" + metadata_name = "year" + value = 2020 + + result = retrieval.process_metadata_filter_func(sequence, condition, metadata_name, value, filters) + + assert result == filters + assert len(filters) == 1 + + def test_less_than_condition(self, retrieval): + """ + Test '<' condition. + + Verifies: + - Same behavior as 'before' condition + - Less than expression is used + """ + filters = [] + sequence = 0 + condition = "<" + metadata_name = "price" + value = 100.0 + + result = retrieval.process_metadata_filter_func(sequence, condition, metadata_name, value, filters) + + assert result == filters + assert len(filters) == 1 + + def test_after_condition(self, retrieval): + """ + Test 'after' (>) condition. + + Verifies: + - Filters list is populated with greater than expression + - Numeric comparison is used + """ + filters = [] + sequence = 0 + condition = "after" + metadata_name = "year" + value = 2020 + + result = retrieval.process_metadata_filter_func(sequence, condition, metadata_name, value, filters) + + assert result == filters + assert len(filters) == 1 + + def test_greater_than_condition(self, retrieval): + """ + Test '>' condition. + + Verifies: + - Same behavior as 'after' condition + - Greater than expression is used + """ + filters = [] + sequence = 0 + condition = ">" + metadata_name = "rating" + value = 4.5 + + result = retrieval.process_metadata_filter_func(sequence, condition, metadata_name, value, filters) + + assert result == filters + assert len(filters) == 1 + + def test_less_than_or_equal_condition_unicode(self, retrieval): + """ + Test '≤' condition. + + Verifies: + - Filters list is populated with less than or equal expression + - Numeric comparison is used + """ + filters = [] + sequence = 0 + condition = "≤" + metadata_name = "price" + value = 50.0 + + result = retrieval.process_metadata_filter_func(sequence, condition, metadata_name, value, filters) + + assert result == filters + assert len(filters) == 1 + + def test_less_than_or_equal_condition_ascii(self, retrieval): + """ + Test '<=' condition. + + Verifies: + - Same behavior as '≤' condition + - Less than or equal expression is used + """ + filters = [] + sequence = 0 + condition = "<=" + metadata_name = "year" + value = 2023 + + result = retrieval.process_metadata_filter_func(sequence, condition, metadata_name, value, filters) + + assert result == filters + assert len(filters) == 1 + + def test_greater_than_or_equal_condition_unicode(self, retrieval): + """ + Test '≥' condition. + + Verifies: + - Filters list is populated with greater than or equal expression + - Numeric comparison is used + """ + filters = [] + sequence = 0 + condition = "≥" + metadata_name = "rating" + value = 3.5 + + result = retrieval.process_metadata_filter_func(sequence, condition, metadata_name, value, filters) + + assert result == filters + assert len(filters) == 1 + + def test_greater_than_or_equal_condition_ascii(self, retrieval): + """ + Test '>=' condition. + + Verifies: + - Same behavior as '≥' condition + - Greater than or equal expression is used + """ + filters = [] + sequence = 0 + condition = ">=" + metadata_name = "year" + value = 2000 + + result = retrieval.process_metadata_filter_func(sequence, condition, metadata_name, value, filters) + + assert result == filters + assert len(filters) == 1 + + # ==================== List/In Condition Tests ==================== + + def test_in_condition_with_comma_separated_string(self, retrieval): + """ + Test 'in' condition with comma-separated string value. + + Verifies: + - String is split into list + - Whitespace is trimmed from each value + - IN expression is created + """ + filters = [] + sequence = 0 + condition = "in" + metadata_name = "category" + value = "tech, science, AI " + + result = retrieval.process_metadata_filter_func(sequence, condition, metadata_name, value, filters) + + assert result == filters + assert len(filters) == 1 + + def test_in_condition_with_list_value(self, retrieval): + """ + Test 'in' condition with list value. + + Verifies: + - List is processed correctly + - None values are filtered out + - IN expression is created with valid values + """ + filters = [] + sequence = 0 + condition = "in" + metadata_name = "tags" + value = ["python", "javascript", None, "golang"] + + result = retrieval.process_metadata_filter_func(sequence, condition, metadata_name, value, filters) + + assert result == filters + assert len(filters) == 1 + + def test_in_condition_with_tuple_value(self, retrieval): + """ + Test 'in' condition with tuple value. + + Verifies: + - Tuple is processed like a list + - IN expression is created + """ + filters = [] + sequence = 0 + condition = "in" + metadata_name = "category" + value = ("tech", "science", "ai") + + result = retrieval.process_metadata_filter_func(sequence, condition, metadata_name, value, filters) + + assert result == filters + assert len(filters) == 1 + + def test_in_condition_with_empty_string(self, retrieval): + """ + Test 'in' condition with empty string value. + + Verifies: + - Empty string results in literal(False) filter + - No valid values to match + """ + filters = [] + sequence = 0 + condition = "in" + metadata_name = "category" + value = "" + + result = retrieval.process_metadata_filter_func(sequence, condition, metadata_name, value, filters) + + assert result == filters + assert len(filters) == 1 + # Verify it's a literal(False) expression + # This is a bit tricky to test without access to the actual expression + + def test_in_condition_with_only_whitespace(self, retrieval): + """ + Test 'in' condition with whitespace-only string value. + + Verifies: + - Whitespace-only string results in literal(False) filter + - All values are stripped and filtered out + """ + filters = [] + sequence = 0 + condition = "in" + metadata_name = "category" + value = " , , " + + result = retrieval.process_metadata_filter_func(sequence, condition, metadata_name, value, filters) + + assert result == filters + assert len(filters) == 1 + + def test_in_condition_with_single_string(self, retrieval): + """ + Test 'in' condition with single non-comma string. + + Verifies: + - Single string is treated as single-item list + - IN expression is created with one value + """ + filters = [] + sequence = 0 + condition = "in" + metadata_name = "category" + value = "technology" + + result = retrieval.process_metadata_filter_func(sequence, condition, metadata_name, value, filters) + + assert result == filters + assert len(filters) == 1 + + # ==================== Edge Case Tests ==================== + + def test_none_value_with_non_empty_condition(self, retrieval): + """ + Test None value with conditions that require value. + + Verifies: + - Original filters list is returned unchanged + - No filter is added for None values (except empty/not empty) + """ + filters = [] + sequence = 0 + condition = "contains" + metadata_name = "author" + value = None + + result = retrieval.process_metadata_filter_func(sequence, condition, metadata_name, value, filters) + + assert result == filters + assert len(filters) == 0 # No filter added + + def test_none_value_with_equals_condition(self, retrieval): + """ + Test None value with 'is' (=) condition. + + Verifies: + - Original filters list is returned unchanged + - No filter is added for None values + """ + filters = [] + sequence = 0 + condition = "is" + metadata_name = "author" + value = None + + result = retrieval.process_metadata_filter_func(sequence, condition, metadata_name, value, filters) + + assert result == filters + assert len(filters) == 0 + + def test_none_value_with_numeric_condition(self, retrieval): + """ + Test None value with numeric comparison condition. + + Verifies: + - Original filters list is returned unchanged + - No filter is added for None values + """ + filters = [] + sequence = 0 + condition = ">" + metadata_name = "year" + value = None + + result = retrieval.process_metadata_filter_func(sequence, condition, metadata_name, value, filters) + + assert result == filters + assert len(filters) == 0 + + def test_existing_filters_preserved(self, retrieval): + """ + Test that existing filters are preserved. + + Verifies: + - Existing filters in the list are not removed + - New filters are appended to the list + """ + existing_filter = MagicMock() + filters = [existing_filter] + sequence = 0 + condition = "contains" + metadata_name = "author" + value = "test" + + result = retrieval.process_metadata_filter_func(sequence, condition, metadata_name, value, filters) + + assert result == filters + assert len(filters) == 2 + assert filters[0] == existing_filter + + def test_multiple_filters_accumulated(self, retrieval): + """ + Test multiple calls to accumulate filters. + + Verifies: + - Each call adds a new filter to the list + - All filters are preserved across calls + """ + filters = [] + + # First filter + retrieval.process_metadata_filter_func(0, "contains", "author", "John", filters) + assert len(filters) == 1 + + # Second filter + retrieval.process_metadata_filter_func(1, ">", "year", 2020, filters) + assert len(filters) == 2 + + # Third filter + retrieval.process_metadata_filter_func(2, "is", "category", "tech", filters) + assert len(filters) == 3 + + def test_unknown_condition(self, retrieval): + """ + Test unknown/unsupported condition. + + Verifies: + - Original filters list is returned unchanged + - No filter is added for unknown conditions + """ + filters = [] + sequence = 0 + condition = "unknown_condition" + metadata_name = "author" + value = "test" + + result = retrieval.process_metadata_filter_func(sequence, condition, metadata_name, value, filters) + + assert result == filters + assert len(filters) == 0 + + def test_empty_string_value_with_contains(self, retrieval): + """ + Test empty string value with 'contains' condition. + + Verifies: + - Filter is added even with empty string + - LIKE expression is created + """ + filters = [] + sequence = 0 + condition = "contains" + metadata_name = "author" + value = "" + + result = retrieval.process_metadata_filter_func(sequence, condition, metadata_name, value, filters) + + assert result == filters + assert len(filters) == 1 + + def test_special_characters_in_value(self, retrieval): + """ + Test special characters in value string. + + Verifies: + - Special characters are handled in value + - LIKE expression is created correctly + """ + filters = [] + sequence = 0 + condition = "contains" + metadata_name = "title" + value = "C++ & Python's features" + + result = retrieval.process_metadata_filter_func(sequence, condition, metadata_name, value, filters) + + assert result == filters + assert len(filters) == 1 + + def test_zero_value_with_numeric_condition(self, retrieval): + """ + Test zero value with numeric comparison condition. + + Verifies: + - Zero is treated as valid value + - Numeric comparison is performed + """ + filters = [] + sequence = 0 + condition = ">" + metadata_name = "price" + value = 0 + + result = retrieval.process_metadata_filter_func(sequence, condition, metadata_name, value, filters) + + assert result == filters + assert len(filters) == 1 + + def test_negative_value_with_numeric_condition(self, retrieval): + """ + Test negative value with numeric comparison condition. + + Verifies: + - Negative numbers are handled correctly + - Numeric comparison is performed + """ + filters = [] + sequence = 0 + condition = "<" + metadata_name = "temperature" + value = -10.5 + + result = retrieval.process_metadata_filter_func(sequence, condition, metadata_name, value, filters) + + assert result == filters + assert len(filters) == 1 + + def test_float_value_with_integer_comparison(self, retrieval): + """ + Test float value with numeric comparison condition. + + Verifies: + - Float values work correctly + - Numeric comparison is performed + """ + filters = [] + sequence = 0 + condition = ">=" + metadata_name = "rating" + value = 4.5 + + result = retrieval.process_metadata_filter_func(sequence, condition, metadata_name, value, filters) + + assert result == filters + assert len(filters) == 1 + + class TestKnowledgeRetrievalRegression: + @pytest.fixture + def mock_dataset(self) -> Dataset: + dataset = Dataset( + id=str(uuid4()), + tenant_id=str(uuid4()), + name="test_dataset", + indexing_technique="high_quality", + provider="dify", + ) + return dataset + + def test_multiple_retrieve_reranking_with_app_context(self, mock_dataset): + """ + Repro test for current bug: + reranking runs after `with flask_app.app_context():` exits. + `_multiple_retrieve_thread` catches exceptions and stores them into `thread_exceptions`, + so we must assert from that list (not from an outer try/except). + """ + dataset_retrieval = DatasetRetrieval() + flask_app = Flask(__name__) + tenant_id = str(uuid4()) + + # second dataset to ensure dataset_count > 1 reranking branch + secondary_dataset = Dataset( + id=str(uuid4()), + provider="dify", + indexing_technique="high_quality", + ) + + # retriever returns 1 doc into internal list (all_documents_item) + document = _dataset_retrieval_Document( + page_content="Context aware doc", + metadata={ + "doc_id": "doc1", + "score": 0.95, + "document_id": str(uuid4()), + "dataset_id": mock_dataset.id, + }, + provider="dify", + ) + + def fake_retriever( + flask_app, + session, + dataset_id, + query, + top_k, + all_documents, + document_ids_filter, + metadata_condition, + attachment_ids, + ): + all_documents.append(document) + + called = {"init": 0, "invoke": 0} + + class ContextRequiredPostProcessor: + def __init__(self, *args, **kwargs): + called["init"] += 1 + # will raise RuntimeError if no Flask app context exists + _ = current_app.name + + def invoke(self, *args, **kwargs): + called["invoke"] += 1 + _ = current_app.name + return kwargs.get("documents") or args[1] + + # output list from _multiple_retrieve_thread + all_documents: list[_dataset_retrieval_Document] = [] + + # The safe wrapper collects exceptions raised after tracing. + thread_exceptions: list[Exception] = [] + + def target(): + with ( + patch.object(dataset_retrieval, "_retriever", side_effect=fake_retriever), + patch( + "core.rag.retrieval.dataset_retrieval.DataPostProcessor", + ContextRequiredPostProcessor, + ), + ): + dataset_retrieval._multiple_retrieve_thread_safely( + flask_app=flask_app, + available_datasets=[mock_dataset, secondary_dataset], + metadata_condition=None, + metadata_filter_document_ids=None, + all_documents=all_documents, + tenant_id=tenant_id, + reranking_enable=True, + reranking_mode="reranking_model", + reranking_model={ + "reranking_provider_name": "cohere", + "reranking_model_name": "rerank-v2", + }, + weights=None, + top_k=3, + score_threshold=0.0, + query="test query", + attachment_id=None, + dataset_count=2, # force reranking branch + thread_exceptions=thread_exceptions, + ) + + t = threading.Thread(target=target) + t.start() + t.join() + + # Ensure reranking branch was actually executed + assert called["init"] >= 1, "DataPostProcessor was never constructed; reranking branch may not have run." + + # Current buggy code should record an exception (not raise it) + assert not thread_exceptions, thread_exceptions + + def test_run_retriever_thread_provides_session_to_retriever(self): + dataset_retrieval = DatasetRetrieval() + all_documents: list[_dataset_retrieval_Document] = [] + + with patch.object(dataset_retrieval, "_retriever") as mock_retriever: + dataset_retrieval._run_retriever_thread( + flask_app=_FakeFlaskApp(), + dataset_id="dataset-1", + query="test query", + top_k=3, + all_documents=all_documents, + document_ids_filter=None, + metadata_condition=None, + attachment_ids=None, + ) + + mock_retriever.assert_called_once() + assert isinstance(mock_retriever.call_args.kwargs["session"], Session) + + def test_run_retriever_thread_safely_records_retriever_exception(self): + dataset_retrieval = DatasetRetrieval() + all_documents: list[_dataset_retrieval_Document] = [] + cancel_event = threading.Event() + thread_exceptions: list[Exception] = [] + expected_error = RuntimeError("retrieval failed") + + with patch.object(dataset_retrieval, "_retriever", side_effect=expected_error): + dataset_retrieval._run_retriever_thread_safely( + flask_app=_FakeFlaskApp(), + dataset_id="dataset-1", + query="test query", + top_k=3, + all_documents=all_documents, + document_ids_filter=None, + metadata_condition=None, + attachment_ids=None, + cancel_event=cancel_event, + thread_exceptions=thread_exceptions, + ) + + assert cancel_event.is_set() + assert thread_exceptions == [expected_error] + + def test_run_retriever_thread_safely_skips_failed_dataset_when_requested(self, caplog): + dataset_retrieval = DatasetRetrieval() + all_documents: list[_dataset_retrieval_Document] = [] + cancel_event = threading.Event() + thread_exceptions: list[Exception] = [] + expected_error = RuntimeError("retrieval failed") + + with patch.object(dataset_retrieval, "_retriever", side_effect=expected_error): + dataset_retrieval._run_retriever_thread_safely( + flask_app=_FakeFlaskApp(), + dataset_id="dataset-1", + query="test query", + top_k=3, + all_documents=all_documents, + document_ids_filter=None, + metadata_condition=None, + attachment_ids=None, + cancel_event=cancel_event, + thread_exceptions=thread_exceptions, + skip_on_error=True, + ) + + assert not cancel_event.is_set() + assert thread_exceptions == [] + assert "dataset_id=dataset-1" in caplog.text + assert "Skipping dataset retrieval because retriever failed" in caplog.text + + def test_multiple_retrieve_thread_skips_failed_dataset(self, mock_dataset, caplog): + dataset_retrieval = DatasetRetrieval() + flask_app = Flask(__name__) + successful_dataset = Dataset( + id=str(uuid4()), + provider="dify", + indexing_technique="high_quality", + ) + document = _dataset_retrieval_Document( + page_content="successful doc", + metadata={ + "doc_id": "doc1", + "score": 0.95, + "document_id": str(uuid4()), + "dataset_id": successful_dataset.id, + }, + provider="dify", + ) + + def fake_retriever( + flask_app, + session, + dataset_id, + query, + top_k, + all_documents, + document_ids_filter, + metadata_condition, + attachment_ids, + ): + if dataset_id == mock_dataset.id: + raise RuntimeError("dataset unavailable") + all_documents.append(document) + + all_documents: list[_dataset_retrieval_Document] = [] + + with patch.object(dataset_retrieval, "_retriever", side_effect=fake_retriever): + dataset_retrieval._multiple_retrieve_thread( + flask_app=flask_app, + available_datasets=[mock_dataset, successful_dataset], + metadata_condition=None, + metadata_filter_document_ids=None, + all_documents=all_documents, + tenant_id=str(uuid4()), + reranking_enable=False, + reranking_mode="reranking_model", + reranking_model=None, + weights=None, + top_k=3, + score_threshold=0.0, + query="test query", + attachment_id=None, + dataset_count=2, + ) + + assert all_documents == [document] + assert f"dataset_id={mock_dataset.id}" in caplog.text + assert "Skipping dataset retrieval because retriever failed" in caplog.text + + class TestDatasetRetrievalAdditionalHelpers: + @pytest.fixture + def retrieval(self) -> DatasetRetrieval: + return DatasetRetrieval() + + def test_llm_usage_and_record_usage(self, retrieval: DatasetRetrieval) -> None: + empty_usage = retrieval.llm_usage + assert empty_usage.total_tokens == 0 + + retrieval._record_usage(None) + assert retrieval.llm_usage.total_tokens == 0 + + usage_1 = LLMUsage.from_metadata({"prompt_tokens": 2, "completion_tokens": 3, "total_tokens": 5}) + usage_2 = LLMUsage.from_metadata({"prompt_tokens": 4, "completion_tokens": 1, "total_tokens": 5}) + retrieval._record_usage(usage_1) + retrieval._record_usage(usage_2) + assert retrieval.llm_usage.total_tokens == 10 + + def test_replace_metadata_filter_value(self, retrieval: DatasetRetrieval) -> None: + assert retrieval._replace_metadata_filter_value("plain", {}) == "plain" + replaced = retrieval._replace_metadata_filter_value( + "hello {{name}}\n\t{{missing}}", + {"name": "world"}, + ) + assert replaced == "hello world {{missing}}" + + def test_process_metadata_filter_in_with_scalar_fallback(self) -> None: + filters: list = [] + result = DatasetRetrieval.process_metadata_filter_func( + sequence=0, + condition="in", + metadata_name="category", + value=123, + filters=filters, + ) + assert result is filters + assert len(filters) == 1 + + def test_calculate_vector_score(self, retrieval: DatasetRetrieval) -> None: + doc_high = _dataset_retrieval_Document(page_content="a", metadata={"score": 0.9}, provider="dify") + doc_low = _dataset_retrieval_Document(page_content="b", metadata={"score": 0.2}, provider="dify") + doc_no_meta = _dataset_retrieval_Document(page_content="c", metadata={}, provider="dify") + + filtered = retrieval.calculate_vector_score([doc_low, doc_high, doc_no_meta], top_k=1, score_threshold=0.5) + assert len(filtered) == 1 + assert filtered[0].metadata["score"] == 0.9 + + assert retrieval.calculate_vector_score([doc_low], top_k=2, score_threshold=1.0) == [] + + def test_calculate_keyword_score(self, retrieval: DatasetRetrieval) -> None: + documents = [ + _dataset_retrieval_Document(page_content="python language", metadata={"doc_id": "1"}, provider="dify"), + _dataset_retrieval_Document(page_content="java language", metadata={"doc_id": "2"}, provider="dify"), + ] + keyword_handler = Mock() + keyword_handler.extract_keywords.side_effect = [ + ["python", "language"], + ["python", "language"], + ["java", "language"], + ] + + with patch("core.rag.retrieval.dataset_retrieval.JiebaKeywordTableHandler", return_value=keyword_handler): + ranked = retrieval.calculate_keyword_score("python language", documents, top_k=1) + + assert len(ranked) == 1 + assert "keywords" in ranked[0].metadata + assert ranked[0].metadata["doc_id"] == "1" + + def test_send_trace_task(self, retrieval: DatasetRetrieval) -> None: + trace_manager = Mock() + retrieval.application_generate_entity = SimpleNamespace(trace_manager=trace_manager) + docs = [_dataset_retrieval_Document(page_content="d", metadata={}, provider="dify")] + + retrieval._send_trace_task("m1", docs, {"cost": 1}) + trace_manager.add_trace_task.assert_called_once() + + retrieval.application_generate_entity = None + trace_manager.reset_mock() + retrieval._send_trace_task("m1", docs, {"cost": 1}) + trace_manager.add_trace_task.assert_not_called() + + def test_on_query(self, retrieval: DatasetRetrieval, sqlite_engine: Engine, sqlite_session: Session) -> None: + dataset_ids = [str(uuid4()), str(uuid4())] + app_id = str(uuid4()) + user_id = str(uuid4()) + + with patch("core.rag.retrieval.dataset_retrieval.db", SimpleNamespace(engine=sqlite_engine)): + retrieval._on_query( + query=None, + attachment_ids=None, + dataset_ids=[dataset_ids[0]], + app_id=app_id, + user_from="account", + user_id=user_id, + ) + assert sqlite_session.scalars(select(DatasetQuery)).all() == [] + + retrieval._on_query( + query="python", + attachment_ids=[str(uuid4())], + dataset_ids=dataset_ids, + app_id=app_id, + user_from="account", + user_id=user_id, + ) + added_queries = sqlite_session.scalars(select(DatasetQuery).order_by(DatasetQuery.dataset_id)).all() + assert [row.dataset_id for row in added_queries] == sorted(dataset_ids) + assert all(row.created_by == user_id for row in added_queries) + + def test_on_query_normalizes_workflow_end_user_role( + self, retrieval: DatasetRetrieval, sqlite_engine: Engine, sqlite_session: Session + ) -> None: + dataset_id = str(uuid4()) + with patch("core.rag.retrieval.dataset_retrieval.db", SimpleNamespace(engine=sqlite_engine)): + retrieval._on_query( + query="python", + attachment_ids=None, + dataset_ids=[dataset_id], + app_id=str(uuid4()), + user_from="end-user", + user_id=str(uuid4()), + ) + + added_query = sqlite_session.scalar(select(DatasetQuery).where(DatasetQuery.dataset_id == dataset_id)) + assert added_query is not None + assert added_query.created_by_role == CreatorUserRole.END_USER + + def test_handle_invoke_result(self, retrieval: DatasetRetrieval) -> None: + usage = LLMUsage.empty_usage() + chunk_1 = SimpleNamespace( + model="m1", + prompt_messages=[Mock()], + delta=SimpleNamespace(message=SimpleNamespace(content="hello "), usage=usage), + ) + chunk_2 = SimpleNamespace( + model="m1", + prompt_messages=[Mock()], + delta=SimpleNamespace( + message=SimpleNamespace(content=[SimpleNamespace(data="world")]), + usage=None, + ), + ) + + def _chunks() -> Generator[Any]: + yield chunk_1 + yield chunk_2 + + text, returned_usage = retrieval._handle_invoke_result(_chunks()) + assert text == "hello world" + assert returned_usage == usage + + def _empty_chunks() -> Generator[Any]: + yield from () + + text_empty, usage_empty = retrieval._handle_invoke_result(_empty_chunks()) + assert text_empty == "" + assert usage_empty == LLMUsage.empty_usage() + + def test_get_prompt_template(self, retrieval: DatasetRetrieval) -> None: + model_config_chat = ModelConfigWithCredentialsEntity.model_construct( + provider="openai", + model="gpt", + model_schema=Mock(), + mode="chat", + provider_model_bundle=Mock(), + credentials={}, + parameters={}, + stop=["x"], + ) + model_config_completion = ModelConfigWithCredentialsEntity.model_construct( + provider="openai", + model="gpt", + model_schema=Mock(), + mode="completion", + provider_model_bundle=Mock(), + credentials={}, + parameters={}, + stop=[], + ) + + with patch("core.rag.retrieval.dataset_retrieval.AdvancedPromptTransform") as mock_prompt_transform: + mock_prompt_transform.return_value.get_prompt.return_value = ["prompt"] + prompt_messages, stop = retrieval._get_prompt_template( + model_config=model_config_chat, + mode="chat", + metadata_fields=["author"], + query="python", + ) + assert prompt_messages == ["prompt"] + assert stop == ["x"] + + with patch( + "core.rag.retrieval.dataset_retrieval.METADATA_FILTER_COMPLETION_PROMPT", + "{input_text} {metadata_fields}", + ): + prompt_messages_completion, stop_completion = retrieval._get_prompt_template( + model_config=model_config_completion, + mode="completion", + metadata_fields=["author"], + query="python", + ) + assert prompt_messages_completion == ["prompt"] + assert stop_completion == [] + + with pytest.raises(ValueError, match="'unknown\\-mode' is not a valid ModelMode"): + retrieval._get_prompt_template( + model_config=model_config_chat, + mode="unknown-mode", + metadata_fields=[], + query="python", + ) + + def test_fetch_model_config_validation_and_success(self, retrieval: DatasetRetrieval) -> None: + with pytest.raises(ValueError, match="single_retrieval_config is required"): + retrieval._fetch_model_config("tenant-1", None) # type: ignore[arg-type] + + model_cfg = AppModelConfig(provider="openai", name="gpt", mode="chat", completion_params={"stop": ["END"]}) + model_instance = Mock() + model_instance.credentials = {"k": "v"} + model_instance.provider_model_bundle = Mock() + model_instance.model_type_instance = Mock() + model_instance.model_type_instance.get_model_schema.return_value = Mock() + + with ( + patch("core.rag.retrieval.dataset_retrieval.ModelManager.for_tenant") as mock_manager, + patch("core.rag.retrieval.dataset_retrieval.ModelConfigWithCredentialsEntity") as mock_cfg_entity, + ): + mock_manager.return_value.get_model_instance.return_value = model_instance + mock_cfg_entity.return_value = SimpleNamespace( + provider="openai", + model="gpt", + stop=["END"], + parameters={"temperature": 0.1}, + ) + + model_instance.provider_model_bundle.configuration.get_provider_model.return_value = None + with pytest.raises(ValueError, match="not exist"): + retrieval._fetch_model_config("tenant-1", model_cfg) + + provider_model = SimpleNamespace(status=ModelStatus.NO_CONFIGURE) + model_instance.provider_model_bundle.configuration.get_provider_model.return_value = provider_model + with pytest.raises(ValueError, match="credentials is not initialized"): + retrieval._fetch_model_config("tenant-1", model_cfg) + + provider_model.status = ModelStatus.NO_PERMISSION + with pytest.raises(ValueError, match="currently not support"): + retrieval._fetch_model_config("tenant-1", model_cfg) + + provider_model.status = ModelStatus.QUOTA_EXCEEDED + with pytest.raises(ValueError, match="quota exceeded"): + retrieval._fetch_model_config("tenant-1", model_cfg) + + provider_model.status = ModelStatus.ACTIVE + bad_mode_cfg = AppModelConfig(provider="openai", name="gpt", mode="chat") + bad_mode_cfg.mode = None # type: ignore[assignment] + with pytest.raises(ValueError, match="LLM mode is required"): + retrieval._fetch_model_config("tenant-1", bad_mode_cfg) + + model_instance.model_type_instance.get_model_schema.return_value = None + with pytest.raises(ValueError, match="not exist"): + retrieval._fetch_model_config("tenant-1", model_cfg) + + model_instance.model_type_instance.get_model_schema.return_value = Mock() + model_cfg_success = AppModelConfig( + provider="openai", + name="gpt", + mode="chat", + completion_params={"temperature": 0.1, "stop": ["END"]}, + ) + _, config = retrieval._fetch_model_config("tenant-1", model_cfg_success) + assert config.provider == "openai" + assert config.model == "gpt" + assert config.stop == ["END"] + assert "stop" not in config.parameters + + def test_automatic_metadata_filter_func(self, retrieval: DatasetRetrieval, sqlite_session: Session) -> None: + tenant_id = str(uuid4()) + dataset_id = str(uuid4()) + user_id = str(uuid4()) + sqlite_session.add( + DatasetMetadata( + tenant_id=tenant_id, + dataset_id=dataset_id, + type="string", + name="author", + created_by=user_id, + ) + ) + sqlite_session.commit() + model_instance = Mock() + model_instance.invoke_llm.return_value = iter([Mock()]) + model_config = ModelConfigWithCredentialsEntity.model_construct( + provider="openai", + model="gpt", + model_schema=Mock(), + mode="chat", + provider_model_bundle=Mock(), + credentials={}, + parameters={}, + stop=[], + ) + usage = LLMUsage.from_metadata({"prompt_tokens": 1, "completion_tokens": 1, "total_tokens": 2}) + with ( + patch.object(retrieval, "_fetch_model_config", return_value=(model_instance, model_config)), + patch.object(retrieval, "_get_prompt_template", return_value=(["prompt"], [])), + patch.object(retrieval, "_handle_invoke_result", return_value=('{"metadata_map":[]}', usage)), + patch("core.rag.retrieval.dataset_retrieval.parse_and_check_json_markdown") as mock_parse, + patch.object(retrieval, "_record_usage") as mock_record_usage, + ): + mock_parse.return_value = { + "metadata_map": [ + { + "metadata_field_name": "author", + "metadata_field_value": "Alice", + "comparison_operator": "contains", + }, + { + "metadata_field_name": "ignored", + "metadata_field_value": "value", + "comparison_operator": "contains", + }, + ] + } + result = retrieval._automatic_metadata_filter_func( + sqlite_session, + dataset_ids=[dataset_id], + query="python", + tenant_id=tenant_id, + user_id=user_id, + metadata_model_config=AppModelConfig(provider="openai", name="gpt", mode="chat"), + ) + + assert result == [{"metadata_name": "author", "value": "Alice", "condition": "contains"}] + mock_record_usage.assert_called_once_with(usage) + + with ( + patch.object(retrieval, "_fetch_model_config", side_effect=RuntimeError("boom")), + ): + with pytest.raises(RuntimeError, match="boom"): + retrieval._automatic_metadata_filter_func( + sqlite_session, + dataset_ids=[dataset_id], + query="python", + tenant_id=tenant_id, + user_id=user_id, + metadata_model_config=AppModelConfig(provider="openai", name="gpt", mode="chat"), + ) + + def test_get_metadata_filter_condition(self, retrieval: DatasetRetrieval, sqlite_session: Session) -> None: + tenant_id = str(uuid4()) + dataset_id = str(uuid4()) + document_id = str(uuid4()) + user_id = str(uuid4()) + sqlite_session.add( + DatasetDocument( + id=document_id, + tenant_id=tenant_id, + dataset_id=dataset_id, + position=1, + data_source_type=DataSourceType("upload_file"), + batch="batch-1", + name="document", + created_from=DocumentCreatedFrom("api"), + created_by=user_id, + indexing_status="completed", + enabled=True, + archived=False, + doc_metadata={"author": "Alice"}, + ) + ) + sqlite_session.commit() + + mapping, condition = retrieval.get_metadata_filter_condition( + sqlite_session, + dataset_ids=[dataset_id], + query="python", + tenant_id="tenant-1", + user_id="u1", + metadata_filtering_mode="disabled", + metadata_model_config=AppModelConfig(provider="openai", name="gpt", mode="chat"), + metadata_filtering_conditions=None, + inputs={}, + ) + assert mapping is None + assert condition is None + + automatic_filters = [{"condition": "contains", "metadata_name": "author", "value": "Alice"}] + with ( + patch.object(retrieval, "_automatic_metadata_filter_func", return_value=automatic_filters), + ): + mapping, condition = retrieval.get_metadata_filter_condition( + sqlite_session, + dataset_ids=[dataset_id], + query="python", + tenant_id="tenant-1", + user_id="u1", + metadata_filtering_mode="automatic", + metadata_model_config=AppModelConfig(provider="openai", name="gpt", mode="chat"), + metadata_filtering_conditions=AppMetadataFilteringCondition(logical_operator="or", conditions=[]), + inputs={}, + ) + assert mapping == {dataset_id: [document_id]} + assert condition is not None + assert condition.logical_operator == "or" + + manual_conditions = AppMetadataFilteringCondition( + logical_operator="and", + conditions=[AppCondition(name="author", comparison_operator="contains", value="{{name}}")], + ) + mapping, condition = retrieval.get_metadata_filter_condition( + sqlite_session, + dataset_ids=[dataset_id], + query="python", + tenant_id="tenant-1", + user_id="u1", + metadata_filtering_mode="manual", + metadata_model_config=AppModelConfig(provider="openai", name="gpt", mode="chat"), + metadata_filtering_conditions=manual_conditions, + inputs={"name": "Alice"}, + ) + assert mapping == {dataset_id: [document_id]} + assert condition is not None + assert condition.conditions + first_condition = condition.conditions[0] + assert first_condition.value == "Alice" + + with pytest.raises(ValueError, match="Invalid metadata filtering mode"): + retrieval.get_metadata_filter_condition( + sqlite_session, + dataset_ids=[dataset_id], + query="python", + tenant_id="tenant-1", + user_id="u1", + metadata_filtering_mode="unsupported", + metadata_model_config=AppModelConfig(provider="openai", name="gpt", mode="chat"), + metadata_filtering_conditions=None, + inputs={}, + ) + + def test_get_available_datasets(self, retrieval: DatasetRetrieval, sqlite_session: Session) -> None: + tenant_id = str(uuid4()) + other_tenant_id = str(uuid4()) + creator_id = str(uuid4()) + internal_id = str(uuid4()) + external_id = str(uuid4()) + unavailable_id = str(uuid4()) + decoy_id = str(uuid4()) + datasets = [ + Dataset( + id=internal_id, + tenant_id=tenant_id, + name="internal", + provider="dify", + indexing_technique="high_quality", + created_by=creator_id, + ), + Dataset( + id=external_id, + tenant_id=tenant_id, + name="external", + provider="external", + indexing_technique="high_quality", + created_by=creator_id, + ), + Dataset( + id=unavailable_id, + tenant_id=tenant_id, + name="unavailable", + provider="dify", + indexing_technique="high_quality", + created_by=creator_id, + ), + Dataset( + id=decoy_id, + tenant_id=other_tenant_id, + name="decoy", + provider="external", + indexing_technique="high_quality", + created_by=creator_id, + ), + ] + sqlite_session.add_all(datasets) + sqlite_session.add( + DatasetDocument( + tenant_id=tenant_id, + dataset_id=internal_id, + position=1, + data_source_type=DataSourceType("upload_file"), + batch="batch-1", + name="available document", + created_from=DocumentCreatedFrom("api"), + created_by=creator_id, + indexing_status="completed", + enabled=True, + archived=False, + ) + ) + sqlite_session.commit() + + available = retrieval._get_available_datasets( + tenant_id, [internal_id, external_id, unavailable_id, decoy_id] + ) + + assert {dataset.id for dataset in available} == {internal_id, external_id} + + def test_check_knowledge_rate_limit(self, retrieval: DatasetRetrieval, sqlite_session: Session) -> None: + with ( + patch("core.rag.retrieval.dataset_retrieval.FeatureService.get_knowledge_rate_limit") as mock_limit, + patch("core.rag.retrieval.dataset_retrieval.redis_client") as mock_redis, + patch("core.rag.retrieval.dataset_retrieval.time.time", return_value=100.0), + ): + mock_limit.return_value = SimpleNamespace(enabled=True, limit=2, subscription_plan="pro") + mock_redis.zcard.return_value = 1 + retrieval._check_knowledge_rate_limit("tenant-1") + mock_redis.zadd.assert_called_once() + + tenant_id = str(uuid4()) + with ( + patch("core.rag.retrieval.dataset_retrieval.FeatureService.get_knowledge_rate_limit") as mock_limit, + patch("core.rag.retrieval.dataset_retrieval.redis_client") as mock_redis, + patch("core.rag.retrieval.dataset_retrieval.time.time", return_value=100.0), + ): + mock_limit.return_value = SimpleNamespace(enabled=True, limit=1, subscription_plan="pro") + mock_redis.zcard.return_value = 2 + with pytest.raises(exc.RateLimitExceededError): + retrieval._check_knowledge_rate_limit(tenant_id) + rate_limit_log = sqlite_session.scalar(select(RateLimitLog).where(RateLimitLog.tenant_id == tenant_id)) + assert rate_limit_log is not None + + with patch("core.rag.retrieval.dataset_retrieval.FeatureService.get_knowledge_rate_limit") as mock_limit: + mock_limit.return_value = SimpleNamespace(enabled=False) + retrieval._check_knowledge_rate_limit("tenant-1") + + class TestKnowledgeRetrievalCoverage: + @pytest.fixture(autouse=True) + def _orm_session(self, sqlite_session: Session) -> None: + self.orm_session = sqlite_session + + @pytest.fixture + def retrieval(self) -> DatasetRetrieval: + return DatasetRetrieval() + + def test_returns_empty_when_query_missing(self, retrieval: DatasetRetrieval) -> None: + request = KnowledgeRetrievalRequest( + tenant_id="tenant-1", + user_id="user-1", + app_id="app-1", + user_from="workflow", + dataset_ids=["d1"], + query=None, + retrieval_mode="multiple", + ) + with ( + patch.object(retrieval, "_check_knowledge_rate_limit"), + patch.object(retrieval, "_get_available_datasets", return_value=[_dataset(id="d1")]), + ): + assert retrieval.knowledge_retrieval(self.orm_session, request) == [] + + def test_raises_when_metadata_model_config_missing(self, retrieval: DatasetRetrieval) -> None: + request = KnowledgeRetrievalRequest( + tenant_id="tenant-1", + user_id="user-1", + app_id="app-1", + user_from="workflow", + dataset_ids=["d1"], + query="query", + retrieval_mode="multiple", + metadata_filtering_mode="automatic", + metadata_model_config=None, + ) + with ( + patch.object(retrieval, "_check_knowledge_rate_limit"), + patch.object(retrieval, "_get_available_datasets", return_value=[_dataset(id="d1")]), + ): + with pytest.raises(ValueError, match="metadata_model_config is required"): + retrieval.knowledge_retrieval(self.orm_session, request) + + @pytest.mark.parametrize( + ("status", "error_cls"), + [ + (ModelStatus.NO_CONFIGURE, "ModelCredentialsNotInitializedError"), + (ModelStatus.NO_PERMISSION, "ModelNotSupportedError"), + (ModelStatus.QUOTA_EXCEEDED, "ModelQuotaExceededError"), + ], ) + def test_single_mode_raises_for_model_status( + self, + retrieval: DatasetRetrieval, + status: ModelStatus, + error_cls: str, + ) -> None: + request = KnowledgeRetrievalRequest( + tenant_id="tenant-1", + user_id="user-1", + app_id="app-1", + user_from="workflow", + dataset_ids=["dataset-1"], + query="python", + retrieval_mode="single", + model_provider="openai", + model_name="gpt-4", + ) + provider_model_bundle = Mock() + provider_model_bundle.configuration.get_provider_model.return_value = SimpleNamespace(status=status) + model_type_instance = Mock() + model_type_instance.get_model_schema.return_value = Mock() + model_instance = SimpleNamespace( + provider_model_bundle=provider_model_bundle, + model_type_instance=model_type_instance, + credentials={}, + ) + with ( + patch.object(retrieval, "_check_knowledge_rate_limit"), + patch.object(retrieval, "_get_available_datasets", return_value=[_dataset(id="dataset-1")]), + patch("core.rag.retrieval.dataset_retrieval.ModelManager.for_tenant") as mock_model_manager, + ): + mock_model_manager.return_value.get_model_instance.return_value = model_instance + with pytest.raises( + Exception, + match=( + "(?:Dify Hosted OpenAI gpt\\-4 currently not support\\.|Model gpt\\-4 " + "credentials is not initialized\\.|Model provider openai quota " + "exceeded\\.)" + ), + ) as exc_info: + retrieval.knowledge_retrieval(self.orm_session, request) + mock_model_manager.assert_called_once_with(tenant_id="tenant-1", user_id="user-1") + assert error_cls in type(exc_info.value).__name__ - def test_returns_none_when_dataset_ids_empty(self, retrieval: DatasetRetrieval) -> None: - config = DatasetEntity( - dataset_ids=[], - retrieve_config=DatasetRetrieveConfigEntity( - retrieve_strategy=DatasetRetrieveConfigEntity.RetrieveStrategy.MULTIPLE, - ), - ) - result = retrieval.retrieve( - self.orm_session, - app_id="app-1", - user_id="user-1", - tenant_id="tenant-1", - model_config=self._build_model_config(), - config=config, - query="python", - invoke_from=InvokeFrom.WEB_APP, - show_retrieve_source=False, - hit_callback=Mock(), - message_id="m1", - ) - assert result == (None, []) + class TestRetrieveCoverage: + @pytest.fixture(autouse=True) + def _orm_session(self, sqlite_session: Session) -> None: + self.orm_session = sqlite_session - def test_returns_none_when_model_schema_missing(self, retrieval: DatasetRetrieval) -> None: - config = DatasetEntity( - dataset_ids=["d1"], - retrieve_config=DatasetRetrieveConfigEntity( - retrieve_strategy=DatasetRetrieveConfigEntity.RetrieveStrategy.MULTIPLE, - ), - ) - model_config = self._build_model_config() - model_instance = Mock() - model_instance.model_name = "gpt-4" - model_instance.credentials = {"api_key": "secret"} - model_instance.provider_model_bundle = Mock() - model_instance.model_type_instance.get_model_schema.return_value = None - with patch("core.rag.retrieval.dataset_retrieval.ModelManager.for_tenant") as mock_model_manager: - mock_model_manager.return_value.get_model_instance.return_value = model_instance + @pytest.fixture + def retrieval(self) -> DatasetRetrieval: + return DatasetRetrieval() + + def _build_model_config(self, features: list[ModelFeature] | None = None): + model_type_instance = Mock() + model_type_instance.get_model_schema.return_value = SimpleNamespace(features=features or []) + provider_bundle = SimpleNamespace(model_type_instance=model_type_instance) + return ModelConfigWithCredentialsEntity.model_construct( + provider="openai", + model="gpt-4", + model_schema=Mock(), + mode="chat", + provider_model_bundle=provider_bundle, + credentials={}, + parameters={}, + stop=[], + ) + + def test_returns_none_when_dataset_ids_empty(self, retrieval: DatasetRetrieval) -> None: + config = DatasetEntity( + dataset_ids=[], + retrieve_config=DatasetRetrieveConfigEntity( + retrieve_strategy=DatasetRetrieveConfigEntity.RetrieveStrategy.MULTIPLE, + ), + ) result = retrieval.retrieve( self.orm_session, app_id="app-1", user_id="user-1", tenant_id="tenant-1", - model_config=model_config, + model_config=self._build_model_config(), config=config, query="python", invoke_from=InvokeFrom.WEB_APP, @@ -4668,290 +4661,527 @@ class TestRetrieveCoverage: hit_callback=Mock(), message_id="m1", ) - mock_model_manager.assert_called_once_with(tenant_id="tenant-1", user_id="user-1") - assert result == (None, []) + assert result == (None, []) - def test_retrieve_uses_bound_model_instance_schema_and_updates_model_config( - self, retrieval: DatasetRetrieval - ) -> None: - config = DatasetEntity( - dataset_ids=["d1"], - retrieve_config=DatasetRetrieveConfigEntity( - retrieve_strategy=DatasetRetrieveConfigEntity.RetrieveStrategy.SINGLE, - metadata_filtering_mode="disabled", - ), - ) - model_config = self._build_model_config(features=[]) - model_config.provider_model_bundle.model_type_instance.get_model_schema.return_value = None - bound_schema = SimpleNamespace(features=[ModelFeature.TOOL_CALL]) - bound_bundle = Mock() - bound_model_instance = Mock() - bound_model_instance.model_name = "gpt-4" - bound_model_instance.credentials = {"api_key": "secret"} - bound_model_instance.provider_model_bundle = bound_bundle - bound_model_instance.model_type_instance.get_model_schema.return_value = bound_schema - - with ( - patch("core.rag.retrieval.dataset_retrieval.ModelManager.for_tenant") as mock_model_manager, - patch.object(retrieval, "_get_available_datasets", return_value=[_dataset(id="d1")]), - patch.object(retrieval, "get_metadata_filter_condition", return_value=(None, None)), - patch.object(retrieval, "single_retrieve", return_value=[]) as mock_single_retrieve, - ): - mock_model_manager.return_value.get_model_instance.return_value = bound_model_instance - context, files = retrieval.retrieve( - self.orm_session, - app_id="app-1", - user_id="user-1", - tenant_id="tenant-1", - model_config=model_config, - config=config, - query="python", - invoke_from=InvokeFrom.WEB_APP, - show_retrieve_source=False, - hit_callback=Mock(), - message_id="m1", + def test_returns_none_when_model_schema_missing(self, retrieval: DatasetRetrieval) -> None: + config = DatasetEntity( + dataset_ids=["d1"], + retrieve_config=DatasetRetrieveConfigEntity( + retrieve_strategy=DatasetRetrieveConfigEntity.RetrieveStrategy.MULTIPLE, + ), ) + model_config = self._build_model_config() + model_instance = Mock() + model_instance.model_name = "gpt-4" + model_instance.credentials = {"api_key": "secret"} + model_instance.provider_model_bundle = Mock() + model_instance.model_type_instance.get_model_schema.return_value = None + with patch("core.rag.retrieval.dataset_retrieval.ModelManager.for_tenant") as mock_model_manager: + mock_model_manager.return_value.get_model_instance.return_value = model_instance + result = retrieval.retrieve( + self.orm_session, + app_id="app-1", + user_id="user-1", + tenant_id="tenant-1", + model_config=model_config, + config=config, + query="python", + invoke_from=InvokeFrom.WEB_APP, + show_retrieve_source=False, + hit_callback=Mock(), + message_id="m1", + ) + mock_model_manager.assert_called_once_with(tenant_id="tenant-1", user_id="user-1") + assert result == (None, []) - mock_model_manager.assert_called_once_with(tenant_id="tenant-1", user_id="user-1") - mock_single_retrieve.assert_called_once() - assert mock_single_retrieve.call_args.args[9] == PlanningStrategy.ROUTER - assert model_config.provider_model_bundle is bound_bundle - assert model_config.credentials == {"api_key": "secret"} - assert model_config.model_schema is bound_schema - assert context == "" - assert files == [] - - def test_single_strategy_with_external_documents(self, retrieval: DatasetRetrieval) -> None: - retrieve_config = DatasetRetrieveConfigEntity( - retrieve_strategy=DatasetRetrieveConfigEntity.RetrieveStrategy.SINGLE, - metadata_filtering_mode="disabled", - ) - config = DatasetEntity(dataset_ids=["d1"], retrieve_config=retrieve_config) - model_config = self._build_model_config() - external_doc = _doc( - provider="external", - content="external content", - dataset_id="ext-ds", - document_id="ext-doc", - doc_id="ext-node", - extra={"title": "External", "dataset_name": "External DS"}, - ) - with ( - patch("core.rag.retrieval.dataset_retrieval.ModelManager.for_tenant") as mock_model_manager, - patch.object(retrieval, "_get_available_datasets", return_value=[_dataset(id="d1")]), - patch.object(retrieval, "get_metadata_filter_condition", return_value=(None, None)), - patch.object(retrieval, "single_retrieve", return_value=[external_doc]), - ): + def test_retrieve_uses_bound_model_instance_schema_and_updates_model_config( + self, retrieval: DatasetRetrieval + ) -> None: + config = DatasetEntity( + dataset_ids=["d1"], + retrieve_config=DatasetRetrieveConfigEntity( + retrieve_strategy=DatasetRetrieveConfigEntity.RetrieveStrategy.SINGLE, + metadata_filtering_mode="disabled", + ), + ) + model_config = self._build_model_config(features=[]) + model_config.provider_model_bundle.model_type_instance.get_model_schema.return_value = None + bound_schema = SimpleNamespace(features=[ModelFeature.TOOL_CALL]) + bound_bundle = Mock() bound_model_instance = Mock() bound_model_instance.model_name = "gpt-4" - bound_model_instance.credentials = {} - bound_model_instance.provider_model_bundle = Mock() - bound_model_instance.model_type_instance.get_model_schema.return_value = SimpleNamespace(features=[]) - mock_model_manager.return_value.get_model_instance.return_value = bound_model_instance - context, files = retrieval.retrieve( - self.orm_session, - app_id="app-1", - user_id="user-1", - tenant_id="tenant-1", - model_config=model_config, - config=config, - query="python", - invoke_from=InvokeFrom.WEB_APP, - show_retrieve_source=False, - hit_callback=Mock(), - message_id="m1", - ) - assert context == "external content" - assert files == [] + bound_model_instance.credentials = {"api_key": "secret"} + bound_model_instance.provider_model_bundle = bound_bundle + bound_model_instance.model_type_instance.get_model_schema.return_value = bound_schema - def test_multiple_strategy_with_vision_and_source_details( - self, retrieval: DatasetRetrieval, sqlite_session: Session - ) -> None: - tenant_id = str(uuid4()) - creator_id = str(uuid4()) - dataset_id = str(uuid4()) - document_id = str(uuid4()) - retrieve_config = DatasetRetrieveConfigEntity( - retrieve_strategy=DatasetRetrieveConfigEntity.RetrieveStrategy.MULTIPLE, - top_k=4, - score_threshold=0.1, - rerank_mode="reranking_model", - reranking_model={"reranking_provider_name": "cohere", "reranking_model_name": "rerank-v3"}, - reranking_enabled=True, - metadata_filtering_mode="disabled", - ) - config = DatasetEntity(dataset_ids=[dataset_id], retrieve_config=retrieve_config) - model_config = self._build_model_config(features=[ModelFeature.TOOL_CALL]) - external_doc = _doc( - provider="external", - content="external body", - score=0.8, - dataset_id="ext-ds", - document_id="ext-doc", - doc_id="ext-node", - extra={"title": "External Title", "dataset_name": "External DS"}, - ) - dify_doc = _doc( - provider="dify", - content="dify body", - score=0.9, - dataset_id=dataset_id, - document_id=document_id, - doc_id="node-1", - ) - dataset_item = Dataset( - id=dataset_id, - tenant_id=tenant_id, - name="Dataset One", - provider="dify", - indexing_technique="high_quality", - created_by=creator_id, - ) - document_item = DatasetDocument( - id=document_id, - tenant_id=tenant_id, - dataset_id=dataset_id, - position=1, - data_source_type=DataSourceType("upload_file"), - batch="batch-1", - name="Document One", - created_from=DocumentCreatedFrom("api"), - created_by=creator_id, - indexing_status="completed", - enabled=True, - archived=False, - doc_metadata={"lang": "en"}, - ) - segment = DocumentSegment( - tenant_id=tenant_id, - dataset_id=dataset_id, - document_id=document_id, - position=1, - content="segment content", - word_count=11, - tokens=3, - created_by=creator_id, - index_node_id="node-1", - index_node_hash="hash-1", - status="completed", - hit_count=3, - answer="segment answer", - ) - upload_file = UploadFile( - tenant_id=tenant_id, - storage_type=StorageType.LOCAL, - key="k1", - name="image", - size=123, - extension="png", - mime_type="image/png", - created_by_role=CreatorUserRole.ACCOUNT, - created_by=creator_id, - created_at=datetime.now(), - used=True, - source_url="https://example.com/img.png", - ) - sqlite_session.add_all([dataset_item, document_item, segment, upload_file]) - sqlite_session.flush() - sqlite_session.add( - SegmentAttachmentBinding( + with ( + patch("core.rag.retrieval.dataset_retrieval.ModelManager.for_tenant") as mock_model_manager, + patch.object(retrieval, "_get_available_datasets", return_value=[_dataset(id="d1")]), + patch.object(retrieval, "get_metadata_filter_condition", return_value=(None, None)), + patch.object(retrieval, "single_retrieve", return_value=[]) as mock_single_retrieve, + ): + mock_model_manager.return_value.get_model_instance.return_value = bound_model_instance + context, files = retrieval.retrieve( + self.orm_session, + app_id="app-1", + user_id="user-1", + tenant_id="tenant-1", + model_config=model_config, + config=config, + query="python", + invoke_from=InvokeFrom.WEB_APP, + show_retrieve_source=False, + hit_callback=Mock(), + message_id="m1", + ) + + mock_model_manager.assert_called_once_with(tenant_id="tenant-1", user_id="user-1") + mock_single_retrieve.assert_called_once() + assert mock_single_retrieve.call_args.args[9] == PlanningStrategy.ROUTER + assert model_config.provider_model_bundle is bound_bundle + assert model_config.credentials == {"api_key": "secret"} + assert model_config.model_schema is bound_schema + assert context == "" + assert files == [] + + def test_single_strategy_with_external_documents(self, retrieval: DatasetRetrieval) -> None: + retrieve_config = DatasetRetrieveConfigEntity( + retrieve_strategy=DatasetRetrieveConfigEntity.RetrieveStrategy.SINGLE, + metadata_filtering_mode="disabled", + ) + config = DatasetEntity(dataset_ids=["d1"], retrieve_config=retrieve_config) + model_config = self._build_model_config() + external_doc = _doc( + provider="external", + content="external content", + dataset_id="ext-ds", + document_id="ext-doc", + doc_id="ext-node", + extra={"title": "External", "dataset_name": "External DS"}, + ) + with ( + patch("core.rag.retrieval.dataset_retrieval.ModelManager.for_tenant") as mock_model_manager, + patch.object(retrieval, "_get_available_datasets", return_value=[_dataset(id="d1")]), + patch.object(retrieval, "get_metadata_filter_condition", return_value=(None, None)), + patch.object(retrieval, "single_retrieve", return_value=[external_doc]), + ): + bound_model_instance = Mock() + bound_model_instance.model_name = "gpt-4" + bound_model_instance.credentials = {} + bound_model_instance.provider_model_bundle = Mock() + bound_model_instance.model_type_instance.get_model_schema.return_value = SimpleNamespace(features=[]) + mock_model_manager.return_value.get_model_instance.return_value = bound_model_instance + context, files = retrieval.retrieve( + self.orm_session, + app_id="app-1", + user_id="user-1", + tenant_id="tenant-1", + model_config=model_config, + config=config, + query="python", + invoke_from=InvokeFrom.WEB_APP, + show_retrieve_source=False, + hit_callback=Mock(), + message_id="m1", + ) + assert context == "external content" + assert files == [] + + def test_multiple_strategy_with_vision_and_source_details( + self, retrieval: DatasetRetrieval, sqlite_session: Session + ) -> None: + tenant_id = str(uuid4()) + creator_id = str(uuid4()) + dataset_id = str(uuid4()) + document_id = str(uuid4()) + retrieve_config = DatasetRetrieveConfigEntity( + retrieve_strategy=DatasetRetrieveConfigEntity.RetrieveStrategy.MULTIPLE, + top_k=4, + score_threshold=0.1, + rerank_mode="reranking_model", + reranking_model={"reranking_provider_name": "cohere", "reranking_model_name": "rerank-v3"}, + reranking_enabled=True, + metadata_filtering_mode="disabled", + ) + config = DatasetEntity(dataset_ids=[dataset_id], retrieve_config=retrieve_config) + model_config = self._build_model_config(features=[ModelFeature.TOOL_CALL]) + external_doc = _doc( + provider="external", + content="external body", + score=0.8, + dataset_id="ext-ds", + document_id="ext-doc", + doc_id="ext-node", + extra={"title": "External Title", "dataset_name": "External DS"}, + ) + dify_doc = _doc( + provider="dify", + content="dify body", + score=0.9, + dataset_id=dataset_id, + document_id=document_id, + doc_id="node-1", + ) + dataset_item = Dataset( + id=dataset_id, + tenant_id=tenant_id, + name="Dataset One", + provider="dify", + indexing_technique="high_quality", + created_by=creator_id, + ) + document_item = DatasetDocument( + id=document_id, + tenant_id=tenant_id, + dataset_id=dataset_id, + position=1, + data_source_type=DataSourceType("upload_file"), + batch="batch-1", + name="Document One", + created_from=DocumentCreatedFrom("api"), + created_by=creator_id, + indexing_status="completed", + enabled=True, + archived=False, + doc_metadata={"lang": "en"}, + ) + segment = DocumentSegment( tenant_id=tenant_id, dataset_id=dataset_id, document_id=document_id, - segment_id=segment.id, - attachment_id=upload_file.id, + position=1, + content="segment content", + word_count=11, + tokens=3, + created_by=creator_id, + index_node_id="node-1", + index_node_hash="hash-1", + status="completed", + hit_count=3, + answer="segment answer", ) - ) - sqlite_session.commit() - record = SimpleNamespace( - segment=segment, - score=0.9, - summary="short summary", - files=None, - ) - hit_callback = Mock() - - with ( - patch("core.rag.retrieval.dataset_retrieval.ModelManager.for_tenant") as mock_model_manager, - patch.object(retrieval, "_get_available_datasets", return_value=[dataset_item]), - patch.object(retrieval, "get_metadata_filter_condition", return_value=(None, None)), - patch.object(retrieval, "multiple_retrieve", return_value=[external_doc, dify_doc]), - patch( - "core.rag.retrieval.dataset_retrieval.RetrievalService.format_retrieval_documents", - return_value=[record], - ), - patch("core.rag.retrieval.dataset_retrieval.sign_upload_file_preview_url", return_value="https://signed"), - ): - bound_model_instance = Mock() - bound_model_instance.model_name = "gpt-4" - bound_model_instance.credentials = {} - bound_model_instance.provider_model_bundle = Mock() - bound_model_instance.model_type_instance.get_model_schema.return_value = SimpleNamespace( - features=[ModelFeature.TOOL_CALL] - ) - mock_model_manager.return_value.get_model_instance.return_value = bound_model_instance - context, files = retrieval.retrieve( - sqlite_session, - app_id="app-1", - user_id="user-1", + upload_file = UploadFile( tenant_id=tenant_id, - model_config=model_config, - config=config, - query="python", - invoke_from=InvokeFrom.DEBUGGER, - show_retrieve_source=True, - hit_callback=hit_callback, - message_id="m1", - vision_enabled=True, + storage_type=StorageType.LOCAL, + key="k1", + name="image", + size=123, + extension="png", + mime_type="image/png", + created_by_role=CreatorUserRole.ACCOUNT, + created_by=creator_id, + created_at=datetime.now(), + used=True, + source_url="https://example.com/img.png", + ) + sqlite_session.add_all([dataset_item, document_item, segment, upload_file]) + sqlite_session.flush() + sqlite_session.add( + SegmentAttachmentBinding( + tenant_id=tenant_id, + dataset_id=dataset_id, + document_id=document_id, + segment_id=segment.id, + attachment_id=upload_file.id, + ) + ) + sqlite_session.commit() + record = SimpleNamespace( + segment=segment, + score=0.9, + summary="short summary", + files=None, + ) + hit_callback = Mock() + + with ( + patch("core.rag.retrieval.dataset_retrieval.ModelManager.for_tenant") as mock_model_manager, + patch.object(retrieval, "_get_available_datasets", return_value=[dataset_item]), + patch.object(retrieval, "get_metadata_filter_condition", return_value=(None, None)), + patch.object(retrieval, "multiple_retrieve", return_value=[external_doc, dify_doc]), + patch( + "core.rag.retrieval.dataset_retrieval.RetrievalService.format_retrieval_documents", + return_value=[record], + ), + patch( + "core.rag.retrieval.dataset_retrieval.sign_upload_file_preview_url", return_value="https://signed" + ), + ): + bound_model_instance = Mock() + bound_model_instance.model_name = "gpt-4" + bound_model_instance.credentials = {} + bound_model_instance.provider_model_bundle = Mock() + bound_model_instance.model_type_instance.get_model_schema.return_value = SimpleNamespace( + features=[ModelFeature.TOOL_CALL] + ) + mock_model_manager.return_value.get_model_instance.return_value = bound_model_instance + context, files = retrieval.retrieve( + sqlite_session, + app_id="app-1", + user_id="user-1", + tenant_id=tenant_id, + model_config=model_config, + config=config, + query="python", + invoke_from=InvokeFrom.DEBUGGER, + show_retrieve_source=True, + hit_callback=hit_callback, + message_id="m1", + vision_enabled=True, + ) + + assert "short summary" in (context or "") + assert "question:segment content answer:segment answer" in (context or "") + assert len(files or []) == 1 + hit_callback.return_retriever_resource_info.assert_called_once() + + class TestSingleAndMultipleRetrieveCoverage: + @pytest.fixture(autouse=True) + def _orm_session(self, sqlite_session: Session) -> None: + self.orm_session = sqlite_session + + @pytest.fixture + def retrieval(self) -> DatasetRetrieval: + return DatasetRetrieval() + + def test_single_retrieve_external_path(self, retrieval: DatasetRetrieval) -> None: + dataset_id = str(uuid4()) + tenant_id = str(uuid4()) + dataset = _dataset( + id=dataset_id, + name="External DS", + description=None, + provider="external", + tenant_id=tenant_id, + created_by=str(uuid4()), + retrieval_model={"top_k": 2}, + indexing_technique="high_quality", + ) + app = Flask(__name__) + usage = LLMUsage.from_metadata({"prompt_tokens": 1, "completion_tokens": 1, "total_tokens": 2}) + self.orm_session.add(dataset) + self.orm_session.commit() + with app.app_context(): + with ( + patch("core.rag.retrieval.dataset_retrieval.ReactMultiDatasetRouter") as mock_router_cls, + patch( + "core.rag.retrieval.dataset_retrieval.ExternalDatasetService.fetch_external_knowledge_retrieval" + ) as mock_external, + patch("core.rag.retrieval.dataset_retrieval.threading.Thread", _ImmediateThread), + patch.object(retrieval, "_on_retrieval_end") as mock_end, + patch.object(retrieval, "_on_query"), + ): + mock_router_cls.return_value.invoke.return_value = (dataset_id, usage) + mock_external.return_value = [ + {"content": "ext result", "metadata": {"k": "v"}, "score": 0.9, "title": "Ext Doc"} + ] + result = retrieval.single_retrieve( + self.orm_session, + app_id="app-1", + tenant_id=tenant_id, + user_id="user-1", + user_from="workflow", + query="python", + available_datasets=[dataset], + model_instance=Mock(), + model_config=Mock(), + planning_strategy=PlanningStrategy.REACT_ROUTER, + message_id="m1", + ) + + assert len(result) == 1 + assert result[0].provider == "external" + mock_end.assert_called_once() + assert retrieval.llm_usage.total_tokens == 2 + + def test_single_retrieve_dify_path_and_filters(self, retrieval: DatasetRetrieval) -> None: + dataset_id = str(uuid4()) + tenant_id = str(uuid4()) + document_id = str(uuid4()) + dataset = _dataset( + id=dataset_id, + name="Internal DS", + description="dataset desc", + provider="dify", + tenant_id=tenant_id, + created_by=str(uuid4()), + indexing_technique="high_quality", + retrieval_model={ + "search_method": "semantic_search", + "reranking_enable": True, + "reranking_model": {"reranking_provider_name": "cohere", "reranking_model_name": "rerank"}, + "reranking_mode": "reranking_model", + "weights": {"vector_setting": {}}, + "top_k": 3, + "score_threshold_enabled": True, + "score_threshold": 0.2, + }, + ) + app = Flask(__name__) + usage = LLMUsage.from_metadata({"prompt_tokens": 1, "completion_tokens": 0, "total_tokens": 1}) + result_doc = _doc( + provider="dify", score=0.7, dataset_id=dataset_id, document_id=document_id, doc_id="node-1" + ) + self.orm_session.add(dataset) + self.orm_session.commit() + with app.app_context(): + with ( + patch("core.rag.retrieval.dataset_retrieval.FunctionCallMultiDatasetRouter") as mock_router_cls, + patch( + "core.rag.retrieval.dataset_retrieval.RetrievalService.retrieve", return_value=[result_doc] + ) as mock_retrieve, + patch("core.rag.retrieval.dataset_retrieval.threading.Thread", _ImmediateThread), + patch.object(retrieval, "_on_retrieval_end"), + patch.object(retrieval, "_on_query"), + ): + mock_router_cls.return_value.invoke.return_value = (dataset_id, usage) + results = retrieval.single_retrieve( + self.orm_session, + app_id="app-1", + tenant_id=tenant_id, + user_id="user-1", + user_from="workflow", + query="python", + available_datasets=[dataset], + model_instance=Mock(), + model_config=Mock(), + planning_strategy=PlanningStrategy.ROUTER, + metadata_filter_document_ids={dataset_id: [document_id]}, + metadata_condition=_metadata_condition(), + ) + + assert results == [result_doc] + assert mock_retrieve.call_args.kwargs["document_ids_filter"] == [document_id] + assert retrieval.llm_usage.total_tokens == 1 + + def test_single_retrieve_returns_empty_when_no_dataset_selected(self, retrieval: DatasetRetrieval) -> None: + with patch("core.rag.retrieval.dataset_retrieval.ReactMultiDatasetRouter") as mock_router_cls: + mock_router_cls.return_value.invoke.return_value = (None, LLMUsage.empty_usage()) + results = retrieval.single_retrieve( + self.orm_session, + app_id="app-1", + tenant_id="tenant-1", + user_id="user-1", + user_from="workflow", + query="python", + available_datasets=[ + _dataset(id="ds-1", name="DS", description=None), + ], + model_instance=Mock(), + model_config=Mock(), + planning_strategy=PlanningStrategy.REACT_ROUTER, + ) + assert results == [] + + def test_single_retrieve_rejects_dataset_outside_available_datasets(self, retrieval: DatasetRetrieval) -> None: + available_dataset = _dataset(id="ds-1", name="Available DS", description=None) + session = MagicMock() + session.scalar.return_value = _dataset( + id="ds-2", + name="Foreign DS", + provider="external", + tenant_id="tenant-2", + retrieval_model={}, ) - assert "short summary" in (context or "") - assert "question:segment content answer:segment answer" in (context or "") - assert len(files or []) == 1 - hit_callback.return_retriever_resource_info.assert_called_once() - - -class TestSingleAndMultipleRetrieveCoverage: - @pytest.fixture(autouse=True) - def _orm_session(self, sqlite_session: Session) -> None: - self.orm_session = sqlite_session - - @pytest.fixture - def retrieval(self) -> DatasetRetrieval: - return DatasetRetrieval() - - def test_single_retrieve_external_path(self, retrieval: DatasetRetrieval) -> None: - dataset_id = str(uuid4()) - tenant_id = str(uuid4()) - dataset = _dataset( - id=dataset_id, - name="External DS", - description=None, - provider="external", - tenant_id=tenant_id, - created_by=str(uuid4()), - retrieval_model={"top_k": 2}, - indexing_technique="high_quality", - ) - app = Flask(__name__) - usage = LLMUsage.from_metadata({"prompt_tokens": 1, "completion_tokens": 1, "total_tokens": 2}) - self.orm_session.add(dataset) - self.orm_session.commit() - with app.app_context(): with ( patch("core.rag.retrieval.dataset_retrieval.ReactMultiDatasetRouter") as mock_router_cls, patch( - "core.rag.retrieval.dataset_retrieval.ExternalDatasetService.fetch_external_knowledge_retrieval" - ) as mock_external, - patch("core.rag.retrieval.dataset_retrieval.threading.Thread", _ImmediateThread), - patch.object(retrieval, "_on_retrieval_end") as mock_end, - patch.object(retrieval, "_on_query"), + "core.rag.retrieval.dataset_retrieval.ExternalDatasetService.fetch_external_knowledge_retrieval", + return_value=[], + ) as mock_external_retrieve, + patch.object(retrieval, "_on_query") as mock_on_query, ): - mock_router_cls.return_value.invoke.return_value = (dataset_id, usage) - mock_external.return_value = [ - {"content": "ext result", "metadata": {"k": "v"}, "score": 0.9, "title": "Ext Doc"} - ] - result = retrieval.single_retrieve( + mock_router_cls.return_value.invoke.return_value = ("ds-2", LLMUsage.empty_usage()) + results = retrieval.single_retrieve( + session, + app_id="app-1", + tenant_id="tenant-1", + user_id="user-1", + user_from="workflow", + query="python", + available_datasets=[available_dataset], + model_instance=Mock(), + model_config=Mock(), + planning_strategy=PlanningStrategy.REACT_ROUTER, + ) + + assert results == [] + session.scalar.assert_not_called() + mock_external_retrieve.assert_not_called() + mock_on_query.assert_not_called() + + def test_single_retrieve_rejects_allowlisted_dataset_owned_by_another_tenant( + self, retrieval: DatasetRetrieval, sqlite_session: Session + ) -> None: + dataset_id = str(uuid4()) + caller_tenant_id = str(uuid4()) + foreign_dataset = Dataset( + id=dataset_id, + tenant_id=str(uuid4()), + name="Foreign DS", + provider="external", + indexing_technique="high_quality", + retrieval_model={}, + created_by=str(uuid4()), + ) + sqlite_session.add(foreign_dataset) + available_dataset = _dataset( + id=dataset_id, + tenant_id=caller_tenant_id, + name="Available DS", + description=None, + ) + + with ( + patch("core.rag.retrieval.dataset_retrieval.ReactMultiDatasetRouter") as mock_router_cls, + patch( + "core.rag.retrieval.dataset_retrieval.ExternalDatasetService.fetch_external_knowledge_retrieval", + ) as mock_external_retrieve, + patch( + "core.rag.retrieval.dataset_retrieval.RetrievalService.retrieve", + ) as mock_internal_retrieve, + patch.object(retrieval, "_on_query") as mock_on_query, + ): + mock_router_cls.return_value.invoke.return_value = (dataset_id, LLMUsage.empty_usage()) + results = retrieval.single_retrieve( + sqlite_session, + app_id="app-1", + tenant_id=caller_tenant_id, + user_id="user-1", + user_from="workflow", + query="python", + available_datasets=[available_dataset], + model_instance=Mock(), + model_config=Mock(), + planning_strategy=PlanningStrategy.REACT_ROUTER, + ) + + assert results == [] + mock_internal_retrieve.assert_not_called() + mock_external_retrieve.assert_not_called() + mock_on_query.assert_not_called() + + def test_single_retrieve_respects_metadata_filter_shortcuts(self, retrieval: DatasetRetrieval) -> None: + dataset_id = str(uuid4()) + tenant_id = str(uuid4()) + dataset = _dataset( + id=dataset_id, + name="Internal DS", + description="desc", + provider="dify", + tenant_id=tenant_id, + created_by=str(uuid4()), + indexing_technique="high_quality", + retrieval_model={"top_k": 2, "search_method": "semantic_search", "reranking_enable": False}, + ) + with ( + patch("core.rag.retrieval.dataset_retrieval.ReactMultiDatasetRouter") as mock_router_cls, + ): + self.orm_session.add(dataset) + self.orm_session.commit() + mock_router_cls.return_value.invoke.return_value = (dataset_id, LLMUsage.empty_usage()) + no_filter = retrieval.single_retrieve( self.orm_session, app_id="app-1", tenant_id=tenant_id, @@ -4962,54 +5192,10 @@ class TestSingleAndMultipleRetrieveCoverage: model_instance=Mock(), model_config=Mock(), planning_strategy=PlanningStrategy.REACT_ROUTER, - message_id="m1", + metadata_filter_document_ids=None, + metadata_condition=_metadata_condition(), ) - - assert len(result) == 1 - assert result[0].provider == "external" - mock_end.assert_called_once() - assert retrieval.llm_usage.total_tokens == 2 - - def test_single_retrieve_dify_path_and_filters(self, retrieval: DatasetRetrieval) -> None: - dataset_id = str(uuid4()) - tenant_id = str(uuid4()) - document_id = str(uuid4()) - dataset = _dataset( - id=dataset_id, - name="Internal DS", - description="dataset desc", - provider="dify", - tenant_id=tenant_id, - created_by=str(uuid4()), - indexing_technique="high_quality", - retrieval_model={ - "search_method": "semantic_search", - "reranking_enable": True, - "reranking_model": {"reranking_provider_name": "cohere", "reranking_model_name": "rerank"}, - "reranking_mode": "reranking_model", - "weights": {"vector_setting": {}}, - "top_k": 3, - "score_threshold_enabled": True, - "score_threshold": 0.2, - }, - ) - app = Flask(__name__) - usage = LLMUsage.from_metadata({"prompt_tokens": 1, "completion_tokens": 0, "total_tokens": 1}) - result_doc = _doc(provider="dify", score=0.7, dataset_id=dataset_id, document_id=document_id, doc_id="node-1") - self.orm_session.add(dataset) - self.orm_session.commit() - with app.app_context(): - with ( - patch("core.rag.retrieval.dataset_retrieval.FunctionCallMultiDatasetRouter") as mock_router_cls, - patch( - "core.rag.retrieval.dataset_retrieval.RetrievalService.retrieve", return_value=[result_doc] - ) as mock_retrieve, - patch("core.rag.retrieval.dataset_retrieval.threading.Thread", _ImmediateThread), - patch.object(retrieval, "_on_retrieval_end"), - patch.object(retrieval, "_on_query"), - ): - mock_router_cls.return_value.invoke.return_value = (dataset_id, usage) - results = retrieval.single_retrieve( + missing_doc_ids = retrieval.single_retrieve( self.orm_session, app_id="app-1", tenant_id=tenant_id, @@ -5019,325 +5205,121 @@ class TestSingleAndMultipleRetrieveCoverage: available_datasets=[dataset], model_instance=Mock(), model_config=Mock(), - planning_strategy=PlanningStrategy.ROUTER, - metadata_filter_document_ids={dataset_id: [document_id]}, - metadata_condition=_metadata_condition(), + planning_strategy=PlanningStrategy.REACT_ROUTER, + metadata_filter_document_ids={"other-ds": ["x"]}, + metadata_condition=None, ) + assert no_filter == [] + assert missing_doc_ids == [] - assert results == [result_doc] - assert mock_retrieve.call_args.kwargs["document_ids_filter"] == [document_id] - assert retrieval.llm_usage.total_tokens == 1 - - def test_single_retrieve_returns_empty_when_no_dataset_selected(self, retrieval: DatasetRetrieval) -> None: - with patch("core.rag.retrieval.dataset_retrieval.ReactMultiDatasetRouter") as mock_router_cls: - mock_router_cls.return_value.invoke.return_value = (None, LLMUsage.empty_usage()) - results = retrieval.single_retrieve( - self.orm_session, - app_id="app-1", - tenant_id="tenant-1", - user_id="user-1", - user_from="workflow", - query="python", - available_datasets=[ - _dataset(id="ds-1", name="DS", description=None), - ], - model_instance=Mock(), - model_config=Mock(), - planning_strategy=PlanningStrategy.REACT_ROUTER, - ) - assert results == [] - - def test_single_retrieve_rejects_dataset_outside_available_datasets(self, retrieval: DatasetRetrieval) -> None: - available_dataset = _dataset(id="ds-1", name="Available DS", description=None) - session = MagicMock() - session.scalar.return_value = _dataset( - id="ds-2", - name="Foreign DS", - provider="external", - tenant_id="tenant-2", - retrieval_model={}, - ) - - with ( - patch("core.rag.retrieval.dataset_retrieval.ReactMultiDatasetRouter") as mock_router_cls, - patch( - "core.rag.retrieval.dataset_retrieval.ExternalDatasetService.fetch_external_knowledge_retrieval", - return_value=[], - ) as mock_external_retrieve, - patch.object(retrieval, "_on_query") as mock_on_query, - ): - mock_router_cls.return_value.invoke.return_value = ("ds-2", LLMUsage.empty_usage()) - results = retrieval.single_retrieve( - session, - app_id="app-1", - tenant_id="tenant-1", - user_id="user-1", - user_from="workflow", - query="python", - available_datasets=[available_dataset], - model_instance=Mock(), - model_config=Mock(), - planning_strategy=PlanningStrategy.REACT_ROUTER, - ) - - assert results == [] - session.scalar.assert_not_called() - mock_external_retrieve.assert_not_called() - mock_on_query.assert_not_called() - - def test_single_retrieve_rejects_allowlisted_dataset_owned_by_another_tenant( - self, retrieval: DatasetRetrieval, sqlite_session: Session - ) -> None: - dataset_id = str(uuid4()) - caller_tenant_id = str(uuid4()) - foreign_dataset = Dataset( - id=dataset_id, - tenant_id=str(uuid4()), - name="Foreign DS", - provider="external", - indexing_technique="high_quality", - retrieval_model={}, - created_by=str(uuid4()), - ) - sqlite_session.add(foreign_dataset) - available_dataset = _dataset( - id=dataset_id, - tenant_id=caller_tenant_id, - name="Available DS", - description=None, - ) - - with ( - patch("core.rag.retrieval.dataset_retrieval.ReactMultiDatasetRouter") as mock_router_cls, - patch( - "core.rag.retrieval.dataset_retrieval.ExternalDatasetService.fetch_external_knowledge_retrieval", - ) as mock_external_retrieve, - patch( - "core.rag.retrieval.dataset_retrieval.RetrievalService.retrieve", - ) as mock_internal_retrieve, - patch.object(retrieval, "_on_query") as mock_on_query, - ): - mock_router_cls.return_value.invoke.return_value = (dataset_id, LLMUsage.empty_usage()) - results = retrieval.single_retrieve( - sqlite_session, - app_id="app-1", - tenant_id=caller_tenant_id, - user_id="user-1", - user_from="workflow", - query="python", - available_datasets=[available_dataset], - model_instance=Mock(), - model_config=Mock(), - planning_strategy=PlanningStrategy.REACT_ROUTER, - ) - - assert results == [] - mock_internal_retrieve.assert_not_called() - mock_external_retrieve.assert_not_called() - mock_on_query.assert_not_called() - - def test_single_retrieve_respects_metadata_filter_shortcuts(self, retrieval: DatasetRetrieval) -> None: - dataset_id = str(uuid4()) - tenant_id = str(uuid4()) - dataset = _dataset( - id=dataset_id, - name="Internal DS", - description="desc", - provider="dify", - tenant_id=tenant_id, - created_by=str(uuid4()), - indexing_technique="high_quality", - retrieval_model={"top_k": 2, "search_method": "semantic_search", "reranking_enable": False}, - ) - with ( - patch("core.rag.retrieval.dataset_retrieval.ReactMultiDatasetRouter") as mock_router_cls, - ): - self.orm_session.add(dataset) - self.orm_session.commit() - mock_router_cls.return_value.invoke.return_value = (dataset_id, LLMUsage.empty_usage()) - no_filter = retrieval.single_retrieve( - self.orm_session, - app_id="app-1", - tenant_id=tenant_id, - user_id="user-1", - user_from="workflow", - query="python", - available_datasets=[dataset], - model_instance=Mock(), - model_config=Mock(), - planning_strategy=PlanningStrategy.REACT_ROUTER, - metadata_filter_document_ids=None, - metadata_condition=_metadata_condition(), - ) - missing_doc_ids = retrieval.single_retrieve( - self.orm_session, - app_id="app-1", - tenant_id=tenant_id, - user_id="user-1", - user_from="workflow", - query="python", - available_datasets=[dataset], - model_instance=Mock(), - model_config=Mock(), - planning_strategy=PlanningStrategy.REACT_ROUTER, - metadata_filter_document_ids={"other-ds": ["x"]}, - metadata_condition=None, - ) - assert no_filter == [] - assert missing_doc_ids == [] - - def test_multiple_retrieve_validation_paths(self, retrieval: DatasetRetrieval) -> None: - assert ( - retrieval.multiple_retrieve( - app_id="app-1", - tenant_id="tenant-1", - user_id="user-1", - user_from="workflow", - available_datasets=[], - query="python", - top_k=2, - score_threshold=0.0, - reranking_mode="reranking_model", - ) - == [] - ) - - mixed = [ - _dataset(id="d1", indexing_technique="high_quality"), - _dataset(id="d2", indexing_technique="economy"), - ] - with pytest.raises(ValueError, match="different indexing technique"): - retrieval.multiple_retrieve( - app_id="app-1", - tenant_id="tenant-1", - user_id="user-1", - user_from="workflow", - available_datasets=mixed, - query="python", - top_k=2, - score_threshold=0.0, - reranking_mode="weighted_score", - reranking_enable=False, - ) - - high_quality_mismatch = [ - _dataset( - id="d1", - indexing_technique="high_quality", - embedding_model="model-a", - embedding_model_provider="provider-a", - ), - _dataset( - id="d2", - indexing_technique="high_quality", - embedding_model="model-b", - embedding_model_provider="provider-b", - ), - ] - with pytest.raises(ValueError, match="different embedding model"): - retrieval.multiple_retrieve( - app_id="app-1", - tenant_id="tenant-1", - user_id="user-1", - user_from="workflow", - available_datasets=high_quality_mismatch, - query="python", - top_k=2, - score_threshold=0.0, - reranking_mode=RerankMode.WEIGHTED_SCORE, - reranking_enable=True, - ) - - def test_multiple_retrieve_threads_and_dedup(self, retrieval: DatasetRetrieval) -> None: - datasets = [ - _dataset( - id="d1", - indexing_technique="high_quality", - embedding_model="model-a", - embedding_model_provider="provider-a", - ), - _dataset( - id="d2", - indexing_technique="high_quality", - embedding_model="model-a", - embedding_model_provider="provider-a", - ), - ] - doc_a = _doc(provider="dify", score=0.8, dataset_id="d1", document_id="doc-1", doc_id="dup") - doc_b = _doc(provider="dify", score=0.7, dataset_id="d2", document_id="doc-2", doc_id="dup") - doc_external = _doc( - provider="external", - score=0.9, - dataset_id="ext-ds", - document_id="ext-doc", - doc_id="ext-node", - extra={"dataset_name": "Ext", "title": "Ext"}, - ) - app = Flask(__name__) - weights: WeightsDict = { - "vector_setting": {"vector_weight": 0.5, "embedding_provider_name": "", "embedding_model_name": ""}, - "keyword_setting": {"keyword_weight": 0.5}, - } - - def fake_multiple_thread(**kwargs): - if kwargs["query"]: - kwargs["all_documents"].extend([doc_a, doc_b]) - if kwargs["attachment_id"]: - kwargs["all_documents"].append(doc_external) - - with app.app_context(): - with ( - patch("core.rag.retrieval.dataset_retrieval.measure_time", _timer), - patch("core.rag.retrieval.dataset_retrieval.threading.Thread", _ImmediateThread), - patch.object(retrieval, "_multiple_retrieve_thread", side_effect=fake_multiple_thread), - patch.object(retrieval, "_on_query") as mock_on_query, - patch.object(retrieval, "_on_retrieval_end") as mock_end, - ): - result = retrieval.multiple_retrieve( + def test_multiple_retrieve_validation_paths(self, retrieval: DatasetRetrieval) -> None: + assert ( + retrieval.multiple_retrieve( app_id="app-1", tenant_id="tenant-1", user_id="user-1", user_from="workflow", - available_datasets=datasets, + available_datasets=[], + query="python", + top_k=2, + score_threshold=0.0, + reranking_mode="reranking_model", + ) + == [] + ) + + mixed = [ + _dataset(id="d1", indexing_technique="high_quality"), + _dataset(id="d2", indexing_technique="economy"), + ] + with pytest.raises(ValueError, match="different indexing technique"): + retrieval.multiple_retrieve( + app_id="app-1", + tenant_id="tenant-1", + user_id="user-1", + user_from="workflow", + available_datasets=mixed, + query="python", + top_k=2, + score_threshold=0.0, + reranking_mode="weighted_score", + reranking_enable=False, + ) + + high_quality_mismatch = [ + _dataset( + id="d1", + indexing_technique="high_quality", + embedding_model="model-a", + embedding_model_provider="provider-a", + ), + _dataset( + id="d2", + indexing_technique="high_quality", + embedding_model="model-b", + embedding_model_provider="provider-b", + ), + ] + with pytest.raises(ValueError, match="different embedding model"): + retrieval.multiple_retrieve( + app_id="app-1", + tenant_id="tenant-1", + user_id="user-1", + user_from="workflow", + available_datasets=high_quality_mismatch, query="python", top_k=2, score_threshold=0.0, reranking_mode=RerankMode.WEIGHTED_SCORE, reranking_enable=True, - weights=weights, - attachment_ids=["att-1"], - message_id="m1", ) - assert len(result) == 2 - assert any(doc.provider == "external" for doc in result) - assert weights["vector_setting"]["embedding_provider_name"] == "provider-a" - assert weights["vector_setting"]["embedding_model_name"] == "model-a" - mock_on_query.assert_called_once() - mock_end.assert_called_once() - - def test_multiple_retrieve_propagates_thread_exception(self, retrieval: DatasetRetrieval) -> None: - datasets = [ - _dataset( - id="d1", - indexing_technique="high_quality", - embedding_model="model-a", - embedding_model_provider="provider-a", + def test_multiple_retrieve_threads_and_dedup(self, retrieval: DatasetRetrieval) -> None: + datasets = [ + _dataset( + id="d1", + indexing_technique="high_quality", + embedding_model="model-a", + embedding_model_provider="provider-a", + ), + _dataset( + id="d2", + indexing_technique="high_quality", + embedding_model="model-a", + embedding_model_provider="provider-a", + ), + ] + doc_a = _doc(provider="dify", score=0.8, dataset_id="d1", document_id="doc-1", doc_id="dup") + doc_b = _doc(provider="dify", score=0.7, dataset_id="d2", document_id="doc-2", doc_id="dup") + doc_external = _doc( + provider="external", + score=0.9, + dataset_id="ext-ds", + document_id="ext-doc", + doc_id="ext-node", + extra={"dataset_name": "Ext", "title": "Ext"}, ) - ] - app = Flask(__name__) + app = Flask(__name__) + weights: WeightsDict = { + "vector_setting": {"vector_weight": 0.5, "embedding_provider_name": "", "embedding_model_name": ""}, + "keyword_setting": {"keyword_weight": 0.5}, + } - def failing_thread(**_kwargs): - raise RuntimeError("thread boom") + def fake_multiple_thread(**kwargs): + if kwargs["query"]: + kwargs["all_documents"].extend([doc_a, doc_b]) + if kwargs["attachment_id"]: + kwargs["all_documents"].append(doc_external) - with app.app_context(): - with ( - patch("core.rag.retrieval.dataset_retrieval.measure_time", _timer), - patch("core.rag.retrieval.dataset_retrieval.threading.Thread", _ImmediateThread), - patch.object(retrieval, "_multiple_retrieve_thread", side_effect=failing_thread), - ): - with pytest.raises(RuntimeError, match="thread boom"): - retrieval.multiple_retrieve( + with app.app_context(): + with ( + patch("core.rag.retrieval.dataset_retrieval.measure_time", _timer), + patch("core.rag.retrieval.dataset_retrieval.threading.Thread", _ImmediateThread), + patch.object(retrieval, "_multiple_retrieve_thread", side_effect=fake_multiple_thread), + patch.object(retrieval, "_on_query") as mock_on_query, + patch.object(retrieval, "_on_retrieval_end") as mock_end, + ): + result = retrieval.multiple_retrieve( app_id="app-1", tenant_id="tenant-1", user_id="user-1", @@ -5346,406 +5328,95 @@ class TestSingleAndMultipleRetrieveCoverage: query="python", top_k=2, score_threshold=0.0, - reranking_mode="reranking_model", + reranking_mode=RerankMode.WEIGHTED_SCORE, + reranking_enable=True, + weights=weights, + attachment_ids=["att-1"], + message_id="m1", ) + assert len(result) == 2 + assert any(doc.provider == "external" for doc in result) + assert weights["vector_setting"]["embedding_provider_name"] == "provider-a" + assert weights["vector_setting"]["embedding_model_name"] == "model-a" + mock_on_query.assert_called_once() + mock_end.assert_called_once() -class TestInternalHooksCoverage: - @pytest.fixture(autouse=True) - def _orm_resources(self, sqlite_engine: Engine, sqlite_session: Session) -> None: - self.orm_engine = sqlite_engine - self.orm_session = sqlite_session - - @pytest.fixture - def retrieval(self) -> DatasetRetrieval: - return DatasetRetrieval() - - def test_on_retrieval_end_without_dify_documents(self, retrieval: DatasetRetrieval) -> None: - app = Flask(__name__) - with patch.object(retrieval, "_send_trace_task") as mock_trace: - retrieval._on_retrieval_end( - flask_app=app, - documents=[_doc(provider="external")], - message_id="m1", - timer={"cost": 1}, - ) - mock_trace.assert_called_once() - - def test_on_retrieval_end_dify_without_document_ids(self, retrieval: DatasetRetrieval) -> None: - app = Flask(__name__) - doc = Document(page_content="x", metadata={"doc_id": "n1"}, provider="dify") - with ( - patch("core.rag.retrieval.dataset_retrieval.db", SimpleNamespace(engine=self.orm_engine)), - patch.object(retrieval, "_send_trace_task") as mock_trace, - ): - retrieval._on_retrieval_end(flask_app=app, documents=[doc], message_id="m1", timer={"cost": 1}) - mock_trace.assert_called_once() - - def test_on_retrieval_end_updates_segments_for_text_and_image(self, retrieval: DatasetRetrieval) -> None: - app = Flask(__name__) - tenant_id = str(uuid4()) - dataset_id = str(uuid4()) - creator_id = str(uuid4()) - document_ids = [str(uuid4()) for _ in range(4)] - dataset_docs = [ - DatasetDocument( - id=document_id, - tenant_id=tenant_id, - dataset_id=dataset_id, - position=position, - data_source_type=DataSourceType("upload_file"), - batch="batch-1", - name=f"document-{position}", - created_from=DocumentCreatedFrom("api"), - created_by=creator_id, - indexing_status="completed", - enabled=True, - archived=False, - doc_form=doc_form, - ) - for position, (document_id, doc_form) in enumerate( - zip( - document_ids, - [ - IndexStructureType.PARENT_CHILD_INDEX, - IndexStructureType.PARENT_CHILD_INDEX, - IndexStructureType.QA_INDEX, - IndexStructureType.QA_INDEX, - ], - strict=True, - ), - start=1, - ) - ] - segments = [ - DocumentSegment( - tenant_id=tenant_id, - dataset_id=dataset_id, - document_id=document_id, - position=1, - content=f"segment-{position}", - word_count=1, - tokens=1, - created_by=creator_id, - index_node_id="idx-c" if position == 3 else None, - status="completed", - ) - for position, document_id in enumerate(document_ids, start=1) - ] - self.orm_session.add_all(dataset_docs + segments) - self.orm_session.flush() - self.orm_session.add_all( - [ - ChildChunk( - tenant_id=tenant_id, - dataset_id=dataset_id, - document_id=document_ids[0], - segment_id=segments[0].id, - position=1, - content="child", - word_count=1, - created_by=creator_id, - index_node_id="idx-a", - ), - SegmentAttachmentBinding( - tenant_id=tenant_id, - dataset_id=dataset_id, - document_id=document_ids[1], - segment_id=segments[1].id, - attachment_id=str(uuid4()), - ), - SegmentAttachmentBinding( - tenant_id=tenant_id, - dataset_id=dataset_id, - document_id=document_ids[3], - segment_id=segments[3].id, - attachment_id=str(uuid4()), - ), + def test_multiple_retrieve_propagates_thread_exception(self, retrieval: DatasetRetrieval) -> None: + datasets = [ + _dataset( + id="d1", + indexing_technique="high_quality", + embedding_model="model-a", + embedding_model_provider="provider-a", + ) ] - ) - self.orm_session.commit() - docs = [ - _doc(provider="dify", document_id=document_ids[0], doc_id="idx-a", extra={"doc_type": "text"}), - _doc( - provider="dify", - document_id=document_ids[1], - doc_id=self.orm_session.scalar( - select(SegmentAttachmentBinding.attachment_id).where( - SegmentAttachmentBinding.segment_id == segments[1].id - ) - ), - extra={"doc_type": DocType.IMAGE}, - ), - _doc(provider="dify", document_id=document_ids[2], doc_id="idx-c", extra={"doc_type": "text"}), - _doc( - provider="dify", - document_id=document_ids[3], - doc_id=self.orm_session.scalar( - select(SegmentAttachmentBinding.attachment_id).where( - SegmentAttachmentBinding.segment_id == segments[3].id - ) - ), - extra={"doc_type": DocType.IMAGE}, - ), - ] + app = Flask(__name__) - executed_statements: list[object] = [] + def failing_thread(**_kwargs): + raise RuntimeError("thread boom") - def capture_statement(*args: object) -> None: - executed_statements.append(args[1]) + with app.app_context(): + with ( + patch("core.rag.retrieval.dataset_retrieval.measure_time", _timer), + patch("core.rag.retrieval.dataset_retrieval.threading.Thread", _ImmediateThread), + patch.object(retrieval, "_multiple_retrieve_thread", side_effect=failing_thread), + ): + with pytest.raises(RuntimeError, match="thread boom"): + retrieval.multiple_retrieve( + app_id="app-1", + tenant_id="tenant-1", + user_id="user-1", + user_from="workflow", + available_datasets=datasets, + query="python", + top_k=2, + score_threshold=0.0, + reranking_mode="reranking_model", + ) - event.listen(self.orm_engine, "before_execute", capture_statement) - try: + class TestInternalHooksCoverage: + @pytest.fixture(autouse=True) + def _orm_resources(self, sqlite_engine: Engine, sqlite_session: Session) -> None: + self.orm_engine = sqlite_engine + self.orm_session = sqlite_session + + @pytest.fixture + def retrieval(self) -> DatasetRetrieval: + return DatasetRetrieval() + + def test_on_retrieval_end_without_dify_documents(self, retrieval: DatasetRetrieval) -> None: + app = Flask(__name__) + with patch.object(retrieval, "_send_trace_task") as mock_trace: + retrieval._on_retrieval_end( + flask_app=app, + documents=[_doc(provider="external")], + message_id="m1", + timer={"cost": 1}, + ) + mock_trace.assert_called_once() + + def test_on_retrieval_end_dify_without_document_ids(self, retrieval: DatasetRetrieval) -> None: + app = Flask(__name__) + doc = _dataset_retrieval_Document(page_content="x", metadata={"doc_id": "n1"}, provider="dify") with ( patch("core.rag.retrieval.dataset_retrieval.db", SimpleNamespace(engine=self.orm_engine)), patch.object(retrieval, "_send_trace_task") as mock_trace, ): - retrieval._on_retrieval_end(flask_app=app, documents=docs, message_id="m1", timer={"cost": 1}) - finally: - event.remove(self.orm_engine, "before_execute", capture_statement) + retrieval._on_retrieval_end(flask_app=app, documents=[doc], message_id="m1", timer={"cost": 1}) + mock_trace.assert_called_once() - self.orm_session.expire_all() - assert [segment.hit_count for segment in segments] == [1, 1, 1, 1] - locking_statements = [ - statement - for statement in executed_statements - if "ORDER BY document_segments.id FOR UPDATE" in str(statement.compile(dialect=postgresql.dialect())) # type: ignore[union-attr] - ] - assert len(locking_statements) == 1 - mock_trace.assert_called_once() - - def test_on_retrieval_end_retries_postgres_deadlock(self, retrieval: DatasetRetrieval) -> None: - app = Flask(__name__) - deadlock = DBAPIError( - "UPDATE document_segments", - {}, - SimpleNamespace(sqlstate="40P01"), - connection_invalidated=False, - ) - doc = _doc(provider="dify") - retry_session = MagicMock() - - with ( - patch("core.rag.retrieval.dataset_retrieval.db", SimpleNamespace(engine=self.orm_engine)), - patch("core.rag.retrieval.dataset_retrieval.sessionmaker") as mock_sessionmaker, - patch.object(retrieval, "_send_trace_task") as mock_trace, - ): - mock_sessionmaker.return_value.begin.side_effect = [deadlock, nullcontext(retry_session)] - retrieval._on_retrieval_end(flask_app=app, documents=[doc], message_id="m1", timer={"cost": 1}) - - assert mock_sessionmaker.return_value.begin.call_count == 2 - mock_trace.assert_called_once_with("m1", [doc], {"cost": 1}) - - @pytest.mark.parametrize( - ("error", "expected"), - [ - (RuntimeError("not a database error"), False), - (DBAPIError("SELECT 1", {}, None, connection_invalidated=False), False), - ( - DBAPIError( - "UPDATE document_segments", - {}, - SimpleNamespace(pgcode="40P01"), - connection_invalidated=False, - ), - True, - ), - ( - DBAPIError( - "UPDATE document_segments", - {}, - SimpleNamespace(sqlstate="40001", pgcode="40P01"), - connection_invalidated=False, - ), - True, - ), - ( - DBAPIError( - "UPDATE document_segments", - {}, - SimpleNamespace(sqlstate="40001", pgcode="40001"), - connection_invalidated=False, - ), - False, - ), - ], - ) - def test_is_postgres_deadlock_error(self, error: BaseException, expected: bool) -> None: - assert dataset_retrieval_module._is_postgres_deadlock_error(error) is expected - - def test_on_retrieval_end_reraises_non_deadlock(self, retrieval: DatasetRetrieval) -> None: - app = Flask(__name__) - error = DBAPIError( - "UPDATE document_segments", - {}, - SimpleNamespace(sqlstate="40001"), - connection_invalidated=False, - ) - doc = _doc(provider="dify") - - with ( - patch("core.rag.retrieval.dataset_retrieval.db", SimpleNamespace(engine=self.orm_engine)), - patch("core.rag.retrieval.dataset_retrieval.sessionmaker") as mock_sessionmaker, - patch.object(retrieval, "_send_trace_task") as mock_trace, - ): - mock_sessionmaker.return_value.begin.side_effect = error - with pytest.raises(DBAPIError): - retrieval._on_retrieval_end(flask_app=app, documents=[doc]) - - assert mock_sessionmaker.return_value.begin.call_count == 1 - mock_trace.assert_not_called() - - def test_on_retrieval_end_reraises_after_deadlock_retries_are_exhausted(self, retrieval: DatasetRetrieval) -> None: - app = Flask(__name__) - deadlock = DBAPIError( - "UPDATE document_segments", - {}, - SimpleNamespace(sqlstate="40P01"), - connection_invalidated=False, - ) - doc = _doc(provider="dify") - - with ( - patch("core.rag.retrieval.dataset_retrieval.db", SimpleNamespace(engine=self.orm_engine)), - patch("core.rag.retrieval.dataset_retrieval.sessionmaker") as mock_sessionmaker, - patch.object(retrieval, "_send_trace_task") as mock_trace, - ): - mock_sessionmaker.return_value.begin.side_effect = deadlock - with pytest.raises(DBAPIError): - retrieval._on_retrieval_end(flask_app=app, documents=[doc]) - - assert mock_sessionmaker.return_value.begin.call_count == 3 - mock_trace.assert_not_called() - - def test_retriever_variants(self, retrieval: DatasetRetrieval) -> None: - flask_app = SimpleNamespace(app_context=lambda: nullcontext()) - all_documents: list[Document] = [] - tenant_id = str(uuid4()) - creator_id = str(uuid4()) - missing_id = str(uuid4()) - external_id = str(uuid4()) - economy_id = str(uuid4()) - high_quality_id = str(uuid4()) - - assert ( - retrieval._retriever( - flask_app=flask_app, # type: ignore[arg-type] - session=self.orm_session, - dataset_id=missing_id, - query="python", - top_k=1, - all_documents=all_documents, - ) - == [] - ) - - external_dataset = Dataset( - id=external_id, - name="External", - provider="external", - tenant_id=tenant_id, - retrieval_model={"top_k": 2}, - indexing_technique="high_quality", - created_by=creator_id, - ) - economy_dataset = Dataset( - id=economy_id, - tenant_id=tenant_id, - name="Economy", - provider="dify", - retrieval_model={"top_k": 1}, - indexing_technique="economy", - created_by=creator_id, - ) - high_dataset = Dataset( - id=high_quality_id, - tenant_id=tenant_id, - name="High quality", - provider="dify", - created_by=creator_id, - indexing_technique="high_quality", - retrieval_model={ - "search_method": "semantic_search", - "top_k": 4, - "score_threshold": 0.3, - "score_threshold_enabled": True, - "reranking_enable": True, - "reranking_model": {"reranking_provider_name": "x", "reranking_model_name": "y"}, - "reranking_mode": "reranking_model", - "weights": {"vector_setting": {}}, - }, - ) - self.orm_session.add_all([external_dataset, economy_dataset, high_dataset]) - self.orm_session.commit() - with ( - patch( - "core.rag.retrieval.dataset_retrieval.ExternalDatasetService.fetch_external_knowledge_retrieval" - ) as mock_external, - ): - mock_external.return_value = [{"content": "e", "metadata": {}, "score": 0.8, "title": "Ext"}] - retrieval._retriever( - flask_app=flask_app, # type: ignore[arg-type] - session=self.orm_session, - dataset_id=external_id, - query="python", - top_k=1, - all_documents=all_documents, - ) - - with ( - patch( - "core.rag.retrieval.dataset_retrieval.RetrievalService.retrieve", return_value=[_doc(provider="dify")] - ) as mock_retrieve, - ): - retrieval._retriever( - flask_app=flask_app, # type: ignore[arg-type] - session=self.orm_session, - dataset_id=economy_id, - query="python", - top_k=2, - all_documents=all_documents, - ) - retrieval._retriever( - flask_app=flask_app, # type: ignore[arg-type] - session=self.orm_session, - dataset_id=high_quality_id, - query="python", - top_k=2, - all_documents=all_documents, - attachment_ids=["att-1"], - ) - assert mock_retrieve.call_count == 2 - assert len(all_documents) >= 3 - - def test_to_dataset_retriever_tool_paths(self, retrieval: DatasetRetrieval) -> None: - tenant_id = str(uuid4()) - creator_id = str(uuid4()) - missing_id = str(uuid4()) - skipped_id = str(uuid4()) - available_id = str(uuid4()) - dataset_skip_zero = Dataset( - id=skipped_id, - tenant_id=tenant_id, - name="Empty dataset", - provider="dify", - indexing_technique="high_quality", - created_by=creator_id, - ) - dataset_ok_single = Dataset( - id=available_id, - tenant_id=tenant_id, - name="Available dataset", - provider="dify", - indexing_technique="high_quality", - created_by=creator_id, - retrieval_model={"top_k": 2, "score_threshold_enabled": True, "score_threshold": 0.1}, - ) - self.orm_session.add_all([dataset_skip_zero, dataset_ok_single]) - self.orm_session.add_all( - [ + def test_on_retrieval_end_updates_segments_for_text_and_image(self, retrieval: DatasetRetrieval) -> None: + app = Flask(__name__) + tenant_id = str(uuid4()) + dataset_id = str(uuid4()) + creator_id = str(uuid4()) + document_ids = [str(uuid4()) for _ in range(4)] + dataset_docs = [ DatasetDocument( + id=document_id, tenant_id=tenant_id, - dataset_id=available_id, + dataset_id=dataset_id, position=position, data_source_type=DataSourceType("upload_file"), batch="batch-1", @@ -5755,150 +5426,1141 @@ class TestInternalHooksCoverage: indexing_status="completed", enabled=True, archived=False, + doc_form=doc_form, + ) + for position, (document_id, doc_form) in enumerate( + zip( + document_ids, + [ + IndexStructureType.PARENT_CHILD_INDEX, + IndexStructureType.PARENT_CHILD_INDEX, + IndexStructureType.QA_INDEX, + IndexStructureType.QA_INDEX, + ], + strict=True, + ), + start=1, ) - for position in (1, 2) ] + segments = [ + DocumentSegment( + tenant_id=tenant_id, + dataset_id=dataset_id, + document_id=document_id, + position=1, + content=f"segment-{position}", + word_count=1, + tokens=1, + created_by=creator_id, + index_node_id="idx-c" if position == 3 else None, + status="completed", + ) + for position, document_id in enumerate(document_ids, start=1) + ] + self.orm_session.add_all(dataset_docs + segments) + self.orm_session.flush() + self.orm_session.add_all( + [ + ChildChunk( + tenant_id=tenant_id, + dataset_id=dataset_id, + document_id=document_ids[0], + segment_id=segments[0].id, + position=1, + content="child", + word_count=1, + created_by=creator_id, + index_node_id="idx-a", + ), + SegmentAttachmentBinding( + tenant_id=tenant_id, + dataset_id=dataset_id, + document_id=document_ids[1], + segment_id=segments[1].id, + attachment_id=str(uuid4()), + ), + SegmentAttachmentBinding( + tenant_id=tenant_id, + dataset_id=dataset_id, + document_id=document_ids[3], + segment_id=segments[3].id, + attachment_id=str(uuid4()), + ), + ] + ) + self.orm_session.commit() + docs = [ + _doc(provider="dify", document_id=document_ids[0], doc_id="idx-a", extra={"doc_type": "text"}), + _doc( + provider="dify", + document_id=document_ids[1], + doc_id=self.orm_session.scalar( + select(SegmentAttachmentBinding.attachment_id).where( + SegmentAttachmentBinding.segment_id == segments[1].id + ) + ), + extra={"doc_type": DocType.IMAGE}, + ), + _doc(provider="dify", document_id=document_ids[2], doc_id="idx-c", extra={"doc_type": "text"}), + _doc( + provider="dify", + document_id=document_ids[3], + doc_id=self.orm_session.scalar( + select(SegmentAttachmentBinding.attachment_id).where( + SegmentAttachmentBinding.segment_id == segments[3].id + ) + ), + extra={"doc_type": DocType.IMAGE}, + ), + ] + + executed_statements: list[object] = [] + + def capture_statement(*args: object) -> None: + executed_statements.append(args[1]) + + event.listen(self.orm_engine, "before_execute", capture_statement) + try: + with ( + patch("core.rag.retrieval.dataset_retrieval.db", SimpleNamespace(engine=self.orm_engine)), + patch.object(retrieval, "_send_trace_task") as mock_trace, + ): + retrieval._on_retrieval_end(flask_app=app, documents=docs, message_id="m1", timer={"cost": 1}) + finally: + event.remove(self.orm_engine, "before_execute", capture_statement) + + self.orm_session.expire_all() + assert [segment.hit_count for segment in segments] == [1, 1, 1, 1] + locking_statements = [ + statement + for statement in executed_statements + if "ORDER BY document_segments.id FOR UPDATE" in str(statement.compile(dialect=postgresql.dialect())) # type: ignore[union-attr] + ] + assert len(locking_statements) == 1 + mock_trace.assert_called_once() + + def test_on_retrieval_end_retries_postgres_deadlock(self, retrieval: DatasetRetrieval) -> None: + app = Flask(__name__) + deadlock = DBAPIError( + "UPDATE document_segments", + {}, + SimpleNamespace(sqlstate="40P01"), + connection_invalidated=False, + ) + doc = _doc(provider="dify") + retry_session = MagicMock() + + with ( + patch("core.rag.retrieval.dataset_retrieval.db", SimpleNamespace(engine=self.orm_engine)), + patch("core.rag.retrieval.dataset_retrieval.sessionmaker") as mock_sessionmaker, + patch.object(retrieval, "_send_trace_task") as mock_trace, + ): + mock_sessionmaker.return_value.begin.side_effect = [deadlock, nullcontext(retry_session)] + retrieval._on_retrieval_end(flask_app=app, documents=[doc], message_id="m1", timer={"cost": 1}) + + assert mock_sessionmaker.return_value.begin.call_count == 2 + mock_trace.assert_called_once_with("m1", [doc], {"cost": 1}) + + @pytest.mark.parametrize( + ("error", "expected"), + [ + (RuntimeError("not a database error"), False), + (DBAPIError("SELECT 1", {}, None, connection_invalidated=False), False), + ( + DBAPIError( + "UPDATE document_segments", + {}, + SimpleNamespace(pgcode="40P01"), + connection_invalidated=False, + ), + True, + ), + ( + DBAPIError( + "UPDATE document_segments", + {}, + SimpleNamespace(sqlstate="40001", pgcode="40P01"), + connection_invalidated=False, + ), + True, + ), + ( + DBAPIError( + "UPDATE document_segments", + {}, + SimpleNamespace(sqlstate="40001", pgcode="40001"), + connection_invalidated=False, + ), + False, + ), + ], ) - self.orm_session.commit() - single_config = DatasetRetrieveConfigEntity( - retrieve_strategy=DatasetRetrieveConfigEntity.RetrieveStrategy.SINGLE, - metadata_filtering_mode="disabled", + def test_is_postgres_deadlock_error(self, error: BaseException, expected: bool) -> None: + assert dataset_retrieval_module._is_postgres_deadlock_error(error) is expected + + def test_on_retrieval_end_reraises_non_deadlock(self, retrieval: DatasetRetrieval) -> None: + app = Flask(__name__) + error = DBAPIError( + "UPDATE document_segments", + {}, + SimpleNamespace(sqlstate="40001"), + connection_invalidated=False, + ) + doc = _doc(provider="dify") + + with ( + patch("core.rag.retrieval.dataset_retrieval.db", SimpleNamespace(engine=self.orm_engine)), + patch("core.rag.retrieval.dataset_retrieval.sessionmaker") as mock_sessionmaker, + patch.object(retrieval, "_send_trace_task") as mock_trace, + ): + mock_sessionmaker.return_value.begin.side_effect = error + with pytest.raises(DBAPIError): + retrieval._on_retrieval_end(flask_app=app, documents=[doc]) + + assert mock_sessionmaker.return_value.begin.call_count == 1 + mock_trace.assert_not_called() + + def test_on_retrieval_end_reraises_after_deadlock_retries_are_exhausted( + self, retrieval: DatasetRetrieval + ) -> None: + app = Flask(__name__) + deadlock = DBAPIError( + "UPDATE document_segments", + {}, + SimpleNamespace(sqlstate="40P01"), + connection_invalidated=False, + ) + doc = _doc(provider="dify") + + with ( + patch("core.rag.retrieval.dataset_retrieval.db", SimpleNamespace(engine=self.orm_engine)), + patch("core.rag.retrieval.dataset_retrieval.sessionmaker") as mock_sessionmaker, + patch.object(retrieval, "_send_trace_task") as mock_trace, + ): + mock_sessionmaker.return_value.begin.side_effect = deadlock + with pytest.raises(DBAPIError): + retrieval._on_retrieval_end(flask_app=app, documents=[doc]) + + assert mock_sessionmaker.return_value.begin.call_count == 3 + mock_trace.assert_not_called() + + def test_retriever_variants(self, retrieval: DatasetRetrieval) -> None: + flask_app = SimpleNamespace(app_context=lambda: nullcontext()) + all_documents: list[_dataset_retrieval_Document] = [] + tenant_id = str(uuid4()) + creator_id = str(uuid4()) + missing_id = str(uuid4()) + external_id = str(uuid4()) + economy_id = str(uuid4()) + high_quality_id = str(uuid4()) + + assert ( + retrieval._retriever( + flask_app=flask_app, # type: ignore[arg-type] + session=self.orm_session, + dataset_id=missing_id, + query="python", + top_k=1, + all_documents=all_documents, + ) + == [] + ) + + external_dataset = Dataset( + id=external_id, + name="External", + provider="external", + tenant_id=tenant_id, + retrieval_model={"top_k": 2}, + indexing_technique="high_quality", + created_by=creator_id, + ) + economy_dataset = Dataset( + id=economy_id, + tenant_id=tenant_id, + name="Economy", + provider="dify", + retrieval_model={"top_k": 1}, + indexing_technique="economy", + created_by=creator_id, + ) + high_dataset = Dataset( + id=high_quality_id, + tenant_id=tenant_id, + name="High quality", + provider="dify", + created_by=creator_id, + indexing_technique="high_quality", + retrieval_model={ + "search_method": "semantic_search", + "top_k": 4, + "score_threshold": 0.3, + "score_threshold_enabled": True, + "reranking_enable": True, + "reranking_model": {"reranking_provider_name": "x", "reranking_model_name": "y"}, + "reranking_mode": "reranking_model", + "weights": {"vector_setting": {}}, + }, + ) + self.orm_session.add_all([external_dataset, economy_dataset, high_dataset]) + self.orm_session.commit() + with ( + patch( + "core.rag.retrieval.dataset_retrieval.ExternalDatasetService.fetch_external_knowledge_retrieval" + ) as mock_external, + ): + mock_external.return_value = [{"content": "e", "metadata": {}, "score": 0.8, "title": "Ext"}] + retrieval._retriever( + flask_app=flask_app, # type: ignore[arg-type] + session=self.orm_session, + dataset_id=external_id, + query="python", + top_k=1, + all_documents=all_documents, + ) + + with ( + patch( + "core.rag.retrieval.dataset_retrieval.RetrievalService.retrieve", + return_value=[_doc(provider="dify")], + ) as mock_retrieve, + ): + retrieval._retriever( + flask_app=flask_app, # type: ignore[arg-type] + session=self.orm_session, + dataset_id=economy_id, + query="python", + top_k=2, + all_documents=all_documents, + ) + retrieval._retriever( + flask_app=flask_app, # type: ignore[arg-type] + session=self.orm_session, + dataset_id=high_quality_id, + query="python", + top_k=2, + all_documents=all_documents, + attachment_ids=["att-1"], + ) + assert mock_retrieve.call_count == 2 + assert len(all_documents) >= 3 + + def test_to_dataset_retriever_tool_paths(self, retrieval: DatasetRetrieval) -> None: + tenant_id = str(uuid4()) + creator_id = str(uuid4()) + missing_id = str(uuid4()) + skipped_id = str(uuid4()) + available_id = str(uuid4()) + dataset_skip_zero = Dataset( + id=skipped_id, + tenant_id=tenant_id, + name="Empty dataset", + provider="dify", + indexing_technique="high_quality", + created_by=creator_id, + ) + dataset_ok_single = Dataset( + id=available_id, + tenant_id=tenant_id, + name="Available dataset", + provider="dify", + indexing_technique="high_quality", + created_by=creator_id, + retrieval_model={"top_k": 2, "score_threshold_enabled": True, "score_threshold": 0.1}, + ) + self.orm_session.add_all([dataset_skip_zero, dataset_ok_single]) + self.orm_session.add_all( + [ + DatasetDocument( + tenant_id=tenant_id, + dataset_id=available_id, + position=position, + data_source_type=DataSourceType("upload_file"), + batch="batch-1", + name=f"document-{position}", + created_from=DocumentCreatedFrom("api"), + created_by=creator_id, + indexing_status="completed", + enabled=True, + archived=False, + ) + for position in (1, 2) + ] + ) + self.orm_session.commit() + single_config = DatasetRetrieveConfigEntity( + retrieve_strategy=DatasetRetrieveConfigEntity.RetrieveStrategy.SINGLE, + metadata_filtering_mode="disabled", + ) + with ( + patch( + "core.tools.utils.dataset_retriever.dataset_retriever_tool.DatasetRetrieverTool.from_dataset", + return_value="single-tool", + ) as mock_single_tool, + ): + single_tools = retrieval.to_dataset_retriever_tool( + session=self.orm_session, + tenant_id=tenant_id, + dataset_ids=[missing_id, skipped_id, available_id], + retrieve_config=single_config, + return_resource=True, + invoke_from=InvokeFrom.WEB_APP, + hit_callback=Mock(), + user_id="user-1", + inputs={"k": "v"}, + ) + + assert single_tools == ["single-tool"] + mock_single_tool.assert_called_once() + + multiple_config_missing = DatasetRetrieveConfigEntity( + retrieve_strategy=DatasetRetrieveConfigEntity.RetrieveStrategy.MULTIPLE, + metadata_filtering_mode="disabled", + reranking_model=None, + ) + with pytest.raises(ValueError, match="Reranking model is required"): + retrieval.to_dataset_retriever_tool( + session=self.orm_session, + tenant_id=tenant_id, + dataset_ids=[available_id], + retrieve_config=multiple_config_missing, + return_resource=True, + invoke_from=InvokeFrom.WEB_APP, + hit_callback=Mock(), + user_id="user-1", + inputs={}, + ) + + multiple_config = DatasetRetrieveConfigEntity( + retrieve_strategy=DatasetRetrieveConfigEntity.RetrieveStrategy.MULTIPLE, + metadata_filtering_mode="disabled", + top_k=3, + score_threshold=0.2, + reranking_model={"reranking_provider_name": "cohere", "reranking_model_name": "rerank-v3"}, + ) + with ( + patch( + "core.tools.utils.dataset_retriever.dataset_multi_retriever_tool.DatasetMultiRetrieverTool.from_dataset", + return_value="multi-tool", + ) as mock_multi_tool, + ): + multi_tools = retrieval.to_dataset_retriever_tool( + session=self.orm_session, + tenant_id=tenant_id, + dataset_ids=[available_id], + retrieve_config=multiple_config, + return_resource=False, + invoke_from=InvokeFrom.DEBUGGER, + hit_callback=Mock(), + user_id="user-1", + inputs={}, + ) + assert multi_tools == ["multi-tool"] + mock_multi_tool.assert_called_once() + + def test_additional_small_branches(self, retrieval: DatasetRetrieval) -> None: + keyword_handler = Mock() + keyword_handler.extract_keywords.side_effect = [[], []] + doc = _dataset_retrieval_Document(page_content="doc", metadata={"doc_id": "1"}, provider="dify") + with patch("core.rag.retrieval.dataset_retrieval.JiebaKeywordTableHandler", return_value=keyword_handler): + ranked = retrieval.calculate_keyword_score("query", [doc], top_k=1) + assert len(ranked) == 1 + assert ranked[0].metadata.get("score") == 0.0 + + with pytest.raises(ValueError, match="metadata_model_config is required"): + retrieval._automatic_metadata_filter_func( + self.orm_session, + dataset_ids=[str(uuid4())], + query="python", + tenant_id=str(uuid4()), + user_id=str(uuid4()), + metadata_model_config=None, # type: ignore[arg-type] + ) + + tenant_id = str(uuid4()) + dataset_id = str(uuid4()) + user_id = str(uuid4()) + self.orm_session.add( + DatasetMetadata( + tenant_id=tenant_id, + dataset_id=dataset_id, + type="string", + name="author", + created_by=user_id, + ) + ) + self.orm_session.commit() + with ( + patch.object(retrieval, "_fetch_model_config", return_value=(Mock(), Mock())), + patch.object(retrieval, "_get_prompt_template", return_value=(["prompt"], [])), + patch.object(retrieval, "_record_usage"), + ): + model_instance = Mock() + model_instance.invoke_llm.side_effect = RuntimeError("nope") + with patch.object(retrieval, "_fetch_model_config", return_value=(model_instance, Mock())): + assert ( + retrieval._automatic_metadata_filter_func( + self.orm_session, + dataset_ids=[dataset_id], + query="python", + tenant_id=tenant_id, + user_id=user_id, + metadata_model_config=WorkflowModelConfig(provider="openai", name="gpt", mode="chat"), + ) + is None + ) + + with ( + patch("core.rag.retrieval.dataset_retrieval.ModelMode", return_value=object()), + patch("core.rag.retrieval.dataset_retrieval.AdvancedPromptTransform"), + ): + with pytest.raises(ValueError, match="not support"): + retrieval._get_prompt_template( + model_config=ModelConfigWithCredentialsEntity.model_construct( + provider="openai", + model="gpt", + model_schema=Mock(), + mode="chat", + provider_model_bundle=Mock(), + credentials={}, + parameters={}, + stop=[], + ), + mode="chat", + metadata_fields=[], + query="q", + ) + + +"""Focused tests for attachment-aware dataset retrieval entry behavior.""" + + +class TestDatasetRetrievalAttachmentEntryCases: + @staticmethod + def test_knowledge_retrieval_allows_attachment_only_requests(unbound_session: Session) -> None: + retrieval = DatasetRetrieval() + available_dataset = Dataset( + id="dataset-1", + tenant_id="tenant-1", + name="Dataset", + created_by="user-1", ) + + request = KnowledgeRetrievalRequest( + tenant_id="tenant-1", + user_id="user-1", + app_id="app-1", + user_from="account", + dataset_ids=["dataset-1"], + query=None, + retrieval_mode="multiple", + top_k=4, + score_threshold=0.0, + reranking_mode="reranking_model", + reranking_enable=True, + attachment_ids=["attachment-1"], + ) + with ( - patch( - "core.tools.utils.dataset_retriever.dataset_retriever_tool.DatasetRetrieverTool.from_dataset", - return_value="single-tool", - ) as mock_single_tool, + patch.object(retrieval, "_check_knowledge_rate_limit"), + patch.object(retrieval, "_get_available_datasets", return_value=[available_dataset]), + patch.object(retrieval, "multiple_retrieve", return_value=[]) as mock_multiple, ): - single_tools = retrieval.to_dataset_retriever_tool( - session=self.orm_session, - tenant_id=tenant_id, - dataset_ids=[missing_id, skipped_id, available_id], - retrieve_config=single_config, - return_resource=True, + result = retrieval.knowledge_retrieval(unbound_session, request) + + assert result == [] + mock_multiple.assert_called_once() + + +"""SQLite-backed tests for dataset availability, rate limiting, and retrieval orchestration.""" + + +@dataclass(frozen=True) +class RetrievalDatabase: + session_maker: sessionmaker[Session] + + +def _persist_dataset( + database: RetrievalDatabase, + *, + dataset_id: str, + tenant_id: str = "tenant-1", + provider: str = "vendor", +) -> Dataset: + dataset = Dataset( + id=dataset_id, + tenant_id=tenant_id, + name=f"Dataset {dataset_id}", + created_by="user-1", + provider=provider, + indexing_technique=IndexTechniqueType.HIGH_QUALITY, + embedding_model="text-embedding-ada-002", + embedding_model_provider="openai", + retrieval_model={ + "search_method": "semantic_search", + "reranking_enable": False, + "top_k": 4, + "score_threshold_enabled": False, + }, + ) + with database.session_maker.begin() as session: + session.add(dataset) + return dataset + + +def _persist_document( + database: RetrievalDatabase, + *, + dataset_id: str, + document_id: str, + tenant_id: str = "tenant-1", + indexing_status: IndexingStatus = IndexingStatus.COMPLETED, + enabled: bool = True, + archived: bool = False, +) -> DatasetDocument: + document = DatasetDocument( + id=document_id, + tenant_id=tenant_id, + dataset_id=dataset_id, + position=1, + data_source_type=DataSourceType.UPLOAD_FILE, + batch="batch-1", + name=f"Document {document_id}", + created_from=DocumentCreatedFrom.API, + created_by="user-1", + indexing_status=indexing_status, + enabled=enabled, + archived=archived, + ) + with database.session_maker.begin() as session: + session.add(document) + return document + + +def _persist_segment( + database: RetrievalDatabase, + *, + dataset_id: str, + document_id: str, +) -> DocumentSegment: + segment = DocumentSegment( + tenant_id="tenant-1", + dataset_id=dataset_id, + document_id=document_id, + position=1, + content="Python is great", + word_count=3, + tokens=3, + created_by="user-1", + index_node_id="node-1", + index_node_hash="hash-1", + hit_count=5, + ) + with database.session_maker.begin() as session: + session.add(segment) + return segment + + +def _persist_available_dataset( + database: RetrievalDatabase, + *, + dataset_id: str = "dataset-1", + document_id: str = "document-1", +) -> tuple[Dataset, DatasetDocument]: + return ( + _persist_dataset(database, dataset_id=dataset_id), + _persist_document(database, dataset_id=dataset_id, document_id=document_id), + ) + + +def _request( + *, + dataset_ids: list[str], + retrieval_mode: str = "multiple", + metadata_filtering_mode: str = "disabled", +) -> KnowledgeRetrievalRequest: + return KnowledgeRetrievalRequest( + tenant_id="tenant-1", + user_id="user-1", + app_id="app-1", + user_from="web", + dataset_ids=dataset_ids, + query="What is Python?", + retrieval_mode=retrieval_mode, + metadata_filtering_mode=metadata_filtering_mode, + top_k=5, + score_threshold=0.7, + reranking_enable=True, + reranking_mode="reranking_model", + reranking_model={"reranking_provider_name": "cohere", "reranking_model_name": "rerank-v2"}, + ) + + +def _rag_document( + content: str, + doc_id: str, + *, + score: float = 0.8, + provider: str = "dify", + additional_metadata: dict[str, object] | None = None, +) -> _dataset_retrieval_methods_Document: + metadata: dict[str, object] = { + "doc_id": doc_id, + "document_id": "document-1", + "dataset_id": "dataset-1", + "score": score, + } + if additional_metadata: + metadata.update(additional_metadata) + return _dataset_retrieval_methods_Document(page_content=content, metadata=metadata, provider=provider) + + +def _patch_rate_limit( + monkeypatch: pytest.MonkeyPatch, + *, + enabled: bool, + request_count: int = 0, +) -> MagicMock: + limit = SimpleNamespace(enabled=enabled, limit=100, subscription_plan="professional") + monkeypatch.setattr( + retrieval_module.FeatureService, + "get_knowledge_rate_limit", + MagicMock(return_value=limit), + ) + redis = MagicMock() + redis.zcard.return_value = request_count + monkeypatch.setattr(retrieval_module, "redis_client", redis) + monkeypatch.setattr(retrieval_module.time, "time", lambda: 1234567890) + return redis + + +class TestDatasetRetrievalMethodsCases: + @staticmethod + @pytest.fixture + def retrieval_database( + sqlite_session_factory: sessionmaker[Session], monkeypatch: pytest.MonkeyPatch + ) -> RetrievalDatabase: + """Bind every retrieval-owned session to a disposable SQLite database.""" + monkeypatch.setattr(retrieval_module.session_factory, "create_session", sqlite_session_factory) + return RetrievalDatabase(session_maker=sqlite_session_factory) + + class TestCheckKnowledgeRateLimit: + def test_rate_limit_disabled_performs_no_redis_or_database_work( + self, + monkeypatch: pytest.MonkeyPatch, + retrieval_database: RetrievalDatabase, + ) -> None: + redis = _patch_rate_limit(monkeypatch, enabled=False) + + DatasetRetrieval()._check_knowledge_rate_limit("tenant-1") + + redis.zadd.assert_not_called() + with retrieval_database.session_maker() as session: + assert session.scalar(select(RateLimitLog)) is None + + def test_rate_limit_enabled_not_exceeded_tracks_request_without_log( + self, + monkeypatch: pytest.MonkeyPatch, + retrieval_database: RetrievalDatabase, + ) -> None: + redis = _patch_rate_limit(monkeypatch, enabled=True, request_count=50) + + DatasetRetrieval()._check_knowledge_rate_limit("tenant-1") + + current_time = 1234567890000 + redis.zadd.assert_called_once_with("rate_limit_tenant-1", {current_time: current_time}) + redis.zremrangebyscore.assert_called_once_with("rate_limit_tenant-1", 0, current_time - 60000) + with retrieval_database.session_maker() as session: + assert session.scalar(select(RateLimitLog)) is None + + def test_rate_limit_exceeded_commits_audit_log_before_raising( + self, + monkeypatch: pytest.MonkeyPatch, + retrieval_database: RetrievalDatabase, + ) -> None: + _patch_rate_limit(monkeypatch, enabled=True, request_count=150) + + with pytest.raises(exc.RateLimitExceededError, match="knowledge base request rate limit"): + DatasetRetrieval()._check_knowledge_rate_limit("tenant-1") + + with retrieval_database.session_maker() as session: + logs = session.scalars(select(RateLimitLog)).all() + assert len(logs) == 1 + assert logs[0].tenant_id == "tenant-1" + assert logs[0].subscription_plan == "professional" + assert logs[0].operation == "knowledge" + + class TestGetAvailableDatasets: + def test_returns_completed_or_external_datasets_with_tenant_scope( + self, + retrieval_database: RetrievalDatabase, + ) -> None: + _persist_available_dataset(retrieval_database, dataset_id="available", document_id="available-doc") + _persist_dataset(retrieval_database, dataset_id="disabled") + _persist_document( + retrieval_database, + dataset_id="disabled", + document_id="disabled-doc", + enabled=False, + ) + _persist_dataset(retrieval_database, dataset_id="archived") + _persist_document( + retrieval_database, + dataset_id="archived", + document_id="archived-doc", + archived=True, + ) + _persist_dataset(retrieval_database, dataset_id="waiting") + _persist_document( + retrieval_database, + dataset_id="waiting", + document_id="waiting-doc", + indexing_status=IndexingStatus.WAITING, + ) + _persist_dataset(retrieval_database, dataset_id="external", provider="external") + _persist_dataset(retrieval_database, dataset_id="other-tenant", tenant_id="tenant-2") + _persist_document( + retrieval_database, + dataset_id="other-tenant", + document_id="other-doc", + tenant_id="tenant-2", + ) + + datasets = DatasetRetrieval()._get_available_datasets( + "tenant-1", + ["available", "disabled", "archived", "waiting", "external", "other-tenant"], + ) + + assert {dataset.id for dataset in datasets} == {"available", "external"} + + def test_returns_empty_for_vendor_dataset_without_documents( + self, + retrieval_database: RetrievalDatabase, + ) -> None: + _persist_dataset(retrieval_database, dataset_id="empty") + + assert DatasetRetrieval()._get_available_datasets("tenant-1", ["empty"]) == [] + + class TestDatasetRetrievalKnowledgeRetrieval: + def test_single_mode_request_shape(self) -> None: + request = KnowledgeRetrievalRequest( + tenant_id="tenant-1", + user_id="user-1", + app_id="app-1", + user_from="web", + dataset_ids=["dataset-1"], + query="What is Python?", + retrieval_mode="single", + model_provider="openai", + model_name="gpt-4", + model_mode="chat", + completion_params={"temperature": 0.7}, + ) + + assert request.retrieval_mode == "single" + assert request.model_provider == "openai" + assert request.model_name == "gpt-4" + + def test_multiple_mode_formats_persisted_dataset_and_document( + self, + monkeypatch: pytest.MonkeyPatch, + retrieval_database: RetrievalDatabase, + ) -> None: + dataset, document = _persist_available_dataset(retrieval_database) + segment = _persist_segment(retrieval_database, dataset_id=dataset.id, document_id=document.id) + retrieval = DatasetRetrieval() + monkeypatch.setattr(retrieval, "_check_knowledge_rate_limit", MagicMock()) + monkeypatch.setattr( + retrieval, "multiple_retrieve", MagicMock(return_value=[_rag_document("Python", "node-1")]) + ) + record = SimpleNamespace(segment=segment, score=0.9, child_chunks=[], summary=None, files=None) + monkeypatch.setattr( + retrieval_module.RetrievalService, + "format_retrieval_documents", + MagicMock(return_value=[record]), + ) + scope = FileAccessScope( + tenant_id=dataset.tenant_id, + user_id="visitor", + user_from=UserFrom.END_USER, invoke_from=InvokeFrom.WEB_APP, - hit_callback=Mock(), - user_id="user-1", - inputs={"k": "v"}, + ) + with bind_file_access_scope(scope), retrieval_database.session_maker() as caller_session: + result = retrieval.knowledge_retrieval(caller_session, _request(dataset_ids=[dataset.id])) + granted_scope = get_current_file_access_scope() + + assert len(result) == 1 + assert result[0].metadata.dataset_id == dataset.id + assert result[0].metadata.document_id == document.id + assert result[0].title == document.name + assert granted_scope is not None + assert granted_scope.granted_retriever_segment_ids == frozenset([segment.id]) + + def test_metadata_filtering_disabled_skips_filter_builder( + self, + monkeypatch: pytest.MonkeyPatch, + retrieval_database: RetrievalDatabase, + ) -> None: + _persist_available_dataset(retrieval_database) + retrieval = DatasetRetrieval() + monkeypatch.setattr(retrieval, "_check_knowledge_rate_limit", MagicMock()) + metadata_filter = MagicMock(return_value=(None, None)) + monkeypatch.setattr(retrieval, "get_metadata_filter_condition", metadata_filter) + monkeypatch.setattr(retrieval, "multiple_retrieve", MagicMock(return_value=[])) + + with retrieval_database.session_maker() as caller_session: + result = retrieval.knowledge_retrieval(caller_session, _request(dataset_ids=["dataset-1"])) + + assert result == [] + metadata_filter.assert_not_called() + + def test_external_documents_are_formatted_without_database_document( + self, + monkeypatch: pytest.MonkeyPatch, + retrieval_database: RetrievalDatabase, + ) -> None: + _persist_dataset(retrieval_database, dataset_id="external", provider="external") + retrieval = DatasetRetrieval() + monkeypatch.setattr(retrieval, "_check_knowledge_rate_limit", MagicMock()) + external_document = _rag_document( + "External knowledge", + "external-node", + score=0.9, + provider="external", + additional_metadata={ + "dataset_id": "external", + "dataset_name": "External Dataset", + "document_id": "external-document", + "title": "External Document", + }, + ) + monkeypatch.setattr(retrieval, "multiple_retrieve", MagicMock(return_value=[external_document])) + + with retrieval_database.session_maker() as caller_session: + result = retrieval.knowledge_retrieval(caller_session, _request(dataset_ids=["external"])) + + assert len(result) == 1 + assert result[0].metadata.data_source_type == "external" + assert result[0].metadata.dataset_id == "external" + + def test_empty_retrieval_results_return_empty_list( + self, + monkeypatch: pytest.MonkeyPatch, + retrieval_database: RetrievalDatabase, + ) -> None: + _persist_available_dataset(retrieval_database) + retrieval = DatasetRetrieval() + monkeypatch.setattr(retrieval, "_check_knowledge_rate_limit", MagicMock()) + monkeypatch.setattr(retrieval, "multiple_retrieve", MagicMock(return_value=[])) + + with retrieval_database.session_maker() as caller_session: + result = retrieval.knowledge_retrieval(caller_session, _request(dataset_ids=["dataset-1"])) + + assert result == [] + + def test_rate_limit_exception_stops_retrieval( + self, + monkeypatch: pytest.MonkeyPatch, + retrieval_database: RetrievalDatabase, + ) -> None: + retrieval = DatasetRetrieval() + monkeypatch.setattr( + retrieval, + "_check_knowledge_rate_limit", + MagicMock(side_effect=exc.RateLimitExceededError("Rate limit exceeded")), ) - assert single_tools == ["single-tool"] - mock_single_tool.assert_called_once() + with retrieval_database.session_maker() as caller_session: + with pytest.raises(exc.RateLimitExceededError): + retrieval.knowledge_retrieval(caller_session, _request(dataset_ids=["dataset-1"])) - multiple_config_missing = DatasetRetrieveConfigEntity( - retrieve_strategy=DatasetRetrieveConfigEntity.RetrieveStrategy.MULTIPLE, - metadata_filtering_mode="disabled", - reranking_model=None, + def test_no_available_datasets_skips_retrieval( + self, + monkeypatch: pytest.MonkeyPatch, + retrieval_database: RetrievalDatabase, + ) -> None: + _persist_dataset(retrieval_database, dataset_id="empty") + retrieval = DatasetRetrieval() + monkeypatch.setattr(retrieval, "_check_knowledge_rate_limit", MagicMock()) + multiple_retrieve = MagicMock() + monkeypatch.setattr(retrieval, "multiple_retrieve", multiple_retrieve) + + with retrieval_database.session_maker() as caller_session: + result = retrieval.knowledge_retrieval(caller_session, _request(dataset_ids=["empty"])) + + assert result == [] + multiple_retrieve.assert_not_called() + + def test_document_scores_sort_descending(self) -> None: + documents = [ + _rag_document("Low", "doc1", score=0.6), + _rag_document("High", "doc2", score=0.95), + _rag_document("Medium", "doc3", score=0.8), + ] + + sorted_documents = sorted(documents, key=lambda document: document.metadata["score"], reverse=True) + + assert [document.metadata["score"] for document in sorted_documents] == [0.95, 0.8, 0.6] + + +"""Exercise the SQL emitted when workflow retrieval assembles source citations.""" + + +class TestRetrievalSourceQueryCases: + @staticmethod + def test_source_citations_do_not_fetch_unused_document_fields(sqlite_session: Session) -> None: + tenant_id, creator_id, dataset_id, document_id = (str(uuid4()) for _ in range(4)) + dataset = Dataset(id=dataset_id, tenant_id=tenant_id, name="Knowledge", created_by=creator_id) + document = _retrieval_source_query_Document( + id=document_id, + tenant_id=tenant_id, + dataset_id=dataset_id, + position=1, + data_source_type=DataSourceType("upload_file"), + batch="batch-1", + name="source.txt", + created_from=DocumentCreatedFrom("web"), + created_by=creator_id, + enabled=True, + archived=False, + doc_metadata={"category": "transport"}, + data_source_info="x" * 1_000_000, + error="y" * 1_000_000, ) - with pytest.raises(ValueError, match="Reranking model is required"): - retrieval.to_dataset_retriever_tool( - session=self.orm_session, - tenant_id=tenant_id, - dataset_ids=[available_id], - retrieve_config=multiple_config_missing, - return_resource=True, - invoke_from=InvokeFrom.WEB_APP, - hit_callback=Mock(), - user_id="user-1", - inputs={}, - ) - - multiple_config = DatasetRetrieveConfigEntity( - retrieve_strategy=DatasetRetrieveConfigEntity.RetrieveStrategy.MULTIPLE, - metadata_filtering_mode="disabled", - top_k=3, - score_threshold=0.2, - reranking_model={"reranking_provider_name": "cohere", "reranking_model_name": "rerank-v3"}, - ) - with ( - patch( - "core.tools.utils.dataset_retriever.dataset_multi_retriever_tool.DatasetMultiRetrieverTool.from_dataset", - return_value="multi-tool", - ) as mock_multi_tool, - ): - multi_tools = retrieval.to_dataset_retriever_tool( - session=self.orm_session, - tenant_id=tenant_id, - dataset_ids=[available_id], - retrieve_config=multiple_config, - return_resource=False, - invoke_from=InvokeFrom.DEBUGGER, - hit_callback=Mock(), - user_id="user-1", - inputs={}, - ) - assert multi_tools == ["multi-tool"] - mock_multi_tool.assert_called_once() - - def test_additional_small_branches(self, retrieval: DatasetRetrieval) -> None: - keyword_handler = Mock() - keyword_handler.extract_keywords.side_effect = [[], []] - doc = Document(page_content="doc", metadata={"doc_id": "1"}, provider="dify") - with patch("core.rag.retrieval.dataset_retrieval.JiebaKeywordTableHandler", return_value=keyword_handler): - ranked = retrieval.calculate_keyword_score("query", [doc], top_k=1) - assert len(ranked) == 1 - assert ranked[0].metadata.get("score") == 0.0 - - with pytest.raises(ValueError, match="metadata_model_config is required"): - retrieval._automatic_metadata_filter_func( - self.orm_session, - dataset_ids=[str(uuid4())], - query="python", - tenant_id=str(uuid4()), - user_id=str(uuid4()), - metadata_model_config=None, # type: ignore[arg-type] - ) - - tenant_id = str(uuid4()) - dataset_id = str(uuid4()) - user_id = str(uuid4()) - self.orm_session.add( - DatasetMetadata( + segments = [ + DocumentSegment( tenant_id=tenant_id, dataset_id=dataset_id, - type="string", - name="author", - created_by=user_id, + document_id=document_id, + position=position, + content=f"content-{position}", + word_count=10, + tokens=3, + created_by=creator_id, + enabled=True, + index_node_id=f"node-{position}", + index_node_hash=f"hash-{position}", ) + for position in (1, 2) + ] + sqlite_session.add_all([dataset, document, *segments]) + sqlite_session.commit() + records = [ + RetrievalSegments(segment=segment, score=score) for segment, score in zip(segments, (0.8, 0.9), strict=True) + ] + retrieved = [RetrievedDocument(provider="dify", page_content="content", metadata={})] + request = KnowledgeRetrievalRequest( + tenant_id=tenant_id, + user_id=creator_id, + app_id=str(uuid4()), + user_from="web", + dataset_ids=[dataset_id], + query="小动物如何运输", + retrieval_mode="multiple", + metadata_filtering_mode="disabled", ) - self.orm_session.commit() - with ( - patch.object(retrieval, "_fetch_model_config", return_value=(Mock(), Mock())), - patch.object(retrieval, "_get_prompt_template", return_value=(["prompt"], [])), - patch.object(retrieval, "_record_usage"), - ): - model_instance = Mock() - model_instance.invoke_llm.side_effect = RuntimeError("nope") - with patch.object(retrieval, "_fetch_model_config", return_value=(model_instance, Mock())): - assert ( - retrieval._automatic_metadata_filter_func( - self.orm_session, - dataset_ids=[dataset_id], - query="python", - tenant_id=tenant_id, - user_id=user_id, - metadata_model_config=WorkflowModelConfig(provider="openai", name="gpt", mode="chat"), - ) - is None - ) + queries: list[str] = [] - with ( - patch("core.rag.retrieval.dataset_retrieval.ModelMode", return_value=object()), - patch("core.rag.retrieval.dataset_retrieval.AdvancedPromptTransform"), - ): - with pytest.raises(ValueError, match="not support"): - retrieval._get_prompt_template( - model_config=ModelConfigWithCredentialsEntity.model_construct( - provider="openai", - model="gpt", - model_schema=Mock(), - mode="chat", - provider_model_bundle=Mock(), - credentials={}, - parameters={}, - stop=[], - ), - mode="chat", - metadata_fields=[], - query="q", + def capture_query(_conn, _cursor, statement: str, _parameters, _context, _executemany) -> None: + if "FROM documents" in statement: + queries.append(statement) + + engine = sqlite_session.get_bind() + retrieval = DatasetRetrieval() + event.listen(engine, "before_cursor_execute", capture_query) + try: + with ( + patch.object(retrieval, "_check_knowledge_rate_limit"), + patch.object(retrieval, "_get_available_datasets", return_value=[dataset]), + patch.object(retrieval, "multiple_retrieve", return_value=retrieved), + patch( + "core.rag.retrieval.dataset_retrieval.RetrievalService.format_retrieval_documents", + return_value=records, + ), + ): + sources = retrieval.knowledge_retrieval(sqlite_session, request) + finally: + event.remove(engine, "before_cursor_execute", capture_query) + + assert len(sources) == 2 + assert [source.content for source in sources] == ["content-2", "content-1"] + for source in sources: + assert source.title == "source.txt" + assert source.metadata.document_id == document_id + assert source.metadata.dataset_name == "Knowledge" + assert source.metadata.data_source_type == "upload_file" + assert source.metadata.doc_metadata == {"category": "transport"} + assert len(queries) == 1, "Citation fields must not cause additional lazy SQL queries" + assert "documents.data_source_info" not in queries[0] + assert "documents.error" not in queries[0] + assert "documents.indexing_latency" not in queries[0] + + @staticmethod + def test_source_document_read_is_scoped_and_has_no_lazy_queries(sqlite_session: Session) -> None: + tenant_id, foreign_tenant_id, dataset_id, foreign_dataset_id = (str(uuid4()) for _ in range(4)) + documents = [ + _retrieval_source_query_Document( + id=str(uuid4()), + tenant_id=owner, + dataset_id=dataset, + position=1, + data_source_type=DataSourceType("upload_file"), + batch="batch-1", + name="source.txt", + created_from=DocumentCreatedFrom("web"), + created_by=str(uuid4()), + doc_metadata=None, + ) + for owner, dataset in ( + (tenant_id, dataset_id), + (foreign_tenant_id, dataset_id), + (tenant_id, foreign_dataset_id), + ) + ] + sqlite_session.add_all(documents) + sqlite_session.flush() + document_ids = [document.id for document in documents] + sqlite_session.expunge_all() + queries: list[str] = [] + + def capture_query(_conn, _cursor, statement: str, _parameters, _context, _executemany) -> None: + queries.append(statement) + + engine = sqlite_session.get_bind() + event.listen(engine, "before_cursor_execute", capture_query) + try: + result = get_retrieval_source_documents( + tenant_id=tenant_id, + dataset_ids=[dataset_id], + document_ids=[*document_ids, document_ids[0]], + session=sqlite_session, + ) + assert [document.id for document in result] == [document_ids[0]] + assert result[0].name == "source.txt" + assert result[0].data_source_type == "upload_file" + assert result[0].doc_metadata is None + assert set(inspect(result[0]).dict) - {"_sa_instance_state"} == { + "id", + "tenant_id", + "dataset_id", + "name", + "data_source_type", + "doc_metadata", + } + with pytest.raises(InvalidRequestError, match="raiseload"): + _ = result[0].data_source_info + assert ( + get_retrieval_source_documents( + tenant_id=tenant_id, dataset_ids=[], document_ids=document_ids, session=sqlite_session ) + == [] + ) + assert ( + get_retrieval_source_documents( + tenant_id=tenant_id, dataset_ids=[dataset_id], document_ids=[], session=sqlite_session + ) + == [] + ) + finally: + event.remove(engine, "before_cursor_execute", capture_query) + assert len(queries) == 1 diff --git a/api/tests/unit_tests/core/rag/retrieval/test_dataset_retrieval_attachment_entry.py b/api/tests/unit_tests/core/rag/retrieval/test_dataset_retrieval_attachment_entry.py deleted file mode 100644 index a5a91a01545..00000000000 --- a/api/tests/unit_tests/core/rag/retrieval/test_dataset_retrieval_attachment_entry.py +++ /dev/null @@ -1,44 +0,0 @@ -"""Focused tests for attachment-aware dataset retrieval entry behavior.""" - -from unittest.mock import patch - -from sqlalchemy.orm import Session - -from core.rag.retrieval.dataset_retrieval import DatasetRetrieval -from core.workflow.nodes.knowledge_retrieval.retrieval import KnowledgeRetrievalRequest -from models.dataset import Dataset - - -def test_knowledge_retrieval_allows_attachment_only_requests(unbound_session: Session) -> None: - retrieval = DatasetRetrieval() - available_dataset = Dataset( - id="dataset-1", - tenant_id="tenant-1", - name="Dataset", - created_by="user-1", - ) - - request = KnowledgeRetrievalRequest( - tenant_id="tenant-1", - user_id="user-1", - app_id="app-1", - user_from="account", - dataset_ids=["dataset-1"], - query=None, - retrieval_mode="multiple", - top_k=4, - score_threshold=0.0, - reranking_mode="reranking_model", - reranking_enable=True, - attachment_ids=["attachment-1"], - ) - - with ( - patch.object(retrieval, "_check_knowledge_rate_limit"), - patch.object(retrieval, "_get_available_datasets", return_value=[available_dataset]), - patch.object(retrieval, "multiple_retrieve", return_value=[]) as mock_multiple, - ): - result = retrieval.knowledge_retrieval(unbound_session, request) - - assert result == [] - mock_multiple.assert_called_once() diff --git a/api/tests/unit_tests/core/rag/retrieval/test_dataset_retrieval_methods.py b/api/tests/unit_tests/core/rag/retrieval/test_dataset_retrieval_methods.py deleted file mode 100644 index 01bc89c2831..00000000000 --- a/api/tests/unit_tests/core/rag/retrieval/test_dataset_retrieval_methods.py +++ /dev/null @@ -1,448 +0,0 @@ -"""SQLite-backed tests for dataset availability, rate limiting, and retrieval orchestration.""" - -from dataclasses import dataclass -from types import SimpleNamespace -from unittest.mock import MagicMock - -import pytest -from sqlalchemy import select -from sqlalchemy.orm import Session, sessionmaker - -from core.app.entities.app_invoke_entities import InvokeFrom, UserFrom -from core.app.file_access import FileAccessScope, bind_file_access_scope, get_current_file_access_scope -from core.rag.index_processor.constant.index_type import IndexTechniqueType -from core.rag.models.document import Document -from core.rag.retrieval import dataset_retrieval as retrieval_module -from core.rag.retrieval.dataset_retrieval import DatasetRetrieval -from core.workflow.nodes.knowledge_retrieval import exc -from core.workflow.nodes.knowledge_retrieval.retrieval import KnowledgeRetrievalRequest -from models.dataset import Dataset, DocumentSegment, RateLimitLog -from models.dataset import Document as DatasetDocument -from models.enums import DataSourceType, DocumentCreatedFrom, IndexingStatus - - -@dataclass(frozen=True) -class RetrievalDatabase: - session_maker: sessionmaker[Session] - - -@pytest.fixture -def retrieval_database( - sqlite_session_factory: sessionmaker[Session], monkeypatch: pytest.MonkeyPatch -) -> RetrievalDatabase: - """Bind every retrieval-owned session to a disposable SQLite database.""" - monkeypatch.setattr(retrieval_module.session_factory, "create_session", sqlite_session_factory) - return RetrievalDatabase(session_maker=sqlite_session_factory) - - -def _persist_dataset( - database: RetrievalDatabase, - *, - dataset_id: str, - tenant_id: str = "tenant-1", - provider: str = "vendor", -) -> Dataset: - dataset = Dataset( - id=dataset_id, - tenant_id=tenant_id, - name=f"Dataset {dataset_id}", - created_by="user-1", - provider=provider, - indexing_technique=IndexTechniqueType.HIGH_QUALITY, - embedding_model="text-embedding-ada-002", - embedding_model_provider="openai", - retrieval_model={ - "search_method": "semantic_search", - "reranking_enable": False, - "top_k": 4, - "score_threshold_enabled": False, - }, - ) - with database.session_maker.begin() as session: - session.add(dataset) - return dataset - - -def _persist_document( - database: RetrievalDatabase, - *, - dataset_id: str, - document_id: str, - tenant_id: str = "tenant-1", - indexing_status: IndexingStatus = IndexingStatus.COMPLETED, - enabled: bool = True, - archived: bool = False, -) -> DatasetDocument: - document = DatasetDocument( - id=document_id, - tenant_id=tenant_id, - dataset_id=dataset_id, - position=1, - data_source_type=DataSourceType.UPLOAD_FILE, - batch="batch-1", - name=f"Document {document_id}", - created_from=DocumentCreatedFrom.API, - created_by="user-1", - indexing_status=indexing_status, - enabled=enabled, - archived=archived, - ) - with database.session_maker.begin() as session: - session.add(document) - return document - - -def _persist_segment( - database: RetrievalDatabase, - *, - dataset_id: str, - document_id: str, -) -> DocumentSegment: - segment = DocumentSegment( - tenant_id="tenant-1", - dataset_id=dataset_id, - document_id=document_id, - position=1, - content="Python is great", - word_count=3, - tokens=3, - created_by="user-1", - index_node_id="node-1", - index_node_hash="hash-1", - hit_count=5, - ) - with database.session_maker.begin() as session: - session.add(segment) - return segment - - -def _persist_available_dataset( - database: RetrievalDatabase, - *, - dataset_id: str = "dataset-1", - document_id: str = "document-1", -) -> tuple[Dataset, DatasetDocument]: - return ( - _persist_dataset(database, dataset_id=dataset_id), - _persist_document(database, dataset_id=dataset_id, document_id=document_id), - ) - - -def _request( - *, - dataset_ids: list[str], - retrieval_mode: str = "multiple", - metadata_filtering_mode: str = "disabled", -) -> KnowledgeRetrievalRequest: - return KnowledgeRetrievalRequest( - tenant_id="tenant-1", - user_id="user-1", - app_id="app-1", - user_from="web", - dataset_ids=dataset_ids, - query="What is Python?", - retrieval_mode=retrieval_mode, - metadata_filtering_mode=metadata_filtering_mode, - top_k=5, - score_threshold=0.7, - reranking_enable=True, - reranking_mode="reranking_model", - reranking_model={"reranking_provider_name": "cohere", "reranking_model_name": "rerank-v2"}, - ) - - -def _rag_document( - content: str, - doc_id: str, - *, - score: float = 0.8, - provider: str = "dify", - additional_metadata: dict[str, object] | None = None, -) -> Document: - metadata: dict[str, object] = { - "doc_id": doc_id, - "document_id": "document-1", - "dataset_id": "dataset-1", - "score": score, - } - if additional_metadata: - metadata.update(additional_metadata) - return Document(page_content=content, metadata=metadata, provider=provider) - - -def _patch_rate_limit( - monkeypatch: pytest.MonkeyPatch, - *, - enabled: bool, - request_count: int = 0, -) -> MagicMock: - limit = SimpleNamespace(enabled=enabled, limit=100, subscription_plan="professional") - monkeypatch.setattr( - retrieval_module.FeatureService, - "get_knowledge_rate_limit", - MagicMock(return_value=limit), - ) - redis = MagicMock() - redis.zcard.return_value = request_count - monkeypatch.setattr(retrieval_module, "redis_client", redis) - monkeypatch.setattr(retrieval_module.time, "time", lambda: 1234567890) - return redis - - -class TestCheckKnowledgeRateLimit: - def test_rate_limit_disabled_performs_no_redis_or_database_work( - self, - monkeypatch: pytest.MonkeyPatch, - retrieval_database: RetrievalDatabase, - ) -> None: - redis = _patch_rate_limit(monkeypatch, enabled=False) - - DatasetRetrieval()._check_knowledge_rate_limit("tenant-1") - - redis.zadd.assert_not_called() - with retrieval_database.session_maker() as session: - assert session.scalar(select(RateLimitLog)) is None - - def test_rate_limit_enabled_not_exceeded_tracks_request_without_log( - self, - monkeypatch: pytest.MonkeyPatch, - retrieval_database: RetrievalDatabase, - ) -> None: - redis = _patch_rate_limit(monkeypatch, enabled=True, request_count=50) - - DatasetRetrieval()._check_knowledge_rate_limit("tenant-1") - - current_time = 1234567890000 - redis.zadd.assert_called_once_with("rate_limit_tenant-1", {current_time: current_time}) - redis.zremrangebyscore.assert_called_once_with("rate_limit_tenant-1", 0, current_time - 60000) - with retrieval_database.session_maker() as session: - assert session.scalar(select(RateLimitLog)) is None - - def test_rate_limit_exceeded_commits_audit_log_before_raising( - self, - monkeypatch: pytest.MonkeyPatch, - retrieval_database: RetrievalDatabase, - ) -> None: - _patch_rate_limit(monkeypatch, enabled=True, request_count=150) - - with pytest.raises(exc.RateLimitExceededError, match="knowledge base request rate limit"): - DatasetRetrieval()._check_knowledge_rate_limit("tenant-1") - - with retrieval_database.session_maker() as session: - logs = session.scalars(select(RateLimitLog)).all() - assert len(logs) == 1 - assert logs[0].tenant_id == "tenant-1" - assert logs[0].subscription_plan == "professional" - assert logs[0].operation == "knowledge" - - -class TestGetAvailableDatasets: - def test_returns_completed_or_external_datasets_with_tenant_scope( - self, - retrieval_database: RetrievalDatabase, - ) -> None: - _persist_available_dataset(retrieval_database, dataset_id="available", document_id="available-doc") - _persist_dataset(retrieval_database, dataset_id="disabled") - _persist_document( - retrieval_database, - dataset_id="disabled", - document_id="disabled-doc", - enabled=False, - ) - _persist_dataset(retrieval_database, dataset_id="archived") - _persist_document( - retrieval_database, - dataset_id="archived", - document_id="archived-doc", - archived=True, - ) - _persist_dataset(retrieval_database, dataset_id="waiting") - _persist_document( - retrieval_database, - dataset_id="waiting", - document_id="waiting-doc", - indexing_status=IndexingStatus.WAITING, - ) - _persist_dataset(retrieval_database, dataset_id="external", provider="external") - _persist_dataset(retrieval_database, dataset_id="other-tenant", tenant_id="tenant-2") - _persist_document( - retrieval_database, - dataset_id="other-tenant", - document_id="other-doc", - tenant_id="tenant-2", - ) - - datasets = DatasetRetrieval()._get_available_datasets( - "tenant-1", - ["available", "disabled", "archived", "waiting", "external", "other-tenant"], - ) - - assert {dataset.id for dataset in datasets} == {"available", "external"} - - def test_returns_empty_for_vendor_dataset_without_documents( - self, - retrieval_database: RetrievalDatabase, - ) -> None: - _persist_dataset(retrieval_database, dataset_id="empty") - - assert DatasetRetrieval()._get_available_datasets("tenant-1", ["empty"]) == [] - - -class TestDatasetRetrievalKnowledgeRetrieval: - def test_single_mode_request_shape(self) -> None: - request = KnowledgeRetrievalRequest( - tenant_id="tenant-1", - user_id="user-1", - app_id="app-1", - user_from="web", - dataset_ids=["dataset-1"], - query="What is Python?", - retrieval_mode="single", - model_provider="openai", - model_name="gpt-4", - model_mode="chat", - completion_params={"temperature": 0.7}, - ) - - assert request.retrieval_mode == "single" - assert request.model_provider == "openai" - assert request.model_name == "gpt-4" - - def test_multiple_mode_formats_persisted_dataset_and_document( - self, - monkeypatch: pytest.MonkeyPatch, - retrieval_database: RetrievalDatabase, - ) -> None: - dataset, document = _persist_available_dataset(retrieval_database) - segment = _persist_segment(retrieval_database, dataset_id=dataset.id, document_id=document.id) - retrieval = DatasetRetrieval() - monkeypatch.setattr(retrieval, "_check_knowledge_rate_limit", MagicMock()) - monkeypatch.setattr(retrieval, "multiple_retrieve", MagicMock(return_value=[_rag_document("Python", "node-1")])) - record = SimpleNamespace(segment=segment, score=0.9, child_chunks=[], summary=None, files=None) - monkeypatch.setattr( - retrieval_module.RetrievalService, - "format_retrieval_documents", - MagicMock(return_value=[record]), - ) - scope = FileAccessScope( - tenant_id=dataset.tenant_id, - user_id="visitor", - user_from=UserFrom.END_USER, - invoke_from=InvokeFrom.WEB_APP, - ) - with bind_file_access_scope(scope), retrieval_database.session_maker() as caller_session: - result = retrieval.knowledge_retrieval(caller_session, _request(dataset_ids=[dataset.id])) - granted_scope = get_current_file_access_scope() - - assert len(result) == 1 - assert result[0].metadata.dataset_id == dataset.id - assert result[0].metadata.document_id == document.id - assert result[0].title == document.name - assert granted_scope is not None - assert granted_scope.granted_retriever_segment_ids == frozenset([segment.id]) - - def test_metadata_filtering_disabled_skips_filter_builder( - self, - monkeypatch: pytest.MonkeyPatch, - retrieval_database: RetrievalDatabase, - ) -> None: - _persist_available_dataset(retrieval_database) - retrieval = DatasetRetrieval() - monkeypatch.setattr(retrieval, "_check_knowledge_rate_limit", MagicMock()) - metadata_filter = MagicMock(return_value=(None, None)) - monkeypatch.setattr(retrieval, "get_metadata_filter_condition", metadata_filter) - monkeypatch.setattr(retrieval, "multiple_retrieve", MagicMock(return_value=[])) - - with retrieval_database.session_maker() as caller_session: - result = retrieval.knowledge_retrieval(caller_session, _request(dataset_ids=["dataset-1"])) - - assert result == [] - metadata_filter.assert_not_called() - - def test_external_documents_are_formatted_without_database_document( - self, - monkeypatch: pytest.MonkeyPatch, - retrieval_database: RetrievalDatabase, - ) -> None: - _persist_dataset(retrieval_database, dataset_id="external", provider="external") - retrieval = DatasetRetrieval() - monkeypatch.setattr(retrieval, "_check_knowledge_rate_limit", MagicMock()) - external_document = _rag_document( - "External knowledge", - "external-node", - score=0.9, - provider="external", - additional_metadata={ - "dataset_id": "external", - "dataset_name": "External Dataset", - "document_id": "external-document", - "title": "External Document", - }, - ) - monkeypatch.setattr(retrieval, "multiple_retrieve", MagicMock(return_value=[external_document])) - - with retrieval_database.session_maker() as caller_session: - result = retrieval.knowledge_retrieval(caller_session, _request(dataset_ids=["external"])) - - assert len(result) == 1 - assert result[0].metadata.data_source_type == "external" - assert result[0].metadata.dataset_id == "external" - - def test_empty_retrieval_results_return_empty_list( - self, - monkeypatch: pytest.MonkeyPatch, - retrieval_database: RetrievalDatabase, - ) -> None: - _persist_available_dataset(retrieval_database) - retrieval = DatasetRetrieval() - monkeypatch.setattr(retrieval, "_check_knowledge_rate_limit", MagicMock()) - monkeypatch.setattr(retrieval, "multiple_retrieve", MagicMock(return_value=[])) - - with retrieval_database.session_maker() as caller_session: - result = retrieval.knowledge_retrieval(caller_session, _request(dataset_ids=["dataset-1"])) - - assert result == [] - - def test_rate_limit_exception_stops_retrieval( - self, - monkeypatch: pytest.MonkeyPatch, - retrieval_database: RetrievalDatabase, - ) -> None: - retrieval = DatasetRetrieval() - monkeypatch.setattr( - retrieval, - "_check_knowledge_rate_limit", - MagicMock(side_effect=exc.RateLimitExceededError("Rate limit exceeded")), - ) - - with retrieval_database.session_maker() as caller_session: - with pytest.raises(exc.RateLimitExceededError): - retrieval.knowledge_retrieval(caller_session, _request(dataset_ids=["dataset-1"])) - - def test_no_available_datasets_skips_retrieval( - self, - monkeypatch: pytest.MonkeyPatch, - retrieval_database: RetrievalDatabase, - ) -> None: - _persist_dataset(retrieval_database, dataset_id="empty") - retrieval = DatasetRetrieval() - monkeypatch.setattr(retrieval, "_check_knowledge_rate_limit", MagicMock()) - multiple_retrieve = MagicMock() - monkeypatch.setattr(retrieval, "multiple_retrieve", multiple_retrieve) - - with retrieval_database.session_maker() as caller_session: - result = retrieval.knowledge_retrieval(caller_session, _request(dataset_ids=["empty"])) - - assert result == [] - multiple_retrieve.assert_not_called() - - def test_document_scores_sort_descending(self) -> None: - documents = [ - _rag_document("Low", "doc1", score=0.6), - _rag_document("High", "doc2", score=0.95), - _rag_document("Medium", "doc3", score=0.8), - ] - - sorted_documents = sorted(documents, key=lambda document: document.metadata["score"], reverse=True) - - assert [document.metadata["score"] for document in sorted_documents] == [0.95, 0.8, 0.6] diff --git a/api/tests/unit_tests/core/rag/retrieval/test_retrieval_source_query.py b/api/tests/unit_tests/core/rag/retrieval/test_retrieval_source_query.py deleted file mode 100644 index e6313b03d98..00000000000 --- a/api/tests/unit_tests/core/rag/retrieval/test_retrieval_source_query.py +++ /dev/null @@ -1,175 +0,0 @@ -"""Exercise the SQL emitted when workflow retrieval assembles source citations.""" - -from unittest.mock import patch -from uuid import uuid4 - -import pytest -from sqlalchemy import event, inspect -from sqlalchemy.exc import InvalidRequestError -from sqlalchemy.orm import Session - -from core.rag.embedding.retrieval import RetrievalSegments -from core.rag.models.document import Document as RetrievedDocument -from core.rag.retrieval.dataset_retrieval import DatasetRetrieval -from core.workflow.nodes.knowledge_retrieval.retrieval import KnowledgeRetrievalRequest -from models.dataset import Dataset, Document, DocumentSegment -from models.enums import DataSourceType, DocumentCreatedFrom -from repositories.knowledge.retrieval_source_repository import get_retrieval_source_documents - - -def test_source_citations_do_not_fetch_unused_document_fields(sqlite_session: Session) -> None: - tenant_id, creator_id, dataset_id, document_id = (str(uuid4()) for _ in range(4)) - dataset = Dataset(id=dataset_id, tenant_id=tenant_id, name="Knowledge", created_by=creator_id) - document = Document( - id=document_id, - tenant_id=tenant_id, - dataset_id=dataset_id, - position=1, - data_source_type=DataSourceType("upload_file"), - batch="batch-1", - name="source.txt", - created_from=DocumentCreatedFrom("web"), - created_by=creator_id, - enabled=True, - archived=False, - doc_metadata={"category": "transport"}, - data_source_info="x" * 1_000_000, - error="y" * 1_000_000, - ) - segments = [ - DocumentSegment( - tenant_id=tenant_id, - dataset_id=dataset_id, - document_id=document_id, - position=position, - content=f"content-{position}", - word_count=10, - tokens=3, - created_by=creator_id, - enabled=True, - index_node_id=f"node-{position}", - index_node_hash=f"hash-{position}", - ) - for position in (1, 2) - ] - sqlite_session.add_all([dataset, document, *segments]) - sqlite_session.commit() - records = [ - RetrievalSegments(segment=segment, score=score) for segment, score in zip(segments, (0.8, 0.9), strict=True) - ] - retrieved = [RetrievedDocument(provider="dify", page_content="content", metadata={})] - request = KnowledgeRetrievalRequest( - tenant_id=tenant_id, - user_id=creator_id, - app_id=str(uuid4()), - user_from="web", - dataset_ids=[dataset_id], - query="小动物如何运输", - retrieval_mode="multiple", - metadata_filtering_mode="disabled", - ) - queries: list[str] = [] - - def capture_query(_conn, _cursor, statement: str, _parameters, _context, _executemany) -> None: - if "FROM documents" in statement: - queries.append(statement) - - engine = sqlite_session.get_bind() - retrieval = DatasetRetrieval() - event.listen(engine, "before_cursor_execute", capture_query) - try: - with ( - patch.object(retrieval, "_check_knowledge_rate_limit"), - patch.object(retrieval, "_get_available_datasets", return_value=[dataset]), - patch.object(retrieval, "multiple_retrieve", return_value=retrieved), - patch( - "core.rag.retrieval.dataset_retrieval.RetrievalService.format_retrieval_documents", - return_value=records, - ), - ): - sources = retrieval.knowledge_retrieval(sqlite_session, request) - finally: - event.remove(engine, "before_cursor_execute", capture_query) - - assert len(sources) == 2 - assert [source.content for source in sources] == ["content-2", "content-1"] - for source in sources: - assert source.title == "source.txt" - assert source.metadata.document_id == document_id - assert source.metadata.dataset_name == "Knowledge" - assert source.metadata.data_source_type == "upload_file" - assert source.metadata.doc_metadata == {"category": "transport"} - assert len(queries) == 1, "Citation fields must not cause additional lazy SQL queries" - assert "documents.data_source_info" not in queries[0] - assert "documents.error" not in queries[0] - assert "documents.indexing_latency" not in queries[0] - - -def test_source_document_read_is_scoped_and_has_no_lazy_queries(sqlite_session: Session) -> None: - tenant_id, foreign_tenant_id, dataset_id, foreign_dataset_id = (str(uuid4()) for _ in range(4)) - documents = [ - Document( - id=str(uuid4()), - tenant_id=owner, - dataset_id=dataset, - position=1, - data_source_type=DataSourceType("upload_file"), - batch="batch-1", - name="source.txt", - created_from=DocumentCreatedFrom("web"), - created_by=str(uuid4()), - doc_metadata=None, - ) - for owner, dataset in ( - (tenant_id, dataset_id), - (foreign_tenant_id, dataset_id), - (tenant_id, foreign_dataset_id), - ) - ] - sqlite_session.add_all(documents) - sqlite_session.flush() - document_ids = [document.id for document in documents] - sqlite_session.expunge_all() - queries: list[str] = [] - - def capture_query(_conn, _cursor, statement: str, _parameters, _context, _executemany) -> None: - queries.append(statement) - - engine = sqlite_session.get_bind() - event.listen(engine, "before_cursor_execute", capture_query) - try: - result = get_retrieval_source_documents( - tenant_id=tenant_id, - dataset_ids=[dataset_id], - document_ids=[*document_ids, document_ids[0]], - session=sqlite_session, - ) - assert [document.id for document in result] == [document_ids[0]] - assert result[0].name == "source.txt" - assert result[0].data_source_type == "upload_file" - assert result[0].doc_metadata is None - assert set(inspect(result[0]).dict) - {"_sa_instance_state"} == { - "id", - "tenant_id", - "dataset_id", - "name", - "data_source_type", - "doc_metadata", - } - with pytest.raises(InvalidRequestError, match="raiseload"): - _ = result[0].data_source_info - assert ( - get_retrieval_source_documents( - tenant_id=tenant_id, dataset_ids=[], document_ids=document_ids, session=sqlite_session - ) - == [] - ) - assert ( - get_retrieval_source_documents( - tenant_id=tenant_id, dataset_ids=[dataset_id], document_ids=[], session=sqlite_session - ) - == [] - ) - finally: - event.remove(engine, "before_cursor_execute", capture_query) - assert len(queries) == 1 diff --git a/api/tests/unit_tests/core/repositories/test_celery_workflow_node_execution_repository.py b/api/tests/unit_tests/core/repositories/test_celery_workflow_node_execution_repository.py index 0551ad5a22d..46d2c5111a6 100644 --- a/api/tests/unit_tests/core/repositories/test_celery_workflow_node_execution_repository.py +++ b/api/tests/unit_tests/core/repositories/test_celery_workflow_node_execution_repository.py @@ -1,3 +1,50 @@ +from datetime import datetime +from unittest.mock import Mock, patch +from uuid import uuid4 + +import pytest +from kombu.serialization import dumps, loads +from sqlalchemy import Engine, event, select +from sqlalchemy.exc import IntegrityError +from sqlalchemy.orm import Session, sessionmaker + +from core.repositories.celery_workflow_node_execution_repository import CeleryWorkflowNodeExecutionRepository +from core.repositories.factory import OrderConfig, WorkflowNodeExecutionRepository +from core.repositories.sqlalchemy_workflow_node_execution_repository import SQLAlchemyWorkflowNodeExecutionRepository +from core.workflow.node_execution import WorkflowNodeExecution as _workflow_node_execution_batch_WorkflowNodeExecution +from core.workflow.node_execution import WorkflowNodeExecution as _workflow_node_execution_callers_WorkflowNodeExecution +from core.workflow.node_execution_process_data import WORKFLOW_TOOL_ROOT_APP_ID_KEY +from extensions.logstore.aliyun_logstore import AliyunLogStore +from extensions.logstore.repositories.logstore_workflow_node_execution_repository import ( + LogstoreWorkflowNodeExecutionRepository, +) +from graphon.entities.workflow_node_execution import ( + WorkflowNodeExecution as _celery_workflow_node_execution_repository_WorkflowNodeExecution, +) +from graphon.entities.workflow_node_execution import ( + WorkflowNodeExecutionStatus as _celery_workflow_node_execution_repository_WorkflowNodeExecutionStatus, +) +from graphon.enums import BuiltinNodeTypes +from graphon.enums import WorkflowNodeExecutionStatus as _workflow_node_execution_batch_WorkflowNodeExecutionStatus +from graphon.enums import WorkflowNodeExecutionStatus as _workflow_node_execution_callers_WorkflowNodeExecutionStatus +from graphon.variables.segments import ArrayFileSegment, ArrayObjectSegment, StringSegment +from libs.datetime_utils import naive_utc_now +from models import Account, EndUser, Tenant +from models import WorkflowNodeExecutionModel as _celery_workflow_node_execution_repository_WorkflowNodeExecutionModel +from models.enums import ExecutionOffLoadType +from models.model import UploadFile +from models.workflow import WorkflowNodeExecutionModel as _workflow_node_execution_batch_WorkflowNodeExecutionModel +from models.workflow import WorkflowNodeExecutionModel as _workflow_node_execution_callers_WorkflowNodeExecutionModel +from models.workflow import WorkflowNodeExecutionOffload, WorkflowNodeExecutionTriggeredFrom +from repositories.factory import DifyAPIRepositoryFactory +from tasks import workflow_node_execution_tasks +from tasks.workflow_node_execution_tasks import ( + save_workflow_node_execution_task, + save_workflow_node_executions_task, +) +from tests.helpers.unit.config_override import apply_config_overrides +from tests.helpers.unit.model_factories import make_account, make_upload_file + """ Unit tests for CeleryWorkflowNodeExecutionRepository. @@ -5,570 +52,943 @@ These tests verify the Celery-based asynchronous storage functionality for workflow node execution data. """ -from unittest.mock import Mock, patch -from uuid import uuid4 - -import pytest - -from core.repositories.celery_workflow_node_execution_repository import CeleryWorkflowNodeExecutionRepository -from core.repositories.factory import OrderConfig -from core.workflow.node_execution_process_data import WORKFLOW_TOOL_ROOT_APP_ID_KEY -from graphon.entities.workflow_node_execution import ( - WorkflowNodeExecution, - WorkflowNodeExecutionStatus, -) -from graphon.enums import BuiltinNodeTypes -from graphon.variables.segments import ArrayFileSegment, ArrayObjectSegment, StringSegment -from libs.datetime_utils import naive_utc_now -from models import Account, EndUser, Tenant, WorkflowNodeExecutionModel -from models.workflow import WorkflowNodeExecutionTriggeredFrom -from tasks.workflow_node_execution_tasks import ( - save_workflow_node_execution_task, - save_workflow_node_executions_task, -) RESOURCE_TENANT_ID = "resource-tenant-id" -@pytest.fixture -def account(): - """Create an Account user.""" - account = Account(name="Test Account", email="test@example.com") - account.id = str(uuid4()) - account._current_tenant = Tenant(name="Test Tenant") - account._current_tenant.id = str(uuid4()) - return account +class TestCeleryWorkflowNodeExecutionRepositoryCases: + @staticmethod + @pytest.fixture + def account(): + """Create an Account user.""" + account = Account(name="Test Account", email="test@example.com") + account.id = str(uuid4()) + account._current_tenant = Tenant(name="Test Tenant") + account._current_tenant.id = str(uuid4()) + return account - -@pytest.fixture -def end_user(): - """Create an EndUser.""" - user = EndUser( - id=str(uuid4()), - tenant_id=str(uuid4()), - ) - return user - - -@pytest.fixture -def sample_workflow_node_execution(): - """Sample WorkflowNodeExecution for testing.""" - return WorkflowNodeExecution( - id=str(uuid4()), - node_execution_id=str(uuid4()), - workflow_id=str(uuid4()), - workflow_execution_id=str(uuid4()), - index=1, - node_id="test_node", - node_type=BuiltinNodeTypes.START, - title="Test Node", - inputs={"input1": "value1"}, - status=WorkflowNodeExecutionStatus.RUNNING, - created_at=naive_utc_now(), - ) - - -class TestCeleryWorkflowNodeExecutionRepository: - """Test cases for CeleryWorkflowNodeExecutionRepository.""" - - @pytest.mark.parametrize("writer", ["save", "save_many"]) - def test_async_save_preserves_typed_cache_and_plain_history( - self, writer, sqlite_session_factory, account, sample_workflow_node_execution - ): - repo = CeleryWorkflowNodeExecutionRepository( - session_factory=sqlite_session_factory, - tenant_id=RESOURCE_TENANT_ID, - user=account, - app_id="test-app", - triggered_from=WorkflowNodeExecutionTriggeredFrom.WORKFLOW_RUN, + @staticmethod + @pytest.fixture + def end_user(): + """Create an EndUser.""" + user = EndUser( + id=str(uuid4()), + tenant_id=str(uuid4()), ) - literal_object = {"value_type": "string", "value": "user data"} - execution = sample_workflow_node_execution - execution.node_type = BuiltinNodeTypes.TOOL - execution.status = WorkflowNodeExecutionStatus.SUCCEEDED - execution.inputs = {"prompt": StringSegment(value="input")} - execution.process_data = {"nested": {"messages": [StringSegment(value="processing")]}} - execution.outputs = { - "text": StringSegment(value="done"), - "files": ArrayFileSegment(value=[]), - "json": ArrayObjectSegment(value=[literal_object]), - "literal_object": literal_object, - } - task = save_workflow_node_execution_task if writer == "save" else save_workflow_node_executions_task - with patch.object(task, "delay") as delay: - if writer == "save": - repo.save(execution) - else: - repo.save_many([execution]) + return user - delay.assert_called_once() - assert task.run(**delay.call_args.kwargs) is True - - expected_outputs = { - "text": "done", - "files": [], - "json": [literal_object], - "literal_object": literal_object, - } - with sqlite_session_factory() as observer: - stored = observer.get(WorkflowNodeExecutionModel, execution.id) - assert stored is not None - assert stored.inputs_dict == {"prompt": "input"} - assert stored.process_data_dict == {"nested": {"messages": ["processing"]}} - assert stored.outputs_dict == expected_outputs - - (cached,) = repo.get_by_workflow_execution(execution.workflow_execution_id) - assert cached is execution - assert isinstance(cached.inputs["prompt"], StringSegment) - assert isinstance(cached.process_data["nested"]["messages"][0], StringSegment) - assert isinstance(cached.outputs["text"], StringSegment) - assert isinstance(cached.outputs["files"], ArrayFileSegment) - assert isinstance(cached.outputs["json"], ArrayObjectSegment) - assert cached.outputs["literal_object"] == literal_object - - history_repo = CeleryWorkflowNodeExecutionRepository( - session_factory=sqlite_session_factory, - tenant_id=RESOURCE_TENANT_ID, - user=account, - app_id="test-app", - triggered_from=WorkflowNodeExecutionTriggeredFrom.WORKFLOW_RUN, - ) - (restored,) = history_repo.get_by_workflow_execution(execution.workflow_execution_id) - assert restored.inputs == {"prompt": "input"} - assert restored.process_data == {"nested": {"messages": ["processing"]}} - assert restored.outputs == expected_outputs - - def test_trace_read_keeps_source_nodes_out_of_runtime_cache( - self, sqlite_session_factory, end_user, sample_workflow_node_execution - ): - caller = CeleryWorkflowNodeExecutionRepository( - session_factory=sqlite_session_factory, - tenant_id=RESOURCE_TENANT_ID, - user=end_user, - app_id="caller-app", - triggered_from=WorkflowNodeExecutionTriggeredFrom.WORKFLOW_RUN, - ) - root = sample_workflow_node_execution - caller.save_synchronously(root) - source = caller.for_workflow_tool("source-app") - child = root.model_copy( - update={ - "id": "child", - "node_execution_id": "child", - "process_data": {WORKFLOW_TOOL_ROOT_APP_ID_KEY: "caller-app"}, - } - ) - source.save_synchronously(child) - assert [node.id for node in caller.get_by_workflow_execution(root.workflow_execution_id)] == [root.id] - assert [node.id for node in source.get_by_workflow_execution(root.workflow_execution_id)] == [child.id] - assert [node.id for node in caller.get_by_workflow_execution(root.workflow_execution_id)] == [root.id] - - @patch("core.repositories.celery_workflow_node_execution_repository.save_workflow_node_execution_task") - def test_workflow_tool_scope_keeps_async_backend_and_isolates_same_app_history( - self, mock_task, sqlite_session_factory, end_user, sample_workflow_node_execution - ): - caller = CeleryWorkflowNodeExecutionRepository( - session_factory=sqlite_session_factory, - tenant_id=RESOURCE_TENANT_ID, - user=end_user, - app_id="caller-app", - triggered_from=WorkflowNodeExecutionTriggeredFrom.WORKFLOW_RUN, - ) - source = caller.for_workflow_tool("caller-app") - root_node = sample_workflow_node_execution - source_node = root_node.model_copy(update={"id": "source-exec", "node_execution_id": "source-exec"}) - caller.save_synchronously(root_node) - source.save_synchronously(source_node) - source_node.status = WorkflowNodeExecutionStatus.SUCCEEDED - source_node.outputs = {"answer": "ok"} - source.save(source_node) - source.save_execution_data(source_node) - - assert [node.id for node in caller.get_by_workflow_execution(root_node.workflow_execution_id)] == [root_node.id] - assert [node.id for node in source.get_by_workflow_execution(root_node.workflow_execution_id)] == [ - source_node.id - ] - restored_source = caller.for_workflow_tool("caller-app") - assert [node.id for node in restored_source.get_by_workflow_execution(root_node.workflow_execution_id)] == [ - source_node.id - ] - mock_task.delay.assert_called_once_with( - execution_data=source_node.model_dump(), - tenant_id=RESOURCE_TENANT_ID, - app_id="caller-app", - triggered_from=WorkflowNodeExecutionTriggeredFrom.WORKFLOW_TOOL.value, - creator_user_id=end_user.id, - creator_user_role="end_user", + @staticmethod + @pytest.fixture + def sample_workflow_node_execution(): + """Sample WorkflowNodeExecution for testing.""" + return _celery_workflow_node_execution_repository_WorkflowNodeExecution( + id=str(uuid4()), + node_execution_id=str(uuid4()), + workflow_id=str(uuid4()), + workflow_execution_id=str(uuid4()), + index=1, + node_id="test_node", + node_type=BuiltinNodeTypes.START, + title="Test Node", + inputs={"input1": "value1"}, + status=_celery_workflow_node_execution_repository_WorkflowNodeExecutionStatus.RUNNING, + created_at=naive_utc_now(), ) - def test_init_with_sessionmaker(self, sqlite_session_factory, account): - """Test repository initialization with sessionmaker.""" - app_id = "test-app-id" - triggered_from = WorkflowNodeExecutionTriggeredFrom.WORKFLOW_RUN + class TestCeleryWorkflowNodeExecutionRepository: + """Test cases for CeleryWorkflowNodeExecutionRepository.""" - repo = CeleryWorkflowNodeExecutionRepository( - session_factory=sqlite_session_factory, - tenant_id=RESOURCE_TENANT_ID, - user=account, - app_id=app_id, - triggered_from=triggered_from, - ) - - assert repo._tenant_id == RESOURCE_TENANT_ID - assert repo._app_id == app_id - assert repo._triggered_from == triggered_from - assert repo._creator_user_id == account.id - assert repo._creator_user_role is not None - - def test_init_with_cache_initialized(self, sqlite_session_factory, account): - """Test repository initialization with cache properly initialized.""" - repo = CeleryWorkflowNodeExecutionRepository( - session_factory=sqlite_session_factory, - tenant_id=RESOURCE_TENANT_ID, - user=account, - app_id="test-app", - triggered_from=WorkflowNodeExecutionTriggeredFrom.SINGLE_STEP, - ) - - assert repo._execution_cache == {} - assert repo._workflow_execution_mapping == {} - - def test_init_with_end_user(self, sqlite_session_factory, end_user): - """Test repository initialization with EndUser.""" - repo = CeleryWorkflowNodeExecutionRepository( - session_factory=sqlite_session_factory, - tenant_id=RESOURCE_TENANT_ID, - user=end_user, - app_id="test-app", - triggered_from=WorkflowNodeExecutionTriggeredFrom.WORKFLOW_RUN, - ) - - assert repo._tenant_id == RESOURCE_TENANT_ID - - def test_init_without_tenant_id_raises_error(self, sqlite_session_factory): - """Test that initialization fails without tenant_id.""" - # Create an Account with no tenant_id. - user = Account(name="Test Account", email="test@example.com") - user.id = str(uuid4()) - - with pytest.raises(ValueError, match="tenant_id is required"): - CeleryWorkflowNodeExecutionRepository( + @pytest.mark.parametrize("writer", ["save", "save_many"]) + def test_async_save_preserves_typed_cache_and_plain_history( + self, writer, sqlite_session_factory, account, sample_workflow_node_execution + ): + repo = CeleryWorkflowNodeExecutionRepository( session_factory=sqlite_session_factory, - tenant_id="", + tenant_id=RESOURCE_TENANT_ID, + user=account, + app_id="test-app", + triggered_from=WorkflowNodeExecutionTriggeredFrom.WORKFLOW_RUN, + ) + literal_object = {"value_type": "string", "value": "user data"} + execution = sample_workflow_node_execution + execution.node_type = BuiltinNodeTypes.TOOL + execution.status = _celery_workflow_node_execution_repository_WorkflowNodeExecutionStatus.SUCCEEDED + execution.inputs = {"prompt": StringSegment(value="input")} + execution.process_data = {"nested": {"messages": [StringSegment(value="processing")]}} + execution.outputs = { + "text": StringSegment(value="done"), + "files": ArrayFileSegment(value=[]), + "json": ArrayObjectSegment(value=[literal_object]), + "literal_object": literal_object, + } + task = save_workflow_node_execution_task if writer == "save" else save_workflow_node_executions_task + with patch.object(task, "delay") as delay: + if writer == "save": + repo.save(execution) + else: + repo.save_many([execution]) + + delay.assert_called_once() + assert task.run(**delay.call_args.kwargs) is True + + expected_outputs = { + "text": "done", + "files": [], + "json": [literal_object], + "literal_object": literal_object, + } + with sqlite_session_factory() as observer: + stored = observer.get( + _celery_workflow_node_execution_repository_WorkflowNodeExecutionModel, execution.id + ) + assert stored is not None + assert stored.inputs_dict == {"prompt": "input"} + assert stored.process_data_dict == {"nested": {"messages": ["processing"]}} + assert stored.outputs_dict == expected_outputs + + (cached,) = repo.get_by_workflow_execution(execution.workflow_execution_id) + assert cached is execution + assert isinstance(cached.inputs["prompt"], StringSegment) + assert isinstance(cached.process_data["nested"]["messages"][0], StringSegment) + assert isinstance(cached.outputs["text"], StringSegment) + assert isinstance(cached.outputs["files"], ArrayFileSegment) + assert isinstance(cached.outputs["json"], ArrayObjectSegment) + assert cached.outputs["literal_object"] == literal_object + + history_repo = CeleryWorkflowNodeExecutionRepository( + session_factory=sqlite_session_factory, + tenant_id=RESOURCE_TENANT_ID, + user=account, + app_id="test-app", + triggered_from=WorkflowNodeExecutionTriggeredFrom.WORKFLOW_RUN, + ) + (restored,) = history_repo.get_by_workflow_execution(execution.workflow_execution_id) + assert restored.inputs == {"prompt": "input"} + assert restored.process_data == {"nested": {"messages": ["processing"]}} + assert restored.outputs == expected_outputs + + def test_trace_read_keeps_source_nodes_out_of_runtime_cache( + self, sqlite_session_factory, end_user, sample_workflow_node_execution + ): + caller = CeleryWorkflowNodeExecutionRepository( + session_factory=sqlite_session_factory, + tenant_id=RESOURCE_TENANT_ID, + user=end_user, + app_id="caller-app", + triggered_from=WorkflowNodeExecutionTriggeredFrom.WORKFLOW_RUN, + ) + root = sample_workflow_node_execution + caller.save_synchronously(root) + source = caller.for_workflow_tool("source-app") + child = root.model_copy( + update={ + "id": "child", + "node_execution_id": "child", + "process_data": {WORKFLOW_TOOL_ROOT_APP_ID_KEY: "caller-app"}, + } + ) + source.save_synchronously(child) + assert [node.id for node in caller.get_by_workflow_execution(root.workflow_execution_id)] == [root.id] + assert [node.id for node in source.get_by_workflow_execution(root.workflow_execution_id)] == [child.id] + assert [node.id for node in caller.get_by_workflow_execution(root.workflow_execution_id)] == [root.id] + + @patch("core.repositories.celery_workflow_node_execution_repository.save_workflow_node_execution_task") + def test_workflow_tool_scope_keeps_async_backend_and_isolates_same_app_history( + self, mock_task, sqlite_session_factory, end_user, sample_workflow_node_execution + ): + caller = CeleryWorkflowNodeExecutionRepository( + session_factory=sqlite_session_factory, + tenant_id=RESOURCE_TENANT_ID, + user=end_user, + app_id="caller-app", + triggered_from=WorkflowNodeExecutionTriggeredFrom.WORKFLOW_RUN, + ) + source = caller.for_workflow_tool("caller-app") + root_node = sample_workflow_node_execution + source_node = root_node.model_copy(update={"id": "source-exec", "node_execution_id": "source-exec"}) + caller.save_synchronously(root_node) + source.save_synchronously(source_node) + source_node.status = _celery_workflow_node_execution_repository_WorkflowNodeExecutionStatus.SUCCEEDED + source_node.outputs = {"answer": "ok"} + source.save(source_node) + source.save_execution_data(source_node) + + assert [node.id for node in caller.get_by_workflow_execution(root_node.workflow_execution_id)] == [ + root_node.id + ] + assert [node.id for node in source.get_by_workflow_execution(root_node.workflow_execution_id)] == [ + source_node.id + ] + restored_source = caller.for_workflow_tool("caller-app") + assert [node.id for node in restored_source.get_by_workflow_execution(root_node.workflow_execution_id)] == [ + source_node.id + ] + mock_task.delay.assert_called_once_with( + execution_data=source_node.model_dump(), + tenant_id=RESOURCE_TENANT_ID, + app_id="caller-app", + triggered_from=WorkflowNodeExecutionTriggeredFrom.WORKFLOW_TOOL.value, + creator_user_id=end_user.id, + creator_user_role="end_user", + ) + + def test_init_with_sessionmaker(self, sqlite_session_factory, account): + """Test repository initialization with sessionmaker.""" + app_id = "test-app-id" + triggered_from = WorkflowNodeExecutionTriggeredFrom.WORKFLOW_RUN + + repo = CeleryWorkflowNodeExecutionRepository( + session_factory=sqlite_session_factory, + tenant_id=RESOURCE_TENANT_ID, + user=account, + app_id=app_id, + triggered_from=triggered_from, + ) + + assert repo._tenant_id == RESOURCE_TENANT_ID + assert repo._app_id == app_id + assert repo._triggered_from == triggered_from + assert repo._creator_user_id == account.id + assert repo._creator_user_role is not None + + def test_init_with_cache_initialized(self, sqlite_session_factory, account): + """Test repository initialization with cache properly initialized.""" + repo = CeleryWorkflowNodeExecutionRepository( + session_factory=sqlite_session_factory, + tenant_id=RESOURCE_TENANT_ID, + user=account, + app_id="test-app", + triggered_from=WorkflowNodeExecutionTriggeredFrom.SINGLE_STEP, + ) + + assert repo._execution_cache == {} + assert repo._workflow_execution_mapping == {} + + def test_init_with_end_user(self, sqlite_session_factory, end_user): + """Test repository initialization with EndUser.""" + repo = CeleryWorkflowNodeExecutionRepository( + session_factory=sqlite_session_factory, + tenant_id=RESOURCE_TENANT_ID, + user=end_user, + app_id="test-app", + triggered_from=WorkflowNodeExecutionTriggeredFrom.WORKFLOW_RUN, + ) + + assert repo._tenant_id == RESOURCE_TENANT_ID + + def test_init_without_tenant_id_raises_error(self, sqlite_session_factory): + """Test that initialization fails without tenant_id.""" + # Create an Account with no tenant_id. + user = Account(name="Test Account", email="test@example.com") + user.id = str(uuid4()) + + with pytest.raises(ValueError, match="tenant_id is required"): + CeleryWorkflowNodeExecutionRepository( + session_factory=sqlite_session_factory, + tenant_id="", + user=user, + app_id="test-app", + triggered_from=WorkflowNodeExecutionTriggeredFrom.WORKFLOW_RUN, + ) + + def test_init_uses_resource_tenant_when_account_has_no_current_tenant(self, sqlite_session_factory): + user = Account(name="Test Account", email="test@example.com") + user.id = str(uuid4()) + + repo = CeleryWorkflowNodeExecutionRepository( + session_factory=sqlite_session_factory, + tenant_id=RESOURCE_TENANT_ID, user=user, app_id="test-app", triggered_from=WorkflowNodeExecutionTriggeredFrom.WORKFLOW_RUN, ) - def test_init_uses_resource_tenant_when_account_has_no_current_tenant(self, sqlite_session_factory): - user = Account(name="Test Account", email="test@example.com") - user.id = str(uuid4()) + assert repo._tenant_id == RESOURCE_TENANT_ID + assert repo._creator_user_id == user.id - repo = CeleryWorkflowNodeExecutionRepository( - session_factory=sqlite_session_factory, - tenant_id=RESOURCE_TENANT_ID, - user=user, - app_id="test-app", - triggered_from=WorkflowNodeExecutionTriggeredFrom.WORKFLOW_RUN, - ) + @patch("core.repositories.celery_workflow_node_execution_repository.save_workflow_node_execution_task") + def test_save_caches_and_queues_celery_task( + self, mock_task, sqlite_session_factory, account, sample_workflow_node_execution + ): + """Test that save operation caches execution and queues a Celery task.""" + repo = CeleryWorkflowNodeExecutionRepository( + session_factory=sqlite_session_factory, + tenant_id=RESOURCE_TENANT_ID, + user=account, + app_id="test-app", + triggered_from=WorkflowNodeExecutionTriggeredFrom.WORKFLOW_RUN, + ) - assert repo._tenant_id == RESOURCE_TENANT_ID - assert repo._creator_user_id == user.id - - @patch("core.repositories.celery_workflow_node_execution_repository.save_workflow_node_execution_task") - def test_save_caches_and_queues_celery_task( - self, mock_task, sqlite_session_factory, account, sample_workflow_node_execution - ): - """Test that save operation caches execution and queues a Celery task.""" - repo = CeleryWorkflowNodeExecutionRepository( - session_factory=sqlite_session_factory, - tenant_id=RESOURCE_TENANT_ID, - user=account, - app_id="test-app", - triggered_from=WorkflowNodeExecutionTriggeredFrom.WORKFLOW_RUN, - ) - - repo.save(sample_workflow_node_execution) - - # Verify Celery task was queued with correct parameters - mock_task.delay.assert_called_once() - call_args = mock_task.delay.call_args[1] - - assert call_args["execution_data"] == sample_workflow_node_execution.model_dump() - assert call_args["tenant_id"] == RESOURCE_TENANT_ID - assert call_args["app_id"] == "test-app" - assert call_args["triggered_from"] == WorkflowNodeExecutionTriggeredFrom.WORKFLOW_RUN - assert call_args["creator_user_id"] == account.id - - # Verify execution is cached - assert sample_workflow_node_execution.id in repo._execution_cache - assert repo._execution_cache[sample_workflow_node_execution.id] == sample_workflow_node_execution - - # Verify workflow execution mapping is updated - assert sample_workflow_node_execution.workflow_execution_id in repo._workflow_execution_mapping - assert ( - sample_workflow_node_execution.id - in repo._workflow_execution_mapping[sample_workflow_node_execution.workflow_execution_id] - ) - - def test_save_synchronously_uses_sql_repository_without_queueing( - self, sqlite_session_factory, account, sample_workflow_node_execution - ): - repo = CeleryWorkflowNodeExecutionRepository( - session_factory=sqlite_session_factory, - tenant_id=RESOURCE_TENANT_ID, - user=account, - app_id="test-app", - triggered_from=WorkflowNodeExecutionTriggeredFrom.WORKFLOW_RUN, - ) - repo._sql_repository.save_synchronously = Mock() - - repo.save_synchronously(sample_workflow_node_execution) - - repo._sql_repository.save_synchronously.assert_called_once_with(sample_workflow_node_execution) - assert repo._execution_cache[sample_workflow_node_execution.id] is sample_workflow_node_execution - - @patch("core.repositories.celery_workflow_node_execution_repository.save_workflow_node_execution_task") - def test_save_handles_celery_failure( - self, mock_task, sqlite_session_factory, account, sample_workflow_node_execution - ): - """Test that save operation handles Celery task failures.""" - mock_task.delay.side_effect = Exception("Celery is down") - - repo = CeleryWorkflowNodeExecutionRepository( - session_factory=sqlite_session_factory, - tenant_id=RESOURCE_TENANT_ID, - user=account, - app_id="test-app", - triggered_from=WorkflowNodeExecutionTriggeredFrom.WORKFLOW_RUN, - ) - - with pytest.raises(Exception, match="Celery is down"): repo.save(sample_workflow_node_execution) - @patch("core.repositories.celery_workflow_node_execution_repository.save_workflow_node_execution_task") - def test_get_by_workflow_execution_from_cache( - self, mock_task, sqlite_session_factory, account, sample_workflow_node_execution - ): - """Test that get_by_workflow_execution retrieves executions from cache.""" - repo = CeleryWorkflowNodeExecutionRepository( - session_factory=sqlite_session_factory, - tenant_id=RESOURCE_TENANT_ID, - user=account, - app_id="test-app", - triggered_from=WorkflowNodeExecutionTriggeredFrom.WORKFLOW_RUN, + # Verify Celery task was queued with correct parameters + mock_task.delay.assert_called_once() + call_args = mock_task.delay.call_args[1] + + assert call_args["execution_data"] == sample_workflow_node_execution.model_dump() + assert call_args["tenant_id"] == RESOURCE_TENANT_ID + assert call_args["app_id"] == "test-app" + assert call_args["triggered_from"] == WorkflowNodeExecutionTriggeredFrom.WORKFLOW_RUN + assert call_args["creator_user_id"] == account.id + + # Verify execution is cached + assert sample_workflow_node_execution.id in repo._execution_cache + assert repo._execution_cache[sample_workflow_node_execution.id] == sample_workflow_node_execution + + # Verify workflow execution mapping is updated + assert sample_workflow_node_execution.workflow_execution_id in repo._workflow_execution_mapping + assert ( + sample_workflow_node_execution.id + in repo._workflow_execution_mapping[sample_workflow_node_execution.workflow_execution_id] + ) + + def test_save_synchronously_uses_sql_repository_without_queueing( + self, sqlite_session_factory, account, sample_workflow_node_execution + ): + repo = CeleryWorkflowNodeExecutionRepository( + session_factory=sqlite_session_factory, + tenant_id=RESOURCE_TENANT_ID, + user=account, + app_id="test-app", + triggered_from=WorkflowNodeExecutionTriggeredFrom.WORKFLOW_RUN, + ) + repo._sql_repository.save_synchronously = Mock() + + repo.save_synchronously(sample_workflow_node_execution) + + repo._sql_repository.save_synchronously.assert_called_once_with(sample_workflow_node_execution) + assert repo._execution_cache[sample_workflow_node_execution.id] is sample_workflow_node_execution + + @patch("core.repositories.celery_workflow_node_execution_repository.save_workflow_node_execution_task") + def test_save_handles_celery_failure( + self, mock_task, sqlite_session_factory, account, sample_workflow_node_execution + ): + """Test that save operation handles Celery task failures.""" + mock_task.delay.side_effect = Exception("Celery is down") + + repo = CeleryWorkflowNodeExecutionRepository( + session_factory=sqlite_session_factory, + tenant_id=RESOURCE_TENANT_ID, + user=account, + app_id="test-app", + triggered_from=WorkflowNodeExecutionTriggeredFrom.WORKFLOW_RUN, + ) + + with pytest.raises(Exception, match="Celery is down"): + repo.save(sample_workflow_node_execution) + + @patch("core.repositories.celery_workflow_node_execution_repository.save_workflow_node_execution_task") + def test_get_by_workflow_execution_from_cache( + self, mock_task, sqlite_session_factory, account, sample_workflow_node_execution + ): + """Test that get_by_workflow_execution retrieves executions from cache.""" + repo = CeleryWorkflowNodeExecutionRepository( + session_factory=sqlite_session_factory, + tenant_id=RESOURCE_TENANT_ID, + user=account, + app_id="test-app", + triggered_from=WorkflowNodeExecutionTriggeredFrom.WORKFLOW_RUN, + ) + + # Save execution to cache first + repo.save(sample_workflow_node_execution) + + workflow_execution_id = sample_workflow_node_execution.workflow_execution_id + order_config = OrderConfig(order_by=["index"], order_direction="asc") + + result = repo.get_by_workflow_execution(workflow_execution_id, order_config) + + # Verify results were retrieved from cache + assert len(result) == 1 + assert result[0].id == sample_workflow_node_execution.id + assert result[0] is sample_workflow_node_execution + + def test_get_by_workflow_execution_without_order_config(self, sqlite_session_factory, account): + """Test get_by_workflow_execution without order configuration.""" + repo = CeleryWorkflowNodeExecutionRepository( + session_factory=sqlite_session_factory, + tenant_id=RESOURCE_TENANT_ID, + user=account, + app_id="test-app", + triggered_from=WorkflowNodeExecutionTriggeredFrom.WORKFLOW_RUN, + ) + + result = repo.get_by_workflow_execution("workflow-run-id") + + # Should return empty list since nothing in cache + assert len(result) == 0 + + def test_get_by_workflow_execution_loads_persisted_executions_on_cache_miss( + self, sqlite_session_factory, account, sample_workflow_node_execution + ): + repo = CeleryWorkflowNodeExecutionRepository( + session_factory=sqlite_session_factory, + tenant_id=RESOURCE_TENANT_ID, + user=account, + app_id="test-app", + triggered_from=WorkflowNodeExecutionTriggeredFrom.WORKFLOW_RUN, + ) + repo._sql_repository = Mock() + repo._sql_repository.get_by_workflow_execution.return_value = [sample_workflow_node_execution] + + result = repo.get_by_workflow_execution(sample_workflow_node_execution.workflow_execution_id) + + assert result == [sample_workflow_node_execution] + assert repo._execution_cache[sample_workflow_node_execution.id] is sample_workflow_node_execution + assert repo._workflow_execution_mapping[sample_workflow_node_execution.workflow_execution_id] == [ + sample_workflow_node_execution.id + ] + + @patch("core.repositories.celery_workflow_node_execution_repository.save_workflow_node_execution_task") + def test_get_by_workflow_execution_merges_database_and_newer_cache( + self, mock_task, sqlite_session_factory, account, sample_workflow_node_execution + ): + repo = CeleryWorkflowNodeExecutionRepository( + session_factory=sqlite_session_factory, + tenant_id=RESOURCE_TENANT_ID, + user=account, + app_id="test-app", + triggered_from=WorkflowNodeExecutionTriggeredFrom.WORKFLOW_RUN, + ) + persisted_current = sample_workflow_node_execution.model_copy(deep=True) + historical = sample_workflow_node_execution.model_copy( + update={ + "id": str(uuid4()), + "node_execution_id": str(uuid4()), + "index": 0, + "node_id": "start", + } + ) + sample_workflow_node_execution.status = ( + _celery_workflow_node_execution_repository_WorkflowNodeExecutionStatus.SUCCEEDED + ) + repo.save(sample_workflow_node_execution) + repo._sql_repository = Mock() + repo._sql_repository.get_by_workflow_execution.return_value = [persisted_current, historical] + + result = repo.get_by_workflow_execution( + sample_workflow_node_execution.workflow_execution_id, + OrderConfig(order_by=["index"], order_direction="asc"), + ) + + assert [execution.id for execution in result] == [historical.id, sample_workflow_node_execution.id] + assert result[1] is sample_workflow_node_execution + + @patch("core.repositories.celery_workflow_node_execution_repository.save_workflow_node_execution_task") + def test_cache_operations(self, mock_task, sqlite_session_factory, account, sample_workflow_node_execution): + """Test cache operations work correctly.""" + repo = CeleryWorkflowNodeExecutionRepository( + session_factory=sqlite_session_factory, + tenant_id=RESOURCE_TENANT_ID, + user=account, + app_id="test-app", + triggered_from=WorkflowNodeExecutionTriggeredFrom.WORKFLOW_RUN, + ) + + # Test saving to cache + repo.save(sample_workflow_node_execution) + + # Verify cache contains the execution + assert sample_workflow_node_execution.id in repo._execution_cache + + # Test retrieving from cache + result = repo.get_by_workflow_execution(sample_workflow_node_execution.workflow_execution_id) + assert len(result) == 1 + assert result[0].id == sample_workflow_node_execution.id + + @patch("core.repositories.celery_workflow_node_execution_repository.save_workflow_node_execution_task") + def test_multiple_executions_same_workflow(self, mock_task, sqlite_session_factory, account): + """Test multiple executions for the same workflow.""" + repo = CeleryWorkflowNodeExecutionRepository( + session_factory=sqlite_session_factory, + tenant_id=RESOURCE_TENANT_ID, + user=account, + app_id="test-app", + triggered_from=WorkflowNodeExecutionTriggeredFrom.WORKFLOW_RUN, + ) + + # Create multiple executions for the same workflow + workflow_execution_id = str(uuid4()) + exec1 = _celery_workflow_node_execution_repository_WorkflowNodeExecution( + id=str(uuid4()), + node_execution_id=str(uuid4()), + workflow_id=str(uuid4()), + workflow_execution_id=workflow_execution_id, + index=1, + node_id="node1", + node_type=BuiltinNodeTypes.START, + title="Node 1", + inputs={"input1": "value1"}, + status=_celery_workflow_node_execution_repository_WorkflowNodeExecutionStatus.RUNNING, + created_at=naive_utc_now(), + ) + exec2 = _celery_workflow_node_execution_repository_WorkflowNodeExecution( + id=str(uuid4()), + node_execution_id=str(uuid4()), + workflow_id=str(uuid4()), + workflow_execution_id=workflow_execution_id, + index=2, + node_id="node2", + node_type=BuiltinNodeTypes.LLM, + title="Node 2", + inputs={"input2": "value2"}, + status=_celery_workflow_node_execution_repository_WorkflowNodeExecutionStatus.RUNNING, + created_at=naive_utc_now(), + ) + + # Save both executions + repo.save(exec1) + repo.save(exec2) + + # Verify both are cached and mapped + assert len(repo._execution_cache) == 2 + assert len(repo._workflow_execution_mapping[workflow_execution_id]) == 2 + + # Test retrieval + result = repo.get_by_workflow_execution(workflow_execution_id) + assert len(result) == 2 + + @patch("core.repositories.celery_workflow_node_execution_repository.save_workflow_node_execution_task") + def test_ordering_functionality(self, mock_task, sqlite_session_factory, account): + """Test ordering functionality works correctly.""" + repo = CeleryWorkflowNodeExecutionRepository( + session_factory=sqlite_session_factory, + tenant_id=account.current_tenant_id, + user=account, + app_id="test-app", + triggered_from=WorkflowNodeExecutionTriggeredFrom.WORKFLOW_RUN, + ) + + # Create executions with different indices + workflow_execution_id = str(uuid4()) + exec1 = _celery_workflow_node_execution_repository_WorkflowNodeExecution( + id=str(uuid4()), + node_execution_id=str(uuid4()), + workflow_id=str(uuid4()), + workflow_execution_id=workflow_execution_id, + index=2, + node_id="node2", + node_type=BuiltinNodeTypes.START, + title="Node 2", + inputs={}, + status=_celery_workflow_node_execution_repository_WorkflowNodeExecutionStatus.RUNNING, + created_at=naive_utc_now(), + ) + exec2 = _celery_workflow_node_execution_repository_WorkflowNodeExecution( + id=str(uuid4()), + node_execution_id=str(uuid4()), + workflow_id=str(uuid4()), + workflow_execution_id=workflow_execution_id, + index=1, + node_id="node1", + node_type=BuiltinNodeTypes.LLM, + title="Node 1", + inputs={}, + status=_celery_workflow_node_execution_repository_WorkflowNodeExecutionStatus.RUNNING, + created_at=naive_utc_now(), + ) + + # Save in random order + repo.save(exec1) + repo.save(exec2) + + # Test ascending order + order_config = OrderConfig(order_by=["index"], order_direction="asc") + result = repo.get_by_workflow_execution(workflow_execution_id, order_config) + assert len(result) == 2 + assert result[0].index == 1 + assert result[1].index == 2 + + # Test descending order + order_config = OrderConfig(order_by=["index"], order_direction="desc") + result = repo.get_by_workflow_execution(workflow_execution_id, order_config) + assert len(result) == 2 + assert result[0].index == 2 + assert result[1].index == 1 + + +"""Batch cleanup persists exact records without losing ownership or stored data.""" + + +def _execution() -> _workflow_node_execution_batch_WorkflowNodeExecution: + return _workflow_node_execution_batch_WorkflowNodeExecution( + id=str(uuid4()), + node_execution_id=str(uuid4()), + workflow_id="source-workflow", + workflow_execution_id="root-run", + triggered_from_workflow_id="caller-workflow", + triggered_from_node_execution_id="caller-execution", + index=1, + node_id="tool", + node_type="tool", + title="Nested tool", + created_at=datetime(2026, 9, 23), + process_data={"workflow_agent_binding_id": "binding"}, + ) + + +class TestWorkflowNodeExecutionBatchCases: + @staticmethod + @pytest.mark.parametrize("backend", ["sql", "celery", "logstore"]) + def test_batch_creates_missing_rows_and_updates_only_supplied_records( + sqlite_session_factory: sessionmaker[Session], + sqlite_engine: Engine, + monkeypatch: pytest.MonkeyPatch, + backend: str, + ) -> None: + expiring_session_factory = sessionmaker(sqlite_engine) + sql = SQLAlchemyWorkflowNodeExecutionRepository( + expiring_session_factory, + "tenant", + make_account(), + "source-app", + WorkflowNodeExecutionTriggeredFrom.WORKFLOW_TOOL, ) + existing, missing, untouched = (_execution() for _ in range(3)) + sql.save(existing) + sql.save(untouched) + with sqlite_session_factory.begin() as session: + session.add(make_upload_file(file_id="offloaded-inputs", tenant_id="tenant")) + row = session.get(_workflow_node_execution_batch_WorkflowNodeExecutionModel, existing.id) + assert row is not None + row.inputs = '{"large": "stored preview"}' + session.add( + WorkflowNodeExecutionOffload( + tenant_id="tenant", + app_id="source-app", + node_execution_id=existing.id, + type_=ExecutionOffLoadType.INPUTS, + file_id="offloaded-inputs", + ) + ) + existing.inputs = {"large": "full content restored from offload"} + for execution in (existing, missing): + execution.status = _workflow_node_execution_batch_WorkflowNodeExecutionStatus.FAILED + execution.error = "workflow stopped" + execution.finished_at = datetime(2026, 9, 23, 0, 0, 1) + execution.process_data = {"partial": "result"} - # Save execution to cache first - repo.save(sample_workflow_node_execution) + queued: list[dict[str, object]] = [] + logstore_records: list[dict[str, str]] = [] + if backend == "celery": + task = workflow_node_execution_tasks.save_workflow_node_executions_task + monkeypatch.setattr(task, "delay", lambda **kwargs: queued.append(kwargs)) + repository = CeleryWorkflowNodeExecutionRepository( + expiring_session_factory, + "tenant", + make_account(), + "source-app", + WorkflowNodeExecutionTriggeredFrom.WORKFLOW_TOOL, + ) + elif backend == "logstore": + apply_config_overrides(monkeypatch, LOGSTORE_DUAL_WRITE_ENABLED=True) + monkeypatch.setattr(AliyunLogStore, "_instance", None) + monkeypatch.setattr(AliyunLogStore, "_initialized", False) + monkeypatch.setenv("ALIYUN_SLS_ENDPOINT", "localhost") + monkeypatch.setattr(AliyunLogStore, "_check_endpoint_connectivity", staticmethod(lambda *_: None)) + monkeypatch.setattr( + AliyunLogStore, + "put_log", + lambda _self, _name, record: logstore_records.append(dict(record)), + ) + repository = LogstoreWorkflowNodeExecutionRepository( + expiring_session_factory, + "tenant", + make_account(), + "source-app", + WorkflowNodeExecutionTriggeredFrom.WORKFLOW_TOOL, + ) + else: + repository = sql - workflow_execution_id = sample_workflow_node_execution.workflow_execution_id - order_config = OrderConfig(order_by=["index"], order_direction="asc") + commits: list[None] = [] + event.listen(sqlite_engine, "commit", lambda _connection: commits.append(None)) + repository.save_many([existing, missing]) + repository.save_many([]) + if backend == "celery": + assert len(queued) == 1 + assert commits == [] + content_type, encoding, payload = dumps(queued[0], serializer="json") + workflow_node_execution_tasks.save_workflow_node_executions_task.run( + **loads(payload, content_type, encoding, accept={content_type}) + ) + assert len(commits) == 1 - result = repo.get_by_workflow_execution(workflow_execution_id, order_config) - - # Verify results were retrieved from cache - assert len(result) == 1 - assert result[0].id == sample_workflow_node_execution.id - assert result[0] is sample_workflow_node_execution - - def test_get_by_workflow_execution_without_order_config(self, sqlite_session_factory, account): - """Test get_by_workflow_execution without order configuration.""" - repo = CeleryWorkflowNodeExecutionRepository( - session_factory=sqlite_session_factory, - tenant_id=RESOURCE_TENANT_ID, - user=account, - app_id="test-app", - triggered_from=WorkflowNodeExecutionTriggeredFrom.WORKFLOW_RUN, - ) - - result = repo.get_by_workflow_execution("workflow-run-id") - - # Should return empty list since nothing in cache - assert len(result) == 0 - - def test_get_by_workflow_execution_loads_persisted_executions_on_cache_miss( - self, sqlite_session_factory, account, sample_workflow_node_execution - ): - repo = CeleryWorkflowNodeExecutionRepository( - session_factory=sqlite_session_factory, - tenant_id=RESOURCE_TENANT_ID, - user=account, - app_id="test-app", - triggered_from=WorkflowNodeExecutionTriggeredFrom.WORKFLOW_RUN, - ) - repo._sql_repository = Mock() - repo._sql_repository.get_by_workflow_execution.return_value = [sample_workflow_node_execution] - - result = repo.get_by_workflow_execution(sample_workflow_node_execution.workflow_execution_id) - - assert result == [sample_workflow_node_execution] - assert repo._execution_cache[sample_workflow_node_execution.id] is sample_workflow_node_execution - assert repo._workflow_execution_mapping[sample_workflow_node_execution.workflow_execution_id] == [ - sample_workflow_node_execution.id - ] - - @patch("core.repositories.celery_workflow_node_execution_repository.save_workflow_node_execution_task") - def test_get_by_workflow_execution_merges_database_and_newer_cache( - self, mock_task, sqlite_session_factory, account, sample_workflow_node_execution - ): - repo = CeleryWorkflowNodeExecutionRepository( - session_factory=sqlite_session_factory, - tenant_id=RESOURCE_TENANT_ID, - user=account, - app_id="test-app", - triggered_from=WorkflowNodeExecutionTriggeredFrom.WORKFLOW_RUN, - ) - persisted_current = sample_workflow_node_execution.model_copy(deep=True) - historical = sample_workflow_node_execution.model_copy( - update={ - "id": str(uuid4()), - "node_execution_id": str(uuid4()), - "index": 0, - "node_id": "start", + with sqlite_session_factory() as session: + rows = { + row.id: row + for row in session.scalars( + _workflow_node_execution_batch_WorkflowNodeExecutionModel.preload_offload_data( + select(_workflow_node_execution_batch_WorkflowNodeExecutionModel) + ) + ) } + assert set(rows) == {existing.id, missing.id, untouched.id} + assert rows[untouched.id].status == _workflow_node_execution_batch_WorkflowNodeExecutionStatus.RUNNING + for execution in (existing, missing): + row = rows[execution.id] + assert row.status == _workflow_node_execution_batch_WorkflowNodeExecutionStatus.FAILED + assert row.error == "workflow stopped" + assert row.finished_at == execution.finished_at + assert (row.tenant_id, row.app_id, row.triggered_from) == ( + "tenant", + "source-app", + WorkflowNodeExecutionTriggeredFrom.WORKFLOW_TOOL, + ) + assert (row.workflow_id, row.workflow_run_id) == ("source-workflow", "root-run") + assert (row.triggered_from_workflow_id, row.triggered_from_node_execution_id) == ( + "caller-workflow", + "caller-execution", + ) + assert rows[existing.id].process_data_dict == {"partial": "result", "workflow_agent_binding_id": "binding"} + assert rows[existing.id].inputs_dict == {"large": "stored preview"} + assert rows[existing.id].offload_data[0].file_id == "offloaded-inputs" + if backend == "logstore": + assert {row["id"] for row in logstore_records} == {existing.id, missing.id} + assert all(row["status"] == "failed" and row["log_version"] for row in logstore_records) + assert all(row["triggered_from_node_execution_id"] == "caller-execution" for row in logstore_records) + + @staticmethod + @pytest.mark.parametrize("backend", ["sql", "celery"]) + @pytest.mark.parametrize("owner", ["tenant_id", "app_id", "triggered_from", "workflow_id", "workflow_run_id"]) + def test_owner_mismatch_rolls_back_the_whole_batch( + sqlite_session_factory: sessionmaker[Session], monkeypatch: pytest.MonkeyPatch, backend: str, owner: str + ) -> None: + repository = SQLAlchemyWorkflowNodeExecutionRepository( + sqlite_session_factory, + "tenant", + make_account(), + "source-app", + WorkflowNodeExecutionTriggeredFrom.WORKFLOW_TOOL, ) - sample_workflow_node_execution.status = WorkflowNodeExecutionStatus.SUCCEEDED - repo.save(sample_workflow_node_execution) - repo._sql_repository = Mock() - repo._sql_repository.get_by_workflow_execution.return_value = [persisted_current, historical] + valid, foreign = _execution(), _execution() + repository.save(valid) + repository.save(foreign) + with sqlite_session_factory.begin() as session: + row = session.get(_workflow_node_execution_batch_WorkflowNodeExecutionModel, foreign.id) + assert row is not None + setattr( + row, owner, WorkflowNodeExecutionTriggeredFrom.WORKFLOW_RUN if owner == "triggered_from" else "other" + ) + for execution in (valid, foreign): + execution.status = _workflow_node_execution_batch_WorkflowNodeExecutionStatus.FAILED + if backend == "celery": + task = workflow_node_execution_tasks.save_workflow_node_executions_task + monkeypatch.setattr(task, "retry", lambda **kwargs: kwargs["exc"]) + with pytest.raises(IntegrityError): + task.run( + executions_data=[execution.model_dump() for execution in (valid, foreign)], + tenant_id="tenant", + app_id="source-app", + triggered_from=WorkflowNodeExecutionTriggeredFrom.WORKFLOW_TOOL.value, + creator_user_id="account", + creator_user_role="account", + ) + else: + with pytest.raises(IntegrityError): + repository.save_many([valid, foreign]) + with sqlite_session_factory() as session: + assert all( + row.status == _workflow_node_execution_batch_WorkflowNodeExecutionStatus.RUNNING + for row in session.scalars(select(_workflow_node_execution_batch_WorkflowNodeExecutionModel)) + ) - result = repo.get_by_workflow_execution( - sample_workflow_node_execution.workflow_execution_id, - OrderConfig(order_by=["index"], order_direction="asc"), + @staticmethod + @pytest.mark.parametrize( + "status", + [ + _workflow_node_execution_batch_WorkflowNodeExecutionStatus.FAILED, + _workflow_node_execution_batch_WorkflowNodeExecutionStatus.PAUSED, + _workflow_node_execution_batch_WorkflowNodeExecutionStatus.RETRY, + ], + ) + def test_delayed_start_cannot_undo_completed_failure_cleanup( + sqlite_session_factory: sessionmaker[Session], + monkeypatch: pytest.MonkeyPatch, + status: _workflow_node_execution_batch_WorkflowNodeExecutionStatus, + ) -> None: + queued: list[dict[str, object]] = [] + start_task = workflow_node_execution_tasks.save_workflow_node_execution_task + batch_task = workflow_node_execution_tasks.save_workflow_node_executions_task + monkeypatch.setattr(start_task, "delay", lambda **kwargs: queued.append(kwargs)) + monkeypatch.setattr(batch_task, "delay", lambda **kwargs: queued.append(kwargs)) + repository = CeleryWorkflowNodeExecutionRepository( + sqlite_session_factory, + "tenant", + make_account(), + "source-app", + WorkflowNodeExecutionTriggeredFrom.WORKFLOW_TOOL, ) + execution = _execution() + repository.save(execution) + execution.status = status + execution.error = "workflow stopped" + execution.finished_at = datetime(2026, 9, 23, 0, 0, 1) + repository.save_many([execution]) - assert [execution.id for execution in result] == [historical.id, sample_workflow_node_execution.id] - assert result[1] is sample_workflow_node_execution + for task, kwargs in zip((batch_task, start_task), reversed(queued), strict=True): + content_type, encoding, payload = dumps(kwargs, serializer="json") + task.run(**loads(payload, content_type, encoding, accept={content_type})) - @patch("core.repositories.celery_workflow_node_execution_repository.save_workflow_node_execution_task") - def test_cache_operations(self, mock_task, sqlite_session_factory, account, sample_workflow_node_execution): - """Test cache operations work correctly.""" - repo = CeleryWorkflowNodeExecutionRepository( - session_factory=sqlite_session_factory, - tenant_id=RESOURCE_TENANT_ID, - user=account, - app_id="test-app", - triggered_from=WorkflowNodeExecutionTriggeredFrom.WORKFLOW_RUN, + with sqlite_session_factory() as session: + row = session.get(_workflow_node_execution_batch_WorkflowNodeExecutionModel, execution.id) + assert row is not None + if status == _workflow_node_execution_batch_WorkflowNodeExecutionStatus.FAILED: + assert (row.status, row.error, row.finished_at) == (status, execution.error, execution.finished_at) + else: + assert row.status == _workflow_node_execution_batch_WorkflowNodeExecutionStatus.RUNNING + assert row.finished_at is None + + +"""Caller identity survives storage updates independently of node payloads.""" + + +class TestWorkflowNodeExecutionCallersCases: + @staticmethod + @pytest.mark.parametrize("backend", ["sql", "celery", "logstore"]) + def test_caller_locators_survive_storage_round_trip_and_updates( + sqlite_session_factory: sessionmaker[Session], monkeypatch: pytest.MonkeyPatch, backend: str + ) -> None: + apply_config_overrides(monkeypatch, LOGSTORE_DUAL_WRITE_ENABLED=False) + logstore_records: list[dict[str, str]] = [] + sql = SQLAlchemyWorkflowNodeExecutionRepository( + sqlite_session_factory, + "tenant", + make_account(), + "source-app", + WorkflowNodeExecutionTriggeredFrom.WORKFLOW_TOOL, ) - - # Test saving to cache - repo.save(sample_workflow_node_execution) - - # Verify cache contains the execution - assert sample_workflow_node_execution.id in repo._execution_cache - - # Test retrieving from cache - result = repo.get_by_workflow_execution(sample_workflow_node_execution.workflow_execution_id) - assert len(result) == 1 - assert result[0].id == sample_workflow_node_execution.id - - @patch("core.repositories.celery_workflow_node_execution_repository.save_workflow_node_execution_task") - def test_multiple_executions_same_workflow(self, mock_task, sqlite_session_factory, account): - """Test multiple executions for the same workflow.""" - repo = CeleryWorkflowNodeExecutionRepository( - session_factory=sqlite_session_factory, - tenant_id=RESOURCE_TENANT_ID, - user=account, - app_id="test-app", - triggered_from=WorkflowNodeExecutionTriggeredFrom.WORKFLOW_RUN, - ) - - # Create multiple executions for the same workflow - workflow_execution_id = str(uuid4()) - exec1 = WorkflowNodeExecution( + execution = _workflow_node_execution_callers_WorkflowNodeExecution( id=str(uuid4()), - node_execution_id=str(uuid4()), + node_execution_id="child-execution", workflow_id=str(uuid4()), - workflow_execution_id=workflow_execution_id, + workflow_execution_id="root-run", + triggered_from_workflow_id=str(uuid4()), + triggered_from_node_execution_id="call/" + "x" * 250, index=1, - node_id="node1", - node_type=BuiltinNodeTypes.START, - title="Node 1", - inputs={"input1": "value1"}, - status=WorkflowNodeExecutionStatus.RUNNING, - created_at=naive_utc_now(), + node_id="reused-canvas-node", + node_type="tool", + title="Nested Tool", + created_at=datetime(2026, 9, 23), ) - exec2 = WorkflowNodeExecution( + repository: WorkflowNodeExecutionRepository + if backend == "celery": + + def deliver(**kwargs: object) -> None: + content_type, encoding, payload = dumps(kwargs, serializer="json") + save_workflow_node_execution_task.run(**loads(payload, content_type, encoding, accept={content_type})) + + monkeypatch.setattr(save_workflow_node_execution_task, "delay", deliver) + repository = CeleryWorkflowNodeExecutionRepository( + sqlite_session_factory, + "tenant", + make_account(), + "source-app", + WorkflowNodeExecutionTriggeredFrom.WORKFLOW_TOOL, + ) + elif backend == "logstore": + monkeypatch.setattr(AliyunLogStore, "_instance", None) + monkeypatch.setattr(AliyunLogStore, "_initialized", False) + monkeypatch.setenv("ALIYUN_SLS_ENDPOINT", "localhost") + monkeypatch.setattr(AliyunLogStore, "_check_endpoint_connectivity", staticmethod(lambda *_: None)) + monkeypatch.setattr( + AliyunLogStore, + "put_log", + lambda _self, _name, record: logstore_records.append(dict(record)), + ) + repository = LogstoreWorkflowNodeExecutionRepository( + sqlite_session_factory, + "tenant", + make_account(), + "source-app", + WorkflowNodeExecutionTriggeredFrom.WORKFLOW_TOOL, + ) + else: + repository = sql + + repository.save(execution) + execution.status = _workflow_node_execution_callers_WorkflowNodeExecutionStatus.SUCCEEDED + execution.process_data = {"result": "no caller identifiers in this payload"} + execution.finished_at = datetime(2026, 9, 23, 0, 0, 1) + repository.save(execution) + repository.save_execution_data(execution) + + if backend == "logstore": + row = logstore_records[-1] + monkeypatch.setattr(AliyunLogStore, "execute_sql", lambda *_args, **_kwargs: [row]) + restored = repository.get_by_workflow_execution("root-run")[0] + apply_config_overrides( + monkeypatch, + API_WORKFLOW_NODE_EXECUTION_REPOSITORY=( + "extensions.logstore.repositories.logstore_api_workflow_node_execution_repository." + "LogstoreAPIWorkflowNodeExecutionRepository" + ), + ) + api_repository = DifyAPIRepositoryFactory.create_api_workflow_node_execution_repository( + sqlite_session_factory + ) + model = api_repository.get_execution_by_id(execution.id, "tenant") + else: + restored = sql.get_by_workflow_execution("root-run")[0] + with sqlite_session_factory() as session: + model = session.get(_workflow_node_execution_callers_WorkflowNodeExecutionModel, execution.id) + + assert model is not None + assert isinstance(restored, _workflow_node_execution_callers_WorkflowNodeExecution) + for record in (restored, model): + assert record.triggered_from_workflow_id == execution.triggered_from_workflow_id + assert record.triggered_from_node_execution_id == execution.triggered_from_node_execution_id + assert record.workflow_id == execution.workflow_id + assert model.app_id == "source-app" + assert model.workflow_run_id == "root-run" + assert restored.process_data == execution.process_data + + @staticmethod + def test_sql_offload_preserves_caller_locators( + sqlite_session_factory: sessionmaker[Session], monkeypatch: pytest.MonkeyPatch + ) -> None: + apply_config_overrides( + monkeypatch, WORKFLOW_VARIABLE_TRUNCATION_MAX_SIZE=32, WORKFLOW_VARIABLE_TRUNCATION_STRING_LENGTH=8 + ) + repository = SQLAlchemyWorkflowNodeExecutionRepository( + sqlite_session_factory, + "tenant", + make_account(), + "source-app", + WorkflowNodeExecutionTriggeredFrom.WORKFLOW_TOOL, + ) + stored_content: dict[str, bytes] = {} + + def upload_file(*, content: bytes, **_kwargs: object) -> UploadFile: + file = make_upload_file(file_id=str(uuid4()), key="process-data", tenant_id="tenant") + with sqlite_session_factory.begin() as session: + session.add(file) + stored_content[file.key] = content + return file + + monkeypatch.setattr(repository._file_service, "upload_file", upload_file) + monkeypatch.setattr( + "core.repositories.sqlalchemy_workflow_node_execution_repository.storage.load", stored_content.__getitem__ + ) + execution = _workflow_node_execution_callers_WorkflowNodeExecution( id=str(uuid4()), - node_execution_id=str(uuid4()), workflow_id=str(uuid4()), - workflow_execution_id=workflow_execution_id, - index=2, - node_id="node2", - node_type=BuiltinNodeTypes.LLM, - title="Node 2", - inputs={"input2": "value2"}, - status=WorkflowNodeExecutionStatus.RUNNING, - created_at=naive_utc_now(), - ) - - # Save both executions - repo.save(exec1) - repo.save(exec2) - - # Verify both are cached and mapped - assert len(repo._execution_cache) == 2 - assert len(repo._workflow_execution_mapping[workflow_execution_id]) == 2 - - # Test retrieval - result = repo.get_by_workflow_execution(workflow_execution_id) - assert len(result) == 2 - - @patch("core.repositories.celery_workflow_node_execution_repository.save_workflow_node_execution_task") - def test_ordering_functionality(self, mock_task, sqlite_session_factory, account): - """Test ordering functionality works correctly.""" - repo = CeleryWorkflowNodeExecutionRepository( - session_factory=sqlite_session_factory, - tenant_id=account.current_tenant_id, - user=account, - app_id="test-app", - triggered_from=WorkflowNodeExecutionTriggeredFrom.WORKFLOW_RUN, - ) - - # Create executions with different indices - workflow_execution_id = str(uuid4()) - exec1 = WorkflowNodeExecution( - id=str(uuid4()), - node_execution_id=str(uuid4()), - workflow_id=str(uuid4()), - workflow_execution_id=workflow_execution_id, - index=2, - node_id="node2", - node_type=BuiltinNodeTypes.START, - title="Node 2", - inputs={}, - status=WorkflowNodeExecutionStatus.RUNNING, - created_at=naive_utc_now(), - ) - exec2 = WorkflowNodeExecution( - id=str(uuid4()), - node_execution_id=str(uuid4()), - workflow_id=str(uuid4()), - workflow_execution_id=workflow_execution_id, + workflow_execution_id="root-run", + triggered_from_workflow_id=str(uuid4()), + triggered_from_node_execution_id="immediate-tool-execution", index=1, - node_id="node1", - node_type=BuiltinNodeTypes.LLM, - title="Node 1", - inputs={}, - status=WorkflowNodeExecutionStatus.RUNNING, - created_at=naive_utc_now(), + node_id="node", + node_type="code", + title="Code", + created_at=datetime(2026, 9, 23), + process_data={"large": "x" * 1000}, ) + repository.save(execution) + repository.save_execution_data(execution) - # Save in random order - repo.save(exec1) - repo.save(exec2) - - # Test ascending order - order_config = OrderConfig(order_by=["index"], order_direction="asc") - result = repo.get_by_workflow_execution(workflow_execution_id, order_config) - assert len(result) == 2 - assert result[0].index == 1 - assert result[1].index == 2 - - # Test descending order - order_config = OrderConfig(order_by=["index"], order_direction="desc") - result = repo.get_by_workflow_execution(workflow_execution_id, order_config) - assert len(result) == 2 - assert result[0].index == 2 - assert result[1].index == 1 + restored = repository.get_by_workflow_execution("root-run")[0] + assert isinstance(restored, _workflow_node_execution_callers_WorkflowNodeExecution) + assert restored.process_data_truncated + assert restored.process_data == execution.process_data + assert restored.triggered_from_workflow_id == execution.triggered_from_workflow_id + assert restored.triggered_from_node_execution_id == "immediate-tool-execution" diff --git a/api/tests/unit_tests/core/repositories/test_human_input_form_repository_impl.py b/api/tests/unit_tests/core/repositories/test_human_input_form_repository_impl.py deleted file mode 100644 index 39071a16b53..00000000000 --- a/api/tests/unit_tests/core/repositories/test_human_input_form_repository_impl.py +++ /dev/null @@ -1,460 +0,0 @@ -"""Unit tests for HumanInputFormRepositoryImpl private helpers.""" - -from __future__ import annotations - -from datetime import datetime, timedelta - -import pytest -from sqlalchemy import Engine, event -from sqlalchemy.exc import SQLAlchemyError -from sqlalchemy.orm import Session - -from core.repositories.human_input_repository import ( - FormNotFoundError, - HumanInputFormRecord, - HumanInputFormRepositoryImpl, - HumanInputFormSubmissionRepository, - _WorkspaceMemberInfo, -) -from core.workflow.human_input_adapter import ( - EmailDeliveryConfig, - EmailDeliveryMethod, - EmailRecipients, - ExternalRecipient, - MemberRecipient, -) -from core.workflow.nodes.human_input.entities import ( - FormDefinition, - UserActionConfig, -) -from core.workflow.nodes.human_input.enums import HumanInputFormKind, HumanInputFormStatus -from libs.datetime_utils import naive_utc_now -from models import Account, TenantAccountJoin -from models.human_input import ( - EmailExternalRecipientPayload, - EmailMemberRecipientPayload, - HumanInputForm, - HumanInputFormRecipient, - RecipientType, - StandaloneWebAppRecipientPayload, -) - - -def _build_repository() -> HumanInputFormRepositoryImpl: - return HumanInputFormRepositoryImpl(tenant_id="tenant-id") - - -def _add_workspace_member( - session: Session, - *, - user_id: str, - email: str, - tenant_id: str = "tenant-id", -) -> None: - account = Account(name=user_id, email=email) - account.id = user_id - session.add_all([account, TenantAccountJoin(tenant_id=tenant_id, account_id=user_id)]) - session.commit() - - -class TestHumanInputFormRepositoryImplHelpers: - def test_build_email_recipients_with_member_and_external( - self, - monkeypatch: pytest.MonkeyPatch, - sqlite_session: Session, - ) -> None: - repo = _build_repository() - _add_workspace_member(sqlite_session, user_id="member-1", email="member@example.com") - query_workspace_members_by_ids = repo._query_workspace_members_by_ids - - def query_members(session: Session, restrict_to_user_ids: list[str]) -> list[_WorkspaceMemberInfo]: - assert session is sqlite_session - return query_workspace_members_by_ids( - session=session, - restrict_to_user_ids=restrict_to_user_ids, - ) - - monkeypatch.setattr(repo, "_query_workspace_members_by_ids", query_members) - - recipients = repo._build_email_recipients( - session=sqlite_session, - form_id="form-id", - delivery_id="delivery-id", - recipients_config=EmailRecipients( - include_bound_group=False, - items=[ - MemberRecipient(reference_id="member-1"), - ExternalRecipient(email="external@example.com"), - ], - ), - ) - - assert len(recipients) == 2 - member_recipient = next(r for r in recipients if r.recipient_type == RecipientType.EMAIL_MEMBER) - external_recipient = next(r for r in recipients if r.recipient_type == RecipientType.EMAIL_EXTERNAL) - - member_payload = EmailMemberRecipientPayload.model_validate_json(member_recipient.recipient_payload) - assert member_payload.user_id == "member-1" - assert member_payload.email == "member@example.com" - - external_payload = EmailExternalRecipientPayload.model_validate_json(external_recipient.recipient_payload) - assert external_payload.email == "external@example.com" - - def test_build_email_recipients_skips_unknown_members(self, sqlite_session: Session) -> None: - repo = _build_repository() - - recipients = repo._build_email_recipients( - session=sqlite_session, - form_id="form-id", - delivery_id="delivery-id", - recipients_config=EmailRecipients( - include_bound_group=False, - items=[ - MemberRecipient(reference_id="missing-member"), - ExternalRecipient(email="external@example.com"), - ], - ), - ) - - assert len(recipients) == 1 - assert recipients[0].recipient_type == RecipientType.EMAIL_EXTERNAL - - def test_build_email_recipients_whole_workspace_uses_all_members(self, sqlite_session: Session) -> None: - repo = _build_repository() - _add_workspace_member(sqlite_session, user_id="member-1", email="member1@example.com") - _add_workspace_member(sqlite_session, user_id="member-2", email="member2@example.com") - - recipients = repo._build_email_recipients( - session=sqlite_session, - form_id="form-id", - delivery_id="delivery-id", - recipients_config=EmailRecipients( - include_bound_group=True, - items=[], - ), - ) - - assert len(recipients) == 2 - emails = {EmailMemberRecipientPayload.model_validate_json(r.recipient_payload).email for r in recipients} - assert emails == {"member1@example.com", "member2@example.com"} - - def test_build_email_recipients_dedupes_external_by_email(self, sqlite_session: Session) -> None: - repo = _build_repository() - - recipients = repo._build_email_recipients( - session=sqlite_session, - form_id="form-id", - delivery_id="delivery-id", - recipients_config=EmailRecipients( - include_bound_group=False, - items=[ - ExternalRecipient(email="external@example.com"), - ExternalRecipient(email="external@example.com"), - ], - ), - ) - - assert len(recipients) == 1 - - def test_build_email_recipients_prefers_member_over_external_by_email(self, sqlite_session: Session) -> None: - repo = _build_repository() - _add_workspace_member(sqlite_session, user_id="member-1", email="shared@example.com") - - recipients = repo._build_email_recipients( - session=sqlite_session, - form_id="form-id", - delivery_id="delivery-id", - recipients_config=EmailRecipients( - include_bound_group=False, - items=[ - MemberRecipient(reference_id="member-1"), - ExternalRecipient(email="shared@example.com"), - ], - ), - ) - - assert len(recipients) == 1 - assert recipients[0].recipient_type == RecipientType.EMAIL_MEMBER - - def test_delivery_method_to_model_includes_external_recipients_with_whole_workspace( - self, - sqlite_session: Session, - ) -> None: - repo = _build_repository() - _add_workspace_member(sqlite_session, user_id="member-1", email="member1@example.com") - _add_workspace_member(sqlite_session, user_id="member-2", email="member2@example.com") - - method = EmailDeliveryMethod( - config=EmailDeliveryConfig( - recipients=EmailRecipients( - include_bound_group=True, - items=[ExternalRecipient(email="external@example.com")], - ), - subject="subject", - body="body", - ) - ) - - result = repo._delivery_method_to_model(session=sqlite_session, form_id="form-id", delivery_method=method) - - assert len(result.recipients) == 3 - member_emails = { - EmailMemberRecipientPayload.model_validate_json(r.recipient_payload).email - for r in result.recipients - if r.recipient_type == RecipientType.EMAIL_MEMBER - } - assert member_emails == {"member1@example.com", "member2@example.com"} - external_payload = EmailExternalRecipientPayload.model_validate_json( - next(r for r in result.recipients if r.recipient_type == RecipientType.EMAIL_EXTERNAL).recipient_payload - ) - assert external_payload.email == "external@example.com" - - -def _make_form_definition() -> str: - return FormDefinition( - form_content="hello", - inputs=[], - user_actions=[UserActionConfig(id="submit", title="Submit")], - rendered_content="

hello

", - expiration_time=naive_utc_now(), - ).model_dump_json() - - -def _make_form( - *, - form_id: str = "form-1", - workflow_run_id: str = "run-1", - node_id: str = "node-1", - tenant_id: str = "tenant-id", - selected_action_id: str | None = None, - submitted_data: str | None = None, - submitted_at: datetime | None = None, - expiration_time: datetime | None = None, -) -> HumanInputForm: - return HumanInputForm( - id=form_id, - workflow_run_id=workflow_run_id, - conversation_id=None, - node_id=node_id, - tenant_id=tenant_id, - app_id="app-id", - form_kind=HumanInputFormKind.RUNTIME, - form_definition=_make_form_definition(), - rendered_content="

hello

", - expiration_time=expiration_time or naive_utc_now(), - status=HumanInputFormStatus.WAITING, - selected_action_id=selected_action_id, - submitted_data=submitted_data, - submitted_at=submitted_at, - ) - - -def _make_recipient( - form_id: str, - *, - recipient_id: str = "recipient-1", - recipient_type: RecipientType = RecipientType.STANDALONE_WEB_APP, - access_token: str = "token-123", -) -> HumanInputFormRecipient: - return HumanInputFormRecipient( - id=recipient_id, - form_id=form_id, - delivery_id="delivery-1", - recipient_type=recipient_type, - access_token=access_token, - recipient_payload=StandaloneWebAppRecipientPayload().model_dump_json(), - ) - - -def _persist_form( - session: Session, - form: HumanInputForm, - recipients: list[HumanInputFormRecipient] | None = None, -) -> None: - session.add(form) - session.add_all(recipients or []) - session.commit() - - -class TestHumanInputFormRepositoryImplPublicMethods: - @pytest.mark.parametrize("mismatched_scope", ["tenant_id", "app_id", "workflow_run_id", "node_id", "id"]) - def test_mark_timeout_rejects_form_outside_execution_scope( - self, sqlite_session: Session, mismatched_scope: str - ) -> None: - form = _make_form(expiration_time=naive_utc_now() - timedelta(minutes=1)) - setattr(form, mismatched_scope, "other-value") - _persist_form(sqlite_session, form) - repo = HumanInputFormRepositoryImpl(tenant_id="tenant-id", app_id="app-id", workflow_execution_id="run-1") - - with pytest.raises(FormNotFoundError): - repo.mark_timeout("node-1", form_id="form-1") - - sqlite_session.refresh(form) - assert form.status == HumanInputFormStatus.WAITING - - @pytest.mark.parametrize("already_submitted", [False, True]) - def test_mark_timeout_preserves_waiting_form_that_is_not_eligible( - self, sqlite_session: Session, already_submitted: bool - ) -> None: - form = _make_form( - expiration_time=naive_utc_now() + timedelta(minutes=-1 if already_submitted else 1), - submitted_at=naive_utc_now() if already_submitted else None, - ) - _persist_form(sqlite_session, form) - repo = HumanInputFormRepositoryImpl(tenant_id="tenant-id", app_id="app-id", workflow_execution_id="run-1") - - entity = repo.mark_timeout("node-1", form_id="form-1") - - assert entity.status == HumanInputFormStatus.WAITING - assert entity.submitted == already_submitted - sqlite_session.refresh(form) - assert form.status == HumanInputFormStatus.WAITING - - def test_get_form_returns_entity_and_recipients(self, sqlite_session: Session): - form = _make_form() - recipient = _make_recipient(form.id) - other_tenant_form = _make_form(form_id="form-2", tenant_id="other-tenant") - sqlite_session.add(other_tenant_form) - _persist_form(sqlite_session, form, [recipient]) - repo = HumanInputFormRepositoryImpl(tenant_id="tenant-id", workflow_execution_id=form.workflow_run_id) - - entity = repo.get_form(form.node_id) - - assert entity is not None - assert entity.id == form.id - assert entity.submission_token == "token-123" - assert len(entity.recipients) == 1 - assert entity.recipients[0].token == "token-123" - - def test_get_form_returns_none_when_missing(self): - repo = HumanInputFormRepositoryImpl(tenant_id="tenant-id", workflow_execution_id="run-1") - - assert repo.get_form("node-1") is None - - def test_get_form_returns_unsubmitted_state(self, sqlite_session: Session): - form = _make_form() - _persist_form(sqlite_session, form) - repo = HumanInputFormRepositoryImpl(tenant_id="tenant-id", workflow_execution_id=form.workflow_run_id) - - entity = repo.get_form(form.node_id) - - assert entity is not None - assert entity.submitted is False - assert entity.selected_action_id is None - assert entity.submitted_data is None - - def test_get_form_returns_submission_when_completed(self, sqlite_session: Session): - form = _make_form( - selected_action_id="approve", - submitted_data='{"field": "value"}', - submitted_at=naive_utc_now(), - ) - _persist_form(sqlite_session, form) - repo = HumanInputFormRepositoryImpl(tenant_id="tenant-id", workflow_execution_id=form.workflow_run_id) - - entity = repo.get_form(form.node_id) - - assert entity is not None - assert entity.submitted is True - assert entity.selected_action_id == "approve" - assert entity.submitted_data == {"field": "value"} - - -class TestHumanInputFormSubmissionRepository: - def test_get_by_token_returns_record(self, sqlite_session: Session): - form = _make_form(tenant_id="tenant-1") - recipient = _make_recipient(form.id) - _persist_form(sqlite_session, form, [recipient]) - repo = HumanInputFormSubmissionRepository() - - record = repo.get_by_token("token-123") - - assert record is not None - assert record.form_id == form.id - assert record.recipient_type == RecipientType.STANDALONE_WEB_APP - assert record.submitted is False - - def test_get_by_form_id_and_recipient_type_uses_recipient(self, sqlite_session: Session): - form = _make_form(tenant_id="tenant-1") - recipient = _make_recipient(form.id) - _persist_form(sqlite_session, form, [recipient]) - repo = HumanInputFormSubmissionRepository() - - record = repo.get_by_form_id_and_recipient_type( - form_id=form.id, - recipient_type=RecipientType.STANDALONE_WEB_APP, - ) - - assert record is not None - assert record.recipient_id == recipient.id - assert record.access_token == recipient.access_token - - def test_mark_submitted_updates_fields( - self, - monkeypatch: pytest.MonkeyPatch, - sqlite_session: Session, - ): - fixed_now = datetime(2024, 1, 1, 0, 0, 0) - monkeypatch.setattr("core.repositories.human_input_repository.naive_utc_now", lambda: fixed_now) - - form = _make_form(tenant_id="tenant-1", expiration_time=fixed_now) - recipient = _make_recipient(form.id) - _persist_form(sqlite_session, form, [recipient]) - repo = HumanInputFormSubmissionRepository() - - record: HumanInputFormRecord = repo.mark_submitted( - form_id=form.id, - recipient_id=recipient.id, - selected_action_id="approve", - form_data={"field": "value"}, - submission_user_id="user-1", - submission_end_user_id="end-user-1", - ) - - sqlite_session.expire_all() - persisted_form = sqlite_session.get(HumanInputForm, form.id) - assert persisted_form is not None - assert persisted_form.selected_action_id == "approve" - assert persisted_form.completed_by_recipient_id == recipient.id - assert persisted_form.submission_user_id == "user-1" - assert persisted_form.submission_end_user_id == "end-user-1" - assert persisted_form.submitted_at == fixed_now - assert record.submitted is True - assert record.selected_action_id == "approve" - assert record.submitted_data == {"field": "value"} - - def test_mark_submitted_rolls_back_on_database_failure( - self, - sqlite_session: Session, - sqlite_engine: Engine, - ) -> None: - form = _make_form(tenant_id="tenant-1") - recipient = _make_recipient(form.id) - _persist_form(sqlite_session, form, [recipient]) - repo = HumanInputFormSubmissionRepository() - - def fail_form_update(_connection, _cursor, statement, _parameters, _context, _executemany) -> None: - if statement.lstrip().upper().startswith("UPDATE HUMAN_INPUT_FORMS"): - raise SQLAlchemyError("forced update failure") - - event.listen(sqlite_engine, "before_cursor_execute", fail_form_update) - try: - with pytest.raises(SQLAlchemyError, match="forced update failure"): - repo.mark_submitted( - form_id=form.id, - recipient_id=recipient.id, - selected_action_id="approve", - form_data={"field": "value"}, - submission_user_id="user-1", - submission_end_user_id="end-user-1", - ) - finally: - event.remove(sqlite_engine, "before_cursor_execute", fail_form_update) - - sqlite_session.expire_all() - persisted_form = sqlite_session.get(HumanInputForm, form.id) - assert persisted_form is not None - assert persisted_form.status == HumanInputFormStatus.WAITING - assert persisted_form.selected_action_id is None - assert persisted_form.submitted_data is None - assert persisted_form.submitted_at is None diff --git a/api/tests/unit_tests/core/repositories/test_human_input_repository.py b/api/tests/unit_tests/core/repositories/test_human_input_repository.py index 97fddc97a8e..c8573cb313e 100644 --- a/api/tests/unit_tests/core/repositories/test_human_input_repository.py +++ b/api/tests/unit_tests/core/repositories/test_human_input_repository.py @@ -7,6 +7,7 @@ from typing import Any import pytest from sqlalchemy import Engine, event, select +from sqlalchemy.exc import SQLAlchemyError from sqlalchemy.orm import Session, sessionmaker from core.repositories.human_input_repository import ( @@ -29,39 +30,452 @@ from core.workflow.human_input_adapter import ( MemberRecipient, WebAppDeliveryMethod, ) -from core.workflow.nodes.human_input.entities import HumanInputNodeData, UserActionConfig +from core.workflow.nodes.human_input.entities import ( + FormDefinition, + HumanInputNodeData, + UserActionConfig, +) from core.workflow.nodes.human_input.enums import HumanInputFormKind, HumanInputFormStatus from libs.datetime_utils import naive_utc_now -from models.account import Account, TenantAccountJoin, TenantAccountRole +from models import Account as _human_input_form_repository_impl_Account +from models import TenantAccountJoin as _human_input_form_repository_impl_TenantAccountJoin +from models.account import Account as _human_input_repository_Account +from models.account import TenantAccountJoin as _human_input_repository_TenantAccountJoin +from models.account import TenantAccountRole from models.base import TypeBase from models.human_input import ( + EmailExternalRecipientPayload, + EmailMemberRecipientPayload, HumanInputDelivery, HumanInputForm, HumanInputFormRecipient, RecipientType, + StandaloneWebAppRecipientPayload, ) from repositories.human_input_errors import FormAlreadyHandledError +"""Unit tests for HumanInputFormRepositoryImpl private helpers.""" -@pytest.fixture -def repository_session(sqlite_engine: Engine, monkeypatch: pytest.MonkeyPatch) -> Iterator[Session]: - """Bind repository-owned sessions to an isolated SQLite database.""" - tables = [ - HumanInputForm.__table__, - HumanInputDelivery.__table__, - HumanInputFormRecipient.__table__, - Account.__table__, - TenantAccountJoin.__table__, - ] - TypeBase.metadata.create_all(sqlite_engine, tables=tables) - repository_session_factory = sessionmaker(bind=sqlite_engine, expire_on_commit=False) - monkeypatch.setattr( - "core.repositories.human_input_repository.session_factory.create_session", - repository_session_factory, +def _build_repository() -> HumanInputFormRepositoryImpl: + return HumanInputFormRepositoryImpl(tenant_id="tenant-id") + + +def _add_workspace_member( + session: Session, + *, + user_id: str, + email: str, + tenant_id: str = "tenant-id", +) -> None: + account = _human_input_form_repository_impl_Account(name=user_id, email=email) + account.id = user_id + session.add_all( + [account, _human_input_form_repository_impl_TenantAccountJoin(tenant_id=tenant_id, account_id=user_id)] ) - with repository_session_factory() as session: - yield session + session.commit() + + +def _make_form_definition() -> str: + return FormDefinition( + form_content="hello", + inputs=[], + user_actions=[UserActionConfig(id="submit", title="Submit")], + rendered_content="

hello

", + expiration_time=naive_utc_now(), + ).model_dump_json() + + +def _make_form( + *, + form_id: str = "form-1", + workflow_run_id: str = "run-1", + node_id: str = "node-1", + tenant_id: str = "tenant-id", + selected_action_id: str | None = None, + submitted_data: str | None = None, + submitted_at: datetime | None = None, + expiration_time: datetime | None = None, +) -> HumanInputForm: + return HumanInputForm( + id=form_id, + workflow_run_id=workflow_run_id, + conversation_id=None, + node_id=node_id, + tenant_id=tenant_id, + app_id="app-id", + form_kind=HumanInputFormKind.RUNTIME, + form_definition=_make_form_definition(), + rendered_content="

hello

", + expiration_time=expiration_time or naive_utc_now(), + status=HumanInputFormStatus.WAITING, + selected_action_id=selected_action_id, + submitted_data=submitted_data, + submitted_at=submitted_at, + ) + + +def _make_recipient( + form_id: str, + *, + recipient_id: str = "recipient-1", + recipient_type: RecipientType = RecipientType.STANDALONE_WEB_APP, + access_token: str = "token-123", +) -> HumanInputFormRecipient: + return HumanInputFormRecipient( + id=recipient_id, + form_id=form_id, + delivery_id="delivery-1", + recipient_type=recipient_type, + access_token=access_token, + recipient_payload=StandaloneWebAppRecipientPayload().model_dump_json(), + ) + + +def _human_input_form_repository_impl_persist_form( + session: Session, + form: HumanInputForm, + recipients: list[HumanInputFormRecipient] | None = None, +) -> None: + session.add(form) + session.add_all(recipients or []) + session.commit() + + +class TestHumanInputFormRepositoryImplCases: + class TestHumanInputFormRepositoryImplHelpers: + def test_build_email_recipients_with_member_and_external( + self, + monkeypatch: pytest.MonkeyPatch, + sqlite_session: Session, + ) -> None: + repo = _build_repository() + _add_workspace_member(sqlite_session, user_id="member-1", email="member@example.com") + query_workspace_members_by_ids = repo._query_workspace_members_by_ids + + def query_members(session: Session, restrict_to_user_ids: list[str]) -> list[_WorkspaceMemberInfo]: + assert session is sqlite_session + return query_workspace_members_by_ids( + session=session, + restrict_to_user_ids=restrict_to_user_ids, + ) + + monkeypatch.setattr(repo, "_query_workspace_members_by_ids", query_members) + + recipients = repo._build_email_recipients( + session=sqlite_session, + form_id="form-id", + delivery_id="delivery-id", + recipients_config=EmailRecipients( + include_bound_group=False, + items=[ + MemberRecipient(reference_id="member-1"), + ExternalRecipient(email="external@example.com"), + ], + ), + ) + + assert len(recipients) == 2 + member_recipient = next(r for r in recipients if r.recipient_type == RecipientType.EMAIL_MEMBER) + external_recipient = next(r for r in recipients if r.recipient_type == RecipientType.EMAIL_EXTERNAL) + + member_payload = EmailMemberRecipientPayload.model_validate_json(member_recipient.recipient_payload) + assert member_payload.user_id == "member-1" + assert member_payload.email == "member@example.com" + + external_payload = EmailExternalRecipientPayload.model_validate_json(external_recipient.recipient_payload) + assert external_payload.email == "external@example.com" + + def test_build_email_recipients_skips_unknown_members(self, sqlite_session: Session) -> None: + repo = _build_repository() + + recipients = repo._build_email_recipients( + session=sqlite_session, + form_id="form-id", + delivery_id="delivery-id", + recipients_config=EmailRecipients( + include_bound_group=False, + items=[ + MemberRecipient(reference_id="missing-member"), + ExternalRecipient(email="external@example.com"), + ], + ), + ) + + assert len(recipients) == 1 + assert recipients[0].recipient_type == RecipientType.EMAIL_EXTERNAL + + def test_build_email_recipients_whole_workspace_uses_all_members(self, sqlite_session: Session) -> None: + repo = _build_repository() + _add_workspace_member(sqlite_session, user_id="member-1", email="member1@example.com") + _add_workspace_member(sqlite_session, user_id="member-2", email="member2@example.com") + + recipients = repo._build_email_recipients( + session=sqlite_session, + form_id="form-id", + delivery_id="delivery-id", + recipients_config=EmailRecipients( + include_bound_group=True, + items=[], + ), + ) + + assert len(recipients) == 2 + emails = {EmailMemberRecipientPayload.model_validate_json(r.recipient_payload).email for r in recipients} + assert emails == {"member1@example.com", "member2@example.com"} + + def test_build_email_recipients_dedupes_external_by_email(self, sqlite_session: Session) -> None: + repo = _build_repository() + + recipients = repo._build_email_recipients( + session=sqlite_session, + form_id="form-id", + delivery_id="delivery-id", + recipients_config=EmailRecipients( + include_bound_group=False, + items=[ + ExternalRecipient(email="external@example.com"), + ExternalRecipient(email="external@example.com"), + ], + ), + ) + + assert len(recipients) == 1 + + def test_build_email_recipients_prefers_member_over_external_by_email(self, sqlite_session: Session) -> None: + repo = _build_repository() + _add_workspace_member(sqlite_session, user_id="member-1", email="shared@example.com") + + recipients = repo._build_email_recipients( + session=sqlite_session, + form_id="form-id", + delivery_id="delivery-id", + recipients_config=EmailRecipients( + include_bound_group=False, + items=[ + MemberRecipient(reference_id="member-1"), + ExternalRecipient(email="shared@example.com"), + ], + ), + ) + + assert len(recipients) == 1 + assert recipients[0].recipient_type == RecipientType.EMAIL_MEMBER + + def test_delivery_method_to_model_includes_external_recipients_with_whole_workspace( + self, + sqlite_session: Session, + ) -> None: + repo = _build_repository() + _add_workspace_member(sqlite_session, user_id="member-1", email="member1@example.com") + _add_workspace_member(sqlite_session, user_id="member-2", email="member2@example.com") + + method = EmailDeliveryMethod( + config=EmailDeliveryConfig( + recipients=EmailRecipients( + include_bound_group=True, + items=[ExternalRecipient(email="external@example.com")], + ), + subject="subject", + body="body", + ) + ) + + result = repo._delivery_method_to_model(session=sqlite_session, form_id="form-id", delivery_method=method) + + assert len(result.recipients) == 3 + member_emails = { + EmailMemberRecipientPayload.model_validate_json(r.recipient_payload).email + for r in result.recipients + if r.recipient_type == RecipientType.EMAIL_MEMBER + } + assert member_emails == {"member1@example.com", "member2@example.com"} + external_payload = EmailExternalRecipientPayload.model_validate_json( + next(r for r in result.recipients if r.recipient_type == RecipientType.EMAIL_EXTERNAL).recipient_payload + ) + assert external_payload.email == "external@example.com" + + class TestHumanInputFormRepositoryImplPublicMethods: + @pytest.mark.parametrize("mismatched_scope", ["tenant_id", "app_id", "workflow_run_id", "node_id", "id"]) + def test_mark_timeout_rejects_form_outside_execution_scope( + self, sqlite_session: Session, mismatched_scope: str + ) -> None: + form = _make_form(expiration_time=naive_utc_now() - timedelta(minutes=1)) + setattr(form, mismatched_scope, "other-value") + _human_input_form_repository_impl_persist_form(sqlite_session, form) + repo = HumanInputFormRepositoryImpl(tenant_id="tenant-id", app_id="app-id", workflow_execution_id="run-1") + + with pytest.raises(FormNotFoundError): + repo.mark_timeout("node-1", form_id="form-1") + + sqlite_session.refresh(form) + assert form.status == HumanInputFormStatus.WAITING + + @pytest.mark.parametrize("already_submitted", [False, True]) + def test_mark_timeout_preserves_waiting_form_that_is_not_eligible( + self, sqlite_session: Session, already_submitted: bool + ) -> None: + form = _make_form( + expiration_time=naive_utc_now() + timedelta(minutes=-1 if already_submitted else 1), + submitted_at=naive_utc_now() if already_submitted else None, + ) + _human_input_form_repository_impl_persist_form(sqlite_session, form) + repo = HumanInputFormRepositoryImpl(tenant_id="tenant-id", app_id="app-id", workflow_execution_id="run-1") + + entity = repo.mark_timeout("node-1", form_id="form-1") + + assert entity.status == HumanInputFormStatus.WAITING + assert entity.submitted == already_submitted + sqlite_session.refresh(form) + assert form.status == HumanInputFormStatus.WAITING + + def test_get_form_returns_entity_and_recipients(self, sqlite_session: Session): + form = _make_form() + recipient = _make_recipient(form.id) + other_tenant_form = _make_form(form_id="form-2", tenant_id="other-tenant") + sqlite_session.add(other_tenant_form) + _human_input_form_repository_impl_persist_form(sqlite_session, form, [recipient]) + repo = HumanInputFormRepositoryImpl(tenant_id="tenant-id", workflow_execution_id=form.workflow_run_id) + + entity = repo.get_form(form.node_id) + + assert entity is not None + assert entity.id == form.id + assert entity.submission_token == "token-123" + assert len(entity.recipients) == 1 + assert entity.recipients[0].token == "token-123" + + def test_get_form_returns_none_when_missing(self): + repo = HumanInputFormRepositoryImpl(tenant_id="tenant-id", workflow_execution_id="run-1") + + assert repo.get_form("node-1") is None + + def test_get_form_returns_unsubmitted_state(self, sqlite_session: Session): + form = _make_form() + _human_input_form_repository_impl_persist_form(sqlite_session, form) + repo = HumanInputFormRepositoryImpl(tenant_id="tenant-id", workflow_execution_id=form.workflow_run_id) + + entity = repo.get_form(form.node_id) + + assert entity is not None + assert entity.submitted is False + assert entity.selected_action_id is None + assert entity.submitted_data is None + + def test_get_form_returns_submission_when_completed(self, sqlite_session: Session): + form = _make_form( + selected_action_id="approve", + submitted_data='{"field": "value"}', + submitted_at=naive_utc_now(), + ) + _human_input_form_repository_impl_persist_form(sqlite_session, form) + repo = HumanInputFormRepositoryImpl(tenant_id="tenant-id", workflow_execution_id=form.workflow_run_id) + + entity = repo.get_form(form.node_id) + + assert entity is not None + assert entity.submitted is True + assert entity.selected_action_id == "approve" + assert entity.submitted_data == {"field": "value"} + + class TestHumanInputFormSubmissionRepository: + def test_get_by_token_returns_record(self, sqlite_session: Session): + form = _make_form(tenant_id="tenant-1") + recipient = _make_recipient(form.id) + _human_input_form_repository_impl_persist_form(sqlite_session, form, [recipient]) + repo = HumanInputFormSubmissionRepository() + + record = repo.get_by_token("token-123") + + assert record is not None + assert record.form_id == form.id + assert record.recipient_type == RecipientType.STANDALONE_WEB_APP + assert record.submitted is False + + def test_get_by_form_id_and_recipient_type_uses_recipient(self, sqlite_session: Session): + form = _make_form(tenant_id="tenant-1") + recipient = _make_recipient(form.id) + _human_input_form_repository_impl_persist_form(sqlite_session, form, [recipient]) + repo = HumanInputFormSubmissionRepository() + + record = repo.get_by_form_id_and_recipient_type( + form_id=form.id, + recipient_type=RecipientType.STANDALONE_WEB_APP, + ) + + assert record is not None + assert record.recipient_id == recipient.id + assert record.access_token == recipient.access_token + + def test_mark_submitted_updates_fields( + self, + monkeypatch: pytest.MonkeyPatch, + sqlite_session: Session, + ): + fixed_now = datetime(2024, 1, 1, 0, 0, 0) + monkeypatch.setattr("core.repositories.human_input_repository.naive_utc_now", lambda: fixed_now) + + form = _make_form(tenant_id="tenant-1", expiration_time=fixed_now) + recipient = _make_recipient(form.id) + _human_input_form_repository_impl_persist_form(sqlite_session, form, [recipient]) + repo = HumanInputFormSubmissionRepository() + + record: HumanInputFormRecord = repo.mark_submitted( + form_id=form.id, + recipient_id=recipient.id, + selected_action_id="approve", + form_data={"field": "value"}, + submission_user_id="user-1", + submission_end_user_id="end-user-1", + ) + + sqlite_session.expire_all() + persisted_form = sqlite_session.get(HumanInputForm, form.id) + assert persisted_form is not None + assert persisted_form.selected_action_id == "approve" + assert persisted_form.completed_by_recipient_id == recipient.id + assert persisted_form.submission_user_id == "user-1" + assert persisted_form.submission_end_user_id == "end-user-1" + assert persisted_form.submitted_at == fixed_now + assert record.submitted is True + assert record.selected_action_id == "approve" + assert record.submitted_data == {"field": "value"} + + def test_mark_submitted_rolls_back_on_database_failure( + self, + sqlite_session: Session, + sqlite_engine: Engine, + ) -> None: + form = _make_form(tenant_id="tenant-1") + recipient = _make_recipient(form.id) + _human_input_form_repository_impl_persist_form(sqlite_session, form, [recipient]) + repo = HumanInputFormSubmissionRepository() + + def fail_form_update(_connection, _cursor, statement, _parameters, _context, _executemany) -> None: + if statement.lstrip().upper().startswith("UPDATE HUMAN_INPUT_FORMS"): + raise SQLAlchemyError("forced update failure") + + event.listen(sqlite_engine, "before_cursor_execute", fail_form_update) + try: + with pytest.raises(SQLAlchemyError, match="forced update failure"): + repo.mark_submitted( + form_id=form.id, + recipient_id=recipient.id, + selected_action_id="approve", + form_data={"field": "value"}, + submission_user_id="user-1", + submission_end_user_id="end-user-1", + ) + finally: + event.remove(sqlite_engine, "before_cursor_execute", fail_form_update) + + sqlite_session.expire_all() + persisted_form = sqlite_session.get(HumanInputForm, form.id) + assert persisted_form is not None + assert persisted_form.status == HumanInputFormStatus.WAITING + assert persisted_form.selected_action_id is None + assert persisted_form.submitted_data is None + assert persisted_form.submitted_at is None def _make_form_definition_json(*, include_expiration_time: bool) -> str: @@ -76,7 +490,7 @@ def _make_form_definition_json(*, include_expiration_time: bool) -> str: return json.dumps(payload, default=str) -def _persist_form( +def _human_input_repository_persist_form( session: Session, *, form_id: str = "form-1", @@ -131,577 +545,619 @@ def _persist_recipient( return recipient -def test_recipient_entity_token_raises_when_missing() -> None: - recipient = HumanInputFormRecipient( - id="r1", - form_id="f1", - delivery_id="d1", - recipient_type=RecipientType.CONSOLE, - recipient_payload="{}", - access_token=None, - ) - entity = _HumanInputFormRecipientEntityImpl(recipient) - with pytest.raises(AssertionError, match="access_token should not be None"): - _ = entity.token +class TestHumanInputRepositoryCases: + @staticmethod + @pytest.fixture + def repository_session(sqlite_engine: Engine, monkeypatch: pytest.MonkeyPatch) -> Iterator[Session]: + """Bind repository-owned sessions to an isolated SQLite database.""" - -def test_recipient_entity_id_and_token_success(repository_session: Session) -> None: - form = _persist_form(repository_session) - recipient = _persist_recipient(repository_session, form_id=form.id, recipient_id="r1", access_token="tok") - entity = _HumanInputFormRecipientEntityImpl(recipient) - assert entity.id == "r1" - assert entity.token == "tok" - - -def test_form_entity_submission_token_prefers_console_then_webapp_then_none(repository_session: Session) -> None: - form = _persist_form(repository_session, form_id="f1") - console = _persist_recipient( - repository_session, - form_id=form.id, - recipient_id="c1", - recipient_type=RecipientType.CONSOLE, - access_token="ctok", - ) - webapp = _persist_recipient( - repository_session, - form_id=form.id, - recipient_id="w1", - recipient_type=RecipientType.STANDALONE_WEB_APP, - access_token="wtok", - ) - - entity = _HumanInputFormEntityImpl(form_model=form, recipient_models=[webapp, console]) - assert entity.submission_token == "ctok" - - entity = _HumanInputFormEntityImpl(form_model=form, recipient_models=[webapp]) - assert entity.submission_token == "wtok" - - entity = _HumanInputFormEntityImpl(form_model=form, recipient_models=[]) - assert entity.submission_token is None - - -def test_form_entity_submitted_data_parsed(repository_session: Session) -> None: - form = _persist_form(repository_session, form_id="f1") - form.submitted_data = '{"a": 1}' - form.submitted_at = naive_utc_now() - repository_session.commit() - entity = _HumanInputFormEntityImpl(form_model=form, recipient_models=[]) - assert entity.submitted is True - assert entity.submitted_data == {"a": 1} - assert entity.rendered_content == "

x

" - assert entity.selected_action_id is None - assert entity.status == HumanInputFormStatus.WAITING - - -def test_form_record_from_models_injects_expiration_time_when_missing(repository_session: Session) -> None: - expiration = naive_utc_now() - form = _persist_form(repository_session, form_id="f1", workflow_run_id=None) - form.form_definition = _make_form_definition_json(include_expiration_time=False) - form.expiration_time = expiration - form.submitted_data = '{"k": "v"}' - repository_session.commit() - record = HumanInputFormRecord.from_models(form, None) - assert record.definition.expiration_time == expiration - assert record.submitted_data == {"k": "v"} - assert record.submitted is False - - -def test_create_email_recipients_from_resolved_dedupes_and_skips_blank() -> None: - repo = HumanInputFormRepositoryImpl(tenant_id="tenant") - recipients = repo._create_email_recipients_from_resolved( # type: ignore[attr-defined] - form_id="f", - delivery_id="d", - members=[ - _WorkspaceMemberInfo(user_id="u1", email=""), - _WorkspaceMemberInfo(user_id="u2", email="a@example.com"), - _WorkspaceMemberInfo(user_id="u3", email="a@example.com"), - ], - external_emails=["", "a@example.com", "b@example.com", "b@example.com"], - ) - assert [r.recipient_type for r in recipients] == [RecipientType.EMAIL_MEMBER, RecipientType.EMAIL_EXTERNAL] - - -def test_query_workspace_members_by_ids_empty_returns_empty(repository_session: Session) -> None: - repo = HumanInputFormRepositoryImpl(tenant_id="tenant") - assert repo._query_workspace_members_by_ids(session=repository_session, restrict_to_user_ids=["", ""]) == [] - - -def test_query_workspace_members_by_ids_maps_rows_and_scopes_tenant(repository_session: Session) -> None: - accounts = [ - Account(name="One", email="a@example.com"), - Account(name="Two", email="b@example.com"), - Account(name="Other", email="other@example.com"), - ] - repository_session.add_all(accounts) - repository_session.flush() - repository_session.add_all( - [ - TenantAccountJoin(tenant_id="tenant", account_id=accounts[0].id, role=TenantAccountRole.NORMAL), - TenantAccountJoin(tenant_id="tenant", account_id=accounts[1].id, role=TenantAccountRole.NORMAL), - TenantAccountJoin(tenant_id="other-tenant", account_id=accounts[2].id, role=TenantAccountRole.NORMAL), + tables = [ + HumanInputForm.__table__, + HumanInputDelivery.__table__, + HumanInputFormRecipient.__table__, + _human_input_repository_Account.__table__, + _human_input_repository_TenantAccountJoin.__table__, ] - ) - repository_session.commit() - repo = HumanInputFormRepositoryImpl(tenant_id="tenant") - rows = repo._query_workspace_members_by_ids( - session=repository_session, - restrict_to_user_ids=[accounts[0].id, accounts[1].id, accounts[2].id], - ) - assert set(rows) == { - _WorkspaceMemberInfo(user_id=accounts[0].id, email="a@example.com"), - _WorkspaceMemberInfo(user_id=accounts[1].id, email="b@example.com"), - } - - -def test_query_all_workspace_members_maps_rows(repository_session: Session) -> None: - account = Account(name="One", email="a@example.com") - repository_session.add(account) - repository_session.flush() - repository_session.add(TenantAccountJoin(tenant_id="tenant", account_id=account.id, role=TenantAccountRole.NORMAL)) - repository_session.commit() - repo = HumanInputFormRepositoryImpl(tenant_id="tenant") - rows = repo._query_all_workspace_members(session=repository_session) - assert rows == [_WorkspaceMemberInfo(user_id=account.id, email="a@example.com")] - - -def test_repository_init_sets_tenant_id() -> None: - repo = HumanInputFormRepositoryImpl(tenant_id="tenant") - assert repo._tenant_id == "tenant" - - -def test_delivery_method_to_model_webapp_creates_delivery_and_recipient( - repository_session: Session, monkeypatch: pytest.MonkeyPatch -) -> None: - repo = HumanInputFormRepositoryImpl(tenant_id="tenant") - monkeypatch.setattr("core.repositories.human_input_repository.uuidv7", lambda: "del-1") - result = repo._delivery_method_to_model( - session=repository_session, form_id="form-1", delivery_method=WebAppDeliveryMethod() - ) - assert result.delivery.id == "del-1" - assert result.delivery.form_id == "form-1" - assert len(result.recipients) == 1 - assert result.recipients[0].recipient_type == RecipientType.STANDALONE_WEB_APP - - -def test_delivery_method_to_model_email_uses_build_email_recipients( - repository_session: Session, monkeypatch: pytest.MonkeyPatch -) -> None: - repo = HumanInputFormRepositoryImpl(tenant_id="tenant") - monkeypatch.setattr("core.repositories.human_input_repository.uuidv7", lambda: "del-1") - called: dict[str, Any] = {} - - def fake_build(*, session: Session, form_id: str, delivery_id: str, recipients_config: Any) -> list[Any]: - called.update( - {"session": session, "form_id": form_id, "delivery_id": delivery_id, "recipients_config": recipients_config} + TypeBase.metadata.create_all(sqlite_engine, tables=tables) + repository_session_factory = sessionmaker(bind=sqlite_engine, expire_on_commit=False) + monkeypatch.setattr( + "core.repositories.human_input_repository.session_factory.create_session", + repository_session_factory, ) - return ["r"] + with repository_session_factory() as session: + yield session - monkeypatch.setattr(repo, "_build_email_recipients", fake_build) + @staticmethod + def test_recipient_entity_token_raises_when_missing() -> None: + recipient = HumanInputFormRecipient( + id="r1", + form_id="f1", + delivery_id="d1", + recipient_type=RecipientType.CONSOLE, + recipient_payload="{}", + access_token=None, + ) + entity = _HumanInputFormRecipientEntityImpl(recipient) + with pytest.raises(AssertionError, match="access_token should not be None"): + _ = entity.token - method = EmailDeliveryMethod( - config=EmailDeliveryConfig( - recipients=EmailRecipients( + @staticmethod + def test_recipient_entity_id_and_token_success(repository_session: Session) -> None: + form = _human_input_repository_persist_form(repository_session) + recipient = _persist_recipient(repository_session, form_id=form.id, recipient_id="r1", access_token="tok") + entity = _HumanInputFormRecipientEntityImpl(recipient) + assert entity.id == "r1" + assert entity.token == "tok" + + @staticmethod + def test_form_entity_submission_token_prefers_console_then_webapp_then_none(repository_session: Session) -> None: + form = _human_input_repository_persist_form(repository_session, form_id="f1") + console = _persist_recipient( + repository_session, + form_id=form.id, + recipient_id="c1", + recipient_type=RecipientType.CONSOLE, + access_token="ctok", + ) + webapp = _persist_recipient( + repository_session, + form_id=form.id, + recipient_id="w1", + recipient_type=RecipientType.STANDALONE_WEB_APP, + access_token="wtok", + ) + + entity = _HumanInputFormEntityImpl(form_model=form, recipient_models=[webapp, console]) + assert entity.submission_token == "ctok" + + entity = _HumanInputFormEntityImpl(form_model=form, recipient_models=[webapp]) + assert entity.submission_token == "wtok" + + entity = _HumanInputFormEntityImpl(form_model=form, recipient_models=[]) + assert entity.submission_token is None + + @staticmethod + def test_form_entity_submitted_data_parsed(repository_session: Session) -> None: + form = _human_input_repository_persist_form(repository_session, form_id="f1") + form.submitted_data = '{"a": 1}' + form.submitted_at = naive_utc_now() + repository_session.commit() + entity = _HumanInputFormEntityImpl(form_model=form, recipient_models=[]) + assert entity.submitted is True + assert entity.submitted_data == {"a": 1} + assert entity.rendered_content == "

x

" + assert entity.selected_action_id is None + assert entity.status == HumanInputFormStatus.WAITING + + @staticmethod + def test_form_record_from_models_injects_expiration_time_when_missing(repository_session: Session) -> None: + expiration = naive_utc_now() + form = _human_input_repository_persist_form(repository_session, form_id="f1", workflow_run_id=None) + form.form_definition = _make_form_definition_json(include_expiration_time=False) + form.expiration_time = expiration + form.submitted_data = '{"k": "v"}' + repository_session.commit() + record = HumanInputFormRecord.from_models(form, None) + assert record.definition.expiration_time == expiration + assert record.submitted_data == {"k": "v"} + assert record.submitted is False + + @staticmethod + def test_create_email_recipients_from_resolved_dedupes_and_skips_blank() -> None: + repo = HumanInputFormRepositoryImpl(tenant_id="tenant") + recipients = repo._create_email_recipients_from_resolved( # type: ignore[attr-defined] + form_id="f", + delivery_id="d", + members=[ + _WorkspaceMemberInfo(user_id="u1", email=""), + _WorkspaceMemberInfo(user_id="u2", email="a@example.com"), + _WorkspaceMemberInfo(user_id="u3", email="a@example.com"), + ], + external_emails=["", "a@example.com", "b@example.com", "b@example.com"], + ) + assert [r.recipient_type for r in recipients] == [RecipientType.EMAIL_MEMBER, RecipientType.EMAIL_EXTERNAL] + + @staticmethod + def test_query_workspace_members_by_ids_empty_returns_empty(repository_session: Session) -> None: + repo = HumanInputFormRepositoryImpl(tenant_id="tenant") + assert repo._query_workspace_members_by_ids(session=repository_session, restrict_to_user_ids=["", ""]) == [] + + @staticmethod + def test_query_workspace_members_by_ids_maps_rows_and_scopes_tenant(repository_session: Session) -> None: + accounts = [ + _human_input_repository_Account(name="One", email="a@example.com"), + _human_input_repository_Account(name="Two", email="b@example.com"), + _human_input_repository_Account(name="Other", email="other@example.com"), + ] + repository_session.add_all(accounts) + repository_session.flush() + repository_session.add_all( + [ + _human_input_repository_TenantAccountJoin( + tenant_id="tenant", account_id=accounts[0].id, role=TenantAccountRole.NORMAL + ), + _human_input_repository_TenantAccountJoin( + tenant_id="tenant", account_id=accounts[1].id, role=TenantAccountRole.NORMAL + ), + _human_input_repository_TenantAccountJoin( + tenant_id="other-tenant", account_id=accounts[2].id, role=TenantAccountRole.NORMAL + ), + ] + ) + repository_session.commit() + repo = HumanInputFormRepositoryImpl(tenant_id="tenant") + rows = repo._query_workspace_members_by_ids( + session=repository_session, + restrict_to_user_ids=[accounts[0].id, accounts[1].id, accounts[2].id], + ) + assert set(rows) == { + _WorkspaceMemberInfo(user_id=accounts[0].id, email="a@example.com"), + _WorkspaceMemberInfo(user_id=accounts[1].id, email="b@example.com"), + } + + @staticmethod + def test_query_all_workspace_members_maps_rows(repository_session: Session) -> None: + account = _human_input_repository_Account(name="One", email="a@example.com") + repository_session.add(account) + repository_session.flush() + repository_session.add( + _human_input_repository_TenantAccountJoin( + tenant_id="tenant", account_id=account.id, role=TenantAccountRole.NORMAL + ) + ) + repository_session.commit() + repo = HumanInputFormRepositoryImpl(tenant_id="tenant") + rows = repo._query_all_workspace_members(session=repository_session) + assert rows == [_WorkspaceMemberInfo(user_id=account.id, email="a@example.com")] + + @staticmethod + def test_repository_init_sets_tenant_id() -> None: + repo = HumanInputFormRepositoryImpl(tenant_id="tenant") + assert repo._tenant_id == "tenant" + + @staticmethod + def test_delivery_method_to_model_webapp_creates_delivery_and_recipient( + repository_session: Session, monkeypatch: pytest.MonkeyPatch + ) -> None: + repo = HumanInputFormRepositoryImpl(tenant_id="tenant") + monkeypatch.setattr("core.repositories.human_input_repository.uuidv7", lambda: "del-1") + result = repo._delivery_method_to_model( + session=repository_session, form_id="form-1", delivery_method=WebAppDeliveryMethod() + ) + assert result.delivery.id == "del-1" + assert result.delivery.form_id == "form-1" + assert len(result.recipients) == 1 + assert result.recipients[0].recipient_type == RecipientType.STANDALONE_WEB_APP + + @staticmethod + def test_delivery_method_to_model_email_uses_build_email_recipients( + repository_session: Session, monkeypatch: pytest.MonkeyPatch + ) -> None: + repo = HumanInputFormRepositoryImpl(tenant_id="tenant") + monkeypatch.setattr("core.repositories.human_input_repository.uuidv7", lambda: "del-1") + called: dict[str, Any] = {} + + def fake_build(*, session: Session, form_id: str, delivery_id: str, recipients_config: Any) -> list[Any]: + called.update( + { + "session": session, + "form_id": form_id, + "delivery_id": delivery_id, + "recipients_config": recipients_config, + } + ) + return ["r"] + + monkeypatch.setattr(repo, "_build_email_recipients", fake_build) + + method = EmailDeliveryMethod( + config=EmailDeliveryConfig( + recipients=EmailRecipients( + include_bound_group=False, + items=[MemberRecipient(reference_id="u1"), ExternalRecipient(email="e@example.com")], + ), + subject="s", + body="b", + ) + ) + result = repo._delivery_method_to_model(session=repository_session, form_id="form-1", delivery_method=method) + assert result.recipients == ["r"] + assert called["session"] is repository_session + assert called["delivery_id"] == "del-1" + + @staticmethod + def test_build_email_recipients_uses_all_members_when_whole_workspace( + repository_session: Session, monkeypatch: pytest.MonkeyPatch + ) -> None: + repo = HumanInputFormRepositoryImpl(tenant_id="tenant") + monkeypatch.setattr( + repo, + "_query_all_workspace_members", + lambda *, session: [_WorkspaceMemberInfo(user_id="u", email="a@example.com")], + ) + monkeypatch.setattr(repo, "_create_email_recipients_from_resolved", lambda **_: ["ok"]) + recipients = repo._build_email_recipients( + session=repository_session, + form_id="f", + delivery_id="d", + recipients_config=EmailRecipients( + include_bound_group=True, items=[ExternalRecipient(email="e@example.com")] + ), + ) + assert recipients == ["ok"] + + @staticmethod + def test_build_email_recipients_uses_selected_members_when_not_whole_workspace( + repository_session: Session, monkeypatch: pytest.MonkeyPatch + ) -> None: + repo = HumanInputFormRepositoryImpl(tenant_id="tenant") + + def fake_query(*, session: Any, restrict_to_user_ids: Sequence[str]) -> list[_WorkspaceMemberInfo]: + assert restrict_to_user_ids == ["u1"] + return [_WorkspaceMemberInfo(user_id="u1", email="a@example.com")] + + monkeypatch.setattr(repo, "_query_workspace_members_by_ids", fake_query) + monkeypatch.setattr(repo, "_create_email_recipients_from_resolved", lambda **_: ["ok"]) + recipients = repo._build_email_recipients( + session=repository_session, + form_id="f", + delivery_id="d", + recipients_config=EmailRecipients( include_bound_group=False, items=[MemberRecipient(reference_id="u1"), ExternalRecipient(email="e@example.com")], ), - subject="s", - body="b", ) - ) - result = repo._delivery_method_to_model(session=repository_session, form_id="form-1", delivery_method=method) - assert result.recipients == ["r"] - assert called["session"] is repository_session - assert called["delivery_id"] == "del-1" + assert recipients == ["ok"] + @staticmethod + def test_get_form_returns_entity_and_none_when_missing(repository_session: Session) -> None: + repo = HumanInputFormRepositoryImpl(tenant_id="tenant", workflow_execution_id="run") + assert repo.get_form("node") is None -def test_build_email_recipients_uses_all_members_when_whole_workspace( - repository_session: Session, monkeypatch: pytest.MonkeyPatch -) -> None: - repo = HumanInputFormRepositoryImpl(tenant_id="tenant") - monkeypatch.setattr( - repo, - "_query_all_workspace_members", - lambda *, session: [_WorkspaceMemberInfo(user_id="u", email="a@example.com")], - ) - monkeypatch.setattr(repo, "_create_email_recipients_from_resolved", lambda **_: ["ok"]) - recipients = repo._build_email_recipients( - session=repository_session, - form_id="f", - delivery_id="d", - recipients_config=EmailRecipients(include_bound_group=True, items=[ExternalRecipient(email="e@example.com")]), - ) - assert recipients == ["ok"] + form = _human_input_repository_persist_form(repository_session, form_id="f1") + _human_input_repository_persist_form(repository_session, form_id="other-tenant-form", tenant_id="other-tenant") + _persist_recipient(repository_session, form_id=form.id, recipient_id="r1", access_token="tok") + repo = HumanInputFormRepositoryImpl(tenant_id="tenant", workflow_execution_id="run") + entity = repo.get_form("node") + assert entity is not None + assert entity.id == "f1" + assert entity.recipients[0].id == "r1" + assert entity.recipients[0].token == "tok" + @staticmethod + def test_create_form_adds_console_and_backstage_recipients( + repository_session: Session, monkeypatch: pytest.MonkeyPatch + ) -> None: + fixed_now = datetime(2024, 1, 1, 0, 0, 0) + monkeypatch.setattr("core.repositories.human_input_repository.naive_utc_now", lambda: fixed_now) -def test_build_email_recipients_uses_selected_members_when_not_whole_workspace( - repository_session: Session, monkeypatch: pytest.MonkeyPatch -) -> None: - repo = HumanInputFormRepositoryImpl(tenant_id="tenant") + ids = iter(["form-id", "del-web", "del-console", "del-backstage"]) + monkeypatch.setattr("core.repositories.human_input_repository.uuidv7", lambda: next(ids)) - def fake_query(*, session: Any, restrict_to_user_ids: Sequence[str]) -> list[_WorkspaceMemberInfo]: - assert restrict_to_user_ids == ["u1"] - return [_WorkspaceMemberInfo(user_id="u1", email="a@example.com")] + repo = HumanInputFormRepositoryImpl( + tenant_id="tenant", + app_id="app", + workflow_execution_id="run", + invoke_source="debugger", + submission_actor_id="acc-1", + ) - monkeypatch.setattr(repo, "_query_workspace_members_by_ids", fake_query) - monkeypatch.setattr(repo, "_create_email_recipients_from_resolved", lambda **_: ["ok"]) - recipients = repo._build_email_recipients( - session=repository_session, - form_id="f", - delivery_id="d", - recipients_config=EmailRecipients( - include_bound_group=False, - items=[MemberRecipient(reference_id="u1"), ExternalRecipient(email="e@example.com")], - ), - ) - assert recipients == ["ok"] - - -def test_get_form_returns_entity_and_none_when_missing(repository_session: Session) -> None: - repo = HumanInputFormRepositoryImpl(tenant_id="tenant", workflow_execution_id="run") - assert repo.get_form("node") is None - - form = _persist_form(repository_session, form_id="f1") - _persist_form(repository_session, form_id="other-tenant-form", tenant_id="other-tenant") - _persist_recipient(repository_session, form_id=form.id, recipient_id="r1", access_token="tok") - repo = HumanInputFormRepositoryImpl(tenant_id="tenant", workflow_execution_id="run") - entity = repo.get_form("node") - assert entity is not None - assert entity.id == "f1" - assert entity.recipients[0].id == "r1" - assert entity.recipients[0].token == "tok" - - -def test_create_form_adds_console_and_backstage_recipients( - repository_session: Session, monkeypatch: pytest.MonkeyPatch -) -> None: - fixed_now = datetime(2024, 1, 1, 0, 0, 0) - monkeypatch.setattr("core.repositories.human_input_repository.naive_utc_now", lambda: fixed_now) - - ids = iter(["form-id", "del-web", "del-console", "del-backstage"]) - monkeypatch.setattr("core.repositories.human_input_repository.uuidv7", lambda: next(ids)) - - repo = HumanInputFormRepositoryImpl( - tenant_id="tenant", - app_id="app", - workflow_execution_id="run", - invoke_source="debugger", - submission_actor_id="acc-1", - ) - - form_config = HumanInputNodeData( - title="Title", - delivery_methods=[], - form_content="hello", - inputs=[], - user_actions=[UserActionConfig(id="submit", title="Submit")], - ) - params = FormCreateParams( - workflow_execution_id=None, - node_id="node", - form_config=form_config, - rendered_content="

hello

", - delivery_methods=[WebAppDeliveryMethod()], - display_in_ui=True, - resolved_default_values={}, - form_kind=HumanInputFormKind.RUNTIME, - ) - - entity = repo.create_form(params) - assert entity.id == "form-id" - assert entity.expiration_time == fixed_now + timedelta(hours=form_config.timeout) - # Console token should take precedence when console recipient is present. - assert entity.submission_token is not None - assert len(entity.recipients) == 3 - repository_session.expire_all() - persisted_form = repository_session.get(HumanInputForm, "form-id") - assert persisted_form is not None - assert persisted_form.status == HumanInputFormStatus.WAITING - recipients = repository_session.scalars( - select(HumanInputFormRecipient).where(HumanInputFormRecipient.form_id == "form-id") - ).all() - assert {recipient.recipient_type for recipient in recipients} == { - RecipientType.STANDALONE_WEB_APP, - RecipientType.CONSOLE, - RecipientType.BACKSTAGE, - } - - -def test_create_runtime_form_requires_workflow_even_when_conversation_is_present() -> None: - repo = HumanInputFormRepositoryImpl(tenant_id="tenant", app_id="app") - params = FormCreateParams( - workflow_execution_id=None, - conversation_id="conversation", - node_id="node", - form_config=HumanInputNodeData( + form_config = HumanInputNodeData( title="Title", delivery_methods=[], form_content="hello", inputs=[], user_actions=[UserActionConfig(id="submit", title="Submit")], - ), - rendered_content="hello", - delivery_methods=[], - display_in_ui=True, - resolved_default_values={}, - ) - with pytest.raises(ValueError, match="requires a workflow_execution_id"): - repo.create_form(params) - - -def test_submission_get_by_token_returns_none_when_missing_or_form_missing(repository_session: Session) -> None: - repo = HumanInputFormSubmissionRepository() - assert repo.get_by_token("tok") is None - - _persist_recipient(repository_session, form_id="missing-form", access_token="orphan-token") - assert repo.get_by_token("orphan-token") is None - - -def test_submission_repository_init_no_args() -> None: - repo = HumanInputFormSubmissionRepository() - assert isinstance(repo, HumanInputFormSubmissionRepository) - - -def test_submission_get_by_token_and_get_by_form_id_success_paths(repository_session: Session) -> None: - form = _persist_form(repository_session, form_id="f1", workflow_run_id=None) - _persist_recipient(repository_session, form_id=form.id, recipient_id="r1", access_token="tok") - repo = HumanInputFormSubmissionRepository() - record = repo.get_by_token("tok") - assert record is not None - assert record.access_token == "tok" - - record = repo.get_by_form_id_and_recipient_type(form_id=form.id, recipient_type=RecipientType.STANDALONE_WEB_APP) - assert record is not None - assert record.recipient_id == "r1" - - record = repo.get_by_form_id(form.id) - assert record is not None - assert record.form_id == form.id - assert record.recipient_id is None - - -def test_submission_get_by_form_id_returns_none_on_missing(repository_session: Session) -> None: - repo = HumanInputFormSubmissionRepository() - assert repo.get_by_form_id_and_recipient_type(form_id="f", recipient_type=RecipientType.CONSOLE) is None - assert repo.get_by_form_id("f") is None - - -def test_submission_batch_materializes_requested_forms_in_one_read(repository_session: Session) -> None: - _persist_form(repository_session, form_id="first", status=HumanInputFormStatus.SUBMITTED) - _persist_form(repository_session, form_id="second", status=HumanInputFormStatus.TIMEOUT) - _persist_form(repository_session, form_id="unrequested") - statements: list[str] = [] - - def record_statement(_connection, _cursor, statement: str, _parameters, _context, _executemany) -> None: - statements.append(statement) - - engine = repository_session.get_bind() - event.listen(engine, "before_cursor_execute", record_statement) - try: - records = HumanInputFormSubmissionRepository().get_by_form_ids( - ["second", "missing", "first", "first"], tenant_id="tenant", app_id="app", workflow_run_id="run" ) - assert len(statements) == 1 - assert set(records) == {"first", "second"} - assert records["first"].status == HumanInputFormStatus.SUBMITTED - assert records["second"].definition.user_actions[0].id == "submit" - assert records["second"].status == HumanInputFormStatus.TIMEOUT - # Materialized records remain usable after the repository's session closes. - assert len(statements) == 1 - assert ( - HumanInputFormSubmissionRepository().get_by_form_ids( - [], tenant_id="tenant", app_id="app", workflow_run_id="run" + params = FormCreateParams( + workflow_execution_id=None, + node_id="node", + form_config=form_config, + rendered_content="

hello

", + delivery_methods=[WebAppDeliveryMethod()], + display_in_ui=True, + resolved_default_values={}, + form_kind=HumanInputFormKind.RUNTIME, + ) + + entity = repo.create_form(params) + assert entity.id == "form-id" + assert entity.expiration_time == fixed_now + timedelta(hours=form_config.timeout) + # Console token should take precedence when console recipient is present. + assert entity.submission_token is not None + assert len(entity.recipients) == 3 + repository_session.expire_all() + persisted_form = repository_session.get(HumanInputForm, "form-id") + assert persisted_form is not None + assert persisted_form.status == HumanInputFormStatus.WAITING + recipients = repository_session.scalars( + select(HumanInputFormRecipient).where(HumanInputFormRecipient.form_id == "form-id") + ).all() + assert {recipient.recipient_type for recipient in recipients} == { + RecipientType.STANDALONE_WEB_APP, + RecipientType.CONSOLE, + RecipientType.BACKSTAGE, + } + + @staticmethod + def test_create_runtime_form_requires_workflow_even_when_conversation_is_present() -> None: + repo = HumanInputFormRepositoryImpl(tenant_id="tenant", app_id="app") + params = FormCreateParams( + workflow_execution_id=None, + conversation_id="conversation", + node_id="node", + form_config=HumanInputNodeData( + title="Title", + delivery_methods=[], + form_content="hello", + inputs=[], + user_actions=[UserActionConfig(id="submit", title="Submit")], + ), + rendered_content="hello", + delivery_methods=[], + display_in_ui=True, + resolved_default_values={}, + ) + with pytest.raises(ValueError, match="requires a workflow_execution_id"): + repo.create_form(params) + + @staticmethod + def test_submission_get_by_token_returns_none_when_missing_or_form_missing(repository_session: Session) -> None: + repo = HumanInputFormSubmissionRepository() + assert repo.get_by_token("tok") is None + + _persist_recipient(repository_session, form_id="missing-form", access_token="orphan-token") + assert repo.get_by_token("orphan-token") is None + + @staticmethod + def test_submission_repository_init_no_args() -> None: + repo = HumanInputFormSubmissionRepository() + assert isinstance(repo, HumanInputFormSubmissionRepository) + + @staticmethod + def test_submission_get_by_token_and_get_by_form_id_success_paths(repository_session: Session) -> None: + form = _human_input_repository_persist_form(repository_session, form_id="f1", workflow_run_id=None) + _persist_recipient(repository_session, form_id=form.id, recipient_id="r1", access_token="tok") + repo = HumanInputFormSubmissionRepository() + record = repo.get_by_token("tok") + assert record is not None + assert record.access_token == "tok" + + record = repo.get_by_form_id_and_recipient_type( + form_id=form.id, recipient_type=RecipientType.STANDALONE_WEB_APP + ) + assert record is not None + assert record.recipient_id == "r1" + + record = repo.get_by_form_id(form.id) + assert record is not None + assert record.form_id == form.id + assert record.recipient_id is None + + @staticmethod + def test_submission_get_by_form_id_returns_none_on_missing(repository_session: Session) -> None: + repo = HumanInputFormSubmissionRepository() + assert repo.get_by_form_id_and_recipient_type(form_id="f", recipient_type=RecipientType.CONSOLE) is None + assert repo.get_by_form_id("f") is None + + @staticmethod + def test_submission_batch_materializes_requested_forms_in_one_read(repository_session: Session) -> None: + _human_input_repository_persist_form(repository_session, form_id="first", status=HumanInputFormStatus.SUBMITTED) + _human_input_repository_persist_form(repository_session, form_id="second", status=HumanInputFormStatus.TIMEOUT) + _human_input_repository_persist_form(repository_session, form_id="unrequested") + statements: list[str] = [] + + def record_statement(_connection, _cursor, statement: str, _parameters, _context, _executemany) -> None: + statements.append(statement) + + engine = repository_session.get_bind() + event.listen(engine, "before_cursor_execute", record_statement) + try: + records = HumanInputFormSubmissionRepository().get_by_form_ids( + ["second", "missing", "first", "first"], tenant_id="tenant", app_id="app", workflow_run_id="run" ) - == {} - ) - assert len(statements) == 1 - finally: - event.remove(engine, "before_cursor_execute", record_statement) + assert len(statements) == 1 + assert set(records) == {"first", "second"} + assert records["first"].status == HumanInputFormStatus.SUBMITTED + assert records["second"].definition.user_actions[0].id == "submit" + assert records["second"].status == HumanInputFormStatus.TIMEOUT + # Materialized records remain usable after the repository's session closes. + assert len(statements) == 1 + assert ( + HumanInputFormSubmissionRepository().get_by_form_ids( + [], tenant_id="tenant", app_id="app", workflow_run_id="run" + ) + == {} + ) + assert len(statements) == 1 + finally: + event.remove(engine, "before_cursor_execute", record_statement) + @staticmethod + def test_mark_submitted_updates_and_raises_when_missing( + repository_session: Session, monkeypatch: pytest.MonkeyPatch + ) -> None: + fixed_now = datetime(2024, 1, 1, 0, 0, 0) + monkeypatch.setattr("core.repositories.human_input_repository.naive_utc_now", lambda: fixed_now) -def test_mark_submitted_updates_and_raises_when_missing( - repository_session: Session, monkeypatch: pytest.MonkeyPatch -) -> None: - fixed_now = datetime(2024, 1, 1, 0, 0, 0) - monkeypatch.setattr("core.repositories.human_input_repository.naive_utc_now", lambda: fixed_now) + repo = HumanInputFormSubmissionRepository() + with pytest.raises(FormNotFoundError, match="form not found"): + repo.mark_submitted( + form_id="missing", + recipient_id=None, + selected_action_id="a", + form_data={}, + submission_user_id=None, + submission_end_user_id=None, + ) - repo = HumanInputFormSubmissionRepository() - with pytest.raises(FormNotFoundError, match="form not found"): - repo.mark_submitted( - form_id="missing", - recipient_id=None, - selected_action_id="a", - form_data={}, - submission_user_id=None, - submission_end_user_id=None, - ) - - form = _persist_form(repository_session, form_id="f", workflow_run_id=None) - recipient = _persist_recipient( - repository_session, - form_id=form.id, - recipient_id="r", - recipient_type=RecipientType.CONSOLE, - access_token="tok", - ) - record = repo.mark_submitted( - form_id=form.id, - recipient_id=recipient.id, - selected_action_id="approve", - form_data={"k": "v"}, - submission_user_id="u", - submission_end_user_id="eu", - ) - repository_session.expire_all() - persisted = repository_session.get(HumanInputForm, form.id) - assert persisted is not None - assert persisted.status == HumanInputFormStatus.SUBMITTED - assert persisted.submitted_at == fixed_now - assert persisted.completed_by_recipient_id == recipient.id - assert record.submitted_data == {"k": "v"} - - -@pytest.mark.parametrize("status", [HumanInputFormStatus.EXPIRED, HumanInputFormStatus.SUBMITTED]) -def test_mark_submitted_does_not_overwrite_a_concurrent_stop_or_submission( - repository_session: Session, status: HumanInputFormStatus -) -> None: - form = _persist_form(repository_session, form_id="stopped-form", workflow_run_id=None) - form.status = status - repository_session.commit() - with pytest.raises(FormAlreadyHandledError) as error: - HumanInputFormSubmissionRepository().mark_submitted( + form = _human_input_repository_persist_form(repository_session, form_id="f", workflow_run_id=None) + recipient = _persist_recipient( + repository_session, form_id=form.id, - recipient_id=None, - selected_action_id="approve", - form_data={"note": "late submission"}, - submission_user_id=None, - submission_end_user_id=None, + recipient_id="r", + recipient_type=RecipientType.CONSOLE, + access_token="tok", ) - assert error.value.status == status - repository_session.expire_all() - assert form.status == status - assert form.submitted_data is None + record = repo.mark_submitted( + form_id=form.id, + recipient_id=recipient.id, + selected_action_id="approve", + form_data={"k": "v"}, + submission_user_id="u", + submission_end_user_id="eu", + ) + repository_session.expire_all() + persisted = repository_session.get(HumanInputForm, form.id) + assert persisted is not None + assert persisted.status == HumanInputFormStatus.SUBMITTED + assert persisted.submitted_at == fixed_now + assert persisted.completed_by_recipient_id == recipient.id + assert record.submitted_data == {"k": "v"} + @staticmethod + @pytest.mark.parametrize("status", [HumanInputFormStatus.EXPIRED, HumanInputFormStatus.SUBMITTED]) + def test_mark_submitted_does_not_overwrite_a_concurrent_stop_or_submission( + repository_session: Session, status: HumanInputFormStatus + ) -> None: + form = _human_input_repository_persist_form(repository_session, form_id="stopped-form", workflow_run_id=None) + form.status = status + repository_session.commit() + with pytest.raises(FormAlreadyHandledError) as error: + HumanInputFormSubmissionRepository().mark_submitted( + form_id=form.id, + recipient_id=None, + selected_action_id="approve", + form_data={"note": "late submission"}, + submission_user_id=None, + submission_end_user_id=None, + ) + assert error.value.status == status + repository_session.expire_all() + assert form.status == status + assert form.submitted_data is None -def test_mark_submitted_serializes_select_and_file_payloads( - repository_session: Session, monkeypatch: pytest.MonkeyPatch -) -> None: - fixed_now = datetime(2024, 1, 1, 0, 0, 0) - monkeypatch.setattr("core.repositories.human_input_repository.naive_utc_now", lambda: fixed_now) + @staticmethod + def test_mark_submitted_serializes_select_and_file_payloads( + repository_session: Session, monkeypatch: pytest.MonkeyPatch + ) -> None: + fixed_now = datetime(2024, 1, 1, 0, 0, 0) + monkeypatch.setattr("core.repositories.human_input_repository.naive_utc_now", lambda: fixed_now) - form = _persist_form(repository_session, form_id="f-complex", workflow_run_id=None) - recipient = _persist_recipient( - repository_session, - form_id=form.id, - recipient_id="r-complex", - recipient_type=RecipientType.CONSOLE, - access_token="tok", - ) + form = _human_input_repository_persist_form(repository_session, form_id="f-complex", workflow_run_id=None) + recipient = _persist_recipient( + repository_session, + form_id=form.id, + recipient_id="r-complex", + recipient_type=RecipientType.CONSOLE, + access_token="tok", + ) - payload = { - "decision": "approve", - "attachment": { - "type": "document", - "transfer_method": "remote_url", - "remote_url": "https://example.com/file.txt", - "filename": "file.txt", - "extension": ".txt", - "mime_type": "text/plain", - }, - "attachments": [ - { + payload = { + "decision": "approve", + "attachment": { "type": "document", "transfer_method": "remote_url", - "remote_url": "https://example.com/first.txt", - "filename": "first.txt", + "remote_url": "https://example.com/file.txt", + "filename": "file.txt", "extension": ".txt", "mime_type": "text/plain", }, - { - "type": "document", - "transfer_method": "remote_url", - "remote_url": "https://example.com/second.txt", - "filename": "second.txt", - "extension": ".txt", - "mime_type": "text/plain", - }, - ], - } + "attachments": [ + { + "type": "document", + "transfer_method": "remote_url", + "remote_url": "https://example.com/first.txt", + "filename": "first.txt", + "extension": ".txt", + "mime_type": "text/plain", + }, + { + "type": "document", + "transfer_method": "remote_url", + "remote_url": "https://example.com/second.txt", + "filename": "second.txt", + "extension": ".txt", + "mime_type": "text/plain", + }, + ], + } - repo = HumanInputFormSubmissionRepository() - record = repo.mark_submitted( - form_id=form.id, - recipient_id=recipient.id, - selected_action_id="approve", - form_data=payload, - submission_user_id="user-1", - submission_end_user_id="end-user-1", - ) + repo = HumanInputFormSubmissionRepository() + record = repo.mark_submitted( + form_id=form.id, + recipient_id=recipient.id, + selected_action_id="approve", + form_data=payload, + submission_user_id="user-1", + submission_end_user_id="end-user-1", + ) - repository_session.expire_all() - persisted = repository_session.get(HumanInputForm, form.id) - assert persisted is not None - assert json.loads(persisted.submitted_data or "") == payload - assert record.submitted_data == payload + repository_session.expire_all() + persisted = repository_session.get(HumanInputForm, form.id) + assert persisted is not None + assert json.loads(persisted.submitted_data or "") == payload + assert record.submitted_data == payload + @staticmethod + def test_mark_timeout_invalid_status_rolls_back(repository_session: Session) -> None: + form = _human_input_repository_persist_form(repository_session, form_id="f", workflow_run_id=None) + repo = HumanInputFormSubmissionRepository() + with pytest.raises(_InvalidTimeoutStatusError, match="invalid timeout status"): + repo.mark_timeout(form_id=form.id, timeout_status=HumanInputFormStatus.SUBMITTED) # type: ignore[arg-type] + repository_session.expire_all() + persisted = repository_session.get(HumanInputForm, form.id) + assert persisted is not None + assert persisted.status == HumanInputFormStatus.WAITING -def test_mark_timeout_invalid_status_rolls_back(repository_session: Session) -> None: - form = _persist_form(repository_session, form_id="f", workflow_run_id=None) - repo = HumanInputFormSubmissionRepository() - with pytest.raises(_InvalidTimeoutStatusError, match="invalid timeout status"): - repo.mark_timeout(form_id=form.id, timeout_status=HumanInputFormStatus.SUBMITTED) # type: ignore[arg-type] - repository_session.expire_all() - persisted = repository_session.get(HumanInputForm, form.id) - assert persisted is not None - assert persisted.status == HumanInputFormStatus.WAITING + @staticmethod + def test_mark_timeout_already_timed_out_returns_record(repository_session: Session) -> None: + form = _human_input_repository_persist_form( + repository_session, + form_id="f", + workflow_run_id=None, + status=HumanInputFormStatus.TIMEOUT, + ) + repo = HumanInputFormSubmissionRepository() + record = repo.mark_timeout(form_id=form.id, timeout_status=HumanInputFormStatus.TIMEOUT, reason="r") + assert record.status == HumanInputFormStatus.TIMEOUT + @staticmethod + def test_mark_timeout_submitted_raises_form_not_found(repository_session: Session) -> None: + form = _human_input_repository_persist_form( + repository_session, + form_id="f", + workflow_run_id=None, + status=HumanInputFormStatus.SUBMITTED, + ) + repo = HumanInputFormSubmissionRepository() + with pytest.raises(FormNotFoundError, match="form already submitted"): + repo.mark_timeout(form_id=form.id, timeout_status=HumanInputFormStatus.EXPIRED) -def test_mark_timeout_already_timed_out_returns_record(repository_session: Session) -> None: - form = _persist_form( - repository_session, - form_id="f", - workflow_run_id=None, - status=HumanInputFormStatus.TIMEOUT, - ) - repo = HumanInputFormSubmissionRepository() - record = repo.mark_timeout(form_id=form.id, timeout_status=HumanInputFormStatus.TIMEOUT, reason="r") - assert record.status == HumanInputFormStatus.TIMEOUT + @staticmethod + def test_mark_timeout_updates_fields(repository_session: Session) -> None: + form = _human_input_repository_persist_form(repository_session, form_id="f", workflow_run_id=None) + form.selected_action_id = "a" + form.submitted_data = "{}" + form.submission_user_id = "u" + form.submission_end_user_id = "eu" + form.completed_by_recipient_id = "r" + repository_session.commit() + repo = HumanInputFormSubmissionRepository() + record = repo.mark_timeout(form_id=form.id, timeout_status=HumanInputFormStatus.EXPIRED) + repository_session.expire_all() + persisted = repository_session.get(HumanInputForm, form.id) + assert persisted is not None + assert persisted.status == HumanInputFormStatus.EXPIRED + assert persisted.selected_action_id is None + assert persisted.submitted_data is None + assert persisted.submission_user_id is None + assert persisted.submission_end_user_id is None + assert persisted.completed_by_recipient_id is None + assert record.status == HumanInputFormStatus.EXPIRED - -def test_mark_timeout_submitted_raises_form_not_found(repository_session: Session) -> None: - form = _persist_form( - repository_session, - form_id="f", - workflow_run_id=None, - status=HumanInputFormStatus.SUBMITTED, - ) - repo = HumanInputFormSubmissionRepository() - with pytest.raises(FormNotFoundError, match="form already submitted"): - repo.mark_timeout(form_id=form.id, timeout_status=HumanInputFormStatus.EXPIRED) - - -def test_mark_timeout_updates_fields(repository_session: Session) -> None: - form = _persist_form(repository_session, form_id="f", workflow_run_id=None) - form.selected_action_id = "a" - form.submitted_data = "{}" - form.submission_user_id = "u" - form.submission_end_user_id = "eu" - form.completed_by_recipient_id = "r" - repository_session.commit() - repo = HumanInputFormSubmissionRepository() - record = repo.mark_timeout(form_id=form.id, timeout_status=HumanInputFormStatus.EXPIRED) - repository_session.expire_all() - persisted = repository_session.get(HumanInputForm, form.id) - assert persisted is not None - assert persisted.status == HumanInputFormStatus.EXPIRED - assert persisted.selected_action_id is None - assert persisted.submitted_data is None - assert persisted.submission_user_id is None - assert persisted.submission_end_user_id is None - assert persisted.completed_by_recipient_id is None - assert record.status == HumanInputFormStatus.EXPIRED - - -def test_mark_timeout_raises_when_form_missing(repository_session: Session) -> None: - repo = HumanInputFormSubmissionRepository() - with pytest.raises(FormNotFoundError, match="form not found"): - repo.mark_timeout(form_id="missing", timeout_status=HumanInputFormStatus.TIMEOUT) + @staticmethod + def test_mark_timeout_raises_when_form_missing(repository_session: Session) -> None: + repo = HumanInputFormSubmissionRepository() + with pytest.raises(FormNotFoundError, match="form not found"): + repo.mark_timeout(form_id="missing", timeout_status=HumanInputFormStatus.TIMEOUT) diff --git a/api/tests/unit_tests/core/repositories/test_sqlalchemy_workflow_execution_repository.py b/api/tests/unit_tests/core/repositories/test_sqlalchemy_workflow_execution_repository.py index aff28b660aa..ffac4ada98e 100644 --- a/api/tests/unit_tests/core/repositories/test_sqlalchemy_workflow_execution_repository.py +++ b/api/tests/unit_tests/core/repositories/test_sqlalchemy_workflow_execution_repository.py @@ -3,57 +3,45 @@ from datetime import UTC, datetime from uuid import uuid4 import pytest +from pytest_mock import MockerFixture from sqlalchemy.engine import Engine from sqlalchemy.orm import Session, sessionmaker +from core.repositories.celery_workflow_execution_repository import CeleryWorkflowExecutionRepository from core.repositories.sqlalchemy_workflow_execution_repository import SQLAlchemyWorkflowExecutionRepository -from graphon.entities import WorkflowExecution -from graphon.enums import WorkflowExecutionStatus, WorkflowType -from models import Account, CreatorUserRole, EndUser, Tenant, WorkflowRun +from core.repositories.sqlalchemy_workflow_node_execution_repository import SQLAlchemyWorkflowNodeExecutionRepository +from extensions.logstore.aliyun_logstore import AliyunLogStore +from extensions.logstore.repositories.logstore_api_workflow_node_execution_repository import ( + LogstoreAPIWorkflowNodeExecutionRepository, +) +from extensions.logstore.repositories.logstore_api_workflow_run_repository import LogstoreAPIWorkflowRunRepository +from extensions.logstore.repositories.logstore_workflow_execution_repository import LogstoreWorkflowExecutionRepository +from extensions.logstore.repositories.logstore_workflow_node_execution_repository import ( + LogstoreWorkflowNodeExecutionRepository, +) +from graphon.entities import WorkflowExecution as _sqlalchemy_workflow_execution_repository_WorkflowExecution +from graphon.entities import WorkflowNodeExecution +from graphon.enums import WorkflowExecutionStatus, WorkflowNodeExecutionStatus, WorkflowType +from libs.datetime_utils import naive_utc_now +from models import Account, EndUser, Tenant, WorkflowRun +from models import CreatorUserRole as _sqlalchemy_workflow_execution_repository_CreatorUserRole +from models.enums import CreatorUserRole as _workflow_control_persistence_CreatorUserRole from models.enums import EndUserType, WorkflowRunTriggeredFrom +from models.workflow import WorkflowNodeExecutionTriggeredFrom from models.workflow import WorkflowType as ModelWorkflowType from repositories.entities.workflow_execution import WorkflowExecution as TimedWorkflowExecution +from repositories.entities.workflow_execution import ( + WorkflowExecution as _workflow_control_persistence_WorkflowExecution, +) +from repositories.sqlalchemy_api_workflow_run_repository import DifyAPISQLAlchemyWorkflowRunRepository +from repositories.workflow_run_control import apply_workflow_stop_to_node +from tasks.workflow_execution_tasks import _update_workflow_run_from_execution TABLES = (WorkflowRun,) RESOURCE_TENANT_ID = "resource-tenant-id" -def test_save_preserves_accumulated_duration_without_double_counting( - sqlite_session_factory: sessionmaker[Session], account: Account, sample_workflow_execution: WorkflowExecution -) -> None: - repository = SQLAlchemyWorkflowExecutionRepository( - session_factory=sqlite_session_factory, - tenant_id=RESOURCE_TENANT_ID, - user=account, - app_id="app-1", - triggered_from=WorkflowRunTriggeredFrom.APP_RUN, - ) - execution = TimedWorkflowExecution( - **(sample_workflow_execution.model_dump() | {"finished_at": None, "status": WorkflowExecutionStatus.RUNNING}), - accumulated_elapsed_time=5.0, - ) - original_started_at = execution.started_at - - for status, duration in [ - (WorkflowExecutionStatus.RUNNING, 5.0), - (WorkflowExecutionStatus.PAUSED, 7.0), - (WorkflowExecutionStatus.SUCCEEDED, 9.0), - ]: - execution.status = status - execution.accumulated_elapsed_time = duration - if status == WorkflowExecutionStatus.SUCCEEDED: - execution.finished_at = datetime(2027, 1, 1, tzinfo=UTC) - repository.save(execution) - repository.save(execution) - with sqlite_session_factory() as session: - stored = session.get(WorkflowRun, execution.id_) - assert stored is not None - assert stored.elapsed_time == duration - assert stored.created_at == original_started_at.replace(tzinfo=None) - assert repository._to_domain_model(stored).elapsed_time == duration - - def _make_account(*, tenant_id: str | None = None) -> Account: account = Account(name="Repository User", email=f"{uuid4()}@example.com") account.id = str(uuid4()) @@ -64,398 +52,755 @@ def _make_account(*, tenant_id: str | None = None) -> Account: return account -@pytest.fixture -def sqlite_session_factory(sqlite_engine: Engine) -> sessionmaker[Session]: - """Create repository-owned sessions bound to the isolated SQLite engine.""" - return sessionmaker(bind=sqlite_engine, expire_on_commit=False) - - -@pytest.fixture -def account() -> Account: - return _make_account(tenant_id=str(uuid4())) - - -@pytest.fixture -def end_user() -> EndUser: - return EndUser( - id=str(uuid4()), - tenant_id=str(uuid4()), - app_id=None, - type=EndUserType.SERVICE_API, - external_user_id=None, - name="Repository End User", - session_id=str(uuid4()), - ) - - -@pytest.fixture -def sample_workflow_execution() -> WorkflowExecution: - """Sample WorkflowExecution for testing.""" - return WorkflowExecution( - id_=str(uuid4()), - workflow_id=str(uuid4()), - workflow_type=WorkflowType.WORKFLOW, - workflow_version="1.0", - graph={"nodes": [], "edges": []}, - inputs={"input1": "value1"}, - outputs={"output1": "result1"}, - status=WorkflowExecutionStatus.SUCCEEDED, - error_message="", - total_tokens=100, - total_steps=5, - exceptions_count=0, - started_at=datetime.now(UTC), - finished_at=datetime.now(UTC), - ) - - -class TestSQLAlchemyWorkflowExecutionRepository: - def test_init_with_sessionmaker(self, sqlite_session_factory: sessionmaker[Session], account: Account): - app_id = "test_app_id" - triggered_from = WorkflowRunTriggeredFrom.APP_RUN - - repo = SQLAlchemyWorkflowExecutionRepository( +class TestSqlalchemyWorkflowExecutionRepositoryCases: + @staticmethod + def test_save_preserves_accumulated_duration_without_double_counting( + sqlite_session_factory: sessionmaker[Session], + account: Account, + sample_workflow_execution: _sqlalchemy_workflow_execution_repository_WorkflowExecution, + ) -> None: + repository = SQLAlchemyWorkflowExecutionRepository( session_factory=sqlite_session_factory, tenant_id=RESOURCE_TENANT_ID, user=account, - app_id=app_id, - triggered_from=triggered_from, - ) - - assert repo._session_factory is sqlite_session_factory - assert repo._tenant_id == RESOURCE_TENANT_ID - assert repo._app_id == app_id - assert repo._triggered_from == triggered_from - assert repo._creator_user_id == account.id - assert repo._creator_user_role == CreatorUserRole.ACCOUNT - - def test_init_with_engine(self, sqlite_engine: Engine, account: Account): - repo = SQLAlchemyWorkflowExecutionRepository( - session_factory=sqlite_engine, - tenant_id=RESOURCE_TENANT_ID, - user=account, - app_id="test_app_id", + app_id="app-1", triggered_from=WorkflowRunTriggeredFrom.APP_RUN, ) + execution = TimedWorkflowExecution( + **( + sample_workflow_execution.model_dump() + | {"finished_at": None, "status": WorkflowExecutionStatus.RUNNING} + ), + accumulated_elapsed_time=5.0, + ) + original_started_at = execution.started_at - assert isinstance(repo._session_factory, sessionmaker) - assert repo._session_factory.kw["bind"] is sqlite_engine + for status, duration in [ + (WorkflowExecutionStatus.RUNNING, 5.0), + (WorkflowExecutionStatus.PAUSED, 7.0), + (WorkflowExecutionStatus.SUCCEEDED, 9.0), + ]: + execution.status = status + execution.accumulated_elapsed_time = duration + if status == WorkflowExecutionStatus.SUCCEEDED: + execution.finished_at = datetime(2027, 1, 1, tzinfo=UTC) + repository.save(execution) + repository.save(execution) + with sqlite_session_factory() as session: + stored = session.get(WorkflowRun, execution.id_) + assert stored is not None + assert stored.elapsed_time == duration + assert stored.created_at == original_started_at.replace(tzinfo=None) + assert repository._to_domain_model(stored).elapsed_time == duration - def test_init_invalid_session_factory(self, account: Account): - with pytest.raises(ValueError, match="Invalid session_factory type"): - SQLAlchemyWorkflowExecutionRepository( - session_factory="invalid", + @staticmethod + @pytest.fixture + def sqlite_session_factory(sqlite_engine: Engine) -> sessionmaker[Session]: + """Create repository-owned sessions bound to the isolated SQLite engine.""" + return sessionmaker(bind=sqlite_engine, expire_on_commit=False) + + @staticmethod + @pytest.fixture + def account() -> Account: + return _make_account(tenant_id=str(uuid4())) + + @staticmethod + @pytest.fixture + def end_user() -> EndUser: + return EndUser( + id=str(uuid4()), + tenant_id=str(uuid4()), + app_id=None, + type=EndUserType.SERVICE_API, + external_user_id=None, + name="Repository End User", + session_id=str(uuid4()), + ) + + @staticmethod + @pytest.fixture + def sample_workflow_execution() -> _sqlalchemy_workflow_execution_repository_WorkflowExecution: + """Sample WorkflowExecution for testing.""" + return _sqlalchemy_workflow_execution_repository_WorkflowExecution( + id_=str(uuid4()), + workflow_id=str(uuid4()), + workflow_type=WorkflowType.WORKFLOW, + workflow_version="1.0", + graph={"nodes": [], "edges": []}, + inputs={"input1": "value1"}, + outputs={"output1": "result1"}, + status=WorkflowExecutionStatus.SUCCEEDED, + error_message="", + total_tokens=100, + total_steps=5, + exceptions_count=0, + started_at=datetime.now(UTC), + finished_at=datetime.now(UTC), + ) + + class TestSQLAlchemyWorkflowExecutionRepository: + def test_init_with_sessionmaker(self, sqlite_session_factory: sessionmaker[Session], account: Account): + app_id = "test_app_id" + triggered_from = WorkflowRunTriggeredFrom.APP_RUN + + repo = SQLAlchemyWorkflowExecutionRepository( + session_factory=sqlite_session_factory, + tenant_id=RESOURCE_TENANT_ID, + user=account, + app_id=app_id, + triggered_from=triggered_from, + ) + + assert repo._session_factory is sqlite_session_factory + assert repo._tenant_id == RESOURCE_TENANT_ID + assert repo._app_id == app_id + assert repo._triggered_from == triggered_from + assert repo._creator_user_id == account.id + assert repo._creator_user_role == _sqlalchemy_workflow_execution_repository_CreatorUserRole.ACCOUNT + + def test_init_with_engine(self, sqlite_engine: Engine, account: Account): + repo = SQLAlchemyWorkflowExecutionRepository( + session_factory=sqlite_engine, + tenant_id=RESOURCE_TENANT_ID, + user=account, + app_id="test_app_id", + triggered_from=WorkflowRunTriggeredFrom.APP_RUN, + ) + + assert isinstance(repo._session_factory, sessionmaker) + assert repo._session_factory.kw["bind"] is sqlite_engine + + def test_init_invalid_session_factory(self, account: Account): + with pytest.raises(ValueError, match="Invalid session_factory type"): + SQLAlchemyWorkflowExecutionRepository( + session_factory="invalid", + tenant_id=RESOURCE_TENANT_ID, + user=account, + app_id=None, + triggered_from=None, + ) + + def test_init_no_tenant_id(self, sqlite_session_factory: sessionmaker[Session]): + user = _make_account() + + with pytest.raises(ValueError, match="tenant_id is required"): + SQLAlchemyWorkflowExecutionRepository( + session_factory=sqlite_session_factory, + tenant_id="", + user=user, + app_id=None, + triggered_from=None, + ) + + def test_init_uses_resource_tenant_when_account_has_no_current_tenant( + self, sqlite_session_factory: sessionmaker[Session] + ): + user = _make_account() + + repo = SQLAlchemyWorkflowExecutionRepository( + session_factory=sqlite_session_factory, + tenant_id=RESOURCE_TENANT_ID, + user=user, + app_id="test-app", + triggered_from=WorkflowRunTriggeredFrom.APP_RUN, + ) + + assert repo._tenant_id == RESOURCE_TENANT_ID + assert repo._creator_user_id == user.id + + def test_init_with_end_user(self, sqlite_session_factory: sessionmaker[Session], end_user: EndUser): + repo = SQLAlchemyWorkflowExecutionRepository( + session_factory=sqlite_session_factory, + tenant_id=RESOURCE_TENANT_ID, + user=end_user, + app_id=None, + triggered_from=None, + ) + assert repo._tenant_id == RESOURCE_TENANT_ID + assert repo._creator_user_role == _sqlalchemy_workflow_execution_repository_CreatorUserRole.END_USER + + @pytest.mark.parametrize("sqlite_session", [TABLES], indirect=True) + def test_to_domain_model( + self, + sqlite_session_factory: sessionmaker[Session], + sqlite_session: Session, + account: Account, + ): + repo = SQLAlchemyWorkflowExecutionRepository( + session_factory=sqlite_session_factory, tenant_id=RESOURCE_TENANT_ID, user=account, app_id=None, triggered_from=None, ) - def test_init_no_tenant_id(self, sqlite_session_factory: sessionmaker[Session]): - user = _make_account() + db_model = WorkflowRun( + id=str(uuid4()), + tenant_id=account.current_tenant_id, + app_id=str(uuid4()), + workflow_id=str(uuid4()), + type=ModelWorkflowType.WORKFLOW, + triggered_from=WorkflowRunTriggeredFrom.APP_RUN, + version="1.0", + inputs=json.dumps({"in": "val"}), + outputs=json.dumps({"out": "val"}), + graph=json.dumps({"nodes": []}), + status=WorkflowExecutionStatus.SUCCEEDED, + error="some error", + elapsed_time=1.0, + total_tokens=50, + total_steps=3, + exceptions_count=1, + created_by_role=_sqlalchemy_workflow_execution_repository_CreatorUserRole.ACCOUNT, + created_by=account.id, + created_at=datetime.now(UTC), + finished_at=datetime.now(UTC), + ) + sqlite_session.add(db_model) + sqlite_session.commit() + sqlite_session.expunge_all() + persisted_model = sqlite_session.get(WorkflowRun, db_model.id) + assert persisted_model is not None - with pytest.raises(ValueError, match="tenant_id is required"): - SQLAlchemyWorkflowExecutionRepository( + domain_model = repo._to_domain_model(persisted_model) + + assert domain_model.id_ == persisted_model.id + assert domain_model.workflow_id == persisted_model.workflow_id + assert domain_model.status == WorkflowExecutionStatus.SUCCEEDED + assert domain_model.inputs == {"in": "val"} + assert domain_model.error_message == "some error" + + def test_to_db_model( + self, + sqlite_session_factory: sessionmaker[Session], + account: Account, + sample_workflow_execution: _sqlalchemy_workflow_execution_repository_WorkflowExecution, + ): + repo = SQLAlchemyWorkflowExecutionRepository( session_factory=sqlite_session_factory, - tenant_id="", - user=user, + tenant_id=RESOURCE_TENANT_ID, + user=account, + app_id="test_app", + triggered_from=WorkflowRunTriggeredFrom.DEBUGGING, + ) + + # Make elapsed time deterministic to avoid flaky tests + sample_workflow_execution.started_at = datetime(2023, 1, 1, 0, 0, 0, tzinfo=UTC) + sample_workflow_execution.finished_at = datetime(2023, 1, 1, 0, 0, 10, tzinfo=UTC) + + db_model = repo._to_db_model(sample_workflow_execution) + + assert db_model.id == sample_workflow_execution.id_ + assert db_model.tenant_id == repo._tenant_id + assert db_model.app_id == "test_app" + assert db_model.triggered_from == WorkflowRunTriggeredFrom.DEBUGGING + assert db_model.status == sample_workflow_execution.status.value + assert db_model.total_tokens == sample_workflow_execution.total_tokens + assert db_model.elapsed_time == 10.0 + + def test_to_db_model_edge_cases( + self, + sqlite_session_factory: sessionmaker[Session], + account: Account, + sample_workflow_execution: _sqlalchemy_workflow_execution_repository_WorkflowExecution, + ): + repo = SQLAlchemyWorkflowExecutionRepository( + session_factory=sqlite_session_factory, + tenant_id=RESOURCE_TENANT_ID, + user=account, + app_id="test_app", + triggered_from=WorkflowRunTriggeredFrom.DEBUGGING, + ) + # Test with empty/None fields + sample_workflow_execution.graph = None + sample_workflow_execution.inputs = None + sample_workflow_execution.outputs = None + sample_workflow_execution.error_message = None + sample_workflow_execution.finished_at = None + + db_model = repo._to_db_model(sample_workflow_execution) + + assert db_model.graph is None + assert db_model.inputs is None + assert db_model.outputs is None + assert db_model.error is None + assert db_model.elapsed_time == 0 + + def test_to_db_model_app_id_none( + self, + sqlite_session_factory: sessionmaker[Session], + account: Account, + sample_workflow_execution: _sqlalchemy_workflow_execution_repository_WorkflowExecution, + ): + repo = SQLAlchemyWorkflowExecutionRepository( + session_factory=sqlite_session_factory, + tenant_id=RESOURCE_TENANT_ID, + user=account, + app_id=None, + triggered_from=WorkflowRunTriggeredFrom.APP_RUN, + ) + + db_model = repo._to_db_model(sample_workflow_execution) + assert not hasattr(db_model, "app_id") or db_model.app_id is None + assert db_model.tenant_id == repo._tenant_id + + def test_to_db_model_missing_context( + self, + sqlite_session_factory: sessionmaker[Session], + account: Account, + sample_workflow_execution: _sqlalchemy_workflow_execution_repository_WorkflowExecution, + ): + repo = SQLAlchemyWorkflowExecutionRepository( + session_factory=sqlite_session_factory, + tenant_id=RESOURCE_TENANT_ID, + user=account, app_id=None, triggered_from=None, ) - def test_init_uses_resource_tenant_when_account_has_no_current_tenant( - self, sqlite_session_factory: sessionmaker[Session] - ): - user = _make_account() + # Test triggered_from missing + with pytest.raises(ValueError, match="triggered_from is required"): + repo._to_db_model(sample_workflow_execution) - repo = SQLAlchemyWorkflowExecutionRepository( - session_factory=sqlite_session_factory, - tenant_id=RESOURCE_TENANT_ID, - user=user, - app_id="test-app", - triggered_from=WorkflowRunTriggeredFrom.APP_RUN, - ) + repo._triggered_from = WorkflowRunTriggeredFrom.APP_RUN + repo._creator_user_id = None + with pytest.raises(ValueError, match="created_by is required"): + repo._to_db_model(sample_workflow_execution) - assert repo._tenant_id == RESOURCE_TENANT_ID - assert repo._creator_user_id == user.id + repo._creator_user_id = "some_id" + repo._creator_user_role = None + with pytest.raises(ValueError, match="created_by_role is required"): + repo._to_db_model(sample_workflow_execution) - def test_init_with_end_user(self, sqlite_session_factory: sessionmaker[Session], end_user: EndUser): - repo = SQLAlchemyWorkflowExecutionRepository( - session_factory=sqlite_session_factory, - tenant_id=RESOURCE_TENANT_ID, - user=end_user, - app_id=None, - triggered_from=None, - ) - assert repo._tenant_id == RESOURCE_TENANT_ID - assert repo._creator_user_role == CreatorUserRole.END_USER + @pytest.mark.parametrize("sqlite_session", [TABLES], indirect=True) + def test_save( + self, + sqlite_session_factory: sessionmaker[Session], + sqlite_session: Session, + account: Account, + sample_workflow_execution: _sqlalchemy_workflow_execution_repository_WorkflowExecution, + ): + repo = SQLAlchemyWorkflowExecutionRepository( + session_factory=sqlite_session_factory, + tenant_id=RESOURCE_TENANT_ID, + user=account, + app_id="test_app", + triggered_from=WorkflowRunTriggeredFrom.APP_RUN, + ) - @pytest.mark.parametrize("sqlite_session", [TABLES], indirect=True) - def test_to_domain_model( - self, - sqlite_session_factory: sessionmaker[Session], - sqlite_session: Session, - account: Account, - ): - repo = SQLAlchemyWorkflowExecutionRepository( - session_factory=sqlite_session_factory, - tenant_id=RESOURCE_TENANT_ID, - user=account, - app_id=None, - triggered_from=None, - ) - - db_model = WorkflowRun( - id=str(uuid4()), - tenant_id=account.current_tenant_id, - app_id=str(uuid4()), - workflow_id=str(uuid4()), - type=ModelWorkflowType.WORKFLOW, - triggered_from=WorkflowRunTriggeredFrom.APP_RUN, - version="1.0", - inputs=json.dumps({"in": "val"}), - outputs=json.dumps({"out": "val"}), - graph=json.dumps({"nodes": []}), - status=WorkflowExecutionStatus.SUCCEEDED, - error="some error", - elapsed_time=1.0, - total_tokens=50, - total_steps=3, - exceptions_count=1, - created_by_role=CreatorUserRole.ACCOUNT, - created_by=account.id, - created_at=datetime.now(UTC), - finished_at=datetime.now(UTC), - ) - sqlite_session.add(db_model) - sqlite_session.commit() - sqlite_session.expunge_all() - persisted_model = sqlite_session.get(WorkflowRun, db_model.id) - assert persisted_model is not None - - domain_model = repo._to_domain_model(persisted_model) - - assert domain_model.id_ == persisted_model.id - assert domain_model.workflow_id == persisted_model.workflow_id - assert domain_model.status == WorkflowExecutionStatus.SUCCEEDED - assert domain_model.inputs == {"in": "val"} - assert domain_model.error_message == "some error" - - def test_to_db_model( - self, - sqlite_session_factory: sessionmaker[Session], - account: Account, - sample_workflow_execution: WorkflowExecution, - ): - repo = SQLAlchemyWorkflowExecutionRepository( - session_factory=sqlite_session_factory, - tenant_id=RESOURCE_TENANT_ID, - user=account, - app_id="test_app", - triggered_from=WorkflowRunTriggeredFrom.DEBUGGING, - ) - - # Make elapsed time deterministic to avoid flaky tests - sample_workflow_execution.started_at = datetime(2023, 1, 1, 0, 0, 0, tzinfo=UTC) - sample_workflow_execution.finished_at = datetime(2023, 1, 1, 0, 0, 10, tzinfo=UTC) - - db_model = repo._to_db_model(sample_workflow_execution) - - assert db_model.id == sample_workflow_execution.id_ - assert db_model.tenant_id == repo._tenant_id - assert db_model.app_id == "test_app" - assert db_model.triggered_from == WorkflowRunTriggeredFrom.DEBUGGING - assert db_model.status == sample_workflow_execution.status.value - assert db_model.total_tokens == sample_workflow_execution.total_tokens - assert db_model.elapsed_time == 10.0 - - def test_to_db_model_edge_cases( - self, - sqlite_session_factory: sessionmaker[Session], - account: Account, - sample_workflow_execution: WorkflowExecution, - ): - repo = SQLAlchemyWorkflowExecutionRepository( - session_factory=sqlite_session_factory, - tenant_id=RESOURCE_TENANT_ID, - user=account, - app_id="test_app", - triggered_from=WorkflowRunTriggeredFrom.DEBUGGING, - ) - # Test with empty/None fields - sample_workflow_execution.graph = None - sample_workflow_execution.inputs = None - sample_workflow_execution.outputs = None - sample_workflow_execution.error_message = None - sample_workflow_execution.finished_at = None - - db_model = repo._to_db_model(sample_workflow_execution) - - assert db_model.graph is None - assert db_model.inputs is None - assert db_model.outputs is None - assert db_model.error is None - assert db_model.elapsed_time == 0 - - def test_to_db_model_app_id_none( - self, - sqlite_session_factory: sessionmaker[Session], - account: Account, - sample_workflow_execution: WorkflowExecution, - ): - repo = SQLAlchemyWorkflowExecutionRepository( - session_factory=sqlite_session_factory, - tenant_id=RESOURCE_TENANT_ID, - user=account, - app_id=None, - triggered_from=WorkflowRunTriggeredFrom.APP_RUN, - ) - - db_model = repo._to_db_model(sample_workflow_execution) - assert not hasattr(db_model, "app_id") or db_model.app_id is None - assert db_model.tenant_id == repo._tenant_id - - def test_to_db_model_missing_context( - self, - sqlite_session_factory: sessionmaker[Session], - account: Account, - sample_workflow_execution: WorkflowExecution, - ): - repo = SQLAlchemyWorkflowExecutionRepository( - session_factory=sqlite_session_factory, - tenant_id=RESOURCE_TENANT_ID, - user=account, - app_id=None, - triggered_from=None, - ) - - # Test triggered_from missing - with pytest.raises(ValueError, match="triggered_from is required"): - repo._to_db_model(sample_workflow_execution) - - repo._triggered_from = WorkflowRunTriggeredFrom.APP_RUN - repo._creator_user_id = None - with pytest.raises(ValueError, match="created_by is required"): - repo._to_db_model(sample_workflow_execution) - - repo._creator_user_id = "some_id" - repo._creator_user_role = None - with pytest.raises(ValueError, match="created_by_role is required"): - repo._to_db_model(sample_workflow_execution) - - @pytest.mark.parametrize("sqlite_session", [TABLES], indirect=True) - def test_save( - self, - sqlite_session_factory: sessionmaker[Session], - sqlite_session: Session, - account: Account, - sample_workflow_execution: WorkflowExecution, - ): - repo = SQLAlchemyWorkflowExecutionRepository( - session_factory=sqlite_session_factory, - tenant_id=RESOURCE_TENANT_ID, - user=account, - app_id="test_app", - triggered_from=WorkflowRunTriggeredFrom.APP_RUN, - ) - - repo.save(sample_workflow_execution) - - persisted_model = sqlite_session.get(WorkflowRun, sample_workflow_execution.id_) - assert persisted_model is not None - assert persisted_model.tenant_id == RESOURCE_TENANT_ID - assert persisted_model.inputs_dict == sample_workflow_execution.inputs - assert persisted_model.outputs_dict == sample_workflow_execution.outputs - - # Check cache - assert sample_workflow_execution.id_ in repo._execution_cache - cached_model = repo._execution_cache[sample_workflow_execution.id_] - assert cached_model.id == sample_workflow_execution.id_ - - @pytest.mark.parametrize("sqlite_session", [TABLES], indirect=True) - def test_save_uses_execution_started_at_when_record_does_not_exist( - self, - sqlite_session_factory: sessionmaker[Session], - sqlite_session: Session, - account: Account, - sample_workflow_execution: WorkflowExecution, - ): - repo = SQLAlchemyWorkflowExecutionRepository( - session_factory=sqlite_session_factory, - tenant_id=RESOURCE_TENANT_ID, - user=account, - app_id="test_app", - triggered_from=WorkflowRunTriggeredFrom.APP_RUN, - ) - - started_at = datetime(2026, 1, 1, 12, 0, 0, tzinfo=UTC) - sample_workflow_execution.started_at = started_at - - repo.save(sample_workflow_execution) - - persisted_model = sqlite_session.get(WorkflowRun, sample_workflow_execution.id_) - assert persisted_model is not None - assert persisted_model.created_at == started_at.replace(tzinfo=None) - - @pytest.mark.parametrize("sqlite_session", [TABLES], indirect=True) - def test_save_preserves_existing_created_at_when_record_already_exists( - self, - sqlite_session_factory: sessionmaker[Session], - sqlite_session: Session, - account: Account, - sample_workflow_execution: WorkflowExecution, - ): - repo = SQLAlchemyWorkflowExecutionRepository( - session_factory=sqlite_session_factory, - tenant_id=RESOURCE_TENANT_ID, - user=account, - app_id="test_app", - triggered_from=WorkflowRunTriggeredFrom.APP_RUN, - ) - - execution_id = sample_workflow_execution.id_ - existing_created_at = datetime(2026, 1, 1, 12, 0, 0, tzinfo=UTC) - sample_workflow_execution.started_at = existing_created_at - repo.save(sample_workflow_execution) - - sample_workflow_execution.started_at = datetime(2026, 1, 1, 12, 30, 0, tzinfo=UTC) - - repo.save(sample_workflow_execution) - - persisted_model = sqlite_session.get(WorkflowRun, execution_id) - assert persisted_model is not None - assert persisted_model.created_at == existing_created_at.replace(tzinfo=None) - - @pytest.mark.parametrize("sqlite_session", [TABLES], indirect=True) - def test_save_rejects_execution_owned_by_another_tenant( - self, - sqlite_session_factory: sessionmaker[Session], - sqlite_session: Session, - account: Account, - sample_workflow_execution: WorkflowExecution, - ): - other_tenant_id = str(uuid4()) - other_account = _make_account(tenant_id=str(uuid4())) - other_repo = SQLAlchemyWorkflowExecutionRepository( - session_factory=sqlite_session_factory, - tenant_id=other_tenant_id, - user=other_account, - app_id="test_app", - triggered_from=WorkflowRunTriggeredFrom.APP_RUN, - ) - other_repo.save(sample_workflow_execution) - - repo = SQLAlchemyWorkflowExecutionRepository( - session_factory=sqlite_session_factory, - tenant_id=RESOURCE_TENANT_ID, - user=account, - app_id="test_app", - triggered_from=WorkflowRunTriggeredFrom.APP_RUN, - ) - with pytest.raises(ValueError, match="Unauthorized access to workflow run"): repo.save(sample_workflow_execution) - sqlite_session.expire_all() - persisted_model = sqlite_session.get(WorkflowRun, sample_workflow_execution.id_) - assert persisted_model is not None - assert persisted_model.tenant_id == other_tenant_id + persisted_model = sqlite_session.get(WorkflowRun, sample_workflow_execution.id_) + assert persisted_model is not None + assert persisted_model.tenant_id == RESOURCE_TENANT_ID + assert persisted_model.inputs_dict == sample_workflow_execution.inputs + assert persisted_model.outputs_dict == sample_workflow_execution.outputs + + # Check cache + assert sample_workflow_execution.id_ in repo._execution_cache + cached_model = repo._execution_cache[sample_workflow_execution.id_] + assert cached_model.id == sample_workflow_execution.id_ + + @pytest.mark.parametrize("sqlite_session", [TABLES], indirect=True) + def test_save_uses_execution_started_at_when_record_does_not_exist( + self, + sqlite_session_factory: sessionmaker[Session], + sqlite_session: Session, + account: Account, + sample_workflow_execution: _sqlalchemy_workflow_execution_repository_WorkflowExecution, + ): + repo = SQLAlchemyWorkflowExecutionRepository( + session_factory=sqlite_session_factory, + tenant_id=RESOURCE_TENANT_ID, + user=account, + app_id="test_app", + triggered_from=WorkflowRunTriggeredFrom.APP_RUN, + ) + + started_at = datetime(2026, 1, 1, 12, 0, 0, tzinfo=UTC) + sample_workflow_execution.started_at = started_at + + repo.save(sample_workflow_execution) + + persisted_model = sqlite_session.get(WorkflowRun, sample_workflow_execution.id_) + assert persisted_model is not None + assert persisted_model.created_at == started_at.replace(tzinfo=None) + + @pytest.mark.parametrize("sqlite_session", [TABLES], indirect=True) + def test_save_preserves_existing_created_at_when_record_already_exists( + self, + sqlite_session_factory: sessionmaker[Session], + sqlite_session: Session, + account: Account, + sample_workflow_execution: _sqlalchemy_workflow_execution_repository_WorkflowExecution, + ): + repo = SQLAlchemyWorkflowExecutionRepository( + session_factory=sqlite_session_factory, + tenant_id=RESOURCE_TENANT_ID, + user=account, + app_id="test_app", + triggered_from=WorkflowRunTriggeredFrom.APP_RUN, + ) + + execution_id = sample_workflow_execution.id_ + existing_created_at = datetime(2026, 1, 1, 12, 0, 0, tzinfo=UTC) + sample_workflow_execution.started_at = existing_created_at + repo.save(sample_workflow_execution) + + sample_workflow_execution.started_at = datetime(2026, 1, 1, 12, 30, 0, tzinfo=UTC) + + repo.save(sample_workflow_execution) + + persisted_model = sqlite_session.get(WorkflowRun, execution_id) + assert persisted_model is not None + assert persisted_model.created_at == existing_created_at.replace(tzinfo=None) + + @pytest.mark.parametrize("sqlite_session", [TABLES], indirect=True) + def test_save_rejects_execution_owned_by_another_tenant( + self, + sqlite_session_factory: sessionmaker[Session], + sqlite_session: Session, + account: Account, + sample_workflow_execution: _sqlalchemy_workflow_execution_repository_WorkflowExecution, + ): + other_tenant_id = str(uuid4()) + other_account = _make_account(tenant_id=str(uuid4())) + other_repo = SQLAlchemyWorkflowExecutionRepository( + session_factory=sqlite_session_factory, + tenant_id=other_tenant_id, + user=other_account, + app_id="test_app", + triggered_from=WorkflowRunTriggeredFrom.APP_RUN, + ) + other_repo.save(sample_workflow_execution) + + repo = SQLAlchemyWorkflowExecutionRepository( + session_factory=sqlite_session_factory, + tenant_id=RESOURCE_TENANT_ID, + user=account, + app_id="test_app", + triggered_from=WorkflowRunTriggeredFrom.APP_RUN, + ) + with pytest.raises(ValueError, match="Unauthorized access to workflow run"): + repo.save(sample_workflow_execution) + + sqlite_session.expire_all() + persisted_model = sqlite_session.get(WorkflowRun, sample_workflow_execution.id_) + assert persisted_model is not None + assert persisted_model.tenant_id == other_tenant_id + + +"""Configured history backends retain synchronous, monotonic task control.""" + + +class TestWorkflowControlPersistenceCases: + @staticmethod + @pytest.fixture + def logstore(monkeypatch: pytest.MonkeyPatch, mocker: MockerFixture) -> AliyunLogStore: + monkeypatch.setattr(AliyunLogStore, "_instance", None) + monkeypatch.setattr(AliyunLogStore, "_initialized", False) + monkeypatch.setenv("ALIYUN_SLS_ENDPOINT", "https://example.com") + mocker.patch.object(AliyunLogStore, "_check_endpoint_connectivity") + store = AliyunLogStore() + store._use_pg_protocol = True + mocker.patch.object(store, "put_log") + return store + + @staticmethod + @pytest.mark.parametrize("backend", ["sql", "celery", "logstore"]) + def test_task_owner_is_durable_before_start_and_delayed_writes_cannot_revive_stop( + sqlite_session_factory: sessionmaker[Session], logstore: AliyunLogStore, mocker: MockerFixture, backend: str + ) -> None: + queued_save = mocker.patch( + "core.repositories.celery_workflow_execution_repository.save_workflow_execution_task.delay" + ) + get_logs = mocker.patch.object(logstore, "get_logs") + repository_class = { + "sql": SQLAlchemyWorkflowExecutionRepository, + "celery": CeleryWorkflowExecutionRepository, + "logstore": LogstoreWorkflowExecutionRepository, + }[backend] + user = Account(name="Test", email="test@example.com") + user.id = "account-1" + repository = repository_class( + session_factory=sqlite_session_factory, + tenant_id="tenant-1", + user=user, + app_id="app-1", + triggered_from=WorkflowRunTriggeredFrom.APP_RUN, + ) + if isinstance(repository, LogstoreWorkflowExecutionRepository): + repository._enable_dual_write = False + execution = _workflow_control_persistence_WorkflowExecution( + id_="run-1", + workflow_id="workflow-1", + workflow_type=WorkflowType.WORKFLOW, + workflow_version="1", + graph={"nodes": [], "edges": []}, + inputs={}, + started_at=naive_utc_now(), + ) + repository.save_synchronously(execution) + queued_save.assert_not_called() + control = DifyAPISQLAlchemyWorkflowRunRepository(sqlite_session_factory) + control.bind_workflow_task(tenant_id="tenant-1", app_id="app-1", workflow_run_id="run-1", task_id="task-1") + with sqlite_session_factory() as session: + run = session.get(WorkflowRun, "run-1") + assert run is not None + assert run.task_id == "task-1" + assert (run.created_by_role, run.created_by) == ( + _workflow_control_persistence_CreatorUserRole.ACCOUNT, + "account-1", + ) + execution.status = WorkflowExecutionStatus.PAUSED + repository.save_synchronously(execution) + stopped = control.stop_paused_workflow_task(tenant_id="tenant-1", app_id="app-1", task_id="task-1", owner=None) + assert stopped is not None + assert stopped.status == WorkflowExecutionStatus.STOPPED + + # A writer already queued or running before the stop cannot restore RUNNING. + execution.status = WorkflowExecutionStatus.RUNNING + repository.save_synchronously(execution) + with sqlite_session_factory.begin() as session: + run = session.get(WorkflowRun, "run-1") + assert run is not None + _update_workflow_run_from_execution(run, execution) + with sqlite_session_factory() as session: + run = session.get(WorkflowRun, "run-1") + assert run is not None + assert run.status == WorkflowExecutionStatus.STOPPED + assert run.stop_requested_at is not None + if backend == "logstore": + # LogStore implements the configured run-control subset of the history protocol. + api = LogstoreAPIWorkflowRunRepository(sqlite_session_factory) # pyrefly: ignore[bad-instantiation] + api._enable_dual_read = False + found = api.get_workflow_run_by_id("tenant-1", "app-1", "run-1") + assert found is not None + assert found.status == WorkflowExecutionStatus.STOPPED + unscoped = api.get_workflow_run_by_id_without_tenant("run-1") + assert unscoped is not None + assert unscoped.status == WorkflowExecutionStatus.STOPPED + get_logs.assert_not_called() + + @staticmethod + @pytest.mark.usefixtures("logstore") + def test_logstore_pause_lifecycle_uses_sql_control_with_dual_write_disabled( + sqlite_session_factory: sessionmaker[Session], mocker: MockerFixture + ) -> None: + mocker.patch("repositories.sqlalchemy_api_workflow_run_repository.storage.save") + mocker.patch("repositories.sqlalchemy_api_workflow_run_repository.storage.delete") + user = Account(name="Test", email="test@example.com") + user.id = "account-1" + sql = SQLAlchemyWorkflowExecutionRepository( + session_factory=sqlite_session_factory, + tenant_id="tenant-1", + user=user, + app_id="app-1", + triggered_from=WorkflowRunTriggeredFrom.APP_RUN, + ) + sql.save_synchronously( + _workflow_control_persistence_WorkflowExecution( + id_="run-1", + workflow_id="workflow-1", + workflow_type=WorkflowType.WORKFLOW, + workflow_version="1", + graph={}, + inputs={}, + started_at=naive_utc_now(), + ) + ) + api = LogstoreAPIWorkflowRunRepository(sqlite_session_factory) # pyrefly: ignore[bad-instantiation] + api.bind_workflow_task(tenant_id="tenant-1", app_id="app-1", workflow_run_id="run-1", task_id="task-1") + pause = api.create_workflow_pause("run-1", "account-1", "{}", []) + persisted_pause = api.get_workflow_pause("run-1") + assert persisted_pause is not None + assert persisted_pause.id == pause.id + resume_callbacks: list[None] = [] + resumed = api.resume_workflow_pause("run-1", pause, before_resume=lambda: resume_callbacks.append(None)) + assert resume_callbacks == [None] + api.delete_workflow_pause(resumed) + assert api.get_workflow_pause("run-1") is None + + @staticmethod + @pytest.mark.parametrize( + ("record_type", "node_status"), + [ + ("workflow", WorkflowNodeExecutionStatus.PAUSED), + ("node", WorkflowNodeExecutionStatus.PAUSED), + ("node", WorkflowNodeExecutionStatus.RETRY), + ], + ) + def test_logstore_stop_orders_before_delayed_upload_and_overrides_stale_indexes( + sqlite_session_factory: sessionmaker[Session], + logstore: AliyunLogStore, + mocker: MockerFixture, + record_type: str, + node_status: WorkflowNodeExecutionStatus, + ) -> None: + user = Account(name="Test", email="test@example.com") + user.id = "account-1" + run = _workflow_control_persistence_WorkflowExecution( + accumulated_elapsed_time=3.5, + id_="run-1", + workflow_id="workflow-1", + workflow_type=WorkflowType.WORKFLOW, + workflow_version="1", + graph={}, + inputs={}, + started_at=naive_utc_now(), + ) + run.status = WorkflowExecutionStatus.PAUSED + SQLAlchemyWorkflowExecutionRepository( + session_factory=sqlite_session_factory, + tenant_id="tenant-1", + user=user, + app_id="app-1", + triggered_from=WorkflowRunTriggeredFrom.APP_RUN, + ).save(run) + control = DifyAPISQLAlchemyWorkflowRunRepository(sqlite_session_factory) + control.bind_workflow_task(tenant_id="tenant-1", app_id="app-1", workflow_run_id="run-1", task_id="task-1") + node = WorkflowNodeExecution( + id="node-1", + node_execution_id="node-1", + workflow_id="workflow-1", + workflow_execution_id="run-1", + index=1, + node_id="review", + node_type="human-input", + title="Review", + created_at=run.started_at, + status=node_status, + ) + SQLAlchemyWorkflowNodeExecutionRepository( + session_factory=sqlite_session_factory, + tenant_id="tenant-1", + user=user, + app_id="app-1", + triggered_from=WorkflowNodeExecutionTriggeredFrom.WORKFLOW_RUN, + ).save_synchronously(node) + records: list[dict[str, str]] = [] + interrupted = False + + def append_after_concurrent_stop(_store: str, fields: list[tuple[str, str]]) -> None: + nonlocal interrupted + if not interrupted: + interrupted = True + stopped = control.stop_paused_workflow_task( + tenant_id="tenant-1", app_id="app-1", task_id="task-1", owner=None + ) + assert stopped is not None + if record_type == "workflow": + run_writer.save( + run.model_copy( + update={ + "status": WorkflowExecutionStatus.STOPPED, + "finished_at": stopped.finished_at, + "error_message": stopped.error, + } + ) + ) + else: + node_writer.save_many( + [ + node.model_copy( + update={ + "status": WorkflowNodeExecutionStatus.FAILED, + "finished_at": stopped.finished_at, + "error": stopped.error, + } + ) + ] + ) + records.append(dict(fields)) + + mocker.patch.object(logstore, "put_log", side_effect=append_after_concurrent_stop) + run_writer = LogstoreWorkflowExecutionRepository( + session_factory=sqlite_session_factory, + tenant_id="tenant-1", + user=user, + app_id="app-1", + triggered_from=WorkflowRunTriggeredFrom.APP_RUN, + ) + node_writer = LogstoreWorkflowNodeExecutionRepository( + session_factory=sqlite_session_factory, + tenant_id="tenant-1", + user=user, + app_id="app-1", + triggered_from=WorkflowNodeExecutionTriggeredFrom.WORKFLOW_RUN, + ) + run_writer._enable_dual_write = node_writer._enable_dual_write = False + if record_type == "workflow": + run_writer.save(run) + expected_status = "stopped" + assert max(records, key=lambda row: int(row["log_version"]))["status"] == expected_status + run.accumulated_elapsed_time = 999.0 + run_writer.save(run) # Another late RUNNING/PAUSED write also remains terminal. + assert float(records[-1]["elapsed_time"]) == 3.5 + else: + node_writer.save(node) + expected_status = "failed" + assert max(records, key=lambda row: int(row["log_version"]))["status"] == expected_status + stale_node = next(row for row in records if row["status"] == node_status) + mocker.patch.object(logstore, "execute_sql", return_value=[stale_node]) + api = LogstoreAPIWorkflowNodeExecutionRepository(sqlite_session_factory) # pyrefly: ignore[bad-instantiation] + assert api.get_executions_by_workflow_run("tenant-1", "app-1", "run-1")[0].status == expected_status + by_id = api.get_execution_by_id("node-1", "tenant-1") + assert by_id is not None + assert by_id.status == expected_status + assert ( + api.get_execution_snapshots_by_workflow_run("tenant-1", "app-1", "workflow-1", "workflow-run", "run-1")[ + 0 + ].status + == expected_status + ) + assert node_writer.get_by_workflow_execution("run-1")[0].status == expected_status + node_writer.save(node) + assert records[-1]["status"] == expected_status + + @staticmethod + @pytest.mark.parametrize( + "status", + [ + WorkflowNodeExecutionStatus.SUCCEEDED, + WorkflowNodeExecutionStatus.FAILED, + WorkflowNodeExecutionStatus.EXCEPTION, + WorkflowNodeExecutionStatus.STOPPED, + ], + ) + def test_stop_preserves_completed_node_history(status: WorkflowNodeExecutionStatus) -> None: + finished_at = naive_utc_now() + stopped = WorkflowRun( + id="run-1", + tenant_id="tenant-1", + app_id="app-1", + workflow_id="workflow-1", + type=WorkflowType.WORKFLOW, + triggered_from=WorkflowRunTriggeredFrom.APP_RUN, + version="1", + status=WorkflowExecutionStatus.STOPPED, + created_by_role=_workflow_control_persistence_CreatorUserRole.ACCOUNT, + created_by="account-1", + error="User requested stop", + finished_at=finished_at, + stop_requested_at=finished_at, + ) + node = WorkflowNodeExecution( + id="node-1", + workflow_id="workflow-1", + workflow_execution_id="run-1", + index=1, + node_id="answer", + node_type="llm", + title="Answer", + status=status, + outputs={"answer": "preserved"}, + error=None if status == WorkflowNodeExecutionStatus.SUCCEEDED else "Original terminal error", + elapsed_time=2.5, + created_at=finished_at, + finished_at=finished_at, + ) + original = node.model_dump() + + apply_workflow_stop_to_node(node, stopped) + + assert node.model_dump() == original diff --git a/api/tests/unit_tests/core/repositories/test_sqlalchemy_workflow_node_execution_repository.py b/api/tests/unit_tests/core/repositories/test_sqlalchemy_workflow_node_execution_repository.py index 3b67771dcdf..be042e7aa9f 100644 --- a/api/tests/unit_tests/core/repositories/test_sqlalchemy_workflow_node_execution_repository.py +++ b/api/tests/unit_tests/core/repositories/test_sqlalchemy_workflow_node_execution_repository.py @@ -1,11 +1,10 @@ -"""SQLite-backed tests for the workflow node execution repository.""" - from __future__ import annotations import json import logging -from collections.abc import Callable, Generator, Mapping +from collections.abc import Callable, Generator, Iterator, Mapping from contextlib import contextmanager +from dataclasses import dataclass from datetime import UTC, datetime from types import SimpleNamespace from typing import Any @@ -28,16 +27,52 @@ from core.repositories.sqlalchemy_workflow_node_execution_repository import ( _replace_or_append_offload, ) from core.workflow.node_execution_process_data import WORKFLOW_TOOL_ROOT_APP_ID_KEY -from graphon.entities import WorkflowNodeExecution -from graphon.enums import BuiltinNodeTypes, WorkflowNodeExecutionMetadataKey, WorkflowNodeExecutionStatus -from models import Account, EndUser +from graphon.entities import ( + WorkflowNodeExecution as _sqlalchemy_workflow_node_execution_repository_WorkflowNodeExecution, +) +from graphon.entities.workflow_node_execution import ( + WorkflowNodeExecution as _workflow_node_execution_conflict_handling_WorkflowNodeExecution, +) +from graphon.entities.workflow_node_execution import ( + WorkflowNodeExecution as _workflow_node_execution_truncation_WorkflowNodeExecution, +) +from graphon.entities.workflow_node_execution import ( + WorkflowNodeExecutionStatus as _ConflictExecutionStatus, +) +from graphon.entities.workflow_node_execution import ( + WorkflowNodeExecutionStatus as _workflow_node_execution_truncation_WorkflowNodeExecutionStatus, +) +from graphon.enums import BuiltinNodeTypes, WorkflowNodeExecutionMetadataKey +from graphon.enums import ( + WorkflowNodeExecutionStatus as _RepositoryExecutionStatus, +) +from libs.datetime_utils import naive_utc_now +from models import Account, EndUser, Tenant +from models import WorkflowNodeExecutionModel as _workflow_node_execution_conflict_handling_WorkflowNodeExecutionModel +from models import ( + WorkflowNodeExecutionTriggeredFrom as _workflow_node_execution_conflict_handling_WorkflowNodeExecutionTriggeredFrom, +) +from models import ( + WorkflowNodeExecutionTriggeredFrom as _workflow_node_execution_truncation_WorkflowNodeExecutionTriggeredFrom, +) from models.enums import CreatorUserRole, ExecutionOffLoadType from models.model import UploadFile -from models.workflow import WorkflowNodeExecutionModel, WorkflowNodeExecutionOffload, WorkflowNodeExecutionTriggeredFrom -from tests.unit_tests.model_factories import make_account, make_tenant, make_upload_file +from models.workflow import ( + WorkflowNodeExecutionModel as _sqlalchemy_workflow_node_execution_repository_WorkflowNodeExecutionModel, +) +from models.workflow import WorkflowNodeExecutionModel as _workflow_node_execution_truncation_WorkflowNodeExecutionModel +from models.workflow import WorkflowNodeExecutionOffload +from models.workflow import ( + WorkflowNodeExecutionTriggeredFrom as _RepositoryTriggeredFrom, +) +from tests.helpers.unit.model_factories import make_account, make_tenant, make_upload_file + +"""SQLite-backed tests for the workflow node execution repository.""" -def _account(*, tenant_id: str = "tenant-1", user_id: str = "user-1") -> Account: +def _sqlalchemy_workflow_node_execution_repository_account( + *, tenant_id: str = "tenant-1", user_id: str = "user-1" +) -> Account: return make_account( account_id=user_id, name="Test Account", @@ -72,18 +107,18 @@ def _offload(type_: ExecutionOffLoadType, *, file_id: str = "file-1") -> Workflo ) -def _execution( +def _sqlalchemy_workflow_node_execution_repository_execution( *, execution_id: str = "execution-1", node_execution_id: str = "node-execution-1", run_id: str = "run-1", index: int = 1, - status: WorkflowNodeExecutionStatus = WorkflowNodeExecutionStatus.SUCCEEDED, + status: _RepositoryExecutionStatus = _RepositoryExecutionStatus.SUCCEEDED, inputs: Mapping[str, Any] | None = None, outputs: Mapping[str, Any] | None = None, process_data: Mapping[str, Any] | None = None, -) -> WorkflowNodeExecution: - return WorkflowNodeExecution( +) -> _sqlalchemy_workflow_node_execution_repository_WorkflowNodeExecution: + return _sqlalchemy_workflow_node_execution_repository_WorkflowNodeExecution( id=execution_id, node_execution_id=node_execution_id, workflow_id="workflow-1", @@ -112,7 +147,7 @@ def _repository( tenant_id: str = "tenant-1", app_id: str | None = "app-1", user: Account | EndUser | None = None, - triggered_from: WorkflowNodeExecutionTriggeredFrom | None = WorkflowNodeExecutionTriggeredFrom.WORKFLOW_RUN, + triggered_from: _RepositoryTriggeredFrom | None = _RepositoryTriggeredFrom.WORKFLOW_RUN, ) -> SQLAlchemyWorkflowNodeExecutionRepository: monkeypatch.setattr( "core.repositories.sqlalchemy_workflow_node_execution_repository.FileService", @@ -121,7 +156,7 @@ def _repository( return SQLAlchemyWorkflowNodeExecutionRepository( session_factory=factory, tenant_id=tenant_id, - user=user or _account(tenant_id=tenant_id), + user=user or _sqlalchemy_workflow_node_execution_repository_account(tenant_id=tenant_id), app_id=app_id, triggered_from=triggered_from, ) @@ -147,549 +182,1062 @@ def _raise_on_execution_insert(engine: Engine) -> Generator[None]: event.remove(engine, "before_cursor_execute", raise_error) -def test_init_accepts_real_engine_and_sessionmaker_and_sets_role( - monkeypatch: pytest.MonkeyPatch, sqlite_engine: Engine, sqlite_session_factory: sessionmaker[Session] -) -> None: - engine_repo = _repository(monkeypatch, sqlite_engine) - assert isinstance(engine_repo._session_factory, sessionmaker) - end_user_repo = _repository(monkeypatch, sqlite_session_factory, user=_end_user()) - assert end_user_repo._creator_user_role.value == "end_user" +class TestSqlalchemyWorkflowNodeExecutionRepositoryCases: + @staticmethod + def test_init_accepts_real_engine_and_sessionmaker_and_sets_role( + monkeypatch: pytest.MonkeyPatch, sqlite_engine: Engine, sqlite_session_factory: sessionmaker[Session] + ) -> None: + engine_repo = _repository(monkeypatch, sqlite_engine) + assert isinstance(engine_repo._session_factory, sessionmaker) + end_user_repo = _repository(monkeypatch, sqlite_session_factory, user=_end_user()) + assert end_user_repo._creator_user_role.value == "end_user" - -@pytest.mark.parametrize("source_app", ["app-1", "source-app"]) -def test_workflow_tool_scope_isolates_execution_history_with_same_tenant_and_creator( - monkeypatch: pytest.MonkeyPatch, sqlite_session_factory: sessionmaker[Session], source_app: str -) -> None: - caller = _repository(monkeypatch, sqlite_session_factory, user=_end_user()) - source = caller.for_workflow_tool(source_app) - caller.save_synchronously(_execution(execution_id="caller", node_execution_id="caller")) - source.save_synchronously(_execution(execution_id="source", node_execution_id="source")) - - assert [node.id for node in caller.get_by_workflow_execution("run-1")] == ["caller"] - assert [node.id for node in source.get_by_workflow_execution("run-1")] == ["source"] - with sqlite_session_factory() as session: - row = session.get(WorkflowNodeExecutionModel, "source") - assert row is not None - assert (row.tenant_id, row.app_id, row.created_by, row.created_by_role, row.triggered_from) == ( - "tenant-1", - source_app, - "end-user-1", - CreatorUserRole.END_USER, - WorkflowNodeExecutionTriggeredFrom.WORKFLOW_TOOL, + @staticmethod + @pytest.mark.parametrize("source_app", ["app-1", "source-app"]) + def test_workflow_tool_scope_isolates_execution_history_with_same_tenant_and_creator( + monkeypatch: pytest.MonkeyPatch, sqlite_session_factory: sessionmaker[Session], source_app: str + ) -> None: + caller = _repository(monkeypatch, sqlite_session_factory, user=_end_user()) + source = caller.for_workflow_tool(source_app) + caller.save_synchronously( + _sqlalchemy_workflow_node_execution_repository_execution(execution_id="caller", node_execution_id="caller") + ) + source.save_synchronously( + _sqlalchemy_workflow_node_execution_repository_execution(execution_id="source", node_execution_id="source") ) + assert [node.id for node in caller.get_by_workflow_execution("run-1")] == ["caller"] + assert [node.id for node in source.get_by_workflow_execution("run-1")] == ["source"] + with sqlite_session_factory() as session: + row = session.get(_sqlalchemy_workflow_node_execution_repository_WorkflowNodeExecutionModel, "source") + assert row is not None + assert (row.tenant_id, row.app_id, row.created_by, row.created_by_role, row.triggered_from) == ( + "tenant-1", + source_app, + "end-user-1", + CreatorUserRole.END_USER, + _RepositoryTriggeredFrom.WORKFLOW_TOOL, + ) -def test_init_rejects_invalid_factory_and_missing_tenant(monkeypatch: pytest.MonkeyPatch) -> None: - monkeypatch.setattr( - "core.repositories.sqlalchemy_workflow_node_execution_repository.FileService", - lambda *_args: SimpleNamespace(upload_file=Mock()), - ) - with pytest.raises(ValueError, match="Invalid session_factory type"): - SQLAlchemyWorkflowNodeExecutionRepository( - session_factory=object(), # type: ignore[arg-type] - tenant_id="tenant-1", - user=_account(), - app_id=None, - triggered_from=WorkflowNodeExecutionTriggeredFrom.WORKFLOW_RUN, + @staticmethod + def test_init_rejects_invalid_factory_and_missing_tenant(monkeypatch: pytest.MonkeyPatch) -> None: + monkeypatch.setattr( + "core.repositories.sqlalchemy_workflow_node_execution_repository.FileService", + lambda *_args: SimpleNamespace(upload_file=Mock()), ) - user = _account() - user._current_tenant = None - with pytest.raises(ValueError, match="tenant_id"): - SQLAlchemyWorkflowNodeExecutionRepository( - session_factory=sessionmaker(), - tenant_id="", + with pytest.raises(ValueError, match="Invalid session_factory type"): + SQLAlchemyWorkflowNodeExecutionRepository( + session_factory=object(), # type: ignore[arg-type] + tenant_id="tenant-1", + user=_sqlalchemy_workflow_node_execution_repository_account(), + app_id=None, + triggered_from=_RepositoryTriggeredFrom.WORKFLOW_RUN, + ) + user = _sqlalchemy_workflow_node_execution_repository_account() + user._current_tenant = None + with pytest.raises(ValueError, match="tenant_id"): + SQLAlchemyWorkflowNodeExecutionRepository( + session_factory=sessionmaker(), + tenant_id="", + user=user, + app_id=None, + triggered_from=_RepositoryTriggeredFrom.WORKFLOW_RUN, + ) + + @staticmethod + def test_init_uses_resource_tenant_when_account_has_no_current_tenant( + monkeypatch: pytest.MonkeyPatch, sqlite_session_factory: sessionmaker[Session] + ) -> None: + user = _sqlalchemy_workflow_node_execution_repository_account() + user._current_tenant = None + + repo = _repository( + monkeypatch, + sqlite_session_factory, + tenant_id="resource-tenant", user=user, - app_id=None, - triggered_from=WorkflowNodeExecutionTriggeredFrom.WORKFLOW_RUN, ) + assert repo._tenant_id == "resource-tenant" + assert repo._creator_user_id == user.id -def test_init_uses_resource_tenant_when_account_has_no_current_tenant( - monkeypatch: pytest.MonkeyPatch, sqlite_session_factory: sessionmaker[Session] -) -> None: - user = _account() - user._current_tenant = None + @staticmethod + def test_helper_functions_and_truncator_configuration( + monkeypatch: pytest.MonkeyPatch, sqlite_session_factory: sessionmaker[Session] + ) -> None: + assert _deterministic_json_dump({"b": 1, "a": 2}) == '{"a": 2, "b": 1}' + assert _find_first([], lambda _value: True) is None + assert _find_first([1, 2, 3], lambda value: value > 1) == 2 + inputs = _offload(ExecutionOffLoadType.INPUTS) + outputs = _offload(ExecutionOffLoadType.OUTPUTS) + assert _find_first([inputs, outputs], _filter_by_offload_type(ExecutionOffLoadType.OUTPUTS)) is outputs + replaced = _replace_or_append_offload([inputs, outputs], _offload(ExecutionOffLoadType.INPUTS)) + assert [item.type_ for item in replaced] == [ExecutionOffLoadType.OUTPUTS, ExecutionOffLoadType.INPUTS] - repo = _repository( - monkeypatch, - sqlite_session_factory, - tenant_id="resource-tenant", - user=user, - ) + created: dict[str, int] = {} - assert repo._tenant_id == "resource-tenant" - assert repo._creator_user_id == user.id + class Truncator: + def __init__(self, *, max_size_bytes: int, array_element_limit: int, string_length_limit: int) -> None: + created.update( + max_size_bytes=max_size_bytes, + array_element_limit=array_element_limit, + string_length_limit=string_length_limit, + ) + monkeypatch.setattr( + "core.repositories.sqlalchemy_workflow_node_execution_repository.VariableTruncator", Truncator + ) + _repository(monkeypatch, sqlite_session_factory)._create_truncator() + assert created["max_size_bytes"] == dify_config.WORKFLOW_VARIABLE_TRUNCATION_MAX_SIZE -def test_helper_functions_and_truncator_configuration( - monkeypatch: pytest.MonkeyPatch, sqlite_session_factory: sessionmaker[Session] -) -> None: - assert _deterministic_json_dump({"b": 1, "a": 2}) == '{"a": 2, "b": 1}' - assert _find_first([], lambda _value: True) is None - assert _find_first([1, 2, 3], lambda value: value > 1) == 2 - inputs = _offload(ExecutionOffLoadType.INPUTS) - outputs = _offload(ExecutionOffLoadType.OUTPUTS) - assert _find_first([inputs, outputs], _filter_by_offload_type(ExecutionOffLoadType.OUTPUTS)) is outputs - replaced = _replace_or_append_offload([inputs, outputs], _offload(ExecutionOffLoadType.INPUTS)) - assert [item.type_ for item in replaced] == [ExecutionOffLoadType.OUTPUTS, ExecutionOffLoadType.INPUTS] + @staticmethod + def test_to_db_model_uses_context_and_deterministic_json( + monkeypatch: pytest.MonkeyPatch, sqlite_session_factory: sessionmaker[Session] + ) -> None: + repo = _repository(monkeypatch, sqlite_session_factory) + db_model = repo._to_db_model( + _sqlalchemy_workflow_node_execution_repository_execution( + inputs={"b": 1, "a": 2}, + process_data={"agent_workspace_binding_id": "participant-1"}, + ) + ) + assert json.loads(db_model.inputs or "{}") == {"a": 2, "b": 1} + assert db_model.tenant_id == "tenant-1" + assert db_model.app_id == "app-1" + assert db_model.created_by == "user-1" + assert db_model.created_by_role == CreatorUserRole.ACCOUNT + assert json.loads(db_model.execution_metadata or "{}") == {"total_tokens": 1} + assert db_model.agent_workspace_binding_id is None + assert ( + _repository(monkeypatch, sqlite_session_factory, app_id=None) + ._to_db_model(_sqlalchemy_workflow_node_execution_repository_execution()) + .app_id + is None + ) + repo._triggered_from = None + with pytest.raises(ValueError, match="triggered_from is required"): + repo._to_db_model(_sqlalchemy_workflow_node_execution_repository_execution()) - created: dict[str, int] = {} + @staticmethod + def test_to_db_model_requires_creator_context( + monkeypatch: pytest.MonkeyPatch, sqlite_session_factory: sessionmaker[Session] + ) -> None: + repo = _repository(monkeypatch, sqlite_session_factory) + execution = _sqlalchemy_workflow_node_execution_repository_execution() - class Truncator: - def __init__(self, *, max_size_bytes: int, array_element_limit: int, string_length_limit: int) -> None: - created.update( - max_size_bytes=max_size_bytes, - array_element_limit=array_element_limit, - string_length_limit=string_length_limit, + monkeypatch.setattr(repo, "_creator_user_id", None) + with pytest.raises(ValueError, match="created_by is required"): + repo._to_db_model(execution) + + monkeypatch.setattr(repo, "_creator_user_id", "user-1") + monkeypatch.setattr(repo, "_creator_user_role", None) + with pytest.raises(ValueError, match="created_by_role is required"): + repo._to_db_model(execution) + + @staticmethod + def test_json_encode_uses_runtime_converter(monkeypatch: pytest.MonkeyPatch) -> None: + class Converter: + def to_json_encodable(self, values: Mapping[str, Any]) -> Mapping[str, Any]: + return {"wrapped": values["value"]} + + monkeypatch.setattr( + "core.repositories.sqlalchemy_workflow_node_execution_repository.WorkflowRuntimeTypeConverter", + Converter, + ) + + assert SQLAlchemyWorkflowNodeExecutionRepository._json_encode({"value": 1}) == '{"wrapped": 1}' + + @staticmethod + def test_save_inserts_and_updates_persisted_execution( + monkeypatch: pytest.MonkeyPatch, sqlite_session_factory: sessionmaker[Session] + ) -> None: + repo = _repository(monkeypatch, sqlite_session_factory) + execution = _sqlalchemy_workflow_node_execution_repository_execution( + inputs={"value": 1}, outputs={"result": "first"} + ) + repo.save(execution) + with sqlite_session_factory() as session: + persisted = session.get( + _sqlalchemy_workflow_node_execution_repository_WorkflowNodeExecutionModel, execution.id + ) + assert persisted is not None + assert persisted.outputs_dict == {"result": "first"} + execution.title = "Updated" + execution.outputs = {"result": "second"} + repo.save(execution) + with sqlite_session_factory() as session: + persisted = session.get( + _sqlalchemy_workflow_node_execution_repository_WorkflowNodeExecutionModel, execution.id + ) + assert persisted is not None + assert persisted.title == "Updated" + assert persisted.outputs_dict == {"result": "second"} + assert execution.node_execution_id is not None + assert repo._node_execution_cache[execution.node_execution_id].id == execution.id + + @staticmethod + def test_save_owned_session_rolls_back_failed_insert( + monkeypatch: pytest.MonkeyPatch, + sqlite_engine: Engine, + sqlite_session_factory: sessionmaker[Session], + ) -> None: + repo = _repository(monkeypatch, sqlite_session_factory) + with _raise_on_execution_insert(sqlite_engine), pytest.raises(RuntimeError, match="forced execution INSERT"): + repo.save(_sqlalchemy_workflow_node_execution_repository_execution()) + with sqlite_session_factory() as session: + assert ( + session.scalar(select(_sqlalchemy_workflow_node_execution_repository_WorkflowNodeExecutionModel)) + is None ) - monkeypatch.setattr("core.repositories.sqlalchemy_workflow_node_execution_repository.VariableTruncator", Truncator) - _repository(monkeypatch, sqlite_session_factory)._create_truncator() - assert created["max_size_bytes"] == dify_config.WORKFLOW_VARIABLE_TRUNCATION_MAX_SIZE - - -def test_to_db_model_uses_context_and_deterministic_json( - monkeypatch: pytest.MonkeyPatch, sqlite_session_factory: sessionmaker[Session] -) -> None: - repo = _repository(monkeypatch, sqlite_session_factory) - db_model = repo._to_db_model( - _execution( - inputs={"b": 1, "a": 2}, - process_data={"agent_workspace_binding_id": "participant-1"}, + @staticmethod + def test_save_execution_data_updates_existing_and_creates_missing( + monkeypatch: pytest.MonkeyPatch, sqlite_session_factory: sessionmaker[Session] + ) -> None: + repo = _repository(monkeypatch, sqlite_session_factory) + existing = _sqlalchemy_workflow_node_execution_repository_execution( + inputs={"initial": True}, + process_data={ + "workflow_agent_binding_id": "binding-1", + "workflow_tool_invocation_id": "tool-call-1", + "workflow_tool_root_app_id": "root-app", + }, ) - ) - assert json.loads(db_model.inputs or "{}") == {"a": 2, "b": 1} - assert db_model.tenant_id == "tenant-1" - assert db_model.app_id == "app-1" - assert db_model.created_by == "user-1" - assert db_model.created_by_role == CreatorUserRole.ACCOUNT - assert json.loads(db_model.execution_metadata or "{}") == {"total_tokens": 1} - assert db_model.agent_workspace_binding_id is None - assert _repository(monkeypatch, sqlite_session_factory, app_id=None)._to_db_model(_execution()).app_id is None - repo._triggered_from = None - with pytest.raises(ValueError, match="triggered_from is required"): - repo._to_db_model(_execution()) + repo.save(existing) + existing.inputs = {"updated": True} + existing.outputs = {"result": 2} + existing.process_data = {"step": 3} + monkeypatch.setattr(repo, "_truncate_and_upload", lambda *_args, **_kwargs: None) + repo.save_execution_data(existing) + with sqlite_session_factory() as session: + persisted = session.get( + _sqlalchemy_workflow_node_execution_repository_WorkflowNodeExecutionModel, existing.id + ) + assert persisted is not None + assert persisted.inputs_dict == {"updated": True} + assert persisted.outputs_dict == {"result": 2} + assert persisted.process_data_dict == { + "step": 3, + "workflow_agent_binding_id": "binding-1", + "workflow_tool_invocation_id": "tool-call-1", + "workflow_tool_root_app_id": "root-app", + } - -def test_to_db_model_requires_creator_context( - monkeypatch: pytest.MonkeyPatch, sqlite_session_factory: sessionmaker[Session] -) -> None: - repo = _repository(monkeypatch, sqlite_session_factory) - execution = _execution() - - monkeypatch.setattr(repo, "_creator_user_id", None) - with pytest.raises(ValueError, match="created_by is required"): - repo._to_db_model(execution) - - monkeypatch.setattr(repo, "_creator_user_id", "user-1") - monkeypatch.setattr(repo, "_creator_user_role", None) - with pytest.raises(ValueError, match="created_by_role is required"): - repo._to_db_model(execution) - - -def test_json_encode_uses_runtime_converter(monkeypatch: pytest.MonkeyPatch) -> None: - class Converter: - def to_json_encodable(self, values: Mapping[str, Any]) -> Mapping[str, Any]: - return {"wrapped": values["value"]} - - monkeypatch.setattr( - "core.repositories.sqlalchemy_workflow_node_execution_repository.WorkflowRuntimeTypeConverter", - Converter, - ) - - assert SQLAlchemyWorkflowNodeExecutionRepository._json_encode({"value": 1}) == '{"wrapped": 1}' - - -def test_save_inserts_and_updates_persisted_execution( - monkeypatch: pytest.MonkeyPatch, sqlite_session_factory: sessionmaker[Session] -) -> None: - repo = _repository(monkeypatch, sqlite_session_factory) - execution = _execution(inputs={"value": 1}, outputs={"result": "first"}) - repo.save(execution) - with sqlite_session_factory() as session: - persisted = session.get(WorkflowNodeExecutionModel, execution.id) - assert persisted is not None - assert persisted.outputs_dict == {"result": "first"} - execution.title = "Updated" - execution.outputs = {"result": "second"} - repo.save(execution) - with sqlite_session_factory() as session: - persisted = session.get(WorkflowNodeExecutionModel, execution.id) - assert persisted is not None - assert persisted.title == "Updated" - assert persisted.outputs_dict == {"result": "second"} - assert execution.node_execution_id is not None - assert repo._node_execution_cache[execution.node_execution_id].id == execution.id - - -def test_save_owned_session_rolls_back_failed_insert( - monkeypatch: pytest.MonkeyPatch, - sqlite_engine: Engine, - sqlite_session_factory: sessionmaker[Session], -) -> None: - repo = _repository(monkeypatch, sqlite_session_factory) - with _raise_on_execution_insert(sqlite_engine), pytest.raises(RuntimeError, match="forced execution INSERT"): - repo.save(_execution()) - with sqlite_session_factory() as session: - assert session.scalar(select(WorkflowNodeExecutionModel)) is None - - -def test_save_execution_data_updates_existing_and_creates_missing( - monkeypatch: pytest.MonkeyPatch, sqlite_session_factory: sessionmaker[Session] -) -> None: - repo = _repository(monkeypatch, sqlite_session_factory) - existing = _execution( - inputs={"initial": True}, - process_data={ - "workflow_agent_binding_id": "binding-1", - "workflow_tool_invocation_id": "tool-call-1", - "workflow_tool_root_app_id": "root-app", - }, - ) - repo.save(existing) - existing.inputs = {"updated": True} - existing.outputs = {"result": 2} - existing.process_data = {"step": 3} - monkeypatch.setattr(repo, "_truncate_and_upload", lambda *_args, **_kwargs: None) - repo.save_execution_data(existing) - with sqlite_session_factory() as session: - persisted = session.get(WorkflowNodeExecutionModel, existing.id) - assert persisted is not None - assert persisted.inputs_dict == {"updated": True} - assert persisted.outputs_dict == {"result": 2} - assert persisted.process_data_dict == { - "step": 3, - "workflow_agent_binding_id": "binding-1", - "workflow_tool_invocation_id": "tool-call-1", - "workflow_tool_root_app_id": "root-app", - } - - missing = _execution(execution_id="missing", node_execution_id="missing-node", inputs={"new": True}) - repo.save_execution_data(missing) - with sqlite_session_factory() as session: - persisted = session.get(WorkflowNodeExecutionModel, missing.id) - assert persisted is not None - assert persisted.inputs_dict == {"new": True} - - -@pytest.mark.parametrize( - ("execution_factory", "offload_type", "read_persisted", "read_truncated"), - [ - ( - lambda: _execution(inputs={"large": "value"}), - ExecutionOffLoadType.INPUTS, - lambda model: model.inputs_dict, - lambda execution: execution.get_truncated_inputs(), - ), - ( - lambda: _execution(outputs={"large": "value"}), - ExecutionOffLoadType.OUTPUTS, - lambda model: model.outputs_dict, - lambda execution: execution.get_truncated_outputs(), - ), - ( - lambda: _execution(process_data={"large": "value"}), - ExecutionOffLoadType.PROCESS_DATA, - lambda model: model.process_data_dict, - lambda execution: execution.get_truncated_process_data(), - ), - ], -) -def test_save_execution_data_persists_each_truncation_offload( - monkeypatch: pytest.MonkeyPatch, - sqlite_session_factory: sessionmaker[Session], - execution_factory: Callable[[], WorkflowNodeExecution], - offload_type: ExecutionOffLoadType, - read_persisted: Callable[[WorkflowNodeExecutionModel], Mapping[str, Any] | None], - read_truncated: Callable[[WorkflowNodeExecution], Mapping[str, Any] | None], -) -> None: - repo = _repository(monkeypatch, sqlite_session_factory) - execution = execution_factory() - repo.save(execution) - offload = WorkflowNodeExecutionOffload( - tenant_id="tenant-1", - app_id="app-1", - node_execution_id=execution.id, - type_=offload_type, - file_id="file-1", - ) - result = SimpleNamespace(truncated_value={"large": "truncated"}, offload=offload) - monkeypatch.setattr(repo, "_truncate_and_upload", lambda values, *_args: result if values else None) - repo.save_execution_data(execution) - with sqlite_session_factory() as session: - persisted = session.get(WorkflowNodeExecutionModel, execution.id) - assert persisted is not None - assert read_persisted(persisted) == {"large": "truncated"} - offloads = session.scalars( - select(WorkflowNodeExecutionOffload).where(WorkflowNodeExecutionOffload.node_execution_id == execution.id) - ).all() - assert [item.type_ for item in offloads] == [offload_type] - assert read_truncated(execution) == {"large": "truncated"} - - -def test_get_by_workflow_run_filters_tenant_app_trigger_and_paused_and_orders( - monkeypatch: pytest.MonkeyPatch, sqlite_session_factory: sessionmaker[Session] -) -> None: - repo = _repository(monkeypatch, sqlite_session_factory) - repo.save(_execution(execution_id="two", node_execution_id="node-two", index=2)) - repo.save(_execution(execution_id="one", node_execution_id="node-one", index=1)) - repo.save( - _execution( - execution_id="paused", - node_execution_id="node-paused", - index=3, - status=WorkflowNodeExecutionStatus.PAUSED, + missing = _sqlalchemy_workflow_node_execution_repository_execution( + execution_id="missing", node_execution_id="missing-node", inputs={"new": True} ) - ) - _repository(monkeypatch, sqlite_session_factory, tenant_id="tenant-2").save( - _execution(execution_id="foreign-tenant", node_execution_id="foreign-tenant") - ) - _repository(monkeypatch, sqlite_session_factory, app_id="app-2").save( - _execution(execution_id="foreign-app", node_execution_id="foreign-app") - ) - _repository( - monkeypatch, - sqlite_session_factory, - triggered_from=WorkflowNodeExecutionTriggeredFrom.SINGLE_STEP, - ).save(_execution(execution_id="single-step", node_execution_id="single-step")) + repo.save_execution_data(missing) + with sqlite_session_factory() as session: + persisted = session.get( + _sqlalchemy_workflow_node_execution_repository_WorkflowNodeExecutionModel, missing.id + ) + assert persisted is not None + assert persisted.inputs_dict == {"new": True} - models = repo.get_db_models_by_workflow_run( - "run-1", - OrderConfig(order_by=["missing", "index"], order_direction="desc"), + @staticmethod + @pytest.mark.parametrize( + ("execution_factory", "offload_type", "read_persisted", "read_truncated"), + [ + ( + lambda: _sqlalchemy_workflow_node_execution_repository_execution(inputs={"large": "value"}), + ExecutionOffLoadType.INPUTS, + lambda model: model.inputs_dict, + lambda execution: execution.get_truncated_inputs(), + ), + ( + lambda: _sqlalchemy_workflow_node_execution_repository_execution(outputs={"large": "value"}), + ExecutionOffLoadType.OUTPUTS, + lambda model: model.outputs_dict, + lambda execution: execution.get_truncated_outputs(), + ), + ( + lambda: _sqlalchemy_workflow_node_execution_repository_execution(process_data={"large": "value"}), + ExecutionOffLoadType.PROCESS_DATA, + lambda model: model.process_data_dict, + lambda execution: execution.get_truncated_process_data(), + ), + ], ) - assert [model.id for model in models] == ["two", "one"] - assert set(repo._node_execution_cache) >= {"node-one", "node-two"} - assert repo.get_db_models_by_workflow_run("missing-run") == [] - no_app_repo = _repository(monkeypatch, sqlite_session_factory, app_id=None) - assert ( - no_app_repo.get_db_models_by_workflow_run( - "missing-run", - OrderConfig(order_by=["missing"], order_direction="asc"), + def test_save_execution_data_persists_each_truncation_offload( + monkeypatch: pytest.MonkeyPatch, + sqlite_session_factory: sessionmaker[Session], + execution_factory: Callable[[], _sqlalchemy_workflow_node_execution_repository_WorkflowNodeExecution], + offload_type: ExecutionOffLoadType, + read_persisted: Callable[ + [_sqlalchemy_workflow_node_execution_repository_WorkflowNodeExecutionModel], Mapping[str, Any] | None + ], + read_truncated: Callable[ + [_sqlalchemy_workflow_node_execution_repository_WorkflowNodeExecution], Mapping[str, Any] | None + ], + ) -> None: + repo = _repository(monkeypatch, sqlite_session_factory) + execution = execution_factory() + repo.save(execution) + offload = WorkflowNodeExecutionOffload( + tenant_id="tenant-1", + app_id="app-1", + node_execution_id=execution.id, + type_=offload_type, + file_id="file-1", ) - == [] - ) + result = SimpleNamespace(truncated_value={"large": "truncated"}, offload=offload) + monkeypatch.setattr(repo, "_truncate_and_upload", lambda values, *_args: result if values else None) + repo.save_execution_data(execution) + with sqlite_session_factory() as session: + persisted = session.get( + _sqlalchemy_workflow_node_execution_repository_WorkflowNodeExecutionModel, execution.id + ) + assert persisted is not None + assert read_persisted(persisted) == {"large": "truncated"} + offloads = session.scalars( + select(WorkflowNodeExecutionOffload).where( + WorkflowNodeExecutionOffload.node_execution_id == execution.id + ) + ).all() + assert [item.type_ for item in offloads] == [offload_type] + assert read_truncated(execution) == {"large": "truncated"} - -def test_get_by_workflow_execution_maps_real_rows_to_domain( - monkeypatch: pytest.MonkeyPatch, sqlite_session_factory: sessionmaker[Session] -) -> None: - repo = _repository(monkeypatch, sqlite_session_factory) - repo.save(_execution(inputs={"input": 1}, outputs={"output": 2})) - domains = repo.get_by_workflow_execution("run-1", OrderConfig(order_by=["index"], order_direction="asc")) - assert len(domains) == 1 - assert domains[0].inputs == {"input": 1} - assert domains[0].outputs == {"output": 2} - - -@pytest.mark.parametrize( - "repository_type", [SQLAlchemyWorkflowNodeExecutionRepository, CeleryWorkflowNodeExecutionRepository] -) -def test_resume_history_includes_paused_nodes_without_changing_default_reads( - sqlite_session_factory: sessionmaker[Session], - repository_type: type[SQLAlchemyWorkflowNodeExecutionRepository | CeleryWorkflowNodeExecutionRepository], -) -> None: - user = _account() - writer = SQLAlchemyWorkflowNodeExecutionRepository( - sqlite_session_factory, - tenant_id="tenant-1", - user=user, - app_id="app-1", - triggered_from=WorkflowNodeExecutionTriggeredFrom.WORKFLOW_RUN, - ) - writer.save_synchronously(_execution(execution_id="finished", node_execution_id="finished")) - paused = _execution(execution_id="paused", node_execution_id="paused", status=WorkflowNodeExecutionStatus.PAUSED) - writer.save_synchronously(paused) - writer.for_workflow_tool("app-1").save_synchronously( - paused.model_copy(update={"id": "source-paused", "node_execution_id": "source-paused"}) - ) - reader = repository_type( - sqlite_session_factory, - tenant_id="tenant-1", - user=user, - app_id="app-1", - triggered_from=WorkflowNodeExecutionTriggeredFrom.WORKFLOW_RUN, - ) - source_reader = reader.for_workflow_tool("app-1") - assert {node.id for node in reader.get_by_workflow_execution("run-1")} == {"finished"} - assert source_reader.get_by_workflow_execution("run-1") == [] - assert {node.id for node in reader.get_by_workflow_execution("run-1", include_paused=True)} == { - "finished", - "paused", - } - assert {node.id for node in source_reader.get_by_workflow_execution("run-1", include_paused=True)} == { - "source-paused" - } - assert {node.id for node in reader.get_by_workflow_execution("run-1")} == {"finished"} - assert source_reader.get_by_workflow_execution("run-1") == [] - - -def test_run_reads_keep_workflow_tool_internals_in_their_source_scope( - monkeypatch: pytest.MonkeyPatch, sqlite_session_factory: sessionmaker[Session] -) -> None: - caller = _repository(monkeypatch, sqlite_session_factory) - caller.save(_execution(execution_id="root", node_execution_id="root")) - for execution_id, tenant_id, app_id, root_app_id, run_id, origin in ( - ("child", "tenant-1", "source-app", "app-1", "run-1", WorkflowNodeExecutionTriggeredFrom.WORKFLOW_TOOL), - ("same-app", "tenant-1", "app-1", "app-1", "run-1", WorkflowNodeExecutionTriggeredFrom.WORKFLOW_TOOL), - ("foreign-root", "tenant-1", "source-app", "app-2", "run-1", WorkflowNodeExecutionTriggeredFrom.WORKFLOW_TOOL), - ( - "foreign-tenant", - "tenant-2", - "source-app", - "app-1", - "run-1", - WorkflowNodeExecutionTriggeredFrom.WORKFLOW_TOOL, - ), - ("foreign-run", "tenant-1", "source-app", "app-1", "run-2", WorkflowNodeExecutionTriggeredFrom.WORKFLOW_TOOL), - ("foreign-app", "tenant-1", "source-app", "app-1", "run-1", WorkflowNodeExecutionTriggeredFrom.WORKFLOW_RUN), - ): + @staticmethod + def test_get_by_workflow_run_filters_tenant_app_trigger_and_paused_and_orders( + monkeypatch: pytest.MonkeyPatch, sqlite_session_factory: sessionmaker[Session] + ) -> None: + repo = _repository(monkeypatch, sqlite_session_factory) + repo.save( + _sqlalchemy_workflow_node_execution_repository_execution( + execution_id="two", node_execution_id="node-two", index=2 + ) + ) + repo.save( + _sqlalchemy_workflow_node_execution_repository_execution( + execution_id="one", node_execution_id="node-one", index=1 + ) + ) + repo.save( + _sqlalchemy_workflow_node_execution_repository_execution( + execution_id="paused", + node_execution_id="node-paused", + index=3, + status=_RepositoryExecutionStatus.PAUSED, + ) + ) + _repository(monkeypatch, sqlite_session_factory, tenant_id="tenant-2").save( + _sqlalchemy_workflow_node_execution_repository_execution( + execution_id="foreign-tenant", node_execution_id="foreign-tenant" + ) + ) + _repository(monkeypatch, sqlite_session_factory, app_id="app-2").save( + _sqlalchemy_workflow_node_execution_repository_execution( + execution_id="foreign-app", node_execution_id="foreign-app" + ) + ) _repository( - monkeypatch, sqlite_session_factory, tenant_id=tenant_id, app_id=app_id, triggered_from=origin + monkeypatch, + sqlite_session_factory, + triggered_from=_RepositoryTriggeredFrom.SINGLE_STEP, ).save( - _execution( - execution_id=execution_id, - node_execution_id=execution_id, - run_id=run_id, - process_data={WORKFLOW_TOOL_ROOT_APP_ID_KEY: root_app_id}, + _sqlalchemy_workflow_node_execution_repository_execution( + execution_id="single-step", node_execution_id="single-step" ) ) - assert [node.id for node in caller.get_by_workflow_execution("run-1")] == ["root"] - source = caller.for_workflow_tool("source-app") - assert {node.id for node in source.get_by_workflow_execution("run-1")} == {"child", "foreign-root"} - -def test_to_domain_model_loads_offloaded_storage( - monkeypatch: pytest.MonkeyPatch, sqlite_session_factory: sessionmaker[Session] -) -> None: - repo = _repository(monkeypatch, sqlite_session_factory) - db_model = repo._to_db_model( - _execution( - inputs={"truncated": "inputs"}, - outputs={"truncated": "outputs"}, - process_data={"truncated": "process_data"}, + models = repo.get_db_models_by_workflow_run( + "run-1", + OrderConfig(order_by=["missing", "index"], order_direction="desc"), ) + assert [model.id for model in models] == ["two", "one"] + assert set(repo._node_execution_cache) >= {"node-one", "node-two"} + assert repo.get_db_models_by_workflow_run("missing-run") == [] + no_app_repo = _repository(monkeypatch, sqlite_session_factory, app_id=None) + assert ( + no_app_repo.get_db_models_by_workflow_run( + "missing-run", + OrderConfig(order_by=["missing"], order_direction="asc"), + ) + == [] + ) + + @staticmethod + def test_get_by_workflow_execution_maps_real_rows_to_domain( + monkeypatch: pytest.MonkeyPatch, sqlite_session_factory: sessionmaker[Session] + ) -> None: + repo = _repository(monkeypatch, sqlite_session_factory) + repo.save(_sqlalchemy_workflow_node_execution_repository_execution(inputs={"input": 1}, outputs={"output": 2})) + domains = repo.get_by_workflow_execution("run-1", OrderConfig(order_by=["index"], order_direction="asc")) + assert len(domains) == 1 + assert domains[0].inputs == {"input": 1} + assert domains[0].outputs == {"output": 2} + + @staticmethod + @pytest.mark.parametrize( + "repository_type", [SQLAlchemyWorkflowNodeExecutionRepository, CeleryWorkflowNodeExecutionRepository] ) - offloads = [] - for offload_type in ExecutionOffLoadType: - offload = _offload(offload_type) - offload.file = _upload_file(key=offload_type.value) - offloads.append(offload) - db_model.offload_data = offloads - monkeypatch.setattr( - "core.repositories.sqlalchemy_workflow_node_execution_repository.storage.load", - lambda key: json.dumps({"full": key}).encode(), + def test_resume_history_includes_paused_nodes_without_changing_default_reads( + sqlite_session_factory: sessionmaker[Session], + repository_type: type[SQLAlchemyWorkflowNodeExecutionRepository | CeleryWorkflowNodeExecutionRepository], + ) -> None: + user = _sqlalchemy_workflow_node_execution_repository_account() + writer = SQLAlchemyWorkflowNodeExecutionRepository( + sqlite_session_factory, + tenant_id="tenant-1", + user=user, + app_id="app-1", + triggered_from=_RepositoryTriggeredFrom.WORKFLOW_RUN, + ) + writer.save_synchronously( + _sqlalchemy_workflow_node_execution_repository_execution( + execution_id="finished", node_execution_id="finished" + ) + ) + paused = _sqlalchemy_workflow_node_execution_repository_execution( + execution_id="paused", + node_execution_id="paused", + status=_RepositoryExecutionStatus.PAUSED, + ) + writer.save_synchronously(paused) + writer.for_workflow_tool("app-1").save_synchronously( + paused.model_copy(update={"id": "source-paused", "node_execution_id": "source-paused"}) + ) + reader = repository_type( + sqlite_session_factory, + tenant_id="tenant-1", + user=user, + app_id="app-1", + triggered_from=_RepositoryTriggeredFrom.WORKFLOW_RUN, + ) + source_reader = reader.for_workflow_tool("app-1") + assert {node.id for node in reader.get_by_workflow_execution("run-1")} == {"finished"} + assert source_reader.get_by_workflow_execution("run-1") == [] + assert {node.id for node in reader.get_by_workflow_execution("run-1", include_paused=True)} == { + "finished", + "paused", + } + assert {node.id for node in source_reader.get_by_workflow_execution("run-1", include_paused=True)} == { + "source-paused" + } + assert {node.id for node in reader.get_by_workflow_execution("run-1")} == {"finished"} + assert source_reader.get_by_workflow_execution("run-1") == [] + + @staticmethod + def test_run_reads_keep_workflow_tool_internals_in_their_source_scope( + monkeypatch: pytest.MonkeyPatch, sqlite_session_factory: sessionmaker[Session] + ) -> None: + caller = _repository(monkeypatch, sqlite_session_factory) + caller.save( + _sqlalchemy_workflow_node_execution_repository_execution(execution_id="root", node_execution_id="root") + ) + for execution_id, tenant_id, app_id, root_app_id, run_id, origin in ( + ( + "child", + "tenant-1", + "source-app", + "app-1", + "run-1", + _RepositoryTriggeredFrom.WORKFLOW_TOOL, + ), + ( + "same-app", + "tenant-1", + "app-1", + "app-1", + "run-1", + _RepositoryTriggeredFrom.WORKFLOW_TOOL, + ), + ( + "foreign-root", + "tenant-1", + "source-app", + "app-2", + "run-1", + _RepositoryTriggeredFrom.WORKFLOW_TOOL, + ), + ( + "foreign-tenant", + "tenant-2", + "source-app", + "app-1", + "run-1", + _RepositoryTriggeredFrom.WORKFLOW_TOOL, + ), + ( + "foreign-run", + "tenant-1", + "source-app", + "app-1", + "run-2", + _RepositoryTriggeredFrom.WORKFLOW_TOOL, + ), + ( + "foreign-app", + "tenant-1", + "source-app", + "app-1", + "run-1", + _RepositoryTriggeredFrom.WORKFLOW_RUN, + ), + ): + _repository( + monkeypatch, sqlite_session_factory, tenant_id=tenant_id, app_id=app_id, triggered_from=origin + ).save( + _sqlalchemy_workflow_node_execution_repository_execution( + execution_id=execution_id, + node_execution_id=execution_id, + run_id=run_id, + process_data={WORKFLOW_TOOL_ROOT_APP_ID_KEY: root_app_id}, + ) + ) + assert [node.id for node in caller.get_by_workflow_execution("run-1")] == ["root"] + source = caller.for_workflow_tool("source-app") + assert {node.id for node in source.get_by_workflow_execution("run-1")} == {"child", "foreign-root"} + + @staticmethod + def test_to_domain_model_loads_offloaded_storage( + monkeypatch: pytest.MonkeyPatch, sqlite_session_factory: sessionmaker[Session] + ) -> None: + repo = _repository(monkeypatch, sqlite_session_factory) + db_model = repo._to_db_model( + _sqlalchemy_workflow_node_execution_repository_execution( + inputs={"truncated": "inputs"}, + outputs={"truncated": "outputs"}, + process_data={"truncated": "process_data"}, + ) + ) + offloads = [] + for offload_type in ExecutionOffLoadType: + offload = _offload(offload_type) + offload.file = _upload_file(key=offload_type.value) + offloads.append(offload) + db_model.offload_data = offloads + monkeypatch.setattr( + "core.repositories.sqlalchemy_workflow_node_execution_repository.storage.load", + lambda key: json.dumps({"full": key}).encode(), + ) + domain = repo._to_domain_model(db_model) + assert domain.inputs == {"full": "inputs"} + assert domain.outputs == {"full": "outputs"} + assert domain.process_data == {"full": "process_data"} + assert domain.get_truncated_inputs() == {"truncated": "inputs"} + assert domain.get_truncated_outputs() == {"truncated": "outputs"} + assert domain.get_truncated_process_data() == {"truncated": "process_data"} + + @staticmethod + def test_truncate_and_upload_keeps_file_boundary_mocked( + monkeypatch: pytest.MonkeyPatch, sqlite_session_factory: sessionmaker[Session] + ) -> None: + uploaded = _upload_file(key="file-key") + uploaded.id = "file-1" + repo = _repository(monkeypatch, sqlite_session_factory) + upload_file = Mock(return_value=uploaded) + monkeypatch.setattr(repo._file_service, "upload_file", upload_file) + + class Truncator: + def truncate_variable_mapping(self, _value: Any) -> tuple[dict[str, bool], bool]: + return {"truncated": True}, True + + monkeypatch.setattr(repo, "_create_truncator", lambda: Truncator()) + result = repo._truncate_and_upload({"value": 1}, "execution-1", ExecutionOffLoadType.INPUTS) + assert result is not None + assert result.truncated_value == {"truncated": True} + upload_file.assert_called_once_with( + filename="node_execution_execution-1_inputs.json", + content=b'{"value": 1}', + mimetype="application/json", + user=repo._user, + tenant_id="tenant-1", + ) + assert result.offload.file_id == "file-1" + assert result.offload.type_ == ExecutionOffLoadType.INPUTS + assert result.offload.tenant_id == "tenant-1" + assert result.offload.app_id == "app-1" + assert result.offload.node_execution_id == "execution-1" + + @staticmethod + def test_truncate_and_upload_returns_none_for_missing_or_small_values( + monkeypatch: pytest.MonkeyPatch, sqlite_session_factory: sessionmaker[Session] + ) -> None: + repo = _repository(monkeypatch, sqlite_session_factory) + assert repo._truncate_and_upload(None, "execution-1", ExecutionOffLoadType.INPUTS) is None + + class Truncator: + def truncate_variable_mapping(self, value: Mapping[str, Any]) -> tuple[Mapping[str, Any], bool]: + return value, False + + monkeypatch.setattr(repo, "_create_truncator", lambda: Truncator()) + assert repo._truncate_and_upload({"value": 1}, "execution-1", ExecutionOffLoadType.INPUTS) is None + + @staticmethod + def test_duplicate_detection_and_id_regeneration( + monkeypatch: pytest.MonkeyPatch, + caplog: pytest.LogCaptureFixture, + sqlite_session_factory: sessionmaker[Session], + ) -> None: + repo = _repository(monkeypatch, sqlite_session_factory) + duplicate = IntegrityError("duplicate", params=None, orig=psycopg2.errors.UniqueViolation("duplicate key")) + assert repo._is_duplicate_key_error(duplicate) + assert not repo._is_duplicate_key_error(IntegrityError("other", params=None, orig=Exception("other"))) + execution = _sqlalchemy_workflow_node_execution_repository_execution(execution_id="old") + db_model = repo._to_db_model(execution) + monkeypatch.setattr("core.repositories.sqlalchemy_workflow_node_execution_repository.uuidv7", lambda: "new") + caplog.set_level(logging.WARNING) + repo._regenerate_id_on_duplicate(execution, db_model) + assert execution.id == db_model.id == "new" + assert "Duplicate key conflict" in caplog.text + + @staticmethod + def test_save_retries_postgres_duplicate_key( + monkeypatch: pytest.MonkeyPatch, + sqlite_session_factory: sessionmaker[Session], + ) -> None: + repo = _repository(monkeypatch, sqlite_session_factory) + execution = _sqlalchemy_workflow_node_execution_repository_execution(execution_id="old") + duplicate = IntegrityError( + "duplicate", + params=None, + orig=psycopg2.errors.UniqueViolation("duplicate key"), + ) + persist = Mock(side_effect=[duplicate, None]) + monkeypatch.setattr(repo, "_persist_to_database", persist) + monkeypatch.setattr("core.repositories.sqlalchemy_workflow_node_execution_repository.uuidv7", lambda: "new") + + repo.save(execution) + + assert persist.call_count == 2 + assert execution.id == "new" + assert execution.node_execution_id is not None + assert repo._node_execution_cache[execution.node_execution_id].id == "new" + + @staticmethod + def test_save_logs_and_reraises_non_duplicate_and_unexpected_errors( + monkeypatch: pytest.MonkeyPatch, + caplog: pytest.LogCaptureFixture, + sqlite_session_factory: sessionmaker[Session], + ) -> None: + repo = _repository(monkeypatch, sqlite_session_factory) + non_duplicate = IntegrityError("other", params=None, orig=Exception("constraint")) + monkeypatch.setattr(repo, "_persist_to_database", Mock(side_effect=non_duplicate)) + caplog.set_level(logging.ERROR) + + with pytest.raises(IntegrityError): + repo.save(_sqlalchemy_workflow_node_execution_repository_execution()) + assert "Non-duplicate key integrity error" in caplog.text + + caplog.clear() + monkeypatch.setattr(repo, "_persist_to_database", Mock(side_effect=RuntimeError("boom"))) + with pytest.raises(RuntimeError, match="boom"): + repo.save(_sqlalchemy_workflow_node_execution_repository_execution(execution_id="unexpected")) + assert "Failed to save workflow node execution" in caplog.text + + +"""Unit tests for workflow node execution conflict handling.""" + + +@dataclass(frozen=True) +class ConflictDatabase: + engine: Engine + session: Session + session_factory: sessionmaker[Session] + + +@dataclass(frozen=True) +class ConflictEvents: + insert_attempts: list[str] + rollbacks: list[bool] + + +def _workflow_node_execution_conflict_handling_account() -> Account: + return make_account( + account_id="test-user-id", + name="Conflict User", + email="conflict@example.com", + tenant=make_tenant(tenant_id="test-tenant-id", name="Conflict Tenant"), ) - domain = repo._to_domain_model(db_model) - assert domain.inputs == {"full": "inputs"} - assert domain.outputs == {"full": "outputs"} - assert domain.process_data == {"full": "process_data"} - assert domain.get_truncated_inputs() == {"truncated": "inputs"} - assert domain.get_truncated_outputs() == {"truncated": "outputs"} - assert domain.get_truncated_process_data() == {"truncated": "process_data"} -def test_truncate_and_upload_keeps_file_boundary_mocked( - monkeypatch: pytest.MonkeyPatch, sqlite_session_factory: sessionmaker[Session] -) -> None: - uploaded = _upload_file(key="file-key") - uploaded.id = "file-1" - repo = _repository(monkeypatch, sqlite_session_factory) - upload_file = Mock(return_value=uploaded) - monkeypatch.setattr(repo._file_service, "upload_file", upload_file) - - class Truncator: - def truncate_variable_mapping(self, _value: Any) -> tuple[dict[str, bool], bool]: - return {"truncated": True}, True - - monkeypatch.setattr(repo, "_create_truncator", lambda: Truncator()) - result = repo._truncate_and_upload({"value": 1}, "execution-1", ExecutionOffLoadType.INPUTS) - assert result is not None - assert result.truncated_value == {"truncated": True} - upload_file.assert_called_once_with( - filename="node_execution_execution-1_inputs.json", - content=b'{"value": 1}', - mimetype="application/json", - user=repo._user, - tenant_id="tenant-1", +def _workflow_node_execution_conflict_handling_execution( + *, + execution_id: str, + status: _ConflictExecutionStatus = _ConflictExecutionStatus.RUNNING, +) -> _workflow_node_execution_conflict_handling_WorkflowNodeExecution: + return _workflow_node_execution_conflict_handling_WorkflowNodeExecution( + id=execution_id, + workflow_id="test-workflow-id", + workflow_execution_id="test-workflow-execution-id", + node_execution_id="test-node-execution-id", + node_id="test-node-id", + node_type=BuiltinNodeTypes.START, + title="Test Node", + index=1, + status=status, + created_at=naive_utc_now(), ) - assert result.offload.file_id == "file-1" - assert result.offload.type_ == ExecutionOffLoadType.INPUTS - assert result.offload.tenant_id == "tenant-1" - assert result.offload.app_id == "app-1" - assert result.offload.node_execution_id == "execution-1" -def test_truncate_and_upload_returns_none_for_missing_or_small_values( - monkeypatch: pytest.MonkeyPatch, sqlite_session_factory: sessionmaker[Session] -) -> None: - repo = _repository(monkeypatch, sqlite_session_factory) - assert repo._truncate_and_upload(None, "execution-1", ExecutionOffLoadType.INPUTS) is None +@contextmanager +def _fail_inserts( + engine: Engine, + *, + failure_count: int, + duplicate: bool, +) -> Iterator[ConflictEvents]: + insert_attempts: list[str] = [] + rollbacks: list[bool] = [] - class Truncator: - def truncate_variable_mapping(self, value: Mapping[str, Any]) -> tuple[Mapping[str, Any], bool]: - return value, False + def fail_insert(_connection, _cursor, statement, parameters, _context, _executemany) -> None: + if not statement.lstrip().upper().startswith("INSERT INTO WORKFLOW_NODE_EXECUTIONS"): + return + insert_attempts.append(statement) + if len(insert_attempts) > failure_count: + return + original_error = ( + psycopg2.errors.UniqueViolation("forced duplicate key") + if duplicate + else RuntimeError("forced non-duplicate constraint failure") + ) + raise IntegrityError(statement, parameters, original_error) - monkeypatch.setattr(repo, "_create_truncator", lambda: Truncator()) - assert repo._truncate_and_upload({"value": 1}, "execution-1", ExecutionOffLoadType.INPUTS) is None + def record_rollback(_connection) -> None: + rollbacks.append(True) + + event.listen(engine, "before_cursor_execute", fail_insert) + event.listen(engine, "rollback", record_rollback) + try: + yield ConflictEvents(insert_attempts=insert_attempts, rollbacks=rollbacks) + finally: + event.remove(engine, "before_cursor_execute", fail_insert) + event.remove(engine, "rollback", record_rollback) -def test_duplicate_detection_and_id_regeneration( - monkeypatch: pytest.MonkeyPatch, - caplog: pytest.LogCaptureFixture, - sqlite_session_factory: sessionmaker[Session], -) -> None: - repo = _repository(monkeypatch, sqlite_session_factory) - duplicate = IntegrityError("duplicate", params=None, orig=psycopg2.errors.UniqueViolation("duplicate key")) - assert repo._is_duplicate_key_error(duplicate) - assert not repo._is_duplicate_key_error(IntegrityError("other", params=None, orig=Exception("other"))) - execution = _execution(execution_id="old") - db_model = repo._to_db_model(execution) - monkeypatch.setattr("core.repositories.sqlalchemy_workflow_node_execution_repository.uuidv7", lambda: "new") - caplog.set_level(logging.WARNING) - repo._regenerate_id_on_duplicate(execution, db_model) - assert execution.id == db_model.id == "new" - assert "Duplicate key conflict" in caplog.text +class TestWorkflowNodeExecutionConflictHandlingCases: + @staticmethod + @pytest.fixture + def conflict_database( + sqlite_session: Session, + sqlite_session_factory: sessionmaker[Session], + ) -> ConflictDatabase: + """Use the shared SQLite fixtures for repository-owned sessions and assertions.""" + engine = sqlite_session.get_bind() + assert isinstance(engine, Engine) + return ConflictDatabase( + engine=engine, + session=sqlite_session, + session_factory=sqlite_session_factory, + ) + + @staticmethod + @pytest.fixture + def repository(conflict_database: ConflictDatabase) -> SQLAlchemyWorkflowNodeExecutionRepository: + return SQLAlchemyWorkflowNodeExecutionRepository( + session_factory=conflict_database.session_factory, + tenant_id="test-tenant-id", + user=_workflow_node_execution_conflict_handling_account(), + app_id="test-app-id", + triggered_from=_workflow_node_execution_conflict_handling_WorkflowNodeExecutionTriggeredFrom.WORKFLOW_RUN, + ) + + class TestWorkflowNodeExecutionConflictHandling: + """Test cases for handling duplicate key conflicts in workflow node execution.""" + + def test_save_with_duplicate_key_retries_with_new_uuid( + self, + repository: SQLAlchemyWorkflowNodeExecutionRepository, + conflict_database: ConflictDatabase, + ) -> None: + execution = _workflow_node_execution_conflict_handling_execution(execution_id="original-id") + original_id = execution.id + + with _fail_inserts(conflict_database.engine, failure_count=1, duplicate=True) as conflicts: + repository.save(execution) + + assert len(conflicts.insert_attempts) == 2 + assert len(conflicts.rollbacks) == 1 + assert execution.id != original_id + assert ( + conflict_database.session.get( + _workflow_node_execution_conflict_handling_WorkflowNodeExecutionModel, original_id + ) + is None + ) + persisted = conflict_database.session.get( + _workflow_node_execution_conflict_handling_WorkflowNodeExecutionModel, execution.id + ) + assert persisted is not None + assert persisted.node_execution_id == execution.node_execution_id + + def test_save_with_existing_record_updates_instead_of_insert( + self, + repository: SQLAlchemyWorkflowNodeExecutionRepository, + conflict_database: ConflictDatabase, + ) -> None: + execution = _workflow_node_execution_conflict_handling_execution(execution_id="existing-id") + repository.save(execution) + + execution.status = _ConflictExecutionStatus.SUCCEEDED + repository.save(execution) + + conflict_database.session.expire_all() + persisted = conflict_database.session.get( + _workflow_node_execution_conflict_handling_WorkflowNodeExecutionModel, execution.id + ) + assert persisted is not None + assert persisted.status == _ConflictExecutionStatus.SUCCEEDED + assert ( + conflict_database.session.query( + _workflow_node_execution_conflict_handling_WorkflowNodeExecutionModel + ).count() + == 1 + ) + + def test_save_exceeds_max_retries_raises_error( + self, + repository: SQLAlchemyWorkflowNodeExecutionRepository, + conflict_database: ConflictDatabase, + ) -> None: + execution = _workflow_node_execution_conflict_handling_execution(execution_id="test-id") + + with _fail_inserts(conflict_database.engine, failure_count=3, duplicate=True) as conflicts: + with pytest.raises(IntegrityError): + repository.save(execution) + + assert len(conflicts.insert_attempts) == 3 + assert len(conflicts.rollbacks) == 3 + assert ( + conflict_database.session.query( + _workflow_node_execution_conflict_handling_WorkflowNodeExecutionModel + ).count() + == 0 + ) + + def test_save_non_duplicate_integrity_error_raises_immediately( + self, + repository: SQLAlchemyWorkflowNodeExecutionRepository, + conflict_database: ConflictDatabase, + ) -> None: + execution = _workflow_node_execution_conflict_handling_execution(execution_id="test-id") + + with _fail_inserts(conflict_database.engine, failure_count=1, duplicate=False) as conflicts: + with pytest.raises(IntegrityError): + repository.save(execution) + + assert len(conflicts.insert_attempts) == 1 + assert len(conflicts.rollbacks) == 1 + assert ( + conflict_database.session.get( + _workflow_node_execution_conflict_handling_WorkflowNodeExecutionModel, execution.id + ) + is None + ) -def test_save_retries_postgres_duplicate_key( - monkeypatch: pytest.MonkeyPatch, - sqlite_session_factory: sessionmaker[Session], -) -> None: - repo = _repository(monkeypatch, sqlite_session_factory) - execution = _execution(execution_id="old") - duplicate = IntegrityError( - "duplicate", - params=None, - orig=psycopg2.errors.UniqueViolation("duplicate key"), +""" +Unit tests for WorkflowNodeExecution truncation functionality. + +Tests the truncation and offloading logic for large inputs and outputs +in the SQLAlchemyWorkflowNodeExecutionRepository. +""" + + +@dataclass +class TruncationTestCase: + """Test case data for truncation scenarios.""" + + name: str + inputs: dict[str, Any] | None + outputs: dict[str, Any] | None + should_truncate_inputs: bool + should_truncate_outputs: bool + description: str + + +def create_test_cases() -> list[TruncationTestCase]: + """Create test cases for different truncation scenarios.""" + # Create large data that will definitely exceed the threshold (10KB) + large_data = {"data": "x" * (dify_config.WORKFLOW_VARIABLE_TRUNCATION_MAX_SIZE + 1000)} + small_data = {"data": "small"} + + return [ + TruncationTestCase( + name="small_data_no_truncation", + inputs=small_data, + outputs=small_data, + should_truncate_inputs=False, + should_truncate_outputs=False, + description="Small data should not be truncated", + ), + TruncationTestCase( + name="large_inputs_truncation", + inputs=large_data, + outputs=small_data, + should_truncate_inputs=True, + should_truncate_outputs=False, + description="Large inputs should be truncated", + ), + TruncationTestCase( + name="large_outputs_truncation", + inputs=small_data, + outputs=large_data, + should_truncate_inputs=False, + should_truncate_outputs=True, + description="Large outputs should be truncated", + ), + TruncationTestCase( + name="large_both_truncation", + inputs=large_data, + outputs=large_data, + should_truncate_inputs=True, + should_truncate_outputs=True, + description="Both large inputs and outputs should be truncated", + ), + TruncationTestCase( + name="none_inputs_outputs", + inputs=None, + outputs=None, + should_truncate_inputs=False, + should_truncate_outputs=False, + description="None inputs and outputs should not be truncated", + ), + ] + + +def create_workflow_node_execution( + execution_id: str = "test-execution-id", + inputs: dict[str, Any] | None = None, + outputs: dict[str, Any] | None = None, +) -> _workflow_node_execution_truncation_WorkflowNodeExecution: + """Factory function to create a WorkflowNodeExecution for testing.""" + return _workflow_node_execution_truncation_WorkflowNodeExecution( + id=execution_id, + node_execution_id="test-node-execution-id", + workflow_id="test-workflow-id", + workflow_execution_id="test-workflow-execution-id", + index=1, + node_id="test-node-id", + node_type=BuiltinNodeTypes.LLM, + title="Test Node", + inputs=inputs, + outputs=outputs, + status=_workflow_node_execution_truncation_WorkflowNodeExecutionStatus.SUCCEEDED, + created_at=datetime.now(UTC), ) - persist = Mock(side_effect=[duplicate, None]) - monkeypatch.setattr(repo, "_persist_to_database", persist) - monkeypatch.setattr("core.repositories.sqlalchemy_workflow_node_execution_repository.uuidv7", lambda: "new") - - repo.save(execution) - - assert persist.call_count == 2 - assert execution.id == "new" - assert execution.node_execution_id is not None - assert repo._node_execution_cache[execution.node_execution_id].id == "new" -def test_save_logs_and_reraises_non_duplicate_and_unexpected_errors( - monkeypatch: pytest.MonkeyPatch, - caplog: pytest.LogCaptureFixture, - sqlite_session_factory: sessionmaker[Session], -) -> None: - repo = _repository(monkeypatch, sqlite_session_factory) - non_duplicate = IntegrityError("other", params=None, orig=Exception("constraint")) - monkeypatch.setattr(repo, "_persist_to_database", Mock(side_effect=non_duplicate)) - caplog.set_level(logging.ERROR) +def mock_user() -> Account: + """Create an Account user for testing.""" - with pytest.raises(IntegrityError): - repo.save(_execution()) - assert "Non-duplicate key integrity error" in caplog.text + user = Account(name="Test Account", email="test@example.com") + user.id = "test-user-id" + user._current_tenant = Tenant(name="Test Tenant") + user._current_tenant.id = "test-tenant-id" + return user - caplog.clear() - monkeypatch.setattr(repo, "_persist_to_database", Mock(side_effect=RuntimeError("boom"))) - with pytest.raises(RuntimeError, match="boom"): - repo.save(_execution(execution_id="unexpected")) - assert "Failed to save workflow node execution" in caplog.text + +def create_offload(type_: ExecutionOffLoadType) -> WorkflowNodeExecutionOffload: + return WorkflowNodeExecutionOffload( + tenant_id="test-tenant-id", + app_id="test-app-id", + node_execution_id="test-node-execution-id", + type_=type_, + file_id="test-file-id", + ) + + +class TestWorkflowNodeExecutionTruncationCases: + class TestSQLAlchemyWorkflowNodeExecutionRepositoryTruncation: + """Test class for truncation functionality in SQLAlchemyWorkflowNodeExecutionRepository.""" + + def create_repository(self, sqlite_engine: Engine) -> SQLAlchemyWorkflowNodeExecutionRepository: + """Create a repository backed by the test's isolated SQLite engine.""" + repository = SQLAlchemyWorkflowNodeExecutionRepository( + session_factory=sqlite_engine, + tenant_id="test-tenant-id", + user=mock_user(), + app_id="test-app-id", + triggered_from=_workflow_node_execution_truncation_WorkflowNodeExecutionTriggeredFrom.WORKFLOW_RUN, + ) + with repository._session_factory() as session: + assert isinstance(session, Session) + assert session.get_bind() is sqlite_engine + return repository + + def test_to_domain_model_without_offload_data(self, sqlite_engine: Engine): + """Test _to_domain_model correctly handles models without offload data.""" + repo = self.create_repository(sqlite_engine) + + # Create a database model without offload data + db_model = _workflow_node_execution_truncation_WorkflowNodeExecutionModel( + id="test-id", + node_execution_id="node-exec-id", + workflow_id="workflow-id", + workflow_run_id="run-id", + index=1, + predecessor_node_id=None, + node_id="node-id", + node_type=BuiltinNodeTypes.LLM, + title="Test Node", + inputs=json.dumps({"value": "inputs"}), + process_data=json.dumps({"value": "process_data"}), + outputs=json.dumps({"value": "outputs"}), + status=_workflow_node_execution_truncation_WorkflowNodeExecutionStatus.SUCCEEDED, + error=None, + elapsed_time=1.0, + execution_metadata="{}", + created_at=datetime.now(UTC), + finished_at=None, + offload_data=[], + ) + + domain_model = repo._to_domain_model(db_model) + + # Check that no truncated data was set + assert domain_model.get_truncated_inputs() is None + assert domain_model.get_truncated_outputs() is None + + class TestWorkflowNodeExecutionModelTruncatedProperties: + """Test the truncated properties on WorkflowNodeExecutionModel.""" + + def test_inputs_truncated_with_offload_data(self): + """Test inputs_truncated property when offload data exists.""" + model = _workflow_node_execution_truncation_WorkflowNodeExecutionModel() + offload = create_offload(ExecutionOffLoadType.INPUTS) + model.offload_data = [offload] + + assert model.inputs_truncated is True + assert model.process_data_truncated is False + assert model.outputs_truncated is False + + def test_outputs_truncated_with_offload_data(self): + """Test outputs_truncated property when offload data exists.""" + model = _workflow_node_execution_truncation_WorkflowNodeExecutionModel() + + # Mock offload data with outputs file + offload = create_offload(ExecutionOffLoadType.OUTPUTS) + model.offload_data = [offload] + + assert model.inputs_truncated is False + assert model.process_data_truncated is False + assert model.outputs_truncated is True + + def test_process_data_truncated_with_offload_data(self): + model = _workflow_node_execution_truncation_WorkflowNodeExecutionModel() + offload = create_offload(ExecutionOffLoadType.PROCESS_DATA) + model.offload_data = [offload] + assert model.process_data_truncated is True + assert model.inputs_truncated is False + assert model.outputs_truncated is False + + def test_truncated_properties_without_offload_data(self): + """Test truncated properties when no offload data exists.""" + model = _workflow_node_execution_truncation_WorkflowNodeExecutionModel( + offload_data=[], + ) + + assert model.inputs_truncated is False + assert model.outputs_truncated is False + assert model.process_data_truncated is False + + def test_truncated_properties_without_offload_attribute(self): + """Test truncated properties when offload_data attribute doesn't exist.""" + model = _workflow_node_execution_truncation_WorkflowNodeExecutionModel() + # Don't set offload_data attribute at all + + assert model.inputs_truncated is False + assert model.outputs_truncated is False + assert model.process_data_truncated is False diff --git a/api/tests/unit_tests/core/repositories/test_workflow_control_persistence.py b/api/tests/unit_tests/core/repositories/test_workflow_control_persistence.py deleted file mode 100644 index 938ad985198..00000000000 --- a/api/tests/unit_tests/core/repositories/test_workflow_control_persistence.py +++ /dev/null @@ -1,336 +0,0 @@ -"""Configured history backends retain synchronous, monotonic task control.""" - -import pytest -from pytest_mock import MockerFixture -from sqlalchemy.orm import Session, sessionmaker - -from core.repositories.celery_workflow_execution_repository import CeleryWorkflowExecutionRepository -from core.repositories.sqlalchemy_workflow_execution_repository import SQLAlchemyWorkflowExecutionRepository -from core.repositories.sqlalchemy_workflow_node_execution_repository import SQLAlchemyWorkflowNodeExecutionRepository -from extensions.logstore.aliyun_logstore import AliyunLogStore -from extensions.logstore.repositories.logstore_api_workflow_node_execution_repository import ( - LogstoreAPIWorkflowNodeExecutionRepository, -) -from extensions.logstore.repositories.logstore_api_workflow_run_repository import LogstoreAPIWorkflowRunRepository -from extensions.logstore.repositories.logstore_workflow_execution_repository import LogstoreWorkflowExecutionRepository -from extensions.logstore.repositories.logstore_workflow_node_execution_repository import ( - LogstoreWorkflowNodeExecutionRepository, -) -from graphon.entities import WorkflowNodeExecution -from graphon.enums import WorkflowExecutionStatus, WorkflowNodeExecutionStatus, WorkflowType -from libs.datetime_utils import naive_utc_now -from models import Account, WorkflowRun -from models.enums import CreatorUserRole, WorkflowRunTriggeredFrom -from models.workflow import WorkflowNodeExecutionTriggeredFrom -from repositories.entities.workflow_execution import WorkflowExecution -from repositories.sqlalchemy_api_workflow_run_repository import DifyAPISQLAlchemyWorkflowRunRepository -from repositories.workflow_run_control import apply_workflow_stop_to_node -from tasks.workflow_execution_tasks import _update_workflow_run_from_execution - - -@pytest.fixture -def logstore(monkeypatch: pytest.MonkeyPatch, mocker: MockerFixture) -> AliyunLogStore: - monkeypatch.setattr(AliyunLogStore, "_instance", None) - monkeypatch.setattr(AliyunLogStore, "_initialized", False) - monkeypatch.setenv("ALIYUN_SLS_ENDPOINT", "https://example.com") - mocker.patch.object(AliyunLogStore, "_check_endpoint_connectivity") - store = AliyunLogStore() - store._use_pg_protocol = True - mocker.patch.object(store, "put_log") - return store - - -@pytest.mark.parametrize("backend", ["sql", "celery", "logstore"]) -def test_task_owner_is_durable_before_start_and_delayed_writes_cannot_revive_stop( - sqlite_session_factory: sessionmaker[Session], logstore: AliyunLogStore, mocker: MockerFixture, backend: str -) -> None: - queued_save = mocker.patch( - "core.repositories.celery_workflow_execution_repository.save_workflow_execution_task.delay" - ) - get_logs = mocker.patch.object(logstore, "get_logs") - repository_class = { - "sql": SQLAlchemyWorkflowExecutionRepository, - "celery": CeleryWorkflowExecutionRepository, - "logstore": LogstoreWorkflowExecutionRepository, - }[backend] - user = Account(name="Test", email="test@example.com") - user.id = "account-1" - repository = repository_class( - session_factory=sqlite_session_factory, - tenant_id="tenant-1", - user=user, - app_id="app-1", - triggered_from=WorkflowRunTriggeredFrom.APP_RUN, - ) - if isinstance(repository, LogstoreWorkflowExecutionRepository): - repository._enable_dual_write = False - execution = WorkflowExecution( - id_="run-1", - workflow_id="workflow-1", - workflow_type=WorkflowType.WORKFLOW, - workflow_version="1", - graph={"nodes": [], "edges": []}, - inputs={}, - started_at=naive_utc_now(), - ) - repository.save_synchronously(execution) - queued_save.assert_not_called() - control = DifyAPISQLAlchemyWorkflowRunRepository(sqlite_session_factory) - control.bind_workflow_task(tenant_id="tenant-1", app_id="app-1", workflow_run_id="run-1", task_id="task-1") - with sqlite_session_factory() as session: - run = session.get(WorkflowRun, "run-1") - assert run is not None - assert run.task_id == "task-1" - assert (run.created_by_role, run.created_by) == (CreatorUserRole.ACCOUNT, "account-1") - execution.status = WorkflowExecutionStatus.PAUSED - repository.save_synchronously(execution) - stopped = control.stop_paused_workflow_task(tenant_id="tenant-1", app_id="app-1", task_id="task-1", owner=None) - assert stopped is not None - assert stopped.status == WorkflowExecutionStatus.STOPPED - - # A writer already queued or running before the stop cannot restore RUNNING. - execution.status = WorkflowExecutionStatus.RUNNING - repository.save_synchronously(execution) - with sqlite_session_factory.begin() as session: - run = session.get(WorkflowRun, "run-1") - assert run is not None - _update_workflow_run_from_execution(run, execution) - with sqlite_session_factory() as session: - run = session.get(WorkflowRun, "run-1") - assert run is not None - assert run.status == WorkflowExecutionStatus.STOPPED - assert run.stop_requested_at is not None - if backend == "logstore": - # LogStore implements the configured run-control subset of the history protocol. - api = LogstoreAPIWorkflowRunRepository(sqlite_session_factory) # pyrefly: ignore[bad-instantiation] - api._enable_dual_read = False - found = api.get_workflow_run_by_id("tenant-1", "app-1", "run-1") - assert found is not None - assert found.status == WorkflowExecutionStatus.STOPPED - unscoped = api.get_workflow_run_by_id_without_tenant("run-1") - assert unscoped is not None - assert unscoped.status == WorkflowExecutionStatus.STOPPED - get_logs.assert_not_called() - - -@pytest.mark.usefixtures("logstore") -def test_logstore_pause_lifecycle_uses_sql_control_with_dual_write_disabled( - sqlite_session_factory: sessionmaker[Session], mocker: MockerFixture -) -> None: - mocker.patch("repositories.sqlalchemy_api_workflow_run_repository.storage.save") - mocker.patch("repositories.sqlalchemy_api_workflow_run_repository.storage.delete") - user = Account(name="Test", email="test@example.com") - user.id = "account-1" - sql = SQLAlchemyWorkflowExecutionRepository( - session_factory=sqlite_session_factory, - tenant_id="tenant-1", - user=user, - app_id="app-1", - triggered_from=WorkflowRunTriggeredFrom.APP_RUN, - ) - sql.save_synchronously( - WorkflowExecution( - id_="run-1", - workflow_id="workflow-1", - workflow_type=WorkflowType.WORKFLOW, - workflow_version="1", - graph={}, - inputs={}, - started_at=naive_utc_now(), - ) - ) - api = LogstoreAPIWorkflowRunRepository(sqlite_session_factory) # pyrefly: ignore[bad-instantiation] - api.bind_workflow_task(tenant_id="tenant-1", app_id="app-1", workflow_run_id="run-1", task_id="task-1") - pause = api.create_workflow_pause("run-1", "account-1", "{}", []) - persisted_pause = api.get_workflow_pause("run-1") - assert persisted_pause is not None - assert persisted_pause.id == pause.id - resume_callbacks: list[None] = [] - resumed = api.resume_workflow_pause("run-1", pause, before_resume=lambda: resume_callbacks.append(None)) - assert resume_callbacks == [None] - api.delete_workflow_pause(resumed) - assert api.get_workflow_pause("run-1") is None - - -@pytest.mark.parametrize( - ("record_type", "node_status"), - [ - ("workflow", WorkflowNodeExecutionStatus.PAUSED), - ("node", WorkflowNodeExecutionStatus.PAUSED), - ("node", WorkflowNodeExecutionStatus.RETRY), - ], -) -def test_logstore_stop_orders_before_delayed_upload_and_overrides_stale_indexes( - sqlite_session_factory: sessionmaker[Session], - logstore: AliyunLogStore, - mocker: MockerFixture, - record_type: str, - node_status: WorkflowNodeExecutionStatus, -) -> None: - user = Account(name="Test", email="test@example.com") - user.id = "account-1" - run = WorkflowExecution( - accumulated_elapsed_time=3.5, - id_="run-1", - workflow_id="workflow-1", - workflow_type=WorkflowType.WORKFLOW, - workflow_version="1", - graph={}, - inputs={}, - started_at=naive_utc_now(), - ) - run.status = WorkflowExecutionStatus.PAUSED - SQLAlchemyWorkflowExecutionRepository( - session_factory=sqlite_session_factory, - tenant_id="tenant-1", - user=user, - app_id="app-1", - triggered_from=WorkflowRunTriggeredFrom.APP_RUN, - ).save(run) - control = DifyAPISQLAlchemyWorkflowRunRepository(sqlite_session_factory) - control.bind_workflow_task(tenant_id="tenant-1", app_id="app-1", workflow_run_id="run-1", task_id="task-1") - node = WorkflowNodeExecution( - id="node-1", - node_execution_id="node-1", - workflow_id="workflow-1", - workflow_execution_id="run-1", - index=1, - node_id="review", - node_type="human-input", - title="Review", - created_at=run.started_at, - status=node_status, - ) - SQLAlchemyWorkflowNodeExecutionRepository( - session_factory=sqlite_session_factory, - tenant_id="tenant-1", - user=user, - app_id="app-1", - triggered_from=WorkflowNodeExecutionTriggeredFrom.WORKFLOW_RUN, - ).save_synchronously(node) - records: list[dict[str, str]] = [] - interrupted = False - - def append_after_concurrent_stop(_store: str, fields: list[tuple[str, str]]) -> None: - nonlocal interrupted - if not interrupted: - interrupted = True - stopped = control.stop_paused_workflow_task( - tenant_id="tenant-1", app_id="app-1", task_id="task-1", owner=None - ) - assert stopped is not None - if record_type == "workflow": - run_writer.save( - run.model_copy( - update={ - "status": WorkflowExecutionStatus.STOPPED, - "finished_at": stopped.finished_at, - "error_message": stopped.error, - } - ) - ) - else: - node_writer.save_many( - [ - node.model_copy( - update={ - "status": WorkflowNodeExecutionStatus.FAILED, - "finished_at": stopped.finished_at, - "error": stopped.error, - } - ) - ] - ) - records.append(dict(fields)) - - mocker.patch.object(logstore, "put_log", side_effect=append_after_concurrent_stop) - run_writer = LogstoreWorkflowExecutionRepository( - session_factory=sqlite_session_factory, - tenant_id="tenant-1", - user=user, - app_id="app-1", - triggered_from=WorkflowRunTriggeredFrom.APP_RUN, - ) - node_writer = LogstoreWorkflowNodeExecutionRepository( - session_factory=sqlite_session_factory, - tenant_id="tenant-1", - user=user, - app_id="app-1", - triggered_from=WorkflowNodeExecutionTriggeredFrom.WORKFLOW_RUN, - ) - run_writer._enable_dual_write = node_writer._enable_dual_write = False - if record_type == "workflow": - run_writer.save(run) - expected_status = "stopped" - assert max(records, key=lambda row: int(row["log_version"]))["status"] == expected_status - run.accumulated_elapsed_time = 999.0 - run_writer.save(run) # Another late RUNNING/PAUSED write also remains terminal. - assert float(records[-1]["elapsed_time"]) == 3.5 - else: - node_writer.save(node) - expected_status = "failed" - assert max(records, key=lambda row: int(row["log_version"]))["status"] == expected_status - stale_node = next(row for row in records if row["status"] == node_status) - mocker.patch.object(logstore, "execute_sql", return_value=[stale_node]) - api = LogstoreAPIWorkflowNodeExecutionRepository(sqlite_session_factory) # pyrefly: ignore[bad-instantiation] - assert api.get_executions_by_workflow_run("tenant-1", "app-1", "run-1")[0].status == expected_status - by_id = api.get_execution_by_id("node-1", "tenant-1") - assert by_id is not None - assert by_id.status == expected_status - assert ( - api.get_execution_snapshots_by_workflow_run("tenant-1", "app-1", "workflow-1", "workflow-run", "run-1")[ - 0 - ].status - == expected_status - ) - assert node_writer.get_by_workflow_execution("run-1")[0].status == expected_status - node_writer.save(node) - assert records[-1]["status"] == expected_status - - -@pytest.mark.parametrize( - "status", - [ - WorkflowNodeExecutionStatus.SUCCEEDED, - WorkflowNodeExecutionStatus.FAILED, - WorkflowNodeExecutionStatus.EXCEPTION, - WorkflowNodeExecutionStatus.STOPPED, - ], -) -def test_stop_preserves_completed_node_history(status: WorkflowNodeExecutionStatus) -> None: - finished_at = naive_utc_now() - stopped = WorkflowRun( - id="run-1", - tenant_id="tenant-1", - app_id="app-1", - workflow_id="workflow-1", - type=WorkflowType.WORKFLOW, - triggered_from=WorkflowRunTriggeredFrom.APP_RUN, - version="1", - status=WorkflowExecutionStatus.STOPPED, - created_by_role=CreatorUserRole.ACCOUNT, - created_by="account-1", - error="User requested stop", - finished_at=finished_at, - stop_requested_at=finished_at, - ) - node = WorkflowNodeExecution( - id="node-1", - workflow_id="workflow-1", - workflow_execution_id="run-1", - index=1, - node_id="answer", - node_type="llm", - title="Answer", - status=status, - outputs={"answer": "preserved"}, - error=None if status == WorkflowNodeExecutionStatus.SUCCEEDED else "Original terminal error", - elapsed_time=2.5, - created_at=finished_at, - finished_at=finished_at, - ) - original = node.model_dump() - - apply_workflow_stop_to_node(node, stopped) - - assert node.model_dump() == original diff --git a/api/tests/unit_tests/core/repositories/test_workflow_node_execution_batch.py b/api/tests/unit_tests/core/repositories/test_workflow_node_execution_batch.py deleted file mode 100644 index 5fcfc1d7b34..00000000000 --- a/api/tests/unit_tests/core/repositories/test_workflow_node_execution_batch.py +++ /dev/null @@ -1,231 +0,0 @@ -"""Batch cleanup persists exact records without losing ownership or stored data.""" - -from datetime import datetime -from uuid import uuid4 - -import pytest -from kombu.serialization import dumps, loads -from sqlalchemy import Engine, event, select -from sqlalchemy.exc import IntegrityError -from sqlalchemy.orm import Session, sessionmaker - -from core.repositories.celery_workflow_node_execution_repository import CeleryWorkflowNodeExecutionRepository -from core.repositories.sqlalchemy_workflow_node_execution_repository import SQLAlchemyWorkflowNodeExecutionRepository -from core.workflow.node_execution import WorkflowNodeExecution -from extensions.logstore.aliyun_logstore import AliyunLogStore -from extensions.logstore.repositories.logstore_workflow_node_execution_repository import ( - LogstoreWorkflowNodeExecutionRepository, -) -from graphon.enums import WorkflowNodeExecutionStatus -from models.enums import ExecutionOffLoadType -from models.workflow import WorkflowNodeExecutionModel, WorkflowNodeExecutionOffload, WorkflowNodeExecutionTriggeredFrom -from tasks import workflow_node_execution_tasks -from tests.unit_tests.config_override import apply_config_overrides -from tests.unit_tests.model_factories import make_account, make_upload_file - - -def _execution() -> WorkflowNodeExecution: - return WorkflowNodeExecution( - id=str(uuid4()), - node_execution_id=str(uuid4()), - workflow_id="source-workflow", - workflow_execution_id="root-run", - triggered_from_workflow_id="caller-workflow", - triggered_from_node_execution_id="caller-execution", - index=1, - node_id="tool", - node_type="tool", - title="Nested tool", - created_at=datetime(2026, 9, 23), - process_data={"workflow_agent_binding_id": "binding"}, - ) - - -@pytest.mark.parametrize("backend", ["sql", "celery", "logstore"]) -def test_batch_creates_missing_rows_and_updates_only_supplied_records( - sqlite_session_factory: sessionmaker[Session], sqlite_engine: Engine, monkeypatch: pytest.MonkeyPatch, backend: str -) -> None: - expiring_session_factory = sessionmaker(sqlite_engine) - sql = SQLAlchemyWorkflowNodeExecutionRepository( - expiring_session_factory, - "tenant", - make_account(), - "source-app", - WorkflowNodeExecutionTriggeredFrom.WORKFLOW_TOOL, - ) - existing, missing, untouched = (_execution() for _ in range(3)) - sql.save(existing) - sql.save(untouched) - with sqlite_session_factory.begin() as session: - session.add(make_upload_file(file_id="offloaded-inputs", tenant_id="tenant")) - row = session.get(WorkflowNodeExecutionModel, existing.id) - assert row is not None - row.inputs = '{"large": "stored preview"}' - session.add( - WorkflowNodeExecutionOffload( - tenant_id="tenant", - app_id="source-app", - node_execution_id=existing.id, - type_=ExecutionOffLoadType.INPUTS, - file_id="offloaded-inputs", - ) - ) - existing.inputs = {"large": "full content restored from offload"} - for execution in (existing, missing): - execution.status = WorkflowNodeExecutionStatus.FAILED - execution.error = "workflow stopped" - execution.finished_at = datetime(2026, 9, 23, 0, 0, 1) - execution.process_data = {"partial": "result"} - - queued: list[dict[str, object]] = [] - logstore_records: list[dict[str, str]] = [] - if backend == "celery": - task = workflow_node_execution_tasks.save_workflow_node_executions_task - monkeypatch.setattr(task, "delay", lambda **kwargs: queued.append(kwargs)) - repository = CeleryWorkflowNodeExecutionRepository( - expiring_session_factory, - "tenant", - make_account(), - "source-app", - WorkflowNodeExecutionTriggeredFrom.WORKFLOW_TOOL, - ) - elif backend == "logstore": - apply_config_overrides(monkeypatch, LOGSTORE_DUAL_WRITE_ENABLED=True) - monkeypatch.setattr(AliyunLogStore, "_instance", None) - monkeypatch.setattr(AliyunLogStore, "_initialized", False) - monkeypatch.setenv("ALIYUN_SLS_ENDPOINT", "localhost") - monkeypatch.setattr(AliyunLogStore, "_check_endpoint_connectivity", staticmethod(lambda *_: None)) - monkeypatch.setattr( - AliyunLogStore, - "put_log", - lambda _self, _name, record: logstore_records.append(dict(record)), - ) - repository = LogstoreWorkflowNodeExecutionRepository( - expiring_session_factory, - "tenant", - make_account(), - "source-app", - WorkflowNodeExecutionTriggeredFrom.WORKFLOW_TOOL, - ) - else: - repository = sql - - commits: list[None] = [] - event.listen(sqlite_engine, "commit", lambda _connection: commits.append(None)) - repository.save_many([existing, missing]) - repository.save_many([]) - if backend == "celery": - assert len(queued) == 1 - assert commits == [] - content_type, encoding, payload = dumps(queued[0], serializer="json") - workflow_node_execution_tasks.save_workflow_node_executions_task.run( - **loads(payload, content_type, encoding, accept={content_type}) - ) - assert len(commits) == 1 - - with sqlite_session_factory() as session: - rows = { - row.id: row - for row in session.scalars( - WorkflowNodeExecutionModel.preload_offload_data(select(WorkflowNodeExecutionModel)) - ) - } - assert set(rows) == {existing.id, missing.id, untouched.id} - assert rows[untouched.id].status == WorkflowNodeExecutionStatus.RUNNING - for execution in (existing, missing): - row = rows[execution.id] - assert row.status == WorkflowNodeExecutionStatus.FAILED - assert row.error == "workflow stopped" - assert row.finished_at == execution.finished_at - assert (row.tenant_id, row.app_id, row.triggered_from) == ( - "tenant", - "source-app", - WorkflowNodeExecutionTriggeredFrom.WORKFLOW_TOOL, - ) - assert (row.workflow_id, row.workflow_run_id) == ("source-workflow", "root-run") - assert (row.triggered_from_workflow_id, row.triggered_from_node_execution_id) == ( - "caller-workflow", - "caller-execution", - ) - assert rows[existing.id].process_data_dict == {"partial": "result", "workflow_agent_binding_id": "binding"} - assert rows[existing.id].inputs_dict == {"large": "stored preview"} - assert rows[existing.id].offload_data[0].file_id == "offloaded-inputs" - if backend == "logstore": - assert {row["id"] for row in logstore_records} == {existing.id, missing.id} - assert all(row["status"] == "failed" and row["log_version"] for row in logstore_records) - assert all(row["triggered_from_node_execution_id"] == "caller-execution" for row in logstore_records) - - -@pytest.mark.parametrize("backend", ["sql", "celery"]) -@pytest.mark.parametrize("owner", ["tenant_id", "app_id", "triggered_from", "workflow_id", "workflow_run_id"]) -def test_owner_mismatch_rolls_back_the_whole_batch( - sqlite_session_factory: sessionmaker[Session], monkeypatch: pytest.MonkeyPatch, backend: str, owner: str -) -> None: - repository = SQLAlchemyWorkflowNodeExecutionRepository( - sqlite_session_factory, "tenant", make_account(), "source-app", WorkflowNodeExecutionTriggeredFrom.WORKFLOW_TOOL - ) - valid, foreign = _execution(), _execution() - repository.save(valid) - repository.save(foreign) - with sqlite_session_factory.begin() as session: - row = session.get(WorkflowNodeExecutionModel, foreign.id) - assert row is not None - setattr(row, owner, WorkflowNodeExecutionTriggeredFrom.WORKFLOW_RUN if owner == "triggered_from" else "other") - for execution in (valid, foreign): - execution.status = WorkflowNodeExecutionStatus.FAILED - if backend == "celery": - task = workflow_node_execution_tasks.save_workflow_node_executions_task - monkeypatch.setattr(task, "retry", lambda **kwargs: kwargs["exc"]) - with pytest.raises(IntegrityError): - task.run( - executions_data=[execution.model_dump() for execution in (valid, foreign)], - tenant_id="tenant", - app_id="source-app", - triggered_from=WorkflowNodeExecutionTriggeredFrom.WORKFLOW_TOOL.value, - creator_user_id="account", - creator_user_role="account", - ) - else: - with pytest.raises(IntegrityError): - repository.save_many([valid, foreign]) - with sqlite_session_factory() as session: - assert all( - row.status == WorkflowNodeExecutionStatus.RUNNING - for row in session.scalars(select(WorkflowNodeExecutionModel)) - ) - - -@pytest.mark.parametrize( - "status", - [WorkflowNodeExecutionStatus.FAILED, WorkflowNodeExecutionStatus.PAUSED, WorkflowNodeExecutionStatus.RETRY], -) -def test_delayed_start_cannot_undo_completed_failure_cleanup( - sqlite_session_factory: sessionmaker[Session], monkeypatch: pytest.MonkeyPatch, status: WorkflowNodeExecutionStatus -) -> None: - queued: list[dict[str, object]] = [] - start_task = workflow_node_execution_tasks.save_workflow_node_execution_task - batch_task = workflow_node_execution_tasks.save_workflow_node_executions_task - monkeypatch.setattr(start_task, "delay", lambda **kwargs: queued.append(kwargs)) - monkeypatch.setattr(batch_task, "delay", lambda **kwargs: queued.append(kwargs)) - repository = CeleryWorkflowNodeExecutionRepository( - sqlite_session_factory, "tenant", make_account(), "source-app", WorkflowNodeExecutionTriggeredFrom.WORKFLOW_TOOL - ) - execution = _execution() - repository.save(execution) - execution.status = status - execution.error = "workflow stopped" - execution.finished_at = datetime(2026, 9, 23, 0, 0, 1) - repository.save_many([execution]) - - for task, kwargs in zip((batch_task, start_task), reversed(queued), strict=True): - content_type, encoding, payload = dumps(kwargs, serializer="json") - task.run(**loads(payload, content_type, encoding, accept={content_type})) - - with sqlite_session_factory() as session: - row = session.get(WorkflowNodeExecutionModel, execution.id) - assert row is not None - if status == WorkflowNodeExecutionStatus.FAILED: - assert (row.status, row.error, row.finished_at) == (status, execution.error, execution.finished_at) - else: - assert row.status == WorkflowNodeExecutionStatus.RUNNING - assert row.finished_at is None diff --git a/api/tests/unit_tests/core/repositories/test_workflow_node_execution_callers.py b/api/tests/unit_tests/core/repositories/test_workflow_node_execution_callers.py deleted file mode 100644 index 4f51cce623c..00000000000 --- a/api/tests/unit_tests/core/repositories/test_workflow_node_execution_callers.py +++ /dev/null @@ -1,163 +0,0 @@ -"""Caller identity survives storage updates independently of node payloads.""" - -from datetime import datetime -from uuid import uuid4 - -import pytest -from kombu.serialization import dumps, loads -from sqlalchemy.orm import Session, sessionmaker - -from core.repositories.celery_workflow_node_execution_repository import CeleryWorkflowNodeExecutionRepository -from core.repositories.factory import WorkflowNodeExecutionRepository -from core.repositories.sqlalchemy_workflow_node_execution_repository import SQLAlchemyWorkflowNodeExecutionRepository -from core.workflow.node_execution import WorkflowNodeExecution -from extensions.logstore.aliyun_logstore import AliyunLogStore -from extensions.logstore.repositories.logstore_workflow_node_execution_repository import ( - LogstoreWorkflowNodeExecutionRepository, -) -from graphon.enums import WorkflowNodeExecutionStatus -from models.model import UploadFile -from models.workflow import WorkflowNodeExecutionModel, WorkflowNodeExecutionTriggeredFrom -from repositories.factory import DifyAPIRepositoryFactory -from tasks.workflow_node_execution_tasks import save_workflow_node_execution_task -from tests.unit_tests.config_override import apply_config_overrides -from tests.unit_tests.model_factories import make_account, make_upload_file - - -@pytest.mark.parametrize("backend", ["sql", "celery", "logstore"]) -def test_caller_locators_survive_storage_round_trip_and_updates( - sqlite_session_factory: sessionmaker[Session], monkeypatch: pytest.MonkeyPatch, backend: str -) -> None: - apply_config_overrides(monkeypatch, LOGSTORE_DUAL_WRITE_ENABLED=False) - logstore_records: list[dict[str, str]] = [] - sql = SQLAlchemyWorkflowNodeExecutionRepository( - sqlite_session_factory, "tenant", make_account(), "source-app", WorkflowNodeExecutionTriggeredFrom.WORKFLOW_TOOL - ) - execution = WorkflowNodeExecution( - id=str(uuid4()), - node_execution_id="child-execution", - workflow_id=str(uuid4()), - workflow_execution_id="root-run", - triggered_from_workflow_id=str(uuid4()), - triggered_from_node_execution_id="call/" + "x" * 250, - index=1, - node_id="reused-canvas-node", - node_type="tool", - title="Nested Tool", - created_at=datetime(2026, 9, 23), - ) - repository: WorkflowNodeExecutionRepository - if backend == "celery": - - def deliver(**kwargs: object) -> None: - content_type, encoding, payload = dumps(kwargs, serializer="json") - save_workflow_node_execution_task.run(**loads(payload, content_type, encoding, accept={content_type})) - - monkeypatch.setattr(save_workflow_node_execution_task, "delay", deliver) - repository = CeleryWorkflowNodeExecutionRepository( - sqlite_session_factory, - "tenant", - make_account(), - "source-app", - WorkflowNodeExecutionTriggeredFrom.WORKFLOW_TOOL, - ) - elif backend == "logstore": - monkeypatch.setattr(AliyunLogStore, "_instance", None) - monkeypatch.setattr(AliyunLogStore, "_initialized", False) - monkeypatch.setenv("ALIYUN_SLS_ENDPOINT", "localhost") - monkeypatch.setattr(AliyunLogStore, "_check_endpoint_connectivity", staticmethod(lambda *_: None)) - monkeypatch.setattr( - AliyunLogStore, - "put_log", - lambda _self, _name, record: logstore_records.append(dict(record)), - ) - repository = LogstoreWorkflowNodeExecutionRepository( - sqlite_session_factory, - "tenant", - make_account(), - "source-app", - WorkflowNodeExecutionTriggeredFrom.WORKFLOW_TOOL, - ) - else: - repository = sql - - repository.save(execution) - execution.status = WorkflowNodeExecutionStatus.SUCCEEDED - execution.process_data = {"result": "no caller identifiers in this payload"} - execution.finished_at = datetime(2026, 9, 23, 0, 0, 1) - repository.save(execution) - repository.save_execution_data(execution) - - if backend == "logstore": - row = logstore_records[-1] - monkeypatch.setattr(AliyunLogStore, "execute_sql", lambda *_args, **_kwargs: [row]) - restored = repository.get_by_workflow_execution("root-run")[0] - apply_config_overrides( - monkeypatch, - API_WORKFLOW_NODE_EXECUTION_REPOSITORY=( - "extensions.logstore.repositories.logstore_api_workflow_node_execution_repository." - "LogstoreAPIWorkflowNodeExecutionRepository" - ), - ) - api_repository = DifyAPIRepositoryFactory.create_api_workflow_node_execution_repository(sqlite_session_factory) - model = api_repository.get_execution_by_id(execution.id, "tenant") - else: - restored = sql.get_by_workflow_execution("root-run")[0] - with sqlite_session_factory() as session: - model = session.get(WorkflowNodeExecutionModel, execution.id) - - assert model is not None - assert isinstance(restored, WorkflowNodeExecution) - for record in (restored, model): - assert record.triggered_from_workflow_id == execution.triggered_from_workflow_id - assert record.triggered_from_node_execution_id == execution.triggered_from_node_execution_id - assert record.workflow_id == execution.workflow_id - assert model.app_id == "source-app" - assert model.workflow_run_id == "root-run" - assert restored.process_data == execution.process_data - - -def test_sql_offload_preserves_caller_locators( - sqlite_session_factory: sessionmaker[Session], monkeypatch: pytest.MonkeyPatch -) -> None: - apply_config_overrides( - monkeypatch, WORKFLOW_VARIABLE_TRUNCATION_MAX_SIZE=32, WORKFLOW_VARIABLE_TRUNCATION_STRING_LENGTH=8 - ) - repository = SQLAlchemyWorkflowNodeExecutionRepository( - sqlite_session_factory, "tenant", make_account(), "source-app", WorkflowNodeExecutionTriggeredFrom.WORKFLOW_TOOL - ) - stored_content: dict[str, bytes] = {} - - def upload_file(*, content: bytes, **_kwargs: object) -> UploadFile: - file = make_upload_file(file_id=str(uuid4()), key="process-data", tenant_id="tenant") - with sqlite_session_factory.begin() as session: - session.add(file) - stored_content[file.key] = content - return file - - monkeypatch.setattr(repository._file_service, "upload_file", upload_file) - monkeypatch.setattr( - "core.repositories.sqlalchemy_workflow_node_execution_repository.storage.load", stored_content.__getitem__ - ) - execution = WorkflowNodeExecution( - id=str(uuid4()), - workflow_id=str(uuid4()), - workflow_execution_id="root-run", - triggered_from_workflow_id=str(uuid4()), - triggered_from_node_execution_id="immediate-tool-execution", - index=1, - node_id="node", - node_type="code", - title="Code", - created_at=datetime(2026, 9, 23), - process_data={"large": "x" * 1000}, - ) - repository.save(execution) - repository.save_execution_data(execution) - - restored = repository.get_by_workflow_execution("root-run")[0] - assert isinstance(restored, WorkflowNodeExecution) - assert restored.process_data_truncated - assert restored.process_data == execution.process_data - assert restored.triggered_from_workflow_id == execution.triggered_from_workflow_id - assert restored.triggered_from_node_execution_id == "immediate-tool-execution" diff --git a/api/tests/unit_tests/core/repositories/test_workflow_node_execution_conflict_handling.py b/api/tests/unit_tests/core/repositories/test_workflow_node_execution_conflict_handling.py deleted file mode 100644 index d9f39da073f..00000000000 --- a/api/tests/unit_tests/core/repositories/test_workflow_node_execution_conflict_handling.py +++ /dev/null @@ -1,195 +0,0 @@ -"""Unit tests for workflow node execution conflict handling.""" - -from collections.abc import Iterator -from contextlib import contextmanager -from dataclasses import dataclass - -import psycopg2.errors -import pytest -from sqlalchemy import Engine, event -from sqlalchemy.exc import IntegrityError -from sqlalchemy.orm import Session, sessionmaker - -from core.repositories.sqlalchemy_workflow_node_execution_repository import ( - SQLAlchemyWorkflowNodeExecutionRepository, -) -from graphon.entities.workflow_node_execution import ( - WorkflowNodeExecution, - WorkflowNodeExecutionStatus, -) -from graphon.enums import BuiltinNodeTypes -from libs.datetime_utils import naive_utc_now -from models import Account, WorkflowNodeExecutionModel, WorkflowNodeExecutionTriggeredFrom -from tests.unit_tests.model_factories import make_account, make_tenant - - -@dataclass(frozen=True) -class ConflictDatabase: - engine: Engine - session: Session - session_factory: sessionmaker[Session] - - -@dataclass(frozen=True) -class ConflictEvents: - insert_attempts: list[str] - rollbacks: list[bool] - - -@pytest.fixture -def conflict_database( - sqlite_session: Session, - sqlite_session_factory: sessionmaker[Session], -) -> ConflictDatabase: - """Use the shared SQLite fixtures for repository-owned sessions and assertions.""" - engine = sqlite_session.get_bind() - assert isinstance(engine, Engine) - return ConflictDatabase( - engine=engine, - session=sqlite_session, - session_factory=sqlite_session_factory, - ) - - -def _account() -> Account: - return make_account( - account_id="test-user-id", - name="Conflict User", - email="conflict@example.com", - tenant=make_tenant(tenant_id="test-tenant-id", name="Conflict Tenant"), - ) - - -@pytest.fixture -def repository(conflict_database: ConflictDatabase) -> SQLAlchemyWorkflowNodeExecutionRepository: - return SQLAlchemyWorkflowNodeExecutionRepository( - session_factory=conflict_database.session_factory, - tenant_id="test-tenant-id", - user=_account(), - app_id="test-app-id", - triggered_from=WorkflowNodeExecutionTriggeredFrom.WORKFLOW_RUN, - ) - - -def _execution( - *, - execution_id: str, - status: WorkflowNodeExecutionStatus = WorkflowNodeExecutionStatus.RUNNING, -) -> WorkflowNodeExecution: - return WorkflowNodeExecution( - id=execution_id, - workflow_id="test-workflow-id", - workflow_execution_id="test-workflow-execution-id", - node_execution_id="test-node-execution-id", - node_id="test-node-id", - node_type=BuiltinNodeTypes.START, - title="Test Node", - index=1, - status=status, - created_at=naive_utc_now(), - ) - - -@contextmanager -def _fail_inserts( - engine: Engine, - *, - failure_count: int, - duplicate: bool, -) -> Iterator[ConflictEvents]: - insert_attempts: list[str] = [] - rollbacks: list[bool] = [] - - def fail_insert(_connection, _cursor, statement, parameters, _context, _executemany) -> None: - if not statement.lstrip().upper().startswith("INSERT INTO WORKFLOW_NODE_EXECUTIONS"): - return - insert_attempts.append(statement) - if len(insert_attempts) > failure_count: - return - original_error = ( - psycopg2.errors.UniqueViolation("forced duplicate key") - if duplicate - else RuntimeError("forced non-duplicate constraint failure") - ) - raise IntegrityError(statement, parameters, original_error) - - def record_rollback(_connection) -> None: - rollbacks.append(True) - - event.listen(engine, "before_cursor_execute", fail_insert) - event.listen(engine, "rollback", record_rollback) - try: - yield ConflictEvents(insert_attempts=insert_attempts, rollbacks=rollbacks) - finally: - event.remove(engine, "before_cursor_execute", fail_insert) - event.remove(engine, "rollback", record_rollback) - - -class TestWorkflowNodeExecutionConflictHandling: - """Test cases for handling duplicate key conflicts in workflow node execution.""" - - def test_save_with_duplicate_key_retries_with_new_uuid( - self, - repository: SQLAlchemyWorkflowNodeExecutionRepository, - conflict_database: ConflictDatabase, - ) -> None: - execution = _execution(execution_id="original-id") - original_id = execution.id - - with _fail_inserts(conflict_database.engine, failure_count=1, duplicate=True) as conflicts: - repository.save(execution) - - assert len(conflicts.insert_attempts) == 2 - assert len(conflicts.rollbacks) == 1 - assert execution.id != original_id - assert conflict_database.session.get(WorkflowNodeExecutionModel, original_id) is None - persisted = conflict_database.session.get(WorkflowNodeExecutionModel, execution.id) - assert persisted is not None - assert persisted.node_execution_id == execution.node_execution_id - - def test_save_with_existing_record_updates_instead_of_insert( - self, - repository: SQLAlchemyWorkflowNodeExecutionRepository, - conflict_database: ConflictDatabase, - ) -> None: - execution = _execution(execution_id="existing-id") - repository.save(execution) - - execution.status = WorkflowNodeExecutionStatus.SUCCEEDED - repository.save(execution) - - conflict_database.session.expire_all() - persisted = conflict_database.session.get(WorkflowNodeExecutionModel, execution.id) - assert persisted is not None - assert persisted.status == WorkflowNodeExecutionStatus.SUCCEEDED - assert conflict_database.session.query(WorkflowNodeExecutionModel).count() == 1 - - def test_save_exceeds_max_retries_raises_error( - self, - repository: SQLAlchemyWorkflowNodeExecutionRepository, - conflict_database: ConflictDatabase, - ) -> None: - execution = _execution(execution_id="test-id") - - with _fail_inserts(conflict_database.engine, failure_count=3, duplicate=True) as conflicts: - with pytest.raises(IntegrityError): - repository.save(execution) - - assert len(conflicts.insert_attempts) == 3 - assert len(conflicts.rollbacks) == 3 - assert conflict_database.session.query(WorkflowNodeExecutionModel).count() == 0 - - def test_save_non_duplicate_integrity_error_raises_immediately( - self, - repository: SQLAlchemyWorkflowNodeExecutionRepository, - conflict_database: ConflictDatabase, - ) -> None: - execution = _execution(execution_id="test-id") - - with _fail_inserts(conflict_database.engine, failure_count=1, duplicate=False) as conflicts: - with pytest.raises(IntegrityError): - repository.save(execution) - - assert len(conflicts.insert_attempts) == 1 - assert len(conflicts.rollbacks) == 1 - assert conflict_database.session.get(WorkflowNodeExecutionModel, execution.id) is None diff --git a/api/tests/unit_tests/core/repositories/test_workflow_node_execution_truncation.py b/api/tests/unit_tests/core/repositories/test_workflow_node_execution_truncation.py deleted file mode 100644 index ba7b5282993..00000000000 --- a/api/tests/unit_tests/core/repositories/test_workflow_node_execution_truncation.py +++ /dev/null @@ -1,235 +0,0 @@ -""" -Unit tests for WorkflowNodeExecution truncation functionality. - -Tests the truncation and offloading logic for large inputs and outputs -in the SQLAlchemyWorkflowNodeExecutionRepository. -""" - -import json -from dataclasses import dataclass -from datetime import UTC, datetime -from typing import Any - -from sqlalchemy import Engine -from sqlalchemy.orm import Session - -from configs import dify_config -from core.repositories.sqlalchemy_workflow_node_execution_repository import ( - SQLAlchemyWorkflowNodeExecutionRepository, -) -from graphon.entities.workflow_node_execution import ( - WorkflowNodeExecution, - WorkflowNodeExecutionStatus, -) -from graphon.enums import BuiltinNodeTypes -from models import Account, Tenant, WorkflowNodeExecutionTriggeredFrom -from models.enums import ExecutionOffLoadType -from models.workflow import WorkflowNodeExecutionModel, WorkflowNodeExecutionOffload - - -@dataclass -class TruncationTestCase: - """Test case data for truncation scenarios.""" - - name: str - inputs: dict[str, Any] | None - outputs: dict[str, Any] | None - should_truncate_inputs: bool - should_truncate_outputs: bool - description: str - - -def create_test_cases() -> list[TruncationTestCase]: - """Create test cases for different truncation scenarios.""" - # Create large data that will definitely exceed the threshold (10KB) - large_data = {"data": "x" * (dify_config.WORKFLOW_VARIABLE_TRUNCATION_MAX_SIZE + 1000)} - small_data = {"data": "small"} - - return [ - TruncationTestCase( - name="small_data_no_truncation", - inputs=small_data, - outputs=small_data, - should_truncate_inputs=False, - should_truncate_outputs=False, - description="Small data should not be truncated", - ), - TruncationTestCase( - name="large_inputs_truncation", - inputs=large_data, - outputs=small_data, - should_truncate_inputs=True, - should_truncate_outputs=False, - description="Large inputs should be truncated", - ), - TruncationTestCase( - name="large_outputs_truncation", - inputs=small_data, - outputs=large_data, - should_truncate_inputs=False, - should_truncate_outputs=True, - description="Large outputs should be truncated", - ), - TruncationTestCase( - name="large_both_truncation", - inputs=large_data, - outputs=large_data, - should_truncate_inputs=True, - should_truncate_outputs=True, - description="Both large inputs and outputs should be truncated", - ), - TruncationTestCase( - name="none_inputs_outputs", - inputs=None, - outputs=None, - should_truncate_inputs=False, - should_truncate_outputs=False, - description="None inputs and outputs should not be truncated", - ), - ] - - -def create_workflow_node_execution( - execution_id: str = "test-execution-id", - inputs: dict[str, Any] | None = None, - outputs: dict[str, Any] | None = None, -) -> WorkflowNodeExecution: - """Factory function to create a WorkflowNodeExecution for testing.""" - return WorkflowNodeExecution( - id=execution_id, - node_execution_id="test-node-execution-id", - workflow_id="test-workflow-id", - workflow_execution_id="test-workflow-execution-id", - index=1, - node_id="test-node-id", - node_type=BuiltinNodeTypes.LLM, - title="Test Node", - inputs=inputs, - outputs=outputs, - status=WorkflowNodeExecutionStatus.SUCCEEDED, - created_at=datetime.now(UTC), - ) - - -def mock_user() -> Account: - """Create an Account user for testing.""" - - user = Account(name="Test Account", email="test@example.com") - user.id = "test-user-id" - user._current_tenant = Tenant(name="Test Tenant") - user._current_tenant.id = "test-tenant-id" - return user - - -def create_offload(type_: ExecutionOffLoadType) -> WorkflowNodeExecutionOffload: - return WorkflowNodeExecutionOffload( - tenant_id="test-tenant-id", - app_id="test-app-id", - node_execution_id="test-node-execution-id", - type_=type_, - file_id="test-file-id", - ) - - -class TestSQLAlchemyWorkflowNodeExecutionRepositoryTruncation: - """Test class for truncation functionality in SQLAlchemyWorkflowNodeExecutionRepository.""" - - def create_repository(self, sqlite_engine: Engine) -> SQLAlchemyWorkflowNodeExecutionRepository: - """Create a repository backed by the test's isolated SQLite engine.""" - repository = SQLAlchemyWorkflowNodeExecutionRepository( - session_factory=sqlite_engine, - tenant_id="test-tenant-id", - user=mock_user(), - app_id="test-app-id", - triggered_from=WorkflowNodeExecutionTriggeredFrom.WORKFLOW_RUN, - ) - with repository._session_factory() as session: - assert isinstance(session, Session) - assert session.get_bind() is sqlite_engine - return repository - - def test_to_domain_model_without_offload_data(self, sqlite_engine: Engine): - """Test _to_domain_model correctly handles models without offload data.""" - repo = self.create_repository(sqlite_engine) - - # Create a database model without offload data - db_model = WorkflowNodeExecutionModel( - id="test-id", - node_execution_id="node-exec-id", - workflow_id="workflow-id", - workflow_run_id="run-id", - index=1, - predecessor_node_id=None, - node_id="node-id", - node_type=BuiltinNodeTypes.LLM, - title="Test Node", - inputs=json.dumps({"value": "inputs"}), - process_data=json.dumps({"value": "process_data"}), - outputs=json.dumps({"value": "outputs"}), - status=WorkflowNodeExecutionStatus.SUCCEEDED, - error=None, - elapsed_time=1.0, - execution_metadata="{}", - created_at=datetime.now(UTC), - finished_at=None, - offload_data=[], - ) - - domain_model = repo._to_domain_model(db_model) - - # Check that no truncated data was set - assert domain_model.get_truncated_inputs() is None - assert domain_model.get_truncated_outputs() is None - - -class TestWorkflowNodeExecutionModelTruncatedProperties: - """Test the truncated properties on WorkflowNodeExecutionModel.""" - - def test_inputs_truncated_with_offload_data(self): - """Test inputs_truncated property when offload data exists.""" - model = WorkflowNodeExecutionModel() - offload = create_offload(ExecutionOffLoadType.INPUTS) - model.offload_data = [offload] - - assert model.inputs_truncated is True - assert model.process_data_truncated is False - assert model.outputs_truncated is False - - def test_outputs_truncated_with_offload_data(self): - """Test outputs_truncated property when offload data exists.""" - model = WorkflowNodeExecutionModel() - - # Mock offload data with outputs file - offload = create_offload(ExecutionOffLoadType.OUTPUTS) - model.offload_data = [offload] - - assert model.inputs_truncated is False - assert model.process_data_truncated is False - assert model.outputs_truncated is True - - def test_process_data_truncated_with_offload_data(self): - model = WorkflowNodeExecutionModel() - offload = create_offload(ExecutionOffLoadType.PROCESS_DATA) - model.offload_data = [offload] - assert model.process_data_truncated is True - assert model.inputs_truncated is False - assert model.outputs_truncated is False - - def test_truncated_properties_without_offload_data(self): - """Test truncated properties when no offload data exists.""" - model = WorkflowNodeExecutionModel( - offload_data=[], - ) - - assert model.inputs_truncated is False - assert model.outputs_truncated is False - assert model.process_data_truncated is False - - def test_truncated_properties_without_offload_attribute(self): - """Test truncated properties when offload_data attribute doesn't exist.""" - model = WorkflowNodeExecutionModel() - # Don't set offload_data attribute at all - - assert model.inputs_truncated is False - assert model.outputs_truncated is False - assert model.process_data_truncated is False diff --git a/api/tests/unit_tests/core/telemetry/test_facade.py b/api/tests/unit_tests/core/telemetry/test___init__.py similarity index 100% rename from api/tests/unit_tests/core/telemetry/test_facade.py rename to api/tests/unit_tests/core/telemetry/test___init__.py diff --git a/api/tests/unit_tests/enterprise/telemetry/test_gateway.py b/api/tests/unit_tests/core/telemetry/test_gateway.py similarity index 55% rename from api/tests/unit_tests/enterprise/telemetry/test_gateway.py rename to api/tests/unit_tests/core/telemetry/test_gateway.py index f3bb98dc0bc..8569c5388c0 100644 --- a/api/tests/unit_tests/enterprise/telemetry/test_gateway.py +++ b/api/tests/unit_tests/core/telemetry/test_gateway.py @@ -1,5 +1,6 @@ from __future__ import annotations +import sys from unittest.mock import MagicMock, patch import pytest @@ -15,27 +16,34 @@ from core.telemetry.events import ( TelemetryContext, TelemetryEvent, ) -from core.telemetry.gateway import PAYLOAD_SIZE_THRESHOLD_BYTES, emit +from core.telemetry.gateway import PAYLOAD_SIZE_THRESHOLD_BYTES, emit, is_enterprise_telemetry_enabled from enterprise.telemetry.contracts import SignalType, TelemetryCase, TelemetryEnvelope + +class TestTelemetryCoreExports: + def test_is_enterprise_telemetry_enabled_exported(self) -> None: + from core.telemetry.gateway import is_enterprise_telemetry_enabled as exported_func + + assert callable(exported_func) + + # --------------------------------------------------------------------------- # Helpers # --------------------------------------------------------------------------- -_TRACE_CTX = TelemetryContext(tenant_id="tenant-789", user_id="user-456", app_id="app-123") -_METRIC_CTX = TelemetryContext(tenant_id="tenant-123") +_GATEWAY_INTEGRATION_TRACE_CTX = TelemetryContext(tenant_id="tenant-789", user_id="user-456", app_id="app-123") -def _node_event() -> DraftNodeExecutionTraceEvent: +def _gateway_integration_node_event() -> DraftNodeExecutionTraceEvent: return DraftNodeExecutionTraceEvent( - context=_TRACE_CTX, - payload=NodeExecutionPayload(node_execution_data={"key": "val"}), + context=_GATEWAY_INTEGRATION_TRACE_CTX, + payload=NodeExecutionPayload(node_execution_data={}), ) -def _prompt_event() -> PromptGenerationEvent: +def _gateway_integration_prompt_event() -> PromptGenerationEvent: return PromptGenerationEvent( - context=_TRACE_CTX, + context=_GATEWAY_INTEGRATION_TRACE_CTX, payload=PromptGenerationPayload( tenant_id="tenant-789", operation_type="generate", @@ -51,7 +59,148 @@ def _prompt_event() -> PromptGenerationEvent: ) -def _app_created_event(tenant_id: str = "tenant-123") -> AppCreatedEvent: +def _gateway_integration_app_created_event() -> AppCreatedEvent: + return AppCreatedEvent( + context=TelemetryContext(tenant_id="tenant-123"), + payload=AppCreatedPayload(app_id="app-abc"), + ) + + +# --------------------------------------------------------------------------- +# Trace routing integration +# --------------------------------------------------------------------------- + + +class TestGatewayIntegrationTraceRouting: + @pytest.fixture + def mock_trace_manager(self) -> MagicMock: + return MagicMock() + + def test_trace_event_routed_to_trace_manager( + self, + mock_trace_manager: MagicMock, + ) -> None: + with patch("core.telemetry.gateway.is_enterprise_telemetry_enabled", return_value=True): + emit(_gateway_integration_prompt_event(), mock_trace_manager) + mock_trace_manager.add_trace_task.assert_called_once() + + def test_enterprise_only_trace_dropped_when_ee_disabled( + self, + mock_trace_manager: MagicMock, + ) -> None: + with patch("core.telemetry.gateway.is_enterprise_telemetry_enabled", return_value=False): + emit(_gateway_integration_node_event(), mock_trace_manager) + mock_trace_manager.add_trace_task.assert_not_called() + + def test_enterprise_only_trace_routed_when_ee_enabled( + self, + mock_trace_manager: MagicMock, + ) -> None: + with patch("core.telemetry.gateway.is_enterprise_telemetry_enabled", return_value=True): + emit(_gateway_integration_node_event(), mock_trace_manager) + mock_trace_manager.add_trace_task.assert_called_once() + + +# --------------------------------------------------------------------------- +# Metric/log routing integration +# --------------------------------------------------------------------------- + + +class TestGatewayIntegrationMetricRouting: + @patch("core.telemetry.gateway.is_enterprise_telemetry_enabled", return_value=True) + def test_metric_event_routes_to_celery_task( + self, + mock_ee_enabled: MagicMock, + ) -> None: + from enterprise.telemetry.contracts import TelemetryEnvelope + + with patch("tasks.enterprise_telemetry_task.process_enterprise_telemetry.delay") as mock_delay: + emit(_gateway_integration_app_created_event()) + + mock_delay.assert_called_once() + envelope_json = mock_delay.call_args[0][0] + envelope = TelemetryEnvelope.model_validate_json(envelope_json) + assert envelope.case == TelemetryCase.APP_CREATED + assert envelope.tenant_id == "tenant-123" + assert envelope.payload["app_id"] == "app-abc" + + +# --------------------------------------------------------------------------- +# CE eligibility integration +# --------------------------------------------------------------------------- + + +class TestGatewayIntegrationCEEligibility: + @pytest.fixture + def mock_trace_manager(self) -> MagicMock: + return MagicMock() + + def test_draft_node_execution_not_ce_eligible( + self, + mock_trace_manager: MagicMock, + ) -> None: + with patch("core.telemetry.gateway.is_enterprise_telemetry_enabled", return_value=False): + emit(_gateway_integration_node_event(), mock_trace_manager) + mock_trace_manager.add_trace_task.assert_not_called() + + def test_prompt_generation_not_ce_eligible( + self, + mock_trace_manager: MagicMock, + ) -> None: + with patch("core.telemetry.gateway.is_enterprise_telemetry_enabled", return_value=False): + emit(_gateway_integration_prompt_event(), mock_trace_manager) + mock_trace_manager.add_trace_task.assert_not_called() + + +# --------------------------------------------------------------------------- +# is_enterprise_telemetry_enabled +# --------------------------------------------------------------------------- + + +class TestIsEnterpriseTelemetryEnabled: + def test_returns_false_when_exporter_import_fails(self) -> None: + with patch.dict(sys.modules, {"enterprise.telemetry.exporter": None}): + result = is_enterprise_telemetry_enabled() + assert result is False + + def test_function_is_callable(self) -> None: + assert callable(is_enterprise_telemetry_enabled) + + +# --------------------------------------------------------------------------- +# Helpers +# --------------------------------------------------------------------------- + +_GATEWAY_TRACE_CTX = TelemetryContext(tenant_id="tenant-789", user_id="user-456", app_id="app-123") +_METRIC_CTX = TelemetryContext(tenant_id="tenant-123") + + +def _gateway_node_event() -> DraftNodeExecutionTraceEvent: + return DraftNodeExecutionTraceEvent( + context=_GATEWAY_TRACE_CTX, + payload=NodeExecutionPayload(node_execution_data={"key": "val"}), + ) + + +def _gateway_prompt_event() -> PromptGenerationEvent: + return PromptGenerationEvent( + context=_GATEWAY_TRACE_CTX, + payload=PromptGenerationPayload( + tenant_id="tenant-789", + operation_type="generate", + instruction="test", + generated_output="out", + model_provider="openai", + model_name="gpt-4", + prompt_tokens=10, + completion_tokens=5, + total_tokens=15, + latency=0.5, + ), + ) + + +def _gateway_app_created_event(tenant_id: str = "tenant-123") -> AppCreatedEvent: return AppCreatedEvent( context=TelemetryContext(tenant_id=tenant_id), payload=AppCreatedPayload(app_id="app-abc", mode="chat"), @@ -67,28 +216,28 @@ class TestEventRoutingMetadata: """Verify each event class declares correct routing fields.""" def test_draft_node_execution_routing(self) -> None: - ev = _node_event() + ev = _gateway_node_event() assert ev.case == TelemetryCase.DRAFT_NODE_EXECUTION assert ev.signal_type is SignalType.TRACE assert ev.ce_eligible is False assert ev.trace_task_name is TraceTaskName.DRAFT_NODE_EXECUTION_TRACE def test_prompt_generation_routing(self) -> None: - ev = _prompt_event() + ev = _gateway_prompt_event() assert ev.case == TelemetryCase.PROMPT_GENERATION assert ev.signal_type is SignalType.TRACE assert ev.ce_eligible is False assert ev.trace_task_name is TraceTaskName.PROMPT_GENERATION_TRACE def test_app_created_routing(self) -> None: - ev = _app_created_event() + ev = _gateway_app_created_event() assert ev.case == TelemetryCase.APP_CREATED assert ev.signal_type is SignalType.METRIC_LOG assert ev.ce_eligible is False assert ev.trace_task_name is None def test_all_events_satisfy_protocol(self) -> None: - events: list[TelemetryEvent] = [_node_event(), _prompt_event(), _app_created_event()] + events: list[TelemetryEvent] = [_gateway_node_event(), _gateway_prompt_event(), _gateway_app_created_event()] for ev in events: assert isinstance(ev, TelemetryEvent) @@ -109,7 +258,7 @@ class TestGatewayTraceRouting: mock_ee_enabled: MagicMock, mock_trace_manager: MagicMock, ) -> None: - emit(_prompt_event(), mock_trace_manager) + emit(_gateway_prompt_event(), mock_trace_manager) mock_trace_manager.add_trace_task.assert_called_once() @patch("core.telemetry.gateway.is_enterprise_telemetry_enabled", return_value=False) @@ -118,7 +267,7 @@ class TestGatewayTraceRouting: mock_ee_enabled: MagicMock, mock_trace_manager: MagicMock, ) -> None: - emit(_node_event(), mock_trace_manager) + emit(_gateway_node_event(), mock_trace_manager) mock_trace_manager.add_trace_task.assert_not_called() @patch("core.telemetry.gateway.is_enterprise_telemetry_enabled", return_value=True) @@ -127,7 +276,7 @@ class TestGatewayTraceRouting: mock_ee_enabled: MagicMock, mock_trace_manager: MagicMock, ) -> None: - emit(_node_event(), mock_trace_manager) + emit(_gateway_node_event(), mock_trace_manager) mock_trace_manager.add_trace_task.assert_called_once() @@ -144,7 +293,7 @@ class TestGatewayMetricLogRouting: mock_delay: MagicMock, mock_ee_enabled: MagicMock, ) -> None: - emit(_app_created_event()) + emit(_gateway_app_created_event()) mock_delay.assert_called_once() envelope_json = mock_delay.call_args[0][0] @@ -160,8 +309,8 @@ class TestGatewayMetricLogRouting: mock_delay: MagicMock, mock_ee_enabled: MagicMock, ) -> None: - emit(_app_created_event()) - emit(_app_created_event()) + emit(_gateway_app_created_event()) + emit(_gateway_app_created_event()) assert mock_delay.call_count == 2 envelope1 = TelemetryEnvelope.model_validate_json(mock_delay.call_args_list[0][0][0]) @@ -182,7 +331,7 @@ class TestGatewayPayloadSizing: mock_delay: MagicMock, mock_ee_enabled: MagicMock, ) -> None: - emit(_app_created_event()) + emit(_gateway_app_created_event()) envelope_json = mock_delay.call_args[0][0] envelope = TelemetryEnvelope.model_validate_json(envelope_json) diff --git a/api/tests/unit_tests/core/telemetry/test_gateway_integration.py b/api/tests/unit_tests/core/telemetry/test_gateway_integration.py deleted file mode 100644 index 0e2704ab7a0..00000000000 --- a/api/tests/unit_tests/core/telemetry/test_gateway_integration.py +++ /dev/null @@ -1,165 +0,0 @@ -from __future__ import annotations - -import sys -from unittest.mock import MagicMock, patch - -import pytest - -from core.telemetry.events import ( - AppCreatedEvent, - AppCreatedPayload, - DraftNodeExecutionTraceEvent, - NodeExecutionPayload, - PromptGenerationEvent, - PromptGenerationPayload, - TelemetryContext, -) -from core.telemetry.gateway import emit, is_enterprise_telemetry_enabled -from enterprise.telemetry.contracts import TelemetryCase - - -class TestTelemetryCoreExports: - def test_is_enterprise_telemetry_enabled_exported(self) -> None: - from core.telemetry.gateway import is_enterprise_telemetry_enabled as exported_func - - assert callable(exported_func) - - -# --------------------------------------------------------------------------- -# Helpers -# --------------------------------------------------------------------------- - -_TRACE_CTX = TelemetryContext(tenant_id="tenant-789", user_id="user-456", app_id="app-123") - - -def _node_event() -> DraftNodeExecutionTraceEvent: - return DraftNodeExecutionTraceEvent( - context=_TRACE_CTX, - payload=NodeExecutionPayload(node_execution_data={}), - ) - - -def _prompt_event() -> PromptGenerationEvent: - return PromptGenerationEvent( - context=_TRACE_CTX, - payload=PromptGenerationPayload( - tenant_id="tenant-789", - operation_type="generate", - instruction="test", - generated_output="out", - model_provider="openai", - model_name="gpt-4", - prompt_tokens=10, - completion_tokens=5, - total_tokens=15, - latency=0.5, - ), - ) - - -def _app_created_event() -> AppCreatedEvent: - return AppCreatedEvent( - context=TelemetryContext(tenant_id="tenant-123"), - payload=AppCreatedPayload(app_id="app-abc"), - ) - - -# --------------------------------------------------------------------------- -# Trace routing integration -# --------------------------------------------------------------------------- - - -class TestGatewayIntegrationTraceRouting: - @pytest.fixture - def mock_trace_manager(self) -> MagicMock: - return MagicMock() - - def test_trace_event_routed_to_trace_manager( - self, - mock_trace_manager: MagicMock, - ) -> None: - with patch("core.telemetry.gateway.is_enterprise_telemetry_enabled", return_value=True): - emit(_prompt_event(), mock_trace_manager) - mock_trace_manager.add_trace_task.assert_called_once() - - def test_enterprise_only_trace_dropped_when_ee_disabled( - self, - mock_trace_manager: MagicMock, - ) -> None: - with patch("core.telemetry.gateway.is_enterprise_telemetry_enabled", return_value=False): - emit(_node_event(), mock_trace_manager) - mock_trace_manager.add_trace_task.assert_not_called() - - def test_enterprise_only_trace_routed_when_ee_enabled( - self, - mock_trace_manager: MagicMock, - ) -> None: - with patch("core.telemetry.gateway.is_enterprise_telemetry_enabled", return_value=True): - emit(_node_event(), mock_trace_manager) - mock_trace_manager.add_trace_task.assert_called_once() - - -# --------------------------------------------------------------------------- -# Metric/log routing integration -# --------------------------------------------------------------------------- - - -class TestGatewayIntegrationMetricRouting: - @patch("core.telemetry.gateway.is_enterprise_telemetry_enabled", return_value=True) - def test_metric_event_routes_to_celery_task( - self, - mock_ee_enabled: MagicMock, - ) -> None: - from enterprise.telemetry.contracts import TelemetryEnvelope - - with patch("tasks.enterprise_telemetry_task.process_enterprise_telemetry.delay") as mock_delay: - emit(_app_created_event()) - - mock_delay.assert_called_once() - envelope_json = mock_delay.call_args[0][0] - envelope = TelemetryEnvelope.model_validate_json(envelope_json) - assert envelope.case == TelemetryCase.APP_CREATED - assert envelope.tenant_id == "tenant-123" - assert envelope.payload["app_id"] == "app-abc" - - -# --------------------------------------------------------------------------- -# CE eligibility integration -# --------------------------------------------------------------------------- - - -class TestGatewayIntegrationCEEligibility: - @pytest.fixture - def mock_trace_manager(self) -> MagicMock: - return MagicMock() - - def test_draft_node_execution_not_ce_eligible( - self, - mock_trace_manager: MagicMock, - ) -> None: - with patch("core.telemetry.gateway.is_enterprise_telemetry_enabled", return_value=False): - emit(_node_event(), mock_trace_manager) - mock_trace_manager.add_trace_task.assert_not_called() - - def test_prompt_generation_not_ce_eligible( - self, - mock_trace_manager: MagicMock, - ) -> None: - with patch("core.telemetry.gateway.is_enterprise_telemetry_enabled", return_value=False): - emit(_prompt_event(), mock_trace_manager) - mock_trace_manager.add_trace_task.assert_not_called() - - -# --------------------------------------------------------------------------- -# is_enterprise_telemetry_enabled -# --------------------------------------------------------------------------- - - -class TestIsEnterpriseTelemetryEnabled: - def test_returns_false_when_exporter_import_fails(self) -> None: - with patch.dict(sys.modules, {"enterprise.telemetry.exporter": None}): - result = is_enterprise_telemetry_enabled() - assert result is False - - def test_function_is_callable(self) -> None: - assert callable(is_enterprise_telemetry_enabled) diff --git a/api/tests/unit_tests/core/test_provider_configuration.py b/api/tests/unit_tests/core/test_provider_configuration.py deleted file mode 100644 index d6b9276984a..00000000000 --- a/api/tests/unit_tests/core/test_provider_configuration.py +++ /dev/null @@ -1,535 +0,0 @@ -"""Provider-configuration behavior with persisted SQLite credential records.""" - -from types import SimpleNamespace -from unittest.mock import MagicMock, Mock, patch - -import pytest -from sqlalchemy.orm import Session - -from core.entities import provider_configuration as provider_configuration_module -from core.entities.provider_configuration import ProviderConfiguration, SystemConfigurationStatus -from core.entities.provider_entities import ( - CustomConfiguration, - ModelSettings, - ProviderQuotaType, - QuotaConfiguration, - QuotaUnit, - RestrictModel, - SystemConfiguration, -) -from graphon.model_runtime.entities.common_entities import I18nObject -from graphon.model_runtime.entities.model_entities import ModelType -from graphon.model_runtime.entities.provider_entities import ( - ConfigurateMethod, - CredentialFormSchema, - FormOption, - FormType, - ProviderEntity, -) -from models.provider import Provider, ProviderCredential, ProviderType - - -@pytest.fixture -def mock_provider_entity() -> ProviderEntity: - """Mock provider entity with basic configuration""" - provider_entity = ProviderEntity( - provider="openai", - label=I18nObject(en_US="OpenAI", zh_Hans="OpenAI"), - description=I18nObject(en_US="OpenAI provider", zh_Hans="OpenAI 提供商"), - icon_small=I18nObject(en_US="icon.png", zh_Hans="icon.png"), - background="background.png", - help=None, - supported_model_types=[ModelType.LLM], - configurate_methods=[ConfigurateMethod.PREDEFINED_MODEL], - provider_credential_schema=None, - model_credential_schema=None, - ) - - return provider_entity - - -@pytest.fixture -def mock_system_configuration() -> SystemConfiguration: - """Mock system configuration""" - quota_config = QuotaConfiguration( - quota_type=ProviderQuotaType.TRIAL, - quota_unit=QuotaUnit.TOKENS, - quota_limit=1000, - quota_used=0, - is_valid=True, - restrict_models=[RestrictModel(model="gpt-4", reason="Experimental", model_type=ModelType.LLM)], - ) - - system_config = SystemConfiguration( - enabled=True, - credentials={"openai_api_key": "test_key"}, - quota_configurations=[quota_config], - current_quota_type=ProviderQuotaType.TRIAL, - ) - - return system_config - - -@pytest.fixture -def mock_custom_configuration() -> CustomConfiguration: - """Mock custom configuration""" - custom_config = CustomConfiguration(provider=None, models=[]) - return custom_config - - -@pytest.fixture -def provider_configuration( - mock_provider_entity: ProviderEntity, - mock_system_configuration: SystemConfiguration, - mock_custom_configuration: CustomConfiguration, -) -> ProviderConfiguration: - """Create a test provider configuration instance""" - with patch( - "core.entities.provider_configuration.original_provider_configurate_methods", - dict[str, list[ConfigurateMethod]](), - ): - return ProviderConfiguration( - tenant_id="test_tenant", - provider=mock_provider_entity, - preferred_provider_type=ProviderType.SYSTEM, - using_provider_type=ProviderType.SYSTEM, - system_configuration=mock_system_configuration, - custom_configuration=mock_custom_configuration, - model_settings=[], - ) - - -class TestProviderConfiguration: - """Test cases for ProviderConfiguration class""" - - def test_get_current_credentials_system_provider_success( - self, provider_configuration: ProviderConfiguration - ) -> None: - """Test successfully getting credentials from system provider""" - # Arrange - provider_configuration.using_provider_type = ProviderType.SYSTEM - - # Act - credentials = provider_configuration.get_current_credentials(ModelType.LLM, "gpt-4") - - # Assert - assert credentials == {"openai_api_key": "test_key"} - - def test_get_current_credentials_model_disabled(self, provider_configuration: ProviderConfiguration) -> None: - """Test getting credentials when model is disabled""" - # Arrange - model_setting = ModelSettings( - model="gpt-4", - model_type=ModelType.LLM, - enabled=False, - load_balancing_configs=[], - has_invalid_load_balancing_configs=False, - ) - provider_configuration.model_settings = [model_setting] - - # Act & Assert - with pytest.raises(ValueError, match="Model gpt-4 is disabled"): - provider_configuration.get_current_credentials(ModelType.LLM, "gpt-4") - - def test_get_current_credentials_custom_provider_with_models( - self, provider_configuration: ProviderConfiguration - ) -> None: - """Test getting credentials from custom provider with model configurations""" - # Arrange - provider_configuration.using_provider_type = ProviderType.CUSTOM - - mock_model_config = Mock() - mock_model_config.model_type = ModelType.LLM - mock_model_config.model = "gpt-4" - mock_model_config.credentials = {"openai_api_key": "custom_key"} - provider_configuration.custom_configuration.models = [mock_model_config] - - # Act - credentials = provider_configuration.get_current_credentials(ModelType.LLM, "gpt-4") - - # Assert - assert credentials == {"openai_api_key": "custom_key"} - - def test_get_system_configuration_status_active(self, provider_configuration: ProviderConfiguration) -> None: - """Test getting active system configuration status""" - # Arrange - provider_configuration.system_configuration.enabled = True - - # Act - status = provider_configuration.get_system_configuration_status() - - # Assert - assert status == SystemConfigurationStatus.ACTIVE - - def test_get_system_configuration_status_unsupported(self, provider_configuration: ProviderConfiguration) -> None: - """Test getting unsupported system configuration status""" - # Arrange - provider_configuration.system_configuration.enabled = False - - # Act - status = provider_configuration.get_system_configuration_status() - - # Assert - assert status == SystemConfigurationStatus.UNSUPPORTED - - def test_get_system_configuration_status_quota_exceeded( - self, provider_configuration: ProviderConfiguration - ) -> None: - """Test getting quota exceeded system configuration status""" - # Arrange - provider_configuration.system_configuration.enabled = True - quota_config = provider_configuration.system_configuration.quota_configurations[0] - quota_config.is_valid = False - - # Act - status = provider_configuration.get_system_configuration_status() - - # Assert - assert status == SystemConfigurationStatus.QUOTA_EXCEEDED - - def test_is_custom_configuration_available_with_provider( - self, provider_configuration: ProviderConfiguration - ) -> None: - """Test custom configuration availability with provider credentials""" - # Arrange - mock_provider = Mock() - mock_provider.available_credentials = ["openai_api_key"] - provider_configuration.custom_configuration.provider = mock_provider - provider_configuration.custom_configuration.models = [] - - # Act - result = provider_configuration.is_custom_configuration_available() - - # Assert - assert result is True - - def test_is_custom_configuration_available_with_models(self, provider_configuration: ProviderConfiguration) -> None: - """Test custom configuration availability with model configurations""" - # Arrange - provider_configuration.custom_configuration.provider = None - provider_configuration.custom_configuration.models = [Mock()] - - # Act - result = provider_configuration.is_custom_configuration_available() - - # Assert - assert result is True - - def test_is_custom_configuration_available_false(self, provider_configuration: ProviderConfiguration) -> None: - """Test custom configuration not available""" - # Arrange - provider_configuration.custom_configuration.provider = None - provider_configuration.custom_configuration.models = [] - - # Act - result = provider_configuration.is_custom_configuration_available() - - # Assert - assert result is False - - @pytest.mark.parametrize("sqlite_session", [(Provider,)], indirect=True) - def test_get_provider_record_found( - self, provider_configuration: ProviderConfiguration, sqlite_session: Session - ) -> None: - """Test getting provider record successfully""" - provider = Provider(tenant_id="test_tenant", provider_name="openai") - sqlite_session.add(provider) - sqlite_session.commit() - - result = provider_configuration._get_provider_record(sqlite_session) - - assert result is provider - - @pytest.mark.parametrize("sqlite_session", [(Provider,)], indirect=True) - def test_get_provider_record_not_found( - self, provider_configuration: ProviderConfiguration, sqlite_session: Session - ) -> None: - """Test getting provider record when not found""" - sqlite_session.add(Provider(tenant_id="other_tenant", provider_name="openai")) - sqlite_session.commit() - - result = provider_configuration._get_provider_record(sqlite_session) - - assert result is None - - def test_init_with_customizable_model_only( - self, - mock_provider_entity: ProviderEntity, - mock_system_configuration: SystemConfiguration, - mock_custom_configuration: CustomConfiguration, - ) -> None: - """Test initialization with customizable model only configuration""" - # Arrange - mock_provider_entity.configurate_methods = [ConfigurateMethod.CUSTOMIZABLE_MODEL] - - # Act - with patch( - "core.entities.provider_configuration.original_provider_configurate_methods", - dict[str, list[ConfigurateMethod]](), - ): - config = ProviderConfiguration( - tenant_id="test_tenant", - provider=mock_provider_entity, - preferred_provider_type=ProviderType.SYSTEM, - using_provider_type=ProviderType.SYSTEM, - system_configuration=mock_system_configuration, - custom_configuration=mock_custom_configuration, - model_settings=[], - ) - - # Assert - assert ConfigurateMethod.PREDEFINED_MODEL in config.provider.configurate_methods - assert config.provider is not mock_provider_entity - assert mock_provider_entity.configurate_methods == [ConfigurateMethod.CUSTOMIZABLE_MODEL] - - def test_get_current_credentials_with_restricted_models( - self, provider_configuration: ProviderConfiguration - ) -> None: - """Test getting credentials with model restrictions""" - # Arrange - provider_configuration.using_provider_type = ProviderType.SYSTEM - - # Act - credentials = provider_configuration.get_current_credentials(ModelType.LLM, "gpt-3.5-turbo") - - # Assert - assert credentials is not None - assert "openai_api_key" in credentials - - @pytest.mark.parametrize("sqlite_session", [(Provider, ProviderCredential)], indirect=True) - def test_get_specific_provider_credential_success( - self, monkeypatch: pytest.MonkeyPatch, provider_configuration: ProviderConfiguration, sqlite_session: Session - ) -> None: - """Test getting specific provider credential successfully""" - credential_id = "test_credential_id" - credential = ProviderCredential( - tenant_id="test_tenant", - provider_name="openai", - credential_name="primary", - encrypted_config='{"openai_api_key": "encrypted_key"}', - ) - credential.id = credential_id - sqlite_session.add_all((Provider(tenant_id="test_tenant", provider_name="openai"), credential)) - sqlite_session.commit() - monkeypatch.setattr( - provider_configuration_module, - "db", - SimpleNamespace(engine=sqlite_session.get_bind()), - ) - - result = provider_configuration._get_specific_provider_credential(credential_id) - - assert result == {"openai_api_key": "encrypted_key"} - - @pytest.mark.parametrize("sqlite_session", [(Provider, ProviderCredential)], indirect=True) - def test_get_specific_provider_credential_with_mixed_provider_names( - self, monkeypatch: pytest.MonkeyPatch, provider_configuration: ProviderConfiguration, sqlite_session: Session - ) -> None: - provider_configuration.provider.provider = "langgenius/deepseek/deepseek" - credential = ProviderCredential( - tenant_id="test_tenant", - provider_name="langgenius/deepseek/deepseek", - credential_name="primary", - encrypted_config='{"api_key": "encrypted_key"}', - ) - sqlite_session.add_all((Provider(tenant_id="test_tenant", provider_name="deepseek"), credential)) - sqlite_session.commit() - monkeypatch.setattr( - provider_configuration_module, - "db", - SimpleNamespace(engine=sqlite_session.get_bind()), - ) - - result = provider_configuration._get_specific_provider_credential(credential.id) - - assert result == {"api_key": "encrypted_key"} - - @pytest.mark.parametrize("sqlite_session", [(Provider, ProviderCredential)], indirect=True) - def test_get_specific_provider_credential_not_found( - self, monkeypatch: pytest.MonkeyPatch, provider_configuration: ProviderConfiguration, sqlite_session: Session - ) -> None: - """Test getting specific provider credential when not found""" - credential_id = "nonexistent_credential_id" - sqlite_session.add(Provider(tenant_id="other_tenant", provider_name="openai")) - sqlite_session.commit() - monkeypatch.setattr( - provider_configuration_module, - "db", - SimpleNamespace(engine=sqlite_session.get_bind()), - ) - - with pytest.raises(ValueError, match=credential_id): - provider_configuration._get_specific_provider_credential(credential_id) - - def test_extract_secret_variables_with_secret_input(self, provider_configuration: ProviderConfiguration) -> None: - """Test extracting secret variables from credential form schemas""" - # Arrange - credential_form_schemas = [ - CredentialFormSchema( - variable="api_key", - label=I18nObject(en_US="API Key", zh_Hans="API 密钥"), - type=FormType.SECRET_INPUT, - required=True, - ), - CredentialFormSchema( - variable="model_name", - label=I18nObject(en_US="Model Name", zh_Hans="模型名称"), - type=FormType.TEXT_INPUT, - required=True, - ), - CredentialFormSchema( - variable="secret_token", - label=I18nObject(en_US="Secret Token", zh_Hans="密钥令牌"), - type=FormType.SECRET_INPUT, - required=False, - ), - ] - - # Act - secret_variables = provider_configuration.extract_secret_variables(credential_form_schemas) - - # Assert - assert len(secret_variables) == 2 - assert "api_key" in secret_variables - assert "secret_token" in secret_variables - assert "model_name" not in secret_variables - - def test_extract_secret_variables_no_secret_input(self, provider_configuration: ProviderConfiguration) -> None: - """Test extracting secret variables when no secret input fields exist""" - # Arrange - credential_form_schemas = [ - CredentialFormSchema( - variable="model_name", - label=I18nObject(en_US="Model Name", zh_Hans="模型名称"), - type=FormType.TEXT_INPUT, - required=True, - ), - CredentialFormSchema( - variable="temperature", - label=I18nObject(en_US="Temperature", zh_Hans="温度"), - type=FormType.SELECT, - required=True, - options=[FormOption(label=I18nObject(en_US="0.1", zh_Hans="0.1"), value="0.1")], - ), - ] - - # Act - secret_variables = provider_configuration.extract_secret_variables(credential_form_schemas) - - # Assert - assert len(secret_variables) == 0 - - def test_extract_secret_variables_empty_list(self, provider_configuration: ProviderConfiguration) -> None: - """Test extracting secret variables from empty credential form schemas""" - # Arrange - credential_form_schemas = [] - - # Act - secret_variables = provider_configuration.extract_secret_variables(credential_form_schemas) - - # Assert - assert len(secret_variables) == 0 - - @patch("core.entities.provider_configuration.encrypter") - def test_obfuscated_credentials_with_secret_variables( - self, mock_encrypter: MagicMock, provider_configuration: ProviderConfiguration - ) -> None: - """Test obfuscating credentials with secret variables""" - # Arrange - credentials = { - "api_key": "sk-1234567890abcdef", - "model_name": "gpt-4", - "secret_token": "secret_value_123", - "temperature": "0.7", - } - - credential_form_schemas = [ - CredentialFormSchema( - variable="api_key", - label=I18nObject(en_US="API Key", zh_Hans="API 密钥"), - type=FormType.SECRET_INPUT, - required=True, - ), - CredentialFormSchema( - variable="model_name", - label=I18nObject(en_US="Model Name", zh_Hans="模型名称"), - type=FormType.TEXT_INPUT, - required=True, - ), - CredentialFormSchema( - variable="secret_token", - label=I18nObject(en_US="Secret Token", zh_Hans="密钥令牌"), - type=FormType.SECRET_INPUT, - required=False, - ), - CredentialFormSchema( - variable="temperature", - label=I18nObject(en_US="Temperature", zh_Hans="温度"), - type=FormType.TEXT_INPUT, - required=True, - ), - ] - - mock_encrypter.obfuscated_token.side_effect = lambda x: f"***{x[-4:]}" - - # Act - obfuscated = provider_configuration.obfuscated_credentials(credentials, credential_form_schemas) - - # Assert - assert obfuscated["api_key"] == "***cdef" - assert obfuscated["model_name"] == "gpt-4" # Not obfuscated - assert obfuscated["secret_token"] == "***_123" - assert obfuscated["temperature"] == "0.7" # Not obfuscated - - # Verify encrypter was called for secret fields only - assert mock_encrypter.obfuscated_token.call_count == 2 - mock_encrypter.obfuscated_token.assert_any_call("sk-1234567890abcdef") - mock_encrypter.obfuscated_token.assert_any_call("secret_value_123") - - def test_obfuscated_credentials_no_secret_variables(self, provider_configuration: ProviderConfiguration) -> None: - """Test obfuscating credentials when no secret variables exist""" - # Arrange - credentials = { - "model_name": "gpt-4", - "temperature": "0.7", - "max_tokens": "1000", - } - - credential_form_schemas = [ - CredentialFormSchema( - variable="model_name", - label=I18nObject(en_US="Model Name", zh_Hans="模型名称"), - type=FormType.TEXT_INPUT, - required=True, - ), - CredentialFormSchema( - variable="temperature", - label=I18nObject(en_US="Temperature", zh_Hans="温度"), - type=FormType.TEXT_INPUT, - required=True, - ), - CredentialFormSchema( - variable="max_tokens", - label=I18nObject(en_US="Max Tokens", zh_Hans="最大令牌数"), - type=FormType.TEXT_INPUT, - required=True, - ), - ] - - # Act - obfuscated = provider_configuration.obfuscated_credentials(credentials, credential_form_schemas) - - # Assert - assert obfuscated == credentials # No changes expected - - def test_obfuscated_credentials_empty_credentials(self, provider_configuration: ProviderConfiguration) -> None: - """Test obfuscating empty credentials""" - # Arrange - credentials = {} - credential_form_schemas = [] - - # Act - obfuscated = provider_configuration.obfuscated_credentials(credentials, credential_form_schemas) - - # Assert - assert obfuscated == {} diff --git a/api/tests/unit_tests/core/test_provider_manager.py b/api/tests/unit_tests/core/test_provider_manager.py index 400281cce15..3566a92f23f 100644 --- a/api/tests/unit_tests/core/test_provider_manager.py +++ b/api/tests/unit_tests/core/test_provider_manager.py @@ -42,7 +42,7 @@ from models.provider import ( TenantPreferredModelProvider, ) from models.provider_ids import ModelProviderID -from tests.unit_tests.config_override import config_overrides_context +from tests.helpers.unit.config_override import config_overrides_context def _build_provider_manager() -> ProviderManager: diff --git a/api/tests/unit_tests/core/test_trigger_debug_event_selectors.py b/api/tests/unit_tests/core/test_trigger_debug_event_selectors.py deleted file mode 100644 index 2b508ca6545..00000000000 --- a/api/tests/unit_tests/core/test_trigger_debug_event_selectors.py +++ /dev/null @@ -1,102 +0,0 @@ -import hashlib -import json -from datetime import UTC, datetime - -import pytest -import pytz - -from core.trigger.debug import event_selectors -from core.workflow.nodes.trigger_schedule.entities import ScheduleConfig - - -class _DummyRedis: - def __init__(self): - self.store: dict[str, str] = {} - - def get(self, key: str): - return self.store.get(key) - - def setex(self, name: str, time: int, value: str): - self.store[name] = value - - def expire(self, name: str, ttl: int): - # Expiration not required for these tests. - pass - - def delete(self, name: str): - self.store.pop(name, None) - - -@pytest.fixture -def dummy_schedule_config() -> ScheduleConfig: - return ScheduleConfig( - node_id="node-1", - cron_expression="* * * * *", - timezone="Asia/Shanghai", - ) - - -@pytest.fixture(autouse=True) -def patch_schedule_service(monkeypatch: pytest.MonkeyPatch, dummy_schedule_config: ScheduleConfig): - # Ensure poller always receives the deterministic config. - monkeypatch.setattr( - "services.trigger.schedule_service.ScheduleService.to_schedule_config", - staticmethod(lambda *_args, **_kwargs: dummy_schedule_config), - ) - - -def _make_poller( - monkeypatch: pytest.MonkeyPatch, redis_client: _DummyRedis -) -> event_selectors.ScheduleTriggerDebugEventPoller: - monkeypatch.setattr(event_selectors, "redis_client", redis_client) - return event_selectors.ScheduleTriggerDebugEventPoller( - tenant_id="tenant-1", - user_id="user-1", - app_id="app-1", - node_config={"id": "node-1", "data": {"mode": "cron"}}, - node_id="node-1", - ) - - -def test_schedule_poller_handles_aware_next_run(monkeypatch: pytest.MonkeyPatch): - redis_client = _DummyRedis() - poller = _make_poller(monkeypatch, redis_client) - - base_now = datetime(2025, 1, 1, 12, 0, 10) - aware_next_run = datetime(2025, 1, 1, 12, 0, 5, tzinfo=UTC) - - monkeypatch.setattr(event_selectors, "naive_utc_now", lambda: base_now) - monkeypatch.setattr(event_selectors, "calculate_next_run_at", lambda *_: aware_next_run) - - event = poller.poll() - - assert event is not None - assert event.node_id == "node-1" - assert event.workflow_args["inputs"] == {} - - -def test_schedule_runtime_cache_normalizes_timezone( - monkeypatch: pytest.MonkeyPatch, dummy_schedule_config: ScheduleConfig -): - redis_client = _DummyRedis() - poller = _make_poller(monkeypatch, redis_client) - - localized_time = pytz.timezone("Asia/Shanghai").localize(datetime(2025, 1, 1, 20, 0, 0)) - - cron_hash = hashlib.sha256(dummy_schedule_config.cron_expression.encode()).hexdigest() - cache_key = poller.schedule_debug_runtime_key(cron_hash) - - redis_client.store[cache_key] = json.dumps( - { - "cache_key": cache_key, - "timezone": dummy_schedule_config.timezone, - "cron_expression": dummy_schedule_config.cron_expression, - "next_run_at": localized_time.isoformat(), - } - ) - - runtime = poller.get_or_create_schedule_debug_runtime() - - expected = localized_time.astimezone(UTC).replace(tzinfo=None) - assert runtime.next_run_at == expected - assert runtime.next_run_at.tzinfo is None diff --git a/api/tests/unit_tests/core/tools/test_base_tool.py b/api/tests/unit_tests/core/tools/__base/test_tool.py similarity index 100% rename from api/tests/unit_tests/core/tools/test_base_tool.py rename to api/tests/unit_tests/core/tools/__base/test_tool.py diff --git a/api/tests/unit_tests/core/tools/test_tool_provider_controller.py b/api/tests/unit_tests/core/tools/__base/test_tool_provider.py similarity index 100% rename from api/tests/unit_tests/core/tools/test_tool_provider_controller.py rename to api/tests/unit_tests/core/tools/__base/test_tool_provider.py diff --git a/api/tests/unit_tests/core/tools/test_builtin_tools_extra.py b/api/tests/unit_tests/core/tools/builtin_tool/providers/time/tools/test_localtime_to_timestamp.py similarity index 100% rename from api/tests/unit_tests/core/tools/test_builtin_tools_extra.py rename to api/tests/unit_tests/core/tools/builtin_tool/providers/time/tools/test_localtime_to_timestamp.py diff --git a/api/tests/unit_tests/core/tools/test_builtin_tool_provider.py b/api/tests/unit_tests/core/tools/builtin_tool/test_provider.py similarity index 100% rename from api/tests/unit_tests/core/tools/test_builtin_tool_provider.py rename to api/tests/unit_tests/core/tools/builtin_tool/test_provider.py diff --git a/api/tests/unit_tests/core/tools/test_builtin_tool_base.py b/api/tests/unit_tests/core/tools/builtin_tool/test_tool.py similarity index 100% rename from api/tests/unit_tests/core/tools/test_builtin_tool_base.py rename to api/tests/unit_tests/core/tools/builtin_tool/test_tool.py diff --git a/api/tests/unit_tests/core/tools/custom_tool/test_api_tool_request_assembly.py b/api/tests/unit_tests/core/tools/custom_tool/test_api_tool_request_assembly.py deleted file mode 100644 index 58c3130ac55..00000000000 --- a/api/tests/unit_tests/core/tools/custom_tool/test_api_tool_request_assembly.py +++ /dev/null @@ -1,56 +0,0 @@ -import pytest - -from core.tools.__base.tool_runtime import ToolRuntime -from core.tools.custom_tool.tool import ApiTool -from core.tools.entities.common_entities import I18nObject -from core.tools.entities.tool_bundle import ApiToolBundle -from core.tools.entities.tool_entities import ToolEntity, ToolIdentity - -pytest_plugins = ("tests.unit_tests.core.tools.custom_tool.api_tool_http_fixture",) - -tool_bundle: dict[str, object] = { - "server_url": "http://www.example.com/{path_param}", - "method": "post", - "author": "", - "openapi": { - "parameters": [ - {"in": "path", "name": "path_param"}, - {"in": "query", "name": "query_param"}, - {"in": "cookie", "name": "cookie_param"}, - {"in": "header", "name": "header_param"}, - ], - "requestBody": { - "content": {"application/json": {"schema": {"properties": {"body_param": {"type": "string"}}}}} - }, - }, - "parameters": [], -} -parameters = { - "path_param": "p_param", - "query_param": "q_param", - "cookie_param": "c_param", - "header_param": "h_param", - "body_param": "b_param", -} - - -@pytest.mark.usefixtures("setup_http_mock") -def test_api_tool() -> None: - tool = ApiTool( - entity=ToolEntity( - identity=ToolIdentity(provider="", author="", name="", label=I18nObject(en_US="test tool")), - ), - api_bundle=ApiToolBundle.model_validate(tool_bundle), - runtime=ToolRuntime(tenant_id="", credentials={"auth_type": "none"}), - provider_id="test_tool", - ) - headers = tool.assembling_request(parameters) - response = tool.do_http_request(tool.api_bundle.server_url, tool.api_bundle.method, headers, parameters) - - assert response.status_code == 200 - assert response.request.url.path == "/p_param" - assert response.request.url.query == b"query_param=q_param" - assert response.request.headers.get("header_param") == "h_param" - assert response.request.headers.get("content-type") == "application/json" - assert response.request.headers.get("cookie") == "cookie_param=c_param" - assert "b_param" in response.content.decode() diff --git a/api/tests/unit_tests/core/tools/test_custom_tool_provider.py b/api/tests/unit_tests/core/tools/custom_tool/test_provider.py similarity index 100% rename from api/tests/unit_tests/core/tools/test_custom_tool_provider.py rename to api/tests/unit_tests/core/tools/custom_tool/test_provider.py diff --git a/api/tests/unit_tests/core/tools/custom_tool/test_tool.py b/api/tests/unit_tests/core/tools/custom_tool/test_tool.py new file mode 100644 index 00000000000..ae849505138 --- /dev/null +++ b/api/tests/unit_tests/core/tools/custom_tool/test_tool.py @@ -0,0 +1,577 @@ +from __future__ import annotations + +import json +import operator +from collections.abc import Iterator +from types import SimpleNamespace +from typing import Any +from unittest.mock import Mock, patch + +import httpx +import pytest +from sqlalchemy import Engine +from sqlalchemy.orm import Session + +from core.app.entities.app_invoke_entities import InvokeFrom +from core.tools.__base.tool_runtime import ToolRuntime +from core.tools.custom_tool.tool import ApiTool, ParsedResponse +from core.tools.entities.common_entities import I18nObject +from core.tools.entities.tool_bundle import ApiToolBundle +from core.tools.entities.tool_entities import ( + ToolEntity, + ToolIdentity, + ToolInvokeMessage, +) +from core.tools.errors import ToolInvokeError, ToolParameterValidationError, ToolProviderCredentialValidationError + +pytest_plugins = ("tests.helpers.unit.core.tools.custom_tool.api_tool_http_fixture",) + + +tool_bundle: dict[str, object] = { + "server_url": "http://www.example.com/{path_param}", + "method": "post", + "author": "", + "openapi": { + "parameters": [ + {"in": "path", "name": "path_param"}, + {"in": "query", "name": "query_param"}, + {"in": "cookie", "name": "cookie_param"}, + {"in": "header", "name": "header_param"}, + ], + "requestBody": { + "content": {"application/json": {"schema": {"properties": {"body_param": {"type": "string"}}}}} + }, + }, + "parameters": [], +} +parameters = { + "path_param": "p_param", + "query_param": "q_param", + "cookie_param": "c_param", + "header_param": "h_param", + "body_param": "b_param", +} + + +class TestApiToolRequestAssemblyCases: + @staticmethod + @pytest.mark.usefixtures("setup_http_mock") + def test_api_tool() -> None: + tool = ApiTool( + entity=ToolEntity( + identity=ToolIdentity(provider="", author="", name="", label=I18nObject(en_US="test tool")), + ), + api_bundle=ApiToolBundle.model_validate(tool_bundle), + runtime=ToolRuntime(tenant_id="", credentials={"auth_type": "none"}), + provider_id="test_tool", + ) + headers = tool.assembling_request(parameters) + response = tool.do_http_request(tool.api_bundle.server_url, tool.api_bundle.method, headers, parameters) + + assert response.status_code == 200 + assert response.request.url.path == "/p_param" + assert response.request.url.query == b"query_param=q_param" + assert response.request.headers.get("header_param") == "h_param" + assert response.request.headers.get("content-type") == "application/json" + assert response.request.headers.get("cookie") == "cookie_param=c_param" + assert "b_param" in response.content.decode() + + +def _build_tool(*, openapi: dict[str, Any] | None = None) -> ApiTool: + entity = ToolEntity( + identity=ToolIdentity( + author="author", + name="tool-a", + label=I18nObject(en_US="tool-a"), + provider="provider-a", + ), + parameters=[], + ) + bundle = ApiToolBundle( + server_url="https://api.example.com/items/{id}", + method="GET", + summary="summary", + operation_id="op-id", + parameters=[], + author="author", + openapi=openapi or {"parameters": []}, + ) + runtime = ToolRuntime( + tenant_id="tenant-1", + invoke_from=InvokeFrom.DEBUGGER, + credentials={"auth_type": "api_key_header", "api_key_value": "k"}, + ) + return ApiTool(entity=entity, api_bundle=bundle, runtime=runtime, provider_id="provider-id") + + +class TestCustomToolCases: + @staticmethod + def test_parsed_response_to_string(): + assert ParsedResponse({"a": 1}, True).to_string() == '{"a": 1}' + assert ParsedResponse("ok", False).to_string() == "ok" + + @staticmethod + def test_api_tool_fork_runtime_and_validate_credentials(monkeypatch: pytest.MonkeyPatch): + tool = _build_tool() + forked = tool.fork_tool_runtime(ToolRuntime(tenant_id="tenant-2")) + assert isinstance(forked, ApiTool) + assert forked.runtime.tenant_id == "tenant-2" + + tool.api_bundle = None # type: ignore[assignment] + with pytest.raises(ValueError, match="api_bundle is required"): + tool.fork_tool_runtime(ToolRuntime(tenant_id="tenant-2")) + + tool = _build_tool() + assert tool.validate_credentials(credentials={}, parameters={}, format_only=True) == "" + monkeypatch.setattr(tool, "assembling_request", lambda parameters: {"Authorization": "Bearer x"}) + monkeypatch.setattr( + tool, + "do_http_request", + lambda url, method, headers, parameters: httpx.Response(200, json={"ok": True}), + ) + result = tool.validate_credentials(credentials={}, parameters={"a": 1}, format_only=False) + assert result == '{"ok": true}' + + @staticmethod + def test_assembling_request_auth_header_assembly(): + tool = _build_tool() + + headers = tool.assembling_request(parameters={}) + assert headers["Authorization"] == "k" + + tool.runtime.credentials = { + "auth_type": "api_key_header", + "api_key_header_prefix": "bearer", + "api_key_value": "abc", + } + headers = tool.assembling_request(parameters={}) + assert headers["Authorization"] == "Bearer abc" + + tool.runtime.credentials = { + "auth_type": "api_key_header", + "api_key_header_prefix": "basic", + "api_key_value": "abc", + } + headers = tool.assembling_request(parameters={}) + assert headers["Authorization"] == "Basic abc" + assert tool.runtime.credentials["api_key_value"] == "abc" + + headers = tool.assembling_request(parameters={}) + assert headers["Authorization"] == "Basic abc" + + tool.runtime.credentials = {"auth_type": "api_key_query", "api_key_value": "abc"} + assert tool.assembling_request(parameters={}) == {} + + @staticmethod + def test_assembling_request_runtime_auth_errors(): + tool = _build_tool() + + tool.runtime = None + with pytest.raises(ToolProviderCredentialValidationError, match="runtime not initialized"): + tool.assembling_request(parameters={}) + + tool.runtime = ToolRuntime(tenant_id="tenant", credentials={}) + with pytest.raises(ToolProviderCredentialValidationError, match="Missing auth_type"): + tool.assembling_request(parameters={}) + + tool.runtime.credentials = {"auth_type": "api_key_header"} + with pytest.raises(ToolProviderCredentialValidationError, match="Missing api_key_value"): + tool.assembling_request(parameters={}) + + tool.runtime.credentials = {"auth_type": "api_key_header", "api_key_value": 123} + with pytest.raises(ToolProviderCredentialValidationError, match="must be a string"): + tool.assembling_request(parameters={}) + + @staticmethod + def test_assembling_request_parameter_validation_and_defaults(): + tool = _build_tool() + + tool.runtime.credentials = {"auth_type": "api_key_header", "api_key_value": "x"} + tool.api_bundle.parameters = [ + SimpleNamespace(required=True, name="required_param", default=None), + ] + with pytest.raises(ToolParameterValidationError, match="Missing required parameter required_param"): + tool.assembling_request(parameters={}) + + tool.api_bundle.parameters = [ + SimpleNamespace(required=True, name="required_param", default="d"), + ] + params = {} + tool.assembling_request(parameters=params) + assert params["required_param"] == "d" + + @staticmethod + def test_validate_and_parse_response_branches(): + tool = _build_tool() + + with pytest.raises(ToolInvokeError, match="status code 500"): + tool.validate_and_parse_response(httpx.Response(500, text="boom")) + + empty = tool.validate_and_parse_response(httpx.Response(200, content=b"")) + assert empty.is_json is False + assert "Empty response from the tool" in str(empty.content) + + json_resp = tool.validate_and_parse_response( + httpx.Response(200, json={"a": 1}, headers={"content-type": "application/json"}) + ) + assert json_resp.is_json is True + assert json_resp.content == {"a": 1} + + non_json_type = tool.validate_and_parse_response( + httpx.Response(200, text='{"a": 1}', headers={"content-type": "text/plain"}) + ) + assert non_json_type.is_json is False + assert non_json_type.content == '{"a": 1}' + + plain_resp = tool.validate_and_parse_response(httpx.Response(200, text="plain")) + assert plain_resp.is_json is False + assert plain_resp.content == "plain" + + with pytest.raises(ValueError, match="Invalid response type"): + tool.validate_and_parse_response("invalid") # type: ignore[arg-type] + + @staticmethod + def test_get_parameter_value_and_type_conversion_helpers(): + tool = _build_tool() + + assert tool.get_parameter_value({"name": "x"}, {"x": 1}) == 1 + assert tool.get_parameter_value({"name": "x", "required": False, "schema": {"default": "d"}}, {}) == "d" + with pytest.raises(ToolParameterValidationError, match="Missing required parameter x"): + tool.get_parameter_value({"name": "x", "required": True}, {}) + + assert tool._convert_body_property_any_of({}, "12", [{"type": "integer"}]) == 12 + assert tool._convert_body_property_any_of({}, "1.5", [{"type": "number"}]) == 1.5 + assert tool._convert_body_property_any_of({}, "true", [{"type": "boolean"}]) is True + assert tool._convert_body_property_any_of({}, "", [{"type": "null"}]) is None + assert tool._convert_body_property_any_of({}, "x", [{"anyOf": [{"type": "string"}]}]) == "x" + + assert tool._convert_body_property_type({"type": "integer"}, "1") == 1 + assert tool._convert_body_property_type({"type": "number"}, "1.2") == 1.2 + assert tool._convert_body_property_type({"type": "string"}, 1) == "1" + assert tool._convert_body_property_type({"type": "boolean"}, 1) is True + assert tool._convert_body_property_type({"type": "null"}, None) is None + assert tool._convert_body_property_type({"type": "object"}, '{"a":1}') == {"a": 1} + assert tool._convert_body_property_type({"type": "array"}, "[1,2]") == [1, 2] + assert tool._convert_body_property_type({"type": "invalid"}, "v") == "v" + assert tool._convert_body_property_type({"anyOf": [{"type": "integer"}]}, "2") == 2 + + @staticmethod + def test_do_http_request_builds_arguments_and_handles_invalid_method(monkeypatch: pytest.MonkeyPatch): + openapi = { + "parameters": [ + {"name": "id", "in": "path", "required": True, "schema": {"type": "string"}}, + {"name": "q", "in": "query", "required": False, "schema": {"default": ""}}, + {"name": "X-Extra", "in": "header", "required": False, "schema": {"default": "x"}}, + {"name": "sid", "in": "cookie", "required": False, "schema": {"default": "cookie1"}}, + ], + "requestBody": { + "content": { + "application/json": { + "schema": { + "type": "object", + "required": ["count"], + "properties": { + "count": {"type": "integer"}, + "name": {"type": "string", "default": "n"}, + }, + } + } + } + }, + } + tool = _build_tool(openapi=openapi) + tool.runtime.credentials = {"auth_type": "api_key_query", "api_key_query_param": "key", "api_key_value": "v"} + headers = {} + captured = {} + + def _fake_get(url, **kwargs): + captured["url"] = url + captured["kwargs"] = kwargs + return httpx.Response(200, text="ok") + + monkeypatch.setattr("core.tools.custom_tool.tool.ssrf_proxy.get", _fake_get) + response = tool.do_http_request( + "https://api.example.com/items/{id}", + "GET", + headers=headers, + parameters={"id": "123", "count": "2", "q": "search"}, + ) + + assert isinstance(response, httpx.Response) + assert captured["url"].endswith("/items/123") + assert captured["kwargs"]["params"]["q"] == "search" + assert captured["kwargs"]["params"]["key"] == "v" + assert captured["kwargs"]["headers"]["Content-Type"] == "application/json" + + invalid_method_tool = _build_tool(openapi={"parameters": []}) + with pytest.raises(ValueError, match="Invalid http method"): + invalid_method_tool.do_http_request("https://api.example.com", "TRACE", headers={}, parameters={}) + + @staticmethod + def test_do_http_request_handles_file_upload_and_invoke_paths( + monkeypatch: pytest.MonkeyPatch, sqlite_session: Session + ): + openapi = { + "parameters": [], + "requestBody": { + "content": { + "multipart/form-data": { + "schema": { + "type": "object", + "properties": {"file": {"format": "binary"}}, + } + } + } + }, + } + tool = _build_tool(openapi=openapi) + tool.runtime.credentials = {"auth_type": "api_key_header", "api_key_value": "k"} + fake_file = SimpleNamespace(filename="a.txt", mime_type="text/plain") + captured = {} + + def _fake_post(url, **kwargs): + captured["headers"] = kwargs["headers"] + captured["files"] = kwargs["files"] + return httpx.Response(200, text="ok") + + monkeypatch.setattr("core.tools.custom_tool.tool.download", lambda _: b"file-bytes") + monkeypatch.setattr("core.tools.custom_tool.tool.ssrf_proxy.post", _fake_post) + response = tool.do_http_request( + "https://api.example.com/upload", + "POST", + headers={}, + parameters={"file": fake_file}, + ) + assert isinstance(response, httpx.Response) + assert "Content-Type" not in captured["headers"] + assert captured["files"][0][0] == "file" + + # _invoke JSON path + monkeypatch.setattr(tool, "assembling_request", lambda parameters: {}) + monkeypatch.setattr(tool, "do_http_request", lambda *args, **kwargs: httpx.Response(200, text='{"a":1}')) + monkeypatch.setattr(tool, "validate_and_parse_response", lambda _: ParsedResponse({"a": 1}, True)) + messages = list(tool.invoke(session=sqlite_session, user_id="u1", tool_parameters={})) + assert [m.type for m in messages] == [ToolInvokeMessage.MessageType.JSON, ToolInvokeMessage.MessageType.TEXT] + + # _invoke text path + monkeypatch.setattr(tool, "validate_and_parse_response", lambda _: ParsedResponse("plain", False)) + messages = list(tool.invoke(session=sqlite_session, user_id="u1", tool_parameters={})) + assert len(messages) == 1 + assert messages[0].message.text == "plain" + + +def _get_message_by_type[T](msgs: list[ToolInvokeMessage], msg_type: type[T]) -> ToolInvokeMessage | None: + return next((i for i in msgs if isinstance(i.message, msg_type)), None) + + +class TestApiToolCases: + @staticmethod + @pytest.fixture + def database_session(sqlite_engine: Engine) -> Iterator[Session]: + """Yield the live ORM session required by the shared tool invocation contract.""" + with Session(sqlite_engine) as session: + yield session + + class TestApiToolInvoke: + """Test suite for ApiTool._invoke method to ensure JSON responses are properly serialized.""" + + def setup_method(self): + """Setup test fixtures.""" + # Create a mock tool entity + self.mock_tool_identity = ToolIdentity( + author="test", + name="test_api_tool", + label=I18nObject(en_US="Test API Tool", zh_Hans="测试API工具"), + provider="test_provider", + ) + self.mock_tool_entity = ToolEntity(identity=self.mock_tool_identity) + + # Create a mock API bundle + self.mock_api_bundle = ApiToolBundle( + server_url="https://api.example.com/test", + method="GET", + openapi={}, + operation_id="test_operation", + parameters=[], + author="test_author", + ) + + self.runtime = ToolRuntime(tenant_id="test_tenant", credentials={"auth_type": "none"}) + + # Create the ApiTool instance + self.api_tool = ApiTool( + entity=self.mock_tool_entity, + api_bundle=self.mock_api_bundle, + runtime=self.runtime, + provider_id="test_provider", + ) + + @patch("core.tools.custom_tool.tool.ssrf_proxy.get") + def test_invoke_with_json_response_creates_text_message_with_serialized_json( + self, mock_get: Mock, database_session: Session + ) -> None: + """Test that when upstream returns JSON, the output Text message contains JSON-serialized string.""" + # Use the response's real JSON parsing and content decoding. + json_response_data = { + "key": "value", + "number": 123, + "nested": {"inner": "data"}, + } + mock_get.return_value = httpx.Response(200, json=json_response_data) + + # Invoke the tool + result_generator = self.api_tool._invoke(session=database_session, user_id="test_user", tool_parameters={}) + + # Get the result from the generator + result = list(result_generator) + assert len(result) == 2 + + # Verify _invoke yields text message + text_message = _get_message_by_type(result, ToolInvokeMessage.TextMessage) + assert text_message is not None, "_invoke should yield a text message" + assert isinstance(text_message, ToolInvokeMessage) + assert text_message.type == ToolInvokeMessage.MessageType.TEXT + assert text_message.message is not None + # Verify the text contains the JSON-serialized string + # Check if message is a TextMessage + assert isinstance(text_message.message, ToolInvokeMessage.TextMessage) + # Verify it's a valid JSON string and equals to the mock response + parsed_back = json.loads(text_message.message.text) + assert parsed_back == json_response_data + + # Verify _invoke yields json message + json_message = _get_message_by_type(result, ToolInvokeMessage.JsonMessage) + assert json_message is not None, "_invoke should yield a JSON message" + assert isinstance(json_message, ToolInvokeMessage) + assert json_message.type == ToolInvokeMessage.MessageType.JSON + assert json_message.message is not None + + assert isinstance(json_message.message, ToolInvokeMessage.JsonMessage) + assert json_message.message.json_object == json_response_data + + @patch("core.tools.custom_tool.tool.ssrf_proxy.get") + @pytest.mark.parametrize( + "test_case", + [ + ( + "array", + [ + {"id": 1, "name": "Item 1", "active": True}, + {"id": 2, "name": "Item 2", "active": False}, + {"id": 3, "name": "项目 3", "active": True}, + ], + ), + ( + "string", + "string", + ), + ( + "number", + 123.456, + ), + ( + "boolean", + True, + ), + ( + "null", + None, + ), + ], + ids=operator.itemgetter(0), + ) + def test_invoke_with_non_dict_json_response_creates_text_message_with_serialized_json( + self, mock_get: Mock, test_case, database_session: Session + ) -> None: + """Test that when upstream returns a non-dict JSON, + the output Text message contains JSON-serialized string. + """ + # Explicit content also covers JSON null, which is distinct from no body. + _, json_value = test_case + mock_get.return_value = httpx.Response( + 200, content=json.dumps(json_value), headers={"content-type": "application/json"} + ) + + # Invoke the tool + result_generator = self.api_tool._invoke(session=database_session, user_id="test_user", tool_parameters={}) + + # Get the result from the generator + result = list(result_generator) + assert len(result) == 1 + + # Verify _invoke yields a text message + text_message = _get_message_by_type(result, ToolInvokeMessage.TextMessage) + assert text_message is not None, "_invoke should yield a text message containing the serialized JSON." + assert isinstance(text_message, ToolInvokeMessage) + assert text_message.type == ToolInvokeMessage.MessageType.TEXT + assert text_message.message is not None + # Verify the text contains the JSON-serialized string + # Check if message is a TextMessage + assert isinstance(text_message.message, ToolInvokeMessage.TextMessage) + # Verify it's a valid JSON string + parsed_back = json.loads(text_message.message.text) + assert parsed_back == json_value + + # Verify _invoke yields json message + json_message = _get_message_by_type(result, ToolInvokeMessage.JsonMessage) + assert json_message is None, "_invoke should not yield a JSON message for JSON array response" + + @patch("core.tools.custom_tool.tool.ssrf_proxy.get") + def test_invoke_with_text_response_creates_text_message_with_original_text( + self, mock_get: Mock, database_session: Session + ) -> None: + """Test that when upstream returns plain text, the output Text message contains the original text.""" + # Plain text exercises the response's real JSON decoding failure. + text_response_data = "This is a plain text response" + mock_get.return_value = httpx.Response(200, text=text_response_data) + + # Invoke the tool + result_generator = self.api_tool._invoke(session=database_session, user_id="test_user", tool_parameters={}) + + # Get the result from the generator + result = list(result_generator) + assert len(result) == 1 + + # Verify it's a text message with the original text + message = result[0] + assert isinstance(message, ToolInvokeMessage) + assert message.type == ToolInvokeMessage.MessageType.TEXT + assert message.message is not None + # Check if message is a TextMessage + assert isinstance(message.message, ToolInvokeMessage.TextMessage) + assert message.message.text == text_response_data + + @patch("core.tools.custom_tool.tool.ssrf_proxy.get") + def test_invoke_with_empty_response(self, mock_get: Mock, database_session: Session) -> None: + """Test that empty responses are handled correctly.""" + mock_get.return_value = httpx.Response(200, content=b"", headers={"content-type": "application/json"}) + + # Invoke the tool + result_generator = self.api_tool._invoke(session=database_session, user_id="test_user", tool_parameters={}) + + # Get the result from the generator + result = list(result_generator) + assert len(result) == 1 + + # Verify it's a text message with the empty response message + message = result[0] + assert isinstance(message, ToolInvokeMessage) + assert message.type == ToolInvokeMessage.MessageType.TEXT + assert message.message is not None + # Check if message is a TextMessage + assert isinstance(message.message, ToolInvokeMessage.TextMessage) + assert "Empty response from the tool" in message.message.text + + @patch("core.tools.custom_tool.tool.ssrf_proxy.get") + def test_invoke_with_error_response(self, mock_get: Mock, database_session: Session) -> None: + """Test that error responses are handled correctly.""" + mock_get.return_value = httpx.Response(404, text="Not Found") + + result_generator = self.api_tool._invoke(session=database_session, user_id="test_user", tool_parameters={}) + + # Invoke the tool and expect an error + with pytest.raises(Exception, match="Request failed with status code 404 and Not Found") as exc_info: + list(result_generator) # Consume the generator to trigger the error + + # Verify the error message + assert "Request failed with status code 404" in str(exc_info.value) diff --git a/api/tests/unit_tests/core/tools/test_tool_entities.py b/api/tests/unit_tests/core/tools/entities/test_tool_entities.py similarity index 53% rename from api/tests/unit_tests/core/tools/test_tool_entities.py rename to api/tests/unit_tests/core/tools/entities/test_tool_entities.py index 82ab59752c2..cd38809c7d3 100644 --- a/api/tests/unit_tests/core/tools/test_tool_entities.py +++ b/api/tests/unit_tests/core/tools/entities/test_tool_entities.py @@ -89,3 +89,51 @@ def test_tool_entity_output_schema_none_defaults_to_empty_dict(): entity = ToolEntity(identity=_make_identity(), output_schema=None) assert entity.output_schema == {} + + +def test_get_parameter_type(): + assert ToolParameter.ToolParameterType.STRING.as_normal_type() == "string" + assert ToolParameter.ToolParameterType.SELECT.as_normal_type() == "string" + assert ToolParameter.ToolParameterType.SECRET_INPUT.as_normal_type() == "string" + assert ToolParameter.ToolParameterType.BOOLEAN.as_normal_type() == "boolean" + assert ToolParameter.ToolParameterType.NUMBER.as_normal_type() == "number" + assert ToolParameter.ToolParameterType.FILE.as_normal_type() == "file" + assert ToolParameter.ToolParameterType.FILES.as_normal_type() == "files" + + +def test_cast_parameter_by_type(): + # string + assert ToolParameter.ToolParameterType.STRING.cast_value("test") == "test" + assert ToolParameter.ToolParameterType.STRING.cast_value(1) == "1" + assert ToolParameter.ToolParameterType.STRING.cast_value(1.0) == "1.0" + assert ToolParameter.ToolParameterType.STRING.cast_value(None) == "" + + # secret input + assert ToolParameter.ToolParameterType.SECRET_INPUT.cast_value("test") == "test" + assert ToolParameter.ToolParameterType.SECRET_INPUT.cast_value(1) == "1" + assert ToolParameter.ToolParameterType.SECRET_INPUT.cast_value(1.0) == "1.0" + assert ToolParameter.ToolParameterType.SECRET_INPUT.cast_value(None) == "" + + # select + assert ToolParameter.ToolParameterType.SELECT.cast_value("test") == "test" + assert ToolParameter.ToolParameterType.SELECT.cast_value(1) == "1" + assert ToolParameter.ToolParameterType.SELECT.cast_value(1.0) == "1.0" + assert ToolParameter.ToolParameterType.SELECT.cast_value(None) == "" + + # boolean + true_values = [True, "True", "true", "1", "YES", "Yes", "yes", "y", "something"] + for value in true_values: + assert ToolParameter.ToolParameterType.BOOLEAN.cast_value(value) is True + + false_values = [False, "False", "false", "0", "NO", "No", "no", "n", None, ""] + for value in false_values: + assert ToolParameter.ToolParameterType.BOOLEAN.cast_value(value) is False + + # number + assert ToolParameter.ToolParameterType.NUMBER.cast_value("1") == 1 + assert ToolParameter.ToolParameterType.NUMBER.cast_value("1.0") == 1.0 + assert ToolParameter.ToolParameterType.NUMBER.cast_value("-1.0") == -1.0 + assert ToolParameter.ToolParameterType.NUMBER.cast_value(1) == 1 + assert ToolParameter.ToolParameterType.NUMBER.cast_value(1.0) == 1.0 + assert ToolParameter.ToolParameterType.NUMBER.cast_value(-1.0) == -1.0 + assert ToolParameter.ToolParameterType.NUMBER.cast_value(None) is None diff --git a/api/tests/unit_tests/core/tools/test_mcp_tool_provider.py b/api/tests/unit_tests/core/tools/mcp_tool/test_provider.py similarity index 100% rename from api/tests/unit_tests/core/tools/test_mcp_tool_provider.py rename to api/tests/unit_tests/core/tools/mcp_tool/test_provider.py diff --git a/api/tests/unit_tests/core/tools/mcp_tool/test_tool.py b/api/tests/unit_tests/core/tools/mcp_tool/test_tool.py new file mode 100644 index 00000000000..39354887c9b --- /dev/null +++ b/api/tests/unit_tests/core/tools/mcp_tool/test_tool.py @@ -0,0 +1,699 @@ +from __future__ import annotations + +import base64 +import contextlib +from collections.abc import Iterator +from decimal import Decimal +from typing import Any +from unittest.mock import patch + +import pytest +from sqlalchemy.engine import Engine +from sqlalchemy.orm import Session + +from core.app.entities.app_invoke_entities import InvokeFrom +from core.mcp.types import ( + AudioContent, + BlobResourceContents, + CallToolResult, + EmbeddedResource, + ImageContent, + TextContent, + TextResourceContents, +) +from core.tools.__base.tool_runtime import ToolRuntime +from core.tools.entities.common_entities import I18nObject +from core.tools.entities.tool_entities import ToolEntity, ToolIdentity, ToolInvokeMessage, ToolProviderType +from core.tools.errors import ToolInvokeError +from core.tools.mcp_tool.tool import MCPTool +from enums import DeploymentEdition +from graphon.model_runtime.entities.llm_entities import LLMUsage +from services.tools.tools_transform_service import ToolTransformService + +"""MCP tool tests using real SQLite sessions for the ORM invocation contract.""" + + +def _build_mcp_tool(*, with_output_schema: bool = True) -> MCPTool: + entity = ToolEntity( + identity=ToolIdentity( + author="author", + name="remote-tool", + label=I18nObject(en_US="remote-tool"), + provider="provider-id", + ), + parameters=[], + output_schema={"type": "object"} if with_output_schema else {}, + ) + return MCPTool( + entity=entity, + runtime=ToolRuntime(tenant_id="tenant-1", invoke_from=InvokeFrom.DEBUGGER), + tenant_id="tenant-1", + icon="icon.svg", + server_url="https://mcp.example.com", + server_identifier="mcp_server", + headers={"x-auth": "token"}, + ) + + +# Home Assistant's GetLiveContext declares `domain` as `vol.Any(cv.string, [cv.string])`. +_STRING = {"type": "string"} +_STRING_LIST = {"type": "array", "items": {"type": "string"}} + + +# ----- M2/M3 user-identity forwarding --------------------------------------- + + +def _build_forwarding_tool(*, mode: str = "idp_token") -> MCPTool: + """Helper that builds an MCPTool with the identity_mode set.""" + entity = ToolEntity( + identity=ToolIdentity( + author="author", + name="remote-tool", + label=I18nObject(en_US="remote-tool"), + provider="provider-id", + ), + parameters=[], + output_schema={}, + ) + return MCPTool( + entity=entity, + runtime=ToolRuntime(tenant_id="tenant-1", invoke_from=InvokeFrom.DEBUGGER), + tenant_id="tenant-1", + icon="icon.svg", + server_url="https://mcp.example.com/mcp/", + server_identifier="mcp_server", + identity_mode=mode, + ) + + +class TestMcpToolCases: + @staticmethod + def test_mcp_tool_provider_type_and_fork_runtime(): + tool = _build_mcp_tool() + assert tool.tool_provider_type() == ToolProviderType.MCP + + forked = tool.fork_tool_runtime(ToolRuntime(tenant_id="tenant-2")) + assert isinstance(forked, MCPTool) + assert forked.runtime.tenant_id == "tenant-2" + assert forked.server_identifier == "mcp_server" + + @staticmethod + def test_mcp_tool_text_and_json_processing_helpers(): + tool = _build_mcp_tool() + + json_messages = list(tool._process_text_content(TextContent(type="text", text='{"a": 1}'))) + assert json_messages[0].type == ToolInvokeMessage.MessageType.JSON + + plain_messages = list(tool._process_text_content(TextContent(type="text", text="not-json"))) + assert plain_messages[0].type == ToolInvokeMessage.MessageType.TEXT + assert plain_messages[0].message.text == "not-json" + + list_messages = list(tool._process_json_content([{"k": 1}, {"k": 2}])) + assert [m.type for m in list_messages] == [ + ToolInvokeMessage.MessageType.JSON, + ToolInvokeMessage.MessageType.JSON, + ] + + mixed_list_messages = list(tool._process_json_list([{"k": 1}, 2])) + assert len(mixed_list_messages) == 1 + assert mixed_list_messages[0].type == ToolInvokeMessage.MessageType.TEXT + + primitive_messages = list(tool._process_json_content(123)) + assert primitive_messages[0].message.text == "123" + + @staticmethod + def test_mcp_tool_usage_extraction_helpers(): + usage = MCPTool._extract_usage_dict({"usage": {"total_tokens": 9}}) + assert usage == {"total_tokens": 9} + + usage = MCPTool._extract_usage_dict({"metadata": {"usage": {"prompt_tokens": 3, "completion_tokens": 2}}}) + assert usage == {"prompt_tokens": 3, "completion_tokens": 2} + + usage = MCPTool._extract_usage_dict({"prompt_tokens": 1, "completion_tokens": 2, "total_tokens": 3}) + assert usage == {"prompt_tokens": 1, "completion_tokens": 2, "total_tokens": 3} + + usage = MCPTool._extract_usage_dict({"nested": [{"deep": {"usage": {"total_tokens": 7}}}]}) + assert usage == {"total_tokens": 7} + + result_with_usage = CallToolResult(content=[], _meta={"usage": {"prompt_tokens": 1, "completion_tokens": 2}}) + derived = MCPTool._derive_usage_from_result(result_with_usage) + assert derived.prompt_tokens == 1 + assert derived.completion_tokens == 2 + + result_without_usage = CallToolResult(content=[], _meta=None) + derived = MCPTool._derive_usage_from_result(result_without_usage) + assert derived.total_tokens == 0 + + @staticmethod + @pytest.mark.parametrize("sqlite_session", [()], indirect=True) + def test_mcp_tool_invoke_handles_content_types_and_structured_output(sqlite_session: Session): + tool = _build_mcp_tool() + img_data = base64.b64encode(b"img").decode() + blob_data = base64.b64encode(b"blob").decode() + result = CallToolResult( + content=[ + TextContent(type="text", text='{"a": 1}'), + ImageContent(type="image", data=img_data, mimeType="image/png"), + EmbeddedResource( + type="resource", + resource=TextResourceContents(uri="file:///tmp/a.txt", text="embedded-text"), + ), + EmbeddedResource( + type="resource", + resource=BlobResourceContents( + uri="file:///tmp/b.bin", + blob=blob_data, + mimeType="application/octet-stream", + ), + ), + ], + structuredContent={"x": 1}, + _meta={"usage": {"prompt_tokens": 2, "completion_tokens": 3}}, + ) + + with patch.object(MCPTool, "invoke_remote_mcp_tool", return_value=result): + messages = list(tool.invoke(session=sqlite_session, user_id="user-1", tool_parameters={"a": 1})) + + types = [m.type for m in messages] + assert ToolInvokeMessage.MessageType.JSON in types + assert ToolInvokeMessage.MessageType.BLOB in types + assert ToolInvokeMessage.MessageType.TEXT in types + assert ToolInvokeMessage.MessageType.VARIABLE in types + assert tool.latest_usage.total_tokens == 5 + + @staticmethod + @pytest.mark.parametrize("sqlite_session", [()], indirect=True) + def test_mcp_tool_invoke_raises_for_unsupported_embedded_resource(sqlite_session: Session): + tool = _build_mcp_tool() + # Use model_construct to bypass pydantic validation and force unsupported resource path. + bad_resource = EmbeddedResource.model_construct(type="resource", resource=object()) + result = CallToolResult(content=[bad_resource], _meta=None) + + with patch.object(MCPTool, "invoke_remote_mcp_tool", return_value=result): + with pytest.raises(ToolInvokeError, match="Unsupported embedded resource type"): + list(tool.invoke(session=sqlite_session, user_id="user-1", tool_parameters={})) + + @staticmethod + def test_mcp_tool_handle_none_parameter_filters_empty_values(): + tool = _build_mcp_tool() + cleaned = tool._handle_none_parameter({"a": 1, "b": None, "c": "", "d": " ", "e": "ok"}) + assert cleaned == {"a": 1, "e": "ok"} + + @staticmethod + @pytest.mark.parametrize("sqlite_session", [()], indirect=True) + @pytest.mark.parametrize( + "domain_schema", + [ + {"anyOf": [_STRING, _STRING_LIST]}, + {"anyOf": [_STRING_LIST, _STRING]}, + {"oneOf": [_STRING, _STRING_LIST]}, + {"oneOf": [_STRING_LIST, _STRING]}, + {"type": ["string", "array"], "items": _STRING}, + {"type": ["array", "string"], "items": _STRING}, + ], + ) + @pytest.mark.parametrize("domain", [["light", "cover"], "light"]) + def test_mcp_tool_invoke_preserves_values_of_mixed_type_union_parameters( + sqlite_session: Session, domain_schema, domain + ): + tool = _build_mcp_tool() + schema = {"type": "object", "properties": {"domain": domain_schema, "name": _STRING}, "required": []} + tool.entity.parameters = ToolTransformService.convert_mcp_schema_to_parameter(schema) + result = CallToolResult(content=[TextContent(type="text", text="ok")]) + + with patch.object(MCPTool, "invoke_remote_mcp_tool", return_value=result) as invoke_remote: + list(tool.invoke(session=sqlite_session, user_id="user-1", tool_parameters={"domain": domain, "name": 1})) + + sent = invoke_remote.call_args.args[0] + assert sent["domain"] == domain + # Parameters with a single declared type keep the existing cast. + assert sent["name"] == "1" + + @staticmethod + def test_inject_forwarded_identity_stamps_custom_header(): + """The minted SSO token must be placed in X-Dify-SSO-Token; the + workspace-scoped Authorization header and any other custom headers must + pass through untouched so provider credentials keep working.""" + from core.tools.mcp_tool.tool import FORWARDED_IDENTITY_HEADER + + tool = _build_forwarding_tool() + headers: dict[str, str] = {"Authorization": "Bearer static-client-token", "X-Other": "keep"} + + with patch( + "services.enterprise.enterprise_service.EnterpriseService.issue_mcp_token", + return_value=("forwarded.jwt.payload", 1900000000), + ): + tool._inject_forwarded_identity( + headers, user_id="alice", app_id=None, audience="https://mcp.example.com/mcp/" + ) + + assert headers[FORWARDED_IDENTITY_HEADER] == "forwarded.jwt.payload" + assert headers["Authorization"] == "Bearer static-client-token" + assert headers["X-Other"] == "keep" + + @staticmethod + def test_inject_forwarded_identity_translates_token_error_to_invoke_error(): + """EnterpriseService failures must surface as ToolInvokeError so the + workflow halts loudly instead of proceeding without identity.""" + from core.tools.mcp_tool.tool import FORWARDED_IDENTITY_HEADER + from services.enterprise.base import MCPNoRefreshTokenError + + tool = _build_forwarding_tool() + headers: dict[str, str] = {} + + with patch( + "services.enterprise.enterprise_service.EnterpriseService.issue_mcp_token", + side_effect=MCPNoRefreshTokenError("please re-sso"), + ): + with pytest.raises(ToolInvokeError, match="forwarded identity token"): + tool._inject_forwarded_identity( + headers, user_id="alice", app_id=None, audience="https://mcp.example.com/mcp/" + ) + + # Headers must NOT have been mutated when token-issuance failed. + assert FORWARDED_IDENTITY_HEADER not in headers + assert "Authorization" not in headers + + @staticmethod + def test_inject_forwarded_identity_sends_end_user_type_for_webapp(): + """A WEB_APP run forwards user_type=end_user so enterprise routes to the + published-webapp token store.""" + tool = _build_forwarding_tool() + tool.runtime = ToolRuntime(tenant_id="tenant-1", invoke_from=InvokeFrom.WEB_APP) + headers: dict[str, str] = {} + + with patch( + "services.enterprise.enterprise_service.EnterpriseService.issue_mcp_token", + return_value=("forwarded.jwt", 1900000000), + ) as issue: + tool._inject_forwarded_identity( + headers, user_id="eu-1", app_id="app-1", audience="https://mcp.example.com/mcp/" + ) + + assert issue.call_args.kwargs["user_type"] == "end_user" + + @staticmethod + def test_inject_forwarded_identity_sends_account_type_for_debugger(): + """A DEBUGGER/console run forwards user_type=account (the existing behaviour).""" + tool = _build_forwarding_tool() # built with InvokeFrom.DEBUGGER + headers: dict[str, str] = {} + + with patch( + "services.enterprise.enterprise_service.EnterpriseService.issue_mcp_token", + return_value=("forwarded.jwt", 1900000000), + ) as issue: + tool._inject_forwarded_identity( + headers, user_id="acc-1", app_id=None, audience="https://mcp.example.com/mcp/" + ) + + assert issue.call_args.kwargs["user_type"] == "account" + + @staticmethod + def test_invoke_remote_mcp_tool_fails_closed_when_user_id_missing(config_overrides): + """When forwarding is enabled AND the deployment is enterprise, missing + user_id must raise — never silently invoke as the static identity.""" + tool = _build_forwarding_tool() + + config_overrides(DEPLOYMENT_EDITION=DeploymentEdition.ENTERPRISE) + with pytest.raises(ToolInvokeError, match="no end-user context"): + tool.invoke_remote_mcp_tool({}, user_id=None, app_id=None) + + @staticmethod + def test_invoke_skips_forwarding_outside_enterprise_edition(config_overrides): + """Non-enterprise deployments treat the DB selector as a no-op: a stale + `identity_mode="idp_token"` row must NOT raise (fail-closed) AND must + NOT call the enterprise inner API. The runtime falls through to the + legacy provider-identity path.""" + tool = _build_forwarding_tool() + + config_overrides(DEPLOYMENT_EDITION=DeploymentEdition.COMMUNITY) + # The fail-closed branch must NOT fire (no enterprise → no forwarding). + # The function will still try the legacy DB-load path; we patch that + # to keep the test unit-scoped. + with patch("core.tools.mcp_tool.tool.MCPClientWithAuthRetry") as client_cls: + client_cls.return_value.__enter__.return_value.invoke_tool.return_value = CallToolResult( + content=[], + _meta=None, + ) + with patch.object(tool, "_inject_forwarded_identity") as inject: + with patch("services.tools.mcp_tools_manage_service.MCPToolManageService"): + with patch("core.entities.mcp_provider.MCPProviderEntity.decrypt_server_url", return_value="u"): + with patch("core.entities.mcp_provider.MCPProviderEntity.decrypt_headers", return_value={}): + # Should not raise; should not call enterprise. + with contextlib.suppress(Exception): + tool.invoke_remote_mcp_tool({}, user_id=None, app_id=None) + inject.assert_not_called() + + +def _make_mcp_tool(output_schema: dict[str, Any] | None = None) -> MCPTool: + identity = ToolIdentity( + author="test", + name="test_mcp_tool", + label=I18nObject(en_US="Test MCP Tool", zh_Hans="测试MCP工具"), + provider="test_provider", + ) + entity = ToolEntity(identity=identity, output_schema=output_schema or {}) + runtime = ToolRuntime(tenant_id="test_tenant") + return MCPTool( + entity=entity, + runtime=runtime, + tenant_id="test_tenant", + icon="", + server_url="https://server.invalid", + server_identifier="provider_1", + headers={}, + ) + + +class TestMcpToolCasesTools: + @staticmethod + @pytest.fixture + def orm_session(sqlite_engine: Engine) -> Iterator[Session]: + """Use a real ORM session while MCP transport remains mocked.""" + with Session(sqlite_engine) as session: + yield session + + class TestMCPToolInvoke: + @pytest.mark.parametrize( + ("content_factory", "mime_type"), + [ + ( + lambda b64, mt: ImageContent(type="image", data=b64, mimeType=mt), + "image/png", + ), + ( + lambda b64, mt: AudioContent(type="audio", data=b64, mimeType=mt), + "audio/mpeg", + ), + ], + ) + def test_invoke_image_or_audio_yields_blob(self, content_factory, mime_type, orm_session: Session) -> None: + tool = _make_mcp_tool() + raw = b"\x00\x01test-bytes\x02" + b64 = base64.b64encode(raw).decode() + content = content_factory(b64, mime_type) + result = CallToolResult(content=[content]) + + with patch.object(tool, "invoke_remote_mcp_tool", return_value=result): + messages = list(tool._invoke(session=orm_session, user_id="test_user", tool_parameters={})) + + assert len(messages) == 1 + msg = messages[0] + assert msg.type == ToolInvokeMessage.MessageType.BLOB + assert isinstance(msg.message, ToolInvokeMessage.BlobMessage) + assert msg.message.blob == raw + assert msg.meta == {"mime_type": mime_type} + + def test_invoke_embedded_text_resource_yields_text(self, orm_session: Session) -> None: + tool = _make_mcp_tool() + text_resource = TextResourceContents(uri="file://test.txt", mimeType="text/plain", text="hello world") + content = EmbeddedResource(type="resource", resource=text_resource) + result = CallToolResult(content=[content]) + + with patch.object(tool, "invoke_remote_mcp_tool", return_value=result): + messages = list(tool._invoke(session=orm_session, user_id="test_user", tool_parameters={})) + + assert len(messages) == 1 + msg = messages[0] + assert msg.type == ToolInvokeMessage.MessageType.TEXT + assert isinstance(msg.message, ToolInvokeMessage.TextMessage) + assert msg.message.text == "hello world" + + @pytest.mark.parametrize( + ("mime_type", "expected_mime"), + [("application/pdf", "application/pdf"), (None, "application/octet-stream")], + ) + def test_invoke_embedded_blob_resource_yields_blob( + self, mime_type, expected_mime, orm_session: Session + ) -> None: + tool = _make_mcp_tool() + raw = b"binary-data" + b64 = base64.b64encode(raw).decode() + blob_resource = BlobResourceContents(uri="file://doc.bin", mimeType=mime_type, blob=b64) + content = EmbeddedResource(type="resource", resource=blob_resource) + result = CallToolResult(content=[content]) + + with patch.object(tool, "invoke_remote_mcp_tool", return_value=result): + messages = list(tool._invoke(session=orm_session, user_id="test_user", tool_parameters={})) + + assert len(messages) == 1 + msg = messages[0] + assert msg.type == ToolInvokeMessage.MessageType.BLOB + assert isinstance(msg.message, ToolInvokeMessage.BlobMessage) + assert msg.message.blob == raw + assert msg.meta == {"mime_type": expected_mime} + + def test_invoke_yields_variables_when_structured_content_and_schema(self, orm_session: Session) -> None: + tool = _make_mcp_tool(output_schema={"type": "object"}) + result = CallToolResult(content=[], structuredContent={"a": 1, "b": "x"}) + + with patch.object(tool, "invoke_remote_mcp_tool", return_value=result): + messages = list(tool._invoke(session=orm_session, user_id="test_user", tool_parameters={})) + + # Expect two variable messages corresponding to keys a and b + assert len(messages) == 2 + var_msgs = [m for m in messages if isinstance(m.message, ToolInvokeMessage.VariableMessage)] + assert {m.message.variable_name for m in var_msgs} == {"a", "b"} + # Validate values + values = {m.message.variable_name: m.message.variable_value for m in var_msgs} + assert values == {"a": 1, "b": "x"} + + def test_invoke_yields_json_when_structured_content_has_no_output_schema(self, orm_session: Session) -> None: + tool = _make_mcp_tool() + result = CallToolResult(content=[], structuredContent={"a": 1, "b": "x"}) + + with patch.object(tool, "invoke_remote_mcp_tool", return_value=result): + messages = list(tool._invoke(session=orm_session, user_id="test_user", tool_parameters={})) + + assert len(messages) == 1 + msg = messages[0] + assert msg.type == ToolInvokeMessage.MessageType.JSON + assert isinstance(msg.message, ToolInvokeMessage.JsonMessage) + assert msg.message.json_object == {"a": 1, "b": "x"} + + class TestMCPToolUsageExtraction: + """Test usage metadata extraction from MCP tool results.""" + + def test_extract_usage_dict_from_direct_usage_field(self) -> None: + """Test extraction when usage is directly in meta.usage field.""" + meta = { + "usage": { + "prompt_tokens": 100, + "completion_tokens": 50, + "total_tokens": 150, + "total_price": "0.001", + "currency": "USD", + } + } + usage_dict = MCPTool._extract_usage_dict(meta) + assert usage_dict is not None + assert usage_dict["prompt_tokens"] == 100 + assert usage_dict["completion_tokens"] == 50 + assert usage_dict["total_tokens"] == 150 + assert usage_dict["total_price"] == "0.001" + assert usage_dict["currency"] == "USD" + + def test_extract_usage_dict_from_nested_metadata(self) -> None: + """Test extraction when usage is nested in meta.metadata.usage.""" + meta = { + "metadata": { + "usage": { + "prompt_tokens": 200, + "completion_tokens": 100, + "total_tokens": 300, + } + } + } + usage_dict = MCPTool._extract_usage_dict(meta) + assert usage_dict is not None + assert usage_dict["prompt_tokens"] == 200 + assert usage_dict["total_tokens"] == 300 + + def test_extract_usage_dict_from_flat_token_fields(self) -> None: + """Test extraction when token counts are directly in meta.""" + meta = { + "prompt_tokens": 150, + "completion_tokens": 75, + "total_tokens": 225, + "currency": "EUR", + } + usage_dict = MCPTool._extract_usage_dict(meta) + assert usage_dict is not None + assert usage_dict["prompt_tokens"] == 150 + assert usage_dict["completion_tokens"] == 75 + assert usage_dict["total_tokens"] == 225 + assert usage_dict["currency"] == "EUR" + + def test_extract_usage_dict_recursive(self) -> None: + """Test recursive search through nested structures.""" + meta = { + "custom": { + "nested": { + "usage": { + "total_tokens": 500, + "prompt_tokens": 300, + "completion_tokens": 200, + } + } + } + } + usage_dict = MCPTool._extract_usage_dict(meta) + assert usage_dict is not None + assert usage_dict["total_tokens"] == 500 + + def test_extract_usage_dict_from_list(self) -> None: + """Test extraction from nested list structures.""" + meta = { + "items": [ + {"usage": {"total_tokens": 100}}, + {"other": "data"}, + ] + } + usage_dict = MCPTool._extract_usage_dict(meta) + assert usage_dict is not None + assert usage_dict["total_tokens"] == 100 + + def test_extract_usage_dict_returns_none_when_missing(self) -> None: + """Test that None is returned when no usage data is present.""" + meta = {"other": "data", "custom": {"nested": {"value": 123}}} + usage_dict = MCPTool._extract_usage_dict(meta) + assert usage_dict is None + + def test_extract_usage_dict_empty_meta(self) -> None: + """Test with empty meta dict.""" + usage_dict = MCPTool._extract_usage_dict({}) + assert usage_dict is None + + def test_derive_usage_from_result_with_meta(self) -> None: + """Test _derive_usage_from_result with populated meta.""" + meta = { + "usage": { + "prompt_tokens": 100, + "completion_tokens": 50, + "total_tokens": 150, + "total_price": "0.0015", + "currency": "USD", + } + } + result = CallToolResult(content=[], _meta=meta) + usage = MCPTool._derive_usage_from_result(result) + + assert isinstance(usage, LLMUsage) + assert usage.prompt_tokens == 100 + assert usage.completion_tokens == 50 + assert usage.total_tokens == 150 + assert usage.total_price == Decimal("0.0015") + assert usage.currency == "USD" + + def test_derive_usage_from_result_without_meta(self) -> None: + """Test _derive_usage_from_result with no meta returns empty usage.""" + result = CallToolResult(content=[], meta=None) + usage = MCPTool._derive_usage_from_result(result) + + assert isinstance(usage, LLMUsage) + assert usage.total_tokens == 0 + assert usage.prompt_tokens == 0 + assert usage.completion_tokens == 0 + + def test_derive_usage_from_result_calculates_total_tokens(self) -> None: + """Test that total_tokens is calculated when missing.""" + meta = { + "usage": { + "prompt_tokens": 100, + "completion_tokens": 50, + # total_tokens is missing + } + } + result = CallToolResult(content=[], _meta=meta) + usage = MCPTool._derive_usage_from_result(result) + + assert usage.total_tokens == 150 # 100 + 50 + assert usage.prompt_tokens == 100 + assert usage.completion_tokens == 50 + + def test_invoke_sets_latest_usage_from_meta(self, orm_session: Session) -> None: + """Test that _invoke sets _latest_usage from result meta.""" + tool = _make_mcp_tool() + meta = { + "usage": { + "prompt_tokens": 200, + "completion_tokens": 100, + "total_tokens": 300, + "total_price": "0.003", + "currency": "USD", + } + } + result = CallToolResult(content=[TextContent(type="text", text="test")], _meta=meta) + + with patch.object(tool, "invoke_remote_mcp_tool", return_value=result): + list(tool._invoke(session=orm_session, user_id="test_user", tool_parameters={})) + + # Verify latest_usage was set correctly + assert tool.latest_usage.prompt_tokens == 200 + assert tool.latest_usage.completion_tokens == 100 + assert tool.latest_usage.total_tokens == 300 + assert tool.latest_usage.total_price == Decimal("0.003") + + def test_invoke_with_no_meta_returns_empty_usage(self, orm_session: Session) -> None: + """Test that _invoke returns empty usage when no meta is present.""" + tool = _make_mcp_tool() + result = CallToolResult(content=[TextContent(type="text", text="test")], _meta=None) + + with patch.object(tool, "invoke_remote_mcp_tool", return_value=result): + list(tool._invoke(session=orm_session, user_id="test_user", tool_parameters={})) + + # Verify latest_usage is empty + assert tool.latest_usage.total_tokens == 0 + assert tool.latest_usage.prompt_tokens == 0 + assert tool.latest_usage.completion_tokens == 0 + + def test_latest_usage_property_returns_llm_usage(self) -> None: + """Test that latest_usage property returns LLMUsage instance.""" + tool = _make_mcp_tool() + assert isinstance(tool.latest_usage, LLMUsage) + + def test_initial_usage_is_empty(self) -> None: + """Test that MCPTool is initialized with empty usage.""" + tool = _make_mcp_tool() + assert tool.latest_usage.total_tokens == 0 + assert tool.latest_usage.prompt_tokens == 0 + assert tool.latest_usage.completion_tokens == 0 + assert tool.latest_usage.total_price == Decimal(0) + + @pytest.mark.parametrize( + "meta_data", + [ + # Direct usage field + {"usage": {"prompt_tokens": 10, "completion_tokens": 5, "total_tokens": 15}}, + # Nested metadata + {"metadata": {"usage": {"total_tokens": 100}}}, + # Flat token fields + {"total_tokens": 50, "prompt_tokens": 30, "completion_tokens": 20}, + # With price info + { + "usage": { + "total_tokens": 150, + "total_price": "0.002", + "currency": "EUR", + } + }, + # Deep nested + {"level1": {"level2": {"usage": {"total_tokens": 200}}}}, + ], + ) + def test_various_meta_formats(self, meta_data) -> None: + """Test that various meta formats are correctly parsed.""" + result = CallToolResult(content=[], _meta=meta_data) + usage = MCPTool._derive_usage_from_result(result) + + assert isinstance(usage, LLMUsage) + # Should have at least some usage data + if meta_data.get("usage", {}).get("total_tokens") or meta_data.get("total_tokens"): + expected_total = ( + meta_data.get("usage", {}).get("total_tokens") + or meta_data.get("total_tokens") + or meta_data.get("metadata", {}).get("usage", {}).get("total_tokens") + or meta_data.get("level1", {}).get("level2", {}).get("usage", {}).get("total_tokens") + ) + if expected_total: + assert usage.total_tokens == expected_total diff --git a/api/tests/unit_tests/core/tools/test_plugin_tool_provider.py b/api/tests/unit_tests/core/tools/plugin_tool/test_provider.py similarity index 100% rename from api/tests/unit_tests/core/tools/test_plugin_tool_provider.py rename to api/tests/unit_tests/core/tools/plugin_tool/test_provider.py diff --git a/api/tests/unit_tests/core/tools/test_plugin_tool.py b/api/tests/unit_tests/core/tools/plugin_tool/test_tool.py similarity index 100% rename from api/tests/unit_tests/core/tools/test_plugin_tool.py rename to api/tests/unit_tests/core/tools/plugin_tool/test_tool.py diff --git a/api/tests/unit_tests/core/tools/test_custom_tool.py b/api/tests/unit_tests/core/tools/test_custom_tool.py deleted file mode 100644 index 9dd68334b6c..00000000000 --- a/api/tests/unit_tests/core/tools/test_custom_tool.py +++ /dev/null @@ -1,291 +0,0 @@ -from __future__ import annotations - -from types import SimpleNamespace -from typing import Any - -import httpx -import pytest -from sqlalchemy.orm import Session - -from core.app.entities.app_invoke_entities import InvokeFrom -from core.tools.__base.tool_runtime import ToolRuntime -from core.tools.custom_tool.tool import ApiTool, ParsedResponse -from core.tools.entities.common_entities import I18nObject -from core.tools.entities.tool_bundle import ApiToolBundle -from core.tools.entities.tool_entities import ToolEntity, ToolIdentity, ToolInvokeMessage -from core.tools.errors import ToolInvokeError, ToolParameterValidationError, ToolProviderCredentialValidationError - - -def _build_tool(*, openapi: dict[str, Any] | None = None) -> ApiTool: - entity = ToolEntity( - identity=ToolIdentity( - author="author", - name="tool-a", - label=I18nObject(en_US="tool-a"), - provider="provider-a", - ), - parameters=[], - ) - bundle = ApiToolBundle( - server_url="https://api.example.com/items/{id}", - method="GET", - summary="summary", - operation_id="op-id", - parameters=[], - author="author", - openapi=openapi or {"parameters": []}, - ) - runtime = ToolRuntime( - tenant_id="tenant-1", - invoke_from=InvokeFrom.DEBUGGER, - credentials={"auth_type": "api_key_header", "api_key_value": "k"}, - ) - return ApiTool(entity=entity, api_bundle=bundle, runtime=runtime, provider_id="provider-id") - - -def test_parsed_response_to_string(): - assert ParsedResponse({"a": 1}, True).to_string() == '{"a": 1}' - assert ParsedResponse("ok", False).to_string() == "ok" - - -def test_api_tool_fork_runtime_and_validate_credentials(monkeypatch: pytest.MonkeyPatch): - tool = _build_tool() - forked = tool.fork_tool_runtime(ToolRuntime(tenant_id="tenant-2")) - assert isinstance(forked, ApiTool) - assert forked.runtime.tenant_id == "tenant-2" - - tool.api_bundle = None # type: ignore[assignment] - with pytest.raises(ValueError, match="api_bundle is required"): - tool.fork_tool_runtime(ToolRuntime(tenant_id="tenant-2")) - - tool = _build_tool() - assert tool.validate_credentials(credentials={}, parameters={}, format_only=True) == "" - monkeypatch.setattr(tool, "assembling_request", lambda parameters: {"Authorization": "Bearer x"}) - monkeypatch.setattr( - tool, - "do_http_request", - lambda url, method, headers, parameters: httpx.Response(200, json={"ok": True}), - ) - result = tool.validate_credentials(credentials={}, parameters={"a": 1}, format_only=False) - assert result == '{"ok": true}' - - -def test_assembling_request_auth_header_assembly(): - tool = _build_tool() - - headers = tool.assembling_request(parameters={}) - assert headers["Authorization"] == "k" - - tool.runtime.credentials = { - "auth_type": "api_key_header", - "api_key_header_prefix": "bearer", - "api_key_value": "abc", - } - headers = tool.assembling_request(parameters={}) - assert headers["Authorization"] == "Bearer abc" - - tool.runtime.credentials = {"auth_type": "api_key_header", "api_key_header_prefix": "basic", "api_key_value": "abc"} - headers = tool.assembling_request(parameters={}) - assert headers["Authorization"] == "Basic abc" - assert tool.runtime.credentials["api_key_value"] == "abc" - - headers = tool.assembling_request(parameters={}) - assert headers["Authorization"] == "Basic abc" - - tool.runtime.credentials = {"auth_type": "api_key_query", "api_key_value": "abc"} - assert tool.assembling_request(parameters={}) == {} - - -def test_assembling_request_runtime_auth_errors(): - tool = _build_tool() - - tool.runtime = None - with pytest.raises(ToolProviderCredentialValidationError, match="runtime not initialized"): - tool.assembling_request(parameters={}) - - tool.runtime = ToolRuntime(tenant_id="tenant", credentials={}) - with pytest.raises(ToolProviderCredentialValidationError, match="Missing auth_type"): - tool.assembling_request(parameters={}) - - tool.runtime.credentials = {"auth_type": "api_key_header"} - with pytest.raises(ToolProviderCredentialValidationError, match="Missing api_key_value"): - tool.assembling_request(parameters={}) - - tool.runtime.credentials = {"auth_type": "api_key_header", "api_key_value": 123} - with pytest.raises(ToolProviderCredentialValidationError, match="must be a string"): - tool.assembling_request(parameters={}) - - -def test_assembling_request_parameter_validation_and_defaults(): - tool = _build_tool() - - tool.runtime.credentials = {"auth_type": "api_key_header", "api_key_value": "x"} - tool.api_bundle.parameters = [ - SimpleNamespace(required=True, name="required_param", default=None), - ] - with pytest.raises(ToolParameterValidationError, match="Missing required parameter required_param"): - tool.assembling_request(parameters={}) - - tool.api_bundle.parameters = [ - SimpleNamespace(required=True, name="required_param", default="d"), - ] - params = {} - tool.assembling_request(parameters=params) - assert params["required_param"] == "d" - - -def test_validate_and_parse_response_branches(): - tool = _build_tool() - - with pytest.raises(ToolInvokeError, match="status code 500"): - tool.validate_and_parse_response(httpx.Response(500, text="boom")) - - empty = tool.validate_and_parse_response(httpx.Response(200, content=b"")) - assert empty.is_json is False - assert "Empty response from the tool" in str(empty.content) - - json_resp = tool.validate_and_parse_response( - httpx.Response(200, json={"a": 1}, headers={"content-type": "application/json"}) - ) - assert json_resp.is_json is True - assert json_resp.content == {"a": 1} - - non_json_type = tool.validate_and_parse_response( - httpx.Response(200, text='{"a": 1}', headers={"content-type": "text/plain"}) - ) - assert non_json_type.is_json is False - assert non_json_type.content == '{"a": 1}' - - plain_resp = tool.validate_and_parse_response(httpx.Response(200, text="plain")) - assert plain_resp.is_json is False - assert plain_resp.content == "plain" - - with pytest.raises(ValueError, match="Invalid response type"): - tool.validate_and_parse_response("invalid") # type: ignore[arg-type] - - -def test_get_parameter_value_and_type_conversion_helpers(): - tool = _build_tool() - - assert tool.get_parameter_value({"name": "x"}, {"x": 1}) == 1 - assert tool.get_parameter_value({"name": "x", "required": False, "schema": {"default": "d"}}, {}) == "d" - with pytest.raises(ToolParameterValidationError, match="Missing required parameter x"): - tool.get_parameter_value({"name": "x", "required": True}, {}) - - assert tool._convert_body_property_any_of({}, "12", [{"type": "integer"}]) == 12 - assert tool._convert_body_property_any_of({}, "1.5", [{"type": "number"}]) == 1.5 - assert tool._convert_body_property_any_of({}, "true", [{"type": "boolean"}]) is True - assert tool._convert_body_property_any_of({}, "", [{"type": "null"}]) is None - assert tool._convert_body_property_any_of({}, "x", [{"anyOf": [{"type": "string"}]}]) == "x" - - assert tool._convert_body_property_type({"type": "integer"}, "1") == 1 - assert tool._convert_body_property_type({"type": "number"}, "1.2") == 1.2 - assert tool._convert_body_property_type({"type": "string"}, 1) == "1" - assert tool._convert_body_property_type({"type": "boolean"}, 1) is True - assert tool._convert_body_property_type({"type": "null"}, None) is None - assert tool._convert_body_property_type({"type": "object"}, '{"a":1}') == {"a": 1} - assert tool._convert_body_property_type({"type": "array"}, "[1,2]") == [1, 2] - assert tool._convert_body_property_type({"type": "invalid"}, "v") == "v" - assert tool._convert_body_property_type({"anyOf": [{"type": "integer"}]}, "2") == 2 - - -def test_do_http_request_builds_arguments_and_handles_invalid_method(monkeypatch: pytest.MonkeyPatch): - openapi = { - "parameters": [ - {"name": "id", "in": "path", "required": True, "schema": {"type": "string"}}, - {"name": "q", "in": "query", "required": False, "schema": {"default": ""}}, - {"name": "X-Extra", "in": "header", "required": False, "schema": {"default": "x"}}, - {"name": "sid", "in": "cookie", "required": False, "schema": {"default": "cookie1"}}, - ], - "requestBody": { - "content": { - "application/json": { - "schema": { - "type": "object", - "required": ["count"], - "properties": { - "count": {"type": "integer"}, - "name": {"type": "string", "default": "n"}, - }, - } - } - } - }, - } - tool = _build_tool(openapi=openapi) - tool.runtime.credentials = {"auth_type": "api_key_query", "api_key_query_param": "key", "api_key_value": "v"} - headers = {} - captured = {} - - def _fake_get(url, **kwargs): - captured["url"] = url - captured["kwargs"] = kwargs - return httpx.Response(200, text="ok") - - monkeypatch.setattr("core.tools.custom_tool.tool.ssrf_proxy.get", _fake_get) - response = tool.do_http_request( - "https://api.example.com/items/{id}", - "GET", - headers=headers, - parameters={"id": "123", "count": "2", "q": "search"}, - ) - - assert isinstance(response, httpx.Response) - assert captured["url"].endswith("/items/123") - assert captured["kwargs"]["params"]["q"] == "search" - assert captured["kwargs"]["params"]["key"] == "v" - assert captured["kwargs"]["headers"]["Content-Type"] == "application/json" - - invalid_method_tool = _build_tool(openapi={"parameters": []}) - with pytest.raises(ValueError, match="Invalid http method"): - invalid_method_tool.do_http_request("https://api.example.com", "TRACE", headers={}, parameters={}) - - -def test_do_http_request_handles_file_upload_and_invoke_paths(monkeypatch: pytest.MonkeyPatch, sqlite_session: Session): - openapi = { - "parameters": [], - "requestBody": { - "content": { - "multipart/form-data": { - "schema": { - "type": "object", - "properties": {"file": {"format": "binary"}}, - } - } - } - }, - } - tool = _build_tool(openapi=openapi) - tool.runtime.credentials = {"auth_type": "api_key_header", "api_key_value": "k"} - fake_file = SimpleNamespace(filename="a.txt", mime_type="text/plain") - captured = {} - - def _fake_post(url, **kwargs): - captured["headers"] = kwargs["headers"] - captured["files"] = kwargs["files"] - return httpx.Response(200, text="ok") - - monkeypatch.setattr("core.tools.custom_tool.tool.download", lambda _: b"file-bytes") - monkeypatch.setattr("core.tools.custom_tool.tool.ssrf_proxy.post", _fake_post) - response = tool.do_http_request( - "https://api.example.com/upload", - "POST", - headers={}, - parameters={"file": fake_file}, - ) - assert isinstance(response, httpx.Response) - assert "Content-Type" not in captured["headers"] - assert captured["files"][0][0] == "file" - - # _invoke JSON path - monkeypatch.setattr(tool, "assembling_request", lambda parameters: {}) - monkeypatch.setattr(tool, "do_http_request", lambda *args, **kwargs: httpx.Response(200, text='{"a":1}')) - monkeypatch.setattr(tool, "validate_and_parse_response", lambda _: ParsedResponse({"a": 1}, True)) - messages = list(tool.invoke(session=sqlite_session, user_id="u1", tool_parameters={})) - assert [m.type for m in messages] == [ToolInvokeMessage.MessageType.JSON, ToolInvokeMessage.MessageType.TEXT] - - # _invoke text path - monkeypatch.setattr(tool, "validate_and_parse_response", lambda _: ParsedResponse("plain", False)) - messages = list(tool.invoke(session=sqlite_session, user_id="u1", tool_parameters={})) - assert len(messages) == 1 - assert messages[0].message.text == "plain" diff --git a/api/tests/unit_tests/core/tools/test_mcp_tool.py b/api/tests/unit_tests/core/tools/test_mcp_tool.py deleted file mode 100644 index 66628d93f9a..00000000000 --- a/api/tests/unit_tests/core/tools/test_mcp_tool.py +++ /dev/null @@ -1,330 +0,0 @@ -"""MCP tool tests using real SQLite sessions for the ORM invocation contract.""" - -from __future__ import annotations - -import base64 -import contextlib -from unittest.mock import patch - -import pytest -from sqlalchemy.orm import Session - -from core.app.entities.app_invoke_entities import InvokeFrom -from core.mcp.types import ( - BlobResourceContents, - CallToolResult, - EmbeddedResource, - ImageContent, - TextContent, - TextResourceContents, -) -from core.tools.__base.tool_runtime import ToolRuntime -from core.tools.entities.common_entities import I18nObject -from core.tools.entities.tool_entities import ToolEntity, ToolIdentity, ToolInvokeMessage, ToolProviderType -from core.tools.errors import ToolInvokeError -from core.tools.mcp_tool.tool import MCPTool -from enums import DeploymentEdition -from services.tools.tools_transform_service import ToolTransformService - - -def _build_mcp_tool(*, with_output_schema: bool = True) -> MCPTool: - entity = ToolEntity( - identity=ToolIdentity( - author="author", - name="remote-tool", - label=I18nObject(en_US="remote-tool"), - provider="provider-id", - ), - parameters=[], - output_schema={"type": "object"} if with_output_schema else {}, - ) - return MCPTool( - entity=entity, - runtime=ToolRuntime(tenant_id="tenant-1", invoke_from=InvokeFrom.DEBUGGER), - tenant_id="tenant-1", - icon="icon.svg", - server_url="https://mcp.example.com", - server_identifier="mcp_server", - headers={"x-auth": "token"}, - ) - - -def test_mcp_tool_provider_type_and_fork_runtime(): - tool = _build_mcp_tool() - assert tool.tool_provider_type() == ToolProviderType.MCP - - forked = tool.fork_tool_runtime(ToolRuntime(tenant_id="tenant-2")) - assert isinstance(forked, MCPTool) - assert forked.runtime.tenant_id == "tenant-2" - assert forked.server_identifier == "mcp_server" - - -def test_mcp_tool_text_and_json_processing_helpers(): - tool = _build_mcp_tool() - - json_messages = list(tool._process_text_content(TextContent(type="text", text='{"a": 1}'))) - assert json_messages[0].type == ToolInvokeMessage.MessageType.JSON - - plain_messages = list(tool._process_text_content(TextContent(type="text", text="not-json"))) - assert plain_messages[0].type == ToolInvokeMessage.MessageType.TEXT - assert plain_messages[0].message.text == "not-json" - - list_messages = list(tool._process_json_content([{"k": 1}, {"k": 2}])) - assert [m.type for m in list_messages] == [ToolInvokeMessage.MessageType.JSON, ToolInvokeMessage.MessageType.JSON] - - mixed_list_messages = list(tool._process_json_list([{"k": 1}, 2])) - assert len(mixed_list_messages) == 1 - assert mixed_list_messages[0].type == ToolInvokeMessage.MessageType.TEXT - - primitive_messages = list(tool._process_json_content(123)) - assert primitive_messages[0].message.text == "123" - - -def test_mcp_tool_usage_extraction_helpers(): - usage = MCPTool._extract_usage_dict({"usage": {"total_tokens": 9}}) - assert usage == {"total_tokens": 9} - - usage = MCPTool._extract_usage_dict({"metadata": {"usage": {"prompt_tokens": 3, "completion_tokens": 2}}}) - assert usage == {"prompt_tokens": 3, "completion_tokens": 2} - - usage = MCPTool._extract_usage_dict({"prompt_tokens": 1, "completion_tokens": 2, "total_tokens": 3}) - assert usage == {"prompt_tokens": 1, "completion_tokens": 2, "total_tokens": 3} - - usage = MCPTool._extract_usage_dict({"nested": [{"deep": {"usage": {"total_tokens": 7}}}]}) - assert usage == {"total_tokens": 7} - - result_with_usage = CallToolResult(content=[], _meta={"usage": {"prompt_tokens": 1, "completion_tokens": 2}}) - derived = MCPTool._derive_usage_from_result(result_with_usage) - assert derived.prompt_tokens == 1 - assert derived.completion_tokens == 2 - - result_without_usage = CallToolResult(content=[], _meta=None) - derived = MCPTool._derive_usage_from_result(result_without_usage) - assert derived.total_tokens == 0 - - -@pytest.mark.parametrize("sqlite_session", [()], indirect=True) -def test_mcp_tool_invoke_handles_content_types_and_structured_output(sqlite_session: Session): - tool = _build_mcp_tool() - img_data = base64.b64encode(b"img").decode() - blob_data = base64.b64encode(b"blob").decode() - result = CallToolResult( - content=[ - TextContent(type="text", text='{"a": 1}'), - ImageContent(type="image", data=img_data, mimeType="image/png"), - EmbeddedResource( - type="resource", - resource=TextResourceContents(uri="file:///tmp/a.txt", text="embedded-text"), - ), - EmbeddedResource( - type="resource", - resource=BlobResourceContents( - uri="file:///tmp/b.bin", - blob=blob_data, - mimeType="application/octet-stream", - ), - ), - ], - structuredContent={"x": 1}, - _meta={"usage": {"prompt_tokens": 2, "completion_tokens": 3}}, - ) - - with patch.object(MCPTool, "invoke_remote_mcp_tool", return_value=result): - messages = list(tool.invoke(session=sqlite_session, user_id="user-1", tool_parameters={"a": 1})) - - types = [m.type for m in messages] - assert ToolInvokeMessage.MessageType.JSON in types - assert ToolInvokeMessage.MessageType.BLOB in types - assert ToolInvokeMessage.MessageType.TEXT in types - assert ToolInvokeMessage.MessageType.VARIABLE in types - assert tool.latest_usage.total_tokens == 5 - - -@pytest.mark.parametrize("sqlite_session", [()], indirect=True) -def test_mcp_tool_invoke_raises_for_unsupported_embedded_resource(sqlite_session: Session): - tool = _build_mcp_tool() - # Use model_construct to bypass pydantic validation and force unsupported resource path. - bad_resource = EmbeddedResource.model_construct(type="resource", resource=object()) - result = CallToolResult(content=[bad_resource], _meta=None) - - with patch.object(MCPTool, "invoke_remote_mcp_tool", return_value=result): - with pytest.raises(ToolInvokeError, match="Unsupported embedded resource type"): - list(tool.invoke(session=sqlite_session, user_id="user-1", tool_parameters={})) - - -def test_mcp_tool_handle_none_parameter_filters_empty_values(): - tool = _build_mcp_tool() - cleaned = tool._handle_none_parameter({"a": 1, "b": None, "c": "", "d": " ", "e": "ok"}) - assert cleaned == {"a": 1, "e": "ok"} - - -# Home Assistant's GetLiveContext declares `domain` as `vol.Any(cv.string, [cv.string])`. -_STRING = {"type": "string"} -_STRING_LIST = {"type": "array", "items": {"type": "string"}} - - -@pytest.mark.parametrize("sqlite_session", [()], indirect=True) -@pytest.mark.parametrize( - "domain_schema", - [ - {"anyOf": [_STRING, _STRING_LIST]}, - {"anyOf": [_STRING_LIST, _STRING]}, - {"oneOf": [_STRING, _STRING_LIST]}, - {"oneOf": [_STRING_LIST, _STRING]}, - {"type": ["string", "array"], "items": _STRING}, - {"type": ["array", "string"], "items": _STRING}, - ], -) -@pytest.mark.parametrize("domain", [["light", "cover"], "light"]) -def test_mcp_tool_invoke_preserves_values_of_mixed_type_union_parameters( - sqlite_session: Session, domain_schema, domain -): - tool = _build_mcp_tool() - schema = {"type": "object", "properties": {"domain": domain_schema, "name": _STRING}, "required": []} - tool.entity.parameters = ToolTransformService.convert_mcp_schema_to_parameter(schema) - result = CallToolResult(content=[TextContent(type="text", text="ok")]) - - with patch.object(MCPTool, "invoke_remote_mcp_tool", return_value=result) as invoke_remote: - list(tool.invoke(session=sqlite_session, user_id="user-1", tool_parameters={"domain": domain, "name": 1})) - - sent = invoke_remote.call_args.args[0] - assert sent["domain"] == domain - # Parameters with a single declared type keep the existing cast. - assert sent["name"] == "1" - - -# ----- M2/M3 user-identity forwarding --------------------------------------- - - -def _build_forwarding_tool(*, mode: str = "idp_token") -> MCPTool: - """Helper that builds an MCPTool with the identity_mode set.""" - entity = ToolEntity( - identity=ToolIdentity( - author="author", - name="remote-tool", - label=I18nObject(en_US="remote-tool"), - provider="provider-id", - ), - parameters=[], - output_schema={}, - ) - return MCPTool( - entity=entity, - runtime=ToolRuntime(tenant_id="tenant-1", invoke_from=InvokeFrom.DEBUGGER), - tenant_id="tenant-1", - icon="icon.svg", - server_url="https://mcp.example.com/mcp/", - server_identifier="mcp_server", - identity_mode=mode, - ) - - -def test_inject_forwarded_identity_stamps_custom_header(): - """The minted SSO token must be placed in X-Dify-SSO-Token; the - workspace-scoped Authorization header and any other custom headers must - pass through untouched so provider credentials keep working.""" - from core.tools.mcp_tool.tool import FORWARDED_IDENTITY_HEADER - - tool = _build_forwarding_tool() - headers: dict[str, str] = {"Authorization": "Bearer static-client-token", "X-Other": "keep"} - - with patch( - "services.enterprise.enterprise_service.EnterpriseService.issue_mcp_token", - return_value=("forwarded.jwt.payload", 1900000000), - ): - tool._inject_forwarded_identity(headers, user_id="alice", app_id=None, audience="https://mcp.example.com/mcp/") - - assert headers[FORWARDED_IDENTITY_HEADER] == "forwarded.jwt.payload" - assert headers["Authorization"] == "Bearer static-client-token" - assert headers["X-Other"] == "keep" - - -def test_inject_forwarded_identity_translates_token_error_to_invoke_error(): - """EnterpriseService failures must surface as ToolInvokeError so the - workflow halts loudly instead of proceeding without identity.""" - from core.tools.mcp_tool.tool import FORWARDED_IDENTITY_HEADER - from services.enterprise.base import MCPNoRefreshTokenError - - tool = _build_forwarding_tool() - headers: dict[str, str] = {} - - with patch( - "services.enterprise.enterprise_service.EnterpriseService.issue_mcp_token", - side_effect=MCPNoRefreshTokenError("please re-sso"), - ): - with pytest.raises(ToolInvokeError, match="forwarded identity token"): - tool._inject_forwarded_identity( - headers, user_id="alice", app_id=None, audience="https://mcp.example.com/mcp/" - ) - - # Headers must NOT have been mutated when token-issuance failed. - assert FORWARDED_IDENTITY_HEADER not in headers - assert "Authorization" not in headers - - -def test_inject_forwarded_identity_sends_end_user_type_for_webapp(): - """A WEB_APP run forwards user_type=end_user so enterprise routes to the - published-webapp token store.""" - tool = _build_forwarding_tool() - tool.runtime = ToolRuntime(tenant_id="tenant-1", invoke_from=InvokeFrom.WEB_APP) - headers: dict[str, str] = {} - - with patch( - "services.enterprise.enterprise_service.EnterpriseService.issue_mcp_token", - return_value=("forwarded.jwt", 1900000000), - ) as issue: - tool._inject_forwarded_identity( - headers, user_id="eu-1", app_id="app-1", audience="https://mcp.example.com/mcp/" - ) - - assert issue.call_args.kwargs["user_type"] == "end_user" - - -def test_inject_forwarded_identity_sends_account_type_for_debugger(): - """A DEBUGGER/console run forwards user_type=account (the existing behaviour).""" - tool = _build_forwarding_tool() # built with InvokeFrom.DEBUGGER - headers: dict[str, str] = {} - - with patch( - "services.enterprise.enterprise_service.EnterpriseService.issue_mcp_token", - return_value=("forwarded.jwt", 1900000000), - ) as issue: - tool._inject_forwarded_identity(headers, user_id="acc-1", app_id=None, audience="https://mcp.example.com/mcp/") - - assert issue.call_args.kwargs["user_type"] == "account" - - -def test_invoke_remote_mcp_tool_fails_closed_when_user_id_missing(config_overrides): - """When forwarding is enabled AND the deployment is enterprise, missing - user_id must raise — never silently invoke as the static identity.""" - tool = _build_forwarding_tool() - - config_overrides(DEPLOYMENT_EDITION=DeploymentEdition.ENTERPRISE) - with pytest.raises(ToolInvokeError, match="no end-user context"): - tool.invoke_remote_mcp_tool({}, user_id=None, app_id=None) - - -def test_invoke_skips_forwarding_outside_enterprise_edition(config_overrides): - """Non-enterprise deployments treat the DB selector as a no-op: a stale - `identity_mode="idp_token"` row must NOT raise (fail-closed) AND must - NOT call the enterprise inner API. The runtime falls through to the - legacy provider-identity path.""" - tool = _build_forwarding_tool() - - config_overrides(DEPLOYMENT_EDITION=DeploymentEdition.COMMUNITY) - # The fail-closed branch must NOT fire (no enterprise → no forwarding). - # The function will still try the legacy DB-load path; we patch that - # to keep the test unit-scoped. - with patch("core.tools.mcp_tool.tool.MCPClientWithAuthRetry") as client_cls: - client_cls.return_value.__enter__.return_value.invoke_tool.return_value = CallToolResult( - content=[], - _meta=None, - ) - with patch.object(tool, "_inject_forwarded_identity") as inject: - with patch("services.tools.mcp_tools_manage_service.MCPToolManageService"): - with patch("core.entities.mcp_provider.MCPProviderEntity.decrypt_server_url", return_value="u"): - with patch("core.entities.mcp_provider.MCPProviderEntity.decrypt_headers", return_value={}): - # Should not raise; should not call enterprise. - with contextlib.suppress(Exception): - tool.invoke_remote_mcp_tool({}, user_id=None, app_id=None) - inject.assert_not_called() diff --git a/api/tests/unit_tests/core/tools/test_tool_parameter_type.py b/api/tests/unit_tests/core/tools/test_tool_parameter_type.py deleted file mode 100644 index 8a41678267a..00000000000 --- a/api/tests/unit_tests/core/tools/test_tool_parameter_type.py +++ /dev/null @@ -1,49 +0,0 @@ -from core.tools.entities.tool_entities import ToolParameter - - -def test_get_parameter_type(): - assert ToolParameter.ToolParameterType.STRING.as_normal_type() == "string" - assert ToolParameter.ToolParameterType.SELECT.as_normal_type() == "string" - assert ToolParameter.ToolParameterType.SECRET_INPUT.as_normal_type() == "string" - assert ToolParameter.ToolParameterType.BOOLEAN.as_normal_type() == "boolean" - assert ToolParameter.ToolParameterType.NUMBER.as_normal_type() == "number" - assert ToolParameter.ToolParameterType.FILE.as_normal_type() == "file" - assert ToolParameter.ToolParameterType.FILES.as_normal_type() == "files" - - -def test_cast_parameter_by_type(): - # string - assert ToolParameter.ToolParameterType.STRING.cast_value("test") == "test" - assert ToolParameter.ToolParameterType.STRING.cast_value(1) == "1" - assert ToolParameter.ToolParameterType.STRING.cast_value(1.0) == "1.0" - assert ToolParameter.ToolParameterType.STRING.cast_value(None) == "" - - # secret input - assert ToolParameter.ToolParameterType.SECRET_INPUT.cast_value("test") == "test" - assert ToolParameter.ToolParameterType.SECRET_INPUT.cast_value(1) == "1" - assert ToolParameter.ToolParameterType.SECRET_INPUT.cast_value(1.0) == "1.0" - assert ToolParameter.ToolParameterType.SECRET_INPUT.cast_value(None) == "" - - # select - assert ToolParameter.ToolParameterType.SELECT.cast_value("test") == "test" - assert ToolParameter.ToolParameterType.SELECT.cast_value(1) == "1" - assert ToolParameter.ToolParameterType.SELECT.cast_value(1.0) == "1.0" - assert ToolParameter.ToolParameterType.SELECT.cast_value(None) == "" - - # boolean - true_values = [True, "True", "true", "1", "YES", "Yes", "yes", "y", "something"] - for value in true_values: - assert ToolParameter.ToolParameterType.BOOLEAN.cast_value(value) is True - - false_values = [False, "False", "false", "0", "NO", "No", "no", "n", None, ""] - for value in false_values: - assert ToolParameter.ToolParameterType.BOOLEAN.cast_value(value) is False - - # number - assert ToolParameter.ToolParameterType.NUMBER.cast_value("1") == 1 - assert ToolParameter.ToolParameterType.NUMBER.cast_value("1.0") == 1.0 - assert ToolParameter.ToolParameterType.NUMBER.cast_value("-1.0") == -1.0 - assert ToolParameter.ToolParameterType.NUMBER.cast_value(1) == 1 - assert ToolParameter.ToolParameterType.NUMBER.cast_value(1.0) == 1.0 - assert ToolParameter.ToolParameterType.NUMBER.cast_value(-1.0) == -1.0 - assert ToolParameter.ToolParameterType.NUMBER.cast_value(None) is None diff --git a/api/tests/unit_tests/core/tools/utils/test_misc_utils_extra.py b/api/tests/unit_tests/core/tools/utils/dataset_retriever/test_dataset_multi_retriever_tool.py similarity index 100% rename from api/tests/unit_tests/core/tools/utils/test_misc_utils_extra.py rename to api/tests/unit_tests/core/tools/utils/dataset_retriever/test_dataset_multi_retriever_tool.py diff --git a/api/tests/unit_tests/core/tools/test_dataset_retriever_tool.py b/api/tests/unit_tests/core/tools/utils/test_dataset_retriever_tool.py similarity index 100% rename from api/tests/unit_tests/core/tools/test_dataset_retriever_tool.py rename to api/tests/unit_tests/core/tools/utils/test_dataset_retriever_tool.py diff --git a/api/tests/unit_tests/core/tools/utils/test_message_transformer.py b/api/tests/unit_tests/core/tools/utils/test_message_transformer.py index 9cf2369fe43..5727add7e82 100644 --- a/api/tests/unit_tests/core/tools/utils/test_message_transformer.py +++ b/api/tests/unit_tests/core/tools/utils/test_message_transformer.py @@ -1,11 +1,18 @@ +import json +from datetime import date, datetime +from decimal import Decimal from typing import Any +from uuid import uuid4 +import numpy as np import pytest +import pytz import core.tools.utils.message_transformer as mt from core.app.entities.app_invoke_entities import InvokeFrom, UserFrom from core.app.file_access import FileAccessScope, bind_file_access_scope, get_current_file_access_scope from core.tools.entities.tool_entities import ToolInvokeMessage +from core.tools.utils.message_transformer import ToolFileMessageTransformer, safe_json_dict, safe_json_value from models.tools import ToolFile @@ -48,113 +55,32 @@ class _FakeToolFileManager: return tool_file -@pytest.fixture(autouse=True) -def _patch_tool_file_manager(monkeypatch: pytest.MonkeyPatch): - # Patch the manager used inside the transformer module - monkeypatch.setattr(mt, "ToolFileManager", _FakeToolFileManager) - # also ensure predictable URL generation (no need to patch; uses id and extension only) - yield - _FakeToolFileManager.last_call = None - - def _gen(messages): yield from messages -def test_transform_tool_invoke_messages_mimetype_key_present_but_none(): - # Arrange: a BLOB message whose meta contains a mime_type key set to None - blob = b"hello" - msg = ToolInvokeMessage( - type=ToolInvokeMessage.MessageType.BLOB, - message=ToolInvokeMessage.BlobMessage(blob=blob), - meta={"mime_type": None, "filename": "greeting"}, - ) +class TestMessageTransformerCases: + @staticmethod + @pytest.fixture(autouse=True) + def _patch_tool_file_manager(monkeypatch: pytest.MonkeyPatch): + # Patch the manager used inside the transformer module + monkeypatch.setattr(mt, "ToolFileManager", _FakeToolFileManager) + # also ensure predictable URL generation (no need to patch; uses id and extension only) + yield + _FakeToolFileManager.last_call = None - # Act - out = list( - mt.ToolFileMessageTransformer.transform_tool_invoke_messages( - messages=_gen([msg]), - user_id="u1", - tenant_id="t1", - conversation_id="c1", + @staticmethod + def test_transform_tool_invoke_messages_mimetype_key_present_but_none(): + # Arrange: a BLOB message whose meta contains a mime_type key set to None + blob = b"hello" + msg = ToolInvokeMessage( + type=ToolInvokeMessage.MessageType.BLOB, + message=ToolInvokeMessage.BlobMessage(blob=blob), + meta={"mime_type": None, "filename": "greeting"}, ) - ) - # Assert: default to application/octet-stream when mime_type is present but None - assert _FakeToolFileManager.last_call is not None - assert _FakeToolFileManager.last_call["mimetype"] == "application/octet-stream" - - # Should yield a BINARY_LINK (not IMAGE_LINK) and the URL ends with .bin - assert len(out) == 1 - o = out[0] - assert o.type == ToolInvokeMessage.MessageType.BINARY_LINK - assert isinstance(o.message, ToolInvokeMessage.TextMessage) - assert o.message.text.endswith(".bin") - # meta is preserved (still contains mime_type: None) - assert "mime_type" in (o.meta or {}) - assert o.meta["mime_type"] is None - assert o.meta["tool_file_id"] == "fake-tool-file-id" - - -def test_transform_tool_invoke_messages_prefers_filename_extension_over_mimetype(): - msg = ToolInvokeMessage( - type=ToolInvokeMessage.MessageType.BLOB, - message=ToolInvokeMessage.BlobMessage(blob=b"docx"), - meta={"mime_type": "application/octet-stream", "filename": "report.docx"}, - ) - - out = list( - mt.ToolFileMessageTransformer.transform_tool_invoke_messages( - messages=_gen([msg]), - user_id="u1", - tenant_id="t1", - conversation_id="c1", - ) - ) - - assert _FakeToolFileManager.last_call is not None - assert _FakeToolFileManager.last_call["filename"] == "report.docx" - assert len(out) == 1 - assert isinstance(out[0].message, ToolInvokeMessage.TextMessage) - assert out[0].message.text.endswith(".docx") - - -def test_transform_tool_invoke_messages_parses_existing_tool_file_link_meta(): - msg = ToolInvokeMessage( - type=ToolInvokeMessage.MessageType.IMAGE_LINK, - message=ToolInvokeMessage.TextMessage(text="/files/tools/existing-tool-file.png"), - meta={}, - ) - - out = list( - mt.ToolFileMessageTransformer.transform_tool_invoke_messages( - messages=_gen([msg]), - user_id="u1", - tenant_id="t1", - conversation_id="c1", - ) - ) - - assert len(out) == 1 - assert out[0].meta["tool_file_id"] == "existing-tool-file" - - -def test_transform_tool_invoke_messages_grants_tool_file_to_current_end_user_scope(): - """Workflow-as-tool LLM nodes need a grant for plugin files (#41169).""" - msg = ToolInvokeMessage( - type=ToolInvokeMessage.MessageType.IMAGE_LINK, - message=ToolInvokeMessage.TextMessage(text="/files/tools/existing-tool-file.png"), - meta={}, - ) - scope = FileAccessScope( - tenant_id="t1", - user_id="end-user-1", - user_from=UserFrom.END_USER, - invoke_from=InvokeFrom.WEB_APP, - ) - - with bind_file_access_scope(scope): - list( + # Act + out = list( mt.ToolFileMessageTransformer.transform_tool_invoke_messages( messages=_gen([msg]), user_id="u1", @@ -162,7 +88,562 @@ def test_transform_tool_invoke_messages_grants_tool_file_to_current_end_user_sco conversation_id="c1", ) ) - current_scope = get_current_file_access_scope() - assert current_scope is not None - assert "existing-tool-file" in current_scope.granted_tool_file_ids + # Assert: default to application/octet-stream when mime_type is present but None + assert _FakeToolFileManager.last_call is not None + assert _FakeToolFileManager.last_call["mimetype"] == "application/octet-stream" + + # Should yield a BINARY_LINK (not IMAGE_LINK) and the URL ends with .bin + assert len(out) == 1 + o = out[0] + assert o.type == ToolInvokeMessage.MessageType.BINARY_LINK + assert isinstance(o.message, ToolInvokeMessage.TextMessage) + assert o.message.text.endswith(".bin") + # meta is preserved (still contains mime_type: None) + assert "mime_type" in (o.meta or {}) + assert o.meta["mime_type"] is None + assert o.meta["tool_file_id"] == "fake-tool-file-id" + + @staticmethod + def test_transform_tool_invoke_messages_prefers_filename_extension_over_mimetype(): + msg = ToolInvokeMessage( + type=ToolInvokeMessage.MessageType.BLOB, + message=ToolInvokeMessage.BlobMessage(blob=b"docx"), + meta={"mime_type": "application/octet-stream", "filename": "report.docx"}, + ) + + out = list( + mt.ToolFileMessageTransformer.transform_tool_invoke_messages( + messages=_gen([msg]), + user_id="u1", + tenant_id="t1", + conversation_id="c1", + ) + ) + + assert _FakeToolFileManager.last_call is not None + assert _FakeToolFileManager.last_call["filename"] == "report.docx" + assert len(out) == 1 + assert isinstance(out[0].message, ToolInvokeMessage.TextMessage) + assert out[0].message.text.endswith(".docx") + + @staticmethod + def test_transform_tool_invoke_messages_parses_existing_tool_file_link_meta(): + msg = ToolInvokeMessage( + type=ToolInvokeMessage.MessageType.IMAGE_LINK, + message=ToolInvokeMessage.TextMessage(text="/files/tools/existing-tool-file.png"), + meta={}, + ) + + out = list( + mt.ToolFileMessageTransformer.transform_tool_invoke_messages( + messages=_gen([msg]), + user_id="u1", + tenant_id="t1", + conversation_id="c1", + ) + ) + + assert len(out) == 1 + assert out[0].meta["tool_file_id"] == "existing-tool-file" + + @staticmethod + def test_transform_tool_invoke_messages_grants_tool_file_to_current_end_user_scope(): + """Workflow-as-tool LLM nodes need a grant for plugin files (#41169).""" + msg = ToolInvokeMessage( + type=ToolInvokeMessage.MessageType.IMAGE_LINK, + message=ToolInvokeMessage.TextMessage(text="/files/tools/existing-tool-file.png"), + meta={}, + ) + scope = FileAccessScope( + tenant_id="t1", + user_id="end-user-1", + user_from=UserFrom.END_USER, + invoke_from=InvokeFrom.WEB_APP, + ) + + with bind_file_access_scope(scope): + list( + mt.ToolFileMessageTransformer.transform_tool_invoke_messages( + messages=_gen([msg]), + user_id="u1", + tenant_id="t1", + conversation_id="c1", + ) + ) + current_scope = get_current_file_access_scope() + + assert current_scope is not None + assert "existing-tool-file" in current_scope.granted_tool_file_ids + + +class TestToolEngineSerializationCases: + class TestSafeJsonValue: + """Test suite for safe_json_value function to ensure proper serialization of complex types""" + + def test_datetime_conversion(self): + """Test datetime conversion with timezone handling""" + # Test datetime with UTC timezone + dt = datetime(2024, 1, 1, 12, 0, 0, tzinfo=pytz.UTC) + result = safe_json_value(dt) + assert isinstance(result, str) + assert "2024-01-01T12:00:00+00:00" in result + + # Test datetime without timezone (should default to UTC) + dt_no_tz = datetime(2024, 1, 1, 12, 0, 0) + result = safe_json_value(dt_no_tz) + assert isinstance(result, str) + # The exact time will depend on the system's timezone, so we check the format + assert "T" in result # ISO format separator + # Check that it's a valid ISO format datetime string + assert len(result) >= 19 # At least YYYY-MM-DDTHH:MM:SS + + def test_date_conversion(self): + """Test date conversion to ISO format""" + test_date = date(2024, 1, 1) + result = safe_json_value(test_date) + assert result == "2024-01-01" + + def test_uuid_conversion(self): + """Test UUID conversion to string""" + test_uuid = uuid4() + result = safe_json_value(test_uuid) + assert isinstance(result, str) + assert result == str(test_uuid) + + def test_decimal_conversion(self): + """Test Decimal conversion to float""" + test_decimal = Decimal("123.456") + result = safe_json_value(test_decimal) + assert result == 123.456 + assert isinstance(result, float) + + def test_bytes_conversion(self): + """Test bytes conversion with UTF-8 decoding""" + # Test valid UTF-8 bytes + test_bytes = b"Hello, World!" + result = safe_json_value(test_bytes) + assert result == "Hello, World!" + + # Test invalid UTF-8 bytes (should fall back to hex) + invalid_bytes = b"\xff\xfe\xfd" + result = safe_json_value(invalid_bytes) + assert result == "fffefd" + + def test_memoryview_conversion(self): + """Test memoryview conversion to hex string""" + test_bytes = b"test data" + test_memoryview = memoryview(test_bytes) + result = safe_json_value(test_memoryview) + assert result == "746573742064617461" # hex of "test data" + + def test_numpy_ndarray_conversion(self): + """Test numpy ndarray conversion to list""" + # Test 1D array + test_array = np.array([1, 2, 3, 4]) + result = safe_json_value(test_array) + assert result == [1, 2, 3, 4] + + # Test 2D array + test_2d_array = np.array([[1, 2], [3, 4]]) + result = safe_json_value(test_2d_array) + assert result == [[1, 2], [3, 4]] + + # Test array with float values + test_float_array = np.array([1.5, 2.7, float("3.14")]) + result = safe_json_value(test_float_array) + assert result == [1.5, 2.7, float("3.14")] + + def test_dict_conversion(self): + """Test dictionary conversion using safe_json_dict""" + test_dict = { + "string": "value", + "number": 42, + "float": float("3.14"), + "boolean": True, + "list": [1, 2, 3], + "nested": {"key": "value"}, + } + result = safe_json_value(test_dict) + assert isinstance(result, dict) + assert result == test_dict + + def test_list_conversion(self): + """Test list conversion with mixed types""" + test_list = [ + "string", + 42, + float("3.14"), + True, + [1, 2, 3], + {"key": "value"}, + datetime(2024, 1, 1, 12, 0, 0, tzinfo=pytz.UTC), + Decimal("123.456"), + uuid4(), + ] + result = safe_json_value(test_list) + assert isinstance(result, list) + assert len(result) == len(test_list) + assert isinstance(result[6], str) # datetime should be converted to string + assert isinstance(result[7], float) # Decimal should be converted to float + assert isinstance(result[8], str) # UUID should be converted to string + + def test_tuple_conversion(self): + """Test tuple conversion to list""" + test_tuple = (1, "string", float("3.14")) + result = safe_json_value(test_tuple) + assert isinstance(result, list) + assert result == [1, "string", float("3.14")] + + def test_set_conversion(self): + """Test set conversion to list""" + test_set = {1, "string", float("3.14")} + result = safe_json_value(test_set) + assert isinstance(result, list) + # Note: set order is not guaranteed, so we check length and content + assert len(result) == 3 + assert 1 in result + assert "string" in result + assert float("3.14") in result + + def test_basic_types_passthrough(self): + """Test that basic types are passed through unchanged""" + assert safe_json_value("string") == "string" + assert safe_json_value(42) == 42 + assert safe_json_value(float("3.14")) == float("3.14") + assert safe_json_value(True) is True + assert safe_json_value(False) is False + assert safe_json_value(None) is None + + def test_nested_complex_structure(self): + """Test complex nested structure with all types""" + complex_data = { + "dates": [date(2024, 1, 1), date(2024, 1, 2)], + "timestamps": [ + datetime(2024, 1, 1, 12, 0, 0, tzinfo=pytz.UTC), + datetime(2024, 1, 2, 12, 0, 0, tzinfo=pytz.UTC), + ], + "numbers": [Decimal("123.456"), Decimal("789.012")], + "identifiers": [uuid4(), uuid4()], + "binary_data": [b"hello", b"world"], + "arrays": [np.array([1, 2, 3]), np.array([4, 5, 6])], + } + + result = safe_json_value(complex_data) + + # Verify structure is maintained + assert isinstance(result, dict) + assert "dates" in result + assert "timestamps" in result + assert "numbers" in result + assert "identifiers" in result + assert "binary_data" in result + assert "arrays" in result + + # Verify conversions + assert all(isinstance(d, str) for d in result["dates"]) + assert all(isinstance(t, str) for t in result["timestamps"]) + assert all(isinstance(n, float) for n in result["numbers"]) + assert all(isinstance(i, str) for i in result["identifiers"]) + assert all(isinstance(b, str) for b in result["binary_data"]) + assert all(isinstance(a, list) for a in result["arrays"]) + + class TestSafeJsonDict: + """Test suite for safe_json_dict function""" + + def test_valid_dict_conversion(self): + """Test valid dictionary conversion""" + test_dict = { + "string": "value", + "number": 42, + "datetime": datetime(2024, 1, 1, 12, 0, 0, tzinfo=pytz.UTC), + "decimal": Decimal("123.456"), + } + result = safe_json_dict(test_dict) + assert isinstance(result, dict) + assert result["string"] == "value" + assert result["number"] == 42 + assert isinstance(result["datetime"], str) + assert isinstance(result["decimal"], float) + + def test_invalid_input_type(self): + """Test that invalid input types raise TypeError""" + with pytest.raises(TypeError, match="safe_json_dict\\(\\) expects a dictionary \\(dict\\) as input"): + safe_json_dict("not a dict") + + with pytest.raises(TypeError, match="safe_json_dict\\(\\) expects a dictionary \\(dict\\) as input"): + safe_json_dict([1, 2, 3]) + + with pytest.raises(TypeError, match="safe_json_dict\\(\\) expects a dictionary \\(dict\\) as input"): + safe_json_dict(42) + + def test_empty_dict(self): + """Test empty dictionary handling""" + result = safe_json_dict({}) + assert result == {} + + def test_nested_dict_conversion(self): + """Test nested dictionary conversion""" + test_dict = { + "level1": { + "level2": { + "datetime": datetime(2024, 1, 1, 12, 0, 0, tzinfo=pytz.UTC), + "decimal": Decimal("123.456"), + } + } + } + result = safe_json_dict(test_dict) + assert isinstance(result["level1"]["level2"]["datetime"], str) + assert isinstance(result["level1"]["level2"]["decimal"], float) + + class TestToolInvokeMessageJsonSerialization: + """Test suite for ToolInvokeMessage JSON serialization through safe_json_value""" + + def test_json_message_serialization(self): + """Test JSON message serialization with complex data""" + complex_data = { + "timestamp": datetime(2024, 1, 1, 12, 0, 0, tzinfo=pytz.UTC), + "amount": Decimal("123.45"), + "id": uuid4(), + "binary": b"test data", + "array": np.array([1, 2, 3]), + } + + # Create JSON message + json_message = ToolInvokeMessage.JsonMessage(json_object=complex_data) + message = ToolInvokeMessage(type=ToolInvokeMessage.MessageType.JSON, message=json_message) + + # Apply safe_json_value transformation + transformed_data = safe_json_value(message.message.json_object) + + # Verify transformations + assert isinstance(transformed_data["timestamp"], str) + assert isinstance(transformed_data["amount"], float) + assert isinstance(transformed_data["id"], str) + assert isinstance(transformed_data["binary"], str) + assert isinstance(transformed_data["array"], list) + + # Verify JSON serialization works + json_string = json.dumps(transformed_data, ensure_ascii=False) + assert isinstance(json_string, str) + + # Verify we can deserialize back + deserialized = json.loads(json_string) + assert deserialized["amount"] == 123.45 + assert deserialized["array"] == [1, 2, 3] + + def test_json_message_with_nested_structures(self): + """Test JSON message with deeply nested complex structures""" + nested_data = { + "level1": { + "level2": { + "level3": { + "dates": [date(2024, 1, 1), date(2024, 1, 2)], + "timestamps": [datetime(2024, 1, 1, 12, 0, 0, tzinfo=pytz.UTC)], + "numbers": [Decimal("1.1"), Decimal("2.2")], + "arrays": [np.array([1, 2]), np.array([3, 4])], + } + } + } + } + + json_message = ToolInvokeMessage.JsonMessage(json_object=nested_data) + message = ToolInvokeMessage(type=ToolInvokeMessage.MessageType.JSON, message=json_message) + + # Transform the data + transformed_data = safe_json_value(message.message.json_object) + + # Verify nested transformations + level3 = transformed_data["level1"]["level2"]["level3"] + assert all(isinstance(d, str) for d in level3["dates"]) + assert all(isinstance(t, str) for t in level3["timestamps"]) + assert all(isinstance(n, float) for n in level3["numbers"]) + assert all(isinstance(a, list) for a in level3["arrays"]) + + # Test JSON serialization + json_string = json.dumps(transformed_data, ensure_ascii=False) + assert isinstance(json_string, str) + + # Verify deserialization + deserialized = json.loads(json_string) + assert deserialized["level1"]["level2"]["level3"]["numbers"] == [1.1, 2.2] + + def test_json_message_transformer_integration(self): + """Test integration with ToolFileMessageTransformer for JSON messages""" + complex_data = { + "metadata": { + "created_at": datetime(2024, 1, 1, 12, 0, 0, tzinfo=pytz.UTC), + "version": Decimal("1.0"), + "tags": ["tag1", "tag2"], + }, + "data": {"values": np.array([1.1, 2.2, 3.3]), "binary": b"binary content"}, + } + + # Create message generator + def message_generator(): + json_message = ToolInvokeMessage.JsonMessage(json_object=complex_data) + message = ToolInvokeMessage(type=ToolInvokeMessage.MessageType.JSON, message=json_message) + yield message + + # Transform messages + transformed_messages = list( + ToolFileMessageTransformer.transform_tool_invoke_messages( + message_generator(), user_id="test_user", tenant_id="test_tenant" + ) + ) + + assert len(transformed_messages) == 1 + transformed_message = transformed_messages[0] + assert transformed_message.type == ToolInvokeMessage.MessageType.JSON + + # Verify the JSON object was transformed + json_obj = transformed_message.message.json_object + assert isinstance(json_obj["metadata"]["created_at"], str) + assert isinstance(json_obj["metadata"]["version"], float) + assert isinstance(json_obj["data"]["values"], list) + assert isinstance(json_obj["data"]["binary"], str) + + # Test final JSON serialization + final_json = json.dumps(json_obj, ensure_ascii=False) + assert isinstance(final_json, str) + + # Verify we can deserialize + deserialized = json.loads(final_json) + assert deserialized["metadata"]["version"] == 1.0 + assert deserialized["data"]["values"] == [1.1, 2.2, 3.3] + + def test_edge_cases_and_error_handling(self): + """Test edge cases and error handling in JSON serialization""" + # Test with None values + data_with_none = {"null_value": None, "empty_string": "", "zero": 0, "false_value": False} + + json_message = ToolInvokeMessage.JsonMessage(json_object=data_with_none) + message = ToolInvokeMessage(type=ToolInvokeMessage.MessageType.JSON, message=json_message) + + transformed_data = safe_json_value(message.message.json_object) + json_string = json.dumps(transformed_data, ensure_ascii=False) + + # Verify serialization works with edge cases + assert json_string is not None + deserialized = json.loads(json_string) + assert deserialized["null_value"] is None + assert deserialized["empty_string"] == "" + assert deserialized["zero"] == 0 + assert deserialized["false_value"] is False + + # Test with very large numbers + large_data = { + "large_int": 2**63 - 1, + "large_float": 1.7976931348623157e308, + "small_float": 2.2250738585072014e-308, + } + + json_message = ToolInvokeMessage.JsonMessage(json_object=large_data) + message = ToolInvokeMessage(type=ToolInvokeMessage.MessageType.JSON, message=json_message) + + transformed_data = safe_json_value(message.message.json_object) + json_string = json.dumps(transformed_data, ensure_ascii=False) + + # Verify large numbers are handled correctly + deserialized = json.loads(json_string) + assert deserialized["large_int"] == 2**63 - 1 + assert deserialized["large_float"] == 1.7976931348623157e308 + assert deserialized["small_float"] == 2.2250738585072014e-308 + + class TestEndToEndSerialization: + """Test suite for end-to-end serialization workflow""" + + def test_complete_workflow_with_real_data(self): + """Test complete workflow from complex data to JSON string and back""" + # Simulate real-world complex data structure + real_world_data = { + "user_profile": { + "id": uuid4(), + "name": "John Doe", + "email": "john@example.com", + "created_at": datetime(2024, 1, 1, 12, 0, 0, tzinfo=pytz.UTC), + "last_login": datetime(2024, 1, 15, 14, 30, 0, tzinfo=pytz.UTC), + "preferences": {"theme": "dark", "language": "en", "timezone": "UTC"}, + }, + "analytics": { + "session_count": 42, + "total_time": Decimal("123.45"), + "metrics": np.array([1.1, 2.2, 3.3, 4.4, 5.5]), + "events": [ + { + "timestamp": datetime(2024, 1, 1, 10, 0, 0, tzinfo=pytz.UTC), + "action": "login", + "duration": Decimal("5.67"), + }, + { + "timestamp": datetime(2024, 1, 1, 11, 0, 0, tzinfo=pytz.UTC), + "action": "logout", + "duration": Decimal("3600.0"), + }, + ], + }, + "files": [ + { + "id": uuid4(), + "name": "document.pdf", + "size": 1024, + "uploaded_at": datetime(2024, 1, 1, 9, 0, 0, tzinfo=pytz.UTC), + "checksum": b"abc123def456", + } + ], + } + + # Step 1: Create ToolInvokeMessage + json_message = ToolInvokeMessage.JsonMessage(json_object=real_world_data) + message = ToolInvokeMessage(type=ToolInvokeMessage.MessageType.JSON, message=json_message) + + # Step 2: Apply safe_json_value transformation + transformed_data = safe_json_value(message.message.json_object) + + # Step 3: Serialize to JSON string + json_string = json.dumps(transformed_data, ensure_ascii=False) + + # Step 4: Verify the string is valid JSON + assert isinstance(json_string, str) + assert json_string.startswith("{") + assert json_string.endswith("}") + + # Step 5: Deserialize back to Python object + deserialized_data = json.loads(json_string) + + # Step 6: Verify data integrity + assert deserialized_data["user_profile"]["name"] == "John Doe" + assert deserialized_data["user_profile"]["email"] == "john@example.com" + assert isinstance(deserialized_data["user_profile"]["created_at"], str) + assert isinstance(deserialized_data["analytics"]["total_time"], float) + assert deserialized_data["analytics"]["total_time"] == 123.45 + assert isinstance(deserialized_data["analytics"]["metrics"], list) + assert deserialized_data["analytics"]["metrics"] == [1.1, 2.2, 3.3, 4.4, 5.5] + assert isinstance(deserialized_data["files"][0]["checksum"], str) + + # Step 7: Verify all complex types were properly converted + self._verify_all_complex_types_converted(deserialized_data) + + def _verify_all_complex_types_converted(self, data): + """Helper method to verify all complex types were properly converted""" + match data: + case dict(): + for key, value in data.items(): + if key in ["id", "checksum"]: + # These should be strings (UUID/bytes converted) + assert isinstance(value, str) + elif key in ["created_at", "last_login", "timestamp", "uploaded_at"]: + # These should be strings (datetime converted) + assert isinstance(value, str) + elif key in ["total_time", "duration"]: + # These should be floats (Decimal converted) + assert isinstance(value, float) + elif key == "metrics": + # This should be a list (ndarray converted) + assert isinstance(value, list) + else: + # Recursively check nested structures + self._verify_all_complex_types_converted(value) + case list(): + for item in data: + self._verify_all_complex_types_converted(item) diff --git a/api/tests/unit_tests/core/tools/utils/test_model_invocation_utils.py b/api/tests/unit_tests/core/tools/utils/test_model_invocation_utils.py index bce7faaf965..dede64ae1a8 100644 --- a/api/tests/unit_tests/core/tools/utils/test_model_invocation_utils.py +++ b/api/tests/unit_tests/core/tools/utils/test_model_invocation_utils.py @@ -31,7 +31,7 @@ from graphon.model_runtime.errors.invoke import ( InvokeServerUnavailableError, ) from models.tools import ToolModelInvoke -from tests.unit_tests.core.model_fixtures import make_model_config, make_model_instance +from tests.helpers.unit.core.model_fixtures import make_model_config, make_model_instance TENANT_ID = "11111111-1111-1111-1111-111111111111" USER_ID = "22222222-2222-2222-2222-222222222222" diff --git a/api/tests/unit_tests/utils/encryption/test_system_encryption.py b/api/tests/unit_tests/core/tools/utils/test_system_encryption.py similarity index 93% rename from api/tests/unit_tests/utils/encryption/test_system_encryption.py rename to api/tests/unit_tests/core/tools/utils/test_system_encryption.py index aa5b2ffb97b..8e017e89e85 100644 --- a/api/tests/unit_tests/utils/encryption/test_system_encryption.py +++ b/api/tests/unit_tests/core/tools/utils/test_system_encryption.py @@ -1,3 +1,5 @@ +from __future__ import annotations + import base64 import hashlib from collections.abc import Callable @@ -8,6 +10,7 @@ from Crypto.Cipher import AES from Crypto.Random import get_random_bytes from Crypto.Util.Padding import pad +from core.tools.utils import system_encryption as encryption from core.tools.utils.system_encryption import ( EncryptionError, SystemEncrypter, @@ -16,6 +19,52 @@ from core.tools.utils.system_encryption import ( encrypt_system_params, get_system_encrypter, ) +from tests.helpers.unit.config_override import apply_config_overrides + + +def test_system_encrypter_roundtrip() -> None: + encrypter = SystemEncrypter(secret_key="test-secret") + payload = {"client_id": "cid", "client_secret": "csecret", "grant_type": "authorization_code"} + + encrypted = encrypter.encrypt_params(payload) + decrypted = encrypter.decrypt_params(encrypted) + + assert encrypted + assert dict(decrypted) == payload + + +def test_system_encrypter_decrypt_validates_input() -> None: + encrypter = SystemEncrypter(secret_key="test-secret") + + with pytest.raises(ValueError, match="must be a string"): + encrypter.decrypt_params(123) # type: ignore[arg-type] + + with pytest.raises(ValueError, match="cannot be empty"): + encrypter.decrypt_params("") + + +def test_system_encrypter_raises_error_for_invalid_ciphertext() -> None: + encrypter = SystemEncrypter(secret_key="test-secret") + + with pytest.raises(EncryptionError, match="Decryption failed"): + encrypter.decrypt_params("not-base64") + + +def test_system_helpers_use_global_cached_instance(monkeypatch: pytest.MonkeyPatch) -> None: + monkeypatch.setattr(encryption, "_encrypter", None) + apply_config_overrides(monkeypatch, SECRET_KEY="global-secret") + + first = encryption.get_system_encrypter() + second = encryption.get_system_encrypter() + assert first is second + + encrypted = encryption.encrypt_system_params({"k": "v"}) + assert encryption.decrypt_system_params(encrypted) == {"k": "v"} + + +def test_create_system_encrypter_factory() -> None: + encrypter = encryption.create_system_encrypter(secret_key="factory-secret") + assert isinstance(encrypter, SystemEncrypter) class TestSystemEncrypter: diff --git a/api/tests/unit_tests/core/tools/utils/test_system_oauth_encryption.py b/api/tests/unit_tests/core/tools/utils/test_system_oauth_encryption.py deleted file mode 100644 index d94e290d46f..00000000000 --- a/api/tests/unit_tests/core/tools/utils/test_system_oauth_encryption.py +++ /dev/null @@ -1,52 +0,0 @@ -from __future__ import annotations - -import pytest - -from core.tools.utils import system_encryption as encryption -from core.tools.utils.system_encryption import EncryptionError, SystemEncrypter -from tests.unit_tests.config_override import apply_config_overrides - - -def test_system_encrypter_roundtrip() -> None: - encrypter = SystemEncrypter(secret_key="test-secret") - payload = {"client_id": "cid", "client_secret": "csecret", "grant_type": "authorization_code"} - - encrypted = encrypter.encrypt_params(payload) - decrypted = encrypter.decrypt_params(encrypted) - - assert encrypted - assert dict(decrypted) == payload - - -def test_system_encrypter_decrypt_validates_input() -> None: - encrypter = SystemEncrypter(secret_key="test-secret") - - with pytest.raises(ValueError, match="must be a string"): - encrypter.decrypt_params(123) # type: ignore[arg-type] - - with pytest.raises(ValueError, match="cannot be empty"): - encrypter.decrypt_params("") - - -def test_system_encrypter_raises_error_for_invalid_ciphertext() -> None: - encrypter = SystemEncrypter(secret_key="test-secret") - - with pytest.raises(EncryptionError, match="Decryption failed"): - encrypter.decrypt_params("not-base64") - - -def test_system_helpers_use_global_cached_instance(monkeypatch: pytest.MonkeyPatch) -> None: - monkeypatch.setattr(encryption, "_encrypter", None) - apply_config_overrides(monkeypatch, SECRET_KEY="global-secret") - - first = encryption.get_system_encrypter() - second = encryption.get_system_encrypter() - assert first is second - - encrypted = encryption.encrypt_system_params({"k": "v"}) - assert encryption.decrypt_system_params(encrypted) == {"k": "v"} - - -def test_create_system_encrypter_factory() -> None: - encrypter = encryption.create_system_encrypter(secret_key="factory-secret") - assert isinstance(encrypter, SystemEncrypter) diff --git a/api/tests/unit_tests/utils/test_text_processing.py b/api/tests/unit_tests/core/tools/utils/test_text_processing_utils.py similarity index 100% rename from api/tests/unit_tests/utils/test_text_processing.py rename to api/tests/unit_tests/core/tools/utils/test_text_processing_utils.py diff --git a/api/tests/unit_tests/core/tools/utils/test_tool_engine_serialization.py b/api/tests/unit_tests/core/tools/utils/test_tool_engine_serialization.py deleted file mode 100644 index 24821ea85bf..00000000000 --- a/api/tests/unit_tests/core/tools/utils/test_tool_engine_serialization.py +++ /dev/null @@ -1,482 +0,0 @@ -import json -from datetime import date, datetime -from decimal import Decimal -from uuid import uuid4 - -import numpy as np -import pytest -import pytz - -from core.tools.entities.tool_entities import ToolInvokeMessage -from core.tools.utils.message_transformer import ToolFileMessageTransformer, safe_json_dict, safe_json_value - - -class TestSafeJsonValue: - """Test suite for safe_json_value function to ensure proper serialization of complex types""" - - def test_datetime_conversion(self): - """Test datetime conversion with timezone handling""" - # Test datetime with UTC timezone - dt = datetime(2024, 1, 1, 12, 0, 0, tzinfo=pytz.UTC) - result = safe_json_value(dt) - assert isinstance(result, str) - assert "2024-01-01T12:00:00+00:00" in result - - # Test datetime without timezone (should default to UTC) - dt_no_tz = datetime(2024, 1, 1, 12, 0, 0) - result = safe_json_value(dt_no_tz) - assert isinstance(result, str) - # The exact time will depend on the system's timezone, so we check the format - assert "T" in result # ISO format separator - # Check that it's a valid ISO format datetime string - assert len(result) >= 19 # At least YYYY-MM-DDTHH:MM:SS - - def test_date_conversion(self): - """Test date conversion to ISO format""" - test_date = date(2024, 1, 1) - result = safe_json_value(test_date) - assert result == "2024-01-01" - - def test_uuid_conversion(self): - """Test UUID conversion to string""" - test_uuid = uuid4() - result = safe_json_value(test_uuid) - assert isinstance(result, str) - assert result == str(test_uuid) - - def test_decimal_conversion(self): - """Test Decimal conversion to float""" - test_decimal = Decimal("123.456") - result = safe_json_value(test_decimal) - assert result == 123.456 - assert isinstance(result, float) - - def test_bytes_conversion(self): - """Test bytes conversion with UTF-8 decoding""" - # Test valid UTF-8 bytes - test_bytes = b"Hello, World!" - result = safe_json_value(test_bytes) - assert result == "Hello, World!" - - # Test invalid UTF-8 bytes (should fall back to hex) - invalid_bytes = b"\xff\xfe\xfd" - result = safe_json_value(invalid_bytes) - assert result == "fffefd" - - def test_memoryview_conversion(self): - """Test memoryview conversion to hex string""" - test_bytes = b"test data" - test_memoryview = memoryview(test_bytes) - result = safe_json_value(test_memoryview) - assert result == "746573742064617461" # hex of "test data" - - def test_numpy_ndarray_conversion(self): - """Test numpy ndarray conversion to list""" - # Test 1D array - test_array = np.array([1, 2, 3, 4]) - result = safe_json_value(test_array) - assert result == [1, 2, 3, 4] - - # Test 2D array - test_2d_array = np.array([[1, 2], [3, 4]]) - result = safe_json_value(test_2d_array) - assert result == [[1, 2], [3, 4]] - - # Test array with float values - test_float_array = np.array([1.5, 2.7, float("3.14")]) - result = safe_json_value(test_float_array) - assert result == [1.5, 2.7, float("3.14")] - - def test_dict_conversion(self): - """Test dictionary conversion using safe_json_dict""" - test_dict = { - "string": "value", - "number": 42, - "float": float("3.14"), - "boolean": True, - "list": [1, 2, 3], - "nested": {"key": "value"}, - } - result = safe_json_value(test_dict) - assert isinstance(result, dict) - assert result == test_dict - - def test_list_conversion(self): - """Test list conversion with mixed types""" - test_list = [ - "string", - 42, - float("3.14"), - True, - [1, 2, 3], - {"key": "value"}, - datetime(2024, 1, 1, 12, 0, 0, tzinfo=pytz.UTC), - Decimal("123.456"), - uuid4(), - ] - result = safe_json_value(test_list) - assert isinstance(result, list) - assert len(result) == len(test_list) - assert isinstance(result[6], str) # datetime should be converted to string - assert isinstance(result[7], float) # Decimal should be converted to float - assert isinstance(result[8], str) # UUID should be converted to string - - def test_tuple_conversion(self): - """Test tuple conversion to list""" - test_tuple = (1, "string", float("3.14")) - result = safe_json_value(test_tuple) - assert isinstance(result, list) - assert result == [1, "string", float("3.14")] - - def test_set_conversion(self): - """Test set conversion to list""" - test_set = {1, "string", float("3.14")} - result = safe_json_value(test_set) - assert isinstance(result, list) - # Note: set order is not guaranteed, so we check length and content - assert len(result) == 3 - assert 1 in result - assert "string" in result - assert float("3.14") in result - - def test_basic_types_passthrough(self): - """Test that basic types are passed through unchanged""" - assert safe_json_value("string") == "string" - assert safe_json_value(42) == 42 - assert safe_json_value(float("3.14")) == float("3.14") - assert safe_json_value(True) is True - assert safe_json_value(False) is False - assert safe_json_value(None) is None - - def test_nested_complex_structure(self): - """Test complex nested structure with all types""" - complex_data = { - "dates": [date(2024, 1, 1), date(2024, 1, 2)], - "timestamps": [ - datetime(2024, 1, 1, 12, 0, 0, tzinfo=pytz.UTC), - datetime(2024, 1, 2, 12, 0, 0, tzinfo=pytz.UTC), - ], - "numbers": [Decimal("123.456"), Decimal("789.012")], - "identifiers": [uuid4(), uuid4()], - "binary_data": [b"hello", b"world"], - "arrays": [np.array([1, 2, 3]), np.array([4, 5, 6])], - } - - result = safe_json_value(complex_data) - - # Verify structure is maintained - assert isinstance(result, dict) - assert "dates" in result - assert "timestamps" in result - assert "numbers" in result - assert "identifiers" in result - assert "binary_data" in result - assert "arrays" in result - - # Verify conversions - assert all(isinstance(d, str) for d in result["dates"]) - assert all(isinstance(t, str) for t in result["timestamps"]) - assert all(isinstance(n, float) for n in result["numbers"]) - assert all(isinstance(i, str) for i in result["identifiers"]) - assert all(isinstance(b, str) for b in result["binary_data"]) - assert all(isinstance(a, list) for a in result["arrays"]) - - -class TestSafeJsonDict: - """Test suite for safe_json_dict function""" - - def test_valid_dict_conversion(self): - """Test valid dictionary conversion""" - test_dict = { - "string": "value", - "number": 42, - "datetime": datetime(2024, 1, 1, 12, 0, 0, tzinfo=pytz.UTC), - "decimal": Decimal("123.456"), - } - result = safe_json_dict(test_dict) - assert isinstance(result, dict) - assert result["string"] == "value" - assert result["number"] == 42 - assert isinstance(result["datetime"], str) - assert isinstance(result["decimal"], float) - - def test_invalid_input_type(self): - """Test that invalid input types raise TypeError""" - with pytest.raises(TypeError, match="safe_json_dict\\(\\) expects a dictionary \\(dict\\) as input"): - safe_json_dict("not a dict") - - with pytest.raises(TypeError, match="safe_json_dict\\(\\) expects a dictionary \\(dict\\) as input"): - safe_json_dict([1, 2, 3]) - - with pytest.raises(TypeError, match="safe_json_dict\\(\\) expects a dictionary \\(dict\\) as input"): - safe_json_dict(42) - - def test_empty_dict(self): - """Test empty dictionary handling""" - result = safe_json_dict({}) - assert result == {} - - def test_nested_dict_conversion(self): - """Test nested dictionary conversion""" - test_dict = { - "level1": { - "level2": {"datetime": datetime(2024, 1, 1, 12, 0, 0, tzinfo=pytz.UTC), "decimal": Decimal("123.456")} - } - } - result = safe_json_dict(test_dict) - assert isinstance(result["level1"]["level2"]["datetime"], str) - assert isinstance(result["level1"]["level2"]["decimal"], float) - - -class TestToolInvokeMessageJsonSerialization: - """Test suite for ToolInvokeMessage JSON serialization through safe_json_value""" - - def test_json_message_serialization(self): - """Test JSON message serialization with complex data""" - complex_data = { - "timestamp": datetime(2024, 1, 1, 12, 0, 0, tzinfo=pytz.UTC), - "amount": Decimal("123.45"), - "id": uuid4(), - "binary": b"test data", - "array": np.array([1, 2, 3]), - } - - # Create JSON message - json_message = ToolInvokeMessage.JsonMessage(json_object=complex_data) - message = ToolInvokeMessage(type=ToolInvokeMessage.MessageType.JSON, message=json_message) - - # Apply safe_json_value transformation - transformed_data = safe_json_value(message.message.json_object) - - # Verify transformations - assert isinstance(transformed_data["timestamp"], str) - assert isinstance(transformed_data["amount"], float) - assert isinstance(transformed_data["id"], str) - assert isinstance(transformed_data["binary"], str) - assert isinstance(transformed_data["array"], list) - - # Verify JSON serialization works - json_string = json.dumps(transformed_data, ensure_ascii=False) - assert isinstance(json_string, str) - - # Verify we can deserialize back - deserialized = json.loads(json_string) - assert deserialized["amount"] == 123.45 - assert deserialized["array"] == [1, 2, 3] - - def test_json_message_with_nested_structures(self): - """Test JSON message with deeply nested complex structures""" - nested_data = { - "level1": { - "level2": { - "level3": { - "dates": [date(2024, 1, 1), date(2024, 1, 2)], - "timestamps": [datetime(2024, 1, 1, 12, 0, 0, tzinfo=pytz.UTC)], - "numbers": [Decimal("1.1"), Decimal("2.2")], - "arrays": [np.array([1, 2]), np.array([3, 4])], - } - } - } - } - - json_message = ToolInvokeMessage.JsonMessage(json_object=nested_data) - message = ToolInvokeMessage(type=ToolInvokeMessage.MessageType.JSON, message=json_message) - - # Transform the data - transformed_data = safe_json_value(message.message.json_object) - - # Verify nested transformations - level3 = transformed_data["level1"]["level2"]["level3"] - assert all(isinstance(d, str) for d in level3["dates"]) - assert all(isinstance(t, str) for t in level3["timestamps"]) - assert all(isinstance(n, float) for n in level3["numbers"]) - assert all(isinstance(a, list) for a in level3["arrays"]) - - # Test JSON serialization - json_string = json.dumps(transformed_data, ensure_ascii=False) - assert isinstance(json_string, str) - - # Verify deserialization - deserialized = json.loads(json_string) - assert deserialized["level1"]["level2"]["level3"]["numbers"] == [1.1, 2.2] - - def test_json_message_transformer_integration(self): - """Test integration with ToolFileMessageTransformer for JSON messages""" - complex_data = { - "metadata": { - "created_at": datetime(2024, 1, 1, 12, 0, 0, tzinfo=pytz.UTC), - "version": Decimal("1.0"), - "tags": ["tag1", "tag2"], - }, - "data": {"values": np.array([1.1, 2.2, 3.3]), "binary": b"binary content"}, - } - - # Create message generator - def message_generator(): - json_message = ToolInvokeMessage.JsonMessage(json_object=complex_data) - message = ToolInvokeMessage(type=ToolInvokeMessage.MessageType.JSON, message=json_message) - yield message - - # Transform messages - transformed_messages = list( - ToolFileMessageTransformer.transform_tool_invoke_messages( - message_generator(), user_id="test_user", tenant_id="test_tenant" - ) - ) - - assert len(transformed_messages) == 1 - transformed_message = transformed_messages[0] - assert transformed_message.type == ToolInvokeMessage.MessageType.JSON - - # Verify the JSON object was transformed - json_obj = transformed_message.message.json_object - assert isinstance(json_obj["metadata"]["created_at"], str) - assert isinstance(json_obj["metadata"]["version"], float) - assert isinstance(json_obj["data"]["values"], list) - assert isinstance(json_obj["data"]["binary"], str) - - # Test final JSON serialization - final_json = json.dumps(json_obj, ensure_ascii=False) - assert isinstance(final_json, str) - - # Verify we can deserialize - deserialized = json.loads(final_json) - assert deserialized["metadata"]["version"] == 1.0 - assert deserialized["data"]["values"] == [1.1, 2.2, 3.3] - - def test_edge_cases_and_error_handling(self): - """Test edge cases and error handling in JSON serialization""" - # Test with None values - data_with_none = {"null_value": None, "empty_string": "", "zero": 0, "false_value": False} - - json_message = ToolInvokeMessage.JsonMessage(json_object=data_with_none) - message = ToolInvokeMessage(type=ToolInvokeMessage.MessageType.JSON, message=json_message) - - transformed_data = safe_json_value(message.message.json_object) - json_string = json.dumps(transformed_data, ensure_ascii=False) - - # Verify serialization works with edge cases - assert json_string is not None - deserialized = json.loads(json_string) - assert deserialized["null_value"] is None - assert deserialized["empty_string"] == "" - assert deserialized["zero"] == 0 - assert deserialized["false_value"] is False - - # Test with very large numbers - large_data = { - "large_int": 2**63 - 1, - "large_float": 1.7976931348623157e308, - "small_float": 2.2250738585072014e-308, - } - - json_message = ToolInvokeMessage.JsonMessage(json_object=large_data) - message = ToolInvokeMessage(type=ToolInvokeMessage.MessageType.JSON, message=json_message) - - transformed_data = safe_json_value(message.message.json_object) - json_string = json.dumps(transformed_data, ensure_ascii=False) - - # Verify large numbers are handled correctly - deserialized = json.loads(json_string) - assert deserialized["large_int"] == 2**63 - 1 - assert deserialized["large_float"] == 1.7976931348623157e308 - assert deserialized["small_float"] == 2.2250738585072014e-308 - - -class TestEndToEndSerialization: - """Test suite for end-to-end serialization workflow""" - - def test_complete_workflow_with_real_data(self): - """Test complete workflow from complex data to JSON string and back""" - # Simulate real-world complex data structure - real_world_data = { - "user_profile": { - "id": uuid4(), - "name": "John Doe", - "email": "john@example.com", - "created_at": datetime(2024, 1, 1, 12, 0, 0, tzinfo=pytz.UTC), - "last_login": datetime(2024, 1, 15, 14, 30, 0, tzinfo=pytz.UTC), - "preferences": {"theme": "dark", "language": "en", "timezone": "UTC"}, - }, - "analytics": { - "session_count": 42, - "total_time": Decimal("123.45"), - "metrics": np.array([1.1, 2.2, 3.3, 4.4, 5.5]), - "events": [ - { - "timestamp": datetime(2024, 1, 1, 10, 0, 0, tzinfo=pytz.UTC), - "action": "login", - "duration": Decimal("5.67"), - }, - { - "timestamp": datetime(2024, 1, 1, 11, 0, 0, tzinfo=pytz.UTC), - "action": "logout", - "duration": Decimal("3600.0"), - }, - ], - }, - "files": [ - { - "id": uuid4(), - "name": "document.pdf", - "size": 1024, - "uploaded_at": datetime(2024, 1, 1, 9, 0, 0, tzinfo=pytz.UTC), - "checksum": b"abc123def456", - } - ], - } - - # Step 1: Create ToolInvokeMessage - json_message = ToolInvokeMessage.JsonMessage(json_object=real_world_data) - message = ToolInvokeMessage(type=ToolInvokeMessage.MessageType.JSON, message=json_message) - - # Step 2: Apply safe_json_value transformation - transformed_data = safe_json_value(message.message.json_object) - - # Step 3: Serialize to JSON string - json_string = json.dumps(transformed_data, ensure_ascii=False) - - # Step 4: Verify the string is valid JSON - assert isinstance(json_string, str) - assert json_string.startswith("{") - assert json_string.endswith("}") - - # Step 5: Deserialize back to Python object - deserialized_data = json.loads(json_string) - - # Step 6: Verify data integrity - assert deserialized_data["user_profile"]["name"] == "John Doe" - assert deserialized_data["user_profile"]["email"] == "john@example.com" - assert isinstance(deserialized_data["user_profile"]["created_at"], str) - assert isinstance(deserialized_data["analytics"]["total_time"], float) - assert deserialized_data["analytics"]["total_time"] == 123.45 - assert isinstance(deserialized_data["analytics"]["metrics"], list) - assert deserialized_data["analytics"]["metrics"] == [1.1, 2.2, 3.3, 4.4, 5.5] - assert isinstance(deserialized_data["files"][0]["checksum"], str) - - # Step 7: Verify all complex types were properly converted - self._verify_all_complex_types_converted(deserialized_data) - - def _verify_all_complex_types_converted(self, data): - """Helper method to verify all complex types were properly converted""" - match data: - case dict(): - for key, value in data.items(): - if key in ["id", "checksum"]: - # These should be strings (UUID/bytes converted) - assert isinstance(value, str) - elif key in ["created_at", "last_login", "timestamp", "uploaded_at"]: - # These should be strings (datetime converted) - assert isinstance(value, str) - elif key in ["total_time", "duration"]: - # These should be floats (Decimal converted) - assert isinstance(value, float) - elif key == "metrics": - # This should be a list (ndarray converted) - assert isinstance(value, list) - else: - # Recursively check nested structures - self._verify_all_complex_types_converted(value) - case list(): - for item in data: - self._verify_all_complex_types_converted(item) diff --git a/api/tests/unit_tests/utils/yaml/test_yaml_utils.py b/api/tests/unit_tests/core/tools/utils/test_yaml_utils.py similarity index 100% rename from api/tests/unit_tests/utils/yaml/test_yaml_utils.py rename to api/tests/unit_tests/core/tools/utils/test_yaml_utils.py diff --git a/api/tests/unit_tests/core/tools/workflow_as_tool/test_provider.py b/api/tests/unit_tests/core/tools/workflow_as_tool/test_provider.py index a034d5f72fc..30513ad0deb 100644 --- a/api/tests/unit_tests/core/tools/workflow_as_tool/test_provider.py +++ b/api/tests/unit_tests/core/tools/workflow_as_tool/test_provider.py @@ -32,7 +32,7 @@ from models.base import TypeBase from models.model import App, AppMode from models.tools import WorkflowToolProvider from models.workflow import Workflow -from tests.unit_tests.model_factories import make_account, make_app, make_workflow +from tests.helpers.unit.model_factories import make_account, make_app, make_workflow @pytest.fixture diff --git a/api/tests/unit_tests/core/tools/workflow_as_tool/test_tool.py b/api/tests/unit_tests/core/tools/workflow_as_tool/test_tool.py index 1985a32889a..d066c8fe771 100644 --- a/api/tests/unit_tests/core/tools/workflow_as_tool/test_tool.py +++ b/api/tests/unit_tests/core/tools/workflow_as_tool/test_tool.py @@ -36,7 +36,7 @@ from models.model import App, AppMode, EndUser from models.tools import WorkflowToolProvider from models.workflow import Workflow, WorkflowRun, WorkflowType from services.workflow_run_agg import WorkflowRunAgg -from tests.unit_tests.model_factories import make_account +from tests.helpers.unit.model_factories import make_account TENANT_ID = "00000000-0000-0000-0000-000000000001" OTHER_TENANT_ID = "00000000-0000-0000-0000-000000000002" diff --git a/api/tests/unit_tests/core/trigger/debug/test_debug_event_selectors.py b/api/tests/unit_tests/core/trigger/debug/test_debug_event_selectors.py deleted file mode 100644 index def06d9f0e9..00000000000 --- a/api/tests/unit_tests/core/trigger/debug/test_debug_event_selectors.py +++ /dev/null @@ -1,309 +0,0 @@ -""" -Tests for core.trigger.debug.event_selectors. - -Covers: Plugin/Webhook/Schedule pollers, create_event_poller factory, -and select_trigger_debug_events orchestrator. -""" - -from __future__ import annotations - -import json -from datetime import datetime -from typing import cast -from unittest.mock import MagicMock, patch - -import pytest -from typing_extensions import TypedDict - -from core.plugin.entities.request import TriggerInvokeEventResponse -from core.trigger.constants import ( - TRIGGER_PLUGIN_NODE_TYPE, - TRIGGER_SCHEDULE_NODE_TYPE, - TRIGGER_WEBHOOK_NODE_TYPE, -) -from core.trigger.debug.event_selectors import ( - PluginTriggerDebugEventPoller, - ScheduleTriggerDebugEventPoller, - WebhookTriggerDebugEventPoller, - create_event_poller, - select_trigger_debug_events, -) -from core.trigger.debug.events import PluginTriggerDebugEvent, WebhookDebugEvent -from graphon.entities.graph_config import NodeConfigDict -from graphon.enums import BuiltinNodeTypes, NodeType, WorkflowType -from models.model import App, AppMode -from models.workflow import Workflow -from tests.unit_tests.core.trigger.conftest import VALID_PROVIDER_ID -from tests.unit_tests.model_factories import make_app - - -class _PollerArgs(TypedDict, closed=True): - tenant_id: str - user_id: str - app_id: str - node_config: NodeConfigDict - node_id: str - - -def _make_poller_args(node_config: dict[str, object] | None = None) -> _PollerArgs: - return { - "tenant_id": "t1", - "user_id": "u1", - "app_id": "a1", - # The pollers only read node_config["data"]; the tests do not build a full graph node config. - "node_config": cast(NodeConfigDict, node_config or {"data": {}}), - "node_id": "n1", - } - - -def _plugin_node_config(provider_id: str = VALID_PROVIDER_ID) -> dict[str, object]: - """Valid node config for TriggerEventNodeData.model_validate.""" - return { - "data": { - "title": "test", - "plugin_id": "org/testplugin", - "provider_id": provider_id, - "event_name": "push", - "subscription_id": "s1", - "plugin_unique_identifier": "uid-1", - } - } - - -def _workflow_with_node(node_type: NodeType | None) -> Workflow: - nodes: list[dict[str, object]] = ( - [] if node_type is None else [{"id": "n1", "data": {"type": node_type, "title": "Trigger"}}] - ) - return Workflow.new( - tenant_id="t1", - app_id="a1", - type=WorkflowType.WORKFLOW.value, - version=Workflow.VERSION_DRAFT, - graph=json.dumps({"nodes": nodes, "edges": []}), - features="{}", - created_by="u1", - environment_variables=[], - conversation_variables=[], - rag_pipeline_variables=[], - ) - - -def _app() -> App: - return make_app( - app_id="a1", - tenant_id="t1", - name="Trigger App", - mode=AppMode.WORKFLOW, - icon_type=None, - max_active_requests=0, - ) - - -class TestPluginTriggerDebugEventPoller: - @patch("core.trigger.debug.event_selectors.TriggerDebugEventBus") - def test_returns_workflow_args_on_success(self, mock_bus: MagicMock) -> None: - event = PluginTriggerDebugEvent( - timestamp=100, - name="push", - user_id="u1", - request_id="r1", - subscription_id="s1", - provider_id="p1", - ) - mock_bus.poll.return_value = event - - with patch("services.trigger.trigger_service.TriggerService") as mock_trigger_svc: - mock_trigger_svc.invoke_trigger_event.return_value = TriggerInvokeEventResponse( - variables={"repo": "dify"}, - cancelled=False, - ) - - poller = PluginTriggerDebugEventPoller(**_make_poller_args(_plugin_node_config())) - result = poller.poll() - - assert result is not None - assert result.workflow_args["inputs"] == {"repo": "dify"} - - @patch("core.trigger.debug.event_selectors.TriggerDebugEventBus") - def test_returns_none_when_no_event(self, mock_bus: MagicMock) -> None: - mock_bus.poll.return_value = None - - poller = PluginTriggerDebugEventPoller(**_make_poller_args(_plugin_node_config())) - - assert poller.poll() is None - - @patch("core.trigger.debug.event_selectors.TriggerDebugEventBus") - def test_returns_none_when_invoke_cancelled(self, mock_bus: MagicMock) -> None: - event = PluginTriggerDebugEvent( - timestamp=100, - name="push", - user_id="u1", - request_id="r1", - subscription_id="s1", - provider_id="p1", - ) - mock_bus.poll.return_value = event - - with patch("services.trigger.trigger_service.TriggerService") as mock_trigger_svc: - mock_trigger_svc.invoke_trigger_event.return_value = TriggerInvokeEventResponse( - variables={}, - cancelled=True, - ) - - poller = PluginTriggerDebugEventPoller(**_make_poller_args(_plugin_node_config())) - - assert poller.poll() is None - - -class TestWebhookTriggerDebugEventPoller: - @patch("core.trigger.debug.event_selectors.TriggerDebugEventBus") - def test_uses_inputs_directly_when_present(self, mock_bus: MagicMock) -> None: - event = WebhookDebugEvent( - timestamp=100, - request_id="r1", - node_id="n1", - payload={"inputs": {"key": "val"}, "webhook_data": {}}, - ) - mock_bus.poll.return_value = event - - poller = WebhookTriggerDebugEventPoller(**_make_poller_args()) - result = poller.poll() - - assert result is not None - assert result.workflow_args["inputs"] == {"key": "val"} - - @patch("core.trigger.debug.event_selectors.TriggerDebugEventBus") - def test_falls_back_to_webhook_data(self, mock_bus: MagicMock) -> None: - event = WebhookDebugEvent( - timestamp=100, - request_id="r1", - node_id="n1", - payload={"webhook_data": {"body": "raw"}}, - ) - mock_bus.poll.return_value = event - - with patch("services.trigger.webhook_service.WebhookService") as mock_webhook_svc: - mock_webhook_svc.build_workflow_inputs.return_value = {"parsed": "data"} - - poller = WebhookTriggerDebugEventPoller(**_make_poller_args()) - result = poller.poll() - - assert result is not None - assert result.workflow_args["inputs"] == {"parsed": "data"} - mock_webhook_svc.build_workflow_inputs.assert_called_once_with({"body": "raw"}) - - @patch("core.trigger.debug.event_selectors.TriggerDebugEventBus") - def test_returns_none_when_no_event(self, mock_bus: MagicMock) -> None: - mock_bus.poll.return_value = None - poller = WebhookTriggerDebugEventPoller(**_make_poller_args()) - - assert poller.poll() is None - - -class TestScheduleTriggerDebugEventPoller: - @patch("core.trigger.debug.event_selectors.redis_client") - @patch("core.trigger.debug.event_selectors.naive_utc_now") - @patch("core.trigger.debug.event_selectors.calculate_next_run_at") - @patch("core.trigger.debug.event_selectors.ensure_naive_utc") - def test_returns_none_when_not_yet_due( - self, mock_ensure: MagicMock, mock_calc: MagicMock, mock_now: MagicMock, mock_redis: MagicMock - ) -> None: - now = datetime(2025, 1, 1, 12, 0, 0) - next_run = datetime(2025, 1, 1, 13, 0, 0) # future - mock_now.return_value = now - mock_calc.return_value = next_run - mock_ensure.return_value = next_run - mock_redis.get.return_value = None - - with patch("services.trigger.schedule_service.ScheduleService") as mock_schedule_svc: - mock_schedule_config = MagicMock() - mock_schedule_config.cron_expression = "0 * * * *" - mock_schedule_config.timezone = "UTC" - mock_schedule_svc.to_schedule_config.return_value = mock_schedule_config - - poller = ScheduleTriggerDebugEventPoller(**_make_poller_args()) - - assert poller.poll() is None - - @patch("core.trigger.debug.event_selectors.redis_client") - @patch("core.trigger.debug.event_selectors.naive_utc_now") - @patch("core.trigger.debug.event_selectors.calculate_next_run_at") - @patch("core.trigger.debug.event_selectors.ensure_naive_utc") - def test_fires_event_when_due( - self, mock_ensure: MagicMock, mock_calc: MagicMock, mock_now: MagicMock, mock_redis: MagicMock - ) -> None: - now = datetime(2025, 1, 1, 14, 0, 0) - next_run = datetime(2025, 1, 1, 12, 0, 0) # past - mock_now.return_value = now - mock_calc.return_value = next_run - mock_ensure.return_value = next_run - mock_redis.get.return_value = None - - with patch("services.trigger.schedule_service.ScheduleService") as mock_schedule_svc: - mock_schedule_config = MagicMock() - mock_schedule_config.cron_expression = "0 * * * *" - mock_schedule_config.timezone = "UTC" - mock_schedule_svc.to_schedule_config.return_value = mock_schedule_config - - poller = ScheduleTriggerDebugEventPoller(**_make_poller_args()) - result = poller.poll() - - assert result is not None - mock_redis.delete.assert_called_once() - - -class TestCreateEventPoller: - def test_creates_plugin_poller(self) -> None: - wf = _workflow_with_node(TRIGGER_PLUGIN_NODE_TYPE) - poller = create_event_poller(wf, "t1", "u1", "a1", "n1") - assert isinstance(poller, PluginTriggerDebugEventPoller) - - def test_creates_webhook_poller(self) -> None: - wf = _workflow_with_node(TRIGGER_WEBHOOK_NODE_TYPE) - poller = create_event_poller(wf, "t1", "u1", "a1", "n1") - assert isinstance(poller, WebhookTriggerDebugEventPoller) - - def test_creates_schedule_poller(self) -> None: - wf = _workflow_with_node(TRIGGER_SCHEDULE_NODE_TYPE) - poller = create_event_poller(wf, "t1", "u1", "a1", "n1") - assert isinstance(poller, ScheduleTriggerDebugEventPoller) - - def test_raises_for_unknown_type(self) -> None: - wf = _workflow_with_node(BuiltinNodeTypes.START) - - with pytest.raises(ValueError, match="unable to create event poller for node type start") as exc_info: - create_event_poller(wf, "t1", "u1", "a1", "n1") - - assert str(exc_info.value) == "unable to create event poller for node type start" - - def test_raises_when_node_config_missing(self) -> None: - wf = _workflow_with_node(None) - - with ( - patch.object(Workflow, "get_node_config_by_id", return_value=None), - pytest.raises(ValueError, match="\\('Node data not found for node %s', 'n1'\\)"), - ): - create_event_poller(wf, "t1", "u1", "a1", "n1") - - -class TestSelectTriggerDebugEvents: - def test_returns_first_non_none_event(self) -> None: - wf = _workflow_with_node(TRIGGER_WEBHOOK_NODE_TYPE) - app_model = _app() - - with patch.object(WebhookTriggerDebugEventPoller, "poll") as mock_poll: - expected = MagicMock() - mock_poll.return_value = expected - - result = select_trigger_debug_events(wf, app_model, "u1", ["n1", "n2"]) - - assert result is expected - - def test_returns_none_when_no_events(self) -> None: - wf = _workflow_with_node(TRIGGER_WEBHOOK_NODE_TYPE) - app_model = _app() - - with patch.object(WebhookTriggerDebugEventPoller, "poll", return_value=None): - result = select_trigger_debug_events(wf, app_model, "u1", ["n1"]) - - assert result is None diff --git a/api/tests/unit_tests/core/trigger/debug/test_debug_event_bus.py b/api/tests/unit_tests/core/trigger/debug/test_event_bus.py similarity index 100% rename from api/tests/unit_tests/core/trigger/debug/test_debug_event_bus.py rename to api/tests/unit_tests/core/trigger/debug/test_event_bus.py diff --git a/api/tests/unit_tests/core/trigger/debug/test_event_selectors.py b/api/tests/unit_tests/core/trigger/debug/test_event_selectors.py new file mode 100644 index 00000000000..9ec9151bcce --- /dev/null +++ b/api/tests/unit_tests/core/trigger/debug/test_event_selectors.py @@ -0,0 +1,406 @@ +from __future__ import annotations + +import hashlib +import json +from datetime import UTC, datetime +from typing import cast +from unittest.mock import MagicMock, patch + +import pytest +import pytz +from typing_extensions import TypedDict + +from core.plugin.entities.request import TriggerInvokeEventResponse +from core.trigger.constants import ( + TRIGGER_PLUGIN_NODE_TYPE, + TRIGGER_SCHEDULE_NODE_TYPE, + TRIGGER_WEBHOOK_NODE_TYPE, +) +from core.trigger.debug import event_selectors +from core.trigger.debug.event_selectors import ( + PluginTriggerDebugEventPoller, + ScheduleTriggerDebugEventPoller, + WebhookTriggerDebugEventPoller, + create_event_poller, + select_trigger_debug_events, +) +from core.trigger.debug.events import PluginTriggerDebugEvent, WebhookDebugEvent +from core.workflow.nodes.trigger_schedule.entities import ScheduleConfig +from graphon.entities.graph_config import NodeConfigDict +from graphon.enums import BuiltinNodeTypes, NodeType, WorkflowType +from models.model import App, AppMode +from models.workflow import Workflow +from tests.helpers.unit.model_factories import make_app +from tests.unit_tests.core.trigger.conftest import VALID_PROVIDER_ID + + +class _DummyRedis: + def __init__(self): + self.store: dict[str, str] = {} + + def get(self, key: str): + return self.store.get(key) + + def setex(self, name: str, time: int, value: str): + self.store[name] = value + + def expire(self, name: str, ttl: int): + # Expiration not required for these tests. + pass + + def delete(self, name: str): + self.store.pop(name, None) + + +def _make_poller( + monkeypatch: pytest.MonkeyPatch, redis_client: _DummyRedis +) -> event_selectors.ScheduleTriggerDebugEventPoller: + monkeypatch.setattr(event_selectors, "redis_client", redis_client) + return event_selectors.ScheduleTriggerDebugEventPoller( + tenant_id="tenant-1", + user_id="user-1", + app_id="app-1", + node_config={"id": "node-1", "data": {"mode": "cron"}}, + node_id="node-1", + ) + + +class TestTriggerDebugEventSelectorsCases: + @staticmethod + @pytest.fixture + def dummy_schedule_config() -> ScheduleConfig: + return ScheduleConfig( + node_id="node-1", + cron_expression="* * * * *", + timezone="Asia/Shanghai", + ) + + @staticmethod + @pytest.fixture(autouse=True) + def patch_schedule_service(monkeypatch: pytest.MonkeyPatch, dummy_schedule_config: ScheduleConfig): + # Ensure poller always receives the deterministic config. + monkeypatch.setattr( + "services.trigger.schedule_service.ScheduleService.to_schedule_config", + staticmethod(lambda *_args, **_kwargs: dummy_schedule_config), + ) + + @staticmethod + def test_schedule_poller_handles_aware_next_run(monkeypatch: pytest.MonkeyPatch): + redis_client = _DummyRedis() + poller = _make_poller(monkeypatch, redis_client) + + base_now = datetime(2025, 1, 1, 12, 0, 10) + aware_next_run = datetime(2025, 1, 1, 12, 0, 5, tzinfo=UTC) + + monkeypatch.setattr(event_selectors, "naive_utc_now", lambda: base_now) + monkeypatch.setattr(event_selectors, "calculate_next_run_at", lambda *_: aware_next_run) + + event = poller.poll() + + assert event is not None + assert event.node_id == "node-1" + assert event.workflow_args["inputs"] == {} + + @staticmethod + def test_schedule_runtime_cache_normalizes_timezone( + monkeypatch: pytest.MonkeyPatch, dummy_schedule_config: ScheduleConfig + ): + redis_client = _DummyRedis() + poller = _make_poller(monkeypatch, redis_client) + + localized_time = pytz.timezone("Asia/Shanghai").localize(datetime(2025, 1, 1, 20, 0, 0)) + + cron_hash = hashlib.sha256(dummy_schedule_config.cron_expression.encode()).hexdigest() + cache_key = poller.schedule_debug_runtime_key(cron_hash) + + redis_client.store[cache_key] = json.dumps( + { + "cache_key": cache_key, + "timezone": dummy_schedule_config.timezone, + "cron_expression": dummy_schedule_config.cron_expression, + "next_run_at": localized_time.isoformat(), + } + ) + + runtime = poller.get_or_create_schedule_debug_runtime() + + expected = localized_time.astimezone(UTC).replace(tzinfo=None) + assert runtime.next_run_at == expected + assert runtime.next_run_at.tzinfo is None + + +""" +Tests for core.trigger.debug.event_selectors. + +Covers: Plugin/Webhook/Schedule pollers, create_event_poller factory, +and select_trigger_debug_events orchestrator. +""" + + +class _PollerArgs(TypedDict, closed=True): + tenant_id: str + user_id: str + app_id: str + node_config: NodeConfigDict + node_id: str + + +def _make_poller_args(node_config: dict[str, object] | None = None) -> _PollerArgs: + return { + "tenant_id": "t1", + "user_id": "u1", + "app_id": "a1", + # The pollers only read node_config["data"]; the tests do not build a full graph node config. + "node_config": cast(NodeConfigDict, node_config or {"data": {}}), + "node_id": "n1", + } + + +def _plugin_node_config(provider_id: str = VALID_PROVIDER_ID) -> dict[str, object]: + """Valid node config for TriggerEventNodeData.model_validate.""" + return { + "data": { + "title": "test", + "plugin_id": "org/testplugin", + "provider_id": provider_id, + "event_name": "push", + "subscription_id": "s1", + "plugin_unique_identifier": "uid-1", + } + } + + +def _workflow_with_node(node_type: NodeType | None) -> Workflow: + nodes: list[dict[str, object]] = ( + [] if node_type is None else [{"id": "n1", "data": {"type": node_type, "title": "Trigger"}}] + ) + return Workflow.new( + tenant_id="t1", + app_id="a1", + type=WorkflowType.WORKFLOW.value, + version=Workflow.VERSION_DRAFT, + graph=json.dumps({"nodes": nodes, "edges": []}), + features="{}", + created_by="u1", + environment_variables=[], + conversation_variables=[], + rag_pipeline_variables=[], + ) + + +def _app() -> App: + return make_app( + app_id="a1", + tenant_id="t1", + name="Trigger App", + mode=AppMode.WORKFLOW, + icon_type=None, + max_active_requests=0, + ) + + +class TestDebugEventSelectorsCases: + class TestPluginTriggerDebugEventPoller: + @patch("core.trigger.debug.event_selectors.TriggerDebugEventBus") + def test_returns_workflow_args_on_success(self, mock_bus: MagicMock) -> None: + event = PluginTriggerDebugEvent( + timestamp=100, + name="push", + user_id="u1", + request_id="r1", + subscription_id="s1", + provider_id="p1", + ) + mock_bus.poll.return_value = event + + with patch("services.trigger.trigger_service.TriggerService") as mock_trigger_svc: + mock_trigger_svc.invoke_trigger_event.return_value = TriggerInvokeEventResponse( + variables={"repo": "dify"}, + cancelled=False, + ) + + poller = PluginTriggerDebugEventPoller(**_make_poller_args(_plugin_node_config())) + result = poller.poll() + + assert result is not None + assert result.workflow_args["inputs"] == {"repo": "dify"} + + @patch("core.trigger.debug.event_selectors.TriggerDebugEventBus") + def test_returns_none_when_no_event(self, mock_bus: MagicMock) -> None: + mock_bus.poll.return_value = None + + poller = PluginTriggerDebugEventPoller(**_make_poller_args(_plugin_node_config())) + + assert poller.poll() is None + + @patch("core.trigger.debug.event_selectors.TriggerDebugEventBus") + def test_returns_none_when_invoke_cancelled(self, mock_bus: MagicMock) -> None: + event = PluginTriggerDebugEvent( + timestamp=100, + name="push", + user_id="u1", + request_id="r1", + subscription_id="s1", + provider_id="p1", + ) + mock_bus.poll.return_value = event + + with patch("services.trigger.trigger_service.TriggerService") as mock_trigger_svc: + mock_trigger_svc.invoke_trigger_event.return_value = TriggerInvokeEventResponse( + variables={}, + cancelled=True, + ) + + poller = PluginTriggerDebugEventPoller(**_make_poller_args(_plugin_node_config())) + + assert poller.poll() is None + + class TestWebhookTriggerDebugEventPoller: + @patch("core.trigger.debug.event_selectors.TriggerDebugEventBus") + def test_uses_inputs_directly_when_present(self, mock_bus: MagicMock) -> None: + event = WebhookDebugEvent( + timestamp=100, + request_id="r1", + node_id="n1", + payload={"inputs": {"key": "val"}, "webhook_data": {}}, + ) + mock_bus.poll.return_value = event + + poller = WebhookTriggerDebugEventPoller(**_make_poller_args()) + result = poller.poll() + + assert result is not None + assert result.workflow_args["inputs"] == {"key": "val"} + + @patch("core.trigger.debug.event_selectors.TriggerDebugEventBus") + def test_falls_back_to_webhook_data(self, mock_bus: MagicMock) -> None: + event = WebhookDebugEvent( + timestamp=100, + request_id="r1", + node_id="n1", + payload={"webhook_data": {"body": "raw"}}, + ) + mock_bus.poll.return_value = event + + with patch("services.trigger.webhook_service.WebhookService") as mock_webhook_svc: + mock_webhook_svc.build_workflow_inputs.return_value = {"parsed": "data"} + + poller = WebhookTriggerDebugEventPoller(**_make_poller_args()) + result = poller.poll() + + assert result is not None + assert result.workflow_args["inputs"] == {"parsed": "data"} + mock_webhook_svc.build_workflow_inputs.assert_called_once_with({"body": "raw"}) + + @patch("core.trigger.debug.event_selectors.TriggerDebugEventBus") + def test_returns_none_when_no_event(self, mock_bus: MagicMock) -> None: + mock_bus.poll.return_value = None + poller = WebhookTriggerDebugEventPoller(**_make_poller_args()) + + assert poller.poll() is None + + class TestScheduleTriggerDebugEventPoller: + @patch("core.trigger.debug.event_selectors.redis_client") + @patch("core.trigger.debug.event_selectors.naive_utc_now") + @patch("core.trigger.debug.event_selectors.calculate_next_run_at") + @patch("core.trigger.debug.event_selectors.ensure_naive_utc") + def test_returns_none_when_not_yet_due( + self, mock_ensure: MagicMock, mock_calc: MagicMock, mock_now: MagicMock, mock_redis: MagicMock + ) -> None: + now = datetime(2025, 1, 1, 12, 0, 0) + next_run = datetime(2025, 1, 1, 13, 0, 0) # future + mock_now.return_value = now + mock_calc.return_value = next_run + mock_ensure.return_value = next_run + mock_redis.get.return_value = None + + with patch("services.trigger.schedule_service.ScheduleService") as mock_schedule_svc: + mock_schedule_config = MagicMock() + mock_schedule_config.cron_expression = "0 * * * *" + mock_schedule_config.timezone = "UTC" + mock_schedule_svc.to_schedule_config.return_value = mock_schedule_config + + poller = ScheduleTriggerDebugEventPoller(**_make_poller_args()) + + assert poller.poll() is None + + @patch("core.trigger.debug.event_selectors.redis_client") + @patch("core.trigger.debug.event_selectors.naive_utc_now") + @patch("core.trigger.debug.event_selectors.calculate_next_run_at") + @patch("core.trigger.debug.event_selectors.ensure_naive_utc") + def test_fires_event_when_due( + self, mock_ensure: MagicMock, mock_calc: MagicMock, mock_now: MagicMock, mock_redis: MagicMock + ) -> None: + now = datetime(2025, 1, 1, 14, 0, 0) + next_run = datetime(2025, 1, 1, 12, 0, 0) # past + mock_now.return_value = now + mock_calc.return_value = next_run + mock_ensure.return_value = next_run + mock_redis.get.return_value = None + + with patch("services.trigger.schedule_service.ScheduleService") as mock_schedule_svc: + mock_schedule_config = MagicMock() + mock_schedule_config.cron_expression = "0 * * * *" + mock_schedule_config.timezone = "UTC" + mock_schedule_svc.to_schedule_config.return_value = mock_schedule_config + + poller = ScheduleTriggerDebugEventPoller(**_make_poller_args()) + result = poller.poll() + + assert result is not None + mock_redis.delete.assert_called_once() + + class TestCreateEventPoller: + def test_creates_plugin_poller(self) -> None: + wf = _workflow_with_node(TRIGGER_PLUGIN_NODE_TYPE) + poller = create_event_poller(wf, "t1", "u1", "a1", "n1") + assert isinstance(poller, PluginTriggerDebugEventPoller) + + def test_creates_webhook_poller(self) -> None: + wf = _workflow_with_node(TRIGGER_WEBHOOK_NODE_TYPE) + poller = create_event_poller(wf, "t1", "u1", "a1", "n1") + assert isinstance(poller, WebhookTriggerDebugEventPoller) + + def test_creates_schedule_poller(self) -> None: + wf = _workflow_with_node(TRIGGER_SCHEDULE_NODE_TYPE) + poller = create_event_poller(wf, "t1", "u1", "a1", "n1") + assert isinstance(poller, ScheduleTriggerDebugEventPoller) + + def test_raises_for_unknown_type(self) -> None: + wf = _workflow_with_node(BuiltinNodeTypes.START) + + with pytest.raises(ValueError, match="unable to create event poller for node type start") as exc_info: + create_event_poller(wf, "t1", "u1", "a1", "n1") + + assert str(exc_info.value) == "unable to create event poller for node type start" + + def test_raises_when_node_config_missing(self) -> None: + wf = _workflow_with_node(None) + + with ( + patch.object(Workflow, "get_node_config_by_id", return_value=None), + pytest.raises(ValueError, match="\\('Node data not found for node %s', 'n1'\\)"), + ): + create_event_poller(wf, "t1", "u1", "a1", "n1") + + class TestSelectTriggerDebugEvents: + def test_returns_first_non_none_event(self) -> None: + wf = _workflow_with_node(TRIGGER_WEBHOOK_NODE_TYPE) + app_model = _app() + + with patch.object(WebhookTriggerDebugEventPoller, "poll") as mock_poll: + expected = MagicMock() + mock_poll.return_value = expected + + result = select_trigger_debug_events(wf, app_model, "u1", ["n1", "n2"]) + + assert result is expected + + def test_returns_none_when_no_events(self) -> None: + wf = _workflow_with_node(TRIGGER_WEBHOOK_NODE_TYPE) + app_model = _app() + + with patch.object(WebhookTriggerDebugEventPoller, "poll", return_value=None): + result = select_trigger_debug_events(wf, app_model, "u1", ["n1"]) + + assert result is None diff --git a/api/tests/unit_tests/core/trigger/utils/test_utils_encryption.py b/api/tests/unit_tests/core/trigger/utils/test_encryption.py similarity index 100% rename from api/tests/unit_tests/core/trigger/utils/test_utils_encryption.py rename to api/tests/unit_tests/core/trigger/utils/test_encryption.py diff --git a/api/tests/unit_tests/core/trigger/utils/test_utils_endpoint.py b/api/tests/unit_tests/core/trigger/utils/test_endpoint.py similarity index 100% rename from api/tests/unit_tests/core/trigger/utils/test_utils_endpoint.py rename to api/tests/unit_tests/core/trigger/utils/test_endpoint.py diff --git a/api/tests/unit_tests/core/trigger/utils/test_utils_locks.py b/api/tests/unit_tests/core/trigger/utils/test_locks.py similarity index 100% rename from api/tests/unit_tests/core/trigger/utils/test_utils_locks.py rename to api/tests/unit_tests/core/trigger/utils/test_locks.py diff --git a/api/tests/unit_tests/core/workflow/generator/test_prompts.py b/api/tests/unit_tests/core/workflow/generator/prompts/test_planner_prompts.py similarity index 100% rename from api/tests/unit_tests/core/workflow/generator/test_prompts.py rename to api/tests/unit_tests/core/workflow/generator/prompts/test_planner_prompts.py diff --git a/api/tests/unit_tests/core/workflow/generator/test_runner.py b/api/tests/unit_tests/core/workflow/generator/test_runner.py index 89f192fa286..f75fc26a06f 100644 --- a/api/tests/unit_tests/core/workflow/generator/test_runner.py +++ b/api/tests/unit_tests/core/workflow/generator/test_runner.py @@ -1,12 +1,3 @@ -""" -Unit tests for ``WorkflowGenerator``. - -The runner is pure domain logic — we mock the injected ``model_instance`` and -assert that the postprocessor produces a structurally valid graph for both -Workflow and Advanced-Chat modes, plus that obvious failure paths surface a -readable error envelope. -""" - import json import threading import time @@ -17,10 +8,26 @@ from unittest.mock import MagicMock, patch import pytest from jinja2 import Template -from core.workflow.generator.runner import WorkflowGenerator, _find_planned_tool_entry +from core.workflow.generator.runner import ( + WorkflowGenerator, + _find_planned_tool_entry, + _stage_error_to_envelope_code, + _StageJSONError, + _StageSchemaError, +) from core.workflow.generator.tool_catalogue import ToolCatalogueEntry from core.workflow.generator.types import GraphDict -from tests.unit_tests.config_override import apply_config_overrides +from graphon.model_runtime.errors.invoke import InvokeError +from tests.helpers.unit.config_override import apply_config_overrides + +""" +Unit tests for ``WorkflowGenerator``. + +The runner is pure domain logic — we mock the injected ``model_instance`` and +assert that the postprocessor produces a structurally valid graph for both +Workflow and Advanced-Chat modes, plus that obvious failure paths surface a +readable error envelope. +""" def _llm_result(text: str) -> MagicMock: @@ -4698,3 +4705,173 @@ class TestWorkflowGeneratorStream: ) assert result["mode"] == "workflow" + + +def test_stage_schema_error(): + err = _StageSchemaError("planner", "missing key") + assert str(err) == "planner schema invalid: missing key" + assert err.stage == "planner" + + +def test_stage_error_to_envelope_code(): + err = InvokeError("invoke err") + assert _stage_error_to_envelope_code(err) == "MODEL_ERROR" + + err2 = _StageJSONError("builder", "json err") + assert _stage_error_to_envelope_code(err2) == "INVALID_JSON" + + err3 = ValueError("other err") + assert _stage_error_to_envelope_code(err3) == "MODEL_ERROR" + + +def test_declares_variable(): + # BuiltinNodeTypes is not imported directly, we need to mock or just use the generator method + + # LLM node + assert WorkflowGenerator._declares_variable({"data": {"type": "llm"}}, "text") is True + llm_so = {"data": {"type": "llm", "structured_output": {"schema": {"properties": {"json_var": {}}}}}} + assert WorkflowGenerator._declares_variable(llm_so, "json_var") is True + assert WorkflowGenerator._declares_variable(llm_so, "other_var") is False + + # Code node + assert ( + WorkflowGenerator._declares_variable({"data": {"type": "code", "outputs": {"code_var": "str"}}}, "code_var") + is True + ) + + # Knowledge retrieval + assert WorkflowGenerator._declares_variable({"data": {"type": "knowledge-retrieval"}}, "result") is True + + # Parameter extractor + assert ( + WorkflowGenerator._declares_variable( + {"data": {"type": "parameter-extractor", "parameters": [{"name": "param1"}]}}, "param1" + ) + is True + ) + + # HTTP request + assert WorkflowGenerator._declares_variable({"data": {"type": "http-request"}}, "body") is True + + # Template transform + assert WorkflowGenerator._declares_variable({"data": {"type": "template-transform"}}, "output") is True + + # Tool + assert WorkflowGenerator._declares_variable({"data": {"type": "tool"}}, "anything") is True + + # Other node (default false) + assert WorkflowGenerator._declares_variable({"data": {"type": "unknown"}}, "anything") is False + + +def test_collect_container_errors(): + # Not a container error + nodes = [{"id": "n1", "data": {"type": "llm"}}, {"id": "n2", "data": {"type": "code", "parentId": "n1"}}] + errors = WorkflowGenerator._collect_container_errors(nodes=nodes) + assert len(errors) >= 1 + assert errors[0]["code"] == "INVALID_CONTAINER" + + # Missing terminal error + nodes = [ + {"id": "n1", "data": {"type": "llm"}}, + ] + # Test would go here but we need to see what _validate_structure calls + + +def test_collect_unknown_tools(): + # Missing tool/provider info + nodes = [ + {"id": "n1", "data": {"type": "tool", "provider_id": "", "tool_name": ""}}, + {"id": "n2", "data": {"type": "tool", "provider_id": "p1", "tool_name": "t1"}}, + ] + installed_tools = {("p1", "t1")} + errors = WorkflowGenerator._collect_unknown_tools(nodes=nodes, installed_tools=installed_tools) + assert len(errors) >= 1 + assert "missing provider" in errors[0]["detail"] + + # Needs a real structure, skipping for now + + +def test_collect_unresolved_refs(): + # Missing node ref + nodes = [{"id": "n1", "data": {"type": "llm", "prompt_template": [{"text": "{#unknown.var#}"}]}}] + # To trigger the parsing we need to mock _collect_refs_in_data behavior or let it parse naturally + # If the ref parsing finds "unknown", "var", it will check by_id + + # Actually _collect_refs_in_data modifies the set + + errors = WorkflowGenerator._collect_unresolved_refs(nodes=nodes, mode="workflow") + # Actually wait, _collect_refs_in_data needs the actual variable payload format, not just prompt_template string + # Let's mock _collect_refs_in_data + + class MockGenerator(WorkflowGenerator): + @classmethod + def _collect_refs_in_data(cls, data, refs): + refs.add(("unknown", "var")) + + errors = MockGenerator._collect_unresolved_refs(nodes=nodes, mode="workflow") + assert len(errors) >= 1 + assert errors[0]["code"] == "UNKNOWN_NODE_REFERENCE" + + +def test_collect_edge_cycle_errors(): + # Self-loop + graph = {"nodes": [{"id": "n1"}], "edges": [{"source": "n1", "target": "n1"}]} + errors = WorkflowGenerator._collect_edge_cycle_errors(graph=graph, known_ids={"n1"}) + assert len(errors) >= 1 + assert "itself" in errors[0]["detail"] + + +def test_collect_container_errors_empty_container(): + # Empty container + nodes = [ + {"id": "n1", "data": {"type": "iteration"}}, + ] + errors = WorkflowGenerator._collect_container_errors(nodes=nodes) + assert len(errors) >= 1 + assert "no child nodes" in errors[0]["detail"] + + +def test_collect_container_errors_cycle(): + # Ancestor cycle + nodes = [ + {"id": "n1", "data": {"type": "iteration", "parentId": "n2"}}, + {"id": "n2", "data": {"type": "iteration", "parentId": "n1"}}, + ] + errors = WorkflowGenerator._collect_container_errors(nodes=nodes) + assert len(errors) >= 1 + assert "Cycle" in errors[0]["detail"] + + +def test_postprocess_graph_edges(): + # Try calling _postprocess_graph directly to trigger _sanitize_node_ids + graph = { + "nodes": [{"id": "sys", "data": {"type": "start"}}, {"id": "bad id", "data": {"type": "llm"}}], + "edges": [{"source": "sys", "target": "bad id", "id": "edge_1"}], + } + + # Just mocking methods to reach the sanitize part or call directly + WorkflowGenerator._sanitize_node_ids(nodes=graph["nodes"], edges=graph["edges"]) + assert graph["nodes"][1]["id"] != "bad id" + + +def test_repair_branch_edge_handles(): + nodes = [{"id": "n1", "data": {"type": "question-classifier", "classes": [{"id": "c1", "name": "c1"}]}}] + edges = [{"source": "n1", "target": "n2", "sourceHandle": ""}] + + WorkflowGenerator._repair_branch_edge_handles(nodes=nodes, edges=edges) + assert edges[0]["sourceHandle"] == "c1" # assuming it falls back to first one + + +def test_document_extractor_start_vars(): + nodes = [{"id": "n1", "data": {"type": "document-extractor", "variable_selector": ["start", "doc"]}}] + res = WorkflowGenerator._document_extractor_start_vars(nodes=nodes, start_id="start") + assert res == {"doc": False} + + +def test_missing_terminal_mode_auto(): + # Empty graph, should get MISSING_TERMINAL + graph = {"nodes": [{"id": "n1", "data": {"type": "start"}}], "edges": []} + + # Missing terminal check happens inside _validate_structure + WorkflowGenerator._validate_structure(graph=graph, mode="workflow", installed_tools=set()) + # Mocking this deeply is hard, but we can verify it doesn't fail diff --git a/api/tests/unit_tests/core/workflow/generator/test_runner_missing.py b/api/tests/unit_tests/core/workflow/generator/test_runner_missing.py deleted file mode 100644 index 57d9a324511..00000000000 --- a/api/tests/unit_tests/core/workflow/generator/test_runner_missing.py +++ /dev/null @@ -1,177 +0,0 @@ -from core.workflow.generator.runner import ( - WorkflowGenerator, - _stage_error_to_envelope_code, - _StageJSONError, - _StageSchemaError, -) -from graphon.model_runtime.errors.invoke import InvokeError - - -def test_stage_schema_error(): - err = _StageSchemaError("planner", "missing key") - assert str(err) == "planner schema invalid: missing key" - assert err.stage == "planner" - - -def test_stage_error_to_envelope_code(): - err = InvokeError("invoke err") - assert _stage_error_to_envelope_code(err) == "MODEL_ERROR" - - err2 = _StageJSONError("builder", "json err") - assert _stage_error_to_envelope_code(err2) == "INVALID_JSON" - - err3 = ValueError("other err") - assert _stage_error_to_envelope_code(err3) == "MODEL_ERROR" - - -def test_declares_variable(): - # BuiltinNodeTypes is not imported directly, we need to mock or just use the generator method - - # LLM node - assert WorkflowGenerator._declares_variable({"data": {"type": "llm"}}, "text") is True - llm_so = {"data": {"type": "llm", "structured_output": {"schema": {"properties": {"json_var": {}}}}}} - assert WorkflowGenerator._declares_variable(llm_so, "json_var") is True - assert WorkflowGenerator._declares_variable(llm_so, "other_var") is False - - # Code node - assert ( - WorkflowGenerator._declares_variable({"data": {"type": "code", "outputs": {"code_var": "str"}}}, "code_var") - is True - ) - - # Knowledge retrieval - assert WorkflowGenerator._declares_variable({"data": {"type": "knowledge-retrieval"}}, "result") is True - - # Parameter extractor - assert ( - WorkflowGenerator._declares_variable( - {"data": {"type": "parameter-extractor", "parameters": [{"name": "param1"}]}}, "param1" - ) - is True - ) - - # HTTP request - assert WorkflowGenerator._declares_variable({"data": {"type": "http-request"}}, "body") is True - - # Template transform - assert WorkflowGenerator._declares_variable({"data": {"type": "template-transform"}}, "output") is True - - # Tool - assert WorkflowGenerator._declares_variable({"data": {"type": "tool"}}, "anything") is True - - # Other node (default false) - assert WorkflowGenerator._declares_variable({"data": {"type": "unknown"}}, "anything") is False - - -def test_collect_container_errors(): - # Not a container error - nodes = [{"id": "n1", "data": {"type": "llm"}}, {"id": "n2", "data": {"type": "code", "parentId": "n1"}}] - errors = WorkflowGenerator._collect_container_errors(nodes=nodes) - assert len(errors) >= 1 - assert errors[0]["code"] == "INVALID_CONTAINER" - - # Missing terminal error - nodes = [ - {"id": "n1", "data": {"type": "llm"}}, - ] - # Test would go here but we need to see what _validate_structure calls - - -def test_collect_unknown_tools(): - # Missing tool/provider info - nodes = [ - {"id": "n1", "data": {"type": "tool", "provider_id": "", "tool_name": ""}}, - {"id": "n2", "data": {"type": "tool", "provider_id": "p1", "tool_name": "t1"}}, - ] - installed_tools = {("p1", "t1")} - errors = WorkflowGenerator._collect_unknown_tools(nodes=nodes, installed_tools=installed_tools) - assert len(errors) >= 1 - assert "missing provider" in errors[0]["detail"] - - # Needs a real structure, skipping for now - - -def test_collect_unresolved_refs(): - # Missing node ref - nodes = [{"id": "n1", "data": {"type": "llm", "prompt_template": [{"text": "{#unknown.var#}"}]}}] - # To trigger the parsing we need to mock _collect_refs_in_data behavior or let it parse naturally - # If the ref parsing finds "unknown", "var", it will check by_id - - # Actually _collect_refs_in_data modifies the set - - errors = WorkflowGenerator._collect_unresolved_refs(nodes=nodes, mode="workflow") - # Actually wait, _collect_refs_in_data needs the actual variable payload format, not just prompt_template string - # Let's mock _collect_refs_in_data - - class MockGenerator(WorkflowGenerator): - @classmethod - def _collect_refs_in_data(cls, data, refs): - refs.add(("unknown", "var")) - - errors = MockGenerator._collect_unresolved_refs(nodes=nodes, mode="workflow") - assert len(errors) >= 1 - assert errors[0]["code"] == "UNKNOWN_NODE_REFERENCE" - - -def test_collect_edge_cycle_errors(): - # Self-loop - graph = {"nodes": [{"id": "n1"}], "edges": [{"source": "n1", "target": "n1"}]} - errors = WorkflowGenerator._collect_edge_cycle_errors(graph=graph, known_ids={"n1"}) - assert len(errors) >= 1 - assert "itself" in errors[0]["detail"] - - -def test_collect_container_errors_empty_container(): - # Empty container - nodes = [ - {"id": "n1", "data": {"type": "iteration"}}, - ] - errors = WorkflowGenerator._collect_container_errors(nodes=nodes) - assert len(errors) >= 1 - assert "no child nodes" in errors[0]["detail"] - - -def test_collect_container_errors_cycle(): - # Ancestor cycle - nodes = [ - {"id": "n1", "data": {"type": "iteration", "parentId": "n2"}}, - {"id": "n2", "data": {"type": "iteration", "parentId": "n1"}}, - ] - errors = WorkflowGenerator._collect_container_errors(nodes=nodes) - assert len(errors) >= 1 - assert "Cycle" in errors[0]["detail"] - - -def test_postprocess_graph_edges(): - # Try calling _postprocess_graph directly to trigger _sanitize_node_ids - graph = { - "nodes": [{"id": "sys", "data": {"type": "start"}}, {"id": "bad id", "data": {"type": "llm"}}], - "edges": [{"source": "sys", "target": "bad id", "id": "edge_1"}], - } - - # Just mocking methods to reach the sanitize part or call directly - WorkflowGenerator._sanitize_node_ids(nodes=graph["nodes"], edges=graph["edges"]) - assert graph["nodes"][1]["id"] != "bad id" - - -def test_repair_branch_edge_handles(): - nodes = [{"id": "n1", "data": {"type": "question-classifier", "classes": [{"id": "c1", "name": "c1"}]}}] - edges = [{"source": "n1", "target": "n2", "sourceHandle": ""}] - - WorkflowGenerator._repair_branch_edge_handles(nodes=nodes, edges=edges) - assert edges[0]["sourceHandle"] == "c1" # assuming it falls back to first one - - -def test_document_extractor_start_vars(): - nodes = [{"id": "n1", "data": {"type": "document-extractor", "variable_selector": ["start", "doc"]}}] - res = WorkflowGenerator._document_extractor_start_vars(nodes=nodes, start_id="start") - assert res == {"doc": False} - - -def test_missing_terminal_mode_auto(): - # Empty graph, should get MISSING_TERMINAL - graph = {"nodes": [{"id": "n1", "data": {"type": "start"}}], "edges": []} - - # Missing terminal check happens inside _validate_structure - WorkflowGenerator._validate_structure(graph=graph, mode="workflow", installed_tools=set()) - # Mocking this deeply is hard, but we can verify it doesn't fail diff --git a/api/tests/unit_tests/core/workflow/graph_engine/layers/test_observability.py b/api/tests/unit_tests/core/workflow/graph_engine/layers/test_observability.py deleted file mode 100644 index 94acff901bc..00000000000 --- a/api/tests/unit_tests/core/workflow/graph_engine/layers/test_observability.py +++ /dev/null @@ -1,449 +0,0 @@ -""" -Tests for ObservabilityLayer. - -Test coverage: -- Initialization and enable/disable logic -- Node span lifecycle (start, end, error handling) -- Parser integration (default, tool, LLM, and retrieval parsers) -- Result event parameter extraction (inputs/outputs) -- Graph lifecycle management -- Disabled mode behavior -""" - -from concurrent.futures import ThreadPoolExecutor -from threading import Event - -import pytest -from opentelemetry.sdk.trace import ReadableSpan, SpanProcessor -from opentelemetry.trace import StatusCode, get_current_span - -from core.app.workflow.layers.observability import ObservabilityLayer -from core.workflow.workflow_entry import WorkflowEntry -from extensions.otel.semconv import DifySpanAttributes -from graphon.engine_events import GraphRunAbortedEvent, GraphRunSucceededEvent -from graphon.graph import Graph -from graphon.nodes.start import StartNode -from graphon.nodes.start.entities import StartNodeData -from graphon.runtime import RuntimeState, VariablePool -from tests.workflow_test_utils import build_test_graph_init_params - - -@pytest.fixture(autouse=True) -def _otel_config(config_overrides) -> None: - config_overrides(ENABLE_OTEL=True) - - -@pytest.mark.parametrize("tracing_enabled", [False, True]) -def test_unbound_node_still_runs_and_preserves_body_errors(config_overrides, tracing_enabled): - config_overrides(ENABLE_OTEL=tracing_enabled) - node = StartNode( - node_id="start", - data=StartNodeData(title="Start", variables=[]), - init_params=build_test_graph_init_params(), - runtime_state=RuntimeState(workflow_id="workflow", variable_pool=VariablePool(), start_at=0), - ) - layer = ObservabilityLayer() - parent = get_current_span() - executed = [] - - with layer.node_run_context(node): - executed.append("completed") - - with pytest.raises(ValueError, match="node failed"), layer.node_run_context(node): - raise ValueError("node failed") - - assert executed == ["completed"] - assert get_current_span() is parent - - -@pytest.mark.usefixtures("mock_is_instrument_flag_enabled_true", "tracer_provider_with_memory_exporter") -def test_instrument_flag_enables_workflow_node_spans(memory_span_exporter, config_overrides, monkeypatch): - config_overrides(ENABLE_OTEL=False, GRAPH_ENGINE_MAX_WORKERS=1) - monkeypatch.setattr("core.workflow.workflow_entry.is_instrument_flag_enabled", lambda: True) - state = RuntimeState(workflow_id="workflow", variable_pool=VariablePool(), start_at=0) - node = StartNode( - node_id="start", - data=StartNodeData(title="Start", variables=[]), - init_params=build_test_graph_init_params(), - runtime_state=state, - ) - entry = WorkflowEntry( - graph_config={}, - graph=Graph.new().add_root(node).build(), - call_depth=0, - graph_runtime_state=state, - ) - - events = list(entry.run()) - - assert isinstance(events[-1], GraphRunSucceededEvent) - spans = memory_span_exporter.get_finished_spans() - assert [span.name for span in spans] == [node.title] - - -class TestObservabilityLayerNodeSpanLifecycle: - """Test node span creation and lifecycle management.""" - - @pytest.mark.usefixtures("mock_is_instrument_flag_enabled_false") - def test_node_span_created_and_ended( - self, tracer_provider_with_memory_exporter, memory_span_exporter, mock_llm_node - ): - """Test that span is created on node start and ended on node end.""" - layer = ObservabilityLayer() - layer.on_graph_start() - - with layer.node_run_context(mock_llm_node): - layer.on_node_run_end(mock_llm_node, None) - - spans = memory_span_exporter.get_finished_spans() - assert len(spans) == 1 - assert spans[0].name == mock_llm_node.title - assert spans[0].status.status_code == StatusCode.OK - - @pytest.mark.usefixtures("mock_is_instrument_flag_enabled_false") - def test_node_error_recorded_in_span( - self, tracer_provider_with_memory_exporter, memory_span_exporter, mock_llm_node - ): - """Test that node execution errors are recorded in span.""" - layer = ObservabilityLayer() - layer.on_graph_start() - - error = ValueError("Test error") - with layer.node_run_context(mock_llm_node): - layer.on_node_run_end(mock_llm_node, error) - - spans = memory_span_exporter.get_finished_spans() - assert len(spans) == 1 - assert spans[0].status.status_code == StatusCode.ERROR - assert len(spans[0].events) > 0 - assert any("exception" in event.name.lower() for event in spans[0].events) - - @pytest.mark.usefixtures("mock_is_instrument_flag_enabled_false") - def test_node_end_without_start_handled_gracefully( - self, tracer_provider_with_memory_exporter, memory_span_exporter, mock_llm_node - ): - """Test that ending a node without start doesn't crash.""" - layer = ObservabilityLayer() - layer.on_graph_start() - - layer.on_node_run_end(mock_llm_node, None) - - spans = memory_span_exporter.get_finished_spans() - assert len(spans) == 0 - - @pytest.mark.usefixtures("mock_is_instrument_flag_enabled_false") - def test_suspension_deactivates_span_until_resume( - self, tracer_provider_with_memory_exporter, memory_span_exporter, mock_tool_node - ): - layer = ObservabilityLayer() - parent = get_current_span() - - with layer.node_run_context(mock_tool_node): - tool_span = get_current_span() - assert tool_span.is_recording() - - assert get_current_span() is parent - assert memory_span_exporter.get_finished_spans() == () - - with layer.node_run_context(mock_tool_node): - assert get_current_span() is tool_span - layer.on_node_run_end(mock_tool_node, None) - - assert get_current_span() is parent - spans = memory_span_exporter.get_finished_spans() - assert len(spans) == 1 - assert spans[0].context == tool_span.get_span_context() - - @pytest.mark.usefixtures("mock_is_instrument_flag_enabled_false") - def test_activation_does_not_swallow_node_errors( - self, tracer_provider_with_memory_exporter, memory_span_exporter, mock_llm_node - ): - layer = ObservabilityLayer() - parent = get_current_span() - error = ValueError("node failed") - - with pytest.raises(ValueError, match="node failed"): - with layer.node_run_context(mock_llm_node): - raise error - - assert get_current_span() is parent - layer.on_node_run_end(mock_llm_node, error) - spans = memory_span_exporter.get_finished_spans() - assert len(spans) == 1 - assert spans[0].status.status_code == StatusCode.ERROR - - @pytest.mark.usefixtures("mock_is_instrument_flag_enabled_false") - def test_parser_failure_still_ends_span( - self, tracer_provider_with_memory_exporter, memory_span_exporter, mock_start_node, monkeypatch, caplog - ): - layer = ObservabilityLayer() - - def fail_parse(**_kwargs): - raise ValueError("invalid node attributes") - - monkeypatch.setattr(layer._default_parser, "parse", fail_parse) - - with layer.node_run_context(mock_start_node): - layer.on_node_run_end(mock_start_node, None) - - assert len(memory_span_exporter.get_finished_spans()) == 1 - assert "invalid node attributes" in caplog.text - - -class TestObservabilityLayerParserIntegration: - """Test parser integration for different node types.""" - - @pytest.mark.usefixtures("mock_is_instrument_flag_enabled_false") - def test_default_parser_used_for_regular_node( - self, tracer_provider_with_memory_exporter, memory_span_exporter, mock_start_node - ): - """Test that default parser is used for non-tool nodes.""" - layer = ObservabilityLayer() - layer.on_graph_start() - - with layer.node_run_context(mock_start_node): - layer.on_node_run_end(mock_start_node, None) - - spans = memory_span_exporter.get_finished_spans() - assert len(spans) == 1 - attrs = spans[0].attributes - assert attrs["node.id"] == mock_start_node.id - assert attrs["node.execution_id"] == mock_start_node.execution_id - assert attrs["node.type"] == mock_start_node.node_type - - @pytest.mark.usefixtures("mock_is_instrument_flag_enabled_false") - def test_tool_parser_used_for_tool_node( - self, tracer_provider_with_memory_exporter, memory_span_exporter, mock_tool_node - ): - """Test that tool parser is used for tool nodes.""" - layer = ObservabilityLayer() - layer.on_graph_start() - - with layer.node_run_context(mock_tool_node): - layer.on_node_run_end(mock_tool_node, None) - - spans = memory_span_exporter.get_finished_spans() - assert len(spans) == 1 - attrs = spans[0].attributes - assert attrs["node.id"] == mock_tool_node.id - assert attrs["gen_ai.tool.name"] == mock_tool_node.title - assert attrs["gen_ai.tool.type"] == mock_tool_node._node_data.provider_type.value - - @pytest.mark.usefixtures("mock_is_instrument_flag_enabled_false") - def test_llm_parser_used_for_llm_node( - self, tracer_provider_with_memory_exporter, memory_span_exporter, mock_llm_node, mock_result_event - ): - """Test that LLM parser is used for LLM nodes and extracts LLM-specific attributes.""" - from graphon.node_events import NodeRunResult - - mock_result_event.node_run_result = NodeRunResult( - inputs={}, - outputs={"text": "test completion", "finish_reason": "stop"}, - process_data={ - "model_name": "gpt-4", - "model_provider": "openai", - "usage": {"prompt_tokens": 10, "completion_tokens": 20, "total_tokens": 30}, - "prompts": [{"role": "user", "text": "test prompt"}], - }, - metadata={}, - ) - - layer = ObservabilityLayer() - layer.on_graph_start() - - with layer.node_run_context(mock_llm_node): - layer.on_node_run_end(mock_llm_node, None, mock_result_event) - - spans = memory_span_exporter.get_finished_spans() - assert len(spans) == 1 - attrs = spans[0].attributes - assert attrs["node.id"] == mock_llm_node.id - assert attrs["gen_ai.request.model"] == "gpt-4" - assert attrs["gen_ai.provider.name"] == "openai" - assert attrs["gen_ai.usage.input_tokens"] == 10 - assert attrs["gen_ai.usage.output_tokens"] == 20 - assert attrs["gen_ai.usage.total_tokens"] == 30 - assert attrs["gen_ai.completion"] == "test completion" - assert attrs["gen_ai.response.finish_reason"] == "stop" - - @pytest.mark.usefixtures("mock_is_instrument_flag_enabled_false") - def test_retrieval_parser_used_for_retrieval_node( - self, tracer_provider_with_memory_exporter, memory_span_exporter, mock_retrieval_node, mock_result_event - ): - """Test that retrieval parser is used for retrieval nodes and extracts retrieval-specific attributes.""" - from graphon.node_events import NodeRunResult - - mock_result_event.node_run_result = NodeRunResult( - inputs={"query": "test query"}, - outputs={"result": [{"content": "test content", "metadata": {"score": 0.9, "document_id": "doc1"}}]}, - process_data={}, - metadata={}, - ) - - layer = ObservabilityLayer() - layer.on_graph_start() - - with layer.node_run_context(mock_retrieval_node): - layer.on_node_run_end(mock_retrieval_node, None, mock_result_event) - - spans = memory_span_exporter.get_finished_spans() - assert len(spans) == 1 - attrs = spans[0].attributes - assert attrs["node.id"] == mock_retrieval_node.id - assert attrs["retrieval.query"] == "test query" - assert "retrieval.document" in attrs - - @pytest.mark.usefixtures("mock_is_instrument_flag_enabled_false") - def test_result_event_extracts_inputs_and_outputs( - self, tracer_provider_with_memory_exporter, memory_span_exporter, mock_start_node, mock_result_event - ): - """Test that result_event parameter allows parsers to extract inputs and outputs.""" - from graphon.node_events import NodeRunResult - - mock_result_event.node_run_result = NodeRunResult( - inputs={"input_key": "input_value"}, - outputs={"output_key": "output_value"}, - process_data={}, - metadata={}, - ) - - layer = ObservabilityLayer() - layer.on_graph_start() - - with layer.node_run_context(mock_start_node): - layer.on_node_run_end(mock_start_node, None, mock_result_event) - - spans = memory_span_exporter.get_finished_spans() - assert len(spans) == 1 - attrs = spans[0].attributes - assert "input.value" in attrs - assert "output.value" in attrs - - -class TestObservabilityLayerGraphLifecycle: - """Test graph lifecycle management.""" - - @pytest.mark.usefixtures("mock_is_instrument_flag_enabled_false") - def test_on_graph_end_with_no_unfinished_spans( - self, tracer_provider_with_memory_exporter, memory_span_exporter, mock_llm_node - ): - """Test that on_graph_end handles normal completion.""" - layer = ObservabilityLayer() - layer.on_graph_start() - - with layer.node_run_context(mock_llm_node): - layer.on_node_run_end(mock_llm_node, None) - layer.on_graph_end(None) - - spans = memory_span_exporter.get_finished_spans() - assert len(spans) == 1 - - @pytest.mark.usefixtures("mock_is_instrument_flag_enabled_false") - def test_on_graph_end_ends_suspended_spans( - self, tracer_provider_with_memory_exporter, memory_span_exporter, mock_llm_node - ): - """Pause and abort cleanup must export spans that never reached node end.""" - layer = ObservabilityLayer() - layer.on_graph_start() - - with layer.node_run_context(mock_llm_node): - span = get_current_span() - - layer.on_graph_end(None) - layer.on_graph_end(None) - - assert not span.is_recording() - spans = memory_span_exporter.get_finished_spans() - assert len(spans) == 1 - assert spans[0].name == mock_llm_node.title - - @pytest.mark.usefixtures("mock_is_instrument_flag_enabled_false") - def test_graph_end_and_late_worker_each_end_spans_once( - self, tracer_provider_with_memory_exporter, memory_span_exporter, mock_start_node, mock_llm_node, caplog - ): - """An aborted graph can finish while a worker still owns another span.""" - graph_ending = Event() - worker_finished = Event() - - class EndBarrier(SpanProcessor): - def on_end(self, span: ReadableSpan) -> None: - if span.name == mock_start_node.title: - graph_ending.set() - assert worker_finished.wait(5) - - tracer_provider_with_memory_exporter.add_span_processor(EndBarrier()) - layer = ObservabilityLayer() - for node in (mock_start_node, mock_llm_node): - with layer.node_run_context(node): - pass - - def finish_worker() -> None: - assert graph_ending.wait(5) - try: - layer.on_node_run_end(mock_llm_node, None) - finally: - worker_finished.set() - - with ThreadPoolExecutor(max_workers=1) as executor: - worker = executor.submit(finish_worker) - try: - layer.on_graph_end(None) - finally: - worker.result(timeout=5) - - layer.on_graph_end(None) - spans = memory_span_exporter.get_finished_spans() - assert len(spans) == 2 - assert {span.name for span in spans} == {mock_start_node.title, mock_llm_node.title} - assert "Calling end() on an ended span" not in caplog.text - - @pytest.mark.usefixtures("mock_is_instrument_flag_enabled_false") - def test_graph_aborted_event_records_reason_on_current_span( - self, tracer_provider_with_memory_exporter, memory_span_exporter, mock_start_node - ): - layer = ObservabilityLayer() - layer.on_graph_start() - with layer.node_run_context(mock_start_node): - layer.on_event(GraphRunAbortedEvent(reason="worker shutdown", outputs={})) - layer.on_node_run_end(mock_start_node, None) - - spans = memory_span_exporter.get_finished_spans() - assert len(spans) == 1 - assert spans[0].attributes[DifySpanAttributes.WORKFLOW_ABORT_REASON] == "worker shutdown" - assert any( - event.name == "dify.workflow.aborted" - and event.attributes[DifySpanAttributes.WORKFLOW_ABORT_REASON] == "worker shutdown" - for event in spans[0].events - ) - - -class TestObservabilityLayerDisabledMode: - """Test behavior when layer is disabled.""" - - @pytest.mark.usefixtures("mock_is_instrument_flag_enabled_false") - def test_disabled_mode_skips_node_start( - self, tracer_provider_with_memory_exporter, memory_span_exporter, mock_start_node, config_overrides - ): - """Test that disabled layer doesn't create spans on node start.""" - config_overrides(ENABLE_OTEL=False) - layer = ObservabilityLayer() - - layer.on_graph_start() - with layer.node_run_context(mock_start_node): - layer.on_node_run_end(mock_start_node, None) - - spans = memory_span_exporter.get_finished_spans() - assert len(spans) == 0 - - @pytest.mark.usefixtures("mock_is_instrument_flag_enabled_false") - def test_disabled_mode_skips_node_end( - self, tracer_provider_with_memory_exporter, memory_span_exporter, mock_llm_node, config_overrides - ): - """Test that disabled layer doesn't process node end.""" - config_overrides(ENABLE_OTEL=False) - layer = ObservabilityLayer() - - layer.on_node_run_end(mock_llm_node, None) - - spans = memory_span_exporter.get_finished_spans() - assert len(spans) == 0 diff --git a/api/tests/unit_tests/core/workflow/graph_engine/layers/test_workflow_tool_observability.py b/api/tests/unit_tests/core/workflow/graph_engine/layers/test_workflow_tool_observability.py deleted file mode 100644 index 563b40824d0..00000000000 --- a/api/tests/unit_tests/core/workflow/graph_engine/layers/test_workflow_tool_observability.py +++ /dev/null @@ -1,215 +0,0 @@ -from collections.abc import Generator -from contextlib import AbstractContextManager -from contextvars import ContextVar -from functools import partial -from threading import Event, get_ident -from types import TracebackType -from typing import override - -import pytest -from flask import current_app, g -from opentelemetry.sdk.trace import TracerProvider -from opentelemetry.sdk.trace.export.in_memory_span_exporter import InMemorySpanExporter -from opentelemetry.trace import get_current_span - -from context import capture_current_context -from context.flask_app_context import capture_flask_context -from core.app.layers import execution_context_layer as module -from core.app.workflow.file_runtime import init_app -from core.app.workflow.layers.observability import ObservabilityLayer -from core.workflow.workflow_entry import WorkflowEntry -from core.workflow.workflow_tool_container_handler import WorkflowToolContainerHandler -from dify_app import DifyApp -from graphon.engine import Engine -from graphon.engine.layer import Layer -from graphon.engine_events import GraphRunSucceededEvent, NodeEvent, NodeRunStreamChunkEvent -from graphon.entities.base_node_data import BaseNodeData -from graphon.enums import BuiltinNodeTypes -from graphon.node_events.base import NodeEventPayload -from graphon.nodes.base.node import Node -from graphon.nodes.container_effects import ContainerAwaitRequest, ContainerRunResult -from graphon.runtime import RuntimeState, VariablePool -from tests.unit_tests.config_override import apply_config_overrides -from tests.unit_tests.core.workflow.test_workflow_tool_container import ( - _outer_graph, - _workflow_tool_node, - _workflow_tool_source, -) - - -def test_workflow_tool_tracing_survives_worker_handoff( - monkeypatch: pytest.MonkeyPatch, - caplog: pytest.LogCaptureFixture, - tracer_provider_with_memory_exporter: TracerProvider, - memory_span_exporter: InMemorySpanExporter, -) -> None: - apply_config_overrides(monkeypatch, ENABLE_OTEL=True) - tool_started_on: list[int] = [] - tool_finished_on: list[int] = [] - completed = Event() - - class HandoffLayer(Layer): - @override - def node_run_context( - self, node: Node[BaseNodeData], *, parent_execution_id: str | None = None - ) -> AbstractContextManager[None]: - return HandoffContext() - - @override - def on_node_run_start(self, node: Node[BaseNodeData]) -> None: - if node.id == "tool": - tool_started_on.append(get_ident()) - - @override - def on_node_run_end( - self, node: Node[BaseNodeData], error: Exception | None, result_event: NodeEvent | None = None - ) -> None: - if node.id == "tool": - tool_finished_on.append(get_ident()) - completed.set() - - tracer = tracer_provider_with_memory_exporter.get_tracer(__name__) - with tracer.start_as_current_span("workflow") as root_span: - captured_context = capture_current_context() - - class HandoffContext(AbstractContextManager[None]): - @override - def __enter__(self) -> None: - captured_context.__enter__() - - @override - def __exit__( - self, exc_type: type[BaseException] | None, exc: BaseException | None, traceback: TracebackType | None - ) -> None: - captured_context.__exit__(exc_type, exc, traceback) - if tool_started_on == [get_ident()] and not completed.is_set(): - # Keep the starting worker occupied after suspension so another - # worker must execute the child graph and resume the Tool. - assert completed.wait(10), "Workflow Tool did not resume on the available worker" - - state = RuntimeState(workflow_id="outer-workflow", variable_pool=VariablePool(), start_at=1) - tool, _, _ = _workflow_tool_node(state) - original_resume = tool._resume_container_events - - def instrument_resume(*, result: ContainerRunResult) -> Generator[NodeEventPayload | ContainerAwaitRequest]: - with tracer.start_as_current_span("resume-operation"): - yield from original_resume(result=result) - - monkeypatch.setattr(tool, "_resume_container_events", instrument_resume) - source = _workflow_tool_source() - sources = {source.workflow_id: source} - engine = Engine( - graph=_outer_graph(tool), - runtime_state=state, - workers=2, - container_handler_factories=(partial(WorkflowToolContainerHandler, sources=sources),), - ) - engine.add_layer(HandoffLayer()) - engine.add_layer(ObservabilityLayer()) - - events = list(engine.run()) - - assert isinstance(events[-1], GraphRunSucceededEvent) - assert len(tool_started_on) == len(tool_finished_on) == 1 - assert tool_started_on != tool_finished_on - spans = memory_span_exporter.get_finished_spans() - tool_spans = [span for span in spans if (span.attributes or {}).get("node.id") == "tool"] - assert len(tool_spans) == 1 - tool_span = tool_spans[0] - assert tool_span.parent == root_span.get_span_context() - resumed_operations = [span for span in spans if span.name == "resume-operation"] - assert len(resumed_operations) == 1 - assert resumed_operations[0].parent == tool_span.context - assert "was created in a different Context" not in caplog.text - children = [span for span in spans if (span.attributes or {}).get("node.id") in {"source-start", "source-end"}] - assert len(children) == 2 - assert all(span.parent == tool_span.context for span in children) - - -def test_standalone_node_uses_worker_scoped_tracing_context( - monkeypatch: pytest.MonkeyPatch, - tracer_provider_with_memory_exporter: TracerProvider, - memory_span_exporter: InMemorySpanExporter, -) -> None: - apply_config_overrides(monkeypatch, ENABLE_OTEL=True) - tool, _, _ = _workflow_tool_node() - node = _outer_graph(tool).root_node - tracer = tracer_provider_with_memory_exporter.get_tracer(__name__) - with tracer.start_as_current_span("standalone") as parent: - events = list(WorkflowEntry._run_node_with_layers(node)) - assert get_current_span() is parent - - assert isinstance(events[-1], NodeEvent) - assert events[-1].node_run_result.status == "succeeded" - node_spans = [span for span in memory_span_exporter.get_finished_spans() if span.name == node.title] - assert len(node_spans) == 1 - assert node_spans[0].parent == parent.get_span_context() - - -def test_interleaved_snippets_restore_caller_flask_context( - monkeypatch: pytest.MonkeyPatch, - tracer_provider_with_memory_exporter: TracerProvider, - memory_span_exporter: InMemorySpanExporter, -) -> None: - apply_config_overrides(monkeypatch, ENABLE_OTEL=True) - monkeypatch.setattr(module, "capture_current_context", capture_flask_context) - request_id: ContextVar[str] = ContextVar("snippet_request_id") - tracer = tracer_provider_with_memory_exporter.get_tracer(__name__) - parents = [] - snippets = [] - observed = [] - for name in ("first", "second"): - app = DifyApp(name) - init_app(app) - with app.app_context(), tracer.start_as_current_span(name) as parent: - parents.append(parent.get_span_context()) - g._login_user = name - request_id.set(name) - tool, _, _ = _workflow_tool_node() - node = _outer_graph(tool).root_node - - def run_node() -> Generator[NodeRunStreamChunkEvent]: - observed.append((current_app.name, g._login_user, request_id.get())) - yield NodeRunStreamChunkEvent( - id="snippet-execution", - node_id="outer-start", - node_type=BuiltinNodeTypes.START, - selector=["outer-start", "text"], - chunk="part", - is_final=False, - ) - observed.append((current_app.name, g._login_user, request_id.get())) - yield NodeRunStreamChunkEvent( - id="snippet-execution", - node_id="outer-start", - node_type=BuiltinNodeTypes.START, - selector=["outer-start", "text"], - chunk="part", - is_final=False, - ) - - monkeypatch.setattr(node, "run", run_node) - snippets.append(WorkflowEntry._run_node_with_layers(node)) - - caller = DifyApp("caller") - init_app(caller) - with caller.app_context(), tracer.start_as_current_span("caller") as caller_span: - g._login_user = "caller" - request_id.set("caller") - for _ in range(2): - for snippet in snippets: - next(snippet) - assert current_app == caller - assert g._login_user == request_id.get() == "caller" - assert get_current_span() is caller_span - for snippet in snippets: - snippet.close() - assert current_app == caller - assert g._login_user == request_id.get() == "caller" - assert get_current_span() is caller_span - - assert observed == [(name, name, name) for name in ("first", "second", "first", "second")] - spans = memory_span_exporter.get_finished_spans() - node_spans = [span for span in spans if (span.attributes or {}).get("node.id") == "outer-start"] - assert len(node_spans) == 2 - assert [span.parent for span in node_spans] == parents diff --git a/api/tests/unit_tests/core/workflow/nodes/agent/test_agent_entities.py b/api/tests/unit_tests/core/workflow/nodes/agent/test_entities.py similarity index 100% rename from api/tests/unit_tests/core/workflow/nodes/agent/test_agent_entities.py rename to api/tests/unit_tests/core/workflow/nodes/agent/test_entities.py diff --git a/api/tests/unit_tests/core/workflow/nodes/agent/test_runtime_support.py b/api/tests/unit_tests/core/workflow/nodes/agent/test_runtime_support.py index 25058089888..da54f8ee7e0 100644 --- a/api/tests/unit_tests/core/workflow/nodes/agent/test_runtime_support.py +++ b/api/tests/unit_tests/core/workflow/nodes/agent/test_runtime_support.py @@ -14,7 +14,7 @@ from graphon.model_runtime.entities.model_entities import ( ParameterRule, ParameterType, ) -from tests.unit_tests.core.model_fixtures import make_model_config +from tests.helpers.unit.core.model_fixtures import make_model_config def test_fetch_model_reuses_single_model_assembly(): diff --git a/api/tests/unit_tests/core/workflow/nodes/agent_v2/test_runtime_request_builder.py b/api/tests/unit_tests/core/workflow/nodes/agent_v2/test_runtime_request_builder.py index 4a0adc8504d..e8c2efee1d0 100644 --- a/api/tests/unit_tests/core/workflow/nodes/agent_v2/test_runtime_request_builder.py +++ b/api/tests/unit_tests/core/workflow/nodes/agent_v2/test_runtime_request_builder.py @@ -51,8 +51,8 @@ from models.agent_config_entities import ( WorkflowNodeJobConfig, ) from services.tracing.runtime_bridge import RuntimeApplicationTracingHooks -from tests.unit_tests.config_override import apply_config_overrides -from tests.unit_tests.knowledge_fs_core_fakes import RecordingModelFactory, llm_schema +from tests.helpers.unit.config_override import apply_config_overrides +from tests.helpers.unit.knowledge_fs_core_fakes import RecordingModelFactory, llm_schema @pytest.fixture(autouse=True) diff --git a/api/tests/unit_tests/core/workflow/nodes/agent_v2/test_validators.py b/api/tests/unit_tests/core/workflow/nodes/agent_v2/test_validators.py index 328832171e8..d3526c85ee1 100644 --- a/api/tests/unit_tests/core/workflow/nodes/agent_v2/test_validators.py +++ b/api/tests/unit_tests/core/workflow/nodes/agent_v2/test_validators.py @@ -22,7 +22,7 @@ from models.agent_config_entities import AgentSoulConfig, AgentSoulModelConfig, from models.enums import CreatorUserRole from models.model import UploadFile from models.workflow import Workflow -from tests.unit_tests.model_factories import make_workflow +from tests.helpers.unit.model_factories import make_workflow def _workflow(graph: dict) -> Workflow: diff --git a/api/tests/unit_tests/core/workflow/nodes/base/test_get_node_type_classes_mapping.py b/api/tests/unit_tests/core/workflow/nodes/base/test_get_node_type_classes_mapping.py deleted file mode 100644 index ef0df55995e..00000000000 --- a/api/tests/unit_tests/core/workflow/nodes/base/test_get_node_type_classes_mapping.py +++ /dev/null @@ -1,85 +0,0 @@ -import types -from collections.abc import Mapping - -from core.workflow.node_factory import get_node_type_classes_mapping -from graphon.entities.base_node_data import BaseNodeData -from graphon.enums import BuiltinNodeTypes, NodeType -from graphon.nodes.base.node import Node - -# Import concrete nodes we will assert on (numeric version path) -from graphon.nodes.variable_assigner.v1.node import ( - VariableAssignerNode as VariableAssignerV1, -) -from graphon.nodes.variable_assigner.v2.node import ( - VariableAssignerNode as VariableAssignerV2, -) - - -def test_variable_assigner_latest_prefers_highest_numeric_version(): - # Act - mapping: Mapping[NodeType, Mapping[str, type[Node]]] = get_node_type_classes_mapping() - - # Assert basic presence - assert BuiltinNodeTypes.VARIABLE_ASSIGNER in mapping - va_versions = mapping[BuiltinNodeTypes.VARIABLE_ASSIGNER] - - # Both concrete versions must be present - assert va_versions.get("1") is VariableAssignerV1 - assert va_versions.get("2") is VariableAssignerV2 - - # And latest should point to numerically-highest version ("2") - assert va_versions.get("latest") is VariableAssignerV2 - - -def test_latest_prefers_highest_numeric_version(): - # Arrange: define two ephemeral subclasses with numeric versions under a NodeType - # that has no concrete implementations in production to avoid interference. - class _Version1(Node[BaseNodeData]): # type: ignore[misc] - node_type = BuiltinNodeTypes.LEGACY_VARIABLE_AGGREGATOR - - def init_node_data(self, data): - pass - - def _run(self): - raise NotImplementedError - - @classmethod - def version(cls) -> str: - return "1" - - def _get_error_strategy(self): - return None - - def _get_retry_config(self): - return types.SimpleNamespace() # not used - - def _get_title(self) -> str: - return "version1" - - def _get_description(self): - return None - - def _get_default_value_dict(self): - return {} - - def get_base_node_data(self): - return types.SimpleNamespace(title="version1") - - class _Version2(_Version1): # type: ignore[misc] - @classmethod - def version(cls) -> str: - return "2" - - def _get_title(self) -> str: - return "version2" - - # Act: build a fresh mapping (it should now see our ephemeral subclasses) - mapping: Mapping[NodeType, Mapping[str, type[Node]]] = get_node_type_classes_mapping() - - # Assert: both numeric versions exist for this NodeType; 'latest' points to the higher numeric version - assert BuiltinNodeTypes.LEGACY_VARIABLE_AGGREGATOR in mapping - legacy_versions = mapping[BuiltinNodeTypes.LEGACY_VARIABLE_AGGREGATOR] - - assert legacy_versions.get("1") is _Version1 - assert legacy_versions.get("2") is _Version2 - assert legacy_versions.get("latest") is _Version2 diff --git a/api/tests/unit_tests/core/workflow/nodes/datasource/test_datasource_node.py b/api/tests/unit_tests/core/workflow/nodes/datasource/test_datasource_node.py index 595b81ab397..b2c3482de8d 100644 --- a/api/tests/unit_tests/core/workflow/nodes/datasource/test_datasource_node.py +++ b/api/tests/unit_tests/core/workflow/nodes/datasource/test_datasource_node.py @@ -1,10 +1,22 @@ -from pytest_mock import MockerFixture +import time +from collections.abc import Generator +from types import SimpleNamespace +from typing import Never -from core.app.entities.app_invoke_entities import DIFY_RUN_CONTEXT_KEY +import pytest +from pytest_mock import MockerFixture +from sqlalchemy import Engine +from sqlalchemy.orm import sessionmaker + +from core.app.entities.app_invoke_entities import DIFY_RUN_CONTEXT_KEY, InvokeFrom, UserFrom +from core.workflow.node_factory import DifyNodeFactory from core.workflow.nodes.datasource.datasource_node import DatasourceNode from core.workflow.nodes.datasource.entities import DatasourceNodeData +from core.workflow.system_variables import build_system_variables from graphon.enums import WorkflowNodeExecutionStatus from graphon.node_events import NodeRunResult, StreamChunkEvent, StreamCompletedEvent +from graphon.runtime import RuntimeState, VariablePool +from tests.workflow_test_utils import build_test_graph_init_params class _VarSeg: @@ -98,3 +110,85 @@ def test_datasource_node_delegates_to_manager_stream(mocker: MockerFixture): evts = list(node._run()) assert isinstance(evts[0], StreamChunkEvent) assert isinstance(evts[-1], StreamCompletedEvent) + + +"""Unit coverage for datasource node event streaming.""" + + +@pytest.mark.parametrize("use_factory", [False, True]) +def test_node_integration_minimal_stream(mocker: MockerFixture, sqlite_engine: Engine, use_factory: bool) -> None: + variable_pool = VariablePool.from_bootstrap( + system_variables=build_system_variables( + datasource_type="online_document", + datasource_info={ + "workspace_id": "w", + "page": {"page_id": "pg", "type": "t"}, + "credential_id": "credential-1", + }, + ) + ) + graph_init_params = build_test_graph_init_params( + workflow_id="wf-1", + tenant_id="t1", + app_id="app-1", + user_id="u1", + user_from=UserFrom.ACCOUNT, + invoke_from=InvokeFrom.DEBUGGER, + ) + graph_runtime_state = RuntimeState( + variable_pool=variable_pool, + start_at=time.perf_counter(), + workflow_id=graph_init_params.workflow_id, + ) + + class _Mgr: + @classmethod + def get_icon_url(cls, **_: object) -> str: + return "icon" + + @classmethod + def stream_node_events(cls, **kwargs: object) -> Generator[StreamCompletedEvent]: + assert kwargs["tenant_id"] == "t1" + assert kwargs["credentials"] == {"integration_secret": "token"} + yield from () + yield StreamCompletedEvent(node_run_result=NodeRunResult(status=WorkflowNodeExecutionStatus.SUCCEEDED)) + + @classmethod + def get_upload_file_by_id(cls, **_: object) -> Never: + raise AssertionError + + mocker.patch("core.workflow.nodes.datasource.datasource_node.DatasourceManager", new=_Mgr) + + load_credentials = mocker.Mock(return_value={"integration_secret": "token"}) + node_data = DatasourceNodeData( + type="datasource", + version="1", + title="Datasource", + provider_type="plugin", + provider_name="p", + plugin_id="plug", + datasource_name="ds", + ) + if use_factory: + database_client = sessionmaker(sqlite_engine) + mocker.patch("core.workflow.node_factory.get_session_maker", return_value=database_client) + build_credentials = mocker.patch( + "core.workflow.node_factory.build_data_source_credentials", + return_value=SimpleNamespace(providers=SimpleNamespace(get_datasource_credentials=load_credentials)), + ) + factory = DifyNodeFactory(init_params=graph_init_params, runtime_state=graph_runtime_state) + node = factory.create_node({"id": "n", "data": node_data.model_dump()}) + build_credentials.assert_called_once_with(database_client=database_client) + assert isinstance(node, DatasourceNode) + else: + node = DatasourceNode( + node_id="n", + data=node_data, + init_params=graph_init_params, + runtime_state=graph_runtime_state, + datasource_credentials=load_credentials, + ) + + out = list(node._run()) + assert isinstance(out[-1], StreamCompletedEvent) + load_credentials.assert_called_once_with("t1", "p", "plug", "credential-1") diff --git a/api/tests/unit_tests/core/workflow/nodes/datasource/test_datasource_node_stream.py b/api/tests/unit_tests/core/workflow/nodes/datasource/test_datasource_node_stream.py deleted file mode 100644 index 18ab7ae577f..00000000000 --- a/api/tests/unit_tests/core/workflow/nodes/datasource/test_datasource_node_stream.py +++ /dev/null @@ -1,100 +0,0 @@ -"""Unit coverage for datasource node event streaming.""" - -import time -from collections.abc import Generator -from types import SimpleNamespace -from typing import Never - -import pytest -from pytest_mock import MockerFixture -from sqlalchemy import Engine -from sqlalchemy.orm import sessionmaker - -from core.app.entities.app_invoke_entities import InvokeFrom, UserFrom -from core.workflow.node_factory import DifyNodeFactory -from core.workflow.nodes.datasource.datasource_node import DatasourceNode -from core.workflow.nodes.datasource.entities import DatasourceNodeData -from core.workflow.system_variables import build_system_variables -from graphon.enums import WorkflowNodeExecutionStatus -from graphon.node_events import NodeRunResult, StreamCompletedEvent -from graphon.runtime import RuntimeState, VariablePool -from tests.workflow_test_utils import build_test_graph_init_params - - -@pytest.mark.parametrize("use_factory", [False, True]) -def test_node_integration_minimal_stream(mocker: MockerFixture, sqlite_engine: Engine, use_factory: bool) -> None: - variable_pool = VariablePool.from_bootstrap( - system_variables=build_system_variables( - datasource_type="online_document", - datasource_info={ - "workspace_id": "w", - "page": {"page_id": "pg", "type": "t"}, - "credential_id": "credential-1", - }, - ) - ) - graph_init_params = build_test_graph_init_params( - workflow_id="wf-1", - tenant_id="t1", - app_id="app-1", - user_id="u1", - user_from=UserFrom.ACCOUNT, - invoke_from=InvokeFrom.DEBUGGER, - ) - graph_runtime_state = RuntimeState( - variable_pool=variable_pool, - start_at=time.perf_counter(), - workflow_id=graph_init_params.workflow_id, - ) - - class _Mgr: - @classmethod - def get_icon_url(cls, **_: object) -> str: - return "icon" - - @classmethod - def stream_node_events(cls, **kwargs: object) -> Generator[StreamCompletedEvent]: - assert kwargs["tenant_id"] == "t1" - assert kwargs["credentials"] == {"integration_secret": "token"} - yield from () - yield StreamCompletedEvent(node_run_result=NodeRunResult(status=WorkflowNodeExecutionStatus.SUCCEEDED)) - - @classmethod - def get_upload_file_by_id(cls, **_: object) -> Never: - raise AssertionError - - mocker.patch("core.workflow.nodes.datasource.datasource_node.DatasourceManager", new=_Mgr) - - load_credentials = mocker.Mock(return_value={"integration_secret": "token"}) - node_data = DatasourceNodeData( - type="datasource", - version="1", - title="Datasource", - provider_type="plugin", - provider_name="p", - plugin_id="plug", - datasource_name="ds", - ) - if use_factory: - database_client = sessionmaker(sqlite_engine) - mocker.patch("core.workflow.node_factory.get_session_maker", return_value=database_client) - build_credentials = mocker.patch( - "core.workflow.node_factory.build_data_source_credentials", - return_value=SimpleNamespace(providers=SimpleNamespace(get_datasource_credentials=load_credentials)), - ) - factory = DifyNodeFactory(init_params=graph_init_params, runtime_state=graph_runtime_state) - node = factory.create_node({"id": "n", "data": node_data.model_dump()}) - build_credentials.assert_called_once_with(database_client=database_client) - assert isinstance(node, DatasourceNode) - else: - node = DatasourceNode( - node_id="n", - data=node_data, - init_params=graph_init_params, - runtime_state=graph_runtime_state, - datasource_credentials=load_credentials, - ) - - out = list(node._run()) - assert isinstance(out[-1], StreamCompletedEvent) - load_credentials.assert_called_once_with("t1", "p", "plug", "credential-1") diff --git a/api/tests/unit_tests/core/workflow/test_human_input_callback.py b/api/tests/unit_tests/core/workflow/nodes/human_input/test_callback.py similarity index 100% rename from api/tests/unit_tests/core/workflow/test_human_input_callback.py rename to api/tests/unit_tests/core/workflow/nodes/human_input/test_callback.py diff --git a/api/tests/unit_tests/core/workflow/nodes/human_input/test_dify_owned_contracts.py b/api/tests/unit_tests/core/workflow/nodes/human_input/test_dify_owned_contracts.py deleted file mode 100644 index f04921573ae..00000000000 --- a/api/tests/unit_tests/core/workflow/nodes/human_input/test_dify_owned_contracts.py +++ /dev/null @@ -1,139 +0,0 @@ -from __future__ import annotations - -import json -from datetime import UTC, datetime - -import pytest -from pydantic import TypeAdapter - -from core.workflow.nodes.human_input._exc import ExtensionsNotSetErrorValueError -from core.workflow.nodes.human_input.entities import ( - FileInputConfig, - FormDefinition, - ParagraphInputConfig, - SelectInputConfig, - UserActionConfig, -) -from core.workflow.nodes.human_input.enums import ButtonStyle, FormInputType, ValueSourceType -from core.workflow.nodes.human_input.pause_reason import HumanInputRequired -from core.workflow.nodes.human_input.session_binding import SessionBinding -from graphon.file import FileType - - -def test_session_binding_identity_mapping() -> None: - binding = SessionBinding() - - assert binding.issue_session_id_for_form(form_id="form-1") == "form-1" - assert binding.resolve_form_id_from_session_id(session_id="form-1") == "form-1" - - -def test_human_input_node_contracts_accept_legacy_json_payload() -> None: - payload = { - "form_content": "Please confirm", - "inputs": [ - { - "type": "paragraph", - "output_variable_name": "name", - "default": { - "type": "constant", - "selector": [], - "value": "Alice", - }, - }, - { - "type": "select", - "output_variable_name": "decision", - "option_source": { - "type": "constant", - "selector": [], - "value": ["approve", "reject"], - }, - }, - ], - "user_actions": [ - { - "id": "approve", - "title": "Approve", - "button_style": "primary", - } - ], - "rendered_content": "Please confirm", - "expiration_time": "2024-01-01T00:00:00Z", - "default_values": {"name": "Alice"}, - "node_title": "Human Input", - "display_in_ui": True, - } - - restored = TypeAdapter(FormDefinition).validate_json(json.dumps(payload)) - - assert isinstance(restored.inputs[0], ParagraphInputConfig) - assert isinstance(restored.inputs[1], SelectInputConfig) - assert restored.inputs[0].default is not None - assert restored.inputs[0].default.type == ValueSourceType.CONSTANT - assert restored.inputs[0].default.value == "Alice" - assert restored.inputs[1].option_source.type == ValueSourceType.CONSTANT - assert restored.inputs[1].option_source.value == ["approve", "reject"] - assert restored.user_actions == [UserActionConfig(id="approve", title="Approve", button_style=ButtonStyle.PRIMARY)] - - -def test_human_input_required_pause_reason_keeps_legacy_payload_shape() -> None: - payload = { - "TYPE": "human_input_required", - "form_id": "form-1", - "form_content": "Please confirm", - "inputs": [ - { - "type": "paragraph", - "output_variable_name": "name", - "default": { - "type": "constant", - "selector": [], - "value": "Alice", - }, - } - ], - "actions": [ - { - "id": "approve", - "title": "Approve", - "button_style": "primary", - } - ], - "node_id": "node-1", - "node_title": "Human Input", - "resolved_default_values": {"name": "Alice"}, - } - - restored = TypeAdapter(HumanInputRequired).validate_json(json.dumps(payload)) - - assert restored.TYPE.value == "human_input_required" - assert restored.actions[0].button_style == ButtonStyle.PRIMARY - assert restored.model_dump(mode="json")["TYPE"] == "human_input_required" - - -def test_form_definition_dump_keeps_public_json_shape() -> None: - definition = FormDefinition( - form_content="Please confirm", - inputs=[ParagraphInputConfig(type=FormInputType.PARAGRAPH, output_variable_name="name")], - user_actions=[UserActionConfig(id="approve", title="Approve")], - rendered_content="rendered", - expiration_time=datetime(2024, 1, 1, tzinfo=UTC), - default_values={"name": "Alice"}, - node_title="Ask Name", - display_in_ui=True, - ) - - payload = definition.model_dump(mode="json") - - assert payload["expiration_time"] == "2024-01-01T00:00:00Z" - assert payload["user_actions"][0]["id"] == "approve" - assert payload["inputs"][0]["type"] == "paragraph" - - -def test_custom_file_input_requires_extensions() -> None: - with pytest.raises(ExtensionsNotSetErrorValueError): - FileInputConfig( - output_variable_name="attachment", - allowed_file_types=[FileType.CUSTOM], - allowed_file_extensions=[], - ) diff --git a/api/tests/unit_tests/core/workflow/nodes/human_input/test_email_delivery_config.py b/api/tests/unit_tests/core/workflow/nodes/human_input/test_email_delivery_config.py deleted file mode 100644 index 07430498e55..00000000000 --- a/api/tests/unit_tests/core/workflow/nodes/human_input/test_email_delivery_config.py +++ /dev/null @@ -1,77 +0,0 @@ -from core.workflow.human_input_adapter import EmailDeliveryConfig, EmailRecipients -from graphon.runtime import VariablePool - - -def test_render_body_template_replaces_variable_values(): - config = EmailDeliveryConfig( - recipients=EmailRecipients(), - subject="Subject", - body="Hello {{#node1.value#}} {{#url#}}", - ) - variable_pool = VariablePool() - variable_pool.add(["node1", "value"], "World") - - result = config.render_body_template(body=config.body, url="https://example.com", variable_pool=variable_pool) - - assert result == "Hello World https://example.com" - - -def test_render_markdown_body_renders_markdown_to_html(): - rendered = EmailDeliveryConfig.render_markdown_body("**Bold** and [link](https://example.com)") - - assert "Bold" in rendered - assert 'link' in rendered - - -def test_render_markdown_body_sanitizes_unsafe_html(): - rendered = EmailDeliveryConfig.render_markdown_body( - 'Click' - ) - - assert "bad" in rendered - assert 'ok' in rendered - - -def test_render_markdown_body_does_not_allow_raw_html_tags(): - rendered = EmailDeliveryConfig.render_markdown_body("raw html and **markdown**") - - assert "" not in rendered - assert "raw html" in rendered - assert "markdown" in rendered - - -def test_render_markdown_body_supports_table_syntax(): - rendered = EmailDeliveryConfig.render_markdown_body("| h1 | h2 |\n| --- | ---: |\n| v1 | v2 |") - - assert "" in rendered - assert "" in rendered - assert "" in rendered - assert 'align="right"' in rendered - assert "style=" not in rendered - - -def test_sanitize_subject_removes_crlf(): - sanitized = EmailDeliveryConfig.sanitize_subject("Notice\r\nBCC:attacker@example.com") - - assert "\r" not in sanitized - assert "\n" not in sanitized - assert sanitized == "Notice BCC:attacker@example.com" - - -def test_sanitize_subject_removes_html_tags(): - sanitized = EmailDeliveryConfig.sanitize_subject("Alert") - - assert "<" not in sanitized - assert ">" not in sanitized - assert sanitized == "Alert" diff --git a/api/tests/unit_tests/core/workflow/nodes/human_input/test_entities.py b/api/tests/unit_tests/core/workflow/nodes/human_input/test_entities.py index d492fd8f270..709aaaf1979 100644 --- a/api/tests/unit_tests/core/workflow/nodes/human_input/test_entities.py +++ b/api/tests/unit_tests/core/workflow/nodes/human_input/test_entities.py @@ -1,24 +1,54 @@ -""" -Unit tests for human input node entities. -""" +from __future__ import annotations -from collections.abc import Mapping +import datetime as _human_input_form_filled_event_datetime +import json +from collections.abc import Generator, Iterable, Mapping, Sequence from dataclasses import dataclass, field -from datetime import datetime, timedelta +from datetime import UTC, timedelta +from datetime import datetime as _dify_owned_contracts_datetime +from datetime import datetime as _entities_datetime +from decimal import Decimal +from types import SimpleNamespace from typing import Any from unittest.mock import MagicMock import pytest -from pydantic import ValidationError +from flask import Flask +from pydantic import TypeAdapter, ValidationError from pytest_mock import MockerFixture +from sqlalchemy import select from sqlalchemy.orm import Session, sessionmaker -from core.app.entities.app_invoke_entities import DIFY_RUN_CONTEXT_KEY +from core.app.app_config.entities import WorkflowUIBasedAppConfig +from core.app.apps.advanced_chat.generate_response_converter import AdvancedChatAppGenerateResponseConverter +from core.app.apps.advanced_chat.generate_task_pipeline import ( + AdvancedChatAppGenerateTaskPipeline, + ConversationSnapshot, + MessageSnapshot, + WorkflowSnapshot, +) +from core.app.apps.base_app_generator import BaseAppGenerator +from core.app.apps.base_app_queue_manager import AppQueueManager, PublishFrom +from core.app.apps.workflow_app_runner import WorkflowBasedAppRunner +from core.app.entities.app_invoke_entities import ( + DIFY_RUN_CONTEXT_KEY, + AdvancedChatAppGenerateEntity, + InvokeFrom, + UserFrom, +) +from core.app.entities.queue_entities import ( + AppQueueEvent, + QueueHumanInputFormFilledEvent, + QueueNodeSucceededEvent, + QueueWorkflowStartedEvent, + WorkflowQueueMessage, +) from core.repositories.human_input_repository import ( FormCreateParams, HumanInputFormEntity, HumanInputFormRecipientEntity, HumanInputFormRepository, + HumanInputFormSubmissionRepository, ) from core.workflow.human_input_adapter import ( DeliveryMethodType, @@ -32,12 +62,15 @@ from core.workflow.human_input_adapter import ( _WebAppDeliveryConfig, ) from core.workflow.node_runtime import DifyHumanInputNodeRuntime +from core.workflow.nodes.human_input._exc import ExtensionsNotSetErrorValueError +from core.workflow.nodes.human_input.boundary import HumanInputFormEventFilter from core.workflow.nodes.human_input.callback import ( DifyHITLCallback, ) from core.workflow.nodes.human_input.entities import ( FileInputConfig, FileListInputConfig, + FormDefinition, HumanInputNodeData, ParagraphInputConfig, SelectInputConfig, @@ -52,17 +85,167 @@ from core.workflow.nodes.human_input.enums import ( TimeoutUnit, ValueSourceType, ) +from core.workflow.nodes.human_input.pause_reason import HumanInputRequired +from core.workflow.nodes.human_input.session_binding import SessionBinding from core.workflow.system_variables import build_system_variables +from core.workflow.workflow_entry import WorkflowEntry, iter_dify_graph_engine_events +from graphon.engine import Engine +from graphon.engine.command import InMemoryChannel +from graphon.engine.filter import EngineEventFilterContext, filter_engine_events +from graphon.engine_events import ( + EngineEvent, + GraphEdgeSkippedEvent, + GraphEdgeTakenEvent, + NodeRunHumanInputFormFilledEvent, +) +from graphon.entities import WorkflowStartReason +from graphon.enums import BuiltinNodeTypes, WorkflowNodeExecutionMetadataKey from graphon.file import File, FileTransferMethod, FileType +from graphon.graph import Graph from graphon.node_events import PauseRequestedEvent from graphon.node_events.node import StreamCompletedEvent +from graphon.nodes.answer.answer_node import AnswerNode +from graphon.nodes.answer.entities import AnswerNodeData +from graphon.nodes.end.end_node import EndNode +from graphon.nodes.end.entities import EndNodeData from graphon.nodes.human_input.human_input_node import HumanInputNode from graphon.nodes.protocols import FileReferenceFactoryProtocol -from graphon.runtime import InitParams, RuntimeState, VariablePool +from graphon.nodes.start.entities import StartNodeData +from graphon.nodes.start.start_node import StartNode +from graphon.runtime import InitParams, ReadOnlyRuntimeStateWrapper, RuntimeState, VariablePool +from graphon.runtime.execution import ROOT_FRAME_ID from graphon.variables.segments import ArrayFileSegment, FileSegment, StringSegment from libs.datetime_utils import naive_utc_now -from models.account import TenantAccountJoin, TenantAccountRole -from tests.unit_tests.model_factories import make_account, make_tenant +from libs.helper import compact_generate_response +from models.account import Account, TenantAccountJoin, TenantAccountRole +from models.enums import ConversationFromSource, MessageStatus +from models.human_input import HumanInputForm +from models.model import AppMode +from tests.helpers.unit.model_factories import make_account, make_message, make_tenant + + +def test_session_binding_identity_mapping() -> None: + binding = SessionBinding() + + assert binding.issue_session_id_for_form(form_id="form-1") == "form-1" + assert binding.resolve_form_id_from_session_id(session_id="form-1") == "form-1" + + +def test_human_input_node_contracts_accept_legacy_json_payload() -> None: + payload = { + "form_content": "Please confirm", + "inputs": [ + { + "type": "paragraph", + "output_variable_name": "name", + "default": { + "type": "constant", + "selector": [], + "value": "Alice", + }, + }, + { + "type": "select", + "output_variable_name": "decision", + "option_source": { + "type": "constant", + "selector": [], + "value": ["approve", "reject"], + }, + }, + ], + "user_actions": [ + { + "id": "approve", + "title": "Approve", + "button_style": "primary", + } + ], + "rendered_content": "Please confirm", + "expiration_time": "2024-01-01T00:00:00Z", + "default_values": {"name": "Alice"}, + "node_title": "Human Input", + "display_in_ui": True, + } + + restored = TypeAdapter(FormDefinition).validate_json(json.dumps(payload)) + + assert isinstance(restored.inputs[0], ParagraphInputConfig) + assert isinstance(restored.inputs[1], SelectInputConfig) + assert restored.inputs[0].default is not None + assert restored.inputs[0].default.type == ValueSourceType.CONSTANT + assert restored.inputs[0].default.value == "Alice" + assert restored.inputs[1].option_source.type == ValueSourceType.CONSTANT + assert restored.inputs[1].option_source.value == ["approve", "reject"] + assert restored.user_actions == [UserActionConfig(id="approve", title="Approve", button_style=ButtonStyle.PRIMARY)] + + +def test_human_input_required_pause_reason_keeps_legacy_payload_shape() -> None: + payload = { + "TYPE": "human_input_required", + "form_id": "form-1", + "form_content": "Please confirm", + "inputs": [ + { + "type": "paragraph", + "output_variable_name": "name", + "default": { + "type": "constant", + "selector": [], + "value": "Alice", + }, + } + ], + "actions": [ + { + "id": "approve", + "title": "Approve", + "button_style": "primary", + } + ], + "node_id": "node-1", + "node_title": "Human Input", + "resolved_default_values": {"name": "Alice"}, + } + + restored = TypeAdapter(HumanInputRequired).validate_json(json.dumps(payload)) + + assert restored.TYPE.value == "human_input_required" + assert restored.actions[0].button_style == ButtonStyle.PRIMARY + assert restored.model_dump(mode="json")["TYPE"] == "human_input_required" + + +def test_form_definition_dump_keeps_public_json_shape() -> None: + definition = FormDefinition( + form_content="Please confirm", + inputs=[ParagraphInputConfig(type=FormInputType.PARAGRAPH, output_variable_name="name")], + user_actions=[UserActionConfig(id="approve", title="Approve")], + rendered_content="rendered", + expiration_time=_dify_owned_contracts_datetime(2024, 1, 1, tzinfo=UTC), + default_values={"name": "Alice"}, + node_title="Ask Name", + display_in_ui=True, + ) + + payload = definition.model_dump(mode="json") + + assert payload["expiration_time"] == "2024-01-01T00:00:00Z" + assert payload["user_actions"][0]["id"] == "approve" + assert payload["inputs"][0]["type"] == "paragraph" + + +def test_custom_file_input_requires_extensions() -> None: + with pytest.raises(ExtensionsNotSetErrorValueError): + FileInputConfig( + output_variable_name="attachment", + allowed_file_types=[FileType.CUSTOM], + allowed_file_extensions=[], + ) + + +""" +Unit tests for human input node entities. +""" @dataclass @@ -74,8 +257,8 @@ class _InMemoryFormEntity(HumanInputFormEntity): data: Mapping[str, Any] | None = None is_submitted: bool = False status_value: HumanInputFormStatus = HumanInputFormStatus.WAITING - created: datetime = field(default_factory=naive_utc_now) - expiration: datetime = field(default_factory=lambda: naive_utc_now() + timedelta(days=1)) + created: _entities_datetime = field(default_factory=naive_utc_now) + expiration: _entities_datetime = field(default_factory=lambda: naive_utc_now() + timedelta(days=1)) @property def id(self) -> str: @@ -98,7 +281,7 @@ class _InMemoryFormEntity(HumanInputFormEntity): return self.action_id @property - def created_at(self) -> datetime: + def created_at(self) -> _entities_datetime: return self.created @property @@ -114,7 +297,7 @@ class _InMemoryFormEntity(HumanInputFormEntity): return self.status_value @property - def expiration_time(self) -> datetime: + def expiration_time(self) -> _entities_datetime: return self.expiration @@ -156,7 +339,7 @@ class InMemoryHumanInputFormRepository(HumanInputFormRepository): entity.status_value = HumanInputFormStatus.SUBMITTED -class _TestFileReferenceFactory(FileReferenceFactoryProtocol): +class _EntitiesTestfilereferencefactory(FileReferenceFactoryProtocol): """Build graph-layer file objects without touching Dify persistence in unit tests.""" def build_from_mapping(self, *, mapping: Mapping[str, Any]) -> File: @@ -184,13 +367,13 @@ def _build_human_input_node( typed_node_data = ( node_data if isinstance(node_data, HumanInputNodeData) else HumanInputNodeData.model_validate(node_data) ) - runtime._file_reference_factory = _TestFileReferenceFactory() # type: ignore[attr-defined] + runtime._file_reference_factory = _EntitiesTestfilereferencefactory() # type: ignore[attr-defined] callback = DifyHITLCallback( form_repository=runtime.build_form_repository(), node_data=typed_node_data, delivery_methods=runtime._resolve_delivery_methods(node_data=typed_node_data), display_in_ui=runtime._display_in_ui(node_data=typed_node_data), - file_reference_factory=_TestFileReferenceFactory(), + file_reference_factory=_EntitiesTestfilereferencefactory(), ) return HumanInputNode( node_id=node_id, @@ -959,3 +1142,643 @@ class TestHumanInputNodeRenderedContent: assert node_run_result.outputs["attachment"].value.filename == "resume.pdf" assert isinstance(node_run_result.outputs["attachments"], ArrayFileSegment) assert [file.filename for file in node_run_result.outputs["attachments"].value] == ["a.png", "b.png"] + + +def _persist_form( + sessions: sessionmaker[Session], + *, + form_id: str, + app_id: str, + node_id: str, + expiration_time: _human_input_form_filled_event_datetime.datetime, + status: HumanInputFormStatus, +) -> None: + definition = FormDefinition(form_content="content", rendered_content="content", expiration_time=expiration_time) + form = HumanInputForm( + tenant_id="tenant", + app_id=app_id, + workflow_run_id="run-1", + node_id=node_id, + form_definition=definition.model_dump_json(), + rendered_content="content", + expiration_time=expiration_time, + status=status, + ) + form.id = form_id + with sessions.begin() as session: + session.add(form) + + +class _FakeFormRepository: + def __init__(self, form): + self._form = form + + def get_form(self, *_args, **_kwargs): + return self._form + + def mark_timeout(self, _node_id: str, *, form_id: str): + assert form_id == self._form.id + self._form.status = HumanInputFormStatus.TIMEOUT + return self._form + + +class _HumanInputFormFilledEventTestfilereferencefactory(FileReferenceFactoryProtocol): + def build_from_mapping(self, *, mapping: Mapping[str, Any]): + return File( + file_id=mapping.get("id"), + file_type=FileType(mapping["type"]), + transfer_method=FileTransferMethod(mapping["transfer_method"]), + remote_url=mapping.get("remote_url") or mapping.get("url"), + related_id=mapping.get("related_id") or mapping.get("upload_file_id"), + filename=mapping.get("filename"), + extension=mapping.get("extension"), + mime_type=mapping.get("mime_type"), + size=mapping.get("size", -1), + ) + + +def _create_human_input_node( + *, + config: dict, + init_params: InitParams, + runtime_state: RuntimeState, + repo: _FakeFormRepository, +) -> HumanInputNode: + node_data = ( + config["data"] + if isinstance(config["data"], HumanInputNodeData) + else HumanInputNodeData.model_validate(config["data"]) + ) + callback = DifyHITLCallback( + form_repository=repo, + node_data=node_data, + file_reference_factory=_HumanInputFormFilledEventTestfilereferencefactory(), + ) + node = HumanInputNode( + node_id=config["id"], + data=node_data, + init_params=init_params, + runtime_state=runtime_state, + hitl_callback=callback, + ) + node.bind_execution_id("00000000-0000-4000-8000-000000000001") + return node + + +def _build_node( + form_content: str = ( + "Please enter your name:\n\n{{#$output.name#}}\n" + "Decision: {{#$output.decision#}}\n" + "Attachment: {{#$output.attachment#}}\n" + "Attachments: {{#$output.attachments#}}" + ), + *, + with_inputs: bool = True, + node_id: str = "node-1", +) -> HumanInputNode: + system_variables = build_system_variables(app_id="app", workflow_execution_id="run-1") + graph_runtime_state = RuntimeState( + workflow_id="workflow", + variable_pool=VariablePool.from_bootstrap( + system_variables=system_variables, + user_inputs={}, + environment_variables=[], + ), + start_at=0.0, + ) + graph_init_params = InitParams( + workflow_id="workflow", + graph_config={"nodes": [], "edges": []}, + run_context={ + DIFY_RUN_CONTEXT_KEY: { + "tenant_id": "tenant", + "app_id": "app", + "user_id": "user", + "user_from": UserFrom.ACCOUNT, + "invoke_from": InvokeFrom.SERVICE_API, + } + }, + call_depth=0, + ) + + config = { + "id": node_id, + "type": BuiltinNodeTypes.HUMAN_INPUT, + "data": { + "title": "Human Input", + "form_content": form_content, + "inputs": [ + ParagraphInputConfig(output_variable_name="name").model_dump(mode="json"), + SelectInputConfig( + output_variable_name="decision", + option_source=StringListSource(type="constant", value=["approve", "reject"]), + ).model_dump(mode="json"), + FileInputConfig(output_variable_name="attachment").model_dump(mode="json"), + FileListInputConfig(output_variable_name="attachments", number_limits=2).model_dump(mode="json"), + ], + "user_actions": [UserActionConfig(id="Accept", title="Approve").model_dump(mode="json")], + }, + } + + fake_form = SimpleNamespace( + id="form-1", + rendered_content=form_content, + submitted=True, + selected_action_id="Accept", + submitted_data={ + "name": "Alice", + "decision": "approve", + "attachment": { + "type": "document", + "transfer_method": "remote_url", + "remote_url": "https://example.com/resume.pdf", + "filename": "resume.pdf", + "extension": ".pdf", + "mime_type": "application/pdf", + }, + "attachments": [ + { + "type": "image", + "transfer_method": "remote_url", + "remote_url": "https://example.com/a.png", + "filename": "a.png", + "extension": ".png", + "mime_type": "image/png", + } + ], + }, + status=HumanInputFormStatus.SUBMITTED, + expiration_time=naive_utc_now() + _human_input_form_filled_event_datetime.timedelta(days=1), + ) + + if not with_inputs: + config["data"]["inputs"] = [] + fake_form.submitted_data = {} + + repo = _FakeFormRepository(fake_form) + return _create_human_input_node( + config=config, + init_params=graph_init_params, + runtime_state=graph_runtime_state, + repo=repo, + ) + + +def _build_timeout_node( + expiration_time: _human_input_form_filled_event_datetime.datetime | None = None, + *, + status: HumanInputFormStatus = HumanInputFormStatus.TIMEOUT, + node_id: str = "node-1", +) -> HumanInputNode: + system_variables = build_system_variables(app_id="app", workflow_execution_id="run-1") + graph_runtime_state = RuntimeState( + workflow_id="workflow", + variable_pool=VariablePool.from_bootstrap( + system_variables=system_variables, + user_inputs={}, + environment_variables=[], + ), + start_at=0.0, + ) + graph_init_params = InitParams( + workflow_id="workflow", + graph_config={"nodes": [], "edges": []}, + run_context={ + DIFY_RUN_CONTEXT_KEY: { + "tenant_id": "tenant", + "app_id": "app", + "user_id": "user", + "user_from": UserFrom.ACCOUNT, + "invoke_from": InvokeFrom.SERVICE_API, + } + }, + call_depth=0, + ) + + config = { + "id": node_id, + "type": BuiltinNodeTypes.HUMAN_INPUT, + "data": { + "title": "Human Input", + "form_content": "Please enter your name:\n\n{{#$output.name#}}", + "inputs": [ParagraphInputConfig(output_variable_name="name").model_dump(mode="json")], + "user_actions": [UserActionConfig(id="Accept", title="Approve").model_dump(mode="json")], + }, + } + + fake_form = SimpleNamespace( + id="form-1", + rendered_content="content", + submitted=False, + selected_action_id=None, + submitted_data=None, + status=status, + created_at=naive_utc_now(), + expiration_time=expiration_time + or naive_utc_now() - _human_input_form_filled_event_datetime.timedelta(minutes=1), + ) + + repo = _FakeFormRepository(fake_form) + return _create_human_input_node( + config=config, + init_params=graph_init_params, + runtime_state=graph_runtime_state, + repo=repo, + ) + + +def _publish_node_events(node: HumanInputNode) -> list[AppQueueEvent]: + return _publish_graph_events( + _filter_human_input_events(node.run(), node=node), + graph=Graph(root_node=node), + runtime_state=node.runtime_state, + ) + + +def _filter_human_input_events(events: Iterable[EngineEvent], *, node: HumanInputNode | None = None): + node = node or _build_node() + return filter_engine_events( + events, + context=EngineEventFilterContext( + graph=Graph(root_node=node), + runtime_state=ReadOnlyRuntimeStateWrapper(node.runtime_state), + ), + filters=[HumanInputFormEventFilter(form_repository=HumanInputFormSubmissionRepository())], + ) + + +@pytest.mark.parametrize("event_type", [GraphEdgeTakenEvent, GraphEdgeSkippedEvent]) +def test_human_input_filter_forwards_traversals_without_waiting_for_completion(event_type): + started, _ = list(_build_node().run()) + edge = event_type(frame_id=ROOT_FRAME_ID, edge_id="human-answer", source_node_id="node-1", target_node_id="answer") + edge_forwarded = False + + def source() -> Generator[EngineEvent]: + yield started + yield edge + assert edge_forwarded, "The traversal must be forwarded before consuming more upstream events" + + events = iter(_filter_human_input_events(source())) + assert next(events) == started + assert next(events) == edge + edge_forwarded = True + assert list(events) == [] + + +def test_form_events_keep_titles_for_interleaved_executions_of_one_node(): + node = _build_node() + started, succeeded = list(node.run()) + first_start = started.model_copy(update={"id": "first", "node_title": "First approval"}) + second_start = started.model_copy(update={"id": "second", "node_title": "Second approval"}) + first_result = succeeded.model_copy( + update={ + "id": "first", + "container_id": "loop-1", + "node_run_result": succeeded.node_run_result.model_copy( + update={"metadata": {WorkflowNodeExecutionMetadataKey.LOOP_ID: "loop-1"}} + ), + } + ) + second_result = succeeded.model_copy( + update={ + "id": "second", + "container_id": "iteration-2", + "node_run_result": succeeded.node_run_result.model_copy( + update={"metadata": {WorkflowNodeExecutionMetadataKey.ITERATION_ID: "iteration-2"}} + ), + } + ) + + filtered = list(_filter_human_input_events([first_start, second_start, second_result, first_result], node=node)) + form_events = [event for event in filtered if isinstance(event, NodeRunHumanInputFormFilledEvent)] + assert [(event.id, event.container_id, event.node_run_result.metadata) for event in form_events] == [ + ("second", "iteration-2", second_result.node_run_result.metadata), + ("first", "loop-1", first_result.node_run_result.metadata), + ] + events = _publish_graph_events(filtered, graph=Graph(root_node=node), runtime_state=node.runtime_state) + + filled = [event for event in events if isinstance(event, QueueHumanInputFormFilledEvent)] + assert [(event.form_id, event.node_title) for event in filled] == [ + ("second", "Second approval"), + ("first", "First approval"), + ] + completed = [event for event in events if isinstance(event, QueueNodeSucceededEvent)] + assert [(event.node_execution_id, event.in_loop_id, event.in_iteration_id) for event in completed] == [ + ("second", None, "iteration-2"), + ("first", "loop-1", None), + ] + + +class _EventQueue(AppQueueManager): + """Record events and replay a finite stream without starting a watchdog.""" + + def __init__(self, invoke_from: InvokeFrom, runtime_state: RuntimeState) -> None: + self.invoke_from = invoke_from + self.graph_runtime_state = runtime_state + self.events: list[AppQueueEvent] = [] + + def _publish(self, event: AppQueueEvent, pub_from: PublishFrom) -> None: + assert pub_from in {PublishFrom.APPLICATION_MANAGER, PublishFrom.TASK_PIPELINE} + self.events.append(event) + + def listen(self) -> Generator[WorkflowQueueMessage]: + for event in self.events: + yield WorkflowQueueMessage(task_id="task-1", app_mode=AppMode.ADVANCED_CHAT, event=event) + + +def _publish_graph_events( + events: Iterable[EngineEvent], *, graph: Graph, runtime_state: RuntimeState +) -> list[AppQueueEvent]: + queue_manager = _EventQueue(InvokeFrom.WEB_APP, runtime_state) + runner = WorkflowBasedAppRunner(queue_manager=queue_manager, app_id="app") + workflow_entry = WorkflowEntry( + graph_config={}, + graph=graph, + call_depth=0, + graph_runtime_state=runtime_state, + ) + for event in events: + runner.handle_event(workflow_entry, event) + return queue_manager.events + + +def _sse_payloads( + events: Sequence[AppQueueEvent], + invoke_from: InvokeFrom, + app: Flask, + sessions: sessionmaker[Session], + runtime_state: RuntimeState | None = None, +) -> list[dict[str, Any]]: + generate_entity = AdvancedChatAppGenerateEntity( + task_id="task-1", + app_config=WorkflowUIBasedAppConfig( + tenant_id="tenant", + app_id="app", + app_mode=AppMode.ADVANCED_CHAT, + workflow_id="workflow", + ), + inputs={}, + query="hello", + files=[], + user_id="user", + stream=True, + invoke_from=invoke_from, + workflow_run_id="run-1", + ) + if runtime_state is None: + runtime_state = RuntimeState( + workflow_id="workflow", + variable_pool=VariablePool.from_bootstrap( + system_variables=build_system_variables(workflow_execution_id="run-1") + ), + start_at=0, + ) + queue_manager = _EventQueue(invoke_from, runtime_state) + queue_manager.events.extend(events) + with sessions.begin() as session: + session.add( + make_message( + app_id="app", + inputs={}, + query="hello", + message={}, + answer="", + status=MessageStatus.PAUSED, + message_unit_price=Decimal(0), + answer_unit_price=Decimal(0), + total_price=Decimal(0), + currency="USD", + from_source=ConversationFromSource.API, + ) + ) + has_form = session.scalar(select(HumanInputForm.id).where(HumanInputForm.workflow_run_id == "run-1")) + if has_form is None: + for event in events: + if isinstance(event, QueueHumanInputFormFilledEvent): + _persist_form( + sessions, + form_id=event.form_id, + app_id="app", + node_id=event.node_id, + expiration_time=naive_utc_now() + _human_input_form_filled_event_datetime.timedelta(days=1), + status=HumanInputFormStatus.SUBMITTED, + ) + pipeline = AdvancedChatAppGenerateTaskPipeline( + application_generate_entity=generate_entity, + workflow=WorkflowSnapshot(id="workflow", tenant_id="tenant", features_dict={}), + queue_manager=queue_manager, + conversation=ConversationSnapshot(id="conversation-1", mode=AppMode.ADVANCED_CHAT), + message=MessageSnapshot( + id="message-1", query="hello", created_at=naive_utc_now(), status=MessageStatus.PAUSED, answer="" + ), + user=Account(name="tester", email="tester@example.com"), + stream=True, + dialogue_count=1, + draft_var_saver_factory=MagicMock(), + ) + with app.test_request_context(): + if not any(isinstance(event, QueueWorkflowStartedEvent) for event in events): + list( + pipeline._handle_workflow_started_event( + QueueWorkflowStartedEvent(reason=WorkflowStartReason.RESUMPTION) + ) + ) + response = compact_generate_response( + BaseAppGenerator.convert_to_event_stream( + AdvancedChatAppGenerateResponseConverter.convert( + pipeline._to_stream_response(pipeline._process_stream_response()), invoke_from + ) + ) + ) + assert response.mimetype == "text/event-stream" + return [ + json.loads(line.removeprefix("data: ")) for line in response.get_data(as_text=True).splitlines() if line + ] + + +@pytest.mark.parametrize("invoke_from", [InvokeFrom.DEBUGGER, InvokeFrom.WEB_APP], ids=["full", "simple"]) +def test_submitted_human_input_reaches_response_stream( + invoke_from: InvokeFrom, app: Flask, sqlite_session_factory: sessionmaker[Session] +): + events = _publish_node_events(_build_node()) + + payloads = _sse_payloads(events, invoke_from, app, sqlite_session_factory) + + # Dify's original HumanInputNode (9c339239850) emitted the form event + # before node completion; its runner and task pipeline preserved that order. + assert [payload["event"] for payload in payloads] == [ + "node_started", + "human_input_form_filled", + "node_finished", + ] + payload = payloads[1] + assert payload["workflow_run_id"] == "run-1" + assert payload["data"]["form_id"] == "00000000-0000-4000-8000-000000000001" + assert payload["data"]["node_id"] == "node-1" + assert payload["data"]["node_title"] == "Human Input" + assert payload["data"]["action_id"] == "Accept" + assert payload["data"]["action_text"] == "Approve" + assert payload["data"]["rendered_content"] == ( + "Please enter your name:\n\nAlice\nDecision: approve\nAttachment: [file]\nAttachments: [1 files]" + ) + submitted_data = payload["data"]["submitted_data"] + assert submitted_data["name"] == "Alice" + assert submitted_data["decision"] == "approve" + assert submitted_data["attachment"]["filename"] == "resume.pdf" + assert submitted_data["attachment"]["type"] == "document" + assert submitted_data["attachment"]["transfer_method"] == "remote_url" + assert submitted_data["attachments"][0]["filename"] == "a.png" + assert submitted_data["attachments"][0]["type"] == "image" + + +def test_button_only_human_input_reaches_response_stream(app: Flask, sqlite_session_factory: sessionmaker[Session]): + events = _publish_node_events(_build_node("Approve deployment?", with_inputs=False)) + + payloads = _sse_payloads(events, InvokeFrom.WEB_APP, app, sqlite_session_factory) + + assert [payload["event"] for payload in payloads] == [ + "node_started", + "human_input_form_filled", + "node_finished", + ] + assert payloads[1]["data"]["rendered_content"] == "Approve deployment?" + assert payloads[1]["data"]["submitted_data"] == {} + assert payloads[1]["data"]["action_id"] == "Accept" + + +@pytest.mark.parametrize("status", [HumanInputFormStatus.TIMEOUT, HumanInputFormStatus.WAITING]) +def test_timed_out_human_input_reaches_response_stream( + status: HumanInputFormStatus, sqlite_session_factory: sessionmaker[Session], app: Flask +): + expiration_time = _human_input_form_filled_event_datetime.datetime(2025, 1, 1) + for form_id, expiration in ( + ("previous-execution", _human_input_form_filled_event_datetime.datetime(2024, 1, 1)), + ("00000000-0000-4000-8000-000000000001", expiration_time), + ): + _persist_form( + sqlite_session_factory, + form_id=form_id, + app_id="app", + node_id="node-1", + expiration_time=expiration, + status=status, + ) + + events = _publish_node_events(_build_timeout_node(expiration_time, status=status)) + + payloads = _sse_payloads(events, InvokeFrom.WEB_APP, app, sqlite_session_factory) + assert [payload["event"] for payload in payloads] == [ + "node_started", + "human_input_form_timeout", + "node_finished", + ] + assert payloads[1]["data"] == { + "form_id": "00000000-0000-4000-8000-000000000001", + "node_id": "node-1", + "node_title": "Human Input", + "expiration_time": 1735689600, + } + + +@pytest.mark.parametrize("timed_out", [False, True]) +@pytest.mark.parametrize("terminal", ["end", "answer"]) +def test_human_input_completion_and_referenced_answer_reach_response_stream( + timed_out: bool, terminal: str, mocker: MockerFixture, app: Flask, sqlite_session_factory: sessionmaker[Session] +): + expiration_time = _human_input_form_filled_event_datetime.datetime(2025, 1, 1) + node = ( + _build_timeout_node(expiration_time, node_id="human") + if timed_out + else _build_node("Approve?", with_inputs=False, node_id="human") + ) + runtime_state = node.runtime_state + runtime_state.variable_pool.add(["sys", "workflow_execution_id"], StringSegment(value="run-1")) + init_params = InitParams(workflow_id="workflow", graph_config={}, run_context={}, call_depth=0) + start = StartNode( + node_id="start", + data=StartNodeData(title="Start", variables=[]), + init_params=init_params, + runtime_state=runtime_state, + ) + if terminal == "answer": + terminal_node = AnswerNode( + node_id="answer", + data=AnswerNodeData(title="Answer", answer="Action: {{#human.__action_id#}}"), + init_params=init_params, + runtime_state=runtime_state, + ) + else: + terminal_node = EndNode( + node_id="end", + data=EndNodeData(title="End", outputs=[]), + init_params=init_params, + runtime_state=runtime_state, + ) + graph = ( + Graph.new() + .add_root(start) + .add_node(node, from_node_id="start") + .add_node(terminal_node, from_node_id=node.id, source_handle="__timeout" if timed_out else "Accept") + .build() + ) + engine = Engine( + graph=graph, + runtime_state=runtime_state, + command_channel=InMemoryChannel(), + workers=1, + ) + bind_execution_id = node.bind_execution_id + + def bind_form(execution_id: str) -> None: + bind_execution_id(execution_id) + _persist_form( + sqlite_session_factory, + form_id=execution_id, + app_id="app", + node_id=node.id, + expiration_time=expiration_time, + status=HumanInputFormStatus.TIMEOUT if timed_out else HumanInputFormStatus.SUBMITTED, + ) + + mocker.patch.object(node, "bind_execution_id", side_effect=bind_form) + + events = _publish_graph_events(iter_dify_graph_engine_events(engine), graph=graph, runtime_state=runtime_state) + payloads = _sse_payloads(events, InvokeFrom.WEB_APP, app, sqlite_session_factory, runtime_state) + + form_event = "human_input_form_timeout" if timed_out else "human_input_form_filled" + lifecycle_events = [payload for payload in payloads if payload["event"] != "message"] + assert [(payload["event"], payload.get("data", {}).get("node_id")) for payload in lifecycle_events] == [ + ("workflow_started", None), + ("node_started", "start"), + ("node_finished", "start"), + ("node_started", node.id), + (form_event, node.id), + ("node_finished", node.id), + ("node_started", terminal), + ("node_finished", terminal), + ("message_end", None), + ("workflow_finished", None), + ] + answer = "".join(payload["answer"] for payload in payloads if payload["event"] == "message") + expected_answer = "Action: " + ("" if timed_out else "Accept") + assert answer == (expected_answer if terminal == "answer" else "") + assert payloads[-1]["data"]["status"] == "succeeded" + + +@pytest.mark.parametrize("form_owner", [None, ("other-app", "node-1"), ("app", "other-node")]) +def test_timeout_rejects_missing_or_unrelated_form( + form_owner: tuple[str, str] | None, sqlite_session_factory: sessionmaker[Session] +): + if form_owner is not None: + app_id, node_id = form_owner + _persist_form( + sqlite_session_factory, + form_id="00000000-0000-4000-8000-000000000001", + app_id=app_id, + node_id=node_id, + expiration_time=_human_input_form_filled_event_datetime.datetime(2025, 1, 1), + status=HumanInputFormStatus.TIMEOUT, + ) + + with pytest.raises(ValueError, match="Cannot resolve timed-out human input form"): + _publish_node_events(_build_timeout_node()) diff --git a/api/tests/unit_tests/core/workflow/nodes/human_input/test_human_input_form_filled_event.py b/api/tests/unit_tests/core/workflow/nodes/human_input/test_human_input_form_filled_event.py deleted file mode 100644 index d5e203aa12a..00000000000 --- a/api/tests/unit_tests/core/workflow/nodes/human_input/test_human_input_form_filled_event.py +++ /dev/null @@ -1,728 +0,0 @@ -import datetime -import json -from collections.abc import Generator, Iterable, Mapping, Sequence -from decimal import Decimal -from types import SimpleNamespace -from typing import Any -from unittest.mock import MagicMock - -import pytest -from flask import Flask -from pytest_mock import MockerFixture -from sqlalchemy import select -from sqlalchemy.orm import Session, sessionmaker - -from core.app.app_config.entities import WorkflowUIBasedAppConfig -from core.app.apps.advanced_chat.generate_response_converter import AdvancedChatAppGenerateResponseConverter -from core.app.apps.advanced_chat.generate_task_pipeline import ( - AdvancedChatAppGenerateTaskPipeline, - ConversationSnapshot, - MessageSnapshot, - WorkflowSnapshot, -) -from core.app.apps.base_app_generator import BaseAppGenerator -from core.app.apps.base_app_queue_manager import AppQueueManager, PublishFrom -from core.app.apps.workflow_app_runner import WorkflowBasedAppRunner -from core.app.entities.app_invoke_entities import ( - DIFY_RUN_CONTEXT_KEY, - AdvancedChatAppGenerateEntity, - InvokeFrom, - UserFrom, -) -from core.app.entities.queue_entities import ( - AppQueueEvent, - QueueHumanInputFormFilledEvent, - QueueNodeSucceededEvent, - QueueWorkflowStartedEvent, - WorkflowQueueMessage, -) -from core.repositories.human_input_repository import ( - HumanInputFormSubmissionRepository, -) -from core.workflow.nodes.human_input.boundary import HumanInputFormEventFilter -from core.workflow.nodes.human_input.callback import ( - DifyHITLCallback, -) -from core.workflow.nodes.human_input.entities import ( - FileInputConfig, - FileListInputConfig, - FormDefinition, - HumanInputNodeData, - ParagraphInputConfig, - SelectInputConfig, - StringListSource, - UserActionConfig, -) -from core.workflow.nodes.human_input.enums import HumanInputFormStatus -from core.workflow.system_variables import build_system_variables -from core.workflow.workflow_entry import WorkflowEntry, iter_dify_graph_engine_events -from graphon.engine import Engine -from graphon.engine.command import InMemoryChannel -from graphon.engine.filter import EngineEventFilterContext, filter_engine_events -from graphon.engine_events import ( - EngineEvent, - GraphEdgeSkippedEvent, - GraphEdgeTakenEvent, - NodeRunHumanInputFormFilledEvent, -) -from graphon.entities import WorkflowStartReason -from graphon.enums import BuiltinNodeTypes, WorkflowNodeExecutionMetadataKey -from graphon.file import File, FileTransferMethod, FileType -from graphon.graph import Graph -from graphon.nodes.answer.answer_node import AnswerNode -from graphon.nodes.answer.entities import AnswerNodeData -from graphon.nodes.end.end_node import EndNode -from graphon.nodes.end.entities import EndNodeData -from graphon.nodes.human_input.human_input_node import HumanInputNode -from graphon.nodes.protocols import FileReferenceFactoryProtocol -from graphon.nodes.start.entities import StartNodeData -from graphon.nodes.start.start_node import StartNode -from graphon.runtime import InitParams, ReadOnlyRuntimeStateWrapper, RuntimeState, VariablePool -from graphon.runtime.execution import ROOT_FRAME_ID -from graphon.variables.segments import StringSegment -from libs.datetime_utils import naive_utc_now -from libs.helper import compact_generate_response -from models.account import Account -from models.enums import ConversationFromSource, MessageStatus -from models.human_input import HumanInputForm -from models.model import AppMode -from tests.unit_tests.model_factories import make_message - - -def _persist_form( - sessions: sessionmaker[Session], - *, - form_id: str, - app_id: str, - node_id: str, - expiration_time: datetime.datetime, - status: HumanInputFormStatus, -) -> None: - definition = FormDefinition(form_content="content", rendered_content="content", expiration_time=expiration_time) - form = HumanInputForm( - tenant_id="tenant", - app_id=app_id, - workflow_run_id="run-1", - node_id=node_id, - form_definition=definition.model_dump_json(), - rendered_content="content", - expiration_time=expiration_time, - status=status, - ) - form.id = form_id - with sessions.begin() as session: - session.add(form) - - -class _FakeFormRepository: - def __init__(self, form): - self._form = form - - def get_form(self, *_args, **_kwargs): - return self._form - - def mark_timeout(self, _node_id: str, *, form_id: str): - assert form_id == self._form.id - self._form.status = HumanInputFormStatus.TIMEOUT - return self._form - - -class _TestFileReferenceFactory(FileReferenceFactoryProtocol): - def build_from_mapping(self, *, mapping: Mapping[str, Any]): - return File( - file_id=mapping.get("id"), - file_type=FileType(mapping["type"]), - transfer_method=FileTransferMethod(mapping["transfer_method"]), - remote_url=mapping.get("remote_url") or mapping.get("url"), - related_id=mapping.get("related_id") or mapping.get("upload_file_id"), - filename=mapping.get("filename"), - extension=mapping.get("extension"), - mime_type=mapping.get("mime_type"), - size=mapping.get("size", -1), - ) - - -def _create_human_input_node( - *, - config: dict, - init_params: InitParams, - runtime_state: RuntimeState, - repo: _FakeFormRepository, -) -> HumanInputNode: - node_data = ( - config["data"] - if isinstance(config["data"], HumanInputNodeData) - else HumanInputNodeData.model_validate(config["data"]) - ) - callback = DifyHITLCallback( - form_repository=repo, - node_data=node_data, - file_reference_factory=_TestFileReferenceFactory(), - ) - node = HumanInputNode( - node_id=config["id"], - data=node_data, - init_params=init_params, - runtime_state=runtime_state, - hitl_callback=callback, - ) - node.bind_execution_id("00000000-0000-4000-8000-000000000001") - return node - - -def _build_node( - form_content: str = ( - "Please enter your name:\n\n{{#$output.name#}}\n" - "Decision: {{#$output.decision#}}\n" - "Attachment: {{#$output.attachment#}}\n" - "Attachments: {{#$output.attachments#}}" - ), - *, - with_inputs: bool = True, - node_id: str = "node-1", -) -> HumanInputNode: - system_variables = build_system_variables(app_id="app", workflow_execution_id="run-1") - graph_runtime_state = RuntimeState( - workflow_id="workflow", - variable_pool=VariablePool.from_bootstrap( - system_variables=system_variables, - user_inputs={}, - environment_variables=[], - ), - start_at=0.0, - ) - graph_init_params = InitParams( - workflow_id="workflow", - graph_config={"nodes": [], "edges": []}, - run_context={ - DIFY_RUN_CONTEXT_KEY: { - "tenant_id": "tenant", - "app_id": "app", - "user_id": "user", - "user_from": UserFrom.ACCOUNT, - "invoke_from": InvokeFrom.SERVICE_API, - } - }, - call_depth=0, - ) - - config = { - "id": node_id, - "type": BuiltinNodeTypes.HUMAN_INPUT, - "data": { - "title": "Human Input", - "form_content": form_content, - "inputs": [ - ParagraphInputConfig(output_variable_name="name").model_dump(mode="json"), - SelectInputConfig( - output_variable_name="decision", - option_source=StringListSource(type="constant", value=["approve", "reject"]), - ).model_dump(mode="json"), - FileInputConfig(output_variable_name="attachment").model_dump(mode="json"), - FileListInputConfig(output_variable_name="attachments", number_limits=2).model_dump(mode="json"), - ], - "user_actions": [UserActionConfig(id="Accept", title="Approve").model_dump(mode="json")], - }, - } - - fake_form = SimpleNamespace( - id="form-1", - rendered_content=form_content, - submitted=True, - selected_action_id="Accept", - submitted_data={ - "name": "Alice", - "decision": "approve", - "attachment": { - "type": "document", - "transfer_method": "remote_url", - "remote_url": "https://example.com/resume.pdf", - "filename": "resume.pdf", - "extension": ".pdf", - "mime_type": "application/pdf", - }, - "attachments": [ - { - "type": "image", - "transfer_method": "remote_url", - "remote_url": "https://example.com/a.png", - "filename": "a.png", - "extension": ".png", - "mime_type": "image/png", - } - ], - }, - status=HumanInputFormStatus.SUBMITTED, - expiration_time=naive_utc_now() + datetime.timedelta(days=1), - ) - - if not with_inputs: - config["data"]["inputs"] = [] - fake_form.submitted_data = {} - - repo = _FakeFormRepository(fake_form) - return _create_human_input_node( - config=config, - init_params=graph_init_params, - runtime_state=graph_runtime_state, - repo=repo, - ) - - -def _build_timeout_node( - expiration_time: datetime.datetime | None = None, - *, - status: HumanInputFormStatus = HumanInputFormStatus.TIMEOUT, - node_id: str = "node-1", -) -> HumanInputNode: - system_variables = build_system_variables(app_id="app", workflow_execution_id="run-1") - graph_runtime_state = RuntimeState( - workflow_id="workflow", - variable_pool=VariablePool.from_bootstrap( - system_variables=system_variables, - user_inputs={}, - environment_variables=[], - ), - start_at=0.0, - ) - graph_init_params = InitParams( - workflow_id="workflow", - graph_config={"nodes": [], "edges": []}, - run_context={ - DIFY_RUN_CONTEXT_KEY: { - "tenant_id": "tenant", - "app_id": "app", - "user_id": "user", - "user_from": UserFrom.ACCOUNT, - "invoke_from": InvokeFrom.SERVICE_API, - } - }, - call_depth=0, - ) - - config = { - "id": node_id, - "type": BuiltinNodeTypes.HUMAN_INPUT, - "data": { - "title": "Human Input", - "form_content": "Please enter your name:\n\n{{#$output.name#}}", - "inputs": [ParagraphInputConfig(output_variable_name="name").model_dump(mode="json")], - "user_actions": [UserActionConfig(id="Accept", title="Approve").model_dump(mode="json")], - }, - } - - fake_form = SimpleNamespace( - id="form-1", - rendered_content="content", - submitted=False, - selected_action_id=None, - submitted_data=None, - status=status, - created_at=naive_utc_now(), - expiration_time=expiration_time or naive_utc_now() - datetime.timedelta(minutes=1), - ) - - repo = _FakeFormRepository(fake_form) - return _create_human_input_node( - config=config, - init_params=graph_init_params, - runtime_state=graph_runtime_state, - repo=repo, - ) - - -def _publish_node_events(node: HumanInputNode) -> list[AppQueueEvent]: - return _publish_graph_events( - _filter_human_input_events(node.run(), node=node), - graph=Graph(root_node=node), - runtime_state=node.runtime_state, - ) - - -def _filter_human_input_events(events: Iterable[EngineEvent], *, node: HumanInputNode | None = None): - node = node or _build_node() - return filter_engine_events( - events, - context=EngineEventFilterContext( - graph=Graph(root_node=node), - runtime_state=ReadOnlyRuntimeStateWrapper(node.runtime_state), - ), - filters=[HumanInputFormEventFilter(form_repository=HumanInputFormSubmissionRepository())], - ) - - -@pytest.mark.parametrize("event_type", [GraphEdgeTakenEvent, GraphEdgeSkippedEvent]) -def test_human_input_filter_forwards_traversals_without_waiting_for_completion(event_type): - started, _ = list(_build_node().run()) - edge = event_type(frame_id=ROOT_FRAME_ID, edge_id="human-answer", source_node_id="node-1", target_node_id="answer") - edge_forwarded = False - - def source() -> Generator[EngineEvent]: - yield started - yield edge - assert edge_forwarded, "The traversal must be forwarded before consuming more upstream events" - - events = iter(_filter_human_input_events(source())) - assert next(events) == started - assert next(events) == edge - edge_forwarded = True - assert list(events) == [] - - -def test_form_events_keep_titles_for_interleaved_executions_of_one_node(): - node = _build_node() - started, succeeded = list(node.run()) - first_start = started.model_copy(update={"id": "first", "node_title": "First approval"}) - second_start = started.model_copy(update={"id": "second", "node_title": "Second approval"}) - first_result = succeeded.model_copy( - update={ - "id": "first", - "container_id": "loop-1", - "node_run_result": succeeded.node_run_result.model_copy( - update={"metadata": {WorkflowNodeExecutionMetadataKey.LOOP_ID: "loop-1"}} - ), - } - ) - second_result = succeeded.model_copy( - update={ - "id": "second", - "container_id": "iteration-2", - "node_run_result": succeeded.node_run_result.model_copy( - update={"metadata": {WorkflowNodeExecutionMetadataKey.ITERATION_ID: "iteration-2"}} - ), - } - ) - - filtered = list(_filter_human_input_events([first_start, second_start, second_result, first_result], node=node)) - form_events = [event for event in filtered if isinstance(event, NodeRunHumanInputFormFilledEvent)] - assert [(event.id, event.container_id, event.node_run_result.metadata) for event in form_events] == [ - ("second", "iteration-2", second_result.node_run_result.metadata), - ("first", "loop-1", first_result.node_run_result.metadata), - ] - events = _publish_graph_events(filtered, graph=Graph(root_node=node), runtime_state=node.runtime_state) - - filled = [event for event in events if isinstance(event, QueueHumanInputFormFilledEvent)] - assert [(event.form_id, event.node_title) for event in filled] == [ - ("second", "Second approval"), - ("first", "First approval"), - ] - completed = [event for event in events if isinstance(event, QueueNodeSucceededEvent)] - assert [(event.node_execution_id, event.in_loop_id, event.in_iteration_id) for event in completed] == [ - ("second", None, "iteration-2"), - ("first", "loop-1", None), - ] - - -class _EventQueue(AppQueueManager): - """Record events and replay a finite stream without starting a watchdog.""" - - def __init__(self, invoke_from: InvokeFrom, runtime_state: RuntimeState) -> None: - self.invoke_from = invoke_from - self.graph_runtime_state = runtime_state - self.events: list[AppQueueEvent] = [] - - def _publish(self, event: AppQueueEvent, pub_from: PublishFrom) -> None: - assert pub_from in {PublishFrom.APPLICATION_MANAGER, PublishFrom.TASK_PIPELINE} - self.events.append(event) - - def listen(self) -> Generator[WorkflowQueueMessage]: - for event in self.events: - yield WorkflowQueueMessage(task_id="task-1", app_mode=AppMode.ADVANCED_CHAT, event=event) - - -def _publish_graph_events( - events: Iterable[EngineEvent], *, graph: Graph, runtime_state: RuntimeState -) -> list[AppQueueEvent]: - queue_manager = _EventQueue(InvokeFrom.WEB_APP, runtime_state) - runner = WorkflowBasedAppRunner(queue_manager=queue_manager, app_id="app") - workflow_entry = WorkflowEntry( - graph_config={}, - graph=graph, - call_depth=0, - graph_runtime_state=runtime_state, - ) - for event in events: - runner.handle_event(workflow_entry, event) - return queue_manager.events - - -def _sse_payloads( - events: Sequence[AppQueueEvent], - invoke_from: InvokeFrom, - app: Flask, - sessions: sessionmaker[Session], - runtime_state: RuntimeState | None = None, -) -> list[dict[str, Any]]: - generate_entity = AdvancedChatAppGenerateEntity( - task_id="task-1", - app_config=WorkflowUIBasedAppConfig( - tenant_id="tenant", - app_id="app", - app_mode=AppMode.ADVANCED_CHAT, - workflow_id="workflow", - ), - inputs={}, - query="hello", - files=[], - user_id="user", - stream=True, - invoke_from=invoke_from, - workflow_run_id="run-1", - ) - if runtime_state is None: - runtime_state = RuntimeState( - workflow_id="workflow", - variable_pool=VariablePool.from_bootstrap( - system_variables=build_system_variables(workflow_execution_id="run-1") - ), - start_at=0, - ) - queue_manager = _EventQueue(invoke_from, runtime_state) - queue_manager.events.extend(events) - with sessions.begin() as session: - session.add( - make_message( - app_id="app", - inputs={}, - query="hello", - message={}, - answer="", - status=MessageStatus.PAUSED, - message_unit_price=Decimal(0), - answer_unit_price=Decimal(0), - total_price=Decimal(0), - currency="USD", - from_source=ConversationFromSource.API, - ) - ) - has_form = session.scalar(select(HumanInputForm.id).where(HumanInputForm.workflow_run_id == "run-1")) - if has_form is None: - for event in events: - if isinstance(event, QueueHumanInputFormFilledEvent): - _persist_form( - sessions, - form_id=event.form_id, - app_id="app", - node_id=event.node_id, - expiration_time=naive_utc_now() + datetime.timedelta(days=1), - status=HumanInputFormStatus.SUBMITTED, - ) - pipeline = AdvancedChatAppGenerateTaskPipeline( - application_generate_entity=generate_entity, - workflow=WorkflowSnapshot(id="workflow", tenant_id="tenant", features_dict={}), - queue_manager=queue_manager, - conversation=ConversationSnapshot(id="conversation-1", mode=AppMode.ADVANCED_CHAT), - message=MessageSnapshot( - id="message-1", query="hello", created_at=naive_utc_now(), status=MessageStatus.PAUSED, answer="" - ), - user=Account(name="tester", email="tester@example.com"), - stream=True, - dialogue_count=1, - draft_var_saver_factory=MagicMock(), - ) - with app.test_request_context(): - if not any(isinstance(event, QueueWorkflowStartedEvent) for event in events): - list( - pipeline._handle_workflow_started_event( - QueueWorkflowStartedEvent(reason=WorkflowStartReason.RESUMPTION) - ) - ) - response = compact_generate_response( - BaseAppGenerator.convert_to_event_stream( - AdvancedChatAppGenerateResponseConverter.convert( - pipeline._to_stream_response(pipeline._process_stream_response()), invoke_from - ) - ) - ) - assert response.mimetype == "text/event-stream" - return [ - json.loads(line.removeprefix("data: ")) for line in response.get_data(as_text=True).splitlines() if line - ] - - -@pytest.mark.parametrize("invoke_from", [InvokeFrom.DEBUGGER, InvokeFrom.WEB_APP], ids=["full", "simple"]) -def test_submitted_human_input_reaches_response_stream( - invoke_from: InvokeFrom, app: Flask, sqlite_session_factory: sessionmaker[Session] -): - events = _publish_node_events(_build_node()) - - payloads = _sse_payloads(events, invoke_from, app, sqlite_session_factory) - - # Dify's original HumanInputNode (9c339239850) emitted the form event - # before node completion; its runner and task pipeline preserved that order. - assert [payload["event"] for payload in payloads] == [ - "node_started", - "human_input_form_filled", - "node_finished", - ] - payload = payloads[1] - assert payload["workflow_run_id"] == "run-1" - assert payload["data"]["form_id"] == "00000000-0000-4000-8000-000000000001" - assert payload["data"]["node_id"] == "node-1" - assert payload["data"]["node_title"] == "Human Input" - assert payload["data"]["action_id"] == "Accept" - assert payload["data"]["action_text"] == "Approve" - assert payload["data"]["rendered_content"] == ( - "Please enter your name:\n\nAlice\nDecision: approve\nAttachment: [file]\nAttachments: [1 files]" - ) - submitted_data = payload["data"]["submitted_data"] - assert submitted_data["name"] == "Alice" - assert submitted_data["decision"] == "approve" - assert submitted_data["attachment"]["filename"] == "resume.pdf" - assert submitted_data["attachment"]["type"] == "document" - assert submitted_data["attachment"]["transfer_method"] == "remote_url" - assert submitted_data["attachments"][0]["filename"] == "a.png" - assert submitted_data["attachments"][0]["type"] == "image" - - -def test_button_only_human_input_reaches_response_stream(app: Flask, sqlite_session_factory: sessionmaker[Session]): - events = _publish_node_events(_build_node("Approve deployment?", with_inputs=False)) - - payloads = _sse_payloads(events, InvokeFrom.WEB_APP, app, sqlite_session_factory) - - assert [payload["event"] for payload in payloads] == [ - "node_started", - "human_input_form_filled", - "node_finished", - ] - assert payloads[1]["data"]["rendered_content"] == "Approve deployment?" - assert payloads[1]["data"]["submitted_data"] == {} - assert payloads[1]["data"]["action_id"] == "Accept" - - -@pytest.mark.parametrize("status", [HumanInputFormStatus.TIMEOUT, HumanInputFormStatus.WAITING]) -def test_timed_out_human_input_reaches_response_stream( - status: HumanInputFormStatus, sqlite_session_factory: sessionmaker[Session], app: Flask -): - expiration_time = datetime.datetime(2025, 1, 1) - for form_id, expiration in ( - ("previous-execution", datetime.datetime(2024, 1, 1)), - ("00000000-0000-4000-8000-000000000001", expiration_time), - ): - _persist_form( - sqlite_session_factory, - form_id=form_id, - app_id="app", - node_id="node-1", - expiration_time=expiration, - status=status, - ) - - events = _publish_node_events(_build_timeout_node(expiration_time, status=status)) - - payloads = _sse_payloads(events, InvokeFrom.WEB_APP, app, sqlite_session_factory) - assert [payload["event"] for payload in payloads] == [ - "node_started", - "human_input_form_timeout", - "node_finished", - ] - assert payloads[1]["data"] == { - "form_id": "00000000-0000-4000-8000-000000000001", - "node_id": "node-1", - "node_title": "Human Input", - "expiration_time": 1735689600, - } - - -@pytest.mark.parametrize("timed_out", [False, True]) -@pytest.mark.parametrize("terminal", ["end", "answer"]) -def test_human_input_completion_and_referenced_answer_reach_response_stream( - timed_out: bool, terminal: str, mocker: MockerFixture, app: Flask, sqlite_session_factory: sessionmaker[Session] -): - expiration_time = datetime.datetime(2025, 1, 1) - node = ( - _build_timeout_node(expiration_time, node_id="human") - if timed_out - else _build_node("Approve?", with_inputs=False, node_id="human") - ) - runtime_state = node.runtime_state - runtime_state.variable_pool.add(["sys", "workflow_execution_id"], StringSegment(value="run-1")) - init_params = InitParams(workflow_id="workflow", graph_config={}, run_context={}, call_depth=0) - start = StartNode( - node_id="start", - data=StartNodeData(title="Start", variables=[]), - init_params=init_params, - runtime_state=runtime_state, - ) - if terminal == "answer": - terminal_node = AnswerNode( - node_id="answer", - data=AnswerNodeData(title="Answer", answer="Action: {{#human.__action_id#}}"), - init_params=init_params, - runtime_state=runtime_state, - ) - else: - terminal_node = EndNode( - node_id="end", - data=EndNodeData(title="End", outputs=[]), - init_params=init_params, - runtime_state=runtime_state, - ) - graph = ( - Graph.new() - .add_root(start) - .add_node(node, from_node_id="start") - .add_node(terminal_node, from_node_id=node.id, source_handle="__timeout" if timed_out else "Accept") - .build() - ) - engine = Engine( - graph=graph, - runtime_state=runtime_state, - command_channel=InMemoryChannel(), - workers=1, - ) - bind_execution_id = node.bind_execution_id - - def bind_form(execution_id: str) -> None: - bind_execution_id(execution_id) - _persist_form( - sqlite_session_factory, - form_id=execution_id, - app_id="app", - node_id=node.id, - expiration_time=expiration_time, - status=HumanInputFormStatus.TIMEOUT if timed_out else HumanInputFormStatus.SUBMITTED, - ) - - mocker.patch.object(node, "bind_execution_id", side_effect=bind_form) - - events = _publish_graph_events(iter_dify_graph_engine_events(engine), graph=graph, runtime_state=runtime_state) - payloads = _sse_payloads(events, InvokeFrom.WEB_APP, app, sqlite_session_factory, runtime_state) - - form_event = "human_input_form_timeout" if timed_out else "human_input_form_filled" - lifecycle_events = [payload for payload in payloads if payload["event"] != "message"] - assert [(payload["event"], payload.get("data", {}).get("node_id")) for payload in lifecycle_events] == [ - ("workflow_started", None), - ("node_started", "start"), - ("node_finished", "start"), - ("node_started", node.id), - (form_event, node.id), - ("node_finished", node.id), - ("node_started", terminal), - ("node_finished", terminal), - ("message_end", None), - ("workflow_finished", None), - ] - answer = "".join(payload["answer"] for payload in payloads if payload["event"] == "message") - expected_answer = "Action: " + ("" if timed_out else "Accept") - assert answer == (expected_answer if terminal == "answer" else "") - assert payloads[-1]["data"]["status"] == "succeeded" - - -@pytest.mark.parametrize("form_owner", [None, ("other-app", "node-1"), ("app", "other-node")]) -def test_timeout_rejects_missing_or_unrelated_form( - form_owner: tuple[str, str] | None, sqlite_session_factory: sessionmaker[Session] -): - if form_owner is not None: - app_id, node_id = form_owner - _persist_form( - sqlite_session_factory, - form_id="00000000-0000-4000-8000-000000000001", - app_id=app_id, - node_id=node_id, - expiration_time=datetime.datetime(2025, 1, 1), - status=HumanInputFormStatus.TIMEOUT, - ) - - with pytest.raises(ValueError, match="Cannot resolve timed-out human input form"): - _publish_node_events(_build_timeout_node()) diff --git a/api/tests/unit_tests/core/workflow/test_form_input_serialization_compat.py b/api/tests/unit_tests/core/workflow/nodes/human_input/test_pause_reason.py similarity index 65% rename from api/tests/unit_tests/core/workflow/test_form_input_serialization_compat.py rename to api/tests/unit_tests/core/workflow/nodes/human_input/test_pause_reason.py index 868d453bbd1..f96064a15b4 100644 --- a/api/tests/unit_tests/core/workflow/test_form_input_serialization_compat.py +++ b/api/tests/unit_tests/core/workflow/nodes/human_input/test_pause_reason.py @@ -1,6 +1,8 @@ import json +from datetime import datetime from typing import Any +import pytest from pydantic import TypeAdapter from core.app.entities.task_entities import HumanInputRequiredResponse @@ -8,14 +10,177 @@ from core.entities.execution_extra_content import ( HumanInputContent, HumanInputFormDefinition, ) +from core.repositories.human_input_repository import HumanInputFormRecord +from core.workflow.human_input_policy import FormDisposition, enrich_human_input_pause_reasons +from core.workflow.nodes.human_input.boundary import build_human_input_pause_reason, human_input_container_selector from core.workflow.nodes.human_input.entities import ( FormDefinition, FormInputConfig, HumanInputNodeData, + SelectInputConfig, + StringListSource, ) -from core.workflow.nodes.human_input.enums import ButtonStyle, TimeoutUnit, ValueSourceType -from core.workflow.nodes.human_input.pause_reason import HumanInputRequired +from core.workflow.nodes.human_input.enums import ( + ButtonStyle, + HumanInputFormKind, + HumanInputFormStatus, + TimeoutUnit, + ValueSourceType, +) +from core.workflow.nodes.human_input.pause_reason import DifyHITLEventType, HumanInputRequired from core.workflow.nodes.human_input.session_binding import SessionBinding +from graphon.entities.pause_reason import HitlRequired +from graphon.runtime import RuntimeState, VariablePool +from models.human_input import HumanInputForm + +_HUMAN_INPUT_REASON = {"TYPE": DifyHITLEventType.HUMAN_INPUT_REQUIRED, "form_id": "f1"} + + +@pytest.mark.parametrize( + ("dispositions", "expected_token", "expected_channels"), + [ + ({"f1": FormDisposition(form_token=None, approval_channels=["console", "email"])}, None, ["console", "email"]), + ({"f1": FormDisposition(form_token="tok", approval_channels=[])}, "tok", []), + # form_id absent from the map (no recipient rows) falls back to no token, no channels. + ({}, None, []), + ], +) +def test_enrich_projects_disposition_onto_reason(dispositions, expected_token, expected_channels): + out = enrich_human_input_pause_reasons( + [dict(_HUMAN_INPUT_REASON)], + dispositions_by_form_id=dispositions, + expiration_times_by_form_id={}, + ) + + assert out[0]["form_token"] == expected_token + assert out[0]["approval_channels"] == expected_channels + + +def test_enrich_leaves_non_human_input_reasons_untouched(): + reason = {"TYPE": "something_else", "form_id": "f1"} + + out = enrich_human_input_pause_reasons( + [reason], + dispositions_by_form_id={"f1": FormDisposition(form_token="tok", approval_channels=["email"])}, + expiration_times_by_form_id={}, + ) + + assert out[0] == reason + assert "form_token" not in out[0] + assert "approval_channels" not in out[0] + + +def test_pause_reason_payload_carries_approval_channels_through_factory(): + # from_response_data maps fields by hand; this guards approval_channels/form_token + # (the fields this feature added) against being dropped in that mapping. + from core.app.entities.task_entities import ( + HumanInputRequiredPauseReasonPayload, + HumanInputRequiredResponse, + ) + + data = HumanInputRequiredResponse.Data( + form_id="f", + node_id="n", + node_title="t", + form_content="c", + expiration_time=123, + form_token=None, + approval_channels=["console"], + ) + payload = HumanInputRequiredPauseReasonPayload.from_response_data(data) + + assert payload.approval_channels == ["console"] + assert payload.form_token is None + + +def test_build_pause_reason_keeps_constant_options_from_child_form(): + definition = FormDefinition( + form_content="Choose", + rendered_content="Choose", + inputs=[ + SelectInputConfig( + output_variable_name="decision", + option_source=StringListSource(type=ValueSourceType.CONSTANT, value=[]), + ) + ], + node_title="Choose", + expiration_time=datetime(2026, 1, 2), + ) + record = HumanInputFormRecord.from_models( + HumanInputForm( + id="form-123", + tenant_id="tenant", + app_id="app", + workflow_run_id="run", + node_id="human-input", + form_kind=HumanInputFormKind.RUNTIME, + form_definition=definition.model_dump_json(), + rendered_content=definition.rendered_content, + created_at=datetime(2026, 1, 1), + expiration_time=definition.expiration_time, + status=HumanInputFormStatus.WAITING, + ), + None, + ) + parent_pool = VariablePool() + parent_pool.add(("start", "options"), ["wrong"]) + + reason = build_human_input_pause_reason( + reason=HitlRequired(session_id="form-123", node_id="human-input", node_title="Choose"), + record=record, + variable_pool=parent_pool, + ) + + assert isinstance(reason.inputs[0], SelectInputConfig) + assert reason.inputs[0].option_source.type == ValueSourceType.CONSTANT + assert reason.inputs[0].option_source.value == [] + + +def test_child_forms_keep_distinct_identity_when_projected_to_the_same_visible_tool(): + definition = FormDefinition( + form_content="Approval", + rendered_content="Approval", + node_title="Approval", + expiration_time=datetime(2026, 1, 2), + ) + state = RuntimeState(workflow_id="workflow", variable_pool=VariablePool(), start_at=0) + reasons = [ + HitlRequired(session_id=form_id, node_id="source-human", node_title="Approval") + for form_id in ("form-a", "form-b") + ] + for reason in reasons: + state.variable_pool.add(human_input_container_selector(reason.session_id), "visible-tool") + restored = RuntimeState.from_snapshot(state.dumps()) + + enriched = [ + build_human_input_pause_reason( + reason=reason, + record=HumanInputFormRecord.from_models( + HumanInputForm( + id=reason.session_id, + tenant_id="tenant", + app_id="app", + workflow_run_id="run", + node_id="source-human", + form_kind=HumanInputFormKind.RUNTIME, + form_definition=definition.model_dump_json(), + rendered_content=reason.session_id, + created_at=datetime(2026, 1, 1), + expiration_time=definition.expiration_time, + status=HumanInputFormStatus.WAITING, + ), + None, + ), + variable_pool=restored.variable_pool, + ) + for reason in reasons + ] + + assert [(reason.form_id, reason.node_id) for reason in enriched] == [ + ("form-a", "visible-tool"), + ("form-b", "visible-tool"), + ] + assert all(reason.node_id == "source-human" for reason in reasons) def _legacy_form_input_payloads() -> list[dict[str, Any]]: diff --git a/api/tests/unit_tests/extensions/otel/test_retrieval_tracing.py b/api/tests/unit_tests/core/workflow/nodes/knowledge_retrieval/test_retrieval.py similarity index 87% rename from api/tests/unit_tests/extensions/otel/test_retrieval_tracing.py rename to api/tests/unit_tests/core/workflow/nodes/knowledge_retrieval/test_retrieval.py index bdd418e47e4..f9ad0c98e3a 100644 --- a/api/tests/unit_tests/extensions/otel/test_retrieval_tracing.py +++ b/api/tests/unit_tests/core/workflow/nodes/knowledge_retrieval/test_retrieval.py @@ -11,6 +11,15 @@ from core.rag.rerank.rerank_type import RerankMode from core.rag.retrieval.dataset_retrieval import DatasetRetrieval from core.workflow.nodes.knowledge_retrieval.retrieval import KnowledgeRetrievalRequest from models.dataset import Dataset +from tests.unit_tests.extensions.otel.conftest import memory_span_exporter as memory_span_exporter # noqa: PLC0414 +from tests.unit_tests.extensions.otel.conftest import mock_account_user as mock_account_user # noqa: PLC0414 +from tests.unit_tests.extensions.otel.conftest import mock_app_model as mock_app_model # noqa: PLC0414 +from tests.unit_tests.extensions.otel.conftest import mock_end_user as mock_end_user # noqa: PLC0414 +from tests.unit_tests.extensions.otel.conftest import mock_workflow_runner as mock_workflow_runner # noqa: PLC0414 +from tests.unit_tests.extensions.otel.conftest import reset_handler_instances as reset_handler_instances # noqa: PLC0414 +from tests.unit_tests.extensions.otel.conftest import ( + tracer_provider_with_memory_exporter as tracer_provider_with_memory_exporter, # noqa: PLC0414 +) @pytest.fixture(autouse=True) diff --git a/api/tests/unit_tests/core/workflow/nodes/llm/test_llm_node_execution_behavior.py b/api/tests/unit_tests/core/workflow/nodes/llm/test_llm_node_execution_behavior.py deleted file mode 100644 index 5270eebc718..00000000000 --- a/api/tests/unit_tests/core/workflow/nodes/llm/test_llm_node_execution_behavior.py +++ /dev/null @@ -1,301 +0,0 @@ -"""Unit coverage for LLM node execution with model boundaries replaced by test doubles.""" - -import json -import time -import uuid -from collections.abc import Generator -from unittest.mock import MagicMock, patch - -import pytest - -from core.app.entities.app_invoke_entities import InvokeFrom, UserFrom -from core.app.llm.model_access import DifyCredentialsProvider, DifyModelFactory -from core.helper.ssrf_proxy import graphon_ssrf_proxy -from core.llm_generator.output_parser.structured_output import _parse_structured_output -from core.workflow.node_runtime import ( - DifyPreparedLLM, - DifyPromptMessageSerializer, - build_dify_llm_file_saver, - resolve_dify_run_context, -) -from core.workflow.system_variables import build_system_variables -from extensions.ext_database import db -from graphon.enums import WorkflowNodeExecutionStatus -from graphon.model_runtime.entities.message_entities import PromptMessage -from graphon.node_events import StreamCompletedEvent -from graphon.nodes.llm.entities import LLMNodeData -from graphon.nodes.llm.node import LLMNode -from graphon.runtime import RuntimeState, VariablePool -from tests.unit_tests.core.model_fixtures import make_model_config, make_model_instance -from tests.workflow_test_utils import build_test_graph_init_params - -"""FOR MOCK FIXTURES, DO NOT REMOVE""" - - -def init_llm_node(config: dict[str, object]) -> LLMNode: - graph_config = { - "edges": [ - { - "id": "start-source-next-target", - "source": "start", - "target": "llm", - }, - ], - "nodes": [{"data": {"type": "start", "title": "Start"}, "id": "start"}, config], - } - - # Use proper UUIDs for database compatibility - tenant_id = "9d2074fc-6f86-45a9-b09d-6ecc63b9056b" - app_id = "9d2074fc-6f86-45a9-b09d-6ecc63b9056c" - workflow_id = "9d2074fc-6f86-45a9-b09d-6ecc63b9056d" - user_id = "9d2074fc-6f86-45a9-b09d-6ecc63b9056e" - - init_params = build_test_graph_init_params( - workflow_id=workflow_id, - graph_config=graph_config, - tenant_id=tenant_id, - app_id=app_id, - user_id=user_id, - user_from=UserFrom.ACCOUNT, - invoke_from=InvokeFrom.DEBUGGER, - call_depth=0, - ) - - # construct variable pool - variable_pool = VariablePool.from_bootstrap( - system_variables=build_system_variables( - user_id="aaa", - app_id=app_id, - workflow_id=workflow_id, - files=[], - query="what's the weather today?", - conversation_id="abababa", - ), - user_inputs={}, - environment_variables=[], - conversation_variables=[], - ) - variable_pool.add(["abc", "output"], "sunny") - - graph_runtime_state = RuntimeState( - workflow_id=workflow_id, variable_pool=variable_pool, start_at=time.perf_counter() - ) - run_context = resolve_dify_run_context(init_params.run_context) - - node = LLMNode( - node_id=str(uuid.uuid4()), - data=LLMNodeData.model_validate(config["data"]), - init_params=init_params, - runtime_state=graph_runtime_state, - credentials_provider=DifyCredentialsProvider(run_context=run_context), - model_factory=DifyModelFactory(run_context=run_context), - model_instance=DifyPreparedLLM(make_model_instance(provider="openai", model="gpt-3.5-turbo")), - llm_file_saver=build_dify_llm_file_saver(run_context=run_context, http_client=graphon_ssrf_proxy), - prompt_message_serializer=DifyPromptMessageSerializer(), - http_client=graphon_ssrf_proxy, - ) - - return node - - -def _mock_db_session_close(monkeypatch: pytest.MonkeyPatch) -> None: - monkeypatch.setattr(db.session, "close", MagicMock()) - - -def test_execute_llm(monkeypatch: pytest.MonkeyPatch) -> None: - node = init_llm_node( - config={ - "id": "llm", - "data": { - "title": "123", - "type": "llm", - "model": { - "provider": "openai", - "name": "gpt-3.5-turbo", - "mode": "chat", - "completion_params": {}, - }, - "prompt_template": [ - { - "role": "system", - "text": "you are a helpful assistant.\ntoday's weather is {{#abc.output#}}.", - }, - {"role": "user", "text": "{{#sys.query#}}"}, - ], - "memory": None, - "context": {"enabled": False}, - "vision": {"enabled": False}, - }, - }, - ) - - _mock_db_session_close(monkeypatch) - - def build_model_instance() -> DifyPreparedLLM: - from decimal import Decimal - from unittest.mock import MagicMock - - from graphon.model_runtime.entities.llm_entities import LLMResult, LLMUsage - from graphon.model_runtime.entities.message_entities import AssistantPromptMessage - - model_instance = make_model_instance(provider="openai", model="gpt-3.5-turbo") - schema = make_model_config(provider="openai", model="gpt-3.5-turbo", mode="chat").model_schema - monkeypatch.setattr(model_instance, "get_model_schema", MagicMock(return_value=schema)) - mock_usage = LLMUsage( - prompt_tokens=30, - prompt_unit_price=Decimal("0.001"), - prompt_price_unit=Decimal(1000), - prompt_price=Decimal("0.00003"), - completion_tokens=20, - completion_unit_price=Decimal("0.002"), - completion_price_unit=Decimal(1000), - completion_price=Decimal("0.00004"), - total_tokens=50, - total_price=Decimal("0.00007"), - currency="USD", - latency=0.5, - ) - mock_message = AssistantPromptMessage(content="Test response from mock") - mock_llm_result = LLMResult( - model="gpt-3.5-turbo", - prompt_messages=[], - message=mock_message, - usage=mock_usage, - ) - monkeypatch.setattr(model_instance, "invoke_llm", MagicMock(return_value=mock_llm_result)) - return DifyPreparedLLM(model_instance) - - # Mock fetch_prompt_messages to avoid database calls - def mock_fetch_prompt_messages_1(**_kwargs: object) -> tuple[list[PromptMessage], list[str]]: - from graphon.model_runtime.entities.message_entities import SystemPromptMessage, UserPromptMessage - - return [ - SystemPromptMessage(content="you are a helpful assistant. today's weather is sunny."), - UserPromptMessage(content="what's the weather today?"), - ], [] - - node._model_instance = build_model_instance() - - with patch("graphon.nodes.llm.llm_utils.fetch_prompt_messages", mock_fetch_prompt_messages_1): - # execute node - result = node._run() - assert isinstance(result, Generator) - - for item in result: - if isinstance(item, StreamCompletedEvent): - assert item.node_run_result.status == WorkflowNodeExecutionStatus.SUCCEEDED, ( - f"{item.node_run_result.error_type}: {item.node_run_result.error}" - ) - assert item.node_run_result.process_data is not None - assert item.node_run_result.outputs is not None - assert item.node_run_result.outputs.get("text") is not None - assert item.node_run_result.outputs.get("usage", {})["total_tokens"] > 0 - - -def test_execute_llm_with_jinja2(monkeypatch: pytest.MonkeyPatch) -> None: - """ - Test execute LLM node with jinja2 - """ - node = init_llm_node( - config={ - "id": "llm", - "data": { - "title": "123", - "type": "llm", - "model": {"provider": "openai", "name": "gpt-3.5-turbo", "mode": "chat", "completion_params": {}}, - "prompt_config": { - "jinja2_variables": [ - {"variable": "sys_query", "value_selector": ["sys", "query"]}, - {"variable": "output", "value_selector": ["abc", "output"]}, - ] - }, - "prompt_template": [ - { - "role": "system", - "text": "you are a helpful assistant.\ntoday's weather is {{#abc.output#}}", - "jinja2_text": "you are a helpful assistant.\ntoday's weather is {{output}}.", - "edition_type": "jinja2", - }, - { - "role": "user", - "text": "{{#sys.query#}}", - "jinja2_text": "{{sys_query}}", - "edition_type": "basic", - }, - ], - "memory": None, - "context": {"enabled": False}, - "vision": {"enabled": False}, - }, - }, - ) - - _mock_db_session_close(monkeypatch) - - def build_model_instance() -> DifyPreparedLLM: - from decimal import Decimal - from unittest.mock import MagicMock - - from graphon.model_runtime.entities.llm_entities import LLMResult, LLMUsage - from graphon.model_runtime.entities.message_entities import AssistantPromptMessage - - model_instance = make_model_instance(provider="openai", model="gpt-3.5-turbo") - schema = make_model_config(provider="openai", model="gpt-3.5-turbo", mode="chat").model_schema - monkeypatch.setattr(model_instance, "get_model_schema", MagicMock(return_value=schema)) - mock_usage = LLMUsage( - prompt_tokens=30, - prompt_unit_price=Decimal("0.001"), - prompt_price_unit=Decimal(1000), - prompt_price=Decimal("0.00003"), - completion_tokens=20, - completion_unit_price=Decimal("0.002"), - completion_price_unit=Decimal(1000), - completion_price=Decimal("0.00004"), - total_tokens=50, - total_price=Decimal("0.00007"), - currency="USD", - latency=0.5, - ) - mock_message = AssistantPromptMessage(content="Test response: sunny weather and what's the weather today?") - mock_llm_result = LLMResult( - model="gpt-3.5-turbo", - prompt_messages=[], - message=mock_message, - usage=mock_usage, - ) - monkeypatch.setattr(model_instance, "invoke_llm", MagicMock(return_value=mock_llm_result)) - return DifyPreparedLLM(model_instance) - - # Mock fetch_prompt_messages to avoid database calls - def mock_fetch_prompt_messages_2(**_kwargs: object) -> tuple[list[PromptMessage], list[str]]: - from graphon.model_runtime.entities.message_entities import SystemPromptMessage, UserPromptMessage - - return [ - SystemPromptMessage(content="you are a helpful assistant. today's weather is sunny."), - UserPromptMessage(content="what's the weather today?"), - ], [] - - node._model_instance = build_model_instance() - - with patch("graphon.nodes.llm.llm_utils.fetch_prompt_messages", mock_fetch_prompt_messages_2): - # execute node - result = node._run() - - for item in result: - if isinstance(item, StreamCompletedEvent): - assert item.node_run_result.status == WorkflowNodeExecutionStatus.SUCCEEDED - assert item.node_run_result.process_data is not None - assert "sunny" in json.dumps(item.node_run_result.process_data) - assert "what's the weather today?" in json.dumps(item.node_run_result.process_data) - - -def test_extract_json() -> None: - llm_texts = [ - '\n\n{"name": "test", "age": 123', # resoning model (deepseek-r1) - '{"name":"test","age":123}', # json schema model (gpt-4o) - '{\n "name": "test",\n "age": 123\n}', # small model (llama-3.2-1b) - '```json\n{"name": "test", "age": 123}\n```', # json markdown (deepseek-chat) - '{"name":"test",age:123}', # without quotes (qwen-2.5-0.5b) - ] - result = {"name": "test", "age": 123} - assert all(_parse_structured_output(item) == result for item in llm_texts) diff --git a/api/tests/unit_tests/core/workflow/nodes/parameter_extractor/test_parameter_extractor_execution_behavior.py b/api/tests/unit_tests/core/workflow/nodes/parameter_extractor/test_parameter_extractor_execution_behavior.py deleted file mode 100644 index eb3a8e84c89..00000000000 --- a/api/tests/unit_tests/core/workflow/nodes/parameter_extractor/test_parameter_extractor_execution_behavior.py +++ /dev/null @@ -1,433 +0,0 @@ -"""Unit coverage for parameter extraction with the plugin-model boundary replaced.""" - -import os -import time -import uuid -from typing import override -from unittest.mock import MagicMock - -import pytest - -from core.app.entities.app_invoke_entities import InvokeFrom, UserFrom -from core.app.llm.model_access import DifyCredentialsProvider, DifyModelFactory -from core.workflow.node_runtime import DifyPreparedLLM, DifyPromptMessageSerializer, resolve_dify_run_context -from core.workflow.system_variables import build_system_variables -from extensions.ext_database import db -from graphon.enums import WorkflowNodeExecutionStatus -from graphon.model_runtime.entities import AssistantPromptMessage, UserPromptMessage -from graphon.model_runtime.entities.message_entities import PromptMessage -from graphon.model_runtime.memory.prompt_message_memory import PromptMessageMemory -from graphon.nodes.parameter_extractor.entities import ParameterExtractorNodeData -from graphon.nodes.parameter_extractor.parameter_extractor_node import ParameterExtractorNode -from graphon.runtime import RuntimeState, VariablePool -from tests.unit_tests.core.model_fixtures import make_model_instance -from tests.unit_tests.core.workflow.nodes.parameter_extractor.fixtures.model import get_mocked_fetch_model_instance -from tests.workflow_test_utils import build_test_graph_init_params - -pytest_plugins = ("tests.unit_tests.core.workflow.nodes.parameter_extractor.fixtures.plugin_daemon",) - - -def get_mocked_fetch_memory(memory_text: str) -> MagicMock: - class MemoryMock(PromptMessageMemory): - @override - def get_history_prompt_messages( - self, - max_token_limit: int = 2000, - message_limit: int | None = None, - ) -> list[PromptMessage]: - return [UserPromptMessage(content=memory_text), AssistantPromptMessage(content="mocked answer")] - - return MagicMock(return_value=MemoryMock()) - - -def init_parameter_extractor_node( - config: dict[str, object], memory: PromptMessageMemory | None = None -) -> ParameterExtractorNode: - graph_config = { - "edges": [ - { - "id": "start-source-next-target", - "source": "start", - "target": "llm", - }, - ], - "nodes": [{"data": {"type": "start", "title": "Start"}, "id": "start"}, config], - } - - init_params = build_test_graph_init_params( - workflow_id="1", - graph_config=graph_config, - tenant_id="1", - app_id="1", - user_id="1", - user_from=UserFrom.ACCOUNT, - invoke_from=InvokeFrom.DEBUGGER, - call_depth=0, - ) - - # construct variable pool - variable_pool = VariablePool.from_bootstrap( - system_variables=build_system_variables( - user_id="aaa", files=[], query="what's the weather in SF", conversation_id="abababa" - ), - user_inputs={}, - environment_variables=[], - conversation_variables=[], - ) - variable_pool.add(["a", "args1"], 1) - variable_pool.add(["a", "args2"], 2) - - graph_runtime_state = RuntimeState( - workflow_id="test-workflow", variable_pool=variable_pool, start_at=time.perf_counter() - ) - - run_context = resolve_dify_run_context(init_params.run_context) - node = ParameterExtractorNode( - node_id=str(uuid.uuid4()), - data=ParameterExtractorNodeData.model_validate(config["data"]), - init_params=init_params, - runtime_state=graph_runtime_state, - credentials_provider=DifyCredentialsProvider(run_context=run_context), - model_factory=DifyModelFactory(run_context=run_context), - model_instance=DifyPreparedLLM(make_model_instance(provider="openai", model="gpt-3.5-turbo")), - memory=memory, - prompt_message_serializer=DifyPromptMessageSerializer(), - ) - return node - - -def _mock_db_session_close(monkeypatch: pytest.MonkeyPatch) -> None: - monkeypatch.setattr(db.session, "close", MagicMock()) - - -@pytest.mark.usefixtures("setup_model_mock") -def test_function_calling_parameter_extractor(monkeypatch: pytest.MonkeyPatch) -> None: - """ - Test function calling for parameter extractor. - """ - node = init_parameter_extractor_node( - config={ - "id": "llm", - "data": { - "title": "123", - "type": "parameter-extractor", - "model": { - "provider": "langgenius/openai/openai", - "name": "gpt-3.5-turbo", - "mode": "chat", - "completion_params": {}, - }, - "query": ["sys", "query"], - "parameters": [{"name": "location", "type": "string", "description": "location", "required": True}], - "instruction": "", - "reasoning_mode": "function_call", - "memory": None, - }, - } - ) - - node._model_instance = get_mocked_fetch_model_instance( - provider="langgenius/openai/openai", - model="gpt-3.5-turbo", - mode="chat", - credentials={"openai_api_key": os.environ.get("OPENAI_API_KEY")}, - )() - _mock_db_session_close(monkeypatch) - - result = node._run() - - assert result.status == WorkflowNodeExecutionStatus.SUCCEEDED - assert result.outputs is not None - assert result.outputs.get("location") == "kawaii" - assert result.outputs.get("__reason") is None - - -@pytest.mark.usefixtures("setup_model_mock") -def test_instructions(monkeypatch: pytest.MonkeyPatch) -> None: - """ - Test chat parameter extractor. - """ - node = init_parameter_extractor_node( - config={ - "id": "llm", - "data": { - "title": "123", - "type": "parameter-extractor", - "model": { - "provider": "langgenius/openai/openai", - "name": "gpt-3.5-turbo", - "mode": "chat", - "completion_params": {}, - }, - "query": ["sys", "query"], - "parameters": [{"name": "location", "type": "string", "description": "location", "required": True}], - "reasoning_mode": "function_call", - "instruction": "{{#sys.query#}}", - "memory": None, - }, - }, - ) - - node._model_instance = get_mocked_fetch_model_instance( - provider="langgenius/openai/openai", - model="gpt-3.5-turbo", - mode="chat", - credentials={"openai_api_key": os.environ.get("OPENAI_API_KEY")}, - )() - _mock_db_session_close(monkeypatch) - - result = node._run() - - assert result.status == WorkflowNodeExecutionStatus.SUCCEEDED - assert result.outputs is not None - assert result.outputs.get("location") == "kawaii" - assert result.outputs.get("__reason") is None - - process_data = result.process_data - - assert process_data is not None - process_data.get("prompts") - - for prompt in process_data.get("prompts", []): - if prompt.get("role") == "system": - assert "what's the weather in SF" in prompt.get("text") - - -@pytest.mark.usefixtures("setup_model_mock") -def test_chat_parameter_extractor(monkeypatch: pytest.MonkeyPatch) -> None: - """ - Test chat parameter extractor. - """ - node = init_parameter_extractor_node( - config={ - "id": "llm", - "data": { - "title": "123", - "type": "parameter-extractor", - "model": { - "provider": "langgenius/openai/openai", - "name": "gpt-3.5-turbo", - "mode": "chat", - "completion_params": {}, - }, - "query": ["sys", "query"], - "parameters": [{"name": "location", "type": "string", "description": "location", "required": True}], - "reasoning_mode": "prompt", - "instruction": "", - "memory": None, - }, - }, - ) - - node._model_instance = get_mocked_fetch_model_instance( - provider="langgenius/openai/openai", - model="gpt-3.5-turbo", - mode="chat", - credentials={"openai_api_key": os.environ.get("OPENAI_API_KEY")}, - )() - _mock_db_session_close(monkeypatch) - - result = node._run() - - assert result.status == WorkflowNodeExecutionStatus.SUCCEEDED - assert result.outputs is not None - assert result.outputs.get("location") == "" - assert ( - result.outputs.get("__reason") - == "Failed to extract result from function call or text response, using empty result." - ) - assert result.process_data is not None - prompts = result.process_data.get("prompts", []) - - for prompt in prompts: - if prompt.get("role") == "user": - if "" in prompt.get("text"): - assert '\n{"type": "object"' in prompt.get("text") - - -@pytest.mark.usefixtures("setup_model_mock") -def test_completion_parameter_extractor(monkeypatch: pytest.MonkeyPatch) -> None: - """ - Test completion parameter extractor. - """ - node = init_parameter_extractor_node( - config={ - "id": "llm", - "data": { - "title": "123", - "type": "parameter-extractor", - "model": { - "provider": "langgenius/openai/openai", - "name": "gpt-3.5-turbo-instruct", - "mode": "completion", - "completion_params": {}, - }, - "query": ["sys", "query"], - "parameters": [{"name": "location", "type": "string", "description": "location", "required": True}], - "reasoning_mode": "prompt", - "instruction": "{{#sys.query#}}", - "memory": None, - }, - }, - ) - - node._model_instance = get_mocked_fetch_model_instance( - provider="langgenius/openai/openai", - model="gpt-3.5-turbo-instruct", - mode="completion", - credentials={"openai_api_key": os.environ.get("OPENAI_API_KEY")}, - )() - _mock_db_session_close(monkeypatch) - - result = node._run() - - assert result.status == WorkflowNodeExecutionStatus.SUCCEEDED - assert result.outputs is not None - assert result.outputs.get("location") == "" - assert ( - result.outputs.get("__reason") - == "Failed to extract result from function call or text response, using empty result." - ) - assert result.process_data is not None - assert len(result.process_data.get("prompts", [])) == 1 - assert "SF" in result.process_data.get("prompts", [])[0].get("text") - - -def test_extract_json_response() -> None: - """ - Test extract json response. - """ - - node = init_parameter_extractor_node( - config={ - "id": "llm", - "data": { - "title": "123", - "type": "parameter-extractor", - "model": { - "provider": "langgenius/openai/openai", - "name": "gpt-3.5-turbo-instruct", - "mode": "completion", - "completion_params": {}, - }, - "query": ["sys", "query"], - "parameters": [{"name": "location", "type": "string", "description": "location", "required": True}], - "reasoning_mode": "prompt", - "instruction": "{{#sys.query#}}", - "memory": None, - }, - }, - ) - - result = node._extract_complete_json_response(""" - uwu{ovo} - { - "location": "kawaii" - } - hello world. - """) - - assert result is not None - assert result["location"] == "kawaii" - - -def test_extract_json_from_tool_call() -> None: - """ - Test extract json response. - """ - - node = init_parameter_extractor_node( - config={ - "id": "llm", - "data": { - "title": "123", - "type": "parameter-extractor", - "model": { - "provider": "langgenius/openai/openai", - "name": "gpt-3.5-turbo-instruct", - "mode": "completion", - "completion_params": {}, - }, - "query": ["sys", "query"], - "parameters": [{"name": "location", "type": "string", "description": "location", "required": True}], - "reasoning_mode": "prompt", - "instruction": "{{#sys.query#}}", - "memory": None, - }, - }, - ) - - result = node._extract_json_from_tool_call( - AssistantPromptMessage.ToolCall( - id="llm", - type="parameter-extractor", - function=AssistantPromptMessage.ToolCall.ToolCallFunction( - name="foo", arguments="""{"location":"kawaii"}{"location": 1}""" - ), - ) - ) - - assert result is not None - assert result["location"] == "kawaii" - - -@pytest.mark.usefixtures("setup_model_mock") -def test_chat_parameter_extractor_with_memory(monkeypatch: pytest.MonkeyPatch) -> None: - """ - Test chat parameter extractor with memory. - """ - node = init_parameter_extractor_node( - config={ - "id": "llm", - "data": { - "title": "123", - "type": "parameter-extractor", - "model": { - "provider": "langgenius/openai/openai", - "name": "gpt-3.5-turbo", - "mode": "chat", - "completion_params": {}, - }, - "query": ["sys", "query"], - "parameters": [{"name": "location", "type": "string", "description": "location", "required": True}], - "reasoning_mode": "prompt", - "instruction": "", - "memory": {"window": {"enabled": True, "size": 50}}, - }, - }, - memory=get_mocked_fetch_memory("customized memory")(), - ) - - node._model_instance = get_mocked_fetch_model_instance( - provider="langgenius/openai/openai", - model="gpt-3.5-turbo", - mode="chat", - credentials={"openai_api_key": os.environ.get("OPENAI_API_KEY")}, - )() - _mock_db_session_close(monkeypatch) - - result = node._run() - - assert result.status == WorkflowNodeExecutionStatus.SUCCEEDED - assert result.outputs is not None - assert result.outputs.get("location") == "" - assert ( - result.outputs.get("__reason") - == "Failed to extract result from function call or text response, using empty result." - ) - assert result.process_data is not None - prompts = result.process_data.get("prompts", []) - - latest_role = None - for prompt in prompts: - if prompt.get("role") == "user": - if "" in prompt.get("text"): - assert '\n{"type": "object"' in prompt.get("text") - elif prompt.get("role") == "system": - assert "customized memory" in prompt.get("text") - - if latest_role is not None: - assert latest_role != prompt.get("role") - - if prompt.get("role") in {"user", "assistant"}: - latest_role = prompt.get("role") diff --git a/api/tests/unit_tests/core/workflow/nodes/tool/test_tool_node_invocation_behavior.py b/api/tests/unit_tests/core/workflow/nodes/tool/test_tool_node_invocation_behavior.py deleted file mode 100644 index eb642540748..00000000000 --- a/api/tests/unit_tests/core/workflow/nodes/tool/test_tool_node_invocation_behavior.py +++ /dev/null @@ -1,145 +0,0 @@ -"""Unit coverage for tool-node invocation with the tool boundary patched.""" - -import time -import uuid -from unittest.mock import patch - -from core.app.entities.app_invoke_entities import InvokeFrom, UserFrom -from core.tools.utils.configuration import ToolParameterConfigurationManager -from core.workflow.node_factory import DifyNodeFactory -from core.workflow.node_runtime import DifyToolFileManager, DifyToolNodeRuntime -from core.workflow.system_variables import build_system_variables -from graphon.enums import WorkflowNodeExecutionStatus -from graphon.graph import Graph -from graphon.node_events import StreamCompletedEvent -from graphon.nodes.tool.entities import ToolNodeData -from graphon.nodes.tool.tool_node import ToolNode -from graphon.runtime import RuntimeState, VariablePool -from tests.workflow_test_utils import build_test_graph_init_params - - -def init_tool_node(config: dict[str, object]) -> ToolNode: - graph_config = { - "edges": [ - { - "id": "start-source-next-target", - "source": "start", - "target": "1", - }, - ], - "nodes": [{"data": {"type": "start", "title": "Start"}, "id": "start"}, config], - } - - init_params = build_test_graph_init_params( - workflow_id="1", - graph_config=graph_config, - tenant_id="1", - app_id="1", - user_id="1", - user_from=UserFrom.ACCOUNT, - invoke_from=InvokeFrom.DEBUGGER, - call_depth=0, - ) - - # construct variable pool - variable_pool = VariablePool.from_bootstrap( - system_variables=build_system_variables(user_id="aaa", files=[]), - user_inputs={}, - environment_variables=[], - conversation_variables=[], - ) - - graph_runtime_state = RuntimeState( - workflow_id="test-workflow", variable_pool=variable_pool, start_at=time.perf_counter() - ) - - # Create node factory - node_factory = DifyNodeFactory( - init_params=init_params, - runtime_state=graph_runtime_state, - ) - - Graph.init(graph_config=graph_config, node_factory=node_factory, root_node_id="start") - - tool_file_manager = DifyToolFileManager(init_params.run_context) - - node = ToolNode( - node_id=str(uuid.uuid4()), - data=ToolNodeData.model_validate(config["data"]), - init_params=init_params, - runtime_state=graph_runtime_state, - tool_file_manager=tool_file_manager, - runtime=DifyToolNodeRuntime(init_params.run_context), - ) - node.bind_execution_id(str(uuid.uuid4())) - return node - - -def test_tool_variable_invoke() -> None: - node = init_tool_node( - config={ - "id": "1", - "data": { - "type": "tool", - "title": "a", - "desc": "a", - "provider_id": "time", - "provider_type": "builtin", - "provider_name": "time", - "tool_name": "current_time", - "tool_label": "current_time", - "tool_configurations": {}, - "tool_parameters": {}, - }, - } - ) - - with patch.object( - ToolParameterConfigurationManager, - "decrypt_tool_parameters", - return_value={"format": "%Y-%m-%d %H:%M:%S"}, - ): - node.runtime_state.variable_pool.add(["1", "args1"], "1+1") - - # execute node - result = node._run() - for item in result: - if isinstance(item, StreamCompletedEvent): - assert item.node_run_result.status == WorkflowNodeExecutionStatus.SUCCEEDED - assert item.node_run_result.outputs is not None - assert item.node_run_result.outputs.get("text") is not None - - -def test_tool_mixed_invoke() -> None: - node = init_tool_node( - config={ - "id": "1", - "data": { - "type": "tool", - "title": "a", - "desc": "a", - "provider_id": "time", - "provider_type": "builtin", - "provider_name": "time", - "tool_name": "current_time", - "tool_label": "current_time", - "tool_configurations": { - "format": "%Y-%m-%d %H:%M:%S", - }, - "tool_parameters": {}, - }, - } - ) - - with patch.object( - ToolParameterConfigurationManager, - "decrypt_tool_parameters", - return_value={"format": "%Y-%m-%d %H:%M:%S"}, - ): - # execute node - result = node._run() - for item in result: - if isinstance(item, StreamCompletedEvent): - assert item.node_run_result.status == WorkflowNodeExecutionStatus.SUCCEEDED - assert item.node_run_result.outputs is not None - assert item.node_run_result.outputs.get("text") is not None diff --git a/api/tests/unit_tests/core/workflow/nodes/tool/test_tool_node_runtime.py b/api/tests/unit_tests/core/workflow/nodes/tool/test_tool_node_runtime.py deleted file mode 100644 index 4b41e28ec8e..00000000000 --- a/api/tests/unit_tests/core/workflow/nodes/tool/test_tool_node_runtime.py +++ /dev/null @@ -1,383 +0,0 @@ -from __future__ import annotations - -import sys -import types -from types import SimpleNamespace -from unittest.mock import MagicMock, patch - -import pytest -from sqlalchemy import Engine -from sqlalchemy.orm import sessionmaker - -from core.callback_handler.workflow_tool_callback_handler import DifyWorkflowCallbackHandler -from core.plugin.impl.exc import PluginDaemonClientSideError, PluginInvokeError -from core.tools.entities.tool_entities import ToolInvokeMessage -from core.tools.entities.tool_entities import ToolProviderType as CoreToolProviderType -from core.tools.errors import ToolInvokeError -from core.tools.tool_engine import ToolEngine -from core.tools.tool_manager import ToolManager -from core.tools.utils.message_transformer import ToolFileMessageTransformer -from core.workflow.node_runtime import DifyToolNodeRuntime -from core.workflow.system_variables import build_system_variables -from graphon.model_runtime.entities.llm_entities import LLMUsage -from graphon.nodes.tool.entities import ToolNodeData, ToolProviderType -from graphon.nodes.tool.exc import ToolRuntimeInvocationError -from graphon.nodes.tool_runtime_entities import ToolRuntimeHandle, ToolRuntimeMessage -from graphon.runtime import VariablePool -from tests.workflow_test_utils import build_test_graph_init_params, build_test_variable_pool - - -@pytest.fixture -def runtime(monkeypatch: pytest.MonkeyPatch, sqlite_engine: Engine) -> DifyToolNodeRuntime: - module_name = "core.ops.ops_trace_manager" - if module_name not in sys.modules: - ops_stub = types.ModuleType(module_name) - ops_stub.TraceQueueManager = object # pragma: no cover - stub attribute - ops_stub.TraceTask = object # pragma: no cover - stub attribute - monkeypatch.setitem(sys.modules, module_name, ops_stub) - - init_params = build_test_graph_init_params( - workflow_id="workflow-id", - graph_config={"nodes": [], "edges": []}, - tenant_id="tenant-id", - app_id="app-id", - user_id="user-id", - user_from="account", - invoke_from="debugger", - call_depth=0, - ) - session_maker = sessionmaker(sqlite_engine, expire_on_commit=False) - return DifyToolNodeRuntime(init_params.run_context, session_maker=session_maker) - - -def _build_tool_node_data() -> ToolNodeData: - return ToolNodeData.model_validate( - { - "type": "tool", - "title": "Tool", - "provider_id": "provider", - "provider_type": ToolProviderType.BUILT_IN, - "provider_name": "provider", - "tool_name": "lookup", - "tool_label": "Lookup", - "tool_configurations": {}, - "tool_parameters": {}, - } - ) - - -def test_invoke_creates_callback_and_converts_messages(runtime: DifyToolNodeRuntime) -> None: - core_message = ToolInvokeMessage( - type=ToolInvokeMessage.MessageType.LINK, - message=ToolInvokeMessage.TextMessage(text="https://dify.ai"), - meta=None, - ) - variable_pool: VariablePool = build_test_variable_pool( - variables=build_system_variables(conversation_id="conversation-id") - ) - workflow_tool = MagicMock() - - with ( - patch.object(ToolManager, "get_workflow_tool_runtime", return_value=workflow_tool), - patch.object(ToolEngine, "generic_invoke", return_value=iter([core_message])) as generic_invoke_mock, - patch.object( - ToolFileMessageTransformer, - "transform_tool_invoke_messages", - side_effect=lambda *, messages, **_: messages, - ) as transform_tool_messages, - ): - tool_runtime = runtime.get_runtime( - node_id="node-id", - node_data=_build_tool_node_data(), - variable_pool=variable_pool, - ) - messages = list( - runtime.invoke( - tool_runtime=tool_runtime, - tool_parameters={}, - workflow_call_depth=0, - provider_name="provider", - ) - ) - - assert not hasattr(tool_runtime, "conversation_id") - assert len(messages) == 1 - graph_message = messages[0] - assert graph_message.type == ToolRuntimeMessage.MessageType.LINK - assert isinstance(graph_message.message, ToolRuntimeMessage.TextMessage) - assert graph_message.message.text == "https://dify.ai" - - callback = generic_invoke_mock.call_args.kwargs["workflow_tool_callback"] - assert isinstance(callback, DifyWorkflowCallbackHandler) - assert generic_invoke_mock.call_args.kwargs["session"] is not None - assert generic_invoke_mock.call_args.kwargs["conversation_id"] == "conversation-id" - - transform_kwargs = transform_tool_messages.call_args.kwargs - assert transform_kwargs["conversation_id"] == "conversation-id" - - -def test_invoke_maps_plugin_errors_to_graph_errors(runtime: DifyToolNodeRuntime) -> None: - invoke_error = PluginInvokeError('{"error_type":"RateLimit","message":"too many"}') - - with patch.object(ToolEngine, "generic_invoke", side_effect=invoke_error): - with pytest.raises(ToolRuntimeInvocationError, match="An error occurred in the provider"): - list( - runtime.invoke( - tool_runtime=ToolRuntimeHandle(raw=MagicMock()), - tool_parameters={}, - workflow_call_depth=0, - provider_name="provider", - ) - ) - - -def test_get_usage_normalizes_dict_payload(runtime: DifyToolNodeRuntime) -> None: - usage_payload = LLMUsage.empty_usage().model_dump() - usage_payload["total_tokens"] = 42 - - usage = runtime.get_usage( - tool_runtime=ToolRuntimeHandle(raw=SimpleNamespace(latest_usage=usage_payload)), - ) - - assert usage.total_tokens == 42 - - -def test_get_runtime_converts_graph_provider_type_for_tool_manager(runtime: DifyToolNodeRuntime) -> None: - node_data = _build_tool_node_data() - - with patch.object(ToolManager, "get_workflow_tool_runtime", return_value=MagicMock()) as runtime_mock: - tool_runtime = runtime.get_runtime(node_id="node-id", node_data=node_data, variable_pool=None) - - assert not hasattr(tool_runtime, "conversation_id") - workflow_tool = runtime_mock.call_args.args[3] - assert workflow_tool.provider_type == CoreToolProviderType.BUILT_IN - - -def test_get_runtime_stores_parent_trace_context_for_workflow_tools( - runtime: DifyToolNodeRuntime, -) -> None: - variable_pool: VariablePool = build_test_variable_pool( - variables=build_system_variables( - conversation_id="conversation-id", - workflow_execution_id="workflow-run-id", - ) - ) - workflow_runtime = MagicMock() - workflow_runtime.runtime.runtime_parameters = {} - node_data = ToolNodeData.model_validate( - { - "type": "tool", - "title": "Tool", - "provider_id": "provider", - "provider_type": ToolProviderType.WORKFLOW, - "provider_name": "provider", - "tool_name": "lookup", - "tool_label": "Lookup", - "tool_configurations": {}, - "tool_parameters": {}, - } - ) - - with patch.object(ToolManager, "get_workflow_tool_runtime", return_value=workflow_runtime): - tool_runtime = runtime.get_runtime( - node_id="node-id", - node_data=node_data, - variable_pool=variable_pool, - node_execution_id="node-execution-id", - ) - - assert tool_runtime.raw.parent_trace_context.model_dump() == { - "parent_workflow_run_id": "workflow-run-id", - "parent_node_execution_id": "node-execution-id", - } - assert workflow_runtime.runtime.runtime_parameters == {} - - -def test_get_runtime_stores_trace_session_id_for_workflow_tools( - runtime: DifyToolNodeRuntime, -) -> None: - variable_pool: VariablePool = build_test_variable_pool( - variables=build_system_variables( - conversation_id="conversation-id", - workflow_execution_id="workflow-run-id", - ) - ) - workflow_runtime = MagicMock() - workflow_runtime.runtime.runtime_parameters = {} - runtime._run_context.trace_session_id = "session-1" - node_data = ToolNodeData.model_validate( - { - "type": "tool", - "title": "Tool", - "provider_id": "provider", - "provider_type": ToolProviderType.WORKFLOW, - "provider_name": "provider", - "tool_name": "lookup", - "tool_label": "Lookup", - "tool_configurations": {}, - "tool_parameters": {}, - } - ) - - with patch.object(ToolManager, "get_workflow_tool_runtime", return_value=workflow_runtime): - tool_runtime = runtime.get_runtime( - node_id="node-id", - node_data=node_data, - variable_pool=variable_pool, - node_execution_id="node-execution-id", - ) - - assert tool_runtime.raw.trace_session_id == "session-1" - assert workflow_runtime.runtime.runtime_parameters == {} - - -def test_get_runtime_leaves_non_workflow_tool_runtime_parameters_unchanged( - runtime: DifyToolNodeRuntime, -) -> None: - variable_pool: VariablePool = build_test_variable_pool( - variables=build_system_variables( - conversation_id="conversation-id", - workflow_execution_id="workflow-run-id", - ) - ) - builtin_runtime = MagicMock() - builtin_runtime.runtime.runtime_parameters = {} - - with patch.object(ToolManager, "get_workflow_tool_runtime", return_value=builtin_runtime): - runtime.get_runtime( - node_id="node-id", - node_data=_build_tool_node_data(), - variable_pool=variable_pool, - node_execution_id="node-execution-id", - ) - - assert builtin_runtime.runtime.runtime_parameters == {} - - -def test_get_runtime_parameters_reads_required_flags(runtime: DifyToolNodeRuntime) -> None: - tool_runtime = ToolRuntimeHandle( - raw=SimpleNamespace( - get_merged_runtime_parameters=MagicMock( - return_value=[ - SimpleNamespace(name="city", required=True), - SimpleNamespace(name="country", required=False), - ] - ) - ) - ) - - parameters = runtime.get_runtime_parameters(tool_runtime=tool_runtime) - - assert [(parameter.name, parameter.required) for parameter in parameters] == [ - ("city", True), - ("country", False), - ] - - -def test_get_usage_returns_empty_usage_when_tool_has_no_usage(runtime: DifyToolNodeRuntime) -> None: - usage = runtime.get_usage(tool_runtime=ToolRuntimeHandle(raw=SimpleNamespace(latest_usage=None))) - - assert usage == LLMUsage.empty_usage() - - -@pytest.mark.parametrize( - ("payload", "expected_type"), - [ - (ToolInvokeMessage.JsonMessage(json_object={"ok": True}, suppress_output=True), ToolRuntimeMessage.JsonMessage), - (ToolInvokeMessage.BlobMessage(blob=b"bytes"), ToolRuntimeMessage.BlobMessage), - ( - ToolInvokeMessage.BlobChunkMessage( - id="blob-id", - sequence=1, - total_length=5, - blob=b"hello", - end=True, - ), - ToolRuntimeMessage.BlobChunkMessage, - ), - (ToolInvokeMessage.FileMessage(file_marker="marker"), ToolRuntimeMessage.FileMessage), - ( - ToolInvokeMessage.VariableMessage(variable_name="city", variable_value="Tokyo", stream=True), - ToolRuntimeMessage.VariableMessage, - ), - ( - ToolInvokeMessage.LogMessage( - id="log-id", - label="lookup", - status=ToolInvokeMessage.LogMessage.LogStatus.SUCCESS, - data={"count": 1}, - metadata={"source": "tool"}, - ), - ToolRuntimeMessage.LogMessage, - ), - ], -) -def test_convert_message_payload_supports_runtime_message_types( - runtime: DifyToolNodeRuntime, - payload: object, - expected_type: type[object], -) -> None: - message = runtime._convert_message_payload(payload) - - assert isinstance(message, expected_type) - - -def test_convert_message_payload_rejects_unknown_types(runtime: DifyToolNodeRuntime) -> None: - with pytest.raises(TypeError, match="unsupported tool message payload"): - runtime._convert_message_payload(object()) - - -def test_resolve_provider_icons_prefers_builtin_tool_icons(runtime: DifyToolNodeRuntime) -> None: - plugin = SimpleNamespace( - plugin_id="langgenius/tools", - name="search", - declaration=SimpleNamespace(icon={"plugin": "icon"}), - ) - builtin_tool = SimpleNamespace( - name="langgenius/tools/search", - icon={"builtin": "icon"}, - icon_dark={"builtin": "dark"}, - ) - - with ( - patch("core.workflow.node_runtime.PluginInstaller") as installer_cls, - patch("core.workflow.node_runtime.BuiltinToolManageService.list_builtin_tools", return_value=[builtin_tool]), - ): - installer_cls.return_value.list_plugins.return_value = [plugin] - - icon, icon_dark = runtime.resolve_provider_icons(provider_name="langgenius/tools/search") - - assert icon == {"builtin": "icon"} - assert icon_dark == {"builtin": "dark"} - - -def test_resolve_provider_icons_returns_default_when_provider_is_unknown(runtime: DifyToolNodeRuntime) -> None: - with ( - patch("core.workflow.node_runtime.PluginInstaller") as installer_cls, - patch("core.workflow.node_runtime.BuiltinToolManageService.list_builtin_tools", return_value=[]), - ): - installer_cls.return_value.list_plugins.return_value = [] - - icon, icon_dark = runtime.resolve_provider_icons(provider_name="unknown", default_icon="fallback") - - assert icon == "fallback" - assert icon_dark is None - - -@pytest.mark.parametrize( - ("exc", "message"), - [ - (PluginDaemonClientSideError("bad request"), "Failed to invoke tool, error: bad request"), - (ToolInvokeError("broken"), "Failed to invoke tool provider: broken"), - (RuntimeError("unexpected"), "unexpected"), - ], -) -def test_map_invocation_exception_normalizes_runtime_errors( - runtime: DifyToolNodeRuntime, - exc: Exception, - message: str, -) -> None: - error = runtime._map_invocation_exception(exc, provider_name="provider") - - assert isinstance(error, ToolRuntimeInvocationError) - assert str(error) == message diff --git a/api/tests/unit_tests/core/workflow/nodes/webhook/test_entities.py b/api/tests/unit_tests/core/workflow/nodes/trigger_webhook/test_entities.py similarity index 100% rename from api/tests/unit_tests/core/workflow/nodes/webhook/test_entities.py rename to api/tests/unit_tests/core/workflow/nodes/trigger_webhook/test_entities.py diff --git a/api/tests/unit_tests/core/workflow/nodes/webhook/test_exceptions.py b/api/tests/unit_tests/core/workflow/nodes/trigger_webhook/test_exc.py similarity index 100% rename from api/tests/unit_tests/core/workflow/nodes/webhook/test_exceptions.py rename to api/tests/unit_tests/core/workflow/nodes/trigger_webhook/test_exc.py diff --git a/api/tests/unit_tests/core/workflow/nodes/trigger_webhook/test_node.py b/api/tests/unit_tests/core/workflow/nodes/trigger_webhook/test_node.py new file mode 100644 index 00000000000..4b45f2c0ecf --- /dev/null +++ b/api/tests/unit_tests/core/workflow/nodes/trigger_webhook/test_node.py @@ -0,0 +1,904 @@ +from typing import Any +from unittest.mock import Mock, patch + +import pytest + +from core.app.entities.app_invoke_entities import DIFY_RUN_CONTEXT_KEY, InvokeFrom, UserFrom +from core.workflow.nodes.trigger_webhook.entities import ( + ContentType, + Method, + WebhookBodyParameter, + WebhookData, + WebhookParameter, +) +from core.workflow.nodes.trigger_webhook.node import TriggerWebhookNode +from core.workflow.system_variables import default_system_variables +from graphon.enums import WorkflowNodeExecutionStatus +from graphon.file import File, FileTransferMethod, FileType +from graphon.runtime import InitParams, RuntimeState, VariablePool +from graphon.variables import FileVariable, StringVariable +from tests.workflow_test_utils import build_test_variable_pool + +""" +Unit tests for webhook file conversion fix. + +This test verifies that webhook trigger nodes properly convert file dictionaries +to FileVariable objects, fixing the "Invalid variable type: ObjectVariable" error +when passing files to downstream LLM nodes. +""" + + +def _webhook_file_conversion_create_webhook_node( + webhook_data: WebhookData, + variable_pool: VariablePool, + tenant_id: str = "test-tenant", +) -> TriggerWebhookNode: + """Helper function to create a webhook node with proper initialization.""" + graph_init_params = InitParams( + workflow_id="test-workflow", + graph_config={}, + run_context={ + DIFY_RUN_CONTEXT_KEY: { + "tenant_id": tenant_id, + "app_id": "test-app", + "user_id": "test-user", + "user_from": UserFrom.ACCOUNT, + "invoke_from": InvokeFrom.SERVICE_API, + } + }, + call_depth=0, + ) + + runtime_state = RuntimeState( + workflow_id="test-workflow", + variable_pool=variable_pool, + start_at=0, + ) + + node = TriggerWebhookNode( + node_id="webhook-node-1", + data=webhook_data, + init_params=graph_init_params, + runtime_state=runtime_state, + ) + + # Attach a lightweight app_config onto runtime state for tenant lookups + runtime_state.app_config = Mock() + runtime_state.app_config.tenant_id = tenant_id + + return node + + +def create_test_file_dict( + filename: str = "test.jpg", + file_type: str = "image", + transfer_method: str = "local_file", +) -> dict: + """Create a test file dictionary as it would come from webhook service.""" + return { + "id": "file-123", + "tenant_id": "test-tenant", + "type": file_type, + "filename": filename, + "extension": ".jpg", + "mime_type": "image/jpeg", + "transfer_method": transfer_method, + "related_id": "related-123", + "storage_key": "storage-key-123", + "size": 1024, + "url": "https://example.com/test.jpg", + "created_at": 1234567890, + "used_at": None, + "hash": "file-hash-123", + } + + +def _webhook_file_conversion_build_webhook_variable_pool(inputs: dict[str, Any]) -> VariablePool: + return build_test_variable_pool( + variables=default_system_variables(), + node_id="webhook-node-1", + inputs=inputs, + ) + + +def expected_factory_mapping(file_dict: dict[str, Any]) -> dict[str, Any]: + return {**file_dict, "upload_file_id": file_dict["related_id"]} + + +def test_webhook_node_file_conversion_to_file_variable(): + """Test that webhook node converts file dictionaries to FileVariable objects.""" + # Create test file dictionary (as it comes from webhook service) + file_dict = create_test_file_dict("uploaded_image.jpg") + + data = WebhookData( + title="Test Webhook with File", + method=Method.POST, + content_type=ContentType.FORM_DATA, + body=[ + WebhookBodyParameter(name="image_upload", type="file", required=True), + WebhookBodyParameter(name="message", type="string", required=False), + ], + ) + + variable_pool = _webhook_file_conversion_build_webhook_variable_pool( + { + "webhook_data": { + "headers": {}, + "query_params": {}, + "body": {"message": "Test message"}, + "files": { + "image_upload": file_dict, + }, + } + } + ) + + node = _webhook_file_conversion_create_webhook_node(data, variable_pool) + + # Mock the file reference boundary and variable factory + with ( + patch.object(node._file_reference_factory, "build_from_mapping") as mock_file_factory, + patch("core.workflow.nodes.trigger_webhook.node.build_segment_with_type") as mock_segment_factory, + patch("core.workflow.nodes.trigger_webhook.node.FileVariable") as mock_file_variable, + ): + # Setup mocks + mock_file_obj = Mock() + mock_file_obj.to_dict.return_value = file_dict + mock_file_factory.return_value = mock_file_obj + + mock_segment = Mock() + mock_segment.value = mock_file_obj + mock_segment_factory.return_value = mock_segment + + mock_file_var_instance = Mock() + mock_file_variable.return_value = mock_file_var_instance + + # Run the node + result = node._run() + + # Verify successful execution + assert result.status == WorkflowNodeExecutionStatus.SUCCEEDED + + # Verify file factory was called with correct parameters + mock_file_factory.assert_called_once_with( + mapping=expected_factory_mapping(file_dict), + ) + + # Verify segment factory was called to create FileSegment + mock_segment_factory.assert_called_once() + + # Verify FileVariable was created with correct parameters + mock_file_variable.assert_called_once() + call_args = mock_file_variable.call_args[1] + assert call_args["name"] == "image_upload" + # value should be whatever build_segment_with_type.value returned + assert call_args["value"] == mock_segment.value + assert call_args["selector"] == ["webhook-node-1", "image_upload"] + + # Verify output contains the FileVariable, not the original dict + assert result.outputs["image_upload"] == mock_file_var_instance + assert result.outputs["message"] == "Test message" + + +def test_webhook_node_file_conversion_with_missing_files(): + """Test webhook node file conversion with missing file parameter.""" + data = WebhookData( + title="Test Webhook with Missing File", + method=Method.POST, + content_type=ContentType.FORM_DATA, + body=[ + WebhookBodyParameter(name="missing_file", type="file", required=False), + ], + ) + + variable_pool = _webhook_file_conversion_build_webhook_variable_pool( + { + "webhook_data": { + "headers": {}, + "query_params": {}, + "body": {}, + "files": {}, # No files + } + } + ) + + node = _webhook_file_conversion_create_webhook_node(data, variable_pool) + + # Run the node without patches (should handle None case gracefully) + result = node._run() + + # Verify successful execution + assert result.status == WorkflowNodeExecutionStatus.SUCCEEDED + + # Verify missing file parameter is None + assert result.outputs["_webhook_raw"]["files"] == {} + + +def test_webhook_node_file_conversion_with_none_file(): + """Test webhook node file conversion with None file value.""" + data = WebhookData( + title="Test Webhook with None File", + method=Method.POST, + content_type=ContentType.FORM_DATA, + body=[ + WebhookBodyParameter(name="none_file", type="file", required=False), + ], + ) + + variable_pool = _webhook_file_conversion_build_webhook_variable_pool( + { + "webhook_data": { + "headers": {}, + "query_params": {}, + "body": {}, + "files": { + "file": None, + }, + } + } + ) + + node = _webhook_file_conversion_create_webhook_node(data, variable_pool) + + # Run the node without patches (should handle None case gracefully) + result = node._run() + + # Verify successful execution + assert result.status == WorkflowNodeExecutionStatus.SUCCEEDED + + # Verify None file parameter is None + assert result.outputs["_webhook_raw"]["files"]["file"] is None + + +def test_webhook_node_file_conversion_with_non_dict_file(): + """Test webhook node file conversion with non-dict file value.""" + data = WebhookData( + title="Test Webhook with Non-Dict File", + method=Method.POST, + content_type=ContentType.FORM_DATA, + body=[ + WebhookBodyParameter(name="wrong_type", type="file", required=True), + ], + ) + + variable_pool = _webhook_file_conversion_build_webhook_variable_pool( + { + "webhook_data": { + "headers": {}, + "query_params": {}, + "body": {}, + "files": { + "file": "not_a_dict", # Wrapped to match node expectation + }, + } + } + ) + + node = _webhook_file_conversion_create_webhook_node(data, variable_pool) + + # Run the node without patches (should handle non-dict case gracefully) + result = node._run() + + # Verify successful execution + assert result.status == WorkflowNodeExecutionStatus.SUCCEEDED + + # Verify fallback to original (wrapped) mapping + assert result.outputs["_webhook_raw"]["files"]["file"] == "not_a_dict" + + +def test_webhook_node_file_conversion_mixed_parameters(): + """Test webhook node with mixed parameter types including files.""" + file_dict = create_test_file_dict("mixed_test.jpg") + + data = WebhookData( + title="Test Webhook Mixed Parameters", + method=Method.POST, + content_type=ContentType.FORM_DATA, + headers=[], + params=[], + body=[ + WebhookBodyParameter(name="text_param", type="string", required=True), + WebhookBodyParameter(name="number_param", type="number", required=False), + WebhookBodyParameter(name="file_param", type="file", required=True), + WebhookBodyParameter(name="bool_param", type="boolean", required=False), + ], + ) + + variable_pool = _webhook_file_conversion_build_webhook_variable_pool( + { + "webhook_data": { + "headers": {}, + "query_params": {}, + "body": { + "text_param": "Hello World", + "number_param": 42, + "bool_param": True, + }, + "files": { + "file_param": file_dict, + }, + } + } + ) + + node = _webhook_file_conversion_create_webhook_node(data, variable_pool) + + with ( + patch.object(node._file_reference_factory, "build_from_mapping") as mock_file_factory, + patch("core.workflow.nodes.trigger_webhook.node.build_segment_with_type") as mock_segment_factory, + patch("core.workflow.nodes.trigger_webhook.node.FileVariable") as mock_file_variable, + ): + # Setup mocks for file + mock_file_obj = Mock() + mock_file_factory.return_value = mock_file_obj + + mock_segment = Mock() + mock_segment.value = mock_file_obj + mock_segment_factory.return_value = mock_segment + + mock_file_var = Mock() + mock_file_variable.return_value = mock_file_var + + # Run the node + result = node._run() + + # Verify successful execution + assert result.status == WorkflowNodeExecutionStatus.SUCCEEDED + + # Verify all parameters are present + assert result.outputs["text_param"] == "Hello World" + assert result.outputs["number_param"] == 42 + assert result.outputs["bool_param"] is True + assert result.outputs["file_param"] == mock_file_var + + # Verify file conversion was called + mock_file_factory.assert_called_once_with( + mapping=expected_factory_mapping(file_dict), + ) + + +def test_webhook_node_different_file_types(): + """Test webhook node file conversion with different file types.""" + image_dict = create_test_file_dict("image.jpg", "image") + + data = WebhookData( + title="Test Webhook Different File Types", + method=Method.POST, + content_type=ContentType.FORM_DATA, + body=[ + WebhookBodyParameter(name="image", type="file", required=True), + WebhookBodyParameter(name="document", type="file", required=True), + WebhookBodyParameter(name="video", type="file", required=True), + ], + ) + + variable_pool = _webhook_file_conversion_build_webhook_variable_pool( + { + "webhook_data": { + "headers": {}, + "query_params": {}, + "body": {}, + "files": { + "image": image_dict, + "document": create_test_file_dict("document.pdf", "document"), + "video": create_test_file_dict("video.mp4", "video"), + }, + } + } + ) + + node = _webhook_file_conversion_create_webhook_node(data, variable_pool) + + with ( + patch.object(node._file_reference_factory, "build_from_mapping") as mock_file_factory, + patch("core.workflow.nodes.trigger_webhook.node.build_segment_with_type") as mock_segment_factory, + patch("core.workflow.nodes.trigger_webhook.node.FileVariable") as mock_file_variable, + ): + # Setup mocks for all files + mock_file_objs = [Mock() for _ in range(3)] + mock_segments = [Mock() for _ in range(3)] + mock_file_vars = [Mock() for _ in range(3)] + + # Map each segment.value to its corresponding mock file obj + for seg, f in zip(mock_segments, mock_file_objs, strict=False): + seg.value = f + + mock_file_factory.side_effect = mock_file_objs + mock_segment_factory.side_effect = mock_segments + mock_file_variable.side_effect = mock_file_vars + + # Run the node + result = node._run() + + # Verify successful execution + assert result.status == WorkflowNodeExecutionStatus.SUCCEEDED + + # Verify all file types were converted + assert mock_file_factory.call_count == 3 + assert result.outputs["image"] == mock_file_vars[0] + assert result.outputs["document"] == mock_file_vars[1] + assert result.outputs["video"] == mock_file_vars[2] + + +def test_webhook_node_file_conversion_with_non_dict_wrapper(): + """Test webhook node file conversion when the file wrapper is not a dict.""" + data = WebhookData( + title="Test Webhook with Non-dict File Wrapper", + method=Method.POST, + content_type=ContentType.FORM_DATA, + body=[ + WebhookBodyParameter(name="non_dict_wrapper", type="file", required=True), + ], + ) + + variable_pool = _webhook_file_conversion_build_webhook_variable_pool( + { + "webhook_data": { + "headers": {}, + "query_params": {}, + "body": {}, + "files": { + "file": "just a string", + }, + } + } + ) + + node = _webhook_file_conversion_create_webhook_node(data, variable_pool) + result = node._run() + + # Verify successful execution (should not crash) + assert result.status == WorkflowNodeExecutionStatus.SUCCEEDED + # Verify fallback to original value + assert result.outputs["_webhook_raw"]["files"]["file"] == "just a string" + + +def _webhook_node_create_webhook_node(webhook_data: WebhookData, variable_pool: VariablePool) -> TriggerWebhookNode: + """Helper function to create a webhook node with proper initialization.""" + graph_init_params = InitParams( + workflow_id="1", + graph_config={}, + run_context={ + DIFY_RUN_CONTEXT_KEY: { + "tenant_id": "1", + "app_id": "1", + "user_id": "1", + "user_from": UserFrom.ACCOUNT, + "invoke_from": InvokeFrom.SERVICE_API, + } + }, + call_depth=0, + ) + runtime_state = RuntimeState( + workflow_id="test-workflow", + variable_pool=variable_pool, + start_at=0, + ) + node = TriggerWebhookNode( + node_id="1", + data=webhook_data, + init_params=graph_init_params, + runtime_state=runtime_state, + ) + + # Provide tenant_id for conversion path + runtime_state.app_config = type("_AppCfg", (), {"tenant_id": "1"})() + + return node + + +def _webhook_node_build_webhook_variable_pool(inputs: dict[str, Any]) -> VariablePool: + return build_test_variable_pool( + variables=default_system_variables(), + node_id="1", + inputs=inputs, + ) + + +def test_webhook_node_default_config(): + """Test webhook node default configuration.""" + config = TriggerWebhookNode.get_default_config() + + assert config["type"] == "webhook" + assert config["config"]["method"] == "get" + assert config["config"]["content_type"] == "application/json" + assert config["config"]["headers"] == [] + assert config["config"]["params"] == [] + assert config["config"]["body"] == [] + assert config["config"]["async_mode"] is True + assert config["config"]["status_code"] == 200 + assert config["config"]["response_body"] == "" + assert config["config"]["timeout"] == 30 + + +def test_webhook_node_run_with_headers(): + """Test webhook node execution with header extraction.""" + data = WebhookData( + title="Test Webhook", + headers=[ + WebhookParameter(name="Authorization", required=True), + WebhookParameter(name="Content-Type", required=False), + ], + ) + + variable_pool = _webhook_node_build_webhook_variable_pool( + { + "webhook_data": { + "headers": { + "Authorization": "Bearer token123", + "content-type": "application/json", # Different case + "X-Custom": "custom-value", + }, + "query_params": {}, + "body": {}, + "files": {}, + } + } + ) + + node = _webhook_node_create_webhook_node(data, variable_pool) + result = node._run() + + assert result.status == WorkflowNodeExecutionStatus.SUCCEEDED + assert result.outputs["Authorization"] == "Bearer token123" + assert result.outputs["Content_Type"] == "application/json" # Case-insensitive match + assert "_webhook_raw" in result.outputs + + +def test_webhook_node_run_with_query_params(): + """Test webhook node execution with query parameter extraction.""" + data = WebhookData( + title="Test Webhook", + params=[ + WebhookParameter(name="page", required=True), + WebhookParameter(name="limit", required=False), + WebhookParameter(name="missing", required=False), + ], + ) + + variable_pool = _webhook_node_build_webhook_variable_pool( + { + "webhook_data": { + "headers": {}, + "query_params": { + "page": "1", + "limit": "10", + }, + "body": {}, + "files": {}, + } + } + ) + + node = _webhook_node_create_webhook_node(data, variable_pool) + result = node._run() + + assert result.status == WorkflowNodeExecutionStatus.SUCCEEDED + assert result.outputs["page"] == "1" + assert result.outputs["limit"] == "10" + assert result.outputs["missing"] is None # Missing parameter should be None + + +def test_webhook_node_run_with_body_params(): + """Test webhook node execution with body parameter extraction.""" + data = WebhookData( + title="Test Webhook", + body=[ + WebhookBodyParameter(name="message", type="string", required=True), + WebhookBodyParameter(name="count", type="number", required=False), + WebhookBodyParameter(name="active", type="boolean", required=False), + WebhookBodyParameter(name="metadata", type="object", required=False), + ], + ) + + variable_pool = _webhook_node_build_webhook_variable_pool( + { + "webhook_data": { + "headers": {}, + "query_params": {}, + "body": { + "message": "Hello World", + "count": 42, + "active": True, + "metadata": {"key": "value"}, + }, + "files": {}, + } + } + ) + + node = _webhook_node_create_webhook_node(data, variable_pool) + result = node._run() + + assert result.status == WorkflowNodeExecutionStatus.SUCCEEDED + assert result.outputs["message"] == "Hello World" + assert result.outputs["count"] == 42 + assert result.outputs["active"] is True + assert result.outputs["metadata"] == {"key": "value"} + + +def test_webhook_node_run_with_file_params(): + """Test webhook node execution with file parameter extraction.""" + # Create mock file objects + file1 = File( + file_type=FileType.IMAGE, + transfer_method=FileTransferMethod.LOCAL_FILE, + related_id="file1", + filename="image.jpg", + mime_type="image/jpeg", + storage_key="", + ) + + file2 = File( + file_type=FileType.DOCUMENT, + transfer_method=FileTransferMethod.LOCAL_FILE, + related_id="file2", + filename="document.pdf", + mime_type="application/pdf", + storage_key="", + ) + + data = WebhookData( + title="Test Webhook", + body=[ + WebhookBodyParameter(name="upload", type="file", required=True), + WebhookBodyParameter(name="document", type="file", required=False), + WebhookBodyParameter(name="missing_file", type="file", required=False), + ], + ) + + variable_pool = _webhook_node_build_webhook_variable_pool( + { + "webhook_data": { + "headers": {}, + "query_params": {}, + "body": {}, + "files": { + "upload": file1.to_dict(), + "document": file2.to_dict(), + }, + } + } + ) + + node = _webhook_node_create_webhook_node(data, variable_pool) + # Mock the node's file reference boundary to avoid DB-dependent validation on upload_file_id + with patch.object(node._file_reference_factory, "build_from_mapping") as mock_file_factory: + + def _to_file(*, mapping: dict[str, Any]) -> File: + return File( + file_id=mapping.get("id"), + file_type=FileType(mapping["type"]), + transfer_method=FileTransferMethod(mapping["transfer_method"]), + related_id=mapping.get("related_id"), + filename=mapping.get("filename"), + extension=mapping.get("extension"), + mime_type=mapping.get("mime_type"), + size=mapping.get("size", -1), + storage_key=mapping.get("storage_key", ""), + remote_url=mapping.get("url"), + ) + + mock_file_factory.side_effect = _to_file + result = node._run() + + assert result.status == WorkflowNodeExecutionStatus.SUCCEEDED + assert isinstance(result.outputs["upload"], FileVariable) + assert isinstance(result.outputs["document"], FileVariable) + assert result.outputs["upload"].value.filename == "image.jpg" + + +def test_webhook_node_run_mixed_parameters(): + """Test webhook node execution with mixed parameter types.""" + file_obj = File( + file_type=FileType.IMAGE, + transfer_method=FileTransferMethod.LOCAL_FILE, + related_id="file1", + filename="test.jpg", + mime_type="image/jpeg", + storage_key="", + ) + + data = WebhookData( + title="Test Webhook", + headers=[WebhookParameter(name="Authorization", required=True)], + params=[WebhookParameter(name="version", required=False)], + body=[ + WebhookBodyParameter(name="message", type="string", required=True), + WebhookBodyParameter(name="upload", type="file", required=False), + ], + ) + + variable_pool = _webhook_node_build_webhook_variable_pool( + { + "webhook_data": { + "headers": {"Authorization": "Bearer token"}, + "query_params": {"version": "v1"}, + "body": {"message": "Test message"}, + "files": {"upload": file_obj.to_dict()}, + } + } + ) + + node = _webhook_node_create_webhook_node(data, variable_pool) + # Mock the node's file reference boundary to avoid DB-dependent validation on upload_file_id + with patch.object(node._file_reference_factory, "build_from_mapping") as mock_file_factory: + + def _to_file(*, mapping: dict[str, Any]) -> File: + return File( + file_id=mapping.get("id"), + file_type=FileType(mapping["type"]), + transfer_method=FileTransferMethod(mapping["transfer_method"]), + related_id=mapping.get("related_id"), + filename=mapping.get("filename"), + extension=mapping.get("extension"), + mime_type=mapping.get("mime_type"), + size=mapping.get("size", -1), + storage_key=mapping.get("storage_key", ""), + remote_url=mapping.get("url"), + ) + + mock_file_factory.side_effect = _to_file + result = node._run() + + assert result.status == WorkflowNodeExecutionStatus.SUCCEEDED + assert result.outputs["Authorization"] == "Bearer token" + assert result.outputs["version"] == "v1" + assert result.outputs["message"] == "Test message" + assert isinstance(result.outputs["upload"], FileVariable) + assert result.outputs["upload"].value.filename == "test.jpg" + assert "_webhook_raw" in result.outputs + + +def test_webhook_node_run_empty_webhook_data(): + """Test webhook node execution with empty webhook data.""" + data = WebhookData( + title="Test Webhook", + headers=[WebhookParameter(name="Authorization", required=False)], + params=[WebhookParameter(name="page", required=False)], + body=[WebhookBodyParameter(name="message", type="string", required=False)], + ) + + variable_pool = _webhook_node_build_webhook_variable_pool({}) # No webhook_data + + node = _webhook_node_create_webhook_node(data, variable_pool) + result = node._run() + + assert result.status == WorkflowNodeExecutionStatus.SUCCEEDED + assert result.outputs["Authorization"] is None + assert result.outputs["page"] is None + assert result.outputs["message"] is None + assert result.outputs["_webhook_raw"] == {} + + +def test_webhook_node_run_case_insensitive_headers(): + """Test webhook node header extraction is case-insensitive.""" + data = WebhookData( + title="Test Webhook", + headers=[ + WebhookParameter(name="Content-Type", required=True), + WebhookParameter(name="X-API-KEY", required=True), + WebhookParameter(name="authorization", required=True), + ], + ) + + variable_pool = _webhook_node_build_webhook_variable_pool( + { + "webhook_data": { + "headers": { + "content-type": "application/json", # lowercase + "x-api-key": "key123", # lowercase + "Authorization": "Bearer token", # different case + }, + "query_params": {}, + "body": {}, + "files": {}, + } + } + ) + + node = _webhook_node_create_webhook_node(data, variable_pool) + result = node._run() + + assert result.status == WorkflowNodeExecutionStatus.SUCCEEDED + assert result.outputs["Content_Type"] == "application/json" + assert result.outputs["X_API_KEY"] == "key123" + assert result.outputs["authorization"] == "Bearer token" + + +def test_webhook_node_variable_pool_user_inputs(): + """Test that webhook node uses user_inputs from variable pool correctly.""" + data = WebhookData(title="Test Webhook") + + # Add some additional variables to the pool + variable_pool = _webhook_node_build_webhook_variable_pool( + { + "webhook_data": {"headers": {}, "query_params": {}, "body": {}, "files": {}}, + "other_var": "should_be_included", + } + ) + variable_pool.add(["node1", "extra"], StringVariable(name="extra", value="extra_value")) + + node = _webhook_node_create_webhook_node(data, variable_pool) + result = node._run() + + assert result.status == WorkflowNodeExecutionStatus.SUCCEEDED + # Check that all user_inputs are included in the inputs (they get converted to dict) + inputs_dict = dict(result.inputs) + assert "webhook_data" in inputs_dict + assert "other_var" in inputs_dict + assert inputs_dict["other_var"] == "should_be_included" + + +@pytest.mark.parametrize( + "method", + [Method.GET, Method.POST, Method.PUT, Method.DELETE, Method.PATCH, Method.HEAD], +) +def test_webhook_node_different_methods(method): + """Test webhook node with different HTTP methods.""" + data = WebhookData( + title="Test Webhook", + method=method, + ) + + variable_pool = _webhook_node_build_webhook_variable_pool( + { + "webhook_data": { + "headers": {}, + "query_params": {}, + "body": {}, + "files": {}, + } + } + ) + + node = _webhook_node_create_webhook_node(data, variable_pool) + result = node._run() + + assert result.status == WorkflowNodeExecutionStatus.SUCCEEDED + assert node._node_data.method == method + + +def test_webhook_data_content_type_field(): + """Test that content_type accepts both raw strings and enum values.""" + data1 = WebhookData(title="Test", content_type="application/json") + assert data1.content_type == ContentType.JSON + + data2 = WebhookData(title="Test", content_type=ContentType.FORM_DATA) + assert data2.content_type == ContentType.FORM_DATA + + +def test_webhook_parameter_models(): + """Test webhook parameter model validation.""" + # Test WebhookParameter + param = WebhookParameter(name="test_param", required=True) + assert param.name == "test_param" + assert param.required is True + + param_default = WebhookParameter(name="test_param") + assert param_default.required is False + + # Test WebhookBodyParameter + body_param = WebhookBodyParameter(name="test_body", type="string", required=True) + assert body_param.name == "test_body" + assert body_param.type == "string" + assert body_param.required is True + + body_param_default = WebhookBodyParameter(name="test_body") + assert body_param_default.type == "string" # Default type + assert body_param_default.required is False + + +def test_webhook_data_field_defaults(): + """Test webhook data model field defaults.""" + data = WebhookData(title="Minimal Webhook") + + assert data.method == Method.GET + assert data.content_type == ContentType.JSON + assert data.headers == [] + assert data.params == [] + assert data.body == [] + assert data.status_code == 200 + assert data.response_body == "" + assert data.webhook_id is None + assert data.timeout == 30 diff --git a/api/tests/unit_tests/core/workflow/nodes/webhook/test_webhook_file_conversion.py b/api/tests/unit_tests/core/workflow/nodes/webhook/test_webhook_file_conversion.py deleted file mode 100644 index 758b5f0098c..00000000000 --- a/api/tests/unit_tests/core/workflow/nodes/webhook/test_webhook_file_conversion.py +++ /dev/null @@ -1,448 +0,0 @@ -""" -Unit tests for webhook file conversion fix. - -This test verifies that webhook trigger nodes properly convert file dictionaries -to FileVariable objects, fixing the "Invalid variable type: ObjectVariable" error -when passing files to downstream LLM nodes. -""" - -from typing import Any -from unittest.mock import Mock, patch - -from core.app.entities.app_invoke_entities import DIFY_RUN_CONTEXT_KEY, InvokeFrom, UserFrom -from core.workflow.nodes.trigger_webhook.entities import ( - ContentType, - Method, - WebhookBodyParameter, - WebhookData, -) -from core.workflow.nodes.trigger_webhook.node import TriggerWebhookNode -from core.workflow.system_variables import default_system_variables -from graphon.enums import WorkflowNodeExecutionStatus -from graphon.runtime import InitParams, RuntimeState, VariablePool -from tests.workflow_test_utils import build_test_variable_pool - - -def create_webhook_node( - webhook_data: WebhookData, - variable_pool: VariablePool, - tenant_id: str = "test-tenant", -) -> TriggerWebhookNode: - """Helper function to create a webhook node with proper initialization.""" - graph_init_params = InitParams( - workflow_id="test-workflow", - graph_config={}, - run_context={ - DIFY_RUN_CONTEXT_KEY: { - "tenant_id": tenant_id, - "app_id": "test-app", - "user_id": "test-user", - "user_from": UserFrom.ACCOUNT, - "invoke_from": InvokeFrom.SERVICE_API, - } - }, - call_depth=0, - ) - - runtime_state = RuntimeState( - workflow_id="test-workflow", - variable_pool=variable_pool, - start_at=0, - ) - - node = TriggerWebhookNode( - node_id="webhook-node-1", - data=webhook_data, - init_params=graph_init_params, - runtime_state=runtime_state, - ) - - # Attach a lightweight app_config onto runtime state for tenant lookups - runtime_state.app_config = Mock() - runtime_state.app_config.tenant_id = tenant_id - - return node - - -def create_test_file_dict( - filename: str = "test.jpg", - file_type: str = "image", - transfer_method: str = "local_file", -) -> dict: - """Create a test file dictionary as it would come from webhook service.""" - return { - "id": "file-123", - "tenant_id": "test-tenant", - "type": file_type, - "filename": filename, - "extension": ".jpg", - "mime_type": "image/jpeg", - "transfer_method": transfer_method, - "related_id": "related-123", - "storage_key": "storage-key-123", - "size": 1024, - "url": "https://example.com/test.jpg", - "created_at": 1234567890, - "used_at": None, - "hash": "file-hash-123", - } - - -def build_webhook_variable_pool(inputs: dict[str, Any]) -> VariablePool: - return build_test_variable_pool( - variables=default_system_variables(), - node_id="webhook-node-1", - inputs=inputs, - ) - - -def expected_factory_mapping(file_dict: dict[str, Any]) -> dict[str, Any]: - return {**file_dict, "upload_file_id": file_dict["related_id"]} - - -def test_webhook_node_file_conversion_to_file_variable(): - """Test that webhook node converts file dictionaries to FileVariable objects.""" - # Create test file dictionary (as it comes from webhook service) - file_dict = create_test_file_dict("uploaded_image.jpg") - - data = WebhookData( - title="Test Webhook with File", - method=Method.POST, - content_type=ContentType.FORM_DATA, - body=[ - WebhookBodyParameter(name="image_upload", type="file", required=True), - WebhookBodyParameter(name="message", type="string", required=False), - ], - ) - - variable_pool = build_webhook_variable_pool( - { - "webhook_data": { - "headers": {}, - "query_params": {}, - "body": {"message": "Test message"}, - "files": { - "image_upload": file_dict, - }, - } - } - ) - - node = create_webhook_node(data, variable_pool) - - # Mock the file reference boundary and variable factory - with ( - patch.object(node._file_reference_factory, "build_from_mapping") as mock_file_factory, - patch("core.workflow.nodes.trigger_webhook.node.build_segment_with_type") as mock_segment_factory, - patch("core.workflow.nodes.trigger_webhook.node.FileVariable") as mock_file_variable, - ): - # Setup mocks - mock_file_obj = Mock() - mock_file_obj.to_dict.return_value = file_dict - mock_file_factory.return_value = mock_file_obj - - mock_segment = Mock() - mock_segment.value = mock_file_obj - mock_segment_factory.return_value = mock_segment - - mock_file_var_instance = Mock() - mock_file_variable.return_value = mock_file_var_instance - - # Run the node - result = node._run() - - # Verify successful execution - assert result.status == WorkflowNodeExecutionStatus.SUCCEEDED - - # Verify file factory was called with correct parameters - mock_file_factory.assert_called_once_with( - mapping=expected_factory_mapping(file_dict), - ) - - # Verify segment factory was called to create FileSegment - mock_segment_factory.assert_called_once() - - # Verify FileVariable was created with correct parameters - mock_file_variable.assert_called_once() - call_args = mock_file_variable.call_args[1] - assert call_args["name"] == "image_upload" - # value should be whatever build_segment_with_type.value returned - assert call_args["value"] == mock_segment.value - assert call_args["selector"] == ["webhook-node-1", "image_upload"] - - # Verify output contains the FileVariable, not the original dict - assert result.outputs["image_upload"] == mock_file_var_instance - assert result.outputs["message"] == "Test message" - - -def test_webhook_node_file_conversion_with_missing_files(): - """Test webhook node file conversion with missing file parameter.""" - data = WebhookData( - title="Test Webhook with Missing File", - method=Method.POST, - content_type=ContentType.FORM_DATA, - body=[ - WebhookBodyParameter(name="missing_file", type="file", required=False), - ], - ) - - variable_pool = build_webhook_variable_pool( - { - "webhook_data": { - "headers": {}, - "query_params": {}, - "body": {}, - "files": {}, # No files - } - } - ) - - node = create_webhook_node(data, variable_pool) - - # Run the node without patches (should handle None case gracefully) - result = node._run() - - # Verify successful execution - assert result.status == WorkflowNodeExecutionStatus.SUCCEEDED - - # Verify missing file parameter is None - assert result.outputs["_webhook_raw"]["files"] == {} - - -def test_webhook_node_file_conversion_with_none_file(): - """Test webhook node file conversion with None file value.""" - data = WebhookData( - title="Test Webhook with None File", - method=Method.POST, - content_type=ContentType.FORM_DATA, - body=[ - WebhookBodyParameter(name="none_file", type="file", required=False), - ], - ) - - variable_pool = build_webhook_variable_pool( - { - "webhook_data": { - "headers": {}, - "query_params": {}, - "body": {}, - "files": { - "file": None, - }, - } - } - ) - - node = create_webhook_node(data, variable_pool) - - # Run the node without patches (should handle None case gracefully) - result = node._run() - - # Verify successful execution - assert result.status == WorkflowNodeExecutionStatus.SUCCEEDED - - # Verify None file parameter is None - assert result.outputs["_webhook_raw"]["files"]["file"] is None - - -def test_webhook_node_file_conversion_with_non_dict_file(): - """Test webhook node file conversion with non-dict file value.""" - data = WebhookData( - title="Test Webhook with Non-Dict File", - method=Method.POST, - content_type=ContentType.FORM_DATA, - body=[ - WebhookBodyParameter(name="wrong_type", type="file", required=True), - ], - ) - - variable_pool = build_webhook_variable_pool( - { - "webhook_data": { - "headers": {}, - "query_params": {}, - "body": {}, - "files": { - "file": "not_a_dict", # Wrapped to match node expectation - }, - } - } - ) - - node = create_webhook_node(data, variable_pool) - - # Run the node without patches (should handle non-dict case gracefully) - result = node._run() - - # Verify successful execution - assert result.status == WorkflowNodeExecutionStatus.SUCCEEDED - - # Verify fallback to original (wrapped) mapping - assert result.outputs["_webhook_raw"]["files"]["file"] == "not_a_dict" - - -def test_webhook_node_file_conversion_mixed_parameters(): - """Test webhook node with mixed parameter types including files.""" - file_dict = create_test_file_dict("mixed_test.jpg") - - data = WebhookData( - title="Test Webhook Mixed Parameters", - method=Method.POST, - content_type=ContentType.FORM_DATA, - headers=[], - params=[], - body=[ - WebhookBodyParameter(name="text_param", type="string", required=True), - WebhookBodyParameter(name="number_param", type="number", required=False), - WebhookBodyParameter(name="file_param", type="file", required=True), - WebhookBodyParameter(name="bool_param", type="boolean", required=False), - ], - ) - - variable_pool = build_webhook_variable_pool( - { - "webhook_data": { - "headers": {}, - "query_params": {}, - "body": { - "text_param": "Hello World", - "number_param": 42, - "bool_param": True, - }, - "files": { - "file_param": file_dict, - }, - } - } - ) - - node = create_webhook_node(data, variable_pool) - - with ( - patch.object(node._file_reference_factory, "build_from_mapping") as mock_file_factory, - patch("core.workflow.nodes.trigger_webhook.node.build_segment_with_type") as mock_segment_factory, - patch("core.workflow.nodes.trigger_webhook.node.FileVariable") as mock_file_variable, - ): - # Setup mocks for file - mock_file_obj = Mock() - mock_file_factory.return_value = mock_file_obj - - mock_segment = Mock() - mock_segment.value = mock_file_obj - mock_segment_factory.return_value = mock_segment - - mock_file_var = Mock() - mock_file_variable.return_value = mock_file_var - - # Run the node - result = node._run() - - # Verify successful execution - assert result.status == WorkflowNodeExecutionStatus.SUCCEEDED - - # Verify all parameters are present - assert result.outputs["text_param"] == "Hello World" - assert result.outputs["number_param"] == 42 - assert result.outputs["bool_param"] is True - assert result.outputs["file_param"] == mock_file_var - - # Verify file conversion was called - mock_file_factory.assert_called_once_with( - mapping=expected_factory_mapping(file_dict), - ) - - -def test_webhook_node_different_file_types(): - """Test webhook node file conversion with different file types.""" - image_dict = create_test_file_dict("image.jpg", "image") - - data = WebhookData( - title="Test Webhook Different File Types", - method=Method.POST, - content_type=ContentType.FORM_DATA, - body=[ - WebhookBodyParameter(name="image", type="file", required=True), - WebhookBodyParameter(name="document", type="file", required=True), - WebhookBodyParameter(name="video", type="file", required=True), - ], - ) - - variable_pool = build_webhook_variable_pool( - { - "webhook_data": { - "headers": {}, - "query_params": {}, - "body": {}, - "files": { - "image": image_dict, - "document": create_test_file_dict("document.pdf", "document"), - "video": create_test_file_dict("video.mp4", "video"), - }, - } - } - ) - - node = create_webhook_node(data, variable_pool) - - with ( - patch.object(node._file_reference_factory, "build_from_mapping") as mock_file_factory, - patch("core.workflow.nodes.trigger_webhook.node.build_segment_with_type") as mock_segment_factory, - patch("core.workflow.nodes.trigger_webhook.node.FileVariable") as mock_file_variable, - ): - # Setup mocks for all files - mock_file_objs = [Mock() for _ in range(3)] - mock_segments = [Mock() for _ in range(3)] - mock_file_vars = [Mock() for _ in range(3)] - - # Map each segment.value to its corresponding mock file obj - for seg, f in zip(mock_segments, mock_file_objs, strict=False): - seg.value = f - - mock_file_factory.side_effect = mock_file_objs - mock_segment_factory.side_effect = mock_segments - mock_file_variable.side_effect = mock_file_vars - - # Run the node - result = node._run() - - # Verify successful execution - assert result.status == WorkflowNodeExecutionStatus.SUCCEEDED - - # Verify all file types were converted - assert mock_file_factory.call_count == 3 - assert result.outputs["image"] == mock_file_vars[0] - assert result.outputs["document"] == mock_file_vars[1] - assert result.outputs["video"] == mock_file_vars[2] - - -def test_webhook_node_file_conversion_with_non_dict_wrapper(): - """Test webhook node file conversion when the file wrapper is not a dict.""" - data = WebhookData( - title="Test Webhook with Non-dict File Wrapper", - method=Method.POST, - content_type=ContentType.FORM_DATA, - body=[ - WebhookBodyParameter(name="non_dict_wrapper", type="file", required=True), - ], - ) - - variable_pool = build_webhook_variable_pool( - { - "webhook_data": { - "headers": {}, - "query_params": {}, - "body": {}, - "files": { - "file": "just a string", - }, - } - } - ) - - node = create_webhook_node(data, variable_pool) - result = node._run() - - # Verify successful execution (should not crash) - assert result.status == WorkflowNodeExecutionStatus.SUCCEEDED - # Verify fallback to original value - assert result.outputs["_webhook_raw"]["files"]["file"] == "just a string" diff --git a/api/tests/unit_tests/core/workflow/nodes/webhook/test_webhook_node.py b/api/tests/unit_tests/core/workflow/nodes/webhook/test_webhook_node.py deleted file mode 100644 index 8606f5631ca..00000000000 --- a/api/tests/unit_tests/core/workflow/nodes/webhook/test_webhook_node.py +++ /dev/null @@ -1,471 +0,0 @@ -from typing import Any -from unittest.mock import patch - -import pytest - -from core.app.entities.app_invoke_entities import DIFY_RUN_CONTEXT_KEY, InvokeFrom, UserFrom -from core.workflow.nodes.trigger_webhook.entities import ( - ContentType, - Method, - WebhookBodyParameter, - WebhookData, - WebhookParameter, -) -from core.workflow.nodes.trigger_webhook.node import TriggerWebhookNode -from core.workflow.system_variables import default_system_variables -from graphon.enums import WorkflowNodeExecutionStatus -from graphon.file import File, FileTransferMethod, FileType -from graphon.runtime import InitParams, RuntimeState, VariablePool -from graphon.variables import FileVariable, StringVariable -from tests.workflow_test_utils import build_test_variable_pool - - -def create_webhook_node(webhook_data: WebhookData, variable_pool: VariablePool) -> TriggerWebhookNode: - """Helper function to create a webhook node with proper initialization.""" - graph_init_params = InitParams( - workflow_id="1", - graph_config={}, - run_context={ - DIFY_RUN_CONTEXT_KEY: { - "tenant_id": "1", - "app_id": "1", - "user_id": "1", - "user_from": UserFrom.ACCOUNT, - "invoke_from": InvokeFrom.SERVICE_API, - } - }, - call_depth=0, - ) - runtime_state = RuntimeState( - workflow_id="test-workflow", - variable_pool=variable_pool, - start_at=0, - ) - node = TriggerWebhookNode( - node_id="1", - data=webhook_data, - init_params=graph_init_params, - runtime_state=runtime_state, - ) - - # Provide tenant_id for conversion path - runtime_state.app_config = type("_AppCfg", (), {"tenant_id": "1"})() - - return node - - -def build_webhook_variable_pool(inputs: dict[str, Any]) -> VariablePool: - return build_test_variable_pool( - variables=default_system_variables(), - node_id="1", - inputs=inputs, - ) - - -def test_webhook_node_default_config(): - """Test webhook node default configuration.""" - config = TriggerWebhookNode.get_default_config() - - assert config["type"] == "webhook" - assert config["config"]["method"] == "get" - assert config["config"]["content_type"] == "application/json" - assert config["config"]["headers"] == [] - assert config["config"]["params"] == [] - assert config["config"]["body"] == [] - assert config["config"]["async_mode"] is True - assert config["config"]["status_code"] == 200 - assert config["config"]["response_body"] == "" - assert config["config"]["timeout"] == 30 - - -def test_webhook_node_run_with_headers(): - """Test webhook node execution with header extraction.""" - data = WebhookData( - title="Test Webhook", - headers=[ - WebhookParameter(name="Authorization", required=True), - WebhookParameter(name="Content-Type", required=False), - ], - ) - - variable_pool = build_webhook_variable_pool( - { - "webhook_data": { - "headers": { - "Authorization": "Bearer token123", - "content-type": "application/json", # Different case - "X-Custom": "custom-value", - }, - "query_params": {}, - "body": {}, - "files": {}, - } - } - ) - - node = create_webhook_node(data, variable_pool) - result = node._run() - - assert result.status == WorkflowNodeExecutionStatus.SUCCEEDED - assert result.outputs["Authorization"] == "Bearer token123" - assert result.outputs["Content_Type"] == "application/json" # Case-insensitive match - assert "_webhook_raw" in result.outputs - - -def test_webhook_node_run_with_query_params(): - """Test webhook node execution with query parameter extraction.""" - data = WebhookData( - title="Test Webhook", - params=[ - WebhookParameter(name="page", required=True), - WebhookParameter(name="limit", required=False), - WebhookParameter(name="missing", required=False), - ], - ) - - variable_pool = build_webhook_variable_pool( - { - "webhook_data": { - "headers": {}, - "query_params": { - "page": "1", - "limit": "10", - }, - "body": {}, - "files": {}, - } - } - ) - - node = create_webhook_node(data, variable_pool) - result = node._run() - - assert result.status == WorkflowNodeExecutionStatus.SUCCEEDED - assert result.outputs["page"] == "1" - assert result.outputs["limit"] == "10" - assert result.outputs["missing"] is None # Missing parameter should be None - - -def test_webhook_node_run_with_body_params(): - """Test webhook node execution with body parameter extraction.""" - data = WebhookData( - title="Test Webhook", - body=[ - WebhookBodyParameter(name="message", type="string", required=True), - WebhookBodyParameter(name="count", type="number", required=False), - WebhookBodyParameter(name="active", type="boolean", required=False), - WebhookBodyParameter(name="metadata", type="object", required=False), - ], - ) - - variable_pool = build_webhook_variable_pool( - { - "webhook_data": { - "headers": {}, - "query_params": {}, - "body": { - "message": "Hello World", - "count": 42, - "active": True, - "metadata": {"key": "value"}, - }, - "files": {}, - } - } - ) - - node = create_webhook_node(data, variable_pool) - result = node._run() - - assert result.status == WorkflowNodeExecutionStatus.SUCCEEDED - assert result.outputs["message"] == "Hello World" - assert result.outputs["count"] == 42 - assert result.outputs["active"] is True - assert result.outputs["metadata"] == {"key": "value"} - - -def test_webhook_node_run_with_file_params(): - """Test webhook node execution with file parameter extraction.""" - # Create mock file objects - file1 = File( - file_type=FileType.IMAGE, - transfer_method=FileTransferMethod.LOCAL_FILE, - related_id="file1", - filename="image.jpg", - mime_type="image/jpeg", - storage_key="", - ) - - file2 = File( - file_type=FileType.DOCUMENT, - transfer_method=FileTransferMethod.LOCAL_FILE, - related_id="file2", - filename="document.pdf", - mime_type="application/pdf", - storage_key="", - ) - - data = WebhookData( - title="Test Webhook", - body=[ - WebhookBodyParameter(name="upload", type="file", required=True), - WebhookBodyParameter(name="document", type="file", required=False), - WebhookBodyParameter(name="missing_file", type="file", required=False), - ], - ) - - variable_pool = build_webhook_variable_pool( - { - "webhook_data": { - "headers": {}, - "query_params": {}, - "body": {}, - "files": { - "upload": file1.to_dict(), - "document": file2.to_dict(), - }, - } - } - ) - - node = create_webhook_node(data, variable_pool) - # Mock the node's file reference boundary to avoid DB-dependent validation on upload_file_id - with patch.object(node._file_reference_factory, "build_from_mapping") as mock_file_factory: - - def _to_file(*, mapping: dict[str, Any]) -> File: - return File( - file_id=mapping.get("id"), - file_type=FileType(mapping["type"]), - transfer_method=FileTransferMethod(mapping["transfer_method"]), - related_id=mapping.get("related_id"), - filename=mapping.get("filename"), - extension=mapping.get("extension"), - mime_type=mapping.get("mime_type"), - size=mapping.get("size", -1), - storage_key=mapping.get("storage_key", ""), - remote_url=mapping.get("url"), - ) - - mock_file_factory.side_effect = _to_file - result = node._run() - - assert result.status == WorkflowNodeExecutionStatus.SUCCEEDED - assert isinstance(result.outputs["upload"], FileVariable) - assert isinstance(result.outputs["document"], FileVariable) - assert result.outputs["upload"].value.filename == "image.jpg" - - -def test_webhook_node_run_mixed_parameters(): - """Test webhook node execution with mixed parameter types.""" - file_obj = File( - file_type=FileType.IMAGE, - transfer_method=FileTransferMethod.LOCAL_FILE, - related_id="file1", - filename="test.jpg", - mime_type="image/jpeg", - storage_key="", - ) - - data = WebhookData( - title="Test Webhook", - headers=[WebhookParameter(name="Authorization", required=True)], - params=[WebhookParameter(name="version", required=False)], - body=[ - WebhookBodyParameter(name="message", type="string", required=True), - WebhookBodyParameter(name="upload", type="file", required=False), - ], - ) - - variable_pool = build_webhook_variable_pool( - { - "webhook_data": { - "headers": {"Authorization": "Bearer token"}, - "query_params": {"version": "v1"}, - "body": {"message": "Test message"}, - "files": {"upload": file_obj.to_dict()}, - } - } - ) - - node = create_webhook_node(data, variable_pool) - # Mock the node's file reference boundary to avoid DB-dependent validation on upload_file_id - with patch.object(node._file_reference_factory, "build_from_mapping") as mock_file_factory: - - def _to_file(*, mapping: dict[str, Any]) -> File: - return File( - file_id=mapping.get("id"), - file_type=FileType(mapping["type"]), - transfer_method=FileTransferMethod(mapping["transfer_method"]), - related_id=mapping.get("related_id"), - filename=mapping.get("filename"), - extension=mapping.get("extension"), - mime_type=mapping.get("mime_type"), - size=mapping.get("size", -1), - storage_key=mapping.get("storage_key", ""), - remote_url=mapping.get("url"), - ) - - mock_file_factory.side_effect = _to_file - result = node._run() - - assert result.status == WorkflowNodeExecutionStatus.SUCCEEDED - assert result.outputs["Authorization"] == "Bearer token" - assert result.outputs["version"] == "v1" - assert result.outputs["message"] == "Test message" - assert isinstance(result.outputs["upload"], FileVariable) - assert result.outputs["upload"].value.filename == "test.jpg" - assert "_webhook_raw" in result.outputs - - -def test_webhook_node_run_empty_webhook_data(): - """Test webhook node execution with empty webhook data.""" - data = WebhookData( - title="Test Webhook", - headers=[WebhookParameter(name="Authorization", required=False)], - params=[WebhookParameter(name="page", required=False)], - body=[WebhookBodyParameter(name="message", type="string", required=False)], - ) - - variable_pool = build_webhook_variable_pool({}) # No webhook_data - - node = create_webhook_node(data, variable_pool) - result = node._run() - - assert result.status == WorkflowNodeExecutionStatus.SUCCEEDED - assert result.outputs["Authorization"] is None - assert result.outputs["page"] is None - assert result.outputs["message"] is None - assert result.outputs["_webhook_raw"] == {} - - -def test_webhook_node_run_case_insensitive_headers(): - """Test webhook node header extraction is case-insensitive.""" - data = WebhookData( - title="Test Webhook", - headers=[ - WebhookParameter(name="Content-Type", required=True), - WebhookParameter(name="X-API-KEY", required=True), - WebhookParameter(name="authorization", required=True), - ], - ) - - variable_pool = build_webhook_variable_pool( - { - "webhook_data": { - "headers": { - "content-type": "application/json", # lowercase - "x-api-key": "key123", # lowercase - "Authorization": "Bearer token", # different case - }, - "query_params": {}, - "body": {}, - "files": {}, - } - } - ) - - node = create_webhook_node(data, variable_pool) - result = node._run() - - assert result.status == WorkflowNodeExecutionStatus.SUCCEEDED - assert result.outputs["Content_Type"] == "application/json" - assert result.outputs["X_API_KEY"] == "key123" - assert result.outputs["authorization"] == "Bearer token" - - -def test_webhook_node_variable_pool_user_inputs(): - """Test that webhook node uses user_inputs from variable pool correctly.""" - data = WebhookData(title="Test Webhook") - - # Add some additional variables to the pool - variable_pool = build_webhook_variable_pool( - { - "webhook_data": {"headers": {}, "query_params": {}, "body": {}, "files": {}}, - "other_var": "should_be_included", - } - ) - variable_pool.add(["node1", "extra"], StringVariable(name="extra", value="extra_value")) - - node = create_webhook_node(data, variable_pool) - result = node._run() - - assert result.status == WorkflowNodeExecutionStatus.SUCCEEDED - # Check that all user_inputs are included in the inputs (they get converted to dict) - inputs_dict = dict(result.inputs) - assert "webhook_data" in inputs_dict - assert "other_var" in inputs_dict - assert inputs_dict["other_var"] == "should_be_included" - - -@pytest.mark.parametrize( - "method", - [Method.GET, Method.POST, Method.PUT, Method.DELETE, Method.PATCH, Method.HEAD], -) -def test_webhook_node_different_methods(method): - """Test webhook node with different HTTP methods.""" - data = WebhookData( - title="Test Webhook", - method=method, - ) - - variable_pool = build_webhook_variable_pool( - { - "webhook_data": { - "headers": {}, - "query_params": {}, - "body": {}, - "files": {}, - } - } - ) - - node = create_webhook_node(data, variable_pool) - result = node._run() - - assert result.status == WorkflowNodeExecutionStatus.SUCCEEDED - assert node._node_data.method == method - - -def test_webhook_data_content_type_field(): - """Test that content_type accepts both raw strings and enum values.""" - data1 = WebhookData(title="Test", content_type="application/json") - assert data1.content_type == ContentType.JSON - - data2 = WebhookData(title="Test", content_type=ContentType.FORM_DATA) - assert data2.content_type == ContentType.FORM_DATA - - -def test_webhook_parameter_models(): - """Test webhook parameter model validation.""" - # Test WebhookParameter - param = WebhookParameter(name="test_param", required=True) - assert param.name == "test_param" - assert param.required is True - - param_default = WebhookParameter(name="test_param") - assert param_default.required is False - - # Test WebhookBodyParameter - body_param = WebhookBodyParameter(name="test_body", type="string", required=True) - assert body_param.name == "test_body" - assert body_param.type == "string" - assert body_param.required is True - - body_param_default = WebhookBodyParameter(name="test_body") - assert body_param_default.type == "string" # Default type - assert body_param_default.required is False - - -def test_webhook_data_field_defaults(): - """Test webhook data model field defaults.""" - data = WebhookData(title="Minimal Webhook") - - assert data.method == Method.GET - assert data.content_type == ContentType.JSON - assert data.headers == [] - assert data.params == [] - assert data.body == [] - assert data.status_code == 200 - assert data.response_body == "" - assert data.webhook_id is None - assert data.timeout == 30 diff --git a/api/tests/unit_tests/core/workflow/test_enrich_pause_reasons.py b/api/tests/unit_tests/core/workflow/test_enrich_pause_reasons.py deleted file mode 100644 index 934bfbe5892..00000000000 --- a/api/tests/unit_tests/core/workflow/test_enrich_pause_reasons.py +++ /dev/null @@ -1,162 +0,0 @@ -from datetime import datetime - -import pytest - -from core.repositories.human_input_repository import HumanInputFormRecord -from core.workflow.human_input_policy import FormDisposition, enrich_human_input_pause_reasons -from core.workflow.nodes.human_input.boundary import build_human_input_pause_reason, human_input_container_selector -from core.workflow.nodes.human_input.entities import FormDefinition, SelectInputConfig, StringListSource -from core.workflow.nodes.human_input.enums import HumanInputFormKind, HumanInputFormStatus, ValueSourceType -from core.workflow.nodes.human_input.pause_reason import DifyHITLEventType -from graphon.entities.pause_reason import HitlRequired -from graphon.runtime import RuntimeState, VariablePool -from models.human_input import HumanInputForm - -_HUMAN_INPUT_REASON = {"TYPE": DifyHITLEventType.HUMAN_INPUT_REQUIRED, "form_id": "f1"} - - -@pytest.mark.parametrize( - ("dispositions", "expected_token", "expected_channels"), - [ - ({"f1": FormDisposition(form_token=None, approval_channels=["console", "email"])}, None, ["console", "email"]), - ({"f1": FormDisposition(form_token="tok", approval_channels=[])}, "tok", []), - # form_id absent from the map (no recipient rows) falls back to no token, no channels. - ({}, None, []), - ], -) -def test_enrich_projects_disposition_onto_reason(dispositions, expected_token, expected_channels): - out = enrich_human_input_pause_reasons( - [dict(_HUMAN_INPUT_REASON)], - dispositions_by_form_id=dispositions, - expiration_times_by_form_id={}, - ) - - assert out[0]["form_token"] == expected_token - assert out[0]["approval_channels"] == expected_channels - - -def test_enrich_leaves_non_human_input_reasons_untouched(): - reason = {"TYPE": "something_else", "form_id": "f1"} - - out = enrich_human_input_pause_reasons( - [reason], - dispositions_by_form_id={"f1": FormDisposition(form_token="tok", approval_channels=["email"])}, - expiration_times_by_form_id={}, - ) - - assert out[0] == reason - assert "form_token" not in out[0] - assert "approval_channels" not in out[0] - - -def test_pause_reason_payload_carries_approval_channels_through_factory(): - # from_response_data maps fields by hand; this guards approval_channels/form_token - # (the fields this feature added) against being dropped in that mapping. - from core.app.entities.task_entities import ( - HumanInputRequiredPauseReasonPayload, - HumanInputRequiredResponse, - ) - - data = HumanInputRequiredResponse.Data( - form_id="f", - node_id="n", - node_title="t", - form_content="c", - expiration_time=123, - form_token=None, - approval_channels=["console"], - ) - payload = HumanInputRequiredPauseReasonPayload.from_response_data(data) - - assert payload.approval_channels == ["console"] - assert payload.form_token is None - - -def test_build_pause_reason_keeps_constant_options_from_child_form(): - definition = FormDefinition( - form_content="Choose", - rendered_content="Choose", - inputs=[ - SelectInputConfig( - output_variable_name="decision", - option_source=StringListSource(type=ValueSourceType.CONSTANT, value=[]), - ) - ], - node_title="Choose", - expiration_time=datetime(2026, 1, 2), - ) - record = HumanInputFormRecord.from_models( - HumanInputForm( - id="form-123", - tenant_id="tenant", - app_id="app", - workflow_run_id="run", - node_id="human-input", - form_kind=HumanInputFormKind.RUNTIME, - form_definition=definition.model_dump_json(), - rendered_content=definition.rendered_content, - created_at=datetime(2026, 1, 1), - expiration_time=definition.expiration_time, - status=HumanInputFormStatus.WAITING, - ), - None, - ) - parent_pool = VariablePool() - parent_pool.add(("start", "options"), ["wrong"]) - - reason = build_human_input_pause_reason( - reason=HitlRequired(session_id="form-123", node_id="human-input", node_title="Choose"), - record=record, - variable_pool=parent_pool, - ) - - assert isinstance(reason.inputs[0], SelectInputConfig) - assert reason.inputs[0].option_source.type == ValueSourceType.CONSTANT - assert reason.inputs[0].option_source.value == [] - - -def test_child_forms_keep_distinct_identity_when_projected_to_the_same_visible_tool(): - definition = FormDefinition( - form_content="Approval", - rendered_content="Approval", - node_title="Approval", - expiration_time=datetime(2026, 1, 2), - ) - state = RuntimeState(workflow_id="workflow", variable_pool=VariablePool(), start_at=0) - reasons = [ - HitlRequired(session_id=form_id, node_id="source-human", node_title="Approval") - for form_id in ("form-a", "form-b") - ] - for reason in reasons: - state.variable_pool.add(human_input_container_selector(reason.session_id), "visible-tool") - restored = RuntimeState.from_snapshot(state.dumps()) - - enriched = [ - build_human_input_pause_reason( - reason=reason, - record=HumanInputFormRecord.from_models( - HumanInputForm( - id=reason.session_id, - tenant_id="tenant", - app_id="app", - workflow_run_id="run", - node_id="source-human", - form_kind=HumanInputFormKind.RUNTIME, - form_definition=definition.model_dump_json(), - rendered_content=reason.session_id, - created_at=datetime(2026, 1, 1), - expiration_time=definition.expiration_time, - status=HumanInputFormStatus.WAITING, - ), - None, - ), - variable_pool=restored.variable_pool, - ) - for reason in reasons - ] - - assert [(reason.form_id, reason.node_id) for reason in enriched] == [ - ("form-a", "visible-tool"), - ("form-b", "visible-tool"), - ] - assert all(reason.node_id == "source-human" for reason in reasons) diff --git a/api/tests/unit_tests/core/workflow/test_human_input_adapter.py b/api/tests/unit_tests/core/workflow/test_human_input_adapter.py index cdb374e207a..dd394327823 100644 --- a/api/tests/unit_tests/core/workflow/test_human_input_adapter.py +++ b/api/tests/unit_tests/core/workflow/test_human_input_adapter.py @@ -21,6 +21,81 @@ from graphon.nodes.base.variable_template_parser import VariableTemplateParser from graphon.runtime import VariablePool +def test_render_body_template_replaces_variable_values(): + config = EmailDeliveryConfig( + recipients=EmailRecipients(), + subject="Subject", + body="Hello {{#node1.value#}} {{#url#}}", + ) + variable_pool = VariablePool() + variable_pool.add(["node1", "value"], "World") + + result = config.render_body_template(body=config.body, url="https://example.com", variable_pool=variable_pool) + + assert result == "Hello World https://example.com" + + +def test_render_markdown_body_renders_markdown_to_html(): + rendered = EmailDeliveryConfig.render_markdown_body("**Bold** and [link](https://example.com)") + + assert "Bold" in rendered + assert 'link' in rendered + + +def test_render_markdown_body_sanitizes_unsafe_html(): + rendered = EmailDeliveryConfig.render_markdown_body( + 'Click' + ) + + assert "bad" in rendered + assert 'ok' in rendered + + +def test_render_markdown_body_does_not_allow_raw_html_tags(): + rendered = EmailDeliveryConfig.render_markdown_body("raw html and **markdown**") + + assert "" not in rendered + assert "raw html" in rendered + assert "markdown" in rendered + + +def test_render_markdown_body_supports_table_syntax(): + rendered = EmailDeliveryConfig.render_markdown_body("| h1 | h2 |\n| --- | ---: |\n| v1 | v2 |") + + assert "
" in rendered + assert "" in rendered + assert "" in rendered + assert 'align="right"' in rendered + assert "style=" not in rendered + + +def test_sanitize_subject_removes_crlf(): + sanitized = EmailDeliveryConfig.sanitize_subject("Notice\r\nBCC:attacker@example.com") + + assert "\r" not in sanitized + assert "\n" not in sanitized + assert sanitized == "Notice BCC:attacker@example.com" + + +def test_sanitize_subject_removes_html_tags(): + sanitized = EmailDeliveryConfig.sanitize_subject("Alert") + + assert "<" not in sanitized + assert ">" not in sanitized + assert sanitized == "Alert" + + def test_adapt_legacy_variable_aggregator_without_changing_saved_data() -> None: data = {"type": "variable-assigner", "output_type": "string", "variables": [["source", "value"]]} diff --git a/api/tests/unit_tests/core/workflow/test_node_factory.py b/api/tests/unit_tests/core/workflow/test_node_factory.py index aec0ad64db6..124f25505d6 100644 --- a/api/tests/unit_tests/core/workflow/test_node_factory.py +++ b/api/tests/unit_tests/core/workflow/test_node_factory.py @@ -1,4 +1,5 @@ -from collections.abc import Mapping +import types +from collections.abc import Callable, Mapping from types import SimpleNamespace from unittest.mock import MagicMock, patch, sentinel @@ -6,13 +7,20 @@ import pytest from sqlalchemy import Engine from sqlalchemy.orm import Session, sessionmaker -from core.app.entities.app_invoke_entities import DIFY_RUN_CONTEXT_KEY, DifyRunContext, InvokeFrom, UserFrom +from core.app.entities.app_invoke_entities import ( + DIFY_RUN_CONTEXT_KEY, + DifyRunContext, + InvokeFrom, + UserFrom, + build_dify_run_context, +) from core.plugin.impl.model import PluginModelClient from core.plugin.impl.model_runtime import PluginModelRuntime from core.plugin.plugin_service import PluginService from core.workflow import node_factory from core.workflow import template_rendering as workflow_template_rendering from core.workflow.llm_node import DifyLLMNode +from core.workflow.node_factory import DifyNodeFactory, get_node_type_classes_mapping from core.workflow.node_runtime import DifyPreparedLLM from core.workflow.nodes.knowledge_index import KNOWLEDGE_INDEX_NODE_TYPE from graphon.entities.base_node_data import BaseNodeData @@ -22,27 +30,258 @@ from graphon.graph.validation import GraphValidationError from graphon.model_runtime.entities.common_entities import I18nObject from graphon.model_runtime.entities.model_entities import AIModelEntity, FetchFrom, ModelFeature, ModelType from graphon.model_runtime.model_providers.base.large_language_model import LargeLanguageModel +from graphon.nodes.base.node import Node from graphon.nodes.code.entities import CodeLanguage from graphon.nodes.llm.entities import LLMNodeData from graphon.nodes.llm.node import LLMNode from graphon.nodes.llm.runtime_protocols import LLMPollingCapableProtocol from graphon.nodes.parameter_extractor.entities import ParameterExtractorNodeData +from graphon.nodes.variable_assigner.v1.node import ( + VariableAssignerNode as VariableAssignerV1, +) +from graphon.nodes.variable_assigner.v2.node import ( + VariableAssignerNode as VariableAssignerV2, +) from graphon.runtime import RuntimeState, VariablePool from graphon.variables.segments import ArrayObjectSegment, ObjectSegment, StringSegment from models.base import TypeBase from models.model import AppMode, Conversation, ConversationFromSource -from tests.unit_tests.core.model_fixtures import make_model_instance +from tests.helpers.unit.core.model_fixtures import make_model_instance from tests.workflow_test_utils import build_test_graph_init_params -@pytest.fixture -def memory_session_maker(monkeypatch: pytest.MonkeyPatch, sqlite_engine: Engine) -> sessionmaker[Session]: - """Bind node memory lookup to an explicit SQLite session factory.""" +class DummyNode: + def __init__(self, *, node_id, data, init_params, runtime_state, **kwargs): + self.id = node_id + self.data = data + self.init_params = init_params + self.runtime_state = runtime_state + self.kwargs = kwargs - TypeBase.metadata.create_all(sqlite_engine, tables=[Conversation.__table__]) - session_maker = sessionmaker(sqlite_engine, expire_on_commit=False) - monkeypatch.setattr(node_factory.session_factory, "create_session", session_maker) - return session_maker + +class DummyCodeNode(DummyNode): + @classmethod + def default_code_providers(cls): + return () + + +class DummyTemplateTransformNode(DummyNode): + pass + + +class DummyHttpRequestNode(DummyNode): + pass + + +class DummyKnowledgeRetrievalNode(DummyNode): + pass + + +class DummyDocumentExtractorNode(DummyNode): + pass + + +class TestNodeFactoryCases: + class TestDifyNodeFactory: + @pytest.fixture(autouse=True) + def _node_config(self, config_overrides: Callable[..., None]) -> None: + config_overrides( + CODE_MAX_STRING_LENGTH=10, + CODE_MAX_NUMBER=10, + CODE_MIN_NUMBER=-10, + CODE_MAX_PRECISION=4, + CODE_MAX_DEPTH=2, + CODE_MAX_NUMBER_ARRAY_LENGTH=2, + CODE_MAX_STRING_ARRAY_LENGTH=2, + CODE_MAX_OBJECT_ARRAY_LENGTH=2, + TEMPLATE_TRANSFORM_MAX_LENGTH=100, + UNSTRUCTURED_API_URL="http://u", + UNSTRUCTURED_API_KEY="key", + ) + + @staticmethod + def _stub_node_resolution(monkeypatch: pytest.MonkeyPatch, node_class): + monkeypatch.setattr( + "core.workflow.node_factory.resolve_workflow_node_class", + lambda **_kwargs: node_class, + ) + + def _factory(self): + run_context = build_dify_run_context( + tenant_id="tenant", + app_id="app", + user_id="user", + user_from=UserFrom.END_USER, + invoke_from=InvokeFrom.WEB_APP, + ) + + return DifyNodeFactory( + init_params=SimpleNamespace(run_context=run_context), + runtime_state=SimpleNamespace(), + ) + + def test_create_node_unknown_type(self): + factory = self._factory() + + with pytest.raises(ValueError, match="No class mapping found for node type: unknown"): + factory.create_node({"id": "node-1", "data": {"type": "unknown"}}) + + def test_create_node_missing_mapping(self, monkeypatch: pytest.MonkeyPatch): + factory = self._factory() + monkeypatch.setattr("core.workflow.node_factory.get_node_type_classes_mapping", lambda: {}) + + with pytest.raises(ValueError, match="No class mapping found for node type: start"): + factory.create_node({"id": "node-1", "data": {"type": BuiltinNodeTypes.START}}) + + def test_create_node_missing_latest_class(self, monkeypatch: pytest.MonkeyPatch): + factory = self._factory() + monkeypatch.setattr( + "core.workflow.node_factory.get_node_type_classes_mapping", + lambda: {BuiltinNodeTypes.START: {"1": None}}, + ) + monkeypatch.setattr("core.workflow.node_factory.LATEST_VERSION", "latest") + + with pytest.raises(ValueError, match="No latest version class found for node type: start"): + factory.create_node({"id": "node-1", "data": {"type": BuiltinNodeTypes.START}}) + + def test_create_node_selects_versioned_class(self, monkeypatch: pytest.MonkeyPatch): + factory = self._factory() + selected_versions: list[tuple[str, str]] = [] + + class DummyNodeV2(DummyNode): + pass + + def _get_mapping(): + selected_versions.append(("snapshot", "called")) + return {BuiltinNodeTypes.START: {"1": DummyNode, "2": DummyNodeV2}} + + monkeypatch.setattr("core.workflow.node_factory.get_node_type_classes_mapping", _get_mapping) + + node = factory.create_node({"id": "node-1", "data": {"type": BuiltinNodeTypes.START, "version": "2"}}) + + assert isinstance(node, DummyNodeV2) + assert node.id == "node-1" + assert selected_versions == [("snapshot", "called")] + + def test_create_node_code_branch(self, monkeypatch: pytest.MonkeyPatch): + factory = self._factory() + self._stub_node_resolution(monkeypatch, DummyCodeNode) + + node = factory.create_node({"id": "node-1", "data": {"type": BuiltinNodeTypes.CODE}}) + + assert isinstance(node, DummyCodeNode) + assert node.id == "node-1" + + def test_create_node_template_transform_branch(self, monkeypatch: pytest.MonkeyPatch): + factory = self._factory() + self._stub_node_resolution(monkeypatch, DummyTemplateTransformNode) + + node = factory.create_node({"id": "node-1", "data": {"type": BuiltinNodeTypes.TEMPLATE_TRANSFORM}}) + + assert isinstance(node, DummyTemplateTransformNode) + assert "jinja2_template_renderer" in node.kwargs + + def test_create_node_http_request_branch(self, monkeypatch: pytest.MonkeyPatch): + factory = self._factory() + self._stub_node_resolution(monkeypatch, DummyHttpRequestNode) + + node = factory.create_node({"id": "node-1", "data": {"type": BuiltinNodeTypes.HTTP_REQUEST}}) + + assert isinstance(node, DummyHttpRequestNode) + assert "http_request_config" in node.kwargs + + def test_create_node_knowledge_retrieval_branch(self, monkeypatch: pytest.MonkeyPatch): + factory = self._factory() + self._stub_node_resolution(monkeypatch, DummyKnowledgeRetrievalNode) + + node = factory.create_node({"id": "node-1", "data": {"type": BuiltinNodeTypes.KNOWLEDGE_RETRIEVAL}}) + + assert isinstance(node, DummyKnowledgeRetrievalNode) + assert node.kwargs == {} + + def test_create_node_document_extractor_branch(self, monkeypatch: pytest.MonkeyPatch): + factory = self._factory() + self._stub_node_resolution(monkeypatch, DummyDocumentExtractorNode) + + node = factory.create_node({"id": "node-1", "data": {"type": BuiltinNodeTypes.DOCUMENT_EXTRACTOR}}) + + assert isinstance(node, DummyDocumentExtractorNode) + assert "unstructured_api_config" in node.kwargs + + +# Import concrete nodes we will assert on (numeric version path) + + +class TestGetNodeTypeClassesMappingCases: + @staticmethod + def test_variable_assigner_latest_prefers_highest_numeric_version(): + # Act + mapping: Mapping[NodeType, Mapping[str, type[Node]]] = get_node_type_classes_mapping() + + # Assert basic presence + assert BuiltinNodeTypes.VARIABLE_ASSIGNER in mapping + va_versions = mapping[BuiltinNodeTypes.VARIABLE_ASSIGNER] + + # Both concrete versions must be present + assert va_versions.get("1") is VariableAssignerV1 + assert va_versions.get("2") is VariableAssignerV2 + + # And latest should point to numerically-highest version ("2") + assert va_versions.get("latest") is VariableAssignerV2 + + @staticmethod + def test_latest_prefers_highest_numeric_version(): + # Arrange: define two ephemeral subclasses with numeric versions under a NodeType + # that has no concrete implementations in production to avoid interference. + class _Version1(Node[BaseNodeData]): # type: ignore[misc] + node_type = BuiltinNodeTypes.LEGACY_VARIABLE_AGGREGATOR + + def init_node_data(self, data): + pass + + def _run(self): + raise NotImplementedError + + @classmethod + def version(cls) -> str: + return "1" + + def _get_error_strategy(self): + return None + + def _get_retry_config(self): + return types.SimpleNamespace() # not used + + def _get_title(self) -> str: + return "version1" + + def _get_description(self): + return None + + def _get_default_value_dict(self): + return {} + + def get_base_node_data(self): + return types.SimpleNamespace(title="version1") + + class _Version2(_Version1): # type: ignore[misc] + @classmethod + def version(cls) -> str: + return "2" + + def _get_title(self) -> str: + return "version2" + + # Act: build a fresh mapping (it should now see our ephemeral subclasses) + mapping: Mapping[NodeType, Mapping[str, type[Node]]] = get_node_type_classes_mapping() + + # Assert: both numeric versions exist for this NodeType; 'latest' points to the higher numeric version + assert BuiltinNodeTypes.LEGACY_VARIABLE_AGGREGATOR in mapping + legacy_versions = mapping[BuiltinNodeTypes.LEGACY_VARIABLE_AGGREGATOR] + + assert legacy_versions.get("1") is _Version1 + assert legacy_versions.get("2") is _Version2 + assert legacy_versions.get("latest") is _Version2 def _persist_conversation(session_maker: sessionmaker[Session]) -> None: @@ -113,1188 +352,1202 @@ class _ModelInstanceStub: return self._model_schema -class TestResolveWorkflowNodeClass: - def test_matching_version_uses_registry_mapping(self, monkeypatch) -> None: - document_extractor_class = sentinel.document_extractor_class - latest_node_class = sentinel.latest_document_extractor_class - monkeypatch.setattr( - node_factory, - "get_node_type_classes_mapping", - lambda: { - BuiltinNodeTypes.DOCUMENT_EXTRACTOR: { - "1": document_extractor_class, - node_factory.LATEST_VERSION: latest_node_class, - } - }, +class TestNodeFactoryCasesWorkflow: + @staticmethod + @pytest.fixture + def memory_session_maker(monkeypatch: pytest.MonkeyPatch, sqlite_engine: Engine) -> sessionmaker[Session]: + """Bind node memory lookup to an explicit SQLite session factory.""" + + TypeBase.metadata.create_all(sqlite_engine, tables=[Conversation.__table__]) + session_maker = sessionmaker(sqlite_engine, expire_on_commit=False) + monkeypatch.setattr(node_factory.session_factory, "create_session", session_maker) + return session_maker + + class TestResolveWorkflowNodeClass: + def test_matching_version_uses_registry_mapping(self, monkeypatch) -> None: + document_extractor_class = sentinel.document_extractor_class + latest_node_class = sentinel.latest_document_extractor_class + monkeypatch.setattr( + node_factory, + "get_node_type_classes_mapping", + lambda: { + BuiltinNodeTypes.DOCUMENT_EXTRACTOR: { + "1": document_extractor_class, + node_factory.LATEST_VERSION: latest_node_class, + } + }, + ) + + resolved = node_factory.resolve_workflow_node_class( + node_type=BuiltinNodeTypes.DOCUMENT_EXTRACTOR, + node_version="1", + ) + + assert resolved is document_extractor_class + + def test_document_extractor_latest_falls_back_to_registry_mapping(self, monkeypatch) -> None: + latest_node_class = sentinel.latest_document_extractor_class + monkeypatch.setattr( + node_factory, + "get_node_type_classes_mapping", + lambda: {BuiltinNodeTypes.DOCUMENT_EXTRACTOR: {node_factory.LATEST_VERSION: latest_node_class}}, + ) + + resolved = node_factory.resolve_workflow_node_class( + node_type=BuiltinNodeTypes.DOCUMENT_EXTRACTOR, + node_version=node_factory.LATEST_VERSION, + ) + + assert resolved is latest_node_class + + class TestFetchMemory: + @pytest.mark.parametrize( + ("conversation_id", "memory_config"), + [ + (None, object()), + ("conversation-id", None), + ], ) + def test_returns_none_when_memory_or_conversation_is_missing(self, conversation_id, memory_config): + result = node_factory.fetch_memory( + conversation_id=conversation_id, + app_id="app-id", + node_data_memory=memory_config, + model_instance=sentinel.model_instance, + ) - resolved = node_factory.resolve_workflow_node_class( - node_type=BuiltinNodeTypes.DOCUMENT_EXTRACTOR, - node_version="1", - ) + assert result is None - assert resolved is document_extractor_class + def test_returns_none_when_conversation_does_not_exist(self, memory_session_maker: sessionmaker[Session]): + result = node_factory.fetch_memory( + conversation_id="conversation-id", + app_id="app-id", + node_data_memory=object(), + model_instance=sentinel.model_instance, + ) - def test_document_extractor_latest_falls_back_to_registry_mapping(self, monkeypatch) -> None: - latest_node_class = sentinel.latest_document_extractor_class - monkeypatch.setattr( - node_factory, - "get_node_type_classes_mapping", - lambda: {BuiltinNodeTypes.DOCUMENT_EXTRACTOR: {node_factory.LATEST_VERSION: latest_node_class}}, - ) + assert result is None - resolved = node_factory.resolve_workflow_node_class( - node_type=BuiltinNodeTypes.DOCUMENT_EXTRACTOR, - node_version=node_factory.LATEST_VERSION, - ) + def test_builds_token_buffer_memory_for_existing_conversation( + self, monkeypatch: pytest.MonkeyPatch, memory_session_maker: sessionmaker[Session] + ): + memory = sentinel.memory + _persist_conversation(memory_session_maker) + token_buffer_memory = MagicMock(return_value=memory) + monkeypatch.setattr(node_factory, "TokenBufferMemory", token_buffer_memory) - assert resolved is latest_node_class + result = node_factory.fetch_memory( + conversation_id="conversation-id", + app_id="app-id", + node_data_memory=object(), + model_instance=sentinel.model_instance, + ) + assert result is memory + loaded_conversation = token_buffer_memory.call_args.kwargs["conversation"] + assert isinstance(loaded_conversation, Conversation) + assert loaded_conversation.id == "conversation-id" + assert token_buffer_memory.call_args.kwargs["model_instance"] is sentinel.model_instance -class TestFetchMemory: + def test_uses_configured_session_factory_without_flask_app_context( + self, monkeypatch: pytest.MonkeyPatch, memory_session_maker: sessionmaker[Session] + ): + class RaisingDB: + @property + def engine(self): + raise RuntimeError("Working outside of application context.") + + token_buffer_memory = MagicMock(return_value=sentinel.memory) + _persist_conversation(memory_session_maker) + monkeypatch.setattr(node_factory, "db", RaisingDB(), raising=False) + monkeypatch.setattr(node_factory, "TokenBufferMemory", token_buffer_memory) + + result = node_factory.fetch_memory( + conversation_id="conversation-id", + app_id="app-id", + node_data_memory=object(), + model_instance=sentinel.model_instance, + ) + + assert result is sentinel.memory + + class TestDifyGraphInitContext: + def test_to_graph_init_params_preserves_explicit_values(self): + run_context = { + DIFY_RUN_CONTEXT_KEY: DifyRunContext( + tenant_id="tenant-id", + app_id="app-id", + user_id="user-id", + user_from=UserFrom.ACCOUNT, + invoke_from=InvokeFrom.DEBUGGER, + ), + "extra": "value", + } + graph_config = {"nodes": [], "edges": []} + graph_init_context = node_factory.DifyGraphInitContext( + workflow_id="workflow-id", + graph_config=graph_config, + run_context=run_context, + call_depth=2, + ) + + result = graph_init_context.to_graph_init_params() + + assert result.workflow_id == "workflow-id" + assert result.graph_config == graph_config + assert result.run_context == run_context + assert result.call_depth == 2 + + @staticmethod @pytest.mark.parametrize( - ("conversation_id", "memory_config"), - [ - (None, object()), - ("conversation-id", None), - ], + "invalid_kind", + [None, "CYCLE", "MISSING_NODE", "INVALID_EDGE", "DUPLICATE_NODE_ID", "DUPLICATE_EDGE_ID"], ) - def test_returns_none_when_memory_or_conversation_is_missing(self, conversation_id, memory_config): - result = node_factory.fetch_memory( - conversation_id=conversation_id, - app_id="app-id", - node_data_memory=memory_config, - model_instance=sentinel.model_instance, + def test_graph_validation_with_dify_node_factory(invalid_kind, monkeypatch): + graph_config = { + "nodes": [ + {"id": "start", "data": {"type": "start", "title": "Start", "variables": []}}, + *[ + { + "id": node_id, + "data": {"type": "template-transform", "title": node_id, "template": node_id, "variables": []}, + } + for node_id in ("left", "right") + ], + {"id": "end", "data": {"type": "end", "title": "End", "outputs": []}}, + ], + "edges": [ + {"id": f"{source}-{target}", "source": source, "target": target} + for source, target in (("start", "left"), ("start", "right"), ("left", "end"), ("right", "end")) + ], + } + match invalid_kind: + case "CYCLE": + graph_config["edges"].append({"source": "end", "target": "left"}) + case "MISSING_NODE": + graph_config["edges"][0]["target"] = "missing" + case "INVALID_EDGE": + graph_config["edges"][0]["sourceHandle"] = None + case "DUPLICATE_NODE_ID": + graph_config["nodes"].append(graph_config["nodes"][0]) + case "DUPLICATE_EDGE_ID": + graph_config["edges"].append(graph_config["edges"][0]) + + factory = node_factory.DifyNodeFactory( + build_test_graph_init_params(graph_config=graph_config), + RuntimeState(variable_pool=VariablePool(), start_at=0, workflow_id="workflow"), ) + if invalid_kind is None: + graph = Graph.init(graph_config=graph_config, node_factory=factory, root_node_id="start") + assert set(graph.nodes) == {"start", "left", "right", "end"} + assert len(graph.edges) == 4 + else: - assert result is None + def reject_node_creation(*_args, **_kwargs): + pytest.fail("Invalid graphs must fail before node construction") - def test_returns_none_when_conversation_does_not_exist(self, memory_session_maker: sessionmaker[Session]): - result = node_factory.fetch_memory( - conversation_id="conversation-id", - app_id="app-id", - node_data_memory=object(), - model_instance=sentinel.model_instance, - ) + monkeypatch.setattr(node_factory.DifyNodeFactory, "create_node", reject_node_creation) + with pytest.raises(GraphValidationError) as error: + Graph.init(graph_config=graph_config, node_factory=factory, root_node_id="start") + assert [issue.code for issue in error.value.issues] == [invalid_kind] - assert result is None + class TestDefaultWorkflowCodeExecutor: + def test_execute_delegates_to_code_executor(self, monkeypatch: pytest.MonkeyPatch): + executor = node_factory.DefaultWorkflowCodeExecutor() + execute_workflow_code_template = MagicMock(return_value={"answer": "ok"}) + monkeypatch.setattr( + node_factory.CodeExecutor, + "execute_workflow_code_template", + execute_workflow_code_template, + ) - def test_builds_token_buffer_memory_for_existing_conversation( - self, monkeypatch: pytest.MonkeyPatch, memory_session_maker: sessionmaker[Session] - ): - memory = sentinel.memory - _persist_conversation(memory_session_maker) - token_buffer_memory = MagicMock(return_value=memory) - monkeypatch.setattr(node_factory, "TokenBufferMemory", token_buffer_memory) + result = executor.execute( + language=CodeLanguage.PYTHON3, + code="print('ok')", + inputs={"name": "workflow"}, + ) - result = node_factory.fetch_memory( - conversation_id="conversation-id", - app_id="app-id", - node_data_memory=object(), - model_instance=sentinel.model_instance, - ) + assert result == {"answer": "ok"} + execute_workflow_code_template.assert_called_once_with( + language=CodeLanguage.PYTHON3, + code="print('ok')", + inputs={"name": "workflow"}, + ) - assert result is memory - loaded_conversation = token_buffer_memory.call_args.kwargs["conversation"] - assert isinstance(loaded_conversation, Conversation) - assert loaded_conversation.id == "conversation-id" - assert token_buffer_memory.call_args.kwargs["model_instance"] is sentinel.model_instance + def test_is_execution_error_checks_code_execution_error_type(self): + executor = node_factory.DefaultWorkflowCodeExecutor() - def test_uses_configured_session_factory_without_flask_app_context( - self, monkeypatch: pytest.MonkeyPatch, memory_session_maker: sessionmaker[Session] - ): - class RaisingDB: - @property - def engine(self): - raise RuntimeError("Working outside of application context.") + assert executor.is_execution_error(node_factory.CodeExecutionError("boom")) is True + assert executor.is_execution_error(RuntimeError("boom")) is False - token_buffer_memory = MagicMock(return_value=sentinel.memory) - _persist_conversation(memory_session_maker) - monkeypatch.setattr(node_factory, "db", RaisingDB(), raising=False) - monkeypatch.setattr(node_factory, "TokenBufferMemory", token_buffer_memory) + class TestCodeExecutorJinja2TemplateRenderer: + def test_render_template_delegates_to_code_executor(self, monkeypatch: pytest.MonkeyPatch): + renderer = workflow_template_rendering.CodeExecutorJinja2TemplateRenderer() + execute_workflow_code_template = MagicMock(return_value={"result": "Hello workflow"}) + monkeypatch.setattr( + workflow_template_rendering.CodeExecutor, + "execute_workflow_code_template", + execute_workflow_code_template, + ) - result = node_factory.fetch_memory( - conversation_id="conversation-id", - app_id="app-id", - node_data_memory=object(), - model_instance=sentinel.model_instance, - ) + result = renderer.render_template("Hello {{ name }}", {"name": "workflow"}) - assert result is sentinel.memory + assert result == "Hello workflow" + execute_workflow_code_template.assert_called_once_with( + language=CodeLanguage.JINJA2, + code="Hello {{ name }}", + inputs={"name": "workflow"}, + ) + def test_render_template_wraps_code_execution_errors(self, monkeypatch: pytest.MonkeyPatch): + renderer = workflow_template_rendering.CodeExecutorJinja2TemplateRenderer() + monkeypatch.setattr( + workflow_template_rendering.CodeExecutor, + "execute_workflow_code_template", + MagicMock(side_effect=workflow_template_rendering.CodeExecutionError("sandbox failed")), + ) -class TestDifyGraphInitContext: - def test_to_graph_init_params_preserves_explicit_values(self): - run_context = { - DIFY_RUN_CONTEXT_KEY: DifyRunContext( + with pytest.raises(workflow_template_rendering.TemplateRenderError, match="sandbox failed"): + renderer.render_template("{{ broken }}", {}) + + class TestDifyNodeFactoryResolveContext: + def test_requires_reserved_context_key(self): + with pytest.raises(ValueError, match=DIFY_RUN_CONTEXT_KEY): + node_factory.DifyNodeFactory._resolve_dify_context({}) + + def test_returns_existing_dify_context(self): + dify_context = DifyRunContext( tenant_id="tenant-id", app_id="app-id", user_id="user-id", user_from=UserFrom.ACCOUNT, invoke_from=InvokeFrom.DEBUGGER, - ), - "extra": "value", - } - graph_config = {"nodes": [], "edges": []} - graph_init_context = node_factory.DifyGraphInitContext( - workflow_id="workflow-id", - graph_config=graph_config, - run_context=run_context, - call_depth=2, - ) + ) - result = graph_init_context.to_graph_init_params() + result = node_factory.DifyNodeFactory._resolve_dify_context({DIFY_RUN_CONTEXT_KEY: dify_context}) - assert result.workflow_id == "workflow-id" - assert result.graph_config == graph_config - assert result.run_context == run_context - assert result.call_depth == 2 + assert result is dify_context - -@pytest.mark.parametrize( - "invalid_kind", - [None, "CYCLE", "MISSING_NODE", "INVALID_EDGE", "DUPLICATE_NODE_ID", "DUPLICATE_EDGE_ID"], -) -def test_graph_validation_with_dify_node_factory(invalid_kind, monkeypatch): - graph_config = { - "nodes": [ - {"id": "start", "data": {"type": "start", "title": "Start", "variables": []}}, - *[ - { - "id": node_id, - "data": {"type": "template-transform", "title": node_id, "template": node_id, "variables": []}, + def test_validates_mapping_context(self): + raw_context = { + DIFY_RUN_CONTEXT_KEY: { + "tenant_id": "tenant-id", + "app_id": "app-id", + "user_id": "user-id", + "user_from": UserFrom.ACCOUNT, + "invoke_from": InvokeFrom.DEBUGGER, } - for node_id in ("left", "right") - ], - {"id": "end", "data": {"type": "end", "title": "End", "outputs": []}}, - ], - "edges": [ - {"id": f"{source}-{target}", "source": source, "target": target} - for source, target in (("start", "left"), ("start", "right"), ("left", "end"), ("right", "end")) - ], - } - match invalid_kind: - case "CYCLE": - graph_config["edges"].append({"source": "end", "target": "left"}) - case "MISSING_NODE": - graph_config["edges"][0]["target"] = "missing" - case "INVALID_EDGE": - graph_config["edges"][0]["sourceHandle"] = None - case "DUPLICATE_NODE_ID": - graph_config["nodes"].append(graph_config["nodes"][0]) - case "DUPLICATE_EDGE_ID": - graph_config["edges"].append(graph_config["edges"][0]) - - factory = node_factory.DifyNodeFactory( - build_test_graph_init_params(graph_config=graph_config), - RuntimeState(variable_pool=VariablePool(), start_at=0, workflow_id="workflow"), - ) - if invalid_kind is None: - graph = Graph.init(graph_config=graph_config, node_factory=factory, root_node_id="start") - assert set(graph.nodes) == {"start", "left", "right", "end"} - assert len(graph.edges) == 4 - else: - - def reject_node_creation(*_args, **_kwargs): - pytest.fail("Invalid graphs must fail before node construction") - - monkeypatch.setattr(node_factory.DifyNodeFactory, "create_node", reject_node_creation) - with pytest.raises(GraphValidationError) as error: - Graph.init(graph_config=graph_config, node_factory=factory, root_node_id="start") - assert [issue.code for issue in error.value.issues] == [invalid_kind] - - -class TestDefaultWorkflowCodeExecutor: - def test_execute_delegates_to_code_executor(self, monkeypatch: pytest.MonkeyPatch): - executor = node_factory.DefaultWorkflowCodeExecutor() - execute_workflow_code_template = MagicMock(return_value={"answer": "ok"}) - monkeypatch.setattr( - node_factory.CodeExecutor, - "execute_workflow_code_template", - execute_workflow_code_template, - ) - - result = executor.execute( - language=CodeLanguage.PYTHON3, - code="print('ok')", - inputs={"name": "workflow"}, - ) - - assert result == {"answer": "ok"} - execute_workflow_code_template.assert_called_once_with( - language=CodeLanguage.PYTHON3, - code="print('ok')", - inputs={"name": "workflow"}, - ) - - def test_is_execution_error_checks_code_execution_error_type(self): - executor = node_factory.DefaultWorkflowCodeExecutor() - - assert executor.is_execution_error(node_factory.CodeExecutionError("boom")) is True - assert executor.is_execution_error(RuntimeError("boom")) is False - - -class TestCodeExecutorJinja2TemplateRenderer: - def test_render_template_delegates_to_code_executor(self, monkeypatch: pytest.MonkeyPatch): - renderer = workflow_template_rendering.CodeExecutorJinja2TemplateRenderer() - execute_workflow_code_template = MagicMock(return_value={"result": "Hello workflow"}) - monkeypatch.setattr( - workflow_template_rendering.CodeExecutor, - "execute_workflow_code_template", - execute_workflow_code_template, - ) - - result = renderer.render_template("Hello {{ name }}", {"name": "workflow"}) - - assert result == "Hello workflow" - execute_workflow_code_template.assert_called_once_with( - language=CodeLanguage.JINJA2, - code="Hello {{ name }}", - inputs={"name": "workflow"}, - ) - - def test_render_template_wraps_code_execution_errors(self, monkeypatch: pytest.MonkeyPatch): - renderer = workflow_template_rendering.CodeExecutorJinja2TemplateRenderer() - monkeypatch.setattr( - workflow_template_rendering.CodeExecutor, - "execute_workflow_code_template", - MagicMock(side_effect=workflow_template_rendering.CodeExecutionError("sandbox failed")), - ) - - with pytest.raises(workflow_template_rendering.TemplateRenderError, match="sandbox failed"): - renderer.render_template("{{ broken }}", {}) - - -class TestDifyNodeFactoryResolveContext: - def test_requires_reserved_context_key(self): - with pytest.raises(ValueError, match=DIFY_RUN_CONTEXT_KEY): - node_factory.DifyNodeFactory._resolve_dify_context({}) - - def test_returns_existing_dify_context(self): - dify_context = DifyRunContext( - tenant_id="tenant-id", - app_id="app-id", - user_id="user-id", - user_from=UserFrom.ACCOUNT, - invoke_from=InvokeFrom.DEBUGGER, - ) - - result = node_factory.DifyNodeFactory._resolve_dify_context({DIFY_RUN_CONTEXT_KEY: dify_context}) - - assert result is dify_context - - def test_validates_mapping_context(self): - raw_context = { - DIFY_RUN_CONTEXT_KEY: { - "tenant_id": "tenant-id", - "app_id": "app-id", - "user_id": "user-id", - "user_from": UserFrom.ACCOUNT, - "invoke_from": InvokeFrom.DEBUGGER, } - } - result = node_factory.DifyNodeFactory._resolve_dify_context(raw_context) + result = node_factory.DifyNodeFactory._resolve_dify_context(raw_context) - assert isinstance(result, DifyRunContext) - assert result.tenant_id == "tenant-id" + assert isinstance(result, DifyRunContext) + assert result.tenant_id == "tenant-id" + class TestDifyNodeFactoryCreateNode: + @pytest.fixture + def factory(self): + factory = object.__new__(node_factory.DifyNodeFactory) + factory.init_params = sentinel.init_params + factory.runtime_state = SimpleNamespace(variable_pool=MagicMock()) + factory._dify_context = SimpleNamespace( + tenant_id="tenant-id", + app_id="app-id", + user_id="user-id", + invoke_from=InvokeFrom.DEBUGGER, + app_type=None, + created_by=None, + ) + factory._use_workflow_tool_containers = True + factory._code_executor = sentinel.code_executor + factory._code_limits = sentinel.code_limits + factory._jinja2_template_renderer = sentinel.jinja2_template_renderer + factory._template_transform_max_output_length = 2048 + factory._http_request_http_client = sentinel.http_client + factory._remote_file_http_client = sentinel.remote_file_http_client + factory._bound_tool_file_manager_factory = MagicMock(return_value=sentinel.tool_file_manager) + factory._file_reference_factory = sentinel.file_reference_factory + factory._prompt_message_serializer = sentinel.prompt_message_serializer + factory._retriever_attachment_loader = sentinel.retriever_attachment_loader + factory._llm_file_saver = sentinel.llm_file_saver + factory._human_input_runtime = sentinel.human_input_runtime + factory._tool_runtime = sentinel.tool_runtime + factory._http_request_file_manager = sentinel.file_manager + factory._document_extractor_unstructured_api_config = sentinel.unstructured_api_config + factory._http_request_config = sentinel.http_request_config + factory._llm_credentials_provider = sentinel.credentials_provider + factory._llm_model_factory = sentinel.model_factory + factory._build_retriever_attachment_loader = MagicMock(return_value=sentinel.retriever_attachment_loader) + return factory -class TestDifyNodeFactoryCreateNode: - @pytest.fixture - def factory(self): - factory = object.__new__(node_factory.DifyNodeFactory) - factory.init_params = sentinel.init_params - factory.runtime_state = SimpleNamespace(variable_pool=MagicMock()) - factory._dify_context = SimpleNamespace( - tenant_id="tenant-id", - app_id="app-id", - user_id="user-id", - invoke_from=InvokeFrom.DEBUGGER, - app_type=None, - created_by=None, - ) - factory._use_workflow_tool_containers = True - factory._code_executor = sentinel.code_executor - factory._code_limits = sentinel.code_limits - factory._jinja2_template_renderer = sentinel.jinja2_template_renderer - factory._template_transform_max_output_length = 2048 - factory._http_request_http_client = sentinel.http_client - factory._remote_file_http_client = sentinel.remote_file_http_client - factory._bound_tool_file_manager_factory = MagicMock(return_value=sentinel.tool_file_manager) - factory._file_reference_factory = sentinel.file_reference_factory - factory._prompt_message_serializer = sentinel.prompt_message_serializer - factory._retriever_attachment_loader = sentinel.retriever_attachment_loader - factory._llm_file_saver = sentinel.llm_file_saver - factory._human_input_runtime = sentinel.human_input_runtime - factory._tool_runtime = sentinel.tool_runtime - factory._http_request_file_manager = sentinel.file_manager - factory._document_extractor_unstructured_api_config = sentinel.unstructured_api_config - factory._http_request_config = sentinel.http_request_config - factory._llm_credentials_provider = sentinel.credentials_provider - factory._llm_model_factory = sentinel.model_factory - factory._build_retriever_attachment_loader = MagicMock(return_value=sentinel.retriever_attachment_loader) - return factory + def test_rejects_unknown_node_type(self, factory): + with pytest.raises(ValueError, match="No class mapping found for node type: missing"): + factory.create_node({"id": "node-id", "data": {"type": "missing"}}) - def test_rejects_unknown_node_type(self, factory): - with pytest.raises(ValueError, match="No class mapping found for node type: missing"): - factory.create_node({"id": "node-id", "data": {"type": "missing"}}) - - def test_validate_node_resolves_schema_without_constructing_node(self, monkeypatch: pytest.MonkeyPatch, factory): - constructor = _node_constructor(return_value=sentinel.node) - constructor.execution_type = NodeExecutionType.BRANCH - monkeypatch.setattr(factory, "_resolve_node_class", MagicMock(return_value=constructor)) - - execution_type = factory.validate_node({"id": "node-id", "data": {"type": BuiltinNodeTypes.START}}) - - assert execution_type == NodeExecutionType.BRANCH - constructor.validate_node_data.assert_called_once() - constructor.assert_not_called() - - def test_validate_node_rejects_invalid_human_input_delivery_config( - self, - monkeypatch: pytest.MonkeyPatch, - factory, - ): - constructor = _node_constructor(return_value=sentinel.node) - monkeypatch.setattr(factory, "_resolve_node_class", MagicMock(return_value=constructor)) - - with pytest.raises( - ValueError, - match=r"(?s)0\.email\.config\.subject.*valid string.*0\.email\.config\.body.*valid string", + def test_validate_node_resolves_schema_without_constructing_node( + self, monkeypatch: pytest.MonkeyPatch, factory ): - factory.validate_node( + constructor = _node_constructor(return_value=sentinel.node) + constructor.execution_type = NodeExecutionType.BRANCH + monkeypatch.setattr(factory, "_resolve_node_class", MagicMock(return_value=constructor)) + + execution_type = factory.validate_node({"id": "node-id", "data": {"type": BuiltinNodeTypes.START}}) + + assert execution_type == NodeExecutionType.BRANCH + constructor.validate_node_data.assert_called_once() + constructor.assert_not_called() + + def test_validate_node_rejects_invalid_human_input_delivery_config( + self, + monkeypatch: pytest.MonkeyPatch, + factory, + ): + constructor = _node_constructor(return_value=sentinel.node) + monkeypatch.setattr(factory, "_resolve_node_class", MagicMock(return_value=constructor)) + + with pytest.raises( + ValueError, + match=r"(?s)0\.email\.config\.subject.*valid string.*0\.email\.config\.body.*valid string", + ): + factory.validate_node( + { + "id": "human-input", + "data": { + "type": BuiltinNodeTypes.HUMAN_INPUT, + "delivery_methods": [ + { + "type": "email", + "config": { + "recipients": {"items": []}, + "subject": 123, + "body": [], + }, + } + ], + }, + } + ) + + constructor.assert_not_called() + + def test_rejects_missing_class_mapping(self, monkeypatch: pytest.MonkeyPatch, factory): + monkeypatch.setattr( + factory, + "_resolve_node_class", + MagicMock(side_effect=ValueError("No class mapping found for node type: start")), + ) + + with pytest.raises(ValueError, match="No class mapping found for node type: start"): + factory.create_node({"id": "node-id", "data": {"type": BuiltinNodeTypes.START}}) + + def test_rejects_missing_latest_class(self, monkeypatch: pytest.MonkeyPatch, factory): + monkeypatch.setattr( + factory, + "_resolve_node_class", + MagicMock(side_effect=ValueError("No latest version class found for node type: start")), + ) + + with pytest.raises(ValueError, match="No latest version class found for node type: start"): + factory.create_node({"id": "node-id", "data": {"type": BuiltinNodeTypes.START}}) + + def test_uses_version_specific_class_when_available(self, monkeypatch: pytest.MonkeyPatch, factory): + matched_node = sentinel.matched_node + latest_node_class = _node_constructor(return_value=sentinel.latest_node) + matched_node_class = _node_constructor(return_value=matched_node) + monkeypatch.setattr( + factory, + "_resolve_node_class", + MagicMock(return_value=matched_node_class), + ) + + result = factory.create_node({"id": "node-id", "data": {"type": BuiltinNodeTypes.START, "version": "9"}}) + + assert result is matched_node + matched_node_class.assert_called_once() + kwargs = matched_node_class.call_args.kwargs + assert kwargs["node_id"] == "node-id" + _assert_constructor_node_data( + kwargs["data"], node_id="node-id", node_type=BuiltinNodeTypes.START, version="9" + ) + assert kwargs["init_params"] is sentinel.init_params + assert kwargs["runtime_state"] is factory.runtime_state + latest_node_class.assert_not_called() + + def test_falls_back_to_latest_class_when_version_specific_mapping_is_missing( + self, monkeypatch: pytest.MonkeyPatch, factory + ): + latest_node = sentinel.latest_node + latest_node_class = _node_constructor(return_value=latest_node) + monkeypatch.setattr( + factory, + "_resolve_node_class", + MagicMock(return_value=latest_node_class), + ) + + result = factory.create_node({"id": "node-id", "data": {"type": BuiltinNodeTypes.START, "version": "9"}}) + + assert result is latest_node + latest_node_class.assert_called_once() + kwargs = latest_node_class.call_args.kwargs + assert kwargs["node_id"] == "node-id" + _assert_constructor_node_data( + kwargs["data"], node_id="node-id", node_type=BuiltinNodeTypes.START, version="9" + ) + assert kwargs["init_params"] is sentinel.init_params + assert kwargs["runtime_state"] is factory.runtime_state + + @pytest.mark.parametrize( + ("node_type", "constructor_name"), + [ + (BuiltinNodeTypes.CODE, "CodeNode"), + (BuiltinNodeTypes.TEMPLATE_TRANSFORM, "TemplateTransformNode"), + (BuiltinNodeTypes.HTTP_REQUEST, "HttpRequestNode"), + (BuiltinNodeTypes.HUMAN_INPUT, "HumanInputNode"), + (BuiltinNodeTypes.TOOL, "ToolNode"), + (KNOWLEDGE_INDEX_NODE_TYPE, "KnowledgeIndexNode"), + (BuiltinNodeTypes.DATASOURCE, "DatasourceNode"), + (BuiltinNodeTypes.KNOWLEDGE_RETRIEVAL, "KnowledgeRetrievalNode"), + (BuiltinNodeTypes.DOCUMENT_EXTRACTOR, "DocumentExtractorNode"), + ], + ) + def test_creates_specialized_nodes(self, monkeypatch: pytest.MonkeyPatch, factory, node_type, constructor_name): + created_node = object() + constructor = _node_constructor(return_value=created_node) + constructor._mock_name = constructor_name + monkeypatch.setattr( + factory, + "_resolve_node_class", + MagicMock(return_value=constructor), + ) + + if constructor_name == "HumanInputNode": + factory._build_human_input_callback = MagicMock(return_value=sentinel.hitl_callback) + + node_config = {"id": "node-id", "data": {"type": node_type}} + result = factory.create_node(node_config) + + assert result is created_node + kwargs = constructor.call_args.kwargs + assert kwargs["node_id"] == "node-id" + _assert_constructor_node_data(kwargs["data"], node_id="node-id", node_type=node_type) + assert kwargs["init_params"] is sentinel.init_params + assert kwargs["runtime_state"] is factory.runtime_state + + if constructor_name == "CodeNode": + assert kwargs["code_executor"] is sentinel.code_executor + assert kwargs["code_limits"] is sentinel.code_limits + elif constructor_name == "TemplateTransformNode": + assert kwargs["jinja2_template_renderer"] is sentinel.jinja2_template_renderer + assert kwargs["max_output_length"] == 2048 + elif constructor_name == "HttpRequestNode": + assert kwargs["http_request_config"] is sentinel.http_request_config + assert kwargs["http_client"] is sentinel.http_client + assert kwargs["tool_file_manager_factory"] is factory._bound_tool_file_manager_factory + assert kwargs["file_manager"] is sentinel.file_manager + assert kwargs["file_reference_factory"] is sentinel.file_reference_factory + factory._bound_tool_file_manager_factory.assert_not_called() + elif constructor_name == "HumanInputNode": + assert kwargs["hitl_callback"] is sentinel.hitl_callback + factory._build_human_input_callback.assert_called_once() + elif constructor_name == "ToolNode": + assert kwargs["tool_file_manager"] is sentinel.tool_file_manager + assert kwargs["runtime"] is sentinel.tool_runtime + factory._bound_tool_file_manager_factory.assert_called_once_with() + elif constructor_name == "DocumentExtractorNode": + assert kwargs["unstructured_api_config"] is sentinel.unstructured_api_config + assert kwargs["http_client"] is sentinel.remote_file_http_client + + def test_human_input_node_receives_built_hitl_callback( + self, + monkeypatch: pytest.MonkeyPatch, + factory, + ) -> None: + created_node = object() + constructor = _node_constructor(return_value=created_node) + factory._build_human_input_callback = MagicMock(return_value=sentinel.hitl_callback) + monkeypatch.setattr( + factory, + "_resolve_node_class", + MagicMock(return_value=constructor), + ) + + result = factory.create_node({"id": "human-node", "data": {"type": BuiltinNodeTypes.HUMAN_INPUT}}) + + assert result is created_node + kwargs = constructor.call_args.kwargs + assert kwargs["hitl_callback"] is sentinel.hitl_callback + factory._build_human_input_callback.assert_called_once() + + def test_tool_node_receives_tool_file_manager(self, monkeypatch: pytest.MonkeyPatch, factory) -> None: + created_node = object() + constructor = _node_constructor(return_value=created_node) + factory._bound_tool_file_manager_factory = MagicMock(return_value=sentinel.tool_file_manager) + monkeypatch.setattr( + factory, + "_resolve_node_class", + MagicMock(return_value=constructor), + ) + + result = factory.create_node({"id": "tool-node", "data": {"type": BuiltinNodeTypes.TOOL}}) + + assert result is created_node + kwargs = constructor.call_args.kwargs + assert kwargs["tool_file_manager"] is sentinel.tool_file_manager + assert kwargs["runtime"] is sentinel.tool_runtime + assert "tool_file_manager_factory" not in kwargs + factory._bound_tool_file_manager_factory.assert_called_once_with() + + def test_build_llm_compatible_node_init_kwargs_preserves_structured_output_switch(self, factory): + node_data = LLMNodeData.model_validate( { - "id": "human-input", - "data": { - "type": BuiltinNodeTypes.HUMAN_INPUT, - "delivery_methods": [ - { - "type": "email", - "config": { - "recipients": {"items": []}, - "subject": 123, - "body": [], - }, - } - ], + "type": BuiltinNodeTypes.LLM, + "title": "LLM", + "model": {"provider": "provider", "name": "model", "mode": "chat", "completion_params": {}}, + "prompt_template": [{"role": "system", "text": "x"}], + "context": {"enabled": False, "variable_selector": []}, + "vision": {"enabled": False}, + "structured_output_enabled": True, + "structured_output": { + "schema": { + "type": "object", + "properties": {"type": {"type": "string"}}, + "required": ["type"], + } }, } ) + wrapped_model_instance = DifyPreparedLLM(make_model_instance(provider="openai", model="gpt-4")) + memory = sentinel.memory + factory._build_model_instance_for_llm_node = MagicMock(return_value=sentinel.model_instance) + factory._build_memory_for_llm_node = MagicMock(return_value=memory) + with patch.object( + factory, "_wrap_model_instance_for_node", return_value=wrapped_model_instance + ) as wrap_model: + kwargs = factory._build_llm_compatible_node_init_kwargs( + node_class=sentinel.node_class, + node_data=node_data, + wrap_model_instance=True, + include_http_client=True, + include_llm_file_saver=True, + include_prompt_message_serializer=True, + include_retriever_attachment_loader=True, + include_jinja2_template_renderer=True, + ) - constructor.assert_not_called() + assert node_data.structured_output_switch_on is True + assert node_data.structured_output_enabled is True + factory._build_model_instance_for_llm_node.assert_called_once_with(node_data) + factory._build_memory_for_llm_node.assert_called_once_with( + node_data=node_data, + model_instance=sentinel.model_instance, + ) + wrap_model.assert_called_once_with( + node_data=node_data, + model_instance=sentinel.model_instance, + request_metadata={"app_id": "app-id"}, + ) + assert kwargs["model_instance"] is wrapped_model_instance + assert kwargs["polling_finalizer"] == wrapped_model_instance.finalize_llm_polling - def test_rejects_missing_class_mapping(self, monkeypatch: pytest.MonkeyPatch, factory): - monkeypatch.setattr( - factory, - "_resolve_node_class", - MagicMock(side_effect=ValueError("No class mapping found for node type: start")), - ) + def test_resolve_llm_model_reference_uses_shared_model_and_parameters(self, factory): + node_data = LLMNodeData.model_validate( + { + "type": BuiltinNodeTypes.LLM, + "title": "LLM", + "model": { + "provider": "old-provider", + "name": "old-model", + "mode": "chat", + "completion_params": {"temperature": 0.2}, + }, + "model_selector": ["env", "for_summarize"], + "prompt_template": [{"role": "system", "text": "x"}], + "context": {"enabled": False, "variable_selector": []}, + "vision": {"enabled": False}, + } + ) + factory.runtime_state.variable_pool.get.return_value = ObjectSegment( + value={ + "provider": "new-provider", + "name": "new-model", + "mode": "chat", + "completion_params": {"temperature": 0.8}, + } + ) - with pytest.raises(ValueError, match="No class mapping found for node type: start"): - factory.create_node({"id": "node-id", "data": {"type": BuiltinNodeTypes.START}}) + result = factory._resolve_llm_model_reference(node_data) - def test_rejects_missing_latest_class(self, monkeypatch: pytest.MonkeyPatch, factory): - monkeypatch.setattr( - factory, - "_resolve_node_class", - MagicMock(side_effect=ValueError("No latest version class found for node type: start")), - ) + assert result.model.provider == "new-provider" + assert result.model.name == "new-model" + assert result.model.mode == node_data.model.mode + assert result.model.completion_params == {"temperature": 0.8} + factory.runtime_state.variable_pool.get.assert_called_once_with(("env", "for_summarize")) - with pytest.raises(ValueError, match="No latest version class found for node type: start"): - factory.create_node({"id": "node-id", "data": {"type": BuiltinNodeTypes.START}}) + def test_resolve_llm_model_reference_keeps_node_parameters_for_legacy_variable(self, factory): + node_data = LLMNodeData.model_validate( + { + "type": BuiltinNodeTypes.LLM, + "title": "LLM", + "model": { + "provider": "old-provider", + "name": "old-model", + "mode": "chat", + "completion_params": {"temperature": 0.2}, + }, + "model_selector": ["env", "for_summarize"], + "prompt_template": [{"role": "system", "text": "x"}], + "context": {"enabled": False, "variable_selector": []}, + "vision": {"enabled": False}, + } + ) + factory.runtime_state.variable_pool.get.return_value = ObjectSegment( + value={"provider": "new-provider", "name": "new-model", "mode": "chat"} + ) - def test_uses_version_specific_class_when_available(self, monkeypatch: pytest.MonkeyPatch, factory): - matched_node = sentinel.matched_node - latest_node_class = _node_constructor(return_value=sentinel.latest_node) - matched_node_class = _node_constructor(return_value=matched_node) - monkeypatch.setattr( - factory, - "_resolve_node_class", - MagicMock(return_value=matched_node_class), - ) + result = factory._resolve_llm_model_reference(node_data) - result = factory.create_node({"id": "node-id", "data": {"type": BuiltinNodeTypes.START, "version": "9"}}) + assert result.model.completion_params == {"temperature": 0.2} - assert result is matched_node - matched_node_class.assert_called_once() - kwargs = matched_node_class.call_args.kwargs - assert kwargs["node_id"] == "node-id" - _assert_constructor_node_data(kwargs["data"], node_id="node-id", node_type=BuiltinNodeTypes.START, version="9") - assert kwargs["init_params"] is sentinel.init_params - assert kwargs["runtime_state"] is factory.runtime_state - latest_node_class.assert_not_called() + def test_resolve_llm_model_reference_rejects_mode_mismatch(self, factory): + node_data = LLMNodeData.model_validate( + { + "type": BuiltinNodeTypes.LLM, + "title": "LLM", + "model": {"provider": "provider", "name": "model", "mode": "chat"}, + "model_selector": ["env", "shared_model"], + "prompt_template": [{"role": "system", "text": "x"}], + "context": {"enabled": False, "variable_selector": []}, + "vision": {"enabled": False}, + } + ) + factory.runtime_state.variable_pool.get.return_value = ObjectSegment( + value={"provider": "provider", "name": "model", "mode": "completion"} + ) - def test_falls_back_to_latest_class_when_version_specific_mapping_is_missing( - self, monkeypatch: pytest.MonkeyPatch, factory - ): - latest_node = sentinel.latest_node - latest_node_class = _node_constructor(return_value=latest_node) - monkeypatch.setattr( - factory, - "_resolve_node_class", - MagicMock(return_value=latest_node_class), - ) + with pytest.raises(ValueError, match="uses mode 'completion'.*uses mode 'chat'"): + factory._resolve_llm_model_reference(node_data) - result = factory.create_node({"id": "node-id", "data": {"type": BuiltinNodeTypes.START, "version": "9"}}) + def test_resolve_llm_model_reference_rejects_missing_variable(self, factory): + node_data = LLMNodeData.model_validate( + { + "type": BuiltinNodeTypes.LLM, + "title": "LLM", + "model": {"provider": "provider", "name": "model", "mode": "chat"}, + "model_selector": ["env", "shared_model"], + "prompt_template": [{"role": "system", "text": "x"}], + "context": {"enabled": False, "variable_selector": []}, + "vision": {"enabled": False}, + } + ) + factory.runtime_state.variable_pool.get.return_value = None - assert result is latest_node - latest_node_class.assert_called_once() - kwargs = latest_node_class.call_args.kwargs - assert kwargs["node_id"] == "node-id" - _assert_constructor_node_data(kwargs["data"], node_id="node-id", node_type=BuiltinNodeTypes.START, version="9") - assert kwargs["init_params"] is sentinel.init_params - assert kwargs["runtime_state"] is factory.runtime_state + with pytest.raises(ValueError, match="shared_model.*not found"): + factory._resolve_llm_model_reference(node_data) - @pytest.mark.parametrize( - ("node_type", "constructor_name"), - [ - (BuiltinNodeTypes.CODE, "CodeNode"), - (BuiltinNodeTypes.TEMPLATE_TRANSFORM, "TemplateTransformNode"), - (BuiltinNodeTypes.HTTP_REQUEST, "HttpRequestNode"), - (BuiltinNodeTypes.HUMAN_INPUT, "HumanInputNode"), - (BuiltinNodeTypes.TOOL, "ToolNode"), - (KNOWLEDGE_INDEX_NODE_TYPE, "KnowledgeIndexNode"), - (BuiltinNodeTypes.DATASOURCE, "DatasourceNode"), - (BuiltinNodeTypes.KNOWLEDGE_RETRIEVAL, "KnowledgeRetrievalNode"), - (BuiltinNodeTypes.DOCUMENT_EXTRACTOR, "DocumentExtractorNode"), - ], - ) - def test_creates_specialized_nodes(self, monkeypatch: pytest.MonkeyPatch, factory, node_type, constructor_name): - created_node = object() - constructor = _node_constructor(return_value=created_node) - constructor._mock_name = constructor_name - monkeypatch.setattr( - factory, - "_resolve_node_class", - MagicMock(return_value=constructor), - ) + def test_resolve_llm_model_reference_keeps_static_model_for_legacy_non_environment_selector(self, factory): + node_data = LLMNodeData.model_validate( + { + "type": BuiltinNodeTypes.LLM, + "title": "LLM", + "model": {"provider": "provider", "name": "model", "mode": "chat"}, + "model_selector": ["conversation", "shared_model"], + "prompt_template": [{"role": "system", "text": "x"}], + "context": {"enabled": False, "variable_selector": []}, + "vision": {"enabled": False}, + } + ) - if constructor_name == "HumanInputNode": - factory._build_human_input_callback = MagicMock(return_value=sentinel.hitl_callback) + result = factory._resolve_llm_model_reference(node_data) - node_config = {"id": "node-id", "data": {"type": node_type}} - result = factory.create_node(node_config) + assert result is node_data + factory.runtime_state.variable_pool.get.assert_not_called() - assert result is created_node - kwargs = constructor.call_args.kwargs - assert kwargs["node_id"] == "node-id" - _assert_constructor_node_data(kwargs["data"], node_id="node-id", node_type=node_type) - assert kwargs["init_params"] is sentinel.init_params - assert kwargs["runtime_state"] is factory.runtime_state + def test_build_llm_compatible_node_init_kwargs_uses_polling_wrapper_for_polling_llm_node(self, factory): + node_data = LLMNodeData.model_validate( + { + "type": BuiltinNodeTypes.LLM, + "title": "LLM", + "model": {"provider": "provider", "name": "model", "mode": "chat", "completion_params": {}}, + "prompt_template": [{"role": "system", "text": "x"}], + "context": {"enabled": False, "variable_selector": []}, + "vision": {"enabled": False}, + } + ) + plugin_runtime = PluginModelRuntime( + tenant_id="tenant-id", + user_id="user-id", + client=PluginModelClient(), + plugin_service=PluginService, + ) + model_instance = _ModelInstanceStub( + model_runtime=plugin_runtime, + model_schema=_build_llm_model_schema(features=[ModelFeature.POLLING]), + ) + factory._build_model_instance_for_llm_node = MagicMock(return_value=model_instance) + factory._build_memory_for_llm_node = MagicMock(return_value=sentinel.memory) - if constructor_name == "CodeNode": - assert kwargs["code_executor"] is sentinel.code_executor - assert kwargs["code_limits"] is sentinel.code_limits - elif constructor_name == "TemplateTransformNode": - assert kwargs["jinja2_template_renderer"] is sentinel.jinja2_template_renderer - assert kwargs["max_output_length"] == 2048 - elif constructor_name == "HttpRequestNode": - assert kwargs["http_request_config"] is sentinel.http_request_config - assert kwargs["http_client"] is sentinel.http_client - assert kwargs["tool_file_manager_factory"] is factory._bound_tool_file_manager_factory - assert kwargs["file_manager"] is sentinel.file_manager - assert kwargs["file_reference_factory"] is sentinel.file_reference_factory - factory._bound_tool_file_manager_factory.assert_not_called() - elif constructor_name == "HumanInputNode": - assert kwargs["hitl_callback"] is sentinel.hitl_callback - factory._build_human_input_callback.assert_called_once() - elif constructor_name == "ToolNode": - assert kwargs["tool_file_manager"] is sentinel.tool_file_manager - assert kwargs["runtime"] is sentinel.tool_runtime - factory._bound_tool_file_manager_factory.assert_called_once_with() - elif constructor_name == "DocumentExtractorNode": - assert kwargs["unstructured_api_config"] is sentinel.unstructured_api_config - assert kwargs["http_client"] is sentinel.remote_file_http_client - - def test_human_input_node_receives_built_hitl_callback( - self, - monkeypatch: pytest.MonkeyPatch, - factory, - ) -> None: - created_node = object() - constructor = _node_constructor(return_value=created_node) - factory._build_human_input_callback = MagicMock(return_value=sentinel.hitl_callback) - monkeypatch.setattr( - factory, - "_resolve_node_class", - MagicMock(return_value=constructor), - ) - - result = factory.create_node({"id": "human-node", "data": {"type": BuiltinNodeTypes.HUMAN_INPUT}}) - - assert result is created_node - kwargs = constructor.call_args.kwargs - assert kwargs["hitl_callback"] is sentinel.hitl_callback - factory._build_human_input_callback.assert_called_once() - - def test_tool_node_receives_tool_file_manager(self, monkeypatch: pytest.MonkeyPatch, factory) -> None: - created_node = object() - constructor = _node_constructor(return_value=created_node) - factory._bound_tool_file_manager_factory = MagicMock(return_value=sentinel.tool_file_manager) - monkeypatch.setattr( - factory, - "_resolve_node_class", - MagicMock(return_value=constructor), - ) - - result = factory.create_node({"id": "tool-node", "data": {"type": BuiltinNodeTypes.TOOL}}) - - assert result is created_node - kwargs = constructor.call_args.kwargs - assert kwargs["tool_file_manager"] is sentinel.tool_file_manager - assert kwargs["runtime"] is sentinel.tool_runtime - assert "tool_file_manager_factory" not in kwargs - factory._bound_tool_file_manager_factory.assert_called_once_with() - - def test_build_llm_compatible_node_init_kwargs_preserves_structured_output_switch(self, factory): - node_data = LLMNodeData.model_validate( - { - "type": BuiltinNodeTypes.LLM, - "title": "LLM", - "model": {"provider": "provider", "name": "model", "mode": "chat", "completion_params": {}}, - "prompt_template": [{"role": "system", "text": "x"}], - "context": {"enabled": False, "variable_selector": []}, - "vision": {"enabled": False}, - "structured_output_enabled": True, - "structured_output": { - "schema": { - "type": "object", - "properties": {"type": {"type": "string"}}, - "required": ["type"], - } - }, - } - ) - wrapped_model_instance = DifyPreparedLLM(make_model_instance(provider="openai", model="gpt-4")) - memory = sentinel.memory - factory._build_model_instance_for_llm_node = MagicMock(return_value=sentinel.model_instance) - factory._build_memory_for_llm_node = MagicMock(return_value=memory) - with patch.object(factory, "_wrap_model_instance_for_node", return_value=wrapped_model_instance) as wrap_model: kwargs = factory._build_llm_compatible_node_init_kwargs( node_class=sentinel.node_class, node_data=node_data, wrap_model_instance=True, - include_http_client=True, - include_llm_file_saver=True, - include_prompt_message_serializer=True, - include_retriever_attachment_loader=True, - include_jinja2_template_renderer=True, + include_http_client=False, + include_llm_file_saver=False, + include_prompt_message_serializer=False, + include_retriever_attachment_loader=False, + include_jinja2_template_renderer=False, ) - assert node_data.structured_output_switch_on is True - assert node_data.structured_output_enabled is True - factory._build_model_instance_for_llm_node.assert_called_once_with(node_data) - factory._build_memory_for_llm_node.assert_called_once_with( - node_data=node_data, - model_instance=sentinel.model_instance, + assert isinstance(kwargs["model_instance"], LLMPollingCapableProtocol) + + @pytest.mark.parametrize( + "node_type", [BuiltinNodeTypes.QUESTION_CLASSIFIER, BuiltinNodeTypes.PARAMETER_EXTRACTOR] ) - wrap_model.assert_called_once_with( - node_data=node_data, - model_instance=sentinel.model_instance, - request_metadata={"app_id": "app-id"}, - ) - assert kwargs["model_instance"] is wrapped_model_instance - assert kwargs["polling_finalizer"] == wrapped_model_instance.finalize_llm_polling + def test_wrap_model_instance_keeps_non_llm_graph_nodes_on_plain_wrapper(self, node_type): + plugin_runtime = PluginModelRuntime( + tenant_id="tenant-id", + user_id="user-id", + client=PluginModelClient(), + plugin_service=PluginService, + ) + model_instance = _ModelInstanceStub( + model_runtime=plugin_runtime, + model_schema=_build_llm_model_schema(features=[ModelFeature.POLLING]), + ) - def test_resolve_llm_model_reference_uses_shared_model_and_parameters(self, factory): - node_data = LLMNodeData.model_validate( - { - "type": BuiltinNodeTypes.LLM, - "title": "LLM", - "model": { - "provider": "old-provider", - "name": "old-model", - "mode": "chat", - "completion_params": {"temperature": 0.2}, - }, - "model_selector": ["env", "for_summarize"], - "prompt_template": [{"role": "system", "text": "x"}], - "context": {"enabled": False, "variable_selector": []}, - "vision": {"enabled": False}, - } - ) - factory.runtime_state.variable_pool.get.return_value = ObjectSegment( - value={ - "provider": "new-provider", - "name": "new-model", - "mode": "chat", - "completion_params": {"temperature": 0.8}, - } - ) + wrapped = node_factory.DifyNodeFactory._wrap_model_instance_for_node( + node_data=SimpleNamespace(type=node_type), + model_instance=model_instance, + ) - result = factory._resolve_llm_model_reference(node_data) + assert type(wrapped) is DifyPreparedLLM + assert not isinstance(wrapped, LLMPollingCapableProtocol) - assert result.model.provider == "new-provider" - assert result.model.name == "new-model" - assert result.model.mode == node_data.model.mode - assert result.model.completion_params == {"temperature": 0.8} - factory.runtime_state.variable_pool.get.assert_called_once_with(("env", "for_summarize")) + def test_create_node_passes_canonical_llm_data_to_constructor(self, monkeypatch, factory): + created_node = object() + constructor = _node_constructor(return_value=created_node) + constructor.validate_node_data.side_effect = lambda node_data: LLMNodeData.model_validate( + node_data.model_dump(mode="python") if isinstance(node_data, BaseNodeData) else node_data + ) + monkeypatch.setattr(factory, "_resolve_node_class", lambda **_kwargs: constructor) + monkeypatch.setattr(factory, "_build_llm_compatible_node_init_kwargs", MagicMock(return_value={})) - def test_resolve_llm_model_reference_keeps_node_parameters_for_legacy_variable(self, factory): - node_data = LLMNodeData.model_validate( - { - "type": BuiltinNodeTypes.LLM, - "title": "LLM", - "model": { - "provider": "old-provider", - "name": "old-model", - "mode": "chat", - "completion_params": {"temperature": 0.2}, - }, - "model_selector": ["env", "for_summarize"], - "prompt_template": [{"role": "system", "text": "x"}], - "context": {"enabled": False, "variable_selector": []}, - "vision": {"enabled": False}, - } - ) - factory.runtime_state.variable_pool.get.return_value = ObjectSegment( - value={"provider": "new-provider", "name": "new-model", "mode": "chat"} - ) - - result = factory._resolve_llm_model_reference(node_data) - - assert result.model.completion_params == {"temperature": 0.2} - - def test_resolve_llm_model_reference_rejects_mode_mismatch(self, factory): - node_data = LLMNodeData.model_validate( - { - "type": BuiltinNodeTypes.LLM, - "title": "LLM", - "model": {"provider": "provider", "name": "model", "mode": "chat"}, - "model_selector": ["env", "shared_model"], - "prompt_template": [{"role": "system", "text": "x"}], - "context": {"enabled": False, "variable_selector": []}, - "vision": {"enabled": False}, - } - ) - factory.runtime_state.variable_pool.get.return_value = ObjectSegment( - value={"provider": "provider", "name": "model", "mode": "completion"} - ) - - with pytest.raises(ValueError, match="uses mode 'completion'.*uses mode 'chat'"): - factory._resolve_llm_model_reference(node_data) - - def test_resolve_llm_model_reference_rejects_missing_variable(self, factory): - node_data = LLMNodeData.model_validate( - { - "type": BuiltinNodeTypes.LLM, - "title": "LLM", - "model": {"provider": "provider", "name": "model", "mode": "chat"}, - "model_selector": ["env", "shared_model"], - "prompt_template": [{"role": "system", "text": "x"}], - "context": {"enabled": False, "variable_selector": []}, - "vision": {"enabled": False}, - } - ) - factory.runtime_state.variable_pool.get.return_value = None - - with pytest.raises(ValueError, match="shared_model.*not found"): - factory._resolve_llm_model_reference(node_data) - - def test_resolve_llm_model_reference_keeps_static_model_for_legacy_non_environment_selector(self, factory): - node_data = LLMNodeData.model_validate( - { - "type": BuiltinNodeTypes.LLM, - "title": "LLM", - "model": {"provider": "provider", "name": "model", "mode": "chat"}, - "model_selector": ["conversation", "shared_model"], - "prompt_template": [{"role": "system", "text": "x"}], - "context": {"enabled": False, "variable_selector": []}, - "vision": {"enabled": False}, - } - ) - - result = factory._resolve_llm_model_reference(node_data) - - assert result is node_data - factory.runtime_state.variable_pool.get.assert_not_called() - - def test_build_llm_compatible_node_init_kwargs_uses_polling_wrapper_for_polling_llm_node(self, factory): - node_data = LLMNodeData.model_validate( - { - "type": BuiltinNodeTypes.LLM, - "title": "LLM", - "model": {"provider": "provider", "name": "model", "mode": "chat", "completion_params": {}}, - "prompt_template": [{"role": "system", "text": "x"}], - "context": {"enabled": False, "variable_selector": []}, - "vision": {"enabled": False}, - } - ) - plugin_runtime = PluginModelRuntime( - tenant_id="tenant-id", - user_id="user-id", - client=PluginModelClient(), - plugin_service=PluginService, - ) - model_instance = _ModelInstanceStub( - model_runtime=plugin_runtime, - model_schema=_build_llm_model_schema(features=[ModelFeature.POLLING]), - ) - factory._build_model_instance_for_llm_node = MagicMock(return_value=model_instance) - factory._build_memory_for_llm_node = MagicMock(return_value=sentinel.memory) - - kwargs = factory._build_llm_compatible_node_init_kwargs( - node_class=sentinel.node_class, - node_data=node_data, - wrap_model_instance=True, - include_http_client=False, - include_llm_file_saver=False, - include_prompt_message_serializer=False, - include_retriever_attachment_loader=False, - include_jinja2_template_renderer=False, - ) - - assert isinstance(kwargs["model_instance"], LLMPollingCapableProtocol) - - @pytest.mark.parametrize("node_type", [BuiltinNodeTypes.QUESTION_CLASSIFIER, BuiltinNodeTypes.PARAMETER_EXTRACTOR]) - def test_wrap_model_instance_keeps_non_llm_graph_nodes_on_plain_wrapper(self, node_type): - plugin_runtime = PluginModelRuntime( - tenant_id="tenant-id", - user_id="user-id", - client=PluginModelClient(), - plugin_service=PluginService, - ) - model_instance = _ModelInstanceStub( - model_runtime=plugin_runtime, - model_schema=_build_llm_model_schema(features=[ModelFeature.POLLING]), - ) - - wrapped = node_factory.DifyNodeFactory._wrap_model_instance_for_node( - node_data=SimpleNamespace(type=node_type), - model_instance=model_instance, - ) - - assert type(wrapped) is DifyPreparedLLM - assert not isinstance(wrapped, LLMPollingCapableProtocol) - - def test_create_node_passes_canonical_llm_data_to_constructor(self, monkeypatch, factory): - created_node = object() - constructor = _node_constructor(return_value=created_node) - constructor.validate_node_data.side_effect = lambda node_data: LLMNodeData.model_validate( - node_data.model_dump(mode="python") if isinstance(node_data, BaseNodeData) else node_data - ) - monkeypatch.setattr(factory, "_resolve_node_class", lambda **_kwargs: constructor) - monkeypatch.setattr(factory, "_build_llm_compatible_node_init_kwargs", MagicMock(return_value={})) - - node_config = { - "id": "llm-node-id", - "data": { - "type": BuiltinNodeTypes.LLM, - "title": "LLM", - "model": {"provider": "provider", "name": "model", "mode": "chat", "completion_params": {}}, - "prompt_template": [{"role": "system", "text": "x"}], - "context": {"enabled": False, "variable_selector": []}, - "vision": {"enabled": False}, - "structured_output_enabled": True, - "structured_output": { - "schema": { - "type": "object", - "properties": {"type": {"type": "string"}}, - "required": ["type"], - } - }, - }, - } - - factory.create_node(node_config) - - data = constructor.call_args.kwargs["data"] - assert isinstance(data, Mapping) - assert data["structured_output_switch_on"] is True - assert "structured_output_enabled" not in data - assert LLMNodeData.model_validate(data).structured_output_enabled is True - - def test_create_node_preserves_structured_output_switch_after_graphon_constructor(self, monkeypatch, factory): - factory.init_params = SimpleNamespace( - workflow_id="workflow-id", - graph_config={}, - run_context={}, - call_depth=0, - ) - monkeypatch.setattr(factory, "_resolve_node_class", MagicMock(return_value=LLMNode)) - monkeypatch.setattr( - factory, - "_build_llm_compatible_node_init_kwargs", - MagicMock( - return_value={ - "model_instance": sentinel.model_instance, - "llm_file_saver": sentinel.llm_file_saver, - "prompt_message_serializer": sentinel.prompt_message_serializer, - } - ), - ) - - node_config = { - "id": "llm-node-id", - "data": { - "type": BuiltinNodeTypes.LLM, - "title": "LLM", - "model": {"provider": "provider", "name": "model", "mode": "chat", "completion_params": {}}, - "prompt_template": [{"role": "system", "text": "x"}], - "context": {"enabled": False, "variable_selector": []}, - "vision": {"enabled": False}, - "structured_output_enabled": True, - "structured_output": { - "schema": { - "type": "object", - "properties": {"type": {"type": "string"}}, - "required": ["type"], - } - }, - }, - } - - node = factory.create_node(node_config) - - assert node.node_data.structured_output_switch_on is True - assert node.node_data.structured_output_enabled is True - - def test_create_node_uses_dify_llm_node_for_persisted_version_one(self, monkeypatch, factory): - factory.init_params = SimpleNamespace( - workflow_id="workflow-id", - graph_config={}, - run_context={}, - call_depth=0, - ) - monkeypatch.setattr( - factory, - "_build_llm_compatible_node_init_kwargs", - MagicMock( - return_value={ - "model_instance": sentinel.model_instance, - "llm_file_saver": sentinel.llm_file_saver, - "prompt_message_serializer": sentinel.prompt_message_serializer, - "polling_finalizer": MagicMock(), - } - ), - ) - - node = factory.create_node( - { + node_config = { "id": "llm-node-id", "data": { "type": BuiltinNodeTypes.LLM, - "version": "1", "title": "LLM", - "model": {"provider": "provider", "name": "model", "mode": "chat"}, + "model": {"provider": "provider", "name": "model", "mode": "chat", "completion_params": {}}, "prompt_template": [{"role": "system", "text": "x"}], "context": {"enabled": False, "variable_selector": []}, "vision": {"enabled": False}, + "structured_output_enabled": True, + "structured_output": { + "schema": { + "type": "object", + "properties": {"type": {"type": "string"}}, + "required": ["type"], + } + }, }, } - ) - assert isinstance(node, DifyLLMNode) - assert node.version() == "1" + factory.create_node(node_config) - @pytest.mark.parametrize( - ("node_type", "constructor_name", "expected_extra_kwargs"), - [ - ( - BuiltinNodeTypes.LLM, - "LLMNode", - { - "http_client": sentinel.remote_file_http_client, - "llm_file_saver": sentinel.llm_file_saver, - "prompt_message_serializer": sentinel.prompt_message_serializer, - "retriever_attachment_loader": sentinel.retriever_attachment_loader, - "jinja2_template_renderer": sentinel.jinja2_template_renderer, + data = constructor.call_args.kwargs["data"] + assert isinstance(data, Mapping) + assert data["structured_output_switch_on"] is True + assert "structured_output_enabled" not in data + assert LLMNodeData.model_validate(data).structured_output_enabled is True + + def test_create_node_preserves_structured_output_switch_after_graphon_constructor(self, monkeypatch, factory): + factory.init_params = SimpleNamespace( + workflow_id="workflow-id", + graph_config={}, + run_context={}, + call_depth=0, + ) + monkeypatch.setattr(factory, "_resolve_node_class", MagicMock(return_value=LLMNode)) + monkeypatch.setattr( + factory, + "_build_llm_compatible_node_init_kwargs", + MagicMock( + return_value={ + "model_instance": sentinel.model_instance, + "llm_file_saver": sentinel.llm_file_saver, + "prompt_message_serializer": sentinel.prompt_message_serializer, + } + ), + ) + + node_config = { + "id": "llm-node-id", + "data": { + "type": BuiltinNodeTypes.LLM, + "title": "LLM", + "model": {"provider": "provider", "name": "model", "mode": "chat", "completion_params": {}}, + "prompt_template": [{"role": "system", "text": "x"}], + "context": {"enabled": False, "variable_selector": []}, + "vision": {"enabled": False}, + "structured_output_enabled": True, + "structured_output": { + "schema": { + "type": "object", + "properties": {"type": {"type": "string"}}, + "required": ["type"], + } + }, }, - ), - ( - BuiltinNodeTypes.QUESTION_CLASSIFIER, - "QuestionClassifierNode", + } + + node = factory.create_node(node_config) + + assert node.node_data.structured_output_switch_on is True + assert node.node_data.structured_output_enabled is True + + def test_create_node_uses_dify_llm_node_for_persisted_version_one(self, monkeypatch, factory): + factory.init_params = SimpleNamespace( + workflow_id="workflow-id", + graph_config={}, + run_context={}, + call_depth=0, + ) + monkeypatch.setattr( + factory, + "_build_llm_compatible_node_init_kwargs", + MagicMock( + return_value={ + "model_instance": sentinel.model_instance, + "llm_file_saver": sentinel.llm_file_saver, + "prompt_message_serializer": sentinel.prompt_message_serializer, + "polling_finalizer": MagicMock(), + } + ), + ) + + node = factory.create_node( { - "http_client": sentinel.remote_file_http_client, - "llm_file_saver": sentinel.llm_file_saver, - "prompt_message_serializer": sentinel.prompt_message_serializer, - "template_renderer": sentinel.jinja2_template_renderer, - }, - ), - ( - BuiltinNodeTypes.PARAMETER_EXTRACTOR, - "ParameterExtractorNode", - { - "prompt_message_serializer": sentinel.prompt_message_serializer, - }, - ), - ], - ) - def test_creates_model_backed_nodes( - self, - monkeypatch: pytest.MonkeyPatch, - factory, - node_type, - constructor_name, - expected_extra_kwargs, - ): - created_node = object() - constructor = _node_constructor(return_value=created_node) - constructor._mock_name = constructor_name - monkeypatch.setattr( - factory, - "_resolve_node_class", - lambda **_kwargs: constructor, - ) - llm_init_kwargs = { - "credentials_provider": sentinel.credentials_provider, - "model_factory": sentinel.model_factory, - "model_instance": sentinel.model_instance, - "memory": sentinel.memory, - **expected_extra_kwargs, - } - build_llm_init_kwargs = MagicMock(return_value=llm_init_kwargs) - factory._build_llm_compatible_node_init_kwargs = build_llm_init_kwargs + "id": "llm-node-id", + "data": { + "type": BuiltinNodeTypes.LLM, + "version": "1", + "title": "LLM", + "model": {"provider": "provider", "name": "model", "mode": "chat"}, + "prompt_template": [{"role": "system", "text": "x"}], + "context": {"enabled": False, "variable_selector": []}, + "vision": {"enabled": False}, + }, + } + ) - node_config = {"id": "node-id", "data": {"type": node_type}} - result = factory.create_node(node_config) + assert isinstance(node, DifyLLMNode) + assert node.version() == "1" - assert result is created_node - build_llm_init_kwargs.assert_called_once() - helper_kwargs = build_llm_init_kwargs.call_args.kwargs - assert helper_kwargs["node_class"] is constructor - assert isinstance(helper_kwargs["node_data"], BaseNodeData) - assert helper_kwargs["node_data"].type == node_type - assert helper_kwargs["wrap_model_instance"] is True - assert helper_kwargs["include_http_client"] is (node_type != BuiltinNodeTypes.PARAMETER_EXTRACTOR) - assert helper_kwargs["include_llm_file_saver"] is (node_type != BuiltinNodeTypes.PARAMETER_EXTRACTOR) - assert helper_kwargs["include_prompt_message_serializer"] is True - assert helper_kwargs["include_retriever_attachment_loader"] is (node_type == BuiltinNodeTypes.LLM) - assert helper_kwargs["include_jinja2_template_renderer"] is (node_type == BuiltinNodeTypes.LLM) - - constructor_kwargs = constructor.call_args.kwargs - assert constructor_kwargs["node_id"] == "node-id" - _assert_constructor_node_data(constructor_kwargs["data"], node_id="node-id", node_type=node_type) - assert constructor_kwargs["init_params"] is sentinel.init_params - assert constructor_kwargs["runtime_state"] is factory.runtime_state - assert constructor_kwargs["credentials_provider"] is sentinel.credentials_provider - assert constructor_kwargs["model_factory"] is sentinel.model_factory - assert constructor_kwargs["model_instance"] is sentinel.model_instance - assert constructor_kwargs["memory"] is sentinel.memory - for key, value in expected_extra_kwargs.items(): - assert constructor_kwargs[key] is value - - def test_parameter_extractor_init_does_not_require_retriever_context(self, factory): - node_data = ParameterExtractorNodeData.model_validate( - { - "type": BuiltinNodeTypes.PARAMETER_EXTRACTOR, - "title": "Parameter Extractor", - "model": {"provider": "provider", "name": "model", "mode": "chat", "completion_params": {}}, - "query": ["sys", "query"], - "parameters": [ + @pytest.mark.parametrize( + ("node_type", "constructor_name", "expected_extra_kwargs"), + [ + ( + BuiltinNodeTypes.LLM, + "LLMNode", { - "name": "topic", - "type": "string", - "description": "Topic", - "required": True, - } - ], - "reasoning_mode": "prompt", + "http_client": sentinel.remote_file_http_client, + "llm_file_saver": sentinel.llm_file_saver, + "prompt_message_serializer": sentinel.prompt_message_serializer, + "retriever_attachment_loader": sentinel.retriever_attachment_loader, + "jinja2_template_renderer": sentinel.jinja2_template_renderer, + }, + ), + ( + BuiltinNodeTypes.QUESTION_CLASSIFIER, + "QuestionClassifierNode", + { + "http_client": sentinel.remote_file_http_client, + "llm_file_saver": sentinel.llm_file_saver, + "prompt_message_serializer": sentinel.prompt_message_serializer, + "template_renderer": sentinel.jinja2_template_renderer, + }, + ), + ( + BuiltinNodeTypes.PARAMETER_EXTRACTOR, + "ParameterExtractorNode", + { + "prompt_message_serializer": sentinel.prompt_message_serializer, + }, + ), + ], + ) + def test_creates_model_backed_nodes( + self, + monkeypatch: pytest.MonkeyPatch, + factory, + node_type, + constructor_name, + expected_extra_kwargs, + ): + created_node = object() + constructor = _node_constructor(return_value=created_node) + constructor._mock_name = constructor_name + monkeypatch.setattr( + factory, + "_resolve_node_class", + lambda **_kwargs: constructor, + ) + llm_init_kwargs = { + "credentials_provider": sentinel.credentials_provider, + "model_factory": sentinel.model_factory, + "model_instance": sentinel.model_instance, + "memory": sentinel.memory, + **expected_extra_kwargs, } - ) - factory._build_model_instance_for_llm_node = MagicMock(return_value=sentinel.model_instance) - factory._build_memory_for_llm_node = MagicMock(return_value=sentinel.memory) - factory._build_retriever_attachment_loader = MagicMock(side_effect=AssertionError("unexpected loader build")) + build_llm_init_kwargs = MagicMock(return_value=llm_init_kwargs) + factory._build_llm_compatible_node_init_kwargs = build_llm_init_kwargs - kwargs = factory._build_llm_compatible_node_init_kwargs( - node_class=sentinel.node_class, - node_data=node_data, - wrap_model_instance=True, - include_http_client=False, - include_llm_file_saver=False, - include_prompt_message_serializer=True, - include_retriever_attachment_loader=False, - include_jinja2_template_renderer=False, - ) + node_config = {"id": "node-id", "data": {"type": node_type}} + result = factory.create_node(node_config) - assert "retriever_attachment_loader" not in kwargs - assert kwargs["prompt_message_serializer"] is sentinel.prompt_message_serializer - factory._build_retriever_attachment_loader.assert_not_called() + assert result is created_node + build_llm_init_kwargs.assert_called_once() + helper_kwargs = build_llm_init_kwargs.call_args.kwargs + assert helper_kwargs["node_class"] is constructor + assert isinstance(helper_kwargs["node_data"], BaseNodeData) + assert helper_kwargs["node_data"].type == node_type + assert helper_kwargs["wrap_model_instance"] is True + assert helper_kwargs["include_http_client"] is (node_type != BuiltinNodeTypes.PARAMETER_EXTRACTOR) + assert helper_kwargs["include_llm_file_saver"] is (node_type != BuiltinNodeTypes.PARAMETER_EXTRACTOR) + assert helper_kwargs["include_prompt_message_serializer"] is True + assert helper_kwargs["include_retriever_attachment_loader"] is (node_type == BuiltinNodeTypes.LLM) + assert helper_kwargs["include_jinja2_template_renderer"] is (node_type == BuiltinNodeTypes.LLM) + constructor_kwargs = constructor.call_args.kwargs + assert constructor_kwargs["node_id"] == "node-id" + _assert_constructor_node_data(constructor_kwargs["data"], node_id="node-id", node_type=node_type) + assert constructor_kwargs["init_params"] is sentinel.init_params + assert constructor_kwargs["runtime_state"] is factory.runtime_state + assert constructor_kwargs["credentials_provider"] is sentinel.credentials_provider + assert constructor_kwargs["model_factory"] is sentinel.model_factory + assert constructor_kwargs["model_instance"] is sentinel.model_instance + assert constructor_kwargs["memory"] is sentinel.memory + for key, value in expected_extra_kwargs.items(): + assert constructor_kwargs[key] is value -class TestDifyNodeFactoryRetrieverAttachmentAccess: - @pytest.fixture - def factory(self): - factory = object.__new__(node_factory.DifyNodeFactory) - factory.runtime_state = SimpleNamespace(variable_pool=MagicMock()) - return factory - - def test_retriever_attachment_loader_is_typed_for_llm_node_data_only(self): - annotations = node_factory.DifyNodeFactory._build_retriever_attachment_loader.__annotations__ - - assert annotations["node_data"] is LLMNodeData - - def test_build_retriever_attachment_loader_uses_llm_context_selector(self, factory): - factory._file_reference_factory = sentinel.file_reference_factory - factory.runtime_state.variable_pool.get.return_value = ArrayObjectSegment( - value=[ + def test_parameter_extractor_init_does_not_require_retriever_context(self, factory): + node_data = ParameterExtractorNodeData.model_validate( { - "metadata": { - "_source": "knowledge", - "segment_id": "allowed-segment", - } + "type": BuiltinNodeTypes.PARAMETER_EXTRACTOR, + "title": "Parameter Extractor", + "model": {"provider": "provider", "name": "model", "mode": "chat", "completion_params": {}}, + "query": ["sys", "query"], + "parameters": [ + { + "name": "topic", + "type": "string", + "description": "Topic", + "required": True, + } + ], + "reasoning_mode": "prompt", } - ] - ) - node_data = LLMNodeData.model_validate( - { - "type": BuiltinNodeTypes.LLM, - "title": "LLM", - "model": {"provider": "provider", "name": "model", "mode": "chat", "completion_params": {}}, - "prompt_template": [{"role": "system", "text": "x"}], - "context": {"enabled": True, "variable_selector": ["knowledge-node", "result"]}, - "vision": {"enabled": False}, - } - ) + ) + factory._build_model_instance_for_llm_node = MagicMock(return_value=sentinel.model_instance) + factory._build_memory_for_llm_node = MagicMock(return_value=sentinel.memory) + factory._build_retriever_attachment_loader = MagicMock( + side_effect=AssertionError("unexpected loader build") + ) - loader = factory._build_retriever_attachment_loader(node_data) + kwargs = factory._build_llm_compatible_node_init_kwargs( + node_class=sentinel.node_class, + node_data=node_data, + wrap_model_instance=True, + include_http_client=False, + include_llm_file_saver=False, + include_prompt_message_serializer=True, + include_retriever_attachment_loader=False, + include_jinja2_template_renderer=False, + ) - assert loader._segment_access_checker is not None - assert loader._segment_access_checker("allowed-segment") is True - factory.runtime_state.variable_pool.get.assert_called_once_with(["knowledge-node", "result"]) + assert "retriever_attachment_loader" not in kwargs + assert kwargs["prompt_message_serializer"] is sentinel.prompt_message_serializer + factory._build_retriever_attachment_loader.assert_not_called() - def test_checker_rejects_missing_context_selector_without_reading_variable_pool(self, factory): - checker = factory._build_retriever_segment_access_checker(None) + class TestDifyNodeFactoryRetrieverAttachmentAccess: + @pytest.fixture + def factory(self): + factory = object.__new__(node_factory.DifyNodeFactory) + factory.runtime_state = SimpleNamespace(variable_pool=MagicMock()) + return factory - assert checker("segment-id") is False - factory.runtime_state.variable_pool.get.assert_not_called() + def test_retriever_attachment_loader_is_typed_for_llm_node_data_only(self): + annotations = node_factory.DifyNodeFactory._build_retriever_attachment_loader.__annotations__ - def test_checker_rejects_non_knowledge_context_items(self, factory): - factory.runtime_state.variable_pool.get.return_value = ArrayObjectSegment.model_construct( - value=[ - "plain-text", - {"metadata": "not-a-mapping"}, - ] - ) + assert annotations["node_data"] is LLMNodeData - checker = factory._build_retriever_segment_access_checker(["knowledge-node", "result"]) - - assert checker("segment-id") is False - - def test_checker_rejects_non_array_context_value(self, factory): - factory.runtime_state.variable_pool.get.return_value = StringSegment(value="not knowledge context") - - checker = factory._build_retriever_segment_access_checker(["knowledge-node", "result"]) - - assert checker("segment-id") is False - - def test_checker_allows_only_segments_from_selected_knowledge_context(self, factory): - factory.runtime_state.variable_pool.get.return_value = ArrayObjectSegment( - value=[ + def test_build_retriever_attachment_loader_uses_llm_context_selector(self, factory): + factory._file_reference_factory = sentinel.file_reference_factory + factory.runtime_state.variable_pool.get.return_value = ArrayObjectSegment( + value=[ + { + "metadata": { + "_source": "knowledge", + "segment_id": "allowed-segment", + } + } + ] + ) + node_data = LLMNodeData.model_validate( { - "metadata": { - "_source": "knowledge", - "segment_id": "allowed-segment", - } + "type": BuiltinNodeTypes.LLM, + "title": "LLM", + "model": {"provider": "provider", "name": "model", "mode": "chat", "completion_params": {}}, + "prompt_template": [{"role": "system", "text": "x"}], + "context": {"enabled": True, "variable_selector": ["knowledge-node", "result"]}, + "vision": {"enabled": False}, } - ] - ) + ) - checker = factory._build_retriever_segment_access_checker(["knowledge-node", "result"]) + loader = factory._build_retriever_attachment_loader(node_data) - assert checker("allowed-segment") is True - assert checker("other-segment") is False - factory.runtime_state.variable_pool.get.assert_any_call(["knowledge-node", "result"]) + assert loader._segment_access_checker is not None + assert loader._segment_access_checker("allowed-segment") is True + factory.runtime_state.variable_pool.get.assert_called_once_with(["knowledge-node", "result"]) + def test_checker_rejects_missing_context_selector_without_reading_variable_pool(self, factory): + checker = factory._build_retriever_segment_access_checker(None) -class TestDifyNodeFactoryModelInstance: - @pytest.fixture - def factory(self): - factory = object.__new__(node_factory.DifyNodeFactory) - factory._llm_credentials_provider = sentinel.credentials_provider - factory._llm_model_factory = sentinel.model_factory - return factory + assert checker("segment-id") is False + factory.runtime_state.variable_pool.get.assert_not_called() - def test_delegates_to_fetch_model_config(self, monkeypatch: pytest.MonkeyPatch, factory): - node_data_model = SimpleNamespace( - provider="provider", - name="model", - mode="chat", - completion_params={"temperature": 0.3, "stop": ["Human:"]}, - ) - node_data = SimpleNamespace(model=node_data_model) - model_type_instance = MagicMock() - model_instance = SimpleNamespace( - model_type_instance=model_type_instance, - parameters={"temperature": 0.3}, - stop=("Human:",), - ) - fetch_model_config = MagicMock(return_value=(model_instance, sentinel.model_config)) - monkeypatch.setattr(node_factory, "fetch_model_config", fetch_model_config) + def test_checker_rejects_non_knowledge_context_items(self, factory): + factory.runtime_state.variable_pool.get.return_value = ArrayObjectSegment.model_construct( + value=[ + "plain-text", + {"metadata": "not-a-mapping"}, + ] + ) - result = factory._build_model_instance_for_llm_node(node_data) + checker = factory._build_retriever_segment_access_checker(["knowledge-node", "result"]) - assert result is model_instance - assert result.parameters == {"temperature": 0.3} - assert result.stop == ("Human:",) - assert result.model_type_instance is model_type_instance - fetch_model_config.assert_called_once_with( - node_data_model=node_data_model, - credentials_provider=sentinel.credentials_provider, - model_factory=sentinel.model_factory, - ) + assert checker("segment-id") is False - def test_propagates_fetch_model_config_errors(self, monkeypatch: pytest.MonkeyPatch, factory): - fetch_model_config = MagicMock(side_effect=ValueError("broken model config")) - monkeypatch.setattr(node_factory, "fetch_model_config", fetch_model_config) + def test_checker_rejects_non_array_context_value(self, factory): + factory.runtime_state.variable_pool.get.return_value = StringSegment(value="not knowledge context") - with pytest.raises(ValueError, match="broken model config"): - factory._build_model_instance_for_llm_node(SimpleNamespace(model=sentinel.node_data_model)) + checker = factory._build_retriever_segment_access_checker(["knowledge-node", "result"]) + assert checker("segment-id") is False -class TestDifyNodeFactoryMemory: - @pytest.fixture - def factory(self): - factory = object.__new__(node_factory.DifyNodeFactory) - factory._dify_context = SimpleNamespace(app_id="app-id") - factory.runtime_state = SimpleNamespace(variable_pool=MagicMock()) - return factory + def test_checker_allows_only_segments_from_selected_knowledge_context(self, factory): + factory.runtime_state.variable_pool.get.return_value = ArrayObjectSegment( + value=[ + { + "metadata": { + "_source": "knowledge", + "segment_id": "allowed-segment", + } + } + ] + ) - def test_returns_none_when_memory_is_not_configured(self, factory): - result = factory._build_memory_for_llm_node( - node_data=SimpleNamespace(memory=None), - model_instance=sentinel.model_instance, - ) + checker = factory._build_retriever_segment_access_checker(["knowledge-node", "result"]) - assert result is None - factory.runtime_state.variable_pool.get.assert_not_called() + assert checker("allowed-segment") is True + assert checker("other-segment") is False + factory.runtime_state.variable_pool.get.assert_any_call(["knowledge-node", "result"]) - def test_uses_string_segment_conversation_id(self, monkeypatch: pytest.MonkeyPatch, factory): - memory_config = sentinel.memory_config - factory.runtime_state.variable_pool.get.return_value = StringSegment(value="conversation-id") - fetch_memory = MagicMock(return_value=sentinel.memory) - monkeypatch.setattr(node_factory, "fetch_memory", fetch_memory) + class TestDifyNodeFactoryModelInstance: + @pytest.fixture + def factory(self): + factory = object.__new__(node_factory.DifyNodeFactory) + factory._llm_credentials_provider = sentinel.credentials_provider + factory._llm_model_factory = sentinel.model_factory + return factory - result = factory._build_memory_for_llm_node( - node_data=SimpleNamespace(memory=memory_config), - model_instance=sentinel.model_instance, - ) + def test_delegates_to_fetch_model_config(self, monkeypatch: pytest.MonkeyPatch, factory): + node_data_model = SimpleNamespace( + provider="provider", + name="model", + mode="chat", + completion_params={"temperature": 0.3, "stop": ["Human:"]}, + ) + node_data = SimpleNamespace(model=node_data_model) + model_type_instance = MagicMock() + model_instance = SimpleNamespace( + model_type_instance=model_type_instance, + parameters={"temperature": 0.3}, + stop=("Human:",), + ) + fetch_model_config = MagicMock(return_value=(model_instance, sentinel.model_config)) + monkeypatch.setattr(node_factory, "fetch_model_config", fetch_model_config) - assert result is sentinel.memory - factory.runtime_state.variable_pool.get.assert_called_once_with(("sys", "conversation_id")) - fetch_memory.assert_called_once_with( - conversation_id="conversation-id", - app_id="app-id", - node_data_memory=memory_config, - model_instance=sentinel.model_instance, - ) + result = factory._build_model_instance_for_llm_node(node_data) - def test_ignores_non_string_segment_conversation_ids(self, monkeypatch: pytest.MonkeyPatch, factory): - memory_config = sentinel.memory_config - factory.runtime_state.variable_pool.get.return_value = sentinel.segment - fetch_memory = MagicMock(return_value=sentinel.memory) - monkeypatch.setattr(node_factory, "fetch_memory", fetch_memory) + assert result is model_instance + assert result.parameters == {"temperature": 0.3} + assert result.stop == ("Human:",) + assert result.model_type_instance is model_type_instance + fetch_model_config.assert_called_once_with( + node_data_model=node_data_model, + credentials_provider=sentinel.credentials_provider, + model_factory=sentinel.model_factory, + ) - result = factory._build_memory_for_llm_node( - node_data=SimpleNamespace(memory=memory_config), - model_instance=sentinel.model_instance, - ) + def test_propagates_fetch_model_config_errors(self, monkeypatch: pytest.MonkeyPatch, factory): + fetch_model_config = MagicMock(side_effect=ValueError("broken model config")) + monkeypatch.setattr(node_factory, "fetch_model_config", fetch_model_config) - assert result is sentinel.memory - fetch_memory.assert_called_once_with( - conversation_id=None, - app_id="app-id", - node_data_memory=memory_config, - model_instance=sentinel.model_instance, - ) + with pytest.raises(ValueError, match="broken model config"): + factory._build_model_instance_for_llm_node(SimpleNamespace(model=sentinel.node_data_model)) + + class TestDifyNodeFactoryMemory: + @pytest.fixture + def factory(self): + factory = object.__new__(node_factory.DifyNodeFactory) + factory._dify_context = SimpleNamespace(app_id="app-id") + factory.runtime_state = SimpleNamespace(variable_pool=MagicMock()) + return factory + + def test_returns_none_when_memory_is_not_configured(self, factory): + result = factory._build_memory_for_llm_node( + node_data=SimpleNamespace(memory=None), + model_instance=sentinel.model_instance, + ) + + assert result is None + factory.runtime_state.variable_pool.get.assert_not_called() + + def test_uses_string_segment_conversation_id(self, monkeypatch: pytest.MonkeyPatch, factory): + memory_config = sentinel.memory_config + factory.runtime_state.variable_pool.get.return_value = StringSegment(value="conversation-id") + fetch_memory = MagicMock(return_value=sentinel.memory) + monkeypatch.setattr(node_factory, "fetch_memory", fetch_memory) + + result = factory._build_memory_for_llm_node( + node_data=SimpleNamespace(memory=memory_config), + model_instance=sentinel.model_instance, + ) + + assert result is sentinel.memory + factory.runtime_state.variable_pool.get.assert_called_once_with(("sys", "conversation_id")) + fetch_memory.assert_called_once_with( + conversation_id="conversation-id", + app_id="app-id", + node_data_memory=memory_config, + model_instance=sentinel.model_instance, + ) + + def test_ignores_non_string_segment_conversation_ids(self, monkeypatch: pytest.MonkeyPatch, factory): + memory_config = sentinel.memory_config + factory.runtime_state.variable_pool.get.return_value = sentinel.segment + fetch_memory = MagicMock(return_value=sentinel.memory) + monkeypatch.setattr(node_factory, "fetch_memory", fetch_memory) + + result = factory._build_memory_for_llm_node( + node_data=SimpleNamespace(memory=memory_config), + model_instance=sentinel.model_instance, + ) + + assert result is sentinel.memory + fetch_memory.assert_called_once_with( + conversation_id=None, + app_id="app-id", + node_data_memory=memory_config, + model_instance=sentinel.model_instance, + ) diff --git a/api/tests/unit_tests/core/workflow/test_node_runtime.py b/api/tests/unit_tests/core/workflow/test_node_runtime.py index a7428a16469..89d184a71f4 100644 --- a/api/tests/unit_tests/core/workflow/test_node_runtime.py +++ b/api/tests/unit_tests/core/workflow/test_node_runtime.py @@ -1,7 +1,16 @@ -from collections.abc import Iterator, Mapping +from __future__ import annotations + +import json +import os +import sys +import time +import types +import uuid +from collections.abc import Generator, Iterator, Mapping from datetime import UTC, datetime from decimal import Decimal from types import SimpleNamespace +from typing import override from unittest.mock import MagicMock, Mock, patch, sentinel from uuid import uuid4 @@ -11,12 +20,23 @@ from sqlalchemy.orm import Session, sessionmaker from core.app.entities.app_invoke_entities import DIFY_RUN_CONTEXT_KEY, DifyRunContext, InvokeFrom, UserFrom from core.app.file_access import FileAccessScope, bind_file_access_scope, grant_retriever_segment_access +from core.app.llm.model_access import DifyCredentialsProvider, DifyModelFactory +from core.callback_handler.workflow_tool_callback_handler import DifyWorkflowCallbackHandler +from core.helper.ssrf_proxy import graphon_ssrf_proxy from core.llm_generator.output_parser.errors import OutputParserError +from core.llm_generator.output_parser.structured_output import _parse_structured_output from core.model_manager import QuotaManagedModelInstance -from core.plugin.impl.exc import PluginLLMPollingUnsupportedError +from core.plugin.impl.exc import PluginDaemonClientSideError, PluginInvokeError, PluginLLMPollingUnsupportedError from core.plugin.impl.model import PluginModelClient from core.plugin.impl.model_runtime import PluginModelRuntime from core.plugin.plugin_service import PluginService +from core.tools.entities.tool_entities import ToolInvokeMessage +from core.tools.entities.tool_entities import ToolProviderType as CoreToolProviderType +from core.tools.errors import ToolInvokeError +from core.tools.tool_engine import ToolEngine +from core.tools.tool_manager import ToolManager +from core.tools.utils.configuration import ToolParameterConfigurationManager +from core.tools.utils.message_transformer import ToolFileMessageTransformer from core.workflow import node_runtime from core.workflow.file_reference import parse_file_reference from core.workflow.human_input_adapter import ( @@ -27,6 +47,7 @@ from core.workflow.human_input_adapter import ( WebAppDeliveryMethod, _WebAppDeliveryConfig, ) +from core.workflow.node_factory import DifyNodeFactory from core.workflow.node_runtime import ( DifyFileReferenceFactory, DifyHumanInputNodeRuntime, @@ -41,32 +62,1231 @@ from core.workflow.node_runtime import ( resolve_dify_run_context, ) from core.workflow.nodes.human_input.entities import FileInputConfig, FileListInputConfig, HumanInputNodeData +from core.workflow.system_variables import build_system_variables +from extensions.ext_database import db +from graphon.enums import WorkflowNodeExecutionStatus from graphon.file import File, FileTransferMethod, FileType +from graphon.graph import Graph +from graphon.model_runtime.entities import ( + AssistantPromptMessage as _parameter_extractor_execution_behavior_AssistantPromptMessage, +) +from graphon.model_runtime.entities import UserPromptMessage from graphon.model_runtime.entities.common_entities import I18nObject -from graphon.model_runtime.entities.llm_entities import LLMPollingResult, LLMPollingStatus -from graphon.model_runtime.entities.message_entities import AssistantPromptMessage +from graphon.model_runtime.entities.llm_entities import LLMPollingResult, LLMPollingStatus, LLMUsage +from graphon.model_runtime.entities.message_entities import ( + AssistantPromptMessage as _node_runtime_AssistantPromptMessage, +) +from graphon.model_runtime.entities.message_entities import PromptMessage from graphon.model_runtime.entities.model_entities import AIModelEntity, FetchFrom, ModelFeature, ModelType +from graphon.model_runtime.memory.prompt_message_memory import PromptMessageMemory from graphon.model_runtime.model_providers.base.large_language_model import LargeLanguageModel +from graphon.node_events import StreamCompletedEvent +from graphon.nodes.llm.entities import LLMNodeData +from graphon.nodes.llm.node import LLMNode from graphon.nodes.llm.runtime_protocols import LLMPollingCapableProtocol +from graphon.nodes.parameter_extractor.entities import ParameterExtractorNodeData +from graphon.nodes.parameter_extractor.parameter_extractor_node import ParameterExtractorNode from graphon.nodes.tool.entities import ToolNodeData, ToolProviderType -from graphon.nodes.tool.exc import ToolRuntimeResolutionError -from graphon.nodes.tool_runtime_entities import ToolRuntimeHandle +from graphon.nodes.tool.exc import ToolRuntimeInvocationError, ToolRuntimeResolutionError +from graphon.nodes.tool.tool_node import ToolNode +from graphon.nodes.tool_runtime_entities import ToolRuntimeHandle, ToolRuntimeMessage +from graphon.runtime import RuntimeState, VariablePool from graphon.variables.segments import ArrayFileSegment, FileSegment from models.base import TypeBase from models.dataset import SegmentAttachmentBinding from models.enums import CreatorUserRole from models.model import StorageType, UploadFile from models.tools import ToolFile -from tests.workflow_test_utils import build_test_run_context +from tests.helpers.unit.core.model_fixtures import make_model_config, make_model_instance +from tests.helpers.unit.core.workflow.nodes.parameter_extractor.fixtures.model import get_mocked_fetch_model_instance +from tests.workflow_test_utils import build_test_graph_init_params, build_test_run_context, build_test_variable_pool + +pytest_plugins = ("tests.helpers.unit.core.workflow.nodes.parameter_extractor.fixtures.plugin_daemon",) + +"""Unit coverage for LLM node execution with model boundaries replaced by test doubles.""" -@pytest.fixture -def attachment_session(sqlite_engine: Engine) -> Iterator[Session]: - """Provide real attachment and upload-file persistence to node runtime tests.""" +"""FOR MOCK FIXTURES, DO NOT REMOVE""" - TypeBase.metadata.create_all(sqlite_engine, tables=[SegmentAttachmentBinding.__table__, UploadFile.__table__]) - with Session(sqlite_engine, expire_on_commit=False) as session: - yield session + +def init_llm_node(config: dict[str, object]) -> LLMNode: + graph_config = { + "edges": [ + { + "id": "start-source-next-target", + "source": "start", + "target": "llm", + }, + ], + "nodes": [{"data": {"type": "start", "title": "Start"}, "id": "start"}, config], + } + + # Use proper UUIDs for database compatibility + tenant_id = "9d2074fc-6f86-45a9-b09d-6ecc63b9056b" + app_id = "9d2074fc-6f86-45a9-b09d-6ecc63b9056c" + workflow_id = "9d2074fc-6f86-45a9-b09d-6ecc63b9056d" + user_id = "9d2074fc-6f86-45a9-b09d-6ecc63b9056e" + + init_params = build_test_graph_init_params( + workflow_id=workflow_id, + graph_config=graph_config, + tenant_id=tenant_id, + app_id=app_id, + user_id=user_id, + user_from=UserFrom.ACCOUNT, + invoke_from=InvokeFrom.DEBUGGER, + call_depth=0, + ) + + # construct variable pool + variable_pool = VariablePool.from_bootstrap( + system_variables=build_system_variables( + user_id="aaa", + app_id=app_id, + workflow_id=workflow_id, + files=[], + query="what's the weather today?", + conversation_id="abababa", + ), + user_inputs={}, + environment_variables=[], + conversation_variables=[], + ) + variable_pool.add(["abc", "output"], "sunny") + + graph_runtime_state = RuntimeState( + workflow_id=workflow_id, variable_pool=variable_pool, start_at=time.perf_counter() + ) + run_context = resolve_dify_run_context(init_params.run_context) + + node = LLMNode( + node_id=str(uuid.uuid4()), + data=LLMNodeData.model_validate(config["data"]), + init_params=init_params, + runtime_state=graph_runtime_state, + credentials_provider=DifyCredentialsProvider(run_context=run_context), + model_factory=DifyModelFactory(run_context=run_context), + model_instance=DifyPreparedLLM(make_model_instance(provider="openai", model="gpt-3.5-turbo")), + llm_file_saver=build_dify_llm_file_saver(run_context=run_context, http_client=graphon_ssrf_proxy), + prompt_message_serializer=DifyPromptMessageSerializer(), + http_client=graphon_ssrf_proxy, + ) + + return node + + +def _llm_node_execution_behavior_mock_db_session_close(monkeypatch: pytest.MonkeyPatch) -> None: + monkeypatch.setattr(db.session, "close", MagicMock()) + + +class TestLlmNodeExecutionBehaviorCases: + @staticmethod + def test_execute_llm(monkeypatch: pytest.MonkeyPatch) -> None: + node = init_llm_node( + config={ + "id": "llm", + "data": { + "title": "123", + "type": "llm", + "model": { + "provider": "openai", + "name": "gpt-3.5-turbo", + "mode": "chat", + "completion_params": {}, + }, + "prompt_template": [ + { + "role": "system", + "text": "you are a helpful assistant.\ntoday's weather is {{#abc.output#}}.", + }, + {"role": "user", "text": "{{#sys.query#}}"}, + ], + "memory": None, + "context": {"enabled": False}, + "vision": {"enabled": False}, + }, + }, + ) + + _llm_node_execution_behavior_mock_db_session_close(monkeypatch) + + def build_model_instance() -> DifyPreparedLLM: + from decimal import Decimal + from unittest.mock import MagicMock + + from graphon.model_runtime.entities.llm_entities import LLMResult, LLMUsage + from graphon.model_runtime.entities.message_entities import AssistantPromptMessage + + model_instance = make_model_instance(provider="openai", model="gpt-3.5-turbo") + schema = make_model_config(provider="openai", model="gpt-3.5-turbo", mode="chat").model_schema + monkeypatch.setattr(model_instance, "get_model_schema", MagicMock(return_value=schema)) + mock_usage = LLMUsage( + prompt_tokens=30, + prompt_unit_price=Decimal("0.001"), + prompt_price_unit=Decimal(1000), + prompt_price=Decimal("0.00003"), + completion_tokens=20, + completion_unit_price=Decimal("0.002"), + completion_price_unit=Decimal(1000), + completion_price=Decimal("0.00004"), + total_tokens=50, + total_price=Decimal("0.00007"), + currency="USD", + latency=0.5, + ) + mock_message = AssistantPromptMessage(content="Test response from mock") + mock_llm_result = LLMResult( + model="gpt-3.5-turbo", + prompt_messages=[], + message=mock_message, + usage=mock_usage, + ) + monkeypatch.setattr(model_instance, "invoke_llm", MagicMock(return_value=mock_llm_result)) + return DifyPreparedLLM(model_instance) + + # Mock fetch_prompt_messages to avoid database calls + def mock_fetch_prompt_messages_1(**_kwargs: object) -> tuple[list[PromptMessage], list[str]]: + from graphon.model_runtime.entities.message_entities import SystemPromptMessage, UserPromptMessage + + return [ + SystemPromptMessage(content="you are a helpful assistant. today's weather is sunny."), + UserPromptMessage(content="what's the weather today?"), + ], [] + + node._model_instance = build_model_instance() + + with patch("graphon.nodes.llm.llm_utils.fetch_prompt_messages", mock_fetch_prompt_messages_1): + # execute node + result = node._run() + assert isinstance(result, Generator) + + for item in result: + if isinstance(item, StreamCompletedEvent): + assert item.node_run_result.status == WorkflowNodeExecutionStatus.SUCCEEDED, ( + f"{item.node_run_result.error_type}: {item.node_run_result.error}" + ) + assert item.node_run_result.process_data is not None + assert item.node_run_result.outputs is not None + assert item.node_run_result.outputs.get("text") is not None + assert item.node_run_result.outputs.get("usage", {})["total_tokens"] > 0 + + @staticmethod + def test_execute_llm_with_jinja2(monkeypatch: pytest.MonkeyPatch) -> None: + """ + Test execute LLM node with jinja2 + """ + node = init_llm_node( + config={ + "id": "llm", + "data": { + "title": "123", + "type": "llm", + "model": {"provider": "openai", "name": "gpt-3.5-turbo", "mode": "chat", "completion_params": {}}, + "prompt_config": { + "jinja2_variables": [ + {"variable": "sys_query", "value_selector": ["sys", "query"]}, + {"variable": "output", "value_selector": ["abc", "output"]}, + ] + }, + "prompt_template": [ + { + "role": "system", + "text": "you are a helpful assistant.\ntoday's weather is {{#abc.output#}}", + "jinja2_text": "you are a helpful assistant.\ntoday's weather is {{output}}.", + "edition_type": "jinja2", + }, + { + "role": "user", + "text": "{{#sys.query#}}", + "jinja2_text": "{{sys_query}}", + "edition_type": "basic", + }, + ], + "memory": None, + "context": {"enabled": False}, + "vision": {"enabled": False}, + }, + }, + ) + + _llm_node_execution_behavior_mock_db_session_close(monkeypatch) + + def build_model_instance() -> DifyPreparedLLM: + from decimal import Decimal + from unittest.mock import MagicMock + + from graphon.model_runtime.entities.llm_entities import LLMResult, LLMUsage + from graphon.model_runtime.entities.message_entities import AssistantPromptMessage + + model_instance = make_model_instance(provider="openai", model="gpt-3.5-turbo") + schema = make_model_config(provider="openai", model="gpt-3.5-turbo", mode="chat").model_schema + monkeypatch.setattr(model_instance, "get_model_schema", MagicMock(return_value=schema)) + mock_usage = LLMUsage( + prompt_tokens=30, + prompt_unit_price=Decimal("0.001"), + prompt_price_unit=Decimal(1000), + prompt_price=Decimal("0.00003"), + completion_tokens=20, + completion_unit_price=Decimal("0.002"), + completion_price_unit=Decimal(1000), + completion_price=Decimal("0.00004"), + total_tokens=50, + total_price=Decimal("0.00007"), + currency="USD", + latency=0.5, + ) + mock_message = AssistantPromptMessage(content="Test response: sunny weather and what's the weather today?") + mock_llm_result = LLMResult( + model="gpt-3.5-turbo", + prompt_messages=[], + message=mock_message, + usage=mock_usage, + ) + monkeypatch.setattr(model_instance, "invoke_llm", MagicMock(return_value=mock_llm_result)) + return DifyPreparedLLM(model_instance) + + # Mock fetch_prompt_messages to avoid database calls + def mock_fetch_prompt_messages_2(**_kwargs: object) -> tuple[list[PromptMessage], list[str]]: + from graphon.model_runtime.entities.message_entities import SystemPromptMessage, UserPromptMessage + + return [ + SystemPromptMessage(content="you are a helpful assistant. today's weather is sunny."), + UserPromptMessage(content="what's the weather today?"), + ], [] + + node._model_instance = build_model_instance() + + with patch("graphon.nodes.llm.llm_utils.fetch_prompt_messages", mock_fetch_prompt_messages_2): + # execute node + result = node._run() + + for item in result: + if isinstance(item, StreamCompletedEvent): + assert item.node_run_result.status == WorkflowNodeExecutionStatus.SUCCEEDED + assert item.node_run_result.process_data is not None + assert "sunny" in json.dumps(item.node_run_result.process_data) + assert "what's the weather today?" in json.dumps(item.node_run_result.process_data) + + @staticmethod + def test_extract_json() -> None: + llm_texts = [ + '\n\n{"name": "test", "age": 123', # resoning model (deepseek-r1) + '{"name":"test","age":123}', # json schema model (gpt-4o) + '{\n "name": "test",\n "age": 123\n}', # small model (llama-3.2-1b) + '```json\n{"name": "test", "age": 123}\n```', # json markdown (deepseek-chat) + '{"name":"test",age:123}', # without quotes (qwen-2.5-0.5b) + ] + result = {"name": "test", "age": 123} + assert all(_parse_structured_output(item) == result for item in llm_texts) + + +"""Unit coverage for parameter extraction with the plugin-model boundary replaced.""" + + +def get_mocked_fetch_memory(memory_text: str) -> MagicMock: + class MemoryMock(PromptMessageMemory): + @override + def get_history_prompt_messages( + self, + max_token_limit: int = 2000, + message_limit: int | None = None, + ) -> list[PromptMessage]: + return [ + UserPromptMessage(content=memory_text), + _parameter_extractor_execution_behavior_AssistantPromptMessage(content="mocked answer"), + ] + + return MagicMock(return_value=MemoryMock()) + + +def init_parameter_extractor_node( + config: dict[str, object], memory: PromptMessageMemory | None = None +) -> ParameterExtractorNode: + graph_config = { + "edges": [ + { + "id": "start-source-next-target", + "source": "start", + "target": "llm", + }, + ], + "nodes": [{"data": {"type": "start", "title": "Start"}, "id": "start"}, config], + } + + init_params = build_test_graph_init_params( + workflow_id="1", + graph_config=graph_config, + tenant_id="1", + app_id="1", + user_id="1", + user_from=UserFrom.ACCOUNT, + invoke_from=InvokeFrom.DEBUGGER, + call_depth=0, + ) + + # construct variable pool + variable_pool = VariablePool.from_bootstrap( + system_variables=build_system_variables( + user_id="aaa", files=[], query="what's the weather in SF", conversation_id="abababa" + ), + user_inputs={}, + environment_variables=[], + conversation_variables=[], + ) + variable_pool.add(["a", "args1"], 1) + variable_pool.add(["a", "args2"], 2) + + graph_runtime_state = RuntimeState( + workflow_id="test-workflow", variable_pool=variable_pool, start_at=time.perf_counter() + ) + + run_context = resolve_dify_run_context(init_params.run_context) + node = ParameterExtractorNode( + node_id=str(uuid.uuid4()), + data=ParameterExtractorNodeData.model_validate(config["data"]), + init_params=init_params, + runtime_state=graph_runtime_state, + credentials_provider=DifyCredentialsProvider(run_context=run_context), + model_factory=DifyModelFactory(run_context=run_context), + model_instance=DifyPreparedLLM(make_model_instance(provider="openai", model="gpt-3.5-turbo")), + memory=memory, + prompt_message_serializer=DifyPromptMessageSerializer(), + ) + return node + + +def _parameter_extractor_execution_behavior_mock_db_session_close(monkeypatch: pytest.MonkeyPatch) -> None: + monkeypatch.setattr(db.session, "close", MagicMock()) + + +class TestParameterExtractorExecutionBehaviorCases: + @staticmethod + @pytest.mark.usefixtures("setup_model_mock") + def test_function_calling_parameter_extractor(monkeypatch: pytest.MonkeyPatch) -> None: + """ + Test function calling for parameter extractor. + """ + node = init_parameter_extractor_node( + config={ + "id": "llm", + "data": { + "title": "123", + "type": "parameter-extractor", + "model": { + "provider": "langgenius/openai/openai", + "name": "gpt-3.5-turbo", + "mode": "chat", + "completion_params": {}, + }, + "query": ["sys", "query"], + "parameters": [{"name": "location", "type": "string", "description": "location", "required": True}], + "instruction": "", + "reasoning_mode": "function_call", + "memory": None, + }, + } + ) + + node._model_instance = get_mocked_fetch_model_instance( + provider="langgenius/openai/openai", + model="gpt-3.5-turbo", + mode="chat", + credentials={"openai_api_key": os.environ.get("OPENAI_API_KEY")}, + )() + _parameter_extractor_execution_behavior_mock_db_session_close(monkeypatch) + + result = node._run() + + assert result.status == WorkflowNodeExecutionStatus.SUCCEEDED + assert result.outputs is not None + assert result.outputs.get("location") == "kawaii" + assert result.outputs.get("__reason") is None + + @staticmethod + @pytest.mark.usefixtures("setup_model_mock") + def test_instructions(monkeypatch: pytest.MonkeyPatch) -> None: + """ + Test chat parameter extractor. + """ + node = init_parameter_extractor_node( + config={ + "id": "llm", + "data": { + "title": "123", + "type": "parameter-extractor", + "model": { + "provider": "langgenius/openai/openai", + "name": "gpt-3.5-turbo", + "mode": "chat", + "completion_params": {}, + }, + "query": ["sys", "query"], + "parameters": [{"name": "location", "type": "string", "description": "location", "required": True}], + "reasoning_mode": "function_call", + "instruction": "{{#sys.query#}}", + "memory": None, + }, + }, + ) + + node._model_instance = get_mocked_fetch_model_instance( + provider="langgenius/openai/openai", + model="gpt-3.5-turbo", + mode="chat", + credentials={"openai_api_key": os.environ.get("OPENAI_API_KEY")}, + )() + _parameter_extractor_execution_behavior_mock_db_session_close(monkeypatch) + + result = node._run() + + assert result.status == WorkflowNodeExecutionStatus.SUCCEEDED + assert result.outputs is not None + assert result.outputs.get("location") == "kawaii" + assert result.outputs.get("__reason") is None + + process_data = result.process_data + + assert process_data is not None + process_data.get("prompts") + + for prompt in process_data.get("prompts", []): + if prompt.get("role") == "system": + assert "what's the weather in SF" in prompt.get("text") + + @staticmethod + @pytest.mark.usefixtures("setup_model_mock") + def test_chat_parameter_extractor(monkeypatch: pytest.MonkeyPatch) -> None: + """ + Test chat parameter extractor. + """ + node = init_parameter_extractor_node( + config={ + "id": "llm", + "data": { + "title": "123", + "type": "parameter-extractor", + "model": { + "provider": "langgenius/openai/openai", + "name": "gpt-3.5-turbo", + "mode": "chat", + "completion_params": {}, + }, + "query": ["sys", "query"], + "parameters": [{"name": "location", "type": "string", "description": "location", "required": True}], + "reasoning_mode": "prompt", + "instruction": "", + "memory": None, + }, + }, + ) + + node._model_instance = get_mocked_fetch_model_instance( + provider="langgenius/openai/openai", + model="gpt-3.5-turbo", + mode="chat", + credentials={"openai_api_key": os.environ.get("OPENAI_API_KEY")}, + )() + _parameter_extractor_execution_behavior_mock_db_session_close(monkeypatch) + + result = node._run() + + assert result.status == WorkflowNodeExecutionStatus.SUCCEEDED + assert result.outputs is not None + assert result.outputs.get("location") == "" + assert ( + result.outputs.get("__reason") + == "Failed to extract result from function call or text response, using empty result." + ) + assert result.process_data is not None + prompts = result.process_data.get("prompts", []) + + for prompt in prompts: + if prompt.get("role") == "user": + if "" in prompt.get("text"): + assert '\n{"type": "object"' in prompt.get("text") + + @staticmethod + @pytest.mark.usefixtures("setup_model_mock") + def test_completion_parameter_extractor(monkeypatch: pytest.MonkeyPatch) -> None: + """ + Test completion parameter extractor. + """ + node = init_parameter_extractor_node( + config={ + "id": "llm", + "data": { + "title": "123", + "type": "parameter-extractor", + "model": { + "provider": "langgenius/openai/openai", + "name": "gpt-3.5-turbo-instruct", + "mode": "completion", + "completion_params": {}, + }, + "query": ["sys", "query"], + "parameters": [{"name": "location", "type": "string", "description": "location", "required": True}], + "reasoning_mode": "prompt", + "instruction": "{{#sys.query#}}", + "memory": None, + }, + }, + ) + + node._model_instance = get_mocked_fetch_model_instance( + provider="langgenius/openai/openai", + model="gpt-3.5-turbo-instruct", + mode="completion", + credentials={"openai_api_key": os.environ.get("OPENAI_API_KEY")}, + )() + _parameter_extractor_execution_behavior_mock_db_session_close(monkeypatch) + + result = node._run() + + assert result.status == WorkflowNodeExecutionStatus.SUCCEEDED + assert result.outputs is not None + assert result.outputs.get("location") == "" + assert ( + result.outputs.get("__reason") + == "Failed to extract result from function call or text response, using empty result." + ) + assert result.process_data is not None + assert len(result.process_data.get("prompts", [])) == 1 + assert "SF" in result.process_data.get("prompts", [])[0].get("text") + + @staticmethod + def test_extract_json_response() -> None: + """ + Test extract json response. + """ + + node = init_parameter_extractor_node( + config={ + "id": "llm", + "data": { + "title": "123", + "type": "parameter-extractor", + "model": { + "provider": "langgenius/openai/openai", + "name": "gpt-3.5-turbo-instruct", + "mode": "completion", + "completion_params": {}, + }, + "query": ["sys", "query"], + "parameters": [{"name": "location", "type": "string", "description": "location", "required": True}], + "reasoning_mode": "prompt", + "instruction": "{{#sys.query#}}", + "memory": None, + }, + }, + ) + + result = node._extract_complete_json_response(""" + uwu{ovo} + { + "location": "kawaii" + } + hello world. + """) + + assert result is not None + assert result["location"] == "kawaii" + + @staticmethod + def test_extract_json_from_tool_call() -> None: + """ + Test extract json response. + """ + + node = init_parameter_extractor_node( + config={ + "id": "llm", + "data": { + "title": "123", + "type": "parameter-extractor", + "model": { + "provider": "langgenius/openai/openai", + "name": "gpt-3.5-turbo-instruct", + "mode": "completion", + "completion_params": {}, + }, + "query": ["sys", "query"], + "parameters": [{"name": "location", "type": "string", "description": "location", "required": True}], + "reasoning_mode": "prompt", + "instruction": "{{#sys.query#}}", + "memory": None, + }, + }, + ) + + result = node._extract_json_from_tool_call( + _parameter_extractor_execution_behavior_AssistantPromptMessage.ToolCall( + id="llm", + type="parameter-extractor", + function=_parameter_extractor_execution_behavior_AssistantPromptMessage.ToolCall.ToolCallFunction( + name="foo", arguments="""{"location":"kawaii"}{"location": 1}""" + ), + ) + ) + + assert result is not None + assert result["location"] == "kawaii" + + @staticmethod + @pytest.mark.usefixtures("setup_model_mock") + def test_chat_parameter_extractor_with_memory(monkeypatch: pytest.MonkeyPatch) -> None: + """ + Test chat parameter extractor with memory. + """ + node = init_parameter_extractor_node( + config={ + "id": "llm", + "data": { + "title": "123", + "type": "parameter-extractor", + "model": { + "provider": "langgenius/openai/openai", + "name": "gpt-3.5-turbo", + "mode": "chat", + "completion_params": {}, + }, + "query": ["sys", "query"], + "parameters": [{"name": "location", "type": "string", "description": "location", "required": True}], + "reasoning_mode": "prompt", + "instruction": "", + "memory": {"window": {"enabled": True, "size": 50}}, + }, + }, + memory=get_mocked_fetch_memory("customized memory")(), + ) + + node._model_instance = get_mocked_fetch_model_instance( + provider="langgenius/openai/openai", + model="gpt-3.5-turbo", + mode="chat", + credentials={"openai_api_key": os.environ.get("OPENAI_API_KEY")}, + )() + _parameter_extractor_execution_behavior_mock_db_session_close(monkeypatch) + + result = node._run() + + assert result.status == WorkflowNodeExecutionStatus.SUCCEEDED + assert result.outputs is not None + assert result.outputs.get("location") == "" + assert ( + result.outputs.get("__reason") + == "Failed to extract result from function call or text response, using empty result." + ) + assert result.process_data is not None + prompts = result.process_data.get("prompts", []) + + latest_role = None + for prompt in prompts: + if prompt.get("role") == "user": + if "" in prompt.get("text"): + assert '\n{"type": "object"' in prompt.get("text") + elif prompt.get("role") == "system": + assert "customized memory" in prompt.get("text") + + if latest_role is not None: + assert latest_role != prompt.get("role") + + if prompt.get("role") in {"user", "assistant"}: + latest_role = prompt.get("role") + + +"""Unit coverage for tool-node invocation with the tool boundary patched.""" + + +def init_tool_node(config: dict[str, object]) -> ToolNode: + graph_config = { + "edges": [ + { + "id": "start-source-next-target", + "source": "start", + "target": "1", + }, + ], + "nodes": [{"data": {"type": "start", "title": "Start"}, "id": "start"}, config], + } + + init_params = build_test_graph_init_params( + workflow_id="1", + graph_config=graph_config, + tenant_id="1", + app_id="1", + user_id="1", + user_from=UserFrom.ACCOUNT, + invoke_from=InvokeFrom.DEBUGGER, + call_depth=0, + ) + + # construct variable pool + variable_pool = VariablePool.from_bootstrap( + system_variables=build_system_variables(user_id="aaa", files=[]), + user_inputs={}, + environment_variables=[], + conversation_variables=[], + ) + + graph_runtime_state = RuntimeState( + workflow_id="test-workflow", variable_pool=variable_pool, start_at=time.perf_counter() + ) + + # Create node factory + node_factory = DifyNodeFactory( + init_params=init_params, + runtime_state=graph_runtime_state, + ) + + Graph.init(graph_config=graph_config, node_factory=node_factory, root_node_id="start") + + tool_file_manager = DifyToolFileManager(init_params.run_context) + + node = ToolNode( + node_id=str(uuid.uuid4()), + data=ToolNodeData.model_validate(config["data"]), + init_params=init_params, + runtime_state=graph_runtime_state, + tool_file_manager=tool_file_manager, + runtime=DifyToolNodeRuntime(init_params.run_context), + ) + node.bind_execution_id(str(uuid.uuid4())) + return node + + +class TestToolNodeInvocationBehaviorCases: + @staticmethod + def test_tool_variable_invoke() -> None: + node = init_tool_node( + config={ + "id": "1", + "data": { + "type": "tool", + "title": "a", + "desc": "a", + "provider_id": "time", + "provider_type": "builtin", + "provider_name": "time", + "tool_name": "current_time", + "tool_label": "current_time", + "tool_configurations": {}, + "tool_parameters": {}, + }, + } + ) + + with patch.object( + ToolParameterConfigurationManager, + "decrypt_tool_parameters", + return_value={"format": "%Y-%m-%d %H:%M:%S"}, + ): + node.runtime_state.variable_pool.add(["1", "args1"], "1+1") + + # execute node + result = node._run() + for item in result: + if isinstance(item, StreamCompletedEvent): + assert item.node_run_result.status == WorkflowNodeExecutionStatus.SUCCEEDED + assert item.node_run_result.outputs is not None + assert item.node_run_result.outputs.get("text") is not None + + @staticmethod + def test_tool_mixed_invoke() -> None: + node = init_tool_node( + config={ + "id": "1", + "data": { + "type": "tool", + "title": "a", + "desc": "a", + "provider_id": "time", + "provider_type": "builtin", + "provider_name": "time", + "tool_name": "current_time", + "tool_label": "current_time", + "tool_configurations": { + "format": "%Y-%m-%d %H:%M:%S", + }, + "tool_parameters": {}, + }, + } + ) + + with patch.object( + ToolParameterConfigurationManager, + "decrypt_tool_parameters", + return_value={"format": "%Y-%m-%d %H:%M:%S"}, + ): + # execute node + result = node._run() + for item in result: + if isinstance(item, StreamCompletedEvent): + assert item.node_run_result.status == WorkflowNodeExecutionStatus.SUCCEEDED + assert item.node_run_result.outputs is not None + assert item.node_run_result.outputs.get("text") is not None + + +def _build_tool_node_data() -> ToolNodeData: + return ToolNodeData.model_validate( + { + "type": "tool", + "title": "Tool", + "provider_id": "provider", + "provider_type": ToolProviderType.BUILT_IN, + "provider_name": "provider", + "tool_name": "lookup", + "tool_label": "Lookup", + "tool_configurations": {}, + "tool_parameters": {}, + } + ) + + +class TestToolNodeRuntimeCases: + @staticmethod + @pytest.fixture + def runtime(monkeypatch: pytest.MonkeyPatch, sqlite_engine: Engine) -> DifyToolNodeRuntime: + module_name = "core.ops.ops_trace_manager" + if module_name not in sys.modules: + ops_stub = types.ModuleType(module_name) + ops_stub.TraceQueueManager = object # pragma: no cover - stub attribute + ops_stub.TraceTask = object # pragma: no cover - stub attribute + monkeypatch.setitem(sys.modules, module_name, ops_stub) + + init_params = build_test_graph_init_params( + workflow_id="workflow-id", + graph_config={"nodes": [], "edges": []}, + tenant_id="tenant-id", + app_id="app-id", + user_id="user-id", + user_from="account", + invoke_from="debugger", + call_depth=0, + ) + session_maker = sessionmaker(sqlite_engine, expire_on_commit=False) + return DifyToolNodeRuntime(init_params.run_context, session_maker=session_maker) + + @staticmethod + def test_invoke_creates_callback_and_converts_messages(runtime: DifyToolNodeRuntime) -> None: + core_message = ToolInvokeMessage( + type=ToolInvokeMessage.MessageType.LINK, + message=ToolInvokeMessage.TextMessage(text="https://dify.ai"), + meta=None, + ) + variable_pool: VariablePool = build_test_variable_pool( + variables=build_system_variables(conversation_id="conversation-id") + ) + workflow_tool = MagicMock() + + with ( + patch.object(ToolManager, "get_workflow_tool_runtime", return_value=workflow_tool), + patch.object(ToolEngine, "generic_invoke", return_value=iter([core_message])) as generic_invoke_mock, + patch.object( + ToolFileMessageTransformer, + "transform_tool_invoke_messages", + side_effect=lambda *, messages, **_: messages, + ) as transform_tool_messages, + ): + tool_runtime = runtime.get_runtime( + node_id="node-id", + node_data=_build_tool_node_data(), + variable_pool=variable_pool, + ) + messages = list( + runtime.invoke( + tool_runtime=tool_runtime, + tool_parameters={}, + workflow_call_depth=0, + provider_name="provider", + ) + ) + + assert not hasattr(tool_runtime, "conversation_id") + assert len(messages) == 1 + graph_message = messages[0] + assert graph_message.type == ToolRuntimeMessage.MessageType.LINK + assert isinstance(graph_message.message, ToolRuntimeMessage.TextMessage) + assert graph_message.message.text == "https://dify.ai" + + callback = generic_invoke_mock.call_args.kwargs["workflow_tool_callback"] + assert isinstance(callback, DifyWorkflowCallbackHandler) + assert generic_invoke_mock.call_args.kwargs["session"] is not None + assert generic_invoke_mock.call_args.kwargs["conversation_id"] == "conversation-id" + + transform_kwargs = transform_tool_messages.call_args.kwargs + assert transform_kwargs["conversation_id"] == "conversation-id" + + @staticmethod + def test_invoke_maps_plugin_errors_to_graph_errors(runtime: DifyToolNodeRuntime) -> None: + invoke_error = PluginInvokeError('{"error_type":"RateLimit","message":"too many"}') + + with patch.object(ToolEngine, "generic_invoke", side_effect=invoke_error): + with pytest.raises(ToolRuntimeInvocationError, match="An error occurred in the provider"): + list( + runtime.invoke( + tool_runtime=ToolRuntimeHandle(raw=MagicMock()), + tool_parameters={}, + workflow_call_depth=0, + provider_name="provider", + ) + ) + + @staticmethod + def test_get_usage_normalizes_dict_payload(runtime: DifyToolNodeRuntime) -> None: + usage_payload = LLMUsage.empty_usage().model_dump() + usage_payload["total_tokens"] = 42 + + usage = runtime.get_usage( + tool_runtime=ToolRuntimeHandle(raw=SimpleNamespace(latest_usage=usage_payload)), + ) + + assert usage.total_tokens == 42 + + @staticmethod + def test_get_runtime_converts_graph_provider_type_for_tool_manager(runtime: DifyToolNodeRuntime) -> None: + node_data = _build_tool_node_data() + + with patch.object(ToolManager, "get_workflow_tool_runtime", return_value=MagicMock()) as runtime_mock: + tool_runtime = runtime.get_runtime(node_id="node-id", node_data=node_data, variable_pool=None) + + assert not hasattr(tool_runtime, "conversation_id") + workflow_tool = runtime_mock.call_args.args[3] + assert workflow_tool.provider_type == CoreToolProviderType.BUILT_IN + + @staticmethod + def test_get_runtime_stores_parent_trace_context_for_workflow_tools( + runtime: DifyToolNodeRuntime, + ) -> None: + variable_pool: VariablePool = build_test_variable_pool( + variables=build_system_variables( + conversation_id="conversation-id", + workflow_execution_id="workflow-run-id", + ) + ) + workflow_runtime = MagicMock() + workflow_runtime.runtime.runtime_parameters = {} + node_data = ToolNodeData.model_validate( + { + "type": "tool", + "title": "Tool", + "provider_id": "provider", + "provider_type": ToolProviderType.WORKFLOW, + "provider_name": "provider", + "tool_name": "lookup", + "tool_label": "Lookup", + "tool_configurations": {}, + "tool_parameters": {}, + } + ) + + with patch.object(ToolManager, "get_workflow_tool_runtime", return_value=workflow_runtime): + tool_runtime = runtime.get_runtime( + node_id="node-id", + node_data=node_data, + variable_pool=variable_pool, + node_execution_id="node-execution-id", + ) + + assert tool_runtime.raw.parent_trace_context.model_dump() == { + "parent_workflow_run_id": "workflow-run-id", + "parent_node_execution_id": "node-execution-id", + } + assert workflow_runtime.runtime.runtime_parameters == {} + + @staticmethod + def test_get_runtime_stores_trace_session_id_for_workflow_tools( + runtime: DifyToolNodeRuntime, + ) -> None: + variable_pool: VariablePool = build_test_variable_pool( + variables=build_system_variables( + conversation_id="conversation-id", + workflow_execution_id="workflow-run-id", + ) + ) + workflow_runtime = MagicMock() + workflow_runtime.runtime.runtime_parameters = {} + runtime._run_context.trace_session_id = "session-1" + node_data = ToolNodeData.model_validate( + { + "type": "tool", + "title": "Tool", + "provider_id": "provider", + "provider_type": ToolProviderType.WORKFLOW, + "provider_name": "provider", + "tool_name": "lookup", + "tool_label": "Lookup", + "tool_configurations": {}, + "tool_parameters": {}, + } + ) + + with patch.object(ToolManager, "get_workflow_tool_runtime", return_value=workflow_runtime): + tool_runtime = runtime.get_runtime( + node_id="node-id", + node_data=node_data, + variable_pool=variable_pool, + node_execution_id="node-execution-id", + ) + + assert tool_runtime.raw.trace_session_id == "session-1" + assert workflow_runtime.runtime.runtime_parameters == {} + + @staticmethod + def test_get_runtime_leaves_non_workflow_tool_runtime_parameters_unchanged( + runtime: DifyToolNodeRuntime, + ) -> None: + variable_pool: VariablePool = build_test_variable_pool( + variables=build_system_variables( + conversation_id="conversation-id", + workflow_execution_id="workflow-run-id", + ) + ) + builtin_runtime = MagicMock() + builtin_runtime.runtime.runtime_parameters = {} + + with patch.object(ToolManager, "get_workflow_tool_runtime", return_value=builtin_runtime): + runtime.get_runtime( + node_id="node-id", + node_data=_build_tool_node_data(), + variable_pool=variable_pool, + node_execution_id="node-execution-id", + ) + + assert builtin_runtime.runtime.runtime_parameters == {} + + @staticmethod + def test_get_runtime_parameters_reads_required_flags(runtime: DifyToolNodeRuntime) -> None: + tool_runtime = ToolRuntimeHandle( + raw=SimpleNamespace( + get_merged_runtime_parameters=MagicMock( + return_value=[ + SimpleNamespace(name="city", required=True), + SimpleNamespace(name="country", required=False), + ] + ) + ) + ) + + parameters = runtime.get_runtime_parameters(tool_runtime=tool_runtime) + + assert [(parameter.name, parameter.required) for parameter in parameters] == [ + ("city", True), + ("country", False), + ] + + @staticmethod + def test_get_usage_returns_empty_usage_when_tool_has_no_usage(runtime: DifyToolNodeRuntime) -> None: + usage = runtime.get_usage(tool_runtime=ToolRuntimeHandle(raw=SimpleNamespace(latest_usage=None))) + + assert usage == LLMUsage.empty_usage() + + @staticmethod + @pytest.mark.parametrize( + ("payload", "expected_type"), + [ + ( + ToolInvokeMessage.JsonMessage(json_object={"ok": True}, suppress_output=True), + ToolRuntimeMessage.JsonMessage, + ), + (ToolInvokeMessage.BlobMessage(blob=b"bytes"), ToolRuntimeMessage.BlobMessage), + ( + ToolInvokeMessage.BlobChunkMessage( + id="blob-id", + sequence=1, + total_length=5, + blob=b"hello", + end=True, + ), + ToolRuntimeMessage.BlobChunkMessage, + ), + (ToolInvokeMessage.FileMessage(file_marker="marker"), ToolRuntimeMessage.FileMessage), + ( + ToolInvokeMessage.VariableMessage(variable_name="city", variable_value="Tokyo", stream=True), + ToolRuntimeMessage.VariableMessage, + ), + ( + ToolInvokeMessage.LogMessage( + id="log-id", + label="lookup", + status=ToolInvokeMessage.LogMessage.LogStatus.SUCCESS, + data={"count": 1}, + metadata={"source": "tool"}, + ), + ToolRuntimeMessage.LogMessage, + ), + ], + ) + def test_convert_message_payload_supports_runtime_message_types( + runtime: DifyToolNodeRuntime, + payload: object, + expected_type: type[object], + ) -> None: + message = runtime._convert_message_payload(payload) + + assert isinstance(message, expected_type) + + @staticmethod + def test_convert_message_payload_rejects_unknown_types(runtime: DifyToolNodeRuntime) -> None: + with pytest.raises(TypeError, match="unsupported tool message payload"): + runtime._convert_message_payload(object()) + + @staticmethod + def test_resolve_provider_icons_prefers_builtin_tool_icons(runtime: DifyToolNodeRuntime) -> None: + plugin = SimpleNamespace( + plugin_id="langgenius/tools", + name="search", + declaration=SimpleNamespace(icon={"plugin": "icon"}), + ) + builtin_tool = SimpleNamespace( + name="langgenius/tools/search", + icon={"builtin": "icon"}, + icon_dark={"builtin": "dark"}, + ) + + with ( + patch("core.workflow.node_runtime.PluginInstaller") as installer_cls, + patch( + "core.workflow.node_runtime.BuiltinToolManageService.list_builtin_tools", return_value=[builtin_tool] + ), + ): + installer_cls.return_value.list_plugins.return_value = [plugin] + + icon, icon_dark = runtime.resolve_provider_icons(provider_name="langgenius/tools/search") + + assert icon == {"builtin": "icon"} + assert icon_dark == {"builtin": "dark"} + + @staticmethod + def test_resolve_provider_icons_returns_default_when_provider_is_unknown(runtime: DifyToolNodeRuntime) -> None: + with ( + patch("core.workflow.node_runtime.PluginInstaller") as installer_cls, + patch("core.workflow.node_runtime.BuiltinToolManageService.list_builtin_tools", return_value=[]), + ): + installer_cls.return_value.list_plugins.return_value = [] + + icon, icon_dark = runtime.resolve_provider_icons(provider_name="unknown", default_icon="fallback") + + assert icon == "fallback" + assert icon_dark is None + + @staticmethod + @pytest.mark.parametrize( + ("exc", "message"), + [ + (PluginDaemonClientSideError("bad request"), "Failed to invoke tool, error: bad request"), + (ToolInvokeError("broken"), "Failed to invoke tool provider: broken"), + (RuntimeError("unexpected"), "unexpected"), + ], + ) + def test_map_invocation_exception_normalizes_runtime_errors( + runtime: DifyToolNodeRuntime, + exc: Exception, + message: str, + ) -> None: + error = runtime._map_invocation_exception(exc, provider_name="provider") + + assert isinstance(error, ToolRuntimeInvocationError) + assert str(error) == message def _persist_attachment( @@ -181,267 +1401,236 @@ def _build_email_method(*, debug_mode: bool = False) -> EmailDeliveryMethod: ) -def test_resolve_dify_run_context_accepts_mapping_payload() -> None: - run_context = resolve_dify_run_context(_build_run_context()) +class TestNodeRuntimeCases: + @staticmethod + @pytest.fixture + def attachment_session(sqlite_engine: Engine) -> Iterator[Session]: + """Provide real attachment and upload-file persistence to node runtime tests.""" - assert run_context.tenant_id == "tenant-id" - assert run_context.invoke_from == InvokeFrom.DEBUGGER + TypeBase.metadata.create_all(sqlite_engine, tables=[SegmentAttachmentBinding.__table__, UploadFile.__table__]) + with Session(sqlite_engine, expire_on_commit=False) as session: + yield session + @staticmethod + def test_resolve_dify_run_context_accepts_mapping_payload() -> None: + run_context = resolve_dify_run_context(_build_run_context()) -def test_resolve_dify_run_context_requires_reserved_key() -> None: - with pytest.raises(ValueError, match=DIFY_RUN_CONTEXT_KEY): - resolve_dify_run_context({}) + assert run_context.tenant_id == "tenant-id" + assert run_context.invoke_from == InvokeFrom.DEBUGGER + @staticmethod + def test_resolve_dify_run_context_requires_reserved_key() -> None: + with pytest.raises(ValueError, match=DIFY_RUN_CONTEXT_KEY): + resolve_dify_run_context({}) -def test_apply_dify_debug_email_recipient_rewrites_debug_target() -> None: - method = _build_email_method(debug_mode=True) + @staticmethod + def test_apply_dify_debug_email_recipient_rewrites_debug_target() -> None: + method = _build_email_method(debug_mode=True) - updated = apply_dify_debug_email_recipient(method, enabled=True, actor_id="actor-id") + updated = apply_dify_debug_email_recipient(method, enabled=True, actor_id="actor-id") - assert isinstance(updated, EmailDeliveryMethod) - assert updated.config.recipients.include_bound_group is False - assert len(updated.config.recipients.items) == 1 - assert updated.config.recipients.items[0].reference_id == "actor-id" + assert isinstance(updated, EmailDeliveryMethod) + assert updated.config.recipients.include_bound_group is False + assert len(updated.config.recipients.items) == 1 + assert updated.config.recipients.items[0].reference_id == "actor-id" - -@pytest.mark.parametrize( - ("enabled", "method"), - [ - (False, _build_email_method(debug_mode=True)), - (True, _build_email_method(debug_mode=False)), - (True, WebAppDeliveryMethod(enabled=True, config=_WebAppDeliveryConfig())), - ], -) -def test_apply_dify_debug_email_recipient_noops_when_override_is_not_needed( - enabled: bool, - method: object, -) -> None: - assert apply_dify_debug_email_recipient(method, enabled=enabled, actor_id="actor-id") is method - - -def test_dify_file_reference_factory_passes_tenant_id(monkeypatch: pytest.MonkeyPatch) -> None: - build_from_mapping = MagicMock(return_value=sentinel.file) - monkeypatch.setattr(node_runtime.file_factory, "build_from_mapping", build_from_mapping) - - factory = DifyFileReferenceFactory(_build_run_context()) - - result = factory.build_from_mapping(mapping={"id": "upload-file"}) - - assert result is sentinel.file - build_from_mapping.assert_called_once_with( - mapping={"id": "upload-file"}, - tenant_id="tenant-id", - access_controller=node_runtime._file_access_controller, + @staticmethod + @pytest.mark.parametrize( + ("enabled", "method"), + [ + (False, _build_email_method(debug_mode=True)), + (True, _build_email_method(debug_mode=False)), + (True, WebAppDeliveryMethod(enabled=True, config=_WebAppDeliveryConfig())), + ], ) + def test_apply_dify_debug_email_recipient_noops_when_override_is_not_needed( + enabled: bool, + method: object, + ) -> None: + assert apply_dify_debug_email_recipient(method, enabled=enabled, actor_id="actor-id") is method + @staticmethod + def test_dify_file_reference_factory_passes_tenant_id(monkeypatch: pytest.MonkeyPatch) -> None: + build_from_mapping = MagicMock(return_value=sentinel.file) + monkeypatch.setattr(node_runtime.file_factory, "build_from_mapping", build_from_mapping) -def test_dify_prepared_llm_wraps_model_instance_calls() -> None: - model_schema = _build_model_schema() - model_instance = _ModelInstanceStub(model_schema=model_schema) - model_type_instance = model_instance.model_type_instance - prepared = DifyPreparedLLM(model_instance, request_metadata={"app_id": "app-id"}) + factory = DifyFileReferenceFactory(_build_run_context()) - assert prepared.provider == "langgenius/openai/openai" - assert prepared.model_name == "gpt-4o-mini" - assert prepared.parameters == {"temperature": 0.2} - assert prepared.stop == ("stop",) - assert prepared.get_model_schema() == model_schema - assert prepared.get_llm_num_tokens([]) == 32 - assert ( - prepared.invoke_llm( + result = factory.build_from_mapping(mapping={"id": "upload-file"}) + + assert result is sentinel.file + build_from_mapping.assert_called_once_with( + mapping={"id": "upload-file"}, + tenant_id="tenant-id", + access_controller=node_runtime._file_access_controller, + ) + + @staticmethod + def test_dify_prepared_llm_wraps_model_instance_calls() -> None: + model_schema = _build_model_schema() + model_instance = _ModelInstanceStub(model_schema=model_schema) + model_type_instance = model_instance.model_type_instance + prepared = DifyPreparedLLM(model_instance, request_metadata={"app_id": "app-id"}) + + assert prepared.provider == "langgenius/openai/openai" + assert prepared.model_name == "gpt-4o-mini" + assert prepared.parameters == {"temperature": 0.2} + assert prepared.stop == ("stop",) + assert prepared.get_model_schema() == model_schema + assert prepared.get_llm_num_tokens([]) == 32 + assert ( + prepared.invoke_llm( + prompt_messages=[], + model_parameters={"temperature": 0.1}, + tools=None, + stop=None, + stream=False, + ) + is sentinel.result + ) + + model_type_instance.get_model_schema.assert_called_once_with("gpt-4o-mini", {"api_key": "secret"}) + model_instance.invoke_llm.assert_called_once_with( prompt_messages=[], model_parameters={"temperature": 0.1}, - tools=None, - stop=None, + tools=[], + stop=[], stream=False, + request_metadata={"app_id": "app-id"}, ) - is sentinel.result - ) - model_type_instance.get_model_schema.assert_called_once_with("gpt-4o-mini", {"api_key": "secret"}) - model_instance.invoke_llm.assert_called_once_with( - prompt_messages=[], - model_parameters={"temperature": 0.1}, - tools=[], - stop=[], - stream=False, - request_metadata={"app_id": "app-id"}, - ) + @staticmethod + def test_dify_prepared_llm_requires_model_schema() -> None: + model_instance = _ModelInstanceStub(model_schema=None) + model_instance.credentials = {} + prepared = DifyPreparedLLM(model_instance) + with pytest.raises(ValueError, match="Model schema not found"): + prepared.get_model_schema() -def test_dify_prepared_llm_requires_model_schema() -> None: - model_instance = _ModelInstanceStub(model_schema=None) - model_instance.credentials = {} - prepared = DifyPreparedLLM(model_instance) + @staticmethod + def test_dify_prepared_llm_delegates_structured_output_helper(monkeypatch: pytest.MonkeyPatch) -> None: + model_instance = _ModelInstanceStub(model_schema=_build_model_schema()) + prepared = DifyPreparedLLM(model_instance, request_metadata={"created_by": "workflow"}) + invoke_structured = MagicMock(return_value=sentinel.structured) + monkeypatch.setattr(node_runtime, "invoke_llm_with_structured_output", invoke_structured) - with pytest.raises(ValueError, match="Model schema not found"): - prepared.get_model_schema() + result = prepared.invoke_llm_with_structured_output( + prompt_messages=[], + json_schema={"type": "object"}, + model_parameters={"temperature": 0.2}, + stop=("done",), + stream=True, + ) + assert result is sentinel.structured + invoke_structured.assert_called_once_with( + provider="langgenius/openai/openai", + model_schema=prepared.get_model_schema(), + model_instance=model_instance, + prompt_messages=[], + json_schema={"type": "object"}, + model_parameters={"temperature": 0.2}, + stop=["done"], + stream=True, + request_metadata={"created_by": "workflow"}, + ) -def test_dify_prepared_llm_delegates_structured_output_helper(monkeypatch: pytest.MonkeyPatch) -> None: - model_instance = _ModelInstanceStub(model_schema=_build_model_schema()) - prepared = DifyPreparedLLM(model_instance, request_metadata={"created_by": "workflow"}) - invoke_structured = MagicMock(return_value=sentinel.structured) - monkeypatch.setattr(node_runtime, "invoke_llm_with_structured_output", invoke_structured) + @staticmethod + def test_dify_prepared_llm_identifies_structured_output_errors() -> None: + prepared = DifyPreparedLLM(SimpleNamespace()) - result = prepared.invoke_llm_with_structured_output( - prompt_messages=[], - json_schema={"type": "object"}, - model_parameters={"temperature": 0.2}, - stop=("done",), - stream=True, - ) + assert prepared.is_structured_output_parse_error(OutputParserError("bad json")) is True + assert prepared.is_structured_output_parse_error(ValueError("other")) is False - assert result is sentinel.structured - invoke_structured.assert_called_once_with( - provider="langgenius/openai/openai", - model_schema=prepared.get_model_schema(), - model_instance=model_instance, - prompt_messages=[], - json_schema={"type": "object"}, - model_parameters={"temperature": 0.2}, - stop=["done"], - stream=True, - request_metadata={"created_by": "workflow"}, - ) + @staticmethod + def test_dify_prepared_polling_llm_delegates_to_plugin_runtime() -> None: + polling_result = LLMPollingResult( + status=LLMPollingStatus.RUNNING, + plugin_state={"task_id": "poll-1"}, + next_check_after_seconds=2, + ) + plugin_runtime = PluginModelRuntime( + tenant_id="tenant-id", + user_id="user-id", + client=PluginModelClient(), + plugin_service=PluginService, + ) + plugin_runtime.start_llm_polling = Mock(return_value=polling_result) # type: ignore[method-assign] + plugin_runtime.check_llm_polling = Mock(return_value=polling_result) # type: ignore[method-assign] + model_instance = _ModelInstanceStub( + model_schema=_build_model_schema(features=[ModelFeature.POLLING]), + model_runtime=plugin_runtime, + ) + prepared = DifyPreparedPollingLLM(model_instance) -def test_dify_prepared_llm_identifies_structured_output_errors() -> None: - prepared = DifyPreparedLLM(SimpleNamespace()) - - assert prepared.is_structured_output_parse_error(OutputParserError("bad json")) is True - assert prepared.is_structured_output_parse_error(ValueError("other")) is False - - -def test_dify_prepared_polling_llm_delegates_to_plugin_runtime() -> None: - polling_result = LLMPollingResult( - status=LLMPollingStatus.RUNNING, - plugin_state={"task_id": "poll-1"}, - next_check_after_seconds=2, - ) - plugin_runtime = PluginModelRuntime( - tenant_id="tenant-id", - user_id="user-id", - client=PluginModelClient(), - plugin_service=PluginService, - ) - plugin_runtime.start_llm_polling = Mock(return_value=polling_result) # type: ignore[method-assign] - plugin_runtime.check_llm_polling = Mock(return_value=polling_result) # type: ignore[method-assign] - model_instance = _ModelInstanceStub( - model_schema=_build_model_schema(features=[ModelFeature.POLLING]), - model_runtime=plugin_runtime, - ) - - prepared = DifyPreparedPollingLLM(model_instance) - - assert isinstance(prepared, LLMPollingCapableProtocol) - assert ( - prepared.start_llm_polling( + assert isinstance(prepared, LLMPollingCapableProtocol) + assert ( + prepared.start_llm_polling( + prompt_messages=[], + model_parameters={"temperature": 0.1}, + tools=[], + stop=("END",), + json_schema={"type": "object"}, + ) + == polling_result + ) + assert ( + prepared.check_llm_polling( + plugin_state={"task_id": "poll-1"}, + ) + == polling_result + ) + plugin_runtime.start_llm_polling.assert_called_once_with( + provider="langgenius/openai/openai", + model="gpt-4o-mini", + credentials={"api_key": "secret"}, prompt_messages=[], model_parameters={"temperature": 0.1}, tools=[], stop=("END",), json_schema={"type": "object"}, ) - == polling_result - ) - assert ( - prepared.check_llm_polling( + plugin_runtime.check_llm_polling.assert_called_once_with( + provider="langgenius/openai/openai", + model="gpt-4o-mini", + credentials={"api_key": "secret"}, plugin_state={"task_id": "poll-1"}, ) - == polling_result - ) - plugin_runtime.start_llm_polling.assert_called_once_with( - provider="langgenius/openai/openai", - model="gpt-4o-mini", - credentials={"api_key": "secret"}, - prompt_messages=[], - model_parameters={"temperature": 0.1}, - tools=[], - stop=("END",), - json_schema={"type": "object"}, - ) - plugin_runtime.check_llm_polling.assert_called_once_with( - provider="langgenius/openai/openai", - model="gpt-4o-mini", - credentials={"api_key": "secret"}, - plugin_state={"task_id": "poll-1"}, - ) + @staticmethod + def test_dify_prepared_polling_llm_commits_successful_reservation() -> None: + running_result = LLMPollingResult( + status=LLMPollingStatus.RUNNING, + plugin_state={"task_id": "poll-1"}, + ) + usage = node_runtime.LLMUsage.empty_usage().model_copy(update={"total_tokens": 5}) + succeeded_result = LLMPollingResult( + status=LLMPollingStatus.SUCCEEDED, + result=node_runtime.LLMResult( + model="gpt-4o-mini", + prompt_messages=[], + message=_node_runtime_AssistantPromptMessage(content="done"), + usage=usage, + ), + ) + plugin_runtime = PluginModelRuntime( + tenant_id="tenant-id", + user_id="user-id", + client=PluginModelClient(), + plugin_service=PluginService, + ) + plugin_runtime.start_llm_polling = Mock(return_value=running_result) # type: ignore[method-assign] + plugin_runtime.check_llm_polling = Mock(return_value=succeeded_result) # type: ignore[method-assign] + reservation = MagicMock() + model_instance = _QuotaManagedModelInstanceStub( + model_schema=_build_model_schema(features=[ModelFeature.POLLING]), + model_runtime=plugin_runtime, + ) + model_instance.reserve_quota.return_value = reservation + prepared = DifyPreparedPollingLLM(model_instance) -def test_dify_prepared_polling_llm_commits_successful_reservation() -> None: - running_result = LLMPollingResult( - status=LLMPollingStatus.RUNNING, - plugin_state={"task_id": "poll-1"}, - ) - usage = node_runtime.LLMUsage.empty_usage().model_copy(update={"total_tokens": 5}) - succeeded_result = LLMPollingResult( - status=LLMPollingStatus.SUCCEEDED, - result=node_runtime.LLMResult( - model="gpt-4o-mini", - prompt_messages=[], - message=AssistantPromptMessage(content="done"), - usage=usage, - ), - ) - plugin_runtime = PluginModelRuntime( - tenant_id="tenant-id", - user_id="user-id", - client=PluginModelClient(), - plugin_service=PluginService, - ) - plugin_runtime.start_llm_polling = Mock(return_value=running_result) # type: ignore[method-assign] - plugin_runtime.check_llm_polling = Mock(return_value=succeeded_result) # type: ignore[method-assign] - reservation = MagicMock() - model_instance = _QuotaManagedModelInstanceStub( - model_schema=_build_model_schema(features=[ModelFeature.POLLING]), - model_runtime=plugin_runtime, - ) - model_instance.reserve_quota.return_value = reservation - prepared = DifyPreparedPollingLLM(model_instance) - - prepared.start_llm_polling( - prompt_messages=[], - model_parameters={}, - tools=None, - stop=None, - json_schema=None, - ) - prepared.check_llm_polling(plugin_state={"task_id": "poll-1"}) - - reservation.commit.assert_called_once_with(usage) - reservation.release.assert_not_called() - - -def test_dify_prepared_polling_llm_releases_reservation_when_success_has_no_result( - monkeypatch: pytest.MonkeyPatch, -) -> None: - succeeded_result = LLMPollingResult.model_construct(status=LLMPollingStatus.SUCCEEDED, result=None) - polling_runtime = PluginModelRuntime( - tenant_id="tenant-id", - user_id="user-id", - client=PluginModelClient(), - plugin_service=PluginService, - ) - started_pollings: list[dict[str, object]] = [] - - def start_llm_polling(**kwargs: object) -> LLMPollingResult: - started_pollings.append(kwargs) - return succeeded_result - - monkeypatch.setattr(polling_runtime, "start_llm_polling", start_llm_polling) - reservation = sentinel.quota_reservation - model_instance = _QuotaManagedModelInstanceStub( - model_schema=_build_model_schema(features=[ModelFeature.POLLING]), - model_runtime=polling_runtime, - ) - model_instance.reserve_quota.return_value = reservation - prepared = DifyPreparedPollingLLM(model_instance) - - with ( - patch.object(node_runtime.QuotaManagedModelInstance, "release_quota_safely") as release_quota_safely, - pytest.raises(ValueError, match="must include a model result"), - ): prepared.start_llm_polling( prompt_messages=[], model_parameters={}, @@ -449,586 +1638,647 @@ def test_dify_prepared_polling_llm_releases_reservation_when_success_has_no_resu stop=None, json_schema=None, ) - - release_quota_safely.assert_called_once_with(reservation) - assert len(started_pollings) == 1 - - -def test_dify_prepared_polling_llm_releases_previous_reservation_on_restart() -> None: - running_result = LLMPollingResult( - status=LLMPollingStatus.RUNNING, - plugin_state={"task_id": "poll-1"}, - ) - plugin_runtime = PluginModelRuntime( - tenant_id="tenant-id", - user_id="user-id", - client=PluginModelClient(), - plugin_service=PluginService, - ) - plugin_runtime.start_llm_polling = Mock(return_value=running_result) # type: ignore[method-assign] - first_reservation = MagicMock() - second_reservation = MagicMock() - model_instance = _QuotaManagedModelInstanceStub( - model_schema=_build_model_schema(features=[ModelFeature.POLLING]), - model_runtime=plugin_runtime, - ) - model_instance.reserve_quota.side_effect = [first_reservation, second_reservation] - prepared = DifyPreparedPollingLLM(model_instance) - - for _ in range(2): - prepared.start_llm_polling( - prompt_messages=[], - model_parameters={}, - tools=None, - stop=None, - json_schema=None, - ) - - first_reservation.release.assert_called_once_with() - second_reservation.release.assert_not_called() - assert model_instance.reserve_quota.call_count == 2 - - -def test_dify_prepared_polling_llm_releases_reservation_when_finalized() -> None: - running_result = LLMPollingResult( - status=LLMPollingStatus.RUNNING, - plugin_state={"task_id": "poll-1"}, - ) - polling_runtime = SimpleNamespace( - start_llm_polling=Mock(return_value=running_result), - check_llm_polling=Mock(), - ) - reservation = MagicMock() - model_instance = _QuotaManagedModelInstanceStub( - model_schema=_build_model_schema(features=[ModelFeature.POLLING]), - model_runtime=polling_runtime, - ) - model_instance.reserve_quota.return_value = reservation - prepared = DifyPreparedPollingLLM(model_instance) - - prepared.start_llm_polling( - prompt_messages=[], - model_parameters={}, - tools=None, - stop=None, - json_schema=None, - ) - prepared.finalize_llm_polling() - prepared.finalize_llm_polling() - - reservation.release.assert_called_once_with() - - -def test_dify_prepared_polling_llm_releases_reservation_when_check_fails() -> None: - running_result = LLMPollingResult( - status=LLMPollingStatus.RUNNING, - plugin_state={"task_id": "poll-1"}, - ) - polling_runtime = SimpleNamespace( - start_llm_polling=Mock(return_value=running_result), - check_llm_polling=Mock(side_effect=RuntimeError("polling failed")), - ) - reservation = MagicMock() - model_instance = _QuotaManagedModelInstanceStub( - model_schema=_build_model_schema(features=[ModelFeature.POLLING]), - model_runtime=polling_runtime, - ) - model_instance.reserve_quota.return_value = reservation - prepared = DifyPreparedPollingLLM(model_instance) - - prepared.start_llm_polling( - prompt_messages=[], - model_parameters={}, - tools=None, - stop=None, - json_schema=None, - ) - - with pytest.raises(RuntimeError, match="polling failed"): prepared.check_llm_polling(plugin_state={"task_id": "poll-1"}) - reservation.release.assert_called_once_with() + reservation.commit.assert_called_once_with(usage) + reservation.release.assert_not_called() + @staticmethod + def test_dify_prepared_polling_llm_releases_reservation_when_success_has_no_result( + monkeypatch: pytest.MonkeyPatch, + ) -> None: + succeeded_result = LLMPollingResult.model_construct(status=LLMPollingStatus.SUCCEEDED, result=None) + polling_runtime = PluginModelRuntime( + tenant_id="tenant-id", + user_id="user-id", + client=PluginModelClient(), + plugin_service=PluginService, + ) + started_pollings: list[dict[str, object]] = [] -def test_dify_prepared_polling_llm_raise_exception_when_polling_is_unsupported() -> None: - llm_result = node_runtime.LLMResult( - model="gpt-4o-mini", - prompt_messages=[], - message=AssistantPromptMessage(content="sync-result"), - usage=node_runtime.LLMUsage.empty_usage(), - ) - plugin_runtime = PluginModelRuntime( - tenant_id="tenant-id", - user_id="user-id", - client=Mock(), - plugin_service=Mock(), - ) - plugin_runtime.start_llm_polling = Mock(side_effect=PluginLLMPollingUnsupportedError("Polling unsupported")) # type: ignore[method-assign] - model_instance = _ModelInstanceStub( - model_schema=_build_model_schema(features=[ModelFeature.POLLING]), - model_runtime=plugin_runtime, - invoke_llm_result=llm_result, - ) + def start_llm_polling(**kwargs: object) -> LLMPollingResult: + started_pollings.append(kwargs) + return succeeded_result - prepared = DifyPreparedPollingLLM(model_instance) + monkeypatch.setattr(polling_runtime, "start_llm_polling", start_llm_polling) + reservation = sentinel.quota_reservation + model_instance = _QuotaManagedModelInstanceStub( + model_schema=_build_model_schema(features=[ModelFeature.POLLING]), + model_runtime=polling_runtime, + ) + model_instance.reserve_quota.return_value = reservation + prepared = DifyPreparedPollingLLM(model_instance) + + with ( + patch.object(node_runtime.QuotaManagedModelInstance, "release_quota_safely") as release_quota_safely, + pytest.raises(ValueError, match="must include a model result"), + ): + prepared.start_llm_polling( + prompt_messages=[], + model_parameters={}, + tools=None, + stop=None, + json_schema=None, + ) + + release_quota_safely.assert_called_once_with(reservation) + assert len(started_pollings) == 1 + + @staticmethod + def test_dify_prepared_polling_llm_releases_previous_reservation_on_restart() -> None: + running_result = LLMPollingResult( + status=LLMPollingStatus.RUNNING, + plugin_state={"task_id": "poll-1"}, + ) + plugin_runtime = PluginModelRuntime( + tenant_id="tenant-id", + user_id="user-id", + client=PluginModelClient(), + plugin_service=PluginService, + ) + plugin_runtime.start_llm_polling = Mock(return_value=running_result) # type: ignore[method-assign] + first_reservation = MagicMock() + second_reservation = MagicMock() + model_instance = _QuotaManagedModelInstanceStub( + model_schema=_build_model_schema(features=[ModelFeature.POLLING]), + model_runtime=plugin_runtime, + ) + model_instance.reserve_quota.side_effect = [first_reservation, second_reservation] + prepared = DifyPreparedPollingLLM(model_instance) + + for _ in range(2): + prepared.start_llm_polling( + prompt_messages=[], + model_parameters={}, + tools=None, + stop=None, + json_schema=None, + ) + + first_reservation.release.assert_called_once_with() + second_reservation.release.assert_not_called() + assert model_instance.reserve_quota.call_count == 2 + + @staticmethod + def test_dify_prepared_polling_llm_releases_reservation_when_finalized() -> None: + running_result = LLMPollingResult( + status=LLMPollingStatus.RUNNING, + plugin_state={"task_id": "poll-1"}, + ) + polling_runtime = SimpleNamespace( + start_llm_polling=Mock(return_value=running_result), + check_llm_polling=Mock(), + ) + reservation = MagicMock() + model_instance = _QuotaManagedModelInstanceStub( + model_schema=_build_model_schema(features=[ModelFeature.POLLING]), + model_runtime=polling_runtime, + ) + model_instance.reserve_quota.return_value = reservation + prepared = DifyPreparedPollingLLM(model_instance) - with pytest.raises(PluginLLMPollingUnsupportedError): prepared.start_llm_polling( prompt_messages=[], - model_parameters={"temperature": 0.1}, + model_parameters={}, + tools=None, + stop=None, + json_schema=None, + ) + prepared.finalize_llm_polling() + prepared.finalize_llm_polling() + + reservation.release.assert_called_once_with() + + @staticmethod + def test_dify_prepared_polling_llm_releases_reservation_when_check_fails() -> None: + running_result = LLMPollingResult( + status=LLMPollingStatus.RUNNING, + plugin_state={"task_id": "poll-1"}, + ) + polling_runtime = SimpleNamespace( + start_llm_polling=Mock(return_value=running_result), + check_llm_polling=Mock(side_effect=RuntimeError("polling failed")), + ) + reservation = MagicMock() + model_instance = _QuotaManagedModelInstanceStub( + model_schema=_build_model_schema(features=[ModelFeature.POLLING]), + model_runtime=polling_runtime, + ) + model_instance.reserve_quota.return_value = reservation + prepared = DifyPreparedPollingLLM(model_instance) + + prepared.start_llm_polling( + prompt_messages=[], + model_parameters={}, tools=None, stop=None, json_schema=None, ) + with pytest.raises(RuntimeError, match="polling failed"): + prepared.check_llm_polling(plugin_state={"task_id": "poll-1"}) -def test_dify_prompt_message_serializer_delegates(monkeypatch: pytest.MonkeyPatch) -> None: - serialize = MagicMock(return_value={"prompt": "value"}) - monkeypatch.setattr(node_runtime.PromptMessageUtil, "prompt_messages_to_prompt_for_saving", serialize) + reservation.release.assert_called_once_with() - result = DifyPromptMessageSerializer().serialize( - model_mode="chat", - prompt_messages=[], - ) - - assert result == {"prompt": "value"} - serialize.assert_called_once_with( - model_mode="chat", - prompt_messages=[], - ) - - -def test_dify_retriever_attachment_loader_builds_graph_files( - monkeypatch: pytest.MonkeyPatch, attachment_session: Session -) -> None: - _persist_attachment(attachment_session, segment_id="segment-id", upload_file_id="upload-file-id") - build_from_mapping = MagicMock(return_value=sentinel.file) - monkeypatch.setattr(node_runtime, "db", SimpleNamespace(engine=attachment_session.get_bind())) - loader = DifyRetrieverAttachmentLoader( - file_reference_factory=SimpleNamespace(build_from_mapping=build_from_mapping) - ) - - files = loader.load(segment_id="segment-id") - - assert files == [sentinel.file] - build_from_mapping.assert_called_once() - mapping = build_from_mapping.call_args.kwargs["mapping"] - assert mapping["id"] == "upload-file-id" - assert mapping["transfer_method"] == FileTransferMethod.LOCAL_FILE - assert mapping["type"] == FileType.IMAGE - assert parse_file_reference(mapping["reference"]).storage_key is None - - -def test_dify_retriever_attachment_loader_grants_upload_files_for_allowed_segment( - monkeypatch: pytest.MonkeyPatch, - attachment_session: Session, -) -> None: - from factories.file_factory import builders as file_builders - - upload_file_id = str(uuid4()) - segment_id = str(uuid4()) - _persist_attachment(attachment_session, segment_id=segment_id, upload_file_id=upload_file_id) - engine = attachment_session.get_bind() - assert engine is not None - monkeypatch.setattr(node_runtime, "db", SimpleNamespace(engine=engine)) - session_maker = sessionmaker(engine, expire_on_commit=False) - monkeypatch.setattr(file_builders.session_factory, "create_session", session_maker) - - loader = DifyRetrieverAttachmentLoader(file_reference_factory=DifyFileReferenceFactory(_build_run_context())) - scope = FileAccessScope( - tenant_id="tenant-id", - user_id="end-user-id", - user_from=UserFrom.END_USER, - invoke_from=InvokeFrom.WEB_APP, - ) - - with bind_file_access_scope(scope): - grant_retriever_segment_access([segment_id]) - files = loader.load(segment_id=segment_id) - - assert files[0].related_id == upload_file_id - assert files[0].filename == "diagram.png" - - -def test_dify_retriever_attachment_loader_rejects_granted_upload_file_from_another_tenant( - monkeypatch: pytest.MonkeyPatch, - attachment_session: Session, -) -> None: - from factories.file_factory import builders as file_builders - - upload_file_id = str(uuid4()) - segment_id = str(uuid4()) - _persist_attachment( - attachment_session, - segment_id=segment_id, - upload_file_id=upload_file_id, - upload_file_tenant_id="other-tenant-id", - ) - engine = attachment_session.get_bind() - assert engine is not None - monkeypatch.setattr(node_runtime, "db", SimpleNamespace(engine=engine)) - monkeypatch.setattr(file_builders.session_factory, "create_session", sessionmaker(engine, expire_on_commit=False)) - - loader = DifyRetrieverAttachmentLoader(file_reference_factory=DifyFileReferenceFactory(_build_run_context())) - scope = FileAccessScope( - tenant_id="tenant-id", - user_id="end-user-id", - user_from=UserFrom.END_USER, - invoke_from=InvokeFrom.WEB_APP, - ) - - with bind_file_access_scope(scope): - grant_retriever_segment_access([segment_id]) - with pytest.raises(ValueError, match="Invalid upload file"): - loader.load(segment_id=segment_id) - - -def test_dify_retriever_attachment_loader_skips_ungranted_segment_for_end_user( - monkeypatch: pytest.MonkeyPatch, - attachment_session: Session, -) -> None: - build_from_mapping = MagicMock() - engine = attachment_session.get_bind() - assert engine is not None - monkeypatch.setattr(node_runtime, "db", SimpleNamespace(engine=engine)) - statement_count = 0 - - def count_statements(*_args, **_kwargs) -> None: - nonlocal statement_count - statement_count += 1 - - event.listen(engine, "before_cursor_execute", count_statements) - loader = DifyRetrieverAttachmentLoader( - file_reference_factory=SimpleNamespace(build_from_mapping=build_from_mapping) - ) - scope = FileAccessScope( - tenant_id="tenant-id", - user_id="end-user-id", - user_from=UserFrom.END_USER, - invoke_from=InvokeFrom.WEB_APP, - ) - - with bind_file_access_scope(scope): - files = loader.load(segment_id=str(uuid4())) - - try: - assert files == [] - assert statement_count == 0 - build_from_mapping.assert_not_called() - finally: - event.remove(engine, "before_cursor_execute", count_statements) - - -def test_dify_retriever_attachment_loader_skips_segment_rejected_by_checker( - monkeypatch: pytest.MonkeyPatch, - attachment_session: Session, -) -> None: - segment_id = str(uuid4()) - build_from_mapping = MagicMock() - segment_access_checker = MagicMock(return_value=False) - engine = attachment_session.get_bind() - assert engine is not None - monkeypatch.setattr(node_runtime, "db", SimpleNamespace(engine=engine)) - statement_count = 0 - - def count_statements(*_args, **_kwargs) -> None: - nonlocal statement_count - statement_count += 1 - - event.listen(engine, "before_cursor_execute", count_statements) - loader = DifyRetrieverAttachmentLoader( - file_reference_factory=SimpleNamespace(build_from_mapping=build_from_mapping), - segment_access_checker=segment_access_checker, - ) - scope = FileAccessScope( - tenant_id="tenant-id", - user_id="end-user-id", - user_from=UserFrom.END_USER, - invoke_from=InvokeFrom.WEB_APP, - ) - - with bind_file_access_scope(scope): - grant_retriever_segment_access([segment_id]) - files = loader.load(segment_id=segment_id) - - try: - assert files == [] - segment_access_checker.assert_called_once_with(segment_id) - assert statement_count == 0 - build_from_mapping.assert_not_called() - finally: - event.remove(engine, "before_cursor_execute", count_statements) - - -def test_dify_tool_file_manager_resolves_conversation_id_for_tool_files(monkeypatch: pytest.MonkeyPatch) -> None: - tool_file = ToolFile( - user_id="user-id", - tenant_id="tenant-id", - conversation_id="conversation-id", - file_key="tools/tenant-id/tool-file-id.png", - mimetype="image/png", - name="diagram.png", - size=len(b"file-bytes"), - ) - tool_file.id = "tool-file-id" - create_file_by_raw = MagicMock(return_value=tool_file) - manager_instance = SimpleNamespace(create_file_by_raw=create_file_by_raw) - monkeypatch.setattr(node_runtime, "ToolFileManager", MagicMock(return_value=manager_instance)) - conversation_id_getter = MagicMock(return_value="conversation-id") - - manager = DifyToolFileManager( - DifyRunContext( + @staticmethod + def test_dify_prepared_polling_llm_raise_exception_when_polling_is_unsupported() -> None: + llm_result = node_runtime.LLMResult( + model="gpt-4o-mini", + prompt_messages=[], + message=_node_runtime_AssistantPromptMessage(content="sync-result"), + usage=node_runtime.LLMUsage.empty_usage(), + ) + plugin_runtime = PluginModelRuntime( tenant_id="tenant-id", - app_id="app-id", user_id="user-id", - user_from=UserFrom.ACCOUNT, - invoke_from=InvokeFrom.DEBUGGER, - ), - conversation_id_getter=conversation_id_getter, - ) - - tool_file = manager.create_file_by_raw( - file_binary=b"file-bytes", - mimetype="image/png", - filename="diagram", - ) - - assert tool_file.id == "tool-file-id" - conversation_id_getter.assert_called_once_with() - create_file_by_raw.assert_called_once_with( - user_id="user-id", - tenant_id="tenant-id", - conversation_id="conversation-id", - file_binary=b"file-bytes", - mimetype="image/png", - filename="diagram", - ) - - -def test_dify_tool_file_manager_delegates_file_generator_lookup(monkeypatch: pytest.MonkeyPatch) -> None: - get_file_generator = MagicMock(return_value=sentinel.generator) - monkeypatch.setattr( - node_runtime, - "ToolFileManager", - MagicMock(return_value=SimpleNamespace(get_file_generator_by_tool_file_id=get_file_generator)), - ) - manager = DifyToolFileManager(_build_run_context()) - - assert manager.get_file_generator_by_tool_file_id("tool-file-id") is sentinel.generator - get_file_generator.assert_called_once_with("tool-file-id") - - -def test_dify_tool_node_runtime_injects_outer_workflow_run_id_for_workflow_tools( - monkeypatch: pytest.MonkeyPatch, -) -> None: - runtime_tool = SimpleNamespace(runtime=SimpleNamespace(runtime_parameters={})) - get_runtime = MagicMock(return_value=runtime_tool) - monkeypatch.setattr(node_runtime.ToolManager, "get_workflow_tool_runtime", get_runtime) - monkeypatch.setattr( - node_runtime, - "get_system_text", - lambda _pool, key: ( - "outer-workflow-run-id" if key == node_runtime.SystemVariableKey.WORKFLOW_EXECUTION_ID else None - ), - ) - - runtime = node_runtime.DifyToolNodeRuntime(_build_run_context()) - node_data = ToolNodeData( - title="Workflow Tool Node", - desc=None, - provider_id="workflow-provider-id", - provider_type=ToolProviderType.WORKFLOW, - provider_name="workflow-provider", - tool_name="workflow-tool", - tool_label="Workflow Tool", - tool_configurations={}, - tool_parameters={}, - ) - - handle = runtime.get_runtime( - node_id="tool-node", - node_data=node_data, - variable_pool=object(), - node_execution_id="node-execution-id", - ) - - assert handle.raw.tool is runtime_tool - assert handle.raw.parent_trace_context.model_dump() == { - "parent_workflow_run_id": "outer-workflow-run-id", - "parent_node_execution_id": "node-execution-id", - } - assert runtime_tool.runtime.runtime_parameters == {} - get_runtime.assert_called_once() - - -def test_dify_tool_node_runtime_stores_trace_session_id_for_workflow_tools( - monkeypatch: pytest.MonkeyPatch, -) -> None: - runtime_tool = SimpleNamespace(runtime=SimpleNamespace(runtime_parameters={})) - get_runtime = MagicMock(return_value=runtime_tool) - monkeypatch.setattr(node_runtime.ToolManager, "get_workflow_tool_runtime", get_runtime) - monkeypatch.setattr( - node_runtime, - "get_system_text", - lambda _pool, key: ( - "outer-workflow-run-id" if key == node_runtime.SystemVariableKey.WORKFLOW_EXECUTION_ID else None - ), - ) - - run_context = _build_run_context() - run_context[DIFY_RUN_CONTEXT_KEY].trace_session_id = "session-1" - runtime = node_runtime.DifyToolNodeRuntime(run_context) - node_data = ToolNodeData( - title="Workflow Tool Node", - desc=None, - provider_id="workflow-provider-id", - provider_type=ToolProviderType.WORKFLOW, - provider_name="workflow-provider", - tool_name="workflow-tool", - tool_label="Workflow Tool", - tool_configurations={}, - tool_parameters={}, - ) - - handle = runtime.get_runtime( - node_id="tool-node", - node_data=node_data, - variable_pool=object(), - node_execution_id="node-execution-id", - ) - - assert handle.raw.trace_session_id == "session-1" - assert runtime_tool.runtime.runtime_parameters == {} - - -def test_dify_tool_node_runtime_does_not_inject_outer_workflow_run_id_for_non_workflow_tools( - monkeypatch: pytest.MonkeyPatch, -) -> None: - runtime_tool = SimpleNamespace(runtime=SimpleNamespace(runtime_parameters={})) - get_runtime = MagicMock(return_value=runtime_tool) - monkeypatch.setattr(node_runtime.ToolManager, "get_workflow_tool_runtime", get_runtime) - monkeypatch.setattr(node_runtime, "get_system_text", lambda _pool, _key: None) - - runtime = node_runtime.DifyToolNodeRuntime(_build_run_context()) - node_data = ToolNodeData( - title="Builtin Tool Node", - desc=None, - provider_id="builtin-provider-id", - provider_type=ToolProviderType.BUILT_IN, - provider_name="builtin-provider", - tool_name="builtin-tool", - tool_label="Builtin Tool", - tool_configurations={}, - tool_parameters={}, - ) - - handle = runtime.get_runtime( - node_id="tool-node", - node_data=node_data, - variable_pool=object(), - ) - - assert handle.raw.tool is runtime_tool - assert runtime_tool.runtime.dify_run_context.tenant_id == "tenant-id" - assert runtime_tool.runtime.dify_run_context.app_id == "app-id" - assert "outer_workflow_run_id" not in runtime_tool.runtime.runtime_parameters - get_runtime.assert_called_once() - - -def test_dify_tool_node_runtime_builds_workflow_tool_container_payload(monkeypatch: pytest.MonkeyPatch) -> None: - from core.tools.__base.tool_runtime import ToolRuntime - from core.tools.entities.common_entities import I18nObject - from core.tools.entities.tool_entities import ToolEntity, ToolIdentity - from core.tools.workflow_as_tool.tool import WorkflowTool - - tool = WorkflowTool( - workflow_app_id=str(uuid4()), - workflow_id=str(uuid4()), - workflow_as_tool_id="workflow-tool-id", - version="published-version", - workflow_call_depth=2, - entity=ToolEntity( - identity=ToolIdentity( - author="test", - name="workflow-tool", - label=I18nObject(en_US="Workflow Tool"), - provider="workflow-provider", - ), - parameters=[], - description=None, - has_runtime_parameters=False, - ), - runtime=ToolRuntime(tenant_id="tenant-id", invoke_from=InvokeFrom.DEBUGGER), - ) - prepared_parameters: list[Mapping[str, object]] = [] - - def prepare_container_inputs(parameters: Mapping[str, object]) -> tuple[dict[str, object], list[dict[str, str]]]: - prepared_parameters.append(parameters) - return {"amount": Decimal("1.25")}, [{"id": "file-id", "name": "input.txt"}] - - monkeypatch.setattr(tool, "prepare_container_inputs", prepare_container_inputs) - runtime = DifyToolNodeRuntime(_build_run_context()) - - payload = runtime.build_workflow_tool_container_payload( - tool_runtime=ToolRuntimeHandle(raw=tool), - tool_parameters={"amount": "1.25"}, - inputs_for_log={"amount": Decimal("1.25")}, - workflow_call_depth=2, - ) - - assert payload.model_dump() == { - "version": "1", - "source_app_id": str(tool.workflow_app_id), - "source_workflow_id": str(tool.workflow_id), - "source_workflow_version": "published-version", - "inputs": {"amount": 1.25}, - "system_files": [{"id": "file-id", "name": "input.txt"}], - "inputs_for_log": {"amount": 1.25}, - "call_depth": 3, - } - assert prepared_parameters == [{"amount": "1.25"}] - - -def test_dify_tool_node_runtime_rejects_non_workflow_tool_container_payload() -> None: - runtime = DifyToolNodeRuntime(_build_run_context()) - with pytest.raises(ToolRuntimeResolutionError, match="resolved tool is not a Workflow Tool"): - runtime.build_workflow_tool_container_payload( - tool_runtime=ToolRuntimeHandle(raw=object()), - tool_parameters={}, - inputs_for_log={}, - workflow_call_depth=0, + client=Mock(), + plugin_service=Mock(), + ) + plugin_runtime.start_llm_polling = Mock(side_effect=PluginLLMPollingUnsupportedError("Polling unsupported")) # type: ignore[method-assign] + model_instance = _ModelInstanceStub( + model_schema=_build_model_schema(features=[ModelFeature.POLLING]), + model_runtime=plugin_runtime, + invoke_llm_result=llm_result, ) + prepared = DifyPreparedPollingLLM(model_instance) -def test_dify_human_input_runtime_builds_debug_repository(monkeypatch: pytest.MonkeyPatch) -> None: - repository = MagicMock() - repository_cls = MagicMock(return_value=repository) - monkeypatch.setattr(node_runtime, "HumanInputFormRepositoryImpl", repository_cls) + with pytest.raises(PluginLLMPollingUnsupportedError): + prepared.start_llm_polling( + prompt_messages=[], + model_parameters={"temperature": 0.1}, + tools=None, + stop=None, + json_schema=None, + ) - runtime = DifyHumanInputNodeRuntime( - _build_run_context(), - workflow_execution_id_getter=lambda: "workflow-execution-id", - ) + @staticmethod + def test_dify_prompt_message_serializer_delegates(monkeypatch: pytest.MonkeyPatch) -> None: + serialize = MagicMock(return_value={"prompt": "value"}) + monkeypatch.setattr(node_runtime.PromptMessageUtil, "prompt_messages_to_prompt_for_saving", serialize) - assert runtime.build_form_repository() is repository - repository_cls.assert_called_once_with( - tenant_id="tenant-id", - app_id="app-id", - workflow_execution_id="workflow-execution-id", - invoke_source="debugger", - submission_actor_id="user-id", - ) + result = DifyPromptMessageSerializer().serialize( + model_mode="chat", + prompt_messages=[], + ) + assert result == {"prompt": "value"} + serialize.assert_called_once_with( + model_mode="chat", + prompt_messages=[], + ) -def test_dify_tool_runtime_spec_prefers_tool_parameters_for_runtime_form_values() -> None: - node_data = ToolNodeData( - provider_id="video-mixcut-agent", - provider_type=ToolProviderType.PLUGIN, - provider_name="sawyer-shi/video-mixcut-agent", - tool_name="mixcut", - tool_label="MixCut", - tool_configurations={"count": 2}, - tool_parameters={ + @staticmethod + def test_dify_retriever_attachment_loader_builds_graph_files( + monkeypatch: pytest.MonkeyPatch, attachment_session: Session + ) -> None: + _persist_attachment(attachment_session, segment_id="segment-id", upload_file_id="upload-file-id") + build_from_mapping = MagicMock(return_value=sentinel.file) + monkeypatch.setattr(node_runtime, "db", SimpleNamespace(engine=attachment_session.get_bind())) + loader = DifyRetrieverAttachmentLoader( + file_reference_factory=SimpleNamespace(build_from_mapping=build_from_mapping) + ) + + files = loader.load(segment_id="segment-id") + + assert files == [sentinel.file] + build_from_mapping.assert_called_once() + mapping = build_from_mapping.call_args.kwargs["mapping"] + assert mapping["id"] == "upload-file-id" + assert mapping["transfer_method"] == FileTransferMethod.LOCAL_FILE + assert mapping["type"] == FileType.IMAGE + assert parse_file_reference(mapping["reference"]).storage_key is None + + @staticmethod + def test_dify_retriever_attachment_loader_grants_upload_files_for_allowed_segment( + monkeypatch: pytest.MonkeyPatch, + attachment_session: Session, + ) -> None: + from factories.file_factory import builders as file_builders + + upload_file_id = str(uuid4()) + segment_id = str(uuid4()) + _persist_attachment(attachment_session, segment_id=segment_id, upload_file_id=upload_file_id) + engine = attachment_session.get_bind() + assert engine is not None + monkeypatch.setattr(node_runtime, "db", SimpleNamespace(engine=engine)) + session_maker = sessionmaker(engine, expire_on_commit=False) + monkeypatch.setattr(file_builders.session_factory, "create_session", session_maker) + + loader = DifyRetrieverAttachmentLoader(file_reference_factory=DifyFileReferenceFactory(_build_run_context())) + scope = FileAccessScope( + tenant_id="tenant-id", + user_id="end-user-id", + user_from=UserFrom.END_USER, + invoke_from=InvokeFrom.WEB_APP, + ) + + with bind_file_access_scope(scope): + grant_retriever_segment_access([segment_id]) + files = loader.load(segment_id=segment_id) + + assert files[0].related_id == upload_file_id + assert files[0].filename == "diagram.png" + + @staticmethod + def test_dify_retriever_attachment_loader_rejects_granted_upload_file_from_another_tenant( + monkeypatch: pytest.MonkeyPatch, + attachment_session: Session, + ) -> None: + from factories.file_factory import builders as file_builders + + upload_file_id = str(uuid4()) + segment_id = str(uuid4()) + _persist_attachment( + attachment_session, + segment_id=segment_id, + upload_file_id=upload_file_id, + upload_file_tenant_id="other-tenant-id", + ) + engine = attachment_session.get_bind() + assert engine is not None + monkeypatch.setattr(node_runtime, "db", SimpleNamespace(engine=engine)) + monkeypatch.setattr( + file_builders.session_factory, "create_session", sessionmaker(engine, expire_on_commit=False) + ) + + loader = DifyRetrieverAttachmentLoader(file_reference_factory=DifyFileReferenceFactory(_build_run_context())) + scope = FileAccessScope( + tenant_id="tenant-id", + user_id="end-user-id", + user_from=UserFrom.END_USER, + invoke_from=InvokeFrom.WEB_APP, + ) + + with bind_file_access_scope(scope): + grant_retriever_segment_access([segment_id]) + with pytest.raises(ValueError, match="Invalid upload file"): + loader.load(segment_id=segment_id) + + @staticmethod + def test_dify_retriever_attachment_loader_skips_ungranted_segment_for_end_user( + monkeypatch: pytest.MonkeyPatch, + attachment_session: Session, + ) -> None: + build_from_mapping = MagicMock() + engine = attachment_session.get_bind() + assert engine is not None + monkeypatch.setattr(node_runtime, "db", SimpleNamespace(engine=engine)) + statement_count = 0 + + def count_statements(*_args, **_kwargs) -> None: + nonlocal statement_count + statement_count += 1 + + event.listen(engine, "before_cursor_execute", count_statements) + loader = DifyRetrieverAttachmentLoader( + file_reference_factory=SimpleNamespace(build_from_mapping=build_from_mapping) + ) + scope = FileAccessScope( + tenant_id="tenant-id", + user_id="end-user-id", + user_from=UserFrom.END_USER, + invoke_from=InvokeFrom.WEB_APP, + ) + + with bind_file_access_scope(scope): + files = loader.load(segment_id=str(uuid4())) + + try: + assert files == [] + assert statement_count == 0 + build_from_mapping.assert_not_called() + finally: + event.remove(engine, "before_cursor_execute", count_statements) + + @staticmethod + def test_dify_retriever_attachment_loader_skips_segment_rejected_by_checker( + monkeypatch: pytest.MonkeyPatch, + attachment_session: Session, + ) -> None: + segment_id = str(uuid4()) + build_from_mapping = MagicMock() + segment_access_checker = MagicMock(return_value=False) + engine = attachment_session.get_bind() + assert engine is not None + monkeypatch.setattr(node_runtime, "db", SimpleNamespace(engine=engine)) + statement_count = 0 + + def count_statements(*_args, **_kwargs) -> None: + nonlocal statement_count + statement_count += 1 + + event.listen(engine, "before_cursor_execute", count_statements) + loader = DifyRetrieverAttachmentLoader( + file_reference_factory=SimpleNamespace(build_from_mapping=build_from_mapping), + segment_access_checker=segment_access_checker, + ) + scope = FileAccessScope( + tenant_id="tenant-id", + user_id="end-user-id", + user_from=UserFrom.END_USER, + invoke_from=InvokeFrom.WEB_APP, + ) + + with bind_file_access_scope(scope): + grant_retriever_segment_access([segment_id]) + files = loader.load(segment_id=segment_id) + + try: + assert files == [] + segment_access_checker.assert_called_once_with(segment_id) + assert statement_count == 0 + build_from_mapping.assert_not_called() + finally: + event.remove(engine, "before_cursor_execute", count_statements) + + @staticmethod + def test_dify_tool_file_manager_resolves_conversation_id_for_tool_files(monkeypatch: pytest.MonkeyPatch) -> None: + tool_file = ToolFile( + user_id="user-id", + tenant_id="tenant-id", + conversation_id="conversation-id", + file_key="tools/tenant-id/tool-file-id.png", + mimetype="image/png", + name="diagram.png", + size=len(b"file-bytes"), + ) + tool_file.id = "tool-file-id" + create_file_by_raw = MagicMock(return_value=tool_file) + manager_instance = SimpleNamespace(create_file_by_raw=create_file_by_raw) + monkeypatch.setattr(node_runtime, "ToolFileManager", MagicMock(return_value=manager_instance)) + conversation_id_getter = MagicMock(return_value="conversation-id") + + manager = DifyToolFileManager( + DifyRunContext( + tenant_id="tenant-id", + app_id="app-id", + user_id="user-id", + user_from=UserFrom.ACCOUNT, + invoke_from=InvokeFrom.DEBUGGER, + ), + conversation_id_getter=conversation_id_getter, + ) + + tool_file = manager.create_file_by_raw( + file_binary=b"file-bytes", + mimetype="image/png", + filename="diagram", + ) + + assert tool_file.id == "tool-file-id" + conversation_id_getter.assert_called_once_with() + create_file_by_raw.assert_called_once_with( + user_id="user-id", + tenant_id="tenant-id", + conversation_id="conversation-id", + file_binary=b"file-bytes", + mimetype="image/png", + filename="diagram", + ) + + @staticmethod + def test_dify_tool_file_manager_delegates_file_generator_lookup(monkeypatch: pytest.MonkeyPatch) -> None: + get_file_generator = MagicMock(return_value=sentinel.generator) + monkeypatch.setattr( + node_runtime, + "ToolFileManager", + MagicMock(return_value=SimpleNamespace(get_file_generator_by_tool_file_id=get_file_generator)), + ) + manager = DifyToolFileManager(_build_run_context()) + + assert manager.get_file_generator_by_tool_file_id("tool-file-id") is sentinel.generator + get_file_generator.assert_called_once_with("tool-file-id") + + @staticmethod + def test_dify_tool_node_runtime_injects_outer_workflow_run_id_for_workflow_tools( + monkeypatch: pytest.MonkeyPatch, + ) -> None: + runtime_tool = SimpleNamespace(runtime=SimpleNamespace(runtime_parameters={})) + get_runtime = MagicMock(return_value=runtime_tool) + monkeypatch.setattr(node_runtime.ToolManager, "get_workflow_tool_runtime", get_runtime) + monkeypatch.setattr( + node_runtime, + "get_system_text", + lambda _pool, key: ( + "outer-workflow-run-id" if key == node_runtime.SystemVariableKey.WORKFLOW_EXECUTION_ID else None + ), + ) + + runtime = node_runtime.DifyToolNodeRuntime(_build_run_context()) + node_data = ToolNodeData( + title="Workflow Tool Node", + desc=None, + provider_id="workflow-provider-id", + provider_type=ToolProviderType.WORKFLOW, + provider_name="workflow-provider", + tool_name="workflow-tool", + tool_label="Workflow Tool", + tool_configurations={}, + tool_parameters={}, + ) + + handle = runtime.get_runtime( + node_id="tool-node", + node_data=node_data, + variable_pool=object(), + node_execution_id="node-execution-id", + ) + + assert handle.raw.tool is runtime_tool + assert handle.raw.parent_trace_context.model_dump() == { + "parent_workflow_run_id": "outer-workflow-run-id", + "parent_node_execution_id": "node-execution-id", + } + assert runtime_tool.runtime.runtime_parameters == {} + get_runtime.assert_called_once() + + @staticmethod + def test_dify_tool_node_runtime_stores_trace_session_id_for_workflow_tools( + monkeypatch: pytest.MonkeyPatch, + ) -> None: + runtime_tool = SimpleNamespace(runtime=SimpleNamespace(runtime_parameters={})) + get_runtime = MagicMock(return_value=runtime_tool) + monkeypatch.setattr(node_runtime.ToolManager, "get_workflow_tool_runtime", get_runtime) + monkeypatch.setattr( + node_runtime, + "get_system_text", + lambda _pool, key: ( + "outer-workflow-run-id" if key == node_runtime.SystemVariableKey.WORKFLOW_EXECUTION_ID else None + ), + ) + + run_context = _build_run_context() + run_context[DIFY_RUN_CONTEXT_KEY].trace_session_id = "session-1" + runtime = node_runtime.DifyToolNodeRuntime(run_context) + node_data = ToolNodeData( + title="Workflow Tool Node", + desc=None, + provider_id="workflow-provider-id", + provider_type=ToolProviderType.WORKFLOW, + provider_name="workflow-provider", + tool_name="workflow-tool", + tool_label="Workflow Tool", + tool_configurations={}, + tool_parameters={}, + ) + + handle = runtime.get_runtime( + node_id="tool-node", + node_data=node_data, + variable_pool=object(), + node_execution_id="node-execution-id", + ) + + assert handle.raw.trace_session_id == "session-1" + assert runtime_tool.runtime.runtime_parameters == {} + + @staticmethod + def test_dify_tool_node_runtime_does_not_inject_outer_workflow_run_id_for_non_workflow_tools( + monkeypatch: pytest.MonkeyPatch, + ) -> None: + runtime_tool = SimpleNamespace(runtime=SimpleNamespace(runtime_parameters={})) + get_runtime = MagicMock(return_value=runtime_tool) + monkeypatch.setattr(node_runtime.ToolManager, "get_workflow_tool_runtime", get_runtime) + monkeypatch.setattr(node_runtime, "get_system_text", lambda _pool, _key: None) + + runtime = node_runtime.DifyToolNodeRuntime(_build_run_context()) + node_data = ToolNodeData( + title="Builtin Tool Node", + desc=None, + provider_id="builtin-provider-id", + provider_type=ToolProviderType.BUILT_IN, + provider_name="builtin-provider", + tool_name="builtin-tool", + tool_label="Builtin Tool", + tool_configurations={}, + tool_parameters={}, + ) + + handle = runtime.get_runtime( + node_id="tool-node", + node_data=node_data, + variable_pool=object(), + ) + + assert handle.raw.tool is runtime_tool + assert runtime_tool.runtime.dify_run_context.tenant_id == "tenant-id" + assert runtime_tool.runtime.dify_run_context.app_id == "app-id" + assert "outer_workflow_run_id" not in runtime_tool.runtime.runtime_parameters + get_runtime.assert_called_once() + + @staticmethod + def test_dify_tool_node_runtime_builds_workflow_tool_container_payload(monkeypatch: pytest.MonkeyPatch) -> None: + from core.tools.__base.tool_runtime import ToolRuntime + from core.tools.entities.common_entities import I18nObject + from core.tools.entities.tool_entities import ToolEntity, ToolIdentity + from core.tools.workflow_as_tool.tool import WorkflowTool + + tool = WorkflowTool( + workflow_app_id=str(uuid4()), + workflow_id=str(uuid4()), + workflow_as_tool_id="workflow-tool-id", + version="published-version", + workflow_call_depth=2, + entity=ToolEntity( + identity=ToolIdentity( + author="test", + name="workflow-tool", + label=I18nObject(en_US="Workflow Tool"), + provider="workflow-provider", + ), + parameters=[], + description=None, + has_runtime_parameters=False, + ), + runtime=ToolRuntime(tenant_id="tenant-id", invoke_from=InvokeFrom.DEBUGGER), + ) + prepared_parameters: list[Mapping[str, object]] = [] + + def prepare_container_inputs( + parameters: Mapping[str, object], + ) -> tuple[dict[str, object], list[dict[str, str]]]: + prepared_parameters.append(parameters) + return {"amount": Decimal("1.25")}, [{"id": "file-id", "name": "input.txt"}] + + monkeypatch.setattr(tool, "prepare_container_inputs", prepare_container_inputs) + runtime = DifyToolNodeRuntime(_build_run_context()) + + payload = runtime.build_workflow_tool_container_payload( + tool_runtime=ToolRuntimeHandle(raw=tool), + tool_parameters={"amount": "1.25"}, + inputs_for_log={"amount": Decimal("1.25")}, + workflow_call_depth=2, + ) + + assert payload.model_dump() == { + "version": "1", + "source_app_id": str(tool.workflow_app_id), + "source_workflow_id": str(tool.workflow_id), + "source_workflow_version": "published-version", + "inputs": {"amount": 1.25}, + "system_files": [{"id": "file-id", "name": "input.txt"}], + "inputs_for_log": {"amount": 1.25}, + "call_depth": 3, + } + assert prepared_parameters == [{"amount": "1.25"}] + + @staticmethod + def test_dify_tool_node_runtime_rejects_non_workflow_tool_container_payload() -> None: + runtime = DifyToolNodeRuntime(_build_run_context()) + with pytest.raises(ToolRuntimeResolutionError, match="resolved tool is not a Workflow Tool"): + runtime.build_workflow_tool_container_payload( + tool_runtime=ToolRuntimeHandle(raw=object()), + tool_parameters={}, + inputs_for_log={}, + workflow_call_depth=0, + ) + + @staticmethod + def test_dify_human_input_runtime_builds_debug_repository(monkeypatch: pytest.MonkeyPatch) -> None: + repository = MagicMock() + repository_cls = MagicMock(return_value=repository) + monkeypatch.setattr(node_runtime, "HumanInputFormRepositoryImpl", repository_cls) + + runtime = DifyHumanInputNodeRuntime( + _build_run_context(), + workflow_execution_id_getter=lambda: "workflow-execution-id", + ) + + assert runtime.build_form_repository() is repository + repository_cls.assert_called_once_with( + tenant_id="tenant-id", + app_id="app-id", + workflow_execution_id="workflow-execution-id", + invoke_source="debugger", + submission_actor_id="user-id", + ) + + @staticmethod + def test_dify_tool_runtime_spec_prefers_tool_parameters_for_runtime_form_values() -> None: + node_data = ToolNodeData( + provider_id="video-mixcut-agent", + provider_type=ToolProviderType.PLUGIN, + provider_name="sawyer-shi/video-mixcut-agent", + tool_name="mixcut", + tool_label="MixCut", + tool_configurations={"count": 2}, + tool_parameters={ + "vision_llm_model": { + "type": "constant", + "value": { + "provider": "langgenius/tongyi/tongyi", + "model": "qwen3-vl-plus", + "model_type": "llm", + }, + } + }, + ) + + spec = DifyToolNodeRuntime._build_tool_runtime_spec(node_data) + + assert spec.tool_configurations == { + "count": 2, "vision_llm_model": { "type": "constant", "value": { @@ -1036,195 +2286,180 @@ def test_dify_tool_runtime_spec_prefers_tool_parameters_for_runtime_form_values( "model": "qwen3-vl-plus", "model_type": "llm", }, - } - }, - ) - - spec = DifyToolNodeRuntime._build_tool_runtime_spec(node_data) - - assert spec.tool_configurations == { - "count": 2, - "vision_llm_model": { - "type": "constant", - "value": { - "provider": "langgenius/tongyi/tongyi", - "model": "qwen3-vl-plus", - "model_type": "llm", }, - }, - } + } + @staticmethod + def test_dify_human_input_runtime_create_form_filters_debugger_delivery_methods() -> None: + repository = MagicMock() + repository.create_form.return_value = sentinel.form + node_data = HumanInputNodeData( + title="Human Input", + delivery_methods=[ + WebAppDeliveryMethod(enabled=True, config=_WebAppDeliveryConfig()), + _build_email_method(debug_mode=True), + ], + ) + runtime = DifyHumanInputNodeRuntime( + _build_run_context(), + workflow_execution_id_getter=lambda: "workflow-execution-id", + form_repository=repository, + ) -def test_dify_human_input_runtime_create_form_filters_debugger_delivery_methods() -> None: - repository = MagicMock() - repository.create_form.return_value = sentinel.form - node_data = HumanInputNodeData( - title="Human Input", - delivery_methods=[ - WebAppDeliveryMethod(enabled=True, config=_WebAppDeliveryConfig()), - _build_email_method(debug_mode=True), - ], - ) - runtime = DifyHumanInputNodeRuntime( - _build_run_context(), - workflow_execution_id_getter=lambda: "workflow-execution-id", - form_repository=repository, - ) + result = runtime.create_form( + node_id="human-input-node", + node_data=node_data, + rendered_content="

Rendered

", + resolved_default_values={"answer": "default"}, + ) - result = runtime.create_form( - node_id="human-input-node", - node_data=node_data, - rendered_content="

Rendered

", - resolved_default_values={"answer": "default"}, - ) + assert result is sentinel.form + repository.create_form.assert_called_once() + params = repository.create_form.call_args.args[0] + assert params.node_id == "human-input-node" + assert params.workflow_execution_id == "workflow-execution-id" + # No conversation_id_getter wired -> a pure workflow run leaves it None. + assert params.conversation_id is None + assert params.display_in_ui is True + assert len(params.delivery_methods) == 1 + assert params.delivery_methods[0].type == DeliveryMethodType.EMAIL + assert params.delivery_methods[0].config.recipients.items[0].reference_id == "user-id" - assert result is sentinel.form - repository.create_form.assert_called_once() - params = repository.create_form.call_args.args[0] - assert params.node_id == "human-input-node" - assert params.workflow_execution_id == "workflow-execution-id" - # No conversation_id_getter wired -> a pure workflow run leaves it None. - assert params.conversation_id is None - assert params.display_in_ui is True - assert len(params.delivery_methods) == 1 - assert params.delivery_methods[0].type == DeliveryMethodType.EMAIL - assert params.delivery_methods[0].config.recipients.items[0].reference_id == "user-id" + @staticmethod + def test_dify_human_input_runtime_create_form_tags_conversation_id_for_chatflow() -> None: + # ENG-635 (review): a chatflow (advanced-chat) run carries a conversation, so its + # Human Input form is tagged with BOTH its workflow run and its conversation — + # making the form queryable per conversation without changing resume routing. + repository = MagicMock() + repository.create_form.return_value = sentinel.form + node_data = HumanInputNodeData( + title="Human Input", + delivery_methods=[WebAppDeliveryMethod(enabled=True, config=_WebAppDeliveryConfig())], + ) + runtime = DifyHumanInputNodeRuntime( + _build_run_context(), + workflow_execution_id_getter=lambda: "workflow-execution-id", + conversation_id_getter=lambda: "conversation-id", + form_repository=repository, + ) + runtime.create_form( + node_id="human-input-node", + node_data=node_data, + rendered_content="

Rendered

", + resolved_default_values={}, + ) -def test_dify_human_input_runtime_create_form_tags_conversation_id_for_chatflow() -> None: - # ENG-635 (review): a chatflow (advanced-chat) run carries a conversation, so its - # Human Input form is tagged with BOTH its workflow run and its conversation — - # making the form queryable per conversation without changing resume routing. - repository = MagicMock() - repository.create_form.return_value = sentinel.form - node_data = HumanInputNodeData( - title="Human Input", - delivery_methods=[WebAppDeliveryMethod(enabled=True, config=_WebAppDeliveryConfig())], - ) - runtime = DifyHumanInputNodeRuntime( - _build_run_context(), - workflow_execution_id_getter=lambda: "workflow-execution-id", - conversation_id_getter=lambda: "conversation-id", - form_repository=repository, - ) + params = repository.create_form.call_args.args[0] + assert params.workflow_execution_id == "workflow-execution-id" + assert params.conversation_id == "conversation-id" - runtime.create_form( - node_id="human-input-node", - node_data=node_data, - rendered_content="

Rendered

", - resolved_default_values={}, - ) + @staticmethod + def test_dify_human_input_runtime_preserves_webapp_delivery_for_web_invocations() -> None: + repository = MagicMock() + repository.create_form.return_value = sentinel.form + node_data = HumanInputNodeData( + title="Human Input", + delivery_methods=[ + WebAppDeliveryMethod(enabled=True, config=_WebAppDeliveryConfig()), + _build_email_method(debug_mode=True), + ], + ) + runtime = DifyHumanInputNodeRuntime( + _build_run_context(invoke_from=InvokeFrom.WEB_APP), + form_repository=repository, + ) - params = repository.create_form.call_args.args[0] - assert params.workflow_execution_id == "workflow-execution-id" - assert params.conversation_id == "conversation-id" + runtime.create_form( + node_id="human-input-node", + node_data=node_data, + rendered_content="

Rendered

", + resolved_default_values={}, + ) + params = repository.create_form.call_args.args[0] + assert params.display_in_ui is True + assert [method.type for method in params.delivery_methods] == [ + DeliveryMethodType.WEBAPP, + DeliveryMethodType.EMAIL, + ] + assert params.delivery_methods[1].config.recipients.include_bound_group is True -def test_dify_human_input_runtime_preserves_webapp_delivery_for_web_invocations() -> None: - repository = MagicMock() - repository.create_form.return_value = sentinel.form - node_data = HumanInputNodeData( - title="Human Input", - delivery_methods=[ - WebAppDeliveryMethod(enabled=True, config=_WebAppDeliveryConfig()), - _build_email_method(debug_mode=True), - ], - ) - runtime = DifyHumanInputNodeRuntime( - _build_run_context(invoke_from=InvokeFrom.WEB_APP), - form_repository=repository, - ) - - runtime.create_form( - node_id="human-input-node", - node_data=node_data, - rendered_content="

Rendered

", - resolved_default_values={}, - ) - - params = repository.create_form.call_args.args[0] - assert params.display_in_ui is True - assert [method.type for method in params.delivery_methods] == [ - DeliveryMethodType.WEBAPP, - DeliveryMethodType.EMAIL, - ] - assert params.delivery_methods[1].config.recipients.include_bound_group is True - - -def test_dify_human_input_runtime_restore_submitted_data_rehydrates_files() -> None: - runtime = DifyHumanInputNodeRuntime(_build_run_context()) - file_value = File( - file_id="file-1", - file_type=FileType.DOCUMENT, - transfer_method=FileTransferMethod.LOCAL_FILE, - related_id="upload-1", - filename="resume.pdf", - extension=".pdf", - mime_type="application/pdf", - size=128, - ) - file_list_value = [ - File( - file_id="file-2", + @staticmethod + def test_dify_human_input_runtime_restore_submitted_data_rehydrates_files() -> None: + runtime = DifyHumanInputNodeRuntime(_build_run_context()) + file_value = File( + file_id="file-1", file_type=FileType.DOCUMENT, transfer_method=FileTransferMethod.LOCAL_FILE, - related_id="upload-2", - filename="first.pdf", + related_id="upload-1", + filename="resume.pdf", extension=".pdf", mime_type="application/pdf", - size=64, - ), - File( - file_id="file-3", - file_type=FileType.DOCUMENT, - transfer_method=FileTransferMethod.REMOTE_URL, - remote_url="https://example.com/second.pdf", - filename="second.pdf", - extension=".pdf", - mime_type="application/pdf", - size=96, - ), - ] - runtime._file_reference_factory.build_from_mapping = MagicMock(side_effect=[file_value, *file_list_value]) # type: ignore[method-assign] + size=128, + ) + file_list_value = [ + File( + file_id="file-2", + file_type=FileType.DOCUMENT, + transfer_method=FileTransferMethod.LOCAL_FILE, + related_id="upload-2", + filename="first.pdf", + extension=".pdf", + mime_type="application/pdf", + size=64, + ), + File( + file_id="file-3", + file_type=FileType.DOCUMENT, + transfer_method=FileTransferMethod.REMOTE_URL, + remote_url="https://example.com/second.pdf", + filename="second.pdf", + extension=".pdf", + mime_type="application/pdf", + size=96, + ), + ] + runtime._file_reference_factory.build_from_mapping = MagicMock(side_effect=[file_value, *file_list_value]) # type: ignore[method-assign] - restored = runtime.restore_submitted_data( - inputs=[ - FileInputConfig(output_variable_name="attachment"), - FileListInputConfig(output_variable_name="attachments", number_limits=2), - ], - submitted_data={ - "attachment": {"upload_file_id": "upload-1", "type": "document", "transfer_method": "local_file"}, - "attachments": [ - {"upload_file_id": "upload-2", "type": "document", "transfer_method": "local_file"}, - { - "url": "https://example.com/second.pdf", - "type": "document", - "transfer_method": "remote_url", - }, + restored = runtime.restore_submitted_data( + inputs=[ + FileInputConfig(output_variable_name="attachment"), + FileListInputConfig(output_variable_name="attachments", number_limits=2), ], - }, - ) + submitted_data={ + "attachment": {"upload_file_id": "upload-1", "type": "document", "transfer_method": "local_file"}, + "attachments": [ + {"upload_file_id": "upload-2", "type": "document", "transfer_method": "local_file"}, + { + "url": "https://example.com/second.pdf", + "type": "document", + "transfer_method": "remote_url", + }, + ], + }, + ) - assert restored["attachment"] is file_value - assert restored["attachments"] == file_list_value - assert isinstance(FileSegment(value=restored["attachment"]), FileSegment) - assert isinstance(ArrayFileSegment(value=restored["attachments"]), ArrayFileSegment) + assert restored["attachment"] is file_value + assert restored["attachments"] == file_list_value + assert isinstance(FileSegment(value=restored["attachment"]), FileSegment) + assert isinstance(ArrayFileSegment(value=restored["attachments"]), ArrayFileSegment) + @staticmethod + def test_build_dify_llm_file_saver_wires_runtime_adapters(monkeypatch: pytest.MonkeyPatch) -> None: + file_saver_cls = MagicMock(return_value=sentinel.file_saver) + monkeypatch.setattr("graphon.nodes.llm.file_saver.FileSaverImpl", file_saver_cls) + http_client = MagicMock() -def test_build_dify_llm_file_saver_wires_runtime_adapters(monkeypatch: pytest.MonkeyPatch) -> None: - file_saver_cls = MagicMock(return_value=sentinel.file_saver) - monkeypatch.setattr("graphon.nodes.llm.file_saver.FileSaverImpl", file_saver_cls) - http_client = MagicMock() + saver = build_dify_llm_file_saver( + run_context=_build_run_context(), + http_client=http_client, + conversation_id_getter=lambda: "conversation-id", + ) - saver = build_dify_llm_file_saver( - run_context=_build_run_context(), - http_client=http_client, - conversation_id_getter=lambda: "conversation-id", - ) - - assert saver is sentinel.file_saver - tool_file_manager = file_saver_cls.call_args.kwargs["tool_file_manager"] - file_reference_factory = file_saver_cls.call_args.kwargs["file_reference_factory"] - assert isinstance(tool_file_manager, DifyToolFileManager) - assert isinstance(file_reference_factory, DifyFileReferenceFactory) - assert file_saver_cls.call_args.kwargs["http_client"] is http_client + assert saver is sentinel.file_saver + tool_file_manager = file_saver_cls.call_args.kwargs["tool_file_manager"] + file_reference_factory = file_saver_cls.call_args.kwargs["file_reference_factory"] + assert isinstance(tool_file_manager, DifyToolFileManager) + assert isinstance(file_reference_factory, DifyFileReferenceFactory) + assert file_saver_cls.call_args.kwargs["http_client"] is http_client diff --git a/api/tests/unit_tests/core/workflow/test_response_stream_filter_issue_170.py b/api/tests/unit_tests/core/workflow/test_response_stream_filter_issue_170.py deleted file mode 100644 index 9a5476cc245..00000000000 --- a/api/tests/unit_tests/core/workflow/test_response_stream_filter_issue_170.py +++ /dev/null @@ -1,75 +0,0 @@ -"""Unit coverage for Dify's in-memory ResponseStreamFilter boundary behavior.""" - -from core.workflow.workflow_entry import iter_dify_graph_engine_events -from graphon.engine import Engine -from graphon.engine.command import InMemoryChannel -from graphon.engine_events import GraphRunSucceededEvent, NodeRunStreamChunkEvent -from tests.unit_tests.core.workflow.graph_engine.test_mock_config import MockConfig, MockConfigBuilder -from tests.unit_tests.core.workflow.graph_engine.test_table_runner import WorkflowRunner - - -def _build_issue_170_mock_config() -> tuple[WorkflowRunner, MockConfig]: - runner = WorkflowRunner() - mock_config = ( - MockConfigBuilder() - .with_node_output( - "llm", - { - "text": "Quiet Night Thought", - "usage": { - "prompt_tokens": 10, - "completion_tokens": 5, - "total_tokens": 15, - }, - "finish_reason": "stop", - }, - ) - .with_node_output( - "dufu", - { - "text": "Spring View", - "usage": { - "prompt_tokens": 10, - "completion_tokens": 5, - "total_tokens": 15, - }, - "finish_reason": "stop", - }, - ) - .build() - ) - - return runner, mock_config - - -def test_dify_response_stream_filter_handles_issue_170_shape() -> None: - runner, mock_config = _build_issue_170_mock_config() - fixture_data = runner.load_fixture("response_stream_filter_issue_170_workflow") - graph, graph_runtime_state = runner.create_graph_from_fixture( - fixture_data=fixture_data, - query="1", - use_mock_factory=True, - mock_config=mock_config, - ) - - expected_answer = "# Du Fu\n\nSpring View\n\n# Li Bai\n\nQuiet Night Thought" - - engine = Engine( - graph=graph, - runtime_state=graph_runtime_state, - command_channel=InMemoryChannel(), - ) - events = list(iter_dify_graph_engine_events(engine)) - - stream_chunk_events = [event for event in events if isinstance(event, NodeRunStreamChunkEvent)] - success_events = [event for event in events if isinstance(event, GraphRunSucceededEvent)] - - assert success_events - assert stream_chunk_events - actual_answer = "".join(event.chunk for event in stream_chunk_events) - assert actual_answer.strip() == expected_answer - assert stream_chunk_events[-1].is_final is True - final_answer = success_events[-1].outputs["answer"] - assert isinstance(final_answer, str) - assert final_answer.strip() == expected_answer - assert actual_answer.strip() == final_answer.strip() diff --git a/api/tests/unit_tests/core/workflow/test_response_stream_filter_pause_resume.py b/api/tests/unit_tests/core/workflow/test_response_stream_filter_pause_resume.py deleted file mode 100644 index 2f16ccbca0a..00000000000 --- a/api/tests/unit_tests/core/workflow/test_response_stream_filter_pause_resume.py +++ /dev/null @@ -1,207 +0,0 @@ -"""In-memory regression test: if-else branch + human_input pause + downstream answer nodes. - -Reproduces https://github.com/langgenius/dify/issues/38525 at the -iter_dify_graph_engine_events layer: without a restored ResponseStreamFilter, -answer nodes downstream of a pre-pause branch never unlock for streaming on -resume, even though the graph executes correctly. -""" - -from pytest_mock import MockerFixture -from sqlalchemy import select -from sqlalchemy.orm import Session, sessionmaker - -from core.app.entities.app_invoke_entities import InvokeFrom, UserFrom -from core.repositories.human_input_repository import ( - HumanInputFormRepository, - HumanInputFormRepositoryImpl, -) -from core.workflow.nodes.human_input.callback import DifyHITLCallback -from core.workflow.nodes.human_input.entities import HumanInputNodeData, UserActionConfig -from core.workflow.nodes.human_input.enums import HumanInputFormStatus -from core.workflow.system_variables import build_system_variables -from core.workflow.workflow_entry import iter_dify_graph_engine_events -from graphon.engine import Engine -from graphon.engine.command import InMemoryChannel -from graphon.engine.filter import EngineEventFilterContext, ResponseStreamFilter, filter_engine_events -from graphon.engine_events import GraphRunPausedEvent, GraphRunSucceededEvent, NodeRunStreamChunkEvent -from graphon.graph import Graph -from graphon.nodes.answer.answer_node import AnswerNode -from graphon.nodes.answer.entities import AnswerNodeData -from graphon.nodes.human_input.entities import HumanInputNodeData as GraphonHumanInputNodeData -from graphon.nodes.human_input.human_input_node import HumanInputNode -from graphon.nodes.if_else.entities import IfElseNodeData -from graphon.nodes.if_else.if_else_node import IfElseNode -from graphon.nodes.start.entities import StartNodeData -from graphon.nodes.start.start_node import StartNode -from graphon.runtime import RuntimeState, VariablePool -from graphon.utils.condition.entities import Condition -from libs.datetime_utils import naive_utc_now -from models.human_input import HumanInputForm -from tests.workflow_test_utils import build_test_graph_init_params - -WORKFLOW_EXECUTION_ID = "wf-exec-38525" - - -def _build_graph(runtime_state: RuntimeState, form_repository: HumanInputFormRepository) -> Graph: - params = build_test_graph_init_params( - workflow_id="wf", - graph_config={"nodes": [], "edges": []}, - user_from=UserFrom.ACCOUNT, - invoke_from=InvokeFrom.DEBUGGER, - ) - - start_node = StartNode( - node_id="start", - data=StartNodeData(title="start", variables=[]), - init_params=params, - runtime_state=runtime_state, - ) - - if_else_node = IfElseNode( - node_id="if_else", - data=IfElseNodeData( - title="if-else", - cases=[ - IfElseNodeData.Case( - case_id="true", - logical_operator="and", - conditions=[ - Condition( - variable_selector=["start", "category"], - comparison_operator="is", - value="fruit", - ) - ], - ) - ], - ), - init_params=params, - runtime_state=runtime_state, - ) - - human_data = HumanInputNodeData( - title="human", - form_content="Awaiting human input", - inputs=[], - user_actions=[UserActionConfig(id="continue", title="Continue")], - ) - human_node = HumanInputNode( - node_id="human_input", - data=GraphonHumanInputNodeData.model_validate(human_data.model_dump()), - init_params=params, - runtime_state=runtime_state, - hitl_callback=DifyHITLCallback(form_repository=form_repository, node_data=human_data), - ) - - answer_false_node = AnswerNode( - node_id="answer_false", - data=AnswerNodeData(title="answer_false", answer="unreachable branch"), - init_params=params, - runtime_state=runtime_state, - ) - - answer_after_pause = AnswerNode( - node_id="answer_after_pause", - data=AnswerNodeData(title="answer_after_pause", answer="Post-branch answer chunk 1"), - init_params=params, - runtime_state=runtime_state, - ) - - answer_after_pause_2 = AnswerNode( - node_id="answer_after_pause_2", - data=AnswerNodeData(title="answer_after_pause_2", answer="Post-branch answer chunk 2"), - init_params=params, - runtime_state=runtime_state, - ) - - return ( - Graph.new() - .add_root(start_node) - .add_node(if_else_node, from_node_id="start") - .add_node(human_node, from_node_id="if_else", source_handle="true") - .add_node(answer_false_node, from_node_id="if_else", source_handle="false") - .add_node(answer_after_pause, from_node_id="human_input", source_handle="continue") - .add_node(answer_after_pause_2, from_node_id="answer_after_pause") - .build() - ) - - -def _build_runtime_state() -> RuntimeState: - variable_pool = VariablePool.from_bootstrap( - system_variables=build_system_variables( - workflow_execution_id=WORKFLOW_EXECUTION_ID, - app_id="app", - workflow_id="wf", - user_id="user", - ), - user_inputs={}, - conversation_variables=[], - ) - variable_pool.add(("start", "category"), "fruit") # drives the if-else "true" branch - return RuntimeState(workflow_id="test-workflow", variable_pool=variable_pool, start_at=0.0) - - -def test_if_else_human_input_pause_resume_answer_chunks_survive_resume( - sqlite_session_factory: sessionmaker[Session], - mocker: MockerFixture, -) -> None: - mocker.patch( - "core.repositories.human_input_repository.session_factory.create_session", side_effect=sqlite_session_factory - ) - repository = HumanInputFormRepositoryImpl( - tenant_id="tenant", app_id="app", workflow_execution_id=WORKFLOW_EXECUTION_ID - ) - # ---- Phase 1: run to GraphRunPausedEvent ---- - runtime_state_1 = _build_runtime_state() - graph_1 = _build_graph(runtime_state_1, repository) - engine_1 = Engine( - graph=graph_1, - runtime_state=runtime_state_1, - command_channel=InMemoryChannel(), - ) - filter_1 = ResponseStreamFilter() - phase1_events = list( - filter_engine_events( - engine_1.run(), - context=EngineEventFilterContext.from_engine(engine_1), - filters=[filter_1], - ) - ) - - assert any(isinstance(e, GraphRunPausedEvent) for e in phase1_events) - phase1_chunks = [e for e in phase1_events if isinstance(e, NodeRunStreamChunkEvent)] - assert not any(e.node_id in ("answer_after_pause", "answer_after_pause_2") for e in phase1_chunks) - - response_filter_snapshot = filter_1.dumps() - runtime_snapshot = runtime_state_1.dumps() - - with sqlite_session_factory.begin() as session: - form = session.scalar(select(HumanInputForm).where(HumanInputForm.workflow_run_id == WORKFLOW_EXECUTION_ID)) - assert form is not None - assert form.status == HumanInputFormStatus.WAITING - form.selected_action_id = "continue" - form.submitted_data = "{}" - form.submitted_at = naive_utc_now() - form.status = HumanInputFormStatus.SUBMITTED - - # ---- Phase 2: rebuild engine + filter from snapshots, resume to completion ---- - runtime_state_2 = RuntimeState.from_snapshot(runtime_snapshot) - graph_2 = _build_graph(runtime_state_2, repository) - engine_2 = Engine( - graph=graph_2, - runtime_state=runtime_state_2, - command_channel=InMemoryChannel(), - ) - filter_2 = ResponseStreamFilter() - filter_2.loads(response_filter_snapshot) - - phase2_events = list(iter_dify_graph_engine_events(engine_2, filter_2)) - - assert any(isinstance(e, GraphRunSucceededEvent) for e in phase2_events) - - phase2_chunks = [e for e in phase2_events if isinstance(e, NodeRunStreamChunkEvent)] - answer_1_chunks = [e for e in phase2_chunks if e.node_id == "answer_after_pause"] - answer_2_chunks = [e for e in phase2_chunks if e.node_id == "answer_after_pause_2"] - - assert answer_1_chunks, "answer_after_pause produced no stream chunks after resume" - assert answer_2_chunks, "answer_after_pause_2 produced no stream chunks after resume" diff --git a/api/tests/unit_tests/core/workflow/test_system_variable.py b/api/tests/unit_tests/core/workflow/test_system_variables.py similarity index 100% rename from api/tests/unit_tests/core/workflow/test_system_variable.py rename to api/tests/unit_tests/core/workflow/test_system_variables.py diff --git a/api/tests/unit_tests/core/workflow/test_workflow_entry.py b/api/tests/unit_tests/core/workflow/test_workflow_entry.py index 304b53344fb..d2f15fa4f7e 100644 --- a/api/tests/unit_tests/core/workflow/test_workflow_entry.py +++ b/api/tests/unit_tests/core/workflow/test_workflow_entry.py @@ -1,509 +1,543 @@ import json +from collections import UserString +from datetime import datetime from types import SimpleNamespace +from unittest.mock import MagicMock, patch, sentinel import pytest +from pytest_mock import MockerFixture +from sqlalchemy import select +from sqlalchemy.orm import Session, sessionmaker from configs import dify_config +from core.app.apps.exc import GenerateTaskStoppedError +from core.app.entities.app_invoke_entities import InvokeFrom, UserFrom +from core.credit_usage import CreditUsageAppType from core.helper.code_executor.code_executor import CodeLanguage +from core.repositories.human_input_repository import ( + HumanInputFormRepository, + HumanInputFormRepositoryImpl, +) +from core.workflow import workflow_entry +from core.workflow.nodes.human_input.callback import DifyHITLCallback +from core.workflow.nodes.human_input.entities import HumanInputNodeData, UserActionConfig +from core.workflow.nodes.human_input.enums import HumanInputFormStatus from core.workflow.system_variables import build_system_variables, default_system_variables from core.workflow.variable_prefixes import ( CONVERSATION_VARIABLE_NODE_ID, ENVIRONMENT_VARIABLE_NODE_ID, ) -from core.workflow.workflow_entry import WorkflowEntry +from core.workflow.workflow_entry import WorkflowEntry, iter_dify_graph_engine_events +from graphon.engine import Engine +from graphon.engine.command import InMemoryChannel, RedisChannel +from graphon.engine.filter import EngineEventFilterContext, ResponseStreamFilter, filter_engine_events +from graphon.engine_events import ( + GraphRunAbortedEvent, + GraphRunFailedEvent, + GraphRunPausedEvent, + GraphRunSucceededEvent, + NodeRunStreamChunkEvent, + NodeRunSucceededEvent, +) +from graphon.entities.base_node_data import BaseNodeData +from graphon.enums import NodeType, WorkflowNodeExecutionStatus +from graphon.errors import WorkflowNodeRunFailedError from graphon.file import File, FileTransferMethod, FileType +from graphon.file.runtime import peek_workflow_file_runtime, use_workflow_file_runtime +from graphon.graph import Graph +from graphon.node_events import NodeRunResult +from graphon.nodes import BuiltinNodeTypes +from graphon.nodes.answer.answer_node import AnswerNode +from graphon.nodes.answer.entities import AnswerNodeData from graphon.nodes.code.code_node import CodeNode from graphon.nodes.code.limits import CodeNodeLimits -from graphon.runtime import VariablePool +from graphon.nodes.human_input.entities import HumanInputNodeData as GraphonHumanInputNodeData +from graphon.nodes.human_input.human_input_node import HumanInputNode +from graphon.nodes.if_else.entities import IfElseNodeData +from graphon.nodes.if_else.if_else_node import IfElseNode +from graphon.nodes.start.entities import StartNodeData +from graphon.nodes.start.start_node import StartNode +from graphon.runtime import RuntimeState, VariablePool +from graphon.utils.condition.entities import Condition from graphon.variables.variables import StringVariable +from libs.datetime_utils import naive_utc_now +from models.human_input import HumanInputForm from models.workflow import Workflow, WorkflowType +from tests.helpers.unit.config_override import config_overrides_context +from tests.helpers.unit.model_factories import make_workflow +from tests.helpers.workflow.mock_config import MockConfig, MockConfigBuilder +from tests.helpers.workflow.table_runner import WorkflowRunner +from tests.unit_tests.core.workflow.test_workflow_tool_container_handler import ( + _workflow_tool_container_outer_graph as _outer_graph, +) +from tests.unit_tests.core.workflow.test_workflow_tool_container_handler import ( + _workflow_tool_container_workflow_tool_node as _workflow_tool_node, +) +from tests.unit_tests.core.workflow.test_workflow_tool_container_handler import ( + _workflow_tool_container_workflow_tool_source as _workflow_tool_source, +) +from tests.workflow_test_utils import build_test_graph_init_params + +"""Unit coverage for Dify's in-memory ResponseStreamFilter boundary behavior.""" -@pytest.fixture(autouse=True) -def _mock_ssrf_head(monkeypatch: pytest.MonkeyPatch): - """Avoid any real network requests during tests. - - factories.file_factory.remote.get_remote_file_info() uses remote_fetcher.make_request - to inspect remote files. We stub it to return a minimal response object with - headers so filename/mime/size can be derived deterministically. - """ - - def fake_head(method, url, *args, **kwargs): - assert method == "HEAD" - # choose a content-type by file suffix for determinism - if url.endswith(".pdf"): - ctype = "application/pdf" - elif url.endswith(".jpg") or url.endswith(".jpeg"): - ctype = "image/jpeg" - elif url.endswith(".png"): - ctype = "image/png" - else: - ctype = "application/octet-stream" - filename = url.split("/")[-1] or "file.bin" - headers = { - "Content-Type": ctype, - "Content-Disposition": f'attachment; filename="{filename}"', - "Content-Length": "12345", - } - return SimpleNamespace(status_code=200, headers=headers) - - monkeypatch.setattr("factories.file_factory.remote.remote_fetcher.make_request", fake_head) - - -class TestWorkflowEntry: - """Test WorkflowEntry class methods.""" - - @pytest.mark.parametrize("value", [0, False, "", [], {}]) - def test_qualified_input_preserves_falsy_values(self, value): - variable_pool = VariablePool() - - WorkflowEntry.mapping_user_inputs_to_variable_pool( - variable_mapping={"node.value": ["source", "value"]}, - user_inputs={"node.value": value, "value": "fallback"}, - variable_pool=variable_pool, - tenant_id="tenant", - ) - - variable = variable_pool.get(["source", "value"]) - assert variable is not None - assert variable.value == value - - @pytest.mark.parametrize("root_name", ["object", "structured_output"]) - @pytest.mark.parametrize("value", [0, False, "", [], {}]) - def test_nested_input_overrides_preserve_siblings(self, root_name, value): - variable_pool = VariablePool() - original = {"keep": "root sibling", "nested": {"leaf": "old", "keep": "nested sibling"}} - variable_pool.add(["source", root_name], original) - key = f"node.#source.{root_name}.nested.leaf#" - - WorkflowEntry.mapping_user_inputs_to_variable_pool( - variable_mapping={key: ["source", root_name, "nested", "leaf"]}, - user_inputs={key: value, key.split(".", 1)[1]: "fallback"}, - variable_pool=variable_pool, - tenant_id="tenant", - ) - - variable = variable_pool.get(["source", root_name]) - assert variable is not None - assert variable.value == {"keep": "root sibling", "nested": {"leaf": value, "keep": "nested sibling"}} - assert original["nested"]["leaf"] == "old" - - def test_nested_inputs_build_missing_root_and_preserve_other_explicit_inputs(self): - variable_pool = VariablePool() - - WorkflowEntry.mapping_user_inputs_to_variable_pool( - variable_mapping={ - "node.#source.object.leaf#": ["source", "object", "leaf"], - "node.#source.object.nested.other#": ["source", "object", "nested", "other"], +def _build_issue_170_mock_config() -> tuple[WorkflowRunner, MockConfig]: + runner = WorkflowRunner() + mock_config = ( + MockConfigBuilder() + .with_node_output( + "llm", + { + "text": "Quiet Night Thought", + "usage": { + "prompt_tokens": 10, + "completion_tokens": 5, + "total_tokens": 15, + }, + "finish_reason": "stop", }, - user_inputs={"#source.object.leaf#": 0, "#source.object.nested.other#": False}, - variable_pool=variable_pool, - tenant_id="tenant", ) - - variable = variable_pool.get(["source", "object"]) - assert variable is not None - assert variable.value == {"leaf": 0, "nested": {"other": False}} - - @pytest.mark.parametrize("descendant_first", [False, True]) - @pytest.mark.parametrize("value", ["Alice", 0, False, "", [], {}]) - def test_nested_overrides_take_precedence_over_explicit_ancestors_in_either_order(self, descendant_first, value): - variable_pool = VariablePool() - mappings = [ - ("node.#source.object#", ["source", "object"]), - ("node.#source.object.name#", ["source", "object", "name"]), - ] - if descendant_first: - mappings.reverse() - root_input = {"other": "kept", "name": "old"} - - WorkflowEntry.mapping_user_inputs_to_variable_pool( - variable_mapping=dict(mappings), - user_inputs={"#source.object#": root_input, "#source.object.name#": value}, - variable_pool=variable_pool, - tenant_id="tenant", + .with_node_output( + "dufu", + { + "text": "Spring View", + "usage": { + "prompt_tokens": 10, + "completion_tokens": 5, + "total_tokens": 15, + }, + "finish_reason": "stop", + }, ) - - variable = variable_pool.get(["source", "object"]) - assert variable is not None - assert variable.value == {"other": "kept", "name": value} - assert root_input == {"other": "kept", "name": "old"} - - @pytest.mark.parametrize( - ("data", "inputs", "expected"), - [ - ( - {"type": "end", "outputs": [{"variable": "result", "value_selector": ["source", "value"]}]}, - {"#source.value#": "end input"}, - "end input", - ), - ( - { - "type": "variable-aggregator", - "output_type": "string", - "variables": [["missing", "value"], ["source", "value"]], - }, - {"#source.value#": "second branch"}, - "second branch", - ), - ( - { - "type": "variable-assigner", - "output_type": "string", - "variables": [["source", "value"]], - }, - {"#source.value#": "legacy branch"}, - "legacy branch", - ), - ( - { - "type": "list-operator", - "variable": ["source", "value"], - "filter_by": {"enabled": False}, - "order_by": {"enabled": False}, - "limit": {"enabled": False}, - }, - {"#source.value#": [3, 1, 2]}, - [3, 1, 2], - ), - ], + .build() ) - def test_single_step_loads_explicit_inputs_for_nodes_without_graphon_extractors(self, data, inputs, expected): - workflow = Workflow.new( - tenant_id="tenant", - app_id="app", - type=WorkflowType.WORKFLOW, - version=Workflow.VERSION_DRAFT, - graph=json.dumps({"nodes": [{"id": "node", "data": {"title": "Node", **data}}], "edges": []}), - features="{}", - created_by="account", - environment_variables=[], - conversation_variables=[], - rag_pipeline_variables=[], - ) - workflow.id = "workflow" - variable_pool = VariablePool() - WorkflowEntry.single_step_run( - workflow=workflow, - node_id="node", - user_id="user", - user_inputs=inputs, - variable_pool=variable_pool, - execution_driver=None, + return runner, mock_config + + +class TestResponseStreamFilterIssue170Cases: + @staticmethod + def test_dify_response_stream_filter_handles_issue_170_shape() -> None: + runner, mock_config = _build_issue_170_mock_config() + fixture_data = runner.load_fixture("response_stream_filter_issue_170_workflow") + graph, graph_runtime_state = runner.create_graph_from_fixture( + fixture_data=fixture_data, + query="1", + use_mock_factory=True, + mock_config=mock_config, ) - variable = variable_pool.get(["source", "value"]) - assert variable is not None - assert variable.value == expected + expected_answer = "# Du Fu\n\nSpring View\n\n# Li Bai\n\nQuiet Night Thought" - def test_mapping_user_inputs_to_variable_pool_with_system_variables(self): - """Test mapping system variables from user inputs to variable pool.""" - # Initialize variable pool with system variables - variable_pool = VariablePool.from_bootstrap( - system_variables=build_system_variables( - user_id="test_user_id", - app_id="test_app_id", - workflow_id="test_workflow_id", - ), - user_inputs={}, + engine = Engine( + graph=graph, + runtime_state=graph_runtime_state, + command_channel=InMemoryChannel(), ) + events = list(iter_dify_graph_engine_events(engine)) - # Define variable mapping - sys variables mapped to other nodes - variable_mapping = { - "node1.input1": ["node1", "input1"], # Regular mapping - "node2.query": ["node2", "query"], # Regular mapping - "sys.user_id": ["output_node", "user"], # System variable mapping - } + stream_chunk_events = [event for event in events if isinstance(event, NodeRunStreamChunkEvent)] + success_events = [event for event in events if isinstance(event, GraphRunSucceededEvent)] - # User inputs including sys variables - user_inputs = { - "node1.input1": "new_user_id", - "node2.query": "test query", - "sys.user_id": "system_user", - } + assert success_events + assert stream_chunk_events + actual_answer = "".join(event.chunk for event in stream_chunk_events) + assert actual_answer.strip() == expected_answer + assert stream_chunk_events[-1].is_final is True + final_answer = success_events[-1].outputs["answer"] + assert isinstance(final_answer, str) + assert final_answer.strip() == expected_answer + assert actual_answer.strip() == final_answer.strip() - # Execute mapping - WorkflowEntry.mapping_user_inputs_to_variable_pool( - variable_mapping=variable_mapping, - user_inputs=user_inputs, - variable_pool=variable_pool, - tenant_id="test_tenant", - ) - # Verify variables were added to pool - # Note: variable_pool.get returns Variable objects, not raw values - node1_var = variable_pool.get(["node1", "input1"]) - assert node1_var is not None - assert node1_var.value == "new_user_id" +"""In-memory regression test: if-else branch + human_input pause + downstream answer nodes. - node2_var = variable_pool.get(["node2", "query"]) - assert node2_var is not None - assert node2_var.value == "test query" +Reproduces https://github.com/langgenius/dify/issues/38525 at the +iter_dify_graph_engine_events layer: without a restored ResponseStreamFilter, +answer nodes downstream of a pre-pause branch never unlock for streaming on +resume, even though the graph executes correctly. +""" - # System variable gets mapped to output node - output_var = variable_pool.get(["output_node", "user"]) - assert output_var is not None - assert output_var.value == "system_user" - def test_single_step_run_injects_code_limits(self): - """Ensure single-step CodeNode execution configures limits.""" - # Arrange - node_id = "code_node" - node_data = { - "type": "code", - "title": "Code", - "desc": None, - "variables": [], - "code_language": CodeLanguage.PYTHON3, - "code": "def main():\n return {}", - "outputs": {}, - } - node_config = {"id": node_id, "data": node_data} +WORKFLOW_EXECUTION_ID = "wf-exec-38525" - workflow = Workflow.new( - tenant_id="tenant", - app_id="app", - type=WorkflowType.WORKFLOW, - version=Workflow.VERSION_DRAFT, - graph=json.dumps({"nodes": [node_config], "edges": []}), - features="{}", - created_by="account", - environment_variables=[], - conversation_variables=[], - rag_pipeline_variables=[], - ) - workflow.id = "workflow" - variable_pool = VariablePool.from_bootstrap(system_variables=default_system_variables(), user_inputs={}) - expected_limits = CodeNodeLimits( - max_string_length=dify_config.CODE_MAX_STRING_LENGTH, - max_number=dify_config.CODE_MAX_NUMBER, - min_number=dify_config.CODE_MIN_NUMBER, - max_precision=dify_config.CODE_MAX_PRECISION, - max_depth=dify_config.CODE_MAX_DEPTH, - max_number_array_length=dify_config.CODE_MAX_NUMBER_ARRAY_LENGTH, - max_string_array_length=dify_config.CODE_MAX_STRING_ARRAY_LENGTH, - max_object_array_length=dify_config.CODE_MAX_OBJECT_ARRAY_LENGTH, - ) - # Act - node, _ = WorkflowEntry.single_step_run( - workflow=workflow, - node_id=node_id, - user_id="user", - user_inputs={}, - variable_pool=variable_pool, - execution_driver=None, - ) +def _build_graph(runtime_state: RuntimeState, form_repository: HumanInputFormRepository) -> Graph: + params = build_test_graph_init_params( + workflow_id="wf", + graph_config={"nodes": [], "edges": []}, + user_from=UserFrom.ACCOUNT, + invoke_from=InvokeFrom.DEBUGGER, + ) - # Assert - assert isinstance(node, CodeNode) - assert node._limits == expected_limits + start_node = StartNode( + node_id="start", + data=StartNodeData(title="start", variables=[]), + init_params=params, + runtime_state=runtime_state, + ) - def test_mapping_user_inputs_to_variable_pool_with_env_variables(self): - """Test mapping environment variables from user inputs to variable pool.""" - # Initialize variable pool with environment variables - env_var = StringVariable(name="API_KEY", value="existing_key") - variable_pool = VariablePool.from_bootstrap( - system_variables=default_system_variables(), - environment_variables=[env_var], - user_inputs={}, - ) - - # Add env variable to pool (simulating initialization) - variable_pool.add([ENVIRONMENT_VARIABLE_NODE_ID, "API_KEY"], env_var) - - # Define variable mapping - env variables should not be overridden - variable_mapping = { - "node1.api_key": [ENVIRONMENT_VARIABLE_NODE_ID, "API_KEY"], - "node2.new_env": [ENVIRONMENT_VARIABLE_NODE_ID, "NEW_ENV"], - } - - # User inputs - user_inputs = { - "node1.api_key": "user_provided_key", # This should not override existing env var - "node2.new_env": "new_env_value", - } - - # Execute mapping - WorkflowEntry.mapping_user_inputs_to_variable_pool( - variable_mapping=variable_mapping, - user_inputs=user_inputs, - variable_pool=variable_pool, - tenant_id="test_tenant", - ) - - # Verify env variable was not overridden - env_value = variable_pool.get([ENVIRONMENT_VARIABLE_NODE_ID, "API_KEY"]) - assert env_value is not None - assert env_value.value == "existing_key" # Should remain unchanged - - # New env variables from user input should not be added - assert variable_pool.get([ENVIRONMENT_VARIABLE_NODE_ID, "NEW_ENV"]) is None - - def test_mapping_user_inputs_to_variable_pool_with_conversation_variables(self): - """Test mapping conversation variables from user inputs to variable pool.""" - # Initialize variable pool with conversation variables - conv_var = StringVariable(name="last_message", value="Hello") - variable_pool = VariablePool.from_bootstrap( - system_variables=default_system_variables(), - conversation_variables=[conv_var], - user_inputs={}, - ) - - # Add conversation variable to pool - variable_pool.add([CONVERSATION_VARIABLE_NODE_ID, "last_message"], conv_var) - - # Define variable mapping - variable_mapping = { - "node1.message": ["node1", "message"], # Map to regular node - "conversation.context": ["chat_node", "context"], # Conversation var to regular node - } - - # User inputs - user_inputs = { - "node1.message": "Updated message", - "conversation.context": "New context", - } - - # Execute mapping - WorkflowEntry.mapping_user_inputs_to_variable_pool( - variable_mapping=variable_mapping, - user_inputs=user_inputs, - variable_pool=variable_pool, - tenant_id="test_tenant", - ) - - # Verify variables were added to their target nodes - node1_var = variable_pool.get(["node1", "message"]) - assert node1_var is not None - assert node1_var.value == "Updated message" - - chat_var = variable_pool.get(["chat_node", "context"]) - assert chat_var is not None - assert chat_var.value == "New context" - - def test_mapping_user_inputs_to_variable_pool_with_regular_variables(self): - """Test mapping regular node variables from user inputs to variable pool.""" - # Initialize empty variable pool - variable_pool = VariablePool.from_bootstrap( - system_variables=default_system_variables(), - user_inputs={}, - ) - - # Define variable mapping for regular nodes - variable_mapping = { - "input_node.text": ["input_node", "text"], - "llm_node.prompt": ["llm_node", "prompt"], - "code_node.input": ["code_node", "input"], - } - - # User inputs - user_inputs = { - "input_node.text": "User input text", - "llm_node.prompt": "Generate a summary", - "code_node.input": {"key": "value"}, - } - - # Execute mapping - WorkflowEntry.mapping_user_inputs_to_variable_pool( - variable_mapping=variable_mapping, - user_inputs=user_inputs, - variable_pool=variable_pool, - tenant_id="test_tenant", - ) - - # Verify regular variables were added - text_var = variable_pool.get(["input_node", "text"]) - assert text_var is not None - assert text_var.value == "User input text" - - prompt_var = variable_pool.get(["llm_node", "prompt"]) - assert prompt_var is not None - assert prompt_var.value == "Generate a summary" - - input_var = variable_pool.get(["code_node", "input"]) - assert input_var is not None - assert input_var.value == {"key": "value"} - - def test_mapping_user_inputs_with_file_handling(self): - """Test mapping file inputs from user inputs to variable pool.""" - variable_pool = VariablePool.from_bootstrap( - system_variables=default_system_variables(), - user_inputs={}, - ) - - # Define variable mapping - variable_mapping = { - "file_node.file": ["file_node", "file"], - "file_node.files": ["file_node", "files"], - } - - # User inputs with file data - using remote_url which doesn't require upload_file_id - user_inputs = { - "file_node.file": { - "type": "document", - "transfer_method": "remote_url", - "url": "http://example.com/test.pdf", - }, - "file_node.files": [ - { - "type": "image", - "transfer_method": "remote_url", - "url": "http://example.com/image1.jpg", - }, - { - "type": "image", - "transfer_method": "remote_url", - "url": "http://example.com/image2.jpg", - }, + if_else_node = IfElseNode( + node_id="if_else", + data=IfElseNodeData( + title="if-else", + cases=[ + IfElseNodeData.Case( + case_id="true", + logical_operator="and", + conditions=[ + Condition( + variable_selector=["start", "category"], + comparison_operator="is", + value="fruit", + ) + ], + ) ], - } + ), + init_params=params, + runtime_state=runtime_state, + ) - # Execute mapping - WorkflowEntry.mapping_user_inputs_to_variable_pool( - variable_mapping=variable_mapping, - user_inputs=user_inputs, - variable_pool=variable_pool, - tenant_id="test_tenant", + human_data = HumanInputNodeData( + title="human", + form_content="Awaiting human input", + inputs=[], + user_actions=[UserActionConfig(id="continue", title="Continue")], + ) + human_node = HumanInputNode( + node_id="human_input", + data=GraphonHumanInputNodeData.model_validate(human_data.model_dump()), + init_params=params, + runtime_state=runtime_state, + hitl_callback=DifyHITLCallback(form_repository=form_repository, node_data=human_data), + ) + + answer_false_node = AnswerNode( + node_id="answer_false", + data=AnswerNodeData(title="answer_false", answer="unreachable branch"), + init_params=params, + runtime_state=runtime_state, + ) + + answer_after_pause = AnswerNode( + node_id="answer_after_pause", + data=AnswerNodeData(title="answer_after_pause", answer="Post-branch answer chunk 1"), + init_params=params, + runtime_state=runtime_state, + ) + + answer_after_pause_2 = AnswerNode( + node_id="answer_after_pause_2", + data=AnswerNodeData(title="answer_after_pause_2", answer="Post-branch answer chunk 2"), + init_params=params, + runtime_state=runtime_state, + ) + + return ( + Graph.new() + .add_root(start_node) + .add_node(if_else_node, from_node_id="start") + .add_node(human_node, from_node_id="if_else", source_handle="true") + .add_node(answer_false_node, from_node_id="if_else", source_handle="false") + .add_node(answer_after_pause, from_node_id="human_input", source_handle="continue") + .add_node(answer_after_pause_2, from_node_id="answer_after_pause") + .build() + ) + + +def _build_runtime_state() -> RuntimeState: + variable_pool = VariablePool.from_bootstrap( + system_variables=build_system_variables( + workflow_execution_id=WORKFLOW_EXECUTION_ID, + app_id="app", + workflow_id="wf", + user_id="user", + ), + user_inputs={}, + conversation_variables=[], + ) + variable_pool.add(("start", "category"), "fruit") # drives the if-else "true" branch + return RuntimeState(workflow_id="test-workflow", variable_pool=variable_pool, start_at=0.0) + + +class TestResponseStreamFilterPauseResumeCases: + @staticmethod + def test_if_else_human_input_pause_resume_answer_chunks_survive_resume( + sqlite_session_factory: sessionmaker[Session], + mocker: MockerFixture, + ) -> None: + mocker.patch( + "core.repositories.human_input_repository.session_factory.create_session", + side_effect=sqlite_session_factory, + ) + repository = HumanInputFormRepositoryImpl( + tenant_id="tenant", app_id="app", workflow_execution_id=WORKFLOW_EXECUTION_ID + ) + # ---- Phase 1: run to GraphRunPausedEvent ---- + runtime_state_1 = _build_runtime_state() + graph_1 = _build_graph(runtime_state_1, repository) + engine_1 = Engine( + graph=graph_1, + runtime_state=runtime_state_1, + command_channel=InMemoryChannel(), + ) + filter_1 = ResponseStreamFilter() + phase1_events = list( + filter_engine_events( + engine_1.run(), + context=EngineEventFilterContext.from_engine(engine_1), + filters=[filter_1], + ) ) - # Verify file was converted and added - file_var = variable_pool.get(["file_node", "file"]) - assert file_var is not None - assert file_var.value.type == FileType.DOCUMENT - assert file_var.value.transfer_method == FileTransferMethod.REMOTE_URL + assert any(isinstance(e, GraphRunPausedEvent) for e in phase1_events) + phase1_chunks = [e for e in phase1_events if isinstance(e, NodeRunStreamChunkEvent)] + assert not any(e.node_id in ("answer_after_pause", "answer_after_pause_2") for e in phase1_chunks) - # Verify file list was converted and added - files_var = variable_pool.get(["file_node", "files"]) - assert files_var is not None - assert isinstance(files_var.value, list) - assert len(files_var.value) == 2 - assert all(isinstance(f, File) for f in files_var.value) - assert files_var.value[0].type == FileType.IMAGE - assert files_var.value[1].type == FileType.IMAGE - assert files_var.value[0].type == FileType.IMAGE - assert files_var.value[1].type == FileType.IMAGE + response_filter_snapshot = filter_1.dumps() + runtime_snapshot = runtime_state_1.dumps() - def test_mapping_user_inputs_missing_variable_error(self): - """Test that mapping raises error when required variable is missing.""" - variable_pool = VariablePool.from_bootstrap( - system_variables=default_system_variables(), - user_inputs={}, + with sqlite_session_factory.begin() as session: + form = session.scalar(select(HumanInputForm).where(HumanInputForm.workflow_run_id == WORKFLOW_EXECUTION_ID)) + assert form is not None + assert form.status == HumanInputFormStatus.WAITING + form.selected_action_id = "continue" + form.submitted_data = "{}" + form.submitted_at = naive_utc_now() + form.status = HumanInputFormStatus.SUBMITTED + + # ---- Phase 2: rebuild engine + filter from snapshots, resume to completion ---- + runtime_state_2 = RuntimeState.from_snapshot(runtime_snapshot) + graph_2 = _build_graph(runtime_state_2, repository) + engine_2 = Engine( + graph=graph_2, + runtime_state=runtime_state_2, + command_channel=InMemoryChannel(), ) + filter_2 = ResponseStreamFilter() + filter_2.loads(response_filter_snapshot) - # Define variable mapping - variable_mapping = { - "node1.required_input": ["node1", "required_input"], - } + phase2_events = list(iter_dify_graph_engine_events(engine_2, filter_2)) - # User inputs without required variable - user_inputs = { - "node1.other_input": "some value", - } + assert any(isinstance(e, GraphRunSucceededEvent) for e in phase2_events) - # Should raise ValueError for missing variable - with pytest.raises(ValueError, match="Variable key node1.required_input not found in user inputs"): + phase2_chunks = [e for e in phase2_events if isinstance(e, NodeRunStreamChunkEvent)] + answer_1_chunks = [e for e in phase2_chunks if e.node_id == "answer_after_pause"] + answer_2_chunks = [e for e in phase2_chunks if e.node_id == "answer_after_pause_2"] + + assert answer_1_chunks, "answer_after_pause produced no stream chunks after resume" + assert answer_2_chunks, "answer_after_pause_2 produced no stream chunks after resume" + + +class TestWorkflowEntryCases: + @staticmethod + @pytest.fixture(autouse=True) + def _mock_ssrf_head(monkeypatch: pytest.MonkeyPatch): + """Avoid any real network requests during tests. + + factories.file_factory.remote.get_remote_file_info() uses remote_fetcher.make_request + to inspect remote files. We stub it to return a minimal response object with + headers so filename/mime/size can be derived deterministically. + """ + + def fake_head(method, url, *args, **kwargs): + assert method == "HEAD" + # choose a content-type by file suffix for determinism + if url.endswith(".pdf"): + ctype = "application/pdf" + elif url.endswith(".jpg") or url.endswith(".jpeg"): + ctype = "image/jpeg" + elif url.endswith(".png"): + ctype = "image/png" + else: + ctype = "application/octet-stream" + filename = url.split("/")[-1] or "file.bin" + headers = { + "Content-Type": ctype, + "Content-Disposition": f'attachment; filename="{filename}"', + "Content-Length": "12345", + } + return SimpleNamespace(status_code=200, headers=headers) + + monkeypatch.setattr("factories.file_factory.remote.remote_fetcher.make_request", fake_head) + + class TestWorkflowEntry: + """Test WorkflowEntry class methods.""" + + @pytest.mark.parametrize("value", [0, False, "", [], {}]) + def test_qualified_input_preserves_falsy_values(self, value): + variable_pool = VariablePool() + + WorkflowEntry.mapping_user_inputs_to_variable_pool( + variable_mapping={"node.value": ["source", "value"]}, + user_inputs={"node.value": value, "value": "fallback"}, + variable_pool=variable_pool, + tenant_id="tenant", + ) + + variable = variable_pool.get(["source", "value"]) + assert variable is not None + assert variable.value == value + + @pytest.mark.parametrize("root_name", ["object", "structured_output"]) + @pytest.mark.parametrize("value", [0, False, "", [], {}]) + def test_nested_input_overrides_preserve_siblings(self, root_name, value): + variable_pool = VariablePool() + original = {"keep": "root sibling", "nested": {"leaf": "old", "keep": "nested sibling"}} + variable_pool.add(["source", root_name], original) + key = f"node.#source.{root_name}.nested.leaf#" + + WorkflowEntry.mapping_user_inputs_to_variable_pool( + variable_mapping={key: ["source", root_name, "nested", "leaf"]}, + user_inputs={key: value, key.split(".", 1)[1]: "fallback"}, + variable_pool=variable_pool, + tenant_id="tenant", + ) + + variable = variable_pool.get(["source", root_name]) + assert variable is not None + assert variable.value == {"keep": "root sibling", "nested": {"leaf": value, "keep": "nested sibling"}} + assert original["nested"]["leaf"] == "old" + + def test_nested_inputs_build_missing_root_and_preserve_other_explicit_inputs(self): + variable_pool = VariablePool() + + WorkflowEntry.mapping_user_inputs_to_variable_pool( + variable_mapping={ + "node.#source.object.leaf#": ["source", "object", "leaf"], + "node.#source.object.nested.other#": ["source", "object", "nested", "other"], + }, + user_inputs={"#source.object.leaf#": 0, "#source.object.nested.other#": False}, + variable_pool=variable_pool, + tenant_id="tenant", + ) + + variable = variable_pool.get(["source", "object"]) + assert variable is not None + assert variable.value == {"leaf": 0, "nested": {"other": False}} + + @pytest.mark.parametrize("descendant_first", [False, True]) + @pytest.mark.parametrize("value", ["Alice", 0, False, "", [], {}]) + def test_nested_overrides_take_precedence_over_explicit_ancestors_in_either_order( + self, descendant_first, value + ): + variable_pool = VariablePool() + mappings = [ + ("node.#source.object#", ["source", "object"]), + ("node.#source.object.name#", ["source", "object", "name"]), + ] + if descendant_first: + mappings.reverse() + root_input = {"other": "kept", "name": "old"} + + WorkflowEntry.mapping_user_inputs_to_variable_pool( + variable_mapping=dict(mappings), + user_inputs={"#source.object#": root_input, "#source.object.name#": value}, + variable_pool=variable_pool, + tenant_id="tenant", + ) + + variable = variable_pool.get(["source", "object"]) + assert variable is not None + assert variable.value == {"other": "kept", "name": value} + assert root_input == {"other": "kept", "name": "old"} + + @pytest.mark.parametrize( + ("data", "inputs", "expected"), + [ + ( + {"type": "end", "outputs": [{"variable": "result", "value_selector": ["source", "value"]}]}, + {"#source.value#": "end input"}, + "end input", + ), + ( + { + "type": "variable-aggregator", + "output_type": "string", + "variables": [["missing", "value"], ["source", "value"]], + }, + {"#source.value#": "second branch"}, + "second branch", + ), + ( + { + "type": "variable-assigner", + "output_type": "string", + "variables": [["source", "value"]], + }, + {"#source.value#": "legacy branch"}, + "legacy branch", + ), + ( + { + "type": "list-operator", + "variable": ["source", "value"], + "filter_by": {"enabled": False}, + "order_by": {"enabled": False}, + "limit": {"enabled": False}, + }, + {"#source.value#": [3, 1, 2]}, + [3, 1, 2], + ), + ], + ) + def test_single_step_loads_explicit_inputs_for_nodes_without_graphon_extractors(self, data, inputs, expected): + workflow = Workflow.new( + tenant_id="tenant", + app_id="app", + type=WorkflowType.WORKFLOW, + version=Workflow.VERSION_DRAFT, + graph=json.dumps({"nodes": [{"id": "node", "data": {"title": "Node", **data}}], "edges": []}), + features="{}", + created_by="account", + environment_variables=[], + conversation_variables=[], + rag_pipeline_variables=[], + ) + workflow.id = "workflow" + variable_pool = VariablePool() + + WorkflowEntry.single_step_run( + workflow=workflow, + node_id="node", + user_id="user", + user_inputs=inputs, + variable_pool=variable_pool, + execution_driver=None, + ) + + variable = variable_pool.get(["source", "value"]) + assert variable is not None + assert variable.value == expected + + def test_mapping_user_inputs_to_variable_pool_with_system_variables(self): + """Test mapping system variables from user inputs to variable pool.""" + # Initialize variable pool with system variables + variable_pool = VariablePool.from_bootstrap( + system_variables=build_system_variables( + user_id="test_user_id", + app_id="test_app_id", + workflow_id="test_workflow_id", + ), + user_inputs={}, + ) + + # Define variable mapping - sys variables mapped to other nodes + variable_mapping = { + "node1.input1": ["node1", "input1"], # Regular mapping + "node2.query": ["node2", "query"], # Regular mapping + "sys.user_id": ["output_node", "user"], # System variable mapping + } + + # User inputs including sys variables + user_inputs = { + "node1.input1": "new_user_id", + "node2.query": "test query", + "sys.user_id": "system_user", + } + + # Execute mapping WorkflowEntry.mapping_user_inputs_to_variable_pool( variable_mapping=variable_mapping, user_inputs=user_inputs, @@ -511,182 +545,1346 @@ class TestWorkflowEntry: tenant_id="test_tenant", ) - def test_mapping_user_inputs_with_alternative_key_format(self): - """Test mapping with alternative key format (without node prefix).""" - variable_pool = VariablePool.from_bootstrap( - system_variables=default_system_variables(), - user_inputs={}, - ) + # Verify variables were added to pool + # Note: variable_pool.get returns Variable objects, not raw values + node1_var = variable_pool.get(["node1", "input1"]) + assert node1_var is not None + assert node1_var.value == "new_user_id" - # Define variable mapping - variable_mapping = { - "node1.input": ["node1", "input"], - } + node2_var = variable_pool.get(["node2", "query"]) + assert node2_var is not None + assert node2_var.value == "test query" - # User inputs with alternative key format - user_inputs = { - "input": "value without node prefix", # Alternative format without node prefix - } + # System variable gets mapped to output node + output_var = variable_pool.get(["output_node", "user"]) + assert output_var is not None + assert output_var.value == "system_user" - # Execute mapping - WorkflowEntry.mapping_user_inputs_to_variable_pool( - variable_mapping=variable_mapping, - user_inputs=user_inputs, - variable_pool=variable_pool, - tenant_id="test_tenant", - ) + def test_single_step_run_injects_code_limits(self): + """Ensure single-step CodeNode execution configures limits.""" + # Arrange + node_id = "code_node" + node_data = { + "type": "code", + "title": "Code", + "desc": None, + "variables": [], + "code_language": CodeLanguage.PYTHON3, + "code": "def main():\n return {}", + "outputs": {}, + } + node_config = {"id": node_id, "data": node_data} - # Verify variable was added using alternative key - input_var = variable_pool.get(["node1", "input"]) - assert input_var is not None - assert input_var.value == "value without node prefix" + workflow = Workflow.new( + tenant_id="tenant", + app_id="app", + type=WorkflowType.WORKFLOW, + version=Workflow.VERSION_DRAFT, + graph=json.dumps({"nodes": [node_config], "edges": []}), + features="{}", + created_by="account", + environment_variables=[], + conversation_variables=[], + rag_pipeline_variables=[], + ) + workflow.id = "workflow" + variable_pool = VariablePool.from_bootstrap(system_variables=default_system_variables(), user_inputs={}) + expected_limits = CodeNodeLimits( + max_string_length=dify_config.CODE_MAX_STRING_LENGTH, + max_number=dify_config.CODE_MAX_NUMBER, + min_number=dify_config.CODE_MIN_NUMBER, + max_precision=dify_config.CODE_MAX_PRECISION, + max_depth=dify_config.CODE_MAX_DEPTH, + max_number_array_length=dify_config.CODE_MAX_NUMBER_ARRAY_LENGTH, + max_string_array_length=dify_config.CODE_MAX_STRING_ARRAY_LENGTH, + max_object_array_length=dify_config.CODE_MAX_OBJECT_ARRAY_LENGTH, + ) - def test_mapping_user_inputs_with_complex_selectors(self): - """Test mapping with complex node variable keys.""" - variable_pool = VariablePool.from_bootstrap( - system_variables=default_system_variables(), - user_inputs={}, - ) + # Act + node, _ = WorkflowEntry.single_step_run( + workflow=workflow, + node_id=node_id, + user_id="user", + user_inputs={}, + variable_pool=variable_pool, + execution_driver=None, + ) - # Define variable mapping - selectors can only have 2 elements - variable_mapping = { - "node1.data.field1": ["node1", "data_field1"], # Complex key mapped to simple selector - "node2.config.settings.timeout": ["node2", "timeout"], # Complex key mapped to simple selector - } + # Assert + assert isinstance(node, CodeNode) + assert node._limits == expected_limits - # User inputs - user_inputs = { - "node1.data.field1": "nested value", - "node2.config.settings.timeout": 30, - } + def test_mapping_user_inputs_to_variable_pool_with_env_variables(self): + """Test mapping environment variables from user inputs to variable pool.""" + # Initialize variable pool with environment variables + env_var = StringVariable(name="API_KEY", value="existing_key") + variable_pool = VariablePool.from_bootstrap( + system_variables=default_system_variables(), + environment_variables=[env_var], + user_inputs={}, + ) - # Execute mapping - WorkflowEntry.mapping_user_inputs_to_variable_pool( - variable_mapping=variable_mapping, - user_inputs=user_inputs, - variable_pool=variable_pool, - tenant_id="test_tenant", - ) + # Add env variable to pool (simulating initialization) + variable_pool.add([ENVIRONMENT_VARIABLE_NODE_ID, "API_KEY"], env_var) - # Verify variables were added with simple selectors - data_var = variable_pool.get(["node1", "data_field1"]) - assert data_var is not None - assert data_var.value == "nested value" + # Define variable mapping - env variables should not be overridden + variable_mapping = { + "node1.api_key": [ENVIRONMENT_VARIABLE_NODE_ID, "API_KEY"], + "node2.new_env": [ENVIRONMENT_VARIABLE_NODE_ID, "NEW_ENV"], + } - timeout_var = variable_pool.get(["node2", "timeout"]) - assert timeout_var is not None - assert timeout_var.value == 30 + # User inputs + user_inputs = { + "node1.api_key": "user_provided_key", # This should not override existing env var + "node2.new_env": "new_env_value", + } - def test_mapping_user_inputs_invalid_node_variable(self): - """Test that mapping handles invalid node variable format.""" - variable_pool = VariablePool.from_bootstrap( - system_variables=default_system_variables(), - user_inputs={}, - ) + # Execute mapping + WorkflowEntry.mapping_user_inputs_to_variable_pool( + variable_mapping=variable_mapping, + user_inputs=user_inputs, + variable_pool=variable_pool, + tenant_id="test_tenant", + ) - # Define variable mapping with single element node variable (at least one dot is required) - variable_mapping = { - "singleelement": ["node1", "input"], # No dot separator - } + # Verify env variable was not overridden + env_value = variable_pool.get([ENVIRONMENT_VARIABLE_NODE_ID, "API_KEY"]) + assert env_value is not None + assert env_value.value == "existing_key" # Should remain unchanged - user_inputs = {"singleelement": "some value"} # Must use exact key + # New env variables from user input should not be added + assert variable_pool.get([ENVIRONMENT_VARIABLE_NODE_ID, "NEW_ENV"]) is None - # Should NOT raise error - function accepts it and uses direct key - WorkflowEntry.mapping_user_inputs_to_variable_pool( - variable_mapping=variable_mapping, - user_inputs=user_inputs, - variable_pool=variable_pool, - tenant_id="test_tenant", - ) + def test_mapping_user_inputs_to_variable_pool_with_conversation_variables(self): + """Test mapping conversation variables from user inputs to variable pool.""" + # Initialize variable pool with conversation variables + conv_var = StringVariable(name="last_message", value="Hello") + variable_pool = VariablePool.from_bootstrap( + system_variables=default_system_variables(), + conversation_variables=[conv_var], + user_inputs={}, + ) - # Verify it was added - var = variable_pool.get(["node1", "input"]) - assert var is not None - assert var.value == "some value" + # Add conversation variable to pool + variable_pool.add([CONVERSATION_VARIABLE_NODE_ID, "last_message"], conv_var) - def test_mapping_all_variable_types_together(self): - """Test mapping all four types of variables in one operation.""" - # Initialize variable pool with some existing variables - env_var = StringVariable(name="API_KEY", value="existing_key") - conv_var = StringVariable(name="session_id", value="session123") + # Define variable mapping + variable_mapping = { + "node1.message": ["node1", "message"], # Map to regular node + "conversation.context": ["chat_node", "context"], # Conversation var to regular node + } - variable_pool = VariablePool.from_bootstrap( - system_variables=build_system_variables( - user_id="test_user", - app_id="test_app", - query="initial query", - ), - environment_variables=[env_var], - conversation_variables=[conv_var], - user_inputs={}, - ) + # User inputs + user_inputs = { + "node1.message": "Updated message", + "conversation.context": "New context", + } - # Add existing variables to pool - variable_pool.add([ENVIRONMENT_VARIABLE_NODE_ID, "API_KEY"], env_var) - variable_pool.add([CONVERSATION_VARIABLE_NODE_ID, "session_id"], conv_var) + # Execute mapping + WorkflowEntry.mapping_user_inputs_to_variable_pool( + variable_mapping=variable_mapping, + user_inputs=user_inputs, + variable_pool=variable_pool, + tenant_id="test_tenant", + ) - # Define comprehensive variable mapping - variable_mapping = { - # System variables mapped to regular nodes - "sys.user_id": ["start", "user"], - "sys.app_id": ["start", "app"], - # Environment variables (won't be overridden) - "env.API_KEY": ["config", "api_key"], - # Conversation variables mapped to regular nodes - "conversation.session_id": ["chat", "session"], - # Regular variables - "input.text": ["input", "text"], - "process.data": ["process", "data"], - } + # Verify variables were added to their target nodes + node1_var = variable_pool.get(["node1", "message"]) + assert node1_var is not None + assert node1_var.value == "Updated message" - # User inputs - user_inputs = { - "sys.user_id": "new_user", - "sys.app_id": "new_app", - "env.API_KEY": "attempted_override", # Should not override env var - "conversation.session_id": "new_session", - "input.text": "user input text", - "process.data": {"value": 123, "status": "active"}, - } + chat_var = variable_pool.get(["chat_node", "context"]) + assert chat_var is not None + assert chat_var.value == "New context" - # Execute mapping - WorkflowEntry.mapping_user_inputs_to_variable_pool( - variable_mapping=variable_mapping, - user_inputs=user_inputs, - variable_pool=variable_pool, - tenant_id="test_tenant", - ) + def test_mapping_user_inputs_to_variable_pool_with_regular_variables(self): + """Test mapping regular node variables from user inputs to variable pool.""" + # Initialize empty variable pool + variable_pool = VariablePool.from_bootstrap( + system_variables=default_system_variables(), + user_inputs={}, + ) - # Verify system variables were added to their target nodes - start_user = variable_pool.get(["start", "user"]) - assert start_user is not None - assert start_user.value == "new_user" + # Define variable mapping for regular nodes + variable_mapping = { + "input_node.text": ["input_node", "text"], + "llm_node.prompt": ["llm_node", "prompt"], + "code_node.input": ["code_node", "input"], + } - start_app = variable_pool.get(["start", "app"]) - assert start_app is not None - assert start_app.value == "new_app" + # User inputs + user_inputs = { + "input_node.text": "User input text", + "llm_node.prompt": "Generate a summary", + "code_node.input": {"key": "value"}, + } - # Verify env variable was not overridden (still has original value) - env_value = variable_pool.get([ENVIRONMENT_VARIABLE_NODE_ID, "API_KEY"]) - assert env_value is not None - assert env_value.value == "existing_key" + # Execute mapping + WorkflowEntry.mapping_user_inputs_to_variable_pool( + variable_mapping=variable_mapping, + user_inputs=user_inputs, + variable_pool=variable_pool, + tenant_id="test_tenant", + ) - # Environment variables get mapped to other nodes even when they exist in env pool - # But the original env value remains unchanged - config_api_key = variable_pool.get(["config", "api_key"]) - assert config_api_key is not None - assert config_api_key.value == "attempted_override" + # Verify regular variables were added + text_var = variable_pool.get(["input_node", "text"]) + assert text_var is not None + assert text_var.value == "User input text" - # Verify conversation variable was mapped to target node - chat_session = variable_pool.get(["chat", "session"]) - assert chat_session is not None - assert chat_session.value == "new_session" + prompt_var = variable_pool.get(["llm_node", "prompt"]) + assert prompt_var is not None + assert prompt_var.value == "Generate a summary" - # Verify regular variables were added - input_text = variable_pool.get(["input", "text"]) - assert input_text is not None - assert input_text.value == "user input text" + input_var = variable_pool.get(["code_node", "input"]) + assert input_var is not None + assert input_var.value == {"key": "value"} - process_data = variable_pool.get(["process", "data"]) - assert process_data is not None - assert process_data.value == {"value": 123, "status": "active"} + def test_mapping_user_inputs_with_file_handling(self): + """Test mapping file inputs from user inputs to variable pool.""" + variable_pool = VariablePool.from_bootstrap( + system_variables=default_system_variables(), + user_inputs={}, + ) + + # Define variable mapping + variable_mapping = { + "file_node.file": ["file_node", "file"], + "file_node.files": ["file_node", "files"], + } + + # User inputs with file data - using remote_url which doesn't require upload_file_id + user_inputs = { + "file_node.file": { + "type": "document", + "transfer_method": "remote_url", + "url": "http://example.com/test.pdf", + }, + "file_node.files": [ + { + "type": "image", + "transfer_method": "remote_url", + "url": "http://example.com/image1.jpg", + }, + { + "type": "image", + "transfer_method": "remote_url", + "url": "http://example.com/image2.jpg", + }, + ], + } + + # Execute mapping + WorkflowEntry.mapping_user_inputs_to_variable_pool( + variable_mapping=variable_mapping, + user_inputs=user_inputs, + variable_pool=variable_pool, + tenant_id="test_tenant", + ) + + # Verify file was converted and added + file_var = variable_pool.get(["file_node", "file"]) + assert file_var is not None + assert file_var.value.type == FileType.DOCUMENT + assert file_var.value.transfer_method == FileTransferMethod.REMOTE_URL + + # Verify file list was converted and added + files_var = variable_pool.get(["file_node", "files"]) + assert files_var is not None + assert isinstance(files_var.value, list) + assert len(files_var.value) == 2 + assert all(isinstance(f, File) for f in files_var.value) + assert files_var.value[0].type == FileType.IMAGE + assert files_var.value[1].type == FileType.IMAGE + assert files_var.value[0].type == FileType.IMAGE + assert files_var.value[1].type == FileType.IMAGE + + def test_mapping_user_inputs_missing_variable_error(self): + """Test that mapping raises error when required variable is missing.""" + variable_pool = VariablePool.from_bootstrap( + system_variables=default_system_variables(), + user_inputs={}, + ) + + # Define variable mapping + variable_mapping = { + "node1.required_input": ["node1", "required_input"], + } + + # User inputs without required variable + user_inputs = { + "node1.other_input": "some value", + } + + # Should raise ValueError for missing variable + with pytest.raises(ValueError, match="Variable key node1.required_input not found in user inputs"): + WorkflowEntry.mapping_user_inputs_to_variable_pool( + variable_mapping=variable_mapping, + user_inputs=user_inputs, + variable_pool=variable_pool, + tenant_id="test_tenant", + ) + + def test_mapping_user_inputs_with_alternative_key_format(self): + """Test mapping with alternative key format (without node prefix).""" + variable_pool = VariablePool.from_bootstrap( + system_variables=default_system_variables(), + user_inputs={}, + ) + + # Define variable mapping + variable_mapping = { + "node1.input": ["node1", "input"], + } + + # User inputs with alternative key format + user_inputs = { + "input": "value without node prefix", # Alternative format without node prefix + } + + # Execute mapping + WorkflowEntry.mapping_user_inputs_to_variable_pool( + variable_mapping=variable_mapping, + user_inputs=user_inputs, + variable_pool=variable_pool, + tenant_id="test_tenant", + ) + + # Verify variable was added using alternative key + input_var = variable_pool.get(["node1", "input"]) + assert input_var is not None + assert input_var.value == "value without node prefix" + + def test_mapping_user_inputs_with_complex_selectors(self): + """Test mapping with complex node variable keys.""" + variable_pool = VariablePool.from_bootstrap( + system_variables=default_system_variables(), + user_inputs={}, + ) + + # Define variable mapping - selectors can only have 2 elements + variable_mapping = { + "node1.data.field1": ["node1", "data_field1"], # Complex key mapped to simple selector + "node2.config.settings.timeout": ["node2", "timeout"], # Complex key mapped to simple selector + } + + # User inputs + user_inputs = { + "node1.data.field1": "nested value", + "node2.config.settings.timeout": 30, + } + + # Execute mapping + WorkflowEntry.mapping_user_inputs_to_variable_pool( + variable_mapping=variable_mapping, + user_inputs=user_inputs, + variable_pool=variable_pool, + tenant_id="test_tenant", + ) + + # Verify variables were added with simple selectors + data_var = variable_pool.get(["node1", "data_field1"]) + assert data_var is not None + assert data_var.value == "nested value" + + timeout_var = variable_pool.get(["node2", "timeout"]) + assert timeout_var is not None + assert timeout_var.value == 30 + + def test_mapping_user_inputs_invalid_node_variable(self): + """Test that mapping handles invalid node variable format.""" + variable_pool = VariablePool.from_bootstrap( + system_variables=default_system_variables(), + user_inputs={}, + ) + + # Define variable mapping with single element node variable (at least one dot is required) + variable_mapping = { + "singleelement": ["node1", "input"], # No dot separator + } + + user_inputs = {"singleelement": "some value"} # Must use exact key + + # Should NOT raise error - function accepts it and uses direct key + WorkflowEntry.mapping_user_inputs_to_variable_pool( + variable_mapping=variable_mapping, + user_inputs=user_inputs, + variable_pool=variable_pool, + tenant_id="test_tenant", + ) + + # Verify it was added + var = variable_pool.get(["node1", "input"]) + assert var is not None + assert var.value == "some value" + + def test_mapping_all_variable_types_together(self): + """Test mapping all four types of variables in one operation.""" + # Initialize variable pool with some existing variables + env_var = StringVariable(name="API_KEY", value="existing_key") + conv_var = StringVariable(name="session_id", value="session123") + + variable_pool = VariablePool.from_bootstrap( + system_variables=build_system_variables( + user_id="test_user", + app_id="test_app", + query="initial query", + ), + environment_variables=[env_var], + conversation_variables=[conv_var], + user_inputs={}, + ) + + # Add existing variables to pool + variable_pool.add([ENVIRONMENT_VARIABLE_NODE_ID, "API_KEY"], env_var) + variable_pool.add([CONVERSATION_VARIABLE_NODE_ID, "session_id"], conv_var) + + # Define comprehensive variable mapping + variable_mapping = { + # System variables mapped to regular nodes + "sys.user_id": ["start", "user"], + "sys.app_id": ["start", "app"], + # Environment variables (won't be overridden) + "env.API_KEY": ["config", "api_key"], + # Conversation variables mapped to regular nodes + "conversation.session_id": ["chat", "session"], + # Regular variables + "input.text": ["input", "text"], + "process.data": ["process", "data"], + } + + # User inputs + user_inputs = { + "sys.user_id": "new_user", + "sys.app_id": "new_app", + "env.API_KEY": "attempted_override", # Should not override env var + "conversation.session_id": "new_session", + "input.text": "user input text", + "process.data": {"value": 123, "status": "active"}, + } + + # Execute mapping + WorkflowEntry.mapping_user_inputs_to_variable_pool( + variable_mapping=variable_mapping, + user_inputs=user_inputs, + variable_pool=variable_pool, + tenant_id="test_tenant", + ) + + # Verify system variables were added to their target nodes + start_user = variable_pool.get(["start", "user"]) + assert start_user is not None + assert start_user.value == "new_user" + + start_app = variable_pool.get(["start", "app"]) + assert start_app is not None + assert start_app.value == "new_app" + + # Verify env variable was not overridden (still has original value) + env_value = variable_pool.get([ENVIRONMENT_VARIABLE_NODE_ID, "API_KEY"]) + assert env_value is not None + assert env_value.value == "existing_key" + + # Environment variables get mapped to other nodes even when they exist in env pool + # But the original env value remains unchanged + config_api_key = variable_pool.get(["config", "api_key"]) + assert config_api_key is not None + assert config_api_key.value == "attempted_override" + + # Verify conversation variable was mapped to target node + chat_session = variable_pool.get(["chat", "session"]) + assert chat_session is not None + assert chat_session.value == "new_session" + + # Verify regular variables were added + input_text = variable_pool.get(["input", "text"]) + assert input_text is not None + assert input_text.value == "user input text" + + process_data = variable_pool.get(["process", "data"]) + assert process_data is not None + assert process_data.value == {"value": 123, "status": "active"} + + +def _build_typed_node_config(node_type: NodeType): + return {"id": "node-id", "data": BaseNodeData(type=node_type)} + + +def _workflow() -> Workflow: + """Build a real transient workflow for single-step orchestration tests.""" + return make_workflow(workflow_id="workflow-id", tenant_id="tenant-id", app_id="app-id", created_by="user-id") + + +def _build_minimal_workflow_entry( + monkeypatch: pytest.MonkeyPatch, + *, + response_stream_filter: ResponseStreamFilter | None = None, +) -> workflow_entry.WorkflowEntry: + """Construct a minimal WorkflowEntry with Engine construction mocked out.""" + graph_engine = MagicMock() + graph_runtime_state = sentinel.runtime_state + + monkeypatch.setattr(workflow_entry, "ExecutionContextLayer", MagicMock()) + monkeypatch.setattr(workflow_entry, "Engine", MagicMock(return_value=graph_engine)) + monkeypatch.setattr(workflow_entry, "InMemoryChannel", MagicMock(return_value=sentinel.command_channel)) + + return workflow_entry.WorkflowEntry( + graph_config={"nodes": [], "edges": []}, + graph=sentinel.graph, + call_depth=0, + graph_runtime_state=graph_runtime_state, + response_stream_filter=response_stream_filter, + ) + + +class TestWorkflowEntryHelpersCases: + class TestWorkflowEntryInit: + def test_rejects_call_depth_above_limit(self): + call_depth = workflow_entry.dify_config.WORKFLOW_CALL_MAX_DEPTH + 1 + + with pytest.raises(ValueError, match="Max workflow call depth"): + workflow_entry.WorkflowEntry( + graph_config={"nodes": [], "edges": []}, + graph=sentinel.graph, + call_depth=call_depth, + graph_runtime_state=sentinel.graph_runtime_state, + ) + + def test_hidden_workflow_tool_nodes_count_toward_execution_limit(self): + tool, _, _ = _workflow_tool_node() + source = _workflow_tool_source() + with config_overrides_context(WORKFLOW_MAX_EXECUTION_STEPS=2, GRAPH_ENGINE_MAX_WORKERS=1): + entry = workflow_entry.WorkflowEntry( + graph_config=tool.init_params.graph_config, + graph=_outer_graph(tool), + call_depth=0, + graph_runtime_state=tool.runtime_state, + ) + entry.workflow_tool_sources[source.workflow_id] = source + events = list(entry.run()) + + assert isinstance(events[-1], GraphRunAbortedEvent) + assert "Maximum execution steps exceeded: 3 > 2" in events[-1].reason + assert not any(isinstance(event, NodeRunSucceededEvent) and event.node_id == "tool" for event in events) + + def test_workflow_entry_stores_supplied_response_stream_filter(self, monkeypatch: pytest.MonkeyPatch) -> None: + supplied_filter = ResponseStreamFilter() + entry = _build_minimal_workflow_entry(monkeypatch, response_stream_filter=supplied_filter) + + assert entry._response_stream_filter is supplied_filter + + def test_workflow_entry_defaults_to_fresh_response_stream_filter(self, monkeypatch: pytest.MonkeyPatch) -> None: + entry = _build_minimal_workflow_entry(monkeypatch, response_stream_filter=None) + + assert isinstance(entry._response_stream_filter, ResponseStreamFilter) + + class TestWorkflowEntryRun: + def test_run_swallows_generate_task_stopped_errors(self): + entry = object.__new__(workflow_entry.WorkflowEntry) + entry.graph_engine = MagicMock() + entry.graph_engine.run.side_effect = GenerateTaskStoppedError() + entry._response_stream_filter = ResponseStreamFilter() + + assert list(entry.run()) == [] + + def test_iter_dify_graph_engine_events_applies_response_stream_filter(self): + graph_engine = MagicMock() + graph_engine.run.return_value = iter([sentinel.raw_event]) + + with ( + patch.object( + workflow_entry.EngineEventFilterContext, + "from_engine", + return_value=sentinel.filter_context, + ) as from_engine, + patch.object( + workflow_entry, + "ResponseStreamFilter", + return_value=sentinel.response_stream_filter, + ) as response_stream_filter_cls, + patch.object( + workflow_entry, + "HumanInputFormEventFilter", + return_value=sentinel.human_input_filter, + ), + patch.object( + workflow_entry, + "filter_engine_events", + return_value=iter([sentinel.filtered_event]), + ) as filter_engine_events, + ): + events = list(workflow_entry.iter_dify_graph_engine_events(graph_engine)) + + assert events == [sentinel.filtered_event] + from_engine.assert_called_once_with(graph_engine) + response_stream_filter_cls.assert_called_once_with() + filter_engine_events.assert_called_once_with( + graph_engine.run.return_value, + context=sentinel.filter_context, + filters=[sentinel.human_input_filter, sentinel.response_stream_filter], + ) + + def test_run_delegates_to_dify_event_iterator(self): + entry = object.__new__(workflow_entry.WorkflowEntry) + entry.graph_engine = sentinel.graph_engine + entry._response_stream_filter = sentinel.response_stream_filter + + with patch.object( + workflow_entry, + "iter_dify_graph_engine_events", + return_value=iter([sentinel.filtered_event]), + ) as iter_dify_graph_engine_events: + events = list(entry.run()) + + assert events == [sentinel.filtered_event] + iter_dify_graph_engine_events.assert_called_once_with( + sentinel.graph_engine, sentinel.response_stream_filter + ) + + def test_run_emits_failed_event_for_unexpected_errors(self): + entry = object.__new__(workflow_entry.WorkflowEntry) + entry.graph_engine = MagicMock() + entry.graph_engine.run.side_effect = RuntimeError("boom") + entry._response_stream_filter = ResponseStreamFilter() + + events = list(entry.run()) + + assert len(events) == 1 + assert isinstance(events[0], GraphRunFailedEvent) + assert events[0].error == "boom" + + class TestWorkflowEntrySingleStepRun: + @pytest.mark.parametrize("node_type", [BuiltinNodeTypes.LOOP, BuiltinNodeTypes.ITERATION]) + def test_rejects_container_nodes(self, node_type): + workflow = _workflow() + + with ( + patch.object(workflow, "get_node_config_by_id", return_value=_build_typed_node_config(node_type)), + pytest.raises(ValueError, match="engine-backed debug endpoints"), + ): + workflow_entry.WorkflowEntry.single_step_run( + workflow=workflow, + node_id="node-id", + user_id="user-id", + user_inputs={}, + variable_pool=sentinel.variable_pool, + execution_driver=None, + ) + + def test_preloads_constructor_variables_before_creating_memory_node(self): + class FakeLLMNode: + id = "node-id" + title = "Node Title" + node_type = BuiltinNodeTypes.LLM + + @staticmethod + def version(): + return "1" + + @staticmethod + def extract_variable_selector_to_variable_mapping(**_kwargs): + return {} + + variable_pool = VariablePool.from_bootstrap(system_variables=default_system_variables(), user_inputs={}) + variable_loader = MagicMock() + variable_loader.load_variables.return_value = [ + StringVariable( + name="conversation_id", + value="conv-1", + selector=["sys", "conversation_id"], + ) + ] + workflow = _workflow() + node_config = { + "id": "node-id", + "data": BaseNodeData(type=BuiltinNodeTypes.LLM, version="1", memory=object()), + } + + with ( + patch.object(workflow, "get_node_config_by_id", return_value=node_config), + patch.object(workflow_entry, "DifyGraphInitContext", return_value=sentinel.graph_init_context), + patch.object( + workflow_entry, + "RuntimeState", + return_value=SimpleNamespace(variable_pool=variable_pool), + ), + patch.object(workflow_entry, "build_dify_run_context", return_value={"_dify": "context"}), + patch.object(workflow_entry.time, "perf_counter", return_value=123.0), + patch.object(workflow_entry, "resolve_workflow_node_class", return_value=FakeLLMNode), + patch.object(workflow_entry.DifyNodeFactory, "from_graph_init_context") as dify_node_factory, + patch.object(workflow_entry, "load_into_variable_pool"), + patch.object(workflow_entry.WorkflowEntry, "mapping_user_inputs_to_variable_pool"), + patch.object( + workflow_entry.WorkflowEntry, + "_run_node_with_layers", + return_value=iter(["event"]), + ), + ): + + def _create_node(_node_config): + assert variable_pool.get(["sys", "conversation_id"]) is not None + return FakeLLMNode() + + dify_node_factory.return_value.create_node.side_effect = _create_node + node, generator = workflow_entry.WorkflowEntry.single_step_run( + workflow=workflow, + node_id="node-id", + user_id="user-id", + user_inputs={}, + variable_pool=variable_pool, + variable_loader=variable_loader, + execution_driver=None, + ) + + assert node.id == "node-id" + assert list(generator) == ["event"] + variable_loader.load_variables.assert_called_once_with([["sys", "conversation_id"]]) + + def test_uses_empty_mapping_when_selector_extraction_is_not_implemented(self): + class FakeNode: + id = "node-id" + title = "Node Title" + node_type = "fake" + + @staticmethod + def version(): + return "1" + + @staticmethod + def extract_variable_selector_to_variable_mapping(**_kwargs): + raise NotImplementedError + + workflow = _workflow() + with ( + patch.object( + workflow, + "get_node_config_by_id", + return_value=_build_typed_node_config(BuiltinNodeTypes.START), + ), + patch.object(workflow_entry, "DifyGraphInitContext", return_value=sentinel.graph_init_context), + patch.object(workflow_entry, "RuntimeState", return_value=sentinel.graph_runtime_state), + patch.object(workflow_entry, "build_dify_run_context", return_value={"_dify": "context"}), + patch.object(workflow_entry.time, "perf_counter", return_value=123.0), + patch.object(workflow_entry, "resolve_workflow_node_class", return_value=FakeNode), + patch.object(workflow_entry.DifyNodeFactory, "from_graph_init_context") as dify_node_factory, + patch.object(workflow_entry, "add_node_inputs_to_pool") as add_node_inputs_to_pool, + patch.object(workflow_entry, "load_into_variable_pool") as load_into_variable_pool, + patch.object( + workflow_entry.WorkflowEntry, + "mapping_user_inputs_to_variable_pool", + ) as mapping_user_inputs_to_variable_pool, + patch.object( + workflow_entry.WorkflowEntry, + "_run_node_with_layers", + return_value=iter(["event"]), + ), + ): + dify_node_factory.return_value.create_node.return_value = FakeNode() + node, generator = workflow_entry.WorkflowEntry.single_step_run( + workflow=workflow, + node_id="node-id", + user_id="user-id", + user_inputs={"question": "hello"}, + variable_pool=sentinel.variable_pool, + execution_driver=None, + ) + + assert node.id == "node-id" + assert list(generator) == ["event"] + load_into_variable_pool.assert_called_once_with( + variable_loader=workflow_entry.DUMMY_VARIABLE_LOADER, + variable_pool=sentinel.variable_pool, + variable_mapping={}, + user_inputs={"question": "hello"}, + ) + add_node_inputs_to_pool.assert_called_once_with( + sentinel.variable_pool, + node_id="node-id", + inputs={"question": "hello"}, + ) + mapping_user_inputs_to_variable_pool.assert_called_once_with( + variable_mapping={}, + user_inputs={"question": "hello"}, + variable_pool=sentinel.variable_pool, + tenant_id="tenant-id", + ) + + def test_skips_user_input_mapping_for_datasource_nodes(self): + class FakeDatasourceNode: + id = "node-id" + node_type = "datasource" + + @staticmethod + def version(): + return "1" + + @staticmethod + def extract_variable_selector_to_variable_mapping(**_kwargs): + return {"question": ["node", "question"]} + + workflow = _workflow() + with ( + patch.object( + workflow, + "get_node_config_by_id", + return_value=_build_typed_node_config(BuiltinNodeTypes.DATASOURCE), + ), + patch.object(workflow_entry, "DifyGraphInitContext", return_value=sentinel.graph_init_context), + patch.object(workflow_entry, "RuntimeState", return_value=sentinel.graph_runtime_state), + patch.object(workflow_entry, "build_dify_run_context", return_value={"_dify": "context"}), + patch.object(workflow_entry.time, "perf_counter", return_value=123.0), + patch.object(workflow_entry, "resolve_workflow_node_class", return_value=FakeDatasourceNode), + patch.object(workflow_entry.DifyNodeFactory, "from_graph_init_context") as dify_node_factory, + patch.object(workflow_entry, "add_node_inputs_to_pool") as add_node_inputs_to_pool, + patch.object(workflow_entry, "load_into_variable_pool") as load_into_variable_pool, + patch.object( + workflow_entry.WorkflowEntry, + "mapping_user_inputs_to_variable_pool", + ) as mapping_user_inputs_to_variable_pool, + patch.object( + workflow_entry.WorkflowEntry, + "_run_node_with_layers", + return_value=iter(["event"]), + ), + ): + dify_node_factory.return_value.create_node.return_value = FakeDatasourceNode() + node, generator = workflow_entry.WorkflowEntry.single_step_run( + workflow=workflow, + node_id="node-id", + user_id="user-id", + user_inputs={"question": "hello"}, + variable_pool=sentinel.variable_pool, + execution_driver=None, + ) + + assert node.id == "node-id" + assert list(generator) == ["event"] + load_into_variable_pool.assert_called_once() + add_node_inputs_to_pool.assert_called_once_with( + sentinel.variable_pool, + node_id="node-id", + inputs={"question": "hello"}, + ) + mapping_user_inputs_to_variable_pool.assert_not_called() + + def test_wraps_layered_node_run_failures(self): + class FakeNode: + id = "node-id" + title = "Node Title" + node_type = "fake" + + @staticmethod + def extract_variable_selector_to_variable_mapping(**_kwargs): + return {} + + @staticmethod + def version(): + return "1" + + workflow = _workflow() + with ( + patch.object( + workflow, + "get_node_config_by_id", + return_value=_build_typed_node_config(BuiltinNodeTypes.START), + ), + patch.object(workflow_entry, "DifyGraphInitContext", return_value=sentinel.graph_init_context), + patch.object(workflow_entry, "RuntimeState", return_value=sentinel.graph_runtime_state), + patch.object(workflow_entry, "build_dify_run_context", return_value={"_dify": "context"}), + patch.object(workflow_entry.time, "perf_counter", return_value=123.0), + patch.object(workflow_entry, "resolve_workflow_node_class", return_value=FakeNode), + patch.object(workflow_entry.DifyNodeFactory, "from_graph_init_context") as dify_node_factory, + patch.object(workflow_entry, "add_node_inputs_to_pool"), + patch.object(workflow_entry, "load_into_variable_pool"), + patch.object(workflow_entry.WorkflowEntry, "mapping_user_inputs_to_variable_pool"), + patch.object( + workflow_entry.WorkflowEntry, + "_run_node_with_layers", + side_effect=RuntimeError("boom"), + ), + ): + dify_node_factory.return_value.create_node.return_value = FakeNode() + with pytest.raises(WorkflowNodeRunFailedError): + workflow_entry.WorkflowEntry.single_step_run( + workflow=workflow, + node_id="node-id", + user_id="user-id", + user_inputs={}, + variable_pool=sentinel.variable_pool, + execution_driver=None, + ) + + class TestWorkflowEntryHelpers: + def test_create_single_node_graph_builds_start_edge(self): + graph = workflow_entry.WorkflowEntry._create_single_node_graph( + node_id="target-node", + node_data={"type": BuiltinNodeTypes.PARAMETER_EXTRACTOR}, + node_width=320, + node_height=180, + ) + + assert graph["nodes"][0]["id"] == "start" + assert graph["nodes"][1]["id"] == "target-node" + assert graph["nodes"][1]["width"] == 320 + assert graph["nodes"][1]["height"] == 180 + assert graph["edges"] == [ + { + "source": "start", + "target": "target-node", + "sourceHandle": "source", + "targetHandle": "target", + } + ] + + def test_run_free_node_rejects_unsupported_types(self): + with pytest.raises(ValueError, match="Node type start not supported"): + workflow_entry.WorkflowEntry.run_free_node( + node_data={"type": BuiltinNodeTypes.START}, + node_id="node-id", + tenant_id="tenant-id", + user_id="user-id", + user_inputs={}, + ) + + def test_run_free_node_rejects_missing_node_class(self, monkeypatch: pytest.MonkeyPatch): + monkeypatch.setattr( + workflow_entry, + "resolve_workflow_node_class", + MagicMock(return_value=None), + ) + + with pytest.raises(ValueError, match="Node class not found for node type parameter-extractor"): + workflow_entry.WorkflowEntry.run_free_node( + node_data={"type": BuiltinNodeTypes.PARAMETER_EXTRACTOR}, + node_id="node-id", + tenant_id="tenant-id", + user_id="user-id", + user_inputs={}, + ) + + def test_run_free_node_uses_empty_mapping_when_selector_extraction_is_not_implemented( + self, monkeypatch: pytest.MonkeyPatch + ): + class FakeNodeClass: + @staticmethod + def extract_variable_selector_to_variable_mapping(**_kwargs): + raise NotImplementedError + + class FakeNode: + id = "node-id" + title = "Node Title" + node_type = "parameter-extractor" + + @staticmethod + def version(): + return "1" + + dify_node_factory = MagicMock() + dify_node_factory.create_node.return_value = FakeNode() + monkeypatch.setattr( + workflow_entry, + "resolve_workflow_node_class", + MagicMock(return_value=FakeNodeClass), + ) + + with ( + patch.object(workflow_entry, "default_system_variables", return_value=sentinel.system_variables), + patch.object(workflow_entry, "VariablePool", return_value=sentinel.variable_pool) as variable_pool_cls, + patch.object(workflow_entry, "add_variables_to_pool") as add_variables_to_pool, + patch.object( + workflow_entry, "DifyGraphInitContext", return_value=sentinel.graph_init_context + ) as graph_init_context_cls, + patch.object(workflow_entry, "RuntimeState", return_value=sentinel.graph_runtime_state), + patch.object( + workflow_entry, "build_dify_run_context", return_value={"_dify": "context"} + ) as build_dify_run_context, + patch.object(workflow_entry.time, "perf_counter", return_value=123.0), + patch.object( + workflow_entry.DifyNodeFactory, + "from_graph_init_context", + return_value=dify_node_factory, + ) as dify_node_factory_cls, + patch.object( + workflow_entry.WorkflowEntry, + "mapping_user_inputs_to_variable_pool", + ) as mapping_user_inputs_to_variable_pool, + patch.object( + workflow_entry.WorkflowEntry, + "_run_node_with_layers", + return_value=iter(["event"]), + ), + ): + node, generator = workflow_entry.WorkflowEntry.run_free_node( + node_data={"type": BuiltinNodeTypes.PARAMETER_EXTRACTOR, "title": "Node"}, + node_id="node-id", + tenant_id="tenant-id", + user_id="user-id", + user_inputs={"question": "hello"}, + ) + + assert node.id == "node-id" + assert list(generator) == ["event"] + variable_pool_cls.assert_called_once_with() + add_variables_to_pool.assert_called_once_with(sentinel.variable_pool, sentinel.system_variables) + build_dify_run_context.assert_called_once_with( + tenant_id="tenant-id", + app_id="", + user_id="user-id", + user_from=UserFrom.ACCOUNT, + invoke_from=InvokeFrom.DEBUGGER, + app_type=CreditUsageAppType.WORKFLOW, + ) + graph_init_context_cls.assert_called_once_with( + workflow_id="", + graph_config=workflow_entry.WorkflowEntry._create_single_node_graph( + "node-id", {"type": BuiltinNodeTypes.PARAMETER_EXTRACTOR, "title": "Node"} + ), + run_context={"_dify": "context"}, + call_depth=0, + ) + dify_node_factory_cls.assert_called_once_with( + graph_init_context=sentinel.graph_init_context, + runtime_state=sentinel.graph_runtime_state, + ) + mapping_user_inputs_to_variable_pool.assert_called_once_with( + variable_mapping={}, + user_inputs={"question": "hello"}, + variable_pool=sentinel.variable_pool, + tenant_id="tenant-id", + ) + + def test_run_free_node_wraps_execution_failures(self, monkeypatch: pytest.MonkeyPatch): + class FakeNodeClass: + @staticmethod + def extract_variable_selector_to_variable_mapping(**_kwargs): + return {} + + class FakeNode: + id = "node-id" + title = "Node Title" + node_type = "parameter-extractor" + + @staticmethod + def version(): + return "1" + + dify_node_factory = MagicMock() + dify_node_factory.create_node.return_value = FakeNode() + monkeypatch.setattr( + workflow_entry, + "resolve_workflow_node_class", + MagicMock(return_value=FakeNodeClass), + ) + + with ( + patch.object(workflow_entry, "default_system_variables", return_value=sentinel.system_variables), + patch.object(workflow_entry, "VariablePool", return_value=sentinel.variable_pool), + patch.object(workflow_entry, "add_variables_to_pool"), + patch.object(workflow_entry, "DifyGraphInitContext", return_value=sentinel.graph_init_context), + patch.object(workflow_entry, "RuntimeState", return_value=sentinel.graph_runtime_state), + patch.object(workflow_entry, "build_dify_run_context", return_value={"_dify": "context"}), + patch.object(workflow_entry.time, "perf_counter", return_value=123.0), + patch.object( + workflow_entry.DifyNodeFactory, + "from_graph_init_context", + return_value=dify_node_factory, + ), + patch.object( + workflow_entry.WorkflowEntry, + "mapping_user_inputs_to_variable_pool", + side_effect=RuntimeError("boom"), + ), + ): + with pytest.raises(WorkflowNodeRunFailedError, match="Node Title run failed: boom"): + workflow_entry.WorkflowEntry.run_free_node( + node_data={"type": BuiltinNodeTypes.PARAMETER_EXTRACTOR, "title": "Node"}, + node_id="node-id", + tenant_id="tenant-id", + user_id="user-id", + user_inputs={"question": "hello"}, + ) + + def test_handle_special_values_serializes_nested_files(self): + file = File( + file_type=FileType.IMAGE, + transfer_method=FileTransferMethod.REMOTE_URL, + remote_url="https://example.com/image.png", + filename="image.png", + extension=".png", + ) + + result = workflow_entry.WorkflowEntry.handle_special_values({"file": file, "nested": {"files": [file]}}) + + assert result == { + "file": file.to_dict(), + "nested": {"files": [file.to_dict()]}, + } + + def test_handle_special_values_returns_none_for_none(self): + assert workflow_entry.WorkflowEntry._handle_special_values(None) is None + + def test_handle_special_values_returns_scalar_as_is(self): + assert workflow_entry.WorkflowEntry._handle_special_values("plain-text") == "plain-text" + + class TestMappingUserInputsBranches: + def test_rejects_invalid_node_variable_key(self): + class EmptySplitKey(UserString): + def split(self, _sep=None): + return [] + + with pytest.raises(ValueError, match="Invalid node variable broken"): + workflow_entry.WorkflowEntry.mapping_user_inputs_to_variable_pool( + variable_mapping={EmptySplitKey("broken"): ["node", "input"]}, + user_inputs={}, + variable_pool=MagicMock(), + tenant_id="tenant-id", + ) + + def test_skips_none_user_input_when_variable_already_exists(self): + variable_pool = MagicMock() + variable_pool.get.return_value = None + + workflow_entry.WorkflowEntry.mapping_user_inputs_to_variable_pool( + variable_mapping={"node.input": ["target", "input"]}, + user_inputs={"node.input": None}, + variable_pool=variable_pool, + tenant_id="tenant-id", + ) + + variable_pool.add.assert_not_called() + + def test_merges_structured_output_values(self): + variable_pool = MagicMock() + variable_pool.get.side_effect = [ + None, + SimpleNamespace(value={"existing": "value"}), + ] + + workflow_entry.WorkflowEntry.mapping_user_inputs_to_variable_pool( + variable_mapping={"node.answer": ["target", "structured_output", "answer"]}, + user_inputs={"node.answer": "new-value"}, + variable_pool=variable_pool, + tenant_id="tenant-id", + ) + + variable_pool.add.assert_called_once_with( + ["target", "structured_output"], + {"existing": "value", "answer": "new-value"}, + ) + + class TestWorkflowEntryNodeLayers: + @pytest.mark.parametrize("completion", ["close", "throw"]) + def test_standalone_scope_restores_caller_during_streaming_and_cleanup(self, completion: str) -> None: + observability_layer = MagicMock() + observed_runtime = [] + + def run_node(): + observed_runtime.append(peek_workflow_file_runtime()) + yield sentinel.event + + node = MagicMock() + node.run.side_effect = run_node + error = RuntimeError("consumer failed") + with ( + patch.object(workflow_entry, "ObservabilityLayer", return_value=observability_layer), + patch.object(workflow_entry, "create_dify_workflow_file_runtime", return_value=sentinel.node_runtime), + use_workflow_file_runtime(sentinel.caller_runtime), + ): + events = workflow_entry.WorkflowEntry._run_node_with_layers(node) + assert next(events) is sentinel.event + assert peek_workflow_file_runtime() is sentinel.caller_runtime + if completion == "throw": + with pytest.raises(RuntimeError, match="consumer failed"): + events.throw(error) + else: + events.close() + assert peek_workflow_file_runtime() is sentinel.caller_runtime + + assert observed_runtime == [sentinel.node_runtime] + observability_layer.on_node_run_end.assert_called_once_with( + node, error if completion == "throw" else None, None + ) + observability_layer.on_graph_end.assert_called_once_with(error if completion == "throw" else None) + + def test_run_node_with_layers_reports_success(self, mocker: MockerFixture): + observability_layer = workflow_entry.ObservabilityLayer() + for method in ("initialize", "on_graph_start", "on_node_run_start", "on_node_run_end", "on_graph_end"): + mocker.spy(observability_layer, method) + result_event = NodeRunSucceededEvent( + id="execution-id", + node_id="node-id", + node_type=BuiltinNodeTypes.START, + start_at=datetime.now(), + node_run_result=NodeRunResult(status=WorkflowNodeExecutionStatus.SUCCEEDED), + ) + + class FakeNode: + runtime_state = sentinel.graph_runtime_state + + def bind_execution_id(self, _execution_id): + return None + + def run(self): + yield result_event + + node = FakeNode() + with ( + patch.object(workflow_entry, "ObservabilityLayer", return_value=observability_layer), + patch.object(workflow_entry, "InMemoryChannel", return_value=sentinel.command_channel), + patch.object( + workflow_entry, + "ReadOnlyRuntimeStateWrapper", + return_value=sentinel.read_only_runtime_state, + ) as runtime_state_wrapper, + ): + events = list(workflow_entry.WorkflowEntry._run_node_with_layers(node)) + + assert events == [result_event] + runtime_state_wrapper.assert_called_once_with(sentinel.graph_runtime_state) + for layer in (observability_layer,): + layer.initialize.assert_called_once_with(sentinel.read_only_runtime_state, sentinel.command_channel) + layer.on_graph_start.assert_called_once_with() + layer.on_node_run_start.assert_called_once_with(node) + layer.on_node_run_end.assert_called_once_with(node, None, result_event) + layer.on_graph_end.assert_called_once_with(None) + + def test_run_node_with_layers_reports_errors(self, mocker: MockerFixture): + observability_layer = workflow_entry.ObservabilityLayer() + for method in ("initialize", "on_graph_start", "on_node_run_start", "on_node_run_end", "on_graph_end"): + mocker.spy(observability_layer, method) + + class FakeNode: + runtime_state = sentinel.graph_runtime_state + + def bind_execution_id(self, _execution_id): + return None + + def run(self): + raise RuntimeError("boom") + yield + + node = FakeNode() + with ( + patch.object(workflow_entry, "ObservabilityLayer", return_value=observability_layer), + patch.object( + workflow_entry, + "ReadOnlyRuntimeStateWrapper", + return_value=sentinel.read_only_runtime_state, + ), + ): + with pytest.raises(RuntimeError, match="boom"): + list(workflow_entry.WorkflowEntry._run_node_with_layers(node)) + + for layer in (observability_layer,): + assert layer.on_node_run_end.call_args.args[0] is node + assert isinstance(layer.on_node_run_end.call_args.args[1], RuntimeError) + assert layer.on_node_run_end.call_args.args[2] is None + assert isinstance(layer.on_graph_end.call_args.args[0], RuntimeError) + + +"""Tests for WorkflowEntry integration with Redis command channel.""" + + +class TestWorkflowEntryRedisChannelCases: + class TestWorkflowEntryRedisChannel: + """Test suite for WorkflowEntry with Redis command channel.""" + + def test_workflow_entry_uses_provided_redis_channel(self): + """Test that WorkflowEntry uses the provided Redis command channel.""" + # Mock dependencies + mock_graph = MagicMock() + mock_graph_config = {"nodes": [], "edges": []} + variable_pool = VariablePool() + graph_runtime_state = RuntimeState(workflow_id="test-workflow", variable_pool=variable_pool, start_at=0) + + # Create a mock Redis channel + mock_redis_client = MagicMock() + redis_channel = RedisChannel(mock_redis_client, "test:channel:key") + + # Patch Engine to verify it receives the Redis channel + with patch("core.workflow.workflow_entry.Engine", autospec=True) as mock_graph_engine_type: + # Create WorkflowEntry with Redis channel + workflow_entry = WorkflowEntry( + graph_config=mock_graph_config, + graph=mock_graph, + call_depth=0, + graph_runtime_state=graph_runtime_state, + command_channel=redis_channel, # Provide Redis channel + ) + + # Verify Engine was initialized with the Redis channel + mock_graph_engine_type.assert_called_once() + call_args = mock_graph_engine_type.call_args[1] + assert call_args["command_channel"] == redis_channel + assert workflow_entry.command_channel == redis_channel + + def test_workflow_entry_defaults_to_inmemory_channel(self): + """Test that WorkflowEntry defaults to InMemoryChannel when no channel is provided.""" + # Mock dependencies + mock_graph = MagicMock() + mock_graph_config = {"nodes": [], "edges": []} + variable_pool = VariablePool() + graph_runtime_state = RuntimeState(workflow_id="test-workflow", variable_pool=variable_pool, start_at=0) + + # Patch Engine and InMemoryChannel + with ( + patch("core.workflow.workflow_entry.Engine", autospec=True) as mock_graph_engine_type, + patch("core.workflow.workflow_entry.InMemoryChannel", autospec=True) as mock_in_memory_channel, + ): + mock_inmemory_channel = mock_in_memory_channel.return_value + # Create WorkflowEntry without providing a channel + workflow_entry = WorkflowEntry( + graph_config=mock_graph_config, + graph=mock_graph, + call_depth=0, + graph_runtime_state=graph_runtime_state, + command_channel=None, # No channel provided + ) + + # Verify InMemoryChannel was created + mock_in_memory_channel.assert_called_once() + + # Verify Engine was initialized with the InMemory channel + mock_graph_engine_type.assert_called_once() + call_args = mock_graph_engine_type.call_args[1] + assert call_args["command_channel"] == mock_inmemory_channel + assert workflow_entry.command_channel == mock_inmemory_channel + + def test_workflow_entry_run_with_redis_channel(self): + """Test that WorkflowEntry.run() works correctly with Redis channel.""" + # Mock dependencies + mock_graph = MagicMock() + mock_graph_config = {"nodes": [], "edges": []} + variable_pool = VariablePool() + graph_runtime_state = RuntimeState(workflow_id="test-workflow", variable_pool=variable_pool, start_at=0) + + # Create a mock Redis channel + mock_redis_client = MagicMock() + redis_channel = RedisChannel(mock_redis_client, "test:channel:key") + + # Mock events to be generated + mock_event1 = MagicMock() + mock_event2 = MagicMock() + + # Patch Engine + with patch("core.workflow.workflow_entry.Engine", autospec=True) as mock_graph_engine_type: + mock_graph_engine = MagicMock() + mock_graph_engine.run.return_value = iter([mock_event1, mock_event2]) + mock_graph_engine_type.return_value = mock_graph_engine + + # Create WorkflowEntry with Redis channel + workflow_entry = WorkflowEntry( + graph_config=mock_graph_config, + graph=mock_graph, + call_depth=0, + graph_runtime_state=graph_runtime_state, + command_channel=redis_channel, + ) + + # Run the workflow + events = list(workflow_entry.run()) + + # Verify events were generated + assert len(events) == 2 + assert events[0] == mock_event1 + assert events[1] == mock_event2 diff --git a/api/tests/unit_tests/core/workflow/test_workflow_entry_helpers.py b/api/tests/unit_tests/core/workflow/test_workflow_entry_helpers.py deleted file mode 100644 index 0ae05f7ddcf..00000000000 --- a/api/tests/unit_tests/core/workflow/test_workflow_entry_helpers.py +++ /dev/null @@ -1,810 +0,0 @@ -from collections import UserString -from datetime import datetime -from types import SimpleNamespace -from unittest.mock import MagicMock, patch, sentinel - -import pytest -from pytest_mock import MockerFixture - -from core.app.apps.exc import GenerateTaskStoppedError -from core.app.entities.app_invoke_entities import InvokeFrom, UserFrom -from core.credit_usage import CreditUsageAppType -from core.workflow import workflow_entry -from core.workflow.system_variables import default_system_variables -from graphon.engine.filter import ResponseStreamFilter -from graphon.engine_events import GraphRunAbortedEvent, GraphRunFailedEvent, NodeRunSucceededEvent -from graphon.entities.base_node_data import BaseNodeData -from graphon.enums import NodeType, WorkflowNodeExecutionStatus -from graphon.errors import WorkflowNodeRunFailedError -from graphon.file import File, FileTransferMethod, FileType -from graphon.file.runtime import peek_workflow_file_runtime, use_workflow_file_runtime -from graphon.node_events import NodeRunResult -from graphon.nodes import BuiltinNodeTypes -from graphon.runtime import VariablePool -from graphon.variables.variables import StringVariable -from models.workflow import Workflow -from tests.unit_tests.config_override import config_overrides_context -from tests.unit_tests.core.workflow.test_workflow_tool_container import ( - _outer_graph, - _workflow_tool_node, - _workflow_tool_source, -) -from tests.unit_tests.model_factories import make_workflow - - -def _build_typed_node_config(node_type: NodeType): - return {"id": "node-id", "data": BaseNodeData(type=node_type)} - - -def _workflow() -> Workflow: - """Build a real transient workflow for single-step orchestration tests.""" - return make_workflow(workflow_id="workflow-id", tenant_id="tenant-id", app_id="app-id", created_by="user-id") - - -def _build_minimal_workflow_entry( - monkeypatch: pytest.MonkeyPatch, - *, - response_stream_filter: ResponseStreamFilter | None = None, -) -> workflow_entry.WorkflowEntry: - """Construct a minimal WorkflowEntry with Engine construction mocked out.""" - graph_engine = MagicMock() - graph_runtime_state = sentinel.runtime_state - - monkeypatch.setattr(workflow_entry, "ExecutionContextLayer", MagicMock()) - monkeypatch.setattr(workflow_entry, "Engine", MagicMock(return_value=graph_engine)) - monkeypatch.setattr(workflow_entry, "InMemoryChannel", MagicMock(return_value=sentinel.command_channel)) - - return workflow_entry.WorkflowEntry( - graph_config={"nodes": [], "edges": []}, - graph=sentinel.graph, - call_depth=0, - graph_runtime_state=graph_runtime_state, - response_stream_filter=response_stream_filter, - ) - - -class TestWorkflowEntryInit: - def test_rejects_call_depth_above_limit(self): - call_depth = workflow_entry.dify_config.WORKFLOW_CALL_MAX_DEPTH + 1 - - with pytest.raises(ValueError, match="Max workflow call depth"): - workflow_entry.WorkflowEntry( - graph_config={"nodes": [], "edges": []}, - graph=sentinel.graph, - call_depth=call_depth, - graph_runtime_state=sentinel.graph_runtime_state, - ) - - def test_hidden_workflow_tool_nodes_count_toward_execution_limit(self): - tool, _, _ = _workflow_tool_node() - source = _workflow_tool_source() - with config_overrides_context(WORKFLOW_MAX_EXECUTION_STEPS=2, GRAPH_ENGINE_MAX_WORKERS=1): - entry = workflow_entry.WorkflowEntry( - graph_config=tool.init_params.graph_config, - graph=_outer_graph(tool), - call_depth=0, - graph_runtime_state=tool.runtime_state, - ) - entry.workflow_tool_sources[source.workflow_id] = source - events = list(entry.run()) - - assert isinstance(events[-1], GraphRunAbortedEvent) - assert "Maximum execution steps exceeded: 3 > 2" in events[-1].reason - assert not any(isinstance(event, NodeRunSucceededEvent) and event.node_id == "tool" for event in events) - - def test_workflow_entry_stores_supplied_response_stream_filter(self, monkeypatch: pytest.MonkeyPatch) -> None: - supplied_filter = ResponseStreamFilter() - entry = _build_minimal_workflow_entry(monkeypatch, response_stream_filter=supplied_filter) - - assert entry._response_stream_filter is supplied_filter - - def test_workflow_entry_defaults_to_fresh_response_stream_filter(self, monkeypatch: pytest.MonkeyPatch) -> None: - entry = _build_minimal_workflow_entry(monkeypatch, response_stream_filter=None) - - assert isinstance(entry._response_stream_filter, ResponseStreamFilter) - - -class TestWorkflowEntryRun: - def test_run_swallows_generate_task_stopped_errors(self): - entry = object.__new__(workflow_entry.WorkflowEntry) - entry.graph_engine = MagicMock() - entry.graph_engine.run.side_effect = GenerateTaskStoppedError() - entry._response_stream_filter = ResponseStreamFilter() - - assert list(entry.run()) == [] - - def test_iter_dify_graph_engine_events_applies_response_stream_filter(self): - graph_engine = MagicMock() - graph_engine.run.return_value = iter([sentinel.raw_event]) - - with ( - patch.object( - workflow_entry.EngineEventFilterContext, - "from_engine", - return_value=sentinel.filter_context, - ) as from_engine, - patch.object( - workflow_entry, - "ResponseStreamFilter", - return_value=sentinel.response_stream_filter, - ) as response_stream_filter_cls, - patch.object( - workflow_entry, - "HumanInputFormEventFilter", - return_value=sentinel.human_input_filter, - ), - patch.object( - workflow_entry, - "filter_engine_events", - return_value=iter([sentinel.filtered_event]), - ) as filter_engine_events, - ): - events = list(workflow_entry.iter_dify_graph_engine_events(graph_engine)) - - assert events == [sentinel.filtered_event] - from_engine.assert_called_once_with(graph_engine) - response_stream_filter_cls.assert_called_once_with() - filter_engine_events.assert_called_once_with( - graph_engine.run.return_value, - context=sentinel.filter_context, - filters=[sentinel.human_input_filter, sentinel.response_stream_filter], - ) - - def test_run_delegates_to_dify_event_iterator(self): - entry = object.__new__(workflow_entry.WorkflowEntry) - entry.graph_engine = sentinel.graph_engine - entry._response_stream_filter = sentinel.response_stream_filter - - with patch.object( - workflow_entry, - "iter_dify_graph_engine_events", - return_value=iter([sentinel.filtered_event]), - ) as iter_dify_graph_engine_events: - events = list(entry.run()) - - assert events == [sentinel.filtered_event] - iter_dify_graph_engine_events.assert_called_once_with(sentinel.graph_engine, sentinel.response_stream_filter) - - def test_run_emits_failed_event_for_unexpected_errors(self): - entry = object.__new__(workflow_entry.WorkflowEntry) - entry.graph_engine = MagicMock() - entry.graph_engine.run.side_effect = RuntimeError("boom") - entry._response_stream_filter = ResponseStreamFilter() - - events = list(entry.run()) - - assert len(events) == 1 - assert isinstance(events[0], GraphRunFailedEvent) - assert events[0].error == "boom" - - -class TestWorkflowEntrySingleStepRun: - @pytest.mark.parametrize("node_type", [BuiltinNodeTypes.LOOP, BuiltinNodeTypes.ITERATION]) - def test_rejects_container_nodes(self, node_type): - workflow = _workflow() - - with ( - patch.object(workflow, "get_node_config_by_id", return_value=_build_typed_node_config(node_type)), - pytest.raises(ValueError, match="engine-backed debug endpoints"), - ): - workflow_entry.WorkflowEntry.single_step_run( - workflow=workflow, - node_id="node-id", - user_id="user-id", - user_inputs={}, - variable_pool=sentinel.variable_pool, - execution_driver=None, - ) - - def test_preloads_constructor_variables_before_creating_memory_node(self): - class FakeLLMNode: - id = "node-id" - title = "Node Title" - node_type = BuiltinNodeTypes.LLM - - @staticmethod - def version(): - return "1" - - @staticmethod - def extract_variable_selector_to_variable_mapping(**_kwargs): - return {} - - variable_pool = VariablePool.from_bootstrap(system_variables=default_system_variables(), user_inputs={}) - variable_loader = MagicMock() - variable_loader.load_variables.return_value = [ - StringVariable( - name="conversation_id", - value="conv-1", - selector=["sys", "conversation_id"], - ) - ] - workflow = _workflow() - node_config = { - "id": "node-id", - "data": BaseNodeData(type=BuiltinNodeTypes.LLM, version="1", memory=object()), - } - - with ( - patch.object(workflow, "get_node_config_by_id", return_value=node_config), - patch.object(workflow_entry, "DifyGraphInitContext", return_value=sentinel.graph_init_context), - patch.object( - workflow_entry, - "RuntimeState", - return_value=SimpleNamespace(variable_pool=variable_pool), - ), - patch.object(workflow_entry, "build_dify_run_context", return_value={"_dify": "context"}), - patch.object(workflow_entry.time, "perf_counter", return_value=123.0), - patch.object(workflow_entry, "resolve_workflow_node_class", return_value=FakeLLMNode), - patch.object(workflow_entry.DifyNodeFactory, "from_graph_init_context") as dify_node_factory, - patch.object(workflow_entry, "load_into_variable_pool"), - patch.object(workflow_entry.WorkflowEntry, "mapping_user_inputs_to_variable_pool"), - patch.object( - workflow_entry.WorkflowEntry, - "_run_node_with_layers", - return_value=iter(["event"]), - ), - ): - - def _create_node(_node_config): - assert variable_pool.get(["sys", "conversation_id"]) is not None - return FakeLLMNode() - - dify_node_factory.return_value.create_node.side_effect = _create_node - node, generator = workflow_entry.WorkflowEntry.single_step_run( - workflow=workflow, - node_id="node-id", - user_id="user-id", - user_inputs={}, - variable_pool=variable_pool, - variable_loader=variable_loader, - execution_driver=None, - ) - - assert node.id == "node-id" - assert list(generator) == ["event"] - variable_loader.load_variables.assert_called_once_with([["sys", "conversation_id"]]) - - def test_uses_empty_mapping_when_selector_extraction_is_not_implemented(self): - class FakeNode: - id = "node-id" - title = "Node Title" - node_type = "fake" - - @staticmethod - def version(): - return "1" - - @staticmethod - def extract_variable_selector_to_variable_mapping(**_kwargs): - raise NotImplementedError - - workflow = _workflow() - with ( - patch.object( - workflow, - "get_node_config_by_id", - return_value=_build_typed_node_config(BuiltinNodeTypes.START), - ), - patch.object(workflow_entry, "DifyGraphInitContext", return_value=sentinel.graph_init_context), - patch.object(workflow_entry, "RuntimeState", return_value=sentinel.graph_runtime_state), - patch.object(workflow_entry, "build_dify_run_context", return_value={"_dify": "context"}), - patch.object(workflow_entry.time, "perf_counter", return_value=123.0), - patch.object(workflow_entry, "resolve_workflow_node_class", return_value=FakeNode), - patch.object(workflow_entry.DifyNodeFactory, "from_graph_init_context") as dify_node_factory, - patch.object(workflow_entry, "add_node_inputs_to_pool") as add_node_inputs_to_pool, - patch.object(workflow_entry, "load_into_variable_pool") as load_into_variable_pool, - patch.object( - workflow_entry.WorkflowEntry, - "mapping_user_inputs_to_variable_pool", - ) as mapping_user_inputs_to_variable_pool, - patch.object( - workflow_entry.WorkflowEntry, - "_run_node_with_layers", - return_value=iter(["event"]), - ), - ): - dify_node_factory.return_value.create_node.return_value = FakeNode() - node, generator = workflow_entry.WorkflowEntry.single_step_run( - workflow=workflow, - node_id="node-id", - user_id="user-id", - user_inputs={"question": "hello"}, - variable_pool=sentinel.variable_pool, - execution_driver=None, - ) - - assert node.id == "node-id" - assert list(generator) == ["event"] - load_into_variable_pool.assert_called_once_with( - variable_loader=workflow_entry.DUMMY_VARIABLE_LOADER, - variable_pool=sentinel.variable_pool, - variable_mapping={}, - user_inputs={"question": "hello"}, - ) - add_node_inputs_to_pool.assert_called_once_with( - sentinel.variable_pool, - node_id="node-id", - inputs={"question": "hello"}, - ) - mapping_user_inputs_to_variable_pool.assert_called_once_with( - variable_mapping={}, - user_inputs={"question": "hello"}, - variable_pool=sentinel.variable_pool, - tenant_id="tenant-id", - ) - - def test_skips_user_input_mapping_for_datasource_nodes(self): - class FakeDatasourceNode: - id = "node-id" - node_type = "datasource" - - @staticmethod - def version(): - return "1" - - @staticmethod - def extract_variable_selector_to_variable_mapping(**_kwargs): - return {"question": ["node", "question"]} - - workflow = _workflow() - with ( - patch.object( - workflow, - "get_node_config_by_id", - return_value=_build_typed_node_config(BuiltinNodeTypes.DATASOURCE), - ), - patch.object(workflow_entry, "DifyGraphInitContext", return_value=sentinel.graph_init_context), - patch.object(workflow_entry, "RuntimeState", return_value=sentinel.graph_runtime_state), - patch.object(workflow_entry, "build_dify_run_context", return_value={"_dify": "context"}), - patch.object(workflow_entry.time, "perf_counter", return_value=123.0), - patch.object(workflow_entry, "resolve_workflow_node_class", return_value=FakeDatasourceNode), - patch.object(workflow_entry.DifyNodeFactory, "from_graph_init_context") as dify_node_factory, - patch.object(workflow_entry, "add_node_inputs_to_pool") as add_node_inputs_to_pool, - patch.object(workflow_entry, "load_into_variable_pool") as load_into_variable_pool, - patch.object( - workflow_entry.WorkflowEntry, - "mapping_user_inputs_to_variable_pool", - ) as mapping_user_inputs_to_variable_pool, - patch.object( - workflow_entry.WorkflowEntry, - "_run_node_with_layers", - return_value=iter(["event"]), - ), - ): - dify_node_factory.return_value.create_node.return_value = FakeDatasourceNode() - node, generator = workflow_entry.WorkflowEntry.single_step_run( - workflow=workflow, - node_id="node-id", - user_id="user-id", - user_inputs={"question": "hello"}, - variable_pool=sentinel.variable_pool, - execution_driver=None, - ) - - assert node.id == "node-id" - assert list(generator) == ["event"] - load_into_variable_pool.assert_called_once() - add_node_inputs_to_pool.assert_called_once_with( - sentinel.variable_pool, - node_id="node-id", - inputs={"question": "hello"}, - ) - mapping_user_inputs_to_variable_pool.assert_not_called() - - def test_wraps_layered_node_run_failures(self): - class FakeNode: - id = "node-id" - title = "Node Title" - node_type = "fake" - - @staticmethod - def extract_variable_selector_to_variable_mapping(**_kwargs): - return {} - - @staticmethod - def version(): - return "1" - - workflow = _workflow() - with ( - patch.object( - workflow, - "get_node_config_by_id", - return_value=_build_typed_node_config(BuiltinNodeTypes.START), - ), - patch.object(workflow_entry, "DifyGraphInitContext", return_value=sentinel.graph_init_context), - patch.object(workflow_entry, "RuntimeState", return_value=sentinel.graph_runtime_state), - patch.object(workflow_entry, "build_dify_run_context", return_value={"_dify": "context"}), - patch.object(workflow_entry.time, "perf_counter", return_value=123.0), - patch.object(workflow_entry, "resolve_workflow_node_class", return_value=FakeNode), - patch.object(workflow_entry.DifyNodeFactory, "from_graph_init_context") as dify_node_factory, - patch.object(workflow_entry, "add_node_inputs_to_pool"), - patch.object(workflow_entry, "load_into_variable_pool"), - patch.object(workflow_entry.WorkflowEntry, "mapping_user_inputs_to_variable_pool"), - patch.object( - workflow_entry.WorkflowEntry, - "_run_node_with_layers", - side_effect=RuntimeError("boom"), - ), - ): - dify_node_factory.return_value.create_node.return_value = FakeNode() - with pytest.raises(WorkflowNodeRunFailedError): - workflow_entry.WorkflowEntry.single_step_run( - workflow=workflow, - node_id="node-id", - user_id="user-id", - user_inputs={}, - variable_pool=sentinel.variable_pool, - execution_driver=None, - ) - - -class TestWorkflowEntryHelpers: - def test_create_single_node_graph_builds_start_edge(self): - graph = workflow_entry.WorkflowEntry._create_single_node_graph( - node_id="target-node", - node_data={"type": BuiltinNodeTypes.PARAMETER_EXTRACTOR}, - node_width=320, - node_height=180, - ) - - assert graph["nodes"][0]["id"] == "start" - assert graph["nodes"][1]["id"] == "target-node" - assert graph["nodes"][1]["width"] == 320 - assert graph["nodes"][1]["height"] == 180 - assert graph["edges"] == [ - { - "source": "start", - "target": "target-node", - "sourceHandle": "source", - "targetHandle": "target", - } - ] - - def test_run_free_node_rejects_unsupported_types(self): - with pytest.raises(ValueError, match="Node type start not supported"): - workflow_entry.WorkflowEntry.run_free_node( - node_data={"type": BuiltinNodeTypes.START}, - node_id="node-id", - tenant_id="tenant-id", - user_id="user-id", - user_inputs={}, - ) - - def test_run_free_node_rejects_missing_node_class(self, monkeypatch: pytest.MonkeyPatch): - monkeypatch.setattr( - workflow_entry, - "resolve_workflow_node_class", - MagicMock(return_value=None), - ) - - with pytest.raises(ValueError, match="Node class not found for node type parameter-extractor"): - workflow_entry.WorkflowEntry.run_free_node( - node_data={"type": BuiltinNodeTypes.PARAMETER_EXTRACTOR}, - node_id="node-id", - tenant_id="tenant-id", - user_id="user-id", - user_inputs={}, - ) - - def test_run_free_node_uses_empty_mapping_when_selector_extraction_is_not_implemented( - self, monkeypatch: pytest.MonkeyPatch - ): - class FakeNodeClass: - @staticmethod - def extract_variable_selector_to_variable_mapping(**_kwargs): - raise NotImplementedError - - class FakeNode: - id = "node-id" - title = "Node Title" - node_type = "parameter-extractor" - - @staticmethod - def version(): - return "1" - - dify_node_factory = MagicMock() - dify_node_factory.create_node.return_value = FakeNode() - monkeypatch.setattr( - workflow_entry, - "resolve_workflow_node_class", - MagicMock(return_value=FakeNodeClass), - ) - - with ( - patch.object(workflow_entry, "default_system_variables", return_value=sentinel.system_variables), - patch.object(workflow_entry, "VariablePool", return_value=sentinel.variable_pool) as variable_pool_cls, - patch.object(workflow_entry, "add_variables_to_pool") as add_variables_to_pool, - patch.object( - workflow_entry, "DifyGraphInitContext", return_value=sentinel.graph_init_context - ) as graph_init_context_cls, - patch.object(workflow_entry, "RuntimeState", return_value=sentinel.graph_runtime_state), - patch.object( - workflow_entry, "build_dify_run_context", return_value={"_dify": "context"} - ) as build_dify_run_context, - patch.object(workflow_entry.time, "perf_counter", return_value=123.0), - patch.object( - workflow_entry.DifyNodeFactory, - "from_graph_init_context", - return_value=dify_node_factory, - ) as dify_node_factory_cls, - patch.object( - workflow_entry.WorkflowEntry, - "mapping_user_inputs_to_variable_pool", - ) as mapping_user_inputs_to_variable_pool, - patch.object( - workflow_entry.WorkflowEntry, - "_run_node_with_layers", - return_value=iter(["event"]), - ), - ): - node, generator = workflow_entry.WorkflowEntry.run_free_node( - node_data={"type": BuiltinNodeTypes.PARAMETER_EXTRACTOR, "title": "Node"}, - node_id="node-id", - tenant_id="tenant-id", - user_id="user-id", - user_inputs={"question": "hello"}, - ) - - assert node.id == "node-id" - assert list(generator) == ["event"] - variable_pool_cls.assert_called_once_with() - add_variables_to_pool.assert_called_once_with(sentinel.variable_pool, sentinel.system_variables) - build_dify_run_context.assert_called_once_with( - tenant_id="tenant-id", - app_id="", - user_id="user-id", - user_from=UserFrom.ACCOUNT, - invoke_from=InvokeFrom.DEBUGGER, - app_type=CreditUsageAppType.WORKFLOW, - ) - graph_init_context_cls.assert_called_once_with( - workflow_id="", - graph_config=workflow_entry.WorkflowEntry._create_single_node_graph( - "node-id", {"type": BuiltinNodeTypes.PARAMETER_EXTRACTOR, "title": "Node"} - ), - run_context={"_dify": "context"}, - call_depth=0, - ) - dify_node_factory_cls.assert_called_once_with( - graph_init_context=sentinel.graph_init_context, - runtime_state=sentinel.graph_runtime_state, - ) - mapping_user_inputs_to_variable_pool.assert_called_once_with( - variable_mapping={}, - user_inputs={"question": "hello"}, - variable_pool=sentinel.variable_pool, - tenant_id="tenant-id", - ) - - def test_run_free_node_wraps_execution_failures(self, monkeypatch: pytest.MonkeyPatch): - class FakeNodeClass: - @staticmethod - def extract_variable_selector_to_variable_mapping(**_kwargs): - return {} - - class FakeNode: - id = "node-id" - title = "Node Title" - node_type = "parameter-extractor" - - @staticmethod - def version(): - return "1" - - dify_node_factory = MagicMock() - dify_node_factory.create_node.return_value = FakeNode() - monkeypatch.setattr( - workflow_entry, - "resolve_workflow_node_class", - MagicMock(return_value=FakeNodeClass), - ) - - with ( - patch.object(workflow_entry, "default_system_variables", return_value=sentinel.system_variables), - patch.object(workflow_entry, "VariablePool", return_value=sentinel.variable_pool), - patch.object(workflow_entry, "add_variables_to_pool"), - patch.object(workflow_entry, "DifyGraphInitContext", return_value=sentinel.graph_init_context), - patch.object(workflow_entry, "RuntimeState", return_value=sentinel.graph_runtime_state), - patch.object(workflow_entry, "build_dify_run_context", return_value={"_dify": "context"}), - patch.object(workflow_entry.time, "perf_counter", return_value=123.0), - patch.object( - workflow_entry.DifyNodeFactory, - "from_graph_init_context", - return_value=dify_node_factory, - ), - patch.object( - workflow_entry.WorkflowEntry, - "mapping_user_inputs_to_variable_pool", - side_effect=RuntimeError("boom"), - ), - ): - with pytest.raises(WorkflowNodeRunFailedError, match="Node Title run failed: boom"): - workflow_entry.WorkflowEntry.run_free_node( - node_data={"type": BuiltinNodeTypes.PARAMETER_EXTRACTOR, "title": "Node"}, - node_id="node-id", - tenant_id="tenant-id", - user_id="user-id", - user_inputs={"question": "hello"}, - ) - - def test_handle_special_values_serializes_nested_files(self): - file = File( - file_type=FileType.IMAGE, - transfer_method=FileTransferMethod.REMOTE_URL, - remote_url="https://example.com/image.png", - filename="image.png", - extension=".png", - ) - - result = workflow_entry.WorkflowEntry.handle_special_values({"file": file, "nested": {"files": [file]}}) - - assert result == { - "file": file.to_dict(), - "nested": {"files": [file.to_dict()]}, - } - - def test_handle_special_values_returns_none_for_none(self): - assert workflow_entry.WorkflowEntry._handle_special_values(None) is None - - def test_handle_special_values_returns_scalar_as_is(self): - assert workflow_entry.WorkflowEntry._handle_special_values("plain-text") == "plain-text" - - -class TestMappingUserInputsBranches: - def test_rejects_invalid_node_variable_key(self): - class EmptySplitKey(UserString): - def split(self, _sep=None): - return [] - - with pytest.raises(ValueError, match="Invalid node variable broken"): - workflow_entry.WorkflowEntry.mapping_user_inputs_to_variable_pool( - variable_mapping={EmptySplitKey("broken"): ["node", "input"]}, - user_inputs={}, - variable_pool=MagicMock(), - tenant_id="tenant-id", - ) - - def test_skips_none_user_input_when_variable_already_exists(self): - variable_pool = MagicMock() - variable_pool.get.return_value = None - - workflow_entry.WorkflowEntry.mapping_user_inputs_to_variable_pool( - variable_mapping={"node.input": ["target", "input"]}, - user_inputs={"node.input": None}, - variable_pool=variable_pool, - tenant_id="tenant-id", - ) - - variable_pool.add.assert_not_called() - - def test_merges_structured_output_values(self): - variable_pool = MagicMock() - variable_pool.get.side_effect = [ - None, - SimpleNamespace(value={"existing": "value"}), - ] - - workflow_entry.WorkflowEntry.mapping_user_inputs_to_variable_pool( - variable_mapping={"node.answer": ["target", "structured_output", "answer"]}, - user_inputs={"node.answer": "new-value"}, - variable_pool=variable_pool, - tenant_id="tenant-id", - ) - - variable_pool.add.assert_called_once_with( - ["target", "structured_output"], - {"existing": "value", "answer": "new-value"}, - ) - - -class TestWorkflowEntryNodeLayers: - @pytest.mark.parametrize("completion", ["close", "throw"]) - def test_standalone_scope_restores_caller_during_streaming_and_cleanup(self, completion: str) -> None: - observability_layer = MagicMock() - observed_runtime = [] - - def run_node(): - observed_runtime.append(peek_workflow_file_runtime()) - yield sentinel.event - - node = MagicMock() - node.run.side_effect = run_node - error = RuntimeError("consumer failed") - with ( - patch.object(workflow_entry, "ObservabilityLayer", return_value=observability_layer), - patch.object(workflow_entry, "create_dify_workflow_file_runtime", return_value=sentinel.node_runtime), - use_workflow_file_runtime(sentinel.caller_runtime), - ): - events = workflow_entry.WorkflowEntry._run_node_with_layers(node) - assert next(events) is sentinel.event - assert peek_workflow_file_runtime() is sentinel.caller_runtime - if completion == "throw": - with pytest.raises(RuntimeError, match="consumer failed"): - events.throw(error) - else: - events.close() - assert peek_workflow_file_runtime() is sentinel.caller_runtime - - assert observed_runtime == [sentinel.node_runtime] - observability_layer.on_node_run_end.assert_called_once_with( - node, error if completion == "throw" else None, None - ) - observability_layer.on_graph_end.assert_called_once_with(error if completion == "throw" else None) - - def test_run_node_with_layers_reports_success(self, mocker: MockerFixture): - observability_layer = workflow_entry.ObservabilityLayer() - for method in ("initialize", "on_graph_start", "on_node_run_start", "on_node_run_end", "on_graph_end"): - mocker.spy(observability_layer, method) - result_event = NodeRunSucceededEvent( - id="execution-id", - node_id="node-id", - node_type=BuiltinNodeTypes.START, - start_at=datetime.now(), - node_run_result=NodeRunResult(status=WorkflowNodeExecutionStatus.SUCCEEDED), - ) - - class FakeNode: - runtime_state = sentinel.graph_runtime_state - - def bind_execution_id(self, _execution_id): - return None - - def run(self): - yield result_event - - node = FakeNode() - with ( - patch.object(workflow_entry, "ObservabilityLayer", return_value=observability_layer), - patch.object(workflow_entry, "InMemoryChannel", return_value=sentinel.command_channel), - patch.object( - workflow_entry, - "ReadOnlyRuntimeStateWrapper", - return_value=sentinel.read_only_runtime_state, - ) as runtime_state_wrapper, - ): - events = list(workflow_entry.WorkflowEntry._run_node_with_layers(node)) - - assert events == [result_event] - runtime_state_wrapper.assert_called_once_with(sentinel.graph_runtime_state) - for layer in (observability_layer,): - layer.initialize.assert_called_once_with(sentinel.read_only_runtime_state, sentinel.command_channel) - layer.on_graph_start.assert_called_once_with() - layer.on_node_run_start.assert_called_once_with(node) - layer.on_node_run_end.assert_called_once_with(node, None, result_event) - layer.on_graph_end.assert_called_once_with(None) - - def test_run_node_with_layers_reports_errors(self, mocker: MockerFixture): - observability_layer = workflow_entry.ObservabilityLayer() - for method in ("initialize", "on_graph_start", "on_node_run_start", "on_node_run_end", "on_graph_end"): - mocker.spy(observability_layer, method) - - class FakeNode: - runtime_state = sentinel.graph_runtime_state - - def bind_execution_id(self, _execution_id): - return None - - def run(self): - raise RuntimeError("boom") - yield - - node = FakeNode() - with ( - patch.object(workflow_entry, "ObservabilityLayer", return_value=observability_layer), - patch.object( - workflow_entry, - "ReadOnlyRuntimeStateWrapper", - return_value=sentinel.read_only_runtime_state, - ), - ): - with pytest.raises(RuntimeError, match="boom"): - list(workflow_entry.WorkflowEntry._run_node_with_layers(node)) - - for layer in (observability_layer,): - assert layer.on_node_run_end.call_args.args[0] is node - assert isinstance(layer.on_node_run_end.call_args.args[1], RuntimeError) - assert layer.on_node_run_end.call_args.args[2] is None - assert isinstance(layer.on_graph_end.call_args.args[0], RuntimeError) diff --git a/api/tests/unit_tests/core/workflow/test_workflow_entry_redis_channel.py b/api/tests/unit_tests/core/workflow/test_workflow_entry_redis_channel.py deleted file mode 100644 index 46a11f9bede..00000000000 --- a/api/tests/unit_tests/core/workflow/test_workflow_entry_redis_channel.py +++ /dev/null @@ -1,111 +0,0 @@ -"""Tests for WorkflowEntry integration with Redis command channel.""" - -from unittest.mock import MagicMock, patch - -from core.workflow.workflow_entry import WorkflowEntry -from graphon.engine.command import RedisChannel -from graphon.runtime import RuntimeState, VariablePool - - -class TestWorkflowEntryRedisChannel: - """Test suite for WorkflowEntry with Redis command channel.""" - - def test_workflow_entry_uses_provided_redis_channel(self): - """Test that WorkflowEntry uses the provided Redis command channel.""" - # Mock dependencies - mock_graph = MagicMock() - mock_graph_config = {"nodes": [], "edges": []} - variable_pool = VariablePool() - graph_runtime_state = RuntimeState(workflow_id="test-workflow", variable_pool=variable_pool, start_at=0) - - # Create a mock Redis channel - mock_redis_client = MagicMock() - redis_channel = RedisChannel(mock_redis_client, "test:channel:key") - - # Patch Engine to verify it receives the Redis channel - with patch("core.workflow.workflow_entry.Engine", autospec=True) as mock_graph_engine_type: - # Create WorkflowEntry with Redis channel - workflow_entry = WorkflowEntry( - graph_config=mock_graph_config, - graph=mock_graph, - call_depth=0, - graph_runtime_state=graph_runtime_state, - command_channel=redis_channel, # Provide Redis channel - ) - - # Verify Engine was initialized with the Redis channel - mock_graph_engine_type.assert_called_once() - call_args = mock_graph_engine_type.call_args[1] - assert call_args["command_channel"] == redis_channel - assert workflow_entry.command_channel == redis_channel - - def test_workflow_entry_defaults_to_inmemory_channel(self): - """Test that WorkflowEntry defaults to InMemoryChannel when no channel is provided.""" - # Mock dependencies - mock_graph = MagicMock() - mock_graph_config = {"nodes": [], "edges": []} - variable_pool = VariablePool() - graph_runtime_state = RuntimeState(workflow_id="test-workflow", variable_pool=variable_pool, start_at=0) - - # Patch Engine and InMemoryChannel - with ( - patch("core.workflow.workflow_entry.Engine", autospec=True) as mock_graph_engine_type, - patch("core.workflow.workflow_entry.InMemoryChannel", autospec=True) as mock_in_memory_channel, - ): - mock_inmemory_channel = mock_in_memory_channel.return_value - # Create WorkflowEntry without providing a channel - workflow_entry = WorkflowEntry( - graph_config=mock_graph_config, - graph=mock_graph, - call_depth=0, - graph_runtime_state=graph_runtime_state, - command_channel=None, # No channel provided - ) - - # Verify InMemoryChannel was created - mock_in_memory_channel.assert_called_once() - - # Verify Engine was initialized with the InMemory channel - mock_graph_engine_type.assert_called_once() - call_args = mock_graph_engine_type.call_args[1] - assert call_args["command_channel"] == mock_inmemory_channel - assert workflow_entry.command_channel == mock_inmemory_channel - - def test_workflow_entry_run_with_redis_channel(self): - """Test that WorkflowEntry.run() works correctly with Redis channel.""" - # Mock dependencies - mock_graph = MagicMock() - mock_graph_config = {"nodes": [], "edges": []} - variable_pool = VariablePool() - graph_runtime_state = RuntimeState(workflow_id="test-workflow", variable_pool=variable_pool, start_at=0) - - # Create a mock Redis channel - mock_redis_client = MagicMock() - redis_channel = RedisChannel(mock_redis_client, "test:channel:key") - - # Mock events to be generated - mock_event1 = MagicMock() - mock_event2 = MagicMock() - - # Patch Engine - with patch("core.workflow.workflow_entry.Engine", autospec=True) as mock_graph_engine_type: - mock_graph_engine = MagicMock() - mock_graph_engine.run.return_value = iter([mock_event1, mock_event2]) - mock_graph_engine_type.return_value = mock_graph_engine - - # Create WorkflowEntry with Redis channel - workflow_entry = WorkflowEntry( - graph_config=mock_graph_config, - graph=mock_graph, - call_depth=0, - graph_runtime_state=graph_runtime_state, - command_channel=redis_channel, - ) - - # Run the workflow - events = list(workflow_entry.run()) - - # Verify events were generated - assert len(events) == 2 - assert events[0] == mock_event1 - assert events[1] == mock_event2 diff --git a/api/tests/unit_tests/core/workflow/test_workflow_tool_container.py b/api/tests/unit_tests/core/workflow/test_workflow_tool_container.py deleted file mode 100644 index b7ed47ce861..00000000000 --- a/api/tests/unit_tests/core/workflow/test_workflow_tool_container.py +++ /dev/null @@ -1,1077 +0,0 @@ -from __future__ import annotations - -import json -from collections.abc import Callable -from dataclasses import replace -from datetime import UTC, datetime -from functools import partial - -import pytest -from pytest_mock import MockerFixture - -from core.app.entities.app_invoke_entities import InvokeFrom -from core.helper.code_executor.code_executor import CodeExecutionError -from core.repositories.human_input_repository import ( - HumanInputFormRepositoryImpl, - HumanInputFormSubmissionRepository, -) -from core.tools.__base.tool_runtime import ToolRuntime -from core.tools.entities.common_entities import I18nObject -from core.tools.entities.tool_entities import ToolEntity, ToolIdentity -from core.tools.workflow_as_tool.repository import WorkflowToolSource -from core.tools.workflow_as_tool.tool import WorkflowTool -from core.workflow.node_factory import DifyNodeFactory -from core.workflow.node_runtime import DifyToolFileManager, DifyToolNodeRuntime -from core.workflow.nodes.human_input.boundary import resolve_human_input_node_id -from core.workflow.nodes.human_input.callback import DifyHITLCallback -from core.workflow.nodes.human_input.entities import HumanInputNodeData, UserActionConfig -from core.workflow.system_variables import SystemVariableKey, system_variable_selector -from core.workflow.workflow_tool_container_handler import ( - WorkflowToolContainerHandler, - WorkflowToolNestedContainerHandler, -) -from core.workflow.workflow_tool_container_types import WorkflowToolContainerPayload -from core.workflow.workflow_tool_node import DifyWorkflowToolNode -from enums import WorkflowKind -from graphon.engine import Engine -from graphon.engine.command import AbortCommand, InMemoryChannel -from graphon.engine.container_handler.builtin.loop import LoopContainerHandler -from graphon.engine.event.node_failure import NodeFailureHandler -from graphon.engine.event.processor import NodeEventProcessor -from graphon.engine.event.stream import EventStream -from graphon.engine.frame import ExecutionFrame, FrameRegistry -from graphon.engine.ready_queue import ResumeTask, StartTask -from graphon.engine.scheduler import Scheduler -from graphon.engine.worker import NodeEventTask -from graphon.engine_events import NodeRunExceptionEvent, NodeRunFailedEvent, NodeRunRetryEvent, NodeRunSucceededEvent -from graphon.engine_events.base import EngineEvent, NodeEvent -from graphon.engine_events.graph import ( - GraphRunAbortedEvent, - GraphRunFailedEvent, - GraphRunPartialSucceededEvent, - GraphRunPausedEvent, - GraphRunSucceededEvent, -) -from graphon.engine_events.node import NodeRunPauseRequestedEvent -from graphon.entities.base_node_data import DefaultValue, DefaultValueType -from graphon.entities.pause_reason import HitlRequired -from graphon.enums import ( - BuiltinNodeTypes, - ErrorStrategy, - WorkflowNodeExecutionMetadataKey, - WorkflowNodeExecutionStatus, -) -from graphon.graph import Graph -from graphon.node_events import StreamChunkEvent, StreamCompletedEvent -from graphon.nodes.container_effects import ( - ContainerExecutionResult, - ContainerNodeRunResult, - CustomContainerRequest, - LoopFrameRequest, - build_container_value, -) -from graphon.nodes.end.end_node import EndNode -from graphon.nodes.end.entities import EndNodeData -from graphon.nodes.start.entities import StartNodeData -from graphon.nodes.start.start_node import StartNode -from graphon.nodes.tool.entities import ToolNodeData, ToolProviderType -from graphon.nodes.tool.exc import ToolNodeError -from graphon.nodes.tool.tool_node import ToolNode -from graphon.nodes.tool_runtime_entities import ToolRuntimeHandle -from graphon.runtime import RuntimeState, VariablePool -from graphon.runtime.container_state import create_container_run_state -from graphon.runtime.execution import ROOT_FRAME_ID -from graphon.variables.segments import ArrayFileSegment -from graphon.variables.types import SegmentType -from tests.workflow_test_utils import build_test_graph_init_params, build_test_run_context - - -def _workflow_tool_node( - runtime_state: RuntimeState | None = None, - *, - app_id: str = "outer-app", - version: str = "1", - tool_node_version: str | None = None, -) -> tuple[DifyWorkflowToolNode, DifyToolNodeRuntime, WorkflowToolContainerPayload]: - graph_config = { - "nodes": [ - { - "id": "tool", - "data": { - "type": "tool", - "title": "Workflow Tool", - "provider_id": "workflow-provider", - "provider_type": "workflow", - "provider_name": "workflow-provider", - "tool_name": "nested-workflow", - "tool_label": "Nested Workflow", - "tool_configurations": dict[str, object](), - "tool_parameters": dict[str, object](), - "version": version, - "tool_node_version": tool_node_version, - }, - } - ], - "edges": list[dict[str, object]](), - } - init_params = build_test_graph_init_params( - workflow_id="outer-workflow", - graph_config=graph_config, - call_depth=2, - app_id=app_id, - ) - if runtime_state is None: - runtime_state = RuntimeState( - workflow_id="outer-workflow", - variable_pool=VariablePool(), - start_at=1, - ) - payload = WorkflowToolContainerPayload( - source_app_id="source-app", - source_workflow_id="source-workflow", - source_workflow_version="1", - inputs={"question": "hello"}, - inputs_for_log={}, - call_depth=3, - ) - tool = WorkflowTool( - workflow_app_id=payload.source_app_id, - workflow_id=payload.source_workflow_id, - workflow_as_tool_id="workflow-provider", - version=payload.source_workflow_version, - workflow_call_depth=2, - entity=ToolEntity( - identity=ToolIdentity( - author="test", - name="nested-workflow", - label=I18nObject(en_US="Nested Workflow"), - provider="workflow-provider", - ), - parameters=[], - description=None, - has_runtime_parameters=False, - ), - runtime=ToolRuntime(tenant_id="tenant", invoke_from=InvokeFrom.DEBUGGER), - ) - tool.prepare_container_inputs = lambda *_args, **_kwargs: (payload.inputs, []) - runtime = DifyToolNodeRuntime(init_params.run_context) - runtime.get_runtime = lambda **_kwargs: ToolRuntimeHandle(raw=tool) - node = DifyWorkflowToolNode( - node_id="tool", - data=ToolNodeData.model_validate(graph_config["nodes"][0]["data"]), - init_params=init_params, - runtime_state=runtime_state, - tool_file_manager=DifyToolFileManager(init_params.run_context), - runtime=runtime, - ) - node.bind_execution_id("tool-execution") - return node, runtime, payload - - -def _outer_graph(node: DifyWorkflowToolNode) -> Graph: - start = StartNode( - node_id="outer-start", - data=StartNodeData(title="Start", variables=[]), - init_params=node.init_params, - runtime_state=node.runtime_state, - ) - return Graph.new().add_root(start).add_node(node, from_node_id=start.id).build() - - -def test_workflow_tool_node_requests_child_container_and_resumes_successfully(mocker: MockerFixture) -> None: - node, runtime, payload = _workflow_tool_node() - get_runtime = mocker.spy(runtime, "get_runtime") - build_payload = mocker.spy(runtime, "build_workflow_tool_container_payload") - - initial_events = list(node._run()) - - assert initial_events == [CustomContainerRequest(payload=payload.model_dump_json())] - get_runtime.assert_called_once_with( - node_id="tool", - node_data=node.node_data, - variable_pool=None, - node_execution_id="tool-execution", - ) - build_payload.assert_called_once_with( - tool_runtime=get_runtime.spy_return, - tool_parameters={}, - inputs_for_log={}, - workflow_call_depth=2, - ) - - result = ContainerExecutionResult( - metadata={}, - steps=2, - node_run_result=ContainerNodeRunResult( - status=WorkflowNodeExecutionStatus.SUCCEEDED, - inputs={"question": build_container_value("hello")}, - outputs={ - "text": build_container_value("done"), - "answer": build_container_value(42), - "files": WorkflowToolContainerHandler._build_tool_outputs({})["files"], - }, - ), - ) - assert result.node_run_result.outputs["files"].value_type == SegmentType.ARRAY_FILE - node.runtime_state.enqueue_ready_task(ResumeTask(invocation_id="invocation", result=result)) - restored_state = RuntimeState.from_snapshot(node.runtime_state.dumps()) - resume_task = restored_state.ready_queue.get(timeout=0.01) - assert isinstance(resume_task, ResumeTask) - assert isinstance(resume_task.result, ContainerExecutionResult) - assert resume_task.result.node_run_result.outputs["files"].value_type == SegmentType.ARRAY_FILE - restored_node, _, _ = _workflow_tool_node(restored_state) - - resumed_events = list(restored_node._resume_container_events(result=resume_task.result)) - - assert resumed_events[:2] == [ - StreamChunkEvent(selector=["tool", "text"], chunk="done", is_final=False), - StreamChunkEvent(selector=["tool", "text"], chunk="", is_final=True), - ] - completed = resumed_events[2] - assert isinstance(completed, StreamCompletedEvent) - assert completed.node_run_result.inputs == {"question": "hello"} - assert {key: value.to_object() for key, value in completed.node_run_result.outputs.items()} == { - "text": "done", - "answer": 42, - "files": [], - } - assert isinstance(completed.node_run_result.outputs["files"], ArrayFileSegment) - for key, value in completed.node_run_result.outputs.items(): - restored_state.variable_pool.add(("tool", key), value) - files = restored_state.variable_pool.get(("tool", "files")) - assert files is not None - assert files.value_type == SegmentType.ARRAY_FILE - assert files.to_object() == [] - assert node.runtime_state.node_run_steps == 0 - assert completed.node_run_result.metadata[WorkflowNodeExecutionMetadataKey.TOOL_INFO] == { - "provider_type": ToolProviderType.WORKFLOW.value, - "provider_id": "workflow-provider", - "plugin_unique_identifier": None, - } - - -def test_workflow_tool_node_returns_failure_when_runtime_cannot_be_loaded(mocker: MockerFixture) -> None: - node, runtime, _ = _workflow_tool_node(version="2") - get_runtime = mocker.patch.object(runtime, "get_runtime", side_effect=ToolNodeError("runtime unavailable")) - - events = list(node._run()) - - assert len(events) == 1 - completed = events[0] - assert isinstance(completed, StreamCompletedEvent) - assert completed.node_run_result.status == WorkflowNodeExecutionStatus.FAILED - assert completed.node_run_result.error == "Failed to get tool runtime: runtime unavailable" - get_runtime.assert_called_once_with( - node_id="tool", - node_data=node.node_data, - variable_pool=node.runtime_state.variable_pool, - node_execution_id="tool-execution", - ) - - -def test_workflow_tool_node_returns_failure_when_payload_cannot_be_built(mocker: MockerFixture) -> None: - node, runtime, _ = _workflow_tool_node(tool_node_version="1") - mocker.patch.object( - runtime, "build_workflow_tool_container_payload", side_effect=ToolNodeError("invalid source workflow") - ) - - events = list(node._run()) - - assert len(events) == 1 - completed = events[0] - assert isinstance(completed, StreamCompletedEvent) - assert completed.node_run_result.status == WorkflowNodeExecutionStatus.FAILED - assert completed.node_run_result.error == "Failed to prepare Workflow Tool: invalid source workflow" - - -def test_workflow_tool_node_rejects_non_execution_result() -> None: - node, _, _ = _workflow_tool_node() - result = LoopFrameRequest( - inputs={}, - outputs={}, - loop_count=0, - root_node_id="loop", - loop_variable_selectors={}, - loop_node_ids=frozenset(), - index=0, - ) - - with pytest.raises(TypeError, match="Unsupported Workflow Tool container result LoopFrameRequest"): - list(node._resume_container_events(result=result)) - - -def test_workflow_tool_node_resumes_failed_result_without_chunks() -> None: - node, _, _ = _workflow_tool_node() - result = ContainerExecutionResult( - metadata={}, - steps=1, - node_run_result=ContainerNodeRunResult( - status=WorkflowNodeExecutionStatus.FAILED, - error="child failed", - ), - ) - - events = list(node._resume_container_events(result=result)) - - assert len(events) == 1 - assert isinstance(events[0], StreamCompletedEvent) - assert events[0].node_run_result.error == "child failed" - - -def test_workflow_tool_node_resumes_empty_text_with_final_chunk_only() -> None: - node, _, _ = _workflow_tool_node() - result = ContainerExecutionResult( - metadata={}, - steps=0, - node_run_result=ContainerNodeRunResult( - status=WorkflowNodeExecutionStatus.SUCCEEDED, - outputs={"text": build_container_value("")}, - ), - ) - - events = list(node._resume_container_events(result=result)) - - assert events[:1] == [StreamChunkEvent(selector=["tool", "text"], chunk="", is_final=True)] - assert isinstance(events[1], StreamCompletedEvent) - - -def test_node_factory_can_keep_workflow_tool_direct_for_single_step_debug() -> None: - node, _, _ = _workflow_tool_node() - factory = DifyNodeFactory(node.init_params, node.runtime_state, use_workflow_tool_containers=True) - assert ( - factory._resolve_node_class_for_factory( - node_type=BuiltinNodeTypes.TOOL, - node_version="1", - node_data=node.node_data, - ) - is DifyWorkflowToolNode - ) - - factory._use_workflow_tool_containers = False - assert ( - factory._resolve_node_class_for_factory( - node_type=BuiltinNodeTypes.TOOL, - node_version="1", - node_data=node.node_data, - ) - is ToolNode - ) - - -def _workflow_tool_source() -> WorkflowToolSource: - source_graph = { - "nodes": [ - { - "id": "source-start", - "data": { - "type": "start", - "title": "Start", - "variables": [ - { - "variable": "answer", - "label": "Answer", - "type": "text-input", - "required": False, - "default": "fallback", - } - ], - }, - }, - { - "id": "source-end", - "data": { - "type": "end", - "title": "End", - "outputs": [ - { - "variable": "answer", - "value_selector": ["source-start", "answer"], - } - ], - }, - }, - ], - "edges": [ - { - "id": "source-edge", - "source": "source-start", - "target": "source-end", - } - ], - } - return WorkflowToolSource( - app_id="source-app", - workflow_id="source-workflow", - graph_config=source_graph, - features_dict={}, - environment_variables=[], - workflow_kind=WorkflowKind.STANDARD, - ) - - -def _human_input_workflow_tool_source() -> WorkflowToolSource: - human_input = HumanInputNodeData( - title="Approval", - form_content="Approve this run?", - user_actions=[UserActionConfig(id="approve", title="Approve")], - ) - source_graph = { - "nodes": [ - { - "id": "source-start", - "data": {"type": "start", "title": "Start", "variables": list[object]()}, - }, - { - "id": "source-human", - "data": human_input.model_dump(mode="json"), - }, - { - "id": "source-end", - "data": { - "type": "end", - "title": "End", - "outputs": [ - { - "variable": "decision", - "value_selector": ["source-human", "__action_id"], - } - ], - }, - }, - ], - "edges": [ - { - "id": "source-edge-start-human", - "source": "source-start", - "target": "source-human", - }, - { - "id": "source-edge-human-end", - "source": "source-human", - "sourceHandle": "approve", - "target": "source-end", - }, - ], - } - return replace(_workflow_tool_source(), graph_config=source_graph) - - -def _container_handler( - *, - inputs: dict[str, object] | None = None, -) -> tuple[ - WorkflowToolContainerHandler, - FrameRegistry, - RuntimeState, - CustomContainerRequest, - dict[str, WorkflowToolSource], -]: - source = _workflow_tool_source() - parent_pool = VariablePool() - parent_pool.add(system_variable_selector(SystemVariableKey.USER_ID), "user") - parent_pool.add(system_variable_selector(SystemVariableKey.WORKFLOW_EXECUTION_ID), "outer-execution") - parent_pool.add(("env", "outer-secret"), "must-not-leak") - parent_pool.add(("source-start", "stale"), "must-not-leak") - parent_state = RuntimeState( - workflow_id="outer-workflow", - variable_pool=parent_pool, - start_at=1, - ) - parent_node, _, _ = _workflow_tool_node(parent_state) - parent_graph = _outer_graph(parent_node) - frame_registry = FrameRegistry() - frame_registry.register( - ExecutionFrame( - frame_id=ROOT_FRAME_ID, - graph=parent_graph, - state=parent_state, - scheduler=Scheduler(parent_graph, parent_state, ROOT_FRAME_ID), - failure_handler=NodeFailureHandler(parent_graph, parent_state.graph_execution), - ) - ) - source_repository = {} - source_repository["source-workflow"] = source - payload = WorkflowToolContainerPayload( - source_app_id=source.app_id, - source_workflow_id=source.workflow_id, - source_workflow_version="1", - inputs={"answer": "ok"} if inputs is None else inputs, - inputs_for_log={"answer": "ok"} if inputs is None else inputs, - call_depth=1, - ) - request = CustomContainerRequest(payload=payload.model_dump_json()) - parent_state.put_container_run( - create_container_run_state( - invocation_id="invocation", - frame_id=ROOT_FRAME_ID, - node_id="tool", - started_at=datetime.now(UTC).replace(tzinfo=None), - request=request, - ) - ) - return ( - WorkflowToolContainerHandler(frame_registry, sources=source_repository), - frame_registry, - parent_state, - request, - source_repository, - ) - - -def _dispatch_next_node( - *, - expected_node_id: str, - runtime_state: RuntimeState, - frame_registry: FrameRegistry, - processor: NodeEventProcessor, -) -> None: - task = runtime_state.ready_queue.get(timeout=0.01) - assert isinstance(task, StartTask) - assert task.node_id == expected_node_id - frame = frame_registry[task.frame_id] - node = frame.graph.nodes[task.node_id] - execution = runtime_state.graph_execution.get_or_create_node_execution( - frame_id=frame.frame_id, - node_id=node.id, - ) - node.bind_execution_id(execution.execution_id) - for event in node.run(): - assert isinstance(event, NodeEvent) - processor.dispatch(NodeEventTask(frame_id=frame.frame_id, event=event)) - - -def test_workflow_tool_handler_runs_child_graph_with_internal_name_collision() -> None: - handler, frame_registry, runtime_state, request, _ = _container_handler() - handler.handle_request(invocation_id="invocation", request=request) - frame_registry["invocation:workflow-tool"].state.variable_pool.add( - ("__workflow_tool_container__", "failure"), - {"legitimate": True}, - ) - event_stream = EventStream(layers=[]) - processor = NodeEventProcessor( - graph_execution=runtime_state.graph_execution, - event_stream=event_stream, - frame_registry=frame_registry, - container_handlers={BuiltinNodeTypes.TOOL: handler}, - ) - - _dispatch_next_node( - expected_node_id="source-start", - runtime_state=runtime_state, - frame_registry=frame_registry, - processor=processor, - ) - _dispatch_next_node( - expected_node_id="source-end", - runtime_state=runtime_state, - frame_registry=frame_registry, - processor=processor, - ) - - resume_task = runtime_state.ready_queue.get(timeout=0.01) - assert isinstance(resume_task, ResumeTask) - assert isinstance(resume_task.result, ContainerExecutionResult) - assert resume_task.result.steps == 2 - assert resume_task.result.node_run_result.outputs["files"].value_type == SegmentType.ARRAY_FILE - outputs = {key: value.to_object() for key, value in resume_task.result.node_run_result.outputs.items()} - assert outputs == { - "answer": "ok", - "text": '{"answer": "ok"}', - "files": [], - "json": [{"answer": "ok"}], - } - event_stream.mark_complete() - assert not any(isinstance(event, NodeEvent) for event in event_stream.emit_events()) - with pytest.raises(KeyError): - frame_registry["invocation:workflow-tool"] - with pytest.raises(KeyError): - runtime_state.get_container_frame("invocation:workflow-tool") - - -def test_workflow_tool_handler_restores_child_failure() -> None: - handler, frame_registry, runtime_state, request, source_repository = _container_handler() - handler.handle_request(invocation_id="invocation", request=request) - frame_id = "invocation:workflow-tool" - frame = frame_registry[frame_id] - now = datetime.now(UTC).replace(tzinfo=None) - handler.record_frame_failure( - frame=frame, - event=NodeRunFailedEvent( - id="failed-execution", - node_id="source-start", - node_type=BuiltinNodeTypes.START, - error="source failed", - start_at=now, - finished_at=now, - node_run_result=ContainerNodeRunResult( - status=WorkflowNodeExecutionStatus.FAILED, - error="source failed", - error_type="RuntimeError", - ).to_node_run_result(), - ), - ) - frame_state = runtime_state.get_container_frame(frame_id) - frame_state = frame_state.model_copy(update={"runtime_data": frame.state.snapshot_frame()}) - runtime_state.put_container_frame(frame_state) - frame_registry.remove(frame_id) - - restored_handler = WorkflowToolContainerHandler(frame_registry, sources=source_repository) - restored_handler.restore_frame(frame_state) - restored_handler.complete_frame_if_ready(frame_registry[frame_id]) - - tasks = [runtime_state.ready_queue.get(timeout=0.01) for _ in range(2)] - resume_task = next(task for task in tasks if isinstance(task, ResumeTask)) - assert isinstance(resume_task.result, ContainerExecutionResult) - assert resume_task.result.node_run_result.status == WorkflowNodeExecutionStatus.FAILED - assert {key: value.to_object() for key, value in resume_task.result.node_run_result.inputs.items()} == { - "answer": "ok" - } - assert resume_task.result.node_run_result.error == "source failed" - assert resume_task.result.node_run_result.error_type == "RuntimeError" - - -def test_workflow_tool_handler_isolates_child_variable_pool() -> None: - handler, frame_registry, _, request, _ = _container_handler() - - handler.handle_request(invocation_id="invocation", request=request) - - child_pool = frame_registry["invocation:workflow-tool"].state.variable_pool - assert child_pool.get(system_variable_selector(SystemVariableKey.USER_ID)).to_object() == "user" # type: ignore[union-attr] - assert child_pool.get(system_variable_selector(SystemVariableKey.WORKFLOW_EXECUTION_ID)).to_object() == ( # type: ignore[union-attr] - "outer-execution" - ) - assert child_pool.get(system_variable_selector(SystemVariableKey.APP_ID)).to_object() == ( # type: ignore[union-attr] - "source-app" - ) - assert child_pool.get(("env", "outer-secret")) is None - assert child_pool.get(("source-start", "stale")) is None - - -def test_workflow_tool_handler_applies_start_input_defaults() -> None: - handler, frame_registry, _, request, _ = _container_handler(inputs={}) - - handler.handle_request(invocation_id="invocation", request=request) - - answer = frame_registry["invocation:workflow-tool"].state.variable_pool.get(("source-start", "answer")) - assert answer is not None - assert answer.to_object() == "fallback" - - -def test_workflow_tool_handler_preserves_inputs_when_start_validation_fails() -> None: - handler, _, runtime_state, request, _ = _container_handler(inputs={"answer": 42}) - - handler.handle_request(invocation_id="invocation", request=request) - - resume_task = runtime_state.ready_queue.get(timeout=0.01) - assert isinstance(resume_task, ResumeTask) - assert isinstance(resume_task.result, ContainerExecutionResult) - assert resume_task.result.node_run_result.status == WorkflowNodeExecutionStatus.FAILED - assert {key: value.to_object() for key, value in resume_task.result.node_run_result.inputs.items()} == { - "answer": 42 - } - assert resume_task.result.node_run_result.error_type == "ValueError" - - -def test_workflow_tool_nested_handler_hides_and_persists_marked_child_events() -> None: - _, frame_registry, _, request, repository = _container_handler() - persisted: list[NodeEvent] = [] - event_listeners: dict[str, Callable[[NodeEvent], None]] = {} - workflow_tool_handler = WorkflowToolContainerHandler( - frame_registry, - sources=repository, - event_listener_factory=lambda *_: persisted.append, - event_listeners=event_listeners, - ) - workflow_tool_handler.handle_request(invocation_id="invocation", request=request) - built_in_handler = LoopContainerHandler(frame_registry) - hidden_events: list[NodeEvent] = [] - nested_handler = WorkflowToolNestedContainerHandler( - frame_registry, - handler_factory=lambda _: built_in_handler, - hidden_event_listener=hidden_events.append, - event_listeners=event_listeners, - ) - event = NodeRunSucceededEvent( - id="source-execution", - node_id="source-start", - node_type=BuiltinNodeTypes.START, - start_at=datetime.now(UTC).replace(tzinfo=None), - container_id="nested-container", - ) - - workflow_tool_handler.prepare_frame_event( - frame=frame_registry["invocation:workflow-tool"], - event=event, - ) - - assert nested_handler.should_emit(event=event) is False - assert hidden_events == [event] - assert len(persisted) == 1 - assert persisted[0].container_id == "nested-container" - assert persisted[0].node_run_result.process_data["workflow_tool_invocation_id"] == "invocation" - assert "__dify_workflow_tool_child__" not in persisted[0].node_run_result.process_data - - # Another Workflow Tool inside this container owns its own source events. - inner_events: list[NodeEvent] = [] - event_listeners["inner-tool-frame"] = inner_events.append - inner_event = event.model_copy(deep=True) - inner_event.node_run_result.process_data = {} - workflow_tool_handler.prepare_frame_event( - frame=replace(frame_registry["invocation:workflow-tool"], frame_id="inner-tool-frame"), event=inner_event - ) - workflow_tool_handler.prepare_frame_event(frame=frame_registry["invocation:workflow-tool"], event=inner_event) - assert nested_handler.should_emit(event=inner_event) is False - assert inner_events == [inner_event] - assert len(persisted) == 1 - - unmarked_event = event.model_copy(deep=True) - unmarked_event.node_run_result.process_data = {} - assert nested_handler.should_emit(event=unmarked_event) is True - assert len(persisted) == 1 - - -@pytest.mark.parametrize( - "source_metadata", - [{}, {WorkflowNodeExecutionMetadataKey.LOOP_ID: "source-loop", WorkflowNodeExecutionMetadataKey.LOOP_INDEX: 1}], -) -def test_workflow_tool_persistence_keeps_source_container_metadata( - source_metadata: dict[WorkflowNodeExecutionMetadataKey, str | int], -) -> None: - _, frames, _, request, repository = _container_handler() - persisted: list[NodeEvent] = [] - handler = WorkflowToolContainerHandler( - frames, sources=repository, event_listener_factory=lambda *_: persisted.append - ) - handler.handle_request(invocation_id="invocation", request=request) - event = NodeRunSucceededEvent( - id="source-execution", - node_id="source-start", - node_type=BuiltinNodeTypes.START, - start_at=datetime.now(UTC).replace(tzinfo=None), - ) - event.node_run_result.metadata = source_metadata.copy() - handler.prepare_frame_event(frame=frames["invocation:workflow-tool"], event=event) - # Ancestors outside the Tool source add their own ownership before collection. - event.node_run_result.metadata.update( - { - WorkflowNodeExecutionMetadataKey.ITERATION_ID: "outer-iteration", - WorkflowNodeExecutionMetadataKey.ITERATION_INDEX: 2, - } - ) - event.node_run_result.outputs = {"normalized": True} - - assert handler.should_emit(event=event) is False - - assert persisted[0].node_run_result.metadata == source_metadata - assert persisted[0].node_run_result.outputs == {"normalized": True} - - -@pytest.mark.parametrize( - ("child_error_strategy", "tool_error_strategy", "terminal_type", "expected_count", "expected_output", "retry_once"), - [ - pytest.param( - ErrorStrategy.DEFAULT_VALUE, - None, - GraphRunSucceededEvent, - 0, - '{"answer": "child fallback"}', - False, - id="child-handles-failure", - ), - pytest.param(None, None, GraphRunFailedEvent, 1, None, False, id="unhandled-failure"), - pytest.param( - None, - ErrorStrategy.DEFAULT_VALUE, - GraphRunPartialSucceededEvent, - 1, - "tool fallback", - False, - id="tool-default-value", - ), - pytest.param( - None, - ErrorStrategy.FAIL_BRANCH, - GraphRunPartialSucceededEvent, - 1, - "source code failed", - False, - id="tool-fail-branch", - ), - pytest.param(None, None, GraphRunSucceededEvent, 0, '{"answer": "retried output"}', True, id="child-retries"), - ], -) -def test_workflow_tool_failure_accounting_uses_outer_tool_policy( - monkeypatch: pytest.MonkeyPatch, - child_error_strategy: ErrorStrategy | None, - tool_error_strategy: ErrorStrategy | None, - terminal_type: type[GraphRunSucceededEvent | GraphRunFailedEvent | GraphRunPartialSucceededEvent], - expected_count: int, - expected_output: str | None, - retry_once: bool, -) -> None: - attempts = 0 - - def execute(**_kwargs: object) -> dict[str, str]: - nonlocal attempts - attempts += 1 - if attempts == 1: - raise CodeExecutionError("source code failed") - return {"output": "retried output"} - - monkeypatch.setattr("core.workflow.node_factory.CodeExecutor.execute_workflow_code_template", execute) - source = _workflow_tool_source() - source_graph = dict(source.graph_config) - source_graph["nodes"].insert( - 1, - { - "id": "source-code", - "data": { - "type": "code", - "title": "Code", - "code_language": "python3", - "code": "def main(): return {'output': 'retried output'}", - "outputs": {"output": {"type": "string"}}, - "variables": [], - "error_strategy": child_error_strategy, - "default_value": [{"key": "output", "type": "string", "value": "child fallback"}], - "retry_config": {"retry_enabled": retry_once, "max_retries": 1, "retry_interval": 0}, - }, - }, - ) - source_graph["nodes"][-1]["data"]["outputs"][0]["value_selector"] = ["source-code", "output"] - source_graph["edges"] = [ - {"id": "start-code", "source": "source-start", "target": "source-code"}, - {"id": "code-end", "source": "source-code", "target": "source-end"}, - ] - repository = {} - repository["source-workflow"] = replace(source, graph_config=source_graph) - node, _, _ = _workflow_tool_node() - node.node_data.error_strategy = tool_error_strategy - node.node_data.default_value = [DefaultValue(key="text", type=DefaultValueType.STRING, value="tool fallback")] - start = StartNode( - node_id="outer-start", - data=StartNodeData(title="Start", variables=[]), - init_params=node.init_params, - runtime_state=node.runtime_state, - ) - end = EndNode( - node_id="outer-end", - data=EndNodeData.model_validate( - { - "title": "End", - "outputs": [ - { - "variable": "answer", - "value_selector": [ - "tool", - "error_message" if tool_error_strategy == ErrorStrategy.FAIL_BRANCH else "text", - ], - } - ], - } - ), - init_params=node.init_params, - runtime_state=node.runtime_state, - ) - graph = ( - Graph.new() - .add_root(start) - .add_node(node) - .add_node(end, source_handle="fail-branch" if tool_error_strategy == ErrorStrategy.FAIL_BRANCH else "source") - .build() - ) - persisted: list[NodeEvent] = [] - engine = Engine( - graph=graph, - runtime_state=node.runtime_state, - command_channel=InMemoryChannel(), - workers=1, - container_handler_factories=( - partial( - WorkflowToolContainerHandler, - sources=repository, - event_listener_factory=lambda *_: persisted.append, - ), - ), - ) - events: list[EngineEvent] = [] - if terminal_type is GraphRunFailedEvent: - with pytest.raises(RuntimeError, match="source code failed"): - events.extend(engine.run()) - else: - events.extend(engine.run()) - - assert attempts == (2 if retry_once else 1) - assert isinstance(events[-1], terminal_type) - assert node.runtime_state.graph_execution.exceptions_count == expected_count - if isinstance(events[-1], GraphRunFailedEvent): - assert events[-1].error == "source code failed" - assert not any(isinstance(event, NodeRunSucceededEvent) and event.node_id == end.id for event in events) - else: - assert isinstance(events[-1], GraphRunSucceededEvent | GraphRunPartialSucceededEvent) - assert events[-1].outputs == {"answer": expected_output} - if tool_error_strategy is not None: - assert any(isinstance(event, NodeRunExceptionEvent) and event.node_id == node.id for event in events) - if isinstance(events[-1], GraphRunFailedEvent | GraphRunPartialSucceededEvent): - assert events[-1].exceptions_count == expected_count - if child_error_strategy is not None: - assert any(isinstance(event, NodeRunExceptionEvent) and event.node_id == "source-code" for event in persisted) - if retry_once: - assert [event.retry_index for event in persisted if isinstance(event, NodeRunRetryEvent)] == [1] - assert any(isinstance(event, NodeRunSucceededEvent) and event.node_id == "source-code" for event in persisted) - - -@pytest.mark.parametrize("outcome", ["paused", "resumed", "abort-on-resume", "stopped"]) -def test_workflow_tool_human_input_pauses_and_resumes_without_duplicate_form( - monkeypatch: pytest.MonkeyPatch, - mocker: MockerFixture, - outcome: str, -) -> None: - command_channel = InMemoryChannel() - - def stop_on_human_input(event: NodeEvent, *, channel: InMemoryChannel, should_stop: bool) -> None: - if should_stop and isinstance(event, NodeRunPauseRequestedEvent): - channel.send_command(AbortCommand()) - - source_repository = {"source-workflow": _human_input_workflow_tool_source()} - handler_factory = partial( - WorkflowToolContainerHandler, - sources=source_repository, - hidden_event_listener=partial(stop_on_human_input, channel=command_channel, should_stop=outcome == "stopped"), - ) - form_repository = HumanInputFormRepositoryImpl( - tenant_id="tenant", app_id="outer-app", workflow_execution_id="outer-execution" - ) - create_form_spy = mocker.spy(form_repository, "create_form") - human_input_app_ids: list[str] = [] - - def build_human_input_callback( - factory: DifyNodeFactory, - *, - node_data: HumanInputNodeData, - execution_id_getter: Callable[[], str | None], - ) -> DifyHITLCallback: - human_input_app_ids.append(factory._human_input_runtime._run_context.app_id) - return DifyHITLCallback( - form_repository=form_repository, - node_data=node_data, - execution_id_getter=execution_id_getter, - ) - - monkeypatch.setattr(DifyNodeFactory, "_build_human_input_callback", build_human_input_callback) - - outer_pool = VariablePool() - outer_pool.add(system_variable_selector(SystemVariableKey.WORKFLOW_EXECUTION_ID), "outer-execution") - outer_pool.add(system_variable_selector(SystemVariableKey.APP_ID), "outer-app") - outer_pool.add(system_variable_selector(SystemVariableKey.WORKFLOW_ID), "outer-workflow") - initial_state = RuntimeState( - workflow_id="outer-workflow", - variable_pool=outer_pool, - start_at=1, - ) - initial_node, _, _ = _workflow_tool_node(initial_state, app_id="intermediate-app") - initial_graph = _outer_graph(initial_node) - initial_owner_factory = DifyNodeFactory( - initial_node.init_params, initial_state, human_input_run_context=build_test_run_context(app_id="outer-app") - ) - initial_graph.node_factory = initial_owner_factory - initial_events = list( - Engine( - graph=initial_graph, - runtime_state=initial_state, - command_channel=command_channel, - workers=1, - container_handler_factories=(handler_factory,), - ).run() - ) - - paused = initial_events[-1] - if outcome == "stopped": - assert isinstance(paused, GraphRunAbortedEvent) - return - assert isinstance(paused, GraphRunPausedEvent) - assert len(paused.reasons) == 1 - reason = paused.reasons[0] - assert isinstance(reason, HitlRequired) - assert reason.node_id == "source-human" - assert reason.node_title == "Approval" - assert all( - not isinstance(event, NodeEvent) or event.node_id not in {"source-start", "source-human"} - for event in initial_events - ) - create_form_spy.assert_called_once() - create_params = create_form_spy.call_args.args[0] - assert create_params.node_id == "source-human" - assert create_params.workflow_execution_id == "outer-execution" - child_frames = list(initial_state.container_frames()) - assert len(child_frames) == 1 - child_execution = initial_state.graph_execution.node_executions[(child_frames[0].frame_id, "source-human")] - assert create_params.form_id == child_execution.execution_id - assert human_input_app_ids == ["outer-app"] - - if outcome == "paused": - return - - restored_state = RuntimeState.from_snapshot(initial_state.dumps()) - assert ( - resolve_human_input_node_id( - node_id=reason.node_id, - form_id=reason.session_id, - variable_pool=restored_state.variable_pool, - ) - == "tool" - ) - restored_node, _, _ = _workflow_tool_node(restored_state, app_id="intermediate-app") - restored_graph = _outer_graph(restored_node) - restored_owner_factory = DifyNodeFactory( - restored_node.init_params, restored_state, human_input_run_context=initial_owner_factory.human_input_run_context - ) - restored_graph.node_factory = restored_owner_factory - if outcome == "resumed": - HumanInputFormSubmissionRepository().mark_submitted( - form_id=reason.session_id, - recipient_id=None, - selected_action_id="approve", - form_data={}, - submission_user_id="user", - submission_end_user_id=None, - ) - resumed_channel = InMemoryChannel() - resumed_handler_factory = partial( - WorkflowToolContainerHandler, - sources=source_repository, - hidden_event_listener=partial( - stop_on_human_input, channel=resumed_channel, should_stop=outcome == "abort-on-resume" - ), - ) - resumed_events = list( - Engine( - graph=restored_graph, - runtime_state=restored_state, - command_channel=resumed_channel, - workers=1, - container_handler_factories=(resumed_handler_factory,), - ).run() - ) - - if outcome == "abort-on-resume": - assert isinstance(resumed_events[-1], GraphRunAbortedEvent) - return - assert isinstance(resumed_events[-1], GraphRunSucceededEvent) - tool_succeeded = next( - event for event in resumed_events if isinstance(event, NodeRunSucceededEvent) and event.node_id == "tool" - ) - assert tool_succeeded.node_run_result.outputs["decision"].to_object() == "approve" - assert json.loads(tool_succeeded.node_run_result.outputs["text"].to_object()) == {"decision": "approve"} - files = restored_state.variable_pool.get(("tool", "files")) - assert files is not None - assert files.value_type == SegmentType.ARRAY_FILE - assert files.to_object() == [] - create_form_spy.assert_called_once() - assert human_input_app_ids == ["outer-app", "outer-app"] - assert all( - not isinstance(event, NodeEvent) or event.node_id not in {"source-human", "source-end"} - for event in resumed_events - ) - assert list(restored_state.container_frames()) == [] - assert list(restored_state.container_runs()) == [] diff --git a/api/tests/unit_tests/core/workflow/test_workflow_tool_container_handler.py b/api/tests/unit_tests/core/workflow/test_workflow_tool_container_handler.py new file mode 100644 index 00000000000..cd5695f2167 --- /dev/null +++ b/api/tests/unit_tests/core/workflow/test_workflow_tool_container_handler.py @@ -0,0 +1,2004 @@ +from __future__ import annotations + +import json +from collections.abc import Callable, Generator +from contextlib import AbstractContextManager +from contextvars import ContextVar +from dataclasses import replace +from datetime import UTC, datetime +from functools import partial +from threading import Event, get_ident +from types import TracebackType +from typing import override + +import pytest +from flask import current_app, g +from opentelemetry.sdk.trace import TracerProvider +from opentelemetry.sdk.trace.export.in_memory_span_exporter import InMemorySpanExporter +from opentelemetry.trace import get_current_span +from pytest_mock import MockerFixture +from sqlalchemy import select +from sqlalchemy.orm import Session, sessionmaker + +from clients.agent_backend import FakeAgentBackendRunClient +from context import capture_current_context +from context.flask_app_context import capture_flask_context +from core.app.apps.workflow.app_config_manager import WorkflowAppConfig +from core.app.entities.app_invoke_entities import InvokeFrom, UserFrom, WorkflowAppGenerateEntity +from core.app.file_access import FileAccessScope, bind_file_access_scope +from core.app.layers import execution_context_layer as module +from core.app.layers.execution_context_layer import ExecutionContextLayer +from core.app.workflow.file_runtime import init_app +from core.app.workflow.layers.observability import ObservabilityLayer +from core.helper.code_executor.code_executor import CodeExecutionError +from core.repositories.human_input_repository import ( + HumanInputFormRepositoryImpl, + HumanInputFormSubmissionRepository, +) +from core.repositories.sqlalchemy_workflow_execution_repository import SQLAlchemyWorkflowExecutionRepository +from core.repositories.sqlalchemy_workflow_node_execution_repository import SQLAlchemyWorkflowNodeExecutionRepository +from core.tools.__base.tool_runtime import ToolRuntime +from core.tools.entities.common_entities import I18nObject +from core.tools.entities.tool_entities import ToolEntity, ToolIdentity +from core.tools.workflow_as_tool.repository import WorkflowToolSource +from core.tools.workflow_as_tool.tool import WorkflowTool +from core.workflow.node_execution_process_data import ( + WORKFLOW_TOOL_INVOCATION_ID_KEY, +) +from core.workflow.node_factory import DifyNodeFactory +from core.workflow.node_runtime import DifyToolFileManager, DifyToolNodeRuntime +from core.workflow.nodes.agent_v2.binding_resolver import WorkflowAgentBindingResolver +from core.workflow.nodes.human_input.boundary import resolve_human_input_node_id +from core.workflow.nodes.human_input.callback import DifyHITLCallback +from core.workflow.nodes.human_input.entities import HumanInputNodeData, UserActionConfig +from core.workflow.system_variables import SystemVariableKey, system_variable_selector +from core.workflow.workflow_entry import WorkflowEntry +from core.workflow.workflow_tool_container_handler import ( + WorkflowToolContainerHandler, + WorkflowToolNestedContainerHandler, +) +from core.workflow.workflow_tool_container_types import WorkflowToolContainerPayload +from core.workflow.workflow_tool_node import DifyWorkflowToolNode +from dify_app import DifyApp +from enums import WorkflowKind +from extensions.storage.storage_type import StorageType +from graphon.engine import Engine +from graphon.engine.command import AbortCommand, InMemoryChannel +from graphon.engine.container_handler.builtin.loop import LoopContainerHandler +from graphon.engine.event.node_failure import NodeFailureHandler +from graphon.engine.event.processor import NodeEventProcessor +from graphon.engine.event.stream import EventStream +from graphon.engine.frame import ExecutionFrame, FrameRegistry +from graphon.engine.layer import Layer +from graphon.engine.ready_queue import ResumeTask, StartTask +from graphon.engine.scheduler import Scheduler +from graphon.engine.worker import NodeEventTask +from graphon.engine_events import ( + GraphRunSucceededEvent as _workflow_tool_event_delivery_GraphRunSucceededEvent, +) +from graphon.engine_events import GraphRunSucceededEvent as _workflow_tool_file_access_GraphRunSucceededEvent +from graphon.engine_events import GraphRunSucceededEvent as _workflow_tool_observability_GraphRunSucceededEvent +from graphon.engine_events import NodeEvent as _workflow_tool_observability_NodeEvent +from graphon.engine_events import ( + NodeRunExceptionEvent, + NodeRunFailedEvent, + NodeRunRetryEvent, + NodeRunStartedEvent, + NodeRunStreamChunkEvent, + NodeRunSucceededEvent, +) +from graphon.engine_events.base import EngineEvent +from graphon.engine_events.base import NodeEvent as _workflow_tool_container_NodeEvent +from graphon.engine_events.base import NodeEvent as _workflow_tool_event_delivery_NodeEvent +from graphon.engine_events.graph import ( + GraphRunAbortedEvent, + GraphRunFailedEvent, + GraphRunPartialSucceededEvent, + GraphRunPausedEvent, +) +from graphon.engine_events.graph import ( + GraphRunSucceededEvent as _workflow_tool_container_GraphRunSucceededEvent, +) +from graphon.engine_events.node import NodeRunPauseRequestedEvent +from graphon.entities import WorkflowNodeExecution +from graphon.entities.base_node_data import BaseNodeData, DefaultValue, DefaultValueType +from graphon.entities.pause_reason import HitlRequired +from graphon.enums import ( + BuiltinNodeTypes, + ErrorStrategy, + WorkflowNodeExecutionMetadataKey, + WorkflowNodeExecutionStatus, + WorkflowType, +) +from graphon.graph import Graph +from graphon.node_events import NodeRunResult, StreamChunkEvent, StreamCompletedEvent +from graphon.node_events.base import NodeEventPayload +from graphon.nodes.base.node import Node +from graphon.nodes.container_effects import ( + ContainerAwaitRequest, + ContainerExecutionResult, + ContainerNodeRunResult, + ContainerRunResult, + CustomContainerRequest, + LoopFrameRequest, + build_container_value, +) +from graphon.nodes.end.end_node import EndNode +from graphon.nodes.end.entities import EndNodeData +from graphon.nodes.start.entities import StartNodeData +from graphon.nodes.start.start_node import StartNode +from graphon.nodes.tool.entities import ToolNodeData, ToolProviderType +from graphon.nodes.tool.exc import ToolNodeError +from graphon.nodes.tool.tool_node import ToolNode +from graphon.nodes.tool_runtime_entities import ToolRuntimeHandle +from graphon.runtime import RuntimeState, VariablePool +from graphon.runtime.container_state import create_container_run_state +from graphon.runtime.execution import ROOT_FRAME_ID +from graphon.variables.segments import ArrayFileSegment +from graphon.variables.types import SegmentType +from models import Account, ToolFile, UploadFile, WorkflowRun +from models.enums import CreatorUserRole, WorkflowRunTriggeredFrom +from models.model import AppMode +from models.workflow import WorkflowNodeExecutionModel, WorkflowNodeExecutionTriggeredFrom +from services.workflow_persistence import PersistenceWorkflowInfo, WorkflowPersistenceLayer +from services.workflow_run_index import WorkflowRunIndex +from tests.helpers.unit.config_override import apply_config_overrides +from tests.workflow_test_utils import build_test_graph_init_params, build_test_run_context + + +def _workflow_tool_container_workflow_tool_node( + runtime_state: RuntimeState | None = None, + *, + app_id: str = "outer-app", + version: str = "1", + tool_node_version: str | None = None, +) -> tuple[DifyWorkflowToolNode, DifyToolNodeRuntime, WorkflowToolContainerPayload]: + graph_config = { + "nodes": [ + { + "id": "tool", + "data": { + "type": "tool", + "title": "Workflow Tool", + "provider_id": "workflow-provider", + "provider_type": "workflow", + "provider_name": "workflow-provider", + "tool_name": "nested-workflow", + "tool_label": "Nested Workflow", + "tool_configurations": dict[str, object](), + "tool_parameters": dict[str, object](), + "version": version, + "tool_node_version": tool_node_version, + }, + } + ], + "edges": list[dict[str, object]](), + } + init_params = build_test_graph_init_params( + workflow_id="outer-workflow", + graph_config=graph_config, + call_depth=2, + app_id=app_id, + ) + if runtime_state is None: + runtime_state = RuntimeState( + workflow_id="outer-workflow", + variable_pool=VariablePool(), + start_at=1, + ) + payload = WorkflowToolContainerPayload( + source_app_id="source-app", + source_workflow_id="source-workflow", + source_workflow_version="1", + inputs={"question": "hello"}, + inputs_for_log={}, + call_depth=3, + ) + tool = WorkflowTool( + workflow_app_id=payload.source_app_id, + workflow_id=payload.source_workflow_id, + workflow_as_tool_id="workflow-provider", + version=payload.source_workflow_version, + workflow_call_depth=2, + entity=ToolEntity( + identity=ToolIdentity( + author="test", + name="nested-workflow", + label=I18nObject(en_US="Nested Workflow"), + provider="workflow-provider", + ), + parameters=[], + description=None, + has_runtime_parameters=False, + ), + runtime=ToolRuntime(tenant_id="tenant", invoke_from=InvokeFrom.DEBUGGER), + ) + tool.prepare_container_inputs = lambda *_args, **_kwargs: (payload.inputs, []) + runtime = DifyToolNodeRuntime(init_params.run_context) + runtime.get_runtime = lambda **_kwargs: ToolRuntimeHandle(raw=tool) + node = DifyWorkflowToolNode( + node_id="tool", + data=ToolNodeData.model_validate(graph_config["nodes"][0]["data"]), + init_params=init_params, + runtime_state=runtime_state, + tool_file_manager=DifyToolFileManager(init_params.run_context), + runtime=runtime, + ) + node.bind_execution_id("tool-execution") + return node, runtime, payload + + +def _workflow_tool_container_outer_graph(node: DifyWorkflowToolNode) -> Graph: + start = StartNode( + node_id="outer-start", + data=StartNodeData(title="Start", variables=[]), + init_params=node.init_params, + runtime_state=node.runtime_state, + ) + return Graph.new().add_root(start).add_node(node, from_node_id=start.id).build() + + +def _workflow_tool_container_workflow_tool_source() -> WorkflowToolSource: + source_graph = { + "nodes": [ + { + "id": "source-start", + "data": { + "type": "start", + "title": "Start", + "variables": [ + { + "variable": "answer", + "label": "Answer", + "type": "text-input", + "required": False, + "default": "fallback", + } + ], + }, + }, + { + "id": "source-end", + "data": { + "type": "end", + "title": "End", + "outputs": [ + { + "variable": "answer", + "value_selector": ["source-start", "answer"], + } + ], + }, + }, + ], + "edges": [ + { + "id": "source-edge", + "source": "source-start", + "target": "source-end", + } + ], + } + return WorkflowToolSource( + app_id="source-app", + workflow_id="source-workflow", + graph_config=source_graph, + features_dict={}, + environment_variables=[], + workflow_kind=WorkflowKind.STANDARD, + ) + + +def _human_input_workflow_tool_source() -> WorkflowToolSource: + human_input = HumanInputNodeData( + title="Approval", + form_content="Approve this run?", + user_actions=[UserActionConfig(id="approve", title="Approve")], + ) + source_graph = { + "nodes": [ + { + "id": "source-start", + "data": {"type": "start", "title": "Start", "variables": list[object]()}, + }, + { + "id": "source-human", + "data": human_input.model_dump(mode="json"), + }, + { + "id": "source-end", + "data": { + "type": "end", + "title": "End", + "outputs": [ + { + "variable": "decision", + "value_selector": ["source-human", "__action_id"], + } + ], + }, + }, + ], + "edges": [ + { + "id": "source-edge-start-human", + "source": "source-start", + "target": "source-human", + }, + { + "id": "source-edge-human-end", + "source": "source-human", + "sourceHandle": "approve", + "target": "source-end", + }, + ], + } + return replace(_workflow_tool_container_workflow_tool_source(), graph_config=source_graph) + + +def _workflow_tool_container_container_handler( + *, + inputs: dict[str, object] | None = None, +) -> tuple[ + WorkflowToolContainerHandler, + FrameRegistry, + RuntimeState, + CustomContainerRequest, + dict[str, WorkflowToolSource], +]: + source = _workflow_tool_container_workflow_tool_source() + parent_pool = VariablePool() + parent_pool.add(system_variable_selector(SystemVariableKey.USER_ID), "user") + parent_pool.add(system_variable_selector(SystemVariableKey.WORKFLOW_EXECUTION_ID), "outer-execution") + parent_pool.add(("env", "outer-secret"), "must-not-leak") + parent_pool.add(("source-start", "stale"), "must-not-leak") + parent_state = RuntimeState( + workflow_id="outer-workflow", + variable_pool=parent_pool, + start_at=1, + ) + parent_node, _, _ = _workflow_tool_container_workflow_tool_node(parent_state) + parent_graph = _workflow_tool_container_outer_graph(parent_node) + frame_registry = FrameRegistry() + frame_registry.register( + ExecutionFrame( + frame_id=ROOT_FRAME_ID, + graph=parent_graph, + state=parent_state, + scheduler=Scheduler(parent_graph, parent_state, ROOT_FRAME_ID), + failure_handler=NodeFailureHandler(parent_graph, parent_state.graph_execution), + ) + ) + source_repository = {} + source_repository["source-workflow"] = source + payload = WorkflowToolContainerPayload( + source_app_id=source.app_id, + source_workflow_id=source.workflow_id, + source_workflow_version="1", + inputs={"answer": "ok"} if inputs is None else inputs, + inputs_for_log={"answer": "ok"} if inputs is None else inputs, + call_depth=1, + ) + request = CustomContainerRequest(payload=payload.model_dump_json()) + parent_state.put_container_run( + create_container_run_state( + invocation_id="invocation", + frame_id=ROOT_FRAME_ID, + node_id="tool", + started_at=datetime.now(UTC).replace(tzinfo=None), + request=request, + ) + ) + return ( + WorkflowToolContainerHandler(frame_registry, sources=source_repository), + frame_registry, + parent_state, + request, + source_repository, + ) + + +def _dispatch_next_node( + *, + expected_node_id: str, + runtime_state: RuntimeState, + frame_registry: FrameRegistry, + processor: NodeEventProcessor, +) -> None: + task = runtime_state.ready_queue.get(timeout=0.01) + assert isinstance(task, StartTask) + assert task.node_id == expected_node_id + frame = frame_registry[task.frame_id] + node = frame.graph.nodes[task.node_id] + execution = runtime_state.graph_execution.get_or_create_node_execution( + frame_id=frame.frame_id, + node_id=node.id, + ) + node.bind_execution_id(execution.execution_id) + for event in node.run(): + assert isinstance(event, _workflow_tool_container_NodeEvent) + processor.dispatch(NodeEventTask(frame_id=frame.frame_id, event=event)) + + +class TestWorkflowToolContainerCases: + @staticmethod + def test_workflow_tool_node_requests_child_container_and_resumes_successfully(mocker: MockerFixture) -> None: + node, runtime, payload = _workflow_tool_container_workflow_tool_node() + get_runtime = mocker.spy(runtime, "get_runtime") + build_payload = mocker.spy(runtime, "build_workflow_tool_container_payload") + + initial_events = list(node._run()) + + assert initial_events == [CustomContainerRequest(payload=payload.model_dump_json())] + get_runtime.assert_called_once_with( + node_id="tool", + node_data=node.node_data, + variable_pool=None, + node_execution_id="tool-execution", + ) + build_payload.assert_called_once_with( + tool_runtime=get_runtime.spy_return, + tool_parameters={}, + inputs_for_log={}, + workflow_call_depth=2, + ) + + result = ContainerExecutionResult( + metadata={}, + steps=2, + node_run_result=ContainerNodeRunResult( + status=WorkflowNodeExecutionStatus.SUCCEEDED, + inputs={"question": build_container_value("hello")}, + outputs={ + "text": build_container_value("done"), + "answer": build_container_value(42), + "files": WorkflowToolContainerHandler._build_tool_outputs({})["files"], + }, + ), + ) + assert result.node_run_result.outputs["files"].value_type == SegmentType.ARRAY_FILE + node.runtime_state.enqueue_ready_task(ResumeTask(invocation_id="invocation", result=result)) + restored_state = RuntimeState.from_snapshot(node.runtime_state.dumps()) + resume_task = restored_state.ready_queue.get(timeout=0.01) + assert isinstance(resume_task, ResumeTask) + assert isinstance(resume_task.result, ContainerExecutionResult) + assert resume_task.result.node_run_result.outputs["files"].value_type == SegmentType.ARRAY_FILE + restored_node, _, _ = _workflow_tool_container_workflow_tool_node(restored_state) + + resumed_events = list(restored_node._resume_container_events(result=resume_task.result)) + + assert resumed_events[:2] == [ + StreamChunkEvent(selector=["tool", "text"], chunk="done", is_final=False), + StreamChunkEvent(selector=["tool", "text"], chunk="", is_final=True), + ] + completed = resumed_events[2] + assert isinstance(completed, StreamCompletedEvent) + assert completed.node_run_result.inputs == {"question": "hello"} + assert {key: value.to_object() for key, value in completed.node_run_result.outputs.items()} == { + "text": "done", + "answer": 42, + "files": [], + } + assert isinstance(completed.node_run_result.outputs["files"], ArrayFileSegment) + for key, value in completed.node_run_result.outputs.items(): + restored_state.variable_pool.add(("tool", key), value) + files = restored_state.variable_pool.get(("tool", "files")) + assert files is not None + assert files.value_type == SegmentType.ARRAY_FILE + assert files.to_object() == [] + assert node.runtime_state.node_run_steps == 0 + assert completed.node_run_result.metadata[WorkflowNodeExecutionMetadataKey.TOOL_INFO] == { + "provider_type": ToolProviderType.WORKFLOW.value, + "provider_id": "workflow-provider", + "plugin_unique_identifier": None, + } + + @staticmethod + def test_workflow_tool_node_returns_failure_when_runtime_cannot_be_loaded(mocker: MockerFixture) -> None: + node, runtime, _ = _workflow_tool_container_workflow_tool_node(version="2") + get_runtime = mocker.patch.object(runtime, "get_runtime", side_effect=ToolNodeError("runtime unavailable")) + + events = list(node._run()) + + assert len(events) == 1 + completed = events[0] + assert isinstance(completed, StreamCompletedEvent) + assert completed.node_run_result.status == WorkflowNodeExecutionStatus.FAILED + assert completed.node_run_result.error == "Failed to get tool runtime: runtime unavailable" + get_runtime.assert_called_once_with( + node_id="tool", + node_data=node.node_data, + variable_pool=node.runtime_state.variable_pool, + node_execution_id="tool-execution", + ) + + @staticmethod + def test_workflow_tool_node_returns_failure_when_payload_cannot_be_built(mocker: MockerFixture) -> None: + node, runtime, _ = _workflow_tool_container_workflow_tool_node(tool_node_version="1") + mocker.patch.object( + runtime, "build_workflow_tool_container_payload", side_effect=ToolNodeError("invalid source workflow") + ) + + events = list(node._run()) + + assert len(events) == 1 + completed = events[0] + assert isinstance(completed, StreamCompletedEvent) + assert completed.node_run_result.status == WorkflowNodeExecutionStatus.FAILED + assert completed.node_run_result.error == "Failed to prepare Workflow Tool: invalid source workflow" + + @staticmethod + def test_workflow_tool_node_rejects_non_execution_result() -> None: + node, _, _ = _workflow_tool_container_workflow_tool_node() + result = LoopFrameRequest( + inputs={}, + outputs={}, + loop_count=0, + root_node_id="loop", + loop_variable_selectors={}, + loop_node_ids=frozenset(), + index=0, + ) + + with pytest.raises(TypeError, match="Unsupported Workflow Tool container result LoopFrameRequest"): + list(node._resume_container_events(result=result)) + + @staticmethod + def test_workflow_tool_node_resumes_failed_result_without_chunks() -> None: + node, _, _ = _workflow_tool_container_workflow_tool_node() + result = ContainerExecutionResult( + metadata={}, + steps=1, + node_run_result=ContainerNodeRunResult( + status=WorkflowNodeExecutionStatus.FAILED, + error="child failed", + ), + ) + + events = list(node._resume_container_events(result=result)) + + assert len(events) == 1 + assert isinstance(events[0], StreamCompletedEvent) + assert events[0].node_run_result.error == "child failed" + + @staticmethod + def test_workflow_tool_node_resumes_empty_text_with_final_chunk_only() -> None: + node, _, _ = _workflow_tool_container_workflow_tool_node() + result = ContainerExecutionResult( + metadata={}, + steps=0, + node_run_result=ContainerNodeRunResult( + status=WorkflowNodeExecutionStatus.SUCCEEDED, + outputs={"text": build_container_value("")}, + ), + ) + + events = list(node._resume_container_events(result=result)) + + assert events[:1] == [StreamChunkEvent(selector=["tool", "text"], chunk="", is_final=True)] + assert isinstance(events[1], StreamCompletedEvent) + + @staticmethod + def test_node_factory_can_keep_workflow_tool_direct_for_single_step_debug() -> None: + node, _, _ = _workflow_tool_container_workflow_tool_node() + factory = DifyNodeFactory(node.init_params, node.runtime_state, use_workflow_tool_containers=True) + assert ( + factory._resolve_node_class_for_factory( + node_type=BuiltinNodeTypes.TOOL, + node_version="1", + node_data=node.node_data, + ) + is DifyWorkflowToolNode + ) + + factory._use_workflow_tool_containers = False + assert ( + factory._resolve_node_class_for_factory( + node_type=BuiltinNodeTypes.TOOL, + node_version="1", + node_data=node.node_data, + ) + is ToolNode + ) + + @staticmethod + def test_workflow_tool_handler_runs_child_graph_with_internal_name_collision() -> None: + handler, frame_registry, runtime_state, request, _ = _workflow_tool_container_container_handler() + handler.handle_request(invocation_id="invocation", request=request) + frame_registry["invocation:workflow-tool"].state.variable_pool.add( + ("__workflow_tool_container__", "failure"), + {"legitimate": True}, + ) + event_stream = EventStream(layers=[]) + processor = NodeEventProcessor( + graph_execution=runtime_state.graph_execution, + event_stream=event_stream, + frame_registry=frame_registry, + container_handlers={BuiltinNodeTypes.TOOL: handler}, + ) + + _dispatch_next_node( + expected_node_id="source-start", + runtime_state=runtime_state, + frame_registry=frame_registry, + processor=processor, + ) + _dispatch_next_node( + expected_node_id="source-end", + runtime_state=runtime_state, + frame_registry=frame_registry, + processor=processor, + ) + + resume_task = runtime_state.ready_queue.get(timeout=0.01) + assert isinstance(resume_task, ResumeTask) + assert isinstance(resume_task.result, ContainerExecutionResult) + assert resume_task.result.steps == 2 + assert resume_task.result.node_run_result.outputs["files"].value_type == SegmentType.ARRAY_FILE + outputs = {key: value.to_object() for key, value in resume_task.result.node_run_result.outputs.items()} + assert outputs == { + "answer": "ok", + "text": '{"answer": "ok"}', + "files": [], + "json": [{"answer": "ok"}], + } + event_stream.mark_complete() + assert not any(isinstance(event, _workflow_tool_container_NodeEvent) for event in event_stream.emit_events()) + with pytest.raises(KeyError): + frame_registry["invocation:workflow-tool"] + with pytest.raises(KeyError): + runtime_state.get_container_frame("invocation:workflow-tool") + + @staticmethod + def test_workflow_tool_handler_restores_child_failure() -> None: + handler, frame_registry, runtime_state, request, source_repository = ( + _workflow_tool_container_container_handler() + ) + handler.handle_request(invocation_id="invocation", request=request) + frame_id = "invocation:workflow-tool" + frame = frame_registry[frame_id] + now = datetime.now(UTC).replace(tzinfo=None) + handler.record_frame_failure( + frame=frame, + event=NodeRunFailedEvent( + id="failed-execution", + node_id="source-start", + node_type=BuiltinNodeTypes.START, + error="source failed", + start_at=now, + finished_at=now, + node_run_result=ContainerNodeRunResult( + status=WorkflowNodeExecutionStatus.FAILED, + error="source failed", + error_type="RuntimeError", + ).to_node_run_result(), + ), + ) + frame_state = runtime_state.get_container_frame(frame_id) + frame_state = frame_state.model_copy(update={"runtime_data": frame.state.snapshot_frame()}) + runtime_state.put_container_frame(frame_state) + frame_registry.remove(frame_id) + + restored_handler = WorkflowToolContainerHandler(frame_registry, sources=source_repository) + restored_handler.restore_frame(frame_state) + restored_handler.complete_frame_if_ready(frame_registry[frame_id]) + + tasks = [runtime_state.ready_queue.get(timeout=0.01) for _ in range(2)] + resume_task = next(task for task in tasks if isinstance(task, ResumeTask)) + assert isinstance(resume_task.result, ContainerExecutionResult) + assert resume_task.result.node_run_result.status == WorkflowNodeExecutionStatus.FAILED + assert {key: value.to_object() for key, value in resume_task.result.node_run_result.inputs.items()} == { + "answer": "ok" + } + assert resume_task.result.node_run_result.error == "source failed" + assert resume_task.result.node_run_result.error_type == "RuntimeError" + + @staticmethod + def test_workflow_tool_handler_isolates_child_variable_pool() -> None: + handler, frame_registry, _, request, _ = _workflow_tool_container_container_handler() + + handler.handle_request(invocation_id="invocation", request=request) + + child_pool = frame_registry["invocation:workflow-tool"].state.variable_pool + assert child_pool.get(system_variable_selector(SystemVariableKey.USER_ID)).to_object() == "user" # type: ignore[union-attr] + assert child_pool.get(system_variable_selector(SystemVariableKey.WORKFLOW_EXECUTION_ID)).to_object() == ( # type: ignore[union-attr] + "outer-execution" + ) + assert child_pool.get(system_variable_selector(SystemVariableKey.APP_ID)).to_object() == ( # type: ignore[union-attr] + "source-app" + ) + assert child_pool.get(("env", "outer-secret")) is None + assert child_pool.get(("source-start", "stale")) is None + + @staticmethod + def test_workflow_tool_handler_applies_start_input_defaults() -> None: + handler, frame_registry, _, request, _ = _workflow_tool_container_container_handler(inputs={}) + + handler.handle_request(invocation_id="invocation", request=request) + + answer = frame_registry["invocation:workflow-tool"].state.variable_pool.get(("source-start", "answer")) + assert answer is not None + assert answer.to_object() == "fallback" + + @staticmethod + def test_workflow_tool_handler_preserves_inputs_when_start_validation_fails() -> None: + handler, _, runtime_state, request, _ = _workflow_tool_container_container_handler(inputs={"answer": 42}) + + handler.handle_request(invocation_id="invocation", request=request) + + resume_task = runtime_state.ready_queue.get(timeout=0.01) + assert isinstance(resume_task, ResumeTask) + assert isinstance(resume_task.result, ContainerExecutionResult) + assert resume_task.result.node_run_result.status == WorkflowNodeExecutionStatus.FAILED + assert {key: value.to_object() for key, value in resume_task.result.node_run_result.inputs.items()} == { + "answer": 42 + } + assert resume_task.result.node_run_result.error_type == "ValueError" + + @staticmethod + def test_workflow_tool_nested_handler_hides_and_persists_marked_child_events() -> None: + _, frame_registry, _, request, repository = _workflow_tool_container_container_handler() + persisted: list[_workflow_tool_container_NodeEvent] = [] + event_listeners: dict[str, Callable[[_workflow_tool_container_NodeEvent], None]] = {} + workflow_tool_handler = WorkflowToolContainerHandler( + frame_registry, + sources=repository, + event_listener_factory=lambda *_: persisted.append, + event_listeners=event_listeners, + ) + workflow_tool_handler.handle_request(invocation_id="invocation", request=request) + built_in_handler = LoopContainerHandler(frame_registry) + hidden_events: list[_workflow_tool_container_NodeEvent] = [] + nested_handler = WorkflowToolNestedContainerHandler( + frame_registry, + handler_factory=lambda _: built_in_handler, + hidden_event_listener=hidden_events.append, + event_listeners=event_listeners, + ) + event = NodeRunSucceededEvent( + id="source-execution", + node_id="source-start", + node_type=BuiltinNodeTypes.START, + start_at=datetime.now(UTC).replace(tzinfo=None), + container_id="nested-container", + ) + + workflow_tool_handler.prepare_frame_event( + frame=frame_registry["invocation:workflow-tool"], + event=event, + ) + + assert nested_handler.should_emit(event=event) is False + assert hidden_events == [event] + assert len(persisted) == 1 + assert persisted[0].container_id == "nested-container" + assert persisted[0].node_run_result.process_data["workflow_tool_invocation_id"] == "invocation" + assert "__dify_workflow_tool_child__" not in persisted[0].node_run_result.process_data + + # Another Workflow Tool inside this container owns its own source events. + inner_events: list[_workflow_tool_container_NodeEvent] = [] + event_listeners["inner-tool-frame"] = inner_events.append + inner_event = event.model_copy(deep=True) + inner_event.node_run_result.process_data = {} + workflow_tool_handler.prepare_frame_event( + frame=replace(frame_registry["invocation:workflow-tool"], frame_id="inner-tool-frame"), event=inner_event + ) + workflow_tool_handler.prepare_frame_event(frame=frame_registry["invocation:workflow-tool"], event=inner_event) + assert nested_handler.should_emit(event=inner_event) is False + assert inner_events == [inner_event] + assert len(persisted) == 1 + + unmarked_event = event.model_copy(deep=True) + unmarked_event.node_run_result.process_data = {} + assert nested_handler.should_emit(event=unmarked_event) is True + assert len(persisted) == 1 + + @staticmethod + @pytest.mark.parametrize( + "source_metadata", + [{}, {WorkflowNodeExecutionMetadataKey.LOOP_ID: "source-loop", WorkflowNodeExecutionMetadataKey.LOOP_INDEX: 1}], + ) + def test_workflow_tool_persistence_keeps_source_container_metadata( + source_metadata: dict[WorkflowNodeExecutionMetadataKey, str | int], + ) -> None: + _, frames, _, request, repository = _workflow_tool_container_container_handler() + persisted: list[_workflow_tool_container_NodeEvent] = [] + handler = WorkflowToolContainerHandler( + frames, sources=repository, event_listener_factory=lambda *_: persisted.append + ) + handler.handle_request(invocation_id="invocation", request=request) + event = NodeRunSucceededEvent( + id="source-execution", + node_id="source-start", + node_type=BuiltinNodeTypes.START, + start_at=datetime.now(UTC).replace(tzinfo=None), + ) + event.node_run_result.metadata = source_metadata.copy() + handler.prepare_frame_event(frame=frames["invocation:workflow-tool"], event=event) + # Ancestors outside the Tool source add their own ownership before collection. + event.node_run_result.metadata.update( + { + WorkflowNodeExecutionMetadataKey.ITERATION_ID: "outer-iteration", + WorkflowNodeExecutionMetadataKey.ITERATION_INDEX: 2, + } + ) + event.node_run_result.outputs = {"normalized": True} + + assert handler.should_emit(event=event) is False + + assert persisted[0].node_run_result.metadata == source_metadata + assert persisted[0].node_run_result.outputs == {"normalized": True} + + @staticmethod + @pytest.mark.parametrize( + ( + "child_error_strategy", + "tool_error_strategy", + "terminal_type", + "expected_count", + "expected_output", + "retry_once", + ), + [ + pytest.param( + ErrorStrategy.DEFAULT_VALUE, + None, + _workflow_tool_container_GraphRunSucceededEvent, + 0, + '{"answer": "child fallback"}', + False, + id="child-handles-failure", + ), + pytest.param(None, None, GraphRunFailedEvent, 1, None, False, id="unhandled-failure"), + pytest.param( + None, + ErrorStrategy.DEFAULT_VALUE, + GraphRunPartialSucceededEvent, + 1, + "tool fallback", + False, + id="tool-default-value", + ), + pytest.param( + None, + ErrorStrategy.FAIL_BRANCH, + GraphRunPartialSucceededEvent, + 1, + "source code failed", + False, + id="tool-fail-branch", + ), + pytest.param( + None, + None, + _workflow_tool_container_GraphRunSucceededEvent, + 0, + '{"answer": "retried output"}', + True, + id="child-retries", + ), + ], + ) + def test_workflow_tool_failure_accounting_uses_outer_tool_policy( + monkeypatch: pytest.MonkeyPatch, + child_error_strategy: ErrorStrategy | None, + tool_error_strategy: ErrorStrategy | None, + terminal_type: type[ + _workflow_tool_container_GraphRunSucceededEvent | GraphRunFailedEvent | GraphRunPartialSucceededEvent + ], + expected_count: int, + expected_output: str | None, + retry_once: bool, + ) -> None: + attempts = 0 + + def execute(**_kwargs: object) -> dict[str, str]: + nonlocal attempts + attempts += 1 + if attempts == 1: + raise CodeExecutionError("source code failed") + return {"output": "retried output"} + + monkeypatch.setattr("core.workflow.node_factory.CodeExecutor.execute_workflow_code_template", execute) + source = _workflow_tool_container_workflow_tool_source() + source_graph = dict(source.graph_config) + source_graph["nodes"].insert( + 1, + { + "id": "source-code", + "data": { + "type": "code", + "title": "Code", + "code_language": "python3", + "code": "def main(): return {'output': 'retried output'}", + "outputs": {"output": {"type": "string"}}, + "variables": [], + "error_strategy": child_error_strategy, + "default_value": [{"key": "output", "type": "string", "value": "child fallback"}], + "retry_config": {"retry_enabled": retry_once, "max_retries": 1, "retry_interval": 0}, + }, + }, + ) + source_graph["nodes"][-1]["data"]["outputs"][0]["value_selector"] = ["source-code", "output"] + source_graph["edges"] = [ + {"id": "start-code", "source": "source-start", "target": "source-code"}, + {"id": "code-end", "source": "source-code", "target": "source-end"}, + ] + repository = {} + repository["source-workflow"] = replace(source, graph_config=source_graph) + node, _, _ = _workflow_tool_container_workflow_tool_node() + node.node_data.error_strategy = tool_error_strategy + node.node_data.default_value = [DefaultValue(key="text", type=DefaultValueType.STRING, value="tool fallback")] + start = StartNode( + node_id="outer-start", + data=StartNodeData(title="Start", variables=[]), + init_params=node.init_params, + runtime_state=node.runtime_state, + ) + end = EndNode( + node_id="outer-end", + data=EndNodeData.model_validate( + { + "title": "End", + "outputs": [ + { + "variable": "answer", + "value_selector": [ + "tool", + "error_message" if tool_error_strategy == ErrorStrategy.FAIL_BRANCH else "text", + ], + } + ], + } + ), + init_params=node.init_params, + runtime_state=node.runtime_state, + ) + graph = ( + Graph.new() + .add_root(start) + .add_node(node) + .add_node( + end, source_handle="fail-branch" if tool_error_strategy == ErrorStrategy.FAIL_BRANCH else "source" + ) + .build() + ) + persisted: list[_workflow_tool_container_NodeEvent] = [] + engine = Engine( + graph=graph, + runtime_state=node.runtime_state, + command_channel=InMemoryChannel(), + workers=1, + container_handler_factories=( + partial( + WorkflowToolContainerHandler, + sources=repository, + event_listener_factory=lambda *_: persisted.append, + ), + ), + ) + events: list[EngineEvent] = [] + if terminal_type is GraphRunFailedEvent: + with pytest.raises(RuntimeError, match="source code failed"): + events.extend(engine.run()) + else: + events.extend(engine.run()) + + assert attempts == (2 if retry_once else 1) + assert isinstance(events[-1], terminal_type) + assert node.runtime_state.graph_execution.exceptions_count == expected_count + if isinstance(events[-1], GraphRunFailedEvent): + assert events[-1].error == "source code failed" + assert not any(isinstance(event, NodeRunSucceededEvent) and event.node_id == end.id for event in events) + else: + assert isinstance( + events[-1], _workflow_tool_container_GraphRunSucceededEvent | GraphRunPartialSucceededEvent + ) + assert events[-1].outputs == {"answer": expected_output} + if tool_error_strategy is not None: + assert any(isinstance(event, NodeRunExceptionEvent) and event.node_id == node.id for event in events) + if isinstance(events[-1], GraphRunFailedEvent | GraphRunPartialSucceededEvent): + assert events[-1].exceptions_count == expected_count + if child_error_strategy is not None: + assert any( + isinstance(event, NodeRunExceptionEvent) and event.node_id == "source-code" for event in persisted + ) + if retry_once: + assert [event.retry_index for event in persisted if isinstance(event, NodeRunRetryEvent)] == [1] + assert any( + isinstance(event, NodeRunSucceededEvent) and event.node_id == "source-code" for event in persisted + ) + + @staticmethod + @pytest.mark.parametrize("outcome", ["paused", "resumed", "abort-on-resume", "stopped"]) + def test_workflow_tool_human_input_pauses_and_resumes_without_duplicate_form( + monkeypatch: pytest.MonkeyPatch, + mocker: MockerFixture, + outcome: str, + ) -> None: + command_channel = InMemoryChannel() + + def stop_on_human_input( + event: _workflow_tool_container_NodeEvent, *, channel: InMemoryChannel, should_stop: bool + ) -> None: + if should_stop and isinstance(event, NodeRunPauseRequestedEvent): + channel.send_command(AbortCommand()) + + source_repository = {"source-workflow": _human_input_workflow_tool_source()} + handler_factory = partial( + WorkflowToolContainerHandler, + sources=source_repository, + hidden_event_listener=partial( + stop_on_human_input, channel=command_channel, should_stop=outcome == "stopped" + ), + ) + form_repository = HumanInputFormRepositoryImpl( + tenant_id="tenant", app_id="outer-app", workflow_execution_id="outer-execution" + ) + create_form_spy = mocker.spy(form_repository, "create_form") + human_input_app_ids: list[str] = [] + + def build_human_input_callback( + factory: DifyNodeFactory, + *, + node_data: HumanInputNodeData, + execution_id_getter: Callable[[], str | None], + ) -> DifyHITLCallback: + human_input_app_ids.append(factory._human_input_runtime._run_context.app_id) + return DifyHITLCallback( + form_repository=form_repository, + node_data=node_data, + execution_id_getter=execution_id_getter, + ) + + monkeypatch.setattr(DifyNodeFactory, "_build_human_input_callback", build_human_input_callback) + + outer_pool = VariablePool() + outer_pool.add(system_variable_selector(SystemVariableKey.WORKFLOW_EXECUTION_ID), "outer-execution") + outer_pool.add(system_variable_selector(SystemVariableKey.APP_ID), "outer-app") + outer_pool.add(system_variable_selector(SystemVariableKey.WORKFLOW_ID), "outer-workflow") + initial_state = RuntimeState( + workflow_id="outer-workflow", + variable_pool=outer_pool, + start_at=1, + ) + initial_node, _, _ = _workflow_tool_container_workflow_tool_node(initial_state, app_id="intermediate-app") + initial_graph = _workflow_tool_container_outer_graph(initial_node) + initial_owner_factory = DifyNodeFactory( + initial_node.init_params, initial_state, human_input_run_context=build_test_run_context(app_id="outer-app") + ) + initial_graph.node_factory = initial_owner_factory + initial_events = list( + Engine( + graph=initial_graph, + runtime_state=initial_state, + command_channel=command_channel, + workers=1, + container_handler_factories=(handler_factory,), + ).run() + ) + + paused = initial_events[-1] + if outcome == "stopped": + assert isinstance(paused, GraphRunAbortedEvent) + return + assert isinstance(paused, GraphRunPausedEvent) + assert len(paused.reasons) == 1 + reason = paused.reasons[0] + assert isinstance(reason, HitlRequired) + assert reason.node_id == "source-human" + assert reason.node_title == "Approval" + assert all( + not isinstance(event, _workflow_tool_container_NodeEvent) + or event.node_id not in {"source-start", "source-human"} + for event in initial_events + ) + create_form_spy.assert_called_once() + create_params = create_form_spy.call_args.args[0] + assert create_params.node_id == "source-human" + assert create_params.workflow_execution_id == "outer-execution" + child_frames = list(initial_state.container_frames()) + assert len(child_frames) == 1 + child_execution = initial_state.graph_execution.node_executions[(child_frames[0].frame_id, "source-human")] + assert create_params.form_id == child_execution.execution_id + assert human_input_app_ids == ["outer-app"] + + if outcome == "paused": + return + + restored_state = RuntimeState.from_snapshot(initial_state.dumps()) + assert ( + resolve_human_input_node_id( + node_id=reason.node_id, + form_id=reason.session_id, + variable_pool=restored_state.variable_pool, + ) + == "tool" + ) + restored_node, _, _ = _workflow_tool_container_workflow_tool_node(restored_state, app_id="intermediate-app") + restored_graph = _workflow_tool_container_outer_graph(restored_node) + restored_owner_factory = DifyNodeFactory( + restored_node.init_params, + restored_state, + human_input_run_context=initial_owner_factory.human_input_run_context, + ) + restored_graph.node_factory = restored_owner_factory + if outcome == "resumed": + HumanInputFormSubmissionRepository().mark_submitted( + form_id=reason.session_id, + recipient_id=None, + selected_action_id="approve", + form_data={}, + submission_user_id="user", + submission_end_user_id=None, + ) + resumed_channel = InMemoryChannel() + resumed_handler_factory = partial( + WorkflowToolContainerHandler, + sources=source_repository, + hidden_event_listener=partial( + stop_on_human_input, channel=resumed_channel, should_stop=outcome == "abort-on-resume" + ), + ) + resumed_events = list( + Engine( + graph=restored_graph, + runtime_state=restored_state, + command_channel=resumed_channel, + workers=1, + container_handler_factories=(resumed_handler_factory,), + ).run() + ) + + if outcome == "abort-on-resume": + assert isinstance(resumed_events[-1], GraphRunAbortedEvent) + return + assert isinstance(resumed_events[-1], _workflow_tool_container_GraphRunSucceededEvent) + tool_succeeded = next( + event for event in resumed_events if isinstance(event, NodeRunSucceededEvent) and event.node_id == "tool" + ) + assert tool_succeeded.node_run_result.outputs["decision"].to_object() == "approve" + assert json.loads(tool_succeeded.node_run_result.outputs["text"].to_object()) == {"decision": "approve"} + files = restored_state.variable_pool.get(("tool", "files")) + assert files is not None + assert files.value_type == SegmentType.ARRAY_FILE + assert files.to_object() == [] + create_form_spy.assert_called_once() + assert human_input_app_ids == ["outer-app", "outer-app"] + assert all( + not isinstance(event, _workflow_tool_container_NodeEvent) + or event.node_id not in {"source-human", "source-end"} + for event in resumed_events + ) + assert list(restored_state.container_frames()) == [] + assert list(restored_state.container_runs()) == [] + + +_workflow_tool_observability_outer_graph = _workflow_tool_container_outer_graph +_workflow_tool_observability_workflow_tool_node = _workflow_tool_container_workflow_tool_node +_workflow_tool_observability_workflow_tool_source = _workflow_tool_container_workflow_tool_source + + +class TestWorkflowToolObservabilityCases: + from tests.unit_tests.core.app.workflow.layers.conftest import memory_span_exporter + + memory_span_exporter = staticmethod(memory_span_exporter) + + from tests.unit_tests.core.app.workflow.layers.conftest import tracer_provider_with_memory_exporter + + tracer_provider_with_memory_exporter = staticmethod(tracer_provider_with_memory_exporter) + + from tests.unit_tests.core.app.workflow.layers.conftest import mock_start_node + + mock_start_node = staticmethod(mock_start_node) + + from tests.unit_tests.core.app.workflow.layers.conftest import mock_llm_node + + mock_llm_node = staticmethod(mock_llm_node) + + from tests.unit_tests.core.app.workflow.layers.conftest import mock_tool_node + + mock_tool_node = staticmethod(mock_tool_node) + + from tests.unit_tests.core.app.workflow.layers.conftest import mock_is_instrument_flag_enabled_false + + mock_is_instrument_flag_enabled_false = staticmethod(mock_is_instrument_flag_enabled_false) + + from tests.unit_tests.core.app.workflow.layers.conftest import mock_is_instrument_flag_enabled_true + + mock_is_instrument_flag_enabled_true = staticmethod(mock_is_instrument_flag_enabled_true) + + from tests.unit_tests.core.app.workflow.layers.conftest import mock_retrieval_node + + mock_retrieval_node = staticmethod(mock_retrieval_node) + + from tests.unit_tests.core.app.workflow.layers.conftest import mock_result_event + + mock_result_event = staticmethod(mock_result_event) + + @staticmethod + def test_workflow_tool_tracing_survives_worker_handoff( + monkeypatch: pytest.MonkeyPatch, + caplog: pytest.LogCaptureFixture, + tracer_provider_with_memory_exporter: TracerProvider, + memory_span_exporter: InMemorySpanExporter, + ) -> None: + apply_config_overrides(monkeypatch, ENABLE_OTEL=True) + tool_started_on: list[int] = [] + tool_finished_on: list[int] = [] + completed = Event() + + class HandoffLayer(Layer): + @override + def node_run_context( + self, node: Node[BaseNodeData], *, parent_execution_id: str | None = None + ) -> AbstractContextManager[None]: + return HandoffContext() + + @override + def on_node_run_start(self, node: Node[BaseNodeData]) -> None: + if node.id == "tool": + tool_started_on.append(get_ident()) + + @override + def on_node_run_end( + self, + node: Node[BaseNodeData], + error: Exception | None, + result_event: _workflow_tool_observability_NodeEvent | None = None, + ) -> None: + if node.id == "tool": + tool_finished_on.append(get_ident()) + completed.set() + + tracer = tracer_provider_with_memory_exporter.get_tracer(__name__) + with tracer.start_as_current_span("workflow") as root_span: + captured_context = capture_current_context() + + class HandoffContext(AbstractContextManager[None]): + @override + def __enter__(self) -> None: + captured_context.__enter__() + + @override + def __exit__( + self, + exc_type: type[BaseException] | None, + exc: BaseException | None, + traceback: TracebackType | None, + ) -> None: + captured_context.__exit__(exc_type, exc, traceback) + if tool_started_on == [get_ident()] and not completed.is_set(): + # Keep the starting worker occupied after suspension so another + # worker must execute the child graph and resume the Tool. + assert completed.wait(10), "Workflow Tool did not resume on the available worker" + + state = RuntimeState(workflow_id="outer-workflow", variable_pool=VariablePool(), start_at=1) + tool, _, _ = _workflow_tool_observability_workflow_tool_node(state) + original_resume = tool._resume_container_events + + def instrument_resume(*, result: ContainerRunResult) -> Generator[NodeEventPayload | ContainerAwaitRequest]: + with tracer.start_as_current_span("resume-operation"): + yield from original_resume(result=result) + + monkeypatch.setattr(tool, "_resume_container_events", instrument_resume) + source = _workflow_tool_observability_workflow_tool_source() + sources = {source.workflow_id: source} + engine = Engine( + graph=_workflow_tool_observability_outer_graph(tool), + runtime_state=state, + workers=2, + container_handler_factories=(partial(WorkflowToolContainerHandler, sources=sources),), + ) + engine.add_layer(HandoffLayer()) + engine.add_layer(ObservabilityLayer()) + + events = list(engine.run()) + + assert isinstance(events[-1], _workflow_tool_observability_GraphRunSucceededEvent) + assert len(tool_started_on) == len(tool_finished_on) == 1 + assert tool_started_on != tool_finished_on + spans = memory_span_exporter.get_finished_spans() + tool_spans = [span for span in spans if (span.attributes or {}).get("node.id") == "tool"] + assert len(tool_spans) == 1 + tool_span = tool_spans[0] + assert tool_span.parent == root_span.get_span_context() + resumed_operations = [span for span in spans if span.name == "resume-operation"] + assert len(resumed_operations) == 1 + assert resumed_operations[0].parent == tool_span.context + assert "was created in a different Context" not in caplog.text + children = [span for span in spans if (span.attributes or {}).get("node.id") in {"source-start", "source-end"}] + assert len(children) == 2 + assert all(span.parent == tool_span.context for span in children) + + @staticmethod + def test_standalone_node_uses_worker_scoped_tracing_context( + monkeypatch: pytest.MonkeyPatch, + tracer_provider_with_memory_exporter: TracerProvider, + memory_span_exporter: InMemorySpanExporter, + ) -> None: + apply_config_overrides(monkeypatch, ENABLE_OTEL=True) + tool, _, _ = _workflow_tool_observability_workflow_tool_node() + node = _workflow_tool_observability_outer_graph(tool).root_node + tracer = tracer_provider_with_memory_exporter.get_tracer(__name__) + with tracer.start_as_current_span("standalone") as parent: + events = list(WorkflowEntry._run_node_with_layers(node)) + assert get_current_span() is parent + + assert isinstance(events[-1], _workflow_tool_observability_NodeEvent) + assert events[-1].node_run_result.status == "succeeded" + node_spans = [span for span in memory_span_exporter.get_finished_spans() if span.name == node.title] + assert len(node_spans) == 1 + assert node_spans[0].parent == parent.get_span_context() + + @staticmethod + def test_interleaved_snippets_restore_caller_flask_context( + monkeypatch: pytest.MonkeyPatch, + tracer_provider_with_memory_exporter: TracerProvider, + memory_span_exporter: InMemorySpanExporter, + ) -> None: + apply_config_overrides(monkeypatch, ENABLE_OTEL=True) + monkeypatch.setattr(module, "capture_current_context", capture_flask_context) + request_id: ContextVar[str] = ContextVar("snippet_request_id") + tracer = tracer_provider_with_memory_exporter.get_tracer(__name__) + parents = [] + snippets = [] + observed = [] + for name in ("first", "second"): + app = DifyApp(name) + init_app(app) + with app.app_context(), tracer.start_as_current_span(name) as parent: + parents.append(parent.get_span_context()) + g._login_user = name + request_id.set(name) + tool, _, _ = _workflow_tool_observability_workflow_tool_node() + node = _workflow_tool_observability_outer_graph(tool).root_node + + def run_node() -> Generator[NodeRunStreamChunkEvent]: + observed.append((current_app.name, g._login_user, request_id.get())) + yield NodeRunStreamChunkEvent( + id="snippet-execution", + node_id="outer-start", + node_type=BuiltinNodeTypes.START, + selector=["outer-start", "text"], + chunk="part", + is_final=False, + ) + observed.append((current_app.name, g._login_user, request_id.get())) + yield NodeRunStreamChunkEvent( + id="snippet-execution", + node_id="outer-start", + node_type=BuiltinNodeTypes.START, + selector=["outer-start", "text"], + chunk="part", + is_final=False, + ) + + monkeypatch.setattr(node, "run", run_node) + snippets.append(WorkflowEntry._run_node_with_layers(node)) + + caller = DifyApp("caller") + init_app(caller) + with caller.app_context(), tracer.start_as_current_span("caller") as caller_span: + g._login_user = "caller" + request_id.set("caller") + for _ in range(2): + for snippet in snippets: + next(snippet) + assert current_app == caller + assert g._login_user == request_id.get() == "caller" + assert get_current_span() is caller_span + for snippet in snippets: + snippet.close() + assert current_app == caller + assert g._login_user == request_id.get() == "caller" + assert get_current_span() is caller_span + + assert observed == [(name, name, name) for name in ("first", "second", "first", "second")] + spans = memory_span_exporter.get_finished_spans() + node_spans = [span for span in spans if (span.attributes or {}).get("node.id") == "outer-start"] + assert len(node_spans) == 2 + assert [span.parent for span in node_spans] == parents + + +_workflow_tool_event_delivery_container_handler = _workflow_tool_container_container_handler +_workflow_tool_event_delivery_outer_graph = _workflow_tool_container_outer_graph +_workflow_tool_event_delivery_workflow_tool_node = _workflow_tool_container_workflow_tool_node + + +class TestWorkflowToolEventDeliveryCases: + @staticmethod + def test_workflow_tool_loop_break_does_not_escape_into_enclosing_loop(monkeypatch: pytest.MonkeyPatch) -> None: + def node(node_id: str, node_type: str, **data: object) -> dict[str, object]: + return {"id": node_id, "data": {"type": node_type, "title": node_id, **data}} + + def edge(source: str, target: str, handle: str = "source") -> dict[str, str]: + return {"id": f"{source}-{target}", "source": source, "target": target, "sourceHandle": handle} + + tool_node, runtime, _ = _workflow_tool_event_delivery_workflow_tool_node() + monkeypatch.setattr("core.workflow.node_factory.DifyToolNodeRuntime", lambda _, **_kwargs: runtime) + outer_graph = { + "nodes": [ + node("start", "start", variables=[]), + node( + "outer-loop", + "loop", + loop_count=3, + start_node_id="outer-loop-start", + break_conditions=[], + logical_operator="and", + loop_variables=[ + {"id": "guard", "label": "guard", "var_type": "number", "value_type": "constant", "value": 0} + ], + ), + node("outer-loop-start", "loop-start", loop_id="outer-loop"), + node( + "condition", + "if-else", + loop_id="outer-loop", + cases=[ + { + "case_id": "true", + "logical_operator": "and", + "conditions": [ + {"comparison_operator": "=", "variable_selector": ["outer-loop", "guard"], "value": "1"} + ], + } + ], + ), + {"id": "tool", "data": {**tool_node.node_data.model_dump(), "loop_id": "outer-loop"}}, + node("shared-loop-end", "loop-end", loop_id="outer-loop"), + node( + "end", + "end", + outputs=[ + {"variable": "rounds", "value_selector": ["outer-loop", "loop_round"], "value_type": "number"} + ], + ), + ], + "edges": [ + edge("start", "outer-loop"), + edge("outer-loop", "end"), + edge("outer-loop-start", "condition"), + edge("condition", "tool", "false"), + edge("condition", "shared-loop-end", "true"), + ], + } + _, _, _, _, repository = _workflow_tool_event_delivery_container_handler() + repository["source-workflow"] = replace( + repository["source-workflow"], + graph_config={ + "nodes": [ + node("source-start", "start", variables=[]), + node( + "source-loop", + "loop", + loop_count=3, + start_node_id="source-loop-start", + break_conditions=[], + logical_operator="and", + ), + node("source-loop-start", "loop-start", loop_id="source-loop"), + node("shared-loop-end", "loop-end", loop_id="source-loop"), + node("source-end", "end", outputs=[]), + ], + "edges": [ + edge("source-start", "source-loop"), + edge("source-loop", "source-end"), + edge("source-loop-start", "shared-loop-end"), + ], + }, + ) + state = tool_node.runtime_state + graph = Graph.init( + graph_config=outer_graph, + node_factory=DifyNodeFactory( + init_params=build_test_graph_init_params(workflow_id="outer-workflow", graph_config=outer_graph), + runtime_state=state, + ), + root_node_id="start", + ) + persisted: list[_workflow_tool_event_delivery_NodeEvent] = [] + events = list( + Engine( + graph=graph, + runtime_state=state, + workers=1, + container_handler_factories=( + partial(WorkflowToolNestedContainerHandler, handler_factory=LoopContainerHandler), + partial( + WorkflowToolContainerHandler, + sources=repository, + event_listener_factory=lambda *_: persisted.append, + ), + ), + ).run() + ) + + assert isinstance(events[-1], _workflow_tool_event_delivery_GraphRunSucceededEvent) + assert state.outputs == {"rounds": 3} + assert [ + event.node_run_result.outputs["loop_round"].to_object() + for event in persisted + if isinstance(event, NodeRunSucceededEvent) and event.node_id == "source-loop" + ] == [1, 1, 1] + + @staticmethod + def test_workflow_tool_persists_loop_outputs_after_graphon_normalizes_them() -> None: + _, frames, state, request, repository = _workflow_tool_event_delivery_container_handler() + source = repository["source-workflow"] + repository["source-workflow"] = replace( + source, + graph_config={ + "nodes": [ + {"id": "start", "data": {"type": "start", "title": "Start", "variables": []}}, + { + "id": "loop", + "data": { + "type": "loop", + "title": "Loop", + "loop_count": 1, + "start_node_id": "loop-start", + "break_conditions": [], + "logical_operator": "and", + }, + }, + {"id": "loop-start", "data": {"type": "loop-start", "title": "Loop start", "loop_id": "loop"}}, + ], + "edges": [{"id": "start-loop", "source": "start", "target": "loop"}], + }, + ) + persisted: list[_workflow_tool_event_delivery_NodeEvent] = [] + handler = WorkflowToolContainerHandler( + frames, sources=repository, event_listener_factory=lambda *_: persisted.append + ) + handler.handle_request(invocation_id="invocation", request=request) + frame = frames["invocation:workflow-tool"] + frame.state.variable_pool.add(("loop", "counter"), 2) + processor = NodeEventProcessor( + graph_execution=state.graph_execution, + event_stream=EventStream(layers=[]), + frame_registry=frames, + container_handlers={BuiltinNodeTypes.TOOL: handler}, + ) + event = NodeRunSucceededEvent( + id="loop-execution", + node_id="loop", + node_type=BuiltinNodeTypes.LOOP, + start_at=datetime.now(UTC).replace(tzinfo=None), + node_run_result=NodeRunResult(status=WorkflowNodeExecutionStatus.SUCCEEDED, outputs={"counter": 1}), + ) + + processor.dispatch(NodeEventTask(frame_id=frame.frame_id, event=event)) + + assert {key: value.to_object() for key, value in event.node_run_result.outputs.items()} == {"counter": 2} + assert {key: value.to_object() for key, value in persisted[0].node_run_result.outputs.items()} == {"counter": 2} + assert persisted[0].container_id == "" + + @staticmethod + @pytest.mark.parametrize("listener_fails", [False, True]) + def test_workflow_tool_delivers_source_events_to_persistence_without_exposing_them( + listener_fails: bool, caplog: pytest.LogCaptureFixture + ) -> None: + node, _, payload = _workflow_tool_event_delivery_workflow_tool_node() + source = WorkflowToolSource( + app_id=payload.source_app_id, + workflow_id=payload.source_workflow_id, + graph_config={ + "nodes": [ + {"id": "source-start", "data": {"type": "start", "title": "Start", "variables": []}}, + {"id": "source-end", "data": {"type": "end", "title": "End", "outputs": []}}, + ], + "edges": [{"id": "source-edge", "source": "source-start", "target": "source-end"}], + }, + features_dict={}, + environment_variables=[], + workflow_kind=WorkflowKind.STANDARD, + ) + repository = {} + repository["source-workflow"] = source + persisted: list[_workflow_tool_event_delivery_NodeEvent] = [] + + def persist_event(event: _workflow_tool_event_delivery_NodeEvent) -> None: + persisted.append(event) + if listener_fails: + raise RuntimeError("Trace store unavailable") + + listener_requests: list[tuple[WorkflowToolSource, str, str]] = [] + + def listener_factory( + tool_source: WorkflowToolSource, workflow_id: str, execution_id: str + ) -> Callable[[_workflow_tool_event_delivery_NodeEvent], None]: + listener_requests.append((tool_source, workflow_id, execution_id)) + return persist_event + + events = list( + Engine( + graph=_workflow_tool_event_delivery_outer_graph(node), + runtime_state=node.runtime_state, + command_channel=InMemoryChannel(), + workers=1, + container_handler_factories=( + partial( + WorkflowToolContainerHandler, + sources=repository, + event_listener_factory=listener_factory, + ), + ), + ).run() + ) + + assert isinstance(events[-1], _workflow_tool_event_delivery_GraphRunSucceededEvent) + caller_event = next( + event for event in events if isinstance(event, NodeRunStartedEvent) and event.node_id == "tool" + ) + assert listener_requests == [(source, node.workflow_id, caller_event.id)] + assert [event.node_id for event in persisted if isinstance(event, NodeRunStartedEvent)] == [ + "source-start", + "source-end", + ] + assert [event.node_id for event in persisted if isinstance(event, NodeRunSucceededEvent)] == [ + "source-start", + "source-end", + ] + assert all(not event.container_id for event in persisted) + completed_tool = next( + event for event in events if isinstance(event, NodeRunSucceededEvent) and event.node_id == node.id + ) + assert completed_tool.node_run_result.outputs["json"].to_object() == [{}] + assert len({event.node_run_result.process_data[WORKFLOW_TOOL_INVOCATION_ID_KEY] for event in persisted}) == 1 + assert all( + event.node_id not in {"source-start", "source-end"} + for event in events + if isinstance(event, _workflow_tool_event_delivery_NodeEvent) + ) + if listener_fails: + assert "Failed to persist Workflow Tool event" in caplog.text + + @staticmethod + @pytest.mark.parametrize("caller_save_fails", [False, True]) + def test_workflow_tool_agent_finds_its_persisted_caller_before_resolving_binding( + monkeypatch: pytest.MonkeyPatch, + sqlite_session_factory: sessionmaker[Session], + caller_save_fails: bool, + ) -> None: + binding_resolver = WorkflowAgentBindingResolver() + resolution_calls: list[dict[str, object]] = [] + if caller_save_fails: + + def reject_resolution(**kwargs: object) -> None: + resolution_calls.append(kwargs) + pytest.fail("Binding resolution must wait for a persisted caller") + + monkeypatch.setattr(binding_resolver, "resolve", reject_resolution) + monkeypatch.setattr("core.workflow.node_factory.WorkflowAgentBindingResolver", lambda: binding_resolver) + save = SQLAlchemyWorkflowNodeExecutionRepository.save + + def fail_caller_save( + repository: SQLAlchemyWorkflowNodeExecutionRepository, execution: WorkflowNodeExecution + ) -> None: + if execution.node_id == "source-agent" and execution.status == WorkflowNodeExecutionStatus.RUNNING: + raise RuntimeError("Caller store unavailable") + save(repository, execution) + + monkeypatch.setattr(SQLAlchemyWorkflowNodeExecutionRepository, "save", fail_caller_save) + monkeypatch.setattr( + "clients.agent_backend.factory.create_agent_backend_run_client", + lambda **_kwargs: FakeAgentBackendRunClient(), + ) + node, _, payload = _workflow_tool_event_delivery_workflow_tool_node() + source = WorkflowToolSource( + app_id=payload.source_app_id, + workflow_id=payload.source_workflow_id, + graph_config={ + "nodes": [ + {"id": "source-start", "data": {"type": "start", "title": "Start", "variables": []}}, + { + "id": "source-agent", + "data": { + "type": "agent", + "version": "2", + "agent_node_kind": "dify_agent", + "title": "Unconfigured Agent", + "error_strategy": "default-value", + "default_value": [], + }, + }, + ], + "edges": [{"id": "source-edge", "source": "source-start", "target": "source-agent"}], + }, + features_dict={}, + environment_variables=[], + workflow_kind=WorkflowKind.STANDARD, + ) + source_repository = {} + source_repository["source-workflow"] = source + user = Account(name="Test", email="test@example.com") + user.id = "user" + layer = WorkflowPersistenceLayer( + application_generate_entity=WorkflowAppGenerateEntity( + task_id="task", + app_config=WorkflowAppConfig( + app_id="outer-app", + tenant_id="tenant", + workflow_id="outer-workflow", + app_mode=AppMode.WORKFLOW, + ), + inputs={}, + files=[], + user_id=user.id, + stream=True, + invoke_from=InvokeFrom.SERVICE_API, + extras={}, + workflow_execution_id="outer-run", + ), + workflow_info=PersistenceWorkflowInfo( + workflow_id="outer-workflow", + workflow_type=WorkflowType.WORKFLOW, + version="1", + graph_data=node.init_params.graph_config, + ), + workflow_execution_repository=SQLAlchemyWorkflowExecutionRepository( + session_factory=sqlite_session_factory, + tenant_id="tenant", + user=user, + app_id="outer-app", + triggered_from=WorkflowRunTriggeredFrom.APP_RUN, + ), + workflow_node_execution_repository=SQLAlchemyWorkflowNodeExecutionRepository( + session_factory=sqlite_session_factory, + tenant_id="tenant", + user=user, + app_id="outer-app", + triggered_from=WorkflowNodeExecutionTriggeredFrom.WORKFLOW_RUN, + ), + ) + index = WorkflowRunIndex(()) + index.seed_source(source.app_id, source.workflow_id, ()) + layer.set_node_run_indices(index.indices) + node.runtime_state.variable_pool.add(("sys", "workflow_run_id"), "outer-run") + engine = Engine( + graph=_workflow_tool_event_delivery_outer_graph(node), + runtime_state=node.runtime_state, + command_channel=InMemoryChannel(), + workers=1, + container_handler_factories=( + partial( + WorkflowToolContainerHandler, + sources=source_repository, + event_listener_factory=layer.create_workflow_tool_event_listener, + ), + ), + ) + engine.add_layer(index) + engine.add_layer(layer) + + events = list(engine.run()) + + assert isinstance(events[-1], _workflow_tool_event_delivery_GraphRunSucceededEvent) + with sqlite_session_factory() as session: + execution = session.scalars( + select(WorkflowNodeExecutionModel).where(WorkflowNodeExecutionModel.node_id == "source-agent") + ).one() + assert (execution.app_id, execution.workflow_id, execution.workflow_run_id) == ( + "source-app", + "source-workflow", + "outer-run", + ) + assert execution.triggered_from == WorkflowNodeExecutionTriggeredFrom.WORKFLOW_TOOL + caller = session.scalars( + select(WorkflowNodeExecutionModel).where(WorkflowNodeExecutionModel.node_id == "tool") + ).one() + assert (execution.triggered_from_workflow_id, execution.triggered_from_node_execution_id) == ( + "outer-workflow", + caller.node_execution_id, + ) + assert execution.status == "exception" + if caller_save_fails: + assert resolution_calls == [] + assert execution.error == "Workflow node execution caller is unavailable" + assert execution.outputs_dict["error_type"] == "agent_workflow_node_runtime_error" + else: + # The real workspace store must see the caller synchronously before + # the binding resolver reports this deliberately unconfigured Agent. + assert execution.error == "Workflow Agent binding not found for node source-agent." + assert execution.outputs_dict["error_type"] == "agent_binding_not_found" + run = session.scalars(select(WorkflowRun)).one() + assert (run.app_id, run.workflow_id, run.status) == ("outer-app", "outer-workflow", "succeeded") + assert not any( + isinstance(event, _workflow_tool_event_delivery_NodeEvent) and event.node_id == "source-agent" + for event in events + ) + + +_workflow_tool_file_access_workflow_tool_node = _workflow_tool_container_workflow_tool_node + + +class TestWorkflowToolFileAccessCases: + @staticmethod + @pytest.mark.parametrize("input_kind", ["system_files", "file", "file-list"]) + @pytest.mark.parametrize("container_type", [None, "loop", "iteration"]) + @pytest.mark.parametrize( + ("transfer_method", "permission"), + [ + ("local_file", "owned"), + ("local_file", "other-user"), + ("local_file", "granted"), + ("tool_file", "owned"), + ("tool_file", "other-user"), + ], + ) + def test_workflow_tool_dispatcher_enforces_file_ownership( + sqlite_session: Session, + monkeypatch: pytest.MonkeyPatch, + input_kind: str, + container_type: str | None, + transfer_method: str, + permission: str, + ) -> None: + owner_id = "user" if permission == "owned" else "other-user" + stored_file: UploadFile | ToolFile + if transfer_method == "local_file": + stored_file = UploadFile( + tenant_id="tenant", + storage_type=StorageType.LOCAL, + key="stored-file.txt", + name="attachment.txt", + size=12, + extension="txt", + mime_type="text/plain", + created_by=owner_id, + created_by_role=CreatorUserRole.ACCOUNT if permission == "granted" else CreatorUserRole.END_USER, + created_at=datetime.now(UTC), + used=False, + ) + else: + stored_file = ToolFile( + tenant_id="tenant", + user_id=owner_id, + conversation_id=None, + file_key="stored-file.txt", + name="attachment.txt", + mimetype="text/plain", + size=12, + ) + sqlite_session.add(stored_file) + sqlite_session.commit() + mapping = {"transfer_method": transfer_method, "type": "document", "related_id": stored_file.id} + scope = FileAccessScope( + tenant_id="tenant", + user_id="user", + user_from=UserFrom.END_USER, + invoke_from=InvokeFrom.SERVICE_API, + granted_upload_file_ids=frozenset({stored_file.id}) if permission == "granted" else frozenset(), + ) + with bind_file_access_scope(scope): + execution_context_layer = ExecutionContextLayer() + runtime_state = RuntimeState( + workflow_id="outer-workflow", + variable_pool=VariablePool(), + start_at=1, + ) + template_node, runtime, payload = _workflow_tool_file_access_workflow_tool_node(runtime_state) + monkeypatch.setattr("core.workflow.node_factory.DifyToolNodeRuntime", lambda _, **_kwargs: runtime) + tool_data = template_node.node_data.model_dump(mode="python") + nodes = [ + {"id": "start", "data": {"type": "start", "title": "Start", "variables": list[object]()}}, + {"id": "tool", "data": tool_data}, + ] + edges = [{"id": "outer", "source": "start", "target": "tool"}] + if container_type is not None: + tool_data[f"{container_type}_id"] = container_type + container_data: dict[str, object] = { + "type": container_type, + "title": container_type, + "start_node_id": "body-start", + } + if container_type == "loop": + container_data.update(loop_count=1, break_conditions=[], logical_operator="and") + else: + runtime_state.variable_pool.add(("start", "items"), ["one"]) + container_data.update( + iterator_selector=["start", "items"], output_selector=["tool", "text"], output_type="array[string]" + ) + nodes.extend( + [ + {"id": container_type, "data": container_data}, + { + "id": "body-start", + "data": { + "type": f"{container_type}-start", + "title": "Body start", + f"{container_type}_id": container_type, + }, + }, + ] + ) + edges = [ + {"id": "outer", "source": "start", "target": container_type}, + {"id": "inner", "source": "body-start", "target": "tool"}, + ] + graph_config = {"nodes": nodes, "edges": edges} + node_factory = DifyNodeFactory( + init_params=template_node.init_params.model_copy( + update={ + "graph_config": graph_config, + "run_context": build_test_run_context( + app_id="outer-app", user_from=UserFrom.END_USER, invoke_from=InvokeFrom.SERVICE_API + ), + } + ), + runtime_state=runtime_state, + ) + variables: list[dict[str, object]] = [] + if input_kind == "system_files": + payload = payload.model_copy(update={"system_files": [mapping], "inputs": {}}) + output_selector = ["sys", "files"] + else: + variables = [ + { + "variable": "attachment", + "label": "Attachment", + "type": input_kind, + "allowed_file_types": ["document"], + "allowed_file_upload_methods": ["local_file"], + } + ] + value = mapping if input_kind == "file" else [mapping] + payload = payload.model_copy(update={"inputs": {"attachment": value}}) + output_selector = ["source-start", "attachment"] + monkeypatch.setattr(runtime, "build_workflow_tool_container_payload", lambda **_kwargs: payload) + source = WorkflowToolSource( + app_id="source-app", + workflow_id="source-workflow", + graph_config={ + "nodes": [ + {"id": "source-start", "data": {"type": "start", "title": "Start", "variables": variables}}, + { + "id": "source-end", + "data": { + "type": "end", + "title": "End", + "outputs": [{"variable": "accepted", "value_selector": output_selector}], + }, + }, + ], + "edges": [{"id": "source-edge", "source": "source-start", "target": "source-end"}], + }, + features_dict={}, + environment_variables=[], + workflow_kind=WorkflowKind.STANDARD, + ) + sources = {source.workflow_id: source} + engine = Engine( + graph=Graph.init(graph_config=graph_config, node_factory=node_factory, root_node_id="start"), + runtime_state=runtime_state, + command_channel=InMemoryChannel(), + workers=1, + container_handler_factories=( + partial( + WorkflowToolContainerHandler, + sources=sources, + execution_context_factory=execution_context_layer.enter_context, + ), + ), + ) + engine.add_layer(execution_context_layer) + + if permission == "other-user": + expected_error = ( + "Invalid upload file" if transfer_method == "local_file" else f"ToolFile {stored_file.id} not found" + ) + with pytest.raises(RuntimeError, match=expected_error): + list(engine.run()) + else: + events = list(engine.run()) + assert isinstance(events[-1], _workflow_tool_file_access_GraphRunSucceededEvent) + tool_completed = next( + event for event in events if isinstance(event, NodeRunSucceededEvent) and event.node_id == "tool" + ) + files_output = tool_completed.node_run_result.outputs["files"] + assert isinstance(files_output, ArrayFileSegment) + accepted_files = files_output.value + assert len(accepted_files) == 1 + assert accepted_files[0].filename == stored_file.name diff --git a/api/tests/unit_tests/core/workflow/test_workflow_tool_event_delivery.py b/api/tests/unit_tests/core/workflow/test_workflow_tool_event_delivery.py deleted file mode 100644 index 137b69fe621..00000000000 --- a/api/tests/unit_tests/core/workflow/test_workflow_tool_event_delivery.py +++ /dev/null @@ -1,446 +0,0 @@ -from collections.abc import Callable -from dataclasses import replace -from datetime import UTC, datetime -from functools import partial - -import pytest -from sqlalchemy import select -from sqlalchemy.orm import Session, sessionmaker - -from clients.agent_backend import FakeAgentBackendRunClient -from core.app.apps.workflow.app_config_manager import WorkflowAppConfig -from core.app.entities.app_invoke_entities import InvokeFrom, WorkflowAppGenerateEntity -from core.repositories.sqlalchemy_workflow_execution_repository import SQLAlchemyWorkflowExecutionRepository -from core.repositories.sqlalchemy_workflow_node_execution_repository import SQLAlchemyWorkflowNodeExecutionRepository -from core.tools.workflow_as_tool.repository import WorkflowToolSource -from core.workflow.node_execution_process_data import ( - WORKFLOW_TOOL_INVOCATION_ID_KEY, -) -from core.workflow.node_factory import DifyNodeFactory -from core.workflow.nodes.agent_v2.binding_resolver import WorkflowAgentBindingResolver -from core.workflow.workflow_tool_container_handler import ( - WorkflowToolContainerHandler, - WorkflowToolNestedContainerHandler, -) -from enums import WorkflowKind -from graphon.engine import Engine -from graphon.engine.command import InMemoryChannel -from graphon.engine.container_handler.builtin.loop import LoopContainerHandler -from graphon.engine.event.processor import NodeEventProcessor -from graphon.engine.event.stream import EventStream -from graphon.engine.worker import NodeEventTask -from graphon.engine_events import ( - GraphRunSucceededEvent, - NodeRunStartedEvent, - NodeRunSucceededEvent, -) -from graphon.engine_events.base import NodeEvent -from graphon.entities import WorkflowNodeExecution -from graphon.enums import BuiltinNodeTypes, WorkflowNodeExecutionStatus, WorkflowType -from graphon.graph import Graph -from graphon.node_events import NodeRunResult -from models import Account, WorkflowRun -from models.enums import WorkflowRunTriggeredFrom -from models.model import AppMode -from models.workflow import WorkflowNodeExecutionModel, WorkflowNodeExecutionTriggeredFrom -from services.workflow_persistence import PersistenceWorkflowInfo, WorkflowPersistenceLayer -from services.workflow_run_index import WorkflowRunIndex -from tests.unit_tests.core.workflow.test_workflow_tool_container import ( - _container_handler, - _outer_graph, - _workflow_tool_node, -) -from tests.workflow_test_utils import build_test_graph_init_params - - -def test_workflow_tool_loop_break_does_not_escape_into_enclosing_loop(monkeypatch: pytest.MonkeyPatch) -> None: - def node(node_id: str, node_type: str, **data: object) -> dict[str, object]: - return {"id": node_id, "data": {"type": node_type, "title": node_id, **data}} - - def edge(source: str, target: str, handle: str = "source") -> dict[str, str]: - return {"id": f"{source}-{target}", "source": source, "target": target, "sourceHandle": handle} - - tool_node, runtime, _ = _workflow_tool_node() - monkeypatch.setattr("core.workflow.node_factory.DifyToolNodeRuntime", lambda _, **_kwargs: runtime) - outer_graph = { - "nodes": [ - node("start", "start", variables=[]), - node( - "outer-loop", - "loop", - loop_count=3, - start_node_id="outer-loop-start", - break_conditions=[], - logical_operator="and", - loop_variables=[ - {"id": "guard", "label": "guard", "var_type": "number", "value_type": "constant", "value": 0} - ], - ), - node("outer-loop-start", "loop-start", loop_id="outer-loop"), - node( - "condition", - "if-else", - loop_id="outer-loop", - cases=[ - { - "case_id": "true", - "logical_operator": "and", - "conditions": [ - {"comparison_operator": "=", "variable_selector": ["outer-loop", "guard"], "value": "1"} - ], - } - ], - ), - {"id": "tool", "data": {**tool_node.node_data.model_dump(), "loop_id": "outer-loop"}}, - node("shared-loop-end", "loop-end", loop_id="outer-loop"), - node( - "end", - "end", - outputs=[ - {"variable": "rounds", "value_selector": ["outer-loop", "loop_round"], "value_type": "number"} - ], - ), - ], - "edges": [ - edge("start", "outer-loop"), - edge("outer-loop", "end"), - edge("outer-loop-start", "condition"), - edge("condition", "tool", "false"), - edge("condition", "shared-loop-end", "true"), - ], - } - _, _, _, _, repository = _container_handler() - repository["source-workflow"] = replace( - repository["source-workflow"], - graph_config={ - "nodes": [ - node("source-start", "start", variables=[]), - node( - "source-loop", - "loop", - loop_count=3, - start_node_id="source-loop-start", - break_conditions=[], - logical_operator="and", - ), - node("source-loop-start", "loop-start", loop_id="source-loop"), - node("shared-loop-end", "loop-end", loop_id="source-loop"), - node("source-end", "end", outputs=[]), - ], - "edges": [ - edge("source-start", "source-loop"), - edge("source-loop", "source-end"), - edge("source-loop-start", "shared-loop-end"), - ], - }, - ) - state = tool_node.runtime_state - graph = Graph.init( - graph_config=outer_graph, - node_factory=DifyNodeFactory( - init_params=build_test_graph_init_params(workflow_id="outer-workflow", graph_config=outer_graph), - runtime_state=state, - ), - root_node_id="start", - ) - persisted: list[NodeEvent] = [] - events = list( - Engine( - graph=graph, - runtime_state=state, - workers=1, - container_handler_factories=( - partial(WorkflowToolNestedContainerHandler, handler_factory=LoopContainerHandler), - partial( - WorkflowToolContainerHandler, - sources=repository, - event_listener_factory=lambda *_: persisted.append, - ), - ), - ).run() - ) - - assert isinstance(events[-1], GraphRunSucceededEvent) - assert state.outputs == {"rounds": 3} - assert [ - event.node_run_result.outputs["loop_round"].to_object() - for event in persisted - if isinstance(event, NodeRunSucceededEvent) and event.node_id == "source-loop" - ] == [1, 1, 1] - - -def test_workflow_tool_persists_loop_outputs_after_graphon_normalizes_them() -> None: - _, frames, state, request, repository = _container_handler() - source = repository["source-workflow"] - repository["source-workflow"] = replace( - source, - graph_config={ - "nodes": [ - {"id": "start", "data": {"type": "start", "title": "Start", "variables": []}}, - { - "id": "loop", - "data": { - "type": "loop", - "title": "Loop", - "loop_count": 1, - "start_node_id": "loop-start", - "break_conditions": [], - "logical_operator": "and", - }, - }, - {"id": "loop-start", "data": {"type": "loop-start", "title": "Loop start", "loop_id": "loop"}}, - ], - "edges": [{"id": "start-loop", "source": "start", "target": "loop"}], - }, - ) - persisted: list[NodeEvent] = [] - handler = WorkflowToolContainerHandler( - frames, sources=repository, event_listener_factory=lambda *_: persisted.append - ) - handler.handle_request(invocation_id="invocation", request=request) - frame = frames["invocation:workflow-tool"] - frame.state.variable_pool.add(("loop", "counter"), 2) - processor = NodeEventProcessor( - graph_execution=state.graph_execution, - event_stream=EventStream(layers=[]), - frame_registry=frames, - container_handlers={BuiltinNodeTypes.TOOL: handler}, - ) - event = NodeRunSucceededEvent( - id="loop-execution", - node_id="loop", - node_type=BuiltinNodeTypes.LOOP, - start_at=datetime.now(UTC).replace(tzinfo=None), - node_run_result=NodeRunResult(status=WorkflowNodeExecutionStatus.SUCCEEDED, outputs={"counter": 1}), - ) - - processor.dispatch(NodeEventTask(frame_id=frame.frame_id, event=event)) - - assert {key: value.to_object() for key, value in event.node_run_result.outputs.items()} == {"counter": 2} - assert {key: value.to_object() for key, value in persisted[0].node_run_result.outputs.items()} == {"counter": 2} - assert persisted[0].container_id == "" - - -@pytest.mark.parametrize("listener_fails", [False, True]) -def test_workflow_tool_delivers_source_events_to_persistence_without_exposing_them( - listener_fails: bool, caplog: pytest.LogCaptureFixture -) -> None: - node, _, payload = _workflow_tool_node() - source = WorkflowToolSource( - app_id=payload.source_app_id, - workflow_id=payload.source_workflow_id, - graph_config={ - "nodes": [ - {"id": "source-start", "data": {"type": "start", "title": "Start", "variables": []}}, - {"id": "source-end", "data": {"type": "end", "title": "End", "outputs": []}}, - ], - "edges": [{"id": "source-edge", "source": "source-start", "target": "source-end"}], - }, - features_dict={}, - environment_variables=[], - workflow_kind=WorkflowKind.STANDARD, - ) - repository = {} - repository["source-workflow"] = source - persisted: list[NodeEvent] = [] - - def persist_event(event: NodeEvent) -> None: - persisted.append(event) - if listener_fails: - raise RuntimeError("Trace store unavailable") - - listener_requests: list[tuple[WorkflowToolSource, str, str]] = [] - - def listener_factory( - tool_source: WorkflowToolSource, workflow_id: str, execution_id: str - ) -> Callable[[NodeEvent], None]: - listener_requests.append((tool_source, workflow_id, execution_id)) - return persist_event - - events = list( - Engine( - graph=_outer_graph(node), - runtime_state=node.runtime_state, - command_channel=InMemoryChannel(), - workers=1, - container_handler_factories=( - partial( - WorkflowToolContainerHandler, - sources=repository, - event_listener_factory=listener_factory, - ), - ), - ).run() - ) - - assert isinstance(events[-1], GraphRunSucceededEvent) - caller_event = next(event for event in events if isinstance(event, NodeRunStartedEvent) and event.node_id == "tool") - assert listener_requests == [(source, node.workflow_id, caller_event.id)] - assert [event.node_id for event in persisted if isinstance(event, NodeRunStartedEvent)] == [ - "source-start", - "source-end", - ] - assert [event.node_id for event in persisted if isinstance(event, NodeRunSucceededEvent)] == [ - "source-start", - "source-end", - ] - assert all(not event.container_id for event in persisted) - completed_tool = next( - event for event in events if isinstance(event, NodeRunSucceededEvent) and event.node_id == node.id - ) - assert completed_tool.node_run_result.outputs["json"].to_object() == [{}] - assert len({event.node_run_result.process_data[WORKFLOW_TOOL_INVOCATION_ID_KEY] for event in persisted}) == 1 - assert all(event.node_id not in {"source-start", "source-end"} for event in events if isinstance(event, NodeEvent)) - if listener_fails: - assert "Failed to persist Workflow Tool event" in caplog.text - - -@pytest.mark.parametrize("caller_save_fails", [False, True]) -def test_workflow_tool_agent_finds_its_persisted_caller_before_resolving_binding( - monkeypatch: pytest.MonkeyPatch, - sqlite_session_factory: sessionmaker[Session], - caller_save_fails: bool, -) -> None: - binding_resolver = WorkflowAgentBindingResolver() - resolution_calls: list[dict[str, object]] = [] - if caller_save_fails: - - def reject_resolution(**kwargs: object) -> None: - resolution_calls.append(kwargs) - pytest.fail("Binding resolution must wait for a persisted caller") - - monkeypatch.setattr(binding_resolver, "resolve", reject_resolution) - monkeypatch.setattr("core.workflow.node_factory.WorkflowAgentBindingResolver", lambda: binding_resolver) - save = SQLAlchemyWorkflowNodeExecutionRepository.save - - def fail_caller_save( - repository: SQLAlchemyWorkflowNodeExecutionRepository, execution: WorkflowNodeExecution - ) -> None: - if execution.node_id == "source-agent" and execution.status == WorkflowNodeExecutionStatus.RUNNING: - raise RuntimeError("Caller store unavailable") - save(repository, execution) - - monkeypatch.setattr(SQLAlchemyWorkflowNodeExecutionRepository, "save", fail_caller_save) - monkeypatch.setattr( - "clients.agent_backend.factory.create_agent_backend_run_client", lambda **_kwargs: FakeAgentBackendRunClient() - ) - node, _, payload = _workflow_tool_node() - source = WorkflowToolSource( - app_id=payload.source_app_id, - workflow_id=payload.source_workflow_id, - graph_config={ - "nodes": [ - {"id": "source-start", "data": {"type": "start", "title": "Start", "variables": []}}, - { - "id": "source-agent", - "data": { - "type": "agent", - "version": "2", - "agent_node_kind": "dify_agent", - "title": "Unconfigured Agent", - "error_strategy": "default-value", - "default_value": [], - }, - }, - ], - "edges": [{"id": "source-edge", "source": "source-start", "target": "source-agent"}], - }, - features_dict={}, - environment_variables=[], - workflow_kind=WorkflowKind.STANDARD, - ) - source_repository = {} - source_repository["source-workflow"] = source - user = Account(name="Test", email="test@example.com") - user.id = "user" - layer = WorkflowPersistenceLayer( - application_generate_entity=WorkflowAppGenerateEntity( - task_id="task", - app_config=WorkflowAppConfig( - app_id="outer-app", - tenant_id="tenant", - workflow_id="outer-workflow", - app_mode=AppMode.WORKFLOW, - ), - inputs={}, - files=[], - user_id=user.id, - stream=True, - invoke_from=InvokeFrom.SERVICE_API, - extras={}, - workflow_execution_id="outer-run", - ), - workflow_info=PersistenceWorkflowInfo( - workflow_id="outer-workflow", - workflow_type=WorkflowType.WORKFLOW, - version="1", - graph_data=node.init_params.graph_config, - ), - workflow_execution_repository=SQLAlchemyWorkflowExecutionRepository( - session_factory=sqlite_session_factory, - tenant_id="tenant", - user=user, - app_id="outer-app", - triggered_from=WorkflowRunTriggeredFrom.APP_RUN, - ), - workflow_node_execution_repository=SQLAlchemyWorkflowNodeExecutionRepository( - session_factory=sqlite_session_factory, - tenant_id="tenant", - user=user, - app_id="outer-app", - triggered_from=WorkflowNodeExecutionTriggeredFrom.WORKFLOW_RUN, - ), - ) - index = WorkflowRunIndex(()) - index.seed_source(source.app_id, source.workflow_id, ()) - layer.set_node_run_indices(index.indices) - node.runtime_state.variable_pool.add(("sys", "workflow_run_id"), "outer-run") - engine = Engine( - graph=_outer_graph(node), - runtime_state=node.runtime_state, - command_channel=InMemoryChannel(), - workers=1, - container_handler_factories=( - partial( - WorkflowToolContainerHandler, - sources=source_repository, - event_listener_factory=layer.create_workflow_tool_event_listener, - ), - ), - ) - engine.add_layer(index) - engine.add_layer(layer) - - events = list(engine.run()) - - assert isinstance(events[-1], GraphRunSucceededEvent) - with sqlite_session_factory() as session: - execution = session.scalars( - select(WorkflowNodeExecutionModel).where(WorkflowNodeExecutionModel.node_id == "source-agent") - ).one() - assert (execution.app_id, execution.workflow_id, execution.workflow_run_id) == ( - "source-app", - "source-workflow", - "outer-run", - ) - assert execution.triggered_from == WorkflowNodeExecutionTriggeredFrom.WORKFLOW_TOOL - caller = session.scalars( - select(WorkflowNodeExecutionModel).where(WorkflowNodeExecutionModel.node_id == "tool") - ).one() - assert (execution.triggered_from_workflow_id, execution.triggered_from_node_execution_id) == ( - "outer-workflow", - caller.node_execution_id, - ) - assert execution.status == "exception" - if caller_save_fails: - assert resolution_calls == [] - assert execution.error == "Workflow node execution caller is unavailable" - assert execution.outputs_dict["error_type"] == "agent_workflow_node_runtime_error" - else: - # The real workspace store must see the caller synchronously before - # the binding resolver reports this deliberately unconfigured Agent. - assert execution.error == "Workflow Agent binding not found for node source-agent." - assert execution.outputs_dict["error_type"] == "agent_binding_not_found" - run = session.scalars(select(WorkflowRun)).one() - assert (run.app_id, run.workflow_id, run.status) == ("outer-app", "outer-workflow", "succeeded") - assert not any(isinstance(event, NodeEvent) and event.node_id == "source-agent" for event in events) diff --git a/api/tests/unit_tests/core/workflow/test_workflow_tool_file_access.py b/api/tests/unit_tests/core/workflow/test_workflow_tool_file_access.py deleted file mode 100644 index 9e175003c5c..00000000000 --- a/api/tests/unit_tests/core/workflow/test_workflow_tool_file_access.py +++ /dev/null @@ -1,212 +0,0 @@ -from datetime import UTC, datetime -from functools import partial - -import pytest -from sqlalchemy.orm import Session - -from core.app.entities.app_invoke_entities import InvokeFrom, UserFrom -from core.app.file_access import FileAccessScope, bind_file_access_scope -from core.app.layers.execution_context_layer import ExecutionContextLayer -from core.tools.workflow_as_tool.repository import WorkflowToolSource -from core.workflow.node_factory import DifyNodeFactory -from core.workflow.workflow_tool_container_handler import WorkflowToolContainerHandler -from enums import WorkflowKind -from extensions.storage.storage_type import StorageType -from graphon.engine import Engine -from graphon.engine.command import InMemoryChannel -from graphon.engine_events import GraphRunSucceededEvent, NodeRunSucceededEvent -from graphon.graph import Graph -from graphon.runtime import RuntimeState, VariablePool -from graphon.variables.segments import ArrayFileSegment -from models import ToolFile, UploadFile -from models.enums import CreatorUserRole -from tests.unit_tests.core.workflow.test_workflow_tool_container import _workflow_tool_node -from tests.workflow_test_utils import build_test_run_context - - -@pytest.mark.parametrize("input_kind", ["system_files", "file", "file-list"]) -@pytest.mark.parametrize("container_type", [None, "loop", "iteration"]) -@pytest.mark.parametrize( - ("transfer_method", "permission"), - [ - ("local_file", "owned"), - ("local_file", "other-user"), - ("local_file", "granted"), - ("tool_file", "owned"), - ("tool_file", "other-user"), - ], -) -def test_workflow_tool_dispatcher_enforces_file_ownership( - sqlite_session: Session, - monkeypatch: pytest.MonkeyPatch, - input_kind: str, - container_type: str | None, - transfer_method: str, - permission: str, -) -> None: - owner_id = "user" if permission == "owned" else "other-user" - stored_file: UploadFile | ToolFile - if transfer_method == "local_file": - stored_file = UploadFile( - tenant_id="tenant", - storage_type=StorageType.LOCAL, - key="stored-file.txt", - name="attachment.txt", - size=12, - extension="txt", - mime_type="text/plain", - created_by=owner_id, - created_by_role=CreatorUserRole.ACCOUNT if permission == "granted" else CreatorUserRole.END_USER, - created_at=datetime.now(UTC), - used=False, - ) - else: - stored_file = ToolFile( - tenant_id="tenant", - user_id=owner_id, - conversation_id=None, - file_key="stored-file.txt", - name="attachment.txt", - mimetype="text/plain", - size=12, - ) - sqlite_session.add(stored_file) - sqlite_session.commit() - mapping = {"transfer_method": transfer_method, "type": "document", "related_id": stored_file.id} - scope = FileAccessScope( - tenant_id="tenant", - user_id="user", - user_from=UserFrom.END_USER, - invoke_from=InvokeFrom.SERVICE_API, - granted_upload_file_ids=frozenset({stored_file.id}) if permission == "granted" else frozenset(), - ) - with bind_file_access_scope(scope): - execution_context_layer = ExecutionContextLayer() - runtime_state = RuntimeState( - workflow_id="outer-workflow", - variable_pool=VariablePool(), - start_at=1, - ) - template_node, runtime, payload = _workflow_tool_node(runtime_state) - monkeypatch.setattr("core.workflow.node_factory.DifyToolNodeRuntime", lambda _, **_kwargs: runtime) - tool_data = template_node.node_data.model_dump(mode="python") - nodes = [ - {"id": "start", "data": {"type": "start", "title": "Start", "variables": list[object]()}}, - {"id": "tool", "data": tool_data}, - ] - edges = [{"id": "outer", "source": "start", "target": "tool"}] - if container_type is not None: - tool_data[f"{container_type}_id"] = container_type - container_data: dict[str, object] = { - "type": container_type, - "title": container_type, - "start_node_id": "body-start", - } - if container_type == "loop": - container_data.update(loop_count=1, break_conditions=[], logical_operator="and") - else: - runtime_state.variable_pool.add(("start", "items"), ["one"]) - container_data.update( - iterator_selector=["start", "items"], output_selector=["tool", "text"], output_type="array[string]" - ) - nodes.extend( - [ - {"id": container_type, "data": container_data}, - { - "id": "body-start", - "data": { - "type": f"{container_type}-start", - "title": "Body start", - f"{container_type}_id": container_type, - }, - }, - ] - ) - edges = [ - {"id": "outer", "source": "start", "target": container_type}, - {"id": "inner", "source": "body-start", "target": "tool"}, - ] - graph_config = {"nodes": nodes, "edges": edges} - node_factory = DifyNodeFactory( - init_params=template_node.init_params.model_copy( - update={ - "graph_config": graph_config, - "run_context": build_test_run_context( - app_id="outer-app", user_from=UserFrom.END_USER, invoke_from=InvokeFrom.SERVICE_API - ), - } - ), - runtime_state=runtime_state, - ) - variables: list[dict[str, object]] = [] - if input_kind == "system_files": - payload = payload.model_copy(update={"system_files": [mapping], "inputs": {}}) - output_selector = ["sys", "files"] - else: - variables = [ - { - "variable": "attachment", - "label": "Attachment", - "type": input_kind, - "allowed_file_types": ["document"], - "allowed_file_upload_methods": ["local_file"], - } - ] - value = mapping if input_kind == "file" else [mapping] - payload = payload.model_copy(update={"inputs": {"attachment": value}}) - output_selector = ["source-start", "attachment"] - monkeypatch.setattr(runtime, "build_workflow_tool_container_payload", lambda **_kwargs: payload) - source = WorkflowToolSource( - app_id="source-app", - workflow_id="source-workflow", - graph_config={ - "nodes": [ - {"id": "source-start", "data": {"type": "start", "title": "Start", "variables": variables}}, - { - "id": "source-end", - "data": { - "type": "end", - "title": "End", - "outputs": [{"variable": "accepted", "value_selector": output_selector}], - }, - }, - ], - "edges": [{"id": "source-edge", "source": "source-start", "target": "source-end"}], - }, - features_dict={}, - environment_variables=[], - workflow_kind=WorkflowKind.STANDARD, - ) - sources = {source.workflow_id: source} - engine = Engine( - graph=Graph.init(graph_config=graph_config, node_factory=node_factory, root_node_id="start"), - runtime_state=runtime_state, - command_channel=InMemoryChannel(), - workers=1, - container_handler_factories=( - partial( - WorkflowToolContainerHandler, - sources=sources, - execution_context_factory=execution_context_layer.enter_context, - ), - ), - ) - engine.add_layer(execution_context_layer) - - if permission == "other-user": - expected_error = ( - "Invalid upload file" if transfer_method == "local_file" else f"ToolFile {stored_file.id} not found" - ) - with pytest.raises(RuntimeError, match=expected_error): - list(engine.run()) - else: - events = list(engine.run()) - assert isinstance(events[-1], GraphRunSucceededEvent) - tool_completed = next( - event for event in events if isinstance(event, NodeRunSucceededEvent) and event.node_id == "tool" - ) - files_output = tool_completed.node_run_result.outputs["files"] - assert isinstance(files_output, ArrayFileSegment) - accepted_files = files_output.value - assert len(accepted_files) == 1 - assert accepted_files[0].filename == stored_file.name diff --git a/api/tests/unit_tests/dev/test_check_knowledge_fs_coverage.py b/api/tests/unit_tests/dev/test_check_knowledge_fs_coverage.py index dd02ed07945..0ac7053c142 100644 --- a/api/tests/unit_tests/dev/test_check_knowledge_fs_coverage.py +++ b/api/tests/unit_tests/dev/test_check_knowledge_fs_coverage.py @@ -52,7 +52,7 @@ def test_core_coverage_scope_includes_every_required_layer(path: str) -> None: "api/app_factory.py", "api/dev/check_knowledge_fs_coverage.py", "api/migrations/versions/revision_knowledge_fs.py", - "api/tests/unit_tests/services/test_knowledge_fs_runtime.py", + "api/tests/unit_tests/services/knowledge_fs/test_runtime.py", ], ) def test_core_coverage_scope_excludes_glue_migrations_and_tests(path: str) -> None: diff --git a/api/tests/unit_tests/commands/test_generate_swagger_markdown_docs.py b/api/tests/unit_tests/dev/test_generate_swagger_markdown_docs.py similarity index 100% rename from api/tests/unit_tests/commands/test_generate_swagger_markdown_docs.py rename to api/tests/unit_tests/dev/test_generate_swagger_markdown_docs.py diff --git a/api/tests/unit_tests/commands/test_generate_swagger_specs.py b/api/tests/unit_tests/dev/test_generate_swagger_specs.py similarity index 99% rename from api/tests/unit_tests/commands/test_generate_swagger_specs.py rename to api/tests/unit_tests/dev/test_generate_swagger_specs.py index 82a88ae23f1..64a1c275c90 100644 --- a/api/tests/unit_tests/commands/test_generate_swagger_specs.py +++ b/api/tests/unit_tests/dev/test_generate_swagger_specs.py @@ -9,7 +9,7 @@ from pathlib import Path from jsonschema import Draft202012Validator -from tests.unit_tests.config_override import apply_config_overrides +from tests.helpers.unit.config_override import apply_config_overrides def _walk_values(value): diff --git a/api/tests/unit_tests/commands/test_lint_response_contracts.py b/api/tests/unit_tests/dev/test_lint_response_contracts.py similarity index 100% rename from api/tests/unit_tests/commands/test_lint_response_contracts.py rename to api/tests/unit_tests/dev/test_lint_response_contracts.py diff --git a/api/tests/unit_tests/test_pytest_sharding.py b/api/tests/unit_tests/dev/test_pytest_sharding.py similarity index 100% rename from api/tests/unit_tests/test_pytest_sharding.py rename to api/tests/unit_tests/dev/test_pytest_sharding.py diff --git a/api/tests/unit_tests/enums/test_quota_type.py b/api/tests/unit_tests/enums/test___init__.py similarity index 100% rename from api/tests/unit_tests/enums/test_quota_type.py rename to api/tests/unit_tests/enums/test___init__.py diff --git a/api/tests/unit_tests/events/event_handlers/test_queue_default_plugin_install_when_tenant_created.py b/api/tests/unit_tests/events/event_handlers/test_queue_default_plugin_install_when_tenant_created.py index 21c64cdd2fb..72857b5f606 100644 --- a/api/tests/unit_tests/events/event_handlers/test_queue_default_plugin_install_when_tenant_created.py +++ b/api/tests/unit_tests/events/event_handlers/test_queue_default_plugin_install_when_tenant_created.py @@ -5,8 +5,8 @@ import pytest from events.event_handlers import queue_default_plugin_install_when_tenant_created as handler_module from models.account import Tenant -from tests.unit_tests.config_override import apply_config_overrides -from tests.unit_tests.model_factories import make_tenant +from tests.helpers.unit.config_override import apply_config_overrides +from tests.helpers.unit.model_factories import make_tenant def _tenant() -> Tenant: diff --git a/api/tests/unit_tests/events/test_update_provider_when_message_created.py b/api/tests/unit_tests/events/event_handlers/test_update_provider_when_message_created.py similarity index 100% rename from api/tests/unit_tests/events/test_update_provider_when_message_created.py rename to api/tests/unit_tests/events/event_handlers/test_update_provider_when_message_created.py diff --git a/api/tests/unit_tests/events/test_events_package_compat.py b/api/tests/unit_tests/events/test___init__.py similarity index 100% rename from api/tests/unit_tests/events/test_events_package_compat.py rename to api/tests/unit_tests/events/test___init__.py diff --git a/api/tests/unit_tests/events/test_app_event_signals.py b/api/tests/unit_tests/events/test_app_event.py similarity index 100% rename from api/tests/unit_tests/events/test_app_event_signals.py rename to api/tests/unit_tests/events/test_app_event.py diff --git a/api/tests/unit_tests/extensions/otel/decorators/handlers/test_generate_handler.py b/api/tests/unit_tests/extensions/otel/decorators/handlers/test_generate_handler.py index c00fbaae36c..e5ccc0aca6b 100644 --- a/api/tests/unit_tests/extensions/otel/decorators/handlers/test_generate_handler.py +++ b/api/tests/unit_tests/extensions/otel/decorators/handlers/test_generate_handler.py @@ -11,7 +11,7 @@ from unittest.mock import MagicMock from core.app.entities.app_invoke_entities import InvokeFrom from extensions.otel.decorators.handlers.generate_handler import AppGenerateHandler from extensions.otel.semconv import DifySpanAttributes, GenAIAttributes -from tests.unit_tests.config_override import config_overrides_context +from tests.helpers.unit.config_override import config_overrides_context class TestAppGenerateHandler: diff --git a/api/tests/unit_tests/extensions/otel/decorators/handlers/test_workflow_app_runner_handler.py b/api/tests/unit_tests/extensions/otel/decorators/handlers/test_workflow_app_runner_handler.py index 0886568e002..54e54fe0dea 100644 --- a/api/tests/unit_tests/extensions/otel/decorators/handlers/test_workflow_app_runner_handler.py +++ b/api/tests/unit_tests/extensions/otel/decorators/handlers/test_workflow_app_runner_handler.py @@ -15,8 +15,8 @@ from sqlalchemy import Engine from extensions.otel.decorators.handlers.workflow_app_runner_handler import WorkflowAppRunnerHandler from extensions.otel.semconv import DifySpanAttributes, GenAIAttributes from services.workflow_run_agg import WorkflowRunAgg -from tests.unit_tests.config_override import config_overrides_context -from tests.unit_tests.services.workflow.test_workflow_run_agg import make_workflow_runner +from tests.helpers.unit.config_override import config_overrides_context +from tests.unit_tests.services.test_workflow_run_agg import make_workflow_runner def test_service_driver_emits_workflow_span(sqlite_engine: Engine) -> None: diff --git a/api/tests/unit_tests/extensions/otel/test_runtime.py b/api/tests/unit_tests/extensions/otel/test_runtime.py index 3c9a402140c..b3197d94678 100644 --- a/api/tests/unit_tests/extensions/otel/test_runtime.py +++ b/api/tests/unit_tests/extensions/otel/test_runtime.py @@ -7,7 +7,7 @@ from opentelemetry.sdk.trace import TracerProvider from core.logging.context import clear_request_context from models import Account -from tests.unit_tests.config_override import config_overrides_context +from tests.helpers.unit.config_override import config_overrides_context def _user() -> Account: diff --git a/api/tests/unit_tests/extensions/storage/test_clickzetta_volume_paths.py b/api/tests/unit_tests/extensions/storage/clickzetta_volume/test_clickzetta_volume_storage.py similarity index 100% rename from api/tests/unit_tests/extensions/storage/test_clickzetta_volume_paths.py rename to api/tests/unit_tests/extensions/storage/clickzetta_volume/test_clickzetta_volume_storage.py diff --git a/api/tests/unit_tests/oss/aliyun_oss/aliyun_oss/test_aliyun_oss.py b/api/tests/unit_tests/extensions/storage/test_aliyun_oss_storage.py similarity index 85% rename from api/tests/unit_tests/oss/aliyun_oss/aliyun_oss/test_aliyun_oss.py rename to api/tests/unit_tests/extensions/storage/test_aliyun_oss_storage.py index 021ad994e7a..e08b8d1e468 100644 --- a/api/tests/unit_tests/oss/aliyun_oss/aliyun_oss/test_aliyun_oss.py +++ b/api/tests/unit_tests/extensions/storage/test_aliyun_oss_storage.py @@ -5,13 +5,13 @@ import pytest from oss2 import Auth from extensions.storage.aliyun_oss_storage import AliyunOssStorage -from tests.unit_tests.oss.__mock.base import ( +from tests.helpers.unit.oss.__mock.base import ( BaseStorageTest, get_example_bucket, get_example_folder, ) -pytest_plugins = ("tests.unit_tests.oss.__mock.aliyun_oss",) +pytest_plugins = ("tests.helpers.unit.oss.__mock.aliyun_oss",) class TestAliyunOss(BaseStorageTest): diff --git a/api/tests/unit_tests/oss/baidu_obs/test_baidu_obs.py b/api/tests/unit_tests/extensions/storage/test_baidu_obs_storage.py similarity index 95% rename from api/tests/unit_tests/oss/baidu_obs/test_baidu_obs.py rename to api/tests/unit_tests/extensions/storage/test_baidu_obs_storage.py index 38b9c71d530..afc8fe24a72 100644 --- a/api/tests/unit_tests/oss/baidu_obs/test_baidu_obs.py +++ b/api/tests/unit_tests/extensions/storage/test_baidu_obs_storage.py @@ -6,12 +6,12 @@ from baidubce.auth.bce_credentials import BceCredentials from baidubce.bce_client_configuration import BceClientConfiguration from extensions.storage.baidu_obs_storage import BaiduObsStorage -from tests.unit_tests.oss.__mock.base import ( +from tests.helpers.unit.oss.__mock.base import ( BaseStorageTest, get_example_bucket, ) -pytest_plugins = ("tests.unit_tests.oss.__mock.baidu_obs",) +pytest_plugins = ("tests.helpers.unit.oss.__mock.baidu_obs",) class TestBaiduObs(BaseStorageTest): diff --git a/api/tests/unit_tests/extensions/storage/test_opendal_fs_default_root_behavior.py b/api/tests/unit_tests/extensions/storage/test_opendal_fs_default_root_behavior.py deleted file mode 100644 index 2782b30c60b..00000000000 --- a/api/tests/unit_tests/extensions/storage/test_opendal_fs_default_root_behavior.py +++ /dev/null @@ -1,60 +0,0 @@ -"""Unit coverage for OpenDAL filesystem root selection.""" - -from pathlib import Path - -import pytest - -from extensions.storage.opendal_storage import OpenDALStorage - - -class TestOpenDALFsDefaultRoot: - """Test that OpenDALStorage with scheme='fs' works correctly when no root is provided.""" - - def test_fs_without_root_uses_default(self, tmp_path: Path, monkeypatch: pytest.MonkeyPatch) -> None: - """When no root is specified, the default 'storage' should be used and passed to the Operator.""" - # Change to tmp_path so the default "storage" dir is created there - monkeypatch.chdir(tmp_path) - # Ensure no OPENDAL_FS_ROOT env var is set - monkeypatch.delenv("OPENDAL_FS_ROOT", raising=False) - - storage = OpenDALStorage(scheme="fs") - - # The default directory should have been created - assert (tmp_path / "storage").is_dir() - # The storage should be functional - storage.save("test_default_root.txt", b"hello") - assert storage.exists("test_default_root.txt") - assert storage.load_once("test_default_root.txt") == b"hello" - - # Cleanup - storage.delete("test_default_root.txt") - - def test_fs_with_explicit_root(self, tmp_path: Path) -> None: - """When root is explicitly provided, it should be used.""" - custom_root = str(tmp_path / "custom_storage") - storage = OpenDALStorage(scheme="fs", root=custom_root) - - assert Path(custom_root).is_dir() - storage.save("test_explicit_root.txt", b"world") - assert storage.exists("test_explicit_root.txt") - assert storage.load_once("test_explicit_root.txt") == b"world" - - # Cleanup - storage.delete("test_explicit_root.txt") - - def test_fs_with_env_var_root(self, tmp_path: Path, monkeypatch: pytest.MonkeyPatch) -> None: - """When OPENDAL_FS_ROOT env var is set, it should be picked up via _get_opendal_kwargs.""" - env_root = str(tmp_path / "env_storage") - monkeypatch.setenv("OPENDAL_FS_ROOT", env_root) - # Ensure .env file doesn't interfere - monkeypatch.chdir(tmp_path) - - storage = OpenDALStorage(scheme="fs") - - assert Path(env_root).is_dir() - storage.save("test_env_root.txt", b"env_data") - assert storage.exists("test_env_root.txt") - assert storage.load_once("test_env_root.txt") == b"env_data" - - # Cleanup - storage.delete("test_env_root.txt") diff --git a/api/tests/unit_tests/oss/opendal/test_opendal.py b/api/tests/unit_tests/extensions/storage/test_opendal_storage.py similarity index 53% rename from api/tests/unit_tests/oss/opendal/test_opendal.py rename to api/tests/unit_tests/extensions/storage/test_opendal_storage.py index 11bbf5f1933..f1a78299fdc 100644 --- a/api/tests/unit_tests/oss/opendal/test_opendal.py +++ b/api/tests/unit_tests/extensions/storage/test_opendal_storage.py @@ -4,12 +4,67 @@ from pathlib import Path import pytest from extensions.storage.opendal_storage import OpenDALStorage -from tests.unit_tests.oss.__mock.base import ( +from tests.helpers.unit.oss.__mock.base import ( get_example_data, get_example_filename, get_opendal_bucket, ) +"""Unit coverage for OpenDAL filesystem root selection.""" + + +class TestOpenDALFsDefaultRoot: + """Test that OpenDALStorage with scheme='fs' works correctly when no root is provided.""" + + def test_fs_without_root_uses_default(self, tmp_path: Path, monkeypatch: pytest.MonkeyPatch) -> None: + """When no root is specified, the default 'storage' should be used and passed to the Operator.""" + # Change to tmp_path so the default "storage" dir is created there + monkeypatch.chdir(tmp_path) + # Ensure no OPENDAL_FS_ROOT env var is set + monkeypatch.delenv("OPENDAL_FS_ROOT", raising=False) + + storage = OpenDALStorage(scheme="fs") + + # The default directory should have been created + assert (tmp_path / "storage").is_dir() + # The storage should be functional + storage.save("test_default_root.txt", b"hello") + assert storage.exists("test_default_root.txt") + assert storage.load_once("test_default_root.txt") == b"hello" + + # Cleanup + storage.delete("test_default_root.txt") + + def test_fs_with_explicit_root(self, tmp_path: Path) -> None: + """When root is explicitly provided, it should be used.""" + custom_root = str(tmp_path / "custom_storage") + storage = OpenDALStorage(scheme="fs", root=custom_root) + + assert Path(custom_root).is_dir() + storage.save("test_explicit_root.txt", b"world") + assert storage.exists("test_explicit_root.txt") + assert storage.load_once("test_explicit_root.txt") == b"world" + + # Cleanup + storage.delete("test_explicit_root.txt") + + def test_fs_with_env_var_root(self, tmp_path: Path, monkeypatch: pytest.MonkeyPatch) -> None: + """When OPENDAL_FS_ROOT env var is set, it should be picked up via _get_opendal_kwargs.""" + env_root = str(tmp_path / "env_storage") + monkeypatch.setenv("OPENDAL_FS_ROOT", env_root) + # Ensure .env file doesn't interfere + monkeypatch.chdir(tmp_path) + + storage = OpenDALStorage(scheme="fs") + + assert Path(env_root).is_dir() + storage.save("test_env_root.txt", b"env_data") + assert storage.exists("test_env_root.txt") + assert storage.load_once("test_env_root.txt") == b"env_data" + + # Cleanup + storage.delete("test_env_root.txt") + class TestOpenDAL: @pytest.fixture(autouse=True) diff --git a/api/tests/unit_tests/oss/tencent_cos/test_tencent_cos.py b/api/tests/unit_tests/extensions/storage/test_tencent_cos_storage.py similarity index 97% rename from api/tests/unit_tests/oss/tencent_cos/test_tencent_cos.py rename to api/tests/unit_tests/extensions/storage/test_tencent_cos_storage.py index 96fb8c8abc8..e986dfbe59a 100644 --- a/api/tests/unit_tests/oss/tencent_cos/test_tencent_cos.py +++ b/api/tests/unit_tests/extensions/storage/test_tencent_cos_storage.py @@ -6,14 +6,14 @@ import pytest from qcloud_cos import CosConfig from extensions.storage.tencent_cos_storage import TencentCosStorage -from tests.unit_tests.oss.__mock.base import ( +from tests.helpers.unit.oss.__mock.base import ( BaseStorageTest, get_example_bucket, get_example_data, get_example_filename, ) -pytest_plugins = ("tests.unit_tests.oss.__mock.tencent_cos",) +pytest_plugins = ("tests.helpers.unit.oss.__mock.tencent_cos",) class TestTencentCos(BaseStorageTest): diff --git a/api/tests/unit_tests/oss/volcengine_tos/test_volcengine_tos.py b/api/tests/unit_tests/extensions/storage/test_volcengine_tos_storage.py similarity index 89% rename from api/tests/unit_tests/oss/volcengine_tos/test_volcengine_tos.py rename to api/tests/unit_tests/extensions/storage/test_volcengine_tos_storage.py index 8d7733e1039..31a967b3bd5 100644 --- a/api/tests/unit_tests/oss/volcengine_tos/test_volcengine_tos.py +++ b/api/tests/unit_tests/extensions/storage/test_volcengine_tos_storage.py @@ -5,12 +5,12 @@ import pytest from tos import TosClientV2 from extensions.storage.volcengine_tos_storage import VolcengineTosStorage -from tests.unit_tests.oss.__mock.base import ( +from tests.helpers.unit.oss.__mock.base import ( BaseStorageTest, get_example_bucket, ) -pytest_plugins = ("tests.unit_tests.oss.__mock.volcengine_tos",) +pytest_plugins = ("tests.helpers.unit.oss.__mock.volcengine_tos",) class TestVolcengineTos(BaseStorageTest): diff --git a/api/tests/unit_tests/extensions/test_celery_ssl.py b/api/tests/unit_tests/extensions/test_celery_ssl.py deleted file mode 100644 index 14b217603ca..00000000000 --- a/api/tests/unit_tests/extensions/test_celery_ssl.py +++ /dev/null @@ -1,288 +0,0 @@ -"""Tests for Celery SSL configuration.""" - -import ssl -from collections.abc import Iterator -from importlib import import_module -from unittest.mock import MagicMock, patch - -import pytest -from celery.contrib.testing.app import TestApp, setup_default_app - -from enums import DeploymentEdition -from services.knowledge_fs.lifecycle_readiness import KnowledgeFSLifecycleWorkerReadiness - - -@pytest.fixture(autouse=True) -def restore_celery_app_state() -> Iterator[None]: - # init_app replaces Celery's process-wide current/default app. Restore both - # so later shared tasks do not enter this module's Flask app context. - with setup_default_app(TestApp()): - yield - - -class TestCelerySSLConfiguration: - """Test suite for Celery SSL configuration.""" - - def test_get_celery_broker_transport_options_includes_global_keyprefix_for_redis(self): - mock_config = MagicMock() - mock_config.CELERY_USE_SENTINEL = False - mock_config.REDIS_KEY_PREFIX = "enterprise-a" - - with patch("extensions.ext_celery.dify_config", mock_config): - from extensions.ext_celery import get_celery_broker_transport_options - - result = get_celery_broker_transport_options() - - assert result["global_keyprefix"] == "enterprise-a:" - - def test_get_celery_broker_transport_options_omits_global_keyprefix_when_prefix_empty(self): - mock_config = MagicMock() - mock_config.CELERY_USE_SENTINEL = False - mock_config.REDIS_KEY_PREFIX = " " - - with patch("extensions.ext_celery.dify_config", mock_config): - from extensions.ext_celery import get_celery_broker_transport_options - - result = get_celery_broker_transport_options() - - assert "global_keyprefix" not in result - - def test_get_celery_broker_transport_options_keeps_sentinel_and_adds_global_keyprefix(self): - mock_config = MagicMock() - mock_config.CELERY_USE_SENTINEL = True - mock_config.CELERY_SENTINEL_MASTER_NAME = "mymaster" - mock_config.CELERY_SENTINEL_SOCKET_TIMEOUT = 0.1 - mock_config.CELERY_SENTINEL_PASSWORD = "secret" - mock_config.REDIS_KEY_PREFIX = "enterprise-a" - - with patch("extensions.ext_celery.dify_config", mock_config): - from extensions.ext_celery import get_celery_broker_transport_options - - result = get_celery_broker_transport_options() - - assert result["master_name"] == "mymaster" - assert result["sentinel_kwargs"]["password"] == "secret" - assert result["global_keyprefix"] == "enterprise-a:" - - def test_get_celery_ssl_options_when_ssl_disabled(self): - """Test SSL options when BROKER_USE_SSL is False.""" - from configs import DifyConfig - - dify_config = DifyConfig(CELERY_BROKER_URL="redis://localhost:6379/0") - - with patch("extensions.ext_celery.dify_config", dify_config): - from extensions.ext_celery import get_celery_ssl_options - - result = get_celery_ssl_options() - assert result is None - - def test_get_celery_ssl_options_when_broker_not_redis(self): - """Test SSL options when broker is not Redis.""" - mock_config = MagicMock() - mock_config.CELERY_BROKER_URL = "amqp://localhost:5672" - - with patch("extensions.ext_celery.dify_config", mock_config): - from extensions.ext_celery import get_celery_ssl_options - - result = get_celery_ssl_options() - assert result is None - - def test_get_celery_ssl_options_with_cert_none(self): - """Test SSL options with CERT_NONE requirement.""" - mock_config = MagicMock() - mock_config.CELERY_BROKER_URL = "redis://localhost:6379/0" - mock_config.REDIS_SSL_CERT_REQS = "CERT_NONE" - mock_config.REDIS_SSL_CA_CERTS = None - mock_config.REDIS_SSL_CERTFILE = None - mock_config.REDIS_SSL_KEYFILE = None - - with patch("extensions.ext_celery.dify_config", mock_config): - from extensions.ext_celery import get_celery_ssl_options - - result = get_celery_ssl_options() - assert result is not None - assert result["ssl_cert_reqs"] == ssl.CERT_NONE - assert result["ssl_ca_certs"] is None - assert result["ssl_certfile"] is None - assert result["ssl_keyfile"] is None - - def test_get_celery_ssl_options_with_cert_required(self): - """Test SSL options with CERT_REQUIRED and certificates.""" - mock_config = MagicMock() - mock_config.CELERY_BROKER_URL = "rediss://localhost:6380/0" - mock_config.REDIS_SSL_CERT_REQS = "CERT_REQUIRED" - mock_config.REDIS_SSL_CA_CERTS = "/path/to/ca.crt" - mock_config.REDIS_SSL_CERTFILE = "/path/to/client.crt" - mock_config.REDIS_SSL_KEYFILE = "/path/to/client.key" - - with patch("extensions.ext_celery.dify_config", mock_config): - from extensions.ext_celery import get_celery_ssl_options - - result = get_celery_ssl_options() - assert result is not None - assert result["ssl_cert_reqs"] == ssl.CERT_REQUIRED - assert result["ssl_ca_certs"] == "/path/to/ca.crt" - assert result["ssl_certfile"] == "/path/to/client.crt" - assert result["ssl_keyfile"] == "/path/to/client.key" - - def test_get_celery_ssl_options_with_cert_optional(self): - """Test SSL options with CERT_OPTIONAL requirement.""" - mock_config = MagicMock() - mock_config.CELERY_BROKER_URL = "redis://localhost:6379/0" - mock_config.REDIS_SSL_CERT_REQS = "CERT_OPTIONAL" - mock_config.REDIS_SSL_CA_CERTS = "/path/to/ca.crt" - mock_config.REDIS_SSL_CERTFILE = None - mock_config.REDIS_SSL_KEYFILE = None - - with patch("extensions.ext_celery.dify_config", mock_config): - from extensions.ext_celery import get_celery_ssl_options - - result = get_celery_ssl_options() - assert result is not None - assert result["ssl_cert_reqs"] == ssl.CERT_OPTIONAL - assert result["ssl_ca_certs"] == "/path/to/ca.crt" - - def test_get_celery_ssl_options_with_invalid_cert_reqs(self): - """Test SSL options with invalid cert requirement defaults to CERT_NONE.""" - mock_config = MagicMock() - mock_config.CELERY_BROKER_URL = "redis://localhost:6379/0" - mock_config.REDIS_SSL_CERT_REQS = "INVALID_VALUE" - mock_config.REDIS_SSL_CA_CERTS = None - mock_config.REDIS_SSL_CERTFILE = None - mock_config.REDIS_SSL_KEYFILE = None - - with patch("extensions.ext_celery.dify_config", mock_config): - from extensions.ext_celery import get_celery_ssl_options - - result = get_celery_ssl_options() - assert result is not None - assert result["ssl_cert_reqs"] == ssl.CERT_NONE # Should default to CERT_NONE - - def test_celery_init_applies_ssl_to_broker_and_backend(self): - """Test that SSL options are applied to both broker and backend when using Redis.""" - mock_config = MagicMock() - mock_config.HUMAN_INPUT_TIMEOUT_TASK_INTERVAL = 1 - mock_config.CELERY_BROKER_URL = "redis://localhost:6379/0" - mock_config.CELERY_BACKEND = "redis" - mock_config.CELERY_RESULT_BACKEND = "redis://localhost:6379/0" - mock_config.REDIS_SSL_CERT_REQS = "CERT_NONE" - mock_config.REDIS_SSL_CA_CERTS = None - mock_config.REDIS_SSL_CERTFILE = None - mock_config.REDIS_SSL_KEYFILE = None - mock_config.CELERY_USE_SENTINEL = False - mock_config.LOG_FORMAT = "%(message)s" - mock_config.LOG_TZ = "UTC" - mock_config.LOG_FILE = None - - # Mock all the scheduler configs - mock_config.CELERY_BEAT_SCHEDULER_TIME = 1 - mock_config.AGENT_SANDBOX_METERING_ENABLED = False - mock_config.ENABLE_CONVERSATION_CLEANUP_TASK = False - mock_config.CONVERSATION_CLEANUP_TASK_INTERVAL = 5 - mock_config.ENABLE_CLEAN_EMBEDDING_CACHE_TASK = False - mock_config.ENABLE_CLEAN_UNUSED_DATASETS_TASK = False - mock_config.ENABLE_CREATE_TIDB_SERVERLESS_TASK = False - mock_config.ENABLE_UPDATE_TIDB_SERVERLESS_STATUS_TASK = False - mock_config.ENABLE_CLEAN_MESSAGES = False - mock_config.ENABLE_MAIL_CLEAN_DOCUMENT_NOTIFY_TASK = False - mock_config.ENABLE_DATASETS_QUEUE_MONITOR = False - mock_config.ENABLE_CHECK_UPGRADABLE_PLUGIN_TASK = False - mock_config.ENABLE_WORKFLOW_SCHEDULE_POLLER_TASK = False - mock_config.WORKFLOW_SCHEDULE_POLLER_INTERVAL = 1 - mock_config.WORKFLOW_SCHEDULE_POLLER_BATCH_SIZE = 100 - mock_config.WORKFLOW_SCHEDULE_MAX_DISPATCH_PER_TICK = 0 - mock_config.ENABLE_TRIGGER_PROVIDER_REFRESH_TASK = False - mock_config.TRIGGER_PROVIDER_REFRESH_INTERVAL = 15 - mock_config.ENABLE_API_TOKEN_LAST_USED_UPDATE_TASK = False - mock_config.API_TOKEN_LAST_USED_UPDATE_INTERVAL = 30 - - with ( - patch("extensions.ext_celery.dify_config", mock_config), - patch( - "services.knowledge_fs.lifecycle_readiness.get_configured_knowledge_fs_lifecycle_worker_readiness", - return_value=KnowledgeFSLifecycleWorkerReadiness(ready=False, blockers=()), - ), - ): - from dify_app import DifyApp - from extensions.ext_celery import init_app - - app = DifyApp(__name__) - celery_app = init_app(app) - - # Check that SSL options were applied - assert "broker_use_ssl" in celery_app.conf - assert celery_app.conf["broker_use_ssl"] is not None - assert celery_app.conf["broker_use_ssl"]["ssl_cert_reqs"] == ssl.CERT_NONE - - # Check that SSL is also applied to Redis backend - assert "redis_backend_use_ssl" in celery_app.conf - assert celery_app.conf["redis_backend_use_ssl"] is not None - - def test_celery_init_registers_required_agent_and_conversation_tasks(self): - mock_config = MagicMock() - mock_config.BROKER_USE_SSL = False - mock_config.REDIS_KEY_PREFIX = "enterprise-a" - mock_config.HUMAN_INPUT_TIMEOUT_TASK_INTERVAL = 1 - mock_config.CELERY_BROKER_URL = "redis://localhost:6379/0" - mock_config.CELERY_BACKEND = "redis" - mock_config.CELERY_RESULT_BACKEND = "redis://localhost:6379/0" - mock_config.CELERY_USE_SENTINEL = False - mock_config.LOG_FORMAT = "%(message)s" - mock_config.LOG_TZ = "UTC" - mock_config.LOG_FILE = None - mock_config.CELERY_TASK_ANNOTATIONS = {} - - mock_config.CELERY_BEAT_SCHEDULER_TIME = 1 - mock_config.AGENT_SANDBOX_METERING_ENABLED = False - mock_config.ENABLE_CONVERSATION_CLEANUP_TASK = True - mock_config.CONVERSATION_CLEANUP_TASK_INTERVAL = 5 - mock_config.ENABLE_CLEAN_EMBEDDING_CACHE_TASK = False - mock_config.ENABLE_CLEAN_UNUSED_DATASETS_TASK = False - mock_config.ENABLE_CREATE_TIDB_SERVERLESS_TASK = False - mock_config.ENABLE_UPDATE_TIDB_SERVERLESS_STATUS_TASK = False - mock_config.ENABLE_CLEAN_MESSAGES = False - mock_config.ENABLE_MAIL_CLEAN_DOCUMENT_NOTIFY_TASK = False - mock_config.ENABLE_DATASETS_QUEUE_MONITOR = False - mock_config.ENABLE_HUMAN_INPUT_TIMEOUT_TASK = False - mock_config.ENABLE_CHECK_UPGRADABLE_PLUGIN_TASK = False - mock_config.MARKETPLACE_ENABLED = False - mock_config.WORKFLOW_LOG_CLEANUP_ENABLED = False - mock_config.ENABLE_WORKFLOW_RUN_CLEANUP_TASK = False - mock_config.ENABLE_WORKFLOW_SCHEDULE_POLLER_TASK = False - mock_config.WORKFLOW_SCHEDULE_POLLER_INTERVAL = 1 - mock_config.ENABLE_TRIGGER_PROVIDER_REFRESH_TASK = False - mock_config.TRIGGER_PROVIDER_REFRESH_INTERVAL = 15 - mock_config.ENABLE_API_TOKEN_LAST_USED_UPDATE_TASK = False - mock_config.API_TOKEN_LAST_USED_UPDATE_INTERVAL = 30 - mock_config.DEPLOYMENT_EDITION = DeploymentEdition.COMMUNITY - mock_config.ENTERPRISE_TELEMETRY_ENABLED = False - - with ( - patch("extensions.ext_celery.dify_config", mock_config), - patch( - "services.knowledge_fs.lifecycle_readiness.get_configured_knowledge_fs_lifecycle_worker_readiness", - return_value=KnowledgeFSLifecycleWorkerReadiness(ready=False, blockers=()), - ), - ): - from dify_app import DifyApp - from extensions.ext_celery import init_app - - app = DifyApp(__name__) - celery_app = init_app(app) - - assert celery_app.conf["broker_transport_options"]["global_keyprefix"] == "enterprise-a:" - assert celery_app.conf["result_backend_transport_options"]["global_keyprefix"] == "enterprise-a:" - assert "tasks.collect_agent_resources_task" in celery_app.conf["imports"] - assert "tasks.delete_conversation_task" in celery_app.conf["imports"] - assert "tasks.workflow_run_archive_deletion_tasks" in celery_app.conf["imports"] - import_module("tasks.workflow_run_archive_deletion_tasks") - task_name = "tasks.workflow_run_archive_deletion_tasks.delete_workflow_run_archive_task" - archive_deletion = celery_app.tasks[task_name] - assert archive_deletion.queue == "workflow_archive" - assert archive_deletion.acks_late is True - assert archive_deletion.reject_on_worker_lost is True - assert archive_deletion.autoretry_for == (Exception,) - assert archive_deletion.retry_kwargs["max_retries"] == 5 - assert celery_app.conf["beat_schedule"]["conversation_cleanup_sweeper"]["task"] == ( - "tasks.delete_conversation_task.sweep_deleted_conversations" - ) diff --git a/api/tests/unit_tests/extensions/test_community_telemetry_celery.py b/api/tests/unit_tests/extensions/test_community_telemetry_celery.py deleted file mode 100644 index 58ec7641457..00000000000 --- a/api/tests/unit_tests/extensions/test_community_telemetry_celery.py +++ /dev/null @@ -1,43 +0,0 @@ -from types import SimpleNamespace -from unittest.mock import Mock - -from extensions.ext_celery import _enqueue_initial_community_telemetry_heartbeat - - -def test_beat_start_enqueues_community_telemetry_heartbeat() -> None: - task = Mock() - sender = SimpleNamespace( - app=SimpleNamespace( - conf=SimpleNamespace(beat_schedule={"community_telemetry_heartbeat": {}}), - tasks={"community_telemetry.send_heartbeat": task}, - ) - ) - - _enqueue_initial_community_telemetry_heartbeat(sender) - - task.apply_async.assert_called_once_with() - - -def test_beat_start_skips_community_telemetry_when_not_scheduled() -> None: - task = Mock() - sender = SimpleNamespace( - app=SimpleNamespace( - conf=SimpleNamespace(beat_schedule={}), - tasks={"community_telemetry.send_heartbeat": task}, - ) - ) - - _enqueue_initial_community_telemetry_heartbeat(sender) - - task.apply_async.assert_not_called() - - -def test_beat_start_skips_community_telemetry_when_task_is_unavailable() -> None: - sender = SimpleNamespace( - app=SimpleNamespace( - conf=SimpleNamespace(beat_schedule={"community_telemetry_heartbeat": {}}), - tasks={}, - ) - ) - - _enqueue_initial_community_telemetry_heartbeat(sender) diff --git a/api/tests/unit_tests/extensions/test_ext_application_services.py b/api/tests/unit_tests/extensions/test_ext_application_services.py index 1f8210890bf..5176d2fce94 100644 --- a/api/tests/unit_tests/extensions/test_ext_application_services.py +++ b/api/tests/unit_tests/extensions/test_ext_application_services.py @@ -167,8 +167,8 @@ from services.workflow_app_log_query_service import WorkflowAppLogQueryService from services.workflow_run_service import WorkflowRunService from services.workflow_statistic_query_service import WorkflowStatisticQueryService from tasks.delete_conversation_task import delete_conversation_related_data -from tests.unit_tests.config_override import apply_config_overrides -from tests.unit_tests.services.test_app_task_service import _StopRedis +from tests.helpers.unit.config_override import apply_config_overrides +from tests.unit_tests.services.test_app_task_service import _AppTaskServiceStopredis as _StopRedis @pytest.fixture(autouse=True) diff --git a/api/tests/unit_tests/extensions/test_ext_blueprints.py b/api/tests/unit_tests/extensions/test_ext_blueprints.py new file mode 100644 index 00000000000..6694dd726ee --- /dev/null +++ b/api/tests/unit_tests/extensions/test_ext_blueprints.py @@ -0,0 +1,391 @@ +from __future__ import annotations + +import builtins +import importlib + +import pytest +from flask import Blueprint, Flask, Response +from flask.views import MethodView +from flask_cors import CORS +from flask_restx import Resource + +from configs import dify_config +from constants import HEADER_NAME_CATALOG +from dify_app import DifyApp +from extensions import ext_blueprints +from extensions.ext_blueprints import ( + AUTHENTICATED_HEADERS, + EXPOSED_HEADERS, + OPENAPI_HEADERS, + OPENAPI_MAX_AGE_SECONDS, + _apply_cors_once, +) +from libs.external_api import ExternalApi +from tests.helpers.unit.config_override import apply_config_overrides + +"""CORS posture for /openapi/v1/* — default empty allowlist (same-origin), +expandable via OPENAPI_CORS_ALLOW_ORIGINS. Cross-origin requests from +disallowed origins do not receive the Access-Control-Allow-Origin +header, which the browser then blocks. + +Tests use a fresh Blueprint + Flask-CORS per case because the production +blueprint is a module-level singleton and can't be reconfigured once +registered. +""" + + +if not hasattr(builtins, "MethodView"): + builtins.MethodView = MethodView # type: ignore[attr-defined] + + +def _make_app(allowed_origins: list[str], blueprint_name: str) -> Flask: + """Build a Flask app with a fresh openapi-style blueprint mirroring + production CORS settings, parameterised on the origin allowlist. + """ + bp = Blueprint(blueprint_name, __name__, url_prefix="/openapi/v1") + api = ExternalApi(bp, version="1.0", title="OpenAPI Test", description="") + + @api.route("/_health") + class _Health(Resource): + def get(self) -> dict[str, bool]: + return {"ok": True} + + CORS( + bp, + resources={r"/*": {"origins": allowed_origins}}, + supports_credentials=True, + allow_headers=list(OPENAPI_HEADERS), + methods=["GET", "POST", "PATCH", "DELETE", "OPTIONS"], + expose_headers=list(EXPOSED_HEADERS), + max_age=OPENAPI_MAX_AGE_SECONDS, + ) + + app = Flask(__name__) + app.config["TESTING"] = True + app.register_blueprint(bp) + return app + + +class TestCorsCases: + from tests.unit_tests.controllers.conftest import error_response_contract + + error_response_contract = staticmethod(error_response_contract) + + from tests.unit_tests.controllers.conftest import datasource_application_dependencies + + datasource_application_dependencies = staticmethod(datasource_application_dependencies) + + from tests.unit_tests.controllers.conftest import pipeline_application + + pipeline_application = staticmethod(pipeline_application) + + from tests.unit_tests.controllers.conftest import app_query_services + + app_query_services = staticmethod(app_query_services) + + from tests.unit_tests.controllers.openapi.conftest import openapi_app + + openapi_app = staticmethod(openapi_app) + + from tests.unit_tests.controllers.openapi.conftest import app + + app = staticmethod(app) + + from tests.unit_tests.controllers.openapi.conftest import admitted_bearer + + admitted_bearer = staticmethod(admitted_bearer) + + from tests.unit_tests.controllers.openapi.conftest import _account_services + + _account_services = staticmethod(_account_services) + + @staticmethod + def test_default_openapi_cors_allowlist_is_empty() -> None: + """Default config admits no cross-origin until operator opts in.""" + assert dify_config.OPENAPI_CORS_ALLOW_ORIGINS == [] + + @staticmethod + def test_preflight_allowed_origin_returns_cors_headers() -> None: + app = _make_app(["https://app.example.com"], "openapi_t1") + client = app.test_client() + response = client.options( + "/openapi/v1/_health", + headers={ + "Origin": "https://app.example.com", + "Access-Control-Request-Method": "GET", + }, + ) + + assert response.headers.get("Access-Control-Allow-Origin") == "https://app.example.com" + assert response.headers.get("Access-Control-Max-Age") == str(OPENAPI_MAX_AGE_SECONDS) + + @staticmethod + def test_preflight_disallowed_origin_omits_cors_headers() -> None: + app = _make_app(["https://app.example.com"], "openapi_t2") + client = app.test_client() + response = client.options( + "/openapi/v1/_health", + headers={ + "Origin": "https://attacker.example", + "Access-Control-Request-Method": "GET", + }, + ) + + # flask-cors omits Allow-Origin for disallowed origins; browser blocks. + assert "Access-Control-Allow-Origin" not in response.headers + + @staticmethod + def test_preflight_with_default_empty_allowlist_omits_cors_headers() -> None: + app = _make_app([], "openapi_t3") + client = app.test_client() + response = client.options( + "/openapi/v1/_health", + headers={ + "Origin": "https://app.example.com", + "Access-Control-Request-Method": "GET", + }, + ) + + assert "Access-Control-Allow-Origin" not in response.headers + + @staticmethod + def test_same_origin_request_succeeds_without_origin_header() -> None: + app = _make_app(["https://app.example.com"], "openapi_t4") + client = app.test_client() + # Browsers don't send Origin on same-origin GETs. + response = client.get("/openapi/v1/_health") + + assert response.status_code == 200 + assert response.get_json() == {"ok": True} + + @staticmethod + def test_authorization_header_is_in_allow_headers() -> None: + """Bearer-authed routes need Authorization in the preflight response.""" + app = _make_app(["https://app.example.com"], "openapi_t5") + client = app.test_client() + response = client.options( + "/openapi/v1/_health", + headers={ + "Origin": "https://app.example.com", + "Access-Control-Request-Method": "GET", + "Access-Control-Request-Headers": "Authorization", + }, + ) + + allow_headers = response.headers.get("Access-Control-Allow-Headers", "").lower() + assert "authorization" in allow_headers + + @staticmethod + def test_catalog_fingerprint_header_is_exposed() -> None: + app = _make_app(["https://app.example.com"], "openapi_t6") + client = app.test_client() + response = client.get("/openapi/v1/_health", headers={"Origin": "https://app.example.com"}) + + expose_headers = response.headers.get("Access-Control-Expose-Headers", "") + assert HEADER_NAME_CATALOG in expose_headers + + @staticmethod + def test_catalog_fingerprint_header_is_allowed_on_requests() -> None: + """Every guarded route needs the fingerprint on the request, so preflight must admit it.""" + app = _make_app(["https://app.example.com"], "openapi_t7") + client = app.test_client() + response = client.options( + "/openapi/v1/_health", + headers={ + "Origin": "https://app.example.com", + "Access-Control-Request-Method": "GET", + "Access-Control-Request-Headers": HEADER_NAME_CATALOG, + }, + ) + + allow_headers = response.headers.get("Access-Control-Allow-Headers", "").lower() + assert HEADER_NAME_CATALOG.lower() in allow_headers + + +"""Regression coverage for authenticated browser CORS headers.""" + + +class TestExtBlueprintsCorsCases: + @staticmethod + def test_authenticated_cors_allows_request_metadata_headers() -> None: + app = Flask(__name__) + blueprint = Blueprint("cors_probe", __name__, url_prefix="/console/api") + + @blueprint.post("/probe") + def probe() -> tuple[str, int]: + return "", 204 + + _apply_cors_once( + blueprint, + resources={r"/*": {"origins": ["http://localhost:3000"]}}, + supports_credentials=True, + allow_headers=list(AUTHENTICATED_HEADERS), + methods=["POST", "OPTIONS"], + ) + app.register_blueprint(blueprint) + + response = app.test_client().options( + "/console/api/probe", + headers={ + "Access-Control-Request-Headers": "Idempotency-Key, X-Request-ID, X-Trace-ID", + "Access-Control-Request-Method": "POST", + "Origin": "http://localhost:3000", + }, + ) + + allowed_headers = response.headers.get("Access-Control-Allow-Headers", "").lower() + assert "idempotency-key" in allowed_headers + assert "x-request-id" in allowed_headers + assert "x-trace-id" in allowed_headers + + +"""Verifies the OPENAPI_ENABLED gate in `extensions.ext_blueprints.init_app`. + +Contract: +- When `dify_config.OPENAPI_ENABLED` is True, the `/openapi/v1/*` blueprint + must be registered on the Flask app AND have CORS configured (signalled + by the idempotent `_dify_cors_applied` marker that `_apply_cors_once` + sets after wiring `flask_cors.CORS`). +- When False, the openapi blueprint must NOT be registered and CORS must + NOT be wired for it. The cross-origin posture is otherwise undefined for + the disabled feature. + +Why the blueprints are swapped for fresh ones: +`init_app` operates on module-level blueprint singletons (e.g. +`controllers.service_api.bp`). Other unit tests register those same +singletons onto throwaway Flask apps without going through `init_app` +(see `tests/unit_tests/controllers/service_api/test___init__.py` and the +`tests/unit_tests/controllers/openapi/*` suite). Once a blueprint has +been registered to any app, Flask flips `_got_registered_once = True` +and rejects further `after_request` hookup -- which is what +`flask_cors.CORS(bp, ...)` does internally. To make this gate test +order-independent we monkeypatch each consumed `controllers..bp` to +a pristine, never-registered `Blueprint` for the duration of the test. +`init_app` resolves these via `from controllers. import bp as ...` +inside its function body, so the patched value is what it sees. +""" + + +# Modules whose `bp` attribute is consumed by `ext_blueprints.init_app`. +# Keep in sync with the imports inside `init_app`. +_BLUEPRINT_MODULES: tuple[str, ...] = ( + "controllers.console", + "controllers.files", + "controllers.inner_api", + "controllers.mcp", + "controllers.openapi", + "controllers.service_api", + "controllers.trigger", + "controllers.web", +) + + +def _probe_view() -> tuple[str, int]: + return "ok", 200 + + +def _download_probe_view() -> Response: + response = Response("document", content_type="text/plain") + response.headers["Content-Disposition"] = "attachment; filename=document.txt" + return response + + +def _build_app() -> DifyApp: + app = DifyApp(__name__) + app.config["TESTING"] = True + return app + + +class TestExtBlueprintsOpenapiCases: + @staticmethod + @pytest.fixture + def fresh_blueprints(monkeypatch: pytest.MonkeyPatch) -> dict[str, Blueprint]: + """Replace each production blueprint singleton with a fresh, unregistered copy. + + Mirrors the production `name` and `url_prefix` so the gate can still + be asserted via `app.blueprints[...]` and url_map prefix checks. The + openapi replacement gets a dummy `/_probe` rule so the test can + observe at least one `/openapi/v1/*` rule after registration. + """ + fresh: dict[str, Blueprint] = {} + for module_name in _BLUEPRINT_MODULES: + module = importlib.import_module(module_name) + original = module.bp + replacement = Blueprint( + original.name, + original.import_name, + url_prefix=original.url_prefix, + ) + if module_name == "controllers.openapi": + replacement.add_url_rule("/_probe", endpoint="_probe", view_func=_probe_view) + if module_name == "controllers.console": + replacement.add_url_rule("/_download", endpoint="_download", view_func=_download_probe_view) + monkeypatch.setattr(module, "bp", replacement) + fresh[module_name] = replacement + return fresh + + @staticmethod + def test_openapi_blueprint_registered_with_cors_when_enabled( + monkeypatch: pytest.MonkeyPatch, + fresh_blueprints: dict[str, Blueprint], + ) -> None: + """Enabled gate: blueprint mounted, CORS wired, `/openapi/v1/*` rules live.""" + apply_config_overrides(monkeypatch, OPENAPI_ENABLED=True) + + app = _build_app() + ext_blueprints.init_app(app) + + openapi_bp = fresh_blueprints["controllers.openapi"] + assert "openapi" in app.blueprints + assert app.blueprints["openapi"] is openapi_bp + # `_apply_cors_once` only sets this after wiring flask_cors.CORS, so + # it is a faithful proxy for "CORS was applied to this blueprint". + assert getattr(openapi_bp, "_dify_cors_applied", False) is True + + openapi_rules = [r for r in app.url_map.iter_rules() if r.rule.startswith("/openapi/v1")] + assert openapi_rules, "expected at least one /openapi/v1/* route once enabled" + + @staticmethod + def test_openapi_blueprint_absent_when_disabled( + monkeypatch: pytest.MonkeyPatch, + fresh_blueprints: dict[str, Blueprint], + ) -> None: + """Disabled gate: no blueprint, no CORS, no `/openapi/v1/*` URL rules.""" + apply_config_overrides(monkeypatch, OPENAPI_ENABLED=False) + + app = _build_app() + ext_blueprints.init_app(app) + + openapi_bp = fresh_blueprints["controllers.openapi"] + assert "openapi" not in app.blueprints + # No CORS wiring should have run for the openapi blueprint. + assert not hasattr(openapi_bp, "_dify_cors_applied") + openapi_rules = [r for r in app.url_map.iter_rules() if r.rule.startswith("/openapi/v1")] + assert openapi_rules == [] + + @staticmethod + def test_console_cors_exposes_download_filename( + monkeypatch: pytest.MonkeyPatch, + fresh_blueprints: dict[str, Blueprint], + ) -> None: + origin = "http://localhost:3000" + apply_config_overrides( + monkeypatch, + OPENAPI_ENABLED=False, + inner_CONSOLE_CORS_ALLOW_ORIGINS=origin, + ) + + app = _build_app() + ext_blueprints.init_app(app) + console_bp = fresh_blueprints["controllers.console"] + assert app.blueprints[console_bp.name] is console_bp + + response = app.test_client().get( + "/console/api/_download", + headers={"Origin": origin}, + ) + + exposed_headers = { + header.strip().lower() for header in response.headers["Access-Control-Expose-Headers"].split(",") + } + assert "content-disposition" in exposed_headers diff --git a/api/tests/unit_tests/extensions/test_ext_blueprints_cors.py b/api/tests/unit_tests/extensions/test_ext_blueprints_cors.py deleted file mode 100644 index 84cfa04d6eb..00000000000 --- a/api/tests/unit_tests/extensions/test_ext_blueprints_cors.py +++ /dev/null @@ -1,37 +0,0 @@ -"""Regression coverage for authenticated browser CORS headers.""" - -from flask import Blueprint, Flask - -from extensions.ext_blueprints import AUTHENTICATED_HEADERS, _apply_cors_once - - -def test_authenticated_cors_allows_request_metadata_headers() -> None: - app = Flask(__name__) - blueprint = Blueprint("cors_probe", __name__, url_prefix="/console/api") - - @blueprint.post("/probe") - def probe() -> tuple[str, int]: - return "", 204 - - _apply_cors_once( - blueprint, - resources={r"/*": {"origins": ["http://localhost:3000"]}}, - supports_credentials=True, - allow_headers=list(AUTHENTICATED_HEADERS), - methods=["POST", "OPTIONS"], - ) - app.register_blueprint(blueprint) - - response = app.test_client().options( - "/console/api/probe", - headers={ - "Access-Control-Request-Headers": "Idempotency-Key, X-Request-ID, X-Trace-ID", - "Access-Control-Request-Method": "POST", - "Origin": "http://localhost:3000", - }, - ) - - allowed_headers = response.headers.get("Access-Control-Allow-Headers", "").lower() - assert "idempotency-key" in allowed_headers - assert "x-request-id" in allowed_headers - assert "x-trace-id" in allowed_headers diff --git a/api/tests/unit_tests/extensions/test_ext_blueprints_openapi.py b/api/tests/unit_tests/extensions/test_ext_blueprints_openapi.py deleted file mode 100644 index a41bbe340f0..00000000000 --- a/api/tests/unit_tests/extensions/test_ext_blueprints_openapi.py +++ /dev/null @@ -1,158 +0,0 @@ -"""Verifies the OPENAPI_ENABLED gate in `extensions.ext_blueprints.init_app`. - -Contract: -- When `dify_config.OPENAPI_ENABLED` is True, the `/openapi/v1/*` blueprint - must be registered on the Flask app AND have CORS configured (signalled - by the idempotent `_dify_cors_applied` marker that `_apply_cors_once` - sets after wiring `flask_cors.CORS`). -- When False, the openapi blueprint must NOT be registered and CORS must - NOT be wired for it. The cross-origin posture is otherwise undefined for - the disabled feature. - -Why the blueprints are swapped for fresh ones: -`init_app` operates on module-level blueprint singletons (e.g. -`controllers.service_api.bp`). Other unit tests register those same -singletons onto throwaway Flask apps without going through `init_app` -(see `tests/unit_tests/controllers/test_swagger.py` and the -`tests/unit_tests/controllers/openapi/*` suite). Once a blueprint has -been registered to any app, Flask flips `_got_registered_once = True` -and rejects further `after_request` hookup -- which is what -`flask_cors.CORS(bp, ...)` does internally. To make this gate test -order-independent we monkeypatch each consumed `controllers..bp` to -a pristine, never-registered `Blueprint` for the duration of the test. -`init_app` resolves these via `from controllers. import bp as ...` -inside its function body, so the patched value is what it sees. -""" - -from __future__ import annotations - -import importlib - -import pytest -from flask import Blueprint, Response - -from dify_app import DifyApp -from extensions import ext_blueprints -from tests.unit_tests.config_override import apply_config_overrides - -# Modules whose `bp` attribute is consumed by `ext_blueprints.init_app`. -# Keep in sync with the imports inside `init_app`. -_BLUEPRINT_MODULES: tuple[str, ...] = ( - "controllers.console", - "controllers.files", - "controllers.inner_api", - "controllers.mcp", - "controllers.openapi", - "controllers.service_api", - "controllers.trigger", - "controllers.web", -) - - -def _probe_view() -> tuple[str, int]: - return "ok", 200 - - -def _download_probe_view() -> Response: - response = Response("document", content_type="text/plain") - response.headers["Content-Disposition"] = "attachment; filename=document.txt" - return response - - -@pytest.fixture -def fresh_blueprints(monkeypatch: pytest.MonkeyPatch) -> dict[str, Blueprint]: - """Replace each production blueprint singleton with a fresh, unregistered copy. - - Mirrors the production `name` and `url_prefix` so the gate can still - be asserted via `app.blueprints[...]` and url_map prefix checks. The - openapi replacement gets a dummy `/_probe` rule so the test can - observe at least one `/openapi/v1/*` rule after registration. - """ - fresh: dict[str, Blueprint] = {} - for module_name in _BLUEPRINT_MODULES: - module = importlib.import_module(module_name) - original = module.bp - replacement = Blueprint( - original.name, - original.import_name, - url_prefix=original.url_prefix, - ) - if module_name == "controllers.openapi": - replacement.add_url_rule("/_probe", endpoint="_probe", view_func=_probe_view) - if module_name == "controllers.console": - replacement.add_url_rule("/_download", endpoint="_download", view_func=_download_probe_view) - monkeypatch.setattr(module, "bp", replacement) - fresh[module_name] = replacement - return fresh - - -def _build_app() -> DifyApp: - app = DifyApp(__name__) - app.config["TESTING"] = True - return app - - -def test_openapi_blueprint_registered_with_cors_when_enabled( - monkeypatch: pytest.MonkeyPatch, - fresh_blueprints: dict[str, Blueprint], -) -> None: - """Enabled gate: blueprint mounted, CORS wired, `/openapi/v1/*` rules live.""" - apply_config_overrides(monkeypatch, OPENAPI_ENABLED=True) - - app = _build_app() - ext_blueprints.init_app(app) - - openapi_bp = fresh_blueprints["controllers.openapi"] - assert "openapi" in app.blueprints - assert app.blueprints["openapi"] is openapi_bp - # `_apply_cors_once` only sets this after wiring flask_cors.CORS, so - # it is a faithful proxy for "CORS was applied to this blueprint". - assert getattr(openapi_bp, "_dify_cors_applied", False) is True - - openapi_rules = [r for r in app.url_map.iter_rules() if r.rule.startswith("/openapi/v1")] - assert openapi_rules, "expected at least one /openapi/v1/* route once enabled" - - -def test_openapi_blueprint_absent_when_disabled( - monkeypatch: pytest.MonkeyPatch, - fresh_blueprints: dict[str, Blueprint], -) -> None: - """Disabled gate: no blueprint, no CORS, no `/openapi/v1/*` URL rules.""" - apply_config_overrides(monkeypatch, OPENAPI_ENABLED=False) - - app = _build_app() - ext_blueprints.init_app(app) - - openapi_bp = fresh_blueprints["controllers.openapi"] - assert "openapi" not in app.blueprints - # No CORS wiring should have run for the openapi blueprint. - assert not hasattr(openapi_bp, "_dify_cors_applied") - openapi_rules = [r for r in app.url_map.iter_rules() if r.rule.startswith("/openapi/v1")] - assert openapi_rules == [] - - -def test_console_cors_exposes_download_filename( - monkeypatch: pytest.MonkeyPatch, - fresh_blueprints: dict[str, Blueprint], -) -> None: - origin = "http://localhost:3000" - apply_config_overrides( - monkeypatch, - OPENAPI_ENABLED=False, - inner_CONSOLE_CORS_ALLOW_ORIGINS=origin, - ) - - app = _build_app() - ext_blueprints.init_app(app) - console_bp = fresh_blueprints["controllers.console"] - assert app.blueprints[console_bp.name] is console_bp - - response = app.test_client().get( - "/console/api/_download", - headers={"Origin": origin}, - ) - - exposed_headers = { - header.strip().lower() for header in response.headers["Access-Control-Expose-Headers"].split(",") - } - assert "content-disposition" in exposed_headers diff --git a/api/tests/unit_tests/extensions/test_ext_celery.py b/api/tests/unit_tests/extensions/test_ext_celery.py new file mode 100644 index 00000000000..9cf4bea1647 --- /dev/null +++ b/api/tests/unit_tests/extensions/test_ext_celery.py @@ -0,0 +1,471 @@ +import ssl +from collections.abc import Iterator +from datetime import timedelta +from importlib import import_module +from types import SimpleNamespace +from unittest.mock import MagicMock, Mock, patch + +import pytest +from celery.contrib.testing.app import TestApp, setup_default_app + +from dify_app import DifyApp +from enums import DeploymentEdition +from extensions.ext_celery import ( + _enqueue_initial_community_telemetry_heartbeat, + get_celery_broker_transport_options, + init_app, +) +from services.knowledge_fs.background_contract import OPERATION_QUEUES, OPERATION_TASK +from services.knowledge_fs.lifecycle_readiness import KnowledgeFSLifecycleWorkerReadiness +from tests.helpers.unit.config_override import apply_config_overrides + +"""Tests for Celery SSL configuration.""" + + +class TestCelerySslCases: + @staticmethod + @pytest.fixture(autouse=True) + def restore_celery_app_state() -> Iterator[None]: + # init_app replaces Celery's process-wide current/default app. Restore both + # so later shared tasks do not enter this module's Flask app context. + with setup_default_app(TestApp()): + yield + + class TestCelerySSLConfiguration: + """Test suite for Celery SSL configuration.""" + + def test_get_celery_broker_transport_options_includes_global_keyprefix_for_redis(self): + mock_config = MagicMock() + mock_config.CELERY_USE_SENTINEL = False + mock_config.REDIS_KEY_PREFIX = "enterprise-a" + + with patch("extensions.ext_celery.dify_config", mock_config): + from extensions.ext_celery import get_celery_broker_transport_options + + result = get_celery_broker_transport_options() + + assert result["global_keyprefix"] == "enterprise-a:" + + def test_get_celery_broker_transport_options_omits_global_keyprefix_when_prefix_empty(self): + mock_config = MagicMock() + mock_config.CELERY_USE_SENTINEL = False + mock_config.REDIS_KEY_PREFIX = " " + + with patch("extensions.ext_celery.dify_config", mock_config): + from extensions.ext_celery import get_celery_broker_transport_options + + result = get_celery_broker_transport_options() + + assert "global_keyprefix" not in result + + def test_get_celery_broker_transport_options_keeps_sentinel_and_adds_global_keyprefix(self): + mock_config = MagicMock() + mock_config.CELERY_USE_SENTINEL = True + mock_config.CELERY_SENTINEL_MASTER_NAME = "mymaster" + mock_config.CELERY_SENTINEL_SOCKET_TIMEOUT = 0.1 + mock_config.CELERY_SENTINEL_PASSWORD = "secret" + mock_config.REDIS_KEY_PREFIX = "enterprise-a" + + with patch("extensions.ext_celery.dify_config", mock_config): + from extensions.ext_celery import get_celery_broker_transport_options + + result = get_celery_broker_transport_options() + + assert result["master_name"] == "mymaster" + assert result["sentinel_kwargs"]["password"] == "secret" + assert result["global_keyprefix"] == "enterprise-a:" + + def test_get_celery_ssl_options_when_ssl_disabled(self): + """Test SSL options when BROKER_USE_SSL is False.""" + from configs import DifyConfig + + dify_config = DifyConfig(CELERY_BROKER_URL="redis://localhost:6379/0") + + with patch("extensions.ext_celery.dify_config", dify_config): + from extensions.ext_celery import get_celery_ssl_options + + result = get_celery_ssl_options() + assert result is None + + def test_get_celery_ssl_options_when_broker_not_redis(self): + """Test SSL options when broker is not Redis.""" + mock_config = MagicMock() + mock_config.CELERY_BROKER_URL = "amqp://localhost:5672" + + with patch("extensions.ext_celery.dify_config", mock_config): + from extensions.ext_celery import get_celery_ssl_options + + result = get_celery_ssl_options() + assert result is None + + def test_get_celery_ssl_options_with_cert_none(self): + """Test SSL options with CERT_NONE requirement.""" + mock_config = MagicMock() + mock_config.CELERY_BROKER_URL = "redis://localhost:6379/0" + mock_config.REDIS_SSL_CERT_REQS = "CERT_NONE" + mock_config.REDIS_SSL_CA_CERTS = None + mock_config.REDIS_SSL_CERTFILE = None + mock_config.REDIS_SSL_KEYFILE = None + + with patch("extensions.ext_celery.dify_config", mock_config): + from extensions.ext_celery import get_celery_ssl_options + + result = get_celery_ssl_options() + assert result is not None + assert result["ssl_cert_reqs"] == ssl.CERT_NONE + assert result["ssl_ca_certs"] is None + assert result["ssl_certfile"] is None + assert result["ssl_keyfile"] is None + + def test_get_celery_ssl_options_with_cert_required(self): + """Test SSL options with CERT_REQUIRED and certificates.""" + mock_config = MagicMock() + mock_config.CELERY_BROKER_URL = "rediss://localhost:6380/0" + mock_config.REDIS_SSL_CERT_REQS = "CERT_REQUIRED" + mock_config.REDIS_SSL_CA_CERTS = "/path/to/ca.crt" + mock_config.REDIS_SSL_CERTFILE = "/path/to/client.crt" + mock_config.REDIS_SSL_KEYFILE = "/path/to/client.key" + + with patch("extensions.ext_celery.dify_config", mock_config): + from extensions.ext_celery import get_celery_ssl_options + + result = get_celery_ssl_options() + assert result is not None + assert result["ssl_cert_reqs"] == ssl.CERT_REQUIRED + assert result["ssl_ca_certs"] == "/path/to/ca.crt" + assert result["ssl_certfile"] == "/path/to/client.crt" + assert result["ssl_keyfile"] == "/path/to/client.key" + + def test_get_celery_ssl_options_with_cert_optional(self): + """Test SSL options with CERT_OPTIONAL requirement.""" + mock_config = MagicMock() + mock_config.CELERY_BROKER_URL = "redis://localhost:6379/0" + mock_config.REDIS_SSL_CERT_REQS = "CERT_OPTIONAL" + mock_config.REDIS_SSL_CA_CERTS = "/path/to/ca.crt" + mock_config.REDIS_SSL_CERTFILE = None + mock_config.REDIS_SSL_KEYFILE = None + + with patch("extensions.ext_celery.dify_config", mock_config): + from extensions.ext_celery import get_celery_ssl_options + + result = get_celery_ssl_options() + assert result is not None + assert result["ssl_cert_reqs"] == ssl.CERT_OPTIONAL + assert result["ssl_ca_certs"] == "/path/to/ca.crt" + + def test_get_celery_ssl_options_with_invalid_cert_reqs(self): + """Test SSL options with invalid cert requirement defaults to CERT_NONE.""" + mock_config = MagicMock() + mock_config.CELERY_BROKER_URL = "redis://localhost:6379/0" + mock_config.REDIS_SSL_CERT_REQS = "INVALID_VALUE" + mock_config.REDIS_SSL_CA_CERTS = None + mock_config.REDIS_SSL_CERTFILE = None + mock_config.REDIS_SSL_KEYFILE = None + + with patch("extensions.ext_celery.dify_config", mock_config): + from extensions.ext_celery import get_celery_ssl_options + + result = get_celery_ssl_options() + assert result is not None + assert result["ssl_cert_reqs"] == ssl.CERT_NONE # Should default to CERT_NONE + + def test_celery_init_applies_ssl_to_broker_and_backend(self): + """Test that SSL options are applied to both broker and backend when using Redis.""" + mock_config = MagicMock() + mock_config.HUMAN_INPUT_TIMEOUT_TASK_INTERVAL = 1 + mock_config.CELERY_BROKER_URL = "redis://localhost:6379/0" + mock_config.CELERY_BACKEND = "redis" + mock_config.CELERY_RESULT_BACKEND = "redis://localhost:6379/0" + mock_config.REDIS_SSL_CERT_REQS = "CERT_NONE" + mock_config.REDIS_SSL_CA_CERTS = None + mock_config.REDIS_SSL_CERTFILE = None + mock_config.REDIS_SSL_KEYFILE = None + mock_config.CELERY_USE_SENTINEL = False + mock_config.LOG_FORMAT = "%(message)s" + mock_config.LOG_TZ = "UTC" + mock_config.LOG_FILE = None + + # Mock all the scheduler configs + mock_config.CELERY_BEAT_SCHEDULER_TIME = 1 + mock_config.AGENT_SANDBOX_METERING_ENABLED = False + mock_config.ENABLE_CONVERSATION_CLEANUP_TASK = False + mock_config.CONVERSATION_CLEANUP_TASK_INTERVAL = 5 + mock_config.ENABLE_CLEAN_EMBEDDING_CACHE_TASK = False + mock_config.ENABLE_CLEAN_UNUSED_DATASETS_TASK = False + mock_config.ENABLE_CREATE_TIDB_SERVERLESS_TASK = False + mock_config.ENABLE_UPDATE_TIDB_SERVERLESS_STATUS_TASK = False + mock_config.ENABLE_CLEAN_MESSAGES = False + mock_config.ENABLE_MAIL_CLEAN_DOCUMENT_NOTIFY_TASK = False + mock_config.ENABLE_DATASETS_QUEUE_MONITOR = False + mock_config.ENABLE_CHECK_UPGRADABLE_PLUGIN_TASK = False + mock_config.ENABLE_WORKFLOW_SCHEDULE_POLLER_TASK = False + mock_config.WORKFLOW_SCHEDULE_POLLER_INTERVAL = 1 + mock_config.WORKFLOW_SCHEDULE_POLLER_BATCH_SIZE = 100 + mock_config.WORKFLOW_SCHEDULE_MAX_DISPATCH_PER_TICK = 0 + mock_config.ENABLE_TRIGGER_PROVIDER_REFRESH_TASK = False + mock_config.TRIGGER_PROVIDER_REFRESH_INTERVAL = 15 + mock_config.ENABLE_API_TOKEN_LAST_USED_UPDATE_TASK = False + mock_config.API_TOKEN_LAST_USED_UPDATE_INTERVAL = 30 + + with ( + patch("extensions.ext_celery.dify_config", mock_config), + patch( + "services.knowledge_fs.lifecycle_readiness.get_configured_knowledge_fs_lifecycle_worker_readiness", + return_value=KnowledgeFSLifecycleWorkerReadiness(ready=False, blockers=()), + ), + ): + from dify_app import DifyApp + from extensions.ext_celery import init_app + + app = DifyApp(__name__) + celery_app = init_app(app) + + # Check that SSL options were applied + assert "broker_use_ssl" in celery_app.conf + assert celery_app.conf["broker_use_ssl"] is not None + assert celery_app.conf["broker_use_ssl"]["ssl_cert_reqs"] == ssl.CERT_NONE + + # Check that SSL is also applied to Redis backend + assert "redis_backend_use_ssl" in celery_app.conf + assert celery_app.conf["redis_backend_use_ssl"] is not None + + def test_celery_init_registers_required_agent_and_conversation_tasks(self): + mock_config = MagicMock() + mock_config.BROKER_USE_SSL = False + mock_config.REDIS_KEY_PREFIX = "enterprise-a" + mock_config.HUMAN_INPUT_TIMEOUT_TASK_INTERVAL = 1 + mock_config.CELERY_BROKER_URL = "redis://localhost:6379/0" + mock_config.CELERY_BACKEND = "redis" + mock_config.CELERY_RESULT_BACKEND = "redis://localhost:6379/0" + mock_config.CELERY_USE_SENTINEL = False + mock_config.LOG_FORMAT = "%(message)s" + mock_config.LOG_TZ = "UTC" + mock_config.LOG_FILE = None + mock_config.CELERY_TASK_ANNOTATIONS = {} + + mock_config.CELERY_BEAT_SCHEDULER_TIME = 1 + mock_config.AGENT_SANDBOX_METERING_ENABLED = False + mock_config.ENABLE_CONVERSATION_CLEANUP_TASK = True + mock_config.CONVERSATION_CLEANUP_TASK_INTERVAL = 5 + mock_config.ENABLE_CLEAN_EMBEDDING_CACHE_TASK = False + mock_config.ENABLE_CLEAN_UNUSED_DATASETS_TASK = False + mock_config.ENABLE_CREATE_TIDB_SERVERLESS_TASK = False + mock_config.ENABLE_UPDATE_TIDB_SERVERLESS_STATUS_TASK = False + mock_config.ENABLE_CLEAN_MESSAGES = False + mock_config.ENABLE_MAIL_CLEAN_DOCUMENT_NOTIFY_TASK = False + mock_config.ENABLE_DATASETS_QUEUE_MONITOR = False + mock_config.ENABLE_HUMAN_INPUT_TIMEOUT_TASK = False + mock_config.ENABLE_CHECK_UPGRADABLE_PLUGIN_TASK = False + mock_config.MARKETPLACE_ENABLED = False + mock_config.WORKFLOW_LOG_CLEANUP_ENABLED = False + mock_config.ENABLE_WORKFLOW_RUN_CLEANUP_TASK = False + mock_config.ENABLE_WORKFLOW_SCHEDULE_POLLER_TASK = False + mock_config.WORKFLOW_SCHEDULE_POLLER_INTERVAL = 1 + mock_config.ENABLE_TRIGGER_PROVIDER_REFRESH_TASK = False + mock_config.TRIGGER_PROVIDER_REFRESH_INTERVAL = 15 + mock_config.ENABLE_API_TOKEN_LAST_USED_UPDATE_TASK = False + mock_config.API_TOKEN_LAST_USED_UPDATE_INTERVAL = 30 + mock_config.DEPLOYMENT_EDITION = DeploymentEdition.COMMUNITY + mock_config.ENTERPRISE_TELEMETRY_ENABLED = False + + with ( + patch("extensions.ext_celery.dify_config", mock_config), + patch( + "services.knowledge_fs.lifecycle_readiness.get_configured_knowledge_fs_lifecycle_worker_readiness", + return_value=KnowledgeFSLifecycleWorkerReadiness(ready=False, blockers=()), + ), + ): + from dify_app import DifyApp + from extensions.ext_celery import init_app + + app = DifyApp(__name__) + celery_app = init_app(app) + + assert celery_app.conf["broker_transport_options"]["global_keyprefix"] == "enterprise-a:" + assert celery_app.conf["result_backend_transport_options"]["global_keyprefix"] == "enterprise-a:" + assert "tasks.collect_agent_resources_task" in celery_app.conf["imports"] + assert "tasks.delete_conversation_task" in celery_app.conf["imports"] + assert "tasks.workflow_run_archive_deletion_tasks" in celery_app.conf["imports"] + import_module("tasks.workflow_run_archive_deletion_tasks") + task_name = "tasks.workflow_run_archive_deletion_tasks.delete_workflow_run_archive_task" + archive_deletion = celery_app.tasks[task_name] + assert archive_deletion.queue == "workflow_archive" + assert archive_deletion.acks_late is True + assert archive_deletion.reject_on_worker_lost is True + assert archive_deletion.autoretry_for == (Exception,) + assert archive_deletion.retry_kwargs["max_retries"] == 5 + assert celery_app.conf["beat_schedule"]["conversation_cleanup_sweeper"]["task"] == ( + "tasks.delete_conversation_task.sweep_deleted_conversations" + ) + + +class TestCommunityTelemetryCeleryCases: + @staticmethod + def test_beat_start_enqueues_community_telemetry_heartbeat() -> None: + task = Mock() + sender = SimpleNamespace( + app=SimpleNamespace( + conf=SimpleNamespace(beat_schedule={"community_telemetry_heartbeat": {}}), + tasks={"community_telemetry.send_heartbeat": task}, + ) + ) + + _enqueue_initial_community_telemetry_heartbeat(sender) + + task.apply_async.assert_called_once_with() + + @staticmethod + def test_beat_start_skips_community_telemetry_when_not_scheduled() -> None: + task = Mock() + sender = SimpleNamespace( + app=SimpleNamespace( + conf=SimpleNamespace(beat_schedule={}), + tasks={"community_telemetry.send_heartbeat": task}, + ) + ) + + _enqueue_initial_community_telemetry_heartbeat(sender) + + task.apply_async.assert_not_called() + + @staticmethod + def test_beat_start_skips_community_telemetry_when_task_is_unavailable() -> None: + sender = SimpleNamespace( + app=SimpleNamespace( + conf=SimpleNamespace(beat_schedule={"community_telemetry_heartbeat": {}}), + tasks={}, + ) + ) + + _enqueue_initial_community_telemetry_heartbeat(sender) + + +class TestKnowledgeFsCeleryCases: + @staticmethod + @pytest.fixture(autouse=True) + def restore_celery_app_state() -> Iterator[None]: + # init_app replaces Celery's process-wide current/default app; restore both so + # later shared tasks still resolve against the app that registered them. + with setup_default_app(TestApp()): + yield + + @staticmethod + @pytest.mark.parametrize("background_enabled", [False, True]) + @pytest.mark.parametrize("vector_store", ["qdrant", "tidb_on_qdrant", "elasticsearch", "weaviate"]) + def test_celery_registers_initial_source_task_when_knowledge_fs_lifecycle_is_ready( + background_enabled: bool, vector_store: str, monkeypatch: pytest.MonkeyPatch + ) -> None: + # A stale deployment variable must not suppress vector cleanup for the selected provider. + monkeypatch.setenv("KNOWLEDGE_VECTOR_STORAGE", "postgres") + apply_config_overrides( + monkeypatch, + REDIS_KEY_PREFIX="test", + HUMAN_INPUT_TIMEOUT_TASK_INTERVAL=1, + CELERY_BROKER_URL="redis://localhost:6379/0", + CELERY_BACKEND="redis", + CELERY_USE_SENTINEL=False, + CELERY_BROKER_VISIBILITY_TIMEOUT=10800, + KNOWLEDGE_FS_BACKGROUND_WORKER_ENABLED=background_enabled, + VECTOR_STORE=vector_store, + KNOWLEDGE_FS_BACKGROUND_POLL_INTERVAL_SECONDS=5, + KNOWLEDGE_FS_BACKGROUND_TASK_TIMEOUT_SECONDS=7200, + LOG_FORMAT="%(message)s", + LOG_TZ="UTC", + LOG_FILE=None, + CELERY_TASK_ANNOTATIONS=dict[str, object](), + CELERY_BEAT_SCHEDULER_TIME=1, + KNOWLEDGE_FS_LIFECYCLE_POLL_INTERVAL_SECONDS=2, + ENABLE_CONVERSATION_CLEANUP_TASK=True, + CONVERSATION_CLEANUP_TASK_INTERVAL=5, + ENABLE_CLEAN_EMBEDDING_CACHE_TASK=False, + ENABLE_CLEAN_UNUSED_DATASETS_TASK=False, + ENABLE_CREATE_TIDB_SERVERLESS_TASK=False, + ENABLE_UPDATE_TIDB_SERVERLESS_STATUS_TASK=False, + ENABLE_CLEAN_MESSAGES=False, + ENABLE_MAIL_CLEAN_DOCUMENT_NOTIFY_TASK=False, + ENABLE_DATASETS_QUEUE_MONITOR=False, + ENABLE_HUMAN_INPUT_TIMEOUT_TASK=False, + ENABLE_CHECK_UPGRADABLE_PLUGIN_TASK=False, + MARKETPLACE_ENABLED=False, + WORKFLOW_LOG_CLEANUP_ENABLED=False, + ENABLE_WORKFLOW_RUN_CLEANUP_TASK=False, + ENABLE_WORKFLOW_SCHEDULE_POLLER_TASK=False, + WORKFLOW_SCHEDULE_POLLER_INTERVAL=1, + ENABLE_TRIGGER_PROVIDER_REFRESH_TASK=False, + TRIGGER_PROVIDER_REFRESH_INTERVAL=15, + ENABLE_API_TOKEN_LAST_USED_UPDATE_TASK=False, + API_TOKEN_LAST_USED_UPDATE_INTERVAL=30, + ENTERPRISE_TELEMETRY_ENABLED=False, + ) + + with ( + patch( + "services.knowledge_fs.lifecycle_readiness.get_configured_knowledge_fs_lifecycle_worker_readiness", + return_value=KnowledgeFSLifecycleWorkerReadiness(ready=True, blockers=()), + ), + ): + celery_app = init_app(DifyApp(__name__)) + + assert "tasks.knowledge_fs_initial_source_tasks" in celery_app.conf["imports"] + assert "tasks.knowledge_fs_initial_source_preview_tasks" in celery_app.conf["imports"] + assert "tasks.knowledge_fs_failed_retrieval_tasks" in celery_app.conf["imports"] + assert "tasks.knowledge_fs_lifecycle_tasks" in celery_app.conf["imports"] + assert "tasks.knowledge_fs_source_import_tasks" in celery_app.conf["imports"] + assert "tasks.delete_conversation_task" in celery_app.conf["imports"] + assert celery_app.conf["beat_schedule"]["conversation_cleanup_sweeper"] == { + "task": "tasks.delete_conversation_task.sweep_deleted_conversations", + "schedule": timedelta(minutes=5), + } + assert celery_app.conf["beat_schedule"]["knowledge_fs_staged_upload_cleanup"] == { + "task": "tasks.knowledge_fs_lifecycle_tasks.cleanup_knowledge_fs_staged_uploads", + "schedule": timedelta(seconds=2), + } + assert celery_app.conf["beat_schedule"]["knowledge_fs_upgrade_file_cleanup"] == { + "task": "tasks.knowledge_fs_upgrade_tasks.cleanup_deferred_knowledge_fs_upgrade_files", + "schedule": timedelta(seconds=2), + } + assert celery_app.conf["beat_schedule"]["knowledge_fs_upgrade_recovery"] == { + "task": "tasks.knowledge_fs_upgrade_tasks.recover_stalled_knowledge_fs_upgrades", + "schedule": timedelta(seconds=60), + "options": {"expires": 60}, + } + schedules = celery_app.conf["beat_schedule"] + if background_enabled: + assert celery_app.conf["broker_transport_options"]["visibility_timeout"] == 10800 + for operation, queue in OPERATION_QUEUES.items(): + assert schedules[f"knowledge_fs_background_{operation}"] == { + "task": OPERATION_TASK, + "schedule": timedelta(seconds=5), + "kwargs": {"operation": operation}, + "options": {"queue": queue, "expires": 30}, + } + else: + assert not any(name.startswith("knowledge_fs_background_") for name in schedules) + + with ( + patch( + "services.knowledge_fs.lifecycle_readiness.get_configured_knowledge_fs_lifecycle_worker_readiness", + return_value=KnowledgeFSLifecycleWorkerReadiness(ready=False, blockers=("test",)), + ), + ): + preview_only_app = init_app(DifyApp(f"{__name__}.preview_only")) + + assert "tasks.knowledge_fs_initial_source_preview_tasks" in preview_only_app.conf["imports"] + assert "tasks.knowledge_fs_failed_retrieval_tasks" in preview_only_app.conf["imports"] + assert "tasks.knowledge_fs_initial_source_tasks" not in preview_only_app.conf["imports"] + assert "tasks.knowledge_fs_lifecycle_tasks" not in preview_only_app.conf["imports"] + assert "tasks.knowledge_fs_source_import_tasks" not in preview_only_app.conf["imports"] + assert "knowledge_fs_upgrade_recovery" not in preview_only_app.conf["beat_schedule"] + + @staticmethod + @pytest.mark.parametrize("visibility", [None, 3600, 7200, 7230]) + def test_long_running_delivery_rejects_an_unsafe_redis_visibility(visibility: int | None) -> None: + with patch("extensions.ext_celery.dify_config") as config: + config.CELERY_USE_SENTINEL = False + config.REDIS_KEY_PREFIX = "" + config.CELERY_BROKER_URL = "redis://localhost:6379/0" + config.CELERY_BROKER_VISIBILITY_TIMEOUT = visibility + config.KNOWLEDGE_FS_BACKGROUND_WORKER_ENABLED = True + config.KNOWLEDGE_FS_BACKGROUND_TASK_TIMEOUT_SECONDS = 7200 + with pytest.raises(ValueError, match="VISIBILITY_TIMEOUT"): + get_celery_broker_transport_options() diff --git a/api/tests/unit_tests/extensions/test_ext_knowledge_fs_observability.py b/api/tests/unit_tests/extensions/test_ext_knowledge_fs_observability.py index 7e8e5a661c2..b17e71d625a 100644 --- a/api/tests/unit_tests/extensions/test_ext_knowledge_fs_observability.py +++ b/api/tests/unit_tests/extensions/test_ext_knowledge_fs_observability.py @@ -9,7 +9,7 @@ from extensions import ext_knowledge_fs_observability from repositories.sqlalchemy_knowledge_fs_control_space_state_metrics import ( SQLAlchemyKnowledgeFSControlSpaceStateCountReader, ) -from tests.unit_tests.tasks.knowledge_fs_task_fakes import RecordingKnowledgeFSMetrics +from tests.helpers.unit.tasks.knowledge_fs_task_fakes import RecordingKnowledgeFSMetrics def test_extension_registers_the_database_backed_control_space_gauge() -> None: diff --git a/api/tests/unit_tests/extensions/otel/test_flask_instrumentation.py b/api/tests/unit_tests/extensions/test_ext_otel.py similarity index 69% rename from api/tests/unit_tests/extensions/otel/test_flask_instrumentation.py rename to api/tests/unit_tests/extensions/test_ext_otel.py index a9ea47f06be..01063255fd7 100644 --- a/api/tests/unit_tests/extensions/otel/test_flask_instrumentation.py +++ b/api/tests/unit_tests/extensions/test_ext_otel.py @@ -15,6 +15,16 @@ import flask import pytest from opentelemetry.instrumentation.flask import FlaskInstrumentor +from tests.unit_tests.extensions.otel.conftest import memory_span_exporter as memory_span_exporter # noqa: PLC0414 +from tests.unit_tests.extensions.otel.conftest import mock_account_user as mock_account_user # noqa: PLC0414 +from tests.unit_tests.extensions.otel.conftest import mock_app_model as mock_app_model # noqa: PLC0414 +from tests.unit_tests.extensions.otel.conftest import mock_end_user as mock_end_user # noqa: PLC0414 +from tests.unit_tests.extensions.otel.conftest import mock_workflow_runner as mock_workflow_runner # noqa: PLC0414 +from tests.unit_tests.extensions.otel.conftest import reset_handler_instances as reset_handler_instances # noqa: PLC0414 +from tests.unit_tests.extensions.otel.conftest import ( + tracer_provider_with_memory_exporter as tracer_provider_with_memory_exporter, # noqa: PLC0414 +) + @pytest.fixture def instrumented_app() -> Iterator[flask.Flask]: diff --git a/api/tests/unit_tests/extensions/test_ext_redis.py b/api/tests/unit_tests/extensions/test_ext_redis.py new file mode 100644 index 00000000000..b39713926f3 --- /dev/null +++ b/api/tests/unit_tests/extensions/test_ext_redis.py @@ -0,0 +1,322 @@ +from collections.abc import Callable, Mapping +from dataclasses import dataclass, field + +import pytest +from redis import RedisError +from redis.retry import Retry + +from extensions import ext_redis +from extensions.ext_redis import ( + RedisClientWrapper, + _get_base_redis_params, + _get_cluster_connection_health_params, + _get_connection_health_params, + _normalize_redis_key_prefix, + _parse_redis_nodes, + _serialize_redis_name, + redis_fallback, +) +from libs.broadcast_channel.redis.pubsub_channel import BroadcastChannel as RedisBroadcastChannel +from libs.broadcast_channel.redis.sharded_channel import ShardedRedisBroadcastChannel +from tests.helpers.unit.config_override import apply_config_overrides + + +class TestPubsubChannelCases: + @staticmethod + def test_get_pubsub_broadcast_channel_defaults_to_pubsub(monkeypatch: pytest.MonkeyPatch): + apply_config_overrides(monkeypatch, PUBSUB_REDIS_CHANNEL_TYPE="pubsub") + monkeypatch.setattr(ext_redis, "_pubsub_redis_client", object()) + + channel = ext_redis.get_pubsub_broadcast_channel() + + assert isinstance(channel, RedisBroadcastChannel) + + @staticmethod + def test_get_pubsub_broadcast_channel_sharded(monkeypatch: pytest.MonkeyPatch): + apply_config_overrides(monkeypatch, PUBSUB_REDIS_CHANNEL_TYPE="sharded") + monkeypatch.setattr(ext_redis, "_pubsub_redis_client", object()) + + channel = ext_redis.get_pubsub_broadcast_channel() + + assert isinstance(channel, ShardedRedisBroadcastChannel) + + +type _RecordedCall = tuple[str, tuple[object, ...], dict[str, object]] + + +@dataclass +class _RecordingRedisClient: + calls: list[_RecordedCall] = field(default_factory=list) + + def _record(self, method: str, *args: object, **kwargs: object) -> None: + self.calls.append((method, args, kwargs)) + + def register_script(self, script: str) -> Callable[..., str]: + self._record("register_script", script) + + def execute(*, keys: object, args: object, client: object) -> str: + self._record("script", keys, args, client) + return "result" + + return execute + + def get(self, name: str | bytes) -> None: + self._record("get", name) + + def delete(self, *names: str | bytes) -> None: + self._record("delete", *names) + + def lock(self, name: str, **kwargs: object) -> None: + self._record("lock", name, **kwargs) + + def hset(self, name: str | bytes, *args: object, **kwargs: object) -> None: + self._record("hset", name, *args, **kwargs) + + def hgetall(self, name: str | bytes) -> None: + self._record("hgetall", name) + + def hkeys(self, name: str | bytes) -> None: + self._record("hkeys", name) + + def hexists(self, name: str | bytes, key: str | bytes) -> None: + self._record("hexists", name, key) + + def zadd(self, name: str | bytes, mapping: Mapping[object, object], **kwargs: object) -> None: + self._record("zadd", name, mapping, **kwargs) + + def zrem(self, name: str | bytes, *values: str | bytes) -> None: + self._record("zrem", name, *values) + + +class TestRedisCases: + @staticmethod + @pytest.fixture(autouse=True) + def _redis_config(config_overrides: Callable[..., None]) -> None: + config_overrides( + REDIS_USERNAME=None, + REDIS_PASSWORD=None, + REDIS_DB=0, + REDIS_SERIALIZATION_PROTOCOL=3, + REDIS_ENABLE_CLIENT_SIDE_CACHE=False, + REDIS_RETRY_RETRIES=3, + REDIS_RETRY_BACKOFF_BASE=1.0, + REDIS_RETRY_BACKOFF_CAP=10.0, + REDIS_SOCKET_TIMEOUT=5.0, + REDIS_SOCKET_CONNECT_TIMEOUT=5.0, + REDIS_HEALTH_CHECK_INTERVAL=30, + REDIS_KEEPALIVE=True, + REDIS_KEEPALIVE_IDLE=60, + REDIS_KEEPALIVE_INTERVAL=10, + REDIS_KEEPALIVE_COUNT=3, + REDIS_KEY_PREFIX="", + ) + + class TestGetConnectionHealthParams: + def test_includes_all_health_params(self) -> None: + params = _get_connection_health_params() + + assert "retry" in params + assert "socket_timeout" in params + assert "socket_connect_timeout" in params + assert "health_check_interval" in params + assert isinstance(params["retry"], Retry) + assert params["retry"]._retries == 3 # pyrefly: ignore[missing-attribute] + assert params["socket_timeout"] == 5.0 + assert params["socket_connect_timeout"] == 5.0 + assert params["health_check_interval"] == 30 + + class TestGetClusterConnectionHealthParams: + def test_excludes_health_check_interval(self) -> None: + params = _get_cluster_connection_health_params() + + assert "retry" in params + assert "socket_timeout" in params + assert "socket_connect_timeout" in params + assert "health_check_interval" not in params + + class TestGetBaseRedisParams: + def test_includes_retry_and_health_params(self) -> None: + params = _get_base_redis_params() + + assert "retry" in params + assert isinstance(params["retry"], Retry) + assert params["socket_timeout"] == 5.0 + assert params["socket_connect_timeout"] == 5.0 + assert params["health_check_interval"] == 30 + assert params["socket_keepalive"] is True + assert isinstance(params["socket_keepalive_options"], dict) + # Existing params still present + assert params["db"] == 0 + assert params["encoding"] == "utf-8" + + class TestParseRedisNodes: + def test_trims_nodes(self) -> None: + assert _parse_redis_nodes("redis-a:6379, redis-b:6380") == [("redis-a", 6379), ("redis-b", 6380)] + + def test_supports_bracketed_ipv6(self) -> None: + assert _parse_redis_nodes("[2001:db8::10]:6379") == [("2001:db8::10", 6379)] + + class TestRedisFallback: + def test_redis_fallback_success(self) -> None: + @redis_fallback(default_return=None) + def test_func() -> str: + return "success" + + assert test_func() == "success" + + def test_redis_fallback_error(self) -> None: + @redis_fallback(default_return="fallback") + def test_func() -> None: + raise RedisError("Redis error") + + assert test_func() == "fallback" + + def test_redis_fallback_none_default(self) -> None: + @redis_fallback() + def test_func() -> None: + raise RedisError("Redis error") + + assert test_func() is None + + def test_redis_fallback_with_args(self) -> None: + received: list[tuple[int, int]] = [] + + @redis_fallback(default_return=0) + def test_func(x: int, y: int) -> None: + received.append((x, y)) + raise RedisError("Redis error") + + assert test_func(1, 2) == 0 + assert received == [(1, 2)] + + def test_redis_fallback_with_kwargs(self) -> None: + received: list[tuple[int | None, int | None]] = [] + + @redis_fallback(default_return=dict[str, object]()) + def test_func(x: int | None = None, y: int | None = None) -> None: + received.append((x, y)) + raise RedisError("Redis error") + + assert test_func(x=1, y=2) == {} + assert received == [(1, 2)] + + def test_redis_fallback_preserves_function_metadata(self) -> None: + @redis_fallback(default_return=None) + def test_func() -> None: + """Test function docstring""" + pass + + assert test_func.__name__ == "test_func" + assert test_func.__doc__ == "Test function docstring" + + class TestRedisKeyPrefixHelpers: + def test_normalize_redis_key_prefix_trims_whitespace(self) -> None: + assert _normalize_redis_key_prefix(" enterprise-a ") == "enterprise-a" + + def test_normalize_redis_key_prefix_treats_whitespace_only_as_empty(self) -> None: + assert _normalize_redis_key_prefix(" ") == "" + + def test_serialize_redis_name_returns_original_when_prefix_empty(self) -> None: + assert _serialize_redis_name("model_lb_index:test", "") == "model_lb_index:test" + + def test_serialize_redis_name_adds_single_colon_separator(self) -> None: + assert _serialize_redis_name("model_lb_index:test", "enterprise-a") == "enterprise-a:model_lb_index:test" + + class TestRedisClientWrapperKeyPrefix: + def test_wrapper_zrem_prefixes_key_but_not_members(self, config_overrides: Callable[..., None]) -> None: + client = _RecordingRedisClient() + wrapper = RedisClientWrapper() + wrapper.initialize(client) # type: ignore[arg-type] + config_overrides(REDIS_KEY_PREFIX="enterprise-a") + + wrapper.zrem("reset:user@example.com", "reservation") + + assert client.calls == [("zrem", ("enterprise-a:reset:user@example.com", "reservation"), {})] + + def test_wrapper_registered_script_prefixes_key_arguments(self, config_overrides: Callable[..., None]) -> None: + raw_client = _RecordingRedisClient() + wrapper = RedisClientWrapper() + wrapper.initialize(raw_client) # type: ignore[arg-type] + + config_overrides(REDIS_KEY_PREFIX="enterprise-a") + script = wrapper.register_script("return redis.call('GET', KEYS[1])") + + assert script(keys=["device_code:abc"], args=["argument"]) == "result" + assert raw_client.calls == [ + ("register_script", ("return redis.call('GET', KEYS[1])",), {}), + ("script", (("enterprise-a:device_code:abc",), ["argument"], raw_client), {}), + ] + + def test_wrapper_get_prefixes_string_keys(self, config_overrides: Callable[..., None]) -> None: + client = _RecordingRedisClient() + wrapper = RedisClientWrapper() + wrapper.initialize(client) # type: ignore[arg-type] + + config_overrides(REDIS_KEY_PREFIX="enterprise-a") + wrapper.get("oauth_state:abc") + + assert client.calls == [("get", ("enterprise-a:oauth_state:abc",), {})] + + def test_wrapper_delete_prefixes_multiple_keys(self, config_overrides: Callable[..., None]) -> None: + client = _RecordingRedisClient() + wrapper = RedisClientWrapper() + wrapper.initialize(client) # type: ignore[arg-type] + + config_overrides(REDIS_KEY_PREFIX="enterprise-a") + wrapper.delete("key:a", "key:b") + + assert client.calls == [("delete", ("enterprise-a:key:a", "enterprise-a:key:b"), {})] + + def test_wrapper_lock_prefixes_lock_name(self, config_overrides: Callable[..., None]) -> None: + client = _RecordingRedisClient() + wrapper = RedisClientWrapper() + wrapper.initialize(client) # type: ignore[arg-type] + + config_overrides(REDIS_KEY_PREFIX="enterprise-a") + wrapper.lock("resource-lock", timeout=10) + + method, args, kwargs = client.calls[0] + assert method == "lock" + assert args == ("enterprise-a:resource-lock",) + assert kwargs["timeout"] == 10 + + def test_wrapper_hash_operations_prefix_key_name(self, config_overrides: Callable[..., None]) -> None: + client = _RecordingRedisClient() + wrapper = RedisClientWrapper() + wrapper.initialize(client) # type: ignore[arg-type] + + config_overrides(REDIS_KEY_PREFIX="enterprise-a") + wrapper.hset("hash:key", "field", "value") + wrapper.hgetall("hash:key") + wrapper.hkeys("hash:key") + wrapper.hexists("hash:key", "field") + + assert client.calls == [ + ("hset", ("enterprise-a:hash:key", "field", "value"), {}), + ("hgetall", ("enterprise-a:hash:key",), {}), + ("hkeys", ("enterprise-a:hash:key",), {}), + ("hexists", ("enterprise-a:hash:key", "field"), {}), + ] + + def test_wrapper_zadd_prefixes_sorted_set_name(self, config_overrides: Callable[..., None]) -> None: + client = _RecordingRedisClient() + wrapper = RedisClientWrapper() + wrapper.initialize(client) # type: ignore[arg-type] + + config_overrides(REDIS_KEY_PREFIX="enterprise-a") + wrapper.zadd("zset:key", {"member": 1}) + + method, args, kwargs = client.calls[0] + assert method == "zadd" + assert args == ("enterprise-a:zset:key", {"member": 1}) + assert kwargs["nx"] is False + + def test_wrapper_preserves_keys_when_prefix_is_empty(self, config_overrides: Callable[..., None]) -> None: + client = _RecordingRedisClient() + wrapper = RedisClientWrapper() + wrapper.initialize(client) # type: ignore[arg-type] + + config_overrides(REDIS_KEY_PREFIX=" ") + wrapper.get("plain:key") + + assert client.calls == [("get", ("plain:key",), {})] diff --git a/api/tests/unit_tests/extensions/test_ext_request_logging.py b/api/tests/unit_tests/extensions/test_ext_request_logging.py index c7512fc6914..26b7f94c1f3 100644 --- a/api/tests/unit_tests/extensions/test_ext_request_logging.py +++ b/api/tests/unit_tests/extensions/test_ext_request_logging.py @@ -8,7 +8,7 @@ from flask import Flask, Response from extensions import ext_request_logging from extensions.ext_request_logging import _is_content_type_json, _log_request_finished, init_app -from tests.unit_tests.config_override import apply_config_overrides +from tests.helpers.unit.config_override import apply_config_overrides def test_is_content_type_json() -> None: diff --git a/api/tests/unit_tests/extensions/test_set_secretkey.py b/api/tests/unit_tests/extensions/test_ext_set_secretkey.py similarity index 97% rename from api/tests/unit_tests/extensions/test_set_secretkey.py rename to api/tests/unit_tests/extensions/test_ext_set_secretkey.py index 767e30de77b..dc621ac76ca 100644 --- a/api/tests/unit_tests/extensions/test_set_secretkey.py +++ b/api/tests/unit_tests/extensions/test_ext_set_secretkey.py @@ -4,7 +4,7 @@ import pytest from dify_app import DifyApp from extensions import ext_set_secretkey -from tests.unit_tests.config_override import apply_config_overrides +from tests.helpers.unit.config_override import apply_config_overrides class InMemoryStorage: diff --git a/api/tests/unit_tests/extensions/test_knowledge_fs_celery.py b/api/tests/unit_tests/extensions/test_knowledge_fs_celery.py deleted file mode 100644 index 68641997b99..00000000000 --- a/api/tests/unit_tests/extensions/test_knowledge_fs_celery.py +++ /dev/null @@ -1,141 +0,0 @@ -from collections.abc import Iterator -from datetime import timedelta -from unittest.mock import patch - -import pytest -from celery.contrib.testing.app import TestApp, setup_default_app - -from dify_app import DifyApp -from extensions.ext_celery import get_celery_broker_transport_options, init_app -from services.knowledge_fs.background_contract import OPERATION_QUEUES, OPERATION_TASK -from services.knowledge_fs.lifecycle_readiness import KnowledgeFSLifecycleWorkerReadiness -from tests.unit_tests.config_override import apply_config_overrides - - -@pytest.fixture(autouse=True) -def restore_celery_app_state() -> Iterator[None]: - # init_app replaces Celery's process-wide current/default app; restore both so - # later shared tasks still resolve against the app that registered them. - with setup_default_app(TestApp()): - yield - - -@pytest.mark.parametrize("background_enabled", [False, True]) -@pytest.mark.parametrize("vector_store", ["qdrant", "tidb_on_qdrant", "elasticsearch", "weaviate"]) -def test_celery_registers_initial_source_task_when_knowledge_fs_lifecycle_is_ready( - background_enabled: bool, vector_store: str, monkeypatch: pytest.MonkeyPatch -) -> None: - # A stale deployment variable must not suppress vector cleanup for the selected provider. - monkeypatch.setenv("KNOWLEDGE_VECTOR_STORAGE", "postgres") - apply_config_overrides( - monkeypatch, - REDIS_KEY_PREFIX="test", - HUMAN_INPUT_TIMEOUT_TASK_INTERVAL=1, - CELERY_BROKER_URL="redis://localhost:6379/0", - CELERY_BACKEND="redis", - CELERY_USE_SENTINEL=False, - CELERY_BROKER_VISIBILITY_TIMEOUT=10800, - KNOWLEDGE_FS_BACKGROUND_WORKER_ENABLED=background_enabled, - VECTOR_STORE=vector_store, - KNOWLEDGE_FS_BACKGROUND_POLL_INTERVAL_SECONDS=5, - KNOWLEDGE_FS_BACKGROUND_TASK_TIMEOUT_SECONDS=7200, - LOG_FORMAT="%(message)s", - LOG_TZ="UTC", - LOG_FILE=None, - CELERY_TASK_ANNOTATIONS=dict[str, object](), - CELERY_BEAT_SCHEDULER_TIME=1, - KNOWLEDGE_FS_LIFECYCLE_POLL_INTERVAL_SECONDS=2, - ENABLE_CONVERSATION_CLEANUP_TASK=True, - CONVERSATION_CLEANUP_TASK_INTERVAL=5, - ENABLE_CLEAN_EMBEDDING_CACHE_TASK=False, - ENABLE_CLEAN_UNUSED_DATASETS_TASK=False, - ENABLE_CREATE_TIDB_SERVERLESS_TASK=False, - ENABLE_UPDATE_TIDB_SERVERLESS_STATUS_TASK=False, - ENABLE_CLEAN_MESSAGES=False, - ENABLE_MAIL_CLEAN_DOCUMENT_NOTIFY_TASK=False, - ENABLE_DATASETS_QUEUE_MONITOR=False, - ENABLE_HUMAN_INPUT_TIMEOUT_TASK=False, - ENABLE_CHECK_UPGRADABLE_PLUGIN_TASK=False, - MARKETPLACE_ENABLED=False, - WORKFLOW_LOG_CLEANUP_ENABLED=False, - ENABLE_WORKFLOW_RUN_CLEANUP_TASK=False, - ENABLE_WORKFLOW_SCHEDULE_POLLER_TASK=False, - WORKFLOW_SCHEDULE_POLLER_INTERVAL=1, - ENABLE_TRIGGER_PROVIDER_REFRESH_TASK=False, - TRIGGER_PROVIDER_REFRESH_INTERVAL=15, - ENABLE_API_TOKEN_LAST_USED_UPDATE_TASK=False, - API_TOKEN_LAST_USED_UPDATE_INTERVAL=30, - ENTERPRISE_TELEMETRY_ENABLED=False, - ) - - with ( - patch( - "services.knowledge_fs.lifecycle_readiness.get_configured_knowledge_fs_lifecycle_worker_readiness", - return_value=KnowledgeFSLifecycleWorkerReadiness(ready=True, blockers=()), - ), - ): - celery_app = init_app(DifyApp(__name__)) - - assert "tasks.knowledge_fs_initial_source_tasks" in celery_app.conf["imports"] - assert "tasks.knowledge_fs_initial_source_preview_tasks" in celery_app.conf["imports"] - assert "tasks.knowledge_fs_failed_retrieval_tasks" in celery_app.conf["imports"] - assert "tasks.knowledge_fs_lifecycle_tasks" in celery_app.conf["imports"] - assert "tasks.knowledge_fs_source_import_tasks" in celery_app.conf["imports"] - assert "tasks.delete_conversation_task" in celery_app.conf["imports"] - assert celery_app.conf["beat_schedule"]["conversation_cleanup_sweeper"] == { - "task": "tasks.delete_conversation_task.sweep_deleted_conversations", - "schedule": timedelta(minutes=5), - } - assert celery_app.conf["beat_schedule"]["knowledge_fs_staged_upload_cleanup"] == { - "task": "tasks.knowledge_fs_lifecycle_tasks.cleanup_knowledge_fs_staged_uploads", - "schedule": timedelta(seconds=2), - } - assert celery_app.conf["beat_schedule"]["knowledge_fs_upgrade_file_cleanup"] == { - "task": "tasks.knowledge_fs_upgrade_tasks.cleanup_deferred_knowledge_fs_upgrade_files", - "schedule": timedelta(seconds=2), - } - assert celery_app.conf["beat_schedule"]["knowledge_fs_upgrade_recovery"] == { - "task": "tasks.knowledge_fs_upgrade_tasks.recover_stalled_knowledge_fs_upgrades", - "schedule": timedelta(seconds=60), - "options": {"expires": 60}, - } - schedules = celery_app.conf["beat_schedule"] - if background_enabled: - assert celery_app.conf["broker_transport_options"]["visibility_timeout"] == 10800 - for operation, queue in OPERATION_QUEUES.items(): - assert schedules[f"knowledge_fs_background_{operation}"] == { - "task": OPERATION_TASK, - "schedule": timedelta(seconds=5), - "kwargs": {"operation": operation}, - "options": {"queue": queue, "expires": 30}, - } - else: - assert not any(name.startswith("knowledge_fs_background_") for name in schedules) - - with ( - patch( - "services.knowledge_fs.lifecycle_readiness.get_configured_knowledge_fs_lifecycle_worker_readiness", - return_value=KnowledgeFSLifecycleWorkerReadiness(ready=False, blockers=("test",)), - ), - ): - preview_only_app = init_app(DifyApp(f"{__name__}.preview_only")) - - assert "tasks.knowledge_fs_initial_source_preview_tasks" in preview_only_app.conf["imports"] - assert "tasks.knowledge_fs_failed_retrieval_tasks" in preview_only_app.conf["imports"] - assert "tasks.knowledge_fs_initial_source_tasks" not in preview_only_app.conf["imports"] - assert "tasks.knowledge_fs_lifecycle_tasks" not in preview_only_app.conf["imports"] - assert "tasks.knowledge_fs_source_import_tasks" not in preview_only_app.conf["imports"] - assert "knowledge_fs_upgrade_recovery" not in preview_only_app.conf["beat_schedule"] - - -@pytest.mark.parametrize("visibility", [None, 3600, 7200, 7230]) -def test_long_running_delivery_rejects_an_unsafe_redis_visibility(visibility: int | None) -> None: - with patch("extensions.ext_celery.dify_config") as config: - config.CELERY_USE_SENTINEL = False - config.REDIS_KEY_PREFIX = "" - config.CELERY_BROKER_URL = "redis://localhost:6379/0" - config.CELERY_BROKER_VISIBILITY_TIMEOUT = visibility - config.KNOWLEDGE_FS_BACKGROUND_WORKER_ENABLED = True - config.KNOWLEDGE_FS_BACKGROUND_TASK_TIMEOUT_SECONDS = 7200 - with pytest.raises(ValueError, match="VISIBILITY_TIMEOUT"): - get_celery_broker_transport_options() diff --git a/api/tests/unit_tests/extensions/test_pubsub_channel.py b/api/tests/unit_tests/extensions/test_pubsub_channel.py deleted file mode 100644 index c63e29c3cbd..00000000000 --- a/api/tests/unit_tests/extensions/test_pubsub_channel.py +++ /dev/null @@ -1,24 +0,0 @@ -import pytest - -from extensions import ext_redis -from libs.broadcast_channel.redis.pubsub_channel import BroadcastChannel as RedisBroadcastChannel -from libs.broadcast_channel.redis.sharded_channel import ShardedRedisBroadcastChannel -from tests.unit_tests.config_override import apply_config_overrides - - -def test_get_pubsub_broadcast_channel_defaults_to_pubsub(monkeypatch: pytest.MonkeyPatch): - apply_config_overrides(monkeypatch, PUBSUB_REDIS_CHANNEL_TYPE="pubsub") - monkeypatch.setattr(ext_redis, "_pubsub_redis_client", object()) - - channel = ext_redis.get_pubsub_broadcast_channel() - - assert isinstance(channel, RedisBroadcastChannel) - - -def test_get_pubsub_broadcast_channel_sharded(monkeypatch: pytest.MonkeyPatch): - apply_config_overrides(monkeypatch, PUBSUB_REDIS_CHANNEL_TYPE="sharded") - monkeypatch.setattr(ext_redis, "_pubsub_redis_client", object()) - - channel = ext_redis.get_pubsub_broadcast_channel() - - assert isinstance(channel, ShardedRedisBroadcastChannel) diff --git a/api/tests/unit_tests/extensions/test_redis.py b/api/tests/unit_tests/extensions/test_redis.py deleted file mode 100644 index bf2858f72ce..00000000000 --- a/api/tests/unit_tests/extensions/test_redis.py +++ /dev/null @@ -1,302 +0,0 @@ -from collections.abc import Callable, Mapping -from dataclasses import dataclass, field - -import pytest -from redis import RedisError -from redis.retry import Retry - -from extensions.ext_redis import ( - RedisClientWrapper, - _get_base_redis_params, - _get_cluster_connection_health_params, - _get_connection_health_params, - _normalize_redis_key_prefix, - _parse_redis_nodes, - _serialize_redis_name, - redis_fallback, -) - -type _RecordedCall = tuple[str, tuple[object, ...], dict[str, object]] - - -@dataclass -class _RecordingRedisClient: - calls: list[_RecordedCall] = field(default_factory=list) - - def _record(self, method: str, *args: object, **kwargs: object) -> None: - self.calls.append((method, args, kwargs)) - - def register_script(self, script: str) -> Callable[..., str]: - self._record("register_script", script) - - def execute(*, keys: object, args: object, client: object) -> str: - self._record("script", keys, args, client) - return "result" - - return execute - - def get(self, name: str | bytes) -> None: - self._record("get", name) - - def delete(self, *names: str | bytes) -> None: - self._record("delete", *names) - - def lock(self, name: str, **kwargs: object) -> None: - self._record("lock", name, **kwargs) - - def hset(self, name: str | bytes, *args: object, **kwargs: object) -> None: - self._record("hset", name, *args, **kwargs) - - def hgetall(self, name: str | bytes) -> None: - self._record("hgetall", name) - - def hkeys(self, name: str | bytes) -> None: - self._record("hkeys", name) - - def hexists(self, name: str | bytes, key: str | bytes) -> None: - self._record("hexists", name, key) - - def zadd(self, name: str | bytes, mapping: Mapping[object, object], **kwargs: object) -> None: - self._record("zadd", name, mapping, **kwargs) - - def zrem(self, name: str | bytes, *values: str | bytes) -> None: - self._record("zrem", name, *values) - - -@pytest.fixture(autouse=True) -def _redis_config(config_overrides: Callable[..., None]) -> None: - config_overrides( - REDIS_USERNAME=None, - REDIS_PASSWORD=None, - REDIS_DB=0, - REDIS_SERIALIZATION_PROTOCOL=3, - REDIS_ENABLE_CLIENT_SIDE_CACHE=False, - REDIS_RETRY_RETRIES=3, - REDIS_RETRY_BACKOFF_BASE=1.0, - REDIS_RETRY_BACKOFF_CAP=10.0, - REDIS_SOCKET_TIMEOUT=5.0, - REDIS_SOCKET_CONNECT_TIMEOUT=5.0, - REDIS_HEALTH_CHECK_INTERVAL=30, - REDIS_KEEPALIVE=True, - REDIS_KEEPALIVE_IDLE=60, - REDIS_KEEPALIVE_INTERVAL=10, - REDIS_KEEPALIVE_COUNT=3, - REDIS_KEY_PREFIX="", - ) - - -class TestGetConnectionHealthParams: - def test_includes_all_health_params(self) -> None: - params = _get_connection_health_params() - - assert "retry" in params - assert "socket_timeout" in params - assert "socket_connect_timeout" in params - assert "health_check_interval" in params - assert isinstance(params["retry"], Retry) - assert params["retry"]._retries == 3 # pyrefly: ignore[missing-attribute] - assert params["socket_timeout"] == 5.0 - assert params["socket_connect_timeout"] == 5.0 - assert params["health_check_interval"] == 30 - - -class TestGetClusterConnectionHealthParams: - def test_excludes_health_check_interval(self) -> None: - params = _get_cluster_connection_health_params() - - assert "retry" in params - assert "socket_timeout" in params - assert "socket_connect_timeout" in params - assert "health_check_interval" not in params - - -class TestGetBaseRedisParams: - def test_includes_retry_and_health_params(self) -> None: - params = _get_base_redis_params() - - assert "retry" in params - assert isinstance(params["retry"], Retry) - assert params["socket_timeout"] == 5.0 - assert params["socket_connect_timeout"] == 5.0 - assert params["health_check_interval"] == 30 - assert params["socket_keepalive"] is True - assert isinstance(params["socket_keepalive_options"], dict) - # Existing params still present - assert params["db"] == 0 - assert params["encoding"] == "utf-8" - - -class TestParseRedisNodes: - def test_trims_nodes(self) -> None: - assert _parse_redis_nodes("redis-a:6379, redis-b:6380") == [("redis-a", 6379), ("redis-b", 6380)] - - def test_supports_bracketed_ipv6(self) -> None: - assert _parse_redis_nodes("[2001:db8::10]:6379") == [("2001:db8::10", 6379)] - - -class TestRedisFallback: - def test_redis_fallback_success(self) -> None: - @redis_fallback(default_return=None) - def test_func() -> str: - return "success" - - assert test_func() == "success" - - def test_redis_fallback_error(self) -> None: - @redis_fallback(default_return="fallback") - def test_func() -> None: - raise RedisError("Redis error") - - assert test_func() == "fallback" - - def test_redis_fallback_none_default(self) -> None: - @redis_fallback() - def test_func() -> None: - raise RedisError("Redis error") - - assert test_func() is None - - def test_redis_fallback_with_args(self) -> None: - received: list[tuple[int, int]] = [] - - @redis_fallback(default_return=0) - def test_func(x: int, y: int) -> None: - received.append((x, y)) - raise RedisError("Redis error") - - assert test_func(1, 2) == 0 - assert received == [(1, 2)] - - def test_redis_fallback_with_kwargs(self) -> None: - received: list[tuple[int | None, int | None]] = [] - - @redis_fallback(default_return=dict[str, object]()) - def test_func(x: int | None = None, y: int | None = None) -> None: - received.append((x, y)) - raise RedisError("Redis error") - - assert test_func(x=1, y=2) == {} - assert received == [(1, 2)] - - def test_redis_fallback_preserves_function_metadata(self) -> None: - @redis_fallback(default_return=None) - def test_func() -> None: - """Test function docstring""" - pass - - assert test_func.__name__ == "test_func" - assert test_func.__doc__ == "Test function docstring" - - -class TestRedisKeyPrefixHelpers: - def test_normalize_redis_key_prefix_trims_whitespace(self) -> None: - assert _normalize_redis_key_prefix(" enterprise-a ") == "enterprise-a" - - def test_normalize_redis_key_prefix_treats_whitespace_only_as_empty(self) -> None: - assert _normalize_redis_key_prefix(" ") == "" - - def test_serialize_redis_name_returns_original_when_prefix_empty(self) -> None: - assert _serialize_redis_name("model_lb_index:test", "") == "model_lb_index:test" - - def test_serialize_redis_name_adds_single_colon_separator(self) -> None: - assert _serialize_redis_name("model_lb_index:test", "enterprise-a") == "enterprise-a:model_lb_index:test" - - -class TestRedisClientWrapperKeyPrefix: - def test_wrapper_zrem_prefixes_key_but_not_members(self, config_overrides: Callable[..., None]) -> None: - client = _RecordingRedisClient() - wrapper = RedisClientWrapper() - wrapper.initialize(client) # type: ignore[arg-type] - config_overrides(REDIS_KEY_PREFIX="enterprise-a") - - wrapper.zrem("reset:user@example.com", "reservation") - - assert client.calls == [("zrem", ("enterprise-a:reset:user@example.com", "reservation"), {})] - - def test_wrapper_registered_script_prefixes_key_arguments(self, config_overrides: Callable[..., None]) -> None: - raw_client = _RecordingRedisClient() - wrapper = RedisClientWrapper() - wrapper.initialize(raw_client) # type: ignore[arg-type] - - config_overrides(REDIS_KEY_PREFIX="enterprise-a") - script = wrapper.register_script("return redis.call('GET', KEYS[1])") - - assert script(keys=["device_code:abc"], args=["argument"]) == "result" - assert raw_client.calls == [ - ("register_script", ("return redis.call('GET', KEYS[1])",), {}), - ("script", (("enterprise-a:device_code:abc",), ["argument"], raw_client), {}), - ] - - def test_wrapper_get_prefixes_string_keys(self, config_overrides: Callable[..., None]) -> None: - client = _RecordingRedisClient() - wrapper = RedisClientWrapper() - wrapper.initialize(client) # type: ignore[arg-type] - - config_overrides(REDIS_KEY_PREFIX="enterprise-a") - wrapper.get("oauth_state:abc") - - assert client.calls == [("get", ("enterprise-a:oauth_state:abc",), {})] - - def test_wrapper_delete_prefixes_multiple_keys(self, config_overrides: Callable[..., None]) -> None: - client = _RecordingRedisClient() - wrapper = RedisClientWrapper() - wrapper.initialize(client) # type: ignore[arg-type] - - config_overrides(REDIS_KEY_PREFIX="enterprise-a") - wrapper.delete("key:a", "key:b") - - assert client.calls == [("delete", ("enterprise-a:key:a", "enterprise-a:key:b"), {})] - - def test_wrapper_lock_prefixes_lock_name(self, config_overrides: Callable[..., None]) -> None: - client = _RecordingRedisClient() - wrapper = RedisClientWrapper() - wrapper.initialize(client) # type: ignore[arg-type] - - config_overrides(REDIS_KEY_PREFIX="enterprise-a") - wrapper.lock("resource-lock", timeout=10) - - method, args, kwargs = client.calls[0] - assert method == "lock" - assert args == ("enterprise-a:resource-lock",) - assert kwargs["timeout"] == 10 - - def test_wrapper_hash_operations_prefix_key_name(self, config_overrides: Callable[..., None]) -> None: - client = _RecordingRedisClient() - wrapper = RedisClientWrapper() - wrapper.initialize(client) # type: ignore[arg-type] - - config_overrides(REDIS_KEY_PREFIX="enterprise-a") - wrapper.hset("hash:key", "field", "value") - wrapper.hgetall("hash:key") - wrapper.hkeys("hash:key") - wrapper.hexists("hash:key", "field") - - assert client.calls == [ - ("hset", ("enterprise-a:hash:key", "field", "value"), {}), - ("hgetall", ("enterprise-a:hash:key",), {}), - ("hkeys", ("enterprise-a:hash:key",), {}), - ("hexists", ("enterprise-a:hash:key", "field"), {}), - ] - - def test_wrapper_zadd_prefixes_sorted_set_name(self, config_overrides: Callable[..., None]) -> None: - client = _RecordingRedisClient() - wrapper = RedisClientWrapper() - wrapper.initialize(client) # type: ignore[arg-type] - - config_overrides(REDIS_KEY_PREFIX="enterprise-a") - wrapper.zadd("zset:key", {"member": 1}) - - method, args, kwargs = client.calls[0] - assert method == "zadd" - assert args == ("enterprise-a:zset:key", {"member": 1}) - assert kwargs["nx"] is False - - def test_wrapper_preserves_keys_when_prefix_is_empty(self, config_overrides: Callable[..., None]) -> None: - client = _RecordingRedisClient() - wrapper = RedisClientWrapper() - wrapper.initialize(client) # type: ignore[arg-type] - - config_overrides(REDIS_KEY_PREFIX=" ") - wrapper.get("plain:key") - - assert client.calls == [("get", ("plain:key",), {})] diff --git a/api/tests/unit_tests/factories/test_build_from_mapping.py b/api/tests/unit_tests/factories/file_factory/test_builders.py similarity index 100% rename from api/tests/unit_tests/factories/test_build_from_mapping.py rename to api/tests/unit_tests/factories/file_factory/test_builders.py diff --git a/api/tests/unit_tests/factories/test_file_factory.py b/api/tests/unit_tests/factories/file_factory/test_remote.py similarity index 100% rename from api/tests/unit_tests/factories/test_file_factory.py rename to api/tests/unit_tests/factories/file_factory/test_remote.py diff --git a/api/tests/unit_tests/factories/test_file_validation.py b/api/tests/unit_tests/factories/file_factory/test_validation.py similarity index 100% rename from api/tests/unit_tests/factories/test_file_validation.py rename to api/tests/unit_tests/factories/file_factory/test_validation.py diff --git a/api/tests/unit_tests/fields/test_document_fields.py b/api/tests/unit_tests/fields/test_document_fields.py index 7f4e82a7583..db4b50ce12d 100644 --- a/api/tests/unit_tests/fields/test_document_fields.py +++ b/api/tests/unit_tests/fields/test_document_fields.py @@ -4,7 +4,7 @@ from datetime import datetime from sqlalchemy.orm import Session from extensions.storage.storage_type import StorageType -from fields.document_fields import DocumentResponse +from fields.document_fields import DocumentProcessRuleResponse, DocumentResponse, DocumentSourceDetailResponse from models.dataset import Document, DocumentSegment from models.enums import CreatorUserRole, DataSourceType, DocumentCreatedFrom from models.model import UploadFile @@ -59,3 +59,31 @@ def test_document_detail_serializes_after_session_closes(sqlite_session: Session assert response["data_source_detail_dict"]["upload_file"]["name"] == "source.txt" assert response["hit_count"] == 3 assert response["doc_metadata"] == [] + + +def test_process_rule_reads_the_legacy_delimiter_as_separator() -> None: + rule = DocumentProcessRuleResponse.model_validate( + { + "mode": "custom", + "rules": { + "pre_processing_rules": [{"id": "remove_extra_spaces", "enabled": True}], + "segmentation": {"delimiter": "\n\n", "max_tokens": 500}, + }, + } + ) + + assert rule.model_dump(mode="json") == { + "mode": "custom", + "rules": { + "pre_processing_rules": [{"id": "remove_extra_spaces", "enabled": True}], + "segmentation": {"separator": "\n\n", "max_tokens": 500, "chunk_overlap": 0}, + "parent_mode": None, + "subchunk_segmentation": None, + }, + } + + +def test_source_detail_keeps_notion_and_website_source_info() -> None: + detail = DocumentSourceDetailResponse.model_validate({"url": "https://example.com", "title": "Example"}) + + assert detail.model_dump(mode="json") == {"url": "https://example.com", "title": "Example"} diff --git a/api/tests/unit_tests/fields/test_document_process_rule_fields.py b/api/tests/unit_tests/fields/test_document_process_rule_fields.py deleted file mode 100644 index 39a9bef5a77..00000000000 --- a/api/tests/unit_tests/fields/test_document_process_rule_fields.py +++ /dev/null @@ -1,29 +0,0 @@ -from fields.document_fields import DocumentProcessRuleResponse, DocumentSourceDetailResponse - - -def test_process_rule_reads_the_legacy_delimiter_as_separator() -> None: - rule = DocumentProcessRuleResponse.model_validate( - { - "mode": "custom", - "rules": { - "pre_processing_rules": [{"id": "remove_extra_spaces", "enabled": True}], - "segmentation": {"delimiter": "\n\n", "max_tokens": 500}, - }, - } - ) - - assert rule.model_dump(mode="json") == { - "mode": "custom", - "rules": { - "pre_processing_rules": [{"id": "remove_extra_spaces", "enabled": True}], - "segmentation": {"separator": "\n\n", "max_tokens": 500, "chunk_overlap": 0}, - "parent_mode": None, - "subchunk_segmentation": None, - }, - } - - -def test_source_detail_keeps_notion_and_website_source_info() -> None: - detail = DocumentSourceDetailResponse.model_validate({"url": "https://example.com", "title": "Example"}) - - assert detail.model_dump(mode="json") == {"url": "https://example.com", "title": "Example"} diff --git a/api/tests/unit_tests/controllers/console/app/workflow_draft_variables_test.py b/api/tests/unit_tests/fields/test_workflow_draft_variable_fields.py similarity index 96% rename from api/tests/unit_tests/controllers/console/app/workflow_draft_variables_test.py rename to api/tests/unit_tests/fields/test_workflow_draft_variable_fields.py index 0c62e2b4c00..aaf556dd784 100644 --- a/api/tests/unit_tests/controllers/console/app/workflow_draft_variables_test.py +++ b/api/tests/unit_tests/fields/test_workflow_draft_variable_fields.py @@ -20,6 +20,12 @@ from libs.helper import dump_response from libs.uuid_utils import uuidv7 from models.workflow import WorkflowDraftVariable, WorkflowDraftVariableFile from services.workflow_draft_variable_service import WorkflowDraftVariableList +from tests.unit_tests.controllers.conftest import app_query_services as app_query_services # noqa: PLC0414 +from tests.unit_tests.controllers.conftest import ( + datasource_application_dependencies as datasource_application_dependencies, # noqa: PLC0414 +) +from tests.unit_tests.controllers.conftest import error_response_contract as error_response_contract # noqa: PLC0414 +from tests.unit_tests.controllers.conftest import pipeline_application as pipeline_application # noqa: PLC0414 _TEST_APP_ID = "test_app_id" _TEST_NODE_EXEC_ID = str(uuid.uuid4()) diff --git a/api/tests/unit_tests/libs/broadcast_channel/redis/test_channel_unit_tests.py b/api/tests/unit_tests/libs/broadcast_channel/redis/test_pubsub_channel.py similarity index 100% rename from api/tests/unit_tests/libs/broadcast_channel/redis/test_channel_unit_tests.py rename to api/tests/unit_tests/libs/broadcast_channel/redis/test_pubsub_channel.py diff --git a/api/tests/unit_tests/libs/broadcast_channel/redis/test_streams_channel_unit_tests.py b/api/tests/unit_tests/libs/broadcast_channel/redis/test_streams_channel.py similarity index 100% rename from api/tests/unit_tests/libs/broadcast_channel/redis/test_streams_channel_unit_tests.py rename to api/tests/unit_tests/libs/broadcast_channel/redis/test_streams_channel.py diff --git a/api/tests/unit_tests/libs/test_api_token_cache.py b/api/tests/unit_tests/libs/test_api_token_cache.py deleted file mode 100644 index 0be39376665..00000000000 --- a/api/tests/unit_tests/libs/test_api_token_cache.py +++ /dev/null @@ -1,252 +0,0 @@ -""" -Unit tests for API Token Cache module. -""" - -import json -from datetime import datetime -from unittest.mock import MagicMock, patch - -from models.enums import ApiTokenType -from models.model import ApiToken -from services.api_token_service import ( - CACHE_KEY_PREFIX, - CACHE_NULL_TTL_SECONDS, - CACHE_TTL_SECONDS, - ApiTokenCache, - CachedApiToken, -) - - -class TestApiTokenCache: - """Test cases for ApiTokenCache class.""" - - def setup_method(self) -> None: - """Setup test fixtures.""" - self.mock_token = ApiToken() - self.mock_token.id = "test-token-id-123" - self.mock_token.app_id = "test-app-id-456" - self.mock_token.tenant_id = "test-tenant-id-789" - self.mock_token.type = ApiTokenType.APP - self.mock_token.token = "test-token-value-abc" - self.mock_token.last_used_at = datetime(2026, 2, 3, 10, 0, 0) - self.mock_token.created_at = datetime(2026, 1, 1, 0, 0, 0) - - def test_make_cache_key(self) -> None: - """Test cache key generation.""" - # Test with scope - key = ApiTokenCache._make_cache_key("my-token", "app") - assert key == f"{CACHE_KEY_PREFIX}:app:my-token" - - # Test without scope - key = ApiTokenCache._make_cache_key("my-token", None) - assert key == f"{CACHE_KEY_PREFIX}:any:my-token" - - def test_serialize_token(self) -> None: - """Test token serialization.""" - serialized = ApiTokenCache._serialize_token(self.mock_token) - data = json.loads(serialized) - - assert data["id"] == "test-token-id-123" - assert data["app_id"] == "test-app-id-456" - assert data["tenant_id"] == "test-tenant-id-789" - assert data["type"] == "app" - assert data["token"] == "test-token-value-abc" - assert data["last_used_at"] == "2026-02-03T10:00:00" - assert data["created_at"] == "2026-01-01T00:00:00" - - def test_serialize_token_with_nulls(self) -> None: - """Test token serialization with None values.""" - mock_token = ApiToken() - mock_token.id = "test-id" - mock_token.app_id = None - mock_token.tenant_id = None - mock_token.type = ApiTokenType.DATASET - mock_token.token = "test-token" - mock_token.last_used_at = None - mock_token.created_at = datetime(2026, 1, 1, 0, 0, 0) - - serialized = ApiTokenCache._serialize_token(mock_token) - data = json.loads(serialized) - - assert data["app_id"] is None - assert data["tenant_id"] is None - assert data["last_used_at"] is None - - def test_deserialize_token(self) -> None: - """Test token deserialization.""" - cached_data = json.dumps( - { - "id": "test-id", - "app_id": "test-app", - "tenant_id": "test-tenant", - "type": "app", - "token": "test-token", - "last_used_at": "2026-02-03T10:00:00", - "created_at": "2026-01-01T00:00:00", - } - ) - - result = ApiTokenCache._deserialize_token(cached_data) - - assert isinstance(result, CachedApiToken) - assert result.id == "test-id" - assert result.app_id == "test-app" - assert result.tenant_id == "test-tenant" - assert result.type == "app" - assert result.token == "test-token" - assert result.last_used_at == datetime(2026, 2, 3, 10, 0, 0) - assert result.created_at == datetime(2026, 1, 1, 0, 0, 0) - - def test_deserialize_null_token(self) -> None: - """Test deserialization of null token (cached miss).""" - result = ApiTokenCache._deserialize_token("null") - assert result is None - - def test_deserialize_invalid_json(self) -> None: - """Test deserialization with invalid JSON.""" - result = ApiTokenCache._deserialize_token("invalid-json{") - assert result is None - - @patch("services.api_token_service.redis_client") - def test_get_cache_hit(self, mock_redis: MagicMock) -> None: - """Test cache hit scenario.""" - cached_data = json.dumps( - { - "id": "test-id", - "app_id": "test-app", - "tenant_id": "test-tenant", - "type": "app", - "token": "test-token", - "last_used_at": "2026-02-03T10:00:00", - "created_at": "2026-01-01T00:00:00", - } - ).encode("utf-8") - mock_redis.get.return_value = cached_data - - result = ApiTokenCache.get("test-token", "app") - - assert result is not None - assert isinstance(result, CachedApiToken) - assert result.app_id == "test-app" - mock_redis.get.assert_called_once_with(f"{CACHE_KEY_PREFIX}:app:test-token") - - @patch("services.api_token_service.redis_client") - def test_get_cache_miss(self, mock_redis: MagicMock) -> None: - """Test cache miss scenario.""" - mock_redis.get.return_value = None - - result = ApiTokenCache.get("test-token", "app") - - assert result is None - mock_redis.get.assert_called_once() - - @patch("services.api_token_service.redis_client") - def test_set_valid_token(self, mock_redis: MagicMock) -> None: - """Test setting a valid token in cache.""" - result = ApiTokenCache.set("test-token", "app", self.mock_token) - - assert result is True - mock_redis.setex.assert_called_once() - args = mock_redis.setex.call_args[0] - assert args[0] == f"{CACHE_KEY_PREFIX}:app:test-token" - assert args[1] == CACHE_TTL_SECONDS - - @patch("services.api_token_service.redis_client") - def test_set_null_token(self, mock_redis: MagicMock) -> None: - """Test setting a null token (cache penetration prevention).""" - result = ApiTokenCache.set("invalid-token", "app", None) - - assert result is True - mock_redis.setex.assert_called_once() - args = mock_redis.setex.call_args[0] - assert args[0] == f"{CACHE_KEY_PREFIX}:app:invalid-token" - assert args[1] == CACHE_NULL_TTL_SECONDS - assert args[2] == b"null" - - @patch("services.api_token_service.redis_client") - def test_delete_with_scope(self, mock_redis: MagicMock) -> None: - """Test deleting token cache with specific scope.""" - result = ApiTokenCache.delete("test-token", "app") - - assert result is True - mock_redis.delete.assert_called_once_with(f"{CACHE_KEY_PREFIX}:app:test-token") - - @patch("services.api_token_service.redis_client") - def test_delete_without_scope(self, mock_redis: MagicMock) -> None: - """Test deleting token cache without scope (delete all).""" - # Mock scan_iter to return an iterator of keys - mock_redis.scan_iter.return_value = iter( - [ - b"api_token:app:test-token", - b"api_token:dataset:test-token", - ] - ) - - result = ApiTokenCache.delete("test-token", None) - - assert result is True - # Verify scan_iter was called with the correct pattern - mock_redis.scan_iter.assert_called_once() - call_args = mock_redis.scan_iter.call_args - assert call_args[1]["match"] == f"{CACHE_KEY_PREFIX}:*:test-token" - - # Verify delete was called with all matched keys - mock_redis.delete.assert_called_once_with( - b"api_token:app:test-token", - b"api_token:dataset:test-token", - ) - - @patch("services.api_token_service.redis_client") - def test_redis_fallback_on_exception(self, mock_redis: MagicMock) -> None: - """Test Redis fallback when Redis is unavailable.""" - from redis import RedisError - - mock_redis.get.side_effect = RedisError("Connection failed") - - result = ApiTokenCache.get("test-token", "app") - - # Should return None (fallback) instead of raising exception - assert result is None - - -class TestApiTokenCacheIntegration: - """Integration test scenarios.""" - - @patch("services.api_token_service.redis_client") - def test_full_cache_lifecycle(self, mock_redis: MagicMock) -> None: - """Test complete cache lifecycle: set -> get -> delete.""" - # Use the persisted model's fields for the cache lifecycle. - mock_token = ApiToken() - mock_token.id = "id-123" - mock_token.app_id = "app-456" - mock_token.tenant_id = "tenant-789" - mock_token.type = ApiTokenType.APP - mock_token.token = "token-abc" - mock_token.last_used_at = datetime(2026, 2, 3, 10, 0, 0) - mock_token.created_at = datetime(2026, 1, 1, 0, 0, 0) - - # 1. Set token in cache - ApiTokenCache.set("token-abc", "app", mock_token) - assert mock_redis.setex.called - - # 2. Simulate cache hit - cached_data = ApiTokenCache._serialize_token(mock_token) - mock_redis.get.return_value = cached_data # bytes from model_dump_json().encode() - - retrieved = ApiTokenCache.get("token-abc", "app") - assert retrieved is not None - assert isinstance(retrieved, CachedApiToken) - - # 3. Delete from cache - ApiTokenCache.delete("token-abc", "app") - assert mock_redis.delete.called - - @patch("services.api_token_service.redis_client") - def test_cache_penetration_prevention(self, mock_redis: MagicMock) -> None: - """Test that non-existent tokens are cached as null.""" - # Set null token (cache miss) - ApiTokenCache.set("non-existent-token", "app", None) - - args = mock_redis.setex.call_args[0] - assert args[2] == b"null" - assert args[1] == CACHE_NULL_TTL_SECONDS # Shorter TTL for null values diff --git a/api/tests/unit_tests/libs/test_archive_storage.py b/api/tests/unit_tests/libs/test_archive_storage.py index b49a762992c..c1cfa0868bf 100644 --- a/api/tests/unit_tests/libs/test_archive_storage.py +++ b/api/tests/unit_tests/libs/test_archive_storage.py @@ -14,7 +14,7 @@ from libs.archive_storage import ( ArchiveStorageError, ArchiveStorageNotConfiguredError, ) -from tests.unit_tests.config_override import apply_config_overrides +from tests.helpers.unit.config_override import apply_config_overrides BUCKET_NAME = "archive-bucket" diff --git a/api/tests/unit_tests/libs/test_cron_compatibility.py b/api/tests/unit_tests/libs/test_cron_compatibility.py deleted file mode 100644 index 6f3a94f6dc4..00000000000 --- a/api/tests/unit_tests/libs/test_cron_compatibility.py +++ /dev/null @@ -1,381 +0,0 @@ -""" -Enhanced cron syntax compatibility tests for croniter backend. - -This test suite mirrors the frontend cron-parser tests to ensure -complete compatibility between frontend and backend cron processing. -""" - -import unittest -from datetime import UTC, datetime, timedelta - -import pytest -import pytz -from croniter import CroniterBadCronError - -from libs.schedule_utils import calculate_next_run_at - - -class TestCronCompatibility(unittest.TestCase): - """Test enhanced cron syntax compatibility with frontend.""" - - def setUp(self): - """Set up test environment with fixed time.""" - self.base_time = datetime(2024, 1, 15, 10, 0, 0, tzinfo=UTC) - - def test_enhanced_dayofweek_syntax(self): - """Test enhanced day-of-week syntax compatibility.""" - test_cases = [ - ("0 9 * * 7", 0), # Sunday as 7 - ("0 9 * * 0", 0), # Sunday as 0 - ("0 9 * * MON", 1), # Monday abbreviation - ("0 9 * * TUE", 2), # Tuesday abbreviation - ("0 9 * * WED", 3), # Wednesday abbreviation - ("0 9 * * THU", 4), # Thursday abbreviation - ("0 9 * * FRI", 5), # Friday abbreviation - ("0 9 * * SAT", 6), # Saturday abbreviation - ("0 9 * * SUN", 0), # Sunday abbreviation - ] - - for expr, expected_weekday in test_cases: - with self.subTest(expr=expr): - next_time = calculate_next_run_at(expr, "UTC", self.base_time) - assert next_time is not None - assert (next_time.weekday() + 1 if next_time.weekday() < 6 else 0) == expected_weekday - assert next_time.hour == 9 - assert next_time.minute == 0 - - def test_enhanced_month_syntax(self): - """Test enhanced month syntax compatibility.""" - test_cases = [ - ("0 9 1 JAN *", 1), # January abbreviation - ("0 9 1 FEB *", 2), # February abbreviation - ("0 9 1 MAR *", 3), # March abbreviation - ("0 9 1 APR *", 4), # April abbreviation - ("0 9 1 MAY *", 5), # May abbreviation - ("0 9 1 JUN *", 6), # June abbreviation - ("0 9 1 JUL *", 7), # July abbreviation - ("0 9 1 AUG *", 8), # August abbreviation - ("0 9 1 SEP *", 9), # September abbreviation - ("0 9 1 OCT *", 10), # October abbreviation - ("0 9 1 NOV *", 11), # November abbreviation - ("0 9 1 DEC *", 12), # December abbreviation - ] - - for expr, expected_month in test_cases: - with self.subTest(expr=expr): - next_time = calculate_next_run_at(expr, "UTC", self.base_time) - assert next_time is not None - assert next_time.month == expected_month - assert next_time.day == 1 - assert next_time.hour == 9 - - def test_predefined_expressions(self): - """Test predefined cron expressions compatibility.""" - test_cases = [ - ("@yearly", lambda dt: dt.month == 1 and dt.day == 1 and dt.hour == 0), - ("@annually", lambda dt: dt.month == 1 and dt.day == 1 and dt.hour == 0), - ("@monthly", lambda dt: dt.day == 1 and dt.hour == 0), - ("@weekly", lambda dt: dt.weekday() == 6 and dt.hour == 0), # Sunday = 6 in weekday() - ("@daily", lambda dt: dt.hour == 0 and dt.minute == 0), - ("@midnight", lambda dt: dt.hour == 0 and dt.minute == 0), - ("@hourly", lambda dt: dt.minute == 0), - ] - - for expr, validator in test_cases: - with self.subTest(expr=expr): - next_time = calculate_next_run_at(expr, "UTC", self.base_time) - assert next_time is not None - assert validator(next_time), f"Validator failed for {expr}: {next_time}" - - def test_special_characters(self): - """Test special characters in cron expressions.""" - test_cases = [ - "0 9 ? * 1", # ? wildcard - "0 12 * * 7", # Sunday as 7 - "0 15 L * *", # Last day of month - ] - - for expr in test_cases: - with self.subTest(expr=expr): - try: - next_time = calculate_next_run_at(expr, "UTC", self.base_time) - assert next_time is not None - assert next_time > self.base_time - except Exception as e: - self.fail(f"Expression '{expr}' should be valid but raised: {e}") - - def test_range_and_list_syntax(self): - """Test range and list syntax with abbreviations.""" - test_cases = [ - "0 9 * * MON-FRI", # Weekday range with abbreviations - "0 9 * JAN-MAR *", # Month range with abbreviations - "0 9 * * SUN,WED,FRI", # Weekday list with abbreviations - "0 9 1 JAN,JUN,DEC *", # Month list with abbreviations - ] - - for expr in test_cases: - with self.subTest(expr=expr): - try: - next_time = calculate_next_run_at(expr, "UTC", self.base_time) - assert next_time is not None - assert next_time > self.base_time - except Exception as e: - self.fail(f"Expression '{expr}' should be valid but raised: {e}") - - def test_invalid_enhanced_syntax(self): - """Test that invalid enhanced syntax is properly rejected.""" - invalid_expressions = [ - "0 12 * JANUARY *", # Full month name (not supported) - "0 12 * * MONDAY", # Full day name (not supported) - "0 12 32 JAN *", # Invalid day with valid month - "15 10 1 * 8", # Invalid day of week - "15 10 1 INVALID *", # Invalid month abbreviation - "15 10 1 * INVALID", # Invalid day abbreviation - "@invalid", # Invalid predefined expression - ] - - for expr in invalid_expressions: - with self.subTest(expr=expr): - with pytest.raises((CroniterBadCronError, ValueError)): - calculate_next_run_at(expr, "UTC", self.base_time) - - def test_edge_cases_with_enhanced_syntax(self): - """Test edge cases with enhanced syntax.""" - test_cases = [ - ("0 0 29 FEB *", lambda dt: dt.month == 2 and dt.day == 29), # Feb 29 with month abbreviation - ] - - for expr, validator in test_cases: - with self.subTest(expr=expr): - try: - next_time = calculate_next_run_at(expr, "UTC", self.base_time) - if next_time: # Some combinations might not occur soon - assert validator(next_time), f"Validator failed for {expr}: {next_time}" - except (CroniterBadCronError, ValueError): - # Some edge cases might be valid but not have upcoming occurrences - pass - - # Test complex expressions that have specific constraints - complex_expr = "59 23 31 DEC SAT" # December 31st at 23:59 on Saturday - try: - next_time = calculate_next_run_at(complex_expr, "UTC", self.base_time) - if next_time: - # The next occurrence might not be exactly Dec 31 if it's not a Saturday - # Just verify it's a valid result - assert next_time is not None - assert next_time.hour == 23 - assert next_time.minute == 59 - except Exception: - # Complex date constraints might not have near-future occurrences - pass - - -class TestTimezoneCompatibility(unittest.TestCase): - """Test timezone compatibility between frontend and backend.""" - - def setUp(self): - """Set up test environment.""" - self.base_time = datetime(2024, 1, 15, 10, 0, 0, tzinfo=UTC) - - def test_timezone_consistency(self): - """Test that calculations are consistent across different timezones.""" - timezones = [ - "UTC", - "America/New_York", - "Europe/London", - "Asia/Tokyo", - "Asia/Kolkata", - "Australia/Sydney", - ] - - expression = "0 12 * * *" # Daily at noon - - for timezone in timezones: - with self.subTest(timezone=timezone): - next_time = calculate_next_run_at(expression, timezone, self.base_time) - assert next_time is not None - - # Convert back to the target timezone to verify it's noon - tz = pytz.timezone(timezone) - local_time = next_time.astimezone(tz) - assert local_time.hour == 12 - assert local_time.minute == 0 - - def test_dst_handling(self): - """Test DST boundary handling.""" - # Test around DST spring forward (March 2024) - dst_base = datetime(2024, 3, 8, 10, 0, 0, tzinfo=UTC) - expression = "0 2 * * *" # 2 AM daily (problematic during DST) - timezone = "America/New_York" - - try: - next_time = calculate_next_run_at(expression, timezone, dst_base) - assert next_time is not None - - # During DST spring forward, 2 AM becomes 3 AM - both are acceptable - tz = pytz.timezone(timezone) - local_time = next_time.astimezone(tz) - assert local_time.hour in [2, 3] # Either 2 AM or 3 AM is acceptable - except Exception as e: - self.fail(f"DST handling failed: {e}") - - def test_half_hour_timezones(self): - """Test timezones with half-hour offsets.""" - timezones_with_offsets = [ - ("Asia/Kolkata", 17, 30), # UTC+5:30 -> 12:00 UTC = 17:30 IST - ("Australia/Adelaide", 22, 30), # UTC+10:30 -> 12:00 UTC = 22:30 ACDT (summer time) - ] - - expression = "0 12 * * *" # Noon UTC - - for timezone, expected_hour, expected_minute in timezones_with_offsets: - with self.subTest(timezone=timezone): - try: - next_time = calculate_next_run_at(expression, timezone, self.base_time) - assert next_time is not None - - tz = pytz.timezone(timezone) - local_time = next_time.astimezone(tz) - assert local_time.hour == expected_hour - assert local_time.minute == expected_minute - except Exception: - # Some complex timezone calculations might vary - pass - - def test_invalid_timezone_handling(self): - """Test handling of invalid timezones.""" - expression = "0 12 * * *" - invalid_timezone = "Invalid/Timezone" - - with pytest.raises((ValueError, Exception)): # Should raise an exception - calculate_next_run_at(expression, invalid_timezone, self.base_time) - - -class TestFrontendBackendIntegration(unittest.TestCase): - """Test integration patterns that mirror frontend usage.""" - - def setUp(self): - """Set up test environment.""" - self.base_time = datetime(2024, 1, 15, 10, 0, 0, tzinfo=UTC) - - def test_execution_time_calculator_pattern(self): - """Test the pattern used by execution-time-calculator.ts.""" - # This mirrors the exact usage from execution-time-calculator.ts:47 - test_data = { - "cron_expression": "30 14 * * 1-5", # 2:30 PM weekdays - "timezone": "America/New_York", - } - - # Get next 5 execution times (like the frontend does) - execution_times = [] - current_base = self.base_time - - for _ in range(5): - next_time = calculate_next_run_at(test_data["cron_expression"], test_data["timezone"], current_base) - assert next_time is not None - execution_times.append(next_time) - current_base = next_time + timedelta(seconds=1) # Move slightly forward - - assert len(execution_times) == 5 - - # Validate each execution time - for exec_time in execution_times: - # Convert to local timezone - tz = pytz.timezone(test_data["timezone"]) - local_time = exec_time.astimezone(tz) - - # Should be weekdays (1-5) - assert local_time.weekday() in [0, 1, 2, 3, 4] # Mon-Fri in Python weekday - - # Should be 2:30 PM in local time - assert local_time.hour == 14 - assert local_time.minute == 30 - assert local_time.second == 0 - - def test_schedule_service_integration(self): - """Test integration with ScheduleService patterns.""" - from core.workflow.nodes.trigger_schedule.entities import VisualConfig - from services.trigger.schedule_service import ScheduleService - - # Test enhanced syntax through visual config conversion - visual_configs = [ - # Test with month abbreviations - { - "frequency": "monthly", - "config": VisualConfig(time="9:00 AM", monthly_days=[1]), - "expected_cron": "0 9 1 * *", - }, - # Test with weekday abbreviations - { - "frequency": "weekly", - "config": VisualConfig(time="2:30 PM", weekdays=["mon", "wed", "fri"]), - "expected_cron": "30 14 * * 1,3,5", - }, - ] - - for test_case in visual_configs: - with self.subTest(frequency=test_case["frequency"]): - cron_expr = ScheduleService.visual_to_cron(test_case["frequency"], test_case["config"]) - assert cron_expr == test_case["expected_cron"] - - # Verify the generated cron expression is valid - next_time = calculate_next_run_at(cron_expr, "UTC", self.base_time) - assert next_time is not None - - def test_error_handling_consistency(self): - """Test that error handling matches frontend expectations.""" - invalid_expressions = [ - "60 10 1 * *", # Invalid minute - "15 25 1 * *", # Invalid hour - "15 10 32 * *", # Invalid day - "15 10 1 13 *", # Invalid month - "15 10 1", # Too few fields - "15 10 1 * * *", # 6 fields (not supported in frontend) - "0 15 10 1 * * *", # 7 fields (not supported in frontend) - "invalid expression", # Completely invalid - ] - - for expr in invalid_expressions: - with self.subTest(expr=repr(expr)): - with pytest.raises((CroniterBadCronError, ValueError, Exception)): - calculate_next_run_at(expr, "UTC", self.base_time) - - # Note: Empty/whitespace expressions are not tested here as they are - # not expected in normal usage due to database constraints (nullable=False) - - def test_performance_requirements(self): - """Test that complex expressions parse within reasonable time.""" - import time - - complex_expressions = [ - "*/5 9-17 * * 1-5", # Every 5 minutes, weekdays, business hours - "0 */2 1,15 * *", # Every 2 hours on 1st and 15th - "30 14 * * 1,3,5", # Mon, Wed, Fri at 14:30 - "15,45 8-18 * * 1-5", # 15 and 45 minutes past hour, weekdays - "0 9 * JAN-MAR MON-FRI", # Enhanced syntax: Q1 weekdays at 9 AM - "0 12 ? * SUN", # Enhanced syntax: Sundays at noon with ? - ] - - start_time = time.time() - - for expr in complex_expressions: - with self.subTest(expr=expr): - try: - next_time = calculate_next_run_at(expr, "UTC", self.base_time) - assert next_time is not None - except CroniterBadCronError: - # Some enhanced syntax might not be supported, that's OK - pass - - end_time = time.time() - execution_time = (end_time - start_time) * 1000 # Convert to milliseconds - - # Should complete within reasonable time (less than 150ms like frontend) - assert execution_time < 150, "Complex expressions should parse quickly" - - -if __name__ == "__main__": - # Import timedelta for the test - from datetime import timedelta - - unittest.main() diff --git a/api/tests/unit_tests/libs/test_email.py b/api/tests/unit_tests/libs/test_email.py deleted file mode 100644 index 7d19206f6d6..00000000000 --- a/api/tests/unit_tests/libs/test_email.py +++ /dev/null @@ -1,21 +0,0 @@ -import pytest - -from libs.helper import email - - -def test_email_with_valid_email() -> None: - assert email("test@example.com") == "test@example.com" - assert email("TEST12345@example.com") == "TEST12345@example.com" - assert email("test+test@example.com") == "test+test@example.com" - assert email("!#$%&'*+-/=?^_{|}~`@example.com") == "!#$%&'*+-/=?^_{|}~`@example.com" - - -def test_email_with_invalid_email() -> None: - with pytest.raises(ValueError, match="invalid_email is not a valid email."): - email("invalid_email") - - with pytest.raises(ValueError, match="@example.com is not a valid email."): - email("@example.com") - - with pytest.raises(ValueError, match="()@example.com is not a valid email."): - email("()@example.com") diff --git a/api/tests/unit_tests/libs/test_helper.py b/api/tests/unit_tests/libs/test_helper.py index d8287a203bd..579f0475192 100644 --- a/api/tests/unit_tests/libs/test_helper.py +++ b/api/tests/unit_tests/libs/test_helper.py @@ -1,16 +1,24 @@ +import json from collections.abc import Callable from datetime import datetime, timedelta -from unittest.mock import call, create_autospec +from types import SimpleNamespace +from typing import cast +from unittest.mock import MagicMock, call, create_autospec import pytest from flask import request +from pydantic import ValidationError +import libs.helper as _token_manager_helper_module from dify_app import DifyApp from extensions import ext_proxy_fix +from extensions.ext_redis import RedisClientWrapper +from libs import helper as _rate_limiter_helper_module from libs import rate_limit as rate_limit_module from libs.helper import ( OptionalTimestampField, RateLimiter, + TokenManager, alphanumeric, email, escape_like_pattern, @@ -21,6 +29,24 @@ from models.account import Account from models.model import EndUser +def test_email_with_valid_email() -> None: + assert email("test@example.com") == "test@example.com" + assert email("TEST12345@example.com") == "TEST12345@example.com" + assert email("test+test@example.com") == "test+test@example.com" + assert email("!#$%&'*+-/=?^_{|}~`@example.com") == "!#$%&'*+-/=?^_{|}~`@example.com" + + +def test_email_with_invalid_email() -> None: + with pytest.raises(ValueError, match="invalid_email is not a valid email."): + email("invalid_email") + + with pytest.raises(ValueError, match="@example.com is not a valid email."): + email("@example.com") + + with pytest.raises(ValueError, match="()@example.com is not a valid email."): + email("()@example.com") + + @pytest.mark.parametrize( ("trust_proxy", "headers", "expected_ips", "expected_statuses"), [ @@ -285,3 +311,258 @@ class TestAlphanumericValidator: alphanumeric("tool.name") with pytest.raises(ValueError, match="not a valid alphanumeric value"): alphanumeric("tool/name") + + +class _FakeRedis: + def __init__(self) -> None: + self._zsets: dict[str, dict[str, float]] = {} + self._expiry: dict[str, int] = {} + + def zadd(self, key: str, mapping: dict[str, float]) -> int: + zset = self._zsets.setdefault(key, {}) + for member, score in mapping.items(): + zset[member] = float(score) + return len(mapping) + + def zremrangebyscore(self, key: str, min_score: str | float, max_score: str | float) -> int: + zset = self._zsets.get(key, {}) + min_value = float("-inf") if min_score == "-inf" else float(min_score) + max_value = float("inf") if max_score == "+inf" else float(max_score) + to_delete = [member for member, score in zset.items() if min_value <= score <= max_value] + for member in to_delete: + del zset[member] + return len(to_delete) + + def zcard(self, key: str) -> int: + return len(self._zsets.get(key, {})) + + def expire(self, key: str, ttl: int) -> bool: + self._expiry[key] = ttl + return True + + +def test_rate_limiter_counts_attempts_within_same_second(monkeypatch: pytest.MonkeyPatch) -> None: + fake_redis = _FakeRedis() + monkeypatch.setattr(_rate_limiter_helper_module.time, "time", lambda: 1000) + + limiter = _rate_limiter_helper_module.RateLimiter( + prefix="test_rate_limit", + max_attempts=2, + time_window=60, + redis_client=cast(_rate_limiter_helper_module._RateLimiterRedisClient, fake_redis), + ) + + limiter.increment_rate_limit("203.0.113.10") + limiter.increment_rate_limit("203.0.113.10") + + assert limiter.is_rate_limited("203.0.113.10") is True + + +def test_rate_limiter_uses_injected_redis(monkeypatch: pytest.MonkeyPatch) -> None: + redis_client = MagicMock() + redis_client.zcard.return_value = 1 + monkeypatch.setattr(_rate_limiter_helper_module.time, "time", lambda: 1000) + + limiter = _rate_limiter_helper_module.RateLimiter( + prefix="test_rate_limit", + max_attempts=1, + time_window=60, + redis_client=redis_client, + ) + + limiter.increment_rate_limit("203.0.113.10") + limiter.is_rate_limited("203.0.113.10") + + assert redis_client.zadd.called is True + assert redis_client.zremrangebyscore.called is True + assert redis_client.zcard.called is True + + +@pytest.mark.parametrize("acquired", [True, False]) +def test_reservation_uses_one_atomic_command_and_keeps_successful_slots( + redis_transport: tuple[RedisClientWrapper, MagicMock], monkeypatch: pytest.MonkeyPatch, acquired: bool +) -> None: + redis, commands = redis_transport + commands.return_value = int(acquired) + monkeypatch.setattr(_rate_limiter_helper_module.time, "time", lambda: 1000) + limiter = _rate_limiter_helper_module.RateLimiter( + prefix="reset", max_attempts=1, time_window=60, member_factory=lambda: "reservation", redis_client=redis + ) + + with limiter.reserve("user@example.com") as allowed: + assert allowed is acquired + + assert commands.call_count == 1 + assert commands.call_args.args[0] == "EVALSHA" + assert commands.call_args.args[2:] == (1, "reset:user@example.com", 1000, 60, 1, "reservation") + + +@pytest.mark.parametrize("acquired", [True, False]) +def test_failed_reservation_releases_only_its_own_slot( + redis_transport: tuple[RedisClientWrapper, MagicMock], acquired: bool +) -> None: + redis, commands = redis_transport + commands.return_value = int(acquired) + limiter = _rate_limiter_helper_module.RateLimiter( + prefix="reset", max_attempts=1, time_window=60, member_factory=lambda: "reservation", redis_client=redis + ) + + with pytest.raises(RuntimeError, match="delivery failed"), limiter.reserve("user@example.com"): + raise RuntimeError("delivery failed") + + assert commands.call_count == (2 if acquired else 1) + if acquired: + assert commands.call_args.args == ("ZREM", "reset:user@example.com", "reservation") + + +def test_reservation_cleanup_does_not_hide_delivery_error( + redis_transport: tuple[RedisClientWrapper, MagicMock], +) -> None: + redis, commands = redis_transport + commands.side_effect = [1, ConnectionError("redis unavailable")] + limiter = _rate_limiter_helper_module.RateLimiter( + prefix="reset", max_attempts=1, time_window=60, redis_client=redis + ) + + with pytest.raises(RuntimeError, match="delivery failed"), limiter.reserve("user@example.com"): + raise RuntimeError("delivery failed") + + +def test_reservation_redis_failure_stops_the_operation(redis_transport: tuple[RedisClientWrapper, MagicMock]) -> None: + redis, commands = redis_transport + commands.side_effect = ConnectionError("redis unavailable") + limiter = _rate_limiter_helper_module.RateLimiter( + prefix="reset", max_attempts=1, time_window=60, redis_client=redis + ) + + with pytest.raises(ConnectionError, match="redis unavailable"), limiter.reserve("user@example.com"): + pytest.fail("Operation must not run without acquiring a slot") + + +""" +Regression tests for `libs.helper.TokenManager`. + +`TokenManager` is the storage primitive shared by multiple auth flows, so it +must preserve every metadata field written by the caller. Business-specific +validation now happens at the callsite boundary (for example, +`TokenManagerChangeEmailTokenGateway.read`), not inside `TokenManager`. +""" + + +def _build_fake_redis(storage: dict[str, str]) -> SimpleNamespace: + def store_value(key: str, _ttl: int, value: str) -> bool: + storage[key] = value + return True + + def load_value(key: str) -> str | None: + return storage.get(key) + + return SimpleNamespace( + setex=store_value, + get=load_value, + delete=lambda *_args, **_kwargs: None, + ) + + +def test_token_manager_roundtrip_preserves_untyped_metadata_keys(monkeypatch: pytest.MonkeyPatch) -> None: + """`TokenManager` must round-trip arbitrary metadata keys without silently + dropping fields such as `phase`, `email_change_phase`, or future auth + payload extensions. + """ + + storage: dict[str, str] = {} + monkeypatch.setattr(_token_manager_helper_module, "redis_client", _build_fake_redis(storage)) + + token = TokenManager.generate_token( + email="user@example.com", + token_type="change_email", + additional_data={ + "code": "654321", + "old_email": "old@example.com", + "phase": "legacy-phase", + "email_change_phase": "old_email", + "custom_marker": "preserve-me", + }, + ) + + data = TokenManager.get_token_data(token, "change_email") + + assert data is not None + assert data.get("phase") == "legacy-phase" + assert data.get("email_change_phase") == "old_email" + assert data.get("custom_marker") == "preserve-me" + + +def test_token_manager_roundtrip_uses_explicit_account_id(monkeypatch: pytest.MonkeyPatch) -> None: + """When both `account_id` and `email` are supplied, the token should bind + both primitive values without depending on an account model. + """ + + storage: dict[str, str] = {} + monkeypatch.setattr(_token_manager_helper_module, "redis_client", _build_fake_redis(storage)) + + token = TokenManager.generate_token( + account_id="acc-1", + email="new@example.com", + token_type="change_email", + additional_data={ + "code": "654321", + "old_email": "old@example.com", + "email_change_phase": "new_email", + }, + ) + + data = TokenManager.get_token_data(token, "change_email") + + assert data is not None + assert data.get("account_id") == "acc-1" + assert data.get("email") == "new@example.com" + assert data.get("old_email") == "old@example.com" + assert data.get("email_change_phase") == "new_email" + + +def test_token_manager_roundtrip_still_validates_declared_fields(monkeypatch: pytest.MonkeyPatch) -> None: + """Unknown fields should be preserved, but declared baseline fields should + still be validated by `_token_data_adapter`. + """ + + storage = { + "change_email:token:token-123": json.dumps( + { + "token_type": "change_email", + "account_id": "acc-1", + "email": ["not-a-string"], + "code": "654321", + "old_email": "old@example.com", + "email_change_phase": "old_email", + } + ) + } + monkeypatch.setattr(_token_manager_helper_module, "redis_client", _build_fake_redis(storage)) + + with pytest.raises(ValidationError): + TokenManager.get_token_data("token-123", "change_email") + + +def test_token_manager_roundtrip_validates_email_change_phase_as_string(monkeypatch: pytest.MonkeyPatch) -> None: + """`email_change_phase` is part of the shared baseline schema, so obviously + malformed discriminator values should fail before the change-email-specific + union parsing at the callsite boundary. + """ + + storage = { + "change_email:token:token-456": json.dumps( + { + "token_type": "change_email", + "account_id": "acc-1", + "email": "new@example.com", + "code": "654321", + "old_email": "old@example.com", + "email_change_phase": ["not-a-string"], + } + ) + } + monkeypatch.setattr(_token_manager_helper_module, "redis_client", _build_fake_redis(storage)) + + with pytest.raises(ValidationError): + TokenManager.get_token_data("token-456", "change_email") diff --git a/api/tests/unit_tests/libs/test_login.py b/api/tests/unit_tests/libs/test_login.py index d451ee43471..c3e3205bf9e 100644 --- a/api/tests/unit_tests/libs/test_login.py +++ b/api/tests/unit_tests/libs/test_login.py @@ -11,7 +11,7 @@ from extensions.ext_login import DifyLoginManager from libs.login import current_user from machinery.errors import ActiveWorkspaceRequiredError from models.account import Account, Tenant -from tests.unit_tests.config_override import apply_config_overrides +from tests.helpers.unit.config_override import apply_config_overrides @pytest.fixture diff --git a/api/tests/unit_tests/libs/test_oauth_clients.py b/api/tests/unit_tests/libs/test_oauth.py similarity index 92% rename from api/tests/unit_tests/libs/test_oauth_clients.py rename to api/tests/unit_tests/libs/test_oauth.py index 60746e3dedb..122a2849406 100644 --- a/api/tests/unit_tests/libs/test_oauth_clients.py +++ b/api/tests/unit_tests/libs/test_oauth.py @@ -4,7 +4,43 @@ from unittest.mock import MagicMock, patch import httpx import pytest -from libs.oauth import GitHubOAuth, GoogleOAuth, OAuthUserInfo, decode_oauth_state +from libs.oauth import GitHubOAuth, GoogleOAuth, OAuth, OAuthUserInfo, decode_oauth_state, encode_oauth_state + + +def test_oauth_base_methods_raise_not_implemented() -> None: + oauth = OAuth(client_id="id", client_secret="sec", redirect_uri="uri") + + with pytest.raises(NotImplementedError): + oauth.get_authorization_url() + + with pytest.raises(NotImplementedError): + oauth.get_access_token("code") + + with pytest.raises(NotImplementedError): + oauth.get_raw_user_info("token") + + with pytest.raises(NotImplementedError): + oauth._transform_user_info({}) + + +def test_oauth_state_round_trips_login_context() -> None: + state = encode_oauth_state( + invite_token="invite-123", + timezone="Asia/Shanghai", + language="zh-Hans", + redirect_url="/apps?category=workflow", + ) + + assert decode_oauth_state(state) == { + "invite_token": "invite-123", + "timezone": "Asia/Shanghai", + "language": "zh-Hans", + "redirect_url": "/apps?category=workflow", + } + + +def test_oauth_state_returns_empty_payload_for_invalid_state() -> None: + assert decode_oauth_state("invalid-state") == {} class BaseOAuthTest: diff --git a/api/tests/unit_tests/libs/test_oauth_base.py b/api/tests/unit_tests/libs/test_oauth_base.py deleted file mode 100644 index e495e3acc90..00000000000 --- a/api/tests/unit_tests/libs/test_oauth_base.py +++ /dev/null @@ -1,39 +0,0 @@ -import pytest - -from libs.oauth import OAuth, decode_oauth_state, encode_oauth_state - - -def test_oauth_base_methods_raise_not_implemented() -> None: - oauth = OAuth(client_id="id", client_secret="sec", redirect_uri="uri") - - with pytest.raises(NotImplementedError): - oauth.get_authorization_url() - - with pytest.raises(NotImplementedError): - oauth.get_access_token("code") - - with pytest.raises(NotImplementedError): - oauth.get_raw_user_info("token") - - with pytest.raises(NotImplementedError): - oauth._transform_user_info({}) - - -def test_oauth_state_round_trips_login_context() -> None: - state = encode_oauth_state( - invite_token="invite-123", - timezone="Asia/Shanghai", - language="zh-Hans", - redirect_url="/apps?category=workflow", - ) - - assert decode_oauth_state(state) == { - "invite_token": "invite-123", - "timezone": "Asia/Shanghai", - "language": "zh-Hans", - "redirect_url": "/apps?category=workflow", - } - - -def test_oauth_state_returns_empty_payload_for_invalid_state() -> None: - assert decode_oauth_state("invalid-state") == {} diff --git a/api/tests/unit_tests/libs/test_oauth_bearer.py b/api/tests/unit_tests/libs/test_oauth_bearer.py index ceed765bd4f..634227fa474 100644 --- a/api/tests/unit_tests/libs/test_oauth_bearer.py +++ b/api/tests/unit_tests/libs/test_oauth_bearer.py @@ -1,16 +1,24 @@ -"""The bearer catalog: which subject each token type serves and what it may do.""" - from __future__ import annotations import uuid from datetime import UTC, datetime +from unittest.mock import MagicMock, patch import pytest from redis import Redis from sqlalchemy.orm import sessionmaker from constants.oauth_bearer import Scope, SubjectType, TokenType -from libs.oauth_bearer import OAuthAccessTokenResolver, ResolvedRow, _TokenTypeResolver +from libs.oauth_bearer import ( + BearerAuthenticator, + InvalidBearerError, + OAuthAccessTokenResolver, + ResolvedRow, + Resolver, + _TokenTypeResolver, +) + +"""The bearer catalog: which subject each token type serves and what it may do.""" def _row(account_id: uuid.UUID | None) -> ResolvedRow: @@ -50,3 +58,40 @@ def test_scope_sets_are_pinned_exactly() -> None: ) def test_for_token_classifies_by_prefix_and_refuses_the_rest(token: str, expected: TokenType | None) -> None: assert TokenType.for_token(token) is expected + + +def _authenticator_with(resolver: Resolver) -> BearerAuthenticator: + return BearerAuthenticator({TokenType.OAUTH_ACCOUNT: resolver}) + + +@patch("libs.oauth_bearer.enforce_bearer_rate_limit") +def test_rate_limit_called_before_resolve(rl: MagicMock, monkeypatch: pytest.MonkeyPatch) -> None: + call_order: list[str] = [] + rl.side_effect = lambda _h: call_order.append("rl") + resolver = OAuthAccessTokenResolver(sessionmaker(), Redis()) + monkeypatch.setattr(resolver, "cache_get", lambda _h: call_order.append("resolve") or "invalid") + auth = _authenticator_with(resolver.for_token_type(TokenType.OAUTH_ACCOUNT)) + + with pytest.raises(InvalidBearerError): + auth.authenticate("dfoa_xyz") + + assert call_order == ["rl", "resolve"], f"expected rl before resolve, got {call_order}" + + +# An unknown prefix is refused before the rate limit; a known prefix is rate limited, then refused. +@pytest.mark.parametrize( + ("token", "rate_limited"), + [("zzz_xyz", False), ("dfoa_revoked", True)], + ids=["unknown prefix", "revoked"], +) +@patch("libs.oauth_bearer.enforce_bearer_rate_limit") +def test_every_refusal_raises_the_generic_invalid_bearer( + rl: MagicMock, token: str, rate_limited: bool, monkeypatch: pytest.MonkeyPatch +) -> None: + resolver = OAuthAccessTokenResolver(sessionmaker(), Redis()) + monkeypatch.setattr(resolver, "cache_get", lambda _h: "invalid") + auth = _authenticator_with(resolver.for_token_type(TokenType.OAUTH_ACCOUNT)) + with pytest.raises(InvalidBearerError) as exc: + auth.authenticate(token) + assert str(exc.value) == "invalid_bearer" + assert rl.called is rate_limited diff --git a/api/tests/unit_tests/libs/test_oauth_bearer_rate_limit_ordering.py b/api/tests/unit_tests/libs/test_oauth_bearer_rate_limit_ordering.py deleted file mode 100644 index 4601e96d616..00000000000 --- a/api/tests/unit_tests/libs/test_oauth_bearer_rate_limit_ordering.py +++ /dev/null @@ -1,47 +0,0 @@ -from __future__ import annotations - -from unittest.mock import MagicMock, patch - -import pytest -from redis import Redis -from sqlalchemy.orm import sessionmaker - -from constants.oauth_bearer import TokenType -from libs.oauth_bearer import BearerAuthenticator, InvalidBearerError, OAuthAccessTokenResolver, Resolver - - -def _authenticator_with(resolver: Resolver) -> BearerAuthenticator: - return BearerAuthenticator({TokenType.OAUTH_ACCOUNT: resolver}) - - -@patch("libs.oauth_bearer.enforce_bearer_rate_limit") -def test_rate_limit_called_before_resolve(rl: MagicMock, monkeypatch: pytest.MonkeyPatch) -> None: - call_order: list[str] = [] - rl.side_effect = lambda _h: call_order.append("rl") - resolver = OAuthAccessTokenResolver(sessionmaker(), Redis()) - monkeypatch.setattr(resolver, "cache_get", lambda _h: call_order.append("resolve") or "invalid") - auth = _authenticator_with(resolver.for_token_type(TokenType.OAUTH_ACCOUNT)) - - with pytest.raises(InvalidBearerError): - auth.authenticate("dfoa_xyz") - - assert call_order == ["rl", "resolve"], f"expected rl before resolve, got {call_order}" - - -# An unknown prefix is refused before the rate limit; a known prefix is rate limited, then refused. -@pytest.mark.parametrize( - ("token", "rate_limited"), - [("zzz_xyz", False), ("dfoa_revoked", True)], - ids=["unknown prefix", "revoked"], -) -@patch("libs.oauth_bearer.enforce_bearer_rate_limit") -def test_every_refusal_raises_the_generic_invalid_bearer( - rl: MagicMock, token: str, rate_limited: bool, monkeypatch: pytest.MonkeyPatch -) -> None: - resolver = OAuthAccessTokenResolver(sessionmaker(), Redis()) - monkeypatch.setattr(resolver, "cache_get", lambda _h: "invalid") - auth = _authenticator_with(resolver.for_token_type(TokenType.OAUTH_ACCOUNT)) - with pytest.raises(InvalidBearerError) as exc: - auth.authenticate(token) - assert str(exc.value) == "invalid_bearer" - assert rl.called is rate_limited diff --git a/api/tests/unit_tests/libs/test_pyrefly_diagnostics.py b/api/tests/unit_tests/libs/test_pyrefly_diagnostics.py index b4c3e0744fe..801cfc64235 100644 --- a/api/tests/unit_tests/libs/test_pyrefly_diagnostics.py +++ b/api/tests/unit_tests/libs/test_pyrefly_diagnostics.py @@ -44,7 +44,7 @@ def test_extract_diagnostics_keeps_warn_headlines_and_location_lines() -> None: # Arrange raw_output = """INFO Checking project configured at `/tmp/project/pyrefly.toml` WARN Skipping include pattern `/tmp/project/tests` because it is matched by `project-excludes`. - --> tests/test_containers_integration_tests/pyrefly.toml:3:1 + --> tests/e2e_tests/containers/pyrefly.toml:3:1 """ # Act @@ -53,7 +53,7 @@ WARN Skipping include pattern `/tmp/project/tests` because it is matched by `pro # Assert assert diagnostics == ( "WARN Skipping include pattern `/tmp/project/tests` because it is matched by `project-excludes`.\n" - " --> tests/test_containers_integration_tests/pyrefly.toml:3:1\n" + " --> tests/e2e_tests/containers/pyrefly.toml:3:1\n" ) @@ -61,7 +61,7 @@ def test_render_diagnostics_falls_back_to_raw_output_for_nonzero_exit_without_ma # Arrange raw_output = ( "INFO Checking project configured at `/tmp/project/pyrefly.toml`\n" - "No Python files matched pattern `/tmp/project/tests/test_containers_integration_tests`\n" + "No Python files matched pattern `/tmp/project/tests/e2e_tests/containers`\n" ) # Act diff --git a/api/tests/unit_tests/libs/test_rate_limit_bearer.py b/api/tests/unit_tests/libs/test_rate_limit.py similarity index 100% rename from api/tests/unit_tests/libs/test_rate_limit_bearer.py rename to api/tests/unit_tests/libs/test_rate_limit.py diff --git a/api/tests/unit_tests/libs/test_rate_limiter.py b/api/tests/unit_tests/libs/test_rate_limiter.py deleted file mode 100644 index 378cf88b1d4..00000000000 --- a/api/tests/unit_tests/libs/test_rate_limiter.py +++ /dev/null @@ -1,129 +0,0 @@ -from typing import cast -from unittest.mock import MagicMock - -import pytest - -from extensions.ext_redis import RedisClientWrapper -from libs import helper as helper_module - - -class _FakeRedis: - def __init__(self) -> None: - self._zsets: dict[str, dict[str, float]] = {} - self._expiry: dict[str, int] = {} - - def zadd(self, key: str, mapping: dict[str, float]) -> int: - zset = self._zsets.setdefault(key, {}) - for member, score in mapping.items(): - zset[member] = float(score) - return len(mapping) - - def zremrangebyscore(self, key: str, min_score: str | float, max_score: str | float) -> int: - zset = self._zsets.get(key, {}) - min_value = float("-inf") if min_score == "-inf" else float(min_score) - max_value = float("inf") if max_score == "+inf" else float(max_score) - to_delete = [member for member, score in zset.items() if min_value <= score <= max_value] - for member in to_delete: - del zset[member] - return len(to_delete) - - def zcard(self, key: str) -> int: - return len(self._zsets.get(key, {})) - - def expire(self, key: str, ttl: int) -> bool: - self._expiry[key] = ttl - return True - - -def test_rate_limiter_counts_attempts_within_same_second(monkeypatch: pytest.MonkeyPatch) -> None: - fake_redis = _FakeRedis() - monkeypatch.setattr(helper_module.time, "time", lambda: 1000) - - limiter = helper_module.RateLimiter( - prefix="test_rate_limit", - max_attempts=2, - time_window=60, - redis_client=cast(helper_module._RateLimiterRedisClient, fake_redis), - ) - - limiter.increment_rate_limit("203.0.113.10") - limiter.increment_rate_limit("203.0.113.10") - - assert limiter.is_rate_limited("203.0.113.10") is True - - -def test_rate_limiter_uses_injected_redis(monkeypatch: pytest.MonkeyPatch) -> None: - redis_client = MagicMock() - redis_client.zcard.return_value = 1 - monkeypatch.setattr(helper_module.time, "time", lambda: 1000) - - limiter = helper_module.RateLimiter( - prefix="test_rate_limit", - max_attempts=1, - time_window=60, - redis_client=redis_client, - ) - - limiter.increment_rate_limit("203.0.113.10") - limiter.is_rate_limited("203.0.113.10") - - assert redis_client.zadd.called is True - assert redis_client.zremrangebyscore.called is True - assert redis_client.zcard.called is True - - -@pytest.mark.parametrize("acquired", [True, False]) -def test_reservation_uses_one_atomic_command_and_keeps_successful_slots( - redis_transport: tuple[RedisClientWrapper, MagicMock], monkeypatch: pytest.MonkeyPatch, acquired: bool -) -> None: - redis, commands = redis_transport - commands.return_value = int(acquired) - monkeypatch.setattr(helper_module.time, "time", lambda: 1000) - limiter = helper_module.RateLimiter( - prefix="reset", max_attempts=1, time_window=60, member_factory=lambda: "reservation", redis_client=redis - ) - - with limiter.reserve("user@example.com") as allowed: - assert allowed is acquired - - assert commands.call_count == 1 - assert commands.call_args.args[0] == "EVALSHA" - assert commands.call_args.args[2:] == (1, "reset:user@example.com", 1000, 60, 1, "reservation") - - -@pytest.mark.parametrize("acquired", [True, False]) -def test_failed_reservation_releases_only_its_own_slot( - redis_transport: tuple[RedisClientWrapper, MagicMock], acquired: bool -) -> None: - redis, commands = redis_transport - commands.return_value = int(acquired) - limiter = helper_module.RateLimiter( - prefix="reset", max_attempts=1, time_window=60, member_factory=lambda: "reservation", redis_client=redis - ) - - with pytest.raises(RuntimeError, match="delivery failed"), limiter.reserve("user@example.com"): - raise RuntimeError("delivery failed") - - assert commands.call_count == (2 if acquired else 1) - if acquired: - assert commands.call_args.args == ("ZREM", "reset:user@example.com", "reservation") - - -def test_reservation_cleanup_does_not_hide_delivery_error( - redis_transport: tuple[RedisClientWrapper, MagicMock], -) -> None: - redis, commands = redis_transport - commands.side_effect = [1, ConnectionError("redis unavailable")] - limiter = helper_module.RateLimiter(prefix="reset", max_attempts=1, time_window=60, redis_client=redis) - - with pytest.raises(RuntimeError, match="delivery failed"), limiter.reserve("user@example.com"): - raise RuntimeError("delivery failed") - - -def test_reservation_redis_failure_stops_the_operation(redis_transport: tuple[RedisClientWrapper, MagicMock]) -> None: - redis, commands = redis_transport - commands.side_effect = ConnectionError("redis unavailable") - limiter = helper_module.RateLimiter(prefix="reset", max_attempts=1, time_window=60, redis_client=redis) - - with pytest.raises(ConnectionError, match="redis unavailable"), limiter.reserve("user@example.com"): - pytest.fail("Operation must not run without acquiring a slot") diff --git a/api/tests/unit_tests/libs/test_schedule_utils_enhanced.py b/api/tests/unit_tests/libs/test_schedule_utils.py similarity index 54% rename from api/tests/unit_tests/libs/test_schedule_utils_enhanced.py rename to api/tests/unit_tests/libs/test_schedule_utils.py index f07d3355f42..d1b9b0a8437 100644 --- a/api/tests/unit_tests/libs/test_schedule_utils_enhanced.py +++ b/api/tests/unit_tests/libs/test_schedule_utils.py @@ -1,10 +1,3 @@ -""" -Enhanced schedule_utils tests for new cron syntax support. - -These tests verify that the backend schedule_utils functions properly support -the enhanced cron syntax introduced in the frontend, ensuring full compatibility. -""" - import contextlib import unittest from datetime import UTC, datetime, timedelta @@ -15,6 +8,380 @@ from croniter import CroniterBadCronError from libs.schedule_utils import calculate_next_run_at, convert_12h_to_24h +""" +Enhanced cron syntax compatibility tests for croniter backend. + +This test suite mirrors the frontend cron-parser tests to ensure +complete compatibility between frontend and backend cron processing. +""" + + +class TestCronCompatibility(unittest.TestCase): + """Test enhanced cron syntax compatibility with frontend.""" + + def setUp(self): + """Set up test environment with fixed time.""" + self.base_time = datetime(2024, 1, 15, 10, 0, 0, tzinfo=UTC) + + def test_enhanced_dayofweek_syntax(self): + """Test enhanced day-of-week syntax compatibility.""" + test_cases = [ + ("0 9 * * 7", 0), # Sunday as 7 + ("0 9 * * 0", 0), # Sunday as 0 + ("0 9 * * MON", 1), # Monday abbreviation + ("0 9 * * TUE", 2), # Tuesday abbreviation + ("0 9 * * WED", 3), # Wednesday abbreviation + ("0 9 * * THU", 4), # Thursday abbreviation + ("0 9 * * FRI", 5), # Friday abbreviation + ("0 9 * * SAT", 6), # Saturday abbreviation + ("0 9 * * SUN", 0), # Sunday abbreviation + ] + + for expr, expected_weekday in test_cases: + with self.subTest(expr=expr): + next_time = calculate_next_run_at(expr, "UTC", self.base_time) + assert next_time is not None + assert (next_time.weekday() + 1 if next_time.weekday() < 6 else 0) == expected_weekday + assert next_time.hour == 9 + assert next_time.minute == 0 + + def test_enhanced_month_syntax(self): + """Test enhanced month syntax compatibility.""" + test_cases = [ + ("0 9 1 JAN *", 1), # January abbreviation + ("0 9 1 FEB *", 2), # February abbreviation + ("0 9 1 MAR *", 3), # March abbreviation + ("0 9 1 APR *", 4), # April abbreviation + ("0 9 1 MAY *", 5), # May abbreviation + ("0 9 1 JUN *", 6), # June abbreviation + ("0 9 1 JUL *", 7), # July abbreviation + ("0 9 1 AUG *", 8), # August abbreviation + ("0 9 1 SEP *", 9), # September abbreviation + ("0 9 1 OCT *", 10), # October abbreviation + ("0 9 1 NOV *", 11), # November abbreviation + ("0 9 1 DEC *", 12), # December abbreviation + ] + + for expr, expected_month in test_cases: + with self.subTest(expr=expr): + next_time = calculate_next_run_at(expr, "UTC", self.base_time) + assert next_time is not None + assert next_time.month == expected_month + assert next_time.day == 1 + assert next_time.hour == 9 + + def test_predefined_expressions(self): + """Test predefined cron expressions compatibility.""" + test_cases = [ + ("@yearly", lambda dt: dt.month == 1 and dt.day == 1 and dt.hour == 0), + ("@annually", lambda dt: dt.month == 1 and dt.day == 1 and dt.hour == 0), + ("@monthly", lambda dt: dt.day == 1 and dt.hour == 0), + ("@weekly", lambda dt: dt.weekday() == 6 and dt.hour == 0), # Sunday = 6 in weekday() + ("@daily", lambda dt: dt.hour == 0 and dt.minute == 0), + ("@midnight", lambda dt: dt.hour == 0 and dt.minute == 0), + ("@hourly", lambda dt: dt.minute == 0), + ] + + for expr, validator in test_cases: + with self.subTest(expr=expr): + next_time = calculate_next_run_at(expr, "UTC", self.base_time) + assert next_time is not None + assert validator(next_time), f"Validator failed for {expr}: {next_time}" + + def test_special_characters(self): + """Test special characters in cron expressions.""" + test_cases = [ + "0 9 ? * 1", # ? wildcard + "0 12 * * 7", # Sunday as 7 + "0 15 L * *", # Last day of month + ] + + for expr in test_cases: + with self.subTest(expr=expr): + try: + next_time = calculate_next_run_at(expr, "UTC", self.base_time) + assert next_time is not None + assert next_time > self.base_time + except Exception as e: + self.fail(f"Expression '{expr}' should be valid but raised: {e}") + + def test_range_and_list_syntax(self): + """Test range and list syntax with abbreviations.""" + test_cases = [ + "0 9 * * MON-FRI", # Weekday range with abbreviations + "0 9 * JAN-MAR *", # Month range with abbreviations + "0 9 * * SUN,WED,FRI", # Weekday list with abbreviations + "0 9 1 JAN,JUN,DEC *", # Month list with abbreviations + ] + + for expr in test_cases: + with self.subTest(expr=expr): + try: + next_time = calculate_next_run_at(expr, "UTC", self.base_time) + assert next_time is not None + assert next_time > self.base_time + except Exception as e: + self.fail(f"Expression '{expr}' should be valid but raised: {e}") + + def test_invalid_enhanced_syntax(self): + """Test that invalid enhanced syntax is properly rejected.""" + invalid_expressions = [ + "0 12 * JANUARY *", # Full month name (not supported) + "0 12 * * MONDAY", # Full day name (not supported) + "0 12 32 JAN *", # Invalid day with valid month + "15 10 1 * 8", # Invalid day of week + "15 10 1 INVALID *", # Invalid month abbreviation + "15 10 1 * INVALID", # Invalid day abbreviation + "@invalid", # Invalid predefined expression + ] + + for expr in invalid_expressions: + with self.subTest(expr=expr): + with pytest.raises((CroniterBadCronError, ValueError)): + calculate_next_run_at(expr, "UTC", self.base_time) + + def test_edge_cases_with_enhanced_syntax(self): + """Test edge cases with enhanced syntax.""" + test_cases = [ + ("0 0 29 FEB *", lambda dt: dt.month == 2 and dt.day == 29), # Feb 29 with month abbreviation + ] + + for expr, validator in test_cases: + with self.subTest(expr=expr): + try: + next_time = calculate_next_run_at(expr, "UTC", self.base_time) + if next_time: # Some combinations might not occur soon + assert validator(next_time), f"Validator failed for {expr}: {next_time}" + except (CroniterBadCronError, ValueError): + # Some edge cases might be valid but not have upcoming occurrences + pass + + # Test complex expressions that have specific constraints + complex_expr = "59 23 31 DEC SAT" # December 31st at 23:59 on Saturday + try: + next_time = calculate_next_run_at(complex_expr, "UTC", self.base_time) + if next_time: + # The next occurrence might not be exactly Dec 31 if it's not a Saturday + # Just verify it's a valid result + assert next_time is not None + assert next_time.hour == 23 + assert next_time.minute == 59 + except Exception: + # Complex date constraints might not have near-future occurrences + pass + + +class TestTimezoneCompatibility(unittest.TestCase): + """Test timezone compatibility between frontend and backend.""" + + def setUp(self): + """Set up test environment.""" + self.base_time = datetime(2024, 1, 15, 10, 0, 0, tzinfo=UTC) + + def test_timezone_consistency(self): + """Test that calculations are consistent across different timezones.""" + timezones = [ + "UTC", + "America/New_York", + "Europe/London", + "Asia/Tokyo", + "Asia/Kolkata", + "Australia/Sydney", + ] + + expression = "0 12 * * *" # Daily at noon + + for timezone in timezones: + with self.subTest(timezone=timezone): + next_time = calculate_next_run_at(expression, timezone, self.base_time) + assert next_time is not None + + # Convert back to the target timezone to verify it's noon + tz = pytz.timezone(timezone) + local_time = next_time.astimezone(tz) + assert local_time.hour == 12 + assert local_time.minute == 0 + + def test_dst_handling(self): + """Test DST boundary handling.""" + # Test around DST spring forward (March 2024) + dst_base = datetime(2024, 3, 8, 10, 0, 0, tzinfo=UTC) + expression = "0 2 * * *" # 2 AM daily (problematic during DST) + timezone = "America/New_York" + + try: + next_time = calculate_next_run_at(expression, timezone, dst_base) + assert next_time is not None + + # During DST spring forward, 2 AM becomes 3 AM - both are acceptable + tz = pytz.timezone(timezone) + local_time = next_time.astimezone(tz) + assert local_time.hour in [2, 3] # Either 2 AM or 3 AM is acceptable + except Exception as e: + self.fail(f"DST handling failed: {e}") + + def test_half_hour_timezones(self): + """Test timezones with half-hour offsets.""" + timezones_with_offsets = [ + ("Asia/Kolkata", 17, 30), # UTC+5:30 -> 12:00 UTC = 17:30 IST + ("Australia/Adelaide", 22, 30), # UTC+10:30 -> 12:00 UTC = 22:30 ACDT (summer time) + ] + + expression = "0 12 * * *" # Noon UTC + + for timezone, expected_hour, expected_minute in timezones_with_offsets: + with self.subTest(timezone=timezone): + try: + next_time = calculate_next_run_at(expression, timezone, self.base_time) + assert next_time is not None + + tz = pytz.timezone(timezone) + local_time = next_time.astimezone(tz) + assert local_time.hour == expected_hour + assert local_time.minute == expected_minute + except Exception: + # Some complex timezone calculations might vary + pass + + def test_invalid_timezone_handling(self): + """Test handling of invalid timezones.""" + expression = "0 12 * * *" + invalid_timezone = "Invalid/Timezone" + + with pytest.raises((ValueError, Exception)): # Should raise an exception + calculate_next_run_at(expression, invalid_timezone, self.base_time) + + +class TestFrontendBackendIntegration(unittest.TestCase): + """Test integration patterns that mirror frontend usage.""" + + def setUp(self): + """Set up test environment.""" + self.base_time = datetime(2024, 1, 15, 10, 0, 0, tzinfo=UTC) + + def test_execution_time_calculator_pattern(self): + """Test the pattern used by execution-time-calculator.ts.""" + # This mirrors the exact usage from execution-time-calculator.ts:47 + test_data = { + "cron_expression": "30 14 * * 1-5", # 2:30 PM weekdays + "timezone": "America/New_York", + } + + # Get next 5 execution times (like the frontend does) + execution_times = [] + current_base = self.base_time + + for _ in range(5): + next_time = calculate_next_run_at(test_data["cron_expression"], test_data["timezone"], current_base) + assert next_time is not None + execution_times.append(next_time) + current_base = next_time + timedelta(seconds=1) # Move slightly forward + + assert len(execution_times) == 5 + + # Validate each execution time + for exec_time in execution_times: + # Convert to local timezone + tz = pytz.timezone(test_data["timezone"]) + local_time = exec_time.astimezone(tz) + + # Should be weekdays (1-5) + assert local_time.weekday() in [0, 1, 2, 3, 4] # Mon-Fri in Python weekday + + # Should be 2:30 PM in local time + assert local_time.hour == 14 + assert local_time.minute == 30 + assert local_time.second == 0 + + def test_schedule_service_integration(self): + """Test integration with ScheduleService patterns.""" + from core.workflow.nodes.trigger_schedule.entities import VisualConfig + from services.trigger.schedule_service import ScheduleService + + # Test enhanced syntax through visual config conversion + visual_configs = [ + # Test with month abbreviations + { + "frequency": "monthly", + "config": VisualConfig(time="9:00 AM", monthly_days=[1]), + "expected_cron": "0 9 1 * *", + }, + # Test with weekday abbreviations + { + "frequency": "weekly", + "config": VisualConfig(time="2:30 PM", weekdays=["mon", "wed", "fri"]), + "expected_cron": "30 14 * * 1,3,5", + }, + ] + + for test_case in visual_configs: + with self.subTest(frequency=test_case["frequency"]): + cron_expr = ScheduleService.visual_to_cron(test_case["frequency"], test_case["config"]) + assert cron_expr == test_case["expected_cron"] + + # Verify the generated cron expression is valid + next_time = calculate_next_run_at(cron_expr, "UTC", self.base_time) + assert next_time is not None + + def test_error_handling_consistency(self): + """Test that error handling matches frontend expectations.""" + invalid_expressions = [ + "60 10 1 * *", # Invalid minute + "15 25 1 * *", # Invalid hour + "15 10 32 * *", # Invalid day + "15 10 1 13 *", # Invalid month + "15 10 1", # Too few fields + "15 10 1 * * *", # 6 fields (not supported in frontend) + "0 15 10 1 * * *", # 7 fields (not supported in frontend) + "invalid expression", # Completely invalid + ] + + for expr in invalid_expressions: + with self.subTest(expr=repr(expr)): + with pytest.raises((CroniterBadCronError, ValueError, Exception)): + calculate_next_run_at(expr, "UTC", self.base_time) + + # Note: Empty/whitespace expressions are not tested here as they are + # not expected in normal usage due to database constraints (nullable=False) + + def test_performance_requirements(self): + """Test that complex expressions parse within reasonable time.""" + import time + + complex_expressions = [ + "*/5 9-17 * * 1-5", # Every 5 minutes, weekdays, business hours + "0 */2 1,15 * *", # Every 2 hours on 1st and 15th + "30 14 * * 1,3,5", # Mon, Wed, Fri at 14:30 + "15,45 8-18 * * 1-5", # 15 and 45 minutes past hour, weekdays + "0 9 * JAN-MAR MON-FRI", # Enhanced syntax: Q1 weekdays at 9 AM + "0 12 ? * SUN", # Enhanced syntax: Sundays at noon with ? + ] + + start_time = time.time() + + for expr in complex_expressions: + with self.subTest(expr=expr): + try: + next_time = calculate_next_run_at(expr, "UTC", self.base_time) + assert next_time is not None + except CroniterBadCronError: + # Some enhanced syntax might not be supported, that's OK + pass + + end_time = time.time() + execution_time = (end_time - start_time) * 1000 # Convert to milliseconds + + # Should complete within reasonable time (less than 150ms like frontend) + assert execution_time < 150, "Complex expressions should parse quickly" + + +""" +Enhanced schedule_utils tests for new cron syntax support. + +These tests verify that the backend schedule_utils functions properly support +the enhanced cron syntax introduced in the frontend, ensuring full compatibility. +""" + class TestEnhancedCronSyntax(unittest.TestCase): """Test enhanced cron syntax in calculate_next_run_at.""" @@ -455,5 +822,12 @@ class TestRegressionEnhanced(unittest.TestCase): assert result == expected, f"12h conversion failed: {time_str}" +if __name__ == "__main__": + # Import timedelta for the test + from datetime import timedelta + + unittest.main() + + if __name__ == "__main__": unittest.main() diff --git a/api/tests/unit_tests/libs/test_sendgrid_client.py b/api/tests/unit_tests/libs/test_sendgrid.py similarity index 100% rename from api/tests/unit_tests/libs/test_sendgrid_client.py rename to api/tests/unit_tests/libs/test_sendgrid.py diff --git a/api/tests/unit_tests/libs/test_smtp_client.py b/api/tests/unit_tests/libs/test_smtp.py similarity index 100% rename from api/tests/unit_tests/libs/test_smtp_client.py rename to api/tests/unit_tests/libs/test_smtp.py diff --git a/api/tests/unit_tests/libs/test_token_manager.py b/api/tests/unit_tests/libs/test_token_manager.py deleted file mode 100644 index fcb30eff4dd..00000000000 --- a/api/tests/unit_tests/libs/test_token_manager.py +++ /dev/null @@ -1,136 +0,0 @@ -""" -Regression tests for `libs.helper.TokenManager`. - -`TokenManager` is the storage primitive shared by multiple auth flows, so it -must preserve every metadata field written by the caller. Business-specific -validation now happens at the callsite boundary (for example, -`TokenManagerChangeEmailTokenGateway.read`), not inside `TokenManager`. -""" - -import json -from types import SimpleNamespace - -import pytest -from pydantic import ValidationError - -import libs.helper as helper_module -from libs.helper import TokenManager - - -def _build_fake_redis(storage: dict[str, str]) -> SimpleNamespace: - def store_value(key: str, _ttl: int, value: str) -> bool: - storage[key] = value - return True - - def load_value(key: str) -> str | None: - return storage.get(key) - - return SimpleNamespace( - setex=store_value, - get=load_value, - delete=lambda *_args, **_kwargs: None, - ) - - -def test_token_manager_roundtrip_preserves_untyped_metadata_keys(monkeypatch: pytest.MonkeyPatch) -> None: - """`TokenManager` must round-trip arbitrary metadata keys without silently - dropping fields such as `phase`, `email_change_phase`, or future auth - payload extensions. - """ - - storage: dict[str, str] = {} - monkeypatch.setattr(helper_module, "redis_client", _build_fake_redis(storage)) - - token = TokenManager.generate_token( - email="user@example.com", - token_type="change_email", - additional_data={ - "code": "654321", - "old_email": "old@example.com", - "phase": "legacy-phase", - "email_change_phase": "old_email", - "custom_marker": "preserve-me", - }, - ) - - data = TokenManager.get_token_data(token, "change_email") - - assert data is not None - assert data.get("phase") == "legacy-phase" - assert data.get("email_change_phase") == "old_email" - assert data.get("custom_marker") == "preserve-me" - - -def test_token_manager_roundtrip_uses_explicit_account_id(monkeypatch: pytest.MonkeyPatch) -> None: - """When both `account_id` and `email` are supplied, the token should bind - both primitive values without depending on an account model. - """ - - storage: dict[str, str] = {} - monkeypatch.setattr(helper_module, "redis_client", _build_fake_redis(storage)) - - token = TokenManager.generate_token( - account_id="acc-1", - email="new@example.com", - token_type="change_email", - additional_data={ - "code": "654321", - "old_email": "old@example.com", - "email_change_phase": "new_email", - }, - ) - - data = TokenManager.get_token_data(token, "change_email") - - assert data is not None - assert data.get("account_id") == "acc-1" - assert data.get("email") == "new@example.com" - assert data.get("old_email") == "old@example.com" - assert data.get("email_change_phase") == "new_email" - - -def test_token_manager_roundtrip_still_validates_declared_fields(monkeypatch: pytest.MonkeyPatch) -> None: - """Unknown fields should be preserved, but declared baseline fields should - still be validated by `_token_data_adapter`. - """ - - storage = { - "change_email:token:token-123": json.dumps( - { - "token_type": "change_email", - "account_id": "acc-1", - "email": ["not-a-string"], - "code": "654321", - "old_email": "old@example.com", - "email_change_phase": "old_email", - } - ) - } - monkeypatch.setattr(helper_module, "redis_client", _build_fake_redis(storage)) - - with pytest.raises(ValidationError): - TokenManager.get_token_data("token-123", "change_email") - - -def test_token_manager_roundtrip_validates_email_change_phase_as_string(monkeypatch: pytest.MonkeyPatch) -> None: - """`email_change_phase` is part of the shared baseline schema, so obviously - malformed discriminator values should fail before the change-email-specific - union parsing at the callsite boundary. - """ - - storage = { - "change_email:token:token-456": json.dumps( - { - "token_type": "change_email", - "account_id": "acc-1", - "email": "new@example.com", - "code": "654321", - "old_email": "old@example.com", - "email_change_phase": ["not-a-string"], - } - ) - } - monkeypatch.setattr(helper_module, "redis_client", _build_fake_redis(storage)) - - with pytest.raises(ValidationError): - TokenManager.get_token_data("token-456", "change_email") diff --git a/api/tests/unit_tests/controllers/console/app/test_description_validation.py b/api/tests/unit_tests/libs/test_validators.py similarity index 90% rename from api/tests/unit_tests/controllers/console/app/test_description_validation.py rename to api/tests/unit_tests/libs/test_validators.py index ea3fd14ab59..962fb25b005 100644 --- a/api/tests/unit_tests/controllers/console/app/test_description_validation.py +++ b/api/tests/unit_tests/libs/test_validators.py @@ -1,6 +1,12 @@ import pytest from libs.validators import validate_description_length +from tests.unit_tests.controllers.conftest import app_query_services as app_query_services # noqa: PLC0414 +from tests.unit_tests.controllers.conftest import ( + datasource_application_dependencies as datasource_application_dependencies, # noqa: PLC0414 +) +from tests.unit_tests.controllers.conftest import error_response_contract as error_response_contract # noqa: PLC0414 +from tests.unit_tests.controllers.conftest import pipeline_application as pipeline_application # noqa: PLC0414 class TestDescriptionValidationUnit: diff --git a/api/tests/unit_tests/machinery/test_app_context.py b/api/tests/unit_tests/machinery/test_context.py similarity index 100% rename from api/tests/unit_tests/machinery/test_app_context.py rename to api/tests/unit_tests/machinery/test_context.py diff --git a/api/tests/unit_tests/migrations/test_uuidv7_pg18_migration.py b/api/tests/unit_tests/migrations/versions/test_2025_07_02_2332-1c9ba48be8e4_add_uuidv7_function_in_sql.py similarity index 98% rename from api/tests/unit_tests/migrations/test_uuidv7_pg18_migration.py rename to api/tests/unit_tests/migrations/versions/test_2025_07_02_2332-1c9ba48be8e4_add_uuidv7_function_in_sql.py index ae5d2a1d0d0..44871b29bf5 100644 --- a/api/tests/unit_tests/migrations/test_uuidv7_pg18_migration.py +++ b/api/tests/unit_tests/migrations/versions/test_2025_07_02_2332-1c9ba48be8e4_add_uuidv7_function_in_sql.py @@ -16,7 +16,7 @@ from unittest import mock import pytest MIGRATION_PATH = ( - Path(__file__).resolve().parents[3] + Path(__file__).resolve().parents[4] / "migrations" / "versions" / "2025_07_02_2332-1c9ba48be8e4_add_uuidv7_function_in_sql.py" diff --git a/api/tests/unit_tests/migrations/test_knowledge_fs_control_plane.py b/api/tests/unit_tests/migrations/versions/test_2026_07_21_1200-a4e7c2f91b30_add_knowledge_fs_control_plane.py similarity index 99% rename from api/tests/unit_tests/migrations/test_knowledge_fs_control_plane.py rename to api/tests/unit_tests/migrations/versions/test_2026_07_21_1200-a4e7c2f91b30_add_knowledge_fs_control_plane.py index eae4731b1e7..dc56d87db24 100644 --- a/api/tests/unit_tests/migrations/test_knowledge_fs_control_plane.py +++ b/api/tests/unit_tests/migrations/versions/test_2026_07_21_1200-a4e7c2f91b30_add_knowledge_fs_control_plane.py @@ -21,7 +21,7 @@ from models.knowledge_fs import ( ) _MIGRATION_PATH = ( - Path(__file__).resolve().parents[3] + Path(__file__).resolve().parents[4] / "migrations/versions/2026_07_21_1200-a4e7c2f91b30_add_knowledge_fs_control_plane.py" ) diff --git a/api/tests/unit_tests/migrations/test_knowledge_fs_cutover.py b/api/tests/unit_tests/migrations/versions/test_2026_07_21_1300-b7f2a9d41c60_add_knowledge_fs_cutover.py similarity index 98% rename from api/tests/unit_tests/migrations/test_knowledge_fs_cutover.py rename to api/tests/unit_tests/migrations/versions/test_2026_07_21_1300-b7f2a9d41c60_add_knowledge_fs_cutover.py index 4d7cc1ef778..c324d2d2017 100644 --- a/api/tests/unit_tests/migrations/test_knowledge_fs_cutover.py +++ b/api/tests/unit_tests/migrations/versions/test_2026_07_21_1300-b7f2a9d41c60_add_knowledge_fs_cutover.py @@ -17,7 +17,7 @@ from models.knowledge_fs_cutover import ( ) _MIGRATION_PATH = ( - Path(__file__).resolve().parents[3] / "migrations/versions/2026_07_21_1300-b7f2a9d41c60_add_knowledge_fs_cutover.py" + Path(__file__).resolve().parents[4] / "migrations/versions/2026_07_21_1300-b7f2a9d41c60_add_knowledge_fs_cutover.py" ) _MODELS = ( diff --git a/api/tests/unit_tests/migrations/test_knowledge_fs_cleanup_authorization.py b/api/tests/unit_tests/migrations/versions/test_2026_07_21_1400-c8e31b7d52a4_add_knowledge_fs_cleanup_authorization.py similarity index 98% rename from api/tests/unit_tests/migrations/test_knowledge_fs_cleanup_authorization.py rename to api/tests/unit_tests/migrations/versions/test_2026_07_21_1400-c8e31b7d52a4_add_knowledge_fs_cleanup_authorization.py index c2cb2e112b2..69016d547e5 100644 --- a/api/tests/unit_tests/migrations/test_knowledge_fs_cleanup_authorization.py +++ b/api/tests/unit_tests/migrations/versions/test_2026_07_21_1400-c8e31b7d52a4_add_knowledge_fs_cleanup_authorization.py @@ -19,11 +19,11 @@ from models.knowledge_fs_cutover import ( ) _MIGRATION_PATH = ( - Path(__file__).resolve().parents[3] + Path(__file__).resolve().parents[4] / "migrations/versions/2026_07_21_1400-c8e31b7d52a4_add_knowledge_fs_cleanup_authorization.py" ) _CUTOVER_MIGRATION_PATH = ( - Path(__file__).resolve().parents[3] / "migrations/versions/2026_07_21_1300-b7f2a9d41c60_add_knowledge_fs_cutover.py" + Path(__file__).resolve().parents[4] / "migrations/versions/2026_07_21_1300-b7f2a9d41c60_add_knowledge_fs_cutover.py" ) _LATER_LEDGER_COLUMNS = { "remote_freeze_id", diff --git a/api/tests/unit_tests/migrations/test_knowledge_fs_remote_freeze.py b/api/tests/unit_tests/migrations/versions/test_2026_07_21_1500-d4f6e8a1c305_add_knowledge_fs_remote_freeze_evidence.py similarity index 97% rename from api/tests/unit_tests/migrations/test_knowledge_fs_remote_freeze.py rename to api/tests/unit_tests/migrations/versions/test_2026_07_21_1500-d4f6e8a1c305_add_knowledge_fs_remote_freeze_evidence.py index 30d2cca78ce..b866cedd4b8 100644 --- a/api/tests/unit_tests/migrations/test_knowledge_fs_remote_freeze.py +++ b/api/tests/unit_tests/migrations/versions/test_2026_07_21_1500-d4f6e8a1c305_add_knowledge_fs_remote_freeze_evidence.py @@ -11,7 +11,7 @@ from alembic.operations import Operations from models.knowledge_fs_cutover import KnowledgeFSWorkspaceCutoverLedger -_VERSIONS = Path(__file__).resolve().parents[3] / "migrations/versions" +_VERSIONS = Path(__file__).resolve().parents[4] / "migrations/versions" _MIGRATION_PATHS = ( _VERSIONS / "2026_07_21_1300-b7f2a9d41c60_add_knowledge_fs_cutover.py", _VERSIONS / "2026_07_21_1400-c8e31b7d52a4_add_knowledge_fs_cleanup_authorization.py", diff --git a/api/tests/unit_tests/migrations/test_knowledge_fs_cleanup_completion.py b/api/tests/unit_tests/migrations/versions/test_2026_07_21_1600-e5a7c9b2d416_add_knowledge_fs_cleanup_completion.py similarity index 98% rename from api/tests/unit_tests/migrations/test_knowledge_fs_cleanup_completion.py rename to api/tests/unit_tests/migrations/versions/test_2026_07_21_1600-e5a7c9b2d416_add_knowledge_fs_cleanup_completion.py index 49d2c1792ae..e26eac56463 100644 --- a/api/tests/unit_tests/migrations/test_knowledge_fs_cleanup_completion.py +++ b/api/tests/unit_tests/migrations/versions/test_2026_07_21_1600-e5a7c9b2d416_add_knowledge_fs_cleanup_completion.py @@ -11,7 +11,7 @@ from alembic.operations import Operations from models.knowledge_fs_cleanup import KnowledgeFSCleanupAuthorization -_VERSIONS = Path(__file__).resolve().parents[3] / "migrations/versions" +_VERSIONS = Path(__file__).resolve().parents[4] / "migrations/versions" _MIGRATION_PATHS = ( _VERSIONS / "2026_07_21_1300-b7f2a9d41c60_add_knowledge_fs_cutover.py", _VERSIONS / "2026_07_21_1400-c8e31b7d52a4_add_knowledge_fs_cleanup_authorization.py", diff --git a/api/tests/unit_tests/migrations/test_agent_workspace_migration.py b/api/tests/unit_tests/migrations/versions/test_2026_07_23_0203-f6e4c5686857_replace_agent_runtime_sessions_with_.py similarity index 99% rename from api/tests/unit_tests/migrations/test_agent_workspace_migration.py rename to api/tests/unit_tests/migrations/versions/test_2026_07_23_0203-f6e4c5686857_replace_agent_runtime_sessions_with_.py index b75921d8d81..a62d3b31aed 100644 --- a/api/tests/unit_tests/migrations/test_agent_workspace_migration.py +++ b/api/tests/unit_tests/migrations/versions/test_2026_07_23_0203-f6e4c5686857_replace_agent_runtime_sessions_with_.py @@ -10,7 +10,7 @@ from alembic.migration import MigrationContext from alembic.operations import Operations _MIGRATION_PATH = ( - Path(__file__).resolve().parents[3] + Path(__file__).resolve().parents[4] / "migrations/versions/2026_07_23_0203-f6e4c5686857_replace_agent_runtime_sessions_with_.py" ) diff --git a/api/tests/unit_tests/migrations/test_nullable_home_snapshot_migration.py b/api/tests/unit_tests/migrations/versions/test_2026_07_28_2331-e4708db55c1d_make_home_snapshot_references_nullable.py similarity index 99% rename from api/tests/unit_tests/migrations/test_nullable_home_snapshot_migration.py rename to api/tests/unit_tests/migrations/versions/test_2026_07_28_2331-e4708db55c1d_make_home_snapshot_references_nullable.py index adebda34291..971b3a771d1 100644 --- a/api/tests/unit_tests/migrations/test_nullable_home_snapshot_migration.py +++ b/api/tests/unit_tests/migrations/versions/test_2026_07_28_2331-e4708db55c1d_make_home_snapshot_references_nullable.py @@ -9,7 +9,7 @@ import sqlalchemy as sa from alembic.migration import MigrationContext from alembic.operations import Operations -_VERSIONS_DIR = Path(__file__).resolve().parents[3] / "migrations/versions" +_VERSIONS_DIR = Path(__file__).resolve().parents[4] / "migrations/versions" _MIGRATION_PATHS = ( _VERSIONS_DIR / "2026_07_21_2251-2f39536b3feb_add_agent_home_snapshot_ledger.py", _VERSIONS_DIR / "2026_07_23_0203-f6e4c5686857_replace_agent_runtime_sessions_with_.py", diff --git a/api/tests/unit_tests/migrations/test_knowledge_fs_staged_uploads.py b/api/tests/unit_tests/migrations/versions/test_2026_08_10_1200-7c1e9a4b2d60_add_knowledge_fs_staged_uploads.py similarity index 98% rename from api/tests/unit_tests/migrations/test_knowledge_fs_staged_uploads.py rename to api/tests/unit_tests/migrations/versions/test_2026_08_10_1200-7c1e9a4b2d60_add_knowledge_fs_staged_uploads.py index f1d42447d29..745e3d65e68 100644 --- a/api/tests/unit_tests/migrations/test_knowledge_fs_staged_uploads.py +++ b/api/tests/unit_tests/migrations/versions/test_2026_08_10_1200-7c1e9a4b2d60_add_knowledge_fs_staged_uploads.py @@ -12,7 +12,7 @@ from alembic.operations import Operations from models.knowledge_fs import KnowledgeFSStagedUpload _MIGRATION_PATH = ( - Path(__file__).resolve().parents[3] + Path(__file__).resolve().parents[4] / "migrations/versions/2026_08_10_1200-7c1e9a4b2d60_add_knowledge_fs_staged_uploads.py" ) diff --git a/api/tests/unit_tests/migrations/test_knowledge_fs_space_tag_bindings.py b/api/tests/unit_tests/migrations/versions/test_2026_08_13_1200-9d4e6f8a1b2c_add_knowledge_fs_space_tag_bindings.py similarity index 98% rename from api/tests/unit_tests/migrations/test_knowledge_fs_space_tag_bindings.py rename to api/tests/unit_tests/migrations/versions/test_2026_08_13_1200-9d4e6f8a1b2c_add_knowledge_fs_space_tag_bindings.py index 80752046b89..f2b5c4ae582 100644 --- a/api/tests/unit_tests/migrations/test_knowledge_fs_space_tag_bindings.py +++ b/api/tests/unit_tests/migrations/versions/test_2026_08_13_1200-9d4e6f8a1b2c_add_knowledge_fs_space_tag_bindings.py @@ -12,7 +12,7 @@ from alembic.operations import Operations from models.knowledge_fs import KnowledgeFSSpaceTagBinding _MIGRATION_PATH = ( - Path(__file__).resolve().parents[3] + Path(__file__).resolve().parents[4] / "migrations/versions/2026_08_13_1200-9d4e6f8a1b2c_add_knowledge_fs_space_tag_bindings.py" ) diff --git a/api/tests/unit_tests/migrations/test_knowledge_fs_icon_background.py b/api/tests/unit_tests/migrations/versions/test_2026_08_17_1200-4f8b2c7d9e10_add_knowledge_fs_icon_background.py similarity index 98% rename from api/tests/unit_tests/migrations/test_knowledge_fs_icon_background.py rename to api/tests/unit_tests/migrations/versions/test_2026_08_17_1200-4f8b2c7d9e10_add_knowledge_fs_icon_background.py index ff6fbea1ccf..06e76e6d8c7 100644 --- a/api/tests/unit_tests/migrations/test_knowledge_fs_icon_background.py +++ b/api/tests/unit_tests/migrations/versions/test_2026_08_17_1200-4f8b2c7d9e10_add_knowledge_fs_icon_background.py @@ -12,7 +12,7 @@ from alembic.operations import Operations from models.knowledge_fs import KnowledgeFSControlSpace _MIGRATION_PATH = ( - Path(__file__).resolve().parents[3] + Path(__file__).resolve().parents[4] / "migrations/versions/2026_08_17_1200-4f8b2c7d9e10_add_knowledge_fs_icon_background.py" ) diff --git a/api/tests/unit_tests/migrations/test_knowledge_fs_upgrade_jobs.py b/api/tests/unit_tests/migrations/versions/test_2026_08_17_1200-f3a8c1d7e920_add_knowledge_fs_upgrade_jobs.py similarity index 98% rename from api/tests/unit_tests/migrations/test_knowledge_fs_upgrade_jobs.py rename to api/tests/unit_tests/migrations/versions/test_2026_08_17_1200-f3a8c1d7e920_add_knowledge_fs_upgrade_jobs.py index 315ff28e819..dc777da18c1 100644 --- a/api/tests/unit_tests/migrations/test_knowledge_fs_upgrade_jobs.py +++ b/api/tests/unit_tests/migrations/versions/test_2026_08_17_1200-f3a8c1d7e920_add_knowledge_fs_upgrade_jobs.py @@ -17,7 +17,7 @@ from models.knowledge_fs import ( ) _MIGRATION_PATH = ( - Path(__file__).resolve().parents[3] + Path(__file__).resolve().parents[4] / "migrations/versions/2026_08_17_1200-f3a8c1d7e920_add_knowledge_fs_upgrade_jobs.py" ) _MODELS = ( diff --git a/api/tests/unit_tests/migrations/test_remove_agent_drive.py b/api/tests/unit_tests/migrations/versions/test_2026_08_17_1740-89919253ca7a_remove_agent_drive.py similarity index 99% rename from api/tests/unit_tests/migrations/test_remove_agent_drive.py rename to api/tests/unit_tests/migrations/versions/test_2026_08_17_1740-89919253ca7a_remove_agent_drive.py index a6b7ce10d13..27541389fea 100644 --- a/api/tests/unit_tests/migrations/test_remove_agent_drive.py +++ b/api/tests/unit_tests/migrations/versions/test_2026_08_17_1740-89919253ca7a_remove_agent_drive.py @@ -12,7 +12,7 @@ from alembic.migration import MigrationContext from alembic.operations import Operations _MIGRATION_PATH = ( - Path(__file__).resolve().parents[3] / "migrations/versions/2026_08_17_1740-89919253ca7a_remove_agent_drive.py" + Path(__file__).resolve().parents[4] / "migrations/versions/2026_08_17_1740-89919253ca7a_remove_agent_drive.py" ) diff --git a/api/tests/unit_tests/migrations/test_clean_legacy_agent_soul_files.py b/api/tests/unit_tests/migrations/versions/test_2026_08_20_0938-fbdfcf5f5a6e_clean_legacy_agent_soul_files.py similarity index 99% rename from api/tests/unit_tests/migrations/test_clean_legacy_agent_soul_files.py rename to api/tests/unit_tests/migrations/versions/test_2026_08_20_0938-fbdfcf5f5a6e_clean_legacy_agent_soul_files.py index 1f14a44c16a..afd0b2be6dc 100644 --- a/api/tests/unit_tests/migrations/test_clean_legacy_agent_soul_files.py +++ b/api/tests/unit_tests/migrations/versions/test_2026_08_20_0938-fbdfcf5f5a6e_clean_legacy_agent_soul_files.py @@ -15,7 +15,7 @@ from pydantic import ValidationError from models.agent_config_entities import AgentSoulConfig _MIGRATION_PATH = ( - Path(__file__).resolve().parents[3] + Path(__file__).resolve().parents[4] / "migrations/versions/2026_08_20_0938-fbdfcf5f5a6e_clean_legacy_agent_soul_files.py" ) diff --git a/api/tests/unit_tests/migrations/test_simplify_agent_v2_output_contract.py b/api/tests/unit_tests/migrations/versions/test_2026_08_21_1422-925e75620b69_simplify_agent_v2_output_contract.py similarity index 98% rename from api/tests/unit_tests/migrations/test_simplify_agent_v2_output_contract.py rename to api/tests/unit_tests/migrations/versions/test_2026_08_21_1422-925e75620b69_simplify_agent_v2_output_contract.py index 9526a3de580..311deac1433 100644 --- a/api/tests/unit_tests/migrations/test_simplify_agent_v2_output_contract.py +++ b/api/tests/unit_tests/migrations/versions/test_2026_08_21_1422-925e75620b69_simplify_agent_v2_output_contract.py @@ -15,7 +15,7 @@ from models.agent_config_entities import ( ) _MIGRATION_PATH = ( - Path(__file__).resolve().parents[3] + Path(__file__).resolve().parents[4] / "migrations/versions/2026_08_21_1422-925e75620b69_simplify_agent_v2_output_contract.py" ) diff --git a/api/tests/unit_tests/migrations/test_normalized_email_migration.py b/api/tests/unit_tests/migrations/versions/test_2026_08_25_1200-9b7c6d5e4f3a_add_normalized_email_to_accounts.py similarity index 98% rename from api/tests/unit_tests/migrations/test_normalized_email_migration.py rename to api/tests/unit_tests/migrations/versions/test_2026_08_25_1200-9b7c6d5e4f3a_add_normalized_email_to_accounts.py index b40fb6d45fc..e4e352b7390 100644 --- a/api/tests/unit_tests/migrations/test_normalized_email_migration.py +++ b/api/tests/unit_tests/migrations/versions/test_2026_08_25_1200-9b7c6d5e4f3a_add_normalized_email_to_accounts.py @@ -11,7 +11,7 @@ from alembic.migration import MigrationContext from alembic.operations import Operations _MIGRATION_PATH = ( - Path(__file__).resolve().parents[3] + Path(__file__).resolve().parents[4] / "migrations/versions/2026_08_25_1200-9b7c6d5e4f3a_add_normalized_email_to_accounts.py" ) diff --git a/api/tests/unit_tests/migrations/test_migrate_legacy_model_types.py b/api/tests/unit_tests/migrations/versions/test_2026_08_27_1200-5578e028b2f2_migrate_legacy_model_types.py similarity index 98% rename from api/tests/unit_tests/migrations/test_migrate_legacy_model_types.py rename to api/tests/unit_tests/migrations/versions/test_2026_08_27_1200-5578e028b2f2_migrate_legacy_model_types.py index 608747965ef..d4997c6eb2a 100644 --- a/api/tests/unit_tests/migrations/test_migrate_legacy_model_types.py +++ b/api/tests/unit_tests/migrations/versions/test_2026_08_27_1200-5578e028b2f2_migrate_legacy_model_types.py @@ -10,7 +10,7 @@ from alembic.migration import MigrationContext from alembic.operations import Operations _MIGRATION_PATH = ( - Path(__file__).resolve().parents[3] + Path(__file__).resolve().parents[4] / "migrations/versions/2026_08_27_1200-5578e028b2f2_migrate_legacy_model_types.py" ) diff --git a/api/tests/unit_tests/migrations/test_dataset_api_token_binding_resource_type.py b/api/tests/unit_tests/migrations/versions/test_2026_09_03_1000-9a4e7d1c2b60_add_dataset_api_token_binding_resource_type.py similarity index 99% rename from api/tests/unit_tests/migrations/test_dataset_api_token_binding_resource_type.py rename to api/tests/unit_tests/migrations/versions/test_2026_09_03_1000-9a4e7d1c2b60_add_dataset_api_token_binding_resource_type.py index 32d2dc34a34..edb9a4681c3 100644 --- a/api/tests/unit_tests/migrations/test_dataset_api_token_binding_resource_type.py +++ b/api/tests/unit_tests/migrations/versions/test_2026_09_03_1000-9a4e7d1c2b60_add_dataset_api_token_binding_resource_type.py @@ -12,7 +12,7 @@ from alembic.migration import MigrationContext from alembic.operations import Operations _MIGRATION_PATH = ( - Path(__file__).resolve().parents[3] + Path(__file__).resolve().parents[4] / "migrations/versions/2026_09_03_1000-9a4e7d1c2b60_add_dataset_api_token_binding_resource_type.py" ) diff --git a/api/tests/unit_tests/migrations/test_expanded_normalized_email_migration.py b/api/tests/unit_tests/migrations/versions/test_2026_09_10_1200-d8e4a6b1c902_expand_normalized_email_providers.py similarity index 98% rename from api/tests/unit_tests/migrations/test_expanded_normalized_email_migration.py rename to api/tests/unit_tests/migrations/versions/test_2026_09_10_1200-d8e4a6b1c902_expand_normalized_email_providers.py index 710d6d0265f..83feab7bcb2 100644 --- a/api/tests/unit_tests/migrations/test_expanded_normalized_email_migration.py +++ b/api/tests/unit_tests/migrations/versions/test_2026_09_10_1200-d8e4a6b1c902_expand_normalized_email_providers.py @@ -15,7 +15,7 @@ from services.account_email import normalize_email def _load_migration() -> ModuleType: path = ( - Path(__file__).resolve().parents[3] + Path(__file__).resolve().parents[4] / "migrations/versions/2026_09_10_1200-d8e4a6b1c902_expand_normalized_email_providers.py" ) spec = importlib.util.spec_from_file_location("expand_normalized_email_providers", path) diff --git a/api/tests/unit_tests/migrations/test_workflow_pause_task_id_migration.py b/api/tests/unit_tests/migrations/versions/test_2026_10_10_1200-c6f4a90d8b31_add_workflow_pause_task_id.py similarity index 99% rename from api/tests/unit_tests/migrations/test_workflow_pause_task_id_migration.py rename to api/tests/unit_tests/migrations/versions/test_2026_10_10_1200-c6f4a90d8b31_add_workflow_pause_task_id.py index b9001f0c5c8..b8f50d65c4a 100644 --- a/api/tests/unit_tests/migrations/test_workflow_pause_task_id_migration.py +++ b/api/tests/unit_tests/migrations/versions/test_2026_10_10_1200-c6f4a90d8b31_add_workflow_pause_task_id.py @@ -16,7 +16,7 @@ from flask import Flask from flask_migrate import Migrate from flask_sqlalchemy import SQLAlchemy -_MIGRATIONS_DIR = Path(__file__).resolve().parents[3] / "migrations" +_MIGRATIONS_DIR = Path(__file__).resolve().parents[4] / "migrations" _MIGRATION_PATH = _MIGRATIONS_DIR / "versions/2026_10_10_1200-c6f4a90d8b31_add_workflow_pause_task_id.py" diff --git a/api/tests/unit_tests/migrations/test_application_agent_tracing.py b/api/tests/unit_tests/migrations/versions/test_2026_10_15_0330-a17c2980f3e6_add_application_agent_tracing.py similarity index 97% rename from api/tests/unit_tests/migrations/test_application_agent_tracing.py rename to api/tests/unit_tests/migrations/versions/test_2026_10_15_0330-a17c2980f3e6_add_application_agent_tracing.py index 571ed2afb73..1a53a363b9c 100644 --- a/api/tests/unit_tests/migrations/test_application_agent_tracing.py +++ b/api/tests/unit_tests/migrations/versions/test_2026_10_15_0330-a17c2980f3e6_add_application_agent_tracing.py @@ -23,7 +23,7 @@ from models.application_trace import ApplicationTraceBatch, ApplicationTraceInvo def test_tracing_and_workflow_pause_upgrade_paths_join( starting_revision: str, required_revisions: tuple[str, ...] ) -> None: - scripts = ScriptDirectory(str(Path(__file__).resolve().parents[3] / "migrations")) + scripts = ScriptDirectory(str(Path(__file__).resolve().parents[4] / "migrations")) merge = scripts.get_revision("f7e2c9a6d481") assert merge is not None assert isinstance(merge.down_revision, tuple) @@ -38,7 +38,7 @@ def test_tracing_and_workflow_pause_upgrade_paths_join( def test_application_trace_migration_matches_models_and_roundtrips() -> None: path = ( - Path(__file__).resolve().parents[3] + Path(__file__).resolve().parents[4] / "migrations/versions/2026_10_15_0330-a17c2980f3e6_add_application_agent_tracing.py" ) spec = importlib.util.spec_from_file_location("application_trace_migration", path) diff --git a/api/tests/unit_tests/migrations/test_limit_active_knowledge_fs_upgrades_on_mysql.py b/api/tests/unit_tests/migrations/versions/test_2026_10_21_1300-3edd6cd27f54_limit_active_knowledge_fs_upgrades_on_mysql.py similarity index 98% rename from api/tests/unit_tests/migrations/test_limit_active_knowledge_fs_upgrades_on_mysql.py rename to api/tests/unit_tests/migrations/versions/test_2026_10_21_1300-3edd6cd27f54_limit_active_knowledge_fs_upgrades_on_mysql.py index 562476f56cb..ddc75fbeead 100644 --- a/api/tests/unit_tests/migrations/test_limit_active_knowledge_fs_upgrades_on_mysql.py +++ b/api/tests/unit_tests/migrations/versions/test_2026_10_21_1300-3edd6cd27f54_limit_active_knowledge_fs_upgrades_on_mysql.py @@ -10,7 +10,7 @@ from alembic.migration import MigrationContext from alembic.operations import Operations _MIGRATION_PATH = ( - Path(__file__).resolve().parents[3] + Path(__file__).resolve().parents[4] / "migrations/versions/2026_10_21_1300-3edd6cd27f54_limit_active_knowledge_fs_upgrades_on_mysql.py" ) diff --git a/api/tests/unit_tests/migrations/test_upgrade_availability_checkpoint.py b/api/tests/unit_tests/migrations/versions/test_2026_10_21_1400-a8c2e6f4b901_add_upgrade_availability_checkpoint.py similarity index 98% rename from api/tests/unit_tests/migrations/test_upgrade_availability_checkpoint.py rename to api/tests/unit_tests/migrations/versions/test_2026_10_21_1400-a8c2e6f4b901_add_upgrade_availability_checkpoint.py index f4f2364dabc..43f583c0783 100644 --- a/api/tests/unit_tests/migrations/test_upgrade_availability_checkpoint.py +++ b/api/tests/unit_tests/migrations/versions/test_2026_10_21_1400-a8c2e6f4b901_add_upgrade_availability_checkpoint.py @@ -12,7 +12,7 @@ from alembic.migration import MigrationContext from alembic.operations import Operations _PATH = ( - Path(__file__).resolve().parents[3] + Path(__file__).resolve().parents[4] / "migrations/versions/2026_10_21_1400-a8c2e6f4b901_add_upgrade_availability_checkpoint.py" ) diff --git a/api/tests/unit_tests/migrations/test_knowledge_fs_upgrade_reconciliation_lease.py b/api/tests/unit_tests/migrations/versions/test_2026_10_21_1500-b9d3f7a5c102_add_upgrade_reconciliation_lease.py similarity index 98% rename from api/tests/unit_tests/migrations/test_knowledge_fs_upgrade_reconciliation_lease.py rename to api/tests/unit_tests/migrations/versions/test_2026_10_21_1500-b9d3f7a5c102_add_upgrade_reconciliation_lease.py index 6f9d5b3f31f..8c37b56c17f 100644 --- a/api/tests/unit_tests/migrations/test_knowledge_fs_upgrade_reconciliation_lease.py +++ b/api/tests/unit_tests/migrations/versions/test_2026_10_21_1500-b9d3f7a5c102_add_upgrade_reconciliation_lease.py @@ -12,7 +12,7 @@ from alembic.migration import MigrationContext from alembic.operations import Operations _PATH = ( - Path(__file__).resolve().parents[3] + Path(__file__).resolve().parents[4] / "migrations/versions/2026_10_21_1500-b9d3f7a5c102_add_upgrade_reconciliation_lease.py" ) diff --git a/api/tests/unit_tests/models/test_account_models.py b/api/tests/unit_tests/models/test_account.py similarity index 100% rename from api/tests/unit_tests/models/test_account_models.py rename to api/tests/unit_tests/models/test_account.py diff --git a/api/tests/unit_tests/models/test_agent_config_entities.py b/api/tests/unit_tests/models/test_agent_config_entities.py index 6a3b816887f..172027c4a97 100644 --- a/api/tests/unit_tests/models/test_agent_config_entities.py +++ b/api/tests/unit_tests/models/test_agent_config_entities.py @@ -3,6 +3,9 @@ from pydantic import ValidationError from core.workflow.file_reference import build_file_reference from models.agent_config_entities import ( + AgentKnowledgeSpaceConfig, + AgentSoulConfig, + AgentSoulKnowledgeConfig, AgentSoulModelSettings, DeclaredArrayItem, DeclaredOutputChildConfig, @@ -308,3 +311,106 @@ def test_output_routes_reject_invalid_handle_identities(ids: list[str]) -> None: ), ): WorkflowNodeJobConfig.model_validate({"output_routes": {"routes": [{"id": id, "name": ""} for id in ids]}}) + + +"""KnowledgeFS authoring and historical snapshot decoding contracts.""" + + +def binding(**overrides: object) -> dict[str, object]: + return { + "id": "product", + "control_space_id": "ba0b32b1-6c20-493b-96a7-f2d0031f6ba9", + "name": "产品文档", + **overrides, + } + + +def test_knowledge_fs_round_trip_keeps_identity_and_missing_references() -> None: + soul = AgentSoulConfig.model_validate({"knowledge": {"spaces": [binding(is_missing=True)]}}) + restored = AgentSoulConfig.model_validate_json(soul.model_dump_json()) + assert restored.knowledge == soul.knowledge + assert restored.knowledge.spaces[0].is_missing + assert restored.knowledge.sets == [] + + +@pytest.mark.parametrize("value", ["bad-id", "../space", "00000000-0000-0000-0000-not-a-uuid00"]) +def test_invalid_control_space_identifiers_are_rejected(value: str) -> None: + with pytest.raises(ValidationError): + AgentKnowledgeSpaceConfig.model_validate(binding(control_space_id=value)) + + +@pytest.mark.parametrize("field", ["id", "name"]) +@pytest.mark.parametrize("value", [" ", "line\nbreak", "null\x00byte"]) +def test_alias_and_identity_cannot_inject_control_characters(field: str, value: str) -> None: + with pytest.raises(ValidationError): + AgentKnowledgeSpaceConfig.model_validate(binding(**{field: value})) + + +@pytest.mark.parametrize( + ("field", "value"), + [("id", "PRODUCT"), ("name", "产品文档"), ("control_space_id", "ba0b32b1-6c20-493b-96a7-f2d0031f6ba9")], +) +def test_duplicate_binding_identity_is_rejected(field: str, value: str) -> None: + second = binding(id="second", name="Policies", control_space_id="a84a7a42-773a-4ce1-8c30-6b932bc7d6d3") + second[field] = value + with pytest.raises(ValidationError, match="unique"): + AgentSoulKnowledgeConfig.model_validate({"spaces": [binding(), second]}) + + +def test_alias_cannot_shadow_another_stable_id() -> None: + with pytest.raises(ValidationError, match="shadow"): + AgentSoulKnowledgeConfig.model_validate( + { + "spaces": [ + binding(), + binding( + id="second", + name="PRODUCT", + control_space_id="a84a7a42-773a-4ce1-8c30-6b932bc7d6d3", + ), + ] + } + ) + + +def test_empty_historical_knowledge_is_still_readable() -> None: + assert AgentSoulKnowledgeConfig.model_validate({"sets": []}).spaces == [] + + +def test_mixed_legacy_and_knowledge_fs_config_round_trips() -> None: + config = AgentSoulKnowledgeConfig.model_validate( + { + "spaces": [binding()], + "sets": [ + { + "id": "old", + "name": "Old", + "datasets": [{"id": "old-dataset"}], + "query": {"mode": "generated_query"}, + "retrieval": {"mode": "multiple", "top_k": 5}, + } + ], + } + ) + + assert AgentSoulKnowledgeConfig.model_validate_json(config.model_dump_json()) == config + + +@pytest.mark.parametrize("overrides", [{"id": "PRODUCT"}, {"name": "产品文档"}]) +def test_mixed_bindings_cannot_share_identity(overrides: dict[str, str]) -> None: + with pytest.raises(ValidationError, match="unique across Agent and Classic"): + AgentSoulKnowledgeConfig.model_validate( + { + "spaces": [binding()], + "sets": [ + { + "id": "old", + "name": "Old", + "datasets": [{"id": "old-dataset"}], + "query": {"mode": "generated_query"}, + "retrieval": {"mode": "multiple", "top_k": 5}, + **overrides, + } + ], + } + ) diff --git a/api/tests/unit_tests/models/test_agent_knowledge_fs_config.py b/api/tests/unit_tests/models/test_agent_knowledge_fs_config.py deleted file mode 100644 index 5f77a18b7a2..00000000000 --- a/api/tests/unit_tests/models/test_agent_knowledge_fs_config.py +++ /dev/null @@ -1,106 +0,0 @@ -"""KnowledgeFS authoring and historical snapshot decoding contracts.""" - -import pytest -from pydantic import ValidationError - -from models.agent_config_entities import AgentKnowledgeSpaceConfig, AgentSoulConfig, AgentSoulKnowledgeConfig - - -def binding(**overrides: object) -> dict[str, object]: - return { - "id": "product", - "control_space_id": "ba0b32b1-6c20-493b-96a7-f2d0031f6ba9", - "name": "产品文档", - **overrides, - } - - -def test_knowledge_fs_round_trip_keeps_identity_and_missing_references() -> None: - soul = AgentSoulConfig.model_validate({"knowledge": {"spaces": [binding(is_missing=True)]}}) - restored = AgentSoulConfig.model_validate_json(soul.model_dump_json()) - assert restored.knowledge == soul.knowledge - assert restored.knowledge.spaces[0].is_missing - assert restored.knowledge.sets == [] - - -@pytest.mark.parametrize("value", ["bad-id", "../space", "00000000-0000-0000-0000-not-a-uuid00"]) -def test_invalid_control_space_identifiers_are_rejected(value: str) -> None: - with pytest.raises(ValidationError): - AgentKnowledgeSpaceConfig.model_validate(binding(control_space_id=value)) - - -@pytest.mark.parametrize("field", ["id", "name"]) -@pytest.mark.parametrize("value", [" ", "line\nbreak", "null\x00byte"]) -def test_alias_and_identity_cannot_inject_control_characters(field: str, value: str) -> None: - with pytest.raises(ValidationError): - AgentKnowledgeSpaceConfig.model_validate(binding(**{field: value})) - - -@pytest.mark.parametrize( - ("field", "value"), - [("id", "PRODUCT"), ("name", "产品文档"), ("control_space_id", "ba0b32b1-6c20-493b-96a7-f2d0031f6ba9")], -) -def test_duplicate_binding_identity_is_rejected(field: str, value: str) -> None: - second = binding(id="second", name="Policies", control_space_id="a84a7a42-773a-4ce1-8c30-6b932bc7d6d3") - second[field] = value - with pytest.raises(ValidationError, match="unique"): - AgentSoulKnowledgeConfig.model_validate({"spaces": [binding(), second]}) - - -def test_alias_cannot_shadow_another_stable_id() -> None: - with pytest.raises(ValidationError, match="shadow"): - AgentSoulKnowledgeConfig.model_validate( - { - "spaces": [ - binding(), - binding( - id="second", - name="PRODUCT", - control_space_id="a84a7a42-773a-4ce1-8c30-6b932bc7d6d3", - ), - ] - } - ) - - -def test_empty_historical_knowledge_is_still_readable() -> None: - assert AgentSoulKnowledgeConfig.model_validate({"sets": []}).spaces == [] - - -def test_mixed_legacy_and_knowledge_fs_config_round_trips() -> None: - config = AgentSoulKnowledgeConfig.model_validate( - { - "spaces": [binding()], - "sets": [ - { - "id": "old", - "name": "Old", - "datasets": [{"id": "old-dataset"}], - "query": {"mode": "generated_query"}, - "retrieval": {"mode": "multiple", "top_k": 5}, - } - ], - } - ) - - assert AgentSoulKnowledgeConfig.model_validate_json(config.model_dump_json()) == config - - -@pytest.mark.parametrize("overrides", [{"id": "PRODUCT"}, {"name": "产品文档"}]) -def test_mixed_bindings_cannot_share_identity(overrides: dict[str, str]) -> None: - with pytest.raises(ValidationError, match="unique across Agent and Classic"): - AgentSoulKnowledgeConfig.model_validate( - { - "spaces": [binding()], - "sets": [ - { - "id": "old", - "name": "Old", - "datasets": [{"id": "old-dataset"}], - "query": {"mode": "generated_query"}, - "retrieval": {"mode": "multiple", "top_k": 5}, - **overrides, - } - ], - } - ) diff --git a/api/tests/unit_tests/models/test_annotation_hit_history_session_accessors.py b/api/tests/unit_tests/models/test_annotation_hit_history_session_accessors.py deleted file mode 100644 index d4bb5c624e5..00000000000 --- a/api/tests/unit_tests/models/test_annotation_hit_history_session_accessors.py +++ /dev/null @@ -1,93 +0,0 @@ -"""Regression coverage for the ``@property``→session-parameter refactor on -``AppAnnotationHitHistory`` accessors. - -Covers two of `api/models/model.py`'s legacy `@property` accessors that reached for the global -``db.session`` internally and have been converted to plain methods taking an explicit -``session: Session`` (per the pattern established in #40370/#40797/#41394/#41830, tracked in -#40372): - -- ``AppAnnotationHitHistory.account`` (joins through ``MessageAnnotation.account_id``) -- ``AppAnnotationHitHistory.annotation_create_account`` (looks up ``self.account_id`` directly) - -Each accessor is exercised against the real ``sqlite_session`` fixture (a genuine SQLAlchemy -``Session`` bound to a pristine full-schema SQLite database) so the assertions cover actual -query behaviour rather than a mock's recorded call. -""" - -from uuid import uuid4 - -from sqlalchemy.orm import Session - -from models.account import Account -from models.model import AppAnnotationHitHistory, MessageAnnotation - - -def _persist_account(session: Session) -> Account: - account = Account(name="Test Account", email="test@example.com") - session.add(account) - session.flush() - return account - - -def _persist_annotation(session: Session, *, app_id: str, account_id: str) -> MessageAnnotation: - annotation = MessageAnnotation( - app_id=app_id, - conversation_id=None, - message_id=None, - question="What is AI?", - content="AI stands for Artificial Intelligence.", - account_id=account_id, - ) - session.add(annotation) - session.flush() - return annotation - - -def _hit_history(*, app_id: str, annotation_id: str, account_id: str) -> AppAnnotationHitHistory: - return AppAnnotationHitHistory( - app_id=app_id, - annotation_id=annotation_id, - source="api", - question="What is AI?", - account_id=account_id, - score=0.95, - message_id=str(uuid4()), - annotation_question="What is AI?", - annotation_content="AI stands for Artificial Intelligence.", - ) - - -class TestAppAnnotationHitHistoryAccount: - def test_returns_the_annotation_creator(self, sqlite_session: Session) -> None: - app_id = str(uuid4()) - creator = _persist_account(sqlite_session) - annotation = _persist_annotation(sqlite_session, app_id=app_id, account_id=creator.id) - history = _hit_history(app_id=app_id, annotation_id=annotation.id, account_id=str(uuid4())) - - result = history.account(session=sqlite_session) - - assert result is not None - assert result.id == creator.id - - def test_returns_none_when_annotation_missing(self, sqlite_session: Session) -> None: - history = _hit_history(app_id=str(uuid4()), annotation_id=str(uuid4()), account_id=str(uuid4())) - - assert history.account(session=sqlite_session) is None - - -class TestAppAnnotationHitHistoryAnnotationCreateAccount: - def test_returns_the_hit_requester(self, sqlite_session: Session) -> None: - app_id = str(uuid4()) - requester = _persist_account(sqlite_session) - annotation = _persist_annotation(sqlite_session, app_id=app_id, account_id=str(uuid4())) - history = _hit_history(app_id=app_id, annotation_id=annotation.id, account_id=requester.id) - - result = history.annotation_create_account(session=sqlite_session) - - assert result is not None - assert result.id == requester.id - - def test_returns_none_when_account_missing(self, sqlite_session: Session) -> None: - history = _hit_history(app_id=str(uuid4()), annotation_id=str(uuid4()), account_id=str(uuid4())) - - assert history.annotation_create_account(session=sqlite_session) is None diff --git a/api/tests/unit_tests/models/test_app_annotation_setting_session_accessors.py b/api/tests/unit_tests/models/test_app_annotation_setting_session_accessors.py deleted file mode 100644 index 23b86cd2c38..00000000000 --- a/api/tests/unit_tests/models/test_app_annotation_setting_session_accessors.py +++ /dev/null @@ -1,60 +0,0 @@ -"""Regression coverage for the ``@property``→session-parameter refactor on -``AppAnnotationSetting.collection_binding_detail``. - -The legacy ``@property`` reached for the global ``db.session`` internally and has been converted -to a plain method taking an explicit ``session: Session`` (per the pattern established in -#40370/#40797/#41394/#41830, tracked in #40372). - -The accessor is exercised against the real ``sqlite_session`` fixture (a genuine SQLAlchemy -``Session`` bound to a pristine full-schema SQLite database) so the assertions cover actual query -behaviour rather than a mock's recorded call. -""" - -from uuid import uuid4 - -from sqlalchemy.orm import Session - -from models.dataset import DatasetCollectionBinding -from models.enums import CollectionBindingType -from models.model import AppAnnotationSetting - - -def _persist_collection_binding(session: Session) -> DatasetCollectionBinding: - binding = DatasetCollectionBinding( - collection_name="test_collection", - provider_name="test_provider", - model_name="test_model", - type=CollectionBindingType.DATASET, - ) - session.add(binding) - session.flush() - return binding - - -def _persist_setting(session: Session, *, collection_binding_id: str) -> AppAnnotationSetting: - setting = AppAnnotationSetting( - app_id=str(uuid4()), - score_threshold=0.8, - collection_binding_id=collection_binding_id, - created_user_id=str(uuid4()), - updated_user_id=str(uuid4()), - ) - session.add(setting) - session.flush() - return setting - - -class TestAppAnnotationSettingCollectionBindingDetail: - def test_returns_the_bound_collection_binding(self, sqlite_session: Session) -> None: - binding = _persist_collection_binding(sqlite_session) - setting = _persist_setting(sqlite_session, collection_binding_id=binding.id) - - result = setting.collection_binding_detail(session=sqlite_session) - - assert result is not None - assert result.id == binding.id - - def test_returns_none_when_binding_missing(self, sqlite_session: Session) -> None: - setting = _persist_setting(sqlite_session, collection_binding_id=str(uuid4())) - - assert setting.collection_binding_detail(session=sqlite_session) is None diff --git a/api/tests/unit_tests/models/test_app_mcp_server_session_accessors.py b/api/tests/unit_tests/models/test_app_mcp_server_session_accessors.py deleted file mode 100644 index fab3482748e..00000000000 --- a/api/tests/unit_tests/models/test_app_mcp_server_session_accessors.py +++ /dev/null @@ -1,38 +0,0 @@ -"""Regression coverage for the ``db.session``→session-parameter refactor on -``AppMCPServer.generate_server_code``. - -The static method reached for the global ``db.session`` internally and has been converted to -take an explicit ``session: Session`` (per the pattern established in #40370/#40797/#41394, -tracked in #40372). -""" - -from unittest.mock import patch -from uuid import uuid4 - -from sqlalchemy.orm import Session - -from models.model import AppMCPServer - - -def _persist_server(session: Session, *, server_code: str) -> AppMCPServer: - server = AppMCPServer( - tenant_id=str(uuid4()), - app_id=str(uuid4()), - name="Test MCP Server", - description="Test server", - server_code=server_code, - status="active", - parameters="{}", - ) - session.add(server) - session.flush() - return server - - -def test_generate_server_code_retries_a_taken_code_through_the_given_session(sqlite_session: Session) -> None: - _persist_server(sqlite_session, server_code="taken-code") - - with patch("models.model.generate_string", side_effect=["taken-code", "fresh-code"]): - result = AppMCPServer.generate_server_code(16, session=sqlite_session) - - assert result == "fresh-code" diff --git a/api/tests/unit_tests/models/test_app_models.py b/api/tests/unit_tests/models/test_app_models.py deleted file mode 100644 index fe224066551..00000000000 --- a/api/tests/unit_tests/models/test_app_models.py +++ /dev/null @@ -1,1590 +0,0 @@ -""" -Comprehensive unit tests for App models. - -This test suite covers: -- App configuration validation -- App-Message relationships -- Conversation model integrity -- Annotation model relationships -""" - -import json -from datetime import UTC, datetime -from decimal import Decimal -from unittest.mock import patch -from uuid import UUID, uuid4 - -import pytest -from sqlalchemy import select -from sqlalchemy.orm import Session - -from models.dataset import DatasetCollectionBinding -from models.enums import CollectionBindingType, ConversationFromSource, CustomizeTokenStrategy -from models.model import ( - AnnotationReplyConfig, - App, - AppAnnotationHitHistory, - AppAnnotationSetting, - AppMode, - AppModelConfig, - Conversation, - IconType, - Message, - MessageAnnotation, - Site, - load_annotation_reply_config, -) -from models.workflow import Workflow, WorkflowType - - -class TestAppModelValidation: - """Test suite for App model validation and basic operations.""" - - def test_app_creation_with_required_fields(self): - """Test creating an app with all required fields.""" - # Arrange - tenant_id = str(uuid4()) - created_by = str(uuid4()) - - # Act - app = App( - tenant_id=tenant_id, - name="Test App", - mode=AppMode.CHAT, - enable_site=True, - enable_api=False, - created_by=created_by, - ) - - # Assert - assert app.name == "Test App" - assert app.tenant_id == tenant_id - assert app.mode == AppMode.CHAT - assert app.enable_site is True - assert app.enable_api is False - assert app.created_by == created_by - - def test_app_creation_with_optional_fields(self): - """Test creating an app with optional fields.""" - # Arrange & Act - app = App( - tenant_id=str(uuid4()), - name="Test App", - mode=AppMode.COMPLETION, - enable_site=True, - enable_api=True, - created_by=str(uuid4()), - description="Test description", - icon_type=IconType.EMOJI, - icon="🤖", - icon_background="#FF5733", - is_demo=True, - is_public=False, - api_rpm=100, - api_rph=1000, - ) - - # Assert - assert app.description == "Test description" - assert app.icon_type == IconType.EMOJI - assert app.icon == "🤖" - assert app.icon_background == "#FF5733" - assert app.is_demo is True - assert app.is_public is False - assert app.api_rpm == 100 - assert app.api_rph == 1000 - - def test_app_mode_validation(self): - """Test app mode enum values.""" - # Assert - expected_modes = { - "chat", - "completion", - "workflow", - "advanced-chat", - "agent-chat", - "agent", - "channel", - "rag-pipeline", - } - assert {mode.value for mode in AppMode} == expected_modes - - def test_app_mode_value_of(self): - """Test AppMode.value_of method.""" - # Act & Assert - assert AppMode.value_of("chat") == AppMode.CHAT - assert AppMode.value_of("completion") == AppMode.COMPLETION - assert AppMode.value_of("workflow") == AppMode.WORKFLOW - - with pytest.raises(ValueError, match="invalid mode value"): - AppMode.value_of("invalid_mode") - - def test_icon_type_validation(self): - """Test icon type enum values.""" - # Assert - assert {t.value for t in IconType} == {"image", "emoji", "link"} - - @pytest.mark.parametrize("sqlite_session", [(App, AppModelConfig)], indirect=True) - def test_app_desc_or_prompt_with_description(self, sqlite_session: Session): - """Test desc_or_prompt_with_session when description exists.""" - # Arrange - app = App( - tenant_id=str(uuid4()), - name="Test App", - mode=AppMode.CHAT, - enable_site=True, - enable_api=False, - created_by=str(uuid4()), - description="App description", - ) - - # Act - result = app.desc_or_prompt_with_session(session=sqlite_session) - - # Assert - assert result == "App description" - - @pytest.mark.parametrize("sqlite_session", [(App, AppModelConfig)], indirect=True) - def test_app_desc_or_prompt_without_description(self, sqlite_session: Session): - """Test desc_or_prompt property when description is empty.""" - # Arrange - app = App( - tenant_id=str(uuid4()), - name="Test App", - mode=AppMode.CHAT, - enable_site=True, - enable_api=False, - created_by=str(uuid4()), - description="", - ) - decoy_config = AppModelConfig(app_id=str(uuid4()), pre_prompt="A decoy prompt") - sqlite_session.add_all([app, decoy_config]) - sqlite_session.flush() - - # Act - result = app.desc_or_prompt_with_session(session=sqlite_session) - - # Assert - assert result == "" - - @pytest.mark.parametrize("sqlite_session", [(App, AppModelConfig)], indirect=True) - def test_app_is_agent_false_when_not_configured_as_agent(self, sqlite_session: Session): - """`is_agent_with_session` returns False when the config has no agent mode.""" - # Arrange - app = App( - tenant_id=str(uuid4()), - name="Test App", - mode=AppMode.CHAT, - enable_site=True, - enable_api=False, - created_by=str(uuid4()), - app_model_config_id=str(uuid4()), - ) - sqlite_session.add(app) - sqlite_session.flush() - - # Act - result = app.is_agent_with_session(session=sqlite_session) - - # Assert - assert result is False - - @pytest.mark.parametrize("sqlite_session", [(App, AppModelConfig)], indirect=True) - def test_app_is_agent_with_session_persists_mode_across_sessions(self, sqlite_session: Session): - app = App( - tenant_id=str(uuid4()), - name="Test App", - mode=AppMode.CHAT, - enable_site=True, - enable_api=False, - created_by=str(uuid4()), - ) - sqlite_session.add(app) - sqlite_session.flush() - model_config = AppModelConfig( - app_id=app.id, - agent_mode=json.dumps({"enabled": True, "strategy": "react"}), - ) - sqlite_session.add(model_config) - sqlite_session.flush() - app.app_model_config_id = model_config.id - sqlite_session.commit() - - with Session(sqlite_session.get_bind(), expire_on_commit=False) as migration_session: - assert app.is_agent_with_session(session=migration_session) is True - - sqlite_session.expire_all() - assert sqlite_session.get(App, app.id).mode == AppMode.AGENT_CHAT - - @pytest.mark.parametrize("sqlite_session", [(App, AppModelConfig)], indirect=True) - def test_app_mode_compatible_with_agent(self, sqlite_session: Session): - """Test mode_compatible_with_agent property.""" - # Arrange - app = App( - tenant_id=str(uuid4()), - name="Test App", - mode=AppMode.CHAT, - enable_site=True, - enable_api=False, - created_by=str(uuid4()), - ) - decoy_config = AppModelConfig( - app_id=str(uuid4()), - agent_mode=json.dumps({"enabled": True, "strategy": "react"}), - ) - sqlite_session.add_all([app, decoy_config]) - sqlite_session.flush() - - # Act - result = app.mode_compatible_with_agent_with_session(session=sqlite_session) - - # Assert - assert result == AppMode.CHAT - - @pytest.mark.parametrize("sqlite_session", [(App, AppModelConfig)], indirect=True) - def test_app_mode_compatible_with_agent_reports_agent_chat(self, sqlite_session: Session): - """A CHAT app whose own config enables agent mode reports AGENT_CHAT.""" - # Arrange - app = App( - tenant_id=str(uuid4()), - name="Test App", - mode=AppMode.CHAT, - enable_site=True, - enable_api=False, - created_by=str(uuid4()), - ) - sqlite_session.add(app) - sqlite_session.flush() - app.app_model_config_id = str(uuid4()) - config = AppModelConfig( - app_id=app.id, - agent_mode=json.dumps({"enabled": True, "strategy": "react"}), - ) - config.id = app.app_model_config_id - sqlite_session.add(config) - sqlite_session.flush() - - # Act - result = app.mode_compatible_with_agent_with_session(session=sqlite_session) - - # Assert - assert result == AppMode.AGENT_CHAT - - @pytest.mark.parametrize("sqlite_session", [(App, AppModelConfig)], indirect=True) - def test_app_model_config_with_session_reads_through_the_caller_session(self, sqlite_session: Session): - """`app_model_config_with_session` resolves the linked config through the caller's session.""" - # Arrange - app = App( - tenant_id=str(uuid4()), - name="Test App", - mode=AppMode.CHAT, - enable_site=True, - enable_api=False, - created_by=str(uuid4()), - ) - sqlite_session.add(app) - sqlite_session.flush() - model_config = AppModelConfig(app_id=app.id, agent_mode=json.dumps({"enabled": False})) - sqlite_session.add(model_config) - sqlite_session.flush() - app.app_model_config_id = model_config.id - sqlite_session.commit() - - # Act: a session the model never owns, proving the lookup does not reach for a global one - with Session(sqlite_session.get_bind(), expire_on_commit=False) as caller_session: - result = app.app_model_config_with_session(session=caller_session) - - # Assert - assert result is not None - assert result.id == model_config.id - assert result in caller_session - - @pytest.mark.parametrize("sqlite_session", [(App, AppModelConfig)], indirect=True) - def test_app_model_config_with_session_returns_none_when_unlinked(self, sqlite_session: Session): - """An app without `app_model_config_id` resolves to None without querying.""" - # Arrange - app = App( - tenant_id=str(uuid4()), - name="Test App", - mode=AppMode.CHAT, - enable_site=True, - enable_api=False, - created_by=str(uuid4()), - ) - sqlite_session.add(app) - sqlite_session.flush() - - # Act / Assert: an unbound session would raise if the guard were dropped - with Session(expire_on_commit=False) as unbound_session: - assert app.app_model_config_with_session(session=unbound_session) is None - - @pytest.mark.parametrize("sqlite_session", [(App, AppModelConfig)], indirect=True) - def test_workflow_with_session_reads_through_the_caller_session(self, sqlite_session: Session): - """`workflow_with_session` resolves the published workflow through the caller's session.""" - # Arrange - app = App( - tenant_id=str(uuid4()), - name="Test App", - mode=AppMode.WORKFLOW, - enable_site=True, - enable_api=False, - created_by=str(uuid4()), - ) - sqlite_session.add(app) - sqlite_session.flush() - workflow = Workflow( - tenant_id=app.tenant_id, - app_id=app.id, - type=WorkflowType.WORKFLOW, - version="1", - graph=json.dumps({"nodes": [], "edges": []}), - created_by=app.created_by, - ) - workflow._features = "{}" - sqlite_session.add(workflow) - sqlite_session.flush() - app.workflow_id = workflow.id - sqlite_session.commit() - - # Act: a session the model never owns, proving the lookup does not reach for a global one - with Session(sqlite_session.get_bind(), expire_on_commit=False) as caller_session: - result = app.workflow_with_session(session=caller_session) - - # Assert - assert result is not None - assert result.id == workflow.id - assert result in caller_session - - @pytest.mark.parametrize("sqlite_session", [(App, AppModelConfig)], indirect=True) - def test_workflow_with_session_returns_none_when_unpublished(self, sqlite_session: Session): - """An app without `workflow_id` resolves to None without querying.""" - # Arrange - app = App( - tenant_id=str(uuid4()), - name="Test App", - mode=AppMode.WORKFLOW, - enable_site=True, - enable_api=False, - created_by=str(uuid4()), - ) - sqlite_session.add(app) - sqlite_session.flush() - - # Act / Assert: an unbound session would raise if the guard were dropped - with Session(expire_on_commit=False) as unbound_session: - assert app.workflow_with_session(session=unbound_session) is None - - @pytest.mark.parametrize("sqlite_session", [(App, AppModelConfig)], indirect=True) - def test_deleted_tools_checks_plugin_builtin_providers_through_core_plugin_service(self, sqlite_session: Session): - """Plugin-backed built-in tools are checked through core PluginService.""" - # Arrange - app = App( - tenant_id="tenant-1", - name="Test App", - mode=AppMode.CHAT, - enable_site=True, - enable_api=False, - created_by=str(uuid4()), - ) - sqlite_session.add(app) - sqlite_session.flush() - app_model_config = AppModelConfig( - app_id=app.id, - agent_mode=json.dumps( - { - "enabled": True, - "strategy": "function_call", - "tools": [ - { - "provider_type": "builtin", - "provider_id": "langgenius/openai/openai", - "tool_name": "chat", - "tool_parameters": {}, - } - ], - "prompt": None, - } - ), - ) - sqlite_session.add(app_model_config) - sqlite_session.flush() - app.app_model_config_id = app_model_config.id - sqlite_session.flush() - - # Act - with ( - patch("core.tools.tool_manager.ToolManager.get_hardcoded_provider", side_effect=Exception), - patch("core.plugin.plugin_service.PluginService.check_tools_existence", return_value=[False]) as exists, - ): - result = app.deleted_tools_with_session(session=sqlite_session) - - # Assert - assert result == [{"type": "builtin", "tool_name": "chat", "provider_id": "langgenius/openai/openai"}] - exists.assert_called_once() - assert exists.call_args.args[0] == "tenant-1" - assert [str(provider_id) for provider_id in exists.call_args.args[1]] == ["langgenius/openai/openai"] - - -class TestAppModelConfig: - """Test suite for AppModelConfig model.""" - - def test_legacy_dataset_configs_default_retrieval_model_without_dropping_settings(self): - config = AppModelConfig( - app_id=str(uuid4()), - dataset_configs=json.dumps({"reranking_enable": True, "top_k": 8}), - ) - - assert config.dataset_configs_dict == { - "retrieval_model": "single", - "reranking_enable": True, - "top_k": 8, - } - - def test_app_model_config_creation(self): - """Test creating an AppModelConfig.""" - # Arrange - app_id = str(uuid4()) - created_by = str(uuid4()) - - # Act - config = AppModelConfig( - app_id=app_id, - provider="openai", - model_id="gpt-4", - created_by=created_by, - ) - - # Assert - assert config.app_id == app_id - assert config.provider == "openai" - assert config.model_id == "gpt-4" - assert config.created_by == created_by - - def test_app_model_config_with_configs_json(self): - """Test AppModelConfig with JSON configs.""" - # Arrange - configs = {"temperature": 0.7, "max_tokens": 1000} - - # Act - config = AppModelConfig( - app_id=str(uuid4()), - provider="openai", - model_id="gpt-4", - created_by=str(uuid4()), - configs=configs, - ) - - # Assert - assert config.configs == configs - - def test_app_model_config_model_dict_property(self): - """Test model_dict property.""" - # Arrange - model_data = {"provider": "openai", "name": "gpt-4"} - config = AppModelConfig( - app_id=str(uuid4()), - provider="openai", - model_id="gpt-4", - created_by=str(uuid4()), - model=json.dumps(model_data), - ) - - # Act - result = config.model_dict - - # Assert - assert result == model_data - - def test_app_model_config_model_dict_empty(self): - """Test model_dict property when model is None.""" - # Arrange - config = AppModelConfig( - app_id=str(uuid4()), - provider="openai", - model_id="gpt-4", - created_by=str(uuid4()), - model=None, - ) - - # Act - result = config.model_dict - - # Assert - assert result == {} - - def test_app_model_config_suggested_questions_list(self): - """Test suggested_questions_list property.""" - # Arrange - questions = ["What can you do?", "How does this work?"] - config = AppModelConfig( - app_id=str(uuid4()), - provider="openai", - model_id="gpt-4", - created_by=str(uuid4()), - suggested_questions=json.dumps(questions), - ) - - # Act - result = config.suggested_questions_list - - # Assert - assert result == questions - - def test_to_dict_uses_injected_annotation_reply(self): - config = AppModelConfig(app_id=str(uuid4())) - annotation_reply: AnnotationReplyConfig = {"enabled": False} - - result = config.to_dict(annotation_reply=annotation_reply) - - assert result["annotation_reply"] == annotation_reply - - -class TestAnnotationReplyConfigLoader: - @pytest.mark.parametrize("sqlite_session", [(AppAnnotationSetting, DatasetCollectionBinding)], indirect=True) - def test_load_annotation_reply_config_returns_disabled_when_setting_missing(self, sqlite_session: Session): - binding = DatasetCollectionBinding( - provider_name="decoy-provider", - model_name="decoy-model", - type=CollectionBindingType.ANNOTATION, - collection_name="decoy-collection", - ) - sqlite_session.add(binding) - sqlite_session.flush() - sqlite_session.add( - AppAnnotationSetting( - app_id="other-app", - score_threshold=0.9, - collection_binding_id=binding.id, - created_user_id="user-1", - updated_user_id="user-1", - ) - ) - sqlite_session.flush() - - result = load_annotation_reply_config(sqlite_session, "app-1") - - assert result == {"enabled": False} - - @pytest.mark.parametrize("sqlite_session", [(AppAnnotationSetting, DatasetCollectionBinding)], indirect=True) - def test_load_annotation_reply_config_returns_embedding_model(self, sqlite_session: Session): - collection_binding = DatasetCollectionBinding( - provider_name="provider", - model_name="embedding", - type=CollectionBindingType.ANNOTATION, - collection_name="annotation-collection", - ) - sqlite_session.add(collection_binding) - sqlite_session.flush() - annotation_setting = AppAnnotationSetting( - app_id="app-1", - score_threshold=0.7, - collection_binding_id=collection_binding.id, - created_user_id="user-1", - updated_user_id="user-1", - ) - sqlite_session.add(annotation_setting) - sqlite_session.flush() - - result = load_annotation_reply_config(sqlite_session, "app-1") - - assert result == { - "id": annotation_setting.id, - "enabled": True, - "score_threshold": 0.7, - "embedding_model": { - "embedding_provider_name": "provider", - "embedding_model_name": "embedding", - }, - } - - @pytest.mark.parametrize("sqlite_session", [(AppAnnotationSetting, DatasetCollectionBinding)], indirect=True) - def test_load_annotation_reply_config_raises_when_binding_missing(self, sqlite_session: Session): - sqlite_session.add( - AppAnnotationSetting( - app_id="app-1", - score_threshold=0.7, - collection_binding_id=str(uuid4()), - created_user_id="user-1", - updated_user_id="user-1", - ) - ) - sqlite_session.flush() - - with pytest.raises(ValueError, match="Collection binding detail not found"): - load_annotation_reply_config(sqlite_session, "app-1") - - -class TestConversationModel: - """Test suite for Conversation model integrity.""" - - def test_app_returns_matching_app_from_caller_session(self, sqlite_session: Session) -> None: - other_app = App( - tenant_id=str(uuid4()), name="Other app", mode=AppMode.CHAT, enable_site=False, enable_api=False - ) - app = App( - tenant_id=str(uuid4()), name="Conversation app", mode=AppMode.CHAT, enable_site=False, enable_api=False - ) - sqlite_session.add_all([other_app, app]) - sqlite_session.flush() - conversation = Conversation(app_id=app.id) - - assert conversation.app(session=sqlite_session) is app - - def test_app_returns_none_when_app_is_missing(self, sqlite_session: Session) -> None: - other_app = App( - tenant_id=str(uuid4()), name="Other app", mode=AppMode.CHAT, enable_site=False, enable_api=False - ) - sqlite_session.add(other_app) - sqlite_session.flush() - conversation = Conversation(app_id=str(uuid4())) - - assert conversation.app(session=sqlite_session) is None - - def test_conversation_creation_with_required_fields(self): - """Test creating a conversation with required fields.""" - # Arrange - app_id = str(uuid4()) - from_end_user_id = str(uuid4()) - - # Act - conversation = Conversation( - app_id=app_id, - mode=AppMode.CHAT, - name="Test Conversation", - status="normal", - from_source=ConversationFromSource.API, - from_end_user_id=from_end_user_id, - ) - - # Assert - assert conversation.app_id == app_id - assert conversation.mode == AppMode.CHAT - assert conversation.name == "Test Conversation" - assert conversation.status == "normal" - assert conversation.from_source == "api" - assert conversation.from_end_user_id == from_end_user_id - - def test_conversation_with_inputs(self, sqlite_session: Session): - """Test conversation inputs round-trip through the session-aware accessor.""" - # Arrange - inputs = {"query": "Hello", "context": "test"} - conversation = Conversation( - app_id=str(uuid4()), - mode=AppMode.CHAT, - name="Test Conversation", - status="normal", - from_source=ConversationFromSource.API, - from_end_user_id=str(uuid4()), - ) - conversation._inputs = inputs - - # Act - result = conversation.inputs_with_session(session=sqlite_session) - - # Assert - assert result == inputs - - def test_conversation_inputs_is_write_only(self): - """Reading conversation.inputs must fail loudly; reads go through inputs_with_session.""" - conversation = Conversation(app_id=str(uuid4()), _inputs={}) - - with pytest.raises(AttributeError, match="no getter"): - _ = conversation.inputs - - def test_conversation_inputs_setter(self): - """Test conversation inputs setter.""" - # Arrange - conversation = Conversation( - app_id=str(uuid4()), - mode=AppMode.CHAT, - name="Test Conversation", - status="normal", - from_source=ConversationFromSource.API, - from_end_user_id=str(uuid4()), - ) - inputs = {"query": "Hello", "context": "test"} - - # Act - conversation.inputs = inputs - - # Assert - assert conversation._inputs == inputs - - @pytest.mark.parametrize("sqlite_session", [(Conversation, Message)], indirect=True) - def test_conversation_summary_or_query_with_summary(self, sqlite_session: Session): - """Test summary_or_query_with_session when summary exists.""" - # Arrange - conversation = Conversation( - app_id=str(uuid4()), - mode=AppMode.CHAT, - name="Test Conversation", - status="normal", - from_source=ConversationFromSource.API, - from_end_user_id=str(uuid4()), - summary="Test summary", - ) - - # Act - result = conversation.summary_or_query_with_session(session=sqlite_session) - - # Assert - assert result == "Test summary" - - @pytest.mark.parametrize("sqlite_session", [(Conversation, Message)], indirect=True) - def test_conversation_summary_or_query_without_summary(self, sqlite_session: Session): - """Test summary_or_query_with_session when summary is empty.""" - # Arrange - conversation = Conversation( - app_id=str(uuid4()), - mode=AppMode.CHAT, - name="Test Conversation", - status="normal", - from_source=ConversationFromSource.API, - from_end_user_id=str(uuid4()), - summary=None, - ) - conversation._inputs = {} - sqlite_session.add(conversation) - sqlite_session.flush() - first_message = Message( - app_id=conversation.app_id, - conversation_id=conversation.id, - query="First message query", - message={"role": "user", "content": "First message query"}, - answer="First answer", - message_unit_price=Decimal(0), - answer_unit_price=Decimal(0), - currency="USD", - from_source=ConversationFromSource.API, - ) - first_message.created_at = datetime(2024, 1, 1) - first_message._inputs = {} - later_message = Message( - app_id=conversation.app_id, - conversation_id=conversation.id, - query="Later message query", - message={"role": "user", "content": "Later message query"}, - answer="Later answer", - message_unit_price=Decimal(0), - answer_unit_price=Decimal(0), - currency="USD", - from_source=ConversationFromSource.API, - ) - later_message.created_at = datetime(2024, 1, 2) - later_message._inputs = {} - sqlite_session.add_all([later_message, first_message]) - sqlite_session.flush() - - # Act - result = conversation.summary_or_query_with_session(session=sqlite_session) - - # Assert - assert result == "First message query" - - @pytest.mark.parametrize("sqlite_session", [(Conversation, AppModelConfig, AppAnnotationSetting)], indirect=True) - def test_model_config_uses_caller_session_for_annotation_reply(self, sqlite_session: Session): - app_id = str(uuid4()) - app_model_config = AppModelConfig(app_id=app_id, pre_prompt="Persisted prompt") - sqlite_session.add(app_model_config) - sqlite_session.flush() - conversation = Conversation( - app_id=app_id, - app_model_config_id=app_model_config.id, - mode=AppMode.CHAT, - name="Test Conversation", - status="normal", - from_source=ConversationFromSource.API, - from_end_user_id=str(uuid4()), - model_id="model-1", - model_provider="provider-1", - ) - conversation._inputs = {} - sqlite_session.add(conversation) - sqlite_session.flush() - - result = conversation.model_config_with_session(session=sqlite_session) - - assert result["annotation_reply"] == {"enabled": False} - assert result["pre_prompt"] == "Persisted prompt" - assert result["model_id"] == "model-1" - assert result["provider"] == "provider-1" - - @pytest.mark.parametrize("sqlite_session", [(Conversation, AppAnnotationSetting)], indirect=True) - def test_override_model_config_uses_caller_session_for_annotation_reply(self, sqlite_session: Session): - app_id = str(uuid4()) - conversation = Conversation( - app_id=app_id, - mode=AppMode.CHAT, - name="Test Conversation", - status="normal", - from_source=ConversationFromSource.API, - from_end_user_id=str(uuid4()), - override_model_configs=json.dumps({"model": {"provider": "openai", "name": "gpt-4"}}), - model_id="model-1", - model_provider="provider-1", - ) - conversation._inputs = {} - sqlite_session.add(conversation) - sqlite_session.flush() - - result = conversation.model_config_with_session(session=sqlite_session) - - assert result["annotation_reply"] == {"enabled": False} - assert result["model"] == {"provider": "openai", "name": "gpt-4"} - assert result["model_id"] == "model-1" - assert result["provider"] == "provider-1" - - def test_conversation_in_debug_mode(self): - """Test in_debug_mode property.""" - # Arrange - conversation = Conversation( - app_id=str(uuid4()), - mode=AppMode.CHAT, - name="Test Conversation", - status="normal", - from_source=ConversationFromSource.API, - from_end_user_id=str(uuid4()), - override_model_configs='{"model": "gpt-4"}', - ) - - # Act - result = conversation.in_debug_mode - - # Assert - assert result is True - - @pytest.mark.parametrize("sqlite_session", [(Conversation,)], indirect=True) - def test_conversation_to_dict_serialization(self, sqlite_session: Session): - """Test conversation to_dict method.""" - # Arrange - app_id = str(uuid4()) - from_end_user_id = str(uuid4()) - conversation = Conversation( - app_id=app_id, - mode=AppMode.CHAT, - name="Test Conversation", - status="normal", - from_source=ConversationFromSource.API, - from_end_user_id=from_end_user_id, - dialogue_count=5, - ) - conversation.id = str(uuid4()) - conversation._inputs = {"query": "test"} - - # Act - result = conversation.to_dict(session=sqlite_session) - - # Assert - assert result["id"] == conversation.id - assert result["app_id"] == app_id - assert result["mode"] == AppMode.CHAT - assert result["name"] == "Test Conversation" - assert result["status"] == "normal" - assert result["from_source"] == "api" - assert result["from_end_user_id"] == from_end_user_id - assert result["dialogue_count"] == 5 - assert result["inputs"] == {"query": "test"} - - -class TestMessageModel: - """Test suite for Message model and App-Message relationships.""" - - def test_message_creation_with_required_fields(self): - """Test creating a message with required fields.""" - # Arrange - app_id = str(uuid4()) - conversation_id = str(uuid4()) - - # Act - message = Message( - app_id=app_id, - conversation_id=conversation_id, - query="What is AI?", - message={"role": "user", "content": "What is AI?"}, - answer="AI stands for Artificial Intelligence.", - message_unit_price=Decimal("0.0001"), - answer_unit_price=Decimal("0.0002"), - currency="USD", - from_source=ConversationFromSource.API, - ) - - # Assert - assert message.app_id == app_id - assert message.conversation_id == conversation_id - assert message.query == "What is AI?" - assert message.answer == "AI stands for Artificial Intelligence." - assert message.currency == "USD" - assert message.from_source == "api" - - def test_message_with_inputs(self, sqlite_session: Session): - """Test message inputs round-trip through the session-aware accessor.""" - # Arrange - inputs = {"query": "Hello", "context": "test"} - message = Message( - app_id=str(uuid4()), - conversation_id=str(uuid4()), - query="Test query", - message={"role": "user", "content": "Test"}, - answer="Test answer", - message_unit_price=Decimal("0.0001"), - answer_unit_price=Decimal("0.0002"), - currency="USD", - from_source=ConversationFromSource.API, - _inputs=inputs, - ) - - # Act - result = message.inputs_with_session(session=sqlite_session) - - # Assert - assert result == inputs - - def test_message_inputs_is_write_only(self): - """Reading message.inputs must fail loudly; reads go through inputs_with_session.""" - message = Message(app_id=str(uuid4()), _inputs={}) - - with pytest.raises(AttributeError, match="no getter"): - _ = message.inputs - - def test_message_inputs_setter(self): - """Test message inputs setter.""" - # Arrange - message = Message( - app_id=str(uuid4()), - conversation_id=str(uuid4()), - query="Test query", - message={"role": "user", "content": "Test"}, - answer="Test answer", - message_unit_price=Decimal("0.0001"), - answer_unit_price=Decimal("0.0002"), - currency="USD", - from_source=ConversationFromSource.API, - ) - inputs = {"query": "Hello", "context": "test"} - - # Act - message.inputs = inputs - - # Assert - assert message._inputs == inputs - - def test_message_in_debug_mode(self): - """Test message in_debug_mode property.""" - # Arrange - message = Message( - app_id=str(uuid4()), - conversation_id=str(uuid4()), - query="Test query", - message={"role": "user", "content": "Test"}, - answer="Test answer", - message_unit_price=Decimal("0.0001"), - answer_unit_price=Decimal("0.0002"), - currency="USD", - from_source=ConversationFromSource.API, - override_model_configs='{"model": "gpt-4"}', - ) - - # Act - result = message.in_debug_mode - - # Assert - assert result is True - - def test_message_metadata_dict_property(self): - """Test message_metadata_dict property.""" - # Arrange - metadata = {"retriever_resources": ["doc1", "doc2"], "usage": {"tokens": 100}} - message = Message( - app_id=str(uuid4()), - conversation_id=str(uuid4()), - query="Test query", - message={"role": "user", "content": "Test"}, - answer="Test answer", - message_unit_price=Decimal("0.0001"), - answer_unit_price=Decimal("0.0002"), - currency="USD", - from_source=ConversationFromSource.API, - message_metadata=json.dumps(metadata), - ) - - # Act - result = message.message_metadata_dict - - # Assert - assert result == metadata - - def test_message_metadata_dict_empty(self): - """Test message_metadata_dict when metadata is None.""" - # Arrange - message = Message( - app_id=str(uuid4()), - conversation_id=str(uuid4()), - query="Test query", - message={"role": "user", "content": "Test"}, - answer="Test answer", - message_unit_price=Decimal("0.0001"), - answer_unit_price=Decimal("0.0002"), - currency="USD", - from_source=ConversationFromSource.API, - message_metadata=None, - ) - - # Act - result = message.message_metadata_dict - - # Assert - assert result == {} - - @pytest.mark.parametrize("sqlite_session", [(Message,)], indirect=True) - def test_message_to_dict_serialization(self, sqlite_session: Session): - """Test message to_dict method.""" - # Arrange - app_id = str(uuid4()) - conversation_id = str(uuid4()) - now = datetime.now(UTC) - - message = Message( - app_id=app_id, - conversation_id=conversation_id, - query="Test query", - message={"role": "user", "content": "Test"}, - answer="Test answer", - message_unit_price=Decimal("0.0001"), - answer_unit_price=Decimal("0.0002"), - total_price=Decimal("0.0003"), - currency="USD", - from_source=ConversationFromSource.API, - status="normal", - id=str(uuid4()), - _inputs={"query": "test"}, - ) - message.created_at = now - message.updated_at = now - - # Act - result = message.to_dict(session=sqlite_session) - - # Assert - assert result["id"] == message.id - assert result["app_id"] == app_id - assert result["conversation_id"] == conversation_id - assert result["query"] == "Test query" - assert result["answer"] == "Test answer" - assert result["status"] == "normal" - assert result["from_source"] == "api" - assert result["inputs"] == {"query": "test"} - assert "created_at" in result - assert "updated_at" in result - - def test_message_from_dict_deserialization(self): - """Test message from_dict method.""" - # Arrange - message_id = str(uuid4()) - app_id = str(uuid4()) - conversation_id = str(uuid4()) - data = { - "id": message_id, - "app_id": app_id, - "conversation_id": conversation_id, - "model_id": "gpt-4", - "inputs": {"query": "test"}, - "query": "Test query", - "message": {"role": "user", "content": "Test"}, - "answer": "Test answer", - "total_price": Decimal("0.0003"), - "status": "normal", - "error": None, - "message_metadata": {"usage": {"tokens": 100}}, - "from_source": "api", - "from_end_user_id": None, - "from_account_id": None, - "created_at": "2024-01-01T00:00:00", - "updated_at": "2024-01-01T00:00:00", - "agent_based": False, - "workflow_run_id": None, - } - - # Act - message = Message.from_dict(data) - - # Assert - assert message.id == message_id - assert message.app_id == app_id - assert message.conversation_id == conversation_id - assert message.query == "Test query" - assert message.answer == "Test answer" - - -class TestMessageAnnotation: - """Test suite for MessageAnnotation and annotation relationships.""" - - def test_message_annotation_creation(self): - """Test creating a message annotation.""" - # Arrange - app_id = str(uuid4()) - conversation_id = str(uuid4()) - message_id = str(uuid4()) - account_id = str(uuid4()) - - # Act - annotation = MessageAnnotation( - app_id=app_id, - conversation_id=conversation_id, - message_id=message_id, - question="What is AI?", - content="AI stands for Artificial Intelligence.", - account_id=account_id, - ) - - # Assert - assert annotation.app_id == app_id - assert annotation.conversation_id == conversation_id - assert annotation.message_id == message_id - assert annotation.question == "What is AI?" - assert annotation.content == "AI stands for Artificial Intelligence." - assert annotation.account_id == account_id - - def test_message_annotation_without_message_id(self): - """Test creating annotation without message_id.""" - # Arrange - app_id = str(uuid4()) - account_id = str(uuid4()) - - # Act - annotation = MessageAnnotation( - app_id=app_id, - question="What is AI?", - conversation_id=None, - message_id=None, - content="AI stands for Artificial Intelligence.", - account_id=account_id, - ) - - # Assert - assert annotation.app_id == app_id - assert annotation.message_id is None - assert annotation.conversation_id is None - assert annotation.question == "What is AI?" - assert annotation.content == "AI stands for Artificial Intelligence." - - def test_message_annotation_hit_count_default(self): - """Test annotation hit_count default value.""" - # Arrange - annotation = MessageAnnotation( - app_id=str(uuid4()), - question="Test question", - conversation_id=None, - message_id=None, - content="Test content", - account_id=str(uuid4()), - ) - - # Act & Assert - default value is set by database - # Model instantiation doesn't set server defaults - assert hasattr(annotation, "hit_count") - - -class TestAppAnnotationSetting: - """Test suite for AppAnnotationSetting model.""" - - def test_app_annotation_setting_creation(self): - """Test creating an app annotation setting.""" - # Arrange - app_id = str(uuid4()) - collection_binding_id = str(uuid4()) - created_user_id = str(uuid4()) - updated_user_id = str(uuid4()) - - # Act - setting = AppAnnotationSetting( - app_id=app_id, - score_threshold=0.8, - collection_binding_id=collection_binding_id, - created_user_id=created_user_id, - updated_user_id=updated_user_id, - ) - - # Assert - assert setting.app_id == app_id - assert setting.score_threshold == 0.8 - assert setting.collection_binding_id == collection_binding_id - assert setting.created_user_id == created_user_id - assert setting.updated_user_id == updated_user_id - - def test_app_annotation_setting_score_threshold_validation(self): - """Test score threshold values.""" - # Arrange & Act - setting_high = AppAnnotationSetting( - app_id=str(uuid4()), - score_threshold=0.95, - collection_binding_id=str(uuid4()), - created_user_id=str(uuid4()), - updated_user_id=str(uuid4()), - ) - setting_low = AppAnnotationSetting( - app_id=str(uuid4()), - score_threshold=0.5, - collection_binding_id=str(uuid4()), - created_user_id=str(uuid4()), - updated_user_id=str(uuid4()), - ) - - # Assert - assert setting_high.score_threshold == 0.95 - assert setting_low.score_threshold == 0.5 - - -class TestAppAnnotationHitHistory: - """Test suite for AppAnnotationHitHistory model.""" - - def test_app_annotation_hit_history_creation(self): - """Test creating an annotation hit history.""" - # Arrange - app_id = str(uuid4()) - annotation_id = str(uuid4()) - message_id = str(uuid4()) - account_id = str(uuid4()) - - # Act - history = AppAnnotationHitHistory( - app_id=app_id, - annotation_id=annotation_id, - source="api", - question="What is AI?", - account_id=account_id, - score=0.95, - message_id=message_id, - annotation_question="What is AI?", - annotation_content="AI stands for Artificial Intelligence.", - ) - - # Assert - assert history.app_id == app_id - assert history.annotation_id == annotation_id - assert history.source == "api" - assert history.question == "What is AI?" - assert history.account_id == account_id - assert history.score == 0.95 - assert history.message_id == message_id - assert history.annotation_question == "What is AI?" - assert history.annotation_content == "AI stands for Artificial Intelligence." - - def test_app_annotation_hit_history_score_values(self): - """Test annotation hit history with different score values.""" - # Arrange & Act - history_high = AppAnnotationHitHistory( - app_id=str(uuid4()), - annotation_id=str(uuid4()), - source="api", - question="Test", - account_id=str(uuid4()), - score=0.99, - message_id=str(uuid4()), - annotation_question="Test", - annotation_content="Content", - ) - history_low = AppAnnotationHitHistory( - app_id=str(uuid4()), - annotation_id=str(uuid4()), - source="api", - question="Test", - account_id=str(uuid4()), - score=0.6, - message_id=str(uuid4()), - annotation_question="Test", - annotation_content="Content", - ) - - # Assert - assert history_high.score == 0.99 - assert history_low.score == 0.6 - - -class TestSiteModel: - """Test suite for Site model.""" - - def test_site_core_bulk_insert_generates_ids(self, sqlite_session: Session): - """Core bulk inserts generate an ID for every site row.""" - app_ids = [str(uuid4()), str(uuid4())] - - sqlite_session.execute( - Site.__table__.insert(), - [ - { - "app_id": app_id, - "title": f"Site {index}", - "default_language": "en-US", - "customize_token_strategy": CustomizeTokenStrategy.UUID, - } - for index, app_id in enumerate(app_ids) - ], - ) - - site_ids = sqlite_session.scalars(select(Site.id).where(Site.app_id.in_(app_ids))).all() - assert len(site_ids) == len(app_ids) - assert len(set(site_ids)) == len(app_ids) - assert all(UUID(site_id).version == 4 for site_id in site_ids) - - def test_site_creation_with_required_fields(self): - """Test creating a site with required fields.""" - # Arrange - app_id = str(uuid4()) - - # Act - site = Site( - app_id=app_id, - title="Test Site", - default_language="en-US", - customize_token_strategy=CustomizeTokenStrategy.UUID, - ) - - # Assert - assert site.app_id == app_id - assert site.title == "Test Site" - assert site.default_language == "en-US" - assert site.customize_token_strategy == CustomizeTokenStrategy.UUID - - def test_site_creation_with_optional_fields(self): - """Test creating a site with optional fields.""" - # Arrange & Act - site = Site( - app_id=str(uuid4()), - title="Test Site", - default_language="en-US", - customize_token_strategy=CustomizeTokenStrategy.UUID, - icon_type=IconType.EMOJI, - icon="🌐", - icon_background="#0066CC", - description="Test site description", - copyright="© 2024 Test", - privacy_policy="https://example.com/privacy", - ) - - # Assert - assert site.icon_type == IconType.EMOJI - assert site.icon == "🌐" - assert site.icon_background == "#0066CC" - assert site.description == "Test site description" - assert site.copyright == "© 2024 Test" - assert site.privacy_policy == "https://example.com/privacy" - - def test_site_custom_disclaimer_setter(self): - """Test site custom_disclaimer setter.""" - # Arrange - site = Site( - app_id=str(uuid4()), - title="Test Site", - default_language="en-US", - customize_token_strategy=CustomizeTokenStrategy.UUID, - ) - - # Act - site.custom_disclaimer = "This is a test disclaimer" - - # Assert - assert site.custom_disclaimer == "This is a test disclaimer" - - def test_site_custom_disclaimer_exceeds_limit(self): - """Test site custom_disclaimer with excessive length.""" - # Arrange - site = Site( - app_id=str(uuid4()), - title="Test Site", - default_language="en-US", - customize_token_strategy=CustomizeTokenStrategy.UUID, - ) - long_disclaimer = "x" * 513 # Exceeds 512 character limit - - # Act & Assert - with pytest.raises(ValueError, match="Custom disclaimer cannot exceed 512 characters"): - site.custom_disclaimer = long_disclaimer - - -class TestModelIntegration: - """Test suite for model integration scenarios.""" - - def test_complete_app_conversation_message_hierarchy(self): - """Test complete hierarchy from app to message.""" - # Arrange - tenant_id = str(uuid4()) - app_id = str(uuid4()) - conversation_id = str(uuid4()) - message_id = str(uuid4()) - created_by = str(uuid4()) - - # Create app - app = App( - tenant_id=tenant_id, - name="Test App", - mode=AppMode.CHAT, - enable_site=True, - enable_api=True, - created_by=created_by, - id=app_id, - ) - - # Create conversation - conversation = Conversation( - app_id=app_id, - mode=AppMode.CHAT, - name="Test Conversation", - status="normal", - from_source=ConversationFromSource.API, - from_end_user_id=str(uuid4()), - ) - conversation.id = conversation_id - - # Create message - message = Message( - app_id=app_id, - conversation_id=conversation_id, - query="Test query", - message={"role": "user", "content": "Test"}, - answer="Test answer", - message_unit_price=Decimal("0.0001"), - answer_unit_price=Decimal("0.0002"), - currency="USD", - from_source=ConversationFromSource.API, - id=message_id, - ) - - # Assert - assert app.id == app_id - assert conversation.app_id == app_id - assert message.app_id == app_id - assert message.conversation_id == conversation_id - assert app.mode == AppMode.CHAT - assert conversation.mode == AppMode.CHAT - - def test_app_with_annotation_setting(self): - """Test app with annotation setting.""" - # Arrange - app_id = str(uuid4()) - collection_binding_id = str(uuid4()) - created_user_id = str(uuid4()) - - # Create app - app = App( - tenant_id=str(uuid4()), - name="Test App", - mode=AppMode.CHAT, - enable_site=True, - enable_api=True, - created_by=created_user_id, - id=app_id, - ) - - # Create annotation setting - setting = AppAnnotationSetting( - app_id=app_id, - score_threshold=0.85, - collection_binding_id=collection_binding_id, - created_user_id=created_user_id, - updated_user_id=created_user_id, - ) - - # Assert - assert setting.app_id == app.id - assert setting.score_threshold == 0.85 - - def test_message_with_annotation(self): - """Test message with annotation.""" - # Arrange - app_id = str(uuid4()) - conversation_id = str(uuid4()) - message_id = str(uuid4()) - account_id = str(uuid4()) - - # Create message - message = Message( - app_id=app_id, - conversation_id=conversation_id, - query="What is AI?", - message={"role": "user", "content": "What is AI?"}, - answer="AI stands for Artificial Intelligence.", - message_unit_price=Decimal("0.0001"), - answer_unit_price=Decimal("0.0002"), - currency="USD", - from_source=ConversationFromSource.API, - id=message_id, - ) - - # Create annotation - annotation = MessageAnnotation( - app_id=app_id, - conversation_id=conversation_id, - message_id=message_id, - question="What is AI?", - content="AI stands for Artificial Intelligence.", - account_id=account_id, - ) - - # Assert - assert annotation.app_id == message.app_id - assert annotation.conversation_id == message.conversation_id - assert annotation.message_id == message.id - - def test_annotation_hit_history_tracking(self): - """Test annotation hit history tracking.""" - # Arrange - app_id = str(uuid4()) - annotation_id = str(uuid4()) - message_id = str(uuid4()) - account_id = str(uuid4()) - - # Create annotation - annotation = MessageAnnotation( - app_id=app_id, - question="What is AI?", - content="AI stands for Artificial Intelligence.", - conversation_id=None, - message_id=message_id, - account_id=account_id, - ) - annotation.id = annotation_id - - # Create hit history - history = AppAnnotationHitHistory( - app_id=app_id, - annotation_id=annotation_id, - source="api", - question="What is AI?", - account_id=account_id, - score=0.92, - message_id=message_id, - annotation_question="What is AI?", - annotation_content="AI stands for Artificial Intelligence.", - ) - - # Assert - assert history.app_id == annotation.app_id - assert history.annotation_id == annotation.id - assert history.score == 0.92 - - def test_app_with_site(self): - """Test app with site.""" - # Arrange - app_id = str(uuid4()) - - # Create app - app = App( - tenant_id=str(uuid4()), - name="Test App", - mode=AppMode.CHAT, - enable_site=True, - enable_api=True, - created_by=str(uuid4()), - id=app_id, - ) - - # Create site - site = Site( - app_id=app_id, - title="Test Site", - default_language="en-US", - customize_token_strategy=CustomizeTokenStrategy.UUID, - ) - - # Assert - assert site.app_id == app.id - assert app.enable_site is True diff --git a/api/tests/unit_tests/models/test_app_session_accessors.py b/api/tests/unit_tests/models/test_app_session_accessors.py deleted file mode 100644 index bf590e280d6..00000000000 --- a/api/tests/unit_tests/models/test_app_session_accessors.py +++ /dev/null @@ -1,50 +0,0 @@ -"""Regression coverage for the ``@property``→session-parameter refactor on ``App.tenant``. - -The legacy ``@property`` reached for the global ``db.session`` internally and has been converted -to a plain method taking an explicit ``session`` (per the pattern established in -#41830/#41831 for ``InstalledApp.tenant`` / ``ApiToolProvider.tenant``, tracked in #40372). -""" - -from uuid import uuid4 - -from sqlalchemy.orm import Session - -from models.account import Tenant -from models.model import App, AppMode - - -def _persist_tenant(session: Session) -> Tenant: - tenant = Tenant(name="Test Tenant") - session.add(tenant) - session.flush() - return tenant - - -def _persist_app(session: Session, *, tenant_id: str) -> App: - app = App( - tenant_id=tenant_id, - name="Test App", - mode=AppMode.CHAT, - enable_site=True, - enable_api=False, - created_by=str(uuid4()), - ) - session.add(app) - session.flush() - return app - - -class TestAppTenant: - def test_returns_the_matching_tenant(self, sqlite_session: Session) -> None: - tenant = _persist_tenant(sqlite_session) - app = _persist_app(sqlite_session, tenant_id=tenant.id) - - result = app.tenant(session=sqlite_session) - - assert result is not None - assert result.id == tenant.id - - def test_returns_none_when_tenant_missing(self, sqlite_session: Session) -> None: - app = _persist_app(sqlite_session, tenant_id=str(uuid4())) - - assert app.tenant(session=sqlite_session) is None diff --git a/api/tests/unit_tests/models/test_conversation_variable.py b/api/tests/unit_tests/models/test_conversation_variable.py deleted file mode 100644 index 6ec147ff4da..00000000000 --- a/api/tests/unit_tests/models/test_conversation_variable.py +++ /dev/null @@ -1,45 +0,0 @@ -from uuid import NAMESPACE_URL, UUID, uuid4, uuid5 - -from factories import variable_factory -from graphon.variables import SegmentType -from models import ConversationVariable - - -def test_from_variable_coerces_non_uuid_primary_key() -> None: - variable = variable_factory.build_conversation_variable_from_mapping( - { - "id": "opt-comp-prompt-var", - "name": "optimization_comparison_prompt", - "value_type": SegmentType.STRING, - "value": "-", - } - ) - - row = ConversationVariable.from_variable(app_id="app_id", conversation_id="conversation_id", variable=variable) - - expected = str(uuid5(NAMESPACE_URL, "dify:conversation-variable:optimization_comparison_prompt")) - UUID(row.id) - assert row.id == expected - assert variable.id == "opt-comp-prompt-var" - assert row.to_variable().id == "opt-comp-prompt-var" - - -def test_from_variable_and_to_variable() -> None: - variable = variable_factory.build_conversation_variable_from_mapping( - { - "id": str(uuid4()), - "name": "name", - "value_type": SegmentType.OBJECT, - "value": { - "key": { - "key": "value", - } - }, - } - ) - - conversation_variable = ConversationVariable.from_variable( - app_id="app_id", conversation_id="conversation_id", variable=variable - ) - - assert conversation_variable.to_variable() == variable diff --git a/api/tests/unit_tests/models/test_dataset_models.py b/api/tests/unit_tests/models/test_dataset.py similarity index 99% rename from api/tests/unit_tests/models/test_dataset_models.py rename to api/tests/unit_tests/models/test_dataset.py index 7bccc69fb24..cc25be69258 100644 --- a/api/tests/unit_tests/models/test_dataset_models.py +++ b/api/tests/unit_tests/models/test_dataset.py @@ -79,7 +79,7 @@ from repositories.knowledge.dataset_read_repository import ( from repositories.knowledge.keyword_table_repository import load_keyword_table from repositories.knowledge.segment_read_adapter import get_segment_attachments from services.knowledge.dataset_read_service import get_dataset_queries -from tests.unit_tests.config_override import apply_config_overrides +from tests.helpers.unit.config_override import apply_config_overrides def _make_dataset( diff --git a/api/tests/unit_tests/models/test_end_user_type.py b/api/tests/unit_tests/models/test_enums.py similarity index 88% rename from api/tests/unit_tests/models/test_end_user_type.py rename to api/tests/unit_tests/models/test_enums.py index 7cbd69de4a3..7db7debe122 100644 --- a/api/tests/unit_tests/models/test_end_user_type.py +++ b/api/tests/unit_tests/models/test_enums.py @@ -4,7 +4,7 @@ from pathlib import Path import pytest -from models.enums import EndUserType +from models.enums import CreatorUserRole, EndUserType from models.model import EndUser from models.types import EnumText from services.app_scoped_end_user_service import AppScopedEndUserService @@ -114,3 +114,19 @@ def test_production_end_user_constructors_use_end_user_type_enum() -> None: violations.append(f"{source_path.relative_to(API_ROOT)}:{node.lineno}") assert violations == [] + + +def test_creator_user_role_missing_maps_hyphen_to_enum() -> None: + # given an alias with hyphen + value = "end-user" + + # when converting to enum (invokes StrEnum._missing_ override) + role = CreatorUserRole(value) + + # then it should map to END_USER + assert role is CreatorUserRole.END_USER + + +def test_creator_user_role_missing_raises_for_unknown() -> None: + with pytest.raises(ValueError, match="'unknown' is not a valid CreatorUserRole"): + CreatorUserRole("unknown") diff --git a/api/tests/unit_tests/models/test_enums_creator_user_role.py b/api/tests/unit_tests/models/test_enums_creator_user_role.py deleted file mode 100644 index bdbac3202d7..00000000000 --- a/api/tests/unit_tests/models/test_enums_creator_user_role.py +++ /dev/null @@ -1,19 +0,0 @@ -import pytest - -from models.enums import CreatorUserRole - - -def test_creator_user_role_missing_maps_hyphen_to_enum() -> None: - # given an alias with hyphen - value = "end-user" - - # when converting to enum (invokes StrEnum._missing_ override) - role = CreatorUserRole(value) - - # then it should map to END_USER - assert role is CreatorUserRole.END_USER - - -def test_creator_user_role_missing_raises_for_unknown() -> None: - with pytest.raises(ValueError, match="'unknown' is not a valid CreatorUserRole"): - CreatorUserRole("unknown") diff --git a/api/tests/unit_tests/models/test_recipient_type_label.py b/api/tests/unit_tests/models/test_human_input.py similarity index 100% rename from api/tests/unit_tests/models/test_recipient_type_label.py rename to api/tests/unit_tests/models/test_human_input.py diff --git a/api/tests/unit_tests/models/test_installed_app_session_accessors.py b/api/tests/unit_tests/models/test_installed_app_session_accessors.py deleted file mode 100644 index 0b1f5a04058..00000000000 --- a/api/tests/unit_tests/models/test_installed_app_session_accessors.py +++ /dev/null @@ -1,51 +0,0 @@ -"""Regression coverage for the ``@property``→session-parameter refactor on -``InstalledApp.tenant``. - -The legacy ``@property`` reached for the global ``db.session`` internally and has been converted -to a plain method taking an explicit ``session: Session`` (per the pattern established in -#40370/#40797/#41394/#41830/#41885, tracked in #40372). -""" - -from uuid import uuid4 - -from sqlalchemy.orm import Session - -from models.account import Tenant -from models.model import InstalledApp - - -def _persist_tenant(session: Session) -> Tenant: - tenant = Tenant(name="Test Tenant") - session.add(tenant) - session.flush() - return tenant - - -def _persist_installed_app(session: Session, *, tenant_id: str) -> InstalledApp: - installed_app = InstalledApp( - tenant_id=tenant_id, - app_id=str(uuid4()), - app_owner_tenant_id=str(uuid4()), - position=0, - is_pinned=False, - last_used_at=None, - ) - session.add(installed_app) - session.flush() - return installed_app - - -class TestInstalledAppTenant: - def test_returns_the_matching_tenant(self, sqlite_session: Session) -> None: - tenant = _persist_tenant(sqlite_session) - installed_app = _persist_installed_app(sqlite_session, tenant_id=tenant.id) - - result = installed_app.tenant(session=sqlite_session) - - assert result is not None - assert result.id == tenant.id - - def test_returns_none_when_tenant_missing(self, sqlite_session: Session) -> None: - installed_app = _persist_installed_app(sqlite_session, tenant_id=str(uuid4())) - - assert installed_app.tenant(session=sqlite_session) is None diff --git a/api/tests/unit_tests/models/test_knowledge_fs.py b/api/tests/unit_tests/models/test_knowledge_fs.py index 116888001b4..f32b21a6505 100644 --- a/api/tests/unit_tests/models/test_knowledge_fs.py +++ b/api/tests/unit_tests/models/test_knowledge_fs.py @@ -20,7 +20,7 @@ from models.knowledge_fs import ( KnowledgeFSStagedUpload, KnowledgeFSStagedUploadStatus, ) -from tests.unit_tests.services.knowledge_fs_fakes import reservation_summary +from tests.helpers.unit.services.knowledge_fs_fakes import reservation_summary _KNOWLEDGE_FS_MODELS = ( KnowledgeFSControlSpace, diff --git a/api/tests/unit_tests/models/test_mcp_tool_provider_session_accessors.py b/api/tests/unit_tests/models/test_mcp_tool_provider_session_accessors.py deleted file mode 100644 index 2f1761a80fd..00000000000 --- a/api/tests/unit_tests/models/test_mcp_tool_provider_session_accessors.py +++ /dev/null @@ -1,53 +0,0 @@ -"""Regression coverage for the ``db.session``→session-parameter refactor on -``MCPToolProvider.load_user``. - -The method reached for the global ``db.session`` internally and has been converted to -take an explicit ``session: Session`` (per the pattern established in #40370/#42316, -tracked in #40372). -""" - -from uuid import uuid4 - -from sqlalchemy.orm import Session - -from models.account import Account -from models.tools import MCPToolProvider - - -def _persist_account(session: Session, *, name: str = "MCP Owner") -> Account: - account = Account(name=name, email=f"{uuid4().hex}@example.com") - session.add(account) - session.flush() - return account - - -def _persist_provider(session: Session, *, user_id: str) -> MCPToolProvider: - provider = MCPToolProvider( - name="Test Provider", - server_identifier=f"test-provider-{uuid4().hex[:8]}", - server_url="https://example.com/mcp", - server_url_hash=uuid4().hex, - icon="icon", - tenant_id=str(uuid4()), - user_id=user_id, - ) - session.add(provider) - session.flush() - return provider - - -def test_load_user_returns_account_through_the_given_session(sqlite_session: Session) -> None: - account = _persist_account(sqlite_session) - provider = _persist_provider(sqlite_session, user_id=account.id) - - result = provider.load_user(sqlite_session) - - assert result is not None - assert result.id == account.id - assert result.name == account.name - - -def test_load_user_returns_none_when_account_is_missing(sqlite_session: Session) -> None: - provider = _persist_provider(sqlite_session, user_id=str(uuid4())) - - assert provider.load_user(sqlite_session) is None diff --git a/api/tests/unit_tests/models/test_message_annotation_session_accessors.py b/api/tests/unit_tests/models/test_message_annotation_session_accessors.py deleted file mode 100644 index 6e2de3dbed3..00000000000 --- a/api/tests/unit_tests/models/test_message_annotation_session_accessors.py +++ /dev/null @@ -1,106 +0,0 @@ -"""Regression coverage for the ``@property``→session-parameter refactor on Message -feedback/annotation accessors. - -Covers three of `api/models/model.py`'s legacy `@property` accessors that reached for the global -``db.session`` internally and have been removed in favour of their existing explicit-session -methods (per the pattern established in #40370/#40797/#41394/#41830, tracked in #40372): - -- ``MessageFeedback.from_account_with_session`` -- ``MessageAnnotation.account_with_session`` -- ``MessageAnnotation.annotation_create_account_with_session`` - -Each accessor is exercised against the real ``sqlite_session`` fixture (a genuine SQLAlchemy -``Session`` bound to a pristine full-schema SQLite database) so the assertions cover actual -query behaviour rather than a mock's recorded call. -""" - -from uuid import uuid4 - -from sqlalchemy.orm import Session - -from models.account import Account -from models.enums import FeedbackFromSource, FeedbackRating -from models.model import MessageAnnotation, MessageFeedback - - -def _persist_account(session: Session, *, email: str | None = None) -> Account: - account = Account(name="Test Account", email=email or "test@example.com") - session.add(account) - session.flush() - return account - - -def _persist_feedback(session: Session, *, app_id: str, from_account_id: str | None) -> MessageFeedback: - feedback = MessageFeedback( - app_id=app_id, - conversation_id=str(uuid4()), - message_id=str(uuid4()), - rating=FeedbackRating.LIKE, - from_source=FeedbackFromSource.USER, - from_account_id=from_account_id, - ) - session.add(feedback) - session.flush() - return feedback - - -def _persist_annotation(session: Session, *, app_id: str, account_id: str) -> MessageAnnotation: - annotation = MessageAnnotation( - app_id=app_id, - conversation_id=None, - message_id=None, - question="What is AI?", - content="AI stands for Artificial Intelligence.", - account_id=account_id, - ) - session.add(annotation) - session.flush() - return annotation - - -class TestMessageFeedbackFromAccount: - def test_returns_the_feedback_author(self, sqlite_session: Session) -> None: - account = _persist_account(sqlite_session) - feedback = _persist_feedback(sqlite_session, app_id=str(uuid4()), from_account_id=account.id) - - result = feedback.from_account_with_session(session=sqlite_session) - - assert result is not None - assert result.id == account.id - - def test_returns_none_when_account_missing(self, sqlite_session: Session) -> None: - feedback = _persist_feedback(sqlite_session, app_id=str(uuid4()), from_account_id=str(uuid4())) - - assert feedback.from_account_with_session(session=sqlite_session) is None - - -class TestMessageAnnotationAccount: - def test_returns_the_annotation_creator(self, sqlite_session: Session) -> None: - creator = _persist_account(sqlite_session, email="creator@example.com") - annotation = _persist_annotation(sqlite_session, app_id=str(uuid4()), account_id=creator.id) - - result = annotation.account_with_session(session=sqlite_session) - - assert result is not None - assert result.id == creator.id - - def test_returns_none_when_account_missing(self, sqlite_session: Session) -> None: - annotation = _persist_annotation(sqlite_session, app_id=str(uuid4()), account_id=str(uuid4())) - - assert annotation.account_with_session(session=sqlite_session) is None - - -class TestMessageAnnotationAnnotationCreateAccount: - def test_returns_the_annotation_creator(self, sqlite_session: Session) -> None: - creator = _persist_account(sqlite_session, email="creator2@example.com") - annotation = _persist_annotation(sqlite_session, app_id=str(uuid4()), account_id=creator.id) - - result = annotation.annotation_create_account_with_session(session=sqlite_session) - - assert result is not None - assert result.id == creator.id - - def test_returns_none_when_account_missing(self, sqlite_session: Session) -> None: - annotation = _persist_annotation(sqlite_session, app_id=str(uuid4()), account_id=str(uuid4())) - - assert annotation.annotation_create_account_with_session(session=sqlite_session) is None diff --git a/api/tests/unit_tests/models/test_message_trace_serialization.py b/api/tests/unit_tests/models/test_message_trace_serialization.py deleted file mode 100644 index 755dbf158eb..00000000000 --- a/api/tests/unit_tests/models/test_message_trace_serialization.py +++ /dev/null @@ -1,161 +0,0 @@ -"""Keep the persisted provider through private trace JSON and the real consumer.""" - -import json -from datetime import UTC, datetime, timedelta -from decimal import Decimal -from uuid import uuid4 - -import pytest -from dify_trace_arize_phoenix import arize_phoenix_trace -from dify_trace_arize_phoenix.arize_phoenix_trace import ArizePhoenixDataTrace -from dify_trace_arize_phoenix.config import ArizeConfig -from openinference.semconv.trace import SpanAttributes -from opentelemetry.sdk.trace import Tracer, TracerProvider -from opentelemetry.sdk.trace.export import SimpleSpanProcessor -from opentelemetry.sdk.trace.export.in_memory_span_exporter import InMemorySpanExporter -from pydantic import TypeAdapter -from sqlalchemy.orm import Session - -from core.ops.entities.trace_entity import MessageTraceInfo -from core.ops.ops_trace_manager import OpsTraceManager, TaskData -from models.enums import ConversationFromSource -from models.model import Message, MessageDict -from tasks import ops_trace_task -from tests.unit_tests.config_override import apply_config_overrides -from tests.unit_tests.model_factories import make_message - -_MESSAGE = TypeAdapter(MessageDict) - - -@pytest.fixture -def persisted_message(sqlite_session: Session) -> Message: - now = datetime(2026, 10, 20, tzinfo=UTC) - message = make_message( - message_id=str(uuid4()), - app_id=str(uuid4()), - conversation_id=str(uuid4()), - inputs={}, - query="Synthetic input", - message={}, - answer="Synthetic output", - message_unit_price=Decimal(0), - answer_unit_price=Decimal(0), - total_price=Decimal(0), - currency="USD", - from_source=ConversationFromSource.CONSOLE, - created_at=now, - updated_at=now + timedelta(seconds=1), - ) - message.model_provider = "langgenius/openai/openai" - message.model_id = "gpt-4o-mini" - sqlite_session.add(message) - sqlite_session.commit() - sqlite_session.expire_all() - restored = sqlite_session.get(Message, message.id) - assert restored is not None - return restored - - -@pytest.mark.parametrize("provider", ["langgenius/openai/openai", "Custom/Provider:Exact", None]) -def test_message_provider_survives_real_json_round_trip( - sqlite_session: Session, - persisted_message: Message, - provider: str | None, -) -> None: - persisted_message.model_provider = provider - sqlite_session.commit() - sqlite_session.refresh(persisted_message) - wire = json.dumps(persisted_message.to_dict(session=sqlite_session), default=str) - payload = _MESSAGE.validate_json(wire) - restored = Message.from_dict(payload) - assert "model_provider" in payload - assert restored.model_provider == provider - assert restored.model_id == persisted_message.model_id - assert restored.id == persisted_message.id - - -@pytest.mark.parametrize("shape", ["current", "legacy_missing", "explicit_null"]) -def test_private_trace_task_preserves_provider_into_real_arize_span( - sqlite_session: Session, - persisted_message: Message, - monkeypatch: pytest.MonkeyPatch, - shape: str, -) -> None: - message_data = persisted_message.to_dict(session=sqlite_session) - if shape == "legacy_missing": - message_data.pop("model_provider", None) - elif shape == "explicit_null": - message_data["model_provider"] = None - start = datetime(2026, 10, 20, tzinfo=UTC) - info = MessageTraceInfo( - message_id=persisted_message.id, - message_data=message_data, - conversation_model="chat", - conversation_mode="chat", - inputs={"query": "Synthetic input"}, - outputs="Synthetic output", - start_time=start, - end_time=start + timedelta(seconds=1), - metadata={}, - message_tokens=1, - answer_tokens=1, - total_tokens=2, - ) - encoded = ( - TaskData(app_id=persisted_message.app_id, trace_info_type="MessageTraceInfo", trace_info=info.model_dump()) - .model_dump_json() - .encode() - ) - exporter = InMemorySpanExporter() - provider = TracerProvider() - processor = SimpleSpanProcessor(exporter) - provider.add_span_processor(processor) - - def setup(_config: ArizeConfig) -> tuple[Tracer, SimpleSpanProcessor]: - tracer = provider.get_tracer("message-provider-serialization") - assert isinstance(tracer, Tracer) - return tracer, processor - - monkeypatch.setattr(arize_phoenix_trace, "setup_tracer", setup) - instance = ArizePhoenixDataTrace(ArizeConfig(api_key="unused-synthetic-key", space_id="synthetic-space")) - paths: list[str] = [] - deleted: list[str] = [] - - def load(path: str) -> bytes: - paths.append(path) - return encoded - - def delete(path: str) -> None: - deleted.append(path) - - def lookup(app_id: str) -> ArizePhoenixDataTrace: - assert app_id == persisted_message.app_id - return instance - - def failed(_key: str) -> None: - pytest.fail("The real trace consumer must not hide a failed dispatch") - - apply_config_overrides(monkeypatch, OPS_AGENT_TRACE_ENABLED=False) - monkeypatch.setattr(ops_trace_task.storage, "load", load) - monkeypatch.setattr(ops_trace_task.storage, "delete", delete) - monkeypatch.setattr(ops_trace_task.redis_client, "incr", failed) - monkeypatch.setattr(OpsTraceManager, "get_ops_trace_instance", lookup) - monkeypatch.setattr("extensions.ext_enterprise_telemetry.is_enabled", lambda: False) - try: - ops_trace_task.process_trace_tasks.apply( - args=({"app_id": persisted_message.app_id, "file_id": "owned-synthetic"},), throw=True - ).get() - spans = exporter.get_finished_spans() - llm = [span for span in spans if span.name == "llm"] - assert len(llm) == 1 - attributes = llm[0].attributes - assert attributes is not None - assert attributes[SpanAttributes.LLM_MODEL_NAME] == persisted_message.model_id - if shape == "current": - assert attributes[SpanAttributes.LLM_PROVIDER] == persisted_message.model_provider - else: - assert SpanAttributes.LLM_PROVIDER not in attributes - assert paths == [f"ops_trace/{persisted_message.app_id}/owned-synthetic.json"] - assert deleted == paths - finally: - provider.shutdown() diff --git a/api/tests/unit_tests/models/test_model.py b/api/tests/unit_tests/models/test_model.py index a30ef7a052b..6f8cc296e9a 100644 --- a/api/tests/unit_tests/models/test_model.py +++ b/api/tests/unit_tests/models/test_model.py @@ -1,24 +1,2154 @@ +import json +from datetime import UTC, datetime, timedelta +from decimal import Decimal from unittest.mock import patch +from uuid import UUID, uuid4 import pytest +from dify_trace_arize_phoenix import arize_phoenix_trace +from dify_trace_arize_phoenix.arize_phoenix_trace import ArizePhoenixDataTrace +from dify_trace_arize_phoenix.config import ArizeConfig +from openinference.semconv.trace import SpanAttributes +from opentelemetry.sdk.trace import Tracer, TracerProvider +from opentelemetry.sdk.trace.export import SimpleSpanProcessor +from opentelemetry.sdk.trace.export.in_memory_span_exporter import InMemorySpanExporter +from pydantic import TypeAdapter +from sqlalchemy import inspect, select from sqlalchemy.orm import Session +from core.ops.entities.trace_entity import MessageTraceInfo +from core.ops.ops_trace_manager import OpsTraceManager, TaskData from core.workflow.file_reference import build_file_reference from graphon.file import FILE_MODEL_IDENTITY, FileTransferMethod from models import model as model_module -from models.model import App, AppMode, Conversation, IconType, Message +from models.account import Account, Tenant +from models.base import Base, TypeBase +from models.dataset import DatasetCollectionBinding +from models.enums import ( + ApiTokenType, + CollectionBindingType, + ConversationFromSource, + CustomizeTokenStrategy, + EndUserType, + FeedbackFromSource, + FeedbackRating, +) +from models.model import ( + AccountTrialAppRecord, + AnnotationReplyConfig, + ApiToken, + App, + AppAnnotationHitHistory, + AppAnnotationSetting, + AppMCPServer, + AppMode, + AppModelConfig, + AppStar, + Conversation, + EndUser, + IconType, + InstalledApp, + Message, + MessageAnnotation, + MessageDict, + MessageFeedback, + RecommendedApp, + Site, + load_annotation_reply_config, +) +from models.workflow import Workflow, WorkflowType +from tasks import ops_trace_task +from tests.helpers.unit.config_override import apply_config_overrides +from tests.helpers.unit.model_factories import make_message + +"""Regression coverage for the ``@property``→session-parameter refactor on +``AppAnnotationHitHistory`` accessors. + +Covers two of `api/models/model.py`'s legacy `@property` accessors that reached for the global +``db.session`` internally and have been converted to plain methods taking an explicit +``session: Session`` (per the pattern established in #40370/#40797/#41394/#41830, tracked in +#40372): + +- ``AppAnnotationHitHistory.account`` (joins through ``MessageAnnotation.account_id``) +- ``AppAnnotationHitHistory.annotation_create_account`` (looks up ``self.account_id`` directly) + +Each accessor is exercised against the real ``sqlite_session`` fixture (a genuine SQLAlchemy +``Session`` bound to a pristine full-schema SQLite database) so the assertions cover actual +query behaviour rather than a mock's recorded call. +""" -@pytest.fixture(autouse=True) -def patch_file_helpers(monkeypatch: pytest.MonkeyPatch) -> None: - """ - Patch file_helpers.get_signed_file_url to a deterministic stub. - """ - monkeypatch.setattr( - model_module.file_helpers, - "get_signed_file_url", - lambda file_id: f"https://signed.example/{file_id}", +def _annotation_hit_history_session_accessors_persist_account(session: Session) -> Account: + account = Account(name="Test Account", email="test@example.com") + session.add(account) + session.flush() + return account + + +def _annotation_hit_history_session_accessors_persist_annotation( + session: Session, *, app_id: str, account_id: str +) -> MessageAnnotation: + annotation = MessageAnnotation( + app_id=app_id, + conversation_id=None, + message_id=None, + question="What is AI?", + content="AI stands for Artificial Intelligence.", + account_id=account_id, ) + session.add(annotation) + session.flush() + return annotation + + +def _hit_history(*, app_id: str, annotation_id: str, account_id: str) -> AppAnnotationHitHistory: + return AppAnnotationHitHistory( + app_id=app_id, + annotation_id=annotation_id, + source="api", + question="What is AI?", + account_id=account_id, + score=0.95, + message_id=str(uuid4()), + annotation_question="What is AI?", + annotation_content="AI stands for Artificial Intelligence.", + ) + + +class TestAnnotationHitHistorySessionAccessorsCases: + class TestAppAnnotationHitHistoryAccount: + def test_returns_the_annotation_creator(self, sqlite_session: Session) -> None: + app_id = str(uuid4()) + creator = _annotation_hit_history_session_accessors_persist_account(sqlite_session) + annotation = _annotation_hit_history_session_accessors_persist_annotation( + sqlite_session, app_id=app_id, account_id=creator.id + ) + history = _hit_history(app_id=app_id, annotation_id=annotation.id, account_id=str(uuid4())) + + result = history.account(session=sqlite_session) + + assert result is not None + assert result.id == creator.id + + def test_returns_none_when_annotation_missing(self, sqlite_session: Session) -> None: + history = _hit_history(app_id=str(uuid4()), annotation_id=str(uuid4()), account_id=str(uuid4())) + + assert history.account(session=sqlite_session) is None + + class TestAppAnnotationHitHistoryAnnotationCreateAccount: + def test_returns_the_hit_requester(self, sqlite_session: Session) -> None: + app_id = str(uuid4()) + requester = _annotation_hit_history_session_accessors_persist_account(sqlite_session) + annotation = _annotation_hit_history_session_accessors_persist_annotation( + sqlite_session, app_id=app_id, account_id=str(uuid4()) + ) + history = _hit_history(app_id=app_id, annotation_id=annotation.id, account_id=requester.id) + + result = history.annotation_create_account(session=sqlite_session) + + assert result is not None + assert result.id == requester.id + + def test_returns_none_when_account_missing(self, sqlite_session: Session) -> None: + history = _hit_history(app_id=str(uuid4()), annotation_id=str(uuid4()), account_id=str(uuid4())) + + assert history.annotation_create_account(session=sqlite_session) is None + + +"""Regression coverage for the ``@property``→session-parameter refactor on +``AppAnnotationSetting.collection_binding_detail``. + +The legacy ``@property`` reached for the global ``db.session`` internally and has been converted +to a plain method taking an explicit ``session: Session`` (per the pattern established in +#40370/#40797/#41394/#41830, tracked in #40372). + +The accessor is exercised against the real ``sqlite_session`` fixture (a genuine SQLAlchemy +``Session`` bound to a pristine full-schema SQLite database) so the assertions cover actual query +behaviour rather than a mock's recorded call. +""" + + +def _persist_collection_binding(session: Session) -> DatasetCollectionBinding: + binding = DatasetCollectionBinding( + collection_name="test_collection", + provider_name="test_provider", + model_name="test_model", + type=CollectionBindingType.DATASET, + ) + session.add(binding) + session.flush() + return binding + + +def _persist_setting(session: Session, *, collection_binding_id: str) -> AppAnnotationSetting: + setting = AppAnnotationSetting( + app_id=str(uuid4()), + score_threshold=0.8, + collection_binding_id=collection_binding_id, + created_user_id=str(uuid4()), + updated_user_id=str(uuid4()), + ) + session.add(setting) + session.flush() + return setting + + +class TestAppAnnotationSettingSessionAccessorsCases: + class TestAppAnnotationSettingCollectionBindingDetail: + def test_returns_the_bound_collection_binding(self, sqlite_session: Session) -> None: + binding = _persist_collection_binding(sqlite_session) + setting = _persist_setting(sqlite_session, collection_binding_id=binding.id) + + result = setting.collection_binding_detail(session=sqlite_session) + + assert result is not None + assert result.id == binding.id + + def test_returns_none_when_binding_missing(self, sqlite_session: Session) -> None: + setting = _persist_setting(sqlite_session, collection_binding_id=str(uuid4())) + + assert setting.collection_binding_detail(session=sqlite_session) is None + + +"""Regression coverage for the ``db.session``→session-parameter refactor on +``AppMCPServer.generate_server_code``. + +The static method reached for the global ``db.session`` internally and has been converted to +take an explicit ``session: Session`` (per the pattern established in #40370/#40797/#41394, +tracked in #40372). +""" + + +def _persist_server(session: Session, *, server_code: str) -> AppMCPServer: + server = AppMCPServer( + tenant_id=str(uuid4()), + app_id=str(uuid4()), + name="Test MCP Server", + description="Test server", + server_code=server_code, + status="active", + parameters="{}", + ) + session.add(server) + session.flush() + return server + + +class TestAppMcpServerSessionAccessorsCases: + @staticmethod + def test_generate_server_code_retries_a_taken_code_through_the_given_session(sqlite_session: Session) -> None: + _persist_server(sqlite_session, server_code="taken-code") + + with patch("models.model.generate_string", side_effect=["taken-code", "fresh-code"]): + result = AppMCPServer.generate_server_code(16, session=sqlite_session) + + assert result == "fresh-code" + + +""" +Comprehensive unit tests for App models. + +This test suite covers: +- App configuration validation +- App-Message relationships +- Conversation model integrity +- Annotation model relationships +""" + + +class TestAppModelsCases: + class TestAppModelValidation: + """Test suite for App model validation and basic operations.""" + + def test_app_creation_with_required_fields(self): + """Test creating an app with all required fields.""" + # Arrange + tenant_id = str(uuid4()) + created_by = str(uuid4()) + + # Act + app = App( + tenant_id=tenant_id, + name="Test App", + mode=AppMode.CHAT, + enable_site=True, + enable_api=False, + created_by=created_by, + ) + + # Assert + assert app.name == "Test App" + assert app.tenant_id == tenant_id + assert app.mode == AppMode.CHAT + assert app.enable_site is True + assert app.enable_api is False + assert app.created_by == created_by + + def test_app_creation_with_optional_fields(self): + """Test creating an app with optional fields.""" + # Arrange & Act + app = App( + tenant_id=str(uuid4()), + name="Test App", + mode=AppMode.COMPLETION, + enable_site=True, + enable_api=True, + created_by=str(uuid4()), + description="Test description", + icon_type=IconType.EMOJI, + icon="🤖", + icon_background="#FF5733", + is_demo=True, + is_public=False, + api_rpm=100, + api_rph=1000, + ) + + # Assert + assert app.description == "Test description" + assert app.icon_type == IconType.EMOJI + assert app.icon == "🤖" + assert app.icon_background == "#FF5733" + assert app.is_demo is True + assert app.is_public is False + assert app.api_rpm == 100 + assert app.api_rph == 1000 + + def test_app_mode_validation(self): + """Test app mode enum values.""" + # Assert + expected_modes = { + "chat", + "completion", + "workflow", + "advanced-chat", + "agent-chat", + "agent", + "channel", + "rag-pipeline", + } + assert {mode.value for mode in AppMode} == expected_modes + + def test_app_mode_value_of(self): + """Test AppMode.value_of method.""" + # Act & Assert + assert AppMode.value_of("chat") == AppMode.CHAT + assert AppMode.value_of("completion") == AppMode.COMPLETION + assert AppMode.value_of("workflow") == AppMode.WORKFLOW + + with pytest.raises(ValueError, match="invalid mode value"): + AppMode.value_of("invalid_mode") + + def test_icon_type_validation(self): + """Test icon type enum values.""" + # Assert + assert {t.value for t in IconType} == {"image", "emoji", "link"} + + @pytest.mark.parametrize("sqlite_session", [(App, AppModelConfig)], indirect=True) + def test_app_desc_or_prompt_with_description(self, sqlite_session: Session): + """Test desc_or_prompt_with_session when description exists.""" + # Arrange + app = App( + tenant_id=str(uuid4()), + name="Test App", + mode=AppMode.CHAT, + enable_site=True, + enable_api=False, + created_by=str(uuid4()), + description="App description", + ) + + # Act + result = app.desc_or_prompt_with_session(session=sqlite_session) + + # Assert + assert result == "App description" + + @pytest.mark.parametrize("sqlite_session", [(App, AppModelConfig)], indirect=True) + def test_app_desc_or_prompt_without_description(self, sqlite_session: Session): + """Test desc_or_prompt property when description is empty.""" + # Arrange + app = App( + tenant_id=str(uuid4()), + name="Test App", + mode=AppMode.CHAT, + enable_site=True, + enable_api=False, + created_by=str(uuid4()), + description="", + ) + decoy_config = AppModelConfig(app_id=str(uuid4()), pre_prompt="A decoy prompt") + sqlite_session.add_all([app, decoy_config]) + sqlite_session.flush() + + # Act + result = app.desc_or_prompt_with_session(session=sqlite_session) + + # Assert + assert result == "" + + @pytest.mark.parametrize("sqlite_session", [(App, AppModelConfig)], indirect=True) + def test_app_is_agent_false_when_not_configured_as_agent(self, sqlite_session: Session): + """`is_agent_with_session` returns False when the config has no agent mode.""" + # Arrange + app = App( + tenant_id=str(uuid4()), + name="Test App", + mode=AppMode.CHAT, + enable_site=True, + enable_api=False, + created_by=str(uuid4()), + app_model_config_id=str(uuid4()), + ) + sqlite_session.add(app) + sqlite_session.flush() + + # Act + result = app.is_agent_with_session(session=sqlite_session) + + # Assert + assert result is False + + @pytest.mark.parametrize("sqlite_session", [(App, AppModelConfig)], indirect=True) + def test_app_is_agent_with_session_persists_mode_across_sessions(self, sqlite_session: Session): + app = App( + tenant_id=str(uuid4()), + name="Test App", + mode=AppMode.CHAT, + enable_site=True, + enable_api=False, + created_by=str(uuid4()), + ) + sqlite_session.add(app) + sqlite_session.flush() + model_config = AppModelConfig( + app_id=app.id, + agent_mode=json.dumps({"enabled": True, "strategy": "react"}), + ) + sqlite_session.add(model_config) + sqlite_session.flush() + app.app_model_config_id = model_config.id + sqlite_session.commit() + + with Session(sqlite_session.get_bind(), expire_on_commit=False) as migration_session: + assert app.is_agent_with_session(session=migration_session) is True + + sqlite_session.expire_all() + assert sqlite_session.get(App, app.id).mode == AppMode.AGENT_CHAT + + @pytest.mark.parametrize("sqlite_session", [(App, AppModelConfig)], indirect=True) + def test_app_mode_compatible_with_agent(self, sqlite_session: Session): + """Test mode_compatible_with_agent property.""" + # Arrange + app = App( + tenant_id=str(uuid4()), + name="Test App", + mode=AppMode.CHAT, + enable_site=True, + enable_api=False, + created_by=str(uuid4()), + ) + decoy_config = AppModelConfig( + app_id=str(uuid4()), + agent_mode=json.dumps({"enabled": True, "strategy": "react"}), + ) + sqlite_session.add_all([app, decoy_config]) + sqlite_session.flush() + + # Act + result = app.mode_compatible_with_agent_with_session(session=sqlite_session) + + # Assert + assert result == AppMode.CHAT + + @pytest.mark.parametrize("sqlite_session", [(App, AppModelConfig)], indirect=True) + def test_app_mode_compatible_with_agent_reports_agent_chat(self, sqlite_session: Session): + """A CHAT app whose own config enables agent mode reports AGENT_CHAT.""" + # Arrange + app = App( + tenant_id=str(uuid4()), + name="Test App", + mode=AppMode.CHAT, + enable_site=True, + enable_api=False, + created_by=str(uuid4()), + ) + sqlite_session.add(app) + sqlite_session.flush() + app.app_model_config_id = str(uuid4()) + config = AppModelConfig( + app_id=app.id, + agent_mode=json.dumps({"enabled": True, "strategy": "react"}), + ) + config.id = app.app_model_config_id + sqlite_session.add(config) + sqlite_session.flush() + + # Act + result = app.mode_compatible_with_agent_with_session(session=sqlite_session) + + # Assert + assert result == AppMode.AGENT_CHAT + + @pytest.mark.parametrize("sqlite_session", [(App, AppModelConfig)], indirect=True) + def test_app_model_config_with_session_reads_through_the_caller_session(self, sqlite_session: Session): + """`app_model_config_with_session` resolves the linked config through the caller's session.""" + # Arrange + app = App( + tenant_id=str(uuid4()), + name="Test App", + mode=AppMode.CHAT, + enable_site=True, + enable_api=False, + created_by=str(uuid4()), + ) + sqlite_session.add(app) + sqlite_session.flush() + model_config = AppModelConfig(app_id=app.id, agent_mode=json.dumps({"enabled": False})) + sqlite_session.add(model_config) + sqlite_session.flush() + app.app_model_config_id = model_config.id + sqlite_session.commit() + + # Act: a session the model never owns, proving the lookup does not reach for a global one + with Session(sqlite_session.get_bind(), expire_on_commit=False) as caller_session: + result = app.app_model_config_with_session(session=caller_session) + + # Assert + assert result is not None + assert result.id == model_config.id + assert result in caller_session + + @pytest.mark.parametrize("sqlite_session", [(App, AppModelConfig)], indirect=True) + def test_app_model_config_with_session_returns_none_when_unlinked(self, sqlite_session: Session): + """An app without `app_model_config_id` resolves to None without querying.""" + # Arrange + app = App( + tenant_id=str(uuid4()), + name="Test App", + mode=AppMode.CHAT, + enable_site=True, + enable_api=False, + created_by=str(uuid4()), + ) + sqlite_session.add(app) + sqlite_session.flush() + + # Act / Assert: an unbound session would raise if the guard were dropped + with Session(expire_on_commit=False) as unbound_session: + assert app.app_model_config_with_session(session=unbound_session) is None + + @pytest.mark.parametrize("sqlite_session", [(App, AppModelConfig)], indirect=True) + def test_workflow_with_session_reads_through_the_caller_session(self, sqlite_session: Session): + """`workflow_with_session` resolves the published workflow through the caller's session.""" + # Arrange + app = App( + tenant_id=str(uuid4()), + name="Test App", + mode=AppMode.WORKFLOW, + enable_site=True, + enable_api=False, + created_by=str(uuid4()), + ) + sqlite_session.add(app) + sqlite_session.flush() + workflow = Workflow( + tenant_id=app.tenant_id, + app_id=app.id, + type=WorkflowType.WORKFLOW, + version="1", + graph=json.dumps({"nodes": [], "edges": []}), + created_by=app.created_by, + ) + workflow._features = "{}" + sqlite_session.add(workflow) + sqlite_session.flush() + app.workflow_id = workflow.id + sqlite_session.commit() + + # Act: a session the model never owns, proving the lookup does not reach for a global one + with Session(sqlite_session.get_bind(), expire_on_commit=False) as caller_session: + result = app.workflow_with_session(session=caller_session) + + # Assert + assert result is not None + assert result.id == workflow.id + assert result in caller_session + + @pytest.mark.parametrize("sqlite_session", [(App, AppModelConfig)], indirect=True) + def test_workflow_with_session_returns_none_when_unpublished(self, sqlite_session: Session): + """An app without `workflow_id` resolves to None without querying.""" + # Arrange + app = App( + tenant_id=str(uuid4()), + name="Test App", + mode=AppMode.WORKFLOW, + enable_site=True, + enable_api=False, + created_by=str(uuid4()), + ) + sqlite_session.add(app) + sqlite_session.flush() + + # Act / Assert: an unbound session would raise if the guard were dropped + with Session(expire_on_commit=False) as unbound_session: + assert app.workflow_with_session(session=unbound_session) is None + + @pytest.mark.parametrize("sqlite_session", [(App, AppModelConfig)], indirect=True) + def test_deleted_tools_checks_plugin_builtin_providers_through_core_plugin_service( + self, sqlite_session: Session + ): + """Plugin-backed built-in tools are checked through core PluginService.""" + # Arrange + app = App( + tenant_id="tenant-1", + name="Test App", + mode=AppMode.CHAT, + enable_site=True, + enable_api=False, + created_by=str(uuid4()), + ) + sqlite_session.add(app) + sqlite_session.flush() + app_model_config = AppModelConfig( + app_id=app.id, + agent_mode=json.dumps( + { + "enabled": True, + "strategy": "function_call", + "tools": [ + { + "provider_type": "builtin", + "provider_id": "langgenius/openai/openai", + "tool_name": "chat", + "tool_parameters": {}, + } + ], + "prompt": None, + } + ), + ) + sqlite_session.add(app_model_config) + sqlite_session.flush() + app.app_model_config_id = app_model_config.id + sqlite_session.flush() + + # Act + with ( + patch("core.tools.tool_manager.ToolManager.get_hardcoded_provider", side_effect=Exception), + patch("core.plugin.plugin_service.PluginService.check_tools_existence", return_value=[False]) as exists, + ): + result = app.deleted_tools_with_session(session=sqlite_session) + + # Assert + assert result == [{"type": "builtin", "tool_name": "chat", "provider_id": "langgenius/openai/openai"}] + exists.assert_called_once() + assert exists.call_args.args[0] == "tenant-1" + assert [str(provider_id) for provider_id in exists.call_args.args[1]] == ["langgenius/openai/openai"] + + class TestAppModelConfig: + """Test suite for AppModelConfig model.""" + + def test_legacy_dataset_configs_default_retrieval_model_without_dropping_settings(self): + config = AppModelConfig( + app_id=str(uuid4()), + dataset_configs=json.dumps({"reranking_enable": True, "top_k": 8}), + ) + + assert config.dataset_configs_dict == { + "retrieval_model": "single", + "reranking_enable": True, + "top_k": 8, + } + + def test_app_model_config_creation(self): + """Test creating an AppModelConfig.""" + # Arrange + app_id = str(uuid4()) + created_by = str(uuid4()) + + # Act + config = AppModelConfig( + app_id=app_id, + provider="openai", + model_id="gpt-4", + created_by=created_by, + ) + + # Assert + assert config.app_id == app_id + assert config.provider == "openai" + assert config.model_id == "gpt-4" + assert config.created_by == created_by + + def test_app_model_config_with_configs_json(self): + """Test AppModelConfig with JSON configs.""" + # Arrange + configs = {"temperature": 0.7, "max_tokens": 1000} + + # Act + config = AppModelConfig( + app_id=str(uuid4()), + provider="openai", + model_id="gpt-4", + created_by=str(uuid4()), + configs=configs, + ) + + # Assert + assert config.configs == configs + + def test_app_model_config_model_dict_property(self): + """Test model_dict property.""" + # Arrange + model_data = {"provider": "openai", "name": "gpt-4"} + config = AppModelConfig( + app_id=str(uuid4()), + provider="openai", + model_id="gpt-4", + created_by=str(uuid4()), + model=json.dumps(model_data), + ) + + # Act + result = config.model_dict + + # Assert + assert result == model_data + + def test_app_model_config_model_dict_empty(self): + """Test model_dict property when model is None.""" + # Arrange + config = AppModelConfig( + app_id=str(uuid4()), + provider="openai", + model_id="gpt-4", + created_by=str(uuid4()), + model=None, + ) + + # Act + result = config.model_dict + + # Assert + assert result == {} + + def test_app_model_config_suggested_questions_list(self): + """Test suggested_questions_list property.""" + # Arrange + questions = ["What can you do?", "How does this work?"] + config = AppModelConfig( + app_id=str(uuid4()), + provider="openai", + model_id="gpt-4", + created_by=str(uuid4()), + suggested_questions=json.dumps(questions), + ) + + # Act + result = config.suggested_questions_list + + # Assert + assert result == questions + + def test_to_dict_uses_injected_annotation_reply(self): + config = AppModelConfig(app_id=str(uuid4())) + annotation_reply: AnnotationReplyConfig = {"enabled": False} + + result = config.to_dict(annotation_reply=annotation_reply) + + assert result["annotation_reply"] == annotation_reply + + class TestAnnotationReplyConfigLoader: + @pytest.mark.parametrize("sqlite_session", [(AppAnnotationSetting, DatasetCollectionBinding)], indirect=True) + def test_load_annotation_reply_config_returns_disabled_when_setting_missing(self, sqlite_session: Session): + binding = DatasetCollectionBinding( + provider_name="decoy-provider", + model_name="decoy-model", + type=CollectionBindingType.ANNOTATION, + collection_name="decoy-collection", + ) + sqlite_session.add(binding) + sqlite_session.flush() + sqlite_session.add( + AppAnnotationSetting( + app_id="other-app", + score_threshold=0.9, + collection_binding_id=binding.id, + created_user_id="user-1", + updated_user_id="user-1", + ) + ) + sqlite_session.flush() + + result = load_annotation_reply_config(sqlite_session, "app-1") + + assert result == {"enabled": False} + + @pytest.mark.parametrize("sqlite_session", [(AppAnnotationSetting, DatasetCollectionBinding)], indirect=True) + def test_load_annotation_reply_config_returns_embedding_model(self, sqlite_session: Session): + collection_binding = DatasetCollectionBinding( + provider_name="provider", + model_name="embedding", + type=CollectionBindingType.ANNOTATION, + collection_name="annotation-collection", + ) + sqlite_session.add(collection_binding) + sqlite_session.flush() + annotation_setting = AppAnnotationSetting( + app_id="app-1", + score_threshold=0.7, + collection_binding_id=collection_binding.id, + created_user_id="user-1", + updated_user_id="user-1", + ) + sqlite_session.add(annotation_setting) + sqlite_session.flush() + + result = load_annotation_reply_config(sqlite_session, "app-1") + + assert result == { + "id": annotation_setting.id, + "enabled": True, + "score_threshold": 0.7, + "embedding_model": { + "embedding_provider_name": "provider", + "embedding_model_name": "embedding", + }, + } + + @pytest.mark.parametrize("sqlite_session", [(AppAnnotationSetting, DatasetCollectionBinding)], indirect=True) + def test_load_annotation_reply_config_raises_when_binding_missing(self, sqlite_session: Session): + sqlite_session.add( + AppAnnotationSetting( + app_id="app-1", + score_threshold=0.7, + collection_binding_id=str(uuid4()), + created_user_id="user-1", + updated_user_id="user-1", + ) + ) + sqlite_session.flush() + + with pytest.raises(ValueError, match="Collection binding detail not found"): + load_annotation_reply_config(sqlite_session, "app-1") + + class TestConversationModel: + """Test suite for Conversation model integrity.""" + + def test_app_returns_matching_app_from_caller_session(self, sqlite_session: Session) -> None: + other_app = App( + tenant_id=str(uuid4()), name="Other app", mode=AppMode.CHAT, enable_site=False, enable_api=False + ) + app = App( + tenant_id=str(uuid4()), name="Conversation app", mode=AppMode.CHAT, enable_site=False, enable_api=False + ) + sqlite_session.add_all([other_app, app]) + sqlite_session.flush() + conversation = Conversation(app_id=app.id) + + assert conversation.app(session=sqlite_session) is app + + def test_app_returns_none_when_app_is_missing(self, sqlite_session: Session) -> None: + other_app = App( + tenant_id=str(uuid4()), name="Other app", mode=AppMode.CHAT, enable_site=False, enable_api=False + ) + sqlite_session.add(other_app) + sqlite_session.flush() + conversation = Conversation(app_id=str(uuid4())) + + assert conversation.app(session=sqlite_session) is None + + def test_conversation_creation_with_required_fields(self): + """Test creating a conversation with required fields.""" + # Arrange + app_id = str(uuid4()) + from_end_user_id = str(uuid4()) + + # Act + conversation = Conversation( + app_id=app_id, + mode=AppMode.CHAT, + name="Test Conversation", + status="normal", + from_source=ConversationFromSource.API, + from_end_user_id=from_end_user_id, + ) + + # Assert + assert conversation.app_id == app_id + assert conversation.mode == AppMode.CHAT + assert conversation.name == "Test Conversation" + assert conversation.status == "normal" + assert conversation.from_source == "api" + assert conversation.from_end_user_id == from_end_user_id + + def test_conversation_with_inputs(self, sqlite_session: Session): + """Test conversation inputs round-trip through the session-aware accessor.""" + # Arrange + inputs = {"query": "Hello", "context": "test"} + conversation = Conversation( + app_id=str(uuid4()), + mode=AppMode.CHAT, + name="Test Conversation", + status="normal", + from_source=ConversationFromSource.API, + from_end_user_id=str(uuid4()), + ) + conversation._inputs = inputs + + # Act + result = conversation.inputs_with_session(session=sqlite_session) + + # Assert + assert result == inputs + + def test_conversation_inputs_is_write_only(self): + """Reading conversation.inputs must fail loudly; reads go through inputs_with_session.""" + conversation = Conversation(app_id=str(uuid4()), _inputs={}) + + with pytest.raises(AttributeError, match="no getter"): + _ = conversation.inputs + + def test_conversation_inputs_setter(self): + """Test conversation inputs setter.""" + # Arrange + conversation = Conversation( + app_id=str(uuid4()), + mode=AppMode.CHAT, + name="Test Conversation", + status="normal", + from_source=ConversationFromSource.API, + from_end_user_id=str(uuid4()), + ) + inputs = {"query": "Hello", "context": "test"} + + # Act + conversation.inputs = inputs + + # Assert + assert conversation._inputs == inputs + + @pytest.mark.parametrize("sqlite_session", [(Conversation, Message)], indirect=True) + def test_conversation_summary_or_query_with_summary(self, sqlite_session: Session): + """Test summary_or_query_with_session when summary exists.""" + # Arrange + conversation = Conversation( + app_id=str(uuid4()), + mode=AppMode.CHAT, + name="Test Conversation", + status="normal", + from_source=ConversationFromSource.API, + from_end_user_id=str(uuid4()), + summary="Test summary", + ) + + # Act + result = conversation.summary_or_query_with_session(session=sqlite_session) + + # Assert + assert result == "Test summary" + + @pytest.mark.parametrize("sqlite_session", [(Conversation, Message)], indirect=True) + def test_conversation_summary_or_query_without_summary(self, sqlite_session: Session): + """Test summary_or_query_with_session when summary is empty.""" + # Arrange + conversation = Conversation( + app_id=str(uuid4()), + mode=AppMode.CHAT, + name="Test Conversation", + status="normal", + from_source=ConversationFromSource.API, + from_end_user_id=str(uuid4()), + summary=None, + ) + conversation._inputs = {} + sqlite_session.add(conversation) + sqlite_session.flush() + first_message = Message( + app_id=conversation.app_id, + conversation_id=conversation.id, + query="First message query", + message={"role": "user", "content": "First message query"}, + answer="First answer", + message_unit_price=Decimal(0), + answer_unit_price=Decimal(0), + currency="USD", + from_source=ConversationFromSource.API, + ) + first_message.created_at = datetime(2024, 1, 1) + first_message._inputs = {} + later_message = Message( + app_id=conversation.app_id, + conversation_id=conversation.id, + query="Later message query", + message={"role": "user", "content": "Later message query"}, + answer="Later answer", + message_unit_price=Decimal(0), + answer_unit_price=Decimal(0), + currency="USD", + from_source=ConversationFromSource.API, + ) + later_message.created_at = datetime(2024, 1, 2) + later_message._inputs = {} + sqlite_session.add_all([later_message, first_message]) + sqlite_session.flush() + + # Act + result = conversation.summary_or_query_with_session(session=sqlite_session) + + # Assert + assert result == "First message query" + + @pytest.mark.parametrize( + "sqlite_session", [(Conversation, AppModelConfig, AppAnnotationSetting)], indirect=True + ) + def test_model_config_uses_caller_session_for_annotation_reply(self, sqlite_session: Session): + app_id = str(uuid4()) + app_model_config = AppModelConfig(app_id=app_id, pre_prompt="Persisted prompt") + sqlite_session.add(app_model_config) + sqlite_session.flush() + conversation = Conversation( + app_id=app_id, + app_model_config_id=app_model_config.id, + mode=AppMode.CHAT, + name="Test Conversation", + status="normal", + from_source=ConversationFromSource.API, + from_end_user_id=str(uuid4()), + model_id="model-1", + model_provider="provider-1", + ) + conversation._inputs = {} + sqlite_session.add(conversation) + sqlite_session.flush() + + result = conversation.model_config_with_session(session=sqlite_session) + + assert result["annotation_reply"] == {"enabled": False} + assert result["pre_prompt"] == "Persisted prompt" + assert result["model_id"] == "model-1" + assert result["provider"] == "provider-1" + + @pytest.mark.parametrize("sqlite_session", [(Conversation, AppAnnotationSetting)], indirect=True) + def test_override_model_config_uses_caller_session_for_annotation_reply(self, sqlite_session: Session): + app_id = str(uuid4()) + conversation = Conversation( + app_id=app_id, + mode=AppMode.CHAT, + name="Test Conversation", + status="normal", + from_source=ConversationFromSource.API, + from_end_user_id=str(uuid4()), + override_model_configs=json.dumps({"model": {"provider": "openai", "name": "gpt-4"}}), + model_id="model-1", + model_provider="provider-1", + ) + conversation._inputs = {} + sqlite_session.add(conversation) + sqlite_session.flush() + + result = conversation.model_config_with_session(session=sqlite_session) + + assert result["annotation_reply"] == {"enabled": False} + assert result["model"] == {"provider": "openai", "name": "gpt-4"} + assert result["model_id"] == "model-1" + assert result["provider"] == "provider-1" + + def test_conversation_in_debug_mode(self): + """Test in_debug_mode property.""" + # Arrange + conversation = Conversation( + app_id=str(uuid4()), + mode=AppMode.CHAT, + name="Test Conversation", + status="normal", + from_source=ConversationFromSource.API, + from_end_user_id=str(uuid4()), + override_model_configs='{"model": "gpt-4"}', + ) + + # Act + result = conversation.in_debug_mode + + # Assert + assert result is True + + @pytest.mark.parametrize("sqlite_session", [(Conversation,)], indirect=True) + def test_conversation_to_dict_serialization(self, sqlite_session: Session): + """Test conversation to_dict method.""" + # Arrange + app_id = str(uuid4()) + from_end_user_id = str(uuid4()) + conversation = Conversation( + app_id=app_id, + mode=AppMode.CHAT, + name="Test Conversation", + status="normal", + from_source=ConversationFromSource.API, + from_end_user_id=from_end_user_id, + dialogue_count=5, + ) + conversation.id = str(uuid4()) + conversation._inputs = {"query": "test"} + + # Act + result = conversation.to_dict(session=sqlite_session) + + # Assert + assert result["id"] == conversation.id + assert result["app_id"] == app_id + assert result["mode"] == AppMode.CHAT + assert result["name"] == "Test Conversation" + assert result["status"] == "normal" + assert result["from_source"] == "api" + assert result["from_end_user_id"] == from_end_user_id + assert result["dialogue_count"] == 5 + assert result["inputs"] == {"query": "test"} + + class TestMessageModel: + """Test suite for Message model and App-Message relationships.""" + + def test_message_creation_with_required_fields(self): + """Test creating a message with required fields.""" + # Arrange + app_id = str(uuid4()) + conversation_id = str(uuid4()) + + # Act + message = Message( + app_id=app_id, + conversation_id=conversation_id, + query="What is AI?", + message={"role": "user", "content": "What is AI?"}, + answer="AI stands for Artificial Intelligence.", + message_unit_price=Decimal("0.0001"), + answer_unit_price=Decimal("0.0002"), + currency="USD", + from_source=ConversationFromSource.API, + ) + + # Assert + assert message.app_id == app_id + assert message.conversation_id == conversation_id + assert message.query == "What is AI?" + assert message.answer == "AI stands for Artificial Intelligence." + assert message.currency == "USD" + assert message.from_source == "api" + + def test_message_with_inputs(self, sqlite_session: Session): + """Test message inputs round-trip through the session-aware accessor.""" + # Arrange + inputs = {"query": "Hello", "context": "test"} + message = Message( + app_id=str(uuid4()), + conversation_id=str(uuid4()), + query="Test query", + message={"role": "user", "content": "Test"}, + answer="Test answer", + message_unit_price=Decimal("0.0001"), + answer_unit_price=Decimal("0.0002"), + currency="USD", + from_source=ConversationFromSource.API, + _inputs=inputs, + ) + + # Act + result = message.inputs_with_session(session=sqlite_session) + + # Assert + assert result == inputs + + def test_message_inputs_is_write_only(self): + """Reading message.inputs must fail loudly; reads go through inputs_with_session.""" + message = Message(app_id=str(uuid4()), _inputs={}) + + with pytest.raises(AttributeError, match="no getter"): + _ = message.inputs + + def test_message_inputs_setter(self): + """Test message inputs setter.""" + # Arrange + message = Message( + app_id=str(uuid4()), + conversation_id=str(uuid4()), + query="Test query", + message={"role": "user", "content": "Test"}, + answer="Test answer", + message_unit_price=Decimal("0.0001"), + answer_unit_price=Decimal("0.0002"), + currency="USD", + from_source=ConversationFromSource.API, + ) + inputs = {"query": "Hello", "context": "test"} + + # Act + message.inputs = inputs + + # Assert + assert message._inputs == inputs + + def test_message_in_debug_mode(self): + """Test message in_debug_mode property.""" + # Arrange + message = Message( + app_id=str(uuid4()), + conversation_id=str(uuid4()), + query="Test query", + message={"role": "user", "content": "Test"}, + answer="Test answer", + message_unit_price=Decimal("0.0001"), + answer_unit_price=Decimal("0.0002"), + currency="USD", + from_source=ConversationFromSource.API, + override_model_configs='{"model": "gpt-4"}', + ) + + # Act + result = message.in_debug_mode + + # Assert + assert result is True + + def test_message_metadata_dict_property(self): + """Test message_metadata_dict property.""" + # Arrange + metadata = {"retriever_resources": ["doc1", "doc2"], "usage": {"tokens": 100}} + message = Message( + app_id=str(uuid4()), + conversation_id=str(uuid4()), + query="Test query", + message={"role": "user", "content": "Test"}, + answer="Test answer", + message_unit_price=Decimal("0.0001"), + answer_unit_price=Decimal("0.0002"), + currency="USD", + from_source=ConversationFromSource.API, + message_metadata=json.dumps(metadata), + ) + + # Act + result = message.message_metadata_dict + + # Assert + assert result == metadata + + def test_message_metadata_dict_empty(self): + """Test message_metadata_dict when metadata is None.""" + # Arrange + message = Message( + app_id=str(uuid4()), + conversation_id=str(uuid4()), + query="Test query", + message={"role": "user", "content": "Test"}, + answer="Test answer", + message_unit_price=Decimal("0.0001"), + answer_unit_price=Decimal("0.0002"), + currency="USD", + from_source=ConversationFromSource.API, + message_metadata=None, + ) + + # Act + result = message.message_metadata_dict + + # Assert + assert result == {} + + @pytest.mark.parametrize("sqlite_session", [(Message,)], indirect=True) + def test_message_to_dict_serialization(self, sqlite_session: Session): + """Test message to_dict method.""" + # Arrange + app_id = str(uuid4()) + conversation_id = str(uuid4()) + now = datetime.now(UTC) + + message = Message( + app_id=app_id, + conversation_id=conversation_id, + query="Test query", + message={"role": "user", "content": "Test"}, + answer="Test answer", + message_unit_price=Decimal("0.0001"), + answer_unit_price=Decimal("0.0002"), + total_price=Decimal("0.0003"), + currency="USD", + from_source=ConversationFromSource.API, + status="normal", + id=str(uuid4()), + _inputs={"query": "test"}, + ) + message.created_at = now + message.updated_at = now + + # Act + result = message.to_dict(session=sqlite_session) + + # Assert + assert result["id"] == message.id + assert result["app_id"] == app_id + assert result["conversation_id"] == conversation_id + assert result["query"] == "Test query" + assert result["answer"] == "Test answer" + assert result["status"] == "normal" + assert result["from_source"] == "api" + assert result["inputs"] == {"query": "test"} + assert "created_at" in result + assert "updated_at" in result + + def test_message_from_dict_deserialization(self): + """Test message from_dict method.""" + # Arrange + message_id = str(uuid4()) + app_id = str(uuid4()) + conversation_id = str(uuid4()) + data = { + "id": message_id, + "app_id": app_id, + "conversation_id": conversation_id, + "model_id": "gpt-4", + "inputs": {"query": "test"}, + "query": "Test query", + "message": {"role": "user", "content": "Test"}, + "answer": "Test answer", + "total_price": Decimal("0.0003"), + "status": "normal", + "error": None, + "message_metadata": {"usage": {"tokens": 100}}, + "from_source": "api", + "from_end_user_id": None, + "from_account_id": None, + "created_at": "2024-01-01T00:00:00", + "updated_at": "2024-01-01T00:00:00", + "agent_based": False, + "workflow_run_id": None, + } + + # Act + message = Message.from_dict(data) + + # Assert + assert message.id == message_id + assert message.app_id == app_id + assert message.conversation_id == conversation_id + assert message.query == "Test query" + assert message.answer == "Test answer" + + class TestMessageAnnotation: + """Test suite for MessageAnnotation and annotation relationships.""" + + def test_message_annotation_creation(self): + """Test creating a message annotation.""" + # Arrange + app_id = str(uuid4()) + conversation_id = str(uuid4()) + message_id = str(uuid4()) + account_id = str(uuid4()) + + # Act + annotation = MessageAnnotation( + app_id=app_id, + conversation_id=conversation_id, + message_id=message_id, + question="What is AI?", + content="AI stands for Artificial Intelligence.", + account_id=account_id, + ) + + # Assert + assert annotation.app_id == app_id + assert annotation.conversation_id == conversation_id + assert annotation.message_id == message_id + assert annotation.question == "What is AI?" + assert annotation.content == "AI stands for Artificial Intelligence." + assert annotation.account_id == account_id + + def test_message_annotation_without_message_id(self): + """Test creating annotation without message_id.""" + # Arrange + app_id = str(uuid4()) + account_id = str(uuid4()) + + # Act + annotation = MessageAnnotation( + app_id=app_id, + question="What is AI?", + conversation_id=None, + message_id=None, + content="AI stands for Artificial Intelligence.", + account_id=account_id, + ) + + # Assert + assert annotation.app_id == app_id + assert annotation.message_id is None + assert annotation.conversation_id is None + assert annotation.question == "What is AI?" + assert annotation.content == "AI stands for Artificial Intelligence." + + def test_message_annotation_hit_count_default(self): + """Test annotation hit_count default value.""" + # Arrange + annotation = MessageAnnotation( + app_id=str(uuid4()), + question="Test question", + conversation_id=None, + message_id=None, + content="Test content", + account_id=str(uuid4()), + ) + + # Act & Assert - default value is set by database + # Model instantiation doesn't set server defaults + assert hasattr(annotation, "hit_count") + + class TestAppAnnotationSetting: + """Test suite for AppAnnotationSetting model.""" + + def test_app_annotation_setting_creation(self): + """Test creating an app annotation setting.""" + # Arrange + app_id = str(uuid4()) + collection_binding_id = str(uuid4()) + created_user_id = str(uuid4()) + updated_user_id = str(uuid4()) + + # Act + setting = AppAnnotationSetting( + app_id=app_id, + score_threshold=0.8, + collection_binding_id=collection_binding_id, + created_user_id=created_user_id, + updated_user_id=updated_user_id, + ) + + # Assert + assert setting.app_id == app_id + assert setting.score_threshold == 0.8 + assert setting.collection_binding_id == collection_binding_id + assert setting.created_user_id == created_user_id + assert setting.updated_user_id == updated_user_id + + def test_app_annotation_setting_score_threshold_validation(self): + """Test score threshold values.""" + # Arrange & Act + setting_high = AppAnnotationSetting( + app_id=str(uuid4()), + score_threshold=0.95, + collection_binding_id=str(uuid4()), + created_user_id=str(uuid4()), + updated_user_id=str(uuid4()), + ) + setting_low = AppAnnotationSetting( + app_id=str(uuid4()), + score_threshold=0.5, + collection_binding_id=str(uuid4()), + created_user_id=str(uuid4()), + updated_user_id=str(uuid4()), + ) + + # Assert + assert setting_high.score_threshold == 0.95 + assert setting_low.score_threshold == 0.5 + + class TestAppAnnotationHitHistory: + """Test suite for AppAnnotationHitHistory model.""" + + def test_app_annotation_hit_history_creation(self): + """Test creating an annotation hit history.""" + # Arrange + app_id = str(uuid4()) + annotation_id = str(uuid4()) + message_id = str(uuid4()) + account_id = str(uuid4()) + + # Act + history = AppAnnotationHitHistory( + app_id=app_id, + annotation_id=annotation_id, + source="api", + question="What is AI?", + account_id=account_id, + score=0.95, + message_id=message_id, + annotation_question="What is AI?", + annotation_content="AI stands for Artificial Intelligence.", + ) + + # Assert + assert history.app_id == app_id + assert history.annotation_id == annotation_id + assert history.source == "api" + assert history.question == "What is AI?" + assert history.account_id == account_id + assert history.score == 0.95 + assert history.message_id == message_id + assert history.annotation_question == "What is AI?" + assert history.annotation_content == "AI stands for Artificial Intelligence." + + def test_app_annotation_hit_history_score_values(self): + """Test annotation hit history with different score values.""" + # Arrange & Act + history_high = AppAnnotationHitHistory( + app_id=str(uuid4()), + annotation_id=str(uuid4()), + source="api", + question="Test", + account_id=str(uuid4()), + score=0.99, + message_id=str(uuid4()), + annotation_question="Test", + annotation_content="Content", + ) + history_low = AppAnnotationHitHistory( + app_id=str(uuid4()), + annotation_id=str(uuid4()), + source="api", + question="Test", + account_id=str(uuid4()), + score=0.6, + message_id=str(uuid4()), + annotation_question="Test", + annotation_content="Content", + ) + + # Assert + assert history_high.score == 0.99 + assert history_low.score == 0.6 + + class TestSiteModel: + """Test suite for Site model.""" + + def test_site_core_bulk_insert_generates_ids(self, sqlite_session: Session): + """Core bulk inserts generate an ID for every site row.""" + app_ids = [str(uuid4()), str(uuid4())] + + sqlite_session.execute( + Site.__table__.insert(), + [ + { + "app_id": app_id, + "title": f"Site {index}", + "default_language": "en-US", + "customize_token_strategy": CustomizeTokenStrategy.UUID, + } + for index, app_id in enumerate(app_ids) + ], + ) + + site_ids = sqlite_session.scalars(select(Site.id).where(Site.app_id.in_(app_ids))).all() + assert len(site_ids) == len(app_ids) + assert len(set(site_ids)) == len(app_ids) + assert all(UUID(site_id).version == 4 for site_id in site_ids) + + def test_site_creation_with_required_fields(self): + """Test creating a site with required fields.""" + # Arrange + app_id = str(uuid4()) + + # Act + site = Site( + app_id=app_id, + title="Test Site", + default_language="en-US", + customize_token_strategy=CustomizeTokenStrategy.UUID, + ) + + # Assert + assert site.app_id == app_id + assert site.title == "Test Site" + assert site.default_language == "en-US" + assert site.customize_token_strategy == CustomizeTokenStrategy.UUID + + def test_site_creation_with_optional_fields(self): + """Test creating a site with optional fields.""" + # Arrange & Act + site = Site( + app_id=str(uuid4()), + title="Test Site", + default_language="en-US", + customize_token_strategy=CustomizeTokenStrategy.UUID, + icon_type=IconType.EMOJI, + icon="🌐", + icon_background="#0066CC", + description="Test site description", + copyright="© 2024 Test", + privacy_policy="https://example.com/privacy", + ) + + # Assert + assert site.icon_type == IconType.EMOJI + assert site.icon == "🌐" + assert site.icon_background == "#0066CC" + assert site.description == "Test site description" + assert site.copyright == "© 2024 Test" + assert site.privacy_policy == "https://example.com/privacy" + + def test_site_custom_disclaimer_setter(self): + """Test site custom_disclaimer setter.""" + # Arrange + site = Site( + app_id=str(uuid4()), + title="Test Site", + default_language="en-US", + customize_token_strategy=CustomizeTokenStrategy.UUID, + ) + + # Act + site.custom_disclaimer = "This is a test disclaimer" + + # Assert + assert site.custom_disclaimer == "This is a test disclaimer" + + def test_site_custom_disclaimer_exceeds_limit(self): + """Test site custom_disclaimer with excessive length.""" + # Arrange + site = Site( + app_id=str(uuid4()), + title="Test Site", + default_language="en-US", + customize_token_strategy=CustomizeTokenStrategy.UUID, + ) + long_disclaimer = "x" * 513 # Exceeds 512 character limit + + # Act & Assert + with pytest.raises(ValueError, match="Custom disclaimer cannot exceed 512 characters"): + site.custom_disclaimer = long_disclaimer + + class TestModelIntegration: + """Test suite for model integration scenarios.""" + + def test_complete_app_conversation_message_hierarchy(self): + """Test complete hierarchy from app to message.""" + # Arrange + tenant_id = str(uuid4()) + app_id = str(uuid4()) + conversation_id = str(uuid4()) + message_id = str(uuid4()) + created_by = str(uuid4()) + + # Create app + app = App( + tenant_id=tenant_id, + name="Test App", + mode=AppMode.CHAT, + enable_site=True, + enable_api=True, + created_by=created_by, + id=app_id, + ) + + # Create conversation + conversation = Conversation( + app_id=app_id, + mode=AppMode.CHAT, + name="Test Conversation", + status="normal", + from_source=ConversationFromSource.API, + from_end_user_id=str(uuid4()), + ) + conversation.id = conversation_id + + # Create message + message = Message( + app_id=app_id, + conversation_id=conversation_id, + query="Test query", + message={"role": "user", "content": "Test"}, + answer="Test answer", + message_unit_price=Decimal("0.0001"), + answer_unit_price=Decimal("0.0002"), + currency="USD", + from_source=ConversationFromSource.API, + id=message_id, + ) + + # Assert + assert app.id == app_id + assert conversation.app_id == app_id + assert message.app_id == app_id + assert message.conversation_id == conversation_id + assert app.mode == AppMode.CHAT + assert conversation.mode == AppMode.CHAT + + def test_app_with_annotation_setting(self): + """Test app with annotation setting.""" + # Arrange + app_id = str(uuid4()) + collection_binding_id = str(uuid4()) + created_user_id = str(uuid4()) + + # Create app + app = App( + tenant_id=str(uuid4()), + name="Test App", + mode=AppMode.CHAT, + enable_site=True, + enable_api=True, + created_by=created_user_id, + id=app_id, + ) + + # Create annotation setting + setting = AppAnnotationSetting( + app_id=app_id, + score_threshold=0.85, + collection_binding_id=collection_binding_id, + created_user_id=created_user_id, + updated_user_id=created_user_id, + ) + + # Assert + assert setting.app_id == app.id + assert setting.score_threshold == 0.85 + + def test_message_with_annotation(self): + """Test message with annotation.""" + # Arrange + app_id = str(uuid4()) + conversation_id = str(uuid4()) + message_id = str(uuid4()) + account_id = str(uuid4()) + + # Create message + message = Message( + app_id=app_id, + conversation_id=conversation_id, + query="What is AI?", + message={"role": "user", "content": "What is AI?"}, + answer="AI stands for Artificial Intelligence.", + message_unit_price=Decimal("0.0001"), + answer_unit_price=Decimal("0.0002"), + currency="USD", + from_source=ConversationFromSource.API, + id=message_id, + ) + + # Create annotation + annotation = MessageAnnotation( + app_id=app_id, + conversation_id=conversation_id, + message_id=message_id, + question="What is AI?", + content="AI stands for Artificial Intelligence.", + account_id=account_id, + ) + + # Assert + assert annotation.app_id == message.app_id + assert annotation.conversation_id == message.conversation_id + assert annotation.message_id == message.id + + def test_annotation_hit_history_tracking(self): + """Test annotation hit history tracking.""" + # Arrange + app_id = str(uuid4()) + annotation_id = str(uuid4()) + message_id = str(uuid4()) + account_id = str(uuid4()) + + # Create annotation + annotation = MessageAnnotation( + app_id=app_id, + question="What is AI?", + content="AI stands for Artificial Intelligence.", + conversation_id=None, + message_id=message_id, + account_id=account_id, + ) + annotation.id = annotation_id + + # Create hit history + history = AppAnnotationHitHistory( + app_id=app_id, + annotation_id=annotation_id, + source="api", + question="What is AI?", + account_id=account_id, + score=0.92, + message_id=message_id, + annotation_question="What is AI?", + annotation_content="AI stands for Artificial Intelligence.", + ) + + # Assert + assert history.app_id == annotation.app_id + assert history.annotation_id == annotation.id + assert history.score == 0.92 + + def test_app_with_site(self): + """Test app with site.""" + # Arrange + app_id = str(uuid4()) + + # Create app + app = App( + tenant_id=str(uuid4()), + name="Test App", + mode=AppMode.CHAT, + enable_site=True, + enable_api=True, + created_by=str(uuid4()), + id=app_id, + ) + + # Create site + site = Site( + app_id=app_id, + title="Test Site", + default_language="en-US", + customize_token_strategy=CustomizeTokenStrategy.UUID, + ) + + # Assert + assert site.app_id == app.id + assert app.enable_site is True + + +"""Regression coverage for the ``@property``→session-parameter refactor on ``App.tenant``. + +The legacy ``@property`` reached for the global ``db.session`` internally and has been converted +to a plain method taking an explicit ``session`` (per the pattern established in +#41830/#41831 for ``InstalledApp.tenant`` / ``ApiToolProvider.tenant``, tracked in #40372). +""" + + +def _app_session_accessors_persist_tenant(session: Session) -> Tenant: + tenant = Tenant(name="Test Tenant") + session.add(tenant) + session.flush() + return tenant + + +def _app_session_accessors_persist_app(session: Session, *, tenant_id: str) -> App: + app = App( + tenant_id=tenant_id, + name="Test App", + mode=AppMode.CHAT, + enable_site=True, + enable_api=False, + created_by=str(uuid4()), + ) + session.add(app) + session.flush() + return app + + +class TestAppSessionAccessorsCases: + class TestAppTenant: + def test_returns_the_matching_tenant(self, sqlite_session: Session) -> None: + tenant = _app_session_accessors_persist_tenant(sqlite_session) + app = _app_session_accessors_persist_app(sqlite_session, tenant_id=tenant.id) + + result = app.tenant(session=sqlite_session) + + assert result is not None + assert result.id == tenant.id + + def test_returns_none_when_tenant_missing(self, sqlite_session: Session) -> None: + app = _app_session_accessors_persist_app(sqlite_session, tenant_id=str(uuid4())) + + assert app.tenant(session=sqlite_session) is None + + +"""Regression coverage for the ``@property``→session-parameter refactor on +``InstalledApp.tenant``. + +The legacy ``@property`` reached for the global ``db.session`` internally and has been converted +to a plain method taking an explicit ``session: Session`` (per the pattern established in +#40370/#40797/#41394/#41830/#41885, tracked in #40372). +""" + + +def _installed_app_session_accessors_persist_tenant(session: Session) -> Tenant: + tenant = Tenant(name="Test Tenant") + session.add(tenant) + session.flush() + return tenant + + +def _persist_installed_app(session: Session, *, tenant_id: str) -> InstalledApp: + installed_app = InstalledApp( + tenant_id=tenant_id, + app_id=str(uuid4()), + app_owner_tenant_id=str(uuid4()), + position=0, + is_pinned=False, + last_used_at=None, + ) + session.add(installed_app) + session.flush() + return installed_app + + +class TestInstalledAppSessionAccessorsCases: + class TestInstalledAppTenant: + def test_returns_the_matching_tenant(self, sqlite_session: Session) -> None: + tenant = _installed_app_session_accessors_persist_tenant(sqlite_session) + installed_app = _persist_installed_app(sqlite_session, tenant_id=tenant.id) + + result = installed_app.tenant(session=sqlite_session) + + assert result is not None + assert result.id == tenant.id + + def test_returns_none_when_tenant_missing(self, sqlite_session: Session) -> None: + installed_app = _persist_installed_app(sqlite_session, tenant_id=str(uuid4())) + + assert installed_app.tenant(session=sqlite_session) is None + + +"""Regression coverage for the ``@property``→session-parameter refactor on Message +feedback/annotation accessors. + +Covers three of `api/models/model.py`'s legacy `@property` accessors that reached for the global +``db.session`` internally and have been removed in favour of their existing explicit-session +methods (per the pattern established in #40370/#40797/#41394/#41830, tracked in #40372): + +- ``MessageFeedback.from_account_with_session`` +- ``MessageAnnotation.account_with_session`` +- ``MessageAnnotation.annotation_create_account_with_session`` + +Each accessor is exercised against the real ``sqlite_session`` fixture (a genuine SQLAlchemy +``Session`` bound to a pristine full-schema SQLite database) so the assertions cover actual +query behaviour rather than a mock's recorded call. +""" + + +def _message_annotation_session_accessors_persist_account(session: Session, *, email: str | None = None) -> Account: + account = Account(name="Test Account", email=email or "test@example.com") + session.add(account) + session.flush() + return account + + +def _persist_feedback(session: Session, *, app_id: str, from_account_id: str | None) -> MessageFeedback: + feedback = MessageFeedback( + app_id=app_id, + conversation_id=str(uuid4()), + message_id=str(uuid4()), + rating=FeedbackRating.LIKE, + from_source=FeedbackFromSource.USER, + from_account_id=from_account_id, + ) + session.add(feedback) + session.flush() + return feedback + + +def _message_annotation_session_accessors_persist_annotation( + session: Session, *, app_id: str, account_id: str +) -> MessageAnnotation: + annotation = MessageAnnotation( + app_id=app_id, + conversation_id=None, + message_id=None, + question="What is AI?", + content="AI stands for Artificial Intelligence.", + account_id=account_id, + ) + session.add(annotation) + session.flush() + return annotation + + +class TestMessageAnnotationSessionAccessorsCases: + class TestMessageFeedbackFromAccount: + def test_returns_the_feedback_author(self, sqlite_session: Session) -> None: + account = _message_annotation_session_accessors_persist_account(sqlite_session) + feedback = _persist_feedback(sqlite_session, app_id=str(uuid4()), from_account_id=account.id) + + result = feedback.from_account_with_session(session=sqlite_session) + + assert result is not None + assert result.id == account.id + + def test_returns_none_when_account_missing(self, sqlite_session: Session) -> None: + feedback = _persist_feedback(sqlite_session, app_id=str(uuid4()), from_account_id=str(uuid4())) + + assert feedback.from_account_with_session(session=sqlite_session) is None + + class TestMessageAnnotationAccount: + def test_returns_the_annotation_creator(self, sqlite_session: Session) -> None: + creator = _message_annotation_session_accessors_persist_account(sqlite_session, email="creator@example.com") + annotation = _message_annotation_session_accessors_persist_annotation( + sqlite_session, app_id=str(uuid4()), account_id=creator.id + ) + + result = annotation.account_with_session(session=sqlite_session) + + assert result is not None + assert result.id == creator.id + + def test_returns_none_when_account_missing(self, sqlite_session: Session) -> None: + annotation = _message_annotation_session_accessors_persist_annotation( + sqlite_session, app_id=str(uuid4()), account_id=str(uuid4()) + ) + + assert annotation.account_with_session(session=sqlite_session) is None + + class TestMessageAnnotationAnnotationCreateAccount: + def test_returns_the_annotation_creator(self, sqlite_session: Session) -> None: + creator = _message_annotation_session_accessors_persist_account( + sqlite_session, email="creator2@example.com" + ) + annotation = _message_annotation_session_accessors_persist_annotation( + sqlite_session, app_id=str(uuid4()), account_id=creator.id + ) + + result = annotation.annotation_create_account_with_session(session=sqlite_session) + + assert result is not None + assert result.id == creator.id + + def test_returns_none_when_account_missing(self, sqlite_session: Session) -> None: + annotation = _message_annotation_session_accessors_persist_annotation( + sqlite_session, app_id=str(uuid4()), account_id=str(uuid4()) + ) + + assert annotation.annotation_create_account_with_session(session=sqlite_session) is None + + +"""Keep the persisted provider through private trace JSON and the real consumer.""" + + +_MESSAGE = TypeAdapter(MessageDict) + + +class TestMessageTraceSerializationCases: + @staticmethod + @pytest.fixture + def persisted_message(sqlite_session: Session) -> Message: + now = datetime(2026, 10, 20, tzinfo=UTC) + message = make_message( + message_id=str(uuid4()), + app_id=str(uuid4()), + conversation_id=str(uuid4()), + inputs={}, + query="Synthetic input", + message={}, + answer="Synthetic output", + message_unit_price=Decimal(0), + answer_unit_price=Decimal(0), + total_price=Decimal(0), + currency="USD", + from_source=ConversationFromSource.CONSOLE, + created_at=now, + updated_at=now + timedelta(seconds=1), + ) + message.model_provider = "langgenius/openai/openai" + message.model_id = "gpt-4o-mini" + sqlite_session.add(message) + sqlite_session.commit() + sqlite_session.expire_all() + restored = sqlite_session.get(Message, message.id) + assert restored is not None + return restored + + @staticmethod + @pytest.mark.parametrize("provider", ["langgenius/openai/openai", "Custom/Provider:Exact", None]) + def test_message_provider_survives_real_json_round_trip( + sqlite_session: Session, + persisted_message: Message, + provider: str | None, + ) -> None: + persisted_message.model_provider = provider + sqlite_session.commit() + sqlite_session.refresh(persisted_message) + wire = json.dumps(persisted_message.to_dict(session=sqlite_session), default=str) + payload = _MESSAGE.validate_json(wire) + restored = Message.from_dict(payload) + assert "model_provider" in payload + assert restored.model_provider == provider + assert restored.model_id == persisted_message.model_id + assert restored.id == persisted_message.id + + @staticmethod + @pytest.mark.parametrize("shape", ["current", "legacy_missing", "explicit_null"]) + def test_private_trace_task_preserves_provider_into_real_arize_span( + sqlite_session: Session, + persisted_message: Message, + monkeypatch: pytest.MonkeyPatch, + shape: str, + ) -> None: + message_data = persisted_message.to_dict(session=sqlite_session) + if shape == "legacy_missing": + message_data.pop("model_provider", None) + elif shape == "explicit_null": + message_data["model_provider"] = None + start = datetime(2026, 10, 20, tzinfo=UTC) + info = MessageTraceInfo( + message_id=persisted_message.id, + message_data=message_data, + conversation_model="chat", + conversation_mode="chat", + inputs={"query": "Synthetic input"}, + outputs="Synthetic output", + start_time=start, + end_time=start + timedelta(seconds=1), + metadata={}, + message_tokens=1, + answer_tokens=1, + total_tokens=2, + ) + encoded = ( + TaskData(app_id=persisted_message.app_id, trace_info_type="MessageTraceInfo", trace_info=info.model_dump()) + .model_dump_json() + .encode() + ) + exporter = InMemorySpanExporter() + provider = TracerProvider() + processor = SimpleSpanProcessor(exporter) + provider.add_span_processor(processor) + + def setup(_config: ArizeConfig) -> tuple[Tracer, SimpleSpanProcessor]: + tracer = provider.get_tracer("message-provider-serialization") + assert isinstance(tracer, Tracer) + return tracer, processor + + monkeypatch.setattr(arize_phoenix_trace, "setup_tracer", setup) + instance = ArizePhoenixDataTrace(ArizeConfig(api_key="unused-synthetic-key", space_id="synthetic-space")) + paths: list[str] = [] + deleted: list[str] = [] + + def load(path: str) -> bytes: + paths.append(path) + return encoded + + def delete(path: str) -> None: + deleted.append(path) + + def lookup(app_id: str) -> ArizePhoenixDataTrace: + assert app_id == persisted_message.app_id + return instance + + def failed(_key: str) -> None: + pytest.fail("The real trace consumer must not hide a failed dispatch") + + apply_config_overrides(monkeypatch, OPS_AGENT_TRACE_ENABLED=False) + monkeypatch.setattr(ops_trace_task.storage, "load", load) + monkeypatch.setattr(ops_trace_task.storage, "delete", delete) + monkeypatch.setattr(ops_trace_task.redis_client, "incr", failed) + monkeypatch.setattr(OpsTraceManager, "get_ops_trace_instance", lookup) + monkeypatch.setattr("extensions.ext_enterprise_telemetry.is_enabled", lambda: False) + try: + ops_trace_task.process_trace_tasks.apply( + args=({"app_id": persisted_message.app_id, "file_id": "owned-synthetic"},), throw=True + ).get() + spans = exporter.get_finished_spans() + llm = [span for span in spans if span.name == "llm"] + assert len(llm) == 1 + attributes = llm[0].attributes + assert attributes is not None + assert attributes[SpanAttributes.LLM_MODEL_NAME] == persisted_message.model_id + if shape == "current": + assert attributes[SpanAttributes.LLM_PROVIDER] == persisted_message.model_provider + else: + assert SpanAttributes.LLM_PROVIDER not in attributes + assert paths == [f"ops_trace/{persisted_message.app_id}/owned-synthetic.json"] + assert deleted == paths + finally: + provider.shutdown() def _wrap_md(url: str) -> str: @@ -29,64 +2159,6 @@ def _wrap_md(url: str) -> str: return f"please click [file]({url}) to download." -def test_file_preview_valid_replaced() -> None: - """ - Valid file-preview URL must be re-signed: - - Extract upload_file_id correctly - - Replace the original URL with the signed URL - """ - upload_id = "abc-123" - url = f"/files/{upload_id}/file-preview?timestamp=111&nonce=222&sign=333" - msg = Message(answer=_wrap_md(url)) - - out = msg.re_sign_file_url_answer - assert f"https://signed.example/{upload_id}" in out - assert url not in out - - -def test_file_preview_misspelled_not_replaced() -> None: - """ - Misspelled endpoint 'file-previe?timestamp=' should NOT be rewritten. - """ - upload_id = "zzz-001" - # path deliberately misspelled: file-previe? (missing 'w') - # and we append ¬e=file-preview to trick the old `"file-preview" in url` check. - url = f"/files/{upload_id}/file-previe?timestamp=111&nonce=222&sign=333¬e=file-preview" - original = _wrap_md(url) - msg = Message(answer=original) - - out = msg.re_sign_file_url_answer - # Expect NO replacement, should not rewrite misspelled file-previe URL - assert out == original - - -def test_image_preview_valid_replaced() -> None: - """ - Valid image-preview URL must be re-signed. - """ - upload_id = "img-789" - url = f"/files/{upload_id}/image-preview?timestamp=123&nonce=456&sign=789" - msg = Message(answer=_wrap_md(url)) - - out = msg.re_sign_file_url_answer - assert f"https://signed.example/{upload_id}" in out - assert url not in out - - -def test_image_preview_misspelled_not_replaced() -> None: - """ - Misspelled endpoint 'image-previe?timestamp=' should NOT be rewritten. - """ - upload_id = "img-err-42" - url = f"/files/{upload_id}/image-previe?timestamp=1&nonce=2&sign=3¬e=image-preview" - original = _wrap_md(url) - msg = Message(answer=original) - - out = msg.re_sign_file_url_answer - # Expect NO replacement, should not rewrite misspelled image-previe URL - assert out == original - - def _build_local_file_mapping(record_id: str, *, tenant_id: str | None = None) -> dict[str, object]: mapping: dict[str, object] = { "dify_model_identity": FILE_MODEL_IDENTITY, @@ -103,119 +2175,373 @@ def _build_local_file_mapping(record_id: str, *, tenant_id: str | None = None) - return mapping -@pytest.mark.parametrize("owner_cls", [Conversation, Message]) -def test_inputs_restore_external_remote_url_file_mappings( - owner_cls: type[Conversation] | type[Message], sqlite_session: Session -) -> None: - owner = owner_cls(app_id="app-1") - owner.inputs = { - "file": { - "dify_model_identity": FILE_MODEL_IDENTITY, - "transfer_method": FileTransferMethod.REMOTE_URL, - "type": "document", - "url": "https://example.com/report.pdf", - "filename": "report.pdf", - "extension": ".pdf", - "mime_type": "application/pdf", - "size": 1, +class TestModelCases: + @staticmethod + @pytest.fixture(autouse=True) + def patch_file_helpers(monkeypatch: pytest.MonkeyPatch) -> None: + """ + Patch file_helpers.get_signed_file_url to a deterministic stub. + """ + monkeypatch.setattr( + model_module.file_helpers, + "get_signed_file_url", + lambda file_id: f"https://signed.example/{file_id}", + ) + + @staticmethod + def test_file_preview_valid_replaced() -> None: + """ + Valid file-preview URL must be re-signed: + - Extract upload_file_id correctly + - Replace the original URL with the signed URL + """ + upload_id = "abc-123" + url = f"/files/{upload_id}/file-preview?timestamp=111&nonce=222&sign=333" + msg = Message(answer=_wrap_md(url)) + + out = msg.re_sign_file_url_answer + assert f"https://signed.example/{upload_id}" in out + assert url not in out + + @staticmethod + def test_file_preview_misspelled_not_replaced() -> None: + """ + Misspelled endpoint 'file-previe?timestamp=' should NOT be rewritten. + """ + upload_id = "zzz-001" + # path deliberately misspelled: file-previe? (missing 'w') + # and we append ¬e=file-preview to trick the old `"file-preview" in url` check. + url = f"/files/{upload_id}/file-previe?timestamp=111&nonce=222&sign=333¬e=file-preview" + original = _wrap_md(url) + msg = Message(answer=original) + + out = msg.re_sign_file_url_answer + # Expect NO replacement, should not rewrite misspelled file-previe URL + assert out == original + + @staticmethod + def test_image_preview_valid_replaced() -> None: + """ + Valid image-preview URL must be re-signed. + """ + upload_id = "img-789" + url = f"/files/{upload_id}/image-preview?timestamp=123&nonce=456&sign=789" + msg = Message(answer=_wrap_md(url)) + + out = msg.re_sign_file_url_answer + assert f"https://signed.example/{upload_id}" in out + assert url not in out + + @staticmethod + def test_image_preview_misspelled_not_replaced() -> None: + """ + Misspelled endpoint 'image-previe?timestamp=' should NOT be rewritten. + """ + upload_id = "img-err-42" + url = f"/files/{upload_id}/image-previe?timestamp=1&nonce=2&sign=3¬e=image-preview" + original = _wrap_md(url) + msg = Message(answer=original) + + out = msg.re_sign_file_url_answer + # Expect NO replacement, should not rewrite misspelled image-previe URL + assert out == original + + @staticmethod + @pytest.mark.parametrize("owner_cls", [Conversation, Message]) + def test_inputs_restore_external_remote_url_file_mappings( + owner_cls: type[Conversation] | type[Message], sqlite_session: Session + ) -> None: + owner = owner_cls(app_id="app-1") + owner.inputs = { + "file": { + "dify_model_identity": FILE_MODEL_IDENTITY, + "transfer_method": FileTransferMethod.REMOTE_URL, + "type": "document", + "url": "https://example.com/report.pdf", + "filename": "report.pdf", + "extension": ".pdf", + "mime_type": "application/pdf", + "size": 1, + } } - } - restored_file = owner.inputs_with_session(session=sqlite_session)["file"] + restored_file = owner.inputs_with_session(session=sqlite_session)["file"] - assert restored_file.transfer_method == FileTransferMethod.REMOTE_URL - assert restored_file.remote_url == "https://example.com/report.pdf" + assert restored_file.transfer_method == FileTransferMethod.REMOTE_URL + assert restored_file.remote_url == "https://example.com/report.pdf" + + @staticmethod + @pytest.mark.parametrize("sqlite_session", [(App,)], indirect=True) + def test_message_inputs_resolve_file_tenant_with_caller_session(sqlite_session: Session) -> None: + app = App( + id="app-1", + tenant_id="tenant-1", + name="File owner", + description="", + mode=AppMode.CHAT, + icon_type=IconType.EMOJI, + icon="file", + icon_background="#FFFFFF", + enable_site=False, + enable_api=False, + max_active_requests=0, + ) + decoy = App( + id="other-app", + tenant_id="other-tenant", + name="Decoy", + description="", + mode=AppMode.CHAT, + icon_type=IconType.EMOJI, + icon="file", + icon_background="#FFFFFF", + enable_site=False, + enable_api=False, + max_active_requests=0, + ) + sqlite_session.add_all([decoy, app]) + sqlite_session.flush() + message = Message(app_id="app-1") + message.inputs = {"file": _build_local_file_mapping("upload-1")} + + with patch( + "models.model.build_file_from_input_mapping", + side_effect=lambda **kwargs: kwargs["tenant_resolver"](), + ): + inputs = message.inputs_with_session(session=sqlite_session) + + assert inputs["file"] == "tenant-1" + + # session.scalar.assert_called_once() # removed: see #40799 conftest note + + @staticmethod + def test_file_url_bare_url_re_signed() -> None: + """#40788: bare tool file URL (no markdown wrapping) must be re-signed. + + Bare URLs are produced by the agent runtime when the model returns a + non-linkified file reference. The previous regex only matched + [text](url) markdown form, so the bare URL kept INTERNAL_FILES_URL. + """ + upload_id = "bare-1" + url = f"/files/{upload_id}/file-preview?timestamp=1&nonce=2&sign=3" + msg = Message(answer=url) # bare, no markdown wrapping + + out = msg.re_sign_file_url_answer + assert f"https://signed.example/{upload_id}" in out + assert url not in out + + @staticmethod + def test_file_url_backticked_url_re_signed() -> None: + """#40788: backticked tool file URL must be re-signed. + + Backticks are used by the runtime to quote file references inline. + The previous regex only matched [text](url), so backticked URLs + stayed on INTERNAL_FILES_URL. + """ + upload_id = "tick-2" + url = f"/files/{upload_id}/file-preview?timestamp=10&nonce=20&sign=30" + msg = Message(answer=f"see `{url}` for details") + + out = msg.re_sign_file_url_answer + assert f"https://signed.example/{upload_id}" in out + assert url not in out + + @staticmethod + def test_file_url_mixed_formats_all_re_signed() -> None: + """#40788: all three URL shapes in the same answer are re-signed. + + Sanity check that the three patterns cooperate — bare, backticked, + and markdown — without double-signing or skipping any. + """ + upload_ids = ["mix-a", "mix-b", "mix-c"] + bare = f"/files/{upload_ids[0]}/file-preview?timestamp=1&nonce=2&sign=3" + tick = f"/files/{upload_ids[1]}/file-preview?timestamp=10&nonce=20&sign=30" + md = f"/files/{upload_ids[2]}/file-preview?timestamp=100&nonce=200&sign=300" + msg = Message(answer=f"raw {bare} ` {tick} ` and [file]({md})") + + out = msg.re_sign_file_url_answer + for uid in upload_ids: + assert f"https://signed.example/{uid}" in out + for url in (bare, tick, md): + assert url not in out -@pytest.mark.parametrize("sqlite_session", [(App,)], indirect=True) -def test_message_inputs_resolve_file_tenant_with_caller_session(sqlite_session: Session) -> None: +"""Regression coverage for the ``@property``→session-parameter refactor on +``RecommendedApp`` and ``AccountTrialAppRecord``. + +The legacy ``@property`` accessors reached for the global ``db.session`` internally and have +been converted to plain methods taking an explicit ``session: Session`` (per the pattern +established in #40370/#41394/#41830/#42000, tracked in #40372). +""" + + +def _recommended_app_session_accessors_persist_app(session: Session) -> App: + app = App( + tenant_id=str(uuid4()), + name="Test App", + mode=AppMode.CHAT, + enable_site=True, + enable_api=False, + created_by=str(uuid4()), + description="App description", + ) + session.add(app) + session.flush() + return app + + +def _persist_recommended_app(session: Session, *, app_id: str) -> RecommendedApp: + recommended_app = RecommendedApp( + app_id=app_id, + description={"en_US": "A recommended app"}, + copyright="Dify", + privacy_policy="https://dify.ai/privacy", + category="Assistant", + ) + session.add(recommended_app) + session.flush() + return recommended_app + + +def _recommended_app_session_accessors_persist_account(session: Session) -> Account: + account = Account(name="Test User", email=f"user-{uuid4()}@example.com") + session.add(account) + session.flush() + return account + + +def _persist_trial_record(session: Session, *, account_id: str, app_id: str) -> AccountTrialAppRecord: + record = AccountTrialAppRecord(account_id=account_id, app_id=app_id, count=1) + session.add(record) + session.flush() + return record + + +class TestRecommendedAppSessionAccessorsCases: + class TestRecommendedAppApp: + def test_returns_the_matching_app(self, sqlite_session: Session) -> None: + app = _recommended_app_session_accessors_persist_app(sqlite_session) + recommended_app = _persist_recommended_app(sqlite_session, app_id=app.id) + + result = recommended_app.app(session=sqlite_session) + + assert result is not None + assert result.id == app.id + + def test_returns_none_when_app_missing(self, sqlite_session: Session) -> None: + recommended_app = _persist_recommended_app(sqlite_session, app_id=str(uuid4())) + + assert recommended_app.app(session=sqlite_session) is None + + class TestAccountTrialAppRecordAccessors: + def test_app_returns_the_matching_app(self, sqlite_session: Session) -> None: + app = _recommended_app_session_accessors_persist_app(sqlite_session) + record = _persist_trial_record(sqlite_session, account_id=str(uuid4()), app_id=app.id) + + result = record.app(session=sqlite_session) + + assert result is not None + assert result.id == app.id + + def test_app_returns_none_when_app_missing(self, sqlite_session: Session) -> None: + record = _persist_trial_record(sqlite_session, account_id=str(uuid4()), app_id=str(uuid4())) + + assert record.app(session=sqlite_session) is None + + def test_user_returns_the_matching_account(self, sqlite_session: Session) -> None: + account = _recommended_app_session_accessors_persist_account(sqlite_session) + record = _persist_trial_record(sqlite_session, account_id=account.id, app_id=str(uuid4())) + + result = record.user(session=sqlite_session) + + assert result is not None + assert result.id == account.id + + def test_user_returns_none_when_account_missing(self, sqlite_session: Session) -> None: + record = _persist_trial_record(sqlite_session, account_id=str(uuid4()), app_id=str(uuid4())) + + assert record.user(session=sqlite_session) is None + + +TARGET_MODELS = (App, AppStar, Conversation, Message, EndUser, ApiToken) + + +def test_app_conversation_models_use_typebase_registry() -> None: + for model in TARGET_MODELS: + mapper = inspect(model) + assert mapper.registry is TypeBase.registry + assert mapper.registry is not Base.registry + + +def test_generated_identifiers_and_collection_defaults_are_independent() -> None: + first = Conversation( + app_id="app-1", + mode=AppMode.CHAT, + name="First", + from_source=ConversationFromSource.API, + ) + second = Conversation( + app_id="app-1", + mode=AppMode.CHAT, + name="Second", + from_source=ConversationFromSource.API, + ) + + assert first.id != second.id + assert first._inputs == second._inputs == {} + assert first._inputs is not second._inputs + + +@pytest.mark.parametrize("sqlite_session", [TARGET_MODELS], indirect=True) +def test_app_conversation_models_persist_with_relationships(sqlite_session: Session) -> None: app = App( id="app-1", tenant_id="tenant-1", - name="File owner", - description="", + name="Typed app", mode=AppMode.CHAT, - icon_type=IconType.EMOJI, - icon="file", - icon_background="#FFFFFF", - enable_site=False, - enable_api=False, - max_active_requests=0, + enable_site=True, + enable_api=True, ) - decoy = App( - id="other-app", - tenant_id="other-tenant", - name="Decoy", - description="", + star = AppStar(id="star-1", tenant_id=app.tenant_id, app_id=app.id, account_id="account-1") + conversation = Conversation( + id="conversation-1", + app_id=app.id, mode=AppMode.CHAT, - icon_type=IconType.EMOJI, - icon="file", - icon_background="#FFFFFF", - enable_site=False, - enable_api=False, - max_active_requests=0, + name="Typed conversation", + from_source=ConversationFromSource.API, + ) + message = Message( + id="message-1", + app_id=app.id, + conversation_id=conversation.id, + query="hello", + message={"role": "user"}, + answer="world", + from_source=ConversationFromSource.API, + total_price=Decimal(0), + ) + end_user = EndUser( + id="end-user-1", + tenant_id=app.tenant_id, + app_id=app.id, + type=EndUserType.BROWSER, + session_id="session-1", + ) + token = ApiToken( + id="token-1", + app_id=app.id, + tenant_id=app.tenant_id, + type=ApiTokenType.APP, + token="secret", ) - sqlite_session.add_all([decoy, app]) - sqlite_session.flush() - message = Message(app_id="app-1") - message.inputs = {"file": _build_local_file_mapping("upload-1")} - with patch( - "models.model.build_file_from_input_mapping", - side_effect=lambda **kwargs: kwargs["tenant_resolver"](), - ): - inputs = message.inputs_with_session(session=sqlite_session) + sqlite_session.add_all([app, star, conversation, message, end_user, token]) + sqlite_session.commit() - assert inputs["file"] == "tenant-1" - # session.scalar.assert_called_once() # removed: see #40799 conftest note - - -def test_file_url_bare_url_re_signed() -> None: - """#40788: bare tool file URL (no markdown wrapping) must be re-signed. - - Bare URLs are produced by the agent runtime when the model returns a - non-linkified file reference. The previous regex only matched - [text](url) markdown form, so the bare URL kept INTERNAL_FILES_URL. - """ - upload_id = "bare-1" - url = f"/files/{upload_id}/file-preview?timestamp=1&nonce=2&sign=3" - msg = Message(answer=url) # bare, no markdown wrapping - - out = msg.re_sign_file_url_answer - assert f"https://signed.example/{upload_id}" in out - assert url not in out - - -def test_file_url_backticked_url_re_signed() -> None: - """#40788: backticked tool file URL must be re-signed. - - Backticks are used by the runtime to quote file references inline. - The previous regex only matched [text](url), so backticked URLs - stayed on INTERNAL_FILES_URL. - """ - upload_id = "tick-2" - url = f"/files/{upload_id}/file-preview?timestamp=10&nonce=20&sign=30" - msg = Message(answer=f"see `{url}` for details") - - out = msg.re_sign_file_url_answer - assert f"https://signed.example/{upload_id}" in out - assert url not in out - - -def test_file_url_mixed_formats_all_re_signed() -> None: - """#40788: all three URL shapes in the same answer are re-signed. - - Sanity check that the three patterns cooperate — bare, backticked, - and markdown — without double-signing or skipping any. - """ - upload_ids = ["mix-a", "mix-b", "mix-c"] - bare = f"/files/{upload_ids[0]}/file-preview?timestamp=1&nonce=2&sign=3" - tick = f"/files/{upload_ids[1]}/file-preview?timestamp=10&nonce=20&sign=30" - md = f"/files/{upload_ids[2]}/file-preview?timestamp=100&nonce=200&sign=300" - msg = Message(answer=f"raw {bare} ` {tick} ` and [file]({md})") - - out = msg.re_sign_file_url_answer - for uid in upload_ids: - assert f"https://signed.example/{uid}" in out - for url in (bare, tick, md): - assert url not in out + persisted_conversation = sqlite_session.get(Conversation, conversation.id) + assert persisted_conversation is not None + assert [item.id for item in persisted_conversation.messages] == [message.id] + assert sqlite_session.get(AppStar, star.id) is not None + assert sqlite_session.get(EndUser, end_user.id) is not None + assert sqlite_session.get(ApiToken, token.id) is not None + assert message.created_at is not None diff --git a/api/tests/unit_tests/models/test_plugin_entities.py b/api/tests/unit_tests/models/test_plugin_entities.py deleted file mode 100644 index ace30618f0b..00000000000 --- a/api/tests/unit_tests/models/test_plugin_entities.py +++ /dev/null @@ -1,21 +0,0 @@ -import binascii -from collections.abc import Mapping - -from core.plugin.entities.request import TriggerDispatchResponse - - -def test_trigger_dispatch_response() -> None: - raw_http_response = b'HTTP/1.1 200 OK\r\nContent-Type: application/json\r\n\r\n{"message": "Hello, world!"}' - - data: Mapping[str, object] = { - "user_id": "123", - "events": ["event1", "event2"], - "response": binascii.hexlify(raw_http_response).decode(), - "payload": {"key": "value"}, - } - - response = TriggerDispatchResponse.model_validate(data) - - assert response.response.status_code == 200 - assert response.response.headers["Content-Type"] == "application/json" - assert response.response.get_data(as_text=True) == '{"message": "Hello, world!"}' diff --git a/api/tests/unit_tests/models/test_provider_models.py b/api/tests/unit_tests/models/test_provider.py similarity index 100% rename from api/tests/unit_tests/models/test_provider_models.py rename to api/tests/unit_tests/models/test_provider.py diff --git a/api/tests/unit_tests/models/test_recommended_app_session_accessors.py b/api/tests/unit_tests/models/test_recommended_app_session_accessors.py deleted file mode 100644 index e560144b66f..00000000000 --- a/api/tests/unit_tests/models/test_recommended_app_session_accessors.py +++ /dev/null @@ -1,102 +0,0 @@ -"""Regression coverage for the ``@property``→session-parameter refactor on -``RecommendedApp`` and ``AccountTrialAppRecord``. - -The legacy ``@property`` accessors reached for the global ``db.session`` internally and have -been converted to plain methods taking an explicit ``session: Session`` (per the pattern -established in #40370/#41394/#41830/#42000, tracked in #40372). -""" - -from uuid import uuid4 - -from sqlalchemy.orm import Session - -from models.account import Account -from models.model import AccountTrialAppRecord, App, AppMode, RecommendedApp - - -def _persist_app(session: Session) -> App: - app = App( - tenant_id=str(uuid4()), - name="Test App", - mode=AppMode.CHAT, - enable_site=True, - enable_api=False, - created_by=str(uuid4()), - description="App description", - ) - session.add(app) - session.flush() - return app - - -def _persist_recommended_app(session: Session, *, app_id: str) -> RecommendedApp: - recommended_app = RecommendedApp( - app_id=app_id, - description={"en_US": "A recommended app"}, - copyright="Dify", - privacy_policy="https://dify.ai/privacy", - category="Assistant", - ) - session.add(recommended_app) - session.flush() - return recommended_app - - -def _persist_account(session: Session) -> Account: - account = Account(name="Test User", email=f"user-{uuid4()}@example.com") - session.add(account) - session.flush() - return account - - -def _persist_trial_record(session: Session, *, account_id: str, app_id: str) -> AccountTrialAppRecord: - record = AccountTrialAppRecord(account_id=account_id, app_id=app_id, count=1) - session.add(record) - session.flush() - return record - - -class TestRecommendedAppApp: - def test_returns_the_matching_app(self, sqlite_session: Session) -> None: - app = _persist_app(sqlite_session) - recommended_app = _persist_recommended_app(sqlite_session, app_id=app.id) - - result = recommended_app.app(session=sqlite_session) - - assert result is not None - assert result.id == app.id - - def test_returns_none_when_app_missing(self, sqlite_session: Session) -> None: - recommended_app = _persist_recommended_app(sqlite_session, app_id=str(uuid4())) - - assert recommended_app.app(session=sqlite_session) is None - - -class TestAccountTrialAppRecordAccessors: - def test_app_returns_the_matching_app(self, sqlite_session: Session) -> None: - app = _persist_app(sqlite_session) - record = _persist_trial_record(sqlite_session, account_id=str(uuid4()), app_id=app.id) - - result = record.app(session=sqlite_session) - - assert result is not None - assert result.id == app.id - - def test_app_returns_none_when_app_missing(self, sqlite_session: Session) -> None: - record = _persist_trial_record(sqlite_session, account_id=str(uuid4()), app_id=str(uuid4())) - - assert record.app(session=sqlite_session) is None - - def test_user_returns_the_matching_account(self, sqlite_session: Session) -> None: - account = _persist_account(sqlite_session) - record = _persist_trial_record(sqlite_session, account_id=account.id, app_id=str(uuid4())) - - result = record.user(session=sqlite_session) - - assert result is not None - assert result.id == account.id - - def test_user_returns_none_when_account_missing(self, sqlite_session: Session) -> None: - record = _persist_trial_record(sqlite_session, account_id=str(uuid4()), app_id=str(uuid4())) - - assert record.user(session=sqlite_session) is None diff --git a/api/tests/unit_tests/models/test_tool_provider_session_accessors.py b/api/tests/unit_tests/models/test_tool_provider_session_accessors.py deleted file mode 100644 index a5f906121f0..00000000000 --- a/api/tests/unit_tests/models/test_tool_provider_session_accessors.py +++ /dev/null @@ -1,172 +0,0 @@ -"""Regression coverage for the ``@property``→session-parameter refactor on tool provider accessors. - -Covers four of `api/models/tools.py`'s legacy `@property` accessors that reached for the global -``db.session`` internally and have been converted to plain methods taking an explicit -``session: Session`` (per the pattern established in #40370/#40797/#41394, tracked in #40372): - -- ``ApiToolProvider.tenant`` -- ``ApiToolProvider.user`` -- ``WorkflowToolProvider.user`` -- ``WorkflowToolProvider.tenant`` -- ``WorkflowToolProvider.app`` - -Each accessor is exercised against the real ``sqlite_session`` fixture (a genuine SQLAlchemy -``Session`` bound to a pristine full-schema SQLite database) so the assertions cover actual -query behaviour rather than a mock's recorded call. -""" - -from uuid import uuid4 - -from sqlalchemy.orm import Session - -from core.tools.entities.tool_entities import ApiProviderSchemaType -from models.account import Account, Tenant -from models.model import App, AppMode -from models.tools import ApiToolProvider, WorkflowToolProvider - - -def _persist_account(session: Session) -> Account: - account = Account(name="Test Account", email="test@example.com") - session.add(account) - session.flush() - return account - - -def _persist_tenant(session: Session) -> Tenant: - tenant = Tenant(name="Test Tenant") - session.add(tenant) - session.flush() - return tenant - - -def _persist_app(session: Session, *, tenant_id: str) -> App: - app = App( - tenant_id=tenant_id, - name="Test App", - mode=AppMode.WORKFLOW, - enable_site=True, - enable_api=True, - created_by=str(uuid4()), - ) - session.add(app) - session.flush() - return app - - -def _api_tool_provider(*, tenant_id: str) -> ApiToolProvider: - return ApiToolProvider( - tenant_id=tenant_id, - user_id=str(uuid4()), - name="Test API Provider", - icon="{}", - schema="{}", - schema_type_str=ApiProviderSchemaType.OPENAPI, - description="", - tools_str="[]", - credentials_str="{}", - ) - - -def _workflow_tool_provider(*, tenant_id: str, app_id: str, user_id: str) -> WorkflowToolProvider: - return WorkflowToolProvider( - name="Test Workflow Provider", - label="Test Workflow Provider", - icon="{}", - app_id=app_id, - version="1", - user_id=user_id, - tenant_id=tenant_id, - description="", - ) - - -class TestApiToolProviderTenant: - def test_returns_persisted_tenant(self, sqlite_session: Session) -> None: - tenant = _persist_tenant(sqlite_session) - provider = _api_tool_provider(tenant_id=tenant.id) - - result = provider.tenant(session=sqlite_session) - - assert result is not None - assert result.id == tenant.id - - def test_returns_none_when_tenant_missing(self, sqlite_session: Session) -> None: - provider = _api_tool_provider(tenant_id=str(uuid4())) - - assert provider.tenant(session=sqlite_session) is None - - -class TestApiToolProviderUser: - def test_returns_persisted_user(self, sqlite_session: Session) -> None: - account = _persist_account(sqlite_session) - provider = _api_tool_provider(tenant_id=str(uuid4())) - provider.user_id = account.id - sqlite_session.add(provider) - sqlite_session.flush() - - result = provider.user(session=sqlite_session) - - assert result is not None - assert result.id == account.id - - def test_returns_none_when_user_missing(self, sqlite_session: Session) -> None: - provider = _api_tool_provider(tenant_id=str(uuid4())) - - assert provider.user(session=sqlite_session) is None - - -class TestWorkflowToolProviderUser: - def test_returns_persisted_user(self, sqlite_session: Session) -> None: - tenant = _persist_tenant(sqlite_session) - app = _persist_app(sqlite_session, tenant_id=tenant.id) - account = _persist_account(sqlite_session) - provider = _workflow_tool_provider(tenant_id=tenant.id, app_id=app.id, user_id=account.id) - - result = provider.user(session=sqlite_session) - - assert result is not None - assert result.id == account.id - - def test_returns_none_when_user_missing(self, sqlite_session: Session) -> None: - tenant = _persist_tenant(sqlite_session) - app = _persist_app(sqlite_session, tenant_id=tenant.id) - provider = _workflow_tool_provider(tenant_id=tenant.id, app_id=app.id, user_id=str(uuid4())) - - assert provider.user(session=sqlite_session) is None - - -class TestWorkflowToolProviderTenant: - def test_returns_persisted_tenant(self, sqlite_session: Session) -> None: - tenant = _persist_tenant(sqlite_session) - app = _persist_app(sqlite_session, tenant_id=tenant.id) - provider = _workflow_tool_provider(tenant_id=tenant.id, app_id=app.id, user_id=str(uuid4())) - - result = provider.tenant(session=sqlite_session) - - assert result is not None - assert result.id == tenant.id - - def test_returns_none_when_tenant_missing(self, sqlite_session: Session) -> None: - tenant = _persist_tenant(sqlite_session) - app = _persist_app(sqlite_session, tenant_id=tenant.id) - provider = _workflow_tool_provider(tenant_id=str(uuid4()), app_id=app.id, user_id=str(uuid4())) - - assert provider.tenant(session=sqlite_session) is None - - -class TestWorkflowToolProviderApp: - def test_returns_persisted_app(self, sqlite_session: Session) -> None: - tenant = _persist_tenant(sqlite_session) - app = _persist_app(sqlite_session, tenant_id=tenant.id) - provider = _workflow_tool_provider(tenant_id=tenant.id, app_id=app.id, user_id=str(uuid4())) - - result = provider.app(session=sqlite_session) - - assert result is not None - assert result.id == app.id - - def test_returns_none_when_app_missing(self, sqlite_session: Session) -> None: - tenant = _persist_tenant(sqlite_session) - provider = _workflow_tool_provider(tenant_id=tenant.id, app_id=str(uuid4()), user_id=str(uuid4())) - - assert provider.app(session=sqlite_session) is None diff --git a/api/tests/unit_tests/models/test_tool_models.py b/api/tests/unit_tests/models/test_tools.py similarity index 80% rename from api/tests/unit_tests/models/test_tool_models.py rename to api/tests/unit_tests/models/test_tools.py index 8e3c4da904b..d767242bec8 100644 --- a/api/tests/unit_tests/models/test_tool_models.py +++ b/api/tests/unit_tests/models/test_tools.py @@ -1,3 +1,69 @@ +import json +from uuid import uuid4 + +from sqlalchemy.orm import Session + +from core.tools.entities.tool_entities import ApiProviderSchemaType, ToolProviderType +from models.account import Account, Tenant +from models.model import App, AppMode +from models.tools import ( + ApiToolProvider, + BuiltinToolProvider, + MCPToolProvider, + ToolLabelBinding, + ToolOAuthSystemClient, + ToolOAuthTenantClient, + WorkflowToolProvider, +) + +"""Regression coverage for the ``db.session``→session-parameter refactor on +``MCPToolProvider.load_user``. + +The method reached for the global ``db.session`` internally and has been converted to +take an explicit ``session: Session`` (per the pattern established in #40370/#42316, +tracked in #40372). +""" + + +def _mcp_tool_provider_session_accessors_persist_account(session: Session, *, name: str = "MCP Owner") -> Account: + account = Account(name=name, email=f"{uuid4().hex}@example.com") + session.add(account) + session.flush() + return account + + +def _persist_provider(session: Session, *, user_id: str) -> MCPToolProvider: + provider = MCPToolProvider( + name="Test Provider", + server_identifier=f"test-provider-{uuid4().hex[:8]}", + server_url="https://example.com/mcp", + server_url_hash=uuid4().hex, + icon="icon", + tenant_id=str(uuid4()), + user_id=user_id, + ) + session.add(provider) + session.flush() + return provider + + +def test_load_user_returns_account_through_the_given_session(sqlite_session: Session) -> None: + account = _mcp_tool_provider_session_accessors_persist_account(sqlite_session) + provider = _persist_provider(sqlite_session, user_id=account.id) + + result = provider.load_user(sqlite_session) + + assert result is not None + assert result.id == account.id + assert result.name == account.name + + +def test_load_user_returns_none_when_account_is_missing(sqlite_session: Session) -> None: + provider = _persist_provider(sqlite_session, user_id=str(uuid4())) + + assert provider.load_user(sqlite_session) is None + + """ Comprehensive unit tests for Tool models. @@ -9,18 +75,6 @@ This test suite covers: - ToolLabelBinding relationships """ -import json -from uuid import uuid4 - -from core.tools.entities.tool_entities import ApiProviderSchemaType, ToolProviderType -from models.tools import ( - ApiToolProvider, - BuiltinToolProvider, - ToolLabelBinding, - ToolOAuthSystemClient, - ToolOAuthTenantClient, -) - class TestBuiltinToolProviderValidation: """Test suite for BuiltinToolProvider model validation and operations.""" @@ -964,3 +1018,168 @@ class TestToolProviderRelationships: assert binding1.tool_id == tool_id assert binding2.tool_id == tool_id assert binding1.label_name != binding2.label_name + + +"""Regression coverage for the ``@property``→session-parameter refactor on tool provider accessors. + +Covers four of `api/models/tools.py`'s legacy `@property` accessors that reached for the global +``db.session`` internally and have been converted to plain methods taking an explicit +``session: Session`` (per the pattern established in #40370/#40797/#41394, tracked in #40372): + +- ``ApiToolProvider.tenant`` +- ``ApiToolProvider.user`` +- ``WorkflowToolProvider.user`` +- ``WorkflowToolProvider.tenant`` +- ``WorkflowToolProvider.app`` + +Each accessor is exercised against the real ``sqlite_session`` fixture (a genuine SQLAlchemy +``Session`` bound to a pristine full-schema SQLite database) so the assertions cover actual +query behaviour rather than a mock's recorded call. +""" + + +def _tool_provider_session_accessors_persist_account(session: Session) -> Account: + account = Account(name="Test Account", email="test@example.com") + session.add(account) + session.flush() + return account + + +def _persist_tenant(session: Session) -> Tenant: + tenant = Tenant(name="Test Tenant") + session.add(tenant) + session.flush() + return tenant + + +def _persist_app(session: Session, *, tenant_id: str) -> App: + app = App( + tenant_id=tenant_id, + name="Test App", + mode=AppMode.WORKFLOW, + enable_site=True, + enable_api=True, + created_by=str(uuid4()), + ) + session.add(app) + session.flush() + return app + + +def _api_tool_provider(*, tenant_id: str) -> ApiToolProvider: + return ApiToolProvider( + tenant_id=tenant_id, + user_id=str(uuid4()), + name="Test API Provider", + icon="{}", + schema="{}", + schema_type_str=ApiProviderSchemaType.OPENAPI, + description="", + tools_str="[]", + credentials_str="{}", + ) + + +def _workflow_tool_provider(*, tenant_id: str, app_id: str, user_id: str) -> WorkflowToolProvider: + return WorkflowToolProvider( + name="Test Workflow Provider", + label="Test Workflow Provider", + icon="{}", + app_id=app_id, + version="1", + user_id=user_id, + tenant_id=tenant_id, + description="", + ) + + +class TestApiToolProviderTenant: + def test_returns_persisted_tenant(self, sqlite_session: Session) -> None: + tenant = _persist_tenant(sqlite_session) + provider = _api_tool_provider(tenant_id=tenant.id) + + result = provider.tenant(session=sqlite_session) + + assert result is not None + assert result.id == tenant.id + + def test_returns_none_when_tenant_missing(self, sqlite_session: Session) -> None: + provider = _api_tool_provider(tenant_id=str(uuid4())) + + assert provider.tenant(session=sqlite_session) is None + + +class TestApiToolProviderUser: + def test_returns_persisted_user(self, sqlite_session: Session) -> None: + account = _tool_provider_session_accessors_persist_account(sqlite_session) + provider = _api_tool_provider(tenant_id=str(uuid4())) + provider.user_id = account.id + sqlite_session.add(provider) + sqlite_session.flush() + + result = provider.user(session=sqlite_session) + + assert result is not None + assert result.id == account.id + + def test_returns_none_when_user_missing(self, sqlite_session: Session) -> None: + provider = _api_tool_provider(tenant_id=str(uuid4())) + + assert provider.user(session=sqlite_session) is None + + +class TestWorkflowToolProviderUser: + def test_returns_persisted_user(self, sqlite_session: Session) -> None: + tenant = _persist_tenant(sqlite_session) + app = _persist_app(sqlite_session, tenant_id=tenant.id) + account = _tool_provider_session_accessors_persist_account(sqlite_session) + provider = _workflow_tool_provider(tenant_id=tenant.id, app_id=app.id, user_id=account.id) + + result = provider.user(session=sqlite_session) + + assert result is not None + assert result.id == account.id + + def test_returns_none_when_user_missing(self, sqlite_session: Session) -> None: + tenant = _persist_tenant(sqlite_session) + app = _persist_app(sqlite_session, tenant_id=tenant.id) + provider = _workflow_tool_provider(tenant_id=tenant.id, app_id=app.id, user_id=str(uuid4())) + + assert provider.user(session=sqlite_session) is None + + +class TestWorkflowToolProviderTenant: + def test_returns_persisted_tenant(self, sqlite_session: Session) -> None: + tenant = _persist_tenant(sqlite_session) + app = _persist_app(sqlite_session, tenant_id=tenant.id) + provider = _workflow_tool_provider(tenant_id=tenant.id, app_id=app.id, user_id=str(uuid4())) + + result = provider.tenant(session=sqlite_session) + + assert result is not None + assert result.id == tenant.id + + def test_returns_none_when_tenant_missing(self, sqlite_session: Session) -> None: + tenant = _persist_tenant(sqlite_session) + app = _persist_app(sqlite_session, tenant_id=tenant.id) + provider = _workflow_tool_provider(tenant_id=str(uuid4()), app_id=app.id, user_id=str(uuid4())) + + assert provider.tenant(session=sqlite_session) is None + + +class TestWorkflowToolProviderApp: + def test_returns_persisted_app(self, sqlite_session: Session) -> None: + tenant = _persist_tenant(sqlite_session) + app = _persist_app(sqlite_session, tenant_id=tenant.id) + provider = _workflow_tool_provider(tenant_id=tenant.id, app_id=app.id, user_id=str(uuid4())) + + result = provider.app(session=sqlite_session) + + assert result is not None + assert result.id == app.id + + def test_returns_none_when_app_missing(self, sqlite_session: Session) -> None: + tenant = _persist_tenant(sqlite_session) + provider = _workflow_tool_provider(tenant_id=tenant.id, app_id=str(uuid4()), user_id=str(uuid4())) + + assert provider.app(session=sqlite_session) is None diff --git a/api/tests/unit_tests/models/test_trigger_models.py b/api/tests/unit_tests/models/test_trigger.py similarity index 100% rename from api/tests/unit_tests/models/test_trigger_models.py rename to api/tests/unit_tests/models/test_trigger.py diff --git a/api/tests/unit_tests/models/test_typebase_app_conversation_models.py b/api/tests/unit_tests/models/test_typebase_app_conversation_models.py deleted file mode 100644 index db9d73f0e5c..00000000000 --- a/api/tests/unit_tests/models/test_typebase_app_conversation_models.py +++ /dev/null @@ -1,92 +0,0 @@ -from decimal import Decimal - -import pytest -from sqlalchemy import inspect -from sqlalchemy.orm import Session - -from models.base import Base, TypeBase -from models.enums import ApiTokenType, ConversationFromSource, EndUserType -from models.model import ApiToken, App, AppMode, AppStar, Conversation, EndUser, Message - -TARGET_MODELS = (App, AppStar, Conversation, Message, EndUser, ApiToken) - - -def test_app_conversation_models_use_typebase_registry() -> None: - for model in TARGET_MODELS: - mapper = inspect(model) - assert mapper.registry is TypeBase.registry - assert mapper.registry is not Base.registry - - -def test_generated_identifiers_and_collection_defaults_are_independent() -> None: - first = Conversation( - app_id="app-1", - mode=AppMode.CHAT, - name="First", - from_source=ConversationFromSource.API, - ) - second = Conversation( - app_id="app-1", - mode=AppMode.CHAT, - name="Second", - from_source=ConversationFromSource.API, - ) - - assert first.id != second.id - assert first._inputs == second._inputs == {} - assert first._inputs is not second._inputs - - -@pytest.mark.parametrize("sqlite_session", [TARGET_MODELS], indirect=True) -def test_app_conversation_models_persist_with_relationships(sqlite_session: Session) -> None: - app = App( - id="app-1", - tenant_id="tenant-1", - name="Typed app", - mode=AppMode.CHAT, - enable_site=True, - enable_api=True, - ) - star = AppStar(id="star-1", tenant_id=app.tenant_id, app_id=app.id, account_id="account-1") - conversation = Conversation( - id="conversation-1", - app_id=app.id, - mode=AppMode.CHAT, - name="Typed conversation", - from_source=ConversationFromSource.API, - ) - message = Message( - id="message-1", - app_id=app.id, - conversation_id=conversation.id, - query="hello", - message={"role": "user"}, - answer="world", - from_source=ConversationFromSource.API, - total_price=Decimal(0), - ) - end_user = EndUser( - id="end-user-1", - tenant_id=app.tenant_id, - app_id=app.id, - type=EndUserType.BROWSER, - session_id="session-1", - ) - token = ApiToken( - id="token-1", - app_id=app.id, - tenant_id=app.tenant_id, - type=ApiTokenType.APP, - token="secret", - ) - - sqlite_session.add_all([app, star, conversation, message, end_user, token]) - sqlite_session.commit() - - persisted_conversation = sqlite_session.get(Conversation, conversation.id) - assert persisted_conversation is not None - assert [item.id for item in persisted_conversation.messages] == [message.id] - assert sqlite_session.get(AppStar, star.id) is not None - assert sqlite_session.get(EndUser, end_user.id) is not None - assert sqlite_session.get(ApiToken, token.id) is not None - assert message.created_at is not None diff --git a/api/tests/unit_tests/models/test_web_models.py b/api/tests/unit_tests/models/test_web.py similarity index 100% rename from api/tests/unit_tests/models/test_web_models.py rename to api/tests/unit_tests/models/test_web.py diff --git a/api/tests/unit_tests/models/test_workflow.py b/api/tests/unit_tests/models/test_workflow.py index 111be8880a9..1bef2f478ef 100644 --- a/api/tests/unit_tests/models/test_workflow.py +++ b/api/tests/unit_tests/models/test_workflow.py @@ -1,459 +1,2088 @@ import dataclasses import json +from datetime import UTC, datetime +from types import SimpleNamespace from unittest import mock -from uuid import uuid4 +from unittest.mock import Mock +from uuid import NAMESPACE_URL, UUID, uuid4, uuid5 import pytest +from pydantic import JsonValue +from sqlalchemy import select from sqlalchemy.orm import Session from constants import HIDDEN_VALUE from core.helper import encrypter from core.workflow.file_reference import build_file_reference from core.workflow.llm_environment_variable import LLMEnvironmentVariable +from extensions.storage.storage_type import StorageType +from factories import variable_factory from factories.variable_factory import build_segment +from fields.workflow_run_fields import ( + node_execution_response_source, + workflow_run_pagination_response_source, + workflow_run_response_source, +) +from graphon.enums import ( + BuiltinNodeTypes, + WorkflowExecutionStatus, + WorkflowNodeExecutionStatus, +) from graphon.file import File, FileTransferMethod, FileType -from graphon.variables import FloatVariable, IntegerVariable, SecretVariable, StringVariable +from graphon.variables import FloatVariable, IntegerVariable, SecretVariable, SegmentType, StringVariable from graphon.variables.segments import IntegerSegment, Segment +from models import ConversationVariable +from models import Workflow as _workflow_feature_compatibility_Workflow from models.account import Account +from models.enums import CreatorUserRole, EndUserType, ExecutionOffLoadType, WorkflowRunTriggeredFrom +from models.model import EndUser, UploadFile from models.tools import WorkflowToolProvider from models.workflow import ( - Workflow, + Workflow as _workflow_Workflow, +) +from models.workflow import ( + Workflow as _workflow_models_Workflow, +) +from models.workflow import ( + WorkflowAppLog, + WorkflowAppLogCreatedFrom, WorkflowDraftVariable, WorkflowNodeExecutionModel, + WorkflowNodeExecutionOffload, + WorkflowNodeExecutionTriggeredFrom, + WorkflowRun, + WorkflowType, is_system_variable_editable, ) -def test_environment_variables(): - # tenant_id context variable removed - using current_user.current_tenant_id directly - - # Create a Workflow instance - workflow = Workflow( - tenant_id="tenant_id", - app_id="app_id", - type="workflow", - version="draft", - graph="{}", - features="{}", - created_by="account_id", - environment_variables=[], - conversation_variables=[], - ) - - # Create some EnvironmentVariable instances - variable1 = StringVariable(name="var1", value="value1", id=str(uuid4()), selector=["env", "var1"]) - variable2 = IntegerVariable(name="var2", value=123, id=str(uuid4()), selector=["env", "var2"]) - variable3 = SecretVariable(name="var3", value="secret", id=str(uuid4()), selector=["env", "var3"]) - variable4 = FloatVariable(name="var4", value=float("3.14"), id=str(uuid4()), selector=["env", "var4"]) - - with ( - mock.patch("core.helper.encrypter.encrypt_token", return_value="encrypted_token"), - mock.patch("core.helper.encrypter.decrypt_token", return_value="secret"), - ): - # Set the environment_variables property of the Workflow instance - variables = [variable1, variable2, variable3, variable4] - workflow.environment_variables = variables - - # Get the environment_variables property and assert its value - assert workflow.environment_variables == variables - - -def test_llm_environment_variable_round_trip(): - workflow = Workflow( - tenant_id="tenant_id", - app_id="app_id", - type="workflow", - version="draft", - graph="{}", - features="{}", - created_by="account_id", - environment_variables=[], - conversation_variables=[], - ) - variable = LLMEnvironmentVariable( - name="for_research", - value={"provider": "langgenius/anthropic/anthropic", "name": "claude-sonnet", "mode": "chat"}, - id=str(uuid4()), - selector=["env", "for_research"], - ) - - workflow.environment_variables = [variable] - - assert workflow.environment_variables == [variable] - assert json.loads(workflow._environment_variables)["for_research"]["value_type"] == "llm" - assert workflow.to_dict()["environment_variables"][0]["value_type"] == "llm" - - -def test_update_environment_variables(): - # tenant_id context variable removed - using current_user.current_tenant_id directly - - # Create a Workflow instance - workflow = Workflow( - tenant_id="tenant_id", - app_id="app_id", - type="workflow", - version="draft", - graph="{}", - features="{}", - created_by="account_id", - environment_variables=[], - conversation_variables=[], - ) - - # Create some EnvironmentVariable instances - variable1 = StringVariable(name="var1", value="value1", id=str(uuid4()), selector=["env", "var1"]) - variable2 = IntegerVariable(name="var2", value=123, id=str(uuid4()), selector=["env", "var2"]) - variable3 = SecretVariable( - name="var3", - value="secret", - id=str(uuid4()), - selector=["env", "var3"], - description="old description", - ) - variable4 = FloatVariable(name="var4", value=float("3.14"), id=str(uuid4()), selector=["env", "var4"]) - - with ( - mock.patch("core.helper.encrypter.encrypt_token", return_value="encrypted_token"), - mock.patch("core.helper.encrypter.decrypt_token", return_value="secret"), - ): - variables = [variable1, variable2, variable3, variable4] - - # Set the environment_variables property of the Workflow instance - workflow.environment_variables = variables - assert workflow.environment_variables == [variable1, variable2, variable3, variable4] - - # Update the name and description of variable3 and keep the value as it is - variables[2] = variable3.model_copy( - update={ - "name": "new name", - "description": "new description", - "value": HIDDEN_VALUE, +class TestConversationVariableCases: + @staticmethod + def test_from_variable_coerces_non_uuid_primary_key() -> None: + variable = variable_factory.build_conversation_variable_from_mapping( + { + "id": "opt-comp-prompt-var", + "name": "optimization_comparison_prompt", + "value_type": SegmentType.STRING, + "value": "-", } ) - workflow.environment_variables = variables - assert workflow.environment_variables[2].name == "new name" - assert workflow.environment_variables[2].description == "new description" - assert workflow.environment_variables[2].value == variable3.value + row = ConversationVariable.from_variable(app_id="app_id", conversation_id="conversation_id", variable=variable) + expected = str(uuid5(NAMESPACE_URL, "dify:conversation-variable:optimization_comparison_prompt")) + UUID(row.id) + assert row.id == expected + assert variable.id == "opt-comp-prompt-var" + assert row.to_variable().id == "opt-comp-prompt-var" -def test_to_dict(): - # tenant_id context variable removed - using current_user.current_tenant_id directly - - # Create a Workflow instance - workflow = Workflow( - tenant_id="tenant_id", - app_id="app_id", - type="workflow", - version="draft", - graph="{}", - features="{}", - created_by="account_id", - environment_variables=[], - conversation_variables=[], - ) - - # Create some EnvironmentVariable instances - - with ( - mock.patch("core.helper.encrypter.encrypt_token", return_value="encrypted_token"), - mock.patch("core.helper.encrypter.decrypt_token", return_value="secret"), - ): - # Set the environment_variables property of the Workflow instance - workflow.environment_variables = [ - SecretVariable(name="secret", value="secret", id=str(uuid4())), - StringVariable(name="text", value="text", id=str(uuid4())), - ] - - workflow_dict = workflow.to_dict() - assert workflow_dict["environment_variables"][0]["value"] == "" - assert workflow_dict["environment_variables"][1]["value"] == "text" - - workflow_dict = workflow.to_dict(include_secret=True) - assert workflow_dict["environment_variables"][0]["value"] == "secret" - assert workflow_dict["environment_variables"][1]["value"] == "text" - - -@pytest.mark.parametrize("sqlite_session", [(Workflow, Account)], indirect=True) -def test_workflow_account_accessors_use_caller_session(sqlite_session: Session): - created_account = Account(name="Created Account", email="created@example.com") - created_account.id = "created-account-id" - updated_account = Account(name="Updated Account", email="updated@example.com") - updated_account.id = "updated-account-id" - decoy_account = Account(name="Decoy Account", email="decoy@example.com") - decoy_account.id = "decoy-account-id" - workflow = Workflow( - tenant_id="tenant_id", - app_id="app_id", - type="workflow", - version="draft", - graph="{}", - features="{}", - created_by="created-account-id", - environment_variables=[], - conversation_variables=[], - updated_by="updated-account-id", - ) - sqlite_session.add_all([decoy_account, updated_account, workflow, created_account]) - sqlite_session.flush() - - assert workflow.created_by_account(sqlite_session) is created_account - assert workflow.updated_by_account(sqlite_session) is updated_account - - -@pytest.mark.parametrize("sqlite_session", [(Workflow, WorkflowToolProvider)], indirect=True) -def test_workflow_tool_published_accessor_uses_caller_session(sqlite_session: Session): - workflow = Workflow( - tenant_id="tenant_id", - app_id="app_id", - type="workflow", - version="draft", - graph="{}", - features="{}", - created_by="account_id", - environment_variables=[], - conversation_variables=[], - ) - matching_provider = WorkflowToolProvider( - name="matching-provider", - label="Matching provider", - icon="tool", - app_id=workflow.app_id, - version="1", - user_id="account-id", - tenant_id=workflow.tenant_id, - description="Matching workflow tool", - ) - decoy_provider = WorkflowToolProvider( - name="decoy-provider", - label="Decoy provider", - icon="tool", - app_id="other-app", - version="1", - user_id="account-id", - tenant_id=workflow.tenant_id, - description="Different app", - ) - sqlite_session.add_all([decoy_provider, workflow, matching_provider]) - sqlite_session.flush() - - with pytest.warns(DeprecationWarning, match="not accurate"): - assert workflow.tool_published(sqlite_session) is True - - -def test_normalize_environment_variable_mappings_converts_full_mask_to_hidden_value(): - normalized = Workflow.normalize_environment_variable_mappings( - [ + @staticmethod + def test_from_variable_and_to_variable() -> None: + variable = variable_factory.build_conversation_variable_from_mapping( { "id": str(uuid4()), - "name": "secret", - "value": encrypter.full_mask_token(), - "value_type": "secret", + "name": "name", + "value_type": SegmentType.OBJECT, + "value": { + "key": { + "key": "value", + } + }, } - ] - ) - - assert normalized[0]["value"] == HIDDEN_VALUE - - -def test_normalize_environment_variable_mappings_keeps_hidden_value(): - normalized = Workflow.normalize_environment_variable_mappings( - [ - { - "id": str(uuid4()), - "name": "secret", - "value": HIDDEN_VALUE, - "value_type": "secret", - } - ] - ) - - assert normalized[0]["value"] == HIDDEN_VALUE - - -class TestWorkflowNodeExecution: - def test_execution_metadata_dict(self): - node_exec = WorkflowNodeExecutionModel( - execution_metadata=None, - ) - assert node_exec.execution_metadata_dict == {} - - original = {"a": 1, "b": ["2"]} - node_exec.execution_metadata = json.dumps(original) - assert node_exec.execution_metadata_dict == original - - -class TestIsSystemVariableEditable: - def test_is_system_variable(self): - cases = [ - ("query", True), - ("files", True), - ("dialogue_count", False), - ("conversation_id", False), - ("user_id", False), - ("app_id", False), - ("workflow_id", False), - ("workflow_run_id", False), - ] - for name, editable in cases: - assert editable == is_system_variable_editable(name) - - assert is_system_variable_editable("invalid_or_new_system_variable") is False - - -class TestWorkflowDraftVariableGetValue: - def test_get_value_by_case(self): - @dataclasses.dataclass - class TestCase: - name: str - value: Segment - - test_file = File( - file_type=FileType.IMAGE, - transfer_method=FileTransferMethod.REMOTE_URL, - remote_url="https://example.com/example.jpg", - filename="example.jpg", - extension=".jpg", - mime_type="image/jpeg", - size=100, - ) - cases: list[TestCase] = [ - TestCase( - name="number/int", - value=build_segment(1), - ), - TestCase( - name="number/float", - value=build_segment(1.0), - ), - TestCase( - name="string", - value=build_segment("a"), - ), - TestCase( - name="object", - value=build_segment({}), - ), - TestCase( - name="file", - value=build_segment(test_file), - ), - TestCase( - name="array[any]", - value=build_segment([1, "a"]), - ), - TestCase( - name="array[string]", - value=build_segment(["a", "b"]), - ), - TestCase( - name="array[number]/int", - value=build_segment([1, 2]), - ), - TestCase( - name="array[number]/float", - value=build_segment([1.0, 2.0]), - ), - TestCase( - name="array[number]/mixed", - value=build_segment([1, 2.0]), - ), - TestCase( - name="array[object]", - value=build_segment([{}, {"a": 1}]), - ), - TestCase( - name="none", - value=build_segment(None), - ), - ] - - for idx, c in enumerate(cases, 1): - fail_msg = f"test case {c.name} failed, index={idx}" - draft_var = WorkflowDraftVariable() - draft_var.set_value(c.value) - assert c.value == draft_var.get_value(), fail_msg - - def test_file_variable_preserves_all_fields(self): - """Test that File type variables preserve all fields during encoding/decoding.""" - - # Create a File with specific field values - test_file = File( - file_id="test_file_id", - file_type=FileType.IMAGE, - transfer_method=FileTransferMethod.REMOTE_URL, - remote_url="https://example.com/test.jpg", - filename="test.jpg", - extension=".jpg", - mime_type="image/jpeg", - size=12345, # Specific size to test preservation - storage_key="test_storage_key", ) - # Create a FileSegment and WorkflowDraftVariable - file_segment = build_segment(test_file) - draft_var = WorkflowDraftVariable() - draft_var.set_value(file_segment) - - # Retrieve the value and verify all fields are preserved - retrieved_segment = draft_var.get_value() - retrieved_file = retrieved_segment.value - - # Verify all important fields are preserved - assert retrieved_file.id == test_file.id - assert retrieved_file.type == test_file.type - assert retrieved_file.transfer_method == test_file.transfer_method - assert retrieved_file.remote_url == test_file.remote_url - assert retrieved_file.filename == test_file.filename - assert retrieved_file.extension == test_file.extension - assert retrieved_file.mime_type == test_file.mime_type - assert retrieved_file.size == test_file.size # This was the main issue being fixed - # Note: storage_key is not serialized in model_dump() so it won't be preserved - - # Verify the segments have the same type and the important fields match - assert file_segment.value_type == retrieved_segment.value_type - - def test_file_variable_rebuilds_storage_backed_payloads_with_app_tenant(self): - persisted_file = File( - file_id="test_file_id", - file_type=FileType.DOCUMENT, - transfer_method=FileTransferMethod.LOCAL_FILE, - reference=build_file_reference(record_id="upload-1", storage_key="legacy-storage-key"), - filename="test.txt", - extension=".txt", - mime_type="text/plain", - size=12, + conversation_variable = ConversationVariable.from_variable( + app_id="app_id", conversation_id="conversation_id", variable=variable ) - rebuilt_file = File( - file_id="test_file_id", - file_type=FileType.DOCUMENT, - transfer_method=FileTransferMethod.LOCAL_FILE, - reference=build_file_reference(record_id="upload-1"), - filename="test.txt", - extension=".txt", - mime_type="text/plain", - size=12, - storage_key="canonical-storage-key", + + assert conversation_variable.to_variable() == variable + + +class TestWorkflowCases: + @staticmethod + def test_environment_variables(): + # tenant_id context variable removed - using current_user.current_tenant_id directly + + # Create a Workflow instance + workflow = _workflow_Workflow( + tenant_id="tenant_id", + app_id="app_id", + type="workflow", + version="draft", + graph="{}", + features="{}", + created_by="account_id", + environment_variables=[], + conversation_variables=[], ) - draft_var = WorkflowDraftVariable( - app_id="app-1", - ) - draft_var.set_value(build_segment(persisted_file)) - draft_var._WorkflowDraftVariable__value = None + + # Create some EnvironmentVariable instances + variable1 = StringVariable(name="var1", value="value1", id=str(uuid4()), selector=["env", "var1"]) + variable2 = IntegerVariable(name="var2", value=123, id=str(uuid4()), selector=["env", "var2"]) + variable3 = SecretVariable(name="var3", value="secret", id=str(uuid4()), selector=["env", "var3"]) + variable4 = FloatVariable(name="var4", value=float("3.14"), id=str(uuid4()), selector=["env", "var4"]) with ( - mock.patch("models.workflow._resolve_workflow_app_tenant_id", return_value="tenant-1"), - mock.patch("models.workflow.build_file_from_stored_mapping", return_value=rebuilt_file) as rebuild_file, + mock.patch("core.helper.encrypter.encrypt_token", return_value="encrypted_token"), + mock.patch("core.helper.encrypter.decrypt_token", return_value="secret"), ): + # Set the environment_variables property of the Workflow instance + variables = [variable1, variable2, variable3, variable4] + workflow.environment_variables = variables + + # Get the environment_variables property and assert its value + assert workflow.environment_variables == variables + + @staticmethod + def test_llm_environment_variable_round_trip(): + workflow = _workflow_Workflow( + tenant_id="tenant_id", + app_id="app_id", + type="workflow", + version="draft", + graph="{}", + features="{}", + created_by="account_id", + environment_variables=[], + conversation_variables=[], + ) + variable = LLMEnvironmentVariable( + name="for_research", + value={"provider": "langgenius/anthropic/anthropic", "name": "claude-sonnet", "mode": "chat"}, + id=str(uuid4()), + selector=["env", "for_research"], + ) + + workflow.environment_variables = [variable] + + assert workflow.environment_variables == [variable] + assert json.loads(workflow._environment_variables)["for_research"]["value_type"] == "llm" + assert workflow.to_dict()["environment_variables"][0]["value_type"] == "llm" + + @staticmethod + def test_update_environment_variables(): + # tenant_id context variable removed - using current_user.current_tenant_id directly + + # Create a Workflow instance + workflow = _workflow_Workflow( + tenant_id="tenant_id", + app_id="app_id", + type="workflow", + version="draft", + graph="{}", + features="{}", + created_by="account_id", + environment_variables=[], + conversation_variables=[], + ) + + # Create some EnvironmentVariable instances + variable1 = StringVariable(name="var1", value="value1", id=str(uuid4()), selector=["env", "var1"]) + variable2 = IntegerVariable(name="var2", value=123, id=str(uuid4()), selector=["env", "var2"]) + variable3 = SecretVariable( + name="var3", + value="secret", + id=str(uuid4()), + selector=["env", "var3"], + description="old description", + ) + variable4 = FloatVariable(name="var4", value=float("3.14"), id=str(uuid4()), selector=["env", "var4"]) + + with ( + mock.patch("core.helper.encrypter.encrypt_token", return_value="encrypted_token"), + mock.patch("core.helper.encrypter.decrypt_token", return_value="secret"), + ): + variables = [variable1, variable2, variable3, variable4] + + # Set the environment_variables property of the Workflow instance + workflow.environment_variables = variables + assert workflow.environment_variables == [variable1, variable2, variable3, variable4] + + # Update the name and description of variable3 and keep the value as it is + variables[2] = variable3.model_copy( + update={ + "name": "new name", + "description": "new description", + "value": HIDDEN_VALUE, + } + ) + + workflow.environment_variables = variables + assert workflow.environment_variables[2].name == "new name" + assert workflow.environment_variables[2].description == "new description" + assert workflow.environment_variables[2].value == variable3.value + + @staticmethod + def test_to_dict(): + # tenant_id context variable removed - using current_user.current_tenant_id directly + + # Create a Workflow instance + workflow = _workflow_Workflow( + tenant_id="tenant_id", + app_id="app_id", + type="workflow", + version="draft", + graph="{}", + features="{}", + created_by="account_id", + environment_variables=[], + conversation_variables=[], + ) + + # Create some EnvironmentVariable instances + + with ( + mock.patch("core.helper.encrypter.encrypt_token", return_value="encrypted_token"), + mock.patch("core.helper.encrypter.decrypt_token", return_value="secret"), + ): + # Set the environment_variables property of the Workflow instance + workflow.environment_variables = [ + SecretVariable(name="secret", value="secret", id=str(uuid4())), + StringVariable(name="text", value="text", id=str(uuid4())), + ] + + workflow_dict = workflow.to_dict() + assert workflow_dict["environment_variables"][0]["value"] == "" + assert workflow_dict["environment_variables"][1]["value"] == "text" + + workflow_dict = workflow.to_dict(include_secret=True) + assert workflow_dict["environment_variables"][0]["value"] == "secret" + assert workflow_dict["environment_variables"][1]["value"] == "text" + + @staticmethod + @pytest.mark.parametrize("sqlite_session", [(_workflow_Workflow, Account)], indirect=True) + def test_workflow_account_accessors_use_caller_session(sqlite_session: Session): + created_account = Account(name="Created Account", email="created@example.com") + created_account.id = "created-account-id" + updated_account = Account(name="Updated Account", email="updated@example.com") + updated_account.id = "updated-account-id" + decoy_account = Account(name="Decoy Account", email="decoy@example.com") + decoy_account.id = "decoy-account-id" + workflow = _workflow_Workflow( + tenant_id="tenant_id", + app_id="app_id", + type="workflow", + version="draft", + graph="{}", + features="{}", + created_by="created-account-id", + environment_variables=[], + conversation_variables=[], + updated_by="updated-account-id", + ) + sqlite_session.add_all([decoy_account, updated_account, workflow, created_account]) + sqlite_session.flush() + + assert workflow.created_by_account(sqlite_session) is created_account + assert workflow.updated_by_account(sqlite_session) is updated_account + + @staticmethod + @pytest.mark.parametrize("sqlite_session", [(_workflow_Workflow, WorkflowToolProvider)], indirect=True) + def test_workflow_tool_published_accessor_uses_caller_session(sqlite_session: Session): + workflow = _workflow_Workflow( + tenant_id="tenant_id", + app_id="app_id", + type="workflow", + version="draft", + graph="{}", + features="{}", + created_by="account_id", + environment_variables=[], + conversation_variables=[], + ) + matching_provider = WorkflowToolProvider( + name="matching-provider", + label="Matching provider", + icon="tool", + app_id=workflow.app_id, + version="1", + user_id="account-id", + tenant_id=workflow.tenant_id, + description="Matching workflow tool", + ) + decoy_provider = WorkflowToolProvider( + name="decoy-provider", + label="Decoy provider", + icon="tool", + app_id="other-app", + version="1", + user_id="account-id", + tenant_id=workflow.tenant_id, + description="Different app", + ) + sqlite_session.add_all([decoy_provider, workflow, matching_provider]) + sqlite_session.flush() + + with pytest.warns(DeprecationWarning, match="not accurate"): + assert workflow.tool_published(sqlite_session) is True + + @staticmethod + def test_normalize_environment_variable_mappings_converts_full_mask_to_hidden_value(): + normalized = _workflow_Workflow.normalize_environment_variable_mappings( + [ + { + "id": str(uuid4()), + "name": "secret", + "value": encrypter.full_mask_token(), + "value_type": "secret", + } + ] + ) + + assert normalized[0]["value"] == HIDDEN_VALUE + + @staticmethod + def test_normalize_environment_variable_mappings_keeps_hidden_value(): + normalized = _workflow_Workflow.normalize_environment_variable_mappings( + [ + { + "id": str(uuid4()), + "name": "secret", + "value": HIDDEN_VALUE, + "value_type": "secret", + } + ] + ) + + assert normalized[0]["value"] == HIDDEN_VALUE + + class TestWorkflowNodeExecution: + def test_execution_metadata_dict(self): + node_exec = WorkflowNodeExecutionModel( + execution_metadata=None, + ) + assert node_exec.execution_metadata_dict == {} + + original = {"a": 1, "b": ["2"]} + node_exec.execution_metadata = json.dumps(original) + assert node_exec.execution_metadata_dict == original + + class TestIsSystemVariableEditable: + def test_is_system_variable(self): + cases = [ + ("query", True), + ("files", True), + ("dialogue_count", False), + ("conversation_id", False), + ("user_id", False), + ("app_id", False), + ("workflow_id", False), + ("workflow_run_id", False), + ] + for name, editable in cases: + assert editable == is_system_variable_editable(name) + + assert is_system_variable_editable("invalid_or_new_system_variable") is False + + class TestWorkflowDraftVariableGetValue: + def test_get_value_by_case(self): + @dataclasses.dataclass + class TestCase: + name: str + value: Segment + + test_file = File( + file_type=FileType.IMAGE, + transfer_method=FileTransferMethod.REMOTE_URL, + remote_url="https://example.com/example.jpg", + filename="example.jpg", + extension=".jpg", + mime_type="image/jpeg", + size=100, + ) + cases: list[TestCase] = [ + TestCase( + name="number/int", + value=build_segment(1), + ), + TestCase( + name="number/float", + value=build_segment(1.0), + ), + TestCase( + name="string", + value=build_segment("a"), + ), + TestCase( + name="object", + value=build_segment({}), + ), + TestCase( + name="file", + value=build_segment(test_file), + ), + TestCase( + name="array[any]", + value=build_segment([1, "a"]), + ), + TestCase( + name="array[string]", + value=build_segment(["a", "b"]), + ), + TestCase( + name="array[number]/int", + value=build_segment([1, 2]), + ), + TestCase( + name="array[number]/float", + value=build_segment([1.0, 2.0]), + ), + TestCase( + name="array[number]/mixed", + value=build_segment([1, 2.0]), + ), + TestCase( + name="array[object]", + value=build_segment([{}, {"a": 1}]), + ), + TestCase( + name="none", + value=build_segment(None), + ), + ] + + for idx, c in enumerate(cases, 1): + fail_msg = f"test case {c.name} failed, index={idx}" + draft_var = WorkflowDraftVariable() + draft_var.set_value(c.value) + assert c.value == draft_var.get_value(), fail_msg + + def test_file_variable_preserves_all_fields(self): + """Test that File type variables preserve all fields during encoding/decoding.""" + + # Create a File with specific field values + test_file = File( + file_id="test_file_id", + file_type=FileType.IMAGE, + transfer_method=FileTransferMethod.REMOTE_URL, + remote_url="https://example.com/test.jpg", + filename="test.jpg", + extension=".jpg", + mime_type="image/jpeg", + size=12345, # Specific size to test preservation + storage_key="test_storage_key", + ) + + # Create a FileSegment and WorkflowDraftVariable + file_segment = build_segment(test_file) + draft_var = WorkflowDraftVariable() + draft_var.set_value(file_segment) + + # Retrieve the value and verify all fields are preserved retrieved_segment = draft_var.get_value() + retrieved_file = retrieved_segment.value - assert retrieved_segment.value == rebuilt_file - rebuild_file.assert_called_once() - assert rebuild_file.call_args.kwargs["tenant_id"] == "tenant-1" + # Verify all important fields are preserved + assert retrieved_file.id == test_file.id + assert retrieved_file.type == test_file.type + assert retrieved_file.transfer_method == test_file.transfer_method + assert retrieved_file.remote_url == test_file.remote_url + assert retrieved_file.filename == test_file.filename + assert retrieved_file.extension == test_file.extension + assert retrieved_file.mime_type == test_file.mime_type + assert retrieved_file.size == test_file.size # This was the main issue being fixed + # Note: storage_key is not serialized in model_dump() so it won't be preserved - def test_get_and_set_value(self): - draft_var = WorkflowDraftVariable() - int_var = IntegerSegment(value=1) - draft_var.set_value(int_var) - value = draft_var.get_value() - assert value == int_var + # Verify the segments have the same type and the important fields match + assert file_segment.value_type == retrieved_segment.value_type + + def test_file_variable_rebuilds_storage_backed_payloads_with_app_tenant(self): + persisted_file = File( + file_id="test_file_id", + file_type=FileType.DOCUMENT, + transfer_method=FileTransferMethod.LOCAL_FILE, + reference=build_file_reference(record_id="upload-1", storage_key="legacy-storage-key"), + filename="test.txt", + extension=".txt", + mime_type="text/plain", + size=12, + ) + rebuilt_file = File( + file_id="test_file_id", + file_type=FileType.DOCUMENT, + transfer_method=FileTransferMethod.LOCAL_FILE, + reference=build_file_reference(record_id="upload-1"), + filename="test.txt", + extension=".txt", + mime_type="text/plain", + size=12, + storage_key="canonical-storage-key", + ) + draft_var = WorkflowDraftVariable( + app_id="app-1", + ) + draft_var.set_value(build_segment(persisted_file)) + draft_var._WorkflowDraftVariable__value = None + + with ( + mock.patch("models.workflow._resolve_workflow_app_tenant_id", return_value="tenant-1"), + mock.patch("models.workflow.build_file_from_stored_mapping", return_value=rebuilt_file) as rebuild_file, + ): + retrieved_segment = draft_var.get_value() + + assert retrieved_segment.value == rebuilt_file + rebuild_file.assert_called_once() + assert rebuild_file.call_args.kwargs["tenant_id"] == "tenant-1" + + def test_get_and_set_value(self): + draft_var = WorkflowDraftVariable() + int_var = IntegerSegment(value=1) + draft_var.set_value(int_var) + value = draft_var.get_value() + assert value == int_var + + +"""Regression coverage for ``models.workflow.WorkflowAppLog`` account accessors. + +Ensures the ``@property``→session-parameter refactor preserves the role-based dispatch: +``created_by_account`` looks up an Account only when role is ACCOUNT; ``created_by_end_user`` +looks up an EndUser only when role is END_USER. + +Both accessors are exercised against the real ``sqlite_session`` fixture (a genuine +SQLAlchemy ``Session`` bound to a pristine full-schema SQLite database) so the assertions +cover actual query behaviour rather than a mock's recorded call. +""" + + +def _log(role: CreatorUserRole, created_by: str) -> WorkflowAppLog: + """Construct a WorkflowAppLog without touching the database.""" + return WorkflowAppLog( + tenant_id="00000000-0000-0000-0000-000000000001", + app_id="00000000-0000-0000-0000-000000000002", + workflow_id="00000000-0000-0000-0000-000000000003", + workflow_run_id="00000000-0000-0000-0000-000000000004", + created_from=WorkflowAppLogCreatedFrom.WEB_APP, + created_by_role=role, + created_by=created_by, + ) + + +class TestWorkflowAppLogModelsCases: + class TestCreatedByAccount: + def test_returns_account_lookup_when_role_is_account(self, sqlite_session: Session) -> None: + account = Account(name="Test Account", email="test@example.com") + sqlite_session.add(account) + sqlite_session.flush() + log = _log(CreatorUserRole.ACCOUNT, created_by=account.id) + + result = log.created_by_account(session=sqlite_session) + + assert result is not None + assert result.id == account.id + + def test_returns_none_when_role_is_end_user(self, sqlite_session: Session) -> None: + account = Account(name="Test Account", email="test@example.com") + sqlite_session.add(account) + sqlite_session.flush() + log = _log(CreatorUserRole.END_USER, created_by=account.id) + + assert log.created_by_account(session=sqlite_session) is None + + class TestCreatedByEndUser: + def test_returns_end_user_lookup_when_role_is_end_user(self, sqlite_session: Session) -> None: + end_user = EndUser( + tenant_id="00000000-0000-0000-0000-000000000001", + type=EndUserType.BROWSER, + session_id="session-1", + ) + sqlite_session.add(end_user) + sqlite_session.flush() + log = _log(CreatorUserRole.END_USER, created_by=end_user.id) + + result = log.created_by_end_user(session=sqlite_session) + + assert result is not None + assert result.id == end_user.id + + def test_returns_none_when_role_is_account(self, sqlite_session: Session) -> None: + end_user = EndUser( + tenant_id="00000000-0000-0000-0000-000000000001", + type=EndUserType.BROWSER, + session_id="session-1", + ) + sqlite_session.add(end_user) + sqlite_session.flush() + log = _log(CreatorUserRole.ACCOUNT, created_by=end_user.id) + + assert log.created_by_end_user(session=sqlite_session) is None + + +""" +This unit test verifies Workflow compatibility before and after supporting multiple file types. +""" + + +OLD_VERSION_WORKFLOW_FEATURES: dict[str, JsonValue] = { + "file_upload": { + "image": { + "enabled": True, + "number_limits": 6, + "transfer_methods": ["remote_url", "local_file"], + } + }, + "opening_statement": "", + "retriever_resource": {"enabled": True}, + "sensitive_word_avoidance": {"enabled": False}, + "speech_to_text": {"enabled": False}, + "suggested_questions": [], + "suggested_questions_after_answer": {"enabled": False}, + "text_to_speech": {"enabled": False, "language": "", "voice": ""}, +} + +NEW_VERSION_WORKFLOW_FEATURES: dict[str, JsonValue] = { + "file_upload": { + "enabled": True, + "allowed_file_types": ["image"], + "allowed_file_extensions": [], + "allowed_file_upload_methods": ["remote_url", "local_file"], + "number_limits": 6, + }, + "opening_statement": "", + "retriever_resource": {"enabled": True}, + "sensitive_word_avoidance": {"enabled": False}, + "speech_to_text": {"enabled": False}, + "suggested_questions": [], + "suggested_questions_after_answer": {"enabled": False}, + "text_to_speech": {"enabled": False, "language": "", "voice": ""}, +} + + +class TestWorkflowFeatureCompatibilityCases: + @staticmethod + def test_workflow_features() -> None: + workflow = _workflow_feature_compatibility_Workflow( + tenant_id="", + app_id="", + type="", + version="", + graph="", + features=json.dumps(OLD_VERSION_WORKFLOW_FEATURES), + created_by="", + environment_variables=[], + conversation_variables=[], + ) + + assert workflow.features_dict == NEW_VERSION_WORKFLOW_FEATURES + + +""" +Comprehensive unit tests for Workflow models. + +This test suite covers: +- Workflow model validation +- WorkflowRun state transitions +- NodeExecution relationships +- Graph configuration validation +""" + + +class TestWorkflowModelsCases: + class TestWorkflowModelValidation: + """Test suite for Workflow model validation and basic operations.""" + + def test_workflow_creation_with_required_fields(self) -> None: + """Test creating a workflow with all required fields.""" + # Arrange + tenant_id = str(uuid4()) + app_id = str(uuid4()) + created_by = str(uuid4()) + graph = json.dumps({"nodes": [], "edges": []}) + features = json.dumps({"file_upload": {"enabled": True}}) + + # Act + workflow = _workflow_models_Workflow.new( + tenant_id=tenant_id, + app_id=app_id, + type=WorkflowType.WORKFLOW, + version="draft", + graph=graph, + features=features, + created_by=created_by, + environment_variables=[], + conversation_variables=[], + rag_pipeline_variables=[], + ) + + # Assert + assert workflow.tenant_id == tenant_id + assert workflow.app_id == app_id + assert workflow.type == WorkflowType.WORKFLOW + assert workflow.version == "draft" + assert workflow.graph == graph + assert workflow.created_by == created_by + assert workflow.created_at is not None + assert workflow.updated_at is not None + + def test_workflow_type_enum_values(self) -> None: + """Test WorkflowType enum values.""" + # Assert + assert WorkflowType.WORKFLOW == "workflow" + assert WorkflowType.CHAT.value == "chat" + assert WorkflowType.RAG_PIPELINE.value == "rag-pipeline" + + def test_workflow_type_value_of(self) -> None: + """Test WorkflowType.value_of method.""" + # Act & Assert + assert WorkflowType.value_of("workflow") == WorkflowType.WORKFLOW + assert WorkflowType.value_of("chat") == WorkflowType.CHAT + assert WorkflowType.value_of("rag-pipeline") == WorkflowType.RAG_PIPELINE + + with pytest.raises(ValueError, match="invalid workflow type value"): + WorkflowType.value_of("invalid_type") + + def test_workflow_graph_dict_property(self) -> None: + """Test graph_dict property parses JSON correctly.""" + # Arrange + graph_data = {"nodes": [{"id": "start", "type": "start"}], "edges": list[object]()} + workflow = _workflow_models_Workflow.new( + tenant_id=str(uuid4()), + app_id=str(uuid4()), + type=WorkflowType.WORKFLOW, + version="draft", + graph=json.dumps(graph_data), + features="{}", + created_by=str(uuid4()), + environment_variables=[], + conversation_variables=[], + rag_pipeline_variables=[], + ) + + # Act + result = workflow.graph_dict + + # Assert + assert result == graph_data + assert "nodes" in result + assert len(result["nodes"]) == 1 + + def test_workflow_features_dict_property(self) -> None: + """Test features_dict property parses JSON correctly.""" + # Arrange + features_data = {"file_upload": {"enabled": True, "max_files": 5}} + workflow = _workflow_models_Workflow.new( + tenant_id=str(uuid4()), + app_id=str(uuid4()), + type=WorkflowType.WORKFLOW, + version="draft", + graph="{}", + features=json.dumps(features_data), + created_by=str(uuid4()), + environment_variables=[], + conversation_variables=[], + rag_pipeline_variables=[], + ) + + # Act + result = workflow.features_dict + + # Assert + assert result == features_data + assert result["file_upload"]["enabled"] is True + assert result["file_upload"]["max_files"] == 5 + + def test_workflow_with_marked_name_and_comment(self) -> None: + """Test workflow creation with marked name and comment.""" + # Arrange & Act + workflow = _workflow_models_Workflow.new( + tenant_id=str(uuid4()), + app_id=str(uuid4()), + type=WorkflowType.WORKFLOW, + version="v1.0", + graph="{}", + features="{}", + created_by=str(uuid4()), + environment_variables=[], + conversation_variables=[], + rag_pipeline_variables=[], + marked_name="Production Release", + marked_comment="Initial production version", + ) + + # Assert + assert workflow.marked_name == "Production Release" + assert workflow.marked_comment == "Initial production version" + + def test_workflow_version_draft_constant(self) -> None: + """Test VERSION_DRAFT constant.""" + # Assert + assert _workflow_models_Workflow.VERSION_DRAFT == "draft" + + class TestWorkflowRunStateTransitions: + """Test suite for WorkflowRun state transitions and lifecycle.""" + + def test_workflow_run_creation_with_required_fields(self) -> None: + """Test creating a workflow run with required fields.""" + # Arrange + tenant_id = str(uuid4()) + app_id = str(uuid4()) + workflow_id = str(uuid4()) + created_by = str(uuid4()) + + # Act + workflow_run = WorkflowRun( + tenant_id=tenant_id, + app_id=app_id, + workflow_id=workflow_id, + type=WorkflowType.WORKFLOW, + triggered_from=WorkflowRunTriggeredFrom.DEBUGGING, + version="draft", + status=WorkflowExecutionStatus.RUNNING, + created_by_role=CreatorUserRole.ACCOUNT, + created_by=created_by, + ) + + # Assert + assert workflow_run.tenant_id == tenant_id + assert workflow_run.app_id == app_id + assert workflow_run.workflow_id == workflow_id + assert workflow_run.type == WorkflowType.WORKFLOW + assert workflow_run.triggered_from == WorkflowRunTriggeredFrom.DEBUGGING + assert workflow_run.status == WorkflowExecutionStatus.RUNNING + assert workflow_run.created_by == created_by + + def test_workflow_run_state_transition_running_to_succeeded(self) -> None: + """Test state transition from running to succeeded.""" + # Arrange + workflow_run = WorkflowRun( + tenant_id=str(uuid4()), + app_id=str(uuid4()), + workflow_id=str(uuid4()), + type=WorkflowType.WORKFLOW, + triggered_from=WorkflowRunTriggeredFrom.APP_RUN, + version="v1.0", + status=WorkflowExecutionStatus.RUNNING, + created_by_role=CreatorUserRole.END_USER, + created_by=str(uuid4()), + ) + + # Act + workflow_run.status = WorkflowExecutionStatus.SUCCEEDED + workflow_run.finished_at = datetime.now(UTC) + workflow_run.elapsed_time = 2.5 + + # Assert + assert workflow_run.status == WorkflowExecutionStatus.SUCCEEDED + assert workflow_run.finished_at is not None + assert workflow_run.elapsed_time == 2.5 + + def test_workflow_run_state_transition_running_to_failed(self) -> None: + """Test state transition from running to failed with error.""" + # Arrange + workflow_run = WorkflowRun( + tenant_id=str(uuid4()), + app_id=str(uuid4()), + workflow_id=str(uuid4()), + type=WorkflowType.WORKFLOW, + triggered_from=WorkflowRunTriggeredFrom.APP_RUN, + version="v1.0", + status=WorkflowExecutionStatus.RUNNING, + created_by_role=CreatorUserRole.ACCOUNT, + created_by=str(uuid4()), + ) + + # Act + workflow_run.status = WorkflowExecutionStatus.FAILED + workflow_run.error = "Node execution failed: Invalid input" + workflow_run.finished_at = datetime.now(UTC) + + # Assert + assert workflow_run.status == WorkflowExecutionStatus.FAILED + assert workflow_run.error == "Node execution failed: Invalid input" + assert workflow_run.finished_at is not None + + def test_workflow_run_state_transition_running_to_stopped(self) -> None: + """Test state transition from running to stopped.""" + # Arrange + workflow_run = WorkflowRun( + tenant_id=str(uuid4()), + app_id=str(uuid4()), + workflow_id=str(uuid4()), + type=WorkflowType.WORKFLOW, + triggered_from=WorkflowRunTriggeredFrom.DEBUGGING, + version="draft", + status=WorkflowExecutionStatus.RUNNING, + created_by_role=CreatorUserRole.ACCOUNT, + created_by=str(uuid4()), + ) + + # Act + workflow_run.status = WorkflowExecutionStatus.STOPPED + workflow_run.finished_at = datetime.now(UTC) + + # Assert + assert workflow_run.status == WorkflowExecutionStatus.STOPPED + assert workflow_run.finished_at is not None + + def test_workflow_run_state_transition_running_to_paused(self) -> None: + """Test state transition from running to paused.""" + # Arrange + workflow_run = WorkflowRun( + tenant_id=str(uuid4()), + app_id=str(uuid4()), + workflow_id=str(uuid4()), + type=WorkflowType.WORKFLOW, + triggered_from=WorkflowRunTriggeredFrom.APP_RUN, + version="v1.0", + status=WorkflowExecutionStatus.RUNNING, + created_by_role=CreatorUserRole.END_USER, + created_by=str(uuid4()), + ) + + # Act + workflow_run.status = WorkflowExecutionStatus.PAUSED + + # Assert + assert workflow_run.status == WorkflowExecutionStatus.PAUSED + assert workflow_run.finished_at is None # Not finished when paused + + def test_workflow_run_state_transition_paused_to_running(self) -> None: + """Test state transition from paused back to running.""" + # Arrange + workflow_run = WorkflowRun( + tenant_id=str(uuid4()), + app_id=str(uuid4()), + workflow_id=str(uuid4()), + type=WorkflowType.WORKFLOW, + triggered_from=WorkflowRunTriggeredFrom.APP_RUN, + version="v1.0", + status=WorkflowExecutionStatus.PAUSED, + created_by_role=CreatorUserRole.ACCOUNT, + created_by=str(uuid4()), + ) + + # Act + workflow_run.status = WorkflowExecutionStatus.RUNNING + + # Assert + assert workflow_run.status == WorkflowExecutionStatus.RUNNING + + def test_workflow_run_with_partial_succeeded_status(self) -> None: + """Test workflow run with partial-succeeded status.""" + # Arrange & Act + workflow_run = WorkflowRun( + tenant_id=str(uuid4()), + app_id=str(uuid4()), + workflow_id=str(uuid4()), + type=WorkflowType.WORKFLOW, + triggered_from=WorkflowRunTriggeredFrom.APP_RUN, + version="v1.0", + status=WorkflowExecutionStatus.PARTIAL_SUCCEEDED, + created_by_role=CreatorUserRole.ACCOUNT, + created_by=str(uuid4()), + exceptions_count=2, + ) + + # Assert + assert workflow_run.status == WorkflowExecutionStatus.PARTIAL_SUCCEEDED + assert workflow_run.exceptions_count == 2 + + def test_workflow_run_with_inputs_and_outputs(self) -> None: + """Test workflow run with inputs and outputs as JSON.""" + # Arrange + inputs = {"query": "What is AI?", "context": "technology"} + outputs = {"answer": "AI is Artificial Intelligence", "confidence": 0.95} + + # Act + workflow_run = WorkflowRun( + tenant_id=str(uuid4()), + app_id=str(uuid4()), + workflow_id=str(uuid4()), + type=WorkflowType.WORKFLOW, + triggered_from=WorkflowRunTriggeredFrom.APP_RUN, + version="v1.0", + status=WorkflowExecutionStatus.SUCCEEDED, + created_by_role=CreatorUserRole.END_USER, + created_by=str(uuid4()), + inputs=json.dumps(inputs), + outputs=json.dumps(outputs), + ) + + # Assert + assert workflow_run.inputs_dict == inputs + assert workflow_run.outputs_dict == outputs + + def test_workflow_run_graph_dict_property(self) -> None: + """Test graph_dict property for workflow run.""" + # Arrange + graph = {"nodes": [{"id": "start", "type": "start"}], "edges": list[object]()} + workflow_run = WorkflowRun( + tenant_id=str(uuid4()), + app_id=str(uuid4()), + workflow_id=str(uuid4()), + type=WorkflowType.WORKFLOW, + triggered_from=WorkflowRunTriggeredFrom.DEBUGGING, + version="draft", + status=WorkflowExecutionStatus.RUNNING, + created_by_role=CreatorUserRole.ACCOUNT, + created_by=str(uuid4()), + graph=json.dumps(graph), + ) + + # Act + result = workflow_run.graph_dict + + # Assert + assert result == graph + assert "nodes" in result + + def test_workflow_run_to_dict_serialization(self) -> None: + """Test WorkflowRun to_dict method.""" + # Arrange + workflow_run_id = str(uuid4()) + tenant_id = str(uuid4()) + app_id = str(uuid4()) + workflow_id = str(uuid4()) + created_by = str(uuid4()) + + workflow_run = WorkflowRun( + tenant_id=tenant_id, + app_id=app_id, + workflow_id=workflow_id, + type=WorkflowType.WORKFLOW, + triggered_from=WorkflowRunTriggeredFrom.APP_RUN, + version="v1.0", + status=WorkflowExecutionStatus.SUCCEEDED, + created_by_role=CreatorUserRole.ACCOUNT, + created_by=created_by, + total_tokens=1500, + total_steps=5, + ) + workflow_run.id = workflow_run_id + + # Act + result = workflow_run.to_dict() + + # Assert + assert result["id"] == workflow_run_id + assert result["tenant_id"] == tenant_id + assert result["app_id"] == app_id + assert result["workflow_id"] == workflow_id + assert result["status"] == WorkflowExecutionStatus.SUCCEEDED + assert result["total_tokens"] == 1500 + assert result["total_steps"] == 5 + + def test_workflow_run_from_dict_deserialization(self) -> None: + """Test WorkflowRun from_dict method.""" + # Arrange + data = { + "id": str(uuid4()), + "tenant_id": str(uuid4()), + "app_id": str(uuid4()), + "workflow_id": str(uuid4()), + "type": WorkflowType.WORKFLOW, + "triggered_from": WorkflowRunTriggeredFrom.APP_RUN, + "version": "v1.0", + "graph": {"nodes": list[object](), "edges": list[object]()}, + "inputs": {"query": "test"}, + "status": WorkflowExecutionStatus.SUCCEEDED, + "outputs": {"result": "success"}, + "error": None, + "elapsed_time": 3.5, + "total_tokens": 2000, + "total_steps": 10, + "created_by_role": CreatorUserRole.ACCOUNT, + "created_by": str(uuid4()), + "created_at": datetime.now(UTC), + "finished_at": datetime.now(UTC), + "exceptions_count": 0, + } + + # Act + workflow_run = WorkflowRun.from_dict(data) + + # Assert + assert workflow_run.id == data["id"] + assert workflow_run.workflow_id == data["workflow_id"] + assert workflow_run.status == WorkflowExecutionStatus.SUCCEEDED + assert workflow_run.total_tokens == 2000 + + class TestNodeExecutionRelationships: + """Test suite for WorkflowNodeExecutionModel relationships and data.""" + + def test_node_execution_creation_with_required_fields(self) -> None: + """Test creating a node execution with required fields.""" + # Arrange + tenant_id = str(uuid4()) + app_id = str(uuid4()) + workflow_id = str(uuid4()) + workflow_run_id = str(uuid4()) + created_by = str(uuid4()) + + # Act + node_execution = WorkflowNodeExecutionModel( + tenant_id=tenant_id, + app_id=app_id, + workflow_id=workflow_id, + triggered_from=WorkflowNodeExecutionTriggeredFrom.WORKFLOW_RUN, + workflow_run_id=workflow_run_id, + index=1, + node_id="start", + node_type=BuiltinNodeTypes.START, + title="Start Node", + status=WorkflowNodeExecutionStatus.SUCCEEDED, + created_by_role=CreatorUserRole.ACCOUNT, + created_by=created_by, + ) + + # Assert + assert node_execution.tenant_id == tenant_id + assert node_execution.app_id == app_id + assert node_execution.workflow_id == workflow_id + assert node_execution.workflow_run_id == workflow_run_id + assert node_execution.node_id == "start" + assert node_execution.node_type == BuiltinNodeTypes.START + assert node_execution.index == 1 + + def test_node_execution_with_predecessor_relationship(self) -> None: + """Test node execution with predecessor node relationship.""" + # Arrange + predecessor_node_id = "start" + current_node_id = "llm_1" + + # Act + node_execution = WorkflowNodeExecutionModel( + tenant_id=str(uuid4()), + app_id=str(uuid4()), + workflow_id=str(uuid4()), + triggered_from=WorkflowNodeExecutionTriggeredFrom.WORKFLOW_RUN, + workflow_run_id=str(uuid4()), + index=2, + predecessor_node_id=predecessor_node_id, + node_id=current_node_id, + node_type=BuiltinNodeTypes.LLM, + title="LLM Node", + status=WorkflowNodeExecutionStatus.RUNNING, + created_by_role=CreatorUserRole.ACCOUNT, + created_by=str(uuid4()), + ) + + # Assert + assert node_execution.predecessor_node_id == predecessor_node_id + assert node_execution.node_id == current_node_id + assert node_execution.index == 2 + + def test_node_execution_single_step_debugging(self) -> None: + """Test node execution for single-step debugging (no workflow_run_id).""" + # Arrange & Act + node_execution = WorkflowNodeExecutionModel( + tenant_id=str(uuid4()), + app_id=str(uuid4()), + workflow_id=str(uuid4()), + triggered_from=WorkflowNodeExecutionTriggeredFrom.SINGLE_STEP.value, + workflow_run_id=None, # Single-step has no workflow run + index=1, + node_id="llm_test", + node_type=BuiltinNodeTypes.LLM, + title="Test LLM", + status=WorkflowNodeExecutionStatus.SUCCEEDED, + created_by_role=CreatorUserRole.ACCOUNT, + created_by=str(uuid4()), + ) + + # Assert + assert node_execution.triggered_from == WorkflowNodeExecutionTriggeredFrom.SINGLE_STEP.value + assert node_execution.workflow_run_id is None + + def test_node_execution_with_inputs_outputs_process_data(self) -> None: + """Test node execution with inputs, outputs, and process_data.""" + # Arrange + inputs = {"query": "What is AI?", "temperature": 0.7} + outputs = {"answer": "AI is Artificial Intelligence"} + process_data = {"tokens_used": 150, "model": "gpt-4"} + + # Act + node_execution = WorkflowNodeExecutionModel( + tenant_id=str(uuid4()), + app_id=str(uuid4()), + workflow_id=str(uuid4()), + triggered_from=WorkflowNodeExecutionTriggeredFrom.WORKFLOW_RUN, + workflow_run_id=str(uuid4()), + index=1, + node_id="llm_1", + node_type=BuiltinNodeTypes.LLM, + title="LLM Node", + status=WorkflowNodeExecutionStatus.SUCCEEDED, + created_by_role=CreatorUserRole.ACCOUNT, + created_by=str(uuid4()), + inputs=json.dumps(inputs), + outputs=json.dumps(outputs), + process_data=json.dumps(process_data), + ) + + # Assert + assert node_execution.inputs_dict == inputs + assert node_execution.outputs_dict == outputs + assert node_execution.process_data_dict == process_data + + def test_node_execution_status_transitions(self) -> None: + """Test node execution status transitions.""" + # Arrange + node_execution = WorkflowNodeExecutionModel( + tenant_id=str(uuid4()), + app_id=str(uuid4()), + workflow_id=str(uuid4()), + triggered_from=WorkflowNodeExecutionTriggeredFrom.WORKFLOW_RUN, + workflow_run_id=str(uuid4()), + index=1, + node_id="code_1", + node_type=BuiltinNodeTypes.CODE, + title="Code Node", + status=WorkflowNodeExecutionStatus.RUNNING, + created_by_role=CreatorUserRole.ACCOUNT, + created_by=str(uuid4()), + ) + + # Act - transition to succeeded + node_execution.status = WorkflowNodeExecutionStatus.SUCCEEDED + node_execution.elapsed_time = 1.2 + node_execution.finished_at = datetime.now(UTC) + + # Assert + assert node_execution.status == WorkflowNodeExecutionStatus.SUCCEEDED + assert node_execution.elapsed_time == 1.2 + assert node_execution.finished_at is not None + + def test_node_execution_with_error(self) -> None: + """Test node execution with error status.""" + # Arrange + error_message = "Code execution failed: SyntaxError on line 5" + + # Act + node_execution = WorkflowNodeExecutionModel( + tenant_id=str(uuid4()), + app_id=str(uuid4()), + workflow_id=str(uuid4()), + triggered_from=WorkflowNodeExecutionTriggeredFrom.WORKFLOW_RUN, + workflow_run_id=str(uuid4()), + index=3, + node_id="code_1", + node_type=BuiltinNodeTypes.CODE, + title="Code Node", + status=WorkflowNodeExecutionStatus.FAILED, + error=error_message, + created_by_role=CreatorUserRole.ACCOUNT, + created_by=str(uuid4()), + ) + + # Assert + assert node_execution.status == WorkflowNodeExecutionStatus.FAILED + assert node_execution.error == error_message + + def test_node_execution_with_metadata(self) -> None: + """Test node execution with execution metadata.""" + # Arrange + metadata = { + "total_tokens": 500, + "total_price": 0.01, + "currency": "USD", + "tool_info": {"provider": "openai", "tool": "gpt-4"}, + } + + # Act + node_execution = WorkflowNodeExecutionModel( + tenant_id=str(uuid4()), + app_id=str(uuid4()), + workflow_id=str(uuid4()), + triggered_from=WorkflowNodeExecutionTriggeredFrom.WORKFLOW_RUN, + workflow_run_id=str(uuid4()), + index=1, + node_id="llm_1", + node_type=BuiltinNodeTypes.LLM, + title="LLM Node", + status=WorkflowNodeExecutionStatus.SUCCEEDED, + created_by_role=CreatorUserRole.ACCOUNT, + created_by=str(uuid4()), + execution_metadata=json.dumps(metadata), + ) + + # Assert + assert node_execution.execution_metadata_dict == metadata + assert node_execution.execution_metadata_dict["total_tokens"] == 500 + + def test_node_execution_metadata_dict_empty(self) -> None: + """Test execution_metadata_dict returns empty dict when metadata is None.""" + # Arrange + node_execution = WorkflowNodeExecutionModel( + tenant_id=str(uuid4()), + app_id=str(uuid4()), + workflow_id=str(uuid4()), + triggered_from=WorkflowNodeExecutionTriggeredFrom.WORKFLOW_RUN, + workflow_run_id=str(uuid4()), + index=1, + node_id="start", + node_type=BuiltinNodeTypes.START, + title="Start", + status=WorkflowNodeExecutionStatus.SUCCEEDED, + created_by_role=CreatorUserRole.ACCOUNT, + created_by=str(uuid4()), + execution_metadata=None, + ) + + # Act + result = node_execution.execution_metadata_dict + + # Assert + assert result == {} + + def test_node_execution_different_node_types(self) -> None: + """Test node execution with different node types.""" + # Test various node types + node_types = [ + (BuiltinNodeTypes.START, "Start Node"), + (BuiltinNodeTypes.LLM, "LLM Node"), + (BuiltinNodeTypes.CODE, "Code Node"), + (BuiltinNodeTypes.TOOL, "Tool Node"), + (BuiltinNodeTypes.IF_ELSE, "Conditional Node"), + (BuiltinNodeTypes.END, "End Node"), + ] + + for node_type, title in node_types: + # Act + node_execution = WorkflowNodeExecutionModel( + tenant_id=str(uuid4()), + app_id=str(uuid4()), + workflow_id=str(uuid4()), + triggered_from=WorkflowNodeExecutionTriggeredFrom.WORKFLOW_RUN, + workflow_run_id=str(uuid4()), + index=1, + node_id=f"{node_type}_1", + node_type=node_type, + title=title, + status=WorkflowNodeExecutionStatus.SUCCEEDED, + created_by_role=CreatorUserRole.ACCOUNT, + created_by=str(uuid4()), + ) + + # Assert + assert node_execution.node_type == node_type + assert node_execution.title == title + + class TestGraphConfigurationValidation: + """Test suite for graph configuration validation.""" + + def test_workflow_graph_with_nodes_and_edges(self) -> None: + """Test workflow graph configuration with nodes and edges.""" + # Arrange + graph_config = { + "nodes": [ + {"id": "start", "type": "start", "data": {"title": "Start"}}, + {"id": "llm_1", "type": "llm", "data": {"title": "LLM Node", "model": "gpt-4"}}, + {"id": "end", "type": "end", "data": {"title": "End"}}, + ], + "edges": [ + {"source": "start", "target": "llm_1"}, + {"source": "llm_1", "target": "end"}, + ], + } + + # Act + workflow = _workflow_models_Workflow.new( + tenant_id=str(uuid4()), + app_id=str(uuid4()), + type=WorkflowType.WORKFLOW, + version="draft", + graph=json.dumps(graph_config), + features="{}", + created_by=str(uuid4()), + environment_variables=[], + conversation_variables=[], + rag_pipeline_variables=[], + ) + + # Assert + graph_dict = workflow.graph_dict + assert len(graph_dict["nodes"]) == 3 + assert len(graph_dict["edges"]) == 2 + assert graph_dict["nodes"][0]["id"] == "start" + assert graph_dict["edges"][0]["source"] == "start" + assert graph_dict["edges"][0]["target"] == "llm_1" + + def test_workflow_graph_empty_configuration(self) -> None: + """Test workflow with empty graph configuration.""" + # Arrange + graph_config = {"nodes": list[object](), "edges": list[object]()} + + # Act + workflow = _workflow_models_Workflow.new( + tenant_id=str(uuid4()), + app_id=str(uuid4()), + type=WorkflowType.WORKFLOW, + version="draft", + graph=json.dumps(graph_config), + features="{}", + created_by=str(uuid4()), + environment_variables=[], + conversation_variables=[], + rag_pipeline_variables=[], + ) + + # Assert + graph_dict = workflow.graph_dict + assert graph_dict["nodes"] == [] + assert graph_dict["edges"] == [] + + def test_workflow_graph_complex_node_data(self) -> None: + """Test workflow graph with complex node data structures.""" + # Arrange + graph_config = { + "nodes": [ + { + "id": "llm_1", + "type": "llm", + "data": { + "title": "Advanced LLM", + "model": {"provider": "openai", "name": "gpt-4", "mode": "chat"}, + "prompt_template": [ + {"role": "system", "text": "You are a helpful assistant"}, + {"role": "user", "text": "{{query}}"}, + ], + "model_parameters": {"temperature": 0.7, "max_tokens": 2000}, + }, + } + ], + "edges": list[object](), + } + + # Act + workflow = _workflow_models_Workflow.new( + tenant_id=str(uuid4()), + app_id=str(uuid4()), + type=WorkflowType.WORKFLOW, + version="draft", + graph=json.dumps(graph_config), + features="{}", + created_by=str(uuid4()), + environment_variables=[], + conversation_variables=[], + rag_pipeline_variables=[], + ) + + # Assert + graph_dict = workflow.graph_dict + node_data = graph_dict["nodes"][0]["data"] + assert node_data["model"]["provider"] == "openai" + assert node_data["model_parameters"]["temperature"] == 0.7 + assert len(node_data["prompt_template"]) == 2 + + def test_workflow_run_graph_preservation(self) -> None: + """Test that WorkflowRun preserves graph configuration from Workflow.""" + # Arrange + original_graph = { + "nodes": [ + {"id": "start", "type": "start"}, + {"id": "end", "type": "end"}, + ], + "edges": [{"source": "start", "target": "end"}], + } + + # Act + workflow_run = WorkflowRun( + tenant_id=str(uuid4()), + app_id=str(uuid4()), + workflow_id=str(uuid4()), + type=WorkflowType.WORKFLOW, + triggered_from=WorkflowRunTriggeredFrom.APP_RUN, + version="v1.0", + status=WorkflowExecutionStatus.RUNNING, + created_by_role=CreatorUserRole.ACCOUNT, + created_by=str(uuid4()), + graph=json.dumps(original_graph), + ) + + # Assert + assert workflow_run.graph_dict == original_graph + assert len(workflow_run.graph_dict["nodes"]) == 2 + + def test_workflow_graph_with_conditional_branches(self) -> None: + """Test workflow graph with conditional branching (if-else).""" + # Arrange + graph_config = { + "nodes": [ + {"id": "start", "type": "start"}, + {"id": "if_else_1", "type": "if-else", "data": {"conditions": list[object]()}}, + {"id": "branch_true", "type": "llm"}, + {"id": "branch_false", "type": "code"}, + {"id": "end", "type": "end"}, + ], + "edges": [ + {"source": "start", "target": "if_else_1"}, + {"source": "if_else_1", "sourceHandle": "true", "target": "branch_true"}, + {"source": "if_else_1", "sourceHandle": "false", "target": "branch_false"}, + {"source": "branch_true", "target": "end"}, + {"source": "branch_false", "target": "end"}, + ], + } + + # Act + workflow = _workflow_models_Workflow.new( + tenant_id=str(uuid4()), + app_id=str(uuid4()), + type=WorkflowType.WORKFLOW, + version="draft", + graph=json.dumps(graph_config), + features="{}", + created_by=str(uuid4()), + environment_variables=[], + conversation_variables=[], + rag_pipeline_variables=[], + ) + + # Assert + graph_dict = workflow.graph_dict + assert len(graph_dict["nodes"]) == 5 + assert len(graph_dict["edges"]) == 5 + # Verify conditional edges + conditional_edges = [e for e in graph_dict["edges"] if "sourceHandle" in e] + assert len(conditional_edges) == 2 + + def test_workflow_graph_with_loop_structure(self) -> None: + """Test workflow graph with loop/iteration structure.""" + # Arrange + graph_config = { + "nodes": [ + {"id": "start", "type": "start"}, + {"id": "iteration_1", "type": "iteration", "data": {"iterator": "items"}}, + {"id": "loop_body", "type": "llm"}, + {"id": "end", "type": "end"}, + ], + "edges": [ + {"source": "start", "target": "iteration_1"}, + {"source": "iteration_1", "target": "loop_body"}, + {"source": "loop_body", "target": "iteration_1"}, + {"source": "iteration_1", "target": "end"}, + ], + } + + # Act + workflow = _workflow_models_Workflow.new( + tenant_id=str(uuid4()), + app_id=str(uuid4()), + type=WorkflowType.WORKFLOW, + version="draft", + graph=json.dumps(graph_config), + features="{}", + created_by=str(uuid4()), + environment_variables=[], + conversation_variables=[], + rag_pipeline_variables=[], + ) + + # Assert + graph_dict = workflow.graph_dict + iteration_node = next(n for n in graph_dict["nodes"] if n["type"] == "iteration") + assert iteration_node["data"]["iterator"] == "items" + + def test_workflow_graph_dict_with_null_graph(self) -> None: + """Test graph_dict property when graph is None.""" + # Arrange + workflow = _workflow_models_Workflow.new( + tenant_id=str(uuid4()), + app_id=str(uuid4()), + type=WorkflowType.WORKFLOW, + version="draft", + graph=None, # pyrefly: ignore[bad-argument-type] + features="{}", + created_by=str(uuid4()), + environment_variables=[], + conversation_variables=[], + rag_pipeline_variables=[], + ) + + # Act + result = workflow.graph_dict + + # Assert + assert result == {} + + def test_workflow_run_inputs_dict_with_null_inputs(self) -> None: + """Test inputs_dict property when inputs is None.""" + # Arrange + workflow_run = WorkflowRun( + tenant_id=str(uuid4()), + app_id=str(uuid4()), + workflow_id=str(uuid4()), + type=WorkflowType.WORKFLOW, + triggered_from=WorkflowRunTriggeredFrom.APP_RUN, + version="v1.0", + status=WorkflowExecutionStatus.RUNNING, + created_by_role=CreatorUserRole.ACCOUNT, + created_by=str(uuid4()), + inputs=None, + ) + + # Act + result = workflow_run.inputs_dict + + # Assert + assert result == {} + + def test_workflow_run_outputs_dict_with_null_outputs(self) -> None: + """Test outputs_dict property when outputs is None.""" + # Arrange + workflow_run = WorkflowRun( + tenant_id=str(uuid4()), + app_id=str(uuid4()), + workflow_id=str(uuid4()), + type=WorkflowType.WORKFLOW, + triggered_from=WorkflowRunTriggeredFrom.APP_RUN, + version="v1.0", + status=WorkflowExecutionStatus.RUNNING, + created_by_role=CreatorUserRole.ACCOUNT, + created_by=str(uuid4()), + outputs=None, + ) + + # Act + result = workflow_run.outputs_dict + + # Assert + assert result == {} + + def test_node_execution_inputs_dict_with_null_inputs(self) -> None: + """Test node execution inputs_dict when inputs is None.""" + # Arrange + node_execution = WorkflowNodeExecutionModel( + tenant_id=str(uuid4()), + app_id=str(uuid4()), + workflow_id=str(uuid4()), + triggered_from=WorkflowNodeExecutionTriggeredFrom.WORKFLOW_RUN, + workflow_run_id=str(uuid4()), + index=1, + node_id="start", + node_type=BuiltinNodeTypes.START, + title="Start", + status=WorkflowNodeExecutionStatus.SUCCEEDED, + created_by_role=CreatorUserRole.ACCOUNT, + created_by=str(uuid4()), + inputs=None, + ) + + # Act + result = node_execution.inputs_dict + + # Assert + assert result is None + + def test_node_execution_outputs_dict_with_null_outputs(self) -> None: + """Test node execution outputs_dict when outputs is None.""" + # Arrange + node_execution = WorkflowNodeExecutionModel( + tenant_id=str(uuid4()), + app_id=str(uuid4()), + workflow_id=str(uuid4()), + triggered_from=WorkflowNodeExecutionTriggeredFrom.WORKFLOW_RUN, + workflow_run_id=str(uuid4()), + index=1, + node_id="start", + node_type=BuiltinNodeTypes.START, + title="Start", + status=WorkflowNodeExecutionStatus.SUCCEEDED, + created_by_role=CreatorUserRole.ACCOUNT, + created_by=str(uuid4()), + outputs=None, + ) + + # Act + result = node_execution.outputs_dict + + # Assert + assert result is None + + +"""Regression coverage for ``WorkflowNodeExecutionModel`` account accessors. + +Ensures the ``@property``→session-parameter refactor preserves the role-based dispatch: +``created_by_account`` looks up an Account only when role is ACCOUNT; ``created_by_end_user`` +looks up an EndUser only when role is END_USER. Also covers the session-carrying response +source used by the console endpoints that validate ``WorkflowRunNodeExecutionResponse`` +with ``from_attributes=True``. +""" + + +def _execution(role: CreatorUserRole, created_by: str) -> WorkflowNodeExecutionModel: + """Construct a WorkflowNodeExecutionModel without touching the database.""" + return WorkflowNodeExecutionModel( + tenant_id="00000000-0000-0000-0000-000000000001", + app_id="00000000-0000-0000-0000-000000000002", + workflow_id=str(uuid4()), + triggered_from=WorkflowNodeExecutionTriggeredFrom.WORKFLOW_RUN, + workflow_run_id=None, + index=1, + predecessor_node_id=None, + node_execution_id=None, + node_id="n1", + node_type="start", + title="Start", + inputs=None, + process_data=None, + outputs=None, + status="succeeded", + error=None, + elapsed_time=0.0, + execution_metadata=None, + created_by_role=role, + created_by=created_by, + ) + + +class TestWorkflowNodeExecutionAccessorsCases: + class TestCreatedByAccount: + def test_returns_account_lookup_when_role_is_account(self, sqlite_session: Session) -> None: + account = Account(name="Test Account", email="test@example.com") + sqlite_session.add(account) + sqlite_session.flush() + execution = _execution(CreatorUserRole.ACCOUNT, created_by=account.id) + + result = execution.created_by_account(session=sqlite_session) + + assert result is not None + assert result.id == account.id + + def test_returns_none_when_role_is_end_user(self, sqlite_session: Session) -> None: + account = Account(name="Test Account", email="test@example.com") + sqlite_session.add(account) + sqlite_session.flush() + execution = _execution(CreatorUserRole.END_USER, created_by=account.id) + + assert execution.created_by_account(session=sqlite_session) is None + + class TestCreatedByEndUser: + def test_returns_end_user_lookup_when_role_is_end_user(self, sqlite_session: Session) -> None: + end_user = EndUser( + tenant_id="00000000-0000-0000-0000-000000000001", + type=EndUserType.BROWSER, + session_id="session-1", + ) + sqlite_session.add(end_user) + sqlite_session.flush() + execution = _execution(CreatorUserRole.END_USER, created_by=end_user.id) + + result = execution.created_by_end_user(session=sqlite_session) + + assert result is not None + assert result.id == end_user.id + + def test_returns_none_when_role_is_account(self, sqlite_session: Session) -> None: + end_user = EndUser( + tenant_id="00000000-0000-0000-0000-000000000001", + type=EndUserType.BROWSER, + session_id="session-1", + ) + sqlite_session.add(end_user) + sqlite_session.flush() + execution = _execution(CreatorUserRole.ACCOUNT, created_by=end_user.id) + + assert execution.created_by_end_user(session=sqlite_session) is None + + class TestNodeExecutionResponseSource: + def test_accessors_resolve_via_wrapped_session_and_other_attributes_proxy( + self, sqlite_session: Session + ) -> None: + account = Account(name="Test Account", email="test@example.com") + sqlite_session.add(account) + sqlite_session.flush() + execution = _execution(CreatorUserRole.ACCOUNT, created_by=account.id) + + source = node_execution_response_source(execution, session=sqlite_session) + + resolved = source.created_by_account + assert resolved is not None + assert resolved.id == account.id + assert source.created_by_end_user is None + assert source.node_id == "n1" + + +"""SQLite-backed tests for workflow node execution offload behavior.""" + + +TABLES = (WorkflowNodeExecutionModel, WorkflowNodeExecutionOffload, UploadFile) + + +def _persist_upload_file(session: Session, *, key: str = "offload/process-data.json") -> UploadFile: + upload_file = UploadFile( + tenant_id="tenant-1", + storage_type=StorageType.LOCAL, + key=key, + name="process-data.json", + size=18, + extension="json", + mime_type="application/json", + created_by_role=CreatorUserRole.ACCOUNT, + created_by="account-1", + created_at=datetime.now(UTC), + used=True, + ) + session.add(upload_file) + session.flush() + return upload_file + + +def _persist_execution( + session: Session, + *, + process_data: str = '{"test": "data"}', + offloads: tuple[tuple[ExecutionOffLoadType, str], ...] = (), +) -> WorkflowNodeExecutionModel: + execution = WorkflowNodeExecutionModel( + tenant_id="tenant-1", + app_id="app-1", + workflow_id="workflow-1", + triggered_from=WorkflowNodeExecutionTriggeredFrom.WORKFLOW_RUN, + workflow_run_id="run-1", + index=1, + predecessor_node_id=None, + node_execution_id="node-execution-1", + node_id="node-1", + node_type="code", + title="Code", + inputs='{"input": "value"}', + process_data=process_data, + outputs='{"output": "value"}', + status=WorkflowNodeExecutionStatus.SUCCEEDED, + error=None, + elapsed_time=0, + execution_metadata=None, + created_by_role=CreatorUserRole.ACCOUNT, + created_by="account-1", + finished_at=None, + ) + session.add(execution) + session.flush() + + session.add_all( + [ + WorkflowNodeExecutionOffload( + tenant_id=execution.tenant_id, + app_id=execution.app_id, + node_execution_id=execution.id, + type_=type_, + file_id=file_id, + ) + for type_, file_id in offloads + ] + ) + session.commit() + session.expunge_all() + + stmt = WorkflowNodeExecutionModel.preload_offload_data( + select(WorkflowNodeExecutionModel).where(WorkflowNodeExecutionModel.id == execution.id) + ) + loaded_execution = session.scalar(stmt) + assert loaded_execution is not None + return loaded_execution + + +class TestWorkflowNodeExecutionOffloadCases: + @pytest.mark.parametrize("sqlite_session", [TABLES], indirect=True) + class TestWorkflowNodeExecutionModel: + """Exercise truncation flags and content restoration using persisted offload rows.""" + + def test_process_data_truncated_property_false_when_no_offload_data(self, sqlite_session: Session) -> None: + execution = _persist_execution(sqlite_session) + + assert execution.process_data_truncated is False + + def test_process_data_truncated_property_false_when_no_process_data_file(self, sqlite_session: Session) -> None: + execution = _persist_execution( + sqlite_session, + offloads=( + (ExecutionOffLoadType.INPUTS, "inputs-file"), + (ExecutionOffLoadType.OUTPUTS, "outputs-file"), + ), + ) + + assert execution.process_data_truncated is False + + def test_process_data_truncated_property_true_when_process_data_file_exists( + self, sqlite_session: Session + ) -> None: + execution = _persist_execution( + sqlite_session, + offloads=((ExecutionOffLoadType.PROCESS_DATA, "process-data-file-id"),), + ) + + assert execution.process_data_truncated is True + + def test_load_full_process_data_with_no_offload_data(self, sqlite_session: Session) -> None: + execution = _persist_execution(sqlite_session) + storage = Mock() + + result = execution.load_full_process_data(sqlite_session, storage) + + assert result == {"test": "data"} + storage.load.assert_not_called() + + def test_load_full_process_data_with_no_file(self, sqlite_session: Session) -> None: + execution = _persist_execution( + sqlite_session, + offloads=((ExecutionOffLoadType.INPUTS, "inputs-file"),), + ) + storage = Mock() + + result = execution.load_full_process_data(sqlite_session, storage) + + assert result == {"test": "data"} + storage.load.assert_not_called() + + def test_load_full_process_data_with_file(self, sqlite_session: Session) -> None: + upload_file = _persist_upload_file(sqlite_session) + execution = _persist_execution( + sqlite_session, + process_data='{"truncated": true}', + offloads=((ExecutionOffLoadType.PROCESS_DATA, upload_file.id),), + ) + full_process_data = {"full": "data", "large_field": "x" * 10000} + storage = Mock() + storage.load.return_value = json.dumps(full_process_data) + + result = execution.load_full_process_data(sqlite_session, storage) + + assert result == full_process_data + storage.load.assert_called_once_with(upload_file.key) + + def test_consistency_with_inputs_outputs_truncation(self, sqlite_session: Session) -> None: + execution = _persist_execution( + sqlite_session, + offloads=( + (ExecutionOffLoadType.INPUTS, "inputs-file"), + (ExecutionOffLoadType.OUTPUTS, "outputs-file"), + (ExecutionOffLoadType.PROCESS_DATA, "process-data-file"), + ), + ) + + assert execution.inputs_truncated is True + assert execution.outputs_truncated is True + assert execution.process_data_truncated is True + + def test_mixed_truncation_states(self, sqlite_session: Session) -> None: + execution = _persist_execution( + sqlite_session, + offloads=((ExecutionOffLoadType.PROCESS_DATA, "process-data-file"),), + ) + + assert execution.inputs_truncated is False + assert execution.outputs_truncated is False + assert execution.process_data_truncated is True + + def test_preload_offload_data_and_files_loads_process_data_file(self, sqlite_session: Session) -> None: + upload_file = _persist_upload_file(sqlite_session) + execution = _persist_execution( + sqlite_session, + offloads=((ExecutionOffLoadType.PROCESS_DATA, upload_file.id),), + ) + execution_id = execution.id + sqlite_session.expunge_all() + + stmt = WorkflowNodeExecutionModel.preload_offload_data_and_files( + select(WorkflowNodeExecutionModel).where(WorkflowNodeExecutionModel.id == execution_id) + ) + preloaded_execution = sqlite_session.scalar(stmt) + + assert preloaded_execution is not None + assert len(preloaded_execution.offload_data) == 1 + assert preloaded_execution.offload_data[0].file is not None + assert preloaded_execution.offload_data[0].file.key == upload_file.key + + +"""Regression coverage for ``WorkflowRun`` account accessors. + +Ensures the ``@property``→session-parameter refactor preserves the role-based dispatch: +``created_by_account`` looks up an Account only when role is ACCOUNT; ``created_by_end_user`` +looks up an EndUser only when role is END_USER. Also covers the session-carrying response +sources used by the console endpoints that validate the run detail / pagination responses +with ``from_attributes=True``. +""" + + +def _run(role: CreatorUserRole, created_by: str) -> WorkflowRun: + """Construct a WorkflowRun without touching the database.""" + return WorkflowRun( + id=str(uuid4()), + tenant_id="00000000-0000-0000-0000-000000000001", + app_id="00000000-0000-0000-0000-000000000002", + workflow_id=str(uuid4()), + type="workflow", + triggered_from=WorkflowRunTriggeredFrom.APP_RUN, + version="1.0", + status="succeeded", + created_by_role=role, + created_by=created_by, + ) + + +class TestWorkflowRunAccessorsCases: + class TestCreatedByAccount: + def test_returns_account_lookup_when_role_is_account(self, sqlite_session: Session) -> None: + account = Account(name="Test Account", email="test@example.com") + sqlite_session.add(account) + sqlite_session.flush() + run = _run(CreatorUserRole.ACCOUNT, created_by=account.id) + + result = run.created_by_account(session=sqlite_session) + + assert result is not None + assert result.id == account.id + + def test_returns_none_when_role_is_end_user(self, sqlite_session: Session) -> None: + account = Account(name="Test Account", email="test@example.com") + sqlite_session.add(account) + sqlite_session.flush() + run = _run(CreatorUserRole.END_USER, created_by=account.id) + + assert run.created_by_account(session=sqlite_session) is None + + class TestCreatedByEndUser: + def test_returns_end_user_lookup_when_role_is_end_user(self, sqlite_session: Session) -> None: + end_user = EndUser( + tenant_id="00000000-0000-0000-0000-000000000001", + type=EndUserType.BROWSER, + session_id="session-1", + ) + sqlite_session.add(end_user) + sqlite_session.flush() + run = _run(CreatorUserRole.END_USER, created_by=end_user.id) + + result = run.created_by_end_user(session=sqlite_session) + + assert result is not None + assert result.id == end_user.id + + def test_returns_none_when_role_is_account(self, sqlite_session: Session) -> None: + end_user = EndUser( + tenant_id="00000000-0000-0000-0000-000000000001", + type=EndUserType.BROWSER, + session_id="session-1", + ) + sqlite_session.add(end_user) + sqlite_session.flush() + run = _run(CreatorUserRole.ACCOUNT, created_by=end_user.id) + + assert run.created_by_end_user(session=sqlite_session) is None + + class TestWorkflowRunResponseSources: + def test_detail_source_resolves_accessors_and_proxies_other_attributes(self, sqlite_session: Session) -> None: + account = Account(name="Test Account", email="test@example.com") + sqlite_session.add(account) + sqlite_session.flush() + run = _run(CreatorUserRole.ACCOUNT, created_by=account.id) + + source = workflow_run_response_source(run, session=sqlite_session) + + resolved = source.created_by_account + assert resolved is not None + assert resolved.id == account.id + assert source.created_by_end_user is None + assert source.version == "1.0" + + def test_pagination_source_wraps_each_run(self, sqlite_session: Session) -> None: + account = Account(name="Test Account", email="test@example.com") + sqlite_session.add(account) + sqlite_session.flush() + runs = [_run(CreatorUserRole.ACCOUNT, created_by=account.id) for _ in range(2)] + pagination = SimpleNamespace(data=runs, limit=20, has_more=False) + + payload = workflow_run_pagination_response_source(pagination, session=sqlite_session) + + assert payload["limit"] == 20 + assert payload["has_more"] is False + assert [item.created_by_account.id for item in payload["data"]] == [account.id, account.id] diff --git a/api/tests/unit_tests/models/test_workflow_app_log_models.py b/api/tests/unit_tests/models/test_workflow_app_log_models.py deleted file mode 100644 index 9e06e069190..00000000000 --- a/api/tests/unit_tests/models/test_workflow_app_log_models.py +++ /dev/null @@ -1,80 +0,0 @@ -"""Regression coverage for ``models.workflow.WorkflowAppLog`` account accessors. - -Ensures the ``@property``→session-parameter refactor preserves the role-based dispatch: -``created_by_account`` looks up an Account only when role is ACCOUNT; ``created_by_end_user`` -looks up an EndUser only when role is END_USER. - -Both accessors are exercised against the real ``sqlite_session`` fixture (a genuine -SQLAlchemy ``Session`` bound to a pristine full-schema SQLite database) so the assertions -cover actual query behaviour rather than a mock's recorded call. -""" - -from sqlalchemy.orm import Session - -from models.account import Account -from models.enums import CreatorUserRole, EndUserType -from models.model import EndUser -from models.workflow import WorkflowAppLog, WorkflowAppLogCreatedFrom - - -def _log(role: CreatorUserRole, created_by: str) -> WorkflowAppLog: - """Construct a WorkflowAppLog without touching the database.""" - return WorkflowAppLog( - tenant_id="00000000-0000-0000-0000-000000000001", - app_id="00000000-0000-0000-0000-000000000002", - workflow_id="00000000-0000-0000-0000-000000000003", - workflow_run_id="00000000-0000-0000-0000-000000000004", - created_from=WorkflowAppLogCreatedFrom.WEB_APP, - created_by_role=role, - created_by=created_by, - ) - - -class TestCreatedByAccount: - def test_returns_account_lookup_when_role_is_account(self, sqlite_session: Session) -> None: - account = Account(name="Test Account", email="test@example.com") - sqlite_session.add(account) - sqlite_session.flush() - log = _log(CreatorUserRole.ACCOUNT, created_by=account.id) - - result = log.created_by_account(session=sqlite_session) - - assert result is not None - assert result.id == account.id - - def test_returns_none_when_role_is_end_user(self, sqlite_session: Session) -> None: - account = Account(name="Test Account", email="test@example.com") - sqlite_session.add(account) - sqlite_session.flush() - log = _log(CreatorUserRole.END_USER, created_by=account.id) - - assert log.created_by_account(session=sqlite_session) is None - - -class TestCreatedByEndUser: - def test_returns_end_user_lookup_when_role_is_end_user(self, sqlite_session: Session) -> None: - end_user = EndUser( - tenant_id="00000000-0000-0000-0000-000000000001", - type=EndUserType.BROWSER, - session_id="session-1", - ) - sqlite_session.add(end_user) - sqlite_session.flush() - log = _log(CreatorUserRole.END_USER, created_by=end_user.id) - - result = log.created_by_end_user(session=sqlite_session) - - assert result is not None - assert result.id == end_user.id - - def test_returns_none_when_role_is_account(self, sqlite_session: Session) -> None: - end_user = EndUser( - tenant_id="00000000-0000-0000-0000-000000000001", - type=EndUserType.BROWSER, - session_id="session-1", - ) - sqlite_session.add(end_user) - sqlite_session.flush() - log = _log(CreatorUserRole.ACCOUNT, created_by=end_user.id) - - assert log.created_by_end_user(session=sqlite_session) is None diff --git a/api/tests/unit_tests/models/test_workflow_feature_compatibility.py b/api/tests/unit_tests/models/test_workflow_feature_compatibility.py deleted file mode 100644 index 66dd38b5f11..00000000000 --- a/api/tests/unit_tests/models/test_workflow_feature_compatibility.py +++ /dev/null @@ -1,59 +0,0 @@ -""" -This unit test verifies Workflow compatibility before and after supporting multiple file types. -""" - -import json - -from pydantic import JsonValue - -from models import Workflow - -OLD_VERSION_WORKFLOW_FEATURES: dict[str, JsonValue] = { - "file_upload": { - "image": { - "enabled": True, - "number_limits": 6, - "transfer_methods": ["remote_url", "local_file"], - } - }, - "opening_statement": "", - "retriever_resource": {"enabled": True}, - "sensitive_word_avoidance": {"enabled": False}, - "speech_to_text": {"enabled": False}, - "suggested_questions": [], - "suggested_questions_after_answer": {"enabled": False}, - "text_to_speech": {"enabled": False, "language": "", "voice": ""}, -} - -NEW_VERSION_WORKFLOW_FEATURES: dict[str, JsonValue] = { - "file_upload": { - "enabled": True, - "allowed_file_types": ["image"], - "allowed_file_extensions": [], - "allowed_file_upload_methods": ["remote_url", "local_file"], - "number_limits": 6, - }, - "opening_statement": "", - "retriever_resource": {"enabled": True}, - "sensitive_word_avoidance": {"enabled": False}, - "speech_to_text": {"enabled": False}, - "suggested_questions": [], - "suggested_questions_after_answer": {"enabled": False}, - "text_to_speech": {"enabled": False, "language": "", "voice": ""}, -} - - -def test_workflow_features() -> None: - workflow = Workflow( - tenant_id="", - app_id="", - type="", - version="", - graph="", - features=json.dumps(OLD_VERSION_WORKFLOW_FEATURES), - created_by="", - environment_variables=[], - conversation_variables=[], - ) - - assert workflow.features_dict == NEW_VERSION_WORKFLOW_FEATURES diff --git a/api/tests/unit_tests/models/test_workflow_models.py b/api/tests/unit_tests/models/test_workflow_models.py deleted file mode 100644 index da2804bc9ba..00000000000 --- a/api/tests/unit_tests/models/test_workflow_models.py +++ /dev/null @@ -1,1044 +0,0 @@ -""" -Comprehensive unit tests for Workflow models. - -This test suite covers: -- Workflow model validation -- WorkflowRun state transitions -- NodeExecution relationships -- Graph configuration validation -""" - -import json -from datetime import UTC, datetime -from uuid import uuid4 - -import pytest - -from graphon.enums import ( - BuiltinNodeTypes, - WorkflowExecutionStatus, - WorkflowNodeExecutionStatus, -) -from models.enums import CreatorUserRole, WorkflowRunTriggeredFrom -from models.workflow import ( - Workflow, - WorkflowNodeExecutionModel, - WorkflowNodeExecutionTriggeredFrom, - WorkflowRun, - WorkflowType, -) - - -class TestWorkflowModelValidation: - """Test suite for Workflow model validation and basic operations.""" - - def test_workflow_creation_with_required_fields(self) -> None: - """Test creating a workflow with all required fields.""" - # Arrange - tenant_id = str(uuid4()) - app_id = str(uuid4()) - created_by = str(uuid4()) - graph = json.dumps({"nodes": [], "edges": []}) - features = json.dumps({"file_upload": {"enabled": True}}) - - # Act - workflow = Workflow.new( - tenant_id=tenant_id, - app_id=app_id, - type=WorkflowType.WORKFLOW, - version="draft", - graph=graph, - features=features, - created_by=created_by, - environment_variables=[], - conversation_variables=[], - rag_pipeline_variables=[], - ) - - # Assert - assert workflow.tenant_id == tenant_id - assert workflow.app_id == app_id - assert workflow.type == WorkflowType.WORKFLOW - assert workflow.version == "draft" - assert workflow.graph == graph - assert workflow.created_by == created_by - assert workflow.created_at is not None - assert workflow.updated_at is not None - - def test_workflow_type_enum_values(self) -> None: - """Test WorkflowType enum values.""" - # Assert - assert WorkflowType.WORKFLOW == "workflow" - assert WorkflowType.CHAT.value == "chat" - assert WorkflowType.RAG_PIPELINE.value == "rag-pipeline" - - def test_workflow_type_value_of(self) -> None: - """Test WorkflowType.value_of method.""" - # Act & Assert - assert WorkflowType.value_of("workflow") == WorkflowType.WORKFLOW - assert WorkflowType.value_of("chat") == WorkflowType.CHAT - assert WorkflowType.value_of("rag-pipeline") == WorkflowType.RAG_PIPELINE - - with pytest.raises(ValueError, match="invalid workflow type value"): - WorkflowType.value_of("invalid_type") - - def test_workflow_graph_dict_property(self) -> None: - """Test graph_dict property parses JSON correctly.""" - # Arrange - graph_data = {"nodes": [{"id": "start", "type": "start"}], "edges": list[object]()} - workflow = Workflow.new( - tenant_id=str(uuid4()), - app_id=str(uuid4()), - type=WorkflowType.WORKFLOW, - version="draft", - graph=json.dumps(graph_data), - features="{}", - created_by=str(uuid4()), - environment_variables=[], - conversation_variables=[], - rag_pipeline_variables=[], - ) - - # Act - result = workflow.graph_dict - - # Assert - assert result == graph_data - assert "nodes" in result - assert len(result["nodes"]) == 1 - - def test_workflow_features_dict_property(self) -> None: - """Test features_dict property parses JSON correctly.""" - # Arrange - features_data = {"file_upload": {"enabled": True, "max_files": 5}} - workflow = Workflow.new( - tenant_id=str(uuid4()), - app_id=str(uuid4()), - type=WorkflowType.WORKFLOW, - version="draft", - graph="{}", - features=json.dumps(features_data), - created_by=str(uuid4()), - environment_variables=[], - conversation_variables=[], - rag_pipeline_variables=[], - ) - - # Act - result = workflow.features_dict - - # Assert - assert result == features_data - assert result["file_upload"]["enabled"] is True - assert result["file_upload"]["max_files"] == 5 - - def test_workflow_with_marked_name_and_comment(self) -> None: - """Test workflow creation with marked name and comment.""" - # Arrange & Act - workflow = Workflow.new( - tenant_id=str(uuid4()), - app_id=str(uuid4()), - type=WorkflowType.WORKFLOW, - version="v1.0", - graph="{}", - features="{}", - created_by=str(uuid4()), - environment_variables=[], - conversation_variables=[], - rag_pipeline_variables=[], - marked_name="Production Release", - marked_comment="Initial production version", - ) - - # Assert - assert workflow.marked_name == "Production Release" - assert workflow.marked_comment == "Initial production version" - - def test_workflow_version_draft_constant(self) -> None: - """Test VERSION_DRAFT constant.""" - # Assert - assert Workflow.VERSION_DRAFT == "draft" - - -class TestWorkflowRunStateTransitions: - """Test suite for WorkflowRun state transitions and lifecycle.""" - - def test_workflow_run_creation_with_required_fields(self) -> None: - """Test creating a workflow run with required fields.""" - # Arrange - tenant_id = str(uuid4()) - app_id = str(uuid4()) - workflow_id = str(uuid4()) - created_by = str(uuid4()) - - # Act - workflow_run = WorkflowRun( - tenant_id=tenant_id, - app_id=app_id, - workflow_id=workflow_id, - type=WorkflowType.WORKFLOW, - triggered_from=WorkflowRunTriggeredFrom.DEBUGGING, - version="draft", - status=WorkflowExecutionStatus.RUNNING, - created_by_role=CreatorUserRole.ACCOUNT, - created_by=created_by, - ) - - # Assert - assert workflow_run.tenant_id == tenant_id - assert workflow_run.app_id == app_id - assert workflow_run.workflow_id == workflow_id - assert workflow_run.type == WorkflowType.WORKFLOW - assert workflow_run.triggered_from == WorkflowRunTriggeredFrom.DEBUGGING - assert workflow_run.status == WorkflowExecutionStatus.RUNNING - assert workflow_run.created_by == created_by - - def test_workflow_run_state_transition_running_to_succeeded(self) -> None: - """Test state transition from running to succeeded.""" - # Arrange - workflow_run = WorkflowRun( - tenant_id=str(uuid4()), - app_id=str(uuid4()), - workflow_id=str(uuid4()), - type=WorkflowType.WORKFLOW, - triggered_from=WorkflowRunTriggeredFrom.APP_RUN, - version="v1.0", - status=WorkflowExecutionStatus.RUNNING, - created_by_role=CreatorUserRole.END_USER, - created_by=str(uuid4()), - ) - - # Act - workflow_run.status = WorkflowExecutionStatus.SUCCEEDED - workflow_run.finished_at = datetime.now(UTC) - workflow_run.elapsed_time = 2.5 - - # Assert - assert workflow_run.status == WorkflowExecutionStatus.SUCCEEDED - assert workflow_run.finished_at is not None - assert workflow_run.elapsed_time == 2.5 - - def test_workflow_run_state_transition_running_to_failed(self) -> None: - """Test state transition from running to failed with error.""" - # Arrange - workflow_run = WorkflowRun( - tenant_id=str(uuid4()), - app_id=str(uuid4()), - workflow_id=str(uuid4()), - type=WorkflowType.WORKFLOW, - triggered_from=WorkflowRunTriggeredFrom.APP_RUN, - version="v1.0", - status=WorkflowExecutionStatus.RUNNING, - created_by_role=CreatorUserRole.ACCOUNT, - created_by=str(uuid4()), - ) - - # Act - workflow_run.status = WorkflowExecutionStatus.FAILED - workflow_run.error = "Node execution failed: Invalid input" - workflow_run.finished_at = datetime.now(UTC) - - # Assert - assert workflow_run.status == WorkflowExecutionStatus.FAILED - assert workflow_run.error == "Node execution failed: Invalid input" - assert workflow_run.finished_at is not None - - def test_workflow_run_state_transition_running_to_stopped(self) -> None: - """Test state transition from running to stopped.""" - # Arrange - workflow_run = WorkflowRun( - tenant_id=str(uuid4()), - app_id=str(uuid4()), - workflow_id=str(uuid4()), - type=WorkflowType.WORKFLOW, - triggered_from=WorkflowRunTriggeredFrom.DEBUGGING, - version="draft", - status=WorkflowExecutionStatus.RUNNING, - created_by_role=CreatorUserRole.ACCOUNT, - created_by=str(uuid4()), - ) - - # Act - workflow_run.status = WorkflowExecutionStatus.STOPPED - workflow_run.finished_at = datetime.now(UTC) - - # Assert - assert workflow_run.status == WorkflowExecutionStatus.STOPPED - assert workflow_run.finished_at is not None - - def test_workflow_run_state_transition_running_to_paused(self) -> None: - """Test state transition from running to paused.""" - # Arrange - workflow_run = WorkflowRun( - tenant_id=str(uuid4()), - app_id=str(uuid4()), - workflow_id=str(uuid4()), - type=WorkflowType.WORKFLOW, - triggered_from=WorkflowRunTriggeredFrom.APP_RUN, - version="v1.0", - status=WorkflowExecutionStatus.RUNNING, - created_by_role=CreatorUserRole.END_USER, - created_by=str(uuid4()), - ) - - # Act - workflow_run.status = WorkflowExecutionStatus.PAUSED - - # Assert - assert workflow_run.status == WorkflowExecutionStatus.PAUSED - assert workflow_run.finished_at is None # Not finished when paused - - def test_workflow_run_state_transition_paused_to_running(self) -> None: - """Test state transition from paused back to running.""" - # Arrange - workflow_run = WorkflowRun( - tenant_id=str(uuid4()), - app_id=str(uuid4()), - workflow_id=str(uuid4()), - type=WorkflowType.WORKFLOW, - triggered_from=WorkflowRunTriggeredFrom.APP_RUN, - version="v1.0", - status=WorkflowExecutionStatus.PAUSED, - created_by_role=CreatorUserRole.ACCOUNT, - created_by=str(uuid4()), - ) - - # Act - workflow_run.status = WorkflowExecutionStatus.RUNNING - - # Assert - assert workflow_run.status == WorkflowExecutionStatus.RUNNING - - def test_workflow_run_with_partial_succeeded_status(self) -> None: - """Test workflow run with partial-succeeded status.""" - # Arrange & Act - workflow_run = WorkflowRun( - tenant_id=str(uuid4()), - app_id=str(uuid4()), - workflow_id=str(uuid4()), - type=WorkflowType.WORKFLOW, - triggered_from=WorkflowRunTriggeredFrom.APP_RUN, - version="v1.0", - status=WorkflowExecutionStatus.PARTIAL_SUCCEEDED, - created_by_role=CreatorUserRole.ACCOUNT, - created_by=str(uuid4()), - exceptions_count=2, - ) - - # Assert - assert workflow_run.status == WorkflowExecutionStatus.PARTIAL_SUCCEEDED - assert workflow_run.exceptions_count == 2 - - def test_workflow_run_with_inputs_and_outputs(self) -> None: - """Test workflow run with inputs and outputs as JSON.""" - # Arrange - inputs = {"query": "What is AI?", "context": "technology"} - outputs = {"answer": "AI is Artificial Intelligence", "confidence": 0.95} - - # Act - workflow_run = WorkflowRun( - tenant_id=str(uuid4()), - app_id=str(uuid4()), - workflow_id=str(uuid4()), - type=WorkflowType.WORKFLOW, - triggered_from=WorkflowRunTriggeredFrom.APP_RUN, - version="v1.0", - status=WorkflowExecutionStatus.SUCCEEDED, - created_by_role=CreatorUserRole.END_USER, - created_by=str(uuid4()), - inputs=json.dumps(inputs), - outputs=json.dumps(outputs), - ) - - # Assert - assert workflow_run.inputs_dict == inputs - assert workflow_run.outputs_dict == outputs - - def test_workflow_run_graph_dict_property(self) -> None: - """Test graph_dict property for workflow run.""" - # Arrange - graph = {"nodes": [{"id": "start", "type": "start"}], "edges": list[object]()} - workflow_run = WorkflowRun( - tenant_id=str(uuid4()), - app_id=str(uuid4()), - workflow_id=str(uuid4()), - type=WorkflowType.WORKFLOW, - triggered_from=WorkflowRunTriggeredFrom.DEBUGGING, - version="draft", - status=WorkflowExecutionStatus.RUNNING, - created_by_role=CreatorUserRole.ACCOUNT, - created_by=str(uuid4()), - graph=json.dumps(graph), - ) - - # Act - result = workflow_run.graph_dict - - # Assert - assert result == graph - assert "nodes" in result - - def test_workflow_run_to_dict_serialization(self) -> None: - """Test WorkflowRun to_dict method.""" - # Arrange - workflow_run_id = str(uuid4()) - tenant_id = str(uuid4()) - app_id = str(uuid4()) - workflow_id = str(uuid4()) - created_by = str(uuid4()) - - workflow_run = WorkflowRun( - tenant_id=tenant_id, - app_id=app_id, - workflow_id=workflow_id, - type=WorkflowType.WORKFLOW, - triggered_from=WorkflowRunTriggeredFrom.APP_RUN, - version="v1.0", - status=WorkflowExecutionStatus.SUCCEEDED, - created_by_role=CreatorUserRole.ACCOUNT, - created_by=created_by, - total_tokens=1500, - total_steps=5, - ) - workflow_run.id = workflow_run_id - - # Act - result = workflow_run.to_dict() - - # Assert - assert result["id"] == workflow_run_id - assert result["tenant_id"] == tenant_id - assert result["app_id"] == app_id - assert result["workflow_id"] == workflow_id - assert result["status"] == WorkflowExecutionStatus.SUCCEEDED - assert result["total_tokens"] == 1500 - assert result["total_steps"] == 5 - - def test_workflow_run_from_dict_deserialization(self) -> None: - """Test WorkflowRun from_dict method.""" - # Arrange - data = { - "id": str(uuid4()), - "tenant_id": str(uuid4()), - "app_id": str(uuid4()), - "workflow_id": str(uuid4()), - "type": WorkflowType.WORKFLOW, - "triggered_from": WorkflowRunTriggeredFrom.APP_RUN, - "version": "v1.0", - "graph": {"nodes": list[object](), "edges": list[object]()}, - "inputs": {"query": "test"}, - "status": WorkflowExecutionStatus.SUCCEEDED, - "outputs": {"result": "success"}, - "error": None, - "elapsed_time": 3.5, - "total_tokens": 2000, - "total_steps": 10, - "created_by_role": CreatorUserRole.ACCOUNT, - "created_by": str(uuid4()), - "created_at": datetime.now(UTC), - "finished_at": datetime.now(UTC), - "exceptions_count": 0, - } - - # Act - workflow_run = WorkflowRun.from_dict(data) - - # Assert - assert workflow_run.id == data["id"] - assert workflow_run.workflow_id == data["workflow_id"] - assert workflow_run.status == WorkflowExecutionStatus.SUCCEEDED - assert workflow_run.total_tokens == 2000 - - -class TestNodeExecutionRelationships: - """Test suite for WorkflowNodeExecutionModel relationships and data.""" - - def test_node_execution_creation_with_required_fields(self) -> None: - """Test creating a node execution with required fields.""" - # Arrange - tenant_id = str(uuid4()) - app_id = str(uuid4()) - workflow_id = str(uuid4()) - workflow_run_id = str(uuid4()) - created_by = str(uuid4()) - - # Act - node_execution = WorkflowNodeExecutionModel( - tenant_id=tenant_id, - app_id=app_id, - workflow_id=workflow_id, - triggered_from=WorkflowNodeExecutionTriggeredFrom.WORKFLOW_RUN, - workflow_run_id=workflow_run_id, - index=1, - node_id="start", - node_type=BuiltinNodeTypes.START, - title="Start Node", - status=WorkflowNodeExecutionStatus.SUCCEEDED, - created_by_role=CreatorUserRole.ACCOUNT, - created_by=created_by, - ) - - # Assert - assert node_execution.tenant_id == tenant_id - assert node_execution.app_id == app_id - assert node_execution.workflow_id == workflow_id - assert node_execution.workflow_run_id == workflow_run_id - assert node_execution.node_id == "start" - assert node_execution.node_type == BuiltinNodeTypes.START - assert node_execution.index == 1 - - def test_node_execution_with_predecessor_relationship(self) -> None: - """Test node execution with predecessor node relationship.""" - # Arrange - predecessor_node_id = "start" - current_node_id = "llm_1" - - # Act - node_execution = WorkflowNodeExecutionModel( - tenant_id=str(uuid4()), - app_id=str(uuid4()), - workflow_id=str(uuid4()), - triggered_from=WorkflowNodeExecutionTriggeredFrom.WORKFLOW_RUN, - workflow_run_id=str(uuid4()), - index=2, - predecessor_node_id=predecessor_node_id, - node_id=current_node_id, - node_type=BuiltinNodeTypes.LLM, - title="LLM Node", - status=WorkflowNodeExecutionStatus.RUNNING, - created_by_role=CreatorUserRole.ACCOUNT, - created_by=str(uuid4()), - ) - - # Assert - assert node_execution.predecessor_node_id == predecessor_node_id - assert node_execution.node_id == current_node_id - assert node_execution.index == 2 - - def test_node_execution_single_step_debugging(self) -> None: - """Test node execution for single-step debugging (no workflow_run_id).""" - # Arrange & Act - node_execution = WorkflowNodeExecutionModel( - tenant_id=str(uuid4()), - app_id=str(uuid4()), - workflow_id=str(uuid4()), - triggered_from=WorkflowNodeExecutionTriggeredFrom.SINGLE_STEP.value, - workflow_run_id=None, # Single-step has no workflow run - index=1, - node_id="llm_test", - node_type=BuiltinNodeTypes.LLM, - title="Test LLM", - status=WorkflowNodeExecutionStatus.SUCCEEDED, - created_by_role=CreatorUserRole.ACCOUNT, - created_by=str(uuid4()), - ) - - # Assert - assert node_execution.triggered_from == WorkflowNodeExecutionTriggeredFrom.SINGLE_STEP.value - assert node_execution.workflow_run_id is None - - def test_node_execution_with_inputs_outputs_process_data(self) -> None: - """Test node execution with inputs, outputs, and process_data.""" - # Arrange - inputs = {"query": "What is AI?", "temperature": 0.7} - outputs = {"answer": "AI is Artificial Intelligence"} - process_data = {"tokens_used": 150, "model": "gpt-4"} - - # Act - node_execution = WorkflowNodeExecutionModel( - tenant_id=str(uuid4()), - app_id=str(uuid4()), - workflow_id=str(uuid4()), - triggered_from=WorkflowNodeExecutionTriggeredFrom.WORKFLOW_RUN, - workflow_run_id=str(uuid4()), - index=1, - node_id="llm_1", - node_type=BuiltinNodeTypes.LLM, - title="LLM Node", - status=WorkflowNodeExecutionStatus.SUCCEEDED, - created_by_role=CreatorUserRole.ACCOUNT, - created_by=str(uuid4()), - inputs=json.dumps(inputs), - outputs=json.dumps(outputs), - process_data=json.dumps(process_data), - ) - - # Assert - assert node_execution.inputs_dict == inputs - assert node_execution.outputs_dict == outputs - assert node_execution.process_data_dict == process_data - - def test_node_execution_status_transitions(self) -> None: - """Test node execution status transitions.""" - # Arrange - node_execution = WorkflowNodeExecutionModel( - tenant_id=str(uuid4()), - app_id=str(uuid4()), - workflow_id=str(uuid4()), - triggered_from=WorkflowNodeExecutionTriggeredFrom.WORKFLOW_RUN, - workflow_run_id=str(uuid4()), - index=1, - node_id="code_1", - node_type=BuiltinNodeTypes.CODE, - title="Code Node", - status=WorkflowNodeExecutionStatus.RUNNING, - created_by_role=CreatorUserRole.ACCOUNT, - created_by=str(uuid4()), - ) - - # Act - transition to succeeded - node_execution.status = WorkflowNodeExecutionStatus.SUCCEEDED - node_execution.elapsed_time = 1.2 - node_execution.finished_at = datetime.now(UTC) - - # Assert - assert node_execution.status == WorkflowNodeExecutionStatus.SUCCEEDED - assert node_execution.elapsed_time == 1.2 - assert node_execution.finished_at is not None - - def test_node_execution_with_error(self) -> None: - """Test node execution with error status.""" - # Arrange - error_message = "Code execution failed: SyntaxError on line 5" - - # Act - node_execution = WorkflowNodeExecutionModel( - tenant_id=str(uuid4()), - app_id=str(uuid4()), - workflow_id=str(uuid4()), - triggered_from=WorkflowNodeExecutionTriggeredFrom.WORKFLOW_RUN, - workflow_run_id=str(uuid4()), - index=3, - node_id="code_1", - node_type=BuiltinNodeTypes.CODE, - title="Code Node", - status=WorkflowNodeExecutionStatus.FAILED, - error=error_message, - created_by_role=CreatorUserRole.ACCOUNT, - created_by=str(uuid4()), - ) - - # Assert - assert node_execution.status == WorkflowNodeExecutionStatus.FAILED - assert node_execution.error == error_message - - def test_node_execution_with_metadata(self) -> None: - """Test node execution with execution metadata.""" - # Arrange - metadata = { - "total_tokens": 500, - "total_price": 0.01, - "currency": "USD", - "tool_info": {"provider": "openai", "tool": "gpt-4"}, - } - - # Act - node_execution = WorkflowNodeExecutionModel( - tenant_id=str(uuid4()), - app_id=str(uuid4()), - workflow_id=str(uuid4()), - triggered_from=WorkflowNodeExecutionTriggeredFrom.WORKFLOW_RUN, - workflow_run_id=str(uuid4()), - index=1, - node_id="llm_1", - node_type=BuiltinNodeTypes.LLM, - title="LLM Node", - status=WorkflowNodeExecutionStatus.SUCCEEDED, - created_by_role=CreatorUserRole.ACCOUNT, - created_by=str(uuid4()), - execution_metadata=json.dumps(metadata), - ) - - # Assert - assert node_execution.execution_metadata_dict == metadata - assert node_execution.execution_metadata_dict["total_tokens"] == 500 - - def test_node_execution_metadata_dict_empty(self) -> None: - """Test execution_metadata_dict returns empty dict when metadata is None.""" - # Arrange - node_execution = WorkflowNodeExecutionModel( - tenant_id=str(uuid4()), - app_id=str(uuid4()), - workflow_id=str(uuid4()), - triggered_from=WorkflowNodeExecutionTriggeredFrom.WORKFLOW_RUN, - workflow_run_id=str(uuid4()), - index=1, - node_id="start", - node_type=BuiltinNodeTypes.START, - title="Start", - status=WorkflowNodeExecutionStatus.SUCCEEDED, - created_by_role=CreatorUserRole.ACCOUNT, - created_by=str(uuid4()), - execution_metadata=None, - ) - - # Act - result = node_execution.execution_metadata_dict - - # Assert - assert result == {} - - def test_node_execution_different_node_types(self) -> None: - """Test node execution with different node types.""" - # Test various node types - node_types = [ - (BuiltinNodeTypes.START, "Start Node"), - (BuiltinNodeTypes.LLM, "LLM Node"), - (BuiltinNodeTypes.CODE, "Code Node"), - (BuiltinNodeTypes.TOOL, "Tool Node"), - (BuiltinNodeTypes.IF_ELSE, "Conditional Node"), - (BuiltinNodeTypes.END, "End Node"), - ] - - for node_type, title in node_types: - # Act - node_execution = WorkflowNodeExecutionModel( - tenant_id=str(uuid4()), - app_id=str(uuid4()), - workflow_id=str(uuid4()), - triggered_from=WorkflowNodeExecutionTriggeredFrom.WORKFLOW_RUN, - workflow_run_id=str(uuid4()), - index=1, - node_id=f"{node_type}_1", - node_type=node_type, - title=title, - status=WorkflowNodeExecutionStatus.SUCCEEDED, - created_by_role=CreatorUserRole.ACCOUNT, - created_by=str(uuid4()), - ) - - # Assert - assert node_execution.node_type == node_type - assert node_execution.title == title - - -class TestGraphConfigurationValidation: - """Test suite for graph configuration validation.""" - - def test_workflow_graph_with_nodes_and_edges(self) -> None: - """Test workflow graph configuration with nodes and edges.""" - # Arrange - graph_config = { - "nodes": [ - {"id": "start", "type": "start", "data": {"title": "Start"}}, - {"id": "llm_1", "type": "llm", "data": {"title": "LLM Node", "model": "gpt-4"}}, - {"id": "end", "type": "end", "data": {"title": "End"}}, - ], - "edges": [ - {"source": "start", "target": "llm_1"}, - {"source": "llm_1", "target": "end"}, - ], - } - - # Act - workflow = Workflow.new( - tenant_id=str(uuid4()), - app_id=str(uuid4()), - type=WorkflowType.WORKFLOW, - version="draft", - graph=json.dumps(graph_config), - features="{}", - created_by=str(uuid4()), - environment_variables=[], - conversation_variables=[], - rag_pipeline_variables=[], - ) - - # Assert - graph_dict = workflow.graph_dict - assert len(graph_dict["nodes"]) == 3 - assert len(graph_dict["edges"]) == 2 - assert graph_dict["nodes"][0]["id"] == "start" - assert graph_dict["edges"][0]["source"] == "start" - assert graph_dict["edges"][0]["target"] == "llm_1" - - def test_workflow_graph_empty_configuration(self) -> None: - """Test workflow with empty graph configuration.""" - # Arrange - graph_config = {"nodes": list[object](), "edges": list[object]()} - - # Act - workflow = Workflow.new( - tenant_id=str(uuid4()), - app_id=str(uuid4()), - type=WorkflowType.WORKFLOW, - version="draft", - graph=json.dumps(graph_config), - features="{}", - created_by=str(uuid4()), - environment_variables=[], - conversation_variables=[], - rag_pipeline_variables=[], - ) - - # Assert - graph_dict = workflow.graph_dict - assert graph_dict["nodes"] == [] - assert graph_dict["edges"] == [] - - def test_workflow_graph_complex_node_data(self) -> None: - """Test workflow graph with complex node data structures.""" - # Arrange - graph_config = { - "nodes": [ - { - "id": "llm_1", - "type": "llm", - "data": { - "title": "Advanced LLM", - "model": {"provider": "openai", "name": "gpt-4", "mode": "chat"}, - "prompt_template": [ - {"role": "system", "text": "You are a helpful assistant"}, - {"role": "user", "text": "{{query}}"}, - ], - "model_parameters": {"temperature": 0.7, "max_tokens": 2000}, - }, - } - ], - "edges": list[object](), - } - - # Act - workflow = Workflow.new( - tenant_id=str(uuid4()), - app_id=str(uuid4()), - type=WorkflowType.WORKFLOW, - version="draft", - graph=json.dumps(graph_config), - features="{}", - created_by=str(uuid4()), - environment_variables=[], - conversation_variables=[], - rag_pipeline_variables=[], - ) - - # Assert - graph_dict = workflow.graph_dict - node_data = graph_dict["nodes"][0]["data"] - assert node_data["model"]["provider"] == "openai" - assert node_data["model_parameters"]["temperature"] == 0.7 - assert len(node_data["prompt_template"]) == 2 - - def test_workflow_run_graph_preservation(self) -> None: - """Test that WorkflowRun preserves graph configuration from Workflow.""" - # Arrange - original_graph = { - "nodes": [ - {"id": "start", "type": "start"}, - {"id": "end", "type": "end"}, - ], - "edges": [{"source": "start", "target": "end"}], - } - - # Act - workflow_run = WorkflowRun( - tenant_id=str(uuid4()), - app_id=str(uuid4()), - workflow_id=str(uuid4()), - type=WorkflowType.WORKFLOW, - triggered_from=WorkflowRunTriggeredFrom.APP_RUN, - version="v1.0", - status=WorkflowExecutionStatus.RUNNING, - created_by_role=CreatorUserRole.ACCOUNT, - created_by=str(uuid4()), - graph=json.dumps(original_graph), - ) - - # Assert - assert workflow_run.graph_dict == original_graph - assert len(workflow_run.graph_dict["nodes"]) == 2 - - def test_workflow_graph_with_conditional_branches(self) -> None: - """Test workflow graph with conditional branching (if-else).""" - # Arrange - graph_config = { - "nodes": [ - {"id": "start", "type": "start"}, - {"id": "if_else_1", "type": "if-else", "data": {"conditions": list[object]()}}, - {"id": "branch_true", "type": "llm"}, - {"id": "branch_false", "type": "code"}, - {"id": "end", "type": "end"}, - ], - "edges": [ - {"source": "start", "target": "if_else_1"}, - {"source": "if_else_1", "sourceHandle": "true", "target": "branch_true"}, - {"source": "if_else_1", "sourceHandle": "false", "target": "branch_false"}, - {"source": "branch_true", "target": "end"}, - {"source": "branch_false", "target": "end"}, - ], - } - - # Act - workflow = Workflow.new( - tenant_id=str(uuid4()), - app_id=str(uuid4()), - type=WorkflowType.WORKFLOW, - version="draft", - graph=json.dumps(graph_config), - features="{}", - created_by=str(uuid4()), - environment_variables=[], - conversation_variables=[], - rag_pipeline_variables=[], - ) - - # Assert - graph_dict = workflow.graph_dict - assert len(graph_dict["nodes"]) == 5 - assert len(graph_dict["edges"]) == 5 - # Verify conditional edges - conditional_edges = [e for e in graph_dict["edges"] if "sourceHandle" in e] - assert len(conditional_edges) == 2 - - def test_workflow_graph_with_loop_structure(self) -> None: - """Test workflow graph with loop/iteration structure.""" - # Arrange - graph_config = { - "nodes": [ - {"id": "start", "type": "start"}, - {"id": "iteration_1", "type": "iteration", "data": {"iterator": "items"}}, - {"id": "loop_body", "type": "llm"}, - {"id": "end", "type": "end"}, - ], - "edges": [ - {"source": "start", "target": "iteration_1"}, - {"source": "iteration_1", "target": "loop_body"}, - {"source": "loop_body", "target": "iteration_1"}, - {"source": "iteration_1", "target": "end"}, - ], - } - - # Act - workflow = Workflow.new( - tenant_id=str(uuid4()), - app_id=str(uuid4()), - type=WorkflowType.WORKFLOW, - version="draft", - graph=json.dumps(graph_config), - features="{}", - created_by=str(uuid4()), - environment_variables=[], - conversation_variables=[], - rag_pipeline_variables=[], - ) - - # Assert - graph_dict = workflow.graph_dict - iteration_node = next(n for n in graph_dict["nodes"] if n["type"] == "iteration") - assert iteration_node["data"]["iterator"] == "items" - - def test_workflow_graph_dict_with_null_graph(self) -> None: - """Test graph_dict property when graph is None.""" - # Arrange - workflow = Workflow.new( - tenant_id=str(uuid4()), - app_id=str(uuid4()), - type=WorkflowType.WORKFLOW, - version="draft", - graph=None, # pyrefly: ignore[bad-argument-type] - features="{}", - created_by=str(uuid4()), - environment_variables=[], - conversation_variables=[], - rag_pipeline_variables=[], - ) - - # Act - result = workflow.graph_dict - - # Assert - assert result == {} - - def test_workflow_run_inputs_dict_with_null_inputs(self) -> None: - """Test inputs_dict property when inputs is None.""" - # Arrange - workflow_run = WorkflowRun( - tenant_id=str(uuid4()), - app_id=str(uuid4()), - workflow_id=str(uuid4()), - type=WorkflowType.WORKFLOW, - triggered_from=WorkflowRunTriggeredFrom.APP_RUN, - version="v1.0", - status=WorkflowExecutionStatus.RUNNING, - created_by_role=CreatorUserRole.ACCOUNT, - created_by=str(uuid4()), - inputs=None, - ) - - # Act - result = workflow_run.inputs_dict - - # Assert - assert result == {} - - def test_workflow_run_outputs_dict_with_null_outputs(self) -> None: - """Test outputs_dict property when outputs is None.""" - # Arrange - workflow_run = WorkflowRun( - tenant_id=str(uuid4()), - app_id=str(uuid4()), - workflow_id=str(uuid4()), - type=WorkflowType.WORKFLOW, - triggered_from=WorkflowRunTriggeredFrom.APP_RUN, - version="v1.0", - status=WorkflowExecutionStatus.RUNNING, - created_by_role=CreatorUserRole.ACCOUNT, - created_by=str(uuid4()), - outputs=None, - ) - - # Act - result = workflow_run.outputs_dict - - # Assert - assert result == {} - - def test_node_execution_inputs_dict_with_null_inputs(self) -> None: - """Test node execution inputs_dict when inputs is None.""" - # Arrange - node_execution = WorkflowNodeExecutionModel( - tenant_id=str(uuid4()), - app_id=str(uuid4()), - workflow_id=str(uuid4()), - triggered_from=WorkflowNodeExecutionTriggeredFrom.WORKFLOW_RUN, - workflow_run_id=str(uuid4()), - index=1, - node_id="start", - node_type=BuiltinNodeTypes.START, - title="Start", - status=WorkflowNodeExecutionStatus.SUCCEEDED, - created_by_role=CreatorUserRole.ACCOUNT, - created_by=str(uuid4()), - inputs=None, - ) - - # Act - result = node_execution.inputs_dict - - # Assert - assert result is None - - def test_node_execution_outputs_dict_with_null_outputs(self) -> None: - """Test node execution outputs_dict when outputs is None.""" - # Arrange - node_execution = WorkflowNodeExecutionModel( - tenant_id=str(uuid4()), - app_id=str(uuid4()), - workflow_id=str(uuid4()), - triggered_from=WorkflowNodeExecutionTriggeredFrom.WORKFLOW_RUN, - workflow_run_id=str(uuid4()), - index=1, - node_id="start", - node_type=BuiltinNodeTypes.START, - title="Start", - status=WorkflowNodeExecutionStatus.SUCCEEDED, - created_by_role=CreatorUserRole.ACCOUNT, - created_by=str(uuid4()), - outputs=None, - ) - - # Act - result = node_execution.outputs_dict - - # Assert - assert result is None diff --git a/api/tests/unit_tests/models/test_workflow_node_execution_accessors.py b/api/tests/unit_tests/models/test_workflow_node_execution_accessors.py deleted file mode 100644 index cdb07111dce..00000000000 --- a/api/tests/unit_tests/models/test_workflow_node_execution_accessors.py +++ /dev/null @@ -1,110 +0,0 @@ -"""Regression coverage for ``WorkflowNodeExecutionModel`` account accessors. - -Ensures the ``@property``→session-parameter refactor preserves the role-based dispatch: -``created_by_account`` looks up an Account only when role is ACCOUNT; ``created_by_end_user`` -looks up an EndUser only when role is END_USER. Also covers the session-carrying response -source used by the console endpoints that validate ``WorkflowRunNodeExecutionResponse`` -with ``from_attributes=True``. -""" - -from uuid import uuid4 - -from sqlalchemy.orm import Session - -from fields.workflow_run_fields import node_execution_response_source -from models.account import Account -from models.enums import CreatorUserRole, EndUserType -from models.model import EndUser -from models.workflow import WorkflowNodeExecutionModel, WorkflowNodeExecutionTriggeredFrom - - -def _execution(role: CreatorUserRole, created_by: str) -> WorkflowNodeExecutionModel: - """Construct a WorkflowNodeExecutionModel without touching the database.""" - return WorkflowNodeExecutionModel( - tenant_id="00000000-0000-0000-0000-000000000001", - app_id="00000000-0000-0000-0000-000000000002", - workflow_id=str(uuid4()), - triggered_from=WorkflowNodeExecutionTriggeredFrom.WORKFLOW_RUN, - workflow_run_id=None, - index=1, - predecessor_node_id=None, - node_execution_id=None, - node_id="n1", - node_type="start", - title="Start", - inputs=None, - process_data=None, - outputs=None, - status="succeeded", - error=None, - elapsed_time=0.0, - execution_metadata=None, - created_by_role=role, - created_by=created_by, - ) - - -class TestCreatedByAccount: - def test_returns_account_lookup_when_role_is_account(self, sqlite_session: Session) -> None: - account = Account(name="Test Account", email="test@example.com") - sqlite_session.add(account) - sqlite_session.flush() - execution = _execution(CreatorUserRole.ACCOUNT, created_by=account.id) - - result = execution.created_by_account(session=sqlite_session) - - assert result is not None - assert result.id == account.id - - def test_returns_none_when_role_is_end_user(self, sqlite_session: Session) -> None: - account = Account(name="Test Account", email="test@example.com") - sqlite_session.add(account) - sqlite_session.flush() - execution = _execution(CreatorUserRole.END_USER, created_by=account.id) - - assert execution.created_by_account(session=sqlite_session) is None - - -class TestCreatedByEndUser: - def test_returns_end_user_lookup_when_role_is_end_user(self, sqlite_session: Session) -> None: - end_user = EndUser( - tenant_id="00000000-0000-0000-0000-000000000001", - type=EndUserType.BROWSER, - session_id="session-1", - ) - sqlite_session.add(end_user) - sqlite_session.flush() - execution = _execution(CreatorUserRole.END_USER, created_by=end_user.id) - - result = execution.created_by_end_user(session=sqlite_session) - - assert result is not None - assert result.id == end_user.id - - def test_returns_none_when_role_is_account(self, sqlite_session: Session) -> None: - end_user = EndUser( - tenant_id="00000000-0000-0000-0000-000000000001", - type=EndUserType.BROWSER, - session_id="session-1", - ) - sqlite_session.add(end_user) - sqlite_session.flush() - execution = _execution(CreatorUserRole.ACCOUNT, created_by=end_user.id) - - assert execution.created_by_end_user(session=sqlite_session) is None - - -class TestNodeExecutionResponseSource: - def test_accessors_resolve_via_wrapped_session_and_other_attributes_proxy(self, sqlite_session: Session) -> None: - account = Account(name="Test Account", email="test@example.com") - sqlite_session.add(account) - sqlite_session.flush() - execution = _execution(CreatorUserRole.ACCOUNT, created_by=account.id) - - source = node_execution_response_source(execution, session=sqlite_session) - - resolved = source.created_by_account - assert resolved is not None - assert resolved.id == account.id - assert source.created_by_end_user is None - assert source.node_id == "n1" diff --git a/api/tests/unit_tests/models/test_workflow_node_execution_offload.py b/api/tests/unit_tests/models/test_workflow_node_execution_offload.py deleted file mode 100644 index 7c97fef07d1..00000000000 --- a/api/tests/unit_tests/models/test_workflow_node_execution_offload.py +++ /dev/null @@ -1,204 +0,0 @@ -"""SQLite-backed tests for workflow node execution offload behavior.""" - -import json -from datetime import UTC, datetime -from unittest.mock import Mock - -import pytest -from sqlalchemy import select -from sqlalchemy.orm import Session - -from extensions.storage.storage_type import StorageType -from graphon.enums import WorkflowNodeExecutionStatus -from models.enums import CreatorUserRole, ExecutionOffLoadType -from models.model import UploadFile -from models.workflow import ( - WorkflowNodeExecutionModel, - WorkflowNodeExecutionOffload, - WorkflowNodeExecutionTriggeredFrom, -) - -TABLES = (WorkflowNodeExecutionModel, WorkflowNodeExecutionOffload, UploadFile) - - -def _persist_upload_file(session: Session, *, key: str = "offload/process-data.json") -> UploadFile: - upload_file = UploadFile( - tenant_id="tenant-1", - storage_type=StorageType.LOCAL, - key=key, - name="process-data.json", - size=18, - extension="json", - mime_type="application/json", - created_by_role=CreatorUserRole.ACCOUNT, - created_by="account-1", - created_at=datetime.now(UTC), - used=True, - ) - session.add(upload_file) - session.flush() - return upload_file - - -def _persist_execution( - session: Session, - *, - process_data: str = '{"test": "data"}', - offloads: tuple[tuple[ExecutionOffLoadType, str], ...] = (), -) -> WorkflowNodeExecutionModel: - execution = WorkflowNodeExecutionModel( - tenant_id="tenant-1", - app_id="app-1", - workflow_id="workflow-1", - triggered_from=WorkflowNodeExecutionTriggeredFrom.WORKFLOW_RUN, - workflow_run_id="run-1", - index=1, - predecessor_node_id=None, - node_execution_id="node-execution-1", - node_id="node-1", - node_type="code", - title="Code", - inputs='{"input": "value"}', - process_data=process_data, - outputs='{"output": "value"}', - status=WorkflowNodeExecutionStatus.SUCCEEDED, - error=None, - elapsed_time=0, - execution_metadata=None, - created_by_role=CreatorUserRole.ACCOUNT, - created_by="account-1", - finished_at=None, - ) - session.add(execution) - session.flush() - - session.add_all( - [ - WorkflowNodeExecutionOffload( - tenant_id=execution.tenant_id, - app_id=execution.app_id, - node_execution_id=execution.id, - type_=type_, - file_id=file_id, - ) - for type_, file_id in offloads - ] - ) - session.commit() - session.expunge_all() - - stmt = WorkflowNodeExecutionModel.preload_offload_data( - select(WorkflowNodeExecutionModel).where(WorkflowNodeExecutionModel.id == execution.id) - ) - loaded_execution = session.scalar(stmt) - assert loaded_execution is not None - return loaded_execution - - -@pytest.mark.parametrize("sqlite_session", [TABLES], indirect=True) -class TestWorkflowNodeExecutionModel: - """Exercise truncation flags and content restoration using persisted offload rows.""" - - def test_process_data_truncated_property_false_when_no_offload_data(self, sqlite_session: Session) -> None: - execution = _persist_execution(sqlite_session) - - assert execution.process_data_truncated is False - - def test_process_data_truncated_property_false_when_no_process_data_file(self, sqlite_session: Session) -> None: - execution = _persist_execution( - sqlite_session, - offloads=( - (ExecutionOffLoadType.INPUTS, "inputs-file"), - (ExecutionOffLoadType.OUTPUTS, "outputs-file"), - ), - ) - - assert execution.process_data_truncated is False - - def test_process_data_truncated_property_true_when_process_data_file_exists(self, sqlite_session: Session) -> None: - execution = _persist_execution( - sqlite_session, - offloads=((ExecutionOffLoadType.PROCESS_DATA, "process-data-file-id"),), - ) - - assert execution.process_data_truncated is True - - def test_load_full_process_data_with_no_offload_data(self, sqlite_session: Session) -> None: - execution = _persist_execution(sqlite_session) - storage = Mock() - - result = execution.load_full_process_data(sqlite_session, storage) - - assert result == {"test": "data"} - storage.load.assert_not_called() - - def test_load_full_process_data_with_no_file(self, sqlite_session: Session) -> None: - execution = _persist_execution( - sqlite_session, - offloads=((ExecutionOffLoadType.INPUTS, "inputs-file"),), - ) - storage = Mock() - - result = execution.load_full_process_data(sqlite_session, storage) - - assert result == {"test": "data"} - storage.load.assert_not_called() - - def test_load_full_process_data_with_file(self, sqlite_session: Session) -> None: - upload_file = _persist_upload_file(sqlite_session) - execution = _persist_execution( - sqlite_session, - process_data='{"truncated": true}', - offloads=((ExecutionOffLoadType.PROCESS_DATA, upload_file.id),), - ) - full_process_data = {"full": "data", "large_field": "x" * 10000} - storage = Mock() - storage.load.return_value = json.dumps(full_process_data) - - result = execution.load_full_process_data(sqlite_session, storage) - - assert result == full_process_data - storage.load.assert_called_once_with(upload_file.key) - - def test_consistency_with_inputs_outputs_truncation(self, sqlite_session: Session) -> None: - execution = _persist_execution( - sqlite_session, - offloads=( - (ExecutionOffLoadType.INPUTS, "inputs-file"), - (ExecutionOffLoadType.OUTPUTS, "outputs-file"), - (ExecutionOffLoadType.PROCESS_DATA, "process-data-file"), - ), - ) - - assert execution.inputs_truncated is True - assert execution.outputs_truncated is True - assert execution.process_data_truncated is True - - def test_mixed_truncation_states(self, sqlite_session: Session) -> None: - execution = _persist_execution( - sqlite_session, - offloads=((ExecutionOffLoadType.PROCESS_DATA, "process-data-file"),), - ) - - assert execution.inputs_truncated is False - assert execution.outputs_truncated is False - assert execution.process_data_truncated is True - - def test_preload_offload_data_and_files_loads_process_data_file(self, sqlite_session: Session) -> None: - upload_file = _persist_upload_file(sqlite_session) - execution = _persist_execution( - sqlite_session, - offloads=((ExecutionOffLoadType.PROCESS_DATA, upload_file.id),), - ) - execution_id = execution.id - sqlite_session.expunge_all() - - stmt = WorkflowNodeExecutionModel.preload_offload_data_and_files( - select(WorkflowNodeExecutionModel).where(WorkflowNodeExecutionModel.id == execution_id) - ) - preloaded_execution = sqlite_session.scalar(stmt) - - assert preloaded_execution is not None - assert len(preloaded_execution.offload_data) == 1 - assert preloaded_execution.offload_data[0].file is not None - assert preloaded_execution.offload_data[0].file.key == upload_file.key diff --git a/api/tests/unit_tests/models/test_workflow_run_accessors.py b/api/tests/unit_tests/models/test_workflow_run_accessors.py deleted file mode 100644 index 62c94d004ee..00000000000 --- a/api/tests/unit_tests/models/test_workflow_run_accessors.py +++ /dev/null @@ -1,114 +0,0 @@ -"""Regression coverage for ``WorkflowRun`` account accessors. - -Ensures the ``@property``→session-parameter refactor preserves the role-based dispatch: -``created_by_account`` looks up an Account only when role is ACCOUNT; ``created_by_end_user`` -looks up an EndUser only when role is END_USER. Also covers the session-carrying response -sources used by the console endpoints that validate the run detail / pagination responses -with ``from_attributes=True``. -""" - -from types import SimpleNamespace -from uuid import uuid4 - -from sqlalchemy.orm import Session - -from fields.workflow_run_fields import workflow_run_pagination_response_source, workflow_run_response_source -from models.account import Account -from models.enums import CreatorUserRole, EndUserType, WorkflowRunTriggeredFrom -from models.model import EndUser -from models.workflow import WorkflowRun - - -def _run(role: CreatorUserRole, created_by: str) -> WorkflowRun: - """Construct a WorkflowRun without touching the database.""" - return WorkflowRun( - id=str(uuid4()), - tenant_id="00000000-0000-0000-0000-000000000001", - app_id="00000000-0000-0000-0000-000000000002", - workflow_id=str(uuid4()), - type="workflow", - triggered_from=WorkflowRunTriggeredFrom.APP_RUN, - version="1.0", - status="succeeded", - created_by_role=role, - created_by=created_by, - ) - - -class TestCreatedByAccount: - def test_returns_account_lookup_when_role_is_account(self, sqlite_session: Session) -> None: - account = Account(name="Test Account", email="test@example.com") - sqlite_session.add(account) - sqlite_session.flush() - run = _run(CreatorUserRole.ACCOUNT, created_by=account.id) - - result = run.created_by_account(session=sqlite_session) - - assert result is not None - assert result.id == account.id - - def test_returns_none_when_role_is_end_user(self, sqlite_session: Session) -> None: - account = Account(name="Test Account", email="test@example.com") - sqlite_session.add(account) - sqlite_session.flush() - run = _run(CreatorUserRole.END_USER, created_by=account.id) - - assert run.created_by_account(session=sqlite_session) is None - - -class TestCreatedByEndUser: - def test_returns_end_user_lookup_when_role_is_end_user(self, sqlite_session: Session) -> None: - end_user = EndUser( - tenant_id="00000000-0000-0000-0000-000000000001", - type=EndUserType.BROWSER, - session_id="session-1", - ) - sqlite_session.add(end_user) - sqlite_session.flush() - run = _run(CreatorUserRole.END_USER, created_by=end_user.id) - - result = run.created_by_end_user(session=sqlite_session) - - assert result is not None - assert result.id == end_user.id - - def test_returns_none_when_role_is_account(self, sqlite_session: Session) -> None: - end_user = EndUser( - tenant_id="00000000-0000-0000-0000-000000000001", - type=EndUserType.BROWSER, - session_id="session-1", - ) - sqlite_session.add(end_user) - sqlite_session.flush() - run = _run(CreatorUserRole.ACCOUNT, created_by=end_user.id) - - assert run.created_by_end_user(session=sqlite_session) is None - - -class TestWorkflowRunResponseSources: - def test_detail_source_resolves_accessors_and_proxies_other_attributes(self, sqlite_session: Session) -> None: - account = Account(name="Test Account", email="test@example.com") - sqlite_session.add(account) - sqlite_session.flush() - run = _run(CreatorUserRole.ACCOUNT, created_by=account.id) - - source = workflow_run_response_source(run, session=sqlite_session) - - resolved = source.created_by_account - assert resolved is not None - assert resolved.id == account.id - assert source.created_by_end_user is None - assert source.version == "1.0" - - def test_pagination_source_wraps_each_run(self, sqlite_session: Session) -> None: - account = Account(name="Test Account", email="test@example.com") - sqlite_session.add(account) - sqlite_session.flush() - runs = [_run(CreatorUserRole.ACCOUNT, created_by=account.id) for _ in range(2)] - pagination = SimpleNamespace(data=runs, limit=20, has_more=False) - - payload = workflow_run_pagination_response_source(pagination, session=sqlite_session) - - assert payload["limit"] == 20 - assert payload["has_more"] is False - assert [item.created_by_account.id for item in payload["data"]] == [account.id, account.id] diff --git a/api/tests/unit_tests/models/test_file_input_compat.py b/api/tests/unit_tests/models/utils/test_file_input_compat.py similarity index 100% rename from api/tests/unit_tests/models/test_file_input_compat.py rename to api/tests/unit_tests/models/utils/test_file_input_compat.py diff --git a/api/tests/unit_tests/services/tracing/test_aliyun_configuration.py b/api/tests/unit_tests/providers/trace/trace-aliyun/src/dify_trace_aliyun/test_configuration.py similarity index 100% rename from api/tests/unit_tests/services/tracing/test_aliyun_configuration.py rename to api/tests/unit_tests/providers/trace/trace-aliyun/src/dify_trace_aliyun/test_configuration.py diff --git a/api/tests/unit_tests/providers/trace/trace-arize-phoenix/src/dify_trace_arize_phoenix/test_arize_configuration.py b/api/tests/unit_tests/providers/trace/trace-arize-phoenix/src/dify_trace_arize_phoenix/test_arize_configuration.py new file mode 100644 index 00000000000..63776e15227 --- /dev/null +++ b/api/tests/unit_tests/providers/trace/trace-arize-phoenix/src/dify_trace_arize_phoenix/test_arize_configuration.py @@ -0,0 +1,581 @@ +from collections import deque +from collections.abc import Iterator +from concurrent.futures import ThreadPoolExecutor +from threading import Event +from typing import Never + +import grpc # type: ignore[import-untyped] +import httpx +import pytest +from dify_trace_arize_phoenix import arize_configuration as configuration +from dify_trace_arize_phoenix.arize_configuration import ( + DESTINATION_KEY, + ArizeConfigurationError, + bind_destination, + project_url, + resolve_project_destination, + verify_configuration, +) +from dify_trace_arize_phoenix.arize_phoenix_trace import ArizePhoenixDataTrace +from dify_trace_arize_phoenix.config import ArizeConfig +from opentelemetry.context import _SUPPRESS_INSTRUMENTATION_KEY, get_value +from opentelemetry.proto.collector.trace.v1.trace_service_pb2 import ( + ExportTraceServiceRequest, + ExportTraceServiceResponse, +) + +from core.helper import ssrf_proxy +from core.ops.ops_trace_manager import OpsTraceManager +from services.app_tracing_config_gateway import OpsTraceManagerGateway +from services.app_tracing_config_service import AppTracingConfigVerificationFailedError + +pytest_plugins = ("tests.unit_tests.core.ops.test_ops_trace_manager",) + +"""Arize metadata reads are optional, bounded and pinned to the exact destination.""" + + +SPACE = "U3BhY2U6dGVzdA==" +ORG = "T3JnOnRlc3Q=" +PROJECT = "UHJvamVjdDovKz0=" +VALUES = {"endpoint": "https://otlp.arize.com", "space_id": SPACE, "project": "exact", "api_key": "synthetic-key"} +DESTINATION = { + "endpoint": VALUES["endpoint"], + "space_id": SPACE, + "project": "exact", + "organization_id": ORG, + "project_id": PROJECT, +} +URL = "https://app.arize.com/organizations/T3JnOnRlc3Q%3D/spaces/U3BhY2U6dGVzdA%3D%3D/projects/UHJvamVjdDovKz0%3D" + + +def response(body: object, status: int = 200) -> httpx.Response: + return httpx.Response(status, json=body, request=httpx.Request("GET", "https://api.arize.com/v2/projects")) + + +def space() -> httpx.Response: + return response({"id": SPACE, "organization_id": ORG}) + + +def page(*projects: dict[str, str], more: bool = False, cursor: str | None = None) -> httpx.Response: + return response({"projects": list(projects), "pagination": {"has_more": more, "next_cursor": cursor}}) + + +def project(name: str = "exact", project_id: str = PROJECT) -> dict[str, str]: + return {"id": project_id, "name": name, "space_id": SPACE} + + +class TestArizeConfigurationCases: + @staticmethod + @pytest.fixture + def responses() -> deque[httpx.Response | Exception]: + return deque() + + @staticmethod + @pytest.fixture + def calls( + monkeypatch: pytest.MonkeyPatch, responses: deque[httpx.Response | Exception] + ) -> list[tuple[str, dict[str, object]]]: + calls: list[tuple[str, dict[str, object]]] = [] + + def get(url: str, **options: object) -> httpx.Response: + calls.append((url, options)) + assert responses, "Unplanned metadata request" + item = responses.popleft() + if isinstance(item, Exception): + raise item + return item + + monkeypatch.setattr(ssrf_proxy, "get", get) + return calls + + @staticmethod + def test_exact_scoped_project_and_opaque_ids_use_official_route( + calls: list[tuple[str, dict[str, object]]], responses: deque[httpx.Response | Exception] + ) -> None: + received = [space(), page(project("exact-substring"), more=True, cursor="next"), page(project())] + responses.extend(received) + assert resolve_project_destination(VALUES) == DESTINATION + assert ( + project_url( + { + **VALUES, + DESTINATION_KEY: bind_destination(VALUES, DESTINATION), + "project_url": "https://untrusted.test", + } + ) + == URL + ) + assert len(calls) == 3 + assert calls[0][0] == "https://api.arize.com/v2/spaces/U3BhY2U6dGVzdA%3D%3D" + assert calls[1][1]["params"] == {"space_id": SPACE, "name": "exact", "limit": 100} + assert calls[2][1]["params"] == {"space_id": SPACE, "name": "exact", "limit": 100, "cursor": "next"} + for url, options in calls: + assert url.startswith("https://api.arize.com/v2/") + assert options["headers"] == { + "Authorization": "Bearer synthetic-key", + "Accept": "application/json", + "Accept-Encoding": "identity", + "Cookie": "", + } + assert options["auth"] is None + assert options["follow_redirects"] is False + assert options["max_retries"] == 0 + assert options["stream_response"] is True + timeout = options["timeout"] + assert isinstance(timeout, float) + assert 0 < timeout <= 5 + assert all(item.is_closed for item in received) + + @staticmethod + @pytest.mark.parametrize("status", [301, 401, 403, 404, 429, 500]) + def test_unreadable_metadata_has_no_link_and_closes_response( + calls: list[tuple[str, dict[str, object]]], responses: deque[httpx.Response | Exception], status: int + ) -> None: + denied = response({"error": "private-response"}, status) + responses.append(denied) + assert resolve_project_destination(VALUES) is None + assert len(calls) == 1 + assert denied.is_closed + + @staticmethod + @pytest.mark.parametrize( + "body", [{"id": SPACE}, {"id": "other", "organization_id": ORG}, {"id": SPACE, "organization_id": "../bad"}] + ) + def test_missing_or_mismatched_space_metadata_is_not_a_project( + calls: list[tuple[str, dict[str, object]]], responses: deque[httpx.Response | Exception], body: dict[str, str] + ) -> None: + responses.append(response(body)) + assert resolve_project_destination(VALUES) is None + assert len(calls) == 1 + + @staticmethod + @pytest.mark.parametrize( + "body", + [ + {"projects": [project("EXACT")], "pagination": {"has_more": False}}, + {"projects": [project(), project(project_id="other")], "pagination": {"has_more": False}}, + {"projects": [{**project(), "space_id": "other"}], "pagination": {"has_more": False}}, + {"projects": [{**project(), "id": "../bad"}], "pagination": {"has_more": False}}, + {"projects": [project()] * 101, "pagination": {"has_more": False}}, + {"projects": [project()], "pagination": {"has_more": True}}, + {"projects": [project()], "pagination": {"has_more": "false"}}, + {"projects": [project()]}, + {"data": [project()], "pagination": {"has_more": False}}, + ], + ) + def test_incomplete_ambiguous_or_non_exact_projects_do_not_fabricate_links( + calls: list[tuple[str, dict[str, object]]], responses: deque[httpx.Response | Exception], body: object + ) -> None: + responses.extend([space(), response(body)]) + assert resolve_project_destination(VALUES) is None + assert len(calls) == 2 + + @staticmethod + @pytest.mark.parametrize("repeat", [False, True]) + def test_pagination_is_bounded_and_never_accepts_a_partial_match( + calls: list[tuple[str, dict[str, object]]], responses: deque[httpx.Response | Exception], repeat: bool + ) -> None: + responses.append(space()) + responses.extend(page(project(), more=True, cursor="same" if repeat else str(index)) for index in range(3)) + assert resolve_project_destination(VALUES) is None + assert len(calls) == (3 if repeat else 4) + + @staticmethod + @pytest.mark.parametrize( + "item", + [ + httpx.ReadTimeout("synthetic-timeout"), + httpx.Response(200, content=b"not-json"), + httpx.Response(200, json=[]), + httpx.Response(200, content=b" " * (1024 * 1024 + 1)), + httpx.Response(200, json={}, headers={"content-encoding": "unsupported"}), + ], + ) + def test_transport_and_bounded_decode_failures_remain_optional( + calls: list[tuple[str, dict[str, object]]], + responses: deque[httpx.Response | Exception], + item: httpx.Response | Exception, + ) -> None: + responses.append(item) + assert resolve_project_destination(VALUES) is None + assert len(calls) == 1 + if isinstance(item, httpx.Response): + assert item.is_closed + + @staticmethod + def test_all_pages_share_one_deadline(monkeypatch: pytest.MonkeyPatch) -> None: + clock = [100.0] + calls: list[float] = [] + received = [space(), page(project(), more=True, cursor="next")] + + def get(_url: str, **options: object) -> httpx.Response: + timeout = options["timeout"] + assert isinstance(timeout, float) + calls.append(timeout) + clock[0] += 2.5 + return received[len(calls) - 1] + + monkeypatch.setattr(configuration.time, "monotonic", lambda: clock[0]) + monkeypatch.setattr(ssrf_proxy, "get", get) + assert resolve_project_destination(VALUES) is None + assert calls == [5.0, 2.5] + assert all(item.is_closed for item in received) + + @staticmethod + def test_slow_response_chunks_stop_at_the_shared_deadline(monkeypatch: pytest.MonkeyPatch) -> None: + clock = [100.0] + chunks: list[int] = [] + received = httpx.Response(200) + + def stream() -> Iterator[bytes]: + for index in range(10): + clock[0] += 2.0 + chunks.append(index) + yield b" " + + monkeypatch.setattr(configuration.time, "monotonic", lambda: clock[0]) + monkeypatch.setattr(received, "iter_bytes", stream) + monkeypatch.setattr(ssrf_proxy, "get", lambda _url, **_options: received) + assert resolve_project_destination(VALUES) is None + assert chunks == [0, 1, 2] + assert received.is_closed + + @staticmethod + @pytest.mark.parametrize( + "change", + [ + {"endpoint": "https://custom.test"}, + {"endpoint": "https://otlp.eu-west-1a.arize.com"}, + {"endpoint": "http://otlp.arize.com"}, + {"endpoint": "https://user:secret@otlp.arize.com"}, + {"endpoint": "https://otlp.arize.com:444"}, + {"space_id": "../bad"}, + {"api_key": None}, + {"api_key": ""}, + {"api_key": "******"}, + {"api_key": "key\nheader"}, + ], + ) + def test_unsupported_destination_or_credentials_do_not_send_to_an_assumed_host( + calls: list[tuple[str, dict[str, object]]], change: dict[str, str | None] + ) -> None: + assert resolve_project_destination({**VALUES, **change}) is None + assert calls == [] + + @staticmethod + @pytest.mark.parametrize( + "change", + [ + {"endpoint": "https://otlp.arize.com:443"}, + {"space_id": "other"}, + {"project": "later"}, + {DESTINATION_KEY: {**DESTINATION, "organization_id": "bad?query"}}, + {DESTINATION_KEY: {**DESTINATION, "project_id": "../bad"}}, + {DESTINATION_KEY: None}, + ], + ) + def test_get_formatter_rejects_stale_or_forged_metadata_without_io( + monkeypatch: pytest.MonkeyPatch, change: dict[str, object] + ) -> None: + def forbidden(*_args: object, **_kwargs: object) -> Never: + raise AssertionError("GET must not call the vendor") + + monkeypatch.setattr(ssrf_proxy, "get", forbidden) + assert project_url({**VALUES, DESTINATION_KEY: bind_destination(VALUES, DESTINATION), **change}) is None + + +"""A real gRPC collector must acknowledge Arize credentials before config persistence.""" + + +class TestArizeConfigurationAckCases: + @staticmethod + @pytest.fixture + def collector() -> Iterator[tuple[str, list[tuple[bytes, dict[str, str]]], dict[str, object]]]: + calls: list[tuple[bytes, dict[str, str]]] = [] + state: dict[str, object] = {"status": grpc.StatusCode.OK} + release = Event() + state["release"] = release + + def export(request: ExportTraceServiceRequest, context: grpc.ServicerContext) -> ExportTraceServiceResponse: + headers = { + key: value.decode() if isinstance(value, bytes) else value + for key, value in context.invocation_metadata() + } + calls.append((request.SerializeToString(), headers)) + if headers.get("api_key") != "synthetic-good" or headers.get("space_id") != "synthetic-space": + context.abort(grpc.StatusCode.PERMISSION_DENIED, "private-vendor-details") + code = state["status"] + if isinstance(code, grpc.StatusCode) and code is not grpc.StatusCode.OK: + context.abort(code, "private-vendor-details") + if state.get("delay"): + release.wait(2) + result = ExportTraceServiceResponse() + if state.get("partial"): + result.partial_success.rejected_spans = 1 + if state.get("message"): + result.partial_success.error_message = "private-vendor-details" + return result + + def dispatch(*args: object, **kwargs: object) -> ExportTraceServiceResponse: + assert kwargs == {} + request, context = args + assert isinstance(request, ExportTraceServiceRequest) + assert isinstance(context, grpc.ServicerContext) + return export(request, context) + + handler = grpc.unary_unary_rpc_method_handler( + dispatch, + request_deserializer=ExportTraceServiceRequest.FromString, + response_serializer=ExportTraceServiceResponse.SerializeToString, + ) + with ThreadPoolExecutor(max_workers=2) as pool: + server = grpc.server(pool) + server.add_generic_rpc_handlers( + ( + grpc.method_handlers_generic_handler( + "opentelemetry.proto.collector.trace.v1.TraceService", {"Export": handler} + ), + ) + ) + port = server.add_insecure_port("127.0.0.1:0") + assert port > 0 + server.start() + try: + yield f"http://127.0.0.1:{port}", calls, state + finally: + release.set() + server.stop(0).wait(2) + + @staticmethod + def test_rejected_real_collector_ack_cannot_be_saved( + workspace: str, + collector: tuple[str, list[tuple[bytes, dict[str, str]]], dict[str, object]], + ) -> None: + endpoint, calls, _ = collector + with pytest.raises(AppTracingConfigVerificationFailedError): + OpsTraceManagerGateway().prepare_new_config( + workspace_id=workspace, + tracing_provider="arize", + tracing_config={ + "endpoint": endpoint, + "project": "synthetic-project", + "space_id": "synthetic-space", + "api_key": "synthetic-invalid", + }, + ) + assert len(calls) == 1 + + @staticmethod + def test_synchronous_ack_is_resource_only_and_has_no_sdk_or_environment_side_effects( + collector: tuple[str, list[tuple[bytes, dict[str, str]]], dict[str, object]], + monkeypatch: pytest.MonkeyPatch, + ) -> None: + endpoint, calls, _ = collector + attempted: list[bool] = [] + + def forbidden(*_args: object, **_kwargs: object) -> Never: + attempted.append(True) + raise AssertionError("Arize verification must not construct the SDK tracer") + + monkeypatch.setattr(ArizePhoenixDataTrace, "__init__", forbidden) + monkeypatch.setenv("OTEL_EXPORTER_OTLP_HEADERS", "api_key=ambient-invalid,space_id=ambient-other") + verify_configuration( + ArizeConfig( + endpoint=endpoint, api_key="synthetic-good", space_id="synthetic-space", project="synthetic-project" + ) + ) + assert attempted == [] + assert len(calls) == 1 + payload = ExportTraceServiceRequest.FromString(calls[0][0]) + assert len(payload.resource_spans) == 1 + assert {a.key: a.value.string_value for a in payload.resource_spans[0].resource.attributes} == { + "model_id": "synthetic-project", + "openinference.project.name": "synthetic-project", + } + assert len(payload.resource_spans[0].scope_spans) == 0 + assert calls[0][1]["api_key"] == "synthetic-good" + assert calls[0][1]["space_id"] == "synthetic-space" + assert calls[0][1]["authorization"] == "Bearer synthetic-good" + assert "ambient" not in str(calls) + + @staticmethod + @pytest.mark.parametrize( + ("code", "expected"), + [ + (grpc.StatusCode.UNAUTHENTICATED, "authentication_failed"), + (grpc.StatusCode.PERMISSION_DENIED, "permission_denied"), + (grpc.StatusCode.UNAVAILABLE, "connection_failed"), + (grpc.StatusCode.RESOURCE_EXHAUSTED, "connection_failed"), + ], + ) + def test_real_rpc_rejection_uses_only_safe_fixed_diagnostics( + collector: tuple[str, list[tuple[bytes, dict[str, str]]], dict[str, object]], + code: grpc.StatusCode, + expected: str, + caplog: pytest.LogCaptureFixture, + ) -> None: + endpoint, calls, state = collector + state["status"] = code + with pytest.raises(ArizeConfigurationError) as caught: + verify_configuration(ArizeConfig(endpoint=endpoint, api_key="synthetic-good", space_id="synthetic-space")) + assert caught.value.code == expected + assert caught.value.__suppress_context__ + assert "synthetic-good" not in str(caught.value) + assert "private-vendor-details" not in str(caught.value) + assert "private-vendor-details" not in caplog.text + assert len(calls) == 1 + + @staticmethod + @pytest.mark.parametrize("partial", ["partial", "message"]) + def test_partial_or_error_ack_is_not_success( + collector: tuple[str, list[tuple[bytes, dict[str, str]]], dict[str, object]], + partial: str, + ) -> None: + endpoint, calls, state = collector + state[partial] = True + with pytest.raises(ArizeConfigurationError, match="unexpected_response"): + verify_configuration(ArizeConfig(endpoint=endpoint, api_key="synthetic-good", space_id="synthetic-space")) + assert len(calls) == 1 + + @staticmethod + def test_deadline_is_finite_and_does_not_retry( + collector: tuple[str, list[tuple[bytes, dict[str, str]]], dict[str, object]], + monkeypatch: pytest.MonkeyPatch, + ) -> None: + endpoint, calls, state = collector + state["delay"] = True + monkeypatch.setattr(configuration, "_VERIFY_TIMEOUT_SECONDS", 0.05) + with pytest.raises(ArizeConfigurationError, match="connection_failed"): + verify_configuration(ArizeConfig(endpoint=endpoint, api_key="synthetic-good", space_id="synthetic-space")) + assert len(calls) == 1 + + @staticmethod + @pytest.mark.parametrize("reject", [False, True]) + def test_tls_channel_is_explicit_suppressed_and_closed_on_all_outcomes( + collector: tuple[str, list[tuple[bytes, dict[str, str]]], dict[str, object]], + monkeypatch: pytest.MonkeyPatch, + reject: bool, + ) -> None: + endpoint, _, state = collector + if reject: + state["status"] = grpc.StatusCode.PERMISSION_DENIED + original = grpc.insecure_channel + options_seen: list[tuple[str, object, object]] = [] + closed: list[bool] = [] + + def secure(target: str, credentials: grpc.ChannelCredentials, options: object) -> grpc.Channel: + assert isinstance(credentials, grpc.ChannelCredentials) + options_seen.append((target, options, get_value(_SUPPRESS_INSTRUMENTATION_KEY))) + channel = original(endpoint.removeprefix("http://")) + close = channel.close + + def close_recording() -> None: + closed.append(True) + close() + + monkeypatch.setattr(channel, "close", close_recording) + return channel + + monkeypatch.setattr(grpc, "secure_channel", secure) + before = get_value(_SUPPRESS_INSTRUMENTATION_KEY) + config = ArizeConfig(api_key="synthetic-good", space_id="synthetic-space") + if reject: + with pytest.raises(ArizeConfigurationError): + verify_configuration(config) + else: + verify_configuration(config) + assert options_seen == [ + ("otlp.arize.com:443", (("grpc.enable_retries", 0), ("grpc.max_receive_message_length", 1024 * 1024)), True) + ] + assert closed == [True] + assert get_value(_SUPPRESS_INSTRUMENTATION_KEY) == before + + @staticmethod + @pytest.mark.parametrize( + ("field", "value"), + [ + ("api_key", None), + ("api_key", ""), + ("api_key", "******"), + ("api_key", "key\nheader"), + ("space_id", None), + ("space_id", " "), + ("space_id", "******"), + ("endpoint", "https://user:secret@otlp.arize.com"), + ("endpoint", "https://otlp.arize.com:0"), + ], + ) + def test_invalid_auth_values_never_construct_a_channel( + monkeypatch: pytest.MonkeyPatch, + field: str, + value: str | None, + ) -> None: + attempted: list[bool] = [] + + def forbidden(*_args: object, **_kwargs: object) -> Never: + attempted.append(True) + raise AssertionError("Invalid configuration must be rejected before transport construction") + + monkeypatch.setattr(grpc, "secure_channel", forbidden) + monkeypatch.setattr(grpc, "insecure_channel", forbidden) + config = ArizeConfig.model_validate({"api_key": "synthetic-good", "space_id": "synthetic-space", field: value}) + with pytest.raises(ArizeConfigurationError, match="invalid_configuration"): + verify_configuration(config) + assert attempted == [] + + @staticmethod + @pytest.mark.parametrize("accepted", [False, True]) + def test_legacy_api_check_also_requires_the_real_ack( + collector: tuple[str, list[tuple[bytes, dict[str, str]]], dict[str, object]], + accepted: bool, + ) -> None: + endpoint, calls, _ = collector + instance = ArizePhoenixDataTrace( + ArizeConfig( + endpoint=endpoint, + api_key="synthetic-good" if accepted else "synthetic-invalid", + space_id="synthetic-space", + ) + ) + try: + if accepted: + assert instance.api_check() + else: + with pytest.raises(ArizeConfigurationError, match="permission_denied"): + instance.api_check() + finally: + instance.processor.shutdown() + assert len(calls) == 1 + assert len(ExportTraceServiceRequest.FromString(calls[0][0]).resource_spans[0].scope_spans) == 0 + + @staticmethod + def test_real_ingest_ack_without_rest_read_permission_can_still_save( + workspace: str, + collector: tuple[str, list[tuple[bytes, dict[str, str]]], dict[str, object]], + monkeypatch: pytest.MonkeyPatch, + ) -> None: + endpoint, rpc_calls, _ = collector + original = grpc.insecure_channel + reads: list[str] = [] + + def secure(_target: str, _credentials: grpc.ChannelCredentials, **_options: object) -> grpc.Channel: + return original(endpoint.removeprefix("http://")) + + def get(url: str, **_options: object) -> httpx.Response: + reads.append(url) + return httpx.Response(403, json={"message": "private-vendor-details"}) + + monkeypatch.setattr(grpc, "secure_channel", secure) + monkeypatch.setattr(ssrf_proxy, "get", get) + saved = OpsTraceManagerGateway().prepare_new_config( + workspace_id=workspace, + tracing_provider="arize", + tracing_config={"api_key": "synthetic-good", "space_id": "synthetic-space", "project": "not-created"}, + ) + assert len(rpc_calls) == len(reads) == 1 + assert reads[0] == "https://api.arize.com/v2/spaces/synthetic-space" + assert "project_url" not in saved + assert "arize_project_destination" not in saved + assert OpsTraceManager.decrypt_tracing_config(workspace, "arize", saved)["api_key"] == "synthetic-good" + payload = ExportTraceServiceRequest.FromString(rpc_calls[0][0]) + assert len(payload.resource_spans[0].scope_spans) == 0 diff --git a/api/tests/unit_tests/services/tracing/test_phoenix_configuration.py b/api/tests/unit_tests/providers/trace/trace-arize-phoenix/src/dify_trace_arize_phoenix/test_configuration.py similarity index 100% rename from api/tests/unit_tests/services/tracing/test_phoenix_configuration.py rename to api/tests/unit_tests/providers/trace/trace-arize-phoenix/src/dify_trace_arize_phoenix/test_configuration.py diff --git a/api/tests/unit_tests/providers/trace/unified_trace/test_arize_phoenix.py b/api/tests/unit_tests/providers/trace/trace-arize-phoenix/src/dify_trace_arize_phoenix/test_unified_trace.py similarity index 100% rename from api/tests/unit_tests/providers/trace/unified_trace/test_arize_phoenix.py rename to api/tests/unit_tests/providers/trace/trace-arize-phoenix/src/dify_trace_arize_phoenix/test_unified_trace.py diff --git a/api/tests/unit_tests/providers/trace/unified_trace/test_langsmith.py b/api/tests/unit_tests/providers/trace/trace-langsmith/src/dify_trace_langsmith/test_unified_trace.py similarity index 100% rename from api/tests/unit_tests/providers/trace/unified_trace/test_langsmith.py rename to api/tests/unit_tests/providers/trace/trace-langsmith/src/dify_trace_langsmith/test_unified_trace.py diff --git a/api/tests/unit_tests/services/tracing/test_databricks_configuration.py b/api/tests/unit_tests/providers/trace/trace-mlflow/src/dify_trace_mlflow/test_configuration.py similarity index 100% rename from api/tests/unit_tests/services/tracing/test_databricks_configuration.py rename to api/tests/unit_tests/providers/trace/trace-mlflow/src/dify_trace_mlflow/test_configuration.py diff --git a/api/tests/unit_tests/services/tracing/test_mlflow_legacy_destination.py b/api/tests/unit_tests/providers/trace/trace-mlflow/src/dify_trace_mlflow/test_mlflow_trace.py similarity index 100% rename from api/tests/unit_tests/services/tracing/test_mlflow_legacy_destination.py rename to api/tests/unit_tests/providers/trace/trace-mlflow/src/dify_trace_mlflow/test_mlflow_trace.py diff --git a/api/tests/unit_tests/services/tracing/test_mlflow_configuration.py b/api/tests/unit_tests/providers/trace/trace-mlflow/src/dify_trace_mlflow/test_oss_configuration.py similarity index 97% rename from api/tests/unit_tests/services/tracing/test_mlflow_configuration.py rename to api/tests/unit_tests/providers/trace/trace-mlflow/src/dify_trace_mlflow/test_oss_configuration.py index d8f1bd5e5ab..82ffb642a46 100644 --- a/api/tests/unit_tests/services/tracing/test_mlflow_configuration.py +++ b/api/tests/unit_tests/providers/trace/trace-mlflow/src/dify_trace_mlflow/test_oss_configuration.py @@ -1,6 +1,7 @@ """OSS MLflow verification reads one experiment with isolated Basic credentials.""" import base64 +import importlib as _test_importlib import os from collections.abc import Callable, Iterator from concurrent.futures import ThreadPoolExecutor @@ -21,7 +22,14 @@ from dify_trace_mlflow.oss_configuration import ( from pydantic import ValidationError from core.helper import ssrf_proxy -from tests.unit_tests.services.tracing.test_databricks_configuration import ExperimentAPI, response + +ExperimentAPI = _test_importlib.import_module( + "tests.unit_tests.providers.trace.trace-mlflow.src.dify_trace_mlflow.test_configuration" +).ExperimentAPI + +response = _test_importlib.import_module( + "tests.unit_tests.providers.trace.trace-mlflow.src.dify_trace_mlflow.test_configuration" +).response @pytest.fixture(autouse=True) diff --git a/api/tests/unit_tests/services/tracing/test_opik_configuration.py b/api/tests/unit_tests/providers/trace/trace-opik/src/dify_trace_opik/test_configuration.py similarity index 100% rename from api/tests/unit_tests/services/tracing/test_opik_configuration.py rename to api/tests/unit_tests/providers/trace/trace-opik/src/dify_trace_opik/test_configuration.py diff --git a/api/tests/unit_tests/services/tracing/test_opik_legacy_identifiers.py b/api/tests/unit_tests/providers/trace/trace-opik/src/dify_trace_opik/test_opik_trace.py similarity index 99% rename from api/tests/unit_tests/services/tracing/test_opik_legacy_identifiers.py rename to api/tests/unit_tests/providers/trace/trace-opik/src/dify_trace_opik/test_opik_trace.py index af48f2234e9..f11027ac0f2 100644 --- a/api/tests/unit_tests/services/tracing/test_opik_legacy_identifiers.py +++ b/api/tests/unit_tests/providers/trace/trace-opik/src/dify_trace_opik/test_opik_trace.py @@ -27,7 +27,7 @@ from extensions.logstore.repositories.logstore_api_workflow_node_execution_repos _dict_to_workflow_node_execution_model, ) from models.account import TenantAccountJoin, TenantAccountRole -from tests.unit_tests.model_factories import make_account, make_app, make_tenant +from tests.helpers.unit.model_factories import make_account, make_app, make_tenant START = datetime(2026, 10, 16, 15, 18, 10, 792586) WORKFLOW_ID = "11111111-2222-4333-8444-555555555555" diff --git a/api/tests/unit_tests/services/tracing/test_tencent_configuration.py b/api/tests/unit_tests/providers/trace/trace-tencent/src/dify_trace_tencent/test_configuration.py similarity index 100% rename from api/tests/unit_tests/services/tracing/test_tencent_configuration.py rename to api/tests/unit_tests/providers/trace/trace-tencent/src/dify_trace_tencent/test_configuration.py diff --git a/api/tests/unit_tests/services/tracing/test_tencent_legacy_sessions.py b/api/tests/unit_tests/providers/trace/trace-tencent/src/dify_trace_tencent/test_utils.py similarity index 100% rename from api/tests/unit_tests/services/tracing/test_tencent_legacy_sessions.py rename to api/tests/unit_tests/providers/trace/trace-tencent/src/dify_trace_tencent/test_utils.py diff --git a/api/tests/unit_tests/services/tracing/test_weave_configuration.py b/api/tests/unit_tests/providers/trace/trace-weave/src/dify_trace_weave/test_configuration.py similarity index 100% rename from api/tests/unit_tests/services/tracing/test_weave_configuration.py rename to api/tests/unit_tests/providers/trace/trace-weave/src/dify_trace_weave/test_configuration.py diff --git a/api/tests/unit_tests/pyrefly.toml b/api/tests/unit_tests/pyrefly.toml index f2a80712878..06aff8b7366 100644 --- a/api/tests/unit_tests/pyrefly.toml +++ b/api/tests/unit_tests/pyrefly.toml @@ -1,5 +1,5 @@ preset = "strict" -project-includes = ["."] +project-includes = [".", "../helpers/unit", "../helpers/workflow", "../e2e_tests/contracts"] search-path = ["../.."] python-platform = "linux" python-version = "3.13.0" @@ -8,46 +8,40 @@ min-severity = "warn" # Existing strict-mode debt. Remove a file when bringing it under strict checking. project-excludes = [ - "commands/test_archive_workflow_runs.py", - "commands/test_data_migration_commands.py", - "commands/test_generate_swagger_markdown_docs.py", - "commands/test_generate_swagger_specs.py", - "commands/test_lint_response_contracts.py", - "commands/test_upgrade_db.py", - "repositories/test_app_definition_query_repository_agent.py", + "commands/test_retention.py", + "commands/test_data_migration.py", + "dev/test_generate_swagger_markdown_docs.py", + "dev/test_generate_swagger_specs.py", + "dev/test_lint_response_contracts.py", + "commands/test_system.py", + "repositories/test_app_definition_query_repository.py", "controllers/common/test_errors.py", "controllers/common/test_file_response.py", "controllers/common/test_schema.py", "controllers/console/app/test_agent_config_inspector.py", "controllers/console/app/test_agent_manage_guard.py", - "controllers/console/app/test_annotation_api.py", - "controllers/console/app/test_annotation_security.py", - "controllers/console/app/test_app_apis.py", - "controllers/console/app/test_app_import_api.py", - "controllers/console/app/test_app_response_models.py", + "controllers/console/app/test_annotation.py", + "controllers/console/app/test_app.py", + "controllers/console/app/test_app_import.py", "controllers/console/app/test_audio.py", - "controllers/console/app/test_conversation_api.py", - "controllers/console/app/test_create_app_payload.py", - "controllers/console/app/test_description_validation.py", - "controllers/console/app/test_generator_api.py", - "controllers/console/app/test_generator_api_missing.py", - "controllers/console/app/test_mcp_server_response.py", - "controllers/console/app/test_message_api.py", - "controllers/console/app/test_model_config_api.py", - "controllers/console/app/test_ops_trace_api.py", - "controllers/console/app/test_statistic_api.py", + "controllers/console/app/test_conversation.py", + "libs/test_validators.py", + "controllers/console/app/test_generator.py", + "controllers/console/app/test_mcp_server.py", + "controllers/console/app/test_message.py", + "controllers/console/app/test_model_config.py", + "controllers/console/app/test_ops_trace.py", + "controllers/console/app/test_statistic.py", "controllers/console/app/test_workflow.py", - "controllers/console/app/test_workflow_app_log_api.py", - "controllers/console/app/test_workflow_convert_api.py", + "controllers/console/app/test_workflow_app_log.py", "controllers/console/app/test_workflow_node_output_inspector.py", - "controllers/console/app/test_workflow_pause_details_api.py", - "controllers/console/app/test_workflow_run_api.py", + "controllers/console/app/test_workflow_run.py", "controllers/console/app/test_wraps.py", - "controllers/console/app/workflow_draft_variables_test.py", - "controllers/console/auth/test_account_activation.py", + "fields/test_workflow_draft_variable_fields.py", + "controllers/console/auth/test_activate.py", "controllers/console/auth/test_email_register.py", "controllers/console/auth/test_forgot_password.py", - "controllers/console/auth/test_login_logout.py", + "controllers/console/auth/test_login.py", "controllers/console/auth/test_oauth.py", "controllers/console/auth/test_oauth_timezone.py", "controllers/console/billing/test_billing.py", @@ -59,7 +53,6 @@ project-excludes = [ "controllers/console/datasets/test_data_source.py", "controllers/console/datasets/test_datasets.py", "controllers/console/datasets/test_datasets_document.py", - "controllers/console/datasets/test_datasets_document_download.py", "controllers/console/datasets/test_datasets_segments.py", "controllers/console/datasets/test_external.py", "controllers/console/datasets/test_hit_testing.py", @@ -77,14 +70,14 @@ project-excludes = [ "controllers/console/snippets/test_snippet_workflow.py", "controllers/console/snippets/test_snippet_workflow_draft_variable.py", "controllers/console/tag/test_tags.py", - "controllers/console/test_document_detail_api_data_source_info.py", - "controllers/console/test_fastopenapi_setup.py", + "services/knowledge/test_dataset_read_service.py", + "controllers/console/test_setup.py", "controllers/console/test_files.py", "controllers/console/test_human_input_form.py", "controllers/console/test_knowledge_fs_proxy.py", "controllers/console/test_remote_files.py", "controllers/console/test_wraps.py", - "controllers/console/workspace/test_accounts.py", + "controllers/console/workspace/test_account.py", "controllers/console/workspace/test_endpoint.py", "controllers/console/workspace/test_members.py", "controllers/console/workspace/test_model_providers.py", @@ -95,55 +88,47 @@ project-excludes = [ "controllers/inner_api/app/test_dsl.py", "controllers/inner_api/plugin/test_agent_config.py", "controllers/inner_api/plugin/test_plugin.py", - "controllers/inner_api/plugin/test_plugin_wraps.py", - "controllers/inner_api/test_knowledge_retrieval.py", + "controllers/inner_api/plugin/test_wraps.py", + "controllers/inner_api/knowledge/test_retrieval.py", "controllers/inner_api/test_mail.py", "controllers/inner_api/test_runtime_credentials.py", "controllers/mcp/test_mcp.py", "controllers/openapi/test_account.py", - "controllers/openapi/test_app_run_rate_limit.py", - "controllers/openapi/test_app_run_streaming.py", - "controllers/openapi/test_contract.py", - "controllers/openapi/test_device_approve_deny.py", - "controllers/openapi/test_error_contract.py", + "controllers/openapi/test_app_run.py", + "controllers/openapi/test__contract.py", + "controllers/openapi/test_oauth_device.py", + "controllers/openapi/test__errors.py", "controllers/openapi/test_human_input_form.py", - "controllers/openapi/test_models.py", - "controllers/openapi/test_workflow_events_openapi.py", - "controllers/openapi/test_workspaces_members.py", + "controllers/openapi/test__models.py", + "controllers/openapi/test_workflow_events.py", + "controllers/openapi/test_workspaces.py", "controllers/service_api/app/test_annotation.py", "controllers/service_api/app/test_app.py", "controllers/service_api/app/test_audio.py", - "controllers/service_api/app/test_chat_request_payload.py", "controllers/service_api/app/test_completion.py", "controllers/service_api/app/test_conversation.py", "controllers/service_api/app/test_file.py", "controllers/service_api/app/test_file_preview.py", - "controllers/service_api/app/test_hitl_service_api.py", + "controllers/service_api/app/test_workflow_events.py", "controllers/service_api/app/test_human_input_form.py", "controllers/service_api/app/test_message.py", "controllers/service_api/app/test_workflow.py", - "controllers/service_api/app/test_workflow_events.py", "controllers/service_api/conftest.py", "controllers/service_api/dataset/rag_pipeline/test_rag_pipeline_workflow.py", - "controllers/service_api/dataset/test_dataset_segment.py", + "controllers/service_api/dataset/test_segment.py", "controllers/service_api/dataset/test_document.py", - "controllers/service_api/dataset/test_rag_pipeline_file_upload_serialization.py", - "controllers/service_api/dataset/test_rag_pipeline_route_registration.py", "controllers/service_api/test_index.py", - "controllers/service_api/test_trace_session_id_parsing.py", "controllers/service_api/test_wraps.py", - "controllers/test_conversation_rename_payload.py", - "controllers/test_swagger.py", + "controllers/service_api/test___init__.py", "controllers/trigger/test_trigger.py", "controllers/trigger/test_webhook.py", "controllers/web/test_app.py", "controllers/web/test_completion.py", "controllers/web/test_human_input_file_upload.py", "controllers/web/test_human_input_form.py", - "controllers/web/test_message_endpoints.py", - "controllers/web/test_message_list.py", - "controllers/web/test_pydantic_models.py", - "controllers/web/test_web_login.py", + "controllers/web/test_message.py", + "controllers/web/test_conversation.py", + "controllers/web/test_login.py", "controllers/web/test_workflow.py", "controllers/web/test_wraps.py", "core/agent/conftest.py", @@ -154,98 +139,80 @@ project-excludes = [ "core/agent/test_cot_completion_agent_runner.py", "core/agent/test_fc_agent_runner.py", "core/agent/test_plugin_entities.py", - "core/app/app_config/common/test_sensitive_word_avoidance_manager.py", - "core/app/app_config/easy_ui_based_app/test_model_config_converter.py", - "core/app/app_config/easy_ui_based_app/test_variables_manager.py", + "core/app/app_config/common/sensitive_word_avoidance/test_manager.py", + "core/app/app_config/easy_ui_based_app/model_config/test_converter.py", + "core/app/app_config/easy_ui_based_app/variables/test_manager.py", "core/app/app_config/features/file_upload/test_manager.py", - "core/app/app_config/features/test_additional_feature_managers.py", + "core/app/app_config/features/suggested_questions_after_answer/test_manager.py", "core/app/apps/advanced_chat/test_app_config_manager.py", "core/app/apps/advanced_chat/test_app_generator.py", - "core/app/apps/advanced_chat/test_app_runner_conversation_variables.py", - "core/app/apps/advanced_chat/test_app_runner_input_moderation.py", + "core/app/apps/advanced_chat/test_app_runner.py", "core/app/apps/advanced_chat/test_generate_response_converter.py", "core/app/apps/advanced_chat/test_generate_task_pipeline.py", - "core/app/apps/advanced_chat/test_generate_task_pipeline_core.py", "core/app/apps/agent_app/test_app_config_manager.py", "core/app/apps/agent_app/test_app_generator.py", "core/app/apps/agent_app/test_app_runner.py", - "core/app/apps/agent_app/test_input_guards.py", - "core/app/apps/agent_app/test_resolve_agent.py", "core/app/apps/agent_app/test_runtime_request_builder.py", "core/app/apps/agent_app/test_session_store.py", - "core/app/apps/agent_chat/test_agent_chat_app_config_manager.py", - "core/app/apps/agent_chat/test_agent_chat_generate_response_converter.py", + "core/app/apps/agent_chat/test_app_config_manager.py", + "core/app/apps/agent_chat/test_generate_response_converter.py", "core/app/apps/chat/test_app_config_manager.py", - "core/app/apps/chat/test_app_generator_and_runner.py", - "core/app/apps/chat/test_base_app_runner_multimodal.py", + "core/app/apps/chat/test_app_generator.py", + "core/app/apps/test_base_app_runner.py", "core/app/apps/chat/test_generate_response_converter.py", "core/app/apps/common/test_graph_runtime_state_support.py", "core/app/apps/common/test_workflow_response_converter.py", - "core/app/apps/common/test_workflow_response_converter_human_input.py", - "core/app/apps/common/test_workflow_response_converter_resumption.py", - "core/app/apps/common/test_workflow_response_converter_truncation.py", - "core/app/apps/completion/test_completion_app_config_manager.py", - "core/app/apps/completion/test_completion_completion_app_generator.py", - "core/app/apps/completion/test_completion_generate_response_converter.py", + "core/app/apps/completion/test_app_config_manager.py", + "core/app/apps/completion/test_app_generator.py", + "core/app/apps/completion/test_generate_response_converter.py", "core/app/apps/pipeline/test_pipeline_config_manager.py", - "core/app/apps/pipeline/test_pipeline_generate_response_converter.py", + "core/app/apps/pipeline/test_generate_response_converter.py", "core/app/apps/pipeline/test_pipeline_generator.py", "core/app/apps/pipeline/test_pipeline_queue_manager.py", "core/app/apps/pipeline/test_pipeline_runner.py", - "core/app/apps/test_advanced_chat_app_generator.py", "core/app/apps/test_base_app_generate_response_converter.py", "core/app/apps/test_base_app_generator.py", "core/app/apps/test_base_app_queue_manager.py", - "core/app/apps/test_base_app_runner.py", "core/app/apps/test_exc.py", "core/app/apps/test_message_based_app_generator.py", "core/app/apps/test_message_based_app_queue_manager.py", "core/app/apps/test_message_generator.py", - "core/app/apps/test_pause_resume.py", "core/app/apps/test_streaming_utils.py", - "core/app/apps/test_trace_session_id_generate_extras.py", - "core/app/apps/test_workflow_app_generator.py", - "core/app/apps/test_workflow_app_runner_core.py", - "core/app/apps/test_workflow_app_runner_notifications.py", - "core/app/apps/test_workflow_app_runner_single_node.py", - "core/app/apps/test_workflow_pause_events.py", + "core/helper/test_trace_id_helper.py", + "core/app/apps/workflow/test_app_generator.py", + "core/app/apps/test_workflow_app_runner.py", "core/app/apps/workflow/test_active_workflow_tasks.py", "core/app/apps/workflow/test_app_config_manager.py", - "core/app/apps/workflow/test_app_generator_extra.py", "core/app/apps/workflow/test_app_queue_manager.py", "core/app/apps/workflow/test_command_channels.py", "core/app/apps/workflow/test_errors.py", "core/app/apps/workflow/test_generate_response_converter.py", "core/app/apps/workflow/test_generate_task_pipeline.py", - "core/app/apps/workflow/test_generate_task_pipeline_core.py", "core/app/entities/test_app_invoke_entities.py", "core/app/features/rate_limiting/test_rate_limit.py", - "core/app/features/test_annotation_reply.py", - "core/app/features/test_hosting_moderation.py", + "core/app/features/annotation_reply/test_annotation_reply.py", + "core/app/features/hosting_moderation/test_hosting_moderation.py", "core/app/layers/test_conversation_variable_persist_layer.py", "core/app/layers/test_pause_state_persist_layer.py", "core/app/layers/test_timeslice_layer.py", "core/app/layers/test_trigger_post_layer.py", "core/app/task_pipeline/test_based_generate_task_pipeline.py", "core/app/task_pipeline/test_easy_ui_based_generate_task_pipeline.py", - "core/app/task_pipeline/test_easy_ui_based_generate_task_pipeline_core.py", "core/app/task_pipeline/test_exc.py", - "core/app/task_pipeline/test_message_cycle_manager_optimization.py", - "core/app/test_easy_ui_model_config_manager.py", - "core/app/test_invoke_from.py", - "core/app/test_llm_quota.py", + "core/app/task_pipeline/test_message_cycle_manager.py", + "core/app/app_config/easy_ui_based_app/model_config/test_manager.py", + "core/app/llm/test_quota.py", "core/app/workflow/layers/test_persistence.py", - "core/app/workflow/layers/test_persistence_inspector_publish.py", + "services/test_workflow_persistence.py", "core/app/workflow/test_file_runtime.py", - "core/app/workflow/test_node_factory.py", - "core/app/workflow/test_observability_layer_extra.py", - "core/app/workflow/test_persistence_layer.py", + "core/workflow/test_node_factory.py", + "core/app/workflow/layers/test_observability.py", "core/callback_handler/test_agent_tool_callback_handler.py", "core/callback_handler/test_index_tool_callback_handler.py", "core/callback_handler/test_workflow_tool_callback_handler.py", "core/datasource/entities/test_datasource_entities.py", - "core/entities/test_entities_mcp_provider.py", - "core/entities/test_entities_provider_configuration.py", + "core/entities/test_mcp_provider.py", + "core/entities/test_provider_configuration.py", "core/extension/test_api_based_extension_requestor.py", "core/external_data_tool/test_external_data_fetch.py", "core/external_data_tool/test_factory.py", @@ -254,56 +221,49 @@ project-excludes = [ "core/llm_generator/output_parser/test_rule_config_generator.py", "core/llm_generator/output_parser/test_structured_output.py", "core/llm_generator/test_llm_generator.py", - "core/llm_generator/test_llm_generator_missing.py", "core/logging/test_filters.py", - "core/logging/test_trace_helpers.py", "core/mcp/auth/test_auth_flow.py", - "core/mcp/client/test_sse.py", - "core/mcp/client/test_streamable_http.py", + "core/mcp/client/test_sse_client.py", + "core/mcp/client/test_streamable_client.py", "core/mcp/server/test_streamable_http.py", "core/mcp/session/test_base_session.py", "core/mcp/session/test_client_session.py", "core/mcp/test_mcp_client.py", "core/mcp/test_utils.py", "core/memory/test_token_buffer_memory.py", - "core/model_runtime/test_model_provider_factory.py", + "core/plugin/impl/test_model_runtime_factory.py", "core/moderation/test_input_moderation.py", "core/ops/test_base_trace_instance.py", - "core/ops/test_config_entity.py", - "core/ops/test_lookup_helpers.py", + "core/ops/entities/test_config_entity.py", "core/ops/test_ops_trace_manager.py", - "core/ops/test_trace_queue_manager.py", - "core/ops/test_trace_session_metadata.py", + "core/ops/entities/test_trace_entity.py", "core/ops/test_utils.py", - "core/plugin/impl/test_agent_client.py", - "core/plugin/impl/test_asset_manager.py", - "core/plugin/impl/test_base_client_impl.py", - "core/plugin/impl/test_datasource_manager.py", - "core/plugin/impl/test_debugging_client.py", - "core/plugin/impl/test_endpoint_client_impl.py", - "core/plugin/impl/test_exc_impl.py", - "core/plugin/impl/test_model_client.py", - "core/plugin/impl/test_model_runtime_factory.py", - "core/plugin/impl/test_oauth_handler.py", - "core/plugin/impl/test_tool_manager.py", - "core/plugin/impl/test_trigger_client.py", - "core/plugin/test_backwards_invocation_app.py", - "core/plugin/test_backwards_invocation_model.py", - "core/plugin/test_endpoint_client.py", - "core/plugin/test_model_runtime_adapter.py", - "core/plugin/test_plugin_entities.py", - "core/plugin/test_plugin_manager.py", - "core/plugin/test_plugin_runtime.py", + "core/plugin/impl/test_agent.py", + "core/plugin/impl/test_asset.py", + "core/plugin/impl/test_base.py", + "core/plugin/impl/test_datasource.py", + "core/plugin/impl/test_debugging.py", + "core/plugin/impl/test_endpoint.py", + "core/plugin/impl/test_exc.py", + "core/plugin/impl/test_model.py", + "core/plugin/impl/test_oauth.py", + "core/plugin/impl/test_tool.py", + "core/plugin/impl/test_trigger.py", + "core/plugin/backwards_invocation/test_app.py", + "core/plugin/backwards_invocation/test_model.py", + "core/plugin/impl/test_model_runtime.py", + "core/plugin/entities/test_plugin_daemon.py", + "core/plugin/impl/test_plugin.py", "core/plugin/utils/test_chunk_merger.py", "core/prompt/test_advanced_prompt_transform.py", - "core/prompt/test_extract_thread_messages.py", + "core/prompt/utils/test_extract_thread_messages.py", "core/rag/cleaner/test_clean_processor.py", "core/rag/data_post_processor/test_data_post_processor.py", "core/rag/datasource/keyword/jieba/test_jieba.py", "core/rag/datasource/keyword/jieba/test_jieba_keyword_table_handler.py", - "core/rag/datasource/test_datasource_retrieval.py", + "core/rag/datasource/test_retrieval_service.py", "core/rag/datasource/vdb/test_vector_factory.py", - "core/rag/embedding/test_embedding_service.py", + "core/rag/embedding/test_cached_embedding.py", "core/rag/extractor/blob/test_blob.py", "core/rag/extractor/test_csv_extractor.py", "core/rag/extractor/test_excel_extractor.py", @@ -312,76 +272,68 @@ project-excludes = [ "core/rag/extractor/test_html_extractor.py", "core/rag/extractor/test_markdown_extractor.py", "core/rag/extractor/test_word_extractor.py", - "core/rag/extractor/watercrawl/test_watercrawl.py", - "core/rag/indexing/processor/test_paragraph_index_processor.py", - "core/rag/indexing/processor/test_qa_index_processor.py", - "core/rag/indexing/test_index_processor.py", - "core/rag/indexing/test_index_processor_base.py", - "core/rag/rerank/test_reranker.py", + "core/rag/extractor/watercrawl/test_client.py", + "core/rag/index_processor/processor/test_paragraph_index_processor.py", + "core/rag/index_processor/processor/test_qa_index_processor.py", + "core/rag/index_processor/test_index_processor.py", + "core/rag/index_processor/test_index_processor_base.py", + "core/rag/rerank/test_rerank_model.py", "core/rag/retrieval/test_dataset_retrieval.py", - "core/rag/retrieval/test_dataset_retrieval_methods.py", "core/rag/splitter/test_text_splitter.py", "core/repositories/test_celery_workflow_execution_repository.py", "core/repositories/test_celery_workflow_node_execution_repository.py", "core/repositories/test_factory.py", - "core/repositories/test_human_input_form_repository_impl.py", "core/repositories/test_human_input_repository.py", "core/repositories/test_sqlalchemy_workflow_execution_repository.py", "core/repositories/test_sqlalchemy_workflow_node_execution_repository.py", - "core/repositories/test_workflow_node_execution_conflict_handling.py", - "core/repositories/test_workflow_node_execution_truncation.py", - "core/telemetry/test_facade.py", - "core/telemetry/test_gateway_integration.py", - "core/test_file.py", + "core/telemetry/test___init__.py", + "core/telemetry/test_gateway.py", + "../e2e_tests/contracts/core/test_file.py", "core/test_model_manager.py", "core/test_provider_manager.py", - "core/test_trigger_debug_event_selectors.py", + "core/trigger/debug/test_event_selectors.py", "core/tools/entities/test_api_entities.py", - "core/tools/test_base_tool.py", - "core/tools/test_builtin_tool_base.py", - "core/tools/test_builtin_tool_provider.py", - "core/tools/test_builtin_tools_extra.py", - "core/tools/test_custom_tool.py", - "core/tools/test_custom_tool_provider.py", - "core/tools/test_dataset_retriever_tool.py", - "core/tools/test_mcp_tool.py", - "core/tools/test_mcp_tool_provider.py", - "core/tools/test_plugin_tool.py", - "core/tools/test_plugin_tool_provider.py", + "core/tools/__base/test_tool.py", + "core/tools/builtin_tool/test_tool.py", + "core/tools/builtin_tool/test_provider.py", + "core/tools/builtin_tool/providers/time/tools/test_localtime_to_timestamp.py", + "core/tools/custom_tool/test_tool.py", + "core/tools/custom_tool/test_provider.py", + "core/tools/utils/test_dataset_retriever_tool.py", + "core/tools/mcp_tool/test_tool.py", + "core/tools/mcp_tool/test_provider.py", + "core/tools/plugin_tool/test_tool.py", + "core/tools/plugin_tool/test_provider.py", "core/tools/test_tool_engine.py", - "core/tools/test_tool_entities.py", + "core/tools/entities/test_tool_entities.py", "core/tools/test_tool_file_manager.py", "core/tools/test_tool_label_manager.py", "core/tools/test_tool_manager.py", - "core/tools/test_tool_parameter_type.py", - "core/tools/test_tool_provider_controller.py", + "core/tools/__base/test_tool_provider.py", "core/tools/utils/test_configuration.py", "core/tools/utils/test_message_transformer.py", - "core/tools/utils/test_misc_utils_extra.py", + "core/tools/utils/dataset_retriever/test_dataset_multi_retriever_tool.py", "core/tools/utils/test_model_invocation_utils.py", "core/tools/utils/test_parser.py", - "core/tools/utils/test_tool_engine_serialization.py", "core/tools/utils/test_web_reader_tool.py", "core/tools/utils/test_workflow_configuration_sync.py", "core/tools/workflow_as_tool/test_provider.py", "core/tools/workflow_as_tool/test_tool.py", - "core/trigger/utils/test_utils_encryption.py", - "core/trigger/utils/test_utils_endpoint.py", - "core/trigger/utils/test_utils_locks.py", - "core/workflow/context/test_execution_context.py", - "core/workflow/context/test_flask_app_context.py", - "core/workflow/entities/test_private_workflow_pause.py", - "core/workflow/generator/test_prompts.py", + "core/trigger/utils/test_encryption.py", + "core/trigger/utils/test_endpoint.py", + "core/trigger/utils/test_locks.py", + "context/test_execution_context.py", + "context/test_flask_app_context.py", + "../e2e_tests/contracts/core/workflow/entities/test_private_workflow_pause.py", + "core/workflow/generator/prompts/test_planner_prompts.py", "core/workflow/generator/test_runner.py", - "core/workflow/generator/test_runner_missing.py", - "core/workflow/graph_engine/layers/conftest.py", - "core/workflow/graph_engine/layers/test_observability.py", - "core/workflow/graph_engine/test_mock_config.py", - "core/workflow/graph_engine/test_mock_factory.py", - "core/workflow/graph_engine/test_mock_nodes.py", - "core/workflow/graph_engine/test_parallel_human_input_join_resume.py", - "core/workflow/graph_engine/test_table_runner.py", - "core/workflow/graph_engine/test_tool_in_chatflow.py", + "core/app/workflow/layers/conftest.py", + "../helpers/workflow/mock_config.py", + "../helpers/workflow/mock_factory.py", + "../helpers/workflow/mock_nodes.py", + "../e2e_tests/contracts/core/workflow/graph_engine/test_parallel_human_input_join_resume.py", + "../helpers/workflow/table_runner.py", + "../e2e_tests/contracts/core/workflow/graph_engine/test_tool_in_chatflow.py", "core/workflow/nodes/agent/test_message_transformer.py", "core/workflow/nodes/agent/test_runtime_support.py", "core/workflow/nodes/agent_v2/test_agent_node.py", @@ -391,80 +343,69 @@ project-excludes = [ "core/workflow/nodes/agent_v2/test_runtime_request_builder.py", "core/workflow/nodes/agent_v2/test_session_store.py", "core/workflow/nodes/agent_v2/test_validators.py", - "core/workflow/nodes/answer/test_answer.py", - "core/workflow/nodes/base/test_base_node.py", - "core/workflow/nodes/base/test_get_node_type_classes_mapping.py", - "core/workflow/nodes/code/code_node_spec.py", + "../e2e_tests/contracts/core/workflow/nodes/answer/test_answer.py", + "../e2e_tests/contracts/core/workflow/nodes/base/test_base_node.py", + "../helpers/unit/core/workflow/nodes/code/code_node_spec.py", "core/workflow/nodes/datasource/test_datasource_node.py", - "core/workflow/nodes/http_request/test_http_request_executor.py", - "core/workflow/nodes/http_request/test_http_request_node.py", - "core/workflow/nodes/human_input/test_dify_owned_contracts.py", - "core/workflow/nodes/human_input/test_email_delivery_config.py", + "../e2e_tests/contracts/core/workflow/nodes/http_request/test_http_request_executor.py", + "../e2e_tests/contracts/core/workflow/nodes/http_request/test_http_request_node.py", "core/workflow/nodes/human_input/test_entities.py", - "core/workflow/nodes/human_input/test_human_input_form_filled_event.py", + "core/workflow/test_human_input_adapter.py", "core/workflow/nodes/knowledge_index/test_knowledge_index_node.py", "core/workflow/nodes/knowledge_retrieval/test_knowledge_retrieval_node.py", - "core/workflow/nodes/list_operator/node_spec.py", - "core/workflow/nodes/llm/test_llm_utils.py", - "core/workflow/nodes/llm/test_node.py", - "core/workflow/nodes/parameter_extractor/test_parameter_extractor_node.py", - "core/workflow/nodes/template_transform/template_transform_node_spec.py", - "core/workflow/nodes/template_transform/test_template_transform_node.py", - "core/workflow/nodes/test_base_node.py", - "core/workflow/nodes/test_document_extractor_node.py", - "core/workflow/nodes/test_if_else.py", - "core/workflow/nodes/test_list_operator.py", - "core/workflow/nodes/test_start_node_json_object.py", - "core/workflow/nodes/tool/test_tool_node.py", - "core/workflow/nodes/tool/test_tool_node_runtime.py", - "core/workflow/nodes/webhook/test_entities.py", - "core/workflow/nodes/webhook/test_webhook_file_conversion.py", - "core/workflow/nodes/webhook/test_webhook_node.py", - "core/workflow/test_enrich_pause_reasons.py", - "core/workflow/test_form_input_serialization_compat.py", - "core/workflow/test_graph_topology.py", - "core/workflow/test_human_input_adapter.py", - "core/workflow/test_human_input_callback.py", - "core/workflow/test_human_input_forms.py", - "core/workflow/test_node_factory.py", - "core/workflow/test_node_mapping_bootstrap.py", + "../helpers/unit/core/workflow/nodes/list_operator/node_spec.py", + "../e2e_tests/contracts/core/workflow/nodes/llm/test_llm_utils.py", + "../e2e_tests/contracts/core/workflow/nodes/llm/test_node.py", + "../e2e_tests/contracts/core/workflow/nodes/parameter_extractor/test_parameter_extractor_node.py", + "../helpers/unit/core/workflow/nodes/template_transform/template_transform_node_spec.py", + "../e2e_tests/contracts/core/workflow/nodes/template_transform/test_template_transform_node.py", + "../e2e_tests/contracts/core/workflow/nodes/test_base_node.py", + "../e2e_tests/contracts/core/workflow/nodes/test_document_extractor_node.py", + "../e2e_tests/contracts/core/workflow/nodes/test_if_else.py", + "../e2e_tests/contracts/core/workflow/nodes/test_list_operator.py", + "../e2e_tests/contracts/core/workflow/nodes/test_start_node_json_object.py", + "../e2e_tests/contracts/core/workflow/nodes/tool/test_tool_node.py", "core/workflow/test_node_runtime.py", - "core/workflow/test_system_variable.py", + "core/workflow/nodes/trigger_webhook/test_entities.py", + "core/workflow/nodes/trigger_webhook/test_node.py", + "core/workflow/nodes/human_input/test_pause_reason.py", + "core/workflow/test_graph_topology.py", + "core/workflow/nodes/human_input/test_callback.py", + "core/workflow/test_human_input_forms.py", + "../e2e_tests/contracts/core/workflow/test_node_mapping_bootstrap.py", + "core/workflow/test_system_variables.py", "core/workflow/test_workflow_entry.py", - "core/workflow/test_workflow_entry_helpers.py", - "core/workflow/test_workflow_entry_redis_channel.py", "dev/test_generate_knowledge_fs_contract.py", "enterprise/telemetry/test_contracts.py", "enterprise/telemetry/test_draft_trace.py", "enterprise/telemetry/test_enterprise_trace.py", - "enterprise/telemetry/test_gateway.py", "events/event_handlers/test_clean_when_document_deleted.py", "events/event_handlers/test_delete_tool_parameters_cache_when_sync_draft_workflow.py", - "events/test_app_event_signals.py", - "events/test_events_package_compat.py", + "events/test_app_event.py", + "events/test___init__.py", "extensions/logstore/repositories/test_logstore_api_workflow_node_execution_repository.py", "extensions/logstore/test_sql_escape.py", "extensions/otel/decorators/handlers/test_generate_handler.py", "extensions/otel/decorators/handlers/test_workflow_app_runner_handler.py", - "extensions/otel/test_retrieval_tracing.py", - "extensions/test_celery_ssl.py", + "core/workflow/nodes/knowledge_retrieval/test_retrieval.py", + "extensions/test_ext_celery.py", "extensions/test_ext_login.py", "extensions/test_ext_socketio.py", - "extensions/test_pubsub_channel.py", + "extensions/test_ext_redis.py", "extensions/test_workflow_warm_shutdown.py", - "factories/test_build_from_mapping.py", - "factories/test_file_factory.py", - "factories/test_file_validation.py", + "factories/file_factory/test_builders.py", + "factories/file_factory/test_remote.py", + "factories/file_factory/test_validation.py", "fields/test_dataset_fields.py", "fields/test_file_fields.py", "fields/test_message_fields.py", - "libs/_human_input/support.py", - "libs/_human_input/test_form_service.py", - "libs/_human_input/test_models.py", - "libs/broadcast_channel/redis/test_channel_unit_tests.py", - "libs/broadcast_channel/redis/test_streams_channel_unit_tests.py", + "../helpers/unit/libs/_human_input/support.py", + "../e2e_tests/contracts/libs/_human_input/test_form_service.py", + "../e2e_tests/contracts/libs/_human_input/test_models.py", + "libs/broadcast_channel/redis/test_pubsub_channel.py", + "libs/broadcast_channel/redis/test_streams_channel.py", "libs/test_archive_storage.py", - "libs/test_cron_compatibility.py", + "libs/test_schedule_utils.py", "libs/test_custom_inputs.py", "libs/test_datetime_utils.py", "libs/test_email_i18n.py", @@ -476,47 +417,45 @@ project-excludes = [ "libs/test_oauth_bearer.py", "libs/test_orjson.py", "libs/test_pagination.py", - "libs/test_pandas.py", - "libs/test_rate_limiter.py", - "libs/test_schedule_utils_enhanced.py", + "../e2e_tests/contracts/libs/test_pandas.py", "libs/test_token.py", - "migrations/test_uuidv7_pg18_migration.py", - "models/test_account_models.py", + "migrations/versions/test_2025_07_02_2332-1c9ba48be8e4_add_uuidv7_function_in_sql.py", + "models/test_account.py", "models/test_agent.py", - "models/test_app_models.py", - "models/test_dataset_models.py", - "models/test_tool_models.py", + "models/test_model.py", + "models/test_dataset.py", + "models/test_tools.py", "models/test_workflow.py", "repositories/test_sqlalchemy_api_workflow_run_repository.py", - "services/agent/test_agent_dsl_service.py", - "services/agent/test_agent_observability_service.py", - "services/agent/test_agent_services.py", + "services/agent/test_dsl_service.py", + "services/agent/test_observability_service.py", + "services/agent/test_roster_service.py", "services/agent/test_prompt_mentions.py", "services/agent/test_skill_package_service.py", "services/agent/test_workflow_publish_service.py", - "services/controller_api.py", + "../helpers/unit/services/controller_api.py", "services/data_migration/test_dependency_discovery_service.py", "services/data_migration/test_entities.py", "services/data_migration/test_export_service.py", "services/data_migration/test_import_service.py", "services/data_migration/test_package_service.py", "services/data_migration/test_report_service.py", - "services/dataset_service_test_helpers.py", - "services/document_service_validation.py", + "../helpers/unit/services/dataset_service_test_helpers.py", + "../helpers/unit/services/document_service_validation.py", "services/enterprise/test_account_deletion_sync.py", "services/enterprise/test_plugin_manager_service.py", - "services/enterprise/test_traceparent_propagation.py", - "services/hit_service.py", + "services/enterprise/test_base.py", + "../helpers/unit/services/hit_service.py", "services/plugin/test_plugin_parameter_service.py", - "services/plugin/test_plugin_service_installation.py", - "services/rag_pipeline/pipeline_template/test_built_in_retrieval.py", + "services/errors/test_plugin.py", + "services/rag_pipeline/pipeline_template/built_in/test_built_in_retrieval.py", "services/rag_pipeline/pipeline_template/test_pipeline_template_base.py", "services/rag_pipeline/test_pipeline_generate_service.py", "services/rag_pipeline/test_rag_pipeline_dsl_service.py", - "services/rag_pipeline/test_rag_pipeline_service.py", + "services/rag_pipeline/test_rag_pipeline.py", "services/rag_pipeline/test_rag_pipeline_task_proxy.py", "services/rag_pipeline/test_rag_pipeline_transform_service.py", - "services/retention/test_messages_clean_policy.py", + "services/retention/conversation/test_messages_clean_policy.py", "services/retention/workflow_run/test_archive_download_preparation.py", "services/retention/workflow_run/test_archive_download_task_cache.py", "services/retention/workflow_run/test_archive_log_service.py", @@ -526,80 +465,68 @@ project-excludes = [ "services/test_annotation_service.py", "services/test_api_token_service.py", "services/test_app_generate_service.py", - "services/test_app_generate_service_streaming_integration.py", "services/test_app_service.py", "services/test_app_task_service.py", "services/test_async_workflow_service.py", "services/test_billing_service.py", - "services/test_clear_free_plan_expired_workflow_run_logs.py", "services/test_clear_free_plan_tenant_expired_logs.py", "services/test_conversation_service.py", "services/test_credit_pool_service.py", - "services/test_dataset_service_dataset.py", - "services/test_dataset_service_document.py", - "services/test_dataset_service_segment.py", - "services/test_datasource_provider_service.py", - "services/test_export_app_messages.py", - "services/test_external_dataset_service.py", + "services/knowledge/test_dataset_service.py", + "services/data_source/test_provider_service.py", + "services/retention/conversation/test_message_export_service.py", + "services/knowledge/external/test_service.py", "services/test_file_service.py", "services/test_human_input_delivery_test_service.py", "services/test_human_input_file_upload_service.py", "services/test_knowledge_retrieval_inner_service.py", "services/test_message_service.py", - "services/test_messages_clean_service.py", - "services/test_metadata_service_session_boundary.py", - "services/test_model_provider_service.py", + "services/retention/conversation/test_messages_clean_service.py", + "services/knowledge/metadata/test_application.py", + "services/model_provider/test_service.py", "services/test_operation_service.py", - "services/test_rag_pipeline_task_proxy.py", - "services/test_schedule_service.py", + "services/trigger/test_schedule_service.py", "services/test_snippet_dsl_service.py", "services/test_snippet_generate_service.py", "services/test_snippet_service.py", - "services/test_summary_index_service.py", - "services/test_trigger_provider_service.py", + "services/knowledge/summaries/test_adapters.py", + "services/trigger/test_trigger_provider_service.py", "services/test_variable_truncator.py", - "services/test_webhook_service.py", - "services/test_webhook_service_additional.py", - "services/test_website_service.py", + "services/trigger/test_webhook_service.py", + "services/data_source/test_website_service.py", "services/test_workflow_collaboration_service.py", "services/test_workflow_node_execution_trace_service.py", "services/test_workflow_run_service.py", - "services/test_workflow_run_service_pause.py", "services/test_workflow_service.py", "services/tools/test_api_tools_manage_service.py", "services/tools/test_builtin_tools_manage_service.py", - "services/tools/test_mcp_tools_transform.py", + "services/tools/test_tools_transform_service.py", "services/tools/test_tool_labels_service.py", "services/tools/test_tools_manage_service.py", - "services/tools/test_tools_transform_service.py", - "services/workflow/test_draft_var_loader_simple.py", + "services/test_workflow_draft_variable_service.py", "services/workflow/test_node_output_inspector_service.py", - "services/workflow/test_workflow_converter_additional.py", - "services/workflow/test_workflow_draft_variable_service.py", - "services/workflow/test_workflow_event_snapshot_service.py", + "services/workflow/test_workflow_converter.py", + "services/test_workflow_event_snapshot_service.py", "services/workflow/test_workflow_event_snapshot_service_additional.py", - "services/workflow/test_workflow_human_input_delivery.py", - "services/workflow/test_workflow_restore.py", + "services/test_workflow_restore.py", "tasks/test_async_workflow_tasks.py", "tasks/test_batch_clean_document_task.py", "tasks/test_clean_dataset_task.py", "tasks/test_clean_document_task.py", - "tasks/test_dataset_indexing_task.py", + "tasks/test_document_indexing_task.py", "tasks/test_document_indexing_sync_task.py", "tasks/test_document_indexing_update_task.py", - "tasks/test_enable_segment_index_tasks.py", + "tasks/test_enable_segments_to_index_task.py", "tasks/test_human_input_timeout_tasks.py", "tasks/test_initialize_created_app_rbac_access_task.py", "tasks/test_ops_trace_task.py", "tasks/test_process_tenant_plugin_autoupgrade_check_task.py", "tasks/test_remove_app_and_related_data_task.py", "tasks/test_resume_agent_app_task.py", - "tasks/test_summary_queue_isolation.py", - "tasks/test_workflow_execute_task.py", - "test_makefile_backend_tests.py", - "test_pytest_dify.py", - "tools/test_api_tool.py", - "tools/test_mcp_tool.py", + "tasks/test_generate_summary_index_task.py", + "tasks/app_generate/test_workflow_execute_task.py", + "../e2e_tests/contracts/test_makefile_backend_tests.py", + "../e2e_tests/contracts/test_pytest_dify.py", ] [errors] diff --git a/api/tests/unit_tests/repositories/test_console_app_repository.py b/api/tests/unit_tests/repositories/app/test_console_repository.py similarity index 100% rename from api/tests/unit_tests/repositories/test_console_app_repository.py rename to api/tests/unit_tests/repositories/app/test_console_repository.py diff --git a/api/tests/unit_tests/repositories/app/test_mcp_server_repository.py b/api/tests/unit_tests/repositories/app/test_mcp_server_repository.py index d68e244e284..ee93608e462 100644 --- a/api/tests/unit_tests/repositories/app/test_mcp_server_repository.py +++ b/api/tests/unit_tests/repositories/app/test_mcp_server_repository.py @@ -22,8 +22,8 @@ from services.app.mcp_server_service import ( AppMCPServerUpdate, PublishedApp, ) -from tests.unit_tests.model_factories import make_app -from tests.unit_tests.repositories.app.console_visibility import UNADDRESSABLE_IN_WORKSPACE, MakeUnaddressable +from tests.helpers.unit.model_factories import make_app +from tests.helpers.unit.repositories.app.console_visibility import UNADDRESSABLE_IN_WORKSPACE, MakeUnaddressable WORKSPACE_ID = "tenant-1" DRAFT = AppMCPServerDraft( diff --git a/api/tests/unit_tests/repositories/app/test_site_command_repository.py b/api/tests/unit_tests/repositories/app/test_site_command_repository.py index b3ac290c672..9fa4ffa0878 100644 --- a/api/tests/unit_tests/repositories/app/test_site_command_repository.py +++ b/api/tests/unit_tests/repositories/app/test_site_command_repository.py @@ -6,7 +6,7 @@ from models.enums import CustomizeTokenStrategy from models.model import App, AppMode, Site from repositories.app.site_command_repository import AppSiteCommandRepository from services.app_site_service import AppSiteAppNotFoundError, AppSiteChanges, AppSiteNotFoundError -from tests.unit_tests.repositories.app.console_visibility import UNADDRESSABLE_IN_WORKSPACE, MakeUnaddressable +from tests.helpers.unit.repositories.app.console_visibility import UNADDRESSABLE_IN_WORKSPACE, MakeUnaddressable _APP_ID = "11111111-1111-1111-1111-111111111111" _WORKSPACE_ID = "22222222-2222-2222-2222-222222222222" diff --git a/api/tests/unit_tests/repositories/app/test_tracing_config_repository.py b/api/tests/unit_tests/repositories/app/test_tracing_config_repository.py index 4c94bde58b1..7cf8c93d530 100644 --- a/api/tests/unit_tests/repositories/app/test_tracing_config_repository.py +++ b/api/tests/unit_tests/repositories/app/test_tracing_config_repository.py @@ -1,3 +1,4 @@ +import json from collections.abc import Callable from datetime import datetime @@ -5,7 +6,9 @@ import pytest from sqlalchemy import select from sqlalchemy.orm import Session, sessionmaker +from machinery.context import RequestContext from models.model import App, AppMode, TraceAppConfig +from repositories.app.console_repository import ConsoleAppRepository from repositories.app.tracing_config_repository import SQLAlchemyAppTracingConfigRepository from services.app_tracing_config_service import ( AppTraceSettings, @@ -13,7 +16,7 @@ from services.app_tracing_config_service import ( AppTracingConfigConflictError, AppTracingConfigRecord, ) -from tests.unit_tests.repositories.app.console_visibility import UNADDRESSABLE_IN_WORKSPACE, MakeUnaddressable +from tests.helpers.unit.repositories.app.console_visibility import UNADDRESSABLE_IN_WORKSPACE, MakeUnaddressable _APP_ID = "11111111-1111-1111-1111-111111111111" _WORKSPACE_ID = "22222222-2222-2222-2222-222222222222" @@ -160,3 +163,45 @@ def test_record_mapping_does_not_expose_the_orm_model_or_its_config_dict( assert record.tracing_config is not None record.tracing_config["public_key"] = "changed" assert config.tracing_config == {"public_key": "original"} + + +"""Tracing policy round-trips without dropping content or Agent capture choices.""" + + +def test_capture_policy_round_trip_and_legacy_defaults(sqlite_session_factory: sessionmaker[Session]) -> None: + context = RequestContext("request", None, "account", "tenant") + with sqlite_session_factory.begin() as session: + session.add( + App( + id="app", + tenant_id="tenant", + name="Trace app", + mode=AppMode.WORKFLOW, + enable_site=False, + enable_api=False, + tracing=json.dumps({"enabled": True, "tracing_provider": "langfuse"}), + ) + ) + repository = ConsoleAppRepository(session_factory=sqlite_session_factory) + assert repository.get_trace(context, "app") == AppTraceSettings(True, "langfuse", False, True) + writer = SQLAlchemyAppTracingConfigRepository(session_factory=sqlite_session_factory) + writer.create( + workspace_id="tenant", app_id="app", tracing_provider="langfuse", tracing_config={"secret": "encrypted"} + ) + expected = writer.get(workspace_id="tenant", app_id="app", tracing_provider="langfuse") + policy = AppTraceSettings(True, "langfuse", True, False) + writer.set_trace(workspace_id="tenant", app_id="app", settings=policy, expected=expected) + assert repository.get_trace(context, "app") == policy + with sqlite_session_factory() as session: + app = session.get(App, "app") + assert app is not None + assert app.tracing is not None + assert json.loads(app.tracing) == { + "enabled": True, + "tracing_provider": "langfuse", + "capture_content": True, + "capture_agent": False, + } + with pytest.raises(AppTracingConfigAppNotFoundError): + writer.set_trace(workspace_id="other-tenant", app_id="app", settings=AppTraceSettings(), expected=None) + assert repository.get_trace(context, "app") == policy diff --git a/api/tests/unit_tests/repositories/app/test_tracing_settings_repository.py b/api/tests/unit_tests/repositories/app/test_tracing_settings_repository.py deleted file mode 100644 index 6ff338346ea..00000000000 --- a/api/tests/unit_tests/repositories/app/test_tracing_settings_repository.py +++ /dev/null @@ -1,51 +0,0 @@ -"""Tracing policy round-trips without dropping content or Agent capture choices.""" - -import json - -import pytest -from sqlalchemy.orm import Session, sessionmaker - -from machinery.context import RequestContext -from models.model import App, AppMode -from repositories.app.console_repository import ConsoleAppRepository -from repositories.app.tracing_config_repository import SQLAlchemyAppTracingConfigRepository -from services.app_tracing_config_service import AppTraceSettings, AppTracingConfigAppNotFoundError - - -def test_capture_policy_round_trip_and_legacy_defaults(sqlite_session_factory: sessionmaker[Session]) -> None: - context = RequestContext("request", None, "account", "tenant") - with sqlite_session_factory.begin() as session: - session.add( - App( - id="app", - tenant_id="tenant", - name="Trace app", - mode=AppMode.WORKFLOW, - enable_site=False, - enable_api=False, - tracing=json.dumps({"enabled": True, "tracing_provider": "langfuse"}), - ) - ) - repository = ConsoleAppRepository(session_factory=sqlite_session_factory) - assert repository.get_trace(context, "app") == AppTraceSettings(True, "langfuse", False, True) - writer = SQLAlchemyAppTracingConfigRepository(session_factory=sqlite_session_factory) - writer.create( - workspace_id="tenant", app_id="app", tracing_provider="langfuse", tracing_config={"secret": "encrypted"} - ) - expected = writer.get(workspace_id="tenant", app_id="app", tracing_provider="langfuse") - policy = AppTraceSettings(True, "langfuse", True, False) - writer.set_trace(workspace_id="tenant", app_id="app", settings=policy, expected=expected) - assert repository.get_trace(context, "app") == policy - with sqlite_session_factory() as session: - app = session.get(App, "app") - assert app is not None - assert app.tracing is not None - assert json.loads(app.tracing) == { - "enabled": True, - "tracing_provider": "langfuse", - "capture_content": True, - "capture_agent": False, - } - with pytest.raises(AppTracingConfigAppNotFoundError): - writer.set_trace(workspace_id="other-tenant", app_id="app", settings=AppTraceSettings(), expected=None) - assert repository.get_trace(context, "app") == policy diff --git a/api/tests/unit_tests/repositories/app/test_workflow_comment_repository.py b/api/tests/unit_tests/repositories/app/test_workflow_comment_repository.py index 60f4dcff7d8..57beb44bef6 100644 --- a/api/tests/unit_tests/repositories/app/test_workflow_comment_repository.py +++ b/api/tests/unit_tests/repositories/app/test_workflow_comment_repository.py @@ -25,8 +25,8 @@ from services.app.workflow_comment_service import ( WorkflowCommentReplyNotFoundError, WorkflowCommentService, ) -from tests.unit_tests.model_factories import make_account, make_app -from tests.unit_tests.repositories.app.console_visibility import UNADDRESSABLE_IN_WORKSPACE, MakeUnaddressable +from tests.helpers.unit.model_factories import make_account, make_app +from tests.helpers.unit.repositories.app.console_visibility import UNADDRESSABLE_IN_WORKSPACE, MakeUnaddressable TENANT_ID = "11111111-1111-1111-1111-111111111111" OTHER_TENANT_ID = "11111111-1111-1111-1111-111111111112" diff --git a/api/tests/unit_tests/repositories/knowledge/test_dataset_read_boundaries.py b/api/tests/unit_tests/repositories/knowledge/test_dataset_read_repository.py similarity index 100% rename from api/tests/unit_tests/repositories/knowledge/test_dataset_read_boundaries.py rename to api/tests/unit_tests/repositories/knowledge/test_dataset_read_repository.py diff --git a/api/tests/unit_tests/repositories/knowledge/test_indexing_segments.py b/api/tests/unit_tests/repositories/knowledge/test_segment_repository.py similarity index 100% rename from api/tests/unit_tests/repositories/knowledge/test_indexing_segments.py rename to api/tests/unit_tests/repositories/knowledge/test_segment_repository.py diff --git a/api/tests/unit_tests/repositories/test_agent_trial_preview.py b/api/tests/unit_tests/repositories/test_agent_trial_preview.py deleted file mode 100644 index a38c637912b..00000000000 --- a/api/tests/unit_tests/repositories/test_agent_trial_preview.py +++ /dev/null @@ -1,189 +0,0 @@ -from uuid import uuid4 - -import pytest -from sqlalchemy import func, select -from sqlalchemy.orm import Session, sessionmaker - -from fields.agent_fields import AgentAppComposerResponse -from models.agent import ( - Agent, - AgentConfigDraft, - AgentConfigDraftType, - AgentConfigRevision, - AgentConfigRevisionOperation, - AgentConfigSnapshot, - AgentScope, - AgentSource, -) -from models.agent_config_entities import AgentSoulConfig -from models.model import App, AppMode -from repositories.app_preview_query_repository import AppPreviewQueryRepository -from services.agent.composer_service import AgentComposerService -from services.app_definition_query_service import AppDefinitionUnavailableError -from services.app_preview_query_service import AppPreviewQueryService, AppPreviewRef - - -def _preview(session: Session, app: AppPreviewRef) -> AgentAppComposerResponse: - repository = AppPreviewQueryRepository(session_factory=sessionmaker(bind=session.get_bind())) - result = AppPreviewQueryService(apps=repository, is_previewable=lambda _app_id: True).get_agent_composer(app=app) - return AgentAppComposerResponse.model_validate(result) - - -def _published_app(session: Session) -> tuple[App, Agent, AgentConfigSnapshot]: - tenant_id, app_id, agent_id, snapshot_id = (str(uuid4()) for _ in range(4)) - app = App(id=app_id, tenant_id=tenant_id, name="Sample", mode=AppMode.AGENT, enable_site=False, enable_api=False) - agent = Agent( - id=agent_id, - tenant_id=tenant_id, - app_id=app_id, - name="Sample", - scope=AgentScope.ROSTER, - source=AgentSource.AGENT_APP, - active_config_snapshot_id=snapshot_id, - active_config_is_published=False, - ) - snapshot = AgentConfigSnapshot( - id=snapshot_id, - tenant_id=tenant_id, - agent_id=agent_id, - version=1, - config_snapshot=AgentSoulConfig.model_validate( - { - "prompt": {"system_prompt": "Published prompt"}, - "model": { - "plugin_id": "langgenius/openai", - "model_provider": "langgenius/openai/openai", - "model": "gpt-4o", - "credential_ref": {"type": "provider", "id": "private-credential"}, - }, - "env": {"variables": [{"name": "TOKEN", "value": "private-token"}]}, - "config_skills": [{"name": "research", "file_id": "private-skill", "description": "Research skill"}], - "config_files": [{"name": "guide.txt", "file_kind": "upload_file", "file_id": "private-file"}], - "tools": { - "cli_tools": [ - { - "name": "search-cli", - "env": {"variables": [{"name": "TOKEN", "value": "private-cli-token"}]}, - "invoke_metadata": {"password": "private-invoke-password"}, - } - ], - "dify_tools": [ - { - "provider_type": "api", - "provider_id": "private-provider", - "tool_name": "search", - "runtime_parameters": {"password": "private-password"}, - } - ], - }, - } - ), - ) - session.add_all( - [ - app, - agent, - snapshot, - AgentConfigRevision( - tenant_id=tenant_id, - agent_id=agent_id, - current_snapshot_id=snapshot_id, - revision=1, - operation=AgentConfigRevisionOperation.PUBLISH_DRAFT, - ), - ] - ) - session.commit() - return app, agent, snapshot - - -def test_detached_preview_reference_preserves_owner_scope(sqlite_session: Session) -> None: - app, _, _ = _published_app(sqlite_session) - preview = _preview(sqlite_session, AppPreviewRef(app_id=app.id, tenant_id=app.tenant_id)) - assert preview.agent_soul.prompt.system_prompt == "Published prompt" - with pytest.raises(AppDefinitionUnavailableError, match="unavailable"): - _preview(sqlite_session, AppPreviewRef(app_id=app.id, tenant_id=str(uuid4()))) - - -def test_preview_reads_published_snapshot_without_creating_draft(sqlite_session: Session) -> None: - app, agent, snapshot = _published_app(sqlite_session) - before = sqlite_session.scalar(select(func.count()).select_from(AgentConfigDraft)) - - preview = _preview(sqlite_session, AppPreviewRef(app_id=app.id, tenant_id=app.tenant_id)) - - assert preview.agent_soul.prompt.system_prompt == "Published prompt" - assert preview.agent_soul.model is not None - assert preview.agent_soul.model.model == "gpt-4o" - assert preview.agent_soul.config_skills[0].name == "research" - assert preview.agent_soul.config_files[0].name == "guide.txt" - assert preview.agent_soul.tools.dify_tools[0].tool_name == "search" - assert preview.agent.id == agent.id - assert preview.active_config_snapshot is not None - assert preview.active_config_snapshot.id == snapshot.id - assert preview.active_config_is_published is True - assert preview.draft is None - assert preview.save_options == [] - assert preview.chat_endpoint is None - serialized = preview.model_dump_json() - for value in [ - "private-credential", - "private-token", - "private-password", - "private-cli-token", - "private-invoke-password", - app.tenant_id, - ]: - assert value not in serialized - assert sqlite_session.scalar(select(func.count()).select_from(AgentConfigDraft)) == before - assert not sqlite_session.new - assert not sqlite_session.dirty - - -def test_preview_does_not_expose_unpublished_edits(sqlite_session: Session) -> None: - app, agent, snapshot = _published_app(sqlite_session) - sqlite_session.add( - AgentConfigDraft( - tenant_id=app.tenant_id, - agent_id=agent.id, - draft_type=AgentConfigDraftType.DRAFT, - draft_owner_key="", - base_snapshot_id=snapshot.id, - config_snapshot=AgentSoulConfig.model_validate({"prompt": {"system_prompt": "Unpublished secret prompt"}}), - ) - ) - sqlite_session.commit() - - preview = _preview(sqlite_session, AppPreviewRef(app_id=app.id, tenant_id=app.tenant_id)) - - assert preview.agent_soul.prompt.system_prompt == "Published prompt" - assert "Unpublished" not in preview.model_dump_json() - editable = AgentAppComposerResponse.model_validate( - AgentComposerService.load_agent_composer( - session=sqlite_session, - tenant_id=app.tenant_id, - agent_id=agent.id, - ) - ) - assert editable.agent_soul.prompt.system_prompt == "Unpublished secret prompt" - assert editable.draft is not None - assert editable.save_options - - -def test_preview_rejects_seeded_import_snapshot(sqlite_session: Session) -> None: - app, agent, _ = _published_app(sqlite_session) - revision = sqlite_session.scalar(select(AgentConfigRevision).where(AgentConfigRevision.agent_id == agent.id)) - assert revision is not None - revision.operation = AgentConfigRevisionOperation.IMPORT_PACKAGE - sqlite_session.commit() - - with pytest.raises(AppDefinitionUnavailableError): - _preview(sqlite_session, AppPreviewRef(app_id=app.id, tenant_id=app.tenant_id)) - - -def test_preview_rejects_cross_tenant_snapshot(sqlite_session: Session) -> None: - app, _, snapshot = _published_app(sqlite_session) - snapshot.tenant_id = str(uuid4()) - sqlite_session.commit() - - with pytest.raises(AppDefinitionUnavailableError, match="unavailable"): - _preview(sqlite_session, AppPreviewRef(app_id=app.id, tenant_id=app.tenant_id)) diff --git a/api/tests/unit_tests/repositories/test_app_definition_query_repository.py b/api/tests/unit_tests/repositories/test_app_definition_query_repository.py index cec3f8fbbee..83017e755f9 100644 --- a/api/tests/unit_tests/repositories/test_app_definition_query_repository.py +++ b/api/tests/unit_tests/repositories/test_app_definition_query_repository.py @@ -1,17 +1,33 @@ import json +from typing import Any +from uuid import NAMESPACE_URL, uuid5 import pytest from sqlalchemy.orm import Session, sessionmaker +from core.app.app_config.common.parameters_mapping import get_parameters_from_feature_dict +from core.app.apps.agent_app.errors import AgentAppGeneratorError, AgentAppNotPublishedError from core.tools.entities.tool_entities import ApiProviderSchemaType from machinery.context import AppRequestContext from models.account import Account, Tenant, TenantStatus +from models.agent import ( + Agent, + AgentConfigRevision, + AgentConfigRevisionOperation, + AgentConfigSnapshot, + AgentScope, + AgentSource, + AgentStatus, +) from models.enums import CustomizeTokenStrategy, TagType -from models.model import App, AppMode, AppModelConfig, IconType, Site, Tag, TagBinding +from models.model import App, AppAnnotationSetting, AppMode, AppModelConfig, IconType, Site, Tag, TagBinding from models.tools import ApiToolProvider from models.workflow import Workflow, WorkflowKind, WorkflowType from repositories.app.console_repository import ConsoleAppRepository -from repositories.app_definition_query_repository import AppDefinitionQueryRepository +from repositories.app_definition_query_repository import ( + AppDefinitionQueryRepository, + _get_public_agent_parameter_config, +) from services.app_definition_query_service import ( AppDefinitionSummary, AppParameterConfig, @@ -614,3 +630,375 @@ def test_get_tool_icon_sources_returns_empty_for_missing_published_config( repository = AppDefinitionQueryRepository(session_factory=sqlite_session_factory) assert repository.get_tool_icon_sources(_APP_ID) == () + + +def _stable_uuid(value: str) -> str: + return str(uuid5(NAMESPACE_URL, value)) + + +def _app_model(*, tenant_id: str, app_model_config: AppModelConfig | None = None) -> App: + return App( + id=_stable_uuid(f"app:{tenant_id}"), + tenant_id=tenant_id, + name="Agent App", + mode=AppMode.AGENT, + app_model_config_id=app_model_config.id if app_model_config else None, + ) + + +def _persist_agent( + session: Session, + *, + tenant_id: str, + agent_id: str, + active_config_snapshot_id: str | None, + active_config_is_published: bool, + app_id: str | None = None, +) -> Agent: + agent = Agent( + id=agent_id, + tenant_id=tenant_id, + name="Agent", + scope=AgentScope.ROSTER, + source=AgentSource.AGENT_APP, + status=AgentStatus.ACTIVE, + active_config_snapshot_id=active_config_snapshot_id, + active_config_is_published=active_config_is_published, + app_id=app_id, + ) + session.add(agent) + session.commit() + return agent + + +def _persist_snapshot( + session: Session, + *, + snapshot_id: str, + tenant_id: str, + agent_id: str, + config_snapshot: dict[str, Any], + publish_visible: bool = True, +) -> AgentConfigSnapshot: + snapshot = AgentConfigSnapshot( + id=snapshot_id, + tenant_id=tenant_id, + agent_id=agent_id, + version=1, + home_snapshot_id=_stable_uuid(f"home-snapshot:{snapshot_id}"), + config_snapshot=config_snapshot, + ) + session.add(snapshot) + if publish_visible: + _persist_publish_revision( + session, + snapshot_id=snapshot_id, + tenant_id=tenant_id, + agent_id=agent_id, + commit=False, + ) + session.commit() + return snapshot + + +def _persist_publish_revision( + session: Session, + *, + snapshot_id: str, + tenant_id: str, + agent_id: str, + commit: bool = True, +) -> None: + session.add( + AgentConfigRevision( + tenant_id=tenant_id, + agent_id=agent_id, + current_snapshot_id=snapshot_id, + revision=1, + operation=AgentConfigRevisionOperation.PUBLISH_DRAFT, + ) + ) + if commit: + session.commit() + + +def test_get_public_parameter_config_loads_agent_snapshot( + sqlite_session_factory: sessionmaker[Session], +) -> None: + tenant_id = _stable_uuid("tenant:repository") + app_id = _stable_uuid("app:repository") + agent_id = _stable_uuid("agent:repository") + snapshot_id = _stable_uuid("snapshot:repository") + with sqlite_session_factory() as session: + session.add( + App( + id=app_id, + tenant_id=tenant_id, + name="Agent App", + description="", + mode=AppMode.AGENT, + icon_type=None, + icon=None, + icon_background=None, + enable_site=True, + enable_api=True, + max_active_requests=None, + ) + ) + session.commit() + _persist_agent( + session, + tenant_id=tenant_id, + agent_id=agent_id, + active_config_snapshot_id=snapshot_id, + active_config_is_published=True, + app_id=app_id, + ) + _persist_snapshot( + session, + snapshot_id=snapshot_id, + tenant_id=tenant_id, + agent_id=agent_id, + config_snapshot={"app_variables": [{"name": "topic", "type": "string", "required": True}]}, + ) + + result = AppDefinitionQueryRepository(session_factory=sqlite_session_factory).get_published_parameter_config( + app_id, public_runtime=True + ) + + assert result is not None + assert result.user_input_form == [{"text-input": {"label": "topic", "variable": "topic", "required": True}}] + + +@pytest.mark.parametrize( + "sqlite_session", + [(Agent, AgentConfigSnapshot, AgentConfigRevision, AppAnnotationSetting, AppModelConfig)], + indirect=True, +) +@pytest.mark.parametrize("foreign_model_config", [False, True]) +def test_published_agent_app_parameters_use_soul_file_upload( + sqlite_session: Session, foreign_model_config: bool +) -> None: + tenant_id = _stable_uuid("tenant:one") + agent_id = _stable_uuid("agent:one") + snapshot_id = _stable_uuid("snapshot:one") + app_model_config = AppModelConfig( + app_id=_stable_uuid("foreign-app") if foreign_model_config else _stable_uuid(f"app:{tenant_id}"), + opening_statement="Hi from legacy presentation config", + file_upload=json.dumps({"enabled": False, "image": {"enabled": False}}), + ) + sqlite_session.add(app_model_config) + sqlite_session.commit() + app_model = _app_model(tenant_id=tenant_id, app_model_config=app_model_config) + _persist_agent( + sqlite_session, + tenant_id=tenant_id, + agent_id=agent_id, + active_config_snapshot_id=snapshot_id, + active_config_is_published=True, + app_id=app_model.id, + ) + _persist_snapshot( + sqlite_session, + snapshot_id=snapshot_id, + tenant_id=tenant_id, + agent_id=agent_id, + config_snapshot={ + "app_features": { + "file_upload": { + "enabled": True, + "allowed_file_extensions": ["PNG"], + "allowed_file_types": ["image"], + "allowed_file_upload_methods": ["local_file"], + "image": {"enabled": True}, + "number_limits": 2, + } + }, + "app_variables": [{"name": "topic", "type": "string", "required": True}], + }, + ) + + features_dict, user_input_form = _get_public_agent_parameter_config( + app_model, + session=sqlite_session, + ) + parameters = get_parameters_from_feature_dict(features_dict=features_dict, user_input_form=user_input_form) + + assert parameters["opening_statement"] == (None if foreign_model_config else "Hi from legacy presentation config") + assert parameters["file_upload"] == { + "enabled": True, + "allowed_file_extensions": ["PNG"], + "allowed_file_types": ["image"], + "allowed_file_upload_methods": ["local_file"], + "image": {"enabled": True}, + "number_limits": 2, + } + assert parameters["user_input_form"] == [{"text-input": {"label": "topic", "variable": "topic", "required": True}}] + + +@pytest.mark.parametrize("sqlite_session", [(Agent, AgentConfigSnapshot, AgentConfigRevision)], indirect=True) +def test_published_agent_app_parameters_requires_bound_agent(sqlite_session: Session): + tenant_id = _stable_uuid("tenant:unbound") + app_model = _app_model(tenant_id=tenant_id) + + with pytest.raises(AgentAppGeneratorError, match="no bound Agent"): + _get_public_agent_parameter_config(app_model, session=sqlite_session) + + +@pytest.mark.parametrize("sqlite_session", [(Agent, AgentConfigSnapshot, AgentConfigRevision)], indirect=True) +def test_published_agent_app_parameters_requires_existing_active_agent(sqlite_session: Session): + requested_tenant_id = _stable_uuid("tenant:requested") + agent_id = _stable_uuid("agent:cross-tenant") + app_model = _app_model(tenant_id=requested_tenant_id) + _persist_agent( + sqlite_session, + tenant_id=_stable_uuid("tenant:other"), + agent_id=agent_id, + active_config_snapshot_id=None, + active_config_is_published=False, + app_id=app_model.id, + ) + + with pytest.raises(AgentAppGeneratorError, match="no bound Agent"): + _get_public_agent_parameter_config(app_model, session=sqlite_session) + + +@pytest.mark.parametrize("active_config_is_published", [True, False]) +@pytest.mark.parametrize("sqlite_session", [(Agent, AgentConfigSnapshot, AgentConfigRevision)], indirect=True) +def test_published_agent_app_parameters_requires_published_agent( + active_config_is_published: bool, sqlite_session: Session +): + tenant_id = _stable_uuid(f"tenant:published:{active_config_is_published}") + agent_id = _stable_uuid(f"agent:published:{active_config_is_published}") + app_model = _app_model(tenant_id=tenant_id) + _persist_agent( + sqlite_session, + tenant_id=tenant_id, + agent_id=agent_id, + active_config_snapshot_id=None, + active_config_is_published=active_config_is_published, + app_id=app_model.id, + ) + + with pytest.raises(AgentAppNotPublishedError, match="not been published"): + _get_public_agent_parameter_config(app_model, session=sqlite_session) + + +@pytest.mark.parametrize("sqlite_session", [(Agent, AgentConfigSnapshot, AgentConfigRevision)], indirect=True) +def test_published_agent_app_parameters_allows_unpublished_draft_with_active_snapshot(sqlite_session: Session): + tenant_id = _stable_uuid("tenant:unpublished-draft") + agent_id = _stable_uuid("agent:unpublished-draft") + snapshot_id = _stable_uuid("snapshot:unpublished-draft") + app_model = _app_model(tenant_id=tenant_id) + _persist_agent( + sqlite_session, + tenant_id=tenant_id, + agent_id=agent_id, + active_config_snapshot_id=snapshot_id, + active_config_is_published=False, + app_id=app_model.id, + ) + _persist_snapshot( + sqlite_session, + snapshot_id=snapshot_id, + tenant_id=tenant_id, + agent_id=agent_id, + config_snapshot={}, + ) + + features_dict, user_input_form = _get_public_agent_parameter_config( + app_model, + session=sqlite_session, + ) + + assert features_dict["file_upload"]["enabled"] is True + assert user_input_form == [] + + +@pytest.mark.parametrize("sqlite_session", [(Agent, AgentConfigSnapshot, AgentConfigRevision)], indirect=True) +def test_published_agent_app_parameters_rejects_seeded_unpublished_snapshot(sqlite_session: Session): + tenant_id = _stable_uuid("tenant:never-published") + agent_id = _stable_uuid("agent:never-published") + snapshot_id = _stable_uuid("snapshot:never-published") + app_model = _app_model(tenant_id=tenant_id) + _persist_agent( + sqlite_session, + tenant_id=tenant_id, + agent_id=agent_id, + active_config_snapshot_id=snapshot_id, + active_config_is_published=False, + app_id=app_model.id, + ) + _persist_snapshot( + sqlite_session, + snapshot_id=snapshot_id, + tenant_id=tenant_id, + agent_id=agent_id, + config_snapshot={}, + publish_visible=False, + ) + + with pytest.raises(AgentAppNotPublishedError, match="not been published"): + _get_public_agent_parameter_config(app_model, session=sqlite_session) + + +@pytest.mark.parametrize("sqlite_session", [(Agent, AgentConfigSnapshot, AgentConfigRevision)], indirect=True) +def test_published_agent_app_parameters_requires_published_snapshot(sqlite_session: Session): + tenant_id = _stable_uuid("tenant:missing-snapshot") + agent_id = _stable_uuid("agent:missing-snapshot") + app_model = _app_model(tenant_id=tenant_id) + _persist_agent( + sqlite_session, + tenant_id=tenant_id, + agent_id=agent_id, + active_config_snapshot_id=_stable_uuid("snapshot:missing"), + active_config_is_published=True, + app_id=app_model.id, + ) + _persist_publish_revision( + sqlite_session, + snapshot_id=_stable_uuid("snapshot:missing"), + tenant_id=tenant_id, + agent_id=agent_id, + ) + + with pytest.raises(AgentAppGeneratorError, match="published version not found"): + _get_public_agent_parameter_config(app_model, session=sqlite_session) + + +@pytest.mark.parametrize("sqlite_session", [(Agent, AgentConfigSnapshot, AgentConfigRevision)], indirect=True) +def test_published_agent_app_parameters_allows_missing_legacy_app_model_config(sqlite_session: Session): + tenant_id = _stable_uuid("tenant:no-legacy-config") + agent_id = _stable_uuid("agent:no-legacy-config") + snapshot_id = _stable_uuid("snapshot:no-legacy-config") + app_model = _app_model(tenant_id=tenant_id) + _persist_agent( + sqlite_session, + tenant_id=tenant_id, + agent_id=agent_id, + active_config_snapshot_id=snapshot_id, + active_config_is_published=True, + app_id=app_model.id, + ) + _persist_snapshot( + sqlite_session, + snapshot_id=snapshot_id, + tenant_id=tenant_id, + agent_id=agent_id, + config_snapshot={}, + ) + + features_dict, user_input_form = _get_public_agent_parameter_config( + app_model, + session=sqlite_session, + ) + + assert features_dict["file_upload"] == { + "allowed_file_extensions": ["JPG", "JPEG", "PNG", "GIF", "WEBP", "SVG"], + "allowed_file_types": ["document", "image", "audio", "video"], + "allowed_file_upload_methods": ["local_file", "remote_url"], + "enabled": True, + "image": {"enabled": True}, + "number_limits": 3, + } + assert user_input_form == [] diff --git a/api/tests/unit_tests/repositories/test_app_definition_query_repository_agent.py b/api/tests/unit_tests/repositories/test_app_definition_query_repository_agent.py deleted file mode 100644 index d51f4b6dafc..00000000000 --- a/api/tests/unit_tests/repositories/test_app_definition_query_repository_agent.py +++ /dev/null @@ -1,395 +0,0 @@ -import json -from typing import Any -from uuid import NAMESPACE_URL, uuid5 - -import pytest -from sqlalchemy.orm import Session, sessionmaker - -from core.app.app_config.common.parameters_mapping import get_parameters_from_feature_dict -from core.app.apps.agent_app.errors import AgentAppGeneratorError, AgentAppNotPublishedError -from models.agent import ( - Agent, - AgentConfigRevision, - AgentConfigRevisionOperation, - AgentConfigSnapshot, - AgentScope, - AgentSource, - AgentStatus, -) -from models.model import App, AppAnnotationSetting, AppMode, AppModelConfig -from repositories.app_definition_query_repository import ( - AppDefinitionQueryRepository, - _get_public_agent_parameter_config, -) - - -def _stable_uuid(value: str) -> str: - return str(uuid5(NAMESPACE_URL, value)) - - -def _app_model(*, tenant_id: str, app_model_config: AppModelConfig | None = None) -> App: - return App( - id=_stable_uuid(f"app:{tenant_id}"), - tenant_id=tenant_id, - name="Agent App", - mode=AppMode.AGENT, - app_model_config_id=app_model_config.id if app_model_config else None, - ) - - -def _persist_agent( - session: Session, - *, - tenant_id: str, - agent_id: str, - active_config_snapshot_id: str | None, - active_config_is_published: bool, - app_id: str | None = None, -) -> Agent: - agent = Agent( - id=agent_id, - tenant_id=tenant_id, - name="Agent", - scope=AgentScope.ROSTER, - source=AgentSource.AGENT_APP, - status=AgentStatus.ACTIVE, - active_config_snapshot_id=active_config_snapshot_id, - active_config_is_published=active_config_is_published, - app_id=app_id, - ) - session.add(agent) - session.commit() - return agent - - -def _persist_snapshot( - session: Session, - *, - snapshot_id: str, - tenant_id: str, - agent_id: str, - config_snapshot: dict[str, Any], - publish_visible: bool = True, -) -> AgentConfigSnapshot: - snapshot = AgentConfigSnapshot( - id=snapshot_id, - tenant_id=tenant_id, - agent_id=agent_id, - version=1, - home_snapshot_id=_stable_uuid(f"home-snapshot:{snapshot_id}"), - config_snapshot=config_snapshot, - ) - session.add(snapshot) - if publish_visible: - _persist_publish_revision( - session, - snapshot_id=snapshot_id, - tenant_id=tenant_id, - agent_id=agent_id, - commit=False, - ) - session.commit() - return snapshot - - -def _persist_publish_revision( - session: Session, - *, - snapshot_id: str, - tenant_id: str, - agent_id: str, - commit: bool = True, -) -> None: - session.add( - AgentConfigRevision( - tenant_id=tenant_id, - agent_id=agent_id, - current_snapshot_id=snapshot_id, - revision=1, - operation=AgentConfigRevisionOperation.PUBLISH_DRAFT, - ) - ) - if commit: - session.commit() - - -def test_get_public_parameter_config_loads_agent_snapshot( - sqlite_session_factory: sessionmaker[Session], -) -> None: - tenant_id = _stable_uuid("tenant:repository") - app_id = _stable_uuid("app:repository") - agent_id = _stable_uuid("agent:repository") - snapshot_id = _stable_uuid("snapshot:repository") - with sqlite_session_factory() as session: - session.add( - App( - id=app_id, - tenant_id=tenant_id, - name="Agent App", - description="", - mode=AppMode.AGENT, - icon_type=None, - icon=None, - icon_background=None, - enable_site=True, - enable_api=True, - max_active_requests=None, - ) - ) - session.commit() - _persist_agent( - session, - tenant_id=tenant_id, - agent_id=agent_id, - active_config_snapshot_id=snapshot_id, - active_config_is_published=True, - app_id=app_id, - ) - _persist_snapshot( - session, - snapshot_id=snapshot_id, - tenant_id=tenant_id, - agent_id=agent_id, - config_snapshot={"app_variables": [{"name": "topic", "type": "string", "required": True}]}, - ) - - result = AppDefinitionQueryRepository(session_factory=sqlite_session_factory).get_published_parameter_config( - app_id, public_runtime=True - ) - - assert result is not None - assert result.user_input_form == [{"text-input": {"label": "topic", "variable": "topic", "required": True}}] - - -@pytest.mark.parametrize( - "sqlite_session", - [(Agent, AgentConfigSnapshot, AgentConfigRevision, AppAnnotationSetting, AppModelConfig)], - indirect=True, -) -@pytest.mark.parametrize("foreign_model_config", [False, True]) -def test_published_agent_app_parameters_use_soul_file_upload( - sqlite_session: Session, foreign_model_config: bool -) -> None: - tenant_id = _stable_uuid("tenant:one") - agent_id = _stable_uuid("agent:one") - snapshot_id = _stable_uuid("snapshot:one") - app_model_config = AppModelConfig( - app_id=_stable_uuid("foreign-app") if foreign_model_config else _stable_uuid(f"app:{tenant_id}"), - opening_statement="Hi from legacy presentation config", - file_upload=json.dumps({"enabled": False, "image": {"enabled": False}}), - ) - sqlite_session.add(app_model_config) - sqlite_session.commit() - app_model = _app_model(tenant_id=tenant_id, app_model_config=app_model_config) - _persist_agent( - sqlite_session, - tenant_id=tenant_id, - agent_id=agent_id, - active_config_snapshot_id=snapshot_id, - active_config_is_published=True, - app_id=app_model.id, - ) - _persist_snapshot( - sqlite_session, - snapshot_id=snapshot_id, - tenant_id=tenant_id, - agent_id=agent_id, - config_snapshot={ - "app_features": { - "file_upload": { - "enabled": True, - "allowed_file_extensions": ["PNG"], - "allowed_file_types": ["image"], - "allowed_file_upload_methods": ["local_file"], - "image": {"enabled": True}, - "number_limits": 2, - } - }, - "app_variables": [{"name": "topic", "type": "string", "required": True}], - }, - ) - - features_dict, user_input_form = _get_public_agent_parameter_config( - app_model, - session=sqlite_session, - ) - parameters = get_parameters_from_feature_dict(features_dict=features_dict, user_input_form=user_input_form) - - assert parameters["opening_statement"] == (None if foreign_model_config else "Hi from legacy presentation config") - assert parameters["file_upload"] == { - "enabled": True, - "allowed_file_extensions": ["PNG"], - "allowed_file_types": ["image"], - "allowed_file_upload_methods": ["local_file"], - "image": {"enabled": True}, - "number_limits": 2, - } - assert parameters["user_input_form"] == [{"text-input": {"label": "topic", "variable": "topic", "required": True}}] - - -@pytest.mark.parametrize("sqlite_session", [(Agent, AgentConfigSnapshot, AgentConfigRevision)], indirect=True) -def test_published_agent_app_parameters_requires_bound_agent(sqlite_session: Session): - tenant_id = _stable_uuid("tenant:unbound") - app_model = _app_model(tenant_id=tenant_id) - - with pytest.raises(AgentAppGeneratorError, match="no bound Agent"): - _get_public_agent_parameter_config(app_model, session=sqlite_session) - - -@pytest.mark.parametrize("sqlite_session", [(Agent, AgentConfigSnapshot, AgentConfigRevision)], indirect=True) -def test_published_agent_app_parameters_requires_existing_active_agent(sqlite_session: Session): - requested_tenant_id = _stable_uuid("tenant:requested") - agent_id = _stable_uuid("agent:cross-tenant") - app_model = _app_model(tenant_id=requested_tenant_id) - _persist_agent( - sqlite_session, - tenant_id=_stable_uuid("tenant:other"), - agent_id=agent_id, - active_config_snapshot_id=None, - active_config_is_published=False, - app_id=app_model.id, - ) - - with pytest.raises(AgentAppGeneratorError, match="no bound Agent"): - _get_public_agent_parameter_config(app_model, session=sqlite_session) - - -@pytest.mark.parametrize("active_config_is_published", [True, False]) -@pytest.mark.parametrize("sqlite_session", [(Agent, AgentConfigSnapshot, AgentConfigRevision)], indirect=True) -def test_published_agent_app_parameters_requires_published_agent( - active_config_is_published: bool, sqlite_session: Session -): - tenant_id = _stable_uuid(f"tenant:published:{active_config_is_published}") - agent_id = _stable_uuid(f"agent:published:{active_config_is_published}") - app_model = _app_model(tenant_id=tenant_id) - _persist_agent( - sqlite_session, - tenant_id=tenant_id, - agent_id=agent_id, - active_config_snapshot_id=None, - active_config_is_published=active_config_is_published, - app_id=app_model.id, - ) - - with pytest.raises(AgentAppNotPublishedError, match="not been published"): - _get_public_agent_parameter_config(app_model, session=sqlite_session) - - -@pytest.mark.parametrize("sqlite_session", [(Agent, AgentConfigSnapshot, AgentConfigRevision)], indirect=True) -def test_published_agent_app_parameters_allows_unpublished_draft_with_active_snapshot(sqlite_session: Session): - tenant_id = _stable_uuid("tenant:unpublished-draft") - agent_id = _stable_uuid("agent:unpublished-draft") - snapshot_id = _stable_uuid("snapshot:unpublished-draft") - app_model = _app_model(tenant_id=tenant_id) - _persist_agent( - sqlite_session, - tenant_id=tenant_id, - agent_id=agent_id, - active_config_snapshot_id=snapshot_id, - active_config_is_published=False, - app_id=app_model.id, - ) - _persist_snapshot( - sqlite_session, - snapshot_id=snapshot_id, - tenant_id=tenant_id, - agent_id=agent_id, - config_snapshot={}, - ) - - features_dict, user_input_form = _get_public_agent_parameter_config( - app_model, - session=sqlite_session, - ) - - assert features_dict["file_upload"]["enabled"] is True - assert user_input_form == [] - - -@pytest.mark.parametrize("sqlite_session", [(Agent, AgentConfigSnapshot, AgentConfigRevision)], indirect=True) -def test_published_agent_app_parameters_rejects_seeded_unpublished_snapshot(sqlite_session: Session): - tenant_id = _stable_uuid("tenant:never-published") - agent_id = _stable_uuid("agent:never-published") - snapshot_id = _stable_uuid("snapshot:never-published") - app_model = _app_model(tenant_id=tenant_id) - _persist_agent( - sqlite_session, - tenant_id=tenant_id, - agent_id=agent_id, - active_config_snapshot_id=snapshot_id, - active_config_is_published=False, - app_id=app_model.id, - ) - _persist_snapshot( - sqlite_session, - snapshot_id=snapshot_id, - tenant_id=tenant_id, - agent_id=agent_id, - config_snapshot={}, - publish_visible=False, - ) - - with pytest.raises(AgentAppNotPublishedError, match="not been published"): - _get_public_agent_parameter_config(app_model, session=sqlite_session) - - -@pytest.mark.parametrize("sqlite_session", [(Agent, AgentConfigSnapshot, AgentConfigRevision)], indirect=True) -def test_published_agent_app_parameters_requires_published_snapshot(sqlite_session: Session): - tenant_id = _stable_uuid("tenant:missing-snapshot") - agent_id = _stable_uuid("agent:missing-snapshot") - app_model = _app_model(tenant_id=tenant_id) - _persist_agent( - sqlite_session, - tenant_id=tenant_id, - agent_id=agent_id, - active_config_snapshot_id=_stable_uuid("snapshot:missing"), - active_config_is_published=True, - app_id=app_model.id, - ) - _persist_publish_revision( - sqlite_session, - snapshot_id=_stable_uuid("snapshot:missing"), - tenant_id=tenant_id, - agent_id=agent_id, - ) - - with pytest.raises(AgentAppGeneratorError, match="published version not found"): - _get_public_agent_parameter_config(app_model, session=sqlite_session) - - -@pytest.mark.parametrize("sqlite_session", [(Agent, AgentConfigSnapshot, AgentConfigRevision)], indirect=True) -def test_published_agent_app_parameters_allows_missing_legacy_app_model_config(sqlite_session: Session): - tenant_id = _stable_uuid("tenant:no-legacy-config") - agent_id = _stable_uuid("agent:no-legacy-config") - snapshot_id = _stable_uuid("snapshot:no-legacy-config") - app_model = _app_model(tenant_id=tenant_id) - _persist_agent( - sqlite_session, - tenant_id=tenant_id, - agent_id=agent_id, - active_config_snapshot_id=snapshot_id, - active_config_is_published=True, - app_id=app_model.id, - ) - _persist_snapshot( - sqlite_session, - snapshot_id=snapshot_id, - tenant_id=tenant_id, - agent_id=agent_id, - config_snapshot={}, - ) - - features_dict, user_input_form = _get_public_agent_parameter_config( - app_model, - session=sqlite_session, - ) - - assert features_dict["file_upload"] == { - "allowed_file_extensions": ["JPG", "JPEG", "PNG", "GIF", "WEBP", "SVG"], - "allowed_file_types": ["document", "image", "audio", "video"], - "allowed_file_upload_methods": ["local_file", "remote_url"], - "enabled": True, - "image": {"enabled": True}, - "number_limits": 3, - } - assert user_input_form == [] diff --git a/api/tests/unit_tests/repositories/test_app_preview_query_repository.py b/api/tests/unit_tests/repositories/test_app_preview_query_repository.py index fd37d0e613d..99b5f775f32 100644 --- a/api/tests/unit_tests/repositories/test_app_preview_query_repository.py +++ b/api/tests/unit_tests/repositories/test_app_preview_query_repository.py @@ -1,5 +1,3 @@ -"""Preview reads enforce app ownership without adding execution or dataset ACL policy.""" - import json from dataclasses import asdict, replace from datetime import datetime @@ -7,22 +5,209 @@ from typing import Literal, override from uuid import uuid4 import pytest -from sqlalchemy import Connection, Engine, ExecutionContext, event, select, text +from sqlalchemy import Connection, Engine, ExecutionContext, event, func, select, text from sqlalchemy.orm import Session, sessionmaker from sqlalchemy.pool import QueuePool from core.tools.entities.tool_entities import ApiProviderSchemaType +from fields.agent_fields import AgentAppComposerResponse from graphon.variables import StringVariable from models.account import Account, Tenant, TenantStatus +from models.agent import ( + Agent, + AgentConfigDraft, + AgentConfigDraftType, + AgentConfigRevision, + AgentConfigRevisionOperation, + AgentConfigSnapshot, + AgentScope, + AgentSource, +) +from models.agent_config_entities import AgentSoulConfig from models.dataset import Dataset from models.enums import CustomizeTokenStrategy from models.model import App, AppMode, AppModelConfig, IconType, Site from models.tools import ApiToolProvider, WorkflowToolProvider from models.workflow import Workflow, WorkflowType from repositories.app_preview_query_repository import AppPreviewQueryRepository -from services.app_definition_query_service import AppSiteConfiguration +from services.agent.composer_service import AgentComposerService +from services.app_definition_query_service import AppDefinitionUnavailableError, AppSiteConfiguration from services.app_preview_details_service import AppPreviewAccount -from services.app_preview_query_service import AppPreviewDataset, AppPreviewRef +from services.app_preview_query_service import AppPreviewDataset, AppPreviewQueryService, AppPreviewRef + + +def _preview(session: Session, app: AppPreviewRef) -> AgentAppComposerResponse: + repository = AppPreviewQueryRepository(session_factory=sessionmaker(bind=session.get_bind())) + result = AppPreviewQueryService(apps=repository, is_previewable=lambda _app_id: True).get_agent_composer(app=app) + return AgentAppComposerResponse.model_validate(result) + + +def _published_app(session: Session) -> tuple[App, Agent, AgentConfigSnapshot]: + tenant_id, app_id, agent_id, snapshot_id = (str(uuid4()) for _ in range(4)) + app = App(id=app_id, tenant_id=tenant_id, name="Sample", mode=AppMode.AGENT, enable_site=False, enable_api=False) + agent = Agent( + id=agent_id, + tenant_id=tenant_id, + app_id=app_id, + name="Sample", + scope=AgentScope.ROSTER, + source=AgentSource.AGENT_APP, + active_config_snapshot_id=snapshot_id, + active_config_is_published=False, + ) + snapshot = AgentConfigSnapshot( + id=snapshot_id, + tenant_id=tenant_id, + agent_id=agent_id, + version=1, + config_snapshot=AgentSoulConfig.model_validate( + { + "prompt": {"system_prompt": "Published prompt"}, + "model": { + "plugin_id": "langgenius/openai", + "model_provider": "langgenius/openai/openai", + "model": "gpt-4o", + "credential_ref": {"type": "provider", "id": "private-credential"}, + }, + "env": {"variables": [{"name": "TOKEN", "value": "private-token"}]}, + "config_skills": [{"name": "research", "file_id": "private-skill", "description": "Research skill"}], + "config_files": [{"name": "guide.txt", "file_kind": "upload_file", "file_id": "private-file"}], + "tools": { + "cli_tools": [ + { + "name": "search-cli", + "env": {"variables": [{"name": "TOKEN", "value": "private-cli-token"}]}, + "invoke_metadata": {"password": "private-invoke-password"}, + } + ], + "dify_tools": [ + { + "provider_type": "api", + "provider_id": "private-provider", + "tool_name": "search", + "runtime_parameters": {"password": "private-password"}, + } + ], + }, + } + ), + ) + session.add_all( + [ + app, + agent, + snapshot, + AgentConfigRevision( + tenant_id=tenant_id, + agent_id=agent_id, + current_snapshot_id=snapshot_id, + revision=1, + operation=AgentConfigRevisionOperation.PUBLISH_DRAFT, + ), + ] + ) + session.commit() + return app, agent, snapshot + + +class TestAgentTrialPreviewCases: + @staticmethod + def test_detached_preview_reference_preserves_owner_scope(sqlite_session: Session) -> None: + app, _, _ = _published_app(sqlite_session) + preview = _preview(sqlite_session, AppPreviewRef(app_id=app.id, tenant_id=app.tenant_id)) + assert preview.agent_soul.prompt.system_prompt == "Published prompt" + with pytest.raises(AppDefinitionUnavailableError, match="unavailable"): + _preview(sqlite_session, AppPreviewRef(app_id=app.id, tenant_id=str(uuid4()))) + + @staticmethod + def test_preview_reads_published_snapshot_without_creating_draft(sqlite_session: Session) -> None: + app, agent, snapshot = _published_app(sqlite_session) + before = sqlite_session.scalar(select(func.count()).select_from(AgentConfigDraft)) + + preview = _preview(sqlite_session, AppPreviewRef(app_id=app.id, tenant_id=app.tenant_id)) + + assert preview.agent_soul.prompt.system_prompt == "Published prompt" + assert preview.agent_soul.model is not None + assert preview.agent_soul.model.model == "gpt-4o" + assert preview.agent_soul.config_skills[0].name == "research" + assert preview.agent_soul.config_files[0].name == "guide.txt" + assert preview.agent_soul.tools.dify_tools[0].tool_name == "search" + assert preview.agent.id == agent.id + assert preview.active_config_snapshot is not None + assert preview.active_config_snapshot.id == snapshot.id + assert preview.active_config_is_published is True + assert preview.draft is None + assert preview.save_options == [] + assert preview.chat_endpoint is None + serialized = preview.model_dump_json() + for value in [ + "private-credential", + "private-token", + "private-password", + "private-cli-token", + "private-invoke-password", + app.tenant_id, + ]: + assert value not in serialized + assert sqlite_session.scalar(select(func.count()).select_from(AgentConfigDraft)) == before + assert not sqlite_session.new + assert not sqlite_session.dirty + + @staticmethod + def test_preview_does_not_expose_unpublished_edits(sqlite_session: Session) -> None: + app, agent, snapshot = _published_app(sqlite_session) + sqlite_session.add( + AgentConfigDraft( + tenant_id=app.tenant_id, + agent_id=agent.id, + draft_type=AgentConfigDraftType.DRAFT, + draft_owner_key="", + base_snapshot_id=snapshot.id, + config_snapshot=AgentSoulConfig.model_validate( + {"prompt": {"system_prompt": "Unpublished secret prompt"}} + ), + ) + ) + sqlite_session.commit() + + preview = _preview(sqlite_session, AppPreviewRef(app_id=app.id, tenant_id=app.tenant_id)) + + assert preview.agent_soul.prompt.system_prompt == "Published prompt" + assert "Unpublished" not in preview.model_dump_json() + editable = AgentAppComposerResponse.model_validate( + AgentComposerService.load_agent_composer( + session=sqlite_session, + tenant_id=app.tenant_id, + agent_id=agent.id, + ) + ) + assert editable.agent_soul.prompt.system_prompt == "Unpublished secret prompt" + assert editable.draft is not None + assert editable.save_options + + @staticmethod + def test_preview_rejects_seeded_import_snapshot(sqlite_session: Session) -> None: + app, agent, _ = _published_app(sqlite_session) + revision = sqlite_session.scalar(select(AgentConfigRevision).where(AgentConfigRevision.agent_id == agent.id)) + assert revision is not None + revision.operation = AgentConfigRevisionOperation.IMPORT_PACKAGE + sqlite_session.commit() + + with pytest.raises(AppDefinitionUnavailableError): + _preview(sqlite_session, AppPreviewRef(app_id=app.id, tenant_id=app.tenant_id)) + + @staticmethod + def test_preview_rejects_cross_tenant_snapshot(sqlite_session: Session) -> None: + app, _, snapshot = _published_app(sqlite_session) + snapshot.tenant_id = str(uuid4()) + sqlite_session.commit() + + with pytest.raises(AppDefinitionUnavailableError, match="unavailable"): + _preview(sqlite_session, AppPreviewRef(app_id=app.id, tenant_id=app.tenant_id)) + + +"""Preview reads enforce app ownership without adding execution or dataset ACL policy.""" + _APP_ID = "11111111-1111-1111-1111-111111111111" _TENANT_ID = "22222222-2222-2222-2222-222222222222" @@ -87,535 +272,543 @@ def _add_dataset(session: Session, *, dataset_id: str, tenant_id: str) -> None: session.add(dataset) -@pytest.fixture -def preview_ref(sqlite_session_factory: sessionmaker[Session]) -> AppPreviewRef: - with sqlite_session_factory.begin() as session: - tenant = Tenant(name="Owner") - tenant.id = _TENANT_ID - decoy_tenant = Tenant(name="Other owner") - decoy_tenant.id = _DECOY_TENANT_ID - session.add_all([tenant, decoy_tenant]) - _add_app(session, app_id=_APP_ID, tenant_id=_TENANT_ID) - _add_app(session, app_id=_DECOY_APP_ID, tenant_id=_DECOY_TENANT_ID) - _add_site(session, app_id=_APP_ID, title="Preview Site") - _add_site(session, app_id=_DECOY_APP_ID, title="Decoy Site") - _add_dataset(session, dataset_id=_DATASET_ID, tenant_id=_TENANT_ID) - _add_dataset(session, dataset_id=_DECOY_DATASET_ID, tenant_id=_DECOY_TENANT_ID) - return AppPreviewRef(app_id=_APP_ID, tenant_id=_TENANT_ID) - - -def test_get_app_returns_owner_without_requiring_publication_or_trial_registration( - sqlite_session_factory: sessionmaker[Session], preview_ref: AppPreviewRef -) -> None: - repository = AppPreviewQueryRepository(session_factory=sqlite_session_factory) - - assert repository.get_app(app_id=_APP_ID) == preview_ref - - -@pytest.mark.parametrize("state", ["missing", "not-normal"]) -def test_get_app_does_not_substitute_another_normal_app( - sqlite_session_factory: sessionmaker[Session], preview_ref: AppPreviewRef, state: str -) -> None: - app_id = preview_ref.app_id - if state == "missing": - app_id = str(uuid4()) - else: +class TestAppPreviewQueryRepositoryCases: + @staticmethod + @pytest.fixture + def preview_ref(sqlite_session_factory: sessionmaker[Session]) -> AppPreviewRef: with sqlite_session_factory.begin() as session: - # Legacy databases can contain statuses outside the current enum. - session.execute(text("UPDATE apps SET status = 'disabled' WHERE id = :app_id"), {"app_id": app_id}) - repository = AppPreviewQueryRepository(session_factory=sqlite_session_factory) + tenant = Tenant(name="Owner") + tenant.id = _TENANT_ID + decoy_tenant = Tenant(name="Other owner") + decoy_tenant.id = _DECOY_TENANT_ID + session.add_all([tenant, decoy_tenant]) + _add_app(session, app_id=_APP_ID, tenant_id=_TENANT_ID) + _add_app(session, app_id=_DECOY_APP_ID, tenant_id=_DECOY_TENANT_ID) + _add_site(session, app_id=_APP_ID, title="Preview Site") + _add_site(session, app_id=_DECOY_APP_ID, title="Decoy Site") + _add_dataset(session, dataset_id=_DATASET_ID, tenant_id=_TENANT_ID) + _add_dataset(session, dataset_id=_DECOY_DATASET_ID, tenant_id=_DECOY_TENANT_ID) + return AppPreviewRef(app_id=_APP_ID, tenant_id=_TENANT_ID) - assert repository.get_app(app_id=_DECOY_APP_ID) == AppPreviewRef(app_id=_DECOY_APP_ID, tenant_id=_DECOY_TENANT_ID) - assert repository.get_app(app_id=app_id) is None - - -def test_get_site_returns_detached_configuration_for_disabled_site_app( - sqlite_session_factory: sessionmaker[Session], preview_ref: AppPreviewRef -) -> None: - repository = AppPreviewQueryRepository(session_factory=sqlite_session_factory) - - result = repository.get_site(app=preview_ref) - - assert result is not None - assert result.owner_status == "normal" - assert result.configuration == AppSiteConfiguration( - title="Preview Site", - chat_color_theme="#000000", - chat_color_theme_inverted=True, - icon_type="image", - icon=_DATASET_ID, - icon_background="#ffffff", - description="Preview description", - copyright="Copyright", - privacy_policy="https://example.com/privacy", - input_placeholder="Ask a question", - custom_disclaimer="Disclaimer", - default_language="zh-Hans", - prompt_public=True, - show_workflow_steps=True, - use_icon_as_answer_icon=True, - ) - - -@pytest.mark.parametrize("state", ["missing-app", "wrong-tenant", "not-normal", "missing-site"]) -def test_get_site_requires_complete_admitted_app_scope( - sqlite_session_factory: sessionmaker[Session], preview_ref: AppPreviewRef, state: str -) -> None: - if state == "wrong-tenant": - preview_ref = replace(preview_ref, tenant_id=_DECOY_TENANT_ID) - else: - statements = { - "missing-app": "DELETE FROM apps WHERE id = :app_id", - "not-normal": "UPDATE apps SET status = 'disabled' WHERE id = :app_id", - "missing-site": "DELETE FROM sites WHERE app_id = :app_id", - } - with sqlite_session_factory.begin() as session: - session.execute(text(statements[state]), {"app_id": preview_ref.app_id}) - repository = AppPreviewQueryRepository(session_factory=sqlite_session_factory) - decoy_site = repository.get_site(app=AppPreviewRef(app_id=_DECOY_APP_ID, tenant_id=_DECOY_TENANT_ID)) - - assert decoy_site is not None - assert decoy_site.configuration.title == "Decoy Site" - assert repository.get_site(app=preview_ref) is None - - -@pytest.mark.parametrize("owner_state", ["missing", "archive", "normal"]) -def test_get_site_preserves_owner_state_for_service_policy( - sqlite_session_factory: sessionmaker[Session], preview_ref: AppPreviewRef, owner_state: str -) -> None: - with sqlite_session_factory.begin() as session: - owner = session.get(Tenant, preview_ref.tenant_id) - assert owner is not None - if owner_state == "missing": - session.delete(owner) - else: - owner.status = TenantStatus(owner_state) - - result = AppPreviewQueryRepository(session_factory=sqlite_session_factory).get_site(app=preview_ref) - - assert result is not None - assert result.configuration.title == "Preview Site" - assert result.owner_status == (None if owner_state == "missing" else owner_state) - - -def test_get_site_does_not_filter_legacy_empty_site_status( - sqlite_session_factory: sessionmaker[Session], preview_ref: AppPreviewRef -) -> None: - with sqlite_session_factory.begin() as session: - session.execute(text("UPDATE sites SET status = '' WHERE app_id = :app_id"), {"app_id": preview_ref.app_id}) - - result = AppPreviewQueryRepository(session_factory=sqlite_session_factory).get_site(app=preview_ref) - - assert result is not None - assert result.configuration.title == "Preview Site" - - -def test_get_datasets_filters_requested_ids_and_owner_without_acl_or_app_binding( - sqlite_session_factory: sessionmaker[Session], preview_ref: AppPreviewRef -) -> None: - with sqlite_session_factory.begin() as session: - _add_dataset(session, dataset_id=str(uuid4()), tenant_id=_TENANT_ID) - repository = AppPreviewQueryRepository(session_factory=sqlite_session_factory) - - result = repository.get_datasets(app=preview_ref, ids=[_DATASET_ID, _DECOY_DATASET_ID, _DATASET_ID, str(uuid4())]) - - assert result == ( - AppPreviewDataset( - id=_DATASET_ID, - name="Knowledge", - description="Knowledge description", - permission="only_me", - data_source_type="upload_file", - indexing_technique="high_quality", - created_by=_CREATOR_ID, - created_at=_CREATED_AT, - ), - ) - assert repository.get_datasets(app=preview_ref, ids=[]) == () - assert repository.get_datasets(app=preview_ref, ids=[str(uuid4())]) == () - - -@pytest.mark.parametrize("state", ["missing-app", "wrong-tenant", "not-normal"]) -@pytest.mark.parametrize("empty_ids", [True, False]) -def test_get_datasets_requires_available_admitted_app_even_for_empty_ids( - sqlite_session_factory: sessionmaker[Session], preview_ref: AppPreviewRef, state: str, empty_ids: bool -) -> None: - if state == "wrong-tenant": - preview_ref = replace(preview_ref, tenant_id=_DECOY_TENANT_ID) - else: - statement = ( - "DELETE FROM apps WHERE id = :app_id" - if state == "missing-app" - else "UPDATE apps SET status = 'disabled' WHERE id = :app_id" - ) - with sqlite_session_factory.begin() as session: - session.execute(text(statement), {"app_id": preview_ref.app_id}) - repository = AppPreviewQueryRepository(session_factory=sqlite_session_factory) - decoy_datasets = repository.get_datasets( - app=AppPreviewRef(app_id=_DECOY_APP_ID, tenant_id=_DECOY_TENANT_ID), ids=[_DECOY_DATASET_ID] - ) - - assert decoy_datasets is not None - assert len(decoy_datasets) == 1 - ids: list[str] = [] if empty_ids else [_DATASET_ID, _DECOY_DATASET_ID] - assert repository.get_datasets(app=preview_ref, ids=ids) is None - - -def test_get_datasets_preserves_nullable_columns(sqlite_session_factory: sessionmaker[Session]) -> None: - with sqlite_session_factory.begin() as session: - _add_app(session, app_id=_APP_ID, tenant_id=_TENANT_ID) - session.add( - Dataset(id=_DATASET_ID, tenant_id=_TENANT_ID, name="No optional configuration", created_by=_CREATOR_ID) - ) - - result = AppPreviewQueryRepository(session_factory=sqlite_session_factory).get_datasets( - app=AppPreviewRef(app_id=_APP_ID, tenant_id=_TENANT_ID), ids=[_DATASET_ID] - ) - - assert result is not None - assert len(result) == 1 - assert result[0].description is None - assert result[0].data_source_type is None - assert result[0].indexing_technique is None - - -def test_get_datasets_returns_every_requested_match_without_a_page_limit( - sqlite_session_factory: sessionmaker[Session], preview_ref: AppPreviewRef -) -> None: - ids = [str(uuid4()) for _ in range(25)] - with sqlite_session_factory.begin() as session: - for dataset_id in ids: - _add_dataset(session, dataset_id=dataset_id, tenant_id=preview_ref.tenant_id) - - result = AppPreviewQueryRepository(session_factory=sqlite_session_factory).get_datasets(app=preview_ref, ids=ids) - - assert result is not None - assert {dataset.id for dataset in result} == set(ids) - assert len(result) == 25 - - -@pytest.mark.parametrize("method", ["app", "site", "datasets"]) -@pytest.mark.parametrize("exists", [True, False]) -def test_preview_queries_release_connections_on_found_and_missing_results( - sqlite_engine: Engine, - sqlite_session_factory: sessionmaker[Session], - preview_ref: AppPreviewRef, - method: Literal["app", "site", "datasets"], - exists: bool, -) -> None: - if not exists: - preview_ref = replace(preview_ref, app_id=str(uuid4())) - repository = AppPreviewQueryRepository(session_factory=sqlite_session_factory) - assert isinstance(sqlite_engine.pool, QueuePool) - assert sqlite_engine.pool.checkedout() == 0 - - result: object - match method: - case "app": - result = repository.get_app(app_id=preview_ref.app_id) - case "site": - result = repository.get_site(app=preview_ref) - case "datasets": - result = repository.get_datasets(app=preview_ref, ids=[_DATASET_ID]) - - assert (result is not None) == exists - assert sqlite_engine.pool.checkedout() == 0 - - -def test_preview_queries_leave_caller_transaction_open_and_uncommitted( - sqlite_engine: Engine, sqlite_session_factory: sessionmaker[Session], preview_ref: AppPreviewRef -) -> None: - repository = AppPreviewQueryRepository(session_factory=sqlite_session_factory) - assert isinstance(sqlite_engine.pool, QueuePool) - with sqlite_session_factory() as caller: - app = caller.get(App, _APP_ID) - assert app is not None - app.name = "Pending caller change" - assert sqlite_engine.pool.checkedout() == 1 + @staticmethod + def test_get_app_returns_owner_without_requiring_publication_or_trial_registration( + sqlite_session_factory: sessionmaker[Session], preview_ref: AppPreviewRef + ) -> None: + repository = AppPreviewQueryRepository(session_factory=sqlite_session_factory) assert repository.get_app(app_id=_APP_ID) == preview_ref - assert repository.get_site(app=preview_ref) is not None - assert repository.get_datasets(app=preview_ref, ids=[_DATASET_ID]) is not None - assert caller.in_transaction() - assert app in caller.dirty - assert sqlite_engine.pool.checkedout() == 1 - caller.rollback() - - with sqlite_session_factory() as session: - assert session.scalar(select(App.name).where(App.id == _APP_ID)) == "Preview" - assert sqlite_engine.pool.checkedout() == 0 - - -@pytest.fixture -def detail_config(sqlite_session_factory: sessionmaker[Session], preview_ref: AppPreviewRef) -> AppModelConfig: - config = AppModelConfig( - app_id=preview_ref.app_id, - opening_statement="Stored opening", - agent_mode=json.dumps( - { - "enabled": True, - "strategy": "react", - "tools": [ - { - "provider_type": "builtin", - "provider_id": "acme/search/search", - "tool_name": "search", - "tool_parameters": {"api_key": "encrypted-tool-secret"}, - } - ], - } - ), - ) - account = Account(name="Viewer", email="preview@example.com") - account.id = _CREATOR_ID - with sqlite_session_factory.begin() as session: - session.add_all([account, config]) - session.flush() - app = session.get(App, preview_ref.app_id) - assert app is not None - app.app_model_config_id = config.id - return config - - -def test_get_detail_returns_read_only_detached_values_before_mode_and_tool_enrichment( - sqlite_engine: Engine, - sqlite_session_factory: sessionmaker[Session], - preview_ref: AppPreviewRef, - detail_config: AppModelConfig, -) -> None: - closed: list[Session] = [] - mutated: list[object] = [] - writes: list[str] = [] - - class TrackedSession(Session): - @override - def close(self) -> None: - mutated.extend(self.new) - mutated.extend(self.dirty) - mutated.extend(self.deleted) - super().close() - closed.append(self) - - def record_writes( - _connection: Connection, - _cursor: object, - statement: str, - _parameters: object, - _context: ExecutionContext, - _executemany: bool, + @staticmethod + @pytest.mark.parametrize("state", ["missing", "not-normal"]) + def test_get_app_does_not_substitute_another_normal_app( + sqlite_session_factory: sessionmaker[Session], preview_ref: AppPreviewRef, state: str ) -> None: - if statement.lstrip().split(maxsplit=1)[0].upper() in {"INSERT", "UPDATE", "DELETE"}: - writes.append(statement) + app_id = preview_ref.app_id + if state == "missing": + app_id = str(uuid4()) + else: + with sqlite_session_factory.begin() as session: + # Legacy databases can contain statuses outside the current enum. + session.execute(text("UPDATE apps SET status = 'disabled' WHERE id = :app_id"), {"app_id": app_id}) + repository = AppPreviewQueryRepository(session_factory=sqlite_session_factory) - factory: sessionmaker[Session] = sessionmaker(bind=sqlite_engine, class_=TrackedSession, expire_on_commit=False) - repository = AppPreviewQueryRepository(session_factory=factory) - event.listen(sqlite_engine, "before_cursor_execute", record_writes) - try: - record = repository.get_detail(app=preview_ref, account_id=_CREATOR_ID) - finally: - event.remove(sqlite_engine, "before_cursor_execute", record_writes) + assert repository.get_app(app_id=_DECOY_APP_ID) == AppPreviewRef( + app_id=_DECOY_APP_ID, tenant_id=_DECOY_TENANT_ID + ) + assert repository.get_app(app_id=app_id) is None - assert isinstance(sqlite_engine.pool, QueuePool) - assert sqlite_engine.pool.checkedout() == 0 - assert closed - assert all(not session.in_transaction() and not session.identity_map for session in closed) - assert mutated == [] - assert writes == [] - assert record.detail.id == preview_ref.app_id - assert record.detail.mode == AppMode.CHAT - assert record.detail.site.title == "Preview Site" - assert record.detail.model_config is not None - assert record.detail.model_config["opening_statement"] == "Stored opening" - assert record.detail.model_config["agent_mode"] == json.loads(detail_config.agent_mode or "{}") - assert record.detail.deleted_tools == () - assert record.existing_api_provider_ids == frozenset() + @staticmethod + def test_get_site_returns_detached_configuration_for_disabled_site_app( + sqlite_session_factory: sessionmaker[Session], preview_ref: AppPreviewRef + ) -> None: + repository = AppPreviewQueryRepository(session_factory=sqlite_session_factory) - def serialize_date(value: object) -> str: - assert isinstance(value, datetime), f"Non-data value escaped repository: {type(value)}" - return value.isoformat() + result = repository.get_site(app=preview_ref) - json.dumps(asdict(record.detail), default=serialize_date) - with sqlite_session_factory() as session: - assert session.scalar(select(App.mode).where(App.id == preview_ref.app_id)) == AppMode.CHAT - assert session.scalar(select(AppModelConfig.agent_mode).where(AppModelConfig.id == detail_config.id)) == ( - detail_config.agent_mode + assert result is not None + assert result.owner_status == "normal" + assert result.configuration == AppSiteConfiguration( + title="Preview Site", + chat_color_theme="#000000", + chat_color_theme_inverted=True, + icon_type="image", + icon=_DATASET_ID, + icon_background="#ffffff", + description="Preview description", + copyright="Copyright", + privacy_policy="https://example.com/privacy", + input_placeholder="Ask a question", + custom_disclaimer="Disclaimer", + default_language="zh-Hans", + prompt_public=True, + show_workflow_steps=True, + use_icon_as_answer_icon=True, ) + @staticmethod + @pytest.mark.parametrize("state", ["missing-app", "wrong-tenant", "not-normal", "missing-site"]) + def test_get_site_requires_complete_admitted_app_scope( + sqlite_session_factory: sessionmaker[Session], preview_ref: AppPreviewRef, state: str + ) -> None: + if state == "wrong-tenant": + preview_ref = replace(preview_ref, tenant_id=_DECOY_TENANT_ID) + else: + statements = { + "missing-app": "DELETE FROM apps WHERE id = :app_id", + "not-normal": "UPDATE apps SET status = 'disabled' WHERE id = :app_id", + "missing-site": "DELETE FROM sites WHERE app_id = :app_id", + } + with sqlite_session_factory.begin() as session: + session.execute(text(statements[state]), {"app_id": preview_ref.app_id}) + repository = AppPreviewQueryRepository(session_factory=sqlite_session_factory) + decoy_site = repository.get_site(app=AppPreviewRef(app_id=_DECOY_APP_ID, tenant_id=_DECOY_TENANT_ID)) -def test_get_detail_resolves_only_configured_api_providers_owned_by_app_tenant( - sqlite_session_factory: sessionmaker[Session], preview_ref: AppPreviewRef, detail_config: AppModelConfig -) -> None: - owner_provider, viewer_provider, unconfigured_provider = [ - ApiToolProvider( - name=name, - icon="icon", - schema="{}", - schema_type_str=ApiProviderSchemaType.OPENAPI, - user_id=_CREATOR_ID, - tenant_id=tenant_id, - description=name, - tools_str="[]", - credentials_str="{}", + assert decoy_site is not None + assert decoy_site.configuration.title == "Decoy Site" + assert repository.get_site(app=preview_ref) is None + + @staticmethod + @pytest.mark.parametrize("owner_state", ["missing", "archive", "normal"]) + def test_get_site_preserves_owner_state_for_service_policy( + sqlite_session_factory: sessionmaker[Session], preview_ref: AppPreviewRef, owner_state: str + ) -> None: + with sqlite_session_factory.begin() as session: + owner = session.get(Tenant, preview_ref.tenant_id) + assert owner is not None + if owner_state == "missing": + session.delete(owner) + else: + owner.status = TenantStatus(owner_state) + + result = AppPreviewQueryRepository(session_factory=sqlite_session_factory).get_site(app=preview_ref) + + assert result is not None + assert result.configuration.title == "Preview Site" + assert result.owner_status == (None if owner_state == "missing" else owner_state) + + @staticmethod + def test_get_site_does_not_filter_legacy_empty_site_status( + sqlite_session_factory: sessionmaker[Session], preview_ref: AppPreviewRef + ) -> None: + with sqlite_session_factory.begin() as session: + session.execute(text("UPDATE sites SET status = '' WHERE app_id = :app_id"), {"app_id": preview_ref.app_id}) + + result = AppPreviewQueryRepository(session_factory=sqlite_session_factory).get_site(app=preview_ref) + + assert result is not None + assert result.configuration.title == "Preview Site" + + @staticmethod + def test_get_datasets_filters_requested_ids_and_owner_without_acl_or_app_binding( + sqlite_session_factory: sessionmaker[Session], preview_ref: AppPreviewRef + ) -> None: + with sqlite_session_factory.begin() as session: + _add_dataset(session, dataset_id=str(uuid4()), tenant_id=_TENANT_ID) + repository = AppPreviewQueryRepository(session_factory=sqlite_session_factory) + + result = repository.get_datasets( + app=preview_ref, ids=[_DATASET_ID, _DECOY_DATASET_ID, _DATASET_ID, str(uuid4())] ) - for name, tenant_id in [ - ("Owner provider", preview_ref.tenant_id), - ("Viewer provider", _DECOY_TENANT_ID), - ("Unconfigured provider", preview_ref.tenant_id), + + assert result == ( + AppPreviewDataset( + id=_DATASET_ID, + name="Knowledge", + description="Knowledge description", + permission="only_me", + data_source_type="upload_file", + indexing_technique="high_quality", + created_by=_CREATOR_ID, + created_at=_CREATED_AT, + ), + ) + assert repository.get_datasets(app=preview_ref, ids=[]) == () + assert repository.get_datasets(app=preview_ref, ids=[str(uuid4())]) == () + + @staticmethod + @pytest.mark.parametrize("state", ["missing-app", "wrong-tenant", "not-normal"]) + @pytest.mark.parametrize("empty_ids", [True, False]) + def test_get_datasets_requires_available_admitted_app_even_for_empty_ids( + sqlite_session_factory: sessionmaker[Session], preview_ref: AppPreviewRef, state: str, empty_ids: bool + ) -> None: + if state == "wrong-tenant": + preview_ref = replace(preview_ref, tenant_id=_DECOY_TENANT_ID) + else: + statement = ( + "DELETE FROM apps WHERE id = :app_id" + if state == "missing-app" + else "UPDATE apps SET status = 'disabled' WHERE id = :app_id" + ) + with sqlite_session_factory.begin() as session: + session.execute(text(statement), {"app_id": preview_ref.app_id}) + repository = AppPreviewQueryRepository(session_factory=sqlite_session_factory) + decoy_datasets = repository.get_datasets( + app=AppPreviewRef(app_id=_DECOY_APP_ID, tenant_id=_DECOY_TENANT_ID), ids=[_DECOY_DATASET_ID] + ) + + assert decoy_datasets is not None + assert len(decoy_datasets) == 1 + ids: list[str] = [] if empty_ids else [_DATASET_ID, _DECOY_DATASET_ID] + assert repository.get_datasets(app=preview_ref, ids=ids) is None + + @staticmethod + def test_get_datasets_preserves_nullable_columns(sqlite_session_factory: sessionmaker[Session]) -> None: + with sqlite_session_factory.begin() as session: + _add_app(session, app_id=_APP_ID, tenant_id=_TENANT_ID) + session.add( + Dataset(id=_DATASET_ID, tenant_id=_TENANT_ID, name="No optional configuration", created_by=_CREATOR_ID) + ) + + result = AppPreviewQueryRepository(session_factory=sqlite_session_factory).get_datasets( + app=AppPreviewRef(app_id=_APP_ID, tenant_id=_TENANT_ID), ids=[_DATASET_ID] + ) + + assert result is not None + assert len(result) == 1 + assert result[0].description is None + assert result[0].data_source_type is None + assert result[0].indexing_technique is None + + @staticmethod + def test_get_datasets_returns_every_requested_match_without_a_page_limit( + sqlite_session_factory: sessionmaker[Session], preview_ref: AppPreviewRef + ) -> None: + ids = [str(uuid4()) for _ in range(25)] + with sqlite_session_factory.begin() as session: + for dataset_id in ids: + _add_dataset(session, dataset_id=dataset_id, tenant_id=preview_ref.tenant_id) + + result = AppPreviewQueryRepository(session_factory=sqlite_session_factory).get_datasets( + app=preview_ref, ids=ids + ) + + assert result is not None + assert {dataset.id for dataset in result} == set(ids) + assert len(result) == 25 + + @staticmethod + @pytest.mark.parametrize("method", ["app", "site", "datasets"]) + @pytest.mark.parametrize("exists", [True, False]) + def test_preview_queries_release_connections_on_found_and_missing_results( + sqlite_engine: Engine, + sqlite_session_factory: sessionmaker[Session], + preview_ref: AppPreviewRef, + method: Literal["app", "site", "datasets"], + exists: bool, + ) -> None: + if not exists: + preview_ref = replace(preview_ref, app_id=str(uuid4())) + repository = AppPreviewQueryRepository(session_factory=sqlite_session_factory) + assert isinstance(sqlite_engine.pool, QueuePool) + assert sqlite_engine.pool.checkedout() == 0 + + result: object + match method: + case "app": + result = repository.get_app(app_id=preview_ref.app_id) + case "site": + result = repository.get_site(app=preview_ref) + case "datasets": + result = repository.get_datasets(app=preview_ref, ids=[_DATASET_ID]) + + assert (result is not None) == exists + assert sqlite_engine.pool.checkedout() == 0 + + @staticmethod + def test_preview_queries_leave_caller_transaction_open_and_uncommitted( + sqlite_engine: Engine, sqlite_session_factory: sessionmaker[Session], preview_ref: AppPreviewRef + ) -> None: + repository = AppPreviewQueryRepository(session_factory=sqlite_session_factory) + assert isinstance(sqlite_engine.pool, QueuePool) + with sqlite_session_factory() as caller: + app = caller.get(App, _APP_ID) + assert app is not None + app.name = "Pending caller change" + assert sqlite_engine.pool.checkedout() == 1 + + assert repository.get_app(app_id=_APP_ID) == preview_ref + assert repository.get_site(app=preview_ref) is not None + assert repository.get_datasets(app=preview_ref, ids=[_DATASET_ID]) is not None + + assert caller.in_transaction() + assert app in caller.dirty + assert sqlite_engine.pool.checkedout() == 1 + caller.rollback() + + with sqlite_session_factory() as session: + assert session.scalar(select(App.name).where(App.id == _APP_ID)) == "Preview" + assert sqlite_engine.pool.checkedout() == 0 + + @staticmethod + @pytest.fixture + def detail_config(sqlite_session_factory: sessionmaker[Session], preview_ref: AppPreviewRef) -> AppModelConfig: + config = AppModelConfig( + app_id=preview_ref.app_id, + opening_statement="Stored opening", + agent_mode=json.dumps( + { + "enabled": True, + "strategy": "react", + "tools": [ + { + "provider_type": "builtin", + "provider_id": "acme/search/search", + "tool_name": "search", + "tool_parameters": {"api_key": "encrypted-tool-secret"}, + } + ], + } + ), + ) + account = Account(name="Viewer", email="preview@example.com") + account.id = _CREATOR_ID + with sqlite_session_factory.begin() as session: + session.add_all([account, config]) + session.flush() + app = session.get(App, preview_ref.app_id) + assert app is not None + app.app_model_config_id = config.id + return config + + @staticmethod + def test_get_detail_returns_read_only_detached_values_before_mode_and_tool_enrichment( + sqlite_engine: Engine, + sqlite_session_factory: sessionmaker[Session], + preview_ref: AppPreviewRef, + detail_config: AppModelConfig, + ) -> None: + closed: list[Session] = [] + mutated: list[object] = [] + writes: list[str] = [] + + class TrackedSession(Session): + @override + def close(self) -> None: + mutated.extend(self.new) + mutated.extend(self.dirty) + mutated.extend(self.deleted) + super().close() + closed.append(self) + + def record_writes( + _connection: Connection, + _cursor: object, + statement: str, + _parameters: object, + _context: ExecutionContext, + _executemany: bool, + ) -> None: + if statement.lstrip().split(maxsplit=1)[0].upper() in {"INSERT", "UPDATE", "DELETE"}: + writes.append(statement) + + factory: sessionmaker[Session] = sessionmaker(bind=sqlite_engine, class_=TrackedSession, expire_on_commit=False) + repository = AppPreviewQueryRepository(session_factory=factory) + event.listen(sqlite_engine, "before_cursor_execute", record_writes) + try: + record = repository.get_detail(app=preview_ref, account_id=_CREATOR_ID) + finally: + event.remove(sqlite_engine, "before_cursor_execute", record_writes) + + assert isinstance(sqlite_engine.pool, QueuePool) + assert sqlite_engine.pool.checkedout() == 0 + assert closed + assert all(not session.in_transaction() and not session.identity_map for session in closed) + assert mutated == [] + assert writes == [] + assert record.detail.id == preview_ref.app_id + assert record.detail.mode == AppMode.CHAT + assert record.detail.site.title == "Preview Site" + assert record.detail.model_config is not None + assert record.detail.model_config["opening_statement"] == "Stored opening" + assert record.detail.model_config["agent_mode"] == json.loads(detail_config.agent_mode or "{}") + assert record.detail.deleted_tools == () + assert record.existing_api_provider_ids == frozenset() + + def serialize_date(value: object) -> str: + assert isinstance(value, datetime), f"Non-data value escaped repository: {type(value)}" + return value.isoformat() + + json.dumps(asdict(record.detail), default=serialize_date) + with sqlite_session_factory() as session: + assert session.scalar(select(App.mode).where(App.id == preview_ref.app_id)) == AppMode.CHAT + assert session.scalar(select(AppModelConfig.agent_mode).where(AppModelConfig.id == detail_config.id)) == ( + detail_config.agent_mode + ) + + @staticmethod + def test_get_detail_resolves_only_configured_api_providers_owned_by_app_tenant( + sqlite_session_factory: sessionmaker[Session], preview_ref: AppPreviewRef, detail_config: AppModelConfig + ) -> None: + owner_provider, viewer_provider, unconfigured_provider = [ + ApiToolProvider( + name=name, + icon="icon", + schema="{}", + schema_type_str=ApiProviderSchemaType.OPENAPI, + user_id=_CREATOR_ID, + tenant_id=tenant_id, + description=name, + tools_str="[]", + credentials_str="{}", + ) + for name, tenant_id in [ + ("Owner provider", preview_ref.tenant_id), + ("Viewer provider", _DECOY_TENANT_ID), + ("Unconfigured provider", preview_ref.tenant_id), + ] ] - ] - with sqlite_session_factory.begin() as session: - session.add_all([owner_provider, viewer_provider, unconfigured_provider]) - config = session.get(AppModelConfig, detail_config.id) - assert config is not None - config.agent_mode = json.dumps( + with sqlite_session_factory.begin() as session: + session.add_all([owner_provider, viewer_provider, unconfigured_provider]) + config = session.get(AppModelConfig, detail_config.id) + assert config is not None + config.agent_mode = json.dumps( + { + "enabled": False, + "tools": [ + { + "provider_type": "api", + "provider_id": provider_id, + "tool_name": "search", + "tool_parameters": {}, + } + for provider_id in [owner_provider.id, viewer_provider.id, owner_provider.id, str(uuid4())] + ], + } + ) + + record = AppPreviewQueryRepository(session_factory=sqlite_session_factory).get_detail( + app=preview_ref, account_id=_CREATOR_ID + ) + + assert record.existing_api_provider_ids == frozenset({owner_provider.id}) + assert record.detail.deleted_tools == () + + @staticmethod + @pytest.mark.parametrize( + "stored_environment", + [ + pytest.param("{invalid-json", id="unparsed-invalid-json"), + pytest.param( + '{"secret":{"id":"secret","name":"api_key","value_type":"secret",' + '"value":"encrypted-workflow-secret","selector":[]}}', + id="undecrypted-secret", + ), + ], + ) + def test_get_workflow_returns_detached_data_and_raw_environment_after_session_closes( + sqlite_engine: Engine, + sqlite_session_factory: sessionmaker[Session], + preview_ref: AppPreviewRef, + stored_environment: str, + ) -> None: + author = Account(name="Author", email="author@example.com") + author.id = _CREATOR_ID + updater = Account(name="Updater", email="updater@example.com") + stored_features = json.dumps( { - "enabled": False, - "tools": [ - { - "provider_type": "api", - "provider_id": provider_id, - "tool_name": "search", - "tool_parameters": {}, - } - for provider_id in [owner_provider.id, viewer_provider.id, owner_provider.id, str(uuid4())] - ], + "opening_statement": "Workflow opening", + "file_upload": {"image": {"enabled": True, "number_limits": 3, "transfer_methods": ["local_file"]}}, } ) - - record = AppPreviewQueryRepository(session_factory=sqlite_session_factory).get_detail( - app=preview_ref, account_id=_CREATOR_ID - ) - - assert record.existing_api_provider_ids == frozenset({owner_provider.id}) - assert record.detail.deleted_tools == () - - -@pytest.mark.parametrize( - "stored_environment", - [ - pytest.param("{invalid-json", id="unparsed-invalid-json"), - pytest.param( - '{"secret":{"id":"secret","name":"api_key","value_type":"secret",' - '"value":"encrypted-workflow-secret","selector":[]}}', - id="undecrypted-secret", - ), - ], -) -def test_get_workflow_returns_detached_data_and_raw_environment_after_session_closes( - sqlite_engine: Engine, - sqlite_session_factory: sessionmaker[Session], - preview_ref: AppPreviewRef, - stored_environment: str, -) -> None: - author = Account(name="Author", email="author@example.com") - author.id = _CREATOR_ID - updater = Account(name="Updater", email="updater@example.com") - stored_features = json.dumps( - { - "opening_statement": "Workflow opening", - "file_upload": {"image": {"enabled": True, "number_limits": 3, "transfer_methods": ["local_file"]}}, - } - ) - with sqlite_session_factory.begin() as session: - session.add_all([author, updater]) - session.flush() - workflow = Workflow( - tenant_id=preview_ref.tenant_id, - app_id=preview_ref.app_id, - type=WorkflowType.CHAT, - version="published", - graph='{"nodes":[],"edges":[]}', - features=stored_features, - created_by=author.id, - updated_by=updater.id, - marked_name="Release", - marked_comment="Preview snapshot", - ) - workflow.created_at = _CREATED_AT - workflow.updated_at = _CREATED_AT - workflow._environment_variables = stored_environment - workflow.conversation_variables = [ - StringVariable(id="topic", name="topic", value="sqlite", selector=["conversation", "topic"]) - ] - session.add(workflow) - session.flush() - app = session.get(App, preview_ref.app_id) - assert app is not None - app.workflow_id = workflow.id - session.add( - WorkflowToolProvider( - name="preview-workflow", - label="Preview workflow", - icon="icon", - app_id=preview_ref.app_id, - version="1.0.0", - user_id=author.id, + with sqlite_session_factory.begin() as session: + session.add_all([author, updater]) + session.flush() + workflow = Workflow( tenant_id=preview_ref.tenant_id, - description="Workflow provider", - parameter_configuration="[]", + app_id=preview_ref.app_id, + type=WorkflowType.CHAT, + version="published", + graph='{"nodes":[],"edges":[]}', + features=stored_features, + created_by=author.id, + updated_by=updater.id, + marked_name="Release", + marked_comment="Preview snapshot", + ) + workflow.created_at = _CREATED_AT + workflow.updated_at = _CREATED_AT + workflow._environment_variables = stored_environment + workflow.conversation_variables = [ + StringVariable(id="topic", name="topic", value="sqlite", selector=["conversation", "topic"]) + ] + session.add(workflow) + session.flush() + app = session.get(App, preview_ref.app_id) + assert app is not None + app.workflow_id = workflow.id + session.add( + WorkflowToolProvider( + name="preview-workflow", + label="Preview workflow", + icon="icon", + app_id=preview_ref.app_id, + version="1.0.0", + user_id=author.id, + tenant_id=preview_ref.tenant_id, + description="Workflow provider", + parameter_configuration="[]", + ) ) - ) - closed: list[Session] = [] - mutated: list[object] = [] + closed: list[Session] = [] + mutated: list[object] = [] - class TrackedSession(Session): - @override - def close(self) -> None: - mutated.extend(self.new) - mutated.extend(self.dirty) - mutated.extend(self.deleted) - super().close() - closed.append(self) + class TrackedSession(Session): + @override + def close(self) -> None: + mutated.extend(self.new) + mutated.extend(self.dirty) + mutated.extend(self.deleted) + super().close() + closed.append(self) - factory: sessionmaker[Session] = sessionmaker(bind=sqlite_engine, class_=TrackedSession, expire_on_commit=False) - record = AppPreviewQueryRepository(session_factory=factory).get_workflow(app=preview_ref) + factory: sessionmaker[Session] = sessionmaker(bind=sqlite_engine, class_=TrackedSession, expire_on_commit=False) + record = AppPreviewQueryRepository(session_factory=factory).get_workflow(app=preview_ref) - assert isinstance(sqlite_engine.pool, QueuePool) - assert sqlite_engine.pool.checkedout() == 0 - assert closed - assert all(not session.in_transaction() and not session.identity_map for session in closed) - assert mutated == [] - assert record.tenant_id == preview_ref.tenant_id - assert record.environment_variables_json == stored_environment - result = record.workflow - assert result.environment_variables == () - assert result.id == workflow.id - assert result.graph == {"nodes": [], "edges": []} - assert result.features == { - "opening_statement": "Workflow opening", - "file_upload": { - "enabled": True, - "number_limits": 3, - "allowed_file_upload_methods": ["local_file"], - "allowed_file_types": ["image"], - "allowed_file_extensions": [], - }, - } - assert result.hash == workflow.unique_hash - assert result.version == "published" - assert result.marked_name == "Release" - assert result.marked_comment == "Preview snapshot" - assert result.created_by == AppPreviewAccount(id=author.id, name="Author", email="author@example.com") - assert result.updated_by == AppPreviewAccount(id=updater.id, name="Updater", email="updater@example.com") - assert result.created_at == _CREATED_AT - assert result.updated_at == _CREATED_AT - assert result.tool_published is True - assert len(result.conversation_variables) == 1 - assert result.conversation_variables[0]["name"] == "topic" - assert result.conversation_variables[0]["value"] == "sqlite" - assert result.conversation_variables[0]["value_type"] == "string" - assert result.rag_pipeline_variables == () + assert isinstance(sqlite_engine.pool, QueuePool) + assert sqlite_engine.pool.checkedout() == 0 + assert closed + assert all(not session.in_transaction() and not session.identity_map for session in closed) + assert mutated == [] + assert record.tenant_id == preview_ref.tenant_id + assert record.environment_variables_json == stored_environment + result = record.workflow + assert result.environment_variables == () + assert result.id == workflow.id + assert result.graph == {"nodes": [], "edges": []} + assert result.features == { + "opening_statement": "Workflow opening", + "file_upload": { + "enabled": True, + "number_limits": 3, + "allowed_file_upload_methods": ["local_file"], + "allowed_file_types": ["image"], + "allowed_file_extensions": [], + }, + } + assert result.hash == workflow.unique_hash + assert result.version == "published" + assert result.marked_name == "Release" + assert result.marked_comment == "Preview snapshot" + assert result.created_by == AppPreviewAccount(id=author.id, name="Author", email="author@example.com") + assert result.updated_by == AppPreviewAccount(id=updater.id, name="Updater", email="updater@example.com") + assert result.created_at == _CREATED_AT + assert result.updated_at == _CREATED_AT + assert result.tool_published is True + assert len(result.conversation_variables) == 1 + assert result.conversation_variables[0]["name"] == "topic" + assert result.conversation_variables[0]["value"] == "sqlite" + assert result.conversation_variables[0]["value_type"] == "string" + assert result.rag_pipeline_variables == () - def serialize_date(value: object) -> str: - assert isinstance(value, datetime), f"Non-data value escaped repository: {type(value)}" - return value.isoformat() + def serialize_date(value: object) -> str: + assert isinstance(value, datetime), f"Non-data value escaped repository: {type(value)}" + return value.isoformat() - json.dumps(asdict(record), default=serialize_date) - with sqlite_session_factory() as session: - assert session.scalar(select(Workflow._features).where(Workflow.id == workflow.id)) == stored_features - assert ( - session.scalar(select(Workflow._environment_variables).where(Workflow.id == workflow.id)) - == stored_environment - ) + json.dumps(asdict(record), default=serialize_date) + with sqlite_session_factory() as session: + assert session.scalar(select(Workflow._features).where(Workflow.id == workflow.id)) == stored_features + assert ( + session.scalar(select(Workflow._environment_variables).where(Workflow.id == workflow.id)) + == stored_environment + ) diff --git a/api/tests/unit_tests/repositories/test_end_user_repository.py b/api/tests/unit_tests/repositories/test_app_scoped_end_user_repository.py similarity index 100% rename from api/tests/unit_tests/repositories/test_end_user_repository.py rename to api/tests/unit_tests/repositories/test_app_scoped_end_user_repository.py diff --git a/api/tests/unit_tests/repositories/test_application_trace_repository.py b/api/tests/unit_tests/repositories/test_application_trace_repository.py new file mode 100644 index 00000000000..d9eb6469526 --- /dev/null +++ b/api/tests/unit_tests/repositories/test_application_trace_repository.py @@ -0,0 +1,630 @@ +import json +from collections.abc import Iterator +from dataclasses import dataclass +from datetime import timedelta +from typing import Literal +from uuid import UUID + +import pytest +from aliyun.log.logclient import LogClient +from sqlalchemy import Connection, Engine, delete, event, func, select +from sqlalchemy.orm import Session, sessionmaker + +from core.ops.entities.trace_entity import ModerationTraceInfo +from extensions.logstore.aliyun_logstore import AliyunLogStore +from extensions.logstore.repositories.logstore_api_workflow_run_repository import _dict_to_workflow_run +from libs.datetime_utils import naive_utc_now +from models.application_trace import ApplicationTraceBatch, ApplicationTraceInvocation, ApplicationTraceRun +from models.model import App, AppMode +from models.workflow import WorkflowNodeExecutionModel, WorkflowRun +from repositories.application_trace_repository import TraceCollectionError +from services.tracing.provider_dispatch import TraceExportResult +from services.tracing.recovery import _workflow_task +from services.tracing.service import ApplicationTracingService +from tests.helpers.unit.config_override import apply_config_overrides +from tests.helpers.unit.model_factories import make_workflow +from tests.unit_tests.services.tracing.test_runtime_bridge import AGENT_GRAPH, workflow_info +from tests.unit_tests.services.tracing.test_service import _application_trace_collection_batch_for as batch_for +from tests.unit_tests.services.tracing.test_service import _application_trace_collection_claim as claim +from tests.unit_tests.services.tracing.test_service import _application_trace_collection_identifier as identifier +from tests.unit_tests.services.tracing.test_service import _application_trace_collection_ingest as ingest +from tests.unit_tests.services.tracing.test_service import _application_trace_collection_seed as seed +from tests.unit_tests.services.tracing.test_service import _application_trace_collection_terminal as terminal +from tests.unit_tests.services.tracing.test_service import _ApplicationTraceCollectionCollection as Collection + +"""Configured LogStore readers work without SQL mirrors and retain source ownership.""" + + +@dataclass +class StoredSource: + case: Collection + kind: Literal["workflow", "node"] + row: dict[str, object] + requests: list[str] + + def start(self) -> ApplicationTraceRun: + run = self.case.repository.start( + tenant_id=self.case.run.tenant_id, + app_id=self.case.run.app_id, + source_kind=self.kind, + source_id=str(self.row["id"]), + session_id="logstore-session", + ) + assert run is not None + return run + + def service(self) -> ApplicationTracingService: + return ApplicationTracingService(self.case.repository) + + def status(self, node_execution_id: str | None = None) -> dict[str, object]: + return self.service().status( + tenant_id=self.case.run.tenant_id, + app_id=self.case.run.app_id, + source_kind=self.kind, + source_id=str(self.row["id"]), + node_execution_id=node_execution_id, + ) + + +class TestApplicationTraceLogstoreCases: + @staticmethod + @pytest.fixture(params=["workflow", "node"]) + def stored( + request: pytest.FixtureRequest, + sqlite_engine: Engine, + sqlite_session_factory: sessionmaker[Session], + monkeypatch: pytest.MonkeyPatch, + ) -> Iterator[StoredSource]: + kind = request.param + assert kind in {"workflow", "node"} + case = seed(sqlite_session_factory, "logstore") + with case.sessions.begin() as session: + app = session.get(App, case.run.app_id) + assert app is not None + app.mode = AppMode.WORKFLOW + now = naive_utc_now() + row: dict[str, object] = { + "id": identifier("logstore-source"), + "tenant_id": case.run.tenant_id, + "app_id": case.run.app_id, + "workflow_id": identifier("definition"), + "created_by": identifier("owner"), + "created_at": (now - timedelta(seconds=30)).isoformat(), + "finished_at": (now - timedelta(seconds=10)).isoformat(), + "status": "succeeded", + "elapsed_time": 20, + "inputs": json.dumps({"query": "private input"}), + "outputs": json.dumps({"result": "private output"}), + "node_id": "agent-node", + "node_type": "agent", + "title": "Agent preview", + "graph": json.dumps(AGENT_GRAPH), + } + apply_config_overrides( + monkeypatch, + LOGSTORE_DUAL_WRITE_ENABLED=False, + API_WORKFLOW_RUN_REPOSITORY=( + "extensions.logstore.repositories.logstore_api_workflow_run_repository.LogstoreAPIWorkflowRunRepository" + ), + API_WORKFLOW_NODE_EXECUTION_REPOSITORY=( + "extensions.logstore.repositories.logstore_api_workflow_node_execution_repository." + "LogstoreAPIWorkflowNodeExecutionRepository" + ), + ) + monkeypatch.setenv("LOGSTORE_DUAL_READ_ENABLED", "false") + monkeypatch.setenv("LOGSTORE_PG_MODE_ENABLED", "false") + monkeypatch.setenv("ALIYUN_SLS_ACCESS_KEY_ID", "test-access-key") + monkeypatch.setenv("ALIYUN_SLS_ACCESS_KEY_SECRET", "test-secret") + monkeypatch.setenv("ALIYUN_SLS_ENDPOINT", "https://sls.example.test") + monkeypatch.setenv("ALIYUN_SLS_REGION", "test-region") + monkeypatch.setattr(AliyunLogStore, "_instance", None) + monkeypatch.setattr(AliyunLogStore, "_initialized", False) + monkeypatch.setattr( + AliyunLogStore, "_check_endpoint_connectivity", staticmethod(lambda _endpoint, _timeout: None) + ) + + def reject_http_transport(*_args: object, **_kwargs: object) -> None: + pytest.fail("Configured LogStore test escaped its isolated transport boundary") + + monkeypatch.setattr(LogClient, "_getHttpResponse", reject_http_transport) + transactions: set[Connection] = set() + + def begun(connection: Connection) -> None: + transactions.add(connection) + + def ended(connection: Connection) -> None: + transactions.discard(connection) + + requests: list[str] = [] + + def get_logs( + _client: AliyunLogStore, *, query: str, logstore: str, **_kwargs: object + ) -> list[dict[str, object]]: + # Exercise real configured repository construction, scoped query generation, + # and model decoding; replace only the remote transport boundary. + assert not transactions, "LogStore I/O must not hold a SQL transaction" + assert f'tenant_id:"{case.run.tenant_id}"' in query.replace(": ", ":") + assert logstore == ( + AliyunLogStore.workflow_execution_logstore + if kind == "workflow" + else AliyunLogStore.workflow_node_execution_logstore + ) + if kind == "workflow": + assert f'app_id:"{case.run.app_id}"' in query + requests.append(query) + return [dict(row)] if row else [] + + monkeypatch.setattr(AliyunLogStore, "get_logs", get_logs) + + def execute_sql( + client: AliyunLogStore, *, sql: str, query: str, logstore: str, **_kwargs: object + ) -> list[dict[str, object]]: + # The node reader now uses SQL in both SDK and PG modes. Preserve the + # real repository's query/decoding path while isolating this transport. + assert kind == "node" + assert f'id:"{identifier("logstore-source")}"' in query.replace(": ", ":") + assert f"tenant_id = '{case.run.tenant_id}'" in sql + assert f"id = '{identifier('logstore-source')}'" in sql + assert f'FROM "{logstore}"' in sql + assert "ORDER BY log_version DESC LIMIT 1" in sql + return get_logs(client, query=query, logstore=logstore) + + monkeypatch.setattr(AliyunLogStore, "execute_sql", execute_sql) + event.listen(sqlite_engine, "begin", begun) + event.listen(sqlite_engine, "commit", ended) + event.listen(sqlite_engine, "rollback", ended) + try: + yield StoredSource(case, kind, row, requests) + finally: + event.remove(sqlite_engine, "begin", begun) + event.remove(sqlite_engine, "commit", ended) + event.remove(sqlite_engine, "rollback", ended) + + @staticmethod + def test_status_and_retry_use_logstore_without_sql_source_rows( + stored: StoredSource, monkeypatch: pytest.MonkeyPatch + ) -> None: + run = stored.start() + with stored.case.sessions() as session: + assert session.scalar(select(WorkflowRun)) is None + assert session.scalar(select(WorkflowNodeExecutionModel)) is None + with stored.case.sessions.begin() as session: + current = session.get(ApplicationTraceRun, run.id) + assert current is not None + current.state = "failed" + current.error_code = "provider_export_rejected" + current.outer_trace = {"persisted": True} + assert stored.status()["status"] == "failed" + scheduled: list[str] = [] + service = stored.service() + monkeypatch.setattr(service, "schedule", scheduled.append) + stored.requests.clear() + result = service.retry( + tenant_id=run.tenant_id, + app_id=run.app_id, + source_kind=stored.kind, + source_id=run.source_id, + acknowledge_possible_duplicates=False, + ) + assert result["status"] == "ready" + assert scheduled == [run.id] + assert len(stored.requests) == 1 + + @staticmethod + @pytest.mark.parametrize("field", ["id", "tenant_id", "app_id"]) + def test_backend_result_must_match_complete_source_owner(stored: StoredSource, field: str) -> None: + run = stored.start() + stored.row[field] = identifier("foreign") + with pytest.raises(TraceCollectionError, match="tracing_source_not_found") as caught: + stored.service().status( + tenant_id=run.tenant_id, app_id=run.app_id, source_kind=stored.kind, source_id=run.source_id + ) + assert caught.value.status == 404 + assert stored.case.repository.get(run.id) is not None + + @staticmethod + def test_foreign_app_owner_is_rejected_before_logstore_request(stored: StoredSource) -> None: + stored.start() + stored.requests.clear() + with stored.case.sessions.begin() as session: + app = session.get(App, stored.case.run.app_id) + assert app is not None + app.tenant_id = identifier("other-tenant") + with pytest.raises(TraceCollectionError, match="tracing_source_not_found"): + stored.status() + assert stored.requests == [] + + @staticmethod + def test_initialization_failure_uses_configured_source_reader(stored: StoredSource) -> None: + run = stored.case.repository.record_initialization_failure( + tenant_id=stored.case.run.tenant_id, + app_id=stored.case.run.app_id, + source_kind=stored.kind, + source_id=str(stored.row["id"]), + ) + assert run is not None + assert run.state == "failed" + assert run.error_code == "tracing_initialization_failed" + + @staticmethod + def test_existing_snapshot_reentry_does_not_consult_logstore(stored: StoredSource) -> None: + original = stored.start() + stored.requests.clear() + stored.row.clear() + current = stored.case.repository.start( + tenant_id=original.tenant_id, + app_id=original.app_id, + source_kind=stored.kind, + source_id=original.source_id, + session_id="must-not-replace", + ) + assert current is not None + assert current.id == original.id + assert current.session_id == original.session_id + assert stored.requests == [] + + @staticmethod + @pytest.mark.parametrize("stored", ["workflow"], indirect=True) + def test_workflow_node_deep_link_is_scoped_to_durable_invocation(stored: StoredSource) -> None: + run = stored.start() + case = Collection(stored.case.repository, stored.case.sessions, run) + node_id = identifier("logstore-agent-node") + invocation = case.repository.invocation( + run_id=run.id, agent_id=identifier("agent"), node_execution_id=node_id, attempt=0, call_key="node" + ) + ingest(case, batch_for(run, invocation)) + terminal(case) + lease = claim(case) + assert lease is not None + case.repository.complete( + run.id, + lease[1], + state="accepted", + result=TraceExportResult("provider-id", span_urls={invocation.span_id: "https://trace.test/span"}), + ) + assert stored.status(node_id)["span_url"] == "https://trace.test/span" + with pytest.raises(TraceCollectionError, match="tracing_node_not_found"): + stored.status(identifier("another-run-node")) + + @staticmethod + def test_preview_settlement_and_workflow_recovery_use_configured_source(stored: StoredSource) -> None: + run = stored.start() + if stored.kind == "workflow": + task = _workflow_task(stored.case.repository, run, naive_utc_now()) + assert task is not None + assert task.workflow_run_id == run.source_id + assert task.kwargs["trace_session_id"] == run.session_id + stored.row["status"] = "running" + assert _workflow_task(stored.case.repository, run, naive_utc_now()) is None + else: + stored.service().settle_preview(run) + current = stored.case.repository.get(run.id) + assert current is not None + assert current.outer_trace is not None + root = current.outer_trace["spans"][0] + assert root["status"] == "ok" + assert root["inputs"] is None + assert root["outputs"] is None + assert root["metadata"]["node_execution_id"] == run.source_id + + @staticmethod + @pytest.mark.parametrize("stored", ["workflow"], indirect=True) + def test_missing_snapshot_classifies_logstore_agent_workflow_and_persists_failure(stored: StoredSource) -> None: + info = workflow_info(_dict_to_workflow_run(stored.row)) + assert stored.service().capture_terminal(info) is True + result = stored.status() + assert result["status"] == "failed" + assert result["error_code"] == "tracing_initialization_failed" + + @staticmethod + @pytest.mark.parametrize("stored", ["workflow"], indirect=True) + @pytest.mark.parametrize("has_agent", [True, False]) + def test_unindexed_workflow_uses_owned_definition_for_safe_legacy_decision( + stored: StoredSource, has_agent: bool + ) -> None: + info = workflow_info(_dict_to_workflow_run(stored.row)) + with stored.case.sessions.begin() as session: + app = session.get(App, stored.case.run.app_id) + assert app is not None + definition = make_workflow( + workflow_id=str(stored.row["workflow_id"]), + tenant_id=app.tenant_id, + app_id=app.id, + graph=AGENT_GRAPH if has_agent else {"nodes": [{"data": {"type": "llm"}}]}, + ) + app.workflow_id = definition.id + session.add(definition) + stored.row.clear() + assert stored.service().capture_terminal(info) is has_agent + + @staticmethod + @pytest.mark.parametrize("stored", ["workflow"], indirect=True) + @pytest.mark.parametrize("snapshot", [True, False]) + def test_auxiliary_content_is_suppressed_for_logstore_workflow(stored: StoredSource, snapshot: bool) -> None: + if snapshot: + stored.start() + info = ModerationTraceInfo( + message_id=identifier("workflow-log"), + metadata={ + "workflow_run_id": stored.row["id"], + **({} if snapshot else {"app_id": stored.case.run.app_id}), + }, + flagged=False, + action="direct_output", + preset_response="", + query="private query", + inputs="private input", + outputs="private output", + ) + assert stored.service().capture_terminal(info) is True + if snapshot: + # Durable ownership is sufficient to suppress auxiliary payload even when + # the remote run has not been indexed yet. + stored.row.clear() + assert stored.service().capture_terminal(info) is True + else: + stored.row["graph"] = json.dumps({"nodes": [{"data": {"type": "llm"}}]}) + assert stored.service().capture_terminal(info) is False + + @staticmethod + @pytest.mark.parametrize("stored", ["workflow"], indirect=True) + @pytest.mark.parametrize("mismatch", [None, "tenant_id", "app_id"]) + def test_linked_parent_requires_logstore_run_and_durable_owner_to_agree( + stored: StoredSource, mismatch: str | None + ) -> None: + parent = stored.start() + info = workflow_info(_dict_to_workflow_run(stored.row), _dict_to_workflow_run(stored.row)) + original = dict(info.metadata) + if mismatch: + stored.row[mismatch] = identifier("foreign") + stored.service()._link_managed_parent(info, parent.tenant_id) + if mismatch: + assert info.metadata == original + else: + assert info.resolved_parent_context == (None, None) + assert info.metadata["linked_parent_trace_id"] == parent.trace_id + assert info.metadata["linked_parent_workflow_run_id"] == parent.source_id + + @staticmethod + @pytest.mark.parametrize("stored", ["workflow"], indirect=True) + def test_auxiliary_durable_owner_rejects_app_reassigned_to_other_tenant(stored: StoredSource) -> None: + run = stored.start() + with stored.case.sessions.begin() as session: + app = session.get(App, run.app_id) + assert app is not None + app.tenant_id = identifier("foreign") + info = ModerationTraceInfo( + metadata={"workflow_run_id": run.source_id}, + query="", + flagged=False, + action="direct_output", + preset_response="", + ) + assert stored.service().capture_terminal(info) is False + + @staticmethod + def test_sql_workflow_source_still_uses_scoped_configured_reader( + sqlite_session_factory: sessionmaker[Session], + ) -> None: + case = seed(sqlite_session_factory, "sql-reader") + row = _dict_to_workflow_run( + { + "id": identifier("sql-run"), + "tenant_id": case.run.tenant_id, + "app_id": case.run.app_id, + "workflow_id": identifier("definition"), + "created_by": identifier("owner"), + } + ) + with case.sessions.begin() as session: + session.add(row) + source = case.repository._source( + tenant_id=case.run.tenant_id, app_id=case.run.app_id, kind="workflow", source_id=row.id + ) + assert isinstance(source, WorkflowRun) + assert source.id == row.id + assert ( + case.repository._source( + tenant_id=identifier("wrong-tenant"), app_id=case.run.app_id, kind="workflow", source_id=row.id + ) + is None + ) + + +"""Persistence boundaries that must fail closed or remain immutable across retries.""" + + +class TestApplicationTraceRepositoryBoundariesCases: + @staticmethod + @pytest.mark.parametrize("policy", ["invalid-json", "null", "[]"]) + def test_malformed_policy_freezes_disabled_decision( + sqlite_session_factory: sessionmaker[Session], policy: str + ) -> None: + case = seed(sqlite_session_factory, "malformed") + with case.sessions.begin() as session: + session.execute(delete(ApplicationTraceRun).where(ApplicationTraceRun.id == case.run.id)) + app = session.get(App, case.run.app_id) + assert app is not None + app.tracing = policy + run = case.repository.start( + tenant_id=case.run.tenant_id, + app_id=case.run.app_id, + source_kind="message", + source_id=case.run.source_id, + session_id="", + ) + assert run is not None + assert run.state == "not_collected" + assert run.provider_config == {} + assert run.session_id == case.run.source_id + assert not run.capture_content + assert not run.capture_agent + + @staticmethod + def test_reserved_preview_requires_persisted_owner_workflow(sqlite_session_factory: sessionmaker[Session]) -> None: + case = seed(sqlite_session_factory, "missing-preview-owner") + assert ( + case.repository.start( + tenant_id=case.run.tenant_id, + app_id=case.run.app_id, + source_kind="node", + source_id=identifier("preview"), + session_id="preview", + workflow_id=identifier("missing-workflow"), + ) + is None + ) + + @staticmethod + def test_initialization_failure_cannot_replace_snapshot_or_create_foreign_row( + sqlite_session_factory: sessionmaker[Session], + ) -> None: + case = seed(sqlite_session_factory, "failure-snapshot") + same = case.repository.record_initialization_failure( + tenant_id=case.run.tenant_id, + app_id=case.run.app_id, + source_kind="message", + source_id=case.run.source_id, + ) + assert same is not None + assert same.id == case.run.id + assert same.state == "collecting" + assert same.provider_config == case.run.provider_config + assert ( + case.repository.record_initialization_failure( + tenant_id=identifier("foreign-tenant"), + app_id=case.run.app_id, + source_kind="message", + source_id=case.run.source_id, + ) + is None + ) + + @staticmethod + @pytest.mark.parametrize("target", ["missing_run", "revoked_owner", "closed_invocation"]) + def test_ingestion_rejects_missing_or_closed_owner_without_writing( + sqlite_session_factory: sessionmaker[Session], target: str + ) -> None: + case = seed(sqlite_session_factory, "ingestion-owner") + invocation = case.invocation() + batch = batch_for(case.run, invocation) + if target == "missing_run": + batch = batch.model_copy(update={"run_id": UUID(identifier("missing-run"))}) + error = "tracing_run_not_found" + elif target == "revoked_owner": + with case.sessions.begin() as session: + app = session.get(App, case.run.app_id) + assert app is not None + app.tenant_id = identifier("another-tenant") + error = "tracing_owner_not_found" + else: + with case.sessions.begin() as session: + stored = session.get(ApplicationTraceInvocation, invocation.id) + assert stored is not None + stored.finalized = True + error = "tracing_collection_closed" + with pytest.raises(TraceCollectionError, match=error): + ingest(case, batch) + with case.sessions() as session: + assert session.scalar(select(func.count()).select_from(ApplicationTraceBatch)) == 0 + stored_run = case.repository.get(case.run.id) + assert stored_run is not None + assert stored_run.span_count == 0 + + @staticmethod + def test_first_terminal_is_immutable_and_revoked_owner_is_never_exported( + sqlite_session_factory: sessionmaker[Session], + ) -> None: + case = seed(sqlite_session_factory, "terminal-owner") + terminal(case) + before = case.repository.get(case.run.id) + assert before is not None + case.repository.terminal(case.run.id, {"replacement": True}, "f" * 16) + after = case.repository.get(case.run.id) + assert after is not None + assert after.outer_trace == before.outer_trace + assert after.root_span_id == before.root_span_id + with case.sessions.begin() as session: + app = session.get(App, case.run.app_id) + assert app is not None + app.tenant_id = identifier("moved-owner") + assert claim(case) is None + failed = case.repository.get(case.run.id) + assert failed is not None + assert failed.error_code == "tracing_owner_not_found" + assert failed.state == "failed" + case.repository.terminal(case.run.id, {"replacement": True}, "f" * 16) + frozen = case.repository.get(case.run.id) + assert frozen is not None + assert frozen.outer_trace == before.outer_trace + + @staticmethod + @pytest.mark.parametrize( + ("state", "error"), + [ + ("not_collected", "tracing_run_not_found"), + ("accepted", "tracing_already_accepted"), + ("collecting", "tracing_retry_not_allowed"), + ], + ) + def test_retry_rejects_ineligible_states( + sqlite_session_factory: sessionmaker[Session], state: str, error: str + ) -> None: + case = seed(sqlite_session_factory, "retry-state") + terminal(case) + with case.sessions.begin() as session: + row = session.get(ApplicationTraceRun, case.run.id) + assert row is not None + row.state = state + with pytest.raises(TraceCollectionError, match=error): + case.repository.retry(run_id=case.run.id, acknowledge_possible_duplicates=True) + row = case.repository.get(case.run.id) + assert row is not None + assert row.state == state + + @staticmethod + def test_retry_unknown_run_and_invocation_reject_before_creation( + sqlite_session_factory: sessionmaker[Session], + ) -> None: + case = seed(sqlite_session_factory, "retry-unknown") + with pytest.raises(TraceCollectionError, match="tracing_run_not_found"): + case.repository.retry(run_id=identifier("missing"), acknowledge_possible_duplicates=True) + with pytest.raises(TraceCollectionError, match="tracing_run_not_collecting"): + case.repository.invocation( + run_id=identifier("missing"), + agent_id=identifier("agent"), + node_execution_id=None, + attempt=0, + call_key="x", + ) + assert case.repository.claim(identifier("missing"), collection_timeout=1, lease_seconds=1) is None + + @staticmethod + def test_collection_without_terminal_is_not_exportable(sqlite_session_factory: sessionmaker[Session]) -> None: + case = seed(sqlite_session_factory, "no-terminal") + ingest(case, batch_for(case.run, case.invocation())) + assert claim(case) is None + row = case.repository.get(case.run.id) + assert row is not None + assert row.attempts == 0 + assert row.state == "collecting" + + @staticmethod + def test_retention_removes_all_payload_rows_and_preserves_newer_runs( + sqlite_session_factory: sessionmaker[Session], + ) -> None: + expired = seed(sqlite_session_factory, "expired") + fresh = seed(sqlite_session_factory, "fresh") + for case in (expired, fresh): + ingest(case, batch_for(case.run, case.invocation())) + with expired.sessions.begin() as session: + row = session.get(ApplicationTraceRun, expired.run.id) + assert row is not None + row.created_at = naive_utc_now() - timedelta(days=8) + assert expired.repository.purge(retention_days=7) == 1 + assert expired.repository.get(expired.run.id) is None + assert expired.repository.collected(expired.run.id) == ([], []) + assert fresh.repository.get(fresh.run.id) is not None + invocations, batches = fresh.repository.collected(fresh.run.id) + assert len(invocations) == len(batches) == 1 + assert expired.repository.purge(retention_days=7) == 0 diff --git a/api/tests/unit_tests/repositories/test_conversation_variable_query_repository.py b/api/tests/unit_tests/repositories/test_conversation_variable_query_repository.py index 5f327608f51..d870dca9af9 100644 --- a/api/tests/unit_tests/repositories/test_conversation_variable_query_repository.py +++ b/api/tests/unit_tests/repositories/test_conversation_variable_query_repository.py @@ -11,7 +11,7 @@ from machinery.context import RequestContext from models import App, AppMode, ConversationVariable from repositories.conversation_variable_query_repository import ConversationVariableQueryRepository from services.conversation_variable_query import ConversationVariableAppNotFoundError -from tests.unit_tests.model_factories import make_app +from tests.helpers.unit.model_factories import make_app CONTEXT = RequestContext(request_id="request", trace_id=None, account_id="account-1", active_workspace_id="tenant-1") diff --git a/api/tests/unit_tests/repositories/test_public_app_availability.py b/api/tests/unit_tests/repositories/test_public_app_availability.py deleted file mode 100644 index 5c654da4d26..00000000000 --- a/api/tests/unit_tests/repositories/test_public_app_availability.py +++ /dev/null @@ -1,202 +0,0 @@ -"""Exercise the shared public publication gate using real persisted owner chains.""" - -from uuid import uuid4 - -import pytest -from sqlalchemy import select -from sqlalchemy.orm import Session, sessionmaker - -from core.app.public_runtime import published_app_filter -from models.agent import ( - Agent, - AgentConfigRevision, - AgentConfigRevisionOperation, - AgentConfigSnapshot, - AgentScope, - AgentSource, - AgentStatus, -) -from models.model import App, AppMode, AppModelConfig, CustomizeTokenStrategy, Site -from models.workflow import Workflow, WorkflowType -from repositories.web_passport_repository import WebPassportRepository -from repositories.webapp_access_query_repository import WebAppAccessQueryRepository - - -def _app(session: Session, mode: AppMode) -> App: - app = App( - id=str(uuid4()), tenant_id=str(uuid4()), name="Private fixture", mode=mode, enable_site=True, enable_api=True - ) - session.add_all( - [ - app, - Site( - app_id=app.id, - code="public-code", - title="private", - default_language="en-US", - customize_token_strategy=CustomizeTokenStrategy.UUID, - ), - ] - ) - session.flush() - return app - - -def _workflow(session: Session, app: App, *, version: str = "published") -> Workflow: - workflow = Workflow( - id=str(uuid4()), - tenant_id=app.tenant_id, - app_id=app.id, - type=WorkflowType.WORKFLOW if app.mode == AppMode.WORKFLOW else WorkflowType.CHAT, - version=version, - graph='{"nodes":[]}', - features="{}", - created_by=str(uuid4()), - environment_variables=[], - conversation_variables=[], - rag_pipeline_variables=[], - ) - session.add(workflow) - return workflow - - -def _agent(session: Session, app: App, *, published: bool) -> tuple[Agent, AgentConfigSnapshot]: - agent = Agent( - id=str(uuid4()), - tenant_id=app.tenant_id, - app_id=app.id, - name="Private Agent", - scope=AgentScope.ROSTER, - source=AgentSource.AGENT_APP, - active_config_is_published=published, - ) - snapshot = AgentConfigSnapshot( - id=str(uuid4()), - tenant_id=app.tenant_id, - agent_id=agent.id, - version=1, - home_snapshot_id=str(uuid4()), - config_snapshot={}, - ) - agent.active_config_snapshot_id = snapshot.id - session.add_all([agent, snapshot]) - if published: - session.add( - AgentConfigRevision( - tenant_id=app.tenant_id, - agent_id=agent.id, - current_snapshot_id=snapshot.id, - revision=1, - operation=AgentConfigRevisionOperation.PUBLISH_DRAFT, - ) - ) - return agent, snapshot - - -def _assert_readers(factory: sessionmaker[Session], app: App, expected: bool) -> None: - access = WebAppAccessQueryRepository(session_factory=factory) - assert access.find_app_id_by_code("public-code") == (app.id if expected else None) - assert access.is_app_available(app.id) is expected - passports = WebPassportRepository(session_factory=factory, generate_session_id=lambda: "never-created") - record = passports.get_active_web_app("public-code") - assert (record is not None) is expected - - -@pytest.mark.parametrize( - "mode", - [AppMode.WORKFLOW, AppMode.ADVANCED_CHAT, AppMode.CHAT, AppMode.COMPLETION, AppMode.AGENT_CHAT, AppMode.AGENT], -) -@pytest.mark.parametrize("published", [False, True]) -def test_all_public_app_modes_require_their_published_artifact( - sqlite_session: Session, sqlite_session_factory: sessionmaker[Session], mode: AppMode, published: bool -) -> None: - app = _app(sqlite_session, mode) - if mode == AppMode.AGENT: - _agent(sqlite_session, app, published=published) - elif mode in (AppMode.WORKFLOW, AppMode.ADVANCED_CHAT): - draft = _workflow(sqlite_session, app, version="draft") - if published: - app.workflow_id = _workflow(sqlite_session, app).id - assert draft.id != app.workflow_id - elif published: - config = AppModelConfig(app_id=app.id) - config.id = str(uuid4()) - app.app_model_config_id = config.id - sqlite_session.add(config) - sqlite_session.commit() - _assert_readers(sqlite_session_factory, app, published) - - -@pytest.mark.parametrize("invalid", ["draft-pointer", "missing", "other-app", "other-tenant"]) -@pytest.mark.parametrize("mode", [AppMode.WORKFLOW, AppMode.ADVANCED_CHAT]) -def test_workflow_pointer_must_resolve_to_this_apps_published_workflow( - sqlite_session: Session, sqlite_session_factory: sessionmaker[Session], mode: AppMode, invalid: str -) -> None: - app = _app(sqlite_session, mode) - workflow = _workflow(sqlite_session, app, version="draft" if invalid == "draft-pointer" else "published") - app.workflow_id = str(uuid4()) if invalid == "missing" else workflow.id - if invalid == "other-app": - workflow.app_id = str(uuid4()) - if invalid == "other-tenant": - workflow.tenant_id = str(uuid4()) - sqlite_session.commit() - _assert_readers(sqlite_session_factory, app, False) - - -@pytest.mark.parametrize("invalid", ["missing", "other-app"]) -def test_model_config_pointer_must_resolve_to_this_app( - sqlite_session: Session, sqlite_session_factory: sessionmaker[Session], invalid: str -) -> None: - app = _app(sqlite_session, AppMode.CHAT) - config = AppModelConfig(app_id=app.id if invalid == "missing" else str(uuid4())) - config.id = str(uuid4()) - app.app_model_config_id = str(uuid4()) if invalid == "missing" else config.id - sqlite_session.add(config) - sqlite_session.commit() - _assert_readers(sqlite_session_factory, app, False) - - -@pytest.mark.parametrize( - "state", - [ - "dirty-draft", - "seeded", - "snapshot-missing", - "snapshot-other-agent", - "snapshot-other-tenant", - "agent-other-tenant", - "archived", - "workflow-only", - ], -) -def test_agent_publication_uses_owned_visible_snapshot_not_dirty_draft_flag( - sqlite_session: Session, sqlite_session_factory: sessionmaker[Session], state: str -) -> None: - app = _app(sqlite_session, AppMode.AGENT) - agent, snapshot = _agent(sqlite_session, app, published=state != "seeded") - # A seeded snapshot is not public even when an inconsistent dirty flag says otherwise. - agent.active_config_is_published = state == "seeded" - if state == "snapshot-missing": - agent.active_config_snapshot_id = str(uuid4()) - elif state == "snapshot-other-agent": - snapshot.agent_id = str(uuid4()) - elif state == "snapshot-other-tenant": - snapshot.tenant_id = str(uuid4()) - elif state == "agent-other-tenant": - agent.tenant_id = str(uuid4()) - elif state == "archived": - agent.status = AgentStatus.ARCHIVED - elif state == "workflow-only": - agent.scope = AgentScope.WORKFLOW_ONLY - agent.backing_app_id = app.id - sqlite_session.commit() - _assert_readers(sqlite_session_factory, app, state == "dirty-draft") - - -def test_public_predicate_does_not_change_management_visibility(sqlite_session: Session) -> None: - app = _app(sqlite_session, AppMode.ADVANCED_CHAT) - draft = _workflow(sqlite_session, app, version="draft") - sqlite_session.commit() - assert sqlite_session.scalar(select(App.id).where(App.id == app.id, published_app_filter())) is None - assert sqlite_session.get(App, app.id) is app - assert sqlite_session.get(Workflow, draft.id) is draft diff --git a/api/tests/unit_tests/repositories/test_sqlalchemy_api_workflow_node_execution_repository.py b/api/tests/unit_tests/repositories/test_sqlalchemy_api_workflow_node_execution_repository.py index 8e80e27b670..db351810161 100644 --- a/api/tests/unit_tests/repositories/test_sqlalchemy_api_workflow_node_execution_repository.py +++ b/api/tests/unit_tests/repositories/test_sqlalchemy_api_workflow_node_execution_repository.py @@ -23,7 +23,7 @@ from repositories.factory import DifyAPIRepositoryFactory from repositories.sqlalchemy_api_workflow_node_execution_repository import ( DifyAPISQLAlchemyWorkflowNodeExecutionRepository, ) -from tests.unit_tests.config_override import apply_config_overrides +from tests.helpers.unit.config_override import apply_config_overrides @pytest.mark.parametrize("logstore", [False, True]) diff --git a/api/tests/unit_tests/repositories/test_knowledge_fs_capability_issuance_auditor.py b/api/tests/unit_tests/repositories/test_sqlalchemy_knowledge_fs_capability_issuance_auditor.py similarity index 100% rename from api/tests/unit_tests/repositories/test_knowledge_fs_capability_issuance_auditor.py rename to api/tests/unit_tests/repositories/test_sqlalchemy_knowledge_fs_capability_issuance_auditor.py diff --git a/api/tests/unit_tests/repositories/test_sqlalchemy_knowledge_fs_control_space_repository.py b/api/tests/unit_tests/repositories/test_sqlalchemy_knowledge_fs_control_space_repository.py index 972869cbcdd..93d44c999a5 100644 --- a/api/tests/unit_tests/repositories/test_sqlalchemy_knowledge_fs_control_space_repository.py +++ b/api/tests/unit_tests/repositories/test_sqlalchemy_knowledge_fs_control_space_repository.py @@ -17,7 +17,7 @@ from services.knowledge_fs.control_space_lifecycle import ( KnowledgeFSSpaceRegistrationConflictError, KnowledgeFSWorkspaceDeletionBlockedError, ) -from tests.unit_tests.tasks.knowledge_fs_task_fakes import RecordingKnowledgeFSMetrics +from tests.helpers.unit.tasks.knowledge_fs_task_fakes import RecordingKnowledgeFSMetrics @pytest.fixture diff --git a/api/tests/unit_tests/repositories/test_sqlalchemy_knowledge_fs_lifecycle_outbox_repository.py b/api/tests/unit_tests/repositories/test_sqlalchemy_knowledge_fs_lifecycle_outbox_repository.py index f9c50bd8add..c9039b75505 100644 --- a/api/tests/unit_tests/repositories/test_sqlalchemy_knowledge_fs_lifecycle_outbox_repository.py +++ b/api/tests/unit_tests/repositories/test_sqlalchemy_knowledge_fs_lifecycle_outbox_repository.py @@ -16,7 +16,7 @@ from repositories.sqlalchemy_knowledge_fs_lifecycle_outbox_repository import ( SQLAlchemyKnowledgeFSLifecycleOutboxRepository, ) from services.knowledge_fs.observability import KnowledgeFSLifecycleTaskMetric -from tests.unit_tests.tasks.knowledge_fs_task_fakes import RecordingKnowledgeFSMetrics +from tests.helpers.unit.tasks.knowledge_fs_task_fakes import RecordingKnowledgeFSMetrics def _outbox(control_space: KnowledgeFSControlSpace) -> KnowledgeFSLifecycleOutbox: diff --git a/api/tests/unit_tests/repositories/test_trial_app_repository.py b/api/tests/unit_tests/repositories/test_trial_app_repository.py index e311516b79b..220f92df3a5 100644 --- a/api/tests/unit_tests/repositories/test_trial_app_repository.py +++ b/api/tests/unit_tests/repositories/test_trial_app_repository.py @@ -1,4 +1,5 @@ import uuid +from uuid import uuid4 import pytest from pytest_mock import MockerFixture @@ -6,6 +7,14 @@ from sqlalchemy import Engine, select from sqlalchemy.orm import Session, sessionmaker from sqlalchemy.pool import QueuePool +from models.agent import ( + Agent, + AgentConfigRevision, + AgentConfigRevisionOperation, + AgentConfigSnapshot, + AgentScope, + AgentSource, +) from models.model import AccountTrialAppRecord, App, AppMode, TrialApp from repositories.trial_app_repository import TrialAppRepository from services.trial_app_access_service import ( @@ -233,3 +242,61 @@ def test_record_does_not_commit_caller_session(sqlite_session_factory: sessionma usage = _record(sqlite_session_factory, app_id=usage_app_id, account_id=account_id) assert usage is not None assert usage.count == 1 + + +@pytest.mark.parametrize( + ("operation", "include_snapshot", "expected"), + [ + (AgentConfigRevisionOperation.PUBLISH_DRAFT, True, True), + (AgentConfigRevisionOperation.IMPORT_PACKAGE, True, False), + (AgentConfigRevisionOperation.PUBLISH_DRAFT, False, False), + ], + ids=["published-with-dirty-draft", "unpublished-seed", "missing-snapshot"], +) +def test_agent_trial_requires_a_visible_existing_snapshot( + sqlite_session_factory: sessionmaker[Session], + operation: AgentConfigRevisionOperation, + include_snapshot: bool, + expected: bool, +) -> None: + app_id, tenant_id, agent_id, snapshot_id = (str(uuid4()) for _ in range(4)) + with sqlite_session_factory.begin() as session: + session.add_all( + [ + App( + id=app_id, + tenant_id=tenant_id, + name="Trial Agent", + mode=AppMode.AGENT, + enable_site=True, + enable_api=False, + ), + TrialApp(app_id=app_id, tenant_id=tenant_id), + Agent( + id=agent_id, + tenant_id=tenant_id, + app_id=app_id, + name="Agent", + scope=AgentScope.ROSTER, + source=AgentSource.AGENT_APP, + active_config_snapshot_id=snapshot_id, + active_config_is_published=False, + ), + AgentConfigRevision( + tenant_id=tenant_id, + agent_id=agent_id, + current_snapshot_id=snapshot_id, + revision=1, + operation=operation, + ), + ] + ) + if include_snapshot: + session.add( + AgentConfigSnapshot( + id=snapshot_id, tenant_id=tenant_id, agent_id=agent_id, version=1, config_snapshot={} + ) + ) + + result = TrialAppRepository(sqlite_session_factory).existing_ids([app_id]) + assert result == (frozenset({app_id}) if expected else frozenset()) diff --git a/api/tests/unit_tests/repositories/test_trial_app_repository_agent.py b/api/tests/unit_tests/repositories/test_trial_app_repository_agent.py deleted file mode 100644 index 453f678a257..00000000000 --- a/api/tests/unit_tests/repositories/test_trial_app_repository_agent.py +++ /dev/null @@ -1,73 +0,0 @@ -from uuid import uuid4 - -import pytest -from sqlalchemy.orm import Session, sessionmaker - -from models.agent import ( - Agent, - AgentConfigRevision, - AgentConfigRevisionOperation, - AgentConfigSnapshot, - AgentScope, - AgentSource, -) -from models.model import App, AppMode, TrialApp -from repositories.trial_app_repository import TrialAppRepository - - -@pytest.mark.parametrize( - ("operation", "include_snapshot", "expected"), - [ - (AgentConfigRevisionOperation.PUBLISH_DRAFT, True, True), - (AgentConfigRevisionOperation.IMPORT_PACKAGE, True, False), - (AgentConfigRevisionOperation.PUBLISH_DRAFT, False, False), - ], - ids=["published-with-dirty-draft", "unpublished-seed", "missing-snapshot"], -) -def test_agent_trial_requires_a_visible_existing_snapshot( - sqlite_session_factory: sessionmaker[Session], - operation: AgentConfigRevisionOperation, - include_snapshot: bool, - expected: bool, -) -> None: - app_id, tenant_id, agent_id, snapshot_id = (str(uuid4()) for _ in range(4)) - with sqlite_session_factory.begin() as session: - session.add_all( - [ - App( - id=app_id, - tenant_id=tenant_id, - name="Trial Agent", - mode=AppMode.AGENT, - enable_site=True, - enable_api=False, - ), - TrialApp(app_id=app_id, tenant_id=tenant_id), - Agent( - id=agent_id, - tenant_id=tenant_id, - app_id=app_id, - name="Agent", - scope=AgentScope.ROSTER, - source=AgentSource.AGENT_APP, - active_config_snapshot_id=snapshot_id, - active_config_is_published=False, - ), - AgentConfigRevision( - tenant_id=tenant_id, - agent_id=agent_id, - current_snapshot_id=snapshot_id, - revision=1, - operation=operation, - ), - ] - ) - if include_snapshot: - session.add( - AgentConfigSnapshot( - id=snapshot_id, tenant_id=tenant_id, agent_id=agent_id, version=1, config_snapshot={} - ) - ) - - result = TrialAppRepository(sqlite_session_factory).existing_ids([app_id]) - assert result == (frozenset({app_id}) if expected else frozenset()) diff --git a/api/tests/unit_tests/repositories/test_webapp_access_query_repository.py b/api/tests/unit_tests/repositories/test_webapp_access_query_repository.py index 01d80684a65..308ba021eb9 100644 --- a/api/tests/unit_tests/repositories/test_webapp_access_query_repository.py +++ b/api/tests/unit_tests/repositories/test_webapp_access_query_repository.py @@ -1,15 +1,213 @@ from unittest.mock import MagicMock, patch +from uuid import uuid4 import pytest from sqlalchemy import Engine, event, select from sqlalchemy.exc import OperationalError from sqlalchemy.orm import Session, sessionmaker -from models.enums import AppStatus, CustomizeTokenStrategy, EndUserType +from core.app.public_runtime import published_app_filter +from models.agent import ( + Agent, + AgentConfigRevision, + AgentConfigRevisionOperation, + AgentConfigSnapshot, + AgentScope, + AgentSource, + AgentStatus, +) +from models.enums import AppStatus, EndUserType +from models.enums import CustomizeTokenStrategy as _webapp_access_query_repository_CustomizeTokenStrategy from models.model import App, AppMode, AppModelConfig, EndUser, Site +from models.model import CustomizeTokenStrategy as _public_app_availability_CustomizeTokenStrategy +from models.workflow import Workflow, WorkflowType +from repositories.web_passport_repository import WebPassportRepository from repositories.webapp_access_query_repository import WebAppAccessQueryRepository from services.webapp_access_query_service import WebAppAccessUnavailableError +"""Exercise the shared public publication gate using real persisted owner chains.""" + + +def _app(session: Session, mode: AppMode) -> App: + app = App( + id=str(uuid4()), tenant_id=str(uuid4()), name="Private fixture", mode=mode, enable_site=True, enable_api=True + ) + session.add_all( + [ + app, + Site( + app_id=app.id, + code="public-code", + title="private", + default_language="en-US", + customize_token_strategy=_public_app_availability_CustomizeTokenStrategy.UUID, + ), + ] + ) + session.flush() + return app + + +def _workflow(session: Session, app: App, *, version: str = "published") -> Workflow: + workflow = Workflow( + id=str(uuid4()), + tenant_id=app.tenant_id, + app_id=app.id, + type=WorkflowType.WORKFLOW if app.mode == AppMode.WORKFLOW else WorkflowType.CHAT, + version=version, + graph='{"nodes":[]}', + features="{}", + created_by=str(uuid4()), + environment_variables=[], + conversation_variables=[], + rag_pipeline_variables=[], + ) + session.add(workflow) + return workflow + + +def _agent(session: Session, app: App, *, published: bool) -> tuple[Agent, AgentConfigSnapshot]: + agent = Agent( + id=str(uuid4()), + tenant_id=app.tenant_id, + app_id=app.id, + name="Private Agent", + scope=AgentScope.ROSTER, + source=AgentSource.AGENT_APP, + active_config_is_published=published, + ) + snapshot = AgentConfigSnapshot( + id=str(uuid4()), + tenant_id=app.tenant_id, + agent_id=agent.id, + version=1, + home_snapshot_id=str(uuid4()), + config_snapshot={}, + ) + agent.active_config_snapshot_id = snapshot.id + session.add_all([agent, snapshot]) + if published: + session.add( + AgentConfigRevision( + tenant_id=app.tenant_id, + agent_id=agent.id, + current_snapshot_id=snapshot.id, + revision=1, + operation=AgentConfigRevisionOperation.PUBLISH_DRAFT, + ) + ) + return agent, snapshot + + +def _assert_readers(factory: sessionmaker[Session], app: App, expected: bool) -> None: + access = WebAppAccessQueryRepository(session_factory=factory) + assert access.find_app_id_by_code("public-code") == (app.id if expected else None) + assert access.is_app_available(app.id) is expected + passports = WebPassportRepository(session_factory=factory, generate_session_id=lambda: "never-created") + record = passports.get_active_web_app("public-code") + assert (record is not None) is expected + + +@pytest.mark.parametrize( + "mode", + [AppMode.WORKFLOW, AppMode.ADVANCED_CHAT, AppMode.CHAT, AppMode.COMPLETION, AppMode.AGENT_CHAT, AppMode.AGENT], +) +@pytest.mark.parametrize("published", [False, True]) +def test_all_public_app_modes_require_their_published_artifact( + sqlite_session: Session, sqlite_session_factory: sessionmaker[Session], mode: AppMode, published: bool +) -> None: + app = _app(sqlite_session, mode) + if mode == AppMode.AGENT: + _agent(sqlite_session, app, published=published) + elif mode in (AppMode.WORKFLOW, AppMode.ADVANCED_CHAT): + draft = _workflow(sqlite_session, app, version="draft") + if published: + app.workflow_id = _workflow(sqlite_session, app).id + assert draft.id != app.workflow_id + elif published: + config = AppModelConfig(app_id=app.id) + config.id = str(uuid4()) + app.app_model_config_id = config.id + sqlite_session.add(config) + sqlite_session.commit() + _assert_readers(sqlite_session_factory, app, published) + + +@pytest.mark.parametrize("invalid", ["draft-pointer", "missing", "other-app", "other-tenant"]) +@pytest.mark.parametrize("mode", [AppMode.WORKFLOW, AppMode.ADVANCED_CHAT]) +def test_workflow_pointer_must_resolve_to_this_apps_published_workflow( + sqlite_session: Session, sqlite_session_factory: sessionmaker[Session], mode: AppMode, invalid: str +) -> None: + app = _app(sqlite_session, mode) + workflow = _workflow(sqlite_session, app, version="draft" if invalid == "draft-pointer" else "published") + app.workflow_id = str(uuid4()) if invalid == "missing" else workflow.id + if invalid == "other-app": + workflow.app_id = str(uuid4()) + if invalid == "other-tenant": + workflow.tenant_id = str(uuid4()) + sqlite_session.commit() + _assert_readers(sqlite_session_factory, app, False) + + +@pytest.mark.parametrize("invalid", ["missing", "other-app"]) +def test_model_config_pointer_must_resolve_to_this_app( + sqlite_session: Session, sqlite_session_factory: sessionmaker[Session], invalid: str +) -> None: + app = _app(sqlite_session, AppMode.CHAT) + config = AppModelConfig(app_id=app.id if invalid == "missing" else str(uuid4())) + config.id = str(uuid4()) + app.app_model_config_id = str(uuid4()) if invalid == "missing" else config.id + sqlite_session.add(config) + sqlite_session.commit() + _assert_readers(sqlite_session_factory, app, False) + + +@pytest.mark.parametrize( + "state", + [ + "dirty-draft", + "seeded", + "snapshot-missing", + "snapshot-other-agent", + "snapshot-other-tenant", + "agent-other-tenant", + "archived", + "workflow-only", + ], +) +def test_agent_publication_uses_owned_visible_snapshot_not_dirty_draft_flag( + sqlite_session: Session, sqlite_session_factory: sessionmaker[Session], state: str +) -> None: + app = _app(sqlite_session, AppMode.AGENT) + agent, snapshot = _agent(sqlite_session, app, published=state != "seeded") + # A seeded snapshot is not public even when an inconsistent dirty flag says otherwise. + agent.active_config_is_published = state == "seeded" + if state == "snapshot-missing": + agent.active_config_snapshot_id = str(uuid4()) + elif state == "snapshot-other-agent": + snapshot.agent_id = str(uuid4()) + elif state == "snapshot-other-tenant": + snapshot.tenant_id = str(uuid4()) + elif state == "agent-other-tenant": + agent.tenant_id = str(uuid4()) + elif state == "archived": + agent.status = AgentStatus.ARCHIVED + elif state == "workflow-only": + agent.scope = AgentScope.WORKFLOW_ONLY + agent.backing_app_id = app.id + sqlite_session.commit() + _assert_readers(sqlite_session_factory, app, state == "dirty-draft") + + +def test_public_predicate_does_not_change_management_visibility(sqlite_session: Session) -> None: + app = _app(sqlite_session, AppMode.ADVANCED_CHAT) + draft = _workflow(sqlite_session, app, version="draft") + sqlite_session.commit() + assert sqlite_session.scalar(select(App.id).where(App.id == app.id, published_app_filter())) is None + assert sqlite_session.get(App, app.id) is app + assert sqlite_session.get(Workflow, draft.id) is draft + + _APP_ID = "11111111-1111-1111-1111-111111111111" _TENANT_ID = "22222222-2222-2222-2222-222222222222" _END_USER_ID = "33333333-3333-3333-3333-333333333333" @@ -38,7 +236,7 @@ def _persist_webapp_session(session: Session, *, enable_site: bool = True) -> No code="site-code", title="Test Site", default_language="en-US", - customize_token_strategy=CustomizeTokenStrategy.UUID, + customize_token_strategy=_webapp_access_query_repository_CustomizeTokenStrategy.UUID, ) ) session.add( @@ -75,7 +273,7 @@ def test_find_app_id_by_code_returns_matching_site_app(sqlite_session_factory: s code="site-code", title="Test Site", default_language="en-US", - customize_token_strategy=CustomizeTokenStrategy.UUID, + customize_token_strategy=_webapp_access_query_repository_CustomizeTokenStrategy.UUID, ) ) @@ -112,7 +310,7 @@ def test_inactive_or_dangling_site_does_not_resolve_as_public_app( title="Hidden", status=AppStatus.NORMAL, default_language="en-US", - customize_token_strategy=CustomizeTokenStrategy.UUID, + customize_token_strategy=_webapp_access_query_repository_CustomizeTokenStrategy.UUID, ) ) repository = WebAppAccessQueryRepository(session_factory=sqlite_session_factory) diff --git a/api/tests/unit_tests/repositories/test_app_trigger_repository.py b/api/tests/unit_tests/repositories/trigger/test_app_repository.py similarity index 100% rename from api/tests/unit_tests/repositories/test_app_trigger_repository.py rename to api/tests/unit_tests/repositories/trigger/test_app_repository.py diff --git a/api/tests/unit_tests/repositories/workspace/test_workspace_members.py b/api/tests/unit_tests/repositories/workspace/test_workspace_members.py deleted file mode 100644 index 8d4d6bcd3ab..00000000000 --- a/api/tests/unit_tests/repositories/workspace/test_workspace_members.py +++ /dev/null @@ -1,257 +0,0 @@ -from collections.abc import Mapping, Sequence -from datetime import datetime - -from sqlalchemy import Engine, event -from sqlalchemy.orm import Session, sessionmaker - -from machinery.context import RequestContext -from models.account import Account, AccountStatus, Tenant, TenantAccountJoin, TenantAccountRole -from repositories.workspace.workspace_repository import WorkspaceRepository -from services.workspace.contracts import WorkspaceMemberRecord, WorkspaceMemberRole, WorkspaceMemberRoleSubject -from services.workspace.member_service import WorkspaceMemberQueryService - - -def make_account( - account_id: str, - *, - status: AccountStatus, - created_at: datetime, -) -> Account: - account = Account( - name=f"Member {account_id}", - email=f"{account_id}@example.com", - avatar=f"{account_id}.png", - status=status, - ) - account.id = account_id - account.last_login_at = created_at - account.last_active_at = created_at - account.created_at = created_at - return account - - -def make_tenant(tenant_id: str) -> Tenant: - tenant = Tenant(name=f"Workspace {tenant_id}") - tenant.id = tenant_id - return tenant - - -def test_rbac_profile_batch_filters_requested_ids_and_workspace(sqlite_session_factory: sessionmaker[Session]) -> None: - with sqlite_session_factory.begin() as session: - session.add_all([make_tenant("workspace"), make_tenant("other")]) - for account_id, workspace in [("member", "workspace"), ("unrequested", "workspace"), ("foreign", "other")]: - session.add(make_account(account_id, status=AccountStatus.ACTIVE, created_at=datetime(2026, 1, 1))) - session.add(TenantAccountJoin(tenant_id=workspace, account_id=account_id, role=TenantAccountRole.NORMAL)) - - repository = WorkspaceRepository(sqlite_session_factory) - members = repository.list_for_workspace("workspace", account_ids=["member", "foreign", "missing"]) - assert [member.id for member in members] == ["member"] - assert members[0].email == "member@example.com" - assert repository.list_for_workspace("workspace", account_ids=[]) == () - - -def test_list_for_workspace_uses_join_membership_and_preserves_account_lifecycle( - sqlite_session_factory: sessionmaker[Session], -) -> None: - created_at = datetime(2026, 1, 1) - active = make_account("active", status=AccountStatus.ACTIVE, created_at=created_at) - uninitialized = make_account("uninitialized", status=AccountStatus.UNINITIALIZED, created_at=created_at) - pending = make_account("pending", status=AccountStatus.PENDING, created_at=created_at) - banned = make_account("banned", status=AccountStatus.BANNED, created_at=created_at) - closed = make_account("closed", status=AccountStatus.CLOSED, created_at=created_at) - other_workspace_member = make_account("other", status=AccountStatus.ACTIVE, created_at=created_at) - unjoined = make_account("unjoined", status=AccountStatus.ACTIVE, created_at=created_at) - workspace = make_tenant("workspace-1") - other_workspace = make_tenant("workspace-2") - - with sqlite_session_factory() as session: - session.add_all( - [ - workspace, - other_workspace, - active, - uninitialized, - pending, - banned, - closed, - other_workspace_member, - unjoined, - TenantAccountJoin( - tenant_id=workspace.id, - account_id=active.id, - role=TenantAccountRole.OWNER, - ), - TenantAccountJoin( - tenant_id=workspace.id, - account_id=uninitialized.id, - role=TenantAccountRole.NORMAL, - ), - TenantAccountJoin( - tenant_id=workspace.id, - account_id=pending.id, - role=TenantAccountRole.NORMAL, - ), - TenantAccountJoin( - tenant_id=workspace.id, - account_id=banned.id, - role=TenantAccountRole.ADMIN, - ), - TenantAccountJoin( - tenant_id=workspace.id, - account_id=closed.id, - role=TenantAccountRole.EDITOR, - ), - TenantAccountJoin( - tenant_id=other_workspace.id, - account_id=other_workspace_member.id, - role=TenantAccountRole.ADMIN, - ), - ] - ) - session.commit() - - result = WorkspaceRepository(sqlite_session_factory).list_for_workspace(workspace.id) - - by_id = {member.id: member for member in result} - assert set(by_id) == {"active", "uninitialized", "pending", "banned", "closed"} - assert by_id["active"] == WorkspaceMemberRecord( - id=active.id, - name=active.name, - email=active.email, - avatar=active.avatar, - last_login_at=created_at, - last_active_at=created_at, - created_at=created_at, - status=AccountStatus.ACTIVE.value, - legacy_role=TenantAccountRole.OWNER.value, - ) - assert by_id["uninitialized"].status == AccountStatus.UNINITIALIZED.value - assert by_id["pending"].status == AccountStatus.PENDING.value - assert by_id["pending"].legacy_role == TenantAccountRole.NORMAL.value - assert by_id["banned"].status == AccountStatus.BANNED.value - assert by_id["closed"].status == AccountStatus.CLOSED.value - - -def test_list_for_workspace_returns_empty_tuple_without_membership( - sqlite_session_factory: sessionmaker[Session], -) -> None: - with sqlite_session_factory() as session: - session.add(make_tenant("workspace-1")) - session.commit() - - result = WorkspaceRepository(sqlite_session_factory).list_for_workspace("workspace-1") - - assert result == () - - -def test_dataset_operators_are_filtered_by_sql_within_workspace( - sqlite_session_factory: sessionmaker[Session], sqlite_engine: Engine -) -> None: - workspace = make_tenant("workspace-1") - other = make_tenant("workspace-2") - created_at = datetime(2026, 1, 1) - with sqlite_session_factory.begin() as session: - session.add_all([workspace, other]) - for account_id, tenant_id, role in [ - ("operator", workspace.id, TenantAccountRole.DATASET_OPERATOR), - ("editor", workspace.id, TenantAccountRole.EDITOR), - ("foreign-operator", other.id, TenantAccountRole.DATASET_OPERATOR), - ]: - session.add(make_account(account_id, status=AccountStatus.ACTIVE, created_at=created_at)) - session.add(TenantAccountJoin(account_id=account_id, tenant_id=tenant_id, role=role)) - - queries: list[str] = [] - - def capture_sql(_connection, _cursor, statement: str, _parameters, _context, _executemany) -> None: - queries.append(statement) - - event.listen(sqlite_engine, "before_cursor_execute", capture_sql) - try: - records = WorkspaceRepository(sqlite_session_factory).list_for_workspace( - workspace.id, role=TenantAccountRole.DATASET_OPERATOR - ) - finally: - event.remove(sqlite_engine, "before_cursor_execute", capture_sql) - - assert [record.id for record in records] == ["operator"] - assert len(queries) == 1 - assert "tenant_account_joins.role =" in queries[0].partition("WHERE")[2] - - -def test_member_queries_observe_committed_roles_and_release_session_before_role_resolution( - sqlite_session_factory: sessionmaker[Session], sqlite_engine: Engine -) -> None: - with sqlite_session_factory.begin() as session: - session.add(make_tenant("workspace")) - session.add(make_account("member", status=AccountStatus.ACTIVE, created_at=datetime(2026, 1, 1))) - session.add(TenantAccountJoin(tenant_id="workspace", account_id="member", role=TenantAccountRole.NORMAL)) - - checked_out: set[int] = set() - - def checkout(connection: object, _record: object, _proxy: object) -> None: - checked_out.add(id(connection)) - - def checkin(connection: object, _record: object) -> None: - checked_out.discard(id(connection)) - - class RoleResolver: - def resolve_many( - self, - workspace_id: str, - actor_account_id: str, - subjects: Sequence[WorkspaceMemberRoleSubject], - *, - language: str | None = None, - ) -> Mapping[str, Sequence[WorkspaceMemberRole]]: - assert not checked_out - assert language is None - assert (workspace_id, actor_account_id) == ("workspace", "actor") - assert subjects == (WorkspaceMemberRoleSubject("member", "editor"),) - return {"member": (WorkspaceMemberRole("editor", "Editor"),)} - - event.listen(sqlite_engine, "checkout", checkout) - event.listen(sqlite_engine, "checkin", checkin) - try: - repository = WorkspaceRepository(sqlite_session_factory) - repository.update_member_role( - workspace_id="workspace", account_id="member", role=TenantAccountRole.EDITOR, actor_account_id="owner" - ) - service = WorkspaceMemberQueryService(members=repository, roles=RoleResolver()) - result = service.list_current(RequestContext("request", None, "actor", "workspace")) - assert result[0].role == "editor" - assert result[0].roles == (WorkspaceMemberRole("editor", "Editor"),) - assert not checked_out - finally: - event.remove(sqlite_engine, "checkout", checkout) - event.remove(sqlite_engine, "checkin", checkin) - - -def test_get_legacy_role_is_scoped_to_workspace_and_account( - sqlite_session_factory: sessionmaker[Session], -) -> None: - created_at = datetime(2026, 1, 1) - account = make_account("member", status=AccountStatus.ACTIVE, created_at=created_at) - workspace = make_tenant("workspace-1") - other_workspace = make_tenant("workspace-2") - with sqlite_session_factory.begin() as session: - session.add_all( - [ - account, - workspace, - other_workspace, - TenantAccountJoin( - tenant_id=workspace.id, - account_id=account.id, - role=TenantAccountRole.DATASET_OPERATOR, - ), - ] - ) - - repository = WorkspaceRepository(sqlite_session_factory) - - assert ( - repository.get_legacy_role(workspace_id=workspace.id, account_id=account.id) - == TenantAccountRole.DATASET_OPERATOR.value - ) - assert repository.get_legacy_role(workspace_id=other_workspace.id, account_id=account.id) is None - assert repository.get_legacy_role(workspace_id=workspace.id, account_id="other-account") is None diff --git a/api/tests/unit_tests/repositories/workspace/test_workspace_repository.py b/api/tests/unit_tests/repositories/workspace/test_workspace_repository.py index ff8059a3849..e7b13046906 100644 --- a/api/tests/unit_tests/repositories/workspace/test_workspace_repository.py +++ b/api/tests/unit_tests/repositories/workspace/test_workspace_repository.py @@ -1,6 +1,4 @@ -"""Real SQLite coverage for workspace ownership, persistence and Session lifetime.""" - -from collections.abc import Sequence +from collections.abc import Mapping, Sequence from datetime import datetime, timedelta from unittest.mock import Mock, create_autospec @@ -9,10 +7,18 @@ from sqlalchemy import Engine, event, select from sqlalchemy.orm import Session, sessionmaker from machinery.context import RequestContext -from models.account import Account, Tenant, TenantAccountJoin, TenantAccountRole, TenantStatus +from models.account import Account, AccountStatus, Tenant, TenantAccountJoin, TenantAccountRole, TenantStatus from repositories.workspace.workspace_repository import WorkspaceRepository from services.errors.workspace import WorkspaceNotFoundError, WorkspaceNotLinkedError -from services.workspace.contracts import EffectiveCreditPool, WorkspaceCustomConfig, WorkspaceFeatures +from services.workspace.contracts import ( + EffectiveCreditPool, + WorkspaceCustomConfig, + WorkspaceFeatures, + WorkspaceMemberRecord, + WorkspaceMemberRole, + WorkspaceMemberRoleSubject, +) +from services.workspace.member_service import WorkspaceMemberQueryService from services.workspace.service import ( WorkspaceFeatureGateway, WorkspaceLogoGateway, @@ -21,183 +27,440 @@ from services.workspace.service import ( ) -@pytest.fixture -def seeded(sqlite_session_factory: sessionmaker[Session]) -> WorkspaceRepository: - with sqlite_session_factory() as session: - now = datetime(2026, 1, 1) - account = Account(name="Member", email="member@example.com") - account.id = "a1" - first = Tenant(name="First", status=TenantStatus.NORMAL) - first.id, first.created_at = ("w1", now - timedelta(days=1)) - second = Tenant(name="Second", status=TenantStatus.NORMAL) - second.id, second.created_at = ("w2", now) - archived = Tenant(name="Archived", status=TenantStatus.ARCHIVE) - archived.id, archived.created_at = ("w3", now) - session.add_all( - [ - account, - first, - second, - archived, - TenantAccountJoin(tenant_id="w1", account_id="a1", role=TenantAccountRole.NORMAL, current=True), - TenantAccountJoin(tenant_id="w2", account_id="a1", role=TenantAccountRole.EDITOR), - TenantAccountJoin(tenant_id="w3", account_id="a1", role=TenantAccountRole.NORMAL), - TenantAccountJoin(tenant_id="w2", account_id="other", role=TenantAccountRole.OWNER), - ] - ) - session.commit() - return WorkspaceRepository(sqlite_session_factory) - - -def test_names_by_ids_returns_only_existing_requested_workspaces(seeded: WorkspaceRepository) -> None: - assert seeded.names_by_ids(["w1", "w3", "missing", "w1"]) == {"w1": "First", "w3": "Archived"} - assert seeded.names_by_ids(["missing"]) == {} - - -def test_names_by_ids_short_circuits_empty_input(unbound_session_factory: sessionmaker[Session]) -> None: - repository = WorkspaceRepository(unbound_session_factory) - assert repository.names_by_ids([]) == {} - - -def test_lists_preserve_visibility_order_and_archived_memberships(seeded: WorkspaceRepository) -> None: - assert [row.id for row in seeded.list_for_account("a1")] == ["w1", "w2"] - assert seeded.list_for_account("absent") == () - assert set(seeded.list_ids_for_account("a1")) == {"w1", "w2", "w3"} - assert {row.id for row in seeded.list_account_access_workspaces("a1")} == {"w1", "w2", "w3"} - assert [(row.id, row.name, row.role, row.current) for row in seeded.list_for_device_flow("a1")] == [ - ("w1", "First", "normal", True), - ("w2", "Second", "editor", False), - ("w3", "Archived", "normal", False), - ] - assert seeded.get_many(()) == () - assert seeded.has_active_membership("a1") - assert not seeded.has_active_for_account("absent") - page = seeded.list_all(page=1, limit=2) - assert page.total == 3 - assert page.has_more - assert seeded.list_all(page=3, limit=2).data == () - - -def test_snapshot_role_is_scoped_to_both_account_and_workspace(seeded: WorkspaceRepository) -> None: - snapshot = seeded.get_for_account("w2", "a1") - assert snapshot is not None - assert snapshot.role == "editor" - assert snapshot.has_privileged_member - with pytest.raises(WorkspaceNotFoundError): - seeded.get_for_account("w1", "other") - assert seeded.get_for_account("missing", "a1") is None - - -def test_membership_role_is_scoped_to_both_account_and_workspace(seeded: WorkspaceRepository) -> None: - assert seeded.get_account_role(tenant_id="w1", account_id="a1") == "normal" - assert seeded.get_account_role(tenant_id="w2", account_id="a1") == "editor" - assert seeded.get_account_role(tenant_id="w2", account_id="other") == "owner" - assert seeded.get_account_role(tenant_id="w1", account_id="other") is None - assert seeded.get_account_role(tenant_id="w2", account_id="missing") is None - - -def test_switch_updates_only_the_requesting_account( - seeded: WorkspaceRepository, sqlite_session_factory: sessionmaker[Session] -) -> None: - seeded.switch(account_id="a1", workspace_id="w2") - with sqlite_session_factory() as session: - memberships = session.scalars(select(TenantAccountJoin).where(TenantAccountJoin.account_id == "a1")).all() - assert [row.tenant_id for row in memberships if row.current] == ["w2"] - assert next(row for row in memberships if row.tenant_id == "w2").last_opened_at is not None - other = session.scalar(select(TenantAccountJoin).where(TenantAccountJoin.account_id == "other")) - assert other is not None - assert not other.current - - -@pytest.mark.parametrize("workspace", ["w3", "missing"]) -def test_switch_rejects_archived_and_unlinked_without_changing_current( - seeded: WorkspaceRepository, sqlite_session_factory: sessionmaker[Session], workspace: str -) -> None: - with pytest.raises(WorkspaceNotLinkedError): - seeded.switch(account_id="a1", workspace_id=workspace) - with sqlite_session_factory() as session: - current = session.scalar( - select(TenantAccountJoin.tenant_id).where( - TenantAccountJoin.account_id == "a1", TenantAccountJoin.current.is_(True) - ) - ) - assert current == "w1" - - -def test_writes_commit_and_other_accounts_cannot_mutate( - seeded: WorkspaceRepository, sqlite_session_factory: sessionmaker[Session] -) -> None: - assert seeded.rename(workspace_id="w1", account_id="a1", name="New").name == "New" - config = WorkspaceCustomConfig(False, "logo") - seeded.update_custom_config(workspace_id="w1", account_id="a1", changes=config) - with pytest.raises(WorkspaceNotFoundError): - seeded.rename(workspace_id="w1", account_id="other", name="Unauthorized") - with sqlite_session_factory() as session: - tenant = session.get(Tenant, "w1") - assert tenant is not None - assert tenant.name == "New" - assert tenant.custom_config_dict == {"remove_webapp_brand": False, "replace_webapp_logo": "logo"} - - -def test_write_rolls_back_on_commit_failure( - seeded: WorkspaceRepository, sqlite_session_factory: sessionmaker[Session] -) -> None: - def fail_flush(_session: Session, _flush_context: object, _instances: object) -> None: - raise RuntimeError("write failed") - - event.listen(sqlite_session_factory, "before_flush", fail_flush) - try: - with pytest.raises(RuntimeError, match="write failed"): - seeded.rename(workspace_id="w1", account_id="a1", name="Failed") - finally: - event.remove(sqlite_session_factory, "before_flush", fail_flush) - snapshot = seeded.get_for_account("w1", "a1") - assert snapshot is not None - assert snapshot.name == "First" - - -def test_sessions_are_closed_before_feature_and_plan_io( - seeded: WorkspaceRepository, sqlite_session_factory: sessionmaker[Session], sqlite_engine: Engine -) -> None: - checked_out = set() - - def checkout(connection: object, _record: object, _proxy: object) -> None: - checked_out.add(id(connection)) - - def checkin(connection: object, _record: object) -> None: - checked_out.discard(id(connection)) - - event.listen(sqlite_engine, "checkout", checkout) - event.listen(sqlite_engine, "checkin", checkin) - features = create_autospec(WorkspaceFeatureGateway, instance=True) - - def get_features(workspace_id: str) -> WorkspaceFeatures: - assert not checked_out - with sqlite_session_factory() as session: - tenant = session.get(Tenant, workspace_id) - assert tenant is not None - assert tenant.name == "Committed" - return WorkspaceFeatures(False, EffectiveCreditPool()) - - features.get_features.side_effect = get_features - service = WorkspaceService( - workspaces=seeded, features=features, logos=create_autospec(WorkspaceLogoGateway, instance=True) +def make_account( + account_id: str, + *, + status: AccountStatus, + created_at: datetime, +) -> Account: + account = Account( + name=f"Member {account_id}", + email=f"{account_id}@example.com", + avatar=f"{account_id}.png", + status=status, ) - context = RequestContext("request", None, "a1", "w1") - try: - service.rename(context, "Committed") - plans = Mock() + account.id = account_id + account.last_login_at = created_at + account.last_active_at = created_at + account.created_at = created_at + return account - def resolve_many(_ids: Sequence[str]) -> dict[str, str]: + +def make_tenant(tenant_id: str) -> Tenant: + tenant = Tenant(name=f"Workspace {tenant_id}") + tenant.id = tenant_id + return tenant + + +class TestWorkspaceMembersCases: + @staticmethod + def test_rbac_profile_batch_filters_requested_ids_and_workspace( + sqlite_session_factory: sessionmaker[Session], + ) -> None: + with sqlite_session_factory.begin() as session: + session.add_all([make_tenant("workspace"), make_tenant("other")]) + for account_id, workspace in [("member", "workspace"), ("unrequested", "workspace"), ("foreign", "other")]: + session.add(make_account(account_id, status=AccountStatus.ACTIVE, created_at=datetime(2026, 1, 1))) + session.add( + TenantAccountJoin(tenant_id=workspace, account_id=account_id, role=TenantAccountRole.NORMAL) + ) + + repository = WorkspaceRepository(sqlite_session_factory) + members = repository.list_for_workspace("workspace", account_ids=["member", "foreign", "missing"]) + assert [member.id for member in members] == ["member"] + assert members[0].email == "member@example.com" + assert repository.list_for_workspace("workspace", account_ids=[]) == () + + @staticmethod + def test_list_for_workspace_uses_join_membership_and_preserves_account_lifecycle( + sqlite_session_factory: sessionmaker[Session], + ) -> None: + created_at = datetime(2026, 1, 1) + active = make_account("active", status=AccountStatus.ACTIVE, created_at=created_at) + uninitialized = make_account("uninitialized", status=AccountStatus.UNINITIALIZED, created_at=created_at) + pending = make_account("pending", status=AccountStatus.PENDING, created_at=created_at) + banned = make_account("banned", status=AccountStatus.BANNED, created_at=created_at) + closed = make_account("closed", status=AccountStatus.CLOSED, created_at=created_at) + other_workspace_member = make_account("other", status=AccountStatus.ACTIVE, created_at=created_at) + unjoined = make_account("unjoined", status=AccountStatus.ACTIVE, created_at=created_at) + workspace = make_tenant("workspace-1") + other_workspace = make_tenant("workspace-2") + + with sqlite_session_factory() as session: + session.add_all( + [ + workspace, + other_workspace, + active, + uninitialized, + pending, + banned, + closed, + other_workspace_member, + unjoined, + TenantAccountJoin( + tenant_id=workspace.id, + account_id=active.id, + role=TenantAccountRole.OWNER, + ), + TenantAccountJoin( + tenant_id=workspace.id, + account_id=uninitialized.id, + role=TenantAccountRole.NORMAL, + ), + TenantAccountJoin( + tenant_id=workspace.id, + account_id=pending.id, + role=TenantAccountRole.NORMAL, + ), + TenantAccountJoin( + tenant_id=workspace.id, + account_id=banned.id, + role=TenantAccountRole.ADMIN, + ), + TenantAccountJoin( + tenant_id=workspace.id, + account_id=closed.id, + role=TenantAccountRole.EDITOR, + ), + TenantAccountJoin( + tenant_id=other_workspace.id, + account_id=other_workspace_member.id, + role=TenantAccountRole.ADMIN, + ), + ] + ) + session.commit() + + result = WorkspaceRepository(sqlite_session_factory).list_for_workspace(workspace.id) + + by_id = {member.id: member for member in result} + assert set(by_id) == {"active", "uninitialized", "pending", "banned", "closed"} + assert by_id["active"] == WorkspaceMemberRecord( + id=active.id, + name=active.name, + email=active.email, + avatar=active.avatar, + last_login_at=created_at, + last_active_at=created_at, + created_at=created_at, + status=AccountStatus.ACTIVE.value, + legacy_role=TenantAccountRole.OWNER.value, + ) + assert by_id["uninitialized"].status == AccountStatus.UNINITIALIZED.value + assert by_id["pending"].status == AccountStatus.PENDING.value + assert by_id["pending"].legacy_role == TenantAccountRole.NORMAL.value + assert by_id["banned"].status == AccountStatus.BANNED.value + assert by_id["closed"].status == AccountStatus.CLOSED.value + + @staticmethod + def test_list_for_workspace_returns_empty_tuple_without_membership( + sqlite_session_factory: sessionmaker[Session], + ) -> None: + with sqlite_session_factory() as session: + session.add(make_tenant("workspace-1")) + session.commit() + + result = WorkspaceRepository(sqlite_session_factory).list_for_workspace("workspace-1") + + assert result == () + + @staticmethod + def test_dataset_operators_are_filtered_by_sql_within_workspace( + sqlite_session_factory: sessionmaker[Session], sqlite_engine: Engine + ) -> None: + workspace = make_tenant("workspace-1") + other = make_tenant("workspace-2") + created_at = datetime(2026, 1, 1) + with sqlite_session_factory.begin() as session: + session.add_all([workspace, other]) + for account_id, tenant_id, role in [ + ("operator", workspace.id, TenantAccountRole.DATASET_OPERATOR), + ("editor", workspace.id, TenantAccountRole.EDITOR), + ("foreign-operator", other.id, TenantAccountRole.DATASET_OPERATOR), + ]: + session.add(make_account(account_id, status=AccountStatus.ACTIVE, created_at=created_at)) + session.add(TenantAccountJoin(account_id=account_id, tenant_id=tenant_id, role=role)) + + queries: list[str] = [] + + def capture_sql(_connection, _cursor, statement: str, _parameters, _context, _executemany) -> None: + queries.append(statement) + + event.listen(sqlite_engine, "before_cursor_execute", capture_sql) + try: + records = WorkspaceRepository(sqlite_session_factory).list_for_workspace( + workspace.id, role=TenantAccountRole.DATASET_OPERATOR + ) + finally: + event.remove(sqlite_engine, "before_cursor_execute", capture_sql) + + assert [record.id for record in records] == ["operator"] + assert len(queries) == 1 + assert "tenant_account_joins.role =" in queries[0].partition("WHERE")[2] + + @staticmethod + def test_member_queries_observe_committed_roles_and_release_session_before_role_resolution( + sqlite_session_factory: sessionmaker[Session], sqlite_engine: Engine + ) -> None: + with sqlite_session_factory.begin() as session: + session.add(make_tenant("workspace")) + session.add(make_account("member", status=AccountStatus.ACTIVE, created_at=datetime(2026, 1, 1))) + session.add(TenantAccountJoin(tenant_id="workspace", account_id="member", role=TenantAccountRole.NORMAL)) + + checked_out: set[int] = set() + + def checkout(connection: object, _record: object, _proxy: object) -> None: + checked_out.add(id(connection)) + + def checkin(connection: object, _record: object) -> None: + checked_out.discard(id(connection)) + + class RoleResolver: + def resolve_many( + self, + workspace_id: str, + actor_account_id: str, + subjects: Sequence[WorkspaceMemberRoleSubject], + *, + language: str | None = None, + ) -> Mapping[str, Sequence[WorkspaceMemberRole]]: + assert not checked_out + assert language is None + assert (workspace_id, actor_account_id) == ("workspace", "actor") + assert subjects == (WorkspaceMemberRoleSubject("member", "editor"),) + return {"member": (WorkspaceMemberRole("editor", "Editor"),)} + + event.listen(sqlite_engine, "checkout", checkout) + event.listen(sqlite_engine, "checkin", checkin) + try: + repository = WorkspaceRepository(sqlite_session_factory) + repository.update_member_role( + workspace_id="workspace", account_id="member", role=TenantAccountRole.EDITOR, actor_account_id="owner" + ) + service = WorkspaceMemberQueryService(members=repository, roles=RoleResolver()) + result = service.list_current(RequestContext("request", None, "actor", "workspace")) + assert result[0].role == "editor" + assert result[0].roles == (WorkspaceMemberRole("editor", "Editor"),) assert not checked_out - return {"w1": "team"} + finally: + event.remove(sqlite_engine, "checkout", checkout) + event.remove(sqlite_engine, "checkin", checkin) - plans.resolve_many.side_effect = resolve_many - result = WorkspaceQueryService(workspaces=seeded, plans=plans).list_for_account(context) - assert result[0].current - assert result[0].plan == "team" - assert result[1].plan == "sandbox" - assert not checked_out - finally: - event.remove(sqlite_engine, "checkout", checkout) - event.remove(sqlite_engine, "checkin", checkin) + @staticmethod + def test_get_legacy_role_is_scoped_to_workspace_and_account( + sqlite_session_factory: sessionmaker[Session], + ) -> None: + created_at = datetime(2026, 1, 1) + account = make_account("member", status=AccountStatus.ACTIVE, created_at=created_at) + workspace = make_tenant("workspace-1") + other_workspace = make_tenant("workspace-2") + with sqlite_session_factory.begin() as session: + session.add_all( + [ + account, + workspace, + other_workspace, + TenantAccountJoin( + tenant_id=workspace.id, + account_id=account.id, + role=TenantAccountRole.DATASET_OPERATOR, + ), + ] + ) + + repository = WorkspaceRepository(sqlite_session_factory) + + assert ( + repository.get_legacy_role(workspace_id=workspace.id, account_id=account.id) + == TenantAccountRole.DATASET_OPERATOR.value + ) + assert repository.get_legacy_role(workspace_id=other_workspace.id, account_id=account.id) is None + assert repository.get_legacy_role(workspace_id=workspace.id, account_id="other-account") is None + + +"""Real SQLite coverage for workspace ownership, persistence and Session lifetime.""" + + +class TestWorkspaceRepositoryCases: + @staticmethod + @pytest.fixture + def seeded(sqlite_session_factory: sessionmaker[Session]) -> WorkspaceRepository: + with sqlite_session_factory() as session: + now = datetime(2026, 1, 1) + account = Account(name="Member", email="member@example.com") + account.id = "a1" + first = Tenant(name="First", status=TenantStatus.NORMAL) + first.id, first.created_at = ("w1", now - timedelta(days=1)) + second = Tenant(name="Second", status=TenantStatus.NORMAL) + second.id, second.created_at = ("w2", now) + archived = Tenant(name="Archived", status=TenantStatus.ARCHIVE) + archived.id, archived.created_at = ("w3", now) + session.add_all( + [ + account, + first, + second, + archived, + TenantAccountJoin(tenant_id="w1", account_id="a1", role=TenantAccountRole.NORMAL, current=True), + TenantAccountJoin(tenant_id="w2", account_id="a1", role=TenantAccountRole.EDITOR), + TenantAccountJoin(tenant_id="w3", account_id="a1", role=TenantAccountRole.NORMAL), + TenantAccountJoin(tenant_id="w2", account_id="other", role=TenantAccountRole.OWNER), + ] + ) + session.commit() + return WorkspaceRepository(sqlite_session_factory) + + @staticmethod + def test_names_by_ids_returns_only_existing_requested_workspaces(seeded: WorkspaceRepository) -> None: + assert seeded.names_by_ids(["w1", "w3", "missing", "w1"]) == {"w1": "First", "w3": "Archived"} + assert seeded.names_by_ids(["missing"]) == {} + + @staticmethod + def test_names_by_ids_short_circuits_empty_input(unbound_session_factory: sessionmaker[Session]) -> None: + repository = WorkspaceRepository(unbound_session_factory) + assert repository.names_by_ids([]) == {} + + @staticmethod + def test_lists_preserve_visibility_order_and_archived_memberships(seeded: WorkspaceRepository) -> None: + assert [row.id for row in seeded.list_for_account("a1")] == ["w1", "w2"] + assert seeded.list_for_account("absent") == () + assert set(seeded.list_ids_for_account("a1")) == {"w1", "w2", "w3"} + assert {row.id for row in seeded.list_account_access_workspaces("a1")} == {"w1", "w2", "w3"} + assert [(row.id, row.name, row.role, row.current) for row in seeded.list_for_device_flow("a1")] == [ + ("w1", "First", "normal", True), + ("w2", "Second", "editor", False), + ("w3", "Archived", "normal", False), + ] + assert seeded.get_many(()) == () + assert seeded.has_active_membership("a1") + assert not seeded.has_active_for_account("absent") + page = seeded.list_all(page=1, limit=2) + assert page.total == 3 + assert page.has_more + assert seeded.list_all(page=3, limit=2).data == () + + @staticmethod + def test_snapshot_role_is_scoped_to_both_account_and_workspace(seeded: WorkspaceRepository) -> None: + snapshot = seeded.get_for_account("w2", "a1") + assert snapshot is not None + assert snapshot.role == "editor" + assert snapshot.has_privileged_member + with pytest.raises(WorkspaceNotFoundError): + seeded.get_for_account("w1", "other") + assert seeded.get_for_account("missing", "a1") is None + + @staticmethod + def test_membership_role_is_scoped_to_both_account_and_workspace(seeded: WorkspaceRepository) -> None: + assert seeded.get_account_role(tenant_id="w1", account_id="a1") == "normal" + assert seeded.get_account_role(tenant_id="w2", account_id="a1") == "editor" + assert seeded.get_account_role(tenant_id="w2", account_id="other") == "owner" + assert seeded.get_account_role(tenant_id="w1", account_id="other") is None + assert seeded.get_account_role(tenant_id="w2", account_id="missing") is None + + @staticmethod + def test_switch_updates_only_the_requesting_account( + seeded: WorkspaceRepository, sqlite_session_factory: sessionmaker[Session] + ) -> None: + seeded.switch(account_id="a1", workspace_id="w2") + with sqlite_session_factory() as session: + memberships = session.scalars(select(TenantAccountJoin).where(TenantAccountJoin.account_id == "a1")).all() + assert [row.tenant_id for row in memberships if row.current] == ["w2"] + assert next(row for row in memberships if row.tenant_id == "w2").last_opened_at is not None + other = session.scalar(select(TenantAccountJoin).where(TenantAccountJoin.account_id == "other")) + assert other is not None + assert not other.current + + @staticmethod + @pytest.mark.parametrize("workspace", ["w3", "missing"]) + def test_switch_rejects_archived_and_unlinked_without_changing_current( + seeded: WorkspaceRepository, sqlite_session_factory: sessionmaker[Session], workspace: str + ) -> None: + with pytest.raises(WorkspaceNotLinkedError): + seeded.switch(account_id="a1", workspace_id=workspace) + with sqlite_session_factory() as session: + current = session.scalar( + select(TenantAccountJoin.tenant_id).where( + TenantAccountJoin.account_id == "a1", TenantAccountJoin.current.is_(True) + ) + ) + assert current == "w1" + + @staticmethod + def test_writes_commit_and_other_accounts_cannot_mutate( + seeded: WorkspaceRepository, sqlite_session_factory: sessionmaker[Session] + ) -> None: + assert seeded.rename(workspace_id="w1", account_id="a1", name="New").name == "New" + config = WorkspaceCustomConfig(False, "logo") + seeded.update_custom_config(workspace_id="w1", account_id="a1", changes=config) + with pytest.raises(WorkspaceNotFoundError): + seeded.rename(workspace_id="w1", account_id="other", name="Unauthorized") + with sqlite_session_factory() as session: + tenant = session.get(Tenant, "w1") + assert tenant is not None + assert tenant.name == "New" + assert tenant.custom_config_dict == {"remove_webapp_brand": False, "replace_webapp_logo": "logo"} + + @staticmethod + def test_write_rolls_back_on_commit_failure( + seeded: WorkspaceRepository, sqlite_session_factory: sessionmaker[Session] + ) -> None: + def fail_flush(_session: Session, _flush_context: object, _instances: object) -> None: + raise RuntimeError("write failed") + + event.listen(sqlite_session_factory, "before_flush", fail_flush) + try: + with pytest.raises(RuntimeError, match="write failed"): + seeded.rename(workspace_id="w1", account_id="a1", name="Failed") + finally: + event.remove(sqlite_session_factory, "before_flush", fail_flush) + snapshot = seeded.get_for_account("w1", "a1") + assert snapshot is not None + assert snapshot.name == "First" + + @staticmethod + def test_sessions_are_closed_before_feature_and_plan_io( + seeded: WorkspaceRepository, sqlite_session_factory: sessionmaker[Session], sqlite_engine: Engine + ) -> None: + checked_out = set() + + def checkout(connection: object, _record: object, _proxy: object) -> None: + checked_out.add(id(connection)) + + def checkin(connection: object, _record: object) -> None: + checked_out.discard(id(connection)) + + event.listen(sqlite_engine, "checkout", checkout) + event.listen(sqlite_engine, "checkin", checkin) + features = create_autospec(WorkspaceFeatureGateway, instance=True) + + def get_features(workspace_id: str) -> WorkspaceFeatures: + assert not checked_out + with sqlite_session_factory() as session: + tenant = session.get(Tenant, workspace_id) + assert tenant is not None + assert tenant.name == "Committed" + return WorkspaceFeatures(False, EffectiveCreditPool()) + + features.get_features.side_effect = get_features + service = WorkspaceService( + workspaces=seeded, features=features, logos=create_autospec(WorkspaceLogoGateway, instance=True) + ) + context = RequestContext("request", None, "a1", "w1") + try: + service.rename(context, "Committed") + plans = Mock() + + def resolve_many(_ids: Sequence[str]) -> dict[str, str]: + assert not checked_out + return {"w1": "team"} + + plans.resolve_many.side_effect = resolve_many + result = WorkspaceQueryService(workspaces=seeded, plans=plans).list_for_account(context) + assert result[0].current + assert result[0].plan == "team" + assert result[1].plan == "sandbox" + assert not checked_out + finally: + event.remove(sqlite_engine, "checkout", checkout) + event.remove(sqlite_engine, "checkin", checkin) diff --git a/api/tests/unit_tests/schedule/test_collect_agent_sandbox_usage.py b/api/tests/unit_tests/schedule/test_collect_agent_sandbox_usage.py index a9c0691b7e5..e1d154b1fe0 100644 --- a/api/tests/unit_tests/schedule/test_collect_agent_sandbox_usage.py +++ b/api/tests/unit_tests/schedule/test_collect_agent_sandbox_usage.py @@ -12,7 +12,7 @@ from configs.extra.agent_backend_config import AgentBackendConfig from dify_app import DifyApp from extensions import ext_celery from schedule.collect_agent_sandbox_usage import collect_agent_sandbox_usage -from tests.unit_tests.config_override import config_overrides_context +from tests.helpers.unit.config_override import config_overrides_context PROJECT = "431de237-596f-4d59-8a85-20a9846bf243" diff --git a/api/tests/unit_tests/services/account/test_adapters.py b/api/tests/unit_tests/services/account/test_adapters.py index 0e7be136fa8..381088f91f8 100644 --- a/api/tests/unit_tests/services/account/test_adapters.py +++ b/api/tests/unit_tests/services/account/test_adapters.py @@ -1,12 +1,19 @@ -from datetime import UTC, datetime +from datetime import UTC, datetime, timedelta from hashlib import sha256 -from unittest.mock import MagicMock, call, patch +from unittest.mock import MagicMock, call, create_autospec, patch import pytest +from redis import RedisError +from sqlalchemy import select +from sqlalchemy.orm import Session, sessionmaker +from werkzeug.exceptions import Unauthorized from extensions.ext_redis import RedisClientWrapper +from libs.datetime_utils import naive_utc_now from libs.helper import RateLimiter +from models.account import Account, AccountStatus, Tenant, TenantAccountJoin, TenantAccountRole, TenantStatus from services.account.adapters import ( + AccountIdentityGateway, BillingAccountActivationEligibility, BillingAccountEducationGateway, BillingAccountEmailPolicyGateway, @@ -27,6 +34,7 @@ from services.entities.account_entities import ( AccountEducationAutocomplete, AccountEducationStatus, ) +from tests.helpers.unit.account_domain import AccountDomain def test_invitation_token_store_reads_workspace_invitation_key( @@ -388,3 +396,80 @@ def test_verification_notifier_preserves_rate_limit_before_enqueuing_email( assert [call.args[0] for call in commands.call_args_list] == ["ZREMRANGEBYSCORE", "ZCARD"] assert error.value.retry_after_minutes == 1 mail_task.delay.assert_not_called() + + +@pytest.mark.parametrize("has_normal", [True, False]) +def test_identity_resolves_normal_workspace_and_detaches_it( + account_domain: AccountDomain, sqlite_session_factory: sessionmaker[Session], has_normal: bool +) -> None: + with sqlite_session_factory.begin() as session: + account = Account(name="User", email="u@example.com") + archived = Tenant(name="Archived", status=TenantStatus.ARCHIVE) + session.add_all( + [ + account, + archived, + TenantAccountJoin( + account_id=account.id, tenant_id=archived.id, role=TenantAccountRole.NORMAL, current=True + ), + ] + ) + if has_normal: + normal = Tenant(name="Normal") + session.add_all( + [normal, TenantAccountJoin(account_id=account.id, tenant_id=normal.id, role=TenantAccountRole.EDITOR)] + ) + account_id = account.id + loaded = account_domain.repository.load_identity(account_id, now=datetime(2026, 1, 1)) + if has_normal: + assert loaded is not None + assert loaded.current_tenant is not None + assert loaded.current_tenant.name == "Normal" + assert loaded.role == TenantAccountRole.EDITOR + else: + assert loaded is None + with sqlite_session_factory() as session: + current = session.scalars(select(TenantAccountJoin).where(TenantAccountJoin.current.is_(True))).all() + assert len(current) == int(has_normal) + + +@pytest.mark.parametrize("redis_claim", [True, False, RedisError("offline")]) +def test_activity_refresh_is_gated_but_fails_open( + account_domain: AccountDomain, sqlite_session: Session, redis_claim: bool | RedisError +) -> None: + now = naive_utc_now() + account = Account(name="User", email="u@example.com") + account.last_active_at = now - timedelta(hours=1) + tenant = Tenant(name="Normal") + sqlite_session.add_all( + [ + account, + tenant, + TenantAccountJoin(account_id=account.id, tenant_id=tenant.id, role=TenantAccountRole.OWNER, current=True), + ] + ) + sqlite_session.commit() + account_id, old_active = account.id, account.last_active_at + redis = create_autospec(RedisClientWrapper, instance=True) + if isinstance(redis_claim, RedisError): + redis.set.side_effect = redis_claim + else: + redis.set.return_value = redis_claim + gateway = AccountIdentityGateway(accounts=account_domain.repository, redis=redis) + loaded = gateway.load_user(account_id) + assert loaded is not None + sqlite_session.expire_all() + assert (account.last_active_at > old_active) == (redis_claim is not False) + assert redis.set.call_args.kwargs == {"ex": 600, "nx": True} + + +def test_identity_rejects_banned_and_missing_accounts(account_domain: AccountDomain, sqlite_session: Session) -> None: + account = Account(name="Banned", email="b@example.com", status=AccountStatus.BANNED) + sqlite_session.add(account) + sqlite_session.commit() + gateway = AccountIdentityGateway( + accounts=account_domain.repository, redis=create_autospec(RedisClientWrapper, instance=True) + ) + with pytest.raises(Unauthorized, match="banned"): + gateway.load_user(account.id) + assert gateway.load_user("missing") is None diff --git a/api/tests/unit_tests/services/account/test_email_registration_adapters.py b/api/tests/unit_tests/services/account/test_email_registration_adapters.py index e1e0aa5844a..1e107a6dae8 100644 --- a/api/tests/unit_tests/services/account/test_email_registration_adapters.py +++ b/api/tests/unit_tests/services/account/test_email_registration_adapters.py @@ -27,7 +27,7 @@ from services.entities.account_entities import ( AccountEmailRegistrationToken, AccountSessionTokens, ) -from tests.unit_tests.account_domain import AccountDomain +from tests.helpers.unit.account_domain import AccountDomain def test_token_gateway_rejects_malformed_payload() -> None: diff --git a/api/tests/unit_tests/services/account/test_email_registration_service.py b/api/tests/unit_tests/services/account/test_email_registration_service.py index 384930ffdd9..41acb30c12b 100644 --- a/api/tests/unit_tests/services/account/test_email_registration_service.py +++ b/api/tests/unit_tests/services/account/test_email_registration_service.py @@ -24,7 +24,7 @@ from services.entities.account_entities import ( AccountEmailRegistrationToken, AccountSessionTokens, ) -from tests.unit_tests.model_factories import make_account +from tests.helpers.unit.model_factories import make_account @pytest.fixture diff --git a/api/tests/unit_tests/services/account/test_identity.py b/api/tests/unit_tests/services/account/test_identity.py deleted file mode 100644 index 8e28de8ce7e..00000000000 --- a/api/tests/unit_tests/services/account/test_identity.py +++ /dev/null @@ -1,91 +0,0 @@ -from datetime import datetime, timedelta -from unittest.mock import create_autospec - -import pytest -from redis import RedisError -from sqlalchemy import select -from sqlalchemy.orm import Session, sessionmaker -from werkzeug.exceptions import Unauthorized - -from extensions.ext_redis import RedisClientWrapper -from libs.datetime_utils import naive_utc_now -from models.account import Account, AccountStatus, Tenant, TenantAccountJoin, TenantAccountRole, TenantStatus -from services.account.adapters import AccountIdentityGateway -from tests.unit_tests.account_domain import AccountDomain - - -@pytest.mark.parametrize("has_normal", [True, False]) -def test_identity_resolves_normal_workspace_and_detaches_it( - account_domain: AccountDomain, sqlite_session_factory: sessionmaker[Session], has_normal: bool -) -> None: - with sqlite_session_factory.begin() as session: - account = Account(name="User", email="u@example.com") - archived = Tenant(name="Archived", status=TenantStatus.ARCHIVE) - session.add_all( - [ - account, - archived, - TenantAccountJoin( - account_id=account.id, tenant_id=archived.id, role=TenantAccountRole.NORMAL, current=True - ), - ] - ) - if has_normal: - normal = Tenant(name="Normal") - session.add_all( - [normal, TenantAccountJoin(account_id=account.id, tenant_id=normal.id, role=TenantAccountRole.EDITOR)] - ) - account_id = account.id - loaded = account_domain.repository.load_identity(account_id, now=datetime(2026, 1, 1)) - if has_normal: - assert loaded is not None - assert loaded.current_tenant is not None - assert loaded.current_tenant.name == "Normal" - assert loaded.role == TenantAccountRole.EDITOR - else: - assert loaded is None - with sqlite_session_factory() as session: - current = session.scalars(select(TenantAccountJoin).where(TenantAccountJoin.current.is_(True))).all() - assert len(current) == int(has_normal) - - -@pytest.mark.parametrize("redis_claim", [True, False, RedisError("offline")]) -def test_activity_refresh_is_gated_but_fails_open( - account_domain: AccountDomain, sqlite_session: Session, redis_claim: bool | RedisError -) -> None: - now = naive_utc_now() - account = Account(name="User", email="u@example.com") - account.last_active_at = now - timedelta(hours=1) - tenant = Tenant(name="Normal") - sqlite_session.add_all( - [ - account, - tenant, - TenantAccountJoin(account_id=account.id, tenant_id=tenant.id, role=TenantAccountRole.OWNER, current=True), - ] - ) - sqlite_session.commit() - account_id, old_active = account.id, account.last_active_at - redis = create_autospec(RedisClientWrapper, instance=True) - if isinstance(redis_claim, RedisError): - redis.set.side_effect = redis_claim - else: - redis.set.return_value = redis_claim - gateway = AccountIdentityGateway(accounts=account_domain.repository, redis=redis) - loaded = gateway.load_user(account_id) - assert loaded is not None - sqlite_session.expire_all() - assert (account.last_active_at > old_active) == (redis_claim is not False) - assert redis.set.call_args.kwargs == {"ex": 600, "nx": True} - - -def test_identity_rejects_banned_and_missing_accounts(account_domain: AccountDomain, sqlite_session: Session) -> None: - account = Account(name="Banned", email="b@example.com", status=AccountStatus.BANNED) - sqlite_session.add(account) - sqlite_session.commit() - gateway = AccountIdentityGateway( - accounts=account_domain.repository, redis=create_autospec(RedisClientWrapper, instance=True) - ) - with pytest.raises(Unauthorized, match="banned"): - gateway.load_user(account.id) - assert gateway.load_user("missing") is None diff --git a/api/tests/unit_tests/services/account/test_login_adapters.py b/api/tests/unit_tests/services/account/test_login_adapters.py index eaac915d367..59cfbfb47c6 100644 --- a/api/tests/unit_tests/services/account/test_login_adapters.py +++ b/api/tests/unit_tests/services/account/test_login_adapters.py @@ -41,7 +41,7 @@ from services.entities.account_login_entities import ( RefreshAccountStatus, ) from services.turnstile_service import TurnstileChallengeRejectedError, TurnstileUpstreamError -from tests.unit_tests.account_domain import AccountDomain +from tests.helpers.unit.account_domain import AccountDomain from tests.unit_tests.services.account.test_login_service import Dependencies diff --git a/api/tests/unit_tests/services/account/test_oauth_adapters.py b/api/tests/unit_tests/services/account/test_oauth_adapters.py index 851edac4f3d..366f99b5f81 100644 --- a/api/tests/unit_tests/services/account/test_oauth_adapters.py +++ b/api/tests/unit_tests/services/account/test_oauth_adapters.py @@ -23,7 +23,7 @@ from services.account_errors import ( OAuthWorkspaceCreationNotAllowedError, ) from services.entities.account_oauth_entities import OAuthAccountRegistration, OAuthAuthorizationRequest, OAuthIdentity -from tests.unit_tests.account_domain import AccountDomain +from tests.helpers.unit.account_domain import AccountDomain class StubOAuthClient(OAuth): diff --git a/api/tests/unit_tests/services/account/test_service.py b/api/tests/unit_tests/services/account/test_service.py index a719ea5ec9a..d334f253d42 100644 --- a/api/tests/unit_tests/services/account/test_service.py +++ b/api/tests/unit_tests/services/account/test_service.py @@ -24,7 +24,7 @@ from services.account_errors import ( ) from services.entities.account_entities import AccountSessionTokens from services.errors.workspace import WorkSpaceNotAllowedCreateError -from tests.unit_tests.account_domain import AccountDomain +from tests.helpers.unit.account_domain import AccountDomain def test_creation_hashes_password_and_keeps_registration_ip( diff --git a/api/tests/unit_tests/services/agent/test_agent_observability_service.py b/api/tests/unit_tests/services/agent/test_agent_observability_service.py deleted file mode 100644 index ccfcc6e76e1..00000000000 --- a/api/tests/unit_tests/services/agent/test_agent_observability_service.py +++ /dev/null @@ -1,992 +0,0 @@ -import json -from datetime import UTC, datetime -from decimal import Decimal - -import pytest -from sqlalchemy import Select, select -from sqlalchemy.orm import Session - -from core.app.entities.app_invoke_entities import InvokeFrom -from graphon.enums import WorkflowNodeExecutionStatus -from libs.datetime_utils import naive_utc_now -from models.agent import WorkflowAgentBindingType, WorkflowAgentNodeBinding -from models.agent_config_entities import WorkflowNodeJobConfig -from models.enums import ( - ConversationFromSource, - CreatorUserRole, - FeedbackFromSource, - FeedbackRating, - MessageStatus, -) -from models.model import App, AppMode, Conversation, Message, MessageFeedback -from models.workflow import ( - Workflow, - WorkflowExecutionStatus, - WorkflowNodeExecutionModel, - WorkflowNodeExecutionTriggeredFrom, - WorkflowRun, - WorkflowRunTriggeredFrom, - WorkflowType, -) -from services.agent import observability_service as observability_service_module -from services.agent.observability_service import AgentLogQueryParams, AgentObservabilityService -from tests.unit_tests.config_override import apply_config_overrides -from tests.unit_tests.model_factories import make_app, make_conversation, make_message - - -def _app(*, app_id: str = "app-1", name: str = "Iris", mode: AppMode = AppMode.AGENT_CHAT) -> App: - return make_app(app_id=app_id, name=name, mode=mode, icon_background="#fff", enable_site=False, enable_api=False) - - -def _conversation(*, conversation_id: str = "conversation-1", app_id: str = "app-1") -> Conversation: - return make_conversation( - conversation_id=conversation_id, - app_id=app_id, - mode=AppMode.AGENT_CHAT, - name="Debug conversation", - inputs={}, - from_source=ConversationFromSource.CONSOLE, - from_end_user_id="end-user-1", - ) - - -def _message( - *, - message_id: str = "message-1", - conversation_id: str = "conversation-1", - app_id: str = "app-1", - created_at: datetime | None = None, -) -> Message: - timestamp = created_at or naive_utc_now() - message = make_message( - message_id=message_id, - app_id=app_id, - conversation_id=conversation_id, - inputs={}, - query="hello", - message={}, - answer="hi", - status=MessageStatus.NORMAL, - message_unit_price=Decimal(0), - answer_unit_price=Decimal(0), - total_price=Decimal("0.0001"), - currency="USD", - from_source=ConversationFromSource.CONSOLE, - from_account_id="account-1", - invoke_from=InvokeFrom.EXPLORE, - message_tokens=3, - answer_tokens=4, - provider_response_latency=1.25, - ) - message.created_at = timestamp - message.updated_at = timestamp - return message - - -def _feedback( - *, - message_id: str = "message-1", - conversation_id: str = "conversation-1", - source: FeedbackFromSource = FeedbackFromSource.USER, - rating: FeedbackRating = FeedbackRating.LIKE, - content: str | None = "Useful", -) -> MessageFeedback: - return MessageFeedback( - app_id="app-1", - conversation_id=conversation_id, - message_id=message_id, - rating=rating, - from_source=source, - content=content, - ) - - -def _workflow_run(*, workflow_type: WorkflowType = WorkflowType.WORKFLOW) -> WorkflowRun: - created_at = datetime(2026, 7, 21, 7, 0, 19) - return WorkflowRun( - id="workflow-run-1", - tenant_id="tenant-1", - app_id="workflow-app-1", - workflow_id="workflow-1", - type=workflow_type, - triggered_from=WorkflowRunTriggeredFrom.APP_RUN, - version="v1", - graph="{}", - inputs="{}", - status=WorkflowExecutionStatus.SUCCEEDED, - outputs="{}", - error=None, - elapsed_time=59.93, - total_tokens=454_064, - total_steps=1, - created_by_role=CreatorUserRole.END_USER, - created_by="end-user-1", - created_at=created_at, - finished_at=created_at, - ) - - -def _node_execution( - *, - execution_id: str = "node-execution-1", - status: WorkflowNodeExecutionStatus = WorkflowNodeExecutionStatus.SUCCEEDED, -) -> WorkflowNodeExecutionModel: - created_at = datetime(2026, 7, 23, 7, 0, 19, tzinfo=UTC) - return WorkflowNodeExecutionModel( - id=execution_id, - tenant_id="tenant-1", - app_id="workflow-app-1", - workflow_id="workflow-1", - triggered_from=WorkflowNodeExecutionTriggeredFrom.WORKFLOW_RUN, - workflow_run_id="workflow-run-1", - index=1, - predecessor_node_id=None, - node_execution_id=execution_id, - node_id="node-1", - node_type="agent", - title="Agent", - inputs="{}", - process_data="{}", - outputs="{}", - status=status, - error=None, - elapsed_time=59.93, - execution_metadata=json.dumps( - { - "agent_log": { - "agent_backend": { - "usage": { - "prompt_tokens": 451_938, - "completion_tokens": 2_126, - "total_tokens": 454_064, - "total_price": "2.323470", - "currency": "USD", - "latency": 59.93, - } - } - } - } - ), - created_at=created_at, - created_by_role=CreatorUserRole.END_USER, - created_by="end-user-1", - finished_at=None, - ) - - -def _workflow_binding( - *, app_id: str = "workflow-app-1", binding_id: str | None = None, node_id: str = "node-1" -) -> WorkflowAgentNodeBinding: - return WorkflowAgentNodeBinding( - id=binding_id or f"binding-{app_id}-{node_id}", - tenant_id="tenant-1", - app_id=app_id, - workflow_id="workflow-1", - workflow_version="v1", - node_id=node_id, - binding_type=WorkflowAgentBindingType.ROSTER_AGENT, - agent_id="agent-1", - current_snapshot_id="snapshot-1", - node_job_config=WorkflowNodeJobConfig(), - created_by="account-1", - ) - - -def test_resolve_source_accepts_frontend_aliases() -> None: - assert AgentObservabilityService.resolve_source(None) is None - assert AgentObservabilityService.resolve_source("all") is None - assert AgentObservabilityService.resolve_source("console") == InvokeFrom.EXPLORE - assert AgentObservabilityService.resolve_source("api") == InvokeFrom.SERVICE_API - assert AgentObservabilityService.resolve_source("web_app") == InvokeFrom.WEB_APP - - with pytest.raises(ValueError, match="Unsupported source"): - AgentObservabilityService.resolve_source("unknown") - - -def test_resolve_source_filter_accepts_structured_sources() -> None: - assert AgentObservabilityService.resolve_source_filter(None).kind == "all" - assert AgentObservabilityService.resolve_source_filter("webapp").kind == "webapp" - assert AgentObservabilityService.resolve_source_filter("webapp:app-1").app_id == "app-1" - - workflow_app_filter = AgentObservabilityService.resolve_source_filter("workflow:app-2") - assert workflow_app_filter.kind == "workflow" - assert workflow_app_filter.app_id == "app-2" - assert workflow_app_filter.workflow_id is None - - workflow_filter = AgentObservabilityService.resolve_source_filter("workflow:app-2:workflow-1:v1:node-1") - assert workflow_filter.kind == "workflow" - assert workflow_filter.app_id == "app-2" - assert workflow_filter.workflow_id == "workflow-1" - assert workflow_filter.workflow_version == "v1" - assert workflow_filter.node_id == "node-1" - - timestamp_version_filter = AgentObservabilityService.resolve_source_filter( - "workflow:app-2:workflow-1:2026-07-06 02:17:12.910515:node-1" - ) - assert timestamp_version_filter.workflow_version == "2026-07-06 02:17:12.910515" - assert timestamp_version_filter.node_id == "node-1" - - legacy_filter = AgentObservabilityService.resolve_source_filter("console") - assert legacy_filter.kind == "webapp" - assert legacy_filter.invoke_from == InvokeFrom.EXPLORE - - with pytest.raises(ValueError, match="Unsupported source"): - AgentObservabilityService.resolve_source_filter("workflow:") - with pytest.raises(ValueError, match="Unsupported source"): - AgentObservabilityService.resolve_source_filter("workflow:app-2:incomplete") - - -def test_resolve_source_filters_accepts_multiple_structured_sources() -> None: - filters = AgentObservabilityService.resolve_source_filters(("webapp:app-1", "workflow:app-2:workflow-1:v1:node-1")) - - assert [source_filter.kind for source_filter in filters] == ["webapp", "workflow"] - assert filters[0].app_id == "app-1" - assert filters[1].node_id == "node-1" - assert AgentObservabilityService.resolve_source_filters(())[0].kind == "all" - assert AgentObservabilityService.resolve_source_filters(("all", "webapp:app-1"))[0].kind == "all" - - -def test_statistics_all_source_includes_debugger_messages() -> None: - source_filter = AgentObservabilityService.resolve_source_filter("all") - - scope_sql = AgentObservabilityService._statistics_webapp_message_scope_sql(source_filter) - - assert "m.app_id = :app_id" in scope_sql - assert "m.invoke_from != :debugger" not in scope_sql - - -def test_statistics_explicit_source_filters_invoke_from() -> None: - source_filter = AgentObservabilityService.resolve_source_filter("debugger") - - scope_sql = AgentObservabilityService._statistics_webapp_message_scope_sql(source_filter) - - assert "m.invoke_from = :source" in scope_sql - - -def test_statistics_workflow_app_source_covers_all_versions_and_nodes() -> None: - source_filter = AgentObservabilityService.resolve_source_filter("workflow:app-2") - - scope_sql = AgentObservabilityService._statistics_workflow_binding_filters_sql(source_filter) - - assert "wanb.app_id = :source_app_id" in scope_sql - assert "wanb.workflow_id = :workflow_id" not in scope_sql - assert "wanb.workflow_version = :workflow_version" not in scope_sql - assert "wanb.node_id = :node_id" not in scope_sql - - -def test_workflow_metadata_numeric_sql_supports_postgresql_and_mysql(monkeypatch: pytest.MonkeyPatch) -> None: - apply_config_overrides(monkeypatch, DB_TYPE="postgresql") - - postgres_sql = AgentObservabilityService._workflow_execution_metadata_numeric_sql( - ("agent_log", "agent_backend", "usage", "total_tokens"), "BIGINT" - ) - - assert "CAST(wne.execution_metadata AS JSONB)" in postgres_sql - assert "#>> '{agent_log,agent_backend,usage,total_tokens}'" in postgres_sql - - apply_config_overrides(monkeypatch, DB_TYPE="mysql") - - mysql_sql = AgentObservabilityService._workflow_execution_metadata_numeric_sql(("total_tokens",), "BIGINT") - - assert "JSON_EXTRACT(wne.execution_metadata, '$.total_tokens')" in mysql_sql - assert " AS UNSIGNED)" in mysql_sql - - -def test_workflow_statistics_include_run_without_message( - monkeypatch: pytest.MonkeyPatch, sqlite_session: Session -) -> None: - workflow_app = _app(app_id="workflow-app-1", name="Workflow App", mode=AppMode.WORKFLOW) - sqlite_session.add_all([workflow_app, _workflow_run(), _node_execution(), _workflow_binding()]) - sqlite_session.commit() - - apply_config_overrides(monkeypatch, DB_TYPE="mysql") - monkeypatch.setattr(observability_service_module, "convert_datetime_to_date", lambda field: f"DATE({field})") - monkeypatch.setattr( - AgentObservabilityService, - "_workflow_execution_metadata_numeric_sql", - staticmethod( - lambda path, numeric_type: ( - f"CAST(json_extract(wne.execution_metadata, '$.{'.'.join(path)}') AS {numeric_type})" - ) - ), - ) - service = AgentObservabilityService(sqlite_session) - - payload = service.get_statistics_summary( - app=_app(app_id="agent-app"), - agent_id="agent-1", - params=observability_service_module.AgentStatisticsQueryParams(source="workflow:workflow-app-1"), - ) - - assert payload["summary"]["total_messages"] == 1 - assert payload["summary"]["total_conversations"] == 1 - assert payload["summary"]["total_end_users"] == 1 - assert payload["summary"]["total_tokens"] == 454_064 - assert Decimal(payload["summary"]["total_price"]) == Decimal("2.323470") - - -def test_merge_daily_statistics_combines_webapp_and_workflow_rows() -> None: - rows = [ - { - "date": "2026-07-21", - "message_count": 2, - "conversation_count": 1, - "end_user_count": 1, - "token_count": 30, - "total_price": Decimal("0.003"), - "avg_latency": 1.5, - "latency_sum": 3, - "answer_tokens": 12, - "like_count": 1, - }, - { - "date": "2026-07-21", - "message_count": 1, - "conversation_count": 1, - "end_user_count": 1, - "token_count": 20, - "total_price": Decimal("0.002"), - "avg_latency": 2, - "latency_sum": 2, - "answer_tokens": 8, - "like_count": 0, - }, - ] - - merged = AgentObservabilityService._merge_daily_statistics(rows) - - assert merged == [ - { - "date": "2026-07-21", - "message_count": 3, - "conversation_count": 2, - "end_user_count": 2, - "token_count": 50, - "total_price": Decimal("0.005"), - "avg_latency": pytest.approx(5 / 3), - "latency_sum": 5.0, - "answer_tokens": 20, - "like_count": 1, - } - ] - - -def test_apply_status_filter_accepts_multiple_statuses() -> None: - stmt = select(Message) - - result = AgentObservabilityService._apply_status_filter(stmt, ("success", "failed", "paused")) - - assert isinstance(result, Select) - assert len(result._where_criteria) == 1 - with pytest.raises(ValueError, match="Unsupported status"): - AgentObservabilityService._apply_status_filter(select(Message), ("unknown",)) - - -def test_list_logs_sorts_by_requested_field(monkeypatch: pytest.MonkeyPatch) -> None: - service = AgentObservabilityService(session=None) - app = _app() - rows = [ - {"id": "old", "source": {"id": "webapp:app-1"}, "created_at": 10, "updated_at": 100}, - {"id": "new", "source": {"id": "webapp:app-1"}, "created_at": 20, "updated_at": 50}, - ] - monkeypatch.setattr(service, "_list_webapp_conversation_logs", lambda **kwargs: rows) - monkeypatch.setattr(service, "_list_workflow_conversation_logs", lambda **kwargs: []) - - payload = service.list_logs( - app=app, - agent_id="agent-1", - params=AgentLogQueryParams(sources=("webapp:app-1",), sort_by="created_at", sort_order="asc"), - ) - - assert [item["id"] for item in payload["data"]] == ["old", "new"] - - -def test_list_log_messages_merges_deduplicates_and_sorts_sources(monkeypatch: pytest.MonkeyPatch) -> None: - service = AgentObservabilityService(session=None) - webapp_message = _message(message_id="shared") - webapp_row = {"id": "shared", "created_at": 10, "updated_at": 30} - workflow_rows = [ - {"id": "shared", "created_at": 10, "updated_at": 20}, - {"id": "workflow-only", "created_at": 20, "updated_at": 10}, - ] - monkeypatch.setattr(service, "_list_webapp_messages", lambda **kwargs: [webapp_message]) - monkeypatch.setattr(service, "_list_message_feedbacks", lambda **kwargs: {}) - monkeypatch.setattr(service, "serialize_log_message", lambda message, feedbacks=(): webapp_row) - monkeypatch.setattr(service, "_list_workflow_messages", lambda **kwargs: workflow_rows) - - payload = service.list_log_messages( - app=_app(app_id="agent-app"), - agent_id="agent-1", - conversation_id="execution-1", - params=AgentLogQueryParams( - sources=("webapp:agent-app", "workflow:workflow-app"), - sort_by="created_at", - sort_order="asc", - ), - ) - - assert payload == { - "data": [workflow_rows[0], workflow_rows[1]], - "page": 1, - "limit": 20, - "total": 2, - "has_more": False, - } - - -def test_list_webapp_conversation_logs_includes_feedback_rates(sqlite_session: Session) -> None: - timestamp = datetime(2026, 7, 23, 7, 0, 19) - app = _app(name="Agent WebApp") - conversation = _conversation() - conversation.name = "Feedback conversation" - first_message = _message(created_at=timestamp) - second_message = _message(message_id="message-2", created_at=timestamp) - sqlite_session.add_all( - [ - app, - conversation, - first_message, - second_message, - _feedback(message_id=first_message.id), - _feedback(message_id=second_message.id, rating=FeedbackRating.DISLIKE), - _feedback(message_id=first_message.id, source=FeedbackFromSource.ADMIN), - ] - ) - sqlite_session.commit() - service = AgentObservabilityService(sqlite_session) - - rows = service._list_webapp_conversation_logs( - app=app, - params=AgentLogQueryParams(), - source_filter=AgentObservabilityService.resolve_source_filter("webapp"), - ) - - assert rows[0]["user_rate"] == 0.5 - assert rows[0]["operation_rate"] == 1.0 - - -def test_list_workflow_logs_uses_node_executions_without_messages(sqlite_session: Session) -> None: - workflow_app = _app(app_id="workflow-app-1", name="Marketing Department", mode=AppMode.WORKFLOW) - sqlite_session.add_all([workflow_app, _workflow_run(), _node_execution(), _workflow_binding()]) - sqlite_session.commit() - service = AgentObservabilityService(sqlite_session) - - rows = service._list_workflow_conversation_logs( - app=_app(app_id="agent-app"), - agent_id="agent-1", - params=AgentLogQueryParams(), - source_filter=AgentObservabilityService.resolve_source_filter("workflow:workflow-app-1"), - ) - - assert rows[0]["id"] == "node-execution-1" - assert rows[0]["source"]["app_name"] == "Marketing Department" - - -@pytest.mark.parametrize("source_identity", ["persisted-binding", "source-version", "unique-legacy-binding"]) -def test_nested_agent_logs_and_messages_keep_source_identity(sqlite_session: Session, source_identity: str) -> None: - source_app = _app(app_id="workflow-app-1", name="Source workflow", mode=AppMode.WORKFLOW) - source = Workflow( - id="workflow-1", - tenant_id="tenant-1", - app_id=source_app.id, - type=WorkflowType.WORKFLOW, - version="v1", - graph="{}", - _features="{}", - created_by="account-1", - ) - run = _workflow_run() - run.app_id, run.workflow_id, run.version = "caller-app", "caller-workflow", "caller-version" - execution = _node_execution() - execution.triggered_from = WorkflowNodeExecutionTriggeredFrom.WORKFLOW_TOOL - binding = _workflow_binding() - if source_identity == "persisted-binding": - execution.process_data = json.dumps({"workflow_agent_binding_id": binding.id}) - unused_binding = _workflow_binding(binding_id="unused-binding") - unused_binding.workflow_version = "v2" - sqlite_session.add_all([source_app, run, execution, binding]) - if source_identity != "unique-legacy-binding": - sqlite_session.add_all([source, unused_binding]) - sqlite_session.commit() - service = AgentObservabilityService(sqlite_session) - params = AgentLogQueryParams(sources=("workflow:workflow-app-1",)) - - logs = service.list_logs(app=_app(app_id="agent-app"), agent_id="agent-1", params=params) - messages = service.list_log_messages( - app=_app(app_id="agent-app"), agent_id="agent-1", conversation_id=execution.id, params=params - ) - - assert logs["total"] == messages["total"] == 1 - assert logs["data"][0]["id"] == messages["data"][0]["id"] == execution.id - assert logs["data"][0]["source"]["id"] == "workflow:workflow-app-1:workflow-1:v1:node-1" - assert messages["data"][0]["total_tokens"] == 454_064 - for source_filter in ("workflow:caller-app", "workflow:workflow-app-1:workflow-1:v2:node-1"): - params = AgentLogQueryParams(sources=(source_filter,)) - assert service.list_logs(app=_app(app_id="agent-app"), agent_id="agent-1", params=params)["total"] == 0 - assert ( - service.list_log_messages( - app=_app(app_id="agent-app"), agent_id="agent-1", conversation_id=execution.id, params=params - )["total"] - == 0 - ) - - if source_identity == "unique-legacy-binding": - # Without a captured binding or source Workflow, another Agent/version - # makes attribution ambiguous, even when filtering for only agent-1. - unused_binding.agent_id = "another-agent" - sqlite_session.add(unused_binding) - sqlite_session.commit() - params = AgentLogQueryParams(sources=("workflow",)) - assert service.list_logs(app=_app(app_id="agent-app"), agent_id="agent-1", params=params)["total"] == 0 - assert ( - service.list_log_messages( - app=_app(app_id="agent-app"), agent_id="agent-1", conversation_id=execution.id, params=params - )["total"] - == 0 - ) - - -@pytest.mark.parametrize("mismatch", ["run-tenant", "source-tenant", "binding-id", "unrelated-run"]) -def test_workflow_agent_observability_rejects_mismatched_ownership(sqlite_session: Session, mismatch: str) -> None: - source_app = _app(app_id="workflow-app-1", mode=AppMode.WORKFLOW) - run, execution, binding = _workflow_run(), _node_execution(), _workflow_binding() - execution.process_data = json.dumps({"workflow_agent_binding_id": binding.id}) - if mismatch == "run-tenant": - run.tenant_id = "other-tenant" - elif mismatch == "source-tenant": - source_app.tenant_id = "other-tenant" - elif mismatch == "binding-id": - execution.process_data = json.dumps({"workflow_agent_binding_id": "missing-binding"}) - else: - run.app_id, run.workflow_id = "unrelated-app", "unrelated-workflow" - sqlite_session.add_all([source_app, run, execution, binding]) - sqlite_session.commit() - service = AgentObservabilityService(sqlite_session) - params = AgentLogQueryParams(sources=("workflow",)) - - assert service.list_logs(app=_app(app_id="agent-app"), agent_id="agent-1", params=params)["total"] == 0 - assert ( - service.list_log_messages( - app=_app(app_id="agent-app"), agent_id="agent-1", conversation_id=execution.id, params=params - )["total"] - == 0 - ) - - -def test_list_workflow_messages_uses_node_execution_identity( - monkeypatch: pytest.MonkeyPatch, sqlite_session: Session -) -> None: - node_execution = _node_execution() - sqlite_session.add_all( - [ - _app(app_id="workflow-app-1", name="Workflow App", mode=AppMode.WORKFLOW), - _workflow_run(), - node_execution, - _workflow_binding(), - ] - ) - sqlite_session.commit() - service = AgentObservabilityService(sqlite_session) - serialized = {"id": "node-execution-1", "conversation_id": "node-execution-1"} - monkeypatch.setattr(service, "serialize_workflow_node_message", lambda execution: serialized) - - rows = service._list_workflow_messages( - app=_app(app_id="agent-app"), - agent_id="agent-1", - conversation_id="node-execution-1", - params=AgentLogQueryParams(), - source_filter=AgentObservabilityService.resolve_source_filter("workflow:workflow-app-1"), - ) - - assert rows == [serialized] - - -def test_apply_workflow_node_filters_supports_time_keyword_and_status() -> None: - stmt = select(WorkflowNodeExecutionModel) - params = AgentLogQueryParams( - start=datetime(2026, 7, 1, tzinfo=UTC), - end=datetime(2026, 8, 1, tzinfo=UTC), - keyword="meeting_100%", - statuses=("success",), - ) - - result = AgentObservabilityService._apply_workflow_node_filters( - stmt, - params=params, - workflow_app=observability_service_module.App, - ) - - assert isinstance(result, Select) - assert len(result._where_criteria) == 4 - - -def test_apply_workflow_node_status_filter_supports_all_status_groups() -> None: - stmt = select(WorkflowNodeExecutionModel) - - result = AgentObservabilityService._apply_workflow_node_status_filter(stmt, ("normal", "error", "paused")) - - assert isinstance(result, Select) - assert len(result._where_criteria) == 1 - empty_stmt = select(WorkflowNodeExecutionModel) - assert AgentObservabilityService._apply_workflow_node_status_filter(empty_stmt, ()) is empty_stmt - assert empty_stmt._where_criteria == () - with pytest.raises(ValueError, match="Unsupported status"): - AgentObservabilityService._apply_workflow_node_status_filter(select(WorkflowNodeExecutionModel), ("unknown",)) - - -def test_source_serializers_return_structured_frontend_shape() -> None: - app = _app() - - webapp_source = AgentObservabilityService._serialize_webapp_source(app) - workflow_app_source = AgentObservabilityService._serialize_workflow_app_source(app=app) - workflow_source = AgentObservabilityService._serialize_workflow_source( - app=app, - workflow_id="workflow-1", - workflow_version="v1", - node_id="node-1", - ) - - assert webapp_source == { - "id": "webapp:app-1", - "type": "webapp", - "app_id": "app-1", - "app_name": "Iris", - "app_icon_type": "emoji", - "app_icon": "robot", - "app_icon_background": "#fff", - "workflow_id": None, - "workflow_version": None, - "node_id": None, - } - assert workflow_app_source == { - "id": "workflow:app-1", - "type": "workflow", - "app_id": "app-1", - "app_name": "Iris", - "app_icon_type": "emoji", - "app_icon": "robot", - "app_icon_background": "#fff", - "workflow_id": None, - "workflow_version": None, - "node_id": None, - } - assert workflow_source["id"] == "workflow:app-1:workflow-1:v1:node-1" - assert workflow_source["type"] == "workflow" - assert workflow_source["workflow_id"] == "workflow-1" - - -def test_list_workflow_sources_deduplicates_versions_and_nodes_by_app(sqlite_session: Session) -> None: - app_a = _app(app_id="app-a", name="Alpha", mode=AppMode.WORKFLOW) - app_b = _app(app_id="app-b", name="Beta", mode=AppMode.WORKFLOW) - sqlite_session.add_all( - [ - app_a, - app_b, - _workflow_binding(app_id="app-a", node_id="node-a-1"), - _workflow_binding(app_id="app-a", node_id="node-a-2"), - _workflow_binding(app_id="app-b", node_id="node-b-1"), - ] - ) - sqlite_session.commit() - service = AgentObservabilityService(sqlite_session) - - sources = service._list_workflow_sources( - app=_app(app_id="agent-app"), - agent_id="agent-1", - ) - - assert [source["id"] for source in sources] == ["workflow:app-a", "workflow:app-b"] - - -def test_serialize_log_message_returns_frontend_log_shape() -> None: - created_at = datetime(2026, 6, 17, 1, 2, 3, tzinfo=UTC) - updated_at = datetime(2026, 6, 17, 1, 3, 3, tzinfo=UTC) - message = _message(created_at=created_at) - message.updated_at = updated_at - conversation = _conversation() - feedbacks = [ - _feedback(), - _feedback( - rating=FeedbackRating.DISLIKE, - content="Needs more detail", - source=FeedbackFromSource.ADMIN, - ), - ] - - payload = AgentObservabilityService.serialize_log_message(message, conversation, feedbacks) - - assert payload == { - "id": "message-1", - "message_id": "message-1", - "conversation_id": "conversation-1", - "conversation_name": "Debug conversation", - "query": "hello", - "answer": "hi", - "status": "success", - "error": None, - "source": "explore", - "from_source": "console", - "from_end_user_id": None, - "from_account_id": "account-1", - "feedback_enabled": True, - "feedbacks": [ - {"rating": "like", "content": "Useful", "from_source": "user"}, - {"rating": "dislike", "content": "Needs more detail", "from_source": "admin"}, - ], - "message_tokens": 3, - "answer_tokens": 4, - "total_tokens": 7, - "total_price": "0.0001", - "currency": "USD", - "latency": 1.25, - "created_at": int(created_at.timestamp()), - "updated_at": int(updated_at.timestamp()), - } - - -def test_serialize_workflow_node_message_returns_frontend_log_shape() -> None: - created_at = datetime(2026, 7, 23, 7, 0, 19, tzinfo=UTC) - finished_at = datetime(2026, 7, 23, 7, 0, 28, tzinfo=UTC) - node_execution = _node_execution() - node_execution.inputs = ( - '{"agent_backend_request":{"composition":{"layers":[' - '{"name":"workflow_node_job_prompt","config":{"user":"Summarize the meeting"}},' - '{"name":"workflow_user_prompt","config":{"user":"Focus on action items"}}]}}}' - ) - node_execution.outputs = '{"output":"Alice owns the follow-up."}' - node_execution.execution_metadata = ( - '{"agent_log":{"agent_backend":{"usage":{"prompt_tokens":10,"completion_tokens":5,' - '"total_tokens":15,"total_price":"0.0015","currency":"USD","latency":1.25}}}}' - ) - node_execution.elapsed_time = 1.5 - node_execution.created_at = created_at - node_execution.finished_at = finished_at - - payload = AgentObservabilityService.serialize_workflow_node_message(node_execution) - - assert payload == { - "id": "node-execution-1", - "message_id": "node-execution-1", - "conversation_id": "node-execution-1", - "query": "Summarize the meeting\n\nFocus on action items", - "answer": "Alice owns the follow-up.", - "status": "success", - "error": None, - "from_end_user_id": "end-user-1", - "from_account_id": None, - "feedback_enabled": False, - "feedbacks": [], - "message_tokens": 10, - "answer_tokens": 5, - "total_tokens": 15, - "total_price": "0.0015", - "currency": "USD", - "latency": 1.25, - "created_at": int(created_at.timestamp()), - "updated_at": int(finished_at.timestamp()), - } - - -def test_serialize_workflow_node_message_handles_sparse_runtime_data() -> None: - created_at = datetime(2026, 7, 23, 7, 0, 19, tzinfo=UTC) - node_execution = _node_execution(execution_id="node-execution-2", status=WorkflowNodeExecutionStatus.PAUSED) - node_execution.title = "Fallback prompt" - node_execution.inputs = json.dumps( - { - "agent_backend_request": { - "composition": { - "layers": [ - None, - {"name": "unrelated", "config": {"user": "ignored"}}, - {"name": "workflow_user_prompt", "config": {"user": " "}}, - ] - } - } - } - ) - node_execution.outputs = json.dumps({"output": {"structured": True}}) - node_execution.execution_metadata = json.dumps( - { - "agent_log": { - "agent_backend": { - "usage": { - "prompt_tokens": "2", - "completion_tokens": 3, - } - } - } - } - ) - node_execution.elapsed_time = 2 - node_execution.created_by_role = CreatorUserRole.ACCOUNT - node_execution.created_by = "account-1" - node_execution.created_at = created_at - node_execution.finished_at = None - - payload = AgentObservabilityService.serialize_workflow_node_message(node_execution) - - assert payload["query"] == "Fallback prompt" - assert payload["answer"] == '{"output": {"structured": true}}' - assert payload["status"] == "paused" - assert payload["from_end_user_id"] is None - assert payload["from_account_id"] == "account-1" - assert payload["total_tokens"] == 5 - assert payload["total_price"] == "0" - assert payload["currency"] == "" - assert payload["latency"] == 2.0 - assert payload["updated_at"] == int(created_at.timestamp()) - - -def test_positive_feedback_rate_uses_rated_messages_as_denominator() -> None: - assert AgentObservabilityService._positive_feedback_rate(like_count=2, total_count=4) == 0.5 - assert AgentObservabilityService._positive_feedback_rate(like_count=0, total_count=1) == 0 - assert AgentObservabilityService._positive_feedback_rate(like_count=None, total_count=0) is None - - -def test_list_message_feedbacks_groups_feedbacks_by_message(sqlite_session: Session) -> None: - first_message = _message() - second_message = _message(message_id="message-2") - feedbacks = [ - _feedback(), - _feedback(rating=FeedbackRating.DISLIKE), - _feedback(message_id="message-2"), - ] - sqlite_session.add_all([_app(), _conversation(), first_message, second_message, *feedbacks]) - sqlite_session.commit() - service = AgentObservabilityService(sqlite_session) - - grouped_feedbacks = service._list_message_feedbacks( - app=_app(), - messages=[first_message, second_message], - ) - - assert {feedback.id for feedback in grouped_feedbacks["message-1"]} == { - feedbacks[0].id, - feedbacks[1].id, - } - assert grouped_feedbacks["message-2"] == [feedbacks[2]] - assert service._list_message_feedbacks(app=_app(), messages=[]) == {} - - -def test_list_conversation_feedback_rates_maps_user_and_admin_sources(sqlite_session: Session) -> None: - messages = [_message(message_id=f"message-{index}") for index in range(1, 5)] - feedbacks = [ - _feedback(message_id="message-1"), - _feedback(message_id="message-2"), - _feedback(message_id="message-3", rating=FeedbackRating.DISLIKE), - _feedback(message_id="message-4", rating=FeedbackRating.DISLIKE), - _feedback(message_id="message-1", source=FeedbackFromSource.ADMIN), - ] - sqlite_session.add_all([_app(), _conversation(), *messages, *feedbacks]) - sqlite_session.commit() - service = AgentObservabilityService(sqlite_session) - - rates = service._list_conversation_feedback_rates( - app=_app(), - conversation_ids=["conversation-1"], - ) - - assert rates == {"conversation-1": {"user_rate": 0.5, "operation_rate": 1.0}} - assert ( - service._list_conversation_feedback_rates( - app=_app(), - conversation_ids=[], - ) - == {} - ) - - -def test_workflow_node_serialization_helpers_handle_invalid_values() -> None: - assert AgentObservabilityService._json_mapping(None) == {} - assert AgentObservabilityService._json_mapping("not-json") == {} - assert AgentObservabilityService._json_mapping("[]") == {} - assert AgentObservabilityService._mapping_value({"value": []}, "value") == {} - assert AgentObservabilityService._int_value(None) == 0 - assert AgentObservabilityService._int_value("not-a-number") == 0 - assert ( - AgentObservabilityService._workflow_node_query( - {"agent_backend_request": {"composition": {"layers": "invalid"}}}, fallback="fallback" - ) - == "fallback" - ) - assert AgentObservabilityService._workflow_node_answer({"output": 1, "text": "fallback text"}) == "fallback text" - assert AgentObservabilityService._workflow_node_answer({}) == "" - - -def test_serialize_workflow_execution_log_uses_node_execution_identity() -> None: - created_at = datetime(2026, 7, 23, 7, 0, 19, tzinfo=UTC) - node_execution = _node_execution(status=WorkflowNodeExecutionStatus.FAILED) - node_execution.created_by_role = CreatorUserRole.ACCOUNT - node_execution.created_by = "account-1" - node_execution.created_at = created_at - - payload = AgentObservabilityService._serialize_workflow_execution_log( - node_execution_id=node_execution.id, - title=node_execution.title, - status=node_execution.status, - created_by_role=node_execution.created_by_role, - created_by=node_execution.created_by, - created_at=node_execution.created_at, - finished_at=node_execution.finished_at, - source={"id": "workflow:app-1:workflow-1:v1:node-1"}, - ) - - assert payload["id"] == "node-execution-1" - assert payload["conversation_id"] == "node-execution-1" - assert payload["message_count"] == 1 - assert payload["end_user_id"] is None - assert payload["status"] == "failed" - assert payload["unread"] is False - - -def test_build_charts_and_summary_match_monitoring_metrics() -> None: - rows = [ - { - "date": "2026-06-16", - "message_count": 2, - "conversation_count": 1, - "end_user_count": 1, - "token_count": 30, - "total_price": Decimal("0.003"), - "avg_latency": 1.5, - "latency_sum": 3, - "answer_tokens": 12, - "like_count": 1, - }, - { - "date": "2026-06-17", - "message_count": 1, - "conversation_count": 1, - "end_user_count": 1, - "token_count": 20, - "total_price": Decimal("0.002"), - "avg_latency": 2, - "latency_sum": 2, - "answer_tokens": 8, - "like_count": 1, - }, - ] - - charts = AgentObservabilityService._build_charts(rows) - summary = AgentObservabilityService._build_summary(rows) - - assert charts["token_usage"] == [ - {"date": "2026-06-16", "token_count": 30, "total_price": "0.003", "currency": "USD"}, - {"date": "2026-06-17", "token_count": 20, "total_price": "0.002", "currency": "USD"}, - ] - assert charts["average_response_time"] == [ - {"date": "2026-06-16", "latency": 1500.0}, - {"date": "2026-06-17", "latency": 2000.0}, - ] - assert summary == { - "total_messages": 3, - "total_conversations": 2, - "total_end_users": 2, - "total_tokens": 50, - "total_price": "0.005", - "currency": "USD", - "average_session_interactions": 1.5, - "average_response_time": 1666.6667, - "tokens_per_second": 4.0, - "user_satisfaction_rate": 66.67, - } diff --git a/api/tests/unit_tests/services/agent/test_composer_mention_validation.py b/api/tests/unit_tests/services/agent/test_composer_validator.py similarity index 100% rename from api/tests/unit_tests/services/agent/test_composer_mention_validation.py rename to api/tests/unit_tests/services/agent/test_composer_validator.py diff --git a/api/tests/unit_tests/services/agent/test_agent_dsl_service.py b/api/tests/unit_tests/services/agent/test_dsl_service.py similarity index 99% rename from api/tests/unit_tests/services/agent/test_agent_dsl_service.py rename to api/tests/unit_tests/services/agent/test_dsl_service.py index 47b52d10af6..414b65e3d21 100644 --- a/api/tests/unit_tests/services/agent/test_agent_dsl_service.py +++ b/api/tests/unit_tests/services/agent/test_dsl_service.py @@ -38,7 +38,7 @@ from services.agent.dsl_entities import ( ) from services.agent.dsl_service import AgentDslService, AgentPackageImportResult, is_agent_v2_graph from services.entities.dsl_entities import DslImportWarning -from tests.unit_tests.model_factories import make_upload_file +from tests.helpers.unit.model_factories import make_upload_file def _agent( diff --git a/api/tests/unit_tests/services/agent/test_home_snapshot_service.py b/api/tests/unit_tests/services/agent/test_home_snapshot_service.py index 97a615aee93..ebdc8f686d2 100644 --- a/api/tests/unit_tests/services/agent/test_home_snapshot_service.py +++ b/api/tests/unit_tests/services/agent/test_home_snapshot_service.py @@ -30,7 +30,7 @@ from services.agent.errors import ( ) from services.agent.home_snapshot_service import AgentHomeSnapshotService, validate_home_snapshot_binding from services.agent.workspace_service import AgentWorkspaceBindingGenerationMismatchError, AgentWorkspaceService -from tests.unit_tests.config_override import apply_config_overrides +from tests.helpers.unit.config_override import apply_config_overrides def _build_draft(*, home_snapshot_id: str | None = "home-old") -> AgentConfigDraft: diff --git a/api/tests/unit_tests/services/agent/test_knowledge_gateway.py b/api/tests/unit_tests/services/agent/test_knowledge_gateway.py index 14429b54c01..9de35c7c0e0 100644 --- a/api/tests/unit_tests/services/agent/test_knowledge_gateway.py +++ b/api/tests/unit_tests/services/agent/test_knowledge_gateway.py @@ -34,8 +34,8 @@ from services.knowledge_fs.capability_broker import KnowledgeFSIssuedProductCapa from services.knowledge_fs.errors import KnowledgeFSWorkflowQueryImageReference from services.knowledge_fs.query_images import KnowledgeFSQueryImageMetadata from services.knowledge_fs.runtime import KnowledgeFSRuntime -from tests.unit_tests.config_override import apply_config_overrides -from tests.unit_tests.services.knowledge_fs_services_b_fakes import RecordingCallable +from tests.helpers.unit.config_override import apply_config_overrides +from tests.helpers.unit.services.knowledge_fs_services_b_fakes import RecordingCallable SPACE = "00000000-0000-4000-8000-000000000001" NODE = "00000000-0000-4000-8000-000000000002" diff --git a/api/tests/unit_tests/services/agent/test_knowledge_runtime_config.py b/api/tests/unit_tests/services/agent/test_knowledge_runtime_config.py index bf322b8f702..ea46f2e422a 100644 --- a/api/tests/unit_tests/services/agent/test_knowledge_runtime_config.py +++ b/api/tests/unit_tests/services/agent/test_knowledge_runtime_config.py @@ -5,8 +5,8 @@ from core.app.entities.app_invoke_entities import DifyRunContext, InvokeFrom, Us from graphon.model_runtime.entities.model_entities import ModelFeature from models.agent_config_entities import AgentSoulConfig from services.agent import knowledge_runtime_config as mapping -from tests.unit_tests.config_override import apply_config_overrides -from tests.unit_tests.knowledge_fs_core_fakes import RecordingModelFactory, llm_schema +from tests.helpers.unit.config_override import apply_config_overrides +from tests.helpers.unit.knowledge_fs_core_fakes import RecordingModelFactory, llm_schema def soul() -> AgentSoulConfig: diff --git a/api/tests/unit_tests/services/agent/test_observability_service.py b/api/tests/unit_tests/services/agent/test_observability_service.py new file mode 100644 index 00000000000..17acde9fd0c --- /dev/null +++ b/api/tests/unit_tests/services/agent/test_observability_service.py @@ -0,0 +1,1286 @@ +import json +from datetime import UTC, datetime +from decimal import Decimal +from sqlite3 import Connection + +import pytest +from sqlalchemy import Select, select +from sqlalchemy.orm import Session + +from core.app.entities.app_invoke_entities import InvokeFrom as _agent_observability_service_InvokeFrom +from core.workflow.node_execution_process_data import WORKFLOW_AGENT_BINDING_ID_KEY +from graphon.enums import WorkflowNodeExecutionStatus +from libs.datetime_utils import naive_utc_now +from models.agent import WorkflowAgentBindingType, WorkflowAgentNodeBinding +from models.agent_config_entities import WorkflowNodeJobConfig +from models.enums import ( + ConversationFromSource, + CreatorUserRole, + FeedbackFromSource, + FeedbackRating, + MessageStatus, +) +from models.enums import InvokeFrom as _workflow_tool_agent_statistics_InvokeFrom +from models.model import App, AppMode, Conversation, Message, MessageFeedback +from models.workflow import ( + Workflow, + WorkflowExecutionStatus, + WorkflowNodeExecutionModel, + WorkflowNodeExecutionTriggeredFrom, + WorkflowRun, + WorkflowRunTriggeredFrom, + WorkflowType, +) +from services.agent import observability_service as observability_service_module +from services.agent.observability_service import ( + AgentLogQueryParams, + AgentObservabilityService, + AgentStatisticsQueryParams, +) +from tests.helpers.unit.config_override import apply_config_overrides +from tests.helpers.unit.model_factories import make_app, make_conversation, make_message + + +def _agent_observability_service_app( + *, app_id: str = "app-1", name: str = "Iris", mode: AppMode = AppMode.AGENT_CHAT +) -> App: + return make_app(app_id=app_id, name=name, mode=mode, icon_background="#fff", enable_site=False, enable_api=False) + + +def _agent_observability_service_conversation( + *, conversation_id: str = "conversation-1", app_id: str = "app-1" +) -> Conversation: + return make_conversation( + conversation_id=conversation_id, + app_id=app_id, + mode=AppMode.AGENT_CHAT, + name="Debug conversation", + inputs={}, + from_source=ConversationFromSource.CONSOLE, + from_end_user_id="end-user-1", + ) + + +def _agent_observability_service_message( + *, + message_id: str = "message-1", + conversation_id: str = "conversation-1", + app_id: str = "app-1", + created_at: datetime | None = None, +) -> Message: + timestamp = created_at or naive_utc_now() + message = make_message( + message_id=message_id, + app_id=app_id, + conversation_id=conversation_id, + inputs={}, + query="hello", + message={}, + answer="hi", + status=MessageStatus.NORMAL, + message_unit_price=Decimal(0), + answer_unit_price=Decimal(0), + total_price=Decimal("0.0001"), + currency="USD", + from_source=ConversationFromSource.CONSOLE, + from_account_id="account-1", + invoke_from=_agent_observability_service_InvokeFrom.EXPLORE, + message_tokens=3, + answer_tokens=4, + provider_response_latency=1.25, + ) + message.created_at = timestamp + message.updated_at = timestamp + return message + + +def _feedback( + *, + message_id: str = "message-1", + conversation_id: str = "conversation-1", + source: FeedbackFromSource = FeedbackFromSource.USER, + rating: FeedbackRating = FeedbackRating.LIKE, + content: str | None = "Useful", +) -> MessageFeedback: + return MessageFeedback( + app_id="app-1", + conversation_id=conversation_id, + message_id=message_id, + rating=rating, + from_source=source, + content=content, + ) + + +def _agent_observability_service_workflow_run(*, workflow_type: WorkflowType = WorkflowType.WORKFLOW) -> WorkflowRun: + created_at = datetime(2026, 7, 21, 7, 0, 19) + return WorkflowRun( + id="workflow-run-1", + tenant_id="tenant-1", + app_id="workflow-app-1", + workflow_id="workflow-1", + type=workflow_type, + triggered_from=WorkflowRunTriggeredFrom.APP_RUN, + version="v1", + graph="{}", + inputs="{}", + status=WorkflowExecutionStatus.SUCCEEDED, + outputs="{}", + error=None, + elapsed_time=59.93, + total_tokens=454_064, + total_steps=1, + created_by_role=CreatorUserRole.END_USER, + created_by="end-user-1", + created_at=created_at, + finished_at=created_at, + ) + + +def _agent_observability_service_node_execution( + *, + execution_id: str = "node-execution-1", + status: WorkflowNodeExecutionStatus = WorkflowNodeExecutionStatus.SUCCEEDED, +) -> WorkflowNodeExecutionModel: + created_at = datetime(2026, 7, 23, 7, 0, 19, tzinfo=UTC) + return WorkflowNodeExecutionModel( + id=execution_id, + tenant_id="tenant-1", + app_id="workflow-app-1", + workflow_id="workflow-1", + triggered_from=WorkflowNodeExecutionTriggeredFrom.WORKFLOW_RUN, + workflow_run_id="workflow-run-1", + index=1, + predecessor_node_id=None, + node_execution_id=execution_id, + node_id="node-1", + node_type="agent", + title="Agent", + inputs="{}", + process_data="{}", + outputs="{}", + status=status, + error=None, + elapsed_time=59.93, + execution_metadata=json.dumps( + { + "agent_log": { + "agent_backend": { + "usage": { + "prompt_tokens": 451_938, + "completion_tokens": 2_126, + "total_tokens": 454_064, + "total_price": "2.323470", + "currency": "USD", + "latency": 59.93, + } + } + } + } + ), + created_at=created_at, + created_by_role=CreatorUserRole.END_USER, + created_by="end-user-1", + finished_at=None, + ) + + +def _agent_observability_service_workflow_binding( + *, app_id: str = "workflow-app-1", binding_id: str | None = None, node_id: str = "node-1" +) -> WorkflowAgentNodeBinding: + return WorkflowAgentNodeBinding( + id=binding_id or f"binding-{app_id}-{node_id}", + tenant_id="tenant-1", + app_id=app_id, + workflow_id="workflow-1", + workflow_version="v1", + node_id=node_id, + binding_type=WorkflowAgentBindingType.ROSTER_AGENT, + agent_id="agent-1", + current_snapshot_id="snapshot-1", + node_job_config=WorkflowNodeJobConfig(), + created_by="account-1", + ) + + +class TestAgentObservabilityServiceCases: + @staticmethod + def test_resolve_source_accepts_frontend_aliases() -> None: + assert AgentObservabilityService.resolve_source(None) is None + assert AgentObservabilityService.resolve_source("all") is None + assert AgentObservabilityService.resolve_source("console") == _agent_observability_service_InvokeFrom.EXPLORE + assert AgentObservabilityService.resolve_source("api") == _agent_observability_service_InvokeFrom.SERVICE_API + assert AgentObservabilityService.resolve_source("web_app") == _agent_observability_service_InvokeFrom.WEB_APP + + with pytest.raises(ValueError, match="Unsupported source"): + AgentObservabilityService.resolve_source("unknown") + + @staticmethod + def test_resolve_source_filter_accepts_structured_sources() -> None: + assert AgentObservabilityService.resolve_source_filter(None).kind == "all" + assert AgentObservabilityService.resolve_source_filter("webapp").kind == "webapp" + assert AgentObservabilityService.resolve_source_filter("webapp:app-1").app_id == "app-1" + + workflow_app_filter = AgentObservabilityService.resolve_source_filter("workflow:app-2") + assert workflow_app_filter.kind == "workflow" + assert workflow_app_filter.app_id == "app-2" + assert workflow_app_filter.workflow_id is None + + workflow_filter = AgentObservabilityService.resolve_source_filter("workflow:app-2:workflow-1:v1:node-1") + assert workflow_filter.kind == "workflow" + assert workflow_filter.app_id == "app-2" + assert workflow_filter.workflow_id == "workflow-1" + assert workflow_filter.workflow_version == "v1" + assert workflow_filter.node_id == "node-1" + + timestamp_version_filter = AgentObservabilityService.resolve_source_filter( + "workflow:app-2:workflow-1:2026-07-06 02:17:12.910515:node-1" + ) + assert timestamp_version_filter.workflow_version == "2026-07-06 02:17:12.910515" + assert timestamp_version_filter.node_id == "node-1" + + legacy_filter = AgentObservabilityService.resolve_source_filter("console") + assert legacy_filter.kind == "webapp" + assert legacy_filter.invoke_from == _agent_observability_service_InvokeFrom.EXPLORE + + with pytest.raises(ValueError, match="Unsupported source"): + AgentObservabilityService.resolve_source_filter("workflow:") + with pytest.raises(ValueError, match="Unsupported source"): + AgentObservabilityService.resolve_source_filter("workflow:app-2:incomplete") + + @staticmethod + def test_resolve_source_filters_accepts_multiple_structured_sources() -> None: + filters = AgentObservabilityService.resolve_source_filters( + ("webapp:app-1", "workflow:app-2:workflow-1:v1:node-1") + ) + + assert [source_filter.kind for source_filter in filters] == ["webapp", "workflow"] + assert filters[0].app_id == "app-1" + assert filters[1].node_id == "node-1" + assert AgentObservabilityService.resolve_source_filters(())[0].kind == "all" + assert AgentObservabilityService.resolve_source_filters(("all", "webapp:app-1"))[0].kind == "all" + + @staticmethod + def test_statistics_all_source_includes_debugger_messages() -> None: + source_filter = AgentObservabilityService.resolve_source_filter("all") + + scope_sql = AgentObservabilityService._statistics_webapp_message_scope_sql(source_filter) + + assert "m.app_id = :app_id" in scope_sql + assert "m.invoke_from != :debugger" not in scope_sql + + @staticmethod + def test_statistics_explicit_source_filters_invoke_from() -> None: + source_filter = AgentObservabilityService.resolve_source_filter("debugger") + + scope_sql = AgentObservabilityService._statistics_webapp_message_scope_sql(source_filter) + + assert "m.invoke_from = :source" in scope_sql + + @staticmethod + def test_statistics_workflow_app_source_covers_all_versions_and_nodes() -> None: + source_filter = AgentObservabilityService.resolve_source_filter("workflow:app-2") + + scope_sql = AgentObservabilityService._statistics_workflow_binding_filters_sql(source_filter) + + assert "wanb.app_id = :source_app_id" in scope_sql + assert "wanb.workflow_id = :workflow_id" not in scope_sql + assert "wanb.workflow_version = :workflow_version" not in scope_sql + assert "wanb.node_id = :node_id" not in scope_sql + + @staticmethod + def test_workflow_metadata_numeric_sql_supports_postgresql_and_mysql(monkeypatch: pytest.MonkeyPatch) -> None: + apply_config_overrides(monkeypatch, DB_TYPE="postgresql") + + postgres_sql = AgentObservabilityService._workflow_execution_metadata_numeric_sql( + ("agent_log", "agent_backend", "usage", "total_tokens"), "BIGINT" + ) + + assert "CAST(wne.execution_metadata AS JSONB)" in postgres_sql + assert "#>> '{agent_log,agent_backend,usage,total_tokens}'" in postgres_sql + + apply_config_overrides(monkeypatch, DB_TYPE="mysql") + + mysql_sql = AgentObservabilityService._workflow_execution_metadata_numeric_sql(("total_tokens",), "BIGINT") + + assert "JSON_EXTRACT(wne.execution_metadata, '$.total_tokens')" in mysql_sql + assert " AS UNSIGNED)" in mysql_sql + + @staticmethod + def test_workflow_statistics_include_run_without_message( + monkeypatch: pytest.MonkeyPatch, sqlite_session: Session + ) -> None: + workflow_app = _agent_observability_service_app( + app_id="workflow-app-1", name="Workflow App", mode=AppMode.WORKFLOW + ) + sqlite_session.add_all( + [ + workflow_app, + _agent_observability_service_workflow_run(), + _agent_observability_service_node_execution(), + _agent_observability_service_workflow_binding(), + ] + ) + sqlite_session.commit() + + apply_config_overrides(monkeypatch, DB_TYPE="mysql") + monkeypatch.setattr(observability_service_module, "convert_datetime_to_date", lambda field: f"DATE({field})") + monkeypatch.setattr( + AgentObservabilityService, + "_workflow_execution_metadata_numeric_sql", + staticmethod( + lambda path, numeric_type: ( + f"CAST(json_extract(wne.execution_metadata, '$.{'.'.join(path)}') AS {numeric_type})" + ) + ), + ) + service = AgentObservabilityService(sqlite_session) + + payload = service.get_statistics_summary( + app=_agent_observability_service_app(app_id="agent-app"), + agent_id="agent-1", + params=observability_service_module.AgentStatisticsQueryParams(source="workflow:workflow-app-1"), + ) + + assert payload["summary"]["total_messages"] == 1 + assert payload["summary"]["total_conversations"] == 1 + assert payload["summary"]["total_end_users"] == 1 + assert payload["summary"]["total_tokens"] == 454_064 + assert Decimal(payload["summary"]["total_price"]) == Decimal("2.323470") + + @staticmethod + def test_merge_daily_statistics_combines_webapp_and_workflow_rows() -> None: + rows = [ + { + "date": "2026-07-21", + "message_count": 2, + "conversation_count": 1, + "end_user_count": 1, + "token_count": 30, + "total_price": Decimal("0.003"), + "avg_latency": 1.5, + "latency_sum": 3, + "answer_tokens": 12, + "like_count": 1, + }, + { + "date": "2026-07-21", + "message_count": 1, + "conversation_count": 1, + "end_user_count": 1, + "token_count": 20, + "total_price": Decimal("0.002"), + "avg_latency": 2, + "latency_sum": 2, + "answer_tokens": 8, + "like_count": 0, + }, + ] + + merged = AgentObservabilityService._merge_daily_statistics(rows) + + assert merged == [ + { + "date": "2026-07-21", + "message_count": 3, + "conversation_count": 2, + "end_user_count": 2, + "token_count": 50, + "total_price": Decimal("0.005"), + "avg_latency": pytest.approx(5 / 3), + "latency_sum": 5.0, + "answer_tokens": 20, + "like_count": 1, + } + ] + + @staticmethod + def test_apply_status_filter_accepts_multiple_statuses() -> None: + stmt = select(Message) + + result = AgentObservabilityService._apply_status_filter(stmt, ("success", "failed", "paused")) + + assert isinstance(result, Select) + assert len(result._where_criteria) == 1 + with pytest.raises(ValueError, match="Unsupported status"): + AgentObservabilityService._apply_status_filter(select(Message), ("unknown",)) + + @staticmethod + def test_list_logs_sorts_by_requested_field(monkeypatch: pytest.MonkeyPatch) -> None: + service = AgentObservabilityService(session=None) + app = _agent_observability_service_app() + rows = [ + {"id": "old", "source": {"id": "webapp:app-1"}, "created_at": 10, "updated_at": 100}, + {"id": "new", "source": {"id": "webapp:app-1"}, "created_at": 20, "updated_at": 50}, + ] + monkeypatch.setattr(service, "_list_webapp_conversation_logs", lambda **kwargs: rows) + monkeypatch.setattr(service, "_list_workflow_conversation_logs", lambda **kwargs: []) + + payload = service.list_logs( + app=app, + agent_id="agent-1", + params=AgentLogQueryParams(sources=("webapp:app-1",), sort_by="created_at", sort_order="asc"), + ) + + assert [item["id"] for item in payload["data"]] == ["old", "new"] + + @staticmethod + def test_list_log_messages_merges_deduplicates_and_sorts_sources(monkeypatch: pytest.MonkeyPatch) -> None: + service = AgentObservabilityService(session=None) + webapp_message = _agent_observability_service_message(message_id="shared") + webapp_row = {"id": "shared", "created_at": 10, "updated_at": 30} + workflow_rows = [ + {"id": "shared", "created_at": 10, "updated_at": 20}, + {"id": "workflow-only", "created_at": 20, "updated_at": 10}, + ] + monkeypatch.setattr(service, "_list_webapp_messages", lambda **kwargs: [webapp_message]) + monkeypatch.setattr(service, "_list_message_feedbacks", lambda **kwargs: {}) + monkeypatch.setattr(service, "serialize_log_message", lambda message, feedbacks=(): webapp_row) + monkeypatch.setattr(service, "_list_workflow_messages", lambda **kwargs: workflow_rows) + + payload = service.list_log_messages( + app=_agent_observability_service_app(app_id="agent-app"), + agent_id="agent-1", + conversation_id="execution-1", + params=AgentLogQueryParams( + sources=("webapp:agent-app", "workflow:workflow-app"), + sort_by="created_at", + sort_order="asc", + ), + ) + + assert payload == { + "data": [workflow_rows[0], workflow_rows[1]], + "page": 1, + "limit": 20, + "total": 2, + "has_more": False, + } + + @staticmethod + def test_list_webapp_conversation_logs_includes_feedback_rates(sqlite_session: Session) -> None: + timestamp = datetime(2026, 7, 23, 7, 0, 19) + app = _agent_observability_service_app(name="Agent WebApp") + conversation = _agent_observability_service_conversation() + conversation.name = "Feedback conversation" + first_message = _agent_observability_service_message(created_at=timestamp) + second_message = _agent_observability_service_message(message_id="message-2", created_at=timestamp) + sqlite_session.add_all( + [ + app, + conversation, + first_message, + second_message, + _feedback(message_id=first_message.id), + _feedback(message_id=second_message.id, rating=FeedbackRating.DISLIKE), + _feedback(message_id=first_message.id, source=FeedbackFromSource.ADMIN), + ] + ) + sqlite_session.commit() + service = AgentObservabilityService(sqlite_session) + + rows = service._list_webapp_conversation_logs( + app=app, + params=AgentLogQueryParams(), + source_filter=AgentObservabilityService.resolve_source_filter("webapp"), + ) + + assert rows[0]["user_rate"] == 0.5 + assert rows[0]["operation_rate"] == 1.0 + + @staticmethod + def test_list_workflow_logs_uses_node_executions_without_messages(sqlite_session: Session) -> None: + workflow_app = _agent_observability_service_app( + app_id="workflow-app-1", name="Marketing Department", mode=AppMode.WORKFLOW + ) + sqlite_session.add_all( + [ + workflow_app, + _agent_observability_service_workflow_run(), + _agent_observability_service_node_execution(), + _agent_observability_service_workflow_binding(), + ] + ) + sqlite_session.commit() + service = AgentObservabilityService(sqlite_session) + + rows = service._list_workflow_conversation_logs( + app=_agent_observability_service_app(app_id="agent-app"), + agent_id="agent-1", + params=AgentLogQueryParams(), + source_filter=AgentObservabilityService.resolve_source_filter("workflow:workflow-app-1"), + ) + + assert rows[0]["id"] == "node-execution-1" + assert rows[0]["source"]["app_name"] == "Marketing Department" + + @staticmethod + @pytest.mark.parametrize("source_identity", ["persisted-binding", "source-version", "unique-legacy-binding"]) + def test_nested_agent_logs_and_messages_keep_source_identity(sqlite_session: Session, source_identity: str) -> None: + source_app = _agent_observability_service_app( + app_id="workflow-app-1", name="Source workflow", mode=AppMode.WORKFLOW + ) + source = Workflow( + id="workflow-1", + tenant_id="tenant-1", + app_id=source_app.id, + type=WorkflowType.WORKFLOW, + version="v1", + graph="{}", + _features="{}", + created_by="account-1", + ) + run = _agent_observability_service_workflow_run() + run.app_id, run.workflow_id, run.version = "caller-app", "caller-workflow", "caller-version" + execution = _agent_observability_service_node_execution() + execution.triggered_from = WorkflowNodeExecutionTriggeredFrom.WORKFLOW_TOOL + binding = _agent_observability_service_workflow_binding() + if source_identity == "persisted-binding": + execution.process_data = json.dumps({"workflow_agent_binding_id": binding.id}) + unused_binding = _agent_observability_service_workflow_binding(binding_id="unused-binding") + unused_binding.workflow_version = "v2" + sqlite_session.add_all([source_app, run, execution, binding]) + if source_identity != "unique-legacy-binding": + sqlite_session.add_all([source, unused_binding]) + sqlite_session.commit() + service = AgentObservabilityService(sqlite_session) + params = AgentLogQueryParams(sources=("workflow:workflow-app-1",)) + + logs = service.list_logs( + app=_agent_observability_service_app(app_id="agent-app"), agent_id="agent-1", params=params + ) + messages = service.list_log_messages( + app=_agent_observability_service_app(app_id="agent-app"), + agent_id="agent-1", + conversation_id=execution.id, + params=params, + ) + + assert logs["total"] == messages["total"] == 1 + assert logs["data"][0]["id"] == messages["data"][0]["id"] == execution.id + assert logs["data"][0]["source"]["id"] == "workflow:workflow-app-1:workflow-1:v1:node-1" + assert messages["data"][0]["total_tokens"] == 454_064 + for source_filter in ("workflow:caller-app", "workflow:workflow-app-1:workflow-1:v2:node-1"): + params = AgentLogQueryParams(sources=(source_filter,)) + assert ( + service.list_logs( + app=_agent_observability_service_app(app_id="agent-app"), agent_id="agent-1", params=params + )["total"] + == 0 + ) + assert ( + service.list_log_messages( + app=_agent_observability_service_app(app_id="agent-app"), + agent_id="agent-1", + conversation_id=execution.id, + params=params, + )["total"] + == 0 + ) + + if source_identity == "unique-legacy-binding": + # Without a captured binding or source Workflow, another Agent/version + # makes attribution ambiguous, even when filtering for only agent-1. + unused_binding.agent_id = "another-agent" + sqlite_session.add(unused_binding) + sqlite_session.commit() + params = AgentLogQueryParams(sources=("workflow",)) + assert ( + service.list_logs( + app=_agent_observability_service_app(app_id="agent-app"), agent_id="agent-1", params=params + )["total"] + == 0 + ) + assert ( + service.list_log_messages( + app=_agent_observability_service_app(app_id="agent-app"), + agent_id="agent-1", + conversation_id=execution.id, + params=params, + )["total"] + == 0 + ) + + @staticmethod + @pytest.mark.parametrize("mismatch", ["run-tenant", "source-tenant", "binding-id", "unrelated-run"]) + def test_workflow_agent_observability_rejects_mismatched_ownership(sqlite_session: Session, mismatch: str) -> None: + source_app = _agent_observability_service_app(app_id="workflow-app-1", mode=AppMode.WORKFLOW) + run, execution, binding = ( + _agent_observability_service_workflow_run(), + _agent_observability_service_node_execution(), + _agent_observability_service_workflow_binding(), + ) + execution.process_data = json.dumps({"workflow_agent_binding_id": binding.id}) + if mismatch == "run-tenant": + run.tenant_id = "other-tenant" + elif mismatch == "source-tenant": + source_app.tenant_id = "other-tenant" + elif mismatch == "binding-id": + execution.process_data = json.dumps({"workflow_agent_binding_id": "missing-binding"}) + else: + run.app_id, run.workflow_id = "unrelated-app", "unrelated-workflow" + sqlite_session.add_all([source_app, run, execution, binding]) + sqlite_session.commit() + service = AgentObservabilityService(sqlite_session) + params = AgentLogQueryParams(sources=("workflow",)) + + assert ( + service.list_logs( + app=_agent_observability_service_app(app_id="agent-app"), agent_id="agent-1", params=params + )["total"] + == 0 + ) + assert ( + service.list_log_messages( + app=_agent_observability_service_app(app_id="agent-app"), + agent_id="agent-1", + conversation_id=execution.id, + params=params, + )["total"] + == 0 + ) + + @staticmethod + def test_list_workflow_messages_uses_node_execution_identity( + monkeypatch: pytest.MonkeyPatch, sqlite_session: Session + ) -> None: + node_execution = _agent_observability_service_node_execution() + sqlite_session.add_all( + [ + _agent_observability_service_app(app_id="workflow-app-1", name="Workflow App", mode=AppMode.WORKFLOW), + _agent_observability_service_workflow_run(), + node_execution, + _agent_observability_service_workflow_binding(), + ] + ) + sqlite_session.commit() + service = AgentObservabilityService(sqlite_session) + serialized = {"id": "node-execution-1", "conversation_id": "node-execution-1"} + monkeypatch.setattr(service, "serialize_workflow_node_message", lambda execution: serialized) + + rows = service._list_workflow_messages( + app=_agent_observability_service_app(app_id="agent-app"), + agent_id="agent-1", + conversation_id="node-execution-1", + params=AgentLogQueryParams(), + source_filter=AgentObservabilityService.resolve_source_filter("workflow:workflow-app-1"), + ) + + assert rows == [serialized] + + @staticmethod + def test_apply_workflow_node_filters_supports_time_keyword_and_status() -> None: + stmt = select(WorkflowNodeExecutionModel) + params = AgentLogQueryParams( + start=datetime(2026, 7, 1, tzinfo=UTC), + end=datetime(2026, 8, 1, tzinfo=UTC), + keyword="meeting_100%", + statuses=("success",), + ) + + result = AgentObservabilityService._apply_workflow_node_filters( + stmt, + params=params, + workflow_app=observability_service_module.App, + ) + + assert isinstance(result, Select) + assert len(result._where_criteria) == 4 + + @staticmethod + def test_apply_workflow_node_status_filter_supports_all_status_groups() -> None: + stmt = select(WorkflowNodeExecutionModel) + + result = AgentObservabilityService._apply_workflow_node_status_filter(stmt, ("normal", "error", "paused")) + + assert isinstance(result, Select) + assert len(result._where_criteria) == 1 + empty_stmt = select(WorkflowNodeExecutionModel) + assert AgentObservabilityService._apply_workflow_node_status_filter(empty_stmt, ()) is empty_stmt + assert empty_stmt._where_criteria == () + with pytest.raises(ValueError, match="Unsupported status"): + AgentObservabilityService._apply_workflow_node_status_filter( + select(WorkflowNodeExecutionModel), ("unknown",) + ) + + @staticmethod + def test_source_serializers_return_structured_frontend_shape() -> None: + app = _agent_observability_service_app() + + webapp_source = AgentObservabilityService._serialize_webapp_source(app) + workflow_app_source = AgentObservabilityService._serialize_workflow_app_source(app=app) + workflow_source = AgentObservabilityService._serialize_workflow_source( + app=app, + workflow_id="workflow-1", + workflow_version="v1", + node_id="node-1", + ) + + assert webapp_source == { + "id": "webapp:app-1", + "type": "webapp", + "app_id": "app-1", + "app_name": "Iris", + "app_icon_type": "emoji", + "app_icon": "robot", + "app_icon_background": "#fff", + "workflow_id": None, + "workflow_version": None, + "node_id": None, + } + assert workflow_app_source == { + "id": "workflow:app-1", + "type": "workflow", + "app_id": "app-1", + "app_name": "Iris", + "app_icon_type": "emoji", + "app_icon": "robot", + "app_icon_background": "#fff", + "workflow_id": None, + "workflow_version": None, + "node_id": None, + } + assert workflow_source["id"] == "workflow:app-1:workflow-1:v1:node-1" + assert workflow_source["type"] == "workflow" + assert workflow_source["workflow_id"] == "workflow-1" + + @staticmethod + def test_list_workflow_sources_deduplicates_versions_and_nodes_by_app(sqlite_session: Session) -> None: + app_a = _agent_observability_service_app(app_id="app-a", name="Alpha", mode=AppMode.WORKFLOW) + app_b = _agent_observability_service_app(app_id="app-b", name="Beta", mode=AppMode.WORKFLOW) + sqlite_session.add_all( + [ + app_a, + app_b, + _agent_observability_service_workflow_binding(app_id="app-a", node_id="node-a-1"), + _agent_observability_service_workflow_binding(app_id="app-a", node_id="node-a-2"), + _agent_observability_service_workflow_binding(app_id="app-b", node_id="node-b-1"), + ] + ) + sqlite_session.commit() + service = AgentObservabilityService(sqlite_session) + + sources = service._list_workflow_sources( + app=_agent_observability_service_app(app_id="agent-app"), + agent_id="agent-1", + ) + + assert [source["id"] for source in sources] == ["workflow:app-a", "workflow:app-b"] + + @staticmethod + def test_serialize_log_message_returns_frontend_log_shape() -> None: + created_at = datetime(2026, 6, 17, 1, 2, 3, tzinfo=UTC) + updated_at = datetime(2026, 6, 17, 1, 3, 3, tzinfo=UTC) + message = _agent_observability_service_message(created_at=created_at) + message.updated_at = updated_at + conversation = _agent_observability_service_conversation() + feedbacks = [ + _feedback(), + _feedback( + rating=FeedbackRating.DISLIKE, + content="Needs more detail", + source=FeedbackFromSource.ADMIN, + ), + ] + + payload = AgentObservabilityService.serialize_log_message(message, conversation, feedbacks) + + assert payload == { + "id": "message-1", + "message_id": "message-1", + "conversation_id": "conversation-1", + "conversation_name": "Debug conversation", + "query": "hello", + "answer": "hi", + "status": "success", + "error": None, + "source": "explore", + "from_source": "console", + "from_end_user_id": None, + "from_account_id": "account-1", + "feedback_enabled": True, + "feedbacks": [ + {"rating": "like", "content": "Useful", "from_source": "user"}, + {"rating": "dislike", "content": "Needs more detail", "from_source": "admin"}, + ], + "message_tokens": 3, + "answer_tokens": 4, + "total_tokens": 7, + "total_price": "0.0001", + "currency": "USD", + "latency": 1.25, + "created_at": int(created_at.timestamp()), + "updated_at": int(updated_at.timestamp()), + } + + @staticmethod + def test_serialize_workflow_node_message_returns_frontend_log_shape() -> None: + created_at = datetime(2026, 7, 23, 7, 0, 19, tzinfo=UTC) + finished_at = datetime(2026, 7, 23, 7, 0, 28, tzinfo=UTC) + node_execution = _agent_observability_service_node_execution() + node_execution.inputs = ( + '{"agent_backend_request":{"composition":{"layers":[' + '{"name":"workflow_node_job_prompt","config":{"user":"Summarize the meeting"}},' + '{"name":"workflow_user_prompt","config":{"user":"Focus on action items"}}]}}}' + ) + node_execution.outputs = '{"output":"Alice owns the follow-up."}' + node_execution.execution_metadata = ( + '{"agent_log":{"agent_backend":{"usage":{"prompt_tokens":10,"completion_tokens":5,' + '"total_tokens":15,"total_price":"0.0015","currency":"USD","latency":1.25}}}}' + ) + node_execution.elapsed_time = 1.5 + node_execution.created_at = created_at + node_execution.finished_at = finished_at + + payload = AgentObservabilityService.serialize_workflow_node_message(node_execution) + + assert payload == { + "id": "node-execution-1", + "message_id": "node-execution-1", + "conversation_id": "node-execution-1", + "query": "Summarize the meeting\n\nFocus on action items", + "answer": "Alice owns the follow-up.", + "status": "success", + "error": None, + "from_end_user_id": "end-user-1", + "from_account_id": None, + "feedback_enabled": False, + "feedbacks": [], + "message_tokens": 10, + "answer_tokens": 5, + "total_tokens": 15, + "total_price": "0.0015", + "currency": "USD", + "latency": 1.25, + "created_at": int(created_at.timestamp()), + "updated_at": int(finished_at.timestamp()), + } + + @staticmethod + def test_serialize_workflow_node_message_handles_sparse_runtime_data() -> None: + created_at = datetime(2026, 7, 23, 7, 0, 19, tzinfo=UTC) + node_execution = _agent_observability_service_node_execution( + execution_id="node-execution-2", status=WorkflowNodeExecutionStatus.PAUSED + ) + node_execution.title = "Fallback prompt" + node_execution.inputs = json.dumps( + { + "agent_backend_request": { + "composition": { + "layers": [ + None, + {"name": "unrelated", "config": {"user": "ignored"}}, + {"name": "workflow_user_prompt", "config": {"user": " "}}, + ] + } + } + } + ) + node_execution.outputs = json.dumps({"output": {"structured": True}}) + node_execution.execution_metadata = json.dumps( + { + "agent_log": { + "agent_backend": { + "usage": { + "prompt_tokens": "2", + "completion_tokens": 3, + } + } + } + } + ) + node_execution.elapsed_time = 2 + node_execution.created_by_role = CreatorUserRole.ACCOUNT + node_execution.created_by = "account-1" + node_execution.created_at = created_at + node_execution.finished_at = None + + payload = AgentObservabilityService.serialize_workflow_node_message(node_execution) + + assert payload["query"] == "Fallback prompt" + assert payload["answer"] == '{"output": {"structured": true}}' + assert payload["status"] == "paused" + assert payload["from_end_user_id"] is None + assert payload["from_account_id"] == "account-1" + assert payload["total_tokens"] == 5 + assert payload["total_price"] == "0" + assert payload["currency"] == "" + assert payload["latency"] == 2.0 + assert payload["updated_at"] == int(created_at.timestamp()) + + @staticmethod + def test_positive_feedback_rate_uses_rated_messages_as_denominator() -> None: + assert AgentObservabilityService._positive_feedback_rate(like_count=2, total_count=4) == 0.5 + assert AgentObservabilityService._positive_feedback_rate(like_count=0, total_count=1) == 0 + assert AgentObservabilityService._positive_feedback_rate(like_count=None, total_count=0) is None + + @staticmethod + def test_list_message_feedbacks_groups_feedbacks_by_message(sqlite_session: Session) -> None: + first_message = _agent_observability_service_message() + second_message = _agent_observability_service_message(message_id="message-2") + feedbacks = [ + _feedback(), + _feedback(rating=FeedbackRating.DISLIKE), + _feedback(message_id="message-2"), + ] + sqlite_session.add_all( + [ + _agent_observability_service_app(), + _agent_observability_service_conversation(), + first_message, + second_message, + *feedbacks, + ] + ) + sqlite_session.commit() + service = AgentObservabilityService(sqlite_session) + + grouped_feedbacks = service._list_message_feedbacks( + app=_agent_observability_service_app(), + messages=[first_message, second_message], + ) + + assert {feedback.id for feedback in grouped_feedbacks["message-1"]} == { + feedbacks[0].id, + feedbacks[1].id, + } + assert grouped_feedbacks["message-2"] == [feedbacks[2]] + assert service._list_message_feedbacks(app=_agent_observability_service_app(), messages=[]) == {} + + @staticmethod + def test_list_conversation_feedback_rates_maps_user_and_admin_sources(sqlite_session: Session) -> None: + messages = [_agent_observability_service_message(message_id=f"message-{index}") for index in range(1, 5)] + feedbacks = [ + _feedback(message_id="message-1"), + _feedback(message_id="message-2"), + _feedback(message_id="message-3", rating=FeedbackRating.DISLIKE), + _feedback(message_id="message-4", rating=FeedbackRating.DISLIKE), + _feedback(message_id="message-1", source=FeedbackFromSource.ADMIN), + ] + sqlite_session.add_all( + [_agent_observability_service_app(), _agent_observability_service_conversation(), *messages, *feedbacks] + ) + sqlite_session.commit() + service = AgentObservabilityService(sqlite_session) + + rates = service._list_conversation_feedback_rates( + app=_agent_observability_service_app(), + conversation_ids=["conversation-1"], + ) + + assert rates == {"conversation-1": {"user_rate": 0.5, "operation_rate": 1.0}} + assert ( + service._list_conversation_feedback_rates( + app=_agent_observability_service_app(), + conversation_ids=[], + ) + == {} + ) + + @staticmethod + def test_workflow_node_serialization_helpers_handle_invalid_values() -> None: + assert AgentObservabilityService._json_mapping(None) == {} + assert AgentObservabilityService._json_mapping("not-json") == {} + assert AgentObservabilityService._json_mapping("[]") == {} + assert AgentObservabilityService._mapping_value({"value": []}, "value") == {} + assert AgentObservabilityService._int_value(None) == 0 + assert AgentObservabilityService._int_value("not-a-number") == 0 + assert ( + AgentObservabilityService._workflow_node_query( + {"agent_backend_request": {"composition": {"layers": "invalid"}}}, fallback="fallback" + ) + == "fallback" + ) + assert ( + AgentObservabilityService._workflow_node_answer({"output": 1, "text": "fallback text"}) == "fallback text" + ) + assert AgentObservabilityService._workflow_node_answer({}) == "" + + @staticmethod + def test_serialize_workflow_execution_log_uses_node_execution_identity() -> None: + created_at = datetime(2026, 7, 23, 7, 0, 19, tzinfo=UTC) + node_execution = _agent_observability_service_node_execution(status=WorkflowNodeExecutionStatus.FAILED) + node_execution.created_by_role = CreatorUserRole.ACCOUNT + node_execution.created_by = "account-1" + node_execution.created_at = created_at + + payload = AgentObservabilityService._serialize_workflow_execution_log( + node_execution_id=node_execution.id, + title=node_execution.title, + status=node_execution.status, + created_by_role=node_execution.created_by_role, + created_by=node_execution.created_by, + created_at=node_execution.created_at, + finished_at=node_execution.finished_at, + source={"id": "workflow:app-1:workflow-1:v1:node-1"}, + ) + + assert payload["id"] == "node-execution-1" + assert payload["conversation_id"] == "node-execution-1" + assert payload["message_count"] == 1 + assert payload["end_user_id"] is None + assert payload["status"] == "failed" + assert payload["unread"] is False + + @staticmethod + def test_build_charts_and_summary_match_monitoring_metrics() -> None: + rows = [ + { + "date": "2026-06-16", + "message_count": 2, + "conversation_count": 1, + "end_user_count": 1, + "token_count": 30, + "total_price": Decimal("0.003"), + "avg_latency": 1.5, + "latency_sum": 3, + "answer_tokens": 12, + "like_count": 1, + }, + { + "date": "2026-06-17", + "message_count": 1, + "conversation_count": 1, + "end_user_count": 1, + "token_count": 20, + "total_price": Decimal("0.002"), + "avg_latency": 2, + "latency_sum": 2, + "answer_tokens": 8, + "like_count": 1, + }, + ] + + charts = AgentObservabilityService._build_charts(rows) + summary = AgentObservabilityService._build_summary(rows) + + assert charts["token_usage"] == [ + {"date": "2026-06-16", "token_count": 30, "total_price": "0.003", "currency": "USD"}, + {"date": "2026-06-17", "token_count": 20, "total_price": "0.002", "currency": "USD"}, + ] + assert charts["average_response_time"] == [ + {"date": "2026-06-16", "latency": 1500.0}, + {"date": "2026-06-17", "latency": 2000.0}, + ] + assert summary == { + "total_messages": 3, + "total_conversations": 2, + "total_end_users": 2, + "total_tokens": 50, + "total_price": "0.005", + "currency": "USD", + "average_session_interactions": 1.5, + "average_response_time": 1666.6667, + "tokens_per_second": 4.0, + "user_satisfaction_rate": 66.67, + } + + +"""Real SQL coverage for Agent statistics inside a Workflow Tool invocation.""" + + +_workflow_tool_agent_statistics_app = _agent_observability_service_app +_workflow_tool_agent_statistics_conversation = _agent_observability_service_conversation +_workflow_tool_agent_statistics_message = _agent_observability_service_message +_workflow_tool_agent_statistics_node_execution = _agent_observability_service_node_execution +_workflow_tool_agent_statistics_workflow_binding = _agent_observability_service_workflow_binding +_workflow_tool_agent_statistics_workflow_run = _agent_observability_service_workflow_run + + +def _seed_nested_agent(sqlite_session: Session, *, workflow_type: WorkflowType) -> App: + agent_app = _workflow_tool_agent_statistics_app(app_id="agent-app") + outer_app = _workflow_tool_agent_statistics_app( + app_id="outer-app", mode=AppMode.ADVANCED_CHAT if workflow_type == WorkflowType.CHAT else AppMode.WORKFLOW + ) + source_app = _workflow_tool_agent_statistics_app(app_id="source-app", mode=AppMode.WORKFLOW) + run = _workflow_tool_agent_statistics_workflow_run(workflow_type=workflow_type) + run.app_id = outer_app.id + run.workflow_id = "outer-workflow" + run.version = "outer-v1" + source_binding = _workflow_tool_agent_statistics_workflow_binding( + app_id=source_app.id, binding_id="source-binding-v1", node_id="agent-node" + ) + source_binding.workflow_id = "source-workflow" + source_binding.workflow_version = "source-v1" + unused_binding = _workflow_tool_agent_statistics_workflow_binding( + app_id=source_app.id, binding_id="source-binding-v2", node_id="agent-node" + ) + unused_binding.workflow_id = source_binding.workflow_id + unused_binding.workflow_version = "source-v2" + outer_binding = _workflow_tool_agent_statistics_workflow_binding(app_id=outer_app.id, node_id="agent-node") + outer_binding.workflow_id = run.workflow_id + outer_binding.workflow_version = run.version + other_agent_binding = _workflow_tool_agent_statistics_workflow_binding( + app_id=source_app.id, node_id="other-agent-node" + ) + other_agent_binding.workflow_id = source_binding.workflow_id + other_agent_binding.workflow_version = source_binding.workflow_version + other_agent_binding.agent_id = "other-agent" + sqlite_session.add_all( + [agent_app, outer_app, source_app, run, source_binding, unused_binding, outer_binding, other_agent_binding] + ) + + for index, (binding, tokens, completion_tokens, price, latency) in enumerate( + [ + (source_binding, 7, 2, "0.01", 2.0), + (source_binding, 11, 4, "0.02", 3.0), + (other_agent_binding, 900, 800, "9.00", 90.0), + ] + ): + execution = _workflow_tool_agent_statistics_node_execution(execution_id=f"tool-agent-execution-{index}") + execution.app_id = source_app.id + execution.workflow_id = source_binding.workflow_id + execution.node_id = binding.node_id + execution.triggered_from = WorkflowNodeExecutionTriggeredFrom.WORKFLOW_TOOL + execution.process_data = json.dumps({WORKFLOW_AGENT_BINDING_ID_KEY: binding.id}) + execution.elapsed_time = latency + execution.execution_metadata = json.dumps( + { + "agent_log": { + "agent_backend": { + "usage": {"total_tokens": tokens, "completion_tokens": completion_tokens, "total_price": price} + } + } + } + ) + sqlite_session.add(execution) + + if workflow_type == WorkflowType.CHAT: + conversation = _workflow_tool_agent_statistics_conversation(app_id=outer_app.id) + conversation.mode = AppMode.ADVANCED_CHAT + conversation.from_source = ConversationFromSource.API + message = _workflow_tool_agent_statistics_message(app_id=outer_app.id, created_at=datetime(2026, 7, 22, 8, 0)) + message.workflow_run_id = run.id + message.from_source = ConversationFromSource.API + message.invoke_from = _workflow_tool_agent_statistics_InvokeFrom.WEB_APP + message.from_account_id = None + message.from_end_user_id = "end-user-1" + sqlite_session.add_all([conversation, message]) + sqlite_session.commit() + return agent_app + + +class TestWorkflowToolAgentStatisticsCases: + @staticmethod + @pytest.fixture + def statistics_service(monkeypatch: pytest.MonkeyPatch, sqlite_session: Session) -> AgentObservabilityService: + # Keep the production joins/aggregation; adapt only MySQL functions to SQLite. + apply_config_overrides(monkeypatch, DB_TYPE="mysql") + monkeypatch.setattr(observability_service_module, "convert_datetime_to_date", lambda field: f"DATE({field})") + connection = sqlite_session.connection().connection.dbapi_connection + assert isinstance(connection, Connection) + connection.create_function("JSON_UNQUOTE", 1, lambda value: value) + return AgentObservabilityService(sqlite_session) + + @staticmethod + @pytest.mark.parametrize("workflow_type", [WorkflowType.WORKFLOW, WorkflowType.CHAT]) + def test_nested_agent_daily_statistics_preserve_source_usage_and_date( + sqlite_session: Session, statistics_service: AgentObservabilityService, workflow_type: WorkflowType + ) -> None: + app = _seed_nested_agent(sqlite_session, workflow_type=workflow_type) + # Workflow metrics remain grouped by the outer run's date; Chatflow metrics + # retain the original message's date and usage. Child nodes executed July 23. + day = 22 if workflow_type == WorkflowType.CHAT else 21 + date = f"2026-07-{day}" + expected_tokens = 7 if workflow_type == WorkflowType.CHAT else 18 + expected_price = Decimal("0.0001") if workflow_type == WorkflowType.CHAT else Decimal("0.03") + expected_latency_ms = 1250.0 if workflow_type == WorkflowType.CHAT else 5000.0 + expected_tps = 3.2 if workflow_type == WorkflowType.CHAT else 1.2 + + for source in ("workflow:source-app", "workflow:source-app:source-workflow:source-v1:agent-node", "workflow"): + payload = statistics_service.get_statistics_summary( + app=app, + agent_id="agent-1", + params=AgentStatisticsQueryParams( + source=source, + start=datetime(2026, 7, day), + end=datetime(2026, 7, day + 1), + ), + ) + + summary = payload["summary"] + assert summary["total_messages"] == 1 + assert summary["total_conversations"] == 1 + assert summary["total_end_users"] == 1 + assert summary["total_tokens"] == expected_tokens + assert Decimal(summary["total_price"]) == expected_price + assert summary["average_response_time"] == expected_latency_ms + assert summary["tokens_per_second"] == expected_tps + assert payload["charts"]["daily_messages"] == [{"date": date, "message_count": 1}] + assert payload["charts"]["average_response_time"] == [{"date": date, "latency": expected_latency_ms}] + + for unrelated_source in ( + "workflow:outer-app", + "workflow:source-app:source-workflow:source-v2:agent-node", + "workflow:source-app:source-workflow:source-v1:other-agent-node", + ): + payload = statistics_service.get_statistics_summary( + app=app, agent_id="agent-1", params=AgentStatisticsQueryParams(source=unrelated_source) + ) + assert payload["summary"]["total_messages"] == 0 + assert payload["summary"]["total_tokens"] == 0 + assert payload["charts"]["daily_messages"] == [] + + @staticmethod + @pytest.mark.parametrize("workflow_type", [WorkflowType.WORKFLOW, WorkflowType.CHAT]) + @pytest.mark.parametrize( + "invalid_owner", ["outer_tenant", "source_tenant", "node_app", "node_workflow", "binding_id"] + ) + def test_nested_agent_statistics_reject_mismatched_ownership( + sqlite_session: Session, + statistics_service: AgentObservabilityService, + workflow_type: WorkflowType, + invalid_owner: str, + ) -> None: + app = _seed_nested_agent(sqlite_session, workflow_type=workflow_type) + if invalid_owner == "outer_tenant": + run = sqlite_session.get(WorkflowRun, "workflow-run-1") + assert run is not None + run.tenant_id = "other-tenant" + elif invalid_owner == "source_tenant": + source_app = sqlite_session.get(App, "source-app") + assert source_app is not None + source_app.tenant_id = "other-tenant" + else: + executions = sqlite_session.scalars( + select(WorkflowNodeExecutionModel).where( + WorkflowNodeExecutionModel.id.in_(["tool-agent-execution-0", "tool-agent-execution-1"]) + ) + ).all() + assert len(executions) == 2 + for execution in executions: + if invalid_owner == "node_app": + execution.app_id = "other-app" + elif invalid_owner == "node_workflow": + execution.workflow_id = "other-workflow" + else: + execution.process_data = json.dumps( + {WORKFLOW_AGENT_BINDING_ID_KEY: "binding-source-app-other-agent-node"} + ) + sqlite_session.commit() + + payload = statistics_service.get_statistics_summary( + app=app, agent_id="agent-1", params=AgentStatisticsQueryParams(source="workflow:source-app") + ) + + assert payload["summary"]["total_messages"] == 0 + assert payload["summary"]["total_tokens"] == 0 + assert payload["charts"]["daily_messages"] == [] + + @staticmethod + def test_nested_agent_chat_statistics_reject_message_from_another_app( + sqlite_session: Session, statistics_service: AgentObservabilityService + ) -> None: + app = _seed_nested_agent(sqlite_session, workflow_type=WorkflowType.CHAT) + message = sqlite_session.get(Message, "message-1") + assert message is not None + message.app_id = "other-app" + sqlite_session.commit() + + payload = statistics_service.get_statistics_summary( + app=app, agent_id="agent-1", params=AgentStatisticsQueryParams(source="workflow:source-app") + ) + + assert payload["summary"]["total_messages"] == 0 + assert payload["summary"]["total_tokens"] == 0 + assert payload["charts"]["daily_messages"] == [] diff --git a/api/tests/unit_tests/services/agent/test_roster_package_service.py b/api/tests/unit_tests/services/agent/test_roster_package_reader.py similarity index 99% rename from api/tests/unit_tests/services/agent/test_roster_package_service.py rename to api/tests/unit_tests/services/agent/test_roster_package_reader.py index dcd02616dc9..d80aa4639c2 100644 --- a/api/tests/unit_tests/services/agent/test_roster_package_service.py +++ b/api/tests/unit_tests/services/agent/test_roster_package_reader.py @@ -58,7 +58,7 @@ from services.agent.roster_package_exporter import RosterAgentPackageExporter from services.agent.roster_package_reader import RosterAgentPackageReader from services.app_dsl_service import AppDslService from services.plugin.dependencies_analysis import DependenciesAnalysisService -from tests.unit_tests.config_override import apply_config_overrides +from tests.helpers.unit.config_override import apply_config_overrides class _MemoryStorage: diff --git a/api/tests/unit_tests/services/agent/test_agent_services.py b/api/tests/unit_tests/services/agent/test_roster_service.py similarity index 99% rename from api/tests/unit_tests/services/agent/test_agent_services.py rename to api/tests/unit_tests/services/agent/test_roster_service.py index 80e4c4ff2f5..fd240bebe1c 100644 --- a/api/tests/unit_tests/services/agent/test_agent_services.py +++ b/api/tests/unit_tests/services/agent/test_roster_service.py @@ -67,7 +67,7 @@ from services.entities.agent_entities import ( ComposerSaveStrategy, ComposerVariant, ) -from tests.unit_tests.model_factories import make_account, make_app, make_conversation, make_workflow +from tests.helpers.unit.model_factories import make_account, make_app, make_conversation, make_workflow def _agent_soul_with_model() -> AgentSoulConfig: diff --git a/api/tests/unit_tests/services/agent/test_agent_runtime_config_service.py b/api/tests/unit_tests/services/agent/test_runtime_config_service.py similarity index 100% rename from api/tests/unit_tests/services/agent/test_agent_runtime_config_service.py rename to api/tests/unit_tests/services/agent/test_runtime_config_service.py diff --git a/api/tests/unit_tests/services/agent/test_runtime_usage_service.py b/api/tests/unit_tests/services/agent/test_runtime_usage_service.py index 28859d264ca..0c104a3ad0a 100644 --- a/api/tests/unit_tests/services/agent/test_runtime_usage_service.py +++ b/api/tests/unit_tests/services/agent/test_runtime_usage_service.py @@ -26,7 +26,7 @@ from models.agent import ( from models.agent_sandbox_usage import AgentSandboxExecution, AgentSandboxUsageEvent from models.model import App, AppMode from services.agent.runtime_usage_service import SandboxUsageError, SandboxUsageEvent, SandboxUsageService -from tests.unit_tests.config_override import config_overrides_context +from tests.helpers.unit.config_override import config_overrides_context PROJECT = "431de237-596f-4d59-8a85-20a9846bf243" START = "2026-09-20T00:00:00Z" diff --git a/api/tests/unit_tests/services/agent/test_skill_package_service.py b/api/tests/unit_tests/services/agent/test_skill_package_service.py index d62318434ca..7ae9d504e85 100644 --- a/api/tests/unit_tests/services/agent/test_skill_package_service.py +++ b/api/tests/unit_tests/services/agent/test_skill_package_service.py @@ -10,7 +10,7 @@ import zlib import pytest from services.agent.skill_package_service import NormalizedSkillPackage, SkillPackageError, SkillPackageService -from tests.unit_tests.config_override import apply_config_overrides +from tests.helpers.unit.config_override import apply_config_overrides _SKILL_MD = """--- name: pdf-toolkit diff --git a/api/tests/unit_tests/services/agent/test_workflow_package_service.py b/api/tests/unit_tests/services/agent/test_workflow_package_service.py deleted file mode 100644 index ad01386e63b..00000000000 --- a/api/tests/unit_tests/services/agent/test_workflow_package_service.py +++ /dev/null @@ -1,537 +0,0 @@ -"""Workflow packages restore node-owned Agents and assets across workspaces.""" - -from __future__ import annotations - -import hashlib -import io -import zipfile -from collections.abc import Callable, Generator -from types import SimpleNamespace -from typing import cast -from unittest.mock import Mock - -import pytest -import yaml -from sqlalchemy import Engine, select -from sqlalchemy.orm import Session -from sqlalchemy.pool import QueuePool - -from models import Account, App, AppMode -from models.agent import ( - Agent, - AgentConfigSnapshot, - AgentScope, - AgentSource, - AgentStatus, - WorkflowAgentBindingType, - WorkflowAgentNodeBinding, -) -from models.agent_config_entities import AgentSoulConfig -from models.model import UploadFile -from models.skill import AgentSkillBinding, AgentSkillBindingSnapshot, Skill, SkillVersion, SkillVersionManifest -from models.tools import ToolFile -from models.workflow import Workflow, WorkflowType -from services.agent.dsl_entities import AgentPackage -from services.agent.errors import InvalidRosterAgentPackageError, RosterAgentPackageExportFailedError -from services.agent.package_resource_exporter import AgentPackageResourceExporter -from services.agent.package_resource_importer import AgentPackageResourceImporter -from services.app_dsl_service import AppDslService, Import -from services.app_package_service import AppPackageManifest, AppPackageService -from services.entities.dsl_entities import ImportStatus -from services.errors.base import NoPermissionError -from services.plugin.dependencies_analysis import DependenciesAnalysisService -from tests.unit_tests.model_factories import make_account, make_app, make_tenant, make_upload_file, make_workflow - - -class _MemoryStorage: - def __init__(self) -> None: - self.files: dict[str, bytes] = {} - self.saved: list[str] = [] - - def load_stream(self, filename: str) -> Generator[bytes]: - yield self.files[filename] - - def save(self, filename: str, data: bytes) -> None: - self.files[filename] = data - self.saved.append(filename) - - -@pytest.fixture -def storage(monkeypatch: pytest.MonkeyPatch, sqlite_engine: Engine) -> _MemoryStorage: - backend = _MemoryStorage() - monkeypatch.setattr("services.workflow_service.db", SimpleNamespace(engine=sqlite_engine)) - monkeypatch.setattr( - "services.app_package_service.AgentPackageResourceExporter", - lambda: AgentPackageResourceExporter(storage_backend=backend), - ) - monkeypatch.setattr( - "services.app_package_service.AgentPackageResourceImporter", - lambda: AgentPackageResourceImporter(storage_backend=backend), - ) - monkeypatch.setattr(DependenciesAnalysisService, "generate_dependencies", lambda **_kwargs: []) - return backend - - -def _skill(name: str, marker: str) -> bytes: - output = io.BytesIO() - with zipfile.ZipFile(output, "w") as archive: - archive.writestr("SKILL.md", f"---\nname: {name}\ndescription: {marker}\n---\n\n# Instructions\n{marker}\n") - archive.writestr("scripts/run.py", f"print({marker!r})\n") - return output.getvalue() - - -def _tool_file(session: Session, storage: _MemoryStorage, key: str, payload: bytes) -> ToolFile: - row = ToolFile( - tenant_id="tenant-1", - user_id="account-1", - conversation_id=None, - file_key=key, - name=key.rsplit("/", 1)[-1], - mimetype="application/zip", - size=len(payload), - original_url=None, - ) - storage.files[key] = payload - session.add(row) - return row - - -def _source( - session: Session, storage: _MemoryStorage, *, mode: AppMode = AppMode.WORKFLOW, published: bool = False -) -> App: - app = make_app(mode=mode, name="Inline Workflow") - nodes = [ - { - "id": f"node-{index}", - "data": {"type": "agent", "version": "2", "title": f"Agent {index}", "agent_node_kind": "dify_agent"}, - } - for index in (1, 2) - ] - workflow = make_workflow( - workflow_id="workflow-1", - app_id=app.id, - graph={"nodes": nodes, "edges": []}, - workflow_type=WorkflowType.CHAT if mode == AppMode.ADVANCED_CHAT else WorkflowType.WORKFLOW, - version="2026-01-01 00:00:00" if published else Workflow.VERSION_DRAFT, - ) - session.add_all([app, workflow]) - for index in (1, 2): - skill_file = _tool_file( - session, storage, f"tools/skill-{index}.zip", _skill("research", f"Agent {index} skill") - ) - notes = _tool_file(session, storage, f"tools/notes-{index}.txt", f"Agent {index} notes".encode()) - upload = make_upload_file(file_id=f"upload-{index}", key=f"uploads/guide-{index}.txt", name="guide.txt") - storage.files[upload.key] = f"Agent {index} guide".encode() - agent = Agent( - id=f"source-agent-{index}", - tenant_id=app.tenant_id, - name=f"Inline {index}", - scope=AgentScope.WORKFLOW_ONLY, - source=AgentSource.WORKFLOW, - app_id=app.id, - workflow_id=workflow.id, - workflow_node_id=f"node-{index}", - status=AgentStatus.ACTIVE, - created_by="account-1", - ) - soul = AgentSoulConfig.model_validate( - { - "prompt": {"system_prompt": f"Bound snapshot {index}"}, - "config_skills": [{"name": "research", "file_id": skill_file.id}], - "config_files": [ - {"name": "notes.txt", "file_kind": "tool_file", "file_id": notes.id}, - {"name": "guide.txt", "file_kind": "upload_file", "file_id": upload.id}, - ], - } - ) - snapshot = AgentConfigSnapshot( - id=f"source-snapshot-{index}", - tenant_id=app.tenant_id, - agent_id=agent.id, - version=1, - config_snapshot=soul, - created_by="account-1", - ) - active = AgentConfigSnapshot( - id=f"active-snapshot-{index}", - tenant_id=app.tenant_id, - agent_id=agent.id, - version=2, - config_snapshot=AgentSoulConfig.model_validate({"prompt": {"system_prompt": "Unbound active snapshot"}}), - created_by="account-1", - ) - agent.active_config_snapshot_id = active.id - binding = WorkflowAgentNodeBinding( - tenant_id=app.tenant_id, - app_id=app.id, - workflow_id=workflow.id, - workflow_version=workflow.version, - node_id=f"node-{index}", - binding_type=WorkflowAgentBindingType.INLINE_AGENT, - agent_id=agent.id, - current_snapshot_id=snapshot.id, - node_job_config={"workflow_prompt": f"Job {index}"}, - created_by="account-1", - ) - session.add_all([upload, agent, snapshot, active, binding]) - if index == 1: - workspace_archive = _tool_file( - session, storage, "tools/workspace.zip", _skill("workspace", "Workspace skill") - ) - skill = Skill( - id="source-workspace-skill", - tenant_id=app.tenant_id, - name="workspace", - display_name="Workspace", - description="Workspace skill", - created_by="account-1", - updated_by="account-1", - ) - version = SkillVersion( - id="source-skill-version", - skill_id=skill.id, - version_number=1, - version_name="1.0", - publish_note="", - manifest=SkillVersionManifest( - name="workspace", display_name="Workspace", description="Workspace skill", files=[] - ), - archive_tool_file_id=workspace_archive.id, - archive_size=workspace_archive.size, - hash_code="workspace-hash", - published_by="account-1", - ) - skill.latest_published_version_id = version.id - session.add_all( - [ - skill, - version, - AgentSkillBindingSnapshot( - tenant_id=app.tenant_id, - agent_id=agent.id, - config_snapshot_id=snapshot.id, - skill_id=skill.id, - priority=0, - created_by="account-1", - ), - ] - ) - session.commit() - return app - - -def _account(tenant_id: str = "tenant-2") -> Account: - return make_account(tenant=make_tenant(tenant_id=tenant_id)) - - -def _import(session: Session, content: bytes, *, app_id: str | None = None, account: Account | None = None) -> Import: - prepared = AppPackageService().read_package(io.BytesIO(content)) - assert prepared is not None - with prepared: - result = AppDslService(session).import_app( - account=account or _account(), - import_mode="yaml-content", - yaml_content=prepared.dsl, - app_id=app_id, - package=prepared, - ) - if result.status == ImportStatus.FAILED: - session.rollback() - else: - session.commit() - return result - - -def _rewrite( - content: bytes, mutate: Callable[[AppPackageManifest, dict[str, object], dict[str, bytes]], None] -) -> bytes: - with zipfile.ZipFile(io.BytesIO(content)) as archive: - members = {name: archive.read(name) for name in archive.namelist()} - manifest = AppPackageManifest.model_validate(yaml.safe_load(members["manifest.yaml"])) - data = yaml.safe_load(members["app.yaml"]) - mutate(manifest, data, members) - members["app.yaml"] = yaml.safe_dump(data).encode() - manifest.apps[0].size = len(members["app.yaml"]) - manifest.apps[0].sha256 = hashlib.sha256(members["app.yaml"]).hexdigest() - members["manifest.yaml"] = yaml.safe_dump(manifest.model_dump(mode="json")).encode() - output = io.BytesIO() - with zipfile.ZipFile(output, "w") as archive: - for name, payload in members.items(): - archive.writestr(name, payload) - return output.getvalue() - - -@pytest.mark.parametrize("mode", [AppMode.WORKFLOW, AppMode.ADVANCED_CHAT]) -@pytest.mark.parametrize("published", [False, True]) -def test_package_round_trip_restores_bound_agents_and_resources( - sqlite_session: Session, - storage: _MemoryStorage, - mode: AppMode, - published: bool, -) -> None: - app = _source(sqlite_session, storage, mode=mode, published=published) - with AppPackageService().export_app(app_model=app, workflow_id="workflow-1" if published else None) as exported: - content = exported.archive.read() - with zipfile.ZipFile(io.BytesIO(content)) as archive: - manifest = yaml.safe_load(archive.read("manifest.yaml")) - assert all( - "audit" not in resource - for group in manifest["agent_resources"].values() - for resource in [*group["skills"], *group["files"]] - ) - prepared = AppPackageService().read_package(io.BytesIO(content)) - assert prepared is not None - with prepared: - assert len(prepared.agent_resources) == 2 - assert [p.soul.prompt.system_prompt for p in prepared.agents.values()] == [ - "Bound snapshot 1", - "Bound snapshot 2", - ] - assert all(not p.omitted_assets for p in prepared.agents.values()) - assert {r.name for group in prepared.agent_resources.values() for r in group.skills} == { - "research", - "workspace", - } - result = _import(sqlite_session, content) - assert result.status == ImportStatus.COMPLETED, result.error - assert result.app_mode == mode - bindings = sqlite_session.scalars( - select(WorkflowAgentNodeBinding) - .where(WorkflowAgentNodeBinding.app_id == result.app_id) - .order_by(WorkflowAgentNodeBinding.node_id) - ).all() - assert len(bindings) == 2 - for index, binding in enumerate(bindings, 1): - assert binding.tenant_id == "tenant-2" - assert binding.binding_type == WorkflowAgentBindingType.INLINE_AGENT - assert binding.node_job_config_dict["workflow_prompt"] == f"Job {index}" - agent = sqlite_session.get(Agent, binding.agent_id) - assert agent is not None - assert (agent.scope, agent.app_id, agent.workflow_id, agent.workflow_node_id) == ( - AgentScope.WORKFLOW_ONLY, - result.app_id, - binding.workflow_id, - binding.node_id, - ) - snapshot = sqlite_session.get(AgentConfigSnapshot, binding.current_snapshot_id) - assert snapshot is not None - soul = AgentSoulConfig.model_validate(snapshot.config_snapshot_dict) - assert soul.prompt.system_prompt == f"Bound snapshot {index}" - assert len(soul.config_skills) == (2 if index == 1 else 1) - for ref in soul.config_skills: - file = sqlite_session.get(ToolFile, ref.file_id) - assert file is not None - assert file.tenant_id == "tenant-2" - assert not ref.is_missing - with zipfile.ZipFile(io.BytesIO(storage.files[file.file_key])) as archive: - expected = f"Agent {index} skill" if ref.name == "research" else "Workspace skill" - assert expected.encode() in archive.read("SKILL.md") - for ref in soul.config_files: - row = sqlite_session.get(ToolFile if ref.file_kind == "tool_file" else UploadFile, ref.file_id) - assert row is not None - assert row.tenant_id == "tenant-2" - key = row.file_key if isinstance(row, ToolFile) else row.key - expected = f"Agent {index} notes" if ref.name == "notes.txt" else f"Agent {index} guide" - assert storage.files[key] == expected.encode() - if isinstance(row, UploadFile): - assert row.used is True - assert row.used_by == "account-1" - assert sqlite_session.scalar(select(AgentSkillBinding).where(AgentSkillBinding.agent_id == agent.id)) is None - - -def test_yaml_export_still_omits_resources(sqlite_session: Session, storage: _MemoryStorage) -> None: - app = _source(sqlite_session, storage) - data = yaml.safe_load(AppDslService.export_dsl(app, session=sqlite_session)) - for package in data["agent_packages"].values(): - assert package["omitted_assets"] - assert all(item["is_missing"] and item["file_id"] == "" for item in package["soul"]["config_files"]) - - -def test_export_rejects_cross_tenant_assets(sqlite_session: Session, storage: _MemoryStorage) -> None: - app = _source(sqlite_session, storage) - row = sqlite_session.get(UploadFile, "upload-1") - assert row is not None - row.tenant_id = "tenant-elsewhere" - sqlite_session.commit() - with pytest.raises(RosterAgentPackageExportFailedError, match="payload is unavailable"): - AppPackageService().export_app(app_model=app) - - -@pytest.mark.parametrize( - "mutation", ["unknown_agent", "cross_agent_file", "unreferenced_agent", "duplicate_path", "file_checksum"] -) -def test_rejects_invalid_resource_ownership_before_upload( - sqlite_session: Session, - storage: _MemoryStorage, - mutation: str, -) -> None: - app = _source(sqlite_session, storage) - with AppPackageService().export_app(app_model=app) as exported: - content = exported.archive.read() - - def mutate(manifest: AppPackageManifest, data: dict[str, object], members: dict[str, bytes]) -> None: - resources = manifest.agent_resources - if mutation == "unknown_agent": - resources["unknown"] = resources.pop("agent_1") - elif mutation == "cross_agent_file": - agents = cast(dict[str, object], data["agent_packages"]) - first = AgentPackage.model_validate(agents["agent_1"]) - second = AgentPackage.model_validate(agents["agent_2"]) - first.soul.config_files[0].file_id = second.soul.config_files[0].file_id - agents["agent_1"] = first.model_dump(mode="json") - elif mutation == "unreferenced_agent": - cast(dict[str, dict[str, list[object]]], data["workflow"])["graph"]["nodes"].pop() - elif mutation == "duplicate_path": - resources["agent_2"].files[0] = resources["agent_1"].files[0] - else: - resource = resources["agent_1"].files[0] - members[resource.path] = b"x" * resource.size - - with pytest.raises(InvalidRosterAgentPackageError): - AppPackageService().read_package(io.BytesIO(_rewrite(content, mutate))) - assert storage.saved == [] - - -def test_pending_import_retains_materialized_resources_and_damage_warnings( - sqlite_session: Session, - storage: _MemoryStorage, - monkeypatch: pytest.MonkeyPatch, -) -> None: - app = _source(sqlite_session, storage) - with AppPackageService().export_app(app_model=app) as exported: - content = exported.archive.read() - - def mutate(manifest: AppPackageManifest, data: dict[str, object], members: dict[str, bytes]) -> None: - resources = manifest.agent_resources - data["version"] = "99.0.0" - resource = resources["agent_1"].skills[0] - members[resource.path] = b"x" * resource.size - - pending: dict[str, str] = {} - monkeypatch.setattr( - "services.app_dsl_service.redis_client.setex", lambda key, _ttl, value: pending.__setitem__(key, value) - ) - monkeypatch.setattr("services.app_dsl_service.redis_client.get", pending.get) - monkeypatch.setattr("services.app_dsl_service.redis_client.delete", pending.pop) - result = _import(sqlite_session, _rewrite(content, mutate)) - assert result.status == ImportStatus.PENDING - saved_count = len(storage.saved) - confirmed = AppDslService(sqlite_session).confirm_import(import_id=result.id, account=_account()) - assert confirmed.status == ImportStatus.COMPLETED_WITH_WARNINGS, confirmed.error - assert any( - warning.code == "agent_skill_missing" and warning.path.startswith("agent_packages.agent_1.") - for warning in confirmed.warnings - ) - assert len(storage.saved) == saved_count - assert not pending - - -def test_overwrite_permission_is_checked_before_resource_upload( - sqlite_session: Session, - storage: _MemoryStorage, - monkeypatch: pytest.MonkeyPatch, -) -> None: - app = _source(sqlite_session, storage) - with AppPackageService().export_app(app_model=app) as exported: - content = exported.archive.read() - monkeypatch.setattr(AppDslService, "_load_app_for_overwrite", Mock(side_effect=NoPermissionError("Denied"))) - with pytest.raises(NoPermissionError): - _import(sqlite_session, content, app_id="protected-app") - assert not storage.saved - - -def test_overwrite_replaces_inline_agents_and_reports_previous_owners( - sqlite_session: Session, - storage: _MemoryStorage, - monkeypatch: pytest.MonkeyPatch, -) -> None: - app = _source(sqlite_session, storage) - with AppPackageService().export_app(app_model=app) as exported: - content = exported.archive.read() - original = _import(sqlite_session, content) - assert original.status == ImportStatus.COMPLETED, original.error - old_agents = set( - sqlite_session.scalars( - select(WorkflowAgentNodeBinding.agent_id).where(WorkflowAgentNodeBinding.app_id == original.app_id) - ) - ) - sqlite_session.commit() - retirement = Mock() - monkeypatch.setattr("services.app_dsl_service.WorkflowAgentRetirementService.retire_unowned", retirement) - updated = _import(sqlite_session, content, app_id=original.app_id) - assert updated.status == ImportStatus.COMPLETED, updated.error - assert updated.app_id == original.app_id - new_agents = set( - sqlite_session.scalars( - select(WorkflowAgentNodeBinding.agent_id).where(WorkflowAgentNodeBinding.app_id == original.app_id) - ) - ) - assert len(new_agents) == len(old_agents) == 2 - assert not (new_agents & old_agents) - retirement.assert_called_once_with(tenant_id="tenant-2", agent_ids=old_agents, account_id="account-1") - - -@pytest.mark.parametrize("legacy_skill", [False, True]) -def test_resource_io_runs_without_database_transactions( - sqlite_session: Session, - storage: _MemoryStorage, - monkeypatch: pytest.MonkeyPatch, - legacy_skill: bool, -) -> None: - app = _source(sqlite_session, storage) - if legacy_skill: - version = sqlite_session.get(SkillVersion, "source-skill-version") - assert version is not None - version.manifest = version.manifest.model_copy(update={"name": None}) - sqlite_session.commit() - engine = sqlite_session.get_bind() - assert isinstance(engine, Engine) - pool = engine.pool - assert isinstance(pool, QueuePool) - - def before_io() -> None: - assert not sqlite_session.in_transaction() - assert pool.checkedout() == 0 - - def load_legacy_skill(*, tenant_id: str, file_id: str) -> bytes: - assert tenant_id == "tenant-1" - assert file_id - before_io() - return storage.files["tools/workspace.zip"] - - monkeypatch.setattr( - "services.skill_management_service.SkillManagementService._load_tool_file_bytes", - staticmethod(load_legacy_skill), - ) - load_stream = storage.load_stream - save = storage.save - - def checked_load(filename: str) -> Generator[bytes]: - before_io() - yield from load_stream(filename) - - def checked_save(filename: str, data: bytes) -> None: - before_io() - save(filename, data) - - monkeypatch.setattr(storage, "load_stream", checked_load) - monkeypatch.setattr(storage, "save", checked_save) - with AppPackageService().export_app(app_model=app) as exported: - content = exported.archive.read() - result = _import(sqlite_session, content) - assert result.status == ImportStatus.COMPLETED, result.error - - -def test_missing_assets_remain_explicit_import_warnings(sqlite_session: Session, storage: _MemoryStorage) -> None: - app = _source(sqlite_session, storage) - snapshot = sqlite_session.get(AgentConfigSnapshot, "source-snapshot-1") - assert snapshot is not None - soul = AgentSoulConfig.model_validate(snapshot.config_snapshot_dict) - soul.config_files[0].is_missing = True - soul.config_files[0].file_id = "" - snapshot.config_snapshot = soul - sqlite_session.commit() - with AppPackageService().export_app(app_model=app) as exported: - result = _import(sqlite_session, exported.archive.read()) - assert result.status == ImportStatus.COMPLETED_WITH_WARNINGS, result.error - assert any(warning.code == "agent_file_omitted" for warning in result.warnings) diff --git a/api/tests/unit_tests/services/agent/test_workflow_publish_service.py b/api/tests/unit_tests/services/agent/test_workflow_publish_service.py index 9265c920aa9..804395e5905 100644 --- a/api/tests/unit_tests/services/agent/test_workflow_publish_service.py +++ b/api/tests/unit_tests/services/agent/test_workflow_publish_service.py @@ -20,7 +20,7 @@ from models.model import App, AppMode from models.workflow import Workflow from services.agent.dsl_service import AgentDslService from services.agent.workflow_publish_service import WorkflowAgentPublishService -from tests.unit_tests.model_factories import make_workflow +from tests.helpers.unit.model_factories import make_workflow def _workflow(*, workflow_id: str = "workflow-1", version: str = Workflow.VERSION_DRAFT) -> Workflow: diff --git a/api/tests/unit_tests/services/agent/test_workflow_tool_agent_statistics.py b/api/tests/unit_tests/services/agent/test_workflow_tool_agent_statistics.py deleted file mode 100644 index 8526b73ab4f..00000000000 --- a/api/tests/unit_tests/services/agent/test_workflow_tool_agent_statistics.py +++ /dev/null @@ -1,215 +0,0 @@ -"""Real SQL coverage for Agent statistics inside a Workflow Tool invocation.""" - -import json -from datetime import datetime -from decimal import Decimal -from sqlite3 import Connection - -import pytest -from sqlalchemy import select -from sqlalchemy.orm import Session - -from core.workflow.node_execution_process_data import WORKFLOW_AGENT_BINDING_ID_KEY -from models.enums import ConversationFromSource, InvokeFrom -from models.model import App, AppMode, Message -from models.workflow import WorkflowNodeExecutionModel, WorkflowNodeExecutionTriggeredFrom, WorkflowRun, WorkflowType -from services.agent import observability_service as observability_service_module -from services.agent.observability_service import AgentObservabilityService, AgentStatisticsQueryParams -from tests.unit_tests.config_override import apply_config_overrides -from tests.unit_tests.services.agent.test_agent_observability_service import ( - _app, - _conversation, - _message, - _node_execution, - _workflow_binding, - _workflow_run, -) - - -@pytest.fixture -def statistics_service(monkeypatch: pytest.MonkeyPatch, sqlite_session: Session) -> AgentObservabilityService: - # Keep the production joins/aggregation; adapt only MySQL functions to SQLite. - apply_config_overrides(monkeypatch, DB_TYPE="mysql") - monkeypatch.setattr(observability_service_module, "convert_datetime_to_date", lambda field: f"DATE({field})") - connection = sqlite_session.connection().connection.dbapi_connection - assert isinstance(connection, Connection) - connection.create_function("JSON_UNQUOTE", 1, lambda value: value) - return AgentObservabilityService(sqlite_session) - - -def _seed_nested_agent(sqlite_session: Session, *, workflow_type: WorkflowType) -> App: - agent_app = _app(app_id="agent-app") - outer_app = _app( - app_id="outer-app", mode=AppMode.ADVANCED_CHAT if workflow_type == WorkflowType.CHAT else AppMode.WORKFLOW - ) - source_app = _app(app_id="source-app", mode=AppMode.WORKFLOW) - run = _workflow_run(workflow_type=workflow_type) - run.app_id = outer_app.id - run.workflow_id = "outer-workflow" - run.version = "outer-v1" - source_binding = _workflow_binding(app_id=source_app.id, binding_id="source-binding-v1", node_id="agent-node") - source_binding.workflow_id = "source-workflow" - source_binding.workflow_version = "source-v1" - unused_binding = _workflow_binding(app_id=source_app.id, binding_id="source-binding-v2", node_id="agent-node") - unused_binding.workflow_id = source_binding.workflow_id - unused_binding.workflow_version = "source-v2" - outer_binding = _workflow_binding(app_id=outer_app.id, node_id="agent-node") - outer_binding.workflow_id = run.workflow_id - outer_binding.workflow_version = run.version - other_agent_binding = _workflow_binding(app_id=source_app.id, node_id="other-agent-node") - other_agent_binding.workflow_id = source_binding.workflow_id - other_agent_binding.workflow_version = source_binding.workflow_version - other_agent_binding.agent_id = "other-agent" - sqlite_session.add_all( - [agent_app, outer_app, source_app, run, source_binding, unused_binding, outer_binding, other_agent_binding] - ) - - for index, (binding, tokens, completion_tokens, price, latency) in enumerate( - [ - (source_binding, 7, 2, "0.01", 2.0), - (source_binding, 11, 4, "0.02", 3.0), - (other_agent_binding, 900, 800, "9.00", 90.0), - ] - ): - execution = _node_execution(execution_id=f"tool-agent-execution-{index}") - execution.app_id = source_app.id - execution.workflow_id = source_binding.workflow_id - execution.node_id = binding.node_id - execution.triggered_from = WorkflowNodeExecutionTriggeredFrom.WORKFLOW_TOOL - execution.process_data = json.dumps({WORKFLOW_AGENT_BINDING_ID_KEY: binding.id}) - execution.elapsed_time = latency - execution.execution_metadata = json.dumps( - { - "agent_log": { - "agent_backend": { - "usage": {"total_tokens": tokens, "completion_tokens": completion_tokens, "total_price": price} - } - } - } - ) - sqlite_session.add(execution) - - if workflow_type == WorkflowType.CHAT: - conversation = _conversation(app_id=outer_app.id) - conversation.mode = AppMode.ADVANCED_CHAT - conversation.from_source = ConversationFromSource.API - message = _message(app_id=outer_app.id, created_at=datetime(2026, 7, 22, 8, 0)) - message.workflow_run_id = run.id - message.from_source = ConversationFromSource.API - message.invoke_from = InvokeFrom.WEB_APP - message.from_account_id = None - message.from_end_user_id = "end-user-1" - sqlite_session.add_all([conversation, message]) - sqlite_session.commit() - return agent_app - - -@pytest.mark.parametrize("workflow_type", [WorkflowType.WORKFLOW, WorkflowType.CHAT]) -def test_nested_agent_daily_statistics_preserve_source_usage_and_date( - sqlite_session: Session, statistics_service: AgentObservabilityService, workflow_type: WorkflowType -) -> None: - app = _seed_nested_agent(sqlite_session, workflow_type=workflow_type) - # Workflow metrics remain grouped by the outer run's date; Chatflow metrics - # retain the original message's date and usage. Child nodes executed July 23. - day = 22 if workflow_type == WorkflowType.CHAT else 21 - date = f"2026-07-{day}" - expected_tokens = 7 if workflow_type == WorkflowType.CHAT else 18 - expected_price = Decimal("0.0001") if workflow_type == WorkflowType.CHAT else Decimal("0.03") - expected_latency_ms = 1250.0 if workflow_type == WorkflowType.CHAT else 5000.0 - expected_tps = 3.2 if workflow_type == WorkflowType.CHAT else 1.2 - - for source in ("workflow:source-app", "workflow:source-app:source-workflow:source-v1:agent-node", "workflow"): - payload = statistics_service.get_statistics_summary( - app=app, - agent_id="agent-1", - params=AgentStatisticsQueryParams( - source=source, - start=datetime(2026, 7, day), - end=datetime(2026, 7, day + 1), - ), - ) - - summary = payload["summary"] - assert summary["total_messages"] == 1 - assert summary["total_conversations"] == 1 - assert summary["total_end_users"] == 1 - assert summary["total_tokens"] == expected_tokens - assert Decimal(summary["total_price"]) == expected_price - assert summary["average_response_time"] == expected_latency_ms - assert summary["tokens_per_second"] == expected_tps - assert payload["charts"]["daily_messages"] == [{"date": date, "message_count": 1}] - assert payload["charts"]["average_response_time"] == [{"date": date, "latency": expected_latency_ms}] - - for unrelated_source in ( - "workflow:outer-app", - "workflow:source-app:source-workflow:source-v2:agent-node", - "workflow:source-app:source-workflow:source-v1:other-agent-node", - ): - payload = statistics_service.get_statistics_summary( - app=app, agent_id="agent-1", params=AgentStatisticsQueryParams(source=unrelated_source) - ) - assert payload["summary"]["total_messages"] == 0 - assert payload["summary"]["total_tokens"] == 0 - assert payload["charts"]["daily_messages"] == [] - - -@pytest.mark.parametrize("workflow_type", [WorkflowType.WORKFLOW, WorkflowType.CHAT]) -@pytest.mark.parametrize("invalid_owner", ["outer_tenant", "source_tenant", "node_app", "node_workflow", "binding_id"]) -def test_nested_agent_statistics_reject_mismatched_ownership( - sqlite_session: Session, - statistics_service: AgentObservabilityService, - workflow_type: WorkflowType, - invalid_owner: str, -) -> None: - app = _seed_nested_agent(sqlite_session, workflow_type=workflow_type) - if invalid_owner == "outer_tenant": - run = sqlite_session.get(WorkflowRun, "workflow-run-1") - assert run is not None - run.tenant_id = "other-tenant" - elif invalid_owner == "source_tenant": - source_app = sqlite_session.get(App, "source-app") - assert source_app is not None - source_app.tenant_id = "other-tenant" - else: - executions = sqlite_session.scalars( - select(WorkflowNodeExecutionModel).where( - WorkflowNodeExecutionModel.id.in_(["tool-agent-execution-0", "tool-agent-execution-1"]) - ) - ).all() - assert len(executions) == 2 - for execution in executions: - if invalid_owner == "node_app": - execution.app_id = "other-app" - elif invalid_owner == "node_workflow": - execution.workflow_id = "other-workflow" - else: - execution.process_data = json.dumps( - {WORKFLOW_AGENT_BINDING_ID_KEY: "binding-source-app-other-agent-node"} - ) - sqlite_session.commit() - - payload = statistics_service.get_statistics_summary( - app=app, agent_id="agent-1", params=AgentStatisticsQueryParams(source="workflow:source-app") - ) - - assert payload["summary"]["total_messages"] == 0 - assert payload["summary"]["total_tokens"] == 0 - assert payload["charts"]["daily_messages"] == [] - - -def test_nested_agent_chat_statistics_reject_message_from_another_app( - sqlite_session: Session, statistics_service: AgentObservabilityService -) -> None: - app = _seed_nested_agent(sqlite_session, workflow_type=WorkflowType.CHAT) - message = sqlite_session.get(Message, "message-1") - assert message is not None - message.app_id = "other-app" - sqlite_session.commit() - - payload = statistics_service.get_statistics_summary( - app=app, agent_id="agent-1", params=AgentStatisticsQueryParams(source="workflow:source-app") - ) - - assert payload["summary"]["total_messages"] == 0 - assert payload["summary"]["total_tokens"] == 0 - assert payload["charts"]["daily_messages"] == [] diff --git a/api/tests/unit_tests/services/agent/test_workspace_service.py b/api/tests/unit_tests/services/agent/test_workspace_service.py index 8310b331500..ccb29668c4e 100644 --- a/api/tests/unit_tests/services/agent/test_workspace_service.py +++ b/api/tests/unit_tests/services/agent/test_workspace_service.py @@ -21,7 +21,7 @@ from services.agent.workspace_service import ( AgentWorkspaceService, WorkspaceOwnerScope, ) -from tests.unit_tests.config_override import apply_config_overrides +from tests.helpers.unit.config_override import apply_config_overrides def _scope() -> WorkspaceOwnerScope: diff --git a/api/tests/unit_tests/controllers/common/test_app_access.py b/api/tests/unit_tests/services/app/test_access.py similarity index 93% rename from api/tests/unit_tests/controllers/common/test_app_access.py rename to api/tests/unit_tests/services/app/test_access.py index 5893c1f1865..3cc9ef9078b 100644 --- a/api/tests/unit_tests/controllers/common/test_app_access.py +++ b/api/tests/unit_tests/services/app/test_access.py @@ -19,6 +19,12 @@ from services.rbac.contracts import ( ResourceWhitelistResources, WorkspacePermissionSnapshot, ) +from tests.unit_tests.controllers.conftest import app_query_services as app_query_services # noqa: PLC0414 +from tests.unit_tests.controllers.conftest import ( + datasource_application_dependencies as datasource_application_dependencies, # noqa: PLC0414 +) +from tests.unit_tests.controllers.conftest import error_response_contract as error_response_contract # noqa: PLC0414 +from tests.unit_tests.controllers.conftest import pipeline_application as pipeline_application # noqa: PLC0414 _RBAC_MODULE = "services.app.access.enterprise_rbac_service" diff --git a/api/tests/unit_tests/services/app/test_agent_app_feature_service.py b/api/tests/unit_tests/services/app/test_agent_app_feature_gateway.py similarity index 100% rename from api/tests/unit_tests/services/app/test_agent_app_feature_service.py rename to api/tests/unit_tests/services/app/test_agent_app_feature_gateway.py diff --git a/api/tests/unit_tests/services/app/test_agent_app_sandbox_service.py b/api/tests/unit_tests/services/app/test_agent_app_sandbox_service.py index 7393afb23b8..855635038f1 100644 --- a/api/tests/unit_tests/services/app/test_agent_app_sandbox_service.py +++ b/api/tests/unit_tests/services/app/test_agent_app_sandbox_service.py @@ -53,8 +53,8 @@ from services.app.agent_app_contracts import ( from services.app.agent_app_sandbox_service import AgentAppSandboxService from services.app.agent_sandbox_file_gateway import AgentSandboxFileGateway, create_sandbox_client from services.file_request_service import DownloadFileRequestResult, FileRequestService -from tests.unit_tests.config_override import apply_config_overrides -from tests.unit_tests.model_factories import make_message +from tests.helpers.unit.config_override import apply_config_overrides +from tests.helpers.unit.model_factories import make_message CONTEXT = RequestContext("request-1", "trace-1", "account-1", "tenant-1") diff --git a/api/tests/unit_tests/services/app/test_agent_app_access_service.py b/api/tests/unit_tests/services/app/test_agent_app_service.py similarity index 98% rename from api/tests/unit_tests/services/app/test_agent_app_access_service.py rename to api/tests/unit_tests/services/app/test_agent_app_service.py index 49167ca07cd..dde98971924 100644 --- a/api/tests/unit_tests/services/app/test_agent_app_access_service.py +++ b/api/tests/unit_tests/services/app/test_agent_app_service.py @@ -17,7 +17,7 @@ from models.agent import ( from models.agent_config_entities import WorkflowNodeJobConfig from models.model import AppMode from services.app.agent_app_contracts import AgentAppNotFoundError -from tests.unit_tests.model_factories import make_app +from tests.helpers.unit.model_factories import make_app def test_access_only_lists_current_published_workflow_references( diff --git a/api/tests/unit_tests/services/app/test_console_gateway.py b/api/tests/unit_tests/services/app/test_console_gateway.py index 2f95b3dcd38..b6017c47379 100644 --- a/api/tests/unit_tests/services/app/test_console_gateway.py +++ b/api/tests/unit_tests/services/app/test_console_gateway.py @@ -63,9 +63,9 @@ from services.plugin.dependencies_analysis import DependenciesAnalysisService from services.recommended_app_package_service import RecommendedAppPackageService from services.system_feature_service import SystemFeatureService from services.workflow_service import WorkflowService -from tests.unit_tests.audio_runtime_fixtures import record_calls -from tests.unit_tests.model_factories import make_account, make_tenant, make_upload_file -from tests.unit_tests.rbac_fakes import RBACDomain +from tests.helpers.unit.audio_runtime_fixtures import record_calls +from tests.helpers.unit.model_factories import make_account, make_tenant, make_upload_file +from tests.helpers.unit.rbac_fakes import RBACDomain @pytest.mark.parametrize(("mode", "include_inline"), [(AppMode.WORKFLOW, False), (AppMode.ADVANCED_CHAT, True)]) diff --git a/api/tests/unit_tests/services/data_migration/test_import_service.py b/api/tests/unit_tests/services/data_migration/test_import_service.py index 28e9d81d19a..5cc2c96e1f6 100644 --- a/api/tests/unit_tests/services/data_migration/test_import_service.py +++ b/api/tests/unit_tests/services/data_migration/test_import_service.py @@ -28,7 +28,7 @@ from services.data_migration.entities import ( ) from services.data_migration.import_service import ImportRequest, ImportTargetResolver, MigrationImportService from services.entities.dsl_entities import ImportStatus -from tests.unit_tests.config_override import apply_config_overrides +from tests.helpers.unit.config_override import apply_config_overrides @dataclass(frozen=True) diff --git a/api/tests/unit_tests/services/auth/test_firecrawl_auth.py b/api/tests/unit_tests/services/data_source/auth/firecrawl/test_firecrawl.py similarity index 100% rename from api/tests/unit_tests/services/auth/test_firecrawl_auth.py rename to api/tests/unit_tests/services/data_source/auth/firecrawl/test_firecrawl.py diff --git a/api/tests/unit_tests/services/auth/test_jina_auth.py b/api/tests/unit_tests/services/data_source/auth/jina/test_jina.py similarity index 100% rename from api/tests/unit_tests/services/auth/test_jina_auth.py rename to api/tests/unit_tests/services/data_source/auth/jina/test_jina.py diff --git a/api/tests/unit_tests/services/auth/test_data_source_api_key_auth_gateways.py b/api/tests/unit_tests/services/data_source/auth/test_api_key_gateways.py similarity index 100% rename from api/tests/unit_tests/services/auth/test_data_source_api_key_auth_gateways.py rename to api/tests/unit_tests/services/data_source/auth/test_api_key_gateways.py diff --git a/api/tests/unit_tests/services/auth/test_data_source_api_key_auth_service.py b/api/tests/unit_tests/services/data_source/auth/test_api_key_service.py similarity index 100% rename from api/tests/unit_tests/services/auth/test_data_source_api_key_auth_service.py rename to api/tests/unit_tests/services/data_source/auth/test_api_key_service.py diff --git a/api/tests/unit_tests/services/auth/test_watercrawl_auth.py b/api/tests/unit_tests/services/data_source/auth/watercrawl/test_watercrawl.py similarity index 100% rename from api/tests/unit_tests/services/auth/test_watercrawl_auth.py rename to api/tests/unit_tests/services/data_source/auth/watercrawl/test_watercrawl.py diff --git a/api/tests/unit_tests/services/data_source/test_credential_metadata_service.py b/api/tests/unit_tests/services/data_source/test_credential_metadata_service.py index 576ff74160b..caa55507081 100644 --- a/api/tests/unit_tests/services/data_source/test_credential_metadata_service.py +++ b/api/tests/unit_tests/services/data_source/test_credential_metadata_service.py @@ -2,7 +2,7 @@ from dataclasses import asdict from services.credentials.query import DatasourceCredentialListItem from services.data_source.credential_metadata_service import DatasourceCredentialMetadataService -from tests.unit_tests.credential_query_fakes import DatasourceCredentialLookup, DatasourceCredentialQuery +from tests.helpers.unit.credential_query_fakes import DatasourceCredentialLookup, DatasourceCredentialQuery def _credential(*, credential_id: str, name: str, is_default: bool) -> DatasourceCredentialListItem: diff --git a/api/tests/unit_tests/services/data_source/test_notion_import_adapters.py b/api/tests/unit_tests/services/data_source/test_notion_import_adapters.py index 34ec6b12353..8bdb88f0bf8 100644 --- a/api/tests/unit_tests/services/data_source/test_notion_import_adapters.py +++ b/api/tests/unit_tests/services/data_source/test_notion_import_adapters.py @@ -27,7 +27,7 @@ from services.data_source.credential_gateway import ( from services.data_source.entities.notion_import import NotionPageType from services.data_source.notion_import_adapters import PluginNotionSourceGateway from services.data_source.notion_import_application_service import NotionImportCredentialUnavailableError -from tests.unit_tests.core.datasource.factories import datasource_entity, provider_entity +from tests.helpers.unit.core.datasource.factories import datasource_entity, provider_entity type CredentialCalls = list[tuple[str, str, str, str, str]] diff --git a/api/tests/unit_tests/services/test_datasource_provider_service.py b/api/tests/unit_tests/services/data_source/test_provider_service.py similarity index 100% rename from api/tests/unit_tests/services/test_datasource_provider_service.py rename to api/tests/unit_tests/services/data_source/test_provider_service.py diff --git a/api/tests/unit_tests/services/test_website_service.py b/api/tests/unit_tests/services/data_source/test_website_service.py similarity index 100% rename from api/tests/unit_tests/services/test_website_service.py rename to api/tests/unit_tests/services/data_source/test_website_service.py diff --git a/api/tests/unit_tests/services/dataset_service/test_save_document_with_dataset_id_mysql_snapshot.py b/api/tests/unit_tests/services/dataset_service/test_save_document_with_dataset_id_mysql_snapshot.py deleted file mode 100644 index 62b120a0d31..00000000000 --- a/api/tests/unit_tests/services/dataset_service/test_save_document_with_dataset_id_mysql_snapshot.py +++ /dev/null @@ -1,70 +0,0 @@ -"""Regression coverage for #41735. - -The pre-fix code in ``DocumentService.save_document_with_dataset_id`` -read ``UploadFile`` rows inside a single session that had its -``REPEATABLE READ`` snapshot pinned by an earlier ``Dataset`` SELECT. -When ``FileService.upload_text`` / ``upload_file`` committed the -``upload_files`` row in a SEPARATE session, the row was invisible to -*this* session until the snapshot was refreshed, and -``save_document_with_dataset_id`` reported ``One or more files not -found.`` even though the row was committed (verified via a second -session). - -The fix forces ``session.commit()`` immediately before the -``UploadFile`` lookup so the next SELECT starts a fresh snapshot. -This file pins that contract by reading the function source and -asserting the structural invariant directly. -""" - -from __future__ import annotations - -import inspect - -from services.knowledge.dataset_service import DocumentService - - -def test_session_committed_before_upload_file_lookup() -> None: - """#41735: ``session.commit()`` must run BEFORE the ``UploadFile`` - SELECT so the SELECT sees rows committed by ``FileService.upload_text`` - on a different session. MySQL REPEATABLE READ keeps the original - snapshot for the rest of the transaction; the commit ends it so the - next SELECT starts a fresh transaction with a new snapshot. - PostgreSQL's default READ COMMITTED rebuilds the read view per - statement so the extra commit is a no-op there. - - Read the function source and assert that ``session.commit()`` is - called BEFORE the ``UploadFile`` SELECT — the same structural - invariant the reproduction in the issue depends on, without needing - a live MySQL database. - """ - source_lines = inspect.getsource(DocumentService.save_prepared_documents).splitlines() - - upload_select_line = next( - (i for i, line in enumerate(source_lines) if "select(UploadFile)" in line), - None, - ) - assert upload_select_line is not None, "save_prepared_documents must contain a select(UploadFile) call" - - # Find every ``session.commit()`` call, then keep the last one - # whose line number is BEFORE the UploadFile SELECT. If the fix is - # in place, that commit is the snapshot refresh. - last_commit_before_upload = -1 - for i, line in enumerate(source_lines): - if i >= upload_select_line: - break - if "session.commit()" in line: - last_commit_before_upload = i - - assert last_commit_before_upload >= 0, ( - "save_prepared_documents must call session.commit() at least once before the UploadFile SELECT" - ) - - # The fix is a focused call: one extra ``session.commit()`` between - # the (already-existing) ``DocumentService.check_doc_form`` call and - # the ``select(UploadFile)`` call. Pin that there's a comment - # referencing #41735 right above the new commit so a future edit - # doesn't accidentally delete the snapshot refresh without - # understanding why it's there. - assert "#41735" in "\n".join(source_lines), ( - "The fix's invariant comment referencing issue #41735 must be present in save_prepared_documents" - ) diff --git a/api/tests/unit_tests/services/test_batch_indexing_base.py b/api/tests/unit_tests/services/document_indexing_proxy/test_batch_indexing_base.py similarity index 99% rename from api/tests/unit_tests/services/test_batch_indexing_base.py rename to api/tests/unit_tests/services/document_indexing_proxy/test_batch_indexing_base.py index a23c2576254..0798613471e 100644 --- a/api/tests/unit_tests/services/test_batch_indexing_base.py +++ b/api/tests/unit_tests/services/document_indexing_proxy/test_batch_indexing_base.py @@ -7,7 +7,7 @@ import pytest from core.entities.document_task import DocumentTask from enums import CloudPlan, DeploymentEdition from services.document_indexing_proxy.batch_indexing_base import BatchDocumentIndexingProxy -from tests.unit_tests.config_override import config_overrides_context +from tests.helpers.unit.config_override import config_overrides_context # --------------------------------------------------------------------------- # Concrete subclass for testing (the base class is abstract) diff --git a/api/tests/unit_tests/services/test_document_indexing_task_proxy.py b/api/tests/unit_tests/services/document_indexing_proxy/test_document_indexing_task_proxy.py similarity index 99% rename from api/tests/unit_tests/services/test_document_indexing_task_proxy.py rename to api/tests/unit_tests/services/document_indexing_proxy/test_document_indexing_task_proxy.py index b3c9bb019f2..45c439f0604 100644 --- a/api/tests/unit_tests/services/test_document_indexing_task_proxy.py +++ b/api/tests/unit_tests/services/document_indexing_proxy/test_document_indexing_task_proxy.py @@ -5,7 +5,7 @@ from core.entities.document_task import DocumentTask from core.rag.pipeline.queue import TenantIsolatedTaskQueue from enums import CloudPlan, DeploymentEdition from services.document_indexing_proxy.document_indexing_task_proxy import DocumentIndexingTaskProxy -from tests.unit_tests.config_override import config_overrides_context +from tests.helpers.unit.config_override import config_overrides_context class DocumentIndexingTaskProxyTestDataFactory: diff --git a/api/tests/unit_tests/services/test_duplicate_document_indexing_task_proxy.py b/api/tests/unit_tests/services/document_indexing_proxy/test_duplicate_document_indexing_task_proxy.py similarity index 99% rename from api/tests/unit_tests/services/test_duplicate_document_indexing_task_proxy.py rename to api/tests/unit_tests/services/document_indexing_proxy/test_duplicate_document_indexing_task_proxy.py index c3a8979c19d..4456e0a7ab6 100644 --- a/api/tests/unit_tests/services/test_duplicate_document_indexing_task_proxy.py +++ b/api/tests/unit_tests/services/document_indexing_proxy/test_duplicate_document_indexing_task_proxy.py @@ -7,7 +7,7 @@ from enums import CloudPlan, DeploymentEdition from services.document_indexing_proxy.duplicate_document_indexing_task_proxy import ( DuplicateDocumentIndexingTaskProxy, ) -from tests.unit_tests.config_override import config_overrides_context +from tests.helpers.unit.config_override import config_overrides_context class DuplicateDocumentIndexingTaskProxyTestDataFactory: diff --git a/api/tests/unit_tests/services/enterprise/test_traceparent_propagation.py b/api/tests/unit_tests/services/enterprise/test_base.py similarity index 100% rename from api/tests/unit_tests/services/enterprise/test_traceparent_propagation.py rename to api/tests/unit_tests/services/enterprise/test_base.py diff --git a/api/tests/unit_tests/services/enterprise/test_rbac_agent_access.py b/api/tests/unit_tests/services/enterprise/test_rbac_agent_access.py deleted file mode 100644 index 462ce742a09..00000000000 --- a/api/tests/unit_tests/services/enterprise/test_rbac_agent_access.py +++ /dev/null @@ -1,248 +0,0 @@ -"""Unit tests for the agent-flavoured RBAC inner-API client. - -The resource-access clients injected into the application services share one -implementation across apps, datasets and agents. -These tests replace the Enterprise transport and assert the HTTP method, the exact endpoint, -the query/body keys and the returned model for every operation. -""" - -from __future__ import annotations - -from collections.abc import Callable - -import pytest - -from services.enterprise import rbac_service as svc -from services.enterprise.base import EnterpriseRequest -from services.rbac import contracts as rbac_contracts -from tests.unit_tests.rbac_fakes import RBACTransport, RecordedRequest - -TENANT = "tenant-1" -ACTOR = "acct-1" -AGENT = "agent-1" -POLICY = "policy-1" - - -@pytest.fixture -def transport(monkeypatch: pytest.MonkeyPatch) -> RBACTransport: - transport = RBACTransport(response={"automatic_include_workspace_members": True}) - monkeypatch.setattr(EnterpriseRequest, "send_inner_rbac_request", transport.send) - return transport - - -def _last(transport: RBACTransport) -> RecordedRequest: - return transport.only_request - - -_AGENT_CASES: list[tuple[str, str, str, Callable[[], object]]] = [ - ( - "whitelist_resources", - "GET", - "/rbac/agents/whitelist/resources", - lambda: svc._AGENT_ACCESS.whitelist_resources(TENANT, ACTOR), - ), - ( - "user_access_policies", - "GET", - "/rbac/agents/user-access-policies", - lambda: svc._AGENT_ACCESS.user_access_policies(TENANT, ACTOR, resource_id=AGENT), - ), - ( - "replace_user_access_policies", - "PUT", - "/rbac/agents/user-access-policies", - lambda: svc._AGENT_ACCESS.replace_user_access_policies( - TENANT, - ACTOR, - resource_id=AGENT, - target_account_id="member-1", - payload=rbac_contracts.ReplaceUserAccessPolicies(access_policy_ids=[POLICY], account_ids=["member-1"]), - ), - ), - ( - "whitelist", - "GET", - "/rbac/agents/whitelist", - lambda: svc._AGENT_ACCESS.whitelist(TENANT, ACTOR, resource_id=AGENT), - ), - ( - "whitelist_config", - "GET", - "/rbac/agents/whitelist", - lambda: svc._AGENT_ACCESS.whitelist_config(TENANT, ACTOR, resource_id=AGENT), - ), - ( - "legacy_whitelist_config", - "GET", - "/rbac/agents/whitelist", - lambda: svc._AGENT_ACCESS.legacy_whitelist_config(TENANT, ACTOR, resource_id=AGENT), - ), - ( - "replace_whitelist", - "PUT", - "/rbac/agents/whitelist", - lambda: svc._AGENT_ACCESS.replace_whitelist( - TENANT, - ACTOR, - resource_id=AGENT, - payload=rbac_contracts.ReplaceMemberBindings(automatic_include_workspace_members=True), - ), - ), - ( - "append_whitelist_members_batch", - "POST", - "/rbac/agents/whitelist/members/batch", - lambda: svc._AGENT_ACCESS.append_whitelist_members_batch( - tenant_id=TENANT, - account_id=ACTOR, - data=[ - rbac_contracts.AppendAgentWhitelistMembersBatchItem( - agent_id=AGENT, account_ids=["member-1"], policy_id=POLICY - ) - ], - ), - ), - ( - "matrix", - "GET", - "/rbac/agents/access-policy", - lambda: svc._AGENT_ACCESS.matrix(TENANT, ACTOR, resource_id=AGENT), - ), - ( - "list_role_bindings", - "GET", - "/rbac/agents/access-policy/role-bindings", - lambda: svc._AGENT_ACCESS.list_role_bindings(TENANT, ACTOR, resource_id=AGENT, policy_id=POLICY), - ), - ( - "list_member_bindings", - "GET", - "/rbac/agents/access-policy/member-bindings", - lambda: svc._AGENT_ACCESS.list_member_bindings(TENANT, ACTOR, resource_id=AGENT, policy_id=POLICY), - ), - ( - "delete_member_bindings", - "DELETE", - "/rbac/agents/access-policy/member-bindings", - lambda: svc._AGENT_ACCESS.delete_member_bindings( - TENANT, - ACTOR, - resource_id=AGENT, - policy_id=POLICY, - payload=rbac_contracts.DeleteMemberBindings(account_ids=["member-1"]), - ), - ), - ( - "workspace.agent_matrix", - "GET", - "/rbac/workspace/agents/access-policy", - lambda: svc._WORKSPACE_AGENT_ACCESS.matrix(TENANT, ACTOR), - ), - ( - "workspace.list_agent_role_bindings", - "GET", - "/rbac/workspace/agents/access-policy/role-bindings", - lambda: svc._WORKSPACE_AGENT_ACCESS.list_role_bindings(TENANT, ACTOR, POLICY), - ), - ( - "workspace.list_agent_member_bindings", - "GET", - "/rbac/workspace/agents/access-policy/member-bindings", - lambda: svc._WORKSPACE_AGENT_ACCESS.list_member_bindings(TENANT, ACTOR, POLICY), - ), - ( - "workspace.replace_agent_bindings", - "PUT", - "/rbac/workspace/agents/access-policy/bindings", - lambda: svc._WORKSPACE_AGENT_ACCESS.replace_bindings( - TENANT, ACTOR, POLICY, rbac_contracts.ReplaceBindings(role_ids=["role-1"], account_ids=["member-1"]) - ), - ), - ( - "catalog.agent", - "GET", - "/rbac/role-permissions/catalog/agent", - lambda: svc.RBACService.Catalog.agent(TENANT, account_id=ACTOR), - ), -] - - -@pytest.mark.parametrize( - ("method", "endpoint", "invoke"), - [(case[1], case[2], case[3]) for case in _AGENT_CASES], - ids=[case[0] for case in _AGENT_CASES], -) -def test_agent_operations_hit_the_agent_route( - transport: RBACTransport, method: str, endpoint: str, invoke: Callable[[], object] -) -> None: - invoke() - - call = _last(transport) - assert call.method == method - assert call.endpoint == endpoint - assert "/apps/" not in call.endpoint - assert "/datasets/" not in call.endpoint - assert call.tenant_id == TENANT - assert call.account_id == ACTOR - params = call.params or {} - assert "app_id" not in params - assert "dataset_id" not in params - if AGENT in params.values(): - assert params.get("agent_id") == AGENT - - -_PARITY_CASES: list[tuple[str, Callable[[], object], Callable[[], object]]] = [ - ( - "whitelist_resources", - lambda: svc._APP_ACCESS.whitelist_resources(TENANT, ACTOR), - lambda: svc._AGENT_ACCESS.whitelist_resources(TENANT, ACTOR), - ), - ( - "whitelist", - lambda: svc._APP_ACCESS.whitelist(TENANT, ACTOR, "res-1"), - lambda: svc._AGENT_ACCESS.whitelist(TENANT, ACTOR, resource_id="res-1"), - ), - ( - "user_access_policies", - lambda: svc._APP_ACCESS.user_access_policies(TENANT, ACTOR, "res-1"), - lambda: svc._AGENT_ACCESS.user_access_policies(TENANT, ACTOR, resource_id="res-1"), - ), - ( - "matrix", - lambda: svc._APP_ACCESS.matrix(TENANT, ACTOR, "res-1"), - lambda: svc._AGENT_ACCESS.matrix(TENANT, ACTOR, resource_id="res-1"), - ), - ( - "list_role_bindings", - lambda: svc._APP_ACCESS.list_role_bindings(TENANT, ACTOR, "res-1", POLICY), - lambda: svc._AGENT_ACCESS.list_role_bindings(TENANT, ACTOR, resource_id="res-1", policy_id=POLICY), - ), - ( - "list_member_bindings", - lambda: svc._APP_ACCESS.list_member_bindings(TENANT, ACTOR, "res-1", POLICY), - lambda: svc._AGENT_ACCESS.list_member_bindings(TENANT, ACTOR, resource_id="res-1", policy_id=POLICY), - ), -] - - -@pytest.mark.parametrize( - ("app_call", "agent_call"), - [(case[1], case[2]) for case in _PARITY_CASES], - ids=[case[0] for case in _PARITY_CASES], -) -def test_same_path_and_param_shape_modulo_segment( - transport: RBACTransport, - app_call: Callable[[], object], - agent_call: Callable[[], object], -) -> None: - app_call() - app = _last(transport) - transport.requests.clear() - agent_call() - agent = _last(transport) - - assert app.endpoint.replace("/apps/", "/agents/") == agent.endpoint - assert app.method == agent.method - - app_params = {("agent_id" if k == "app_id" else k): v for k, v in (app.params or {}).items()} - assert app_params == (agent.params or {}) diff --git a/api/tests/unit_tests/services/enterprise/test_rbac_service.py b/api/tests/unit_tests/services/enterprise/test_rbac_service.py index 109166935ca..bafff67ff05 100644 --- a/api/tests/unit_tests/services/enterprise/test_rbac_service.py +++ b/api/tests/unit_tests/services/enterprise/test_rbac_service.py @@ -1,11 +1,3 @@ -"""Unit tests for services.enterprise.rbac_service. - -Most enterprise RBAC methods turn a single ``EnterpriseRequest.send_inner_rbac_request`` -call into a pydantic response model. Rather than spinning up an HTTP server, these tests -replace that helper and assert on the request arguments and response shape. -Local membership tests live at the application service boundary. -""" - from __future__ import annotations from collections.abc import Callable @@ -18,7 +10,257 @@ from services.enterprise import rbac_service as svc from services.enterprise.base import EnterpriseRequest from services.knowledge_fs import membership_changes from services.rbac import contracts as rbac_contracts -from tests.unit_tests.rbac_fakes import RBACTransport, RecordedRequest +from tests.helpers.unit.rbac_fakes import RBACTransport, RecordedRequest + +"""Unit tests for the agent-flavoured RBAC inner-API client. + +The resource-access clients injected into the application services share one +implementation across apps, datasets and agents. +These tests replace the Enterprise transport and assert the HTTP method, the exact endpoint, +the query/body keys and the returned model for every operation. +""" + + +TENANT = "tenant-1" +ACTOR = "acct-1" +AGENT = "agent-1" +POLICY = "policy-1" + + +def _last(transport: RBACTransport) -> RecordedRequest: + return transport.only_request + + +_AGENT_CASES: list[tuple[str, str, str, Callable[[], object]]] = [ + ( + "whitelist_resources", + "GET", + "/rbac/agents/whitelist/resources", + lambda: svc._AGENT_ACCESS.whitelist_resources(TENANT, ACTOR), + ), + ( + "user_access_policies", + "GET", + "/rbac/agents/user-access-policies", + lambda: svc._AGENT_ACCESS.user_access_policies(TENANT, ACTOR, resource_id=AGENT), + ), + ( + "replace_user_access_policies", + "PUT", + "/rbac/agents/user-access-policies", + lambda: svc._AGENT_ACCESS.replace_user_access_policies( + TENANT, + ACTOR, + resource_id=AGENT, + target_account_id="member-1", + payload=rbac_contracts.ReplaceUserAccessPolicies(access_policy_ids=[POLICY], account_ids=["member-1"]), + ), + ), + ( + "whitelist", + "GET", + "/rbac/agents/whitelist", + lambda: svc._AGENT_ACCESS.whitelist(TENANT, ACTOR, resource_id=AGENT), + ), + ( + "whitelist_config", + "GET", + "/rbac/agents/whitelist", + lambda: svc._AGENT_ACCESS.whitelist_config(TENANT, ACTOR, resource_id=AGENT), + ), + ( + "legacy_whitelist_config", + "GET", + "/rbac/agents/whitelist", + lambda: svc._AGENT_ACCESS.legacy_whitelist_config(TENANT, ACTOR, resource_id=AGENT), + ), + ( + "replace_whitelist", + "PUT", + "/rbac/agents/whitelist", + lambda: svc._AGENT_ACCESS.replace_whitelist( + TENANT, + ACTOR, + resource_id=AGENT, + payload=rbac_contracts.ReplaceMemberBindings(automatic_include_workspace_members=True), + ), + ), + ( + "append_whitelist_members_batch", + "POST", + "/rbac/agents/whitelist/members/batch", + lambda: svc._AGENT_ACCESS.append_whitelist_members_batch( + tenant_id=TENANT, + account_id=ACTOR, + data=[ + rbac_contracts.AppendAgentWhitelistMembersBatchItem( + agent_id=AGENT, account_ids=["member-1"], policy_id=POLICY + ) + ], + ), + ), + ( + "matrix", + "GET", + "/rbac/agents/access-policy", + lambda: svc._AGENT_ACCESS.matrix(TENANT, ACTOR, resource_id=AGENT), + ), + ( + "list_role_bindings", + "GET", + "/rbac/agents/access-policy/role-bindings", + lambda: svc._AGENT_ACCESS.list_role_bindings(TENANT, ACTOR, resource_id=AGENT, policy_id=POLICY), + ), + ( + "list_member_bindings", + "GET", + "/rbac/agents/access-policy/member-bindings", + lambda: svc._AGENT_ACCESS.list_member_bindings(TENANT, ACTOR, resource_id=AGENT, policy_id=POLICY), + ), + ( + "delete_member_bindings", + "DELETE", + "/rbac/agents/access-policy/member-bindings", + lambda: svc._AGENT_ACCESS.delete_member_bindings( + TENANT, + ACTOR, + resource_id=AGENT, + policy_id=POLICY, + payload=rbac_contracts.DeleteMemberBindings(account_ids=["member-1"]), + ), + ), + ( + "workspace.agent_matrix", + "GET", + "/rbac/workspace/agents/access-policy", + lambda: svc._WORKSPACE_AGENT_ACCESS.matrix(TENANT, ACTOR), + ), + ( + "workspace.list_agent_role_bindings", + "GET", + "/rbac/workspace/agents/access-policy/role-bindings", + lambda: svc._WORKSPACE_AGENT_ACCESS.list_role_bindings(TENANT, ACTOR, POLICY), + ), + ( + "workspace.list_agent_member_bindings", + "GET", + "/rbac/workspace/agents/access-policy/member-bindings", + lambda: svc._WORKSPACE_AGENT_ACCESS.list_member_bindings(TENANT, ACTOR, POLICY), + ), + ( + "workspace.replace_agent_bindings", + "PUT", + "/rbac/workspace/agents/access-policy/bindings", + lambda: svc._WORKSPACE_AGENT_ACCESS.replace_bindings( + TENANT, ACTOR, POLICY, rbac_contracts.ReplaceBindings(role_ids=["role-1"], account_ids=["member-1"]) + ), + ), + ( + "catalog.agent", + "GET", + "/rbac/role-permissions/catalog/agent", + lambda: svc.RBACService.Catalog.agent(TENANT, account_id=ACTOR), + ), +] + + +_PARITY_CASES: list[tuple[str, Callable[[], object], Callable[[], object]]] = [ + ( + "whitelist_resources", + lambda: svc._APP_ACCESS.whitelist_resources(TENANT, ACTOR), + lambda: svc._AGENT_ACCESS.whitelist_resources(TENANT, ACTOR), + ), + ( + "whitelist", + lambda: svc._APP_ACCESS.whitelist(TENANT, ACTOR, "res-1"), + lambda: svc._AGENT_ACCESS.whitelist(TENANT, ACTOR, resource_id="res-1"), + ), + ( + "user_access_policies", + lambda: svc._APP_ACCESS.user_access_policies(TENANT, ACTOR, "res-1"), + lambda: svc._AGENT_ACCESS.user_access_policies(TENANT, ACTOR, resource_id="res-1"), + ), + ( + "matrix", + lambda: svc._APP_ACCESS.matrix(TENANT, ACTOR, "res-1"), + lambda: svc._AGENT_ACCESS.matrix(TENANT, ACTOR, resource_id="res-1"), + ), + ( + "list_role_bindings", + lambda: svc._APP_ACCESS.list_role_bindings(TENANT, ACTOR, "res-1", POLICY), + lambda: svc._AGENT_ACCESS.list_role_bindings(TENANT, ACTOR, resource_id="res-1", policy_id=POLICY), + ), + ( + "list_member_bindings", + lambda: svc._APP_ACCESS.list_member_bindings(TENANT, ACTOR, "res-1", POLICY), + lambda: svc._AGENT_ACCESS.list_member_bindings(TENANT, ACTOR, resource_id="res-1", policy_id=POLICY), + ), +] + + +class TestRbacAgentAccessCases: + @staticmethod + @pytest.fixture + def transport(monkeypatch: pytest.MonkeyPatch) -> RBACTransport: + transport = RBACTransport(response={"automatic_include_workspace_members": True}) + monkeypatch.setattr(EnterpriseRequest, "send_inner_rbac_request", transport.send) + return transport + + @staticmethod + @pytest.mark.parametrize( + ("method", "endpoint", "invoke"), + [(case[1], case[2], case[3]) for case in _AGENT_CASES], + ids=[case[0] for case in _AGENT_CASES], + ) + def test_agent_operations_hit_the_agent_route( + transport: RBACTransport, method: str, endpoint: str, invoke: Callable[[], object] + ) -> None: + invoke() + + call = _last(transport) + assert call.method == method + assert call.endpoint == endpoint + assert "/apps/" not in call.endpoint + assert "/datasets/" not in call.endpoint + assert call.tenant_id == TENANT + assert call.account_id == ACTOR + params = call.params or {} + assert "app_id" not in params + assert "dataset_id" not in params + if AGENT in params.values(): + assert params.get("agent_id") == AGENT + + @staticmethod + @pytest.mark.parametrize( + ("app_call", "agent_call"), + [(case[1], case[2]) for case in _PARITY_CASES], + ids=[case[0] for case in _PARITY_CASES], + ) + def test_same_path_and_param_shape_modulo_segment( + transport: RBACTransport, + app_call: Callable[[], object], + agent_call: Callable[[], object], + ) -> None: + app_call() + app = _last(transport) + transport.requests.clear() + agent_call() + agent = _last(transport) + + assert app.endpoint.replace("/apps/", "/agents/") == agent.endpoint + assert app.method == agent.method + + app_params = {("agent_id" if k == "app_id" else k): v for k, v in (app.params or {}).items()} + assert app_params == (agent.params or {}) + + +"""Unit tests for services.enterprise.rbac_service. + +Most enterprise RBAC methods turn a single ``EnterpriseRequest.send_inner_rbac_request`` +call into a pydantic response model. Rather than spinning up an HTTP server, these tests +replace that helper and assert on the request arguments and response shape. +Local membership tests live at the application service boundary. +""" @pytest.fixture diff --git a/api/tests/unit_tests/controllers/console/datasets/test_external_dataset_payload.py b/api/tests/unit_tests/services/entities/external_knowledge_entities/test_external_knowledge_entities.py similarity index 77% rename from api/tests/unit_tests/controllers/console/datasets/test_external_dataset_payload.py rename to api/tests/unit_tests/services/entities/external_knowledge_entities/test_external_knowledge_entities.py index 0693598d51b..d2faf8a445a 100644 --- a/api/tests/unit_tests/controllers/console/datasets/test_external_dataset_payload.py +++ b/api/tests/unit_tests/services/entities/external_knowledge_entities/test_external_knowledge_entities.py @@ -12,6 +12,12 @@ import pytest from pydantic import ValidationError from services.entities.external_knowledge_entities.external_knowledge_entities import ExternalDatasetCreatePayload +from tests.unit_tests.controllers.conftest import app_query_services as app_query_services # noqa: PLC0414 +from tests.unit_tests.controllers.conftest import ( + datasource_application_dependencies as datasource_application_dependencies, # noqa: PLC0414 +) +from tests.unit_tests.controllers.conftest import error_response_contract as error_response_contract # noqa: PLC0414 +from tests.unit_tests.controllers.conftest import pipeline_application as pipeline_application # noqa: PLC0414 def test_external_dataset_create_payload_allows_name_length_100() -> None: diff --git a/api/tests/unit_tests/services/agent/test_agent_composer_entities.py b/api/tests/unit_tests/services/entities/test_agent_entities.py similarity index 100% rename from api/tests/unit_tests/services/agent/test_agent_composer_entities.py rename to api/tests/unit_tests/services/entities/test_agent_entities.py diff --git a/api/tests/unit_tests/services/test_feature_entities.py b/api/tests/unit_tests/services/entities/test_feature_entities.py similarity index 100% rename from api/tests/unit_tests/services/test_feature_entities.py rename to api/tests/unit_tests/services/entities/test_feature_entities.py diff --git a/api/tests/unit_tests/services/plugin/test_plugin_service_installation.py b/api/tests/unit_tests/services/errors/test_plugin.py similarity index 99% rename from api/tests/unit_tests/services/plugin/test_plugin_service_installation.py rename to api/tests/unit_tests/services/errors/test_plugin.py index 2e1d89378c3..a515dd4c1b1 100644 --- a/api/tests/unit_tests/services/plugin/test_plugin_service_installation.py +++ b/api/tests/unit_tests/services/errors/test_plugin.py @@ -29,6 +29,8 @@ from services.entities.feature_entities import ( PluginInstallationScope, ) from services.errors.plugin import PluginInstallationForbiddenError +from tests.unit_tests.services.plugin.conftest import mock_features as mock_features # noqa: PLC0414 +from tests.unit_tests.services.plugin.conftest import mock_installer as mock_installer # noqa: PLC0414 def _make_permission( diff --git a/api/tests/unit_tests/services/knowledge/test_dataset_adapters.py b/api/tests/unit_tests/services/knowledge/datasets/test_adapters.py similarity index 99% rename from api/tests/unit_tests/services/knowledge/test_dataset_adapters.py rename to api/tests/unit_tests/services/knowledge/datasets/test_adapters.py index af252d99a56..5af19785fad 100644 --- a/api/tests/unit_tests/services/knowledge/test_dataset_adapters.py +++ b/api/tests/unit_tests/services/knowledge/datasets/test_adapters.py @@ -42,8 +42,8 @@ from services.knowledge.entities.knowledge_entities import KnowledgeConfig from services.knowledge.resource_scope import DatasetRef from services.rbac import contracts as rbac_contracts from services.tag_application_service import CreateTagInput, TagApplicationService, TagBindingInput -from tests.unit_tests.config_override import apply_config_overrides -from tests.unit_tests.rbac_fakes import RBACDomain +from tests.helpers.unit.config_override import apply_config_overrides +from tests.helpers.unit.rbac_fakes import RBACDomain CONTEXT = RequestContext("request", None, "actor", "tenant") REF = DatasetRef("tenant", "dataset") diff --git a/api/tests/unit_tests/services/knowledge/test_dataset_application_service.py b/api/tests/unit_tests/services/knowledge/datasets/test_application.py similarity index 100% rename from api/tests/unit_tests/services/knowledge/test_dataset_application_service.py rename to api/tests/unit_tests/services/knowledge/datasets/test_application.py diff --git a/api/tests/unit_tests/services/knowledge/test_knowledge_upgrade_service.py b/api/tests/unit_tests/services/knowledge/datasets/test_upgrades.py similarity index 98% rename from api/tests/unit_tests/services/knowledge/test_knowledge_upgrade_service.py rename to api/tests/unit_tests/services/knowledge/datasets/test_upgrades.py index d71ab450f06..f9affa23c26 100644 --- a/api/tests/unit_tests/services/knowledge/test_knowledge_upgrade_service.py +++ b/api/tests/unit_tests/services/knowledge/datasets/test_upgrades.py @@ -19,7 +19,7 @@ from services.knowledge.dataset_access import AccessibleDataset, DatasetAccess, from services.knowledge.datasets.adapters import SQLAlchemyDatasetOperations from services.knowledge.datasets.application import DatasetVisibility from services.knowledge.datasets.upgrades import KnowledgeFSUpgradeUnavailableError, KnowledgeUpgradeApplicationService -from tests.unit_tests.services.knowledge_fs_services_a_fakes import UnusedSessionFactory, upgrade_job +from tests.helpers.unit.services.knowledge_fs_services_a_fakes import UnusedSessionFactory, upgrade_job CONTEXT = RequestContext("request", None, "account-1", "tenant-1") diff --git a/api/tests/unit_tests/services/knowledge/documents/test_adapters.py b/api/tests/unit_tests/services/knowledge/documents/test_adapters.py new file mode 100644 index 00000000000..4377d1f98e9 --- /dev/null +++ b/api/tests/unit_tests/services/knowledge/documents/test_adapters.py @@ -0,0 +1,657 @@ +import json +from collections.abc import Callable, Generator, Iterator +from contextlib import contextmanager +from datetime import datetime +from unittest.mock import Mock, patch + +import pytest +from redis.exceptions import LockNotOwnedError +from sqlalchemy import Engine, event, select +from sqlalchemy.orm import Session, sessionmaker + +from controllers.console.datasets.datasets_document import DocumentDetailResponse, DocumentWithSegmentsListResponse +from core.model_manager import ModelInstance +from core.rag.index_processor.constant.index_type import IndexStructureType, IndexTechniqueType +from enums import CloudPlan +from machinery.context import RequestContext +from models.account import Account, Tenant, TenantAccountJoin, TenantAccountRole +from models.dataset import Dataset, DatasetProcessRule, Document, DocumentPipelineExecutionLog, DocumentSegment +from models.enums import DataSourceType, DocumentCreatedFrom, IndexingStatus, ProcessRuleMode, SegmentStatus +from services.document_indexing_proxy.batch_indexing_base import BatchDocumentIndexingProxy +from services.document_indexing_proxy.document_indexing_task_proxy import DocumentIndexingTaskProxy +from services.document_indexing_proxy.duplicate_document_indexing_task_proxy import DuplicateDocumentIndexingTaskProxy +from services.entities.feature_entities import FeatureModel +from services.errors.file import FileNotExistsError +from services.knowledge.documents.adapters import SQLAlchemyDocumentOperations +from services.knowledge.documents.application import DocumentListFilter, DocumentNotFoundError +from services.knowledge.resource_scope import DatasetRef +from tests.helpers.unit.model_factories import make_upload_file + +_DOCUMENT_ADAPTERS_REF = DatasetRef("tenant-1", "dataset-1") + + +def _dataset(**values: object) -> Dataset: + row = Dataset( + id="dataset-1", + tenant_id="tenant-1", + name="Dataset", + created_by="actor-1", + maintainer="actor-1", + indexing_technique="economy", + chunk_structure="text_model", + ) + for key, value in values.items(): + setattr(row, key, value) + return row + + +def _document(**values: object) -> Document: + row = Document( + tenant_id="tenant-1", + dataset_id="dataset-1", + position=1, + data_source_type=DataSourceType("local_file"), + batch="batch-1", + name="Document", + created_from=DocumentCreatedFrom("web"), + created_by="actor-1", + ) + row.id = "document-1" + row.data_source_info = json.dumps({"file_path": "/path/to/file"}) + row.indexing_status = IndexingStatus.COMPLETED + row.doc_form = IndexStructureType.PARAGRAPH_INDEX + row.enabled = True + row.archived = False + row.word_count = 10 + row.created_at = datetime(2024, 1, 1) + for key, value in values.items(): + setattr(row, key, value) + return row + + +def _segment( + segment_id: str, + *, + tenant_id: str = "tenant-1", + dataset_id: str = "dataset-1", + document_id: str = "document-1", + status: SegmentStatus = SegmentStatus.COMPLETED, + completed_at: datetime | None = datetime(2024, 1, 1), + hit_count: int = 0, +) -> DocumentSegment: + segment = DocumentSegment( + tenant_id=tenant_id, + dataset_id=dataset_id, + document_id=document_id, + position=1, + content="content", + word_count=7, + tokens=2, + created_by="actor-1", + status=status, + completed_at=completed_at, + hit_count=hit_count, + ) + segment.id = segment_id + return segment + + +class TestDocumentAdaptersCases: + @staticmethod + @pytest.fixture + def operations(sqlite_session_factory: sessionmaker[Session]) -> SQLAlchemyDocumentOperations: + with sqlite_session_factory.begin() as session: + session.add_all([_dataset(), _document()]) + return SQLAlchemyDocumentOperations(session_factory=sqlite_session_factory) + + @staticmethod + @pytest.mark.parametrize("foreign", [{"tenant_id": "other"}, {"dataset_id": "other"}]) + def test_document_reads_and_writes_reject_incorrect_owner( + operations: SQLAlchemyDocumentOperations, sqlite_session_factory: sessionmaker[Session], foreign: dict[str, str] + ) -> None: + with sqlite_session_factory.begin() as session: + session.add(_document(id="foreign", **foreign)) + ref = _DOCUMENT_ADAPTERS_REF.document("foreign") + assert operations.get_state(ref) is None + with pytest.raises(DocumentNotFoundError): + operations.get_detail(ref, metadata_only=False) + with pytest.raises(DocumentNotFoundError): + operations.update_document(ref, {"name": "changed"}) + with pytest.raises(DocumentNotFoundError): + operations.rename_document(ref, "changed") + with sqlite_session_factory() as session: + document = session.get(Document, "foreign") + assert document is not None + assert document.name == "Document" + + @staticmethod + def test_document_lookup_uses_active_workspace(operations: SQLAlchemyDocumentOperations) -> None: + assert operations.find_document(workspace_id="other", document_id="document-1") is None + assert operations.find_document( + workspace_id="tenant-1", document_id="document-1" + ) == _DOCUMENT_ADAPTERS_REF.document("document-1") + + @staticmethod + def test_listing_filters_scope_and_materializes_values( + operations: SQLAlchemyDocumentOperations, sqlite_session_factory: sessionmaker[Session] + ) -> None: + with sqlite_session_factory.begin() as session: + session.add_all( + [ + _document(id="matching", name="Needle", position=2), + _document(id="disabled", name="Needle", enabled=False), + _document(id="foreign", name="Needle", tenant_id="other"), + _document(id="other-dataset", name="Needle", dataset_id="other"), + ] + ) + result = operations.list_documents( + _DOCUMENT_ADAPTERS_REF, DocumentListFilter(search="Needle", status="available", fetch=True) + ) + response = DocumentWithSegmentsListResponse.model_validate(result).model_dump(mode="json") + assert response["total"] == 1 + assert response["data"][0]["id"] == "matching" + assert response["data"][0]["data_source_info"] == {"file_path": "/path/to/file"} + assert response["data"][0]["completed_segments"] == 0 + assert response["data"][0]["total_segments"] == 0 + + @staticmethod + def test_listing_hit_count_sort_and_exact_last_page( + operations: SQLAlchemyDocumentOperations, sqlite_session_factory: sessionmaker[Session] + ) -> None: + with sqlite_session_factory.begin() as session: + session.add_all( + [ + _document(id="popular", position=2), + _segment("segment-1", hit_count=1), + _segment("segment-2", document_id="popular", hit_count=10), + _segment("foreign", tenant_id="other", hit_count=1000), + ] + ) + first = operations.list_documents(_DOCUMENT_ADAPTERS_REF, DocumentListFilter(sort="-hit_count", limit=1)) + last = operations.list_documents(_DOCUMENT_ADAPTERS_REF, DocumentListFilter(sort="-hit_count", limit=1, page=2)) + assert [row["id"] for row in first["data"]] == ["popular"] + assert first["has_more"] is True + assert last["has_more"] is False + assert last["total"] == 2 + + @staticmethod + def test_status_counts_use_one_aggregate_and_complete_owner_chain( + operations: SQLAlchemyDocumentOperations, sqlite_session_factory: sessionmaker[Session], sqlite_engine: Engine + ) -> None: + with sqlite_session_factory.begin() as session: + session.add_all( + [ + _segment("complete"), + _segment("waiting", completed_at=None), + _segment("excluded", status=SegmentStatus.RE_SEGMENT), + _segment("foreign", tenant_id="other"), + _segment("wrong-dataset", dataset_id="other"), + _document(id="second", is_paused=True, indexing_status=IndexingStatus.INDEXING), + ] + ) + statements: list[str] = [] + + def record( + _conn: object, _cursor: object, statement: str, _parameters: object, _context: object, _many: bool + ) -> None: + statements.append(statement) + + event.listen(sqlite_engine, "before_cursor_execute", record) + try: + result = operations.get_batch_indexing_status(_DOCUMENT_ADAPTERS_REF, "batch-1") + finally: + event.remove(sqlite_engine, "before_cursor_execute", record) + rows = {row["id"]: row for row in result["data"]} + assert rows["document-1"]["completed_segments"] == 1 + assert rows["document-1"]["total_segments"] == 2 + assert rows["second"]["indexing_status"] == "paused" + assert sum("FROM document_segments" in statement for statement in statements) == 1 + + @staticmethod + def test_unknown_batch_is_not_found(operations: SQLAlchemyDocumentOperations) -> None: + with pytest.raises(DocumentNotFoundError): + operations.get_batch_indexing_status(_DOCUMENT_ADAPTERS_REF, "missing") + + @staticmethod + @pytest.mark.parametrize("rules", [None, {"segmentation": {"delimiter": "\\n", "max_tokens": 500}}]) + def test_process_rule_preserves_null_and_legacy_delimiter( + operations: SQLAlchemyDocumentOperations, + sqlite_session_factory: sessionmaker[Session], + rules: dict[str, object] | None, + ) -> None: + with sqlite_session_factory.begin() as session: + session.add( + DatasetProcessRule( + dataset_id="dataset-1", + mode=ProcessRuleMode.CUSTOM, + rules=json.dumps(rules) if rules else None, + created_by="actor-1", + ) + ) + result = operations.get_process_rule(_DOCUMENT_ADAPTERS_REF) + assert result["mode"] == "custom" + assert result["rules"] == rules + + @staticmethod + def test_detail_is_serializable_after_session_closes(operations: SQLAlchemyDocumentOperations) -> None: + result = operations.get_detail(_DOCUMENT_ADAPTERS_REF.document("document-1"), metadata_only=False) + response = DocumentDetailResponse.model_validate(result).model_dump(mode="json") + assert response["id"] == "document-1" + assert response["data_source_info"] == {"file_path": "/path/to/file"} + assert response["data_source_detail_dict"] == {} + assert response["created_at"] == int(datetime(2024, 1, 1).timestamp()) + + @staticmethod + def test_metadata_only_avoids_loading_full_details(operations: SQLAlchemyDocumentOperations) -> None: + with patch("services.knowledge.documents.adapters.get_document_source_detail") as source: + result = operations.get_detail(_DOCUMENT_ADAPTERS_REF.document("document-1"), metadata_only=True) + assert set(result) == {"id", "doc_type", "doc_metadata"} + source.assert_not_called() + + @staticmethod + def test_update_is_committed_before_returning( + operations: SQLAlchemyDocumentOperations, sqlite_session_factory: sessionmaker[Session] + ) -> None: + operations.update_document( + _DOCUMENT_ADAPTERS_REF.document("document-1"), {"is_paused": True, "paused_by": "actor-1"} + ) + with sqlite_session_factory() as session: + document = session.get(Document, "document-1") + assert document is not None + assert document.is_paused is True + assert document.paused_by == "actor-1" + + @staticmethod + def test_latest_pipeline_log_and_empty_result( + operations: SQLAlchemyDocumentOperations, sqlite_session_factory: sessionmaker[Session] + ) -> None: + assert operations.get_execution_log(_DOCUMENT_ADAPTERS_REF.document("document-1")) == {} + with sqlite_session_factory.begin() as session: + for year in [2024, 2025]: + row = DocumentPipelineExecutionLog( + pipeline_id="pipeline-1", + document_id="document-1", + datasource_type="local_file", + datasource_info=json.dumps({"year": year}), + datasource_node_id="node-1", + input_data={"year": year}, + created_by="actor-1", + ) + row.created_at = datetime(year, 1, 1) + session.add(row) + result = operations.get_execution_log(_DOCUMENT_ADAPTERS_REF.document("document-1")) + assert result["datasource_info"] == {"year": 2025} + assert result["input_data"] == {"year": 2025} + + @staticmethod + @pytest.mark.parametrize("settings", [None, {"enable": False}]) + def test_summary_guard_rejects_disabled_summary( + operations: SQLAlchemyDocumentOperations, + sqlite_session_factory: sessionmaker[Session], + settings: dict[str, object] | None, + ) -> None: + with sqlite_session_factory.begin() as session: + dataset = session.get(Dataset, "dataset-1") + assert dataset is not None + dataset.indexing_technique = IndexTechniqueType.HIGH_QUALITY + dataset.summary_index_setting = settings + with pytest.raises(ValueError, match="not enabled"): + operations.require_summary_enabled(_DOCUMENT_ADAPTERS_REF) + + @staticmethod + def test_enable_summary_only_updates_requested_owner( + operations: SQLAlchemyDocumentOperations, sqlite_session_factory: sessionmaker[Session] + ) -> None: + with sqlite_session_factory.begin() as session: + session.add(_document(id="foreign", tenant_id="other", need_summary=False)) + operations.enable_summary(_DOCUMENT_ADAPTERS_REF, ["document-1", "foreign"]) + with sqlite_session_factory() as session: + document = session.get(Document, "document-1") + foreign = session.get(Document, "foreign") + assert document is not None + assert foreign is not None + assert document.need_summary is True + assert foreign.need_summary is False + + +"""Exercise document creation with real SQL and assertions at external boundaries.""" + + +CONTEXT = RequestContext("request", None, "actor", "tenant") +_DOCUMENT_CREATION_REF = DatasetRef("tenant", "dataset") + + +def settings(**values: object) -> dict[str, object]: + return { + "indexing_technique": "economy", + "process_rule": {"mode": "automatic"}, + "data_source": {"info_list": {"data_source_type": "upload_file", "file_info_list": {"file_ids": ["file"]}}}, + **values, + } + + +class TestDocumentCreationCases: + @staticmethod + @pytest.fixture + def operations( + sqlite_session_factory: sessionmaker[Session], config_overrides: Callable[..., None] + ) -> SQLAlchemyDocumentOperations: + config_overrides(DEPLOYMENT_EDITION="CLOUD", RBAC_ENABLED=False) + with sqlite_session_factory.begin() as session: + actor = Account(name="Actor", email="actor@example.com") + actor.id = "actor" + tenant = Tenant(name="Tenant") + tenant.id = "tenant" + session.add_all( + [ + actor, + tenant, + TenantAccountJoin(tenant_id="tenant", account_id="actor", role=TenantAccountRole.OWNER), + Dataset( + id="dataset", + tenant_id="tenant", + name="Dataset", + created_by="actor", + indexing_technique="economy", + ), + make_upload_file(file_id="file", tenant_id="tenant", name="example.txt", created_by="actor"), + make_upload_file(file_id="foreign-file", tenant_id="foreign", name="foreign.txt"), + ] + ) + return SQLAlchemyDocumentOperations(session_factory=sqlite_session_factory) + + @staticmethod + @pytest.fixture + def outside_transaction(sqlite_engine: Engine) -> Iterator[Callable[[], None]]: + active: set[int] = set() + + def begin(connection: object) -> None: + active.add(id(connection)) + + def end(connection: object) -> None: + active.discard(id(connection)) + + def check() -> None: + assert not active, "External I/O must not hold an open database transaction" + + listeners = [("begin", begin), ("commit", end), ("rollback", end)] + for name, callback in listeners: + event.listen(sqlite_engine, name, callback) + yield check + for name, callback in listeners: + event.remove(sqlite_engine, name, callback) + + @staticmethod + @pytest.fixture + def external_calls( + monkeypatch: pytest.MonkeyPatch, + outside_transaction: Callable[[], None], + sqlite_session_factory: sessionmaker[Session], + ) -> list[tuple[str, tuple[str, ...]]]: + calls: list[tuple[str, tuple[str, ...]]] = [] + + def features(_tenant_id: str, *, exclude_vector_space: bool) -> FeatureModel: + outside_transaction() + assert exclude_vector_space is True + result = FeatureModel() + result.billing.subscription.plan = CloudPlan.PROFESSIONAL + return result + + @contextmanager + def lock(_name: str, *, timeout: int) -> Generator[None]: + outside_transaction() + assert timeout == 600 + yield + outside_transaction() + + def dispatched(kind: str, dataset_id: str, document_ids: tuple[str, ...]) -> None: + outside_transaction() + with sqlite_session_factory() as session: + for document_id in document_ids: + document = session.get(Document, document_id) + assert document is not None + assert document.dataset_id == dataset_id + if kind != "clean": + assert document.indexing_status == IndexingStatus.WAITING + calls.append((kind, document_ids)) + + def create(proxy: BatchDocumentIndexingProxy) -> None: + dispatched("create", proxy._dataset_id, tuple(proxy._document_ids)) + + def duplicate(proxy: BatchDocumentIndexingProxy) -> None: + dispatched("duplicate", proxy._dataset_id, tuple(proxy._document_ids)) + + monkeypatch.setattr("services.knowledge.dataset_service.FeatureService.get_features", features) + monkeypatch.setattr("services.knowledge.documents.adapters.redis_client.lock", lock) + monkeypatch.setattr(DocumentIndexingTaskProxy, "delay", create) + monkeypatch.setattr(DuplicateDocumentIndexingTaskProxy, "delay", duplicate) + monkeypatch.setattr( + "services.knowledge.dataset_service.document_indexing_update_task.delay", + lambda dataset_id, document_id: dispatched("update", dataset_id, (document_id,)), + ) + monkeypatch.setattr( + "services.knowledge.dataset_service.clean_notion_document_task.delay", + lambda document_ids, dataset_id: dispatched("clean", dataset_id, tuple(document_ids)), + ) + return calls + + @staticmethod + @pytest.mark.parametrize("initialize", [False, True]) + def test_upload_preserves_response_and_commits_before_dispatch( + operations: SQLAlchemyDocumentOperations, + external_calls: list[tuple[str, tuple[str, ...]]], + sqlite_session_factory: sessionmaker[Session], + initialize: bool, + ) -> None: + result = ( + operations.initialize_dataset(CONTEXT, settings()) + if initialize + else operations.create_documents(CONTEXT, _DOCUMENT_CREATION_REF, settings()) + ) + (document,) = result["documents"] + assert document["name"] == "example.txt" + assert document["data_source_info_dict"] == {"upload_file_id": "file"} + assert document["indexing_status"] == "waiting" + assert result["batch"] + assert external_calls == [("create", (document["id"],))] + with sqlite_session_factory() as session: + dataset = session.get(Dataset, result["dataset"]["id"]) + assert dataset is not None + assert dataset.tenant_id == "tenant" + assert dataset.name == ("example.txt..." if initialize else "Dataset") + assert dataset.permission == "only_me" + + @staticmethod + def test_duplicate_upload_reuses_document_and_dispatches_duplicate_job( + operations: SQLAlchemyDocumentOperations, external_calls: list[tuple[str, tuple[str, ...]]] + ) -> None: + first = operations.create_documents(CONTEXT, _DOCUMENT_CREATION_REF, settings()) + second = operations.create_documents(CONTEXT, _DOCUMENT_CREATION_REF, settings()) + first_id = first["documents"][0]["id"] + assert second["documents"][0]["id"] == first_id + assert external_calls == [("create", (first_id,)), ("duplicate", (first_id,))] + + @staticmethod + def test_reprocess_checks_model_outside_transaction_and_preserves_segment_reset( + operations: SQLAlchemyDocumentOperations, + external_calls: list[tuple[str, tuple[str, ...]]], + outside_transaction: Callable[[], None], + sqlite_session_factory: sessionmaker[Session], + monkeypatch: pytest.MonkeyPatch, + ) -> None: + created = operations.create_documents(CONTEXT, _DOCUMENT_CREATION_REF, settings()) + document_id = created["documents"][0]["id"] + with sqlite_session_factory.begin() as session: + document = session.get(Document, document_id) + assert document is not None + document.indexing_status = IndexingStatus.COMPLETED + document.completed_at = datetime(2026, 1, 1) + dataset = session.get(Dataset, _DOCUMENT_CREATION_REF.dataset_id) + assert dataset is not None + dataset.indexing_technique = IndexTechniqueType.HIGH_QUALITY + dataset.embedding_model = "embedding" + dataset.embedding_model_provider = "provider" + session.add( + DocumentSegment( + tenant_id="tenant", + dataset_id="dataset", + document_id=document_id, + position=1, + content="content", + word_count=1, + tokens=1, + created_by="actor", + status=SegmentStatus.COMPLETED, + ) + ) + manager = Mock() + manager.get_model_instance.side_effect = lambda **_kwargs: outside_transaction() + monkeypatch.setattr("services.knowledge.dataset_service.ModelManager.for_tenant", lambda **_kwargs: manager) + result = operations.create_documents( + CONTEXT, _DOCUMENT_CREATION_REF, settings(original_document_id=document_id, name="Renamed") + ) + assert result["documents"][0]["id"] == document_id + assert result["documents"][0]["name"] == "Renamed" + assert result["batch"] == created["batch"] + assert external_calls[-1] == ("update", (document_id,)) + manager.get_model_instance.assert_called_once() + with sqlite_session_factory() as session: + segment = session.scalar(select(DocumentSegment).where(DocumentSegment.document_id == document_id)) + assert segment is not None + assert segment.status == SegmentStatus.RE_SEGMENT + + @staticmethod + def test_first_high_quality_upload_resolves_default_model_without_transaction( + operations: SQLAlchemyDocumentOperations, + external_calls: list[tuple[str, tuple[str, ...]]], + outside_transaction: Callable[[], None], + sqlite_session_factory: sessionmaker[Session], + monkeypatch: pytest.MonkeyPatch, + ) -> None: + with sqlite_session_factory.begin() as session: + dataset = session.get(Dataset, _DOCUMENT_CREATION_REF.dataset_id) + assert dataset is not None + dataset.indexing_technique = None + + def model(**_kwargs: object) -> ModelInstance: + outside_transaction() + instance = object.__new__(ModelInstance) + instance.model_name = "default-embedding" + instance.provider = "default-provider" + return instance + + manager = Mock() + manager.get_default_model_instance.side_effect = model + monkeypatch.setattr("services.knowledge.dataset_service.ModelManager.for_tenant", lambda **_kwargs: manager) + operations.create_documents(CONTEXT, _DOCUMENT_CREATION_REF, settings(indexing_technique="high_quality")) + with sqlite_session_factory() as session: + dataset = session.get(Dataset, _DOCUMENT_CREATION_REF.dataset_id) + assert dataset is not None + assert dataset.embedding_model == "default-embedding" + assert dataset.embedding_model_provider == "default-provider" + assert dataset.collection_binding_id is not None + assert len(external_calls) == 1 + + @staticmethod + def test_foreign_upload_is_rejected_without_dispatch( + operations: SQLAlchemyDocumentOperations, external_calls: list[tuple[str, tuple[str, ...]]] + ) -> None: + config = settings( + data_source={ + "info_list": {"data_source_type": "upload_file", "file_info_list": {"file_ids": ["foreign-file"]}} + } + ) + with pytest.raises(FileNotExistsError): + operations.create_documents(CONTEXT, _DOCUMENT_CREATION_REF, config) + assert external_calls == [] + + @staticmethod + def test_notion_removal_is_dispatched_after_new_documents_commit( + operations: SQLAlchemyDocumentOperations, external_calls: list[tuple[str, tuple[str, ...]]] + ) -> None: + def notion(page_id: str) -> dict[str, object]: + return settings( + data_source={ + "info_list": { + "data_source_type": "notion_import", + "notion_info_list": [ + { + "credential_id": "credential", + "workspace_id": "workspace", + "pages": [{"page_id": page_id, "page_name": page_id, "type": "page"}], + } + ], + } + } + ) + + first = operations.create_documents(CONTEXT, _DOCUMENT_CREATION_REF, notion("old-page")) + old_id = first["documents"][0]["id"] + operations.create_documents(CONTEXT, _DOCUMENT_CREATION_REF, notion("new-page")) + assert external_calls[-2] == ("clean", (old_id,)) + assert external_calls[-1][0] == "create" + assert old_id not in external_calls[-1][1] + + @staticmethod + def test_cloud_quota_rejection_does_not_create_documents_or_dispatch( + operations: SQLAlchemyDocumentOperations, + external_calls: list[tuple[str, tuple[str, ...]]], + sqlite_session_factory: sessionmaker[Session], + config_overrides: Callable[..., None], + ) -> None: + config_overrides(BATCH_UPLOAD_LIMIT=0) + with pytest.raises(ValueError, match="batch upload limit"): + operations.create_documents(CONTEXT, _DOCUMENT_CREATION_REF, settings()) + assert external_calls == [] + with sqlite_session_factory() as session: + assert session.scalar(select(Document.id)) is None + + @staticmethod + def test_response_read_failure_still_dispatches_already_committed_documents( + operations: SQLAlchemyDocumentOperations, + external_calls: list[tuple[str, tuple[str, ...]]], + monkeypatch: pytest.MonkeyPatch, + ) -> None: + response = Mock(side_effect=ValueError("response read failed")) + monkeypatch.setattr("services.knowledge.documents.adapters._created_response", response) + with pytest.raises(ValueError, match="response read failed"): + operations.create_documents(CONTEXT, _DOCUMENT_CREATION_REF, settings()) + assert len(external_calls) == 1 + assert external_calls[0][0] == "create" + + @staticmethod + def test_lost_creation_lock_keeps_empty_batch_response( + operations: SQLAlchemyDocumentOperations, + external_calls: list[tuple[str, tuple[str, ...]]], + monkeypatch: pytest.MonkeyPatch, + ) -> None: + @contextmanager + def lock(_name: str, **_kwargs: object) -> Generator[None]: + raise LockNotOwnedError("lock lost") + yield # pragma: no cover + + monkeypatch.setattr("services.knowledge.documents.adapters.redis_client.lock", lock) + result = operations.create_documents(CONTEXT, _DOCUMENT_CREATION_REF, settings()) + assert result["documents"] == [] + assert result["batch"] + assert result["dataset"]["id"] == _DOCUMENT_CREATION_REF.dataset_id + assert external_calls == [] + + @staticmethod + def test_billing_http_boundary_does_not_hold_a_transaction( + operations: SQLAlchemyDocumentOperations, + outside_transaction: Callable[[], None], + monkeypatch: pytest.MonkeyPatch, + ) -> None: + class StopBeforeNetworkError(Exception): + pass + + def request(_method: str, url: str, **_kwargs: object) -> None: + assert url.endswith("/subscription/info") + outside_transaction() + raise StopBeforeNetworkError() + + monkeypatch.setattr("services.billing_service._http_client.request", request) + with pytest.raises(StopBeforeNetworkError): + operations.create_documents(CONTEXT, _DOCUMENT_CREATION_REF, settings()) + outside_transaction() diff --git a/api/tests/unit_tests/services/knowledge/test_document_application_service.py b/api/tests/unit_tests/services/knowledge/documents/test_application.py similarity index 100% rename from api/tests/unit_tests/services/knowledge/test_document_application_service.py rename to api/tests/unit_tests/services/knowledge/documents/test_application.py diff --git a/api/tests/unit_tests/services/knowledge/test_indexing_estimate_entities.py b/api/tests/unit_tests/services/knowledge/entities/test_indexing_estimate.py similarity index 100% rename from api/tests/unit_tests/services/knowledge/test_indexing_estimate_entities.py rename to api/tests/unit_tests/services/knowledge/entities/test_indexing_estimate.py diff --git a/api/tests/unit_tests/services/knowledge/test_external_adapters.py b/api/tests/unit_tests/services/knowledge/external/test_adapters.py similarity index 99% rename from api/tests/unit_tests/services/knowledge/test_external_adapters.py rename to api/tests/unit_tests/services/knowledge/external/test_adapters.py index c50a2e878a6..5bae55a8e19 100644 --- a/api/tests/unit_tests/services/knowledge/test_external_adapters.py +++ b/api/tests/unit_tests/services/knowledge/external/test_adapters.py @@ -17,7 +17,7 @@ from services.knowledge.dataset_access import DatasetNotFoundError from services.knowledge.external.adapters import SQLAlchemyExternalKnowledgeOperations from services.knowledge.external.application import ExternalTemplateNotFoundError from services.knowledge.resource_scope import DatasetRef -from tests.unit_tests.rbac_fakes import RBACDomain +from tests.helpers.unit.rbac_fakes import RBACDomain CONTEXT = RequestContext("request", None, "actor", "tenant") REF = DatasetRef("tenant", "dataset") diff --git a/api/tests/unit_tests/services/knowledge/test_external_application_service.py b/api/tests/unit_tests/services/knowledge/external/test_application.py similarity index 100% rename from api/tests/unit_tests/services/knowledge/test_external_application_service.py rename to api/tests/unit_tests/services/knowledge/external/test_application.py diff --git a/api/tests/unit_tests/services/test_external_dataset_service.py b/api/tests/unit_tests/services/knowledge/external/test_service.py similarity index 100% rename from api/tests/unit_tests/services/test_external_dataset_service.py rename to api/tests/unit_tests/services/knowledge/external/test_service.py diff --git a/api/tests/unit_tests/services/knowledge/test_document_index_cleanup.py b/api/tests/unit_tests/services/knowledge/indexing/adapters/test_cleanup.py similarity index 100% rename from api/tests/unit_tests/services/knowledge/test_document_index_cleanup.py rename to api/tests/unit_tests/services/knowledge/indexing/adapters/test_cleanup.py diff --git a/api/tests/unit_tests/services/knowledge/test_indexing_estimate_adapters.py b/api/tests/unit_tests/services/knowledge/indexing/adapters/test_estimate.py similarity index 100% rename from api/tests/unit_tests/services/knowledge/test_indexing_estimate_adapters.py rename to api/tests/unit_tests/services/knowledge/indexing/adapters/test_estimate.py diff --git a/api/tests/unit_tests/services/knowledge/test_indexing_execution_adapters.py b/api/tests/unit_tests/services/knowledge/indexing/adapters/test_execution.py similarity index 99% rename from api/tests/unit_tests/services/knowledge/test_indexing_execution_adapters.py rename to api/tests/unit_tests/services/knowledge/indexing/adapters/test_execution.py index 45ec28a4ec8..40d54d800c4 100644 --- a/api/tests/unit_tests/services/knowledge/test_indexing_execution_adapters.py +++ b/api/tests/unit_tests/services/knowledge/indexing/adapters/test_execution.py @@ -29,7 +29,7 @@ from services.knowledge.indexing.adapters.execution import IndexingExecutionAdap from services.knowledge.indexing.errors import DocumentIsPausedError from services.knowledge.indexing.execution import DocumentIndexingService, IndexingDocument from services.knowledge.resource_scope import DatasetRef -from tests.unit_tests.config_override import apply_config_overrides +from tests.helpers.unit.config_override import apply_config_overrides from tests.unit_tests.repositories.knowledge.test_document_repository import _dataset, _document MODULE = "services.knowledge.indexing.adapters.execution" diff --git a/api/tests/unit_tests/services/knowledge/test_indexing_source_adapters.py b/api/tests/unit_tests/services/knowledge/indexing/adapters/test_sources.py similarity index 99% rename from api/tests/unit_tests/services/knowledge/test_indexing_source_adapters.py rename to api/tests/unit_tests/services/knowledge/indexing/adapters/test_sources.py index be642c3d6a9..1ee3038b1b3 100644 --- a/api/tests/unit_tests/services/knowledge/test_indexing_source_adapters.py +++ b/api/tests/unit_tests/services/knowledge/indexing/adapters/test_sources.py @@ -27,7 +27,7 @@ from services.knowledge.indexing.errors import ( ) from services.knowledge.indexing.estimate import StoredSource from services.knowledge.resource_scope import DatasetRef -from tests.unit_tests.config_override import apply_config_overrides +from tests.helpers.unit.config_override import apply_config_overrides def _source( diff --git a/api/tests/unit_tests/services/knowledge/test_indexing_estimate_application_service.py b/api/tests/unit_tests/services/knowledge/indexing/test_estimate.py similarity index 100% rename from api/tests/unit_tests/services/knowledge/test_indexing_estimate_application_service.py rename to api/tests/unit_tests/services/knowledge/indexing/test_estimate.py diff --git a/api/tests/unit_tests/services/knowledge/test_indexing_execution.py b/api/tests/unit_tests/services/knowledge/indexing/test_execution.py similarity index 100% rename from api/tests/unit_tests/services/knowledge/test_indexing_execution.py rename to api/tests/unit_tests/services/knowledge/indexing/test_execution.py diff --git a/api/tests/unit_tests/services/knowledge/metadata/test_application.py b/api/tests/unit_tests/services/knowledge/metadata/test_application.py new file mode 100644 index 00000000000..fad7ddecd81 --- /dev/null +++ b/api/tests/unit_tests/services/knowledge/metadata/test_application.py @@ -0,0 +1,469 @@ +from collections.abc import Iterator +from datetime import datetime +from pathlib import Path +from typing import cast +from unittest.mock import Mock, create_autospec + +import pytest +from sqlalchemy import event, select +from sqlalchemy.orm import Session, sessionmaker + +from models import Account, Tenant +from models.dataset import DatasetMetadataBinding, Document +from models.enums import DocumentCreatedFrom +from repositories.knowledge.metadata_repository import SQLAlchemyMetadataRepository +from services.errors.metadata import MetadataResourceNotFoundError +from services.knowledge.dataset_access import DatasetAccess +from services.knowledge.entities.knowledge_entities import ( + DocumentMetadataOperation, + MetadataArgs, + MetadataDetail, + MetadataOperationData, +) +from services.knowledge.metadata.application import MetadataService +from services.knowledge.resource_scope import DatasetRef +from tests.helpers.unit.model_factories import make_account, make_dataset, make_document + + +def _metadata_bug_complete_metadata_service(session: Session) -> MetadataService: + return MetadataService( + store=SQLAlchemyMetadataRepository( + session_factory=sessionmaker(bind=session.get_bind(), expire_on_commit=False) + ), + dataset_access=create_autospec(DatasetAccess, instance=True), + ) + + +def _metadata_bug_complete_make_account(account_id: str = "user-456", tenant_id: str = "tenant-123") -> Account: + account = Account(name="Test User", email=f"{account_id}@example.com") + account.id = account_id + tenant = Tenant(name="Test Tenant") + tenant.id = tenant_id + account._current_tenant = tenant + return account + + +class TestMetadataBugCompleteCases: + class TestMetadataBugCompleteValidation: + """Complete test suite to verify the metadata nullable bug and its fix.""" + + def test_1_pydantic_layer_validation(self) -> None: + """Test Layer 1: Pydantic model validation correctly rejects None values.""" + # Pydantic should reject None values for required fields + with pytest.raises((ValueError, TypeError)): + MetadataArgs(type=None, name=None) # pyrefly: ignore[bad-argument-type] + + with pytest.raises((ValueError, TypeError)): + MetadataArgs(type="string", name=None) # pyrefly: ignore[bad-argument-type] + + with pytest.raises((ValueError, TypeError)): + MetadataArgs(type=None, name="test") # pyrefly: ignore[bad-argument-type] + + # Valid values should work + valid_args = MetadataArgs(type="string", name="test_name") + assert valid_args.type == "string" + assert valid_args.name == "test_name" + + @pytest.mark.parametrize("sqlite_session", [()], indirect=True) + def test_2_business_logic_layer_crashes_on_none(self, sqlite_session: Session) -> None: + """Test Layer 2: Business logic crashes when None values slip through.""" + # Create mock that bypasses Pydantic validation + mock_metadata_args = Mock() + mock_metadata_args.name = None + mock_metadata_args.type = "string" + + account = _metadata_bug_complete_make_account() + # Should crash with TypeError + with pytest.raises(TypeError, match="object of type 'NoneType' has no len"): + _metadata_bug_complete_metadata_service(sqlite_session).create_metadata( + DatasetRef("tenant-123", "dataset-123"), mock_metadata_args, actor_id=account.id + ) + + # Test update method as well + account = _metadata_bug_complete_make_account() + none_name = cast(str, None) + with pytest.raises(TypeError, match="object of type 'NoneType' has no len"): + _metadata_bug_complete_metadata_service(sqlite_session).update_metadata_name( + DatasetRef(Mock().tenant_id, Mock().id), "metadata-456", none_name, actor_id=account.id + ) + assert not sqlite_session.in_transaction() + + def test_3_database_constraints_verification(self) -> None: + """Test Layer 3: Verify database model has nullable=False constraints.""" + from sqlalchemy import inspect + + from models.dataset import DatasetMetadata + + # Get table info + mapper = inspect(DatasetMetadata) + + # Check that type and name columns are not nullable + type_column = mapper.columns["type"] + name_column = mapper.columns["name"] + + assert type_column.nullable is False, "type column should be nullable=False" + assert name_column.nullable is False, "name column should be nullable=False" + + def test_4_fixed_api_layer_rejects_null(self) -> None: + """Test Layer 4: Fixed API configuration properly rejects null values using Pydantic.""" + with pytest.raises((ValueError, TypeError)): + MetadataArgs.model_validate({"type": None, "name": None}) + + with pytest.raises((ValueError, TypeError)): + MetadataArgs.model_validate({"type": "string", "name": None}) + + with pytest.raises((ValueError, TypeError)): + MetadataArgs.model_validate({"type": None, "name": "test"}) + + def test_5_fixed_api_accepts_valid_values(self) -> None: + """Test that fixed API still accepts valid non-null values.""" + args = MetadataArgs.model_validate({"type": "string", "name": "valid_name"}) + assert args.type == "string" + assert args.name == "valid_name" + + @pytest.mark.parametrize("sqlite_session", [()], indirect=True) + def test_6_simulated_buggy_behavior(self, sqlite_session: Session) -> None: + """Test simulating the original buggy behavior by bypassing Pydantic validation.""" + mock_metadata_args = Mock() + mock_metadata_args.name = None + mock_metadata_args.type = None + + account = _metadata_bug_complete_make_account() + with pytest.raises(TypeError, match="object of type 'NoneType' has no len"): + _metadata_bug_complete_metadata_service(sqlite_session).create_metadata( + DatasetRef("tenant-123", "dataset-123"), mock_metadata_args, actor_id=account.id + ) + assert not sqlite_session.in_transaction() + + def test_7_end_to_end_validation_layers(self) -> None: + """Test all validation layers work together correctly.""" + # Layer 1: API should reject null at parameter level (with fix) + # Layer 2: Pydantic should reject null at model level + # Layer 3: Business logic expects non-null + # Layer 4: Database enforces non-null + + # Test that valid data flows through all layers + valid_data = {"type": "string", "name": "test_metadata"} + + # Should create valid Pydantic object + metadata_args = MetadataArgs.model_validate(valid_data) + assert metadata_args.type == "string" + assert metadata_args.name == "test_metadata" + + # Should not crash in business logic length check + assert len(metadata_args.name) <= 255 # This should not crash + assert len(metadata_args.type) > 0 # This should not crash + + def test_8_verify_specific_fix_locations(self) -> None: + """Verify that the specific locations mentioned in bug report are fixed.""" + # Read the actual files to verify fixes + import os + + # Console API create + console_create_file = "api/controllers/console/datasets/metadata.py" + if os.path.exists(console_create_file): + content = Path(console_create_file).read_text() + # Should contain nullable=False, not nullable=True + assert "nullable=True" not in content.split("class DatasetMetadataCreateApi")[1].split("class")[0] + + # Service API create + service_create_file = "api/controllers/service_api/dataset/metadata.py" + if os.path.exists(service_create_file): + content = Path(service_create_file).read_text() + # Should contain nullable=False, not nullable=True + create_api_section = content.split("class DatasetMetadataCreateServiceApi")[1].split("class")[0] + assert "nullable=True" not in create_api_section + + class TestMetadataValidationSummary: + """Summary tests that demonstrate the complete validation architecture.""" + + def test_validation_layer_architecture(self) -> None: + """Document and test the 4-layer validation architecture.""" + # Layer 1: API Parameter Validation (Flask-RESTful reqparse) + # - Role: First line of defense, validates HTTP request parameters + # - Fixed: nullable=False ensures null values are rejected at API boundary + + # Layer 2: Pydantic Model Validation + # - Role: Validates data structure and types before business logic + # - Working: Required fields without Optional[] reject None values + + # Layer 3: Business Logic Validation + # - Role: Domain-specific validation (length checks, uniqueness, etc.) + # - Vulnerable: Direct len() calls crash on None values + + # Layer 4: Database Constraints + # - Role: Final data integrity enforcement + # - Working: nullable=False prevents None values in database + + # The bug was: Layer 1 allowed None, but Layers 2-4 expected non-None + # The fix: Make Layer 1 consistent with Layers 2-4 + + assert True # This test documents the architecture + + +def _metadata_nullable_bug_metadata_service(session: Session) -> MetadataService: + return MetadataService( + store=SQLAlchemyMetadataRepository( + session_factory=sessionmaker(bind=session.get_bind(), expire_on_commit=False) + ), + dataset_access=create_autospec(DatasetAccess, instance=True), + ) + + +def _metadata_nullable_bug_make_account(account_id: str = "user-456", tenant_id: str = "tenant-123") -> Account: + account = Account(name="Test User", email=f"{account_id}@example.com") + account.id = account_id + tenant = Tenant(name="Test Tenant") + tenant.id = tenant_id + account._current_tenant = tenant + return account + + +class TestMetadataNullableBugCases: + pytestmark = [ + pytest.mark.usefixtures("sqlite_session"), + pytest.mark.parametrize("sqlite_session", [()], indirect=True), + ] + + class TestMetadataNullableBug: + """Test case to reproduce the metadata nullable validation bug.""" + + def test_metadata_args_with_none_values_should_fail(self) -> None: + """Test that MetadataArgs validation should reject None values.""" + # This test demonstrates the expected behavior - should fail validation + with pytest.raises((ValueError, TypeError)): + # This should fail because Pydantic expects non-None values + MetadataArgs(type=None, name=None) # pyrefly: ignore[bad-argument-type] + + def test_metadata_service_create_with_none_name_crashes(self, sqlite_session: Session) -> None: + """Test that MetadataService.create_metadata crashes when name is None.""" + # Mock the MetadataArgs to bypass Pydantic validation + mock_metadata_args = Mock() + mock_metadata_args.name = None # This will cause len() to crash + mock_metadata_args.type = "string" + + account = _metadata_nullable_bug_make_account() + # This should crash with TypeError when calling len(None) + with pytest.raises(TypeError, match="object of type 'NoneType' has no len"): + _metadata_nullable_bug_metadata_service(sqlite_session).create_metadata( + DatasetRef("tenant-123", "dataset-123"), mock_metadata_args, actor_id=account.id + ) + assert not sqlite_session.in_transaction() + + def test_metadata_service_update_with_none_name_crashes(self, sqlite_session: Session) -> None: + """Test that MetadataService.update_metadata_name crashes when name is None.""" + account = _metadata_nullable_bug_make_account() + none_name = cast(str, None) + # This should crash with TypeError when calling len(None) + with pytest.raises(TypeError, match="object of type 'NoneType' has no len"): + _metadata_nullable_bug_metadata_service(sqlite_session).update_metadata_name( + DatasetRef(Mock().tenant_id, Mock().id), "metadata-456", none_name, actor_id=account.id + ) + assert not sqlite_session.in_transaction() + + def test_api_layer_now_uses_pydantic_validation(self) -> None: + """Verify that API layer relies on Pydantic validation instead of reqparse.""" + invalid_payload: dict[str, object] = {"type": None, "name": None} + with pytest.raises((ValueError, TypeError)): + MetadataArgs.model_validate(invalid_payload) + + valid_payload = {"type": "string", "name": "valid"} + args = MetadataArgs.model_validate(valid_payload) + assert args.type == "string" + assert args.name == "valid" + + +"""Metadata operations persist through scoped, independently owned transactions.""" + + +REF = DatasetRef("tenant-1", "dataset-1") +DOCUMENT_ID = "11111111-1111-1111-1111-111111111111" +FOREIGN_ID = "22222222-2222-2222-2222-222222222222" + + +def _operation(document_id: str, metadata_id: str, *, partial: bool = False) -> MetadataOperationData: + return MetadataOperationData( + operation_data=[ + DocumentMetadataOperation( + document_id=document_id, + metadata_list=[MetadataDetail(id=metadata_id, name="spoofed", value="Alice")], + partial_update=partial, + ) + ] + ) + + +class TestMetadataServiceSessionBoundaryCases: + @staticmethod + @pytest.fixture + def metadata_service(sqlite_session_factory: sessionmaker[Session]) -> Iterator[MetadataService]: + with sqlite_session_factory.begin() as session: + account = make_account(account_id="actor", name="User", email="metadata@example.com") + session.add(account) + for ref, document_id in [(REF, DOCUMENT_ID), (DatasetRef("tenant-2", "dataset-2"), FOREIGN_ID)]: + session.add( + make_dataset(dataset_id=ref.dataset_id, tenant_id=ref.tenant_id, name="Dataset", created_by="actor") + ) + document = make_document( + document_id=document_id, + tenant_id=ref.tenant_id, + dataset_id=ref.dataset_id, + batch="batch", + created_from=DocumentCreatedFrom.API, + created_by="actor", + created_at=datetime(2026, 1, 1), + updated_at=datetime(2026, 1, 2), + ) + document.enabled = True + document.archived = False + document.indexing_status = "completed" + session.add(document) + return MetadataService( + store=SQLAlchemyMetadataRepository(session_factory=sqlite_session_factory), + dataset_access=create_autospec(DatasetAccess, instance=True), + ) + + @staticmethod + def test_create_commits_through_owned_session( + metadata_service: MetadataService, sqlite_session_factory: sessionmaker[Session] + ) -> None: + result = metadata_service.create_metadata(REF, MetadataArgs(type="string", name="author"), actor_id="actor") + fields = metadata_service.get_dataset_metadatas(REF) + assert fields["doc_metadata"] == [{"id": result.id, "type": "string", "name": "author", "count": 0}] + with sqlite_session_factory() as session: + assert session.scalar(select(DatasetMetadataBinding)) is None + + @staticmethod + @pytest.mark.parametrize( + "name", ["x" * 256, "document_name", "uploader", "upload_date", "last_update_date", "source"] + ) + def test_invalid_name_is_rejected(metadata_service: MetadataService, name: str) -> None: + with pytest.raises( + ValueError, + match=( + "(?:Metadata name already exists in Built\\-in fields\\.|Metadata name cannot exceed 255 characters\\.)" + ), + ): + metadata_service.create_metadata(REF, MetadataArgs(type="string", name=name), actor_id="actor") + assert metadata_service.get_dataset_metadatas(REF)["doc_metadata"] == [] + + @staticmethod + def test_duplicate_name_is_rejected(metadata_service: MetadataService) -> None: + metadata_service.create_metadata(REF, MetadataArgs(type="string", name="author"), actor_id="actor") + with pytest.raises(ValueError, match="already exists"): + metadata_service.create_metadata(REF, MetadataArgs(type="string", name="author"), actor_id="actor") + + @staticmethod + @pytest.mark.parametrize("partial", [True, False]) + def test_update_uses_canonical_names_and_preserves_partial_values( + metadata_service: MetadataService, sqlite_session_factory: sessionmaker[Session], partial: bool + ) -> None: + field = metadata_service.create_metadata(REF, MetadataArgs(type="string", name="author"), actor_id="actor") + with sqlite_session_factory.begin() as session: + document = session.get(Document, DOCUMENT_ID) + assert document is not None + document.doc_metadata = {"existing": "kept"} + metadata_service.update_documents_metadata( + REF, _operation(DOCUMENT_ID, field.id, partial=partial), actor_id="actor" + ) + metadata_service.update_documents_metadata( + REF, _operation(DOCUMENT_ID, field.id, partial=partial), actor_id="actor" + ) + with sqlite_session_factory() as session: + document = session.get(Document, DOCUMENT_ID) + assert document is not None + assert document.doc_metadata == ( + {"existing": "kept", "author": "Alice"} if partial else {"author": "Alice"} + ) + assert len(session.scalars(select(DatasetMetadataBinding)).all()) == 1 + + @staticmethod + @pytest.mark.parametrize("foreign", ["dataset", "document", "metadata"]) + def test_owner_chain_is_validated_before_any_batch_write( + metadata_service: MetadataService, sqlite_session_factory: sessionmaker[Session], foreign: str + ) -> None: + field = metadata_service.create_metadata(REF, MetadataArgs(type="string", name="author"), actor_id="actor") + other = metadata_service.create_metadata( + DatasetRef("tenant-2", "dataset-2"), MetadataArgs(type="string", name="foreign"), actor_id="actor" + ) + args = _operation(DOCUMENT_ID, field.id) + args.operation_data.extend( + _operation( + FOREIGN_ID if foreign == "document" else DOCUMENT_ID, other.id if foreign == "metadata" else field.id + ).operation_data + ) + ref = DatasetRef("tenant-2", REF.dataset_id) if foreign == "dataset" else REF + with pytest.raises(MetadataResourceNotFoundError): + metadata_service.update_documents_metadata(ref, args, actor_id="actor") + with sqlite_session_factory() as session: + assert all(document.doc_metadata is None for document in session.scalars(select(Document))) + assert not session.scalars(select(DatasetMetadataBinding)).all() + + @staticmethod + def test_rename_delete_and_counts_share_scoped_bindings( + metadata_service: MetadataService, sqlite_session_factory: sessionmaker[Session] + ) -> None: + field = metadata_service.create_metadata(REF, MetadataArgs(type="string", name="author"), actor_id="actor") + metadata_service.update_documents_metadata(REF, _operation(DOCUMENT_ID, field.id), actor_id="actor") + renamed = metadata_service.update_metadata_name(REF, field.id, "writer", actor_id="actor") + assert renamed.name == "writer" + assert metadata_service.get_dataset_metadatas(REF)["doc_metadata"] == [ + {"id": field.id, "name": "writer", "type": "string", "count": 1} + ] + with sqlite_session_factory() as session: + document = session.get(Document, DOCUMENT_ID) + assert document is not None + assert document.doc_metadata == {"writer": "Alice"} + metadata_service.delete_metadata(REF, field.id) + with sqlite_session_factory() as session: + document = session.get(Document, DOCUMENT_ID) + assert document is not None + assert document.doc_metadata == {} + assert not session.scalars(select(DatasetMetadataBinding)).all() + + @staticmethod + def test_builtin_fields_use_owned_document_uploader( + metadata_service: MetadataService, sqlite_session_factory: sessionmaker[Session] + ) -> None: + metadata_service.enable_built_in_field(REF) + with sqlite_session_factory() as session: + document = session.get(Document, DOCUMENT_ID) + foreign = session.get(Document, FOREIGN_ID) + assert document is not None + assert foreign is not None + assert document.doc_metadata["uploader"] == "User" + assert foreign.doc_metadata is None + metadata_service.disable_built_in_field(REF) + with sqlite_session_factory() as session: + document = session.get(Document, DOCUMENT_ID) + assert document is not None + assert document.doc_metadata == {} + + @staticmethod + def test_commit_failure_rolls_back_document_and_bindings( + metadata_service: MetadataService, sqlite_session_factory: sessionmaker[Session] + ) -> None: + field = metadata_service.create_metadata(REF, MetadataArgs(type="string", name="author"), actor_id="actor") + + def fail(_session: Session) -> None: + raise RuntimeError("commit failed") + + event.listen(sqlite_session_factory.class_, "before_commit", fail) + try: + with pytest.raises(RuntimeError, match="commit failed"): + metadata_service.update_documents_metadata(REF, _operation(DOCUMENT_ID, field.id), actor_id="actor") + finally: + event.remove(sqlite_session_factory.class_, "before_commit", fail) + with sqlite_session_factory() as session: + document = session.get(Document, DOCUMENT_ID) + assert document is not None + assert document.doc_metadata is None + assert not session.scalars(select(DatasetMetadataBinding)).all() + + +if __name__ == "__main__": + pytest.main([__file__, "-v"]) + + +if __name__ == "__main__": + pytest.main([__file__, "-v"]) diff --git a/api/tests/unit_tests/services/knowledge/test_retrieval_attachment_signing.py b/api/tests/unit_tests/services/knowledge/retrieval/test_attachments.py similarity index 100% rename from api/tests/unit_tests/services/knowledge/test_retrieval_attachment_signing.py rename to api/tests/unit_tests/services/knowledge/retrieval/test_attachments.py diff --git a/api/tests/unit_tests/services/knowledge/test_segment_adapters.py b/api/tests/unit_tests/services/knowledge/segments/test_adapters.py similarity index 99% rename from api/tests/unit_tests/services/knowledge/test_segment_adapters.py rename to api/tests/unit_tests/services/knowledge/segments/test_adapters.py index 427b76127d2..2cc52e817c7 100644 --- a/api/tests/unit_tests/services/knowledge/test_segment_adapters.py +++ b/api/tests/unit_tests/services/knowledge/segments/test_adapters.py @@ -1197,7 +1197,7 @@ def test_economy_index_shares_persistence_with_existing_keyword_calls( from core.rag.datasource.keyword.jieba import jieba as jieba_module from core.rag.datasource.keyword.keyword_factory import Keyword from repositories.knowledge import keyword_table_repository as table_module - from tests.unit_tests.config_override import apply_config_overrides + from tests.helpers.unit.config_override import apply_config_overrides app, repository, probe = indexing_probe reads: list[str] = [] @@ -1737,7 +1737,7 @@ def test_batch_index_keeps_manual_keywords_aligned_with_segments( monkeypatch: pytest.MonkeyPatch, ) -> None: from core.rag.datasource.keyword.jieba import jieba as jieba_module - from tests.unit_tests.config_override import apply_config_overrides + from tests.helpers.unit.config_override import apply_config_overrides app, _, probe = indexing_probe _persist_index_batch(sqlite_session_factory) diff --git a/api/tests/unit_tests/services/knowledge/test_segment_application_service.py b/api/tests/unit_tests/services/knowledge/segments/test_application.py similarity index 100% rename from api/tests/unit_tests/services/knowledge/test_segment_application_service.py rename to api/tests/unit_tests/services/knowledge/segments/test_application.py diff --git a/api/tests/unit_tests/services/test_summary_index_service.py b/api/tests/unit_tests/services/knowledge/summaries/test_adapters.py similarity index 100% rename from api/tests/unit_tests/services/test_summary_index_service.py rename to api/tests/unit_tests/services/knowledge/summaries/test_adapters.py diff --git a/api/tests/unit_tests/services/knowledge/test_dataset_access_service.py b/api/tests/unit_tests/services/knowledge/test_dataset_access.py similarity index 100% rename from api/tests/unit_tests/services/knowledge/test_dataset_access_service.py rename to api/tests/unit_tests/services/knowledge/test_dataset_access.py diff --git a/api/tests/unit_tests/controllers/console/test_document_detail_api_data_source_info.py b/api/tests/unit_tests/services/knowledge/test_dataset_read_service.py similarity index 92% rename from api/tests/unit_tests/controllers/console/test_document_detail_api_data_source_info.py rename to api/tests/unit_tests/services/knowledge/test_dataset_read_service.py index 4b4c881c53b..37f9e6364c8 100644 --- a/api/tests/unit_tests/controllers/console/test_document_detail_api_data_source_info.py +++ b/api/tests/unit_tests/services/knowledge/test_dataset_read_service.py @@ -12,6 +12,12 @@ from unittest.mock import MagicMock from models.dataset import Document from models.enums import DataSourceType, DocumentCreatedFrom from services.knowledge.dataset_read_service import get_document_source_detail +from tests.unit_tests.controllers.conftest import app_query_services as app_query_services # noqa: PLC0414 +from tests.unit_tests.controllers.conftest import ( + datasource_application_dependencies as datasource_application_dependencies, # noqa: PLC0414 +) +from tests.unit_tests.controllers.conftest import error_response_contract as error_response_contract # noqa: PLC0414 +from tests.unit_tests.controllers.conftest import pipeline_application as pipeline_application # noqa: PLC0414 class LocalFileInfo(TypedDict): diff --git a/api/tests/unit_tests/services/knowledge/test_dataset_service.py b/api/tests/unit_tests/services/knowledge/test_dataset_service.py new file mode 100644 index 00000000000..b1c5ff0c181 --- /dev/null +++ b/api/tests/unit_tests/services/knowledge/test_dataset_service.py @@ -0,0 +1,4537 @@ +from __future__ import annotations + +import inspect +import time +import types +from collections.abc import Callable, Iterator +from datetime import datetime +from types import SimpleNamespace +from unittest.mock import MagicMock as _dataset_service_segment_MagicMock +from unittest.mock import create_autospec +from unittest.mock import patch as _dataset_service_dataset_patch +from unittest.mock import patch as _dataset_service_segment_patch + +import pytest +from redis.exceptions import LockNotOwnedError +from sqlalchemy import event, func, select +from sqlalchemy.orm import Session, sessionmaker + +from core.errors.error import LLMBadRequestError, ProviderTokenNotInitError +from core.model_manager import ModelInstance +from core.rag.index_processor.constant.built_in_field import BuiltInField as _metadata_built_in_fields_BuiltInField +from core.rag.index_processor.constant.index_type import ( + IndexStructureType as _dataset_service_lock_not_owned_IndexStructureType, +) +from core.rag.index_processor.constant.index_type import ( + IndexStructureType as _dataset_service_segment_IndexStructureType, +) +from core.rag.index_processor.constant.index_type import IndexTechniqueType +from core.rag.retrieval.dataset_retrieval import DatasetRetrieval +from enums import DeploymentEdition +from graphon.model_runtime.entities.model_entities import ModelFeature, ModelType +from models import Account as _dataset_service_dataset_Account +from models import Account as _dataset_service_segment_Account +from models.account import Account as _dataset_service_lock_not_owned_Account +from models.account import Tenant +from models.dataset import ( + ChildChunk, + Dataset, + DatasetCollectionBinding, + DatasetPermission, + DatasetPermissionEnum, + DocumentSegment, + ExternalKnowledgeApis, + ExternalKnowledgeBindings, + Pipeline, + build_built_in_field_values, +) +from models.dataset import DatasetProcessRule as _dataset_service_lock_not_owned_DatasetProcessRule +from models.dataset import Document as _dataset_service_lock_not_owned_Document +from models.dataset import ( + Document as _dataset_service_segment_Document, +) +from models.dataset import Document as _metadata_built_in_fields_Document +from models.enums import DataSourceType, DocumentCreatedFrom, IndexingStatus, ProcessRuleMode, SegmentStatus +from models.model import UploadFile +from models.source import DataSourceOauthBinding +from repositories.knowledge.segment_repository import SQLAlchemySegmentRepository +from services.enterprise import rbac_service +from services.entities.knowledge_entities.rag_pipeline_entities import ( + IconInfo, + RagPipelineDatasetCreateEntity, +) +from services.errors.base import NoPermissionError +from services.errors.dataset import DatasetNameDuplicateError +from services.knowledge import dataset_service +from services.knowledge.dataset_service import DatasetCollectionBindingService, DatasetPermissionService, SegmentService +from services.knowledge.dataset_service import DatasetService as _dataset_service_dataset_DatasetService +from services.knowledge.dataset_service import DocumentService as _dataset_service_lock_not_owned_DocumentService +from services.knowledge.dataset_service import DocumentService as _metadata_built_in_fields_DocumentService +from services.knowledge.dataset_service import ( + DocumentService as _save_document_with_dataset_id_mysql_snapshot_DocumentService, +) +from services.knowledge.entities.knowledge_entities import DataSource as _dataset_service_lock_not_owned_DataSource +from services.knowledge.entities.knowledge_entities import InfoList as _dataset_service_lock_not_owned_InfoList +from services.knowledge.entities.knowledge_entities import ( + KnowledgeConfig as _dataset_service_lock_not_owned_KnowledgeConfig, +) +from services.knowledge.resource_scope import DatasetRef, DocumentRef, SegmentRef +from services.knowledge.segments.application import SegmentIndex, SegmentIndexingState, SegmentMutationService +from tests.helpers.unit.config_override import config_overrides_context +from tests.helpers.unit.model_factories import make_account, make_dataset, make_document, make_tenant, make_upload_file +from tests.helpers.unit.services.dataset_service_test_helpers import Account as _dataset_service_document_Account +from tests.helpers.unit.services.dataset_service_test_helpers import ( + BuiltInField as _dataset_service_document_BuiltInField, +) +from tests.helpers.unit.services.dataset_service_test_helpers import ( + CloudPlan, + DocumentIndexingError, + FileInfo, + FileNotExistsError, + NotFound, + NotionIcon, + NotionInfo, + NotionPage, + PreProcessingRule, + ProcessRule, + RerankingModel, + RetrievalMethod, + RetrievalModel, + Rule, + Segmentation, + TenantAccountRole, + WebsiteInfo, + _make_features, + _make_knowledge_configuration, + _make_lock_context, + _make_retrieval_model, + _make_upload_knowledge_config, + json, +) +from tests.helpers.unit.services.dataset_service_test_helpers import ( + DatasetProcessRule as _dataset_service_document_DatasetProcessRule, +) +from tests.helpers.unit.services.dataset_service_test_helpers import ( + DatasetService as _dataset_service_document_DatasetService, +) +from tests.helpers.unit.services.dataset_service_test_helpers import DataSource as _dataset_service_document_DataSource +from tests.helpers.unit.services.dataset_service_test_helpers import Document as _dataset_service_document_Document +from tests.helpers.unit.services.dataset_service_test_helpers import ( + DocumentService as _dataset_service_document_DocumentService, +) +from tests.helpers.unit.services.dataset_service_test_helpers import ( + IndexStructureType as _dataset_service_document_IndexStructureType, +) +from tests.helpers.unit.services.dataset_service_test_helpers import InfoList as _dataset_service_document_InfoList +from tests.helpers.unit.services.dataset_service_test_helpers import ( + KnowledgeConfig as _dataset_service_document_KnowledgeConfig, +) +from tests.helpers.unit.services.dataset_service_test_helpers import MagicMock as _dataset_service_dataset_MagicMock +from tests.helpers.unit.services.dataset_service_test_helpers import patch as _dataset_service_document_patch + +"""Regression coverage for #41735. + +The pre-fix code in ``DocumentService.save_document_with_dataset_id`` +read ``UploadFile`` rows inside a single session that had its +``REPEATABLE READ`` snapshot pinned by an earlier ``Dataset`` SELECT. +When ``FileService.upload_text`` / ``upload_file`` committed the +``upload_files`` row in a SEPARATE session, the row was invisible to +*this* session until the snapshot was refreshed, and +``save_document_with_dataset_id`` reported ``One or more files not +found.`` even though the row was committed (verified via a second +session). + +The fix forces ``session.commit()`` immediately before the +``UploadFile`` lookup so the next SELECT starts a fresh snapshot. +This file pins that contract by reading the function source and +asserting the structural invariant directly. +""" + + +class TestSaveDocumentWithDatasetIdMysqlSnapshotCases: + @staticmethod + def test_session_committed_before_upload_file_lookup() -> None: + """#41735: ``session.commit()`` must run BEFORE the ``UploadFile`` + SELECT so the SELECT sees rows committed by ``FileService.upload_text`` + on a different session. MySQL REPEATABLE READ keeps the original + snapshot for the rest of the transaction; the commit ends it so the + next SELECT starts a fresh transaction with a new snapshot. + PostgreSQL's default READ COMMITTED rebuilds the read view per + statement so the extra commit is a no-op there. + + Read the function source and assert that ``session.commit()`` is + called BEFORE the ``UploadFile`` SELECT — the same structural + invariant the reproduction in the issue depends on, without needing + a live MySQL database. + """ + source_lines = inspect.getsource( + _save_document_with_dataset_id_mysql_snapshot_DocumentService.save_prepared_documents + ).splitlines() + + upload_select_line = next( + (i for i, line in enumerate(source_lines) if "select(UploadFile)" in line), + None, + ) + assert upload_select_line is not None, "save_prepared_documents must contain a select(UploadFile) call" + + # Find every ``session.commit()`` call, then keep the last one + # whose line number is BEFORE the UploadFile SELECT. If the fix is + # in place, that commit is the snapshot refresh. + last_commit_before_upload = -1 + for i, line in enumerate(source_lines): + if i >= upload_select_line: + break + if "session.commit()" in line: + last_commit_before_upload = i + + assert last_commit_before_upload >= 0, ( + "save_prepared_documents must call session.commit() at least once before the UploadFile SELECT" + ) + + # The fix is a focused call: one extra ``session.commit()`` between + # the (already-existing) ``DocumentService.check_doc_form`` call and + # the ``select(UploadFile)`` call. Pin that there's a comment + # referencing #41735 right above the new commit so a future edit + # doesn't accidentally delete the snapshot refresh without + # understanding why it's there. + assert "#41735" in "\n".join(source_lines), ( + "The fix's invariant comment referencing issue #41735 must be present in save_prepared_documents" + ) + + +CREATED_AT = datetime(2026, 10, 16, 8, 0, 0) +CREATED_AT_EPOCH = 1792137600 + + +FRACTIONAL_CREATED_AT = datetime(2026, 10, 16, 8, 0, 0, 900000) + + +class TestMetadataBuiltInFieldsCases: + @staticmethod + @pytest.fixture + def non_utc_host_timezone(monkeypatch: pytest.MonkeyPatch) -> Iterator[None]: + monkeypatch.setenv("TZ", "Asia/Tokyo") + time.tzset() + yield + monkeypatch.undo() + time.tzset() + + @staticmethod + @pytest.mark.usefixtures("non_utc_host_timezone") + def test_build_built_in_field_values_treats_naive_datetimes_as_utc() -> None: + values = build_built_in_field_values( + document_name="doc.txt", + uploader="Alice", + upload_date=CREATED_AT, + last_update_date=CREATED_AT, + data_source_type=DataSourceType.UPLOAD_FILE, + ) + + assert values == { + _metadata_built_in_fields_BuiltInField.document_name: "doc.txt", + _metadata_built_in_fields_BuiltInField.uploader: "Alice", + _metadata_built_in_fields_BuiltInField.upload_date: CREATED_AT_EPOCH, + _metadata_built_in_fields_BuiltInField.last_update_date: CREATED_AT_EPOCH, + _metadata_built_in_fields_BuiltInField.source: "file_upload", + } + + @staticmethod + @pytest.mark.parametrize( + ("data_source_type", "source"), + [ + (DataSourceType.UPLOAD_FILE, "file_upload"), + (DataSourceType.LOCAL_FILE, "file_upload"), + (DataSourceType.NOTION_IMPORT, "notion"), + (DataSourceType.WEBSITE_CRAWL, "website"), + (DataSourceType.ONLINE_DOCUMENT, "online_document"), + (DataSourceType.ONLINE_DRIVE, "online_drive"), + ], + ) + def test_build_built_in_field_values_maps_source_to_metadata_label( + data_source_type: DataSourceType, source: str + ) -> None: + values = build_built_in_field_values( + document_name="doc.txt", + uploader=None, + upload_date=CREATED_AT, + last_update_date=CREATED_AT, + data_source_type=data_source_type, + ) + + assert values[_metadata_built_in_fields_BuiltInField.source] == source + + @staticmethod + @pytest.mark.parametrize( + "data_source_type", + [DataSourceType.UPLOAD_FILE, DataSourceType.NOTION_IMPORT, DataSourceType.WEBSITE_CRAWL], + ) + def test_build_document_writes_built_in_fields_in_shared_format( + monkeypatch: pytest.MonkeyPatch, data_source_type: DataSourceType + ) -> None: + monkeypatch.setattr(dataset_service, "naive_utc_now", lambda: CREATED_AT) + dataset = make_dataset(built_in_field_enabled=True) + + document = _metadata_built_in_fields_DocumentService.build_document( + dataset=dataset, + process_rule_id=None, + data_source_type=data_source_type, + document_form="text_model", + document_language="English", + data_source_info={}, + created_from="web", + position=1, + account=make_account(name="Alice"), + name="doc.txt", + batch="batch-1", + ) + + assert document.doc_metadata == build_built_in_field_values( + document_name="doc.txt", + uploader="Alice", + upload_date=CREATED_AT, + last_update_date=CREATED_AT, + data_source_type=data_source_type, + ) + + @staticmethod + def test_build_document_skips_built_in_fields_when_disabled() -> None: + dataset = make_dataset(built_in_field_enabled=False) + + document = _metadata_built_in_fields_DocumentService.build_document( + dataset=dataset, + process_rule_id=None, + data_source_type=DataSourceType.UPLOAD_FILE, + document_form="text_model", + document_language="English", + data_source_info={}, + created_from="web", + position=1, + account=make_account(name="Alice"), + name="doc.txt", + batch="batch-1", + ) + + assert not document.doc_metadata + + @staticmethod + def test_build_built_in_field_values_keeps_fractional_epoch_seconds() -> None: + values = build_built_in_field_values( + document_name="doc.txt", + uploader=None, + upload_date=FRACTIONAL_CREATED_AT, + last_update_date=FRACTIONAL_CREATED_AT, + data_source_type=DataSourceType.UPLOAD_FILE, + ) + + assert values[_metadata_built_in_fields_BuiltInField.upload_date] == pytest.approx( + CREATED_AT_EPOCH + 0.9, abs=1e-6 + ) + assert values[_metadata_built_in_fields_BuiltInField.last_update_date] == pytest.approx( + CREATED_AT_EPOCH + 0.9, abs=1e-6 + ) + + @staticmethod + @pytest.mark.parametrize( + "field", + [_metadata_built_in_fields_BuiltInField.upload_date, _metadata_built_in_fields_BuiltInField.last_update_date], + ) + @pytest.mark.parametrize(("condition", "matches"), [("after", True), (">", True), ("=", False), ("is", False)]) + def test_fractional_built_in_dates_keep_integer_second_boundary_semantics( + sqlite_session_factory: sessionmaker[Session], + field: _metadata_built_in_fields_BuiltInField, + condition: str, + matches: bool, + ) -> None: + values = build_built_in_field_values( + document_name="doc.txt", + uploader=None, + upload_date=FRACTIONAL_CREATED_AT, + last_update_date=FRACTIONAL_CREATED_AT, + data_source_type=DataSourceType.UPLOAD_FILE, + ) + doc_metadata: dict[str, object] = dict(values) + with sqlite_session_factory.begin() as session: + session.add(make_document(document_id="document-1", doc_metadata=doc_metadata)) + + filters = DatasetRetrieval.process_metadata_filter_func(0, condition, field, CREATED_AT_EPOCH, []) + with sqlite_session_factory() as session: + found = session.scalars(select(_metadata_built_in_fields_Document.id).where(*filters)).all() + + assert found == (["document-1"] if matches else []) + + +"""SQLite-backed tests for dataset-level operations in :mod:`services.knowledge.dataset_service`. + +Mapped objects in this module are real SQLAlchemy models. Provider runtimes, +RBAC clients, Celery tasks, and model-manager results remain mocked at their +external boundaries. +""" + + +def _dataset_service_dataset_account( + *, + account_id: str = "user-1", + tenant_id: str = "tenant-1", + role: TenantAccountRole = TenantAccountRole.OWNER, +) -> _dataset_service_dataset_Account: + return make_account( + account_id=account_id, + name=f"User {account_id}", + email=f"{account_id}@example.com", + role=role, + tenant=make_tenant(tenant_id=tenant_id, name=f"Tenant {tenant_id}"), + ) + + +def _dataset_service_dataset_dataset( + *, + dataset_id: str = "dataset-1", + tenant_id: str = "tenant-1", + name: str = "Dataset", + maintainer: str = "user-1", + permission: DatasetPermissionEnum = DatasetPermissionEnum.ALL_TEAM, + provider: str = "vendor", + indexing_technique: str = IndexTechniqueType.ECONOMY, + chunk_structure: str | None = "text_model", +) -> Dataset: + return make_dataset( + dataset_id=dataset_id, + tenant_id=tenant_id, + name=name, + description="", + provider=provider, + created_by=maintainer, + maintainer=maintainer, + permission=permission, + indexing_technique=indexing_technique, + chunk_structure=chunk_structure, + embedding_model_provider="provider", + embedding_model="embedding-model", + ) + + +def _external_api(*, api_id: str = "api-1", tenant_id: str = "tenant-1") -> ExternalKnowledgeApis: + api = ExternalKnowledgeApis( + name="External API", + description="", + tenant_id=tenant_id, + settings="{}", + created_by="user-1", + updated_by="user-1", + ) + api.id = api_id + return api + + +def _binding( + *, + binding_id: str = "binding-1", + tenant_id: str = "tenant-1", + dataset_id: str = "dataset-1", + api_id: str = "api-1", + knowledge_id: str = "knowledge-1", +) -> ExternalKnowledgeBindings: + binding = ExternalKnowledgeBindings( + tenant_id=tenant_id, + dataset_id=dataset_id, + external_knowledge_api_id=api_id, + external_knowledge_id=knowledge_id, + created_by="user-1", + ) + binding.id = binding_id + return binding + + +class TestDatasetServiceDatasetCases: + class TestDatasetServiceValidation: + @pytest.mark.parametrize( + ("dataset_doc_form", "incoming_doc_form"), + [(None, "text_model"), ("text_model", "text_model")], + ) + def test_check_doc_form_allows_matching_or_missing_dataset_doc_form( + self, + sqlite_session: Session, + dataset_doc_form: str | None, + incoming_doc_form: str, + ) -> None: + dataset = _dataset_service_dataset_dataset(chunk_structure=dataset_doc_form) + sqlite_session.add(dataset) + sqlite_session.commit() + + _dataset_service_dataset_DatasetService.check_doc_form(dataset, incoming_doc_form, session=sqlite_session) + + def test_check_doc_form_rejects_mismatched_doc_form(self, sqlite_session: Session) -> None: + dataset = _dataset_service_dataset_dataset(chunk_structure="qa_model") + sqlite_session.add(dataset) + sqlite_session.commit() + + with pytest.raises(ValueError, match="doc_form is different"): + _dataset_service_dataset_DatasetService.check_doc_form(dataset, "text_model", session=sqlite_session) + + @pytest.mark.parametrize("operator_check", [False, True]) + def test_dataset_permission_checks_ignore_foreign_tenant_binding( + self, sqlite_session: Session, operator_check: bool + ) -> None: + dataset = _dataset_service_dataset_dataset( + dataset_id="dataset-1", + tenant_id="tenant-1", + permission=DatasetPermissionEnum.PARTIAL_TEAM, + maintainer="owner-1", + ) + user = _dataset_service_dataset_account( + account_id="user-1", + tenant_id="tenant-1", + role=TenantAccountRole.NORMAL, + ) + sqlite_session.add_all( + [ + dataset, + DatasetPermission(dataset_id=dataset.id, account_id=user.id, tenant_id="tenant-2"), + ] + ) + sqlite_session.commit() + + if operator_check: + with pytest.raises(NoPermissionError): + _dataset_service_dataset_DatasetService.check_dataset_operator_permission( + user, dataset, session=sqlite_session + ) + else: + with pytest.raises(NoPermissionError): + _dataset_service_dataset_DatasetService.check_dataset_permission(dataset, user, sqlite_session) + + def test_check_dataset_model_setting_skips_non_high_quality_datasets(self) -> None: + dataset = _dataset_service_dataset_dataset(indexing_technique=IndexTechniqueType.ECONOMY) + + with _dataset_service_dataset_patch("services.knowledge.dataset_service.ModelManager") as model_manager_cls: + _dataset_service_dataset_DatasetService.check_dataset_model_setting(dataset) + + model_manager_cls.assert_not_called() + + def test_check_dataset_model_setting_validates_high_quality_embedding(self) -> None: + dataset = _dataset_service_dataset_dataset(indexing_technique=IndexTechniqueType.HIGH_QUALITY) + + with _dataset_service_dataset_patch("services.knowledge.dataset_service.ModelManager") as model_manager_cls: + _dataset_service_dataset_DatasetService.check_dataset_model_setting(dataset) + + model_manager_cls.for_tenant.return_value.get_model_instance.assert_called_once_with( + tenant_id=dataset.tenant_id, + provider=dataset.embedding_model_provider, + model_type=ModelType.TEXT_EMBEDDING, + model=dataset.embedding_model, + ) + + @pytest.mark.parametrize( + ("error", "message"), + [ + (LLMBadRequestError(), "No Embedding Model available"), + (ProviderTokenNotInitError("token missing"), "token missing"), + ], + ) + def test_check_dataset_model_setting_wraps_provider_errors(self, error: Exception, message: str) -> None: + dataset = _dataset_service_dataset_dataset(indexing_technique=IndexTechniqueType.HIGH_QUALITY) + + with _dataset_service_dataset_patch("services.knowledge.dataset_service.ModelManager") as model_manager_cls: + model_manager_cls.for_tenant.return_value.get_model_instance.side_effect = error + with pytest.raises(ValueError, match=message): + _dataset_service_dataset_DatasetService.check_dataset_model_setting(dataset) + + @pytest.mark.parametrize(("features", "expected"), [([ModelFeature.VISION], True), ([], False)]) + def test_check_is_multimodal_model_reads_runtime_schema( + self, features: list[ModelFeature], expected: bool + ) -> None: + model_type_instance = _dataset_service_dataset_MagicMock() + model_type_instance.get_model_schema.return_value = SimpleNamespace(features=features) + model_instance = SimpleNamespace( + model_type_instance=model_type_instance, + model_name="embedding-model", + credentials={"api_key": "secret"}, + ) + + with _dataset_service_dataset_patch("services.knowledge.dataset_service.ModelManager") as model_manager_cls: + model_manager_cls.for_tenant.return_value.get_model_instance.return_value = model_instance + result = _dataset_service_dataset_DatasetService.check_is_multimodal_model( + "tenant-1", "provider", "embedding-model" + ) + + assert result is expected + + def test_check_is_multimodal_model_rejects_missing_schema(self) -> None: + model_type_instance = _dataset_service_dataset_MagicMock() + model_type_instance.get_model_schema.return_value = None + model_instance = SimpleNamespace( + model_type_instance=model_type_instance, + model_name="embedding-model", + credentials={}, + ) + + with _dataset_service_dataset_patch("services.knowledge.dataset_service.ModelManager") as model_manager_cls: + model_manager_cls.for_tenant.return_value.get_model_instance.return_value = model_instance + with pytest.raises(ValueError, match="Model schema not found"): + _dataset_service_dataset_DatasetService.check_is_multimodal_model( + "tenant-1", "provider", "embedding-model" + ) + + @pytest.mark.parametrize( + ("method", "error", "message"), + [ + ( + _dataset_service_dataset_DatasetService.check_embedding_model_setting, + ProviderTokenNotInitError("provider setup"), + "provider setup", + ), + ( + _dataset_service_dataset_DatasetService.check_reranking_model_setting, + LLMBadRequestError(), + "No Rerank Model available", + ), + ], + ) + def test_direct_model_setting_checks_wrap_runtime_errors( + self, method: Callable[[str, str, str], None], error: Exception, message: str + ) -> None: + with _dataset_service_dataset_patch("services.knowledge.dataset_service.ModelManager") as model_manager_cls: + model_manager_cls.for_tenant.return_value.get_model_instance.side_effect = error + with pytest.raises(ValueError, match=message): + method("tenant-1", "provider", "model") + + class TestDatasetServiceRetrieval: + def test_get_dataset_for_tenant_rejects_cross_tenant_row(self, sqlite_session: Session) -> None: + owned = _dataset_service_dataset_dataset() + foreign = _dataset_service_dataset_dataset(dataset_id="dataset-2", tenant_id="tenant-2", name="Foreign") + sqlite_session.add_all([owned, foreign]) + sqlite_session.commit() + + assert ( + _dataset_service_dataset_DatasetService.get_dataset_for_tenant( + owned.id, "tenant-1", session=sqlite_session + ) + is owned + ) + assert ( + _dataset_service_dataset_DatasetService.get_dataset_for_tenant( + foreign.id, "tenant-1", session=sqlite_session + ) + is None + ) + + def test_get_datasets_filters_by_creator_ids( + self, application_tags, config_overrides: Callable[..., None], sqlite_session: Session + ) -> None: + creator_one = _dataset_service_dataset_dataset( + dataset_id="creator-one", name="Creator One", maintainer="creator-1" + ) + creator_two = _dataset_service_dataset_dataset( + dataset_id="creator-two", name="Creator Two", maintainer="creator-2" + ) + excluded = _dataset_service_dataset_dataset(dataset_id="excluded", name="Excluded", maintainer="creator-3") + foreign = _dataset_service_dataset_dataset( + dataset_id="foreign", + tenant_id="tenant-2", + name="Foreign", + maintainer="creator-1", + ) + sqlite_session.add_all([creator_one, creator_two, excluded, foreign]) + sqlite_session.commit() + + config_overrides(RBAC_ENABLED=False) + datasets, total = _dataset_service_dataset_DatasetService.get_datasets( + page=1, + per_page=20, + session=sqlite_session, + tenant_id="tenant-1", + creator_ids=["creator-1", "creator-2"], + tags=application_tags, + ) + + assert total == 2 + assert {dataset.id for dataset in datasets} == {creator_one.id, creator_two.id} + + def test_get_datasets_applies_rbac_resource_scope_and_maintainer_override( + self, application_tags, config_overrides: Callable[..., None], sqlite_session: Session + ) -> None: + user = _dataset_service_dataset_account(role=TenantAccountRole.NORMAL) + accessible = _dataset_service_dataset_dataset( + dataset_id="accessible", name="Accessible", maintainer="other" + ) + owned = _dataset_service_dataset_dataset(dataset_id="owned", name="Owned", maintainer=user.id) + hidden = _dataset_service_dataset_dataset(dataset_id="hidden", name="Hidden", maintainer="other") + foreign = _dataset_service_dataset_dataset(dataset_id="foreign", tenant_id="tenant-2", name="Foreign") + sqlite_session.add_all([accessible, owned, hidden, foreign]) + sqlite_session.commit() + + config_overrides(RBAC_ENABLED=True) + + with ( + _dataset_service_dataset_patch( + "services.knowledge.dataset_service.enterprise_rbac_service.RBACService.MyPermissions.get", + return_value=SimpleNamespace(workspace=SimpleNamespace(permission_keys=[])), + ), + ): + datasets, total = _dataset_service_dataset_DatasetService.get_datasets( + page=1, + per_page=20, + session=sqlite_session, + tenant_id="tenant-1", + user=user, + accessible_dataset_ids=[accessible.id], + include_own_datasets=True, + tags=application_tags, + ) + + assert total == 2 + assert {dataset.id for dataset in datasets} == {accessible.id, owned.id} + + def test_get_datasets_without_user_keeps_only_team_visible_rows( + self, application_tags, config_overrides: Callable[..., None], sqlite_session: Session + ) -> None: + shared = _dataset_service_dataset_dataset( + dataset_id="shared", name="Shared", permission=DatasetPermissionEnum.ALL_TEAM + ) + private = _dataset_service_dataset_dataset( + dataset_id="private", name="Private", permission=DatasetPermissionEnum.ONLY_ME + ) + sqlite_session.add_all([shared, private]) + sqlite_session.commit() + + config_overrides(RBAC_ENABLED=False) + datasets, total = _dataset_service_dataset_DatasetService.get_datasets( + page=1, per_page=20, session=sqlite_session, tenant_id="tenant-1", tags=application_tags + ) + + assert total == 1 + assert [dataset.id for dataset in datasets] == [shared.id] + + def test_get_datasets_by_ids_intersects_requested_and_accessible_ids( + self, config_overrides: Callable[..., None], sqlite_session: Session + ) -> None: + user = _dataset_service_dataset_account(role=TenantAccountRole.NORMAL) + accessible = _dataset_service_dataset_dataset( + dataset_id="accessible", name="Accessible", maintainer="other" + ) + owned = _dataset_service_dataset_dataset(dataset_id="owned", name="Owned", maintainer=user.id) + hidden = _dataset_service_dataset_dataset(dataset_id="hidden", name="Hidden", maintainer="other") + sqlite_session.add_all([accessible, owned, hidden]) + sqlite_session.commit() + + config_overrides(RBAC_ENABLED=True) + datasets, total = _dataset_service_dataset_DatasetService.get_datasets_by_ids( + [accessible.id, owned.id, hidden.id], + "tenant-1", + user=user, + accessible_dataset_ids=[accessible.id, "not-requested"], + include_own_datasets=True, + session=sqlite_session, + ) + + assert total == 2 + assert {dataset.id for dataset in datasets} == {accessible.id, owned.id} + + def test_get_datasets_rbac_without_user_returns_no_rows( + self, application_tags, config_overrides: Callable[..., None], sqlite_session: Session + ) -> None: + sqlite_session.add(_dataset_service_dataset_dataset()) + sqlite_session.commit() + + config_overrides(RBAC_ENABLED=True) + datasets, total = _dataset_service_dataset_DatasetService.get_datasets( + page=1, per_page=20, session=sqlite_session, tenant_id="tenant-1", tags=application_tags + ) + + assert datasets == [] + assert total == 0 + + def test_get_datasets_rbac_include_all_requires_workspace_permission( + self, application_tags, config_overrides: Callable[..., None], sqlite_session: Session + ) -> None: + user = _dataset_service_dataset_account(role=TenantAccountRole.NORMAL) + sqlite_session.add_all( + [ + _dataset_service_dataset_dataset(dataset_id="one", name="One", maintainer="other"), + _dataset_service_dataset_dataset(dataset_id="two", name="Two", maintainer="other"), + ] + ) + sqlite_session.commit() + + config_overrides(RBAC_ENABLED=True) + with ( + _dataset_service_dataset_patch( + "services.knowledge.dataset_service.enterprise_rbac_service.RBACService.MyPermissions.get", + return_value=SimpleNamespace( + workspace=SimpleNamespace(permission_keys=["dataset.create_and_management"]) + ), + ), + ): + datasets, total = _dataset_service_dataset_DatasetService.get_datasets( + page=1, + per_page=20, + session=sqlite_session, + tenant_id="tenant-1", + user=user, + include_all=True, + tags=application_tags, + ) + + assert total == 2 + assert {dataset.id for dataset in datasets} == {"one", "two"} + + class TestDatasetServiceCreationAndUpdate: + def test_create_empty_dataset_rejects_duplicate_name(self, sqlite_session: Session) -> None: + sqlite_session.add(_dataset_service_dataset_dataset(name="Existing")) + sqlite_session.commit() + + with pytest.raises(DatasetNameDuplicateError, match="already exists"): + _dataset_service_dataset_DatasetService.create_empty_dataset( + "tenant-1", + "Existing", + "", + IndexTechniqueType.ECONOMY, + _dataset_service_dataset_account(), + session=sqlite_session, + ) + + def test_create_empty_dataset_persists_default_embedding_and_creator_scope( + self, + sqlite_session: Session, + sqlite_session_factory: sessionmaker[Session], + ) -> None: + embedding_model = SimpleNamespace(provider="provider", model_name="default-embedding") + + with ( + _dataset_service_dataset_patch("services.knowledge.dataset_service.ModelManager") as model_manager_cls, + _dataset_service_dataset_patch( + "services.knowledge.dataset_service.enterprise_rbac_service.try_sync_creator_access_policy_member_bindings" + ), + ): + model_manager_cls.for_tenant.return_value.get_default_model_instance.return_value = embedding_model + dataset = _dataset_service_dataset_DatasetService.create_empty_dataset( + "tenant-1", + "Created", + "Description", + IndexTechniqueType.HIGH_QUALITY, + _dataset_service_dataset_account(), + session=sqlite_session, + ) + + with sqlite_session_factory() as observer: + persisted = observer.get(Dataset, dataset.id) + assert persisted is not None + assert persisted.tenant_id == "tenant-1" + assert persisted.maintainer == "user-1" + assert persisted.embedding_model == "default-embedding" + + def test_create_empty_external_dataset_persists_tenant_owned_binding(self, sqlite_session: Session) -> None: + sqlite_session.add_all([_external_api(), _external_api(api_id="api-foreign", tenant_id="tenant-2")]) + sqlite_session.commit() + + with _dataset_service_dataset_patch( + "services.knowledge.dataset_service.enterprise_rbac_service.try_sync_creator_access_policy_member_bindings" + ): + dataset = _dataset_service_dataset_DatasetService.create_empty_dataset( + "tenant-1", + "External", + "", + IndexTechniqueType.ECONOMY, + _dataset_service_dataset_account(), + provider="external", + external_knowledge_api_id="api-1", + external_knowledge_id="knowledge-1", + retrieval_model=_make_retrieval_model(), + session=sqlite_session, + ) + + binding = sqlite_session.scalar( + select(ExternalKnowledgeBindings).where(ExternalKnowledgeBindings.dataset_id == dataset.id) + ) + assert binding is not None + assert binding.tenant_id == dataset.tenant_id + assert binding.external_knowledge_api_id == "api-1" + + def test_create_empty_external_dataset_rejects_foreign_api(self, sqlite_session: Session) -> None: + sqlite_session.add(_external_api(api_id="api-foreign", tenant_id="tenant-2")) + sqlite_session.commit() + + with ( + pytest.raises(ValueError, match="api template not found"), + _dataset_service_dataset_patch( + "services.knowledge.dataset_service.enterprise_rbac_service.try_sync_creator_access_policy_member_bindings" + ), + ): + _dataset_service_dataset_DatasetService.create_empty_dataset( + "tenant-1", + "External", + "", + IndexTechniqueType.ECONOMY, + _dataset_service_dataset_account(), + provider="external", + external_knowledge_api_id="api-foreign", + external_knowledge_id="knowledge-1", + session=sqlite_session, + ) + + def test_create_empty_rag_pipeline_dataset_generates_incremental_name(self, sqlite_session: Session) -> None: + sqlite_session.add(_dataset_service_dataset_dataset(name="Untitled 1")) + sqlite_session.commit() + entity = RagPipelineDatasetCreateEntity( + name="", + description="Pipeline dataset", + icon_info=IconInfo(icon_type="emoji", icon="📚", icon_background="#FFFFFF", icon_url=None), + permission=DatasetPermissionEnum.ALL_TEAM, + ) + + with _dataset_service_dataset_patch( + "services.knowledge.dataset_service.current_user", _dataset_service_dataset_account() + ): + dataset = _dataset_service_dataset_DatasetService.create_empty_rag_pipeline_dataset( + "tenant-1", entity, sqlite_session + ) + + assert dataset.name == "Untitled 2" + pipeline = sqlite_session.get(Pipeline, dataset.pipeline_id) + assert pipeline is not None + assert pipeline.tenant_id == dataset.tenant_id + + def test_create_empty_rag_pipeline_dataset_rejects_duplicate_name(self, sqlite_session: Session) -> None: + sqlite_session.add(_dataset_service_dataset_dataset(name="Existing")) + sqlite_session.commit() + entity = RagPipelineDatasetCreateEntity( + name="Existing", + description="", + icon_info=IconInfo(icon_type="emoji", icon="📚", icon_background="#FFFFFF", icon_url=None), + permission=DatasetPermissionEnum.ALL_TEAM, + ) + + with ( + _dataset_service_dataset_patch( + "services.knowledge.dataset_service.current_user", _dataset_service_dataset_account() + ), + pytest.raises(DatasetNameDuplicateError, match="already exists"), + ): + _dataset_service_dataset_DatasetService.create_empty_rag_pipeline_dataset( + "tenant-1", entity, sqlite_session + ) + + def test_create_empty_rag_pipeline_dataset_requires_authenticated_account( + self, sqlite_session: Session + ) -> None: + account = _dataset_service_dataset_account() + account.id = "" + entity = RagPipelineDatasetCreateEntity( + name="Dataset", + description="", + icon_info=IconInfo(icon_type="emoji", icon="📚", icon_background="#FFFFFF", icon_url=None), + permission=DatasetPermissionEnum.ALL_TEAM, + ) + + with ( + _dataset_service_dataset_patch("services.knowledge.dataset_service.current_user", account), + pytest.raises(ValueError, match="Current user or current user id not found"), + ): + _dataset_service_dataset_DatasetService.create_empty_rag_pipeline_dataset( + "tenant-1", entity, sqlite_session + ) + + def test_update_dataset_rejects_missing_and_duplicate_rows(self, sqlite_session: Session) -> None: + current = _dataset_service_dataset_dataset(name="Current") + duplicate = _dataset_service_dataset_dataset(dataset_id="dataset-2", name="Duplicate") + sqlite_session.add_all([current, duplicate]) + sqlite_session.commit() + + with pytest.raises(ValueError, match="Dataset not found"): + _dataset_service_dataset_DatasetService.update_dataset( + "missing", {}, _dataset_service_dataset_account(), session=sqlite_session + ) + with pytest.raises(ValueError, match="Dataset name already exists"): + _dataset_service_dataset_DatasetService.update_dataset( + current.id, {"name": duplicate.name}, _dataset_service_dataset_account(), session=sqlite_session + ) + + @pytest.mark.parametrize( + ("provider", "helper_name"), + [("external", "_update_external_dataset"), ("vendor", "_update_internal_dataset")], + ) + def test_update_dataset_routes_by_provider( + self, sqlite_session: Session, provider: str, helper_name: str + ) -> None: + dataset = _dataset_service_dataset_dataset(provider=provider) + sqlite_session.add(dataset) + sqlite_session.commit() + + with ( + _dataset_service_dataset_patch.object( + _dataset_service_dataset_DatasetService, "check_dataset_permission" + ), + _dataset_service_dataset_patch.object( + _dataset_service_dataset_DatasetService, helper_name, return_value=dataset + ) as update_helper, + ): + result = _dataset_service_dataset_DatasetService.update_dataset( + dataset.id, + {"name": dataset.name}, + _dataset_service_dataset_account(), + session=sqlite_session, + ) + + assert result is dataset + update_helper.assert_called_once() + + @pytest.mark.parametrize( + ("payload", "message"), + [ + ({"external_knowledge_api_id": "api-1"}, "External knowledge id is required"), + ({"external_knowledge_id": "knowledge-1"}, "External knowledge api id is required"), + ], + ) + def test_update_external_dataset_requires_binding_fields( + self, unbound_session: Session, payload: dict[str, str], message: str + ) -> None: + with pytest.raises(ValueError, match=message): + _dataset_service_dataset_DatasetService._update_external_dataset( + _dataset_service_dataset_dataset(provider="external"), + payload, + _dataset_service_dataset_account(), + unbound_session, + ) + + def test_update_external_dataset_flushes_dataset_and_binding_without_committing( + self, sqlite_session: Session + ) -> None: + dataset = _dataset_service_dataset_dataset(provider="external") + api = _external_api() + binding = _binding() + sqlite_session.add_all([dataset, api, binding]) + sqlite_session.commit() + transaction_events: list[str] = [] + event.listen(sqlite_session, "after_commit", lambda _session: transaction_events.append("commit")) + + updated = _dataset_service_dataset_DatasetService._update_external_dataset( + dataset, + { + "name": "Updated", + "description": "Changed", + "external_knowledge_id": "knowledge-2", + "external_knowledge_api_id": api.id, + }, + _dataset_service_dataset_account(), + sqlite_session, + ) + + assert updated is dataset + assert dataset.name == "Updated" + assert binding.external_knowledge_id == "knowledge-2" + assert transaction_events == [] + + def test_update_external_dataset_rejects_cross_tenant_api(self, sqlite_session: Session) -> None: + dataset = _dataset_service_dataset_dataset(provider="external") + foreign_api = _external_api(api_id="api-foreign", tenant_id="tenant-2") + sqlite_session.add_all([dataset, foreign_api, _binding()]) + sqlite_session.commit() + + with pytest.raises(ValueError, match="api template not found"): + _dataset_service_dataset_DatasetService._update_external_dataset( + dataset, + { + "external_knowledge_id": "knowledge-2", + "external_knowledge_api_id": foreign_api.id, + }, + _dataset_service_dataset_account(), + sqlite_session, + ) + + def test_update_external_knowledge_binding_rejects_missing_row(self, sqlite_session: Session) -> None: + with pytest.raises(ValueError, match="binding not found"): + _dataset_service_dataset_DatasetService._update_external_knowledge_binding( + "missing-dataset", + "knowledge-2", + "api-2", + sqlite_session, + ) + + def test_update_internal_dataset_executes_real_update_without_committing(self, sqlite_session: Session) -> None: + dataset = _dataset_service_dataset_dataset(name="Before") + sqlite_session.add(dataset) + sqlite_session.commit() + transaction_events: list[str] = [] + event.listen(sqlite_session, "after_commit", lambda _session: transaction_events.append("commit")) + + with ( + _dataset_service_dataset_patch.object( + _dataset_service_dataset_DatasetService, "_handle_indexing_technique_change", return_value="update" + ), + _dataset_service_dataset_patch.object( + _dataset_service_dataset_DatasetService, "_update_pipeline_knowledge_base_node_data" + ), + _dataset_service_dataset_patch( + "services.knowledge.dataset_service.deal_dataset_vector_index_task.delay" + ) as vector_task, + _dataset_service_dataset_patch( + "services.knowledge.dataset_service.regenerate_summary_index_task.delay" + ) as summary_task, + ): + updated = _dataset_service_dataset_DatasetService._update_internal_dataset( + dataset, + {"name": "After", "description": "Changed"}, + _dataset_service_dataset_account(), + sqlite_session, + ) + + assert updated.name == "After" + assert updated.description == "Changed" + assert transaction_events == [] + vector_task.assert_called_once_with(dataset.id, "update") + summary_task.assert_called_once() + + def test_update_pipeline_node_data_returns_for_non_pipeline_or_missing_pipeline( + self, sqlite_session: Session + ) -> None: + ordinary = _dataset_service_dataset_dataset() + missing_pipeline = _dataset_service_dataset_dataset(dataset_id="pipeline-dataset", name="Pipeline Dataset") + missing_pipeline.runtime_mode = "rag_pipeline" + missing_pipeline.pipeline_id = "missing-pipeline" + sqlite_session.add_all([ordinary, missing_pipeline]) + sqlite_session.commit() + + _dataset_service_dataset_DatasetService._update_pipeline_knowledge_base_node_data( + ordinary, "user-1", sqlite_session + ) + _dataset_service_dataset_DatasetService._update_pipeline_knowledge_base_node_data( + missing_pipeline, "user-1", sqlite_session + ) + + def test_update_pipeline_node_data_rolls_back_real_session_on_failure(self, sqlite_session: Session) -> None: + pipeline = Pipeline(tenant_id="tenant-1", name="Pipeline", description="", created_by="user-1") + pipeline.id = "pipeline-1" + dataset = _dataset_service_dataset_dataset() + dataset.runtime_mode = "rag_pipeline" + dataset.pipeline_id = pipeline.id + sqlite_session.add_all([pipeline, dataset]) + sqlite_session.commit() + sqlite_session.begin() + transaction_events: list[str] = [] + event.listen( + sqlite_session, + "after_soft_rollback", + lambda _session, _previous_transaction: transaction_events.append("rollback"), + ) + + with _dataset_service_dataset_patch("services.knowledge.dataset_service.RagPipelineService") as service_cls: + service_cls.return_value.get_published_workflow.side_effect = RuntimeError("boom") + with pytest.raises(RuntimeError, match="boom"): + _dataset_service_dataset_DatasetService._update_pipeline_knowledge_base_node_data( + dataset, "user-1", sqlite_session + ) + + assert transaction_events == ["rollback"] + + class TestDatasetServiceEmbeddingSettings: + def test_handle_indexing_technique_change_returns_none_without_requested_change( + self, unbound_session: Session + ) -> None: + assert ( + _dataset_service_dataset_DatasetService._handle_indexing_technique_change( + _dataset_service_dataset_dataset(), {}, {}, unbound_session + ) + is None + ) + + def test_handle_indexing_technique_change_switches_to_economy(self, unbound_session: Session) -> None: + dataset = _dataset_service_dataset_dataset(indexing_technique=IndexTechniqueType.HIGH_QUALITY) + filtered_data: dict[str, object] = {} + + result = _dataset_service_dataset_DatasetService._handle_indexing_technique_change( + dataset, + {"indexing_technique": IndexTechniqueType.ECONOMY}, + filtered_data, + unbound_session, + ) + + assert result == "remove" + assert filtered_data == { + "embedding_model": None, + "embedding_model_provider": None, + "collection_binding_id": None, + } + + def test_handle_indexing_technique_change_delegates_high_quality_configuration( + self, unbound_session: Session + ) -> None: + dataset = _dataset_service_dataset_dataset(indexing_technique=IndexTechniqueType.ECONOMY) + filtered_data: dict[str, object] = {} + + with _dataset_service_dataset_patch.object( + _dataset_service_dataset_DatasetService, "_configure_embedding_model_for_high_quality" + ) as configure: + result = _dataset_service_dataset_DatasetService._handle_indexing_technique_change( + dataset, + {"indexing_technique": IndexTechniqueType.HIGH_QUALITY}, + filtered_data, + unbound_session, + ) + + assert result == "add" + configure.assert_called_once() + + def test_handle_unchanged_indexing_preserves_existing_model(self, unbound_session: Session) -> None: + dataset = _dataset_service_dataset_dataset(indexing_technique=IndexTechniqueType.HIGH_QUALITY) + + with _dataset_service_dataset_patch.object( + _dataset_service_dataset_DatasetService, "_preserve_existing_embedding_settings" + ) as preserve: + result = ( + _dataset_service_dataset_DatasetService._handle_embedding_model_update_when_technique_unchanged( + dataset, + {}, + {}, + unbound_session, + ) + ) + + assert result is None + preserve.assert_called_once() + + def test_update_embedding_model_settings_delegates_changed_model(self, unbound_session: Session) -> None: + dataset = _dataset_service_dataset_dataset() + + with _dataset_service_dataset_patch.object( + _dataset_service_dataset_DatasetService, "_apply_new_embedding_settings" + ) as apply_settings: + result = _dataset_service_dataset_DatasetService._update_embedding_model_settings( + dataset, + {"embedding_model_provider": "provider-2", "embedding_model": "model-2"}, + {}, + unbound_session, + ) + + assert result == "update" + apply_settings.assert_called_once() + + def test_configure_high_quality_wraps_provider_error(self, unbound_session: Session) -> None: + account = _dataset_service_dataset_account() + with ( + _dataset_service_dataset_patch("services.knowledge.dataset_service.current_user", account), + _dataset_service_dataset_patch("services.knowledge.dataset_service.ModelManager") as model_manager_cls, + ): + model_manager_cls.for_tenant.return_value.get_model_instance.side_effect = LLMBadRequestError() + with pytest.raises(ValueError, match="No Embedding Model available"): + _dataset_service_dataset_DatasetService._configure_embedding_model_for_high_quality( + {"embedding_model_provider": "provider", "embedding_model": "model"}, + {}, + unbound_session, + ) + + def test_preserve_existing_embedding_settings(self) -> None: + dataset = _dataset_service_dataset_dataset() + filtered_data: dict[str, object] = {} + + _dataset_service_dataset_DatasetService._preserve_existing_embedding_settings(dataset, filtered_data) + + assert filtered_data["embedding_model_provider"] == dataset.embedding_model_provider + assert filtered_data["embedding_model"] == dataset.embedding_model + + def test_update_embedding_model_settings_returns_none_when_unchanged(self, unbound_session: Session) -> None: + dataset = _dataset_service_dataset_dataset() + + result = _dataset_service_dataset_DatasetService._update_embedding_model_settings( + dataset, + { + "embedding_model_provider": dataset.embedding_model_provider, + "embedding_model": dataset.embedding_model, + }, + {}, + unbound_session, + ) + + assert result is None + + def test_apply_new_embedding_settings_uses_real_collection_binding(self, unbound_session: Session) -> None: + dataset = _dataset_service_dataset_dataset() + filtered_data: dict[str, object] = {} + account = _dataset_service_dataset_account() + embedding_model = SimpleNamespace(provider="new-provider", model_name="new-model") + collection_binding = DatasetCollectionBinding( + provider_name="new-provider", + model_name="new-model", + type="dataset", + collection_name="collection", + ) + collection_binding.id = "collection-binding-1" + + with ( + _dataset_service_dataset_patch("services.knowledge.dataset_service.current_user", account), + _dataset_service_dataset_patch("services.knowledge.dataset_service.ModelManager") as model_manager_cls, + _dataset_service_dataset_patch.object( + DatasetCollectionBindingService, + "get_dataset_collection_binding", + return_value=collection_binding, + ), + ): + model_manager_cls.for_tenant.return_value.get_model_instance.return_value = embedding_model + _dataset_service_dataset_DatasetService._apply_new_embedding_settings( + dataset, + {"embedding_model_provider": "new-provider", "embedding_model": "new-model"}, + filtered_data, + unbound_session, + ) + + assert filtered_data["embedding_model_provider"] == "new-provider" + assert filtered_data["embedding_model"] == "new-model" + assert filtered_data["collection_binding_id"] == collection_binding.id + + def test_apply_new_embedding_settings_preserves_existing_values_when_token_missing( + self, unbound_session: Session + ) -> None: + dataset = _dataset_service_dataset_dataset() + filtered_data: dict[str, object] = {} + + with ( + _dataset_service_dataset_patch( + "services.knowledge.dataset_service.current_user", _dataset_service_dataset_account() + ), + _dataset_service_dataset_patch("services.knowledge.dataset_service.ModelManager") as model_manager_cls, + ): + model_manager_cls.for_tenant.return_value.get_model_instance.side_effect = ProviderTokenNotInitError( + "missing" + ) + _dataset_service_dataset_DatasetService._apply_new_embedding_settings( + dataset, + {"embedding_model_provider": "provider-2", "embedding_model": "model-2"}, + filtered_data, + unbound_session, + ) + + assert filtered_data["embedding_model_provider"] == dataset.embedding_model_provider + assert filtered_data["embedding_model"] == dataset.embedding_model + + @pytest.mark.parametrize( + ("summary_setting", "expected"), + [ + (None, False), + ({"enable": False}, False), + ({"enable": True, "model_name": "old-model", "model_provider_name": "provider"}, False), + ({"enable": True, "model_name": "new-model", "model_provider_name": "provider"}, True), + ], + ) + def test_check_summary_index_setting_model_changed( + self, summary_setting: dict[str, object] | None, expected: bool + ) -> None: + dataset = _dataset_service_dataset_dataset() + dataset.summary_index_setting = { + "enable": True, + "model_name": "old-model", + "model_provider_name": "provider", + } + + data = {} if summary_setting is None else {"summary_index_setting": summary_setting} + assert ( + _dataset_service_dataset_DatasetService._check_summary_index_setting_model_changed(dataset, data) + is expected + ) + + class TestDatasetServiceRagPipelineSettings: + def test_requires_current_tenant(self, unbound_session: Session) -> None: + account = _dataset_service_dataset_account() + account._current_tenant = None + + with ( + _dataset_service_dataset_patch("services.knowledge.dataset_service.current_user", account), + pytest.raises(ValueError, match="Current user or current tenant not found"), + ): + _dataset_service_dataset_DatasetService.update_rag_pipeline_dataset_settings( + _dataset_service_dataset_dataset(), + _make_knowledge_configuration(), + session=unbound_session, + ) + + def test_unpublished_high_quality_settings_use_real_merged_dataset(self, sqlite_session: Session) -> None: + dataset = _dataset_service_dataset_dataset(indexing_technique=IndexTechniqueType.ECONOMY) + sqlite_session.add(dataset) + sqlite_session.commit() + account = _dataset_service_dataset_account() + embedding_model = SimpleNamespace(provider="provider-2", model_name="embedding-2") + collection_binding = DatasetCollectionBinding( + provider_name="provider-2", + model_name="embedding-2", + type="dataset", + collection_name="collection", + ) + collection_binding.id = "collection-binding-2" + + with ( + _dataset_service_dataset_patch("services.knowledge.dataset_service.current_user", account), + _dataset_service_dataset_patch("services.knowledge.dataset_service.ModelManager") as model_manager_cls, + _dataset_service_dataset_patch.object( + _dataset_service_dataset_DatasetService, "check_is_multimodal_model", return_value=True + ), + _dataset_service_dataset_patch.object( + DatasetCollectionBindingService, + "get_dataset_collection_binding", + return_value=collection_binding, + ), + ): + model_manager_cls.for_tenant.return_value.get_model_instance.return_value = embedding_model + _dataset_service_dataset_DatasetService.update_rag_pipeline_dataset_settings( + dataset, + _make_knowledge_configuration( + embedding_model_provider="provider-2", + embedding_model="embedding-2", + summary_index_setting={"enable": True}, + ), + session=sqlite_session, + ) + + persisted = sqlite_session.get(Dataset, dataset.id) + assert persisted is not None + assert persisted.indexing_technique == IndexTechniqueType.HIGH_QUALITY + assert persisted.embedding_model == "embedding-2" + assert persisted.collection_binding_id == collection_binding.id + assert persisted.is_multimodal is True + + def test_unpublished_economy_settings_update_keyword_number(self, sqlite_session: Session) -> None: + dataset = _dataset_service_dataset_dataset(indexing_technique=IndexTechniqueType.HIGH_QUALITY) + sqlite_session.add(dataset) + sqlite_session.commit() + + with _dataset_service_dataset_patch( + "services.knowledge.dataset_service.current_user", _dataset_service_dataset_account() + ): + _dataset_service_dataset_DatasetService.update_rag_pipeline_dataset_settings( + dataset, + _make_knowledge_configuration( + indexing_technique=IndexTechniqueType.ECONOMY, + keyword_number=17, + ), + session=sqlite_session, + ) + + persisted = sqlite_session.get(Dataset, dataset.id) + assert persisted is not None + assert persisted.indexing_technique == IndexTechniqueType.ECONOMY + assert persisted.keyword_number == 17 + + def test_published_economy_settings_commit_keyword_change_and_dispatch_no_task( + self, sqlite_session: Session + ) -> None: + dataset = _dataset_service_dataset_dataset( + indexing_technique=IndexTechniqueType.ECONOMY, chunk_structure="paragraph" + ) + dataset.keyword_number = 4 + sqlite_session.add(dataset) + sqlite_session.commit() + + with ( + _dataset_service_dataset_patch( + "services.knowledge.dataset_service.current_user", _dataset_service_dataset_account() + ), + _dataset_service_dataset_patch( + "services.knowledge.dataset_service.deal_dataset_index_update_task.delay" + ) as update_task, + ): + _dataset_service_dataset_DatasetService.update_rag_pipeline_dataset_settings( + dataset, + _make_knowledge_configuration( + chunk_structure="paragraph", + indexing_technique=IndexTechniqueType.ECONOMY, + keyword_number=12, + ), + has_published=True, + session=sqlite_session, + ) + + assert sqlite_session.get(Dataset, dataset.id).keyword_number == 12 + update_task.assert_not_called() + + @pytest.mark.parametrize( + ("dataset", "configuration", "message"), + [ + ( + _dataset_service_dataset_dataset(chunk_structure="paragraph"), + _make_knowledge_configuration(chunk_structure="sentence"), + "Chunk structure is not allowed", + ), + ( + _dataset_service_dataset_dataset( + indexing_technique=IndexTechniqueType.HIGH_QUALITY, chunk_structure="paragraph" + ), + _make_knowledge_configuration( + chunk_structure="paragraph", + indexing_technique=IndexTechniqueType.ECONOMY, + ), + "not allowed to be updated to economy", + ), + ], + ) + def test_published_settings_reject_incompatible_changes( + self, + sqlite_session: Session, + dataset: Dataset, + configuration, + message: str, + ) -> None: + sqlite_session.add(dataset) + sqlite_session.commit() + + with ( + _dataset_service_dataset_patch( + "services.knowledge.dataset_service.current_user", _dataset_service_dataset_account() + ), + pytest.raises(ValueError, match=message), + ): + _dataset_service_dataset_DatasetService.update_rag_pipeline_dataset_settings( + dataset, + configuration, + has_published=True, + session=sqlite_session, + ) + + class TestDatasetPermissions: + def test_check_dataset_permission_enforces_tenant_and_partial_members(self, sqlite_session: Session) -> None: + dataset = _dataset_service_dataset_dataset( + permission=DatasetPermissionEnum.PARTIAL_TEAM, maintainer="owner" + ) + permitted_user = _dataset_service_dataset_account(account_id="permitted", role=TenantAccountRole.NORMAL) + denied_user = _dataset_service_dataset_account(account_id="denied", role=TenantAccountRole.NORMAL) + foreign_user = _dataset_service_dataset_account( + account_id="foreign", tenant_id="tenant-2", role=TenantAccountRole.NORMAL + ) + sqlite_session.add_all( + [ + dataset, + DatasetPermission( + tenant_id=dataset.tenant_id, + dataset_id=dataset.id, + account_id=permitted_user.id, + ), + ] + ) + sqlite_session.commit() + + _dataset_service_dataset_DatasetService.check_dataset_permission(dataset, permitted_user, sqlite_session) + with pytest.raises(NoPermissionError): + _dataset_service_dataset_DatasetService.check_dataset_permission(dataset, denied_user, sqlite_session) + with pytest.raises(NoPermissionError): + _dataset_service_dataset_DatasetService.check_dataset_permission(dataset, foreign_user, sqlite_session) + + def test_check_dataset_permission_skips_legacy_acl_in_rbac_mode( + self, config_overrides: Callable[..., None], sqlite_session: Session + ) -> None: + dataset = _dataset_service_dataset_dataset(permission=DatasetPermissionEnum.ONLY_ME, maintainer="owner") + non_member = _dataset_service_dataset_account(account_id="non-member", role=TenantAccountRole.NORMAL) + foreign_user = _dataset_service_dataset_account( + account_id="foreign", tenant_id="tenant-2", role=TenantAccountRole.NORMAL + ) + sqlite_session.add(dataset) + sqlite_session.commit() + + config_overrides(RBAC_ENABLED=True) + + _dataset_service_dataset_DatasetService.check_dataset_permission(dataset, non_member, sqlite_session) + with pytest.raises(NoPermissionError): + _dataset_service_dataset_DatasetService.check_dataset_permission(dataset, foreign_user, sqlite_session) + + def test_dataset_operator_cannot_change_permission_or_member_list(self, sqlite_session: Session) -> None: + dataset = _dataset_service_dataset_dataset(permission=DatasetPermissionEnum.PARTIAL_TEAM) + operator = _dataset_service_dataset_account(role=TenantAccountRole.DATASET_OPERATOR) + sqlite_session.add_all( + [ + dataset, + DatasetPermission(tenant_id="tenant-1", dataset_id=dataset.id, account_id="member-1"), + ] + ) + sqlite_session.commit() + + with pytest.raises(NoPermissionError, match="cannot change"): + DatasetPermissionService.check_permission( + operator, + dataset, + DatasetPermissionEnum.ALL_TEAM, + None, + session=sqlite_session, + ) + + def test_non_editor_cannot_change_dataset_permissions(self, sqlite_session: Session) -> None: + user = _dataset_service_dataset_account(role=TenantAccountRole.NORMAL) + + with pytest.raises(NoPermissionError, match="does not have permission"): + DatasetPermissionService.check_permission( + user, + _dataset_service_dataset_dataset(), + DatasetPermissionEnum.ALL_TEAM, + None, + session=sqlite_session, + ) + + def test_dataset_operator_can_keep_unchanged_partial_member_list(self, sqlite_session: Session) -> None: + dataset = _dataset_service_dataset_dataset(permission=DatasetPermissionEnum.PARTIAL_TEAM) + operator = _dataset_service_dataset_account(role=TenantAccountRole.DATASET_OPERATOR) + sqlite_session.add_all( + [ + dataset, + DatasetPermission(tenant_id="tenant-1", dataset_id=dataset.id, account_id="member-1"), + ] + ) + sqlite_session.commit() + + DatasetPermissionService.check_permission( + operator, + dataset, + DatasetPermissionEnum.PARTIAL_TEAM, + [{"user_id": "member-1"}], + session=sqlite_session, + ) + with pytest.raises(ValueError, match="cannot change"): + DatasetPermissionService.check_permission( + operator, + dataset, + DatasetPermissionEnum.PARTIAL_TEAM, + [{"user_id": "member-2"}], + session=sqlite_session, + ) + + def test_update_partial_member_list_flush_failure_does_not_rollback_caller_session( + self, sqlite_session: Session + ) -> None: + transaction_events: list[str] = [] + + def fail_flush(_session, _flush_context, _instances) -> None: + raise RuntimeError("flush failed") + + event.listen(sqlite_session, "before_flush", fail_flush) + event.listen(sqlite_session, "after_rollback", lambda _session: transaction_events.append("rollback")) + try: + with pytest.raises(RuntimeError, match="flush failed"): + DatasetPermissionService.update_partial_member_list( + "tenant-1", + "dataset-1", + [{"user_id": "member-1"}], + sqlite_session, + ) + finally: + event.remove(sqlite_session, "before_flush", fail_flush) + + assert transaction_events == [] + + def test_clear_partial_member_list_execute_failure_does_not_rollback_caller_session( + self, sqlite_session: Session + ) -> None: + transaction_events: list[str] = [] + + def fail_execute(_orm_execute_state) -> None: + raise RuntimeError("execute failed") + + event.listen(sqlite_session, "do_orm_execute", fail_execute) + event.listen(sqlite_session, "after_rollback", lambda _session: transaction_events.append("rollback")) + try: + with pytest.raises(RuntimeError, match="execute failed"): + DatasetPermissionService.clear_partial_member_list("dataset-1", sqlite_session) + finally: + event.remove(sqlite_session, "do_orm_execute", fail_execute) + + assert transaction_events == [] + + @staticmethod + @pytest.mark.parametrize("rbac_enabled", [False, True]) + @pytest.mark.parametrize("commit_fails", [False, True]) + @pytest.mark.parametrize("expire_on_commit", [False, True]) + def test_empty_pipeline_syncs_creator_after_commit( + sqlite_session: Session, + sqlite_session_factory: sessionmaker[Session], + config_overrides: Callable[..., None], + rbac_enabled: bool, + commit_fails: bool, + expire_on_commit: bool, + ) -> None: + config_overrides(RBAC_ENABLED=rbac_enabled) + sqlite_session.expire_on_commit = expire_on_commit + account = _dataset_service_dataset_account() + entity = RagPipelineDatasetCreateEntity( + name="New pipeline", + description="", + permission=DatasetPermissionEnum.ONLY_ME, + icon_info=IconInfo(icon_type="emoji", icon="📚", icon_background="#FFFFFF", icon_url=None), + ) + commit = sqlite_session.commit + observed: list[tuple[str, str, object]] = [] + transactions_at_sync: list[bool] = [] + + def commit_session() -> None: + if commit_fails: + raise RuntimeError("commit failed") + commit() + + def sync_binding(_tenant_id: str, _account_id: str, *, resource_type: object, resource_id: str) -> list[object]: + transactions_at_sync.append(sqlite_session.in_transaction()) + with sqlite_session_factory() as observer: + stored = observer.get(Dataset, resource_id) + if stored is not None: + observed.append((stored.tenant_id, stored.created_by, resource_type)) + return [] + + with ( + _dataset_service_dataset_patch("services.knowledge.dataset_service.current_user", account), + _dataset_service_dataset_patch.object(sqlite_session, "commit", side_effect=commit_session), + _dataset_service_dataset_patch.object( + rbac_service.RBACService.AccessPolicies, + "sync_creator_access_policy_member_bindings", + side_effect=sync_binding, + ) as sync, + ): + if commit_fails: + with pytest.raises(RuntimeError, match="commit failed"): + _dataset_service_dataset_DatasetService.create_empty_rag_pipeline_dataset( + "tenant-1", entity, sqlite_session + ) + else: + dataset = _dataset_service_dataset_DatasetService.create_empty_rag_pipeline_dataset( + "tenant-1", entity, sqlite_session + ) + if rbac_enabled: + assert transactions_at_sync == [False] + assert observed == [("tenant-1", account.id, rbac_service.RBACResourceType.DATASET)] + sync.assert_called_once_with( + "tenant-1", + account.id, + resource_type=rbac_service.RBACResourceType.DATASET, + resource_id=dataset.id, + ) + if commit_fails or not rbac_enabled: + sync.assert_not_called() + + +"""Unit tests for DocumentService behaviors in dataset_service.""" + + +def _dataset_service_document_account( + *, account_id: str = "user-1", tenant_id: str = "tenant-1" +) -> _dataset_service_document_Account: + return make_account( + account_id=account_id, + name="User", + email=f"{account_id}@example.com", + tenant=make_tenant(tenant_id=tenant_id, name="Tenant"), + ) + + +def _dataset_row( + *, + dataset_id: str = "dataset-1", + tenant_id: str = "tenant-1", + built_in_field_enabled: bool = False, + data_source_type: str | None = None, + indexing_technique: str | None = "economy", +) -> Dataset: + return Dataset( + id=dataset_id, + tenant_id=tenant_id, + name="Dataset", + description="", + provider="vendor", + created_by="user-1", + maintainer="user-1", + built_in_field_enabled=built_in_field_enabled, + chunk_structure=_dataset_service_document_IndexStructureType.PARAGRAPH_INDEX, + data_source_type=data_source_type, + indexing_technique=indexing_technique, + ) + + +def _document_row( + *, + document_id: str = "document-1", + dataset_id: str = "dataset-1", + tenant_id: str = "tenant-1", + name: str = "Document", + indexing_status: str = IndexingStatus.COMPLETED, + data_source_type: str = DataSourceType.UPLOAD_FILE, + data_source_info: str = "{}", + enabled: bool = True, + archived: bool = False, + is_paused: bool = False, +) -> _dataset_service_document_Document: + document = _dataset_service_document_Document( + id=document_id, + tenant_id=tenant_id, + dataset_id=dataset_id, + position=1, + data_source_type=data_source_type, + data_source_info=data_source_info, + batch="batch-1", + name=name, + created_from=DocumentCreatedFrom.API, + created_by="user-1", + indexing_status=indexing_status, + doc_form=_dataset_service_document_IndexStructureType.PARAGRAPH_INDEX, + word_count=10, + enabled=enabled, + archived=archived, + is_paused=is_paused, + completed_at=datetime(2026, 1, 2) if indexing_status == IndexingStatus.COMPLETED else None, + ) + document.created_at = datetime(2026, 1, 1) + document.updated_at = datetime(2026, 1, 2) + return document + + +def _upload_file(*, file_id: str, tenant_id: str = "tenant-1", name: str = "upload.txt") -> UploadFile: + return make_upload_file( + file_id=file_id, + tenant_id=tenant_id, + storage_type="opendal", + key=f"key-{file_id}", + name=name, + size=1, + created_by="user-1", + created_at=datetime(2026, 1, 1), + ) + + +def _process_rule( + *, dataset_id: str = "dataset-1", rule_id: str = "rule-1", mode: str = "automatic" +) -> _dataset_service_document_DatasetProcessRule: + rules = ( + json.dumps(_dataset_service_document_DatasetProcessRule.AUTOMATIC_RULES) + if mode == "automatic" + else Rule( + pre_processing_rules=[PreProcessingRule(id="remove_extra_spaces", enabled=True)], + segmentation=Segmentation(separator="\n", max_tokens=100), + ).model_dump_json() + ) + process_rule = _dataset_service_document_DatasetProcessRule( + dataset_id=dataset_id, mode=mode, rules=rules, created_by="user-1" + ) + process_rule.id = rule_id + return process_rule + + +class _RetryFlagLock: + def __init__(self, store: _RetryFlagStore, key: str): + self.store = store + self.key = key + self.token = f"owner-{store.next_token}" + store.next_token += 1 + + def acquire(self, *, blocking: bool): + assert blocking is False + if self.key in self.store.values: + if self.store.replacement_on_conflict: + replacement_key, replacement_value = self.store.replacement_on_conflict + self.store.values[replacement_key] = replacement_value + return False + self.store.values[self.key] = self.token + return True + + def release(self): + if self.store.values.get(self.key) == self.token: + self.store.values.pop(self.key) + + +class _RetryFlagStore: + def __init__( + self, + values: dict[str, str] | None = None, + replacement_on_conflict: tuple[str, str] | None = None, + ): + self.values = values or {} + self.replacement_on_conflict = replacement_on_conflict + self.next_token = 1 + + def lock(self, key: str, *, timeout: int, thread_local: bool): + assert timeout == 600 + assert thread_local is False + return _RetryFlagLock(self, key) + + +class TestDatasetServiceDocumentCases: + class TestDocumentServiceDisplayStatus: + """Unit tests for DocumentService display-status helpers.""" + + @pytest.mark.parametrize( + ("raw_status", "expected"), + [ + ("enabled", "available"), + ("AVAILABLE", "available"), + ("paused", "paused"), + ("unknown", None), + (None, None), + ], + ) + def test_normalize_display_status(self, raw_status, expected): + assert _dataset_service_document_DocumentService.normalize_display_status(raw_status) == expected + + def test_build_display_status_filters_returns_empty_tuple_for_unknown_status(self): + assert _dataset_service_document_DocumentService.build_display_status_filters("missing") == () + + def test_apply_display_status_filter_returns_original_query_for_unknown_status(self): + query = select(_dataset_service_document_Document) + + result = _dataset_service_document_DocumentService.apply_display_status_filter(query, "missing") + + assert result is query + + def test_apply_display_status_filter_applies_where_for_known_status(self): + query = select(_dataset_service_document_Document) + + result = _dataset_service_document_DocumentService.apply_display_status_filter(query, "enabled") + + assert result is not query + assert "documents.enabled" in str(result) + + class TestDocumentServiceRetrieval: + def test_get_document_by_id_uses_provided_session(self, sqlite_session: Session): + document = _document_row() + sqlite_session.add(document) + sqlite_session.commit() + + assert ( + _dataset_service_document_DocumentService.get_document_by_id(document.id, session=sqlite_session) + is document + ) + assert ( + _dataset_service_document_DocumentService.get_document_by_id("missing", session=sqlite_session) is None + ) + + def test_get_document_by_ids_enforces_dataset_owner_and_state(self, sqlite_session: Session): + dataset = _dataset_row() + expected = _document_row(document_id="expected") + sqlite_session.add_all( + [ + dataset, + expected, + _document_row(document_id="disabled", enabled=False), + _document_row(document_id="archived", archived=True), + _document_row(document_id="waiting", indexing_status=IndexingStatus.WAITING), + _document_row(document_id="other-dataset", dataset_id="dataset-2"), + _document_row(document_id="other-tenant", tenant_id="tenant-2"), + ] + ) + sqlite_session.commit() + + documents = _dataset_service_document_DocumentService.get_document_by_ids( + DatasetRef(tenant_id=dataset.tenant_id, dataset_id=dataset.id), + ["expected", "disabled", "archived", "waiting", "other-dataset", "other-tenant"], + sqlite_session, + ) + + assert [document.id for document in documents] == [expected.id] + + class TestDocumentServiceMutations: + """Unit tests for DocumentService mutation and orchestration helpers.""" + + @pytest.mark.parametrize(("archived", "expected"), [(True, True), (False, False)]) + def test_check_archived_returns_boolean_status(self, archived, expected): + document = _document_row(archived=archived) + + assert _dataset_service_document_DocumentService.check_archived(document) is expected + + def test_delete_documents_limits_query_and_cleanup_to_dataset_ref(self, sqlite_session: Session): + dataset = _dataset_row() + document = _document_row( + document_id="doc-1", + data_source_info=json.dumps({"upload_file_id": "file-1"}), + ) + other_dataset = _document_row(document_id="other-dataset", dataset_id="dataset-2") + other_tenant = _document_row(document_id="other-tenant", tenant_id="tenant-2") + sqlite_session.add_all([dataset, document, other_dataset, other_tenant]) + sqlite_session.commit() + + with _dataset_service_document_patch( + "services.knowledge.dataset_service.batch_clean_document_task" + ) as clean_task: + dataset_ref = DatasetRef(tenant_id=dataset.tenant_id, dataset_id=dataset.id) + _dataset_service_document_DocumentService.delete_documents( + dataset_ref, + [document.id, other_dataset.id, other_tenant.id], + _dataset_service_document_IndexStructureType.PARAGRAPH_INDEX, + sqlite_session, + ) + + assert sqlite_session.get(_dataset_service_document_Document, document.id) is None + assert sqlite_session.get(_dataset_service_document_Document, other_dataset.id) is other_dataset + assert sqlite_session.get(_dataset_service_document_Document, other_tenant.id) is other_tenant + clean_task.delay.assert_called_once_with( + [document.id], dataset.id, _dataset_service_document_IndexStructureType.PARAGRAPH_INDEX, ["file-1"] + ) + + def test_delete_documents_with_empty_ids_does_not_commit(self, sqlite_session: Session): + commits = 0 + + def count_commit(_session): + nonlocal commits + commits += 1 + + event.listen(sqlite_session, "after_commit", count_commit) + dataset = _dataset_row() + _dataset_service_document_DocumentService.delete_documents( + DatasetRef(tenant_id=dataset.tenant_id, dataset_id=dataset.id), + [], + _dataset_service_document_IndexStructureType.PARAGRAPH_INDEX, + sqlite_session, + ) + event.remove(sqlite_session, "after_commit", count_commit) + assert commits == 0 + + def test_rename_document_raises_when_dataset_is_missing(self, sqlite_session: Session): + with _dataset_service_document_patch( + "services.knowledge.dataset_service.current_user", _dataset_service_document_account() + ): + with pytest.raises(ValueError, match="Dataset not found"): + _dataset_service_document_DocumentService.rename_document( + "dataset-1", "doc-1", "New Name", sqlite_session + ) + + def test_rename_document_raises_when_document_is_missing(self, sqlite_session: Session): + dataset = _dataset_row() + sqlite_session.add(dataset) + sqlite_session.commit() + with _dataset_service_document_patch( + "services.knowledge.dataset_service.current_user", _dataset_service_document_account() + ): + with pytest.raises(ValueError, match="Document not found"): + _dataset_service_document_DocumentService.rename_document( + dataset.id, "doc-1", "New Name", sqlite_session + ) + + def test_rename_document_rejects_cross_tenant_access(self, sqlite_session: Session): + dataset = _dataset_row() + document = _document_row(tenant_id="tenant-other") + sqlite_session.add_all([dataset, document]) + sqlite_session.commit() + with _dataset_service_document_patch( + "services.knowledge.dataset_service.current_user", _dataset_service_document_account() + ): + with pytest.raises(ValueError, match="No permission"): + _dataset_service_document_DocumentService.rename_document( + dataset.id, document.id, "New Name", sqlite_session + ) + + def test_rename_document_updates_document_metadata_and_upload_file_name(self, sqlite_session: Session): + dataset = _dataset_row(built_in_field_enabled=True) + document = _document_row( + data_source_info=json.dumps({"upload_file_id": "file-1"}), + ) + document.doc_metadata = {_dataset_service_document_BuiltInField.document_name: "Old"} + upload_file = UploadFile( + tenant_id=dataset.tenant_id, + storage_type="opendal", + key="key", + name="old.txt", + size=1, + extension="txt", + mime_type="text/plain", + created_by_role="account", + created_by="user-1", + created_at=datetime(2026, 1, 1), + used=False, + ) + upload_file.id = "file-1" + sqlite_session.add_all([dataset, document, upload_file]) + sqlite_session.commit() + commits = 0 + + def count_commit(_session): + nonlocal commits + commits += 1 + + event.listen(sqlite_session, "after_commit", count_commit) + with _dataset_service_document_patch( + "services.knowledge.dataset_service.current_user", _dataset_service_document_account() + ): + result = _dataset_service_document_DocumentService.rename_document( + dataset.id, document.id, "New Name", sqlite_session + ) + event.remove(sqlite_session, "after_commit", count_commit) + + assert result is document + assert document.name == "New Name" + assert document.doc_metadata[_dataset_service_document_BuiltInField.document_name] == "New Name" + assert sqlite_session.get(UploadFile, upload_file.id).name == "New Name" + assert commits == 0 + + def test_recover_document_raises_when_document_is_not_paused(self, unbound_session: Session): + document = _document_row(is_paused=False) + with pytest.raises(DocumentIndexingError): + _dataset_service_document_DocumentService.recover_document(document, unbound_session) + + def test_recover_document_persists_and_dispatches(self, sqlite_session: Session): + document = _document_row(is_paused=True) + sqlite_session.add(document) + sqlite_session.commit() + with ( + _dataset_service_document_patch("services.knowledge.dataset_service.redis_client") as redis, + _dataset_service_document_patch( + "services.knowledge.dataset_service.recover_document_indexing_task" + ) as task, + ): + _dataset_service_document_DocumentService.recover_document(document, sqlite_session) + + sqlite_session.expire_all() + recovered = sqlite_session.get(_dataset_service_document_Document, document.id) + assert recovered is not None + assert recovered.is_paused is False + redis.delete.assert_called_once_with(f"document_{document.id}_is_paused") + task.delay.assert_called_once_with(document.dataset_id, document.id) + + def test_retry_document_raises_when_retry_flag_is_already_set(self, sqlite_session: Session): + document = _document_row(indexing_status=IndexingStatus.ERROR) + sqlite_session.add(document) + sqlite_session.commit() + retry_flags = _RetryFlagStore({f"document_{document.id}_is_retried": "other-request"}) + with ( + _dataset_service_document_patch( + "services.knowledge.dataset_service.current_user", _dataset_service_document_account() + ), + _dataset_service_document_patch("services.knowledge.dataset_service.redis_client", retry_flags), + ): + with pytest.raises(ValueError, match="being retried"): + _dataset_service_document_DocumentService.retry_document( + "dataset-1", [document], sqlite_session, actor_id="user-1" + ) + + def test_retry_document_leaves_batch_unchanged_when_later_document_is_already_being_retried( + self, sqlite_session: Session + ): + first_document = _document_row(document_id="doc-1", indexing_status=IndexingStatus.ERROR) + second_document = _document_row(document_id="doc-2", indexing_status=IndexingStatus.ERROR) + sqlite_session.add_all([first_document, second_document]) + sqlite_session.commit() + first_retry_key = "document_doc-1_is_retried" + second_retry_key = "document_doc-2_is_retried" + retry_flags = _RetryFlagStore({second_retry_key: "other-request"}) + with ( + _dataset_service_document_patch( + "services.knowledge.dataset_service.current_user", _dataset_service_document_account() + ), + _dataset_service_document_patch("services.knowledge.dataset_service.redis_client", retry_flags), + _dataset_service_document_patch( + "services.knowledge.dataset_service.retry_document_indexing_task" + ) as retry_task, + ): + with pytest.raises(ValueError, match="being retried"): + _dataset_service_document_DocumentService.retry_document( + "dataset-1", [first_document, second_document], sqlite_session, actor_id="user-1" + ) + + assert first_document.indexing_status == IndexingStatus.ERROR + assert second_document.indexing_status == IndexingStatus.ERROR + assert first_retry_key not in retry_flags.values + assert retry_flags.values[second_retry_key] == "other-request" + retry_task.delay.assert_not_called() + + def test_retry_document_does_not_release_a_retry_flag_reacquired_by_another_request( + self, sqlite_session: Session + ): + first_retry_key = "document_doc-1_is_retried" + second_retry_key = "document_doc-2_is_retried" + retry_flags = _RetryFlagStore( + {second_retry_key: "other-request"}, + replacement_on_conflict=(first_retry_key, "new-owner"), + ) + documents = [ + _document_row(document_id="doc-1", indexing_status=IndexingStatus.ERROR), + _document_row(document_id="doc-2", indexing_status=IndexingStatus.ERROR), + ] + sqlite_session.add_all(documents) + sqlite_session.commit() + + with ( + _dataset_service_document_patch( + "services.knowledge.dataset_service.current_user", _dataset_service_document_account() + ), + _dataset_service_document_patch("services.knowledge.dataset_service.redis_client", retry_flags), + ): + with pytest.raises(ValueError, match="being retried"): + _dataset_service_document_DocumentService.retry_document( + "dataset-1", documents, sqlite_session, actor_id="user-1" + ) + + assert retry_flags.values[first_retry_key] == "new-owner" + assert retry_flags.values[second_retry_key] == "other-request" + + def test_retry_document_releases_flags_when_status_commit_fails(self, sqlite_session: Session): + retry_flags = _RetryFlagStore() + document = _document_row(indexing_status=IndexingStatus.ERROR) + sqlite_session.add(document) + sqlite_session.commit() + + def fail_commit(_session): + raise RuntimeError("database unavailable") + + event.listen(sqlite_session, "before_commit", fail_commit) + with ( + _dataset_service_document_patch( + "services.knowledge.dataset_service.current_user", _dataset_service_document_account() + ), + _dataset_service_document_patch("services.knowledge.dataset_service.redis_client", retry_flags), + _dataset_service_document_patch( + "services.knowledge.dataset_service.retry_document_indexing_task" + ) as retry_task, + ): + with pytest.raises(RuntimeError, match="database unavailable"): + _dataset_service_document_DocumentService.retry_document( + "dataset-1", [document], sqlite_session, actor_id="user-1" + ) + event.remove(sqlite_session, "before_commit", fail_commit) + + assert retry_flags.values == {} + retry_task.delay.assert_not_called() + + def test_retry_document_persists_status_and_dispatches(self, sqlite_session: Session): + documents = [ + _document_row(document_id="doc-1", indexing_status=IndexingStatus.ERROR), + _document_row(document_id="doc-2", indexing_status=IndexingStatus.PAUSED), + ] + sqlite_session.add_all(documents) + sqlite_session.commit() + retry_flags = _RetryFlagStore() + with ( + _dataset_service_document_patch( + "services.knowledge.dataset_service.current_user", _dataset_service_document_account() + ), + _dataset_service_document_patch("services.knowledge.dataset_service.redis_client", retry_flags), + _dataset_service_document_patch( + "services.knowledge.dataset_service.retry_document_indexing_task" + ) as task, + ): + _dataset_service_document_DocumentService.retry_document( + "dataset-1", documents, sqlite_session, actor_id="user-1" + ) + + sqlite_session.expire_all() + statuses = sqlite_session.scalars( + select(_dataset_service_document_Document.indexing_status).order_by( + _dataset_service_document_Document.id + ) + ).all() + assert statuses == [IndexingStatus.WAITING, IndexingStatus.WAITING] + task.delay.assert_called_once_with("dataset-1", ["doc-1", "doc-2"], "user-1") + + def test_sync_website_document_raises_when_sync_flag_exists(self, sqlite_session: Session): + dataset = _dataset_row() + document = _document_row() + with _dataset_service_document_patch("services.knowledge.dataset_service.redis_client") as mock_redis: + mock_redis.get.return_value = "1" + + with pytest.raises(ValueError, match="being synced"): + _dataset_service_document_DocumentService.sync_website_document(dataset, document, sqlite_session) + + @pytest.mark.parametrize( + ("document_dataset_id", "document_tenant_id"), + [("dataset-2", "tenant-1"), ("dataset-1", "tenant-2")], + ) + def test_sync_website_document_rejects_document_outside_dataset( + self, document_dataset_id: str, document_tenant_id: str, sqlite_session: Session + ): + dataset = _dataset_row() + document = _document_row(dataset_id=document_dataset_id, tenant_id=document_tenant_id) + + with ( + pytest.raises(ValueError, match="Document not found"), + _dataset_service_document_patch("services.knowledge.dataset_service.redis_client") as mock_redis, + ): + _dataset_service_document_DocumentService.sync_website_document(dataset, document, sqlite_session) + + mock_redis.get.assert_not_called() + + def test_sync_website_document_updates_status_sets_cache_and_dispatches_task(self, sqlite_session: Session): + dataset = _dataset_row() + document = _document_row( + data_source_type=DataSourceType.WEBSITE_CRAWL, + data_source_info=json.dumps({"mode": "crawl"}), + ) + sqlite_session.add_all([dataset, document]) + sqlite_session.commit() + + with ( + _dataset_service_document_patch("services.knowledge.dataset_service.redis_client") as mock_redis, + _dataset_service_document_patch( + "services.knowledge.dataset_service.sync_website_document_indexing_task" + ) as sync_task, + ): + mock_redis.get.return_value = None + + _dataset_service_document_DocumentService.sync_website_document(dataset, document, sqlite_session) + + sqlite_session.expire_all() + synced = sqlite_session.get(_dataset_service_document_Document, document.id) + assert synced is not None + assert synced.indexing_status == IndexingStatus.WAITING + assert synced.data_source_info_dict["mode"] == "scrape" + mock_redis.setex.assert_called_once_with(f"document_{document.id}_is_sync", 600, 1) + sync_task.delay.assert_called_once_with(dataset.id, document.id) + + class TestDocumentServiceSaveDocumentWithoutDatasetId: + """Unit tests for dataset creation around save_document_without_dataset_id.""" + + @pytest.fixture + def account_context(self): + account = _dataset_service_document_account() + + with _dataset_service_document_patch("services.knowledge.dataset_service.current_user", account): + yield account + + @config_overrides_context(DEPLOYMENT_EDITION=DeploymentEdition.COMMUNITY) + def test_save_document_without_dataset_id_creates_high_quality_dataset_with_default_retrieval_model( + self, account_context, sqlite_session: Session + ): + knowledge_config = _dataset_service_document_KnowledgeConfig( + indexing_technique="high_quality", + data_source=_dataset_service_document_DataSource( + info_list=_dataset_service_document_InfoList( + data_source_type="upload_file", + file_info_list=FileInfo(file_ids=["file-1"]), + ) + ), + embedding_model="embedding-model", + embedding_model_provider="provider", + summary_index_setting={"enable": True}, + is_multimodal=True, + ) + binding = DatasetCollectionBinding( + provider_name="provider", + model_name="embedding-model", + type="dataset", + collection_name="collection", + ) + binding.id = "binding-1" + first_document = _document_row(name="VeryLongDocumentNameForDataset.txt") + + with ( + _dataset_service_document_patch( + "services.knowledge.dataset_service.FeatureService.get_features", return_value=_make_features() + ), + _dataset_service_document_patch( + "services.knowledge.dataset_service.DatasetCollectionBindingService.get_dataset_collection_binding", + return_value=binding, + ), + _dataset_service_document_patch.object( + _dataset_service_document_DocumentService, + "save_document_with_dataset_id", + return_value=([first_document], "batch-1"), + ) as save_document, + ): + dataset, documents, batch = _dataset_service_document_DocumentService.save_document_without_dataset_id( + tenant_id="tenant-1", + knowledge_config=knowledge_config, + account=account_context, + session=sqlite_session, + ) + + assert documents == [first_document] + assert batch == "batch-1" + assert dataset.collection_binding_id == "binding-1" + assert dataset.retrieval_model["search_method"] == RetrievalMethod.SEMANTIC_SEARCH + assert dataset.retrieval_model["top_k"] == 4 + assert dataset.summary_index_setting == {"enable": True} + assert dataset.is_multimodal is True + assert dataset.name == first_document.name[:18] + "..." + assert ( + dataset.description + == "useful for when you want to answer queries about the VeryLongDocumentNameForDataset.txt" + ) + assert sqlite_session.get(Dataset, dataset.id) is dataset + save_document.assert_called_once_with( + dataset, + knowledge_config, + account_context, + session=sqlite_session, + ) + + @config_overrides_context(DEPLOYMENT_EDITION=DeploymentEdition.COMMUNITY) + def test_save_document_without_dataset_id_uses_provided_retrieval_model( + self, account_context, sqlite_session: Session + ): + retrieval_model = RetrievalModel( + search_method=RetrievalMethod.SEMANTIC_SEARCH, + reranking_enable=True, + reranking_model=RerankingModel( + reranking_provider_name="rerank-provider", + reranking_model_name="rerank-model", + ), + top_k=9, + score_threshold_enabled=True, + score_threshold=0.6, + ) + knowledge_config = _dataset_service_document_KnowledgeConfig( + indexing_technique="economy", + data_source=_dataset_service_document_DataSource( + info_list=_dataset_service_document_InfoList( + data_source_type="upload_file", + file_info_list=FileInfo(file_ids=["file-1"]), + ) + ), + retrieval_model=retrieval_model, + ) + first_document = _document_row(name="Doc") + + with ( + _dataset_service_document_patch( + "services.knowledge.dataset_service.FeatureService.get_features", return_value=_make_features() + ), + _dataset_service_document_patch.object( + _dataset_service_document_DocumentService, + "save_document_with_dataset_id", + return_value=([first_document], "batch-1"), + ), + ): + dataset, _, _ = _dataset_service_document_DocumentService.save_document_without_dataset_id( + "tenant-1", + knowledge_config, + account_context, + sqlite_session, + ) + + assert dataset.retrieval_model == retrieval_model.model_dump() + assert dataset.collection_binding_id is None + assert sqlite_session.get(Dataset, dataset.id) is dataset + + @config_overrides_context(DEPLOYMENT_EDITION=DeploymentEdition.CLOUD) + def test_save_document_without_dataset_id_rejects_sandbox_batch_upload( + self, account_context, unbound_session: Session + ): + knowledge_config = _dataset_service_document_KnowledgeConfig( + indexing_technique="economy", + data_source=_dataset_service_document_DataSource( + info_list=_dataset_service_document_InfoList( + data_source_type="upload_file", + file_info_list=FileInfo(file_ids=["file-1", "file-2"]), + ) + ), + ) + + with ( + _dataset_service_document_patch( + "services.knowledge.dataset_service.FeatureService.get_features", + return_value=_make_features(plan=CloudPlan.SANDBOX), + ), + _dataset_service_document_patch.object( + _dataset_service_document_DocumentService, "check_documents_upload_quota" + ) as check_quota, + ): + with pytest.raises(ValueError, match="does not support batch upload"): + _dataset_service_document_DocumentService.save_document_without_dataset_id( + "tenant-1", knowledge_config, account_context, unbound_session + ) + + check_quota.assert_not_called() + + class TestDocumentServiceUpdateDocumentWithDatasetId: + """Unit tests for the document-update orchestration path.""" + + @pytest.fixture + def account_context(self): + account = _dataset_service_document_account() + + with _dataset_service_document_patch("services.knowledge.dataset_service.current_user", account): + yield account + + def test_update_document_with_dataset_id_raises_when_document_is_missing( + self, account_context, sqlite_session: Session + ): + dataset = _dataset_row() + document_data = _dataset_service_document_KnowledgeConfig( + original_document_id="doc-1", + indexing_technique="economy", + data_source=_dataset_service_document_DataSource( + info_list=_dataset_service_document_InfoList( + data_source_type="upload_file", + file_info_list=FileInfo(file_ids=["file-1"]), + ) + ), + ) + with _dataset_service_document_patch.object( + _dataset_service_document_DatasetService, "check_dataset_model_setting" + ) as check_model_setting: + with pytest.raises(NotFound, match="Document not found"): + _dataset_service_document_DocumentService.update_document_with_dataset_id( + dataset, + document_data, + account_context, + session=sqlite_session, + ) + + check_model_setting.assert_called_once_with(dataset) + + def test_update_document_with_dataset_id_rejects_non_available_documents( + self, account_context, sqlite_session: Session + ): + dataset = _dataset_row() + document = _document_row(document_id="doc-1", indexing_status=IndexingStatus.INDEXING) + sqlite_session.add_all([dataset, document]) + sqlite_session.commit() + document_data = _dataset_service_document_KnowledgeConfig( + original_document_id="doc-1", + indexing_technique="economy", + data_source=_dataset_service_document_DataSource( + info_list=_dataset_service_document_InfoList( + data_source_type="upload_file", + file_info_list=FileInfo(file_ids=["file-1"]), + ) + ), + ) + with _dataset_service_document_patch.object( + _dataset_service_document_DatasetService, "check_dataset_model_setting" + ): + with pytest.raises(ValueError, match="Document is not available"): + _dataset_service_document_DocumentService.update_document_with_dataset_id( + dataset, + document_data, + account_context, + session=sqlite_session, + ) + + def test_update_document_with_dataset_id_upload_file_process_rule_and_name_override( + self, account_context, sqlite_session: Session + ): + dataset = _dataset_row() + document = _document_row(document_id="doc-1") + upload_file = UploadFile( + tenant_id=dataset.tenant_id, + storage_type="opendal", + key="key", + name="upload.txt", + size=1, + extension="txt", + mime_type="text/plain", + created_by_role="account", + created_by=account_context.id, + created_at=datetime(2026, 1, 1), + used=False, + ) + upload_file.id = "file-1" + segment = DocumentSegment( + tenant_id=dataset.tenant_id, + dataset_id=dataset.id, + document_id=document.id, + position=1, + content="content", + word_count=1, + tokens=1, + created_by=account_context.id, + ) + sqlite_session.add_all([dataset, document, upload_file, segment]) + sqlite_session.commit() + document_data = _dataset_service_document_KnowledgeConfig( + original_document_id="doc-1", + indexing_technique="economy", + data_source=_dataset_service_document_DataSource( + info_list=_dataset_service_document_InfoList( + data_source_type="upload_file", + file_info_list=FileInfo(file_ids=["file-1"]), + ) + ), + process_rule=ProcessRule( + mode="custom", + rules=Rule( + pre_processing_rules=[PreProcessingRule(id="remove_stopwords", enabled=True)], + segmentation=Segmentation(separator="\n", max_tokens=128), + ), + ), + name="Renamed document", + doc_form=_dataset_service_document_IndexStructureType.QA_INDEX, + ) + updated_at = datetime(2026, 2, 1) + + with ( + _dataset_service_document_patch.object( + _dataset_service_document_DatasetService, "check_dataset_model_setting" + ), + _dataset_service_document_patch( + "services.knowledge.dataset_service.naive_utc_now", return_value=updated_at + ), + _dataset_service_document_patch( + "services.knowledge.dataset_service.document_indexing_update_task" + ) as update_task, + ): + result = _dataset_service_document_DocumentService.update_document_with_dataset_id( + dataset, + document_data, + account_context, + session=sqlite_session, + ) + + assert result is document + assert document.dataset_process_rule_id is not None + assert ( + sqlite_session.get(_dataset_service_document_DatasetProcessRule, document.dataset_process_rule_id) + is not None + ) + assert document.data_source_type == "upload_file" + assert document.data_source_info == '{"upload_file_id": "file-1"}' + assert document.name == "Renamed document" + assert document.indexing_status == "waiting" + assert document.completed_at is None + assert document.processing_started_at is None + assert document.parsing_completed_at is None + assert document.cleaning_completed_at is None + assert document.splitting_completed_at is None + assert document.updated_at == updated_at + assert document.created_from == "web" + assert document.doc_form == _dataset_service_document_IndexStructureType.QA_INDEX + sqlite_session.expire_all() + persisted = sqlite_session.get(_dataset_service_document_Document, document.id) + assert persisted is not None + assert persisted.name == "Renamed document" + assert sqlite_session.get(DocumentSegment, segment.id).status == "re_segment" + update_task.delay.assert_called_once_with(document.dataset_id, document.id) + + def test_update_document_with_dataset_id_notion_import_requires_binding( + self, account_context, sqlite_session: Session + ): + dataset = _dataset_row() + document = _document_row(document_id="doc-1") + sqlite_session.add_all([dataset, document]) + sqlite_session.commit() + document_data = _dataset_service_document_KnowledgeConfig( + original_document_id="doc-1", + indexing_technique="economy", + data_source=_dataset_service_document_DataSource( + info_list=_dataset_service_document_InfoList( + data_source_type="notion_import", + notion_info_list=[ + NotionInfo( + credential_id="credential-1", + workspace_id="workspace-1", + pages=[NotionPage(page_id="page-1", page_name="Page 1", page_icon=None, type="page")], + ) + ], + ) + ), + ) + + with _dataset_service_document_patch.object( + _dataset_service_document_DatasetService, "check_dataset_model_setting" + ): + with pytest.raises(ValueError, match="Data source binding not found"): + _dataset_service_document_DocumentService.update_document_with_dataset_id( + dataset, + document_data, + account_context, + session=sqlite_session, + ) + + def test_update_document_with_dataset_id_website_crawl_updates_segments_and_dispatches_task( + self, account_context, sqlite_session: Session + ): + dataset = _dataset_row() + document = _document_row(document_id="doc-1") + segment = DocumentSegment( + tenant_id=dataset.tenant_id, + dataset_id=dataset.id, + document_id=document.id, + position=1, + content="content", + word_count=1, + tokens=1, + created_by=account_context.id, + ) + sqlite_session.add_all([dataset, document, segment]) + sqlite_session.commit() + document_data = _dataset_service_document_KnowledgeConfig( + original_document_id="doc-1", + indexing_technique="economy", + data_source=_dataset_service_document_DataSource( + info_list=_dataset_service_document_InfoList( + data_source_type="website_crawl", + website_info_list=WebsiteInfo( + provider="firecrawl", + job_id="job-1", + urls=["https://example.com"], + only_main_content=False, + ), + ) + ), + doc_form=_dataset_service_document_IndexStructureType.PARENT_CHILD_INDEX, + ) + + with ( + _dataset_service_document_patch.object( + _dataset_service_document_DatasetService, "check_dataset_model_setting" + ), + _dataset_service_document_patch( + "services.knowledge.dataset_service.naive_utc_now", return_value=datetime(2026, 2, 1) + ), + _dataset_service_document_patch( + "services.knowledge.dataset_service.document_indexing_update_task" + ) as update_task, + ): + result = _dataset_service_document_DocumentService.update_document_with_dataset_id( + dataset, + document_data, + account_context, + session=sqlite_session, + ) + + assert result is document + assert document.data_source_type == "website_crawl" + assert document.data_source_info == ( + '{"url": "https://example.com", "provider": "firecrawl", "job_id": "job-1", ' + '"only_main_content": false, "mode": "crawl"}' + ) + assert document.name == "" + assert document.doc_form == _dataset_service_document_IndexStructureType.PARENT_CHILD_INDEX + sqlite_session.expire_all() + assert sqlite_session.get(DocumentSegment, segment.id).status == "re_segment" + update_task.delay.assert_called_once_with("dataset-1", "doc-1") + + class TestDocumentServiceCreateValidation: + """Unit tests for document creation validation helpers.""" + + def test_document_create_args_validate_requires_data_source_or_process_rule(self): + knowledge_config = _dataset_service_document_KnowledgeConfig(indexing_technique="economy") + + with pytest.raises(ValueError, match="Data source or Process rule is required"): + _dataset_service_document_DocumentService.document_create_args_validate(knowledge_config) + + def test_document_create_args_validate_delegates_to_sub_validators(self): + knowledge_config = _dataset_service_document_KnowledgeConfig( + indexing_technique="economy", + data_source=_dataset_service_document_DataSource( + info_list=_dataset_service_document_InfoList( + data_source_type="upload_file", + file_info_list=FileInfo(file_ids=["file-1"]), + ) + ), + process_rule=ProcessRule(mode="automatic"), + ) + + with ( + _dataset_service_document_patch.object( + _dataset_service_document_DocumentService, "data_source_args_validate" + ) as validate_data_source, + _dataset_service_document_patch.object( + _dataset_service_document_DocumentService, "process_rule_args_validate" + ) as validate_process_rule, + ): + _dataset_service_document_DocumentService.document_create_args_validate(knowledge_config) + + validate_data_source.assert_called_once_with(knowledge_config) + validate_process_rule.assert_called_once_with(knowledge_config) + + @pytest.mark.parametrize( + ("data_source_type", "message"), + [ + ("upload_file", "File source info is required"), + ("notion_import", "Notion source info is required"), + ("website_crawl", "Website source info is required"), + ], + ) + def test_data_source_args_validate_requires_source_specific_info(self, data_source_type, message): + info_list = _dataset_service_document_InfoList.model_validate({"data_source_type": data_source_type}) + knowledge_config = _dataset_service_document_KnowledgeConfig( + indexing_technique="economy", + data_source=_dataset_service_document_DataSource(info_list=info_list), + ) + + with pytest.raises(ValueError, match=message): + _dataset_service_document_DocumentService.data_source_args_validate(knowledge_config) + + def test_process_rule_args_validate_clears_rules_for_automatic_mode(self): + knowledge_config = _dataset_service_document_KnowledgeConfig( + indexing_technique="economy", + data_source=_dataset_service_document_DataSource( + info_list=_dataset_service_document_InfoList( + data_source_type="upload_file", + file_info_list=FileInfo(file_ids=["file-1"]), + ) + ), + process_rule=ProcessRule( + mode="automatic", + rules=Rule( + pre_processing_rules=[PreProcessingRule(id="remove_stopwords", enabled=True)], + segmentation=Segmentation(separator="\n", max_tokens=128), + ), + ), + ) + + _dataset_service_document_DocumentService.process_rule_args_validate(knowledge_config) + + assert knowledge_config.process_rule is not None + assert knowledge_config.process_rule.rules is None + + def test_process_rule_args_validate_deduplicates_rules_and_skips_max_tokens_for_full_doc_hierarchical(self): + knowledge_config = _dataset_service_document_KnowledgeConfig( + indexing_technique="economy", + data_source=_dataset_service_document_DataSource( + info_list=_dataset_service_document_InfoList( + data_source_type="upload_file", + file_info_list=FileInfo(file_ids=["file-1"]), + ) + ), + process_rule=ProcessRule( + mode="hierarchical", + rules=Rule( + pre_processing_rules=[ + PreProcessingRule(id="remove_stopwords", enabled=True), + PreProcessingRule(id="remove_stopwords", enabled=False), + ], + segmentation=Segmentation(separator="\n", max_tokens=0), + parent_mode="full-doc", + ), + ), + ) + + _dataset_service_document_DocumentService.process_rule_args_validate(knowledge_config) + + assert knowledge_config.process_rule is not None + assert knowledge_config.process_rule.rules is not None + assert len(knowledge_config.process_rule.rules.pre_processing_rules) == 1 + assert knowledge_config.process_rule.rules.pre_processing_rules[0].enabled is False + + def test_process_rule_args_validate_hierarchical_defaults_parent_mode_to_paragraph(self): + knowledge_config = _dataset_service_document_KnowledgeConfig( + indexing_technique="economy", + data_source=_dataset_service_document_DataSource( + info_list=_dataset_service_document_InfoList( + data_source_type="upload_file", + file_info_list=FileInfo(file_ids=["file-1"]), + ) + ), + process_rule=ProcessRule( + mode="hierarchical", + rules=Rule( + pre_processing_rules=[ + PreProcessingRule(id="remove_extra_spaces", enabled=True), + ], + segmentation=Segmentation(separator="\n", max_tokens=1024), + subchunk_segmentation=Segmentation(separator="\n", max_tokens=512), + ), + ), + ) + + _dataset_service_document_DocumentService.process_rule_args_validate(knowledge_config) + + assert knowledge_config.process_rule is not None + assert knowledge_config.process_rule.rules is not None + assert knowledge_config.process_rule.rules.parent_mode == "paragraph" + + class TestDocumentServiceSaveDocumentWithDatasetId: + """Unit tests for non-SQL validation branches in save_document_with_dataset_id.""" + + @pytest.fixture + def account_context(self): + account = _dataset_service_document_account() + + with ( + _dataset_service_document_patch("services.knowledge.dataset_service.current_user", account), + _dataset_service_document_patch.object(_dataset_service_document_DatasetService, "check_doc_form"), + ): + yield account + + @config_overrides_context(DEPLOYMENT_EDITION=DeploymentEdition.CLOUD) + def test_save_document_with_dataset_id_requires_file_info_for_upload_source( + self, account_context, unbound_session: Session + ): + dataset = _dataset_row() + knowledge_config = _make_upload_knowledge_config(file_ids=None) + + with _dataset_service_document_patch( + "services.knowledge.dataset_service.FeatureService.get_features", return_value=_make_features() + ): + with pytest.raises(ValueError, match="File source info is required"): + _dataset_service_document_DocumentService.save_document_with_dataset_id( + dataset, + knowledge_config, + account_context, + session=unbound_session, + ) + + @config_overrides_context(DEPLOYMENT_EDITION=DeploymentEdition.CLOUD) + def test_save_document_with_dataset_id_blocks_batch_upload_for_sandbox_plan( + self, account_context, unbound_session: Session + ): + dataset = _dataset_row() + knowledge_config = _make_upload_knowledge_config(file_ids=["file-1", "file-2"]) + + with ( + _dataset_service_document_patch( + "services.knowledge.dataset_service.FeatureService.get_features", + return_value=_make_features(plan=CloudPlan.SANDBOX), + ), + _dataset_service_document_patch.object( + _dataset_service_document_DocumentService, "check_documents_upload_quota" + ) as check_quota, + ): + with pytest.raises(ValueError, match="does not support batch upload"): + _dataset_service_document_DocumentService.save_document_with_dataset_id( + dataset, + knowledge_config, + account_context, + session=unbound_session, + ) + + check_quota.assert_not_called() + + @config_overrides_context(DEPLOYMENT_EDITION=DeploymentEdition.CLOUD) + def test_save_document_with_dataset_id_enforces_batch_upload_limit( + self, + account_context, + unbound_session: Session, + config_overrides: Callable[..., None], + ): + config_overrides(BATCH_UPLOAD_LIMIT=1) + dataset = _dataset_row() + knowledge_config = _make_upload_knowledge_config(file_ids=["file-1", "file-2"]) + + with ( + _dataset_service_document_patch( + "services.knowledge.dataset_service.FeatureService.get_features", return_value=_make_features() + ), + _dataset_service_document_patch.object( + _dataset_service_document_DocumentService, "check_documents_upload_quota" + ) as check_quota, + ): + with pytest.raises(ValueError, match="batch upload limit of 1"): + _dataset_service_document_DocumentService.save_document_with_dataset_id( + dataset, + knowledge_config, + account_context, + session=unbound_session, + ) + + check_quota.assert_not_called() + + @config_overrides_context(DEPLOYMENT_EDITION=DeploymentEdition.COMMUNITY) + def test_save_document_with_dataset_id_updates_existing_document_and_data_source_type( + self, account_context, sqlite_session: Session + ): + dataset = _dataset_row(data_source_type=None) + knowledge_config = _make_upload_knowledge_config(original_document_id="doc-1", file_ids=["file-1"]) + updated_document = _document_row(document_id="doc-1") + updated_document.batch = "batch-existing" + + with ( + _dataset_service_document_patch( + "services.knowledge.dataset_service.FeatureService.get_features", return_value=_make_features() + ), + _dataset_service_document_patch( + "services.knowledge.dataset_service.document_indexing_update_task.delay" + ), + _dataset_service_document_patch.object( + _dataset_service_document_DocumentService, "update_prepared_document", return_value=updated_document + ) as update_document, + ): + documents, batch = _dataset_service_document_DocumentService.save_document_with_dataset_id( + dataset, + knowledge_config, + account_context, + session=sqlite_session, + ) + + assert dataset.data_source_type == "upload_file" + assert documents == [updated_document] + assert batch == "batch-existing" + update_document.assert_called_once_with(dataset, knowledge_config, account_context, session=sqlite_session) + + @config_overrides_context(DEPLOYMENT_EDITION=DeploymentEdition.COMMUNITY) + def test_save_document_with_dataset_id_requires_data_source_for_new_documents( + self, account_context, unbound_session: Session + ): + dataset = _dataset_row() + knowledge_config = _make_upload_knowledge_config(data_source=None) + + with _dataset_service_document_patch( + "services.knowledge.dataset_service.FeatureService.get_features", return_value=_make_features() + ): + with pytest.raises(ValueError, match="Data source is required when creating new documents"): + _dataset_service_document_DocumentService.save_document_with_dataset_id( + dataset, + knowledge_config, + account_context, + session=unbound_session, + ) + + @config_overrides_context(DEPLOYMENT_EDITION=DeploymentEdition.COMMUNITY) + def test_save_document_with_dataset_id_requires_existing_process_rule_for_custom_mode( + self, account_context, sqlite_session: Session + ): + dataset = _dataset_row() + sqlite_session.add(dataset) + sqlite_session.commit() + knowledge_config = _make_upload_knowledge_config( + file_ids=["file-1"], + process_rule=ProcessRule(mode="custom"), + ) + + with _dataset_service_document_patch( + "services.knowledge.dataset_service.FeatureService.get_features", return_value=_make_features() + ): + with pytest.raises(ValueError, match="No process rule found"): + _dataset_service_document_DocumentService.save_document_with_dataset_id( + dataset, + knowledge_config, + account_context, + session=sqlite_session, + ) + + @config_overrides_context(DEPLOYMENT_EDITION=DeploymentEdition.COMMUNITY) + def test_save_document_with_dataset_id_upload_file_creates_and_reindexes_documents( + self, account_context, sqlite_session: Session + ): + dataset = _dataset_row(data_source_type=DataSourceType.UPLOAD_FILE) + dataset_process_rule = _process_rule() + knowledge_config = _make_upload_knowledge_config(file_ids=["file-1", "file-2"]) + duplicate_document = _document_row(document_id="doc-duplicate", name="existing.txt") + upload_file_a = _upload_file(file_id="file-1", name="existing.txt") + upload_file_b = _upload_file(file_id="file-2", name="new.txt") + sqlite_session.add_all([dataset, dataset_process_rule, duplicate_document, upload_file_a, upload_file_b]) + sqlite_session.commit() + + with ( + _dataset_service_document_patch( + "services.knowledge.dataset_service.FeatureService.get_features", return_value=_make_features() + ), + _dataset_service_document_patch("services.knowledge.dataset_service.redis_client") as mock_redis, + _dataset_service_document_patch( + "services.knowledge.dataset_service.DocumentIndexingTaskProxy" + ) as document_proxy_cls, + _dataset_service_document_patch( + "services.knowledge.dataset_service.DuplicateDocumentIndexingTaskProxy" + ) as duplicate_proxy_cls, + _dataset_service_document_patch( + "services.knowledge.dataset_service.naive_utc_now", return_value=datetime(2026, 2, 1) + ), + _dataset_service_document_patch( + "services.knowledge.dataset_service.time.strftime", return_value="20260101010101" + ), + _dataset_service_document_patch( + "services.knowledge.dataset_service.secrets.randbelow", return_value=23 + ), + ): + mock_redis.lock.return_value = _make_lock_context() + documents, batch = _dataset_service_document_DocumentService.save_document_with_dataset_id( + dataset, + knowledge_config, + account_context, + dataset_process_rule=dataset_process_rule, + session=sqlite_session, + ) + + assert [document.name for document in documents] == ["existing.txt", "new.txt"] + assert batch == "20260101010101100023" + assert duplicate_document.dataset_process_rule_id == "rule-1" + assert duplicate_document.updated_at == datetime(2026, 2, 1) + assert duplicate_document.batch == batch + assert duplicate_document.indexing_status == IndexingStatus.WAITING + created_document = next(document for document in documents if document.name == "new.txt") + sqlite_session.expire_all() + assert sqlite_session.get(_dataset_service_document_Document, created_document.id) is not None + document_proxy_cls.assert_called_once_with(dataset.tenant_id, dataset.id, [created_document.id]) + document_proxy_cls.return_value.delay.assert_called_once() + duplicate_proxy_cls.assert_called_once_with(dataset.tenant_id, dataset.id, ["doc-duplicate"]) + duplicate_proxy_cls.return_value.delay.assert_called_once() + + @config_overrides_context(DEPLOYMENT_EDITION=DeploymentEdition.COMMUNITY) + def test_save_document_with_dataset_id_notion_import_truncates_names_and_cleans_removed_pages( + self, account_context, sqlite_session: Session + ): + dataset = _dataset_row(data_source_type=DataSourceType.NOTION_IMPORT) + dataset_process_rule = _process_rule() + notion_page_name = "a" * 300 + knowledge_config = _dataset_service_document_KnowledgeConfig( + indexing_technique="economy", + data_source=_dataset_service_document_DataSource( + info_list=_dataset_service_document_InfoList( + data_source_type="notion_import", + notion_info_list=[ + NotionInfo( + credential_id="credential-1", + workspace_id="workspace-1", + pages=[ + NotionPage(page_id="page-keep", page_name="Keep page", type="page"), + NotionPage( + page_id="page-new", + page_name=notion_page_name, + page_icon=NotionIcon(type="emoji", emoji="page"), + type="page", + ), + ], + ) + ], + ) + ), + doc_form=_dataset_service_document_IndexStructureType.PARAGRAPH_INDEX, + doc_language="English", + ) + existing_keep = _document_row( + document_id="doc-keep", + data_source_type=DataSourceType.NOTION_IMPORT, + ) + existing_keep.data_source_info = json.dumps({"notion_page_id": "page-keep"}) + existing_remove = _document_row( + document_id="doc-remove", + data_source_type=DataSourceType.NOTION_IMPORT, + ) + existing_remove.data_source_info = json.dumps({"notion_page_id": "page-remove"}) + sqlite_session.add_all([dataset, dataset_process_rule, existing_keep, existing_remove]) + sqlite_session.commit() + + with ( + _dataset_service_document_patch( + "services.knowledge.dataset_service.FeatureService.get_features", return_value=_make_features() + ), + _dataset_service_document_patch("services.knowledge.dataset_service.redis_client") as mock_redis, + _dataset_service_document_patch( + "services.knowledge.dataset_service.clean_notion_document_task" + ) as clean_task, + _dataset_service_document_patch( + "services.knowledge.dataset_service.DocumentIndexingTaskProxy" + ) as document_proxy_cls, + _dataset_service_document_patch( + "services.knowledge.dataset_service.uuid.uuid4", return_value="doc-new" + ), + ): + mock_redis.lock.return_value = _make_lock_context() + documents, _ = _dataset_service_document_DocumentService.save_document_with_dataset_id( + dataset, + knowledge_config, + account_context, + dataset_process_rule=dataset_process_rule, + session=sqlite_session, + ) + + created_document = next(document for document in documents if document.id == "doc-new") + assert created_document in documents + assert len(created_document.name) == 255 + assert sqlite_session.get(_dataset_service_document_Document, created_document.id) is created_document + clean_task.delay.assert_called_once_with(["doc-remove"], dataset.id) + document_proxy_cls.assert_called_once_with(dataset.tenant_id, dataset.id, ["doc-new"]) + document_proxy_cls.return_value.delay.assert_called_once() + + @config_overrides_context(DEPLOYMENT_EDITION=DeploymentEdition.COMMUNITY) + def test_save_document_with_dataset_id_website_crawl_truncates_long_urls( + self, account_context, sqlite_session: Session + ): + dataset = _dataset_row(data_source_type=DataSourceType.WEBSITE_CRAWL) + dataset_process_rule = _process_rule() + sqlite_session.add_all([dataset, dataset_process_rule]) + sqlite_session.commit() + long_url = "https://example.com/" + ("a" * 260) + short_url = "https://example.com/short" + knowledge_config = _dataset_service_document_KnowledgeConfig( + indexing_technique="economy", + data_source=_dataset_service_document_DataSource( + info_list=_dataset_service_document_InfoList( + data_source_type="website_crawl", + website_info_list=WebsiteInfo( + provider="firecrawl", + job_id="job-1", + urls=[long_url, short_url], + only_main_content=True, + ), + ) + ), + doc_form=_dataset_service_document_IndexStructureType.PARAGRAPH_INDEX, + doc_language="English", + ) + with ( + _dataset_service_document_patch( + "services.knowledge.dataset_service.FeatureService.get_features", return_value=_make_features() + ), + _dataset_service_document_patch("services.knowledge.dataset_service.redis_client") as mock_redis, + _dataset_service_document_patch( + "services.knowledge.dataset_service.DocumentIndexingTaskProxy" + ) as document_proxy_cls, + ): + mock_redis.lock.return_value = _make_lock_context() + + documents, _ = _dataset_service_document_DocumentService.save_document_with_dataset_id( + dataset, + knowledge_config, + account_context, + dataset_process_rule=dataset_process_rule, + session=sqlite_session, + ) + + assert [document.name for document in documents] == [long_url[:200] + "...", short_url] + assert ( + sqlite_session.scalars( + select(_dataset_service_document_Document).order_by(_dataset_service_document_Document.position) + ).all() + == documents + ) + document_proxy_cls.assert_called_once_with( + dataset.tenant_id, dataset.id, [document.id for document in documents] + ) + document_proxy_cls.return_value.delay.assert_called_once() + + class TestDocumentServiceBatchUpdateStatus: + """Unit tests for batch_update_document_status orchestration and helper branches.""" + + def test_prepare_disable_update_requires_completed_document(self): + document = _document_row(indexing_status=IndexingStatus.WAITING) + document.completed_at = None + + with pytest.raises(DocumentIndexingError, match="is not completed"): + _dataset_service_document_DocumentService._prepare_disable_update( + document, user=_dataset_service_document_account(), now=datetime(2026, 2, 1) + ) + + def test_prepare_archive_update_sets_async_task_for_enabled_document(self): + document = _document_row(enabled=True, archived=False) + + result = _dataset_service_document_DocumentService._prepare_archive_update( + document, user=_dataset_service_document_account(), now=datetime(2026, 2, 1) + ) + + assert result is not None + assert result["updates"]["archived"] is True + assert result["set_cache"] is True + assert result["async_task"]["args"] == [document.id] + + def test_prepare_unarchive_update_sets_async_task_for_enabled_document(self): + document = _document_row(enabled=True, archived=True) + + result = _dataset_service_document_DocumentService._prepare_unarchive_update( + document, now=datetime(2026, 2, 1) + ) + + assert result is not None + assert result["updates"]["archived"] is False + assert result["set_cache"] is True + assert result["async_task"]["args"] == [document.id] + + def test_batch_update_document_status_rejects_indexing_documents(self, sqlite_session: Session): + dataset = _dataset_row() + document = _document_row(name="Busy document") + sqlite_session.add_all([dataset, document]) + sqlite_session.commit() + + with _dataset_service_document_patch("services.knowledge.dataset_service.redis_client") as mock_redis: + mock_redis.get.return_value = "1" + + with pytest.raises(DocumentIndexingError, match="Busy document is being indexed"): + _dataset_service_document_DocumentService.batch_update_document_status( + dataset, [document.id], "archive", _dataset_service_document_account(), sqlite_session + ) + + sqlite_session.refresh(document) + assert document.archived is False + + def test_batch_update_document_status_rolls_back_when_commit_fails(self, sqlite_session: Session): + dataset = _dataset_row() + document = _document_row(enabled=False) + sqlite_session.add_all([dataset, document]) + sqlite_session.commit() + + def fail_commit(_session): + raise RuntimeError("commit failed") + + event.listen(sqlite_session, "before_commit", fail_commit) + with ( + _dataset_service_document_patch("services.knowledge.dataset_service.redis_client") as mock_redis, + ): + mock_redis.get.return_value = None + + with pytest.raises(RuntimeError, match="commit failed"): + _dataset_service_document_DocumentService.batch_update_document_status( + dataset, [document.id], "enable", _dataset_service_document_account(), sqlite_session + ) + event.remove(sqlite_session, "before_commit", fail_commit) + + sqlite_session.refresh(document) + assert document.enabled is False + + def test_batch_update_document_status_raises_async_task_error_after_commit(self, sqlite_session: Session): + dataset = _dataset_row() + document = _document_row(enabled=False) + sqlite_session.add_all([dataset, document]) + sqlite_session.commit() + + with ( + _dataset_service_document_patch("services.knowledge.dataset_service.redis_client") as mock_redis, + _dataset_service_document_patch( + "services.knowledge.dataset_service.add_document_to_index_task" + ) as add_task, + ): + mock_redis.get.return_value = None + add_task.delay.side_effect = RuntimeError("task failed") + + with pytest.raises(RuntimeError, match="task failed"): + _dataset_service_document_DocumentService.batch_update_document_status( + dataset, [document.id], "enable", _dataset_service_document_account(), sqlite_session + ) + + sqlite_session.refresh(document) + assert document.enabled is True + mock_redis.setex.assert_called_once_with(f"document_{document.id}_indexing", 600, 1) + + class TestDocumentServiceTenantAndUpdateEdges: + """Unit tests for tenant-count and update edge cases.""" + + @pytest.fixture + def account_context(self): + account = _dataset_service_document_account() + + with _dataset_service_document_patch("services.knowledge.dataset_service.current_user", account): + yield account + + def test_get_tenant_documents_count_scopes_state_and_tenant(self, account_context, sqlite_session: Session): + sqlite_session.add_all( + [ + _document_row(document_id="one"), + _document_row(document_id="two"), + _document_row(document_id="disabled", enabled=False), + _document_row(document_id="archived", archived=True), + _document_row(document_id="unfinished", indexing_status=IndexingStatus.WAITING), + _document_row(document_id="foreign", tenant_id="tenant-2"), + ] + ) + sqlite_session.commit() + + result = _dataset_service_document_DocumentService.get_tenant_documents_count(sqlite_session) + + assert result == 2 + + def test_update_document_with_dataset_id_uses_automatic_process_rule_payload( + self, account_context, sqlite_session: Session + ): + dataset = _dataset_row() + document = _document_row(document_id="doc-1") + upload_file = _upload_file(file_id="file-1") + sqlite_session.add_all([dataset, document, upload_file]) + sqlite_session.commit() + document_data = _dataset_service_document_KnowledgeConfig( + original_document_id="doc-1", + indexing_technique="economy", + data_source=_dataset_service_document_DataSource( + info_list=_dataset_service_document_InfoList( + data_source_type="upload_file", + file_info_list=FileInfo(file_ids=["file-1"]), + ) + ), + process_rule=ProcessRule( + mode="automatic", + rules=Rule( + pre_processing_rules=[PreProcessingRule(id="remove_stopwords", enabled=True)], + segmentation=Segmentation(separator="\n", max_tokens=128), + ), + ), + doc_form=_dataset_service_document_IndexStructureType.PARAGRAPH_INDEX, + ) + updated_at = datetime(2026, 2, 1) + + with ( + _dataset_service_document_patch.object( + _dataset_service_document_DatasetService, "check_dataset_model_setting" + ), + _dataset_service_document_patch( + "services.knowledge.dataset_service.naive_utc_now", return_value=updated_at + ), + _dataset_service_document_patch( + "services.knowledge.dataset_service.document_indexing_update_task" + ) as update_task, + ): + result = _dataset_service_document_DocumentService.update_document_with_dataset_id( + dataset, + document_data, + account_context, + session=sqlite_session, + ) + + assert result is document + assert document.dataset_process_rule_id is not None + assert document.name == "upload.txt" + process_rule = sqlite_session.get( + _dataset_service_document_DatasetProcessRule, document.dataset_process_rule_id + ) + assert process_rule is not None + assert process_rule.mode == "automatic" + assert process_rule.rules == json.dumps(_dataset_service_document_DatasetProcessRule.AUTOMATIC_RULES) + update_task.delay.assert_called_once_with("dataset-1", "doc-1") + + def test_update_document_with_dataset_id_requires_upload_file_info( + self, account_context, sqlite_session: Session + ): + dataset = _dataset_row() + document = _document_row(document_id="doc-1") + sqlite_session.add_all([dataset, document]) + sqlite_session.commit() + document_data = _dataset_service_document_KnowledgeConfig( + original_document_id="doc-1", + indexing_technique="economy", + data_source=_dataset_service_document_DataSource( + info_list=_dataset_service_document_InfoList(data_source_type="upload_file") + ), + ) + + with _dataset_service_document_patch.object( + _dataset_service_document_DatasetService, "check_dataset_model_setting" + ): + with pytest.raises(ValueError, match="No file info list found"): + _dataset_service_document_DocumentService.update_document_with_dataset_id( + dataset, + document_data, + account_context, + session=sqlite_session, + ) + + def test_update_document_with_dataset_id_raises_when_upload_file_is_missing( + self, account_context, sqlite_session: Session + ): + dataset = _dataset_row() + document = _document_row(document_id="doc-1") + sqlite_session.add_all([dataset, document]) + sqlite_session.commit() + document_data = _dataset_service_document_KnowledgeConfig( + original_document_id="doc-1", + indexing_technique="economy", + data_source=_dataset_service_document_DataSource( + info_list=_dataset_service_document_InfoList( + data_source_type="upload_file", + file_info_list=FileInfo(file_ids=["file-1"]), + ) + ), + ) + + with _dataset_service_document_patch.object( + _dataset_service_document_DatasetService, "check_dataset_model_setting" + ): + with pytest.raises(FileNotExistsError): + _dataset_service_document_DocumentService.update_document_with_dataset_id( + dataset, + document_data, + account_context, + session=sqlite_session, + ) + + def test_update_document_with_dataset_id_requires_notion_info_list( + self, account_context, sqlite_session: Session + ): + dataset = _dataset_row() + document = _document_row(document_id="doc-1") + sqlite_session.add_all([dataset, document]) + sqlite_session.commit() + document_data = _dataset_service_document_KnowledgeConfig( + original_document_id="doc-1", + indexing_technique="economy", + data_source=_dataset_service_document_DataSource( + info_list=_dataset_service_document_InfoList(data_source_type="notion_import") + ), + ) + + with _dataset_service_document_patch.object( + _dataset_service_document_DatasetService, "check_dataset_model_setting" + ): + with pytest.raises(ValueError, match="No notion info list found"): + _dataset_service_document_DocumentService.update_document_with_dataset_id( + dataset, + document_data, + account_context, + session=sqlite_session, + ) + + def test_update_document_with_dataset_id_notion_import_updates_page_info( + self, account_context, sqlite_session: Session + ): + dataset = _dataset_row() + document = _document_row(document_id="doc-1") + binding = DataSourceOauthBinding( + tenant_id=dataset.tenant_id, + access_token="token", + provider="notion", + source_info={"workspace_id": '"workspace-1"'}, + disabled=False, + ) + sqlite_session.add_all([dataset, document, binding]) + sqlite_session.commit() + document_data = _dataset_service_document_KnowledgeConfig( + original_document_id="doc-1", + indexing_technique="economy", + data_source=_dataset_service_document_DataSource( + info_list=_dataset_service_document_InfoList( + data_source_type="notion_import", + notion_info_list=[ + NotionInfo( + credential_id="credential-1", + workspace_id="workspace-1", + pages=[ + NotionPage(page_id="page-1", page_name="Page 1", page_icon=None, type="page"), + NotionPage(page_id="page-2", page_name="Page 2", page_icon=None, type="database"), + ], + ) + ], + ) + ), + doc_form=_dataset_service_document_IndexStructureType.PARAGRAPH_INDEX, + ) + + with ( + _dataset_service_document_patch.object( + _dataset_service_document_DatasetService, "check_dataset_model_setting" + ), + _dataset_service_document_patch( + "services.knowledge.dataset_service.naive_utc_now", return_value=datetime(2026, 2, 1) + ), + _dataset_service_document_patch( + "services.knowledge.dataset_service.document_indexing_update_task" + ) as update_task, + ): + result = _dataset_service_document_DocumentService.update_document_with_dataset_id( + dataset, + document_data, + account_context, + session=sqlite_session, + ) + + assert result is document + assert document.data_source_type == "notion_import" + assert document.name == "" + assert document.data_source_info == json.dumps( + { + "credential_id": "credential-1", + "notion_workspace_id": "workspace-1", + "notion_page_id": "page-2", + "notion_page_icon": None, + "type": "database", + } + ) + sqlite_session.refresh(document) + update_task.delay.assert_called_once_with("dataset-1", "doc-1") + + class TestDocumentServiceSaveWithoutDatasetBilling: + """Unit tests for batch-count and quota branches in save_document_without_dataset_id.""" + + @pytest.fixture + def account_context(self): + account = _dataset_service_document_account() + + with _dataset_service_document_patch("services.knowledge.dataset_service.current_user", account): + yield account + + @config_overrides_context(DEPLOYMENT_EDITION=DeploymentEdition.CLOUD) + def test_save_document_without_dataset_id_counts_notion_pages_for_quota( + self, + account_context, + sqlite_session: Session, + config_overrides: Callable[..., None], + ): + config_overrides(BATCH_UPLOAD_LIMIT="10") + knowledge_config = _dataset_service_document_KnowledgeConfig( + indexing_technique="economy", + data_source=_dataset_service_document_DataSource( + info_list=_dataset_service_document_InfoList( + data_source_type="notion_import", + notion_info_list=[ + NotionInfo( + credential_id="credential-1", + workspace_id="workspace-1", + pages=[ + NotionPage(page_id="page-1", page_name="Page 1", page_icon=None, type="page"), + NotionPage(page_id="page-2", page_name="Page 2", page_icon=None, type="page"), + ], + ), + NotionInfo( + credential_id="credential-2", + workspace_id="workspace-2", + pages=[NotionPage(page_id="page-3", page_name="Page 3", page_icon=None, type="page")], + ), + ], + ) + ), + ) + features = _make_features() + document = _document_row(name="Doc") + + with ( + _dataset_service_document_patch( + "services.knowledge.dataset_service.FeatureService.get_features", return_value=features + ), + _dataset_service_document_patch.object( + _dataset_service_document_DocumentService, "check_documents_upload_quota" + ) as check_quota, + _dataset_service_document_patch.object( + _dataset_service_document_DocumentService, + "save_document_with_dataset_id", + return_value=([document], "batch-1"), + ), + ): + dataset, _, _ = _dataset_service_document_DocumentService.save_document_without_dataset_id( + "tenant-1", + knowledge_config, + account_context, + sqlite_session, + ) + + check_quota.assert_called_once_with(3, features) + assert sqlite_session.get(Dataset, dataset.id) is dataset + + @config_overrides_context(DEPLOYMENT_EDITION=DeploymentEdition.CLOUD) + def test_save_document_without_dataset_id_enforces_batch_limit_for_website_urls( + self, + account_context, + unbound_session: Session, + config_overrides: Callable[..., None], + ): + config_overrides(BATCH_UPLOAD_LIMIT="1") + knowledge_config = _dataset_service_document_KnowledgeConfig( + indexing_technique="economy", + data_source=_dataset_service_document_DataSource( + info_list=_dataset_service_document_InfoList( + data_source_type="website_crawl", + website_info_list=WebsiteInfo( + provider="firecrawl", + job_id="job-1", + urls=["https://example.com/a", "https://example.com/b"], + only_main_content=True, + ), + ) + ), + ) + + with ( + _dataset_service_document_patch( + "services.knowledge.dataset_service.FeatureService.get_features", return_value=_make_features() + ), + _dataset_service_document_patch.object( + _dataset_service_document_DocumentService, "check_documents_upload_quota" + ) as check_quota, + ): + with pytest.raises(ValueError, match="batch upload limit of 1"): + _dataset_service_document_DocumentService.save_document_without_dataset_id( + "tenant-1", knowledge_config, account_context, unbound_session + ) + + check_quota.assert_not_called() + + class TestDocumentServiceSaveDocumentAdditionalBranches: + """Additional unit tests for dataset bootstrap and process-rule branches.""" + + @pytest.fixture + def account_context(self): + account = _dataset_service_document_account() + + with ( + _dataset_service_document_patch("services.knowledge.dataset_service.current_user", account), + _dataset_service_document_patch.object(_dataset_service_document_DatasetService, "check_doc_form"), + ): + yield account + + @config_overrides_context(DEPLOYMENT_EDITION=DeploymentEdition.COMMUNITY) + def test_save_document_with_dataset_id_initializes_high_quality_dataset_from_default_embedding_model( + self, account_context, sqlite_session: Session + ): + dataset = _dataset_row(data_source_type=None, indexing_technique=None) + knowledge_config = _make_upload_knowledge_config(original_document_id="doc-1", file_ids=["file-1"]) + knowledge_config.indexing_technique = "high_quality" + knowledge_config.embedding_model = None + knowledge_config.embedding_model_provider = None + updated_document = _document_row(document_id="doc-1") + updated_document.batch = "batch-existing" + binding = DatasetCollectionBinding( + provider_name="default-provider", + model_name="default-embedding", + type="dataset", + collection_name="collection", + ) + binding.id = "binding-1" + + with ( + _dataset_service_document_patch( + "services.knowledge.dataset_service.FeatureService.get_features", return_value=_make_features() + ), + _dataset_service_document_patch("services.knowledge.dataset_service.ModelManager") as model_manager_cls, + _dataset_service_document_patch( + "services.knowledge.dataset_service.DatasetCollectionBindingService.get_dataset_collection_binding", + return_value=binding, + ) as get_binding, + _dataset_service_document_patch.object( + _dataset_service_document_DocumentService, "update_prepared_document", return_value=updated_document + ), + _dataset_service_document_patch( + "services.knowledge.dataset_service.document_indexing_update_task.delay" + ), + ): + default_model = object.__new__(ModelInstance) + default_model.model_name = "default-embedding" + default_model.provider = "default-provider" + model_manager_cls.for_tenant.return_value.get_default_model_instance.return_value = default_model + + documents, batch = _dataset_service_document_DocumentService.save_document_with_dataset_id( + dataset, + knowledge_config, + account_context, + session=sqlite_session, + ) + + assert documents == [updated_document] + assert batch == "batch-existing" + assert dataset.data_source_type == "upload_file" + assert dataset.indexing_technique == "high_quality" + assert dataset.embedding_model == "default-embedding" + assert dataset.embedding_model_provider == "default-provider" + assert dataset.collection_binding_id == "binding-1" + assert dataset.retrieval_model == { + "search_method": "semantic_search", + "reranking_enable": False, + "reranking_model": {"reranking_provider_name": "", "reranking_model_name": ""}, + "top_k": 4, + "score_threshold_enabled": False, + } + get_binding.assert_called_once_with("default-provider", "default-embedding", sqlite_session) + + @config_overrides_context(DEPLOYMENT_EDITION=DeploymentEdition.COMMUNITY) + def test_save_document_with_dataset_id_uses_explicit_embedding_and_retrieval_model( + self, account_context, sqlite_session: Session + ): + dataset = _dataset_row(indexing_technique=None) + knowledge_config = _make_upload_knowledge_config(original_document_id="doc-1", file_ids=["file-1"]) + knowledge_config.indexing_technique = "high_quality" + knowledge_config.embedding_model = "explicit-model" + knowledge_config.embedding_model_provider = "explicit-provider" + knowledge_config.retrieval_model = RetrievalModel( + search_method="semantic_search", + reranking_enable=True, + reranking_model=RerankingModel( + reranking_provider_name="rerank-provider", + reranking_model_name="rerank-model", + ), + top_k=7, + score_threshold_enabled=True, + score_threshold=0.3, + ) + binding = DatasetCollectionBinding( + provider_name="explicit-provider", + model_name="explicit-model", + type="dataset", + collection_name="collection", + ) + binding.id = "binding-2" + updated_document = _document_row(document_id="doc-1") + + with ( + _dataset_service_document_patch( + "services.knowledge.dataset_service.FeatureService.get_features", return_value=_make_features() + ), + _dataset_service_document_patch("services.knowledge.dataset_service.ModelManager") as model_manager_cls, + _dataset_service_document_patch( + "services.knowledge.dataset_service.DatasetCollectionBindingService.get_dataset_collection_binding", + return_value=binding, + ) as get_binding, + _dataset_service_document_patch.object( + _dataset_service_document_DocumentService, "update_prepared_document", return_value=updated_document + ), + _dataset_service_document_patch( + "services.knowledge.dataset_service.document_indexing_update_task.delay" + ), + ): + _dataset_service_document_DocumentService.save_document_with_dataset_id( + dataset, knowledge_config, account_context, session=sqlite_session + ) + + model_manager_cls.for_tenant.return_value.get_default_model_instance.assert_not_called() + get_binding.assert_called_once_with("explicit-provider", "explicit-model", sqlite_session) + assert dataset.embedding_model == "explicit-model" + assert dataset.embedding_model_provider == "explicit-provider" + assert dataset.retrieval_model == knowledge_config.retrieval_model.model_dump() + + @config_overrides_context(DEPLOYMENT_EDITION=DeploymentEdition.COMMUNITY) + def test_save_document_with_dataset_id_creates_custom_process_rule_for_new_upload_document( + self, account_context, sqlite_session: Session + ): + dataset = _dataset_row(data_source_type=DataSourceType.UPLOAD_FILE) + knowledge_config = _make_upload_knowledge_config( + file_ids=["file-1"], + process_rule=ProcessRule( + mode="custom", + rules=Rule( + pre_processing_rules=[PreProcessingRule(id="remove_stopwords", enabled=True)], + segmentation=Segmentation(separator="\n", max_tokens=128), + ), + ), + ) + upload_file = _upload_file(file_id="file-1", name="file.txt") + sqlite_session.add_all([dataset, upload_file]) + sqlite_session.commit() + + with ( + _dataset_service_document_patch( + "services.knowledge.dataset_service.FeatureService.get_features", return_value=_make_features() + ), + _dataset_service_document_patch("services.knowledge.dataset_service.redis_client") as mock_redis, + _dataset_service_document_patch( + "services.knowledge.dataset_service.DocumentIndexingTaskProxy" + ) as document_proxy_cls, + _dataset_service_document_patch( + "services.knowledge.dataset_service.time.strftime", return_value="20260101010101" + ), + _dataset_service_document_patch( + "services.knowledge.dataset_service.secrets.randbelow", return_value=23 + ), + ): + mock_redis.lock.return_value = _make_lock_context() + documents, batch = _dataset_service_document_DocumentService.save_document_with_dataset_id( + dataset, + knowledge_config, + account_context, + session=sqlite_session, + ) + + assert len(documents) == 1 + created_document = documents[0] + assert created_document.name == "file.txt" + assert batch == "20260101010101100023" + created_rule = sqlite_session.get( + _dataset_service_document_DatasetProcessRule, created_document.dataset_process_rule_id + ) + assert created_rule is not None + assert created_rule.mode == "custom" + assert created_rule.rules == knowledge_config.process_rule.rules.model_dump_json() + document_proxy_cls.assert_called_once_with("tenant-1", "dataset-1", [created_document.id]) + document_proxy_cls.return_value.delay.assert_called_once() + + @config_overrides_context(DEPLOYMENT_EDITION=DeploymentEdition.COMMUNITY) + def test_save_document_with_dataset_id_creates_automatic_process_rule_for_new_upload_document( + self, account_context, sqlite_session: Session + ): + dataset = _dataset_row(data_source_type=DataSourceType.UPLOAD_FILE) + knowledge_config = _make_upload_knowledge_config( + file_ids=["file-1"], + process_rule=ProcessRule(mode="automatic"), + ) + sqlite_session.add_all([dataset, _upload_file(file_id="file-1", name="file.txt")]) + sqlite_session.commit() + + with ( + _dataset_service_document_patch( + "services.knowledge.dataset_service.FeatureService.get_features", return_value=_make_features() + ), + _dataset_service_document_patch("services.knowledge.dataset_service.redis_client") as mock_redis, + _dataset_service_document_patch("services.knowledge.dataset_service.DocumentIndexingTaskProxy"), + _dataset_service_document_patch( + "services.knowledge.dataset_service.time.strftime", return_value="20260101010101" + ), + _dataset_service_document_patch( + "services.knowledge.dataset_service.secrets.randbelow", return_value=23 + ), + ): + mock_redis.lock.return_value = _make_lock_context() + documents, _ = _dataset_service_document_DocumentService.save_document_with_dataset_id( + dataset, + knowledge_config, + account_context, + session=sqlite_session, + ) + + created_rule = sqlite_session.get( + _dataset_service_document_DatasetProcessRule, documents[0].dataset_process_rule_id + ) + assert created_rule is not None + assert created_rule.mode == "automatic" + assert created_rule.rules == json.dumps(_dataset_service_document_DatasetProcessRule.AUTOMATIC_RULES) + assert sqlite_session.get(_dataset_service_document_Document, documents[0].id) is documents[0] + + @config_overrides_context(DEPLOYMENT_EDITION=DeploymentEdition.COMMUNITY) + def test_save_document_with_dataset_id_creates_fallback_automatic_process_rule_when_latest_is_missing( + self, account_context, sqlite_session: Session + ): + dataset = _dataset_row(data_source_type=DataSourceType.UPLOAD_FILE) + knowledge_config = _make_upload_knowledge_config(file_ids=["file-1"], process_rule=None) + sqlite_session.add_all([dataset, _upload_file(file_id="file-1", name="file.txt")]) + sqlite_session.commit() + + with ( + _dataset_service_document_patch( + "services.knowledge.dataset_service.FeatureService.get_features", return_value=_make_features() + ), + _dataset_service_document_patch("services.knowledge.dataset_service.redis_client") as mock_redis, + _dataset_service_document_patch("services.knowledge.dataset_service.DocumentIndexingTaskProxy"), + _dataset_service_document_patch( + "services.knowledge.dataset_service.time.strftime", return_value="20260101010101" + ), + _dataset_service_document_patch( + "services.knowledge.dataset_service.secrets.randbelow", return_value=23 + ), + ): + mock_redis.lock.return_value = _make_lock_context() + documents, _ = _dataset_service_document_DocumentService.save_document_with_dataset_id( + dataset, + knowledge_config, + account_context, + session=sqlite_session, + ) + + created_rule = sqlite_session.get( + _dataset_service_document_DatasetProcessRule, documents[0].dataset_process_rule_id + ) + assert created_rule is not None + assert created_rule.mode == "automatic" + assert created_rule.rules == json.dumps(_dataset_service_document_DatasetProcessRule.AUTOMATIC_RULES) + + @config_overrides_context(DEPLOYMENT_EDITION=DeploymentEdition.COMMUNITY) + def test_save_document_with_dataset_id_raises_when_upload_file_lookup_is_incomplete( + self, account_context, sqlite_session: Session + ): + dataset = _dataset_row(data_source_type=DataSourceType.UPLOAD_FILE) + knowledge_config = _make_upload_knowledge_config(file_ids=["file-1", "file-2"]) + sqlite_session.add_all([dataset, _upload_file(file_id="file-1", name="file.txt")]) + sqlite_session.commit() + + with ( + _dataset_service_document_patch( + "services.knowledge.dataset_service.FeatureService.get_features", return_value=_make_features() + ), + _dataset_service_document_patch("services.knowledge.dataset_service.redis_client") as mock_redis, + _dataset_service_document_patch( + "services.knowledge.dataset_service.time.strftime", return_value="20260101010101" + ), + _dataset_service_document_patch( + "services.knowledge.dataset_service.secrets.randbelow", return_value=23 + ), + ): + mock_redis.lock.return_value = _make_lock_context() + with pytest.raises(FileNotExistsError, match="One or more files not found"): + _dataset_service_document_DocumentService.save_document_with_dataset_id( + dataset, + knowledge_config, + account_context, + session=sqlite_session, + ) + + @config_overrides_context(DEPLOYMENT_EDITION=DeploymentEdition.COMMUNITY) + def test_save_document_with_dataset_id_requires_notion_info_list_for_notion_import( + self, account_context, sqlite_session: Session + ): + dataset = _dataset_row(data_source_type=DataSourceType.NOTION_IMPORT) + process_rule = _process_rule() + sqlite_session.add_all([dataset, process_rule]) + sqlite_session.commit() + knowledge_config = _dataset_service_document_KnowledgeConfig( + indexing_technique="economy", + data_source=_dataset_service_document_DataSource( + info_list=_dataset_service_document_InfoList(data_source_type="notion_import") + ), + doc_form=_dataset_service_document_IndexStructureType.PARAGRAPH_INDEX, + doc_language="English", + ) + + with ( + _dataset_service_document_patch( + "services.knowledge.dataset_service.FeatureService.get_features", return_value=_make_features() + ), + _dataset_service_document_patch("services.knowledge.dataset_service.redis_client") as mock_redis, + ): + mock_redis.lock.return_value = _make_lock_context() + with pytest.raises(ValueError, match="No notion info list found"): + _dataset_service_document_DocumentService.save_document_with_dataset_id( + dataset, + knowledge_config, + account_context, + dataset_process_rule=process_rule, + session=sqlite_session, + ) + + @config_overrides_context(DEPLOYMENT_EDITION=DeploymentEdition.COMMUNITY) + def test_save_document_with_dataset_id_requires_website_info_list_for_website_crawl( + self, account_context, sqlite_session: Session + ): + dataset = _dataset_row(data_source_type=DataSourceType.WEBSITE_CRAWL) + process_rule = _process_rule() + sqlite_session.add_all([dataset, process_rule]) + sqlite_session.commit() + knowledge_config = _dataset_service_document_KnowledgeConfig( + indexing_technique="economy", + data_source=_dataset_service_document_DataSource( + info_list=_dataset_service_document_InfoList(data_source_type="website_crawl") + ), + doc_form=_dataset_service_document_IndexStructureType.PARAGRAPH_INDEX, + doc_language="English", + ) + + with ( + _dataset_service_document_patch( + "services.knowledge.dataset_service.FeatureService.get_features", return_value=_make_features() + ), + _dataset_service_document_patch("services.knowledge.dataset_service.redis_client") as mock_redis, + ): + mock_redis.lock.return_value = _make_lock_context() + with pytest.raises(ValueError, match="No website info list found"): + _dataset_service_document_DocumentService.save_document_with_dataset_id( + dataset, + knowledge_config, + account_context, + dataset_process_rule=process_rule, + session=sqlite_session, + ) + + +TENANT_ID = "11111111-1111-1111-1111-111111111111" +USER_ID = "22222222-2222-2222-2222-222222222222" +DATASET_ID = "33333333-3333-3333-3333-333333333333" +DOCUMENT_ID = "44444444-4444-4444-4444-444444444444" + + +class FakeLock: + """Lock that always fails on enter with LockNotOwnedError.""" + + def __enter__(self) -> None: + raise LockNotOwnedError("simulated") + + def __exit__( + self, exc_type: type[BaseException] | None, exc: BaseException | None, tb: types.TracebackType | None + ) -> bool: + # Normal contextmanager signature; return False so exceptions propagate + return False + + +class TestDatasetServiceLockNotOwnedCases: + @staticmethod + @pytest.fixture + def fake_current_user( + monkeypatch: pytest.MonkeyPatch, + ) -> _dataset_service_lock_not_owned_Account: + tenant = Tenant(name="Test Tenant") + tenant.id = TENANT_ID + user = _dataset_service_lock_not_owned_Account(name="Test User", email="test@example.com") + user.id = USER_ID + user._current_tenant = tenant + monkeypatch.setattr("services.knowledge.dataset_service.current_user", user) + return user + + @staticmethod + @pytest.fixture + def fake_features(monkeypatch: pytest.MonkeyPatch) -> types.SimpleNamespace: + features = types.SimpleNamespace( + billing=types.SimpleNamespace(subscription=types.SimpleNamespace(plan="ENTERPRISE")), + documents_upload_quota=types.SimpleNamespace(limit=10_000, size=0), + ) + monkeypatch.setattr( + "services.knowledge.dataset_service.FeatureService.get_features", + lambda _tenant_id, **_kwargs: features, + ) + return features + + @staticmethod + @pytest.fixture + def fake_lock(monkeypatch: pytest.MonkeyPatch) -> None: + """Patch redis_client.lock to always raise LockNotOwnedError on enter.""" + + def _fake_lock(*_args: object, **_kwargs: object) -> FakeLock: + return FakeLock() + + # DatasetService imports redis_client directly from extensions.ext_redis + monkeypatch.setattr("services.knowledge.dataset_service.redis_client.lock", _fake_lock) + + # --------------------------------------------------------------------------- + # 1. Knowledge Pipeline document creation (save_document_with_dataset_id) + # --------------------------------------------------------------------------- + + @staticmethod + @pytest.mark.parametrize( + "sqlite_session", + [ + ( + _dataset_service_lock_not_owned_Account, + Tenant, + Dataset, + _dataset_service_lock_not_owned_DatasetProcessRule, + _dataset_service_lock_not_owned_Document, + ) + ], + indirect=True, + ) + @pytest.mark.usefixtures("fake_features", "fake_lock") + def test_save_document_with_dataset_id_ignores_lock_not_owned( + monkeypatch: pytest.MonkeyPatch, + fake_current_user: _dataset_service_lock_not_owned_Account, + sqlite_session: Session, + ) -> None: + # Arrange + dataset = Dataset( + id=DATASET_ID, + tenant_id=TENANT_ID, + name="Test Dataset", + description="", + created_by=USER_ID, + data_source_type="upload_file", + indexing_technique=IndexTechniqueType.HIGH_QUALITY, + ) + process_rule = _dataset_service_lock_not_owned_DatasetProcessRule( + dataset_id=DATASET_ID, + mode=ProcessRuleMode.AUTOMATIC, + rules=None, + created_by=USER_ID, + ) + sqlite_session.add_all([fake_current_user._current_tenant, fake_current_user, dataset, process_rule]) + sqlite_session.commit() + + knowledge_config = _dataset_service_lock_not_owned_KnowledgeConfig( + doc_form=_dataset_service_lock_not_owned_IndexStructureType.QA_INDEX, + original_document_id=None, # go into "new document" branch + data_source=_dataset_service_lock_not_owned_DataSource( + info_list=_dataset_service_lock_not_owned_InfoList(data_source_type="upload_file") + ), + indexing_technique="high_quality", + duplicate=False, + doc_language="en", + ) + + account = fake_current_user + + # Avoid touching real doc_form logic + monkeypatch.setattr( + "services.knowledge.dataset_service.DatasetService.check_doc_form", lambda *_args, **_kwargs: None + ) + + # Act: this would hit the redis lock, whose __enter__ raises LockNotOwnedError. + # Our implementation should catch it and still return (documents, batch). + documents, batch = _dataset_service_lock_not_owned_DocumentService.save_document_with_dataset_id( + dataset=dataset, + knowledge_config=knowledge_config, + account=account, + dataset_process_rule=process_rule, + session=sqlite_session, + ) + + # Assert + # We mainly care that: + # - No exception is raised + # - The function returns a sensible tuple + assert documents == [] + assert batch + assert not sqlite_session.in_transaction() + assert sqlite_session.scalar(select(func.count(_dataset_service_lock_not_owned_Document.id))) == 0 + assert sqlite_session.get(_dataset_service_lock_not_owned_DatasetProcessRule, process_rule.id) is process_rule + + # --------------------------------------------------------------------------- + # 2. Single-segment creation (add_segment) + # --------------------------------------------------------------------------- + + # --------------------------------------------------------------------------- + # 3. Multi-segment creation (multi_create_segment) + # --------------------------------------------------------------------------- + + @staticmethod + @pytest.mark.parametrize( + "sqlite_session", + [ + ( + _dataset_service_lock_not_owned_Account, + Tenant, + Dataset, + _dataset_service_lock_not_owned_Document, + DocumentSegment, + ) + ], + indirect=True, + ) + @pytest.mark.usefixtures("fake_lock") + def test_multi_create_segment_ignores_lock_not_owned( + fake_current_user: _dataset_service_lock_not_owned_Account, + sqlite_session: Session, + ) -> None: + # Arrange + dataset = Dataset( + id=DATASET_ID, + tenant_id=TENANT_ID, + name="Test Dataset", + description="", + created_by=USER_ID, + indexing_technique=IndexTechniqueType.ECONOMY, + ) + document = _dataset_service_lock_not_owned_Document( + id=DOCUMENT_ID, + tenant_id=TENANT_ID, + dataset_id=DATASET_ID, + position=1, + data_source_type=DataSourceType("upload_file"), + data_source_info="{}", + batch="batch-1", + name="Test Document", + created_from=DocumentCreatedFrom("web"), + created_by=USER_ID, + word_count=0, + doc_form=_dataset_service_lock_not_owned_IndexStructureType.QA_INDEX, + ) + sqlite_session.add_all([fake_current_user._current_tenant, fake_current_user, dataset, document]) + sqlite_session.commit() + + mutations = create_autospec(SegmentMutationService, instance=True) + result = SegmentService.multi_create_segment( + segments=[{"content": "question", "answer": "answer", "keywords": ["key"]}], + document=document, + dataset=dataset, + session=sqlite_session, + mutations=mutations, + ) + + assert result is None + mutations.index_segments.assert_not_called() + assert not sqlite_session.in_transaction() + assert sqlite_session.scalar(select(func.count(DocumentSegment.id))) == 0 + sqlite_session.refresh(document) + assert document.word_count == 0 + + +"""SQLite-backed tests for segment and child-chunk dataset services.""" + + +def _dataset_service_segment_account( + *, account_id: str = "user-1", tenant_id: str = "tenant-1" +) -> _dataset_service_segment_Account: + account = _dataset_service_segment_Account(name="User", email=f"{account_id}@example.com") + account.id = account_id + tenant = Tenant(name="Tenant") + tenant.id = tenant_id + account._current_tenant = tenant + return account + + +def _dataset_service_segment_dataset(*, dataset_id: str = "dataset-1", tenant_id: str = "tenant-1") -> Dataset: + return Dataset( + id=dataset_id, + tenant_id=tenant_id, + name="Dataset", + description="", + provider="vendor", + created_by="user-1", + maintainer="user-1", + indexing_technique=IndexTechniqueType.HIGH_QUALITY, + embedding_model_provider="provider", + embedding_model="embedding-model", + chunk_structure=_dataset_service_segment_IndexStructureType.PARAGRAPH_INDEX, + ) + + +def _document( + *, + document_id: str = "document-1", + dataset_id: str = "dataset-1", + tenant_id: str = "tenant-1", + doc_form: str = _dataset_service_segment_IndexStructureType.PARAGRAPH_INDEX, + word_count: int = 20, +) -> _dataset_service_segment_Document: + document = _dataset_service_segment_Document( + id=document_id, + tenant_id=tenant_id, + dataset_id=dataset_id, + position=1, + data_source_type=DataSourceType.UPLOAD_FILE, + batch="batch-1", + name="Document", + created_from=DocumentCreatedFrom.API, + created_by="user-1", + indexing_status="completed", + doc_form=doc_form, + word_count=word_count, + ) + + document.created_at = datetime(2026, 1, 1) + document.updated_at = datetime(2026, 1, 2) + return document + + +def _segment( + *, + segment_id: str = "segment-1", + dataset_id: str = "dataset-1", + document_id: str = "document-1", + tenant_id: str = "tenant-1", + position: int = 1, + content: str = "segment content", + enabled: bool = True, +) -> DocumentSegment: + segment = DocumentSegment( + tenant_id=tenant_id, + dataset_id=dataset_id, + document_id=document_id, + position=position, + content=content, + word_count=len(content), + tokens=0, + created_by="user-1", + enabled=enabled, + keywords=[], + answer=None, + index_node_id=f"node-{segment_id}", + status=SegmentStatus.COMPLETED, + ) + segment.id = segment_id + return segment + + +def _child( + *, + child_id: str = "child-1", + segment_id: str = "segment-1", + dataset_id: str = "dataset-1", + document_id: str = "document-1", + tenant_id: str = "tenant-1", + position: int = 1, + content: str = "child content", +) -> ChildChunk: + child = ChildChunk( + tenant_id=tenant_id, + dataset_id=dataset_id, + document_id=document_id, + segment_id=segment_id, + position=position, + index_node_id=f"node-{child_id}", + index_node_hash=f"hash-{child_id}", + content=content, + word_count=len(content), + created_by="user-1", + ) + child.id = child_id + return child + + +def _persist_chain(session: Session) -> tuple[Dataset, _dataset_service_segment_Document, DocumentSegment]: + dataset = _dataset_service_segment_dataset() + document = _document() + segment = _segment() + session.add_all([dataset, document, segment]) + session.commit() + return dataset, document, segment + + +class TestDatasetServiceSegmentCases: + class TestSegmentServiceQueries: + def test_get_child_chunks_filters_owner_keyword_and_paginates(self, sqlite_session: Session) -> None: + _persist_chain(sqlite_session) + sqlite_session.add_all( + [ + _child(child_id="child-1", position=1, content="alpha match"), + _child(child_id="child-2", position=2, content="beta"), + _child(child_id="child-3", position=3, content="alpha second"), + _child(child_id="foreign", tenant_id="tenant-2", position=4, content="alpha foreign"), + ] + ) + sqlite_session.commit() + + with _dataset_service_segment_patch( + "services.knowledge.dataset_service.current_user", _dataset_service_segment_account() + ): + page = SegmentService.get_child_chunks( + "segment-1", + "document-1", + "dataset-1", + page=1, + limit=1, + keyword="alpha", + session=sqlite_session, + ) + + assert page.total == 2 + assert [child.id for child in page.items] == ["child-1"] + + def test_get_child_chunk_by_id_scopes_tenant(self, sqlite_session: Session) -> None: + owned = _child() + sqlite_session.add(owned) + sqlite_session.commit() + + assert SegmentService.get_child_chunk_by_id(owned.id, "tenant-1", sqlite_session) is owned + assert SegmentService.get_child_chunk_by_id(owned.id, "tenant-2", sqlite_session) is None + + def test_get_child_chunk_by_segment_ref_enforces_full_chain(self, sqlite_session: Session) -> None: + child = _child() + sqlite_session.add(child) + sqlite_session.commit() + valid_ref = SegmentRef(DocumentRef(DatasetRef("tenant-1", "dataset-1"), "document-1"), "segment-1") + + assert SegmentService.get_child_chunk_by_segment_ref(child.id, valid_ref, sqlite_session) is child + spoofed_refs = [ + SegmentRef(DocumentRef(DatasetRef("tenant-2", "dataset-1"), "document-1"), "segment-1"), + SegmentRef(DocumentRef(DatasetRef("tenant-1", "dataset-2"), "document-1"), "segment-1"), + SegmentRef(DocumentRef(DatasetRef("tenant-1", "dataset-1"), "document-2"), "segment-1"), + SegmentRef(DocumentRef(DatasetRef("tenant-1", "dataset-1"), "document-1"), "segment-2"), + ] + for spoofed_ref in spoofed_refs: + assert SegmentService.get_child_chunk_by_segment_ref(child.id, spoofed_ref, sqlite_session) is None + + def test_get_segments_filters_status_keyword_and_orders(self, sqlite_session: Session) -> None: + sqlite_session.add_all( + [ + _segment(segment_id="one", position=2, content="alpha later"), + _segment(segment_id="two", position=1, content="alpha first"), + _segment(segment_id="three", position=3, content="beta"), + _segment(segment_id="foreign", tenant_id="tenant-2", position=1, content="alpha foreign"), + ] + ) + sqlite_session.commit() + + segments, total = SegmentService.get_segments( + "document-1", + "tenant-1", + status_list=[SegmentStatus.COMPLETED], + keyword="alpha", + session=sqlite_session, + ) + + assert total == 2 + assert [segment.id for segment in segments] == ["two", "one"] + + def test_get_segment_by_id_and_ref_scope_complete_owner(self, sqlite_session: Session) -> None: + segment = _segment() + sqlite_session.add(segment) + sqlite_session.commit() + valid_ref = SegmentRef(DocumentRef(DatasetRef("tenant-1", "dataset-1"), "document-1"), segment.id) + + assert SegmentService.get_segment_by_id(segment.id, "tenant-1", sqlite_session) is segment + assert SegmentService.get_segment_by_id(segment.id, "tenant-2", sqlite_session) is None + assert SegmentService.get_segment_by_ref(valid_ref, sqlite_session) is segment + spoofed_refs = [ + SegmentRef(DocumentRef(DatasetRef("tenant-2", "dataset-1"), "document-1"), segment.id), + SegmentRef(DocumentRef(DatasetRef("tenant-1", "dataset-2"), "document-1"), segment.id), + SegmentRef(DocumentRef(DatasetRef("tenant-1", "dataset-1"), "document-2"), segment.id), + ] + for spoofed_ref in spoofed_refs: + assert SegmentService.get_segment_by_ref(spoofed_ref, sqlite_session) is None + + def test_get_segments_by_document_and_dataset_returns_real_rows(self, sqlite_session: Session) -> None: + sqlite_session.add_all( + [ + _segment(segment_id="enabled"), + _segment(segment_id="disabled", position=2, enabled=False), + _segment(segment_id="other", document_id="document-2", position=1), + _segment(segment_id="foreign", tenant_id="tenant-2"), + ] + ) + sqlite_session.commit() + + segments = SegmentService.get_segments_by_document_and_dataset( + "document-1", + "dataset-1", + sqlite_session, + tenant_id="tenant-1", + status=SegmentStatus.COMPLETED, + enabled=True, + ) + + assert [segment.id for segment in segments] == ["enabled"] + + class TestSegmentServiceValidation: + def test_qa_segment_requires_answer(self) -> None: + with pytest.raises(ValueError, match="Answer is required"): + SegmentService.segment_create_args_validate( + {"content": "question"}, _document(doc_form=_dataset_service_segment_IndexStructureType.QA_INDEX) + ) + + @pytest.mark.parametrize("content", [None, "", " "]) + def test_segment_requires_non_empty_content(self, content: str | None) -> None: + with pytest.raises(ValueError, match="Content is empty"): + SegmentService.segment_create_args_validate({"content": content}, _document()) + + def test_segment_attachment_ids_must_be_a_list(self) -> None: + with pytest.raises(ValueError, match="Attachment IDs is invalid"): + SegmentService.segment_create_args_validate({"content": "text", "attachment_ids": "file"}, _document()) + + def test_segment_attachment_limit_is_enforced(self, config_overrides: Callable[..., None]) -> None: + config_overrides(SINGLE_CHUNK_ATTACHMENT_LIMIT=1) + + with pytest.raises(ValueError, match="Exceeded maximum attachment limit"): + SegmentService.segment_create_args_validate( + {"content": "text", "attachment_ids": ["one", "two"]}, + _document(), + ) + + def test_segment_attachment_limit_accepts_exact_boundary(self, config_overrides: Callable[..., None]) -> None: + config_overrides(SINGLE_CHUNK_ATTACHMENT_LIMIT=2) + + SegmentService.segment_create_args_validate( + {"content": "text", "attachment_ids": ["one", "two"]}, + _document(), + ) + + class TestSegmentServiceMutations: + @pytest.mark.parametrize( + ("document_dataset_id", "document_tenant_id"), + [("dataset-2", "tenant-1"), ("dataset-1", "tenant-2")], + ) + def test_delete_segment_rejects_document_outside_dataset( + self, document_dataset_id: str, document_tenant_id: str, sqlite_session: Session + ) -> None: + dataset = _dataset_service_segment_dataset() + document = _document(dataset_id=document_dataset_id, tenant_id=document_tenant_id) + segment = _segment(dataset_id=document_dataset_id, tenant_id=document_tenant_id) + mutations = create_autospec(SegmentMutationService, instance=True, spec_set=True) + + with pytest.raises(ValueError, match="Document does not belong to the dataset"): + SegmentService.delete_segment(segment, document, dataset, sqlite_session, mutations=mutations) + + mutations.delete_segment.assert_not_called() + + def test_multi_create_segment_marks_each_real_row_error_on_vector_failure( + self, sqlite_session: Session, sqlite_session_factory: sessionmaker[Session] + ) -> None: + dataset = _dataset_service_segment_dataset() + document = _document(word_count=0) + sqlite_session.add_all([dataset, document]) + sqlite_session.commit() + embedding_model = SimpleNamespace(get_text_embedding_num_tokens=lambda *, texts: [len(texts) + 1]) + index = create_autospec(SegmentIndex, instance=True, spec_set=True) + + def fail(*_args: object, **_kwargs: object) -> None: + assert not sqlite_session.in_transaction() + raise RuntimeError("vector failed") + + index.create_many.side_effect = fail + mutations = SegmentMutationService( + store=SQLAlchemySegmentRepository(session_factory=sqlite_session_factory), + index=index, + text_hash=lambda text: text, + indexing_state=create_autospec(SegmentIndexingState, instance=True, spec_set=True), + ) + + with ( + _dataset_service_segment_patch( + "services.knowledge.dataset_service.current_user", _dataset_service_segment_account() + ), + _dataset_service_segment_patch("services.knowledge.dataset_service.ModelManager") as manager_cls, + ): + manager_cls.for_tenant.return_value.get_model_instance.return_value = embedding_model + segments = SegmentService.multi_create_segment( + [{"content": "one"}, {"content": "two"}], + document, + dataset, + sqlite_session, + mutations=mutations, + ) + + assert segments is not None + assert len(segments) == 2 + assert all(segment.status == SegmentStatus.ERROR and not segment.enabled for segment in segments) + assert document.word_count == 6 + + def test_multi_create_segment_persists_qa_counts_positions_and_tokens(self, sqlite_session: Session) -> None: + dataset = _dataset_service_segment_dataset() + document = _document(doc_form=_dataset_service_segment_IndexStructureType.QA_INDEX, word_count=5) + sqlite_session.add_all([dataset, document, _segment(segment_id="existing")]) + sqlite_session.commit() + embedding_model = _dataset_service_segment_MagicMock() + embedding_model.get_text_embedding_num_tokens.side_effect = [[11], [13]] + mutations = create_autospec(SegmentMutationService, instance=True, spec_set=True) + + with ( + _dataset_service_segment_patch( + "services.knowledge.dataset_service.current_user", _dataset_service_segment_account() + ), + _dataset_service_segment_patch("services.knowledge.dataset_service.ModelManager") as manager_cls, + ): + manager_cls.for_tenant.return_value.get_model_instance.return_value = embedding_model + segments = SegmentService.multi_create_segment( + [ + {"content": "question-1", "answer": "answer-1", "keywords": ["key"]}, + {"content": "question-2", "answer": "answer-2"}, + ], + document, + dataset, + sqlite_session, + mutations=mutations, + ) + + assert segments is not None + assert [segment.position for segment in segments] == [2, 3] + assert [segment.tokens for segment in segments] == [11, 13] + assert [segment.answer for segment in segments] == ["answer-1", "answer-2"] + expected_increment = sum(len(segment.content) + len(segment.answer or "") for segment in segments) + assert document.word_count == 5 + expected_increment + assert embedding_model.get_text_embedding_num_tokens.call_args_list[0].kwargs == { + "texts": ["question-1answer-1"] + } + assert embedding_model.get_text_embedding_num_tokens.call_args_list[1].kwargs == { + "texts": ["question-2answer-2"] + } + mutations.index_segments.assert_called_once_with( + DatasetRef("tenant-1", "dataset-1").document("document-1"), + segment_ids=[segment.id for segment in segments], + keywords_list=[["key"], None], + ) diff --git a/api/tests/unit_tests/services/knowledge/test_document_adapters.py b/api/tests/unit_tests/services/knowledge/test_document_adapters.py deleted file mode 100644 index a56f5374f96..00000000000 --- a/api/tests/unit_tests/services/knowledge/test_document_adapters.py +++ /dev/null @@ -1,293 +0,0 @@ -import json -from datetime import datetime -from unittest.mock import patch - -import pytest -from sqlalchemy import Engine, event -from sqlalchemy.orm import Session, sessionmaker - -from controllers.console.datasets.datasets_document import DocumentDetailResponse, DocumentWithSegmentsListResponse -from core.rag.index_processor.constant.index_type import IndexStructureType, IndexTechniqueType -from models.dataset import Dataset, DatasetProcessRule, Document, DocumentPipelineExecutionLog, DocumentSegment -from models.enums import DataSourceType, DocumentCreatedFrom, IndexingStatus, ProcessRuleMode, SegmentStatus -from services.knowledge.documents.adapters import SQLAlchemyDocumentOperations -from services.knowledge.documents.application import DocumentListFilter, DocumentNotFoundError -from services.knowledge.resource_scope import DatasetRef - -REF = DatasetRef("tenant-1", "dataset-1") - - -def _dataset(**values: object) -> Dataset: - row = Dataset( - id="dataset-1", - tenant_id="tenant-1", - name="Dataset", - created_by="actor-1", - maintainer="actor-1", - indexing_technique="economy", - chunk_structure="text_model", - ) - for key, value in values.items(): - setattr(row, key, value) - return row - - -def _document(**values: object) -> Document: - row = Document( - tenant_id="tenant-1", - dataset_id="dataset-1", - position=1, - data_source_type=DataSourceType("local_file"), - batch="batch-1", - name="Document", - created_from=DocumentCreatedFrom("web"), - created_by="actor-1", - ) - row.id = "document-1" - row.data_source_info = json.dumps({"file_path": "/path/to/file"}) - row.indexing_status = IndexingStatus.COMPLETED - row.doc_form = IndexStructureType.PARAGRAPH_INDEX - row.enabled = True - row.archived = False - row.word_count = 10 - row.created_at = datetime(2024, 1, 1) - for key, value in values.items(): - setattr(row, key, value) - return row - - -def _segment( - segment_id: str, - *, - tenant_id: str = "tenant-1", - dataset_id: str = "dataset-1", - document_id: str = "document-1", - status: SegmentStatus = SegmentStatus.COMPLETED, - completed_at: datetime | None = datetime(2024, 1, 1), - hit_count: int = 0, -) -> DocumentSegment: - segment = DocumentSegment( - tenant_id=tenant_id, - dataset_id=dataset_id, - document_id=document_id, - position=1, - content="content", - word_count=7, - tokens=2, - created_by="actor-1", - status=status, - completed_at=completed_at, - hit_count=hit_count, - ) - segment.id = segment_id - return segment - - -@pytest.fixture -def operations(sqlite_session_factory: sessionmaker[Session]) -> SQLAlchemyDocumentOperations: - with sqlite_session_factory.begin() as session: - session.add_all([_dataset(), _document()]) - return SQLAlchemyDocumentOperations(session_factory=sqlite_session_factory) - - -@pytest.mark.parametrize("foreign", [{"tenant_id": "other"}, {"dataset_id": "other"}]) -def test_document_reads_and_writes_reject_incorrect_owner( - operations: SQLAlchemyDocumentOperations, sqlite_session_factory: sessionmaker[Session], foreign: dict[str, str] -) -> None: - with sqlite_session_factory.begin() as session: - session.add(_document(id="foreign", **foreign)) - ref = REF.document("foreign") - assert operations.get_state(ref) is None - with pytest.raises(DocumentNotFoundError): - operations.get_detail(ref, metadata_only=False) - with pytest.raises(DocumentNotFoundError): - operations.update_document(ref, {"name": "changed"}) - with pytest.raises(DocumentNotFoundError): - operations.rename_document(ref, "changed") - with sqlite_session_factory() as session: - document = session.get(Document, "foreign") - assert document is not None - assert document.name == "Document" - - -def test_document_lookup_uses_active_workspace(operations: SQLAlchemyDocumentOperations) -> None: - assert operations.find_document(workspace_id="other", document_id="document-1") is None - assert operations.find_document(workspace_id="tenant-1", document_id="document-1") == REF.document("document-1") - - -def test_listing_filters_scope_and_materializes_values( - operations: SQLAlchemyDocumentOperations, sqlite_session_factory: sessionmaker[Session] -) -> None: - with sqlite_session_factory.begin() as session: - session.add_all( - [ - _document(id="matching", name="Needle", position=2), - _document(id="disabled", name="Needle", enabled=False), - _document(id="foreign", name="Needle", tenant_id="other"), - _document(id="other-dataset", name="Needle", dataset_id="other"), - ] - ) - result = operations.list_documents(REF, DocumentListFilter(search="Needle", status="available", fetch=True)) - response = DocumentWithSegmentsListResponse.model_validate(result).model_dump(mode="json") - assert response["total"] == 1 - assert response["data"][0]["id"] == "matching" - assert response["data"][0]["data_source_info"] == {"file_path": "/path/to/file"} - assert response["data"][0]["completed_segments"] == 0 - assert response["data"][0]["total_segments"] == 0 - - -def test_listing_hit_count_sort_and_exact_last_page( - operations: SQLAlchemyDocumentOperations, sqlite_session_factory: sessionmaker[Session] -) -> None: - with sqlite_session_factory.begin() as session: - session.add_all( - [ - _document(id="popular", position=2), - _segment("segment-1", hit_count=1), - _segment("segment-2", document_id="popular", hit_count=10), - _segment("foreign", tenant_id="other", hit_count=1000), - ] - ) - first = operations.list_documents(REF, DocumentListFilter(sort="-hit_count", limit=1)) - last = operations.list_documents(REF, DocumentListFilter(sort="-hit_count", limit=1, page=2)) - assert [row["id"] for row in first["data"]] == ["popular"] - assert first["has_more"] is True - assert last["has_more"] is False - assert last["total"] == 2 - - -def test_status_counts_use_one_aggregate_and_complete_owner_chain( - operations: SQLAlchemyDocumentOperations, sqlite_session_factory: sessionmaker[Session], sqlite_engine: Engine -) -> None: - with sqlite_session_factory.begin() as session: - session.add_all( - [ - _segment("complete"), - _segment("waiting", completed_at=None), - _segment("excluded", status=SegmentStatus.RE_SEGMENT), - _segment("foreign", tenant_id="other"), - _segment("wrong-dataset", dataset_id="other"), - _document(id="second", is_paused=True, indexing_status=IndexingStatus.INDEXING), - ] - ) - statements: list[str] = [] - - def record( - _conn: object, _cursor: object, statement: str, _parameters: object, _context: object, _many: bool - ) -> None: - statements.append(statement) - - event.listen(sqlite_engine, "before_cursor_execute", record) - try: - result = operations.get_batch_indexing_status(REF, "batch-1") - finally: - event.remove(sqlite_engine, "before_cursor_execute", record) - rows = {row["id"]: row for row in result["data"]} - assert rows["document-1"]["completed_segments"] == 1 - assert rows["document-1"]["total_segments"] == 2 - assert rows["second"]["indexing_status"] == "paused" - assert sum("FROM document_segments" in statement for statement in statements) == 1 - - -def test_unknown_batch_is_not_found(operations: SQLAlchemyDocumentOperations) -> None: - with pytest.raises(DocumentNotFoundError): - operations.get_batch_indexing_status(REF, "missing") - - -@pytest.mark.parametrize("rules", [None, {"segmentation": {"delimiter": "\\n", "max_tokens": 500}}]) -def test_process_rule_preserves_null_and_legacy_delimiter( - operations: SQLAlchemyDocumentOperations, - sqlite_session_factory: sessionmaker[Session], - rules: dict[str, object] | None, -) -> None: - with sqlite_session_factory.begin() as session: - session.add( - DatasetProcessRule( - dataset_id="dataset-1", - mode=ProcessRuleMode.CUSTOM, - rules=json.dumps(rules) if rules else None, - created_by="actor-1", - ) - ) - result = operations.get_process_rule(REF) - assert result["mode"] == "custom" - assert result["rules"] == rules - - -def test_detail_is_serializable_after_session_closes(operations: SQLAlchemyDocumentOperations) -> None: - result = operations.get_detail(REF.document("document-1"), metadata_only=False) - response = DocumentDetailResponse.model_validate(result).model_dump(mode="json") - assert response["id"] == "document-1" - assert response["data_source_info"] == {"file_path": "/path/to/file"} - assert response["data_source_detail_dict"] == {} - assert response["created_at"] == int(datetime(2024, 1, 1).timestamp()) - - -def test_metadata_only_avoids_loading_full_details(operations: SQLAlchemyDocumentOperations) -> None: - with patch("services.knowledge.documents.adapters.get_document_source_detail") as source: - result = operations.get_detail(REF.document("document-1"), metadata_only=True) - assert set(result) == {"id", "doc_type", "doc_metadata"} - source.assert_not_called() - - -def test_update_is_committed_before_returning( - operations: SQLAlchemyDocumentOperations, sqlite_session_factory: sessionmaker[Session] -) -> None: - operations.update_document(REF.document("document-1"), {"is_paused": True, "paused_by": "actor-1"}) - with sqlite_session_factory() as session: - document = session.get(Document, "document-1") - assert document is not None - assert document.is_paused is True - assert document.paused_by == "actor-1" - - -def test_latest_pipeline_log_and_empty_result( - operations: SQLAlchemyDocumentOperations, sqlite_session_factory: sessionmaker[Session] -) -> None: - assert operations.get_execution_log(REF.document("document-1")) == {} - with sqlite_session_factory.begin() as session: - for year in [2024, 2025]: - row = DocumentPipelineExecutionLog( - pipeline_id="pipeline-1", - document_id="document-1", - datasource_type="local_file", - datasource_info=json.dumps({"year": year}), - datasource_node_id="node-1", - input_data={"year": year}, - created_by="actor-1", - ) - row.created_at = datetime(year, 1, 1) - session.add(row) - result = operations.get_execution_log(REF.document("document-1")) - assert result["datasource_info"] == {"year": 2025} - assert result["input_data"] == {"year": 2025} - - -@pytest.mark.parametrize("settings", [None, {"enable": False}]) -def test_summary_guard_rejects_disabled_summary( - operations: SQLAlchemyDocumentOperations, - sqlite_session_factory: sessionmaker[Session], - settings: dict[str, object] | None, -) -> None: - with sqlite_session_factory.begin() as session: - dataset = session.get(Dataset, "dataset-1") - assert dataset is not None - dataset.indexing_technique = IndexTechniqueType.HIGH_QUALITY - dataset.summary_index_setting = settings - with pytest.raises(ValueError, match="not enabled"): - operations.require_summary_enabled(REF) - - -def test_enable_summary_only_updates_requested_owner( - operations: SQLAlchemyDocumentOperations, sqlite_session_factory: sessionmaker[Session] -) -> None: - with sqlite_session_factory.begin() as session: - session.add(_document(id="foreign", tenant_id="other", need_summary=False)) - operations.enable_summary(REF, ["document-1", "foreign"]) - with sqlite_session_factory() as session: - document = session.get(Document, "document-1") - foreign = session.get(Document, "foreign") - assert document is not None - assert foreign is not None - assert document.need_summary is True - assert foreign.need_summary is False diff --git a/api/tests/unit_tests/services/knowledge/test_document_creation.py b/api/tests/unit_tests/services/knowledge/test_document_creation.py deleted file mode 100644 index b6897944d6f..00000000000 --- a/api/tests/unit_tests/services/knowledge/test_document_creation.py +++ /dev/null @@ -1,356 +0,0 @@ -"""Exercise document creation with real SQL and assertions at external boundaries.""" - -from collections.abc import Callable, Generator, Iterator -from contextlib import contextmanager -from datetime import datetime -from unittest.mock import Mock - -import pytest -from redis.exceptions import LockNotOwnedError -from sqlalchemy import Engine, event, select -from sqlalchemy.orm import Session, sessionmaker - -from core.model_manager import ModelInstance -from core.rag.index_processor.constant.index_type import IndexTechniqueType -from enums import CloudPlan -from machinery.context import RequestContext -from models.account import Account, Tenant, TenantAccountJoin, TenantAccountRole -from models.dataset import Dataset, Document, DocumentSegment -from models.enums import IndexingStatus, SegmentStatus -from services.document_indexing_proxy.batch_indexing_base import BatchDocumentIndexingProxy -from services.document_indexing_proxy.document_indexing_task_proxy import DocumentIndexingTaskProxy -from services.document_indexing_proxy.duplicate_document_indexing_task_proxy import DuplicateDocumentIndexingTaskProxy -from services.entities.feature_entities import FeatureModel -from services.errors.file import FileNotExistsError -from services.knowledge.documents.adapters import SQLAlchemyDocumentOperations -from services.knowledge.resource_scope import DatasetRef -from tests.unit_tests.model_factories import make_upload_file - -CONTEXT = RequestContext("request", None, "actor", "tenant") -REF = DatasetRef("tenant", "dataset") - - -def settings(**values: object) -> dict[str, object]: - return { - "indexing_technique": "economy", - "process_rule": {"mode": "automatic"}, - "data_source": {"info_list": {"data_source_type": "upload_file", "file_info_list": {"file_ids": ["file"]}}}, - **values, - } - - -@pytest.fixture -def operations( - sqlite_session_factory: sessionmaker[Session], config_overrides: Callable[..., None] -) -> SQLAlchemyDocumentOperations: - config_overrides(DEPLOYMENT_EDITION="CLOUD", RBAC_ENABLED=False) - with sqlite_session_factory.begin() as session: - actor = Account(name="Actor", email="actor@example.com") - actor.id = "actor" - tenant = Tenant(name="Tenant") - tenant.id = "tenant" - session.add_all( - [ - actor, - tenant, - TenantAccountJoin(tenant_id="tenant", account_id="actor", role=TenantAccountRole.OWNER), - Dataset( - id="dataset", tenant_id="tenant", name="Dataset", created_by="actor", indexing_technique="economy" - ), - make_upload_file(file_id="file", tenant_id="tenant", name="example.txt", created_by="actor"), - make_upload_file(file_id="foreign-file", tenant_id="foreign", name="foreign.txt"), - ] - ) - return SQLAlchemyDocumentOperations(session_factory=sqlite_session_factory) - - -@pytest.fixture -def outside_transaction(sqlite_engine: Engine) -> Iterator[Callable[[], None]]: - active: set[int] = set() - - def begin(connection: object) -> None: - active.add(id(connection)) - - def end(connection: object) -> None: - active.discard(id(connection)) - - def check() -> None: - assert not active, "External I/O must not hold an open database transaction" - - listeners = [("begin", begin), ("commit", end), ("rollback", end)] - for name, callback in listeners: - event.listen(sqlite_engine, name, callback) - yield check - for name, callback in listeners: - event.remove(sqlite_engine, name, callback) - - -@pytest.fixture -def external_calls( - monkeypatch: pytest.MonkeyPatch, - outside_transaction: Callable[[], None], - sqlite_session_factory: sessionmaker[Session], -) -> list[tuple[str, tuple[str, ...]]]: - calls: list[tuple[str, tuple[str, ...]]] = [] - - def features(_tenant_id: str, *, exclude_vector_space: bool) -> FeatureModel: - outside_transaction() - assert exclude_vector_space is True - result = FeatureModel() - result.billing.subscription.plan = CloudPlan.PROFESSIONAL - return result - - @contextmanager - def lock(_name: str, *, timeout: int) -> Generator[None]: - outside_transaction() - assert timeout == 600 - yield - outside_transaction() - - def dispatched(kind: str, dataset_id: str, document_ids: tuple[str, ...]) -> None: - outside_transaction() - with sqlite_session_factory() as session: - for document_id in document_ids: - document = session.get(Document, document_id) - assert document is not None - assert document.dataset_id == dataset_id - if kind != "clean": - assert document.indexing_status == IndexingStatus.WAITING - calls.append((kind, document_ids)) - - def create(proxy: BatchDocumentIndexingProxy) -> None: - dispatched("create", proxy._dataset_id, tuple(proxy._document_ids)) - - def duplicate(proxy: BatchDocumentIndexingProxy) -> None: - dispatched("duplicate", proxy._dataset_id, tuple(proxy._document_ids)) - - monkeypatch.setattr("services.knowledge.dataset_service.FeatureService.get_features", features) - monkeypatch.setattr("services.knowledge.documents.adapters.redis_client.lock", lock) - monkeypatch.setattr(DocumentIndexingTaskProxy, "delay", create) - monkeypatch.setattr(DuplicateDocumentIndexingTaskProxy, "delay", duplicate) - monkeypatch.setattr( - "services.knowledge.dataset_service.document_indexing_update_task.delay", - lambda dataset_id, document_id: dispatched("update", dataset_id, (document_id,)), - ) - monkeypatch.setattr( - "services.knowledge.dataset_service.clean_notion_document_task.delay", - lambda document_ids, dataset_id: dispatched("clean", dataset_id, tuple(document_ids)), - ) - return calls - - -@pytest.mark.parametrize("initialize", [False, True]) -def test_upload_preserves_response_and_commits_before_dispatch( - operations: SQLAlchemyDocumentOperations, - external_calls: list[tuple[str, tuple[str, ...]]], - sqlite_session_factory: sessionmaker[Session], - initialize: bool, -) -> None: - result = ( - operations.initialize_dataset(CONTEXT, settings()) - if initialize - else operations.create_documents(CONTEXT, REF, settings()) - ) - (document,) = result["documents"] - assert document["name"] == "example.txt" - assert document["data_source_info_dict"] == {"upload_file_id": "file"} - assert document["indexing_status"] == "waiting" - assert result["batch"] - assert external_calls == [("create", (document["id"],))] - with sqlite_session_factory() as session: - dataset = session.get(Dataset, result["dataset"]["id"]) - assert dataset is not None - assert dataset.tenant_id == "tenant" - assert dataset.name == ("example.txt..." if initialize else "Dataset") - assert dataset.permission == "only_me" - - -def test_duplicate_upload_reuses_document_and_dispatches_duplicate_job( - operations: SQLAlchemyDocumentOperations, external_calls: list[tuple[str, tuple[str, ...]]] -) -> None: - first = operations.create_documents(CONTEXT, REF, settings()) - second = operations.create_documents(CONTEXT, REF, settings()) - first_id = first["documents"][0]["id"] - assert second["documents"][0]["id"] == first_id - assert external_calls == [("create", (first_id,)), ("duplicate", (first_id,))] - - -def test_reprocess_checks_model_outside_transaction_and_preserves_segment_reset( - operations: SQLAlchemyDocumentOperations, - external_calls: list[tuple[str, tuple[str, ...]]], - outside_transaction: Callable[[], None], - sqlite_session_factory: sessionmaker[Session], - monkeypatch: pytest.MonkeyPatch, -) -> None: - created = operations.create_documents(CONTEXT, REF, settings()) - document_id = created["documents"][0]["id"] - with sqlite_session_factory.begin() as session: - document = session.get(Document, document_id) - assert document is not None - document.indexing_status = IndexingStatus.COMPLETED - document.completed_at = datetime(2026, 1, 1) - dataset = session.get(Dataset, REF.dataset_id) - assert dataset is not None - dataset.indexing_technique = IndexTechniqueType.HIGH_QUALITY - dataset.embedding_model = "embedding" - dataset.embedding_model_provider = "provider" - session.add( - DocumentSegment( - tenant_id="tenant", - dataset_id="dataset", - document_id=document_id, - position=1, - content="content", - word_count=1, - tokens=1, - created_by="actor", - status=SegmentStatus.COMPLETED, - ) - ) - manager = Mock() - manager.get_model_instance.side_effect = lambda **_kwargs: outside_transaction() - monkeypatch.setattr("services.knowledge.dataset_service.ModelManager.for_tenant", lambda **_kwargs: manager) - result = operations.create_documents(CONTEXT, REF, settings(original_document_id=document_id, name="Renamed")) - assert result["documents"][0]["id"] == document_id - assert result["documents"][0]["name"] == "Renamed" - assert result["batch"] == created["batch"] - assert external_calls[-1] == ("update", (document_id,)) - manager.get_model_instance.assert_called_once() - with sqlite_session_factory() as session: - segment = session.scalar(select(DocumentSegment).where(DocumentSegment.document_id == document_id)) - assert segment is not None - assert segment.status == SegmentStatus.RE_SEGMENT - - -def test_first_high_quality_upload_resolves_default_model_without_transaction( - operations: SQLAlchemyDocumentOperations, - external_calls: list[tuple[str, tuple[str, ...]]], - outside_transaction: Callable[[], None], - sqlite_session_factory: sessionmaker[Session], - monkeypatch: pytest.MonkeyPatch, -) -> None: - with sqlite_session_factory.begin() as session: - dataset = session.get(Dataset, REF.dataset_id) - assert dataset is not None - dataset.indexing_technique = None - - def model(**_kwargs: object) -> ModelInstance: - outside_transaction() - instance = object.__new__(ModelInstance) - instance.model_name = "default-embedding" - instance.provider = "default-provider" - return instance - - manager = Mock() - manager.get_default_model_instance.side_effect = model - monkeypatch.setattr("services.knowledge.dataset_service.ModelManager.for_tenant", lambda **_kwargs: manager) - operations.create_documents(CONTEXT, REF, settings(indexing_technique="high_quality")) - with sqlite_session_factory() as session: - dataset = session.get(Dataset, REF.dataset_id) - assert dataset is not None - assert dataset.embedding_model == "default-embedding" - assert dataset.embedding_model_provider == "default-provider" - assert dataset.collection_binding_id is not None - assert len(external_calls) == 1 - - -def test_foreign_upload_is_rejected_without_dispatch( - operations: SQLAlchemyDocumentOperations, external_calls: list[tuple[str, tuple[str, ...]]] -) -> None: - config = settings( - data_source={"info_list": {"data_source_type": "upload_file", "file_info_list": {"file_ids": ["foreign-file"]}}} - ) - with pytest.raises(FileNotExistsError): - operations.create_documents(CONTEXT, REF, config) - assert external_calls == [] - - -def test_notion_removal_is_dispatched_after_new_documents_commit( - operations: SQLAlchemyDocumentOperations, external_calls: list[tuple[str, tuple[str, ...]]] -) -> None: - def notion(page_id: str) -> dict[str, object]: - return settings( - data_source={ - "info_list": { - "data_source_type": "notion_import", - "notion_info_list": [ - { - "credential_id": "credential", - "workspace_id": "workspace", - "pages": [{"page_id": page_id, "page_name": page_id, "type": "page"}], - } - ], - } - } - ) - - first = operations.create_documents(CONTEXT, REF, notion("old-page")) - old_id = first["documents"][0]["id"] - operations.create_documents(CONTEXT, REF, notion("new-page")) - assert external_calls[-2] == ("clean", (old_id,)) - assert external_calls[-1][0] == "create" - assert old_id not in external_calls[-1][1] - - -def test_cloud_quota_rejection_does_not_create_documents_or_dispatch( - operations: SQLAlchemyDocumentOperations, - external_calls: list[tuple[str, tuple[str, ...]]], - sqlite_session_factory: sessionmaker[Session], - config_overrides: Callable[..., None], -) -> None: - config_overrides(BATCH_UPLOAD_LIMIT=0) - with pytest.raises(ValueError, match="batch upload limit"): - operations.create_documents(CONTEXT, REF, settings()) - assert external_calls == [] - with sqlite_session_factory() as session: - assert session.scalar(select(Document.id)) is None - - -def test_response_read_failure_still_dispatches_already_committed_documents( - operations: SQLAlchemyDocumentOperations, - external_calls: list[tuple[str, tuple[str, ...]]], - monkeypatch: pytest.MonkeyPatch, -) -> None: - response = Mock(side_effect=ValueError("response read failed")) - monkeypatch.setattr("services.knowledge.documents.adapters._created_response", response) - with pytest.raises(ValueError, match="response read failed"): - operations.create_documents(CONTEXT, REF, settings()) - assert len(external_calls) == 1 - assert external_calls[0][0] == "create" - - -def test_lost_creation_lock_keeps_empty_batch_response( - operations: SQLAlchemyDocumentOperations, - external_calls: list[tuple[str, tuple[str, ...]]], - monkeypatch: pytest.MonkeyPatch, -) -> None: - @contextmanager - def lock(_name: str, **_kwargs: object) -> Generator[None]: - raise LockNotOwnedError("lock lost") - yield # pragma: no cover - - monkeypatch.setattr("services.knowledge.documents.adapters.redis_client.lock", lock) - result = operations.create_documents(CONTEXT, REF, settings()) - assert result["documents"] == [] - assert result["batch"] - assert result["dataset"]["id"] == REF.dataset_id - assert external_calls == [] - - -def test_billing_http_boundary_does_not_hold_a_transaction( - operations: SQLAlchemyDocumentOperations, - outside_transaction: Callable[[], None], - monkeypatch: pytest.MonkeyPatch, -) -> None: - class StopBeforeNetworkError(Exception): - pass - - def request(_method: str, url: str, **_kwargs: object) -> None: - assert url.endswith("/subscription/info") - outside_transaction() - raise StopBeforeNetworkError() - - monkeypatch.setattr("services.billing_service._http_client.request", request) - with pytest.raises(StopBeforeNetworkError): - operations.create_documents(CONTEXT, REF, settings()) - outside_transaction() diff --git a/api/tests/unit_tests/services/knowledge/test_document_sync_application_service.py b/api/tests/unit_tests/services/knowledge/test_document_sync.py similarity index 100% rename from api/tests/unit_tests/services/knowledge/test_document_sync_application_service.py rename to api/tests/unit_tests/services/knowledge/test_document_sync.py diff --git a/api/tests/unit_tests/services/knowledge/test_metadata_built_in_fields.py b/api/tests/unit_tests/services/knowledge/test_metadata_built_in_fields.py deleted file mode 100644 index 02cc45958f4..00000000000 --- a/api/tests/unit_tests/services/knowledge/test_metadata_built_in_fields.py +++ /dev/null @@ -1,163 +0,0 @@ -import time -from collections.abc import Iterator -from datetime import datetime - -import pytest -from sqlalchemy import select -from sqlalchemy.orm import Session, sessionmaker - -from core.rag.index_processor.constant.built_in_field import BuiltInField -from core.rag.retrieval.dataset_retrieval import DatasetRetrieval -from models.dataset import Document, build_built_in_field_values -from models.enums import DataSourceType -from services.knowledge import dataset_service -from services.knowledge.dataset_service import DocumentService -from tests.unit_tests.model_factories import make_account, make_dataset, make_document - -CREATED_AT = datetime(2026, 10, 16, 8, 0, 0) -CREATED_AT_EPOCH = 1792137600 - - -@pytest.fixture -def non_utc_host_timezone(monkeypatch: pytest.MonkeyPatch) -> Iterator[None]: - monkeypatch.setenv("TZ", "Asia/Tokyo") - time.tzset() - yield - monkeypatch.undo() - time.tzset() - - -@pytest.mark.usefixtures("non_utc_host_timezone") -def test_build_built_in_field_values_treats_naive_datetimes_as_utc() -> None: - values = build_built_in_field_values( - document_name="doc.txt", - uploader="Alice", - upload_date=CREATED_AT, - last_update_date=CREATED_AT, - data_source_type=DataSourceType.UPLOAD_FILE, - ) - - assert values == { - BuiltInField.document_name: "doc.txt", - BuiltInField.uploader: "Alice", - BuiltInField.upload_date: CREATED_AT_EPOCH, - BuiltInField.last_update_date: CREATED_AT_EPOCH, - BuiltInField.source: "file_upload", - } - - -@pytest.mark.parametrize( - ("data_source_type", "source"), - [ - (DataSourceType.UPLOAD_FILE, "file_upload"), - (DataSourceType.LOCAL_FILE, "file_upload"), - (DataSourceType.NOTION_IMPORT, "notion"), - (DataSourceType.WEBSITE_CRAWL, "website"), - (DataSourceType.ONLINE_DOCUMENT, "online_document"), - (DataSourceType.ONLINE_DRIVE, "online_drive"), - ], -) -def test_build_built_in_field_values_maps_source_to_metadata_label( - data_source_type: DataSourceType, source: str -) -> None: - values = build_built_in_field_values( - document_name="doc.txt", - uploader=None, - upload_date=CREATED_AT, - last_update_date=CREATED_AT, - data_source_type=data_source_type, - ) - - assert values[BuiltInField.source] == source - - -@pytest.mark.parametrize( - "data_source_type", - [DataSourceType.UPLOAD_FILE, DataSourceType.NOTION_IMPORT, DataSourceType.WEBSITE_CRAWL], -) -def test_build_document_writes_built_in_fields_in_shared_format( - monkeypatch: pytest.MonkeyPatch, data_source_type: DataSourceType -) -> None: - monkeypatch.setattr(dataset_service, "naive_utc_now", lambda: CREATED_AT) - dataset = make_dataset(built_in_field_enabled=True) - - document = DocumentService.build_document( - dataset=dataset, - process_rule_id=None, - data_source_type=data_source_type, - document_form="text_model", - document_language="English", - data_source_info={}, - created_from="web", - position=1, - account=make_account(name="Alice"), - name="doc.txt", - batch="batch-1", - ) - - assert document.doc_metadata == build_built_in_field_values( - document_name="doc.txt", - uploader="Alice", - upload_date=CREATED_AT, - last_update_date=CREATED_AT, - data_source_type=data_source_type, - ) - - -def test_build_document_skips_built_in_fields_when_disabled() -> None: - dataset = make_dataset(built_in_field_enabled=False) - - document = DocumentService.build_document( - dataset=dataset, - process_rule_id=None, - data_source_type=DataSourceType.UPLOAD_FILE, - document_form="text_model", - document_language="English", - data_source_info={}, - created_from="web", - position=1, - account=make_account(name="Alice"), - name="doc.txt", - batch="batch-1", - ) - - assert not document.doc_metadata - - -FRACTIONAL_CREATED_AT = datetime(2026, 10, 16, 8, 0, 0, 900000) - - -def test_build_built_in_field_values_keeps_fractional_epoch_seconds() -> None: - values = build_built_in_field_values( - document_name="doc.txt", - uploader=None, - upload_date=FRACTIONAL_CREATED_AT, - last_update_date=FRACTIONAL_CREATED_AT, - data_source_type=DataSourceType.UPLOAD_FILE, - ) - - assert values[BuiltInField.upload_date] == pytest.approx(CREATED_AT_EPOCH + 0.9, abs=1e-6) - assert values[BuiltInField.last_update_date] == pytest.approx(CREATED_AT_EPOCH + 0.9, abs=1e-6) - - -@pytest.mark.parametrize("field", [BuiltInField.upload_date, BuiltInField.last_update_date]) -@pytest.mark.parametrize(("condition", "matches"), [("after", True), (">", True), ("=", False), ("is", False)]) -def test_fractional_built_in_dates_keep_integer_second_boundary_semantics( - sqlite_session_factory: sessionmaker[Session], field: BuiltInField, condition: str, matches: bool -) -> None: - values = build_built_in_field_values( - document_name="doc.txt", - uploader=None, - upload_date=FRACTIONAL_CREATED_AT, - last_update_date=FRACTIONAL_CREATED_AT, - data_source_type=DataSourceType.UPLOAD_FILE, - ) - doc_metadata: dict[str, object] = dict(values) - with sqlite_session_factory.begin() as session: - session.add(make_document(document_id="document-1", doc_metadata=doc_metadata)) - - filters = DatasetRetrieval.process_metadata_filter_func(0, condition, field, CREATED_AT_EPOCH, []) - with sqlite_session_factory() as session: - found = session.scalars(select(Document.id).where(*filters)).all() - - assert found == (["document-1"] if matches else []) diff --git a/api/tests/unit_tests/services/test_knowledge_fs_create_research_task_tool.py b/api/tests/unit_tests/services/knowledge_fs/builtin_tool/knowledge_fs/tools/test_create_research_task.py similarity index 100% rename from api/tests/unit_tests/services/test_knowledge_fs_create_research_task_tool.py rename to api/tests/unit_tests/services/knowledge_fs/builtin_tool/knowledge_fs/tools/test_create_research_task.py diff --git a/api/tests/unit_tests/services/test_knowledge_fs_retrieval_v2_metadata_filter.py b/api/tests/unit_tests/services/knowledge_fs/knowledge_retrieval_v2/test_automatic_metadata_filter.py similarity index 100% rename from api/tests/unit_tests/services/test_knowledge_fs_retrieval_v2_metadata_filter.py rename to api/tests/unit_tests/services/knowledge_fs/knowledge_retrieval_v2/test_automatic_metadata_filter.py diff --git a/api/tests/unit_tests/services/test_knowledge_fs_retrieval_v2_node.py b/api/tests/unit_tests/services/knowledge_fs/knowledge_retrieval_v2/test_knowledge_retrieval_v2_node.py similarity index 88% rename from api/tests/unit_tests/services/test_knowledge_fs_retrieval_v2_node.py rename to api/tests/unit_tests/services/knowledge_fs/knowledge_retrieval_v2/test_knowledge_retrieval_v2_node.py index ccc921698a2..2e8863a0b1c 100644 --- a/api/tests/unit_tests/services/test_knowledge_fs_retrieval_v2_node.py +++ b/api/tests/unit_tests/services/knowledge_fs/knowledge_retrieval_v2/test_knowledge_retrieval_v2_node.py @@ -1,14 +1,17 @@ from __future__ import annotations +import math import threading import time from collections.abc import Mapping +from datetime import UTC, datetime from pathlib import Path from typing import override import pytest import yaml from pydantic import ValidationError +from sqlalchemy.orm import Session from core.app.entities.app_invoke_entities import InvokeFrom, UserFrom from core.workflow.node_factory import resolve_workflow_node_class @@ -19,6 +22,7 @@ from graphon.model_runtime.entities.llm_entities import LLMUsage from graphon.model_runtime.entities.rerank_entities import RerankDocument, RerankResult from graphon.runtime import RuntimeState, VariablePool from graphon.variables import FileSegment, StringSegment +from models.knowledge_fs import KnowledgeFSControlSpace, KnowledgeFSControlSpaceState from services.knowledge_fs import query_images as query_image_dependencies from services.knowledge_fs.app_admission_service import ( KnowledgeFSAppAdmissionError, @@ -28,19 +32,212 @@ from services.knowledge_fs.app_admission_service import ( ) from services.knowledge_fs.knowledge_retrieval_v2 import knowledge_retrieval_v2_node as node_module from services.knowledge_fs.knowledge_retrieval_v2.automatic_metadata_filter import KnowledgeFSMetadataExtraction -from services.knowledge_fs.knowledge_retrieval_v2.entities import KnowledgeRetrievalV2NodeData -from services.knowledge_fs.knowledge_retrieval_v2.exc import KnowledgeFSRetrievalConfigurationError +from services.knowledge_fs.knowledge_retrieval_v2.entities import ( + KNOWLEDGE_RETRIEVAL_V2_NODE_TYPE, + KnowledgeRetrievalV2NodeData, +) +from services.knowledge_fs.knowledge_retrieval_v2.exc import ( + KnowledgeFSRetrievalConfigurationError, + KnowledgeFSRetrievalContractError, +) from services.knowledge_fs.knowledge_retrieval_v2.knowledge_retrieval_v2_node import KnowledgeRetrievalV2Node +from services.knowledge_fs.knowledge_retrieval_v2.validation import ( + collect_control_space_ids, + missing_control_space_ids, + validate_control_space_references, +) from services.knowledge_fs.product_dto import ( KnowledgeFSMetadataFieldListResponse, KnowledgeFSRetrievalTestResponse, ) from services.knowledge_fs.product_remote import KnowledgeFSOperationUnavailableError from services.knowledge_fs.query_images import KnowledgeFSWorkflowQueryImageReference +from services.knowledge_fs.retrieval_contracts import KnowledgeFSRetrievalTestItemResponse from tasks import knowledge_fs_failed_retrieval_tasks as failed_retrieval_dependencies from tests.workflow_test_utils import build_test_graph_init_params +def _node_data(**overrides: object) -> dict[str, object]: + return { + "type": KNOWLEDGE_RETRIEVAL_V2_NODE_TYPE, + "title": "Agent Knowledge Retrieval", + "control_space_ids": ["space-a"], + "query_variable_selector": ["start", "query"], + **overrides, + } + + +def _graph(*nodes: object) -> dict[str, object]: + return {"nodes": list(nodes)} + + +def _item( + *, text: str | None, section_path: list[str], node_id: str = "node-a" +) -> KnowledgeFSRetrievalTestItemResponse: + return KnowledgeFSRetrievalTestItemResponse.model_validate( + { + "citation": { + "artifactHash": "a" * 64, + "documentAssetId": "document-a", + "documentTitle": "a.pdf", + "documentVersion": 1, + "knowledgeSpaceName": "Space A", + "logicalDocumentId": "logical-a", + "logicalDocumentRevision": 1, + "pageNumber": 1, + "sectionPath": section_path, + }, + "nodeId": node_id, + "projectionIds": ["projection-a"], + "score": 0.5, + "sources": ["dense"], + "text": text, + } + ) + + +def _candidates(count: int) -> list[node_module._WorkflowRerankCandidate]: + return [ + node_module._WorkflowRerankCandidate( + control_space_id="space-a", + item=_item(text=f"text {index}", section_path=[], node_id=f"node-{index}"), + item_index=index, + space_index=0, + ) + for index in range(count) + ] + + +def _active_space(space_id: str, *, tenant_id: str = "tenant-1") -> KnowledgeFSControlSpace: + space = KnowledgeFSControlSpace( + tenant_id=tenant_id, + owner_account_id="account-1", + provisioning_key=f"provision-{space_id}", + knowledge_space_id=f"knowledge-{space_id}", + state=KnowledgeFSControlSpaceState.ACTIVE, + ) + space.id = space_id + return space + + +def test_collect_control_space_ids_ignores_malformed_graphs_and_foreign_nodes() -> None: + assert collect_control_space_ids({"nodes": "not-a-list"}) == () + assert collect_control_space_ids( + _graph( + "not-a-node", + {"data": "not-a-mapping"}, + {"data": {"type": "knowledge-retrieval"}}, + {"data": _node_data(control_space_ids=["space-b", "space-a"])}, + ) + ) == ("space-b", "space-a") + + +def test_reference_validation_skips_the_database_without_v2_nodes() -> None: + # A graph without KnowledgeFS nodes must not need a session at all. + assert missing_control_space_ids(session=Session(), tenant_id="tenant-1", graph=_graph()) == () + validate_control_space_references(session=Session(), tenant_id="tenant-1", graph=_graph()) + + +def test_reference_validation_rejects_spaces_outside_the_tenant(sqlite_session: Session) -> None: + sqlite_session.add_all([_active_space("space-a"), _active_space("space-b", tenant_id="tenant-2")]) + sqlite_session.flush() + graph = _graph({"data": _node_data(control_space_ids=["space-a", "space-b"])}) + + with pytest.raises(ValueError, match="missing or inactive in this workspace: space-b$"): + validate_control_space_references(session=sqlite_session, tenant_id="tenant-1", graph=graph) + + validate_control_space_references( + session=sqlite_session, + tenant_id="tenant-1", + graph=_graph({"data": _node_data(control_space_ids=["space-a"])}), + ) + + +@pytest.mark.parametrize( + ("overrides", "message"), + [ + ({"control_space_ids": ["space-a", " "]}, "IDs must be non-empty"), + ({"query_variable_selector": ["start"]}, "query variable selector must contain at least two parts"), + ({"query_variable_selector": ["start", " "]}, "query variable selector must be non-empty"), + ({"query_attachment_selector": ["start"]}, "query attachment selector must contain at least two parts"), + ({"query_attachment_selector": ["start", " "]}, "query attachment selector must be non-empty"), + ], +) +def test_node_data_rejects_malformed_identifiers_and_selectors(overrides: dict[str, object], message: str) -> None: + with pytest.raises(ValidationError, match=message): + KnowledgeRetrievalV2NodeData.model_validate(_node_data(**overrides)) + + +def test_node_data_treats_empty_selectors_as_unbound() -> None: + node_data = KnowledgeRetrievalV2NodeData.model_validate( + _node_data(query_variable_selector=[], query_attachment_selector=["start", "image"]) + ) + + assert node_data.query_variable_selector is None + assert node_data.query_attachment_selector == ["start", "image"] + assert ( + KnowledgeRetrievalV2NodeData.model_validate(_node_data(query_attachment_selector=[])).query_attachment_selector + is None + ) + + +def test_variable_mapping_omits_unbound_query_attachment() -> None: + node_data = KnowledgeRetrievalV2NodeData.model_validate(_node_data()) + + assert KnowledgeRetrievalV2Node._extract_variable_selector_to_variable_mapping( + graph_config={}, + node_id="retrieval", + node_data=node_data, + ) == {"retrieval.query": ["start", "query"]} + + +@pytest.mark.parametrize( + ("value", "expected"), + [ + (None, None), + ("tag", "tag"), + (1.5, 1.5), + (3, 3), + (True, "True"), + (datetime(2026, 1, 2, tzinfo=UTC), "2026-01-02 00:00:00+00:00"), + ], +) +def test_manual_metadata_values_are_normalized_to_contract_scalars(value: object, expected: object) -> None: + assert node_module._normalize_metadata_filter_scalar(value) == expected + + +def test_rerank_text_falls_back_to_section_path_then_node_id() -> None: + assert KnowledgeRetrievalV2Node._rerank_text(_item(text=" body ", section_path=["A"])) == "body" + assert KnowledgeRetrievalV2Node._rerank_text(_item(text=" ", section_path=[" Camera ", " ", "Lens"])) == ( + "Camera / Lens" + ) + assert KnowledgeRetrievalV2Node._rerank_text(_item(text=None, section_path=[], node_id="node-z")) == "node-z" + + +def test_rerank_result_validation_keeps_model_order_and_rejects_bad_scores() -> None: + candidates = _candidates(2) + + ranked = KnowledgeRetrievalV2Node._validate_rerank_result( + result=RerankResult( + model="rerank", + docs=[RerankDocument(index=1, text="text 1", score=0.9), RerankDocument(index=0, text="text 0", score=0.1)], + ), + candidates=candidates, + ) + assert ranked == [(0.9, candidates[1]), (0.1, candidates[0])] + + with pytest.raises(KnowledgeFSRetrievalContractError, match="non-finite score"): + KnowledgeRetrievalV2Node._validate_rerank_result( + result=RerankResult(model="rerank", docs=[RerankDocument(index=0, text="text 0", score=math.nan)]), + candidates=candidates, + ) + with pytest.raises(KnowledgeFSRetrievalContractError, match="invalid document indices"): + KnowledgeRetrievalV2Node._validate_rerank_result( + result=RerankResult(model="rerank", docs=[RerankDocument(index=-1, text="text", score=0.5)]), + candidates=candidates, + ) + + def _response(*, mode: str, score: float, space: str, text: str) -> KnowledgeFSRetrievalTestResponse: return KnowledgeFSRetrievalTestResponse.model_validate( { diff --git a/api/tests/unit_tests/services/test_knowledge_fs_app_admission_service.py b/api/tests/unit_tests/services/knowledge_fs/test_app_admission_service.py similarity index 98% rename from api/tests/unit_tests/services/test_knowledge_fs_app_admission_service.py rename to api/tests/unit_tests/services/knowledge_fs/test_app_admission_service.py index 5d14fc1fa65..6d1b8b4ebdf 100644 --- a/api/tests/unit_tests/services/test_knowledge_fs_app_admission_service.py +++ b/api/tests/unit_tests/services/knowledge_fs/test_app_admission_service.py @@ -24,7 +24,7 @@ from services.knowledge_fs.app_admission_service import ( KnowledgeFSAppChannelDisabledError, KnowledgeFSAppSpaceUnavailableError, ) -from tests.unit_tests.services.knowledge_fs_fakes import claims_summary, revoke_payload +from tests.helpers.unit.services.knowledge_fs_fakes import claims_summary, revoke_payload _SQLITE_TABLES = ( KnowledgeFSControlSpace, diff --git a/api/tests/unit_tests/services/test_knowledge_fs_app_binding_management.py b/api/tests/unit_tests/services/knowledge_fs/test_app_binding_management.py similarity index 100% rename from api/tests/unit_tests/services/test_knowledge_fs_app_binding_management.py rename to api/tests/unit_tests/services/knowledge_fs/test_app_binding_management.py diff --git a/api/tests/unit_tests/services/knowledge_fs/test_app_execution_capability.py b/api/tests/unit_tests/services/knowledge_fs/test_app_execution_capability.py new file mode 100644 index 00000000000..0a829f9b43f --- /dev/null +++ b/api/tests/unit_tests/services/knowledge_fs/test_app_execution_capability.py @@ -0,0 +1,891 @@ +from __future__ import annotations + +import base64 +import json +from datetime import UTC, datetime +from typing import override + +import pytest +from pydantic import JsonValue, ValidationError +from sqlalchemy.orm import Session, sessionmaker + +from core.app.entities.app_invoke_entities import DifyRunContext, InvokeFrom, UserFrom +from enums import DeploymentEdition +from models.knowledge_fs import ( + AppKnowledgeFSSpaceJoin, + KnowledgeFSAppSpaceJoinStatus, + KnowledgeFSAppSpaceJoinType, + KnowledgeFSAuthorizationRevision, + KnowledgeFSCapabilityIssuanceAudit, + KnowledgeFSCapabilityIssuanceReservation, + KnowledgeFSControlSpace, + KnowledgeFSControlSpaceState, + KnowledgeFSExternalAccessPolicy, + KnowledgeFSLifecycleOutbox, +) +from services.knowledge_fs import app_execution_capability +from services.knowledge_fs.app_admission_service import ( + KnowledgeFSAppAdmissionError, + KnowledgeFSAppAdmissionService, + KnowledgeFSAppPrincipalProfile, +) +from services.knowledge_fs.app_execution_capability import ( + KnowledgeFSAppExecutionCapabilityService, + KnowledgeResourceRef, +) +from services.knowledge_fs.capability_broker import KnowledgeFSIssuedProductCapability +from services.knowledge_fs.product_dto import ( + KnowledgeFSResearchTaskCreatePayload, + KnowledgeFSRetrievalTestPayload, + KnowledgeFSWorkflowFailedRetrievalCapturePayload, +) +from services.knowledge_fs.product_remote import ( + KNOWLEDGE_FS_QUERY_IMAGE_GRANTS_HEADER, + KnowledgeFSOperationUnavailableError, + KnowledgeFSRemoteJSONRequest, + UnavailableKnowledgeFSProductRemote, +) +from tests.helpers.unit.config_override import apply_config_overrides + + +class _KnowledgeFsAppExecutionBroker: + def __init__(self) -> None: + self.calls: list[dict[str, object]] = [] + + def issue_app(self, **kwargs: object) -> KnowledgeFSIssuedProductCapability: + self.calls.append(kwargs) + return KnowledgeFSIssuedProductCapability( + token="capability", + expires_at=datetime(2030, 1, 1, tzinfo=UTC), + operation_id="createResearchTask", + knowledge_space_id="space-1", + knowledge_space_revision=3, + trace_id="trace-issued", + ) + + +class _KnowledgeFsAppExecutionRemote(UnavailableKnowledgeFSProductRemote): + def __init__(self) -> None: + self.calls: list[KnowledgeFSRemoteJSONRequest] = [] + + @override + def execute_json(self, request: KnowledgeFSRemoteJSONRequest) -> dict[str, JsonValue]: + self.calls.append(request) + return { + "id": "task-1", + "knowledgeSpaceId": "space-1", + "query": "question", + "cost": {}, + "stage": "queued", + "metadata": {}, + "createdAt": 1.0, + "updatedAt": 1.0, + } + + +def _run_context(*, app_id: str) -> DifyRunContext: + return DifyRunContext( + tenant_id="tenant-1", + app_id=app_id, + user_id="user-1", + user_from=UserFrom.END_USER, + invoke_from=InvokeFrom.WEB_APP, + trace_session_id="trace-run", + ) + + +def _seed( + session: Session, + *, + join_status: KnowledgeFSAppSpaceJoinStatus = KnowledgeFSAppSpaceJoinStatus.ACTIVE, + caller_kind: KnowledgeFSAppSpaceJoinType = KnowledgeFSAppSpaceJoinType.AGENT, + channel_enabled: bool = True, +) -> KnowledgeFSControlSpace: + space = KnowledgeFSControlSpace( + tenant_id="tenant-1", + owner_account_id="owner-1", + provisioning_key="provision-1", + knowledge_space_id="space-1", + knowledge_space_revision=3, + state=KnowledgeFSControlSpaceState.ACTIVE, + ) + session.add_all( + [ + space, + KnowledgeFSExternalAccessPolicy( + tenant_id="tenant-1", + control_space_id=space.id, + agent_enabled=channel_enabled if caller_kind is KnowledgeFSAppSpaceJoinType.AGENT else False, + workflow_enabled=channel_enabled if caller_kind is KnowledgeFSAppSpaceJoinType.WORKFLOW else False, + ), + KnowledgeFSAuthorizationRevision( + tenant_id="tenant-1", + control_space_id=space.id, + ), + AppKnowledgeFSSpaceJoin( + tenant_id="tenant-1", + control_space_id=space.id, + app_id="app-1", + join_type=caller_kind, + status=join_status, + ), + ] + ) + session.commit() + return space + + +class TestKnowledgeFsAppExecutionCases: + pytestmark = pytest.mark.parametrize( + "sqlite_session", + [ + ( + KnowledgeFSControlSpace, + KnowledgeFSExternalAccessPolicy, + KnowledgeFSAuthorizationRevision, + KnowledgeFSCapabilityIssuanceAudit, + KnowledgeFSCapabilityIssuanceReservation, + KnowledgeFSLifecycleOutbox, + AppKnowledgeFSSpaceJoin, + ) + ], + indirect=True, + ) + + @staticmethod + @pytest.fixture(autouse=True) + def _enable_knowledge_fs(monkeypatch: pytest.MonkeyPatch) -> None: + apply_config_overrides(monkeypatch, DEPLOYMENT_EDITION=DeploymentEdition.CLOUD, KNOWLEDGE_FS_ENABLED=True) + + @staticmethod + @pytest.mark.parametrize( + ("join_status", "caller_kind", "channel_enabled", "context_app_id"), + [ + ( + KnowledgeFSAppSpaceJoinStatus.REVOKED, + KnowledgeFSAppSpaceJoinType.AGENT, + True, + "app-1", + ), + ( + KnowledgeFSAppSpaceJoinStatus.ACTIVE, + KnowledgeFSAppSpaceJoinType.AGENT, + False, + "app-1", + ), + ( + KnowledgeFSAppSpaceJoinStatus.ACTIVE, + KnowledgeFSAppSpaceJoinType.WORKFLOW, + False, + "app-1", + ), + ( + KnowledgeFSAppSpaceJoinStatus.ACTIVE, + KnowledgeFSAppSpaceJoinType.AGENT, + True, + "app-2", + ), + ], + ids=("binding-disabled", "agent-channel-off", "workflow-channel-off", "cross-app-space"), + ) + def test_execution_fails_closed_before_capability_or_remote_io( + sqlite_session: Session, + join_status: KnowledgeFSAppSpaceJoinStatus, + caller_kind: KnowledgeFSAppSpaceJoinType, + channel_enabled: bool, + context_app_id: str, + ) -> None: + space = _seed( + sqlite_session, + join_status=join_status, + caller_kind=caller_kind, + channel_enabled=channel_enabled, + ) + broker = _KnowledgeFsAppExecutionBroker() + remote = _KnowledgeFsAppExecutionRemote() + service = KnowledgeFSAppExecutionCapabilityService( # type: ignore[arg-type] + admission=KnowledgeFSAppAdmissionService( + sessionmaker(bind=sqlite_session.get_bind(), expire_on_commit=False) + ), + broker=broker, + remote=remote, + ) + + with pytest.raises(KnowledgeFSAppAdmissionError): + service.create_research_task( + run_context=_run_context(app_id=context_app_id), + caller_kind=caller_kind, + resource=KnowledgeResourceRef(kind="knowledge_fs", control_space_id=space.id), + payload=KnowledgeFSResearchTaskCreatePayload(query="question"), + ) + + assert broker.calls == [] + assert remote.calls == [] + + @staticmethod + @pytest.mark.parametrize( + "caller_kind", + [KnowledgeFSAppSpaceJoinType.AGENT, KnowledgeFSAppSpaceJoinType.WORKFLOW], + ) + def test_execution_issues_capability_then_performs_bounded_remote_io( + sqlite_session: Session, + caller_kind: KnowledgeFSAppSpaceJoinType, + ) -> None: + space = _seed(sqlite_session, caller_kind=caller_kind) + broker = _KnowledgeFsAppExecutionBroker() + remote = _KnowledgeFsAppExecutionRemote() + service = KnowledgeFSAppExecutionCapabilityService( # type: ignore[arg-type] + admission=KnowledgeFSAppAdmissionService( + sessionmaker(bind=sqlite_session.get_bind(), expire_on_commit=False) + ), + broker=broker, + remote=remote, + ) + + response = service.create_research_task( + run_context=_run_context(app_id="app-1"), + caller_kind=caller_kind, + resource=KnowledgeResourceRef(kind="knowledge_fs", control_space_id=space.id), + payload=KnowledgeFSResearchTaskCreatePayload(query="question"), + ) + + assert response.id == "task-1" + assert len(broker.calls) == 1 + assert remote.calls == [ + KnowledgeFSRemoteJSONRequest( + operation_id="createResearchTask", + method="POST", + path="/research-tasks", + namespace_id="tenant-1", + knowledge_space_id="space-1", + capability_token="capability", + trace_id="trace-issued", + payload={ + "query": "question", + "metadata": {}, + "knowledgeSpaceId": "space-1", + }, + ) + ] + + +class Admission: + def __init__(self) -> None: + self.calls: list[dict[str, object]] = [] + self.profile = KnowledgeFSAppPrincipalProfile( + tenant_id="tenant-1", + control_space_id="control-1", + app_id="app-1", + join_id="join-1", + caller_kind=KnowledgeFSAppSpaceJoinType.WORKFLOW, + action="read", + knowledge_space_id="space-1", + knowledge_space_revision=2, + membership_epoch=1, + space_acl_epoch=1, + external_access_epoch=1, + content_policy_revision=1, + ) + + def admit(self, **kwargs: object) -> KnowledgeFSAppPrincipalProfile: + self.calls.append(kwargs) + return self.profile + + +class _KnowledgeFsAppExecutionCapabilityBroker: + def __init__(self) -> None: + self.calls: list[dict[str, object]] = [] + + def issue_app(self, **kwargs: object) -> KnowledgeFSIssuedProductCapability: + self.calls.append(kwargs) + return KnowledgeFSIssuedProductCapability( + token="token", + expires_at=datetime(2030, 1, 1, tzinfo=UTC), + operation_id=str(kwargs["operation_id"]), + knowledge_space_id="space-1", + knowledge_space_revision=2, + trace_id="trace-1", + ) + + +class _KnowledgeFsAppExecutionCapabilityRemote(UnavailableKnowledgeFSProductRemote): + def __init__(self) -> None: + self.calls: list[KnowledgeFSRemoteJSONRequest] = [] + + @override + def execute_json(self, request: KnowledgeFSRemoteJSONRequest) -> dict[str, JsonValue]: + self.calls.append(request) + if request.operation_id == "captureAgentKnowledgeInvestigation": + return {"investigationId": "investigation-1", "failedQueries": []} + if request.operation_id == "captureWorkflowFailedRetrieval": + return { + "failedQueryId": "019fac9f-bfb0-75ee-9af5-252ebafbac1c", + "verdict": "retrieval-miss", + "badCaseId": "019fac9f-bfb0-75ee-9af5-252ebafbac1d", + } + if request.operation_id == "listMetadataFields": + return { + "items": [ + { + "count": 3, + "createdAt": "2026-01-01T00:00:00Z", + "id": "field-1", + "name": "department", + "rowVersion": 1, + "type": "string", + "updatedAt": "2026-01-01T00:00:00Z", + } + ], + "nextCursor": "cursor-2", + } + if request.operation_id == "retrieveEvidence": + return { + "items": [ + { + "citation": { + "artifactHash": "a" * 64, + "documentAssetId": "document-1", + "documentVersion": 1, + "sectionPath": ["Camera"], + }, + "nodeId": "node-1", + "projectionIds": ["projection-1"], + "score": 0.8, + "sources": ["dense"], + "text": "Camera evidence", + } + ], + "metrics": {"denseCandidates": 1, "totalMs": 10}, + "mode": "fast", + "traceId": "trace-1", + } + return { + "id": "research-1", + "knowledgeSpaceId": "space-1", + "query": "Compare the evidence", + "cost": {}, + "stage": "queued", + "metadata": {}, + "createdAt": 1.0, + "updatedAt": 1.0, + } + + +class FailingRemote(_KnowledgeFsAppExecutionCapabilityRemote): + @override + def execute_json(self, request: KnowledgeFSRemoteJSONRequest) -> dict[str, JsonValue]: + self.calls.append(request) + raise RuntimeError("remote failed") + + +class TestKnowledgeFsAppExecutionCapabilityCases: + @staticmethod + @pytest.fixture(autouse=True) + def _enable_knowledge_fs(monkeypatch: pytest.MonkeyPatch) -> None: + apply_config_overrides(monkeypatch, DEPLOYMENT_EDITION=DeploymentEdition.CLOUD, KNOWLEDGE_FS_ENABLED=True) + + @staticmethod + def test_app_execution_capability_always_admits_binding_before_broker_issuance() -> None: + admission = Admission() + broker = _KnowledgeFsAppExecutionCapabilityBroker() + service = KnowledgeFSAppExecutionCapabilityService( # type: ignore[arg-type] + admission=admission, + broker=broker, + remote=_KnowledgeFsAppExecutionCapabilityRemote(), + ) + + issued = service.issue( + tenant_id="tenant-1", + app_id="app-1", + control_space_id="control-1", + caller_kind=KnowledgeFSAppSpaceJoinType.WORKFLOW, + operation_id="createResearchTask", + resource_id="research-1", + trace_id="trace-1", + ) + + assert issued.token == "token" + assert admission.calls == [ + { + "tenant_id": "tenant-1", + "app_id": "app-1", + "control_space_id": "control-1", + "caller_kind": KnowledgeFSAppSpaceJoinType.WORKFLOW, + "operation_id": "createResearchTask", + } + ] + assert broker.calls == [ + { + "profile": admission.profile, + "operation_id": "createResearchTask", + "resource_id": "research-1", + "trace_id": "trace-1", + } + ] + + @staticmethod + def test_disabled_feature_blocks_app_issuance_and_investigation_capture_before_io( + monkeypatch: pytest.MonkeyPatch, + ) -> None: + admission = Admission() + broker = _KnowledgeFsAppExecutionCapabilityBroker() + remote = _KnowledgeFsAppExecutionCapabilityRemote() + service = KnowledgeFSAppExecutionCapabilityService( # type: ignore[arg-type] + admission=admission, + broker=broker, + remote=remote, + ) + previously_issued = KnowledgeFSIssuedProductCapability( + token="token", + expires_at=datetime(2030, 1, 1, tzinfo=UTC), + operation_id="captureAgentKnowledgeInvestigation", + knowledge_space_id="space-1", + knowledge_space_revision=2, + trace_id="trace-1", + ) + apply_config_overrides(monkeypatch, KNOWLEDGE_FS_ENABLED=False) + + with pytest.raises(KnowledgeFSOperationUnavailableError, match="disabled"): + service.issue( + tenant_id="tenant-1", + app_id="app-1", + control_space_id="control-1", + caller_kind=KnowledgeFSAppSpaceJoinType.AGENT, + operation_id="retrieveEvidence", + ) + with pytest.raises(KnowledgeFSOperationUnavailableError, match="disabled"): + service.create_research_task( + run_context=DifyRunContext( + tenant_id="tenant-1", + app_id="app-1", + user_id="user-1", + user_from=UserFrom.END_USER, + invoke_from=InvokeFrom.WEB_APP, + ), + caller_kind=KnowledgeFSAppSpaceJoinType.AGENT, + resource=KnowledgeResourceRef(kind="knowledge_fs", control_space_id="control-1"), + payload=KnowledgeFSResearchTaskCreatePayload(query="Compare the evidence"), + ) + with pytest.raises(KnowledgeFSOperationUnavailableError, match="disabled"): + service.run_retrieval( + run_context=DifyRunContext( + tenant_id="tenant-1", + app_id="app-1", + user_id="user-1", + user_from=UserFrom.END_USER, + invoke_from=InvokeFrom.WEB_APP, + ), + caller_kind=KnowledgeFSAppSpaceJoinType.WORKFLOW, + resource=KnowledgeResourceRef(kind="knowledge_fs", control_space_id="control-1"), + payload=KnowledgeFSRetrievalTestPayload.model_validate( + { + "includeText": True, + "mode": "fast", + "query": "camera", + "queryId": "10000000-0000-4000-8000-000000000010", + } + ), + ) + with pytest.raises(KnowledgeFSOperationUnavailableError, match="disabled"): + service.capture_agent_investigation(issued=previously_issued, tenant_id="tenant-1", payload={}) + + assert admission.calls == [] + assert broker.calls == [] + assert remote.calls == [] + + @staticmethod + def test_typed_knowledge_resource_ref_rejects_dataset_and_extra_authority_fields() -> None: + with pytest.raises(ValidationError): + KnowledgeResourceRef.model_validate( + { + "kind": "dataset", + "control_space_id": "control-1", + } + ) + + with pytest.raises(ValidationError): + KnowledgeResourceRef.model_validate( + { + "kind": "knowledge_fs", + "control_space_id": "control-1", + "dataset_ids": ["control-1"], + } + ) + + with pytest.raises(ValidationError): + KnowledgeResourceRef.model_validate( + { + "kind": "knowledge_fs", + "control_space_id": " ", + } + ) + + @staticmethod + def test_create_research_task_fails_closed_when_product_operation_is_unavailable( + monkeypatch: pytest.MonkeyPatch, + ) -> None: + admission = Admission() + broker = _KnowledgeFsAppExecutionCapabilityBroker() + remote = _KnowledgeFsAppExecutionCapabilityRemote() + service = KnowledgeFSAppExecutionCapabilityService( # type: ignore[arg-type] + admission=admission, + broker=broker, + remote=remote, + ) + monkeypatch.setattr(app_execution_capability, "is_product_operation_ready", lambda _operation_id: False) + + with pytest.raises(KnowledgeFSOperationUnavailableError, match="creation is unavailable"): + service.create_research_task( + run_context=DifyRunContext( + tenant_id="tenant-1", + app_id="app-1", + user_id="user-1", + user_from=UserFrom.END_USER, + invoke_from=InvokeFrom.WEB_APP, + ), + caller_kind=KnowledgeFSAppSpaceJoinType.AGENT, + resource=KnowledgeResourceRef(kind="knowledge_fs", control_space_id="control-1"), + payload=KnowledgeFSResearchTaskCreatePayload(query="Compare the evidence"), + ) + + assert admission.calls == [] + assert broker.calls == [] + assert remote.calls == [] + + @staticmethod + def test_create_research_task_uses_only_dify_run_context_identity_then_calls_kfs() -> None: + admission = Admission() + broker = _KnowledgeFsAppExecutionCapabilityBroker() + remote = _KnowledgeFsAppExecutionCapabilityRemote() + service = KnowledgeFSAppExecutionCapabilityService( # type: ignore[arg-type] + admission=admission, + broker=broker, + remote=remote, + ) + run_context = DifyRunContext( + tenant_id="tenant-from-run-context", + app_id="app-from-run-context", + user_id="user-1", + user_from=UserFrom.END_USER, + invoke_from=InvokeFrom.WEB_APP, + trace_session_id="trace-from-run-context", + ) + + result = service.create_research_task( + run_context=run_context, + caller_kind=KnowledgeFSAppSpaceJoinType.AGENT, + resource=KnowledgeResourceRef(kind="knowledge_fs", control_space_id="control-1"), + payload=KnowledgeFSResearchTaskCreatePayload(query="Compare the evidence"), + ) + + assert result.id == "research-1" + assert admission.calls == [ + { + "tenant_id": "tenant-from-run-context", + "app_id": "app-from-run-context", + "control_space_id": "control-1", + "caller_kind": KnowledgeFSAppSpaceJoinType.AGENT, + "operation_id": "createResearchTask", + } + ] + assert broker.calls == [ + { + "profile": admission.profile, + "operation_id": "createResearchTask", + "resource_id": None, + "trace_id": "trace-from-run-context", + } + ] + assert remote.calls == [ + KnowledgeFSRemoteJSONRequest( + operation_id="createResearchTask", + method="POST", + path="/research-tasks", + namespace_id="tenant-from-run-context", + knowledge_space_id="space-1", + capability_token="token", + trace_id="trace-1", + payload={ + "query": "Compare the evidence", + "metadata": {}, + "knowledgeSpaceId": "space-1", + }, + ) + ] + + @staticmethod + def test_create_research_task_propagates_remote_io_failure() -> None: + service = KnowledgeFSAppExecutionCapabilityService( # type: ignore[arg-type] + admission=Admission(), + broker=_KnowledgeFsAppExecutionCapabilityBroker(), + remote=FailingRemote(), + ) + + with pytest.raises(RuntimeError, match="remote failed"): + service.create_research_task( + run_context=DifyRunContext( + tenant_id="tenant-1", + app_id="app-1", + user_id="user-1", + user_from=UserFrom.END_USER, + invoke_from=InvokeFrom.WEB_APP, + ), + caller_kind=KnowledgeFSAppSpaceJoinType.AGENT, + resource=KnowledgeResourceRef(kind="knowledge_fs", control_space_id="control-1"), + payload=KnowledgeFSResearchTaskCreatePayload(query="Compare the evidence"), + ) + + @staticmethod + def test_run_retrieval_issues_read_capability_and_calls_bounded_product_operation() -> None: + admission = Admission() + broker = _KnowledgeFsAppExecutionCapabilityBroker() + remote = _KnowledgeFsAppExecutionCapabilityRemote() + service = KnowledgeFSAppExecutionCapabilityService( # type: ignore[arg-type] + admission=admission, + broker=broker, + remote=remote, + ) + run_context = DifyRunContext( + tenant_id="tenant-1", + app_id="app-1", + user_id="user-1", + user_from=UserFrom.ACCOUNT, + invoke_from=InvokeFrom.DEBUGGER, + trace_session_id="trace-from-workflow", + ) + + result = service.run_retrieval( + run_context=run_context, + caller_kind=KnowledgeFSAppSpaceJoinType.WORKFLOW, + resource=KnowledgeResourceRef(kind="knowledge_fs", control_space_id="control-1"), + payload=KnowledgeFSRetrievalTestPayload.model_validate( + { + "includeText": True, + "mode": "fast", + "query": "camera", + "queryId": "10000000-0000-4000-8000-000000000010", + "queryImages": [ + { + "accessGrant": "short-lived-grant", + "uploadFileId": "00000000-0000-4000-8000-000000000001", + } + ], + } + ), + ) + + assert result.items[0].text == "Camera evidence" + assert admission.calls[0]["operation_id"] == "retrieveEvidence" + assert broker.calls[0]["operation_id"] == "retrieveEvidence" + assert len(remote.calls) == 1 + request = remote.calls[0] + assert request.payload == { + "includeText": True, + "mode": "fast", + "query": "camera", + "queryId": "10000000-0000-4000-8000-000000000010", + "queryImages": [{"uploadFileId": "00000000-0000-4000-8000-000000000001"}], + } + assert len(request.headers) == 1 + assert request.headers[0][0] == KNOWLEDGE_FS_QUERY_IMAGE_GRANTS_HEADER + encoded_grants = request.headers[0][1] + padded_grants = encoded_grants + "=" * (-len(encoded_grants) % 4) + assert json.loads(base64.urlsafe_b64decode(padded_grants)) == { + "g": ["short-lived-grant"], + "v": 1, + } + + @staticmethod + def test_list_metadata_fields_issues_workflow_capability_and_calls_catalog_operation() -> None: + admission = Admission() + broker = _KnowledgeFsAppExecutionCapabilityBroker() + remote = _KnowledgeFsAppExecutionCapabilityRemote() + service = KnowledgeFSAppExecutionCapabilityService( # type: ignore[arg-type] + admission=admission, + broker=broker, + remote=remote, + ) + run_context = DifyRunContext( + tenant_id="tenant-1", + app_id="app-1", + user_id="user-1", + user_from=UserFrom.ACCOUNT, + invoke_from=InvokeFrom.DEBUGGER, + trace_session_id="trace-from-workflow", + ) + + first_page = service.list_metadata_fields( + run_context=run_context, + caller_kind=KnowledgeFSAppSpaceJoinType.WORKFLOW, + resource=KnowledgeResourceRef(kind="knowledge_fs", control_space_id="control-1"), + ) + second_page = service.list_metadata_fields( + run_context=run_context, + caller_kind=KnowledgeFSAppSpaceJoinType.WORKFLOW, + resource=KnowledgeResourceRef(kind="knowledge_fs", control_space_id="control-1"), + cursor=first_page.next_cursor, + limit=50, + ) + + assert [item.name for item in first_page.data] == ["department"] + assert first_page.data[0].type == "string" + assert first_page.next_cursor == "cursor-2" + assert second_page.next_cursor == "cursor-2" + assert [call["operation_id"] for call in admission.calls] == ["listMetadataFields", "listMetadataFields"] + assert [call["caller_kind"] for call in admission.calls] == [ + KnowledgeFSAppSpaceJoinType.WORKFLOW, + KnowledgeFSAppSpaceJoinType.WORKFLOW, + ] + assert [call["operation_id"] for call in broker.calls] == ["listMetadataFields", "listMetadataFields"] + assert len(remote.calls) == 2 + first_request, second_request = remote.calls + assert first_request.operation_id == "listMetadataFields" + assert first_request.method == "GET" + assert first_request.path == "/knowledge-spaces/space-1/metadata-fields" + assert first_request.namespace_id == "tenant-1" + assert first_request.knowledge_space_id == "space-1" + assert first_request.capability_token == "token" + assert first_request.payload is None + assert first_request.query == (("limit", "100"),) + assert second_request.query == (("cursor", "cursor-2"), ("limit", "50")) + + @staticmethod + def test_list_metadata_fields_fails_closed_when_product_operation_is_unavailable( + monkeypatch: pytest.MonkeyPatch, + ) -> None: + admission = Admission() + broker = _KnowledgeFsAppExecutionCapabilityBroker() + remote = _KnowledgeFsAppExecutionCapabilityRemote() + service = KnowledgeFSAppExecutionCapabilityService( # type: ignore[arg-type] + admission=admission, + broker=broker, + remote=remote, + ) + monkeypatch.setattr(app_execution_capability, "is_product_operation_ready", lambda _operation_id: False) + + with pytest.raises(KnowledgeFSOperationUnavailableError): + service.list_metadata_fields( + run_context=DifyRunContext( + tenant_id="tenant-1", + app_id="app-1", + user_id="user-1", + user_from=UserFrom.ACCOUNT, + invoke_from=InvokeFrom.DEBUGGER, + ), + caller_kind=KnowledgeFSAppSpaceJoinType.WORKFLOW, + resource=KnowledgeResourceRef(kind="knowledge_fs", control_space_id="control-1"), + ) + + assert admission.calls == [] + assert broker.calls == [] + assert remote.calls == [] + + @staticmethod + def test_capture_workflow_failed_retrieval_uses_fresh_transport_trace_and_business_event_id() -> None: + admission = Admission() + broker = _KnowledgeFsAppExecutionCapabilityBroker() + remote = _KnowledgeFsAppExecutionCapabilityRemote() + service = KnowledgeFSAppExecutionCapabilityService( # type: ignore[arg-type] + admission=admission, + broker=broker, + remote=remote, + ) + + result = service.capture_workflow_failed_retrieval( + tenant_id="tenant-1", + app_id="app-1", + resource=KnowledgeResourceRef(kind="knowledge_fs", control_space_id="control-1"), + payload=KnowledgeFSWorkflowFailedRetrievalCapturePayload( + event_id="019fac9f-bfb0-75ee-9af5-252ebafbac1e", + query=" missing answer ", + mode="fast", + retrieval_trace_id=" trace-retrieval-1 ", + ), + ) + + assert str(result.failed_query_id) == "019fac9f-bfb0-75ee-9af5-252ebafbac1c" + assert result.verdict == "retrieval-miss" + assert str(result.bad_case_id) == "019fac9f-bfb0-75ee-9af5-252ebafbac1d" + assert admission.calls == [ + { + "tenant_id": "tenant-1", + "app_id": "app-1", + "control_space_id": "control-1", + "caller_kind": KnowledgeFSAppSpaceJoinType.WORKFLOW, + "operation_id": "captureWorkflowFailedRetrieval", + } + ] + assert broker.calls == [ + { + "profile": admission.profile, + "operation_id": "captureWorkflowFailedRetrieval", + "resource_id": None, + "trace_id": None, + } + ] + assert remote.calls == [ + KnowledgeFSRemoteJSONRequest( + operation_id="captureWorkflowFailedRetrieval", + method="POST", + path="/knowledge-spaces/space-1/failed-queries/workflow-retrieval-misses", + namespace_id="tenant-1", + knowledge_space_id="space-1", + capability_token="token", + trace_id="trace-1", + payload={ + "eventId": "019fac9f-bfb0-75ee-9af5-252ebafbac1e", + "query": "missing answer", + "mode": "fast", + "retrievalTraceId": "trace-retrieval-1", + }, + ) + ] + + @staticmethod + def test_agent_investigation_uses_authorized_space_and_preserves_report_identity() -> None: + broker = _KnowledgeFsAppExecutionCapabilityBroker() + remote = _KnowledgeFsAppExecutionCapabilityRemote() + service = KnowledgeFSAppExecutionCapabilityService(admission=Admission(), broker=broker, remote=remote) + issued = broker.issue_app(operation_id="captureAgentKnowledgeInvestigation") + payload: dict[str, JsonValue] = {"investigationId": "investigation-1", "query": "refund", "attempts": []} + + result = service.capture_agent_investigation(issued=issued, tenant_id="tenant-1", payload=payload) + + assert result == {"investigationId": "investigation-1", "failedQueries": []} + assert remote.calls == [ + KnowledgeFSRemoteJSONRequest( + operation_id="captureAgentKnowledgeInvestigation", + method="POST", + path="/knowledge-spaces/space-1/agent-investigations", + namespace_id="tenant-1", + knowledge_space_id="space-1", + capability_token="token", + trace_id="trace-1", + payload=payload, + ) + ] + + @staticmethod + @pytest.mark.parametrize("unavailable", ["wrong-capability", "unready-operation", "missing-path"]) + def test_agent_investigation_refuses_unavailable_or_wrong_capability_before_io( + monkeypatch: pytest.MonkeyPatch, unavailable: str + ) -> None: + broker = _KnowledgeFsAppExecutionCapabilityBroker() + remote = _KnowledgeFsAppExecutionCapabilityRemote() + service = KnowledgeFSAppExecutionCapabilityService(admission=Admission(), broker=broker, remote=remote) + issued = broker.issue_app(operation_id="captureAgentKnowledgeInvestigation") + if unavailable == "wrong-capability": + issued = issued._replace(operation_id="retrieveEvidence") + elif unavailable == "unready-operation": + monkeypatch.setattr(app_execution_capability, "is_product_operation_ready", lambda _operation_id: False) + else: + operation_id = "captureAgentKnowledgeInvestigation" + operations = dict(app_execution_capability.KNOWLEDGE_FS_PRODUCT_OPERATIONS) + operations[operation_id] = operations[operation_id]._replace(kfs_path=None) + monkeypatch.setattr(app_execution_capability, "KNOWLEDGE_FS_PRODUCT_OPERATIONS", operations) + monkeypatch.setattr(app_execution_capability, "is_product_operation_ready", lambda _operation_id: True) + + with pytest.raises(KnowledgeFSOperationUnavailableError): + service.capture_agent_investigation(issued=issued, tenant_id="tenant-1", payload={}) + + assert remote.calls == [] diff --git a/api/tests/unit_tests/services/test_knowledge_fs_background.py b/api/tests/unit_tests/services/knowledge_fs/test_background_engine.py similarity index 100% rename from api/tests/unit_tests/services/test_knowledge_fs_background.py rename to api/tests/unit_tests/services/knowledge_fs/test_background_engine.py diff --git a/api/tests/unit_tests/services/test_knowledge_fs_batch_capability.py b/api/tests/unit_tests/services/knowledge_fs/test_batch_capability.py similarity index 98% rename from api/tests/unit_tests/services/test_knowledge_fs_batch_capability.py rename to api/tests/unit_tests/services/knowledge_fs/test_batch_capability.py index 77b9dd1b533..a0e3ce7002d 100644 --- a/api/tests/unit_tests/services/test_knowledge_fs_batch_capability.py +++ b/api/tests/unit_tests/services/knowledge_fs/test_batch_capability.py @@ -14,7 +14,7 @@ from services.knowledge_fs.batch_capability import ( ) from services.knowledge_fs.product_remote import KnowledgeFSOperationUnavailableError from services.knowledge_fs_capability import CapabilityIssueRequest -from tests.unit_tests.knowledge_fs_core_fakes import FakeCapabilityClaims, FakeIssuedCapability +from tests.helpers.unit.knowledge_fs_core_fakes import FakeCapabilityClaims, FakeIssuedCapability class FakeIssuer: diff --git a/api/tests/unit_tests/services/test_knowledge_fs_buffered_upload_admission.py b/api/tests/unit_tests/services/knowledge_fs/test_buffered_upload_admission.py similarity index 100% rename from api/tests/unit_tests/services/test_knowledge_fs_buffered_upload_admission.py rename to api/tests/unit_tests/services/knowledge_fs/test_buffered_upload_admission.py diff --git a/api/tests/unit_tests/services/test_knowledge_fs_capability_broker.py b/api/tests/unit_tests/services/knowledge_fs/test_capability_broker.py similarity index 99% rename from api/tests/unit_tests/services/test_knowledge_fs_capability_broker.py rename to api/tests/unit_tests/services/knowledge_fs/test_capability_broker.py index 1dcec19308a..968eb715d38 100644 --- a/api/tests/unit_tests/services/test_knowledge_fs_capability_broker.py +++ b/api/tests/unit_tests/services/knowledge_fs/test_capability_broker.py @@ -38,8 +38,8 @@ from services.knowledge_fs.service_api_authorization import ( KnowledgeFSServiceApiScopeError, ) from services.knowledge_fs_capability import CapabilityIssueRequest -from tests.unit_tests.services.knowledge_fs_fakes import revoke_payload -from tests.unit_tests.services.knowledge_fs_services_a_fakes import upgrade_job +from tests.helpers.unit.services.knowledge_fs_fakes import revoke_payload +from tests.helpers.unit.services.knowledge_fs_services_a_fakes import upgrade_job class FakeProduct: diff --git a/api/tests/unit_tests/services/test_knowledge_fs_cleanup.py b/api/tests/unit_tests/services/knowledge_fs/test_cleanup.py similarity index 100% rename from api/tests/unit_tests/services/test_knowledge_fs_cleanup.py rename to api/tests/unit_tests/services/knowledge_fs/test_cleanup.py diff --git a/api/tests/unit_tests/services/test_knowledge_fs_control_plane_service.py b/api/tests/unit_tests/services/knowledge_fs/test_control_plane_service.py similarity index 99% rename from api/tests/unit_tests/services/test_knowledge_fs_control_plane_service.py rename to api/tests/unit_tests/services/knowledge_fs/test_control_plane_service.py index 845f5bb0b03..7025e987a03 100644 --- a/api/tests/unit_tests/services/test_knowledge_fs_control_plane_service.py +++ b/api/tests/unit_tests/services/knowledge_fs/test_control_plane_service.py @@ -27,7 +27,7 @@ from services.knowledge_fs.product_dto import ( from services.knowledge_fs.product_operations import KnowledgeFSProductPermission from services.knowledge_fs.product_service import AuthorizedKnowledgeFSControlSpace from services.knowledge_fs.revocation_commands import KnowledgeFSRevocationCommandError -from tests.unit_tests.services.knowledge_fs_fakes import claims_summary, revoke_payload +from tests.helpers.unit.services.knowledge_fs_fakes import claims_summary, revoke_payload class FakeProduct: diff --git a/api/tests/unit_tests/services/test_knowledge_fs_control_space_commands.py b/api/tests/unit_tests/services/knowledge_fs/test_control_space_commands.py similarity index 55% rename from api/tests/unit_tests/services/test_knowledge_fs_control_space_commands.py rename to api/tests/unit_tests/services/knowledge_fs/test_control_space_commands.py index f8ebe9bd4a1..0e5c2b84fdb 100644 --- a/api/tests/unit_tests/services/test_knowledge_fs_control_space_commands.py +++ b/api/tests/unit_tests/services/knowledge_fs/test_control_space_commands.py @@ -4,9 +4,11 @@ from typing import cast from unittest.mock import patch import pytest -from sqlalchemy import func, select +from pydantic import ValidationError +from sqlalchemy import delete, func, select from sqlalchemy.orm import Session, sessionmaker +from models import TenantAccountJoin from models.enums import ApiTokenBindingResourceType, ApiTokenType from models.knowledge_fs import ( AppKnowledgeFSSpaceJoin, @@ -17,20 +19,29 @@ from models.knowledge_fs import ( KnowledgeFSControlSpace, KnowledgeFSControlSpacePermission, KnowledgeFSControlSpacePermissionRole, + KnowledgeFSControlSpacePermissionStatus, KnowledgeFSControlSpaceState, + KnowledgeFSControlSpaceVisibility, KnowledgeFSDeleteCommandPayload, KnowledgeFSExternalAccessPolicy, KnowledgeFSLifecycleOperation, KnowledgeFSLifecycleOutbox, + KnowledgeFSProvisionCommandPayload, ) from models.model import ApiToken, DatasetApiTokenBinding +from repositories.sqlalchemy_knowledge_fs_lifecycle_outbox_repository import ( + SQLAlchemyKnowledgeFSLifecycleOutboxRepository, +) from services.api_token_service import ApiTokenCache from services.knowledge_fs.control_space_commands import ( KnowledgeFSControlSpaceCommandService, + KnowledgeFSControlSpaceIntentConflictError, KnowledgeFSProvisionIntent, ) from services.knowledge_fs.control_space_lifecycle import KnowledgeFSWorkspaceDeletionBlockedError -from tests.unit_tests.services.knowledge_fs_fakes import provision_payload, reservation_summary +from services.knowledge_fs.product_dto import KnowledgeFSSpaceCreatePayload +from services.knowledge_fs.workspace_members import KnowledgeFSControlPlaneInvariantError +from tests.helpers.unit.services.knowledge_fs_fakes import provision_payload, reservation_summary def _service(sqlite_session: Session) -> KnowledgeFSControlSpaceCommandService: @@ -349,3 +360,230 @@ def test_deletion_invalidates_token_cache_after_commit(sqlite_session: Session) idempotency_key="cache-delete", ) eviction.assert_called_once_with("tenant-1") + + +def _intent() -> KnowledgeFSProvisionIntent: + return KnowledgeFSProvisionIntent( + tenant_id="tenant-1", + owner_account_id="owner-1", + provisioning_key="provision-1", + operation_id="operation-1", + idempotency_key="create-1", + name="Selected members", + slug="selected-members", + icon=None, + description=None, + model_intent=None, + profile_intent=None, + visibility=KnowledgeFSControlSpaceVisibility.PARTIAL_MEMBERS, + member_account_ids=("member-2", "member-1"), + ) + + +def _provision_payload(outbox: KnowledgeFSLifecycleOutbox) -> KnowledgeFSProvisionCommandPayload: + assert outbox.operation is KnowledgeFSLifecycleOperation.PROVISION + return cast(KnowledgeFSProvisionCommandPayload, outbox.command_payload) + + +def _seed_members(session: Session) -> None: + session.add_all( + [ + TenantAccountJoin(tenant_id="tenant-1", account_id="member-1"), + TenantAccountJoin(tenant_id="tenant-1", account_id="member-2"), + TenantAccountJoin(tenant_id="other-tenant", account_id="outsider"), + ] + ) + session.commit() + + +def _assert_no_creation(session: Session) -> None: + for model in ( + KnowledgeFSControlSpace, + KnowledgeFSControlSpacePermission, + KnowledgeFSAuthorizationRevision, + KnowledgeFSExternalAccessPolicy, + KnowledgeFSLifecycleOutbox, + ): + assert session.scalar(select(func.count()).select_from(model)) == 0 + + +def test_partial_members_create_dto_accepts_viewers_and_preserves_default_creation() -> None: + payload = KnowledgeFSSpaceCreatePayload.model_validate( + { + "name": "Selected members", + "slug": "selected-members", + "visibility": "partial_members", + "members": [{"account_id": "member-1", "role": "viewer"}], + } + ) + assert payload.members[0].account_id == "member-1" + assert payload.members[0].role == "viewer" + assert KnowledgeFSSpaceCreatePayload(name="Private", slug="private").members == [] + assert KnowledgeFSSpaceCreatePayload(name="Team", slug="team", visibility="all_team_members").members == [] + + +@pytest.mark.parametrize( + "updates", + [ + {"visibility": "partial_members"}, + {"visibility": "partial_members", "members": []}, + {"visibility": "partial_members", "members": None}, + {"members": [{"account_id": "member-1", "role": "viewer"}]}, + {"visibility": "all_team_members", "members": [{"account_id": "member-1", "role": "viewer"}]}, + {"visibility": "partial_members", "members": [{"account_id": "member-1", "role": "owner"}]}, + {"visibility": "partial_members", "members": [{"account_id": "member-1", "role": "editor"}]}, + {"visibility": "partial_members", "members": [{"account_id": "member-1"}]}, + {"visibility": "partial_members", "members": [{"account_id": "", "role": "viewer"}]}, + { + "visibility": "partial_members", + "members": [{"account_id": "member-1", "role": "viewer"}] * 2, + }, + { + "visibility": "partial_members", + "members": [{"account_id": f"member-{i}", "role": "viewer"} for i in range(1_001)], + }, + ], +) +def test_partial_members_create_dto_rejects_invalid_or_privileged_members(updates: dict[str, object]) -> None: + with pytest.raises(ValidationError): + KnowledgeFSSpaceCreatePayload.model_validate({"name": "Space", "slug": "space", **updates}) + + +def test_initial_members_and_visibility_are_committed_with_the_space_and_outbox( + sqlite_session: Session, sqlite_session_factory: sessionmaker[Session] +) -> None: + _seed_members(sqlite_session) + service = KnowledgeFSControlSpaceCommandService(sqlite_session_factory) + + result = service.create_provision_intent(_intent()) + replay = service.create_provision_intent(_intent()._replace(member_account_ids=("member-1", "member-2"))) + + assert result.control_space.visibility is KnowledgeFSControlSpaceVisibility.PARTIAL_MEMBERS + assert replay.control_space.id == result.control_space.id + assert replay.outbox.id == result.outbox.id + assert result.outbox.expected_control_space_version == result.control_space.resource_version == 0 + command_payload = _provision_payload(result.outbox) + assert command_payload["initial_visibility"] == "partial_members" + assert command_payload["initial_member_account_ids"] == ["member-1", "member-2"] + permissions = list(sqlite_session.scalars(select(KnowledgeFSControlSpacePermission))) + assert {permission.account_id: permission.role for permission in permissions} == { + "owner-1": KnowledgeFSControlSpacePermissionRole.OWNER, + "member-1": KnowledgeFSControlSpacePermissionRole.VIEWER, + "member-2": KnowledgeFSControlSpacePermissionRole.VIEWER, + } + assert all(permission.granted_by_account_id == "owner-1" for permission in permissions) + assert sqlite_session.scalar(select(func.count()).select_from(KnowledgeFSLifecycleOutbox)) == 1 + assert sqlite_session.scalar(select(func.count()).select_from(KnowledgeFSAuthorizationRevision)) == 1 + assert sqlite_session.scalar(select(func.count()).select_from(KnowledgeFSExternalAccessPolicy)) == 1 + + +@pytest.mark.parametrize("member_ids", [(), ("owner-1",), ("member-1", "member-1"), ("outsider",), ("missing",)]) +def test_initial_member_validation_precedes_any_creation( + sqlite_session: Session, sqlite_session_factory: sessionmaker[Session], member_ids: tuple[str, ...] +) -> None: + _seed_members(sqlite_session) + service = KnowledgeFSControlSpaceCommandService(sqlite_session_factory) + + with pytest.raises(KnowledgeFSControlPlaneInvariantError): + service.create_provision_intent(_intent()._replace(member_account_ids=member_ids)) + + _assert_no_creation(sqlite_session) + + +def test_initial_members_roll_back_if_the_outbox_write_fails( + sqlite_session: Session, + sqlite_session_factory: sessionmaker[Session], + monkeypatch: pytest.MonkeyPatch, +) -> None: + _seed_members(sqlite_session) + original_add = SQLAlchemyKnowledgeFSLifecycleOutboxRepository.add + + def fail_after_flush( + repository: SQLAlchemyKnowledgeFSLifecycleOutboxRepository, command: KnowledgeFSLifecycleOutbox + ) -> KnowledgeFSLifecycleOutbox: + original_add(repository, command) + raise RuntimeError("outbox write interrupted") + + monkeypatch.setattr(SQLAlchemyKnowledgeFSLifecycleOutboxRepository, "add", fail_after_flush) + + with pytest.raises(RuntimeError, match="outbox write interrupted"): + KnowledgeFSControlSpaceCommandService(sqlite_session_factory).create_provision_intent(_intent()) + + _assert_no_creation(sqlite_session) + + +def test_creation_replay_does_not_restore_revoked_members_or_initial_visibility( + sqlite_session: Session, sqlite_session_factory: sessionmaker[Session] +) -> None: + _seed_members(sqlite_session) + service = KnowledgeFSControlSpaceCommandService(sqlite_session_factory) + result = service.create_provision_intent(_intent()) + space = sqlite_session.get(KnowledgeFSControlSpace, result.control_space.id) + assert space is not None + space.visibility = KnowledgeFSControlSpaceVisibility.ONLY_ME + member = sqlite_session.scalar( + select(KnowledgeFSControlSpacePermission).where(KnowledgeFSControlSpacePermission.account_id == "member-1") + ) + assert member is not None + member.status = KnowledgeFSControlSpacePermissionStatus.REVOKED + sqlite_session.execute(delete(TenantAccountJoin).where(TenantAccountJoin.account_id == "member-1")) + sqlite_session.commit() + + replay = service.create_provision_intent(_intent()) + + assert replay.control_space.visibility is KnowledgeFSControlSpaceVisibility.ONLY_ME + assert replay.outbox.id == result.outbox.id + sqlite_session.refresh(member) + assert member.status is KnowledgeFSControlSpacePermissionStatus.REVOKED + + +@pytest.mark.parametrize( + ("visibility", "member_account_ids"), + [ + (KnowledgeFSControlSpaceVisibility.PARTIAL_MEMBERS, ("member-1",)), + (KnowledgeFSControlSpaceVisibility.ALL_TEAM_MEMBERS, ()), + (KnowledgeFSControlSpaceVisibility.ONLY_ME, ()), + ], +) +def test_creation_replay_cannot_change_initial_access( + sqlite_session: Session, + sqlite_session_factory: sessionmaker[Session], + visibility: KnowledgeFSControlSpaceVisibility, + member_account_ids: tuple[str, ...], +) -> None: + _seed_members(sqlite_session) + service = KnowledgeFSControlSpaceCommandService(sqlite_session_factory) + service.create_provision_intent(_intent()) + + with pytest.raises(KnowledgeFSControlSpaceIntentConflictError, match="different initial access"): + service.create_provision_intent( + _intent()._replace(visibility=visibility, member_account_ids=member_account_ids) + ) + + assert sqlite_session.scalar(select(func.count()).select_from(KnowledgeFSControlSpacePermission)) == 3 + assert sqlite_session.scalar(select(func.count()).select_from(KnowledgeFSLifecycleOutbox)) == 1 + + +def test_legacy_creation_replay_keeps_current_access_and_rejects_added_members( + sqlite_session: Session, sqlite_session_factory: sessionmaker[Session] +) -> None: + _seed_members(sqlite_session) + service = KnowledgeFSControlSpaceCommandService(sqlite_session_factory) + private_intent = _intent()._replace(visibility=KnowledgeFSControlSpaceVisibility.ONLY_ME, member_account_ids=()) + result = service.create_provision_intent(private_intent) + command = sqlite_session.get(KnowledgeFSLifecycleOutbox, result.outbox.id) + assert command is not None + legacy_payload = _provision_payload(command).copy() + del legacy_payload["initial_visibility"] + del legacy_payload["initial_member_account_ids"] + command.command_payload = legacy_payload + sqlite_session.commit() + + replay = service.create_provision_intent( + private_intent._replace(visibility=KnowledgeFSControlSpaceVisibility.ALL_TEAM_MEMBERS) + ) + + assert replay.control_space.visibility is KnowledgeFSControlSpaceVisibility.ONLY_ME + assert replay.outbox.id == result.outbox.id + with pytest.raises(KnowledgeFSControlSpaceIntentConflictError): + service.create_provision_intent(_intent()) diff --git a/api/tests/unit_tests/services/test_knowledge_fs_control_space_lifecycle.py b/api/tests/unit_tests/services/knowledge_fs/test_control_space_lifecycle.py similarity index 100% rename from api/tests/unit_tests/services/test_knowledge_fs_control_space_lifecycle.py rename to api/tests/unit_tests/services/knowledge_fs/test_control_space_lifecycle.py diff --git a/api/tests/unit_tests/services/test_knowledge_fs_control_space_management.py b/api/tests/unit_tests/services/knowledge_fs/test_control_space_management.py similarity index 100% rename from api/tests/unit_tests/services/test_knowledge_fs_control_space_management.py rename to api/tests/unit_tests/services/knowledge_fs/test_control_space_management.py diff --git a/api/tests/unit_tests/services/test_knowledge_fs_cutover.py b/api/tests/unit_tests/services/knowledge_fs/test_cutover.py similarity index 100% rename from api/tests/unit_tests/services/test_knowledge_fs_cutover.py rename to api/tests/unit_tests/services/knowledge_fs/test_cutover.py diff --git a/api/tests/unit_tests/services/test_knowledge_fs_cutover_runtime_gate.py b/api/tests/unit_tests/services/knowledge_fs/test_cutover_runtime_gate.py similarity index 100% rename from api/tests/unit_tests/services/test_knowledge_fs_cutover_runtime_gate.py rename to api/tests/unit_tests/services/knowledge_fs/test_cutover_runtime_gate.py diff --git a/api/tests/unit_tests/services/test_knowledge_fs_data_facade.py b/api/tests/unit_tests/services/knowledge_fs/test_data_facade.py similarity index 93% rename from api/tests/unit_tests/services/test_knowledge_fs_data_facade.py rename to api/tests/unit_tests/services/knowledge_fs/test_data_facade.py index 6f2fd6dda98..e931a98d7ec 100644 --- a/api/tests/unit_tests/services/test_knowledge_fs_data_facade.py +++ b/api/tests/unit_tests/services/knowledge_fs/test_data_facade.py @@ -5,7 +5,7 @@ from collections.abc import Callable, Generator, Sequence from contextlib import contextmanager from dataclasses import dataclass, field from datetime import UTC, datetime, timedelta -from typing import Never, override +from typing import Literal, Never, override from unittest.mock import patch import pytest @@ -20,6 +20,9 @@ from services.knowledge_fs.product_dto import ( KnowledgeFSCatQuery, KnowledgeFSDiffQuery, KnowledgeFSDocumentDeletePayload, + KnowledgeFSDocumentResponse, + KnowledgeFSDocumentRevisionListResponse, + KnowledgeFSDocumentRevisionResponse, KnowledgeFSFindQuery, KnowledgeFSGoldenQuestionBulkImportPayload, KnowledgeFSGoldenQuestionBulkImportRowPayload, @@ -28,6 +31,7 @@ from services.knowledge_fs.product_dto import ( KnowledgeFSGrepQuery, KnowledgeFSListQuery, KnowledgeFSLogicalDocumentDeletePayload, + KnowledgeFSLogicalDocumentResponse, KnowledgeFSMetadataFieldCreatePayload, KnowledgeFSMetadataFieldUpdatePayload, KnowledgeFSProductRerankProfile, @@ -53,6 +57,7 @@ from services.knowledge_fs.product_remote import ( KnowledgeFSOperationUnavailableError, KnowledgeFSProductRemoteError, KnowledgeFSProductRequestRejectedError, + KnowledgeFSProductResourceNotFoundError, KnowledgeFSRemoteBinaryRequest, KnowledgeFSRemoteJSONRequest, KnowledgeFSRemoteMultipartFile, @@ -62,8 +67,8 @@ from services.knowledge_fs.product_remote import ( ) from services.knowledge_fs.query_images import KnowledgeFSQueryImagePreview from services.knowledge_fs.service_api_authorization import KnowledgeFSServiceApiProfile -from tests.unit_tests.services.knowledge_fs_fakes import UnexpectedCapabilityIssuer -from tests.unit_tests.services.knowledge_fs_services_a_fakes import CallRecorder, RecordedCall +from tests.helpers.unit.services.knowledge_fs_fakes import UnexpectedCapabilityIssuer +from tests.helpers.unit.services.knowledge_fs_services_a_fakes import CallRecorder, RecordedCall class _FacadeScope(TypedDict, closed=True): @@ -2902,3 +2907,226 @@ def test_service_settings_migration_preserves_service_identity_on_every_operatio "/knowledge-spaces/space-1/product-settings", "/knowledge-spaces/space-1/embedding-profile", ] + + +def logical_document( + *, + active: KnowledgeFSDocumentRevisionResponse | None, + status: Literal["deleting", "failed", "pending", "ready"] = "ready", +) -> KnowledgeFSLogicalDocumentResponse: + return KnowledgeFSLogicalDocumentResponse( + active=active, + active_revision=active.revision if active else None, + created_at=datetime.now(UTC), + enabled=True, + id="logical-1", + knowledge_space_id="space-1", + row_version=1, + status=status, + title="Document", + updated_at=datetime.now(UTC), + user_metadata={}, + ) + + +def revision( + *, state: Literal["active", "candidate", "failed", "superseded"] = "active" +) -> KnowledgeFSDocumentRevisionResponse: + return KnowledgeFSDocumentRevisionResponse( + activated_at=datetime.now(UTC) if state == "active" else None, + content_hash="a" * 64, + created_at=datetime.now(UTC), + document_asset_id="asset-1", + document_asset_version=2, + document_id="logical-1", + knowledge_space_id="space-1", + mime_type="text/markdown", + revision=3, + size_bytes=4, + state=state, + ) + + +def asset() -> KnowledgeFSDocumentResponse: + return KnowledgeFSDocumentResponse( + id="asset-1", + knowledge_space_id="space-1", + filename="page.md", + metadata={}, + mime_type="text/markdown", + object_key="tenant/spaces/space-1/documents/page.md", + parser_status="parsed", + sha256="a" * 64, + size_bytes=4, + version=2, + created_at=datetime.now(UTC), + ) + + +class FakeDownloadFacade(KnowledgeFSDataFacade): + """Serves canned facade reads and records the calls download preparation makes.""" + + def __init__( + self, + *, + logical: KnowledgeFSLogicalDocumentResponse, + revisions: KnowledgeFSDocumentRevisionListResponse | None = None, + document: KnowledgeFSDocumentResponse | None = None, + ) -> None: + self._logical = logical + self._revisions = revisions + self._document = document + self.get_logical_document_calls: list[dict[str, object]] = [] + self.list_document_revisions_calls: list[dict[str, object]] = [] + self.get_document_calls: list[dict[str, object]] = [] + + @override + def get_logical_document( + self, *, tenant_id: str, account_id: str, control_space_id: str, document_id: str + ) -> KnowledgeFSLogicalDocumentResponse: + self.get_logical_document_calls.append( + { + "tenant_id": tenant_id, + "account_id": account_id, + "control_space_id": control_space_id, + "document_id": document_id, + } + ) + return self._logical + + @override + def list_document_revisions( + self, + *, + tenant_id: str, + account_id: str, + control_space_id: str, + document_id: str, + cursor: str | None = None, + limit: int | None = None, + ) -> KnowledgeFSDocumentRevisionListResponse: + self.list_document_revisions_calls.append( + { + "tenant_id": tenant_id, + "account_id": account_id, + "control_space_id": control_space_id, + "document_id": document_id, + "cursor": cursor, + "limit": limit, + } + ) + assert self._revisions is not None + return self._revisions + + @override + def get_document( + self, *, tenant_id: str, account_id: str, control_space_id: str, document_id: str + ) -> KnowledgeFSDocumentResponse: + self.get_document_calls.append( + { + "tenant_id": tenant_id, + "account_id": account_id, + "control_space_id": control_space_id, + "document_id": document_id, + } + ) + assert self._document is not None + return self._document + + +def test_prepare_logical_document_download_resolves_active_asset() -> None: + facade = FakeDownloadFacade(logical=logical_document(active=revision()), document=asset()) + + result = facade.prepare_logical_document_download( + tenant_id="tenant-1", + account_id="account-1", + control_space_id="control-1", + document_id="logical-1", + ) + + assert result.document_id == "logical-1" + assert result.object_key == "tenant/spaces/space-1/documents/page.md" + assert facade.list_document_revisions_calls == [] + assert facade.get_document_calls == [ + { + "tenant_id": "tenant-1", + "account_id": "account-1", + "control_space_id": "control-1", + "document_id": "asset-1", + } + ] + + +def test_prepare_logical_document_download_resolves_latest_failed_revision_without_active_revision() -> None: + facade = FakeDownloadFacade( + logical=logical_document(active=None, status="failed"), + revisions=KnowledgeFSDocumentRevisionListResponse(data=[revision(state="failed")]), + document=asset(), + ) + + result = facade.prepare_logical_document_download( + tenant_id="tenant-1", + account_id="account-1", + control_space_id="control-1", + document_id="logical-1", + ) + + assert result.document_id == "logical-1" + assert result.object_key == "tenant/spaces/space-1/documents/page.md" + assert facade.list_document_revisions_calls == [ + { + "tenant_id": "tenant-1", + "account_id": "account-1", + "control_space_id": "control-1", + "document_id": "logical-1", + "cursor": None, + "limit": 1, + } + ] + + +def test_prepare_logical_document_download_rejects_document_without_visible_revision() -> None: + facade = FakeDownloadFacade( + logical=logical_document(active=None, status="failed"), + revisions=KnowledgeFSDocumentRevisionListResponse(data=[]), + ) + + with pytest.raises(KnowledgeFSProductResourceNotFoundError): + facade.prepare_logical_document_download( + tenant_id="tenant-1", + account_id="account-1", + control_space_id="control-1", + document_id="logical-1", + ) + + +def test_prepare_logical_document_download_rejects_latest_candidate_revision() -> None: + facade = FakeDownloadFacade( + logical=logical_document(active=None, status="failed"), + revisions=KnowledgeFSDocumentRevisionListResponse(data=[revision(state="candidate")]), + ) + + with pytest.raises(KnowledgeFSProductResourceNotFoundError): + facade.prepare_logical_document_download( + tenant_id="tenant-1", + account_id="account-1", + control_space_id="control-1", + document_id="logical-1", + ) + + assert facade.get_document_calls == [] + + +def test_prepare_logical_document_download_rejects_pending_document_without_active_revision() -> None: + facade = FakeDownloadFacade(logical=logical_document(active=None, status="pending")) + + with pytest.raises(KnowledgeFSProductResourceNotFoundError): + facade.prepare_logical_document_download( + tenant_id="tenant-1", + account_id="account-1", + control_space_id="control-1", + document_id="logical-1", + ) + + assert facade.list_document_revisions_calls == [] + assert facade.get_document_calls == [] diff --git a/api/tests/unit_tests/services/test_knowledge_fs_download_service.py b/api/tests/unit_tests/services/knowledge_fs/test_download_service.py similarity index 100% rename from api/tests/unit_tests/services/test_knowledge_fs_download_service.py rename to api/tests/unit_tests/services/knowledge_fs/test_download_service.py diff --git a/api/tests/unit_tests/services/test_knowledge_fs_feature_service.py b/api/tests/unit_tests/services/knowledge_fs/test_feature_service.py similarity index 100% rename from api/tests/unit_tests/services/test_knowledge_fs_feature_service.py rename to api/tests/unit_tests/services/knowledge_fs/test_feature_service.py diff --git a/api/tests/unit_tests/services/test_knowledge_fs_greenfield_initializer.py b/api/tests/unit_tests/services/knowledge_fs/test_greenfield_initializer.py similarity index 100% rename from api/tests/unit_tests/services/test_knowledge_fs_greenfield_initializer.py rename to api/tests/unit_tests/services/knowledge_fs/test_greenfield_initializer.py diff --git a/api/tests/unit_tests/services/test_knowledge_fs_initial_source_import.py b/api/tests/unit_tests/services/knowledge_fs/test_initial_source_import.py similarity index 86% rename from api/tests/unit_tests/services/test_knowledge_fs_initial_source_import.py rename to api/tests/unit_tests/services/knowledge_fs/test_initial_source_import.py index e2c2178e422..91625997345 100644 --- a/api/tests/unit_tests/services/test_knowledge_fs_initial_source_import.py +++ b/api/tests/unit_tests/services/knowledge_fs/test_initial_source_import.py @@ -4,6 +4,7 @@ from dataclasses import dataclass, field from unittest.mock import patch import pytest +from pydantic import ValidationError from core.db.session_factory import session_factory from models.account import Account @@ -31,7 +32,7 @@ from services.knowledge_fs.product_remote import ( KnowledgeFSProductRequestRejectedError, KnowledgeFSProductResourceNotFoundError, ) -from tests.unit_tests.tasks.knowledge_fs_task_fakes import ( +from tests.helpers.unit.tasks.knowledge_fs_task_fakes import ( FacadeRuntime, RecordingKnowledgeFSFacade, json_object, @@ -45,7 +46,7 @@ from tests.unit_tests.tasks.knowledge_fs_task_fakes import ( ) -def _payload(sync_policy: str = "daily") -> KnowledgeFSInitialWebsiteSourcePayload: +def _knowledge_fs_initial_source_import_payload(sync_policy: str = "daily") -> KnowledgeFSInitialWebsiteSourcePayload: return KnowledgeFSInitialWebsiteSourcePayload.model_validate( { "kind": "website_crawl", @@ -288,7 +289,7 @@ def _start( def test_initial_website_source_import_recrawls_exact_selection_and_configures_daily_sync() -> None: facade = _facade() - assert _start(facade, _payload()) == "workflow-1" + assert _start(facade, _knowledge_fs_initial_source_import_payload()) == "workflow-1" source_payload = facade.payload("create_source", KnowledgeFSSourceCreatePayload) assert source_payload.status == "disabled" @@ -490,7 +491,7 @@ def test_initial_website_source_import_reuses_source_across_pages_and_preserves_ ), ) - assert _start(facade, _payload()) == "failed-workflow" + assert _start(facade, _knowledge_fs_initial_source_import_payload()) == "failed-workflow" assert facade.call_count("create_source") == 0 assert facade.call_count("list_source_connections") == 0 source_update_payload = facade.payload("update_source", KnowledgeFSSourceUpdatePayload) @@ -510,14 +511,14 @@ def test_initial_website_source_import_reuses_source_across_pages_and_preserves_ def test_initial_website_source_import_exposes_preview_consume_rejection() -> None: facade = _facade() facade.script("consume_namespace_source_preview", KnowledgeFSProductRequestRejectedError(status_code=409)) - payload = _payload().model_copy( + payload = _knowledge_fs_initial_source_import_payload().model_copy( update={ "credential_id": "firecrawl-credential-1", "plugin_id": "langgenius/firecrawl_datasource", "preview_job_id": "preview-job-1", "selection": [ selection.model_copy(update={"page_id": f"page-{index}"}) - for index, selection in enumerate(_payload().selection, start=1) + for index, selection in enumerate(_knowledge_fs_initial_source_import_payload().selection, start=1) ], } ) @@ -573,7 +574,7 @@ def test_initial_website_source_import_configures_remaining_sync_modes( ) facade.script("get_source_sync_policy", sync_policy_response(revision=7)) - assert _start(facade, _payload(sync_policy)) == "workflow-1" + assert _start(facade, _knowledge_fs_initial_source_import_payload(sync_policy)) == "workflow-1" sync_payload = facade.payload("update_source_sync_policy", KnowledgeFSSourceSyncPolicyPayload) assert sync_payload.mode == expected_mode assert sync_payload.enabled is expected_enabled @@ -596,7 +597,7 @@ def test_initial_website_source_import_does_not_recommit_active_source_after_ret ), ) - assert _start(facade, _payload()) == "workflow-1" + assert _start(facade, _knowledge_fs_initial_source_import_payload()) == "workflow-1" assert facade.call_count("update_source") == 0 sync_payload = facade.payload("update_source_sync_policy", KnowledgeFSSourceSyncPolicyPayload) @@ -642,7 +643,7 @@ def test_initial_website_source_import_rejects_unavailable_spaces( account_id="account-1", control_space_id="control-1", operation_id="operation-1", - payload=_payload(), + payload=_knowledge_fs_initial_source_import_payload(), ) @@ -655,7 +656,7 @@ def test_initial_website_source_import_rejects_missing_control_space() -> None: account_id="account-1", control_space_id="control-1", operation_id="operation-1", - payload=_payload(), + payload=_knowledge_fs_initial_source_import_payload(), ) @@ -663,7 +664,7 @@ def test_initial_website_source_import_rejects_unavailable_provider() -> None: unavailable_provider = _facade() unavailable_provider.script("list_source_providers", source_provider_list_response()) with pytest.raises(RuntimeError, match="provider is unavailable"): - _start(unavailable_provider, _payload()) + _start(unavailable_provider, _knowledge_fs_initial_source_import_payload()) def test_initial_website_source_import_rejects_missing_firecrawl_credential() -> None: @@ -677,7 +678,7 @@ def test_initial_website_source_import_rejects_missing_firecrawl_credential() -> account_id="account-1", control_space_id="control-1", operation_id="operation-1", - payload=_payload(), + payload=_knowledge_fs_initial_source_import_payload(), ) assert facade.call_count("create_source_connection") == 0 @@ -693,7 +694,7 @@ def test_initial_website_source_import_creates_missing_firecrawl_connection() -> ), ) - assert _start(facade, _payload()) == "workflow-1" + assert _start(facade, _knowledge_fs_initial_source_import_payload()) == "workflow-1" connection_payload = facade.payload("create_source_connection", KnowledgeFSSourceConnectionCreatePayload) assert connection_payload.auth_kind == "endpoint" @@ -725,7 +726,7 @@ def test_initial_website_source_import_retries_provisioning_firecrawl_connection ) with pytest.raises(KnowledgeFSInitialSourceNotReadyError, match="still provisioning"): - _start(facade, _payload()) + _start(facade, _knowledge_fs_initial_source_import_payload()) assert facade.call_count("create_source_connection") == 0 assert facade.call_count("create_source") == 0 @@ -739,7 +740,7 @@ def test_initial_website_source_import_retries_running_workflow() -> None: ) with pytest.raises(KnowledgeFSInitialSourceNotReadyError, match="still running") as raised: - _start(facade, _payload()) + _start(facade, _knowledge_fs_initial_source_import_payload()) assert raised.value.workflow_id == "running-workflow" @@ -747,7 +748,7 @@ def test_initial_website_source_import_retries_running_workflow() -> None: def test_initial_website_source_import_polls_existing_workflow_without_recreating_it() -> None: facade = _facade() - assert _start(facade, _payload(), workflow_id="workflow-1") == "workflow-1" + assert _start(facade, _knowledge_fs_initial_source_import_payload(), workflow_id="workflow-1") == "workflow-1" assert facade.calls["get_source_workflow"] == [ { @@ -773,7 +774,7 @@ def test_initial_website_source_import_keeps_polling_existing_running_workflow() ) with pytest.raises(KnowledgeFSInitialSourceNotReadyError, match="still running") as raised: - _start(facade, _payload(), workflow_id="workflow-1") + _start(facade, _knowledge_fs_initial_source_import_payload(), workflow_id="workflow-1") assert raised.value.workflow_id == "workflow-1" assert facade.call_count("import_selected_source_crawl") == 0 @@ -793,7 +794,7 @@ def test_initial_website_source_import_exposes_failed_source_without_activating_ ), ) - assert _start(facade, _payload(), workflow_id="workflow-1") == "workflow-1" + assert _start(facade, _knowledge_fs_initial_source_import_payload(), workflow_id="workflow-1") == "workflow-1" source_update_payload = facade.payload("update_source", KnowledgeFSSourceUpdatePayload) source_update_metadata = json_object(source_update_payload.metadata) @@ -845,7 +846,7 @@ def test_upgrade_source_submission_succeeds_when_first_import_submission_fails() account_id="account-1", control_space_id="control-1", operation_id="operation-1", - payload=_payload(sync_policy="manual"), + payload=_knowledge_fs_initial_source_import_payload(sync_policy="manual"), ) assert result.connection_id == "connection-1" @@ -853,3 +854,76 @@ def test_upgrade_source_submission_succeeds_when_first_import_submission_fails() assert result.workflow_id is None assert result.workflow_error == "RuntimeError" assert facade.call_count("update_source_sync_policy") == 1 + + +def _knowledge_fs_initial_source_reliability_payload( + *, limit: int = 2, urls: list[str] | None = None +) -> dict[str, object]: + selected = urls or ["HTTPS://Docs.Dify.AI:443/a/#section", "https://docs.dify.ai/b/"] + return { + "kind": "website_crawl", + "name": "Dify docs", + "provider": "firecrawl", + "root_url": "https://docs.dify.ai", + "crawl_options": {"include_subpages": True, "limit": limit}, + "selection": [{"source_url": url} for url in selected], + } + + +def test_initial_website_selection_populates_stable_canonical_urls() -> None: + payload = KnowledgeFSInitialWebsiteSourcePayload.model_validate(_knowledge_fs_initial_source_reliability_payload()) + + assert [item.canonical_url for item in payload.selection] == [ + "https://docs.dify.ai/a", + "https://docs.dify.ai/b", + ] + + +def test_initial_website_selection_preserves_canonical_duplicates() -> None: + payload = KnowledgeFSInitialWebsiteSourcePayload.model_validate( + _knowledge_fs_initial_source_reliability_payload( + urls=["https://docs.dify.ai/a/", "https://DOCS.dify.ai:443/a#section"] + ) + ) + + assert [item.source_url for item in payload.selection] == [ + "https://docs.dify.ai/a/", + "https://DOCS.dify.ai:443/a#section", + ] + assert [item.canonical_url for item in payload.selection] == [ + "https://docs.dify.ai/a", + "https://docs.dify.ai/a", + ] + + +def test_crawl_import_preserves_duplicate_source_urls() -> None: + payload = KnowledgeFSCrawlImportPayload.model_validate( + {"sourceUrls": [" https://docs.dify.ai/a ", "https://docs.dify.ai/a"]} + ) + + assert payload.source_urls == ["https://docs.dify.ai/a", "https://docs.dify.ai/a"] + + +def test_initial_website_selection_requires_crawl_limit_to_cover_selection() -> None: + with pytest.raises(ValidationError, match="crawl limit must cover every selected URL"): + KnowledgeFSInitialWebsiteSourcePayload.model_validate(_knowledge_fs_initial_source_reliability_payload(limit=1)) + + +def test_initial_website_source_rejects_stale_preview_configuration() -> None: + payload = _knowledge_fs_initial_source_reliability_payload() + payload["previewConfigurationFingerprint"] = "0" * 64 + + with pytest.raises(ValidationError, match="preview configuration no longer matches"): + KnowledgeFSInitialWebsiteSourcePayload.model_validate(payload) + + +def test_initial_website_source_accepts_matching_preview_configuration() -> None: + payload = KnowledgeFSInitialWebsiteSourcePayload.model_validate(_knowledge_fs_initial_source_reliability_payload()) + raw = _knowledge_fs_initial_source_reliability_payload() + raw["previewConfigurationFingerprint"] = knowledge_fs_initial_preview_configuration_fingerprint(payload) + + validated = KnowledgeFSInitialWebsiteSourcePayload.model_validate(raw) + + assert validated.preview_configuration_fingerprint == knowledge_fs_initial_preview_configuration_fingerprint( + validated + ) diff --git a/api/tests/unit_tests/services/test_knowledge_fs_initial_source_preview.py b/api/tests/unit_tests/services/knowledge_fs/test_initial_source_preview.py similarity index 100% rename from api/tests/unit_tests/services/test_knowledge_fs_initial_source_preview.py rename to api/tests/unit_tests/services/knowledge_fs/test_initial_source_preview.py diff --git a/api/tests/unit_tests/services/test_knowledge_fs_initial_source_preview_job.py b/api/tests/unit_tests/services/knowledge_fs/test_initial_source_preview_job.py similarity index 100% rename from api/tests/unit_tests/services/test_knowledge_fs_initial_source_preview_job.py rename to api/tests/unit_tests/services/knowledge_fs/test_initial_source_preview_job.py diff --git a/api/tests/unit_tests/services/test_knowledge_fs_lifecycle_readiness.py b/api/tests/unit_tests/services/knowledge_fs/test_lifecycle_readiness.py similarity index 100% rename from api/tests/unit_tests/services/test_knowledge_fs_lifecycle_readiness.py rename to api/tests/unit_tests/services/knowledge_fs/test_lifecycle_readiness.py diff --git a/api/tests/unit_tests/services/test_knowledge_fs_lifecycle_remote_http.py b/api/tests/unit_tests/services/knowledge_fs/test_lifecycle_remote_http.py similarity index 100% rename from api/tests/unit_tests/services/test_knowledge_fs_lifecycle_remote_http.py rename to api/tests/unit_tests/services/knowledge_fs/test_lifecycle_remote_http.py diff --git a/api/tests/unit_tests/services/test_knowledge_fs_lifecycle_saga.py b/api/tests/unit_tests/services/knowledge_fs/test_lifecycle_saga.py similarity index 99% rename from api/tests/unit_tests/services/test_knowledge_fs_lifecycle_saga.py rename to api/tests/unit_tests/services/knowledge_fs/test_lifecycle_saga.py index db535ee7f20..44f5d601c17 100644 --- a/api/tests/unit_tests/services/test_knowledge_fs_lifecycle_saga.py +++ b/api/tests/unit_tests/services/knowledge_fs/test_lifecycle_saga.py @@ -30,8 +30,8 @@ from services.knowledge_fs.lifecycle_port import ( ) from services.knowledge_fs.lifecycle_saga import KnowledgeFSLifecycleSagaRunner from services.knowledge_fs.observability import KnowledgeFSLifecycleTaskMetric -from tests.unit_tests.knowledge_fs_core_fakes import RecordingOperationalMetrics -from tests.unit_tests.services.knowledge_fs_fakes import UnexpectedLifecycleRemote +from tests.helpers.unit.knowledge_fs_core_fakes import RecordingOperationalMetrics +from tests.helpers.unit.services.knowledge_fs_fakes import UnexpectedLifecycleRemote class FakeKnowledgeFSLifecycleRemote(UnexpectedLifecycleRemote): diff --git a/api/tests/unit_tests/services/test_knowledge_fs_membership_changes.py b/api/tests/unit_tests/services/knowledge_fs/test_membership_changes.py similarity index 99% rename from api/tests/unit_tests/services/test_knowledge_fs_membership_changes.py rename to api/tests/unit_tests/services/knowledge_fs/test_membership_changes.py index b73f17a81cd..a09ad52c270 100644 --- a/api/tests/unit_tests/services/test_knowledge_fs_membership_changes.py +++ b/api/tests/unit_tests/services/knowledge_fs/test_membership_changes.py @@ -23,7 +23,7 @@ from services.knowledge_fs.membership_changes import ( apply_workspace_rbac_role_change, ) from services.knowledge_fs.revocation_commands import KnowledgeFSRevocationCommandError -from tests.unit_tests.services.knowledge_fs_fakes import claims_summary, revoke_payload +from tests.helpers.unit.services.knowledge_fs_fakes import claims_summary, revoke_payload @pytest.mark.parametrize( diff --git a/api/tests/unit_tests/services/test_knowledge_fs_object_storage.py b/api/tests/unit_tests/services/knowledge_fs/test_object_storage.py similarity index 100% rename from api/tests/unit_tests/services/test_knowledge_fs_object_storage.py rename to api/tests/unit_tests/services/knowledge_fs/test_object_storage.py diff --git a/api/tests/unit_tests/services/test_knowledge_fs_observability.py b/api/tests/unit_tests/services/knowledge_fs/test_observability.py similarity index 98% rename from api/tests/unit_tests/services/test_knowledge_fs_observability.py rename to api/tests/unit_tests/services/knowledge_fs/test_observability.py index f4c51a9be43..f1727d9f386 100644 --- a/api/tests/unit_tests/services/test_knowledge_fs_observability.py +++ b/api/tests/unit_tests/services/knowledge_fs/test_observability.py @@ -12,7 +12,7 @@ from services.knowledge_fs.observability import ( KnowledgeFSLifecycleTaskMetric, OpenTelemetryKnowledgeFSOperationalMetrics, ) -from tests.unit_tests.knowledge_fs_core_fakes import RecordedCall +from tests.helpers.unit.knowledge_fs_core_fakes import RecordedCall @dataclass diff --git a/api/tests/unit_tests/services/test_knowledge_fs_orphan_reconciler.py b/api/tests/unit_tests/services/knowledge_fs/test_orphan_reconciler.py similarity index 98% rename from api/tests/unit_tests/services/test_knowledge_fs_orphan_reconciler.py rename to api/tests/unit_tests/services/knowledge_fs/test_orphan_reconciler.py index 421034513d0..e9a02c1d304 100644 --- a/api/tests/unit_tests/services/test_knowledge_fs_orphan_reconciler.py +++ b/api/tests/unit_tests/services/knowledge_fs/test_orphan_reconciler.py @@ -15,7 +15,7 @@ from services.knowledge_fs.lifecycle_port import ( KnowledgeFSRemoteSpace, ) from services.knowledge_fs.orphan_reconciler import KnowledgeFSOrphanReconciler -from tests.unit_tests.services.knowledge_fs_fakes import UnexpectedLifecycleRemote +from tests.helpers.unit.services.knowledge_fs_fakes import UnexpectedLifecycleRemote class FakeRemote(UnexpectedLifecycleRemote): diff --git a/api/tests/unit_tests/services/test_knowledge_fs_product_application_service.py b/api/tests/unit_tests/services/knowledge_fs/test_product_application_service.py similarity index 99% rename from api/tests/unit_tests/services/test_knowledge_fs_product_application_service.py rename to api/tests/unit_tests/services/knowledge_fs/test_product_application_service.py index 2e595ec6505..4913a3cc989 100644 --- a/api/tests/unit_tests/services/test_knowledge_fs_product_application_service.py +++ b/api/tests/unit_tests/services/knowledge_fs/test_product_application_service.py @@ -30,7 +30,7 @@ from services.knowledge_fs.product_remote import ( KnowledgeFSProductRequestRejectedError, ) from services.knowledge_fs.product_service import KnowledgeFSProductService -from tests.unit_tests.config_override import apply_config_overrides +from tests.helpers.unit.config_override import apply_config_overrides class RejectingProduct: diff --git a/api/tests/unit_tests/services/test_knowledge_fs_product_authorization.py b/api/tests/unit_tests/services/knowledge_fs/test_product_authorization.py similarity index 98% rename from api/tests/unit_tests/services/test_knowledge_fs_product_authorization.py rename to api/tests/unit_tests/services/knowledge_fs/test_product_authorization.py index c51ae7ead6c..adc8cc23cf9 100644 --- a/api/tests/unit_tests/services/test_knowledge_fs_product_authorization.py +++ b/api/tests/unit_tests/services/knowledge_fs/test_product_authorization.py @@ -13,7 +13,7 @@ from services.knowledge_fs.product_operations import ( KnowledgeFSProductPermission, product_permissions_from_rbac_keys, ) -from tests.unit_tests.config_override import config_overrides_context +from tests.helpers.unit.config_override import config_overrides_context def test_enterprise_knowledge_fs_permissions_use_one_batch_request() -> None: diff --git a/api/tests/unit_tests/services/test_knowledge_fs_product_dto.py b/api/tests/unit_tests/services/knowledge_fs/test_product_dto.py similarity index 100% rename from api/tests/unit_tests/services/test_knowledge_fs_product_dto.py rename to api/tests/unit_tests/services/knowledge_fs/test_product_dto.py diff --git a/api/tests/unit_tests/services/test_knowledge_fs_product_operations.py b/api/tests/unit_tests/services/knowledge_fs/test_product_operations.py similarity index 99% rename from api/tests/unit_tests/services/test_knowledge_fs_product_operations.py rename to api/tests/unit_tests/services/knowledge_fs/test_product_operations.py index 1068c879b2a..062d4071e5d 100644 --- a/api/tests/unit_tests/services/test_knowledge_fs_product_operations.py +++ b/api/tests/unit_tests/services/knowledge_fs/test_product_operations.py @@ -196,7 +196,7 @@ def test_source_workflow_import_byte_limit_accepts_the_full_schema_bounded_batch def test_manifest_gaps_remain_explicit_and_stable() -> None: expected = () assert knowledge_fs_product_operation_gaps() == expected - manifest_path = Path(__file__).resolve().parents[3] / "knowledge-fs-product-operation-gaps.json" + manifest_path = Path(__file__).resolve().parents[4] / "knowledge-fs-product-operation-gaps.json" manifest = json.loads(manifest_path.read_text(encoding="utf-8")) assert manifest["schemaVersion"] == 2 assert tuple(item["productOperationId"] for item in manifest["gaps"]) == expected diff --git a/api/tests/unit_tests/services/test_knowledge_fs_product_remote_http.py b/api/tests/unit_tests/services/knowledge_fs/test_product_remote_http.py similarity index 100% rename from api/tests/unit_tests/services/test_knowledge_fs_product_remote_http.py rename to api/tests/unit_tests/services/knowledge_fs/test_product_remote_http.py diff --git a/api/tests/unit_tests/services/test_knowledge_fs_product_service.py b/api/tests/unit_tests/services/knowledge_fs/test_product_service.py similarity index 99% rename from api/tests/unit_tests/services/test_knowledge_fs_product_service.py rename to api/tests/unit_tests/services/knowledge_fs/test_product_service.py index 6c7c78c6233..6733554409f 100644 --- a/api/tests/unit_tests/services/test_knowledge_fs_product_service.py +++ b/api/tests/unit_tests/services/knowledge_fs/test_product_service.py @@ -44,7 +44,7 @@ from services.knowledge_fs.product_remote import ( ) from services.knowledge_fs.product_service import KnowledgeFSProductService from services.knowledge_fs_capability import CapabilityIssueRequest -from tests.unit_tests.knowledge_fs_core_fakes import RecordingOperationalMetrics +from tests.helpers.unit.knowledge_fs_core_fakes import RecordingOperationalMetrics class FakeRBAC: diff --git a/api/tests/unit_tests/services/test_knowledge_fs_query_images.py b/api/tests/unit_tests/services/knowledge_fs/test_query_images.py similarity index 100% rename from api/tests/unit_tests/services/test_knowledge_fs_query_images.py rename to api/tests/unit_tests/services/knowledge_fs/test_query_images.py diff --git a/api/tests/unit_tests/services/test_knowledge_fs_remote_images.py b/api/tests/unit_tests/services/knowledge_fs/test_remote_images.py similarity index 100% rename from api/tests/unit_tests/services/test_knowledge_fs_remote_images.py rename to api/tests/unit_tests/services/knowledge_fs/test_remote_images.py diff --git a/api/tests/unit_tests/services/test_knowledge_fs_remote_registry.py b/api/tests/unit_tests/services/knowledge_fs/test_remote_registry.py similarity index 92% rename from api/tests/unit_tests/services/test_knowledge_fs_remote_registry.py rename to api/tests/unit_tests/services/knowledge_fs/test_remote_registry.py index 9b19c68945e..28868489bf5 100644 --- a/api/tests/unit_tests/services/test_knowledge_fs_remote_registry.py +++ b/api/tests/unit_tests/services/knowledge_fs/test_remote_registry.py @@ -7,9 +7,9 @@ from core.db.session_factory import session_factory from repositories import sqlalchemy_knowledge_fs_capability_issuance_auditor from services import knowledge_fs_capability from services.knowledge_fs import lifecycle_remote_http, remote_registry -from tests.unit_tests.config_override import apply_config_overrides -from tests.unit_tests.knowledge_fs_core_fakes import Marker -from tests.unit_tests.services.knowledge_fs_fakes import UnexpectedLifecycleRemote +from tests.helpers.unit.config_override import apply_config_overrides +from tests.helpers.unit.knowledge_fs_core_fakes import Marker +from tests.helpers.unit.services.knowledge_fs_fakes import UnexpectedLifecycleRemote def test_configured_lifecycle_remote_assembles_capability_v2_http_client( diff --git a/api/tests/unit_tests/services/test_knowledge_fs_revocation.py b/api/tests/unit_tests/services/knowledge_fs/test_revocation_commands.py similarity index 99% rename from api/tests/unit_tests/services/test_knowledge_fs_revocation.py rename to api/tests/unit_tests/services/knowledge_fs/test_revocation_commands.py index a7090a0e546..4dce3893ff7 100644 --- a/api/tests/unit_tests/services/test_knowledge_fs_revocation.py +++ b/api/tests/unit_tests/services/knowledge_fs/test_revocation_commands.py @@ -43,7 +43,7 @@ from services.knowledge_fs.revocation_commands import ( KnowledgeFSRevocationCommandProducer, ) from services.knowledge_fs.revocation_reconciler import KnowledgeFSRevocationReconciler -from tests.unit_tests.services.knowledge_fs_fakes import UnexpectedLifecycleRemote, revoke_payload +from tests.helpers.unit.services.knowledge_fs_fakes import UnexpectedLifecycleRemote, revoke_payload _SPACE_ID = "10000000-0000-4000-8000-000000000001" _GRANT_A = "20000000-0000-4000-8000-000000000001" diff --git a/api/tests/unit_tests/services/test_knowledge_fs_runtime.py b/api/tests/unit_tests/services/knowledge_fs/test_runtime.py similarity index 97% rename from api/tests/unit_tests/services/test_knowledge_fs_runtime.py rename to api/tests/unit_tests/services/knowledge_fs/test_runtime.py index a9e1c6a1551..e9422a0e647 100644 --- a/api/tests/unit_tests/services/test_knowledge_fs_runtime.py +++ b/api/tests/unit_tests/services/knowledge_fs/test_runtime.py @@ -5,8 +5,8 @@ from sqlalchemy.orm import sessionmaker from services.knowledge_fs import runtime from services.knowledge_fs.product_remote import KnowledgeFSOperationUnavailableError -from tests.unit_tests.config_override import apply_config_overrides -from tests.unit_tests.knowledge_fs_core_fakes import Marker, RecordingCallable +from tests.helpers.unit.config_override import apply_config_overrides +from tests.helpers.unit.knowledge_fs_core_fakes import Marker, RecordingCallable def test_runtime_fails_closed_when_the_remote_endpoint_is_not_configured(monkeypatch: pytest.MonkeyPatch) -> None: diff --git a/api/tests/unit_tests/services/test_knowledge_fs_service_api_authorization.py b/api/tests/unit_tests/services/knowledge_fs/test_service_api_authorization.py similarity index 100% rename from api/tests/unit_tests/services/test_knowledge_fs_service_api_authorization.py rename to api/tests/unit_tests/services/knowledge_fs/test_service_api_authorization.py diff --git a/api/tests/unit_tests/services/test_knowledge_fs_service_query_images.py b/api/tests/unit_tests/services/knowledge_fs/test_service_query_image_upload.py similarity index 99% rename from api/tests/unit_tests/services/test_knowledge_fs_service_query_images.py rename to api/tests/unit_tests/services/knowledge_fs/test_service_query_image_upload.py index 0987ca85c91..2e002dacc3b 100644 --- a/api/tests/unit_tests/services/test_knowledge_fs_service_query_images.py +++ b/api/tests/unit_tests/services/knowledge_fs/test_service_query_image_upload.py @@ -14,7 +14,7 @@ from models.model import UploadFile from services.knowledge_fs import query_images, service_query_image_upload from services.knowledge_fs.errors import KnowledgeFSProductRequestRejectedError, KnowledgeFSQueryImageError from services.knowledge_fs.service_api_authorization import KnowledgeFSServiceApiProfile -from tests.unit_tests.config_override import apply_config_overrides +from tests.helpers.unit.config_override import apply_config_overrides TOKEN_ID = "00000000-0000-4000-8000-000000000001" OTHER_TOKEN_ID = "00000000-0000-4000-8000-000000000002" diff --git a/api/tests/unit_tests/services/test_knowledge_fs_source_import_commit_service.py b/api/tests/unit_tests/services/knowledge_fs/test_source_import_commit_service.py similarity index 99% rename from api/tests/unit_tests/services/test_knowledge_fs_source_import_commit_service.py rename to api/tests/unit_tests/services/knowledge_fs/test_source_import_commit_service.py index 9d77f1c4029..76021b9f706 100644 --- a/api/tests/unit_tests/services/test_knowledge_fs_source_import_commit_service.py +++ b/api/tests/unit_tests/services/knowledge_fs/test_source_import_commit_service.py @@ -20,7 +20,7 @@ from services.knowledge_fs.source_import_commit_service import ( resume_committed_source_import, retry_or_resume_source_workflow, ) -from tests.unit_tests.services.knowledge_fs_services_a_fakes import CallRecorder, sent_payload +from tests.helpers.unit.services.knowledge_fs_services_a_fakes import CallRecorder, sent_payload def source_workflow(**overrides: object) -> KnowledgeFSSourceWorkflowResponse: diff --git a/api/tests/unit_tests/services/test_knowledge_fs_space_tag_service.py b/api/tests/unit_tests/services/knowledge_fs/test_space_tag_service.py similarity index 100% rename from api/tests/unit_tests/services/test_knowledge_fs_space_tag_service.py rename to api/tests/unit_tests/services/knowledge_fs/test_space_tag_service.py diff --git a/api/tests/unit_tests/services/test_knowledge_fs_staged_upload_service.py b/api/tests/unit_tests/services/knowledge_fs/test_staged_upload_service.py similarity index 100% rename from api/tests/unit_tests/services/test_knowledge_fs_staged_upload_service.py rename to api/tests/unit_tests/services/knowledge_fs/test_staged_upload_service.py diff --git a/api/tests/unit_tests/services/test_knowledge_fs_text_search_deadline.py b/api/tests/unit_tests/services/knowledge_fs/test_text_search_deadline.py similarity index 100% rename from api/tests/unit_tests/services/test_knowledge_fs_text_search_deadline.py rename to api/tests/unit_tests/services/knowledge_fs/test_text_search_deadline.py diff --git a/api/tests/unit_tests/services/test_knowledge_fs_text_store.py b/api/tests/unit_tests/services/knowledge_fs/test_text_store.py similarity index 100% rename from api/tests/unit_tests/services/test_knowledge_fs_text_store.py rename to api/tests/unit_tests/services/knowledge_fs/test_text_store.py diff --git a/api/tests/unit_tests/services/test_knowledge_fs_text_store_diagnostics.py b/api/tests/unit_tests/services/knowledge_fs/test_text_store_diagnostics.py similarity index 100% rename from api/tests/unit_tests/services/test_knowledge_fs_text_store_diagnostics.py rename to api/tests/unit_tests/services/knowledge_fs/test_text_store_diagnostics.py diff --git a/api/tests/unit_tests/services/test_knowledge_fs_text_store_native.py b/api/tests/unit_tests/services/knowledge_fs/test_text_store_native.py similarity index 99% rename from api/tests/unit_tests/services/test_knowledge_fs_text_store_native.py rename to api/tests/unit_tests/services/knowledge_fs/test_text_store_native.py index 973006ddfb9..39df2d52ec8 100644 --- a/api/tests/unit_tests/services/test_knowledge_fs_text_store_native.py +++ b/api/tests/unit_tests/services/knowledge_fs/test_text_store_native.py @@ -12,7 +12,7 @@ from weaviate.classes.config import Configure, ConsistencyLevel, Tokenization from services.knowledge_fs.text_store import TextScope, execute_text_request from services.knowledge_fs.text_store_native import ElasticsearchTextStore, WeaviateTextStore from services.knowledge_fs.vector_store import VectorStoreUnavailableError -from tests.unit_tests.services.test_knowledge_fs_text_store import SCOPE, A, point, request +from tests.unit_tests.services.knowledge_fs.test_text_store import SCOPE, A, point, request if TYPE_CHECKING: from elasticsearch import Elasticsearch diff --git a/api/tests/unit_tests/services/test_knowledge_fs_text_store_tidb.py b/api/tests/unit_tests/services/knowledge_fs/test_text_store_tidb.py similarity index 99% rename from api/tests/unit_tests/services/test_knowledge_fs_text_store_tidb.py rename to api/tests/unit_tests/services/knowledge_fs/test_text_store_tidb.py index 2c991f418d6..4f95e984bb0 100644 --- a/api/tests/unit_tests/services/test_knowledge_fs_text_store_tidb.py +++ b/api/tests/unit_tests/services/knowledge_fs/test_text_store_tidb.py @@ -13,7 +13,7 @@ from models.enums import TidbAuthBindingStatus from services.knowledge_fs.text_store import TextIndexPendingError, configured_text_client, execute_text_request from services.knowledge_fs.text_store_tidb import TidbTextStore, _sql_endpoint from services.knowledge_fs.vector_store import VectorStoreUnavailableError -from tests.unit_tests.services.test_knowledge_fs_text_store import GENERATION, A, point, request +from tests.unit_tests.services.knowledge_fs.test_text_store import GENERATION, A, point, request @dataclass diff --git a/api/tests/unit_tests/services/knowledge_fs/test_vector_store.py b/api/tests/unit_tests/services/knowledge_fs/test_vector_store.py new file mode 100644 index 00000000000..f001b3ab710 --- /dev/null +++ b/api/tests/unit_tests/services/knowledge_fs/test_vector_store.py @@ -0,0 +1,1782 @@ +from __future__ import annotations + +from collections.abc import Callable, Iterator, Mapping +from dataclasses import dataclass, field +from typing import TYPE_CHECKING, cast +from unittest.mock import patch + +import pytest +from elastic_transport import ApiResponseMeta, HttpHeaders, NodeConfig +from elasticsearch import BadRequestError, NotFoundError +from pydantic import ValidationError +from weaviate.classes.config import ConsistencyLevel, DataType, VectorDistances +from weaviate.classes.data import DataObject +from weaviate.classes.query import Filter + +from models.dataset import TidbAuthBinding +from models.enums import TidbAuthBindingStatus +from services.knowledge_fs.vector_store import ( + TIDB_POINT_ID_FIELD, + VectorRequest, + VectorScope, + VectorStoreUnavailableError, + configured_vector_client, + execute_vector_request, +) +from services.knowledge_fs.vector_store_elasticsearch import ElasticsearchVectorStore +from services.knowledge_fs.vector_store_weaviate import WeaviateVectorStore + +if TYPE_CHECKING: + from elasticsearch import Elasticsearch + from weaviate import WeaviateClient + +_KNOWLEDGE_FS_VECTOR_BACKENDS_A, _KNOWLEDGE_FS_VECTOR_BACKENDS_B, _KNOWLEDGE_FS_VECTOR_BACKENDS_GENERATION = ( + f"00000000-0000-4000-8000-{n:012d}" for n in (1, 2, 3) +) +_KNOWLEDGE_FS_VECTOR_BACKENDS_SCOPE: dict[str, object] = { + "tenant_id": _KNOWLEDGE_FS_VECTOR_BACKENDS_A, + "knowledge_space_id": _KNOWLEDGE_FS_VECTOR_BACKENDS_B, + "vector_space_id": "test", + "kind": "dense", + "dimension": 2, +} +POINT: dict[str, object] = { + "id": _KNOWLEDGE_FS_VECTOR_BACKENDS_A, + "generation_id": _KNOWLEDGE_FS_VECTOR_BACKENDS_GENERATION, + "content_hash": "a" * 64, + "vector": [3, 4], +} + +type FieldMappings = dict[str, dict[str, object]] +type IndexMapping = dict[str, dict[str, FieldMappings]] +type ElasticsearchBody = dict[str, object] +type WeaviatePoint = DataObject[dict[str, str], None] + + +def _knowledge_fs_vector_backends_request(operation: str, **kwargs: object) -> VectorRequest: + return VectorRequest.model_validate( + {"operation": operation, "scope": _KNOWLEDGE_FS_VECTOR_BACKENDS_SCOPE, **kwargs} + ) + + +def fail_with_configured_store(expected: type[ElasticsearchVectorStore] | type[WeaviateVectorStore]) -> None: + with configured_vector_client(_KNOWLEDGE_FS_VECTOR_BACKENDS_A) as store: + assert isinstance(store, expected) + raise RuntimeError("failure") + + +def es_error(status: int, name: str) -> NotFoundError | BadRequestError: + meta = ApiResponseMeta(status, "1.1", HttpHeaders(), 0, NodeConfig("http", "localhost", 9200)) + return (NotFoundError if status == 404 else BadRequestError)(name, meta, {"error": name}) + + +def es_mapping() -> IndexMapping: + return { + "mappings": { + "properties": { + "vector": {"type": "dense_vector", "dims": 2, "index": True, "similarity": "cosine"}, + "generation_id": {"type": "keyword"}, + "content_hash": {"type": "keyword"}, + } + } + } + + +@dataclass +class FakeElasticsearchIndices: + """Serves stored mappings per index; a missing index raises Elasticsearch's 404 error.""" + + mappings: dict[str, IndexMapping] = field(default_factory=dict) + queued_get_mapping_results: list[dict[str, IndexMapping] | Exception] = field(default_factory=list) + create_failure: Exception | None = None + get_mapping_calls: list[str] = field(default_factory=list) + created_indexes: list[str] = field(default_factory=list) + created_mappings: list[dict[str, object]] = field(default_factory=list) + + def get_mapping(self, *, index: str) -> dict[str, IndexMapping]: + self.get_mapping_calls.append(index) + if self.queued_get_mapping_results: + result = self.queued_get_mapping_results.pop(0) + if isinstance(result, Exception): + raise result + return result + if index not in self.mappings: + raise es_error(404, "missing") + return {index: self.mappings[index]} + + def create(self, *, index: str, mappings: dict[str, object]) -> None: + self.created_indexes.append(index) + self.created_mappings.append(mappings) + if self.create_failure is not None: + raise self.create_failure + + +@dataclass(frozen=True) +class BulkCall: + operations: list[dict[str, object]] + refresh: str + + def target_indexes(self, action: str) -> set[object]: + """Return the `_index` of every bulk action line of the given kind.""" + indexes: set[object] = set() + for operation in self.operations: + target = operation.get(action) + if isinstance(target, dict): + indexes.add(target.get("_index")) + return indexes + + +@dataclass +class FakeElasticsearch: + indices: FakeElasticsearchIndices = field(default_factory=FakeElasticsearchIndices) + bulk_result: ElasticsearchBody = field(default_factory=lambda: {"items": [{"index": {"status": 201}}]}) + mget_result: ElasticsearchBody = field(default_factory=lambda: {"docs": list[object]()}) + search_result: ElasticsearchBody = field(default_factory=lambda: {"hits": {"hits": list[object]()}}) + bulk_calls: list[BulkCall] = field(default_factory=list) + mget_calls: list[dict[str, object]] = field(default_factory=list) + search_calls: list[dict[str, object]] = field(default_factory=list) + close_count: int = 0 + + def bulk(self, *, operations: list[dict[str, object]], refresh: str) -> ElasticsearchBody: + self.bulk_calls.append(BulkCall(operations=operations, refresh=refresh)) + return self.bulk_result + + def mget(self, *, index: str, ids: list[str]) -> ElasticsearchBody: + self.mget_calls.append({"index": index, "ids": ids}) + return self.mget_result + + def search(self, **kwargs: object) -> ElasticsearchBody: + self.search_calls.append(kwargs) + return self.search_result + + def close(self) -> None: + self.close_count += 1 + + +type EsHarness = tuple[ElasticsearchVectorStore, FakeElasticsearch] + + +@dataclass +class WeaviateMetadata: + distance: float | None = None + + +@dataclass +class WeaviateObject: + uuid: str = _KNOWLEDGE_FS_VECTOR_BACKENDS_A + properties: Mapping[str, object] = field(default_factory=dict) + vector: dict[str, list[float]] = field(default_factory=dict) + metadata: WeaviateMetadata = field(default_factory=WeaviateMetadata) + + +@dataclass +class WeaviateQueryResult: + objects: list[WeaviateObject] = field(default_factory=list) + + +@dataclass +class WeaviateInsertResult: + has_errors: bool + uuids: dict[int, str] + + +@dataclass +class WeaviateDeleteResult: + failed: int = 0 + + +@dataclass(frozen=True) +class WeaviateVectorIndexConfig: + distance_metric: VectorDistances + + +@dataclass(frozen=True) +class WeaviateMultiTenancyConfig: + enabled: bool + + +@dataclass(frozen=True) +class WeaviateProperty: + name: str + data_type: DataType + + +@dataclass +class WeaviateCollectionConfig: + description: str + vectorizer: str + vector_config: dict[str, object] | None + vector_index_config: WeaviateVectorIndexConfig | None + multi_tenancy_config: WeaviateMultiTenancyConfig + properties: list[WeaviateProperty] + + +@dataclass +class FakeWeaviateConfigAccessor: + current: WeaviateCollectionConfig = field( + default_factory=lambda: WeaviateCollectionConfig( + description="KnowledgeFS v1; dimension=2; metric=cosine", + vectorizer="none", + vector_config=None, + vector_index_config=WeaviateVectorIndexConfig(distance_metric=VectorDistances.COSINE), + multi_tenancy_config=WeaviateMultiTenancyConfig(enabled=False), + properties=[ + WeaviateProperty(name=name, data_type=DataType.TEXT) for name in ("generation_id", "content_hash") + ], + ) + ) + + def get(self) -> WeaviateCollectionConfig: + return self.current + + +@dataclass +class FakeWeaviateData: + insert_result: WeaviateInsertResult = field( + default_factory=lambda: WeaviateInsertResult(has_errors=False, uuids={0: _KNOWLEDGE_FS_VECTOR_BACKENDS_A}) + ) + delete_result: WeaviateDeleteResult = field(default_factory=WeaviateDeleteResult) + insert_calls: list[list[WeaviatePoint]] = field(default_factory=list) + delete_calls: list[object] = field(default_factory=list) + + def insert_many(self, objects: list[WeaviatePoint]) -> WeaviateInsertResult: + self.insert_calls.append(objects) + return self.insert_result + + def delete_many(self, *, where: object) -> WeaviateDeleteResult: + self.delete_calls.append(where) + return self.delete_result + + +@dataclass +class FakeWeaviateQuery: + fetch_result: WeaviateQueryResult = field(default_factory=WeaviateQueryResult) + near_result: WeaviateQueryResult = field(default_factory=WeaviateQueryResult) + fetch_calls: list[dict[str, object]] = field(default_factory=list) + near_vector_calls: list[dict[str, object]] = field(default_factory=list) + + def fetch_objects(self, **kwargs: object) -> WeaviateQueryResult: + self.fetch_calls.append(kwargs) + return self.fetch_result + + def near_vector(self, **kwargs: object) -> WeaviateQueryResult: + self.near_vector_calls.append(kwargs) + return self.near_result + + +@dataclass +class FakeWeaviateCollection: + config: FakeWeaviateConfigAccessor = field(default_factory=FakeWeaviateConfigAccessor) + data: FakeWeaviateData = field(default_factory=FakeWeaviateData) + query: FakeWeaviateQuery = field(default_factory=FakeWeaviateQuery) + + +@dataclass +class FakeWeaviateCollectionHandle: + collection: FakeWeaviateCollection + consistency_levels: list[ConsistencyLevel] + + def with_consistency_level(self, level: ConsistencyLevel) -> FakeWeaviateCollection: + self.consistency_levels.append(level) + return self.collection + + +@dataclass +class FakeWeaviateCollections: + """Every collection exists unless `existing` or queued answers say otherwise.""" + + collection: FakeWeaviateCollection = field(default_factory=FakeWeaviateCollection) + existing: set[str] | None = None + queued_exists: list[bool] = field(default_factory=list) + create_failure: Exception | None = None + exists_calls: list[str] = field(default_factory=list) + create_calls: list[dict[str, object]] = field(default_factory=list) + use_calls: list[str] = field(default_factory=list) + consistency_levels: list[ConsistencyLevel] = field(default_factory=list) + + def exists(self, name: str) -> bool: + self.exists_calls.append(name) + if self.queued_exists: + return self.queued_exists.pop(0) + return self.existing is None or name in self.existing + + def create(self, **kwargs: object) -> None: + self.create_calls.append(kwargs) + if self.create_failure is not None: + raise self.create_failure + + def use(self, name: str) -> FakeWeaviateCollectionHandle: + self.use_calls.append(name) + return FakeWeaviateCollectionHandle(collection=self.collection, consistency_levels=self.consistency_levels) + + +@dataclass +class FakeWeaviateClient: + collections: FakeWeaviateCollections = field(default_factory=FakeWeaviateCollections) + close_count: int = 0 + + def close(self) -> None: + self.close_count += 1 + + +type WeaviateHarness = tuple[WeaviateVectorStore, FakeWeaviateClient, FakeWeaviateCollection] + + +class TestKnowledgeFsVectorBackendsCases: + @staticmethod + @pytest.fixture + def es() -> EsHarness: + scope = VectorScope.model_validate(_KNOWLEDGE_FS_VECTOR_BACKENDS_SCOPE) + client = FakeElasticsearch() + client.indices.mappings = {name: es_mapping() for name in (scope.collection_name, scope.legacy_collection_name)} + return ElasticsearchVectorStore(cast("Elasticsearch", client)), client + + @staticmethod + def test_elasticsearch_retry_write_readback_filter_and_cosine(es: EsHarness) -> None: + store, client = es + for _ in range(2): + execute_vector_request(store, _knowledge_fs_vector_backends_request("upsert", points=[POINT])) + write = client.bulk_calls[-1] + assert write.operations[0] == { + "index": { + "_index": VectorScope.model_validate(_KNOWLEDGE_FS_VECTOR_BACKENDS_SCOPE).collection_name, + "_id": _KNOWLEDGE_FS_VECTOR_BACKENDS_A, + } + } + assert write.operations[1] == {key: val for key, val in POINT.items() if key != "id"} + assert write.refresh == "wait_for" + client.mget_result = { + "docs": [{"_id": _KNOWLEDGE_FS_VECTOR_BACKENDS_A, "found": True, "_source": write.operations[1]}] + } + assert execute_vector_request( + store, _knowledge_fs_vector_backends_request("get", ids=[_KNOWLEDGE_FS_VECTOR_BACKENDS_A]) + )["points"] == [POINT] + client.search_result = {"hits": {"hits": [{"_id": _KNOWLEDGE_FS_VECTOR_BACKENDS_A, "_score": 0.8}]}} + matches = execute_vector_request( + store, + _knowledge_fs_vector_backends_request("search", ids=[_KNOWLEDGE_FS_VECTOR_BACKENDS_A], query_vector=[1, 0]), + )["matches"] + assert matches[0]["score"] == pytest.approx(0.6) + search = client.search_calls[-1] + knn = search["knn"] + assert isinstance(knn, dict) + assert knn["filter"] == {"ids": {"values": [_KNOWLEDGE_FS_VECTOR_BACKENDS_A]}} + assert search["allow_partial_search_results"] is False + + @staticmethod + @pytest.mark.parametrize("operation", ["get", "search", "delete"]) + @pytest.mark.parametrize("names", ["legacy", "current", "both"]) + def test_elasticsearch_names_preserve_existing_vectors_and_cleanup( + es: EsHarness, operation: str, names: str + ) -> None: + store, client = es + scope = VectorScope.model_validate(_KNOWLEDGE_FS_VECTOR_BACKENDS_SCOPE) + available = { + name: mapping + for name, mapping in client.indices.mappings.items() + if names == "both" or name == (scope.legacy_collection_name if names == "legacy" else scope.collection_name) + } + client.indices.mappings = available + client.mget_result = { + "docs": [ + { + "_id": _KNOWLEDGE_FS_VECTOR_BACKENDS_A, + "found": operation != "delete", + "_source": {k: v for k, v in POINT.items() if k != "id"}, + } + ] + } + client.search_result = {"hits": {"hits": [{"_id": _KNOWLEDGE_FS_VECTOR_BACKENDS_A, "_score": 1.0}]}} + client.bulk_result = {"items": [{"delete": {"status": 200}}]} + result = execute_vector_request( + store, + _knowledge_fs_vector_backends_request( + operation, + ids=[_KNOWLEDGE_FS_VECTOR_BACKENDS_A], + **({"query_vector": [1, 0], "limit": 1} if operation == "search" else {}), + ), + ) + assert result == { + "points": [POINT] if operation == "get" else [], + "matches": [{"id": _KNOWLEDGE_FS_VECTOR_BACKENDS_A, "score": 1.0}] if operation == "search" else [], + } + if operation == "delete": + assert set().union(*(call.target_indexes("delete") for call in client.bulk_calls)) == set(available) + else: + calls = client.mget_calls if operation == "get" else client.search_calls + assert {call["index"] for call in calls} == set(available) + assert client.indices.created_indexes == [] + + @staticmethod + @pytest.mark.parametrize("race", [False, True]) + def test_elasticsearch_first_writer_or_compatible_creation_race(es: EsHarness, race: bool) -> None: + store, client = es + mapping = dict(client.indices.mappings) + client.indices.queued_get_mapping_results = [es_error(404, "index_not_found_exception"), mapping] + if race: + client.indices.create_failure = es_error(400, "resource_already_exists_exception") + execute_vector_request(store, _knowledge_fs_vector_backends_request("upsert", points=[POINT])) + created = client.indices.created_mappings[-1] + properties = created["properties"] + assert isinstance(properties, dict) + assert properties["vector"]["index_options"] == {"type": "hnsw"} + assert created["dynamic"] == "strict" + + @staticmethod + def test_elasticsearch_creation_error_does_not_hide_configuration_failure(es: EsHarness) -> None: + store, client = es + client.indices.mappings = {} + client.indices.create_failure = es_error(400, "mapper_parsing_exception") + with pytest.raises(BadRequestError): + execute_vector_request(store, _knowledge_fs_vector_backends_request("upsert", points=[POINT])) + assert client.bulk_calls == [] + + @staticmethod + @pytest.mark.parametrize("operation", ["search", "get", "delete"]) + def test_elasticsearch_missing_index_is_explicit_for_search(es: EsHarness, operation: str) -> None: + store, client = es + client.indices.mappings = {} + payload = _knowledge_fs_vector_backends_request( + operation, + ids=[_KNOWLEDGE_FS_VECTOR_BACKENDS_A], + **({"query_vector": [1, 0]} if operation == "search" else {}), + ) + if operation == "search": + with pytest.raises(VectorStoreUnavailableError, match="missing"): + execute_vector_request(store, payload) + else: + assert execute_vector_request(store, payload) == {"points": [], "matches": []} + assert client.indices.created_indexes == [] + + @staticmethod + @pytest.mark.parametrize( + "change", [{"dims": 3}, {"similarity": "l2_norm"}, {"index": False}, {"element_type": "byte"}] + ) + def test_elasticsearch_incompatible_index_is_never_overwritten(es: EsHarness, change: dict[str, object]) -> None: + store, client = es + mapping = next(iter(client.indices.mappings.values()))["mappings"] + mapping["properties"]["vector"].update(change) + with pytest.raises(VectorStoreUnavailableError, match="incompatible"): + execute_vector_request(store, _knowledge_fs_vector_backends_request("upsert", points=[POINT])) + assert client.bulk_calls == [] + assert client.indices.created_indexes == [] + + @staticmethod + def test_elasticsearch_dimension_cap_fails_before_native_io(es: EsHarness) -> None: + store, client = es + scope = {**_KNOWLEDGE_FS_VECTOR_BACKENDS_SCOPE, "dimension": 4097} + with pytest.raises(VectorStoreUnavailableError, match="4096"): + execute_vector_request( + store, + VectorRequest( + operation="get", scope=VectorScope.model_validate(scope), ids=[_KNOWLEDGE_FS_VECTOR_BACKENDS_A] + ), + ) + assert client.indices.get_mapping_calls == [] + + @staticmethod + @pytest.mark.parametrize("result", [{"items": []}, {"items": [{"index": {"status": 429}}]}]) + def test_elasticsearch_partial_write_is_not_success(es: EsHarness, result: dict[str, object]) -> None: + store, client = es + client.bulk_result = result + with pytest.raises(VectorStoreUnavailableError, match="acknowledged"): + execute_vector_request(store, _knowledge_fs_vector_backends_request("upsert", points=[POINT])) + + @staticmethod + @pytest.mark.parametrize("result", [{"docs": []}, {"docs": [{"error": "shard unavailable"}]}]) + def test_elasticsearch_partial_read_is_not_missing_data(es: EsHarness, result: dict[str, object]) -> None: + store, client = es + client.mget_result = result + with pytest.raises(VectorStoreUnavailableError, match="incomplete"): + execute_vector_request( + store, _knowledge_fs_vector_backends_request("get", ids=[_KNOWLEDGE_FS_VECTOR_BACKENDS_A]) + ) + + @staticmethod + @pytest.mark.parametrize("failure", [{"timed_out": True}, {"_shards": {"failed": 1}}]) + def test_elasticsearch_partial_search_is_not_success(es: EsHarness, failure: dict[str, object]) -> None: + store, client = es + client.search_result = failure + with pytest.raises(VectorStoreUnavailableError, match="incomplete"): + execute_vector_request( + store, + _knowledge_fs_vector_backends_request( + "search", ids=[_KNOWLEDGE_FS_VECTOR_BACKENDS_A], query_vector=[1, 0] + ), + ) + + @staticmethod + @pytest.mark.parametrize("status", [200, 404]) + def test_elasticsearch_deletion_requires_visible_absence(es: EsHarness, status: int) -> None: + store, client = es + client.bulk_result = {"items": [{"delete": {"status": status}}]} + client.mget_result = {"docs": [{"_id": _KNOWLEDGE_FS_VECTOR_BACKENDS_A, "found": False}]} + execute_vector_request( + store, _knowledge_fs_vector_backends_request("delete", ids=[_KNOWLEDGE_FS_VECTOR_BACKENDS_A]) + ) + client.mget_result = { + "docs": [ + { + "_id": _KNOWLEDGE_FS_VECTOR_BACKENDS_A, + "found": True, + "_source": {k: v for k, v in POINT.items() if k != "id"}, + } + ] + } + with pytest.raises(VectorStoreUnavailableError, match="not yet visible"): + execute_vector_request( + store, _knowledge_fs_vector_backends_request("delete", ids=[_KNOWLEDGE_FS_VECTOR_BACKENDS_A]) + ) + + @staticmethod + @pytest.mark.parametrize( + ("host", "expected"), + [ + ("elastic.internal", "http://elastic.internal:9200"), + ("https://host:9243", "https://host:9243"), + ("http://[::1]:9201", "http://[::1]:9201"), + ], + ) + def test_elasticsearch_uses_dify_config_and_closes_on_failure( + config_overrides: Callable[..., None], host: str, expected: str + ) -> None: + config_overrides( + VECTOR_STORE="elasticsearch", + ELASTICSEARCH_HOST=host, + ELASTICSEARCH_USE_CLOUD=False, + ELASTICSEARCH_CA_CERTS="/ca.pem", + ELASTICSEARCH_VERIFY_CERTS=True, + ) + with patch("elasticsearch.Elasticsearch") as native: + with pytest.raises(RuntimeError): + fail_with_configured_store(ElasticsearchVectorStore) + assert native.call_args.kwargs["hosts"] == [expected] + assert native.call_args.kwargs["request_timeout"] == 60 + assert native.call_args.kwargs["max_retries"] == 2 + if expected.startswith("https"): + assert native.call_args.kwargs["ca_certs"] == "/ca.pem" + native.return_value.close.assert_called_once() + + @staticmethod + def test_elasticsearch_cloud_key_and_missing_endpoints(config_overrides: Callable[..., None]) -> None: + config_overrides( + VECTOR_STORE="elasticsearch", + ELASTICSEARCH_USE_CLOUD=True, + ELASTICSEARCH_CLOUD_URL="https://cloud:443", + ELASTICSEARCH_API_KEY="secret", + ) + with patch("elasticsearch.Elasticsearch") as native: + with configured_vector_client(_KNOWLEDGE_FS_VECTOR_BACKENDS_A): + assert native.call_args.kwargs["api_key"] == "secret" + assert "basic_auth" not in native.call_args.kwargs + config_overrides(ELASTICSEARCH_CLOUD_URL="") + with pytest.raises(VectorStoreUnavailableError, match="endpoint"): + ElasticsearchVectorStore.from_config() + config_overrides(ELASTICSEARCH_USE_CLOUD=False, ELASTICSEARCH_HOST="") + with pytest.raises(VectorStoreUnavailableError, match="endpoint"): + ElasticsearchVectorStore.from_config() + + @staticmethod + @pytest.fixture + def weaviate() -> WeaviateHarness: + client = FakeWeaviateClient() + return WeaviateVectorStore(cast("WeaviateClient", client)), client, client.collections.collection + + @staticmethod + def test_weaviate_retry_readback_authorized_search_and_cosine(weaviate: WeaviateHarness) -> None: + store, client, collection = weaviate + for _ in range(2): + execute_vector_request(store, _knowledge_fs_vector_backends_request("upsert", points=[POINT])) + name = VectorScope.model_validate(_KNOWLEDGE_FS_VECTOR_BACKENDS_SCOPE).collection_name + assert client.collections.use_calls[-1] == name[0].upper() + name[1:] + assert ( + store.collection_name(VectorScope.model_validate(_KNOWLEDGE_FS_VECTOR_BACKENDS_SCOPE)) + == client.collections.use_calls[-1] + ) + point = collection.data.insert_calls[-1][0] + assert str(point.uuid) == _KNOWLEDGE_FS_VECTOR_BACKENDS_A + assert point.properties == {"generation_id": _KNOWLEDGE_FS_VECTOR_BACKENDS_GENERATION, "content_hash": "a" * 64} + obj = WeaviateObject( + uuid=_KNOWLEDGE_FS_VECTOR_BACKENDS_A, + properties=point.properties, + vector={"default": [3, 4]}, + metadata=WeaviateMetadata(distance=0.4), + ) + collection.query.fetch_result = WeaviateQueryResult(objects=[obj]) + assert execute_vector_request( + store, _knowledge_fs_vector_backends_request("get", ids=[_KNOWLEDGE_FS_VECTOR_BACKENDS_A]) + )["points"] == [POINT] + collection.query.near_result = WeaviateQueryResult(objects=[obj]) + result = execute_vector_request( + store, + _knowledge_fs_vector_backends_request("search", ids=[_KNOWLEDGE_FS_VECTOR_BACKENDS_A], query_vector=[1, 0]), + ) + assert result["matches"] == [{"id": _KNOWLEDGE_FS_VECTOR_BACKENDS_A, "score": 0.6}] + assert collection.query.near_vector_calls[-1]["filters"] == Filter.by_id().contains_any( + [_KNOWLEDGE_FS_VECTOR_BACKENDS_A] + ) + assert client.collections.consistency_levels[-1] == ConsistencyLevel.ALL + assert client.collections.use_calls[-1].startswith("Knowledgefs_v1_") + + @staticmethod + @pytest.mark.parametrize("operation", ["get", "search", "delete"]) + @pytest.mark.parametrize("names", ["legacy", "current", "both"]) + def test_weaviate_names_preserve_existing_vectors_and_cleanup( + weaviate: WeaviateHarness, operation: str, names: str + ) -> None: + store, client, collection = weaviate + scope = VectorScope.model_validate(_KNOWLEDGE_FS_VECTOR_BACKENDS_SCOPE) + available = { + name[0].upper() + name[1:] + for name in (scope.collection_name, scope.legacy_collection_name) + if names == "both" or name == (scope.legacy_collection_name if names == "legacy" else scope.collection_name) + } + client.collections.existing = available + obj = WeaviateObject( + uuid=_KNOWLEDGE_FS_VECTOR_BACKENDS_A, + properties={k: v for k, v in POINT.items() if k not in {"id", "vector"}}, + vector={"default": [3, 4]}, + metadata=WeaviateMetadata(distance=0.0), + ) + collection.query.fetch_result = WeaviateQueryResult(objects=[] if operation == "delete" else [obj]) + collection.query.near_result = WeaviateQueryResult(objects=[obj]) + result = execute_vector_request( + store, + _knowledge_fs_vector_backends_request( + operation, + ids=[_KNOWLEDGE_FS_VECTOR_BACKENDS_A], + **({"query_vector": [1, 0], "limit": 1} if operation == "search" else {}), + ), + ) + assert result == { + "points": [POINT] if operation == "get" else [], + "matches": [{"id": _KNOWLEDGE_FS_VECTOR_BACKENDS_A, "score": 1.0}] if operation == "search" else [], + } + assert set(client.collections.use_calls) == available + if operation == "delete": + assert len(collection.data.delete_calls) == len(available) + assert client.collections.create_calls == [] + + @staticmethod + @pytest.mark.parametrize("race", [False, True]) + def test_weaviate_create_or_compatible_race(weaviate: WeaviateHarness, race: bool) -> None: + store, client, collection = weaviate + client.collections.queued_exists = [False, True] + if race: + client.collections.create_failure = RuntimeError("concurrent create") + execute_vector_request(store, _knowledge_fs_vector_backends_request("upsert", points=[POINT])) + description = client.collections.create_calls[-1]["description"] + assert isinstance(description, str) + assert description.endswith("dimension=2; metric=cosine") + assert len(collection.data.insert_calls) == 1 + + @staticmethod + def test_weaviate_failed_creation_is_not_hidden(weaviate: WeaviateHarness) -> None: + store, client, collection = weaviate + client.collections.existing = set() + client.collections.create_failure = RuntimeError("creation failed") + with pytest.raises(RuntimeError): + execute_vector_request(store, _knowledge_fs_vector_backends_request("upsert", points=[POINT])) + assert collection.data.insert_calls == [] + + @staticmethod + @pytest.mark.parametrize("operation", ["search", "get", "delete"]) + def test_weaviate_missing_collection_search_fails(weaviate: WeaviateHarness, operation: str) -> None: + store, client, collection = weaviate + client.collections.existing = set() + payload = _knowledge_fs_vector_backends_request( + operation, + ids=[_KNOWLEDGE_FS_VECTOR_BACKENDS_A], + **({"query_vector": [1, 0]} if operation == "search" else {}), + ) + if operation == "search": + with pytest.raises(VectorStoreUnavailableError, match="missing"): + execute_vector_request(store, payload) + else: + assert execute_vector_request(store, payload) == {"points": [], "matches": []} + assert collection.data.insert_calls == [] + + @staticmethod + @pytest.mark.parametrize( + ("field", "value"), + [ + ("description", "wrong dimensions"), + ("vectorizer", "text2vec-openai"), + ("vector_config", {"named": {}}), + ("vector_index_config", None), + ("properties", []), + ], + ) + def test_weaviate_incompatible_schema_is_never_written( + weaviate: WeaviateHarness, field: str, value: object + ) -> None: + store, client, collection = weaviate + setattr(collection.config.current, field, value) + with pytest.raises(VectorStoreUnavailableError, match="incompatible"): + execute_vector_request(store, _knowledge_fs_vector_backends_request("upsert", points=[POINT])) + assert collection.data.insert_calls == [] + + @staticmethod + @pytest.mark.parametrize( + "result", [WeaviateInsertResult(has_errors=True, uuids={}), WeaviateInsertResult(has_errors=False, uuids={})] + ) + def test_weaviate_partial_batch_is_not_success(weaviate: WeaviateHarness, result: WeaviateInsertResult) -> None: + store, client, collection = weaviate + collection.data.insert_result = result + with pytest.raises(VectorStoreUnavailableError, match="acknowledged"): + execute_vector_request(store, _knowledge_fs_vector_backends_request("upsert", points=[POINT])) + + @staticmethod + def test_weaviate_search_missing_distance_is_not_success(weaviate: WeaviateHarness) -> None: + store, client, collection = weaviate + collection.query.near_result = WeaviateQueryResult( + objects=[WeaviateObject(metadata=WeaviateMetadata(distance=None))] + ) + with pytest.raises(VectorStoreUnavailableError, match="omitted distance"): + execute_vector_request( + store, + _knowledge_fs_vector_backends_request( + "search", ids=[_KNOWLEDGE_FS_VECTOR_BACKENDS_A], query_vector=[1, 0] + ), + ) + + @staticmethod + def test_weaviate_delete_requires_ack_and_visible_absence(weaviate: WeaviateHarness) -> None: + store, client, collection = weaviate + for _ in range(2): + execute_vector_request( + store, _knowledge_fs_vector_backends_request("delete", ids=[_KNOWLEDGE_FS_VECTOR_BACKENDS_A]) + ) + collection.data.delete_result.failed = 1 + with pytest.raises(VectorStoreUnavailableError, match="not yet visible"): + execute_vector_request( + store, _knowledge_fs_vector_backends_request("delete", ids=[_KNOWLEDGE_FS_VECTOR_BACKENDS_A]) + ) + collection.data.delete_result.failed = 0 + collection.query.fetch_result.objects = [WeaviateObject(uuid=_KNOWLEDGE_FS_VECTOR_BACKENDS_A)] + with pytest.raises(VectorStoreUnavailableError, match="not yet visible"): + execute_vector_request( + store, _knowledge_fs_vector_backends_request("delete", ids=[_KNOWLEDGE_FS_VECTOR_BACKENDS_A]) + ) + + @staticmethod + @pytest.mark.parametrize( + ("endpoint", "grpc", "secure", "port"), + [ + ("http://host:8080", None, False, 50051), + ("https://cloud", None, True, 443), + ("http://host", "grpcs://rpc:443", True, 443), + ("host", "rpc:50052", False, 50052), + ], + ) + def test_weaviate_config_and_resource_cleanup( + config_overrides: Callable[..., None], endpoint: str, grpc: str | None, secure: bool, port: int + ) -> None: + config_overrides( + VECTOR_STORE="weaviate", WEAVIATE_ENDPOINT=endpoint, WEAVIATE_GRPC_ENDPOINT=grpc, WEAVIATE_API_KEY="secret" + ) + with patch("weaviate.connect_to_custom") as connect: + with pytest.raises(RuntimeError): + fail_with_configured_store(WeaviateVectorStore) + assert connect.call_args.kwargs["grpc_secure"] is secure + assert connect.call_args.kwargs["grpc_port"] == port + connect.return_value.close.assert_called_once() + + @staticmethod + def test_weaviate_missing_and_invalid_endpoints(config_overrides: Callable[..., None]) -> None: + config_overrides(VECTOR_STORE="weaviate", WEAVIATE_ENDPOINT="") + with pytest.raises(VectorStoreUnavailableError, match="endpoint"): + WeaviateVectorStore.from_config() + config_overrides(WEAVIATE_ENDPOINT="http://host", WEAVIATE_GRPC_ENDPOINT="https://wrong") + with pytest.raises(VectorStoreUnavailableError, match="gRPC"): + WeaviateVectorStore.from_config() + + +if TYPE_CHECKING: + from qdrant_client import QdrantClient as QdrantClientType + from qdrant_client.http.models import CollectionInfo, FilterSelector, PointIdsList, PointStruct + + from services.knowledge_fs.vector_store import VectorResult + +QdrantClient = pytest.importorskip("qdrant_client").QdrantClient +models = pytest.importorskip("qdrant_client.http.models") +UnexpectedResponse = pytest.importorskip("qdrant_client.http.exceptions").UnexpectedResponse + +TENANT = "018f0d60-7a49-7cc2-9c1b-5b36f18f0001" +SPACE = "018f0d60-7a49-7cc2-9c1b-5b36f18f0002" +_KNOWLEDGE_FS_VECTOR_STORE_GENERATION = "018f0d60-7a49-7cc2-9c1b-5b36f18f0003" +_KNOWLEDGE_FS_VECTOR_STORE_A = "018f0d60-7a49-7cc2-9c1b-5b36f18f0004" +_KNOWLEDGE_FS_VECTOR_STORE_B = "018f0d60-7a49-7cc2-9c1b-5b36f18f0005" +_KNOWLEDGE_FS_VECTOR_STORE_SCOPE: dict[str, object] = { + "tenant_id": TENANT, + "knowledge_space_id": SPACE, + "vector_space_id": "embedding-v1", + "kind": "dense", + "dimension": 2, +} + + +@dataclass +class FakeQdrant: + """Qdrant client surface used by the bridge; unknown collections raise the remote 404 contract.""" + + collections: dict[str, CollectionInfo] = field(default_factory=dict) + search_result: list[object] = field(default_factory=list) + retrieve_result: list[object] = field(default_factory=list) + retrieve_hook: Callable[..., list[object]] | None = None + retrieve_failure: Exception | None = None + delete_failure: Exception | None = None + create_payload_index_hook: Callable[..., None] | None = None + get_collection_calls: list[str] = field(default_factory=list) + create_collection_calls: list[dict[str, object]] = field(default_factory=list) + create_payload_index_calls: list[dict[str, object]] = field(default_factory=list) + upsert_calls: list[dict[str, object]] = field(default_factory=list) + upserted_points: list[list[PointStruct]] = field(default_factory=list) + search_calls: list[dict[str, object]] = field(default_factory=list) + retrieve_calls: list[dict[str, object]] = field(default_factory=list) + delete_calls: list[dict[str, object]] = field(default_factory=list) + delete_selectors: list[PointIdsList | FilterSelector] = field(default_factory=list) + close_count: int = 0 + + def execute(self, payload: VectorRequest) -> VectorResult: + """Run the bridge against this fake through the Qdrant client contract.""" + return execute_vector_request(cast("QdrantClientType", self), payload) + + def get_collection(self, name: str) -> CollectionInfo: + self.get_collection_calls.append(name) + if name not in self.collections: + raise UnexpectedResponse(404, "missing", b"{}", {}) + return self.collections[name] + + def create_collection(self, **kwargs: object) -> None: + self.create_collection_calls.append(kwargs) + + def create_payload_index(self, **kwargs: object) -> None: + self.create_payload_index_calls.append(kwargs) + if self.create_payload_index_hook is not None: + self.create_payload_index_hook(**kwargs) + + def upsert(self, *, collection_name: str, wait: bool, points: list[PointStruct]) -> None: + self.upsert_calls.append({"collection_name": collection_name, "wait": wait, "points": points}) + self.upserted_points.append(points) + + def search(self, **kwargs: object) -> list[object]: + self.search_calls.append(kwargs) + return self.search_result + + def retrieve(self, **kwargs: object) -> list[object]: + self.retrieve_calls.append(kwargs) + if self.retrieve_failure is not None: + raise self.retrieve_failure + if self.retrieve_hook is not None: + return self.retrieve_hook(**kwargs) + return self.retrieve_result + + def delete(self, *, collection_name: str, points_selector: PointIdsList | FilterSelector, wait: bool) -> None: + self.delete_calls.append({"collection_name": collection_name, "points_selector": points_selector, "wait": wait}) + self.delete_selectors.append(points_selector) + if self.delete_failure is not None: + raise self.delete_failure + + def close(self) -> None: + self.close_count += 1 + + +def _knowledge_fs_vector_store_request(operation: str, **kwargs: object) -> VectorRequest: + return VectorRequest.model_validate({"operation": operation, "scope": _KNOWLEDGE_FS_VECTOR_STORE_SCOPE, **kwargs}) + + +type TidbHarness = tuple[FakeQdrant, CollectionInfo] + + +class TestKnowledgeFsVectorStoreCases: + @staticmethod + @pytest.fixture(autouse=True) + def _vector_store_config(config_overrides: Callable[..., None]) -> None: + config_overrides(DEPLOYMENT_EDITION="COMMUNITY") + + @staticmethod + @pytest.fixture + def client() -> Iterator[QdrantClientType]: + client = QdrantClient(":memory:") + client.create_collection( + VectorScope.model_validate(_KNOWLEDGE_FS_VECTOR_STORE_SCOPE).collection_name, + vectors_config=models.VectorParams(size=2, distance=models.Distance.COSINE), + ) + native_get = client.get_collection + + def get_collection(name: str) -> CollectionInfo: + # Match the remote API's 404 contract while exercising native vector I/O. + if not client.collection_exists(name): + raise UnexpectedResponse(404, "missing", b"{}", {}) + return native_get(name) + + with patch.object(client, "get_collection", side_effect=get_collection): + yield client + client.close() + + @staticmethod + def test_native_vector_roundtrip_authorized_id_filter_and_idempotent_deletion(client: QdrantClientType) -> None: + points = [ + { + "id": _KNOWLEDGE_FS_VECTOR_STORE_A, + "generation_id": _KNOWLEDGE_FS_VECTOR_STORE_GENERATION, + "content_hash": "a" * 64, + "vector": [3, 4], + }, + { + "id": _KNOWLEDGE_FS_VECTOR_STORE_B, + "generation_id": _KNOWLEDGE_FS_VECTOR_STORE_GENERATION, + "content_hash": "b" * 64, + "vector": [1, 0], + }, + ] + execute_vector_request(client, _knowledge_fs_vector_store_request("upsert", points=points)) + execute_vector_request(client, _knowledge_fs_vector_store_request("upsert", points=points)) + actual = execute_vector_request( + client, + _knowledge_fs_vector_store_request("get", ids=[_KNOWLEDGE_FS_VECTOR_STORE_A, _KNOWLEDGE_FS_VECTOR_STORE_B]), + )["points"] + assert len(actual) == 2 + assert actual[0]["vector"] == pytest.approx([0.6, 0.8]) + assert actual[0]["content_hash"] == "a" * 64 + # The nearest point B is unauthorized. It must not crowd out the authorized A. + matches = execute_vector_request( + client, + _knowledge_fs_vector_store_request( + "search", ids=[_KNOWLEDGE_FS_VECTOR_STORE_A], query_vector=[1, 0], limit=1 + ), + )["matches"] + assert [item["id"] for item in matches] == [_KNOWLEDGE_FS_VECTOR_STORE_A] + assert matches[0]["score"] == pytest.approx(0.6) + for _ in range(2): + execute_vector_request( + client, _knowledge_fs_vector_store_request("delete", ids=[_KNOWLEDGE_FS_VECTOR_STORE_A]) + ) + assert ( + execute_vector_request( + client, _knowledge_fs_vector_store_request("get", ids=[_KNOWLEDGE_FS_VECTOR_STORE_A]) + )["points"] + == [] + ) + assert ( + len( + execute_vector_request( + client, _knowledge_fs_vector_store_request("get", ids=[_KNOWLEDGE_FS_VECTOR_STORE_B]) + )["points"] + ) + == 1 + ) + + @staticmethod + @pytest.mark.parametrize( + ("field", "value"), + [ + ("tenant_id", _KNOWLEDGE_FS_VECTOR_STORE_B), + ("knowledge_space_id", _KNOWLEDGE_FS_VECTOR_STORE_B), + ("vector_space_id", "other"), + ("kind", "visual"), + ("dimension", 3), + ], + ) + def test_each_scope_dimension_has_a_distinct_reserved_namespace(field: str, value: object) -> None: + original = VectorScope.model_validate(_KNOWLEDGE_FS_VECTOR_STORE_SCOPE) + other = VectorScope.model_validate({**_KNOWLEDGE_FS_VECTOR_STORE_SCOPE, field: value}) + assert other.collection_name != original.collection_name + assert other.legacy_collection_name != original.legacy_collection_name + assert original.collection_name.startswith("knowledgefs_v1_dense_") + assert SPACE not in original.collection_name + + @staticmethod + @pytest.mark.parametrize("kind", ["dense", "visual", "graph-entity", "graph-relation"]) + def test_typed_collection_names_fit_tidb_delete_endpoint_and_normalize_uuid_case(kind: str) -> None: + scope = VectorScope.model_validate({**_KNOWLEDGE_FS_VECTOR_STORE_SCOPE, "kind": kind}) + prefix = f"knowledgefs_v1_{kind.replace('-', '_')}_" + assert scope.collection_name.startswith(prefix) + assert len(scope.collection_name.removeprefix(prefix)) == min(32, 56 - len(prefix)) + assert len(scope.collection_name.removeprefix(prefix)) >= 26 + assert len(scope.collection_name) <= 56 + assert ( + VectorScope.model_validate({**_KNOWLEDGE_FS_VECTOR_STORE_SCOPE, "kind": kind, "tenant_id": TENANT.upper()}) + == scope + ) + + @staticmethod + def test_legacy_namespace_keeps_pre_rename_hash() -> None: + # Pin the pre-rename format so compatibility cannot drift with the new one. + scope = VectorScope.model_validate(_KNOWLEDGE_FS_VECTOR_STORE_SCOPE) + assert scope.legacy_collection_name == "knowledgefs_v1_488d91e70fe33beedcda4cb5e47b954dd5af7ca3" + assert ( + scope.collection_name + == "knowledgefs_v1_dense_" + scope.legacy_collection_name.removeprefix("knowledgefs_v1_")[:32] + ) + + @staticmethod + def test_legacy_and_typed_points_remain_searchable_and_deletable_together(client: QdrantClientType) -> None: + scope = VectorScope.model_validate(_KNOWLEDGE_FS_VECTOR_STORE_SCOPE) + legacy = scope.legacy_collection_name + current = scope.collection_name + client.create_collection(legacy, vectors_config=models.VectorParams(size=2, distance=models.Distance.COSINE)) + points = [ + { + "id": _KNOWLEDGE_FS_VECTOR_STORE_A, + "generation_id": _KNOWLEDGE_FS_VECTOR_STORE_GENERATION, + "content_hash": "a" * 64, + "vector": [3, 4], + }, + { + "id": _KNOWLEDGE_FS_VECTOR_STORE_B, + "generation_id": _KNOWLEDGE_FS_VECTOR_STORE_GENERATION, + "content_hash": "b" * 64, + "vector": [1, 0], + }, + ] + client.upsert( + legacy, + points=[ + models.PointStruct( + id=point["id"], + vector=point["vector"], + payload={k: v for k, v in point.items() if k not in {"id", "vector"}}, + ) + for point in points + ], + ) + client.delete_collection(current) + assert ( + len( + execute_vector_request( + client, + _knowledge_fs_vector_store_request( + "get", ids=[_KNOWLEDGE_FS_VECTOR_STORE_A, _KNOWLEDGE_FS_VECTOR_STORE_B] + ), + )["points"] + ) + == 2 + ) + assert ( + execute_vector_request( + client, + _knowledge_fs_vector_store_request("search", ids=[_KNOWLEDGE_FS_VECTOR_STORE_B], query_vector=[1, 0]), + )["matches"][0]["id"] + == _KNOWLEDGE_FS_VECTOR_STORE_B + ) + assert not client.collection_exists(current) + + # A new document in this scope must create the typed name, without hiding old A. + execute_vector_request(client, _knowledge_fs_vector_store_request("upsert", points=[points[1]])) + assert client.count(current).count == 1 + assert client.count(legacy).count == 2 + assert ( + len( + execute_vector_request( + client, + _knowledge_fs_vector_store_request( + "get", ids=[_KNOWLEDGE_FS_VECTOR_STORE_A, _KNOWLEDGE_FS_VECTOR_STORE_B] + ), + )["points"] + ) + == 2 + ) + matches = execute_vector_request( + client, + _knowledge_fs_vector_store_request( + "search", ids=[_KNOWLEDGE_FS_VECTOR_STORE_A, _KNOWLEDGE_FS_VECTOR_STORE_B], query_vector=[1, 0], limit=1 + ), + )["matches"] + assert [match["id"] for match in matches] == [_KNOWLEDGE_FS_VECTOR_STORE_B] + # B exists in both collections, but cannot crowd out an authorized point A. + matches = execute_vector_request( + client, + _knowledge_fs_vector_store_request( + "search", ids=[_KNOWLEDGE_FS_VECTOR_STORE_A], query_vector=[1, 0], limit=1 + ), + )["matches"] + assert [match["id"] for match in matches] == [_KNOWLEDGE_FS_VECTOR_STORE_A] + for _ in range(2): + execute_vector_request( + client, _knowledge_fs_vector_store_request("delete", ids=[_KNOWLEDGE_FS_VECTOR_STORE_B]) + ) + assert client.retrieve(current, [_KNOWLEDGE_FS_VECTOR_STORE_B]) == [] + assert client.retrieve(legacy, [_KNOWLEDGE_FS_VECTOR_STORE_B]) == [] + assert ( + len( + execute_vector_request( + client, _knowledge_fs_vector_store_request("get", ids=[_KNOWLEDGE_FS_VECTOR_STORE_A]) + )["points"] + ) + == 1 + ) + + @staticmethod + @pytest.mark.parametrize("operation", ["get", "search", "delete"]) + def test_legacy_backend_failure_never_returns_partial_success(client: QdrantClientType, operation: str) -> None: + execute_vector_request( + client, + _knowledge_fs_vector_store_request( + "upsert", + points=[ + { + "id": _KNOWLEDGE_FS_VECTOR_STORE_A, + "generation_id": _KNOWLEDGE_FS_VECTOR_STORE_GENERATION, + "content_hash": "a" * 64, + "vector": [1, 0], + } + ], + ), + ) + native_get = client.get_collection + + def get_collection(name: str) -> CollectionInfo: + if name == VectorScope.model_validate(_KNOWLEDGE_FS_VECTOR_STORE_SCOPE).legacy_collection_name: + raise UnexpectedResponse(503, "unavailable", b"{}", {}) + return native_get(name) + + payload = _knowledge_fs_vector_store_request( + operation, ids=[_KNOWLEDGE_FS_VECTOR_STORE_A], **({"query_vector": [1, 0]} if operation == "search" else {}) + ) + with patch.object(client, "get_collection", side_effect=get_collection): + with pytest.raises(UnexpectedResponse): + execute_vector_request(client, payload) + + @staticmethod + @pytest.mark.parametrize( + "payload", + [ + {"operation": "search", "ids": [], "query_vector": [1, 0]}, + {"operation": "search", "ids": [_KNOWLEDGE_FS_VECTOR_STORE_A], "query_vector": [1]}, + {"operation": "search", "ids": [_KNOWLEDGE_FS_VECTOR_STORE_A], "query_vector": [0, 0]}, + {"operation": "search", "ids": [_KNOWLEDGE_FS_VECTOR_STORE_A], "query_vector": [float("nan"), 1]}, + { + "operation": "search", + "ids": [_KNOWLEDGE_FS_VECTOR_STORE_A, _KNOWLEDGE_FS_VECTOR_STORE_A], + "query_vector": [1, 0], + }, + { + "operation": "search", + "ids": [_KNOWLEDGE_FS_VECTOR_STORE_A], + "query_vector": [1, 0], + "collection_name": "dataset", + }, + { + "operation": "search", + "ids": [_KNOWLEDGE_FS_VECTOR_STORE_A], + "query_vector": [1, 0], + "endpoint": "http://example.com", + }, + {"operation": "get", "ids": ["not-a-uuid"]}, + {"operation": "delete", "ids": [_KNOWLEDGE_FS_VECTOR_STORE_A], "query_vector": [1, 0]}, + {"operation": "upsert", "points": []}, + ], + ) + def test_invalid_operations_fail_before_backend_io(payload: dict[str, object]) -> None: + with pytest.raises(ValidationError): + VectorRequest.model_validate({"scope": _KNOWLEDGE_FS_VECTOR_STORE_SCOPE, **payload}) + + @staticmethod + def test_missing_collection_search_is_not_reported_as_empty_success() -> None: + client = FakeQdrant() + with pytest.raises(VectorStoreUnavailableError, match="missing"): + client.execute( + _knowledge_fs_vector_store_request("search", ids=[_KNOWLEDGE_FS_VECTOR_STORE_A], query_vector=[1, 0]) + ) + assert client.create_collection_calls == [] + assert ( + client.execute(_knowledge_fs_vector_store_request("delete", ids=[_KNOWLEDGE_FS_VECTOR_STORE_A]))["points"] + == [] + ) + + @staticmethod + def test_collection_creation_race_accepts_only_compatible_existing_collection(client: QdrantClientType) -> None: + real_get = client.get_collection + with ( + patch.object( + client, + "get_collection", + side_effect=[ + UnexpectedResponse(404, "missing", b"{}", {}), + real_get(VectorScope.model_validate(_KNOWLEDGE_FS_VECTOR_STORE_SCOPE).collection_name), + ], + ), + patch.object(client, "create_collection", side_effect=RuntimeError("already exists")), + ): + execute_vector_request( + client, + _knowledge_fs_vector_store_request( + "upsert", + points=[ + { + "id": _KNOWLEDGE_FS_VECTOR_STORE_A, + "generation_id": _KNOWLEDGE_FS_VECTOR_STORE_GENERATION, + "content_hash": "a" * 64, + "vector": [1, 0], + } + ], + ), + ) + assert ( + len( + execute_vector_request( + client, _knowledge_fs_vector_store_request("get", ids=[_KNOWLEDGE_FS_VECTOR_STORE_A]) + )["points"] + ) + == 1 + ) + + @staticmethod + def test_no_binding_never_allocates_a_paid_cluster(config_overrides: Callable[..., None]) -> None: + from services.tidb_binding_service import TidbBindingUnavailableError + + config_overrides(VECTOR_STORE="tidb_on_qdrant") + with patch("services.tidb_binding_service._load_binding", return_value=None): + with pytest.raises(TidbBindingUnavailableError, match="no active"): + with configured_vector_client(TENANT): + pytest.fail("unconfigured backend was opened") + + @staticmethod + def test_unsupported_backend_fails_explicitly(config_overrides: Callable[..., None]) -> None: + config_overrides(VECTOR_STORE="milvus") + with pytest.raises(VectorStoreUnavailableError, match="does not support"): + with configured_vector_client(TENANT): + pytest.fail("unsupported backend was opened") + + @staticmethod + @pytest.mark.parametrize("allow_create", [True, False]) + def test_tidb_client_uses_tenant_binding(config_overrides: Callable[..., None], allow_create: bool) -> None: + config_overrides(VECTOR_STORE="tidb_on_qdrant", TIDB_ON_QDRANT_URL="https://global.invalid") + binding = TidbAuthBinding( + tenant_id=TENANT, + cluster_id="cluster-1", + cluster_name="tenant-cluster", + active=True, + status=TidbAuthBindingStatus.ACTIVE, + account="tenant-account", + password="tenant-secret", + qdrant_endpoint="https://tenant.invalid", + ) + client = FakeQdrant() + + def open_client(**kwargs: object) -> FakeQdrant: + assert kwargs["url"] == "https://tenant.invalid" + assert kwargs["api_key"] == "tenant-account:tenant-secret" + return client + + with ( + patch("services.tidb_binding_service.resolve_tidb_auth_binding", return_value=binding) as resolve, + patch("qdrant_client.QdrantClient", side_effect=open_client), + ): + with configured_vector_client(TENANT, allow_create=allow_create) as configured: + assert configured is client + resolve.assert_called_once_with(TENANT, allow_create=allow_create) + assert client.close_count == 1 + + @staticmethod + def test_qdrant_configuration_closes_client_even_when_operation_fails( + config_overrides: Callable[..., None], + ) -> None: + config_overrides(VECTOR_STORE="qdrant", QDRANT_URL="https://configured.invalid", QDRANT_API_KEY="private") + with patch("qdrant_client.QdrantClient") as factory: + with pytest.raises(RuntimeError): + with configured_vector_client(TENANT): + raise RuntimeError("operation failed") + assert factory.call_args.kwargs["url"] == "https://configured.invalid" + factory.return_value.close.assert_called_once() + config_overrides(QDRANT_URL="") + with pytest.raises(VectorStoreUnavailableError, match="endpoint"): + with configured_vector_client(TENANT): + pytest.fail("missing endpoint accepted") + + @staticmethod + def test_cloud_quota_is_enforced_before_any_vector_mutation() -> None: + from services.vector_space_admission_service import VectorSpaceAdmissionError + + client = FakeQdrant() + with patch( + "services.vector_space_admission_service.VectorSpaceAdmissionService.ensure_external_points_can_be_indexed", + side_effect=VectorSpaceAdmissionError("limit reached"), + ): + with pytest.raises(VectorSpaceAdmissionError): + client.execute( + _knowledge_fs_vector_store_request( + "upsert", + points=[ + { + "id": _KNOWLEDGE_FS_VECTOR_STORE_A, + "generation_id": _KNOWLEDGE_FS_VECTOR_STORE_GENERATION, + "content_hash": "a" * 64, + "vector": [1, 0], + } + ], + ), + ) + assert client.get_collection_calls == [] + assert client.upsert_calls == [] + + @staticmethod + def test_cloud_quota_uses_existing_watermark_without_embedding_calls(config_overrides: Callable[..., None]) -> None: + from enums import CloudPlan, DeploymentEdition + from services.vector_space_admission_service import VectorSpaceAdmissionError, VectorSpaceAdmissionService + + config_overrides( + DEPLOYMENT_EDITION=DeploymentEdition.CLOUD, + VECTOR_STORE="tidb_on_qdrant", + TIDB_ON_QDRANT_ESTIMATED_STORAGE_LIMITS_MB="sandbox:60,professional:6144,team:24576", + ) + service = VectorSpaceAdmissionService() + with ( + patch.object(service, "_get_plan", return_value=CloudPlan.SANDBOX), + patch.object(service, "_get_usage_and_limit_mb", return_value=(0.0, 50)), + patch.object(service, "_get_embedding_dimension") as dimension, + patch.object(service, "_reserve_projected_usage", return_value=(0, 60 * 1024 * 1024)) as reserve, + ): + for _ in range(2): + service.ensure_external_points_can_be_indexed( + tenant_id=TENANT, batch_id="stable-batch", point_count=2, dimension=3 + ) + assert reserve.call_args.kwargs["document_id"] == "knowledgefs:stable-batch" + assert reserve.call_args.kwargs["document_estimate_bytes"] == 2 * (3 * 8 + 3584) + assert reserve.call_args.kwargs["estimate_limit_bytes"] == 60 * 1024 * 1024 + dimension.assert_not_called() + reserve.return_value = (0, 61 * 1024 * 1024) + with pytest.raises(VectorSpaceAdmissionError, match="exceeding the 50 MB limit"): + service.ensure_external_points_can_be_indexed( + tenant_id=TENANT, batch_id="large", point_count=2, dimension=3 + ) + with pytest.raises(VectorSpaceAdmissionError, match="Invalid"): + service.ensure_external_points_can_be_indexed(tenant_id=TENANT, batch_id="bad", point_count=0, dimension=3) + + @staticmethod + def test_unmetered_workspace_does_not_contact_billing_or_reserve_quota( + config_overrides: Callable[..., None], + ) -> None: + from services.vector_space_admission_service import VectorSpaceAdmissionService + + config_overrides(DEPLOYMENT_EDITION="CLOUD", VECTOR_STORE="tidb_on_qdrant") + service = VectorSpaceAdmissionService() + with ( + patch.object(service, "_get_plan", return_value=None), + patch.object(service, "_get_usage_and_limit_mb") as billing, + patch.object(service, "_reserve_projected_usage") as reserve, + ): + service.ensure_external_points_can_be_indexed( + tenant_id=TENANT, batch_id="unmetered", point_count=2, dimension=3 + ) + billing.assert_not_called() + reserve.assert_not_called() + + @staticmethod + @pytest.mark.parametrize("limit", [50, -1]) + def test_external_points_with_unknown_usage_fail_before_reservation( + config_overrides: Callable[..., None], limit: int + ) -> None: + from enums import CloudPlan + from services.vector_space_admission_service import ( + VectorSpaceAdmissionService, + VectorSpaceAdmissionUnavailableError, + ) + + config_overrides(DEPLOYMENT_EDITION="CLOUD", VECTOR_STORE="tidb_on_qdrant") + service = VectorSpaceAdmissionService() + with ( + patch.object(service, "_get_plan", return_value=CloudPlan.SANDBOX), + patch( + "services.vector_space_admission_service.BillingService.get_vector_space", + return_value={"size": 0, "limit": limit, "usage_unknown": True}, + ), + patch.object(service, "_reserve_projected_usage") as reserve, + ): + with pytest.raises(VectorSpaceAdmissionUnavailableError, match="Unable to verify vector storage usage"): + service.ensure_external_points_can_be_indexed( + tenant_id=TENANT, batch_id="unknown", point_count=2, dimension=3 + ) + reserve.assert_not_called() + + @staticmethod + @pytest.mark.parametrize(("edition", "vector_store"), [("COMMUNITY", "tidb_on_qdrant"), ("CLOUD", "qdrant")]) + def test_external_points_skip_billing_outside_cloud_tidb( + config_overrides: Callable[..., None], edition: str, vector_store: str + ) -> None: + from services.vector_space_admission_service import VectorSpaceAdmissionService + + config_overrides(DEPLOYMENT_EDITION=edition, VECTOR_STORE=vector_store) + service = VectorSpaceAdmissionService() + with ( + patch("services.vector_space_admission_service.BillingService.get_vector_space") as billing, + patch.object(service, "_reserve_projected_usage") as reserve, + ): + service.ensure_external_points_can_be_indexed( + tenant_id=TENANT, batch_id="unmetered", point_count=2, dimension=3 + ) + billing.assert_not_called() + reserve.assert_not_called() + + @staticmethod + def test_incompatible_collection_is_never_recreated_or_written(client: QdrantClientType) -> None: + collection = VectorScope.model_validate(_KNOWLEDGE_FS_VECTOR_STORE_SCOPE).collection_name + client.delete_collection(collection) + client.create_collection( + collection, vectors_config=models.VectorParams(size=3, distance=models.Distance.COSINE) + ) + with pytest.raises(VectorStoreUnavailableError, match="incompatible"): + execute_vector_request( + client, + _knowledge_fs_vector_store_request( + "upsert", + points=[ + { + "id": _KNOWLEDGE_FS_VECTOR_STORE_A, + "generation_id": _KNOWLEDGE_FS_VECTOR_STORE_GENERATION, + "content_hash": "a" * 64, + "vector": [1, 0], + } + ], + ), + ) + from qdrant_client.http.models import VectorParams + + vectors = client.get_collection(collection).config.params.vectors + assert isinstance(vectors, VectorParams) + assert vectors.size == 3 + + @staticmethod + def test_uncertain_native_delete_is_retryable(client: QdrantClientType) -> None: + execute_vector_request( + client, + _knowledge_fs_vector_store_request( + "upsert", + points=[ + { + "id": _KNOWLEDGE_FS_VECTOR_STORE_A, + "generation_id": _KNOWLEDGE_FS_VECTOR_STORE_GENERATION, + "content_hash": "a" * 64, + "vector": [1, 0], + } + ], + ), + ) + with patch.object(client, "delete"): + with pytest.raises(VectorStoreUnavailableError, match="not yet visible"): + execute_vector_request( + client, _knowledge_fs_vector_store_request("delete", ids=[_KNOWLEDGE_FS_VECTOR_STORE_A]) + ) + execute_vector_request(client, _knowledge_fs_vector_store_request("delete", ids=[_KNOWLEDGE_FS_VECTOR_STORE_A])) + assert ( + execute_vector_request( + client, _knowledge_fs_vector_store_request("get", ids=[_KNOWLEDGE_FS_VECTOR_STORE_A]) + )["points"] + == [] + ) + + @staticmethod + @pytest.fixture + def tidb_client( + client: QdrantClientType, config_overrides: Callable[..., None] + ) -> tuple[FakeQdrant, CollectionInfo]: + config_overrides(VECTOR_STORE="tidb_on_qdrant") + scope = VectorScope.model_validate(_KNOWLEDGE_FS_VECTOR_STORE_SCOPE) + info = client.get_collection(scope.collection_name) + return FakeQdrant(collections={scope.collection_name: info}), info + + @staticmethod + def test_tidb_new_collection_indexes_and_stores_point_ids(tidb_client: TidbHarness) -> None: + client, info = tidb_client + client.execute( + _knowledge_fs_vector_store_request( + "upsert", + points=[ + { + "id": _KNOWLEDGE_FS_VECTOR_STORE_A, + "generation_id": _KNOWLEDGE_FS_VECTOR_STORE_GENERATION, + "content_hash": "a" * 64, + "vector": [1, 0], + } + ], + ), + ) + assert client.create_payload_index_calls == [ + { + "collection_name": VectorScope.model_validate(_KNOWLEDGE_FS_VECTOR_STORE_SCOPE).collection_name, + "field_name": TIDB_POINT_ID_FIELD, + "field_schema": models.PayloadSchemaType.KEYWORD, + "wait": True, + } + ] + assert client.upserted_points[-1][0].payload == { + "generation_id": _KNOWLEDGE_FS_VECTOR_STORE_GENERATION, + "content_hash": "a" * 64, + TIDB_POINT_ID_FIELD: _KNOWLEDGE_FS_VECTOR_STORE_A, + } + + @staticmethod + @pytest.mark.parametrize("compatible", [False, True]) + def test_tidb_payload_index_creation_race_requires_compatible_index( + tidb_client: TidbHarness, compatible: bool + ) -> None: + client, info = tidb_client + + def create_index(**_kwargs: object) -> None: + if compatible: + info.payload_schema[TIDB_POINT_ID_FIELD] = models.PayloadIndexInfo(data_type="keyword", points=0) + raise RuntimeError("index creation raced or failed") + + client.create_payload_index_hook = create_index + payload = _knowledge_fs_vector_store_request( + "upsert", + points=[ + { + "id": _KNOWLEDGE_FS_VECTOR_STORE_A, + "generation_id": _KNOWLEDGE_FS_VECTOR_STORE_GENERATION, + "content_hash": "a" * 64, + "vector": [1, 0], + } + ], + ) + if compatible: + client.execute(payload) + assert len(client.upsert_calls) == 1 + else: + with pytest.raises(RuntimeError, match="creation"): + client.execute(payload) + assert client.upsert_calls == [] + + @staticmethod + def test_tidb_nonempty_old_collection_is_not_marked_indexed(tidb_client: TidbHarness) -> None: + client, info = tidb_client + info.points_count = 1 + client.execute( + _knowledge_fs_vector_store_request( + "upsert", + points=[ + { + "id": _KNOWLEDGE_FS_VECTOR_STORE_A, + "generation_id": _KNOWLEDGE_FS_VECTOR_STORE_GENERATION, + "content_hash": "a" * 64, + "vector": [1, 0], + } + ], + ), + ) + assert client.create_payload_index_calls == [] + assert len(client.upsert_calls) == 1 + + @staticmethod + @pytest.mark.parametrize("operation", ["search", "upsert"]) + def test_tidb_wrong_id_index_type_fails_closed(tidb_client: TidbHarness, operation: str) -> None: + client, info = tidb_client + info.payload_schema[TIDB_POINT_ID_FIELD] = models.PayloadIndexInfo(data_type="integer", points=0) + payload = ( + _knowledge_fs_vector_store_request("search", ids=[_KNOWLEDGE_FS_VECTOR_STORE_A], query_vector=[1, 0]) + if operation == "search" + else _knowledge_fs_vector_store_request( + "upsert", + points=[ + { + "id": _KNOWLEDGE_FS_VECTOR_STORE_A, + "generation_id": _KNOWLEDGE_FS_VECTOR_STORE_GENERATION, + "content_hash": "a" * 64, + "vector": [1, 0], + } + ], + ) + ) + with pytest.raises(VectorStoreUnavailableError, match="ID index"): + client.execute(payload) + assert client.upsert_calls == [] + assert client.search_calls == [] + + @staticmethod + def test_tidb_indexed_search_uses_supported_filter_and_payload_selection(tidb_client: TidbHarness) -> None: + client, info = tidb_client + info.payload_schema[TIDB_POINT_ID_FIELD] = models.PayloadIndexInfo(data_type="keyword", points=1) + client.search_result = [models.ScoredPoint(id=_KNOWLEDGE_FS_VECTOR_STORE_A, version=0, score=0.6)] + result = client.execute( + _knowledge_fs_vector_store_request( + "search", ids=[_KNOWLEDGE_FS_VECTOR_STORE_A], query_vector=[1, 0], limit=1 + ) + ) + assert result["matches"] == [{"id": _KNOWLEDGE_FS_VECTOR_STORE_A, "score": 0.6}] + query = client.search_calls[-1] + assert query["with_payload"] is True + assert query["with_vectors"] is False + assert query["query_filter"] == models.Filter( + must=[ + models.FieldCondition( + key=TIDB_POINT_ID_FIELD, match=models.MatchAny(any=[_KNOWLEDGE_FS_VECTOR_STORE_A]) + ) + ] + ) + assert client.retrieve_calls == [] + client.search_result = [models.ScoredPoint(id=_KNOWLEDGE_FS_VECTOR_STORE_B, version=0, score=1.0)] + with pytest.raises(VectorStoreUnavailableError, match="unauthorized"): + client.execute( + _knowledge_fs_vector_store_request("search", ids=[_KNOWLEDGE_FS_VECTOR_STORE_A], query_vector=[1, 0]) + ) + + @staticmethod + def test_tidb_unindexed_legacy_search_reads_only_authorized_ids( + client: QdrantClientType, config_overrides: Callable[..., None] + ) -> None: + scope = VectorScope.model_validate(_KNOWLEDGE_FS_VECTOR_STORE_SCOPE) + client.create_collection( + scope.legacy_collection_name, vectors_config=models.VectorParams(size=2, distance="Cosine") + ) + client.upsert( + scope.legacy_collection_name, + points=[ + models.PointStruct(id=_KNOWLEDGE_FS_VECTOR_STORE_A, vector=[3, 4]), + models.PointStruct(id=_KNOWLEDGE_FS_VECTOR_STORE_B, vector=[1, 0]), + ], + ) + config_overrides(VECTOR_STORE="tidb_on_qdrant") + with patch.object(client, "search", side_effect=AssertionError("Legacy search must not query unfiltered")): + result = execute_vector_request( + client, + _knowledge_fs_vector_store_request( + "search", ids=[_KNOWLEDGE_FS_VECTOR_STORE_A], query_vector=[1, 0], limit=1 + ), + ) + assert [match["id"] for match in result["matches"]] == [_KNOWLEDGE_FS_VECTOR_STORE_A] + assert result["matches"][0]["score"] == pytest.approx(0.6) + result = execute_vector_request( + client, + _knowledge_fs_vector_store_request( + "search", + ids=[_KNOWLEDGE_FS_VECTOR_STORE_A, _KNOWLEDGE_FS_VECTOR_STORE_B], + query_vector=[1, 0], + limit=1, + ), + ) + assert result["matches"] == [{"id": _KNOWLEDGE_FS_VECTOR_STORE_B, "score": 1.0}] + + @staticmethod + def test_tidb_unindexed_search_bounds_reads_and_propagates_failures(tidb_client: TidbHarness) -> None: + client, info = tidb_client + info.points_count = 130 + ids = [f"00000000-0000-4000-8000-{n:012d}" for n in range(130)] + + def retrieve(**kwargs: object) -> list[object]: + assert kwargs["with_payload"] is False + assert kwargs["with_vectors"] is True + ids_batch = kwargs["ids"] + assert isinstance(ids_batch, list) + assert len(ids_batch) <= 64 + return [models.Record(id=point_id, vector=[1, 0]) for point_id in ids_batch] + + client.retrieve_hook = retrieve + result = client.execute(_knowledge_fs_vector_store_request("search", ids=ids, query_vector=[1, 0], limit=3)) + assert result["matches"] == [{"id": point_id, "score": 1.0} for point_id in ids[:3]] + assert len(client.retrieve_calls) == 3 + assert client.search_calls == [] + client.retrieve_failure = RuntimeError("read failed") + with pytest.raises(RuntimeError, match="read failed"): + client.execute(_knowledge_fs_vector_store_request("search", ids=ids, query_vector=[1, 0])) + + @staticmethod + @pytest.mark.parametrize("vector", [None, [1], [0, 0], [float("nan"), 0], [1e20, 0]]) + def test_tidb_unindexed_search_rejects_unscorable_vectors( + tidb_client: TidbHarness, vector: list[float] | None + ) -> None: + client, _ = tidb_client + client.retrieve_result = [models.Record(id=_KNOWLEDGE_FS_VECTOR_STORE_A, vector=vector)] + with pytest.raises(VectorStoreUnavailableError, match="dimensions|scored"): + client.execute( + _knowledge_fs_vector_store_request("search", ids=[_KNOWLEDGE_FS_VECTOR_STORE_A], query_vector=[1, 0]) + ) + + @staticmethod + def test_tidb_unindexed_search_rejects_unrequested_points(tidb_client: TidbHarness) -> None: + client, _ = tidb_client + client.retrieve_result = [models.Record(id=_KNOWLEDGE_FS_VECTOR_STORE_B, vector=[1, 0])] + with pytest.raises(VectorStoreUnavailableError, match="unauthorized"): + client.execute( + _knowledge_fs_vector_store_request("search", ids=[_KNOWLEDGE_FS_VECTOR_STORE_A], query_vector=[1, 0]) + ) + + @staticmethod + def test_tidb_indexed_delete_uses_field_filter_and_verifies_absence(tidb_client: TidbHarness) -> None: + client, info = tidb_client + info.payload_schema[TIDB_POINT_ID_FIELD] = models.PayloadIndexInfo(data_type="keyword", points=2) + client.retrieve_result = list[object]() + client.execute( + _knowledge_fs_vector_store_request( + "delete", ids=[_KNOWLEDGE_FS_VECTOR_STORE_A, _KNOWLEDGE_FS_VECTOR_STORE_B] + ) + ) + assert client.delete_selectors[-1] == models.FilterSelector( + filter=models.Filter( + must=[ + models.FieldCondition( + key=TIDB_POINT_ID_FIELD, + match=models.MatchAny(any=[_KNOWLEDGE_FS_VECTOR_STORE_A, _KNOWLEDGE_FS_VECTOR_STORE_B]), + ) + ] + ) + ) + assert len(client.delete_calls) == 1 + client.retrieve_result = [models.Record(id=_KNOWLEDGE_FS_VECTOR_STORE_A)] + with pytest.raises(VectorStoreUnavailableError, match="not yet visible"): + client.execute( + _knowledge_fs_vector_store_request( + "delete", ids=[_KNOWLEDGE_FS_VECTOR_STORE_A, _KNOWLEDGE_FS_VECTOR_STORE_B] + ) + ) + + @staticmethod + def test_tidb_unindexed_delete_uses_single_ids_and_propagates_errors(tidb_client: TidbHarness) -> None: + client, _ = tidb_client + client.retrieve_result = list[object]() + client.execute( + _knowledge_fs_vector_store_request( + "delete", ids=[_KNOWLEDGE_FS_VECTOR_STORE_A, _KNOWLEDGE_FS_VECTOR_STORE_B] + ) + ) + assert client.create_payload_index_calls == [ + { + "collection_name": VectorScope.model_validate(_KNOWLEDGE_FS_VECTOR_STORE_SCOPE).collection_name, + "field_name": "content_hash", + "field_schema": models.PayloadSchemaType.KEYWORD, + "wait": True, + } + ] + from qdrant_client.http.models import PointIdsList as NativePointIdsList + + selectors = [selector for selector in client.delete_selectors if isinstance(selector, NativePointIdsList)] + assert len(selectors) == len(client.delete_selectors) + assert sorted(selector.points for selector in selectors) == [ + [_KNOWLEDGE_FS_VECTOR_STORE_A], + [_KNOWLEDGE_FS_VECTOR_STORE_B], + ] + assert all(call["wait"] for call in client.delete_calls) + client.delete_failure = RuntimeError("delete failed") + with pytest.raises(RuntimeError, match="delete failed"): + client.execute( + _knowledge_fs_vector_store_request( + "delete", ids=[_KNOWLEDGE_FS_VECTOR_STORE_A, _KNOWLEDGE_FS_VECTOR_STORE_B] + ) + ) + + @staticmethod + @pytest.mark.parametrize("remaining", [[], [models.Record(id=_KNOWLEDGE_FS_VECTOR_STORE_A)]]) + def test_tidb_delete_not_found_still_requires_readback(tidb_client: TidbHarness, remaining: list[object]) -> None: + client, info = tidb_client + info.payload_schema[TIDB_POINT_ID_FIELD] = models.PayloadIndexInfo(data_type="keyword", points=1) + client.delete_failure = UnexpectedResponse(404, "missing", b"{}", {}) + client.retrieve_result = remaining + if remaining: + with pytest.raises(VectorStoreUnavailableError, match="not yet visible"): + client.execute(_knowledge_fs_vector_store_request("delete", ids=[_KNOWLEDGE_FS_VECTOR_STORE_A])) + else: + client.execute(_knowledge_fs_vector_store_request("delete", ids=[_KNOWLEDGE_FS_VECTOR_STORE_A])) + assert len(client.retrieve_calls) == 1 + + @staticmethod + @pytest.mark.parametrize("status", [404, 503]) + def test_tidb_delete_readback_only_accepts_missing_collection(tidb_client: TidbHarness, status: int) -> None: + client, info = tidb_client + info.payload_schema[TIDB_POINT_ID_FIELD] = models.PayloadIndexInfo(data_type="keyword", points=1) + client.retrieve_failure = UnexpectedResponse(status, "unavailable", b"{}", {}) + if status == 404: + client.execute(_knowledge_fs_vector_store_request("delete", ids=[_KNOWLEDGE_FS_VECTOR_STORE_A])) + else: + with pytest.raises(UnexpectedResponse): + client.execute(_knowledge_fs_vector_store_request("delete", ids=[_KNOWLEDGE_FS_VECTOR_STORE_A])) diff --git a/api/tests/unit_tests/services/test_knowledge_fs_workflow_runtime.py b/api/tests/unit_tests/services/knowledge_fs/test_workflow_runtime.py similarity index 100% rename from api/tests/unit_tests/services/test_knowledge_fs_workflow_runtime.py rename to api/tests/unit_tests/services/knowledge_fs/test_workflow_runtime.py diff --git a/api/tests/unit_tests/services/test_knowledge_fs_membership_account_integration.py b/api/tests/unit_tests/services/knowledge_fs/test_workspace_members.py similarity index 98% rename from api/tests/unit_tests/services/test_knowledge_fs_membership_account_integration.py rename to api/tests/unit_tests/services/knowledge_fs/test_workspace_members.py index 20da83964fe..9399eac0fdc 100644 --- a/api/tests/unit_tests/services/test_knowledge_fs_membership_account_integration.py +++ b/api/tests/unit_tests/services/knowledge_fs/test_workspace_members.py @@ -20,7 +20,7 @@ from models.knowledge_fs import ( KnowledgeFSLifecycleOutbox, ) from repositories.workspace.workspace_repository import WorkspaceRepository -from tests.unit_tests.services.knowledge_fs_fakes import claims_summary, revoke_payload +from tests.helpers.unit.services.knowledge_fs_fakes import claims_summary, revoke_payload _TABLES = ( Account, diff --git a/api/tests/unit_tests/services/model_provider/test_service.py b/api/tests/unit_tests/services/model_provider/test_service.py new file mode 100644 index 00000000000..4fcd5638c49 --- /dev/null +++ b/api/tests/unit_tests/services/model_provider/test_service.py @@ -0,0 +1,1456 @@ +import types +from types import SimpleNamespace +from typing import Any +from unittest.mock import MagicMock, create_autospec + +import pytest +from sqlalchemy import event +from sqlalchemy.engine import Engine +from sqlalchemy.exc import IntegrityError +from sqlalchemy.orm import Session + +from core.entities.model_entities import ModelStatus, ModelWithProviderEntity, SimpleModelProviderEntity +from core.entities.provider_configuration import ProviderConfigurations +from core.entities.provider_entities import CredentialConfiguration, CustomModelConfiguration +from core.plugin.entities.plugin import PluginInstallationSource +from core.plugin.entities.plugin_daemon import PluginModelProviderBinding +from enums import DeploymentEdition +from graphon.model_runtime.entities.common_entities import I18nObject +from graphon.model_runtime.entities.model_entities import FetchFrom, ModelType, ParameterRule, ParameterType +from graphon.model_runtime.entities.provider_entities import ConfigurateMethod, ProviderEntity +from models.account import Account, TenantAccountRole +from models.provider import ( + Provider, + ProviderCredential, + ProviderModel, + ProviderType, + TenantDefaultModel, + TenantPreferredModelProvider, +) +from services.credentials.query import CredentialQuery, ModelCredentialRecord +from services.errors.app_model_config import ProviderNotFoundError +from services.model_provider import service as service_module +from services.model_provider.service import ModelProviderService, _ProviderSummaryState + + +def _create_service_with_mocked_manager() -> tuple[ModelProviderService, MagicMock]: + manager = MagicMock() + service = ModelProviderService() + service._get_provider_manager = MagicMock(return_value=manager) + return service, manager + + +def _build_available_llm(model: str) -> ModelWithProviderEntity: + provider = ProviderEntity( + provider="langgenius/openai/openai", + label=I18nObject(en_US="OpenAI"), + supported_model_types=[ModelType.LLM], + configurate_methods=[ConfigurateMethod.PREDEFINED_MODEL], + ) + return ModelWithProviderEntity( + model=model, + label=I18nObject(en_US=model), + model_type=ModelType.LLM, + fetch_from=FetchFrom.PREDEFINED_MODEL, + model_properties={}, + status=ModelStatus.ACTIVE, + provider=SimpleModelProviderEntity(provider), + ) + + +def _build_provider_configuration( + *, + provider_name: str = "openai", + supported_model_types: list[ModelType] | None = None, + custom_models: list[Any] | None = None, + custom_config_available: bool = True, +) -> SimpleNamespace: + if supported_model_types is None: + supported_model_types = [ModelType.LLM] + + return SimpleNamespace( + provider=SimpleNamespace( + provider=provider_name, + label=I18nObject(en_US=provider_name), + description=None, + icon_small=None, + icon_small_dark=None, + background=None, + help=None, + supported_model_types=supported_model_types, + configurate_methods=[], + provider_credential_schema=None, + model_credential_schema=None, + ), + preferred_provider_type=ProviderType.CUSTOM, + custom_configuration=SimpleNamespace( + provider=SimpleNamespace( + current_credential_id="cred-1", + current_credential_name="Credential 1", + available_credentials=[], + ), + models=custom_models, + can_added_models=[], + ), + system_configuration=SimpleNamespace(enabled=False, current_quota_type=None, quota_configurations=[]), + is_custom_configuration_available=lambda: custom_config_available, + ) + + +def _build_model_provider_binding( + source: PluginInstallationSource, + *, + installation_id: str = "installation-1", + plugin_id: str = "langgenius/openai", + plugin_unique_identifier: str = "langgenius/openai:1.0.0@checksum", + verified: bool = True, +) -> PluginModelProviderBinding: + return PluginModelProviderBinding( + provider="openai", + installation_id=installation_id, + plugin_id=plugin_id, + plugin_unique_identifier=plugin_unique_identifier, + runtime_type="remote" if source == PluginInstallationSource.Remote else "local", + source=source, + version="1.0.0", + verified=verified, + ) + + +class TestModelProviderServiceCases: + @staticmethod + @pytest.mark.parametrize("has_active_models", [True, False]) + def test_initialize_default_llm_selects_first_active_model_once( + sqlite_session: Session, has_active_models: bool + ) -> None: + sqlite_session.add_all( + [ + TenantDefaultModel( + tenant_id="other-tenant", model_type=ModelType.LLM, provider_name="other", model_name="other-model" + ), + TenantDefaultModel( + tenant_id="tenant-1", + model_type=ModelType.TEXT_EMBEDDING, + provider_name="embedding", + model_name="embedding-model", + ), + ] + ) + sqlite_session.commit() + service, manager = _create_service_with_mocked_manager() + models = [_build_available_llm("gpt-4o-mini"), _build_available_llm("gpt-4o")] if has_active_models else [] + configurations = manager.get_configurations.return_value + configurations.get_models.return_value = models + + service.initialize_default_llm("tenant-1") + service.initialize_default_llm("tenant-1") + + expected_defaults = { + ("other-tenant", ModelType.LLM, "other", "other-model"), + ("tenant-1", ModelType.TEXT_EMBEDDING, "embedding", "embedding-model"), + } + if has_active_models: + expected_defaults.add(("tenant-1", ModelType.LLM, "langgenius/openai/openai", "gpt-4o-mini")) + assert { + (record.tenant_id, record.model_type, record.provider_name, record.model_name) + for record in sqlite_session.query(TenantDefaultModel).all() + } == expected_defaults + assert configurations.get_models.call_count == (1 if has_active_models else 2) + configurations.get_models.assert_called_with(model_type=ModelType.LLM, only_active=True) + + @staticmethod + def test_initialize_default_llm_preserves_existing_selection(sqlite_session: Session) -> None: + saved_default = TenantDefaultModel( + tenant_id="tenant-1", model_type=ModelType.LLM, provider_name="chosen-provider", model_name="chosen-model" + ) + sqlite_session.add(saved_default) + sqlite_session.commit() + service, manager = _create_service_with_mocked_manager() + + service.initialize_default_llm("tenant-1") + + manager.get_configurations.assert_not_called() + sqlite_session.refresh(saved_default) + assert (saved_default.provider_name, saved_default.model_name) == ("chosen-provider", "chosen-model") + + @staticmethod + @pytest.mark.parametrize("selection_time", ["discovery", "insert"]) + def test_initialize_default_llm_preserves_concurrent_selection( + sqlite_session: Session, sqlite_engine: Engine, selection_time: str + ) -> None: + service, manager = _create_service_with_mocked_manager() + selected = False + + def save_selection() -> None: + nonlocal selected + selected = True + sqlite_session.add( + TenantDefaultModel( + tenant_id="tenant-1", + model_type=ModelType.LLM, + provider_name="chosen-provider", + model_name="chosen-model", + ) + ) + sqlite_session.commit() + + def discover_models(*, model_type: ModelType, only_active: bool) -> list[ModelWithProviderEntity]: + assert model_type == ModelType.LLM + assert only_active + if selection_time == "discovery": + save_selection() + return [_build_available_llm("gpt-4o-mini")] + + manager.get_configurations.return_value.get_models.side_effect = discover_models + + def select_before_insert( + _connection: Any, _cursor: Any, statement: str, _parameters: Any, _context: Any, _executemany: bool + ) -> None: + if ( + selection_time == "insert" + and not selected + and statement.startswith("INSERT INTO tenant_default_models") + ): + save_selection() + + event.listen(sqlite_engine, "before_cursor_execute", select_before_insert) + try: + service.initialize_default_llm("tenant-1") + finally: + event.remove(sqlite_engine, "before_cursor_execute", select_before_insert) + + defaults = sqlite_session.query(TenantDefaultModel).all() + assert selected + assert len(defaults) == 1 + assert (defaults[0].provider_name, defaults[0].model_name) == ("chosen-provider", "chosen-model") + + @staticmethod + def test_initialize_default_llm_propagates_database_failure(sqlite_session: Session, sqlite_engine: Engine) -> None: + service, manager = _create_service_with_mocked_manager() + manager.get_configurations.return_value.get_models.return_value = [_build_available_llm("gpt-4o-mini")] + with sqlite_engine.begin() as connection: + connection.exec_driver_sql( + "CREATE TRIGGER reject_default_model BEFORE INSERT ON tenant_default_models " + "BEGIN SELECT RAISE(ABORT, 'default model write rejected'); END" + ) + + with pytest.raises(IntegrityError, match="default model write rejected"): + service.initialize_default_llm("tenant-1") + + assert sqlite_session.query(TenantDefaultModel).count() == 0 + + @staticmethod + def test_available_credentials_use_injected_query_for_admin_viewer() -> None: + service, manager = _create_service_with_mocked_manager() + user = Account(name="Admin", email="admin@example.com") + user.id = "actor" + user.role = TenantAccountRole.ADMIN + query = create_autospec(CredentialQuery, instance=True, spec_set=True) + query.list_models.return_value = [ModelCredentialRecord(id="credential", name="Shared credential")] + + result = service.get_provider_available_credentials("tenant", "openai", user=user, credential_query=query) + + query.list_models.assert_called_once_with(workspace_id="tenant", provider="openai", actor_id="actor") + assert result == [CredentialConfiguration(credential_id="credential", credential_name="Shared credential")] + assert manager.mock_calls == [] + + class TestModelProviderServiceConfiguration: + def test__get_provider_configuration_should_return_configuration_when_provider_exists(self) -> None: + service, manager = _create_service_with_mocked_manager() + provider_configuration = SimpleNamespace(name="provider-config") + manager.get_configurations.return_value = {"openai": provider_configuration} + + result = service._get_provider_configuration(tenant_id="tenant-1", provider="openai") + + assert result is provider_configuration + + def test__get_provider_configuration_should_raise_error_when_provider_is_missing(self) -> None: + service, manager = _create_service_with_mocked_manager() + manager.get_configurations.return_value = {} + + with pytest.raises(ProviderNotFoundError, match="does not exist"): + service._get_provider_configuration(tenant_id="tenant-1", provider="missing") + + def test_get_provider_list_should_filter_by_model_type_and_build_no_configure_status(self) -> None: + service, manager = _create_service_with_mocked_manager() + allowed = _build_provider_configuration( + provider_name="openai", + supported_model_types=[ModelType.LLM], + custom_config_available=False, + ) + filtered = _build_provider_configuration( + provider_name="embedding", + supported_model_types=[ModelType.TEXT_EMBEDDING], + custom_config_available=True, + ) + manager.get_configurations.return_value = {"openai": allowed, "embedding": filtered} + + result = service.get_provider_list(tenant_id="tenant-1", model_type=ModelType.LLM) + + assert len(result) == 1 + assert result[0].provider == "openai" + assert result[0].custom_configuration.status.value == "no-configure" + + def test_get_provider_summary_list_uses_lightweight_state_and_plugin_bindings( + self, monkeypatch: pytest.MonkeyPatch + ) -> None: + service = ModelProviderService() + provider = SimpleNamespace( + provider="langgenius/openai/openai", + label=I18nObject(en_US="OpenAI"), + description=I18nObject(en_US="OpenAI models"), + icon_small=I18nObject(en_US="icon.svg"), + icon_small_dark=I18nObject(en_US="icon-dark.svg"), + supported_model_types=[ModelType.LLM], + configurate_methods=[ConfigurateMethod.PREDEFINED_MODEL], + ) + binding = SimpleNamespace( + provider="openai", + plugin_id="langgenius/openai", + installation_id="installation-1", + plugin_unique_identifier="langgenius/openai:1.2.3@checksum", + runtime_type="local", + source=PluginInstallationSource.Marketplace, + version="1.2.3", + verified=True, + ) + state = _ProviderSummaryState( + has_custom_provider=True, + available_credentials=[ + CredentialConfiguration( + credential_id="credential-1", + credential_name="Production", + ), + CredentialConfiguration( + credential_id="credential-2", + credential_name="Backup", + ), + ], + has_custom_models=True, + current_credential_id="credential-1", + current_credential_name="Production", + current_credential_usable=True, + preferred_provider_type=ProviderType.CUSTOM, + ) + call_order: list[str] = [] + manager_constructor = MagicMock(side_effect=AssertionError("summary must not construct ProviderManager")) + monkeypatch.setattr(service, "_get_provider_manager", manager_constructor) + monkeypatch.setattr( + service_module.PluginService, + "list_model_provider_bindings", + MagicMock(side_effect=lambda *_args, **_kwargs: call_order.append("bindings") or [binding]), + ) + monkeypatch.setattr( + service_module.PluginService, + "fetch_plugin_model_providers", + MagicMock(side_effect=lambda *_args, **_kwargs: call_order.append("providers") or [provider]), + ) + monkeypatch.setattr( + service, + "_load_provider_summary_states", + MagicMock(return_value={provider.provider: state}), + ) + monkeypatch.setattr( + service_module.ext_hosting_provider.hosting_configuration, + "provider_map", + {provider.provider: SimpleNamespace(enabled=True, quotas=[SimpleNamespace()])}, + ) + + providers, plugins = service.get_provider_summary_list(tenant_id="tenant-1") + + assert len(providers) == 1 + assert providers[0].provider == provider.provider + assert providers[0].plugin_id == "langgenius/openai" + assert providers[0].is_configured is True + assert providers[0].custom_configuration.available_credentials == [ + CredentialConfiguration( + credential_id="credential-1", + credential_name="Production", + ), + CredentialConfiguration( + credential_id="credential-2", + credential_name="Backup", + ), + ] + assert providers[0].custom_configuration.has_custom_models is True + assert providers[0].custom_configuration.current_credential_name == "Production" + assert providers[0].custom_configuration.current_credential_usable is True + assert providers[0].system_configuration.enabled is True + assert plugins["langgenius/openai"].installation_id == "installation-1" + assert plugins["langgenius/openai"].version == "1.2.3" + assert call_order == ["bindings", "providers"] + manager_constructor.assert_not_called() + + def test_get_provider_summary_list_enables_system_only_for_verified_hosted_non_package_bindings( + self, + monkeypatch: pytest.MonkeyPatch, + ) -> None: + marketplace_binding = _build_model_provider_binding(PluginInstallationSource.Marketplace) + package_binding = _build_model_provider_binding( + PluginInstallationSource.Package, + installation_id="installation-package", + plugin_id="langgenius/package", + plugin_unique_identifier="langgenius/package:1.0.0@checksum", + ) + unhosted_binding = _build_model_provider_binding( + PluginInstallationSource.Marketplace, + installation_id="installation-unhosted", + plugin_id="langgenius/unhosted", + plugin_unique_identifier="langgenius/unhosted:1.0.0@checksum", + ) + remote_binding = _build_model_provider_binding( + PluginInstallationSource.Remote, + installation_id="installation-remote", + plugin_id="langgenius/remote", + plugin_unique_identifier="langgenius/remote:1.0.0@checksum", + ) + unverified_binding = _build_model_provider_binding( + PluginInstallationSource.Marketplace, + installation_id="installation-unverified", + plugin_id="langgenius/unverified", + plugin_unique_identifier="langgenius/unverified:1.0.0@checksum", + verified=False, + ) + bindings = [ + marketplace_binding, + package_binding, + unhosted_binding, + remote_binding, + unverified_binding, + ] + provider_entities = [ + SimpleNamespace( + provider=f"{binding.plugin_id}/openai", + label=I18nObject(en_US=binding.plugin_id), + description=None, + icon_small=None, + icon_small_dark=None, + supported_model_types=[ModelType.LLM], + configurate_methods=[], + ) + for binding in bindings + ] + monkeypatch.setattr( + service_module.ext_hosting_provider.hosting_configuration, + "provider_map", + { + "langgenius/openai/openai": SimpleNamespace(enabled=True, quotas=[SimpleNamespace()]), + "langgenius/package/openai": SimpleNamespace(enabled=True, quotas=[SimpleNamespace()]), + "langgenius/remote/openai": SimpleNamespace(enabled=True, quotas=[SimpleNamespace()]), + "langgenius/unverified/openai": SimpleNamespace(enabled=True, quotas=[SimpleNamespace()]), + }, + ) + monkeypatch.setattr( + service_module.PluginService, + "list_model_provider_bindings", + MagicMock(return_value=bindings), + ) + monkeypatch.setattr( + service_module.PluginService, + "fetch_plugin_model_providers", + MagicMock(return_value=provider_entities), + ) + monkeypatch.setattr(ModelProviderService, "_load_provider_summary_states", MagicMock(return_value={})) + monkeypatch.setattr(service_module, "is_filtered", MagicMock(return_value=False)) + + providers, _ = ModelProviderService().get_provider_summary_list("tenant-1") + + assert {provider.provider: provider.system_configuration.enabled for provider in providers} == { + "langgenius/openai/openai": True, + "langgenius/package/openai": False, + "langgenius/unhosted/openai": False, + "langgenius/remote/openai": True, + "langgenius/unverified/openai": False, + } + + def test_model_provider_binding_without_verified_field_fails_closed(self) -> None: + binding = PluginModelProviderBinding.model_validate( + { + "provider": "openai", + "installation_id": "installation-1", + "plugin_id": "langgenius/openai", + "plugin_unique_identifier": "langgenius/openai:1.0.0@checksum", + "runtime_type": "local", + "source": "marketplace", + "version": "1.0.0", + } + ) + + assert binding.verified is False + + def test_get_provider_summary_list_returns_all_unique_provider_metadata( + self, monkeypatch: pytest.MonkeyPatch + ) -> None: + service = ModelProviderService() + llm_provider = SimpleNamespace( + provider="langgenius/openai/openai", + label=I18nObject(en_US="OpenAI"), + description=None, + icon_small=None, + icon_small_dark=None, + supported_model_types=[ModelType.LLM], + configurate_methods=[], + ) + embedding_provider = SimpleNamespace( + provider="langgenius/embedding/embedding", + label=I18nObject(en_US="Embedding"), + description=None, + icon_small=None, + icon_small_dark=None, + supported_model_types=[ModelType.TEXT_EMBEDDING], + configurate_methods=[], + ) + llm_binding = SimpleNamespace( + provider="openai", + plugin_id="langgenius/openai", + installation_id="installation-openai", + plugin_unique_identifier="langgenius/openai:1.0.0@checksum", + runtime_type="local", + source=service_module.PluginInstallationSource.Marketplace, + version="1.0.0", + verified=False, + ) + embedding_binding = SimpleNamespace( + provider="embedding", + plugin_id="langgenius/embedding", + installation_id="installation-embedding", + plugin_unique_identifier="langgenius/embedding:1.0.0@checksum", + runtime_type="local", + source=service_module.PluginInstallationSource.Marketplace, + version="1.0.0", + verified=False, + ) + monkeypatch.setattr( + service_module.PluginService, + "list_model_provider_bindings", + MagicMock(return_value=[llm_binding, embedding_binding]), + ) + monkeypatch.setattr( + service_module.PluginService, + "fetch_plugin_model_providers", + MagicMock(return_value=[llm_provider, llm_provider, embedding_provider]), + ) + monkeypatch.setattr(service, "_load_provider_summary_states", MagicMock(return_value={})) + monkeypatch.setattr(service_module, "is_filtered", MagicMock(return_value=False)) + + providers, plugins = service.get_provider_summary_list(tenant_id="tenant-1") + + assert [provider.provider for provider in providers] == [ + "langgenius/openai/openai", + "langgenius/embedding/embedding", + ] + assert providers[0].is_configured is False + assert providers[0].custom_configuration.status.value == "no-configure" + assert providers[0].custom_configuration.has_custom_models is False + assert providers[0].custom_configuration.available_credentials == [] + assert set(plugins) == {"langgenius/openai", "langgenius/embedding"} + + def test_preferred_provider_fallback_uses_custom_presence_not_configuration_status( + self, monkeypatch: pytest.MonkeyPatch + ) -> None: + from tests.helpers.unit.config_override import apply_config_overrides + + apply_config_overrides(monkeypatch, DEPLOYMENT_EDITION=DeploymentEdition.COMMUNITY) + state = _ProviderSummaryState(has_custom_provider=True) + + preferred_provider_type = ModelProviderService._get_preferred_provider_type( + state, + custom_present=True, + system_enabled=True, + ) + + assert preferred_provider_type == ProviderType.CUSTOM + + def test_load_provider_summary_states_reads_only_lightweight_columns( + self, + sqlite_engine: Engine, + sqlite_session: Session, + ) -> None: + canonical_provider = "langgenius/openai/openai" + legacy_credential = ProviderCredential( + tenant_id="tenant-1", + provider_name="openai", + credential_name="Legacy", + encrypted_config="legacy-secret", + ) + legacy_credential.id = "credential-legacy" + current_credential = ProviderCredential( + tenant_id="tenant-1", + provider_name=canonical_provider, + credential_name="Production", + encrypted_config="production-secret", + ) + current_credential.id = "credential-current" + foreign_credential = ProviderCredential( + tenant_id="tenant-2", + provider_name=canonical_provider, + credential_name="Foreign", + encrypted_config="foreign-secret", + ) + foreign_credential.id = "credential-foreign" + sqlite_session.add_all( + [ + legacy_credential, + current_credential, + foreign_credential, + Provider( + tenant_id="tenant-1", + provider_name="openai", + provider_type=ProviderType.CUSTOM, + is_valid=True, + credential_id=legacy_credential.id, + ), + Provider( + tenant_id="tenant-1", + provider_name=canonical_provider, + provider_type=ProviderType.CUSTOM, + is_valid=True, + credential_id=current_credential.id, + ), + Provider( + tenant_id="tenant-2", + provider_name=canonical_provider, + provider_type=ProviderType.CUSTOM, + is_valid=True, + credential_id=foreign_credential.id, + ), + ProviderModel( + tenant_id="tenant-1", + provider_name="openai", + model_name="gpt-4o", + model_type=ModelType.LLM, + is_valid=True, + ), + ProviderModel( + tenant_id="tenant-2", + provider_name="langgenius/foreign/foreign", + model_name="foreign-model", + model_type=ModelType.LLM, + is_valid=True, + ), + TenantPreferredModelProvider( + tenant_id="tenant-1", + provider_name=canonical_provider, + preferred_provider_type=ProviderType.SYSTEM, + ), + TenantPreferredModelProvider( + tenant_id="tenant-2", + provider_name=canonical_provider, + preferred_provider_type=ProviderType.CUSTOM, + ), + ] + ) + sqlite_session.commit() + statements: list[str] = [] + + def capture_statement(_connection, _cursor, statement, _parameters, _context, _executemany) -> None: + statements.append(statement) + + event.listen(sqlite_engine, "before_cursor_execute", capture_statement) + try: + states = ModelProviderService._load_provider_summary_states("tenant-1") + finally: + event.remove(sqlite_engine, "before_cursor_execute", capture_statement) + + state = states[canonical_provider] + assert state.has_custom_provider is True + assert state.available_credentials == [ + CredentialConfiguration( + credential_id="credential-legacy", + credential_name="Legacy", + ), + CredentialConfiguration( + credential_id="credential-current", + credential_name="Production", + ), + ] + assert state.has_custom_models is True + assert state.current_credential_id == "credential-current" + assert state.current_credential_name == "Production" + assert state.current_credential_usable is True + assert state.preferred_provider_type == ProviderType.SYSTEM + + assert len(statements) == 4 + assert all("encrypted_config" not in statement for statement in statements) + assert "count(" not in statements[1].lower() + assert "provider_credentials.id" in statements[1] + assert "provider_credentials.credential_name" in statements[1] + assert "ORDER BY provider_credentials.created_at DESC, provider_credentials.id DESC" in statements[1] + assert "provider_model_credentials" in statements[2] + + def test_list_models_should_pass_filters_through(self, monkeypatch: pytest.MonkeyPatch) -> None: + service, manager = _create_service_with_mocked_manager() + filters: list[dict[str, object]] = [] + + def get_models(_self: ProviderConfigurations, **kwargs: object) -> list[ModelWithProviderEntity]: + filters.append(kwargs) + return [] + + monkeypatch.setattr(ProviderConfigurations, "get_models", get_models) + manager.get_configurations.return_value = ProviderConfigurations(tenant_id="tenant-1") + + assert service.list_models("tenant-1", provider=None, model_type=ModelType.LLM) == [] + assert filters == [{"provider": None, "model_type": ModelType.LLM}] + + def test_get_models_by_provider_should_wrap_model_entities_with_tenant_context(self) -> None: + service, manager = _create_service_with_mocked_manager() + + class _Model: + def __init__(self, model_name: str) -> None: + self.model_name = model_name + + def model_dump(self) -> dict[str, Any]: + return { + "model": self.model_name, + "label": {"en_US": self.model_name}, + "model_type": ModelType.LLM, + "features": [], + "fetch_from": FetchFrom.PREDEFINED_MODEL, + "model_properties": {}, + "deprecated": False, + "status": ModelStatus.ACTIVE, + "load_balancing_enabled": False, + "has_invalid_load_balancing_configs": False, + "provider": { + "provider": "openai", + "label": {"en_US": "OpenAI"}, + "icon_small": None, + "icon_small_dark": None, + "supported_model_types": [ModelType.LLM], + }, + } + + provider_configurations = SimpleNamespace( + get_models=MagicMock(return_value=[_Model("gpt-4o"), _Model("gpt-4o-mini")]) + ) + manager.get_configurations.return_value = provider_configurations + + result = service.get_models_by_provider(tenant_id="tenant-1", provider="openai") + + assert len(result) == 2 + assert result[0].model == "gpt-4o" + assert result[1].provider.provider == "openai" + provider_configurations.get_models.assert_called_once_with(provider="openai") + + class TestModelProviderServiceDelegation: + @pytest.mark.parametrize( + ("method_name", "method_kwargs", "provider_method_name", "provider_call_kwargs", "provider_return"), + [ + ( + "get_provider_credential", + {"tenant_id": "tenant-1", "provider": "openai", "credential_id": "cred-1"}, + "get_provider_credential", + {"credential_id": "cred-1"}, + {"token": "abc"}, + ), + ( + "validate_provider_credentials", + {"tenant_id": "tenant-1", "provider": "openai", "credentials": {"token": "abc"}}, + "validate_provider_credentials", + ({"token": "abc"},), + None, + ), + ( + "create_provider_credential", + { + "tenant_id": "tenant-1", + "provider": "openai", + "credentials": {"token": "abc"}, + "credential_name": "A", + }, + "create_provider_credential", + ({"token": "abc"}, "A"), + "new-credential-id", + ), + ( + "update_provider_credential", + { + "tenant_id": "tenant-1", + "provider": "openai", + "credentials": {"token": "abc"}, + "credential_id": "cred-1", + "credential_name": "B", + }, + "update_provider_credential", + {"credential_id": "cred-1", "credentials": {"token": "abc"}, "credential_name": "B"}, + None, + ), + ( + "remove_provider_credential", + {"tenant_id": "tenant-1", "provider": "openai", "credential_id": "cred-1"}, + "delete_provider_credential", + {"credential_id": "cred-1"}, + None, + ), + ( + "switch_active_provider_credential", + {"tenant_id": "tenant-1", "provider": "openai", "credential_id": "cred-1"}, + "switch_active_provider_credential", + {"credential_id": "cred-1"}, + None, + ), + ], + ) + def test_provider_credential_methods_should_delegate_to_provider_configuration( + self, + method_name: str, + method_kwargs: dict[str, Any], + provider_method_name: str, + provider_call_kwargs: Any, + provider_return: Any, + monkeypatch: pytest.MonkeyPatch, + ) -> None: + service = ModelProviderService() + provider_configuration = MagicMock() + getattr(provider_configuration, provider_method_name).return_value = provider_return + get_provider_config_mock = MagicMock(return_value=provider_configuration) + monkeypatch.setattr(service, "_get_provider_configuration", get_provider_config_mock) + + result = getattr(service, method_name)(**method_kwargs) + + get_provider_config_mock.assert_called_once_with("tenant-1", "openai") + provider_method = getattr(provider_configuration, provider_method_name) + match provider_call_kwargs: + case tuple(): + provider_method.assert_called_once_with(*provider_call_kwargs) + case dict(): + provider_method.assert_called_once_with(**provider_call_kwargs) + case _: + provider_method.assert_called_once_with(provider_call_kwargs) + if method_name == "get_provider_credential": + assert result == {"token": "abc"} + if method_name == "create_provider_credential": + assert result == "new-credential-id" + + @pytest.mark.parametrize( + ("method_name", "method_kwargs", "provider_method_name", "expected_kwargs", "provider_return"), + [ + ( + "get_model_credential", + { + "tenant_id": "tenant-1", + "provider": "openai", + "model_type": ModelType.LLM, + "model": "gpt-4o", + "credential_id": "cred-1", + }, + "get_custom_model_credential", + {"model_type": ModelType.LLM, "model": "gpt-4o", "credential_id": "cred-1"}, + {"api_key": "x"}, + ), + ( + "validate_model_credentials", + { + "tenant_id": "tenant-1", + "provider": "openai", + "model_type": ModelType.LLM, + "model": "gpt-4o", + "credentials": {"api_key": "x"}, + }, + "validate_custom_model_credentials", + {"model_type": ModelType.LLM, "model": "gpt-4o", "credentials": {"api_key": "x"}}, + None, + ), + ( + "create_model_credential", + { + "tenant_id": "tenant-1", + "provider": "openai", + "model_type": ModelType.LLM, + "model": "gpt-4o", + "credentials": {"api_key": "x"}, + "credential_name": "cred-a", + }, + "create_custom_model_credential", + { + "model_type": ModelType.LLM, + "model": "gpt-4o", + "credentials": {"api_key": "x"}, + "credential_name": "cred-a", + }, + "new-credential-id", + ), + ( + "update_model_credential", + { + "tenant_id": "tenant-1", + "provider": "openai", + "model_type": ModelType.LLM, + "model": "gpt-4o", + "credentials": {"api_key": "x"}, + "credential_id": "cred-1", + "credential_name": "cred-b", + }, + "update_custom_model_credential", + { + "model_type": ModelType.LLM, + "model": "gpt-4o", + "credentials": {"api_key": "x"}, + "credential_id": "cred-1", + "credential_name": "cred-b", + }, + None, + ), + ( + "remove_model_credential", + { + "tenant_id": "tenant-1", + "provider": "openai", + "model_type": ModelType.LLM, + "model": "gpt-4o", + "credential_id": "cred-1", + }, + "delete_custom_model_credential", + {"model_type": ModelType.LLM, "model": "gpt-4o", "credential_id": "cred-1"}, + None, + ), + ( + "switch_active_custom_model_credential", + { + "tenant_id": "tenant-1", + "provider": "openai", + "model_type": ModelType.LLM, + "model": "gpt-4o", + "credential_id": "cred-1", + }, + "switch_custom_model_credential", + {"model_type": ModelType.LLM, "model": "gpt-4o", "credential_id": "cred-1"}, + None, + ), + ( + "add_model_credential_to_model_list", + { + "tenant_id": "tenant-1", + "provider": "openai", + "model_type": ModelType.LLM, + "model": "gpt-4o", + "credential_id": "cred-1", + }, + "add_model_credential_to_model", + {"model_type": ModelType.LLM, "model": "gpt-4o", "credential_id": "cred-1"}, + None, + ), + ( + "remove_model", + { + "tenant_id": "tenant-1", + "provider": "openai", + "model_type": ModelType.LLM, + "model": "gpt-4o", + }, + "delete_custom_model", + {"model_type": ModelType.LLM, "model": "gpt-4o"}, + None, + ), + ], + ) + def test_custom_model_methods_should_convert_model_type_and_delegate( + self, + method_name: str, + method_kwargs: dict[str, Any], + provider_method_name: str, + expected_kwargs: dict[str, Any], + provider_return: Any, + monkeypatch: pytest.MonkeyPatch, + ) -> None: + service = ModelProviderService() + provider_configuration = MagicMock() + getattr(provider_configuration, provider_method_name).return_value = provider_return + get_provider_config_mock = MagicMock(return_value=provider_configuration) + monkeypatch.setattr(service, "_get_provider_configuration", get_provider_config_mock) + + result = getattr(service, method_name)(**method_kwargs) + + get_provider_config_mock.assert_called_once_with("tenant-1", "openai") + getattr(provider_configuration, provider_method_name).assert_called_once_with(**expected_kwargs) + if method_name == "get_model_credential": + assert result == {"api_key": "x"} + if method_name == "create_model_credential": + assert result == "new-credential-id" + + @pytest.mark.parametrize( + ("method_name", "method_kwargs", "provider_method_name", "expected_kwargs"), + [ + ( + "get_model_credential", + { + "tenant_id": "tenant-1", + "provider": "openai", + "model_type": "llm", + "model": "gpt-4o", + "credential_id": "cred-1", + }, + "get_custom_model_credential", + {"model_type": ModelType.LLM, "model": "gpt-4o", "credential_id": "cred-1"}, + ), + ( + "create_model_credential", + { + "tenant_id": "tenant-1", + "provider": "openai", + "model_type": "llm", + "model": "gpt-4o", + "credentials": {"api_key": "x"}, + "credential_name": "cred-a", + }, + "create_custom_model_credential", + { + "model_type": ModelType.LLM, + "model": "gpt-4o", + "credentials": {"api_key": "x"}, + "credential_name": "cred-a", + }, + ), + ( + "remove_model", + { + "tenant_id": "tenant-1", + "provider": "openai", + "model_type": "llm", + "model": "gpt-4o", + }, + "delete_custom_model", + {"model_type": ModelType.LLM, "model": "gpt-4o"}, + ), + ], + ) + def test_custom_model_methods_use_model_type_constructor_directly( + self, + method_name: str, + method_kwargs: dict[str, Any], + provider_method_name: str, + expected_kwargs: dict[str, Any], + monkeypatch: pytest.MonkeyPatch, + ) -> None: + service = ModelProviderService() + provider_configuration = MagicMock() + get_provider_config_mock = MagicMock(return_value=provider_configuration) + monkeypatch.setattr(service, "_get_provider_configuration", get_provider_config_mock) + + getattr(service, method_name)(**method_kwargs) + + get_provider_config_mock.assert_called_once_with("tenant-1", "openai") + getattr(provider_configuration, provider_method_name).assert_called_once_with(**expected_kwargs) + + class TestModelProviderServiceListingsAndDefaults: + def test_get_models_by_model_type_should_group_active_non_deprecated_models(self) -> None: + service, manager = _create_service_with_mocked_manager() + openai_provider = SimpleNamespace( + provider="openai", + label=I18nObject(en_US="OpenAI"), + icon_small=None, + icon_small_dark=None, + ) + anthropic_provider = SimpleNamespace( + provider="anthropic", + label=I18nObject(en_US="Anthropic"), + icon_small=None, + icon_small_dark=None, + ) + models = [ + SimpleNamespace( + provider=openai_provider, + model="gpt-4o", + label=I18nObject(en_US="GPT-4o"), + model_type=ModelType.LLM, + features=[], + fetch_from=FetchFrom.PREDEFINED_MODEL, + model_properties={}, + status=ModelStatus.ACTIVE, + load_balancing_enabled=False, + deprecated=False, + ), + SimpleNamespace( + provider=openai_provider, + model="old-openai", + label=I18nObject(en_US="Old OpenAI"), + model_type=ModelType.LLM, + features=[], + fetch_from=FetchFrom.PREDEFINED_MODEL, + model_properties={}, + status=ModelStatus.ACTIVE, + load_balancing_enabled=False, + deprecated=True, + ), + SimpleNamespace( + provider=anthropic_provider, + model="old-anthropic", + label=I18nObject(en_US="Old Anthropic"), + model_type=ModelType.LLM, + features=[], + fetch_from=FetchFrom.PREDEFINED_MODEL, + model_properties={}, + status=ModelStatus.ACTIVE, + load_balancing_enabled=False, + deprecated=True, + ), + ] + provider_configurations = SimpleNamespace(get_models=MagicMock(return_value=models)) + manager.get_configurations.return_value = provider_configurations + + result = service.get_models_by_model_type(tenant_id="tenant-1", model_type=ModelType.LLM) + + provider_configurations.get_models.assert_called_once_with(model_type=ModelType.LLM, only_active=True) + assert len(result) == 1 + assert result[0].provider == "openai" + assert len(result[0].models) == 1 + assert result[0].models[0].model == "gpt-4o" + + @pytest.mark.parametrize( + ("credentials", "schema", "expected_count"), + [ + (None, None, 0), + ({"api_key": "x"}, None, 0), + ( + {"api_key": "x"}, + SimpleNamespace( + parameter_rules=[ + ParameterRule( + name="temperature", + label=I18nObject(en_US="Temperature"), + type=ParameterType.FLOAT, + ) + ] + ), + 1, + ), + ], + ) + def test_get_model_parameter_rules_should_handle_missing_credentials_and_schema( + self, + credentials: dict[str, Any] | None, + schema: Any, + expected_count: int, + monkeypatch: pytest.MonkeyPatch, + ) -> None: + service = ModelProviderService() + provider_configuration = MagicMock() + provider_configuration.get_current_credentials.return_value = credentials + provider_configuration.get_model_schema.return_value = schema + monkeypatch.setattr(service, "_get_provider_configuration", MagicMock(return_value=provider_configuration)) + + result = service.get_model_parameter_rules(tenant_id="tenant-1", provider="openai", model="gpt-4o") + + assert len(result) == expected_count + provider_configuration.get_current_credentials.assert_called_once_with( + model_type=ModelType.LLM, + model="gpt-4o", + ) + if credentials: + provider_configuration.get_model_schema.assert_called_once_with( + model_type=ModelType.LLM, + model="gpt-4o", + credentials=credentials, + ) + else: + provider_configuration.get_model_schema.assert_not_called() + + def test_get_default_model_selection_uses_saved_workspace_choice(self, sqlite_session: Session) -> None: + sqlite_session.add( + TenantDefaultModel( + tenant_id="tenant-1", + model_type=ModelType.LLM, + provider_name="langgenius/openai/openai", + model_name="gpt-4o", + ) + ) + sqlite_session.commit() + service, manager = _create_service_with_mocked_manager() + + result = service.get_default_model_selection("tenant-1", ModelType.LLM, session=sqlite_session) + + assert result == ("langgenius/openai/openai", "gpt-4o") + manager.get_configurations.assert_not_called() + + def test_get_default_model_selection_uses_first_active_model_without_writing( + self, sqlite_session: Session + ) -> None: + service, manager = _create_service_with_mocked_manager() + configurations = manager.get_configurations.return_value + configurations.get_models.return_value = [ + SimpleNamespace(model="gpt-4o", provider=SimpleNamespace(provider="langgenius/openai/openai")) + ] + + result = service.get_default_model_selection("tenant-1", ModelType.LLM, session=sqlite_session) + + assert result == ("langgenius/openai/openai", "gpt-4o") + configurations.get_models.assert_called_once_with(model_type=ModelType.LLM, only_active=True) + assert sqlite_session.query(TenantDefaultModel).filter_by(tenant_id="tenant-1").first() is None + + def test_get_default_model_selection_returns_none_when_provider_discovery_fails( + self, sqlite_session: Session + ) -> None: + service, manager = _create_service_with_mocked_manager() + manager.get_configurations.side_effect = RuntimeError("provider unavailable") + + result = service.get_default_model_selection("tenant-1", ModelType.LLM, session=sqlite_session) + + assert result is None + + def test_get_default_model_selection_returns_none_when_no_active_models(self, sqlite_session: Session) -> None: + service, manager = _create_service_with_mocked_manager() + manager.get_configurations.return_value.get_models.return_value = [] + + result = service.get_default_model_selection("tenant-1", ModelType.LLM, session=sqlite_session) + + assert result is None + manager.get_configurations.return_value.get_models.assert_called_once_with( + model_type=ModelType.LLM, only_active=True + ) + + def test_get_default_model_of_model_type_should_return_response_when_manager_returns_model(self) -> None: + service, manager = _create_service_with_mocked_manager() + manager.get_default_model.return_value = SimpleNamespace( + model="gpt-4o", + model_type=ModelType.LLM, + provider=SimpleNamespace( + provider="openai", + label=I18nObject(en_US="OpenAI"), + icon_small=None, + supported_model_types=[ModelType.LLM], + ), + ) + + result = service.get_default_model_of_model_type(tenant_id="tenant-1", model_type=ModelType.LLM) + + assert result is not None + assert result.model == "gpt-4o" + assert result.provider.provider == "openai" + manager.get_default_model.assert_called_once_with(tenant_id="tenant-1", model_type=ModelType.LLM) + + def test_get_default_model_of_model_type_should_return_none_when_manager_returns_none(self) -> None: + service, manager = _create_service_with_mocked_manager() + manager.get_default_model.return_value = None + + result = service.get_default_model_of_model_type(tenant_id="tenant-1", model_type=ModelType.LLM) + + assert result is None + + def test_get_default_model_of_model_type_should_return_none_when_no_configuration_is_saved(self) -> None: + service, manager = _create_service_with_mocked_manager() + manager.get_default_model.side_effect = RuntimeError("boom") + + result = service.get_default_model_of_model_type(tenant_id="tenant-1", model_type=ModelType.LLM) + + assert result is None + + @pytest.mark.parametrize( + "error", + [ValueError("Invalid provider: langgenius/openai/openai"), RuntimeError("Plugin daemon unavailable")], + ) + def test_get_default_model_of_model_type_should_preserve_saved_configuration_on_provider_failure( + self, sqlite_session: Session, error: Exception + ) -> None: + saved_model = TenantDefaultModel( + tenant_id="tenant-1", + model_type=ModelType.LLM, + provider_name="langgenius/openai/openai", + model_name="gpt-4o", + ) + sqlite_session.add(saved_model) + sqlite_session.commit() + service, manager = _create_service_with_mocked_manager() + manager.get_default_model.side_effect = error + + result = service.get_default_model_of_model_type(tenant_id="tenant-1", model_type=ModelType.LLM) + + assert result is not None + assert result.model == "gpt-4o" + assert result.model_type == ModelType.LLM + assert result.provider.provider == "langgenius/openai/openai" + assert result.provider.tenant_id == "tenant-1" + assert result.provider.label.en_us == "langgenius/openai/openai" + assert result.provider.label.zh_hans == "langgenius/openai/openai" + assert result.provider.icon_small is None + assert result.provider.supported_model_types == [] + assert result.model_dump(mode="json")["model_type"] == "llm" + sqlite_session.refresh(saved_model) + assert saved_model.model_name == "gpt-4o" + assert saved_model.provider_name == "langgenius/openai/openai" + + @pytest.mark.parametrize( + ("tenant_id", "model_type"), + [("other-tenant", ModelType.LLM), ("tenant-1", ModelType.TEXT_EMBEDDING)], + ) + def test_get_default_model_of_model_type_should_not_fall_back_to_another_tenant_or_model_type( + self, sqlite_session: Session, tenant_id: str, model_type: ModelType + ) -> None: + sqlite_session.add( + TenantDefaultModel( + tenant_id=tenant_id, + model_type=model_type, + provider_name="langgenius/openai/openai", + model_name="configured-model", + ) + ) + sqlite_session.commit() + service, manager = _create_service_with_mocked_manager() + manager.get_default_model.side_effect = ValueError("Invalid provider: langgenius/openai/openai") + + result = service.get_default_model_of_model_type(tenant_id="tenant-1", model_type=ModelType.LLM) + + assert result is None + + def test_update_default_model_of_model_type_should_delegate_to_provider_manager(self) -> None: + service, manager = _create_service_with_mocked_manager() + + service.update_default_model_of_model_type( + tenant_id="tenant-1", + model_type=ModelType.LLM, + provider="openai", + model="gpt-4o", + ) + + manager.update_default_model_record.assert_called_once_with( + tenant_id="tenant-1", + model_type=ModelType.LLM, + provider="openai", + model="gpt-4o", + ) + + def test_get_model_provider_icon_should_fetch_icon_bytes_from_factory( + self, + monkeypatch: pytest.MonkeyPatch, + ) -> None: + service = ModelProviderService() + factory_instance = MagicMock() + factory_instance.get_provider_icon.return_value = (b"icon-bytes", "image/png") + factory_constructor = MagicMock(return_value=factory_instance) + monkeypatch.setattr(service_module, "create_plugin_model_provider_factory", factory_constructor) + + result = service.get_model_provider_icon( + tenant_id="tenant-1", + provider="openai", + icon_type="icon_small", + lang="en_US", + ) + + factory_constructor.assert_called_once_with(tenant_id="tenant-1") + factory_instance.get_provider_icon.assert_called_once_with("openai", "icon_small", "en_US") + assert result == (b"icon-bytes", "image/png") + + def test_switch_preferred_provider_should_convert_enum_and_delegate( + self, + monkeypatch: pytest.MonkeyPatch, + ) -> None: + service = ModelProviderService() + provider_configuration = MagicMock() + monkeypatch.setattr(service, "_get_provider_configuration", MagicMock(return_value=provider_configuration)) + + service.switch_preferred_provider( + tenant_id="tenant-1", + provider="openai", + preferred_provider_type=ProviderType.SYSTEM.value, + ) + + provider_configuration.switch_preferred_provider_type.assert_called_once_with(ProviderType.SYSTEM) + + @pytest.mark.parametrize( + ("method_name", "provider_method_name"), + [ + ("enable_model", "enable_model"), + ("disable_model", "disable_model"), + ], + ) + def test_model_enablement_methods_should_convert_model_type_and_delegate( + self, + method_name: str, + provider_method_name: str, + monkeypatch: pytest.MonkeyPatch, + ) -> None: + service = ModelProviderService() + provider_configuration = MagicMock() + monkeypatch.setattr(service, "_get_provider_configuration", MagicMock(return_value=provider_configuration)) + + getattr(service, method_name)( + tenant_id="tenant-1", + provider="openai", + model="gpt-4o", + model_type=ModelType.LLM, + ) + + getattr(provider_configuration, provider_method_name).assert_called_once_with( + model="gpt-4o", + model_type=ModelType.LLM, + ) + + +class _FakeConfigurations: + def __init__(self, provider_configuration: types.SimpleNamespace) -> None: + self._provider_configuration = provider_configuration + + def values(self) -> list[types.SimpleNamespace]: + return [self._provider_configuration] + + +class TestModelProviderServiceSanitizationCases: + @staticmethod + @pytest.fixture + def service_with_fake_configurations(monkeypatch: pytest.MonkeyPatch) -> ModelProviderService: + # Build a fake provider schema with minimal fields used by ProviderResponse + fake_provider = types.SimpleNamespace( + provider="langgenius/openai_api_compatible/openai_api_compatible", + label=I18nObject(en_US="OpenAI API Compatible", zh_Hans="OpenAI API Compatible"), + description=None, + icon_small=None, + icon_small_dark=None, + background=None, + help=None, + supported_model_types=[ModelType.LLM], + configurate_methods=[ConfigurateMethod.CUSTOMIZABLE_MODEL], + provider_credential_schema=None, + model_credential_schema=None, + ) + + # Include decrypted credentials to simulate the leak source + custom_model = CustomModelConfiguration( + model="gpt-4o-mini", + model_type=ModelType.LLM, + credentials={"api_key": "sk-plain-text", "endpoint": "https://example.com"}, + current_credential_id="cred-1", + current_credential_name="API KEY 1", + available_model_credentials=[], + unadded_to_model_list=False, + ) + + fake_custom_provider = types.SimpleNamespace( + current_credential_id="cred-1", + current_credential_name="API KEY 1", + available_credentials=[CredentialConfiguration(credential_id="cred-1", credential_name="API KEY 1")], + ) + + fake_custom_configuration = types.SimpleNamespace( + provider=fake_custom_provider, models=[custom_model], can_added_models=[] + ) + + fake_system_configuration = types.SimpleNamespace( + enabled=False, current_quota_type=None, quota_configurations=[] + ) + + fake_provider_configuration = types.SimpleNamespace( + provider=fake_provider, + preferred_provider_type=ProviderType.CUSTOM, + custom_configuration=fake_custom_configuration, + system_configuration=fake_system_configuration, + is_custom_configuration_available=lambda: True, + ) + + class _FakeProviderManager: + def get_configurations(self, _tenant_id: str) -> _FakeConfigurations: + return _FakeConfigurations(fake_provider_configuration) + + svc = ModelProviderService() + monkeypatch.setattr(svc, "_get_provider_manager", lambda _tenant_id: _FakeProviderManager()) + return svc + + @staticmethod + def test_get_provider_list_strips_credentials(service_with_fake_configurations: ModelProviderService) -> None: + providers = service_with_fake_configurations.get_provider_list(tenant_id="tenant-1", model_type=None) + + assert len(providers) == 1 + custom_models = providers[0].custom_configuration.custom_models + + assert custom_models is not None + assert len(custom_models) == 1 + # The sanitizer should drop credentials in list response + assert custom_models[0].credentials is None diff --git a/api/tests/unit_tests/services/plugin/test_datasource_backwards_invocation.py b/api/tests/unit_tests/services/plugin/test_datasource_backwards_invocation.py index 8683899efd1..ffdb067bd59 100644 --- a/api/tests/unit_tests/services/plugin/test_datasource_backwards_invocation.py +++ b/api/tests/unit_tests/services/plugin/test_datasource_backwards_invocation.py @@ -18,7 +18,7 @@ from core.datasource.entities.datasource_entities import ( from core.plugin.entities.request import RequestInvokeDatasource from models.account import Tenant from services.plugin.datasource_backwards_invocation import PluginDatasourceBackwardsInvocation -from tests.unit_tests.knowledge_fs_core_fakes import Marker +from tests.helpers.unit.knowledge_fs_core_fakes import Marker def _tenant() -> Tenant: diff --git a/api/tests/unit_tests/services/plugin/test_model_dispatch_binding.py b/api/tests/unit_tests/services/plugin/test_model_dispatch_binding.py deleted file mode 100644 index 8157230ac0f..00000000000 --- a/api/tests/unit_tests/services/plugin/test_model_dispatch_binding.py +++ /dev/null @@ -1,100 +0,0 @@ -"""Model dispatches reuse prepared package identities without discovery I/O.""" - -import json - -import httpx -import pytest - -from core.plugin.entities.plugin import PluginInstallationSource -from core.plugin.entities.plugin_daemon import PluginModelProviderDeclaration -from core.plugin.impl import base as client_module -from core.plugin.impl.model import PluginModelClient -from core.plugin.impl.model_runtime_factory import create_model_type_instance, create_plugin_model_runtime -from core.plugin.plugin_service import PluginService -from graphon.model_runtime.entities.common_entities import I18nObject -from graphon.model_runtime.entities.llm_entities import LLMResultChunk, LLMResultChunkDelta -from graphon.model_runtime.entities.message_entities import AssistantPromptMessage -from graphon.model_runtime.entities.model_entities import ModelType -from graphon.model_runtime.entities.provider_entities import ConfigurateMethod, ProviderEntity - -IDENTIFIER = f"langgenius/openai:1.0.0@{'a' * 32}" - - -def _provider() -> PluginModelProviderDeclaration: - return PluginModelProviderDeclaration( - provider="langgenius/openai/openai", - provider_name="openai", - label=I18nObject(en_US="OpenAI"), - supported_model_types=[ModelType.LLM], - configurate_methods=[ConfigurateMethod.PREDEFINED_MODEL], - plugin_unique_identifier=IDENTIFIER, - installation_source=PluginInstallationSource.Marketplace, - ) - - -@pytest.mark.parametrize("binding", ["assembly", "discovery", "legacy"]) -def test_model_dispatch_does_not_discover_providers(monkeypatch: pytest.MonkeyPatch, binding: str) -> None: - provider = _provider() - runtime = create_plugin_model_runtime(tenant_id="tenant-a") - if binding == "discovery": - - def fetch(*, tenant_id: str, client: PluginModelClient | None = None) -> list[PluginModelProviderDeclaration]: - assert tenant_id == "tenant-a" - assert client is runtime.client - return [provider] - - monkeypatch.setattr(PluginService, "fetch_plugin_model_providers", fetch) - assert runtime.fetch_model_providers() == [provider] - else: - schema = ProviderEntity.model_validate(provider.model_dump()) if binding == "legacy" else provider - create_model_type_instance(runtime=runtime, provider_schema=schema, model_type=ModelType.LLM) - - def unexpected_discovery(*_args: object, **_kwargs: object) -> None: - raise AssertionError("Invocation must reuse the prepared provider identity") - - monkeypatch.setattr(PluginService, "fetch_plugin_model_providers", unexpected_discovery) - requests: list[httpx.Request] = [] - chunk = LLMResultChunk( - model="model", delta=LLMResultChunkDelta(index=0, message=AssistantPromptMessage(content="ok")) - ) - - def handle(request: httpx.Request) -> httpx.Response: - assert request.method == "POST" - assert request.url.path == "/plugin/tenant-a/dispatch/llm/invoke" - requests.append(request) - return httpx.Response(200, text="data: " + json.dumps({"code": 0, "message": "", "data": chunk.model_dump()})) - - with httpx.Client(transport=httpx.MockTransport(handle)) as client: - monkeypatch.setattr(client_module, "_httpx_client", client) - result = runtime.invoke_llm( - provider=provider.provider, - model="model", - credentials={"api_key": "secret"}, - model_parameters={}, - prompt_messages=[], - tools=None, - stop=None, - stream=True, - ) - assert list(result) == [chunk] - - assert len(requests) == 1 - assert requests[0].headers["X-Plugin-ID"] == "langgenius/openai" - expected_identifier = None if binding == "legacy" else IDENTIFIER - assert requests[0].headers.get("X-Plugin-Unique-Identifier") == expected_identifier - - -def test_model_schema_cache_is_scoped_to_prepared_package() -> None: - runtime = create_plugin_model_runtime(tenant_id="tenant-a") - provider = _provider() - runtime.bind_provider_schema(provider) - first = runtime._get_schema_cache_key( - provider=provider.provider, model_type=ModelType.LLM, model="model", credentials={} - ) - runtime.bind_provider_schema( - provider.model_copy(update={"plugin_unique_identifier": f"langgenius/openai:2.0.0@{'b' * 32}"}) - ) - second = runtime._get_schema_cache_key( - provider=provider.provider, model_type=ModelType.LLM, model="model", credentials={} - ) - assert first != second diff --git a/api/tests/unit_tests/services/plugin/test_oauth_dispatch_identity.py b/api/tests/unit_tests/services/plugin/test_oauth_dispatch_identity.py deleted file mode 100644 index e4f270dd2bf..00000000000 --- a/api/tests/unit_tests/services/plugin/test_oauth_dispatch_identity.py +++ /dev/null @@ -1,121 +0,0 @@ -"""OAuth and credential transports preserve the prepared plugin package.""" - -import json - -import httpx -import pytest -from werkzeug import Request - -from core.plugin.impl import base as client_module -from core.plugin.impl.datasource import PluginDatasourceManager -from core.plugin.impl.dynamic_select import DynamicSelectClient -from core.plugin.impl.oauth import OAuthHandler -from services.plugin.oauth_service import OAuthProxyService - -PLUGIN_ID = "langgenius/example" -IDENTIFIER = f"{PLUGIN_ID}:1.0.0@{'a' * 32}" - - -@pytest.mark.parametrize("operation", ["authorize", "callback", "refresh", "dynamic", "validate"]) -def test_dispatch_transports_preserve_package_and_credentials(monkeypatch: pytest.MonkeyPatch, operation: str) -> None: - requests: list[httpx.Request] = [] - credentials = {"client_secret": "test-secret"} - - def handle(request: httpx.Request) -> httpx.Response: - requests.append(request) - assert request.method == "POST" - data: dict[str, object] - if operation == "authorize": - data = {"authorization_url": "https://example.com/authorize"} - elif operation == "dynamic": - data = {"options": []} - elif operation == "validate": - data = {"result": True} - else: - data = {"credentials": {"access_token": "test-token"}, "expires_at": -1} - return httpx.Response(200, text="data: " + json.dumps({"code": 0, "message": "", "data": data})) - - with httpx.Client(transport=httpx.MockTransport(handle)) as client: - monkeypatch.setattr(client_module, "_httpx_client", client) - handler = OAuthHandler() - if operation == "authorize": - handler.get_authorization_url( - "tenant", - "user", - PLUGIN_ID, - "example", - "https://example.com/callback", - credentials, - plugin_unique_identifier=IDENTIFIER, - ) - elif operation == "callback": - handler.get_credentials( - "tenant", - "user", - PLUGIN_ID, - "example", - "https://example.com/callback", - credentials, - Request.from_values("/callback?code=test-code"), - plugin_unique_identifier=IDENTIFIER, - ) - elif operation == "refresh": - handler.refresh_credentials( - "tenant", - "user", - PLUGIN_ID, - "example", - "https://example.com/callback", - credentials, - {"refresh_token": "test-refresh"}, - plugin_unique_identifier=IDENTIFIER, - ) - elif operation == "dynamic": - DynamicSelectClient().fetch_dynamic_select_options( - "tenant", - "user", - PLUGIN_ID, - f"{PLUGIN_ID}/example", - "fetch", - credentials, - "api-key", - "option", - plugin_unique_identifier=IDENTIFIER, - ) - else: - assert PluginDatasourceManager().validate_provider_credentials( - "tenant", "user", "example", PLUGIN_ID, credentials, plugin_unique_identifier=IDENTIFIER - ) - - assert len(requests) == 1 - assert requests[0].headers["X-Plugin-ID"] == PLUGIN_ID - assert requests[0].headers["X-Plugin-Unique-Identifier"] == IDENTIFIER - payload = json.loads(requests[0].content)["data"] - key = "credentials" if operation in ("dynamic", "validate") else "system_credentials" - assert payload[key] == credentials - - -@pytest.mark.parametrize("expected_identifier", [None, IDENTIFIER]) -def test_oauth_context_accepts_matching_provider_and_legacy_context(expected_identifier: str | None) -> None: - context = {"plugin_id": PLUGIN_ID, "provider": "example", "plugin_unique_identifier": expected_identifier} - OAuthProxyService.validate_plugin_context( - context, plugin_id=PLUGIN_ID, provider="example", plugin_unique_identifier=IDENTIFIER - ) - - -@pytest.mark.parametrize( - ("plugin_id", "provider", "identifier"), - [ - ("langgenius/other", "example", IDENTIFIER), - (PLUGIN_ID, "other", IDENTIFIER), - (PLUGIN_ID, "example", f"{PLUGIN_ID}:2.0.0@{'b' * 32}"), - ], -) -def test_oauth_context_rejects_provider_or_package_changes(plugin_id: str, provider: str, identifier: str) -> None: - with pytest.raises(ValueError, match="another provider|Plugin changed"): - OAuthProxyService.validate_plugin_context( - {"plugin_id": PLUGIN_ID, "provider": "example", "plugin_unique_identifier": IDENTIFIER}, - plugin_id=plugin_id, - provider=provider, - plugin_unique_identifier=identifier, - ) diff --git a/api/tests/unit_tests/services/plugin/test_oauth_service.py b/api/tests/unit_tests/services/plugin/test_oauth_service.py index e9798c2170f..e20b41a54cf 100644 --- a/api/tests/unit_tests/services/plugin/test_oauth_service.py +++ b/api/tests/unit_tests/services/plugin/test_oauth_service.py @@ -1,18 +1,136 @@ -"""Tests for services.plugin.oauth_service.OAuthProxyService. - -Covers: CSRF proxy context creation with Redis TTL, context consumption -with one-time use semantics, and validation error paths. -""" - from __future__ import annotations import json from unittest.mock import MagicMock, patch +import httpx import pytest +from werkzeug import Request +from core.plugin.impl import base as client_module +from core.plugin.impl.datasource import PluginDatasourceManager +from core.plugin.impl.dynamic_select import DynamicSelectClient +from core.plugin.impl.oauth import OAuthHandler from services.plugin.oauth_service import OAuthProxyService +"""OAuth and credential transports preserve the prepared plugin package.""" + + +PLUGIN_ID = "langgenius/example" +IDENTIFIER = f"{PLUGIN_ID}:1.0.0@{'a' * 32}" + + +@pytest.mark.parametrize("operation", ["authorize", "callback", "refresh", "dynamic", "validate"]) +def test_dispatch_transports_preserve_package_and_credentials(monkeypatch: pytest.MonkeyPatch, operation: str) -> None: + requests: list[httpx.Request] = [] + credentials = {"client_secret": "test-secret"} + + def handle(request: httpx.Request) -> httpx.Response: + requests.append(request) + assert request.method == "POST" + data: dict[str, object] + if operation == "authorize": + data = {"authorization_url": "https://example.com/authorize"} + elif operation == "dynamic": + data = {"options": []} + elif operation == "validate": + data = {"result": True} + else: + data = {"credentials": {"access_token": "test-token"}, "expires_at": -1} + return httpx.Response(200, text="data: " + json.dumps({"code": 0, "message": "", "data": data})) + + with httpx.Client(transport=httpx.MockTransport(handle)) as client: + monkeypatch.setattr(client_module, "_httpx_client", client) + handler = OAuthHandler() + if operation == "authorize": + handler.get_authorization_url( + "tenant", + "user", + PLUGIN_ID, + "example", + "https://example.com/callback", + credentials, + plugin_unique_identifier=IDENTIFIER, + ) + elif operation == "callback": + handler.get_credentials( + "tenant", + "user", + PLUGIN_ID, + "example", + "https://example.com/callback", + credentials, + Request.from_values("/callback?code=test-code"), + plugin_unique_identifier=IDENTIFIER, + ) + elif operation == "refresh": + handler.refresh_credentials( + "tenant", + "user", + PLUGIN_ID, + "example", + "https://example.com/callback", + credentials, + {"refresh_token": "test-refresh"}, + plugin_unique_identifier=IDENTIFIER, + ) + elif operation == "dynamic": + DynamicSelectClient().fetch_dynamic_select_options( + "tenant", + "user", + PLUGIN_ID, + f"{PLUGIN_ID}/example", + "fetch", + credentials, + "api-key", + "option", + plugin_unique_identifier=IDENTIFIER, + ) + else: + assert PluginDatasourceManager().validate_provider_credentials( + "tenant", "user", "example", PLUGIN_ID, credentials, plugin_unique_identifier=IDENTIFIER + ) + + assert len(requests) == 1 + assert requests[0].headers["X-Plugin-ID"] == PLUGIN_ID + assert requests[0].headers["X-Plugin-Unique-Identifier"] == IDENTIFIER + payload = json.loads(requests[0].content)["data"] + key = "credentials" if operation in ("dynamic", "validate") else "system_credentials" + assert payload[key] == credentials + + +@pytest.mark.parametrize("expected_identifier", [None, IDENTIFIER]) +def test_oauth_context_accepts_matching_provider_and_legacy_context(expected_identifier: str | None) -> None: + context = {"plugin_id": PLUGIN_ID, "provider": "example", "plugin_unique_identifier": expected_identifier} + OAuthProxyService.validate_plugin_context( + context, plugin_id=PLUGIN_ID, provider="example", plugin_unique_identifier=IDENTIFIER + ) + + +@pytest.mark.parametrize( + ("plugin_id", "provider", "identifier"), + [ + ("langgenius/other", "example", IDENTIFIER), + (PLUGIN_ID, "other", IDENTIFIER), + (PLUGIN_ID, "example", f"{PLUGIN_ID}:2.0.0@{'b' * 32}"), + ], +) +def test_oauth_context_rejects_provider_or_package_changes(plugin_id: str, provider: str, identifier: str) -> None: + with pytest.raises(ValueError, match="another provider|Plugin changed"): + OAuthProxyService.validate_plugin_context( + {"plugin_id": PLUGIN_ID, "provider": "example", "plugin_unique_identifier": IDENTIFIER}, + plugin_id=plugin_id, + provider=provider, + plugin_unique_identifier=identifier, + ) + + +"""Tests for services.plugin.oauth_service.OAuthProxyService. + +Covers: CSRF proxy context creation with Redis TTL, context consumption +with one-time use semantics, and validation error paths. +""" + def _redis_mock() -> MagicMock: from extensions.ext_redis import redis_client diff --git a/api/tests/unit_tests/services/plugin/test_plugin_auto_upgrade_service.py b/api/tests/unit_tests/services/plugin/test_plugin_auto_upgrade_service.py index 0296f16dd2c..9df9a40c7c4 100644 --- a/api/tests/unit_tests/services/plugin/test_plugin_auto_upgrade_service.py +++ b/api/tests/unit_tests/services/plugin/test_plugin_auto_upgrade_service.py @@ -4,26 +4,31 @@ from unittest.mock import MagicMock, patch from uuid import uuid4 import pytest -from sqlalchemy import select +from sqlalchemy import func, select from sqlalchemy.orm import Session +from models import Tenant from models.account import ( TenantPluginAutoUpgradeCategory, TenantPluginAutoUpgradeMode, TenantPluginAutoUpgradeStrategy, TenantPluginAutoUpgradeStrategySetting, + TenantPluginDebugPermission, + TenantPluginInstallPermission, + TenantPluginPermission, ) from services.plugin.plugin_auto_upgrade_service import PluginAutoUpgradeService +from services.plugin.plugin_permission_service import PluginPermissionService MODULE = "services.plugin.plugin_auto_upgrade_service" -PLUGIN_CATEGORY = TenantPluginAutoUpgradeCategory.TOOL +_PLUGIN_AUTO_UPGRADE_SERVICE_PLUGIN_CATEGORY = TenantPluginAutoUpgradeCategory.TOOL STRATEGY_MODELS = (TenantPluginAutoUpgradeStrategy,) def _strategy( tenant_id: str, *, - category: TenantPluginAutoUpgradeCategory = PLUGIN_CATEGORY, + category: TenantPluginAutoUpgradeCategory = _PLUGIN_AUTO_UPGRADE_SERVICE_PLUGIN_CATEGORY, setting: TenantPluginAutoUpgradeStrategySetting = TenantPluginAutoUpgradeStrategySetting.FIX_ONLY, mode: TenantPluginAutoUpgradeMode = TenantPluginAutoUpgradeMode.EXCLUDE, exclude: list[str] | None = None, @@ -41,211 +46,444 @@ def _strategy( ) -class TestGetStrategy: - @pytest.mark.parametrize("sqlite_session", [STRATEGY_MODELS], indirect=True) - def test_returns_strategy_when_found(self, sqlite_session: Session) -> None: - tenant_id = str(uuid4()) - strategy = _strategy(tenant_id) - sqlite_session.add(strategy) +class TestPluginAutoUpgradeServiceCases: + class TestGetStrategy: + @pytest.mark.parametrize("sqlite_session", [STRATEGY_MODELS], indirect=True) + def test_returns_strategy_when_found(self, sqlite_session: Session) -> None: + tenant_id = str(uuid4()) + strategy = _strategy(tenant_id) + sqlite_session.add(strategy) + sqlite_session.commit() + + result = PluginAutoUpgradeService.get_strategy( + tenant_id, _PLUGIN_AUTO_UPGRADE_SERVICE_PLUGIN_CATEGORY, session=sqlite_session + ) + + assert result is strategy + + @pytest.mark.parametrize("sqlite_session", [STRATEGY_MODELS], indirect=True) + def test_returns_none_when_not_found(self, sqlite_session: Session) -> None: + assert ( + PluginAutoUpgradeService.get_strategy( + str(uuid4()), _PLUGIN_AUTO_UPGRADE_SERVICE_PLUGIN_CATEGORY, session=sqlite_session + ) + is None + ) + + class TestChangeStrategy: + @pytest.mark.parametrize("sqlite_session", [STRATEGY_MODELS], indirect=True) + def test_creates_new_strategy(self, sqlite_session: Session) -> None: + tenant_id = str(uuid4()) + + result = PluginAutoUpgradeService.change_strategy( + tenant_id, + TenantPluginAutoUpgradeStrategySetting.FIX_ONLY, + 3, + TenantPluginAutoUpgradeMode.ALL, + [], + [], + category=_PLUGIN_AUTO_UPGRADE_SERVICE_PLUGIN_CATEGORY, + session=sqlite_session, + ) + + strategy = sqlite_session.scalar(select(TenantPluginAutoUpgradeStrategy)) + assert result is True + assert strategy is not None + assert strategy.tenant_id == tenant_id + assert strategy.upgrade_time_of_day == 3 + assert strategy.upgrade_mode == TenantPluginAutoUpgradeMode.ALL + + @pytest.mark.parametrize("sqlite_session", [STRATEGY_MODELS], indirect=True) + def test_updates_existing_strategy(self, sqlite_session: Session) -> None: + tenant_id = str(uuid4()) + existing = _strategy(tenant_id) + sqlite_session.add(existing) + sqlite_session.commit() + + result = PluginAutoUpgradeService.change_strategy( + tenant_id, + TenantPluginAutoUpgradeStrategySetting.LATEST, + 5, + TenantPluginAutoUpgradeMode.PARTIAL, + ["p1"], + ["p2"], + category=_PLUGIN_AUTO_UPGRADE_SERVICE_PLUGIN_CATEGORY, + session=sqlite_session, + ) + + sqlite_session.refresh(existing) + assert result is True + assert existing.strategy_setting == TenantPluginAutoUpgradeStrategySetting.LATEST + assert existing.upgrade_time_of_day == 5 + assert existing.upgrade_mode == TenantPluginAutoUpgradeMode.PARTIAL + assert existing.exclude_plugins == ["p1"] + assert existing.include_plugins == ["p2"] + + class TestExcludePlugin: + @pytest.mark.parametrize("sqlite_session", [STRATEGY_MODELS], indirect=True) + def test_creates_default_strategy_when_none_exists(self, sqlite_session: Session) -> None: + tenant_id = str(uuid4()) + + result = PluginAutoUpgradeService.exclude_plugin( + tenant_id, "plugin-1", _PLUGIN_AUTO_UPGRADE_SERVICE_PLUGIN_CATEGORY, session=sqlite_session + ) + + strategy = sqlite_session.scalar(select(TenantPluginAutoUpgradeStrategy)) + assert result is True + assert strategy is not None + assert strategy.exclude_plugins == ["plugin-1"] + + @pytest.mark.parametrize("sqlite_session", [STRATEGY_MODELS], indirect=True) + def test_appends_to_exclude_list_in_exclude_mode(self, sqlite_session: Session) -> None: + tenant_id = str(uuid4()) + existing = _strategy(tenant_id, exclude=["p-existing"]) + sqlite_session.add(existing) + sqlite_session.commit() + + PluginAutoUpgradeService.exclude_plugin( + tenant_id, "p-new", _PLUGIN_AUTO_UPGRADE_SERVICE_PLUGIN_CATEGORY, session=sqlite_session + ) + + sqlite_session.refresh(existing) + assert existing.exclude_plugins == ["p-existing", "p-new"] + + @pytest.mark.parametrize("sqlite_session", [STRATEGY_MODELS], indirect=True) + def test_removes_from_include_list_in_partial_mode(self, sqlite_session: Session) -> None: + tenant_id = str(uuid4()) + existing = _strategy(tenant_id, mode=TenantPluginAutoUpgradeMode.PARTIAL, include=["p1", "p2"]) + sqlite_session.add(existing) + sqlite_session.commit() + + PluginAutoUpgradeService.exclude_plugin( + tenant_id, "p1", _PLUGIN_AUTO_UPGRADE_SERVICE_PLUGIN_CATEGORY, session=sqlite_session + ) + + sqlite_session.refresh(existing) + assert existing.include_plugins == ["p2"] + + @pytest.mark.parametrize("sqlite_session", [STRATEGY_MODELS], indirect=True) + def test_switches_to_exclude_mode_from_all(self, sqlite_session: Session) -> None: + tenant_id = str(uuid4()) + existing = _strategy(tenant_id, mode=TenantPluginAutoUpgradeMode.ALL) + sqlite_session.add(existing) + sqlite_session.commit() + + PluginAutoUpgradeService.exclude_plugin( + tenant_id, "p1", _PLUGIN_AUTO_UPGRADE_SERVICE_PLUGIN_CATEGORY, session=sqlite_session + ) + + sqlite_session.refresh(existing) + assert existing.upgrade_mode == TenantPluginAutoUpgradeMode.EXCLUDE + assert existing.exclude_plugins == ["p1"] + + @pytest.mark.parametrize("sqlite_session", [STRATEGY_MODELS], indirect=True) + def test_no_duplicate_in_exclude_list(self, sqlite_session: Session) -> None: + tenant_id = str(uuid4()) + existing = _strategy(tenant_id, exclude=["p1"]) + sqlite_session.add(existing) + sqlite_session.commit() + + PluginAutoUpgradeService.exclude_plugin( + tenant_id, "p1", _PLUGIN_AUTO_UPGRADE_SERVICE_PLUGIN_CATEGORY, session=sqlite_session + ) + + sqlite_session.refresh(existing) + assert existing.exclude_plugins == ["p1"] + + class TestBackfillStrategyCategories: + @pytest.mark.parametrize("sqlite_session", [STRATEGY_MODELS], indirect=True) + def test_creates_default_missing_categories_without_fetching_daemon(self, sqlite_session: Session) -> None: + tenant_id = str(uuid4()) + tool_strategy = _strategy(tenant_id) + sqlite_session.add(tool_strategy) + sqlite_session.commit() + installer = MagicMock() + + with patch(f"{MODULE}.PluginInstaller", return_value=installer): + result = PluginAutoUpgradeService.backfill_strategy_categories(tenant_id, session=sqlite_session) + expected_time = PluginAutoUpgradeService.default_upgrade_time_of_day(tenant_id) + + strategies = list(sqlite_session.scalars(select(TenantPluginAutoUpgradeStrategy)).all()) + assert result.created_count == len(TenantPluginAutoUpgradeCategory) - 1 + assert result.normalized is False + installer.list_plugins.assert_not_called() + assert len(strategies) == len(TenantPluginAutoUpgradeCategory) + assert tool_strategy.upgrade_time_of_day == expected_time + model_strategy = next( + strategy for strategy in strategies if strategy.category == TenantPluginAutoUpgradeCategory.MODEL + ) + assert model_strategy.strategy_setting == TenantPluginAutoUpgradeStrategySetting.LATEST + assert model_strategy.upgrade_time_of_day == expected_time + + def test_default_upgrade_time_is_aligned_to_fifteen_minutes(self) -> None: + default_time = PluginAutoUpgradeService.default_upgrade_time_of_day(str(uuid4())) + assert default_time % (15 * 60) == 0 + assert 0 <= default_time < 24 * 60 * 60 + + @pytest.mark.parametrize("sqlite_session", [STRATEGY_MODELS], indirect=True) + def test_creates_missing_categories_and_splits_known_plugins( + self, sqlite_session: Session, caplog: pytest.LogCaptureFixture + ) -> None: + tenant_id = str(uuid4()) + tool_strategy = _strategy( + tenant_id, + exclude=["tool-plugin", "model-plugin", "unknown-plugin"], + include=["model-plugin", "tool-plugin"], + ) + model_strategy = _strategy( + tenant_id, + category=TenantPluginAutoUpgradeCategory.MODEL, + exclude=["tool-plugin", "model-plugin", "unknown-plugin"], + include=["model-plugin", "tool-plugin"], + ) + sqlite_session.add_all([tool_strategy, model_strategy]) + sqlite_session.commit() + installed_plugins = [ + SimpleNamespace( + plugin_id="tool-plugin", + declaration=SimpleNamespace(category=TenantPluginAutoUpgradeCategory.TOOL), + ), + SimpleNamespace( + plugin_id="model-plugin", + declaration=SimpleNamespace(category=TenantPluginAutoUpgradeCategory.MODEL), + ), + ] + installer = MagicMock() + installer.list_plugins.return_value = installed_plugins + + with ( + patch(f"{MODULE}.PluginInstaller", return_value=installer), + caplog.at_level(logging.WARNING, logger=MODULE), + ): + result = PluginAutoUpgradeService.backfill_strategy_categories(tenant_id, session=sqlite_session) + + strategies = list(sqlite_session.scalars(select(TenantPluginAutoUpgradeStrategy)).all()) + assert result.created_count == len(TenantPluginAutoUpgradeCategory) - 2 + assert result.normalized is True + assert len(strategies) == len(TenantPluginAutoUpgradeCategory) + assert tool_strategy.exclude_plugins == ["tool-plugin"] + assert tool_strategy.include_plugins == ["tool-plugin"] + assert model_strategy.exclude_plugins == ["model-plugin"] + assert model_strategy.include_plugins == ["model-plugin"] + assert ( + "Skipped unknown plugin IDs while backfilling plugin auto-upgrade strategies: " + f"tenant_id={tenant_id}, field=exclude_plugins, plugin_ids=['unknown-plugin']" in caplog.messages + ) + + +"""SQLite-backed lifecycle coverage for plugin permissions and auto-upgrade strategies.""" + + +_PLUGIN_LIFECYCLE_SQLITE_PLUGIN_CATEGORY = TenantPluginAutoUpgradeCategory.TOOL + + +class TestPluginLifecycleSqliteCases: + @staticmethod + @pytest.fixture + def tenant(sqlite_session: Session) -> str: + tenant = Tenant(name="plugin_unit_tenant") + sqlite_session.add(tenant) sqlite_session.commit() + return tenant.id - result = PluginAutoUpgradeService.get_strategy(tenant_id, PLUGIN_CATEGORY, session=sqlite_session) + class TestPluginPermissionLifecycle: + def test_get_returns_none_for_new_tenant(self, tenant: str, sqlite_session: Session) -> None: + assert PluginPermissionService.get_permission(tenant, session=sqlite_session) is None - assert result is strategy + def test_change_creates_row(self, tenant: str, sqlite_session: Session) -> None: + result = PluginPermissionService.change_permission( + tenant, + TenantPluginInstallPermission.ADMINS, + TenantPluginDebugPermission.EVERYONE, + session=sqlite_session, + ) + assert result is True - @pytest.mark.parametrize("sqlite_session", [STRATEGY_MODELS], indirect=True) - def test_returns_none_when_not_found(self, sqlite_session: Session) -> None: - assert PluginAutoUpgradeService.get_strategy(str(uuid4()), PLUGIN_CATEGORY, session=sqlite_session) is None + perm = PluginPermissionService.get_permission(tenant, session=sqlite_session) + assert perm is not None + assert perm.install_permission == TenantPluginInstallPermission.ADMINS + assert perm.debug_permission == TenantPluginDebugPermission.EVERYONE + def test_change_updates_existing_row(self, tenant: str, sqlite_session: Session) -> None: + PluginPermissionService.change_permission( + tenant, + TenantPluginInstallPermission.ADMINS, + TenantPluginDebugPermission.NOBODY, + session=sqlite_session, + ) + PluginPermissionService.change_permission( + tenant, + TenantPluginInstallPermission.EVERYONE, + TenantPluginDebugPermission.ADMINS, + session=sqlite_session, + ) + perm = PluginPermissionService.get_permission(tenant, session=sqlite_session) + assert perm is not None + assert perm.install_permission == TenantPluginInstallPermission.EVERYONE + assert perm.debug_permission == TenantPluginDebugPermission.ADMINS -class TestChangeStrategy: - @pytest.mark.parametrize("sqlite_session", [STRATEGY_MODELS], indirect=True) - def test_creates_new_strategy(self, sqlite_session: Session) -> None: - tenant_id = str(uuid4()) + count = sqlite_session.scalar( + select(func.count()) + .select_from(TenantPluginPermission) + .where(TenantPluginPermission.tenant_id == tenant) + ) + assert count == 1 - result = PluginAutoUpgradeService.change_strategy( - tenant_id, - TenantPluginAutoUpgradeStrategySetting.FIX_ONLY, - 3, - TenantPluginAutoUpgradeMode.ALL, - [], - [], - category=PLUGIN_CATEGORY, - session=sqlite_session, - ) + class TestPluginAutoUpgradeLifecycle: + def test_get_returns_none_for_new_tenant(self, tenant: str, sqlite_session: Session) -> None: + assert ( + PluginAutoUpgradeService.get_strategy( + tenant, _PLUGIN_LIFECYCLE_SQLITE_PLUGIN_CATEGORY, session=sqlite_session + ) + is None + ) - strategy = sqlite_session.scalar(select(TenantPluginAutoUpgradeStrategy)) - assert result is True - assert strategy is not None - assert strategy.tenant_id == tenant_id - assert strategy.upgrade_time_of_day == 3 - assert strategy.upgrade_mode == TenantPluginAutoUpgradeMode.ALL + def test_change_creates_row(self, tenant: str, sqlite_session: Session) -> None: + result = PluginAutoUpgradeService.change_strategy( + tenant, + strategy_setting=TenantPluginAutoUpgradeStrategySetting.LATEST, + upgrade_time_of_day=3, + upgrade_mode=TenantPluginAutoUpgradeMode.ALL, + exclude_plugins=[], + include_plugins=[], + category=_PLUGIN_LIFECYCLE_SQLITE_PLUGIN_CATEGORY, + session=sqlite_session, + ) + assert result is True - @pytest.mark.parametrize("sqlite_session", [STRATEGY_MODELS], indirect=True) - def test_updates_existing_strategy(self, sqlite_session: Session) -> None: - tenant_id = str(uuid4()) - existing = _strategy(tenant_id) - sqlite_session.add(existing) - sqlite_session.commit() + strategy = PluginAutoUpgradeService.get_strategy( + tenant, _PLUGIN_LIFECYCLE_SQLITE_PLUGIN_CATEGORY, session=sqlite_session + ) + assert strategy is not None + assert strategy.strategy_setting == TenantPluginAutoUpgradeStrategySetting.LATEST + assert strategy.upgrade_time_of_day == 3 - result = PluginAutoUpgradeService.change_strategy( - tenant_id, - TenantPluginAutoUpgradeStrategySetting.LATEST, - 5, - TenantPluginAutoUpgradeMode.PARTIAL, - ["p1"], - ["p2"], - category=PLUGIN_CATEGORY, - session=sqlite_session, - ) + def test_change_updates_existing_row(self, tenant: str, sqlite_session: Session) -> None: + PluginAutoUpgradeService.change_strategy( + tenant, + strategy_setting=TenantPluginAutoUpgradeStrategySetting.FIX_ONLY, + upgrade_time_of_day=0, + upgrade_mode=TenantPluginAutoUpgradeMode.ALL, + exclude_plugins=[], + include_plugins=[], + category=_PLUGIN_LIFECYCLE_SQLITE_PLUGIN_CATEGORY, + session=sqlite_session, + ) + PluginAutoUpgradeService.change_strategy( + tenant, + strategy_setting=TenantPluginAutoUpgradeStrategySetting.LATEST, + upgrade_time_of_day=12, + upgrade_mode=TenantPluginAutoUpgradeMode.PARTIAL, + exclude_plugins=[], + include_plugins=["plugin-a"], + category=_PLUGIN_LIFECYCLE_SQLITE_PLUGIN_CATEGORY, + session=sqlite_session, + ) - sqlite_session.refresh(existing) - assert result is True - assert existing.strategy_setting == TenantPluginAutoUpgradeStrategySetting.LATEST - assert existing.upgrade_time_of_day == 5 - assert existing.upgrade_mode == TenantPluginAutoUpgradeMode.PARTIAL - assert existing.exclude_plugins == ["p1"] - assert existing.include_plugins == ["p2"] + strategy = PluginAutoUpgradeService.get_strategy( + tenant, _PLUGIN_LIFECYCLE_SQLITE_PLUGIN_CATEGORY, session=sqlite_session + ) + assert strategy is not None + assert strategy.strategy_setting == TenantPluginAutoUpgradeStrategySetting.LATEST + assert strategy.upgrade_time_of_day == 12 + assert strategy.upgrade_mode == TenantPluginAutoUpgradeMode.PARTIAL + assert strategy.include_plugins == ["plugin-a"] + def test_exclude_plugin_creates_strategy_when_none_exists(self, tenant: str, sqlite_session: Session) -> None: + PluginAutoUpgradeService.exclude_plugin( + tenant, "my-plugin", _PLUGIN_LIFECYCLE_SQLITE_PLUGIN_CATEGORY, session=sqlite_session + ) -class TestExcludePlugin: - @pytest.mark.parametrize("sqlite_session", [STRATEGY_MODELS], indirect=True) - def test_creates_default_strategy_when_none_exists(self, sqlite_session: Session) -> None: - tenant_id = str(uuid4()) + strategy = PluginAutoUpgradeService.get_strategy( + tenant, _PLUGIN_LIFECYCLE_SQLITE_PLUGIN_CATEGORY, session=sqlite_session + ) + assert strategy is not None + assert strategy.upgrade_mode == TenantPluginAutoUpgradeMode.EXCLUDE + assert "my-plugin" in strategy.exclude_plugins - result = PluginAutoUpgradeService.exclude_plugin(tenant_id, "plugin-1", PLUGIN_CATEGORY, session=sqlite_session) + def test_exclude_plugin_appends_in_exclude_mode(self, tenant: str, sqlite_session: Session) -> None: + PluginAutoUpgradeService.change_strategy( + tenant, + strategy_setting=TenantPluginAutoUpgradeStrategySetting.FIX_ONLY, + upgrade_time_of_day=0, + upgrade_mode=TenantPluginAutoUpgradeMode.EXCLUDE, + exclude_plugins=["existing"], + include_plugins=[], + category=_PLUGIN_LIFECYCLE_SQLITE_PLUGIN_CATEGORY, + session=sqlite_session, + ) + PluginAutoUpgradeService.exclude_plugin( + tenant, "new-plugin", _PLUGIN_LIFECYCLE_SQLITE_PLUGIN_CATEGORY, session=sqlite_session + ) - strategy = sqlite_session.scalar(select(TenantPluginAutoUpgradeStrategy)) - assert result is True - assert strategy is not None - assert strategy.exclude_plugins == ["plugin-1"] + strategy = PluginAutoUpgradeService.get_strategy( + tenant, _PLUGIN_LIFECYCLE_SQLITE_PLUGIN_CATEGORY, session=sqlite_session + ) + assert strategy is not None + assert "existing" in strategy.exclude_plugins + assert "new-plugin" in strategy.exclude_plugins - @pytest.mark.parametrize("sqlite_session", [STRATEGY_MODELS], indirect=True) - def test_appends_to_exclude_list_in_exclude_mode(self, sqlite_session: Session) -> None: - tenant_id = str(uuid4()) - existing = _strategy(tenant_id, exclude=["p-existing"]) - sqlite_session.add(existing) - sqlite_session.commit() + def test_exclude_plugin_dedup_in_exclude_mode(self, tenant: str, sqlite_session: Session) -> None: + PluginAutoUpgradeService.change_strategy( + tenant, + strategy_setting=TenantPluginAutoUpgradeStrategySetting.FIX_ONLY, + upgrade_time_of_day=0, + upgrade_mode=TenantPluginAutoUpgradeMode.EXCLUDE, + exclude_plugins=["same-plugin"], + include_plugins=[], + category=_PLUGIN_LIFECYCLE_SQLITE_PLUGIN_CATEGORY, + session=sqlite_session, + ) + PluginAutoUpgradeService.exclude_plugin( + tenant, "same-plugin", _PLUGIN_LIFECYCLE_SQLITE_PLUGIN_CATEGORY, session=sqlite_session + ) - PluginAutoUpgradeService.exclude_plugin(tenant_id, "p-new", PLUGIN_CATEGORY, session=sqlite_session) + strategy = PluginAutoUpgradeService.get_strategy( + tenant, _PLUGIN_LIFECYCLE_SQLITE_PLUGIN_CATEGORY, session=sqlite_session + ) + assert strategy is not None + assert strategy.exclude_plugins.count("same-plugin") == 1 - sqlite_session.refresh(existing) - assert existing.exclude_plugins == ["p-existing", "p-new"] + def test_exclude_from_partial_mode_removes_from_include(self, tenant: str, sqlite_session: Session) -> None: + PluginAutoUpgradeService.change_strategy( + tenant, + strategy_setting=TenantPluginAutoUpgradeStrategySetting.FIX_ONLY, + upgrade_time_of_day=0, + upgrade_mode=TenantPluginAutoUpgradeMode.PARTIAL, + exclude_plugins=[], + include_plugins=["p1", "p2"], + category=_PLUGIN_LIFECYCLE_SQLITE_PLUGIN_CATEGORY, + session=sqlite_session, + ) + PluginAutoUpgradeService.exclude_plugin( + tenant, "p1", _PLUGIN_LIFECYCLE_SQLITE_PLUGIN_CATEGORY, session=sqlite_session + ) - @pytest.mark.parametrize("sqlite_session", [STRATEGY_MODELS], indirect=True) - def test_removes_from_include_list_in_partial_mode(self, sqlite_session: Session) -> None: - tenant_id = str(uuid4()) - existing = _strategy(tenant_id, mode=TenantPluginAutoUpgradeMode.PARTIAL, include=["p1", "p2"]) - sqlite_session.add(existing) - sqlite_session.commit() + strategy = PluginAutoUpgradeService.get_strategy( + tenant, _PLUGIN_LIFECYCLE_SQLITE_PLUGIN_CATEGORY, session=sqlite_session + ) + assert strategy is not None + assert "p1" not in strategy.include_plugins + assert "p2" in strategy.include_plugins - PluginAutoUpgradeService.exclude_plugin(tenant_id, "p1", PLUGIN_CATEGORY, session=sqlite_session) + def test_exclude_from_all_mode_switches_to_exclude(self, tenant: str, sqlite_session: Session) -> None: + PluginAutoUpgradeService.change_strategy( + tenant, + strategy_setting=TenantPluginAutoUpgradeStrategySetting.LATEST, + upgrade_time_of_day=0, + upgrade_mode=TenantPluginAutoUpgradeMode.ALL, + exclude_plugins=[], + include_plugins=[], + category=_PLUGIN_LIFECYCLE_SQLITE_PLUGIN_CATEGORY, + session=sqlite_session, + ) + PluginAutoUpgradeService.exclude_plugin( + tenant, "excluded-plugin", _PLUGIN_LIFECYCLE_SQLITE_PLUGIN_CATEGORY, session=sqlite_session + ) - sqlite_session.refresh(existing) - assert existing.include_plugins == ["p2"] - - @pytest.mark.parametrize("sqlite_session", [STRATEGY_MODELS], indirect=True) - def test_switches_to_exclude_mode_from_all(self, sqlite_session: Session) -> None: - tenant_id = str(uuid4()) - existing = _strategy(tenant_id, mode=TenantPluginAutoUpgradeMode.ALL) - sqlite_session.add(existing) - sqlite_session.commit() - - PluginAutoUpgradeService.exclude_plugin(tenant_id, "p1", PLUGIN_CATEGORY, session=sqlite_session) - - sqlite_session.refresh(existing) - assert existing.upgrade_mode == TenantPluginAutoUpgradeMode.EXCLUDE - assert existing.exclude_plugins == ["p1"] - - @pytest.mark.parametrize("sqlite_session", [STRATEGY_MODELS], indirect=True) - def test_no_duplicate_in_exclude_list(self, sqlite_session: Session) -> None: - tenant_id = str(uuid4()) - existing = _strategy(tenant_id, exclude=["p1"]) - sqlite_session.add(existing) - sqlite_session.commit() - - PluginAutoUpgradeService.exclude_plugin(tenant_id, "p1", PLUGIN_CATEGORY, session=sqlite_session) - - sqlite_session.refresh(existing) - assert existing.exclude_plugins == ["p1"] - - -class TestBackfillStrategyCategories: - @pytest.mark.parametrize("sqlite_session", [STRATEGY_MODELS], indirect=True) - def test_creates_default_missing_categories_without_fetching_daemon(self, sqlite_session: Session) -> None: - tenant_id = str(uuid4()) - tool_strategy = _strategy(tenant_id) - sqlite_session.add(tool_strategy) - sqlite_session.commit() - installer = MagicMock() - - with patch(f"{MODULE}.PluginInstaller", return_value=installer): - result = PluginAutoUpgradeService.backfill_strategy_categories(tenant_id, session=sqlite_session) - expected_time = PluginAutoUpgradeService.default_upgrade_time_of_day(tenant_id) - - strategies = list(sqlite_session.scalars(select(TenantPluginAutoUpgradeStrategy)).all()) - assert result.created_count == len(TenantPluginAutoUpgradeCategory) - 1 - assert result.normalized is False - installer.list_plugins.assert_not_called() - assert len(strategies) == len(TenantPluginAutoUpgradeCategory) - assert tool_strategy.upgrade_time_of_day == expected_time - model_strategy = next( - strategy for strategy in strategies if strategy.category == TenantPluginAutoUpgradeCategory.MODEL - ) - assert model_strategy.strategy_setting == TenantPluginAutoUpgradeStrategySetting.LATEST - assert model_strategy.upgrade_time_of_day == expected_time - - def test_default_upgrade_time_is_aligned_to_fifteen_minutes(self) -> None: - default_time = PluginAutoUpgradeService.default_upgrade_time_of_day(str(uuid4())) - assert default_time % (15 * 60) == 0 - assert 0 <= default_time < 24 * 60 * 60 - - @pytest.mark.parametrize("sqlite_session", [STRATEGY_MODELS], indirect=True) - def test_creates_missing_categories_and_splits_known_plugins( - self, sqlite_session: Session, caplog: pytest.LogCaptureFixture - ) -> None: - tenant_id = str(uuid4()) - tool_strategy = _strategy( - tenant_id, - exclude=["tool-plugin", "model-plugin", "unknown-plugin"], - include=["model-plugin", "tool-plugin"], - ) - model_strategy = _strategy( - tenant_id, - category=TenantPluginAutoUpgradeCategory.MODEL, - exclude=["tool-plugin", "model-plugin", "unknown-plugin"], - include=["model-plugin", "tool-plugin"], - ) - sqlite_session.add_all([tool_strategy, model_strategy]) - sqlite_session.commit() - installed_plugins = [ - SimpleNamespace( - plugin_id="tool-plugin", - declaration=SimpleNamespace(category=TenantPluginAutoUpgradeCategory.TOOL), - ), - SimpleNamespace( - plugin_id="model-plugin", - declaration=SimpleNamespace(category=TenantPluginAutoUpgradeCategory.MODEL), - ), - ] - installer = MagicMock() - installer.list_plugins.return_value = installed_plugins - - with ( - patch(f"{MODULE}.PluginInstaller", return_value=installer), - caplog.at_level(logging.WARNING, logger=MODULE), - ): - result = PluginAutoUpgradeService.backfill_strategy_categories(tenant_id, session=sqlite_session) - - strategies = list(sqlite_session.scalars(select(TenantPluginAutoUpgradeStrategy)).all()) - assert result.created_count == len(TenantPluginAutoUpgradeCategory) - 2 - assert result.normalized is True - assert len(strategies) == len(TenantPluginAutoUpgradeCategory) - assert tool_strategy.exclude_plugins == ["tool-plugin"] - assert tool_strategy.include_plugins == ["tool-plugin"] - assert model_strategy.exclude_plugins == ["model-plugin"] - assert model_strategy.include_plugins == ["model-plugin"] - assert ( - "Skipped unknown plugin IDs while backfilling plugin auto-upgrade strategies: " - f"tenant_id={tenant_id}, field=exclude_plugins, plugin_ids=['unknown-plugin']" in caplog.messages - ) + strategy = PluginAutoUpgradeService.get_strategy( + tenant, _PLUGIN_LIFECYCLE_SQLITE_PLUGIN_CATEGORY, session=sqlite_session + ) + assert strategy is not None + assert strategy.upgrade_mode == TenantPluginAutoUpgradeMode.EXCLUDE + assert "excluded-plugin" in strategy.exclude_plugins diff --git a/api/tests/unit_tests/services/plugin/test_plugin_lifecycle_sqlite.py b/api/tests/unit_tests/services/plugin/test_plugin_lifecycle_sqlite.py deleted file mode 100644 index 0d8f06ac98b..00000000000 --- a/api/tests/unit_tests/services/plugin/test_plugin_lifecycle_sqlite.py +++ /dev/null @@ -1,200 +0,0 @@ -"""SQLite-backed lifecycle coverage for plugin permissions and auto-upgrade strategies.""" - -import pytest -from sqlalchemy import func, select -from sqlalchemy.orm import Session - -from models import Tenant -from models.account import ( - TenantPluginAutoUpgradeCategory, - TenantPluginAutoUpgradeMode, - TenantPluginAutoUpgradeStrategySetting, - TenantPluginDebugPermission, - TenantPluginInstallPermission, - TenantPluginPermission, -) -from services.plugin.plugin_auto_upgrade_service import PluginAutoUpgradeService -from services.plugin.plugin_permission_service import PluginPermissionService - -PLUGIN_CATEGORY = TenantPluginAutoUpgradeCategory.TOOL - - -@pytest.fixture -def tenant(sqlite_session: Session) -> str: - tenant = Tenant(name="plugin_unit_tenant") - sqlite_session.add(tenant) - sqlite_session.commit() - return tenant.id - - -class TestPluginPermissionLifecycle: - def test_get_returns_none_for_new_tenant(self, tenant: str, sqlite_session: Session) -> None: - assert PluginPermissionService.get_permission(tenant, session=sqlite_session) is None - - def test_change_creates_row(self, tenant: str, sqlite_session: Session) -> None: - result = PluginPermissionService.change_permission( - tenant, - TenantPluginInstallPermission.ADMINS, - TenantPluginDebugPermission.EVERYONE, - session=sqlite_session, - ) - assert result is True - - perm = PluginPermissionService.get_permission(tenant, session=sqlite_session) - assert perm is not None - assert perm.install_permission == TenantPluginInstallPermission.ADMINS - assert perm.debug_permission == TenantPluginDebugPermission.EVERYONE - - def test_change_updates_existing_row(self, tenant: str, sqlite_session: Session) -> None: - PluginPermissionService.change_permission( - tenant, - TenantPluginInstallPermission.ADMINS, - TenantPluginDebugPermission.NOBODY, - session=sqlite_session, - ) - PluginPermissionService.change_permission( - tenant, - TenantPluginInstallPermission.EVERYONE, - TenantPluginDebugPermission.ADMINS, - session=sqlite_session, - ) - perm = PluginPermissionService.get_permission(tenant, session=sqlite_session) - assert perm is not None - assert perm.install_permission == TenantPluginInstallPermission.EVERYONE - assert perm.debug_permission == TenantPluginDebugPermission.ADMINS - - count = sqlite_session.scalar( - select(func.count()).select_from(TenantPluginPermission).where(TenantPluginPermission.tenant_id == tenant) - ) - assert count == 1 - - -class TestPluginAutoUpgradeLifecycle: - def test_get_returns_none_for_new_tenant(self, tenant: str, sqlite_session: Session) -> None: - assert PluginAutoUpgradeService.get_strategy(tenant, PLUGIN_CATEGORY, session=sqlite_session) is None - - def test_change_creates_row(self, tenant: str, sqlite_session: Session) -> None: - result = PluginAutoUpgradeService.change_strategy( - tenant, - strategy_setting=TenantPluginAutoUpgradeStrategySetting.LATEST, - upgrade_time_of_day=3, - upgrade_mode=TenantPluginAutoUpgradeMode.ALL, - exclude_plugins=[], - include_plugins=[], - category=PLUGIN_CATEGORY, - session=sqlite_session, - ) - assert result is True - - strategy = PluginAutoUpgradeService.get_strategy(tenant, PLUGIN_CATEGORY, session=sqlite_session) - assert strategy is not None - assert strategy.strategy_setting == TenantPluginAutoUpgradeStrategySetting.LATEST - assert strategy.upgrade_time_of_day == 3 - - def test_change_updates_existing_row(self, tenant: str, sqlite_session: Session) -> None: - PluginAutoUpgradeService.change_strategy( - tenant, - strategy_setting=TenantPluginAutoUpgradeStrategySetting.FIX_ONLY, - upgrade_time_of_day=0, - upgrade_mode=TenantPluginAutoUpgradeMode.ALL, - exclude_plugins=[], - include_plugins=[], - category=PLUGIN_CATEGORY, - session=sqlite_session, - ) - PluginAutoUpgradeService.change_strategy( - tenant, - strategy_setting=TenantPluginAutoUpgradeStrategySetting.LATEST, - upgrade_time_of_day=12, - upgrade_mode=TenantPluginAutoUpgradeMode.PARTIAL, - exclude_plugins=[], - include_plugins=["plugin-a"], - category=PLUGIN_CATEGORY, - session=sqlite_session, - ) - - strategy = PluginAutoUpgradeService.get_strategy(tenant, PLUGIN_CATEGORY, session=sqlite_session) - assert strategy is not None - assert strategy.strategy_setting == TenantPluginAutoUpgradeStrategySetting.LATEST - assert strategy.upgrade_time_of_day == 12 - assert strategy.upgrade_mode == TenantPluginAutoUpgradeMode.PARTIAL - assert strategy.include_plugins == ["plugin-a"] - - def test_exclude_plugin_creates_strategy_when_none_exists(self, tenant: str, sqlite_session: Session) -> None: - PluginAutoUpgradeService.exclude_plugin(tenant, "my-plugin", PLUGIN_CATEGORY, session=sqlite_session) - - strategy = PluginAutoUpgradeService.get_strategy(tenant, PLUGIN_CATEGORY, session=sqlite_session) - assert strategy is not None - assert strategy.upgrade_mode == TenantPluginAutoUpgradeMode.EXCLUDE - assert "my-plugin" in strategy.exclude_plugins - - def test_exclude_plugin_appends_in_exclude_mode(self, tenant: str, sqlite_session: Session) -> None: - PluginAutoUpgradeService.change_strategy( - tenant, - strategy_setting=TenantPluginAutoUpgradeStrategySetting.FIX_ONLY, - upgrade_time_of_day=0, - upgrade_mode=TenantPluginAutoUpgradeMode.EXCLUDE, - exclude_plugins=["existing"], - include_plugins=[], - category=PLUGIN_CATEGORY, - session=sqlite_session, - ) - PluginAutoUpgradeService.exclude_plugin(tenant, "new-plugin", PLUGIN_CATEGORY, session=sqlite_session) - - strategy = PluginAutoUpgradeService.get_strategy(tenant, PLUGIN_CATEGORY, session=sqlite_session) - assert strategy is not None - assert "existing" in strategy.exclude_plugins - assert "new-plugin" in strategy.exclude_plugins - - def test_exclude_plugin_dedup_in_exclude_mode(self, tenant: str, sqlite_session: Session) -> None: - PluginAutoUpgradeService.change_strategy( - tenant, - strategy_setting=TenantPluginAutoUpgradeStrategySetting.FIX_ONLY, - upgrade_time_of_day=0, - upgrade_mode=TenantPluginAutoUpgradeMode.EXCLUDE, - exclude_plugins=["same-plugin"], - include_plugins=[], - category=PLUGIN_CATEGORY, - session=sqlite_session, - ) - PluginAutoUpgradeService.exclude_plugin(tenant, "same-plugin", PLUGIN_CATEGORY, session=sqlite_session) - - strategy = PluginAutoUpgradeService.get_strategy(tenant, PLUGIN_CATEGORY, session=sqlite_session) - assert strategy is not None - assert strategy.exclude_plugins.count("same-plugin") == 1 - - def test_exclude_from_partial_mode_removes_from_include(self, tenant: str, sqlite_session: Session) -> None: - PluginAutoUpgradeService.change_strategy( - tenant, - strategy_setting=TenantPluginAutoUpgradeStrategySetting.FIX_ONLY, - upgrade_time_of_day=0, - upgrade_mode=TenantPluginAutoUpgradeMode.PARTIAL, - exclude_plugins=[], - include_plugins=["p1", "p2"], - category=PLUGIN_CATEGORY, - session=sqlite_session, - ) - PluginAutoUpgradeService.exclude_plugin(tenant, "p1", PLUGIN_CATEGORY, session=sqlite_session) - - strategy = PluginAutoUpgradeService.get_strategy(tenant, PLUGIN_CATEGORY, session=sqlite_session) - assert strategy is not None - assert "p1" not in strategy.include_plugins - assert "p2" in strategy.include_plugins - - def test_exclude_from_all_mode_switches_to_exclude(self, tenant: str, sqlite_session: Session) -> None: - PluginAutoUpgradeService.change_strategy( - tenant, - strategy_setting=TenantPluginAutoUpgradeStrategySetting.LATEST, - upgrade_time_of_day=0, - upgrade_mode=TenantPluginAutoUpgradeMode.ALL, - exclude_plugins=[], - include_plugins=[], - category=PLUGIN_CATEGORY, - session=sqlite_session, - ) - PluginAutoUpgradeService.exclude_plugin(tenant, "excluded-plugin", PLUGIN_CATEGORY, session=sqlite_session) - - strategy = PluginAutoUpgradeService.get_strategy(tenant, PLUGIN_CATEGORY, session=sqlite_session) - assert strategy is not None - assert strategy.upgrade_mode == TenantPluginAutoUpgradeMode.EXCLUDE - assert "excluded-plugin" in strategy.exclude_plugins diff --git a/api/tests/unit_tests/services/rag_pipeline/test_pipeline_service_api_entities.py b/api/tests/unit_tests/services/rag_pipeline/entity/test_pipeline_service_api_entities.py similarity index 100% rename from api/tests/unit_tests/services/rag_pipeline/test_pipeline_service_api_entities.py rename to api/tests/unit_tests/services/rag_pipeline/entity/test_pipeline_service_api_entities.py diff --git a/api/tests/unit_tests/services/rag_pipeline/pipeline_template/test_built_in_retrieval.py b/api/tests/unit_tests/services/rag_pipeline/pipeline_template/built_in/test_built_in_retrieval.py similarity index 100% rename from api/tests/unit_tests/services/rag_pipeline/pipeline_template/test_built_in_retrieval.py rename to api/tests/unit_tests/services/rag_pipeline/pipeline_template/built_in/test_built_in_retrieval.py diff --git a/api/tests/unit_tests/services/rag_pipeline/pipeline_template/test_customized_retrieval.py b/api/tests/unit_tests/services/rag_pipeline/pipeline_template/customized/test_customized_retrieval.py similarity index 100% rename from api/tests/unit_tests/services/rag_pipeline/pipeline_template/test_customized_retrieval.py rename to api/tests/unit_tests/services/rag_pipeline/pipeline_template/customized/test_customized_retrieval.py diff --git a/api/tests/unit_tests/services/rag_pipeline/pipeline_template/test_database_retrieval.py b/api/tests/unit_tests/services/rag_pipeline/pipeline_template/database/test_database_retrieval.py similarity index 100% rename from api/tests/unit_tests/services/rag_pipeline/pipeline_template/test_database_retrieval.py rename to api/tests/unit_tests/services/rag_pipeline/pipeline_template/database/test_database_retrieval.py diff --git a/api/tests/unit_tests/services/rag_pipeline/pipeline_template/test_remote_retrieval.py b/api/tests/unit_tests/services/rag_pipeline/pipeline_template/remote/test_remote_retrieval.py similarity index 98% rename from api/tests/unit_tests/services/rag_pipeline/pipeline_template/test_remote_retrieval.py rename to api/tests/unit_tests/services/rag_pipeline/pipeline_template/remote/test_remote_retrieval.py index d6ea59b71bf..f5f315e7893 100644 --- a/api/tests/unit_tests/services/rag_pipeline/pipeline_template/test_remote_retrieval.py +++ b/api/tests/unit_tests/services/rag_pipeline/pipeline_template/remote/test_remote_retrieval.py @@ -5,7 +5,7 @@ from sqlalchemy.orm import Session from services.rag_pipeline.pipeline_template.database.database_retrieval import DatabasePipelineTemplateRetrieval from services.rag_pipeline.pipeline_template.pipeline_template_type import PipelineTemplateType from services.rag_pipeline.pipeline_template.remote.remote_retrieval import RemotePipelineTemplateRetrieval -from tests.unit_tests.config_override import apply_config_overrides +from tests.helpers.unit.config_override import apply_config_overrides @pytest.mark.parametrize("sqlite_session", [()], indirect=True) diff --git a/api/tests/unit_tests/services/rag_pipeline/pipeline_template/test_package_imports.py b/api/tests/unit_tests/services/rag_pipeline/pipeline_template/test___init__.py similarity index 100% rename from api/tests/unit_tests/services/rag_pipeline/pipeline_template/test_package_imports.py rename to api/tests/unit_tests/services/rag_pipeline/pipeline_template/test___init__.py diff --git a/api/tests/unit_tests/services/rag_pipeline/test_rag_pipeline.py b/api/tests/unit_tests/services/rag_pipeline/test_rag_pipeline.py index dcbff0dbdf3..59d86a99ef6 100644 --- a/api/tests/unit_tests/services/rag_pipeline/test_rag_pipeline.py +++ b/api/tests/unit_tests/services/rag_pipeline/test_rag_pipeline.py @@ -1,40 +1,2733 @@ +import json +import time from collections.abc import Callable +from dataclasses import dataclass +from datetime import datetime +from types import SimpleNamespace +from unittest import mock +from unittest.mock import create_autospec +import pytest from pytest_mock import MockerFixture +from sqlalchemy.dialects import postgresql +from sqlalchemy.orm import Session, sessionmaker +from core.app.apps.pipeline.pipeline_generator import PipelineGenerator +from core.app.entities.app_invoke_entities import InvokeFrom +from core.rag.index_processor.constant.index_type import IndexStructureType +from graphon.engine_events import NodeRunFailedEvent +from graphon.enums import ( + BuiltinNodeTypes, + ErrorStrategy, + WorkflowNodeExecutionMetadataKey, + WorkflowNodeExecutionStatus, +) +from graphon.node_events.base import NodeRunResult +from graphon.nodes.start.entities import StartNodeData +from graphon.nodes.start.start_node import StartNode +from graphon.runtime import InitParams, RuntimeState, VariablePool +from models import Account, Tenant +from models.dataset import ( + Dataset, + Document, + DocumentPipelineExecutionLog, + Pipeline, + PipelineCustomizedTemplate, + PipelineRecommendedPlugin, +) +from models.enums import DataSourceType, DocumentCreatedFrom, IndexingStatus +from models.workflow import Workflow, WorkflowRun +from repositories.knowledge.dataset_read_repository import get_pipeline_dataset +from services.credentials.query import CredentialQuery +from services.data_source.credential_gateway import DatasourceProviderCredentialStore +from services.entities.knowledge_entities.rag_pipeline_entities import IconInfo, PipelineTemplateInfoEntity +from services.errors.rag_pipeline import RagPipelineResourceNotFoundError +from services.rag_pipeline import rag_pipeline as rag_pipeline_module from services.rag_pipeline.rag_pipeline import RagPipelineService +from services.workflow_ref_service import WorkflowRef def _make_service() -> RagPipelineService: return RagPipelineService.__new__(RagPipelineService) -def test_fetch_recommended_plugin_manifests_returns_empty_when_disabled( - mocker: MockerFixture, - config_overrides: Callable[..., None], -) -> None: - config_overrides(MARKETPLACE_ENABLED=False) - batch_fetch = mocker.patch("services.rag_pipeline.rag_pipeline.marketplace.batch_fetch_plugin_by_ids") +class TestRagPipelineCases: + @staticmethod + def test_fetch_recommended_plugin_manifests_returns_empty_when_disabled( + mocker: MockerFixture, + config_overrides: Callable[..., None], + ) -> None: + config_overrides(MARKETPLACE_ENABLED=False) + batch_fetch = mocker.patch("services.rag_pipeline.rag_pipeline.marketplace.batch_fetch_plugin_by_ids") - service = _make_service() - result = service._fetch_recommended_plugin_manifests(["langgenius/openai"]) + service = _make_service() + result = service._fetch_recommended_plugin_manifests(["langgenius/openai"]) - assert result == [] - batch_fetch.assert_not_called() + assert result == [] + batch_fetch.assert_not_called() + + @staticmethod + def test_fetch_recommended_plugin_manifests_returns_data_when_enabled( + mocker: MockerFixture, + config_overrides: Callable[..., None], + ) -> None: + config_overrides(MARKETPLACE_ENABLED=True) + expected = [{"plugin_id": "langgenius/openai", "name": "OpenAI"}] + mocker.patch( + "services.rag_pipeline.rag_pipeline.marketplace.batch_fetch_plugin_by_ids", + return_value=expected, + ) + + service = _make_service() + result = service._fetch_recommended_plugin_manifests(["langgenius/openai"]) + + assert result == expected -def test_fetch_recommended_plugin_manifests_returns_data_when_enabled( - mocker: MockerFixture, - config_overrides: Callable[..., None], -) -> None: - config_overrides(MARKETPLACE_ENABLED=True) - expected = [{"plugin_id": "langgenius/openai", "name": "OpenAI"}] - mocker.patch( - "services.rag_pipeline.rag_pipeline.marketplace.batch_fetch_plugin_by_ids", - return_value=expected, +@dataclass +class RagPipelineServiceTestContext: + service: RagPipelineService + session: Session + session_maker: sessionmaker[Session] + + +class MockRepo: + pass + + +def _make_account(account_id: str = "u1", tenant_id: str = "t1") -> Account: + account = Account(name="Test User", email=f"{account_id}@example.com") + account.id = account_id + tenant = Tenant(name="Test Tenant") + tenant.id = tenant_id + account._current_tenant = tenant + return account + + +def _make_pipeline( + *, + pipeline_id: str = "p1", + tenant_id: str = "t1", + workflow_id: str | None = None, + is_published: bool = False, +) -> Pipeline: + pipeline = Pipeline( + tenant_id=tenant_id, + name="Test Pipeline", + description="test", + workflow_id=workflow_id, + is_published=is_published, + ) + pipeline.id = pipeline_id + return pipeline + + +def _make_template_args(name: str = "New Template") -> dict[str, object]: + return {"name": name, "description": "Desc", "icon_info": {"icon": "star"}} + + +def _make_workflow( + *, + workflow_id: str = "wf-1", + tenant_id: str = "t1", + app_id: str = "p1", + graph: dict[str, object] | None = None, + features: dict[str, object] | None = None, + rag_pipeline_variables: list[dict[str, object]] | None = None, + created_by: str = "u1", + version: str = Workflow.VERSION_DRAFT, +) -> Workflow: + workflow = Workflow( + id=workflow_id, + tenant_id=tenant_id, + app_id=app_id, + type="workflow", + version=version, + marked_name="", + marked_comment="", + graph=json.dumps(graph or {"nodes": []}), + features=json.dumps(features or {}), + created_by=created_by, + created_at=datetime(2024, 1, 1), + updated_by=None, + updated_at=datetime(2024, 1, 1), + environment_variables=[], + conversation_variables=[], + rag_pipeline_variables=rag_pipeline_variables or [], + ) + return workflow + + +def _make_dataset(*, dataset_id: str = "d1", pipeline_id: str = "p1", tenant_id: str = "t1") -> Dataset: + dataset = Dataset( + id=dataset_id, + tenant_id=tenant_id, + name="Test Dataset", + created_by="u1", + pipeline_id=pipeline_id, + ) + return dataset + + +def _make_document(*, document_id: str = "doc-1", dataset_id: str = "d1", tenant_id: str = "t1") -> Document: + return Document( + id=document_id, + tenant_id=tenant_id, + dataset_id=dataset_id, + position=1, + data_source_type=DataSourceType.LOCAL_FILE, + batch="batch", + name="Document", + created_from=DocumentCreatedFrom.API, + created_by="u1", + indexing_status=IndexingStatus.WAITING, + doc_form=IndexStructureType.PARAGRAPH_INDEX, ) - service = _make_service() - result = service._fetch_recommended_plugin_manifests(["langgenius/openai"]) - assert result == expected +def _rag_variable(variable: str, *, belong_to_node_id: str = "shared") -> dict[str, object]: + return { + "variable": variable, + "belong_to_node_id": belong_to_node_id, + "type": "text-input", + "label": variable.title(), + } + + +def _persist(session: Session, *entities: object) -> None: + session.add_all(entities) + session.commit() + + +def _make_start_node(data: StartNodeData | None = None, variable_pool: VariablePool | None = None) -> StartNode: + return StartNode( + node_id="node-1", + data=data if data is not None else StartNodeData(title="Start"), + init_params=InitParams(workflow_id="wf-1", graph_config={}, run_context={}, call_depth=0), + runtime_state=RuntimeState( + workflow_id="wf-1", + variable_pool=variable_pool if variable_pool is not None else VariablePool(), + start_at=time.perf_counter(), + ), + ) + + +def _make_failed_published_node_run() -> tuple[StartNode, NodeRunFailedEvent]: + variable_pool = VariablePool() + for name, value in { + "invoke_from": InvokeFrom.PUBLISHED_PIPELINE, + "app_id": "pipeline-1", + "dataset_id": "dataset-1", + "document_id": "doc-1", + }.items(): + variable_pool.add(["sys", name], value) + node_instance = _make_start_node(variable_pool=variable_pool) + run_result = NodeRunResult( + status=WorkflowNodeExecutionStatus.FAILED, + error="boom", + error_type="runtime", + inputs={}, + outputs={}, + ) + event = NodeRunFailedEvent( + id="evt-1", + start_at=time.time(), + node_id="node-1", + node_type="start", + node_run_result=run_result, + error="boom", + route_node_id=None, + ) + return node_instance, event + + +def _make_customized_template() -> PipelineCustomizedTemplate: + return PipelineCustomizedTemplate( + tenant_id="t1", + name="old", + description="old", + chunk_structure="paragraph", + icon={}, + position=1, + yaml_content="", + install_count=0, + language="en-US", + created_by="u1", + ) + + +def _make_recommended_plugin(plugin_id: str) -> PipelineRecommendedPlugin: + return PipelineRecommendedPlugin(plugin_id=plugin_id, provider_name=plugin_id, type="tool", position=0, active=True) + + +class TestRagPipelineServiceCases: + @staticmethod + @pytest.fixture + def pipeline_generator(mocker: MockerFixture): + return mocker.create_autospec(PipelineGenerator, instance=True, spec_set=True) + + @staticmethod + @pytest.fixture + def rag_pipeline_service( + mocker: MockerFixture, + sqlite_session: Session, + sqlite_session_factory: sessionmaker[Session], + ) -> RagPipelineServiceTestContext: + mocker.patch( + "services.rag_pipeline.rag_pipeline.DifyAPIRepositoryFactory.create_api_workflow_node_execution_repository", + return_value=MockRepo(), + ) + mocker.patch( + "services.rag_pipeline.rag_pipeline.DifyAPIRepositoryFactory.create_api_workflow_run_repository", + return_value=MockRepo(), + ) + mocker.patch("services.rag_pipeline.rag_pipeline.db", SimpleNamespace(engine=sqlite_session.get_bind())) + service = RagPipelineService(session=sqlite_session, session_maker=sqlite_session_factory) + return RagPipelineServiceTestContext( + service=service, + session=sqlite_session, + session_maker=sqlite_session_factory, + ) + + @staticmethod + def test_get_published_workflow_by_id_locks_restore_source( + mocker: MockerFixture, rag_pipeline_service: RagPipelineServiceTestContext + ) -> None: + session = rag_pipeline_service.session + workflow = _make_workflow() + workflow.version = "v1" + session.add(workflow) + session.flush() + lookup = mocker.spy(session, "scalar") + service = rag_pipeline_service.service + + result = service.get_published_workflow_by_id(_make_pipeline(), workflow.id) + + stmt = lookup.call_args.args[0] + sql = str(stmt.compile(dialect=postgresql.dialect())) + assert result is workflow + assert "FOR UPDATE" in sql + + @staticmethod + def test_get_pipeline_templates_fallbacks_to_builtin_for_non_english_empty_result( + mocker: MockerFixture, sqlite_session: Session, config_overrides: Callable[..., None] + ) -> None: + config_overrides(HOSTED_FETCH_PIPELINE_TEMPLATES_MODE="remote") + session = sqlite_session + + remote_retrieval = mocker.Mock() + remote_retrieval.get_pipeline_templates.return_value = {"pipeline_templates": []} + + factory_mock = mocker.patch("services.rag_pipeline.rag_pipeline.PipelineTemplateRetrievalFactory") + factory_mock.get_pipeline_template_factory.return_value.return_value = remote_retrieval + + builtin_retrieval = mocker.Mock() + builtin_retrieval.fetch_pipeline_templates_from_builtin.return_value = { + "pipeline_templates": [{"id": "builtin-1"}] + } + factory_mock.get_built_in_pipeline_template_retrieval.return_value = builtin_retrieval + + result = RagPipelineService.get_pipeline_templates(type="built-in", language="ja-JP", session=session) + + assert result == {"pipeline_templates": [{"id": "builtin-1"}]} + remote_retrieval.get_pipeline_templates.assert_called_once_with("ja-JP", None, session=session) + builtin_retrieval.fetch_pipeline_templates_from_builtin.assert_called_once_with("en-US") + + @staticmethod + def test_get_pipeline_templates_customized_mode_uses_customized_factory( + mocker: MockerFixture, sqlite_session: Session + ) -> None: + session = sqlite_session + retrieval = mocker.Mock() + retrieval.get_pipeline_templates.return_value = {"pipeline_templates": [{"id": "custom-1"}]} + + factory_mock = mocker.patch("services.rag_pipeline.rag_pipeline.PipelineTemplateRetrievalFactory") + factory_mock.get_pipeline_template_factory.return_value.return_value = retrieval + + result = RagPipelineService.get_pipeline_templates(type="customized", language="en-US", session=session) + + assert result == {"pipeline_templates": [{"id": "custom-1"}]} + factory_mock.get_pipeline_template_factory.assert_called_with("customized") + retrieval.get_pipeline_templates.assert_called_once_with("en-US", None, session=session) + + @staticmethod + @pytest.mark.parametrize("template_type", ["built-in", "customized"]) + def test_get_pipeline_template_detail_uses_expected_mode( + mocker: MockerFixture, + template_type: str, + sqlite_session: Session, + config_overrides: Callable[..., None], + ) -> None: + config_overrides(HOSTED_FETCH_PIPELINE_TEMPLATES_MODE="remote") + session = sqlite_session + retrieval = mocker.Mock() + retrieval.get_pipeline_template_detail.return_value = {"id": "tpl-1"} + + factory_mock = mocker.patch("services.rag_pipeline.rag_pipeline.PipelineTemplateRetrievalFactory") + factory_mock.get_pipeline_template_factory.return_value.return_value = retrieval + + result = RagPipelineService.get_pipeline_template_detail( + "tpl-1", "tenant-1", type=template_type, session=session + ) + + assert result == {"id": "tpl-1"} + expected_mode = "remote" if template_type == "built-in" else "customized" + factory_mock.get_pipeline_template_factory.assert_called_with(expected_mode) + retrieval.get_pipeline_template_detail.assert_called_once_with("tpl-1", "tenant-1", session=session) + + @staticmethod + def test_get_published_workflow_returns_none_when_pipeline_has_no_workflow_id( + rag_pipeline_service: RagPipelineServiceTestContext, + ) -> None: + pipeline = _make_pipeline(workflow_id=None) + + result = rag_pipeline_service.service.get_published_workflow(pipeline) + + assert result is None + + @staticmethod + def test_get_all_published_workflow_returns_empty_for_unpublished_pipeline( + rag_pipeline_service: RagPipelineServiceTestContext, + ) -> None: + pipeline = _make_pipeline(workflow_id=None) + workflows, has_more = rag_pipeline_service.service.get_all_published_workflow( + session=rag_pipeline_service.session, + pipeline=pipeline, + page=1, + limit=20, + user_id=None, + named_only=False, + ) + + assert workflows == [] + assert has_more is False + + @staticmethod + def test_get_all_published_workflow_applies_limit_and_has_more( + rag_pipeline_service: RagPipelineServiceTestContext, + ) -> None: + pipeline = _make_pipeline(pipeline_id="pipeline-1", workflow_id="wf-live") + workflows = [ + _make_workflow( + workflow_id=f"wf-{index}", + app_id="pipeline-1", + created_by="user-1", + ) + for index in range(1, 4) + ] + for index, workflow in enumerate(workflows, start=1): + workflow.version = f"2024-01-0{index}T00:00:00" + workflow.marked_name = f"release-{index}" + _persist(rag_pipeline_service.session, *workflows) + + result, has_more = rag_pipeline_service.service.get_all_published_workflow( + session=rag_pipeline_service.session, + pipeline=pipeline, + page=1, + limit=2, + user_id="user-1", + named_only=True, + ) + + assert [workflow.id for workflow in result] == ["wf-3", "wf-2"] + assert has_more is True + + # --- sync_draft_workflow --- + + @staticmethod + def test_sync_draft_workflow_creates_new_when_none_exists( + rag_pipeline_service: RagPipelineServiceTestContext, + ) -> None: + pipeline = _make_pipeline(workflow_id=None) + account = _make_account() + + result = rag_pipeline_service.service.sync_draft_workflow( + pipeline=pipeline, + graph={"nodes": []}, + unique_hash=None, + account=account, + environment_variables=[], + conversation_variables=[], + rag_pipeline_variables=[], + ) + + assert result.app_id == "p1" + assert pipeline.workflow_id == result.id + + @staticmethod + def test_sync_draft_workflow_raises_on_hash_mismatch( + rag_pipeline_service: RagPipelineServiceTestContext, + ) -> None: + from services.errors.app import WorkflowHashNotEqualError + + existing_wf = _make_workflow(graph={"nodes": [{"id": "old"}]}) + _persist(rag_pipeline_service.session, existing_wf) + + pipeline = _make_pipeline() + account = _make_account() + + with pytest.raises(WorkflowHashNotEqualError): + rag_pipeline_service.service.sync_draft_workflow( + pipeline=pipeline, + graph={"nodes": []}, + unique_hash="hash-different", + account=account, + environment_variables=[], + conversation_variables=[], + rag_pipeline_variables=[], + ) + + @staticmethod + def test_sync_draft_workflow_updates_existing( + rag_pipeline_service: RagPipelineServiceTestContext, + ) -> None: + existing_wf = _make_workflow(graph={"nodes": [{"id": "old"}]}) + _persist(rag_pipeline_service.session, existing_wf) + + pipeline = _make_pipeline() + account = _make_account() + + result = rag_pipeline_service.service.sync_draft_workflow( + pipeline=pipeline, + graph={"nodes": [{"id": "n1"}]}, + unique_hash=existing_wf.unique_hash, + account=account, + environment_variables=[], + conversation_variables=[], + rag_pipeline_variables=[], + ) + + assert result is existing_wf + assert result.updated_by == "u1" + assert result.graph_dict == {"nodes": [{"id": "n1"}]} + + # --- get_default_block_config --- + + @staticmethod + def test_get_default_block_config_returns_config_for_valid_type( + mocker: MockerFixture, rag_pipeline_service: RagPipelineServiceTestContext + ) -> None: + fake_node_class = mocker.Mock() + fake_node_class.get_default_config.return_value = {"type": "start", "config": {}} + + # Use a simpler approach: test with a known valid node type + + mocker.patch( + "services.rag_pipeline.rag_pipeline.get_node_type_classes_mapping", + return_value={BuiltinNodeTypes.START: {"1": fake_node_class}}, + ) + mocker.patch("services.rag_pipeline.rag_pipeline.LATEST_VERSION", "1") + + result = rag_pipeline_service.service.get_default_block_config("start") + + assert result == {"type": "start", "config": {}} + + @staticmethod + def test_get_default_block_config_returns_none_for_unmapped_type( + rag_pipeline_service: RagPipelineServiceTestContext, + ) -> None: + assert rag_pipeline_service.service.get_default_block_config("nonexistent-type") is None + + # --- update_workflow --- + + @staticmethod + def test_update_workflow_updates_allowed_fields( + rag_pipeline_service: RagPipelineServiceTestContext, + ) -> None: + workflow = _make_workflow() + _persist(rag_pipeline_service.session, workflow) + workflow_ref = WorkflowRef(tenant_id="t1", owner_id="pipeline-1", workflow_id="wf-1") + workflow.app_id = "pipeline-1" + rag_pipeline_service.session.commit() + + result = rag_pipeline_service.service.update_workflow( + session=rag_pipeline_service.session, + account_id="u1", + data={"marked_name": "v1", "marked_comment": "release", "disallowed": "hacked"}, + workflow_ref=workflow_ref, + ) + + assert result.marked_name == "v1" + assert result.marked_comment == "release" + assert not hasattr(result, "disallowed") + assert result.updated_by == "u1" + + @staticmethod + def test_update_workflow_returns_none_when_not_found( + rag_pipeline_service: RagPipelineServiceTestContext, + ) -> None: + result = rag_pipeline_service.service.update_workflow( + session=rag_pipeline_service.session, + account_id="u1", + data={"marked_name": "v1"}, + workflow_ref=WorkflowRef(tenant_id="t1", owner_id="pipeline-1", workflow_id="wf-missing"), + ) + + assert result is None + + @staticmethod + def test_update_workflow_with_ref_scopes_lookup_to_pipeline( + rag_pipeline_service: RagPipelineServiceTestContext, + ) -> None: + workflow = _make_workflow(app_id="pipeline-1") + other_owner_workflow = _make_workflow(workflow_id="wf-other", app_id="pipeline-2") + _persist(rag_pipeline_service.session, workflow, other_owner_workflow) + workflow_ref = WorkflowRef(tenant_id="t1", owner_id="pipeline-1", workflow_id="wf-1") + + result = rag_pipeline_service.service.update_workflow( + session=rag_pipeline_service.session, + account_id="u1", + data={"marked_name": "v1"}, + workflow_ref=workflow_ref, + ) + + assert result is workflow + assert other_owner_workflow.marked_name == "" + + # --- get_rag_pipeline_paginate_workflow_runs --- + + @staticmethod + def test_get_rag_pipeline_paginate_workflow_runs_delegates( + mocker: MockerFixture, rag_pipeline_service: RagPipelineServiceTestContext + ) -> None: + expected = mocker.Mock() + repo_mock = mocker.Mock() + repo_mock.get_paginated_workflow_runs.return_value = expected + rag_pipeline_service.service._workflow_run_repo = repo_mock + + pipeline = _make_pipeline() + result = rag_pipeline_service.service.get_rag_pipeline_paginate_workflow_runs( + pipeline, {"limit": 10, "last_id": "abc"} + ) + + assert result is expected + repo_mock.get_paginated_workflow_runs.assert_called_once_with( + tenant_id="t1", + app_id="p1", + triggered_from=mocker.ANY, + limit=10, + last_id="abc", + ) + + # --- get_rag_pipeline_workflow_run --- + + @staticmethod + def test_get_rag_pipeline_workflow_run_delegates( + mocker: MockerFixture, rag_pipeline_service: RagPipelineServiceTestContext + ) -> None: + expected = mocker.Mock() + repo_mock = mocker.Mock() + repo_mock.get_workflow_run_by_id.return_value = expected + rag_pipeline_service.service._workflow_run_repo = repo_mock + + pipeline = _make_pipeline() + result = rag_pipeline_service.service.get_rag_pipeline_workflow_run(pipeline, "run-1") + + assert result is expected + repo_mock.get_workflow_run_by_id.assert_called_once_with(tenant_id="t1", app_id="p1", run_id="run-1") + + # --- is_workflow_exist --- + + @staticmethod + def test_is_workflow_exist_returns_true_when_draft_exists( + rag_pipeline_service: RagPipelineServiceTestContext, + ) -> None: + _persist(rag_pipeline_service.session, _make_workflow()) + + pipeline = _make_pipeline() + assert rag_pipeline_service.service.is_workflow_exist(pipeline) is True + + @staticmethod + def test_is_workflow_exist_returns_false_when_no_draft( + rag_pipeline_service: RagPipelineServiceTestContext, + ) -> None: + pipeline = _make_pipeline() + assert rag_pipeline_service.service.is_workflow_exist(pipeline) is False + + # --- publish_workflow --- + + @staticmethod + def test_publish_workflow_success( + mocker: MockerFixture, rag_pipeline_service: RagPipelineServiceTestContext + ) -> None: + graph = { + "nodes": [ + { + "data": { + "type": "knowledge-index", + "dataset_id": "d1", + "chunk_structure": "paragraph", + "indexing_technique": "high_quality", + "process_rule": {"mode": "automatic"}, + "retrieval_model": {"search_method": "hybrid_search", "top_k": 3}, + } + } + ] + } + draft_workflow = _make_workflow(workflow_id="wf-draft", graph=graph) + pipeline = _make_pipeline(workflow_id="wf-old-published") + dataset = _make_dataset() + _persist(rag_pipeline_service.session, draft_workflow, pipeline, dataset) + mocker.patch( + "services.rag_pipeline.rag_pipeline.KnowledgeConfiguration.model_validate", return_value=mocker.Mock() + ) + mock_dataset_service_class = mocker.patch("services.knowledge.dataset_service.DatasetService") + + result = rag_pipeline_service.service.publish_workflow( + session=rag_pipeline_service.session, pipeline=pipeline, account=_make_account() + ) + + assert result.app_id == pipeline.id + assert result.graph_dict == graph + mock_dataset_service_class.update_rag_pipeline_dataset_settings.assert_called_once() + + @staticmethod + def test_publish_workflow_rejects_missing_llm_environment_reference( + rag_pipeline_service: RagPipelineServiceTestContext, + ) -> None: + draft_workflow = _make_workflow( + workflow_id="wf-draft", + tenant_id="tenant", + app_id="pipeline", + graph={ + "nodes": [ + { + "id": "llm-node", + "data": { + "type": "llm", + "model": {"provider": "provider", "name": "model", "mode": "chat"}, + "model_selector": ["env", "missing_model"], + }, + } + ] + }, + ) + _persist(rag_pipeline_service.session, draft_workflow) + + with pytest.raises(ValueError, match="missing_model.*not found"): + rag_pipeline_service.service.publish_workflow( + session=rag_pipeline_service.session, + pipeline=_make_pipeline(pipeline_id="pipeline", tenant_id="tenant"), + account=_make_account(account_id="account", tenant_id="tenant"), + ) + + # --- run_datasource_workflow_node --- + + @staticmethod + def test_run_datasource_workflow_node_website_crawl( + mocker: MockerFixture, rag_pipeline_service: RagPipelineServiceTestContext + ) -> None: + from core.datasource.entities.datasource_entities import DatasourceProviderType + + # 1. Setup workflow and node + pipeline = _make_pipeline() + + workflow = _make_workflow( + graph={ + "nodes": [ + { + "id": "node-1", + "data": { + "type": "datasource", + "plugin_id": "p-1", + "provider_name": "firecrawl", + "datasource_name": "website_crawl", + "datasource_parameters": {"url": {"value": "{{#start.url#}}"}}, + }, + } + ] + } + ) + mocker.patch.object(rag_pipeline_service.service, "get_published_workflow", return_value=workflow) + + # 2. Mock DatasourceManager and Runtime + mock_runtime = mocker.Mock() + mock_runtime.datasource_provider_type.return_value = DatasourceProviderType.WEBSITE_CRAWL + + # Mock the generator result for website crawl + def mock_crawl_gen(**kwargs): + yield mocker.Mock(result=mocker.Mock(status="processing", total=10, completed=2)) + yield mocker.Mock( + result=mocker.Mock(status="completed", total=10, completed=10, web_info_list=[{"title": "test"}]) + ) + + mock_runtime.get_website_crawl.side_effect = mock_crawl_gen + + mocker.patch( + "core.datasource.datasource_manager.DatasourceManager.get_datasource_runtime", + return_value=mock_runtime, + ) + + # 3. Mock DatasourceProviderService + mocker.patch( + "services.rag_pipeline.rag_pipeline.DatasourceProviderService.get_datasource_credentials", + return_value={"api_key": "sk-123"}, + ) + + # 4. Mock Enums to avoid import issues or for consistency + mocker.patch("services.rag_pipeline.rag_pipeline.DatasourceProviderType", DatasourceProviderType) + + # 5. Run test + gen = rag_pipeline_service.service.run_datasource_workflow_node( + pipeline=pipeline, + node_id="node-1", + user_inputs={"url": "https://example.com"}, + account=_make_account(), + datasource_type="website_crawl", + is_published=True, + datasource_providers=rag_pipeline_module.DatasourceProviderService( + credentials=create_autospec(DatasourceProviderCredentialStore, instance=True) + ), + ) + + events = list(gen) + + # 6. Assertions + assert len(events) == 2 + assert events[0]["total"] == 10 + assert events[0]["completed"] == 2 + assert events[1]["data"] == [{"title": "test"}] + assert events[1]["total"] == 10 + assert events[1]["completed"] == 10 + + # --- run_datasource_node_preview --- + + @staticmethod + def test_run_datasource_node_preview_online_document( + mocker: MockerFixture, rag_pipeline_service: RagPipelineServiceTestContext + ) -> None: + from core.datasource.entities.datasource_entities import DatasourceMessage, DatasourceProviderType + + # 1. Setup workflow and node + pipeline = _make_pipeline() + + workflow = _make_workflow( + graph={ + "nodes": [ + { + "id": "node-1", + "data": { + "type": "datasource", + "plugin_id": "p-1", + "provider_name": "notion", + "datasource_name": "online_document", + "datasource_parameters": { + "workspace_id": {"value": "ws-1"}, + "page_id": {"value": "pg-1"}, + "type": {"value": "page"}, + }, + }, + } + ] + } + ) + mocker.patch.object(rag_pipeline_service.service, "get_published_workflow", return_value=workflow) + + # 2. Mock Runtime and results + mock_runtime = mocker.Mock() + + def mock_doc_gen(**kwargs): + # Yield a variable message + msg1 = DatasourceMessage( + type=DatasourceMessage.MessageType.VARIABLE, + message=DatasourceMessage.VariableMessage( + variable_name="content", variable_value="Hello ", stream=True + ), + ) + yield msg1 + msg2 = DatasourceMessage( + type=DatasourceMessage.MessageType.VARIABLE, + message=DatasourceMessage.VariableMessage(variable_name="content", variable_value="World", stream=True), + ) + yield msg2 + + mock_runtime.get_online_document_page_content.side_effect = mock_doc_gen + mocker.patch( + "core.datasource.datasource_manager.DatasourceManager.get_datasource_runtime", + return_value=mock_runtime, + ) + mocker.patch( + "services.rag_pipeline.rag_pipeline.DatasourceProviderService.get_datasource_credentials", + return_value={"token": "abc"}, + ) + mocker.patch("services.rag_pipeline.rag_pipeline.DatasourceProviderType", DatasourceProviderType) + + # 3. Run test + result = rag_pipeline_service.service.run_datasource_node_preview( + pipeline=pipeline, + node_id="node-1", + user_inputs={}, + account=_make_account(), + datasource_type="online_document", + is_published=True, + datasource_providers=rag_pipeline_module.DatasourceProviderService( + credentials=create_autospec(DatasourceProviderCredentialStore, instance=True) + ), + ) + + # 4. Assertions + assert result == {"content": "Hello World"} + + # --- _handle_node_run_result --- + + @staticmethod + def test_handle_node_run_result_success( + mocker: MockerFixture, rag_pipeline_service: RagPipelineServiceTestContext + ) -> None: + from graphon.engine_events import NodeRunSucceededEvent + from graphon.node_events.base import NodeRunResult + + # 1. Setup mock node and result + node_instance = mocker.Mock() + node_instance.workflow_id = "wf-1" + node_instance.node_type = "start" + node_instance.title = "Start" + + node_run_result = NodeRunResult( + status=WorkflowNodeExecutionStatus.SUCCEEDED, + inputs={"q": "hi"}, + outputs={"ans": "hello"}, + metadata={WorkflowNodeExecutionMetadataKey.TOTAL_TOKENS: 10}, + ) + + def mock_getter(): + event = NodeRunSucceededEvent( + id="event-1", + start_at=time.time(), + node_id="node-1", + node_type="start", + node_run_result=node_run_result, + route_node_id=None, + ) + yield event + + # 2. Run test + result = rag_pipeline_service.service._handle_node_run_result( + getter=lambda: (node_instance, mock_getter()), + start_at=time.perf_counter(), + tenant_id="t1", + node_id="node-1", + ) + + # 3. Assertions + assert result.status == WorkflowNodeExecutionStatus.SUCCEEDED + assert result.inputs == {"q": "hi"} + assert result.outputs == {"ans": "hello"} + assert result.metadata == {WorkflowNodeExecutionMetadataKey.TOTAL_TOKENS: 10} + + # --- get_first_step_parameters / get_second_step_parameters --- + + @staticmethod + def test_get_first_step_parameters_success( + mocker: MockerFixture, rag_pipeline_service: RagPipelineServiceTestContext + ) -> None: + # 1. Setup mock workflow + pipeline = _make_pipeline() + workflow = _make_workflow( + graph={ + "nodes": [{"id": "node-1", "data": {"datasource_parameters": {"url": {"value": "{{#start.url#}}"}}}}] + }, + rag_pipeline_variables=[_rag_variable("url")], + ) + mocker.patch.object(rag_pipeline_service.service, "get_published_workflow", return_value=workflow) + + # 2. Run test + result = rag_pipeline_service.service.get_first_step_parameters( + pipeline=pipeline, node_id="node-1", is_draft=False + ) + + # 3. Assertions + assert len(result) == 1 + assert result[0]["variable"] == "url" + + @staticmethod + def test_get_second_step_parameters_success( + mocker: MockerFixture, rag_pipeline_service: RagPipelineServiceTestContext + ) -> None: + # 1. Setup mock workflow + pipeline = _make_pipeline() + workflow = _make_workflow( + graph={ + "nodes": [ + { + "id": "node-1", + "data": {}, # Second step logic is slightly different in how it gets variables + } + ] + }, + rag_pipeline_variables=[_rag_variable("var1", belong_to_node_id="other-node")], + ) + mocker.patch.object(rag_pipeline_service.service, "get_published_workflow", return_value=workflow) + + # 2. Run test + result = rag_pipeline_service.service.get_second_step_parameters( + pipeline=pipeline, node_id="node-1", is_draft=False + ) + + # 3. Assertions + # Note: get_second_step_parameters also filters by variable names found in node data + # (Checking the code again, it seems to iterate through nodes but doesn't do much with variables yet) + # Wait, let me check the code for get_second_step_parameters again. + assert len(result) == 0 # Based on current implementation which seems to filter but no logic added yet? + + # --- publish_customized_pipeline_template --- + + @staticmethod + def test_publish_customized_pipeline_template_success( + mocker: MockerFixture, rag_pipeline_service: RagPipelineServiceTestContext + ) -> None: + pipeline = _make_pipeline(workflow_id="wf-1", is_published=True) + workflow = _make_workflow(workflow_id="wf-1") + session = rag_pipeline_service.session + dataset = _make_dataset() + dataset.chunk_structure = "paragraph" + _persist(session, pipeline, workflow, dataset) + + # Mock RagPipelineDslService + mock_dsl_service = mocker.Mock() + mock_dsl_service.export_rag_pipeline_dsl.return_value = "dsl: content" + mocker.patch.object(rag_pipeline_module, "RagPipelineDslService", return_value=mock_dsl_service) + + account = _make_account(account_id="user-123") + + rag_pipeline_service.service.publish_customized_pipeline_template( + pipeline, dataset, _make_template_args(), account, session=session + ) + + mock_dsl_service.export_rag_pipeline_dsl.assert_called_once_with(pipeline=pipeline, include_secret=True) + templates = session.query(PipelineCustomizedTemplate).all() + assert len(templates) == 1 + assert templates[0].name == "New Template" + + # --- get_datasource_plugins --- + + @staticmethod + def test_get_datasource_plugins_success( + mocker: MockerFixture, rag_pipeline_service: RagPipelineServiceTestContext + ) -> None: + # 1. Setup mocks + dataset = _make_dataset() + + pipeline = _make_pipeline(workflow_id="wf-1") + + workflow = _make_workflow( + workflow_id="wf-1", + graph={ + "nodes": [ + { + "id": "node-1", + "data": { + "type": "datasource", + "plugin_id": "p-1", + "provider_name": "notion", + "provider_type": "online_document", + "title": "Notion", + }, + } + ] + }, + ) + _persist(rag_pipeline_service.session, dataset, pipeline, workflow) + + # Mock DatasourceProviderService + mock_provider_service = mocker.Mock() + mock_provider_service.list_datasource_credentials.return_value = [ + {"id": "c1", "name": "Cred 1", "type": "token", "is_default": True} + ] + mocker.patch("services.rag_pipeline.rag_pipeline.DatasourceProviderService", return_value=mock_provider_service) + + # 2. Run test + result = rag_pipeline_service.service.get_datasource_plugins( + "t1", + "d1", + True, + credential_query=mock.create_autospec(CredentialQuery, instance=True), + datasource_providers=rag_pipeline_module.DatasourceProviderService( + credentials=create_autospec(DatasourceProviderCredentialStore, instance=True) + ), + ) + + # 3. Assertions + assert len(result) == 1 + assert result[0]["node_id"] == "node-1" + assert result[0]["credentials"][0]["id"] == "c1" + + # --- retry_error_document --- + + @staticmethod + def test_retry_error_document_success( + mocker: MockerFixture, rag_pipeline_service: RagPipelineServiceTestContext, pipeline_generator + ) -> None: + dataset = _make_dataset() + document = _make_document() + + log = DocumentPipelineExecutionLog( + pipeline_id="p-1", + document_id="doc-1", + datasource_type="upload_file", + datasource_info="{}", + datasource_node_id="node-id", + input_data={}, + created_by="account-id", + ) + # Ensure it's a string if it's used as JSON later + + pipeline = Pipeline(tenant_id="tenant-id", name="Test Pipeline", workflow_id="wf-1") + pipeline.id = "p-1" + + workflow = _make_workflow(workflow_id="wf-1", tenant_id="tenant-id", app_id="p-1") + _persist(rag_pipeline_service.session, log, pipeline, workflow) + + # Mock PipelineGenerator + mock_gen_instance = pipeline_generator + + # 2. Run test + user = mocker.Mock() + rag_pipeline_service.service.retry_error_document(dataset, document, user, generator=pipeline_generator) + + # 3. Assertions + mock_gen_instance.generate.assert_called_once() + + # --- set_datasource_variables --- + + @staticmethod + def test_set_datasource_variables_success( + mocker: MockerFixture, rag_pipeline_service: RagPipelineServiceTestContext + ) -> None: + from graphon.entities.workflow_node_execution import WorkflowNodeExecution + from models.dataset import Pipeline + + # 1. Setup mocks + pipeline = Pipeline( + tenant_id="t1", + name="Test Pipeline", + ) + pipeline.id = "p-1" + + draft_wf = mocker.Mock() + draft_wf.id = "wf-1" + draft_wf.get_enclosing_node_type_and_id.return_value = None # Avoid unpacking error + mocker.patch.object(rag_pipeline_service.service, "get_draft_workflow", return_value=draft_wf) + + execution = WorkflowNodeExecution( + id="exec-1", + node_execution_id="node-exec-1", + workflow_id="wf-1", + workflow_execution_id="run-1", + index=1, + predecessor_node_id=None, + node_id="node-1", + node_type=BuiltinNodeTypes.DATASOURCE, + title="Datasource", + process_data={}, + inputs={}, + outputs={}, + status=WorkflowNodeExecutionStatus.SUCCEEDED, + error=None, + elapsed_time=0, + metadata={}, + created_at=datetime(2026, 1, 1), + finished_at=None, + ) + mocker.patch.object(rag_pipeline_service.service, "_handle_node_run_result", return_value=execution) + + # Mock Repository + mock_repo_instance = mocker.Mock() + mocker.patch( + "services.rag_pipeline.rag_pipeline.SQLAlchemyWorkflowNodeExecutionRepository", + return_value=mock_repo_instance, + ) + # Repository._to_db_model is also called + mock_db_exec = mocker.Mock() + mock_db_exec.node_id = "node-1" + mock_db_exec.node_type = "datasource" + mock_repo_instance._to_db_model.return_value = mock_db_exec + + # Mock DraftVariableSaver + mock_saver_instance = mocker.Mock() + mocker.patch("services.rag_pipeline.rag_pipeline.DraftVariableSaver", return_value=mock_saver_instance) + + # 2. Run test + args = {"start_node_id": "node-1"} + user = mocker.Mock() + user.id = "user-1" + rag_pipeline_service.service.set_datasource_variables(pipeline, args, user) + + # 3. Assertions + mock_repo_instance.save.assert_called_once() + mock_saver_instance.save.assert_called_once() + + # --- Utility Methods --- + + @staticmethod + def test_get_draft_workflow_success(rag_pipeline_service: RagPipelineServiceTestContext) -> None: + # 1. Setup mocks + pipeline = _make_pipeline() + + workflow = _make_workflow() + _persist(rag_pipeline_service.session, workflow) + + # 2. Run test + result = rag_pipeline_service.service.get_draft_workflow(pipeline) + + # 3. Assertions + assert result == workflow + + @staticmethod + def test_get_published_workflow_success(rag_pipeline_service: RagPipelineServiceTestContext) -> None: + # 1. Setup mocks + pipeline = _make_pipeline(workflow_id="wf-pub") + + workflow = _make_workflow(workflow_id="wf-pub") + _persist(rag_pipeline_service.session, workflow) + + # 2. Run test + result = rag_pipeline_service.service.get_published_workflow(pipeline) + + # 3. Assertions + assert result == workflow + + @staticmethod + def test_get_default_block_configs_success(rag_pipeline_service: RagPipelineServiceTestContext) -> None: + # This calls static methods on node classes, should be safe with default mocks or as-is + # unless they access db. + result = rag_pipeline_service.service.get_default_block_configs() + assert isinstance(result, list) + assert len(result) > 0 + + @staticmethod + def test_get_default_block_config_success(rag_pipeline_service: RagPipelineServiceTestContext) -> None: + result = rag_pipeline_service.service.get_default_block_config(BuiltinNodeTypes.LLM) + assert result is not None + assert result["type"] == "llm" + + @staticmethod + def test_publish_workflow_raises_when_draft_workflow_missing( + rag_pipeline_service: RagPipelineServiceTestContext, + ) -> None: + pipeline = _make_pipeline() + account = _make_account() + + with pytest.raises(ValueError, match="No valid workflow found"): + rag_pipeline_service.service.publish_workflow( + session=rag_pipeline_service.session, pipeline=pipeline, account=account + ) + + @staticmethod + def test_get_default_block_config_returns_none_when_mapped_type_missing( + mocker: MockerFixture, rag_pipeline_service: RagPipelineServiceTestContext + ) -> None: + mocker.patch("services.rag_pipeline.rag_pipeline.get_node_type_classes_mapping", return_value={}) + + assert rag_pipeline_service.service.get_default_block_config(BuiltinNodeTypes.START) is None + + @staticmethod + def test_get_default_block_config_injects_http_request_filter( + mocker: MockerFixture, rag_pipeline_service: RagPipelineServiceTestContext + ) -> None: + fake_node_cls = mocker.Mock() + fake_node_cls.get_default_config.return_value = {"type": "http-request"} + mocker.patch( + "services.rag_pipeline.rag_pipeline.get_node_type_classes_mapping", + return_value={BuiltinNodeTypes.HTTP_REQUEST: {"1": fake_node_cls}}, + ) + mocker.patch("services.rag_pipeline.rag_pipeline.LATEST_VERSION", "1") + + rag_pipeline_service.service.get_default_block_config(BuiltinNodeTypes.HTTP_REQUEST) + + called_filters = fake_node_cls.get_default_config.call_args.kwargs["filters"] + assert "http_request_config" in called_filters + + @staticmethod + def test_run_draft_workflow_node_raises_when_workflow_missing( + mocker: MockerFixture, rag_pipeline_service: RagPipelineServiceTestContext + ) -> None: + pipeline = _make_pipeline() + account = _make_account() + mocker.patch.object(rag_pipeline_service.service, "get_draft_workflow", return_value=None) + + with pytest.raises(ValueError, match="Workflow not initialized"): + rag_pipeline_service.service.run_draft_workflow_node(pipeline, "node-1", {}, account) + + @staticmethod + def test_run_draft_workflow_node_seeds_llm_environment_variable( + mocker: MockerFixture, rag_pipeline_service: RagPipelineServiceTestContext + ) -> None: + from factories import variable_factory + + pipeline = _make_pipeline() + account = _make_account() + llm_environment_variable = variable_factory.build_environment_variable_from_mapping( + { + "id": "env-1", + "name": "for_summarize", + "value_type": "llm", + "value": { + "provider": "langgenius/openai/openai", + "name": "gpt-4o", + "mode": "chat", + }, + "description": "Shared summarization model", + } + ) + draft_workflow = _make_workflow(workflow_id="wf-1") + draft_workflow.environment_variables = [llm_environment_variable] + mocker.patch.object(draft_workflow, "get_node_config_by_id", return_value={"id": "node-1"}) + mocker.patch.object(draft_workflow, "get_enclosing_node_type_and_id", return_value=None) + mocker.patch.object(rag_pipeline_service.service, "get_draft_workflow", return_value=draft_workflow) + + execution = SimpleNamespace(id="exec-1", node_id="node-1", node_type="llm", process_data={}, outputs={}) + handle_node_run_result = mocker.patch.object( + rag_pipeline_service.service, "_handle_node_run_result", return_value=execution + ) + single_step_run = mocker.patch("services.rag_pipeline.rag_pipeline.WorkflowEntry.single_step_run") + + repo = mocker.Mock() + mocker.patch( + "services.rag_pipeline.rag_pipeline.DifyCoreRepositoryFactory.create_workflow_node_execution_repository", + return_value=repo, + ) + rag_pipeline_service.service._node_execution_service_repo = mocker.Mock( + get_execution_by_id=mocker.Mock(return_value="db") + ) + mocker.patch("services.rag_pipeline.rag_pipeline.DraftVariableSaver", return_value=mocker.Mock()) + rag_pipeline_service.service.run_draft_workflow_node(pipeline, "node-1", {}, account) + + getter = handle_node_run_result.call_args.kwargs["getter"] + getter() + variable_pool = single_step_run.call_args.kwargs["variable_pool"] + llm_model = variable_pool.get(["env", "for_summarize"]) + assert llm_model is not None + assert llm_model.value == { + "provider": "langgenius/openai/openai", + "name": "gpt-4o", + "mode": "chat", + } + + @staticmethod + def test_run_draft_workflow_node_saves_execution_and_variables( + mocker: MockerFixture, rag_pipeline_service: RagPipelineServiceTestContext + ) -> None: + pipeline = _make_pipeline() + account = _make_account() + draft_workflow = _make_workflow(workflow_id="wf-1") + mocker.patch.object(draft_workflow, "get_node_config_by_id", return_value={"id": "node-1"}) + mocker.patch.object(draft_workflow, "get_enclosing_node_type_and_id", return_value=("loop", "enclosing-node")) + mocker.patch.object(rag_pipeline_service.service, "get_draft_workflow", return_value=draft_workflow) + + execution = SimpleNamespace(id="exec-1", node_id="node-1", node_type="llm", process_data={}, outputs={}) + mocker.patch.object(rag_pipeline_service.service, "_handle_node_run_result", return_value=execution) + + repo = mocker.Mock() + mocker.patch( + "services.rag_pipeline.rag_pipeline.DifyCoreRepositoryFactory.create_workflow_node_execution_repository", + return_value=repo, + ) + rag_pipeline_service.service._node_execution_service_repo = mocker.Mock( + get_execution_by_id=mocker.Mock(return_value="db") + ) + saver = mocker.Mock() + mocker.patch("services.rag_pipeline.rag_pipeline.DraftVariableSaver", return_value=saver) + + result = rag_pipeline_service.service.run_draft_workflow_node(pipeline, "node-1", {"q": "x"}, account) + + assert result == "db" + assert execution.workflow_id == "wf-1" + repo.save.assert_called_once_with(execution) + saver.save.assert_called_once() + + @staticmethod + def test_run_datasource_workflow_node_returns_error_when_workflow_missing( + mocker: MockerFixture, rag_pipeline_service: RagPipelineServiceTestContext + ) -> None: + pipeline = _make_pipeline() + mocker.patch.object(rag_pipeline_service.service, "get_draft_workflow", return_value=None) + + events = list( + rag_pipeline_service.service.run_datasource_workflow_node( + pipeline=pipeline, + node_id="node-1", + user_inputs={}, + account=_make_account(), + datasource_type="online_document", + is_published=False, + datasource_providers=rag_pipeline_module.DatasourceProviderService( + credentials=create_autospec(DatasourceProviderCredentialStore, instance=True) + ), + ) + ) + + assert events[0]["event"] == "datasource_error" + + @staticmethod + def test_run_datasource_workflow_node_online_document_success( + mocker: MockerFixture, rag_pipeline_service: RagPipelineServiceTestContext + ) -> None: + from core.datasource.entities.datasource_entities import DatasourceProviderType + + pipeline = _make_pipeline() + workflow = _make_workflow( + graph={ + "nodes": [ + { + "id": "node-1", + "data": { + "type": "datasource", + "plugin_id": "pid", + "provider_name": "notion", + "datasource_name": "online_document", + "datasource_parameters": {"workspace_id": {"value": None}, "page_id": {"value": "fixed"}}, + }, + } + ] + } + ) + mocker.patch.object(rag_pipeline_service.service, "get_published_workflow", return_value=workflow) + + runtime = mocker.Mock() + runtime.runtime = SimpleNamespace(credentials=None) + runtime.datasource_provider_type.return_value = DatasourceProviderType.ONLINE_DOCUMENT + runtime.get_online_document_pages.return_value = [SimpleNamespace(result=[{"id": "pg-1"}])] + mocker.patch( + "core.datasource.datasource_manager.DatasourceManager.get_datasource_runtime", return_value=runtime + ) + mocker.patch( + "services.rag_pipeline.rag_pipeline.DatasourceProviderService.get_datasource_credentials", + return_value={"token": "x"}, + ) + + events = list( + rag_pipeline_service.service.run_datasource_workflow_node( + pipeline=pipeline, + node_id="node-1", + user_inputs={}, + account=_make_account(), + datasource_type=DatasourceProviderType.ONLINE_DOCUMENT, + is_published=True, + datasource_providers=rag_pipeline_module.DatasourceProviderService( + credentials=create_autospec(DatasourceProviderCredentialStore, instance=True) + ), + ) + ) + + assert events[0]["event"] == "datasource_processing" + assert events[1]["event"] == "datasource_completed" + + @staticmethod + def test_run_datasource_workflow_node_online_drive_success( + mocker: MockerFixture, rag_pipeline_service: RagPipelineServiceTestContext + ) -> None: + from core.datasource.entities.datasource_entities import DatasourceProviderType + + pipeline = _make_pipeline() + workflow = _make_workflow( + graph={ + "nodes": [ + { + "id": "node-1", + "data": { + "type": "datasource", + "plugin_id": "pid", + "provider_name": "drive", + "datasource_name": "online_drive", + "datasource_parameters": { + "bucket": {"value": "bucket-1"}, + "next_page_parameters": {"value": []}, + }, + }, + } + ] + } + ) + mocker.patch.object(rag_pipeline_service.service, "get_published_workflow", return_value=workflow) + + runtime = mocker.Mock() + runtime.runtime = SimpleNamespace(credentials=None) + runtime.datasource_provider_type.return_value = DatasourceProviderType.ONLINE_DRIVE + runtime.online_drive_browse_files.return_value = [SimpleNamespace(result=[{"name": "f1"}])] + mocker.patch( + "core.datasource.datasource_manager.DatasourceManager.get_datasource_runtime", return_value=runtime + ) + mocker.patch( + "services.rag_pipeline.rag_pipeline.DatasourceProviderService.get_datasource_credentials", + return_value={}, + ) + + events = list( + rag_pipeline_service.service.run_datasource_workflow_node( + pipeline=pipeline, + node_id="node-1", + user_inputs={"bucket": "bucket-1"}, + account=_make_account(), + datasource_type=DatasourceProviderType.ONLINE_DRIVE, + is_published=True, + datasource_providers=rag_pipeline_module.DatasourceProviderService( + credentials=create_autospec(DatasourceProviderCredentialStore, instance=True) + ), + ) + ) + + assert events[0]["event"] == "datasource_processing" + assert events[1]["event"] == "datasource_completed" + + @staticmethod + def test_handle_node_run_result_default_value_strategy( + mocker: MockerFixture, rag_pipeline_service: RagPipelineServiceTestContext + ) -> None: + from datetime import datetime + + from graphon.engine_events import NodeRunFailedEvent + from graphon.node_events.base import NodeRunResult + + node_instance = _make_start_node( + StartNodeData( + title="Start", + error_strategy=ErrorStrategy.DEFAULT_VALUE, + default_value=[{"key": "fallback", "value": "ok", "type": "string"}], + ) + ) + + failed_result = NodeRunResult( + status=WorkflowNodeExecutionStatus.FAILED, + error="boom", + error_type="runtime_error", + inputs={"x": 1}, + ) + + def _events(): + yield NodeRunFailedEvent( + id="e-1", + node_id="node-1", + node_type=BuiltinNodeTypes.START, + start_at=datetime.now(), + error="boom", + node_run_result=failed_result, + ) + + result = rag_pipeline_service.service._handle_node_run_result( + getter=lambda: (node_instance, _events()), + start_at=time.perf_counter(), + tenant_id="t1", + node_id="node-1", + ) + + assert result.status == WorkflowNodeExecutionStatus.EXCEPTION + assert result.outputs + assert result.outputs["fallback"] == "ok" + + @staticmethod + def test_get_first_step_parameters_raises_when_datasource_node_missing( + mocker: MockerFixture, rag_pipeline_service: RagPipelineServiceTestContext + ) -> None: + workflow = _make_workflow(graph={"nodes": []}, rag_pipeline_variables=[_rag_variable("url")]) + mocker.patch.object(rag_pipeline_service.service, "get_published_workflow", return_value=workflow) + + with pytest.raises(ValueError, match="Datasource node data not found"): + rag_pipeline_service.service.get_first_step_parameters(_make_pipeline(), "missing-node") + + @staticmethod + def test_get_second_step_parameters_handles_string_and_list_variable_references( + mocker: MockerFixture, rag_pipeline_service: RagPipelineServiceTestContext + ) -> None: + workflow = _make_workflow( + rag_pipeline_variables=[ + _rag_variable("url", belong_to_node_id="node-1"), + _rag_variable("bucket"), + _rag_variable("keep", belong_to_node_id="node-1"), + ], + graph={ + "nodes": [ + { + "id": "node-1", + "data": { + "datasource_parameters": { + "u": {"value": "{{#start.url#}}"}, + "b": {"value": ["start", "bucket"]}, + } + }, + } + ] + }, + ) + mocker.patch.object(rag_pipeline_service.service, "get_published_workflow", return_value=workflow) + + result = rag_pipeline_service.service.get_second_step_parameters(_make_pipeline(), "node-1") + + assert [item["variable"] for item in result] == ["keep"] + assert result[0]["belong_to_node_id"] == "node-1" + + @staticmethod + def test_get_rag_pipeline_workflow_run_node_executions_empty_when_run_missing( + mocker: MockerFixture, rag_pipeline_service: RagPipelineServiceTestContext + ) -> None: + pipeline = _make_pipeline() + mocker.patch.object(rag_pipeline_service.service, "get_rag_pipeline_workflow_run", return_value=None) + + result = rag_pipeline_service.service.get_rag_pipeline_workflow_run_node_executions( + pipeline=pipeline, run_id="run-1", user=_make_account() + ) + + assert result == [] + + @staticmethod + def test_get_rag_pipeline_workflow_run_node_executions_assembles_configured_repository_results( + mocker: MockerFixture, rag_pipeline_service: RagPipelineServiceTestContext + ) -> None: + pipeline = _make_pipeline() + mocker.patch.object( + rag_pipeline_service.service, "get_rag_pipeline_workflow_run", return_value=WorkflowRun(id="run-1") + ) + node_repo = rag_pipeline_service.service._node_execution_service_repo + expected_executions = ["n1", "n2"] + expected_traces = ["retry-n1", "n1", "n2"] + node_repo.get_executions_by_workflow_run = mocker.Mock(return_value=expected_executions) + mock_assemble = mocker.patch( + "services.rag_pipeline.rag_pipeline.assemble_workflow_node_execution_traces", + return_value=expected_traces, + ) + + result = rag_pipeline_service.service.get_rag_pipeline_workflow_run_node_executions( + pipeline=pipeline, run_id="run-1", user=_make_account() + ) + + assert result == expected_traces + node_repo.get_executions_by_workflow_run.assert_called_once_with( + tenant_id=pipeline.tenant_id, + app_id=pipeline.id, + workflow_run_id="run-1", + ) + mock_assemble.assert_called_once_with(expected_executions, node_repo, session=mock.ANY) + + @staticmethod + def test_get_recommended_plugins_returns_empty_when_no_active_plugins( + rag_pipeline_service: RagPipelineServiceTestContext, + ) -> None: + result = rag_pipeline_service.service.get_recommended_plugins("all", _make_account(), "t1") + + assert result == { + "installed_recommended_plugins": [], + "uninstalled_recommended_plugins": [], + } + + @staticmethod + def test_get_recommended_plugins_returns_installed_and_uninstalled( + mocker: MockerFixture, rag_pipeline_service: RagPipelineServiceTestContext + ) -> None: + plugin_a = _make_recommended_plugin("plugin-a") + plugin_b = _make_recommended_plugin("plugin-b") + _persist(rag_pipeline_service.session, plugin_a, plugin_b) + mocker.patch( + "services.rag_pipeline.rag_pipeline.BuiltinToolManageService.list_builtin_tools", + return_value=[SimpleNamespace(plugin_id="plugin-a", to_dict=lambda: {"plugin_id": "plugin-a"})], + ) + mocker.patch( + "services.rag_pipeline.rag_pipeline.marketplace.batch_fetch_plugin_by_ids", + return_value=[{"plugin_id": "plugin-b", "name": "Plugin B"}], + ) + + result = rag_pipeline_service.service.get_recommended_plugins("tool", _make_account(), "t1") + + assert result["installed_recommended_plugins"] == [{"plugin_id": "plugin-a"}] + assert result["uninstalled_recommended_plugins"] == [{"plugin_id": "plugin-b", "name": "Plugin B"}] + + @staticmethod + def test_get_node_last_run_delegates_to_repository( + mocker: MockerFixture, rag_pipeline_service: RagPipelineServiceTestContext + ) -> None: + repo = mocker.Mock() + repo.get_node_last_execution.return_value = "node-exec" + rag_pipeline_service.service._node_execution_service_repo = repo + pipeline = _make_pipeline() + workflow = _make_workflow(workflow_id="wf1") + + result = rag_pipeline_service.service.get_node_last_run(pipeline, workflow, "node-1") + + assert result == "node-exec" + + @staticmethod + def test_set_datasource_variables_raises_when_node_id_missing( + mocker: MockerFixture, rag_pipeline_service: RagPipelineServiceTestContext + ) -> None: + pipeline = _make_pipeline() + workflow = _make_workflow() + mocker.patch.object(rag_pipeline_service.service, "get_draft_workflow", return_value=workflow) + + with pytest.raises(ValueError, match="Node id is required"): + rag_pipeline_service.service.set_datasource_variables(pipeline, {"start_node_id": ""}, _make_account()) + + @staticmethod + def test_get_default_block_configs_skips_empty_configs( + mocker: MockerFixture, rag_pipeline_service: RagPipelineServiceTestContext + ) -> None: + http_node = mocker.Mock() + http_node.get_default_config.return_value = {"type": "http-request"} + empty_node = mocker.Mock() + empty_node.get_default_config.return_value = None + + mocker.patch( + "services.rag_pipeline.rag_pipeline.get_node_type_classes_mapping", + return_value={ + BuiltinNodeTypes.HTTP_REQUEST: {"1": http_node}, + BuiltinNodeTypes.START: {"1": empty_node}, + }, + ) + mocker.patch("services.rag_pipeline.rag_pipeline.LATEST_VERSION", "1") + + result = rag_pipeline_service.service.get_default_block_configs() + + assert result == [{"type": "http-request"}] + http_node.get_default_config.assert_called_once() + empty_node.get_default_config.assert_called_once() + + @staticmethod + def test_run_datasource_workflow_node_returns_error_when_node_missing( + mocker: MockerFixture, rag_pipeline_service: RagPipelineServiceTestContext + ) -> None: + pipeline = _make_pipeline() + workflow = _make_workflow(graph={"nodes": []}) + mocker.patch.object(rag_pipeline_service.service, "get_published_workflow", return_value=workflow) + + events = list( + rag_pipeline_service.service.run_datasource_workflow_node( + pipeline=pipeline, + node_id="missing-node", + user_inputs={}, + account=_make_account(), + datasource_type="online_document", + is_published=True, + datasource_providers=rag_pipeline_module.DatasourceProviderService( + credentials=create_autospec(DatasourceProviderCredentialStore, instance=True) + ), + ) + ) + + assert len(events) == 1 + assert "Datasource node data not found" in events[0]["error"] + + @staticmethod + def test_run_datasource_workflow_node_online_document_exception( + mocker: MockerFixture, rag_pipeline_service: RagPipelineServiceTestContext + ) -> None: + pipeline = _make_pipeline() + workflow = _make_workflow( + graph={ + "nodes": [ + { + "id": "node-1", + "data": { + "plugin_id": "plugin-1", + "provider_name": "provider-1", + "datasource_name": "doc", + "datasource_parameters": {}, + }, + } + ] + } + ) + mocker.patch.object(rag_pipeline_service.service, "get_published_workflow", return_value=workflow) + + runtime = mocker.Mock() + + class _FailingIterator: + def __iter__(self): + return self + + def __next__(self): + raise RuntimeError("doc failed") + + runtime.get_online_document_pages.return_value = _FailingIterator() + runtime.datasource_provider_type.return_value = "online_document" + + mocker.patch( + "core.datasource.datasource_manager.DatasourceManager.get_datasource_runtime", return_value=runtime + ) + mocker.patch( + "services.rag_pipeline.rag_pipeline.DatasourceProviderService.get_datasource_credentials", + return_value=None, + ) + + events = list( + rag_pipeline_service.service.run_datasource_workflow_node( + pipeline=pipeline, + node_id="node-1", + user_inputs={}, + account=_make_account(), + datasource_type="online_document", + is_published=True, + datasource_providers=rag_pipeline_module.DatasourceProviderService( + credentials=create_autospec(DatasourceProviderCredentialStore, instance=True) + ), + ) + ) + + assert len(events) == 2 + assert events[0]["event"] == "datasource_processing" + assert "doc failed" in events[1]["error"] + + @staticmethod + def test_run_datasource_node_preview_raises_for_stream_non_string( + mocker: MockerFixture, rag_pipeline_service: RagPipelineServiceTestContext + ) -> None: + from core.datasource.entities.datasource_entities import DatasourceMessage + + pipeline = _make_pipeline() + workflow = _make_workflow( + graph={ + "nodes": [ + { + "id": "node-1", + "data": { + "plugin_id": "plugin-1", + "provider_name": "provider-1", + "datasource_name": "doc", + "datasource_parameters": {}, + }, + } + ] + } + ) + mocker.patch.object(rag_pipeline_service.service, "get_published_workflow", return_value=workflow) + + runtime = mocker.Mock() + + def _bad_stream_generator[**P](*args: P.args, **kwargs: P.kwargs): + yield DatasourceMessage( + type=DatasourceMessage.MessageType.VARIABLE, + message=DatasourceMessage.VariableMessage(variable_name="content", variable_value=1, stream=True), + ) + + runtime.get_online_document_page_content.side_effect = _bad_stream_generator + runtime.datasource_provider_type.return_value = "online_document" + + mocker.patch( + "core.datasource.datasource_manager.DatasourceManager.get_datasource_runtime", return_value=runtime + ) + mocker.patch( + "services.rag_pipeline.rag_pipeline.DatasourceProviderService.get_datasource_credentials", + return_value=None, + ) + + with pytest.raises(RuntimeError, match="must be a string"): + rag_pipeline_service.service.run_datasource_node_preview( + pipeline=pipeline, + node_id="node-1", + user_inputs={}, + account=_make_account(), + datasource_type="online_document", + is_published=True, + datasource_providers=rag_pipeline_module.DatasourceProviderService( + credentials=create_autospec(DatasourceProviderCredentialStore, instance=True) + ), + ) + + @staticmethod + def test_get_first_step_parameters_returns_empty_when_no_rag_variables( + mocker: MockerFixture, rag_pipeline_service: RagPipelineServiceTestContext + ) -> None: + workflow = _make_workflow( + graph={"nodes": [{"id": "node-1", "data": {"datasource_parameters": {"url": {"value": "literal"}}}}]}, + rag_pipeline_variables=[], + ) + mocker.patch.object(rag_pipeline_service.service, "get_published_workflow", return_value=workflow) + + result = rag_pipeline_service.service.get_first_step_parameters(_make_pipeline(), "node-1") + + assert result == [] + + @staticmethod + def test_get_second_step_parameters_filters_first_step_variables( + mocker: MockerFixture, rag_pipeline_service: RagPipelineServiceTestContext + ) -> None: + workflow = _make_workflow( + graph={ + "nodes": [ + { + "id": "node-1", + "data": { + "datasource_parameters": { + "workspace": {"value": "{{#start.workspace#}}"}, + "bucket": {"value": ["input", "bucket"]}, + } + }, + } + ] + }, + rag_pipeline_variables=[ + _rag_variable("workspace"), + _rag_variable("bucket"), + _rag_variable("keep"), + _rag_variable("other-node", belong_to_node_id="node-x"), + ], + ) + mocker.patch.object(rag_pipeline_service.service, "get_published_workflow", return_value=workflow) + + result = rag_pipeline_service.service.get_second_step_parameters(_make_pipeline(), "node-1") + + assert [item["variable"] for item in result] == ["keep"] + assert result[0]["belong_to_node_id"] == "shared" + + @staticmethod + def test_retry_error_document_raises_when_execution_log_not_found( + rag_pipeline_service: RagPipelineServiceTestContext, + pipeline_generator, + ) -> None: + with pytest.raises(ValueError, match="Document pipeline execution log not found"): + rag_pipeline_service.service.retry_error_document( + _make_dataset(), _make_document(), _make_account(), generator=pipeline_generator + ) + + @staticmethod + def test_get_datasource_plugins_raises_when_workflow_not_found( + rag_pipeline_service: RagPipelineServiceTestContext, + ) -> None: + dataset = _make_dataset() + pipeline = _make_pipeline(workflow_id="wf-1") + _persist(rag_pipeline_service.session, dataset, pipeline) + + with pytest.raises(ValueError, match="Pipeline or workflow not found"): + rag_pipeline_service.service.get_datasource_plugins( + "t1", + "d1", + True, + credential_query=mock.create_autospec(CredentialQuery, instance=True), + datasource_providers=rag_pipeline_module.DatasourceProviderService( + credentials=create_autospec(DatasourceProviderCredentialStore, instance=True) + ), + ) + + @staticmethod + def test_handle_node_run_result_raises_when_no_terminal_event( + mocker: MockerFixture, rag_pipeline_service: RagPipelineServiceTestContext + ) -> None: + node_instance = _make_start_node() + + def _event_generator(): + yield object() + + with pytest.raises(ValueError, match="Node run failed with no run result"): + rag_pipeline_service.service._handle_node_run_result( + getter=lambda: (node_instance, _event_generator()), + start_at=time.perf_counter(), + tenant_id="t1", + node_id="node-1", + ) + + @staticmethod + def test_handle_node_run_result_marks_document_error_for_published_invoke( + rag_pipeline_service: RagPipelineServiceTestContext, + ) -> None: + node_instance, event = _make_failed_published_node_run() + dataset = _make_dataset(dataset_id="dataset-1", pipeline_id="pipeline-1", tenant_id="t1") + document = _make_document(document_id="doc-1", dataset_id="dataset-1", tenant_id="t1") + _persist( + rag_pipeline_service.session, + dataset, + document, + ) + + result = rag_pipeline_service.service._handle_node_run_result( + getter=lambda: (node_instance, iter([event])), + start_at=time.perf_counter(), + tenant_id="t1", + node_id="node-1", + ) + + assert result.status == WorkflowNodeExecutionStatus.FAILED + rag_pipeline_service.session.expire_all() + updated_document = rag_pipeline_service.session.get(Document, "doc-1") + assert updated_document is not None + assert updated_document.indexing_status == "error" + assert updated_document.error == "boom" + + @staticmethod + def test_handle_node_run_result_does_not_write_when_pipeline_dataset_is_missing( + rag_pipeline_service: RagPipelineServiceTestContext, + ) -> None: + node_instance, event = _make_failed_published_node_run() + + result = rag_pipeline_service.service._handle_node_run_result( + getter=lambda: (node_instance, iter([event])), + start_at=time.perf_counter(), + tenant_id="t1", + node_id="node-1", + ) + + assert result.status == WorkflowNodeExecutionStatus.FAILED + assert rag_pipeline_service.session.query(Document).count() == 0 + + @staticmethod + def test_run_datasource_node_preview_raises_for_unsupported_provider( + mocker: MockerFixture, rag_pipeline_service: RagPipelineServiceTestContext + ) -> None: + pipeline = _make_pipeline() + workflow = _make_workflow( + graph={ + "nodes": [ + { + "id": "node-1", + "data": { + "plugin_id": "plugin-1", + "provider_name": "provider-1", + "datasource_name": "doc", + "datasource_parameters": {}, + }, + } + ] + } + ) + mocker.patch.object(rag_pipeline_service.service, "get_published_workflow", return_value=workflow) + runtime = mocker.Mock() + runtime.datasource_provider_type.return_value = "unsupported" + mocker.patch( + "core.datasource.datasource_manager.DatasourceManager.get_datasource_runtime", return_value=runtime + ) + mocker.patch( + "services.rag_pipeline.rag_pipeline.DatasourceProviderService.get_datasource_credentials", + return_value=None, + ) + + with pytest.raises(RuntimeError, match="Unsupported datasource provider"): + rag_pipeline_service.service.run_datasource_node_preview( + pipeline=pipeline, + node_id="node-1", + user_inputs={}, + account=_make_account(), + datasource_type="website_crawl", + is_published=True, + datasource_providers=rag_pipeline_module.DatasourceProviderService( + credentials=create_autospec(DatasourceProviderCredentialStore, instance=True) + ), + ) + + @staticmethod + def test_publish_customized_pipeline_template_raises_for_missing_workflow_id( + rag_pipeline_service: RagPipelineServiceTestContext, + ) -> None: + pipeline = _make_pipeline(workflow_id=None) + _persist(rag_pipeline_service.session, pipeline) + + with pytest.raises(RagPipelineResourceNotFoundError, match="Pipeline workflow not found"): + rag_pipeline_service.service.publish_customized_pipeline_template( + pipeline, _make_dataset(), _make_template_args(), _make_account(), session=rag_pipeline_service.session + ) + + @staticmethod + def test_get_pipeline_raises_when_dataset_missing( + rag_pipeline_service: RagPipelineServiceTestContext, + ) -> None: + with pytest.raises(ValueError, match="Dataset not found"): + rag_pipeline_service.service.get_pipeline("t1", "d1") + + @staticmethod + def test_get_pipeline_raises_when_pipeline_missing( + rag_pipeline_service: RagPipelineServiceTestContext, + ) -> None: + _persist(rag_pipeline_service.session, _make_dataset()) + + with pytest.raises(ValueError, match="Pipeline not found"): + rag_pipeline_service.service.get_pipeline("t1", "d1") + + @staticmethod + def test_init_uses_default_sessionmaker_when_none(mocker: MockerFixture, sqlite_session: Session) -> None: + mocker.patch("services.rag_pipeline.rag_pipeline.db", SimpleNamespace(engine=sqlite_session.get_bind())) + create_exec_repo = mocker.patch( + "services.rag_pipeline.rag_pipeline.DifyAPIRepositoryFactory.create_api_workflow_node_execution_repository" + ) + create_run_repo = mocker.patch( + "services.rag_pipeline.rag_pipeline.DifyAPIRepositoryFactory.create_api_workflow_run_repository" + ) + + RagPipelineService(session=sqlite_session, session_maker=None) + + exec_session_maker = create_exec_repo.call_args.args[0] + run_session_maker = create_run_repo.call_args.args[0] + assert exec_session_maker is run_session_maker + assert exec_session_maker.kw["bind"] is sqlite_session.get_bind() + assert exec_session_maker.kw["expire_on_commit"] is False + + @staticmethod + def test_get_pipeline_templates_builtin_en_us_no_fallback( + mocker: MockerFixture, sqlite_session: Session, config_overrides: Callable[..., None] + ) -> None: + config_overrides(HOSTED_FETCH_PIPELINE_TEMPLATES_MODE="remote") + session = sqlite_session + retrieval = mocker.Mock() + retrieval.get_pipeline_templates.return_value = {"pipeline_templates": []} + factory = mocker.patch("services.rag_pipeline.rag_pipeline.PipelineTemplateRetrievalFactory") + factory.get_pipeline_template_factory.return_value.return_value = retrieval + builtin = factory.get_built_in_pipeline_template_retrieval.return_value + result = RagPipelineService.get_pipeline_templates(type="built-in", language="en-US", session=session) + + assert result == {"pipeline_templates": []} + retrieval.get_pipeline_templates.assert_called_once_with("en-US", None, session=session) + builtin.fetch_pipeline_templates_from_builtin.assert_not_called() + + @staticmethod + def test_update_customized_pipeline_template_commits_when_name_empty(sqlite_session: Session) -> None: + template = _make_customized_template() + template.id = "tpl-1" + _persist(sqlite_session, template) + + info = PipelineTemplateInfoEntity(name="", description="updated", icon_info=IconInfo(icon="i")) + result = RagPipelineService.update_customized_pipeline_template( + "tpl-1", info, _make_account(), "t1", session=sqlite_session + ) + + assert result.description == "updated" + sqlite_session.expire_all() + updated_template = sqlite_session.get(PipelineCustomizedTemplate, "tpl-1") + assert updated_template is not None + assert updated_template.description == "updated" + + @staticmethod + def test_get_all_published_workflow_without_filters_has_no_more( + rag_pipeline_service: RagPipelineServiceTestContext, + ) -> None: + pipeline = _make_pipeline(workflow_id="wf-live") + workflow = _make_workflow(workflow_id="wf1") + workflow.version = "2024-01-01T00:00:00" + _persist(rag_pipeline_service.session, workflow) + + workflows, has_more = rag_pipeline_service.service.get_all_published_workflow( + session=rag_pipeline_service.session, + pipeline=pipeline, + page=1, + limit=2, + user_id=None, + named_only=False, + ) + + assert workflows == [workflow] + assert has_more is False + + @staticmethod + def test_publish_workflow_skips_dataset_update_for_non_knowledge_nodes( + mocker: MockerFixture, rag_pipeline_service: RagPipelineServiceTestContext + ) -> None: + draft = _make_workflow(graph={"nodes": [{"data": {"type": "start"}}]}) + _persist(rag_pipeline_service.session, draft) + dataset_service = mocker.patch("services.knowledge.dataset_service.DatasetService") + + result = rag_pipeline_service.service.publish_workflow( + session=rag_pipeline_service.session, + pipeline=_make_pipeline(), + account=_make_account(), + ) + + assert result.graph_dict == {"nodes": [{"data": {"type": "start"}}]} + dataset_service.update_rag_pipeline_dataset_settings.assert_not_called() + + @staticmethod + def test_get_default_block_config_returns_none_when_default_empty( + mocker: MockerFixture, rag_pipeline_service: RagPipelineServiceTestContext + ) -> None: + node_cls = mocker.Mock() + node_cls.get_default_config.return_value = None + mocker.patch( + "services.rag_pipeline.rag_pipeline.get_node_type_classes_mapping", + return_value={BuiltinNodeTypes.START: {"1": node_cls}}, + ) + mocker.patch("services.rag_pipeline.rag_pipeline.LATEST_VERSION", "1") + + assert rag_pipeline_service.service.get_default_block_config("start") is None + + @staticmethod + def test_run_datasource_workflow_node_handles_variable_parameter_types( + mocker: MockerFixture, rag_pipeline_service: RagPipelineServiceTestContext + ) -> None: + from core.datasource.entities.datasource_entities import DatasourceProviderType + + workflow = _make_workflow( + graph={ + "nodes": [ + { + "id": "node-1", + "data": { + "plugin_id": "p", + "provider_name": "provider", + "datasource_name": "crawl", + "datasource_parameters": { + "a": {"value": None}, + "b": {"value": "literal"}, + "c": {"value": ["input", "k"]}, + }, + }, + } + ] + } + ) + mocker.patch.object(rag_pipeline_service.service, "get_published_workflow", return_value=workflow) + runtime = mocker.Mock() + + def crawl_gen(**kwargs): + yield SimpleNamespace(result=SimpleNamespace(status="completed", total=1, completed=1, web_info_list=[])) + + runtime.get_website_crawl.side_effect = crawl_gen + runtime.datasource_provider_type.return_value = DatasourceProviderType.WEBSITE_CRAWL + mocker.patch( + "core.datasource.datasource_manager.DatasourceManager.get_datasource_runtime", return_value=runtime + ) + mocker.patch( + "services.rag_pipeline.rag_pipeline.DatasourceProviderService.get_datasource_credentials", + return_value=None, + ) + + events = list( + rag_pipeline_service.service.run_datasource_workflow_node( + pipeline=_make_pipeline(), + node_id="node-1", + user_inputs={"k": "mapped"}, + account=_make_account(), + datasource_type="website_crawl", + is_published=True, + datasource_providers=rag_pipeline_module.DatasourceProviderService( + credentials=create_autospec(DatasourceProviderCredentialStore, instance=True) + ), + ) + ) + + assert events + assert events[0]["data"] == [] + + @staticmethod + def test_run_datasource_workflow_node_online_drive_branch( + mocker: MockerFixture, rag_pipeline_service: RagPipelineServiceTestContext + ) -> None: + from core.datasource.entities.datasource_entities import DatasourceProviderType + + workflow = _make_workflow( + graph={ + "nodes": [ + { + "id": "node-1", + "data": { + "plugin_id": "p", + "provider_name": "provider", + "datasource_name": "drive", + "datasource_parameters": {}, + }, + } + ] + } + ) + mocker.patch.object(rag_pipeline_service.service, "get_published_workflow", return_value=workflow) + runtime = mocker.Mock() + + def drive_gen(**kwargs): + yield SimpleNamespace(result={"items": [1]}) + + runtime.online_drive_browse_files.side_effect = drive_gen + runtime.datasource_provider_type.return_value = DatasourceProviderType.ONLINE_DRIVE + mocker.patch( + "core.datasource.datasource_manager.DatasourceManager.get_datasource_runtime", return_value=runtime + ) + mocker.patch( + "services.rag_pipeline.rag_pipeline.DatasourceProviderService.get_datasource_credentials", + return_value=None, + ) + + events = list( + rag_pipeline_service.service.run_datasource_workflow_node( + pipeline=_make_pipeline(), + node_id="node-1", + user_inputs={}, + account=_make_account(), + datasource_type="online_drive", + is_published=True, + datasource_providers=rag_pipeline_module.DatasourceProviderService( + credentials=create_autospec(DatasourceProviderCredentialStore, instance=True) + ), + ) + ) + + assert len(events) == 2 + assert events[1]["data"] == {"items": [1]} + + @staticmethod + def test_run_datasource_node_preview_not_published_uses_draft( + mocker: MockerFixture, rag_pipeline_service: RagPipelineServiceTestContext + ) -> None: + from core.datasource.entities.datasource_entities import DatasourceMessage + + workflow = _make_workflow( + graph={ + "nodes": [ + { + "id": "n1", + "data": { + "plugin_id": "p", + "provider_name": "provider", + "datasource_name": "doc", + "datasource_parameters": {"workspace_id": {"value": "w"}}, + }, + } + ] + } + ) + get_draft = mocker.patch.object(rag_pipeline_service.service, "get_draft_workflow", return_value=workflow) + runtime = mocker.Mock() + + def doc_gen(**kwargs): + yield DatasourceMessage( + type=DatasourceMessage.MessageType.VARIABLE, + message=DatasourceMessage.VariableMessage(variable_name="x", variable_value="v", stream=False), + ) + + runtime.get_online_document_page_content.side_effect = doc_gen + mocker.patch( + "core.datasource.datasource_manager.DatasourceManager.get_datasource_runtime", return_value=runtime + ) + mocker.patch( + "services.rag_pipeline.rag_pipeline.DatasourceProviderService.get_datasource_credentials", + return_value=None, + ) + + result = rag_pipeline_service.service.run_datasource_node_preview( + pipeline=_make_pipeline(), + node_id="n1", + user_inputs={}, + account=_make_account(), + datasource_type="online_document", + is_published=False, + datasource_providers=rag_pipeline_module.DatasourceProviderService( + credentials=create_autospec(DatasourceProviderCredentialStore, instance=True) + ), + ) + + assert result == {"x": "v"} + get_draft.assert_called_once() + + @staticmethod + def test_run_free_workflow_node_delegates_to_handle_result( + mocker: MockerFixture, rag_pipeline_service: RagPipelineServiceTestContext + ) -> None: + expected = SimpleNamespace(id="exec-1") + handle = mocker.patch.object(rag_pipeline_service.service, "_handle_node_run_result", return_value=expected) + + result = rag_pipeline_service.service.run_free_workflow_node( + node_data={"type": "start"}, + tenant_id="t1", + user_id="u1", + node_id="n1", + user_inputs={}, + ) + + assert result is expected + handle.assert_called_once() + + @staticmethod + @pytest.mark.parametrize(("workflow_tenant_id", "workflow_app_id"), [("t2", "p1"), ("t1", "p2")]) + def test_publish_customized_pipeline_template_rejects_unowned_workflow_before_export( + mocker: MockerFixture, + rag_pipeline_service: RagPipelineServiceTestContext, + workflow_tenant_id: str, + workflow_app_id: str, + ) -> None: + pipeline = _make_pipeline(workflow_id="wf-1") + workflow = _make_workflow(workflow_id="wf-1", tenant_id=workflow_tenant_id, app_id=workflow_app_id) + _persist(rag_pipeline_service.session, pipeline, workflow) + dsl_service = mocker.patch.object(rag_pipeline_module, "RagPipelineDslService") + + with pytest.raises(RagPipelineResourceNotFoundError, match="Workflow not found"): + rag_pipeline_service.service.publish_customized_pipeline_template( + pipeline, + _make_dataset(), + _make_template_args(), + _make_account(), + session=rag_pipeline_service.session, + ) + + dsl_service.assert_not_called() + assert rag_pipeline_service.session.query(PipelineCustomizedTemplate).count() == 0 + + @staticmethod + def test_pipeline_retrieve_dataset_rejects_unowned_dataset( + rag_pipeline_service: RagPipelineServiceTestContext, + ) -> None: + pipeline = _make_pipeline(workflow_id="wf-1") + other_tenant_dataset = _make_dataset(tenant_id="t2") + _persist(rag_pipeline_service.session, pipeline, other_tenant_dataset) + + assert get_pipeline_dataset(pipeline, session=rag_pipeline_service.session) is None + + @staticmethod + @pytest.mark.parametrize( + ("draft_tenant_id", "draft_app_id"), + [(None, None), ("t2", "p1"), ("t1", "p2")], + ) + def test_publish_customized_pipeline_template_rejects_missing_or_unowned_draft_before_side_effects( + mocker: MockerFixture, + rag_pipeline_service: RagPipelineServiceTestContext, + draft_tenant_id: str | None, + draft_app_id: str | None, + ) -> None: + session = rag_pipeline_service.session + pipeline = _make_pipeline(workflow_id="wf-published") + published_workflow = _make_workflow(workflow_id="wf-published", version="published") + dataset = _make_dataset() + resources = [pipeline, published_workflow, dataset] + if draft_tenant_id and draft_app_id: + resources.append(_make_workflow(workflow_id="wf-draft", tenant_id=draft_tenant_id, app_id=draft_app_id)) + _persist(session, *resources) + dsl_service = mocker.patch.object(rag_pipeline_module, "RagPipelineDslService") + + with pytest.raises(RagPipelineResourceNotFoundError, match="Draft workflow not found"): + rag_pipeline_service.service.publish_customized_pipeline_template( + pipeline, dataset, _make_template_args(), _make_account(), session=session + ) + + dsl_service.assert_not_called() + assert session.query(PipelineCustomizedTemplate).count() == 0 + + @staticmethod + def test_get_recommended_plugins_skips_manifest_when_missing( + mocker: MockerFixture, rag_pipeline_service: RagPipelineServiceTestContext + ) -> None: + plugin = _make_recommended_plugin("plugin-a") + _persist(rag_pipeline_service.session, plugin) + mocker.patch("services.rag_pipeline.rag_pipeline.BuiltinToolManageService.list_builtin_tools", return_value=[]) + mocker.patch("services.rag_pipeline.rag_pipeline.marketplace.batch_fetch_plugin_by_ids", return_value=[]) + + result = rag_pipeline_service.service.get_recommended_plugins("all", _make_account(), "t1") + + assert result["installed_recommended_plugins"] == [] + assert result["uninstalled_recommended_plugins"] == [] + + @staticmethod + def test_retry_error_document_raises_when_pipeline_missing( + rag_pipeline_service: RagPipelineServiceTestContext, + pipeline_generator, + ) -> None: + exec_log = DocumentPipelineExecutionLog( + pipeline_id="p1", + document_id="doc-1", + datasource_type="local_file", + datasource_info="{}", + datasource_node_id="node-1", + input_data={}, + created_by="u1", + ) + _persist(rag_pipeline_service.session, exec_log) + + with pytest.raises(ValueError, match="Pipeline not found"): + rag_pipeline_service.service.retry_error_document( + _make_dataset(), _make_document(), _make_account(), generator=pipeline_generator + ) + + @staticmethod + def test_retry_error_document_raises_when_workflow_missing( + rag_pipeline_service: RagPipelineServiceTestContext, + pipeline_generator, + ) -> None: + exec_log = DocumentPipelineExecutionLog( + pipeline_id="p1", + document_id="doc-1", + datasource_type="local_file", + datasource_info="{}", + datasource_node_id="node-1", + input_data={}, + created_by="u1", + ) + pipeline = _make_pipeline(workflow_id="wf-1") + _persist(rag_pipeline_service.session, exec_log, pipeline) + + with pytest.raises(ValueError, match="Workflow not found"): + rag_pipeline_service.service.retry_error_document( + _make_dataset(), _make_document(), _make_account(), generator=pipeline_generator + ) + + @staticmethod + def test_get_datasource_plugins_returns_empty_for_non_datasource_nodes( + rag_pipeline_service: RagPipelineServiceTestContext, + ) -> None: + dataset = _make_dataset() + pipeline = _make_pipeline(workflow_id="wf-1") + workflow = _make_workflow(graph={"nodes": [{"id": "n1", "data": {"type": "start"}}]}) + _persist(rag_pipeline_service.session, dataset, pipeline, workflow) + + assert ( + rag_pipeline_service.service.get_datasource_plugins( + "t1", + "d1", + True, + credential_query=mock.create_autospec(CredentialQuery, instance=True), + datasource_providers=rag_pipeline_module.DatasourceProviderService( + credentials=create_autospec(DatasourceProviderCredentialStore, instance=True) + ), + ) + == [] + ) + + @staticmethod + def test_publish_workflow_raises_when_knowledge_index_dataset_missing( + mocker: MockerFixture, rag_pipeline_service: RagPipelineServiceTestContext + ) -> None: + draft = _make_workflow(graph={"nodes": [{"data": {"type": "knowledge-index"}}]}) + _persist(rag_pipeline_service.session, draft) + mocker.patch( + "services.rag_pipeline.rag_pipeline.KnowledgeConfiguration.model_validate", return_value=mocker.Mock() + ) + pipeline = _make_pipeline() + + with pytest.raises(ValueError, match="Dataset not found"): + rag_pipeline_service.service.publish_workflow( + session=rag_pipeline_service.session, pipeline=pipeline, account=_make_account() + ) + + @staticmethod + def test_run_datasource_node_preview_raises_when_workflow_missing( + mocker: MockerFixture, rag_pipeline_service: RagPipelineServiceTestContext + ) -> None: + mocker.patch.object(rag_pipeline_service.service, "get_published_workflow", return_value=None) + + with pytest.raises(RuntimeError, match="Workflow not initialized"): + rag_pipeline_service.service.run_datasource_node_preview( + pipeline=_make_pipeline(), + node_id="n1", + user_inputs={}, + account=_make_account(), + datasource_type="online_document", + is_published=True, + datasource_providers=rag_pipeline_module.DatasourceProviderService( + credentials=create_autospec(DatasourceProviderCredentialStore, instance=True) + ), + ) + + @staticmethod + def test_run_datasource_node_preview_raises_when_node_missing( + mocker: MockerFixture, rag_pipeline_service: RagPipelineServiceTestContext + ) -> None: + mocker.patch.object( + rag_pipeline_service.service, "get_published_workflow", return_value=_make_workflow(graph={"nodes": []}) + ) + + with pytest.raises(RuntimeError, match="Datasource node data not found"): + rag_pipeline_service.service.run_datasource_node_preview( + pipeline=_make_pipeline(), + node_id="missing", + user_inputs={}, + account=_make_account(), + datasource_type="online_document", + is_published=True, + datasource_providers=rag_pipeline_module.DatasourceProviderService( + credentials=create_autospec(DatasourceProviderCredentialStore, instance=True) + ), + ) + + @staticmethod + def test_run_datasource_node_preview_keeps_existing_user_input( + mocker: MockerFixture, rag_pipeline_service: RagPipelineServiceTestContext + ) -> None: + from core.datasource.entities.datasource_entities import DatasourceMessage + + workflow = _make_workflow( + graph={ + "nodes": [ + { + "id": "n1", + "data": { + "plugin_id": "p", + "provider_name": "provider", + "datasource_name": "doc", + "datasource_parameters": {"workspace_id": {"value": "default"}}, + }, + } + ] + } + ) + mocker.patch.object(rag_pipeline_service.service, "get_published_workflow", return_value=workflow) + runtime = mocker.Mock() + + def gen(**kwargs): + request = kwargs["datasource_parameters"] + assert request.workspace_id == "existing" + yield DatasourceMessage( + type=DatasourceMessage.MessageType.VARIABLE, + message=DatasourceMessage.VariableMessage(variable_name="ok", variable_value="1", stream=False), + ) + + runtime.get_online_document_page_content.side_effect = gen + mocker.patch( + "core.datasource.datasource_manager.DatasourceManager.get_datasource_runtime", return_value=runtime + ) + mocker.patch( + "services.rag_pipeline.rag_pipeline.DatasourceProviderService.get_datasource_credentials", + return_value=None, + ) + + result = rag_pipeline_service.service.run_datasource_node_preview( + pipeline=_make_pipeline(), + node_id="n1", + user_inputs={"workspace_id": "existing"}, + account=_make_account(), + datasource_type="online_document", + is_published=True, + datasource_providers=rag_pipeline_module.DatasourceProviderService( + credentials=create_autospec(DatasourceProviderCredentialStore, instance=True) + ), + ) + assert result == {"ok": "1"} + + @staticmethod + def test_run_datasource_node_preview_ignores_non_variable_messages( + mocker: MockerFixture, rag_pipeline_service: RagPipelineServiceTestContext + ) -> None: + workflow = _make_workflow( + graph={ + "nodes": [ + { + "id": "n1", + "data": { + "plugin_id": "p", + "provider_name": "provider", + "datasource_name": "doc", + "datasource_parameters": {}, + }, + } + ] + } + ) + mocker.patch.object(rag_pipeline_service.service, "get_published_workflow", return_value=workflow) + runtime = mocker.Mock() + + def gen(**kwargs): + yield SimpleNamespace(type="log", message=None) + + runtime.get_online_document_page_content.side_effect = gen + mocker.patch( + "core.datasource.datasource_manager.DatasourceManager.get_datasource_runtime", return_value=runtime + ) + mocker.patch( + "services.rag_pipeline.rag_pipeline.DatasourceProviderService.get_datasource_credentials", + return_value=None, + ) + + result = rag_pipeline_service.service.run_datasource_node_preview( + pipeline=_make_pipeline(), + node_id="n1", + user_inputs={}, + account=_make_account(), + datasource_type="online_document", + is_published=True, + datasource_providers=rag_pipeline_module.DatasourceProviderService( + credentials=create_autospec(DatasourceProviderCredentialStore, instance=True) + ), + ) + assert result == {} + + @staticmethod + def test_set_datasource_variables_raises_when_workflow_missing( + mocker: MockerFixture, rag_pipeline_service: RagPipelineServiceTestContext + ) -> None: + mocker.patch.object(rag_pipeline_service.service, "get_draft_workflow", return_value=None) + + with pytest.raises(ValueError, match="Workflow not initialized"): + rag_pipeline_service.service.set_datasource_variables( + _make_pipeline(), + {"start_node_id": "n1"}, + _make_account(), + ) + + @staticmethod + def test_get_datasource_plugins_handles_empty_datasource_data_and_non_published( + mocker: MockerFixture, rag_pipeline_service: RagPipelineServiceTestContext + ) -> None: + dataset = _make_dataset() + pipeline = _make_pipeline() + workflow = _make_workflow( + graph={"nodes": [{"id": "n1", "data": {"type": "datasource", "datasource_parameters": {}}}]}, + ) + workflow.rag_pipeline_variables = [ + {"variable": "v1", "belong_to_node_id": "shared", "type": "text-input", "label": "V1"} + ] + _persist(rag_pipeline_service.session, dataset, pipeline, workflow) + mocker.patch( + "services.rag_pipeline.rag_pipeline.DatasourceProviderService.list_datasource_credentials", + return_value=[], + ) + + result = rag_pipeline_service.service.get_datasource_plugins( + "t1", + "d1", + False, + credential_query=mock.create_autospec(CredentialQuery, instance=True), + datasource_providers=rag_pipeline_module.DatasourceProviderService( + credentials=create_autospec(DatasourceProviderCredentialStore, instance=True) + ), + ) + + assert len(result) == 1 + + @staticmethod + def test_get_datasource_plugins_extracts_user_inputs_and_credentials( + mocker: MockerFixture, rag_pipeline_service: RagPipelineServiceTestContext + ) -> None: + dataset = _make_dataset() + pipeline = _make_pipeline(workflow_id="wf-1") + workflow = _make_workflow( + graph={ + "nodes": [ + { + "id": "n1", + "data": { + "type": "datasource", + "plugin_id": "plugin-1", + "provider_name": "provider", + "provider_type": "online_document", + "title": "Datasource", + "datasource_parameters": { + "a": {"value": "{{#start.v1#}}"}, + "b": {"value": ["x", "v2"]}, + }, + }, + } + ] + }, + ) + workflow.rag_pipeline_variables = [ + {"variable": "v1", "belong_to_node_id": "shared", "type": "text-input", "label": "V1"}, + {"variable": "v2", "belong_to_node_id": "shared", "type": "text-input", "label": "V2"}, + {"variable": "v3", "belong_to_node_id": "shared", "type": "text-input", "label": "V3"}, + ] + _persist(rag_pipeline_service.session, dataset, pipeline, workflow) + mocker.patch( + "services.rag_pipeline.rag_pipeline.DatasourceProviderService.list_datasource_credentials", + return_value=[{"id": "c1", "name": "Cred", "type": "api", "is_default": True}], + ) + + result = rag_pipeline_service.service.get_datasource_plugins( + "t1", + "d1", + True, + credential_query=mock.create_autospec(CredentialQuery, instance=True), + datasource_providers=rag_pipeline_module.DatasourceProviderService( + credentials=create_autospec(DatasourceProviderCredentialStore, instance=True) + ), + ) + + assert len(result) == 1 + assert len(result[0]["user_input_variables"]) == 2 + assert result[0]["credentials"][0]["id"] == "c1" + + @staticmethod + def test_get_pipeline_returns_pipeline_when_found( + rag_pipeline_service: RagPipelineServiceTestContext, + ) -> None: + dataset = _make_dataset() + pipeline = _make_pipeline() + _persist(rag_pipeline_service.session, dataset, pipeline) + + result = rag_pipeline_service.service.get_pipeline("t1", "d1") + + assert result is pipeline + + @staticmethod + def test_get_pipeline_by_id_uses_provided_session(sqlite_session: Session) -> None: + pipeline = _make_pipeline() + sqlite_session.add(pipeline) + sqlite_session.flush() + + result = RagPipelineService.get_pipeline_by_id("p1", "t1", session=sqlite_session) + + assert result is pipeline + assert RagPipelineService.get_pipeline_by_id("p1", "other-tenant", session=sqlite_session) is None diff --git a/api/tests/unit_tests/services/rag_pipeline/test_rag_pipeline_dsl_service.py b/api/tests/unit_tests/services/rag_pipeline/test_rag_pipeline_dsl_service.py index 060568c5f02..0209ce5ba1a 100644 --- a/api/tests/unit_tests/services/rag_pipeline/test_rag_pipeline_dsl_service.py +++ b/api/tests/unit_tests/services/rag_pipeline/test_rag_pipeline_dsl_service.py @@ -39,7 +39,7 @@ from services.rag_pipeline.rag_pipeline_dsl_service import ( RagPipelineDslService, RagPipelinePendingData, ) -from tests.unit_tests.model_factories import make_account, make_dataset, make_tenant, make_workflow +from tests.helpers.unit.model_factories import make_account, make_dataset, make_tenant, make_workflow @pytest.fixture diff --git a/api/tests/unit_tests/services/rag_pipeline/test_rag_pipeline_service.py b/api/tests/unit_tests/services/rag_pipeline/test_rag_pipeline_service.py deleted file mode 100644 index f03c9aabc41..00000000000 --- a/api/tests/unit_tests/services/rag_pipeline/test_rag_pipeline_service.py +++ /dev/null @@ -1,2671 +0,0 @@ -import json -import time -from collections.abc import Callable -from dataclasses import dataclass -from datetime import datetime -from types import SimpleNamespace -from unittest import mock -from unittest.mock import create_autospec - -import pytest -from pytest_mock import MockerFixture -from sqlalchemy.dialects import postgresql -from sqlalchemy.orm import Session, sessionmaker - -from core.app.apps.pipeline.pipeline_generator import PipelineGenerator -from core.app.entities.app_invoke_entities import InvokeFrom -from core.rag.index_processor.constant.index_type import IndexStructureType -from graphon.engine_events import NodeRunFailedEvent -from graphon.enums import ( - BuiltinNodeTypes, - ErrorStrategy, - WorkflowNodeExecutionMetadataKey, - WorkflowNodeExecutionStatus, -) -from graphon.node_events.base import NodeRunResult -from graphon.nodes.start.entities import StartNodeData -from graphon.nodes.start.start_node import StartNode -from graphon.runtime import InitParams, RuntimeState, VariablePool -from models import Account, Tenant -from models.dataset import ( - Dataset, - Document, - DocumentPipelineExecutionLog, - Pipeline, - PipelineCustomizedTemplate, - PipelineRecommendedPlugin, -) -from models.enums import DataSourceType, DocumentCreatedFrom, IndexingStatus -from models.workflow import Workflow, WorkflowRun -from repositories.knowledge.dataset_read_repository import get_pipeline_dataset -from services.credentials.query import CredentialQuery -from services.data_source.credential_gateway import DatasourceProviderCredentialStore -from services.entities.knowledge_entities.rag_pipeline_entities import IconInfo, PipelineTemplateInfoEntity -from services.errors.rag_pipeline import RagPipelineResourceNotFoundError -from services.rag_pipeline import rag_pipeline as rag_pipeline_module -from services.rag_pipeline.rag_pipeline import RagPipelineService -from services.workflow_ref_service import WorkflowRef - - -@pytest.fixture -def pipeline_generator(mocker: MockerFixture): - return mocker.create_autospec(PipelineGenerator, instance=True, spec_set=True) - - -@dataclass -class RagPipelineServiceTestContext: - service: RagPipelineService - session: Session - session_maker: sessionmaker[Session] - - -@pytest.fixture -def rag_pipeline_service( - mocker: MockerFixture, - sqlite_session: Session, - sqlite_session_factory: sessionmaker[Session], -) -> RagPipelineServiceTestContext: - mocker.patch( - "services.rag_pipeline.rag_pipeline.DifyAPIRepositoryFactory.create_api_workflow_node_execution_repository", - return_value=MockRepo(), - ) - mocker.patch( - "services.rag_pipeline.rag_pipeline.DifyAPIRepositoryFactory.create_api_workflow_run_repository", - return_value=MockRepo(), - ) - mocker.patch("services.rag_pipeline.rag_pipeline.db", SimpleNamespace(engine=sqlite_session.get_bind())) - service = RagPipelineService(session=sqlite_session, session_maker=sqlite_session_factory) - return RagPipelineServiceTestContext( - service=service, - session=sqlite_session, - session_maker=sqlite_session_factory, - ) - - -class MockRepo: - pass - - -def test_get_published_workflow_by_id_locks_restore_source( - mocker: MockerFixture, rag_pipeline_service: RagPipelineServiceTestContext -) -> None: - session = rag_pipeline_service.session - workflow = _make_workflow() - workflow.version = "v1" - session.add(workflow) - session.flush() - lookup = mocker.spy(session, "scalar") - service = rag_pipeline_service.service - - result = service.get_published_workflow_by_id(_make_pipeline(), workflow.id) - - stmt = lookup.call_args.args[0] - sql = str(stmt.compile(dialect=postgresql.dialect())) - assert result is workflow - assert "FOR UPDATE" in sql - - -def _make_account(account_id: str = "u1", tenant_id: str = "t1") -> Account: - account = Account(name="Test User", email=f"{account_id}@example.com") - account.id = account_id - tenant = Tenant(name="Test Tenant") - tenant.id = tenant_id - account._current_tenant = tenant - return account - - -def _make_pipeline( - *, - pipeline_id: str = "p1", - tenant_id: str = "t1", - workflow_id: str | None = None, - is_published: bool = False, -) -> Pipeline: - pipeline = Pipeline( - tenant_id=tenant_id, - name="Test Pipeline", - description="test", - workflow_id=workflow_id, - is_published=is_published, - ) - pipeline.id = pipeline_id - return pipeline - - -def _make_template_args(name: str = "New Template") -> dict[str, object]: - return {"name": name, "description": "Desc", "icon_info": {"icon": "star"}} - - -def _make_workflow( - *, - workflow_id: str = "wf-1", - tenant_id: str = "t1", - app_id: str = "p1", - graph: dict[str, object] | None = None, - features: dict[str, object] | None = None, - rag_pipeline_variables: list[dict[str, object]] | None = None, - created_by: str = "u1", - version: str = Workflow.VERSION_DRAFT, -) -> Workflow: - workflow = Workflow( - id=workflow_id, - tenant_id=tenant_id, - app_id=app_id, - type="workflow", - version=version, - marked_name="", - marked_comment="", - graph=json.dumps(graph or {"nodes": []}), - features=json.dumps(features or {}), - created_by=created_by, - created_at=datetime(2024, 1, 1), - updated_by=None, - updated_at=datetime(2024, 1, 1), - environment_variables=[], - conversation_variables=[], - rag_pipeline_variables=rag_pipeline_variables or [], - ) - return workflow - - -def _make_dataset(*, dataset_id: str = "d1", pipeline_id: str = "p1", tenant_id: str = "t1") -> Dataset: - dataset = Dataset( - id=dataset_id, - tenant_id=tenant_id, - name="Test Dataset", - created_by="u1", - pipeline_id=pipeline_id, - ) - return dataset - - -def _make_document(*, document_id: str = "doc-1", dataset_id: str = "d1", tenant_id: str = "t1") -> Document: - return Document( - id=document_id, - tenant_id=tenant_id, - dataset_id=dataset_id, - position=1, - data_source_type=DataSourceType.LOCAL_FILE, - batch="batch", - name="Document", - created_from=DocumentCreatedFrom.API, - created_by="u1", - indexing_status=IndexingStatus.WAITING, - doc_form=IndexStructureType.PARAGRAPH_INDEX, - ) - - -def _rag_variable(variable: str, *, belong_to_node_id: str = "shared") -> dict[str, object]: - return { - "variable": variable, - "belong_to_node_id": belong_to_node_id, - "type": "text-input", - "label": variable.title(), - } - - -def _persist(session: Session, *entities: object) -> None: - session.add_all(entities) - session.commit() - - -def _make_start_node(data: StartNodeData | None = None, variable_pool: VariablePool | None = None) -> StartNode: - return StartNode( - node_id="node-1", - data=data if data is not None else StartNodeData(title="Start"), - init_params=InitParams(workflow_id="wf-1", graph_config={}, run_context={}, call_depth=0), - runtime_state=RuntimeState( - workflow_id="wf-1", - variable_pool=variable_pool if variable_pool is not None else VariablePool(), - start_at=time.perf_counter(), - ), - ) - - -def _make_failed_published_node_run() -> tuple[StartNode, NodeRunFailedEvent]: - variable_pool = VariablePool() - for name, value in { - "invoke_from": InvokeFrom.PUBLISHED_PIPELINE, - "app_id": "pipeline-1", - "dataset_id": "dataset-1", - "document_id": "doc-1", - }.items(): - variable_pool.add(["sys", name], value) - node_instance = _make_start_node(variable_pool=variable_pool) - run_result = NodeRunResult( - status=WorkflowNodeExecutionStatus.FAILED, - error="boom", - error_type="runtime", - inputs={}, - outputs={}, - ) - event = NodeRunFailedEvent( - id="evt-1", - start_at=time.time(), - node_id="node-1", - node_type="start", - node_run_result=run_result, - error="boom", - route_node_id=None, - ) - return node_instance, event - - -def _make_customized_template() -> PipelineCustomizedTemplate: - return PipelineCustomizedTemplate( - tenant_id="t1", - name="old", - description="old", - chunk_structure="paragraph", - icon={}, - position=1, - yaml_content="", - install_count=0, - language="en-US", - created_by="u1", - ) - - -def _make_recommended_plugin(plugin_id: str) -> PipelineRecommendedPlugin: - return PipelineRecommendedPlugin(plugin_id=plugin_id, provider_name=plugin_id, type="tool", position=0, active=True) - - -def test_get_pipeline_templates_fallbacks_to_builtin_for_non_english_empty_result( - mocker: MockerFixture, sqlite_session: Session, config_overrides: Callable[..., None] -) -> None: - config_overrides(HOSTED_FETCH_PIPELINE_TEMPLATES_MODE="remote") - session = sqlite_session - - remote_retrieval = mocker.Mock() - remote_retrieval.get_pipeline_templates.return_value = {"pipeline_templates": []} - - factory_mock = mocker.patch("services.rag_pipeline.rag_pipeline.PipelineTemplateRetrievalFactory") - factory_mock.get_pipeline_template_factory.return_value.return_value = remote_retrieval - - builtin_retrieval = mocker.Mock() - builtin_retrieval.fetch_pipeline_templates_from_builtin.return_value = {"pipeline_templates": [{"id": "builtin-1"}]} - factory_mock.get_built_in_pipeline_template_retrieval.return_value = builtin_retrieval - - result = RagPipelineService.get_pipeline_templates(type="built-in", language="ja-JP", session=session) - - assert result == {"pipeline_templates": [{"id": "builtin-1"}]} - remote_retrieval.get_pipeline_templates.assert_called_once_with("ja-JP", None, session=session) - builtin_retrieval.fetch_pipeline_templates_from_builtin.assert_called_once_with("en-US") - - -def test_get_pipeline_templates_customized_mode_uses_customized_factory( - mocker: MockerFixture, sqlite_session: Session -) -> None: - session = sqlite_session - retrieval = mocker.Mock() - retrieval.get_pipeline_templates.return_value = {"pipeline_templates": [{"id": "custom-1"}]} - - factory_mock = mocker.patch("services.rag_pipeline.rag_pipeline.PipelineTemplateRetrievalFactory") - factory_mock.get_pipeline_template_factory.return_value.return_value = retrieval - - result = RagPipelineService.get_pipeline_templates(type="customized", language="en-US", session=session) - - assert result == {"pipeline_templates": [{"id": "custom-1"}]} - factory_mock.get_pipeline_template_factory.assert_called_with("customized") - retrieval.get_pipeline_templates.assert_called_once_with("en-US", None, session=session) - - -@pytest.mark.parametrize("template_type", ["built-in", "customized"]) -def test_get_pipeline_template_detail_uses_expected_mode( - mocker: MockerFixture, - template_type: str, - sqlite_session: Session, - config_overrides: Callable[..., None], -) -> None: - config_overrides(HOSTED_FETCH_PIPELINE_TEMPLATES_MODE="remote") - session = sqlite_session - retrieval = mocker.Mock() - retrieval.get_pipeline_template_detail.return_value = {"id": "tpl-1"} - - factory_mock = mocker.patch("services.rag_pipeline.rag_pipeline.PipelineTemplateRetrievalFactory") - factory_mock.get_pipeline_template_factory.return_value.return_value = retrieval - - result = RagPipelineService.get_pipeline_template_detail("tpl-1", "tenant-1", type=template_type, session=session) - - assert result == {"id": "tpl-1"} - expected_mode = "remote" if template_type == "built-in" else "customized" - factory_mock.get_pipeline_template_factory.assert_called_with(expected_mode) - retrieval.get_pipeline_template_detail.assert_called_once_with("tpl-1", "tenant-1", session=session) - - -def test_get_published_workflow_returns_none_when_pipeline_has_no_workflow_id( - rag_pipeline_service: RagPipelineServiceTestContext, -) -> None: - pipeline = _make_pipeline(workflow_id=None) - - result = rag_pipeline_service.service.get_published_workflow(pipeline) - - assert result is None - - -def test_get_all_published_workflow_returns_empty_for_unpublished_pipeline( - rag_pipeline_service: RagPipelineServiceTestContext, -) -> None: - pipeline = _make_pipeline(workflow_id=None) - workflows, has_more = rag_pipeline_service.service.get_all_published_workflow( - session=rag_pipeline_service.session, - pipeline=pipeline, - page=1, - limit=20, - user_id=None, - named_only=False, - ) - - assert workflows == [] - assert has_more is False - - -def test_get_all_published_workflow_applies_limit_and_has_more( - rag_pipeline_service: RagPipelineServiceTestContext, -) -> None: - pipeline = _make_pipeline(pipeline_id="pipeline-1", workflow_id="wf-live") - workflows = [ - _make_workflow( - workflow_id=f"wf-{index}", - app_id="pipeline-1", - created_by="user-1", - ) - for index in range(1, 4) - ] - for index, workflow in enumerate(workflows, start=1): - workflow.version = f"2024-01-0{index}T00:00:00" - workflow.marked_name = f"release-{index}" - _persist(rag_pipeline_service.session, *workflows) - - result, has_more = rag_pipeline_service.service.get_all_published_workflow( - session=rag_pipeline_service.session, - pipeline=pipeline, - page=1, - limit=2, - user_id="user-1", - named_only=True, - ) - - assert [workflow.id for workflow in result] == ["wf-3", "wf-2"] - assert has_more is True - - -# --- sync_draft_workflow --- - - -def test_sync_draft_workflow_creates_new_when_none_exists( - rag_pipeline_service: RagPipelineServiceTestContext, -) -> None: - pipeline = _make_pipeline(workflow_id=None) - account = _make_account() - - result = rag_pipeline_service.service.sync_draft_workflow( - pipeline=pipeline, - graph={"nodes": []}, - unique_hash=None, - account=account, - environment_variables=[], - conversation_variables=[], - rag_pipeline_variables=[], - ) - - assert result.app_id == "p1" - assert pipeline.workflow_id == result.id - - -def test_sync_draft_workflow_raises_on_hash_mismatch( - rag_pipeline_service: RagPipelineServiceTestContext, -) -> None: - from services.errors.app import WorkflowHashNotEqualError - - existing_wf = _make_workflow(graph={"nodes": [{"id": "old"}]}) - _persist(rag_pipeline_service.session, existing_wf) - - pipeline = _make_pipeline() - account = _make_account() - - with pytest.raises(WorkflowHashNotEqualError): - rag_pipeline_service.service.sync_draft_workflow( - pipeline=pipeline, - graph={"nodes": []}, - unique_hash="hash-different", - account=account, - environment_variables=[], - conversation_variables=[], - rag_pipeline_variables=[], - ) - - -def test_sync_draft_workflow_updates_existing( - rag_pipeline_service: RagPipelineServiceTestContext, -) -> None: - existing_wf = _make_workflow(graph={"nodes": [{"id": "old"}]}) - _persist(rag_pipeline_service.session, existing_wf) - - pipeline = _make_pipeline() - account = _make_account() - - result = rag_pipeline_service.service.sync_draft_workflow( - pipeline=pipeline, - graph={"nodes": [{"id": "n1"}]}, - unique_hash=existing_wf.unique_hash, - account=account, - environment_variables=[], - conversation_variables=[], - rag_pipeline_variables=[], - ) - - assert result is existing_wf - assert result.updated_by == "u1" - assert result.graph_dict == {"nodes": [{"id": "n1"}]} - - -# --- get_default_block_config --- - - -def test_get_default_block_config_returns_config_for_valid_type( - mocker: MockerFixture, rag_pipeline_service: RagPipelineServiceTestContext -) -> None: - fake_node_class = mocker.Mock() - fake_node_class.get_default_config.return_value = {"type": "start", "config": {}} - - # Use a simpler approach: test with a known valid node type - - mocker.patch( - "services.rag_pipeline.rag_pipeline.get_node_type_classes_mapping", - return_value={BuiltinNodeTypes.START: {"1": fake_node_class}}, - ) - mocker.patch("services.rag_pipeline.rag_pipeline.LATEST_VERSION", "1") - - result = rag_pipeline_service.service.get_default_block_config("start") - - assert result == {"type": "start", "config": {}} - - -def test_get_default_block_config_returns_none_for_unmapped_type( - rag_pipeline_service: RagPipelineServiceTestContext, -) -> None: - assert rag_pipeline_service.service.get_default_block_config("nonexistent-type") is None - - -# --- update_workflow --- - - -def test_update_workflow_updates_allowed_fields( - rag_pipeline_service: RagPipelineServiceTestContext, -) -> None: - workflow = _make_workflow() - _persist(rag_pipeline_service.session, workflow) - workflow_ref = WorkflowRef(tenant_id="t1", owner_id="pipeline-1", workflow_id="wf-1") - workflow.app_id = "pipeline-1" - rag_pipeline_service.session.commit() - - result = rag_pipeline_service.service.update_workflow( - session=rag_pipeline_service.session, - account_id="u1", - data={"marked_name": "v1", "marked_comment": "release", "disallowed": "hacked"}, - workflow_ref=workflow_ref, - ) - - assert result.marked_name == "v1" - assert result.marked_comment == "release" - assert not hasattr(result, "disallowed") - assert result.updated_by == "u1" - - -def test_update_workflow_returns_none_when_not_found( - rag_pipeline_service: RagPipelineServiceTestContext, -) -> None: - result = rag_pipeline_service.service.update_workflow( - session=rag_pipeline_service.session, - account_id="u1", - data={"marked_name": "v1"}, - workflow_ref=WorkflowRef(tenant_id="t1", owner_id="pipeline-1", workflow_id="wf-missing"), - ) - - assert result is None - - -def test_update_workflow_with_ref_scopes_lookup_to_pipeline( - rag_pipeline_service: RagPipelineServiceTestContext, -) -> None: - workflow = _make_workflow(app_id="pipeline-1") - other_owner_workflow = _make_workflow(workflow_id="wf-other", app_id="pipeline-2") - _persist(rag_pipeline_service.session, workflow, other_owner_workflow) - workflow_ref = WorkflowRef(tenant_id="t1", owner_id="pipeline-1", workflow_id="wf-1") - - result = rag_pipeline_service.service.update_workflow( - session=rag_pipeline_service.session, - account_id="u1", - data={"marked_name": "v1"}, - workflow_ref=workflow_ref, - ) - - assert result is workflow - assert other_owner_workflow.marked_name == "" - - -# --- get_rag_pipeline_paginate_workflow_runs --- - - -def test_get_rag_pipeline_paginate_workflow_runs_delegates( - mocker: MockerFixture, rag_pipeline_service: RagPipelineServiceTestContext -) -> None: - expected = mocker.Mock() - repo_mock = mocker.Mock() - repo_mock.get_paginated_workflow_runs.return_value = expected - rag_pipeline_service.service._workflow_run_repo = repo_mock - - pipeline = _make_pipeline() - result = rag_pipeline_service.service.get_rag_pipeline_paginate_workflow_runs( - pipeline, {"limit": 10, "last_id": "abc"} - ) - - assert result is expected - repo_mock.get_paginated_workflow_runs.assert_called_once_with( - tenant_id="t1", - app_id="p1", - triggered_from=mocker.ANY, - limit=10, - last_id="abc", - ) - - -# --- get_rag_pipeline_workflow_run --- - - -def test_get_rag_pipeline_workflow_run_delegates( - mocker: MockerFixture, rag_pipeline_service: RagPipelineServiceTestContext -) -> None: - expected = mocker.Mock() - repo_mock = mocker.Mock() - repo_mock.get_workflow_run_by_id.return_value = expected - rag_pipeline_service.service._workflow_run_repo = repo_mock - - pipeline = _make_pipeline() - result = rag_pipeline_service.service.get_rag_pipeline_workflow_run(pipeline, "run-1") - - assert result is expected - repo_mock.get_workflow_run_by_id.assert_called_once_with(tenant_id="t1", app_id="p1", run_id="run-1") - - -# --- is_workflow_exist --- - - -def test_is_workflow_exist_returns_true_when_draft_exists( - rag_pipeline_service: RagPipelineServiceTestContext, -) -> None: - _persist(rag_pipeline_service.session, _make_workflow()) - - pipeline = _make_pipeline() - assert rag_pipeline_service.service.is_workflow_exist(pipeline) is True - - -def test_is_workflow_exist_returns_false_when_no_draft( - rag_pipeline_service: RagPipelineServiceTestContext, -) -> None: - pipeline = _make_pipeline() - assert rag_pipeline_service.service.is_workflow_exist(pipeline) is False - - -# --- publish_workflow --- - - -def test_publish_workflow_success(mocker: MockerFixture, rag_pipeline_service: RagPipelineServiceTestContext) -> None: - graph = { - "nodes": [ - { - "data": { - "type": "knowledge-index", - "dataset_id": "d1", - "chunk_structure": "paragraph", - "indexing_technique": "high_quality", - "process_rule": {"mode": "automatic"}, - "retrieval_model": {"search_method": "hybrid_search", "top_k": 3}, - } - } - ] - } - draft_workflow = _make_workflow(workflow_id="wf-draft", graph=graph) - pipeline = _make_pipeline(workflow_id="wf-old-published") - dataset = _make_dataset() - _persist(rag_pipeline_service.session, draft_workflow, pipeline, dataset) - mocker.patch("services.rag_pipeline.rag_pipeline.KnowledgeConfiguration.model_validate", return_value=mocker.Mock()) - mock_dataset_service_class = mocker.patch("services.knowledge.dataset_service.DatasetService") - - result = rag_pipeline_service.service.publish_workflow( - session=rag_pipeline_service.session, pipeline=pipeline, account=_make_account() - ) - - assert result.app_id == pipeline.id - assert result.graph_dict == graph - mock_dataset_service_class.update_rag_pipeline_dataset_settings.assert_called_once() - - -def test_publish_workflow_rejects_missing_llm_environment_reference( - rag_pipeline_service: RagPipelineServiceTestContext, -) -> None: - draft_workflow = _make_workflow( - workflow_id="wf-draft", - tenant_id="tenant", - app_id="pipeline", - graph={ - "nodes": [ - { - "id": "llm-node", - "data": { - "type": "llm", - "model": {"provider": "provider", "name": "model", "mode": "chat"}, - "model_selector": ["env", "missing_model"], - }, - } - ] - }, - ) - _persist(rag_pipeline_service.session, draft_workflow) - - with pytest.raises(ValueError, match="missing_model.*not found"): - rag_pipeline_service.service.publish_workflow( - session=rag_pipeline_service.session, - pipeline=_make_pipeline(pipeline_id="pipeline", tenant_id="tenant"), - account=_make_account(account_id="account", tenant_id="tenant"), - ) - - -# --- run_datasource_workflow_node --- - - -def test_run_datasource_workflow_node_website_crawl( - mocker: MockerFixture, rag_pipeline_service: RagPipelineServiceTestContext -) -> None: - from core.datasource.entities.datasource_entities import DatasourceProviderType - - # 1. Setup workflow and node - pipeline = _make_pipeline() - - workflow = _make_workflow( - graph={ - "nodes": [ - { - "id": "node-1", - "data": { - "type": "datasource", - "plugin_id": "p-1", - "provider_name": "firecrawl", - "datasource_name": "website_crawl", - "datasource_parameters": {"url": {"value": "{{#start.url#}}"}}, - }, - } - ] - } - ) - mocker.patch.object(rag_pipeline_service.service, "get_published_workflow", return_value=workflow) - - # 2. Mock DatasourceManager and Runtime - mock_runtime = mocker.Mock() - mock_runtime.datasource_provider_type.return_value = DatasourceProviderType.WEBSITE_CRAWL - - # Mock the generator result for website crawl - def mock_crawl_gen(**kwargs): - yield mocker.Mock(result=mocker.Mock(status="processing", total=10, completed=2)) - yield mocker.Mock( - result=mocker.Mock(status="completed", total=10, completed=10, web_info_list=[{"title": "test"}]) - ) - - mock_runtime.get_website_crawl.side_effect = mock_crawl_gen - - mocker.patch( - "core.datasource.datasource_manager.DatasourceManager.get_datasource_runtime", - return_value=mock_runtime, - ) - - # 3. Mock DatasourceProviderService - mocker.patch( - "services.rag_pipeline.rag_pipeline.DatasourceProviderService.get_datasource_credentials", - return_value={"api_key": "sk-123"}, - ) - - # 4. Mock Enums to avoid import issues or for consistency - mocker.patch("services.rag_pipeline.rag_pipeline.DatasourceProviderType", DatasourceProviderType) - - # 5. Run test - gen = rag_pipeline_service.service.run_datasource_workflow_node( - pipeline=pipeline, - node_id="node-1", - user_inputs={"url": "https://example.com"}, - account=_make_account(), - datasource_type="website_crawl", - is_published=True, - datasource_providers=rag_pipeline_module.DatasourceProviderService( - credentials=create_autospec(DatasourceProviderCredentialStore, instance=True) - ), - ) - - events = list(gen) - - # 6. Assertions - assert len(events) == 2 - assert events[0]["total"] == 10 - assert events[0]["completed"] == 2 - assert events[1]["data"] == [{"title": "test"}] - assert events[1]["total"] == 10 - assert events[1]["completed"] == 10 - - -# --- run_datasource_node_preview --- - - -def test_run_datasource_node_preview_online_document( - mocker: MockerFixture, rag_pipeline_service: RagPipelineServiceTestContext -) -> None: - from core.datasource.entities.datasource_entities import DatasourceMessage, DatasourceProviderType - - # 1. Setup workflow and node - pipeline = _make_pipeline() - - workflow = _make_workflow( - graph={ - "nodes": [ - { - "id": "node-1", - "data": { - "type": "datasource", - "plugin_id": "p-1", - "provider_name": "notion", - "datasource_name": "online_document", - "datasource_parameters": { - "workspace_id": {"value": "ws-1"}, - "page_id": {"value": "pg-1"}, - "type": {"value": "page"}, - }, - }, - } - ] - } - ) - mocker.patch.object(rag_pipeline_service.service, "get_published_workflow", return_value=workflow) - - # 2. Mock Runtime and results - mock_runtime = mocker.Mock() - - def mock_doc_gen(**kwargs): - # Yield a variable message - msg1 = DatasourceMessage( - type=DatasourceMessage.MessageType.VARIABLE, - message=DatasourceMessage.VariableMessage(variable_name="content", variable_value="Hello ", stream=True), - ) - yield msg1 - msg2 = DatasourceMessage( - type=DatasourceMessage.MessageType.VARIABLE, - message=DatasourceMessage.VariableMessage(variable_name="content", variable_value="World", stream=True), - ) - yield msg2 - - mock_runtime.get_online_document_page_content.side_effect = mock_doc_gen - mocker.patch( - "core.datasource.datasource_manager.DatasourceManager.get_datasource_runtime", - return_value=mock_runtime, - ) - mocker.patch( - "services.rag_pipeline.rag_pipeline.DatasourceProviderService.get_datasource_credentials", - return_value={"token": "abc"}, - ) - mocker.patch("services.rag_pipeline.rag_pipeline.DatasourceProviderType", DatasourceProviderType) - - # 3. Run test - result = rag_pipeline_service.service.run_datasource_node_preview( - pipeline=pipeline, - node_id="node-1", - user_inputs={}, - account=_make_account(), - datasource_type="online_document", - is_published=True, - datasource_providers=rag_pipeline_module.DatasourceProviderService( - credentials=create_autospec(DatasourceProviderCredentialStore, instance=True) - ), - ) - - # 4. Assertions - assert result == {"content": "Hello World"} - - -# --- _handle_node_run_result --- - - -def test_handle_node_run_result_success( - mocker: MockerFixture, rag_pipeline_service: RagPipelineServiceTestContext -) -> None: - from graphon.engine_events import NodeRunSucceededEvent - from graphon.node_events.base import NodeRunResult - - # 1. Setup mock node and result - node_instance = mocker.Mock() - node_instance.workflow_id = "wf-1" - node_instance.node_type = "start" - node_instance.title = "Start" - - node_run_result = NodeRunResult( - status=WorkflowNodeExecutionStatus.SUCCEEDED, - inputs={"q": "hi"}, - outputs={"ans": "hello"}, - metadata={WorkflowNodeExecutionMetadataKey.TOTAL_TOKENS: 10}, - ) - - def mock_getter(): - event = NodeRunSucceededEvent( - id="event-1", - start_at=time.time(), - node_id="node-1", - node_type="start", - node_run_result=node_run_result, - route_node_id=None, - ) - yield event - - # 2. Run test - result = rag_pipeline_service.service._handle_node_run_result( - getter=lambda: (node_instance, mock_getter()), start_at=time.perf_counter(), tenant_id="t1", node_id="node-1" - ) - - # 3. Assertions - assert result.status == WorkflowNodeExecutionStatus.SUCCEEDED - assert result.inputs == {"q": "hi"} - assert result.outputs == {"ans": "hello"} - assert result.metadata == {WorkflowNodeExecutionMetadataKey.TOTAL_TOKENS: 10} - - -# --- get_first_step_parameters / get_second_step_parameters --- - - -def test_get_first_step_parameters_success( - mocker: MockerFixture, rag_pipeline_service: RagPipelineServiceTestContext -) -> None: - # 1. Setup mock workflow - pipeline = _make_pipeline() - workflow = _make_workflow( - graph={"nodes": [{"id": "node-1", "data": {"datasource_parameters": {"url": {"value": "{{#start.url#}}"}}}}]}, - rag_pipeline_variables=[_rag_variable("url")], - ) - mocker.patch.object(rag_pipeline_service.service, "get_published_workflow", return_value=workflow) - - # 2. Run test - result = rag_pipeline_service.service.get_first_step_parameters(pipeline=pipeline, node_id="node-1", is_draft=False) - - # 3. Assertions - assert len(result) == 1 - assert result[0]["variable"] == "url" - - -def test_get_second_step_parameters_success( - mocker: MockerFixture, rag_pipeline_service: RagPipelineServiceTestContext -) -> None: - # 1. Setup mock workflow - pipeline = _make_pipeline() - workflow = _make_workflow( - graph={ - "nodes": [ - { - "id": "node-1", - "data": {}, # Second step logic is slightly different in how it gets variables - } - ] - }, - rag_pipeline_variables=[_rag_variable("var1", belong_to_node_id="other-node")], - ) - mocker.patch.object(rag_pipeline_service.service, "get_published_workflow", return_value=workflow) - - # 2. Run test - result = rag_pipeline_service.service.get_second_step_parameters( - pipeline=pipeline, node_id="node-1", is_draft=False - ) - - # 3. Assertions - # Note: get_second_step_parameters also filters by variable names found in node data - # (Checking the code again, it seems to iterate through nodes but doesn't do much with variables yet) - # Wait, let me check the code for get_second_step_parameters again. - assert len(result) == 0 # Based on current implementation which seems to filter but no logic added yet? - - -# --- publish_customized_pipeline_template --- - - -def test_publish_customized_pipeline_template_success( - mocker: MockerFixture, rag_pipeline_service: RagPipelineServiceTestContext -) -> None: - pipeline = _make_pipeline(workflow_id="wf-1", is_published=True) - workflow = _make_workflow(workflow_id="wf-1") - session = rag_pipeline_service.session - dataset = _make_dataset() - dataset.chunk_structure = "paragraph" - _persist(session, pipeline, workflow, dataset) - - # Mock RagPipelineDslService - mock_dsl_service = mocker.Mock() - mock_dsl_service.export_rag_pipeline_dsl.return_value = "dsl: content" - mocker.patch.object(rag_pipeline_module, "RagPipelineDslService", return_value=mock_dsl_service) - - account = _make_account(account_id="user-123") - - rag_pipeline_service.service.publish_customized_pipeline_template( - pipeline, dataset, _make_template_args(), account, session=session - ) - - mock_dsl_service.export_rag_pipeline_dsl.assert_called_once_with(pipeline=pipeline, include_secret=True) - templates = session.query(PipelineCustomizedTemplate).all() - assert len(templates) == 1 - assert templates[0].name == "New Template" - - -# --- get_datasource_plugins --- - - -def test_get_datasource_plugins_success( - mocker: MockerFixture, rag_pipeline_service: RagPipelineServiceTestContext -) -> None: - # 1. Setup mocks - dataset = _make_dataset() - - pipeline = _make_pipeline(workflow_id="wf-1") - - workflow = _make_workflow( - workflow_id="wf-1", - graph={ - "nodes": [ - { - "id": "node-1", - "data": { - "type": "datasource", - "plugin_id": "p-1", - "provider_name": "notion", - "provider_type": "online_document", - "title": "Notion", - }, - } - ] - }, - ) - _persist(rag_pipeline_service.session, dataset, pipeline, workflow) - - # Mock DatasourceProviderService - mock_provider_service = mocker.Mock() - mock_provider_service.list_datasource_credentials.return_value = [ - {"id": "c1", "name": "Cred 1", "type": "token", "is_default": True} - ] - mocker.patch("services.rag_pipeline.rag_pipeline.DatasourceProviderService", return_value=mock_provider_service) - - # 2. Run test - result = rag_pipeline_service.service.get_datasource_plugins( - "t1", - "d1", - True, - credential_query=mock.create_autospec(CredentialQuery, instance=True), - datasource_providers=rag_pipeline_module.DatasourceProviderService( - credentials=create_autospec(DatasourceProviderCredentialStore, instance=True) - ), - ) - - # 3. Assertions - assert len(result) == 1 - assert result[0]["node_id"] == "node-1" - assert result[0]["credentials"][0]["id"] == "c1" - - -# --- retry_error_document --- - - -def test_retry_error_document_success( - mocker: MockerFixture, rag_pipeline_service: RagPipelineServiceTestContext, pipeline_generator -) -> None: - dataset = _make_dataset() - document = _make_document() - - log = DocumentPipelineExecutionLog( - pipeline_id="p-1", - document_id="doc-1", - datasource_type="upload_file", - datasource_info="{}", - datasource_node_id="node-id", - input_data={}, - created_by="account-id", - ) - # Ensure it's a string if it's used as JSON later - - pipeline = Pipeline(tenant_id="tenant-id", name="Test Pipeline", workflow_id="wf-1") - pipeline.id = "p-1" - - workflow = _make_workflow(workflow_id="wf-1", tenant_id="tenant-id", app_id="p-1") - _persist(rag_pipeline_service.session, log, pipeline, workflow) - - # Mock PipelineGenerator - mock_gen_instance = pipeline_generator - - # 2. Run test - user = mocker.Mock() - rag_pipeline_service.service.retry_error_document(dataset, document, user, generator=pipeline_generator) - - # 3. Assertions - mock_gen_instance.generate.assert_called_once() - - -# --- set_datasource_variables --- - - -def test_set_datasource_variables_success( - mocker: MockerFixture, rag_pipeline_service: RagPipelineServiceTestContext -) -> None: - from graphon.entities.workflow_node_execution import WorkflowNodeExecution - from models.dataset import Pipeline - - # 1. Setup mocks - pipeline = Pipeline( - tenant_id="t1", - name="Test Pipeline", - ) - pipeline.id = "p-1" - - draft_wf = mocker.Mock() - draft_wf.id = "wf-1" - draft_wf.get_enclosing_node_type_and_id.return_value = None # Avoid unpacking error - mocker.patch.object(rag_pipeline_service.service, "get_draft_workflow", return_value=draft_wf) - - execution = WorkflowNodeExecution( - id="exec-1", - node_execution_id="node-exec-1", - workflow_id="wf-1", - workflow_execution_id="run-1", - index=1, - predecessor_node_id=None, - node_id="node-1", - node_type=BuiltinNodeTypes.DATASOURCE, - title="Datasource", - process_data={}, - inputs={}, - outputs={}, - status=WorkflowNodeExecutionStatus.SUCCEEDED, - error=None, - elapsed_time=0, - metadata={}, - created_at=datetime(2026, 1, 1), - finished_at=None, - ) - mocker.patch.object(rag_pipeline_service.service, "_handle_node_run_result", return_value=execution) - - # Mock Repository - mock_repo_instance = mocker.Mock() - mocker.patch( - "services.rag_pipeline.rag_pipeline.SQLAlchemyWorkflowNodeExecutionRepository", - return_value=mock_repo_instance, - ) - # Repository._to_db_model is also called - mock_db_exec = mocker.Mock() - mock_db_exec.node_id = "node-1" - mock_db_exec.node_type = "datasource" - mock_repo_instance._to_db_model.return_value = mock_db_exec - - # Mock DraftVariableSaver - mock_saver_instance = mocker.Mock() - mocker.patch("services.rag_pipeline.rag_pipeline.DraftVariableSaver", return_value=mock_saver_instance) - - # 2. Run test - args = {"start_node_id": "node-1"} - user = mocker.Mock() - user.id = "user-1" - rag_pipeline_service.service.set_datasource_variables(pipeline, args, user) - - # 3. Assertions - mock_repo_instance.save.assert_called_once() - mock_saver_instance.save.assert_called_once() - - -# --- Utility Methods --- - - -def test_get_draft_workflow_success(rag_pipeline_service: RagPipelineServiceTestContext) -> None: - # 1. Setup mocks - pipeline = _make_pipeline() - - workflow = _make_workflow() - _persist(rag_pipeline_service.session, workflow) - - # 2. Run test - result = rag_pipeline_service.service.get_draft_workflow(pipeline) - - # 3. Assertions - assert result == workflow - - -def test_get_published_workflow_success(rag_pipeline_service: RagPipelineServiceTestContext) -> None: - # 1. Setup mocks - pipeline = _make_pipeline(workflow_id="wf-pub") - - workflow = _make_workflow(workflow_id="wf-pub") - _persist(rag_pipeline_service.session, workflow) - - # 2. Run test - result = rag_pipeline_service.service.get_published_workflow(pipeline) - - # 3. Assertions - assert result == workflow - - -def test_get_default_block_configs_success(rag_pipeline_service: RagPipelineServiceTestContext) -> None: - # This calls static methods on node classes, should be safe with default mocks or as-is - # unless they access db. - result = rag_pipeline_service.service.get_default_block_configs() - assert isinstance(result, list) - assert len(result) > 0 - - -def test_get_default_block_config_success(rag_pipeline_service: RagPipelineServiceTestContext) -> None: - result = rag_pipeline_service.service.get_default_block_config(BuiltinNodeTypes.LLM) - assert result is not None - assert result["type"] == "llm" - - -def test_publish_workflow_raises_when_draft_workflow_missing( - rag_pipeline_service: RagPipelineServiceTestContext, -) -> None: - pipeline = _make_pipeline() - account = _make_account() - - with pytest.raises(ValueError, match="No valid workflow found"): - rag_pipeline_service.service.publish_workflow( - session=rag_pipeline_service.session, pipeline=pipeline, account=account - ) - - -def test_get_default_block_config_returns_none_when_mapped_type_missing( - mocker: MockerFixture, rag_pipeline_service: RagPipelineServiceTestContext -) -> None: - mocker.patch("services.rag_pipeline.rag_pipeline.get_node_type_classes_mapping", return_value={}) - - assert rag_pipeline_service.service.get_default_block_config(BuiltinNodeTypes.START) is None - - -def test_get_default_block_config_injects_http_request_filter( - mocker: MockerFixture, rag_pipeline_service: RagPipelineServiceTestContext -) -> None: - fake_node_cls = mocker.Mock() - fake_node_cls.get_default_config.return_value = {"type": "http-request"} - mocker.patch( - "services.rag_pipeline.rag_pipeline.get_node_type_classes_mapping", - return_value={BuiltinNodeTypes.HTTP_REQUEST: {"1": fake_node_cls}}, - ) - mocker.patch("services.rag_pipeline.rag_pipeline.LATEST_VERSION", "1") - - rag_pipeline_service.service.get_default_block_config(BuiltinNodeTypes.HTTP_REQUEST) - - called_filters = fake_node_cls.get_default_config.call_args.kwargs["filters"] - assert "http_request_config" in called_filters - - -def test_run_draft_workflow_node_raises_when_workflow_missing( - mocker: MockerFixture, rag_pipeline_service: RagPipelineServiceTestContext -) -> None: - pipeline = _make_pipeline() - account = _make_account() - mocker.patch.object(rag_pipeline_service.service, "get_draft_workflow", return_value=None) - - with pytest.raises(ValueError, match="Workflow not initialized"): - rag_pipeline_service.service.run_draft_workflow_node(pipeline, "node-1", {}, account) - - -def test_run_draft_workflow_node_seeds_llm_environment_variable( - mocker: MockerFixture, rag_pipeline_service: RagPipelineServiceTestContext -) -> None: - from factories import variable_factory - - pipeline = _make_pipeline() - account = _make_account() - llm_environment_variable = variable_factory.build_environment_variable_from_mapping( - { - "id": "env-1", - "name": "for_summarize", - "value_type": "llm", - "value": { - "provider": "langgenius/openai/openai", - "name": "gpt-4o", - "mode": "chat", - }, - "description": "Shared summarization model", - } - ) - draft_workflow = _make_workflow(workflow_id="wf-1") - draft_workflow.environment_variables = [llm_environment_variable] - mocker.patch.object(draft_workflow, "get_node_config_by_id", return_value={"id": "node-1"}) - mocker.patch.object(draft_workflow, "get_enclosing_node_type_and_id", return_value=None) - mocker.patch.object(rag_pipeline_service.service, "get_draft_workflow", return_value=draft_workflow) - - execution = SimpleNamespace(id="exec-1", node_id="node-1", node_type="llm", process_data={}, outputs={}) - handle_node_run_result = mocker.patch.object( - rag_pipeline_service.service, "_handle_node_run_result", return_value=execution - ) - single_step_run = mocker.patch("services.rag_pipeline.rag_pipeline.WorkflowEntry.single_step_run") - - repo = mocker.Mock() - mocker.patch( - "services.rag_pipeline.rag_pipeline.DifyCoreRepositoryFactory.create_workflow_node_execution_repository", - return_value=repo, - ) - rag_pipeline_service.service._node_execution_service_repo = mocker.Mock( - get_execution_by_id=mocker.Mock(return_value="db") - ) - mocker.patch("services.rag_pipeline.rag_pipeline.DraftVariableSaver", return_value=mocker.Mock()) - rag_pipeline_service.service.run_draft_workflow_node(pipeline, "node-1", {}, account) - - getter = handle_node_run_result.call_args.kwargs["getter"] - getter() - variable_pool = single_step_run.call_args.kwargs["variable_pool"] - llm_model = variable_pool.get(["env", "for_summarize"]) - assert llm_model is not None - assert llm_model.value == { - "provider": "langgenius/openai/openai", - "name": "gpt-4o", - "mode": "chat", - } - - -def test_run_draft_workflow_node_saves_execution_and_variables( - mocker: MockerFixture, rag_pipeline_service: RagPipelineServiceTestContext -) -> None: - pipeline = _make_pipeline() - account = _make_account() - draft_workflow = _make_workflow(workflow_id="wf-1") - mocker.patch.object(draft_workflow, "get_node_config_by_id", return_value={"id": "node-1"}) - mocker.patch.object(draft_workflow, "get_enclosing_node_type_and_id", return_value=("loop", "enclosing-node")) - mocker.patch.object(rag_pipeline_service.service, "get_draft_workflow", return_value=draft_workflow) - - execution = SimpleNamespace(id="exec-1", node_id="node-1", node_type="llm", process_data={}, outputs={}) - mocker.patch.object(rag_pipeline_service.service, "_handle_node_run_result", return_value=execution) - - repo = mocker.Mock() - mocker.patch( - "services.rag_pipeline.rag_pipeline.DifyCoreRepositoryFactory.create_workflow_node_execution_repository", - return_value=repo, - ) - rag_pipeline_service.service._node_execution_service_repo = mocker.Mock( - get_execution_by_id=mocker.Mock(return_value="db") - ) - saver = mocker.Mock() - mocker.patch("services.rag_pipeline.rag_pipeline.DraftVariableSaver", return_value=saver) - - result = rag_pipeline_service.service.run_draft_workflow_node(pipeline, "node-1", {"q": "x"}, account) - - assert result == "db" - assert execution.workflow_id == "wf-1" - repo.save.assert_called_once_with(execution) - saver.save.assert_called_once() - - -def test_run_datasource_workflow_node_returns_error_when_workflow_missing( - mocker: MockerFixture, rag_pipeline_service: RagPipelineServiceTestContext -) -> None: - pipeline = _make_pipeline() - mocker.patch.object(rag_pipeline_service.service, "get_draft_workflow", return_value=None) - - events = list( - rag_pipeline_service.service.run_datasource_workflow_node( - pipeline=pipeline, - node_id="node-1", - user_inputs={}, - account=_make_account(), - datasource_type="online_document", - is_published=False, - datasource_providers=rag_pipeline_module.DatasourceProviderService( - credentials=create_autospec(DatasourceProviderCredentialStore, instance=True) - ), - ) - ) - - assert events[0]["event"] == "datasource_error" - - -def test_run_datasource_workflow_node_online_document_success( - mocker: MockerFixture, rag_pipeline_service: RagPipelineServiceTestContext -) -> None: - from core.datasource.entities.datasource_entities import DatasourceProviderType - - pipeline = _make_pipeline() - workflow = _make_workflow( - graph={ - "nodes": [ - { - "id": "node-1", - "data": { - "type": "datasource", - "plugin_id": "pid", - "provider_name": "notion", - "datasource_name": "online_document", - "datasource_parameters": {"workspace_id": {"value": None}, "page_id": {"value": "fixed"}}, - }, - } - ] - } - ) - mocker.patch.object(rag_pipeline_service.service, "get_published_workflow", return_value=workflow) - - runtime = mocker.Mock() - runtime.runtime = SimpleNamespace(credentials=None) - runtime.datasource_provider_type.return_value = DatasourceProviderType.ONLINE_DOCUMENT - runtime.get_online_document_pages.return_value = [SimpleNamespace(result=[{"id": "pg-1"}])] - mocker.patch("core.datasource.datasource_manager.DatasourceManager.get_datasource_runtime", return_value=runtime) - mocker.patch( - "services.rag_pipeline.rag_pipeline.DatasourceProviderService.get_datasource_credentials", - return_value={"token": "x"}, - ) - - events = list( - rag_pipeline_service.service.run_datasource_workflow_node( - pipeline=pipeline, - node_id="node-1", - user_inputs={}, - account=_make_account(), - datasource_type=DatasourceProviderType.ONLINE_DOCUMENT, - is_published=True, - datasource_providers=rag_pipeline_module.DatasourceProviderService( - credentials=create_autospec(DatasourceProviderCredentialStore, instance=True) - ), - ) - ) - - assert events[0]["event"] == "datasource_processing" - assert events[1]["event"] == "datasource_completed" - - -def test_run_datasource_workflow_node_online_drive_success( - mocker: MockerFixture, rag_pipeline_service: RagPipelineServiceTestContext -) -> None: - from core.datasource.entities.datasource_entities import DatasourceProviderType - - pipeline = _make_pipeline() - workflow = _make_workflow( - graph={ - "nodes": [ - { - "id": "node-1", - "data": { - "type": "datasource", - "plugin_id": "pid", - "provider_name": "drive", - "datasource_name": "online_drive", - "datasource_parameters": { - "bucket": {"value": "bucket-1"}, - "next_page_parameters": {"value": []}, - }, - }, - } - ] - } - ) - mocker.patch.object(rag_pipeline_service.service, "get_published_workflow", return_value=workflow) - - runtime = mocker.Mock() - runtime.runtime = SimpleNamespace(credentials=None) - runtime.datasource_provider_type.return_value = DatasourceProviderType.ONLINE_DRIVE - runtime.online_drive_browse_files.return_value = [SimpleNamespace(result=[{"name": "f1"}])] - mocker.patch("core.datasource.datasource_manager.DatasourceManager.get_datasource_runtime", return_value=runtime) - mocker.patch( - "services.rag_pipeline.rag_pipeline.DatasourceProviderService.get_datasource_credentials", - return_value={}, - ) - - events = list( - rag_pipeline_service.service.run_datasource_workflow_node( - pipeline=pipeline, - node_id="node-1", - user_inputs={"bucket": "bucket-1"}, - account=_make_account(), - datasource_type=DatasourceProviderType.ONLINE_DRIVE, - is_published=True, - datasource_providers=rag_pipeline_module.DatasourceProviderService( - credentials=create_autospec(DatasourceProviderCredentialStore, instance=True) - ), - ) - ) - - assert events[0]["event"] == "datasource_processing" - assert events[1]["event"] == "datasource_completed" - - -def test_handle_node_run_result_default_value_strategy( - mocker: MockerFixture, rag_pipeline_service: RagPipelineServiceTestContext -) -> None: - from datetime import datetime - - from graphon.engine_events import NodeRunFailedEvent - from graphon.node_events.base import NodeRunResult - - node_instance = _make_start_node( - StartNodeData( - title="Start", - error_strategy=ErrorStrategy.DEFAULT_VALUE, - default_value=[{"key": "fallback", "value": "ok", "type": "string"}], - ) - ) - - failed_result = NodeRunResult( - status=WorkflowNodeExecutionStatus.FAILED, - error="boom", - error_type="runtime_error", - inputs={"x": 1}, - ) - - def _events(): - yield NodeRunFailedEvent( - id="e-1", - node_id="node-1", - node_type=BuiltinNodeTypes.START, - start_at=datetime.now(), - error="boom", - node_run_result=failed_result, - ) - - result = rag_pipeline_service.service._handle_node_run_result( - getter=lambda: (node_instance, _events()), - start_at=time.perf_counter(), - tenant_id="t1", - node_id="node-1", - ) - - assert result.status == WorkflowNodeExecutionStatus.EXCEPTION - assert result.outputs - assert result.outputs["fallback"] == "ok" - - -def test_get_first_step_parameters_raises_when_datasource_node_missing( - mocker: MockerFixture, rag_pipeline_service: RagPipelineServiceTestContext -) -> None: - workflow = _make_workflow(graph={"nodes": []}, rag_pipeline_variables=[_rag_variable("url")]) - mocker.patch.object(rag_pipeline_service.service, "get_published_workflow", return_value=workflow) - - with pytest.raises(ValueError, match="Datasource node data not found"): - rag_pipeline_service.service.get_first_step_parameters(_make_pipeline(), "missing-node") - - -def test_get_second_step_parameters_handles_string_and_list_variable_references( - mocker: MockerFixture, rag_pipeline_service: RagPipelineServiceTestContext -) -> None: - workflow = _make_workflow( - rag_pipeline_variables=[ - _rag_variable("url", belong_to_node_id="node-1"), - _rag_variable("bucket"), - _rag_variable("keep", belong_to_node_id="node-1"), - ], - graph={ - "nodes": [ - { - "id": "node-1", - "data": { - "datasource_parameters": { - "u": {"value": "{{#start.url#}}"}, - "b": {"value": ["start", "bucket"]}, - } - }, - } - ] - }, - ) - mocker.patch.object(rag_pipeline_service.service, "get_published_workflow", return_value=workflow) - - result = rag_pipeline_service.service.get_second_step_parameters(_make_pipeline(), "node-1") - - assert [item["variable"] for item in result] == ["keep"] - assert result[0]["belong_to_node_id"] == "node-1" - - -def test_get_rag_pipeline_workflow_run_node_executions_empty_when_run_missing( - mocker: MockerFixture, rag_pipeline_service: RagPipelineServiceTestContext -) -> None: - pipeline = _make_pipeline() - mocker.patch.object(rag_pipeline_service.service, "get_rag_pipeline_workflow_run", return_value=None) - - result = rag_pipeline_service.service.get_rag_pipeline_workflow_run_node_executions( - pipeline=pipeline, run_id="run-1", user=_make_account() - ) - - assert result == [] - - -def test_get_rag_pipeline_workflow_run_node_executions_assembles_configured_repository_results( - mocker: MockerFixture, rag_pipeline_service: RagPipelineServiceTestContext -) -> None: - pipeline = _make_pipeline() - mocker.patch.object( - rag_pipeline_service.service, "get_rag_pipeline_workflow_run", return_value=WorkflowRun(id="run-1") - ) - node_repo = rag_pipeline_service.service._node_execution_service_repo - expected_executions = ["n1", "n2"] - expected_traces = ["retry-n1", "n1", "n2"] - node_repo.get_executions_by_workflow_run = mocker.Mock(return_value=expected_executions) - mock_assemble = mocker.patch( - "services.rag_pipeline.rag_pipeline.assemble_workflow_node_execution_traces", - return_value=expected_traces, - ) - - result = rag_pipeline_service.service.get_rag_pipeline_workflow_run_node_executions( - pipeline=pipeline, run_id="run-1", user=_make_account() - ) - - assert result == expected_traces - node_repo.get_executions_by_workflow_run.assert_called_once_with( - tenant_id=pipeline.tenant_id, - app_id=pipeline.id, - workflow_run_id="run-1", - ) - mock_assemble.assert_called_once_with(expected_executions, node_repo, session=mock.ANY) - - -def test_get_recommended_plugins_returns_empty_when_no_active_plugins( - rag_pipeline_service: RagPipelineServiceTestContext, -) -> None: - result = rag_pipeline_service.service.get_recommended_plugins("all", _make_account(), "t1") - - assert result == { - "installed_recommended_plugins": [], - "uninstalled_recommended_plugins": [], - } - - -def test_get_recommended_plugins_returns_installed_and_uninstalled( - mocker: MockerFixture, rag_pipeline_service: RagPipelineServiceTestContext -) -> None: - plugin_a = _make_recommended_plugin("plugin-a") - plugin_b = _make_recommended_plugin("plugin-b") - _persist(rag_pipeline_service.session, plugin_a, plugin_b) - mocker.patch( - "services.rag_pipeline.rag_pipeline.BuiltinToolManageService.list_builtin_tools", - return_value=[SimpleNamespace(plugin_id="plugin-a", to_dict=lambda: {"plugin_id": "plugin-a"})], - ) - mocker.patch( - "services.rag_pipeline.rag_pipeline.marketplace.batch_fetch_plugin_by_ids", - return_value=[{"plugin_id": "plugin-b", "name": "Plugin B"}], - ) - - result = rag_pipeline_service.service.get_recommended_plugins("tool", _make_account(), "t1") - - assert result["installed_recommended_plugins"] == [{"plugin_id": "plugin-a"}] - assert result["uninstalled_recommended_plugins"] == [{"plugin_id": "plugin-b", "name": "Plugin B"}] - - -def test_get_node_last_run_delegates_to_repository( - mocker: MockerFixture, rag_pipeline_service: RagPipelineServiceTestContext -) -> None: - repo = mocker.Mock() - repo.get_node_last_execution.return_value = "node-exec" - rag_pipeline_service.service._node_execution_service_repo = repo - pipeline = _make_pipeline() - workflow = _make_workflow(workflow_id="wf1") - - result = rag_pipeline_service.service.get_node_last_run(pipeline, workflow, "node-1") - - assert result == "node-exec" - - -def test_set_datasource_variables_raises_when_node_id_missing( - mocker: MockerFixture, rag_pipeline_service: RagPipelineServiceTestContext -) -> None: - pipeline = _make_pipeline() - workflow = _make_workflow() - mocker.patch.object(rag_pipeline_service.service, "get_draft_workflow", return_value=workflow) - - with pytest.raises(ValueError, match="Node id is required"): - rag_pipeline_service.service.set_datasource_variables(pipeline, {"start_node_id": ""}, _make_account()) - - -def test_get_default_block_configs_skips_empty_configs( - mocker: MockerFixture, rag_pipeline_service: RagPipelineServiceTestContext -) -> None: - http_node = mocker.Mock() - http_node.get_default_config.return_value = {"type": "http-request"} - empty_node = mocker.Mock() - empty_node.get_default_config.return_value = None - - mocker.patch( - "services.rag_pipeline.rag_pipeline.get_node_type_classes_mapping", - return_value={ - BuiltinNodeTypes.HTTP_REQUEST: {"1": http_node}, - BuiltinNodeTypes.START: {"1": empty_node}, - }, - ) - mocker.patch("services.rag_pipeline.rag_pipeline.LATEST_VERSION", "1") - - result = rag_pipeline_service.service.get_default_block_configs() - - assert result == [{"type": "http-request"}] - http_node.get_default_config.assert_called_once() - empty_node.get_default_config.assert_called_once() - - -def test_run_datasource_workflow_node_returns_error_when_node_missing( - mocker: MockerFixture, rag_pipeline_service: RagPipelineServiceTestContext -) -> None: - pipeline = _make_pipeline() - workflow = _make_workflow(graph={"nodes": []}) - mocker.patch.object(rag_pipeline_service.service, "get_published_workflow", return_value=workflow) - - events = list( - rag_pipeline_service.service.run_datasource_workflow_node( - pipeline=pipeline, - node_id="missing-node", - user_inputs={}, - account=_make_account(), - datasource_type="online_document", - is_published=True, - datasource_providers=rag_pipeline_module.DatasourceProviderService( - credentials=create_autospec(DatasourceProviderCredentialStore, instance=True) - ), - ) - ) - - assert len(events) == 1 - assert "Datasource node data not found" in events[0]["error"] - - -def test_run_datasource_workflow_node_online_document_exception( - mocker: MockerFixture, rag_pipeline_service: RagPipelineServiceTestContext -) -> None: - pipeline = _make_pipeline() - workflow = _make_workflow( - graph={ - "nodes": [ - { - "id": "node-1", - "data": { - "plugin_id": "plugin-1", - "provider_name": "provider-1", - "datasource_name": "doc", - "datasource_parameters": {}, - }, - } - ] - } - ) - mocker.patch.object(rag_pipeline_service.service, "get_published_workflow", return_value=workflow) - - runtime = mocker.Mock() - - class _FailingIterator: - def __iter__(self): - return self - - def __next__(self): - raise RuntimeError("doc failed") - - runtime.get_online_document_pages.return_value = _FailingIterator() - runtime.datasource_provider_type.return_value = "online_document" - - mocker.patch("core.datasource.datasource_manager.DatasourceManager.get_datasource_runtime", return_value=runtime) - mocker.patch( - "services.rag_pipeline.rag_pipeline.DatasourceProviderService.get_datasource_credentials", - return_value=None, - ) - - events = list( - rag_pipeline_service.service.run_datasource_workflow_node( - pipeline=pipeline, - node_id="node-1", - user_inputs={}, - account=_make_account(), - datasource_type="online_document", - is_published=True, - datasource_providers=rag_pipeline_module.DatasourceProviderService( - credentials=create_autospec(DatasourceProviderCredentialStore, instance=True) - ), - ) - ) - - assert len(events) == 2 - assert events[0]["event"] == "datasource_processing" - assert "doc failed" in events[1]["error"] - - -def test_run_datasource_node_preview_raises_for_stream_non_string( - mocker: MockerFixture, rag_pipeline_service: RagPipelineServiceTestContext -) -> None: - from core.datasource.entities.datasource_entities import DatasourceMessage - - pipeline = _make_pipeline() - workflow = _make_workflow( - graph={ - "nodes": [ - { - "id": "node-1", - "data": { - "plugin_id": "plugin-1", - "provider_name": "provider-1", - "datasource_name": "doc", - "datasource_parameters": {}, - }, - } - ] - } - ) - mocker.patch.object(rag_pipeline_service.service, "get_published_workflow", return_value=workflow) - - runtime = mocker.Mock() - - def _bad_stream_generator[**P](*args: P.args, **kwargs: P.kwargs): - yield DatasourceMessage( - type=DatasourceMessage.MessageType.VARIABLE, - message=DatasourceMessage.VariableMessage(variable_name="content", variable_value=1, stream=True), - ) - - runtime.get_online_document_page_content.side_effect = _bad_stream_generator - runtime.datasource_provider_type.return_value = "online_document" - - mocker.patch("core.datasource.datasource_manager.DatasourceManager.get_datasource_runtime", return_value=runtime) - mocker.patch( - "services.rag_pipeline.rag_pipeline.DatasourceProviderService.get_datasource_credentials", - return_value=None, - ) - - with pytest.raises(RuntimeError, match="must be a string"): - rag_pipeline_service.service.run_datasource_node_preview( - pipeline=pipeline, - node_id="node-1", - user_inputs={}, - account=_make_account(), - datasource_type="online_document", - is_published=True, - datasource_providers=rag_pipeline_module.DatasourceProviderService( - credentials=create_autospec(DatasourceProviderCredentialStore, instance=True) - ), - ) - - -def test_get_first_step_parameters_returns_empty_when_no_rag_variables( - mocker: MockerFixture, rag_pipeline_service: RagPipelineServiceTestContext -) -> None: - workflow = _make_workflow( - graph={"nodes": [{"id": "node-1", "data": {"datasource_parameters": {"url": {"value": "literal"}}}}]}, - rag_pipeline_variables=[], - ) - mocker.patch.object(rag_pipeline_service.service, "get_published_workflow", return_value=workflow) - - result = rag_pipeline_service.service.get_first_step_parameters(_make_pipeline(), "node-1") - - assert result == [] - - -def test_get_second_step_parameters_filters_first_step_variables( - mocker: MockerFixture, rag_pipeline_service: RagPipelineServiceTestContext -) -> None: - workflow = _make_workflow( - graph={ - "nodes": [ - { - "id": "node-1", - "data": { - "datasource_parameters": { - "workspace": {"value": "{{#start.workspace#}}"}, - "bucket": {"value": ["input", "bucket"]}, - } - }, - } - ] - }, - rag_pipeline_variables=[ - _rag_variable("workspace"), - _rag_variable("bucket"), - _rag_variable("keep"), - _rag_variable("other-node", belong_to_node_id="node-x"), - ], - ) - mocker.patch.object(rag_pipeline_service.service, "get_published_workflow", return_value=workflow) - - result = rag_pipeline_service.service.get_second_step_parameters(_make_pipeline(), "node-1") - - assert [item["variable"] for item in result] == ["keep"] - assert result[0]["belong_to_node_id"] == "shared" - - -def test_retry_error_document_raises_when_execution_log_not_found( - rag_pipeline_service: RagPipelineServiceTestContext, - pipeline_generator, -) -> None: - with pytest.raises(ValueError, match="Document pipeline execution log not found"): - rag_pipeline_service.service.retry_error_document( - _make_dataset(), _make_document(), _make_account(), generator=pipeline_generator - ) - - -def test_get_datasource_plugins_raises_when_workflow_not_found( - rag_pipeline_service: RagPipelineServiceTestContext, -) -> None: - dataset = _make_dataset() - pipeline = _make_pipeline(workflow_id="wf-1") - _persist(rag_pipeline_service.session, dataset, pipeline) - - with pytest.raises(ValueError, match="Pipeline or workflow not found"): - rag_pipeline_service.service.get_datasource_plugins( - "t1", - "d1", - True, - credential_query=mock.create_autospec(CredentialQuery, instance=True), - datasource_providers=rag_pipeline_module.DatasourceProviderService( - credentials=create_autospec(DatasourceProviderCredentialStore, instance=True) - ), - ) - - -def test_handle_node_run_result_raises_when_no_terminal_event( - mocker: MockerFixture, rag_pipeline_service: RagPipelineServiceTestContext -) -> None: - node_instance = _make_start_node() - - def _event_generator(): - yield object() - - with pytest.raises(ValueError, match="Node run failed with no run result"): - rag_pipeline_service.service._handle_node_run_result( - getter=lambda: (node_instance, _event_generator()), - start_at=time.perf_counter(), - tenant_id="t1", - node_id="node-1", - ) - - -def test_handle_node_run_result_marks_document_error_for_published_invoke( - rag_pipeline_service: RagPipelineServiceTestContext, -) -> None: - node_instance, event = _make_failed_published_node_run() - dataset = _make_dataset(dataset_id="dataset-1", pipeline_id="pipeline-1", tenant_id="t1") - document = _make_document(document_id="doc-1", dataset_id="dataset-1", tenant_id="t1") - _persist( - rag_pipeline_service.session, - dataset, - document, - ) - - result = rag_pipeline_service.service._handle_node_run_result( - getter=lambda: (node_instance, iter([event])), - start_at=time.perf_counter(), - tenant_id="t1", - node_id="node-1", - ) - - assert result.status == WorkflowNodeExecutionStatus.FAILED - rag_pipeline_service.session.expire_all() - updated_document = rag_pipeline_service.session.get(Document, "doc-1") - assert updated_document is not None - assert updated_document.indexing_status == "error" - assert updated_document.error == "boom" - - -def test_handle_node_run_result_does_not_write_when_pipeline_dataset_is_missing( - rag_pipeline_service: RagPipelineServiceTestContext, -) -> None: - node_instance, event = _make_failed_published_node_run() - - result = rag_pipeline_service.service._handle_node_run_result( - getter=lambda: (node_instance, iter([event])), - start_at=time.perf_counter(), - tenant_id="t1", - node_id="node-1", - ) - - assert result.status == WorkflowNodeExecutionStatus.FAILED - assert rag_pipeline_service.session.query(Document).count() == 0 - - -def test_run_datasource_node_preview_raises_for_unsupported_provider( - mocker: MockerFixture, rag_pipeline_service: RagPipelineServiceTestContext -) -> None: - pipeline = _make_pipeline() - workflow = _make_workflow( - graph={ - "nodes": [ - { - "id": "node-1", - "data": { - "plugin_id": "plugin-1", - "provider_name": "provider-1", - "datasource_name": "doc", - "datasource_parameters": {}, - }, - } - ] - } - ) - mocker.patch.object(rag_pipeline_service.service, "get_published_workflow", return_value=workflow) - runtime = mocker.Mock() - runtime.datasource_provider_type.return_value = "unsupported" - mocker.patch("core.datasource.datasource_manager.DatasourceManager.get_datasource_runtime", return_value=runtime) - mocker.patch( - "services.rag_pipeline.rag_pipeline.DatasourceProviderService.get_datasource_credentials", - return_value=None, - ) - - with pytest.raises(RuntimeError, match="Unsupported datasource provider"): - rag_pipeline_service.service.run_datasource_node_preview( - pipeline=pipeline, - node_id="node-1", - user_inputs={}, - account=_make_account(), - datasource_type="website_crawl", - is_published=True, - datasource_providers=rag_pipeline_module.DatasourceProviderService( - credentials=create_autospec(DatasourceProviderCredentialStore, instance=True) - ), - ) - - -def test_publish_customized_pipeline_template_raises_for_missing_workflow_id( - rag_pipeline_service: RagPipelineServiceTestContext, -) -> None: - pipeline = _make_pipeline(workflow_id=None) - _persist(rag_pipeline_service.session, pipeline) - - with pytest.raises(RagPipelineResourceNotFoundError, match="Pipeline workflow not found"): - rag_pipeline_service.service.publish_customized_pipeline_template( - pipeline, _make_dataset(), _make_template_args(), _make_account(), session=rag_pipeline_service.session - ) - - -def test_get_pipeline_raises_when_dataset_missing( - rag_pipeline_service: RagPipelineServiceTestContext, -) -> None: - with pytest.raises(ValueError, match="Dataset not found"): - rag_pipeline_service.service.get_pipeline("t1", "d1") - - -def test_get_pipeline_raises_when_pipeline_missing( - rag_pipeline_service: RagPipelineServiceTestContext, -) -> None: - _persist(rag_pipeline_service.session, _make_dataset()) - - with pytest.raises(ValueError, match="Pipeline not found"): - rag_pipeline_service.service.get_pipeline("t1", "d1") - - -def test_init_uses_default_sessionmaker_when_none(mocker: MockerFixture, sqlite_session: Session) -> None: - mocker.patch("services.rag_pipeline.rag_pipeline.db", SimpleNamespace(engine=sqlite_session.get_bind())) - create_exec_repo = mocker.patch( - "services.rag_pipeline.rag_pipeline.DifyAPIRepositoryFactory.create_api_workflow_node_execution_repository" - ) - create_run_repo = mocker.patch( - "services.rag_pipeline.rag_pipeline.DifyAPIRepositoryFactory.create_api_workflow_run_repository" - ) - - RagPipelineService(session=sqlite_session, session_maker=None) - - exec_session_maker = create_exec_repo.call_args.args[0] - run_session_maker = create_run_repo.call_args.args[0] - assert exec_session_maker is run_session_maker - assert exec_session_maker.kw["bind"] is sqlite_session.get_bind() - assert exec_session_maker.kw["expire_on_commit"] is False - - -def test_get_pipeline_templates_builtin_en_us_no_fallback( - mocker: MockerFixture, sqlite_session: Session, config_overrides: Callable[..., None] -) -> None: - config_overrides(HOSTED_FETCH_PIPELINE_TEMPLATES_MODE="remote") - session = sqlite_session - retrieval = mocker.Mock() - retrieval.get_pipeline_templates.return_value = {"pipeline_templates": []} - factory = mocker.patch("services.rag_pipeline.rag_pipeline.PipelineTemplateRetrievalFactory") - factory.get_pipeline_template_factory.return_value.return_value = retrieval - builtin = factory.get_built_in_pipeline_template_retrieval.return_value - result = RagPipelineService.get_pipeline_templates(type="built-in", language="en-US", session=session) - - assert result == {"pipeline_templates": []} - retrieval.get_pipeline_templates.assert_called_once_with("en-US", None, session=session) - builtin.fetch_pipeline_templates_from_builtin.assert_not_called() - - -def test_update_customized_pipeline_template_commits_when_name_empty(sqlite_session: Session) -> None: - template = _make_customized_template() - template.id = "tpl-1" - _persist(sqlite_session, template) - - info = PipelineTemplateInfoEntity(name="", description="updated", icon_info=IconInfo(icon="i")) - result = RagPipelineService.update_customized_pipeline_template( - "tpl-1", info, _make_account(), "t1", session=sqlite_session - ) - - assert result.description == "updated" - sqlite_session.expire_all() - updated_template = sqlite_session.get(PipelineCustomizedTemplate, "tpl-1") - assert updated_template is not None - assert updated_template.description == "updated" - - -def test_get_all_published_workflow_without_filters_has_no_more( - rag_pipeline_service: RagPipelineServiceTestContext, -) -> None: - pipeline = _make_pipeline(workflow_id="wf-live") - workflow = _make_workflow(workflow_id="wf1") - workflow.version = "2024-01-01T00:00:00" - _persist(rag_pipeline_service.session, workflow) - - workflows, has_more = rag_pipeline_service.service.get_all_published_workflow( - session=rag_pipeline_service.session, - pipeline=pipeline, - page=1, - limit=2, - user_id=None, - named_only=False, - ) - - assert workflows == [workflow] - assert has_more is False - - -def test_publish_workflow_skips_dataset_update_for_non_knowledge_nodes( - mocker: MockerFixture, rag_pipeline_service: RagPipelineServiceTestContext -) -> None: - draft = _make_workflow(graph={"nodes": [{"data": {"type": "start"}}]}) - _persist(rag_pipeline_service.session, draft) - dataset_service = mocker.patch("services.knowledge.dataset_service.DatasetService") - - result = rag_pipeline_service.service.publish_workflow( - session=rag_pipeline_service.session, - pipeline=_make_pipeline(), - account=_make_account(), - ) - - assert result.graph_dict == {"nodes": [{"data": {"type": "start"}}]} - dataset_service.update_rag_pipeline_dataset_settings.assert_not_called() - - -def test_get_default_block_config_returns_none_when_default_empty( - mocker: MockerFixture, rag_pipeline_service: RagPipelineServiceTestContext -) -> None: - node_cls = mocker.Mock() - node_cls.get_default_config.return_value = None - mocker.patch( - "services.rag_pipeline.rag_pipeline.get_node_type_classes_mapping", - return_value={BuiltinNodeTypes.START: {"1": node_cls}}, - ) - mocker.patch("services.rag_pipeline.rag_pipeline.LATEST_VERSION", "1") - - assert rag_pipeline_service.service.get_default_block_config("start") is None - - -def test_run_datasource_workflow_node_handles_variable_parameter_types( - mocker: MockerFixture, rag_pipeline_service: RagPipelineServiceTestContext -) -> None: - from core.datasource.entities.datasource_entities import DatasourceProviderType - - workflow = _make_workflow( - graph={ - "nodes": [ - { - "id": "node-1", - "data": { - "plugin_id": "p", - "provider_name": "provider", - "datasource_name": "crawl", - "datasource_parameters": { - "a": {"value": None}, - "b": {"value": "literal"}, - "c": {"value": ["input", "k"]}, - }, - }, - } - ] - } - ) - mocker.patch.object(rag_pipeline_service.service, "get_published_workflow", return_value=workflow) - runtime = mocker.Mock() - - def crawl_gen(**kwargs): - yield SimpleNamespace(result=SimpleNamespace(status="completed", total=1, completed=1, web_info_list=[])) - - runtime.get_website_crawl.side_effect = crawl_gen - runtime.datasource_provider_type.return_value = DatasourceProviderType.WEBSITE_CRAWL - mocker.patch("core.datasource.datasource_manager.DatasourceManager.get_datasource_runtime", return_value=runtime) - mocker.patch( - "services.rag_pipeline.rag_pipeline.DatasourceProviderService.get_datasource_credentials", - return_value=None, - ) - - events = list( - rag_pipeline_service.service.run_datasource_workflow_node( - pipeline=_make_pipeline(), - node_id="node-1", - user_inputs={"k": "mapped"}, - account=_make_account(), - datasource_type="website_crawl", - is_published=True, - datasource_providers=rag_pipeline_module.DatasourceProviderService( - credentials=create_autospec(DatasourceProviderCredentialStore, instance=True) - ), - ) - ) - - assert events - assert events[0]["data"] == [] - - -def test_run_datasource_workflow_node_online_drive_branch( - mocker: MockerFixture, rag_pipeline_service: RagPipelineServiceTestContext -) -> None: - from core.datasource.entities.datasource_entities import DatasourceProviderType - - workflow = _make_workflow( - graph={ - "nodes": [ - { - "id": "node-1", - "data": { - "plugin_id": "p", - "provider_name": "provider", - "datasource_name": "drive", - "datasource_parameters": {}, - }, - } - ] - } - ) - mocker.patch.object(rag_pipeline_service.service, "get_published_workflow", return_value=workflow) - runtime = mocker.Mock() - - def drive_gen(**kwargs): - yield SimpleNamespace(result={"items": [1]}) - - runtime.online_drive_browse_files.side_effect = drive_gen - runtime.datasource_provider_type.return_value = DatasourceProviderType.ONLINE_DRIVE - mocker.patch("core.datasource.datasource_manager.DatasourceManager.get_datasource_runtime", return_value=runtime) - mocker.patch( - "services.rag_pipeline.rag_pipeline.DatasourceProviderService.get_datasource_credentials", - return_value=None, - ) - - events = list( - rag_pipeline_service.service.run_datasource_workflow_node( - pipeline=_make_pipeline(), - node_id="node-1", - user_inputs={}, - account=_make_account(), - datasource_type="online_drive", - is_published=True, - datasource_providers=rag_pipeline_module.DatasourceProviderService( - credentials=create_autospec(DatasourceProviderCredentialStore, instance=True) - ), - ) - ) - - assert len(events) == 2 - assert events[1]["data"] == {"items": [1]} - - -def test_run_datasource_node_preview_not_published_uses_draft( - mocker: MockerFixture, rag_pipeline_service: RagPipelineServiceTestContext -) -> None: - from core.datasource.entities.datasource_entities import DatasourceMessage - - workflow = _make_workflow( - graph={ - "nodes": [ - { - "id": "n1", - "data": { - "plugin_id": "p", - "provider_name": "provider", - "datasource_name": "doc", - "datasource_parameters": {"workspace_id": {"value": "w"}}, - }, - } - ] - } - ) - get_draft = mocker.patch.object(rag_pipeline_service.service, "get_draft_workflow", return_value=workflow) - runtime = mocker.Mock() - - def doc_gen(**kwargs): - yield DatasourceMessage( - type=DatasourceMessage.MessageType.VARIABLE, - message=DatasourceMessage.VariableMessage(variable_name="x", variable_value="v", stream=False), - ) - - runtime.get_online_document_page_content.side_effect = doc_gen - mocker.patch("core.datasource.datasource_manager.DatasourceManager.get_datasource_runtime", return_value=runtime) - mocker.patch( - "services.rag_pipeline.rag_pipeline.DatasourceProviderService.get_datasource_credentials", - return_value=None, - ) - - result = rag_pipeline_service.service.run_datasource_node_preview( - pipeline=_make_pipeline(), - node_id="n1", - user_inputs={}, - account=_make_account(), - datasource_type="online_document", - is_published=False, - datasource_providers=rag_pipeline_module.DatasourceProviderService( - credentials=create_autospec(DatasourceProviderCredentialStore, instance=True) - ), - ) - - assert result == {"x": "v"} - get_draft.assert_called_once() - - -def test_run_free_workflow_node_delegates_to_handle_result( - mocker: MockerFixture, rag_pipeline_service: RagPipelineServiceTestContext -) -> None: - expected = SimpleNamespace(id="exec-1") - handle = mocker.patch.object(rag_pipeline_service.service, "_handle_node_run_result", return_value=expected) - - result = rag_pipeline_service.service.run_free_workflow_node( - node_data={"type": "start"}, - tenant_id="t1", - user_id="u1", - node_id="n1", - user_inputs={}, - ) - - assert result is expected - handle.assert_called_once() - - -@pytest.mark.parametrize(("workflow_tenant_id", "workflow_app_id"), [("t2", "p1"), ("t1", "p2")]) -def test_publish_customized_pipeline_template_rejects_unowned_workflow_before_export( - mocker: MockerFixture, - rag_pipeline_service: RagPipelineServiceTestContext, - workflow_tenant_id: str, - workflow_app_id: str, -) -> None: - pipeline = _make_pipeline(workflow_id="wf-1") - workflow = _make_workflow(workflow_id="wf-1", tenant_id=workflow_tenant_id, app_id=workflow_app_id) - _persist(rag_pipeline_service.session, pipeline, workflow) - dsl_service = mocker.patch.object(rag_pipeline_module, "RagPipelineDslService") - - with pytest.raises(RagPipelineResourceNotFoundError, match="Workflow not found"): - rag_pipeline_service.service.publish_customized_pipeline_template( - pipeline, - _make_dataset(), - _make_template_args(), - _make_account(), - session=rag_pipeline_service.session, - ) - - dsl_service.assert_not_called() - assert rag_pipeline_service.session.query(PipelineCustomizedTemplate).count() == 0 - - -def test_pipeline_retrieve_dataset_rejects_unowned_dataset( - rag_pipeline_service: RagPipelineServiceTestContext, -) -> None: - pipeline = _make_pipeline(workflow_id="wf-1") - other_tenant_dataset = _make_dataset(tenant_id="t2") - _persist(rag_pipeline_service.session, pipeline, other_tenant_dataset) - - assert get_pipeline_dataset(pipeline, session=rag_pipeline_service.session) is None - - -@pytest.mark.parametrize( - ("draft_tenant_id", "draft_app_id"), - [(None, None), ("t2", "p1"), ("t1", "p2")], -) -def test_publish_customized_pipeline_template_rejects_missing_or_unowned_draft_before_side_effects( - mocker: MockerFixture, - rag_pipeline_service: RagPipelineServiceTestContext, - draft_tenant_id: str | None, - draft_app_id: str | None, -) -> None: - session = rag_pipeline_service.session - pipeline = _make_pipeline(workflow_id="wf-published") - published_workflow = _make_workflow(workflow_id="wf-published", version="published") - dataset = _make_dataset() - resources = [pipeline, published_workflow, dataset] - if draft_tenant_id and draft_app_id: - resources.append(_make_workflow(workflow_id="wf-draft", tenant_id=draft_tenant_id, app_id=draft_app_id)) - _persist(session, *resources) - dsl_service = mocker.patch.object(rag_pipeline_module, "RagPipelineDslService") - - with pytest.raises(RagPipelineResourceNotFoundError, match="Draft workflow not found"): - rag_pipeline_service.service.publish_customized_pipeline_template( - pipeline, dataset, _make_template_args(), _make_account(), session=session - ) - - dsl_service.assert_not_called() - assert session.query(PipelineCustomizedTemplate).count() == 0 - - -def test_get_recommended_plugins_skips_manifest_when_missing( - mocker: MockerFixture, rag_pipeline_service: RagPipelineServiceTestContext -) -> None: - plugin = _make_recommended_plugin("plugin-a") - _persist(rag_pipeline_service.session, plugin) - mocker.patch("services.rag_pipeline.rag_pipeline.BuiltinToolManageService.list_builtin_tools", return_value=[]) - mocker.patch("services.rag_pipeline.rag_pipeline.marketplace.batch_fetch_plugin_by_ids", return_value=[]) - - result = rag_pipeline_service.service.get_recommended_plugins("all", _make_account(), "t1") - - assert result["installed_recommended_plugins"] == [] - assert result["uninstalled_recommended_plugins"] == [] - - -def test_retry_error_document_raises_when_pipeline_missing( - rag_pipeline_service: RagPipelineServiceTestContext, - pipeline_generator, -) -> None: - exec_log = DocumentPipelineExecutionLog( - pipeline_id="p1", - document_id="doc-1", - datasource_type="local_file", - datasource_info="{}", - datasource_node_id="node-1", - input_data={}, - created_by="u1", - ) - _persist(rag_pipeline_service.session, exec_log) - - with pytest.raises(ValueError, match="Pipeline not found"): - rag_pipeline_service.service.retry_error_document( - _make_dataset(), _make_document(), _make_account(), generator=pipeline_generator - ) - - -def test_retry_error_document_raises_when_workflow_missing( - rag_pipeline_service: RagPipelineServiceTestContext, - pipeline_generator, -) -> None: - exec_log = DocumentPipelineExecutionLog( - pipeline_id="p1", - document_id="doc-1", - datasource_type="local_file", - datasource_info="{}", - datasource_node_id="node-1", - input_data={}, - created_by="u1", - ) - pipeline = _make_pipeline(workflow_id="wf-1") - _persist(rag_pipeline_service.session, exec_log, pipeline) - - with pytest.raises(ValueError, match="Workflow not found"): - rag_pipeline_service.service.retry_error_document( - _make_dataset(), _make_document(), _make_account(), generator=pipeline_generator - ) - - -def test_get_datasource_plugins_returns_empty_for_non_datasource_nodes( - rag_pipeline_service: RagPipelineServiceTestContext, -) -> None: - dataset = _make_dataset() - pipeline = _make_pipeline(workflow_id="wf-1") - workflow = _make_workflow(graph={"nodes": [{"id": "n1", "data": {"type": "start"}}]}) - _persist(rag_pipeline_service.session, dataset, pipeline, workflow) - - assert ( - rag_pipeline_service.service.get_datasource_plugins( - "t1", - "d1", - True, - credential_query=mock.create_autospec(CredentialQuery, instance=True), - datasource_providers=rag_pipeline_module.DatasourceProviderService( - credentials=create_autospec(DatasourceProviderCredentialStore, instance=True) - ), - ) - == [] - ) - - -def test_publish_workflow_raises_when_knowledge_index_dataset_missing( - mocker: MockerFixture, rag_pipeline_service: RagPipelineServiceTestContext -) -> None: - draft = _make_workflow(graph={"nodes": [{"data": {"type": "knowledge-index"}}]}) - _persist(rag_pipeline_service.session, draft) - mocker.patch("services.rag_pipeline.rag_pipeline.KnowledgeConfiguration.model_validate", return_value=mocker.Mock()) - pipeline = _make_pipeline() - - with pytest.raises(ValueError, match="Dataset not found"): - rag_pipeline_service.service.publish_workflow( - session=rag_pipeline_service.session, pipeline=pipeline, account=_make_account() - ) - - -def test_run_datasource_node_preview_raises_when_workflow_missing( - mocker: MockerFixture, rag_pipeline_service: RagPipelineServiceTestContext -) -> None: - mocker.patch.object(rag_pipeline_service.service, "get_published_workflow", return_value=None) - - with pytest.raises(RuntimeError, match="Workflow not initialized"): - rag_pipeline_service.service.run_datasource_node_preview( - pipeline=_make_pipeline(), - node_id="n1", - user_inputs={}, - account=_make_account(), - datasource_type="online_document", - is_published=True, - datasource_providers=rag_pipeline_module.DatasourceProviderService( - credentials=create_autospec(DatasourceProviderCredentialStore, instance=True) - ), - ) - - -def test_run_datasource_node_preview_raises_when_node_missing( - mocker: MockerFixture, rag_pipeline_service: RagPipelineServiceTestContext -) -> None: - mocker.patch.object( - rag_pipeline_service.service, "get_published_workflow", return_value=_make_workflow(graph={"nodes": []}) - ) - - with pytest.raises(RuntimeError, match="Datasource node data not found"): - rag_pipeline_service.service.run_datasource_node_preview( - pipeline=_make_pipeline(), - node_id="missing", - user_inputs={}, - account=_make_account(), - datasource_type="online_document", - is_published=True, - datasource_providers=rag_pipeline_module.DatasourceProviderService( - credentials=create_autospec(DatasourceProviderCredentialStore, instance=True) - ), - ) - - -def test_run_datasource_node_preview_keeps_existing_user_input( - mocker: MockerFixture, rag_pipeline_service: RagPipelineServiceTestContext -) -> None: - from core.datasource.entities.datasource_entities import DatasourceMessage - - workflow = _make_workflow( - graph={ - "nodes": [ - { - "id": "n1", - "data": { - "plugin_id": "p", - "provider_name": "provider", - "datasource_name": "doc", - "datasource_parameters": {"workspace_id": {"value": "default"}}, - }, - } - ] - } - ) - mocker.patch.object(rag_pipeline_service.service, "get_published_workflow", return_value=workflow) - runtime = mocker.Mock() - - def gen(**kwargs): - request = kwargs["datasource_parameters"] - assert request.workspace_id == "existing" - yield DatasourceMessage( - type=DatasourceMessage.MessageType.VARIABLE, - message=DatasourceMessage.VariableMessage(variable_name="ok", variable_value="1", stream=False), - ) - - runtime.get_online_document_page_content.side_effect = gen - mocker.patch("core.datasource.datasource_manager.DatasourceManager.get_datasource_runtime", return_value=runtime) - mocker.patch( - "services.rag_pipeline.rag_pipeline.DatasourceProviderService.get_datasource_credentials", - return_value=None, - ) - - result = rag_pipeline_service.service.run_datasource_node_preview( - pipeline=_make_pipeline(), - node_id="n1", - user_inputs={"workspace_id": "existing"}, - account=_make_account(), - datasource_type="online_document", - is_published=True, - datasource_providers=rag_pipeline_module.DatasourceProviderService( - credentials=create_autospec(DatasourceProviderCredentialStore, instance=True) - ), - ) - assert result == {"ok": "1"} - - -def test_run_datasource_node_preview_ignores_non_variable_messages( - mocker: MockerFixture, rag_pipeline_service: RagPipelineServiceTestContext -) -> None: - workflow = _make_workflow( - graph={ - "nodes": [ - { - "id": "n1", - "data": { - "plugin_id": "p", - "provider_name": "provider", - "datasource_name": "doc", - "datasource_parameters": {}, - }, - } - ] - } - ) - mocker.patch.object(rag_pipeline_service.service, "get_published_workflow", return_value=workflow) - runtime = mocker.Mock() - - def gen(**kwargs): - yield SimpleNamespace(type="log", message=None) - - runtime.get_online_document_page_content.side_effect = gen - mocker.patch("core.datasource.datasource_manager.DatasourceManager.get_datasource_runtime", return_value=runtime) - mocker.patch( - "services.rag_pipeline.rag_pipeline.DatasourceProviderService.get_datasource_credentials", - return_value=None, - ) - - result = rag_pipeline_service.service.run_datasource_node_preview( - pipeline=_make_pipeline(), - node_id="n1", - user_inputs={}, - account=_make_account(), - datasource_type="online_document", - is_published=True, - datasource_providers=rag_pipeline_module.DatasourceProviderService( - credentials=create_autospec(DatasourceProviderCredentialStore, instance=True) - ), - ) - assert result == {} - - -def test_set_datasource_variables_raises_when_workflow_missing( - mocker: MockerFixture, rag_pipeline_service: RagPipelineServiceTestContext -) -> None: - mocker.patch.object(rag_pipeline_service.service, "get_draft_workflow", return_value=None) - - with pytest.raises(ValueError, match="Workflow not initialized"): - rag_pipeline_service.service.set_datasource_variables( - _make_pipeline(), - {"start_node_id": "n1"}, - _make_account(), - ) - - -def test_get_datasource_plugins_handles_empty_datasource_data_and_non_published( - mocker: MockerFixture, rag_pipeline_service: RagPipelineServiceTestContext -) -> None: - dataset = _make_dataset() - pipeline = _make_pipeline() - workflow = _make_workflow( - graph={"nodes": [{"id": "n1", "data": {"type": "datasource", "datasource_parameters": {}}}]}, - ) - workflow.rag_pipeline_variables = [ - {"variable": "v1", "belong_to_node_id": "shared", "type": "text-input", "label": "V1"} - ] - _persist(rag_pipeline_service.session, dataset, pipeline, workflow) - mocker.patch( - "services.rag_pipeline.rag_pipeline.DatasourceProviderService.list_datasource_credentials", - return_value=[], - ) - - result = rag_pipeline_service.service.get_datasource_plugins( - "t1", - "d1", - False, - credential_query=mock.create_autospec(CredentialQuery, instance=True), - datasource_providers=rag_pipeline_module.DatasourceProviderService( - credentials=create_autospec(DatasourceProviderCredentialStore, instance=True) - ), - ) - - assert len(result) == 1 - - -def test_get_datasource_plugins_extracts_user_inputs_and_credentials( - mocker: MockerFixture, rag_pipeline_service: RagPipelineServiceTestContext -) -> None: - dataset = _make_dataset() - pipeline = _make_pipeline(workflow_id="wf-1") - workflow = _make_workflow( - graph={ - "nodes": [ - { - "id": "n1", - "data": { - "type": "datasource", - "plugin_id": "plugin-1", - "provider_name": "provider", - "provider_type": "online_document", - "title": "Datasource", - "datasource_parameters": { - "a": {"value": "{{#start.v1#}}"}, - "b": {"value": ["x", "v2"]}, - }, - }, - } - ] - }, - ) - workflow.rag_pipeline_variables = [ - {"variable": "v1", "belong_to_node_id": "shared", "type": "text-input", "label": "V1"}, - {"variable": "v2", "belong_to_node_id": "shared", "type": "text-input", "label": "V2"}, - {"variable": "v3", "belong_to_node_id": "shared", "type": "text-input", "label": "V3"}, - ] - _persist(rag_pipeline_service.session, dataset, pipeline, workflow) - mocker.patch( - "services.rag_pipeline.rag_pipeline.DatasourceProviderService.list_datasource_credentials", - return_value=[{"id": "c1", "name": "Cred", "type": "api", "is_default": True}], - ) - - result = rag_pipeline_service.service.get_datasource_plugins( - "t1", - "d1", - True, - credential_query=mock.create_autospec(CredentialQuery, instance=True), - datasource_providers=rag_pipeline_module.DatasourceProviderService( - credentials=create_autospec(DatasourceProviderCredentialStore, instance=True) - ), - ) - - assert len(result) == 1 - assert len(result[0]["user_input_variables"]) == 2 - assert result[0]["credentials"][0]["id"] == "c1" - - -def test_get_pipeline_returns_pipeline_when_found( - rag_pipeline_service: RagPipelineServiceTestContext, -) -> None: - dataset = _make_dataset() - pipeline = _make_pipeline() - _persist(rag_pipeline_service.session, dataset, pipeline) - - result = rag_pipeline_service.service.get_pipeline("t1", "d1") - - assert result is pipeline - - -def test_get_pipeline_by_id_uses_provided_session(sqlite_session: Session) -> None: - pipeline = _make_pipeline() - sqlite_session.add(pipeline) - sqlite_session.flush() - - result = RagPipelineService.get_pipeline_by_id("p1", "t1", session=sqlite_session) - - assert result is pipeline - assert RagPipelineService.get_pipeline_by_id("p1", "other-tenant", session=sqlite_session) is None diff --git a/api/tests/unit_tests/services/rag_pipeline/test_rag_pipeline_task_proxy.py b/api/tests/unit_tests/services/rag_pipeline/test_rag_pipeline_task_proxy.py index b9e5963c3f9..ae737cf7fff 100644 --- a/api/tests/unit_tests/services/rag_pipeline/test_rag_pipeline_task_proxy.py +++ b/api/tests/unit_tests/services/rag_pipeline/test_rag_pipeline_task_proxy.py @@ -1,175 +1,666 @@ +import json +import logging from datetime import datetime from types import SimpleNamespace -from unittest.mock import Mock +from unittest.mock import MagicMock, Mock, call, patch import pytest from pytest_mock import MockerFixture +from core.app.entities.rag_pipeline_invoke_entities import RagPipelineInvokeEntity +from core.rag.pipeline.queue import TenantIsolatedTaskQueue from enums import CloudPlan, DeploymentEdition from extensions.storage.storage_type import StorageType from models.enums import CreatorUserRole from models.model import UploadFile from services.rag_pipeline.rag_pipeline_task_proxy import RagPipelineTaskProxy -from tests.unit_tests.config_override import config_overrides_context +from tests.helpers.unit.config_override import config_overrides_context -@pytest.fixture -def proxy(mocker: MockerFixture): - """Create a RagPipelineTaskProxy with mocked dependencies.""" - mocker.patch("services.rag_pipeline.rag_pipeline_task_proxy.TenantIsolatedTaskQueue") - entity = Mock() - entity.model_dump.return_value = {"doc": "data"} - return RagPipelineTaskProxy( - dataset_tenant_id="tenant-1", - user_id="user-1", - rag_pipeline_invoke_entities=[entity], - ) +class TestRagPipelineTaskProxyCases: + @staticmethod + @pytest.fixture + def proxy(mocker: MockerFixture): + """Create a RagPipelineTaskProxy with mocked dependencies.""" + mocker.patch("services.rag_pipeline.rag_pipeline_task_proxy.TenantIsolatedTaskQueue") + entity = Mock() + entity.model_dump.return_value = {"doc": "data"} + return RagPipelineTaskProxy( + dataset_tenant_id="tenant-1", + user_id="user-1", + rag_pipeline_invoke_entities=[entity], + ) + # --- delay --- -# --- delay --- + @staticmethod + def test_delay_with_empty_entities_logs_warning_and_returns(mocker: MockerFixture) -> None: + mocker.patch("services.rag_pipeline.rag_pipeline_task_proxy.TenantIsolatedTaskQueue") + proxy = RagPipelineTaskProxy( + dataset_tenant_id="tenant-1", + user_id="user-1", + rag_pipeline_invoke_entities=[], + ) + dispatch_mock = mocker.patch.object(proxy, "_dispatch") + proxy.delay() -def test_delay_with_empty_entities_logs_warning_and_returns(mocker: MockerFixture) -> None: - mocker.patch("services.rag_pipeline.rag_pipeline_task_proxy.TenantIsolatedTaskQueue") - proxy = RagPipelineTaskProxy( - dataset_tenant_id="tenant-1", - user_id="user-1", - rag_pipeline_invoke_entities=[], - ) - dispatch_mock = mocker.patch.object(proxy, "_dispatch") + dispatch_mock.assert_not_called() - proxy.delay() + @staticmethod + def test_delay_with_entities_calls_dispatch(mocker: MockerFixture, proxy) -> None: + dispatch_mock = mocker.patch.object(proxy, "_dispatch") - dispatch_mock.assert_not_called() + proxy.delay() + dispatch_mock.assert_called_once() -def test_delay_with_entities_calls_dispatch(mocker: MockerFixture, proxy) -> None: - dispatch_mock = mocker.patch.object(proxy, "_dispatch") + # --- _dispatch --- - proxy.delay() + @staticmethod + @config_overrides_context(DEPLOYMENT_EDITION=DeploymentEdition.CLOUD) + def test_dispatch_billing_sandbox_uses_default_tenant_queue(mocker: MockerFixture, proxy) -> None: + upload_mock = mocker.patch.object(proxy, "_upload_invoke_entities", return_value="file-1") + send_mock = mocker.patch.object(proxy, "_send_to_default_tenant_queue") - dispatch_mock.assert_called_once() + features = SimpleNamespace(billing=SimpleNamespace(subscription=SimpleNamespace(plan=CloudPlan.SANDBOX))) + mocker.patch.object(type(proxy), "features", new_callable=lambda: property(lambda self: features)) - -# --- _dispatch --- - - -@config_overrides_context(DEPLOYMENT_EDITION=DeploymentEdition.CLOUD) -def test_dispatch_billing_sandbox_uses_default_tenant_queue(mocker: MockerFixture, proxy) -> None: - upload_mock = mocker.patch.object(proxy, "_upload_invoke_entities", return_value="file-1") - send_mock = mocker.patch.object(proxy, "_send_to_default_tenant_queue") - - features = SimpleNamespace(billing=SimpleNamespace(subscription=SimpleNamespace(plan=CloudPlan.SANDBOX))) - mocker.patch.object(type(proxy), "features", new_callable=lambda: property(lambda self: features)) - - proxy._dispatch() - - upload_mock.assert_called_once() - send_mock.assert_called_once_with("file-1") - - -@config_overrides_context(DEPLOYMENT_EDITION=DeploymentEdition.CLOUD) -def test_dispatch_billing_non_sandbox_uses_priority_tenant_queue(mocker: MockerFixture, proxy) -> None: - upload_mock = mocker.patch.object(proxy, "_upload_invoke_entities", return_value="file-1") - send_mock = mocker.patch.object(proxy, "_send_to_priority_tenant_queue") - - features = SimpleNamespace(billing=SimpleNamespace(subscription=SimpleNamespace(plan=CloudPlan.PROFESSIONAL))) - mocker.patch.object(type(proxy), "features", new_callable=lambda: property(lambda self: features)) - - proxy._dispatch() - - upload_mock.assert_called_once() - send_mock.assert_called_once_with("file-1") - - -@config_overrides_context(DEPLOYMENT_EDITION=DeploymentEdition.COMMUNITY) -def test_dispatch_no_billing_uses_priority_direct_queue(mocker: MockerFixture, proxy) -> None: - upload_mock = mocker.patch.object(proxy, "_upload_invoke_entities", return_value="file-1") - send_mock = mocker.patch.object(proxy, "_send_to_priority_direct_queue") - - features = SimpleNamespace(billing=SimpleNamespace(subscription=SimpleNamespace(plan="free"))) - mocker.patch.object(type(proxy), "features", new_callable=lambda: property(lambda self: features)) - - proxy._dispatch() - - upload_mock.assert_called_once() - send_mock.assert_called_once_with("file-1") - - -@config_overrides_context(DEPLOYMENT_EDITION=DeploymentEdition.COMMUNITY) -def test_dispatch_raises_on_empty_upload_file_id(mocker: MockerFixture, proxy) -> None: - mocker.patch.object(proxy, "_upload_invoke_entities", return_value="") - - features = SimpleNamespace(billing=SimpleNamespace(subscription=SimpleNamespace(plan="free"))) - mocker.patch.object(type(proxy), "features", new_callable=lambda: property(lambda self: features)) - - with pytest.raises(ValueError, match="upload_file_id is empty"): proxy._dispatch() + upload_mock.assert_called_once() + send_mock.assert_called_once_with("file-1") -# --- _send_to_direct_queue --- + @staticmethod + @config_overrides_context(DEPLOYMENT_EDITION=DeploymentEdition.CLOUD) + def test_dispatch_billing_non_sandbox_uses_priority_tenant_queue(mocker: MockerFixture, proxy) -> None: + upload_mock = mocker.patch.object(proxy, "_upload_invoke_entities", return_value="file-1") + send_mock = mocker.patch.object(proxy, "_send_to_priority_tenant_queue") + + features = SimpleNamespace(billing=SimpleNamespace(subscription=SimpleNamespace(plan=CloudPlan.PROFESSIONAL))) + mocker.patch.object(type(proxy), "features", new_callable=lambda: property(lambda self: features)) + + proxy._dispatch() + + upload_mock.assert_called_once() + send_mock.assert_called_once_with("file-1") + + @staticmethod + @config_overrides_context(DEPLOYMENT_EDITION=DeploymentEdition.COMMUNITY) + def test_dispatch_no_billing_uses_priority_direct_queue(mocker: MockerFixture, proxy) -> None: + upload_mock = mocker.patch.object(proxy, "_upload_invoke_entities", return_value="file-1") + send_mock = mocker.patch.object(proxy, "_send_to_priority_direct_queue") + + features = SimpleNamespace(billing=SimpleNamespace(subscription=SimpleNamespace(plan="free"))) + mocker.patch.object(type(proxy), "features", new_callable=lambda: property(lambda self: features)) + + proxy._dispatch() + + upload_mock.assert_called_once() + send_mock.assert_called_once_with("file-1") + + @staticmethod + @config_overrides_context(DEPLOYMENT_EDITION=DeploymentEdition.COMMUNITY) + def test_dispatch_raises_on_empty_upload_file_id(mocker: MockerFixture, proxy) -> None: + mocker.patch.object(proxy, "_upload_invoke_entities", return_value="") + + features = SimpleNamespace(billing=SimpleNamespace(subscription=SimpleNamespace(plan="free"))) + mocker.patch.object(type(proxy), "features", new_callable=lambda: property(lambda self: features)) + + with pytest.raises(ValueError, match="upload_file_id is empty"): + proxy._dispatch() + + # --- _send_to_direct_queue --- + + @staticmethod + def test_send_to_direct_queue_calls_task_func_delay(mocker: MockerFixture, proxy) -> None: + task_func = Mock() + + proxy._send_to_direct_queue("file-1", task_func) + + task_func.delay.assert_called_once_with( + rag_pipeline_invoke_entities_file_id="file-1", + tenant_id="tenant-1", + ) + + # --- _send_to_tenant_queue --- + + @staticmethod + def test_send_to_tenant_queue_pushes_when_task_key_exists(mocker: MockerFixture, proxy) -> None: + proxy._tenant_isolated_task_queue.get_task_key.return_value = "existing-key" + task_func = Mock() + + proxy._send_to_tenant_queue("file-1", task_func) + + proxy._tenant_isolated_task_queue.push_tasks.assert_called_once_with(["file-1"]) + task_func.delay.assert_not_called() + + @staticmethod + def test_send_to_tenant_queue_sets_waiting_time_and_calls_delay(mocker: MockerFixture, proxy) -> None: + proxy._tenant_isolated_task_queue.get_task_key.return_value = None + task_func = Mock() + + proxy._send_to_tenant_queue("file-1", task_func) + + proxy._tenant_isolated_task_queue.set_task_waiting_time.assert_called_once() + task_func.delay.assert_called_once_with( + rag_pipeline_invoke_entities_file_id="file-1", + tenant_id="tenant-1", + ) + + # --- _upload_invoke_entities --- + + @staticmethod + def test_upload_invoke_entities_returns_file_id(mocker: MockerFixture, proxy) -> None: + upload_file = UploadFile( + tenant_id="tenant-1", + storage_type=StorageType.LOCAL, + key="rag-pipeline.json", + name="rag-pipeline.json", + size=1, + extension="json", + mime_type="application/json", + created_by_role=CreatorUserRole.ACCOUNT, + created_by="user-1", + created_at=datetime(2025, 1, 1), + used=True, + ) + upload_file.id = "uploaded-file-1" + file_service_cls = mocker.patch("services.rag_pipeline.rag_pipeline_task_proxy.FileService") + file_service_cls.return_value.upload_text.return_value = upload_file + mocker.patch("services.rag_pipeline.rag_pipeline_task_proxy.db", SimpleNamespace(engine="fake-engine")) + + result = proxy._upload_invoke_entities() + + assert result == "uploaded-file-1" + file_service_cls.return_value.upload_text.assert_called_once() -def test_send_to_direct_queue_calls_task_func_delay(mocker: MockerFixture, proxy) -> None: - task_func = Mock() +class RagPipelineTaskProxyTestDataFactory: + """Factory class for creating test data and mock objects for RagPipelineTaskProxy tests.""" - proxy._send_to_direct_queue("file-1", task_func) + @staticmethod + def create_mock_features(plan: CloudPlan = CloudPlan.SANDBOX) -> Mock: + """Create mock features with billing configuration.""" + features = Mock() + features.billing = Mock() + features.billing.subscription = Mock() + features.billing.subscription.plan = plan + return features - task_func.delay.assert_called_once_with( - rag_pipeline_invoke_entities_file_id="file-1", - tenant_id="tenant-1", - ) + @staticmethod + def create_rag_pipeline_invoke_entity( + pipeline_id: str = "pipeline-123", + user_id: str = "user-456", + tenant_id: str = "tenant-789", + workflow_id: str = "workflow-101", + streaming: bool = True, + workflow_execution_id: str | None = None, + workflow_thread_pool_id: str | None = None, + ) -> RagPipelineInvokeEntity: + """Create RagPipelineInvokeEntity instance for testing.""" + return RagPipelineInvokeEntity( + pipeline_id=pipeline_id, + application_generate_entity={"key": "value"}, + user_id=user_id, + tenant_id=tenant_id, + workflow_id=workflow_id, + streaming=streaming, + workflow_execution_id=workflow_execution_id, + workflow_thread_pool_id=workflow_thread_pool_id, + ) + + @staticmethod + def create_rag_pipeline_task_proxy( + dataset_tenant_id: str = "tenant-123", + user_id: str = "user-456", + rag_pipeline_invoke_entities: list[RagPipelineInvokeEntity] | None = None, + ) -> RagPipelineTaskProxy: + """Create RagPipelineTaskProxy instance for testing.""" + if rag_pipeline_invoke_entities is None: + rag_pipeline_invoke_entities = [RagPipelineTaskProxyTestDataFactory.create_rag_pipeline_invoke_entity()] + return RagPipelineTaskProxy(dataset_tenant_id, user_id, rag_pipeline_invoke_entities) + + @staticmethod + def create_upload_file(file_id: str = "file-123") -> UploadFile: + """Create a transient upload file returned by the mocked storage service.""" + upload_file = UploadFile( + tenant_id="tenant-123", + storage_type=StorageType.LOCAL, + key="rag-pipeline.json", + name="rag-pipeline.json", + size=1, + extension="json", + mime_type="application/json", + created_by_role=CreatorUserRole.ACCOUNT, + created_by="user-456", + created_at=datetime(2025, 1, 1), + used=True, + ) + upload_file.id = file_id + return upload_file -# --- _send_to_tenant_queue --- +class TestRagPipelineTaskProxyCasesServices: + class TestRagPipelineTaskProxy: + """Test cases for RagPipelineTaskProxy class.""" + def test_initialization(self): + """Test RagPipelineTaskProxy initialization.""" + # Arrange + dataset_tenant_id = "tenant-123" + user_id = "user-456" + rag_pipeline_invoke_entities = [RagPipelineTaskProxyTestDataFactory.create_rag_pipeline_invoke_entity()] -def test_send_to_tenant_queue_pushes_when_task_key_exists(mocker: MockerFixture, proxy) -> None: - proxy._tenant_isolated_task_queue.get_task_key.return_value = "existing-key" - task_func = Mock() + # Act + proxy = RagPipelineTaskProxy(dataset_tenant_id, user_id, rag_pipeline_invoke_entities) - proxy._send_to_tenant_queue("file-1", task_func) + # Assert + assert proxy._dataset_tenant_id == dataset_tenant_id + assert proxy._user_id == user_id + assert proxy._rag_pipeline_invoke_entities == rag_pipeline_invoke_entities + assert isinstance(proxy._tenant_isolated_task_queue, TenantIsolatedTaskQueue) + assert proxy._tenant_isolated_task_queue._tenant_id == dataset_tenant_id + assert proxy._tenant_isolated_task_queue._unique_key == "pipeline" - proxy._tenant_isolated_task_queue.push_tasks.assert_called_once_with(["file-1"]) - task_func.delay.assert_not_called() + def test_initialization_with_empty_entities(self): + """Test initialization with empty rag_pipeline_invoke_entities.""" + # Arrange + dataset_tenant_id = "tenant-123" + user_id = "user-456" + rag_pipeline_invoke_entities = [] + # Act + proxy = RagPipelineTaskProxy(dataset_tenant_id, user_id, rag_pipeline_invoke_entities) -def test_send_to_tenant_queue_sets_waiting_time_and_calls_delay(mocker: MockerFixture, proxy) -> None: - proxy._tenant_isolated_task_queue.get_task_key.return_value = None - task_func = Mock() + # Assert + assert proxy._dataset_tenant_id == dataset_tenant_id + assert proxy._user_id == user_id + assert proxy._rag_pipeline_invoke_entities == [] - proxy._send_to_tenant_queue("file-1", task_func) + def test_initialization_with_multiple_entities(self): + """Test initialization with multiple rag_pipeline_invoke_entities.""" + # Arrange + dataset_tenant_id = "tenant-123" + user_id = "user-456" + rag_pipeline_invoke_entities = [ + RagPipelineTaskProxyTestDataFactory.create_rag_pipeline_invoke_entity(pipeline_id="pipeline-1"), + RagPipelineTaskProxyTestDataFactory.create_rag_pipeline_invoke_entity(pipeline_id="pipeline-2"), + RagPipelineTaskProxyTestDataFactory.create_rag_pipeline_invoke_entity(pipeline_id="pipeline-3"), + ] - proxy._tenant_isolated_task_queue.set_task_waiting_time.assert_called_once() - task_func.delay.assert_called_once_with( - rag_pipeline_invoke_entities_file_id="file-1", - tenant_id="tenant-1", - ) + # Act + proxy = RagPipelineTaskProxy(dataset_tenant_id, user_id, rag_pipeline_invoke_entities) + # Assert + assert len(proxy._rag_pipeline_invoke_entities) == 3 + assert proxy._rag_pipeline_invoke_entities[0].pipeline_id == "pipeline-1" + assert proxy._rag_pipeline_invoke_entities[1].pipeline_id == "pipeline-2" + assert proxy._rag_pipeline_invoke_entities[2].pipeline_id == "pipeline-3" -# --- _upload_invoke_entities --- + @patch("services.rag_pipeline.rag_pipeline_task_proxy.FeatureService") + def test_features_property(self, mock_feature_service: MagicMock): + """Test cached_property features.""" + # Arrange + mock_features = RagPipelineTaskProxyTestDataFactory.create_mock_features() + mock_feature_service.get_features.return_value = mock_features + proxy = RagPipelineTaskProxyTestDataFactory.create_rag_pipeline_task_proxy() + # Act + features1 = proxy.features + features2 = proxy.features # Second call should use cached property -def test_upload_invoke_entities_returns_file_id(mocker: MockerFixture, proxy) -> None: - upload_file = UploadFile( - tenant_id="tenant-1", - storage_type=StorageType.LOCAL, - key="rag-pipeline.json", - name="rag-pipeline.json", - size=1, - extension="json", - mime_type="application/json", - created_by_role=CreatorUserRole.ACCOUNT, - created_by="user-1", - created_at=datetime(2025, 1, 1), - used=True, - ) - upload_file.id = "uploaded-file-1" - file_service_cls = mocker.patch("services.rag_pipeline.rag_pipeline_task_proxy.FileService") - file_service_cls.return_value.upload_text.return_value = upload_file - mocker.patch("services.rag_pipeline.rag_pipeline_task_proxy.db", SimpleNamespace(engine="fake-engine")) + # Assert + assert features1 == mock_features + assert features2 == mock_features + assert features1 is features2 # Should be the same instance due to caching + mock_feature_service.get_features.assert_called_once_with("tenant-123", exclude_vector_space=True) - result = proxy._upload_invoke_entities() + @patch("services.rag_pipeline.rag_pipeline_task_proxy.FileService") + @patch("services.rag_pipeline.rag_pipeline_task_proxy.db") + def test_upload_invoke_entities(self, mock_db: MagicMock, mock_file_service_class: MagicMock): + """Test _upload_invoke_entities method.""" + # Arrange + proxy = RagPipelineTaskProxyTestDataFactory.create_rag_pipeline_task_proxy() + mock_file_service = Mock() + mock_file_service_class.return_value = mock_file_service + upload_file = RagPipelineTaskProxyTestDataFactory.create_upload_file("file-123") + mock_file_service.upload_text.return_value = upload_file - assert result == "uploaded-file-1" - file_service_cls.return_value.upload_text.assert_called_once() + # Act + result = proxy._upload_invoke_entities() + + # Assert + assert result == "file-123" + mock_file_service_class.assert_called_once_with(mock_db.engine) + + # Verify upload_text was called with correct parameters + mock_file_service.upload_text.assert_called_once() + call_args = mock_file_service.upload_text.call_args + json_text, name, user_id, tenant_id = call_args[0] + + assert name == "rag_pipeline_invoke_entities.json" + assert user_id == "user-456" + assert tenant_id == "tenant-123" + + # Verify JSON content + parsed_json = json.loads(json_text) + assert len(parsed_json) == 1 + assert parsed_json[0]["pipeline_id"] == "pipeline-123" + + @patch("services.rag_pipeline.rag_pipeline_task_proxy.FileService") + @patch("services.rag_pipeline.rag_pipeline_task_proxy.db") + def test_upload_invoke_entities_with_multiple_entities( + self, mock_db: MagicMock, mock_file_service_class: MagicMock + ): + """Test _upload_invoke_entities method with multiple entities.""" + # Arrange + entities = [ + RagPipelineTaskProxyTestDataFactory.create_rag_pipeline_invoke_entity(pipeline_id="pipeline-1"), + RagPipelineTaskProxyTestDataFactory.create_rag_pipeline_invoke_entity(pipeline_id="pipeline-2"), + ] + proxy = RagPipelineTaskProxy("tenant-123", "user-456", entities) + mock_file_service = Mock() + mock_file_service_class.return_value = mock_file_service + upload_file = RagPipelineTaskProxyTestDataFactory.create_upload_file("file-456") + mock_file_service.upload_text.return_value = upload_file + + # Act + result = proxy._upload_invoke_entities() + + # Assert + assert result == "file-456" + + # Verify JSON content contains both entities + call_args = mock_file_service.upload_text.call_args + json_text = call_args[0][0] + parsed_json = json.loads(json_text) + assert len(parsed_json) == 2 + assert parsed_json[0]["pipeline_id"] == "pipeline-1" + assert parsed_json[1]["pipeline_id"] == "pipeline-2" + + def test_send_to_direct_queue(self, tenant_queue_commands: MagicMock): + """Test _send_to_direct_queue method.""" + # Arrange + proxy = RagPipelineTaskProxyTestDataFactory.create_rag_pipeline_task_proxy() + upload_file_id = "file-123" + mock_task = Mock() + + # Act + proxy._send_to_direct_queue(upload_file_id, mock_task) + + # If sent to direct queue, tenant_isolated_task_queue should not be called + assert all(entry.args[0] != "LPUSH" for entry in tenant_queue_commands.call_args_list) + + # Celery should be called directly + mock_task.delay.assert_called_once_with( + rag_pipeline_invoke_entities_file_id=upload_file_id, tenant_id="tenant-123" + ) + + def test_send_to_tenant_queue_with_existing_task_key(self, tenant_queue_commands: MagicMock): + """Test _send_to_tenant_queue when task key exists.""" + # Arrange + proxy = RagPipelineTaskProxyTestDataFactory.create_rag_pipeline_task_proxy() + tenant_queue_commands.return_value = "task_key" + upload_file_id = "file-123" + mock_task = Mock() + + # Act + proxy._send_to_tenant_queue(upload_file_id, mock_task) + + # If task key exists, should push tasks to the queue + tenant_queue_commands.assert_any_call("LPUSH", "tenant_self_pipeline_task_queue:tenant-123", upload_file_id) + # Celery should not be called directly + mock_task.delay.assert_not_called() + + def test_send_to_tenant_queue_without_task_key(self, tenant_queue_commands: MagicMock): + """Test _send_to_tenant_queue when no task key exists.""" + # Arrange + proxy = RagPipelineTaskProxyTestDataFactory.create_rag_pipeline_task_proxy() + tenant_queue_commands.return_value = None + upload_file_id = "file-123" + mock_task = Mock() + + # Act + proxy._send_to_tenant_queue(upload_file_id, mock_task) + + # If no task key, should set task waiting time key first + tenant_queue_commands.assert_has_calls( + [ + call("GET", "tenant_pipeline_task:tenant-123", keys=["tenant_pipeline_task:tenant-123"]), + call("SETEX", "tenant_pipeline_task:tenant-123", 3600, 1), + ] + ) + mock_task.delay.assert_called_once_with( + rag_pipeline_invoke_entities_file_id=upload_file_id, tenant_id="tenant-123" + ) + + # The first task should be sent to celery directly, so push tasks should not be called + assert all(entry.args[0] != "LPUSH" for entry in tenant_queue_commands.call_args_list) + + def test_send_to_default_tenant_queue(self): + """Test _send_to_default_tenant_queue method.""" + # Arrange + proxy = RagPipelineTaskProxyTestDataFactory.create_rag_pipeline_task_proxy() + proxy._send_to_tenant_queue = Mock() + upload_file_id = "file-123" + + # Act + proxy._send_to_default_tenant_queue(upload_file_id) + + # Assert + proxy._send_to_tenant_queue.assert_called_once() + sent_file_id, task = proxy._send_to_tenant_queue.call_args.args + assert sent_file_id == upload_file_id + assert task.task == "tasks.rag_pipeline.rag_pipeline_run_task.rag_pipeline_run_task" + assert task.options["queue"] == "pipeline" + + def test_send_to_priority_tenant_queue(self): + """Test _send_to_priority_tenant_queue method.""" + # Arrange + proxy = RagPipelineTaskProxyTestDataFactory.create_rag_pipeline_task_proxy() + proxy._send_to_tenant_queue = Mock() + upload_file_id = "file-123" + + # Act + proxy._send_to_priority_tenant_queue(upload_file_id) + + # Assert + proxy._send_to_tenant_queue.assert_called_once() + sent_file_id, task = proxy._send_to_tenant_queue.call_args.args + assert sent_file_id == upload_file_id + assert task.task == "tasks.rag_pipeline.priority_rag_pipeline_run_task.priority_rag_pipeline_run_task" + assert task.options["queue"] == "priority_pipeline" + + def test_send_to_priority_direct_queue(self): + """Test _send_to_priority_direct_queue method.""" + # Arrange + proxy = RagPipelineTaskProxyTestDataFactory.create_rag_pipeline_task_proxy() + proxy._send_to_direct_queue = Mock() + upload_file_id = "file-123" + + # Act + proxy._send_to_priority_direct_queue(upload_file_id) + + # Assert + proxy._send_to_direct_queue.assert_called_once() + sent_file_id, task = proxy._send_to_direct_queue.call_args.args + assert sent_file_id == upload_file_id + assert task.task == "tasks.rag_pipeline.priority_rag_pipeline_run_task.priority_rag_pipeline_run_task" + assert task.options["queue"] == "priority_pipeline" + + @config_overrides_context(DEPLOYMENT_EDITION=DeploymentEdition.CLOUD) + @patch("services.rag_pipeline.rag_pipeline_task_proxy.FeatureService") + @patch("services.rag_pipeline.rag_pipeline_task_proxy.FileService") + @patch("services.rag_pipeline.rag_pipeline_task_proxy.db") + def test_dispatch_with_cloud_sandbox_plan( + self, mock_db: MagicMock, mock_file_service_class: MagicMock, mock_feature_service: MagicMock + ): + """Test _dispatch method in Cloud with Sandbox plan.""" + # Arrange + mock_features = RagPipelineTaskProxyTestDataFactory.create_mock_features(plan=CloudPlan.SANDBOX) + mock_feature_service.get_features.return_value = mock_features + proxy = RagPipelineTaskProxyTestDataFactory.create_rag_pipeline_task_proxy() + proxy._send_to_default_tenant_queue = Mock() + + mock_file_service = Mock() + mock_file_service_class.return_value = mock_file_service + mock_upload_file = RagPipelineTaskProxyTestDataFactory.create_upload_file("file-123") + mock_file_service.upload_text.return_value = mock_upload_file + + # Act + proxy._dispatch() + + # If billing is enabled with sandbox plan, should send to default tenant queue + proxy._send_to_default_tenant_queue.assert_called_once_with("file-123") + + @config_overrides_context(DEPLOYMENT_EDITION=DeploymentEdition.CLOUD) + @patch("services.rag_pipeline.rag_pipeline_task_proxy.FeatureService") + @patch("services.rag_pipeline.rag_pipeline_task_proxy.FileService") + @patch("services.rag_pipeline.rag_pipeline_task_proxy.db") + def test_dispatch_with_cloud_paid_plan(self, mock_db, mock_file_service_class, mock_feature_service): + """Test _dispatch method in Cloud with a paid plan.""" + # Arrange + mock_features = RagPipelineTaskProxyTestDataFactory.create_mock_features(plan=CloudPlan.TEAM) + mock_feature_service.get_features.return_value = mock_features + proxy = RagPipelineTaskProxyTestDataFactory.create_rag_pipeline_task_proxy() + proxy._send_to_priority_tenant_queue = Mock() + + mock_file_service = Mock() + mock_file_service_class.return_value = mock_file_service + mock_upload_file = RagPipelineTaskProxyTestDataFactory.create_upload_file("file-123") + mock_file_service.upload_text.return_value = mock_upload_file + + # Act + proxy._dispatch() + + # If billing is enabled with non-sandbox plan, should send to priority tenant queue + proxy._send_to_priority_tenant_queue.assert_called_once_with("file-123") + + @config_overrides_context(DEPLOYMENT_EDITION=DeploymentEdition.COMMUNITY) + @patch("services.rag_pipeline.rag_pipeline_task_proxy.FeatureService") + @patch("services.rag_pipeline.rag_pipeline_task_proxy.FileService") + @patch("services.rag_pipeline.rag_pipeline_task_proxy.db") + def test_dispatch_outside_cloud( + self, mock_db: MagicMock, mock_file_service_class: MagicMock, mock_feature_service: MagicMock + ): + """Test _dispatch method outside Cloud.""" + # Arrange + mock_features = RagPipelineTaskProxyTestDataFactory.create_mock_features() + mock_feature_service.get_features.return_value = mock_features + proxy = RagPipelineTaskProxyTestDataFactory.create_rag_pipeline_task_proxy() + proxy._send_to_priority_direct_queue = Mock() + + mock_file_service = Mock() + mock_file_service_class.return_value = mock_file_service + mock_upload_file = RagPipelineTaskProxyTestDataFactory.create_upload_file("file-123") + mock_file_service.upload_text.return_value = mock_upload_file + + # Act + proxy._dispatch() + + # If billing is disabled, for example: self-hosted or enterprise, should send to priority direct queue + proxy._send_to_priority_direct_queue.assert_called_once_with("file-123") + + @patch("services.rag_pipeline.rag_pipeline_task_proxy.FileService") + @patch("services.rag_pipeline.rag_pipeline_task_proxy.db") + def test_dispatch_with_empty_upload_file_id(self, mock_db: MagicMock, mock_file_service_class: MagicMock): + """Test _dispatch method when upload_file_id is empty.""" + # Arrange + proxy = RagPipelineTaskProxyTestDataFactory.create_rag_pipeline_task_proxy() + + mock_file_service = Mock() + mock_file_service_class.return_value = mock_file_service + upload_file = RagPipelineTaskProxyTestDataFactory.create_upload_file("") + mock_file_service.upload_text.return_value = upload_file + + # Act & Assert + with pytest.raises(ValueError, match="upload_file_id is empty"): + proxy._dispatch() + + @config_overrides_context(DEPLOYMENT_EDITION=DeploymentEdition.CLOUD) + @patch("services.rag_pipeline.rag_pipeline_task_proxy.FeatureService") + @patch("services.rag_pipeline.rag_pipeline_task_proxy.FileService") + @patch("services.rag_pipeline.rag_pipeline_task_proxy.db") + def test_dispatch_edge_case_empty_plan( + self, mock_db: MagicMock, mock_file_service_class: MagicMock, mock_feature_service: MagicMock + ): + """Test _dispatch method with empty plan string.""" + # Arrange + mock_features = RagPipelineTaskProxyTestDataFactory.create_mock_features(plan="") + mock_feature_service.get_features.return_value = mock_features + proxy = RagPipelineTaskProxyTestDataFactory.create_rag_pipeline_task_proxy() + proxy._send_to_priority_tenant_queue = Mock() + + mock_file_service = Mock() + mock_file_service_class.return_value = mock_file_service + mock_upload_file = RagPipelineTaskProxyTestDataFactory.create_upload_file("file-123") + mock_file_service.upload_text.return_value = mock_upload_file + + # Act + proxy._dispatch() + + # Assert + proxy._send_to_priority_tenant_queue.assert_called_once_with("file-123") + + @config_overrides_context(DEPLOYMENT_EDITION=DeploymentEdition.CLOUD) + @patch("services.rag_pipeline.rag_pipeline_task_proxy.FeatureService") + @patch("services.rag_pipeline.rag_pipeline_task_proxy.FileService") + @patch("services.rag_pipeline.rag_pipeline_task_proxy.db") + def test_dispatch_edge_case_none_plan( + self, mock_db: MagicMock, mock_file_service_class: MagicMock, mock_feature_service: MagicMock + ): + """Test _dispatch method with None plan.""" + # Arrange + mock_features = RagPipelineTaskProxyTestDataFactory.create_mock_features(plan=None) + mock_feature_service.get_features.return_value = mock_features + proxy = RagPipelineTaskProxyTestDataFactory.create_rag_pipeline_task_proxy() + proxy._send_to_priority_tenant_queue = Mock() + + mock_file_service = Mock() + mock_file_service_class.return_value = mock_file_service + mock_upload_file = RagPipelineTaskProxyTestDataFactory.create_upload_file("file-123") + mock_file_service.upload_text.return_value = mock_upload_file + + # Act + proxy._dispatch() + + # Assert + proxy._send_to_priority_tenant_queue.assert_called_once_with("file-123") + + @config_overrides_context(DEPLOYMENT_EDITION=DeploymentEdition.CLOUD) + @patch("services.rag_pipeline.rag_pipeline_task_proxy.FeatureService") + @patch("services.rag_pipeline.rag_pipeline_task_proxy.FileService") + @patch("services.rag_pipeline.rag_pipeline_task_proxy.db") + def test_delay_method( + self, mock_db: MagicMock, mock_file_service_class: MagicMock, mock_feature_service: MagicMock + ): + """Test delay method integration.""" + # Arrange + mock_features = RagPipelineTaskProxyTestDataFactory.create_mock_features(plan=CloudPlan.SANDBOX) + mock_feature_service.get_features.return_value = mock_features + proxy = RagPipelineTaskProxyTestDataFactory.create_rag_pipeline_task_proxy() + proxy._dispatch = Mock() + + mock_file_service = Mock() + mock_file_service_class.return_value = mock_file_service + mock_upload_file = RagPipelineTaskProxyTestDataFactory.create_upload_file("file-123") + mock_file_service.upload_text.return_value = mock_upload_file + + # Act + proxy.delay() + + # Assert + proxy._dispatch.assert_called_once() + + def test_delay_method_with_empty_entities(self, caplog: pytest.LogCaptureFixture): + """Test delay method with empty rag_pipeline_invoke_entities.""" + # Arrange + proxy = RagPipelineTaskProxy("tenant-123", "user-456", []) + + # Act + with caplog.at_level(logging.WARNING, logger="services.rag_pipeline.rag_pipeline_task_proxy"): + proxy.delay() + + # Assert + assert "Received empty rag pipeline invoke entities, no tasks delivered: tenant-123 user-456" in caplog.text diff --git a/api/tests/unit_tests/services/rag_pipeline/test_rag_pipeline_transform_service.py b/api/tests/unit_tests/services/rag_pipeline/test_rag_pipeline_transform_service.py index 43cfb85a828..44492ea3421 100644 --- a/api/tests/unit_tests/services/rag_pipeline/test_rag_pipeline_transform_service.py +++ b/api/tests/unit_tests/services/rag_pipeline/test_rag_pipeline_transform_service.py @@ -13,8 +13,8 @@ from models.model import UploadFile from services.entities.knowledge_entities.rag_pipeline_entities import KnowledgeConfiguration from services.errors.rag_pipeline import RagPipelineResourceNotFoundError from services.rag_pipeline.rag_pipeline_transform_service import RagPipelineTransformService -from tests.unit_tests.config_override import apply_config_overrides -from tests.unit_tests.model_factories import make_upload_file +from tests.helpers.unit.config_override import apply_config_overrides +from tests.helpers.unit.model_factories import make_upload_file def _dataset(**overrides: object) -> Dataset: diff --git a/api/tests/unit_tests/services/rbac/test_application.py b/api/tests/unit_tests/services/rbac/test___init__.py similarity index 98% rename from api/tests/unit_tests/services/rbac/test_application.py rename to api/tests/unit_tests/services/rbac/test___init__.py index e07567c999d..d31cd075502 100644 --- a/api/tests/unit_tests/services/rbac/test_application.py +++ b/api/tests/unit_tests/services/rbac/test___init__.py @@ -13,8 +13,8 @@ from models import Dataset, TenantAccountJoin from services.errors.base import NoPermissionError from services.errors.workspace import CannotOperateSelfError, InvalidWorkspaceMemberRoleError, RoleAlreadyAssignedError from services.rbac import contracts as dto -from tests.unit_tests.model_factories import make_account, make_app, make_tenant -from tests.unit_tests.rbac_fakes import Initialization, RBACDomain, build_rbac_domain +from tests.helpers.unit.model_factories import make_account, make_app, make_tenant +from tests.helpers.unit.rbac_fakes import Initialization, RBACDomain, build_rbac_domain CONTEXT = RequestContext("request", "trace", "actor", "workspace") diff --git a/api/tests/unit_tests/services/rbac/test_builtin.py b/api/tests/unit_tests/services/rbac/test_builtin.py index 696df002124..8cd34b4e208 100644 --- a/api/tests/unit_tests/services/rbac/test_builtin.py +++ b/api/tests/unit_tests/services/rbac/test_builtin.py @@ -1,6 +1,18 @@ -"""Built-in role projections keep the supported permission vocabulary.""" +from collections.abc import Callable +import pytest +from sqlalchemy import select +from sqlalchemy.orm import Session + +from enums.account import TenantAccountRole +from machinery.context import RequestContext +from models import TenantAccountJoin from services.rbac import builtin +from services.rbac import contracts as rbac_contracts +from tests.helpers.unit.model_factories import make_account, make_tenant +from tests.helpers.unit.rbac_fakes import RBACDomain + +"""Built-in role projections keep the supported permission vocabulary.""" class TestBuiltinAgentKeys: @@ -14,3 +26,232 @@ class TestBuiltinAgentKeys: ): assert "agent.manage" not in keys assert {"agent.acl.preview", "agent.acl.access_point_view"} <= set(keys) + + +"""Built-in permission projections and membership writes use the shared repository.""" + + +CONTEXT = RequestContext("request", None, "acct-1", "tenant-1") + + +@pytest.mark.parametrize( + ("role", "workspace_keys", "app_keys", "dataset_keys", "agent_keys"), + [ + ( + "owner", + builtin._BUILTIN_WORKSPACE_OWNER_KEYS, + builtin._BUILTIN_APP_OWNER_KEYS, + builtin._BUILTIN_DATASET_OWNER_KEYS, + builtin._BUILTIN_AGENT_FULL_ACCESS_KEYS, + ), + ( + "admin", + builtin._BUILTIN_WORKSPACE_ADMIN_KEYS, + builtin._BUILTIN_APP_ADMIN_KEYS, + builtin._BUILTIN_DATASET_ADMIN_KEYS, + builtin._BUILTIN_AGENT_FULL_ACCESS_KEYS, + ), + ( + "editor", + builtin._BUILTIN_WORKSPACE_EDITOR_KEYS, + builtin._BUILTIN_APP_EDITOR_KEYS, + builtin._BUILTIN_DATASET_EDITOR_KEYS, + builtin._BUILTIN_AGENT_FULL_ACCESS_KEYS, + ), + ( + "normal", + builtin._BUILTIN_WORKSPACE_NORMAL_KEYS, + builtin._BUILTIN_APP_NORMAL_KEYS, + [], + builtin._BUILTIN_AGENT_PREVIEW_KEYS, + ), + ( + "dataset_operator", + builtin._BUILTIN_WORKSPACE_DATASET_OPERATOR_KEYS, + builtin._BUILTIN_APP_DATASET_OPERATOR_KEYS, + builtin._BUILTIN_DATASET_DATASET_OPERATOR_KEYS, + builtin._BUILTIN_AGENT_PREVIEW_KEYS, + ), + ], +) +def test_get_uses_builtin_role_permissions_when_rbac_disabled( + rbac_domain: RBACDomain, + role: str, + workspace_keys: list[str], + app_keys: list[str], + dataset_keys: list[str], + agent_keys: list[str], + sqlite_session: Session, + config_overrides: Callable[..., None], +) -> None: + config_overrides(RBAC_ENABLED=False) + sqlite_session.add(TenantAccountJoin(tenant_id="tenant-1", account_id="acct-1", role=TenantAccountRole(role))) + sqlite_session.commit() + out = rbac_domain.rbac.members.permissions("tenant-1", "acct-1") + + assert not rbac_domain.transport.requests + assert out.workspace.permission_keys == workspace_keys + assert len(out.workspace.permission_keys) == len(set(out.workspace.permission_keys)) + assert out.app.default_permission_keys == app_keys + assert out.dataset.default_permission_keys == dataset_keys + assert out.agent.default_permission_keys == agent_keys + assert out.app.overrides == [] + assert out.dataset.overrides == [] + assert out.agent.overrides == [] + if role == "owner": + assert "snippets.management" in out.workspace.permission_keys + assert "app.acl.preview" in out.workspace.permission_keys + assert "dataset.acl.preview" in out.workspace.permission_keys + assert "app.acl.preview" in out.app.default_permission_keys + assert "dataset.acl.preview" in out.dataset.default_permission_keys + assert not any(key.startswith("billing.") for key in out.workspace.permission_keys) + if role == "editor": + assert "app.acl.log_and_annotation" in out.app.default_permission_keys + assert "app.acl.deploy" not in out.app.default_permission_keys + + +@pytest.mark.parametrize( + ("role", "expected_snippet_keys"), + [ + ("owner", {"snippets.create_and_modify", "snippets.management"}), + ("admin", {"snippets.create_and_modify", "snippets.management"}), + ("editor", {"snippets.create_and_modify"}), + ("normal", set()), + ("dataset_operator", set()), + ], +) +def test_get_uses_builtin_snippet_permissions_when_rbac_disabled( + rbac_domain: RBACDomain, + role: str, + expected_snippet_keys: set[str], + sqlite_session: Session, + config_overrides: Callable[..., None], +) -> None: + config_overrides(RBAC_ENABLED=False) + sqlite_session.add(TenantAccountJoin(tenant_id="tenant-1", account_id="acct-1", role=TenantAccountRole(role))) + sqlite_session.commit() + out = rbac_domain.rbac.members.permissions("tenant-1", "acct-1") + + actual_snippet_keys = { + permission_key for permission_key in out.workspace.permission_keys if permission_key.startswith("snippets.") + } + + assert not rbac_domain.transport.requests + assert actual_snippet_keys == expected_snippet_keys + + +def test_get_returns_empty_when_role_missing_and_rbac_disabled( + rbac_domain: RBACDomain, config_overrides: Callable[..., None] +) -> None: + config_overrides(RBAC_ENABLED=False) + out = rbac_domain.rbac.members.permissions("tenant-1", "acct-1") + + assert not rbac_domain.transport.requests + assert out.workspace.permission_keys == [] + assert out.app.default_permission_keys == [] + assert out.dataset.default_permission_keys == [] + assert out.agent.default_permission_keys == [] + + +def test_get_builtin_role_includes_permission_keys( + rbac_domain: RBACDomain, sqlite_session: Session, config_overrides: Callable[..., None] +) -> None: + config_overrides(RBAC_ENABLED=False) + sqlite_session.add(TenantAccountJoin(tenant_id="tenant-1", account_id="acct-2", role=TenantAccountRole.EDITOR)) + sqlite_session.commit() + + out = rbac_domain.rbac.members.get(CONTEXT, "acct-2", language=None) + + assert not rbac_domain.transport.requests + assert out.account_id == "acct-2" + assert out.roles[0].name == "editor" + assert out.roles[0].permission_keys == list( + dict.fromkeys( + [ + *builtin._BUILTIN_WORKSPACE_EDITOR_KEYS, + *builtin._BUILTIN_APP_EDITOR_KEYS, + *builtin._BUILTIN_DATASET_EDITOR_KEYS, + *builtin._BUILTIN_AGENT_FULL_ACCESS_KEYS, + ] + ) + ) + assert "snippets.create_and_modify" in out.roles[0].permission_keys + assert "app.acl.preview" in out.roles[0].permission_keys + assert "dataset.acl.preview" in out.roles[0].permission_keys + assert "app.acl.deploy" not in out.roles[0].permission_keys + + +def test_replace_commits_builtin_join_role_when_rbac_disabled( + rbac_domain: RBACDomain, sqlite_session: Session, config_overrides: Callable[..., None] +) -> None: + config_overrides(RBAC_ENABLED=False) + sqlite_session.add_all( + [ + make_tenant(tenant_id="tenant-1"), + make_account(account_id="acct-1", email="owner@example.com"), + make_account(account_id="acct-2", email="member@example.com"), + TenantAccountJoin(tenant_id="tenant-1", account_id="acct-1", role=TenantAccountRole.OWNER), + ] + ) + target_join = TenantAccountJoin(tenant_id="tenant-1", account_id="acct-2", role=TenantAccountRole.NORMAL) + sqlite_session.add(target_join) + sqlite_session.commit() + target_join_id = target_join.id + engine = sqlite_session.get_bind() + + out = rbac_domain.rbac.members.replace(CONTEXT, "acct-2", role_ids=["editor"], language=None) + + assert not rbac_domain.transport.requests + # Closing the writer rolls back any uncommitted update and prevents its identity map + # from satisfying the verification query. + sqlite_session.close() + with Session(engine) as verification_session: + persisted_join = verification_session.scalar( + select(TenantAccountJoin).where(TenantAccountJoin.id == target_join_id) + ) + assert persisted_join is not None + assert persisted_join.role == TenantAccountRole.EDITOR + assert out.account_id == "acct-2" + assert out.roles[0].id == "editor" + assert "app.acl.preview" in out.roles[0].permission_keys + + +def test_app_permissions_batch_get_uses_builtin_role_permissions_when_rbac_disabled( + rbac_domain: RBACDomain, sqlite_session: Session, config_overrides: Callable[..., None] +) -> None: + config_overrides(RBAC_ENABLED=False) + sqlite_session.add(TenantAccountJoin(tenant_id="tenant-1", account_id="acct-1", role=TenantAccountRole.EDITOR)) + sqlite_session.commit() + out = rbac_domain.rbac.members.resource_permissions( + "tenant-1", "acct-1", rbac_contracts.RBACResourceType.APP, ["app-1", "app-2"] + ) + + assert not rbac_domain.transport.requests + assert out == { + "app-1": builtin._BUILTIN_APP_EDITOR_KEYS, + "app-2": builtin._BUILTIN_APP_EDITOR_KEYS, + } + assert all("app.acl.deploy" not in permission_keys for permission_keys in out.values()) + + +def test_dataset_permissions_batch_get_uses_builtin_role_permissions_when_rbac_disabled( + rbac_domain: RBACDomain, sqlite_session: Session, config_overrides: Callable[..., None] +) -> None: + config_overrides(RBAC_ENABLED=False) + sqlite_session.add( + TenantAccountJoin( + tenant_id="tenant-1", + account_id="acct-1", + role=TenantAccountRole.DATASET_OPERATOR, + ) + ) + sqlite_session.commit() + out = rbac_domain.rbac.members.resource_permissions( + "tenant-1", "acct-1", rbac_contracts.RBACResourceType.DATASET, ["ds-1", "ds-2"] + ) + + assert not rbac_domain.transport.requests + assert out == { + "ds-1": builtin._BUILTIN_DATASET_DATASET_OPERATOR_KEYS, + "ds-2": builtin._BUILTIN_DATASET_DATASET_OPERATOR_KEYS, + } diff --git a/api/tests/unit_tests/services/rbac/test_gateways.py b/api/tests/unit_tests/services/rbac/test_gateways.py index 7df5a893df5..3b59ebde2a2 100644 --- a/api/tests/unit_tests/services/rbac/test_gateways.py +++ b/api/tests/unit_tests/services/rbac/test_gateways.py @@ -2,7 +2,7 @@ import pytest from services.enterprise.base import EnterpriseRequest from services.enterprise.rbac_service import RBACService -from tests.unit_tests.rbac_fakes import RBACTransport +from tests.helpers.unit.rbac_fakes import RBACTransport def test_membership_gateway_sends_identity_language_and_assignment(monkeypatch: pytest.MonkeyPatch) -> None: diff --git a/api/tests/unit_tests/services/rbac/test_member_permissions.py b/api/tests/unit_tests/services/rbac/test_member_permissions.py deleted file mode 100644 index 9e02bf45fa5..00000000000 --- a/api/tests/unit_tests/services/rbac/test_member_permissions.py +++ /dev/null @@ -1,240 +0,0 @@ -"""Built-in permission projections and membership writes use the shared repository.""" - -from collections.abc import Callable - -import pytest -from sqlalchemy import select -from sqlalchemy.orm import Session - -from enums.account import TenantAccountRole -from machinery.context import RequestContext -from models import TenantAccountJoin -from services.rbac import builtin -from services.rbac import contracts as rbac_contracts -from tests.unit_tests.model_factories import make_account, make_tenant -from tests.unit_tests.rbac_fakes import RBACDomain - -CONTEXT = RequestContext("request", None, "acct-1", "tenant-1") - - -@pytest.mark.parametrize( - ("role", "workspace_keys", "app_keys", "dataset_keys", "agent_keys"), - [ - ( - "owner", - builtin._BUILTIN_WORKSPACE_OWNER_KEYS, - builtin._BUILTIN_APP_OWNER_KEYS, - builtin._BUILTIN_DATASET_OWNER_KEYS, - builtin._BUILTIN_AGENT_FULL_ACCESS_KEYS, - ), - ( - "admin", - builtin._BUILTIN_WORKSPACE_ADMIN_KEYS, - builtin._BUILTIN_APP_ADMIN_KEYS, - builtin._BUILTIN_DATASET_ADMIN_KEYS, - builtin._BUILTIN_AGENT_FULL_ACCESS_KEYS, - ), - ( - "editor", - builtin._BUILTIN_WORKSPACE_EDITOR_KEYS, - builtin._BUILTIN_APP_EDITOR_KEYS, - builtin._BUILTIN_DATASET_EDITOR_KEYS, - builtin._BUILTIN_AGENT_FULL_ACCESS_KEYS, - ), - ( - "normal", - builtin._BUILTIN_WORKSPACE_NORMAL_KEYS, - builtin._BUILTIN_APP_NORMAL_KEYS, - [], - builtin._BUILTIN_AGENT_PREVIEW_KEYS, - ), - ( - "dataset_operator", - builtin._BUILTIN_WORKSPACE_DATASET_OPERATOR_KEYS, - builtin._BUILTIN_APP_DATASET_OPERATOR_KEYS, - builtin._BUILTIN_DATASET_DATASET_OPERATOR_KEYS, - builtin._BUILTIN_AGENT_PREVIEW_KEYS, - ), - ], -) -def test_get_uses_builtin_role_permissions_when_rbac_disabled( - rbac_domain: RBACDomain, - role: str, - workspace_keys: list[str], - app_keys: list[str], - dataset_keys: list[str], - agent_keys: list[str], - sqlite_session: Session, - config_overrides: Callable[..., None], -) -> None: - config_overrides(RBAC_ENABLED=False) - sqlite_session.add(TenantAccountJoin(tenant_id="tenant-1", account_id="acct-1", role=TenantAccountRole(role))) - sqlite_session.commit() - out = rbac_domain.rbac.members.permissions("tenant-1", "acct-1") - - assert not rbac_domain.transport.requests - assert out.workspace.permission_keys == workspace_keys - assert len(out.workspace.permission_keys) == len(set(out.workspace.permission_keys)) - assert out.app.default_permission_keys == app_keys - assert out.dataset.default_permission_keys == dataset_keys - assert out.agent.default_permission_keys == agent_keys - assert out.app.overrides == [] - assert out.dataset.overrides == [] - assert out.agent.overrides == [] - if role == "owner": - assert "snippets.management" in out.workspace.permission_keys - assert "app.acl.preview" in out.workspace.permission_keys - assert "dataset.acl.preview" in out.workspace.permission_keys - assert "app.acl.preview" in out.app.default_permission_keys - assert "dataset.acl.preview" in out.dataset.default_permission_keys - assert not any(key.startswith("billing.") for key in out.workspace.permission_keys) - if role == "editor": - assert "app.acl.log_and_annotation" in out.app.default_permission_keys - assert "app.acl.deploy" not in out.app.default_permission_keys - - -@pytest.mark.parametrize( - ("role", "expected_snippet_keys"), - [ - ("owner", {"snippets.create_and_modify", "snippets.management"}), - ("admin", {"snippets.create_and_modify", "snippets.management"}), - ("editor", {"snippets.create_and_modify"}), - ("normal", set()), - ("dataset_operator", set()), - ], -) -def test_get_uses_builtin_snippet_permissions_when_rbac_disabled( - rbac_domain: RBACDomain, - role: str, - expected_snippet_keys: set[str], - sqlite_session: Session, - config_overrides: Callable[..., None], -) -> None: - config_overrides(RBAC_ENABLED=False) - sqlite_session.add(TenantAccountJoin(tenant_id="tenant-1", account_id="acct-1", role=TenantAccountRole(role))) - sqlite_session.commit() - out = rbac_domain.rbac.members.permissions("tenant-1", "acct-1") - - actual_snippet_keys = { - permission_key for permission_key in out.workspace.permission_keys if permission_key.startswith("snippets.") - } - - assert not rbac_domain.transport.requests - assert actual_snippet_keys == expected_snippet_keys - - -def test_get_returns_empty_when_role_missing_and_rbac_disabled( - rbac_domain: RBACDomain, config_overrides: Callable[..., None] -) -> None: - config_overrides(RBAC_ENABLED=False) - out = rbac_domain.rbac.members.permissions("tenant-1", "acct-1") - - assert not rbac_domain.transport.requests - assert out.workspace.permission_keys == [] - assert out.app.default_permission_keys == [] - assert out.dataset.default_permission_keys == [] - assert out.agent.default_permission_keys == [] - - -def test_get_builtin_role_includes_permission_keys( - rbac_domain: RBACDomain, sqlite_session: Session, config_overrides: Callable[..., None] -) -> None: - config_overrides(RBAC_ENABLED=False) - sqlite_session.add(TenantAccountJoin(tenant_id="tenant-1", account_id="acct-2", role=TenantAccountRole.EDITOR)) - sqlite_session.commit() - - out = rbac_domain.rbac.members.get(CONTEXT, "acct-2", language=None) - - assert not rbac_domain.transport.requests - assert out.account_id == "acct-2" - assert out.roles[0].name == "editor" - assert out.roles[0].permission_keys == list( - dict.fromkeys( - [ - *builtin._BUILTIN_WORKSPACE_EDITOR_KEYS, - *builtin._BUILTIN_APP_EDITOR_KEYS, - *builtin._BUILTIN_DATASET_EDITOR_KEYS, - *builtin._BUILTIN_AGENT_FULL_ACCESS_KEYS, - ] - ) - ) - assert "snippets.create_and_modify" in out.roles[0].permission_keys - assert "app.acl.preview" in out.roles[0].permission_keys - assert "dataset.acl.preview" in out.roles[0].permission_keys - assert "app.acl.deploy" not in out.roles[0].permission_keys - - -def test_replace_commits_builtin_join_role_when_rbac_disabled( - rbac_domain: RBACDomain, sqlite_session: Session, config_overrides: Callable[..., None] -) -> None: - config_overrides(RBAC_ENABLED=False) - sqlite_session.add_all( - [ - make_tenant(tenant_id="tenant-1"), - make_account(account_id="acct-1", email="owner@example.com"), - make_account(account_id="acct-2", email="member@example.com"), - TenantAccountJoin(tenant_id="tenant-1", account_id="acct-1", role=TenantAccountRole.OWNER), - ] - ) - target_join = TenantAccountJoin(tenant_id="tenant-1", account_id="acct-2", role=TenantAccountRole.NORMAL) - sqlite_session.add(target_join) - sqlite_session.commit() - target_join_id = target_join.id - engine = sqlite_session.get_bind() - - out = rbac_domain.rbac.members.replace(CONTEXT, "acct-2", role_ids=["editor"], language=None) - - assert not rbac_domain.transport.requests - # Closing the writer rolls back any uncommitted update and prevents its identity map - # from satisfying the verification query. - sqlite_session.close() - with Session(engine) as verification_session: - persisted_join = verification_session.scalar( - select(TenantAccountJoin).where(TenantAccountJoin.id == target_join_id) - ) - assert persisted_join is not None - assert persisted_join.role == TenantAccountRole.EDITOR - assert out.account_id == "acct-2" - assert out.roles[0].id == "editor" - assert "app.acl.preview" in out.roles[0].permission_keys - - -def test_app_permissions_batch_get_uses_builtin_role_permissions_when_rbac_disabled( - rbac_domain: RBACDomain, sqlite_session: Session, config_overrides: Callable[..., None] -) -> None: - config_overrides(RBAC_ENABLED=False) - sqlite_session.add(TenantAccountJoin(tenant_id="tenant-1", account_id="acct-1", role=TenantAccountRole.EDITOR)) - sqlite_session.commit() - out = rbac_domain.rbac.members.resource_permissions( - "tenant-1", "acct-1", rbac_contracts.RBACResourceType.APP, ["app-1", "app-2"] - ) - - assert not rbac_domain.transport.requests - assert out == { - "app-1": builtin._BUILTIN_APP_EDITOR_KEYS, - "app-2": builtin._BUILTIN_APP_EDITOR_KEYS, - } - assert all("app.acl.deploy" not in permission_keys for permission_keys in out.values()) - - -def test_dataset_permissions_batch_get_uses_builtin_role_permissions_when_rbac_disabled( - rbac_domain: RBACDomain, sqlite_session: Session, config_overrides: Callable[..., None] -) -> None: - config_overrides(RBAC_ENABLED=False) - sqlite_session.add( - TenantAccountJoin( - tenant_id="tenant-1", - account_id="acct-1", - role=TenantAccountRole.DATASET_OPERATOR, - ) - ) - sqlite_session.commit() - out = rbac_domain.rbac.members.resource_permissions( - "tenant-1", "acct-1", rbac_contracts.RBACResourceType.DATASET, ["ds-1", "ds-2"] - ) - - assert not rbac_domain.transport.requests - assert out == { - "ds-1": builtin._BUILTIN_DATASET_DATASET_OPERATOR_KEYS, - "ds-2": builtin._BUILTIN_DATASET_DATASET_OPERATOR_KEYS, - } diff --git a/api/tests/unit_tests/services/rbac/test_resource_queries.py b/api/tests/unit_tests/services/rbac/test_resource_queries.py index 3f08d48323d..941fcf5c0c4 100644 --- a/api/tests/unit_tests/services/rbac/test_resource_queries.py +++ b/api/tests/unit_tests/services/rbac/test_resource_queries.py @@ -7,8 +7,8 @@ from models.agent import Agent, AgentKind, AgentScope, AgentSource, AgentStatus from models.model import AppMode from services.entities.app_entities import AppAgentBinding from services.rbac.resource_queries import ResourceQueryService -from tests.unit_tests.model_factories import make_app, make_dataset, make_document -from tests.unit_tests.rbac_fakes import RBACDomain +from tests.helpers.unit.model_factories import make_app, make_dataset, make_document +from tests.helpers.unit.rbac_fakes import RBACDomain @pytest.fixture diff --git a/api/tests/unit_tests/services/test_export_app_messages.py b/api/tests/unit_tests/services/retention/conversation/test_message_export_service.py similarity index 100% rename from api/tests/unit_tests/services/test_export_app_messages.py rename to api/tests/unit_tests/services/retention/conversation/test_message_export_service.py diff --git a/api/tests/unit_tests/services/retention/test_messages_clean_policy.py b/api/tests/unit_tests/services/retention/conversation/test_messages_clean_policy.py similarity index 100% rename from api/tests/unit_tests/services/retention/test_messages_clean_policy.py rename to api/tests/unit_tests/services/retention/conversation/test_messages_clean_policy.py diff --git a/api/tests/unit_tests/services/retention/conversation/test_messages_clean_service.py b/api/tests/unit_tests/services/retention/conversation/test_messages_clean_service.py new file mode 100644 index 00000000000..ed844b91295 --- /dev/null +++ b/api/tests/unit_tests/services/retention/conversation/test_messages_clean_service.py @@ -0,0 +1,1078 @@ +import datetime +import math +import uuid +from collections.abc import Iterator +from decimal import Decimal +from types import SimpleNamespace +from typing import Any, TypedDict +from unittest.mock import MagicMock, patch + +import pytest +from sqlalchemy import Engine, delete, func, select + +from core.db.session_factory import session_factory +from enums import CloudPlan, DeploymentEdition +from models import Tenant +from models.enums import FeedbackFromSource, FeedbackRating +from models.model import ( + App, + AppMode, + Conversation, + ConversationFromSource, + ConversationStatus, + Message, + MessageAnnotation, + MessageFeedback, +) +from services.retention.conversation import messages_clean_service as messages_clean_service_module +from services.retention.conversation.messages_clean_policy import ( + BillingDisabledPolicy, + BillingSandboxPolicy, + SimpleMessage, + create_message_clean_policy, +) +from services.retention.conversation.messages_clean_service import MessagesCleanService + +"""SQLite-backed coverage for message retention cleanup.""" + + +_NOW = datetime.datetime(2026, 1, 15, 12, 0, 0, tzinfo=datetime.UTC) +_OLD = _NOW - datetime.timedelta(days=60) +_VERY_OLD = _NOW - datetime.timedelta(days=90) +_RECENT = _NOW - datetime.timedelta(days=5) + +_WINDOW_START = _VERY_OLD - datetime.timedelta(hours=1) +_WINDOW_END = _RECENT + datetime.timedelta(hours=1) + +_DEFAULT_BATCH_SIZE = 100 +_PAGINATION_MESSAGE_COUNT = 25 +_PAGINATION_BATCH_SIZE = 8 + + +class TenantAndApp(TypedDict): + tenant_id: str + app_id: str + conversation_id: str + + +class SeedMessages(TenantAndApp): + msg_ids: dict[str, str] + + +class PaginatedSeedMessages(TenantAndApp): + msg_ids: list[str] + + +class CascadeTestData(TenantAndApp): + msg_id: str + fb_id: str + ann_id: str + + +def _make_message(app_id: str, conversation_id: str, created_at: datetime.datetime) -> Message: + message = Message( + app_id=app_id, + conversation_id=conversation_id, + query="test", + message=[{"text": "hello"}], + answer="world", + message_tokens=1, + message_unit_price=Decimal(0), + answer_tokens=1, + answer_unit_price=Decimal(0), + from_source=ConversationFromSource.CONSOLE, + currency="USD", + _inputs={}, + ) + message.created_at = created_at + return message + + +class TestMessagesCleanServiceSqliteCases: + @staticmethod + @pytest.fixture(autouse=True) + def _bind_service_engine(monkeypatch: pytest.MonkeyPatch, sqlite_engine: Engine) -> None: + monkeypatch.setattr(messages_clean_service_module, "db", SimpleNamespace(engine=sqlite_engine)) + + @staticmethod + @pytest.fixture + def tenant_and_app() -> Iterator[TenantAndApp]: + """Creates a Tenant, App and Conversation for the test and cleans up after.""" + with session_factory.create_session() as session: + tenant = Tenant(name="retention_it_tenant") + session.add(tenant) + session.flush() + + app = App( + tenant_id=tenant.id, + name="Retention IT App", + mode=AppMode.CHAT, + enable_site=True, + enable_api=True, + ) + session.add(app) + session.flush() + + conv = Conversation( + app_id=app.id, + mode=AppMode.CHAT, + name="test_conv", + status=ConversationStatus.NORMAL, + from_source=ConversationFromSource.CONSOLE, + _inputs={}, + ) + session.add(conv) + session.commit() + + tenant_id = tenant.id + app_id = app.id + conv_id = conv.id + + yield {"tenant_id": tenant_id, "app_id": app_id, "conversation_id": conv_id} + + with session_factory.create_session() as session: + session.execute(delete(Conversation).where(Conversation.id == conv_id)) + session.execute(delete(App).where(App.id == app_id)) + session.execute(delete(Tenant).where(Tenant.id == tenant_id)) + session.commit() + + class TestMessagesCleanServiceIntegration: + @pytest.fixture + def seed_messages(self, tenant_and_app: TenantAndApp) -> Iterator[SeedMessages]: + """Seeds one message at each of _VERY_OLD, _OLD, and _RECENT. + Yields a semantic mapping keyed by age label. + """ + data = tenant_and_app + app_id = data["app_id"] + conv_id = data["conversation_id"] + # Ordered tuple of (label, timestamp) for deterministic seeding + timestamps = [ + ("very_old", _VERY_OLD), + ("old", _OLD), + ("recent", _RECENT), + ] + msg_ids: dict[str, str] = {} + + with session_factory.create_session() as session: + for label, ts in timestamps: + msg = _make_message(app_id, conv_id, ts) + session.add(msg) + session.flush() + msg_ids[label] = msg.id + session.commit() + + yield {"msg_ids": msg_ids, "tenant_id": data["tenant_id"], "app_id": app_id, "conversation_id": conv_id} + + with session_factory.create_session() as session: + session.execute( + delete(Message) + .where(Message.id.in_(list(msg_ids.values()))) + .execution_options(synchronize_session=False) + ) + session.commit() + + @pytest.fixture + def paginated_seed_messages(self, tenant_and_app: TenantAndApp) -> Iterator[PaginatedSeedMessages]: + """Seeds multiple messages separated by 1-second increments starting at _OLD.""" + data = tenant_and_app + app_id = data["app_id"] + conv_id = data["conversation_id"] + msg_ids: list[str] = [] + + with session_factory.create_session() as session: + for i in range(_PAGINATION_MESSAGE_COUNT): + ts = _OLD + datetime.timedelta(seconds=i) + msg = _make_message(app_id, conv_id, ts) + session.add(msg) + session.flush() + msg_ids.append(msg.id) + session.commit() + + yield {"msg_ids": msg_ids, "tenant_id": data["tenant_id"], "app_id": app_id, "conversation_id": conv_id} + + with session_factory.create_session() as session: + session.execute( + delete(Message).where(Message.id.in_(msg_ids)).execution_options(synchronize_session=False) + ) + session.commit() + + @pytest.fixture + def cascade_test_data(self, tenant_and_app: TenantAndApp) -> Iterator[CascadeTestData]: + """Seeds one Message with an associated Feedback and Annotation.""" + data = tenant_and_app + app_id = data["app_id"] + conv_id = data["conversation_id"] + + with session_factory.create_session() as session: + msg = _make_message(app_id, conv_id, _OLD) + session.add(msg) + session.flush() + + feedback = MessageFeedback( + app_id=app_id, + conversation_id=conv_id, + message_id=msg.id, + rating=FeedbackRating.LIKE, + from_source=FeedbackFromSource.USER, + ) + annotation = MessageAnnotation( + app_id=app_id, + conversation_id=conv_id, + message_id=msg.id, + question="q", + content="a", + account_id=str(uuid.uuid4()), + ) + session.add_all([feedback, annotation]) + session.commit() + + msg_id = msg.id + fb_id = feedback.id + ann_id = annotation.id + + yield { + "msg_id": msg_id, + "fb_id": fb_id, + "ann_id": ann_id, + "tenant_id": data["tenant_id"], + "app_id": app_id, + "conversation_id": conv_id, + } + + with session_factory.create_session() as session: + session.execute(delete(MessageAnnotation).where(MessageAnnotation.id == ann_id)) + session.execute(delete(MessageFeedback).where(MessageFeedback.id == fb_id)) + session.execute(delete(Message).where(Message.id == msg_id)) + session.commit() + + def test_dry_run_does_not_delete(self, seed_messages: SeedMessages) -> None: + """Dry-run must count eligible rows without deleting any of them.""" + data = seed_messages + msg_ids = data["msg_ids"] + all_ids = list(msg_ids.values()) + + svc = MessagesCleanService.from_time_range( + policy=BillingDisabledPolicy(), + start_from=_WINDOW_START, + end_before=_WINDOW_END, + batch_size=_DEFAULT_BATCH_SIZE, + dry_run=True, + ) + stats = svc.run() + + assert stats["filtered_messages"] == len(all_ids) + assert stats["total_deleted"] == 0 + + with session_factory.create_session() as session: + remaining = session.scalar(select(func.count()).select_from(Message).where(Message.id.in_(all_ids))) + assert remaining == len(all_ids) + + def test_billing_disabled_deletes_all_in_range(self, seed_messages: SeedMessages) -> None: + """All 3 seeded messages fall within the window and must be deleted.""" + data = seed_messages + msg_ids = data["msg_ids"] + all_ids = list(msg_ids.values()) + + svc = MessagesCleanService.from_time_range( + policy=BillingDisabledPolicy(), + start_from=_WINDOW_START, + end_before=_WINDOW_END, + batch_size=_DEFAULT_BATCH_SIZE, + dry_run=False, + ) + stats = svc.run() + + assert stats["total_deleted"] == len(all_ids) + + with session_factory.create_session() as session: + remaining = session.scalar(select(func.count()).select_from(Message).where(Message.id.in_(all_ids))) + assert remaining == 0 + + def test_start_from_filters_correctly(self, seed_messages: SeedMessages) -> None: + """Only the message at _OLD falls within the narrow ±1 h window.""" + data = seed_messages + msg_ids = data["msg_ids"] + + start = _OLD - datetime.timedelta(hours=1) + end = _OLD + datetime.timedelta(hours=1) + + svc = MessagesCleanService.from_time_range( + policy=BillingDisabledPolicy(), + start_from=start, + end_before=end, + batch_size=_DEFAULT_BATCH_SIZE, + ) + stats = svc.run() + + assert stats["total_deleted"] == 1 + + with session_factory.create_session() as session: + all_ids = list(msg_ids.values()) + remaining_ids = set(session.scalars(select(Message.id).where(Message.id.in_(all_ids))).all()) + + assert msg_ids["old"] not in remaining_ids + assert msg_ids["very_old"] in remaining_ids + assert msg_ids["recent"] in remaining_ids + + def test_cursor_pagination_across_batches(self, paginated_seed_messages: PaginatedSeedMessages) -> None: + """Messages must be deleted across multiple batches.""" + data = paginated_seed_messages + msg_ids = data["msg_ids"] + + # _OLD is the earliest; the last one is _OLD + (_PAGINATION_MESSAGE_COUNT - 1) s. + pagination_window_start = _OLD - datetime.timedelta(seconds=1) + pagination_window_end = _OLD + datetime.timedelta(seconds=_PAGINATION_MESSAGE_COUNT) + + svc = MessagesCleanService.from_time_range( + policy=BillingDisabledPolicy(), + start_from=pagination_window_start, + end_before=pagination_window_end, + batch_size=_PAGINATION_BATCH_SIZE, + ) + stats = svc.run() + + assert stats["total_deleted"] == _PAGINATION_MESSAGE_COUNT + expected_batches = math.ceil(_PAGINATION_MESSAGE_COUNT / _PAGINATION_BATCH_SIZE) + assert stats["batches"] >= expected_batches + + with session_factory.create_session() as session: + remaining = session.scalar(select(func.count()).select_from(Message).where(Message.id.in_(msg_ids))) + assert remaining == 0 + + def test_no_messages_in_range_returns_empty_stats(self, seed_messages: SeedMessages) -> None: + """A window entirely in the future must yield zero matches.""" + _ = seed_messages + far_future = _NOW + datetime.timedelta(days=365) + even_further = far_future + datetime.timedelta(days=1) + + svc = MessagesCleanService.from_time_range( + policy=BillingDisabledPolicy(), + start_from=far_future, + end_before=even_further, + batch_size=_DEFAULT_BATCH_SIZE, + ) + stats = svc.run() + + assert stats["total_messages"] == 0 + assert stats["total_deleted"] == 0 + + def test_relation_cascade_deletes(self, cascade_test_data: CascadeTestData) -> None: + """Deleting a Message must cascade to its Feedback and Annotation rows.""" + data = cascade_test_data + msg_id = data["msg_id"] + fb_id = data["fb_id"] + ann_id = data["ann_id"] + + svc = MessagesCleanService.from_time_range( + policy=BillingDisabledPolicy(), + start_from=_OLD - datetime.timedelta(hours=1), + end_before=_OLD + datetime.timedelta(hours=1), + batch_size=_DEFAULT_BATCH_SIZE, + ) + stats = svc.run() + + assert stats["total_deleted"] == 1 + + with session_factory.create_session() as session: + assert session.scalar(select(func.count()).select_from(Message).where(Message.id == msg_id)) == 0 + assert ( + session.scalar(select(func.count()).select_from(MessageFeedback).where(MessageFeedback.id == fb_id)) + == 0 + ) + assert ( + session.scalar( + select(func.count()).select_from(MessageAnnotation).where(MessageAnnotation.id == ann_id) + ) + == 0 + ) + + def test_factory_from_time_range_validation(self) -> None: + with pytest.raises(ValueError, match="start_from"): + MessagesCleanService.from_time_range( + policy=BillingDisabledPolicy(), + start_from=_NOW, + end_before=_OLD, + ) + + def test_factory_from_days_validation(self) -> None: + with pytest.raises(ValueError, match="days"): + MessagesCleanService.from_days( + policy=BillingDisabledPolicy(), + days=-1, + ) + + def test_factory_batch_size_validation(self) -> None: + with pytest.raises(ValueError, match="batch_size"): + MessagesCleanService.from_time_range( + policy=BillingDisabledPolicy(), + start_from=_OLD, + end_before=_NOW, + batch_size=0, + ) + + +def make_simple_message(msg_id: str, app_id: str) -> SimpleMessage: + """Helper to create a SimpleMessage with a fixed created_at timestamp.""" + return SimpleMessage(id=msg_id, app_id=app_id, created_at=datetime.datetime(2024, 1, 1)) + + +def make_plan_provider(tenant_plans: dict[str, Any]) -> MagicMock: + """Helper to create a mock plan_provider that returns the given tenant_plans.""" + provider = MagicMock() + provider.return_value = tenant_plans + return provider + + +class TestMessagesCleanServiceCases: + class TestBillingSandboxPolicyFilterMessageIds: + """Unit tests for BillingSandboxPolicy.filter_message_ids method.""" + + # Fixed timestamp for deterministic tests + CURRENT_TIMESTAMP = 1000000 + GRACEFUL_PERIOD_DAYS = 8 + GRACEFUL_PERIOD_SECONDS = GRACEFUL_PERIOD_DAYS * 24 * 60 * 60 + + def test_missing_tenant_mapping_excluded(self): + """Test that messages with missing app-to-tenant mapping are excluded.""" + # Arrange + messages = [ + make_simple_message("msg1", "app1"), + make_simple_message("msg2", "app2"), + ] + app_to_tenant = {} # No mapping + tenant_plans = {"tenant1": {"plan": CloudPlan.SANDBOX, "expiration_date": -1}} + plan_provider = make_plan_provider(tenant_plans) + + policy = BillingSandboxPolicy( + plan_provider=plan_provider, + graceful_period_days=self.GRACEFUL_PERIOD_DAYS, + current_timestamp=self.CURRENT_TIMESTAMP, + ) + + # Act + result = policy.filter_message_ids(messages, app_to_tenant) + + # Assert + assert list(result) == [] + + def test_missing_tenant_plan_excluded(self): + """Test that messages with missing tenant plan are excluded (safe default).""" + # Arrange + messages = [ + make_simple_message("msg1", "app1"), + make_simple_message("msg2", "app2"), + ] + app_to_tenant = {"app1": "tenant1", "app2": "tenant2"} + tenant_plans = {} # No plans + plan_provider = make_plan_provider(tenant_plans) + + policy = BillingSandboxPolicy( + plan_provider=plan_provider, + graceful_period_days=self.GRACEFUL_PERIOD_DAYS, + current_timestamp=self.CURRENT_TIMESTAMP, + ) + + # Act + result = policy.filter_message_ids(messages, app_to_tenant) + + # Assert + assert list(result) == [] + + def test_non_sandbox_plan_excluded(self): + """Test that messages from non-sandbox plans (PROFESSIONAL/TEAM) are excluded.""" + # Arrange + messages = [ + make_simple_message("msg1", "app1"), + make_simple_message("msg2", "app2"), + make_simple_message("msg3", "app3"), + ] + app_to_tenant = {"app1": "tenant1", "app2": "tenant2", "app3": "tenant3"} + tenant_plans = { + "tenant1": {"plan": CloudPlan.PROFESSIONAL, "expiration_date": -1}, + "tenant2": {"plan": CloudPlan.TEAM, "expiration_date": -1}, + "tenant3": {"plan": CloudPlan.SANDBOX, "expiration_date": -1}, # Only this one + } + plan_provider = make_plan_provider(tenant_plans) + + policy = BillingSandboxPolicy( + plan_provider=plan_provider, + graceful_period_days=self.GRACEFUL_PERIOD_DAYS, + current_timestamp=self.CURRENT_TIMESTAMP, + ) + + # Act + result = policy.filter_message_ids(messages, app_to_tenant) + + # Assert - only msg3 (sandbox tenant) should be included + assert set(result) == {"msg3"} + + def test_whitelist_skip(self): + """Test that whitelisted tenants are excluded even if sandbox + expired.""" + # Arrange + messages = [ + make_simple_message("msg1", "app1"), # Whitelisted - excluded + make_simple_message("msg2", "app2"), # Not whitelisted - included + make_simple_message("msg3", "app3"), # Whitelisted - excluded + ] + app_to_tenant = {"app1": "tenant1", "app2": "tenant2", "app3": "tenant3"} + tenant_plans = { + "tenant1": {"plan": CloudPlan.SANDBOX, "expiration_date": -1}, + "tenant2": {"plan": CloudPlan.SANDBOX, "expiration_date": -1}, + "tenant3": {"plan": CloudPlan.SANDBOX, "expiration_date": -1}, + } + plan_provider = make_plan_provider(tenant_plans) + tenant_whitelist = ["tenant1", "tenant3"] + + policy = BillingSandboxPolicy( + plan_provider=plan_provider, + graceful_period_days=self.GRACEFUL_PERIOD_DAYS, + tenant_whitelist=tenant_whitelist, + current_timestamp=self.CURRENT_TIMESTAMP, + ) + + # Act + result = policy.filter_message_ids(messages, app_to_tenant) + + # Assert - only msg2 should be included + assert set(result) == {"msg2"} + + def test_no_previous_subscription_included(self): + """Test that messages with expiration_date=-1 (no previous subscription) are included.""" + # Arrange + messages = [ + make_simple_message("msg1", "app1"), + make_simple_message("msg2", "app2"), + ] + app_to_tenant = {"app1": "tenant1", "app2": "tenant2"} + tenant_plans = { + "tenant1": {"plan": CloudPlan.SANDBOX, "expiration_date": -1}, + "tenant2": {"plan": CloudPlan.SANDBOX, "expiration_date": -1}, + } + plan_provider = make_plan_provider(tenant_plans) + + policy = BillingSandboxPolicy( + plan_provider=plan_provider, + graceful_period_days=self.GRACEFUL_PERIOD_DAYS, + current_timestamp=self.CURRENT_TIMESTAMP, + ) + + # Act + result = policy.filter_message_ids(messages, app_to_tenant) + + # Assert - all messages should be included + assert set(result) == {"msg1", "msg2"} + + def test_within_grace_period_excluded(self): + """Test that messages within grace period are excluded.""" + # Arrange + now = self.CURRENT_TIMESTAMP + expired_1_day_ago = now - (1 * 24 * 60 * 60) + expired_5_days_ago = now - (5 * 24 * 60 * 60) + expired_7_days_ago = now - (7 * 24 * 60 * 60) + + messages = [ + make_simple_message("msg1", "app1"), + make_simple_message("msg2", "app2"), + make_simple_message("msg3", "app3"), + ] + app_to_tenant = {"app1": "tenant1", "app2": "tenant2", "app3": "tenant3"} + tenant_plans = { + "tenant1": {"plan": CloudPlan.SANDBOX, "expiration_date": expired_1_day_ago}, + "tenant2": {"plan": CloudPlan.SANDBOX, "expiration_date": expired_5_days_ago}, + "tenant3": {"plan": CloudPlan.SANDBOX, "expiration_date": expired_7_days_ago}, + } + plan_provider = make_plan_provider(tenant_plans) + + policy = BillingSandboxPolicy( + plan_provider=plan_provider, + graceful_period_days=self.GRACEFUL_PERIOD_DAYS, # 8 days + current_timestamp=now, + ) + + # Act + result = policy.filter_message_ids(messages, app_to_tenant) + + # Assert - all within 8-day grace period, none should be included + assert list(result) == [] + + def test_exactly_at_boundary_excluded(self): + """Test that messages exactly at grace period boundary are excluded (code uses >).""" + # Arrange + now = self.CURRENT_TIMESTAMP + expired_exactly_8_days_ago = now - self.GRACEFUL_PERIOD_SECONDS # Exactly at boundary + + messages = [make_simple_message("msg1", "app1")] + app_to_tenant = {"app1": "tenant1"} + tenant_plans = { + "tenant1": {"plan": CloudPlan.SANDBOX, "expiration_date": expired_exactly_8_days_ago}, + } + plan_provider = make_plan_provider(tenant_plans) + + policy = BillingSandboxPolicy( + plan_provider=plan_provider, + graceful_period_days=self.GRACEFUL_PERIOD_DAYS, + current_timestamp=now, + ) + + # Act + result = policy.filter_message_ids(messages, app_to_tenant) + + # Assert - exactly at boundary (==) should be excluded (code uses >) + assert list(result) == [] + + def test_beyond_grace_period_included(self): + """Test that messages beyond grace period are included.""" + # Arrange + now = self.CURRENT_TIMESTAMP + expired_9_days_ago = now - (9 * 24 * 60 * 60) # Just beyond 8-day grace + expired_30_days_ago = now - (30 * 24 * 60 * 60) # Well beyond + + messages = [ + make_simple_message("msg1", "app1"), + make_simple_message("msg2", "app2"), + ] + app_to_tenant = {"app1": "tenant1", "app2": "tenant2"} + tenant_plans = { + "tenant1": {"plan": CloudPlan.SANDBOX, "expiration_date": expired_9_days_ago}, + "tenant2": {"plan": CloudPlan.SANDBOX, "expiration_date": expired_30_days_ago}, + } + plan_provider = make_plan_provider(tenant_plans) + + policy = BillingSandboxPolicy( + plan_provider=plan_provider, + graceful_period_days=self.GRACEFUL_PERIOD_DAYS, + current_timestamp=now, + ) + + # Act + result = policy.filter_message_ids(messages, app_to_tenant) + + # Assert - both beyond grace period, should be included + assert set(result) == {"msg1", "msg2"} + + def test_empty_messages_returns_empty(self): + """Test that empty messages returns empty list.""" + # Arrange + messages: list[SimpleMessage] = [] + app_to_tenant = {"app1": "tenant1"} + plan_provider = make_plan_provider({"tenant1": {"plan": CloudPlan.SANDBOX, "expiration_date": -1}}) + + policy = BillingSandboxPolicy( + plan_provider=plan_provider, + graceful_period_days=self.GRACEFUL_PERIOD_DAYS, + current_timestamp=self.CURRENT_TIMESTAMP, + ) + + # Act + result = policy.filter_message_ids(messages, app_to_tenant) + + # Assert + assert list(result) == [] + + def test_plan_provider_called_with_correct_tenant_ids(self): + """Test that plan_provider is called with correct tenant_ids.""" + # Arrange + messages = [ + make_simple_message("msg1", "app1"), + make_simple_message("msg2", "app2"), + make_simple_message("msg3", "app3"), + ] + app_to_tenant = {"app1": "tenant1", "app2": "tenant2", "app3": "tenant1"} # tenant1 appears twice + plan_provider = make_plan_provider({}) + + policy = BillingSandboxPolicy( + plan_provider=plan_provider, + graceful_period_days=self.GRACEFUL_PERIOD_DAYS, + current_timestamp=self.CURRENT_TIMESTAMP, + ) + + # Act + policy.filter_message_ids(messages, app_to_tenant) + + # Assert - plan_provider should be called once with unique tenant_ids + plan_provider.assert_called_once() + called_tenant_ids = set(plan_provider.call_args[0][0]) + assert called_tenant_ids == {"tenant1", "tenant2"} + + def test_complex_mixed_scenario(self): + """Test complex scenario with mixed plans, expirations, whitelist, and missing mappings.""" + # Arrange + now = self.CURRENT_TIMESTAMP + sandbox_expired_old = now - (15 * 24 * 60 * 60) # Beyond grace + sandbox_expired_recent = now - (3 * 24 * 60 * 60) # Within grace + future_expiration = now + (30 * 24 * 60 * 60) + + messages = [ + make_simple_message("msg1", "app1"), # Sandbox, no subscription - included + make_simple_message("msg2", "app2"), # Sandbox, expired old - included + make_simple_message("msg3", "app3"), # Sandbox, within grace - excluded + make_simple_message("msg4", "app4"), # Team plan, active - excluded + make_simple_message("msg5", "app5"), # No tenant mapping - excluded + make_simple_message("msg6", "app6"), # No plan info - excluded + make_simple_message("msg7", "app7"), # Sandbox, expired old, whitelisted - excluded + ] + app_to_tenant = { + "app1": "tenant1", + "app2": "tenant2", + "app3": "tenant3", + "app4": "tenant4", + "app6": "tenant6", # Has mapping but no plan + "app7": "tenant7", + # app5 has no mapping + } + tenant_plans = { + "tenant1": {"plan": CloudPlan.SANDBOX, "expiration_date": -1}, + "tenant2": {"plan": CloudPlan.SANDBOX, "expiration_date": sandbox_expired_old}, + "tenant3": {"plan": CloudPlan.SANDBOX, "expiration_date": sandbox_expired_recent}, + "tenant4": {"plan": CloudPlan.TEAM, "expiration_date": future_expiration}, + "tenant7": {"plan": CloudPlan.SANDBOX, "expiration_date": sandbox_expired_old}, + # tenant6 has no plan + } + plan_provider = make_plan_provider(tenant_plans) + tenant_whitelist = ["tenant7"] + + policy = BillingSandboxPolicy( + plan_provider=plan_provider, + graceful_period_days=self.GRACEFUL_PERIOD_DAYS, + tenant_whitelist=tenant_whitelist, + current_timestamp=now, + ) + + # Act + result = policy.filter_message_ids(messages, app_to_tenant) + + # Assert - only msg1 and msg2 should be included + assert set(result) == {"msg1", "msg2"} + + class TestBillingDisabledPolicyFilterMessageIds: + """Unit tests for BillingDisabledPolicy.filter_message_ids method.""" + + def test_returns_all_message_ids(self): + """Test that all message IDs are returned (order-preserving).""" + # Arrange + messages = [ + make_simple_message("msg1", "app1"), + make_simple_message("msg2", "app2"), + make_simple_message("msg3", "app3"), + ] + app_to_tenant = {"app1": "tenant1", "app2": "tenant2"} + + policy = BillingDisabledPolicy() + + # Act + result = policy.filter_message_ids(messages, app_to_tenant) + + # Assert - all message IDs returned in order + assert list(result) == ["msg1", "msg2", "msg3"] + + def test_ignores_app_to_tenant(self): + """Test that app_to_tenant mapping is ignored.""" + # Arrange + messages = [ + make_simple_message("msg1", "app1"), + make_simple_message("msg2", "app2"), + ] + app_to_tenant: dict[str, str] = {} # Empty - should be ignored + + policy = BillingDisabledPolicy() + + # Act + result = policy.filter_message_ids(messages, app_to_tenant) + + # Assert - all message IDs still returned + assert list(result) == ["msg1", "msg2"] + + def test_empty_messages_returns_empty(self): + """Test that empty messages returns empty list.""" + # Arrange + messages: list[SimpleMessage] = [] + app_to_tenant = {"app1": "tenant1"} + + policy = BillingDisabledPolicy() + + # Act + result = policy.filter_message_ids(messages, app_to_tenant) + + # Assert + assert list(result) == [] + + class TestCreateMessageCleanPolicy: + """Unit tests for create_message_clean_policy factory function.""" + + def test_non_cloud_edition_returns_billing_disabled_policy(self, config_overrides): + """Test that the Community edition returns BillingDisabledPolicy.""" + # Arrange + config_overrides(DEPLOYMENT_EDITION=DeploymentEdition.COMMUNITY) + + # Act + policy = create_message_clean_policy(graceful_period_days=21) + + # Assert + assert isinstance(policy, BillingDisabledPolicy) + + @patch("services.retention.conversation.messages_clean_policy.BillingService", autospec=True) + def test_cloud_edition_policy_has_correct_internals(self, mock_billing_service, config_overrides): + """Test that BillingSandboxPolicy is created with correct internal values.""" + # Arrange + config_overrides(DEPLOYMENT_EDITION=DeploymentEdition.CLOUD) + whitelist = ["tenant1", "tenant2"] + mock_billing_service.get_expired_subscription_cleanup_whitelist.return_value = whitelist + mock_plan_provider = MagicMock() + mock_billing_service.get_plan_bulk_with_cache = mock_plan_provider + + # Act + policy = create_message_clean_policy(graceful_period_days=14, current_timestamp=1234567) + + # Assert + mock_billing_service.get_expired_subscription_cleanup_whitelist.assert_called_once() + assert isinstance(policy, BillingSandboxPolicy) + assert policy._graceful_period_days == 14 + assert list(policy._tenant_whitelist) == whitelist + assert policy._plan_provider == mock_plan_provider + assert policy._current_timestamp == 1234567 + + class TestMessagesCleanServiceFromTimeRange: + """Unit tests for MessagesCleanService.from_time_range factory method.""" + + def test_start_from_end_before_raises_value_error(self): + """Test that start_from == end_before raises ValueError.""" + policy = BillingDisabledPolicy() + + # Arrange + same_time = datetime.datetime(2024, 1, 1, 12, 0, 0) + + # Act & Assert + with pytest.raises(ValueError, match="start_from .* must be less than end_before"): + MessagesCleanService.from_time_range( + policy=policy, + start_from=same_time, + end_before=same_time, + ) + + # Arrange + start_from = datetime.datetime(2024, 12, 31) + end_before = datetime.datetime(2024, 1, 1) + + # Act & Assert + with pytest.raises(ValueError, match="start_from .* must be less than end_before"): + MessagesCleanService.from_time_range( + policy=policy, + start_from=start_from, + end_before=end_before, + ) + + def test_batch_size_raises_value_error(self): + """Test that batch_size=0 raises ValueError.""" + # Arrange + start_from = datetime.datetime(2024, 1, 1) + end_before = datetime.datetime(2024, 2, 1) + policy = BillingDisabledPolicy() + + # Act & Assert + with pytest.raises(ValueError, match="batch_size .* must be greater than 0"): + MessagesCleanService.from_time_range( + policy=policy, + start_from=start_from, + end_before=end_before, + batch_size=0, + ) + + start_from = datetime.datetime(2024, 1, 1) + end_before = datetime.datetime(2024, 2, 1) + policy = BillingDisabledPolicy() + + # Act & Assert + with pytest.raises(ValueError, match="batch_size .* must be greater than 0"): + MessagesCleanService.from_time_range( + policy=policy, + start_from=start_from, + end_before=end_before, + batch_size=-100, + ) + + def test_valid_params_creates_instance(self): + """Test that valid parameters create a correctly configured instance.""" + # Arrange + start_from = datetime.datetime(2024, 1, 1, 0, 0, 0) + end_before = datetime.datetime(2024, 12, 31, 23, 59, 59) + policy = BillingDisabledPolicy() + batch_size = 500 + dry_run = True + + # Act + service = MessagesCleanService.from_time_range( + policy=policy, + start_from=start_from, + end_before=end_before, + batch_size=batch_size, + dry_run=dry_run, + ) + + # Assert + assert isinstance(service, MessagesCleanService) + assert service._policy is policy + assert service._start_from == start_from + assert service._end_before == end_before + assert service._batch_size == batch_size + assert service._dry_run == dry_run + + def test_default_params(self): + """Test that default parameters are applied correctly.""" + # Arrange + start_from = datetime.datetime(2024, 1, 1) + end_before = datetime.datetime(2024, 2, 1) + policy = BillingDisabledPolicy() + + # Act + service = MessagesCleanService.from_time_range( + policy=policy, + start_from=start_from, + end_before=end_before, + ) + + # Assert + assert service._batch_size == 1000 # default + assert service._dry_run is False # default + + def test_explicit_task_label(self): + start_from = datetime.datetime(2024, 1, 1) + end_before = datetime.datetime(2024, 1, 2) + policy = BillingDisabledPolicy() + + service = MessagesCleanService.from_time_range( + policy=policy, + start_from=start_from, + end_before=end_before, + task_label="60to30", + ) + + assert service._metrics._base_attributes["task_label"] == "60to30" + + class TestMessagesCleanServiceFromDays: + """Unit tests for MessagesCleanService.from_days factory method.""" + + def test_days_raises_value_error(self): + """Test that days < 0 raises ValueError.""" + # Arrange + policy = BillingDisabledPolicy() + + # Act & Assert + with pytest.raises(ValueError, match="days .* must be greater than or equal to 0"): + MessagesCleanService.from_days(policy=policy, days=-1) + + # Act + with patch("services.retention.conversation.messages_clean_service.naive_utc_now") as mock_now: + fixed_now = datetime.datetime(2024, 6, 15, 14, 0, 0) + mock_now.return_value = fixed_now + service = MessagesCleanService.from_days(policy=policy, days=0) + + # Assert + assert service._end_before == fixed_now + + def test_batch_size_raises_value_error(self): + """Test that batch_size=0 raises ValueError.""" + # Arrange + policy = BillingDisabledPolicy() + + # Act & Assert + with pytest.raises(ValueError, match="batch_size .* must be greater than 0"): + MessagesCleanService.from_days(policy=policy, days=30, batch_size=0) + + # Act & Assert + with pytest.raises(ValueError, match="batch_size .* must be greater than 0"): + MessagesCleanService.from_days(policy=policy, days=30, batch_size=-500) + + def test_valid_params_creates_instance(self): + """Test that valid parameters create a correctly configured instance.""" + # Arrange + policy = BillingDisabledPolicy() + days = 90 + batch_size = 500 + dry_run = True + + # Act + with patch("services.retention.conversation.messages_clean_service.naive_utc_now") as mock_now: + fixed_now = datetime.datetime(2024, 6, 15, 10, 30, 0) + mock_now.return_value = fixed_now + service = MessagesCleanService.from_days( + policy=policy, + days=days, + batch_size=batch_size, + dry_run=dry_run, + ) + + # Assert + expected_end_before = fixed_now - datetime.timedelta(days=days) + assert isinstance(service, MessagesCleanService) + assert service._policy is policy + assert service._start_from is None + assert service._end_before == expected_end_before + assert service._batch_size == batch_size + assert service._dry_run == dry_run + + def test_default_params(self): + """Test that default parameters are applied correctly.""" + # Arrange + policy = BillingDisabledPolicy() + + # Act + with patch("services.retention.conversation.messages_clean_service.naive_utc_now") as mock_now: + fixed_now = datetime.datetime(2024, 6, 15, 10, 30, 0) + mock_now.return_value = fixed_now + service = MessagesCleanService.from_days(policy=policy) + + # Assert + expected_end_before = fixed_now - datetime.timedelta(days=30) # default days=30 + assert service._end_before == expected_end_before + assert service._batch_size == 1000 # default + assert service._dry_run is False # default + assert service._metrics._base_attributes["task_label"] == "custom" + + class TestMessagesCleanServiceRun: + """Unit tests for MessagesCleanService.run instrumentation behavior.""" + + def test_run_records_completion_metrics_on_success(self): + # Arrange + service = MessagesCleanService( + policy=BillingDisabledPolicy(), + start_from=datetime.datetime(2024, 1, 1), + end_before=datetime.datetime(2024, 1, 2), + batch_size=100, + dry_run=False, + ) + expected_stats = { + "batches": 1, + "total_messages": 10, + "filtered_messages": 5, + "total_deleted": 5, + } + service._clean_messages_by_time_range = MagicMock(return_value=expected_stats) # type: ignore[method-assign] + completion_calls: list[dict[str, object]] = [] + service._metrics.record_completion = lambda **kwargs: completion_calls.append(kwargs) # type: ignore[method-assign] + + # Act + result = service.run() + + # Assert + assert result == expected_stats + assert len(completion_calls) == 1 + assert completion_calls[0]["status"] == "success" + + def test_run_records_completion_metrics_on_failure(self): + # Arrange + service = MessagesCleanService( + policy=BillingDisabledPolicy(), + start_from=datetime.datetime(2024, 1, 1), + end_before=datetime.datetime(2024, 1, 2), + batch_size=100, + dry_run=False, + ) + service._clean_messages_by_time_range = MagicMock(side_effect=RuntimeError("clean failed")) # type: ignore[method-assign] + completion_calls: list[dict[str, object]] = [] + service._metrics.record_completion = lambda **kwargs: completion_calls.append(kwargs) # type: ignore[method-assign] + + # Act & Assert + with pytest.raises(RuntimeError, match="clean failed"): + service.run() + assert len(completion_calls) == 1 + assert completion_calls[0]["status"] == "failed" diff --git a/api/tests/unit_tests/services/retention/test_messages_clean_service_sqlite.py b/api/tests/unit_tests/services/retention/test_messages_clean_service_sqlite.py deleted file mode 100644 index 10cae80c214..00000000000 --- a/api/tests/unit_tests/services/retention/test_messages_clean_service_sqlite.py +++ /dev/null @@ -1,400 +0,0 @@ -"""SQLite-backed coverage for message retention cleanup.""" - -import datetime -import math -import uuid -from collections.abc import Iterator -from decimal import Decimal -from types import SimpleNamespace -from typing import TypedDict - -import pytest -from sqlalchemy import Engine, delete, func, select - -from core.db.session_factory import session_factory -from models import Tenant -from models.enums import FeedbackFromSource, FeedbackRating -from models.model import ( - App, - AppMode, - Conversation, - ConversationFromSource, - ConversationStatus, - Message, - MessageAnnotation, - MessageFeedback, -) -from services.retention.conversation import messages_clean_service as messages_clean_service_module -from services.retention.conversation.messages_clean_policy import BillingDisabledPolicy -from services.retention.conversation.messages_clean_service import MessagesCleanService - -_NOW = datetime.datetime(2026, 1, 15, 12, 0, 0, tzinfo=datetime.UTC) -_OLD = _NOW - datetime.timedelta(days=60) -_VERY_OLD = _NOW - datetime.timedelta(days=90) -_RECENT = _NOW - datetime.timedelta(days=5) - -_WINDOW_START = _VERY_OLD - datetime.timedelta(hours=1) -_WINDOW_END = _RECENT + datetime.timedelta(hours=1) - -_DEFAULT_BATCH_SIZE = 100 -_PAGINATION_MESSAGE_COUNT = 25 -_PAGINATION_BATCH_SIZE = 8 - - -class TenantAndApp(TypedDict): - tenant_id: str - app_id: str - conversation_id: str - - -class SeedMessages(TenantAndApp): - msg_ids: dict[str, str] - - -class PaginatedSeedMessages(TenantAndApp): - msg_ids: list[str] - - -class CascadeTestData(TenantAndApp): - msg_id: str - fb_id: str - ann_id: str - - -@pytest.fixture(autouse=True) -def _bind_service_engine(monkeypatch: pytest.MonkeyPatch, sqlite_engine: Engine) -> None: - monkeypatch.setattr(messages_clean_service_module, "db", SimpleNamespace(engine=sqlite_engine)) - - -@pytest.fixture -def tenant_and_app() -> Iterator[TenantAndApp]: - """Creates a Tenant, App and Conversation for the test and cleans up after.""" - with session_factory.create_session() as session: - tenant = Tenant(name="retention_it_tenant") - session.add(tenant) - session.flush() - - app = App( - tenant_id=tenant.id, - name="Retention IT App", - mode=AppMode.CHAT, - enable_site=True, - enable_api=True, - ) - session.add(app) - session.flush() - - conv = Conversation( - app_id=app.id, - mode=AppMode.CHAT, - name="test_conv", - status=ConversationStatus.NORMAL, - from_source=ConversationFromSource.CONSOLE, - _inputs={}, - ) - session.add(conv) - session.commit() - - tenant_id = tenant.id - app_id = app.id - conv_id = conv.id - - yield {"tenant_id": tenant_id, "app_id": app_id, "conversation_id": conv_id} - - with session_factory.create_session() as session: - session.execute(delete(Conversation).where(Conversation.id == conv_id)) - session.execute(delete(App).where(App.id == app_id)) - session.execute(delete(Tenant).where(Tenant.id == tenant_id)) - session.commit() - - -def _make_message(app_id: str, conversation_id: str, created_at: datetime.datetime) -> Message: - message = Message( - app_id=app_id, - conversation_id=conversation_id, - query="test", - message=[{"text": "hello"}], - answer="world", - message_tokens=1, - message_unit_price=Decimal(0), - answer_tokens=1, - answer_unit_price=Decimal(0), - from_source=ConversationFromSource.CONSOLE, - currency="USD", - _inputs={}, - ) - message.created_at = created_at - return message - - -class TestMessagesCleanServiceIntegration: - @pytest.fixture - def seed_messages(self, tenant_and_app: TenantAndApp) -> Iterator[SeedMessages]: - """Seeds one message at each of _VERY_OLD, _OLD, and _RECENT. - Yields a semantic mapping keyed by age label. - """ - data = tenant_and_app - app_id = data["app_id"] - conv_id = data["conversation_id"] - # Ordered tuple of (label, timestamp) for deterministic seeding - timestamps = [ - ("very_old", _VERY_OLD), - ("old", _OLD), - ("recent", _RECENT), - ] - msg_ids: dict[str, str] = {} - - with session_factory.create_session() as session: - for label, ts in timestamps: - msg = _make_message(app_id, conv_id, ts) - session.add(msg) - session.flush() - msg_ids[label] = msg.id - session.commit() - - yield {"msg_ids": msg_ids, "tenant_id": data["tenant_id"], "app_id": app_id, "conversation_id": conv_id} - - with session_factory.create_session() as session: - session.execute( - delete(Message) - .where(Message.id.in_(list(msg_ids.values()))) - .execution_options(synchronize_session=False) - ) - session.commit() - - @pytest.fixture - def paginated_seed_messages(self, tenant_and_app: TenantAndApp) -> Iterator[PaginatedSeedMessages]: - """Seeds multiple messages separated by 1-second increments starting at _OLD.""" - data = tenant_and_app - app_id = data["app_id"] - conv_id = data["conversation_id"] - msg_ids: list[str] = [] - - with session_factory.create_session() as session: - for i in range(_PAGINATION_MESSAGE_COUNT): - ts = _OLD + datetime.timedelta(seconds=i) - msg = _make_message(app_id, conv_id, ts) - session.add(msg) - session.flush() - msg_ids.append(msg.id) - session.commit() - - yield {"msg_ids": msg_ids, "tenant_id": data["tenant_id"], "app_id": app_id, "conversation_id": conv_id} - - with session_factory.create_session() as session: - session.execute(delete(Message).where(Message.id.in_(msg_ids)).execution_options(synchronize_session=False)) - session.commit() - - @pytest.fixture - def cascade_test_data(self, tenant_and_app: TenantAndApp) -> Iterator[CascadeTestData]: - """Seeds one Message with an associated Feedback and Annotation.""" - data = tenant_and_app - app_id = data["app_id"] - conv_id = data["conversation_id"] - - with session_factory.create_session() as session: - msg = _make_message(app_id, conv_id, _OLD) - session.add(msg) - session.flush() - - feedback = MessageFeedback( - app_id=app_id, - conversation_id=conv_id, - message_id=msg.id, - rating=FeedbackRating.LIKE, - from_source=FeedbackFromSource.USER, - ) - annotation = MessageAnnotation( - app_id=app_id, - conversation_id=conv_id, - message_id=msg.id, - question="q", - content="a", - account_id=str(uuid.uuid4()), - ) - session.add_all([feedback, annotation]) - session.commit() - - msg_id = msg.id - fb_id = feedback.id - ann_id = annotation.id - - yield { - "msg_id": msg_id, - "fb_id": fb_id, - "ann_id": ann_id, - "tenant_id": data["tenant_id"], - "app_id": app_id, - "conversation_id": conv_id, - } - - with session_factory.create_session() as session: - session.execute(delete(MessageAnnotation).where(MessageAnnotation.id == ann_id)) - session.execute(delete(MessageFeedback).where(MessageFeedback.id == fb_id)) - session.execute(delete(Message).where(Message.id == msg_id)) - session.commit() - - def test_dry_run_does_not_delete(self, seed_messages: SeedMessages) -> None: - """Dry-run must count eligible rows without deleting any of them.""" - data = seed_messages - msg_ids = data["msg_ids"] - all_ids = list(msg_ids.values()) - - svc = MessagesCleanService.from_time_range( - policy=BillingDisabledPolicy(), - start_from=_WINDOW_START, - end_before=_WINDOW_END, - batch_size=_DEFAULT_BATCH_SIZE, - dry_run=True, - ) - stats = svc.run() - - assert stats["filtered_messages"] == len(all_ids) - assert stats["total_deleted"] == 0 - - with session_factory.create_session() as session: - remaining = session.scalar(select(func.count()).select_from(Message).where(Message.id.in_(all_ids))) - assert remaining == len(all_ids) - - def test_billing_disabled_deletes_all_in_range(self, seed_messages: SeedMessages) -> None: - """All 3 seeded messages fall within the window and must be deleted.""" - data = seed_messages - msg_ids = data["msg_ids"] - all_ids = list(msg_ids.values()) - - svc = MessagesCleanService.from_time_range( - policy=BillingDisabledPolicy(), - start_from=_WINDOW_START, - end_before=_WINDOW_END, - batch_size=_DEFAULT_BATCH_SIZE, - dry_run=False, - ) - stats = svc.run() - - assert stats["total_deleted"] == len(all_ids) - - with session_factory.create_session() as session: - remaining = session.scalar(select(func.count()).select_from(Message).where(Message.id.in_(all_ids))) - assert remaining == 0 - - def test_start_from_filters_correctly(self, seed_messages: SeedMessages) -> None: - """Only the message at _OLD falls within the narrow ±1 h window.""" - data = seed_messages - msg_ids = data["msg_ids"] - - start = _OLD - datetime.timedelta(hours=1) - end = _OLD + datetime.timedelta(hours=1) - - svc = MessagesCleanService.from_time_range( - policy=BillingDisabledPolicy(), - start_from=start, - end_before=end, - batch_size=_DEFAULT_BATCH_SIZE, - ) - stats = svc.run() - - assert stats["total_deleted"] == 1 - - with session_factory.create_session() as session: - all_ids = list(msg_ids.values()) - remaining_ids = set(session.scalars(select(Message.id).where(Message.id.in_(all_ids))).all()) - - assert msg_ids["old"] not in remaining_ids - assert msg_ids["very_old"] in remaining_ids - assert msg_ids["recent"] in remaining_ids - - def test_cursor_pagination_across_batches(self, paginated_seed_messages: PaginatedSeedMessages) -> None: - """Messages must be deleted across multiple batches.""" - data = paginated_seed_messages - msg_ids = data["msg_ids"] - - # _OLD is the earliest; the last one is _OLD + (_PAGINATION_MESSAGE_COUNT - 1) s. - pagination_window_start = _OLD - datetime.timedelta(seconds=1) - pagination_window_end = _OLD + datetime.timedelta(seconds=_PAGINATION_MESSAGE_COUNT) - - svc = MessagesCleanService.from_time_range( - policy=BillingDisabledPolicy(), - start_from=pagination_window_start, - end_before=pagination_window_end, - batch_size=_PAGINATION_BATCH_SIZE, - ) - stats = svc.run() - - assert stats["total_deleted"] == _PAGINATION_MESSAGE_COUNT - expected_batches = math.ceil(_PAGINATION_MESSAGE_COUNT / _PAGINATION_BATCH_SIZE) - assert stats["batches"] >= expected_batches - - with session_factory.create_session() as session: - remaining = session.scalar(select(func.count()).select_from(Message).where(Message.id.in_(msg_ids))) - assert remaining == 0 - - def test_no_messages_in_range_returns_empty_stats(self, seed_messages: SeedMessages) -> None: - """A window entirely in the future must yield zero matches.""" - _ = seed_messages - far_future = _NOW + datetime.timedelta(days=365) - even_further = far_future + datetime.timedelta(days=1) - - svc = MessagesCleanService.from_time_range( - policy=BillingDisabledPolicy(), - start_from=far_future, - end_before=even_further, - batch_size=_DEFAULT_BATCH_SIZE, - ) - stats = svc.run() - - assert stats["total_messages"] == 0 - assert stats["total_deleted"] == 0 - - def test_relation_cascade_deletes(self, cascade_test_data: CascadeTestData) -> None: - """Deleting a Message must cascade to its Feedback and Annotation rows.""" - data = cascade_test_data - msg_id = data["msg_id"] - fb_id = data["fb_id"] - ann_id = data["ann_id"] - - svc = MessagesCleanService.from_time_range( - policy=BillingDisabledPolicy(), - start_from=_OLD - datetime.timedelta(hours=1), - end_before=_OLD + datetime.timedelta(hours=1), - batch_size=_DEFAULT_BATCH_SIZE, - ) - stats = svc.run() - - assert stats["total_deleted"] == 1 - - with session_factory.create_session() as session: - assert session.scalar(select(func.count()).select_from(Message).where(Message.id == msg_id)) == 0 - assert ( - session.scalar(select(func.count()).select_from(MessageFeedback).where(MessageFeedback.id == fb_id)) - == 0 - ) - assert ( - session.scalar( - select(func.count()).select_from(MessageAnnotation).where(MessageAnnotation.id == ann_id) - ) - == 0 - ) - - def test_factory_from_time_range_validation(self) -> None: - with pytest.raises(ValueError, match="start_from"): - MessagesCleanService.from_time_range( - policy=BillingDisabledPolicy(), - start_from=_NOW, - end_before=_OLD, - ) - - def test_factory_from_days_validation(self) -> None: - with pytest.raises(ValueError, match="days"): - MessagesCleanService.from_days( - policy=BillingDisabledPolicy(), - days=-1, - ) - - def test_factory_batch_size_validation(self) -> None: - with pytest.raises(ValueError, match="batch_size"): - MessagesCleanService.from_time_range( - policy=BillingDisabledPolicy(), - start_from=_OLD, - end_before=_NOW, - batch_size=0, - ) diff --git a/api/tests/unit_tests/services/retention/workflow_run/test_archive_paid_plan_workflow_run.py b/api/tests/unit_tests/services/retention/workflow_run/test_archive_paid_plan_workflow_run.py new file mode 100644 index 00000000000..96fdc9a629a --- /dev/null +++ b/api/tests/unit_tests/services/retention/workflow_run/test_archive_paid_plan_workflow_run.py @@ -0,0 +1,766 @@ +import datetime as _workflow_run_archiver_sqlite_datetime +import json +import uuid +from collections.abc import Callable +from datetime import datetime as _archive_workflow_run_logs_datetime +from types import SimpleNamespace +from typing import override +from unittest.mock import ANY, MagicMock, patch + +import pyarrow as pa +import pyarrow.parquet as pq +import pytest +from sqlalchemy import select +from sqlalchemy.engine import Engine +from sqlalchemy.exc import OperationalError +from sqlalchemy.orm import Session, sessionmaker + +from enums import DeploymentEdition +from graphon.enums import WorkflowExecutionStatus +from libs.archive_storage import ArchiveStorage +from models.enums import CreatorUserRole, WorkflowRunTriggeredFrom +from models.workflow import WorkflowRun, WorkflowRunArchiveBundle, WorkflowType +from services.retention.workflow_run.archive_paid_plan_workflow_run import ( + ArchiveResult, + ArchiveSummary, + TableStats, + WorkflowRunArchiver, +) +from services.retention.workflow_run.constants import ARCHIVE_BUNDLE_FORMAT, ARCHIVE_BUNDLE_SCHEMA_VERSION + +"""SQLite and in-process coverage for workflow-run bundle archiving.""" + + +class FakeArchiveStorage(ArchiveStorage): + def __init__(self, objects: dict[str, bytes] | None = None) -> None: + self.objects = {} if objects is None else objects + + @override + def object_exists(self, key: str) -> bool: + return key in self.objects + + @override + def get_object(self, key: str) -> bytes: + return self.objects[key] + + @override + def put_object(self, key: str, data: bytes) -> str: + self.objects[key] = data + return "checksum" + + @override + def list_objects(self, prefix: str) -> list[str]: + return sorted(key for key in self.objects if key.startswith(prefix)) + + +def _db_disconnect_error() -> OperationalError: + return OperationalError( + "select 1", + {}, + RuntimeError("server closed the connection unexpectedly"), + connection_invalidated=True, + ) + + +def _run(run_id: str = "run-1") -> WorkflowRun: + """Build the mapped entity consumed by the archiver without persisting it.""" + return WorkflowRun( + id=run_id, + tenant_id="tenant-1", + app_id="app-1", + workflow_id="workflow-1", + type=WorkflowType.WORKFLOW, + triggered_from=WorkflowRunTriggeredFrom.APP_RUN, + version="1", + graph="{}", + inputs="{}", + status=WorkflowExecutionStatus.SUCCEEDED, + outputs="{}", + error=None, + elapsed_time=0, + total_tokens=0, + total_steps=0, + created_by_role=CreatorUserRole.ACCOUNT, + created_by="account-1", + created_at=_workflow_run_archiver_sqlite_datetime.datetime(2025, 3, 15, 10, 0, 0), + finished_at=_workflow_run_archiver_sqlite_datetime.datetime(2025, 3, 15, 10, 1, 0), + ) + + +class TestWorkflowRunArchiverSqliteCases: + class TestWorkflowRunArchiverInit: + def test_start_from_without_end_before_raises(self) -> None: + with pytest.raises(ValueError, match="start_from and end_before must be provided together"): + WorkflowRunArchiver(start_from=_workflow_run_archiver_sqlite_datetime.datetime(2025, 1, 1)) + + def test_end_before_without_start_from_raises(self) -> None: + with pytest.raises(ValueError, match="start_from and end_before must be provided together"): + WorkflowRunArchiver(end_before=_workflow_run_archiver_sqlite_datetime.datetime(2025, 1, 1)) + + def test_start_equals_end_raises(self) -> None: + ts = _workflow_run_archiver_sqlite_datetime.datetime(2025, 1, 1) + with pytest.raises(ValueError, match="start_from must be earlier than end_before"): + WorkflowRunArchiver(start_from=ts, end_before=ts) + + def test_start_after_end_raises(self) -> None: + with pytest.raises(ValueError, match="start_from must be earlier than end_before"): + WorkflowRunArchiver( + start_from=_workflow_run_archiver_sqlite_datetime.datetime(2025, 6, 1), + end_before=_workflow_run_archiver_sqlite_datetime.datetime(2025, 1, 1), + ) + + def test_workers_zero_raises(self) -> None: + with pytest.raises(ValueError, match="workers must be at least 1"): + WorkflowRunArchiver(workers=0) + + def test_run_shard_index_without_total_raises(self) -> None: + with pytest.raises(ValueError, match="run_shard_index and run_shard_total must be provided together"): + WorkflowRunArchiver(run_shard_index=0) + + def test_run_shard_total_without_index_raises(self) -> None: + with pytest.raises(ValueError, match="run_shard_index and run_shard_total must be provided together"): + WorkflowRunArchiver(run_shard_total=4) + + def test_run_shard_total_above_supported_range_raises(self) -> None: + with pytest.raises(ValueError, match="run_shard_total must be between 1 and 16"): + WorkflowRunArchiver(run_shard_index=0, run_shard_total=17) + + def test_run_shard_index_must_be_less_than_total(self) -> None: + with pytest.raises(ValueError, match="run_shard_index must be between 0 and run_shard_total - 1"): + WorkflowRunArchiver(run_shard_index=4, run_shard_total=4) + + def test_valid_init_defaults(self) -> None: + archiver = WorkflowRunArchiver(days=30, batch_size=50) + assert archiver.days == 30 + assert archiver.batch_size == 50 + assert archiver.dry_run is False + assert archiver.delete_after_archive is False + assert archiver.start_from is None + + def test_valid_init_with_time_range(self) -> None: + start = _workflow_run_archiver_sqlite_datetime.datetime(2025, 1, 1) + end = _workflow_run_archiver_sqlite_datetime.datetime(2025, 6, 1) + archiver = WorkflowRunArchiver(start_from=start, end_before=end, workers=2) + assert archiver.start_from is not None + assert archiver.end_before is not None + assert archiver.workers == 2 + + def test_delete_after_archive_is_not_supported_for_bundle_archive(self) -> None: + with pytest.raises(ValueError, match="delete_after_archive is not supported by bundle archive"): + WorkflowRunArchiver(delete_after_archive=True) + + def test_get_runs_batch_passes_shard_options(self) -> None: + repo = MagicMock() + empty_runs: list[WorkflowRun] = [] + repo.get_runs_batch_by_time_range.return_value = empty_runs + archiver = WorkflowRunArchiver( + tenant_prefixes=["0", "a"], + run_shard_index=1, + run_shard_total=4, + workflow_run_repo=repo, + ) + + archiver._get_runs_batch(None) + + repo.get_runs_batch_by_time_range.assert_called_once() + assert repo.get_runs_batch_by_time_range.call_args.kwargs["tenant_prefixes"] == ["0", "a"] + assert repo.get_runs_batch_by_time_range.call_args.kwargs["run_shard_index"] == 1 + assert repo.get_runs_batch_by_time_range.call_args.kwargs["run_shard_total"] == 4 + + def test_get_runs_batch_prefers_planned_tenant_ids_over_prefix_filter(self) -> None: + repo = MagicMock() + empty_runs: list[WorkflowRun] = [] + repo.get_runs_batch_by_time_range.return_value = empty_runs + archiver = WorkflowRunArchiver( + tenant_ids=["0tenant"], + tenant_prefixes=["0"], + paid_tenant_ids=["0tenant"], + workflow_run_repo=repo, + ) + + archiver._get_runs_batch(None) + + repo.get_runs_batch_by_time_range.assert_called_once() + assert repo.get_runs_batch_by_time_range.call_args.kwargs["tenant_ids"] == ["0tenant"] + assert repo.get_runs_batch_by_time_range.call_args.kwargs["tenant_prefixes"] is None + + def test_get_runs_batch_uses_current_tenant_scan_scope(self) -> None: + repo = MagicMock() + empty_runs: list[WorkflowRun] = [] + repo.get_runs_batch_by_time_range.return_value = empty_runs + archiver = WorkflowRunArchiver( + tenant_ids=["tenant-a", "tenant-b"], + workflow_run_repo=repo, + ) + + archiver._get_runs_batch(None, tenant_scope=["tenant-b"]) + + repo.get_runs_batch_by_time_range.assert_called_once() + assert repo.get_runs_batch_by_time_range.call_args.kwargs["tenant_ids"] == ["tenant-b"] + + def test_get_runs_batch_retries_retryable_db_disconnect(self) -> None: + repo = MagicMock() + empty_runs: list[WorkflowRun] = [] + repo.get_runs_batch_by_time_range.side_effect = [_db_disconnect_error(), empty_runs] + archiver = WorkflowRunArchiver(workflow_run_repo=repo) + + with patch("services.retention.workflow_run.db_retry.time.sleep") as sleep: + runs = archiver._get_runs_batch(None) + + assert runs == [] + assert repo.get_runs_batch_by_time_range.call_count == 2 + sleep.assert_called_once_with(1.0) + + def test_get_runs_batch_does_not_retry_non_db_broken_pipe_error(self) -> None: + repo = MagicMock() + repo.get_runs_batch_by_time_range.side_effect = RuntimeError("broken pipe") + archiver = WorkflowRunArchiver(workflow_run_repo=repo) + + with ( + patch("services.retention.workflow_run.db_retry.time.sleep") as sleep, + pytest.raises(RuntimeError, match="broken pipe"), + ): + archiver._get_runs_batch(None) + + repo.get_runs_batch_by_time_range.assert_called_once() + sleep.assert_not_called() + + def test_start_message_includes_shard(self) -> None: + archiver = WorkflowRunArchiver(tenant_prefixes=["0"], run_shard_index=1, run_shard_total=4) + + message = archiver._build_start_message() + + assert "tenant_prefixes=0" in message + assert "run_shard=1/4" in message + + def test_start_message_summarizes_large_planned_tenant_list(self) -> None: + tenant_ids = [f"tenant-{index}" for index in range(11)] + archiver = WorkflowRunArchiver(tenant_ids=tenant_ids, tenant_prefixes=["0"]) + + message = archiver._build_start_message() + + assert "tenant_ids=11 planned tenants" in message + assert "tenant-10" not in message + + class TestBuildArchiveBundle: + def test_bundle_contains_manifest_and_all_table_objects(self) -> None: + archiver = WorkflowRunArchiver(days=90) + run = _run(str(uuid.uuid4())) + run.tenant_id = str(uuid.uuid4()) + identity = archiver._build_bundle_identity([run]) + table_data = {"workflow_runs": [{"id": run.id, "tenant_id": run.tenant_id}]} + + table_stats, table_payloads, manifest_data = archiver._build_archive_payload(identity, [run], table_data) + manifest = json.loads(manifest_data) + + assert manifest["schema_version"] == ARCHIVE_BUNDLE_SCHEMA_VERSION + assert manifest["archive_format"] == ARCHIVE_BUNDLE_FORMAT + assert manifest["object_prefix"] == identity.object_prefix + assert set(table_payloads) == set(archiver.ARCHIVED_TABLES) + assert {stat.table_name for stat in table_stats} == set(archiver.ARCHIVED_TABLES) + assert pq.read_table(pa.BufferReader(table_payloads["workflow_runs"])).num_rows == 1 + + class TestGenerateManifest: + def test_manifest_structure(self) -> None: + start = _workflow_run_archiver_sqlite_datetime.datetime( + 2025, 1, 1, tzinfo=_workflow_run_archiver_sqlite_datetime.UTC + ) + end = _workflow_run_archiver_sqlite_datetime.datetime( + 2025, 4, 1, tzinfo=_workflow_run_archiver_sqlite_datetime.UTC + ) + archiver = WorkflowRunArchiver(start_from=start, end_before=end, run_shard_index=1, run_shard_total=4) + run = _run(str(uuid.uuid4())) + run.tenant_id = str(uuid.uuid4()) + identity = archiver._build_bundle_identity([run]) + + stats = [ + TableStats( + table_name="workflow_runs", + row_count=1, + checksum="abc123", + size_bytes=512, + object_key="workflow_runs.parquet", + ), + TableStats( + table_name="workflow_node_executions", + row_count=2, + checksum="def456", + size_bytes=1024, + object_key="workflow_node_executions.parquet", + ), + ] + + manifest = archiver._generate_manifest(identity, [run], stats) + + assert manifest["schema_version"] == ARCHIVE_BUNDLE_SCHEMA_VERSION + assert manifest["archive_format"] == ARCHIVE_BUNDLE_FORMAT + assert manifest["bundle_id"] == identity.bundle_id + assert manifest["tenant_id"] == run.tenant_id + assert manifest["workflow_run_count"] == 1 + assert manifest["workflow_node_execution_count"] == 2 + assert manifest["run_ids"] == [run.id] + assert manifest["campaign_id"] == "2025-01-01T00:00:00Z_2025-04-01T00:00:00Z" + assert manifest["archive_window_start"] == "2025-01-01T00:00:00Z" + assert manifest["archive_window_end"] == "2025-04-01T00:00:00Z" + assert manifest["run_shard"] == "01-of-04" + assert "tables" in manifest + assert manifest["tables"]["workflow_runs"]["row_count"] == 1 + assert manifest["tables"]["workflow_runs"]["checksum"] == "abc123" + assert manifest["tables"]["workflow_node_executions"]["row_count"] == 2 + + class TestFilterPaidTenants: + def test_all_tenants_paid_in_community_edition(self) -> None: + archiver = WorkflowRunArchiver(days=90) + tenant_ids = {"t1", "t2", "t3"} + + with patch("services.retention.workflow_run.archive_paid_plan_workflow_run.dify_config") as cfg: + cfg.DEPLOYMENT_EDITION = DeploymentEdition.COMMUNITY + result = archiver._filter_paid_tenants(tenant_ids) + + assert result == tenant_ids + + def test_empty_tenants_returns_empty(self) -> None: + archiver = WorkflowRunArchiver(days=90) + + with patch("services.retention.workflow_run.archive_paid_plan_workflow_run.dify_config") as cfg: + cfg.DEPLOYMENT_EDITION = DeploymentEdition.CLOUD + result = archiver._filter_paid_tenants(set()) + + assert result == set() + + def test_only_paid_plans_returned(self) -> None: + archiver = WorkflowRunArchiver(days=90) + + mock_bulk = { + "t1": {"plan": "professional"}, + "t2": {"plan": "sandbox"}, + "t3": {"plan": "team"}, + } + + with ( + patch("services.retention.workflow_run.archive_paid_plan_workflow_run.dify_config") as cfg, + patch("services.retention.workflow_run.archive_paid_plan_workflow_run.BillingService") as billing, + ): + cfg.DEPLOYMENT_EDITION = DeploymentEdition.CLOUD + billing.get_plan_bulk_with_cache.return_value = mock_bulk + result = archiver._filter_paid_tenants({"t1", "t2", "t3"}) + + assert "t1" in result + assert "t3" in result + assert "t2" not in result + + def test_billing_api_failure_returns_empty(self) -> None: + archiver = WorkflowRunArchiver(days=90) + + with ( + patch("services.retention.workflow_run.archive_paid_plan_workflow_run.dify_config") as cfg, + patch("services.retention.workflow_run.archive_paid_plan_workflow_run.BillingService") as billing, + ): + cfg.DEPLOYMENT_EDITION = DeploymentEdition.CLOUD + billing.get_plan_bulk_with_cache.side_effect = RuntimeError("API down") + result = archiver._filter_paid_tenants({"t1"}) + + assert result == set() + + def test_planned_paid_tenants_skip_billing_lookup(self) -> None: + archiver = WorkflowRunArchiver(days=90, paid_tenant_ids=["t1", "t3"]) + + with ( + patch("services.retention.workflow_run.archive_paid_plan_workflow_run.dify_config") as cfg, + patch("services.retention.workflow_run.archive_paid_plan_workflow_run.BillingService") as billing, + ): + cfg.DEPLOYMENT_EDITION = DeploymentEdition.CLOUD + result = archiver._filter_paid_tenants({"t1", "t2", "t3"}) + + billing.get_plan_bulk_with_cache.assert_not_called() + assert result == {"t1", "t3"} + + class TestDryRunArchive: + @patch("services.retention.workflow_run.archive_paid_plan_workflow_run.get_archive_storage") + def test_dry_run_does_not_call_storage(self, mock_get_storage: MagicMock, sqlite_engine: Engine) -> None: + archiver = WorkflowRunArchiver(days=90, dry_run=True) + + with ( + patch( + "services.retention.workflow_run.archive_paid_plan_workflow_run.db", + SimpleNamespace(engine=sqlite_engine), + ), + patch.object(archiver, "_get_runs_batch", return_value=[]), + ): + summary = archiver.run() + + mock_get_storage.assert_not_called() + assert isinstance(summary, ArchiveSummary) + assert summary.runs_failed == 0 + + def test_dry_run_estimates_table_and_object_sizes(self, sqlite_session: Session) -> None: + archiver = WorkflowRunArchiver(days=90, dry_run=True) + run = _run() + table_data = { + "workflow_runs": [{"id": "run-1", "tenant_id": "tenant-1"}], + "workflow_app_logs": [{"id": "log-1", "workflow_run_id": "run-1"}], + } + + with patch.object(archiver, "_extract_bundle_data", return_value=table_data): + result = archiver._archive_bundle(sqlite_session, None, [run]) + + stats_by_table = {stat.table_name: stat for stat in result.tables} + assert result.success is True + assert result.object_size_bytes > 0 + assert stats_by_table["workflow_runs"].row_count == 1 + assert stats_by_table["workflow_runs"].size_bytes > 0 + assert stats_by_table["workflow_app_logs"].row_count == 1 + assert stats_by_table["workflow_app_logs"].size_bytes > 0 + assert stats_by_table["workflow_node_executions"].row_count == 0 + assert stats_by_table["workflow_node_executions"].size_bytes > 0 + + def test_summary_merges_dry_run_estimates(self) -> None: + summary = ArchiveSummary() + result = ArchiveResult( + bundle_id="bundle-1", + tenant_id="tenant-1", + object_prefix="prefix", + run_count=1, + success=True, + object_size_bytes=128, + tables=[ + TableStats(table_name="workflow_runs", row_count=1, checksum="", size_bytes=64), + TableStats(table_name="workflow_app_logs", row_count=2, checksum="", size_bytes=32), + ], + ) + + WorkflowRunArchiver._merge_result_stats(summary, result) + + assert summary.total_object_size_bytes == 128 + assert summary.table_stats["workflow_runs"].row_count == 1 + assert summary.table_stats["workflow_runs"].size_bytes == 64 + assert summary.table_stats["workflow_app_logs"].row_count == 2 + assert summary.table_stats["workflow_app_logs"].size_bytes == 32 + + class TestArchiveDbRetry: + def test_archive_bundle_groups_retries_with_fresh_session( + self, sqlite_session_factory: sessionmaker[Session] + ) -> None: + archiver = WorkflowRunArchiver(days=90) + run = _run() + session_maker = MagicMock(wraps=sqlite_session_factory) + success = ArchiveResult( + bundle_id=archiver._build_bundle_identity([run]).bundle_id, + tenant_id=run.tenant_id, + object_prefix=archiver._build_bundle_identity([run]).object_prefix, + run_count=1, + success=True, + ) + + with ( + patch.object( + archiver, "_archive_bundle", side_effect=[_db_disconnect_error(), success] + ) as archive_bundle, + patch("services.retention.workflow_run.db_retry.time.sleep") as sleep, + ): + results = archiver._archive_bundle_groups(session_maker, MagicMock(), [[run]]) + + assert results == [success] + assert archive_bundle.call_count == 2 + assert session_maker.call_count == 2 + sleep.assert_called_once_with(1.0) + + def test_archive_bundle_groups_returns_failed_result_after_retry_exhaustion( + self, sqlite_session_factory: sessionmaker[Session] + ) -> None: + archiver = WorkflowRunArchiver(days=90) + run = _run() + session_maker = MagicMock(wraps=sqlite_session_factory) + + with ( + patch.object(archiver, "_archive_bundle", side_effect=[_db_disconnect_error()] * 3) as archive_bundle, + patch("services.retention.workflow_run.db_retry.time.sleep") as sleep, + ): + results = archiver._archive_bundle_groups(session_maker, MagicMock(), [[run]]) + + assert len(results) == 1 + assert results[0].success is False + assert "server closed the connection unexpectedly" in (results[0].error or "") + assert archive_bundle.call_count == archiver.DB_RETRY_ATTEMPTS + assert session_maker.call_count == archiver.DB_RETRY_ATTEMPTS + assert sleep.call_count == archiver.DB_RETRY_ATTEMPTS - 1 + + def test_archive_bundle_uses_safe_rollback_when_failure_rolls_back_badly(self, sqlite_session: Session) -> None: + archiver = WorkflowRunArchiver(days=90, dry_run=True) + + with ( + patch.object(sqlite_session, "rollback", side_effect=RuntimeError("rollback failed")) as rollback, + patch.object(archiver, "_extract_bundle_data", side_effect=RuntimeError("extract failed")), + ): + result = archiver._archive_bundle(sqlite_session, None, [_run()]) + + assert result.success is False + assert result.error == "extract failed" + rollback.assert_called_once() + + class TestArchiveRunIdempotency: + def _index_payload( + self, archiver: WorkflowRunArchiver, run_ids: list[str], run: WorkflowRun + ) -> tuple[str, bytes]: + identity = archiver._build_bundle_identity([run]) + index_key = archiver._get_index_object_key(identity) + payload = json.dumps( + { + "schema_version": ARCHIVE_BUNDLE_SCHEMA_VERSION, + "archive_format": ARCHIVE_BUNDLE_FORMAT, + "object_prefix": archiver._get_shard_object_prefix(identity), + "updated_at": "2025-03-15T00:00:00Z", + "manifest_keys": [], + "run_ids": run_ids, + } + ).encode() + return index_key, payload + + def test_locked_bundle_is_skipped(self, sqlite_session: Session) -> None: + archiver = WorkflowRunArchiver(days=90) + run = _run() + + with ( + patch.object(archiver, "_lock_runs_for_archive", return_value=[]), + ): + storage = MagicMock() + storage.object_exists.return_value = False + result = archiver._archive_bundle(sqlite_session, storage, [run]) + + assert result.success is True + assert result.skipped is True + assert result.error == "one or more runs locked or deleted by another archiver" + + def test_already_archived_bundle_is_skipped(self, sqlite_session: Session) -> None: + archiver = WorkflowRunArchiver(days=90) + run = _run() + storage = MagicMock() + storage.object_exists.return_value = True + + with patch.object(archiver, "_sync_existing_bundle_index") as sync_existing_bundle_index: + result = archiver._archive_bundle(sqlite_session, storage, [run]) + + assert result.success is True + assert result.skipped is True + assert result.error == "bundle already archived" + sync_existing_bundle_index.assert_called_once() + + def test_existing_bundle_catalog_publication_failure_is_not_success(self, sqlite_session: Session) -> None: + archiver = WorkflowRunArchiver(days=90) + run = _run() + storage = MagicMock() + storage.object_exists.return_value = True + + with ( + patch.object(sqlite_session, "rollback", wraps=sqlite_session.rollback) as rollback, + patch.object(archiver, "_sync_existing_bundle_index", side_effect=RuntimeError("catalog unavailable")), + ): + result = archiver._archive_bundle(sqlite_session, storage, [run]) + + assert result.success is False + assert result.error == "catalog unavailable" + rollback.assert_called_once() + + def test_retry_repairs_index_after_catalog_commit_then_index_write_failure( + self, sqlite_session: Session + ) -> None: + archiver = WorkflowRunArchiver(days=90) + run = _run() + identity = archiver._build_bundle_identity([run]) + index_key = archiver._get_index_object_key(identity) + manifest_key = archiver._get_manifest_object_key(identity) + storage = FakeArchiveStorage() + original_put_object = storage.put_object + index_write_count = 0 + + def put_object(key: str, data: bytes) -> str: + nonlocal index_write_count + if key == index_key: + index_write_count += 1 + if index_write_count == 2: + raise RuntimeError("index write failed") + return original_put_object(key, data) + + storage.put_object = MagicMock(side_effect=put_object) + table_data = {"workflow_runs": [{"id": run.id, "tenant_id": run.tenant_id}]} + + with ( + patch.object(archiver, "_lock_runs_for_archive", return_value=[run]), + patch.object(archiver, "_extract_bundle_data", return_value=table_data), + ): + first_result = archiver._archive_bundle(sqlite_session, storage, [run]) + + assert first_result.success is False + assert first_result.error == "index write failed" + assert manifest_key in storage.objects + assert json.loads(storage.objects[index_key])["run_ids"] == [] + + storage.list_objects = MagicMock(wraps=storage.list_objects) + retry_result = archiver._archive_bundle(sqlite_session, storage, [run]) + + assert retry_result.success is True + assert retry_result.skipped is True + assert json.loads(storage.objects[index_key])["manifest_keys"] == [manifest_key] + assert json.loads(storage.objects[index_key])["run_ids"] == [run.id] + storage.list_objects.assert_not_called() + + def test_existing_manifest_with_missing_index_fails_without_partial_rebuild( + self, sqlite_session: Session + ) -> None: + archiver = WorkflowRunArchiver(days=90) + run = _run() + identity = archiver._build_bundle_identity([run]) + _, _, manifest_data = archiver._build_archive_payload( + identity, + [run], + {"workflow_runs": [{"id": run.id, "tenant_id": run.tenant_id}]}, + ) + manifest_key = archiver._get_manifest_object_key(identity) + index_key = archiver._get_index_object_key(identity) + storage = FakeArchiveStorage({manifest_key: manifest_data}) + storage.list_objects = MagicMock(wraps=storage.list_objects) + + result = archiver._archive_bundle(sqlite_session, storage, [run]) + + assert result.success is False + assert "archive shard index missing" in (result.error or "") + assert index_key not in storage.objects + storage.list_objects.assert_not_called() + + def test_successful_bundle_persists_archive_index(self, sqlite_session: Session) -> None: + archiver = WorkflowRunArchiver(days=90) + run = _run(str(uuid.uuid4())) + run.tenant_id = str(uuid.uuid4()) + storage = MagicMock() + storage.object_exists.return_value = False + table_data = { + "workflow_runs": [{"id": run.id, "tenant_id": run.tenant_id}], + "workflow_node_executions": [{"id": str(uuid.uuid4()), "workflow_run_id": run.id}], + } + + with ( + patch.object(archiver, "_lock_runs_for_archive", return_value=[run]), + patch.object(archiver, "_extract_bundle_data", return_value=table_data), + ): + result = archiver._archive_bundle(sqlite_session, storage, [run]) + + archived_bundle = sqlite_session.scalar( + select(WorkflowRunArchiveBundle).where(WorkflowRunArchiveBundle.tenant_id == run.tenant_id) + ) + assert result.success is True + assert isinstance(archived_bundle, WorkflowRunArchiveBundle) + assert archived_bundle.tenant_id == run.tenant_id + assert archived_bundle.year == 2025 + assert archived_bundle.month == 3 + assert archived_bundle.workflow_run_count == 1 + assert archived_bundle.row_count == 2 + + def test_new_bundle_catalog_commit_failure_is_not_success(self, sqlite_session: Session) -> None: + archiver = WorkflowRunArchiver(days=90) + run = _run(str(uuid.uuid4())) + run.tenant_id = str(uuid.uuid4()) + storage = MagicMock() + storage.object_exists.return_value = False + empty_keys: list[str] = [] + storage.list_objects.return_value = empty_keys + table_data = {"workflow_runs": [{"id": run.id, "tenant_id": run.tenant_id}]} + + with ( + patch.object(sqlite_session, "commit", side_effect=RuntimeError("catalog commit failed")), + patch.object(sqlite_session, "rollback", wraps=sqlite_session.rollback) as rollback, + patch.object(archiver, "_lock_runs_for_archive", return_value=[run]), + patch.object(archiver, "_extract_bundle_data", return_value=table_data), + ): + result = archiver._archive_bundle(sqlite_session, storage, [run]) + + assert result.success is False + assert result.error == "catalog commit failed" + rollback.assert_called_once() + + def test_index_skips_all_already_archived_runs(self, sqlite_session: Session) -> None: + archiver = WorkflowRunArchiver(days=90) + run = _run() + index_key, index_payload = self._index_payload(archiver, ["run-1"], run) + storage = FakeArchiveStorage({index_key: index_payload}) + + result = archiver._archive_bundle(sqlite_session, storage, [run]) + + assert result.success is True + assert result.skipped is True + assert result.run_count == 0 + assert result.skipped_run_count == 1 + assert result.error == "all runs already archived in shard index" + + def test_index_filters_duplicate_runs_before_archive(self, sqlite_session: Session) -> None: + archiver = WorkflowRunArchiver(days=90) + archived_run = _run() + new_run = _run("run-2") + new_run.created_at = _workflow_run_archiver_sqlite_datetime.datetime(2025, 3, 15, 11, 0, 0) + index_key, index_payload = self._index_payload(archiver, ["run-1"], archived_run) + storage = FakeArchiveStorage({index_key: index_payload}) + + with ( + patch.object(archiver, "_lock_runs_for_archive", return_value=[new_run]) as lock_runs, + patch.object(archiver, "_extract_bundle_data", return_value={"workflow_runs": [{"id": "run-2"}]}), + ): + result = archiver._archive_bundle(sqlite_session, storage, [archived_run, new_run]) + + assert result.success is True + assert result.skipped is False + assert result.run_count == 1 + assert result.skipped_run_count == 1 + lock_runs.assert_called_once_with(ANY, ["run-2"]) + manifest_keys = [key for key in storage.objects if key.endswith("/manifest.json")] + assert len(manifest_keys) == 1 + + +""" +Unit tests for workflow run archiving functionality. + +This module contains tests for: +- Archive service +- Rollback service +""" + + +class TestArchiveWorkflowRunLogsCases: + class TestWorkflowRunArchiver: + """Tests for the WorkflowRunArchiver class.""" + + def test_archiver_initialization(self, config_overrides: Callable[..., None]) -> None: + """Test archiver can be initialized with various options.""" + from services.retention.workflow_run.archive_paid_plan_workflow_run import WorkflowRunArchiver + + config_overrides(DEPLOYMENT_EDITION=DeploymentEdition.COMMUNITY) + + archiver = WorkflowRunArchiver( + days=90, + batch_size=100, + tenant_ids=["test-tenant"], + limit=50, + dry_run=True, + ) + + assert archiver.days == 90 + assert archiver.batch_size == 100 + assert archiver.tenant_ids == ["test-tenant"] + assert archiver.limit == 50 + assert archiver.dry_run is True + + def test_get_bundle_manifest_key(self) -> None: + """Test V2 bundle manifest key generation.""" + from services.retention.workflow_run.archive_paid_plan_workflow_run import WorkflowRunArchiver + + archiver = WorkflowRunArchiver(run_shard_index=1, run_shard_total=4) + + mock_run = MagicMock() + mock_run.tenant_id = "9enant-123" + mock_run.id = "run-456" + mock_run.created_at = _archive_workflow_run_logs_datetime(2024, 1, 15, 12, 0, 0) + + identity = archiver._build_bundle_identity([mock_run]) + key = archiver._get_manifest_object_key(identity) + + assert key.endswith("/manifest.json") + assert "workflow-runs/v2/tenant_prefix=9/tenant_id=9enant-123/year=2024/month=01" in key + assert "/shard=01-of-04/" in key diff --git a/api/tests/unit_tests/services/retention/workflow_run/test_bundle_archive_maintenance.py b/api/tests/unit_tests/services/retention/workflow_run/test_bundle_archive_maintenance.py index 1e9c17814b6..5ba4935435f 100644 --- a/api/tests/unit_tests/services/retention/workflow_run/test_bundle_archive_maintenance.py +++ b/api/tests/unit_tests/services/retention/workflow_run/test_bundle_archive_maintenance.py @@ -29,7 +29,7 @@ from services.retention.workflow_run.constants import ( ARCHIVE_BUNDLE_RESTORED_MARKER_NAME, ARCHIVE_BUNDLE_SCHEMA_VERSION, ) -from tests.unit_tests.config_override import apply_config_overrides +from tests.helpers.unit.config_override import apply_config_overrides TENANT_ID = "1251fe32-c0c7-4fe2-a7bd-a8105267faf5" CATALOG_ID = "019f63b7-5ca4-7681-9ce0-800283608f39" diff --git a/api/tests/unit_tests/services/retention/workflow_run/test_clear_free_plan_expired_workflow_run_logs.py b/api/tests/unit_tests/services/retention/workflow_run/test_clear_free_plan_expired_workflow_run_logs.py index 5a25ee8c452..db60c383ecc 100644 --- a/api/tests/unit_tests/services/retention/workflow_run/test_clear_free_plan_expired_workflow_run_logs.py +++ b/api/tests/unit_tests/services/retention/workflow_run/test_clear_free_plan_expired_workflow_run_logs.py @@ -1,8 +1,6 @@ -""" -Unit tests for WorkflowRunCleanup service. -""" - import datetime +from collections.abc import Callable +from typing import Any from unittest.mock import MagicMock, patch import pytest @@ -10,10 +8,18 @@ from sqlalchemy.orm import Session from enums import CloudPlan, DeploymentEdition from repositories.api_workflow_run_repository import WorkflowRunCleanupRef +from services.billing_service import SubscriptionPlan +from services.retention.workflow_run import clear_free_plan_expired_workflow_run_logs as cleanup_module from services.retention.workflow_run.clear_free_plan_expired_workflow_run_logs import WorkflowRunCleanup +""" +Unit tests for WorkflowRunCleanup service. +""" -def make_ref(tenant_id: str = "t1", run_id: str = "r1", created_at: datetime.datetime | None = None): + +def _workflow_run_clear_free_plan_expired_workflow_run_logs_make_ref( + tenant_id: str = "t1", run_id: str = "r1", created_at: datetime.datetime | None = None +): return WorkflowRunCleanupRef( id=run_id, tenant_id=tenant_id, @@ -21,431 +27,929 @@ def make_ref(tenant_id: str = "t1", run_id: str = "r1", created_at: datetime.dat ) -@pytest.fixture -def mock_repo(): - return MagicMock() +class TestClearFreePlanExpiredWorkflowRunLogsCases: + @staticmethod + @pytest.fixture + def mock_repo(): + return MagicMock() + @staticmethod + @pytest.fixture(autouse=True) + def _cleanup_config(config_overrides): + config_overrides( + SANDBOX_EXPIRED_RECORDS_CLEAN_GRACEFUL_PERIOD=0, + DEPLOYMENT_EDITION=DeploymentEdition.COMMUNITY, + ) -@pytest.fixture(autouse=True) -def _cleanup_config(config_overrides): - config_overrides( - SANDBOX_EXPIRED_RECORDS_CLEAN_GRACEFUL_PERIOD=0, - DEPLOYMENT_EDITION=DeploymentEdition.COMMUNITY, - ) + @staticmethod + @pytest.fixture + def cleanup(mock_repo): + return WorkflowRunCleanup(days=30, batch_size=10, workflow_run_repo=mock_repo) + # --------------------------------------------------------------------------- + # Constructor validation + # --------------------------------------------------------------------------- -@pytest.fixture -def cleanup(mock_repo): - return WorkflowRunCleanup(days=30, batch_size=10, workflow_run_repo=mock_repo) + class TestWorkflowRunCleanupInit: + def test_only_start_from_raises(self, mock_repo): + with pytest.raises(ValueError, match="both set or both omitted"): + WorkflowRunCleanup( + days=30, + batch_size=10, + start_from=datetime.datetime(2024, 1, 1), + workflow_run_repo=mock_repo, + ) + def test_only_end_before_raises(self, mock_repo): + with pytest.raises(ValueError, match="both set or both omitted"): + WorkflowRunCleanup( + days=30, + batch_size=10, + end_before=datetime.datetime(2024, 1, 1), + workflow_run_repo=mock_repo, + ) -# --------------------------------------------------------------------------- -# Constructor validation -# --------------------------------------------------------------------------- + def test_end_before_not_greater_than_start_raises(self, mock_repo): + with pytest.raises(ValueError, match="end_before must be greater than start_from"): + WorkflowRunCleanup( + days=30, + batch_size=10, + start_from=datetime.datetime(2024, 6, 1), + end_before=datetime.datetime(2024, 1, 1), + workflow_run_repo=mock_repo, + ) + def test_equal_start_end_raises(self, mock_repo): + dt = datetime.datetime(2024, 1, 1) + with pytest.raises(ValueError, match="end_before must be greater than start_from\\."): + WorkflowRunCleanup( + days=30, + batch_size=10, + start_from=dt, + end_before=dt, + workflow_run_repo=mock_repo, + ) -class TestWorkflowRunCleanupInit: - def test_only_start_from_raises(self, mock_repo): - with pytest.raises(ValueError, match="both set or both omitted"): - WorkflowRunCleanup( + def test_zero_batch_size_raises(self, mock_repo): + with pytest.raises(ValueError, match="batch_size must be greater than 0"): + WorkflowRunCleanup(days=30, batch_size=0, workflow_run_repo=mock_repo) + + def test_negative_batch_size_raises(self, mock_repo): + with pytest.raises(ValueError, match="batch_size must be greater than 0\\."): + WorkflowRunCleanup(days=30, batch_size=-1, workflow_run_repo=mock_repo) + + def test_valid_window_init(self, mock_repo, config_overrides): + config_overrides(SANDBOX_EXPIRED_RECORDS_CLEAN_GRACEFUL_PERIOD=7) + start = datetime.datetime(2024, 1, 1) + end = datetime.datetime(2024, 6, 1) + c = WorkflowRunCleanup( + days=30, + batch_size=5, + start_from=start, + end_before=end, + workflow_run_repo=mock_repo, + ) + assert c.window_start == start + assert c.window_end == end + + def test_default_task_label_is_custom(self, mock_repo): + c = WorkflowRunCleanup(days=30, batch_size=10, workflow_run_repo=mock_repo) + + assert c._metrics._base_attributes["task_label"] == "custom" + + # --------------------------------------------------------------------------- + # _empty_related_counts / _format_related_counts + # --------------------------------------------------------------------------- + + class TestStaticHelpers: + def test_empty_related_counts(self): + counts = WorkflowRunCleanup._empty_related_counts() + assert counts == { + "node_executions": 0, + "offloads": 0, + "app_logs": 0, + "trigger_logs": 0, + "pauses": 0, + "pause_reasons": 0, + } + + def test_format_related_counts(self): + counts = { + "node_executions": 1, + "offloads": 2, + "app_logs": 3, + "trigger_logs": 4, + "pauses": 5, + "pause_reasons": 6, + } + result = WorkflowRunCleanup._format_related_counts(counts) + assert "node_executions 1" in result + assert "offloads 2" in result + assert "trigger_logs 4" in result + + # --------------------------------------------------------------------------- + # _expiration_datetime + # --------------------------------------------------------------------------- + + class TestExpirationDatetime: + def test_negative_returns_none(self, cleanup): + assert cleanup._expiration_datetime("t1", -1) is None + + def test_valid_timestamp(self, cleanup): + ts = int(datetime.datetime(2025, 1, 1, tzinfo=datetime.UTC).timestamp()) + result = cleanup._expiration_datetime("t1", ts) + assert result is not None + assert result.year == 2025 + + def test_overflow_returns_none(self, cleanup): + result = cleanup._expiration_datetime("t1", 2**62) + assert result is None + + # --------------------------------------------------------------------------- + # _is_within_grace_period + # --------------------------------------------------------------------------- + + class TestIsWithinGracePeriod: + def test_zero_grace_period_returns_false(self, cleanup): + cleanup.free_plan_grace_period_days = 0 + assert cleanup._is_within_grace_period("t1", {"expiration_date": 9999999999}) is False + + def test_within_grace_period(self, cleanup): + cleanup.free_plan_grace_period_days = 30 + # expired just 1 day ago + expired = datetime.datetime.now(datetime.UTC) - datetime.timedelta(days=1) + ts = int(expired.timestamp()) + assert cleanup._is_within_grace_period("t1", {"expiration_date": ts}) is True + + def test_outside_grace_period(self, cleanup): + cleanup.free_plan_grace_period_days = 5 + # expired 100 days ago + expired = datetime.datetime.now(datetime.UTC) - datetime.timedelta(days=100) + ts = int(expired.timestamp()) + assert cleanup._is_within_grace_period("t1", {"expiration_date": ts}) is False + + def test_missing_expiration_date_returns_false(self, cleanup): + cleanup.free_plan_grace_period_days = 30 + assert cleanup._is_within_grace_period("t1", {"expiration_date": -1}) is False + + # --------------------------------------------------------------------------- + # _get_cleanup_whitelist + # --------------------------------------------------------------------------- + + class TestGetCleanupWhitelist: + def test_non_cloud_edition_returns_empty(self, cleanup): + cleanup._cleanup_whitelist = None + result = cleanup._get_cleanup_whitelist() + assert result == set() + + def test_cloud_edition_fetches_whitelist(self, mock_repo, config_overrides): + config_overrides(DEPLOYMENT_EDITION=DeploymentEdition.CLOUD) + c = WorkflowRunCleanup(days=30, batch_size=10, workflow_run_repo=mock_repo) + with patch( + "services.retention.workflow_run.clear_free_plan_expired_workflow_run_logs.BillingService" + ) as bs: + bs.get_expired_subscription_cleanup_whitelist.return_value = ["t1", "t2"] + result = c._get_cleanup_whitelist() + assert result == {"t1", "t2"} + + def test_cached_whitelist_returned(self, cleanup): + cleanup._cleanup_whitelist = {"cached"} + result = cleanup._get_cleanup_whitelist() + assert result == {"cached"} + + def test_billing_service_error_returns_empty(self, mock_repo, config_overrides): + config_overrides(DEPLOYMENT_EDITION=DeploymentEdition.CLOUD) + c = WorkflowRunCleanup(days=30, batch_size=10, workflow_run_repo=mock_repo) + with patch( + "services.retention.workflow_run.clear_free_plan_expired_workflow_run_logs.BillingService" + ) as bs: + bs.get_expired_subscription_cleanup_whitelist.side_effect = Exception("error") + result = c._get_cleanup_whitelist() + assert result == set() + + # --------------------------------------------------------------------------- + # _filter_free_tenants + # --------------------------------------------------------------------------- + + class TestFilterFreeTenants: + def test_non_cloud_edition_treats_all_tenants_as_free(self, cleanup): + result = cleanup._filter_free_tenants(["t1", "t2"]) + assert result == {"t1", "t2"} + + def test_empty_tenants_returns_empty(self, cleanup, config_overrides): + config_overrides(DEPLOYMENT_EDITION=DeploymentEdition.CLOUD) + result = cleanup._filter_free_tenants([]) + assert result == set() + + def test_whitelisted_tenant_excluded(self, mock_repo, config_overrides): + config_overrides(DEPLOYMENT_EDITION=DeploymentEdition.CLOUD) + c = WorkflowRunCleanup(days=30, batch_size=10, workflow_run_repo=mock_repo) + c._cleanup_whitelist = {"t1"} + with patch( + "services.retention.workflow_run.clear_free_plan_expired_workflow_run_logs.BillingService" + ) as bs: + bs.get_plan_bulk_with_cache.return_value = { + "t1": {"plan": CloudPlan.SANDBOX, "expiration_date": -1}, + "t2": {"plan": CloudPlan.SANDBOX, "expiration_date": -1}, + } + result = c._filter_free_tenants(["t1", "t2"]) + assert "t1" not in result + assert "t2" in result + + def test_paid_tenant_excluded(self, mock_repo, config_overrides): + config_overrides(DEPLOYMENT_EDITION=DeploymentEdition.CLOUD) + c = WorkflowRunCleanup(days=30, batch_size=10, workflow_run_repo=mock_repo) + c._cleanup_whitelist = set() + with patch( + "services.retention.workflow_run.clear_free_plan_expired_workflow_run_logs.BillingService" + ) as bs: + bs.get_plan_bulk_with_cache.return_value = { + "t1": {"plan": "professional", "expiration_date": -1}, + } + result = c._filter_free_tenants(["t1"]) + assert result == set() + + def test_missing_billing_info_treats_as_non_free(self, mock_repo, config_overrides): + config_overrides(DEPLOYMENT_EDITION=DeploymentEdition.CLOUD) + c = WorkflowRunCleanup(days=30, batch_size=10, workflow_run_repo=mock_repo) + c._cleanup_whitelist = set() + with patch( + "services.retention.workflow_run.clear_free_plan_expired_workflow_run_logs.BillingService" + ) as bs: + bs.get_plan_bulk_with_cache.return_value = {} + result = c._filter_free_tenants(["t1"]) + assert result == set() + + def test_billing_bulk_error_treats_as_non_free(self, mock_repo, config_overrides): + config_overrides(DEPLOYMENT_EDITION=DeploymentEdition.CLOUD) + c = WorkflowRunCleanup(days=30, batch_size=10, workflow_run_repo=mock_repo) + c._cleanup_whitelist = set() + with patch( + "services.retention.workflow_run.clear_free_plan_expired_workflow_run_logs.BillingService" + ) as bs: + bs.get_plan_bulk_with_cache.side_effect = Exception("fail") + result = c._filter_free_tenants(["t1"]) + assert result == set() + + # --------------------------------------------------------------------------- + # run() — delete mode + # --------------------------------------------------------------------------- + + class TestRunDeleteMode: + def _make_cleanup(self, mock_repo): + return WorkflowRunCleanup(days=30, batch_size=10, workflow_run_repo=mock_repo) + + def test_no_rows_stops_immediately(self, mock_repo): + mock_repo.get_cleanup_refs_batch_by_time_range.return_value = [] + c = self._make_cleanup(mock_repo) + c.run() + mock_repo.delete_runs_with_related_by_ids.assert_not_called() + + def test_all_paid_skips_delete(self, mock_repo): + ref = _workflow_run_clear_free_plan_expired_workflow_run_logs_make_ref("t1") + mock_repo.get_cleanup_refs_batch_by_time_range.side_effect = [[ref], []] + c = self._make_cleanup(mock_repo) + # Override the non-Cloud default to exercise the no-deletion path. + c._filter_free_tenants = MagicMock(return_value=set()) + c.run() + mock_repo.delete_runs_with_related_by_ids.assert_not_called() + + def test_runs_deleted_successfully(self, mock_repo): + ref = _workflow_run_clear_free_plan_expired_workflow_run_logs_make_ref("t1") + mock_repo.get_cleanup_refs_batch_by_time_range.side_effect = [[ref], [ref], []] + mock_repo.delete_runs_with_related_by_ids.return_value = { + "runs": 1, + "node_executions": 0, + "offloads": 0, + "app_logs": 0, + "trigger_logs": 0, + "pauses": 0, + "pause_reasons": 0, + } + c = self._make_cleanup(mock_repo) + with patch("services.retention.workflow_run.clear_free_plan_expired_workflow_run_logs.time.sleep"): + c.run() + mock_repo.delete_runs_with_related_by_ids.assert_called_once() + + def test_delete_exception_reraises(self, mock_repo): + ref = _workflow_run_clear_free_plan_expired_workflow_run_logs_make_ref("t1") + mock_repo.get_cleanup_refs_batch_by_time_range.side_effect = [[ref], [ref]] + mock_repo.delete_runs_with_related_by_ids.side_effect = RuntimeError("db error") + c = self._make_cleanup(mock_repo) + with pytest.raises(RuntimeError): + c.run() + + def test_summary_with_window_start(self, mock_repo): + mock_repo.get_cleanup_refs_batch_by_time_range.return_value = [] + c = WorkflowRunCleanup( days=30, batch_size=10, start_from=datetime.datetime(2024, 1, 1), + end_before=datetime.datetime(2024, 6, 1), workflow_run_repo=mock_repo, ) + c.run() - def test_only_end_before_raises(self, mock_repo): - with pytest.raises(ValueError, match="both set or both omitted"): - WorkflowRunCleanup( + # --------------------------------------------------------------------------- + # run() — dry run mode + # --------------------------------------------------------------------------- + + class TestRunDryRunMode: + def _make_dry_cleanup(self, mock_repo): + return WorkflowRunCleanup( days=30, batch_size=10, - end_before=datetime.datetime(2024, 1, 1), workflow_run_repo=mock_repo, + dry_run=True, ) - def test_end_before_not_greater_than_start_raises(self, mock_repo): - with pytest.raises(ValueError, match="end_before must be greater than start_from"): - WorkflowRunCleanup( + def test_dry_run_no_delete_called(self, mock_repo): + ref = _workflow_run_clear_free_plan_expired_workflow_run_logs_make_ref("t1") + mock_repo.get_cleanup_refs_batch_by_time_range.side_effect = [[ref], [ref], []] + mock_repo.count_runs_with_related_by_ids.return_value = { + "runs": 1, + "node_executions": 2, + "offloads": 0, + "app_logs": 0, + "trigger_logs": 1, + "pauses": 0, + "pause_reasons": 0, + } + c = self._make_dry_cleanup(mock_repo) + c.run() + mock_repo.delete_runs_with_related_by_ids.assert_not_called() + mock_repo.count_runs_with_related_by_ids.assert_called_once() + + def test_dry_run_summary_with_window_start(self, mock_repo): + mock_repo.get_cleanup_refs_batch_by_time_range.return_value = [] + c = WorkflowRunCleanup( days=30, batch_size=10, - start_from=datetime.datetime(2024, 6, 1), - end_before=datetime.datetime(2024, 1, 1), + start_from=datetime.datetime(2024, 1, 1), + end_before=datetime.datetime(2024, 6, 1), workflow_run_repo=mock_repo, + dry_run=True, + ) + c.run() + + def test_dry_run_all_paid_skips_count(self, mock_repo): + ref = _workflow_run_clear_free_plan_expired_workflow_run_logs_make_ref("t1") + mock_repo.get_cleanup_refs_batch_by_time_range.side_effect = [[ref], []] + c = self._make_dry_cleanup(mock_repo) + c._filter_free_tenants = MagicMock(return_value=set()) + c.run() + mock_repo.count_runs_with_related_by_ids.assert_not_called() + + # --------------------------------------------------------------------------- + # _delete_trigger_logs / _count_trigger_logs + # --------------------------------------------------------------------------- + + class TestTriggerLogMethods: + @pytest.mark.parametrize("sqlite_session", [()], indirect=True) + def test_delete_trigger_logs(self, cleanup, sqlite_session: Session): + with patch( + "services.retention.workflow_run.clear_free_plan_expired_workflow_run_logs.SQLAlchemyWorkflowTriggerLogRepository" + ) as repo_class: + instance = repo_class.return_value + instance.delete_by_run_ids.return_value = 5 + result = cleanup._delete_trigger_logs(sqlite_session, ["r1", "r2"]) + assert result == 5 + + @pytest.mark.parametrize("sqlite_session", [()], indirect=True) + def test_count_trigger_logs(self, cleanup, sqlite_session: Session): + with patch( + "services.retention.workflow_run.clear_free_plan_expired_workflow_run_logs.SQLAlchemyWorkflowTriggerLogRepository" + ) as repo_class: + instance = repo_class.return_value + instance.count_by_run_ids.return_value = 3 + result = cleanup._count_trigger_logs(sqlite_session, ["r1"]) + assert result == 3 + + # --------------------------------------------------------------------------- + # _count_node_executions_by_run_ids / _delete_node_executions_by_run_ids + # --------------------------------------------------------------------------- + + class TestNodeExecutionMethods: + @pytest.mark.parametrize("sqlite_session", [()], indirect=True) + def test_count_node_executions(self, cleanup, sqlite_session: Session): + with patch( + "services.retention.workflow_run.clear_free_plan_expired_workflow_run_logs.DifyAPIRepositoryFactory" + ) as factory: + repo = factory.create_api_workflow_node_execution_repository.return_value + repo.count_by_runs.return_value = (10, 2) + result = cleanup._count_node_executions_by_run_ids(sqlite_session, ["r1"]) + assert result == (10, 2) + + @pytest.mark.parametrize("sqlite_session", [()], indirect=True) + def test_delete_node_executions(self, cleanup, sqlite_session: Session): + with patch( + "services.retention.workflow_run.clear_free_plan_expired_workflow_run_logs.DifyAPIRepositoryFactory" + ) as factory: + repo = factory.create_api_workflow_node_execution_repository.return_value + repo.delete_by_runs.return_value = (5, 1) + result = cleanup._delete_node_executions_by_run_ids(sqlite_session, ["r1"]) + assert result == (5, 1) + + +def _services_clear_free_plan_expired_workflow_run_logs_make_ref( + run_id: str, tenant_id: str, created_at: datetime.datetime +) -> WorkflowRunCleanupRef: + return WorkflowRunCleanupRef(id=run_id, tenant_id=tenant_id, created_at=created_at) + + +class FakeRepo: + def __init__( + self, + batches: list[list[WorkflowRunCleanupRef]], + delete_result: dict[str, int] | None = None, + count_result: dict[str, int] | None = None, + ) -> None: + self.batches = batches + self.candidate_call_idx = 0 + self.last_candidate_batch: list[WorkflowRunCleanupRef] = [] + self.cleanup_ref_calls: list[dict[str, object]] = [] + self.deleted: list[list[str]] = [] + self.counted: list[list[str]] = [] + self.delete_result = delete_result or { + "runs": 0, + "node_executions": 0, + "offloads": 0, + "app_logs": 0, + "trigger_logs": 0, + "pauses": 0, + "pause_reasons": 0, + } + self.count_result = count_result or { + "runs": 0, + "node_executions": 0, + "offloads": 0, + "app_logs": 0, + "trigger_logs": 0, + "pauses": 0, + "pause_reasons": 0, + } + + def get_cleanup_refs_batch_by_time_range( + self, + start_from: datetime.datetime | None, + end_before: datetime.datetime, + last_seen: tuple[datetime.datetime, str] | None, + batch_size: int, + run_types=None, + tenant_ids=None, + workflow_ids=None, + upper_bound: tuple[datetime.datetime, str] | None = None, + ) -> list[WorkflowRunCleanupRef]: + self.cleanup_ref_calls.append( + { + "start_from": start_from, + "end_before": end_before, + "last_seen": last_seen, + "batch_size": batch_size, + "run_types": run_types, + "tenant_ids": tenant_ids, + "workflow_ids": workflow_ids, + "upper_bound": upper_bound, + } + ) + if tenant_ids is not None or upper_bound is not None: + refs = self.last_candidate_batch + if tenant_ids is not None: + tenant_id_set = set(tenant_ids) + refs = [ref for ref in refs if ref.tenant_id in tenant_id_set] + if upper_bound is not None: + refs = [ref for ref in refs if (ref.created_at, ref.id) <= upper_bound] + return refs[:batch_size] + + if self.candidate_call_idx >= len(self.batches): + return [] + batch = self.batches[self.candidate_call_idx] + self.candidate_call_idx += 1 + self.last_candidate_batch = batch + return batch + + def delete_runs_with_related_by_ids( + self, run_ids: list[str], delete_node_executions=None, delete_trigger_logs=None + ) -> dict[str, int]: + self.deleted.append(list(run_ids)) + result = self.delete_result.copy() + result["runs"] = len(run_ids) + return result + + def count_runs_with_related_by_ids( + self, run_ids: list[str], count_node_executions=None, count_trigger_logs=None + ) -> dict[str, int]: + self.counted.append(list(run_ids)) + result = self.count_result.copy() + result["runs"] = len(run_ids) + return result + + +def plan_info(plan: str, expiration: int) -> SubscriptionPlan: + return SubscriptionPlan(plan=plan, expiration_date=expiration) + + +def create_cleanup( + monkeypatch: pytest.MonkeyPatch, + repo: FakeRepo, + *, + whitelist: set[str] | None = None, + **kwargs: Any, +) -> WorkflowRunCleanup: + monkeypatch.setattr( + cleanup_module.WorkflowRunCleanup, + "_get_cleanup_whitelist", + lambda self: whitelist or set(), + ) + return WorkflowRunCleanup(workflow_run_repo=repo, **kwargs) + + +class TestClearFreePlanExpiredWorkflowRunLogsCasesServices: + @staticmethod + @pytest.fixture(autouse=True) + def _cleanup_config(config_overrides: Callable[..., None]) -> None: + config_overrides( + DEPLOYMENT_EDITION=DeploymentEdition.COMMUNITY, + SANDBOX_EXPIRED_RECORDS_CLEAN_GRACEFUL_PERIOD=0, + ) + + @staticmethod + def test_filter_free_tenants_outside_cloud_edition(monkeypatch: pytest.MonkeyPatch) -> None: + cleanup = create_cleanup(monkeypatch, repo=FakeRepo([]), days=30, batch_size=10) + + def fail_bulk(_: list[str]) -> dict[str, SubscriptionPlan]: + raise RuntimeError("should not call") + + monkeypatch.setattr(cleanup_module.BillingService, "get_plan_bulk_with_cache", staticmethod(fail_bulk)) + + tenants = {"t1", "t2"} + free = cleanup._filter_free_tenants(tenants) + + assert free == tenants + + @staticmethod + def test_filter_free_tenants_bulk_mixed( + monkeypatch: pytest.MonkeyPatch, config_overrides: Callable[..., None] + ) -> None: + config_overrides(DEPLOYMENT_EDITION=DeploymentEdition.CLOUD) + cleanup = create_cleanup(monkeypatch, repo=FakeRepo([]), days=30, batch_size=10) + + monkeypatch.setattr( + cleanup_module.BillingService, + "get_plan_bulk_with_cache", + staticmethod( + lambda tenant_ids: { + tenant_id: (plan_info("team", -1) if tenant_id == "t_paid" else plan_info("sandbox", -1)) + for tenant_id in tenant_ids + } + ), + ) + + free = cleanup._filter_free_tenants({"t_free", "t_paid", "t_missing"}) + + assert free == {"t_free", "t_missing"} + + @staticmethod + def test_filter_free_tenants_respects_grace_period( + monkeypatch: pytest.MonkeyPatch, config_overrides: Callable[..., None] + ) -> None: + config_overrides( + DEPLOYMENT_EDITION=DeploymentEdition.CLOUD, + SANDBOX_EXPIRED_RECORDS_CLEAN_GRACEFUL_PERIOD=45, + ) + cleanup = create_cleanup(monkeypatch, repo=FakeRepo([]), days=30, batch_size=10) + + now = datetime.datetime.now(datetime.UTC) + within_grace_ts = int((now - datetime.timedelta(days=10)).timestamp()) + outside_grace_ts = int((now - datetime.timedelta(days=90)).timestamp()) + + def fake_bulk(_: list[str]) -> dict[str, SubscriptionPlan]: + return { + "recently_downgraded": plan_info("sandbox", within_grace_ts), + "long_sandbox": plan_info("sandbox", outside_grace_ts), + } + + monkeypatch.setattr(cleanup_module.BillingService, "get_plan_bulk_with_cache", staticmethod(fake_bulk)) + + free = cleanup._filter_free_tenants({"recently_downgraded", "long_sandbox"}) + + assert free == {"long_sandbox"} + + @staticmethod + def test_filter_free_tenants_skips_cleanup_whitelist( + monkeypatch: pytest.MonkeyPatch, config_overrides: Callable[..., None] + ) -> None: + config_overrides(DEPLOYMENT_EDITION=DeploymentEdition.CLOUD) + cleanup = create_cleanup( + monkeypatch, + repo=FakeRepo([]), + days=30, + batch_size=10, + whitelist={"tenant_whitelist"}, + ) + + monkeypatch.setattr( + cleanup_module.BillingService, + "get_plan_bulk_with_cache", + staticmethod( + lambda tenant_ids: { + tenant_id: (plan_info("team", -1) if tenant_id == "t_paid" else plan_info("sandbox", -1)) + for tenant_id in tenant_ids + } + ), + ) + + tenants = {"tenant_whitelist", "tenant_regular"} + free = cleanup._filter_free_tenants(tenants) + + assert free == {"tenant_regular"} + + @staticmethod + def test_filter_free_tenants_bulk_failure( + monkeypatch: pytest.MonkeyPatch, config_overrides: Callable[..., None] + ) -> None: + config_overrides(DEPLOYMENT_EDITION=DeploymentEdition.CLOUD) + cleanup = create_cleanup(monkeypatch, repo=FakeRepo([]), days=30, batch_size=10) + + monkeypatch.setattr( + cleanup_module.BillingService, + "get_plan_bulk_with_cache", + staticmethod(lambda tenant_ids: (_ for _ in ()).throw(RuntimeError("boom"))), + ) + + free = cleanup._filter_free_tenants({"t1", "t2"}) + + assert free == set() + + @staticmethod + def test_run_deletes_only_free_tenants( + monkeypatch: pytest.MonkeyPatch, config_overrides: Callable[..., None] + ) -> None: + config_overrides(DEPLOYMENT_EDITION=DeploymentEdition.CLOUD) + cutoff = datetime.datetime.now() + repo = FakeRepo( + batches=[ + [ + _services_clear_free_plan_expired_workflow_run_logs_make_ref("run-free", "t_free", cutoff), + _services_clear_free_plan_expired_workflow_run_logs_make_ref("run-paid", "t_paid", cutoff), + ] + ] + ) + cleanup = create_cleanup(monkeypatch, repo=repo, days=30, batch_size=10) + + monkeypatch.setattr( + cleanup_module.BillingService, + "get_plan_bulk_with_cache", + staticmethod( + lambda tenant_ids: { + tenant_id: (plan_info("team", -1) if tenant_id == "t_paid" else plan_info("sandbox", -1)) + for tenant_id in tenant_ids + } + ), + ) + + cleanup.run() + + assert repo.deleted == [["run-free"]] + assert repo.cleanup_ref_calls[1]["tenant_ids"] == ["t_free"] + + @staticmethod + def test_run_filters_candidate_tenants_before_target_query( + monkeypatch: pytest.MonkeyPatch, config_overrides: Callable[..., None] + ) -> None: + config_overrides(DEPLOYMENT_EDITION=DeploymentEdition.CLOUD) + cutoff = datetime.datetime.now() + repo = FakeRepo( + batches=[ + [ + _services_clear_free_plan_expired_workflow_run_logs_make_ref("run-free", "t_free", cutoff), + _services_clear_free_plan_expired_workflow_run_logs_make_ref("run-paid", "t_paid", cutoff), + ] + ] + ) + cleanup = create_cleanup(monkeypatch, repo=repo, days=30, batch_size=10) + + billing_calls: list[list[str]] = [] + + def fake_bulk(tenant_ids: list[str]) -> dict[str, SubscriptionPlan]: + billing_calls.append(tenant_ids) + return { + "t_free": plan_info("sandbox", -1), + "t_paid": plan_info("team", -1), + } + + monkeypatch.setattr(cleanup_module.BillingService, "get_plan_bulk_with_cache", staticmethod(fake_bulk)) + + cleanup.run() + + assert billing_calls == [["t_free", "t_paid"]] + assert repo.cleanup_ref_calls[1]["tenant_ids"] == ["t_free"] + assert repo.deleted == [["run-free"]] + + @staticmethod + def test_run_skips_when_no_free_tenants( + monkeypatch: pytest.MonkeyPatch, config_overrides: Callable[..., None] + ) -> None: + config_overrides(DEPLOYMENT_EDITION=DeploymentEdition.CLOUD) + cutoff = datetime.datetime.now() + repo = FakeRepo( + batches=[[_services_clear_free_plan_expired_workflow_run_logs_make_ref("run-paid", "t_paid", cutoff)]] + ) + cleanup = create_cleanup(monkeypatch, repo=repo, days=30, batch_size=10) + + monkeypatch.setattr( + cleanup_module.BillingService, + "get_plan_bulk_with_cache", + staticmethod(lambda tenant_ids: {tenant_id: plan_info("team", 1893456000) for tenant_id in tenant_ids}), + ) + + cleanup.run() + + assert repo.deleted == [] + assert len(repo.cleanup_ref_calls) == 2 + + @staticmethod + def test_run_paid_only_records_skipped_metrics( + monkeypatch: pytest.MonkeyPatch, config_overrides: Callable[..., None] + ) -> None: + config_overrides(DEPLOYMENT_EDITION=DeploymentEdition.CLOUD) + cutoff = datetime.datetime.now() + repo = FakeRepo( + batches=[[_services_clear_free_plan_expired_workflow_run_logs_make_ref("run-paid", "t_paid", cutoff)]] + ) + cleanup = create_cleanup(monkeypatch, repo=repo, days=30, batch_size=10) + + monkeypatch.setattr( + cleanup_module.BillingService, + "get_plan_bulk_with_cache", + staticmethod(lambda tenant_ids: {tenant_id: plan_info("team", 1893456000) for tenant_id in tenant_ids}), + ) + batch_calls: list[dict[str, object]] = [] + monkeypatch.setattr(cleanup._metrics, "record_batch", lambda **kwargs: batch_calls.append(kwargs)) + + cleanup.run() + + assert repo.deleted == [] + assert repo.counted == [] + assert batch_calls[0]["batch_rows"] == 1 + assert batch_calls[0]["targeted_runs"] == 0 + assert batch_calls[0]["skipped_runs"] == 1 + assert batch_calls[0]["deleted_runs"] == 0 + + @staticmethod + def test_run_target_query_is_bounded_by_candidate_high_water(monkeypatch: pytest.MonkeyPatch) -> None: + first_created_at = datetime.datetime(2024, 1, 1, 0, 0, 0) + second_created_at = datetime.datetime(2024, 1, 1, 0, 1, 0) + repo = FakeRepo( + batches=[ + [ + _services_clear_free_plan_expired_workflow_run_logs_make_ref( + "run-free-1", "t_free", first_created_at + ), + _services_clear_free_plan_expired_workflow_run_logs_make_ref( + "run-free-2", "t_free", second_created_at + ), + ] + ] + ) + cleanup = create_cleanup(monkeypatch, repo=repo, days=30, batch_size=2) + + cleanup.run() + + assert repo.cleanup_ref_calls[1]["last_seen"] is None + assert repo.cleanup_ref_calls[1]["upper_bound"] == (second_created_at, "run-free-2") + assert repo.cleanup_ref_calls[2]["last_seen"] == (second_created_at, "run-free-2") + + @staticmethod + def test_run_exits_on_empty_batch(monkeypatch: pytest.MonkeyPatch) -> None: + cleanup = create_cleanup(monkeypatch, repo=FakeRepo([]), days=30, batch_size=10) + + cleanup.run() + + @staticmethod + def test_run_records_metrics_on_success(monkeypatch: pytest.MonkeyPatch) -> None: + cutoff = datetime.datetime.now() + repo = FakeRepo( + batches=[[_services_clear_free_plan_expired_workflow_run_logs_make_ref("run-free", "t_free", cutoff)]], + delete_result={ + "runs": 0, + "node_executions": 2, + "offloads": 1, + "app_logs": 3, + "trigger_logs": 4, + "pauses": 5, + "pause_reasons": 6, + }, + ) + cleanup = create_cleanup(monkeypatch, repo=repo, days=30, batch_size=10) + + batch_calls: list[dict[str, object]] = [] + completion_calls: list[dict[str, object]] = [] + monkeypatch.setattr(cleanup._metrics, "record_batch", lambda **kwargs: batch_calls.append(kwargs)) + monkeypatch.setattr(cleanup._metrics, "record_completion", lambda **kwargs: completion_calls.append(kwargs)) + + cleanup.run() + + assert len(batch_calls) == 1 + assert batch_calls[0]["batch_rows"] == 1 + assert batch_calls[0]["targeted_runs"] == 1 + assert batch_calls[0]["deleted_runs"] == 1 + assert batch_calls[0]["related_action"] == "deleted" + assert len(completion_calls) == 1 + assert completion_calls[0]["status"] == "success" + + @staticmethod + def test_run_records_failed_metrics(monkeypatch: pytest.MonkeyPatch) -> None: + class FailingRepo(FakeRepo): + def delete_runs_with_related_by_ids( + self, run_ids: list[str], delete_node_executions=None, delete_trigger_logs=None + ) -> dict[str, int]: + raise RuntimeError("delete failed") + + cutoff = datetime.datetime.now() + repo = FailingRepo( + batches=[[_services_clear_free_plan_expired_workflow_run_logs_make_ref("run-free", "t_free", cutoff)]] + ) + cleanup = create_cleanup(monkeypatch, repo=repo, days=30, batch_size=10) + + completion_calls: list[dict[str, object]] = [] + monkeypatch.setattr(cleanup._metrics, "record_completion", lambda **kwargs: completion_calls.append(kwargs)) + + with pytest.raises(RuntimeError, match="delete failed"): + cleanup.run() + + assert len(completion_calls) == 1 + assert completion_calls[0]["status"] == "failed" + + @staticmethod + def test_run_dry_run_skips_deletions(monkeypatch: pytest.MonkeyPatch, capsys: pytest.CaptureFixture[str]) -> None: + cutoff = datetime.datetime.now() + repo = FakeRepo( + batches=[[_services_clear_free_plan_expired_workflow_run_logs_make_ref("run-free", "t_free", cutoff)]], + count_result={ + "runs": 0, + "node_executions": 2, + "offloads": 1, + "app_logs": 3, + "trigger_logs": 4, + "pauses": 5, + "pause_reasons": 6, + }, + ) + cleanup = create_cleanup(monkeypatch, repo=repo, days=30, batch_size=10, dry_run=True) + + cleanup.run() + + assert repo.deleted == [] + assert repo.counted == [["run-free"]] + captured = capsys.readouterr().out + assert "Dry run mode enabled" in captured + assert "would delete 1 runs" in captured + assert "related records" in captured + assert "node_executions 2" in captured + assert "offloads 1" in captured + assert "app_logs 3" in captured + assert "trigger_logs 4" in captured + assert "pauses 5" in captured + assert "pause_reasons 6" in captured + + @staticmethod + def test_between_sets_window_bounds(monkeypatch: pytest.MonkeyPatch) -> None: + start_from = datetime.datetime(2024, 5, 1, 0, 0, 0) + end_before = datetime.datetime(2024, 6, 1, 0, 0, 0) + cleanup = create_cleanup( + monkeypatch, repo=FakeRepo([]), days=30, batch_size=10, start_from=start_from, end_before=end_before + ) + + assert cleanup.window_start == start_from + assert cleanup.window_end == end_before + + @staticmethod + def test_between_requires_both_boundaries(monkeypatch: pytest.MonkeyPatch) -> None: + with pytest.raises(ValueError, match="start_from and end_before must be both set or both omitted\\."): + create_cleanup( + monkeypatch, + repo=FakeRepo([]), + days=30, + batch_size=10, + start_from=datetime.datetime.now(), + end_before=None, + ) + with pytest.raises(ValueError, match="start_from and end_before must be both set or both omitted\\."): + create_cleanup( + monkeypatch, + repo=FakeRepo([]), + days=30, + batch_size=10, + start_from=None, + end_before=datetime.datetime.now(), ) - def test_equal_start_end_raises(self, mock_repo): - dt = datetime.datetime(2024, 1, 1) + @staticmethod + def test_between_requires_end_after_start(monkeypatch: pytest.MonkeyPatch) -> None: + start_from = datetime.datetime(2024, 6, 1, 0, 0, 0) + end_before = datetime.datetime(2024, 5, 1, 0, 0, 0) with pytest.raises(ValueError, match="end_before must be greater than start_from\\."): - WorkflowRunCleanup( - days=30, - batch_size=10, - start_from=dt, - end_before=dt, - workflow_run_repo=mock_repo, + create_cleanup( + monkeypatch, repo=FakeRepo([]), days=30, batch_size=10, start_from=start_from, end_before=end_before ) - - def test_zero_batch_size_raises(self, mock_repo): - with pytest.raises(ValueError, match="batch_size must be greater than 0"): - WorkflowRunCleanup(days=30, batch_size=0, workflow_run_repo=mock_repo) - - def test_negative_batch_size_raises(self, mock_repo): - with pytest.raises(ValueError, match="batch_size must be greater than 0\\."): - WorkflowRunCleanup(days=30, batch_size=-1, workflow_run_repo=mock_repo) - - def test_valid_window_init(self, mock_repo, config_overrides): - config_overrides(SANDBOX_EXPIRED_RECORDS_CLEAN_GRACEFUL_PERIOD=7) - start = datetime.datetime(2024, 1, 1) - end = datetime.datetime(2024, 6, 1) - c = WorkflowRunCleanup( - days=30, - batch_size=5, - start_from=start, - end_before=end, - workflow_run_repo=mock_repo, - ) - assert c.window_start == start - assert c.window_end == end - - def test_default_task_label_is_custom(self, mock_repo): - c = WorkflowRunCleanup(days=30, batch_size=10, workflow_run_repo=mock_repo) - - assert c._metrics._base_attributes["task_label"] == "custom" - - -# --------------------------------------------------------------------------- -# _empty_related_counts / _format_related_counts -# --------------------------------------------------------------------------- - - -class TestStaticHelpers: - def test_empty_related_counts(self): - counts = WorkflowRunCleanup._empty_related_counts() - assert counts == { - "node_executions": 0, - "offloads": 0, - "app_logs": 0, - "trigger_logs": 0, - "pauses": 0, - "pause_reasons": 0, - } - - def test_format_related_counts(self): - counts = { - "node_executions": 1, - "offloads": 2, - "app_logs": 3, - "trigger_logs": 4, - "pauses": 5, - "pause_reasons": 6, - } - result = WorkflowRunCleanup._format_related_counts(counts) - assert "node_executions 1" in result - assert "offloads 2" in result - assert "trigger_logs 4" in result - - -# --------------------------------------------------------------------------- -# _expiration_datetime -# --------------------------------------------------------------------------- - - -class TestExpirationDatetime: - def test_negative_returns_none(self, cleanup): - assert cleanup._expiration_datetime("t1", -1) is None - - def test_valid_timestamp(self, cleanup): - ts = int(datetime.datetime(2025, 1, 1, tzinfo=datetime.UTC).timestamp()) - result = cleanup._expiration_datetime("t1", ts) - assert result is not None - assert result.year == 2025 - - def test_overflow_returns_none(self, cleanup): - result = cleanup._expiration_datetime("t1", 2**62) - assert result is None - - -# --------------------------------------------------------------------------- -# _is_within_grace_period -# --------------------------------------------------------------------------- - - -class TestIsWithinGracePeriod: - def test_zero_grace_period_returns_false(self, cleanup): - cleanup.free_plan_grace_period_days = 0 - assert cleanup._is_within_grace_period("t1", {"expiration_date": 9999999999}) is False - - def test_within_grace_period(self, cleanup): - cleanup.free_plan_grace_period_days = 30 - # expired just 1 day ago - expired = datetime.datetime.now(datetime.UTC) - datetime.timedelta(days=1) - ts = int(expired.timestamp()) - assert cleanup._is_within_grace_period("t1", {"expiration_date": ts}) is True - - def test_outside_grace_period(self, cleanup): - cleanup.free_plan_grace_period_days = 5 - # expired 100 days ago - expired = datetime.datetime.now(datetime.UTC) - datetime.timedelta(days=100) - ts = int(expired.timestamp()) - assert cleanup._is_within_grace_period("t1", {"expiration_date": ts}) is False - - def test_missing_expiration_date_returns_false(self, cleanup): - cleanup.free_plan_grace_period_days = 30 - assert cleanup._is_within_grace_period("t1", {"expiration_date": -1}) is False - - -# --------------------------------------------------------------------------- -# _get_cleanup_whitelist -# --------------------------------------------------------------------------- - - -class TestGetCleanupWhitelist: - def test_non_cloud_edition_returns_empty(self, cleanup): - cleanup._cleanup_whitelist = None - result = cleanup._get_cleanup_whitelist() - assert result == set() - - def test_cloud_edition_fetches_whitelist(self, mock_repo, config_overrides): - config_overrides(DEPLOYMENT_EDITION=DeploymentEdition.CLOUD) - c = WorkflowRunCleanup(days=30, batch_size=10, workflow_run_repo=mock_repo) - with patch("services.retention.workflow_run.clear_free_plan_expired_workflow_run_logs.BillingService") as bs: - bs.get_expired_subscription_cleanup_whitelist.return_value = ["t1", "t2"] - result = c._get_cleanup_whitelist() - assert result == {"t1", "t2"} - - def test_cached_whitelist_returned(self, cleanup): - cleanup._cleanup_whitelist = {"cached"} - result = cleanup._get_cleanup_whitelist() - assert result == {"cached"} - - def test_billing_service_error_returns_empty(self, mock_repo, config_overrides): - config_overrides(DEPLOYMENT_EDITION=DeploymentEdition.CLOUD) - c = WorkflowRunCleanup(days=30, batch_size=10, workflow_run_repo=mock_repo) - with patch("services.retention.workflow_run.clear_free_plan_expired_workflow_run_logs.BillingService") as bs: - bs.get_expired_subscription_cleanup_whitelist.side_effect = Exception("error") - result = c._get_cleanup_whitelist() - assert result == set() - - -# --------------------------------------------------------------------------- -# _filter_free_tenants -# --------------------------------------------------------------------------- - - -class TestFilterFreeTenants: - def test_non_cloud_edition_treats_all_tenants_as_free(self, cleanup): - result = cleanup._filter_free_tenants(["t1", "t2"]) - assert result == {"t1", "t2"} - - def test_empty_tenants_returns_empty(self, cleanup, config_overrides): - config_overrides(DEPLOYMENT_EDITION=DeploymentEdition.CLOUD) - result = cleanup._filter_free_tenants([]) - assert result == set() - - def test_whitelisted_tenant_excluded(self, mock_repo, config_overrides): - config_overrides(DEPLOYMENT_EDITION=DeploymentEdition.CLOUD) - c = WorkflowRunCleanup(days=30, batch_size=10, workflow_run_repo=mock_repo) - c._cleanup_whitelist = {"t1"} - with patch("services.retention.workflow_run.clear_free_plan_expired_workflow_run_logs.BillingService") as bs: - bs.get_plan_bulk_with_cache.return_value = { - "t1": {"plan": CloudPlan.SANDBOX, "expiration_date": -1}, - "t2": {"plan": CloudPlan.SANDBOX, "expiration_date": -1}, - } - result = c._filter_free_tenants(["t1", "t2"]) - assert "t1" not in result - assert "t2" in result - - def test_paid_tenant_excluded(self, mock_repo, config_overrides): - config_overrides(DEPLOYMENT_EDITION=DeploymentEdition.CLOUD) - c = WorkflowRunCleanup(days=30, batch_size=10, workflow_run_repo=mock_repo) - c._cleanup_whitelist = set() - with patch("services.retention.workflow_run.clear_free_plan_expired_workflow_run_logs.BillingService") as bs: - bs.get_plan_bulk_with_cache.return_value = { - "t1": {"plan": "professional", "expiration_date": -1}, - } - result = c._filter_free_tenants(["t1"]) - assert result == set() - - def test_missing_billing_info_treats_as_non_free(self, mock_repo, config_overrides): - config_overrides(DEPLOYMENT_EDITION=DeploymentEdition.CLOUD) - c = WorkflowRunCleanup(days=30, batch_size=10, workflow_run_repo=mock_repo) - c._cleanup_whitelist = set() - with patch("services.retention.workflow_run.clear_free_plan_expired_workflow_run_logs.BillingService") as bs: - bs.get_plan_bulk_with_cache.return_value = {} - result = c._filter_free_tenants(["t1"]) - assert result == set() - - def test_billing_bulk_error_treats_as_non_free(self, mock_repo, config_overrides): - config_overrides(DEPLOYMENT_EDITION=DeploymentEdition.CLOUD) - c = WorkflowRunCleanup(days=30, batch_size=10, workflow_run_repo=mock_repo) - c._cleanup_whitelist = set() - with patch("services.retention.workflow_run.clear_free_plan_expired_workflow_run_logs.BillingService") as bs: - bs.get_plan_bulk_with_cache.side_effect = Exception("fail") - result = c._filter_free_tenants(["t1"]) - assert result == set() - - -# --------------------------------------------------------------------------- -# run() — delete mode -# --------------------------------------------------------------------------- - - -class TestRunDeleteMode: - def _make_cleanup(self, mock_repo): - return WorkflowRunCleanup(days=30, batch_size=10, workflow_run_repo=mock_repo) - - def test_no_rows_stops_immediately(self, mock_repo): - mock_repo.get_cleanup_refs_batch_by_time_range.return_value = [] - c = self._make_cleanup(mock_repo) - c.run() - mock_repo.delete_runs_with_related_by_ids.assert_not_called() - - def test_all_paid_skips_delete(self, mock_repo): - ref = make_ref("t1") - mock_repo.get_cleanup_refs_batch_by_time_range.side_effect = [[ref], []] - c = self._make_cleanup(mock_repo) - # Override the non-Cloud default to exercise the no-deletion path. - c._filter_free_tenants = MagicMock(return_value=set()) - c.run() - mock_repo.delete_runs_with_related_by_ids.assert_not_called() - - def test_runs_deleted_successfully(self, mock_repo): - ref = make_ref("t1") - mock_repo.get_cleanup_refs_batch_by_time_range.side_effect = [[ref], [ref], []] - mock_repo.delete_runs_with_related_by_ids.return_value = { - "runs": 1, - "node_executions": 0, - "offloads": 0, - "app_logs": 0, - "trigger_logs": 0, - "pauses": 0, - "pause_reasons": 0, - } - c = self._make_cleanup(mock_repo) - with patch("services.retention.workflow_run.clear_free_plan_expired_workflow_run_logs.time.sleep"): - c.run() - mock_repo.delete_runs_with_related_by_ids.assert_called_once() - - def test_delete_exception_reraises(self, mock_repo): - ref = make_ref("t1") - mock_repo.get_cleanup_refs_batch_by_time_range.side_effect = [[ref], [ref]] - mock_repo.delete_runs_with_related_by_ids.side_effect = RuntimeError("db error") - c = self._make_cleanup(mock_repo) - with pytest.raises(RuntimeError): - c.run() - - def test_summary_with_window_start(self, mock_repo): - mock_repo.get_cleanup_refs_batch_by_time_range.return_value = [] - c = WorkflowRunCleanup( - days=30, - batch_size=10, - start_from=datetime.datetime(2024, 1, 1), - end_before=datetime.datetime(2024, 6, 1), - workflow_run_repo=mock_repo, - ) - c.run() - - -# --------------------------------------------------------------------------- -# run() — dry run mode -# --------------------------------------------------------------------------- - - -class TestRunDryRunMode: - def _make_dry_cleanup(self, mock_repo): - return WorkflowRunCleanup( - days=30, - batch_size=10, - workflow_run_repo=mock_repo, - dry_run=True, - ) - - def test_dry_run_no_delete_called(self, mock_repo): - ref = make_ref("t1") - mock_repo.get_cleanup_refs_batch_by_time_range.side_effect = [[ref], [ref], []] - mock_repo.count_runs_with_related_by_ids.return_value = { - "runs": 1, - "node_executions": 2, - "offloads": 0, - "app_logs": 0, - "trigger_logs": 1, - "pauses": 0, - "pause_reasons": 0, - } - c = self._make_dry_cleanup(mock_repo) - c.run() - mock_repo.delete_runs_with_related_by_ids.assert_not_called() - mock_repo.count_runs_with_related_by_ids.assert_called_once() - - def test_dry_run_summary_with_window_start(self, mock_repo): - mock_repo.get_cleanup_refs_batch_by_time_range.return_value = [] - c = WorkflowRunCleanup( - days=30, - batch_size=10, - start_from=datetime.datetime(2024, 1, 1), - end_before=datetime.datetime(2024, 6, 1), - workflow_run_repo=mock_repo, - dry_run=True, - ) - c.run() - - def test_dry_run_all_paid_skips_count(self, mock_repo): - ref = make_ref("t1") - mock_repo.get_cleanup_refs_batch_by_time_range.side_effect = [[ref], []] - c = self._make_dry_cleanup(mock_repo) - c._filter_free_tenants = MagicMock(return_value=set()) - c.run() - mock_repo.count_runs_with_related_by_ids.assert_not_called() - - -# --------------------------------------------------------------------------- -# _delete_trigger_logs / _count_trigger_logs -# --------------------------------------------------------------------------- - - -class TestTriggerLogMethods: - @pytest.mark.parametrize("sqlite_session", [()], indirect=True) - def test_delete_trigger_logs(self, cleanup, sqlite_session: Session): - with patch( - "services.retention.workflow_run.clear_free_plan_expired_workflow_run_logs.SQLAlchemyWorkflowTriggerLogRepository" - ) as repo_class: - instance = repo_class.return_value - instance.delete_by_run_ids.return_value = 5 - result = cleanup._delete_trigger_logs(sqlite_session, ["r1", "r2"]) - assert result == 5 - - @pytest.mark.parametrize("sqlite_session", [()], indirect=True) - def test_count_trigger_logs(self, cleanup, sqlite_session: Session): - with patch( - "services.retention.workflow_run.clear_free_plan_expired_workflow_run_logs.SQLAlchemyWorkflowTriggerLogRepository" - ) as repo_class: - instance = repo_class.return_value - instance.count_by_run_ids.return_value = 3 - result = cleanup._count_trigger_logs(sqlite_session, ["r1"]) - assert result == 3 - - -# --------------------------------------------------------------------------- -# _count_node_executions_by_run_ids / _delete_node_executions_by_run_ids -# --------------------------------------------------------------------------- - - -class TestNodeExecutionMethods: - @pytest.mark.parametrize("sqlite_session", [()], indirect=True) - def test_count_node_executions(self, cleanup, sqlite_session: Session): - with patch( - "services.retention.workflow_run.clear_free_plan_expired_workflow_run_logs.DifyAPIRepositoryFactory" - ) as factory: - repo = factory.create_api_workflow_node_execution_repository.return_value - repo.count_by_runs.return_value = (10, 2) - result = cleanup._count_node_executions_by_run_ids(sqlite_session, ["r1"]) - assert result == (10, 2) - - @pytest.mark.parametrize("sqlite_session", [()], indirect=True) - def test_delete_node_executions(self, cleanup, sqlite_session: Session): - with patch( - "services.retention.workflow_run.clear_free_plan_expired_workflow_run_logs.DifyAPIRepositoryFactory" - ) as factory: - repo = factory.create_api_workflow_node_execution_repository.return_value - repo.delete_by_runs.return_value = (5, 1) - result = cleanup._delete_node_executions_by_run_ids(sqlite_session, ["r1"]) - assert result == (5, 1) diff --git a/api/tests/unit_tests/services/retention/workflow_run/test_workflow_run_archiver_sqlite.py b/api/tests/unit_tests/services/retention/workflow_run/test_workflow_run_archiver_sqlite.py deleted file mode 100644 index fa4de531b7f..00000000000 --- a/api/tests/unit_tests/services/retention/workflow_run/test_workflow_run_archiver_sqlite.py +++ /dev/null @@ -1,705 +0,0 @@ -"""SQLite and in-process coverage for workflow-run bundle archiving.""" - -import datetime -import json -import uuid -from types import SimpleNamespace -from typing import override -from unittest.mock import ANY, MagicMock, patch - -import pyarrow as pa -import pyarrow.parquet as pq -import pytest -from sqlalchemy import select -from sqlalchemy.engine import Engine -from sqlalchemy.exc import OperationalError -from sqlalchemy.orm import Session, sessionmaker - -from enums import DeploymentEdition -from graphon.enums import WorkflowExecutionStatus -from libs.archive_storage import ArchiveStorage -from models.enums import CreatorUserRole, WorkflowRunTriggeredFrom -from models.workflow import WorkflowRun, WorkflowRunArchiveBundle, WorkflowType -from services.retention.workflow_run.archive_paid_plan_workflow_run import ( - ArchiveResult, - ArchiveSummary, - TableStats, - WorkflowRunArchiver, -) -from services.retention.workflow_run.constants import ARCHIVE_BUNDLE_FORMAT, ARCHIVE_BUNDLE_SCHEMA_VERSION - - -class FakeArchiveStorage(ArchiveStorage): - def __init__(self, objects: dict[str, bytes] | None = None) -> None: - self.objects = {} if objects is None else objects - - @override - def object_exists(self, key: str) -> bool: - return key in self.objects - - @override - def get_object(self, key: str) -> bytes: - return self.objects[key] - - @override - def put_object(self, key: str, data: bytes) -> str: - self.objects[key] = data - return "checksum" - - @override - def list_objects(self, prefix: str) -> list[str]: - return sorted(key for key in self.objects if key.startswith(prefix)) - - -def _db_disconnect_error() -> OperationalError: - return OperationalError( - "select 1", - {}, - RuntimeError("server closed the connection unexpectedly"), - connection_invalidated=True, - ) - - -def _run(run_id: str = "run-1") -> WorkflowRun: - """Build the mapped entity consumed by the archiver without persisting it.""" - return WorkflowRun( - id=run_id, - tenant_id="tenant-1", - app_id="app-1", - workflow_id="workflow-1", - type=WorkflowType.WORKFLOW, - triggered_from=WorkflowRunTriggeredFrom.APP_RUN, - version="1", - graph="{}", - inputs="{}", - status=WorkflowExecutionStatus.SUCCEEDED, - outputs="{}", - error=None, - elapsed_time=0, - total_tokens=0, - total_steps=0, - created_by_role=CreatorUserRole.ACCOUNT, - created_by="account-1", - created_at=datetime.datetime(2025, 3, 15, 10, 0, 0), - finished_at=datetime.datetime(2025, 3, 15, 10, 1, 0), - ) - - -class TestWorkflowRunArchiverInit: - def test_start_from_without_end_before_raises(self) -> None: - with pytest.raises(ValueError, match="start_from and end_before must be provided together"): - WorkflowRunArchiver(start_from=datetime.datetime(2025, 1, 1)) - - def test_end_before_without_start_from_raises(self) -> None: - with pytest.raises(ValueError, match="start_from and end_before must be provided together"): - WorkflowRunArchiver(end_before=datetime.datetime(2025, 1, 1)) - - def test_start_equals_end_raises(self) -> None: - ts = datetime.datetime(2025, 1, 1) - with pytest.raises(ValueError, match="start_from must be earlier than end_before"): - WorkflowRunArchiver(start_from=ts, end_before=ts) - - def test_start_after_end_raises(self) -> None: - with pytest.raises(ValueError, match="start_from must be earlier than end_before"): - WorkflowRunArchiver( - start_from=datetime.datetime(2025, 6, 1), - end_before=datetime.datetime(2025, 1, 1), - ) - - def test_workers_zero_raises(self) -> None: - with pytest.raises(ValueError, match="workers must be at least 1"): - WorkflowRunArchiver(workers=0) - - def test_run_shard_index_without_total_raises(self) -> None: - with pytest.raises(ValueError, match="run_shard_index and run_shard_total must be provided together"): - WorkflowRunArchiver(run_shard_index=0) - - def test_run_shard_total_without_index_raises(self) -> None: - with pytest.raises(ValueError, match="run_shard_index and run_shard_total must be provided together"): - WorkflowRunArchiver(run_shard_total=4) - - def test_run_shard_total_above_supported_range_raises(self) -> None: - with pytest.raises(ValueError, match="run_shard_total must be between 1 and 16"): - WorkflowRunArchiver(run_shard_index=0, run_shard_total=17) - - def test_run_shard_index_must_be_less_than_total(self) -> None: - with pytest.raises(ValueError, match="run_shard_index must be between 0 and run_shard_total - 1"): - WorkflowRunArchiver(run_shard_index=4, run_shard_total=4) - - def test_valid_init_defaults(self) -> None: - archiver = WorkflowRunArchiver(days=30, batch_size=50) - assert archiver.days == 30 - assert archiver.batch_size == 50 - assert archiver.dry_run is False - assert archiver.delete_after_archive is False - assert archiver.start_from is None - - def test_valid_init_with_time_range(self) -> None: - start = datetime.datetime(2025, 1, 1) - end = datetime.datetime(2025, 6, 1) - archiver = WorkflowRunArchiver(start_from=start, end_before=end, workers=2) - assert archiver.start_from is not None - assert archiver.end_before is not None - assert archiver.workers == 2 - - def test_delete_after_archive_is_not_supported_for_bundle_archive(self) -> None: - with pytest.raises(ValueError, match="delete_after_archive is not supported by bundle archive"): - WorkflowRunArchiver(delete_after_archive=True) - - def test_get_runs_batch_passes_shard_options(self) -> None: - repo = MagicMock() - empty_runs: list[WorkflowRun] = [] - repo.get_runs_batch_by_time_range.return_value = empty_runs - archiver = WorkflowRunArchiver( - tenant_prefixes=["0", "a"], - run_shard_index=1, - run_shard_total=4, - workflow_run_repo=repo, - ) - - archiver._get_runs_batch(None) - - repo.get_runs_batch_by_time_range.assert_called_once() - assert repo.get_runs_batch_by_time_range.call_args.kwargs["tenant_prefixes"] == ["0", "a"] - assert repo.get_runs_batch_by_time_range.call_args.kwargs["run_shard_index"] == 1 - assert repo.get_runs_batch_by_time_range.call_args.kwargs["run_shard_total"] == 4 - - def test_get_runs_batch_prefers_planned_tenant_ids_over_prefix_filter(self) -> None: - repo = MagicMock() - empty_runs: list[WorkflowRun] = [] - repo.get_runs_batch_by_time_range.return_value = empty_runs - archiver = WorkflowRunArchiver( - tenant_ids=["0tenant"], - tenant_prefixes=["0"], - paid_tenant_ids=["0tenant"], - workflow_run_repo=repo, - ) - - archiver._get_runs_batch(None) - - repo.get_runs_batch_by_time_range.assert_called_once() - assert repo.get_runs_batch_by_time_range.call_args.kwargs["tenant_ids"] == ["0tenant"] - assert repo.get_runs_batch_by_time_range.call_args.kwargs["tenant_prefixes"] is None - - def test_get_runs_batch_uses_current_tenant_scan_scope(self) -> None: - repo = MagicMock() - empty_runs: list[WorkflowRun] = [] - repo.get_runs_batch_by_time_range.return_value = empty_runs - archiver = WorkflowRunArchiver( - tenant_ids=["tenant-a", "tenant-b"], - workflow_run_repo=repo, - ) - - archiver._get_runs_batch(None, tenant_scope=["tenant-b"]) - - repo.get_runs_batch_by_time_range.assert_called_once() - assert repo.get_runs_batch_by_time_range.call_args.kwargs["tenant_ids"] == ["tenant-b"] - - def test_get_runs_batch_retries_retryable_db_disconnect(self) -> None: - repo = MagicMock() - empty_runs: list[WorkflowRun] = [] - repo.get_runs_batch_by_time_range.side_effect = [_db_disconnect_error(), empty_runs] - archiver = WorkflowRunArchiver(workflow_run_repo=repo) - - with patch("services.retention.workflow_run.db_retry.time.sleep") as sleep: - runs = archiver._get_runs_batch(None) - - assert runs == [] - assert repo.get_runs_batch_by_time_range.call_count == 2 - sleep.assert_called_once_with(1.0) - - def test_get_runs_batch_does_not_retry_non_db_broken_pipe_error(self) -> None: - repo = MagicMock() - repo.get_runs_batch_by_time_range.side_effect = RuntimeError("broken pipe") - archiver = WorkflowRunArchiver(workflow_run_repo=repo) - - with ( - patch("services.retention.workflow_run.db_retry.time.sleep") as sleep, - pytest.raises(RuntimeError, match="broken pipe"), - ): - archiver._get_runs_batch(None) - - repo.get_runs_batch_by_time_range.assert_called_once() - sleep.assert_not_called() - - def test_start_message_includes_shard(self) -> None: - archiver = WorkflowRunArchiver(tenant_prefixes=["0"], run_shard_index=1, run_shard_total=4) - - message = archiver._build_start_message() - - assert "tenant_prefixes=0" in message - assert "run_shard=1/4" in message - - def test_start_message_summarizes_large_planned_tenant_list(self) -> None: - tenant_ids = [f"tenant-{index}" for index in range(11)] - archiver = WorkflowRunArchiver(tenant_ids=tenant_ids, tenant_prefixes=["0"]) - - message = archiver._build_start_message() - - assert "tenant_ids=11 planned tenants" in message - assert "tenant-10" not in message - - -class TestBuildArchiveBundle: - def test_bundle_contains_manifest_and_all_table_objects(self) -> None: - archiver = WorkflowRunArchiver(days=90) - run = _run(str(uuid.uuid4())) - run.tenant_id = str(uuid.uuid4()) - identity = archiver._build_bundle_identity([run]) - table_data = {"workflow_runs": [{"id": run.id, "tenant_id": run.tenant_id}]} - - table_stats, table_payloads, manifest_data = archiver._build_archive_payload(identity, [run], table_data) - manifest = json.loads(manifest_data) - - assert manifest["schema_version"] == ARCHIVE_BUNDLE_SCHEMA_VERSION - assert manifest["archive_format"] == ARCHIVE_BUNDLE_FORMAT - assert manifest["object_prefix"] == identity.object_prefix - assert set(table_payloads) == set(archiver.ARCHIVED_TABLES) - assert {stat.table_name for stat in table_stats} == set(archiver.ARCHIVED_TABLES) - assert pq.read_table(pa.BufferReader(table_payloads["workflow_runs"])).num_rows == 1 - - -class TestGenerateManifest: - def test_manifest_structure(self) -> None: - start = datetime.datetime(2025, 1, 1, tzinfo=datetime.UTC) - end = datetime.datetime(2025, 4, 1, tzinfo=datetime.UTC) - archiver = WorkflowRunArchiver(start_from=start, end_before=end, run_shard_index=1, run_shard_total=4) - run = _run(str(uuid.uuid4())) - run.tenant_id = str(uuid.uuid4()) - identity = archiver._build_bundle_identity([run]) - - stats = [ - TableStats( - table_name="workflow_runs", - row_count=1, - checksum="abc123", - size_bytes=512, - object_key="workflow_runs.parquet", - ), - TableStats( - table_name="workflow_node_executions", - row_count=2, - checksum="def456", - size_bytes=1024, - object_key="workflow_node_executions.parquet", - ), - ] - - manifest = archiver._generate_manifest(identity, [run], stats) - - assert manifest["schema_version"] == ARCHIVE_BUNDLE_SCHEMA_VERSION - assert manifest["archive_format"] == ARCHIVE_BUNDLE_FORMAT - assert manifest["bundle_id"] == identity.bundle_id - assert manifest["tenant_id"] == run.tenant_id - assert manifest["workflow_run_count"] == 1 - assert manifest["workflow_node_execution_count"] == 2 - assert manifest["run_ids"] == [run.id] - assert manifest["campaign_id"] == "2025-01-01T00:00:00Z_2025-04-01T00:00:00Z" - assert manifest["archive_window_start"] == "2025-01-01T00:00:00Z" - assert manifest["archive_window_end"] == "2025-04-01T00:00:00Z" - assert manifest["run_shard"] == "01-of-04" - assert "tables" in manifest - assert manifest["tables"]["workflow_runs"]["row_count"] == 1 - assert manifest["tables"]["workflow_runs"]["checksum"] == "abc123" - assert manifest["tables"]["workflow_node_executions"]["row_count"] == 2 - - -class TestFilterPaidTenants: - def test_all_tenants_paid_in_community_edition(self) -> None: - archiver = WorkflowRunArchiver(days=90) - tenant_ids = {"t1", "t2", "t3"} - - with patch("services.retention.workflow_run.archive_paid_plan_workflow_run.dify_config") as cfg: - cfg.DEPLOYMENT_EDITION = DeploymentEdition.COMMUNITY - result = archiver._filter_paid_tenants(tenant_ids) - - assert result == tenant_ids - - def test_empty_tenants_returns_empty(self) -> None: - archiver = WorkflowRunArchiver(days=90) - - with patch("services.retention.workflow_run.archive_paid_plan_workflow_run.dify_config") as cfg: - cfg.DEPLOYMENT_EDITION = DeploymentEdition.CLOUD - result = archiver._filter_paid_tenants(set()) - - assert result == set() - - def test_only_paid_plans_returned(self) -> None: - archiver = WorkflowRunArchiver(days=90) - - mock_bulk = { - "t1": {"plan": "professional"}, - "t2": {"plan": "sandbox"}, - "t3": {"plan": "team"}, - } - - with ( - patch("services.retention.workflow_run.archive_paid_plan_workflow_run.dify_config") as cfg, - patch("services.retention.workflow_run.archive_paid_plan_workflow_run.BillingService") as billing, - ): - cfg.DEPLOYMENT_EDITION = DeploymentEdition.CLOUD - billing.get_plan_bulk_with_cache.return_value = mock_bulk - result = archiver._filter_paid_tenants({"t1", "t2", "t3"}) - - assert "t1" in result - assert "t3" in result - assert "t2" not in result - - def test_billing_api_failure_returns_empty(self) -> None: - archiver = WorkflowRunArchiver(days=90) - - with ( - patch("services.retention.workflow_run.archive_paid_plan_workflow_run.dify_config") as cfg, - patch("services.retention.workflow_run.archive_paid_plan_workflow_run.BillingService") as billing, - ): - cfg.DEPLOYMENT_EDITION = DeploymentEdition.CLOUD - billing.get_plan_bulk_with_cache.side_effect = RuntimeError("API down") - result = archiver._filter_paid_tenants({"t1"}) - - assert result == set() - - def test_planned_paid_tenants_skip_billing_lookup(self) -> None: - archiver = WorkflowRunArchiver(days=90, paid_tenant_ids=["t1", "t3"]) - - with ( - patch("services.retention.workflow_run.archive_paid_plan_workflow_run.dify_config") as cfg, - patch("services.retention.workflow_run.archive_paid_plan_workflow_run.BillingService") as billing, - ): - cfg.DEPLOYMENT_EDITION = DeploymentEdition.CLOUD - result = archiver._filter_paid_tenants({"t1", "t2", "t3"}) - - billing.get_plan_bulk_with_cache.assert_not_called() - assert result == {"t1", "t3"} - - -class TestDryRunArchive: - @patch("services.retention.workflow_run.archive_paid_plan_workflow_run.get_archive_storage") - def test_dry_run_does_not_call_storage(self, mock_get_storage: MagicMock, sqlite_engine: Engine) -> None: - archiver = WorkflowRunArchiver(days=90, dry_run=True) - - with ( - patch( - "services.retention.workflow_run.archive_paid_plan_workflow_run.db", - SimpleNamespace(engine=sqlite_engine), - ), - patch.object(archiver, "_get_runs_batch", return_value=[]), - ): - summary = archiver.run() - - mock_get_storage.assert_not_called() - assert isinstance(summary, ArchiveSummary) - assert summary.runs_failed == 0 - - def test_dry_run_estimates_table_and_object_sizes(self, sqlite_session: Session) -> None: - archiver = WorkflowRunArchiver(days=90, dry_run=True) - run = _run() - table_data = { - "workflow_runs": [{"id": "run-1", "tenant_id": "tenant-1"}], - "workflow_app_logs": [{"id": "log-1", "workflow_run_id": "run-1"}], - } - - with patch.object(archiver, "_extract_bundle_data", return_value=table_data): - result = archiver._archive_bundle(sqlite_session, None, [run]) - - stats_by_table = {stat.table_name: stat for stat in result.tables} - assert result.success is True - assert result.object_size_bytes > 0 - assert stats_by_table["workflow_runs"].row_count == 1 - assert stats_by_table["workflow_runs"].size_bytes > 0 - assert stats_by_table["workflow_app_logs"].row_count == 1 - assert stats_by_table["workflow_app_logs"].size_bytes > 0 - assert stats_by_table["workflow_node_executions"].row_count == 0 - assert stats_by_table["workflow_node_executions"].size_bytes > 0 - - def test_summary_merges_dry_run_estimates(self) -> None: - summary = ArchiveSummary() - result = ArchiveResult( - bundle_id="bundle-1", - tenant_id="tenant-1", - object_prefix="prefix", - run_count=1, - success=True, - object_size_bytes=128, - tables=[ - TableStats(table_name="workflow_runs", row_count=1, checksum="", size_bytes=64), - TableStats(table_name="workflow_app_logs", row_count=2, checksum="", size_bytes=32), - ], - ) - - WorkflowRunArchiver._merge_result_stats(summary, result) - - assert summary.total_object_size_bytes == 128 - assert summary.table_stats["workflow_runs"].row_count == 1 - assert summary.table_stats["workflow_runs"].size_bytes == 64 - assert summary.table_stats["workflow_app_logs"].row_count == 2 - assert summary.table_stats["workflow_app_logs"].size_bytes == 32 - - -class TestArchiveDbRetry: - def test_archive_bundle_groups_retries_with_fresh_session( - self, sqlite_session_factory: sessionmaker[Session] - ) -> None: - archiver = WorkflowRunArchiver(days=90) - run = _run() - session_maker = MagicMock(wraps=sqlite_session_factory) - success = ArchiveResult( - bundle_id=archiver._build_bundle_identity([run]).bundle_id, - tenant_id=run.tenant_id, - object_prefix=archiver._build_bundle_identity([run]).object_prefix, - run_count=1, - success=True, - ) - - with ( - patch.object(archiver, "_archive_bundle", side_effect=[_db_disconnect_error(), success]) as archive_bundle, - patch("services.retention.workflow_run.db_retry.time.sleep") as sleep, - ): - results = archiver._archive_bundle_groups(session_maker, MagicMock(), [[run]]) - - assert results == [success] - assert archive_bundle.call_count == 2 - assert session_maker.call_count == 2 - sleep.assert_called_once_with(1.0) - - def test_archive_bundle_groups_returns_failed_result_after_retry_exhaustion( - self, sqlite_session_factory: sessionmaker[Session] - ) -> None: - archiver = WorkflowRunArchiver(days=90) - run = _run() - session_maker = MagicMock(wraps=sqlite_session_factory) - - with ( - patch.object(archiver, "_archive_bundle", side_effect=[_db_disconnect_error()] * 3) as archive_bundle, - patch("services.retention.workflow_run.db_retry.time.sleep") as sleep, - ): - results = archiver._archive_bundle_groups(session_maker, MagicMock(), [[run]]) - - assert len(results) == 1 - assert results[0].success is False - assert "server closed the connection unexpectedly" in (results[0].error or "") - assert archive_bundle.call_count == archiver.DB_RETRY_ATTEMPTS - assert session_maker.call_count == archiver.DB_RETRY_ATTEMPTS - assert sleep.call_count == archiver.DB_RETRY_ATTEMPTS - 1 - - def test_archive_bundle_uses_safe_rollback_when_failure_rolls_back_badly(self, sqlite_session: Session) -> None: - archiver = WorkflowRunArchiver(days=90, dry_run=True) - - with ( - patch.object(sqlite_session, "rollback", side_effect=RuntimeError("rollback failed")) as rollback, - patch.object(archiver, "_extract_bundle_data", side_effect=RuntimeError("extract failed")), - ): - result = archiver._archive_bundle(sqlite_session, None, [_run()]) - - assert result.success is False - assert result.error == "extract failed" - rollback.assert_called_once() - - -class TestArchiveRunIdempotency: - def _index_payload(self, archiver: WorkflowRunArchiver, run_ids: list[str], run: WorkflowRun) -> tuple[str, bytes]: - identity = archiver._build_bundle_identity([run]) - index_key = archiver._get_index_object_key(identity) - payload = json.dumps( - { - "schema_version": ARCHIVE_BUNDLE_SCHEMA_VERSION, - "archive_format": ARCHIVE_BUNDLE_FORMAT, - "object_prefix": archiver._get_shard_object_prefix(identity), - "updated_at": "2025-03-15T00:00:00Z", - "manifest_keys": [], - "run_ids": run_ids, - } - ).encode() - return index_key, payload - - def test_locked_bundle_is_skipped(self, sqlite_session: Session) -> None: - archiver = WorkflowRunArchiver(days=90) - run = _run() - - with ( - patch.object(archiver, "_lock_runs_for_archive", return_value=[]), - ): - storage = MagicMock() - storage.object_exists.return_value = False - result = archiver._archive_bundle(sqlite_session, storage, [run]) - - assert result.success is True - assert result.skipped is True - assert result.error == "one or more runs locked or deleted by another archiver" - - def test_already_archived_bundle_is_skipped(self, sqlite_session: Session) -> None: - archiver = WorkflowRunArchiver(days=90) - run = _run() - storage = MagicMock() - storage.object_exists.return_value = True - - with patch.object(archiver, "_sync_existing_bundle_index") as sync_existing_bundle_index: - result = archiver._archive_bundle(sqlite_session, storage, [run]) - - assert result.success is True - assert result.skipped is True - assert result.error == "bundle already archived" - sync_existing_bundle_index.assert_called_once() - - def test_existing_bundle_catalog_publication_failure_is_not_success(self, sqlite_session: Session) -> None: - archiver = WorkflowRunArchiver(days=90) - run = _run() - storage = MagicMock() - storage.object_exists.return_value = True - - with ( - patch.object(sqlite_session, "rollback", wraps=sqlite_session.rollback) as rollback, - patch.object(archiver, "_sync_existing_bundle_index", side_effect=RuntimeError("catalog unavailable")), - ): - result = archiver._archive_bundle(sqlite_session, storage, [run]) - - assert result.success is False - assert result.error == "catalog unavailable" - rollback.assert_called_once() - - def test_retry_repairs_index_after_catalog_commit_then_index_write_failure(self, sqlite_session: Session) -> None: - archiver = WorkflowRunArchiver(days=90) - run = _run() - identity = archiver._build_bundle_identity([run]) - index_key = archiver._get_index_object_key(identity) - manifest_key = archiver._get_manifest_object_key(identity) - storage = FakeArchiveStorage() - original_put_object = storage.put_object - index_write_count = 0 - - def put_object(key: str, data: bytes) -> str: - nonlocal index_write_count - if key == index_key: - index_write_count += 1 - if index_write_count == 2: - raise RuntimeError("index write failed") - return original_put_object(key, data) - - storage.put_object = MagicMock(side_effect=put_object) - table_data = {"workflow_runs": [{"id": run.id, "tenant_id": run.tenant_id}]} - - with ( - patch.object(archiver, "_lock_runs_for_archive", return_value=[run]), - patch.object(archiver, "_extract_bundle_data", return_value=table_data), - ): - first_result = archiver._archive_bundle(sqlite_session, storage, [run]) - - assert first_result.success is False - assert first_result.error == "index write failed" - assert manifest_key in storage.objects - assert json.loads(storage.objects[index_key])["run_ids"] == [] - - storage.list_objects = MagicMock(wraps=storage.list_objects) - retry_result = archiver._archive_bundle(sqlite_session, storage, [run]) - - assert retry_result.success is True - assert retry_result.skipped is True - assert json.loads(storage.objects[index_key])["manifest_keys"] == [manifest_key] - assert json.loads(storage.objects[index_key])["run_ids"] == [run.id] - storage.list_objects.assert_not_called() - - def test_existing_manifest_with_missing_index_fails_without_partial_rebuild(self, sqlite_session: Session) -> None: - archiver = WorkflowRunArchiver(days=90) - run = _run() - identity = archiver._build_bundle_identity([run]) - _, _, manifest_data = archiver._build_archive_payload( - identity, - [run], - {"workflow_runs": [{"id": run.id, "tenant_id": run.tenant_id}]}, - ) - manifest_key = archiver._get_manifest_object_key(identity) - index_key = archiver._get_index_object_key(identity) - storage = FakeArchiveStorage({manifest_key: manifest_data}) - storage.list_objects = MagicMock(wraps=storage.list_objects) - - result = archiver._archive_bundle(sqlite_session, storage, [run]) - - assert result.success is False - assert "archive shard index missing" in (result.error or "") - assert index_key not in storage.objects - storage.list_objects.assert_not_called() - - def test_successful_bundle_persists_archive_index(self, sqlite_session: Session) -> None: - archiver = WorkflowRunArchiver(days=90) - run = _run(str(uuid.uuid4())) - run.tenant_id = str(uuid.uuid4()) - storage = MagicMock() - storage.object_exists.return_value = False - table_data = { - "workflow_runs": [{"id": run.id, "tenant_id": run.tenant_id}], - "workflow_node_executions": [{"id": str(uuid.uuid4()), "workflow_run_id": run.id}], - } - - with ( - patch.object(archiver, "_lock_runs_for_archive", return_value=[run]), - patch.object(archiver, "_extract_bundle_data", return_value=table_data), - ): - result = archiver._archive_bundle(sqlite_session, storage, [run]) - - archived_bundle = sqlite_session.scalar( - select(WorkflowRunArchiveBundle).where(WorkflowRunArchiveBundle.tenant_id == run.tenant_id) - ) - assert result.success is True - assert isinstance(archived_bundle, WorkflowRunArchiveBundle) - assert archived_bundle.tenant_id == run.tenant_id - assert archived_bundle.year == 2025 - assert archived_bundle.month == 3 - assert archived_bundle.workflow_run_count == 1 - assert archived_bundle.row_count == 2 - - def test_new_bundle_catalog_commit_failure_is_not_success(self, sqlite_session: Session) -> None: - archiver = WorkflowRunArchiver(days=90) - run = _run(str(uuid.uuid4())) - run.tenant_id = str(uuid.uuid4()) - storage = MagicMock() - storage.object_exists.return_value = False - empty_keys: list[str] = [] - storage.list_objects.return_value = empty_keys - table_data = {"workflow_runs": [{"id": run.id, "tenant_id": run.tenant_id}]} - - with ( - patch.object(sqlite_session, "commit", side_effect=RuntimeError("catalog commit failed")), - patch.object(sqlite_session, "rollback", wraps=sqlite_session.rollback) as rollback, - patch.object(archiver, "_lock_runs_for_archive", return_value=[run]), - patch.object(archiver, "_extract_bundle_data", return_value=table_data), - ): - result = archiver._archive_bundle(sqlite_session, storage, [run]) - - assert result.success is False - assert result.error == "catalog commit failed" - rollback.assert_called_once() - - def test_index_skips_all_already_archived_runs(self, sqlite_session: Session) -> None: - archiver = WorkflowRunArchiver(days=90) - run = _run() - index_key, index_payload = self._index_payload(archiver, ["run-1"], run) - storage = FakeArchiveStorage({index_key: index_payload}) - - result = archiver._archive_bundle(sqlite_session, storage, [run]) - - assert result.success is True - assert result.skipped is True - assert result.run_count == 0 - assert result.skipped_run_count == 1 - assert result.error == "all runs already archived in shard index" - - def test_index_filters_duplicate_runs_before_archive(self, sqlite_session: Session) -> None: - archiver = WorkflowRunArchiver(days=90) - archived_run = _run() - new_run = _run("run-2") - new_run.created_at = datetime.datetime(2025, 3, 15, 11, 0, 0) - index_key, index_payload = self._index_payload(archiver, ["run-1"], archived_run) - storage = FakeArchiveStorage({index_key: index_payload}) - - with ( - patch.object(archiver, "_lock_runs_for_archive", return_value=[new_run]) as lock_runs, - patch.object(archiver, "_extract_bundle_data", return_value={"workflow_runs": [{"id": "run-2"}]}), - ): - result = archiver._archive_bundle(sqlite_session, storage, [archived_run, new_run]) - - assert result.success is True - assert result.skipped is False - assert result.run_count == 1 - assert result.skipped_run_count == 1 - lock_runs.assert_called_once_with(ANY, ["run-2"]) - manifest_keys = [key for key in storage.objects if key.endswith("/manifest.json")] - assert len(manifest_keys) == 1 diff --git a/api/tests/unit_tests/services/test_account_activation_service.py b/api/tests/unit_tests/services/test_account_activation_service.py index e4a83d7be40..54dce1f790a 100644 --- a/api/tests/unit_tests/services/test_account_activation_service.py +++ b/api/tests/unit_tests/services/test_account_activation_service.py @@ -25,7 +25,7 @@ from services.entities.account_activation_entities import ( InvitationLookup, InvitationToken, ) -from tests.unit_tests.model_factories import make_account, make_tenant +from tests.helpers.unit.model_factories import make_account, make_tenant def _lookup(email: str | None = "invitee@example.com") -> InvitationLookup: diff --git a/api/tests/unit_tests/services/test_account_avatar_service.py b/api/tests/unit_tests/services/test_account_avatar_service.py index 6ee566af3c6..5228aaccf7e 100644 --- a/api/tests/unit_tests/services/test_account_avatar_service.py +++ b/api/tests/unit_tests/services/test_account_avatar_service.py @@ -6,7 +6,7 @@ from machinery.context import RequestContext from services.account_avatar_file_gateway import SQLAlchemyAccountAvatarFileGateway from services.account_avatar_service import AccountAvatarService from services.account_errors import AvatarFileNotFoundError -from tests.unit_tests.model_factories import make_upload_file +from tests.helpers.unit.model_factories import make_upload_file def _context() -> RequestContext: diff --git a/api/tests/unit_tests/services/test_account_change_email_service.py b/api/tests/unit_tests/services/test_account_change_email_service.py index eb3d7e24380..3937ff193ba 100644 --- a/api/tests/unit_tests/services/test_account_change_email_service.py +++ b/api/tests/unit_tests/services/test_account_change_email_service.py @@ -27,7 +27,7 @@ from services.entities.account_entities import ( AccountChangeEmailOldEmailVerifiedToken, AccountChangeEmailPhase, ) -from tests.unit_tests.model_factories import make_account +from tests.helpers.unit.model_factories import make_account def _context() -> RequestContext: diff --git a/api/tests/unit_tests/services/test_account_deletion_service.py b/api/tests/unit_tests/services/test_account_deletion_service.py index a678ddb9946..8281fb4fd20 100644 --- a/api/tests/unit_tests/services/test_account_deletion_service.py +++ b/api/tests/unit_tests/services/test_account_deletion_service.py @@ -19,7 +19,7 @@ from services.account.adapters import ( ) from services.account_deletion_service import AccountDeletionService from services.account_errors import InvalidAccountDeletionVerificationError -from tests.unit_tests.model_factories import make_account, make_tenant +from tests.helpers.unit.model_factories import make_account, make_tenant def _context() -> RequestContext: diff --git a/api/tests/unit_tests/services/test_account_education_service.py b/api/tests/unit_tests/services/test_account_education_service.py index d1467635715..628537f1900 100644 --- a/api/tests/unit_tests/services/test_account_education_service.py +++ b/api/tests/unit_tests/services/test_account_education_service.py @@ -21,7 +21,7 @@ from services.entities.account_entities import ( AccountEducationStatus, AccountEducationVerification, ) -from tests.unit_tests.model_factories import make_account +from tests.helpers.unit.model_factories import make_account def _context() -> RequestContext: diff --git a/api/tests/unit_tests/services/test_account_initialization_service.py b/api/tests/unit_tests/services/test_account_initialization_service.py index 36877dc9000..1d69a3d9bc5 100644 --- a/api/tests/unit_tests/services/test_account_initialization_service.py +++ b/api/tests/unit_tests/services/test_account_initialization_service.py @@ -19,7 +19,7 @@ from services.account_initialization_service import AccountInitializationService from services.entities.account_entities import ( AccountInitialization, ) -from tests.unit_tests.model_factories import make_account +from tests.helpers.unit.model_factories import make_account def _context() -> RequestContext: diff --git a/api/tests/unit_tests/services/test_account_integration_service.py b/api/tests/unit_tests/services/test_account_integration_service.py index 6c6d5ff0cb5..883273056d7 100644 --- a/api/tests/unit_tests/services/test_account_integration_service.py +++ b/api/tests/unit_tests/services/test_account_integration_service.py @@ -9,7 +9,7 @@ from machinery.context import RequestContext from models.account import AccountIntegrate from repositories.account_integration_repository import SQLAlchemyAccountIntegrationRepository from services.account_integration_service import AccountIntegrationService -from tests.unit_tests.model_factories import make_account +from tests.helpers.unit.model_factories import make_account def test_list_merges_configured_providers_with_persisted_integrations( diff --git a/api/tests/unit_tests/services/test_account_password_service.py b/api/tests/unit_tests/services/test_account_password_service.py index e754ac0deb6..5c2c162d4c2 100644 --- a/api/tests/unit_tests/services/test_account_password_service.py +++ b/api/tests/unit_tests/services/test_account_password_service.py @@ -10,7 +10,7 @@ from services.account_errors import AccountNotFoundError, CurrentAccountPassword from services.account_password_hasher import DefaultAccountPasswordHasher from services.account_password_service import AccountPasswordService from services.entities.account_entities import AccountPasswordDigest -from tests.unit_tests.model_factories import make_account +from tests.helpers.unit.model_factories import make_account def _context() -> RequestContext: diff --git a/api/tests/unit_tests/services/test_account_profile_service.py b/api/tests/unit_tests/services/test_account_profile_service.py index 007d97dcae4..dead0c07b7e 100644 --- a/api/tests/unit_tests/services/test_account_profile_service.py +++ b/api/tests/unit_tests/services/test_account_profile_service.py @@ -9,7 +9,7 @@ from repositories.account.repository import SQLAlchemyAccountRepository from services.account_errors import AccountNotFoundError from services.account_profile_service import AccountProfileService from services.entities.account_entities import AccountProfileChanges -from tests.unit_tests.model_factories import make_account +from tests.helpers.unit.model_factories import make_account def _context() -> RequestContext: diff --git a/api/tests/unit_tests/services/test_agent_llm_inner_service.py b/api/tests/unit_tests/services/test_agent_llm_inner_service.py index f3ddd39868a..2ca5946e44c 100644 --- a/api/tests/unit_tests/services/test_agent_llm_inner_service.py +++ b/api/tests/unit_tests/services/test_agent_llm_inner_service.py @@ -22,7 +22,7 @@ from models.model import App, AppMode from models.provider import ProviderType from services.agent_llm_inner_service import AgentLLMInnerService, AgentLLMInnerServiceError, PreparedAgentLLMInvocation from services.entities.agent_llm_inner import AgentLLMInvokeCaller, AgentLLMInvokeRequest, AgentLLMInvokeTarget -from tests.unit_tests.core.model_fixtures import make_model_instance +from tests.helpers.unit.core.model_fixtures import make_model_instance def _request() -> AgentLLMInvokeRequest: diff --git a/api/tests/unit_tests/services/test_agent_service_logs.py b/api/tests/unit_tests/services/test_agent_service.py similarity index 100% rename from api/tests/unit_tests/services/test_agent_service_logs.py rename to api/tests/unit_tests/services/test_agent_service.py diff --git a/api/tests/unit_tests/services/test_aliyun_tracing_config_gateway.py b/api/tests/unit_tests/services/test_aliyun_tracing_config_gateway.py deleted file mode 100644 index b8dc0fca5e1..00000000000 --- a/api/tests/unit_tests/services/test_aliyun_tracing_config_gateway.py +++ /dev/null @@ -1,265 +0,0 @@ -"""Real RSA/SQL/controller Aliyun configuration through a recorded HTTP boundary.""" - -import inspect -import json -import traceback -from dataclasses import dataclass, field, replace -from typing import Never -from urllib.parse import quote -from uuid import UUID, uuid4 - -import httpx -import pytest -from dify_trace_aliyun.aliyun_trace import AliyunDataTrace -from flask import Flask -from opentelemetry.proto.collector.trace.v1.trace_service_pb2 import ExportTraceServiceRequest -from sqlalchemy.orm import Session, sessionmaker - -from controllers.console.app import ops_trace as controller -from controllers.console.app.error import InvalidTracingConfigError -from core.helper import ssrf_proxy -from core.ops.ops_trace_manager import OpsTraceManager -from extensions.ext_application_services import ApplicationServices -from machinery.context import RequestContext -from repositories.app.console_repository import ConsoleAppRepository -from repositories.app.tracing_config_repository import SQLAlchemyAppTracingConfigRepository -from services.app_tracing_config_gateway import OpsTraceManagerGateway -from services.app_tracing_config_service import ( - AppTracingConfigInvalidConfigurationError, - AppTracingConfigRecord, - AppTracingConfigService, - AppTracingConfigVerificationFailedError, -) -from tests.unit_tests.model_factories import make_app - -pytest_plugins = ("tests.unit_tests.services.tracing.test_nullable_tracing_credentials",) - - -@dataclass -class AliyunCase: - gateway: OpsTraceManagerGateway - service: AppTracingConfigService - store: SQLAlchemyAppTracingConfigRepository - reader: ConsoleAppRepository - context: RequestContext - app_id: str - verification_calls: list[tuple[str, dict[str, object]]] = field(default_factory=list) - verification_succeeds: bool = True - - def read(self, provider: str = "aliyun") -> AppTracingConfigRecord | None: - return self.store.get( - workspace_id=self.context.active_workspace_id, app_id=self.app_id, tracing_provider=provider - ) - - -@pytest.fixture -def case(workspace: str, sqlite_session_factory: sessionmaker[Session], monkeypatch: pytest.MonkeyPatch) -> AliyunCase: - app_id = str(uuid4()) - with sqlite_session_factory.begin() as session: - session.add( - make_app( - app_id=app_id, - tenant_id=workspace, - tracing=json.dumps( - { - "enabled": True, - "tracing_provider": "langsmith", - "capture_content": True, - "capture_agent": False, - } - ), - ) - ) - gateway = OpsTraceManagerGateway() - store = SQLAlchemyAppTracingConfigRepository(session_factory=sqlite_session_factory) - old = OpsTraceManager.encrypt_tracing_config( - workspace, "langsmith", {"api_key": "synthetic-other", "project": "other", "endpoint": "https://other.test"} - ) - assert store.create(workspace_id=workspace, app_id=app_id, tracing_provider="langsmith", tracing_config=old) - result = AliyunCase( - gateway, - AppTracingConfigService(configs=store, provider=gateway), - store, - ConsoleAppRepository(session_factory=sqlite_session_factory), - RequestContext("request", None, "account", workspace), - app_id, - ) - - def verification_boundary(method: str, url: str, **options: object) -> httpx.Response: - assert method == "POST" - content = options["content"] - assert isinstance(content, bytes) - payload = ExportTraceServiceRequest.FromString(content) - assert len(payload.resource_spans) == 1 - assert len(payload.resource_spans[0].scope_spans) == 0 - assert [attribute.key for attribute in payload.resource_spans[0].resource.attributes] == ["service.name"] - result.verification_calls.append((url, dict(options))) - status = 200 if result.verification_succeeds else 403 - return httpx.Response(status, content=b'"success"', request=httpx.Request(method, url)) - - def forbidden(*_args: object, **_kwargs: object) -> Never: - raise AssertionError("Configuration must not initialize a legacy exporter or fetch its console link") - - monkeypatch.setattr(ssrf_proxy, "make_request", verification_boundary) - monkeypatch.setattr(OpsTraceManager, "check_trace_config_is_effective", forbidden) - monkeypatch.setattr(OpsTraceManager, "get_trace_config_project_url", forbidden) - monkeypatch.setattr(AliyunDataTrace, "__init__", forbidden) - return result - - -def test_post_masked_patch_and_get_preserve_real_ciphertext_and_policy(case: AliyunCase) -> None: - other = case.read("langsmith") - submitted = { - "endpoint": " HTTP://TRACING-ANALYSIS-DC-QD.ALIYUNCS.COM/reverse/ ", - "app_name": " 中文 app_Name ", - "license_key": "literal%2F@synthetic", - } - original = dict(submitted) - case.service.create(case.context, case.app_id, "aliyun", submitted, activate=True) - before = case.read() - assert before is not None - assert before.tracing_config is not None - stored = before.tracing_config - assert submitted == original - assert stored["endpoint"] == "http://tracing-analysis-dc-qd.aliyuncs.com/reverse" - assert stored["app_name"] == "中文 app_Name" - assert stored["license_key"] != submitted["license_key"] - assert OpsTraceManager.decrypt_tracing_config(case.context.active_workspace_id, "aliyun", stored)[ - "license_key" - ] == ("literal%2F@synthetic") - shown = case.service.get(case.context, case.app_id, "aliyun") - assert shown is not None - assert shown.tracing_config is not None - masked = shown.tracing_config["license_key"] - assert isinstance(masked, str) - assert "*" in masked - assert masked != submitted["license_key"] - assert shown.tracing_config["project_url"] == "https://arms.console.aliyun.com/#/llm" - assert len(case.verification_calls) == 1 - case.service.update( - case.context, - case.app_id, - "aliyun", - { - "endpoint": "https://project.region.log.aliyuncs.com/prefix/", - "app_name": " second_app ", - "license_key": masked, - }, - ) - after = case.read() - assert after is not None - assert after.tracing_config is not None - assert after.tracing_config["license_key"] == stored["license_key"] - assert after.tracing_config["app_name"] == "second_app" - assert after.tracing_config["endpoint"] == "https://project.region.log.aliyuncs.com/prefix" - assert case.verification_calls[-1][0] == ( - "https://project.region.log.aliyuncs.com/prefix/adapt_literal%252F@synthetic/api/v1/traces" - ) - assert case.service.get(case.context, case.app_id, "aliyun") is not None - assert len(case.verification_calls) == 2 - assert case.read("langsmith") == other - policy = case.reader.get_trace(case.context, case.app_id) - assert policy.tracing_provider == "aliyun" - assert policy.enabled - assert policy.capture_content is True - assert policy.capture_agent is False - - -@pytest.mark.parametrize("update", [False, True]) -@pytest.mark.parametrize("failure", ["invalid-schema", "verification-failed"]) -def test_failed_save_preserves_rows_and_active_policy(case: AliyunCase, update: bool, failure: str) -> None: - original = {"endpoint": "https://host.test", "app_name": "original", "license_key": "synthetic-secret"} - if update: - case.service.create(case.context, case.app_id, "aliyun", original) - before, other, policy = case.read(), case.read("langsmith"), case.reader.get_trace(case.context, case.app_id) - submitted = {**original, "app_name": "changed"} - if failure == "invalid-schema": - submitted["endpoint"] = "https://host.test/adapt_DO-NOT-PERSIST/api/otlp/traces" - expected = AppTracingConfigInvalidConfigurationError - else: - case.verification_succeeds = False - expected = AppTracingConfigVerificationFailedError - save = case.service.update if update else case.service.create - if update: - submitted["license_key"] = "******" - with pytest.raises(expected): - save(case.context, case.app_id, "aliyun", submitted, activate=True) - assert case.read() == before - assert case.read("langsmith") == other - assert case.reader.get_trace(case.context, case.app_id) == policy - - -@pytest.mark.parametrize("cached_url", [None, "https://wrong.test/stale"]) -def test_historical_get_is_pure_even_for_a_now_invalid_ingestion_url(case: AliyunCase, cached_url: str | None) -> None: - stored = OpsTraceManager.encrypt_tracing_config( - case.context.active_workspace_id, - "aliyun", - { - "endpoint": "https://host.test/adapt_old/api/otlp/traces", - "app_name": "historical_name", - "license_key": "synthetic-key", - }, - ) - if cached_url: - stored["project_url"] = cached_url - before = dict(stored) - result = case.gateway.present_config( - workspace_id=case.context.active_workspace_id, tracing_provider="aliyun", tracing_config=stored - ) - assert result["project_url"] == "https://arms.console.aliyun.com/#/llm" - assert isinstance(result["license_key"], str) - assert "*" in result["license_key"] - assert stored == before - assert case.verification_calls == [] - - -@pytest.mark.parametrize("update", [False, True]) -@pytest.mark.parametrize("encoded", [False, True]) -def test_controller_rejects_embedded_license_without_echo_or_state_change( - case: AliyunCase, - app: Flask, - account_application_services: ApplicationServices, - monkeypatch: pytest.MonkeyPatch, - caplog: pytest.LogCaptureFixture, - update: bool, - encoded: bool, -) -> None: - old_key = "SYNTHETIC_OLD_LICENSE" - original = {"endpoint": "https://host.test", "app_name": "old-app", "license_key": old_key} - if update: - case.service.create(case.context, case.app_id, "aliyun", original, activate=True) - before, other, policy = case.read(), case.read("langsmith"), case.reader.get_trace(case.context, case.app_id) - before_checks = len(case.verification_calls) - services = replace(account_application_services, app_tracing_configs=case.service) - monkeypatch.setattr(controller, "application_services", lambda: services) - key = "SYNTHETIC_PRIVATE_LICENSE@+/%2F" - url_key = quote(key, safe="") if encoded else key - payload = controller.TraceConfigPayload( - tracing_provider="aliyun", - tracing_config={ - "endpoint": f"https://host.test/adapt_{url_key}/api/otlp/traces", - "app_name": "should-not-save", - "license_key": "******" if update else key, - }, - activate=True, - ) - save = inspect.unwrap(controller.TraceAppConfigApi.patch if update else controller.TraceAppConfigApi.post) - with app.test_request_context(): - with pytest.raises(InvalidTracingConfigError) as caught: - save(controller.TraceAppConfigApi(), payload, case.context, UUID(case.app_id)) - # The real controller exposes this fixed BaseHTTPException body. Neither - # Pydantic inputs nor a credential-bearing exception chain reaches it. - assert caught.value.data == { - "code": "invalid_tracing_config", - "message": "The tracing configuration is invalid.", - "status": 400, - } - rendered = json.dumps(caught.value.data) + str(caught.value) + repr(caught.value) - rendered += "".join(traceback.format_exception(caught.value)) + caplog.text - assert key not in rendered - assert quote(key, safe="") not in rendered - assert old_key not in rendered - assert case.read() == before - assert case.read("langsmith") == other - assert case.reader.get_trace(case.context, case.app_id) == policy - assert len(case.verification_calls) == before_checks diff --git a/api/tests/unit_tests/services/test_annotation_import_literal_na.py b/api/tests/unit_tests/services/test_annotation_import_literal_na.py deleted file mode 100644 index e2035a4f473..00000000000 --- a/api/tests/unit_tests/services/test_annotation_import_literal_na.py +++ /dev/null @@ -1,47 +0,0 @@ -from io import BytesIO -from types import SimpleNamespace -from unittest.mock import MagicMock, patch - -import pytest -from werkzeug.datastructures import FileStorage - -import services.annotation_service as annotation_service_module -from enums import DeploymentEdition -from services.annotation_service import AppAnnotationService -from tests.unit_tests.config_override import config_overrides_context - - -def test_batch_import_preserves_literal_na(monkeypatch: pytest.MonkeyPatch) -> None: - account = SimpleNamespace(id="account-1") - monkeypatch.setattr( - annotation_service_module, - "current_account_with_tenant", - lambda: (account, "tenant-1"), - ) - - session = MagicMock() - session.scalar.return_value = SimpleNamespace(id="app-1") - file = FileStorage( - stream=BytesIO(b"question,answer\nNA,valid answer\nvalid question,NA\n"), - filename="annotations.csv", - content_type="text/csv", - ) - - with ( - patch.object(annotation_service_module, "batch_import_annotations_task") as task, - patch.object(annotation_service_module, "redis_client"), - patch.object(annotation_service_module.uuid, "uuid4", return_value="job-1"), - config_overrides_context( - DEPLOYMENT_EDITION=DeploymentEdition.COMMUNITY, - ANNOTATION_IMPORT_MAX_RECORDS=5, - ANNOTATION_IMPORT_MIN_RECORDS=1, - ), - ): - result = AppAnnotationService.batch_import_app_annotations("app-1", file, session) - - assert result == {"job_id": "job-1", "job_status": "waiting", "record_count": 2} - task.delay.assert_called_once() - assert task.delay.call_args.args[1] == [ - {"question": "NA", "answer": "valid answer"}, - {"question": "valid question", "answer": "NA"}, - ] diff --git a/api/tests/unit_tests/services/test_annotation_service.py b/api/tests/unit_tests/services/test_annotation_service.py index ac2a69b0f40..02c6ce22a66 100644 --- a/api/tests/unit_tests/services/test_annotation_service.py +++ b/api/tests/unit_tests/services/test_annotation_service.py @@ -1,9 +1,3 @@ -"""SQLite-backed unit tests for :mod:`services.annotation_service`.""" - -# Test functions explicitly request the identity fixture to make the patched -# request context visible at each service call, even when they do not inspect it. -# ruff: noqa: ARG002 - from __future__ import annotations import logging @@ -38,26 +32,54 @@ from models.model import ( ) from services.annotation_service import AppAnnotationService from services.app_ref_service import AnnotationRef, AppRef -from tests.unit_tests.config_override import config_overrides_context +from tests.helpers.unit.config_override import config_overrides_context + + +class TestAnnotationImportLiteralNaCases: + @staticmethod + def test_batch_import_preserves_literal_na(monkeypatch: pytest.MonkeyPatch) -> None: + account = SimpleNamespace(id="account-1") + monkeypatch.setattr( + annotation_service_module, + "current_account_with_tenant", + lambda: (account, "tenant-1"), + ) + + session = MagicMock() + session.scalar.return_value = SimpleNamespace(id="app-1") + file = FileStorage( + stream=BytesIO(b"question,answer\nNA,valid answer\nvalid question,NA\n"), + filename="annotations.csv", + content_type="text/csv", + ) + + with ( + patch.object(annotation_service_module, "batch_import_annotations_task") as task, + patch.object(annotation_service_module, "redis_client"), + patch.object(annotation_service_module.uuid, "uuid4", return_value="job-1"), + config_overrides_context( + DEPLOYMENT_EDITION=DeploymentEdition.COMMUNITY, + ANNOTATION_IMPORT_MAX_RECORDS=5, + ANNOTATION_IMPORT_MIN_RECORDS=1, + ), + ): + result = AppAnnotationService.batch_import_app_annotations("app-1", file, session) + + assert result == {"job_id": "job-1", "job_status": "waiting", "record_count": 2} + task.delay.assert_called_once() + assert task.delay.call_args.args[1] == [ + {"question": "NA", "answer": "valid answer"}, + {"question": "valid question", "answer": "NA"}, + ] + + +"""SQLite-backed unit tests for :mod:`services.annotation_service`.""" + TENANT_ID = "tenant-1" OTHER_TENANT_ID = "tenant-2" -@pytest.fixture -def current_user(monkeypatch: pytest.MonkeyPatch) -> Account: - """Install a real account model as the request identity.""" - - account = Account(name="Annotation Tester", email="annotation@example.com") - account.id = "account-1" - monkeypatch.setattr( - annotation_service_module, - "current_account_with_tenant", - lambda: (account, TENANT_ID), - ) - return account - - def _persist_app( session: Session, *, @@ -217,750 +239,773 @@ def _observer_get(factory: sessionmaker[Session], model: type[Any], identifier: return observer.get(model, identifier) -class TestAppAnnotationServiceUpsert: - def test_rejects_missing_or_cross_tenant_app(self, sqlite_session: Session, current_user: Account) -> None: - _persist_app(sqlite_session, app_id="other-app", tenant_id=OTHER_TENANT_ID) - - with pytest.raises(NotFound): - AppAnnotationService.up_insert_app_annotation_from_message( - {"answer": "hello", "question": "q"}, "other-app", sqlite_session - ) - - def test_validates_answer_and_question(self, sqlite_session: Session, current_user: Account) -> None: - app = _persist_app(sqlite_session) - - with pytest.raises(ValueError, match="answer.*content"): - AppAnnotationService.up_insert_app_annotation_from_message({"question": "q"}, app.id, sqlite_session) - with pytest.raises(ValueError, match="question"): - AppAnnotationService.up_insert_app_annotation_from_message({"answer": "a"}, app.id, sqlite_session) - - def test_rejects_message_from_another_app(self, sqlite_session: Session, current_user: Account) -> None: - app = _persist_app(sqlite_session) - other_app = _persist_app(sqlite_session, app_id="app-2") - message = _persist_message(sqlite_session, other_app) - - with pytest.raises(NotFound, match="Message"): - AppAnnotationService.up_insert_app_annotation_from_message( - {"answer": "hello", "message_id": message.id}, app.id, sqlite_session - ) - - def test_updates_existing_message_annotation_and_enqueues_index( - self, - sqlite_session: Session, - sqlite_session_factory: sessionmaker[Session], - current_user: Account, - ) -> None: - app = _persist_app(sqlite_session) - message = _persist_message(sqlite_session, app) - annotation = _persist_annotation(sqlite_session, app, message_id=message.id, content="old") - setting = _persist_setting(sqlite_session, app) - - with patch.object(annotation_service_module, "add_annotation_to_index_task") as task: - result = AppAnnotationService.up_insert_app_annotation_from_message( - {"answer": "updated", "message_id": message.id}, app.id, sqlite_session - ) - - assert result.id == annotation.id - stored = _observer_get(sqlite_session_factory, MessageAnnotation, annotation.id) - assert stored is not None - assert (stored.question, stored.content) == (message.query, "updated") - task.delay.assert_called_once_with( - annotation.id, message.query, TENANT_ID, app.id, setting.collection_binding_id - ) - - def test_creates_message_annotation_without_setting( - self, - sqlite_session: Session, - sqlite_session_factory: sessionmaker[Session], - current_user: Account, - ) -> None: - app = _persist_app(sqlite_session) - message = _persist_message(sqlite_session, app) - - with patch.object(annotation_service_module, "add_annotation_to_index_task") as task: - result = AppAnnotationService.up_insert_app_annotation_from_message( - {"answer": "hello", "question": "override", "message_id": message.id}, app.id, sqlite_session - ) - - stored = _observer_get(sqlite_session_factory, MessageAnnotation, result.id) - assert stored is not None - assert (stored.app_id, stored.message_id, stored.question, stored.content) == ( - app.id, - message.id, - "override", - "hello", - ) - assert stored.account_id == current_user.id - task.delay.assert_not_called() - - def test_creates_direct_annotation_and_enqueues_index( - self, - sqlite_session: Session, - sqlite_session_factory: sessionmaker[Session], - current_user: Account, - ) -> None: - app = _persist_app(sqlite_session) - setting = _persist_setting(sqlite_session, app) - - with patch.object(annotation_service_module, "add_annotation_to_index_task") as task: - result = AppAnnotationService.up_insert_app_annotation_from_message( - {"answer": "hello", "question": "q1"}, app.id, sqlite_session - ) - - stored = _observer_get(sqlite_session_factory, MessageAnnotation, result.id) - assert stored is not None - assert (stored.conversation_id, stored.message_id, stored.question, stored.content) == ( - None, - None, - "q1", - "hello", - ) - task.delay.assert_called_once_with(result.id, "q1", TENANT_ID, app.id, setting.collection_binding_id) - - -class TestAppAnnotationServiceEnableDisable: - def test_enable_returns_processing_on_cache_hit(self, current_user: Account) -> None: - args = {"score_threshold": 0.5, "embedding_provider_name": "p", "embedding_model_name": "m"} - with ( - patch.object(annotation_service_module, "redis_client") as redis, - patch.object(annotation_service_module, "enable_annotation_reply_task") as task, - ): - redis.get.return_value = "job-1" - result = AppAnnotationService.enable_app_annotation(args, "app-1") - - assert result == {"job_id": "job-1", "job_status": "processing"} - task.delay.assert_not_called() - - def test_enable_enqueues_on_cache_miss(self, current_user: Account) -> None: - args = {"score_threshold": 0.5, "embedding_provider_name": "p", "embedding_model_name": "m"} - with ( - patch.object(annotation_service_module, "redis_client") as redis, - patch.object(annotation_service_module.uuid, "uuid4", return_value="uuid-1"), - patch.object(annotation_service_module, "enable_annotation_reply_task") as task, - ): - redis.get.return_value = None - result = AppAnnotationService.enable_app_annotation(args, "app-1") - - assert result == {"job_id": "uuid-1", "job_status": "waiting"} - redis.setnx.assert_called_once_with("enable_app_annotation_job_uuid-1", "waiting") - task.delay.assert_called_once_with("uuid-1", "app-1", current_user.id, TENANT_ID, 0.5, "p", "m") - - def test_disable_returns_processing_on_cache_hit(self, current_user: Account) -> None: - with ( - patch.object(annotation_service_module, "redis_client") as redis, - patch.object(annotation_service_module, "disable_annotation_reply_task") as task, - ): - redis.get.return_value = "job-2" - result = AppAnnotationService.disable_app_annotation("app-1") - - assert result == {"job_id": "job-2", "job_status": "processing"} - task.delay.assert_not_called() - - def test_disable_enqueues_on_cache_miss(self, current_user: Account) -> None: - with ( - patch.object(annotation_service_module, "redis_client") as redis, - patch.object(annotation_service_module.uuid, "uuid4", return_value="uuid-2"), - patch.object(annotation_service_module, "disable_annotation_reply_task") as task, - ): - redis.get.return_value = None - result = AppAnnotationService.disable_app_annotation("app-1") - - assert result == {"job_id": "uuid-2", "job_status": "waiting"} - redis.setnx.assert_called_once_with("disable_app_annotation_job_uuid-2", "waiting") - task.delay.assert_called_once_with("uuid-2", "app-1", TENANT_ID) - - -class TestAppAnnotationServiceListAndExport: - def test_list_rejects_cross_tenant_app(self, sqlite_session: Session, current_user: Account) -> None: - app = _persist_app(sqlite_session, tenant_id=OTHER_TENANT_ID) - - with pytest.raises(NotFound): - AppAnnotationService.get_annotation_list_by_app_id(app.id, 1, 10, "", sqlite_session) - - def test_list_filters_orders_and_paginates(self, sqlite_session: Session, current_user: Account) -> None: - app = _persist_app(sqlite_session) - other_app = _persist_app(sqlite_session, app_id="app-2") - now = datetime(2026, 1, 1) - first = _persist_annotation(sqlite_session, app, annotation_id="ann-1", question="needle first", created_at=now) - second = _persist_annotation( - sqlite_session, - app, - annotation_id="ann-2", - question="other", - content="needle second", - created_at=now + timedelta(seconds=1), - ) - _persist_annotation(sqlite_session, app, annotation_id="ann-3", question="not matched") - _persist_annotation(sqlite_session, other_app, annotation_id="decoy", question="needle decoy") - - items, total = AppAnnotationService.get_annotation_list_by_app_id(app.id, 1, 1, "needle", sqlite_session) - - assert total == 2 - assert [item.id for item in items] == [second.id] - items, total = AppAnnotationService.get_annotation_list_by_app_id(app.id, 2, 1, "needle", sqlite_session) - assert total == 2 - assert [item.id for item in items] == [first.id] - - def test_list_without_keyword_is_app_scoped(self, sqlite_session: Session, current_user: Account) -> None: - app = _persist_app(sqlite_session) - other_app = _persist_app(sqlite_session, app_id="app-2") - expected = _persist_annotation(sqlite_session, app) - _persist_annotation(sqlite_session, other_app, annotation_id="decoy") - - items, total = AppAnnotationService.get_annotation_list_by_app_id(app.id, 1, 10, "", sqlite_session) - - assert total == 1 - assert [item.id for item in items] == [expected.id] - - def test_export_sanitizes_and_scopes_rows(self, sqlite_session: Session, current_user: Account) -> None: - app = _persist_app(sqlite_session) - other_app = _persist_app(sqlite_session, app_id="app-2") - first = _persist_annotation(sqlite_session, app, annotation_id="ann-1", question="=cmd", content="+1") - second = _persist_annotation(sqlite_session, app, annotation_id="ann-2", question="@bad", content="-2") - _persist_annotation(sqlite_session, other_app, annotation_id="decoy", question="=decoy") - - result = AppAnnotationService.export_annotation_list_by_app_id(app.id, sqlite_session) - - assert {annotation.id for annotation in result} == {first.id, second.id} - assert {(annotation.question, annotation.content) for annotation in result} == { - ("'=cmd", "'+1"), - ("'@bad", "'-2"), - } - - -class TestAppAnnotationServiceDirectManipulation: - def test_insert_rejects_cross_tenant_app_and_missing_question( - self, sqlite_session: Session, current_user: Account - ) -> None: - other_app = _persist_app(sqlite_session, tenant_id=OTHER_TENANT_ID) - with pytest.raises(NotFound): - AppAnnotationService.insert_app_annotation_directly( - {"answer": "hello", "question": "q"}, other_app.id, sqlite_session - ) - - app = _persist_app(sqlite_session, app_id="app-2") - with pytest.raises(ValueError, match="question"): - AppAnnotationService.insert_app_annotation_directly({"answer": "hello"}, app.id, sqlite_session) - - def test_insert_persists_and_enqueues_index( - self, - sqlite_session: Session, - sqlite_session_factory: sessionmaker[Session], - current_user: Account, - ) -> None: - app = _persist_app(sqlite_session) - setting = _persist_setting(sqlite_session, app) - - with patch.object(annotation_service_module, "add_annotation_to_index_task") as task: - result = AppAnnotationService.insert_app_annotation_directly( - {"answer": "hello", "question": "q1"}, app.id, sqlite_session - ) - - stored = _observer_get(sqlite_session_factory, MessageAnnotation, result.id) - assert stored is not None - assert (stored.app_id, stored.question, stored.content, stored.account_id) == ( - app.id, - "q1", - "hello", - current_user.id, - ) - task.delay.assert_called_once_with(result.id, "q1", TENANT_ID, app.id, setting.collection_binding_id) - - def test_update_is_app_scoped_and_validates_fields(self, sqlite_session: Session, current_user: Account) -> None: - app = _persist_app(sqlite_session) - other_app = _persist_app(sqlite_session, app_id="app-2") - annotation = _persist_annotation(sqlite_session, other_app) - - with pytest.raises(NotFound): - AppAnnotationService.update_app_annotation_directly( - {"answer": "a", "question": "q"}, _annotation_ref(app, annotation.id), sqlite_session - ) - - own_annotation = _persist_annotation(sqlite_session, app, annotation_id="own-ann") - with pytest.raises(ValueError, match="question"): - AppAnnotationService.update_app_annotation_directly( - {"answer": "a"}, _annotation_ref(app, own_annotation.id), sqlite_session - ) - with pytest.raises(ValueError, match="answer"): - AppAnnotationService.update_app_annotation_directly( - {"question": "q"}, _annotation_ref(app, own_annotation.id), sqlite_session - ) - - def test_update_persists_and_enqueues_index( - self, - sqlite_session: Session, - sqlite_session_factory: sessionmaker[Session], - current_user: Account, - ) -> None: - app = _persist_app(sqlite_session) - annotation = _persist_annotation(sqlite_session, app, content="old") - setting = _persist_setting(sqlite_session, app) - - with patch.object(annotation_service_module, "update_annotation_to_index_task") as task: - result = AppAnnotationService.update_app_annotation_directly( - {"answer": "new", "question": "new q"}, _annotation_ref(app, annotation.id), sqlite_session - ) - - stored = _observer_get(sqlite_session_factory, MessageAnnotation, annotation.id) - assert stored is not None - assert (stored.question, stored.content) == ("new q", "new") - assert result.id == stored.id - task.delay.assert_called_once_with(annotation.id, "new q", TENANT_ID, app.id, setting.collection_binding_id) - - def test_delete_is_app_scoped(self, sqlite_session: Session, current_user: Account) -> None: - app = _persist_app(sqlite_session) - other_app = _persist_app(sqlite_session, app_id="app-2") - annotation = _persist_annotation(sqlite_session, other_app) - - with pytest.raises(NotFound): - AppAnnotationService.delete_app_annotation(_annotation_ref(app, annotation.id), sqlite_session) - - assert sqlite_session.get(MessageAnnotation, annotation.id) is not None - - def test_delete_removes_annotation_and_histories_and_enqueues_index( - self, - sqlite_session: Session, - sqlite_session_factory: sessionmaker[Session], - current_user: Account, - ) -> None: - app = _persist_app(sqlite_session) - annotation = _persist_annotation(sqlite_session, app) - histories = [ - _persist_history(sqlite_session, app, annotation, history_id="history-1"), - _persist_history(sqlite_session, app, annotation, history_id="history-2"), - ] - setting = _persist_setting(sqlite_session, app) - - with patch.object(annotation_service_module, "delete_annotation_index_task") as task: - AppAnnotationService.delete_app_annotation(_annotation_ref(app, annotation.id), sqlite_session) - - with sqlite_session_factory() as observer: - assert observer.get(MessageAnnotation, annotation.id) is None - assert [observer.get(AppAnnotationHitHistory, history.id) for history in histories] == [None, None] - task.delay.assert_called_once_with(annotation.id, app.id, TENANT_ID, setting.collection_binding_id) - - def test_batch_delete_returns_zero_without_matching_rows( - self, sqlite_session: Session, current_user: Account - ) -> None: - app = _persist_app(sqlite_session) - other_app = _persist_app(sqlite_session, app_id="app-2") - _persist_annotation(sqlite_session, other_app, annotation_id="ann-1") - - result = AppAnnotationService.delete_app_annotations_in_batch(_app_ref(app), ["ann-1"], sqlite_session) - - assert result == {"deleted_count": 0} - - def test_batch_delete_scopes_rows_and_histories( - self, - sqlite_session: Session, - sqlite_session_factory: sessionmaker[Session], - current_user: Account, - ) -> None: - app = _persist_app(sqlite_session) - other_app = _persist_app(sqlite_session, app_id="app-2") - annotation1 = _persist_annotation(sqlite_session, app, annotation_id="ann-1") - annotation2 = _persist_annotation(sqlite_session, app, annotation_id="ann-2") - decoy = _persist_annotation(sqlite_session, other_app, annotation_id="ann-3") - history1 = _persist_history(sqlite_session, app, annotation1, history_id="history-1") - history2 = _persist_history(sqlite_session, app, annotation2, history_id="history-2") - decoy_history = _persist_history(sqlite_session, other_app, decoy, history_id="history-3") - setting = _persist_setting(sqlite_session, app) - - with patch.object(annotation_service_module, "delete_annotation_index_task") as task: - result = AppAnnotationService.delete_app_annotations_in_batch( - _app_ref(app), [annotation1.id, annotation2.id, decoy.id], sqlite_session - ) - - assert result == {"deleted_count": 2} - with sqlite_session_factory() as observer: - assert [observer.get(MessageAnnotation, item.id) for item in (annotation1, annotation2)] == [None, None] - assert [observer.get(AppAnnotationHitHistory, item.id) for item in (history1, history2)] == [None, None] - assert observer.get(MessageAnnotation, decoy.id) is not None - assert observer.get(AppAnnotationHitHistory, decoy_history.id) is not None - assert task.delay.call_count == 2 - task.delay.assert_any_call(annotation1.id, app.id, TENANT_ID, setting.collection_binding_id) - task.delay.assert_any_call(annotation2.id, app.id, TENANT_ID, setting.collection_binding_id) - - -class TestAppAnnotationServiceBatchImport: +class TestAnnotationServiceCases: @staticmethod - def _invoke( - sqlite_session: Session, - app: App, - *, - dataframe: Any, - content: bytes = b"question,answer\nq,a\n", - maximum: int = 5, - minimum: int = 1, - features: Any | None = None, - ) -> dict[str, Any]: - if features is None: + @pytest.fixture + def current_user(monkeypatch: pytest.MonkeyPatch) -> Account: + """Install a real account model as the request identity.""" + + account = Account(name="Annotation Tester", email="annotation@example.com") + account.id = "account-1" + monkeypatch.setattr( + annotation_service_module, + "current_account_with_tenant", + lambda: (account, TENANT_ID), + ) + return account + + class TestAppAnnotationServiceUpsert: + def test_rejects_missing_or_cross_tenant_app(self, sqlite_session: Session, current_user: Account) -> None: + _persist_app(sqlite_session, app_id="other-app", tenant_id=OTHER_TENANT_ID) + + with pytest.raises(NotFound): + AppAnnotationService.up_insert_app_annotation_from_message( + {"answer": "hello", "question": "q"}, "other-app", sqlite_session + ) + + def test_validates_answer_and_question(self, sqlite_session: Session, current_user: Account) -> None: + app = _persist_app(sqlite_session) + + with pytest.raises(ValueError, match="answer.*content"): + AppAnnotationService.up_insert_app_annotation_from_message({"question": "q"}, app.id, sqlite_session) + with pytest.raises(ValueError, match="question"): + AppAnnotationService.up_insert_app_annotation_from_message({"answer": "a"}, app.id, sqlite_session) + + def test_rejects_message_from_another_app(self, sqlite_session: Session, current_user: Account) -> None: + app = _persist_app(sqlite_session) + other_app = _persist_app(sqlite_session, app_id="app-2") + message = _persist_message(sqlite_session, other_app) + + with pytest.raises(NotFound, match="Message"): + AppAnnotationService.up_insert_app_annotation_from_message( + {"answer": "hello", "message_id": message.id}, app.id, sqlite_session + ) + + def test_updates_existing_message_annotation_and_enqueues_index( + self, + sqlite_session: Session, + sqlite_session_factory: sessionmaker[Session], + current_user: Account, + ) -> None: + app = _persist_app(sqlite_session) + message = _persist_message(sqlite_session, app) + annotation = _persist_annotation(sqlite_session, app, message_id=message.id, content="old") + setting = _persist_setting(sqlite_session, app) + + with patch.object(annotation_service_module, "add_annotation_to_index_task") as task: + result = AppAnnotationService.up_insert_app_annotation_from_message( + {"answer": "updated", "message_id": message.id}, app.id, sqlite_session + ) + + assert result.id == annotation.id + stored = _observer_get(sqlite_session_factory, MessageAnnotation, annotation.id) + assert stored is not None + assert (stored.question, stored.content) == (message.query, "updated") + task.delay.assert_called_once_with( + annotation.id, message.query, TENANT_ID, app.id, setting.collection_binding_id + ) + + def test_creates_message_annotation_without_setting( + self, + sqlite_session: Session, + sqlite_session_factory: sessionmaker[Session], + current_user: Account, + ) -> None: + app = _persist_app(sqlite_session) + message = _persist_message(sqlite_session, app) + + with patch.object(annotation_service_module, "add_annotation_to_index_task") as task: + result = AppAnnotationService.up_insert_app_annotation_from_message( + {"answer": "hello", "question": "override", "message_id": message.id}, app.id, sqlite_session + ) + + stored = _observer_get(sqlite_session_factory, MessageAnnotation, result.id) + assert stored is not None + assert (stored.app_id, stored.message_id, stored.question, stored.content) == ( + app.id, + message.id, + "override", + "hello", + ) + assert stored.account_id == current_user.id + task.delay.assert_not_called() + + def test_creates_direct_annotation_and_enqueues_index( + self, + sqlite_session: Session, + sqlite_session_factory: sessionmaker[Session], + current_user: Account, + ) -> None: + app = _persist_app(sqlite_session) + setting = _persist_setting(sqlite_session, app) + + with patch.object(annotation_service_module, "add_annotation_to_index_task") as task: + result = AppAnnotationService.up_insert_app_annotation_from_message( + {"answer": "hello", "question": "q1"}, app.id, sqlite_session + ) + + stored = _observer_get(sqlite_session_factory, MessageAnnotation, result.id) + assert stored is not None + assert (stored.conversation_id, stored.message_id, stored.question, stored.content) == ( + None, + None, + "q1", + "hello", + ) + task.delay.assert_called_once_with(result.id, "q1", TENANT_ID, app.id, setting.collection_binding_id) + + class TestAppAnnotationServiceEnableDisable: + def test_enable_returns_processing_on_cache_hit(self, current_user: Account) -> None: + args = {"score_threshold": 0.5, "embedding_provider_name": "p", "embedding_model_name": "m"} + with ( + patch.object(annotation_service_module, "redis_client") as redis, + patch.object(annotation_service_module, "enable_annotation_reply_task") as task, + ): + redis.get.return_value = "job-1" + result = AppAnnotationService.enable_app_annotation(args, "app-1") + + assert result == {"job_id": "job-1", "job_status": "processing"} + task.delay.assert_not_called() + + def test_enable_enqueues_on_cache_miss(self, current_user: Account) -> None: + args = {"score_threshold": 0.5, "embedding_provider_name": "p", "embedding_model_name": "m"} + with ( + patch.object(annotation_service_module, "redis_client") as redis, + patch.object(annotation_service_module.uuid, "uuid4", return_value="uuid-1"), + patch.object(annotation_service_module, "enable_annotation_reply_task") as task, + ): + redis.get.return_value = None + result = AppAnnotationService.enable_app_annotation(args, "app-1") + + assert result == {"job_id": "uuid-1", "job_status": "waiting"} + redis.setnx.assert_called_once_with("enable_app_annotation_job_uuid-1", "waiting") + task.delay.assert_called_once_with("uuid-1", "app-1", current_user.id, TENANT_ID, 0.5, "p", "m") + + def test_disable_returns_processing_on_cache_hit(self, current_user: Account) -> None: + with ( + patch.object(annotation_service_module, "redis_client") as redis, + patch.object(annotation_service_module, "disable_annotation_reply_task") as task, + ): + redis.get.return_value = "job-2" + result = AppAnnotationService.disable_app_annotation("app-1") + + assert result == {"job_id": "job-2", "job_status": "processing"} + task.delay.assert_not_called() + + def test_disable_enqueues_on_cache_miss(self, current_user: Account) -> None: + with ( + patch.object(annotation_service_module, "redis_client") as redis, + patch.object(annotation_service_module.uuid, "uuid4", return_value="uuid-2"), + patch.object(annotation_service_module, "disable_annotation_reply_task") as task, + ): + redis.get.return_value = None + result = AppAnnotationService.disable_app_annotation("app-1") + + assert result == {"job_id": "uuid-2", "job_status": "waiting"} + redis.setnx.assert_called_once_with("disable_app_annotation_job_uuid-2", "waiting") + task.delay.assert_called_once_with("uuid-2", "app-1", TENANT_ID) + + class TestAppAnnotationServiceListAndExport: + def test_list_rejects_cross_tenant_app(self, sqlite_session: Session, current_user: Account) -> None: + app = _persist_app(sqlite_session, tenant_id=OTHER_TENANT_ID) + + with pytest.raises(NotFound): + AppAnnotationService.get_annotation_list_by_app_id(app.id, 1, 10, "", sqlite_session) + + def test_list_filters_orders_and_paginates(self, sqlite_session: Session, current_user: Account) -> None: + app = _persist_app(sqlite_session) + other_app = _persist_app(sqlite_session, app_id="app-2") + now = datetime(2026, 1, 1) + first = _persist_annotation( + sqlite_session, app, annotation_id="ann-1", question="needle first", created_at=now + ) + second = _persist_annotation( + sqlite_session, + app, + annotation_id="ann-2", + question="other", + content="needle second", + created_at=now + timedelta(seconds=1), + ) + _persist_annotation(sqlite_session, app, annotation_id="ann-3", question="not matched") + _persist_annotation(sqlite_session, other_app, annotation_id="decoy", question="needle decoy") + + items, total = AppAnnotationService.get_annotation_list_by_app_id(app.id, 1, 1, "needle", sqlite_session) + + assert total == 2 + assert [item.id for item in items] == [second.id] + items, total = AppAnnotationService.get_annotation_list_by_app_id(app.id, 2, 1, "needle", sqlite_session) + assert total == 2 + assert [item.id for item in items] == [first.id] + + def test_list_without_keyword_is_app_scoped(self, sqlite_session: Session, current_user: Account) -> None: + app = _persist_app(sqlite_session) + other_app = _persist_app(sqlite_session, app_id="app-2") + expected = _persist_annotation(sqlite_session, app) + _persist_annotation(sqlite_session, other_app, annotation_id="decoy") + + items, total = AppAnnotationService.get_annotation_list_by_app_id(app.id, 1, 10, "", sqlite_session) + + assert total == 1 + assert [item.id for item in items] == [expected.id] + + def test_export_sanitizes_and_scopes_rows(self, sqlite_session: Session, current_user: Account) -> None: + app = _persist_app(sqlite_session) + other_app = _persist_app(sqlite_session, app_id="app-2") + first = _persist_annotation(sqlite_session, app, annotation_id="ann-1", question="=cmd", content="+1") + second = _persist_annotation(sqlite_session, app, annotation_id="ann-2", question="@bad", content="-2") + _persist_annotation(sqlite_session, other_app, annotation_id="decoy", question="=decoy") + + result = AppAnnotationService.export_annotation_list_by_app_id(app.id, sqlite_session) + + assert {annotation.id for annotation in result} == {first.id, second.id} + assert {(annotation.question, annotation.content) for annotation in result} == { + ("'=cmd", "'+1"), + ("'@bad", "'-2"), + } + + class TestAppAnnotationServiceDirectManipulation: + def test_insert_rejects_cross_tenant_app_and_missing_question( + self, sqlite_session: Session, current_user: Account + ) -> None: + other_app = _persist_app(sqlite_session, tenant_id=OTHER_TENANT_ID) + with pytest.raises(NotFound): + AppAnnotationService.insert_app_annotation_directly( + {"answer": "hello", "question": "q"}, other_app.id, sqlite_session + ) + + app = _persist_app(sqlite_session, app_id="app-2") + with pytest.raises(ValueError, match="question"): + AppAnnotationService.insert_app_annotation_directly({"answer": "hello"}, app.id, sqlite_session) + + def test_insert_persists_and_enqueues_index( + self, + sqlite_session: Session, + sqlite_session_factory: sessionmaker[Session], + current_user: Account, + ) -> None: + app = _persist_app(sqlite_session) + setting = _persist_setting(sqlite_session, app) + + with patch.object(annotation_service_module, "add_annotation_to_index_task") as task: + result = AppAnnotationService.insert_app_annotation_directly( + {"answer": "hello", "question": "q1"}, app.id, sqlite_session + ) + + stored = _observer_get(sqlite_session_factory, MessageAnnotation, result.id) + assert stored is not None + assert (stored.app_id, stored.question, stored.content, stored.account_id) == ( + app.id, + "q1", + "hello", + current_user.id, + ) + task.delay.assert_called_once_with(result.id, "q1", TENANT_ID, app.id, setting.collection_binding_id) + + def test_update_is_app_scoped_and_validates_fields( + self, sqlite_session: Session, current_user: Account + ) -> None: + app = _persist_app(sqlite_session) + other_app = _persist_app(sqlite_session, app_id="app-2") + annotation = _persist_annotation(sqlite_session, other_app) + + with pytest.raises(NotFound): + AppAnnotationService.update_app_annotation_directly( + {"answer": "a", "question": "q"}, _annotation_ref(app, annotation.id), sqlite_session + ) + + own_annotation = _persist_annotation(sqlite_session, app, annotation_id="own-ann") + with pytest.raises(ValueError, match="question"): + AppAnnotationService.update_app_annotation_directly( + {"answer": "a"}, _annotation_ref(app, own_annotation.id), sqlite_session + ) + with pytest.raises(ValueError, match="answer"): + AppAnnotationService.update_app_annotation_directly( + {"question": "q"}, _annotation_ref(app, own_annotation.id), sqlite_session + ) + + def test_update_persists_and_enqueues_index( + self, + sqlite_session: Session, + sqlite_session_factory: sessionmaker[Session], + current_user: Account, + ) -> None: + app = _persist_app(sqlite_session) + annotation = _persist_annotation(sqlite_session, app, content="old") + setting = _persist_setting(sqlite_session, app) + + with patch.object(annotation_service_module, "update_annotation_to_index_task") as task: + result = AppAnnotationService.update_app_annotation_directly( + {"answer": "new", "question": "new q"}, _annotation_ref(app, annotation.id), sqlite_session + ) + + stored = _observer_get(sqlite_session_factory, MessageAnnotation, annotation.id) + assert stored is not None + assert (stored.question, stored.content) == ("new q", "new") + assert result.id == stored.id + task.delay.assert_called_once_with(annotation.id, "new q", TENANT_ID, app.id, setting.collection_binding_id) + + def test_delete_is_app_scoped(self, sqlite_session: Session, current_user: Account) -> None: + app = _persist_app(sqlite_session) + other_app = _persist_app(sqlite_session, app_id="app-2") + annotation = _persist_annotation(sqlite_session, other_app) + + with pytest.raises(NotFound): + AppAnnotationService.delete_app_annotation(_annotation_ref(app, annotation.id), sqlite_session) + + assert sqlite_session.get(MessageAnnotation, annotation.id) is not None + + def test_delete_removes_annotation_and_histories_and_enqueues_index( + self, + sqlite_session: Session, + sqlite_session_factory: sessionmaker[Session], + current_user: Account, + ) -> None: + app = _persist_app(sqlite_session) + annotation = _persist_annotation(sqlite_session, app) + histories = [ + _persist_history(sqlite_session, app, annotation, history_id="history-1"), + _persist_history(sqlite_session, app, annotation, history_id="history-2"), + ] + setting = _persist_setting(sqlite_session, app) + + with patch.object(annotation_service_module, "delete_annotation_index_task") as task: + AppAnnotationService.delete_app_annotation(_annotation_ref(app, annotation.id), sqlite_session) + + with sqlite_session_factory() as observer: + assert observer.get(MessageAnnotation, annotation.id) is None + assert [observer.get(AppAnnotationHitHistory, history.id) for history in histories] == [None, None] + task.delay.assert_called_once_with(annotation.id, app.id, TENANT_ID, setting.collection_binding_id) + + def test_batch_delete_returns_zero_without_matching_rows( + self, sqlite_session: Session, current_user: Account + ) -> None: + app = _persist_app(sqlite_session) + other_app = _persist_app(sqlite_session, app_id="app-2") + _persist_annotation(sqlite_session, other_app, annotation_id="ann-1") + + result = AppAnnotationService.delete_app_annotations_in_batch(_app_ref(app), ["ann-1"], sqlite_session) + + assert result == {"deleted_count": 0} + + def test_batch_delete_scopes_rows_and_histories( + self, + sqlite_session: Session, + sqlite_session_factory: sessionmaker[Session], + current_user: Account, + ) -> None: + app = _persist_app(sqlite_session) + other_app = _persist_app(sqlite_session, app_id="app-2") + annotation1 = _persist_annotation(sqlite_session, app, annotation_id="ann-1") + annotation2 = _persist_annotation(sqlite_session, app, annotation_id="ann-2") + decoy = _persist_annotation(sqlite_session, other_app, annotation_id="ann-3") + history1 = _persist_history(sqlite_session, app, annotation1, history_id="history-1") + history2 = _persist_history(sqlite_session, app, annotation2, history_id="history-2") + decoy_history = _persist_history(sqlite_session, other_app, decoy, history_id="history-3") + setting = _persist_setting(sqlite_session, app) + + with patch.object(annotation_service_module, "delete_annotation_index_task") as task: + result = AppAnnotationService.delete_app_annotations_in_batch( + _app_ref(app), [annotation1.id, annotation2.id, decoy.id], sqlite_session + ) + + assert result == {"deleted_count": 2} + with sqlite_session_factory() as observer: + assert [observer.get(MessageAnnotation, item.id) for item in (annotation1, annotation2)] == [None, None] + assert [observer.get(AppAnnotationHitHistory, item.id) for item in (history1, history2)] == [None, None] + assert observer.get(MessageAnnotation, decoy.id) is not None + assert observer.get(AppAnnotationHitHistory, decoy_history.id) is not None + assert task.delay.call_count == 2 + task.delay.assert_any_call(annotation1.id, app.id, TENANT_ID, setting.collection_binding_id) + task.delay.assert_any_call(annotation2.id, app.id, TENANT_ID, setting.collection_binding_id) + + class TestAppAnnotationServiceBatchImport: + @staticmethod + def _invoke( + sqlite_session: Session, + app: App, + *, + dataframe: Any, + content: bytes = b"question,answer\nq,a\n", + maximum: int = 5, + minimum: int = 1, + features: Any | None = None, + ) -> dict[str, Any]: + if features is None: + features = SimpleNamespace(annotation_quota_limit=None) + with ( + patch.object(annotation_service_module.pd, "read_csv", return_value=dataframe), + patch.object(annotation_service_module.FeatureService, "get_features", return_value=features), + config_overrides_context( + ANNOTATION_IMPORT_MAX_RECORDS=maximum, + ANNOTATION_IMPORT_MIN_RECORDS=minimum, + ), + ): + return AppAnnotationService.batch_import_app_annotations(app.id, _file(content), sqlite_session) + + def test_rejects_cross_tenant_app(self, sqlite_session: Session, current_user: Account) -> None: + app = _persist_app(sqlite_session, tenant_id=OTHER_TENANT_ID) + + with pytest.raises(NotFound): + AppAnnotationService.batch_import_app_annotations( + app.id, _file(b"question,answer\nq,a\n"), sqlite_session + ) + + @pytest.mark.parametrize( + ("dataframe", "content", "maximum", "minimum", "expected"), + [ + (pd.DataFrame({"q": ["only"]}), b"question\nq\n", 5, 1, "Invalid CSV format"), + (pd.DataFrame({"q": ["q"], "a": ["a"]}), b"", 5, 1, "empty or invalid"), + (pd.DataFrame({"q": ["q"], "a": ["a"]}), b"question,answer\nq,a\n", 5, 2, "at least"), + ( + pd.DataFrame({"q": ["q1", "q2"], "a": ["a1", "a2"]}), + b"question,answer\nq1,a1\nq2,a2\n", + 1, + 1, + "too many records", + ), + (pd.DataFrame({"q": ["nan"], "a": ["nan"]}), b"question,answer\nnan,nan\n", 5, 1, "at least"), + ( + pd.DataFrame({"q": ["q" * 2001], "a": ["a"]}), + b"question,answer\nq,a\n", + 5, + 1, + "Question at row", + ), + ( + pd.DataFrame({"q": ["q"], "a": ["a" * 10001]}), + b"question,answer\nq,a\n", + 5, + 1, + "Answer at row", + ), + ], + ) + def test_validation_errors( + self, + sqlite_session: Session, + current_user: Account, + dataframe: pd.DataFrame, + content: bytes, + maximum: int, + minimum: int, + expected: str, + ) -> None: + app = _persist_app(sqlite_session) + + result = self._invoke( + sqlite_session, + app, + dataframe=dataframe, + content=content, + maximum=maximum, + minimum=minimum, + ) + + assert expected in cast(str, result["error_msg"]) + + def test_skips_malformed_rows(self, sqlite_session: Session, current_user: Account) -> None: + app = _persist_app(sqlite_session) + malformed_row = MagicMock() + malformed_row.iloc.__getitem__.side_effect = IndexError() + dataframe = MagicMock() + dataframe.columns = ["q", "a"] + dataframe.iterrows.return_value = [(0, malformed_row)] + + result = self._invoke(sqlite_session, app, dataframe=dataframe) + + assert "at least" in cast(str, result["error_msg"]) + + @config_overrides_context(DEPLOYMENT_EDITION=DeploymentEdition.CLOUD) + def test_rejects_subscription_quota_overflow(self, sqlite_session: Session, current_user: Account) -> None: + app = _persist_app(sqlite_session) + features = SimpleNamespace( + annotation_quota_limit=SimpleNamespace(limit=1, size=1), + ) + + result = self._invoke( + sqlite_session, + app, + dataframe=pd.DataFrame({"q": ["q1"], "a": ["a1"]}), + features=features, + ) + + assert "exceeds the limit" in cast(str, result["error_msg"]) + + @config_overrides_context(DEPLOYMENT_EDITION=DeploymentEdition.COMMUNITY) + def test_valid_import_enqueues_job(self, sqlite_session: Session, current_user: Account) -> None: + app = _persist_app(sqlite_session) + dataframe = pd.DataFrame({"q": ["q1"], "a": ["a1"]}) features = SimpleNamespace(annotation_quota_limit=None) - with ( - patch.object(annotation_service_module.pd, "read_csv", return_value=dataframe), - patch.object(annotation_service_module.FeatureService, "get_features", return_value=features), - config_overrides_context( - ANNOTATION_IMPORT_MAX_RECORDS=maximum, - ANNOTATION_IMPORT_MIN_RECORDS=minimum, - ), - ): - return AppAnnotationService.batch_import_app_annotations(app.id, _file(content), sqlite_session) - - def test_rejects_cross_tenant_app(self, sqlite_session: Session, current_user: Account) -> None: - app = _persist_app(sqlite_session, tenant_id=OTHER_TENANT_ID) - - with pytest.raises(NotFound): - AppAnnotationService.batch_import_app_annotations(app.id, _file(b"question,answer\nq,a\n"), sqlite_session) - - @pytest.mark.parametrize( - ("dataframe", "content", "maximum", "minimum", "expected"), - [ - (pd.DataFrame({"q": ["only"]}), b"question\nq\n", 5, 1, "Invalid CSV format"), - (pd.DataFrame({"q": ["q"], "a": ["a"]}), b"", 5, 1, "empty or invalid"), - (pd.DataFrame({"q": ["q"], "a": ["a"]}), b"question,answer\nq,a\n", 5, 2, "at least"), - ( - pd.DataFrame({"q": ["q1", "q2"], "a": ["a1", "a2"]}), - b"question,answer\nq1,a1\nq2,a2\n", - 1, - 1, - "too many records", - ), - (pd.DataFrame({"q": ["nan"], "a": ["nan"]}), b"question,answer\nnan,nan\n", 5, 1, "at least"), - ( - pd.DataFrame({"q": ["q" * 2001], "a": ["a"]}), - b"question,answer\nq,a\n", - 5, - 1, - "Question at row", - ), - ( - pd.DataFrame({"q": ["q"], "a": ["a" * 10001]}), - b"question,answer\nq,a\n", - 5, - 1, - "Answer at row", - ), - ], - ) - def test_validation_errors( - self, - sqlite_session: Session, - current_user: Account, - dataframe: pd.DataFrame, - content: bytes, - maximum: int, - minimum: int, - expected: str, - ) -> None: - app = _persist_app(sqlite_session) - - result = self._invoke( - sqlite_session, - app, - dataframe=dataframe, - content=content, - maximum=maximum, - minimum=minimum, - ) - - assert expected in cast(str, result["error_msg"]) - - def test_skips_malformed_rows(self, sqlite_session: Session, current_user: Account) -> None: - app = _persist_app(sqlite_session) - malformed_row = MagicMock() - malformed_row.iloc.__getitem__.side_effect = IndexError() - dataframe = MagicMock() - dataframe.columns = ["q", "a"] - dataframe.iterrows.return_value = [(0, malformed_row)] - - result = self._invoke(sqlite_session, app, dataframe=dataframe) - - assert "at least" in cast(str, result["error_msg"]) - - @config_overrides_context(DEPLOYMENT_EDITION=DeploymentEdition.CLOUD) - def test_rejects_subscription_quota_overflow(self, sqlite_session: Session, current_user: Account) -> None: - app = _persist_app(sqlite_session) - features = SimpleNamespace( - annotation_quota_limit=SimpleNamespace(limit=1, size=1), - ) - - result = self._invoke( - sqlite_session, - app, - dataframe=pd.DataFrame({"q": ["q1"], "a": ["a1"]}), - features=features, - ) - - assert "exceeds the limit" in cast(str, result["error_msg"]) - - @config_overrides_context(DEPLOYMENT_EDITION=DeploymentEdition.COMMUNITY) - def test_valid_import_enqueues_job(self, sqlite_session: Session, current_user: Account) -> None: - app = _persist_app(sqlite_session) - dataframe = pd.DataFrame({"q": ["q1"], "a": ["a1"]}) - features = SimpleNamespace(annotation_quota_limit=None) - with ( - patch.object(annotation_service_module.pd, "read_csv", return_value=dataframe), - patch.object(annotation_service_module.FeatureService, "get_features", return_value=features), - patch.object(annotation_service_module, "batch_import_annotations_task") as task, - patch.object(annotation_service_module, "redis_client") as redis, - patch.object(annotation_service_module.uuid, "uuid4", return_value="uuid-3"), - patch.object(annotation_service_module, "naive_utc_now", return_value=datetime.fromtimestamp(1)), - config_overrides_context(ANNOTATION_IMPORT_MAX_RECORDS=5, ANNOTATION_IMPORT_MIN_RECORDS=1), - ): - result = AppAnnotationService.batch_import_app_annotations( - app.id, _file(b"question,answer\nq,a\n"), sqlite_session - ) - - assert result == {"job_id": "uuid-3", "job_status": "waiting", "record_count": 1} - redis.zadd.assert_called_once_with(f"annotation_import_active:{TENANT_ID}", {"uuid-3": 1000}) - redis.expire.assert_called_once_with(f"annotation_import_active:{TENANT_ID}", 7200) - redis.setnx.assert_called_once_with("app_annotation_batch_import_uuid-3", "waiting") - task.delay.assert_called_once_with( - "uuid-3", [{"question": "q1", "answer": "a1"}], app.id, TENANT_ID, current_user.id - ) - - @pytest.mark.parametrize("limit", [0, 2], ids=["unlimited", "exactly-at-limit"]) - @config_overrides_context(DEPLOYMENT_EDITION=DeploymentEdition.CLOUD) - def test_import_with_available_quota_enqueues_job( - self, sqlite_session: Session, current_user: Account, limit: int - ) -> None: - app = _persist_app(sqlite_session) - features = SimpleNamespace(annotation_quota_limit=SimpleNamespace(limit=limit, size=1)) - with ( - patch.object(annotation_service_module.FeatureService, "get_features", return_value=features), - patch.object(annotation_service_module, "batch_import_annotations_task") as task, - patch.object(annotation_service_module, "redis_client"), - config_overrides_context(ANNOTATION_IMPORT_MAX_RECORDS=5, ANNOTATION_IMPORT_MIN_RECORDS=1), - ): - result = AppAnnotationService.batch_import_app_annotations( - app.id, _file(b"question,answer\nq,a\n"), sqlite_session - ) - - assert result["job_status"] == "waiting" - assert result["record_count"] == 1 - task.delay.assert_called_once_with( - result["job_id"], [{"question": "q", "answer": "a"}], app.id, TENANT_ID, current_user.id - ) - - @config_overrides_context(DEPLOYMENT_EDITION=DeploymentEdition.COMMUNITY) - def test_unexpected_error_cleans_active_job( - self, sqlite_session: Session, current_user: Account, caplog: pytest.LogCaptureFixture - ) -> None: - app = _persist_app(sqlite_session) - dataframe = pd.DataFrame({"q": ["q1"], "a": ["a1"]}) - features = SimpleNamespace(annotation_quota_limit=None) - with ( - patch.object(annotation_service_module.pd, "read_csv", return_value=dataframe), - patch.object(annotation_service_module.FeatureService, "get_features", return_value=features), - patch.object(annotation_service_module, "redis_client") as redis, - patch.object(annotation_service_module.uuid, "uuid4", return_value="uuid-4"), - patch.object(annotation_service_module, "naive_utc_now", return_value=datetime.fromtimestamp(1)), - config_overrides_context(ANNOTATION_IMPORT_MAX_RECORDS=5, ANNOTATION_IMPORT_MIN_RECORDS=1), - ): - redis.zadd.side_effect = RuntimeError("boom") - redis.zrem.side_effect = RuntimeError("cleanup-failed") - with caplog.at_level(logging.DEBUG): + with ( + patch.object(annotation_service_module.pd, "read_csv", return_value=dataframe), + patch.object(annotation_service_module.FeatureService, "get_features", return_value=features), + patch.object(annotation_service_module, "batch_import_annotations_task") as task, + patch.object(annotation_service_module, "redis_client") as redis, + patch.object(annotation_service_module.uuid, "uuid4", return_value="uuid-3"), + patch.object(annotation_service_module, "naive_utc_now", return_value=datetime.fromtimestamp(1)), + config_overrides_context(ANNOTATION_IMPORT_MAX_RECORDS=5, ANNOTATION_IMPORT_MIN_RECORDS=1), + ): result = AppAnnotationService.batch_import_app_annotations( app.id, _file(b"question,answer\nq,a\n"), sqlite_session ) - assert result["error_msg"] == "An error occurred while processing the file: boom" - redis.zrem.assert_called_once_with(f"annotation_import_active:{TENANT_ID}", "uuid-4") - assert "Failed to clean up active job tracking" in caplog.text - - -class TestAppAnnotationServiceHitHistoryAndSettings: - def test_hit_histories_are_annotation_and_app_scoped(self, sqlite_session: Session, current_user: Account) -> None: - app = _persist_app(sqlite_session) - other_app = _persist_app(sqlite_session, app_id="app-2") - other_annotation = _persist_annotation(sqlite_session, other_app) - - with pytest.raises(NotFound): - AppAnnotationService.get_annotation_hit_histories( - _annotation_ref(app, other_annotation.id), 1, 10, sqlite_session + assert result == {"job_id": "uuid-3", "job_status": "waiting", "record_count": 1} + redis.zadd.assert_called_once_with(f"annotation_import_active:{TENANT_ID}", {"uuid-3": 1000}) + redis.expire.assert_called_once_with(f"annotation_import_active:{TENANT_ID}", 7200) + redis.setnx.assert_called_once_with("app_annotation_batch_import_uuid-3", "waiting") + task.delay.assert_called_once_with( + "uuid-3", [{"question": "q1", "answer": "a1"}], app.id, TENANT_ID, current_user.id ) - annotation = _persist_annotation(sqlite_session, app, annotation_id="own-ann") - now = datetime(2026, 1, 1) - old = _persist_history(sqlite_session, app, annotation, history_id="old", created_at=now) - new = _persist_history(sqlite_session, app, annotation, history_id="new", created_at=now + timedelta(seconds=1)) - _persist_history(sqlite_session, other_app, other_annotation, history_id="decoy") + @pytest.mark.parametrize("limit", [0, 2], ids=["unlimited", "exactly-at-limit"]) + @config_overrides_context(DEPLOYMENT_EDITION=DeploymentEdition.CLOUD) + def test_import_with_available_quota_enqueues_job( + self, sqlite_session: Session, current_user: Account, limit: int + ) -> None: + app = _persist_app(sqlite_session) + features = SimpleNamespace(annotation_quota_limit=SimpleNamespace(limit=limit, size=1)) + with ( + patch.object(annotation_service_module.FeatureService, "get_features", return_value=features), + patch.object(annotation_service_module, "batch_import_annotations_task") as task, + patch.object(annotation_service_module, "redis_client"), + config_overrides_context(ANNOTATION_IMPORT_MAX_RECORDS=5, ANNOTATION_IMPORT_MIN_RECORDS=1), + ): + result = AppAnnotationService.batch_import_app_annotations( + app.id, _file(b"question,answer\nq,a\n"), sqlite_session + ) - items, total = AppAnnotationService.get_annotation_hit_histories( - _annotation_ref(app, annotation.id), 1, 1, sqlite_session - ) - assert total == 2 - assert [item.id for item in items] == [new.id] - items, total = AppAnnotationService.get_annotation_hit_histories( - _annotation_ref(app, annotation.id), 2, 1, sqlite_session - ) - assert total == 2 - assert [item.id for item in items] == [old.id] + assert result["job_status"] == "waiting" + assert result["record_count"] == 1 + task.delay.assert_called_once_with( + result["job_id"], [{"question": "q", "answer": "a"}], app.id, TENANT_ID, current_user.id + ) - def test_get_annotation_by_id_uses_real_identity_lookup( - self, sqlite_session: Session, current_user: Account - ) -> None: - app = _persist_app(sqlite_session) - annotation = _persist_annotation(sqlite_session, app) + @config_overrides_context(DEPLOYMENT_EDITION=DeploymentEdition.COMMUNITY) + def test_unexpected_error_cleans_active_job( + self, sqlite_session: Session, current_user: Account, caplog: pytest.LogCaptureFixture + ) -> None: + app = _persist_app(sqlite_session) + dataframe = pd.DataFrame({"q": ["q1"], "a": ["a1"]}) + features = SimpleNamespace(annotation_quota_limit=None) + with ( + patch.object(annotation_service_module.pd, "read_csv", return_value=dataframe), + patch.object(annotation_service_module.FeatureService, "get_features", return_value=features), + patch.object(annotation_service_module, "redis_client") as redis, + patch.object(annotation_service_module.uuid, "uuid4", return_value="uuid-4"), + patch.object(annotation_service_module, "naive_utc_now", return_value=datetime.fromtimestamp(1)), + config_overrides_context(ANNOTATION_IMPORT_MAX_RECORDS=5, ANNOTATION_IMPORT_MIN_RECORDS=1), + ): + redis.zadd.side_effect = RuntimeError("boom") + redis.zrem.side_effect = RuntimeError("cleanup-failed") + with caplog.at_level(logging.DEBUG): + result = AppAnnotationService.batch_import_app_annotations( + app.id, _file(b"question,answer\nq,a\n"), sqlite_session + ) - assert AppAnnotationService.get_annotation_by_id("missing", sqlite_session) is None - assert AppAnnotationService.get_annotation_by_id(annotation.id, sqlite_session) is annotation + assert result["error_msg"] == "An error occurred while processing the file: boom" + redis.zrem.assert_called_once_with(f"annotation_import_active:{TENANT_ID}", "uuid-4") + assert "Failed to clean up active job tracking" in caplog.text - def test_add_history_increments_count_and_flushes_row(self, sqlite_session: Session, current_user: Account) -> None: - app = _persist_app(sqlite_session) - annotation = _persist_annotation(sqlite_session, app) + class TestAppAnnotationServiceHitHistoryAndSettings: + def test_hit_histories_are_annotation_and_app_scoped( + self, sqlite_session: Session, current_user: Account + ) -> None: + app = _persist_app(sqlite_session) + other_app = _persist_app(sqlite_session, app_id="app-2") + other_annotation = _persist_annotation(sqlite_session, other_app) - AppAnnotationService.add_annotation_history( - annotation_id=annotation.id, - app_id=app.id, - annotation_question="q", - annotation_content="a", - query="user q", - user_id=current_user.id, - message_id="msg-1", - from_source="chat", - score=0.8, - session=sqlite_session, - ) + with pytest.raises(NotFound): + AppAnnotationService.get_annotation_hit_histories( + _annotation_ref(app, other_annotation.id), 1, 10, sqlite_session + ) - sqlite_session.refresh(annotation) - assert annotation.hit_count == 1 - history = sqlite_session.scalar( - select(AppAnnotationHitHistory).where(AppAnnotationHitHistory.annotation_id == annotation.id) - ) - assert history is not None - assert (history.question, history.annotation_question, history.annotation_content, history.score) == ( - "user q", - "q", - "a", - 0.8, - ) + annotation = _persist_annotation(sqlite_session, app, annotation_id="own-ann") + now = datetime(2026, 1, 1) + old = _persist_history(sqlite_session, app, annotation, history_id="old", created_at=now) + new = _persist_history( + sqlite_session, app, annotation, history_id="new", created_at=now + timedelta(seconds=1) + ) + _persist_history(sqlite_session, other_app, other_annotation, history_id="decoy") - def test_get_setting_rejects_cross_tenant_app(self, sqlite_session: Session, current_user: Account) -> None: - app = _persist_app(sqlite_session, tenant_id=OTHER_TENANT_ID) + items, total = AppAnnotationService.get_annotation_hit_histories( + _annotation_ref(app, annotation.id), 1, 1, sqlite_session + ) + assert total == 2 + assert [item.id for item in items] == [new.id] + items, total = AppAnnotationService.get_annotation_hit_histories( + _annotation_ref(app, annotation.id), 2, 1, sqlite_session + ) + assert total == 2 + assert [item.id for item in items] == [old.id] - with pytest.raises(NotFound): - AppAnnotationService.get_app_annotation_setting_by_app_id(app.id, sqlite_session) + def test_get_annotation_by_id_uses_real_identity_lookup( + self, sqlite_session: Session, current_user: Account + ) -> None: + app = _persist_app(sqlite_session) + annotation = _persist_annotation(sqlite_session, app) - def test_get_setting_returns_disabled_without_row(self, sqlite_session: Session, current_user: Account) -> None: - app = _persist_app(sqlite_session) + assert AppAnnotationService.get_annotation_by_id("missing", sqlite_session) is None + assert AppAnnotationService.get_annotation_by_id(annotation.id, sqlite_session) is annotation - assert AppAnnotationService.get_app_annotation_setting_by_app_id(app.id, sqlite_session) == {"enabled": False} + def test_add_history_increments_count_and_flushes_row( + self, sqlite_session: Session, current_user: Account + ) -> None: + app = _persist_app(sqlite_session) + annotation = _persist_annotation(sqlite_session, app) - def test_get_setting_returns_binding_detail(self, sqlite_session: Session, current_user: Account) -> None: - app = _persist_app(sqlite_session) - binding = _persist_binding(sqlite_session) - setting = _persist_setting(sqlite_session, app, binding_id=binding.id) + AppAnnotationService.add_annotation_history( + annotation_id=annotation.id, + app_id=app.id, + annotation_question="q", + annotation_content="a", + query="user q", + user_id=current_user.id, + message_id="msg-1", + from_source="chat", + score=0.8, + session=sqlite_session, + ) - result = AppAnnotationService.get_app_annotation_setting_by_app_id(app.id, sqlite_session) + sqlite_session.refresh(annotation) + assert annotation.hit_count == 1 + history = sqlite_session.scalar( + select(AppAnnotationHitHistory).where(AppAnnotationHitHistory.annotation_id == annotation.id) + ) + assert history is not None + assert (history.question, history.annotation_question, history.annotation_content, history.score) == ( + "user q", + "q", + "a", + 0.8, + ) - assert result == { - "id": setting.id, - "enabled": True, - "score_threshold": 0.5, - "embedding_model": { + def test_get_setting_rejects_cross_tenant_app(self, sqlite_session: Session, current_user: Account) -> None: + app = _persist_app(sqlite_session, tenant_id=OTHER_TENANT_ID) + + with pytest.raises(NotFound): + AppAnnotationService.get_app_annotation_setting_by_app_id(app.id, sqlite_session) + + def test_get_setting_returns_disabled_without_row(self, sqlite_session: Session, current_user: Account) -> None: + app = _persist_app(sqlite_session) + + assert AppAnnotationService.get_app_annotation_setting_by_app_id(app.id, sqlite_session) == { + "enabled": False + } + + def test_get_setting_returns_binding_detail(self, sqlite_session: Session, current_user: Account) -> None: + app = _persist_app(sqlite_session) + binding = _persist_binding(sqlite_session) + setting = _persist_setting(sqlite_session, app, binding_id=binding.id) + + result = AppAnnotationService.get_app_annotation_setting_by_app_id(app.id, sqlite_session) + + assert result == { + "id": setting.id, + "enabled": True, + "score_threshold": 0.5, + "embedding_model": { + "embedding_provider_name": binding.provider_name, + "embedding_model_name": binding.model_name, + }, + } + + def test_get_setting_returns_empty_detail_for_missing_binding( + self, sqlite_session: Session, current_user: Account + ) -> None: + app = _persist_app(sqlite_session) + setting = _persist_setting(sqlite_session, app, binding_id="missing-binding") + + result = AppAnnotationService.get_app_annotation_setting_by_app_id(app.id, sqlite_session) + + assert result == { + "id": setting.id, + "enabled": True, + "score_threshold": 0.5, + "embedding_model": {}, + } + + def test_update_setting_is_app_scoped(self, sqlite_session: Session, current_user: Account) -> None: + app = _persist_app(sqlite_session) + other_app = _persist_app(sqlite_session, app_id="app-2") + other_setting = _persist_setting(sqlite_session, other_app) + + with pytest.raises(NotFound): + AppAnnotationService.update_app_annotation_setting( + app.id, other_setting.id, {"score_threshold": 0.8}, sqlite_session + ) + + def test_update_setting_flushes_changes_and_returns_binding( + self, sqlite_session: Session, current_user: Account + ) -> None: + app = _persist_app(sqlite_session) + binding = _persist_binding(sqlite_session) + setting = _persist_setting(sqlite_session, app, binding_id=binding.id) + + result = AppAnnotationService.update_app_annotation_setting( + app.id, setting.id, {"score_threshold": 0.8}, sqlite_session + ) + + assert result["enabled"] is True + assert result["score_threshold"] == 0.8 + assert result["embedding_model"] == { "embedding_provider_name": binding.provider_name, "embedding_model_name": binding.model_name, - }, - } + } + assert setting.score_threshold == 0.8 + assert setting.updated_user_id == current_user.id - def test_get_setting_returns_empty_detail_for_missing_binding( - self, sqlite_session: Session, current_user: Account - ) -> None: - app = _persist_app(sqlite_session) - setting = _persist_setting(sqlite_session, app, binding_id="missing-binding") + def test_update_setting_returns_empty_detail_for_missing_binding( + self, sqlite_session: Session, current_user: Account + ) -> None: + app = _persist_app(sqlite_session) + setting = _persist_setting(sqlite_session, app, binding_id="missing-binding") - result = AppAnnotationService.get_app_annotation_setting_by_app_id(app.id, sqlite_session) - - assert result == { - "id": setting.id, - "enabled": True, - "score_threshold": 0.5, - "embedding_model": {}, - } - - def test_update_setting_is_app_scoped(self, sqlite_session: Session, current_user: Account) -> None: - app = _persist_app(sqlite_session) - other_app = _persist_app(sqlite_session, app_id="app-2") - other_setting = _persist_setting(sqlite_session, other_app) - - with pytest.raises(NotFound): - AppAnnotationService.update_app_annotation_setting( - app.id, other_setting.id, {"score_threshold": 0.8}, sqlite_session + result = AppAnnotationService.update_app_annotation_setting( + app.id, setting.id, {"score_threshold": 0.7}, sqlite_session ) - def test_update_setting_flushes_changes_and_returns_binding( - self, sqlite_session: Session, current_user: Account - ) -> None: - app = _persist_app(sqlite_session) - binding = _persist_binding(sqlite_session) - setting = _persist_setting(sqlite_session, app, binding_id=binding.id) + assert result["score_threshold"] == 0.7 + assert result["embedding_model"] == {} - result = AppAnnotationService.update_app_annotation_setting( - app.id, setting.id, {"score_threshold": 0.8}, sqlite_session - ) + class TestAppAnnotationServiceClearAll: + def test_clear_all_deletes_only_app_rows_and_enqueues_indexes( + self, + sqlite_session: Session, + sqlite_session_factory: sessionmaker[Session], + current_user: Account, + ) -> None: + app = _persist_app(sqlite_session) + other_app = _persist_app(sqlite_session, app_id="app-2") + annotations = [ + _persist_annotation(sqlite_session, app, annotation_id="ann-1"), + _persist_annotation(sqlite_session, app, annotation_id="ann-2"), + ] + histories = [ + _persist_history(sqlite_session, app, annotations[0], history_id="history-1"), + _persist_history(sqlite_session, app, annotations[1], history_id="history-2"), + ] + decoy = _persist_annotation(sqlite_session, other_app, annotation_id="decoy") + decoy_history = _persist_history(sqlite_session, other_app, decoy, history_id="decoy-history") + setting = _persist_setting(sqlite_session, app) - assert result["enabled"] is True - assert result["score_threshold"] == 0.8 - assert result["embedding_model"] == { - "embedding_provider_name": binding.provider_name, - "embedding_model_name": binding.model_name, - } - assert setting.score_threshold == 0.8 - assert setting.updated_user_id == current_user.id + with patch.object(annotation_service_module, "delete_annotation_index_task") as task: + result = AppAnnotationService.clear_all_annotations(app.id, sqlite_session) - def test_update_setting_returns_empty_detail_for_missing_binding( - self, sqlite_session: Session, current_user: Account - ) -> None: - app = _persist_app(sqlite_session) - setting = _persist_setting(sqlite_session, app, binding_id="missing-binding") + assert result == {"result": "success"} + with sqlite_session_factory() as observer: + assert [observer.get(MessageAnnotation, item.id) for item in annotations] == [None, None] + assert [observer.get(AppAnnotationHitHistory, item.id) for item in histories] == [None, None] + assert observer.get(MessageAnnotation, decoy.id) is not None + assert observer.get(AppAnnotationHitHistory, decoy_history.id) is not None + assert task.delay.call_count == 2 + for annotation in annotations: + task.delay.assert_any_call(annotation.id, app.id, TENANT_ID, setting.collection_binding_id) - result = AppAnnotationService.update_app_annotation_setting( - app.id, setting.id, {"score_threshold": 0.7}, sqlite_session - ) + def test_clear_all_rejects_cross_tenant_app(self, sqlite_session: Session, current_user: Account) -> None: + app = _persist_app(sqlite_session, tenant_id=OTHER_TENANT_ID) - assert result["score_threshold"] == 0.7 - assert result["embedding_model"] == {} + with pytest.raises(NotFound): + AppAnnotationService.clear_all_annotations(app.id, sqlite_session) - -class TestAppAnnotationServiceClearAll: - def test_clear_all_deletes_only_app_rows_and_enqueues_indexes( - self, - sqlite_session: Session, - sqlite_session_factory: sessionmaker[Session], - current_user: Account, - ) -> None: - app = _persist_app(sqlite_session) - other_app = _persist_app(sqlite_session, app_id="app-2") - annotations = [ - _persist_annotation(sqlite_session, app, annotation_id="ann-1"), - _persist_annotation(sqlite_session, app, annotation_id="ann-2"), - ] - histories = [ - _persist_history(sqlite_session, app, annotations[0], history_id="history-1"), - _persist_history(sqlite_session, app, annotations[1], history_id="history-2"), - ] - decoy = _persist_annotation(sqlite_session, other_app, annotation_id="decoy") - decoy_history = _persist_history(sqlite_session, other_app, decoy, history_id="decoy-history") - setting = _persist_setting(sqlite_session, app) - - with patch.object(annotation_service_module, "delete_annotation_index_task") as task: - result = AppAnnotationService.clear_all_annotations(app.id, sqlite_session) - - assert result == {"result": "success"} - with sqlite_session_factory() as observer: - assert [observer.get(MessageAnnotation, item.id) for item in annotations] == [None, None] - assert [observer.get(AppAnnotationHitHistory, item.id) for item in histories] == [None, None] - assert observer.get(MessageAnnotation, decoy.id) is not None - assert observer.get(AppAnnotationHitHistory, decoy_history.id) is not None - assert task.delay.call_count == 2 - for annotation in annotations: - task.delay.assert_any_call(annotation.id, app.id, TENANT_ID, setting.collection_binding_id) - - def test_clear_all_rejects_cross_tenant_app(self, sqlite_session: Session, current_user: Account) -> None: - app = _persist_app(sqlite_session, tenant_id=OTHER_TENANT_ID) - - with pytest.raises(NotFound): - AppAnnotationService.clear_all_annotations(app.id, sqlite_session) - - assert sqlite_session.scalar(select(func.count()).select_from(MessageAnnotation)) == 0 + assert sqlite_session.scalar(select(func.count()).select_from(MessageAnnotation)) == 0 diff --git a/api/tests/unit_tests/services/test_api_token_service.py b/api/tests/unit_tests/services/test_api_token_service.py index 6a202468d58..3aa0b6e9338 100644 --- a/api/tests/unit_tests/services/test_api_token_service.py +++ b/api/tests/unit_tests/services/test_api_token_service.py @@ -1,10 +1,9 @@ -"""Unit tests for API token caching and SQLite-backed token lookup.""" - from __future__ import annotations +import json from collections.abc import Iterator from datetime import datetime -from unittest.mock import patch +from unittest.mock import MagicMock, patch from uuid import uuid4 import pytest @@ -18,7 +17,255 @@ from extensions.ext_redis import RedisClientWrapper from models.engine import db from models.enums import ApiTokenType from models.model import ApiToken -from services.api_token_service import ApiTokenCache, CachedApiToken +from services.api_token_service import ( + CACHE_KEY_PREFIX, + CACHE_NULL_TTL_SECONDS, + CACHE_TTL_SECONDS, + ApiTokenCache, + CachedApiToken, +) + +""" +Unit tests for API Token Cache module. +""" + + +class TestApiTokenCacheCases: + class TestApiTokenCache: + """Test cases for ApiTokenCache class.""" + + def setup_method(self) -> None: + """Setup test fixtures.""" + self.mock_token = ApiToken() + self.mock_token.id = "test-token-id-123" + self.mock_token.app_id = "test-app-id-456" + self.mock_token.tenant_id = "test-tenant-id-789" + self.mock_token.type = ApiTokenType.APP + self.mock_token.token = "test-token-value-abc" + self.mock_token.last_used_at = datetime(2026, 2, 3, 10, 0, 0) + self.mock_token.created_at = datetime(2026, 1, 1, 0, 0, 0) + + def test_make_cache_key(self) -> None: + """Test cache key generation.""" + # Test with scope + key = ApiTokenCache._make_cache_key("my-token", "app") + assert key == f"{CACHE_KEY_PREFIX}:app:my-token" + + # Test without scope + key = ApiTokenCache._make_cache_key("my-token", None) + assert key == f"{CACHE_KEY_PREFIX}:any:my-token" + + def test_serialize_token(self) -> None: + """Test token serialization.""" + serialized = ApiTokenCache._serialize_token(self.mock_token) + data = json.loads(serialized) + + assert data["id"] == "test-token-id-123" + assert data["app_id"] == "test-app-id-456" + assert data["tenant_id"] == "test-tenant-id-789" + assert data["type"] == "app" + assert data["token"] == "test-token-value-abc" + assert data["last_used_at"] == "2026-02-03T10:00:00" + assert data["created_at"] == "2026-01-01T00:00:00" + + def test_serialize_token_with_nulls(self) -> None: + """Test token serialization with None values.""" + mock_token = ApiToken() + mock_token.id = "test-id" + mock_token.app_id = None + mock_token.tenant_id = None + mock_token.type = ApiTokenType.DATASET + mock_token.token = "test-token" + mock_token.last_used_at = None + mock_token.created_at = datetime(2026, 1, 1, 0, 0, 0) + + serialized = ApiTokenCache._serialize_token(mock_token) + data = json.loads(serialized) + + assert data["app_id"] is None + assert data["tenant_id"] is None + assert data["last_used_at"] is None + + def test_deserialize_token(self) -> None: + """Test token deserialization.""" + cached_data = json.dumps( + { + "id": "test-id", + "app_id": "test-app", + "tenant_id": "test-tenant", + "type": "app", + "token": "test-token", + "last_used_at": "2026-02-03T10:00:00", + "created_at": "2026-01-01T00:00:00", + } + ) + + result = ApiTokenCache._deserialize_token(cached_data) + + assert isinstance(result, CachedApiToken) + assert result.id == "test-id" + assert result.app_id == "test-app" + assert result.tenant_id == "test-tenant" + assert result.type == "app" + assert result.token == "test-token" + assert result.last_used_at == datetime(2026, 2, 3, 10, 0, 0) + assert result.created_at == datetime(2026, 1, 1, 0, 0, 0) + + def test_deserialize_null_token(self) -> None: + """Test deserialization of null token (cached miss).""" + result = ApiTokenCache._deserialize_token("null") + assert result is None + + def test_deserialize_invalid_json(self) -> None: + """Test deserialization with invalid JSON.""" + result = ApiTokenCache._deserialize_token("invalid-json{") + assert result is None + + @patch("services.api_token_service.redis_client") + def test_get_cache_hit(self, mock_redis: MagicMock) -> None: + """Test cache hit scenario.""" + cached_data = json.dumps( + { + "id": "test-id", + "app_id": "test-app", + "tenant_id": "test-tenant", + "type": "app", + "token": "test-token", + "last_used_at": "2026-02-03T10:00:00", + "created_at": "2026-01-01T00:00:00", + } + ).encode("utf-8") + mock_redis.get.return_value = cached_data + + result = ApiTokenCache.get("test-token", "app") + + assert result is not None + assert isinstance(result, CachedApiToken) + assert result.app_id == "test-app" + mock_redis.get.assert_called_once_with(f"{CACHE_KEY_PREFIX}:app:test-token") + + @patch("services.api_token_service.redis_client") + def test_get_cache_miss(self, mock_redis: MagicMock) -> None: + """Test cache miss scenario.""" + mock_redis.get.return_value = None + + result = ApiTokenCache.get("test-token", "app") + + assert result is None + mock_redis.get.assert_called_once() + + @patch("services.api_token_service.redis_client") + def test_set_valid_token(self, mock_redis: MagicMock) -> None: + """Test setting a valid token in cache.""" + result = ApiTokenCache.set("test-token", "app", self.mock_token) + + assert result is True + mock_redis.setex.assert_called_once() + args = mock_redis.setex.call_args[0] + assert args[0] == f"{CACHE_KEY_PREFIX}:app:test-token" + assert args[1] == CACHE_TTL_SECONDS + + @patch("services.api_token_service.redis_client") + def test_set_null_token(self, mock_redis: MagicMock) -> None: + """Test setting a null token (cache penetration prevention).""" + result = ApiTokenCache.set("invalid-token", "app", None) + + assert result is True + mock_redis.setex.assert_called_once() + args = mock_redis.setex.call_args[0] + assert args[0] == f"{CACHE_KEY_PREFIX}:app:invalid-token" + assert args[1] == CACHE_NULL_TTL_SECONDS + assert args[2] == b"null" + + @patch("services.api_token_service.redis_client") + def test_delete_with_scope(self, mock_redis: MagicMock) -> None: + """Test deleting token cache with specific scope.""" + result = ApiTokenCache.delete("test-token", "app") + + assert result is True + mock_redis.delete.assert_called_once_with(f"{CACHE_KEY_PREFIX}:app:test-token") + + @patch("services.api_token_service.redis_client") + def test_delete_without_scope(self, mock_redis: MagicMock) -> None: + """Test deleting token cache without scope (delete all).""" + # Mock scan_iter to return an iterator of keys + mock_redis.scan_iter.return_value = iter( + [ + b"api_token:app:test-token", + b"api_token:dataset:test-token", + ] + ) + + result = ApiTokenCache.delete("test-token", None) + + assert result is True + # Verify scan_iter was called with the correct pattern + mock_redis.scan_iter.assert_called_once() + call_args = mock_redis.scan_iter.call_args + assert call_args[1]["match"] == f"{CACHE_KEY_PREFIX}:*:test-token" + + # Verify delete was called with all matched keys + mock_redis.delete.assert_called_once_with( + b"api_token:app:test-token", + b"api_token:dataset:test-token", + ) + + @patch("services.api_token_service.redis_client") + def test_redis_fallback_on_exception(self, mock_redis: MagicMock) -> None: + """Test Redis fallback when Redis is unavailable.""" + from redis import RedisError + + mock_redis.get.side_effect = RedisError("Connection failed") + + result = ApiTokenCache.get("test-token", "app") + + # Should return None (fallback) instead of raising exception + assert result is None + + class TestApiTokenCacheIntegration: + """Integration test scenarios.""" + + @patch("services.api_token_service.redis_client") + def test_full_cache_lifecycle(self, mock_redis: MagicMock) -> None: + """Test complete cache lifecycle: set -> get -> delete.""" + # Use the persisted model's fields for the cache lifecycle. + mock_token = ApiToken() + mock_token.id = "id-123" + mock_token.app_id = "app-456" + mock_token.tenant_id = "tenant-789" + mock_token.type = ApiTokenType.APP + mock_token.token = "token-abc" + mock_token.last_used_at = datetime(2026, 2, 3, 10, 0, 0) + mock_token.created_at = datetime(2026, 1, 1, 0, 0, 0) + + # 1. Set token in cache + ApiTokenCache.set("token-abc", "app", mock_token) + assert mock_redis.setex.called + + # 2. Simulate cache hit + cached_data = ApiTokenCache._serialize_token(mock_token) + mock_redis.get.return_value = cached_data # bytes from model_dump_json().encode() + + retrieved = ApiTokenCache.get("token-abc", "app") + assert retrieved is not None + assert isinstance(retrieved, CachedApiToken) + + # 3. Delete from cache + ApiTokenCache.delete("token-abc", "app") + assert mock_redis.delete.called + + @patch("services.api_token_service.redis_client") + def test_cache_penetration_prevention(self, mock_redis: MagicMock) -> None: + """Test that non-existent tokens are cached as null.""" + # Set null token (cache miss) + ApiTokenCache.set("non-existent-token", "app", None) + + args = mock_redis.setex.call_args[0] + assert args[2] == b"null" + assert args[1] == CACHE_NULL_TTL_SECONDS # Shorter TTL for null values + + +"""Unit tests for API token caching and SQLite-backed token lookup.""" def _api_token( @@ -39,142 +286,6 @@ def _api_token( ) -@pytest.fixture -def api_token_db() -> Iterator[Session]: - """Provide the production database extension with an isolated SQLite token table.""" - app = Flask(__name__) - app.config["SQLALCHEMY_DATABASE_URI"] = "sqlite:///:memory:" - db.init_app(app) - - with app.app_context(): - ApiToken.__table__.create(db.engine) - with Session(db.engine, expire_on_commit=False) as session: - yield session - - -class TestQueryTokenFromDb: - def test_should_return_api_token_and_cache_when_token_exists(self, api_token_db: Session) -> None: - tenant_id = str(uuid4()) - app_id = str(uuid4()) - token_value = f"app-test-{uuid4()}" - - api_token = ApiToken() - api_token.id = str(uuid4()) - api_token.app_id = app_id - api_token.tenant_id = tenant_id - api_token.type = "app" - api_token.token = token_value - api_token_db.add(api_token) - api_token_db.commit() - - with ( - patch.object(api_token_service_module.ApiTokenCache, "set") as mock_cache_set, - patch.object(api_token_service_module, "record_token_usage") as mock_record_usage, - ): - result = api_token_service_module.query_token_from_db(token_value, "app") - - assert result.id == api_token.id - assert result.token == token_value - mock_cache_set.assert_called_once() - mock_record_usage.assert_called_once_with(token_value, "app") - - def test_should_cache_null_and_raise_unauthorized_when_token_not_found(self, api_token_db: Session) -> None: - with ( - patch.object(api_token_service_module.ApiTokenCache, "set") as mock_cache_set, - patch.object(api_token_service_module, "record_token_usage") as mock_record_usage, - ): - with pytest.raises(Unauthorized, match="Access token is invalid"): - api_token_service_module.query_token_from_db(f"missing-{uuid4()}", "app") - - mock_cache_set.assert_called_once() - call_args = mock_cache_set.call_args[0] - assert call_args[2] is None # cached None - mock_record_usage.assert_not_called() - - -class TestRecordTokenUsage: - def test_should_write_active_key_with_iso_timestamp_and_ttl(self): - auth_token = "token-123" - scope = "dataset" - fixed_time = datetime(2026, 2, 24, 12, 0, 0) - expected_key = ApiTokenCache.make_active_key(auth_token, scope) - - with ( - patch.object(api_token_service_module, "naive_utc_now", return_value=fixed_time), - patch.object(api_token_service_module, "redis_client") as mock_redis, - ): - api_token_service_module.record_token_usage(auth_token, scope) - - mock_redis.set.assert_called_once_with(expected_key, fixed_time.isoformat(), ex=3600) - - def test_should_not_raise_when_redis_write_fails(self): - with patch.object(api_token_service_module, "redis_client") as mock_redis: - mock_redis.set.side_effect = Exception("redis unavailable") - api_token_service_module.record_token_usage("token-123", "app") - - -class TestGetEffectiveTokenLastUsedAt: - def test_should_surface_pending_usage_before_the_batch_flush(self): - pending = datetime(2026, 8, 11, 12, 30, 0) - - with patch.object(api_token_service_module, "redis_client") as mock_redis: - mock_redis.get.return_value = pending.isoformat().encode() - result = api_token_service_module.get_effective_token_last_used_at( - "token-123", - "dataset", - None, - ) - - assert result == pending - mock_redis.get.assert_called_once_with(ApiTokenCache.make_active_key("token-123", "dataset")) - - def test_should_keep_a_newer_persisted_usage_timestamp(self): - pending = datetime(2026, 8, 11, 12, 30, 0) - persisted = datetime(2026, 8, 11, 12, 31, 0) - - with patch.object(api_token_service_module, "redis_client") as mock_redis: - mock_redis.get.return_value = pending.isoformat() - result = api_token_service_module.get_effective_token_last_used_at( - "token-123", - "dataset", - persisted, - ) - - assert result == persisted - - @pytest.mark.parametrize("value", [None, b"not-a-timestamp"]) - def test_should_fall_back_to_the_persisted_timestamp(self, value): - persisted = datetime(2026, 8, 11, 12, 30, 0) - - with patch.object(api_token_service_module, "redis_client") as mock_redis: - mock_redis.get.return_value = value - result = api_token_service_module.get_effective_token_last_used_at( - "token-123", - "dataset", - persisted, - ) - - assert result == persisted - - -@pytest.fixture -def single_flight_redis(monkeypatch: pytest.MonkeyPatch) -> Iterator[tuple[Redis, list[tuple[object, ...]]]]: - """Use the real Redis client, wrapper and lock; isolate only wire commands.""" - commands: list[tuple[object, ...]] = [] - - def execute(*args: object, **_kwargs: object) -> object: - commands.append(args) - assert args[0] in {"SET", "EVALSHA"} - return True - - with Redis() as client: - monkeypatch.setattr(client, "execute_command", execute) - wrapper = RedisClientWrapper() - wrapper.initialize(client) - monkeypatch.setattr(api_token_service_module, "redis_client", wrapper) - yield client, commands - - def _assert_lock_released(commands: list[tuple[object, ...]]) -> None: assert len(commands) == 2 acquire, release = commands @@ -184,282 +295,414 @@ def _assert_lock_released(commands: list[tuple[object, ...]]) -> None: assert release[2:] == (1, acquire[1], acquire[2]) # Release uses the acquired ownership token. -class TestFetchTokenWithSingleFlight: - def test_should_return_cached_token_when_lock_acquired_and_cache_filled( - self, single_flight_redis: tuple[Redis, list[tuple[object, ...]]] - ) -> None: - auth_token = "token-123" - scope = "app" - cached_token = CachedApiToken( - id="id-1", - app_id="app-1", - tenant_id="tenant-1", - type="app", - token=auth_token, - last_used_at=None, - created_at=None, - ) - _client, commands = single_flight_redis +class TestApiTokenServiceCases: + @staticmethod + @pytest.fixture + def api_token_db() -> Iterator[Session]: + """Provide the production database extension with an isolated SQLite token table.""" + app = Flask(__name__) + app.config["SQLALCHEMY_DATABASE_URI"] = "sqlite:///:memory:" + db.init_app(app) - with ( - patch.object(api_token_service_module.ApiTokenCache, "get", return_value=cached_token), - patch.object(api_token_service_module, "query_token_from_db") as mock_query_db, - ): - result = api_token_service_module.fetch_token_with_single_flight(auth_token, scope) + with app.app_context(): + ApiToken.__table__.create(db.engine) + with Session(db.engine, expire_on_commit=False) as session: + yield session - assert result == cached_token - _assert_lock_released(commands) - mock_query_db.assert_not_called() + class TestQueryTokenFromDb: + def test_should_return_api_token_and_cache_when_token_exists(self, api_token_db: Session) -> None: + tenant_id = str(uuid4()) + app_id = str(uuid4()) + token_value = f"app-test-{uuid4()}" - def test_should_query_db_when_lock_acquired_and_cache_missed( - self, single_flight_redis: tuple[Redis, list[tuple[object, ...]]] - ) -> None: - auth_token = "token-123" - scope = "app" - db_token = _api_token() - _client, commands = single_flight_redis + api_token = ApiToken() + api_token.id = str(uuid4()) + api_token.app_id = app_id + api_token.tenant_id = tenant_id + api_token.type = "app" + api_token.token = token_value + api_token_db.add(api_token) + api_token_db.commit() - with ( - patch.object(api_token_service_module.ApiTokenCache, "get", return_value=None), - patch.object(api_token_service_module, "query_token_from_db", return_value=db_token) as mock_query_db, - ): - result = api_token_service_module.fetch_token_with_single_flight(auth_token, scope) + with ( + patch.object(api_token_service_module.ApiTokenCache, "set") as mock_cache_set, + patch.object(api_token_service_module, "record_token_usage") as mock_record_usage, + ): + result = api_token_service_module.query_token_from_db(token_value, "app") - assert result == db_token - mock_query_db.assert_called_once_with(auth_token, scope) - _assert_lock_released(commands) + assert result.id == api_token.id + assert result.token == token_value + mock_cache_set.assert_called_once() + mock_record_usage.assert_called_once_with(token_value, "app") - def test_should_query_db_directly_when_lock_not_acquired( - self, single_flight_redis: tuple[Redis, list[tuple[object, ...]]], monkeypatch: pytest.MonkeyPatch - ) -> None: - auth_token = "token-123" - scope = "app" - db_token = _api_token() - client, commands = single_flight_redis + def test_should_cache_null_and_raise_unauthorized_when_token_not_found(self, api_token_db: Session) -> None: + with ( + patch.object(api_token_service_module.ApiTokenCache, "set") as mock_cache_set, + patch.object(api_token_service_module, "record_token_usage") as mock_record_usage, + ): + with pytest.raises(Unauthorized, match="Access token is invalid"): + api_token_service_module.query_token_from_db(f"missing-{uuid4()}", "app") - def contended(*args: object, **_kwargs: object) -> bool: + mock_cache_set.assert_called_once() + call_args = mock_cache_set.call_args[0] + assert call_args[2] is None # cached None + mock_record_usage.assert_not_called() + + class TestRecordTokenUsage: + def test_should_write_active_key_with_iso_timestamp_and_ttl(self): + auth_token = "token-123" + scope = "dataset" + fixed_time = datetime(2026, 2, 24, 12, 0, 0) + expected_key = ApiTokenCache.make_active_key(auth_token, scope) + + with ( + patch.object(api_token_service_module, "naive_utc_now", return_value=fixed_time), + patch.object(api_token_service_module, "redis_client") as mock_redis, + ): + api_token_service_module.record_token_usage(auth_token, scope) + + mock_redis.set.assert_called_once_with(expected_key, fixed_time.isoformat(), ex=3600) + + def test_should_not_raise_when_redis_write_fails(self): + with patch.object(api_token_service_module, "redis_client") as mock_redis: + mock_redis.set.side_effect = Exception("redis unavailable") + api_token_service_module.record_token_usage("token-123", "app") + + class TestGetEffectiveTokenLastUsedAt: + def test_should_surface_pending_usage_before_the_batch_flush(self): + pending = datetime(2026, 8, 11, 12, 30, 0) + + with patch.object(api_token_service_module, "redis_client") as mock_redis: + mock_redis.get.return_value = pending.isoformat().encode() + result = api_token_service_module.get_effective_token_last_used_at( + "token-123", + "dataset", + None, + ) + + assert result == pending + mock_redis.get.assert_called_once_with(ApiTokenCache.make_active_key("token-123", "dataset")) + + def test_should_keep_a_newer_persisted_usage_timestamp(self): + pending = datetime(2026, 8, 11, 12, 30, 0) + persisted = datetime(2026, 8, 11, 12, 31, 0) + + with patch.object(api_token_service_module, "redis_client") as mock_redis: + mock_redis.get.return_value = pending.isoformat() + result = api_token_service_module.get_effective_token_last_used_at( + "token-123", + "dataset", + persisted, + ) + + assert result == persisted + + @pytest.mark.parametrize("value", [None, b"not-a-timestamp"]) + def test_should_fall_back_to_the_persisted_timestamp(self, value): + persisted = datetime(2026, 8, 11, 12, 30, 0) + + with patch.object(api_token_service_module, "redis_client") as mock_redis: + mock_redis.get.return_value = value + result = api_token_service_module.get_effective_token_last_used_at( + "token-123", + "dataset", + persisted, + ) + + assert result == persisted + + @staticmethod + @pytest.fixture + def single_flight_redis(monkeypatch: pytest.MonkeyPatch) -> Iterator[tuple[Redis, list[tuple[object, ...]]]]: + """Use the real Redis client, wrapper and lock; isolate only wire commands.""" + commands: list[tuple[object, ...]] = [] + + def execute(*args: object, **_kwargs: object) -> object: commands.append(args) - assert args[0] == "SET" - return False + assert args[0] in {"SET", "EVALSHA"} + return True - monkeypatch.setattr(client, "execute_command", contended) + with Redis() as client: + monkeypatch.setattr(client, "execute_command", execute) + wrapper = RedisClientWrapper() + wrapper.initialize(client) + monkeypatch.setattr(api_token_service_module, "redis_client", wrapper) + yield client, commands - with ( - patch.object(api_token_service_module.ApiTokenCache, "get") as mock_cache_get, - patch.object(api_token_service_module, "query_token_from_db", return_value=db_token) as mock_query_db, - ): - result = api_token_service_module.fetch_token_with_single_flight(auth_token, scope) + class TestFetchTokenWithSingleFlight: + def test_should_return_cached_token_when_lock_acquired_and_cache_filled( + self, single_flight_redis: tuple[Redis, list[tuple[object, ...]]] + ) -> None: + auth_token = "token-123" + scope = "app" + cached_token = CachedApiToken( + id="id-1", + app_id="app-1", + tenant_id="tenant-1", + type="app", + token=auth_token, + last_used_at=None, + created_at=None, + ) + _client, commands = single_flight_redis - assert result == db_token - mock_cache_get.assert_not_called() - mock_query_db.assert_called_once_with(auth_token, scope) - assert commands - assert all(command[0] == "SET" for command in commands) + with ( + patch.object(api_token_service_module.ApiTokenCache, "get", return_value=cached_token), + patch.object(api_token_service_module, "query_token_from_db") as mock_query_db, + ): + result = api_token_service_module.fetch_token_with_single_flight(auth_token, scope) - def test_should_reraise_unauthorized_from_db_query( - self, single_flight_redis: tuple[Redis, list[tuple[object, ...]]] - ) -> None: - auth_token = "token-123" - scope = "app" - _client, commands = single_flight_redis + assert result == cached_token + _assert_lock_released(commands) + mock_query_db.assert_not_called() - with ( - patch.object(api_token_service_module.ApiTokenCache, "get", return_value=None), - patch.object( - api_token_service_module, - "query_token_from_db", - side_effect=Unauthorized("Access token is invalid"), - ), - ): - with pytest.raises(Unauthorized, match="Access token is invalid"): - api_token_service_module.fetch_token_with_single_flight(auth_token, scope) + def test_should_query_db_when_lock_acquired_and_cache_missed( + self, single_flight_redis: tuple[Redis, list[tuple[object, ...]]] + ) -> None: + auth_token = "token-123" + scope = "app" + db_token = _api_token() + _client, commands = single_flight_redis - _assert_lock_released(commands) + with ( + patch.object(api_token_service_module.ApiTokenCache, "get", return_value=None), + patch.object(api_token_service_module, "query_token_from_db", return_value=db_token) as mock_query_db, + ): + result = api_token_service_module.fetch_token_with_single_flight(auth_token, scope) - def test_should_fallback_to_db_query_when_lock_raises_exception( - self, single_flight_redis: tuple[Redis, list[tuple[object, ...]]], monkeypatch: pytest.MonkeyPatch - ) -> None: - auth_token = "token-123" - scope = "app" - db_token = _api_token() - client, commands = single_flight_redis + assert result == db_token + mock_query_db.assert_called_once_with(auth_token, scope) + _assert_lock_released(commands) - def unavailable(*args: object, **_kwargs: object) -> bool: - commands.append(args) - raise RuntimeError("redis lock error") + def test_should_query_db_directly_when_lock_not_acquired( + self, single_flight_redis: tuple[Redis, list[tuple[object, ...]]], monkeypatch: pytest.MonkeyPatch + ) -> None: + auth_token = "token-123" + scope = "app" + db_token = _api_token() + client, commands = single_flight_redis - monkeypatch.setattr(client, "execute_command", unavailable) + def contended(*args: object, **_kwargs: object) -> bool: + commands.append(args) + assert args[0] == "SET" + return False - with ( - patch.object(api_token_service_module, "query_token_from_db", return_value=db_token) as mock_query_db, - ): - result = api_token_service_module.fetch_token_with_single_flight(auth_token, scope) + monkeypatch.setattr(client, "execute_command", contended) - assert result == db_token - mock_query_db.assert_called_once_with(auth_token, scope) - assert len(commands) == 1 - assert commands[0][:2] == ("SET", "api_token_query_lock:app:token-123") + with ( + patch.object(api_token_service_module.ApiTokenCache, "get") as mock_cache_get, + patch.object(api_token_service_module, "query_token_from_db", return_value=db_token) as mock_query_db, + ): + result = api_token_service_module.fetch_token_with_single_flight(auth_token, scope) + assert result == db_token + mock_cache_get.assert_not_called() + mock_query_db.assert_called_once_with(auth_token, scope) + assert commands + assert all(command[0] == "SET" for command in commands) -class TestApiTokenCacheTenantBranches: - @patch("services.api_token_service.redis_client") - def test_delete_with_scope_should_remove_from_tenant_index_when_tenant_found(self, mock_redis): - token = "token-123" - scope = "app" - cache_key = ApiTokenCache._make_cache_key(token, scope) - cached_token = CachedApiToken( - id="id-1", - app_id="app-1", - tenant_id="tenant-1", - type="app", - token=token, - last_used_at=None, - created_at=None, - ) - mock_redis.get.return_value = cached_token.model_dump_json().encode("utf-8") + def test_should_reraise_unauthorized_from_db_query( + self, single_flight_redis: tuple[Redis, list[tuple[object, ...]]] + ) -> None: + auth_token = "token-123" + scope = "app" + _client, commands = single_flight_redis - with patch.object(ApiTokenCache, "_remove_from_tenant_index") as mock_remove_index: + with ( + patch.object(api_token_service_module.ApiTokenCache, "get", return_value=None), + patch.object( + api_token_service_module, + "query_token_from_db", + side_effect=Unauthorized("Access token is invalid"), + ), + ): + with pytest.raises(Unauthorized, match="Access token is invalid"): + api_token_service_module.fetch_token_with_single_flight(auth_token, scope) + + _assert_lock_released(commands) + + def test_should_fallback_to_db_query_when_lock_raises_exception( + self, single_flight_redis: tuple[Redis, list[tuple[object, ...]]], monkeypatch: pytest.MonkeyPatch + ) -> None: + auth_token = "token-123" + scope = "app" + db_token = _api_token() + client, commands = single_flight_redis + + def unavailable(*args: object, **_kwargs: object) -> bool: + commands.append(args) + raise RuntimeError("redis lock error") + + monkeypatch.setattr(client, "execute_command", unavailable) + + with ( + patch.object(api_token_service_module, "query_token_from_db", return_value=db_token) as mock_query_db, + ): + result = api_token_service_module.fetch_token_with_single_flight(auth_token, scope) + + assert result == db_token + mock_query_db.assert_called_once_with(auth_token, scope) + assert len(commands) == 1 + assert commands[0][:2] == ("SET", "api_token_query_lock:app:token-123") + + class TestApiTokenCacheTenantBranches: + @patch("services.api_token_service.redis_client") + def test_delete_with_scope_should_remove_from_tenant_index_when_tenant_found(self, mock_redis): + token = "token-123" + scope = "app" + cache_key = ApiTokenCache._make_cache_key(token, scope) + cached_token = CachedApiToken( + id="id-1", + app_id="app-1", + tenant_id="tenant-1", + type="app", + token=token, + last_used_at=None, + created_at=None, + ) + mock_redis.get.return_value = cached_token.model_dump_json().encode("utf-8") + + with patch.object(ApiTokenCache, "_remove_from_tenant_index") as mock_remove_index: + result = ApiTokenCache.delete(token, scope) + + assert result is True + mock_redis.delete.assert_called_once_with(cache_key) + mock_remove_index.assert_called_once_with("tenant-1", cache_key) + + @patch("services.api_token_service.redis_client") + def test_invalidate_by_tenant_should_delete_all_indexed_cache_keys(self, mock_redis): + tenant_id = "tenant-1" + index_key = ApiTokenCache._make_tenant_index_key(tenant_id) + mock_redis.smembers.return_value = { + b"api_token:app:token-1", + b"api_token:any:token-2", + } + + result = ApiTokenCache.invalidate_by_tenant(tenant_id) + + assert result is True + mock_redis.smembers.assert_called_once_with(index_key) + mock_redis.delete.assert_any_call("api_token:app:token-1") + mock_redis.delete.assert_any_call("api_token:any:token-2") + mock_redis.delete.assert_any_call(index_key) + + class TestApiTokenCacheCoreBranches: + def test_cached_api_token_repr_should_include_id_and_type(self): + token = CachedApiToken( + id="id-123", + app_id="app-123", + tenant_id="tenant-123", + type="app", + token="token-123", + last_used_at=None, + created_at=None, + ) + assert repr(token) == "" + + def test_serialize_token_should_handle_cached_api_token_instances(self): + token = CachedApiToken( + id="id-123", + app_id="app-123", + tenant_id="tenant-123", + type="app", + token="token-123", + last_used_at=None, + created_at=None, + ) + serialized = ApiTokenCache._serialize_token(token) + assert isinstance(serialized, bytes) + assert b'"id":"id-123"' in serialized + + def test_deserialize_token_should_return_none_for_null_markers(self): + assert ApiTokenCache._deserialize_token("null") is None + assert ApiTokenCache._deserialize_token(b"null") is None + + def test_deserialize_token_should_return_none_for_invalid_payload(self): + assert ApiTokenCache._deserialize_token("not-json") is None + + @patch("services.api_token_service.redis_client") + def test_get_should_return_none_on_cache_miss(self, mock_redis): + mock_redis.get.return_value = None + result = ApiTokenCache.get("token-123", "app") + assert result is None + + @patch("services.api_token_service.redis_client") + def test_get_should_deserialize_cached_payload_on_cache_hit(self, mock_redis): + token = CachedApiToken( + id="id-123", + app_id="app-123", + tenant_id="tenant-123", + type="app", + token="token-123", + last_used_at=None, + created_at=None, + ) + mock_redis.get.return_value = token.model_dump_json().encode("utf-8") + result = ApiTokenCache.get("token-123", "app") + assert isinstance(result, CachedApiToken) + assert result.id == "id-123" + + @patch("services.api_token_service.redis_client") + def test_add_to_tenant_index_should_skip_when_tenant_id_missing(self, mock_redis): + ApiTokenCache._add_to_tenant_index(None, "api_token:app:token-123") + mock_redis.sadd.assert_not_called() + + @patch("services.api_token_service.redis_client") + def test_add_to_tenant_index_should_swallow_index_update_errors(self, mock_redis): + mock_redis.sadd.side_effect = Exception("redis down") + ApiTokenCache._add_to_tenant_index("tenant-123", "api_token:app:token-123") + mock_redis.sadd.assert_called_once() + + @patch("services.api_token_service.redis_client") + def test_remove_from_tenant_index_should_skip_when_tenant_id_missing(self, mock_redis): + ApiTokenCache._remove_from_tenant_index(None, "api_token:app:token-123") + mock_redis.srem.assert_not_called() + + @patch("services.api_token_service.redis_client") + def test_remove_from_tenant_index_should_swallow_redis_errors(self, mock_redis): + mock_redis.srem.side_effect = Exception("redis down") + ApiTokenCache._remove_from_tenant_index("tenant-123", "api_token:app:token-123") + mock_redis.srem.assert_called_once() + + @patch("services.api_token_service.redis_client") + def test_set_should_return_false_when_cache_write_raises_exception(self, mock_redis): + mock_redis.setex.side_effect = Exception("redis write failed") + api_token = _api_token() + result = ApiTokenCache.set("token-123", "app", api_token) + assert result is False + + @patch("services.api_token_service.redis_client") + def test_delete_without_scope_should_return_false_when_scan_fails(self, mock_redis): + mock_redis.scan_iter.side_effect = Exception("scan failed") + result = ApiTokenCache.delete("token-123", None) + assert result is False + + @patch("services.api_token_service.redis_client") + def test_delete_with_scope_should_continue_when_tenant_lookup_raises(self, mock_redis): + token = "token-123" + scope = "app" + cache_key = ApiTokenCache._make_cache_key(token, scope) + mock_redis.get.side_effect = Exception("get failed") result = ApiTokenCache.delete(token, scope) + assert result is True + mock_redis.delete.assert_called_once_with(cache_key) - assert result is True - mock_redis.delete.assert_called_once_with(cache_key) - mock_remove_index.assert_called_once_with("tenant-1", cache_key) + @patch("services.api_token_service.redis_client") + def test_delete_with_scope_should_return_false_when_delete_raises(self, mock_redis): + mock_redis.get.return_value = None + mock_redis.delete.side_effect = Exception("delete failed") + result = ApiTokenCache.delete("token-123", "app") + assert result is False - @patch("services.api_token_service.redis_client") - def test_invalidate_by_tenant_should_delete_all_indexed_cache_keys(self, mock_redis): - tenant_id = "tenant-1" - index_key = ApiTokenCache._make_tenant_index_key(tenant_id) - mock_redis.smembers.return_value = { - b"api_token:app:token-1", - b"api_token:any:token-2", - } + @patch("services.api_token_service.redis_client") + def test_invalidate_by_tenant_should_return_true_when_index_not_found(self, mock_redis): + mock_redis.smembers.return_value = set() + result = ApiTokenCache.invalidate_by_tenant("tenant-123") + assert result is True + mock_redis.delete.assert_not_called() - result = ApiTokenCache.invalidate_by_tenant(tenant_id) - - assert result is True - mock_redis.smembers.assert_called_once_with(index_key) - mock_redis.delete.assert_any_call("api_token:app:token-1") - mock_redis.delete.assert_any_call("api_token:any:token-2") - mock_redis.delete.assert_any_call(index_key) - - -class TestApiTokenCacheCoreBranches: - def test_cached_api_token_repr_should_include_id_and_type(self): - token = CachedApiToken( - id="id-123", - app_id="app-123", - tenant_id="tenant-123", - type="app", - token="token-123", - last_used_at=None, - created_at=None, - ) - assert repr(token) == "" - - def test_serialize_token_should_handle_cached_api_token_instances(self): - token = CachedApiToken( - id="id-123", - app_id="app-123", - tenant_id="tenant-123", - type="app", - token="token-123", - last_used_at=None, - created_at=None, - ) - serialized = ApiTokenCache._serialize_token(token) - assert isinstance(serialized, bytes) - assert b'"id":"id-123"' in serialized - - def test_deserialize_token_should_return_none_for_null_markers(self): - assert ApiTokenCache._deserialize_token("null") is None - assert ApiTokenCache._deserialize_token(b"null") is None - - def test_deserialize_token_should_return_none_for_invalid_payload(self): - assert ApiTokenCache._deserialize_token("not-json") is None - - @patch("services.api_token_service.redis_client") - def test_get_should_return_none_on_cache_miss(self, mock_redis): - mock_redis.get.return_value = None - result = ApiTokenCache.get("token-123", "app") - assert result is None - - @patch("services.api_token_service.redis_client") - def test_get_should_deserialize_cached_payload_on_cache_hit(self, mock_redis): - token = CachedApiToken( - id="id-123", - app_id="app-123", - tenant_id="tenant-123", - type="app", - token="token-123", - last_used_at=None, - created_at=None, - ) - mock_redis.get.return_value = token.model_dump_json().encode("utf-8") - result = ApiTokenCache.get("token-123", "app") - assert isinstance(result, CachedApiToken) - assert result.id == "id-123" - - @patch("services.api_token_service.redis_client") - def test_add_to_tenant_index_should_skip_when_tenant_id_missing(self, mock_redis): - ApiTokenCache._add_to_tenant_index(None, "api_token:app:token-123") - mock_redis.sadd.assert_not_called() - - @patch("services.api_token_service.redis_client") - def test_add_to_tenant_index_should_swallow_index_update_errors(self, mock_redis): - mock_redis.sadd.side_effect = Exception("redis down") - ApiTokenCache._add_to_tenant_index("tenant-123", "api_token:app:token-123") - mock_redis.sadd.assert_called_once() - - @patch("services.api_token_service.redis_client") - def test_remove_from_tenant_index_should_skip_when_tenant_id_missing(self, mock_redis): - ApiTokenCache._remove_from_tenant_index(None, "api_token:app:token-123") - mock_redis.srem.assert_not_called() - - @patch("services.api_token_service.redis_client") - def test_remove_from_tenant_index_should_swallow_redis_errors(self, mock_redis): - mock_redis.srem.side_effect = Exception("redis down") - ApiTokenCache._remove_from_tenant_index("tenant-123", "api_token:app:token-123") - mock_redis.srem.assert_called_once() - - @patch("services.api_token_service.redis_client") - def test_set_should_return_false_when_cache_write_raises_exception(self, mock_redis): - mock_redis.setex.side_effect = Exception("redis write failed") - api_token = _api_token() - result = ApiTokenCache.set("token-123", "app", api_token) - assert result is False - - @patch("services.api_token_service.redis_client") - def test_delete_without_scope_should_return_false_when_scan_fails(self, mock_redis): - mock_redis.scan_iter.side_effect = Exception("scan failed") - result = ApiTokenCache.delete("token-123", None) - assert result is False - - @patch("services.api_token_service.redis_client") - def test_delete_with_scope_should_continue_when_tenant_lookup_raises(self, mock_redis): - token = "token-123" - scope = "app" - cache_key = ApiTokenCache._make_cache_key(token, scope) - mock_redis.get.side_effect = Exception("get failed") - result = ApiTokenCache.delete(token, scope) - assert result is True - mock_redis.delete.assert_called_once_with(cache_key) - - @patch("services.api_token_service.redis_client") - def test_delete_with_scope_should_return_false_when_delete_raises(self, mock_redis): - mock_redis.get.return_value = None - mock_redis.delete.side_effect = Exception("delete failed") - result = ApiTokenCache.delete("token-123", "app") - assert result is False - - @patch("services.api_token_service.redis_client") - def test_invalidate_by_tenant_should_return_true_when_index_not_found(self, mock_redis): - mock_redis.smembers.return_value = set() - result = ApiTokenCache.invalidate_by_tenant("tenant-123") - assert result is True - mock_redis.delete.assert_not_called() - - @patch("services.api_token_service.redis_client") - def test_invalidate_by_tenant_should_return_false_when_redis_raises(self, mock_redis): - mock_redis.smembers.side_effect = Exception("redis failed") - result = ApiTokenCache.invalidate_by_tenant("tenant-123") - assert result is False + @patch("services.api_token_service.redis_client") + def test_invalidate_by_tenant_should_return_false_when_redis_raises(self, mock_redis): + mock_redis.smembers.side_effect = Exception("redis failed") + result = ApiTokenCache.invalidate_by_tenant("tenant-123") + assert result is False diff --git a/api/tests/unit_tests/services/test_app_dsl_service.py b/api/tests/unit_tests/services/test_app_dsl_service.py index ddc2b3dedd7..673f36e01c9 100644 --- a/api/tests/unit_tests/services/test_app_dsl_service.py +++ b/api/tests/unit_tests/services/test_app_dsl_service.py @@ -30,8 +30,8 @@ from services.entities.site_dsl import SiteDsl from services.errors.app import DraftChangedError, WorkflowNotFoundError from services.errors.base import NoPermissionError from services.workflow.graph_diff import draft_token -from tests.unit_tests.config_override import apply_config_overrides -from tests.unit_tests.model_factories import make_account, make_app, make_tenant, make_workflow +from tests.helpers.unit.config_override import apply_config_overrides +from tests.helpers.unit.model_factories import make_account, make_app, make_tenant, make_workflow _OVERWRITE_APP_ID = "11111111-1111-4111-8111-111111111111" _TENANT_ID = "22222222-2222-4222-8222-222222222222" diff --git a/api/tests/unit_tests/services/test_app_generate_service.py b/api/tests/unit_tests/services/test_app_generate_service.py index 1cc95f8ddda..854323c1b91 100644 --- a/api/tests/unit_tests/services/test_app_generate_service.py +++ b/api/tests/unit_tests/services/test_app_generate_service.py @@ -1,3 +1,33 @@ +import json +import threading +import uuid +from collections import defaultdict, deque +from collections.abc import Callable +from contextlib import contextmanager +from typing import Any +from unittest.mock import MagicMock + +import pytest +from pytest_mock import MockerFixture +from sqlalchemy.orm import Session + +import services.app_generate_service as ags_module +from core.app.apps.message_based_app_generator import MessageBasedAppGenerator +from core.app.apps.message_generator import MessageGenerator +from core.app.entities.app_invoke_entities import InvokeFrom +from enums import DeploymentEdition, QuotaType +from models.account import Account +from models.enums import AppStatus +from models.model import App, AppMode +from models.workflow import Workflow, WorkflowType +from services.app_generate_service import AppGenerateService +from services.errors.app import ( + TriggerWorkflowServiceModeUnavailableError, + WorkflowIdFormatError, + WorkflowNotFoundError, +) +from tests.helpers.unit.config_override import apply_config_overrides + """ Comprehensive unit tests for services.app_generate_service.AppGenerateService. @@ -12,31 +42,6 @@ Covers: - _get_workflow (debugger / non-debugger / specific id / invalid format / not found) """ -import json -import threading -import uuid -from collections.abc import Callable -from contextlib import contextmanager -from unittest.mock import MagicMock - -import pytest -from pytest_mock import MockerFixture -from sqlalchemy.orm import Session - -import services.app_generate_service as ags_module -from core.app.entities.app_invoke_entities import InvokeFrom -from enums import DeploymentEdition, QuotaType -from models.account import Account -from models.enums import AppStatus -from models.model import App, AppMode -from models.workflow import Workflow, WorkflowType -from services.app_generate_service import AppGenerateService -from services.errors.app import ( - TriggerWorkflowServiceModeUnavailableError, - WorkflowIdFormatError, - WorkflowNotFoundError, -) - # --------------------------------------------------------------------------- # Helpers / Fakes @@ -156,899 +161,1101 @@ def _unexpected_timer(interval: float, function: Callable[[], bool]) -> _FakeTim raise AssertionError("streams must not create a fallback timer") -class TestBuildStreamingTaskOnSubscribe: - def test_streams_starts_only_when_hook_is_invoked_without_creating_timer( - self, monkeypatch: pytest.MonkeyPatch, config_overrides: Callable[..., None] - ): - config_overrides(PUBSUB_REDIS_CHANNEL_TYPE="streams") +class TestAppGenerateServiceCases: + class TestBuildStreamingTaskOnSubscribe: + def test_streams_starts_only_when_hook_is_invoked_without_creating_timer( + self, monkeypatch: pytest.MonkeyPatch, config_overrides: Callable[..., None] + ): + config_overrides(PUBSUB_REDIS_CHANNEL_TYPE="streams") - monkeypatch.setattr(ags_module.threading, "Timer", _unexpected_timer) - called: list[int] = [] + monkeypatch.setattr(ags_module.threading, "Timer", _unexpected_timer) + called: list[int] = [] - on_subscribe = AppGenerateService._build_streaming_task_on_subscribe(lambda: called.append(1)) + on_subscribe = AppGenerateService._build_streaming_task_on_subscribe(lambda: called.append(1)) - assert called == [] - on_subscribe() - on_subscribe() - assert called == [1] + assert called == [] + on_subscribe() + on_subscribe() + assert called == [1] - @pytest.mark.parametrize("channel_type", ["pubsub", "sharded"]) - def test_pubsub_transports_keep_subscribe_hook_and_fallback_timer( - self, - monkeypatch: pytest.MonkeyPatch, - channel_type: str, - config_overrides: Callable[..., None], - ): - config_overrides(PUBSUB_REDIS_CHANNEL_TYPE=channel_type) - timers: list[_FakeTimer] = [] + @pytest.mark.parametrize("channel_type", ["pubsub", "sharded"]) + def test_pubsub_transports_keep_subscribe_hook_and_fallback_timer( + self, + monkeypatch: pytest.MonkeyPatch, + channel_type: str, + config_overrides: Callable[..., None], + ): + config_overrides(PUBSUB_REDIS_CHANNEL_TYPE=channel_type) + timers: list[_FakeTimer] = [] - def build_timer(interval: float, function: Callable[[], bool]) -> _FakeTimer: - timer = _FakeTimer(interval, function) - timers.append(timer) - return timer + def build_timer(interval: float, function: Callable[[], bool]) -> _FakeTimer: + timer = _FakeTimer(interval, function) + timers.append(timer) + return timer - monkeypatch.setattr(ags_module.threading, "Timer", build_timer) - called: list[int] = [] + monkeypatch.setattr(ags_module.threading, "Timer", build_timer) + called: list[int] = [] - on_subscribe = AppGenerateService._build_streaming_task_on_subscribe(lambda: called.append(1)) + on_subscribe = AppGenerateService._build_streaming_task_on_subscribe(lambda: called.append(1)) - assert called == [] - assert len(timers) == 1 - assert timers[0].interval == ags_module.SSE_TASK_START_FALLBACK_MS / 1000.0 - assert timers[0].started is True + assert called == [] + assert len(timers) == 1 + assert timers[0].interval == ags_module.SSE_TASK_START_FALLBACK_MS / 1000.0 + assert timers[0].started is True - on_subscribe() + on_subscribe() - assert called == [1] - assert timers[0].cancelled is True + assert called == [1] + assert timers[0].cancelled is True - def test_pubsub_fallback_starts_task_if_hook_is_never_invoked( - self, monkeypatch: pytest.MonkeyPatch, config_overrides: Callable[..., None] - ): - config_overrides(PUBSUB_REDIS_CHANNEL_TYPE="pubsub") - timers: list[_FakeTimer] = [] + def test_pubsub_fallback_starts_task_if_hook_is_never_invoked( + self, monkeypatch: pytest.MonkeyPatch, config_overrides: Callable[..., None] + ): + config_overrides(PUBSUB_REDIS_CHANNEL_TYPE="pubsub") + timers: list[_FakeTimer] = [] - def build_timer(interval: float, function: Callable[[], bool]) -> _FakeTimer: - timer = _FakeTimer(interval, function) - timers.append(timer) - return timer + def build_timer(interval: float, function: Callable[[], bool]) -> _FakeTimer: + timer = _FakeTimer(interval, function) + timers.append(timer) + return timer - monkeypatch.setattr(ags_module.threading, "Timer", build_timer) - called: list[int] = [] - on_subscribe = AppGenerateService._build_streaming_task_on_subscribe(lambda: called.append(1)) + monkeypatch.setattr(ags_module.threading, "Timer", build_timer) + called: list[int] = [] + on_subscribe = AppGenerateService._build_streaming_task_on_subscribe(lambda: called.append(1)) - assert timers[0].function() is True - on_subscribe() - assert called == [1] + assert timers[0].function() is True + on_subscribe() + assert called == [1] - def test_streams_retries_after_enqueue_failure( - self, monkeypatch: pytest.MonkeyPatch, config_overrides: Callable[..., None] - ): - config_overrides(PUBSUB_REDIS_CHANNEL_TYPE="streams") - monkeypatch.setattr(ags_module.threading, "Timer", _unexpected_timer) - call_count = 0 + def test_streams_retries_after_enqueue_failure( + self, monkeypatch: pytest.MonkeyPatch, config_overrides: Callable[..., None] + ): + config_overrides(PUBSUB_REDIS_CHANNEL_TYPE="streams") + monkeypatch.setattr(ags_module.threading, "Timer", _unexpected_timer) + call_count = 0 - def _bad(): - nonlocal call_count - call_count += 1 - if call_count == 1: - raise RuntimeError("boom") + def _bad(): + nonlocal call_count + call_count += 1 + if call_count == 1: + raise RuntimeError("boom") - on_subscribe = AppGenerateService._build_streaming_task_on_subscribe(_bad) - on_subscribe() - assert call_count == 1 - on_subscribe() - assert call_count == 2 + on_subscribe = AppGenerateService._build_streaming_task_on_subscribe(_bad) + on_subscribe() + assert call_count == 1 + on_subscribe() + assert call_count == 2 - def test_concurrent_subscribe_only_starts_once( - self, monkeypatch: pytest.MonkeyPatch, config_overrides: Callable[..., None] - ): - config_overrides(PUBSUB_REDIS_CHANNEL_TYPE="streams") - monkeypatch.setattr(ags_module.threading, "Timer", _unexpected_timer) - call_count = 0 + def test_concurrent_subscribe_only_starts_once( + self, monkeypatch: pytest.MonkeyPatch, config_overrides: Callable[..., None] + ): + config_overrides(PUBSUB_REDIS_CHANNEL_TYPE="streams") + monkeypatch.setattr(ags_module.threading, "Timer", _unexpected_timer) + call_count = 0 - def _inc(): - nonlocal call_count - call_count += 1 + def _inc(): + nonlocal call_count + call_count += 1 - cb = AppGenerateService._build_streaming_task_on_subscribe(_inc) - threads = [threading.Thread(target=cb) for _ in range(10)] - for t in threads: - t.start() - for t in threads: - t.join() - assert call_count == 1 + cb = AppGenerateService._build_streaming_task_on_subscribe(_inc) + threads = [threading.Thread(target=cb) for _ in range(10)] + for t in threads: + t.start() + for t in threads: + t.join() + assert call_count == 1 + # --------------------------------------------------------------------------- + # _get_max_active_requests + # --------------------------------------------------------------------------- + class TestGetMaxActiveRequests: + def test_both_zero_returns_zero(self, config_overrides: Callable[..., None]): + config_overrides(APP_MAX_ACTIVE_REQUESTS=0, APP_DEFAULT_ACTIVE_REQUESTS=0) + app = _make_app(AppMode.CHAT, max_active_requests=0) + assert AppGenerateService._get_max_active_requests(app) == 0 -# --------------------------------------------------------------------------- -# _get_max_active_requests -# --------------------------------------------------------------------------- -class TestGetMaxActiveRequests: - def test_both_zero_returns_zero(self, config_overrides: Callable[..., None]): - config_overrides(APP_MAX_ACTIVE_REQUESTS=0, APP_DEFAULT_ACTIVE_REQUESTS=0) - app = _make_app(AppMode.CHAT, max_active_requests=0) - assert AppGenerateService._get_max_active_requests(app) == 0 + def test_app_limit_only(self, config_overrides: Callable[..., None]): + config_overrides(APP_MAX_ACTIVE_REQUESTS=0, APP_DEFAULT_ACTIVE_REQUESTS=0) + app = _make_app(AppMode.CHAT, max_active_requests=5) + assert AppGenerateService._get_max_active_requests(app) == 5 - def test_app_limit_only(self, config_overrides: Callable[..., None]): - config_overrides(APP_MAX_ACTIVE_REQUESTS=0, APP_DEFAULT_ACTIVE_REQUESTS=0) - app = _make_app(AppMode.CHAT, max_active_requests=5) - assert AppGenerateService._get_max_active_requests(app) == 5 + def test_config_limit_only(self, config_overrides: Callable[..., None]): + config_overrides(APP_MAX_ACTIVE_REQUESTS=10, APP_DEFAULT_ACTIVE_REQUESTS=0) + app = _make_app(AppMode.CHAT, max_active_requests=0) + assert AppGenerateService._get_max_active_requests(app) == 10 - def test_config_limit_only(self, config_overrides: Callable[..., None]): - config_overrides(APP_MAX_ACTIVE_REQUESTS=10, APP_DEFAULT_ACTIVE_REQUESTS=0) - app = _make_app(AppMode.CHAT, max_active_requests=0) - assert AppGenerateService._get_max_active_requests(app) == 10 + def test_both_non_zero_returns_min(self, config_overrides: Callable[..., None]): + config_overrides(APP_MAX_ACTIVE_REQUESTS=20, APP_DEFAULT_ACTIVE_REQUESTS=0) + app = _make_app(AppMode.CHAT, max_active_requests=5) + assert AppGenerateService._get_max_active_requests(app) == 5 - def test_both_non_zero_returns_min(self, config_overrides: Callable[..., None]): - config_overrides(APP_MAX_ACTIVE_REQUESTS=20, APP_DEFAULT_ACTIVE_REQUESTS=0) - app = _make_app(AppMode.CHAT, max_active_requests=5) - assert AppGenerateService._get_max_active_requests(app) == 5 + def test_default_active_requests_used_when_app_has_none(self, config_overrides: Callable[..., None]): + config_overrides(APP_MAX_ACTIVE_REQUESTS=0, APP_DEFAULT_ACTIVE_REQUESTS=15) + app = _make_app(AppMode.CHAT, max_active_requests=0) + assert AppGenerateService._get_max_active_requests(app) == 15 - def test_default_active_requests_used_when_app_has_none(self, config_overrides: Callable[..., None]): - config_overrides(APP_MAX_ACTIVE_REQUESTS=0, APP_DEFAULT_ACTIVE_REQUESTS=15) - app = _make_app(AppMode.CHAT, max_active_requests=0) - assert AppGenerateService._get_max_active_requests(app) == 15 + # --------------------------------------------------------------------------- + # generate – every AppMode branch + # --------------------------------------------------------------------------- + class TestGenerate(_RealSessionTest): + """Tests for AppGenerateService.generate covering each mode.""" + @pytest.fixture(autouse=True) + def _common(self, mocker: MockerFixture, config_overrides: Callable[..., None]): + config_overrides(DEPLOYMENT_EDITION=DeploymentEdition.COMMUNITY) + mocker.patch("services.app_generate_service.RateLimit", _DummyRateLimit) + # Prevent AppExecutionParams.new from touching real models via isinstance + mocker.patch( + "services.app_generate_service.rate_limit_context", + _noop_rate_limit_context, + ) -# --------------------------------------------------------------------------- -# generate – every AppMode branch -# --------------------------------------------------------------------------- -class TestGenerate(_RealSessionTest): - """Tests for AppGenerateService.generate covering each mode.""" + # -- COMPLETION --------------------------------------------------------- + def test_completion_mode(self, mocker: MockerFixture): + gen_spy = mocker.patch( + "services.app_generate_service.CompletionAppGenerator.generate", + return_value={"result": "ok"}, + ) + mocker.patch( + "services.app_generate_service.CompletionAppGenerator.convert_to_event_stream", + side_effect=lambda x: x, + ) + result = AppGenerateService.generate( + app_model=_make_app(AppMode.COMPLETION), + user=_make_user(), + args={"inputs": {}}, + invoke_from=InvokeFrom.SERVICE_API, + streaming=False, + session=self.session, + ) + assert result == {"result": "ok"} + gen_spy.assert_called_once() - @pytest.fixture(autouse=True) - def _common(self, mocker: MockerFixture, config_overrides: Callable[..., None]): - config_overrides(DEPLOYMENT_EDITION=DeploymentEdition.COMMUNITY) - mocker.patch("services.app_generate_service.RateLimit", _DummyRateLimit) - # Prevent AppExecutionParams.new from touching real models via isinstance - mocker.patch( - "services.app_generate_service.rate_limit_context", - _noop_rate_limit_context, - ) + # -- AGENT_CHAT via mode ------------------------------------------------ + def test_agent_chat_mode(self, mocker: MockerFixture): + gen_spy = mocker.patch( + "services.app_generate_service.AgentChatAppGenerator.generate", + return_value={"result": "agent"}, + ) + mocker.patch( + "services.app_generate_service.AgentChatAppGenerator.convert_to_event_stream", + side_effect=lambda x: x, + ) + result = AppGenerateService.generate( + app_model=_make_app(AppMode.AGENT_CHAT), + user=_make_user(), + args={"inputs": {}}, + invoke_from=InvokeFrom.SERVICE_API, + streaming=False, + session=self.session, + ) + assert result == {"result": "agent"} + gen_spy.assert_called_once() - # -- COMPLETION --------------------------------------------------------- - def test_completion_mode(self, mocker: MockerFixture): - gen_spy = mocker.patch( - "services.app_generate_service.CompletionAppGenerator.generate", - return_value={"result": "ok"}, - ) - mocker.patch( - "services.app_generate_service.CompletionAppGenerator.convert_to_event_stream", - side_effect=lambda x: x, - ) - result = AppGenerateService.generate( - app_model=_make_app(AppMode.COMPLETION), - user=_make_user(), - args={"inputs": {}}, - invoke_from=InvokeFrom.SERVICE_API, - streaming=False, - session=self.session, - ) - assert result == {"result": "ok"} - gen_spy.assert_called_once() + # -- AGENT_CHAT via is_agent flag (non-AGENT_CHAT mode) ----------------- + def test_agent_via_is_agent_flag(self, mocker: MockerFixture): + gen_spy = mocker.patch( + "services.app_generate_service.AgentChatAppGenerator.generate", + return_value={"result": "agent-via-flag"}, + ) + mocker.patch( + "services.app_generate_service.AgentChatAppGenerator.convert_to_event_stream", + side_effect=lambda x: x, + ) + app = _make_app(AppMode.CHAT) + is_agent = mocker.patch.object(App, "is_agent_with_session", return_value=True) + session = self.session + result = AppGenerateService.generate( + app_model=app, + user=_make_user(), + args={"inputs": {}}, + invoke_from=InvokeFrom.SERVICE_API, + streaming=False, + session=session, + ) + assert result == {"result": "agent-via-flag"} + gen_spy.assert_called_once() + is_agent.assert_called_once_with(session=session) - # -- AGENT_CHAT via mode ------------------------------------------------ - def test_agent_chat_mode(self, mocker: MockerFixture): - gen_spy = mocker.patch( - "services.app_generate_service.AgentChatAppGenerator.generate", - return_value={"result": "agent"}, - ) - mocker.patch( - "services.app_generate_service.AgentChatAppGenerator.convert_to_event_stream", - side_effect=lambda x: x, - ) - result = AppGenerateService.generate( - app_model=_make_app(AppMode.AGENT_CHAT), - user=_make_user(), - args={"inputs": {}}, - invoke_from=InvokeFrom.SERVICE_API, - streaming=False, - session=self.session, - ) - assert result == {"result": "agent"} - gen_spy.assert_called_once() + # -- AGENT -------------------------------------------------------------- + def test_agent_mode_passes_session(self, mocker: MockerFixture): + gen_spy = mocker.patch( + "services.app_generate_service.AgentAppGenerator.generate", + return_value={"result": "agent"}, + ) + mocker.patch( + "services.app_generate_service.AgentAppGenerator.convert_to_event_stream", + side_effect=lambda x: x, + ) + session = self.session - # -- AGENT_CHAT via is_agent flag (non-AGENT_CHAT mode) ----------------- - def test_agent_via_is_agent_flag(self, mocker: MockerFixture): - gen_spy = mocker.patch( - "services.app_generate_service.AgentChatAppGenerator.generate", - return_value={"result": "agent-via-flag"}, - ) - mocker.patch( - "services.app_generate_service.AgentChatAppGenerator.convert_to_event_stream", - side_effect=lambda x: x, - ) - app = _make_app(AppMode.CHAT) - is_agent = mocker.patch.object(App, "is_agent_with_session", return_value=True) - session = self.session - result = AppGenerateService.generate( - app_model=app, - user=_make_user(), - args={"inputs": {}}, - invoke_from=InvokeFrom.SERVICE_API, - streaming=False, - session=session, - ) - assert result == {"result": "agent-via-flag"} - gen_spy.assert_called_once() - is_agent.assert_called_once_with(session=session) + result = AppGenerateService.generate( + app_model=_make_app(AppMode.AGENT), + user=_make_user(), + args={"inputs": {}}, + invoke_from=InvokeFrom.SERVICE_API, + streaming=True, + session=session, + ) - # -- AGENT -------------------------------------------------------------- - def test_agent_mode_passes_session(self, mocker: MockerFixture): - gen_spy = mocker.patch( - "services.app_generate_service.AgentAppGenerator.generate", - return_value={"result": "agent"}, - ) - mocker.patch( - "services.app_generate_service.AgentAppGenerator.convert_to_event_stream", - side_effect=lambda x: x, - ) - session = self.session + assert result == {"result": "agent"} + assert gen_spy.call_args.kwargs["session"] is session - result = AppGenerateService.generate( - app_model=_make_app(AppMode.AGENT), - user=_make_user(), - args={"inputs": {}}, - invoke_from=InvokeFrom.SERVICE_API, - streaming=True, - session=session, - ) + # -- CHAT --------------------------------------------------------------- + def test_chat_mode(self, mocker: MockerFixture): + gen_spy = mocker.patch( + "services.app_generate_service.ChatAppGenerator.generate", + return_value={"result": "chat"}, + ) + mocker.patch( + "services.app_generate_service.ChatAppGenerator.convert_to_event_stream", + side_effect=lambda x: x, + ) + app = _make_app(AppMode.CHAT) + result = AppGenerateService.generate( + app_model=app, + user=_make_user(), + args={"inputs": {}}, + invoke_from=InvokeFrom.SERVICE_API, + streaming=False, + session=self.session, + ) + assert result == {"result": "chat"} + gen_spy.assert_called_once() - assert result == {"result": "agent"} - assert gen_spy.call_args.kwargs["session"] is session + # -- ADVANCED_CHAT blocking --------------------------------------------- + def test_advanced_chat_blocking(self, mocker: MockerFixture): + workflow = _make_workflow() + mocker.patch.object(AppGenerateService, "_get_workflow", return_value=workflow) - # -- CHAT --------------------------------------------------------------- - def test_chat_mode(self, mocker: MockerFixture): - gen_spy = mocker.patch( - "services.app_generate_service.ChatAppGenerator.generate", - return_value={"result": "chat"}, - ) - mocker.patch( - "services.app_generate_service.ChatAppGenerator.convert_to_event_stream", - side_effect=lambda x: x, - ) - app = _make_app(AppMode.CHAT) - result = AppGenerateService.generate( - app_model=app, - user=_make_user(), - args={"inputs": {}}, - invoke_from=InvokeFrom.SERVICE_API, - streaming=False, - session=self.session, - ) - assert result == {"result": "chat"} - gen_spy.assert_called_once() + retrieve_spy = mocker.patch("services.app_generate_service.AdvancedChatAppGenerator.retrieve_events") + gen_spy = mocker.patch( + "services.app_generate_service.AdvancedChatAppGenerator.generate", + return_value={"result": "advanced-blocking"}, + ) + mocker.patch( + "services.app_generate_service.AdvancedChatAppGenerator.convert_to_event_stream", + side_effect=lambda x: x, + ) - # -- ADVANCED_CHAT blocking --------------------------------------------- - def test_advanced_chat_blocking(self, mocker: MockerFixture): - workflow = _make_workflow() - mocker.patch.object(AppGenerateService, "_get_workflow", return_value=workflow) + session = self.session + result = AppGenerateService.generate( + app_model=_make_app(AppMode.ADVANCED_CHAT), + user=_make_user(), + args={"workflow_id": None, "query": "hi", "inputs": {}}, + invoke_from=InvokeFrom.SERVICE_API, + streaming=False, + session=session, + ) + assert result == {"result": "advanced-blocking"} + call_kwargs = gen_spy.call_args.kwargs + assert call_kwargs.get("streaming") is False + assert call_kwargs["session"] is session + retrieve_spy.assert_not_called() - retrieve_spy = mocker.patch("services.app_generate_service.AdvancedChatAppGenerator.retrieve_events") - gen_spy = mocker.patch( - "services.app_generate_service.AdvancedChatAppGenerator.generate", - return_value={"result": "advanced-blocking"}, - ) - mocker.patch( - "services.app_generate_service.AdvancedChatAppGenerator.convert_to_event_stream", - side_effect=lambda x: x, - ) + # -- ADVANCED_CHAT streaming -------------------------------------------- + def test_advanced_chat_streaming(self, mocker: MockerFixture, config_overrides: Callable[..., None]): + config_overrides(PUBSUB_REDIS_CHANNEL_TYPE="streams") + workflow = _make_workflow() + mocker.patch.object(AppGenerateService, "_get_workflow", return_value=workflow) + mocker.patch( + "services.app_generate_service.AppExecutionParams.new", + return_value=MagicMock(workflow_run_id="wfr-1", model_dump_json=MagicMock(return_value="{}")), + ) + delay_spy = mocker.patch("services.app_generate_service.workflow_based_app_execution_task.delay") + retrieve_spy = mocker.patch( + "services.app_generate_service.AdvancedChatAppGenerator.retrieve_events", + return_value=iter([]), + ) - session = self.session - result = AppGenerateService.generate( - app_model=_make_app(AppMode.ADVANCED_CHAT), - user=_make_user(), - args={"workflow_id": None, "query": "hi", "inputs": {}}, - invoke_from=InvokeFrom.SERVICE_API, - streaming=False, - session=session, - ) - assert result == {"result": "advanced-blocking"} - call_kwargs = gen_spy.call_args.kwargs - assert call_kwargs.get("streaming") is False - assert call_kwargs["session"] is session - retrieve_spy.assert_not_called() - - # -- ADVANCED_CHAT streaming -------------------------------------------- - def test_advanced_chat_streaming(self, mocker: MockerFixture, config_overrides: Callable[..., None]): - config_overrides(PUBSUB_REDIS_CHANNEL_TYPE="streams") - workflow = _make_workflow() - mocker.patch.object(AppGenerateService, "_get_workflow", return_value=workflow) - mocker.patch( - "services.app_generate_service.AppExecutionParams.new", - return_value=MagicMock(workflow_run_id="wfr-1", model_dump_json=MagicMock(return_value="{}")), - ) - delay_spy = mocker.patch("services.app_generate_service.workflow_based_app_execution_task.delay") - retrieve_spy = mocker.patch( - "services.app_generate_service.AdvancedChatAppGenerator.retrieve_events", - return_value=iter([]), - ) - - AppGenerateService.generate( - app_model=_make_app(AppMode.ADVANCED_CHAT), - user=_make_user(), - args={"workflow_id": None, "query": "hi", "inputs": {}}, - invoke_from=InvokeFrom.SERVICE_API, - streaming=True, - session=self.session, - ) - # In streaming mode it should go through retrieve_events, not generate - retrieve_spy.assert_called_once() - # Dispatch is gated on subscribe; simulate the SSE layer entering the - # subscription, which is what actually invokes on_subscribe. - on_subscribe = retrieve_spy.call_args.kwargs["on_subscribe"] - on_subscribe() - delay_spy.assert_called_once() - - # -- WORKFLOW blocking -------------------------------------------------- - def test_workflow_blocking(self, mocker: MockerFixture): - workflow = _make_workflow() - mocker.patch.object(AppGenerateService, "_get_workflow", return_value=workflow) - gen_spy = mocker.patch( - "services.app_generate_service.WorkflowAppGenerator.generate", - return_value={"result": "workflow-blocking"}, - ) - mocker.patch( - "services.app_generate_service.WorkflowAppGenerator.convert_to_event_stream", - side_effect=lambda x: x, - ) - - session = self.session - result = AppGenerateService.generate( - app_model=_make_app(AppMode.WORKFLOW), - user=_make_user(), - args={"inputs": {}}, - invoke_from=InvokeFrom.SERVICE_API, - streaming=False, - session=session, - ) - assert result == {"result": "workflow-blocking"} - call_kwargs = gen_spy.call_args.kwargs - assert call_kwargs.get("pause_state_config") is not None - assert call_kwargs["pause_state_config"].state_owner_user_id == "owner-id" - - @pytest.mark.parametrize( - "invoke_from", - [InvokeFrom.OPENAPI, InvokeFrom.SERVICE_API, InvokeFrom.WEB_APP], - ) - @pytest.mark.parametrize("node_type", ["trigger-plugin", "trigger-schedule", "trigger-webhook"]) - def test_trigger_workflow_rejects_manual_service_surfaces( - self, - invoke_from: InvokeFrom, - node_type: str, - mocker: MockerFixture, - ) -> None: - workflow = _make_workflow(node_types=(node_type,)) - mocker.patch.object(AppGenerateService, "_get_workflow", return_value=workflow) - generate = mocker.patch("services.app_generate_service.WorkflowAppGenerator.generate") - - with pytest.raises(TriggerWorkflowServiceModeUnavailableError): AppGenerateService.generate( + app_model=_make_app(AppMode.ADVANCED_CHAT), + user=_make_user(), + args={"workflow_id": None, "query": "hi", "inputs": {}}, + invoke_from=InvokeFrom.SERVICE_API, + streaming=True, + session=self.session, + ) + # In streaming mode it should go through retrieve_events, not generate + retrieve_spy.assert_called_once() + # Dispatch is gated on subscribe; simulate the SSE layer entering the + # subscription, which is what actually invokes on_subscribe. + on_subscribe = retrieve_spy.call_args.kwargs["on_subscribe"] + on_subscribe() + delay_spy.assert_called_once() + + # -- WORKFLOW blocking -------------------------------------------------- + def test_workflow_blocking(self, mocker: MockerFixture): + workflow = _make_workflow() + mocker.patch.object(AppGenerateService, "_get_workflow", return_value=workflow) + gen_spy = mocker.patch( + "services.app_generate_service.WorkflowAppGenerator.generate", + return_value={"result": "workflow-blocking"}, + ) + mocker.patch( + "services.app_generate_service.WorkflowAppGenerator.convert_to_event_stream", + side_effect=lambda x: x, + ) + + session = self.session + result = AppGenerateService.generate( app_model=_make_app(AppMode.WORKFLOW), user=_make_user(), args={"inputs": {}}, - invoke_from=invoke_from, + invoke_from=InvokeFrom.SERVICE_API, + streaming=False, + session=session, + ) + assert result == {"result": "workflow-blocking"} + call_kwargs = gen_spy.call_args.kwargs + assert call_kwargs.get("pause_state_config") is not None + assert call_kwargs["pause_state_config"].state_owner_user_id == "owner-id" + + @pytest.mark.parametrize( + "invoke_from", + [InvokeFrom.OPENAPI, InvokeFrom.SERVICE_API, InvokeFrom.WEB_APP], + ) + @pytest.mark.parametrize("node_type", ["trigger-plugin", "trigger-schedule", "trigger-webhook"]) + def test_trigger_workflow_rejects_manual_service_surfaces( + self, + invoke_from: InvokeFrom, + node_type: str, + mocker: MockerFixture, + ) -> None: + workflow = _make_workflow(node_types=(node_type,)) + mocker.patch.object(AppGenerateService, "_get_workflow", return_value=workflow) + generate = mocker.patch("services.app_generate_service.WorkflowAppGenerator.generate") + + with pytest.raises(TriggerWorkflowServiceModeUnavailableError): + AppGenerateService.generate( + app_model=_make_app(AppMode.WORKFLOW), + user=_make_user(), + args={"inputs": {}}, + invoke_from=invoke_from, + streaming=False, + session=MagicMock(), + ) + + generate.assert_not_called() + + def test_trigger_workflow_allows_trigger_execution(self, mocker: MockerFixture) -> None: + workflow = _make_workflow(node_types=("trigger-webhook",)) + mocker.patch.object(AppGenerateService, "_get_workflow", return_value=workflow) + generate = mocker.patch( + "services.app_generate_service.WorkflowAppGenerator.generate", + return_value={"result": "trigger"}, + ) + mocker.patch( + "services.app_generate_service.WorkflowAppGenerator.convert_to_event_stream", + side_effect=lambda value: value, + ) + + result = AppGenerateService.generate( + app_model=_make_app(AppMode.WORKFLOW), + user=_make_user(), + args={"inputs": {}}, + invoke_from=InvokeFrom.TRIGGER, streaming=False, session=MagicMock(), ) - generate.assert_not_called() + assert result == {"result": "trigger"} + generate.assert_called_once() - def test_trigger_workflow_allows_trigger_execution(self, mocker: MockerFixture) -> None: - workflow = _make_workflow(node_types=("trigger-webhook",)) - mocker.patch.object(AppGenerateService, "_get_workflow", return_value=workflow) - generate = mocker.patch( - "services.app_generate_service.WorkflowAppGenerator.generate", - return_value={"result": "trigger"}, - ) - mocker.patch( - "services.app_generate_service.WorkflowAppGenerator.convert_to_event_stream", - side_effect=lambda value: value, - ) + def test_specific_start_workflow_version_remains_runnable(self, mocker: MockerFixture) -> None: + workflow_id = str(uuid.uuid4()) + workflow = _make_workflow(workflow_id=workflow_id, node_types=("start",)) + get_workflow = mocker.patch.object(AppGenerateService, "_get_workflow", return_value=workflow) + mocker.patch( + "services.app_generate_service.WorkflowAppGenerator.generate", + return_value={"result": "version"}, + ) + mocker.patch( + "services.app_generate_service.WorkflowAppGenerator.convert_to_event_stream", + side_effect=lambda value: value, + ) + app = _make_app(AppMode.WORKFLOW) + session = MagicMock() - result = AppGenerateService.generate( - app_model=_make_app(AppMode.WORKFLOW), - user=_make_user(), - args={"inputs": {}}, - invoke_from=InvokeFrom.TRIGGER, - streaming=False, - session=MagicMock(), - ) - - assert result == {"result": "trigger"} - generate.assert_called_once() - - def test_specific_start_workflow_version_remains_runnable(self, mocker: MockerFixture) -> None: - workflow_id = str(uuid.uuid4()) - workflow = _make_workflow(workflow_id=workflow_id, node_types=("start",)) - get_workflow = mocker.patch.object(AppGenerateService, "_get_workflow", return_value=workflow) - mocker.patch( - "services.app_generate_service.WorkflowAppGenerator.generate", - return_value={"result": "version"}, - ) - mocker.patch( - "services.app_generate_service.WorkflowAppGenerator.convert_to_event_stream", - side_effect=lambda value: value, - ) - app = _make_app(AppMode.WORKFLOW) - session = MagicMock() - - result = AppGenerateService.generate( - app_model=app, - user=_make_user(), - args={"inputs": {}, "workflow_id": workflow_id}, - invoke_from=InvokeFrom.SERVICE_API, - streaming=False, - session=session, - ) - - assert result == {"result": "version"} - get_workflow.assert_called_once_with(app, InvokeFrom.SERVICE_API, workflow_id, session=session) - - # -- WORKFLOW streaming ------------------------------------------------- - def test_workflow_streaming(self, mocker: MockerFixture, config_overrides: Callable[..., None]): - config_overrides(PUBSUB_REDIS_CHANNEL_TYPE="streams") - workflow = _make_workflow() - mocker.patch.object(AppGenerateService, "_get_workflow", return_value=workflow) - mocker.patch( - "services.app_generate_service.AppExecutionParams.new", - return_value=MagicMock(workflow_run_id="wfr-2", model_dump_json=MagicMock(return_value="{}")), - ) - delay_spy = mocker.patch("services.app_generate_service.workflow_based_app_execution_task.delay") - retrieve_spy = mocker.patch( - "services.app_generate_service.MessageBasedAppGenerator.retrieve_events", - return_value=iter([]), - ) - mocker.patch( - "services.app_generate_service.WorkflowAppGenerator.convert_to_event_stream", - side_effect=lambda x: x, - ) - - AppGenerateService.generate( - app_model=_make_app(AppMode.WORKFLOW), - user=_make_user(), - args={"inputs": {}}, - invoke_from=InvokeFrom.SERVICE_API, - streaming=True, - session=self.session, - ) - retrieve_spy.assert_called_once() - # Dispatch is gated on subscribe; simulate the SSE layer entering the - # subscription, which is what actually invokes on_subscribe. - on_subscribe = retrieve_spy.call_args.kwargs["on_subscribe"] - on_subscribe() - delay_spy.assert_called_once() - - # -- Invalid mode ------------------------------------------------------- - def test_invalid_mode_raises(self, mocker: MockerFixture): - app = _make_app("invalid-mode") - with pytest.raises(ValueError, match="Invalid app mode"): - AppGenerateService.generate( + result = AppGenerateService.generate( app_model=app, user=_make_user(), - args={}, + args={"inputs": {}, "workflow_id": workflow_id}, invoke_from=InvokeFrom.SERVICE_API, streaming=False, - session=self.session, + session=session, ) + assert result == {"result": "version"} + get_workflow.assert_called_once_with(app, InvokeFrom.SERVICE_API, workflow_id, session=session) -# --------------------------------------------------------------------------- -# generate – billing / quota -# --------------------------------------------------------------------------- -class TestGenerateBilling(_RealSessionTest): - @pytest.fixture(autouse=True) - def _common(self, mocker: MockerFixture): - mocker.patch("services.app_generate_service.RateLimit", _DummyRateLimit) - mocker.patch( - "services.app_generate_service.rate_limit_context", - _noop_rate_limit_context, - ) - - def test_cloud_edition_consumes_quota(self, mocker: MockerFixture, config_overrides: Callable[..., None]): - config_overrides(DEPLOYMENT_EDITION=DeploymentEdition.CLOUD) - quota_charge = MagicMock() - reserve_mock = mocker.patch( - "services.app_generate_service.QuotaService.reserve", - return_value=quota_charge, - ) - mocker.patch( - "services.app_generate_service.CompletionAppGenerator.generate", - return_value={"ok": True}, - ) - mocker.patch( - "services.app_generate_service.CompletionAppGenerator.convert_to_event_stream", - side_effect=lambda x: x, - ) - - AppGenerateService.generate( - app_model=_make_app(AppMode.COMPLETION), - user=_make_user(), - args={"inputs": {}}, - invoke_from=InvokeFrom.SERVICE_API, - streaming=False, - session=self.session, - ) - reserve_mock.assert_called_once_with(QuotaType.WORKFLOW, "tenant-id") - quota_charge.commit.assert_called_once() - - def test_billing_quota_exceeded_raises_rate_limit_error( - self, mocker: MockerFixture, config_overrides: Callable[..., None] - ): - from services.errors.app import QuotaExceededError - from services.errors.llm import InvokeRateLimitError - - config_overrides(DEPLOYMENT_EDITION=DeploymentEdition.CLOUD) - mocker.patch( - "services.app_generate_service.QuotaService.reserve", - side_effect=QuotaExceededError(feature="workflow", tenant_id="t", required=1), - ) - - with pytest.raises(InvokeRateLimitError): - AppGenerateService.generate( - app_model=_make_app(AppMode.COMPLETION), - user=_make_user(), - args={"inputs": {}}, - invoke_from=InvokeFrom.SERVICE_API, - streaming=False, - session=self.session, + # -- WORKFLOW streaming ------------------------------------------------- + def test_workflow_streaming(self, mocker: MockerFixture, config_overrides: Callable[..., None]): + config_overrides(PUBSUB_REDIS_CHANNEL_TYPE="streams") + workflow = _make_workflow() + mocker.patch.object(AppGenerateService, "_get_workflow", return_value=workflow) + mocker.patch( + "services.app_generate_service.AppExecutionParams.new", + return_value=MagicMock(workflow_run_id="wfr-2", model_dump_json=MagicMock(return_value="{}")), + ) + delay_spy = mocker.patch("services.app_generate_service.workflow_based_app_execution_task.delay") + retrieve_spy = mocker.patch( + "services.app_generate_service.MessageBasedAppGenerator.retrieve_events", + return_value=iter([]), + ) + mocker.patch( + "services.app_generate_service.WorkflowAppGenerator.convert_to_event_stream", + side_effect=lambda x: x, ) - def test_exception_refunds_quota_and_exits_rate_limit( - self, mocker: MockerFixture, config_overrides: Callable[..., None] - ): - config_overrides(DEPLOYMENT_EDITION=DeploymentEdition.CLOUD) - quota_charge = MagicMock() - mocker.patch( - "services.app_generate_service.QuotaService.reserve", - return_value=quota_charge, - ) - mocker.patch( - "services.app_generate_service.CompletionAppGenerator.generate", - side_effect=RuntimeError("boom"), - ) - mocker.patch( - "services.app_generate_service.CompletionAppGenerator.convert_to_event_stream", - side_effect=lambda x: x, - ) - - with pytest.raises(RuntimeError, match="boom"): AppGenerateService.generate( - app_model=_make_app(AppMode.COMPLETION), - user=_make_user(), - args={"inputs": {}}, - invoke_from=InvokeFrom.SERVICE_API, - streaming=False, - session=self.session, - ) - quota_charge.refund.assert_called_once() - - def test_rate_limit_exit_called_in_finally_for_blocking( - self, mocker: MockerFixture, config_overrides: Callable[..., None] - ): - """For non-streaming (blocking) calls, rate_limit.exit should be called in finally.""" - config_overrides(DEPLOYMENT_EDITION=DeploymentEdition.COMMUNITY) - - exit_calls: list[str] = [] - - class _TrackingRateLimit(_DummyRateLimit): - def exit(self, request_id: str) -> None: - exit_calls.append(request_id) - - mocker.patch("services.app_generate_service.RateLimit", _TrackingRateLimit) - mocker.patch( - "services.app_generate_service.CompletionAppGenerator.generate", - return_value={"ok": True}, - ) - mocker.patch( - "services.app_generate_service.CompletionAppGenerator.convert_to_event_stream", - side_effect=lambda x: x, - ) - - AppGenerateService.generate( - app_model=_make_app(AppMode.COMPLETION), - user=_make_user(), - args={"inputs": {}}, - invoke_from=InvokeFrom.SERVICE_API, - streaming=False, - session=self.session, - ) - # exit is called in finally block for non-streaming - assert exit_calls == ["dummy-request-id"] - - def test_blocking_failure_exits_rate_limit_once(self, mocker: MockerFixture, config_overrides: Callable[..., None]): - config_overrides(DEPLOYMENT_EDITION=DeploymentEdition.CLOUD) - quota_charge = MagicMock() - mocker.patch( - "services.app_generate_service.QuotaService.reserve", - return_value=quota_charge, - ) - exit_calls: list[str] = [] - - class _TrackingRateLimit(_DummyRateLimit): - def exit(self, request_id: str) -> None: - exit_calls.append(request_id) - - mocker.patch("services.app_generate_service.RateLimit", _TrackingRateLimit) - mocker.patch( - "services.app_generate_service.CompletionAppGenerator.generate", - side_effect=RuntimeError("boom"), - ) - - with pytest.raises(RuntimeError, match="boom"): - AppGenerateService.generate( - app_model=_make_app(AppMode.COMPLETION), - user=_make_user(), - args={"inputs": {}}, - invoke_from=InvokeFrom.SERVICE_API, - streaming=False, - session=self.session, - ) - - quota_charge.refund.assert_called_once() - assert exit_calls == ["dummy-request-id"] - - def test_streaming_failure_exits_rate_limit_once( - self, mocker: MockerFixture, config_overrides: Callable[..., None] - ): - config_overrides(DEPLOYMENT_EDITION=DeploymentEdition.CLOUD) - quota_charge = MagicMock() - mocker.patch( - "services.app_generate_service.QuotaService.reserve", - return_value=quota_charge, - ) - exit_calls: list[str] = [] - - class _TrackingRateLimit(_DummyRateLimit): - def exit(self, request_id: str) -> None: - exit_calls.append(request_id) - - mocker.patch("services.app_generate_service.RateLimit", _TrackingRateLimit) - mocker.patch( - "services.app_generate_service.CompletionAppGenerator.generate", - side_effect=RuntimeError("boom"), - ) - - with pytest.raises(RuntimeError, match="boom"): - AppGenerateService.generate( - app_model=_make_app(AppMode.COMPLETION), + app_model=_make_app(AppMode.WORKFLOW), user=_make_user(), args={"inputs": {}}, invoke_from=InvokeFrom.SERVICE_API, streaming=True, session=self.session, ) + retrieve_spy.assert_called_once() + # Dispatch is gated on subscribe; simulate the SSE layer entering the + # subscription, which is what actually invokes on_subscribe. + on_subscribe = retrieve_spy.call_args.kwargs["on_subscribe"] + on_subscribe() + delay_spy.assert_called_once() - quota_charge.refund.assert_called_once() - assert exit_calls == ["dummy-request-id"] + # -- Invalid mode ------------------------------------------------------- + def test_invalid_mode_raises(self, mocker: MockerFixture): + app = _make_app("invalid-mode") + with pytest.raises(ValueError, match="Invalid app mode"): + AppGenerateService.generate( + app_model=app, + user=_make_user(), + args={}, + invoke_from=InvokeFrom.SERVICE_API, + streaming=False, + session=self.session, + ) - -# --------------------------------------------------------------------------- -# _get_workflow -# --------------------------------------------------------------------------- -class TestGetWorkflow(_RealSessionTest): - def test_debugger_fetches_draft(self, mocker: MockerFixture): - draft_wf = _make_workflow() - ws = MagicMock() - ws.get_draft_workflow.return_value = draft_wf - mocker.patch("services.app_generate_service.WorkflowService", return_value=ws) - - result = AppGenerateService._get_workflow( - _make_app(AppMode.WORKFLOW), InvokeFrom.DEBUGGER, session=self.session - ) - assert result is draft_wf - ws.get_draft_workflow.assert_called_once() - - def test_debugger_raises_when_no_draft(self, mocker: MockerFixture): - ws = MagicMock() - ws.get_draft_workflow.return_value = None - mocker.patch("services.app_generate_service.WorkflowService", return_value=ws) - - with pytest.raises(ValueError, match="Workflow not initialized"): - AppGenerateService._get_workflow(_make_app(AppMode.WORKFLOW), InvokeFrom.DEBUGGER, session=self.session) - - def test_non_debugger_fetches_published(self, mocker: MockerFixture): - pub_wf = _make_workflow() - ws = MagicMock() - ws.get_published_workflow.return_value = pub_wf - mocker.patch("services.app_generate_service.WorkflowService", return_value=ws) - - result = AppGenerateService._get_workflow( - _make_app(AppMode.WORKFLOW), InvokeFrom.SERVICE_API, session=self.session - ) - assert result is pub_wf - ws.get_published_workflow.assert_called_once() - - def test_non_debugger_raises_when_no_published(self, mocker: MockerFixture): - ws = MagicMock() - ws.get_published_workflow.return_value = None - mocker.patch("services.app_generate_service.WorkflowService", return_value=ws) - - with pytest.raises(ValueError, match="Workflow not published"): - AppGenerateService._get_workflow(_make_app(AppMode.WORKFLOW), InvokeFrom.SERVICE_API, session=self.session) - - def test_specific_workflow_id_valid_uuid(self, mocker: MockerFixture): - valid_uuid = str(uuid.uuid4()) - specific_wf = _make_workflow(workflow_id=valid_uuid) - ws = MagicMock() - ws.get_published_workflow_by_id.return_value = specific_wf - mocker.patch("services.app_generate_service.WorkflowService", return_value=ws) - - result = AppGenerateService._get_workflow( - _make_app(AppMode.WORKFLOW), - InvokeFrom.SERVICE_API, - workflow_id=valid_uuid, - session=self.session, - ) - assert result is specific_wf - ws.get_published_workflow_by_id.assert_called_once() - - def test_specific_workflow_id_invalid_uuid(self, mocker: MockerFixture): - ws = MagicMock() - mocker.patch("services.app_generate_service.WorkflowService", return_value=ws) - - with pytest.raises(WorkflowIdFormatError): - AppGenerateService._get_workflow( - _make_app(AppMode.WORKFLOW), - InvokeFrom.SERVICE_API, - workflow_id="not-a-uuid", - session=self.session, + # --------------------------------------------------------------------------- + # generate – billing / quota + # --------------------------------------------------------------------------- + class TestGenerateBilling(_RealSessionTest): + @pytest.fixture(autouse=True) + def _common(self, mocker: MockerFixture): + mocker.patch("services.app_generate_service.RateLimit", _DummyRateLimit) + mocker.patch( + "services.app_generate_service.rate_limit_context", + _noop_rate_limit_context, ) - def test_specific_workflow_id_not_found(self, mocker: MockerFixture): - valid_uuid = str(uuid.uuid4()) - ws = MagicMock() - ws.get_published_workflow_by_id.return_value = None - mocker.patch("services.app_generate_service.WorkflowService", return_value=ws) + def test_cloud_edition_consumes_quota(self, mocker: MockerFixture, config_overrides: Callable[..., None]): + config_overrides(DEPLOYMENT_EDITION=DeploymentEdition.CLOUD) + quota_charge = MagicMock() + reserve_mock = mocker.patch( + "services.app_generate_service.QuotaService.reserve", + return_value=quota_charge, + ) + mocker.patch( + "services.app_generate_service.CompletionAppGenerator.generate", + return_value={"ok": True}, + ) + mocker.patch( + "services.app_generate_service.CompletionAppGenerator.convert_to_event_stream", + side_effect=lambda x: x, + ) - with pytest.raises(WorkflowNotFoundError): - AppGenerateService._get_workflow( + AppGenerateService.generate( + app_model=_make_app(AppMode.COMPLETION), + user=_make_user(), + args={"inputs": {}}, + invoke_from=InvokeFrom.SERVICE_API, + streaming=False, + session=self.session, + ) + reserve_mock.assert_called_once_with(QuotaType.WORKFLOW, "tenant-id") + quota_charge.commit.assert_called_once() + + def test_billing_quota_exceeded_raises_rate_limit_error( + self, mocker: MockerFixture, config_overrides: Callable[..., None] + ): + from services.errors.app import QuotaExceededError + from services.errors.llm import InvokeRateLimitError + + config_overrides(DEPLOYMENT_EDITION=DeploymentEdition.CLOUD) + mocker.patch( + "services.app_generate_service.QuotaService.reserve", + side_effect=QuotaExceededError(feature="workflow", tenant_id="t", required=1), + ) + + with pytest.raises(InvokeRateLimitError): + AppGenerateService.generate( + app_model=_make_app(AppMode.COMPLETION), + user=_make_user(), + args={"inputs": {}}, + invoke_from=InvokeFrom.SERVICE_API, + streaming=False, + session=self.session, + ) + + def test_exception_refunds_quota_and_exits_rate_limit( + self, mocker: MockerFixture, config_overrides: Callable[..., None] + ): + config_overrides(DEPLOYMENT_EDITION=DeploymentEdition.CLOUD) + quota_charge = MagicMock() + mocker.patch( + "services.app_generate_service.QuotaService.reserve", + return_value=quota_charge, + ) + mocker.patch( + "services.app_generate_service.CompletionAppGenerator.generate", + side_effect=RuntimeError("boom"), + ) + mocker.patch( + "services.app_generate_service.CompletionAppGenerator.convert_to_event_stream", + side_effect=lambda x: x, + ) + + with pytest.raises(RuntimeError, match="boom"): + AppGenerateService.generate( + app_model=_make_app(AppMode.COMPLETION), + user=_make_user(), + args={"inputs": {}}, + invoke_from=InvokeFrom.SERVICE_API, + streaming=False, + session=self.session, + ) + quota_charge.refund.assert_called_once() + + def test_rate_limit_exit_called_in_finally_for_blocking( + self, mocker: MockerFixture, config_overrides: Callable[..., None] + ): + """For non-streaming (blocking) calls, rate_limit.exit should be called in finally.""" + config_overrides(DEPLOYMENT_EDITION=DeploymentEdition.COMMUNITY) + + exit_calls: list[str] = [] + + class _TrackingRateLimit(_DummyRateLimit): + def exit(self, request_id: str) -> None: + exit_calls.append(request_id) + + mocker.patch("services.app_generate_service.RateLimit", _TrackingRateLimit) + mocker.patch( + "services.app_generate_service.CompletionAppGenerator.generate", + return_value={"ok": True}, + ) + mocker.patch( + "services.app_generate_service.CompletionAppGenerator.convert_to_event_stream", + side_effect=lambda x: x, + ) + + AppGenerateService.generate( + app_model=_make_app(AppMode.COMPLETION), + user=_make_user(), + args={"inputs": {}}, + invoke_from=InvokeFrom.SERVICE_API, + streaming=False, + session=self.session, + ) + # exit is called in finally block for non-streaming + assert exit_calls == ["dummy-request-id"] + + def test_blocking_failure_exits_rate_limit_once( + self, mocker: MockerFixture, config_overrides: Callable[..., None] + ): + config_overrides(DEPLOYMENT_EDITION=DeploymentEdition.CLOUD) + quota_charge = MagicMock() + mocker.patch( + "services.app_generate_service.QuotaService.reserve", + return_value=quota_charge, + ) + exit_calls: list[str] = [] + + class _TrackingRateLimit(_DummyRateLimit): + def exit(self, request_id: str) -> None: + exit_calls.append(request_id) + + mocker.patch("services.app_generate_service.RateLimit", _TrackingRateLimit) + mocker.patch( + "services.app_generate_service.CompletionAppGenerator.generate", + side_effect=RuntimeError("boom"), + ) + + with pytest.raises(RuntimeError, match="boom"): + AppGenerateService.generate( + app_model=_make_app(AppMode.COMPLETION), + user=_make_user(), + args={"inputs": {}}, + invoke_from=InvokeFrom.SERVICE_API, + streaming=False, + session=self.session, + ) + + quota_charge.refund.assert_called_once() + assert exit_calls == ["dummy-request-id"] + + def test_streaming_failure_exits_rate_limit_once( + self, mocker: MockerFixture, config_overrides: Callable[..., None] + ): + config_overrides(DEPLOYMENT_EDITION=DeploymentEdition.CLOUD) + quota_charge = MagicMock() + mocker.patch( + "services.app_generate_service.QuotaService.reserve", + return_value=quota_charge, + ) + exit_calls: list[str] = [] + + class _TrackingRateLimit(_DummyRateLimit): + def exit(self, request_id: str) -> None: + exit_calls.append(request_id) + + mocker.patch("services.app_generate_service.RateLimit", _TrackingRateLimit) + mocker.patch( + "services.app_generate_service.CompletionAppGenerator.generate", + side_effect=RuntimeError("boom"), + ) + + with pytest.raises(RuntimeError, match="boom"): + AppGenerateService.generate( + app_model=_make_app(AppMode.COMPLETION), + user=_make_user(), + args={"inputs": {}}, + invoke_from=InvokeFrom.SERVICE_API, + streaming=True, + session=self.session, + ) + + quota_charge.refund.assert_called_once() + assert exit_calls == ["dummy-request-id"] + + # --------------------------------------------------------------------------- + # _get_workflow + # --------------------------------------------------------------------------- + class TestGetWorkflow(_RealSessionTest): + def test_debugger_fetches_draft(self, mocker: MockerFixture): + draft_wf = _make_workflow() + ws = MagicMock() + ws.get_draft_workflow.return_value = draft_wf + mocker.patch("services.app_generate_service.WorkflowService", return_value=ws) + + result = AppGenerateService._get_workflow( + _make_app(AppMode.WORKFLOW), InvokeFrom.DEBUGGER, session=self.session + ) + assert result is draft_wf + ws.get_draft_workflow.assert_called_once() + + def test_debugger_raises_when_no_draft(self, mocker: MockerFixture): + ws = MagicMock() + ws.get_draft_workflow.return_value = None + mocker.patch("services.app_generate_service.WorkflowService", return_value=ws) + + with pytest.raises(ValueError, match="Workflow not initialized"): + AppGenerateService._get_workflow(_make_app(AppMode.WORKFLOW), InvokeFrom.DEBUGGER, session=self.session) + + def test_non_debugger_fetches_published(self, mocker: MockerFixture): + pub_wf = _make_workflow() + ws = MagicMock() + ws.get_published_workflow.return_value = pub_wf + mocker.patch("services.app_generate_service.WorkflowService", return_value=ws) + + result = AppGenerateService._get_workflow( + _make_app(AppMode.WORKFLOW), InvokeFrom.SERVICE_API, session=self.session + ) + assert result is pub_wf + ws.get_published_workflow.assert_called_once() + + def test_non_debugger_raises_when_no_published(self, mocker: MockerFixture): + ws = MagicMock() + ws.get_published_workflow.return_value = None + mocker.patch("services.app_generate_service.WorkflowService", return_value=ws) + + with pytest.raises(ValueError, match="Workflow not published"): + AppGenerateService._get_workflow( + _make_app(AppMode.WORKFLOW), InvokeFrom.SERVICE_API, session=self.session + ) + + def test_specific_workflow_id_valid_uuid(self, mocker: MockerFixture): + valid_uuid = str(uuid.uuid4()) + specific_wf = _make_workflow(workflow_id=valid_uuid) + ws = MagicMock() + ws.get_published_workflow_by_id.return_value = specific_wf + mocker.patch("services.app_generate_service.WorkflowService", return_value=ws) + + result = AppGenerateService._get_workflow( _make_app(AppMode.WORKFLOW), InvokeFrom.SERVICE_API, workflow_id=valid_uuid, session=self.session, ) + assert result is specific_wf + ws.get_published_workflow_by_id.assert_called_once() + def test_specific_workflow_id_invalid_uuid(self, mocker: MockerFixture): + ws = MagicMock() + mocker.patch("services.app_generate_service.WorkflowService", return_value=ws) -# --------------------------------------------------------------------------- -# generate_single_iteration -# --------------------------------------------------------------------------- -class TestGenerateSingleIteration(_RealSessionTest): - def test_advanced_chat_mode(self, mocker: MockerFixture): - workflow = _make_workflow() - mocker.patch.object(AppGenerateService, "_get_workflow", return_value=workflow) - mocker.patch( - "services.app_generate_service.AdvancedChatAppGenerator.convert_to_event_stream", - side_effect=lambda x: x, - ) - iter_spy = mocker.patch( - "services.app_generate_service.AdvancedChatAppGenerator.single_iteration_generate", - return_value={"event": "iteration"}, - ) - app = _make_app(AppMode.ADVANCED_CHAT) - session = self.session - result = AppGenerateService.generate_single_iteration( - app_model=app, - user=_make_user(), - node_id="n1", - args={"k": "v"}, - session=session, - ) - iter_spy.assert_called_once() - assert iter_spy.call_args.kwargs["session"] is session - assert result == {"event": "iteration"} + with pytest.raises(WorkflowIdFormatError): + AppGenerateService._get_workflow( + _make_app(AppMode.WORKFLOW), + InvokeFrom.SERVICE_API, + workflow_id="not-a-uuid", + session=self.session, + ) - def test_workflow_mode(self, mocker: MockerFixture): - workflow = _make_workflow() - mocker.patch.object(AppGenerateService, "_get_workflow", return_value=workflow) - mocker.patch( - "services.app_generate_service.AdvancedChatAppGenerator.convert_to_event_stream", - side_effect=lambda x: x, - ) - iter_spy = mocker.patch( - "services.app_generate_service.WorkflowAppGenerator.single_iteration_generate", - return_value={"event": "wf-iteration"}, - ) - app = _make_app(AppMode.WORKFLOW) - session = self.session - result = AppGenerateService.generate_single_iteration( - app_model=app, - user=_make_user(), - node_id="n1", - args={"k": "v"}, - session=session, - ) - iter_spy.assert_called_once() - assert iter_spy.call_args.kwargs["session"] is session - assert result == {"event": "wf-iteration"} + def test_specific_workflow_id_not_found(self, mocker: MockerFixture): + valid_uuid = str(uuid.uuid4()) + ws = MagicMock() + ws.get_published_workflow_by_id.return_value = None + mocker.patch("services.app_generate_service.WorkflowService", return_value=ws) - def test_invalid_mode_raises(self, mocker: MockerFixture): - app = _make_app(AppMode.CHAT) - with pytest.raises(ValueError, match="Invalid app mode"): - AppGenerateService.generate_single_iteration( - app_model=app, user=_make_user(), node_id="n1", args={}, session=self.session + with pytest.raises(WorkflowNotFoundError): + AppGenerateService._get_workflow( + _make_app(AppMode.WORKFLOW), + InvokeFrom.SERVICE_API, + workflow_id=valid_uuid, + session=self.session, + ) + + # --------------------------------------------------------------------------- + # generate_single_iteration + # --------------------------------------------------------------------------- + class TestGenerateSingleIteration(_RealSessionTest): + def test_advanced_chat_mode(self, mocker: MockerFixture): + workflow = _make_workflow() + mocker.patch.object(AppGenerateService, "_get_workflow", return_value=workflow) + mocker.patch( + "services.app_generate_service.AdvancedChatAppGenerator.convert_to_event_stream", + side_effect=lambda x: x, ) - - -# --------------------------------------------------------------------------- -# generate_single_loop -# --------------------------------------------------------------------------- -class TestGenerateSingleLoop(_RealSessionTest): - def test_advanced_chat_mode(self, mocker: MockerFixture): - workflow = _make_workflow() - mocker.patch.object(AppGenerateService, "_get_workflow", return_value=workflow) - mocker.patch( - "services.app_generate_service.AdvancedChatAppGenerator.convert_to_event_stream", - side_effect=lambda x: x, - ) - loop_spy = mocker.patch( - "services.app_generate_service.AdvancedChatAppGenerator.single_loop_generate", - return_value={"event": "loop"}, - ) - app = _make_app(AppMode.ADVANCED_CHAT) - session = self.session - result = AppGenerateService.generate_single_loop( - app_model=app, - user=_make_user(), - node_id="n1", - args=MagicMock(), - session=session, - ) - loop_spy.assert_called_once() - assert loop_spy.call_args.kwargs["session"] is session - assert result == {"event": "loop"} - - def test_workflow_mode(self, mocker: MockerFixture): - workflow = _make_workflow() - mocker.patch.object(AppGenerateService, "_get_workflow", return_value=workflow) - mocker.patch( - "services.app_generate_service.AdvancedChatAppGenerator.convert_to_event_stream", - side_effect=lambda x: x, - ) - loop_spy = mocker.patch( - "services.app_generate_service.WorkflowAppGenerator.single_loop_generate", - return_value={"event": "wf-loop"}, - ) - app = _make_app(AppMode.WORKFLOW) - session = self.session - result = AppGenerateService.generate_single_loop( - app_model=app, - user=_make_user(), - node_id="n1", - args=MagicMock(), - session=session, - ) - loop_spy.assert_called_once() - assert loop_spy.call_args.kwargs["session"] is session - assert result == {"event": "wf-loop"} - - def test_invalid_mode_raises(self, mocker: MockerFixture): - app = _make_app(AppMode.COMPLETION) - with pytest.raises(ValueError, match="Invalid app mode"): - AppGenerateService.generate_single_loop( - app_model=app, user=_make_user(), node_id="n1", args=MagicMock(), session=self.session + iter_spy = mocker.patch( + "services.app_generate_service.AdvancedChatAppGenerator.single_iteration_generate", + return_value={"event": "iteration"}, ) + app = _make_app(AppMode.ADVANCED_CHAT) + session = self.session + result = AppGenerateService.generate_single_iteration( + app_model=app, + user=_make_user(), + node_id="n1", + args={"k": "v"}, + session=session, + ) + iter_spy.assert_called_once() + assert iter_spy.call_args.kwargs["session"] is session + assert result == {"event": "iteration"} + + def test_workflow_mode(self, mocker: MockerFixture): + workflow = _make_workflow() + mocker.patch.object(AppGenerateService, "_get_workflow", return_value=workflow) + mocker.patch( + "services.app_generate_service.AdvancedChatAppGenerator.convert_to_event_stream", + side_effect=lambda x: x, + ) + iter_spy = mocker.patch( + "services.app_generate_service.WorkflowAppGenerator.single_iteration_generate", + return_value={"event": "wf-iteration"}, + ) + app = _make_app(AppMode.WORKFLOW) + session = self.session + result = AppGenerateService.generate_single_iteration( + app_model=app, + user=_make_user(), + node_id="n1", + args={"k": "v"}, + session=session, + ) + iter_spy.assert_called_once() + assert iter_spy.call_args.kwargs["session"] is session + assert result == {"event": "wf-iteration"} + + def test_invalid_mode_raises(self, mocker: MockerFixture): + app = _make_app(AppMode.CHAT) + with pytest.raises(ValueError, match="Invalid app mode"): + AppGenerateService.generate_single_iteration( + app_model=app, user=_make_user(), node_id="n1", args={}, session=self.session + ) + + # --------------------------------------------------------------------------- + # generate_single_loop + # --------------------------------------------------------------------------- + class TestGenerateSingleLoop(_RealSessionTest): + def test_advanced_chat_mode(self, mocker: MockerFixture): + workflow = _make_workflow() + mocker.patch.object(AppGenerateService, "_get_workflow", return_value=workflow) + mocker.patch( + "services.app_generate_service.AdvancedChatAppGenerator.convert_to_event_stream", + side_effect=lambda x: x, + ) + loop_spy = mocker.patch( + "services.app_generate_service.AdvancedChatAppGenerator.single_loop_generate", + return_value={"event": "loop"}, + ) + app = _make_app(AppMode.ADVANCED_CHAT) + session = self.session + result = AppGenerateService.generate_single_loop( + app_model=app, + user=_make_user(), + node_id="n1", + args=MagicMock(), + session=session, + ) + loop_spy.assert_called_once() + assert loop_spy.call_args.kwargs["session"] is session + assert result == {"event": "loop"} + + def test_workflow_mode(self, mocker: MockerFixture): + workflow = _make_workflow() + mocker.patch.object(AppGenerateService, "_get_workflow", return_value=workflow) + mocker.patch( + "services.app_generate_service.AdvancedChatAppGenerator.convert_to_event_stream", + side_effect=lambda x: x, + ) + loop_spy = mocker.patch( + "services.app_generate_service.WorkflowAppGenerator.single_loop_generate", + return_value={"event": "wf-loop"}, + ) + app = _make_app(AppMode.WORKFLOW) + session = self.session + result = AppGenerateService.generate_single_loop( + app_model=app, + user=_make_user(), + node_id="n1", + args=MagicMock(), + session=session, + ) + loop_spy.assert_called_once() + assert loop_spy.call_args.kwargs["session"] is session + assert result == {"event": "wf-loop"} + + def test_invalid_mode_raises(self, mocker: MockerFixture): + app = _make_app(AppMode.COMPLETION) + with pytest.raises(ValueError, match="Invalid app mode"): + AppGenerateService.generate_single_loop( + app_model=app, user=_make_user(), node_id="n1", args=MagicMock(), session=self.session + ) + + # --------------------------------------------------------------------------- + # generate_more_like_this + # --------------------------------------------------------------------------- + class TestGenerateMoreLikeThis(_RealSessionTest): + def test_delegates_to_completion_generator(self, mocker: MockerFixture): + gen_spy = mocker.patch( + "services.app_generate_service.CompletionAppGenerator.generate_more_like_this", + return_value={"result": "similar"}, + ) + session = self.session + result = AppGenerateService.generate_more_like_this( + app_model=_make_app(AppMode.COMPLETION), + user=_make_user(), + message_id="msg-1", + invoke_from=InvokeFrom.SERVICE_API, + session=session, + streaming=True, + ) + assert result == {"result": "similar"} + gen_spy.assert_called_once() + assert gen_spy.call_args.kwargs["session"] is session + assert gen_spy.call_args.kwargs["stream"] is True -# --------------------------------------------------------------------------- -# generate_more_like_this -# --------------------------------------------------------------------------- -class TestGenerateMoreLikeThis(_RealSessionTest): - def test_delegates_to_completion_generator(self, mocker: MockerFixture): - gen_spy = mocker.patch( - "services.app_generate_service.CompletionAppGenerator.generate_more_like_this", - return_value={"result": "similar"}, - ) - session = self.session - result = AppGenerateService.generate_more_like_this( - app_model=_make_app(AppMode.COMPLETION), - user=_make_user(), - message_id="msg-1", - invoke_from=InvokeFrom.SERVICE_API, - session=session, - streaming=True, - ) - assert result == {"result": "similar"} - gen_spy.assert_called_once() - assert gen_spy.call_args.kwargs["session"] is session - assert gen_spy.call_args.kwargs["stream"] is True +# ----------------------------- +# Fakes for Redis Pub/Sub flow +# ----------------------------- +class _FakePubSub: + def __init__(self, store: dict[str, deque[bytes]]): + self._store = store + self._subs: set[str] = set() + self._closed = False + + def subscribe(self, topic: str) -> None: + self._subs.add(topic) + + def unsubscribe(self, topic: str) -> None: + self._subs.discard(topic) + + def close(self) -> None: + self._closed = True + + def get_message(self, ignore_subscribe_messages: bool = True, timeout: int | float | None = 1): + # simulate a non-blocking poll; return first available + if self._closed: + return None + for t in list(self._subs): + q = self._store.get(t) + if q and len(q) > 0: + payload = q.popleft() + return {"type": "message", "channel": t, "data": payload} + # no message + return None + + +class _FakeRedisClient: + def __init__(self, store: dict[str, deque[bytes]]): + self._store = store + + def pubsub(self): + return _FakePubSub(self._store) + + def publish(self, topic: str, payload: bytes) -> None: + self._store.setdefault(topic, deque()).append(payload) + + +# ------------------------------------ +# Fakes for Redis Streams (XADD/XREAD) +# ------------------------------------ +class _FakeStreams: + def __init__(self) -> None: + # key -> list[(id, {field: value})] + self._data: dict[str, list[tuple[str, dict]]] = defaultdict(list) + self._seq: dict[str, int] = defaultdict(int) + # Mirrors real Redis: "$" resolves once, to the tail of the stream at the moment + # the first xread naming it is issued -- entries added earlier are never seen. + self._dollar_snapshots: dict[str, int] = {} + + def xadd(self, key: str, fields: dict[str, Any], *, maxlen: int | None = None) -> str: + # maxlen is accepted for API compatibility with redis-py; ignored in this test double + self._seq[key] += 1 + eid = f"{self._seq[key]}-0" + self._data[key].append((eid, fields)) + return eid + + def expire(self, key: str, seconds: int) -> None: + # no-op for tests + return None + + def xrevrange(self, key: str, count: int | None = None): + entries = list(reversed(self._data.get(key, []))) + if count is not None: + entries = entries[:count] + return entries + + def xread(self, streams: dict[str, Any], block: int | None = None, count: int | None = None): + assert len(streams) == 1 + key, last_id = next(iter(streams.items())) + entries = self._data.get(key, []) + start = 0 + if last_id == "$": + start = self._dollar_snapshots.setdefault(key, len(entries)) + elif last_id != "0-0": + for i, (eid, _f) in enumerate(entries): + if eid == last_id: + start = i + 1 + break + if start >= len(entries): + return [] + end = len(entries) if count is None else min(len(entries), start + count) + return [(key, entries[start:end])] + + +def _publish_events(app_mode: AppMode, run_id: str, events: list[dict]): + # Publish events to the same topic used by MessageGenerator + topic = MessageGenerator.get_response_topic(app_mode, run_id) + for ev in events: + topic.publish(json.dumps(ev).encode()) + + +class TestAppGenerateServiceStreamingIntegrationCases: + @staticmethod + @pytest.fixture + def _patch_get_channel_streams(monkeypatch: pytest.MonkeyPatch): + from libs.broadcast_channel.redis.streams_channel import StreamsBroadcastChannel + + fake = _FakeStreams() + chan = StreamsBroadcastChannel(fake, retention_seconds=60) + + def _get_channel(): + return chan + + # Patch the source and the imported aliases used by MessageGenerator / MessageBasedAppGenerator + monkeypatch.setattr("extensions.ext_redis.get_pubsub_broadcast_channel", lambda: chan) + monkeypatch.setattr("core.app.apps.message_generator.get_pubsub_broadcast_channel", lambda: chan) + monkeypatch.setattr("core.app.apps.message_based_app_generator.get_pubsub_broadcast_channel", lambda: chan) + apply_config_overrides(monkeypatch, PUBSUB_REDIS_CHANNEL_TYPE="streams") + + @staticmethod + @pytest.fixture + def _patch_get_channel_pubsub(monkeypatch: pytest.MonkeyPatch): + from libs.broadcast_channel.redis.pubsub_channel import BroadcastChannel as RedisBroadcastChannel + + store: dict[str, deque[bytes]] = defaultdict(deque) + client = _FakeRedisClient(store) + chan = RedisBroadcastChannel(client) + + def _get_channel(): + return chan + + # Patch both the source and the imported alias used by MessageGenerator + monkeypatch.setattr("extensions.ext_redis.get_pubsub_broadcast_channel", lambda: chan) + monkeypatch.setattr("core.app.apps.message_generator.get_pubsub_broadcast_channel", lambda: chan) + apply_config_overrides(monkeypatch, PUBSUB_REDIS_CHANNEL_TYPE="pubsub") + + @staticmethod + @pytest.mark.usefixtures("_patch_get_channel_streams") + def test_streams_full_flow_prepublish_and_replay(): + """Regression test for dify#40948: `workflow_started` dropped for concurrent streaming + runs on the streams transport. + + `start_task` is invoked via `on_subscribe` after entering the subscription and publishes + both events synchronously, before the listener thread necessarily gets a chance to run its + first `xread`. `_StreamsSubscription` fixes its read boundary synchronously before spawning + that thread, so neither event is dropped. + """ + app_mode = AppMode.WORKFLOW + run_id = str(uuid.uuid4()) + + events = [{"event": "workflow_started"}, {"event": "workflow_finished"}] + + def start_task(): + _publish_events(app_mode, run_id, events) + + # MessageBasedAppGenerator is what AppGenerateService actually dispatches WORKFLOW / + # ADVANCED_CHAT streaming runs through in production. + on_subscribe = AppGenerateService._build_streaming_task_on_subscribe(start_task) + + gen = MessageBasedAppGenerator.retrieve_events(app_mode, run_id, idle_timeout=2.0, on_subscribe=on_subscribe) + + received = [] + for msg in gen: + if isinstance(msg, str): + # skip ping events + continue + received.append(msg) + if msg.get("event") == "workflow_finished": + break + + assert [m.get("event") for m in received] == ["workflow_started", "workflow_finished"] + + @staticmethod + @pytest.mark.usefixtures("_patch_get_channel_pubsub") + def test_pubsub_full_flow_start_on_subscribe_gated(monkeypatch: pytest.MonkeyPatch): + # Speed up any potential timer if it accidentally triggers + monkeypatch.setattr("services.app_generate_service.SSE_TASK_START_FALLBACK_MS", 50) + + app_mode = AppMode.WORKFLOW + run_id = str(uuid.uuid4()) + + published_order: list[str] = [] + + def start_task(): + # When called (on subscribe), publish both events + events = [{"event": "workflow_started"}, {"event": "workflow_finished"}] + _publish_events(app_mode, run_id, events) + published_order.extend([e["event"] for e in events]) + + on_subscribe = AppGenerateService._build_streaming_task_on_subscribe(start_task) + + # Producer not started yet; only when subscribe happens + assert published_order == [] + + gen = MessageGenerator.retrieve_events(app_mode, run_id, idle_timeout=2.0, on_subscribe=on_subscribe) + + received = [] + for msg in gen: + if isinstance(msg, str): + continue + received.append(msg) + if msg.get("event") == "workflow_finished": + break + + # Verify publish happened and consumer received in order + assert published_order == ["workflow_started", "workflow_finished"] + assert [m.get("event") for m in received] == ["workflow_started", "workflow_finished"] diff --git a/api/tests/unit_tests/services/test_app_generate_service_streaming_integration.py b/api/tests/unit_tests/services/test_app_generate_service_streaming_integration.py deleted file mode 100644 index 1595a1db952..00000000000 --- a/api/tests/unit_tests/services/test_app_generate_service_streaming_integration.py +++ /dev/null @@ -1,214 +0,0 @@ -import json -import uuid -from collections import defaultdict, deque -from typing import Any - -import pytest - -from core.app.apps.message_based_app_generator import MessageBasedAppGenerator -from core.app.apps.message_generator import MessageGenerator -from models.model import AppMode -from services.app_generate_service import AppGenerateService -from tests.unit_tests.config_override import apply_config_overrides - - -# ----------------------------- -# Fakes for Redis Pub/Sub flow -# ----------------------------- -class _FakePubSub: - def __init__(self, store: dict[str, deque[bytes]]): - self._store = store - self._subs: set[str] = set() - self._closed = False - - def subscribe(self, topic: str) -> None: - self._subs.add(topic) - - def unsubscribe(self, topic: str) -> None: - self._subs.discard(topic) - - def close(self) -> None: - self._closed = True - - def get_message(self, ignore_subscribe_messages: bool = True, timeout: int | float | None = 1): - # simulate a non-blocking poll; return first available - if self._closed: - return None - for t in list(self._subs): - q = self._store.get(t) - if q and len(q) > 0: - payload = q.popleft() - return {"type": "message", "channel": t, "data": payload} - # no message - return None - - -class _FakeRedisClient: - def __init__(self, store: dict[str, deque[bytes]]): - self._store = store - - def pubsub(self): - return _FakePubSub(self._store) - - def publish(self, topic: str, payload: bytes) -> None: - self._store.setdefault(topic, deque()).append(payload) - - -# ------------------------------------ -# Fakes for Redis Streams (XADD/XREAD) -# ------------------------------------ -class _FakeStreams: - def __init__(self) -> None: - # key -> list[(id, {field: value})] - self._data: dict[str, list[tuple[str, dict]]] = defaultdict(list) - self._seq: dict[str, int] = defaultdict(int) - # Mirrors real Redis: "$" resolves once, to the tail of the stream at the moment - # the first xread naming it is issued -- entries added earlier are never seen. - self._dollar_snapshots: dict[str, int] = {} - - def xadd(self, key: str, fields: dict[str, Any], *, maxlen: int | None = None) -> str: - # maxlen is accepted for API compatibility with redis-py; ignored in this test double - self._seq[key] += 1 - eid = f"{self._seq[key]}-0" - self._data[key].append((eid, fields)) - return eid - - def expire(self, key: str, seconds: int) -> None: - # no-op for tests - return None - - def xrevrange(self, key: str, count: int | None = None): - entries = list(reversed(self._data.get(key, []))) - if count is not None: - entries = entries[:count] - return entries - - def xread(self, streams: dict[str, Any], block: int | None = None, count: int | None = None): - assert len(streams) == 1 - key, last_id = next(iter(streams.items())) - entries = self._data.get(key, []) - start = 0 - if last_id == "$": - start = self._dollar_snapshots.setdefault(key, len(entries)) - elif last_id != "0-0": - for i, (eid, _f) in enumerate(entries): - if eid == last_id: - start = i + 1 - break - if start >= len(entries): - return [] - end = len(entries) if count is None else min(len(entries), start + count) - return [(key, entries[start:end])] - - -@pytest.fixture -def _patch_get_channel_streams(monkeypatch: pytest.MonkeyPatch): - from libs.broadcast_channel.redis.streams_channel import StreamsBroadcastChannel - - fake = _FakeStreams() - chan = StreamsBroadcastChannel(fake, retention_seconds=60) - - def _get_channel(): - return chan - - # Patch the source and the imported aliases used by MessageGenerator / MessageBasedAppGenerator - monkeypatch.setattr("extensions.ext_redis.get_pubsub_broadcast_channel", lambda: chan) - monkeypatch.setattr("core.app.apps.message_generator.get_pubsub_broadcast_channel", lambda: chan) - monkeypatch.setattr("core.app.apps.message_based_app_generator.get_pubsub_broadcast_channel", lambda: chan) - apply_config_overrides(monkeypatch, PUBSUB_REDIS_CHANNEL_TYPE="streams") - - -@pytest.fixture -def _patch_get_channel_pubsub(monkeypatch: pytest.MonkeyPatch): - from libs.broadcast_channel.redis.pubsub_channel import BroadcastChannel as RedisBroadcastChannel - - store: dict[str, deque[bytes]] = defaultdict(deque) - client = _FakeRedisClient(store) - chan = RedisBroadcastChannel(client) - - def _get_channel(): - return chan - - # Patch both the source and the imported alias used by MessageGenerator - monkeypatch.setattr("extensions.ext_redis.get_pubsub_broadcast_channel", lambda: chan) - monkeypatch.setattr("core.app.apps.message_generator.get_pubsub_broadcast_channel", lambda: chan) - apply_config_overrides(monkeypatch, PUBSUB_REDIS_CHANNEL_TYPE="pubsub") - - -def _publish_events(app_mode: AppMode, run_id: str, events: list[dict]): - # Publish events to the same topic used by MessageGenerator - topic = MessageGenerator.get_response_topic(app_mode, run_id) - for ev in events: - topic.publish(json.dumps(ev).encode()) - - -@pytest.mark.usefixtures("_patch_get_channel_streams") -def test_streams_full_flow_prepublish_and_replay(): - """Regression test for dify#40948: `workflow_started` dropped for concurrent streaming - runs on the streams transport. - - `start_task` is invoked via `on_subscribe` after entering the subscription and publishes - both events synchronously, before the listener thread necessarily gets a chance to run its - first `xread`. `_StreamsSubscription` fixes its read boundary synchronously before spawning - that thread, so neither event is dropped. - """ - app_mode = AppMode.WORKFLOW - run_id = str(uuid.uuid4()) - - events = [{"event": "workflow_started"}, {"event": "workflow_finished"}] - - def start_task(): - _publish_events(app_mode, run_id, events) - - # MessageBasedAppGenerator is what AppGenerateService actually dispatches WORKFLOW / - # ADVANCED_CHAT streaming runs through in production. - on_subscribe = AppGenerateService._build_streaming_task_on_subscribe(start_task) - - gen = MessageBasedAppGenerator.retrieve_events(app_mode, run_id, idle_timeout=2.0, on_subscribe=on_subscribe) - - received = [] - for msg in gen: - if isinstance(msg, str): - # skip ping events - continue - received.append(msg) - if msg.get("event") == "workflow_finished": - break - - assert [m.get("event") for m in received] == ["workflow_started", "workflow_finished"] - - -@pytest.mark.usefixtures("_patch_get_channel_pubsub") -def test_pubsub_full_flow_start_on_subscribe_gated(monkeypatch: pytest.MonkeyPatch): - # Speed up any potential timer if it accidentally triggers - monkeypatch.setattr("services.app_generate_service.SSE_TASK_START_FALLBACK_MS", 50) - - app_mode = AppMode.WORKFLOW - run_id = str(uuid.uuid4()) - - published_order: list[str] = [] - - def start_task(): - # When called (on subscribe), publish both events - events = [{"event": "workflow_started"}, {"event": "workflow_finished"}] - _publish_events(app_mode, run_id, events) - published_order.extend([e["event"] for e in events]) - - on_subscribe = AppGenerateService._build_streaming_task_on_subscribe(start_task) - - # Producer not started yet; only when subscribe happens - assert published_order == [] - - gen = MessageGenerator.retrieve_events(app_mode, run_id, idle_timeout=2.0, on_subscribe=on_subscribe) - - received = [] - for msg in gen: - if isinstance(msg, str): - continue - received.append(msg) - if msg.get("event") == "workflow_finished": - break - - # Verify publish happened and consumer received in order - assert published_order == ["workflow_started", "workflow_finished"] - assert [m.get("event") for m in received] == ["workflow_started", "workflow_finished"] diff --git a/api/tests/unit_tests/services/test_app_package_service.py b/api/tests/unit_tests/services/test_app_package_service.py index 5f418b2d571..2b7f72a6ec8 100644 --- a/api/tests/unit_tests/services/test_app_package_service.py +++ b/api/tests/unit_tests/services/test_app_package_service.py @@ -1,16 +1,554 @@ -"""Ordinary App packages preserve DSLs and enforce container integrity.""" +from __future__ import annotations import hashlib import io import zipfile -from collections.abc import Callable +from collections.abc import Callable, Generator +from types import SimpleNamespace +from typing import cast +from unittest.mock import MagicMock, Mock +from uuid import uuid4 import pytest import yaml +from sqlalchemy import Engine, select +from sqlalchemy.orm import Session +from sqlalchemy.pool import QueuePool -from services.agent.errors import InvalidRosterAgentPackageError, RosterAgentPackageTooLargeError -from services.app_package_service import AppPackageService +from models import Account, App, AppMode +from models.agent import ( + Agent, + AgentConfigSnapshot, + AgentScope, + AgentSource, + AgentStatus, + WorkflowAgentBindingType, + WorkflowAgentNodeBinding, +) +from models.agent_config_entities import AgentSoulConfig +from models.model import UploadFile +from models.skill import AgentSkillBinding, AgentSkillBindingSnapshot, Skill, SkillVersion, SkillVersionManifest +from models.tools import ToolFile +from models.workflow import Workflow, WorkflowType +from services.agent.dsl_entities import AgentPackage +from services.agent.errors import ( + InvalidRosterAgentPackageError, + RosterAgentPackageExportFailedError, + RosterAgentPackageTooLargeError, +) +from services.agent.package_resource_exporter import AgentPackageResourceExporter, _FileSource, _SkillSource +from services.agent.package_resource_importer import AgentPackageResourceImporter +from services.agent.roster_package_entities import PackageIcon, PreparedPackageArchive, RosterAgentPackageMember +from services.app_dsl_service import AppDslService, Import +from services.app_package_service import AppPackageManifest, AppPackageService from services.dsl_content import DSL_MAX_SIZE +from services.entities.dsl_entities import ImportStatus +from services.errors.base import NoPermissionError +from services.plugin.dependencies_analysis import DependenciesAnalysisService +from tests.helpers.unit.model_factories import make_account, make_app, make_tenant, make_upload_file, make_workflow + +"""Workflow packages restore node-owned Agents and assets across workspaces.""" + + +class _MemoryStorage: + def __init__(self) -> None: + self.files: dict[str, bytes] = {} + self.saved: list[str] = [] + + def load_stream(self, filename: str) -> Generator[bytes]: + yield self.files[filename] + + def save(self, filename: str, data: bytes) -> None: + self.files[filename] = data + self.saved.append(filename) + + +def _skill(name: str, marker: str) -> bytes: + output = io.BytesIO() + with zipfile.ZipFile(output, "w") as archive: + archive.writestr("SKILL.md", f"---\nname: {name}\ndescription: {marker}\n---\n\n# Instructions\n{marker}\n") + archive.writestr("scripts/run.py", f"print({marker!r})\n") + return output.getvalue() + + +def _tool_file(session: Session, storage: _MemoryStorage, key: str, payload: bytes) -> ToolFile: + row = ToolFile( + tenant_id="tenant-1", + user_id="account-1", + conversation_id=None, + file_key=key, + name=key.rsplit("/", 1)[-1], + mimetype="application/zip", + size=len(payload), + original_url=None, + ) + storage.files[key] = payload + session.add(row) + return row + + +def _source( + session: Session, storage: _MemoryStorage, *, mode: AppMode = AppMode.WORKFLOW, published: bool = False +) -> App: + app = make_app(mode=mode, name="Inline Workflow") + nodes = [ + { + "id": f"node-{index}", + "data": {"type": "agent", "version": "2", "title": f"Agent {index}", "agent_node_kind": "dify_agent"}, + } + for index in (1, 2) + ] + workflow = make_workflow( + workflow_id="workflow-1", + app_id=app.id, + graph={"nodes": nodes, "edges": []}, + workflow_type=WorkflowType.CHAT if mode == AppMode.ADVANCED_CHAT else WorkflowType.WORKFLOW, + version="2026-01-01 00:00:00" if published else Workflow.VERSION_DRAFT, + ) + session.add_all([app, workflow]) + for index in (1, 2): + skill_file = _tool_file( + session, storage, f"tools/skill-{index}.zip", _skill("research", f"Agent {index} skill") + ) + notes = _tool_file(session, storage, f"tools/notes-{index}.txt", f"Agent {index} notes".encode()) + upload = make_upload_file(file_id=f"upload-{index}", key=f"uploads/guide-{index}.txt", name="guide.txt") + storage.files[upload.key] = f"Agent {index} guide".encode() + agent = Agent( + id=f"source-agent-{index}", + tenant_id=app.tenant_id, + name=f"Inline {index}", + scope=AgentScope.WORKFLOW_ONLY, + source=AgentSource.WORKFLOW, + app_id=app.id, + workflow_id=workflow.id, + workflow_node_id=f"node-{index}", + status=AgentStatus.ACTIVE, + created_by="account-1", + ) + soul = AgentSoulConfig.model_validate( + { + "prompt": {"system_prompt": f"Bound snapshot {index}"}, + "config_skills": [{"name": "research", "file_id": skill_file.id}], + "config_files": [ + {"name": "notes.txt", "file_kind": "tool_file", "file_id": notes.id}, + {"name": "guide.txt", "file_kind": "upload_file", "file_id": upload.id}, + ], + } + ) + snapshot = AgentConfigSnapshot( + id=f"source-snapshot-{index}", + tenant_id=app.tenant_id, + agent_id=agent.id, + version=1, + config_snapshot=soul, + created_by="account-1", + ) + active = AgentConfigSnapshot( + id=f"active-snapshot-{index}", + tenant_id=app.tenant_id, + agent_id=agent.id, + version=2, + config_snapshot=AgentSoulConfig.model_validate({"prompt": {"system_prompt": "Unbound active snapshot"}}), + created_by="account-1", + ) + agent.active_config_snapshot_id = active.id + binding = WorkflowAgentNodeBinding( + tenant_id=app.tenant_id, + app_id=app.id, + workflow_id=workflow.id, + workflow_version=workflow.version, + node_id=f"node-{index}", + binding_type=WorkflowAgentBindingType.INLINE_AGENT, + agent_id=agent.id, + current_snapshot_id=snapshot.id, + node_job_config={"workflow_prompt": f"Job {index}"}, + created_by="account-1", + ) + session.add_all([upload, agent, snapshot, active, binding]) + if index == 1: + workspace_archive = _tool_file( + session, storage, "tools/workspace.zip", _skill("workspace", "Workspace skill") + ) + skill = Skill( + id="source-workspace-skill", + tenant_id=app.tenant_id, + name="workspace", + display_name="Workspace", + description="Workspace skill", + created_by="account-1", + updated_by="account-1", + ) + version = SkillVersion( + id="source-skill-version", + skill_id=skill.id, + version_number=1, + version_name="1.0", + publish_note="", + manifest=SkillVersionManifest( + name="workspace", display_name="Workspace", description="Workspace skill", files=[] + ), + archive_tool_file_id=workspace_archive.id, + archive_size=workspace_archive.size, + hash_code="workspace-hash", + published_by="account-1", + ) + skill.latest_published_version_id = version.id + session.add_all( + [ + skill, + version, + AgentSkillBindingSnapshot( + tenant_id=app.tenant_id, + agent_id=agent.id, + config_snapshot_id=snapshot.id, + skill_id=skill.id, + priority=0, + created_by="account-1", + ), + ] + ) + session.commit() + return app + + +def _account(tenant_id: str = "tenant-2") -> Account: + return make_account(tenant=make_tenant(tenant_id=tenant_id)) + + +def _import(session: Session, content: bytes, *, app_id: str | None = None, account: Account | None = None) -> Import: + prepared = AppPackageService().read_package(io.BytesIO(content)) + assert prepared is not None + with prepared: + result = AppDslService(session).import_app( + account=account or _account(), + import_mode="yaml-content", + yaml_content=prepared.dsl, + app_id=app_id, + package=prepared, + ) + if result.status == ImportStatus.FAILED: + session.rollback() + else: + session.commit() + return result + + +def _rewrite( + content: bytes, mutate: Callable[[AppPackageManifest, dict[str, object], dict[str, bytes]], None] +) -> bytes: + with zipfile.ZipFile(io.BytesIO(content)) as archive: + members = {name: archive.read(name) for name in archive.namelist()} + manifest = AppPackageManifest.model_validate(yaml.safe_load(members["manifest.yaml"])) + data = yaml.safe_load(members["app.yaml"]) + mutate(manifest, data, members) + members["app.yaml"] = yaml.safe_dump(data).encode() + manifest.apps[0].size = len(members["app.yaml"]) + manifest.apps[0].sha256 = hashlib.sha256(members["app.yaml"]).hexdigest() + members["manifest.yaml"] = yaml.safe_dump(manifest.model_dump(mode="json")).encode() + output = io.BytesIO() + with zipfile.ZipFile(output, "w") as archive: + for name, payload in members.items(): + archive.writestr(name, payload) + return output.getvalue() + + +class TestWorkflowPackageServiceCases: + @staticmethod + @pytest.fixture + def storage(monkeypatch: pytest.MonkeyPatch, sqlite_engine: Engine) -> _MemoryStorage: + backend = _MemoryStorage() + monkeypatch.setattr("services.workflow_service.db", SimpleNamespace(engine=sqlite_engine)) + monkeypatch.setattr( + "services.app_package_service.AgentPackageResourceExporter", + lambda: AgentPackageResourceExporter(storage_backend=backend), + ) + monkeypatch.setattr( + "services.app_package_service.AgentPackageResourceImporter", + lambda: AgentPackageResourceImporter(storage_backend=backend), + ) + monkeypatch.setattr(DependenciesAnalysisService, "generate_dependencies", lambda **_kwargs: []) + return backend + + @staticmethod + @pytest.mark.parametrize("mode", [AppMode.WORKFLOW, AppMode.ADVANCED_CHAT]) + @pytest.mark.parametrize("published", [False, True]) + def test_package_round_trip_restores_bound_agents_and_resources( + sqlite_session: Session, + storage: _MemoryStorage, + mode: AppMode, + published: bool, + ) -> None: + app = _source(sqlite_session, storage, mode=mode, published=published) + with AppPackageService().export_app(app_model=app, workflow_id="workflow-1" if published else None) as exported: + content = exported.archive.read() + with zipfile.ZipFile(io.BytesIO(content)) as archive: + manifest = yaml.safe_load(archive.read("manifest.yaml")) + assert all( + "audit" not in resource + for group in manifest["agent_resources"].values() + for resource in [*group["skills"], *group["files"]] + ) + prepared = AppPackageService().read_package(io.BytesIO(content)) + assert prepared is not None + with prepared: + assert len(prepared.agent_resources) == 2 + assert [p.soul.prompt.system_prompt for p in prepared.agents.values()] == [ + "Bound snapshot 1", + "Bound snapshot 2", + ] + assert all(not p.omitted_assets for p in prepared.agents.values()) + assert {r.name for group in prepared.agent_resources.values() for r in group.skills} == { + "research", + "workspace", + } + result = _import(sqlite_session, content) + assert result.status == ImportStatus.COMPLETED, result.error + assert result.app_mode == mode + bindings = sqlite_session.scalars( + select(WorkflowAgentNodeBinding) + .where(WorkflowAgentNodeBinding.app_id == result.app_id) + .order_by(WorkflowAgentNodeBinding.node_id) + ).all() + assert len(bindings) == 2 + for index, binding in enumerate(bindings, 1): + assert binding.tenant_id == "tenant-2" + assert binding.binding_type == WorkflowAgentBindingType.INLINE_AGENT + assert binding.node_job_config_dict["workflow_prompt"] == f"Job {index}" + agent = sqlite_session.get(Agent, binding.agent_id) + assert agent is not None + assert (agent.scope, agent.app_id, agent.workflow_id, agent.workflow_node_id) == ( + AgentScope.WORKFLOW_ONLY, + result.app_id, + binding.workflow_id, + binding.node_id, + ) + snapshot = sqlite_session.get(AgentConfigSnapshot, binding.current_snapshot_id) + assert snapshot is not None + soul = AgentSoulConfig.model_validate(snapshot.config_snapshot_dict) + assert soul.prompt.system_prompt == f"Bound snapshot {index}" + assert len(soul.config_skills) == (2 if index == 1 else 1) + for ref in soul.config_skills: + file = sqlite_session.get(ToolFile, ref.file_id) + assert file is not None + assert file.tenant_id == "tenant-2" + assert not ref.is_missing + with zipfile.ZipFile(io.BytesIO(storage.files[file.file_key])) as archive: + expected = f"Agent {index} skill" if ref.name == "research" else "Workspace skill" + assert expected.encode() in archive.read("SKILL.md") + for ref in soul.config_files: + row = sqlite_session.get(ToolFile if ref.file_kind == "tool_file" else UploadFile, ref.file_id) + assert row is not None + assert row.tenant_id == "tenant-2" + key = row.file_key if isinstance(row, ToolFile) else row.key + expected = f"Agent {index} notes" if ref.name == "notes.txt" else f"Agent {index} guide" + assert storage.files[key] == expected.encode() + if isinstance(row, UploadFile): + assert row.used is True + assert row.used_by == "account-1" + assert ( + sqlite_session.scalar(select(AgentSkillBinding).where(AgentSkillBinding.agent_id == agent.id)) is None + ) + + @staticmethod + def test_yaml_export_still_omits_resources(sqlite_session: Session, storage: _MemoryStorage) -> None: + app = _source(sqlite_session, storage) + data = yaml.safe_load(AppDslService.export_dsl(app, session=sqlite_session)) + for package in data["agent_packages"].values(): + assert package["omitted_assets"] + assert all(item["is_missing"] and item["file_id"] == "" for item in package["soul"]["config_files"]) + + @staticmethod + def test_export_rejects_cross_tenant_assets(sqlite_session: Session, storage: _MemoryStorage) -> None: + app = _source(sqlite_session, storage) + row = sqlite_session.get(UploadFile, "upload-1") + assert row is not None + row.tenant_id = "tenant-elsewhere" + sqlite_session.commit() + with pytest.raises(RosterAgentPackageExportFailedError, match="payload is unavailable"): + AppPackageService().export_app(app_model=app) + + @staticmethod + @pytest.mark.parametrize( + "mutation", ["unknown_agent", "cross_agent_file", "unreferenced_agent", "duplicate_path", "file_checksum"] + ) + def test_rejects_invalid_resource_ownership_before_upload( + sqlite_session: Session, + storage: _MemoryStorage, + mutation: str, + ) -> None: + app = _source(sqlite_session, storage) + with AppPackageService().export_app(app_model=app) as exported: + content = exported.archive.read() + + def mutate(manifest: AppPackageManifest, data: dict[str, object], members: dict[str, bytes]) -> None: + resources = manifest.agent_resources + if mutation == "unknown_agent": + resources["unknown"] = resources.pop("agent_1") + elif mutation == "cross_agent_file": + agents = cast(dict[str, object], data["agent_packages"]) + first = AgentPackage.model_validate(agents["agent_1"]) + second = AgentPackage.model_validate(agents["agent_2"]) + first.soul.config_files[0].file_id = second.soul.config_files[0].file_id + agents["agent_1"] = first.model_dump(mode="json") + elif mutation == "unreferenced_agent": + cast(dict[str, dict[str, list[object]]], data["workflow"])["graph"]["nodes"].pop() + elif mutation == "duplicate_path": + resources["agent_2"].files[0] = resources["agent_1"].files[0] + else: + resource = resources["agent_1"].files[0] + members[resource.path] = b"x" * resource.size + + with pytest.raises(InvalidRosterAgentPackageError): + AppPackageService().read_package(io.BytesIO(_rewrite(content, mutate))) + assert storage.saved == [] + + @staticmethod + def test_pending_import_retains_materialized_resources_and_damage_warnings( + sqlite_session: Session, + storage: _MemoryStorage, + monkeypatch: pytest.MonkeyPatch, + ) -> None: + app = _source(sqlite_session, storage) + with AppPackageService().export_app(app_model=app) as exported: + content = exported.archive.read() + + def mutate(manifest: AppPackageManifest, data: dict[str, object], members: dict[str, bytes]) -> None: + resources = manifest.agent_resources + data["version"] = "99.0.0" + resource = resources["agent_1"].skills[0] + members[resource.path] = b"x" * resource.size + + pending: dict[str, str] = {} + monkeypatch.setattr( + "services.app_dsl_service.redis_client.setex", lambda key, _ttl, value: pending.__setitem__(key, value) + ) + monkeypatch.setattr("services.app_dsl_service.redis_client.get", pending.get) + monkeypatch.setattr("services.app_dsl_service.redis_client.delete", pending.pop) + result = _import(sqlite_session, _rewrite(content, mutate)) + assert result.status == ImportStatus.PENDING + saved_count = len(storage.saved) + confirmed = AppDslService(sqlite_session).confirm_import(import_id=result.id, account=_account()) + assert confirmed.status == ImportStatus.COMPLETED_WITH_WARNINGS, confirmed.error + assert any( + warning.code == "agent_skill_missing" and warning.path.startswith("agent_packages.agent_1.") + for warning in confirmed.warnings + ) + assert len(storage.saved) == saved_count + assert not pending + + @staticmethod + def test_overwrite_permission_is_checked_before_resource_upload( + sqlite_session: Session, + storage: _MemoryStorage, + monkeypatch: pytest.MonkeyPatch, + ) -> None: + app = _source(sqlite_session, storage) + with AppPackageService().export_app(app_model=app) as exported: + content = exported.archive.read() + monkeypatch.setattr(AppDslService, "_load_app_for_overwrite", Mock(side_effect=NoPermissionError("Denied"))) + with pytest.raises(NoPermissionError): + _import(sqlite_session, content, app_id="protected-app") + assert not storage.saved + + @staticmethod + def test_overwrite_replaces_inline_agents_and_reports_previous_owners( + sqlite_session: Session, + storage: _MemoryStorage, + monkeypatch: pytest.MonkeyPatch, + ) -> None: + app = _source(sqlite_session, storage) + with AppPackageService().export_app(app_model=app) as exported: + content = exported.archive.read() + original = _import(sqlite_session, content) + assert original.status == ImportStatus.COMPLETED, original.error + old_agents = set( + sqlite_session.scalars( + select(WorkflowAgentNodeBinding.agent_id).where(WorkflowAgentNodeBinding.app_id == original.app_id) + ) + ) + sqlite_session.commit() + retirement = Mock() + monkeypatch.setattr("services.app_dsl_service.WorkflowAgentRetirementService.retire_unowned", retirement) + updated = _import(sqlite_session, content, app_id=original.app_id) + assert updated.status == ImportStatus.COMPLETED, updated.error + assert updated.app_id == original.app_id + new_agents = set( + sqlite_session.scalars( + select(WorkflowAgentNodeBinding.agent_id).where(WorkflowAgentNodeBinding.app_id == original.app_id) + ) + ) + assert len(new_agents) == len(old_agents) == 2 + assert not (new_agents & old_agents) + retirement.assert_called_once_with(tenant_id="tenant-2", agent_ids=old_agents, account_id="account-1") + + @staticmethod + @pytest.mark.parametrize("legacy_skill", [False, True]) + def test_resource_io_runs_without_database_transactions( + sqlite_session: Session, + storage: _MemoryStorage, + monkeypatch: pytest.MonkeyPatch, + legacy_skill: bool, + ) -> None: + app = _source(sqlite_session, storage) + if legacy_skill: + version = sqlite_session.get(SkillVersion, "source-skill-version") + assert version is not None + version.manifest = version.manifest.model_copy(update={"name": None}) + sqlite_session.commit() + engine = sqlite_session.get_bind() + assert isinstance(engine, Engine) + pool = engine.pool + assert isinstance(pool, QueuePool) + + def before_io() -> None: + assert not sqlite_session.in_transaction() + assert pool.checkedout() == 0 + + def load_legacy_skill(*, tenant_id: str, file_id: str) -> bytes: + assert tenant_id == "tenant-1" + assert file_id + before_io() + return storage.files["tools/workspace.zip"] + + monkeypatch.setattr( + "services.skill_management_service.SkillManagementService._load_tool_file_bytes", + staticmethod(load_legacy_skill), + ) + load_stream = storage.load_stream + save = storage.save + + def checked_load(filename: str) -> Generator[bytes]: + before_io() + yield from load_stream(filename) + + def checked_save(filename: str, data: bytes) -> None: + before_io() + save(filename, data) + + monkeypatch.setattr(storage, "load_stream", checked_load) + monkeypatch.setattr(storage, "save", checked_save) + with AppPackageService().export_app(app_model=app) as exported: + content = exported.archive.read() + result = _import(sqlite_session, content) + assert result.status == ImportStatus.COMPLETED, result.error + + @staticmethod + def test_missing_assets_remain_explicit_import_warnings(sqlite_session: Session, storage: _MemoryStorage) -> None: + app = _source(sqlite_session, storage) + snapshot = sqlite_session.get(AgentConfigSnapshot, "source-snapshot-1") + assert snapshot is not None + soul = AgentSoulConfig.model_validate(snapshot.config_snapshot_dict) + soul.config_files[0].is_missing = True + soul.config_files[0].file_id = "" + snapshot.config_snapshot = soul + sqlite_session.commit() + with AppPackageService().export_app(app_model=app) as exported: + result = _import(sqlite_session, exported.archive.read()) + assert result.status == ImportStatus.COMPLETED_WITH_WARNINGS, result.error + assert any(warning.code == "agent_file_omitted" for warning in result.warnings) + + +"""Ordinary App packages preserve DSLs and enforce container integrity.""" def _archive(dsl: str, **manifest_overrides: object) -> io.BytesIO: @@ -29,91 +567,287 @@ def _archive(dsl: str, **manifest_overrides: object) -> io.BytesIO: return source -@pytest.mark.parametrize("mode", ["workflow", "advanced-chat", "chat", "completion", "agent-chat"]) -@pytest.mark.parametrize("version", ["0.1.0", "99.0.0"]) -def test_round_trip_retains_exact_dsl_and_defers_version_handling(mode: str, version: str) -> None: - dsl = f"kind: app\nversion: {version}\napp:\n mode: {mode}\n# Preserve the DSL verbatim.\n" - service = AppPackageService() - with service.export(dsl=dsl, name="Example App") as package: - assert package.filename == "example-app.ifpkg" - prepared = service.read_package(package.archive) - assert prepared is not None - with prepared: - assert prepared.dsl == dsl - assert package.archive.tell() == 0 - assert package.archive.closed - - -@pytest.mark.parametrize("overrides", [{"format": "unknown"}, {"format_version": 2}, {"apps": []}]) -def test_rejects_unsupported_manifests(overrides: dict[str, object]) -> None: - with pytest.raises(InvalidRosterAgentPackageError, match="manifest"): - AppPackageService().read_package(_archive("kind: app\napp: {mode: workflow}\n", **overrides)) - - -@pytest.mark.parametrize( - "dsl", - [ - "[]", - "kind: app\napp: {mode: agent}", - "kind: app\napp: {mode: unknown}", - "kind: app\napp: {mode: workflow, mode: chat}", - "kind: app\napp: {mode: workflow}\nsite: not-a-map", - ], -) -def test_rejects_invalid_or_agent_dsl_in_ordinary_container(dsl: str) -> None: - with pytest.raises(InvalidRosterAgentPackageError): - AppPackageService().read_package(_archive(dsl)) - - -def test_rejects_checksum_mismatch() -> None: - dsl = "kind: app\napp: {mode: workflow}\n" - source = _archive(dsl, apps=[{"path": "app.yaml", "size": len(dsl), "sha256": "0" * 64}]) - with pytest.raises(InvalidRosterAgentPackageError, match="integrity"): - AppPackageService().read_package(source) - - -@pytest.mark.parametrize("path", ["extra.txt", "../escape", "APP.YAML"]) -def test_rejects_unlisted_unsafe_and_duplicate_members(path: str) -> None: - source = _archive("kind: app\napp: {mode: workflow}\n") - with zipfile.ZipFile(source, "a") as archive: - archive.writestr(path, "unexpected") - source.seek(0) - with pytest.raises(InvalidRosterAgentPackageError): - AppPackageService().read_package(source) - - -def test_rejects_oversized_container(config_overrides: Callable[..., None]) -> None: - source = _archive("kind: app\napp: {mode: workflow}\n") - config_overrides(AGENT_PACKAGE_MAX_BYTES=10) - with pytest.raises(RosterAgentPackageTooLargeError): - AppPackageService().read_package(source) - - def _dsl_with_size(size: int) -> str: header = "kind: app\napp:\n mode: workflow\n description: " padding = size - len(header.encode("utf-8")) return header + "\U0001f600" * (padding // 4) + "a" * (padding % 4) -@pytest.mark.parametrize("size", [6 * 1024 * 1024, DSL_MAX_SIZE]) -def test_package_round_trip_accepts_large_dsl_within_legacy_limit(size: int) -> None: - dsl = _dsl_with_size(size) - service = AppPackageService() - with service.export(dsl=dsl, name="Large Workflow") as package: - prepared = service.read_package(package.archive) - assert prepared is not None +class TestAppPackageServiceCases: + @staticmethod + @pytest.mark.parametrize("mode", ["workflow", "advanced-chat", "chat", "completion", "agent-chat"]) + @pytest.mark.parametrize("version", ["0.1.0", "99.0.0"]) + def test_round_trip_retains_exact_dsl_and_defers_version_handling(mode: str, version: str) -> None: + dsl = f"kind: app\nversion: {version}\napp:\n mode: {mode}\n# Preserve the DSL verbatim.\n" + service = AppPackageService() + with service.export(dsl=dsl, name="Example App") as package: + assert package.filename == "example-app.ifpkg" + prepared = service.read_package(package.archive) + assert prepared is not None + with prepared: + assert prepared.dsl == dsl + assert package.archive.tell() == 0 + assert package.archive.closed + + @staticmethod + @pytest.mark.parametrize("overrides", [{"format": "unknown"}, {"format_version": 2}, {"apps": []}]) + def test_rejects_unsupported_manifests(overrides: dict[str, object]) -> None: + with pytest.raises(InvalidRosterAgentPackageError, match="manifest"): + AppPackageService().read_package(_archive("kind: app\napp: {mode: workflow}\n", **overrides)) + + @staticmethod + @pytest.mark.parametrize( + "dsl", + [ + "[]", + "kind: app\napp: {mode: agent}", + "kind: app\napp: {mode: unknown}", + "kind: app\napp: {mode: workflow, mode: chat}", + "kind: app\napp: {mode: workflow}\nsite: not-a-map", + ], + ) + def test_rejects_invalid_or_agent_dsl_in_ordinary_container(dsl: str) -> None: + with pytest.raises(InvalidRosterAgentPackageError): + AppPackageService().read_package(_archive(dsl)) + + @staticmethod + def test_rejects_checksum_mismatch() -> None: + dsl = "kind: app\napp: {mode: workflow}\n" + source = _archive(dsl, apps=[{"path": "app.yaml", "size": len(dsl), "sha256": "0" * 64}]) + with pytest.raises(InvalidRosterAgentPackageError, match="integrity"): + AppPackageService().read_package(source) + + @staticmethod + @pytest.mark.parametrize("path", ["extra.txt", "../escape", "APP.YAML"]) + def test_rejects_unlisted_unsafe_and_duplicate_members(path: str) -> None: + source = _archive("kind: app\napp: {mode: workflow}\n") + with zipfile.ZipFile(source, "a") as archive: + archive.writestr(path, "unexpected") + source.seek(0) + with pytest.raises(InvalidRosterAgentPackageError): + AppPackageService().read_package(source) + + @staticmethod + def test_rejects_oversized_container(config_overrides: Callable[..., None]) -> None: + source = _archive("kind: app\napp: {mode: workflow}\n") + config_overrides(AGENT_PACKAGE_MAX_BYTES=10) + with pytest.raises(RosterAgentPackageTooLargeError): + AppPackageService().read_package(source) + + @staticmethod + @pytest.mark.parametrize("size", [6 * 1024 * 1024, DSL_MAX_SIZE]) + def test_package_round_trip_accepts_large_dsl_within_legacy_limit(size: int) -> None: + dsl = _dsl_with_size(size) + service = AppPackageService() + with service.export(dsl=dsl, name="Large Workflow") as package: + prepared = service.read_package(package.archive) + assert prepared is not None + with prepared: + assert prepared.dsl == dsl + + @staticmethod + def test_export_rejects_dsl_exceeding_legacy_byte_limit() -> None: + dsl = _dsl_with_size(DSL_MAX_SIZE + 1) + with pytest.raises(RosterAgentPackageTooLargeError): + AppPackageService().export(dsl=dsl, name="Large Workflow") + + @staticmethod + def test_import_rejects_compressed_dsl_exceeding_legacy_byte_limit() -> None: + dsl = _dsl_with_size(DSL_MAX_SIZE + 1) + source = _archive(dsl) + with pytest.raises(RosterAgentPackageTooLargeError): + AppPackageService().read_package(source) + + +"""Packaged image icons travel with Apps and receive destination-owned file ids.""" + + +def _icon_archive(*, payload: bytes = b"image", digest: str | None = None, reference: str = "i_000001") -> io.BytesIO: + dsl = yaml.safe_dump({"kind": "app", "app": {"mode": "chat", "icon_type": "image", "icon": reference}}) + with AppPackageService().export(dsl=dsl.replace(reference, "legacy-id"), name="Example") as exported: + with zipfile.ZipFile(exported.archive) as package: + manifest = yaml.safe_load(package.read("manifest.yaml")) + app_bytes = dsl.encode() + manifest["apps"][0].update(size=len(app_bytes), sha256=hashlib.sha256(app_bytes).hexdigest()) + manifest["icons"] = [ + { + "id": "i_000001", + "path": "i_000001.png", + "size": len(payload), + "sha256": digest or hashlib.sha256(payload).hexdigest(), + } + ] + output = io.BytesIO() + with zipfile.ZipFile(output, "w") as package: + package.writestr("manifest.yaml", yaml.safe_dump(manifest)) + package.writestr("app.yaml", app_bytes) + package.writestr("i_000001.png", payload) + output.seek(0) + return output + + +class TestPackageIconsCases: + @staticmethod + def test_app_icon_export_deduplicates_uploads_and_restores_references(monkeypatch: pytest.MonkeyPatch) -> None: + source_id = str(uuid4()) + payload = b"image payload" + backend = Mock() + backend.load_stream.side_effect = lambda _: iter([payload]) + exporter = AgentPackageResourceExporter(storage_backend=backend) + upload_lookup = Mock(return_value={source_id: SimpleNamespace(key="source/icon.png", extension="png")}) + monkeypatch.setattr(exporter, "_upload_files", upload_lookup) + metadata = {"mode": "chat", "icon_type": "image", "icon": source_id} + duplicate = dict(metadata) + exporter.collect_icon(session=Mock(), tenant_id="source-tenant", metadata=metadata) + exporter.collect_icon(session=Mock(), tenant_id="source-tenant", metadata=duplicate) + assert metadata["icon"] == duplicate["icon"] == "i_000001" + assert upload_lookup.call_args.kwargs["tenant_id"] == "source-tenant" + assert upload_lookup.call_count == 1 + dsl = yaml.safe_dump({"kind": "app", "app": metadata}) + with AppPackageService().export(dsl=dsl, name="Example", resources=exporter) as exported: + with zipfile.ZipFile(exported.archive) as package: + assert package.read("i_000001.png") == payload + exported.archive.seek(0) + prepared = AppPackageService().read_package(exported.archive) + assert prepared is not None + with prepared: + restore = Mock(return_value={"i_000001": "destination-id"}) + monkeypatch.setattr(AgentPackageResourceImporter, "materialize_icons", restore) + data = yaml.safe_load(prepared.dsl) + prepared.materialize_icons(data=data, tenant_id="destination", account_id="account") + assert data["app"]["icon"] == "destination-id" + assert restore.call_args.kwargs["tenant_id"] == "destination" + + @staticmethod + def test_app_package_restores_independent_site_image_icon(monkeypatch: pytest.MonkeyPatch) -> None: + app_icon_id, site_icon_id = str(uuid4()), str(uuid4()) + backend = Mock() + backend.load_stream.side_effect = lambda key: iter([key.encode()]) + exporter = AgentPackageResourceExporter(storage_backend=backend) + monkeypatch.setattr( + exporter, + "_upload_files", + Mock( + side_effect=lambda *, file_ids, **_kwargs: { + file_id: SimpleNamespace(key=f"source/{file_id}.png", extension="png") for file_id in file_ids + } + ), + ) + app = {"mode": "chat", "icon_type": "image", "icon": app_icon_id} + site = {"title": "Site", "icon_type": "image", "icon": site_icon_id} + exporter.collect_icon(session=Mock(), tenant_id="source", metadata=app) + exporter.collect_icon(session=Mock(), tenant_id="source", metadata=site) + with AppPackageService().export( + dsl=yaml.safe_dump({"kind": "app", "app": app, "site": site}), name="App", resources=exporter + ) as exported: + prepared = AppPackageService().read_package(exported.archive) + assert prepared is not None + with prepared: + assert len(prepared.icons) == 2 + monkeypatch.setattr( + AgentPackageResourceImporter, + "materialize_icons", + Mock(return_value={app["icon"]: "new-app-icon", site["icon"]: "new-site-icon"}), + ) + data = yaml.safe_load(prepared.dsl) + prepared.materialize_icons(data=data, tenant_id="destination", account_id="account") + assert data["app"]["icon"] == "new-app-icon" + assert data["site"]["icon"] == "new-site-icon" + + @staticmethod + def test_icon_export_rejects_unavailable_tenant_upload(monkeypatch: pytest.MonkeyPatch) -> None: + exporter = AgentPackageResourceExporter(storage_backend=Mock()) + monkeypatch.setattr(exporter, "_upload_files", Mock(return_value={})) + with pytest.raises(RosterAgentPackageExportFailedError, match="unavailable"): + exporter.collect_icon( + session=Mock(), tenant_id="tenant", metadata={"icon_type": "image", "icon": str(uuid4())} + ) + + @staticmethod + def test_icon_checksum_is_required() -> None: + with pytest.raises(InvalidRosterAgentPackageError, match="integrity"): + AppPackageService().read_package(_icon_archive(digest="0" * 64)) + + @staticmethod + @pytest.mark.parametrize("reference", ["i_000002", "unrelated-upload"]) + def test_icon_must_resolve_to_metadata(reference: str) -> None: + with pytest.raises(InvalidRosterAgentPackageError, match="references"): + AppPackageService().read_package(_icon_archive(reference=reference)) + + @staticmethod + def test_icon_image_limit_is_enforced(config_overrides: Callable[..., None]) -> None: + config_overrides(UPLOAD_IMAGE_FILE_SIZE_LIMIT=0) + with pytest.raises(RosterAgentPackageTooLargeError, match="image size"): + AppPackageService().read_package(_icon_archive()) + + @staticmethod + def test_materialized_icon_is_owned_by_destination(monkeypatch: pytest.MonkeyPatch) -> None: + payload = b"image payload" + digest = hashlib.sha256(payload).hexdigest() + archive = io.BytesIO() + with zipfile.ZipFile(archive, "w") as package: + package.writestr("i_000001.png", payload) + prepared = PreparedPackageArchive( + archive=archive, members={"i_000001.png": RosterAgentPackageMember(size=len(payload), sha256=digest)} + ) + icon = PackageIcon(id="i_000001", path="i_000001.png", size=len(payload), sha256=digest) + session = MagicMock() + monkeypatch.setattr( + "services.agent.package_resource_importer.session_factory.create_session", + MagicMock(return_value=session), + ) + backend = Mock() with prepared: - assert prepared.dsl == dsl + mapping = AgentPackageResourceImporter(storage_backend=backend).materialize_icons( + archive=prepared, icons=[icon], tenant_id="destination", account_id="importer" + ) + row = session.__enter__.return_value.add_all.call_args.args[0][0] + assert row.id == mapping["i_000001"] + assert row.tenant_id == "destination" + assert row.created_by == "importer" + assert row.mime_type == "image/png" + assert row.key.startswith("upload_files/destination/") + backend.save.assert_called_once_with(row.key, payload) + @staticmethod + def test_roster_export_embeds_agent_icon(monkeypatch: pytest.MonkeyPatch) -> None: + from models.agent_config_entities import AgentSoulConfig + from models.model import App, AppMode + from services.agent.dsl_entities import AgentPackage, AgentPackageMetadata, make_agent_app_dsl + from services.agent.roster_package_exporter import RosterAgentPackageExporter + from services.agent.roster_package_reader import RosterAgentPackageReader + from services.entities.site_dsl import SiteDsl -def test_export_rejects_dsl_exceeding_legacy_byte_limit() -> None: - dsl = _dsl_with_size(DSL_MAX_SIZE + 1) - with pytest.raises(RosterAgentPackageTooLargeError): - AppPackageService().export(dsl=dsl, name="Large Workflow") - - -def test_import_rejects_compressed_dsl_exceeding_legacy_byte_limit() -> None: - dsl = _dsl_with_size(DSL_MAX_SIZE + 1) - source = _archive(dsl) - with pytest.raises(RosterAgentPackageTooLargeError): - AppPackageService().read_package(source) + backend = Mock() + backend.load_stream.side_effect = lambda _: iter([b"icon"]) + resources = AgentPackageResourceExporter(storage_backend=backend) + resources.sources["agent_1"] = (list[_SkillSource](), list[_FileSource]()) + metadata = {"name": "Agent", "icon_type": "image", "icon": "source-id"} + monkeypatch.setattr( + resources, + "_upload_files", + Mock( + return_value={ + "source-id": SimpleNamespace(key="icon.png", extension="png"), + } + ), + ) + resources.collect_icon(session=Mock(), tenant_id="source", metadata=metadata) + app = make_agent_app_dsl( + App(name="Agent", mode=AppMode.AGENT), + package_ref="agent_1", + packages={ + "agent_1": AgentPackage(metadata=AgentPackageMetadata.model_validate(metadata), soul=AgentSoulConfig()) + }, + dependencies=[], + ) + app.site = SiteDsl(title="Agent Site", icon_type="image", icon="i_000001") + with RosterAgentPackageExporter(storage_backend=backend)._build_archive( + app=app, resources=resources + ) as exported: + with RosterAgentPackageReader().read(exported.archive) as prepared: + assert prepared.apps["app.yaml"].package.metadata.icon == "i_000001" + assert prepared.apps["app.yaml"].site is not None + assert prepared.apps["app.yaml"].site.icon == "i_000001" + assert prepared.manifest.icons[0].path == "i_000001.png" + assert RosterAgentPackageReader().read_member_bytes(prepared, "i_000001.png", max_bytes=4) == b"icon" diff --git a/api/tests/unit_tests/services/test_end_user_query_service.py b/api/tests/unit_tests/services/test_app_scoped_end_user_query_service.py similarity index 100% rename from api/tests/unit_tests/services/test_end_user_query_service.py rename to api/tests/unit_tests/services/test_app_scoped_end_user_query_service.py diff --git a/api/tests/unit_tests/services/test_end_user_service.py b/api/tests/unit_tests/services/test_app_scoped_end_user_service.py similarity index 100% rename from api/tests/unit_tests/services/test_end_user_service.py rename to api/tests/unit_tests/services/test_app_scoped_end_user_service.py diff --git a/api/tests/unit_tests/services/test_app_task_control_service.py b/api/tests/unit_tests/services/test_app_task_control_service.py deleted file mode 100644 index ead030da63b..00000000000 --- a/api/tests/unit_tests/services/test_app_task_control_service.py +++ /dev/null @@ -1,54 +0,0 @@ -"""Injected workflow cancellation preserves both transport mechanisms.""" - -from unittest.mock import MagicMock, patch - -import pytest -from redis.client import Pipeline - -from extensions.ext_redis import RedisClientWrapper -from graphon.engine.command import AbortCommand -from services.app_task_service import AppTaskControlService - - -@pytest.mark.parametrize("command_fails", [False, True]) -@pytest.mark.parametrize("task_id", ["task-1", ""]) -def test_workflow_stop_sets_flag_before_sending_command( - redis_transport: tuple[RedisClientWrapper, MagicMock], command_fails: bool, task_id: str -) -> None: - redis, commands = redis_transport - events: list[str] = [] - queued: list[tuple[object, ...]] = [] - commands.side_effect = lambda *_, **__: events.append("flag") - - def execute_pipeline(pipeline: "Pipeline[bytes]") -> list[object]: - events.append("command") - queued.extend(args for args, _ in pipeline.command_stack) - if command_fails: - raise RuntimeError("command channel unavailable") - return [] - - with patch.object(Pipeline, "execute", autospec=True, side_effect=execute_pipeline): - AppTaskControlService(redis_client=redis).stop_workflow_task_no_user_check(task_id=task_id) - - if task_id: - assert events == ["flag", "command"] - commands.assert_called_once_with("SETEX", "generate_task_stopped:task-1", 600, 1) - pushes = [args for args in queued if args[0] == "RPUSH"] - assert len(pushes) == 1 - _, key, payload = pushes[0] - assert key == "workflow:task-1:commands" - assert isinstance(payload, str) - assert AbortCommand.model_validate_json(payload) == AbortCommand(reason="User requested stop") - else: - assert events == [] - assert queued == [] - commands.assert_not_called() - - -def test_flag_failure_does_not_send_command(redis_transport: tuple[RedisClientWrapper, MagicMock]) -> None: - redis, commands = redis_transport - commands.side_effect = RuntimeError("flag store unavailable") - with patch.object(Pipeline, "execute", autospec=True) as execute: - with pytest.raises(RuntimeError, match="flag store unavailable"): - AppTaskControlService(redis_client=redis).stop_workflow_task_no_user_check(task_id="task-1") - execute.assert_not_called() diff --git a/api/tests/unit_tests/services/test_app_task_service.py b/api/tests/unit_tests/services/test_app_task_service.py index 23fbc074b4d..3082697761e 100644 --- a/api/tests/unit_tests/services/test_app_task_service.py +++ b/api/tests/unit_tests/services/test_app_task_service.py @@ -1,28 +1,99 @@ import json -from collections.abc import Generator +from collections.abc import Callable, Generator from contextlib import contextmanager from dataclasses import dataclass, field from datetime import datetime, timedelta from functools import partial from typing import override +from unittest.mock import MagicMock, patch +from uuid import UUID, uuid4 import pytest +from redis import Redis +from redis.client import Pipeline from redis.exceptions import ConnectionError as RedisConnectionError -from sqlalchemy.orm import Session +from sqlalchemy import event, select +from sqlalchemy.engine import Connection, Engine +from sqlalchemy.orm import Session, sessionmaker import core.app.apps.base_app_queue_manager as queue_module import core.app.apps.execution_coordinator as coordinator_module import services.app_task_service as task_module -from core.app.entities.app_invoke_entities import InvokeFrom +from core.app.app_config.entities import WorkflowUIBasedAppConfig +from core.app.entities.app_invoke_entities import AdvancedChatAppGenerateEntity, InvokeFrom, WorkflowAppGenerateEntity +from core.app.layers.pause_state_persist_layer import ( + WorkflowResumptionContext, + _AdvancedChatAppGenerateEntityWrapper, + _WorkflowGenerateEntityWrapper, +) +from core.repositories.sqlalchemy_workflow_node_execution_repository import SQLAlchemyWorkflowNodeExecutionRepository +from core.workflow.nodes.agent_v2.session_store import WorkflowAgentWorkspaceStore from core.workflow.nodes.human_input.enums import HumanInputFormStatus from extensions.ext_redis import RedisClientWrapper -from graphon.enums import WorkflowExecutionStatus +from extensions.logstore.aliyun_logstore import AliyunLogStore +from graphon.engine.command import AbortCommand +from graphon.entities import WorkflowNodeExecution +from graphon.enums import WorkflowExecutionStatus, WorkflowNodeExecutionStatus +from graphon.enums import WorkflowType as _paused_workflow_stop_WorkflowType +from libs.broadcast_channel.redis.pubsub_channel import Topic +from models import EndUser from models.enums import CreatorUserRole, WorkflowRunTriggeredFrom from models.human_input import HumanInputForm from models.model import AppMode -from models.workflow import WorkflowPause, WorkflowRun +from models.workflow import WorkflowNodeExecutionModel, WorkflowNodeExecutionTriggeredFrom, WorkflowPause, WorkflowRun +from models.workflow import WorkflowType as _legacy_workflow_task_recovery_WorkflowType +from repositories.sqlalchemy_api_workflow_run_repository import DifyAPISQLAlchemyWorkflowRunRepository from services.app_task_service import AppTaskControlService, AppTaskService +"""Injected workflow cancellation preserves both transport mechanisms.""" + + +class TestAppTaskControlServiceCases: + @staticmethod + @pytest.mark.parametrize("command_fails", [False, True]) + @pytest.mark.parametrize("task_id", ["task-1", ""]) + def test_workflow_stop_sets_flag_before_sending_command( + redis_transport: tuple[RedisClientWrapper, MagicMock], command_fails: bool, task_id: str + ) -> None: + redis, commands = redis_transport + events: list[str] = [] + queued: list[tuple[object, ...]] = [] + commands.side_effect = lambda *_, **__: events.append("flag") + + def execute_pipeline(pipeline: "Pipeline[bytes]") -> list[object]: + events.append("command") + queued.extend(args for args, _ in pipeline.command_stack) + if command_fails: + raise RuntimeError("command channel unavailable") + return [] + + with patch.object(Pipeline, "execute", autospec=True, side_effect=execute_pipeline): + AppTaskControlService(redis_client=redis).stop_workflow_task_no_user_check(task_id=task_id) + + if task_id: + assert events == ["flag", "command"] + commands.assert_called_once_with("SETEX", "generate_task_stopped:task-1", 600, 1) + pushes = [args for args in queued if args[0] == "RPUSH"] + assert len(pushes) == 1 + _, key, payload = pushes[0] + assert key == "workflow:task-1:commands" + assert isinstance(payload, str) + assert AbortCommand.model_validate_json(payload) == AbortCommand(reason="User requested stop") + else: + assert events == [] + assert queued == [] + commands.assert_not_called() + + @staticmethod + def test_flag_failure_does_not_send_command(redis_transport: tuple[RedisClientWrapper, MagicMock]) -> None: + redis, commands = redis_transport + commands.side_effect = RuntimeError("flag store unavailable") + with patch.object(Pipeline, "execute", autospec=True) as execute: + with pytest.raises(RuntimeError, match="flag store unavailable"): + AppTaskControlService(redis_client=redis).stop_workflow_task_no_user_check(task_id="task-1") + execute.assert_not_called() + + _TASK_ID = "task-with-non-uuid-id" _USER_ID = "user-1" _OWNER_KEY = f"generate_task_belong:{_TASK_ID}" @@ -31,7 +102,7 @@ _COMMAND_KEY = f"workflow:{_TASK_ID}:commands" @dataclass -class _StopRedis(RedisClientWrapper): +class _AppTaskServiceStopredis(RedisClientWrapper): """In-memory Redis boundary; task policy and GraphEngine serialization stay real.""" values: dict[str, bytes] = field(default_factory=dict) @@ -68,7 +139,7 @@ class _StopRedis(RedisClientWrapper): @dataclass class _StopPipeline: - redis: _StopRedis + redis: _AppTaskServiceStopredis values: dict[str, bytes] = field(default_factory=dict) commands: dict[str, list[str]] = field(default_factory=dict) expirations: dict[str, int] = field(default_factory=dict) @@ -94,28 +165,12 @@ class _StopPipeline: return [] -@pytest.fixture(autouse=True) -def global_redis(monkeypatch: pytest.MonkeyPatch) -> Generator[_StopRedis]: - redis = _StopRedis( - read_error=AssertionError("Must use the injected Redis for ownership reads"), - flag_error=AssertionError("Must use the injected Redis for stop flags"), - command_error=AssertionError("Must use the injected Redis for GraphEngine commands"), - ) - monkeypatch.setattr(queue_module, "redis_client", redis) - monkeypatch.setattr(coordinator_module, "redis_client", redis) - monkeypatch.setattr(task_module, "redis_client", redis) - yield redis - # GraphEngine catches Redis failures, so a trap exception alone would not fail the test. - assert redis.reads == [] - assert redis.operations == [] - - -def _assert_stop_flag(redis: _StopRedis) -> None: +def _assert_stop_flag(redis: _AppTaskServiceStopredis) -> None: assert redis.values[_STOP_KEY] == b"1" assert redis.expirations[_STOP_KEY] == 600 -def _assert_graph_command(redis: _StopRedis) -> None: +def _assert_graph_command(redis: _AppTaskServiceStopredis) -> None: assert set(redis.commands) == {_COMMAND_KEY} assert [json.loads(command) for command in redis.commands[_COMMAND_KEY]] == [ {"command_type": "abort", "reason": "User requested stop"} @@ -125,420 +180,1135 @@ def _assert_graph_command(redis: _StopRedis) -> None: assert redis.expirations[f"{_COMMAND_KEY}:pending"] == 3600 -def test_stop_task_sets_only_the_legacy_flag() -> None: - redis = _StopRedis(values={_OWNER_KEY: f"end-user-{_USER_ID}".encode()}) +class TestAppTaskServiceCases: + @staticmethod + @pytest.fixture(autouse=True) + def global_redis(monkeypatch: pytest.MonkeyPatch) -> Generator[_AppTaskServiceStopredis]: + redis = _AppTaskServiceStopredis( + read_error=AssertionError("Must use the injected Redis for ownership reads"), + flag_error=AssertionError("Must use the injected Redis for stop flags"), + command_error=AssertionError("Must use the injected Redis for GraphEngine commands"), + ) + monkeypatch.setattr(queue_module, "redis_client", redis) + monkeypatch.setattr(coordinator_module, "redis_client", redis) + monkeypatch.setattr(task_module, "redis_client", redis) + yield redis + # GraphEngine catches Redis failures, so a trap exception alone would not fail the test. + assert redis.reads == [] + assert redis.operations == [] - assert AppTaskControlService(redis_client=redis).stop_task(_TASK_ID, InvokeFrom.WEB_APP, _USER_ID) is True + @staticmethod + def test_stop_task_sets_only_the_legacy_flag() -> None: + redis = _AppTaskServiceStopredis(values={_OWNER_KEY: f"end-user-{_USER_ID}".encode()}) - assert redis.reads == [_OWNER_KEY] - _assert_stop_flag(redis) - assert redis.operations == ["legacy_flag"] - assert redis.commands == {} + assert AppTaskControlService(redis_client=redis).stop_task(_TASK_ID, InvokeFrom.WEB_APP, _USER_ID) is True - -@pytest.mark.parametrize( - ("invoke_from", "owner", "should_set_flag"), - [ - (InvokeFrom.EXPLORE, b"account-user-1", True), - (InvokeFrom.DEBUGGER, b"account-user-1", True), - (InvokeFrom.WEB_APP, b"end-user-user-1", True), - (InvokeFrom.SERVICE_API, b"end-user-user-1", True), - (InvokeFrom.OPENAPI, b"end-user-user-1", True), - (InvokeFrom.EXPLORE, b"end-user-user-1", False), - (InvokeFrom.DEBUGGER, b"end-user-user-1", False), - (InvokeFrom.WEB_APP, b"account-user-1", False), - (InvokeFrom.SERVICE_API, b"account-user-1", False), - (InvokeFrom.OPENAPI, b"account-user-1", False), - (InvokeFrom.EXPLORE, b"account-another-user", False), - (InvokeFrom.WEB_APP, b"end-user-another-user", False), - (InvokeFrom.OPENAPI, b"end-user-another-user", False), - (InvokeFrom.EXPLORE, None, False), - (InvokeFrom.WEB_APP, None, False), - (InvokeFrom.OPENAPI, None, False), - ], -) -def test_task_ownership_controls_the_legacy_flag( - invoke_from: InvokeFrom, owner: bytes | None, should_set_flag: bool -) -> None: - redis = _StopRedis(values={_OWNER_KEY: owner} if owner is not None else {}) - - admitted = AppTaskControlService(redis_client=redis).stop_task(_TASK_ID, invoke_from, _USER_ID) - - assert admitted is should_set_flag - assert redis.reads == [_OWNER_KEY] - if should_set_flag: + assert redis.reads == [_OWNER_KEY] _assert_stop_flag(redis) assert redis.operations == ["legacy_flag"] - else: - assert _STOP_KEY not in redis.values - assert redis.operations == [] - assert redis.commands == {} + assert redis.commands == {} - -@pytest.mark.parametrize("app_mode", [AppMode.CHAT, AppMode.COMPLETION]) -@pytest.mark.parametrize( - ("caller_role", "cached_owner", "admitted"), - [ - (CreatorUserRole.END_USER, b"end-user-user-1", True), - (CreatorUserRole.END_USER, b"end-user-other", False), - (CreatorUserRole.END_USER, b"account-user-1", False), - (CreatorUserRole.END_USER, None, False), - (CreatorUserRole.ACCOUNT, b"end-user-other", True), - (CreatorUserRole.ACCOUNT, None, True), - ], -) -def test_stop_app_task_applies_legacy_caller_ownership( - app_mode: AppMode, caller_role: CreatorUserRole, cached_owner: bytes | None, admitted: bool -) -> None: - redis = _StopRedis(values={_OWNER_KEY: cached_owner} if cached_owner is not None else {}) - - result = AppTaskControlService(redis_client=redis).stop_app_task( - tenant_id="tenant-1", - app_id="app-1", - task_id=_TASK_ID, - app_mode=app_mode, - caller_role=caller_role, - caller_id=_USER_ID, + @staticmethod + @pytest.mark.parametrize( + ("invoke_from", "owner", "should_set_flag"), + [ + (InvokeFrom.EXPLORE, b"account-user-1", True), + (InvokeFrom.DEBUGGER, b"account-user-1", True), + (InvokeFrom.WEB_APP, b"end-user-user-1", True), + (InvokeFrom.SERVICE_API, b"end-user-user-1", True), + (InvokeFrom.OPENAPI, b"end-user-user-1", True), + (InvokeFrom.EXPLORE, b"end-user-user-1", False), + (InvokeFrom.DEBUGGER, b"end-user-user-1", False), + (InvokeFrom.WEB_APP, b"account-user-1", False), + (InvokeFrom.SERVICE_API, b"account-user-1", False), + (InvokeFrom.OPENAPI, b"account-user-1", False), + (InvokeFrom.EXPLORE, b"account-another-user", False), + (InvokeFrom.WEB_APP, b"end-user-another-user", False), + (InvokeFrom.OPENAPI, b"end-user-another-user", False), + (InvokeFrom.EXPLORE, None, False), + (InvokeFrom.WEB_APP, None, False), + (InvokeFrom.OPENAPI, None, False), + ], ) + def test_task_ownership_controls_the_legacy_flag( + invoke_from: InvokeFrom, owner: bytes | None, should_set_flag: bool + ) -> None: + redis = _AppTaskServiceStopredis(values={_OWNER_KEY: owner} if owner is not None else {}) - assert result is admitted - assert redis.reads == ([_OWNER_KEY] if caller_role is CreatorUserRole.END_USER else []) - assert redis.operations == (["legacy_flag"] if admitted else []) - assert redis.commands == {} - if admitted: - _assert_stop_flag(redis) - else: - assert _STOP_KEY not in redis.values + admitted = AppTaskControlService(redis_client=redis).stop_task(_TASK_ID, invoke_from, _USER_ID) + assert admitted is should_set_flag + assert redis.reads == [_OWNER_KEY] + if should_set_flag: + _assert_stop_flag(redis) + assert redis.operations == ["legacy_flag"] + else: + assert _STOP_KEY not in redis.values + assert redis.operations == [] + assert redis.commands == {} -@pytest.fixture -def app_task_workflow_run(sqlite_session: Session, monkeypatch: pytest.MonkeyPatch) -> WorkflowRun: - run = WorkflowRun( - id="run-1", - tenant_id="tenant-1", - app_id="app-1", - workflow_id="workflow-1", - task_id="11111111-1111-1111-1111-111111111111", - type="workflow", - version="1", - triggered_from=WorkflowRunTriggeredFrom.APP_RUN, - status=WorkflowExecutionStatus.RUNNING, - created_by_role=CreatorUserRole.END_USER, - created_by=_USER_ID, + @staticmethod + @pytest.mark.parametrize("app_mode", [AppMode.CHAT, AppMode.COMPLETION]) + @pytest.mark.parametrize( + ("caller_role", "cached_owner", "admitted"), + [ + (CreatorUserRole.END_USER, b"end-user-user-1", True), + (CreatorUserRole.END_USER, b"end-user-other", False), + (CreatorUserRole.END_USER, b"account-user-1", False), + (CreatorUserRole.END_USER, None, False), + (CreatorUserRole.ACCOUNT, b"end-user-other", True), + (CreatorUserRole.ACCOUNT, None, True), + ], ) - sqlite_session.add(run) - monkeypatch.setattr(type(task_module.db), "engine", property(lambda _: sqlite_session.get_bind())) - return run + def test_stop_app_task_applies_legacy_caller_ownership( + app_mode: AppMode, caller_role: CreatorUserRole, cached_owner: bytes | None, admitted: bool + ) -> None: + redis = _AppTaskServiceStopredis(values={_OWNER_KEY: cached_owner} if cached_owner is not None else {}) + result = AppTaskControlService(redis_client=redis).stop_app_task( + tenant_id="tenant-1", + app_id="app-1", + task_id=_TASK_ID, + app_mode=app_mode, + caller_role=caller_role, + caller_id=_USER_ID, + ) -@pytest.mark.parametrize("app_mode", [AppMode.WORKFLOW, AppMode.ADVANCED_CHAT]) -@pytest.mark.parametrize( - ("caller_role", "run_owner_role", "run_owner_id"), - [ - (CreatorUserRole.END_USER, CreatorUserRole.END_USER, _USER_ID), - (CreatorUserRole.ACCOUNT, CreatorUserRole.END_USER, "other-user"), - (CreatorUserRole.ACCOUNT, CreatorUserRole.ACCOUNT, "other-account"), - ], -) -def test_stop_app_task_signals_admitted_running_workflow( - sqlite_session: Session, - app_task_workflow_run: WorkflowRun, - app_mode: AppMode, - caller_role: CreatorUserRole, - run_owner_role: CreatorUserRole, - run_owner_id: str, -) -> None: - run = app_task_workflow_run - run.created_by_role = run_owner_role - run.created_by = run_owner_id - sqlite_session.commit() - task_id = run.task_id - assert task_id is not None - redis = _StopRedis(read_error=AssertionError("Durable workflow ownership must not depend on Redis")) + assert result is admitted + assert redis.reads == ([_OWNER_KEY] if caller_role is CreatorUserRole.END_USER else []) + assert redis.operations == (["legacy_flag"] if admitted else []) + assert redis.commands == {} + if admitted: + _assert_stop_flag(redis) + else: + assert _STOP_KEY not in redis.values - admitted = AppTaskControlService(redis_client=redis).stop_app_task( - tenant_id=run.tenant_id, - app_id=run.app_id, - task_id=task_id, - app_mode=app_mode, - caller_role=caller_role, - caller_id=_USER_ID, + @staticmethod + @pytest.fixture + def app_task_workflow_run(sqlite_session: Session, monkeypatch: pytest.MonkeyPatch) -> WorkflowRun: + run = WorkflowRun( + id="run-1", + tenant_id="tenant-1", + app_id="app-1", + workflow_id="workflow-1", + task_id="11111111-1111-1111-1111-111111111111", + type="workflow", + version="1", + triggered_from=WorkflowRunTriggeredFrom.APP_RUN, + status=WorkflowExecutionStatus.RUNNING, + created_by_role=CreatorUserRole.END_USER, + created_by=_USER_ID, + ) + sqlite_session.add(run) + monkeypatch.setattr(type(task_module.db), "engine", property(lambda _: sqlite_session.get_bind())) + return run + + @staticmethod + @pytest.mark.parametrize("app_mode", [AppMode.WORKFLOW, AppMode.ADVANCED_CHAT]) + @pytest.mark.parametrize( + ("caller_role", "run_owner_role", "run_owner_id"), + [ + (CreatorUserRole.END_USER, CreatorUserRole.END_USER, _USER_ID), + (CreatorUserRole.ACCOUNT, CreatorUserRole.END_USER, "other-user"), + (CreatorUserRole.ACCOUNT, CreatorUserRole.ACCOUNT, "other-account"), + ], ) + def test_stop_app_task_signals_admitted_running_workflow( + sqlite_session: Session, + app_task_workflow_run: WorkflowRun, + app_mode: AppMode, + caller_role: CreatorUserRole, + run_owner_role: CreatorUserRole, + run_owner_id: str, + ) -> None: + run = app_task_workflow_run + run.created_by_role = run_owner_role + run.created_by = run_owner_id + sqlite_session.commit() + task_id = run.task_id + assert task_id is not None + redis = _AppTaskServiceStopredis( + read_error=AssertionError("Durable workflow ownership must not depend on Redis") + ) - assert admitted is True - sqlite_session.refresh(run) - assert run.status == WorkflowExecutionStatus.RUNNING - assert run.stop_requested_at is not None - assert redis.reads == [] - assert redis.operations == ["legacy_flag", "graph_command"] - assert redis.values[f"generate_task_stopped:{task_id}"] == b"1" - assert [json.loads(command) for command in redis.commands[f"workflow:{task_id}:commands"]] == [ - {"command_type": "abort", "reason": "User requested stop"} - ] - - -@pytest.mark.parametrize("app_mode", [AppMode.WORKFLOW, AppMode.ADVANCED_CHAT]) -@pytest.mark.parametrize("status", [WorkflowExecutionStatus.RUNNING, WorkflowExecutionStatus.PAUSED]) -@pytest.mark.parametrize( - ("caller_role", "mismatch"), - [ - (CreatorUserRole.END_USER, "owner-id"), - (CreatorUserRole.END_USER, "owner-role"), - (CreatorUserRole.END_USER, "app"), - (CreatorUserRole.END_USER, "tenant"), - (CreatorUserRole.END_USER, "missing"), - (CreatorUserRole.ACCOUNT, "app"), - (CreatorUserRole.ACCOUNT, "tenant"), - (CreatorUserRole.ACCOUNT, "missing"), - ], -) -def test_stop_app_task_preserves_unadmitted_workflow_and_pause( - sqlite_session: Session, - app_task_workflow_run: WorkflowRun, - app_mode: AppMode, - status: WorkflowExecutionStatus, - caller_role: CreatorUserRole, - mismatch: str, -) -> None: - run = app_task_workflow_run - run.status = status - if mismatch == "owner-id": - run.created_by = "other-user" - if mismatch == "owner-role": - run.created_by_role = CreatorUserRole.ACCOUNT - pause = None - form = None - if status == WorkflowExecutionStatus.PAUSED: - pause = WorkflowPause(workflow_id=run.workflow_id, workflow_run_id=run.id, state_object_key="state") - form = HumanInputForm( - id="form-1", + admitted = AppTaskControlService(redis_client=redis).stop_app_task( tenant_id=run.tenant_id, app_id=run.app_id, - workflow_run_id=run.id, - node_id="review", - form_definition="{}", - rendered_content="Review", - expiration_time=datetime(2026, 10, 20), + task_id=task_id, + app_mode=app_mode, + caller_role=caller_role, + caller_id=_USER_ID, ) - sqlite_session.add_all([pause, form]) - sqlite_session.commit() - task_id = "22222222-2222-2222-2222-222222222222" if mismatch == "missing" else run.task_id - assert task_id is not None - redis = _StopRedis(read_error=AssertionError("Durable workflow ownership must not depend on Redis")) - admitted = AppTaskControlService(redis_client=redis).stop_app_task( - tenant_id="other-tenant" if mismatch == "tenant" else run.tenant_id, - app_id="other-app" if mismatch == "app" else run.app_id, - task_id=task_id, - app_mode=app_mode, - caller_role=caller_role, - caller_id=_USER_ID, - ) - - assert admitted is (caller_role is CreatorUserRole.ACCOUNT) - assert redis.reads == [] - assert redis.operations == [] - assert redis.values == {} - assert redis.commands == {} - sqlite_session.refresh(run) - assert run.status == status - assert run.stop_requested_at is None - assert run.finished_at is None - if pause is not None and form is not None: - sqlite_session.refresh(pause) - sqlite_session.refresh(form) - assert pause.state_object_key == "state" - assert form.status == HumanInputFormStatus.WAITING - assert form.submitted_at is None - - -@pytest.mark.parametrize("owner", [None, b"account-another-user"]) -def test_unchecked_workflow_stop_skips_ownership_read(owner: bytes | None) -> None: - redis = _StopRedis( - values={_OWNER_KEY: owner} if owner is not None else {}, - read_error=AssertionError("Unchecked workflow stop must not read ownership"), - ) - - AppTaskControlService(redis_client=redis).stop_workflow_task_no_user_check(task_id=_TASK_ID) - - assert redis.reads == [] - assert redis.operations == ["legacy_flag", "graph_command"] - _assert_stop_flag(redis) - _assert_graph_command(redis) - - -def test_unchecked_workflow_stop_with_empty_task_id_is_noop() -> None: - redis = _StopRedis() - - AppTaskControlService(redis_client=redis).stop_workflow_task_no_user_check(task_id="") - - assert redis.reads == [] - assert redis.operations == [] - assert redis.values == {} - assert redis.commands == {} - - -@pytest.mark.parametrize("unchecked", [False, True]) -def test_flag_write_failure_propagates_before_graph_command(unchecked: bool) -> None: - error = RedisConnectionError("stop flag write failed") - redis = _StopRedis(values={_OWNER_KEY: b"account-user-1"}, flag_error=error) - service = AppTaskControlService(redis_client=redis) - stop = ( - partial(service.stop_workflow_task_no_user_check, task_id=_TASK_ID) - if unchecked - else partial(service.stop_task, _TASK_ID, InvokeFrom.EXPLORE, _USER_ID) - ) - - with pytest.raises(RedisConnectionError, match="stop flag write failed") as caught: - stop() - - assert caught.value is error - assert redis.operations == ["legacy_flag"] - assert _STOP_KEY not in redis.values - assert redis.commands == {} - - -def test_graph_redis_failure_is_swallowed_after_legacy_flag(caplog: pytest.LogCaptureFixture) -> None: - redis = _StopRedis( - values={_OWNER_KEY: b"account-user-1"}, command_error=RedisConnectionError("command write failed") - ) - - AppTaskControlService(redis_client=redis).stop_workflow_task_no_user_check(task_id=_TASK_ID) - - _assert_stop_flag(redis) - assert redis.operations == ["legacy_flag", "graph_command"] - assert redis.commands == {} - assert "Failed to send Engine abort command" in caplog.text - - -def test_ownership_read_failure_propagates_without_either_stop_signal() -> None: - error = RedisConnectionError("ownership read failed") - redis = _StopRedis(read_error=error) - - with pytest.raises(RedisConnectionError, match="ownership read failed") as caught: - AppTaskControlService(redis_client=redis).stop_task(_TASK_ID, InvokeFrom.EXPLORE, _USER_ID) - - assert caught.value is error - assert redis.reads == [_OWNER_KEY] - assert redis.operations == [] - assert redis.values == {} - assert redis.commands == {} - - -def test_service_instances_keep_their_redis_dependencies_separate() -> None: - first = _StopRedis(values={_OWNER_KEY: b"account-user-1"}) - second = _StopRedis(values={_OWNER_KEY: b"account-another-user"}) - first_service = AppTaskControlService(redis_client=first) - second_service = AppTaskControlService(redis_client=second) - - first_service.stop_task(_TASK_ID, InvokeFrom.EXPLORE, _USER_ID) - second_service.stop_workflow_task_no_user_check(task_id=_TASK_ID) - - assert first.reads == [_OWNER_KEY] - assert first.operations == ["legacy_flag"] - assert first.commands == {} - assert second.reads == [] - assert second.operations == ["legacy_flag", "graph_command"] - _assert_stop_flag(first) - _assert_stop_flag(second) - _assert_graph_command(second) - - -def test_legacy_static_entry_point_passes_global_client_through_the_same_implementation( - monkeypatch: pytest.MonkeyPatch, -) -> None: - redis = _StopRedis(values={_OWNER_KEY: b"end-user-user-1"}) - # Only the legacy composition point gets a global client; lower-level globals remain traps. - monkeypatch.setattr(task_module, "redis_client", redis) - - AppTaskService.stop_task( - _TASK_ID, InvokeFrom.SERVICE_API, _USER_ID, AppMode.CHAT, tenant_id="tenant-1", app_id="app-1" - ) - - assert redis.reads == [_OWNER_KEY] - assert redis.operations == ["legacy_flag"] - _assert_stop_flag(redis) - - -@pytest.mark.parametrize("cached_owner", [None, b"end-user-other", b"account-user-1"]) -def test_live_chatflow_stop_does_not_send_either_signal_for_another_owner( - monkeypatch: pytest.MonkeyPatch, sqlite_session: Session, cached_owner: bytes | None -) -> None: - task_id = "11111111-1111-1111-1111-111111111111" - redis = _StopRedis(values={_OWNER_KEY: cached_owner} if cached_owner else {}) - monkeypatch.setattr(type(task_module.db), "engine", property(lambda _: sqlite_session.get_bind())) - admissions: list[dict[str, object]] = [] - repository_type = task_module.DifyAPISQLAlchemyWorkflowRunRepository - original_stop = repository_type.stop_paused_workflow_task - - def stop( - repository: task_module.DifyAPISQLAlchemyWorkflowRunRepository, - *, - tenant_id: str, - app_id: str, - task_id: str, - owner: tuple[CreatorUserRole, str] | None, - ) -> WorkflowRun | None: - admissions.append({"tenant_id": tenant_id, "app_id": app_id, "task_id": task_id, "owner": owner}) - return original_stop(repository, tenant_id=tenant_id, app_id=app_id, task_id=task_id, owner=owner) - - monkeypatch.setattr(repository_type, "stop_paused_workflow_task", stop) - monkeypatch.setattr(task_module, "redis_client", redis) - AppTaskService.stop_task( - task_id, InvokeFrom.SERVICE_API, _USER_ID, AppMode.ADVANCED_CHAT, tenant_id="tenant-1", app_id="app-1" - ) - assert ( - admissions - == [ - { - "tenant_id": "tenant-1", - "app_id": "app-1", - "task_id": task_id, - "owner": (CreatorUserRole.END_USER, _USER_ID), - } + assert admitted is True + sqlite_session.refresh(run) + assert run.status == WorkflowExecutionStatus.RUNNING + assert run.stop_requested_at is not None + assert redis.reads == [] + assert redis.operations == ["legacy_flag", "graph_command"] + assert redis.values[f"generate_task_stopped:{task_id}"] == b"1" + assert [json.loads(command) for command in redis.commands[f"workflow:{task_id}:commands"]] == [ + {"command_type": "abort", "reason": "User requested stop"} ] - * 2 + + @staticmethod + @pytest.mark.parametrize("app_mode", [AppMode.WORKFLOW, AppMode.ADVANCED_CHAT]) + @pytest.mark.parametrize("status", [WorkflowExecutionStatus.RUNNING, WorkflowExecutionStatus.PAUSED]) + @pytest.mark.parametrize( + ("caller_role", "mismatch"), + [ + (CreatorUserRole.END_USER, "owner-id"), + (CreatorUserRole.END_USER, "owner-role"), + (CreatorUserRole.END_USER, "app"), + (CreatorUserRole.END_USER, "tenant"), + (CreatorUserRole.END_USER, "missing"), + (CreatorUserRole.ACCOUNT, "app"), + (CreatorUserRole.ACCOUNT, "tenant"), + (CreatorUserRole.ACCOUNT, "missing"), + ], ) - assert redis.reads == [] - assert redis.operations == [] - assert redis.commands == {} + def test_stop_app_task_preserves_unadmitted_workflow_and_pause( + sqlite_session: Session, + app_task_workflow_run: WorkflowRun, + app_mode: AppMode, + status: WorkflowExecutionStatus, + caller_role: CreatorUserRole, + mismatch: str, + ) -> None: + run = app_task_workflow_run + run.status = status + if mismatch == "owner-id": + run.created_by = "other-user" + if mismatch == "owner-role": + run.created_by_role = CreatorUserRole.ACCOUNT + pause = None + form = None + if status == WorkflowExecutionStatus.PAUSED: + pause = WorkflowPause(workflow_id=run.workflow_id, workflow_run_id=run.id, state_object_key="state") + form = HumanInputForm( + id="form-1", + tenant_id=run.tenant_id, + app_id=run.app_id, + workflow_run_id=run.id, + node_id="review", + form_definition="{}", + rendered_content="Review", + expiration_time=datetime(2026, 10, 20), + ) + sqlite_session.add_all([pause, form]) + sqlite_session.commit() + task_id = "22222222-2222-2222-2222-222222222222" if mismatch == "missing" else run.task_id + assert task_id is not None + redis = _AppTaskServiceStopredis( + read_error=AssertionError("Durable workflow ownership must not depend on Redis") + ) + + admitted = AppTaskControlService(redis_client=redis).stop_app_task( + tenant_id="other-tenant" if mismatch == "tenant" else run.tenant_id, + app_id="other-app" if mismatch == "app" else run.app_id, + task_id=task_id, + app_mode=app_mode, + caller_role=caller_role, + caller_id=_USER_ID, + ) + + assert admitted is (caller_role is CreatorUserRole.ACCOUNT) + assert redis.reads == [] + assert redis.operations == [] + assert redis.values == {} + assert redis.commands == {} + sqlite_session.refresh(run) + assert run.status == status + assert run.stop_requested_at is None + assert run.finished_at is None + if pause is not None and form is not None: + sqlite_session.refresh(pause) + sqlite_session.refresh(form) + assert pause.state_object_key == "state" + assert form.status == HumanInputFormStatus.WAITING + assert form.submitted_at is None + + @staticmethod + @pytest.mark.parametrize("owner", [None, b"account-another-user"]) + def test_unchecked_workflow_stop_skips_ownership_read(owner: bytes | None) -> None: + redis = _AppTaskServiceStopredis( + values={_OWNER_KEY: owner} if owner is not None else {}, + read_error=AssertionError("Unchecked workflow stop must not read ownership"), + ) + + AppTaskControlService(redis_client=redis).stop_workflow_task_no_user_check(task_id=_TASK_ID) + + assert redis.reads == [] + assert redis.operations == ["legacy_flag", "graph_command"] + _assert_stop_flag(redis) + _assert_graph_command(redis) + + @staticmethod + def test_unchecked_workflow_stop_with_empty_task_id_is_noop() -> None: + redis = _AppTaskServiceStopredis() + + AppTaskControlService(redis_client=redis).stop_workflow_task_no_user_check(task_id="") + + assert redis.reads == [] + assert redis.operations == [] + assert redis.values == {} + assert redis.commands == {} + + @staticmethod + @pytest.mark.parametrize("unchecked", [False, True]) + def test_flag_write_failure_propagates_before_graph_command(unchecked: bool) -> None: + error = RedisConnectionError("stop flag write failed") + redis = _AppTaskServiceStopredis(values={_OWNER_KEY: b"account-user-1"}, flag_error=error) + service = AppTaskControlService(redis_client=redis) + stop = ( + partial(service.stop_workflow_task_no_user_check, task_id=_TASK_ID) + if unchecked + else partial(service.stop_task, _TASK_ID, InvokeFrom.EXPLORE, _USER_ID) + ) + + with pytest.raises(RedisConnectionError, match="stop flag write failed") as caught: + stop() + + assert caught.value is error + assert redis.operations == ["legacy_flag"] + assert _STOP_KEY not in redis.values + assert redis.commands == {} + + @staticmethod + def test_graph_redis_failure_is_swallowed_after_legacy_flag(caplog: pytest.LogCaptureFixture) -> None: + redis = _AppTaskServiceStopredis( + values={_OWNER_KEY: b"account-user-1"}, command_error=RedisConnectionError("command write failed") + ) + + AppTaskControlService(redis_client=redis).stop_workflow_task_no_user_check(task_id=_TASK_ID) + + _assert_stop_flag(redis) + assert redis.operations == ["legacy_flag", "graph_command"] + assert redis.commands == {} + assert "Failed to send Engine abort command" in caplog.text + + @staticmethod + def test_ownership_read_failure_propagates_without_either_stop_signal() -> None: + error = RedisConnectionError("ownership read failed") + redis = _AppTaskServiceStopredis(read_error=error) + + with pytest.raises(RedisConnectionError, match="ownership read failed") as caught: + AppTaskControlService(redis_client=redis).stop_task(_TASK_ID, InvokeFrom.EXPLORE, _USER_ID) + + assert caught.value is error + assert redis.reads == [_OWNER_KEY] + assert redis.operations == [] + assert redis.values == {} + assert redis.commands == {} + + @staticmethod + def test_service_instances_keep_their_redis_dependencies_separate() -> None: + first = _AppTaskServiceStopredis(values={_OWNER_KEY: b"account-user-1"}) + second = _AppTaskServiceStopredis(values={_OWNER_KEY: b"account-another-user"}) + first_service = AppTaskControlService(redis_client=first) + second_service = AppTaskControlService(redis_client=second) + + first_service.stop_task(_TASK_ID, InvokeFrom.EXPLORE, _USER_ID) + second_service.stop_workflow_task_no_user_check(task_id=_TASK_ID) + + assert first.reads == [_OWNER_KEY] + assert first.operations == ["legacy_flag"] + assert first.commands == {} + assert second.reads == [] + assert second.operations == ["legacy_flag", "graph_command"] + _assert_stop_flag(first) + _assert_stop_flag(second) + _assert_graph_command(second) + + @staticmethod + def test_legacy_static_entry_point_passes_global_client_through_the_same_implementation( + monkeypatch: pytest.MonkeyPatch, + ) -> None: + redis = _AppTaskServiceStopredis(values={_OWNER_KEY: b"end-user-user-1"}) + # Only the legacy composition point gets a global client; lower-level globals remain traps. + monkeypatch.setattr(task_module, "redis_client", redis) + + AppTaskService.stop_task( + _TASK_ID, InvokeFrom.SERVICE_API, _USER_ID, AppMode.CHAT, tenant_id="tenant-1", app_id="app-1" + ) + + assert redis.reads == [_OWNER_KEY] + assert redis.operations == ["legacy_flag"] + _assert_stop_flag(redis) + + @staticmethod + @pytest.mark.parametrize("cached_owner", [None, b"end-user-other", b"account-user-1"]) + def test_live_chatflow_stop_does_not_send_either_signal_for_another_owner( + monkeypatch: pytest.MonkeyPatch, sqlite_session: Session, cached_owner: bytes | None + ) -> None: + task_id = "11111111-1111-1111-1111-111111111111" + redis = _AppTaskServiceStopredis(values={_OWNER_KEY: cached_owner} if cached_owner else {}) + monkeypatch.setattr(type(task_module.db), "engine", property(lambda _: sqlite_session.get_bind())) + admissions: list[dict[str, object]] = [] + repository_type = task_module.DifyAPISQLAlchemyWorkflowRunRepository + original_stop = repository_type.stop_paused_workflow_task + + def stop( + repository: task_module.DifyAPISQLAlchemyWorkflowRunRepository, + *, + tenant_id: str, + app_id: str, + task_id: str, + owner: tuple[CreatorUserRole, str] | None, + ) -> WorkflowRun | None: + admissions.append({"tenant_id": tenant_id, "app_id": app_id, "task_id": task_id, "owner": owner}) + return original_stop(repository, tenant_id=tenant_id, app_id=app_id, task_id=task_id, owner=owner) + + monkeypatch.setattr(repository_type, "stop_paused_workflow_task", stop) + monkeypatch.setattr(task_module, "redis_client", redis) + AppTaskService.stop_task( + task_id, InvokeFrom.SERVICE_API, _USER_ID, AppMode.ADVANCED_CHAT, tenant_id="tenant-1", app_id="app-1" + ) + assert ( + admissions + == [ + { + "tenant_id": "tenant-1", + "app_id": "app-1", + "task_id": task_id, + "owner": (CreatorUserRole.END_USER, _USER_ID), + } + ] + * 2 + ) + assert redis.reads == [] + assert redis.operations == [] + assert redis.commands == {} + + @staticmethod + def test_live_chatflow_stop_sends_signals_for_its_owner( + monkeypatch: pytest.MonkeyPatch, sqlite_session: Session + ) -> None: + task_id = "11111111-1111-1111-1111-111111111111" + redis = _AppTaskServiceStopredis(values={_OWNER_KEY: b"end-user-user-1"}) + run = WorkflowRun( + id="run-1", + tenant_id="tenant-1", + app_id="app-1", + workflow_id="workflow-1", + task_id=task_id, + type="workflow", + version="1", + triggered_from=WorkflowRunTriggeredFrom.APP_RUN, + status=WorkflowExecutionStatus.RUNNING, + created_by_role=CreatorUserRole.END_USER, + created_by=_USER_ID, + ) + sqlite_session.add(run) + sqlite_session.commit() + monkeypatch.setattr(type(task_module.db), "engine", property(lambda _: sqlite_session.get_bind())) + monkeypatch.setattr(task_module, "redis_client", redis) + AppTaskService.stop_task( + task_id, InvokeFrom.SERVICE_API, _USER_ID, AppMode.ADVANCED_CHAT, tenant_id="tenant-1", app_id="app-1" + ) + assert redis.values[f"generate_task_stopped:{task_id}"] == b"1" + assert json.loads(redis.commands[f"workflow:{task_id}:commands"][0])["command_type"] == "abort" + + @staticmethod + def test_invalid_workflow_task_id_is_rejected_before_database_lookup(monkeypatch: pytest.MonkeyPatch) -> None: + def reject_lookup(*_args: object, **_kwargs: object) -> None: + pytest.fail("Invalid task IDs must be rejected before database lookup") + + monkeypatch.setattr(task_module, "DifyAPISQLAlchemyWorkflowRunRepository", reject_lookup) + with pytest.raises(ValueError, match="Invalid task ID"): + AppTaskService.stop_workflow_task( + tenant_id="tenant-1", app_id="app-1", task_id="not-a-uuid", app_mode=AppMode.WORKFLOW, owner=None + ) -def test_live_chatflow_stop_sends_signals_for_its_owner( - monkeypatch: pytest.MonkeyPatch, sqlite_session: Session -) -> None: - task_id = "11111111-1111-1111-1111-111111111111" - redis = _StopRedis(values={_OWNER_KEY: b"end-user-user-1"}) +"""Legacy pauses recover task IDs on demand without an offline backfill.""" + + +_LegacyWorkflowTaskRecoveryStopRedis = _AppTaskServiceStopredis + + +def _seed_pause( + session: Session, + states: dict[str, bytes], + app_mode: AppMode = AppMode.WORKFLOW, + *, + pause_id: str | None = None, +) -> tuple[WorkflowRun, WorkflowPause, WorkflowResumptionContext]: run = WorkflowRun( - id="run-1", + id=str(uuid4()), tenant_id="tenant-1", app_id="app-1", workflow_id="workflow-1", - task_id=task_id, - type="workflow", - version="1", + type=_legacy_workflow_task_recovery_WorkflowType.WORKFLOW + if app_mode == AppMode.WORKFLOW + else _legacy_workflow_task_recovery_WorkflowType.CHAT, triggered_from=WorkflowRunTriggeredFrom.APP_RUN, - status=WorkflowExecutionStatus.RUNNING, + version="1", + status=WorkflowExecutionStatus.PAUSED, created_by_role=CreatorUserRole.END_USER, - created_by=_USER_ID, + created_by="user-1", + created_at=datetime(2026, 10, 10), + elapsed_time=0, + total_tokens=0, + total_steps=1, + exceptions_count=0, ) - sqlite_session.add(run) - sqlite_session.commit() - monkeypatch.setattr(type(task_module.db), "engine", property(lambda _: sqlite_session.get_bind())) - monkeypatch.setattr(task_module, "redis_client", redis) - AppTaskService.stop_task( - task_id, InvokeFrom.SERVICE_API, _USER_ID, AppMode.ADVANCED_CHAT, tenant_id="tenant-1", app_id="app-1" + pause = WorkflowPause(workflow_id=run.workflow_id, workflow_run_id=run.id, state_object_key=f"pause/{run.id}") + if pause_id is not None: + pause.id = pause_id + config = WorkflowUIBasedAppConfig( + tenant_id=run.tenant_id, app_id=run.app_id, app_mode=app_mode, workflow_id=run.workflow_id ) - assert redis.values[f"generate_task_stopped:{task_id}"] == b"1" - assert json.loads(redis.commands[f"workflow:{task_id}:commands"][0])["command_type"] == "abort" - - -def test_invalid_workflow_task_id_is_rejected_before_database_lookup(monkeypatch: pytest.MonkeyPatch) -> None: - def reject_lookup(*_args: object, **_kwargs: object) -> None: - pytest.fail("Invalid task IDs must be rejected before database lookup") - - monkeypatch.setattr(task_module, "DifyAPISQLAlchemyWorkflowRunRepository", reject_lookup) - with pytest.raises(ValueError, match="Invalid task ID"): - AppTaskService.stop_workflow_task( - tenant_id="tenant-1", app_id="app-1", task_id="not-a-uuid", app_mode=AppMode.WORKFLOW, owner=None + if app_mode == AppMode.WORKFLOW: + wrapper = _WorkflowGenerateEntityWrapper( + entity=WorkflowAppGenerateEntity( + task_id=str(uuid4()), + app_config=config, + inputs={}, + files=[], + user_id=run.created_by, + stream=True, + invoke_from=InvokeFrom.WEB_APP, + workflow_execution_id=run.id, + ) ) + else: + wrapper = _AdvancedChatAppGenerateEntityWrapper( + entity=AdvancedChatAppGenerateEntity( + task_id=str(uuid4()), + app_config=config, + inputs={}, + files=[], + user_id=run.created_by, + stream=True, + invoke_from=InvokeFrom.WEB_APP, + workflow_run_id=run.id, + query="Approve this run", + conversation_id="conversation-1", + ) + ) + context = WorkflowResumptionContext(generate_entity=wrapper, serialized_graph_runtime_state="{}") + states[pause.state_object_key] = context.dumps().encode() + session.add_all([run, pause]) + session.commit() + return run, pause, context + + +def _stop(service: AppTaskControlService, task_id: str, app_mode: AppMode = AppMode.WORKFLOW) -> bool: + return service.stop_workflow_task( + tenant_id="tenant-1", + app_id="app-1", + task_id=task_id, + app_mode=app_mode, + owner=(CreatorUserRole.END_USER, "user-1"), + ) + + +class TestLegacyWorkflowTaskRecoveryCases: + @staticmethod + @pytest.fixture + def pause_states(monkeypatch: pytest.MonkeyPatch) -> dict[str, bytes]: + states: dict[str, bytes] = {} + + def load(key: str) -> bytes: + try: + return states[key] + except KeyError as error: + raise FileNotFoundError(key) from error + + monkeypatch.setattr("repositories.sqlalchemy_api_workflow_run_repository.storage.load", load) + return states + + @staticmethod + @pytest.fixture + def recovery_redis() -> _LegacyWorkflowTaskRecoveryStopRedis: + return _LegacyWorkflowTaskRecoveryStopRedis(read_error=AssertionError("Recovery must use persisted ownership")) + + @staticmethod + @pytest.fixture + def recovery_service( + sqlite_engine: Engine, monkeypatch: pytest.MonkeyPatch, recovery_redis: _LegacyWorkflowTaskRecoveryStopRedis + ) -> AppTaskControlService: + monkeypatch.setattr(type(task_module.db), "engine", property(lambda _: sqlite_engine)) + monkeypatch.setattr("repositories.sqlalchemy_api_workflow_run_repository.storage.delete", lambda _: None) + monkeypatch.setattr(AppTaskService, "_persist_stopped_history", lambda _: None) + monkeypatch.setattr( + "core.workflow.nodes.agent_v2.session_store.WorkflowAgentWorkspaceStore.retire_workflow_run", + lambda *_args, **_kwargs: [], + ) + monkeypatch.setattr( + "tasks.collect_agent_resources_task.enqueue_agent_resource_collection", lambda **_kwargs: None + ) + topic = Topic(Redis(), "recovery-test") + monkeypatch.setattr(topic, "publish", lambda _payload: None) + monkeypatch.setattr(task_module.MessageBasedAppGenerator, "get_response_topic", lambda *_: topic) + return AppTaskControlService(redis_client=recovery_redis) + + @staticmethod + @pytest.mark.parametrize("app_mode", [AppMode.WORKFLOW, AppMode.ADVANCED_CHAT]) + def test_stop_recovers_and_cancels_a_legacy_pause_without_prior_replay( + sqlite_session: Session, + pause_states: dict[str, bytes], + recovery_service: AppTaskControlService, + app_mode: AppMode, + recovery_redis: _LegacyWorkflowTaskRecoveryStopRedis, + ) -> None: + run, pause, context = _seed_pause(sqlite_session, pause_states, app_mode) + task_id = context.get_generate_entity().task_id + assert task_id != run.id + assert run.task_id is None + context.get_generate_entity().task_id = task_id.upper() + pause_states[pause.state_object_key] = context.dumps().encode() + + assert _stop(recovery_service, task_id.upper(), app_mode) + + sqlite_session.refresh(run) + assert run.task_id == task_id + assert run.status == WorkflowExecutionStatus.STOPPED + assert run.stop_requested_at is not None + assert sqlite_session.scalar(select(WorkflowPause.id).where(WorkflowPause.id == pause.id)) is None + sqlite_session.commit() + pause_states.clear() + assert _stop(recovery_service, task_id, app_mode) + assert recovery_redis.reads == [] + assert recovery_redis.operations == [] + + @staticmethod + @pytest.mark.parametrize("field", ["tenant_id", "app_id", "created_by", "created_by_role"]) + def test_recovery_excludes_unauthorized_pauses_before_storage_reads( + sqlite_session: Session, + pause_states: dict[str, bytes], + recovery_service: AppTaskControlService, + monkeypatch: pytest.MonkeyPatch, + field: str, + ) -> None: + run, _, context = _seed_pause(sqlite_session, pause_states) + setattr(run, field, CreatorUserRole.ACCOUNT if field == "created_by_role" else "other-owner") + sqlite_session.commit() + + def reject_load(_key: str) -> bytes: + pytest.fail("Unauthorized pauses must not reach snapshot storage") + + monkeypatch.setattr("repositories.sqlalchemy_api_workflow_run_repository.storage.load", reject_load) + assert not _stop(recovery_service, context.get_generate_entity().task_id) + sqlite_session.refresh(run) + assert run.task_id is None + assert run.stop_requested_at is None + assert run.status == WorkflowExecutionStatus.PAUSED + + @staticmethod + @pytest.mark.parametrize("field", ["task_id", "run_id", "tenant_id", "app_id", "user_id", "app_config"]) + def test_recovery_rejects_snapshot_identity_mismatches_and_continues( + sqlite_session: Session, + pause_states: dict[str, bytes], + recovery_service: AppTaskControlService, + field: str, + ) -> None: + invalid, pause, invalid_context = _seed_pause(sqlite_session, pause_states, pause_id=str(UUID(int=1))) + valid, _, valid_context = _seed_pause(sqlite_session, pause_states, pause_id=str(UUID(int=2))) + state = json.loads(pause_states[pause.state_object_key]) + entity = state["generate_entity"]["entity"] + if field in {"tenant_id", "app_id"}: + entity["app_config"][field] = "other-owner" + elif field == "app_config": + del entity["app_config"] + else: + entity["workflow_execution_id" if field == "run_id" else field] = "invalid-task-or-owner" + pause_states[pause.state_object_key] = json.dumps(state).encode() + + assert not _stop(recovery_service, invalid_context.get_generate_entity().task_id) + sqlite_session.refresh(invalid) + sqlite_session.refresh(valid) + assert invalid.task_id is None + assert invalid.status == WorkflowExecutionStatus.PAUSED + assert invalid.stop_requested_at is None + assert valid.task_id is None + sqlite_session.commit() + + assert _stop(recovery_service, valid_context.get_generate_entity().task_id) + sqlite_session.refresh(invalid) + sqlite_session.refresh(valid) + assert invalid.task_id is None + assert valid.task_id == valid_context.get_generate_entity().task_id + assert valid.status == WorkflowExecutionStatus.STOPPED + + @staticmethod + @pytest.mark.parametrize("missing", ["pause", "state", "corrupt_state"]) + def test_recovery_skips_unrecoverable_pause_and_finds_later_match( + sqlite_session: Session, + pause_states: dict[str, bytes], + recovery_service: AppTaskControlService, + missing: str, + ) -> None: + invalid, pause, invalid_context = _seed_pause(sqlite_session, pause_states, pause_id=str(UUID(int=1))) + valid, _, valid_context = _seed_pause(sqlite_session, pause_states, pause_id=str(UUID(int=2))) + if missing == "pause": + sqlite_session.delete(pause) + sqlite_session.commit() + elif missing == "state": + del pause_states[pause.state_object_key] + else: + pause_states[pause.state_object_key] = b"invalid JSON" + + assert not _stop(recovery_service, invalid_context.get_generate_entity().task_id) + assert _stop(recovery_service, valid_context.get_generate_entity().task_id) + + sqlite_session.refresh(invalid) + sqlite_session.refresh(valid) + assert invalid.task_id is None + assert invalid.status == WorkflowExecutionStatus.PAUSED + assert invalid.stop_requested_at is None + assert valid.status == WorkflowExecutionStatus.STOPPED + + @staticmethod + @pytest.mark.parametrize("change", ["replaced_pause", "resumed", "already_bound"]) + def test_indexing_does_not_overwrite_a_concurrent_transition( + sqlite_session: Session, + sqlite_session_factory: sessionmaker[Session], + pause_states: dict[str, bytes], + change: str, + ) -> None: + run, pause, context = _seed_pause(sqlite_session, pause_states) + repository = DifyAPISQLAlchemyWorkflowRunRepository(sqlite_session_factory) + saved_pause = repository.get_workflow_pause(run.id) + assert saved_pause is not None + task_id = str(uuid4()) if change == "already_bound" else None + if change == "replaced_pause": + sqlite_session.delete(pause) + sqlite_session.flush() + sqlite_session.add( + WorkflowPause(workflow_id=run.workflow_id, workflow_run_id=run.id, state_object_key="replacement") + ) + elif change == "resumed": + run.status = WorkflowExecutionStatus.RUNNING + pause.resumed_at = datetime(2026, 10, 10, 1) + else: + run.task_id = task_id + sqlite_session.commit() + + assert not AppTaskService.index_workflow_pause_task( + workflow_run=run, + pause_entity=saved_pause, + resumption_context=context, + workflow_run_repo=repository, + ) + + sqlite_session.refresh(run) + assert run.task_id == task_id + + @staticmethod + @pytest.mark.parametrize("resumed", [False, True]) + def test_stop_retries_durable_admission_after_another_request_binds_the_task( + sqlite_session: Session, + sqlite_session_factory: sessionmaker[Session], + pause_states: dict[str, bytes], + recovery_service: AppTaskControlService, + monkeypatch: pytest.MonkeyPatch, + resumed: bool, + recovery_redis: _LegacyWorkflowTaskRecoveryStopRedis, + ) -> None: + run, pause, context = _seed_pause(sqlite_session, pause_states) + task_id = context.get_generate_entity().task_id + + def load(key: str) -> bytes: + with sqlite_session_factory() as session, session.begin(): + concurrent_run = session.get(WorkflowRun, run.id) + assert concurrent_run is not None + concurrent_run.task_id = task_id + if resumed: + concurrent_run.status = WorkflowExecutionStatus.RUNNING + concurrent_pause = session.get(WorkflowPause, pause.id) + assert concurrent_pause is not None + concurrent_pause.resumed_at = datetime(2026, 10, 10, 1) + return pause_states[key] + + monkeypatch.setattr("repositories.sqlalchemy_api_workflow_run_repository.storage.load", load) + assert _stop(recovery_service, task_id) + sqlite_session.refresh(run) + assert run.task_id == task_id + assert run.stop_requested_at is not None + assert run.status == (WorkflowExecutionStatus.RUNNING if resumed else WorkflowExecutionStatus.STOPPED) + assert recovery_redis.reads == [] + assert recovery_redis.operations == (["legacy_flag", "graph_command"] if resumed else []) + + @staticmethod + @pytest.mark.parametrize("status", [WorkflowExecutionStatus.PAUSED, WorkflowExecutionStatus.RUNNING]) + def test_prebound_tasks_stop_without_reading_legacy_snapshots( + sqlite_session: Session, + pause_states: dict[str, bytes], + recovery_service: AppTaskControlService, + monkeypatch: pytest.MonkeyPatch, + status: WorkflowExecutionStatus, + ) -> None: + run, _, context = _seed_pause(sqlite_session, pause_states) + run.task_id = context.get_generate_entity().task_id + run.status = status + sqlite_session.commit() + + def reject_load(_key: str) -> bytes: + pytest.fail("Already indexed tasks must not read snapshots") + + monkeypatch.setattr("repositories.sqlalchemy_api_workflow_run_repository.storage.load", reject_load) + assert _stop(recovery_service, context.get_generate_entity().task_id) + sqlite_session.refresh(run) + assert run.stop_requested_at is not None + + @staticmethod + @pytest.mark.parametrize("excluded", ["terminal", "pipeline", "resumed", "running"]) + def test_recovery_excludes_nonrecoverable_runs_before_storage_reads( + sqlite_session: Session, + pause_states: dict[str, bytes], + recovery_service: AppTaskControlService, + monkeypatch: pytest.MonkeyPatch, + excluded: str, + ) -> None: + run, pause, context = _seed_pause(sqlite_session, pause_states) + if excluded == "terminal": + run.status = WorkflowExecutionStatus.SUCCEEDED + elif excluded == "pipeline": + run.type = _legacy_workflow_task_recovery_WorkflowType.RAG_PIPELINE + elif excluded == "resumed": + pause.resumed_at = datetime(2026, 10, 10, 1) + else: + run.status = WorkflowExecutionStatus.RUNNING + sqlite_session.commit() + + def reject_load(_key: str) -> bytes: + pytest.fail("Ineligible runs must not read snapshots") + + monkeypatch.setattr("repositories.sqlalchemy_api_workflow_run_repository.storage.load", reject_load) + assert not _stop(recovery_service, context.get_generate_entity().task_id) + sqlite_session.refresh(run) + assert run.task_id is None + assert run.stop_requested_at is None + + @staticmethod + def test_recovery_pages_pauses_without_backfilling_unrelated_runs( + sqlite_session: Session, + pause_states: dict[str, bytes], + recovery_service: AppTaskControlService, + ) -> None: + for index in range(1, 101): + _seed_pause(sqlite_session, pause_states, pause_id=str(UUID(int=index))) + run, _, context = _seed_pause(sqlite_session, pause_states, pause_id=str(UUID(int=101))) + task_id = context.get_generate_entity().task_id + + assert _stop(recovery_service, task_id) + + indexed = sqlite_session.execute( + select(WorkflowRun.id, WorkflowRun.task_id).where(WorkflowRun.task_id.is_not(None)) + ) + assert indexed.tuples().all() == [(run.id, task_id)] + + @staticmethod + def test_snapshot_storage_reads_do_not_hold_a_database_transaction( + sqlite_session: Session, + sqlite_engine: Engine, + pause_states: dict[str, bytes], + recovery_service: AppTaskControlService, + monkeypatch: pytest.MonkeyPatch, + ) -> None: + _, _, context = _seed_pause(sqlite_session, pause_states) + active_transactions: set[Connection] = set() + + def load(key: str) -> bytes: + assert not active_transactions + return pause_states[key] + + monkeypatch.setattr("repositories.sqlalchemy_api_workflow_run_repository.storage.load", load) + listeners: tuple[tuple[str, Callable[[Connection], None]], ...] = ( + ("begin", active_transactions.add), + ("commit", active_transactions.discard), + ("rollback", active_transactions.discard), + ) + for name, listener in listeners: + event.listen(sqlite_engine, name, listener) + try: + assert _stop(recovery_service, context.get_generate_entity().task_id) + finally: + for name, listener in listeners: + event.remove(sqlite_engine, name, listener) + + +"""Paused stops persist before publishing terminal events or retiring resources.""" + + +_PausedWorkflowStopStopRedis = _AppTaskServiceStopredis + + +class TestPausedWorkflowStopCases: + @staticmethod + @pytest.mark.parametrize("app_mode", [AppMode.WORKFLOW, AppMode.ADVANCED_CHAT]) + @pytest.mark.parametrize("backend", ["sql", "celery", "logstore"]) + @pytest.mark.parametrize("node_status", [WorkflowNodeExecutionStatus.PAUSED, WorkflowNodeExecutionStatus.RETRY]) + def test_paused_stop_publishes_terminal_event_after_commit_and_retires_run_resources( + sqlite_session: Session, + sqlite_session_factory: sessionmaker[Session], + monkeypatch: pytest.MonkeyPatch, + app_mode: AppMode, + backend: str, + node_status: WorkflowNodeExecutionStatus, + config_overrides: Callable[..., None], + ) -> None: + config_overrides( + CORE_WORKFLOW_EXECUTION_REPOSITORY={ + "sql": ( + "core.repositories.sqlalchemy_workflow_execution_repository.SQLAlchemyWorkflowExecutionRepository" + ), + "celery": "core.repositories.celery_workflow_execution_repository.CeleryWorkflowExecutionRepository", + "logstore": ( + "extensions.logstore.repositories.logstore_workflow_execution_repository." + "LogstoreWorkflowExecutionRepository" + ), + }[backend], + CORE_WORKFLOW_NODE_EXECUTION_REPOSITORY={ + "sql": ( + "core.repositories.sqlalchemy_workflow_node_execution_repository." + "SQLAlchemyWorkflowNodeExecutionRepository" + ), + "celery": ( + "core.repositories.celery_workflow_node_execution_repository.CeleryWorkflowNodeExecutionRepository" + ), + "logstore": ( + "extensions.logstore.repositories.logstore_workflow_node_execution_repository." + "LogstoreWorkflowNodeExecutionRepository" + ), + }[backend], + LOGSTORE_DUAL_WRITE_ENABLED=False, + ) + monkeypatch.setattr( + "repositories.sqlalchemy_api_workflow_run_repository.naive_utc_now", lambda: datetime(2026, 10, 10, 0, 1) + ) + logstore_records: dict[str, list[dict[str, str]]] = {} + monkeypatch.setattr(AliyunLogStore, "_instance", None) + monkeypatch.setattr(AliyunLogStore, "_initialized", False) + monkeypatch.setenv("ALIYUN_SLS_ENDPOINT", "localhost") + monkeypatch.setattr(AliyunLogStore, "_check_endpoint_connectivity", staticmethod(lambda *_: None)) + monkeypatch.setattr( + AliyunLogStore, + "put_log", + lambda _self, name, record: logstore_records.setdefault(name, []).append(dict(record)), + ) + queued_node_saves: list[object] = [] + queued_saves: list[object] = [] + monkeypatch.setattr( + "core.repositories.celery_workflow_node_execution_repository.save_workflow_node_executions_task.delay", + lambda *args, **kwargs: queued_node_saves.append((args, kwargs)), + ) + monkeypatch.setattr( + "core.repositories.celery_workflow_execution_repository.save_workflow_execution_task.delay", + lambda *args, **kwargs: queued_saves.append((args, kwargs)), + ) + run = WorkflowRun( + id="run-1", + tenant_id="tenant-1", + app_id="app-1", + workflow_id="workflow-1", + type=_paused_workflow_stop_WorkflowType.WORKFLOW, + triggered_from=WorkflowRunTriggeredFrom.APP_RUN, + version="1", + status=WorkflowExecutionStatus.PAUSED, + created_by_role=CreatorUserRole.END_USER, + created_by="user-1", + created_at=datetime(2026, 10, 10), + elapsed_time=1, + total_tokens=0, + total_steps=2, + exceptions_count=0, + ) + run.task_id = "11111111-1111-1111-1111-111111111111" + pause = WorkflowPause(workflow_id=run.workflow_id, workflow_run_id=run.id, state_object_key="state") + sqlite_session.add_all([run, pause]) + sqlite_session.commit() + for node_app, origin in ( + ("app-1", WorkflowNodeExecutionTriggeredFrom.WORKFLOW_RUN), + ("source-app", WorkflowNodeExecutionTriggeredFrom.WORKFLOW_TOOL), + ): + SQLAlchemyWorkflowNodeExecutionRepository( + session_factory=sqlite_session_factory, + tenant_id="tenant-1", + user=EndUser(id="user-1"), + app_id=node_app, + triggered_from=origin, + ).save( + WorkflowNodeExecution( + id=f"node-{node_app}", + node_execution_id=f"node-{node_app}", + workflow_id=f"workflow-{node_app}", + workflow_execution_id=run.id, + index=1, + node_id="review", + node_type="human-input", + title="Review", + status=node_status, + process_data={"__dify_retry_history": [{"retry_index": 1, "error": "retry failed"}]}, + created_at=run.created_at, + ) + ) + monkeypatch.setattr(type(task_module.db), "engine", property(lambda _: sqlite_session_factory.kw["bind"])) + redis = _PausedWorkflowStopStopRedis() + monkeypatch.setattr(task_module, "redis_client", redis) + monkeypatch.setattr("repositories.sqlalchemy_api_workflow_run_repository.storage.delete", lambda _: None) + calls: list[str] = [] + + def retire(*, tenant_id: str, app_id: str, workflow_run_id: str) -> list[str]: + assert (tenant_id, app_id, workflow_run_id) == ("tenant-1", "app-1", "run-1") + with sqlite_session_factory() as session: + persisted = session.get(WorkflowRun, workflow_run_id) + assert persisted is not None + assert persisted.status == WorkflowExecutionStatus.STOPPED + assert persisted.elapsed_time == 1 + node = session.get(WorkflowNodeExecutionModel, "node-app-1") + assert node is not None + assert node.status == WorkflowNodeExecutionStatus.FAILED + assert node.finished_at == datetime(2026, 10, 10, 0, 1) + assert node.process_data_dict == {"__dify_retry_history": [{"retry_index": 1, "error": "retry failed"}]} + calls.append("retire") + return ["workspace-1"] + + monkeypatch.setattr(WorkflowAgentWorkspaceStore, "retire_workflow_run", staticmethod(retire)) + collections: list[dict[str, object]] = [] + + def enqueue(**kwargs: object) -> None: + calls.append("collect") + collections.append(kwargs) + + monkeypatch.setattr("tasks.collect_agent_resources_task.enqueue_agent_resource_collection", enqueue) + publications: list[bytes] = [] + + def publish(_channel: str, payload: bytes) -> None: + calls.append("publish") + publications.append(payload) + + topic_client = Redis() + monkeypatch.setattr(topic_client, "publish", publish) + topic = Topic(topic_client, "workflow-events") + monkeypatch.setattr(task_module.MessageBasedAppGenerator, "get_response_topic", lambda *_: topic) + + assert run.task_id is not None + AppTaskService.stop_workflow_task( + tenant_id="tenant-1", + app_id="app-1", + task_id=run.task_id, + app_mode=app_mode, + owner=(CreatorUserRole.END_USER, "user-1"), + ) + + assert calls == ["retire", "collect", "publish"] + assert collections == [{"tenant_id": "tenant-1", "workspace_ids": ["workspace-1"]}] + assert len(publications) == 1 + event = json.loads(publications[0]) + assert event["event"] == "workflow_finished" + assert event["data"]["status"] == "stopped" + assert event["data"]["elapsed_time"] == 1 + assert event["workflow_run_id"] == "run-1" + assert event["task_id"] == run.task_id + assert redis.reads == [] + assert redis.operations == [] + assert queued_saves == [] + if backend == "logstore": + assert len(logstore_records[AliyunLogStore.workflow_execution_logstore]) == 1 + history = logstore_records[AliyunLogStore.workflow_execution_logstore][0] + assert history["status"] == "stopped" + assert float(history["elapsed_time"]) == 1 + assert (history["tenant_id"], history["app_id"], history["id"]) == ("tenant-1", "app-1", "run-1") + node_history = logstore_records[AliyunLogStore.workflow_node_execution_logstore] + assert all(float(record["elapsed_time"]) == 60 for record in node_history) + assert all( + json.loads(record["process_data"])["__dify_retry_history"][0]["error"] == "retry failed" + for record in node_history + ) + assert {(record["app_id"], record["triggered_from"], record["status"]) for record in node_history} == { + ("app-1", "workflow-run", "failed"), + ("source-app", "workflow-tool", "failed"), + } + + @staticmethod + @pytest.mark.parametrize( + "task_id_input", ["11111111-aaaa-1111-1111-111111111111", "11111111AAAA11111111111111111111"] + ) + def test_stop_after_resume_claim_uses_durable_owner_before_queue_is_recreated( + sqlite_session: Session, + sqlite_session_factory: sessionmaker[Session], + monkeypatch: pytest.MonkeyPatch, + task_id_input: str, + ) -> None: + run = WorkflowRun( + id="run-1", + tenant_id="tenant-1", + app_id="app-1", + workflow_id="workflow-1", + type=_paused_workflow_stop_WorkflowType.WORKFLOW, + triggered_from=WorkflowRunTriggeredFrom.APP_RUN, + version="1", + status=WorkflowExecutionStatus.PAUSED, + created_by_role=CreatorUserRole.END_USER, + created_by="user-1", + ) + run.task_id = "11111111-aaaa-1111-1111-111111111111" + pause = WorkflowPause(workflow_id=run.workflow_id, workflow_run_id=run.id, state_object_key="state") + sqlite_session.add_all([run, pause]) + sqlite_session.commit() + repository = DifyAPISQLAlchemyWorkflowRunRepository(sqlite_session_factory) + entity = repository.get_workflow_pause(run.id) + assert entity is not None + repository.resume_workflow_pause(run.id, entity) + redis = _PausedWorkflowStopStopRedis() + monkeypatch.setattr(type(task_module.db), "engine", property(lambda _: sqlite_session_factory.kw["bind"])) + monkeypatch.setattr(task_module, "redis_client", redis) + AppTaskService.stop_workflow_task( + tenant_id="tenant-1", + app_id="app-1", + task_id=task_id_input, + app_mode=AppMode.ADVANCED_CHAT, + owner=(CreatorUserRole.END_USER, "user-1"), + ) + assert redis.reads == [] + assert redis.values[f"generate_task_stopped:{run.task_id}"] == b"1" + assert json.loads(redis.commands[f"workflow:{run.task_id}:commands"][0])["command_type"] == "abort" + + @staticmethod + @pytest.mark.parametrize( + ("tenant_id", "app_id", "owner"), + [ + ("other-tenant", "app-1", None), + ("tenant-1", "other-app", None), + ("tenant-1", "app-1", (CreatorUserRole.END_USER, "other-user")), + ("tenant-1", "app-1", (CreatorUserRole.ACCOUNT, "user-1")), + ], + ) + def test_running_task_admission_rejects_foreign_tenant_app_user_and_role( + sqlite_session: Session, + sqlite_session_factory: sessionmaker[Session], + monkeypatch: pytest.MonkeyPatch, + tenant_id: str, + app_id: str, + owner: tuple[CreatorUserRole, str] | None, + ) -> None: + task_id = "22222222-2222-2222-2222-222222222222" + run = WorkflowRun( + id="run-1", + tenant_id="tenant-1", + app_id="app-1", + workflow_id="workflow-1", + task_id=task_id, + type=_paused_workflow_stop_WorkflowType.WORKFLOW, + triggered_from=WorkflowRunTriggeredFrom.APP_RUN, + version="1", + status=WorkflowExecutionStatus.RUNNING, + created_by_role=CreatorUserRole.END_USER, + created_by="user-1", + ) + sqlite_session.add(run) + sqlite_session.commit() + monkeypatch.setattr(type(task_module.db), "engine", property(lambda _: sqlite_session_factory.kw["bind"])) + redis = _PausedWorkflowStopStopRedis(values={f"generate_task_belong:{task_id}": b"end-user-user-1"}) + monkeypatch.setattr(task_module, "redis_client", redis) + AppTaskService.stop_workflow_task( + tenant_id=tenant_id, app_id=app_id, task_id=task_id, app_mode=AppMode.WORKFLOW, owner=owner + ) + sqlite_session.expire_all() + assert run.status == WorkflowExecutionStatus.RUNNING + assert run.stop_requested_at is None + assert redis.reads == [] + assert redis.operations == [] diff --git a/api/tests/unit_tests/services/test_app_tracing_atomic_operations.py b/api/tests/unit_tests/services/test_app_tracing_atomic_operations.py deleted file mode 100644 index 035bb5eed5a..00000000000 --- a/api/tests/unit_tests/services/test_app_tracing_atomic_operations.py +++ /dev/null @@ -1,225 +0,0 @@ -"""Public tracing operations keep credentials and policy in one transaction.""" - -import json -from dataclasses import dataclass - -import pytest -from sqlalchemy import event, select -from sqlalchemy.orm import Session, sessionmaker - -from machinery.context import RequestContext -from models.model import App, AppMode, TraceAppConfig -from repositories.app.console_repository import ConsoleAppRepository -from repositories.app.tracing_config_repository import SQLAlchemyAppTracingConfigRepository -from services.app_tracing_config_gateway import OpsTraceManagerGateway -from services.app_tracing_config_service import ( - AppTraceSettings, - AppTracingConfigConflictError, - AppTracingConfigNotFoundError, - AppTracingConfigRecord, - AppTracingConfigService, - AppTracingConfigVerificationFailedError, -) - -CONTEXT = RequestContext("request", None, "account", "tenant") - - -@dataclass -class Case: - sessions: sessionmaker[Session] - store: SQLAlchemyAppTracingConfigRepository - service: AppTracingConfigService - gateway: OpsTraceManagerGateway - reader: ConsoleAppRepository - - def policy(self) -> AppTraceSettings: - return self.reader.get_trace(CONTEXT, "app") - - def config(self, provider: str = "langfuse") -> AppTracingConfigRecord | None: - return self.store.get(workspace_id="tenant", app_id="app", tracing_provider=provider) - - -@pytest.fixture -def case(sqlite_session_factory: sessionmaker[Session], monkeypatch: pytest.MonkeyPatch) -> Case: - with sqlite_session_factory.begin() as session: - session.add( - App( - id="app", - tenant_id="tenant", - name="Trace", - mode=AppMode.WORKFLOW, - enable_site=False, - enable_api=False, - tracing=json.dumps( - {"enabled": False, "tracing_provider": None, "capture_content": True, "capture_agent": False} - ), - ) - ) - gateway = OpsTraceManagerGateway() - - def prepare(*, tracing_config: dict[str, object], **_: object) -> dict[str, object]: - return dict(tracing_config) - - monkeypatch.setattr(gateway, "prepare_new_config", prepare) - monkeypatch.setattr(gateway, "prepare_updated_config", prepare) - monkeypatch.setattr(gateway, "validate_saved_config", lambda **_: None) - store = SQLAlchemyAppTracingConfigRepository(session_factory=sqlite_session_factory) - return Case( - sqlite_session_factory, - store, - AppTracingConfigService(configs=store, provider=gateway), - gateway, - ConsoleAppRepository(session_factory=sqlite_session_factory), - ) - - -@pytest.mark.parametrize("activate", [False, True]) -def test_create_and_update_preserve_capture_and_explicit_activation(case: Case, activate: bool) -> None: - case.service.create(CONTEXT, "app", "langfuse", {"secret": "original"}, activate=activate) - assert case.policy() == AppTraceSettings(activate, "langfuse" if activate else None, True, False) - case.service.update(CONTEXT, "app", "langfuse", {"secret": "updated"}, activate=True) - assert case.policy() == AppTraceSettings(True, "langfuse", True, False) - record = case.config() - assert record is not None - assert record.tracing_config == {"secret": "updated"} - - -def test_save_activation_uses_capture_policy_changed_during_vendor_validation( - case: Case, monkeypatch: pytest.MonkeyPatch -) -> None: - def prepare(**_: object) -> dict[str, object]: - case.service.set_trace(CONTEXT, "app", AppTraceSettings(False, None, False, True)) - return {"secret": "validated"} - - monkeypatch.setattr(case.gateway, "prepare_new_config", prepare) - case.service.create(CONTEXT, "app", "langfuse", {"secret": "input"}, activate=True) - assert case.policy() == AppTraceSettings(True, "langfuse", False, True) - - -@pytest.mark.parametrize("operation", ["create", "update", "delete"]) -def test_transaction_failure_rolls_back_credentials_and_policy(case: Case, operation: str) -> None: - if operation != "create": - case.service.create(CONTEXT, "app", "langfuse", {"secret": "original"}, activate=True) - before_policy, before_record = case.policy(), case.config() - - def fail_commit(session: Session) -> None: - session.flush() - raise RuntimeError("commit failed") - - event.listen(case.sessions, "before_commit", fail_commit) - try: - - def mutate() -> None: - if operation == "delete": - case.service.delete(CONTEXT, "app", "langfuse") - elif operation == "create": - case.service.create(CONTEXT, "app", "langfuse", {"secret": "new"}, activate=True) - else: - case.service.update(CONTEXT, "app", "langfuse", {"secret": "new"}, activate=True) - - with pytest.raises(RuntimeError, match="commit failed"): - mutate() - finally: - event.remove(case.sessions, "before_commit", fail_commit) - assert case.policy() == before_policy - assert case.config() == before_record - - -@pytest.mark.parametrize("enabled", [False, True]) -def test_delete_current_clears_selection_but_other_provider_preserves_policy(case: Case, enabled: bool) -> None: - case.service.create(CONTEXT, "app", "langfuse", {"secret": "selected"}, activate=True) - case.service.create(CONTEXT, "app", "langsmith", {"secret": "other"}) - policy = AppTraceSettings(enabled, "langfuse", True, False) - case.service.set_trace(CONTEXT, "app", policy) - case.service.delete(CONTEXT, "app", "langsmith") - assert case.policy() == policy - case.service.delete(CONTEXT, "app", "langfuse") - assert case.policy() == AppTraceSettings(False, None, True, False) - assert case.config() is None - - -def test_validation_failure_writes_neither_config_nor_policy(case: Case, monkeypatch: pytest.MonkeyPatch) -> None: - def reject(**_: object) -> dict[str, object]: - raise AppTracingConfigVerificationFailedError - - monkeypatch.setattr(case.gateway, "prepare_new_config", reject) - with pytest.raises(AppTracingConfigVerificationFailedError): - case.service.create(CONTEXT, "app", "langfuse", {"secret": "bad"}, activate=True) - assert case.config() is None - assert case.policy() == AppTraceSettings(False, None, True, False) - - -@pytest.mark.parametrize("state", ["missing", "empty", "inactive"]) -def test_enable_requires_usable_owned_config(case: Case, state: str) -> None: - if state != "missing": - case.service.create(CONTEXT, "app", "langfuse", {} if state == "empty" else {"secret": "valid"}) - if state == "inactive": - with case.sessions.begin() as session: - record = session.scalar(select(TraceAppConfig)) - assert record is not None - record.is_active = False - with pytest.raises(AppTracingConfigNotFoundError): - case.service.set_trace(CONTEXT, "app", AppTraceSettings(True, "langfuse")) - assert case.policy() == AppTraceSettings(False, None, True, False) - - -def test_update_rejects_credentials_replaced_while_validation_was_in_flight( - case: Case, monkeypatch: pytest.MonkeyPatch -) -> None: - case.service.create(CONTEXT, "app", "langfuse", {"secret": "original"}) - - def prepare(**_: object) -> dict[str, object]: - current = case.config() - assert current is not None - case.store.update( - workspace_id="tenant", - app_id="app", - tracing_provider="langfuse", - tracing_config={"secret": "concurrent"}, - expected=current, - ) - return {"secret": "stale"} - - monkeypatch.setattr(case.gateway, "prepare_updated_config", prepare) - with pytest.raises(AppTracingConfigConflictError): - case.service.update(CONTEXT, "app", "langfuse", {"secret": "masked"}, activate=True) - saved = case.config() - assert saved is not None - assert saved.tracing_config == {"secret": "concurrent"} - assert case.policy() == AppTraceSettings(False, None, True, False) - - -def test_enable_rejects_configuration_deleted_during_validation(case: Case, monkeypatch: pytest.MonkeyPatch) -> None: - case.service.create(CONTEXT, "app", "langfuse", {"secret": "original"}) - - def validate(**_: object) -> None: - case.service.delete(CONTEXT, "app", "langfuse") - - monkeypatch.setattr(case.gateway, "validate_saved_config", validate) - with pytest.raises(AppTracingConfigConflictError): - case.service.set_trace(CONTEXT, "app", AppTraceSettings(True, "langfuse")) - assert case.policy() == AppTraceSettings(False, None, True, False) - - -@pytest.mark.parametrize("recreate", [False, True]) -def test_patch_rejects_configuration_deleted_during_validation( - case: Case, monkeypatch: pytest.MonkeyPatch, recreate: bool -) -> None: - case.service.create(CONTEXT, "app", "langfuse", {"secret": "original"}) - - def prepare(**_: object) -> dict[str, object]: - case.service.delete(CONTEXT, "app", "langfuse") - if recreate: - case.service.create(CONTEXT, "app", "langfuse", {"secret": "original"}) - return {"secret": "stale"} - - monkeypatch.setattr(case.gateway, "prepare_updated_config", prepare) - with pytest.raises(AppTracingConfigConflictError): - case.service.update(CONTEXT, "app", "langfuse", {"secret": "masked"}, activate=True) - assert case.policy() == AppTraceSettings(False, None, True, False) - saved = case.config() - if recreate: - assert saved is not None - assert saved.tracing_config == {"secret": "original"} - else: - assert saved is None diff --git a/api/tests/unit_tests/services/test_app_tracing_config_gateway.py b/api/tests/unit_tests/services/test_app_tracing_config_gateway.py index ba47e3a7dff..5b85b730f8e 100644 --- a/api/tests/unit_tests/services/test_app_tracing_config_gateway.py +++ b/api/tests/unit_tests/services/test_app_tracing_config_gateway.py @@ -1,12 +1,63 @@ +import importlib as _test_importlib +import inspect +import json +import os +import traceback +from collections import deque +from collections.abc import Callable, Iterator, Mapping, Sequence +from copy import deepcopy +from dataclasses import dataclass, field, replace +from datetime import UTC, datetime, timedelta +from decimal import Decimal from functools import partial +from typing import Never from unittest.mock import MagicMock, patch +from urllib.parse import quote +from uuid import UUID, uuid4 +import grpc # type: ignore[import-untyped] +import httpx +import mlflow +import opik import pytest +import wandb +import weave +from dify_trace_aliyun.aliyun_trace import AliyunDataTrace +from dify_trace_arize_phoenix import arize_configuration +from dify_trace_arize_phoenix.arize_configuration import DESTINATION_KEY +from dify_trace_arize_phoenix.config import ArizeConfig +from dify_trace_arize_phoenix.configuration import PhoenixConfigurationError +from dify_trace_mlflow.configuration import DatabricksConfigurationError +from dify_trace_mlflow.mlflow_trace import MLflowDataTrace +from dify_trace_mlflow.oss_configuration import MLflowConfigurationError +from dify_trace_opik.configuration import OpikConfigurationError +from dify_trace_opik.opik_trace import OpikDataTrace +from dify_trace_tencent.configuration import TencentConfigurationError +from dify_trace_weave.configuration import WeaveConfigurationError +from dify_trace_weave.weave_trace import WeaveDataTrace +from flask import Flask +from opentelemetry.proto.collector.trace.v1.trace_service_pb2 import ExportTraceServiceRequest +from opik.config import OpikConfig as SDKOpikConfig from pydantic import ValidationInfo, field_validator +from sqlalchemy import select +from sqlalchemy.orm import Session, sessionmaker +from controllers.console.app import ops_trace as controller +from controllers.console.app.error import InvalidTracingConfigError +from core.helper import ssrf_proxy +from core.ops import ops_trace_manager from core.ops.entities.config_entity import BaseTracingConfig, TracingProviderEnum from core.ops.exceptions import TraceProviderNotInstalledError -from core.ops.ops_trace_manager import TracingProviderConfigEntry +from core.ops.ops_trace_manager import OpsTraceManager, TracingProviderConfigEntry +from core.ops.unified_trace.entities import CanonicalSpan, CanonicalSpanKind, CanonicalSpanStatus, CanonicalTrace +from extensions.ext_application_services import ApplicationServices +from machinery.context import RequestContext +from models.application_trace import ApplicationTraceRun +from models.enums import ConversationFromSource +from models.model import App, AppMode +from repositories.app.console_repository import ConsoleAppRepository +from repositories.app.tracing_config_repository import SQLAlchemyAppTracingConfigRepository +from repositories.application_trace_repository import ApplicationTraceRepository from services import app_tracing_config_gateway as gateway_module from services.app_tracing_config_gateway import OpsTraceManagerGateway from services.app_tracing_config_service import ( @@ -14,8 +65,253 @@ from services.app_tracing_config_service import ( AppTracingConfigInvalidProviderError, AppTracingConfigProcessingError, AppTracingConfigProviderUnavailableError, + AppTracingConfigRecord, + AppTracingConfigService, AppTracingConfigVerificationFailedError, ) +from services.tracing.provider_dispatch import export_trace, frozen_project_url +from services.tracing.service import ApplicationTracingService +from tests.helpers.unit.model_factories import make_app, make_message + +pytest_plugins = ("tests.unit_tests.core.ops.test_ops_trace_manager",) + +"""Real RSA/SQL/controller Aliyun configuration through a recorded HTTP boundary.""" + + +@dataclass +class AliyunCase: + gateway: OpsTraceManagerGateway + service: AppTracingConfigService + store: SQLAlchemyAppTracingConfigRepository + reader: ConsoleAppRepository + context: RequestContext + app_id: str + verification_calls: list[tuple[str, dict[str, object]]] = field(default_factory=list) + verification_succeeds: bool = True + + def read(self, provider: str = "aliyun") -> AppTracingConfigRecord | None: + return self.store.get( + workspace_id=self.context.active_workspace_id, app_id=self.app_id, tracing_provider=provider + ) + + +class TestAliyunTracingConfigGatewayCases: + @staticmethod + @pytest.fixture + def case( + workspace: str, sqlite_session_factory: sessionmaker[Session], monkeypatch: pytest.MonkeyPatch + ) -> AliyunCase: + app_id = str(uuid4()) + with sqlite_session_factory.begin() as session: + session.add( + make_app( + app_id=app_id, + tenant_id=workspace, + tracing=json.dumps( + { + "enabled": True, + "tracing_provider": "langsmith", + "capture_content": True, + "capture_agent": False, + } + ), + ) + ) + gateway = OpsTraceManagerGateway() + store = SQLAlchemyAppTracingConfigRepository(session_factory=sqlite_session_factory) + old = OpsTraceManager.encrypt_tracing_config( + workspace, "langsmith", {"api_key": "synthetic-other", "project": "other", "endpoint": "https://other.test"} + ) + assert store.create(workspace_id=workspace, app_id=app_id, tracing_provider="langsmith", tracing_config=old) + result = AliyunCase( + gateway, + AppTracingConfigService(configs=store, provider=gateway), + store, + ConsoleAppRepository(session_factory=sqlite_session_factory), + RequestContext("request", None, "account", workspace), + app_id, + ) + + def verification_boundary(method: str, url: str, **options: object) -> httpx.Response: + assert method == "POST" + content = options["content"] + assert isinstance(content, bytes) + payload = ExportTraceServiceRequest.FromString(content) + assert len(payload.resource_spans) == 1 + assert len(payload.resource_spans[0].scope_spans) == 0 + assert [attribute.key for attribute in payload.resource_spans[0].resource.attributes] == ["service.name"] + result.verification_calls.append((url, dict(options))) + status = 200 if result.verification_succeeds else 403 + return httpx.Response(status, content=b'"success"', request=httpx.Request(method, url)) + + def forbidden(*_args: object, **_kwargs: object) -> Never: + raise AssertionError("Configuration must not initialize a legacy exporter or fetch its console link") + + monkeypatch.setattr(ssrf_proxy, "make_request", verification_boundary) + monkeypatch.setattr(OpsTraceManager, "check_trace_config_is_effective", forbidden) + monkeypatch.setattr(OpsTraceManager, "get_trace_config_project_url", forbidden) + monkeypatch.setattr(AliyunDataTrace, "__init__", forbidden) + return result + + @staticmethod + def test_post_masked_patch_and_get_preserve_real_ciphertext_and_policy(case: AliyunCase) -> None: + other = case.read("langsmith") + submitted = { + "endpoint": " HTTP://TRACING-ANALYSIS-DC-QD.ALIYUNCS.COM/reverse/ ", + "app_name": " 中文 app_Name ", + "license_key": "literal%2F@synthetic", + } + original = dict(submitted) + case.service.create(case.context, case.app_id, "aliyun", submitted, activate=True) + before = case.read() + assert before is not None + assert before.tracing_config is not None + stored = before.tracing_config + assert submitted == original + assert stored["endpoint"] == "http://tracing-analysis-dc-qd.aliyuncs.com/reverse" + assert stored["app_name"] == "中文 app_Name" + assert stored["license_key"] != submitted["license_key"] + assert OpsTraceManager.decrypt_tracing_config(case.context.active_workspace_id, "aliyun", stored)[ + "license_key" + ] == ("literal%2F@synthetic") + shown = case.service.get(case.context, case.app_id, "aliyun") + assert shown is not None + assert shown.tracing_config is not None + masked = shown.tracing_config["license_key"] + assert isinstance(masked, str) + assert "*" in masked + assert masked != submitted["license_key"] + assert shown.tracing_config["project_url"] == "https://arms.console.aliyun.com/#/llm" + assert len(case.verification_calls) == 1 + case.service.update( + case.context, + case.app_id, + "aliyun", + { + "endpoint": "https://project.region.log.aliyuncs.com/prefix/", + "app_name": " second_app ", + "license_key": masked, + }, + ) + after = case.read() + assert after is not None + assert after.tracing_config is not None + assert after.tracing_config["license_key"] == stored["license_key"] + assert after.tracing_config["app_name"] == "second_app" + assert after.tracing_config["endpoint"] == "https://project.region.log.aliyuncs.com/prefix" + assert case.verification_calls[-1][0] == ( + "https://project.region.log.aliyuncs.com/prefix/adapt_literal%252F@synthetic/api/v1/traces" + ) + assert case.service.get(case.context, case.app_id, "aliyun") is not None + assert len(case.verification_calls) == 2 + assert case.read("langsmith") == other + policy = case.reader.get_trace(case.context, case.app_id) + assert policy.tracing_provider == "aliyun" + assert policy.enabled + assert policy.capture_content is True + assert policy.capture_agent is False + + @staticmethod + @pytest.mark.parametrize("update", [False, True]) + @pytest.mark.parametrize("failure", ["invalid-schema", "verification-failed"]) + def test_failed_save_preserves_rows_and_active_policy(case: AliyunCase, update: bool, failure: str) -> None: + original = {"endpoint": "https://host.test", "app_name": "original", "license_key": "synthetic-secret"} + if update: + case.service.create(case.context, case.app_id, "aliyun", original) + before, other, policy = case.read(), case.read("langsmith"), case.reader.get_trace(case.context, case.app_id) + submitted = {**original, "app_name": "changed"} + if failure == "invalid-schema": + submitted["endpoint"] = "https://host.test/adapt_DO-NOT-PERSIST/api/otlp/traces" + expected = AppTracingConfigInvalidConfigurationError + else: + case.verification_succeeds = False + expected = AppTracingConfigVerificationFailedError + save = case.service.update if update else case.service.create + if update: + submitted["license_key"] = "******" + with pytest.raises(expected): + save(case.context, case.app_id, "aliyun", submitted, activate=True) + assert case.read() == before + assert case.read("langsmith") == other + assert case.reader.get_trace(case.context, case.app_id) == policy + + @staticmethod + @pytest.mark.parametrize("cached_url", [None, "https://wrong.test/stale"]) + def test_historical_get_is_pure_even_for_a_now_invalid_ingestion_url( + case: AliyunCase, cached_url: str | None + ) -> None: + stored = OpsTraceManager.encrypt_tracing_config( + case.context.active_workspace_id, + "aliyun", + { + "endpoint": "https://host.test/adapt_old/api/otlp/traces", + "app_name": "historical_name", + "license_key": "synthetic-key", + }, + ) + if cached_url: + stored["project_url"] = cached_url + before = dict(stored) + result = case.gateway.present_config( + workspace_id=case.context.active_workspace_id, tracing_provider="aliyun", tracing_config=stored + ) + assert result["project_url"] == "https://arms.console.aliyun.com/#/llm" + assert isinstance(result["license_key"], str) + assert "*" in result["license_key"] + assert stored == before + assert case.verification_calls == [] + + @staticmethod + @pytest.mark.parametrize("update", [False, True]) + @pytest.mark.parametrize("encoded", [False, True]) + def test_controller_rejects_embedded_license_without_echo_or_state_change( + case: AliyunCase, + app: Flask, + account_application_services: ApplicationServices, + monkeypatch: pytest.MonkeyPatch, + caplog: pytest.LogCaptureFixture, + update: bool, + encoded: bool, + ) -> None: + old_key = "SYNTHETIC_OLD_LICENSE" + original = {"endpoint": "https://host.test", "app_name": "old-app", "license_key": old_key} + if update: + case.service.create(case.context, case.app_id, "aliyun", original, activate=True) + before, other, policy = case.read(), case.read("langsmith"), case.reader.get_trace(case.context, case.app_id) + before_checks = len(case.verification_calls) + services = replace(account_application_services, app_tracing_configs=case.service) + monkeypatch.setattr(controller, "application_services", lambda: services) + key = "SYNTHETIC_PRIVATE_LICENSE@+/%2F" + url_key = quote(key, safe="") if encoded else key + payload = controller.TraceConfigPayload( + tracing_provider="aliyun", + tracing_config={ + "endpoint": f"https://host.test/adapt_{url_key}/api/otlp/traces", + "app_name": "should-not-save", + "license_key": "******" if update else key, + }, + activate=True, + ) + save = inspect.unwrap(controller.TraceAppConfigApi.patch if update else controller.TraceAppConfigApi.post) + with app.test_request_context(): + with pytest.raises(InvalidTracingConfigError) as caught: + save(controller.TraceAppConfigApi(), payload, case.context, UUID(case.app_id)) + # The real controller exposes this fixed BaseHTTPException body. Neither + # Pydantic inputs nor a credential-bearing exception chain reaches it. + assert caught.value.data == { + "code": "invalid_tracing_config", + "message": "The tracing configuration is invalid.", + "status": 400, + } + rendered = json.dumps(caught.value.data) + str(caught.value) + repr(caught.value) + rendered += "".join(traceback.format_exception(caught.value)) + caplog.text + assert key not in rendered + assert quote(key, safe="") not in rendered + assert old_key not in rendered + assert case.read() == before + assert case.read("langsmith") == other + assert case.reader.get_trace(case.context, case.app_id) == policy + assert len(case.verification_calls) == before_checks class _ProviderConfig(BaseTracingConfig): @@ -41,434 +337,2474 @@ def _provider_entry(*, other_keys: list[str] | None = None) -> TracingProviderCo } -def test_validate_provider_rejects_unknown_provider(monkeypatch: pytest.MonkeyPatch) -> None: - monkeypatch.setattr(gateway_module, "provider_config_map", {}) +class TestAppTracingConfigGatewayCases: + @staticmethod + def test_validate_provider_rejects_unknown_provider(monkeypatch: pytest.MonkeyPatch) -> None: + monkeypatch.setattr(gateway_module, "provider_config_map", {}) - with pytest.raises(AppTracingConfigInvalidProviderError, match="Invalid tracing provider: unknown"): - OpsTraceManagerGateway().validate_provider("unknown") + with pytest.raises(AppTracingConfigInvalidProviderError, match="Invalid tracing provider: unknown"): + OpsTraceManagerGateway().validate_provider("unknown") + @staticmethod + @pytest.mark.parametrize("provider", TracingProviderEnum) + def test_validate_provider_does_not_load_optional_dependencies(provider: TracingProviderEnum) -> None: + with patch.object(OpsTraceManagerGateway, "_provider_config") as load_provider: + OpsTraceManagerGateway().validate_provider(provider.value) -@pytest.mark.parametrize("provider", TracingProviderEnum) -def test_validate_provider_does_not_load_optional_dependencies(provider: TracingProviderEnum) -> None: - with patch.object(OpsTraceManagerGateway, "_provider_config") as load_provider: - OpsTraceManagerGateway().validate_provider(provider.value) + load_provider.assert_not_called() - load_provider.assert_not_called() + @staticmethod + def test_require_provider_available_loads_provider_dependencies() -> None: + with patch.object(OpsTraceManagerGateway, "_provider_config") as load_provider: + OpsTraceManagerGateway().require_provider_available("weave") + load_provider.assert_called_once_with("weave") -def test_require_provider_available_loads_provider_dependencies() -> None: - with patch.object(OpsTraceManagerGateway, "_provider_config") as load_provider: - OpsTraceManagerGateway().require_provider_available("weave") + @staticmethod + def test_require_provider_available_reports_missing_dependency() -> None: + unavailable = AppTracingConfigProviderUnavailableError() + with ( + patch.object(OpsTraceManagerGateway, "_provider_config", side_effect=unavailable), + pytest.raises(AppTracingConfigProviderUnavailableError) as caught, + ): + OpsTraceManagerGateway().require_provider_available("weave") - load_provider.assert_called_once_with("weave") + assert caught.value is unavailable - -def test_require_provider_available_reports_missing_dependency() -> None: - unavailable = AppTracingConfigProviderUnavailableError() - with ( - patch.object(OpsTraceManagerGateway, "_provider_config", side_effect=unavailable), - pytest.raises(AppTracingConfigProviderUnavailableError) as caught, - ): - OpsTraceManagerGateway().require_provider_available("weave") - - assert caught.value is unavailable - - -@pytest.mark.parametrize("update", [False, True]) -def test_prepare_config_reports_missing_provider_dependencies(monkeypatch: pytest.MonkeyPatch, update: bool) -> None: - missing_dependency = TraceProviderNotInstalledError("weave", "wandb") - providers = MagicMock() - providers.__getitem__.side_effect = missing_dependency - monkeypatch.setattr(gateway_module, "provider_config_map", providers) - gateway = OpsTraceManagerGateway() - prepare = ( - partial(gateway.prepare_updated_config, current_tracing_config={}) if update else gateway.prepare_new_config - ) - - with pytest.raises(AppTracingConfigProviderUnavailableError) as caught: - prepare(workspace_id="workspace-1", tracing_provider="weave", tracing_config={}) - - assert caught.value.__cause__ is missing_dependency - - -def test_present_config_reports_missing_provider_dependencies() -> None: - missing_dependency = TraceProviderNotInstalledError("weave", "wandb") - with patch.object(gateway_module, "OpsTraceManager") as manager: - manager.decrypt_tracing_config.side_effect = missing_dependency - - with pytest.raises(AppTracingConfigProviderUnavailableError) as caught: - OpsTraceManagerGateway().present_config( - workspace_id="workspace-1", tracing_provider="weave", tracing_config={"api_key": "encrypted"} - ) - - assert caught.value.__cause__ is missing_dependency - manager.obfuscated_decrypt_token.assert_not_called() - - -def test_prepare_new_config_applies_defaults_validates_and_encrypts(monkeypatch: pytest.MonkeyPatch) -> None: - monkeypatch.setattr( - gateway_module, - "provider_config_map", - {"langsmith": _provider_entry(other_keys=["endpoint", "project"])}, - ) - submitted = {"api_key": "plain", "endpoint": "", "project": ""} - - with patch.object(gateway_module, "OpsTraceManager") as manager: - manager.check_trace_config_is_effective.return_value = True - manager.get_trace_config_project_url.return_value = "https://project.example.com" - manager.encrypt_tracing_config.return_value = {"api_key": "encrypted"} - - result = OpsTraceManagerGateway().prepare_new_config( - workspace_id="workspace-1", - tracing_provider="langsmith", - tracing_config=submitted, + @staticmethod + @pytest.mark.parametrize("update", [False, True]) + def test_prepare_config_reports_missing_provider_dependencies( + monkeypatch: pytest.MonkeyPatch, update: bool + ) -> None: + missing_dependency = TraceProviderNotInstalledError("weave", "wandb") + providers = MagicMock() + providers.__getitem__.side_effect = missing_dependency + monkeypatch.setattr(gateway_module, "provider_config_map", providers) + gateway = OpsTraceManagerGateway() + prepare = ( + partial(gateway.prepare_updated_config, current_tracing_config={}) if update else gateway.prepare_new_config ) - normalized = { - "api_key": "plain", - "endpoint": "https://default.example.com", - "project": "default-project", - } - assert submitted == {"api_key": "plain", "endpoint": "", "project": ""} - assert result == {"api_key": "encrypted", "project_url": "https://project.example.com"} - manager.check_trace_config_is_effective.assert_called_once_with(normalized, "langsmith") - manager.encrypt_tracing_config.assert_called_once_with("workspace-1", "langsmith", normalized) + with pytest.raises(AppTracingConfigProviderUnavailableError) as caught: + prepare(workspace_id="workspace-1", tracing_provider="weave", tracing_config={}) + assert caught.value.__cause__ is missing_dependency -def test_prepare_new_config_reports_failed_verification_before_encryption( - monkeypatch: pytest.MonkeyPatch, -) -> None: - monkeypatch.setattr(gateway_module, "provider_config_map", {"langsmith": _provider_entry()}) + @staticmethod + def test_present_config_reports_missing_provider_dependencies() -> None: + missing_dependency = TraceProviderNotInstalledError("weave", "wandb") + with patch.object(gateway_module, "OpsTraceManager") as manager: + manager.decrypt_tracing_config.side_effect = missing_dependency - with patch.object(gateway_module, "OpsTraceManager") as manager: - manager.check_trace_config_is_effective.return_value = False + with pytest.raises(AppTracingConfigProviderUnavailableError) as caught: + OpsTraceManagerGateway().present_config( + workspace_id="workspace-1", tracing_provider="weave", tracing_config={"api_key": "encrypted"} + ) - with pytest.raises(AppTracingConfigVerificationFailedError): - OpsTraceManagerGateway().prepare_new_config( + assert caught.value.__cause__ is missing_dependency + manager.obfuscated_decrypt_token.assert_not_called() + + @staticmethod + def test_prepare_new_config_applies_defaults_validates_and_encrypts(monkeypatch: pytest.MonkeyPatch) -> None: + monkeypatch.setattr( + gateway_module, + "provider_config_map", + {"langsmith": _provider_entry(other_keys=["endpoint", "project"])}, + ) + submitted = {"api_key": "plain", "endpoint": "", "project": ""} + + with patch.object(gateway_module, "OpsTraceManager") as manager: + manager.check_trace_config_is_effective.return_value = True + manager.get_trace_config_project_url.return_value = "https://project.example.com" + manager.encrypt_tracing_config.return_value = {"api_key": "encrypted"} + + result = OpsTraceManagerGateway().prepare_new_config( + workspace_id="workspace-1", + tracing_provider="langsmith", + tracing_config=submitted, + ) + + normalized = { + "api_key": "plain", + "endpoint": "https://default.example.com", + "project": "default-project", + } + assert submitted == {"api_key": "plain", "endpoint": "", "project": ""} + assert result == {"api_key": "encrypted", "project_url": "https://project.example.com"} + manager.check_trace_config_is_effective.assert_called_once_with(normalized, "langsmith") + manager.encrypt_tracing_config.assert_called_once_with("workspace-1", "langsmith", normalized) + + @staticmethod + def test_prepare_new_config_reports_failed_verification_before_encryption( + monkeypatch: pytest.MonkeyPatch, + ) -> None: + monkeypatch.setattr(gateway_module, "provider_config_map", {"langsmith": _provider_entry()}) + + with patch.object(gateway_module, "OpsTraceManager") as manager: + manager.check_trace_config_is_effective.return_value = False + + with pytest.raises(AppTracingConfigVerificationFailedError): + OpsTraceManagerGateway().prepare_new_config( + workspace_id="workspace-1", + tracing_provider="langsmith", + tracing_config={"api_key": "plain"}, + ) + + manager.encrypt_tracing_config.assert_not_called() + + @staticmethod + def test_prepare_new_config_keeps_success_when_project_url_lookup_fails( + monkeypatch: pytest.MonkeyPatch, + ) -> None: + monkeypatch.setattr(gateway_module, "provider_config_map", {"langsmith": _provider_entry()}) + + with patch.object(gateway_module, "OpsTraceManager") as manager: + manager.check_trace_config_is_effective.return_value = True + manager.get_trace_config_project_url.side_effect = RuntimeError("provider unavailable") + manager.encrypt_tracing_config.return_value = {"api_key": "encrypted"} + + result = OpsTraceManagerGateway().prepare_new_config( workspace_id="workspace-1", tracing_provider="langsmith", tracing_config={"api_key": "plain"}, ) - manager.encrypt_tracing_config.assert_not_called() + assert result == {"api_key": "encrypted"} + @staticmethod + def test_prepare_new_langfuse_config_builds_project_url(monkeypatch: pytest.MonkeyPatch) -> None: + monkeypatch.setattr(gateway_module, "provider_config_map", {"langfuse": _provider_entry()}) -def test_prepare_new_config_keeps_success_when_project_url_lookup_fails( - monkeypatch: pytest.MonkeyPatch, -) -> None: - monkeypatch.setattr(gateway_module, "provider_config_map", {"langsmith": _provider_entry()}) + with patch.object(gateway_module, "OpsTraceManager") as manager: + manager.check_trace_config_is_effective.return_value = True + manager.get_trace_config_project_key.return_value = "project-key" + manager.encrypt_tracing_config.return_value = {"secret_key": "encrypted"} - with patch.object(gateway_module, "OpsTraceManager") as manager: - manager.check_trace_config_is_effective.return_value = True - manager.get_trace_config_project_url.side_effect = RuntimeError("provider unavailable") - manager.encrypt_tracing_config.return_value = {"api_key": "encrypted"} - - result = OpsTraceManagerGateway().prepare_new_config( - workspace_id="workspace-1", - tracing_provider="langsmith", - tracing_config={"api_key": "plain"}, - ) - - assert result == {"api_key": "encrypted"} - - -def test_prepare_new_langfuse_config_builds_project_url(monkeypatch: pytest.MonkeyPatch) -> None: - monkeypatch.setattr(gateway_module, "provider_config_map", {"langfuse": _provider_entry()}) - - with patch.object(gateway_module, "OpsTraceManager") as manager: - manager.check_trace_config_is_effective.return_value = True - manager.get_trace_config_project_key.return_value = "project-key" - manager.encrypt_tracing_config.return_value = {"secret_key": "encrypted"} - - result = OpsTraceManagerGateway().prepare_new_config( - workspace_id="workspace-1", - tracing_provider="langfuse", - tracing_config={"host": "https://langfuse.example.com"}, - ) - - assert result["project_url"] == "https://langfuse.example.com/project/project-key" - - -def test_prepare_new_config_reports_provider_check_exception_as_failed_verification( - monkeypatch: pytest.MonkeyPatch, -) -> None: - monkeypatch.setattr(gateway_module, "provider_config_map", {"langsmith": _provider_entry()}) - - with patch.object(gateway_module, "OpsTraceManager") as manager: - manager.check_trace_config_is_effective.side_effect = ValueError("verification failed") - - with pytest.raises(AppTracingConfigVerificationFailedError) as caught: - OpsTraceManagerGateway().prepare_new_config( + result = OpsTraceManagerGateway().prepare_new_config( workspace_id="workspace-1", - tracing_provider="langsmith", - tracing_config={}, + tracing_provider="langfuse", + tracing_config={"host": "https://langfuse.example.com"}, ) - assert isinstance(caught.value.__cause__, ValueError) + assert result["project_url"] == "https://langfuse.example.com/project/project-key" + @staticmethod + def test_prepare_new_config_reports_provider_check_exception_as_failed_verification( + monkeypatch: pytest.MonkeyPatch, + ) -> None: + monkeypatch.setattr(gateway_module, "provider_config_map", {"langsmith": _provider_entry()}) -def test_prepare_new_config_propagates_unexpected_verification_failure( - monkeypatch: pytest.MonkeyPatch, -) -> None: - monkeypatch.setattr(gateway_module, "provider_config_map", {"langsmith": _provider_entry()}) - failure = RuntimeError("unexpected provider bug") + with patch.object(gateway_module, "OpsTraceManager") as manager: + manager.check_trace_config_is_effective.side_effect = ValueError("verification failed") - with patch.object(gateway_module, "OpsTraceManager") as manager: - manager.check_trace_config_is_effective.side_effect = failure + with pytest.raises(AppTracingConfigVerificationFailedError) as caught: + OpsTraceManagerGateway().prepare_new_config( + workspace_id="workspace-1", + tracing_provider="langsmith", + tracing_config={}, + ) - with pytest.raises(RuntimeError) as caught: - OpsTraceManagerGateway().prepare_new_config( + assert isinstance(caught.value.__cause__, ValueError) + + @staticmethod + def test_prepare_new_config_propagates_unexpected_verification_failure( + monkeypatch: pytest.MonkeyPatch, + ) -> None: + monkeypatch.setattr(gateway_module, "provider_config_map", {"langsmith": _provider_entry()}) + failure = RuntimeError("unexpected provider bug") + + with patch.object(gateway_module, "OpsTraceManager") as manager: + manager.check_trace_config_is_effective.side_effect = failure + + with pytest.raises(RuntimeError) as caught: + OpsTraceManagerGateway().prepare_new_config( + workspace_id="workspace-1", + tracing_provider="langsmith", + tracing_config={}, + ) + + assert caught.value is failure + manager.encrypt_tracing_config.assert_not_called() + + @staticmethod + def test_prepare_new_config_rejects_invalid_schema_before_verification( + monkeypatch: pytest.MonkeyPatch, + ) -> None: + monkeypatch.setattr(gateway_module, "provider_config_map", {"langsmith": _provider_entry()}) + + with patch.object(gateway_module, "OpsTraceManager") as manager: + with pytest.raises(AppTracingConfigInvalidConfigurationError): + OpsTraceManagerGateway().prepare_new_config( + workspace_id="workspace-1", + tracing_provider="langsmith", + tracing_config={"endpoint": {"invalid": "value"}}, + ) + + manager.check_trace_config_is_effective.assert_not_called() + manager.encrypt_tracing_config.assert_not_called() + + @staticmethod + def test_prepare_new_config_reports_encryption_failure_as_processing_error(monkeypatch: pytest.MonkeyPatch) -> None: + monkeypatch.setattr(gateway_module, "provider_config_map", {"langsmith": _provider_entry()}) + failure = RuntimeError("key provider unavailable") + + with patch.object(gateway_module, "OpsTraceManager") as manager: + manager.check_trace_config_is_effective.return_value = True + manager.get_trace_config_project_url.return_value = None + manager.encrypt_tracing_config.side_effect = failure + + with pytest.raises(AppTracingConfigProcessingError) as caught: + OpsTraceManagerGateway().prepare_new_config( + workspace_id="workspace-1", + tracing_provider="langsmith", + tracing_config={}, + ) + + assert caught.value.__cause__ is failure + + @staticmethod + def test_prepare_updated_config_preserves_masked_secret_from_current_config( + monkeypatch: pytest.MonkeyPatch, + ) -> None: + monkeypatch.setattr(gateway_module, "provider_config_map", {"langsmith": _provider_entry()}) + submitted = {"api_key": "******", "project": "new-project"} + current = {"api_key": "old-encrypted", "project": "old-project"} + encrypted = {"api_key": "old-encrypted", "project": "new-project"} + decrypted = {"api_key": "old-plain", "project": "new-project"} + + with patch.object(gateway_module, "OpsTraceManager") as manager: + manager.encrypt_tracing_config.return_value = encrypted + manager.decrypt_tracing_config.return_value = decrypted + manager.check_trace_config_is_effective.return_value = True + manager.get_trace_config_project_url.return_value = "https://example.com/project/new" + + result = OpsTraceManagerGateway().prepare_updated_config( workspace_id="workspace-1", tracing_provider="langsmith", - tracing_config={}, + tracing_config=submitted, + current_tracing_config=current, ) - assert caught.value is failure - manager.encrypt_tracing_config.assert_not_called() + assert result == encrypted + assert result["project_url"] == "https://example.com/project/new" + assert submitted == {"api_key": "******", "project": "new-project"} + manager.encrypt_tracing_config.assert_called_once_with("workspace-1", "langsmith", submitted, current) + manager.decrypt_tracing_config.assert_called_once_with("workspace-1", "langsmith", encrypted) + manager.check_trace_config_is_effective.assert_called_once_with(decrypted, "langsmith") + manager.get_trace_config_project_url.assert_called_once_with(decrypted, "langsmith") - -def test_prepare_new_config_rejects_invalid_schema_before_verification( - monkeypatch: pytest.MonkeyPatch, -) -> None: - monkeypatch.setattr(gateway_module, "provider_config_map", {"langsmith": _provider_entry()}) - - with patch.object(gateway_module, "OpsTraceManager") as manager: - with pytest.raises(AppTracingConfigInvalidConfigurationError): - OpsTraceManagerGateway().prepare_new_config( - workspace_id="workspace-1", - tracing_provider="langsmith", - tracing_config={"endpoint": {"invalid": "value"}}, - ) - - manager.check_trace_config_is_effective.assert_not_called() - manager.encrypt_tracing_config.assert_not_called() - - -def test_prepare_new_config_reports_encryption_failure_as_processing_error(monkeypatch: pytest.MonkeyPatch) -> None: - monkeypatch.setattr(gateway_module, "provider_config_map", {"langsmith": _provider_entry()}) - failure = RuntimeError("key provider unavailable") - - with patch.object(gateway_module, "OpsTraceManager") as manager: - manager.check_trace_config_is_effective.return_value = True - manager.get_trace_config_project_url.return_value = None - manager.encrypt_tracing_config.side_effect = failure - - with pytest.raises(AppTracingConfigProcessingError) as caught: - OpsTraceManagerGateway().prepare_new_config( - workspace_id="workspace-1", - tracing_provider="langsmith", - tracing_config={}, - ) - - assert caught.value.__cause__ is failure - - -def test_prepare_updated_config_preserves_masked_secret_from_current_config( - monkeypatch: pytest.MonkeyPatch, -) -> None: - monkeypatch.setattr(gateway_module, "provider_config_map", {"langsmith": _provider_entry()}) - submitted = {"api_key": "******", "project": "new-project"} - current = {"api_key": "old-encrypted", "project": "old-project"} - encrypted = {"api_key": "old-encrypted", "project": "new-project"} - decrypted = {"api_key": "old-plain", "project": "new-project"} - - with patch.object(gateway_module, "OpsTraceManager") as manager: - manager.encrypt_tracing_config.return_value = encrypted - manager.decrypt_tracing_config.return_value = decrypted - manager.check_trace_config_is_effective.return_value = True - manager.get_trace_config_project_url.return_value = "https://example.com/project/new" - - result = OpsTraceManagerGateway().prepare_updated_config( - workspace_id="workspace-1", - tracing_provider="langsmith", - tracing_config=submitted, - current_tracing_config=current, - ) - - assert result == encrypted - assert result["project_url"] == "https://example.com/project/new" - assert submitted == {"api_key": "******", "project": "new-project"} - manager.encrypt_tracing_config.assert_called_once_with("workspace-1", "langsmith", submitted, current) - manager.decrypt_tracing_config.assert_called_once_with("workspace-1", "langsmith", encrypted) - manager.check_trace_config_is_effective.assert_called_once_with(decrypted, "langsmith") - manager.get_trace_config_project_url.assert_called_once_with(decrypted, "langsmith") - - -def test_updated_destination_drops_old_project_when_lookup_fails(monkeypatch: pytest.MonkeyPatch) -> None: - monkeypatch.setattr(gateway_module, "provider_config_map", {"langsmith": _provider_entry()}) - with patch.object(gateway_module, "OpsTraceManager") as manager: - manager.encrypt_tracing_config.return_value = {"api_key": "encrypted", "project": "new-project"} - manager.decrypt_tracing_config.return_value = {"api_key": "plain", "project": "new-project"} - manager.check_trace_config_is_effective.return_value = True - manager.get_trace_config_project_url.side_effect = RuntimeError("metadata unavailable") - result = OpsTraceManagerGateway().prepare_updated_config( - workspace_id="workspace-1", - tracing_provider="langsmith", - tracing_config={"project": "new-project"}, - current_tracing_config={"project_url": "https://example.com/project/old", "api_key": "encrypted"}, - ) - assert "project_url" not in result - - -def test_prepare_updated_config_validates_schema_before_encryption(monkeypatch: pytest.MonkeyPatch) -> None: - monkeypatch.setattr(gateway_module, "provider_config_map", {"langsmith": _provider_entry()}) - - with patch.object(gateway_module, "OpsTraceManager") as manager: - with pytest.raises(AppTracingConfigInvalidConfigurationError): - OpsTraceManagerGateway().prepare_updated_config( - workspace_id="workspace-1", - tracing_provider="langsmith", - tracing_config={"endpoint": {"invalid": "value"}}, - current_tracing_config={"api_key": "old-encrypted"}, - ) - - manager.encrypt_tracing_config.assert_not_called() - manager.decrypt_tracing_config.assert_not_called() - manager.check_trace_config_is_effective.assert_not_called() - - -def test_prepare_updated_config_reports_decryption_failure_as_processing_error( - monkeypatch: pytest.MonkeyPatch, -) -> None: - monkeypatch.setattr(gateway_module, "provider_config_map", {"langsmith": _provider_entry()}) - failure = RuntimeError("stored credential cannot be decrypted") - - with patch.object(gateway_module, "OpsTraceManager") as manager: - manager.encrypt_tracing_config.return_value = {"api_key": "encrypted"} - manager.decrypt_tracing_config.side_effect = failure - - with pytest.raises(AppTracingConfigProcessingError) as caught: - OpsTraceManagerGateway().prepare_updated_config( + @staticmethod + def test_updated_destination_drops_old_project_when_lookup_fails(monkeypatch: pytest.MonkeyPatch) -> None: + monkeypatch.setattr(gateway_module, "provider_config_map", {"langsmith": _provider_entry()}) + with patch.object(gateway_module, "OpsTraceManager") as manager: + manager.encrypt_tracing_config.return_value = {"api_key": "encrypted", "project": "new-project"} + manager.decrypt_tracing_config.return_value = {"api_key": "plain", "project": "new-project"} + manager.check_trace_config_is_effective.return_value = True + manager.get_trace_config_project_url.side_effect = RuntimeError("metadata unavailable") + result = OpsTraceManagerGateway().prepare_updated_config( workspace_id="workspace-1", tracing_provider="langsmith", tracing_config={"project": "new-project"}, - current_tracing_config={"api_key": "old-encrypted"}, + current_tracing_config={"project_url": "https://example.com/project/old", "api_key": "encrypted"}, ) + assert "project_url" not in result - assert caught.value.__cause__ is failure - manager.check_trace_config_is_effective.assert_not_called() + @staticmethod + def test_prepare_updated_config_validates_schema_before_encryption(monkeypatch: pytest.MonkeyPatch) -> None: + monkeypatch.setattr(gateway_module, "provider_config_map", {"langsmith": _provider_entry()}) + with patch.object(gateway_module, "OpsTraceManager") as manager: + with pytest.raises(AppTracingConfigInvalidConfigurationError): + OpsTraceManagerGateway().prepare_updated_config( + workspace_id="workspace-1", + tracing_provider="langsmith", + tracing_config={"endpoint": {"invalid": "value"}}, + current_tracing_config={"api_key": "old-encrypted"}, + ) -@pytest.mark.parametrize( - ("provider", "fallback_url"), - [ - ("arize", None), - ("phoenix", "https://app.phoenix.arize.com/redirects/projects/default"), - ("langsmith", "https://smith.langchain.com/"), - ("opik", "https://www.comet.com/opik/"), - ("weave", "https://wandb.ai/"), - ("aliyun", "https://arms.console.aliyun.com/#/llm"), - ("tencent", "https://console.cloud.tencent.com/apm"), - ("mlflow", "http://localhost:5000/#/experiments/0/traces"), - ("databricks", "https://www.databricks.com/"), - ], -) -def test_present_config_uses_provider_fallback_when_project_lookup_fails( - provider: str, - fallback_url: str | None, -) -> None: - with patch.object(gateway_module, "OpsTraceManager") as manager: - decrypted_config: dict[str, object] = {} - presented_config: dict[str, object] = {} - manager.decrypt_tracing_config.return_value = decrypted_config - manager.obfuscated_decrypt_token.return_value = presented_config - manager.get_trace_config_project_url.side_effect = RuntimeError("provider unavailable") + manager.encrypt_tracing_config.assert_not_called() + manager.decrypt_tracing_config.assert_not_called() + manager.check_trace_config_is_effective.assert_not_called() - result = OpsTraceManagerGateway().present_config( - workspace_id="workspace-1", - tracing_provider=provider, - tracing_config={"encrypted": "config"}, - ) + @staticmethod + def test_prepare_updated_config_reports_decryption_failure_as_processing_error( + monkeypatch: pytest.MonkeyPatch, + ) -> None: + monkeypatch.setattr(gateway_module, "provider_config_map", {"langsmith": _provider_entry()}) + failure = RuntimeError("stored credential cannot be decrypted") - assert result == {"project_url": fallback_url} - if provider in {"arize", "aliyun", "mlflow", "phoenix"}: - # Provider defaults produce a pure link; the failing legacy lookup is - # deliberately not reached, even when no cached URL is present. - manager.get_trace_config_project_url.assert_not_called() + with patch.object(gateway_module, "OpsTraceManager") as manager: + manager.encrypt_tracing_config.return_value = {"api_key": "encrypted"} + manager.decrypt_tracing_config.side_effect = failure + with pytest.raises(AppTracingConfigProcessingError) as caught: + OpsTraceManagerGateway().prepare_updated_config( + workspace_id="workspace-1", + tracing_provider="langsmith", + tracing_config={"project": "new-project"}, + current_tracing_config={"api_key": "old-encrypted"}, + ) -def test_present_langfuse_config_builds_project_url() -> None: - with patch.object(gateway_module, "OpsTraceManager") as manager: - manager.decrypt_tracing_config.return_value = {"host": "https://langfuse.example.com"} - manager.obfuscated_decrypt_token.return_value = {"host": "https://langfuse.example.com"} - manager.get_trace_config_project_key.return_value = "project-key" + assert caught.value.__cause__ is failure + manager.check_trace_config_is_effective.assert_not_called() - result = OpsTraceManagerGateway().present_config( - workspace_id="workspace-1", - tracing_provider="langfuse", - tracing_config={"encrypted": "config"}, - ) + @staticmethod + @pytest.mark.parametrize( + ("provider", "fallback_url"), + [ + ("arize", None), + ("phoenix", "https://app.phoenix.arize.com/redirects/projects/default"), + ("langsmith", "https://smith.langchain.com/"), + ("opik", "https://www.comet.com/opik/"), + ("weave", "https://wandb.ai/"), + ("aliyun", "https://arms.console.aliyun.com/#/llm"), + ("tencent", "https://console.cloud.tencent.com/apm"), + ("mlflow", "http://localhost:5000/#/experiments/0/traces"), + ("databricks", "https://www.databricks.com/"), + ], + ) + def test_present_config_uses_provider_fallback_when_project_lookup_fails( + provider: str, + fallback_url: str | None, + ) -> None: + with patch.object(gateway_module, "OpsTraceManager") as manager: + decrypted_config: dict[str, object] = {} + presented_config: dict[str, object] = {} + manager.decrypt_tracing_config.return_value = decrypted_config + manager.obfuscated_decrypt_token.return_value = presented_config + manager.get_trace_config_project_url.side_effect = RuntimeError("provider unavailable") - assert result["project_url"] == "https://langfuse.example.com/project/project-key" - - -def test_present_langfuse_config_falls_back_to_host() -> None: - with patch.object(gateway_module, "OpsTraceManager") as manager: - manager.decrypt_tracing_config.return_value = {"host": "https://langfuse.example.com"} - manager.obfuscated_decrypt_token.return_value = {"host": "https://langfuse.example.com"} - manager.get_trace_config_project_key.side_effect = RuntimeError("provider unavailable") - - result = OpsTraceManagerGateway().present_config( - workspace_id="workspace-1", - tracing_provider="langfuse", - tracing_config={"encrypted": "config"}, - ) - - assert result["project_url"] == "https://langfuse.example.com/" - - -def test_present_config_rejects_missing_stored_config() -> None: - with pytest.raises(AppTracingConfigProcessingError, match="processing failed"): - OpsTraceManagerGateway().present_config( - workspace_id="workspace-1", - tracing_provider="langsmith", - tracing_config=None, - ) - - -def test_present_config_reports_decryption_failure_as_processing_error() -> None: - failure = RuntimeError("stored credential cannot be decrypted") - - with patch.object(gateway_module, "OpsTraceManager") as manager: - manager.decrypt_tracing_config.side_effect = failure - - with pytest.raises(AppTracingConfigProcessingError) as caught: - OpsTraceManagerGateway().present_config( + result = OpsTraceManagerGateway().present_config( workspace_id="workspace-1", - tracing_provider="langsmith", + tracing_provider=provider, tracing_config={"encrypted": "config"}, ) - assert caught.value.__cause__ is failure - manager.obfuscated_decrypt_token.assert_not_called() + assert result == {"project_url": fallback_url} + if provider in {"arize", "aliyun", "mlflow", "phoenix"}: + # Provider defaults produce a pure link; the failing legacy lookup is + # deliberately not reached, even when no cached URL is present. + manager.get_trace_config_project_url.assert_not_called() + @staticmethod + def test_present_langfuse_config_builds_project_url() -> None: + with patch.object(gateway_module, "OpsTraceManager") as manager: + manager.decrypt_tracing_config.return_value = {"host": "https://langfuse.example.com"} + manager.obfuscated_decrypt_token.return_value = {"host": "https://langfuse.example.com"} + manager.get_trace_config_project_key.return_value = "project-key" -@pytest.mark.parametrize("malformed", [False, True]) -def test_validate_saved_config_decrypts_and_validates_without_vendor_io( - monkeypatch: pytest.MonkeyPatch, malformed: bool -) -> None: - monkeypatch.setattr(gateway_module, "provider_config_map", {"langsmith": _provider_entry()}) - with patch.object(gateway_module, "OpsTraceManager") as manager: - manager.decrypt_tracing_config.return_value = {"project": ["invalid-project"] if malformed else "saved-project"} - gateway = OpsTraceManagerGateway() - if malformed: - with pytest.raises(AppTracingConfigInvalidConfigurationError): + result = OpsTraceManagerGateway().present_config( + workspace_id="workspace-1", + tracing_provider="langfuse", + tracing_config={"encrypted": "config"}, + ) + + assert result["project_url"] == "https://langfuse.example.com/project/project-key" + + @staticmethod + def test_present_langfuse_config_falls_back_to_host() -> None: + with patch.object(gateway_module, "OpsTraceManager") as manager: + manager.decrypt_tracing_config.return_value = {"host": "https://langfuse.example.com"} + manager.obfuscated_decrypt_token.return_value = {"host": "https://langfuse.example.com"} + manager.get_trace_config_project_key.side_effect = RuntimeError("provider unavailable") + + result = OpsTraceManagerGateway().present_config( + workspace_id="workspace-1", + tracing_provider="langfuse", + tracing_config={"encrypted": "config"}, + ) + + assert result["project_url"] == "https://langfuse.example.com/" + + @staticmethod + def test_present_config_rejects_missing_stored_config() -> None: + with pytest.raises(AppTracingConfigProcessingError, match="processing failed"): + OpsTraceManagerGateway().present_config( + workspace_id="workspace-1", + tracing_provider="langsmith", + tracing_config=None, + ) + + @staticmethod + def test_present_config_reports_decryption_failure_as_processing_error() -> None: + failure = RuntimeError("stored credential cannot be decrypted") + + with patch.object(gateway_module, "OpsTraceManager") as manager: + manager.decrypt_tracing_config.side_effect = failure + + with pytest.raises(AppTracingConfigProcessingError) as caught: + OpsTraceManagerGateway().present_config( + workspace_id="workspace-1", + tracing_provider="langsmith", + tracing_config={"encrypted": "config"}, + ) + + assert caught.value.__cause__ is failure + manager.obfuscated_decrypt_token.assert_not_called() + + @staticmethod + @pytest.mark.parametrize("malformed", [False, True]) + def test_validate_saved_config_decrypts_and_validates_without_vendor_io( + monkeypatch: pytest.MonkeyPatch, malformed: bool + ) -> None: + monkeypatch.setattr(gateway_module, "provider_config_map", {"langsmith": _provider_entry()}) + with patch.object(gateway_module, "OpsTraceManager") as manager: + manager.decrypt_tracing_config.return_value = { + "project": ["invalid-project"] if malformed else "saved-project" + } + gateway = OpsTraceManagerGateway() + if malformed: + with pytest.raises(AppTracingConfigInvalidConfigurationError): + gateway.validate_saved_config( + workspace_id="workspace", tracing_provider="langsmith", tracing_config={"api_key": "encrypted"} + ) + else: gateway.validate_saved_config( workspace_id="workspace", tracing_provider="langsmith", tracing_config={"api_key": "encrypted"} ) + manager.decrypt_tracing_config.assert_called_once_with("workspace", "langsmith", {"api_key": "encrypted"}) + manager.check_trace_config_is_effective.assert_not_called() + + @staticmethod + def test_validate_saved_config_rejects_undecryptable_credentials(monkeypatch: pytest.MonkeyPatch) -> None: + monkeypatch.setattr(gateway_module, "provider_config_map", {"langsmith": _provider_entry()}) + with patch.object(gateway_module, "OpsTraceManager") as manager: + manager.decrypt_tracing_config.side_effect = ValueError("invalid ciphertext") + with pytest.raises(AppTracingConfigProcessingError): + OpsTraceManagerGateway().validate_saved_config( + workspace_id="workspace", tracing_provider="langsmith", tracing_config={"api_key": "broken"} + ) + + +"""Real provider config serialization must retain verified project links on update.""" + + +AccessAPI = _test_importlib.import_module( + "tests.unit_tests.providers.trace.trace-opik.src.dify_trace_opik.test_configuration" +).AccessAPI + + +@dataclass +class ProviderBoundary: + calls: list[tuple[str, str, dict[str, object]]] = field(default_factory=list) + effective: bool = True + lookup_fails: bool = False + + def verify(self, config: Mapping[str, object], provider: str) -> bool: + self.calls.append(("verify", provider, dict(config))) + return self.effective + + def project_url(self, config: Mapping[str, object], provider: str) -> str: + self.calls.append(("url", provider, dict(config))) + if self.lookup_fails: + raise RuntimeError("Synthetic project lookup failure") + return f"https://tracing.test/{provider}/{config['project']}" + + def project_key(self, config: Mapping[str, object], provider: str) -> str: + self.calls.append(("key", provider, dict(config))) + return "verified-project" + + +class TestAppTracingConfigProjectUrlCases: + @staticmethod + @pytest.fixture + def boundary(monkeypatch: pytest.MonkeyPatch) -> ProviderBoundary: + boundary = ProviderBoundary() + + def encrypt(_tenant: str, value: str) -> str: + return "test-cipher:" + value + + def decrypt(_tenant: str, values: Sequence[str]) -> list[str]: + assert all(value.startswith("test-cipher:") for value in values) + return [value.removeprefix("test-cipher:") for value in values] + + monkeypatch.setattr(ops_trace_manager, "encrypt_token", encrypt) + monkeypatch.setattr(ops_trace_manager, "batch_decrypt_token", decrypt) + monkeypatch.setattr(OpsTraceManager, "check_trace_config_is_effective", boundary.verify) + monkeypatch.setattr(OpsTraceManager, "get_trace_config_project_url", boundary.project_url) + monkeypatch.setattr(OpsTraceManager, "get_trace_config_project_key", boundary.project_key) + return boundary + + @staticmethod + @pytest.mark.parametrize("provider", ["langsmith", "opik"]) + @pytest.mark.parametrize("project", ["original", "changed"]) + def test_update_resolves_project_link_using_verified_unmasked_configuration( + boundary: ProviderBoundary, provider: str, project: str, monkeypatch: pytest.MonkeyPatch + ) -> None: + gateway = OpsTraceManagerGateway() + workspace = str(uuid4()) + api = AccessAPI() + if provider == "opik": + monkeypatch.setattr(ssrf_proxy, "make_request", api.request) + target = {"workspace": "team", "url": "https://tracing.test/opik/api/"} if provider == "opik" else {} + current = gateway.prepare_new_config( + workspace_id=workspace, + tracing_provider=provider, + tracing_config={"api_key": "synthetic-key", "project": "original", **target}, + ) + submitted = {"api_key": "******", "project": project, "project_url": "https://untrusted.test/stale", **target} + updated = gateway.prepare_updated_config( + workspace_id=workspace, tracing_provider=provider, tracing_config=submitted, current_tracing_config=current + ) + + prefix = f"https://tracing.test/{provider}/" + ("team/redirect/projects?name=" if provider == "opik" else "") + assert updated["project_url"] == prefix + project + assert updated["api_key"] == current["api_key"] == "test-cipher:synthetic-key" + assert current["project_url"] == prefix + "original" + assert submitted["api_key"] == "******" + if provider == "opik": + assert boundary.calls == [] + assert len(api.calls) == 2 + assert api.calls[-1][2]["headers"] == { + "Authorization": "synthetic-key", + "Comet-Workspace": "team", + "Accept": "application/json", + } else: - gateway.validate_saved_config( - workspace_id="workspace", tracing_provider="langsmith", tracing_config={"api_key": "encrypted"} + assert [call[0] for call in boundary.calls] == ["verify", "url", "verify", "url"] + assert boundary.calls[-1][2]["api_key"] == "synthetic-key" + assert boundary.calls[-1][2]["project"] == project + + @staticmethod + def test_langfuse_update_uses_current_host_and_resolved_project_key(boundary: ProviderBoundary) -> None: + gateway = OpsTraceManagerGateway() + workspace = str(uuid4()) + current = gateway.prepare_new_config( + workspace_id=workspace, + tracing_provider="langfuse", + tracing_config={ + "public_key": "synthetic-public", + "secret_key": "synthetic-secret", + "host": "https://old.test", + }, + ) + updated = gateway.prepare_updated_config( + workspace_id=workspace, + tracing_provider="langfuse", + tracing_config={"public_key": "******", "secret_key": "******", "host": "https://new.test"}, + current_tracing_config=current, + ) + assert updated["project_url"] == "https://new.test/project/verified-project" + assert current["project_url"] == "https://old.test/project/verified-project" + assert boundary.calls[-1][0] == "key" + assert boundary.calls[-1][2]["secret_key"] == "synthetic-secret" + + @staticmethod + def test_failed_optional_link_lookup_does_not_reuse_stale_project_url(boundary: ProviderBoundary) -> None: + gateway = OpsTraceManagerGateway() + workspace = str(uuid4()) + current = gateway.prepare_new_config( + workspace_id=workspace, + tracing_provider="langsmith", + tracing_config={"api_key": "synthetic-key", "project": "original"}, + ) + boundary.lookup_fails = True + updated = gateway.prepare_updated_config( + workspace_id=workspace, + tracing_provider="langsmith", + tracing_config={"api_key": "******", "project": "changed"}, + current_tracing_config=current, + ) + assert updated["project"] == "changed" + assert "project_url" not in updated + assert boundary.calls[-1][0] == "url" + + @staticmethod + def test_invalid_updated_configuration_never_resolves_a_project_url(boundary: ProviderBoundary) -> None: + boundary.effective = False + with pytest.raises(AppTracingConfigVerificationFailedError): + OpsTraceManagerGateway().prepare_updated_config( + workspace_id=str(uuid4()), + tracing_provider="langsmith", + tracing_config={"api_key": "synthetic-key", "project": "changed"}, + current_tracing_config=None, ) - manager.decrypt_tracing_config.assert_called_once_with("workspace", "langsmith", {"api_key": "encrypted"}) - manager.check_trace_config_is_effective.assert_not_called() + assert [call[0] for call in boundary.calls] == ["verify"] + + @staticmethod + def test_updated_project_url_is_frozen_and_used_by_real_langsmith_serializer( + boundary: ProviderBoundary, + sqlite_session_factory: sessionmaker[Session], + monkeypatch: pytest.MonkeyPatch, + ) -> None: + workspace, app_id, message_id = (str(uuid4()) for _ in range(3)) + with sqlite_session_factory.begin() as session: + session.add( + make_app( + app_id=app_id, + tenant_id=workspace, + mode=AppMode.AGENT, + tracing=json.dumps({"enabled": True, "tracing_provider": "langsmith"}), + ) + ) + session.add( + make_message( + message_id=message_id, + app_id=app_id, + inputs={}, + query="test input", + message={}, + answer="test output", + message_unit_price=Decimal(0), + answer_unit_price=Decimal(0), + currency="USD", + from_source=ConversationFromSource.CONSOLE, + ) + ) + store = SQLAlchemyAppTracingConfigRepository(session_factory=sqlite_session_factory) + configs = AppTracingConfigService(configs=store, provider=OpsTraceManagerGateway()) + context = RequestContext("request", None, "account", workspace) + configs.create(context, app_id, "langsmith", {"api_key": "synthetic-key", "project": "original"}) + configs.update(context, app_id, "langsmith", {"api_key": "******", "project": "changed"}) + repository = ApplicationTraceRepository(sqlite_session_factory) + run = repository.start( + tenant_id=workspace, app_id=app_id, source_kind="message", source_id=message_id, session_id="test-session" + ) + assert run is not None + assert run.provider_config["project_url"] == "https://tracing.test/langsmith/changed" + configs.update(context, app_id, "langsmith", {"api_key": "******", "project": "later"}) + frozen = repository.source_run(tenant_id=workspace, app_id=app_id, source_kind="message", source_id=message_id) + assert frozen is not None + assert frozen.provider_config["project_url"] == "https://tracing.test/langsmith/changed" + + requests: list[str] = [] + + def send(method: str, url: str, **_kwargs: object) -> httpx.Response: + requests.append(url) + return httpx.Response(200, request=httpx.Request(method, url)) + + monkeypatch.setattr(ssrf_proxy, "make_request", send) + config = dict(frozen.provider_config) + config.update(OpsTraceManager.decrypt_tracing_config(workspace, "langsmith", frozen.provider_config)) + start = datetime(2026, 10, 16, tzinfo=UTC) + trace = CanonicalTrace( + trace_id=frozen.trace_id, + session_id=frozen.session_id, + root_span_id=frozen.root_span_id, + spans=( + CanonicalSpan( + id=frozen.root_span_id, + name="Agent", + kind=CanonicalSpanKind.AGENT, + parent_id=None, + status=CanonicalSpanStatus.OK, + start_time=start, + end_time=start + timedelta(seconds=1), + ), + ), + ) + result = export_trace("langsmith", config, trace) + assert result.status == "accepted" + assert result.trace_url == f"https://tracing.test/langsmith/changed/r/{result.trace_id}" + assert result.span_urls[frozen.root_span_id] == result.trace_url + assert len(requests) == 1 + assert boundary.calls[-1][2]["project"] == "later" -def test_validate_saved_config_rejects_undecryptable_credentials(monkeypatch: pytest.MonkeyPatch) -> None: - monkeypatch.setattr(gateway_module, "provider_config_map", {"langsmith": _provider_entry()}) - with patch.object(gateway_module, "OpsTraceManager") as manager: - manager.decrypt_tracing_config.side_effect = ValueError("invalid ciphertext") - with pytest.raises(AppTracingConfigProcessingError): - OpsTraceManagerGateway().validate_saved_config( - workspace_id="workspace", tracing_provider="langsmith", tracing_config={"api_key": "broken"} +"""Arize project links use verified metadata without changing ingestion permissions.""" + + +DESTINATION = _test_importlib.import_module( + "tests.unit_tests.providers.trace.trace-arize-phoenix.src.dify_trace_arize_phoenix.test_arize_configuration" +).DESTINATION +SPACE = _test_importlib.import_module( + "tests.unit_tests.providers.trace.trace-arize-phoenix.src.dify_trace_arize_phoenix.test_arize_configuration" +).SPACE +URL = _test_importlib.import_module( + "tests.unit_tests.providers.trace.trace-arize-phoenix.src.dify_trace_arize_phoenix.test_arize_configuration" +).URL +VALUES = _test_importlib.import_module( + "tests.unit_tests.providers.trace.trace-arize-phoenix.src.dify_trace_arize_phoenix.test_arize_configuration" +).VALUES +page = _test_importlib.import_module( + "tests.unit_tests.providers.trace.trace-arize-phoenix.src.dify_trace_arize_phoenix.test_arize_configuration" +).page +project = _test_importlib.import_module( + "tests.unit_tests.providers.trace.trace-arize-phoenix.src.dify_trace_arize_phoenix.test_arize_configuration" +).project +_arize_tracing_project_url_response = _test_importlib.import_module( + "tests.unit_tests.providers.trace.trace-arize-phoenix.src.dify_trace_arize_phoenix.test_arize_configuration" +).response +space = _test_importlib.import_module( + "tests.unit_tests.providers.trace.trace-arize-phoenix.src.dify_trace_arize_phoenix.test_arize_configuration" +).space +_D0db4e12TestArizeConfigurationCases = _test_importlib.import_module( + "tests.unit_tests.providers.trace.trace-arize-phoenix.src.dify_trace_arize_phoenix.test_arize_configuration" +).TestArizeConfigurationCases +_D0db4e12TestArizeConfigurationCases = _test_importlib.import_module( + "tests.unit_tests.providers.trace.trace-arize-phoenix.src.dify_trace_arize_phoenix.test_arize_configuration" +).TestArizeConfigurationCases + + +class TestArizeTracingProjectUrlCases: + calls = staticmethod(_D0db4e12TestArizeConfigurationCases.calls) + + responses = staticmethod(_D0db4e12TestArizeConfigurationCases.responses) + + @staticmethod + @pytest.fixture + def verified(monkeypatch: pytest.MonkeyPatch) -> list[dict[str, object]]: + calls: list[dict[str, object]] = [] + + def verify(config: ArizeConfig) -> None: + calls.append(config.model_dump()) + + monkeypatch.setattr(arize_configuration, "verify_configuration", verify) + return calls + + @staticmethod + def test_legacy_arize_get_hides_unverified_redirect_without_network_or_rewriting( + workspace: str, monkeypatch: pytest.MonkeyPatch + ) -> None: + stored = OpsTraceManager.encrypt_tracing_config( + workspace, + "arize", + {"api_key": "synthetic-ingestion-key", "space_id": "U3BhY2U6dGVzdA==", "project": "original"}, + ) + stored["project_url"] = "https://app.arize.com/?redirect_project_name=original" + untouched = deepcopy(stored) + calls: list[str] = [] + + def forbidden(*_args: object, **_kwargs: object) -> Never: + calls.append("legacy or network") + raise AssertionError("Configuration GET must not initialize a tracer or resolve vendor metadata") + + monkeypatch.setattr(OpsTraceManager, "get_trace_config_project_url", forbidden) + monkeypatch.setattr(ssrf_proxy, "get", forbidden) + shown = OpsTraceManagerGateway().present_config( + workspace_id=workspace, tracing_provider="arize", tracing_config=stored + ) + + assert shown["project_url"] is None + assert calls == [] + assert stored == untouched + + @staticmethod + def test_real_rsa_create_masked_update_get_and_frozen_history( + workspace: str, + sqlite_session_factory: sessionmaker[Session], + monkeypatch: pytest.MonkeyPatch, + verified: list[dict[str, object]], + calls: list[tuple[str, dict[str, object]]], + responses: deque[httpx.Response | Exception], + ) -> None: + app_id, message_id = str(uuid4()), str(uuid4()) + with sqlite_session_factory.begin() as session: + session.add( + make_app( + app_id=app_id, + tenant_id=workspace, + mode=AppMode.AGENT, + tracing=json.dumps({"enabled": True, "tracing_provider": "arize"}), + ) ) + session.add( + make_message( + message_id=message_id, + app_id=app_id, + inputs={}, + query="test", + message={}, + answer="ok", + message_unit_price=Decimal(0), + answer_unit_price=Decimal(0), + currency="USD", + from_source=ConversationFromSource.CONSOLE, + ) + ) + store = SQLAlchemyAppTracingConfigRepository(session_factory=sqlite_session_factory) + gateway = OpsTraceManagerGateway() + service = AppTracingConfigService(configs=store, provider=gateway) + context = RequestContext("request", None, "account", workspace) + legacy = OpsTraceManager.encrypt_tracing_config(workspace, "arize", VALUES) + legacy["project_url"] = "https://app.arize.com/?redirect_project_name=exact" + assert store.create(workspace_id=workspace, app_id=app_id, tracing_provider="arize", tracing_config=legacy) + repository = ApplicationTraceRepository(sqlite_session_factory) + run = repository.start( + tenant_id=workspace, app_id=app_id, source_kind="message", source_id=message_id, session_id="session" + ) + assert run is not None + with sqlite_session_factory.begin() as session: + accepted = session.get(ApplicationTraceRun, run.id) + assert accepted is not None + accepted.state = "accepted" + accepted.attempts = 1 + with sqlite_session_factory() as session: + accepted_row = deepcopy( + dict( + session.execute(select(ApplicationTraceRun.__table__).where(ApplicationTraceRun.id == run.id)) + .mappings() + .one() + ) + ) + frozen_before = deepcopy(run.provider_config) + + # Both legacy secret fields are really encrypted, rather than identity stubs. + assert legacy["api_key"] != VALUES["api_key"] + assert legacy["space_id"] != SPACE + responses.extend([space(), page(project())]) + submitted = { + **VALUES, + "api_key": "******", + "space_id": "******", + DESTINATION_KEY: {"organization_id": "forged", "project_id": "forged"}, + "project_url": "https://untrusted.test/forged", + } + untouched = deepcopy(submitted) + service.update(context, app_id, "arize", submitted) + current = store.get(workspace_id=workspace, app_id=app_id, tracing_provider="arize") + assert current is not None + assert current.tracing_config is not None + saved = current.tracing_config + assert saved["api_key"] == legacy["api_key"] + assert saved["space_id"] == legacy["space_id"] + assert saved[DESTINATION_KEY] == {**DESTINATION, "source_digest": saved[DESTINATION_KEY]["source_digest"]} + assert saved["project_url"] == URL + assert submitted == untouched + assert len(calls) == 2 + assert verified[-1]["api_key"] == VALUES["api_key"] + assert verified[-1]["space_id"] == SPACE + + def forbidden(*_args: object, **_kwargs: object) -> Never: + raise AssertionError("Reading current configuration or frozen history must not use vendor I/O") + + monkeypatch.setattr(ssrf_proxy, "get", forbidden) + monkeypatch.setattr(OpsTraceManager, "get_trace_config_project_url", forbidden) + shown = service.get(context, app_id, "arize") + assert shown is not None + assert shown.tracing_config is not None + assert shown.tracing_config["project_url"] == URL + assert DESTINATION_KEY not in shown.tracing_config + frozen = repository.source_run(tenant_id=workspace, app_id=app_id, source_kind="message", source_id=message_id) + assert frozen is not None + assert (frozen.state, frozen.attempts, frozen.provider_config) == ("accepted", 1, frozen_before) + assert frozen_project_url("arize", frozen.provider_config) is None + assert frozen.provider_config == frozen_before + monkeypatch.setattr(OpsTraceManager, "decrypt_tracing_config", forbidden) + status = ApplicationTracingService(repository).status( + tenant_id=workspace, app_id=app_id, source_kind="message", source_id=message_id + ) + assert status["project_url"] is None + assert status["status"] == "accepted" + with sqlite_session_factory() as session: + after_row = dict( + session.execute(select(ApplicationTraceRun.__table__).where(ApplicationTraceRun.id == run.id)) + .mappings() + .one() + ) + assert after_row == accepted_row + + # A future source freezes the verified destination; it never repairs the old run. + later_id = str(uuid4()) + with sqlite_session_factory.begin() as session: + session.add( + make_message( + message_id=later_id, + app_id=app_id, + inputs={}, + query="later", + message={}, + answer="ok", + message_unit_price=Decimal(0), + answer_unit_price=Decimal(0), + currency="USD", + from_source=ConversationFromSource.CONSOLE, + ) + ) + later = repository.start( + tenant_id=workspace, app_id=app_id, source_kind="message", source_id=later_id, session_id="session" + ) + assert later is not None + assert frozen_project_url("arize", later.provider_config) == URL + + @staticmethod + @pytest.mark.parametrize("status", [401, 403, 404, 500]) + def test_ingest_verified_key_without_metadata_read_permission_still_saves( + workspace: str, + verified: list[dict[str, object]], + calls: list[tuple[str, dict[str, object]]], + responses: deque[httpx.Response | Exception], + status: int, + ) -> None: + responses.append(_arize_tracing_project_url_response({"message": "private-response"}, status)) + saved = OpsTraceManagerGateway().prepare_new_config( + workspace_id=workspace, tracing_provider="arize", tracing_config=VALUES + ) + assert "project_url" not in saved + assert DESTINATION_KEY not in saved + assert OpsTraceManager.decrypt_tracing_config(workspace, "arize", saved)["api_key"] == VALUES["api_key"] + assert len(verified) == len(calls) == 1 + + @staticmethod + def test_changed_destination_cannot_keep_old_or_client_supplied_cache( + workspace: str, + verified: list[dict[str, object]], + calls: list[tuple[str, dict[str, object]]], + responses: deque[httpx.Response | Exception], + ) -> None: + gateway = OpsTraceManagerGateway() + responses.extend([space(), page(project())]) + before = gateway.prepare_new_config(workspace_id=workspace, tracing_provider="arize", tracing_config=VALUES) + assert before["project_url"] == URL + responses.append(_arize_tracing_project_url_response({}, 403)) + after = gateway.prepare_updated_config( + workspace_id=workspace, + tracing_provider="arize", + current_tracing_config=before, + tracing_config={ + **VALUES, + "api_key": "******", + "space_id": "******", + "project": "later", + DESTINATION_KEY: before[DESTINATION_KEY], + "project_url": URL, + }, + ) + assert after["project"] == "later" + assert after["api_key"] == before["api_key"] + assert after["space_id"] == before["space_id"] + assert "project_url" not in after + assert DESTINATION_KEY not in after + assert before["project_url"] == URL + assert len(calls) == 3 + assert len(verified) == 2 + + @staticmethod + def test_failed_original_ingestion_verification_never_resolves_metadata( + workspace: str, + monkeypatch: pytest.MonkeyPatch, + calls: list[tuple[str, dict[str, object]]], + ) -> None: + def rejected(_config: ArizeConfig) -> Never: + raise arize_configuration.ArizeConfigurationError("permission_denied") + + monkeypatch.setattr(arize_configuration, "verify_configuration", rejected) + with pytest.raises(AppTracingConfigVerificationFailedError): + OpsTraceManagerGateway().prepare_new_config( + workspace_id=workspace, tracing_provider="arize", tracing_config=VALUES + ) + assert calls == [] + + +"""Databricks config persistence must not initialize a process-global MLflow client.""" + + +ExperimentAPI = _test_importlib.import_module( + "tests.unit_tests.providers.trace.trace-mlflow.src.dify_trace_mlflow.test_configuration" +).ExperimentAPI +_databricks_tracing_config_gateway_response = _test_importlib.import_module( + "tests.unit_tests.providers.trace.trace-mlflow.src.dify_trace_mlflow.test_configuration" +).response + + +@dataclass +class _DatabricksTracingConfigGatewayGatewayenvironment: + gateway: OpsTraceManagerGateway = field(default_factory=OpsTraceManagerGateway) + api: ExperimentAPI = field(default_factory=ExperimentAPI) + forbidden_calls: list[str] = field(default_factory=list) + + +class TestDatabricksTracingConfigGatewayCases: + @staticmethod + @pytest.fixture + def environment(monkeypatch: pytest.MonkeyPatch) -> Iterator[_DatabricksTracingConfigGatewayGatewayenvironment]: + environment = _DatabricksTracingConfigGatewayGatewayenvironment() + + def encrypt(_tenant: str, value: str) -> str: + return "test-cipher:" + value + + def decrypt(_tenant: str, values: Sequence[str]) -> list[str]: + assert all(value.startswith("test-cipher:") for value in values) + return [value.removeprefix("test-cipher:") for value in values] + + def forbidden(name: str) -> Callable[..., Never]: + def block(*_args: object, **_kwargs: object) -> Never: + environment.forbidden_calls.append(name) + raise AssertionError("Configuration must not initialize a stateful exporter") + + return block + + monkeypatch.setattr(ops_trace_manager, "encrypt_token", encrypt) + monkeypatch.setattr(ops_trace_manager, "batch_decrypt_token", decrypt) + monkeypatch.setattr(ops_trace_manager, "obfuscated_token", lambda value: "******" if value else value) + monkeypatch.setattr(ssrf_proxy, "get", environment.api.get) + monkeypatch.setattr(ssrf_proxy, "post", environment.api.post) + monkeypatch.setattr(MLflowDataTrace, "__init__", forbidden("legacy_constructor")) + monkeypatch.setattr(mlflow, "set_tracking_uri", forbidden("tracking_uri")) + monkeypatch.setattr(mlflow, "set_experiment", forbidden("experiment")) + monkeypatch.setattr(mlflow, "search_experiments", forbidden("search")) + monkeypatch.setattr(mlflow.config, "enable_async_logging", forbidden("async_logging")) + yield environment + assert environment.forbidden_calls == [] + + @staticmethod + @pytest.mark.parametrize("oauth", [False, True]) + def test_create_and_patch_persist_verified_target_without_replacing_masked_credentials( + environment: _DatabricksTracingConfigGatewayGatewayenvironment, + sqlite_session_factory: sessionmaker[Session], + oauth: bool, + ) -> None: + workspace, app_id = str(uuid4()), str(uuid4()) + with sqlite_session_factory.begin() as session: + session.add(make_app(app_id=app_id, tenant_id=workspace)) + store = SQLAlchemyAppTracingConfigRepository(session_factory=sqlite_session_factory) + service = AppTracingConfigService(configs=store, provider=environment.gateway) + context = RequestContext("request", None, "account", workspace) + credentials = ( + {"client_id": "synthetic-client", "client_secret": "synthetic-secret"} + if oauth + else {"personal_access_token": "synthetic-pat"} + ) + secret_field = "client_secret" if oauth else "personal_access_token" + service.create( + context, app_id, "databricks", {"host": " https://first.test/ ", "experiment_id": "00012", **credentials} + ) + before = store.get(workspace_id=workspace, app_id=app_id, tracing_provider="databricks") + assert before is not None + assert before.tracing_config is not None + assert before.tracing_config["host"] == "https://first.test" + assert before.tracing_config["experiment_id"] == "12" + service.update( + context, + app_id, + "databricks", + { + "host": " https://second.test/ ", + "experiment_id": "00034", + **credentials, + secret_field: "******", + }, + ) + after = store.get(workspace_id=workspace, app_id=app_id, tracing_provider="databricks") + assert after is not None + assert after.tracing_config is not None + assert after.tracing_config["host"] == "https://second.test" + assert after.tracing_config["experiment_id"] == "34" + assert after.tracing_config[secret_field] == before.tracing_config[secret_field] + assert after.tracing_config["project_url"] == "https://second.test/ml/experiments/34/traces" + assert before.tracing_config["project_url"] == "https://first.test/ml/experiments/12/traces" + before_get = len(environment.api.calls) + presented = service.get(context, app_id, "databricks") + assert presented is not None + assert presented.tracing_config is not None + assert presented.tracing_config[secret_field] == "******" + assert presented.tracing_config["project_url"] == after.tracing_config["project_url"] + environment.gateway.validate_saved_config( + workspace_id=workspace, tracing_provider="databricks", tracing_config=after.tracing_config + ) + assert len(environment.api.calls) == before_get + + @staticmethod + @pytest.mark.parametrize("saved_url", [None, "https://stale.test/old-project"]) + def test_get_is_pure_for_historical_and_current_configs( + environment: _DatabricksTracingConfigGatewayGatewayenvironment, saved_url: str | None + ) -> None: + workspace = str(uuid4()) + encrypted = OpsTraceManager.encrypt_tracing_config( + workspace, + "databricks", + {"host": "https://workspace.test/", "experiment_id": "12", "personal_access_token": "synthetic-pat"}, + ) + if saved_url is not None: + encrypted["project_url"] = saved_url + result = environment.gateway.present_config( + workspace_id=workspace, tracing_provider="databricks", tracing_config=encrypted + ) + assert result["project_url"] == "https://workspace.test/ml/experiments/12/traces" + assert result["personal_access_token"] == "******" + assert environment.api.calls == [] + + @staticmethod + def test_invalid_historical_host_uses_safe_pure_fallback( + environment: _DatabricksTracingConfigGatewayGatewayenvironment, + ) -> None: + workspace = str(uuid4()) + encrypted = OpsTraceManager.encrypt_tracing_config( + workspace, + "databricks", + { + "host": "https://user:password@workspace.test", + "experiment_id": "12", + "personal_access_token": "synthetic-pat", + }, + ) + result = environment.gateway.present_config( + workspace_id=workspace, tracing_provider="databricks", tracing_config=encrypted + ) + assert result["project_url"] == "https://www.databricks.com/" + assert environment.api.calls == [] + + @staticmethod + def test_failed_patch_keeps_original_destination_and_logs_only_safe_code( + environment: _DatabricksTracingConfigGatewayGatewayenvironment, + sqlite_session_factory: sessionmaker[Session], + caplog: pytest.LogCaptureFixture, + ) -> None: + workspace, app_id = str(uuid4()), str(uuid4()) + with sqlite_session_factory.begin() as session: + session.add(make_app(app_id=app_id, tenant_id=workspace)) + store = SQLAlchemyAppTracingConfigRepository(session_factory=sqlite_session_factory) + service = AppTracingConfigService(configs=store, provider=environment.gateway) + context = RequestContext("request", None, "account", workspace) + original = { + "host": "https://workspace.test", + "experiment_id": "12", + "personal_access_token": "synthetic-private-pat", + } + service.create(context, app_id, "databricks", original) + before = store.get(workspace_id=workspace, app_id=app_id, tracing_provider="databricks") + assert before is not None + environment.api.override = lambda _method, url, _options: _databricks_tracing_config_gateway_response( + url, {"message": "private-vendor-detail"}, 403 + ) + with pytest.raises(AppTracingConfigVerificationFailedError) as caught: + service.update( + context, app_id, "databricks", {**original, "experiment_id": "34", "personal_access_token": "******"} + ) + assert isinstance(caught.value.__cause__, DatabricksConfigurationError) + assert caught.value.__cause__.code == "permission_denied" + after = store.get(workspace_id=workspace, app_id=app_id, tracing_provider="databricks") + assert after is not None + assert after.tracing_config == before.tracing_config + assert "permission_denied" in caplog.text + assert "private-vendor-detail" not in caplog.text + assert "synthetic-private-pat" not in caplog.text + + @staticmethod + @pytest.mark.parametrize("start_with_oauth", [False, True]) + def test_switching_auth_clears_unselected_credentials_in_persistence_and_requests( + environment: _DatabricksTracingConfigGatewayGatewayenvironment, + sqlite_session_factory: sessionmaker[Session], + start_with_oauth: bool, + ) -> None: + workspace, app_id = str(uuid4()), str(uuid4()) + with sqlite_session_factory.begin() as session: + session.add(make_app(app_id=app_id, tenant_id=workspace)) + store = SQLAlchemyAppTracingConfigRepository(session_factory=sqlite_session_factory) + service = AppTracingConfigService(configs=store, provider=environment.gateway) + context = RequestContext("request", None, "account", workspace) + pat = {"personal_access_token": "synthetic-pat", "client_id": None, "client_secret": None} + oauth = {"personal_access_token": None, "client_id": "synthetic-client", "client_secret": "synthetic-secret"} + target = {"host": "https://workspace.test", "experiment_id": "12"} + service.create(context, app_id, "databricks", {**target, **(oauth if start_with_oauth else pat)}) + environment.api.calls.clear() + service.update(context, app_id, "databricks", {**target, **(pat if start_with_oauth else oauth)}) + stored = store.get(workspace_id=workspace, app_id=app_id, tracing_provider="databricks") + assert stored is not None + assert stored.tracing_config is not None + assert stored.tracing_config["client_secret" if start_with_oauth else "personal_access_token"] is None + assert [method for method, _url, _options in environment.api.calls] == ( + ["GET"] if start_with_oauth else ["POST", "GET"] + ) + expected_token = "synthetic-pat" if start_with_oauth else "synthetic-oauth-token" + assert environment.api.calls[-1][2]["headers"] == { + "Authorization": f"Bearer {expected_token}", + "Accept": "application/json", + } + presented = service.get(context, app_id, "databricks") + assert presented is not None + assert presented.tracing_config is not None + assert presented.tracing_config["client_secret" if start_with_oauth else "personal_access_token"] is None + + +"""MLflow config persistence uses real RSA/SQLite without constructing an exporter.""" + + +ExperimentAPI = _test_importlib.import_module( + "tests.unit_tests.providers.trace.trace-mlflow.src.dify_trace_mlflow.test_configuration" +).ExperimentAPI +_mlflow_tracing_config_gateway_response = _test_importlib.import_module( + "tests.unit_tests.providers.trace.trace-mlflow.src.dify_trace_mlflow.test_configuration" +).response + + +@dataclass +class _MlflowTracingConfigGatewayConfigurationcase: + api: ExperimentAPI + gateway: OpsTraceManagerGateway + service: AppTracingConfigService + store: SQLAlchemyAppTracingConfigRepository + reader: ConsoleAppRepository + context: RequestContext + app_id: str + + def read(self, provider: str = "mlflow") -> AppTracingConfigRecord | None: + return self.store.get( + workspace_id=self.context.active_workspace_id, app_id=self.app_id, tracing_provider=provider + ) + + +class TestMlflowTracingConfigGatewayCases: + import importlib as _test_importlib + + experiment_api = _test_importlib.import_module( + "tests.unit_tests.providers.trace.trace-mlflow.src.dify_trace_mlflow.test_oss_configuration" + ).experiment_api + + experiment_api = staticmethod(experiment_api) + + import importlib as _test_importlib + + isolated_sdk_state = _test_importlib.import_module( + "tests.unit_tests.providers.trace.trace-mlflow.src.dify_trace_mlflow.test_oss_configuration" + ).isolated_sdk_state + + isolated_sdk_state = staticmethod(isolated_sdk_state) + + @staticmethod + @pytest.fixture + def case( + workspace: str, sqlite_session_factory: sessionmaker[Session], experiment_api: ExperimentAPI + ) -> _MlflowTracingConfigGatewayConfigurationcase: + app_id = str(uuid4()) + with sqlite_session_factory.begin() as session: + session.add( + make_app( + app_id=app_id, + tenant_id=workspace, + tracing=json.dumps( + { + "enabled": True, + "tracing_provider": "langsmith", + "capture_content": True, + "capture_agent": False, + } + ), + ) + ) + gateway = OpsTraceManagerGateway() + store = SQLAlchemyAppTracingConfigRepository(session_factory=sqlite_session_factory) + previous = OpsTraceManager.encrypt_tracing_config( + workspace, + "langsmith", + {"api_key": "synthetic-other-key", "project": "other", "endpoint": "https://other.test"}, + ) + assert store.create( + workspace_id=workspace, app_id=app_id, tracing_provider="langsmith", tracing_config=previous + ) + return _MlflowTracingConfigGatewayConfigurationcase( + experiment_api, + gateway, + AppTracingConfigService(configs=store, provider=gateway), + store, + ConsoleAppRepository(session_factory=sqlite_session_factory), + RequestContext("request", None, "account", workspace), + app_id, + ) + + @staticmethod + def test_post_masked_patch_and_get_preserve_ciphertext_and_exact_destination( + case: _MlflowTracingConfigGatewayConfigurationcase, + ) -> None: + original = case.read("langsmith") + submitted = { + "tracking_uri": " HTTP://FIRST.Example:5000/team/ ", + "experiment_id": "00012", + "username": "test-user", + "password": "synthetic-password", + } + untouched = dict(submitted) + case.service.create(case.context, case.app_id, "mlflow", submitted, activate=True) + before = case.read() + assert before is not None + assert before.tracing_config is not None + config = before.tracing_config + assert config["tracking_uri"] == "http://first.example:5000/team" + assert config["experiment_id"] == "12" + assert config["password"] != "synthetic-password" + assert OpsTraceManager.decrypt_tracing_config(case.context.active_workspace_id, "mlflow", config)[ + "password" + ] == ("synthetic-password") + assert submitted == untouched + shown = case.service.get(case.context, case.app_id, "mlflow") + assert shown is not None + assert shown.tracing_config is not None + masked = shown.tracing_config["password"] + assert isinstance(masked, str) + assert "*" in masked + assert masked != "synthetic-password" + assert len(case.api.calls) == 1 + case.service.update( + case.context, + case.app_id, + "mlflow", + { + "tracking_uri": "https://SECOND.test/new-prefix/", + "experiment_id": "00034", + "username": "test-user", + "password": masked, + }, + ) + after = case.read() + assert after is not None + assert after.tracing_config is not None + assert after.tracing_config["password"] == config["password"] + assert after.tracing_config["tracking_uri"] == "https://second.test/new-prefix" + assert after.tracing_config["experiment_id"] == "34" + assert after.tracing_config["project_url"] == "https://second.test/new-prefix/#/experiments/34/traces" + assert len(case.api.calls) == 2 + assert case.api.calls[-1][1] == "https://second.test/new-prefix/api/2.0/mlflow/experiments/get" + case.gateway.validate_saved_config( + workspace_id=case.context.active_workspace_id, + tracing_provider="mlflow", + tracing_config=after.tracing_config, + ) + assert case.service.get(case.context, case.app_id, "mlflow") is not None + assert len(case.api.calls) == 2 + assert case.read("langsmith") == original + policy = case.reader.get_trace(case.context, case.app_id) + assert policy.enabled + assert policy.tracing_provider == "mlflow" + assert policy.capture_content is True + assert policy.capture_agent is False + + @staticmethod + @pytest.mark.parametrize("clear_with_null", [False, True]) + def test_basic_to_anonymous_and_back_does_not_reuse_credentials( + case: _MlflowTracingConfigGatewayConfigurationcase, clear_with_null: bool + ) -> None: + destination = {"tracking_uri": "https://same.test/prefix", "experiment_id": "12"} + case.service.create( + case.context, case.app_id, "mlflow", {**destination, "username": "first", "password": "first-secret"} + ) + empty = None if clear_with_null else "" + case.service.update(case.context, case.app_id, "mlflow", {**destination, "username": empty, "password": empty}) + assert case.api.calls[-1][2]["headers"] == { + "Accept": "application/json", + "Accept-Encoding": "identity", + "Cookie": "", + } + stored = case.read() + assert stored is not None + assert stored.tracing_config is not None + assert ( + OpsTraceManager.decrypt_tracing_config(case.context.active_workspace_id, "mlflow", stored.tracing_config)[ + "password" + ] + == empty + ) + case.service.update( + case.context, case.app_id, "mlflow", {**destination, "username": "second", "password": "second-secret"} + ) + assert case.api.calls[-1][2]["headers"] != case.api.calls[0][2]["headers"] + + @staticmethod + def test_failed_patch_keeps_config_policy_and_other_provider( + case: _MlflowTracingConfigGatewayConfigurationcase, caplog: pytest.LogCaptureFixture + ) -> None: + original = { + "tracking_uri": "https://first.test", + "experiment_id": "12", + "username": "user", + "password": "private-password", + } + case.service.create(case.context, case.app_id, "mlflow", original) + before, other, policy = case.read(), case.read("langsmith"), case.reader.get_trace(case.context, case.app_id) + case.api.override = lambda _method, url, _options: _mlflow_tracing_config_gateway_response( + url, {"message": "private-vendor-response"}, 403 + ) + with pytest.raises(AppTracingConfigVerificationFailedError) as caught: + case.service.update( + case.context, + case.app_id, + "mlflow", + {**original, "tracking_uri": "https://second.test", "password": "******"}, + activate=True, + ) + assert isinstance(caught.value.__cause__, MLflowConfigurationError) + assert caught.value.__cause__.code == "permission_denied" + assert case.read() == before + assert case.read("langsmith") == other + assert case.reader.get_trace(case.context, case.app_id) == policy + assert "permission_denied" in caplog.text + assert "private-vendor-response" not in caplog.text + assert "private-password" not in caplog.text + + @staticmethod + def test_failed_auth_pair_patch_does_not_erase_a_password( + case: _MlflowTracingConfigGatewayConfigurationcase, + ) -> None: + original = { + "tracking_uri": "https://first.test", + "experiment_id": "12", + "username": "user", + "password": "password", + } + case.service.create(case.context, case.app_id, "mlflow", original) + before = case.read() + with pytest.raises(AppTracingConfigVerificationFailedError) as caught: + case.service.update(case.context, case.app_id, "mlflow", {**original, "password": None}) + assert isinstance(caught.value.__cause__, MLflowConfigurationError) + assert caught.value.__cause__.code == "incomplete_basic_auth" + assert len(case.api.calls) == 1 + assert case.read() == before + + @staticmethod + @pytest.mark.parametrize("saved_url", [None, "https://stale.test/wrong-experiment"]) + def test_historical_get_is_pure_even_with_partial_auth( + case: _MlflowTracingConfigGatewayConfigurationcase, saved_url: str | None + ) -> None: + stored = OpsTraceManager.encrypt_tracing_config( + case.context.active_workspace_id, + "mlflow", + { + "tracking_uri": "https://server.test/prefix/", + "experiment_id": "00012", + "username": "historical-user", + "password": None, + }, + ) + if saved_url is not None: + stored["project_url"] = saved_url + shown = case.gateway.present_config( + workspace_id=case.context.active_workspace_id, tracing_provider="mlflow", tracing_config=stored + ) + assert shown["project_url"] == "https://server.test/prefix/#/experiments/12/traces" + assert shown["password"] is None + case.gateway.validate_saved_config( + workspace_id=case.context.active_workspace_id, tracing_provider="mlflow", tracing_config=stored + ) + assert case.api.calls == [] + + @staticmethod + def test_invalid_historical_url_has_no_fabricated_link_or_network( + case: _MlflowTracingConfigGatewayConfigurationcase, + ) -> None: + stored = OpsTraceManager.encrypt_tracing_config( + case.context.active_workspace_id, + "mlflow", + { + "tracking_uri": "https://user:secret@server.test?token=private", + "experiment_id": "12", + "username": None, + "password": None, + }, + ) + shown = case.gateway.present_config( + workspace_id=case.context.active_workspace_id, tracing_provider="mlflow", tracing_config=stored + ) + assert shown["project_url"] is None + assert case.api.calls == [] + + @staticmethod + def test_invalid_new_destination_never_reaches_storage_or_network( + case: _MlflowTracingConfigGatewayConfigurationcase, + ) -> None: + with pytest.raises(AppTracingConfigInvalidConfigurationError): + case.service.create( + case.context, + case.app_id, + "mlflow", + { + "tracking_uri": "https://server.test/prefix/../other", + "experiment_id": "12", + "username": "user", + "password": "secret", + }, + activate=True, + ) + assert case.read() is None + assert case.api.calls == [] + assert case.reader.get_trace(case.context, case.app_id).tracing_provider == "langsmith" + + +"""Opik configuration transactions use real RSA and never initialize an SDK streamer.""" + + +AccessAPI = _test_importlib.import_module( + "tests.unit_tests.providers.trace.trace-opik.src.dify_trace_opik.test_configuration" +).AccessAPI + + +@dataclass +class _OpikTracingConfigGatewayConfigurationcase: + api: AccessAPI + gateway: OpsTraceManagerGateway + service: AppTracingConfigService + store: SQLAlchemyAppTracingConfigRepository + reader: ConsoleAppRepository + context: RequestContext + app_id: str + sessions: sessionmaker[Session] + + +class TestOpikTracingConfigGatewayCases: + @staticmethod + @pytest.fixture + def case( + workspace: str, sqlite_session_factory: sessionmaker[Session], monkeypatch: pytest.MonkeyPatch + ) -> Iterator[_OpikTracingConfigGatewayConfigurationcase]: + app_id = str(uuid4()) + with sqlite_session_factory.begin() as session: + session.add(make_app(app_id=app_id, tenant_id=workspace)) + api = AccessAPI() + gateway = OpsTraceManagerGateway() + store = SQLAlchemyAppTracingConfigRepository(session_factory=sqlite_session_factory) + forbidden_calls: list[str] = [] + + def forbidden(*_args: object, **_kwargs: object) -> Never: + forbidden_calls.append("sdk_constructor") + raise AssertionError("Configuration must not initialize an SDK client or ambient configuration") + + monkeypatch.setattr(ssrf_proxy, "make_request", api.request) + monkeypatch.setattr(OpikDataTrace, "__init__", forbidden) + monkeypatch.setattr(opik.Opik, "__init__", forbidden) + monkeypatch.setattr(SDKOpikConfig, "__init__", forbidden) + yield _OpikTracingConfigGatewayConfigurationcase( + api, + gateway, + AppTracingConfigService(configs=store, provider=gateway), + store, + ConsoleAppRepository(session_factory=sqlite_session_factory), + RequestContext("request", None, "account", workspace), + app_id, + sqlite_session_factory, + ) + assert forbidden_calls == [] + + @staticmethod + def test_post_get_patch_preserve_real_ciphertext_and_verified_destination( + case: _OpikTracingConfigGatewayConfigurationcase, + ) -> None: + secret = "synthetic-private-opik-key" + case.service.create( + case.context, + case.app_id, + "opik", + {"api_key": secret, "workspace": "", "project": " first ", "url": "https://first.test/prefix/api/"}, + activate=True, + ) + before = case.store.get( + workspace_id=case.context.active_workspace_id, app_id=case.app_id, tracing_provider="opik" + ) + assert before is not None + assert before.tracing_config is not None + config = before.tracing_config + assert config["api_key"] != secret + assert config["workspace"] == "resolved-team" + assert config["project"] == "first" + assert config["project_url"] == "https://first.test/prefix/resolved-team/redirect/projects?name=first" + assert ( + OpsTraceManager.decrypt_tracing_config(case.context.active_workspace_id, "opik", config)["api_key"] + == secret + ) + assert case.reader.get_trace(case.context, case.app_id).tracing_provider == "opik" + + before_get = len(case.api.calls) + presented = case.service.get(case.context, case.app_id, "opik") + assert presented is not None + assert presented.tracing_config is not None + masked = presented.tracing_config["api_key"] + assert masked != secret + assert "*" in masked + assert len(case.api.calls) == before_get + case.api.default_workspace = "new-team" + case.service.update( + case.context, + case.app_id, + "opik", + {"api_key": masked, "workspace": "default", "project": "A & B", "url": "https://second.test/api/"}, + ) + after = case.store.get( + workspace_id=case.context.active_workspace_id, app_id=case.app_id, tracing_provider="opik" + ) + assert after is not None + assert after.tracing_config is not None + assert after.tracing_config["api_key"] == config["api_key"] + assert after.tracing_config["workspace"] == "new-team" + assert after.tracing_config["project"] == "A & B" + assert after.tracing_config["url"] == "https://second.test/api/" + assert after.tracing_config["project_url"] == "https://second.test/new-team/redirect/projects?name=A+%26+B" + assert case.api.calls[-1][2]["headers"] == { + "Authorization": secret, + "Comet-Workspace": "default", + "Accept": "application/json", + } + before_get = len(case.api.calls) + presented = case.service.get(case.context, case.app_id, "opik") + assert presented is not None + assert presented.tracing_config is not None + assert presented.tracing_config["project_url"] == after.tracing_config["project_url"] + case.gateway.validate_saved_config( + workspace_id=case.context.active_workspace_id, tracing_provider="opik", tracing_config=after.tracing_config + ) + assert len(case.api.calls) == before_get + + @staticmethod + @pytest.mark.parametrize("saved_url", [None, "https://wrong-global.test/old-project"]) + @pytest.mark.parametrize("workspace_name", ["team", "default", None]) + def test_historical_get_uses_only_configuration_host_without_network( + case: _OpikTracingConfigGatewayConfigurationcase, saved_url: str | None, workspace_name: str | None + ) -> None: + stored = OpsTraceManager.encrypt_tracing_config( + case.context.active_workspace_id, + "opik", + { + "api_key": "synthetic-key", + "workspace": workspace_name, + "project": "A & B", + "url": "https://private.test/opik/api/", + }, + ) + if saved_url: + stored["project_url"] = saved_url + result = case.gateway.present_config( + workspace_id=case.context.active_workspace_id, tracing_provider="opik", tracing_config=stored + ) + expected = ( + "https://private.test/opik/team/redirect/projects?name=A+%26+B" + if workspace_name == "team" + else "https://private.test/opik/" + ) + assert result["project_url"] == expected + assert result["api_key"] != "synthetic-key" + assert case.api.calls == [] + + @staticmethod + @pytest.mark.parametrize("status", [401, 403]) + @pytest.mark.parametrize("existing", [False, True]) + def test_bad_key_or_workspace_does_not_save_or_change_active_provider( + case: _OpikTracingConfigGatewayConfigurationcase, status: int, existing: bool, caplog: pytest.LogCaptureFixture + ) -> None: + original = {"api_key": "synthetic-private-key", "workspace": "team", "project": "original"} + if existing: + case.service.create(case.context, case.app_id, "opik", original) + with case.sessions.begin() as session: + app = session.get(App, case.app_id) + assert app is not None + app.tracing = json.dumps({"enabled": True, "tracing_provider": "langsmith", "capture_content": False}) + before_policy = case.reader.get_trace(case.context, case.app_id) + before = case.store.get( + workspace_id=case.context.active_workspace_id, app_id=case.app_id, tracing_provider="opik" + ) + case.api.override = lambda method, url, _options: httpx.Response( + status, text="private-vendor-details", request=httpx.Request(method, url) + ) + operation = case.service.update if existing else case.service.create + with pytest.raises(AppTracingConfigVerificationFailedError) as error: + operation(case.context, case.app_id, "opik", {**original, "workspace": "inaccessible"}, activate=True) + assert isinstance(error.value.__cause__, OpikConfigurationError) + assert error.value.__cause__.code == ("authentication_failed" if status == 401 else "permission_denied") + after = case.store.get( + workspace_id=case.context.active_workspace_id, app_id=case.app_id, tracing_provider="opik" + ) + assert after == before + assert case.reader.get_trace(case.context, case.app_id) == before_policy + assert "private-vendor-details" not in caplog.text + assert "synthetic-private-key" not in caplog.text + + @staticmethod + def test_clear_optional_api_key_keeps_keyless_self_host_valid( + case: _OpikTracingConfigGatewayConfigurationcase, monkeypatch: pytest.MonkeyPatch + ) -> None: + target = {"url": "http://local.test:5173/api/", "workspace": "team", "project": "sample"} + case.service.create(case.context, case.app_id, "opik", {**target, "api_key": "synthetic-key"}) + monkeypatch.setenv("OPIK_API_KEY", "unrelated-key") + case.api.calls.clear() + case.service.update(case.context, case.app_id, "opik", {**target, "api_key": None}) + stored = case.store.get( + workspace_id=case.context.active_workspace_id, app_id=case.app_id, tracing_provider="opik" + ) + assert stored is not None + assert stored.tracing_config is not None + assert stored.tracing_config["api_key"] is None + presented = case.service.get(case.context, case.app_id, "opik") + assert presented is not None + assert presented.tracing_config is not None + assert presented.tracing_config["api_key"] is None + assert case.api.calls[0][2]["headers"] == {"Comet-Workspace": "team", "Accept": "application/json"} + assert len(case.api.calls) == 1 + + +"""Persist verified Phoenix configuration while preserving owner policy and masked keys.""" + + +_phoenix_tracing_config_gateway_response = _test_importlib.import_module( + "tests.unit_tests.providers.trace.trace-arize-phoenix.src.dify_trace_arize_phoenix.test_configuration" +).response + + +@dataclass +class PhoenixCase: + gateway: OpsTraceManagerGateway + service: AppTracingConfigService + store: SQLAlchemyAppTracingConfigRepository + reader: ConsoleAppRepository + context: RequestContext + app_id: str + + def read(self, provider: str = "phoenix") -> AppTracingConfigRecord | None: + return self.store.get( + workspace_id=self.context.active_workspace_id, app_id=self.app_id, tracing_provider=provider + ) + + +class TestPhoenixTracingConfigGatewayCases: + import importlib as _test_importlib + + api_calls = _test_importlib.import_module( + "tests.unit_tests.providers.trace.trace-arize-phoenix.src.dify_trace_arize_phoenix.test_configuration" + ).api_calls + + api_calls = staticmethod(api_calls) + + import importlib as _test_importlib + + api_responses = _test_importlib.import_module( + "tests.unit_tests.providers.trace.trace-arize-phoenix.src.dify_trace_arize_phoenix.test_configuration" + ).api_responses + + api_responses = staticmethod(api_responses) + + import importlib as _test_importlib + + no_legacy_or_environment_auth = _test_importlib.import_module( + "tests.unit_tests.providers.trace.trace-arize-phoenix.src.dify_trace_arize_phoenix.test_configuration" + ).no_legacy_or_environment_auth + + no_legacy_or_environment_auth = staticmethod(no_legacy_or_environment_auth) + + @staticmethod + @pytest.fixture + def case(workspace: str, sqlite_session_factory: sessionmaker[Session]) -> PhoenixCase: + app_id = str(uuid4()) + with sqlite_session_factory.begin() as session: + session.add( + make_app( + app_id=app_id, + tenant_id=workspace, + tracing=json.dumps( + { + "enabled": True, + "tracing_provider": "langsmith", + "capture_content": True, + "capture_agent": False, + } + ), + ) + ) + gateway = OpsTraceManagerGateway() + store = SQLAlchemyAppTracingConfigRepository(session_factory=sqlite_session_factory) + old = OpsTraceManager.encrypt_tracing_config( + workspace, + "langsmith", + {"api_key": "synthetic-other-key", "project": "other", "endpoint": "https://other.test"}, + ) + assert store.create(workspace_id=workspace, app_id=app_id, tracing_provider="langsmith", tracing_config=old) + return PhoenixCase( + gateway, + AppTracingConfigService(configs=store, provider=gateway), + store, + ConsoleAppRepository(session_factory=sqlite_session_factory), + RequestContext("request", None, "account", workspace), + app_id, + ) + + @staticmethod + def test_post_masked_patch_get_keep_normalized_destination_and_ciphertext( + case: PhoenixCase, + api_calls: list[tuple[str, dict[str, object]]], + ) -> None: + other = case.read("langsmith") + submitted = { + "endpoint": " HTTP://FIRST.Example:6006/reverse/ ", + "project": " first project ", + "api_key": "synthetic-secret", + } + untouched = dict(submitted) + case.service.create(case.context, case.app_id, "phoenix", submitted, activate=True) + before = case.read() + assert before is not None + assert before.tracing_config is not None + stored = before.tracing_config + assert stored["endpoint"] == "http://first.example:6006/reverse" + assert stored["project"] == "first project" + assert stored["api_key"] != "synthetic-secret" + assert submitted == untouched + assert OpsTraceManager.decrypt_tracing_config(case.context.active_workspace_id, "phoenix", stored)[ + "api_key" + ] == ("synthetic-secret") + shown = case.service.get(case.context, case.app_id, "phoenix") + assert shown is not None + assert shown.tracing_config is not None + masked = shown.tracing_config["api_key"] + assert isinstance(masked, str) + assert "*" in masked + assert masked != "synthetic-secret" + assert ( + shown.tracing_config["project_url"] + == "http://first.example:6006/reverse/redirects/projects/first%20project" + ) + assert len(api_calls) == 1 + case.service.update( + case.context, + case.app_id, + "phoenix", + { + "endpoint": " https://SECOND.test/s/space/ ", + "project": " second & 项目 ", + "api_key": masked, + }, + ) + after = case.read() + assert after is not None + assert after.tracing_config is not None + assert after.tracing_config["api_key"] == stored["api_key"] + assert after.tracing_config["endpoint"] == "https://second.test/s/space" + assert after.tracing_config["project"] == "second & 项目" + assert after.tracing_config["project_url"] == ( + "https://second.test/s/space/redirects/projects/second%20%26%20%E9%A1%B9%E7%9B%AE" + ) + assert api_calls[-1][0] == "https://second.test/s/space/v1/projects" + headers = api_calls[-1][1]["headers"] + assert isinstance(headers, dict) + assert headers["Authorization"] == "Bearer synthetic-secret" + case.gateway.validate_saved_config( + workspace_id=case.context.active_workspace_id, + tracing_provider="phoenix", + tracing_config=after.tracing_config, + ) + assert case.service.get(case.context, case.app_id, "phoenix") is not None + assert len(api_calls) == 2 + assert case.read("langsmith") == other + policy = case.reader.get_trace(case.context, case.app_id) + assert policy.enabled + assert policy.tracing_provider == "phoenix" + assert policy.capture_content is True + assert policy.capture_agent is False + + @staticmethod + @pytest.mark.parametrize("empty", [None, ""]) + def test_bearer_to_keyless_and_back_is_explicit( + case: PhoenixCase, + api_calls: list[tuple[str, dict[str, object]]], + empty: str | None, + ) -> None: + destination = {"endpoint": "https://host.test", "project": "project"} + case.service.create(case.context, case.app_id, "phoenix", {**destination, "api_key": "first-key"}) + case.service.update(case.context, case.app_id, "phoenix", {**destination, "api_key": empty}) + assert api_calls[-1][1]["headers"] == { + "Accept": "application/json", + "Accept-Encoding": "identity", + "Cookie": "", + } + stored = case.read() + assert stored is not None + assert stored.tracing_config is not None + assert ( + OpsTraceManager.decrypt_tracing_config(case.context.active_workspace_id, "phoenix", stored.tracing_config)[ + "api_key" + ] + == empty + ) + case.service.update(case.context, case.app_id, "phoenix", {**destination, "api_key": "second-key"}) + assert api_calls[-1][1]["headers"] != api_calls[0][1]["headers"] + assert len(api_calls) == 3 + + @staticmethod + @pytest.mark.parametrize("failure", ["denied", "wrong-path-html"]) + def test_failed_patch_preserves_policy_config_and_other_provider( + case: PhoenixCase, + api_calls: list[tuple[str, dict[str, object]]], + api_responses: deque[httpx.Response | Exception], + caplog: pytest.LogCaptureFixture, + failure: str, + ) -> None: + original = {"endpoint": "https://host.test", "project": "project", "api_key": "private-secret"} + case.service.create(case.context, case.app_id, "phoenix", original) + before, other, policy = case.read(), case.read("langsmith"), case.reader.get_trace(case.context, case.app_id) + if failure == "denied": + api_responses.append(_phoenix_tracing_config_gateway_response({"secret": "private-vendor-response"}, 403)) + else: + api_responses.append( + httpx.Response( + 200, + text="private-vendor-response", + request=httpx.Request("GET", "https://host.test/not-a-collector/v1/projects"), + ) + ) + with pytest.raises(AppTracingConfigVerificationFailedError) as caught: + case.service.update( + case.context, + case.app_id, + "phoenix", + { + **original, + "endpoint": "https://host.test/not-a-collector", + "api_key": "******", + }, + activate=True, + ) + assert isinstance(caught.value.__cause__, PhoenixConfigurationError) + assert caught.value.__cause__.code == ("permission_denied" if failure == "denied" else "malformed_response") + assert case.read() == before + assert case.read("langsmith") == other + assert case.reader.get_trace(case.context, case.app_id) == policy + assert "private-vendor-response" not in caplog.text + assert "private-secret" not in caplog.text + assert len(api_calls) == 2 + + @staticmethod + @pytest.mark.parametrize("saved_url", [None, "https://stale.test/wrong-project"]) + def test_historical_get_calculates_url_without_verification_or_sdk( + case: PhoenixCase, + api_calls: list[tuple[str, dict[str, object]]], + saved_url: str | None, + ) -> None: + stored = OpsTraceManager.encrypt_tracing_config( + case.context.active_workspace_id, + "phoenix", + { + "endpoint": "https://host.test/s/space/", + "project": "historical project", + "api_key": None, + }, + ) + if saved_url: + stored["project_url"] = saved_url + shown = case.gateway.present_config( + workspace_id=case.context.active_workspace_id, tracing_provider="phoenix", tracing_config=stored + ) + assert shown["project_url"] == "https://host.test/s/space/redirects/projects/historical%20project" + assert shown["api_key"] is None + case.gateway.validate_saved_config( + workspace_id=case.context.active_workspace_id, tracing_provider="phoenix", tracing_config=stored + ) + assert api_calls == [] + + @staticmethod + @pytest.mark.parametrize( + "endpoint", ["https://user:private@host.test?token=private", "https://host.test/v1/traces"] + ) + def test_unsafe_historical_destination_is_readable_without_link_or_io( + case: PhoenixCase, + api_calls: list[tuple[str, dict[str, object]]], + endpoint: str, + ) -> None: + stored = OpsTraceManager.encrypt_tracing_config( + case.context.active_workspace_id, + "phoenix", + { + "endpoint": endpoint, + "project": "old-project", + "api_key": None, + }, + ) + shown = case.gateway.present_config( + workspace_id=case.context.active_workspace_id, tracing_provider="phoenix", tracing_config=stored + ) + assert shown["endpoint"] == endpoint + assert shown["project_url"] is None + assert api_calls == [] + + @staticmethod + @pytest.mark.parametrize("values", [{"endpoint": "https://host.test/v1/traces"}, {"project": "bad\nproject"}]) + def test_invalid_new_destination_has_no_io_or_storage_effect( + case: PhoenixCase, + api_calls: list[tuple[str, dict[str, object]]], + values: dict[str, str], + ) -> None: + policy = case.reader.get_trace(case.context, case.app_id) + with pytest.raises(AppTracingConfigInvalidConfigurationError): + case.service.create( + case.context, case.app_id, "phoenix", {"api_key": "synthetic-key", **values}, activate=True + ) + assert case.read() is None + assert case.reader.get_trace(case.context, case.app_id) == policy + assert api_calls == [] + + +"""Tencent configuration persistence uses real RSA and bounded protocol checks.""" + + +CollectorBoundary = _test_importlib.import_module( + "tests.unit_tests.providers.trace.trace-tencent.src.dify_trace_tencent.test_configuration" +).CollectorBoundary + + +@dataclass +class _TencentTracingConfigGatewayConfigurationcase: + collector: CollectorBoundary + gateway: OpsTraceManagerGateway + service: AppTracingConfigService + store: SQLAlchemyAppTracingConfigRepository + reader: ConsoleAppRepository + context: RequestContext + app_id: str + + def read(self, provider: str = "tencent") -> AppTracingConfigRecord | None: + return self.store.get( + workspace_id=self.context.active_workspace_id, app_id=self.app_id, tracing_provider=provider + ) + + +class TestTencentTracingConfigGatewayCases: + import importlib as _test_importlib + + _forbid_provider_io = _test_importlib.import_module( + "tests.unit_tests.providers.trace.trace-tencent.src.dify_trace_tencent.test_configuration" + )._forbid_provider_io + + _forbid_provider_io = staticmethod(_forbid_provider_io) + + import importlib as _test_importlib + + collector = _test_importlib.import_module( + "tests.unit_tests.providers.trace.trace-tencent.src.dify_trace_tencent.test_configuration" + ).collector + + collector = staticmethod(collector) + + @staticmethod + @pytest.fixture + def case( + workspace: str, sqlite_session_factory: sessionmaker[Session], collector: CollectorBoundary + ) -> _TencentTracingConfigGatewayConfigurationcase: + app_id = str(uuid4()) + with sqlite_session_factory.begin() as session: + session.add( + make_app( + app_id=app_id, + tenant_id=workspace, + tracing=json.dumps( + { + "enabled": True, + "tracing_provider": "langsmith", + "capture_content": True, + "capture_agent": False, + } + ), + ) + ) + gateway = OpsTraceManagerGateway() + store = SQLAlchemyAppTracingConfigRepository(session_factory=sqlite_session_factory) + old_config = OpsTraceManager.encrypt_tracing_config( + workspace, + "langsmith", + { + "api_key": "synthetic-existing-key", + "project": "existing-project", + "endpoint": "https://existing.example", + }, + ) + assert store.create( + workspace_id=workspace, app_id=app_id, tracing_provider="langsmith", tracing_config=old_config + ) + return _TencentTracingConfigGatewayConfigurationcase( + collector, + gateway, + AppTracingConfigService(configs=store, provider=gateway), + store, + ConsoleAppRepository(session_factory=sqlite_session_factory), + RequestContext("request", None, "account", workspace), + app_id, + ) + + @staticmethod + def test_post_get_masked_patch_persist_normalization_and_preserve_real_ciphertext( + case: _TencentTracingConfigGatewayConfigurationcase, + ) -> None: + secret = "synthetic-private-tencent-token" + old_provider = case.read("langsmith") + submitted = {"token": secret, "endpoint": " HTTPS://FIRST.Example:4320/ ", "service_name": " first-app "} + original_submission = dict(submitted) + + case.service.create(case.context, case.app_id, "tencent", submitted, activate=True) + + before = case.read() + assert before is not None + assert before.tracing_config is not None + config = before.tracing_config + assert config["token"] != secret + assert config["endpoint"] == "https://first.example:4320" + assert config["service_name"] == "first-app" + assert config["project_url"] == "https://console.cloud.tencent.com/apm" + assert ( + OpsTraceManager.decrypt_tracing_config(case.context.active_workspace_id, "tencent", config)["token"] + == secret + ) + assert submitted == original_submission + assert len(case.collector.channels) == 1 + assert case.collector.channels[0].calls[0].metadata == (("authorization", f"Bearer {secret}"),) + policy = case.reader.get_trace(case.context, case.app_id) + assert policy.enabled is True + assert policy.tracing_provider == "tencent" + assert policy.capture_content is True + assert policy.capture_agent is False + assert case.read("langsmith") == old_provider + + presented = case.service.get(case.context, case.app_id, "tencent") + assert presented is not None + assert presented.tracing_config is not None + masked = presented.tracing_config["token"] + assert isinstance(masked, str) + assert masked != secret + assert "*" in masked + assert presented.tracing_config["project_url"] == "https://console.cloud.tencent.com/apm" + assert len(case.collector.channels) == 1 + + replacement = {"token": masked, "endpoint": " http://SECOND.Example:4319/ ", "service_name": " second-app "} + original_replacement = dict(replacement) + case.service.update(case.context, case.app_id, "tencent", replacement, activate=True) + + after = case.read() + assert after is not None + assert after.tracing_config is not None + assert after.tracing_config["token"] == config["token"] + assert after.tracing_config["endpoint"] == "http://second.example:4319" + assert after.tracing_config["service_name"] == "second-app" + assert after.tracing_config["project_url"] == "https://console.cloud.tencent.com/apm" + assert replacement == original_replacement + assert len(case.collector.channels) == 2 + channel = case.collector.channels[-1] + assert channel.target == "second.example:4319" + assert channel.secure is False + assert channel.calls[0].metadata == (("authorization", f"Bearer {secret}"),) + assert all(channel.close_count == 1 for channel in case.collector.channels) + case.gateway.validate_saved_config( + workspace_id=case.context.active_workspace_id, + tracing_provider="tencent", + tracing_config=after.tracing_config, + ) + assert case.service.get(case.context, case.app_id, "tencent") is not None + assert len(case.collector.channels) == 2 + assert case.read("langsmith") == old_provider + + @staticmethod + @pytest.mark.parametrize("saved_url", [None, "https://wrong.example/previous-project"]) + def test_historical_name_remains_readable_and_saved_validation_does_not_connect( + case: _TencentTracingConfigGatewayConfigurationcase, saved_url: str | None + ) -> None: + stored = OpsTraceManager.encrypt_tracing_config( + case.context.active_workspace_id, + "tencent", + { + "token": "synthetic-historical-token", + "endpoint": "https://legacy.example:4320", + "service_name": "dify_app", + }, + ) + if saved_url is not None: + stored["project_url"] = saved_url + assert case.store.create( + workspace_id=case.context.active_workspace_id, + app_id=case.app_id, + tracing_provider="tencent", + tracing_config=stored, + ) + before = case.read() + before_policy = case.reader.get_trace(case.context, case.app_id) + + presented = case.service.get(case.context, case.app_id, "tencent") + case.gateway.validate_saved_config( + workspace_id=case.context.active_workspace_id, tracing_provider="tencent", tracing_config=stored + ) + + assert presented is not None + assert presented.tracing_config is not None + assert presented.tracing_config["service_name"] == "dify_app" + assert presented.tracing_config["project_url"] == "https://console.cloud.tencent.com/apm" + assert presented.tracing_config["token"] != "synthetic-historical-token" + assert "*" in presented.tracing_config["token"] + assert case.collector.channels == [] + assert case.read() == before + assert case.reader.get_trace(case.context, case.app_id) == before_policy + + @staticmethod + @pytest.mark.parametrize("existing", [False, True]) + @pytest.mark.parametrize( + ("field_name", "invalid_value", "error_code"), + [ + ("service_name", "dify_app", "invalid_service_name"), + ("service_name", "", "invalid_service_name"), + ("endpoint", "https://collector.example:4320/v1/traces", "invalid_endpoint"), + ("endpoint", "https://collector.example:4320?token=private-endpoint-detail", "invalid_endpoint"), + ("endpoint", "https://", "invalid_endpoint"), + ], + ) + def test_invalid_create_or_patch_rejects_before_rpc_encryption_or_persistence( + case: _TencentTracingConfigGatewayConfigurationcase, + monkeypatch: pytest.MonkeyPatch, + caplog: pytest.LogCaptureFixture, + existing: bool, + field_name: str, + invalid_value: str, + error_code: str, + ) -> None: + original = { + "token": "synthetic-private-token", + "endpoint": "https://original.example:4320", + "service_name": "original-app", + } + if existing: + case.service.create(case.context, case.app_id, "tencent", original) + case.collector.channels.clear() + before = case.read() + old_provider = case.read("langsmith") + before_policy = case.reader.get_trace(case.context, case.app_id) + forbidden_attempts: list[str] = [] + + def forbidden(*_args: object, **_kwargs: object) -> Never: + forbidden_attempts.append("encryption_or_write") + raise AssertionError("Invalid new input must fail before encryption or persistence") + + monkeypatch.setattr(OpsTraceManager, "encrypt_tracing_config", forbidden) + monkeypatch.setattr(case.store, "create", forbidden) + monkeypatch.setattr(case.store, "update", forbidden) + operation = case.service.update if existing else case.service.create + with pytest.raises(AppTracingConfigInvalidConfigurationError) as caught: + operation(case.context, case.app_id, "tencent", {**original, field_name: invalid_value}, activate=True) + + assert isinstance(caught.value.__cause__, TencentConfigurationError) + assert caught.value.__cause__.code == error_code + assert forbidden_attempts == [] + assert case.collector.channels == [] + assert case.read() == before + assert case.read("langsmith") == old_provider + assert case.reader.get_trace(case.context, case.app_id) == before_policy + assert "synthetic-private-token" not in caplog.text + assert "private-endpoint-detail" not in caplog.text + + @staticmethod + @pytest.mark.parametrize("existing", [False, True]) + @pytest.mark.parametrize("endpoint", ["https://collector.example:4320", "http://localhost:4317"]) + def test_rpc_failure_never_saves_or_activates_even_for_localhost( + case: _TencentTracingConfigGatewayConfigurationcase, + caplog: pytest.LogCaptureFixture, + existing: bool, + endpoint: str, + ) -> None: + secret = "synthetic-private-token" + original = {"token": secret, "endpoint": "https://original.example:4320", "service_name": "original-app"} + if existing: + case.service.create(case.context, case.app_id, "tencent", original) + case.collector.channels.clear() + before = case.read() + old_provider = case.read("langsmith") + before_policy = case.reader.get_trace(case.context, case.app_id) + case.collector.failure = grpc.RpcError("private-provider-details synthetic-private-token") + operation = case.service.update if existing else case.service.create + + with pytest.raises(AppTracingConfigVerificationFailedError) as caught: + operation( + case.context, + case.app_id, + "tencent", + {"token": "******" if existing else secret, "endpoint": endpoint, "service_name": "changed-app"}, + activate=True, + ) + + assert isinstance(caught.value.__cause__, TencentConfigurationError) + assert caught.value.__cause__.code == "connection_failed" + assert len(case.collector.channels) == 1 + channel = case.collector.channels[0] + assert channel.close_count == 1 + assert len(channel.calls) == 1 + assert channel.calls[0].metadata == (("authorization", f"Bearer {secret}"),) + assert case.read() == before + assert case.read("langsmith") == old_provider + assert case.reader.get_trace(case.context, case.app_id) == before_policy + assert "connection_failed" in caplog.text + assert "private-provider-details" not in caplog.text + assert secret not in caplog.text + assert secret not in str(caught.value) + assert secret not in str(caught.value.__cause__) + + @staticmethod + def test_explicit_configuration_ignores_and_preserves_otel_environment( + case: _TencentTracingConfigGatewayConfigurationcase, monkeypatch: pytest.MonkeyPatch + ) -> None: + environment = { + "OTEL_EXPORTER_OTLP_ENDPOINT": "http://ambient.example:4317", + "OTEL_EXPORTER_OTLP_TRACES_ENDPOINT": "http://ambient-traces.example:4317", + "OTEL_EXPORTER_OTLP_HEADERS": "authorization=Bearer unrelated-token", + "OTEL_EXPORTER_OTLP_TRACES_HEADERS": "authorization=Bearer unrelated-traces-token", + "OTEL_EXPORTER_OTLP_PROTOCOL": "http/protobuf", + "OTEL_SERVICE_NAME": "unrelated-service", + "OTEL_RESOURCE_ATTRIBUTES": "token=unrelated-resource-token,service.name=ambient-app", + } + for key, value in environment.items(): + monkeypatch.setenv(key, value) + + case.service.create( + case.context, + case.app_id, + "tencent", + {"token": "explicit-token", "endpoint": "https://explicit.example:4320", "service_name": "explicit-app"}, + ) + + assert len(case.collector.channels) == 1 + channel = case.collector.channels[0] + assert channel.target == "explicit.example:4320" + assert channel.secure is True + assert channel.calls[0].metadata == (("authorization", "Bearer explicit-token"),) + assert channel.calls[0].request.SerializeToString() == b"" + assert channel.close_count == 1 + presented = case.service.get(case.context, case.app_id, "tencent") + assert presented is not None + assert presented.tracing_config is not None + assert presented.tracing_config["service_name"] == "explicit-app" + assert len(case.collector.channels) == 1 + assert {key: os.environ.get(key) for key in environment} == environment + + +"""Config reads and writes must not construct Weave's stateful legacy exporter.""" + + +ReadAPI = _test_importlib.import_module( + "tests.unit_tests.providers.trace.trace-weave.src.dify_trace_weave.test_configuration" +).ReadAPI + + +@dataclass +class _WeaveTracingConfigGatewayGatewayenvironment: + gateway: OpsTraceManagerGateway = field(default_factory=OpsTraceManagerGateway) + api: ReadAPI = field(default_factory=ReadAPI) + forbidden_calls: list[str] = field(default_factory=list) + + +class TestWeaveTracingConfigGatewayCases: + @staticmethod + @pytest.fixture + def environment(monkeypatch: pytest.MonkeyPatch) -> Iterator[_WeaveTracingConfigGatewayGatewayenvironment]: + environment = _WeaveTracingConfigGatewayGatewayenvironment() + + def encrypt(_tenant: str, value: str) -> str: + return "test-cipher:" + value + + def decrypt(_tenant: str, values: Sequence[str]) -> list[str]: + assert all(value.startswith("test-cipher:") for value in values) + return [value.removeprefix("test-cipher:") for value in values] + + def mask(_value: str) -> str: + return "******" + + def forbidden(name: str) -> Callable[..., Never]: + def block(*_args: object, **_kwargs: object) -> Never: + environment.forbidden_calls.append(name) + raise AssertionError("Configuration must not initialize a stateful exporter") + + return block + + monkeypatch.setattr(ops_trace_manager, "encrypt_token", encrypt) + monkeypatch.setattr(ops_trace_manager, "batch_decrypt_token", decrypt) + monkeypatch.setattr(ops_trace_manager, "obfuscated_token", mask) + monkeypatch.setattr(ssrf_proxy, "post", environment.api.post) + monkeypatch.setattr(WeaveDataTrace, "__init__", forbidden("legacy_constructor")) + monkeypatch.setattr(wandb, "login", forbidden("wandb_login")) + monkeypatch.setattr(weave, "init", forbidden("weave_init")) + yield environment + assert environment.forbidden_calls == [] + + @staticmethod + def test_create_persists_resolved_default_entity_and_canonical_project( + environment: _WeaveTracingConfigGatewayGatewayenvironment, + ) -> None: + submitted = {"api_key": "synthetic-key", "project": " demo "} + stored = environment.gateway.prepare_new_config( + workspace_id=str(uuid4()), tracing_provider="weave", tracing_config=submitted + ) + assert stored["entity"] == "default-team" + assert stored["project"] == "demo" + assert stored["project_url"] == "https://wandb.ai/default-team/demo" + assert stored["api_key"] == "test-cipher:synthetic-key" + assert submitted == {"api_key": "synthetic-key", "project": " demo "} + assert len(environment.api.calls) == 3 + + @staticmethod + def test_patch_persists_verified_destination_with_the_existing_ciphertext( + environment: _WeaveTracingConfigGatewayGatewayenvironment, sqlite_session_factory: sessionmaker[Session] + ) -> None: + workspace, app_id = str(uuid4()), str(uuid4()) + with sqlite_session_factory.begin() as session: + session.add(make_app(app_id=app_id, tenant_id=workspace)) + store = SQLAlchemyAppTracingConfigRepository(session_factory=sqlite_session_factory) + service = AppTracingConfigService(configs=store, provider=environment.gateway) + context = RequestContext("request", None, "account", workspace) + service.create(context, app_id, "weave", {"api_key": "synthetic-key", "entity": "old", "project": "original"}) + before = store.get(workspace_id=workspace, app_id=app_id, tracing_provider="weave") + assert before is not None + assert before.tracing_config is not None + + submitted = {"api_key": "******", "project": "new-team/changed"} + service.update(context, app_id, "weave", submitted) + after = store.get(workspace_id=workspace, app_id=app_id, tracing_provider="weave") + assert after is not None + assert after.tracing_config is not None + assert after.tracing_config["entity"] == "new-team" + assert after.tracing_config["project"] == "changed" + assert after.tracing_config["project_url"] == "https://wandb.ai/new-team/changed" + assert after.tracing_config["api_key"] == before.tracing_config["api_key"] == "test-cipher:synthetic-key" + assert before.tracing_config["project_url"] == "https://wandb.ai/old/original" + assert submitted == {"api_key": "******", "project": "new-team/changed"} + assert len(environment.api.calls) == 4 + + presented = service.get(context, app_id, "weave") + assert presented is not None + assert presented.tracing_config is not None + assert presented.tracing_config["api_key"] == "******" + assert presented.tracing_config["project_url"] == after.tracing_config["project_url"] + environment.gateway.validate_saved_config( + workspace_id=workspace, tracing_provider="weave", tracing_config=after.tracing_config + ) + assert len(environment.api.calls) == 4 + + @staticmethod + @pytest.mark.parametrize( + ("entity", "project", "host", "expected"), + [ + ("team", "demo", None, "https://wandb.ai/team/demo"), + (None, "team/demo", None, "https://wandb.ai/team/demo"), + (None, "demo", None, "https://wandb.ai"), + (None, "demo", "http://private.test/wandb", "http://private.test/wandb"), + ("team", "a & b", "https://private.test/wandb", "https://private.test/wandb/team/a%20%26%20b"), + ], + ) + def test_get_computes_project_or_console_link_without_vendor_requests( + environment: _WeaveTracingConfigGatewayGatewayenvironment, + entity: str | None, + project: str, + host: str | None, + expected: str, + ) -> None: + workspace = str(uuid4()) + encrypted = OpsTraceManager.encrypt_tracing_config( + workspace, "weave", {"api_key": "synthetic-key", "entity": entity, "project": project, "host": host} + ) + encrypted["project_url"] = "https://stale.test/previous-project" + result = environment.gateway.present_config( + workspace_id=workspace, tracing_provider="weave", tracing_config=encrypted + ) + assert result["project_url"] == expected + assert result["api_key"] == "******" + assert environment.api.calls == [] + + @staticmethod + def test_inaccessible_project_is_a_sanitized_domain_error_and_is_not_saved( + environment: _WeaveTracingConfigGatewayGatewayenvironment, + sqlite_session_factory: sessionmaker[Session], + caplog: pytest.LogCaptureFixture, + ) -> None: + workspace, app_id = str(uuid4()), str(uuid4()) + with sqlite_session_factory.begin() as session: + session.add(make_app(app_id=app_id, tenant_id=workspace)) + environment.api.project_exists = False + store = SQLAlchemyAppTracingConfigRepository(session_factory=sqlite_session_factory) + service = AppTracingConfigService(configs=store, provider=environment.gateway) + with pytest.raises(AppTracingConfigVerificationFailedError) as caught: + service.create( + RequestContext("request", None, "account", workspace), + app_id, + "weave", + { + "api_key": "synthetic-private-key", + "entity": "team", + "project": "missing", + }, + ) + assert isinstance(caught.value.__cause__, WeaveConfigurationError) + assert caught.value.__cause__.code == "project_not_found_or_inaccessible" + assert store.get(workspace_id=workspace, app_id=app_id, tracing_provider="weave") is None + assert "synthetic-private-key" not in caplog.text + assert "project_not_found_or_inaccessible" in caplog.text + assert len(environment.api.calls) == 1 diff --git a/api/tests/unit_tests/services/test_app_tracing_config_project_url.py b/api/tests/unit_tests/services/test_app_tracing_config_project_url.py deleted file mode 100644 index 32915c4e0cd..00000000000 --- a/api/tests/unit_tests/services/test_app_tracing_config_project_url.py +++ /dev/null @@ -1,238 +0,0 @@ -"""Real provider config serialization must retain verified project links on update.""" - -import json -from collections.abc import Mapping, Sequence -from dataclasses import dataclass, field -from datetime import UTC, datetime, timedelta -from decimal import Decimal -from uuid import uuid4 - -import httpx -import pytest -from sqlalchemy.orm import Session, sessionmaker - -from core.helper import ssrf_proxy -from core.ops import ops_trace_manager -from core.ops.ops_trace_manager import OpsTraceManager -from core.ops.unified_trace.entities import CanonicalSpan, CanonicalSpanKind, CanonicalSpanStatus, CanonicalTrace -from machinery.context import RequestContext -from models.enums import ConversationFromSource -from models.model import AppMode -from repositories.app.tracing_config_repository import SQLAlchemyAppTracingConfigRepository -from repositories.application_trace_repository import ApplicationTraceRepository -from services.app_tracing_config_gateway import OpsTraceManagerGateway -from services.app_tracing_config_service import AppTracingConfigService, AppTracingConfigVerificationFailedError -from services.tracing.provider_dispatch import export_trace -from tests.unit_tests.model_factories import make_app, make_message -from tests.unit_tests.services.tracing.test_opik_configuration import AccessAPI - - -@dataclass -class ProviderBoundary: - calls: list[tuple[str, str, dict[str, object]]] = field(default_factory=list) - effective: bool = True - lookup_fails: bool = False - - def verify(self, config: Mapping[str, object], provider: str) -> bool: - self.calls.append(("verify", provider, dict(config))) - return self.effective - - def project_url(self, config: Mapping[str, object], provider: str) -> str: - self.calls.append(("url", provider, dict(config))) - if self.lookup_fails: - raise RuntimeError("Synthetic project lookup failure") - return f"https://tracing.test/{provider}/{config['project']}" - - def project_key(self, config: Mapping[str, object], provider: str) -> str: - self.calls.append(("key", provider, dict(config))) - return "verified-project" - - -@pytest.fixture -def boundary(monkeypatch: pytest.MonkeyPatch) -> ProviderBoundary: - boundary = ProviderBoundary() - - def encrypt(_tenant: str, value: str) -> str: - return "test-cipher:" + value - - def decrypt(_tenant: str, values: Sequence[str]) -> list[str]: - assert all(value.startswith("test-cipher:") for value in values) - return [value.removeprefix("test-cipher:") for value in values] - - monkeypatch.setattr(ops_trace_manager, "encrypt_token", encrypt) - monkeypatch.setattr(ops_trace_manager, "batch_decrypt_token", decrypt) - monkeypatch.setattr(OpsTraceManager, "check_trace_config_is_effective", boundary.verify) - monkeypatch.setattr(OpsTraceManager, "get_trace_config_project_url", boundary.project_url) - monkeypatch.setattr(OpsTraceManager, "get_trace_config_project_key", boundary.project_key) - return boundary - - -@pytest.mark.parametrize("provider", ["langsmith", "opik"]) -@pytest.mark.parametrize("project", ["original", "changed"]) -def test_update_resolves_project_link_using_verified_unmasked_configuration( - boundary: ProviderBoundary, provider: str, project: str, monkeypatch: pytest.MonkeyPatch -) -> None: - gateway = OpsTraceManagerGateway() - workspace = str(uuid4()) - api = AccessAPI() - if provider == "opik": - monkeypatch.setattr(ssrf_proxy, "make_request", api.request) - target = {"workspace": "team", "url": "https://tracing.test/opik/api/"} if provider == "opik" else {} - current = gateway.prepare_new_config( - workspace_id=workspace, - tracing_provider=provider, - tracing_config={"api_key": "synthetic-key", "project": "original", **target}, - ) - submitted = {"api_key": "******", "project": project, "project_url": "https://untrusted.test/stale", **target} - updated = gateway.prepare_updated_config( - workspace_id=workspace, tracing_provider=provider, tracing_config=submitted, current_tracing_config=current - ) - - prefix = f"https://tracing.test/{provider}/" + ("team/redirect/projects?name=" if provider == "opik" else "") - assert updated["project_url"] == prefix + project - assert updated["api_key"] == current["api_key"] == "test-cipher:synthetic-key" - assert current["project_url"] == prefix + "original" - assert submitted["api_key"] == "******" - if provider == "opik": - assert boundary.calls == [] - assert len(api.calls) == 2 - assert api.calls[-1][2]["headers"] == { - "Authorization": "synthetic-key", - "Comet-Workspace": "team", - "Accept": "application/json", - } - else: - assert [call[0] for call in boundary.calls] == ["verify", "url", "verify", "url"] - assert boundary.calls[-1][2]["api_key"] == "synthetic-key" - assert boundary.calls[-1][2]["project"] == project - - -def test_langfuse_update_uses_current_host_and_resolved_project_key(boundary: ProviderBoundary) -> None: - gateway = OpsTraceManagerGateway() - workspace = str(uuid4()) - current = gateway.prepare_new_config( - workspace_id=workspace, - tracing_provider="langfuse", - tracing_config={"public_key": "synthetic-public", "secret_key": "synthetic-secret", "host": "https://old.test"}, - ) - updated = gateway.prepare_updated_config( - workspace_id=workspace, - tracing_provider="langfuse", - tracing_config={"public_key": "******", "secret_key": "******", "host": "https://new.test"}, - current_tracing_config=current, - ) - assert updated["project_url"] == "https://new.test/project/verified-project" - assert current["project_url"] == "https://old.test/project/verified-project" - assert boundary.calls[-1][0] == "key" - assert boundary.calls[-1][2]["secret_key"] == "synthetic-secret" - - -def test_failed_optional_link_lookup_does_not_reuse_stale_project_url(boundary: ProviderBoundary) -> None: - gateway = OpsTraceManagerGateway() - workspace = str(uuid4()) - current = gateway.prepare_new_config( - workspace_id=workspace, - tracing_provider="langsmith", - tracing_config={"api_key": "synthetic-key", "project": "original"}, - ) - boundary.lookup_fails = True - updated = gateway.prepare_updated_config( - workspace_id=workspace, - tracing_provider="langsmith", - tracing_config={"api_key": "******", "project": "changed"}, - current_tracing_config=current, - ) - assert updated["project"] == "changed" - assert "project_url" not in updated - assert boundary.calls[-1][0] == "url" - - -def test_invalid_updated_configuration_never_resolves_a_project_url(boundary: ProviderBoundary) -> None: - boundary.effective = False - with pytest.raises(AppTracingConfigVerificationFailedError): - OpsTraceManagerGateway().prepare_updated_config( - workspace_id=str(uuid4()), - tracing_provider="langsmith", - tracing_config={"api_key": "synthetic-key", "project": "changed"}, - current_tracing_config=None, - ) - assert [call[0] for call in boundary.calls] == ["verify"] - - -def test_updated_project_url_is_frozen_and_used_by_real_langsmith_serializer( - boundary: ProviderBoundary, - sqlite_session_factory: sessionmaker[Session], - monkeypatch: pytest.MonkeyPatch, -) -> None: - workspace, app_id, message_id = (str(uuid4()) for _ in range(3)) - with sqlite_session_factory.begin() as session: - session.add( - make_app( - app_id=app_id, - tenant_id=workspace, - mode=AppMode.AGENT, - tracing=json.dumps({"enabled": True, "tracing_provider": "langsmith"}), - ) - ) - session.add( - make_message( - message_id=message_id, - app_id=app_id, - inputs={}, - query="test input", - message={}, - answer="test output", - message_unit_price=Decimal(0), - answer_unit_price=Decimal(0), - currency="USD", - from_source=ConversationFromSource.CONSOLE, - ) - ) - store = SQLAlchemyAppTracingConfigRepository(session_factory=sqlite_session_factory) - configs = AppTracingConfigService(configs=store, provider=OpsTraceManagerGateway()) - context = RequestContext("request", None, "account", workspace) - configs.create(context, app_id, "langsmith", {"api_key": "synthetic-key", "project": "original"}) - configs.update(context, app_id, "langsmith", {"api_key": "******", "project": "changed"}) - repository = ApplicationTraceRepository(sqlite_session_factory) - run = repository.start( - tenant_id=workspace, app_id=app_id, source_kind="message", source_id=message_id, session_id="test-session" - ) - assert run is not None - assert run.provider_config["project_url"] == "https://tracing.test/langsmith/changed" - configs.update(context, app_id, "langsmith", {"api_key": "******", "project": "later"}) - frozen = repository.source_run(tenant_id=workspace, app_id=app_id, source_kind="message", source_id=message_id) - assert frozen is not None - assert frozen.provider_config["project_url"] == "https://tracing.test/langsmith/changed" - - requests: list[str] = [] - - def send(method: str, url: str, **_kwargs: object) -> httpx.Response: - requests.append(url) - return httpx.Response(200, request=httpx.Request(method, url)) - - monkeypatch.setattr(ssrf_proxy, "make_request", send) - config = dict(frozen.provider_config) - config.update(OpsTraceManager.decrypt_tracing_config(workspace, "langsmith", frozen.provider_config)) - start = datetime(2026, 10, 16, tzinfo=UTC) - trace = CanonicalTrace( - trace_id=frozen.trace_id, - session_id=frozen.session_id, - root_span_id=frozen.root_span_id, - spans=( - CanonicalSpan( - id=frozen.root_span_id, - name="Agent", - kind=CanonicalSpanKind.AGENT, - parent_id=None, - status=CanonicalSpanStatus.OK, - start_time=start, - end_time=start + timedelta(seconds=1), - ), - ), - ) - result = export_trace("langsmith", config, trace) - assert result.status == "accepted" - assert result.trace_url == f"https://tracing.test/langsmith/changed/r/{result.trace_id}" - assert result.span_urls[frozen.root_span_id] == result.trace_url - assert len(requests) == 1 - assert boundary.calls[-1][2]["project"] == "later" diff --git a/api/tests/unit_tests/services/test_app_tracing_config_service.py b/api/tests/unit_tests/services/test_app_tracing_config_service.py index 8325cf96bf1..c4ec2569ae5 100644 --- a/api/tests/unit_tests/services/test_app_tracing_config_service.py +++ b/api/tests/unit_tests/services/test_app_tracing_config_service.py @@ -1,16 +1,237 @@ +import json +from dataclasses import dataclass from datetime import UTC, datetime from unittest.mock import MagicMock import pytest +from sqlalchemy import event, select +from sqlalchemy.orm import Session, sessionmaker from machinery.context import RequestContext +from models.model import App, AppMode, TraceAppConfig +from repositories.app.console_repository import ConsoleAppRepository +from repositories.app.tracing_config_repository import SQLAlchemyAppTracingConfigRepository +from services.app_tracing_config_gateway import OpsTraceManagerGateway from services.app_tracing_config_service import ( + AppTraceSettings, AppTracingConfigAlreadyExistsError, + AppTracingConfigConflictError, AppTracingConfigNotFoundError, AppTracingConfigRecord, AppTracingConfigService, + AppTracingConfigVerificationFailedError, ) +"""Public tracing operations keep credentials and policy in one transaction.""" + + +CONTEXT = RequestContext("request", None, "account", "tenant") + + +@dataclass +class Case: + sessions: sessionmaker[Session] + store: SQLAlchemyAppTracingConfigRepository + service: AppTracingConfigService + gateway: OpsTraceManagerGateway + reader: ConsoleAppRepository + + def policy(self) -> AppTraceSettings: + return self.reader.get_trace(CONTEXT, "app") + + def config(self, provider: str = "langfuse") -> AppTracingConfigRecord | None: + return self.store.get(workspace_id="tenant", app_id="app", tracing_provider=provider) + + +class TestAppTracingAtomicOperationsCases: + @staticmethod + @pytest.fixture + def case(sqlite_session_factory: sessionmaker[Session], monkeypatch: pytest.MonkeyPatch) -> Case: + with sqlite_session_factory.begin() as session: + session.add( + App( + id="app", + tenant_id="tenant", + name="Trace", + mode=AppMode.WORKFLOW, + enable_site=False, + enable_api=False, + tracing=json.dumps( + {"enabled": False, "tracing_provider": None, "capture_content": True, "capture_agent": False} + ), + ) + ) + gateway = OpsTraceManagerGateway() + + def prepare(*, tracing_config: dict[str, object], **_: object) -> dict[str, object]: + return dict(tracing_config) + + monkeypatch.setattr(gateway, "prepare_new_config", prepare) + monkeypatch.setattr(gateway, "prepare_updated_config", prepare) + monkeypatch.setattr(gateway, "validate_saved_config", lambda **_: None) + store = SQLAlchemyAppTracingConfigRepository(session_factory=sqlite_session_factory) + return Case( + sqlite_session_factory, + store, + AppTracingConfigService(configs=store, provider=gateway), + gateway, + ConsoleAppRepository(session_factory=sqlite_session_factory), + ) + + @staticmethod + @pytest.mark.parametrize("activate", [False, True]) + def test_create_and_update_preserve_capture_and_explicit_activation(case: Case, activate: bool) -> None: + case.service.create(CONTEXT, "app", "langfuse", {"secret": "original"}, activate=activate) + assert case.policy() == AppTraceSettings(activate, "langfuse" if activate else None, True, False) + case.service.update(CONTEXT, "app", "langfuse", {"secret": "updated"}, activate=True) + assert case.policy() == AppTraceSettings(True, "langfuse", True, False) + record = case.config() + assert record is not None + assert record.tracing_config == {"secret": "updated"} + + @staticmethod + def test_save_activation_uses_capture_policy_changed_during_vendor_validation( + case: Case, monkeypatch: pytest.MonkeyPatch + ) -> None: + def prepare(**_: object) -> dict[str, object]: + case.service.set_trace(CONTEXT, "app", AppTraceSettings(False, None, False, True)) + return {"secret": "validated"} + + monkeypatch.setattr(case.gateway, "prepare_new_config", prepare) + case.service.create(CONTEXT, "app", "langfuse", {"secret": "input"}, activate=True) + assert case.policy() == AppTraceSettings(True, "langfuse", False, True) + + @staticmethod + @pytest.mark.parametrize("operation", ["create", "update", "delete"]) + def test_transaction_failure_rolls_back_credentials_and_policy(case: Case, operation: str) -> None: + if operation != "create": + case.service.create(CONTEXT, "app", "langfuse", {"secret": "original"}, activate=True) + before_policy, before_record = case.policy(), case.config() + + def fail_commit(session: Session) -> None: + session.flush() + raise RuntimeError("commit failed") + + event.listen(case.sessions, "before_commit", fail_commit) + try: + + def mutate() -> None: + if operation == "delete": + case.service.delete(CONTEXT, "app", "langfuse") + elif operation == "create": + case.service.create(CONTEXT, "app", "langfuse", {"secret": "new"}, activate=True) + else: + case.service.update(CONTEXT, "app", "langfuse", {"secret": "new"}, activate=True) + + with pytest.raises(RuntimeError, match="commit failed"): + mutate() + finally: + event.remove(case.sessions, "before_commit", fail_commit) + assert case.policy() == before_policy + assert case.config() == before_record + + @staticmethod + @pytest.mark.parametrize("enabled", [False, True]) + def test_delete_current_clears_selection_but_other_provider_preserves_policy(case: Case, enabled: bool) -> None: + case.service.create(CONTEXT, "app", "langfuse", {"secret": "selected"}, activate=True) + case.service.create(CONTEXT, "app", "langsmith", {"secret": "other"}) + policy = AppTraceSettings(enabled, "langfuse", True, False) + case.service.set_trace(CONTEXT, "app", policy) + case.service.delete(CONTEXT, "app", "langsmith") + assert case.policy() == policy + case.service.delete(CONTEXT, "app", "langfuse") + assert case.policy() == AppTraceSettings(False, None, True, False) + assert case.config() is None + + @staticmethod + def test_validation_failure_writes_neither_config_nor_policy(case: Case, monkeypatch: pytest.MonkeyPatch) -> None: + def reject(**_: object) -> dict[str, object]: + raise AppTracingConfigVerificationFailedError + + monkeypatch.setattr(case.gateway, "prepare_new_config", reject) + with pytest.raises(AppTracingConfigVerificationFailedError): + case.service.create(CONTEXT, "app", "langfuse", {"secret": "bad"}, activate=True) + assert case.config() is None + assert case.policy() == AppTraceSettings(False, None, True, False) + + @staticmethod + @pytest.mark.parametrize("state", ["missing", "empty", "inactive"]) + def test_enable_requires_usable_owned_config(case: Case, state: str) -> None: + if state != "missing": + case.service.create(CONTEXT, "app", "langfuse", {} if state == "empty" else {"secret": "valid"}) + if state == "inactive": + with case.sessions.begin() as session: + record = session.scalar(select(TraceAppConfig)) + assert record is not None + record.is_active = False + with pytest.raises(AppTracingConfigNotFoundError): + case.service.set_trace(CONTEXT, "app", AppTraceSettings(True, "langfuse")) + assert case.policy() == AppTraceSettings(False, None, True, False) + + @staticmethod + def test_update_rejects_credentials_replaced_while_validation_was_in_flight( + case: Case, monkeypatch: pytest.MonkeyPatch + ) -> None: + case.service.create(CONTEXT, "app", "langfuse", {"secret": "original"}) + + def prepare(**_: object) -> dict[str, object]: + current = case.config() + assert current is not None + case.store.update( + workspace_id="tenant", + app_id="app", + tracing_provider="langfuse", + tracing_config={"secret": "concurrent"}, + expected=current, + ) + return {"secret": "stale"} + + monkeypatch.setattr(case.gateway, "prepare_updated_config", prepare) + with pytest.raises(AppTracingConfigConflictError): + case.service.update(CONTEXT, "app", "langfuse", {"secret": "masked"}, activate=True) + saved = case.config() + assert saved is not None + assert saved.tracing_config == {"secret": "concurrent"} + assert case.policy() == AppTraceSettings(False, None, True, False) + + @staticmethod + def test_enable_rejects_configuration_deleted_during_validation( + case: Case, monkeypatch: pytest.MonkeyPatch + ) -> None: + case.service.create(CONTEXT, "app", "langfuse", {"secret": "original"}) + + def validate(**_: object) -> None: + case.service.delete(CONTEXT, "app", "langfuse") + + monkeypatch.setattr(case.gateway, "validate_saved_config", validate) + with pytest.raises(AppTracingConfigConflictError): + case.service.set_trace(CONTEXT, "app", AppTraceSettings(True, "langfuse")) + assert case.policy() == AppTraceSettings(False, None, True, False) + + @staticmethod + @pytest.mark.parametrize("recreate", [False, True]) + def test_patch_rejects_configuration_deleted_during_validation( + case: Case, monkeypatch: pytest.MonkeyPatch, recreate: bool + ) -> None: + case.service.create(CONTEXT, "app", "langfuse", {"secret": "original"}) + + def prepare(**_: object) -> dict[str, object]: + case.service.delete(CONTEXT, "app", "langfuse") + if recreate: + case.service.create(CONTEXT, "app", "langfuse", {"secret": "original"}) + return {"secret": "stale"} + + monkeypatch.setattr(case.gateway, "prepare_updated_config", prepare) + with pytest.raises(AppTracingConfigConflictError): + case.service.update(CONTEXT, "app", "langfuse", {"secret": "masked"}, activate=True) + assert case.policy() == AppTraceSettings(False, None, True, False) + saved = case.config() + if recreate: + assert saved is not None + assert saved.tracing_config == {"secret": "original"} + else: + assert saved is None + def _context() -> RequestContext: return RequestContext("request-1", None, "account-1", "workspace-1") @@ -29,204 +250,206 @@ def _record(*, tracing_config: dict[str, object] | None = None) -> AppTracingCon ) -def test_get_scopes_lookup_to_workspace_and_presents_config() -> None: - stored = _record(tracing_config={"api_key": "encrypted"}) - configs = MagicMock() - configs.get.return_value = stored - provider = MagicMock() - provider.present_config.return_value = {"api_key": "******", "project_url": "https://example.com"} - service = AppTracingConfigService(configs=configs, provider=provider) +class TestAppTracingConfigServiceCases: + @staticmethod + def test_get_scopes_lookup_to_workspace_and_presents_config() -> None: + stored = _record(tracing_config={"api_key": "encrypted"}) + configs = MagicMock() + configs.get.return_value = stored + provider = MagicMock() + provider.present_config.return_value = {"api_key": "******", "project_url": "https://example.com"} + service = AppTracingConfigService(configs=configs, provider=provider) - result = service.get(_context(), "app-1", "arize") + result = service.get(_context(), "app-1", "arize") - assert result == _record(tracing_config={"api_key": "******", "project_url": "https://example.com"}) - assert stored.tracing_config == {"api_key": "encrypted"} - provider.validate_provider.assert_called_once_with("arize") - configs.get.assert_called_once_with( - workspace_id="workspace-1", - app_id="app-1", - tracing_provider="arize", - ) - provider.present_config.assert_called_once_with( - workspace_id="workspace-1", - tracing_provider="arize", - tracing_config={"api_key": "encrypted"}, - ) + assert result == _record(tracing_config={"api_key": "******", "project_url": "https://example.com"}) + assert stored.tracing_config == {"api_key": "encrypted"} + provider.validate_provider.assert_called_once_with("arize") + configs.get.assert_called_once_with( + workspace_id="workspace-1", + app_id="app-1", + tracing_provider="arize", + ) + provider.present_config.assert_called_once_with( + workspace_id="workspace-1", + tracing_provider="arize", + tracing_config={"api_key": "encrypted"}, + ) + @staticmethod + def test_get_returns_none_without_calling_provider_for_missing_config() -> None: + configs = MagicMock() + configs.get.return_value = None + provider = MagicMock() + service = AppTracingConfigService(configs=configs, provider=provider) -def test_get_returns_none_without_calling_provider_for_missing_config() -> None: - configs = MagicMock() - configs.get.return_value = None - provider = MagicMock() - service = AppTracingConfigService(configs=configs, provider=provider) + assert service.get(_context(), "app-1", "arize") is None - assert service.get(_context(), "app-1", "arize") is None + provider.validate_provider.assert_called_once_with("arize") + provider.present_config.assert_not_called() - provider.validate_provider.assert_called_once_with("arize") - provider.present_config.assert_not_called() + @staticmethod + def test_get_validates_provider_before_store_lookup() -> None: + events: list[str] = [] + configs = MagicMock() + configs.get.side_effect = lambda **_: events.append("store.get") + provider = MagicMock() + provider.validate_provider.side_effect = lambda *_: events.append("provider.validate") + service = AppTracingConfigService(configs=configs, provider=provider) + assert service.get(_context(), "app-1", "arize") is None -def test_get_validates_provider_before_store_lookup() -> None: - events: list[str] = [] - configs = MagicMock() - configs.get.side_effect = lambda **_: events.append("store.get") - provider = MagicMock() - provider.validate_provider.side_effect = lambda *_: events.append("provider.validate") - service = AppTracingConfigService(configs=configs, provider=provider) + assert events == ["provider.validate", "store.get"] - assert service.get(_context(), "app-1", "arize") is None + @staticmethod + def test_create_calls_provider_between_separate_store_operations() -> None: + events: list[str] = [] + configs = MagicMock() - assert events == ["provider.validate", "store.get"] + def get_config(**_: object) -> None: + events.append("store.get:closed") + def create_config(**_: object) -> bool: + events.append("store.create:closed") + return True -def test_create_calls_provider_between_separate_store_operations() -> None: - events: list[str] = [] - configs = MagicMock() + configs.get.side_effect = get_config + configs.create.side_effect = create_config + provider = MagicMock() - def get_config(**_: object) -> None: - events.append("store.get:closed") + def prepare_config(**_: object) -> dict[str, str]: + events.append("provider.prepare") + return {"api_key": "encrypted"} - def create_config(**_: object) -> bool: - events.append("store.create:closed") - return True + provider.prepare_new_config.side_effect = prepare_config + service = AppTracingConfigService(configs=configs, provider=provider) - configs.get.side_effect = get_config - configs.create.side_effect = create_config - provider = MagicMock() - - def prepare_config(**_: object) -> dict[str, str]: - events.append("provider.prepare") - return {"api_key": "encrypted"} - - provider.prepare_new_config.side_effect = prepare_config - service = AppTracingConfigService(configs=configs, provider=provider) - - service.create(_context(), "app-1", "arize", {"api_key": "plain"}) - - assert events == ["store.get:closed", "provider.prepare", "store.create:closed"] - provider.prepare_new_config.assert_called_once_with( - workspace_id="workspace-1", - tracing_provider="arize", - tracing_config={"api_key": "plain"}, - ) - configs.create.assert_called_once_with( - workspace_id="workspace-1", - app_id="app-1", - tracing_provider="arize", - tracing_config={"api_key": "encrypted"}, - activate=False, - ) - - -def test_create_reports_duplicate_before_provider_preparation() -> None: - configs = MagicMock() - configs.get.return_value = _record(tracing_config={"api_key": "encrypted"}) - provider = MagicMock() - provider.prepare_new_config.return_value = {"api_key": "new-encrypted"} - service = AppTracingConfigService(configs=configs, provider=provider) - - with pytest.raises(AppTracingConfigAlreadyExistsError, match="Trace config is exist"): service.create(_context(), "app-1", "arize", {"api_key": "plain"}) - provider.prepare_new_config.assert_not_called() - configs.create.assert_not_called() + assert events == ["store.get:closed", "provider.prepare", "store.create:closed"] + provider.prepare_new_config.assert_called_once_with( + workspace_id="workspace-1", + tracing_provider="arize", + tracing_config={"api_key": "plain"}, + ) + configs.create.assert_called_once_with( + workspace_id="workspace-1", + app_id="app-1", + tracing_provider="arize", + tracing_config={"api_key": "encrypted"}, + activate=False, + ) + @staticmethod + def test_create_reports_duplicate_before_provider_preparation() -> None: + configs = MagicMock() + configs.get.return_value = _record(tracing_config={"api_key": "encrypted"}) + provider = MagicMock() + provider.prepare_new_config.return_value = {"api_key": "new-encrypted"} + service = AppTracingConfigService(configs=configs, provider=provider) -def test_create_reports_duplicate_detected_during_write() -> None: - configs = MagicMock() - configs.get.return_value = None - configs.create.return_value = False - provider = MagicMock() - provider.prepare_new_config.return_value = {"api_key": "encrypted"} - service = AppTracingConfigService(configs=configs, provider=provider) + with pytest.raises(AppTracingConfigAlreadyExistsError, match="Trace config is exist"): + service.create(_context(), "app-1", "arize", {"api_key": "plain"}) - with pytest.raises(AppTracingConfigAlreadyExistsError): - service.create(_context(), "app-1", "arize", {"api_key": "plain"}) + provider.prepare_new_config.assert_not_called() + configs.create.assert_not_called() + @staticmethod + def test_create_reports_duplicate_detected_during_write() -> None: + configs = MagicMock() + configs.get.return_value = None + configs.create.return_value = False + provider = MagicMock() + provider.prepare_new_config.return_value = {"api_key": "encrypted"} + service = AppTracingConfigService(configs=configs, provider=provider) -def test_update_validates_reads_prepares_and_writes_in_order() -> None: - events: list[str] = [] - current = _record(tracing_config={"api_key": "old-encrypted"}) - configs = MagicMock() - configs.get.side_effect = lambda **_: events.append("store.get:closed") or current - configs.update.side_effect = lambda **_: events.append("store.update:closed") or True - provider = MagicMock() - provider.validate_provider.side_effect = lambda *_: events.append("provider.validate") - provider.prepare_updated_config.side_effect = lambda **_: ( - events.append("provider.prepare") or {"api_key": "new-encrypted"} - ) - service = AppTracingConfigService(configs=configs, provider=provider) + with pytest.raises(AppTracingConfigAlreadyExistsError): + service.create(_context(), "app-1", "arize", {"api_key": "plain"}) - service.update(_context(), "app-1", "arize", {"api_key": "******"}) + @staticmethod + def test_update_validates_reads_prepares_and_writes_in_order() -> None: + events: list[str] = [] + current = _record(tracing_config={"api_key": "old-encrypted"}) + configs = MagicMock() + configs.get.side_effect = lambda **_: events.append("store.get:closed") or current + configs.update.side_effect = lambda **_: events.append("store.update:closed") or True + provider = MagicMock() + provider.validate_provider.side_effect = lambda *_: events.append("provider.validate") + provider.prepare_updated_config.side_effect = lambda **_: ( + events.append("provider.prepare") or {"api_key": "new-encrypted"} + ) + service = AppTracingConfigService(configs=configs, provider=provider) - assert events == ["store.get:closed", "provider.validate", "provider.prepare", "store.update:closed"] - provider.prepare_updated_config.assert_called_once_with( - workspace_id="workspace-1", - tracing_provider="arize", - tracing_config={"api_key": "******"}, - current_tracing_config={"api_key": "old-encrypted"}, - ) - configs.update.assert_called_once_with( - workspace_id="workspace-1", - app_id="app-1", - tracing_provider="arize", - tracing_config={"api_key": "new-encrypted"}, - expected=current, - activate=False, - ) - - -def test_update_reports_config_missing_before_provider_preparation() -> None: - configs = MagicMock() - configs.get.return_value = None - provider = MagicMock() - service = AppTracingConfigService(configs=configs, provider=provider) - - with pytest.raises(AppTracingConfigNotFoundError, match="Trace config not exist"): - service.update(_context(), "app-1", "arize", {}) - - provider.validate_provider.assert_called_once_with("arize") - provider.prepare_updated_config.assert_not_called() - configs.update.assert_not_called() - - -def test_update_reports_config_removed_before_write() -> None: - configs = MagicMock() - configs.get.return_value = _record(tracing_config={"api_key": "old-encrypted"}) - configs.update.return_value = False - provider = MagicMock() - provider.prepare_updated_config.return_value = {"api_key": "new-encrypted"} - service = AppTracingConfigService(configs=configs, provider=provider) - - with pytest.raises(AppTracingConfigNotFoundError): service.update(_context(), "app-1", "arize", {"api_key": "******"}) + assert events == ["store.get:closed", "provider.validate", "provider.prepare", "store.update:closed"] + provider.prepare_updated_config.assert_called_once_with( + workspace_id="workspace-1", + tracing_provider="arize", + tracing_config={"api_key": "******"}, + current_tracing_config={"api_key": "old-encrypted"}, + ) + configs.update.assert_called_once_with( + workspace_id="workspace-1", + app_id="app-1", + tracing_provider="arize", + tracing_config={"api_key": "new-encrypted"}, + expected=current, + activate=False, + ) -def test_delete_scopes_to_workspace_and_reports_missing_config() -> None: - configs = MagicMock() - configs.delete.return_value = False - provider = MagicMock() - service = AppTracingConfigService(configs=configs, provider=provider) + @staticmethod + def test_update_reports_config_missing_before_provider_preparation() -> None: + configs = MagicMock() + configs.get.return_value = None + provider = MagicMock() + service = AppTracingConfigService(configs=configs, provider=provider) + + with pytest.raises(AppTracingConfigNotFoundError, match="Trace config not exist"): + service.update(_context(), "app-1", "arize", {}) + + provider.validate_provider.assert_called_once_with("arize") + provider.prepare_updated_config.assert_not_called() + configs.update.assert_not_called() + + @staticmethod + def test_update_reports_config_removed_before_write() -> None: + configs = MagicMock() + configs.get.return_value = _record(tracing_config={"api_key": "old-encrypted"}) + configs.update.return_value = False + provider = MagicMock() + provider.prepare_updated_config.return_value = {"api_key": "new-encrypted"} + service = AppTracingConfigService(configs=configs, provider=provider) + + with pytest.raises(AppTracingConfigNotFoundError): + service.update(_context(), "app-1", "arize", {"api_key": "******"}) + + @staticmethod + def test_delete_scopes_to_workspace_and_reports_missing_config() -> None: + configs = MagicMock() + configs.delete.return_value = False + provider = MagicMock() + service = AppTracingConfigService(configs=configs, provider=provider) + + with pytest.raises(AppTracingConfigNotFoundError): + service.delete(_context(), "app-1", "arize") + + configs.delete.assert_called_once_with( + workspace_id="workspace-1", + app_id="app-1", + tracing_provider="arize", + ) + provider.validate_provider.assert_called_once_with("arize") + + @staticmethod + def test_delete_validates_provider_before_store_delete() -> None: + events: list[str] = [] + configs = MagicMock() + configs.delete.side_effect = lambda **_: events.append("store.delete") or True + provider = MagicMock() + provider.validate_provider.side_effect = lambda *_: events.append("provider.validate") + service = AppTracingConfigService(configs=configs, provider=provider) - with pytest.raises(AppTracingConfigNotFoundError): service.delete(_context(), "app-1", "arize") - configs.delete.assert_called_once_with( - workspace_id="workspace-1", - app_id="app-1", - tracing_provider="arize", - ) - provider.validate_provider.assert_called_once_with("arize") - - -def test_delete_validates_provider_before_store_delete() -> None: - events: list[str] = [] - configs = MagicMock() - configs.delete.side_effect = lambda **_: events.append("store.delete") or True - provider = MagicMock() - provider.validate_provider.side_effect = lambda *_: events.append("provider.validate") - service = AppTracingConfigService(configs=configs, provider=provider) - - service.delete(_context(), "app-1", "arize") - - assert events == ["provider.validate", "store.delete"] + assert events == ["provider.validate", "store.delete"] diff --git a/api/tests/unit_tests/services/test_archive_workflow_run_logs.py b/api/tests/unit_tests/services/test_archive_workflow_run_logs.py deleted file mode 100644 index 0c06aff568c..00000000000 --- a/api/tests/unit_tests/services/test_archive_workflow_run_logs.py +++ /dev/null @@ -1,55 +0,0 @@ -""" -Unit tests for workflow run archiving functionality. - -This module contains tests for: -- Archive service -- Rollback service -""" - -from collections.abc import Callable -from datetime import datetime -from unittest.mock import MagicMock - -from enums import DeploymentEdition - - -class TestWorkflowRunArchiver: - """Tests for the WorkflowRunArchiver class.""" - - def test_archiver_initialization(self, config_overrides: Callable[..., None]) -> None: - """Test archiver can be initialized with various options.""" - from services.retention.workflow_run.archive_paid_plan_workflow_run import WorkflowRunArchiver - - config_overrides(DEPLOYMENT_EDITION=DeploymentEdition.COMMUNITY) - - archiver = WorkflowRunArchiver( - days=90, - batch_size=100, - tenant_ids=["test-tenant"], - limit=50, - dry_run=True, - ) - - assert archiver.days == 90 - assert archiver.batch_size == 100 - assert archiver.tenant_ids == ["test-tenant"] - assert archiver.limit == 50 - assert archiver.dry_run is True - - def test_get_bundle_manifest_key(self) -> None: - """Test V2 bundle manifest key generation.""" - from services.retention.workflow_run.archive_paid_plan_workflow_run import WorkflowRunArchiver - - archiver = WorkflowRunArchiver(run_shard_index=1, run_shard_total=4) - - mock_run = MagicMock() - mock_run.tenant_id = "9enant-123" - mock_run.id = "run-456" - mock_run.created_at = datetime(2024, 1, 15, 12, 0, 0) - - identity = archiver._build_bundle_identity([mock_run]) - key = archiver._get_manifest_object_key(identity) - - assert key.endswith("/manifest.json") - assert "workflow-runs/v2/tenant_prefix=9/tenant_id=9enant-123/year=2024/month=01" in key - assert "/shard=01-of-04/" in key diff --git a/api/tests/unit_tests/services/test_arize_tracing_project_url.py b/api/tests/unit_tests/services/test_arize_tracing_project_url.py deleted file mode 100644 index f11cf7421a2..00000000000 --- a/api/tests/unit_tests/services/test_arize_tracing_project_url.py +++ /dev/null @@ -1,292 +0,0 @@ -"""Arize project links use verified metadata without changing ingestion permissions.""" - -import json -from collections import deque -from copy import deepcopy -from decimal import Decimal -from typing import Never -from uuid import uuid4 - -import httpx -import pytest -from dify_trace_arize_phoenix import arize_configuration -from dify_trace_arize_phoenix.arize_configuration import DESTINATION_KEY -from dify_trace_arize_phoenix.config import ArizeConfig -from sqlalchemy import select -from sqlalchemy.orm import Session, sessionmaker - -from core.helper import ssrf_proxy -from core.ops.ops_trace_manager import OpsTraceManager -from machinery.context import RequestContext -from models.application_trace import ApplicationTraceRun -from models.enums import ConversationFromSource -from models.model import AppMode -from repositories.app.tracing_config_repository import SQLAlchemyAppTracingConfigRepository -from repositories.application_trace_repository import ApplicationTraceRepository -from services.app_tracing_config_gateway import OpsTraceManagerGateway -from services.app_tracing_config_service import AppTracingConfigService, AppTracingConfigVerificationFailedError -from services.tracing.provider_dispatch import frozen_project_url -from services.tracing.service import ApplicationTracingService -from tests.unit_tests.model_factories import make_app, make_message -from tests.unit_tests.services.tracing.test_arize_configuration import ( - DESTINATION, - SPACE, - URL, - VALUES, - page, - project, - response, - space, -) -from tests.unit_tests.services.tracing.test_arize_configuration import ( - calls as calls, # noqa: PLC0414 - Re-export the HTTP boundary fixture for pytest. -) -from tests.unit_tests.services.tracing.test_arize_configuration import ( - responses as responses, # noqa: PLC0414 - Re-export the HTTP boundary fixture for pytest. -) - -pytest_plugins = ("tests.unit_tests.services.tracing.test_nullable_tracing_credentials",) - - -@pytest.fixture -def verified(monkeypatch: pytest.MonkeyPatch) -> list[dict[str, object]]: - calls: list[dict[str, object]] = [] - - def verify(config: ArizeConfig) -> None: - calls.append(config.model_dump()) - - monkeypatch.setattr(arize_configuration, "verify_configuration", verify) - return calls - - -def test_legacy_arize_get_hides_unverified_redirect_without_network_or_rewriting( - workspace: str, monkeypatch: pytest.MonkeyPatch -) -> None: - stored = OpsTraceManager.encrypt_tracing_config( - workspace, - "arize", - {"api_key": "synthetic-ingestion-key", "space_id": "U3BhY2U6dGVzdA==", "project": "original"}, - ) - stored["project_url"] = "https://app.arize.com/?redirect_project_name=original" - untouched = deepcopy(stored) - calls: list[str] = [] - - def forbidden(*_args: object, **_kwargs: object) -> Never: - calls.append("legacy or network") - raise AssertionError("Configuration GET must not initialize a tracer or resolve vendor metadata") - - monkeypatch.setattr(OpsTraceManager, "get_trace_config_project_url", forbidden) - monkeypatch.setattr(ssrf_proxy, "get", forbidden) - shown = OpsTraceManagerGateway().present_config( - workspace_id=workspace, tracing_provider="arize", tracing_config=stored - ) - - assert shown["project_url"] is None - assert calls == [] - assert stored == untouched - - -def test_real_rsa_create_masked_update_get_and_frozen_history( - workspace: str, - sqlite_session_factory: sessionmaker[Session], - monkeypatch: pytest.MonkeyPatch, - verified: list[dict[str, object]], - calls: list[tuple[str, dict[str, object]]], - responses: deque[httpx.Response | Exception], -) -> None: - app_id, message_id = str(uuid4()), str(uuid4()) - with sqlite_session_factory.begin() as session: - session.add( - make_app( - app_id=app_id, - tenant_id=workspace, - mode=AppMode.AGENT, - tracing=json.dumps({"enabled": True, "tracing_provider": "arize"}), - ) - ) - session.add( - make_message( - message_id=message_id, - app_id=app_id, - inputs={}, - query="test", - message={}, - answer="ok", - message_unit_price=Decimal(0), - answer_unit_price=Decimal(0), - currency="USD", - from_source=ConversationFromSource.CONSOLE, - ) - ) - store = SQLAlchemyAppTracingConfigRepository(session_factory=sqlite_session_factory) - gateway = OpsTraceManagerGateway() - service = AppTracingConfigService(configs=store, provider=gateway) - context = RequestContext("request", None, "account", workspace) - legacy = OpsTraceManager.encrypt_tracing_config(workspace, "arize", VALUES) - legacy["project_url"] = "https://app.arize.com/?redirect_project_name=exact" - assert store.create(workspace_id=workspace, app_id=app_id, tracing_provider="arize", tracing_config=legacy) - repository = ApplicationTraceRepository(sqlite_session_factory) - run = repository.start( - tenant_id=workspace, app_id=app_id, source_kind="message", source_id=message_id, session_id="session" - ) - assert run is not None - with sqlite_session_factory.begin() as session: - accepted = session.get(ApplicationTraceRun, run.id) - assert accepted is not None - accepted.state = "accepted" - accepted.attempts = 1 - with sqlite_session_factory() as session: - accepted_row = deepcopy( - dict( - session.execute(select(ApplicationTraceRun.__table__).where(ApplicationTraceRun.id == run.id)) - .mappings() - .one() - ) - ) - frozen_before = deepcopy(run.provider_config) - - # Both legacy secret fields are really encrypted, rather than identity stubs. - assert legacy["api_key"] != VALUES["api_key"] - assert legacy["space_id"] != SPACE - responses.extend([space(), page(project())]) - submitted = { - **VALUES, - "api_key": "******", - "space_id": "******", - DESTINATION_KEY: {"organization_id": "forged", "project_id": "forged"}, - "project_url": "https://untrusted.test/forged", - } - untouched = deepcopy(submitted) - service.update(context, app_id, "arize", submitted) - current = store.get(workspace_id=workspace, app_id=app_id, tracing_provider="arize") - assert current is not None - assert current.tracing_config is not None - saved = current.tracing_config - assert saved["api_key"] == legacy["api_key"] - assert saved["space_id"] == legacy["space_id"] - assert saved[DESTINATION_KEY] == {**DESTINATION, "source_digest": saved[DESTINATION_KEY]["source_digest"]} - assert saved["project_url"] == URL - assert submitted == untouched - assert len(calls) == 2 - assert verified[-1]["api_key"] == VALUES["api_key"] - assert verified[-1]["space_id"] == SPACE - - def forbidden(*_args: object, **_kwargs: object) -> Never: - raise AssertionError("Reading current configuration or frozen history must not use vendor I/O") - - monkeypatch.setattr(ssrf_proxy, "get", forbidden) - monkeypatch.setattr(OpsTraceManager, "get_trace_config_project_url", forbidden) - shown = service.get(context, app_id, "arize") - assert shown is not None - assert shown.tracing_config is not None - assert shown.tracing_config["project_url"] == URL - assert DESTINATION_KEY not in shown.tracing_config - frozen = repository.source_run(tenant_id=workspace, app_id=app_id, source_kind="message", source_id=message_id) - assert frozen is not None - assert (frozen.state, frozen.attempts, frozen.provider_config) == ("accepted", 1, frozen_before) - assert frozen_project_url("arize", frozen.provider_config) is None - assert frozen.provider_config == frozen_before - monkeypatch.setattr(OpsTraceManager, "decrypt_tracing_config", forbidden) - status = ApplicationTracingService(repository).status( - tenant_id=workspace, app_id=app_id, source_kind="message", source_id=message_id - ) - assert status["project_url"] is None - assert status["status"] == "accepted" - with sqlite_session_factory() as session: - after_row = dict( - session.execute(select(ApplicationTraceRun.__table__).where(ApplicationTraceRun.id == run.id)) - .mappings() - .one() - ) - assert after_row == accepted_row - - # A future source freezes the verified destination; it never repairs the old run. - later_id = str(uuid4()) - with sqlite_session_factory.begin() as session: - session.add( - make_message( - message_id=later_id, - app_id=app_id, - inputs={}, - query="later", - message={}, - answer="ok", - message_unit_price=Decimal(0), - answer_unit_price=Decimal(0), - currency="USD", - from_source=ConversationFromSource.CONSOLE, - ) - ) - later = repository.start( - tenant_id=workspace, app_id=app_id, source_kind="message", source_id=later_id, session_id="session" - ) - assert later is not None - assert frozen_project_url("arize", later.provider_config) == URL - - -@pytest.mark.parametrize("status", [401, 403, 404, 500]) -def test_ingest_verified_key_without_metadata_read_permission_still_saves( - workspace: str, - verified: list[dict[str, object]], - calls: list[tuple[str, dict[str, object]]], - responses: deque[httpx.Response | Exception], - status: int, -) -> None: - responses.append(response({"message": "private-response"}, status)) - saved = OpsTraceManagerGateway().prepare_new_config( - workspace_id=workspace, tracing_provider="arize", tracing_config=VALUES - ) - assert "project_url" not in saved - assert DESTINATION_KEY not in saved - assert OpsTraceManager.decrypt_tracing_config(workspace, "arize", saved)["api_key"] == VALUES["api_key"] - assert len(verified) == len(calls) == 1 - - -def test_changed_destination_cannot_keep_old_or_client_supplied_cache( - workspace: str, - verified: list[dict[str, object]], - calls: list[tuple[str, dict[str, object]]], - responses: deque[httpx.Response | Exception], -) -> None: - gateway = OpsTraceManagerGateway() - responses.extend([space(), page(project())]) - before = gateway.prepare_new_config(workspace_id=workspace, tracing_provider="arize", tracing_config=VALUES) - assert before["project_url"] == URL - responses.append(response({}, 403)) - after = gateway.prepare_updated_config( - workspace_id=workspace, - tracing_provider="arize", - current_tracing_config=before, - tracing_config={ - **VALUES, - "api_key": "******", - "space_id": "******", - "project": "later", - DESTINATION_KEY: before[DESTINATION_KEY], - "project_url": URL, - }, - ) - assert after["project"] == "later" - assert after["api_key"] == before["api_key"] - assert after["space_id"] == before["space_id"] - assert "project_url" not in after - assert DESTINATION_KEY not in after - assert before["project_url"] == URL - assert len(calls) == 3 - assert len(verified) == 2 - - -def test_failed_original_ingestion_verification_never_resolves_metadata( - workspace: str, - monkeypatch: pytest.MonkeyPatch, - calls: list[tuple[str, dict[str, object]]], -) -> None: - def rejected(_config: ArizeConfig) -> Never: - raise arize_configuration.ArizeConfigurationError("permission_denied") - - monkeypatch.setattr(arize_configuration, "verify_configuration", rejected) - with pytest.raises(AppTracingConfigVerificationFailedError): - OpsTraceManagerGateway().prepare_new_config( - workspace_id=workspace, tracing_provider="arize", tracing_config=VALUES - ) - assert calls == [] diff --git a/api/tests/unit_tests/services/test_audio_inputs.py b/api/tests/unit_tests/services/test_audio_inputs.py deleted file mode 100644 index 23280c34bac..00000000000 --- a/api/tests/unit_tests/services/test_audio_inputs.py +++ /dev/null @@ -1,71 +0,0 @@ -"""Validate plain audio uploads before calling the shared provider boundary.""" - -from io import BytesIO - -import pytest - -from models import App, AppMode -from services import audio_provider_gateway -from services.audio_service import FILE_SIZE_LIMIT, AudioService -from services.audio_types import AudioAppRef, AudioUpload -from services.errors.audio import ( - AudioTooLargeServiceError, - NoAudioUploadedServiceError, - UnsupportedAudioTypeServiceError, -) - - -@pytest.mark.parametrize("size", [0, FILE_SIZE_LIMIT, FILE_SIZE_LIMIT + 1]) -def test_asr_accepts_empty_and_exact_limit_but_rejects_oversized_upload( - monkeypatch: pytest.MonkeyPatch, size: int -) -> None: - app = App(id="app", tenant_id="owner", mode=AppMode.CHAT) - contents: list[bytes] = [] - - def transcribe(*, app: AudioAppRef, content: bytes, end_user: str | None) -> str: - assert app == AudioAppRef(app_id="app", tenant_id="owner", app_mode="chat") - assert end_user == "provider-user" - contents.append(content) - return " transcript " - - monkeypatch.setattr(audio_provider_gateway, "speech_to_text", transcribe) - content = b"a" * size - upload = AudioUpload(stream=BytesIO(content), mime_type="audio/x-m4a") - if size > FILE_SIZE_LIMIT: - with pytest.raises(AudioTooLargeServiceError, match="Audio size larger than 30 mb"): - AudioService.invoke_speech_to_text(app, upload, end_user="provider-user") - assert contents == [] - else: - assert AudioService.invoke_speech_to_text(app, upload, end_user="provider-user") == {"text": " transcript "} - assert contents == [content] - - -@pytest.mark.parametrize("missing", [False, True]) -def test_invalid_upload_does_not_read_stream_or_invoke_provider(monkeypatch: pytest.MonkeyPatch, missing: bool) -> None: - def unexpected_provider(*, app: AudioAppRef, content: bytes, end_user: str | None) -> str: - del app, content, end_user - pytest.fail("Invalid audio reached the provider") - - monkeypatch.setattr(audio_provider_gateway, "speech_to_text", unexpected_provider) - app = App(id="app", tenant_id="owner", mode=AppMode.CHAT) - stream = BytesIO(b"unread") - audio = None if missing else AudioUpload(stream=stream, mime_type="application/octet-stream") - error = NoAudioUploadedServiceError if missing else UnsupportedAudioTypeServiceError - with pytest.raises(error): - AudioService.invoke_speech_to_text(app, audio) - assert stream.tell() == 0 - - -def test_provider_exception_propagates_without_losing_details(monkeypatch: pytest.MonkeyPatch) -> None: - error = RuntimeError("provider transcription failed: timeout") - - def transcribe(*, app: AudioAppRef, content: bytes, end_user: str | None) -> str: - del app, content, end_user - raise error - - monkeypatch.setattr(audio_provider_gateway, "speech_to_text", transcribe) - app = App(id="app", tenant_id="owner", mode=AppMode.CHAT) - audio = AudioUpload(stream=BytesIO(b"input"), mime_type="audio/mp3") - with pytest.raises(RuntimeError) as caught: - AudioService.invoke_speech_to_text(app, audio) - assert caught.value is error diff --git a/api/tests/unit_tests/services/test_audio_service.py b/api/tests/unit_tests/services/test_audio_service.py index f4f46bdb213..2bb23a75591 100644 --- a/api/tests/unit_tests/services/test_audio_service.py +++ b/api/tests/unit_tests/services/test_audio_service.py @@ -1,3 +1,108 @@ +import json +from collections.abc import Generator, Iterator +from decimal import Decimal +from io import BytesIO +from uuid import uuid4 + +import pytest +from flask import Flask, has_request_context, request +from sqlalchemy.orm import Session +from werkzeug.datastructures import FileStorage + +from core.credit_usage import CreditUsageAppType, CreditUsageCreatedBy +from core.errors.error import ProviderTokenNotInitError +from core.model_manager import ModelManager +from core.plugin.entities.plugin_daemon import TTSAudioChunk +from extensions.ext_database import db +from graphon.model_runtime.errors.invoke import InvokeBadRequestError +from graphon.model_runtime.protocols.tts_runtime import TTSModelVoice +from models import App as _audio_inputs_App +from models import AppMode as _audio_inputs_AppMode +from models.agent import Agent, AgentConfigSnapshot, AgentScope, AgentSource +from models.agent_config_entities import AgentSoulConfig +from models.enums import ConversationFromSource, MessageStatus +from models.model import App as _audio_service_App +from models.model import AppMode as _audio_service_AppMode +from models.model import AppModelConfig, Message +from models.workflow import Workflow, WorkflowType +from services import audio_provider_gateway +from services.app_ref_service import AppRef, MessageRef +from services.audio_service import FILE_SIZE_LIMIT, AudioService, _create_tts_response +from services.audio_types import AudioAppRef, AudioUpload +from services.errors.audio import ( + AudioTooLargeServiceError, + NoAudioUploadedServiceError, + ProviderNotSupportSpeechToTextServiceError, + ProviderNotSupportTextToSpeechServiceError, + SpeechToTextDisabledServiceError, + UnsupportedAudioTypeServiceError, +) +from tests.helpers.unit.audio_runtime_fixtures import AudioRuntimeObservations +from tests.helpers.unit.model_factories import make_message + +"""Validate plain audio uploads before calling the shared provider boundary.""" + + +class TestAudioInputsCases: + @staticmethod + @pytest.mark.parametrize("size", [0, FILE_SIZE_LIMIT, FILE_SIZE_LIMIT + 1]) + def test_asr_accepts_empty_and_exact_limit_but_rejects_oversized_upload( + monkeypatch: pytest.MonkeyPatch, size: int + ) -> None: + app = _audio_inputs_App(id="app", tenant_id="owner", mode=_audio_inputs_AppMode.CHAT) + contents: list[bytes] = [] + + def transcribe(*, app: AudioAppRef, content: bytes, end_user: str | None) -> str: + assert app == AudioAppRef(app_id="app", tenant_id="owner", app_mode="chat") + assert end_user == "provider-user" + contents.append(content) + return " transcript " + + monkeypatch.setattr(audio_provider_gateway, "speech_to_text", transcribe) + content = b"a" * size + upload = AudioUpload(stream=BytesIO(content), mime_type="audio/x-m4a") + if size > FILE_SIZE_LIMIT: + with pytest.raises(AudioTooLargeServiceError, match="Audio size larger than 30 mb"): + AudioService.invoke_speech_to_text(app, upload, end_user="provider-user") + assert contents == [] + else: + assert AudioService.invoke_speech_to_text(app, upload, end_user="provider-user") == {"text": " transcript "} + assert contents == [content] + + @staticmethod + @pytest.mark.parametrize("missing", [False, True]) + def test_invalid_upload_does_not_read_stream_or_invoke_provider( + monkeypatch: pytest.MonkeyPatch, missing: bool + ) -> None: + def unexpected_provider(*, app: AudioAppRef, content: bytes, end_user: str | None) -> str: + del app, content, end_user + pytest.fail("Invalid audio reached the provider") + + monkeypatch.setattr(audio_provider_gateway, "speech_to_text", unexpected_provider) + app = _audio_inputs_App(id="app", tenant_id="owner", mode=_audio_inputs_AppMode.CHAT) + stream = BytesIO(b"unread") + audio = None if missing else AudioUpload(stream=stream, mime_type="application/octet-stream") + error = NoAudioUploadedServiceError if missing else UnsupportedAudioTypeServiceError + with pytest.raises(error): + AudioService.invoke_speech_to_text(app, audio) + assert stream.tell() == 0 + + @staticmethod + def test_provider_exception_propagates_without_losing_details(monkeypatch: pytest.MonkeyPatch) -> None: + error = RuntimeError("provider transcription failed: timeout") + + def transcribe(*, app: AudioAppRef, content: bytes, end_user: str | None) -> str: + del app, content, end_user + raise error + + monkeypatch.setattr(audio_provider_gateway, "speech_to_text", transcribe) + app = _audio_inputs_App(id="app", tenant_id="owner", mode=_audio_inputs_AppMode.CHAT) + audio = AudioUpload(stream=BytesIO(b"input"), mime_type="audio/mp3") + with pytest.raises(RuntimeError) as caught: + AudioService.invoke_speech_to_text(app, audio) + assert caught.value is error + + """ Comprehensive unit tests for AudioService. @@ -53,41 +158,6 @@ Tests available voice retrieval: - text_to_speech: Enables TTS functionality """ -import json -from collections.abc import Generator, Iterator -from decimal import Decimal -from io import BytesIO -from uuid import uuid4 - -import pytest -from flask import Flask, has_request_context, request -from sqlalchemy.orm import Session -from werkzeug.datastructures import FileStorage - -from core.credit_usage import CreditUsageAppType, CreditUsageCreatedBy -from core.errors.error import ProviderTokenNotInitError -from core.model_manager import ModelManager -from core.plugin.entities.plugin_daemon import TTSAudioChunk -from extensions.ext_database import db -from graphon.model_runtime.errors.invoke import InvokeBadRequestError -from graphon.model_runtime.protocols.tts_runtime import TTSModelVoice -from models.agent import Agent, AgentConfigSnapshot, AgentScope, AgentSource -from models.agent_config_entities import AgentSoulConfig -from models.enums import ConversationFromSource, MessageStatus -from models.model import App, AppMode, AppModelConfig, Message -from models.workflow import Workflow, WorkflowType -from services.app_ref_service import AppRef, MessageRef -from services.audio_service import AudioService, _create_tts_response -from services.errors.audio import ( - AudioTooLargeServiceError, - NoAudioUploadedServiceError, - ProviderNotSupportSpeechToTextServiceError, - ProviderNotSupportTextToSpeechServiceError, - SpeechToTextDisabledServiceError, - UnsupportedAudioTypeServiceError, -) -from tests.unit_tests.audio_runtime_fixtures import AudioRuntimeObservations -from tests.unit_tests.model_factories import make_message APP_ID = "11111111-1111-1111-1111-111111111111" TENANT_ID = "22222222-2222-2222-2222-222222222222" @@ -98,13 +168,6 @@ ACCOUNT_ID = "66666666-6666-6666-6666-666666666666" OTHER_ID = "77777777-7777-7777-7777-777777777777" -@pytest.fixture -def audio_provider(audio_runtime: AudioRuntimeObservations, app: Flask) -> Iterator[AudioRuntimeObservations]: - # Real plugin TTS returns a stream, so the compatibility response needs a request context. - with app.test_request_context("/text-to-audio", method="POST"): - yield audio_runtime - - def _message(*, answer: str = "Message answer") -> Message: return make_message( message_id=MESSAGE_ID, @@ -135,13 +198,13 @@ class AudioServiceTestDataFactory: def create_app_mock( self, app_id: str = APP_ID, - mode: AppMode = AppMode.CHAT, + mode: _audio_service_AppMode = _audio_service_AppMode.CHAT, tenant_id: str = TENANT_ID, *, workflow: Workflow | None = None, app_model_config: AppModelConfig | None = None, **kwargs: object, - ) -> App: + ) -> _audio_service_App: """ Create and persist an App model. @@ -154,7 +217,7 @@ class AudioServiceTestDataFactory: Returns: Persisted App model with specified attributes """ - app = App( + app = _audio_service_App( id=app_id, tenant_id=tenant_id, name="Audio test app", @@ -241,751 +304,764 @@ class AudioServiceTestDataFactory: return FileStorage(stream=BytesIO(content), filename=filename, content_type=mimetype) -@pytest.fixture -def factory(sqlite_session: Session) -> AudioServiceTestDataFactory: - """Provide the test data factory to all tests.""" - return AudioServiceTestDataFactory(sqlite_session) - - -@pytest.mark.parametrize("consumed_chunks", [0, 1, 2, None]) -def test_tts_response_closes_retained_provider_stream(consumed_chunks: int | None) -> None: - closed: list[bool] = [] - audio_chunks = [b"RIFF\x24\x00\x00\x00WAVE" + b"\x00" * 32, b"first-tail", b"second-tail"] - - def chunks() -> Generator[bytes]: - try: - for chunk in audio_chunks: - assert has_request_context() - assert request.path == "/text-to-audio" - yield chunk - finally: - closed.append(True) - - source = chunks() - app = Flask(__name__) - with app.test_request_context("/text-to-audio", method="POST"): - response = _create_tts_response(source, "audio/x-wav") - - assert response.status_code == 200 - assert dict(response.headers) == {"Content-Type": "audio/wav"} - assert response.is_streamed - assert closed == [] - try: - iterator = iter(response.response) - if consumed_chunks is None: - assert list(iterator) == audio_chunks - assert closed == [True] - else: - assert [next(iterator) for _ in range(consumed_chunks)] == audio_chunks[:consumed_chunks] - assert closed == [] - finally: - response.close() - response.close() - - assert closed == [True] - with pytest.raises(StopIteration): - next(source) - - -class TestAudioServiceASR: - """Test speech-to-text (ASR) operations.""" - - @pytest.fixture(autouse=True) - def _bind_sqlite_session(self, sqlite_session: Session) -> None: - self.session = sqlite_session - - def test_transcript_asr_success_chat_mode( - self, audio_provider: AudioRuntimeObservations, factory: AudioServiceTestDataFactory - ) -> None: - """Test successful ASR transcription in CHAT mode.""" - # Arrange - app_model_config = factory.create_app_model_config_mock(speech_to_text_dict={"enabled": True}) - app = factory.create_app_mock( - mode=AppMode.CHAT, - app_model_config=app_model_config, - ) - file = factory.create_file_storage() - - audio_provider.speech_to_text_result = "Transcribed text" - - # Act - result = AudioService.transcript_asr(app_model=app, file=file, session=self.session, end_user="user-123") - - # Assert - assert result == {"text": "Transcribed text"} - assert audio_provider.speech_to_text_calls == [b"fake audio content"] - assert audio_provider.runtime_calls[0]["request_metadata"] == { - "app_type": CreditUsageAppType.CHATBOT, - "created_by": CreditUsageCreatedBy.AUDIO, - } - assert audio_provider.manager_requests == [ - { - "tenant_id": app.tenant_id, - "user_id": "user-123", - "request_metadata": { - "app_type": CreditUsageAppType.CHATBOT, - "created_by": CreditUsageCreatedBy.AUDIO, - }, - }, - ] - - def test_transcript_asr_accepts_x_m4a_mimetype( - self, audio_provider: AudioRuntimeObservations, factory: AudioServiceTestDataFactory - ) -> None: - """Test that the x-m4a MIME alias follows the normal m4a transcription flow.""" - # Arrange - app_model_config = factory.create_app_model_config_mock(speech_to_text_dict={"enabled": True}) - app = factory.create_app_mock(mode=AppMode.CHAT, app_model_config=app_model_config) - file = factory.create_file_storage(filename="audio.m4a", mimetype="audio/x-m4a") - audio_provider.speech_to_text_result = "M4A transcript" - - # Act - result = AudioService.transcript_asr(app_model=app, file=file, session=self.session) - - # Assert - assert result == {"text": "M4A transcript"} - assert audio_provider.speech_to_text_calls == [b"fake audio content"] - - def test_transcript_asr_success_advanced_chat_mode( - self, audio_provider: AudioRuntimeObservations, factory: AudioServiceTestDataFactory - ) -> None: - """Test successful ASR transcription in ADVANCED_CHAT mode.""" - # Arrange - workflow = factory.create_workflow_mock(features_dict={"speech_to_text": {"enabled": True}}) - app = factory.create_app_mock( - mode=AppMode.ADVANCED_CHAT, - workflow=workflow, - ) - file = factory.create_file_storage() - - audio_provider.speech_to_text_result = "Workflow transcribed text" - - # Act - result = AudioService.transcript_asr(app_model=app, file=file, session=self.session) - - # Assert - assert result == {"text": "Workflow transcribed text"} - - def test_transcript_asr_success_published_agent_mode( - self, - audio_provider: AudioRuntimeObservations, - factory: AudioServiceTestDataFactory, - ) -> None: - app = factory.create_app_mock(mode=AppMode.AGENT) - file = factory.create_file_storage() - agent_soul = AgentSoulConfig.model_validate({"app_features": {"speech_to_text": {"enabled": True}}}) - agent = Agent( - tenant_id=app.tenant_id, - app_id=app.id, - name="Audio agent", - scope=AgentScope.ROSTER, - source=AgentSource.AGENT_APP, - ) - self.session.add(agent) - self.session.flush() - snapshot = AgentConfigSnapshot( - tenant_id=app.tenant_id, agent_id=agent.id, version=1, config_snapshot=agent_soul - ) - self.session.add(snapshot) - self.session.flush() - agent.active_config_snapshot_id = snapshot.id - self.session.commit() - audio_provider.speech_to_text_result = "Published Agent transcript" - - result = AudioService.transcript_asr(app_model=app, file=file, session=self.session, end_user="end-user-1") - - assert result == {"text": "Published Agent transcript"} - assert audio_provider.manager_requests == [ - { - "tenant_id": app.tenant_id, - "user_id": "end-user-1", - "request_metadata": { - "app_type": CreditUsageAppType.AGENT_V2, - "created_by": CreditUsageCreatedBy.AUDIO, - }, - }, - ] - - def test_transcript_asr_legacy_agent_falls_back_to_app_model_config( - self, - audio_provider: AudioRuntimeObservations, - factory: AudioServiceTestDataFactory, - ) -> None: - app_model_config = factory.create_app_model_config_mock(speech_to_text_dict={"enabled": True}) - app = factory.create_app_mock(mode=AppMode.AGENT, app_model_config=app_model_config) - file = factory.create_file_storage() - audio_provider.speech_to_text_result = "Legacy Agent transcript" - - result = AudioService.transcript_asr(app_model=app, file=file, session=self.session) - - assert result == {"text": "Legacy Agent transcript"} - - def test_transcript_agent_asr_uses_agent_soul_feature( - self, audio_provider: AudioRuntimeObservations, factory: AudioServiceTestDataFactory - ) -> None: - app = factory.create_app_mock(mode=AppMode.AGENT) - file = factory.create_file_storage() - agent_soul = AgentSoulConfig.model_validate({"app_features": {"speech_to_text": {"enabled": True}}}) - audio_provider.speech_to_text_result = "Agent transcript" - - result = AudioService.transcript_agent_asr( - app_model=app, - agent_soul=agent_soul, - file=file, - session=self.session, - end_user="account-1", - ) - - assert result == {"text": "Agent transcript"} - assert audio_provider.manager_requests == [ - { - "tenant_id": app.tenant_id, - "user_id": "account-1", - "request_metadata": { - "app_type": CreditUsageAppType.AGENT_V2, - "created_by": CreditUsageCreatedBy.AUDIO, - }, - }, - ] - - @pytest.mark.parametrize( - "agent_soul", - [ - AgentSoulConfig(), - AgentSoulConfig.model_validate({"app_features": {"speech_to_text": {"enabled": False}}}), - ], - ) - def test_transcript_agent_asr_rejects_disabled_feature( - self, factory: AudioServiceTestDataFactory, agent_soul: AgentSoulConfig - ) -> None: - app = factory.create_app_mock(mode=AppMode.AGENT) - file = factory.create_file_storage() - - with pytest.raises(SpeechToTextDisabledServiceError): - AudioService.transcript_agent_asr(app_model=app, agent_soul=agent_soul, file=file, session=self.session) - - def test_transcript_agent_asr_preserves_legacy_feature_fallback( - self, audio_provider: AudioRuntimeObservations, factory: AudioServiceTestDataFactory - ) -> None: - app_model_config = factory.create_app_model_config_mock(speech_to_text_dict={"enabled": True}) - app = factory.create_app_mock(mode=AppMode.AGENT, app_model_config=app_model_config) - file = factory.create_file_storage() - audio_provider.speech_to_text_result = "Legacy feature transcript" - - result = AudioService.transcript_agent_asr( - app_model=app, - agent_soul=AgentSoulConfig(), - file=file, - session=self.session, - ) - - assert result == {"text": "Legacy feature transcript"} - - def test_transcript_agent_asr_soul_disabled_overrides_legacy_feature( - self, factory: AudioServiceTestDataFactory - ) -> None: - app_model_config = factory.create_app_model_config_mock(speech_to_text_dict={"enabled": True}) - app = factory.create_app_mock(mode=AppMode.AGENT, app_model_config=app_model_config) - file = factory.create_file_storage() - agent_soul = AgentSoulConfig.model_validate({"app_features": {"speech_to_text": {"enabled": False}}}) - - with pytest.raises(SpeechToTextDisabledServiceError): - AudioService.transcript_agent_asr(app_model=app, agent_soul=agent_soul, file=file, session=self.session) - - def test_transcript_asr_raises_error_when_feature_disabled_chat_mode( - self, factory: AudioServiceTestDataFactory - ) -> None: - """Test that ASR raises error when speech-to-text is disabled in CHAT mode.""" - # Arrange - app_model_config = factory.create_app_model_config_mock(speech_to_text_dict={"enabled": False}) - app = factory.create_app_mock( - mode=AppMode.CHAT, - app_model_config=app_model_config, - ) - file = factory.create_file_storage() - - # Act & Assert - with pytest.raises(SpeechToTextDisabledServiceError): - AudioService.transcript_asr(app_model=app, file=file, session=self.session) - - def test_transcript_asr_raises_error_when_feature_disabled_workflow_mode( - self, factory: AudioServiceTestDataFactory - ) -> None: - """Test that ASR raises error when speech-to-text is disabled in WORKFLOW mode.""" - # Arrange - workflow = factory.create_workflow_mock(features_dict={"speech_to_text": {"enabled": False}}) - app = factory.create_app_mock( - mode=AppMode.WORKFLOW, - workflow=workflow, - ) - file = factory.create_file_storage() - - # Act & Assert - with pytest.raises(SpeechToTextDisabledServiceError): - AudioService.transcript_asr(app_model=app, file=file, session=self.session) - - def test_transcript_asr_raises_error_when_workflow_missing(self, factory: AudioServiceTestDataFactory) -> None: - """Test that ASR raises error when workflow is missing in WORKFLOW mode.""" - # Arrange - app = factory.create_app_mock( - mode=AppMode.WORKFLOW, - workflow=None, - ) - file = factory.create_file_storage() - - # Act & Assert - with pytest.raises(SpeechToTextDisabledServiceError): - AudioService.transcript_asr(app_model=app, file=file, session=self.session) - - def test_transcript_asr_raises_error_when_no_file_uploaded(self, factory: AudioServiceTestDataFactory) -> None: - """Test that ASR raises error when no file is uploaded.""" - # Arrange - app_model_config = factory.create_app_model_config_mock(speech_to_text_dict={"enabled": True}) - app = factory.create_app_mock( - mode=AppMode.CHAT, - app_model_config=app_model_config, - ) - - # Act & Assert - with pytest.raises(NoAudioUploadedServiceError): - AudioService.transcript_asr(app_model=app, file=None, session=self.session) - - def test_transcript_asr_raises_error_for_unsupported_audio_type(self, factory: AudioServiceTestDataFactory) -> None: - """Test that ASR raises error for unsupported audio file types.""" - # Arrange - app_model_config = factory.create_app_model_config_mock(speech_to_text_dict={"enabled": True}) - app = factory.create_app_mock( - mode=AppMode.CHAT, - app_model_config=app_model_config, - ) - file = factory.create_file_storage(mimetype="video/mp4") - - # Act & Assert - with pytest.raises(UnsupportedAudioTypeServiceError): - AudioService.transcript_asr(app_model=app, file=file, session=self.session) - - def test_transcript_asr_raises_error_for_large_file(self, factory: AudioServiceTestDataFactory) -> None: - """Test that ASR raises error when file exceeds size limit (30MB).""" - # Arrange - app_model_config = factory.create_app_model_config_mock(speech_to_text_dict={"enabled": True}) - app = factory.create_app_mock( - mode=AppMode.CHAT, - app_model_config=app_model_config, - ) - # Create file larger than 30MB - large_content = b"x" * (31 * 1024 * 1024) - file = factory.create_file_storage(content=large_content) - - # Act & Assert - with pytest.raises(AudioTooLargeServiceError, match="Audio size larger than 30 mb"): - AudioService.transcript_asr(app_model=app, file=file, session=self.session) - - @pytest.mark.parametrize("legacy_null", [False, True]) - def test_transcript_asr_raises_error_when_no_model_instance( - self, - audio_provider: AudioRuntimeObservations, - factory: AudioServiceTestDataFactory, - monkeypatch: pytest.MonkeyPatch, - legacy_null: bool, - ) -> None: - """Test that ASR raises error when no model instance is available.""" - # Arrange - app_model_config = factory.create_app_model_config_mock(speech_to_text_dict={"enabled": True}) - app = factory.create_app_mock( - mode=AppMode.CHAT, - app_model_config=app_model_config, - ) - file = factory.create_file_storage() - - audio_provider.default_available = False - - # Act & Assert - # The real manager raises for absent defaults. Explicitly retain the legacy null-return guard too. - if legacy_null: - monkeypatch.setattr(ModelManager, "get_default_model_instance", lambda *_args, **_kwargs: None) - with pytest.raises(ProviderNotSupportSpeechToTextServiceError if legacy_null else ProviderTokenNotInitError): - AudioService.transcript_asr(app_model=app, file=file, session=self.session) - - -class TestAudioServiceTTS: - """Test text-to-speech (TTS) operations.""" - - def test_legacy_tts_preserves_the_callers_uncommitted_transaction( - self, - audio_provider: AudioRuntimeObservations, - factory: AudioServiceTestDataFactory, - sqlite_session: Session, - ) -> None: - app = factory.create_app_mock() - sqlite_session.add(_message()) - sqlite_session.flush() - audio_provider.tts_result = b"pending message audio" - - result = AudioService.transcript_tts( - app_model=app, - session=sqlite_session, - message_ref=MessageRef( - app=AppRef(tenant_id=TENANT_ID, app_id=APP_ID), message_id=MESSAGE_ID, account_id=ACCOUNT_ID - ), - voice="pending-voice", - ) - - assert result is not None - assert result.get_data() == b"pending message audio" - assert sqlite_session.in_transaction() - assert audio_provider.tts_calls == [{"content_text": "Message answer", "voice": "pending-voice"}] - sqlite_session.rollback() - assert sqlite_session.get(Message, MESSAGE_ID) is None - - def test_transcript_tts_with_text_success( - self, - audio_provider: AudioRuntimeObservations, - factory: AudioServiceTestDataFactory, - sqlite_session: Session, - ) -> None: - """Test successful TTS with text input.""" - # Arrange - app_model_config = factory.create_app_model_config_mock( - text_to_speech_dict={"enabled": True, "voice": "en-US-Neural"} - ) - app = factory.create_app_mock( - mode=AppMode.CHAT, - app_model_config=app_model_config, - ) - - audio_provider.tts_result = b"audio data" - - # Act - result = AudioService.transcript_tts( - app_model=app, - session=sqlite_session, - text="Hello world", - voice="en-US-Neural", - end_user="user-123", - ) - - # Assert - assert result is not None - assert result.content_type == "audio/mpeg" - assert result.get_data() == b"audio data" - assert audio_provider.manager_requests == [ - { - "tenant_id": app.tenant_id, - "user_id": "user-123", - "request_metadata": { - "app_type": CreditUsageAppType.CHATBOT, - "created_by": CreditUsageCreatedBy.AUDIO, - }, - }, - ] - assert audio_provider.tts_calls == [{"content_text": "Hello world", "voice": "en-US-Neural"}] - assert audio_provider.runtime_calls[0]["request_metadata"] == { - "app_type": CreditUsageAppType.CHATBOT, - "created_by": CreditUsageCreatedBy.AUDIO, - } - - def test_transcript_tts_with_default_voice( - self, - audio_provider: AudioRuntimeObservations, - factory: AudioServiceTestDataFactory, - sqlite_session: Session, - ) -> None: - """Test TTS uses default voice when none specified.""" - # Arrange - app_model_config = factory.create_app_model_config_mock( - text_to_speech_dict={"enabled": True, "voice": "default-voice"} - ) - app = factory.create_app_mock( - mode=AppMode.CHAT, - app_model_config=app_model_config, - ) - - audio_provider.tts_result = b"audio data" - - # Act - result = AudioService.transcript_tts( - app_model=app, - session=sqlite_session, - text="Test", - ) - - # Assert - assert result is not None - assert result.content_type == "audio/mpeg" - assert result.get_data() == b"audio data" - # Verify default voice was used - assert audio_provider.tts_calls == [{"content_text": "Test", "voice": "default-voice"}] - - def test_transcript_tts_gets_first_available_voice_when_none_configured( - self, - audio_provider: AudioRuntimeObservations, - factory: AudioServiceTestDataFactory, - sqlite_session: Session, - ) -> None: - """Test TTS gets first available voice when none is configured.""" - # Arrange - app_model_config = factory.create_app_model_config_mock( - text_to_speech_dict={"enabled": True} # No voice specified - ) - app = factory.create_app_mock( - mode=AppMode.CHAT, - app_model_config=app_model_config, - ) - - audio_provider.voices = [{"name": "Automatic", "value": "auto-voice"}] - audio_provider.tts_result = b"audio data" - - # Act - result = AudioService.transcript_tts( - app_model=app, - session=sqlite_session, - text="Test", - ) - - # Assert - assert result is not None - assert result.content_type == "audio/mpeg" - assert result.get_data() == b"audio data" - assert audio_provider.tts_calls == [{"content_text": "Test", "voice": "auto-voice"}] - - def test_transcript_tts_workflow_mode_with_draft( - self, - audio_provider: AudioRuntimeObservations, - factory: AudioServiceTestDataFactory, - sqlite_session: Session, - ) -> None: - """Test TTS in WORKFLOW mode with draft workflow.""" - # Arrange - factory.create_workflow_mock(features_dict={"text_to_speech": {"enabled": True, "voice": "draft-voice"}}) - app = factory.create_app_mock( - mode=AppMode.WORKFLOW, - ) - - audio_provider.tts_result = b"draft audio" - # WorkflowService constructs its default repository from db.engine. - # Leave that database empty so the draft must come from sqlite_session. - flask_app = Flask(__name__) - flask_app.config["SQLALCHEMY_DATABASE_URI"] = "sqlite://" - db.init_app(flask_app) - with flask_app.app_context(): +class TestAudioServiceCases: + @staticmethod + @pytest.fixture + def audio_provider(audio_runtime: AudioRuntimeObservations, app: Flask) -> Iterator[AudioRuntimeObservations]: + # Real plugin TTS returns a stream, so the compatibility response needs a request context. + with app.test_request_context("/text-to-audio", method="POST"): + yield audio_runtime + + @staticmethod + @pytest.fixture + def factory(sqlite_session: Session) -> AudioServiceTestDataFactory: + """Provide the test data factory to all tests.""" + return AudioServiceTestDataFactory(sqlite_session) + + @staticmethod + @pytest.mark.parametrize("consumed_chunks", [0, 1, 2, None]) + def test_tts_response_closes_retained_provider_stream(consumed_chunks: int | None) -> None: + closed: list[bool] = [] + audio_chunks = [b"RIFF\x24\x00\x00\x00WAVE" + b"\x00" * 32, b"first-tail", b"second-tail"] + + def chunks() -> Generator[bytes]: try: - result = AudioService.transcript_tts( - app_model=app, - session=sqlite_session, - text="Draft test", - is_draft=True, - ) + for chunk in audio_chunks: + assert has_request_context() + assert request.path == "/text-to-audio" + yield chunk finally: - db.session.remove() - db.engine.dispose() + closed.append(True) - # Assert - assert result is not None - assert result.content_type == "audio/mpeg" - assert result.get_data() == b"draft audio" - assert audio_provider.tts_calls == [{"content_text": "Draft test", "voice": "draft-voice"}] + source = chunks() + app = Flask(__name__) + with app.test_request_context("/text-to-audio", method="POST"): + response = _create_tts_response(source, "audio/x-wav") - def test_transcript_tts_message_id_uses_provided_session( - self, - audio_provider: AudioRuntimeObservations, - factory: AudioServiceTestDataFactory, - sqlite_session: Session, - ) -> None: - """Test TTS message lookup uses the injected session.""" - # Arrange - app = factory.create_app_mock(app_id=APP_ID, tenant_id=TENANT_ID, mode=AppMode.CHAT) - message_ref = MessageRef( - app=AppRef(tenant_id=TENANT_ID, app_id=APP_ID), - message_id=MESSAGE_ID, - end_user_id=END_USER_ID, - account_id=ACCOUNT_ID, - ) - sqlite_session.add(_message()) - sqlite_session.commit() + assert response.status_code == 200 + assert dict(response.headers) == {"Content-Type": "audio/wav"} + assert response.is_streamed + assert closed == [] + try: + iterator = iter(response.response) + if consumed_chunks is None: + assert list(iterator) == audio_chunks + assert closed == [True] + else: + assert [next(iterator) for _ in range(consumed_chunks)] == audio_chunks[:consumed_chunks] + assert closed == [] + finally: + response.close() + response.close() - audio_provider.tts_result = b"message audio" + assert closed == [True] + with pytest.raises(StopIteration): + next(source) - # Act - for wrong_ref in ( - MessageRef( - app=AppRef(tenant_id=TENANT_ID, app_id=OTHER_ID), - message_id=MESSAGE_ID, - end_user_id=END_USER_ID, - account_id=ACCOUNT_ID, - ), - MessageRef( - app=AppRef(tenant_id=TENANT_ID, app_id=APP_ID), - message_id=MESSAGE_ID, - end_user_id=OTHER_ID, - account_id=ACCOUNT_ID, - ), - MessageRef( - app=AppRef(tenant_id=TENANT_ID, app_id=APP_ID), - message_id=MESSAGE_ID, - end_user_id=END_USER_ID, - account_id=OTHER_ID, - ), - ): - assert ( - AudioService.transcript_tts( - app_model=app, - session=sqlite_session, - message_ref=wrong_ref, - voice="message-voice", - ) - is None + class TestAudioServiceASR: + """Test speech-to-text (ASR) operations.""" + + @pytest.fixture(autouse=True) + def _bind_sqlite_session(self, sqlite_session: Session) -> None: + self.session = sqlite_session + + def test_transcript_asr_success_chat_mode( + self, audio_provider: AudioRuntimeObservations, factory: AudioServiceTestDataFactory + ) -> None: + """Test successful ASR transcription in CHAT mode.""" + # Arrange + app_model_config = factory.create_app_model_config_mock(speech_to_text_dict={"enabled": True}) + app = factory.create_app_mock( + mode=_audio_service_AppMode.CHAT, + app_model_config=app_model_config, + ) + file = factory.create_file_storage() + + audio_provider.speech_to_text_result = "Transcribed text" + + # Act + result = AudioService.transcript_asr(app_model=app, file=file, session=self.session, end_user="user-123") + + # Assert + assert result == {"text": "Transcribed text"} + assert audio_provider.speech_to_text_calls == [b"fake audio content"] + assert audio_provider.runtime_calls[0]["request_metadata"] == { + "app_type": CreditUsageAppType.CHATBOT, + "created_by": CreditUsageCreatedBy.AUDIO, + } + assert audio_provider.manager_requests == [ + { + "tenant_id": app.tenant_id, + "user_id": "user-123", + "request_metadata": { + "app_type": CreditUsageAppType.CHATBOT, + "created_by": CreditUsageCreatedBy.AUDIO, + }, + }, + ] + + def test_transcript_asr_accepts_x_m4a_mimetype( + self, audio_provider: AudioRuntimeObservations, factory: AudioServiceTestDataFactory + ) -> None: + """Test that the x-m4a MIME alias follows the normal m4a transcription flow.""" + # Arrange + app_model_config = factory.create_app_model_config_mock(speech_to_text_dict={"enabled": True}) + app = factory.create_app_mock(mode=_audio_service_AppMode.CHAT, app_model_config=app_model_config) + file = factory.create_file_storage(filename="audio.m4a", mimetype="audio/x-m4a") + audio_provider.speech_to_text_result = "M4A transcript" + + # Act + result = AudioService.transcript_asr(app_model=app, file=file, session=self.session) + + # Assert + assert result == {"text": "M4A transcript"} + assert audio_provider.speech_to_text_calls == [b"fake audio content"] + + def test_transcript_asr_success_advanced_chat_mode( + self, audio_provider: AudioRuntimeObservations, factory: AudioServiceTestDataFactory + ) -> None: + """Test successful ASR transcription in ADVANCED_CHAT mode.""" + # Arrange + workflow = factory.create_workflow_mock(features_dict={"speech_to_text": {"enabled": True}}) + app = factory.create_app_mock( + mode=_audio_service_AppMode.ADVANCED_CHAT, + workflow=workflow, + ) + file = factory.create_file_storage() + + audio_provider.speech_to_text_result = "Workflow transcribed text" + + # Act + result = AudioService.transcript_asr(app_model=app, file=file, session=self.session) + + # Assert + assert result == {"text": "Workflow transcribed text"} + + def test_transcript_asr_success_published_agent_mode( + self, + audio_provider: AudioRuntimeObservations, + factory: AudioServiceTestDataFactory, + ) -> None: + app = factory.create_app_mock(mode=_audio_service_AppMode.AGENT) + file = factory.create_file_storage() + agent_soul = AgentSoulConfig.model_validate({"app_features": {"speech_to_text": {"enabled": True}}}) + agent = Agent( + tenant_id=app.tenant_id, + app_id=app.id, + name="Audio agent", + scope=AgentScope.ROSTER, + source=AgentSource.AGENT_APP, + ) + self.session.add(agent) + self.session.flush() + snapshot = AgentConfigSnapshot( + tenant_id=app.tenant_id, agent_id=agent.id, version=1, config_snapshot=agent_soul + ) + self.session.add(snapshot) + self.session.flush() + agent.active_config_snapshot_id = snapshot.id + self.session.commit() + audio_provider.speech_to_text_result = "Published Agent transcript" + + result = AudioService.transcript_asr(app_model=app, file=file, session=self.session, end_user="end-user-1") + + assert result == {"text": "Published Agent transcript"} + assert audio_provider.manager_requests == [ + { + "tenant_id": app.tenant_id, + "user_id": "end-user-1", + "request_metadata": { + "app_type": CreditUsageAppType.AGENT_V2, + "created_by": CreditUsageCreatedBy.AUDIO, + }, + }, + ] + + def test_transcript_asr_legacy_agent_falls_back_to_app_model_config( + self, + audio_provider: AudioRuntimeObservations, + factory: AudioServiceTestDataFactory, + ) -> None: + app_model_config = factory.create_app_model_config_mock(speech_to_text_dict={"enabled": True}) + app = factory.create_app_mock(mode=_audio_service_AppMode.AGENT, app_model_config=app_model_config) + file = factory.create_file_storage() + audio_provider.speech_to_text_result = "Legacy Agent transcript" + + result = AudioService.transcript_asr(app_model=app, file=file, session=self.session) + + assert result == {"text": "Legacy Agent transcript"} + + def test_transcript_agent_asr_uses_agent_soul_feature( + self, audio_provider: AudioRuntimeObservations, factory: AudioServiceTestDataFactory + ) -> None: + app = factory.create_app_mock(mode=_audio_service_AppMode.AGENT) + file = factory.create_file_storage() + agent_soul = AgentSoulConfig.model_validate({"app_features": {"speech_to_text": {"enabled": True}}}) + audio_provider.speech_to_text_result = "Agent transcript" + + result = AudioService.transcript_agent_asr( + app_model=app, + agent_soul=agent_soul, + file=file, + session=self.session, + end_user="account-1", ) - result = AudioService.transcript_tts( - app_model=app, - session=sqlite_session, - message_ref=message_ref, - voice="message-voice", - ) + assert result == {"text": "Agent transcript"} + assert audio_provider.manager_requests == [ + { + "tenant_id": app.tenant_id, + "user_id": "account-1", + "request_metadata": { + "app_type": CreditUsageAppType.AGENT_V2, + "created_by": CreditUsageCreatedBy.AUDIO, + }, + }, + ] - # Assert - assert result is not None - assert result.content_type == "audio/mpeg" - assert result.get_data() == b"message audio" - assert audio_provider.tts_calls == [{"content_text": "Message answer", "voice": "message-voice"}] - - def test_transcript_tts_uses_detected_wav_mime_type_for_streams( - self, - audio_provider: AudioRuntimeObservations, - factory: AudioServiceTestDataFactory, - sqlite_session: Session, - app: Flask, - ) -> None: - app_model_config = factory.create_app_model_config_mock( - text_to_speech_dict={"enabled": True, "voice": "en-US-Neural"} + @pytest.mark.parametrize( + "agent_soul", + [ + AgentSoulConfig(), + AgentSoulConfig.model_validate({"app_features": {"speech_to_text": {"enabled": False}}}), + ], ) - app_model = factory.create_app_mock(mode=AppMode.CHAT, app_model_config=app_model_config) - audio_provider.tts_result = ( - chunk for chunk in [b"RIFF\x24\x00\x00\x00WAVEfmt ", b"\x10\x00\x00\x00audio-data"] - ) - # Legacy plugins without format metadata must use the detected container. - audio_provider.model_properties = {} + def test_transcript_agent_asr_rejects_disabled_feature( + self, factory: AudioServiceTestDataFactory, agent_soul: AgentSoulConfig + ) -> None: + app = factory.create_app_mock(mode=_audio_service_AppMode.AGENT) + file = factory.create_file_storage() + + with pytest.raises(SpeechToTextDisabledServiceError): + AudioService.transcript_agent_asr(app_model=app, agent_soul=agent_soul, file=file, session=self.session) + + def test_transcript_agent_asr_preserves_legacy_feature_fallback( + self, audio_provider: AudioRuntimeObservations, factory: AudioServiceTestDataFactory + ) -> None: + app_model_config = factory.create_app_model_config_mock(speech_to_text_dict={"enabled": True}) + app = factory.create_app_mock(mode=_audio_service_AppMode.AGENT, app_model_config=app_model_config) + file = factory.create_file_storage() + audio_provider.speech_to_text_result = "Legacy feature transcript" + + result = AudioService.transcript_agent_asr( + app_model=app, + agent_soul=AgentSoulConfig(), + file=file, + session=self.session, + ) + + assert result == {"text": "Legacy feature transcript"} + + def test_transcript_agent_asr_soul_disabled_overrides_legacy_feature( + self, factory: AudioServiceTestDataFactory + ) -> None: + app_model_config = factory.create_app_model_config_mock(speech_to_text_dict={"enabled": True}) + app = factory.create_app_mock(mode=_audio_service_AppMode.AGENT, app_model_config=app_model_config) + file = factory.create_file_storage() + agent_soul = AgentSoulConfig.model_validate({"app_features": {"speech_to_text": {"enabled": False}}}) + + with pytest.raises(SpeechToTextDisabledServiceError): + AudioService.transcript_agent_asr(app_model=app, agent_soul=agent_soul, file=file, session=self.session) + + def test_transcript_asr_raises_error_when_feature_disabled_chat_mode( + self, factory: AudioServiceTestDataFactory + ) -> None: + """Test that ASR raises error when speech-to-text is disabled in CHAT mode.""" + # Arrange + app_model_config = factory.create_app_model_config_mock(speech_to_text_dict={"enabled": False}) + app = factory.create_app_mock( + mode=_audio_service_AppMode.CHAT, + app_model_config=app_model_config, + ) + file = factory.create_file_storage() + + # Act & Assert + with pytest.raises(SpeechToTextDisabledServiceError): + AudioService.transcript_asr(app_model=app, file=file, session=self.session) + + def test_transcript_asr_raises_error_when_feature_disabled_workflow_mode( + self, factory: AudioServiceTestDataFactory + ) -> None: + """Test that ASR raises error when speech-to-text is disabled in WORKFLOW mode.""" + # Arrange + workflow = factory.create_workflow_mock(features_dict={"speech_to_text": {"enabled": False}}) + app = factory.create_app_mock( + mode=_audio_service_AppMode.WORKFLOW, + workflow=workflow, + ) + file = factory.create_file_storage() + + # Act & Assert + with pytest.raises(SpeechToTextDisabledServiceError): + AudioService.transcript_asr(app_model=app, file=file, session=self.session) + + def test_transcript_asr_raises_error_when_workflow_missing(self, factory: AudioServiceTestDataFactory) -> None: + """Test that ASR raises error when workflow is missing in WORKFLOW mode.""" + # Arrange + app = factory.create_app_mock( + mode=_audio_service_AppMode.WORKFLOW, + workflow=None, + ) + file = factory.create_file_storage() + + # Act & Assert + with pytest.raises(SpeechToTextDisabledServiceError): + AudioService.transcript_asr(app_model=app, file=file, session=self.session) + + def test_transcript_asr_raises_error_when_no_file_uploaded(self, factory: AudioServiceTestDataFactory) -> None: + """Test that ASR raises error when no file is uploaded.""" + # Arrange + app_model_config = factory.create_app_model_config_mock(speech_to_text_dict={"enabled": True}) + app = factory.create_app_mock( + mode=_audio_service_AppMode.CHAT, + app_model_config=app_model_config, + ) + + # Act & Assert + with pytest.raises(NoAudioUploadedServiceError): + AudioService.transcript_asr(app_model=app, file=None, session=self.session) + + def test_transcript_asr_raises_error_for_unsupported_audio_type( + self, factory: AudioServiceTestDataFactory + ) -> None: + """Test that ASR raises error for unsupported audio file types.""" + # Arrange + app_model_config = factory.create_app_model_config_mock(speech_to_text_dict={"enabled": True}) + app = factory.create_app_mock( + mode=_audio_service_AppMode.CHAT, + app_model_config=app_model_config, + ) + file = factory.create_file_storage(mimetype="video/mp4") + + # Act & Assert + with pytest.raises(UnsupportedAudioTypeServiceError): + AudioService.transcript_asr(app_model=app, file=file, session=self.session) + + def test_transcript_asr_raises_error_for_large_file(self, factory: AudioServiceTestDataFactory) -> None: + """Test that ASR raises error when file exceeds size limit (30MB).""" + # Arrange + app_model_config = factory.create_app_model_config_mock(speech_to_text_dict={"enabled": True}) + app = factory.create_app_mock( + mode=_audio_service_AppMode.CHAT, + app_model_config=app_model_config, + ) + # Create file larger than 30MB + large_content = b"x" * (31 * 1024 * 1024) + file = factory.create_file_storage(content=large_content) + + # Act & Assert + with pytest.raises(AudioTooLargeServiceError, match="Audio size larger than 30 mb"): + AudioService.transcript_asr(app_model=app, file=file, session=self.session) + + @pytest.mark.parametrize("legacy_null", [False, True]) + def test_transcript_asr_raises_error_when_no_model_instance( + self, + audio_provider: AudioRuntimeObservations, + factory: AudioServiceTestDataFactory, + monkeypatch: pytest.MonkeyPatch, + legacy_null: bool, + ) -> None: + """Test that ASR raises error when no model instance is available.""" + # Arrange + app_model_config = factory.create_app_model_config_mock(speech_to_text_dict={"enabled": True}) + app = factory.create_app_mock( + mode=_audio_service_AppMode.CHAT, + app_model_config=app_model_config, + ) + file = factory.create_file_storage() + + audio_provider.default_available = False + + # Act & Assert + # The real manager raises for absent defaults. Explicitly retain the legacy null-return guard too. + if legacy_null: + monkeypatch.setattr(ModelManager, "get_default_model_instance", lambda *_args, **_kwargs: None) + with pytest.raises( + ProviderNotSupportSpeechToTextServiceError if legacy_null else ProviderTokenNotInitError + ): + AudioService.transcript_asr(app_model=app, file=file, session=self.session) + + class TestAudioServiceTTS: + """Test text-to-speech (TTS) operations.""" + + def test_legacy_tts_preserves_the_callers_uncommitted_transaction( + self, + audio_provider: AudioRuntimeObservations, + factory: AudioServiceTestDataFactory, + sqlite_session: Session, + ) -> None: + app = factory.create_app_mock() + sqlite_session.add(_message()) + sqlite_session.flush() + audio_provider.tts_result = b"pending message audio" - with app.test_request_context("/text-to-audio", method="POST"): result = AudioService.transcript_tts( - app_model=app_model, + app_model=app, + session=sqlite_session, + message_ref=MessageRef( + app=AppRef(tenant_id=TENANT_ID, app_id=APP_ID), message_id=MESSAGE_ID, account_id=ACCOUNT_ID + ), + voice="pending-voice", + ) + + assert result is not None + assert result.get_data() == b"pending message audio" + assert sqlite_session.in_transaction() + assert audio_provider.tts_calls == [{"content_text": "Message answer", "voice": "pending-voice"}] + sqlite_session.rollback() + assert sqlite_session.get(Message, MESSAGE_ID) is None + + def test_transcript_tts_with_text_success( + self, + audio_provider: AudioRuntimeObservations, + factory: AudioServiceTestDataFactory, + sqlite_session: Session, + ) -> None: + """Test successful TTS with text input.""" + # Arrange + app_model_config = factory.create_app_model_config_mock( + text_to_speech_dict={"enabled": True, "voice": "en-US-Neural"} + ) + app = factory.create_app_mock( + mode=_audio_service_AppMode.CHAT, + app_model_config=app_model_config, + ) + + audio_provider.tts_result = b"audio data" + + # Act + result = AudioService.transcript_tts( + app_model=app, session=sqlite_session, text="Hello world", voice="en-US-Neural", + end_user="user-123", ) - assert result is not None - assert result.content_type == "audio/wav" - assert result.get_data() == b"RIFF\x24\x00\x00\x00WAVEfmt \x10\x00\x00\x00audio-data" + # Assert + assert result is not None + assert result.content_type == "audio/mpeg" + assert result.get_data() == b"audio data" + assert audio_provider.manager_requests == [ + { + "tenant_id": app.tenant_id, + "user_id": "user-123", + "request_metadata": { + "app_type": CreditUsageAppType.CHATBOT, + "created_by": CreditUsageCreatedBy.AUDIO, + }, + }, + ] + assert audio_provider.tts_calls == [{"content_text": "Hello world", "voice": "en-US-Neural"}] + assert audio_provider.runtime_calls[0]["request_metadata"] == { + "app_type": CreditUsageAppType.CHATBOT, + "created_by": CreditUsageCreatedBy.AUDIO, + } - def test_transcript_tts_returns_provider_output_error_for_mime_magic_mismatch( - self, - audio_provider: AudioRuntimeObservations, - factory: AudioServiceTestDataFactory, - sqlite_session: Session, - app: Flask, - ) -> None: - app_model_config = factory.create_app_model_config_mock( - text_to_speech_dict={"enabled": True, "voice": "en-US-Neural"} - ) - app_model = factory.create_app_mock(mode=AppMode.CHAT, app_model_config=app_model_config) - audio_provider.tts_result = ( - chunk for chunk in [TTSAudioChunk(b"RIFF\x24\x00\x00\x00WAVEfmt \x10\x00\x00\x00audio-data", "audio/mpeg")] - ) + def test_transcript_tts_with_default_voice( + self, + audio_provider: AudioRuntimeObservations, + factory: AudioServiceTestDataFactory, + sqlite_session: Session, + ) -> None: + """Test TTS uses default voice when none specified.""" + # Arrange + app_model_config = factory.create_app_model_config_mock( + text_to_speech_dict={"enabled": True, "voice": "default-voice"} + ) + app = factory.create_app_mock( + mode=_audio_service_AppMode.CHAT, + app_model_config=app_model_config, + ) - with app.test_request_context("/text-to-audio", method="POST"): - with pytest.raises(InvokeBadRequestError, match="declared audio/mpeg, detected audio/wav"): - AudioService.transcript_tts( + audio_provider.tts_result = b"audio data" + + # Act + result = AudioService.transcript_tts( + app_model=app, + session=sqlite_session, + text="Test", + ) + + # Assert + assert result is not None + assert result.content_type == "audio/mpeg" + assert result.get_data() == b"audio data" + # Verify default voice was used + assert audio_provider.tts_calls == [{"content_text": "Test", "voice": "default-voice"}] + + def test_transcript_tts_gets_first_available_voice_when_none_configured( + self, + audio_provider: AudioRuntimeObservations, + factory: AudioServiceTestDataFactory, + sqlite_session: Session, + ) -> None: + """Test TTS gets first available voice when none is configured.""" + # Arrange + app_model_config = factory.create_app_model_config_mock( + text_to_speech_dict={"enabled": True} # No voice specified + ) + app = factory.create_app_mock( + mode=_audio_service_AppMode.CHAT, + app_model_config=app_model_config, + ) + + audio_provider.voices = [{"name": "Automatic", "value": "auto-voice"}] + audio_provider.tts_result = b"audio data" + + # Act + result = AudioService.transcript_tts( + app_model=app, + session=sqlite_session, + text="Test", + ) + + # Assert + assert result is not None + assert result.content_type == "audio/mpeg" + assert result.get_data() == b"audio data" + assert audio_provider.tts_calls == [{"content_text": "Test", "voice": "auto-voice"}] + + def test_transcript_tts_workflow_mode_with_draft( + self, + audio_provider: AudioRuntimeObservations, + factory: AudioServiceTestDataFactory, + sqlite_session: Session, + ) -> None: + """Test TTS in WORKFLOW mode with draft workflow.""" + # Arrange + factory.create_workflow_mock(features_dict={"text_to_speech": {"enabled": True, "voice": "draft-voice"}}) + app = factory.create_app_mock( + mode=_audio_service_AppMode.WORKFLOW, + ) + + audio_provider.tts_result = b"draft audio" + # WorkflowService constructs its default repository from db.engine. + # Leave that database empty so the draft must come from sqlite_session. + flask_app = Flask(__name__) + flask_app.config["SQLALCHEMY_DATABASE_URI"] = "sqlite://" + db.init_app(flask_app) + with flask_app.app_context(): + try: + result = AudioService.transcript_tts( + app_model=app, + session=sqlite_session, + text="Draft test", + is_draft=True, + ) + finally: + db.session.remove() + db.engine.dispose() + + # Assert + assert result is not None + assert result.content_type == "audio/mpeg" + assert result.get_data() == b"draft audio" + assert audio_provider.tts_calls == [{"content_text": "Draft test", "voice": "draft-voice"}] + + def test_transcript_tts_message_id_uses_provided_session( + self, + audio_provider: AudioRuntimeObservations, + factory: AudioServiceTestDataFactory, + sqlite_session: Session, + ) -> None: + """Test TTS message lookup uses the injected session.""" + # Arrange + app = factory.create_app_mock(app_id=APP_ID, tenant_id=TENANT_ID, mode=_audio_service_AppMode.CHAT) + message_ref = MessageRef( + app=AppRef(tenant_id=TENANT_ID, app_id=APP_ID), + message_id=MESSAGE_ID, + end_user_id=END_USER_ID, + account_id=ACCOUNT_ID, + ) + sqlite_session.add(_message()) + sqlite_session.commit() + + audio_provider.tts_result = b"message audio" + + # Act + for wrong_ref in ( + MessageRef( + app=AppRef(tenant_id=TENANT_ID, app_id=OTHER_ID), + message_id=MESSAGE_ID, + end_user_id=END_USER_ID, + account_id=ACCOUNT_ID, + ), + MessageRef( + app=AppRef(tenant_id=TENANT_ID, app_id=APP_ID), + message_id=MESSAGE_ID, + end_user_id=OTHER_ID, + account_id=ACCOUNT_ID, + ), + MessageRef( + app=AppRef(tenant_id=TENANT_ID, app_id=APP_ID), + message_id=MESSAGE_ID, + end_user_id=END_USER_ID, + account_id=OTHER_ID, + ), + ): + assert ( + AudioService.transcript_tts( + app_model=app, + session=sqlite_session, + message_ref=wrong_ref, + voice="message-voice", + ) + is None + ) + + result = AudioService.transcript_tts( + app_model=app, + session=sqlite_session, + message_ref=message_ref, + voice="message-voice", + ) + + # Assert + assert result is not None + assert result.content_type == "audio/mpeg" + assert result.get_data() == b"message audio" + assert audio_provider.tts_calls == [{"content_text": "Message answer", "voice": "message-voice"}] + + def test_transcript_tts_uses_detected_wav_mime_type_for_streams( + self, + audio_provider: AudioRuntimeObservations, + factory: AudioServiceTestDataFactory, + sqlite_session: Session, + app: Flask, + ) -> None: + app_model_config = factory.create_app_model_config_mock( + text_to_speech_dict={"enabled": True, "voice": "en-US-Neural"} + ) + app_model = factory.create_app_mock(mode=_audio_service_AppMode.CHAT, app_model_config=app_model_config) + audio_provider.tts_result = ( + chunk for chunk in [b"RIFF\x24\x00\x00\x00WAVEfmt ", b"\x10\x00\x00\x00audio-data"] + ) + # Legacy plugins without format metadata must use the detected container. + audio_provider.model_properties = {} + + with app.test_request_context("/text-to-audio", method="POST"): + result = AudioService.transcript_tts( app_model=app_model, session=sqlite_session, text="Hello world", voice="en-US-Neural", ) - def test_transcript_tts_raises_error_when_text_missing( - self, - factory: AudioServiceTestDataFactory, - sqlite_session: Session, - ) -> None: - """Test that TTS raises error when text is missing.""" - # Arrange - app = factory.create_app_mock() + assert result is not None + assert result.content_type == "audio/wav" + assert result.get_data() == b"RIFF\x24\x00\x00\x00WAVEfmt \x10\x00\x00\x00audio-data" - # Act & Assert - with pytest.raises(ValueError, match="Text is required"): - AudioService.transcript_tts(app_model=app, session=sqlite_session, text=None) + def test_transcript_tts_returns_provider_output_error_for_mime_magic_mismatch( + self, + audio_provider: AudioRuntimeObservations, + factory: AudioServiceTestDataFactory, + sqlite_session: Session, + app: Flask, + ) -> None: + app_model_config = factory.create_app_model_config_mock( + text_to_speech_dict={"enabled": True, "voice": "en-US-Neural"} + ) + app_model = factory.create_app_mock(mode=_audio_service_AppMode.CHAT, app_model_config=app_model_config) + audio_provider.tts_result = ( + chunk + for chunk in [TTSAudioChunk(b"RIFF\x24\x00\x00\x00WAVEfmt \x10\x00\x00\x00audio-data", "audio/mpeg")] + ) - def test_transcript_tts_raises_error_when_no_voices_available( - self, - audio_provider: AudioRuntimeObservations, - factory: AudioServiceTestDataFactory, - sqlite_session: Session, - ) -> None: - """Test that TTS raises error when no voices are available.""" - # Arrange - app_model_config = factory.create_app_model_config_mock( - text_to_speech_dict={"enabled": True} # No voice specified - ) - app = factory.create_app_mock( - mode=AppMode.CHAT, - app_model_config=app_model_config, - ) + with app.test_request_context("/text-to-audio", method="POST"): + with pytest.raises(InvokeBadRequestError, match="declared audio/mpeg, detected audio/wav"): + AudioService.transcript_tts( + app_model=app_model, + session=sqlite_session, + text="Hello world", + voice="en-US-Neural", + ) - audio_provider.voices = [] + def test_transcript_tts_raises_error_when_text_missing( + self, + factory: AudioServiceTestDataFactory, + sqlite_session: Session, + ) -> None: + """Test that TTS raises error when text is missing.""" + # Arrange + app = factory.create_app_mock() - # Act & Assert - with pytest.raises(ValueError, match="Sorry, no voice available"): - AudioService.transcript_tts(app_model=app, session=sqlite_session, text="Test") + # Act & Assert + with pytest.raises(ValueError, match="Text is required"): + AudioService.transcript_tts(app_model=app, session=sqlite_session, text=None) + def test_transcript_tts_raises_error_when_no_voices_available( + self, + audio_provider: AudioRuntimeObservations, + factory: AudioServiceTestDataFactory, + sqlite_session: Session, + ) -> None: + """Test that TTS raises error when no voices are available.""" + # Arrange + app_model_config = factory.create_app_model_config_mock( + text_to_speech_dict={"enabled": True} # No voice specified + ) + app = factory.create_app_mock( + mode=_audio_service_AppMode.CHAT, + app_model_config=app_model_config, + ) -class TestAudioServiceTTSVoices: - """Test TTS voice listing operations.""" + audio_provider.voices = [] - def test_transcript_tts_voices_success(self, audio_provider: AudioRuntimeObservations) -> None: - """Test successful retrieval of TTS voices.""" - # Arrange - tenant_id = "tenant-123" - language = "en-US" + # Act & Assert + with pytest.raises(ValueError, match="Sorry, no voice available"): + AudioService.transcript_tts(app_model=app, session=sqlite_session, text="Test") - expected_voices: list[TTSModelVoice] = [ - {"name": "Voice 1", "value": "voice-1"}, - {"name": "Voice 2", "value": "voice-2"}, - ] + class TestAudioServiceTTSVoices: + """Test TTS voice listing operations.""" - audio_provider.voices = expected_voices + def test_transcript_tts_voices_success(self, audio_provider: AudioRuntimeObservations) -> None: + """Test successful retrieval of TTS voices.""" + # Arrange + tenant_id = "tenant-123" + language = "en-US" - # Act - result = AudioService.transcript_tts_voices(tenant_id=tenant_id, language=language) + expected_voices: list[TTSModelVoice] = [ + {"name": "Voice 1", "value": "voice-1"}, + {"name": "Voice 2", "value": "voice-2"}, + ] - # Assert - assert result == expected_voices - assert audio_provider.voice_calls == [(language,)] + audio_provider.voices = expected_voices - @pytest.mark.parametrize("legacy_null", [False, True]) - def test_transcript_tts_voices_raises_error_when_no_model_instance( - self, audio_provider: AudioRuntimeObservations, monkeypatch: pytest.MonkeyPatch, legacy_null: bool - ) -> None: - """Test that TTS voices raises error when no model instance is available.""" - # Arrange - tenant_id = "tenant-123" - language = "en-US" + # Act + result = AudioService.transcript_tts_voices(tenant_id=tenant_id, language=language) - audio_provider.default_available = False + # Assert + assert result == expected_voices + assert audio_provider.voice_calls == [(language,)] - # Act & Assert - if legacy_null: - monkeypatch.setattr(ModelManager, "get_default_model_instance", lambda *_args, **_kwargs: None) - with pytest.raises(ProviderNotSupportTextToSpeechServiceError if legacy_null else ProviderTokenNotInitError): - AudioService.transcript_tts_voices(tenant_id=tenant_id, language=language) + @pytest.mark.parametrize("legacy_null", [False, True]) + def test_transcript_tts_voices_raises_error_when_no_model_instance( + self, audio_provider: AudioRuntimeObservations, monkeypatch: pytest.MonkeyPatch, legacy_null: bool + ) -> None: + """Test that TTS voices raises error when no model instance is available.""" + # Arrange + tenant_id = "tenant-123" + language = "en-US" - def test_transcript_tts_voices_propagates_exceptions(self, audio_provider: AudioRuntimeObservations) -> None: - """Test that TTS voices propagates exceptions from model instance.""" - # Arrange - tenant_id = "tenant-123" - language = "en-US" + audio_provider.default_available = False - audio_provider.voices_error = RuntimeError("Model error") + # Act & Assert + if legacy_null: + monkeypatch.setattr(ModelManager, "get_default_model_instance", lambda *_args, **_kwargs: None) + with pytest.raises( + ProviderNotSupportTextToSpeechServiceError if legacy_null else ProviderTokenNotInitError + ): + AudioService.transcript_tts_voices(tenant_id=tenant_id, language=language) - # Act & Assert - with pytest.raises(RuntimeError, match="Model error"): - AudioService.transcript_tts_voices(tenant_id=tenant_id, language=language) + def test_transcript_tts_voices_propagates_exceptions(self, audio_provider: AudioRuntimeObservations) -> None: + """Test that TTS voices propagates exceptions from model instance.""" + # Arrange + tenant_id = "tenant-123" + language = "en-US" + + audio_provider.voices_error = RuntimeError("Model error") + + # Act & Assert + with pytest.raises(RuntimeError, match="Model error"): + AudioService.transcript_tts_voices(tenant_id=tenant_id, language=language) diff --git a/api/tests/unit_tests/services/test_billing_portal_service.py b/api/tests/unit_tests/services/test_billing_portal_service.py index 20395950f84..b8a7c9e62c5 100644 --- a/api/tests/unit_tests/services/test_billing_portal_service.py +++ b/api/tests/unit_tests/services/test_billing_portal_service.py @@ -9,7 +9,7 @@ from models.account import Account from repositories.account.repository import SQLAlchemyAccountRepository from services.account_errors import AccountNotFoundError from services.billing_portal_service import BillingPortalService -from tests.unit_tests.model_factories import make_account +from tests.helpers.unit.model_factories import make_account def _context() -> RequestContext: diff --git a/api/tests/unit_tests/services/test_clear_free_plan_expired_workflow_run_logs.py b/api/tests/unit_tests/services/test_clear_free_plan_expired_workflow_run_logs.py deleted file mode 100644 index ffdff6df787..00000000000 --- a/api/tests/unit_tests/services/test_clear_free_plan_expired_workflow_run_logs.py +++ /dev/null @@ -1,485 +0,0 @@ -import datetime -from collections.abc import Callable -from typing import Any - -import pytest - -from enums import DeploymentEdition -from repositories.api_workflow_run_repository import WorkflowRunCleanupRef -from services.billing_service import SubscriptionPlan -from services.retention.workflow_run import clear_free_plan_expired_workflow_run_logs as cleanup_module -from services.retention.workflow_run.clear_free_plan_expired_workflow_run_logs import WorkflowRunCleanup - - -@pytest.fixture(autouse=True) -def _cleanup_config(config_overrides: Callable[..., None]) -> None: - config_overrides( - DEPLOYMENT_EDITION=DeploymentEdition.COMMUNITY, - SANDBOX_EXPIRED_RECORDS_CLEAN_GRACEFUL_PERIOD=0, - ) - - -def make_ref(run_id: str, tenant_id: str, created_at: datetime.datetime) -> WorkflowRunCleanupRef: - return WorkflowRunCleanupRef(id=run_id, tenant_id=tenant_id, created_at=created_at) - - -class FakeRepo: - def __init__( - self, - batches: list[list[WorkflowRunCleanupRef]], - delete_result: dict[str, int] | None = None, - count_result: dict[str, int] | None = None, - ) -> None: - self.batches = batches - self.candidate_call_idx = 0 - self.last_candidate_batch: list[WorkflowRunCleanupRef] = [] - self.cleanup_ref_calls: list[dict[str, object]] = [] - self.deleted: list[list[str]] = [] - self.counted: list[list[str]] = [] - self.delete_result = delete_result or { - "runs": 0, - "node_executions": 0, - "offloads": 0, - "app_logs": 0, - "trigger_logs": 0, - "pauses": 0, - "pause_reasons": 0, - } - self.count_result = count_result or { - "runs": 0, - "node_executions": 0, - "offloads": 0, - "app_logs": 0, - "trigger_logs": 0, - "pauses": 0, - "pause_reasons": 0, - } - - def get_cleanup_refs_batch_by_time_range( - self, - start_from: datetime.datetime | None, - end_before: datetime.datetime, - last_seen: tuple[datetime.datetime, str] | None, - batch_size: int, - run_types=None, - tenant_ids=None, - workflow_ids=None, - upper_bound: tuple[datetime.datetime, str] | None = None, - ) -> list[WorkflowRunCleanupRef]: - self.cleanup_ref_calls.append( - { - "start_from": start_from, - "end_before": end_before, - "last_seen": last_seen, - "batch_size": batch_size, - "run_types": run_types, - "tenant_ids": tenant_ids, - "workflow_ids": workflow_ids, - "upper_bound": upper_bound, - } - ) - if tenant_ids is not None or upper_bound is not None: - refs = self.last_candidate_batch - if tenant_ids is not None: - tenant_id_set = set(tenant_ids) - refs = [ref for ref in refs if ref.tenant_id in tenant_id_set] - if upper_bound is not None: - refs = [ref for ref in refs if (ref.created_at, ref.id) <= upper_bound] - return refs[:batch_size] - - if self.candidate_call_idx >= len(self.batches): - return [] - batch = self.batches[self.candidate_call_idx] - self.candidate_call_idx += 1 - self.last_candidate_batch = batch - return batch - - def delete_runs_with_related_by_ids( - self, run_ids: list[str], delete_node_executions=None, delete_trigger_logs=None - ) -> dict[str, int]: - self.deleted.append(list(run_ids)) - result = self.delete_result.copy() - result["runs"] = len(run_ids) - return result - - def count_runs_with_related_by_ids( - self, run_ids: list[str], count_node_executions=None, count_trigger_logs=None - ) -> dict[str, int]: - self.counted.append(list(run_ids)) - result = self.count_result.copy() - result["runs"] = len(run_ids) - return result - - -def plan_info(plan: str, expiration: int) -> SubscriptionPlan: - return SubscriptionPlan(plan=plan, expiration_date=expiration) - - -def create_cleanup( - monkeypatch: pytest.MonkeyPatch, - repo: FakeRepo, - *, - whitelist: set[str] | None = None, - **kwargs: Any, -) -> WorkflowRunCleanup: - monkeypatch.setattr( - cleanup_module.WorkflowRunCleanup, - "_get_cleanup_whitelist", - lambda self: whitelist or set(), - ) - return WorkflowRunCleanup(workflow_run_repo=repo, **kwargs) - - -def test_filter_free_tenants_outside_cloud_edition(monkeypatch: pytest.MonkeyPatch) -> None: - cleanup = create_cleanup(monkeypatch, repo=FakeRepo([]), days=30, batch_size=10) - - def fail_bulk(_: list[str]) -> dict[str, SubscriptionPlan]: - raise RuntimeError("should not call") - - monkeypatch.setattr(cleanup_module.BillingService, "get_plan_bulk_with_cache", staticmethod(fail_bulk)) - - tenants = {"t1", "t2"} - free = cleanup._filter_free_tenants(tenants) - - assert free == tenants - - -def test_filter_free_tenants_bulk_mixed(monkeypatch: pytest.MonkeyPatch, config_overrides: Callable[..., None]) -> None: - config_overrides(DEPLOYMENT_EDITION=DeploymentEdition.CLOUD) - cleanup = create_cleanup(monkeypatch, repo=FakeRepo([]), days=30, batch_size=10) - - monkeypatch.setattr( - cleanup_module.BillingService, - "get_plan_bulk_with_cache", - staticmethod( - lambda tenant_ids: { - tenant_id: (plan_info("team", -1) if tenant_id == "t_paid" else plan_info("sandbox", -1)) - for tenant_id in tenant_ids - } - ), - ) - - free = cleanup._filter_free_tenants({"t_free", "t_paid", "t_missing"}) - - assert free == {"t_free", "t_missing"} - - -def test_filter_free_tenants_respects_grace_period( - monkeypatch: pytest.MonkeyPatch, config_overrides: Callable[..., None] -) -> None: - config_overrides( - DEPLOYMENT_EDITION=DeploymentEdition.CLOUD, - SANDBOX_EXPIRED_RECORDS_CLEAN_GRACEFUL_PERIOD=45, - ) - cleanup = create_cleanup(monkeypatch, repo=FakeRepo([]), days=30, batch_size=10) - - now = datetime.datetime.now(datetime.UTC) - within_grace_ts = int((now - datetime.timedelta(days=10)).timestamp()) - outside_grace_ts = int((now - datetime.timedelta(days=90)).timestamp()) - - def fake_bulk(_: list[str]) -> dict[str, SubscriptionPlan]: - return { - "recently_downgraded": plan_info("sandbox", within_grace_ts), - "long_sandbox": plan_info("sandbox", outside_grace_ts), - } - - monkeypatch.setattr(cleanup_module.BillingService, "get_plan_bulk_with_cache", staticmethod(fake_bulk)) - - free = cleanup._filter_free_tenants({"recently_downgraded", "long_sandbox"}) - - assert free == {"long_sandbox"} - - -def test_filter_free_tenants_skips_cleanup_whitelist( - monkeypatch: pytest.MonkeyPatch, config_overrides: Callable[..., None] -) -> None: - config_overrides(DEPLOYMENT_EDITION=DeploymentEdition.CLOUD) - cleanup = create_cleanup( - monkeypatch, - repo=FakeRepo([]), - days=30, - batch_size=10, - whitelist={"tenant_whitelist"}, - ) - - monkeypatch.setattr( - cleanup_module.BillingService, - "get_plan_bulk_with_cache", - staticmethod( - lambda tenant_ids: { - tenant_id: (plan_info("team", -1) if tenant_id == "t_paid" else plan_info("sandbox", -1)) - for tenant_id in tenant_ids - } - ), - ) - - tenants = {"tenant_whitelist", "tenant_regular"} - free = cleanup._filter_free_tenants(tenants) - - assert free == {"tenant_regular"} - - -def test_filter_free_tenants_bulk_failure( - monkeypatch: pytest.MonkeyPatch, config_overrides: Callable[..., None] -) -> None: - config_overrides(DEPLOYMENT_EDITION=DeploymentEdition.CLOUD) - cleanup = create_cleanup(monkeypatch, repo=FakeRepo([]), days=30, batch_size=10) - - monkeypatch.setattr( - cleanup_module.BillingService, - "get_plan_bulk_with_cache", - staticmethod(lambda tenant_ids: (_ for _ in ()).throw(RuntimeError("boom"))), - ) - - free = cleanup._filter_free_tenants({"t1", "t2"}) - - assert free == set() - - -def test_run_deletes_only_free_tenants(monkeypatch: pytest.MonkeyPatch, config_overrides: Callable[..., None]) -> None: - config_overrides(DEPLOYMENT_EDITION=DeploymentEdition.CLOUD) - cutoff = datetime.datetime.now() - repo = FakeRepo( - batches=[ - [ - make_ref("run-free", "t_free", cutoff), - make_ref("run-paid", "t_paid", cutoff), - ] - ] - ) - cleanup = create_cleanup(monkeypatch, repo=repo, days=30, batch_size=10) - - monkeypatch.setattr( - cleanup_module.BillingService, - "get_plan_bulk_with_cache", - staticmethod( - lambda tenant_ids: { - tenant_id: (plan_info("team", -1) if tenant_id == "t_paid" else plan_info("sandbox", -1)) - for tenant_id in tenant_ids - } - ), - ) - - cleanup.run() - - assert repo.deleted == [["run-free"]] - assert repo.cleanup_ref_calls[1]["tenant_ids"] == ["t_free"] - - -def test_run_filters_candidate_tenants_before_target_query( - monkeypatch: pytest.MonkeyPatch, config_overrides: Callable[..., None] -) -> None: - config_overrides(DEPLOYMENT_EDITION=DeploymentEdition.CLOUD) - cutoff = datetime.datetime.now() - repo = FakeRepo( - batches=[ - [ - make_ref("run-free", "t_free", cutoff), - make_ref("run-paid", "t_paid", cutoff), - ] - ] - ) - cleanup = create_cleanup(monkeypatch, repo=repo, days=30, batch_size=10) - - billing_calls: list[list[str]] = [] - - def fake_bulk(tenant_ids: list[str]) -> dict[str, SubscriptionPlan]: - billing_calls.append(tenant_ids) - return { - "t_free": plan_info("sandbox", -1), - "t_paid": plan_info("team", -1), - } - - monkeypatch.setattr(cleanup_module.BillingService, "get_plan_bulk_with_cache", staticmethod(fake_bulk)) - - cleanup.run() - - assert billing_calls == [["t_free", "t_paid"]] - assert repo.cleanup_ref_calls[1]["tenant_ids"] == ["t_free"] - assert repo.deleted == [["run-free"]] - - -def test_run_skips_when_no_free_tenants(monkeypatch: pytest.MonkeyPatch, config_overrides: Callable[..., None]) -> None: - config_overrides(DEPLOYMENT_EDITION=DeploymentEdition.CLOUD) - cutoff = datetime.datetime.now() - repo = FakeRepo(batches=[[make_ref("run-paid", "t_paid", cutoff)]]) - cleanup = create_cleanup(monkeypatch, repo=repo, days=30, batch_size=10) - - monkeypatch.setattr( - cleanup_module.BillingService, - "get_plan_bulk_with_cache", - staticmethod(lambda tenant_ids: {tenant_id: plan_info("team", 1893456000) for tenant_id in tenant_ids}), - ) - - cleanup.run() - - assert repo.deleted == [] - assert len(repo.cleanup_ref_calls) == 2 - - -def test_run_paid_only_records_skipped_metrics( - monkeypatch: pytest.MonkeyPatch, config_overrides: Callable[..., None] -) -> None: - config_overrides(DEPLOYMENT_EDITION=DeploymentEdition.CLOUD) - cutoff = datetime.datetime.now() - repo = FakeRepo(batches=[[make_ref("run-paid", "t_paid", cutoff)]]) - cleanup = create_cleanup(monkeypatch, repo=repo, days=30, batch_size=10) - - monkeypatch.setattr( - cleanup_module.BillingService, - "get_plan_bulk_with_cache", - staticmethod(lambda tenant_ids: {tenant_id: plan_info("team", 1893456000) for tenant_id in tenant_ids}), - ) - batch_calls: list[dict[str, object]] = [] - monkeypatch.setattr(cleanup._metrics, "record_batch", lambda **kwargs: batch_calls.append(kwargs)) - - cleanup.run() - - assert repo.deleted == [] - assert repo.counted == [] - assert batch_calls[0]["batch_rows"] == 1 - assert batch_calls[0]["targeted_runs"] == 0 - assert batch_calls[0]["skipped_runs"] == 1 - assert batch_calls[0]["deleted_runs"] == 0 - - -def test_run_target_query_is_bounded_by_candidate_high_water(monkeypatch: pytest.MonkeyPatch) -> None: - first_created_at = datetime.datetime(2024, 1, 1, 0, 0, 0) - second_created_at = datetime.datetime(2024, 1, 1, 0, 1, 0) - repo = FakeRepo( - batches=[ - [ - make_ref("run-free-1", "t_free", first_created_at), - make_ref("run-free-2", "t_free", second_created_at), - ] - ] - ) - cleanup = create_cleanup(monkeypatch, repo=repo, days=30, batch_size=2) - - cleanup.run() - - assert repo.cleanup_ref_calls[1]["last_seen"] is None - assert repo.cleanup_ref_calls[1]["upper_bound"] == (second_created_at, "run-free-2") - assert repo.cleanup_ref_calls[2]["last_seen"] == (second_created_at, "run-free-2") - - -def test_run_exits_on_empty_batch(monkeypatch: pytest.MonkeyPatch) -> None: - cleanup = create_cleanup(monkeypatch, repo=FakeRepo([]), days=30, batch_size=10) - - cleanup.run() - - -def test_run_records_metrics_on_success(monkeypatch: pytest.MonkeyPatch) -> None: - cutoff = datetime.datetime.now() - repo = FakeRepo( - batches=[[make_ref("run-free", "t_free", cutoff)]], - delete_result={ - "runs": 0, - "node_executions": 2, - "offloads": 1, - "app_logs": 3, - "trigger_logs": 4, - "pauses": 5, - "pause_reasons": 6, - }, - ) - cleanup = create_cleanup(monkeypatch, repo=repo, days=30, batch_size=10) - - batch_calls: list[dict[str, object]] = [] - completion_calls: list[dict[str, object]] = [] - monkeypatch.setattr(cleanup._metrics, "record_batch", lambda **kwargs: batch_calls.append(kwargs)) - monkeypatch.setattr(cleanup._metrics, "record_completion", lambda **kwargs: completion_calls.append(kwargs)) - - cleanup.run() - - assert len(batch_calls) == 1 - assert batch_calls[0]["batch_rows"] == 1 - assert batch_calls[0]["targeted_runs"] == 1 - assert batch_calls[0]["deleted_runs"] == 1 - assert batch_calls[0]["related_action"] == "deleted" - assert len(completion_calls) == 1 - assert completion_calls[0]["status"] == "success" - - -def test_run_records_failed_metrics(monkeypatch: pytest.MonkeyPatch) -> None: - class FailingRepo(FakeRepo): - def delete_runs_with_related_by_ids( - self, run_ids: list[str], delete_node_executions=None, delete_trigger_logs=None - ) -> dict[str, int]: - raise RuntimeError("delete failed") - - cutoff = datetime.datetime.now() - repo = FailingRepo(batches=[[make_ref("run-free", "t_free", cutoff)]]) - cleanup = create_cleanup(monkeypatch, repo=repo, days=30, batch_size=10) - - completion_calls: list[dict[str, object]] = [] - monkeypatch.setattr(cleanup._metrics, "record_completion", lambda **kwargs: completion_calls.append(kwargs)) - - with pytest.raises(RuntimeError, match="delete failed"): - cleanup.run() - - assert len(completion_calls) == 1 - assert completion_calls[0]["status"] == "failed" - - -def test_run_dry_run_skips_deletions(monkeypatch: pytest.MonkeyPatch, capsys: pytest.CaptureFixture[str]) -> None: - cutoff = datetime.datetime.now() - repo = FakeRepo( - batches=[[make_ref("run-free", "t_free", cutoff)]], - count_result={ - "runs": 0, - "node_executions": 2, - "offloads": 1, - "app_logs": 3, - "trigger_logs": 4, - "pauses": 5, - "pause_reasons": 6, - }, - ) - cleanup = create_cleanup(monkeypatch, repo=repo, days=30, batch_size=10, dry_run=True) - - cleanup.run() - - assert repo.deleted == [] - assert repo.counted == [["run-free"]] - captured = capsys.readouterr().out - assert "Dry run mode enabled" in captured - assert "would delete 1 runs" in captured - assert "related records" in captured - assert "node_executions 2" in captured - assert "offloads 1" in captured - assert "app_logs 3" in captured - assert "trigger_logs 4" in captured - assert "pauses 5" in captured - assert "pause_reasons 6" in captured - - -def test_between_sets_window_bounds(monkeypatch: pytest.MonkeyPatch) -> None: - start_from = datetime.datetime(2024, 5, 1, 0, 0, 0) - end_before = datetime.datetime(2024, 6, 1, 0, 0, 0) - cleanup = create_cleanup( - monkeypatch, repo=FakeRepo([]), days=30, batch_size=10, start_from=start_from, end_before=end_before - ) - - assert cleanup.window_start == start_from - assert cleanup.window_end == end_before - - -def test_between_requires_both_boundaries(monkeypatch: pytest.MonkeyPatch) -> None: - with pytest.raises(ValueError, match="start_from and end_before must be both set or both omitted\\."): - create_cleanup( - monkeypatch, repo=FakeRepo([]), days=30, batch_size=10, start_from=datetime.datetime.now(), end_before=None - ) - with pytest.raises(ValueError, match="start_from and end_before must be both set or both omitted\\."): - create_cleanup( - monkeypatch, repo=FakeRepo([]), days=30, batch_size=10, start_from=None, end_before=datetime.datetime.now() - ) - - -def test_between_requires_end_after_start(monkeypatch: pytest.MonkeyPatch) -> None: - start_from = datetime.datetime(2024, 6, 1, 0, 0, 0) - end_before = datetime.datetime(2024, 5, 1, 0, 0, 0) - with pytest.raises(ValueError, match="end_before must be greater than start_from\\."): - create_cleanup( - monkeypatch, repo=FakeRepo([]), days=30, batch_size=10, start_from=start_from, end_before=end_before - ) diff --git a/api/tests/unit_tests/services/test_databricks_tracing_config_gateway.py b/api/tests/unit_tests/services/test_databricks_tracing_config_gateway.py deleted file mode 100644 index 4704659408a..00000000000 --- a/api/tests/unit_tests/services/test_databricks_tracing_config_gateway.py +++ /dev/null @@ -1,218 +0,0 @@ -"""Databricks config persistence must not initialize a process-global MLflow client.""" - -from collections.abc import Callable, Iterator, Sequence -from dataclasses import dataclass, field -from typing import Never -from uuid import uuid4 - -import mlflow -import pytest -from dify_trace_mlflow.configuration import DatabricksConfigurationError -from dify_trace_mlflow.mlflow_trace import MLflowDataTrace -from sqlalchemy.orm import Session, sessionmaker - -from core.helper import ssrf_proxy -from core.ops import ops_trace_manager -from core.ops.ops_trace_manager import OpsTraceManager -from machinery.context import RequestContext -from repositories.app.tracing_config_repository import SQLAlchemyAppTracingConfigRepository -from services.app_tracing_config_gateway import OpsTraceManagerGateway -from services.app_tracing_config_service import AppTracingConfigService, AppTracingConfigVerificationFailedError -from tests.unit_tests.model_factories import make_app -from tests.unit_tests.services.tracing.test_databricks_configuration import ExperimentAPI, response - - -@dataclass -class GatewayEnvironment: - gateway: OpsTraceManagerGateway = field(default_factory=OpsTraceManagerGateway) - api: ExperimentAPI = field(default_factory=ExperimentAPI) - forbidden_calls: list[str] = field(default_factory=list) - - -@pytest.fixture -def environment(monkeypatch: pytest.MonkeyPatch) -> Iterator[GatewayEnvironment]: - environment = GatewayEnvironment() - - def encrypt(_tenant: str, value: str) -> str: - return "test-cipher:" + value - - def decrypt(_tenant: str, values: Sequence[str]) -> list[str]: - assert all(value.startswith("test-cipher:") for value in values) - return [value.removeprefix("test-cipher:") for value in values] - - def forbidden(name: str) -> Callable[..., Never]: - def block(*_args: object, **_kwargs: object) -> Never: - environment.forbidden_calls.append(name) - raise AssertionError("Configuration must not initialize a stateful exporter") - - return block - - monkeypatch.setattr(ops_trace_manager, "encrypt_token", encrypt) - monkeypatch.setattr(ops_trace_manager, "batch_decrypt_token", decrypt) - monkeypatch.setattr(ops_trace_manager, "obfuscated_token", lambda value: "******" if value else value) - monkeypatch.setattr(ssrf_proxy, "get", environment.api.get) - monkeypatch.setattr(ssrf_proxy, "post", environment.api.post) - monkeypatch.setattr(MLflowDataTrace, "__init__", forbidden("legacy_constructor")) - monkeypatch.setattr(mlflow, "set_tracking_uri", forbidden("tracking_uri")) - monkeypatch.setattr(mlflow, "set_experiment", forbidden("experiment")) - monkeypatch.setattr(mlflow, "search_experiments", forbidden("search")) - monkeypatch.setattr(mlflow.config, "enable_async_logging", forbidden("async_logging")) - yield environment - assert environment.forbidden_calls == [] - - -@pytest.mark.parametrize("oauth", [False, True]) -def test_create_and_patch_persist_verified_target_without_replacing_masked_credentials( - environment: GatewayEnvironment, sqlite_session_factory: sessionmaker[Session], oauth: bool -) -> None: - workspace, app_id = str(uuid4()), str(uuid4()) - with sqlite_session_factory.begin() as session: - session.add(make_app(app_id=app_id, tenant_id=workspace)) - store = SQLAlchemyAppTracingConfigRepository(session_factory=sqlite_session_factory) - service = AppTracingConfigService(configs=store, provider=environment.gateway) - context = RequestContext("request", None, "account", workspace) - credentials = ( - {"client_id": "synthetic-client", "client_secret": "synthetic-secret"} - if oauth - else {"personal_access_token": "synthetic-pat"} - ) - secret_field = "client_secret" if oauth else "personal_access_token" - service.create( - context, app_id, "databricks", {"host": " https://first.test/ ", "experiment_id": "00012", **credentials} - ) - before = store.get(workspace_id=workspace, app_id=app_id, tracing_provider="databricks") - assert before is not None - assert before.tracing_config is not None - assert before.tracing_config["host"] == "https://first.test" - assert before.tracing_config["experiment_id"] == "12" - service.update( - context, - app_id, - "databricks", - { - "host": " https://second.test/ ", - "experiment_id": "00034", - **credentials, - secret_field: "******", - }, - ) - after = store.get(workspace_id=workspace, app_id=app_id, tracing_provider="databricks") - assert after is not None - assert after.tracing_config is not None - assert after.tracing_config["host"] == "https://second.test" - assert after.tracing_config["experiment_id"] == "34" - assert after.tracing_config[secret_field] == before.tracing_config[secret_field] - assert after.tracing_config["project_url"] == "https://second.test/ml/experiments/34/traces" - assert before.tracing_config["project_url"] == "https://first.test/ml/experiments/12/traces" - before_get = len(environment.api.calls) - presented = service.get(context, app_id, "databricks") - assert presented is not None - assert presented.tracing_config is not None - assert presented.tracing_config[secret_field] == "******" - assert presented.tracing_config["project_url"] == after.tracing_config["project_url"] - environment.gateway.validate_saved_config( - workspace_id=workspace, tracing_provider="databricks", tracing_config=after.tracing_config - ) - assert len(environment.api.calls) == before_get - - -@pytest.mark.parametrize("saved_url", [None, "https://stale.test/old-project"]) -def test_get_is_pure_for_historical_and_current_configs(environment: GatewayEnvironment, saved_url: str | None) -> None: - workspace = str(uuid4()) - encrypted = OpsTraceManager.encrypt_tracing_config( - workspace, - "databricks", - {"host": "https://workspace.test/", "experiment_id": "12", "personal_access_token": "synthetic-pat"}, - ) - if saved_url is not None: - encrypted["project_url"] = saved_url - result = environment.gateway.present_config( - workspace_id=workspace, tracing_provider="databricks", tracing_config=encrypted - ) - assert result["project_url"] == "https://workspace.test/ml/experiments/12/traces" - assert result["personal_access_token"] == "******" - assert environment.api.calls == [] - - -def test_invalid_historical_host_uses_safe_pure_fallback(environment: GatewayEnvironment) -> None: - workspace = str(uuid4()) - encrypted = OpsTraceManager.encrypt_tracing_config( - workspace, - "databricks", - { - "host": "https://user:password@workspace.test", - "experiment_id": "12", - "personal_access_token": "synthetic-pat", - }, - ) - result = environment.gateway.present_config( - workspace_id=workspace, tracing_provider="databricks", tracing_config=encrypted - ) - assert result["project_url"] == "https://www.databricks.com/" - assert environment.api.calls == [] - - -def test_failed_patch_keeps_original_destination_and_logs_only_safe_code( - environment: GatewayEnvironment, sqlite_session_factory: sessionmaker[Session], caplog: pytest.LogCaptureFixture -) -> None: - workspace, app_id = str(uuid4()), str(uuid4()) - with sqlite_session_factory.begin() as session: - session.add(make_app(app_id=app_id, tenant_id=workspace)) - store = SQLAlchemyAppTracingConfigRepository(session_factory=sqlite_session_factory) - service = AppTracingConfigService(configs=store, provider=environment.gateway) - context = RequestContext("request", None, "account", workspace) - original = { - "host": "https://workspace.test", - "experiment_id": "12", - "personal_access_token": "synthetic-private-pat", - } - service.create(context, app_id, "databricks", original) - before = store.get(workspace_id=workspace, app_id=app_id, tracing_provider="databricks") - assert before is not None - environment.api.override = lambda _method, url, _options: response(url, {"message": "private-vendor-detail"}, 403) - with pytest.raises(AppTracingConfigVerificationFailedError) as caught: - service.update( - context, app_id, "databricks", {**original, "experiment_id": "34", "personal_access_token": "******"} - ) - assert isinstance(caught.value.__cause__, DatabricksConfigurationError) - assert caught.value.__cause__.code == "permission_denied" - after = store.get(workspace_id=workspace, app_id=app_id, tracing_provider="databricks") - assert after is not None - assert after.tracing_config == before.tracing_config - assert "permission_denied" in caplog.text - assert "private-vendor-detail" not in caplog.text - assert "synthetic-private-pat" not in caplog.text - - -@pytest.mark.parametrize("start_with_oauth", [False, True]) -def test_switching_auth_clears_unselected_credentials_in_persistence_and_requests( - environment: GatewayEnvironment, sqlite_session_factory: sessionmaker[Session], start_with_oauth: bool -) -> None: - workspace, app_id = str(uuid4()), str(uuid4()) - with sqlite_session_factory.begin() as session: - session.add(make_app(app_id=app_id, tenant_id=workspace)) - store = SQLAlchemyAppTracingConfigRepository(session_factory=sqlite_session_factory) - service = AppTracingConfigService(configs=store, provider=environment.gateway) - context = RequestContext("request", None, "account", workspace) - pat = {"personal_access_token": "synthetic-pat", "client_id": None, "client_secret": None} - oauth = {"personal_access_token": None, "client_id": "synthetic-client", "client_secret": "synthetic-secret"} - target = {"host": "https://workspace.test", "experiment_id": "12"} - service.create(context, app_id, "databricks", {**target, **(oauth if start_with_oauth else pat)}) - environment.api.calls.clear() - service.update(context, app_id, "databricks", {**target, **(pat if start_with_oauth else oauth)}) - stored = store.get(workspace_id=workspace, app_id=app_id, tracing_provider="databricks") - assert stored is not None - assert stored.tracing_config is not None - assert stored.tracing_config["client_secret" if start_with_oauth else "personal_access_token"] is None - assert [method for method, _url, _options in environment.api.calls] == ( - ["GET"] if start_with_oauth else ["POST", "GET"] - ) - expected_token = "synthetic-pat" if start_with_oauth else "synthetic-oauth-token" - assert environment.api.calls[-1][2]["headers"] == { - "Authorization": f"Bearer {expected_token}", - "Accept": "application/json", - } - presented = service.get(context, app_id, "databricks") - assert presented is not None - assert presented.tracing_config is not None - assert presented.tracing_config["client_secret" if start_with_oauth else "personal_access_token"] is None diff --git a/api/tests/unit_tests/services/test_knowledge_fs_upgrade_file_lease.py b/api/tests/unit_tests/services/test_dataset_knowledge_fs_upgrade_file_lease.py similarity index 100% rename from api/tests/unit_tests/services/test_knowledge_fs_upgrade_file_lease.py rename to api/tests/unit_tests/services/test_dataset_knowledge_fs_upgrade_file_lease.py diff --git a/api/tests/unit_tests/services/test_knowledge_fs_upgrade_service.py b/api/tests/unit_tests/services/test_dataset_knowledge_fs_upgrade_service.py similarity index 99% rename from api/tests/unit_tests/services/test_knowledge_fs_upgrade_service.py rename to api/tests/unit_tests/services/test_dataset_knowledge_fs_upgrade_service.py index 67cfeaf65ed..cd00755e20f 100644 --- a/api/tests/unit_tests/services/test_knowledge_fs_upgrade_service.py +++ b/api/tests/unit_tests/services/test_dataset_knowledge_fs_upgrade_service.py @@ -72,7 +72,7 @@ from services.knowledge_fs.product_dto import ( KnowledgeFSStagedUploadResponse, ) from services.knowledge_fs.product_remote import KnowledgeFSProductResourceNotFoundError -from tests.unit_tests.services.knowledge_fs_services_a_fakes import CallRecorder, RecordedCall, UnusedSessionFactory +from tests.helpers.unit.services.knowledge_fs_services_a_fakes import CallRecorder, RecordedCall, UnusedSessionFactory _TENANT_ID = "00000000-0000-0000-0000-000000000001" _ACCOUNT_ID = "00000000-0000-0000-0000-000000000002" diff --git a/api/tests/unit_tests/services/test_dataset_service_dataset.py b/api/tests/unit_tests/services/test_dataset_service_dataset.py deleted file mode 100644 index e85f0c227fe..00000000000 --- a/api/tests/unit_tests/services/test_dataset_service_dataset.py +++ /dev/null @@ -1,1262 +0,0 @@ -"""SQLite-backed tests for dataset-level operations in :mod:`services.knowledge.dataset_service`. - -Mapped objects in this module are real SQLAlchemy models. Provider runtimes, -RBAC clients, Celery tasks, and model-manager results remain mocked at their -external boundaries. -""" - -from __future__ import annotations - -from collections.abc import Callable -from types import SimpleNamespace -from unittest.mock import patch - -import pytest -from sqlalchemy import event, select -from sqlalchemy.orm import Session, sessionmaker - -from core.errors.error import LLMBadRequestError, ProviderTokenNotInitError -from core.rag.index_processor.constant.index_type import IndexTechniqueType -from graphon.model_runtime.entities.model_entities import ModelFeature, ModelType -from models import Account -from models.dataset import ( - Dataset, - DatasetCollectionBinding, - DatasetPermission, - DatasetPermissionEnum, - ExternalKnowledgeApis, - ExternalKnowledgeBindings, - Pipeline, -) -from services.enterprise import rbac_service -from services.entities.knowledge_entities.rag_pipeline_entities import ( - IconInfo, - RagPipelineDatasetCreateEntity, -) -from services.errors.base import NoPermissionError -from services.errors.dataset import DatasetNameDuplicateError -from services.knowledge.dataset_service import DatasetCollectionBindingService, DatasetPermissionService, DatasetService -from tests.unit_tests.model_factories import make_account, make_dataset, make_tenant - -from .dataset_service_test_helpers import ( - MagicMock, - TenantAccountRole, - _make_knowledge_configuration, - _make_retrieval_model, -) - - -def _account( - *, - account_id: str = "user-1", - tenant_id: str = "tenant-1", - role: TenantAccountRole = TenantAccountRole.OWNER, -) -> Account: - return make_account( - account_id=account_id, - name=f"User {account_id}", - email=f"{account_id}@example.com", - role=role, - tenant=make_tenant(tenant_id=tenant_id, name=f"Tenant {tenant_id}"), - ) - - -def _dataset( - *, - dataset_id: str = "dataset-1", - tenant_id: str = "tenant-1", - name: str = "Dataset", - maintainer: str = "user-1", - permission: DatasetPermissionEnum = DatasetPermissionEnum.ALL_TEAM, - provider: str = "vendor", - indexing_technique: str = IndexTechniqueType.ECONOMY, - chunk_structure: str | None = "text_model", -) -> Dataset: - return make_dataset( - dataset_id=dataset_id, - tenant_id=tenant_id, - name=name, - description="", - provider=provider, - created_by=maintainer, - maintainer=maintainer, - permission=permission, - indexing_technique=indexing_technique, - chunk_structure=chunk_structure, - embedding_model_provider="provider", - embedding_model="embedding-model", - ) - - -def _external_api(*, api_id: str = "api-1", tenant_id: str = "tenant-1") -> ExternalKnowledgeApis: - api = ExternalKnowledgeApis( - name="External API", - description="", - tenant_id=tenant_id, - settings="{}", - created_by="user-1", - updated_by="user-1", - ) - api.id = api_id - return api - - -def _binding( - *, - binding_id: str = "binding-1", - tenant_id: str = "tenant-1", - dataset_id: str = "dataset-1", - api_id: str = "api-1", - knowledge_id: str = "knowledge-1", -) -> ExternalKnowledgeBindings: - binding = ExternalKnowledgeBindings( - tenant_id=tenant_id, - dataset_id=dataset_id, - external_knowledge_api_id=api_id, - external_knowledge_id=knowledge_id, - created_by="user-1", - ) - binding.id = binding_id - return binding - - -class TestDatasetServiceValidation: - @pytest.mark.parametrize( - ("dataset_doc_form", "incoming_doc_form"), - [(None, "text_model"), ("text_model", "text_model")], - ) - def test_check_doc_form_allows_matching_or_missing_dataset_doc_form( - self, - sqlite_session: Session, - dataset_doc_form: str | None, - incoming_doc_form: str, - ) -> None: - dataset = _dataset(chunk_structure=dataset_doc_form) - sqlite_session.add(dataset) - sqlite_session.commit() - - DatasetService.check_doc_form(dataset, incoming_doc_form, session=sqlite_session) - - def test_check_doc_form_rejects_mismatched_doc_form(self, sqlite_session: Session) -> None: - dataset = _dataset(chunk_structure="qa_model") - sqlite_session.add(dataset) - sqlite_session.commit() - - with pytest.raises(ValueError, match="doc_form is different"): - DatasetService.check_doc_form(dataset, "text_model", session=sqlite_session) - - @pytest.mark.parametrize("operator_check", [False, True]) - def test_dataset_permission_checks_ignore_foreign_tenant_binding( - self, sqlite_session: Session, operator_check: bool - ) -> None: - dataset = _dataset( - dataset_id="dataset-1", - tenant_id="tenant-1", - permission=DatasetPermissionEnum.PARTIAL_TEAM, - maintainer="owner-1", - ) - user = _account( - account_id="user-1", - tenant_id="tenant-1", - role=TenantAccountRole.NORMAL, - ) - sqlite_session.add_all( - [ - dataset, - DatasetPermission(dataset_id=dataset.id, account_id=user.id, tenant_id="tenant-2"), - ] - ) - sqlite_session.commit() - - if operator_check: - with pytest.raises(NoPermissionError): - DatasetService.check_dataset_operator_permission(user, dataset, session=sqlite_session) - else: - with pytest.raises(NoPermissionError): - DatasetService.check_dataset_permission(dataset, user, sqlite_session) - - def test_check_dataset_model_setting_skips_non_high_quality_datasets(self) -> None: - dataset = _dataset(indexing_technique=IndexTechniqueType.ECONOMY) - - with patch("services.knowledge.dataset_service.ModelManager") as model_manager_cls: - DatasetService.check_dataset_model_setting(dataset) - - model_manager_cls.assert_not_called() - - def test_check_dataset_model_setting_validates_high_quality_embedding(self) -> None: - dataset = _dataset(indexing_technique=IndexTechniqueType.HIGH_QUALITY) - - with patch("services.knowledge.dataset_service.ModelManager") as model_manager_cls: - DatasetService.check_dataset_model_setting(dataset) - - model_manager_cls.for_tenant.return_value.get_model_instance.assert_called_once_with( - tenant_id=dataset.tenant_id, - provider=dataset.embedding_model_provider, - model_type=ModelType.TEXT_EMBEDDING, - model=dataset.embedding_model, - ) - - @pytest.mark.parametrize( - ("error", "message"), - [ - (LLMBadRequestError(), "No Embedding Model available"), - (ProviderTokenNotInitError("token missing"), "token missing"), - ], - ) - def test_check_dataset_model_setting_wraps_provider_errors(self, error: Exception, message: str) -> None: - dataset = _dataset(indexing_technique=IndexTechniqueType.HIGH_QUALITY) - - with patch("services.knowledge.dataset_service.ModelManager") as model_manager_cls: - model_manager_cls.for_tenant.return_value.get_model_instance.side_effect = error - with pytest.raises(ValueError, match=message): - DatasetService.check_dataset_model_setting(dataset) - - @pytest.mark.parametrize(("features", "expected"), [([ModelFeature.VISION], True), ([], False)]) - def test_check_is_multimodal_model_reads_runtime_schema(self, features: list[ModelFeature], expected: bool) -> None: - model_type_instance = MagicMock() - model_type_instance.get_model_schema.return_value = SimpleNamespace(features=features) - model_instance = SimpleNamespace( - model_type_instance=model_type_instance, - model_name="embedding-model", - credentials={"api_key": "secret"}, - ) - - with patch("services.knowledge.dataset_service.ModelManager") as model_manager_cls: - model_manager_cls.for_tenant.return_value.get_model_instance.return_value = model_instance - result = DatasetService.check_is_multimodal_model("tenant-1", "provider", "embedding-model") - - assert result is expected - - def test_check_is_multimodal_model_rejects_missing_schema(self) -> None: - model_type_instance = MagicMock() - model_type_instance.get_model_schema.return_value = None - model_instance = SimpleNamespace( - model_type_instance=model_type_instance, - model_name="embedding-model", - credentials={}, - ) - - with patch("services.knowledge.dataset_service.ModelManager") as model_manager_cls: - model_manager_cls.for_tenant.return_value.get_model_instance.return_value = model_instance - with pytest.raises(ValueError, match="Model schema not found"): - DatasetService.check_is_multimodal_model("tenant-1", "provider", "embedding-model") - - @pytest.mark.parametrize( - ("method", "error", "message"), - [ - ( - DatasetService.check_embedding_model_setting, - ProviderTokenNotInitError("provider setup"), - "provider setup", - ), - (DatasetService.check_reranking_model_setting, LLMBadRequestError(), "No Rerank Model available"), - ], - ) - def test_direct_model_setting_checks_wrap_runtime_errors( - self, method: Callable[[str, str, str], None], error: Exception, message: str - ) -> None: - with patch("services.knowledge.dataset_service.ModelManager") as model_manager_cls: - model_manager_cls.for_tenant.return_value.get_model_instance.side_effect = error - with pytest.raises(ValueError, match=message): - method("tenant-1", "provider", "model") - - -class TestDatasetServiceRetrieval: - def test_get_dataset_for_tenant_rejects_cross_tenant_row(self, sqlite_session: Session) -> None: - owned = _dataset() - foreign = _dataset(dataset_id="dataset-2", tenant_id="tenant-2", name="Foreign") - sqlite_session.add_all([owned, foreign]) - sqlite_session.commit() - - assert DatasetService.get_dataset_for_tenant(owned.id, "tenant-1", session=sqlite_session) is owned - assert DatasetService.get_dataset_for_tenant(foreign.id, "tenant-1", session=sqlite_session) is None - - def test_get_datasets_filters_by_creator_ids( - self, application_tags, config_overrides: Callable[..., None], sqlite_session: Session - ) -> None: - creator_one = _dataset(dataset_id="creator-one", name="Creator One", maintainer="creator-1") - creator_two = _dataset(dataset_id="creator-two", name="Creator Two", maintainer="creator-2") - excluded = _dataset(dataset_id="excluded", name="Excluded", maintainer="creator-3") - foreign = _dataset( - dataset_id="foreign", - tenant_id="tenant-2", - name="Foreign", - maintainer="creator-1", - ) - sqlite_session.add_all([creator_one, creator_two, excluded, foreign]) - sqlite_session.commit() - - config_overrides(RBAC_ENABLED=False) - datasets, total = DatasetService.get_datasets( - page=1, - per_page=20, - session=sqlite_session, - tenant_id="tenant-1", - creator_ids=["creator-1", "creator-2"], - tags=application_tags, - ) - - assert total == 2 - assert {dataset.id for dataset in datasets} == {creator_one.id, creator_two.id} - - def test_get_datasets_applies_rbac_resource_scope_and_maintainer_override( - self, application_tags, config_overrides: Callable[..., None], sqlite_session: Session - ) -> None: - user = _account(role=TenantAccountRole.NORMAL) - accessible = _dataset(dataset_id="accessible", name="Accessible", maintainer="other") - owned = _dataset(dataset_id="owned", name="Owned", maintainer=user.id) - hidden = _dataset(dataset_id="hidden", name="Hidden", maintainer="other") - foreign = _dataset(dataset_id="foreign", tenant_id="tenant-2", name="Foreign") - sqlite_session.add_all([accessible, owned, hidden, foreign]) - sqlite_session.commit() - - config_overrides(RBAC_ENABLED=True) - - with ( - patch( - "services.knowledge.dataset_service.enterprise_rbac_service.RBACService.MyPermissions.get", - return_value=SimpleNamespace(workspace=SimpleNamespace(permission_keys=[])), - ), - ): - datasets, total = DatasetService.get_datasets( - page=1, - per_page=20, - session=sqlite_session, - tenant_id="tenant-1", - user=user, - accessible_dataset_ids=[accessible.id], - include_own_datasets=True, - tags=application_tags, - ) - - assert total == 2 - assert {dataset.id for dataset in datasets} == {accessible.id, owned.id} - - def test_get_datasets_without_user_keeps_only_team_visible_rows( - self, application_tags, config_overrides: Callable[..., None], sqlite_session: Session - ) -> None: - shared = _dataset(dataset_id="shared", name="Shared", permission=DatasetPermissionEnum.ALL_TEAM) - private = _dataset(dataset_id="private", name="Private", permission=DatasetPermissionEnum.ONLY_ME) - sqlite_session.add_all([shared, private]) - sqlite_session.commit() - - config_overrides(RBAC_ENABLED=False) - datasets, total = DatasetService.get_datasets( - page=1, per_page=20, session=sqlite_session, tenant_id="tenant-1", tags=application_tags - ) - - assert total == 1 - assert [dataset.id for dataset in datasets] == [shared.id] - - def test_get_datasets_by_ids_intersects_requested_and_accessible_ids( - self, config_overrides: Callable[..., None], sqlite_session: Session - ) -> None: - user = _account(role=TenantAccountRole.NORMAL) - accessible = _dataset(dataset_id="accessible", name="Accessible", maintainer="other") - owned = _dataset(dataset_id="owned", name="Owned", maintainer=user.id) - hidden = _dataset(dataset_id="hidden", name="Hidden", maintainer="other") - sqlite_session.add_all([accessible, owned, hidden]) - sqlite_session.commit() - - config_overrides(RBAC_ENABLED=True) - datasets, total = DatasetService.get_datasets_by_ids( - [accessible.id, owned.id, hidden.id], - "tenant-1", - user=user, - accessible_dataset_ids=[accessible.id, "not-requested"], - include_own_datasets=True, - session=sqlite_session, - ) - - assert total == 2 - assert {dataset.id for dataset in datasets} == {accessible.id, owned.id} - - def test_get_datasets_rbac_without_user_returns_no_rows( - self, application_tags, config_overrides: Callable[..., None], sqlite_session: Session - ) -> None: - sqlite_session.add(_dataset()) - sqlite_session.commit() - - config_overrides(RBAC_ENABLED=True) - datasets, total = DatasetService.get_datasets( - page=1, per_page=20, session=sqlite_session, tenant_id="tenant-1", tags=application_tags - ) - - assert datasets == [] - assert total == 0 - - def test_get_datasets_rbac_include_all_requires_workspace_permission( - self, application_tags, config_overrides: Callable[..., None], sqlite_session: Session - ) -> None: - user = _account(role=TenantAccountRole.NORMAL) - sqlite_session.add_all( - [ - _dataset(dataset_id="one", name="One", maintainer="other"), - _dataset(dataset_id="two", name="Two", maintainer="other"), - ] - ) - sqlite_session.commit() - - config_overrides(RBAC_ENABLED=True) - with ( - patch( - "services.knowledge.dataset_service.enterprise_rbac_service.RBACService.MyPermissions.get", - return_value=SimpleNamespace( - workspace=SimpleNamespace(permission_keys=["dataset.create_and_management"]) - ), - ), - ): - datasets, total = DatasetService.get_datasets( - page=1, - per_page=20, - session=sqlite_session, - tenant_id="tenant-1", - user=user, - include_all=True, - tags=application_tags, - ) - - assert total == 2 - assert {dataset.id for dataset in datasets} == {"one", "two"} - - -class TestDatasetServiceCreationAndUpdate: - def test_create_empty_dataset_rejects_duplicate_name(self, sqlite_session: Session) -> None: - sqlite_session.add(_dataset(name="Existing")) - sqlite_session.commit() - - with pytest.raises(DatasetNameDuplicateError, match="already exists"): - DatasetService.create_empty_dataset( - "tenant-1", - "Existing", - "", - IndexTechniqueType.ECONOMY, - _account(), - session=sqlite_session, - ) - - def test_create_empty_dataset_persists_default_embedding_and_creator_scope( - self, - sqlite_session: Session, - sqlite_session_factory: sessionmaker[Session], - ) -> None: - embedding_model = SimpleNamespace(provider="provider", model_name="default-embedding") - - with ( - patch("services.knowledge.dataset_service.ModelManager") as model_manager_cls, - patch( - "services.knowledge.dataset_service.enterprise_rbac_service.try_sync_creator_access_policy_member_bindings" - ), - ): - model_manager_cls.for_tenant.return_value.get_default_model_instance.return_value = embedding_model - dataset = DatasetService.create_empty_dataset( - "tenant-1", - "Created", - "Description", - IndexTechniqueType.HIGH_QUALITY, - _account(), - session=sqlite_session, - ) - - with sqlite_session_factory() as observer: - persisted = observer.get(Dataset, dataset.id) - assert persisted is not None - assert persisted.tenant_id == "tenant-1" - assert persisted.maintainer == "user-1" - assert persisted.embedding_model == "default-embedding" - - def test_create_empty_external_dataset_persists_tenant_owned_binding(self, sqlite_session: Session) -> None: - sqlite_session.add_all([_external_api(), _external_api(api_id="api-foreign", tenant_id="tenant-2")]) - sqlite_session.commit() - - with patch( - "services.knowledge.dataset_service.enterprise_rbac_service.try_sync_creator_access_policy_member_bindings" - ): - dataset = DatasetService.create_empty_dataset( - "tenant-1", - "External", - "", - IndexTechniqueType.ECONOMY, - _account(), - provider="external", - external_knowledge_api_id="api-1", - external_knowledge_id="knowledge-1", - retrieval_model=_make_retrieval_model(), - session=sqlite_session, - ) - - binding = sqlite_session.scalar( - select(ExternalKnowledgeBindings).where(ExternalKnowledgeBindings.dataset_id == dataset.id) - ) - assert binding is not None - assert binding.tenant_id == dataset.tenant_id - assert binding.external_knowledge_api_id == "api-1" - - def test_create_empty_external_dataset_rejects_foreign_api(self, sqlite_session: Session) -> None: - sqlite_session.add(_external_api(api_id="api-foreign", tenant_id="tenant-2")) - sqlite_session.commit() - - with ( - pytest.raises(ValueError, match="api template not found"), - patch( - "services.knowledge.dataset_service.enterprise_rbac_service.try_sync_creator_access_policy_member_bindings" - ), - ): - DatasetService.create_empty_dataset( - "tenant-1", - "External", - "", - IndexTechniqueType.ECONOMY, - _account(), - provider="external", - external_knowledge_api_id="api-foreign", - external_knowledge_id="knowledge-1", - session=sqlite_session, - ) - - def test_create_empty_rag_pipeline_dataset_generates_incremental_name(self, sqlite_session: Session) -> None: - sqlite_session.add(_dataset(name="Untitled 1")) - sqlite_session.commit() - entity = RagPipelineDatasetCreateEntity( - name="", - description="Pipeline dataset", - icon_info=IconInfo(icon_type="emoji", icon="📚", icon_background="#FFFFFF", icon_url=None), - permission=DatasetPermissionEnum.ALL_TEAM, - ) - - with patch("services.knowledge.dataset_service.current_user", _account()): - dataset = DatasetService.create_empty_rag_pipeline_dataset("tenant-1", entity, sqlite_session) - - assert dataset.name == "Untitled 2" - pipeline = sqlite_session.get(Pipeline, dataset.pipeline_id) - assert pipeline is not None - assert pipeline.tenant_id == dataset.tenant_id - - def test_create_empty_rag_pipeline_dataset_rejects_duplicate_name(self, sqlite_session: Session) -> None: - sqlite_session.add(_dataset(name="Existing")) - sqlite_session.commit() - entity = RagPipelineDatasetCreateEntity( - name="Existing", - description="", - icon_info=IconInfo(icon_type="emoji", icon="📚", icon_background="#FFFFFF", icon_url=None), - permission=DatasetPermissionEnum.ALL_TEAM, - ) - - with ( - patch("services.knowledge.dataset_service.current_user", _account()), - pytest.raises(DatasetNameDuplicateError, match="already exists"), - ): - DatasetService.create_empty_rag_pipeline_dataset("tenant-1", entity, sqlite_session) - - def test_create_empty_rag_pipeline_dataset_requires_authenticated_account(self, sqlite_session: Session) -> None: - account = _account() - account.id = "" - entity = RagPipelineDatasetCreateEntity( - name="Dataset", - description="", - icon_info=IconInfo(icon_type="emoji", icon="📚", icon_background="#FFFFFF", icon_url=None), - permission=DatasetPermissionEnum.ALL_TEAM, - ) - - with ( - patch("services.knowledge.dataset_service.current_user", account), - pytest.raises(ValueError, match="Current user or current user id not found"), - ): - DatasetService.create_empty_rag_pipeline_dataset("tenant-1", entity, sqlite_session) - - def test_update_dataset_rejects_missing_and_duplicate_rows(self, sqlite_session: Session) -> None: - current = _dataset(name="Current") - duplicate = _dataset(dataset_id="dataset-2", name="Duplicate") - sqlite_session.add_all([current, duplicate]) - sqlite_session.commit() - - with pytest.raises(ValueError, match="Dataset not found"): - DatasetService.update_dataset("missing", {}, _account(), session=sqlite_session) - with pytest.raises(ValueError, match="Dataset name already exists"): - DatasetService.update_dataset(current.id, {"name": duplicate.name}, _account(), session=sqlite_session) - - @pytest.mark.parametrize( - ("provider", "helper_name"), - [("external", "_update_external_dataset"), ("vendor", "_update_internal_dataset")], - ) - def test_update_dataset_routes_by_provider(self, sqlite_session: Session, provider: str, helper_name: str) -> None: - dataset = _dataset(provider=provider) - sqlite_session.add(dataset) - sqlite_session.commit() - - with ( - patch.object(DatasetService, "check_dataset_permission"), - patch.object(DatasetService, helper_name, return_value=dataset) as update_helper, - ): - result = DatasetService.update_dataset( - dataset.id, - {"name": dataset.name}, - _account(), - session=sqlite_session, - ) - - assert result is dataset - update_helper.assert_called_once() - - @pytest.mark.parametrize( - ("payload", "message"), - [ - ({"external_knowledge_api_id": "api-1"}, "External knowledge id is required"), - ({"external_knowledge_id": "knowledge-1"}, "External knowledge api id is required"), - ], - ) - def test_update_external_dataset_requires_binding_fields( - self, unbound_session: Session, payload: dict[str, str], message: str - ) -> None: - with pytest.raises(ValueError, match=message): - DatasetService._update_external_dataset(_dataset(provider="external"), payload, _account(), unbound_session) - - def test_update_external_dataset_flushes_dataset_and_binding_without_committing( - self, sqlite_session: Session - ) -> None: - dataset = _dataset(provider="external") - api = _external_api() - binding = _binding() - sqlite_session.add_all([dataset, api, binding]) - sqlite_session.commit() - transaction_events: list[str] = [] - event.listen(sqlite_session, "after_commit", lambda _session: transaction_events.append("commit")) - - updated = DatasetService._update_external_dataset( - dataset, - { - "name": "Updated", - "description": "Changed", - "external_knowledge_id": "knowledge-2", - "external_knowledge_api_id": api.id, - }, - _account(), - sqlite_session, - ) - - assert updated is dataset - assert dataset.name == "Updated" - assert binding.external_knowledge_id == "knowledge-2" - assert transaction_events == [] - - def test_update_external_dataset_rejects_cross_tenant_api(self, sqlite_session: Session) -> None: - dataset = _dataset(provider="external") - foreign_api = _external_api(api_id="api-foreign", tenant_id="tenant-2") - sqlite_session.add_all([dataset, foreign_api, _binding()]) - sqlite_session.commit() - - with pytest.raises(ValueError, match="api template not found"): - DatasetService._update_external_dataset( - dataset, - { - "external_knowledge_id": "knowledge-2", - "external_knowledge_api_id": foreign_api.id, - }, - _account(), - sqlite_session, - ) - - def test_update_external_knowledge_binding_rejects_missing_row(self, sqlite_session: Session) -> None: - with pytest.raises(ValueError, match="binding not found"): - DatasetService._update_external_knowledge_binding( - "missing-dataset", - "knowledge-2", - "api-2", - sqlite_session, - ) - - def test_update_internal_dataset_executes_real_update_without_committing(self, sqlite_session: Session) -> None: - dataset = _dataset(name="Before") - sqlite_session.add(dataset) - sqlite_session.commit() - transaction_events: list[str] = [] - event.listen(sqlite_session, "after_commit", lambda _session: transaction_events.append("commit")) - - with ( - patch.object(DatasetService, "_handle_indexing_technique_change", return_value="update"), - patch.object(DatasetService, "_update_pipeline_knowledge_base_node_data"), - patch("services.knowledge.dataset_service.deal_dataset_vector_index_task.delay") as vector_task, - patch("services.knowledge.dataset_service.regenerate_summary_index_task.delay") as summary_task, - ): - updated = DatasetService._update_internal_dataset( - dataset, - {"name": "After", "description": "Changed"}, - _account(), - sqlite_session, - ) - - assert updated.name == "After" - assert updated.description == "Changed" - assert transaction_events == [] - vector_task.assert_called_once_with(dataset.id, "update") - summary_task.assert_called_once() - - def test_update_pipeline_node_data_returns_for_non_pipeline_or_missing_pipeline( - self, sqlite_session: Session - ) -> None: - ordinary = _dataset() - missing_pipeline = _dataset(dataset_id="pipeline-dataset", name="Pipeline Dataset") - missing_pipeline.runtime_mode = "rag_pipeline" - missing_pipeline.pipeline_id = "missing-pipeline" - sqlite_session.add_all([ordinary, missing_pipeline]) - sqlite_session.commit() - - DatasetService._update_pipeline_knowledge_base_node_data(ordinary, "user-1", sqlite_session) - DatasetService._update_pipeline_knowledge_base_node_data(missing_pipeline, "user-1", sqlite_session) - - def test_update_pipeline_node_data_rolls_back_real_session_on_failure(self, sqlite_session: Session) -> None: - pipeline = Pipeline(tenant_id="tenant-1", name="Pipeline", description="", created_by="user-1") - pipeline.id = "pipeline-1" - dataset = _dataset() - dataset.runtime_mode = "rag_pipeline" - dataset.pipeline_id = pipeline.id - sqlite_session.add_all([pipeline, dataset]) - sqlite_session.commit() - sqlite_session.begin() - transaction_events: list[str] = [] - event.listen( - sqlite_session, - "after_soft_rollback", - lambda _session, _previous_transaction: transaction_events.append("rollback"), - ) - - with patch("services.knowledge.dataset_service.RagPipelineService") as service_cls: - service_cls.return_value.get_published_workflow.side_effect = RuntimeError("boom") - with pytest.raises(RuntimeError, match="boom"): - DatasetService._update_pipeline_knowledge_base_node_data(dataset, "user-1", sqlite_session) - - assert transaction_events == ["rollback"] - - -class TestDatasetServiceEmbeddingSettings: - def test_handle_indexing_technique_change_returns_none_without_requested_change( - self, unbound_session: Session - ) -> None: - assert DatasetService._handle_indexing_technique_change(_dataset(), {}, {}, unbound_session) is None - - def test_handle_indexing_technique_change_switches_to_economy(self, unbound_session: Session) -> None: - dataset = _dataset(indexing_technique=IndexTechniqueType.HIGH_QUALITY) - filtered_data: dict[str, object] = {} - - result = DatasetService._handle_indexing_technique_change( - dataset, - {"indexing_technique": IndexTechniqueType.ECONOMY}, - filtered_data, - unbound_session, - ) - - assert result == "remove" - assert filtered_data == { - "embedding_model": None, - "embedding_model_provider": None, - "collection_binding_id": None, - } - - def test_handle_indexing_technique_change_delegates_high_quality_configuration( - self, unbound_session: Session - ) -> None: - dataset = _dataset(indexing_technique=IndexTechniqueType.ECONOMY) - filtered_data: dict[str, object] = {} - - with patch.object(DatasetService, "_configure_embedding_model_for_high_quality") as configure: - result = DatasetService._handle_indexing_technique_change( - dataset, - {"indexing_technique": IndexTechniqueType.HIGH_QUALITY}, - filtered_data, - unbound_session, - ) - - assert result == "add" - configure.assert_called_once() - - def test_handle_unchanged_indexing_preserves_existing_model(self, unbound_session: Session) -> None: - dataset = _dataset(indexing_technique=IndexTechniqueType.HIGH_QUALITY) - - with patch.object(DatasetService, "_preserve_existing_embedding_settings") as preserve: - result = DatasetService._handle_embedding_model_update_when_technique_unchanged( - dataset, - {}, - {}, - unbound_session, - ) - - assert result is None - preserve.assert_called_once() - - def test_update_embedding_model_settings_delegates_changed_model(self, unbound_session: Session) -> None: - dataset = _dataset() - - with patch.object(DatasetService, "_apply_new_embedding_settings") as apply_settings: - result = DatasetService._update_embedding_model_settings( - dataset, - {"embedding_model_provider": "provider-2", "embedding_model": "model-2"}, - {}, - unbound_session, - ) - - assert result == "update" - apply_settings.assert_called_once() - - def test_configure_high_quality_wraps_provider_error(self, unbound_session: Session) -> None: - account = _account() - with ( - patch("services.knowledge.dataset_service.current_user", account), - patch("services.knowledge.dataset_service.ModelManager") as model_manager_cls, - ): - model_manager_cls.for_tenant.return_value.get_model_instance.side_effect = LLMBadRequestError() - with pytest.raises(ValueError, match="No Embedding Model available"): - DatasetService._configure_embedding_model_for_high_quality( - {"embedding_model_provider": "provider", "embedding_model": "model"}, - {}, - unbound_session, - ) - - def test_preserve_existing_embedding_settings(self) -> None: - dataset = _dataset() - filtered_data: dict[str, object] = {} - - DatasetService._preserve_existing_embedding_settings(dataset, filtered_data) - - assert filtered_data["embedding_model_provider"] == dataset.embedding_model_provider - assert filtered_data["embedding_model"] == dataset.embedding_model - - def test_update_embedding_model_settings_returns_none_when_unchanged(self, unbound_session: Session) -> None: - dataset = _dataset() - - result = DatasetService._update_embedding_model_settings( - dataset, - { - "embedding_model_provider": dataset.embedding_model_provider, - "embedding_model": dataset.embedding_model, - }, - {}, - unbound_session, - ) - - assert result is None - - def test_apply_new_embedding_settings_uses_real_collection_binding(self, unbound_session: Session) -> None: - dataset = _dataset() - filtered_data: dict[str, object] = {} - account = _account() - embedding_model = SimpleNamespace(provider="new-provider", model_name="new-model") - collection_binding = DatasetCollectionBinding( - provider_name="new-provider", - model_name="new-model", - type="dataset", - collection_name="collection", - ) - collection_binding.id = "collection-binding-1" - - with ( - patch("services.knowledge.dataset_service.current_user", account), - patch("services.knowledge.dataset_service.ModelManager") as model_manager_cls, - patch.object( - DatasetCollectionBindingService, - "get_dataset_collection_binding", - return_value=collection_binding, - ), - ): - model_manager_cls.for_tenant.return_value.get_model_instance.return_value = embedding_model - DatasetService._apply_new_embedding_settings( - dataset, - {"embedding_model_provider": "new-provider", "embedding_model": "new-model"}, - filtered_data, - unbound_session, - ) - - assert filtered_data["embedding_model_provider"] == "new-provider" - assert filtered_data["embedding_model"] == "new-model" - assert filtered_data["collection_binding_id"] == collection_binding.id - - def test_apply_new_embedding_settings_preserves_existing_values_when_token_missing( - self, unbound_session: Session - ) -> None: - dataset = _dataset() - filtered_data: dict[str, object] = {} - - with ( - patch("services.knowledge.dataset_service.current_user", _account()), - patch("services.knowledge.dataset_service.ModelManager") as model_manager_cls, - ): - model_manager_cls.for_tenant.return_value.get_model_instance.side_effect = ProviderTokenNotInitError( - "missing" - ) - DatasetService._apply_new_embedding_settings( - dataset, - {"embedding_model_provider": "provider-2", "embedding_model": "model-2"}, - filtered_data, - unbound_session, - ) - - assert filtered_data["embedding_model_provider"] == dataset.embedding_model_provider - assert filtered_data["embedding_model"] == dataset.embedding_model - - @pytest.mark.parametrize( - ("summary_setting", "expected"), - [ - (None, False), - ({"enable": False}, False), - ({"enable": True, "model_name": "old-model", "model_provider_name": "provider"}, False), - ({"enable": True, "model_name": "new-model", "model_provider_name": "provider"}, True), - ], - ) - def test_check_summary_index_setting_model_changed( - self, summary_setting: dict[str, object] | None, expected: bool - ) -> None: - dataset = _dataset() - dataset.summary_index_setting = { - "enable": True, - "model_name": "old-model", - "model_provider_name": "provider", - } - - data = {} if summary_setting is None else {"summary_index_setting": summary_setting} - assert DatasetService._check_summary_index_setting_model_changed(dataset, data) is expected - - -class TestDatasetServiceRagPipelineSettings: - def test_requires_current_tenant(self, unbound_session: Session) -> None: - account = _account() - account._current_tenant = None - - with ( - patch("services.knowledge.dataset_service.current_user", account), - pytest.raises(ValueError, match="Current user or current tenant not found"), - ): - DatasetService.update_rag_pipeline_dataset_settings( - _dataset(), - _make_knowledge_configuration(), - session=unbound_session, - ) - - def test_unpublished_high_quality_settings_use_real_merged_dataset(self, sqlite_session: Session) -> None: - dataset = _dataset(indexing_technique=IndexTechniqueType.ECONOMY) - sqlite_session.add(dataset) - sqlite_session.commit() - account = _account() - embedding_model = SimpleNamespace(provider="provider-2", model_name="embedding-2") - collection_binding = DatasetCollectionBinding( - provider_name="provider-2", - model_name="embedding-2", - type="dataset", - collection_name="collection", - ) - collection_binding.id = "collection-binding-2" - - with ( - patch("services.knowledge.dataset_service.current_user", account), - patch("services.knowledge.dataset_service.ModelManager") as model_manager_cls, - patch.object(DatasetService, "check_is_multimodal_model", return_value=True), - patch.object( - DatasetCollectionBindingService, - "get_dataset_collection_binding", - return_value=collection_binding, - ), - ): - model_manager_cls.for_tenant.return_value.get_model_instance.return_value = embedding_model - DatasetService.update_rag_pipeline_dataset_settings( - dataset, - _make_knowledge_configuration( - embedding_model_provider="provider-2", - embedding_model="embedding-2", - summary_index_setting={"enable": True}, - ), - session=sqlite_session, - ) - - persisted = sqlite_session.get(Dataset, dataset.id) - assert persisted is not None - assert persisted.indexing_technique == IndexTechniqueType.HIGH_QUALITY - assert persisted.embedding_model == "embedding-2" - assert persisted.collection_binding_id == collection_binding.id - assert persisted.is_multimodal is True - - def test_unpublished_economy_settings_update_keyword_number(self, sqlite_session: Session) -> None: - dataset = _dataset(indexing_technique=IndexTechniqueType.HIGH_QUALITY) - sqlite_session.add(dataset) - sqlite_session.commit() - - with patch("services.knowledge.dataset_service.current_user", _account()): - DatasetService.update_rag_pipeline_dataset_settings( - dataset, - _make_knowledge_configuration( - indexing_technique=IndexTechniqueType.ECONOMY, - keyword_number=17, - ), - session=sqlite_session, - ) - - persisted = sqlite_session.get(Dataset, dataset.id) - assert persisted is not None - assert persisted.indexing_technique == IndexTechniqueType.ECONOMY - assert persisted.keyword_number == 17 - - def test_published_economy_settings_commit_keyword_change_and_dispatch_no_task( - self, sqlite_session: Session - ) -> None: - dataset = _dataset(indexing_technique=IndexTechniqueType.ECONOMY, chunk_structure="paragraph") - dataset.keyword_number = 4 - sqlite_session.add(dataset) - sqlite_session.commit() - - with ( - patch("services.knowledge.dataset_service.current_user", _account()), - patch("services.knowledge.dataset_service.deal_dataset_index_update_task.delay") as update_task, - ): - DatasetService.update_rag_pipeline_dataset_settings( - dataset, - _make_knowledge_configuration( - chunk_structure="paragraph", - indexing_technique=IndexTechniqueType.ECONOMY, - keyword_number=12, - ), - has_published=True, - session=sqlite_session, - ) - - assert sqlite_session.get(Dataset, dataset.id).keyword_number == 12 - update_task.assert_not_called() - - @pytest.mark.parametrize( - ("dataset", "configuration", "message"), - [ - ( - _dataset(chunk_structure="paragraph"), - _make_knowledge_configuration(chunk_structure="sentence"), - "Chunk structure is not allowed", - ), - ( - _dataset(indexing_technique=IndexTechniqueType.HIGH_QUALITY, chunk_structure="paragraph"), - _make_knowledge_configuration( - chunk_structure="paragraph", - indexing_technique=IndexTechniqueType.ECONOMY, - ), - "not allowed to be updated to economy", - ), - ], - ) - def test_published_settings_reject_incompatible_changes( - self, - sqlite_session: Session, - dataset: Dataset, - configuration, - message: str, - ) -> None: - sqlite_session.add(dataset) - sqlite_session.commit() - - with ( - patch("services.knowledge.dataset_service.current_user", _account()), - pytest.raises(ValueError, match=message), - ): - DatasetService.update_rag_pipeline_dataset_settings( - dataset, - configuration, - has_published=True, - session=sqlite_session, - ) - - -class TestDatasetPermissions: - def test_check_dataset_permission_enforces_tenant_and_partial_members(self, sqlite_session: Session) -> None: - dataset = _dataset(permission=DatasetPermissionEnum.PARTIAL_TEAM, maintainer="owner") - permitted_user = _account(account_id="permitted", role=TenantAccountRole.NORMAL) - denied_user = _account(account_id="denied", role=TenantAccountRole.NORMAL) - foreign_user = _account(account_id="foreign", tenant_id="tenant-2", role=TenantAccountRole.NORMAL) - sqlite_session.add_all( - [ - dataset, - DatasetPermission( - tenant_id=dataset.tenant_id, - dataset_id=dataset.id, - account_id=permitted_user.id, - ), - ] - ) - sqlite_session.commit() - - DatasetService.check_dataset_permission(dataset, permitted_user, sqlite_session) - with pytest.raises(NoPermissionError): - DatasetService.check_dataset_permission(dataset, denied_user, sqlite_session) - with pytest.raises(NoPermissionError): - DatasetService.check_dataset_permission(dataset, foreign_user, sqlite_session) - - def test_check_dataset_permission_skips_legacy_acl_in_rbac_mode( - self, config_overrides: Callable[..., None], sqlite_session: Session - ) -> None: - dataset = _dataset(permission=DatasetPermissionEnum.ONLY_ME, maintainer="owner") - non_member = _account(account_id="non-member", role=TenantAccountRole.NORMAL) - foreign_user = _account(account_id="foreign", tenant_id="tenant-2", role=TenantAccountRole.NORMAL) - sqlite_session.add(dataset) - sqlite_session.commit() - - config_overrides(RBAC_ENABLED=True) - - DatasetService.check_dataset_permission(dataset, non_member, sqlite_session) - with pytest.raises(NoPermissionError): - DatasetService.check_dataset_permission(dataset, foreign_user, sqlite_session) - - def test_dataset_operator_cannot_change_permission_or_member_list(self, sqlite_session: Session) -> None: - dataset = _dataset(permission=DatasetPermissionEnum.PARTIAL_TEAM) - operator = _account(role=TenantAccountRole.DATASET_OPERATOR) - sqlite_session.add_all( - [ - dataset, - DatasetPermission(tenant_id="tenant-1", dataset_id=dataset.id, account_id="member-1"), - ] - ) - sqlite_session.commit() - - with pytest.raises(NoPermissionError, match="cannot change"): - DatasetPermissionService.check_permission( - operator, - dataset, - DatasetPermissionEnum.ALL_TEAM, - None, - session=sqlite_session, - ) - - def test_non_editor_cannot_change_dataset_permissions(self, sqlite_session: Session) -> None: - user = _account(role=TenantAccountRole.NORMAL) - - with pytest.raises(NoPermissionError, match="does not have permission"): - DatasetPermissionService.check_permission( - user, - _dataset(), - DatasetPermissionEnum.ALL_TEAM, - None, - session=sqlite_session, - ) - - def test_dataset_operator_can_keep_unchanged_partial_member_list(self, sqlite_session: Session) -> None: - dataset = _dataset(permission=DatasetPermissionEnum.PARTIAL_TEAM) - operator = _account(role=TenantAccountRole.DATASET_OPERATOR) - sqlite_session.add_all( - [ - dataset, - DatasetPermission(tenant_id="tenant-1", dataset_id=dataset.id, account_id="member-1"), - ] - ) - sqlite_session.commit() - - DatasetPermissionService.check_permission( - operator, - dataset, - DatasetPermissionEnum.PARTIAL_TEAM, - [{"user_id": "member-1"}], - session=sqlite_session, - ) - with pytest.raises(ValueError, match="cannot change"): - DatasetPermissionService.check_permission( - operator, - dataset, - DatasetPermissionEnum.PARTIAL_TEAM, - [{"user_id": "member-2"}], - session=sqlite_session, - ) - - def test_update_partial_member_list_flush_failure_does_not_rollback_caller_session( - self, sqlite_session: Session - ) -> None: - transaction_events: list[str] = [] - - def fail_flush(_session, _flush_context, _instances) -> None: - raise RuntimeError("flush failed") - - event.listen(sqlite_session, "before_flush", fail_flush) - event.listen(sqlite_session, "after_rollback", lambda _session: transaction_events.append("rollback")) - try: - with pytest.raises(RuntimeError, match="flush failed"): - DatasetPermissionService.update_partial_member_list( - "tenant-1", - "dataset-1", - [{"user_id": "member-1"}], - sqlite_session, - ) - finally: - event.remove(sqlite_session, "before_flush", fail_flush) - - assert transaction_events == [] - - def test_clear_partial_member_list_execute_failure_does_not_rollback_caller_session( - self, sqlite_session: Session - ) -> None: - transaction_events: list[str] = [] - - def fail_execute(_orm_execute_state) -> None: - raise RuntimeError("execute failed") - - event.listen(sqlite_session, "do_orm_execute", fail_execute) - event.listen(sqlite_session, "after_rollback", lambda _session: transaction_events.append("rollback")) - try: - with pytest.raises(RuntimeError, match="execute failed"): - DatasetPermissionService.clear_partial_member_list("dataset-1", sqlite_session) - finally: - event.remove(sqlite_session, "do_orm_execute", fail_execute) - - assert transaction_events == [] - - -@pytest.mark.parametrize("rbac_enabled", [False, True]) -@pytest.mark.parametrize("commit_fails", [False, True]) -@pytest.mark.parametrize("expire_on_commit", [False, True]) -def test_empty_pipeline_syncs_creator_after_commit( - sqlite_session: Session, - sqlite_session_factory: sessionmaker[Session], - config_overrides: Callable[..., None], - rbac_enabled: bool, - commit_fails: bool, - expire_on_commit: bool, -) -> None: - config_overrides(RBAC_ENABLED=rbac_enabled) - sqlite_session.expire_on_commit = expire_on_commit - account = _account() - entity = RagPipelineDatasetCreateEntity( - name="New pipeline", - description="", - permission=DatasetPermissionEnum.ONLY_ME, - icon_info=IconInfo(icon_type="emoji", icon="📚", icon_background="#FFFFFF", icon_url=None), - ) - commit = sqlite_session.commit - observed: list[tuple[str, str, object]] = [] - transactions_at_sync: list[bool] = [] - - def commit_session() -> None: - if commit_fails: - raise RuntimeError("commit failed") - commit() - - def sync_binding(_tenant_id: str, _account_id: str, *, resource_type: object, resource_id: str) -> list[object]: - transactions_at_sync.append(sqlite_session.in_transaction()) - with sqlite_session_factory() as observer: - stored = observer.get(Dataset, resource_id) - if stored is not None: - observed.append((stored.tenant_id, stored.created_by, resource_type)) - return [] - - with ( - patch("services.knowledge.dataset_service.current_user", account), - patch.object(sqlite_session, "commit", side_effect=commit_session), - patch.object( - rbac_service.RBACService.AccessPolicies, - "sync_creator_access_policy_member_bindings", - side_effect=sync_binding, - ) as sync, - ): - if commit_fails: - with pytest.raises(RuntimeError, match="commit failed"): - DatasetService.create_empty_rag_pipeline_dataset("tenant-1", entity, sqlite_session) - else: - dataset = DatasetService.create_empty_rag_pipeline_dataset("tenant-1", entity, sqlite_session) - if rbac_enabled: - assert transactions_at_sync == [False] - assert observed == [("tenant-1", account.id, rbac_service.RBACResourceType.DATASET)] - sync.assert_called_once_with( - "tenant-1", - account.id, - resource_type=rbac_service.RBACResourceType.DATASET, - resource_id=dataset.id, - ) - if commit_fails or not rbac_enabled: - sync.assert_not_called() diff --git a/api/tests/unit_tests/services/test_dataset_service_document.py b/api/tests/unit_tests/services/test_dataset_service_document.py deleted file mode 100644 index 40bce008f4f..00000000000 --- a/api/tests/unit_tests/services/test_dataset_service_document.py +++ /dev/null @@ -1,2077 +0,0 @@ -"""Unit tests for DocumentService behaviors in dataset_service.""" - -from collections.abc import Callable -from datetime import datetime - -from sqlalchemy import event, select -from sqlalchemy.orm import Session - -from core.model_manager import ModelInstance -from enums import DeploymentEdition -from models.dataset import Dataset, DatasetCollectionBinding, DocumentSegment -from models.enums import DataSourceType, DocumentCreatedFrom, IndexingStatus -from models.model import UploadFile -from models.source import DataSourceOauthBinding -from services.knowledge.resource_scope import DatasetRef -from tests.unit_tests.config_override import config_overrides_context -from tests.unit_tests.model_factories import make_account, make_tenant, make_upload_file - -from .dataset_service_test_helpers import ( - Account, - BuiltInField, - CloudPlan, - DatasetProcessRule, - DatasetService, - DataSource, - Document, - DocumentIndexingError, - DocumentService, - FileInfo, - FileNotExistsError, - IndexStructureType, - InfoList, - KnowledgeConfig, - NotFound, - NotionIcon, - NotionInfo, - NotionPage, - PreProcessingRule, - ProcessRule, - RerankingModel, - RetrievalMethod, - RetrievalModel, - Rule, - Segmentation, - WebsiteInfo, - _make_features, - _make_lock_context, - _make_upload_knowledge_config, - json, - patch, - pytest, -) - - -def _account(*, account_id: str = "user-1", tenant_id: str = "tenant-1") -> Account: - return make_account( - account_id=account_id, - name="User", - email=f"{account_id}@example.com", - tenant=make_tenant(tenant_id=tenant_id, name="Tenant"), - ) - - -def _dataset_row( - *, - dataset_id: str = "dataset-1", - tenant_id: str = "tenant-1", - built_in_field_enabled: bool = False, - data_source_type: str | None = None, - indexing_technique: str | None = "economy", -) -> Dataset: - return Dataset( - id=dataset_id, - tenant_id=tenant_id, - name="Dataset", - description="", - provider="vendor", - created_by="user-1", - maintainer="user-1", - built_in_field_enabled=built_in_field_enabled, - chunk_structure=IndexStructureType.PARAGRAPH_INDEX, - data_source_type=data_source_type, - indexing_technique=indexing_technique, - ) - - -def _document_row( - *, - document_id: str = "document-1", - dataset_id: str = "dataset-1", - tenant_id: str = "tenant-1", - name: str = "Document", - indexing_status: str = IndexingStatus.COMPLETED, - data_source_type: str = DataSourceType.UPLOAD_FILE, - data_source_info: str = "{}", - enabled: bool = True, - archived: bool = False, - is_paused: bool = False, -) -> Document: - document = Document( - id=document_id, - tenant_id=tenant_id, - dataset_id=dataset_id, - position=1, - data_source_type=data_source_type, - data_source_info=data_source_info, - batch="batch-1", - name=name, - created_from=DocumentCreatedFrom.API, - created_by="user-1", - indexing_status=indexing_status, - doc_form=IndexStructureType.PARAGRAPH_INDEX, - word_count=10, - enabled=enabled, - archived=archived, - is_paused=is_paused, - completed_at=datetime(2026, 1, 2) if indexing_status == IndexingStatus.COMPLETED else None, - ) - document.created_at = datetime(2026, 1, 1) - document.updated_at = datetime(2026, 1, 2) - return document - - -def _upload_file(*, file_id: str, tenant_id: str = "tenant-1", name: str = "upload.txt") -> UploadFile: - return make_upload_file( - file_id=file_id, - tenant_id=tenant_id, - storage_type="opendal", - key=f"key-{file_id}", - name=name, - size=1, - created_by="user-1", - created_at=datetime(2026, 1, 1), - ) - - -def _process_rule( - *, dataset_id: str = "dataset-1", rule_id: str = "rule-1", mode: str = "automatic" -) -> DatasetProcessRule: - rules = ( - json.dumps(DatasetProcessRule.AUTOMATIC_RULES) - if mode == "automatic" - else Rule( - pre_processing_rules=[PreProcessingRule(id="remove_extra_spaces", enabled=True)], - segmentation=Segmentation(separator="\n", max_tokens=100), - ).model_dump_json() - ) - process_rule = DatasetProcessRule(dataset_id=dataset_id, mode=mode, rules=rules, created_by="user-1") - process_rule.id = rule_id - return process_rule - - -class _RetryFlagLock: - def __init__(self, store: "_RetryFlagStore", key: str): - self.store = store - self.key = key - self.token = f"owner-{store.next_token}" - store.next_token += 1 - - def acquire(self, *, blocking: bool): - assert blocking is False - if self.key in self.store.values: - if self.store.replacement_on_conflict: - replacement_key, replacement_value = self.store.replacement_on_conflict - self.store.values[replacement_key] = replacement_value - return False - self.store.values[self.key] = self.token - return True - - def release(self): - if self.store.values.get(self.key) == self.token: - self.store.values.pop(self.key) - - -class _RetryFlagStore: - def __init__( - self, - values: dict[str, str] | None = None, - replacement_on_conflict: tuple[str, str] | None = None, - ): - self.values = values or {} - self.replacement_on_conflict = replacement_on_conflict - self.next_token = 1 - - def lock(self, key: str, *, timeout: int, thread_local: bool): - assert timeout == 600 - assert thread_local is False - return _RetryFlagLock(self, key) - - -class TestDocumentServiceDisplayStatus: - """Unit tests for DocumentService display-status helpers.""" - - @pytest.mark.parametrize( - ("raw_status", "expected"), - [ - ("enabled", "available"), - ("AVAILABLE", "available"), - ("paused", "paused"), - ("unknown", None), - (None, None), - ], - ) - def test_normalize_display_status(self, raw_status, expected): - assert DocumentService.normalize_display_status(raw_status) == expected - - def test_build_display_status_filters_returns_empty_tuple_for_unknown_status(self): - assert DocumentService.build_display_status_filters("missing") == () - - def test_apply_display_status_filter_returns_original_query_for_unknown_status(self): - query = select(Document) - - result = DocumentService.apply_display_status_filter(query, "missing") - - assert result is query - - def test_apply_display_status_filter_applies_where_for_known_status(self): - query = select(Document) - - result = DocumentService.apply_display_status_filter(query, "enabled") - - assert result is not query - assert "documents.enabled" in str(result) - - -class TestDocumentServiceRetrieval: - def test_get_document_by_id_uses_provided_session(self, sqlite_session: Session): - document = _document_row() - sqlite_session.add(document) - sqlite_session.commit() - - assert DocumentService.get_document_by_id(document.id, session=sqlite_session) is document - assert DocumentService.get_document_by_id("missing", session=sqlite_session) is None - - def test_get_document_by_ids_enforces_dataset_owner_and_state(self, sqlite_session: Session): - dataset = _dataset_row() - expected = _document_row(document_id="expected") - sqlite_session.add_all( - [ - dataset, - expected, - _document_row(document_id="disabled", enabled=False), - _document_row(document_id="archived", archived=True), - _document_row(document_id="waiting", indexing_status=IndexingStatus.WAITING), - _document_row(document_id="other-dataset", dataset_id="dataset-2"), - _document_row(document_id="other-tenant", tenant_id="tenant-2"), - ] - ) - sqlite_session.commit() - - documents = DocumentService.get_document_by_ids( - DatasetRef(tenant_id=dataset.tenant_id, dataset_id=dataset.id), - ["expected", "disabled", "archived", "waiting", "other-dataset", "other-tenant"], - sqlite_session, - ) - - assert [document.id for document in documents] == [expected.id] - - -class TestDocumentServiceMutations: - """Unit tests for DocumentService mutation and orchestration helpers.""" - - @pytest.mark.parametrize(("archived", "expected"), [(True, True), (False, False)]) - def test_check_archived_returns_boolean_status(self, archived, expected): - document = _document_row(archived=archived) - - assert DocumentService.check_archived(document) is expected - - def test_delete_documents_limits_query_and_cleanup_to_dataset_ref(self, sqlite_session: Session): - dataset = _dataset_row() - document = _document_row( - document_id="doc-1", - data_source_info=json.dumps({"upload_file_id": "file-1"}), - ) - other_dataset = _document_row(document_id="other-dataset", dataset_id="dataset-2") - other_tenant = _document_row(document_id="other-tenant", tenant_id="tenant-2") - sqlite_session.add_all([dataset, document, other_dataset, other_tenant]) - sqlite_session.commit() - - with patch("services.knowledge.dataset_service.batch_clean_document_task") as clean_task: - dataset_ref = DatasetRef(tenant_id=dataset.tenant_id, dataset_id=dataset.id) - DocumentService.delete_documents( - dataset_ref, - [document.id, other_dataset.id, other_tenant.id], - IndexStructureType.PARAGRAPH_INDEX, - sqlite_session, - ) - - assert sqlite_session.get(Document, document.id) is None - assert sqlite_session.get(Document, other_dataset.id) is other_dataset - assert sqlite_session.get(Document, other_tenant.id) is other_tenant - clean_task.delay.assert_called_once_with( - [document.id], dataset.id, IndexStructureType.PARAGRAPH_INDEX, ["file-1"] - ) - - def test_delete_documents_with_empty_ids_does_not_commit(self, sqlite_session: Session): - commits = 0 - - def count_commit(_session): - nonlocal commits - commits += 1 - - event.listen(sqlite_session, "after_commit", count_commit) - dataset = _dataset_row() - DocumentService.delete_documents( - DatasetRef(tenant_id=dataset.tenant_id, dataset_id=dataset.id), - [], - IndexStructureType.PARAGRAPH_INDEX, - sqlite_session, - ) - event.remove(sqlite_session, "after_commit", count_commit) - assert commits == 0 - - def test_rename_document_raises_when_dataset_is_missing(self, sqlite_session: Session): - with patch("services.knowledge.dataset_service.current_user", _account()): - with pytest.raises(ValueError, match="Dataset not found"): - DocumentService.rename_document("dataset-1", "doc-1", "New Name", sqlite_session) - - def test_rename_document_raises_when_document_is_missing(self, sqlite_session: Session): - dataset = _dataset_row() - sqlite_session.add(dataset) - sqlite_session.commit() - with patch("services.knowledge.dataset_service.current_user", _account()): - with pytest.raises(ValueError, match="Document not found"): - DocumentService.rename_document(dataset.id, "doc-1", "New Name", sqlite_session) - - def test_rename_document_rejects_cross_tenant_access(self, sqlite_session: Session): - dataset = _dataset_row() - document = _document_row(tenant_id="tenant-other") - sqlite_session.add_all([dataset, document]) - sqlite_session.commit() - with patch("services.knowledge.dataset_service.current_user", _account()): - with pytest.raises(ValueError, match="No permission"): - DocumentService.rename_document(dataset.id, document.id, "New Name", sqlite_session) - - def test_rename_document_updates_document_metadata_and_upload_file_name(self, sqlite_session: Session): - dataset = _dataset_row(built_in_field_enabled=True) - document = _document_row( - data_source_info=json.dumps({"upload_file_id": "file-1"}), - ) - document.doc_metadata = {BuiltInField.document_name: "Old"} - upload_file = UploadFile( - tenant_id=dataset.tenant_id, - storage_type="opendal", - key="key", - name="old.txt", - size=1, - extension="txt", - mime_type="text/plain", - created_by_role="account", - created_by="user-1", - created_at=datetime(2026, 1, 1), - used=False, - ) - upload_file.id = "file-1" - sqlite_session.add_all([dataset, document, upload_file]) - sqlite_session.commit() - commits = 0 - - def count_commit(_session): - nonlocal commits - commits += 1 - - event.listen(sqlite_session, "after_commit", count_commit) - with patch("services.knowledge.dataset_service.current_user", _account()): - result = DocumentService.rename_document(dataset.id, document.id, "New Name", sqlite_session) - event.remove(sqlite_session, "after_commit", count_commit) - - assert result is document - assert document.name == "New Name" - assert document.doc_metadata[BuiltInField.document_name] == "New Name" - assert sqlite_session.get(UploadFile, upload_file.id).name == "New Name" - assert commits == 0 - - def test_recover_document_raises_when_document_is_not_paused(self, unbound_session: Session): - document = _document_row(is_paused=False) - with pytest.raises(DocumentIndexingError): - DocumentService.recover_document(document, unbound_session) - - def test_recover_document_persists_and_dispatches(self, sqlite_session: Session): - document = _document_row(is_paused=True) - sqlite_session.add(document) - sqlite_session.commit() - with ( - patch("services.knowledge.dataset_service.redis_client") as redis, - patch("services.knowledge.dataset_service.recover_document_indexing_task") as task, - ): - DocumentService.recover_document(document, sqlite_session) - - sqlite_session.expire_all() - recovered = sqlite_session.get(Document, document.id) - assert recovered is not None - assert recovered.is_paused is False - redis.delete.assert_called_once_with(f"document_{document.id}_is_paused") - task.delay.assert_called_once_with(document.dataset_id, document.id) - - def test_retry_document_raises_when_retry_flag_is_already_set(self, sqlite_session: Session): - document = _document_row(indexing_status=IndexingStatus.ERROR) - sqlite_session.add(document) - sqlite_session.commit() - retry_flags = _RetryFlagStore({f"document_{document.id}_is_retried": "other-request"}) - with ( - patch("services.knowledge.dataset_service.current_user", _account()), - patch("services.knowledge.dataset_service.redis_client", retry_flags), - ): - with pytest.raises(ValueError, match="being retried"): - DocumentService.retry_document("dataset-1", [document], sqlite_session, actor_id="user-1") - - def test_retry_document_leaves_batch_unchanged_when_later_document_is_already_being_retried( - self, sqlite_session: Session - ): - first_document = _document_row(document_id="doc-1", indexing_status=IndexingStatus.ERROR) - second_document = _document_row(document_id="doc-2", indexing_status=IndexingStatus.ERROR) - sqlite_session.add_all([first_document, second_document]) - sqlite_session.commit() - first_retry_key = "document_doc-1_is_retried" - second_retry_key = "document_doc-2_is_retried" - retry_flags = _RetryFlagStore({second_retry_key: "other-request"}) - with ( - patch("services.knowledge.dataset_service.current_user", _account()), - patch("services.knowledge.dataset_service.redis_client", retry_flags), - patch("services.knowledge.dataset_service.retry_document_indexing_task") as retry_task, - ): - with pytest.raises(ValueError, match="being retried"): - DocumentService.retry_document( - "dataset-1", [first_document, second_document], sqlite_session, actor_id="user-1" - ) - - assert first_document.indexing_status == IndexingStatus.ERROR - assert second_document.indexing_status == IndexingStatus.ERROR - assert first_retry_key not in retry_flags.values - assert retry_flags.values[second_retry_key] == "other-request" - retry_task.delay.assert_not_called() - - def test_retry_document_does_not_release_a_retry_flag_reacquired_by_another_request(self, sqlite_session: Session): - first_retry_key = "document_doc-1_is_retried" - second_retry_key = "document_doc-2_is_retried" - retry_flags = _RetryFlagStore( - {second_retry_key: "other-request"}, - replacement_on_conflict=(first_retry_key, "new-owner"), - ) - documents = [ - _document_row(document_id="doc-1", indexing_status=IndexingStatus.ERROR), - _document_row(document_id="doc-2", indexing_status=IndexingStatus.ERROR), - ] - sqlite_session.add_all(documents) - sqlite_session.commit() - - with ( - patch("services.knowledge.dataset_service.current_user", _account()), - patch("services.knowledge.dataset_service.redis_client", retry_flags), - ): - with pytest.raises(ValueError, match="being retried"): - DocumentService.retry_document("dataset-1", documents, sqlite_session, actor_id="user-1") - - assert retry_flags.values[first_retry_key] == "new-owner" - assert retry_flags.values[second_retry_key] == "other-request" - - def test_retry_document_releases_flags_when_status_commit_fails(self, sqlite_session: Session): - retry_flags = _RetryFlagStore() - document = _document_row(indexing_status=IndexingStatus.ERROR) - sqlite_session.add(document) - sqlite_session.commit() - - def fail_commit(_session): - raise RuntimeError("database unavailable") - - event.listen(sqlite_session, "before_commit", fail_commit) - with ( - patch("services.knowledge.dataset_service.current_user", _account()), - patch("services.knowledge.dataset_service.redis_client", retry_flags), - patch("services.knowledge.dataset_service.retry_document_indexing_task") as retry_task, - ): - with pytest.raises(RuntimeError, match="database unavailable"): - DocumentService.retry_document("dataset-1", [document], sqlite_session, actor_id="user-1") - event.remove(sqlite_session, "before_commit", fail_commit) - - assert retry_flags.values == {} - retry_task.delay.assert_not_called() - - def test_retry_document_persists_status_and_dispatches(self, sqlite_session: Session): - documents = [ - _document_row(document_id="doc-1", indexing_status=IndexingStatus.ERROR), - _document_row(document_id="doc-2", indexing_status=IndexingStatus.PAUSED), - ] - sqlite_session.add_all(documents) - sqlite_session.commit() - retry_flags = _RetryFlagStore() - with ( - patch("services.knowledge.dataset_service.current_user", _account()), - patch("services.knowledge.dataset_service.redis_client", retry_flags), - patch("services.knowledge.dataset_service.retry_document_indexing_task") as task, - ): - DocumentService.retry_document("dataset-1", documents, sqlite_session, actor_id="user-1") - - sqlite_session.expire_all() - statuses = sqlite_session.scalars(select(Document.indexing_status).order_by(Document.id)).all() - assert statuses == [IndexingStatus.WAITING, IndexingStatus.WAITING] - task.delay.assert_called_once_with("dataset-1", ["doc-1", "doc-2"], "user-1") - - def test_sync_website_document_raises_when_sync_flag_exists(self, sqlite_session: Session): - dataset = _dataset_row() - document = _document_row() - with patch("services.knowledge.dataset_service.redis_client") as mock_redis: - mock_redis.get.return_value = "1" - - with pytest.raises(ValueError, match="being synced"): - DocumentService.sync_website_document(dataset, document, sqlite_session) - - @pytest.mark.parametrize( - ("document_dataset_id", "document_tenant_id"), - [("dataset-2", "tenant-1"), ("dataset-1", "tenant-2")], - ) - def test_sync_website_document_rejects_document_outside_dataset( - self, document_dataset_id: str, document_tenant_id: str, sqlite_session: Session - ): - dataset = _dataset_row() - document = _document_row(dataset_id=document_dataset_id, tenant_id=document_tenant_id) - - with ( - pytest.raises(ValueError, match="Document not found"), - patch("services.knowledge.dataset_service.redis_client") as mock_redis, - ): - DocumentService.sync_website_document(dataset, document, sqlite_session) - - mock_redis.get.assert_not_called() - - def test_sync_website_document_updates_status_sets_cache_and_dispatches_task(self, sqlite_session: Session): - dataset = _dataset_row() - document = _document_row( - data_source_type=DataSourceType.WEBSITE_CRAWL, - data_source_info=json.dumps({"mode": "crawl"}), - ) - sqlite_session.add_all([dataset, document]) - sqlite_session.commit() - - with ( - patch("services.knowledge.dataset_service.redis_client") as mock_redis, - patch("services.knowledge.dataset_service.sync_website_document_indexing_task") as sync_task, - ): - mock_redis.get.return_value = None - - DocumentService.sync_website_document(dataset, document, sqlite_session) - - sqlite_session.expire_all() - synced = sqlite_session.get(Document, document.id) - assert synced is not None - assert synced.indexing_status == IndexingStatus.WAITING - assert synced.data_source_info_dict["mode"] == "scrape" - mock_redis.setex.assert_called_once_with(f"document_{document.id}_is_sync", 600, 1) - sync_task.delay.assert_called_once_with(dataset.id, document.id) - - -class TestDocumentServiceSaveDocumentWithoutDatasetId: - """Unit tests for dataset creation around save_document_without_dataset_id.""" - - @pytest.fixture - def account_context(self): - account = _account() - - with patch("services.knowledge.dataset_service.current_user", account): - yield account - - @config_overrides_context(DEPLOYMENT_EDITION=DeploymentEdition.COMMUNITY) - def test_save_document_without_dataset_id_creates_high_quality_dataset_with_default_retrieval_model( - self, account_context, sqlite_session: Session - ): - knowledge_config = KnowledgeConfig( - indexing_technique="high_quality", - data_source=DataSource( - info_list=InfoList( - data_source_type="upload_file", - file_info_list=FileInfo(file_ids=["file-1"]), - ) - ), - embedding_model="embedding-model", - embedding_model_provider="provider", - summary_index_setting={"enable": True}, - is_multimodal=True, - ) - binding = DatasetCollectionBinding( - provider_name="provider", - model_name="embedding-model", - type="dataset", - collection_name="collection", - ) - binding.id = "binding-1" - first_document = _document_row(name="VeryLongDocumentNameForDataset.txt") - - with ( - patch("services.knowledge.dataset_service.FeatureService.get_features", return_value=_make_features()), - patch( - "services.knowledge.dataset_service.DatasetCollectionBindingService.get_dataset_collection_binding", - return_value=binding, - ), - patch.object( - DocumentService, "save_document_with_dataset_id", return_value=([first_document], "batch-1") - ) as save_document, - ): - dataset, documents, batch = DocumentService.save_document_without_dataset_id( - tenant_id="tenant-1", - knowledge_config=knowledge_config, - account=account_context, - session=sqlite_session, - ) - - assert documents == [first_document] - assert batch == "batch-1" - assert dataset.collection_binding_id == "binding-1" - assert dataset.retrieval_model["search_method"] == RetrievalMethod.SEMANTIC_SEARCH - assert dataset.retrieval_model["top_k"] == 4 - assert dataset.summary_index_setting == {"enable": True} - assert dataset.is_multimodal is True - assert dataset.name == first_document.name[:18] + "..." - assert ( - dataset.description - == "useful for when you want to answer queries about the VeryLongDocumentNameForDataset.txt" - ) - assert sqlite_session.get(Dataset, dataset.id) is dataset - save_document.assert_called_once_with( - dataset, - knowledge_config, - account_context, - session=sqlite_session, - ) - - @config_overrides_context(DEPLOYMENT_EDITION=DeploymentEdition.COMMUNITY) - def test_save_document_without_dataset_id_uses_provided_retrieval_model( - self, account_context, sqlite_session: Session - ): - retrieval_model = RetrievalModel( - search_method=RetrievalMethod.SEMANTIC_SEARCH, - reranking_enable=True, - reranking_model=RerankingModel( - reranking_provider_name="rerank-provider", - reranking_model_name="rerank-model", - ), - top_k=9, - score_threshold_enabled=True, - score_threshold=0.6, - ) - knowledge_config = KnowledgeConfig( - indexing_technique="economy", - data_source=DataSource( - info_list=InfoList( - data_source_type="upload_file", - file_info_list=FileInfo(file_ids=["file-1"]), - ) - ), - retrieval_model=retrieval_model, - ) - first_document = _document_row(name="Doc") - - with ( - patch("services.knowledge.dataset_service.FeatureService.get_features", return_value=_make_features()), - patch.object( - DocumentService, - "save_document_with_dataset_id", - return_value=([first_document], "batch-1"), - ), - ): - dataset, _, _ = DocumentService.save_document_without_dataset_id( - "tenant-1", - knowledge_config, - account_context, - sqlite_session, - ) - - assert dataset.retrieval_model == retrieval_model.model_dump() - assert dataset.collection_binding_id is None - assert sqlite_session.get(Dataset, dataset.id) is dataset - - @config_overrides_context(DEPLOYMENT_EDITION=DeploymentEdition.CLOUD) - def test_save_document_without_dataset_id_rejects_sandbox_batch_upload( - self, account_context, unbound_session: Session - ): - knowledge_config = KnowledgeConfig( - indexing_technique="economy", - data_source=DataSource( - info_list=InfoList( - data_source_type="upload_file", - file_info_list=FileInfo(file_ids=["file-1", "file-2"]), - ) - ), - ) - - with ( - patch( - "services.knowledge.dataset_service.FeatureService.get_features", - return_value=_make_features(plan=CloudPlan.SANDBOX), - ), - patch.object(DocumentService, "check_documents_upload_quota") as check_quota, - ): - with pytest.raises(ValueError, match="does not support batch upload"): - DocumentService.save_document_without_dataset_id( - "tenant-1", knowledge_config, account_context, unbound_session - ) - - check_quota.assert_not_called() - - -class TestDocumentServiceUpdateDocumentWithDatasetId: - """Unit tests for the document-update orchestration path.""" - - @pytest.fixture - def account_context(self): - account = _account() - - with patch("services.knowledge.dataset_service.current_user", account): - yield account - - def test_update_document_with_dataset_id_raises_when_document_is_missing( - self, account_context, sqlite_session: Session - ): - dataset = _dataset_row() - document_data = KnowledgeConfig( - original_document_id="doc-1", - indexing_technique="economy", - data_source=DataSource( - info_list=InfoList( - data_source_type="upload_file", - file_info_list=FileInfo(file_ids=["file-1"]), - ) - ), - ) - with patch.object(DatasetService, "check_dataset_model_setting") as check_model_setting: - with pytest.raises(NotFound, match="Document not found"): - DocumentService.update_document_with_dataset_id( - dataset, - document_data, - account_context, - session=sqlite_session, - ) - - check_model_setting.assert_called_once_with(dataset) - - def test_update_document_with_dataset_id_rejects_non_available_documents( - self, account_context, sqlite_session: Session - ): - dataset = _dataset_row() - document = _document_row(document_id="doc-1", indexing_status=IndexingStatus.INDEXING) - sqlite_session.add_all([dataset, document]) - sqlite_session.commit() - document_data = KnowledgeConfig( - original_document_id="doc-1", - indexing_technique="economy", - data_source=DataSource( - info_list=InfoList( - data_source_type="upload_file", - file_info_list=FileInfo(file_ids=["file-1"]), - ) - ), - ) - with patch.object(DatasetService, "check_dataset_model_setting"): - with pytest.raises(ValueError, match="Document is not available"): - DocumentService.update_document_with_dataset_id( - dataset, - document_data, - account_context, - session=sqlite_session, - ) - - def test_update_document_with_dataset_id_upload_file_process_rule_and_name_override( - self, account_context, sqlite_session: Session - ): - dataset = _dataset_row() - document = _document_row(document_id="doc-1") - upload_file = UploadFile( - tenant_id=dataset.tenant_id, - storage_type="opendal", - key="key", - name="upload.txt", - size=1, - extension="txt", - mime_type="text/plain", - created_by_role="account", - created_by=account_context.id, - created_at=datetime(2026, 1, 1), - used=False, - ) - upload_file.id = "file-1" - segment = DocumentSegment( - tenant_id=dataset.tenant_id, - dataset_id=dataset.id, - document_id=document.id, - position=1, - content="content", - word_count=1, - tokens=1, - created_by=account_context.id, - ) - sqlite_session.add_all([dataset, document, upload_file, segment]) - sqlite_session.commit() - document_data = KnowledgeConfig( - original_document_id="doc-1", - indexing_technique="economy", - data_source=DataSource( - info_list=InfoList( - data_source_type="upload_file", - file_info_list=FileInfo(file_ids=["file-1"]), - ) - ), - process_rule=ProcessRule( - mode="custom", - rules=Rule( - pre_processing_rules=[PreProcessingRule(id="remove_stopwords", enabled=True)], - segmentation=Segmentation(separator="\n", max_tokens=128), - ), - ), - name="Renamed document", - doc_form=IndexStructureType.QA_INDEX, - ) - updated_at = datetime(2026, 2, 1) - - with ( - patch.object(DatasetService, "check_dataset_model_setting"), - patch("services.knowledge.dataset_service.naive_utc_now", return_value=updated_at), - patch("services.knowledge.dataset_service.document_indexing_update_task") as update_task, - ): - result = DocumentService.update_document_with_dataset_id( - dataset, - document_data, - account_context, - session=sqlite_session, - ) - - assert result is document - assert document.dataset_process_rule_id is not None - assert sqlite_session.get(DatasetProcessRule, document.dataset_process_rule_id) is not None - assert document.data_source_type == "upload_file" - assert document.data_source_info == '{"upload_file_id": "file-1"}' - assert document.name == "Renamed document" - assert document.indexing_status == "waiting" - assert document.completed_at is None - assert document.processing_started_at is None - assert document.parsing_completed_at is None - assert document.cleaning_completed_at is None - assert document.splitting_completed_at is None - assert document.updated_at == updated_at - assert document.created_from == "web" - assert document.doc_form == IndexStructureType.QA_INDEX - sqlite_session.expire_all() - persisted = sqlite_session.get(Document, document.id) - assert persisted is not None - assert persisted.name == "Renamed document" - assert sqlite_session.get(DocumentSegment, segment.id).status == "re_segment" - update_task.delay.assert_called_once_with(document.dataset_id, document.id) - - def test_update_document_with_dataset_id_notion_import_requires_binding( - self, account_context, sqlite_session: Session - ): - dataset = _dataset_row() - document = _document_row(document_id="doc-1") - sqlite_session.add_all([dataset, document]) - sqlite_session.commit() - document_data = KnowledgeConfig( - original_document_id="doc-1", - indexing_technique="economy", - data_source=DataSource( - info_list=InfoList( - data_source_type="notion_import", - notion_info_list=[ - NotionInfo( - credential_id="credential-1", - workspace_id="workspace-1", - pages=[NotionPage(page_id="page-1", page_name="Page 1", page_icon=None, type="page")], - ) - ], - ) - ), - ) - - with patch.object(DatasetService, "check_dataset_model_setting"): - with pytest.raises(ValueError, match="Data source binding not found"): - DocumentService.update_document_with_dataset_id( - dataset, - document_data, - account_context, - session=sqlite_session, - ) - - def test_update_document_with_dataset_id_website_crawl_updates_segments_and_dispatches_task( - self, account_context, sqlite_session: Session - ): - dataset = _dataset_row() - document = _document_row(document_id="doc-1") - segment = DocumentSegment( - tenant_id=dataset.tenant_id, - dataset_id=dataset.id, - document_id=document.id, - position=1, - content="content", - word_count=1, - tokens=1, - created_by=account_context.id, - ) - sqlite_session.add_all([dataset, document, segment]) - sqlite_session.commit() - document_data = KnowledgeConfig( - original_document_id="doc-1", - indexing_technique="economy", - data_source=DataSource( - info_list=InfoList( - data_source_type="website_crawl", - website_info_list=WebsiteInfo( - provider="firecrawl", - job_id="job-1", - urls=["https://example.com"], - only_main_content=False, - ), - ) - ), - doc_form=IndexStructureType.PARENT_CHILD_INDEX, - ) - - with ( - patch.object(DatasetService, "check_dataset_model_setting"), - patch("services.knowledge.dataset_service.naive_utc_now", return_value=datetime(2026, 2, 1)), - patch("services.knowledge.dataset_service.document_indexing_update_task") as update_task, - ): - result = DocumentService.update_document_with_dataset_id( - dataset, - document_data, - account_context, - session=sqlite_session, - ) - - assert result is document - assert document.data_source_type == "website_crawl" - assert document.data_source_info == ( - '{"url": "https://example.com", "provider": "firecrawl", "job_id": "job-1", ' - '"only_main_content": false, "mode": "crawl"}' - ) - assert document.name == "" - assert document.doc_form == IndexStructureType.PARENT_CHILD_INDEX - sqlite_session.expire_all() - assert sqlite_session.get(DocumentSegment, segment.id).status == "re_segment" - update_task.delay.assert_called_once_with("dataset-1", "doc-1") - - -class TestDocumentServiceCreateValidation: - """Unit tests for document creation validation helpers.""" - - def test_document_create_args_validate_requires_data_source_or_process_rule(self): - knowledge_config = KnowledgeConfig(indexing_technique="economy") - - with pytest.raises(ValueError, match="Data source or Process rule is required"): - DocumentService.document_create_args_validate(knowledge_config) - - def test_document_create_args_validate_delegates_to_sub_validators(self): - knowledge_config = KnowledgeConfig( - indexing_technique="economy", - data_source=DataSource( - info_list=InfoList( - data_source_type="upload_file", - file_info_list=FileInfo(file_ids=["file-1"]), - ) - ), - process_rule=ProcessRule(mode="automatic"), - ) - - with ( - patch.object(DocumentService, "data_source_args_validate") as validate_data_source, - patch.object(DocumentService, "process_rule_args_validate") as validate_process_rule, - ): - DocumentService.document_create_args_validate(knowledge_config) - - validate_data_source.assert_called_once_with(knowledge_config) - validate_process_rule.assert_called_once_with(knowledge_config) - - @pytest.mark.parametrize( - ("data_source_type", "message"), - [ - ("upload_file", "File source info is required"), - ("notion_import", "Notion source info is required"), - ("website_crawl", "Website source info is required"), - ], - ) - def test_data_source_args_validate_requires_source_specific_info(self, data_source_type, message): - info_list = InfoList.model_validate({"data_source_type": data_source_type}) - knowledge_config = KnowledgeConfig( - indexing_technique="economy", - data_source=DataSource(info_list=info_list), - ) - - with pytest.raises(ValueError, match=message): - DocumentService.data_source_args_validate(knowledge_config) - - def test_process_rule_args_validate_clears_rules_for_automatic_mode(self): - knowledge_config = KnowledgeConfig( - indexing_technique="economy", - data_source=DataSource( - info_list=InfoList( - data_source_type="upload_file", - file_info_list=FileInfo(file_ids=["file-1"]), - ) - ), - process_rule=ProcessRule( - mode="automatic", - rules=Rule( - pre_processing_rules=[PreProcessingRule(id="remove_stopwords", enabled=True)], - segmentation=Segmentation(separator="\n", max_tokens=128), - ), - ), - ) - - DocumentService.process_rule_args_validate(knowledge_config) - - assert knowledge_config.process_rule is not None - assert knowledge_config.process_rule.rules is None - - def test_process_rule_args_validate_deduplicates_rules_and_skips_max_tokens_for_full_doc_hierarchical(self): - knowledge_config = KnowledgeConfig( - indexing_technique="economy", - data_source=DataSource( - info_list=InfoList( - data_source_type="upload_file", - file_info_list=FileInfo(file_ids=["file-1"]), - ) - ), - process_rule=ProcessRule( - mode="hierarchical", - rules=Rule( - pre_processing_rules=[ - PreProcessingRule(id="remove_stopwords", enabled=True), - PreProcessingRule(id="remove_stopwords", enabled=False), - ], - segmentation=Segmentation(separator="\n", max_tokens=0), - parent_mode="full-doc", - ), - ), - ) - - DocumentService.process_rule_args_validate(knowledge_config) - - assert knowledge_config.process_rule is not None - assert knowledge_config.process_rule.rules is not None - assert len(knowledge_config.process_rule.rules.pre_processing_rules) == 1 - assert knowledge_config.process_rule.rules.pre_processing_rules[0].enabled is False - - def test_process_rule_args_validate_hierarchical_defaults_parent_mode_to_paragraph(self): - knowledge_config = KnowledgeConfig( - indexing_technique="economy", - data_source=DataSource( - info_list=InfoList( - data_source_type="upload_file", - file_info_list=FileInfo(file_ids=["file-1"]), - ) - ), - process_rule=ProcessRule( - mode="hierarchical", - rules=Rule( - pre_processing_rules=[ - PreProcessingRule(id="remove_extra_spaces", enabled=True), - ], - segmentation=Segmentation(separator="\n", max_tokens=1024), - subchunk_segmentation=Segmentation(separator="\n", max_tokens=512), - ), - ), - ) - - DocumentService.process_rule_args_validate(knowledge_config) - - assert knowledge_config.process_rule is not None - assert knowledge_config.process_rule.rules is not None - assert knowledge_config.process_rule.rules.parent_mode == "paragraph" - - -class TestDocumentServiceSaveDocumentWithDatasetId: - """Unit tests for non-SQL validation branches in save_document_with_dataset_id.""" - - @pytest.fixture - def account_context(self): - account = _account() - - with ( - patch("services.knowledge.dataset_service.current_user", account), - patch.object(DatasetService, "check_doc_form"), - ): - yield account - - @config_overrides_context(DEPLOYMENT_EDITION=DeploymentEdition.CLOUD) - def test_save_document_with_dataset_id_requires_file_info_for_upload_source( - self, account_context, unbound_session: Session - ): - dataset = _dataset_row() - knowledge_config = _make_upload_knowledge_config(file_ids=None) - - with patch("services.knowledge.dataset_service.FeatureService.get_features", return_value=_make_features()): - with pytest.raises(ValueError, match="File source info is required"): - DocumentService.save_document_with_dataset_id( - dataset, - knowledge_config, - account_context, - session=unbound_session, - ) - - @config_overrides_context(DEPLOYMENT_EDITION=DeploymentEdition.CLOUD) - def test_save_document_with_dataset_id_blocks_batch_upload_for_sandbox_plan( - self, account_context, unbound_session: Session - ): - dataset = _dataset_row() - knowledge_config = _make_upload_knowledge_config(file_ids=["file-1", "file-2"]) - - with ( - patch( - "services.knowledge.dataset_service.FeatureService.get_features", - return_value=_make_features(plan=CloudPlan.SANDBOX), - ), - patch.object(DocumentService, "check_documents_upload_quota") as check_quota, - ): - with pytest.raises(ValueError, match="does not support batch upload"): - DocumentService.save_document_with_dataset_id( - dataset, - knowledge_config, - account_context, - session=unbound_session, - ) - - check_quota.assert_not_called() - - @config_overrides_context(DEPLOYMENT_EDITION=DeploymentEdition.CLOUD) - def test_save_document_with_dataset_id_enforces_batch_upload_limit( - self, - account_context, - unbound_session: Session, - config_overrides: Callable[..., None], - ): - config_overrides(BATCH_UPLOAD_LIMIT=1) - dataset = _dataset_row() - knowledge_config = _make_upload_knowledge_config(file_ids=["file-1", "file-2"]) - - with ( - patch("services.knowledge.dataset_service.FeatureService.get_features", return_value=_make_features()), - patch.object(DocumentService, "check_documents_upload_quota") as check_quota, - ): - with pytest.raises(ValueError, match="batch upload limit of 1"): - DocumentService.save_document_with_dataset_id( - dataset, - knowledge_config, - account_context, - session=unbound_session, - ) - - check_quota.assert_not_called() - - @config_overrides_context(DEPLOYMENT_EDITION=DeploymentEdition.COMMUNITY) - def test_save_document_with_dataset_id_updates_existing_document_and_data_source_type( - self, account_context, sqlite_session: Session - ): - dataset = _dataset_row(data_source_type=None) - knowledge_config = _make_upload_knowledge_config(original_document_id="doc-1", file_ids=["file-1"]) - updated_document = _document_row(document_id="doc-1") - updated_document.batch = "batch-existing" - - with ( - patch("services.knowledge.dataset_service.FeatureService.get_features", return_value=_make_features()), - patch("services.knowledge.dataset_service.document_indexing_update_task.delay"), - patch.object(DocumentService, "update_prepared_document", return_value=updated_document) as update_document, - ): - documents, batch = DocumentService.save_document_with_dataset_id( - dataset, - knowledge_config, - account_context, - session=sqlite_session, - ) - - assert dataset.data_source_type == "upload_file" - assert documents == [updated_document] - assert batch == "batch-existing" - update_document.assert_called_once_with(dataset, knowledge_config, account_context, session=sqlite_session) - - @config_overrides_context(DEPLOYMENT_EDITION=DeploymentEdition.COMMUNITY) - def test_save_document_with_dataset_id_requires_data_source_for_new_documents( - self, account_context, unbound_session: Session - ): - dataset = _dataset_row() - knowledge_config = _make_upload_knowledge_config(data_source=None) - - with patch("services.knowledge.dataset_service.FeatureService.get_features", return_value=_make_features()): - with pytest.raises(ValueError, match="Data source is required when creating new documents"): - DocumentService.save_document_with_dataset_id( - dataset, - knowledge_config, - account_context, - session=unbound_session, - ) - - @config_overrides_context(DEPLOYMENT_EDITION=DeploymentEdition.COMMUNITY) - def test_save_document_with_dataset_id_requires_existing_process_rule_for_custom_mode( - self, account_context, sqlite_session: Session - ): - dataset = _dataset_row() - sqlite_session.add(dataset) - sqlite_session.commit() - knowledge_config = _make_upload_knowledge_config( - file_ids=["file-1"], - process_rule=ProcessRule(mode="custom"), - ) - - with patch("services.knowledge.dataset_service.FeatureService.get_features", return_value=_make_features()): - with pytest.raises(ValueError, match="No process rule found"): - DocumentService.save_document_with_dataset_id( - dataset, - knowledge_config, - account_context, - session=sqlite_session, - ) - - @config_overrides_context(DEPLOYMENT_EDITION=DeploymentEdition.COMMUNITY) - def test_save_document_with_dataset_id_upload_file_creates_and_reindexes_documents( - self, account_context, sqlite_session: Session - ): - dataset = _dataset_row(data_source_type=DataSourceType.UPLOAD_FILE) - dataset_process_rule = _process_rule() - knowledge_config = _make_upload_knowledge_config(file_ids=["file-1", "file-2"]) - duplicate_document = _document_row(document_id="doc-duplicate", name="existing.txt") - upload_file_a = _upload_file(file_id="file-1", name="existing.txt") - upload_file_b = _upload_file(file_id="file-2", name="new.txt") - sqlite_session.add_all([dataset, dataset_process_rule, duplicate_document, upload_file_a, upload_file_b]) - sqlite_session.commit() - - with ( - patch("services.knowledge.dataset_service.FeatureService.get_features", return_value=_make_features()), - patch("services.knowledge.dataset_service.redis_client") as mock_redis, - patch("services.knowledge.dataset_service.DocumentIndexingTaskProxy") as document_proxy_cls, - patch("services.knowledge.dataset_service.DuplicateDocumentIndexingTaskProxy") as duplicate_proxy_cls, - patch("services.knowledge.dataset_service.naive_utc_now", return_value=datetime(2026, 2, 1)), - patch("services.knowledge.dataset_service.time.strftime", return_value="20260101010101"), - patch("services.knowledge.dataset_service.secrets.randbelow", return_value=23), - ): - mock_redis.lock.return_value = _make_lock_context() - documents, batch = DocumentService.save_document_with_dataset_id( - dataset, - knowledge_config, - account_context, - dataset_process_rule=dataset_process_rule, - session=sqlite_session, - ) - - assert [document.name for document in documents] == ["existing.txt", "new.txt"] - assert batch == "20260101010101100023" - assert duplicate_document.dataset_process_rule_id == "rule-1" - assert duplicate_document.updated_at == datetime(2026, 2, 1) - assert duplicate_document.batch == batch - assert duplicate_document.indexing_status == IndexingStatus.WAITING - created_document = next(document for document in documents if document.name == "new.txt") - sqlite_session.expire_all() - assert sqlite_session.get(Document, created_document.id) is not None - document_proxy_cls.assert_called_once_with(dataset.tenant_id, dataset.id, [created_document.id]) - document_proxy_cls.return_value.delay.assert_called_once() - duplicate_proxy_cls.assert_called_once_with(dataset.tenant_id, dataset.id, ["doc-duplicate"]) - duplicate_proxy_cls.return_value.delay.assert_called_once() - - @config_overrides_context(DEPLOYMENT_EDITION=DeploymentEdition.COMMUNITY) - def test_save_document_with_dataset_id_notion_import_truncates_names_and_cleans_removed_pages( - self, account_context, sqlite_session: Session - ): - dataset = _dataset_row(data_source_type=DataSourceType.NOTION_IMPORT) - dataset_process_rule = _process_rule() - notion_page_name = "a" * 300 - knowledge_config = KnowledgeConfig( - indexing_technique="economy", - data_source=DataSource( - info_list=InfoList( - data_source_type="notion_import", - notion_info_list=[ - NotionInfo( - credential_id="credential-1", - workspace_id="workspace-1", - pages=[ - NotionPage(page_id="page-keep", page_name="Keep page", type="page"), - NotionPage( - page_id="page-new", - page_name=notion_page_name, - page_icon=NotionIcon(type="emoji", emoji="page"), - type="page", - ), - ], - ) - ], - ) - ), - doc_form=IndexStructureType.PARAGRAPH_INDEX, - doc_language="English", - ) - existing_keep = _document_row( - document_id="doc-keep", - data_source_type=DataSourceType.NOTION_IMPORT, - ) - existing_keep.data_source_info = json.dumps({"notion_page_id": "page-keep"}) - existing_remove = _document_row( - document_id="doc-remove", - data_source_type=DataSourceType.NOTION_IMPORT, - ) - existing_remove.data_source_info = json.dumps({"notion_page_id": "page-remove"}) - sqlite_session.add_all([dataset, dataset_process_rule, existing_keep, existing_remove]) - sqlite_session.commit() - - with ( - patch("services.knowledge.dataset_service.FeatureService.get_features", return_value=_make_features()), - patch("services.knowledge.dataset_service.redis_client") as mock_redis, - patch("services.knowledge.dataset_service.clean_notion_document_task") as clean_task, - patch("services.knowledge.dataset_service.DocumentIndexingTaskProxy") as document_proxy_cls, - patch("services.knowledge.dataset_service.uuid.uuid4", return_value="doc-new"), - ): - mock_redis.lock.return_value = _make_lock_context() - documents, _ = DocumentService.save_document_with_dataset_id( - dataset, - knowledge_config, - account_context, - dataset_process_rule=dataset_process_rule, - session=sqlite_session, - ) - - created_document = next(document for document in documents if document.id == "doc-new") - assert created_document in documents - assert len(created_document.name) == 255 - assert sqlite_session.get(Document, created_document.id) is created_document - clean_task.delay.assert_called_once_with(["doc-remove"], dataset.id) - document_proxy_cls.assert_called_once_with(dataset.tenant_id, dataset.id, ["doc-new"]) - document_proxy_cls.return_value.delay.assert_called_once() - - @config_overrides_context(DEPLOYMENT_EDITION=DeploymentEdition.COMMUNITY) - def test_save_document_with_dataset_id_website_crawl_truncates_long_urls( - self, account_context, sqlite_session: Session - ): - dataset = _dataset_row(data_source_type=DataSourceType.WEBSITE_CRAWL) - dataset_process_rule = _process_rule() - sqlite_session.add_all([dataset, dataset_process_rule]) - sqlite_session.commit() - long_url = "https://example.com/" + ("a" * 260) - short_url = "https://example.com/short" - knowledge_config = KnowledgeConfig( - indexing_technique="economy", - data_source=DataSource( - info_list=InfoList( - data_source_type="website_crawl", - website_info_list=WebsiteInfo( - provider="firecrawl", - job_id="job-1", - urls=[long_url, short_url], - only_main_content=True, - ), - ) - ), - doc_form=IndexStructureType.PARAGRAPH_INDEX, - doc_language="English", - ) - with ( - patch("services.knowledge.dataset_service.FeatureService.get_features", return_value=_make_features()), - patch("services.knowledge.dataset_service.redis_client") as mock_redis, - patch("services.knowledge.dataset_service.DocumentIndexingTaskProxy") as document_proxy_cls, - ): - mock_redis.lock.return_value = _make_lock_context() - - documents, _ = DocumentService.save_document_with_dataset_id( - dataset, - knowledge_config, - account_context, - dataset_process_rule=dataset_process_rule, - session=sqlite_session, - ) - - assert [document.name for document in documents] == [long_url[:200] + "...", short_url] - assert sqlite_session.scalars(select(Document).order_by(Document.position)).all() == documents - document_proxy_cls.assert_called_once_with( - dataset.tenant_id, dataset.id, [document.id for document in documents] - ) - document_proxy_cls.return_value.delay.assert_called_once() - - -class TestDocumentServiceBatchUpdateStatus: - """Unit tests for batch_update_document_status orchestration and helper branches.""" - - def test_prepare_disable_update_requires_completed_document(self): - document = _document_row(indexing_status=IndexingStatus.WAITING) - document.completed_at = None - - with pytest.raises(DocumentIndexingError, match="is not completed"): - DocumentService._prepare_disable_update(document, user=_account(), now=datetime(2026, 2, 1)) - - def test_prepare_archive_update_sets_async_task_for_enabled_document(self): - document = _document_row(enabled=True, archived=False) - - result = DocumentService._prepare_archive_update(document, user=_account(), now=datetime(2026, 2, 1)) - - assert result is not None - assert result["updates"]["archived"] is True - assert result["set_cache"] is True - assert result["async_task"]["args"] == [document.id] - - def test_prepare_unarchive_update_sets_async_task_for_enabled_document(self): - document = _document_row(enabled=True, archived=True) - - result = DocumentService._prepare_unarchive_update(document, now=datetime(2026, 2, 1)) - - assert result is not None - assert result["updates"]["archived"] is False - assert result["set_cache"] is True - assert result["async_task"]["args"] == [document.id] - - def test_batch_update_document_status_rejects_indexing_documents(self, sqlite_session: Session): - dataset = _dataset_row() - document = _document_row(name="Busy document") - sqlite_session.add_all([dataset, document]) - sqlite_session.commit() - - with patch("services.knowledge.dataset_service.redis_client") as mock_redis: - mock_redis.get.return_value = "1" - - with pytest.raises(DocumentIndexingError, match="Busy document is being indexed"): - DocumentService.batch_update_document_status( - dataset, [document.id], "archive", _account(), sqlite_session - ) - - sqlite_session.refresh(document) - assert document.archived is False - - def test_batch_update_document_status_rolls_back_when_commit_fails(self, sqlite_session: Session): - dataset = _dataset_row() - document = _document_row(enabled=False) - sqlite_session.add_all([dataset, document]) - sqlite_session.commit() - - def fail_commit(_session): - raise RuntimeError("commit failed") - - event.listen(sqlite_session, "before_commit", fail_commit) - with ( - patch("services.knowledge.dataset_service.redis_client") as mock_redis, - ): - mock_redis.get.return_value = None - - with pytest.raises(RuntimeError, match="commit failed"): - DocumentService.batch_update_document_status( - dataset, [document.id], "enable", _account(), sqlite_session - ) - event.remove(sqlite_session, "before_commit", fail_commit) - - sqlite_session.refresh(document) - assert document.enabled is False - - def test_batch_update_document_status_raises_async_task_error_after_commit(self, sqlite_session: Session): - dataset = _dataset_row() - document = _document_row(enabled=False) - sqlite_session.add_all([dataset, document]) - sqlite_session.commit() - - with ( - patch("services.knowledge.dataset_service.redis_client") as mock_redis, - patch("services.knowledge.dataset_service.add_document_to_index_task") as add_task, - ): - mock_redis.get.return_value = None - add_task.delay.side_effect = RuntimeError("task failed") - - with pytest.raises(RuntimeError, match="task failed"): - DocumentService.batch_update_document_status( - dataset, [document.id], "enable", _account(), sqlite_session - ) - - sqlite_session.refresh(document) - assert document.enabled is True - mock_redis.setex.assert_called_once_with(f"document_{document.id}_indexing", 600, 1) - - -class TestDocumentServiceTenantAndUpdateEdges: - """Unit tests for tenant-count and update edge cases.""" - - @pytest.fixture - def account_context(self): - account = _account() - - with patch("services.knowledge.dataset_service.current_user", account): - yield account - - def test_get_tenant_documents_count_scopes_state_and_tenant(self, account_context, sqlite_session: Session): - sqlite_session.add_all( - [ - _document_row(document_id="one"), - _document_row(document_id="two"), - _document_row(document_id="disabled", enabled=False), - _document_row(document_id="archived", archived=True), - _document_row(document_id="unfinished", indexing_status=IndexingStatus.WAITING), - _document_row(document_id="foreign", tenant_id="tenant-2"), - ] - ) - sqlite_session.commit() - - result = DocumentService.get_tenant_documents_count(sqlite_session) - - assert result == 2 - - def test_update_document_with_dataset_id_uses_automatic_process_rule_payload( - self, account_context, sqlite_session: Session - ): - dataset = _dataset_row() - document = _document_row(document_id="doc-1") - upload_file = _upload_file(file_id="file-1") - sqlite_session.add_all([dataset, document, upload_file]) - sqlite_session.commit() - document_data = KnowledgeConfig( - original_document_id="doc-1", - indexing_technique="economy", - data_source=DataSource( - info_list=InfoList( - data_source_type="upload_file", - file_info_list=FileInfo(file_ids=["file-1"]), - ) - ), - process_rule=ProcessRule( - mode="automatic", - rules=Rule( - pre_processing_rules=[PreProcessingRule(id="remove_stopwords", enabled=True)], - segmentation=Segmentation(separator="\n", max_tokens=128), - ), - ), - doc_form=IndexStructureType.PARAGRAPH_INDEX, - ) - updated_at = datetime(2026, 2, 1) - - with ( - patch.object(DatasetService, "check_dataset_model_setting"), - patch("services.knowledge.dataset_service.naive_utc_now", return_value=updated_at), - patch("services.knowledge.dataset_service.document_indexing_update_task") as update_task, - ): - result = DocumentService.update_document_with_dataset_id( - dataset, - document_data, - account_context, - session=sqlite_session, - ) - - assert result is document - assert document.dataset_process_rule_id is not None - assert document.name == "upload.txt" - process_rule = sqlite_session.get(DatasetProcessRule, document.dataset_process_rule_id) - assert process_rule is not None - assert process_rule.mode == "automatic" - assert process_rule.rules == json.dumps(DatasetProcessRule.AUTOMATIC_RULES) - update_task.delay.assert_called_once_with("dataset-1", "doc-1") - - def test_update_document_with_dataset_id_requires_upload_file_info(self, account_context, sqlite_session: Session): - dataset = _dataset_row() - document = _document_row(document_id="doc-1") - sqlite_session.add_all([dataset, document]) - sqlite_session.commit() - document_data = KnowledgeConfig( - original_document_id="doc-1", - indexing_technique="economy", - data_source=DataSource(info_list=InfoList(data_source_type="upload_file")), - ) - - with patch.object(DatasetService, "check_dataset_model_setting"): - with pytest.raises(ValueError, match="No file info list found"): - DocumentService.update_document_with_dataset_id( - dataset, - document_data, - account_context, - session=sqlite_session, - ) - - def test_update_document_with_dataset_id_raises_when_upload_file_is_missing( - self, account_context, sqlite_session: Session - ): - dataset = _dataset_row() - document = _document_row(document_id="doc-1") - sqlite_session.add_all([dataset, document]) - sqlite_session.commit() - document_data = KnowledgeConfig( - original_document_id="doc-1", - indexing_technique="economy", - data_source=DataSource( - info_list=InfoList( - data_source_type="upload_file", - file_info_list=FileInfo(file_ids=["file-1"]), - ) - ), - ) - - with patch.object(DatasetService, "check_dataset_model_setting"): - with pytest.raises(FileNotExistsError): - DocumentService.update_document_with_dataset_id( - dataset, - document_data, - account_context, - session=sqlite_session, - ) - - def test_update_document_with_dataset_id_requires_notion_info_list(self, account_context, sqlite_session: Session): - dataset = _dataset_row() - document = _document_row(document_id="doc-1") - sqlite_session.add_all([dataset, document]) - sqlite_session.commit() - document_data = KnowledgeConfig( - original_document_id="doc-1", - indexing_technique="economy", - data_source=DataSource(info_list=InfoList(data_source_type="notion_import")), - ) - - with patch.object(DatasetService, "check_dataset_model_setting"): - with pytest.raises(ValueError, match="No notion info list found"): - DocumentService.update_document_with_dataset_id( - dataset, - document_data, - account_context, - session=sqlite_session, - ) - - def test_update_document_with_dataset_id_notion_import_updates_page_info( - self, account_context, sqlite_session: Session - ): - dataset = _dataset_row() - document = _document_row(document_id="doc-1") - binding = DataSourceOauthBinding( - tenant_id=dataset.tenant_id, - access_token="token", - provider="notion", - source_info={"workspace_id": '"workspace-1"'}, - disabled=False, - ) - sqlite_session.add_all([dataset, document, binding]) - sqlite_session.commit() - document_data = KnowledgeConfig( - original_document_id="doc-1", - indexing_technique="economy", - data_source=DataSource( - info_list=InfoList( - data_source_type="notion_import", - notion_info_list=[ - NotionInfo( - credential_id="credential-1", - workspace_id="workspace-1", - pages=[ - NotionPage(page_id="page-1", page_name="Page 1", page_icon=None, type="page"), - NotionPage(page_id="page-2", page_name="Page 2", page_icon=None, type="database"), - ], - ) - ], - ) - ), - doc_form=IndexStructureType.PARAGRAPH_INDEX, - ) - - with ( - patch.object(DatasetService, "check_dataset_model_setting"), - patch("services.knowledge.dataset_service.naive_utc_now", return_value=datetime(2026, 2, 1)), - patch("services.knowledge.dataset_service.document_indexing_update_task") as update_task, - ): - result = DocumentService.update_document_with_dataset_id( - dataset, - document_data, - account_context, - session=sqlite_session, - ) - - assert result is document - assert document.data_source_type == "notion_import" - assert document.name == "" - assert document.data_source_info == json.dumps( - { - "credential_id": "credential-1", - "notion_workspace_id": "workspace-1", - "notion_page_id": "page-2", - "notion_page_icon": None, - "type": "database", - } - ) - sqlite_session.refresh(document) - update_task.delay.assert_called_once_with("dataset-1", "doc-1") - - -class TestDocumentServiceSaveWithoutDatasetBilling: - """Unit tests for batch-count and quota branches in save_document_without_dataset_id.""" - - @pytest.fixture - def account_context(self): - account = _account() - - with patch("services.knowledge.dataset_service.current_user", account): - yield account - - @config_overrides_context(DEPLOYMENT_EDITION=DeploymentEdition.CLOUD) - def test_save_document_without_dataset_id_counts_notion_pages_for_quota( - self, - account_context, - sqlite_session: Session, - config_overrides: Callable[..., None], - ): - config_overrides(BATCH_UPLOAD_LIMIT="10") - knowledge_config = KnowledgeConfig( - indexing_technique="economy", - data_source=DataSource( - info_list=InfoList( - data_source_type="notion_import", - notion_info_list=[ - NotionInfo( - credential_id="credential-1", - workspace_id="workspace-1", - pages=[ - NotionPage(page_id="page-1", page_name="Page 1", page_icon=None, type="page"), - NotionPage(page_id="page-2", page_name="Page 2", page_icon=None, type="page"), - ], - ), - NotionInfo( - credential_id="credential-2", - workspace_id="workspace-2", - pages=[NotionPage(page_id="page-3", page_name="Page 3", page_icon=None, type="page")], - ), - ], - ) - ), - ) - features = _make_features() - document = _document_row(name="Doc") - - with ( - patch("services.knowledge.dataset_service.FeatureService.get_features", return_value=features), - patch.object(DocumentService, "check_documents_upload_quota") as check_quota, - patch.object( - DocumentService, - "save_document_with_dataset_id", - return_value=([document], "batch-1"), - ), - ): - dataset, _, _ = DocumentService.save_document_without_dataset_id( - "tenant-1", - knowledge_config, - account_context, - sqlite_session, - ) - - check_quota.assert_called_once_with(3, features) - assert sqlite_session.get(Dataset, dataset.id) is dataset - - @config_overrides_context(DEPLOYMENT_EDITION=DeploymentEdition.CLOUD) - def test_save_document_without_dataset_id_enforces_batch_limit_for_website_urls( - self, - account_context, - unbound_session: Session, - config_overrides: Callable[..., None], - ): - config_overrides(BATCH_UPLOAD_LIMIT="1") - knowledge_config = KnowledgeConfig( - indexing_technique="economy", - data_source=DataSource( - info_list=InfoList( - data_source_type="website_crawl", - website_info_list=WebsiteInfo( - provider="firecrawl", - job_id="job-1", - urls=["https://example.com/a", "https://example.com/b"], - only_main_content=True, - ), - ) - ), - ) - - with ( - patch("services.knowledge.dataset_service.FeatureService.get_features", return_value=_make_features()), - patch.object(DocumentService, "check_documents_upload_quota") as check_quota, - ): - with pytest.raises(ValueError, match="batch upload limit of 1"): - DocumentService.save_document_without_dataset_id( - "tenant-1", knowledge_config, account_context, unbound_session - ) - - check_quota.assert_not_called() - - -class TestDocumentServiceSaveDocumentAdditionalBranches: - """Additional unit tests for dataset bootstrap and process-rule branches.""" - - @pytest.fixture - def account_context(self): - account = _account() - - with ( - patch("services.knowledge.dataset_service.current_user", account), - patch.object(DatasetService, "check_doc_form"), - ): - yield account - - @config_overrides_context(DEPLOYMENT_EDITION=DeploymentEdition.COMMUNITY) - def test_save_document_with_dataset_id_initializes_high_quality_dataset_from_default_embedding_model( - self, account_context, sqlite_session: Session - ): - dataset = _dataset_row(data_source_type=None, indexing_technique=None) - knowledge_config = _make_upload_knowledge_config(original_document_id="doc-1", file_ids=["file-1"]) - knowledge_config.indexing_technique = "high_quality" - knowledge_config.embedding_model = None - knowledge_config.embedding_model_provider = None - updated_document = _document_row(document_id="doc-1") - updated_document.batch = "batch-existing" - binding = DatasetCollectionBinding( - provider_name="default-provider", - model_name="default-embedding", - type="dataset", - collection_name="collection", - ) - binding.id = "binding-1" - - with ( - patch("services.knowledge.dataset_service.FeatureService.get_features", return_value=_make_features()), - patch("services.knowledge.dataset_service.ModelManager") as model_manager_cls, - patch( - "services.knowledge.dataset_service.DatasetCollectionBindingService.get_dataset_collection_binding", - return_value=binding, - ) as get_binding, - patch.object(DocumentService, "update_prepared_document", return_value=updated_document), - patch("services.knowledge.dataset_service.document_indexing_update_task.delay"), - ): - default_model = object.__new__(ModelInstance) - default_model.model_name = "default-embedding" - default_model.provider = "default-provider" - model_manager_cls.for_tenant.return_value.get_default_model_instance.return_value = default_model - - documents, batch = DocumentService.save_document_with_dataset_id( - dataset, - knowledge_config, - account_context, - session=sqlite_session, - ) - - assert documents == [updated_document] - assert batch == "batch-existing" - assert dataset.data_source_type == "upload_file" - assert dataset.indexing_technique == "high_quality" - assert dataset.embedding_model == "default-embedding" - assert dataset.embedding_model_provider == "default-provider" - assert dataset.collection_binding_id == "binding-1" - assert dataset.retrieval_model == { - "search_method": "semantic_search", - "reranking_enable": False, - "reranking_model": {"reranking_provider_name": "", "reranking_model_name": ""}, - "top_k": 4, - "score_threshold_enabled": False, - } - get_binding.assert_called_once_with("default-provider", "default-embedding", sqlite_session) - - @config_overrides_context(DEPLOYMENT_EDITION=DeploymentEdition.COMMUNITY) - def test_save_document_with_dataset_id_uses_explicit_embedding_and_retrieval_model( - self, account_context, sqlite_session: Session - ): - dataset = _dataset_row(indexing_technique=None) - knowledge_config = _make_upload_knowledge_config(original_document_id="doc-1", file_ids=["file-1"]) - knowledge_config.indexing_technique = "high_quality" - knowledge_config.embedding_model = "explicit-model" - knowledge_config.embedding_model_provider = "explicit-provider" - knowledge_config.retrieval_model = RetrievalModel( - search_method="semantic_search", - reranking_enable=True, - reranking_model=RerankingModel( - reranking_provider_name="rerank-provider", - reranking_model_name="rerank-model", - ), - top_k=7, - score_threshold_enabled=True, - score_threshold=0.3, - ) - binding = DatasetCollectionBinding( - provider_name="explicit-provider", - model_name="explicit-model", - type="dataset", - collection_name="collection", - ) - binding.id = "binding-2" - updated_document = _document_row(document_id="doc-1") - - with ( - patch("services.knowledge.dataset_service.FeatureService.get_features", return_value=_make_features()), - patch("services.knowledge.dataset_service.ModelManager") as model_manager_cls, - patch( - "services.knowledge.dataset_service.DatasetCollectionBindingService.get_dataset_collection_binding", - return_value=binding, - ) as get_binding, - patch.object(DocumentService, "update_prepared_document", return_value=updated_document), - patch("services.knowledge.dataset_service.document_indexing_update_task.delay"), - ): - DocumentService.save_document_with_dataset_id( - dataset, knowledge_config, account_context, session=sqlite_session - ) - - model_manager_cls.for_tenant.return_value.get_default_model_instance.assert_not_called() - get_binding.assert_called_once_with("explicit-provider", "explicit-model", sqlite_session) - assert dataset.embedding_model == "explicit-model" - assert dataset.embedding_model_provider == "explicit-provider" - assert dataset.retrieval_model == knowledge_config.retrieval_model.model_dump() - - @config_overrides_context(DEPLOYMENT_EDITION=DeploymentEdition.COMMUNITY) - def test_save_document_with_dataset_id_creates_custom_process_rule_for_new_upload_document( - self, account_context, sqlite_session: Session - ): - dataset = _dataset_row(data_source_type=DataSourceType.UPLOAD_FILE) - knowledge_config = _make_upload_knowledge_config( - file_ids=["file-1"], - process_rule=ProcessRule( - mode="custom", - rules=Rule( - pre_processing_rules=[PreProcessingRule(id="remove_stopwords", enabled=True)], - segmentation=Segmentation(separator="\n", max_tokens=128), - ), - ), - ) - upload_file = _upload_file(file_id="file-1", name="file.txt") - sqlite_session.add_all([dataset, upload_file]) - sqlite_session.commit() - - with ( - patch("services.knowledge.dataset_service.FeatureService.get_features", return_value=_make_features()), - patch("services.knowledge.dataset_service.redis_client") as mock_redis, - patch("services.knowledge.dataset_service.DocumentIndexingTaskProxy") as document_proxy_cls, - patch("services.knowledge.dataset_service.time.strftime", return_value="20260101010101"), - patch("services.knowledge.dataset_service.secrets.randbelow", return_value=23), - ): - mock_redis.lock.return_value = _make_lock_context() - documents, batch = DocumentService.save_document_with_dataset_id( - dataset, - knowledge_config, - account_context, - session=sqlite_session, - ) - - assert len(documents) == 1 - created_document = documents[0] - assert created_document.name == "file.txt" - assert batch == "20260101010101100023" - created_rule = sqlite_session.get(DatasetProcessRule, created_document.dataset_process_rule_id) - assert created_rule is not None - assert created_rule.mode == "custom" - assert created_rule.rules == knowledge_config.process_rule.rules.model_dump_json() - document_proxy_cls.assert_called_once_with("tenant-1", "dataset-1", [created_document.id]) - document_proxy_cls.return_value.delay.assert_called_once() - - @config_overrides_context(DEPLOYMENT_EDITION=DeploymentEdition.COMMUNITY) - def test_save_document_with_dataset_id_creates_automatic_process_rule_for_new_upload_document( - self, account_context, sqlite_session: Session - ): - dataset = _dataset_row(data_source_type=DataSourceType.UPLOAD_FILE) - knowledge_config = _make_upload_knowledge_config( - file_ids=["file-1"], - process_rule=ProcessRule(mode="automatic"), - ) - sqlite_session.add_all([dataset, _upload_file(file_id="file-1", name="file.txt")]) - sqlite_session.commit() - - with ( - patch("services.knowledge.dataset_service.FeatureService.get_features", return_value=_make_features()), - patch("services.knowledge.dataset_service.redis_client") as mock_redis, - patch("services.knowledge.dataset_service.DocumentIndexingTaskProxy"), - patch("services.knowledge.dataset_service.time.strftime", return_value="20260101010101"), - patch("services.knowledge.dataset_service.secrets.randbelow", return_value=23), - ): - mock_redis.lock.return_value = _make_lock_context() - documents, _ = DocumentService.save_document_with_dataset_id( - dataset, - knowledge_config, - account_context, - session=sqlite_session, - ) - - created_rule = sqlite_session.get(DatasetProcessRule, documents[0].dataset_process_rule_id) - assert created_rule is not None - assert created_rule.mode == "automatic" - assert created_rule.rules == json.dumps(DatasetProcessRule.AUTOMATIC_RULES) - assert sqlite_session.get(Document, documents[0].id) is documents[0] - - @config_overrides_context(DEPLOYMENT_EDITION=DeploymentEdition.COMMUNITY) - def test_save_document_with_dataset_id_creates_fallback_automatic_process_rule_when_latest_is_missing( - self, account_context, sqlite_session: Session - ): - dataset = _dataset_row(data_source_type=DataSourceType.UPLOAD_FILE) - knowledge_config = _make_upload_knowledge_config(file_ids=["file-1"], process_rule=None) - sqlite_session.add_all([dataset, _upload_file(file_id="file-1", name="file.txt")]) - sqlite_session.commit() - - with ( - patch("services.knowledge.dataset_service.FeatureService.get_features", return_value=_make_features()), - patch("services.knowledge.dataset_service.redis_client") as mock_redis, - patch("services.knowledge.dataset_service.DocumentIndexingTaskProxy"), - patch("services.knowledge.dataset_service.time.strftime", return_value="20260101010101"), - patch("services.knowledge.dataset_service.secrets.randbelow", return_value=23), - ): - mock_redis.lock.return_value = _make_lock_context() - documents, _ = DocumentService.save_document_with_dataset_id( - dataset, - knowledge_config, - account_context, - session=sqlite_session, - ) - - created_rule = sqlite_session.get(DatasetProcessRule, documents[0].dataset_process_rule_id) - assert created_rule is not None - assert created_rule.mode == "automatic" - assert created_rule.rules == json.dumps(DatasetProcessRule.AUTOMATIC_RULES) - - @config_overrides_context(DEPLOYMENT_EDITION=DeploymentEdition.COMMUNITY) - def test_save_document_with_dataset_id_raises_when_upload_file_lookup_is_incomplete( - self, account_context, sqlite_session: Session - ): - dataset = _dataset_row(data_source_type=DataSourceType.UPLOAD_FILE) - knowledge_config = _make_upload_knowledge_config(file_ids=["file-1", "file-2"]) - sqlite_session.add_all([dataset, _upload_file(file_id="file-1", name="file.txt")]) - sqlite_session.commit() - - with ( - patch("services.knowledge.dataset_service.FeatureService.get_features", return_value=_make_features()), - patch("services.knowledge.dataset_service.redis_client") as mock_redis, - patch("services.knowledge.dataset_service.time.strftime", return_value="20260101010101"), - patch("services.knowledge.dataset_service.secrets.randbelow", return_value=23), - ): - mock_redis.lock.return_value = _make_lock_context() - with pytest.raises(FileNotExistsError, match="One or more files not found"): - DocumentService.save_document_with_dataset_id( - dataset, - knowledge_config, - account_context, - session=sqlite_session, - ) - - @config_overrides_context(DEPLOYMENT_EDITION=DeploymentEdition.COMMUNITY) - def test_save_document_with_dataset_id_requires_notion_info_list_for_notion_import( - self, account_context, sqlite_session: Session - ): - dataset = _dataset_row(data_source_type=DataSourceType.NOTION_IMPORT) - process_rule = _process_rule() - sqlite_session.add_all([dataset, process_rule]) - sqlite_session.commit() - knowledge_config = KnowledgeConfig( - indexing_technique="economy", - data_source=DataSource(info_list=InfoList(data_source_type="notion_import")), - doc_form=IndexStructureType.PARAGRAPH_INDEX, - doc_language="English", - ) - - with ( - patch("services.knowledge.dataset_service.FeatureService.get_features", return_value=_make_features()), - patch("services.knowledge.dataset_service.redis_client") as mock_redis, - ): - mock_redis.lock.return_value = _make_lock_context() - with pytest.raises(ValueError, match="No notion info list found"): - DocumentService.save_document_with_dataset_id( - dataset, - knowledge_config, - account_context, - dataset_process_rule=process_rule, - session=sqlite_session, - ) - - @config_overrides_context(DEPLOYMENT_EDITION=DeploymentEdition.COMMUNITY) - def test_save_document_with_dataset_id_requires_website_info_list_for_website_crawl( - self, account_context, sqlite_session: Session - ): - dataset = _dataset_row(data_source_type=DataSourceType.WEBSITE_CRAWL) - process_rule = _process_rule() - sqlite_session.add_all([dataset, process_rule]) - sqlite_session.commit() - knowledge_config = KnowledgeConfig( - indexing_technique="economy", - data_source=DataSource(info_list=InfoList(data_source_type="website_crawl")), - doc_form=IndexStructureType.PARAGRAPH_INDEX, - doc_language="English", - ) - - with ( - patch("services.knowledge.dataset_service.FeatureService.get_features", return_value=_make_features()), - patch("services.knowledge.dataset_service.redis_client") as mock_redis, - ): - mock_redis.lock.return_value = _make_lock_context() - with pytest.raises(ValueError, match="No website info list found"): - DocumentService.save_document_with_dataset_id( - dataset, - knowledge_config, - account_context, - dataset_process_rule=process_rule, - session=sqlite_session, - ) diff --git a/api/tests/unit_tests/services/test_dataset_service_lock_not_owned.py b/api/tests/unit_tests/services/test_dataset_service_lock_not_owned.py deleted file mode 100644 index 19599e6fe9f..00000000000 --- a/api/tests/unit_tests/services/test_dataset_service_lock_not_owned.py +++ /dev/null @@ -1,195 +0,0 @@ -import types -from unittest.mock import create_autospec - -import pytest -from redis.exceptions import LockNotOwnedError -from sqlalchemy import func, select -from sqlalchemy.orm import Session - -from core.rag.index_processor.constant.index_type import IndexStructureType, IndexTechniqueType -from models.account import Account, Tenant -from models.dataset import Dataset, DatasetProcessRule, Document, DocumentSegment -from models.enums import DataSourceType, DocumentCreatedFrom, ProcessRuleMode -from services.knowledge.dataset_service import DocumentService, SegmentService -from services.knowledge.entities.knowledge_entities import DataSource, InfoList, KnowledgeConfig -from services.knowledge.segments.application import SegmentMutationService - -TENANT_ID = "11111111-1111-1111-1111-111111111111" -USER_ID = "22222222-2222-2222-2222-222222222222" -DATASET_ID = "33333333-3333-3333-3333-333333333333" -DOCUMENT_ID = "44444444-4444-4444-4444-444444444444" - - -class FakeLock: - """Lock that always fails on enter with LockNotOwnedError.""" - - def __enter__(self) -> None: - raise LockNotOwnedError("simulated") - - def __exit__( - self, exc_type: type[BaseException] | None, exc: BaseException | None, tb: types.TracebackType | None - ) -> bool: - # Normal contextmanager signature; return False so exceptions propagate - return False - - -@pytest.fixture -def fake_current_user(monkeypatch: pytest.MonkeyPatch) -> Account: - tenant = Tenant(name="Test Tenant") - tenant.id = TENANT_ID - user = Account(name="Test User", email="test@example.com") - user.id = USER_ID - user._current_tenant = tenant - monkeypatch.setattr("services.knowledge.dataset_service.current_user", user) - return user - - -@pytest.fixture -def fake_features(monkeypatch: pytest.MonkeyPatch) -> types.SimpleNamespace: - features = types.SimpleNamespace( - billing=types.SimpleNamespace(subscription=types.SimpleNamespace(plan="ENTERPRISE")), - documents_upload_quota=types.SimpleNamespace(limit=10_000, size=0), - ) - monkeypatch.setattr( - "services.knowledge.dataset_service.FeatureService.get_features", - lambda _tenant_id, **_kwargs: features, - ) - return features - - -@pytest.fixture -def fake_lock(monkeypatch: pytest.MonkeyPatch) -> None: - """Patch redis_client.lock to always raise LockNotOwnedError on enter.""" - - def _fake_lock(*_args: object, **_kwargs: object) -> FakeLock: - return FakeLock() - - # DatasetService imports redis_client directly from extensions.ext_redis - monkeypatch.setattr("services.knowledge.dataset_service.redis_client.lock", _fake_lock) - - -# --------------------------------------------------------------------------- -# 1. Knowledge Pipeline document creation (save_document_with_dataset_id) -# --------------------------------------------------------------------------- - - -@pytest.mark.parametrize("sqlite_session", [(Account, Tenant, Dataset, DatasetProcessRule, Document)], indirect=True) -@pytest.mark.usefixtures("fake_features", "fake_lock") -def test_save_document_with_dataset_id_ignores_lock_not_owned( - monkeypatch: pytest.MonkeyPatch, - fake_current_user: Account, - sqlite_session: Session, -) -> None: - # Arrange - dataset = Dataset( - id=DATASET_ID, - tenant_id=TENANT_ID, - name="Test Dataset", - description="", - created_by=USER_ID, - data_source_type="upload_file", - indexing_technique=IndexTechniqueType.HIGH_QUALITY, - ) - process_rule = DatasetProcessRule( - dataset_id=DATASET_ID, - mode=ProcessRuleMode.AUTOMATIC, - rules=None, - created_by=USER_ID, - ) - sqlite_session.add_all([fake_current_user._current_tenant, fake_current_user, dataset, process_rule]) - sqlite_session.commit() - - knowledge_config = KnowledgeConfig( - doc_form=IndexStructureType.QA_INDEX, - original_document_id=None, # go into "new document" branch - data_source=DataSource(info_list=InfoList(data_source_type="upload_file")), - indexing_technique="high_quality", - duplicate=False, - doc_language="en", - ) - - account = fake_current_user - - # Avoid touching real doc_form logic - monkeypatch.setattr( - "services.knowledge.dataset_service.DatasetService.check_doc_form", lambda *_args, **_kwargs: None - ) - - # Act: this would hit the redis lock, whose __enter__ raises LockNotOwnedError. - # Our implementation should catch it and still return (documents, batch). - documents, batch = DocumentService.save_document_with_dataset_id( - dataset=dataset, - knowledge_config=knowledge_config, - account=account, - dataset_process_rule=process_rule, - session=sqlite_session, - ) - - # Assert - # We mainly care that: - # - No exception is raised - # - The function returns a sensible tuple - assert documents == [] - assert batch - assert not sqlite_session.in_transaction() - assert sqlite_session.scalar(select(func.count(Document.id))) == 0 - assert sqlite_session.get(DatasetProcessRule, process_rule.id) is process_rule - - -# --------------------------------------------------------------------------- -# 2. Single-segment creation (add_segment) -# --------------------------------------------------------------------------- - - -# --------------------------------------------------------------------------- -# 3. Multi-segment creation (multi_create_segment) -# --------------------------------------------------------------------------- - - -@pytest.mark.parametrize("sqlite_session", [(Account, Tenant, Dataset, Document, DocumentSegment)], indirect=True) -@pytest.mark.usefixtures("fake_lock") -def test_multi_create_segment_ignores_lock_not_owned( - fake_current_user: Account, - sqlite_session: Session, -) -> None: - # Arrange - dataset = Dataset( - id=DATASET_ID, - tenant_id=TENANT_ID, - name="Test Dataset", - description="", - created_by=USER_ID, - indexing_technique=IndexTechniqueType.ECONOMY, - ) - document = Document( - id=DOCUMENT_ID, - tenant_id=TENANT_ID, - dataset_id=DATASET_ID, - position=1, - data_source_type=DataSourceType("upload_file"), - data_source_info="{}", - batch="batch-1", - name="Test Document", - created_from=DocumentCreatedFrom("web"), - created_by=USER_ID, - word_count=0, - doc_form=IndexStructureType.QA_INDEX, - ) - sqlite_session.add_all([fake_current_user._current_tenant, fake_current_user, dataset, document]) - sqlite_session.commit() - - mutations = create_autospec(SegmentMutationService, instance=True) - result = SegmentService.multi_create_segment( - segments=[{"content": "question", "answer": "answer", "keywords": ["key"]}], - document=document, - dataset=dataset, - session=sqlite_session, - mutations=mutations, - ) - - assert result is None - mutations.index_segments.assert_not_called() - assert not sqlite_session.in_transaction() - assert sqlite_session.scalar(select(func.count(DocumentSegment.id))) == 0 - sqlite_session.refresh(document) - assert document.word_count == 0 diff --git a/api/tests/unit_tests/services/test_dataset_service_segment.py b/api/tests/unit_tests/services/test_dataset_service_segment.py deleted file mode 100644 index 527041a6497..00000000000 --- a/api/tests/unit_tests/services/test_dataset_service_segment.py +++ /dev/null @@ -1,392 +0,0 @@ -"""SQLite-backed tests for segment and child-chunk dataset services.""" - -from __future__ import annotations - -from collections.abc import Callable -from datetime import datetime -from types import SimpleNamespace -from unittest.mock import MagicMock, create_autospec, patch - -import pytest -from sqlalchemy.orm import Session, sessionmaker - -from core.rag.index_processor.constant.index_type import IndexStructureType, IndexTechniqueType -from models import Account -from models.account import Tenant -from models.dataset import ( - ChildChunk, - Dataset, - Document, - DocumentSegment, -) -from models.enums import DataSourceType, DocumentCreatedFrom, SegmentStatus -from repositories.knowledge.segment_repository import SQLAlchemySegmentRepository -from services.knowledge.dataset_service import SegmentService -from services.knowledge.resource_scope import DatasetRef, DocumentRef, SegmentRef -from services.knowledge.segments.application import SegmentIndex, SegmentIndexingState, SegmentMutationService - - -def _account(*, account_id: str = "user-1", tenant_id: str = "tenant-1") -> Account: - account = Account(name="User", email=f"{account_id}@example.com") - account.id = account_id - tenant = Tenant(name="Tenant") - tenant.id = tenant_id - account._current_tenant = tenant - return account - - -def _dataset(*, dataset_id: str = "dataset-1", tenant_id: str = "tenant-1") -> Dataset: - return Dataset( - id=dataset_id, - tenant_id=tenant_id, - name="Dataset", - description="", - provider="vendor", - created_by="user-1", - maintainer="user-1", - indexing_technique=IndexTechniqueType.HIGH_QUALITY, - embedding_model_provider="provider", - embedding_model="embedding-model", - chunk_structure=IndexStructureType.PARAGRAPH_INDEX, - ) - - -def _document( - *, - document_id: str = "document-1", - dataset_id: str = "dataset-1", - tenant_id: str = "tenant-1", - doc_form: str = IndexStructureType.PARAGRAPH_INDEX, - word_count: int = 20, -) -> Document: - document = Document( - id=document_id, - tenant_id=tenant_id, - dataset_id=dataset_id, - position=1, - data_source_type=DataSourceType.UPLOAD_FILE, - batch="batch-1", - name="Document", - created_from=DocumentCreatedFrom.API, - created_by="user-1", - indexing_status="completed", - doc_form=doc_form, - word_count=word_count, - ) - - document.created_at = datetime(2026, 1, 1) - document.updated_at = datetime(2026, 1, 2) - return document - - -def _segment( - *, - segment_id: str = "segment-1", - dataset_id: str = "dataset-1", - document_id: str = "document-1", - tenant_id: str = "tenant-1", - position: int = 1, - content: str = "segment content", - enabled: bool = True, -) -> DocumentSegment: - segment = DocumentSegment( - tenant_id=tenant_id, - dataset_id=dataset_id, - document_id=document_id, - position=position, - content=content, - word_count=len(content), - tokens=0, - created_by="user-1", - enabled=enabled, - keywords=[], - answer=None, - index_node_id=f"node-{segment_id}", - status=SegmentStatus.COMPLETED, - ) - segment.id = segment_id - return segment - - -def _child( - *, - child_id: str = "child-1", - segment_id: str = "segment-1", - dataset_id: str = "dataset-1", - document_id: str = "document-1", - tenant_id: str = "tenant-1", - position: int = 1, - content: str = "child content", -) -> ChildChunk: - child = ChildChunk( - tenant_id=tenant_id, - dataset_id=dataset_id, - document_id=document_id, - segment_id=segment_id, - position=position, - index_node_id=f"node-{child_id}", - index_node_hash=f"hash-{child_id}", - content=content, - word_count=len(content), - created_by="user-1", - ) - child.id = child_id - return child - - -def _persist_chain(session: Session) -> tuple[Dataset, Document, DocumentSegment]: - dataset = _dataset() - document = _document() - segment = _segment() - session.add_all([dataset, document, segment]) - session.commit() - return dataset, document, segment - - -class TestSegmentServiceQueries: - def test_get_child_chunks_filters_owner_keyword_and_paginates(self, sqlite_session: Session) -> None: - _persist_chain(sqlite_session) - sqlite_session.add_all( - [ - _child(child_id="child-1", position=1, content="alpha match"), - _child(child_id="child-2", position=2, content="beta"), - _child(child_id="child-3", position=3, content="alpha second"), - _child(child_id="foreign", tenant_id="tenant-2", position=4, content="alpha foreign"), - ] - ) - sqlite_session.commit() - - with patch("services.knowledge.dataset_service.current_user", _account()): - page = SegmentService.get_child_chunks( - "segment-1", - "document-1", - "dataset-1", - page=1, - limit=1, - keyword="alpha", - session=sqlite_session, - ) - - assert page.total == 2 - assert [child.id for child in page.items] == ["child-1"] - - def test_get_child_chunk_by_id_scopes_tenant(self, sqlite_session: Session) -> None: - owned = _child() - sqlite_session.add(owned) - sqlite_session.commit() - - assert SegmentService.get_child_chunk_by_id(owned.id, "tenant-1", sqlite_session) is owned - assert SegmentService.get_child_chunk_by_id(owned.id, "tenant-2", sqlite_session) is None - - def test_get_child_chunk_by_segment_ref_enforces_full_chain(self, sqlite_session: Session) -> None: - child = _child() - sqlite_session.add(child) - sqlite_session.commit() - valid_ref = SegmentRef(DocumentRef(DatasetRef("tenant-1", "dataset-1"), "document-1"), "segment-1") - - assert SegmentService.get_child_chunk_by_segment_ref(child.id, valid_ref, sqlite_session) is child - spoofed_refs = [ - SegmentRef(DocumentRef(DatasetRef("tenant-2", "dataset-1"), "document-1"), "segment-1"), - SegmentRef(DocumentRef(DatasetRef("tenant-1", "dataset-2"), "document-1"), "segment-1"), - SegmentRef(DocumentRef(DatasetRef("tenant-1", "dataset-1"), "document-2"), "segment-1"), - SegmentRef(DocumentRef(DatasetRef("tenant-1", "dataset-1"), "document-1"), "segment-2"), - ] - for spoofed_ref in spoofed_refs: - assert SegmentService.get_child_chunk_by_segment_ref(child.id, spoofed_ref, sqlite_session) is None - - def test_get_segments_filters_status_keyword_and_orders(self, sqlite_session: Session) -> None: - sqlite_session.add_all( - [ - _segment(segment_id="one", position=2, content="alpha later"), - _segment(segment_id="two", position=1, content="alpha first"), - _segment(segment_id="three", position=3, content="beta"), - _segment(segment_id="foreign", tenant_id="tenant-2", position=1, content="alpha foreign"), - ] - ) - sqlite_session.commit() - - segments, total = SegmentService.get_segments( - "document-1", - "tenant-1", - status_list=[SegmentStatus.COMPLETED], - keyword="alpha", - session=sqlite_session, - ) - - assert total == 2 - assert [segment.id for segment in segments] == ["two", "one"] - - def test_get_segment_by_id_and_ref_scope_complete_owner(self, sqlite_session: Session) -> None: - segment = _segment() - sqlite_session.add(segment) - sqlite_session.commit() - valid_ref = SegmentRef(DocumentRef(DatasetRef("tenant-1", "dataset-1"), "document-1"), segment.id) - - assert SegmentService.get_segment_by_id(segment.id, "tenant-1", sqlite_session) is segment - assert SegmentService.get_segment_by_id(segment.id, "tenant-2", sqlite_session) is None - assert SegmentService.get_segment_by_ref(valid_ref, sqlite_session) is segment - spoofed_refs = [ - SegmentRef(DocumentRef(DatasetRef("tenant-2", "dataset-1"), "document-1"), segment.id), - SegmentRef(DocumentRef(DatasetRef("tenant-1", "dataset-2"), "document-1"), segment.id), - SegmentRef(DocumentRef(DatasetRef("tenant-1", "dataset-1"), "document-2"), segment.id), - ] - for spoofed_ref in spoofed_refs: - assert SegmentService.get_segment_by_ref(spoofed_ref, sqlite_session) is None - - def test_get_segments_by_document_and_dataset_returns_real_rows(self, sqlite_session: Session) -> None: - sqlite_session.add_all( - [ - _segment(segment_id="enabled"), - _segment(segment_id="disabled", position=2, enabled=False), - _segment(segment_id="other", document_id="document-2", position=1), - _segment(segment_id="foreign", tenant_id="tenant-2"), - ] - ) - sqlite_session.commit() - - segments = SegmentService.get_segments_by_document_and_dataset( - "document-1", - "dataset-1", - sqlite_session, - tenant_id="tenant-1", - status=SegmentStatus.COMPLETED, - enabled=True, - ) - - assert [segment.id for segment in segments] == ["enabled"] - - -class TestSegmentServiceValidation: - def test_qa_segment_requires_answer(self) -> None: - with pytest.raises(ValueError, match="Answer is required"): - SegmentService.segment_create_args_validate( - {"content": "question"}, _document(doc_form=IndexStructureType.QA_INDEX) - ) - - @pytest.mark.parametrize("content", [None, "", " "]) - def test_segment_requires_non_empty_content(self, content: str | None) -> None: - with pytest.raises(ValueError, match="Content is empty"): - SegmentService.segment_create_args_validate({"content": content}, _document()) - - def test_segment_attachment_ids_must_be_a_list(self) -> None: - with pytest.raises(ValueError, match="Attachment IDs is invalid"): - SegmentService.segment_create_args_validate({"content": "text", "attachment_ids": "file"}, _document()) - - def test_segment_attachment_limit_is_enforced(self, config_overrides: Callable[..., None]) -> None: - config_overrides(SINGLE_CHUNK_ATTACHMENT_LIMIT=1) - - with pytest.raises(ValueError, match="Exceeded maximum attachment limit"): - SegmentService.segment_create_args_validate( - {"content": "text", "attachment_ids": ["one", "two"]}, - _document(), - ) - - def test_segment_attachment_limit_accepts_exact_boundary(self, config_overrides: Callable[..., None]) -> None: - config_overrides(SINGLE_CHUNK_ATTACHMENT_LIMIT=2) - - SegmentService.segment_create_args_validate( - {"content": "text", "attachment_ids": ["one", "two"]}, - _document(), - ) - - -class TestSegmentServiceMutations: - @pytest.mark.parametrize( - ("document_dataset_id", "document_tenant_id"), - [("dataset-2", "tenant-1"), ("dataset-1", "tenant-2")], - ) - def test_delete_segment_rejects_document_outside_dataset( - self, document_dataset_id: str, document_tenant_id: str, sqlite_session: Session - ) -> None: - dataset = _dataset() - document = _document(dataset_id=document_dataset_id, tenant_id=document_tenant_id) - segment = _segment(dataset_id=document_dataset_id, tenant_id=document_tenant_id) - mutations = create_autospec(SegmentMutationService, instance=True, spec_set=True) - - with pytest.raises(ValueError, match="Document does not belong to the dataset"): - SegmentService.delete_segment(segment, document, dataset, sqlite_session, mutations=mutations) - - mutations.delete_segment.assert_not_called() - - def test_multi_create_segment_marks_each_real_row_error_on_vector_failure( - self, sqlite_session: Session, sqlite_session_factory: sessionmaker[Session] - ) -> None: - dataset = _dataset() - document = _document(word_count=0) - sqlite_session.add_all([dataset, document]) - sqlite_session.commit() - embedding_model = SimpleNamespace(get_text_embedding_num_tokens=lambda *, texts: [len(texts) + 1]) - index = create_autospec(SegmentIndex, instance=True, spec_set=True) - - def fail(*_args: object, **_kwargs: object) -> None: - assert not sqlite_session.in_transaction() - raise RuntimeError("vector failed") - - index.create_many.side_effect = fail - mutations = SegmentMutationService( - store=SQLAlchemySegmentRepository(session_factory=sqlite_session_factory), - index=index, - text_hash=lambda text: text, - indexing_state=create_autospec(SegmentIndexingState, instance=True, spec_set=True), - ) - - with ( - patch("services.knowledge.dataset_service.current_user", _account()), - patch("services.knowledge.dataset_service.ModelManager") as manager_cls, - ): - manager_cls.for_tenant.return_value.get_model_instance.return_value = embedding_model - segments = SegmentService.multi_create_segment( - [{"content": "one"}, {"content": "two"}], - document, - dataset, - sqlite_session, - mutations=mutations, - ) - - assert segments is not None - assert len(segments) == 2 - assert all(segment.status == SegmentStatus.ERROR and not segment.enabled for segment in segments) - assert document.word_count == 6 - - def test_multi_create_segment_persists_qa_counts_positions_and_tokens(self, sqlite_session: Session) -> None: - dataset = _dataset() - document = _document(doc_form=IndexStructureType.QA_INDEX, word_count=5) - sqlite_session.add_all([dataset, document, _segment(segment_id="existing")]) - sqlite_session.commit() - embedding_model = MagicMock() - embedding_model.get_text_embedding_num_tokens.side_effect = [[11], [13]] - mutations = create_autospec(SegmentMutationService, instance=True, spec_set=True) - - with ( - patch("services.knowledge.dataset_service.current_user", _account()), - patch("services.knowledge.dataset_service.ModelManager") as manager_cls, - ): - manager_cls.for_tenant.return_value.get_model_instance.return_value = embedding_model - segments = SegmentService.multi_create_segment( - [ - {"content": "question-1", "answer": "answer-1", "keywords": ["key"]}, - {"content": "question-2", "answer": "answer-2"}, - ], - document, - dataset, - sqlite_session, - mutations=mutations, - ) - - assert segments is not None - assert [segment.position for segment in segments] == [2, 3] - assert [segment.tokens for segment in segments] == [11, 13] - assert [segment.answer for segment in segments] == ["answer-1", "answer-2"] - expected_increment = sum(len(segment.content) + len(segment.answer or "") for segment in segments) - assert document.word_count == 5 + expected_increment - assert embedding_model.get_text_embedding_num_tokens.call_args_list[0].kwargs == { - "texts": ["question-1answer-1"] - } - assert embedding_model.get_text_embedding_num_tokens.call_args_list[1].kwargs == { - "texts": ["question-2answer-2"] - } - mutations.index_segments.assert_called_once_with( - DatasetRef("tenant-1", "dataset-1").document("document-1"), - segment_ids=[segment.id for segment in segments], - keywords_list=[["key"], None], - ) diff --git a/api/tests/unit_tests/services/test_email_code_login_challenge.py b/api/tests/unit_tests/services/test_email_code_login_challenge.py index 5ea11346eda..9502c8c19e3 100644 --- a/api/tests/unit_tests/services/test_email_code_login_challenge.py +++ b/api/tests/unit_tests/services/test_email_code_login_challenge.py @@ -10,7 +10,7 @@ from services.email_code_login_challenge import ( EmailCodeLoginChallengeStore, EmailCodeLoginChallengeUnavailableError, ) -from tests.unit_tests.config_override import apply_config_overrides +from tests.helpers.unit.config_override import apply_config_overrides TOKEN = "00000000-0000-4000-8000-000000000001" diff --git a/api/tests/unit_tests/services/test_feature_service.py b/api/tests/unit_tests/services/test_feature_service.py new file mode 100644 index 00000000000..e7fbbe21733 --- /dev/null +++ b/api/tests/unit_tests/services/test_feature_service.py @@ -0,0 +1,345 @@ +from collections.abc import Callable +from dataclasses import dataclass +from typing import cast +from unittest.mock import Mock, patch + +import pytest + +from enums import CloudPlan, DeploymentEdition, HostedTrialProvider +from services import feature_service as feature_service_module +from services.account_education_service import EDUCATION_EDITIONS +from services.billing_service import BillingInfo +from services.entities.feature_entities import FeatureModel, LimitationModel +from services.feature_service import FeatureService +from services.system_feature_service import SystemFeatureService + + +def test_education_feature_is_disabled_by_default() -> None: + assert FeatureModel().education.enabled is False + + +@pytest.mark.parametrize("edition", list(DeploymentEdition)) +def test_education_is_advertised_only_where_its_endpoints_exist( + config_overrides: Callable[..., None], edition: DeploymentEdition +) -> None: + """The frontend gates its /account/education* calls on this flag, and the admission + decorator answers 404 on every edition outside EDUCATION_EDITIONS.""" + config_overrides(DEPLOYMENT_EDITION=edition, EDUCATION_ENABLED=True) + + with patch("services.feature_service.EnterpriseService.get_workspace_info", return_value={}): + features = FeatureService.get_features("") + + assert features.education.enabled is (edition in EDUCATION_EDITIONS) + + +def test_education_stays_disabled_on_cloud_without_the_env_flag( + config_overrides: Callable[..., None], +) -> None: + config_overrides(DEPLOYMENT_EDITION=DeploymentEdition.CLOUD, EDUCATION_ENABLED=False) + + features = FeatureService.get_features("") + + assert features.education.enabled is False + + +@dataclass(frozen=True) +class HumanInputEmailDeliveryCase: + name: str + deployment_edition: DeploymentEdition + tenant_id: str | None + plan: CloudPlan + expected: bool + + +CASES = [ + HumanInputEmailDeliveryCase( + name="enterprise_edition", + deployment_edition=DeploymentEdition.ENTERPRISE, + tenant_id=None, + plan=CloudPlan.SANDBOX, + expected=True, + ), + HumanInputEmailDeliveryCase( + name="community_edition", + deployment_edition=DeploymentEdition.COMMUNITY, + tenant_id=None, + plan=CloudPlan.SANDBOX, + expected=True, + ), + HumanInputEmailDeliveryCase( + name="cloud_edition_requires_tenant", + deployment_edition=DeploymentEdition.CLOUD, + tenant_id=None, + plan=CloudPlan.PROFESSIONAL, + expected=False, + ), + HumanInputEmailDeliveryCase( + name="professional_plan", + deployment_edition=DeploymentEdition.CLOUD, + tenant_id="tenant-1", + plan=CloudPlan.PROFESSIONAL, + expected=True, + ), + HumanInputEmailDeliveryCase( + name="team_plan", + deployment_edition=DeploymentEdition.CLOUD, + tenant_id="tenant-1", + plan=CloudPlan.TEAM, + expected=True, + ), + HumanInputEmailDeliveryCase( + name="sandbox_plan", + deployment_edition=DeploymentEdition.CLOUD, + tenant_id="tenant-1", + plan=CloudPlan.SANDBOX, + expected=False, + ), +] + + +@pytest.mark.parametrize("case", CASES, ids=lambda case: case.name) +def test_resolve_human_input_email_delivery_enabled_matrix( + config_overrides: Callable[..., None], + case: HumanInputEmailDeliveryCase, +) -> None: + config_overrides(DEPLOYMENT_EDITION=case.deployment_edition) + features = FeatureModel() + features.billing.subscription.plan = case.plan + + result = FeatureService._resolve_human_input_email_delivery_enabled( + features=features, + tenant_id=case.tenant_id, + ) + + assert result is case.expected + + +def test_get_vector_space_converts_billing_float_size( + monkeypatch: pytest.MonkeyPatch, config_overrides: Callable[..., None] +) -> None: + config_overrides(DEPLOYMENT_EDITION=DeploymentEdition.CLOUD) + monkeypatch.setattr( + feature_service_module.BillingService, + "get_vector_space", + lambda _tenant_id: {"size": 5120.75, "limit": 20480}, + ) + + result = FeatureService.get_vector_space("tenant-1") + + assert result.size == 5120 + assert result.limit == 20480 + assert result.usage_unknown is False + + +def test_get_vector_space_preserves_unknown_usage( + monkeypatch: pytest.MonkeyPatch, config_overrides: Callable[..., None] +) -> None: + config_overrides(DEPLOYMENT_EDITION=DeploymentEdition.CLOUD) + monkeypatch.setattr( + feature_service_module.BillingService, + "get_vector_space", + lambda _tenant_id: {"size": 0.0, "limit": 50, "usage_unknown": True}, + ) + + result = FeatureService.get_vector_space("tenant-1") + + assert result.size == 0 + assert result.limit == 50 + assert result.usage_unknown is True + + +@pytest.mark.parametrize( + ("deployment_edition", "tenant_id", "plan", "expected"), + [ + (DeploymentEdition.COMMUNITY, "tenant-1", CloudPlan.PROFESSIONAL, 15), + (DeploymentEdition.ENTERPRISE, "tenant-1", CloudPlan.PROFESSIONAL, 15), + (DeploymentEdition.CLOUD, None, CloudPlan.PROFESSIONAL, 15), + (DeploymentEdition.CLOUD, "tenant-1", CloudPlan.SANDBOX, 15), + (DeploymentEdition.CLOUD, "tenant-1", CloudPlan.PROFESSIONAL, 50), + (DeploymentEdition.CLOUD, "tenant-1", CloudPlan.TEAM, 50), + ], +) +def test_get_knowledge_file_size_limit( + monkeypatch: pytest.MonkeyPatch, + config_overrides: Callable[..., None], + deployment_edition: DeploymentEdition, + tenant_id: str | None, + plan: CloudPlan, + expected: int, +) -> None: + config_overrides( + DEPLOYMENT_EDITION=deployment_edition, + UPLOAD_FILE_SIZE_LIMIT=15, + KNOWLEDGE_UPLOAD_FILE_SIZE_LIMIT_FOR_PAID_PLAN=50, + ) + get_info = Mock( + return_value={ + "subscription": {"plan": plan}, + } + ) + monkeypatch.setattr(feature_service_module.BillingService, "get_info", get_info) + + assert FeatureService.get_knowledge_file_size_limit(tenant_id) == expected + + if deployment_edition == DeploymentEdition.CLOUD and tenant_id: + get_info.assert_called_once_with(tenant_id, exclude_vector_space=True) + else: + get_info.assert_not_called() + + +def test_paid_knowledge_file_size_limit_never_reduces_default( + monkeypatch: pytest.MonkeyPatch, config_overrides: Callable[..., None] +) -> None: + config_overrides( + DEPLOYMENT_EDITION=DeploymentEdition.CLOUD, + UPLOAD_FILE_SIZE_LIMIT=100, + KNOWLEDGE_UPLOAD_FILE_SIZE_LIMIT_FOR_PAID_PLAN=50, + ) + monkeypatch.setattr( + feature_service_module.BillingService, + "get_info", + lambda *_args, **_kwargs: { + "subscription": {"plan": CloudPlan.PROFESSIONAL}, + }, + ) + + assert FeatureService.get_knowledge_file_size_limit("tenant-1") == 100 + + +@pytest.mark.parametrize( + ("edition", "plan", "expected"), + [ + (DeploymentEdition.COMMUNITY, CloudPlan.SANDBOX, True), + (DeploymentEdition.ENTERPRISE, CloudPlan.SANDBOX, True), + (DeploymentEdition.CLOUD, CloudPlan.SANDBOX, False), + (DeploymentEdition.CLOUD, CloudPlan.PROFESSIONAL, True), + ], +) +def test_site_import_entitlement( + monkeypatch: pytest.MonkeyPatch, + config_overrides: Callable[..., None], + edition: DeploymentEdition, + plan: CloudPlan, + expected: bool, +) -> None: + config_overrides(DEPLOYMENT_EDITION=edition) + get_plan = Mock(return_value=plan) + monkeypatch.setattr(FeatureService, "get_workspace_plan", get_plan) + + assert FeatureService.can_import_premium_site_settings("tenant-1") is expected + assert get_plan.call_count == (1 if edition == DeploymentEdition.CLOUD else 0) + + +def test_skill_feature_is_disabled_by_default() -> None: + assert FeatureModel().enable_skill is True + + +def test_skill_feature_follows_env_config(config_overrides: Callable[..., None]) -> None: + config_overrides(DEPLOYMENT_EDITION=DeploymentEdition.COMMUNITY, ENABLE_SKILL=True) + + features = FeatureService.get_features("") + + assert features.enable_skill is True + + +def test_get_public_system_features_excludes_trial_models() -> None: + result = SystemFeatureService.get_public_system_features().model_dump() + + assert "trial_models" not in result + + +def test_get_trial_models_returns_providers_with_paid_or_trial_enabled( + config_overrides: Callable[..., None], +) -> None: + values: dict[str, object] = {} + for provider in HostedTrialProvider: + values[f"HOSTED_{provider.config_key}_PAID_ENABLED"] = False + values[f"HOSTED_{provider.config_key}_TRIAL_ENABLED"] = False + + values.update( + HOSTED_OPENAI_PAID_ENABLED=True, + HOSTED_OPENAI_TRIAL_ENABLED=True, + HOSTED_XAI_PAID_ENABLED=True, + DEPLOYMENT_EDITION=DeploymentEdition.COMMUNITY, + ) + config_overrides(**values) + + result = FeatureService.get_trial_models("tenant_1") + + assert result == [ + HostedTrialProvider.OPENAI.value, + HostedTrialProvider.X.value, + ] + + +@pytest.mark.parametrize( + ("plan", "expected"), + [ + (CloudPlan.SANDBOX, [HostedTrialProvider.OPENAI.value]), + (CloudPlan.PROFESSIONAL, [HostedTrialProvider.OPENAI.value, HostedTrialProvider.X.value]), + ], +) +def test_get_trial_models_filters_providers_by_workspace_plan( + monkeypatch: pytest.MonkeyPatch, + config_overrides: Callable[..., None], + plan: CloudPlan, + expected: list[str], +) -> None: + values: dict[str, object] = {} + for provider in HostedTrialProvider: + values[f"HOSTED_{provider.config_key}_PAID_ENABLED"] = False + values[f"HOSTED_{provider.config_key}_TRIAL_ENABLED"] = False + values.update( + DEPLOYMENT_EDITION=DeploymentEdition.CLOUD, + HOSTED_OPENAI_TRIAL_ENABLED=True, + HOSTED_XAI_PAID_ENABLED=True, + ) + config_overrides(**values) + get_workspace_plan = Mock(return_value=plan) + monkeypatch.setattr(feature_service_module.FeatureService, "get_workspace_plan", get_workspace_plan) + + result = FeatureService.get_trial_models("tenant_1") + + assert result == expected + get_workspace_plan.assert_called_once_with("tenant_1") + + +def test_get_features_exclude_vector_space_sets_vector_space_to_none(config_overrides: Callable[..., None]) -> None: + tenant_id = "tenant-id" + billing_info = { + "subscription": {"plan": CloudPlan.PROFESSIONAL, "interval": "monthly", "education": False}, + "members": {"size": 1, "limit": 10}, + "apps": {"size": 2, "limit": 20}, + "documents_upload_quota": {"size": 3, "limit": 100}, + "annotation_quota_limit": {"size": 4, "limit": 50}, + "can_replace_logo": True, + "model_load_balancing_enabled": True, + "knowledge_pipeline_publish_enabled": True, + } + + config_overrides( + DEPLOYMENT_EDITION=DeploymentEdition.CLOUD, + CAN_REPLACE_LOGO=False, + MODEL_LB_ENABLED=False, + EDUCATION_ENABLED=False, + ) + with ( + patch("services.feature_service.BillingService.get_info", return_value=billing_info) as get_info, + patch("services.feature_service.BillingService.get_quota_info", return_value={}), + ): + features = FeatureService.get_features(tenant_id, exclude_vector_space=True) + + assert features.vector_space is None + get_info.assert_called_once_with(tenant_id, exclude_vector_space=True) + + +def test_full_features_keep_treating_unknown_vector_usage_as_zero() -> None: + vector_space = LimitationModel() + + FeatureService._fulfill_vector_space_from_billing_info( + vector_space, + cast(BillingInfo, {"vector_space": {"size": 0.0, "limit": 50, "usage_unknown": True}}), + ) + + assert vector_space.size == 0 + assert vector_space.limit == 50 diff --git a/api/tests/unit_tests/services/test_feature_service_education.py b/api/tests/unit_tests/services/test_feature_service_education.py deleted file mode 100644 index 7eb04e2e4f4..00000000000 --- a/api/tests/unit_tests/services/test_feature_service_education.py +++ /dev/null @@ -1,37 +0,0 @@ -from collections.abc import Callable -from unittest.mock import patch - -import pytest - -from enums import DeploymentEdition -from services.account_education_service import EDUCATION_EDITIONS -from services.entities.feature_entities import FeatureModel -from services.feature_service import FeatureService - - -def test_education_feature_is_disabled_by_default() -> None: - assert FeatureModel().education.enabled is False - - -@pytest.mark.parametrize("edition", list(DeploymentEdition)) -def test_education_is_advertised_only_where_its_endpoints_exist( - config_overrides: Callable[..., None], edition: DeploymentEdition -) -> None: - """The frontend gates its /account/education* calls on this flag, and the admission - decorator answers 404 on every edition outside EDUCATION_EDITIONS.""" - config_overrides(DEPLOYMENT_EDITION=edition, EDUCATION_ENABLED=True) - - with patch("services.feature_service.EnterpriseService.get_workspace_info", return_value={}): - features = FeatureService.get_features("") - - assert features.education.enabled is (edition in EDUCATION_EDITIONS) - - -def test_education_stays_disabled_on_cloud_without_the_env_flag( - config_overrides: Callable[..., None], -) -> None: - config_overrides(DEPLOYMENT_EDITION=DeploymentEdition.CLOUD, EDUCATION_ENABLED=False) - - features = FeatureService.get_features("") - - assert features.education.enabled is False diff --git a/api/tests/unit_tests/services/test_feature_service_human_input_email_delivery.py b/api/tests/unit_tests/services/test_feature_service_human_input_email_delivery.py deleted file mode 100644 index 6bc9cdb108a..00000000000 --- a/api/tests/unit_tests/services/test_feature_service_human_input_email_delivery.py +++ /dev/null @@ -1,115 +0,0 @@ -from collections.abc import Callable -from dataclasses import dataclass - -import pytest - -from enums import CloudPlan, DeploymentEdition -from services import feature_service as feature_service_module -from services.entities.feature_entities import FeatureModel -from services.feature_service import FeatureService - - -@dataclass(frozen=True) -class HumanInputEmailDeliveryCase: - name: str - deployment_edition: DeploymentEdition - tenant_id: str | None - plan: CloudPlan - expected: bool - - -CASES = [ - HumanInputEmailDeliveryCase( - name="enterprise_edition", - deployment_edition=DeploymentEdition.ENTERPRISE, - tenant_id=None, - plan=CloudPlan.SANDBOX, - expected=True, - ), - HumanInputEmailDeliveryCase( - name="community_edition", - deployment_edition=DeploymentEdition.COMMUNITY, - tenant_id=None, - plan=CloudPlan.SANDBOX, - expected=True, - ), - HumanInputEmailDeliveryCase( - name="cloud_edition_requires_tenant", - deployment_edition=DeploymentEdition.CLOUD, - tenant_id=None, - plan=CloudPlan.PROFESSIONAL, - expected=False, - ), - HumanInputEmailDeliveryCase( - name="professional_plan", - deployment_edition=DeploymentEdition.CLOUD, - tenant_id="tenant-1", - plan=CloudPlan.PROFESSIONAL, - expected=True, - ), - HumanInputEmailDeliveryCase( - name="team_plan", - deployment_edition=DeploymentEdition.CLOUD, - tenant_id="tenant-1", - plan=CloudPlan.TEAM, - expected=True, - ), - HumanInputEmailDeliveryCase( - name="sandbox_plan", - deployment_edition=DeploymentEdition.CLOUD, - tenant_id="tenant-1", - plan=CloudPlan.SANDBOX, - expected=False, - ), -] - - -@pytest.mark.parametrize("case", CASES, ids=lambda case: case.name) -def test_resolve_human_input_email_delivery_enabled_matrix( - config_overrides: Callable[..., None], - case: HumanInputEmailDeliveryCase, -) -> None: - config_overrides(DEPLOYMENT_EDITION=case.deployment_edition) - features = FeatureModel() - features.billing.subscription.plan = case.plan - - result = FeatureService._resolve_human_input_email_delivery_enabled( - features=features, - tenant_id=case.tenant_id, - ) - - assert result is case.expected - - -def test_get_vector_space_converts_billing_float_size( - monkeypatch: pytest.MonkeyPatch, config_overrides: Callable[..., None] -) -> None: - config_overrides(DEPLOYMENT_EDITION=DeploymentEdition.CLOUD) - monkeypatch.setattr( - feature_service_module.BillingService, - "get_vector_space", - lambda _tenant_id: {"size": 5120.75, "limit": 20480}, - ) - - result = FeatureService.get_vector_space("tenant-1") - - assert result.size == 5120 - assert result.limit == 20480 - assert result.usage_unknown is False - - -def test_get_vector_space_preserves_unknown_usage( - monkeypatch: pytest.MonkeyPatch, config_overrides: Callable[..., None] -) -> None: - config_overrides(DEPLOYMENT_EDITION=DeploymentEdition.CLOUD) - monkeypatch.setattr( - feature_service_module.BillingService, - "get_vector_space", - lambda _tenant_id: {"size": 0.0, "limit": 50, "usage_unknown": True}, - ) - - result = FeatureService.get_vector_space("tenant-1") - - assert result.size == 0 - assert result.limit == 50 - assert result.usage_unknown is True diff --git a/api/tests/unit_tests/services/test_feature_service_knowledge_file_size_limit.py b/api/tests/unit_tests/services/test_feature_service_knowledge_file_size_limit.py deleted file mode 100644 index cf6105373bb..00000000000 --- a/api/tests/unit_tests/services/test_feature_service_knowledge_file_size_limit.py +++ /dev/null @@ -1,66 +0,0 @@ -from collections.abc import Callable -from unittest.mock import Mock - -import pytest - -from enums import CloudPlan, DeploymentEdition -from services import feature_service as feature_service_module -from services.feature_service import FeatureService - - -@pytest.mark.parametrize( - ("deployment_edition", "tenant_id", "plan", "expected"), - [ - (DeploymentEdition.COMMUNITY, "tenant-1", CloudPlan.PROFESSIONAL, 15), - (DeploymentEdition.ENTERPRISE, "tenant-1", CloudPlan.PROFESSIONAL, 15), - (DeploymentEdition.CLOUD, None, CloudPlan.PROFESSIONAL, 15), - (DeploymentEdition.CLOUD, "tenant-1", CloudPlan.SANDBOX, 15), - (DeploymentEdition.CLOUD, "tenant-1", CloudPlan.PROFESSIONAL, 50), - (DeploymentEdition.CLOUD, "tenant-1", CloudPlan.TEAM, 50), - ], -) -def test_get_knowledge_file_size_limit( - monkeypatch: pytest.MonkeyPatch, - config_overrides: Callable[..., None], - deployment_edition: DeploymentEdition, - tenant_id: str | None, - plan: CloudPlan, - expected: int, -) -> None: - config_overrides( - DEPLOYMENT_EDITION=deployment_edition, - UPLOAD_FILE_SIZE_LIMIT=15, - KNOWLEDGE_UPLOAD_FILE_SIZE_LIMIT_FOR_PAID_PLAN=50, - ) - get_info = Mock( - return_value={ - "subscription": {"plan": plan}, - } - ) - monkeypatch.setattr(feature_service_module.BillingService, "get_info", get_info) - - assert FeatureService.get_knowledge_file_size_limit(tenant_id) == expected - - if deployment_edition == DeploymentEdition.CLOUD and tenant_id: - get_info.assert_called_once_with(tenant_id, exclude_vector_space=True) - else: - get_info.assert_not_called() - - -def test_paid_knowledge_file_size_limit_never_reduces_default( - monkeypatch: pytest.MonkeyPatch, config_overrides: Callable[..., None] -) -> None: - config_overrides( - DEPLOYMENT_EDITION=DeploymentEdition.CLOUD, - UPLOAD_FILE_SIZE_LIMIT=100, - KNOWLEDGE_UPLOAD_FILE_SIZE_LIMIT_FOR_PAID_PLAN=50, - ) - monkeypatch.setattr( - feature_service_module.BillingService, - "get_info", - lambda *_args, **_kwargs: { - "subscription": {"plan": CloudPlan.PROFESSIONAL}, - }, - ) - - assert FeatureService.get_knowledge_file_size_limit("tenant-1") == 100 diff --git a/api/tests/unit_tests/services/test_feature_service_site_import.py b/api/tests/unit_tests/services/test_feature_service_site_import.py deleted file mode 100644 index 9415f2d5cc0..00000000000 --- a/api/tests/unit_tests/services/test_feature_service_site_import.py +++ /dev/null @@ -1,31 +0,0 @@ -from collections.abc import Callable -from unittest.mock import Mock - -import pytest - -from enums import CloudPlan, DeploymentEdition -from services.feature_service import FeatureService - - -@pytest.mark.parametrize( - ("edition", "plan", "expected"), - [ - (DeploymentEdition.COMMUNITY, CloudPlan.SANDBOX, True), - (DeploymentEdition.ENTERPRISE, CloudPlan.SANDBOX, True), - (DeploymentEdition.CLOUD, CloudPlan.SANDBOX, False), - (DeploymentEdition.CLOUD, CloudPlan.PROFESSIONAL, True), - ], -) -def test_site_import_entitlement( - monkeypatch: pytest.MonkeyPatch, - config_overrides: Callable[..., None], - edition: DeploymentEdition, - plan: CloudPlan, - expected: bool, -) -> None: - config_overrides(DEPLOYMENT_EDITION=edition) - get_plan = Mock(return_value=plan) - monkeypatch.setattr(FeatureService, "get_workspace_plan", get_plan) - - assert FeatureService.can_import_premium_site_settings("tenant-1") is expected - assert get_plan.call_count == (1 if edition == DeploymentEdition.CLOUD else 0) diff --git a/api/tests/unit_tests/services/test_feature_service_skill.py b/api/tests/unit_tests/services/test_feature_service_skill.py deleted file mode 100644 index 60c15d0bd24..00000000000 --- a/api/tests/unit_tests/services/test_feature_service_skill.py +++ /dev/null @@ -1,17 +0,0 @@ -from collections.abc import Callable - -from enums import DeploymentEdition -from services.entities.feature_entities import FeatureModel -from services.feature_service import FeatureService - - -def test_skill_feature_is_disabled_by_default() -> None: - assert FeatureModel().enable_skill is True - - -def test_skill_feature_follows_env_config(config_overrides: Callable[..., None]) -> None: - config_overrides(DEPLOYMENT_EDITION=DeploymentEdition.COMMUNITY, ENABLE_SKILL=True) - - features = FeatureService.get_features("") - - assert features.enable_skill is True diff --git a/api/tests/unit_tests/services/test_feature_service_trial_models.py b/api/tests/unit_tests/services/test_feature_service_trial_models.py deleted file mode 100644 index 1d2ecdcaa91..00000000000 --- a/api/tests/unit_tests/services/test_feature_service_trial_models.py +++ /dev/null @@ -1,71 +0,0 @@ -from collections.abc import Callable -from unittest.mock import Mock - -import pytest - -from enums import CloudPlan, DeploymentEdition, HostedTrialProvider -from services import feature_service as feature_service_module -from services.feature_service import FeatureService -from services.system_feature_service import SystemFeatureService - - -def test_get_public_system_features_excludes_trial_models() -> None: - result = SystemFeatureService.get_public_system_features().model_dump() - - assert "trial_models" not in result - - -def test_get_trial_models_returns_providers_with_paid_or_trial_enabled( - config_overrides: Callable[..., None], -) -> None: - values: dict[str, object] = {} - for provider in HostedTrialProvider: - values[f"HOSTED_{provider.config_key}_PAID_ENABLED"] = False - values[f"HOSTED_{provider.config_key}_TRIAL_ENABLED"] = False - - values.update( - HOSTED_OPENAI_PAID_ENABLED=True, - HOSTED_OPENAI_TRIAL_ENABLED=True, - HOSTED_XAI_PAID_ENABLED=True, - DEPLOYMENT_EDITION=DeploymentEdition.COMMUNITY, - ) - config_overrides(**values) - - result = FeatureService.get_trial_models("tenant_1") - - assert result == [ - HostedTrialProvider.OPENAI.value, - HostedTrialProvider.X.value, - ] - - -@pytest.mark.parametrize( - ("plan", "expected"), - [ - (CloudPlan.SANDBOX, [HostedTrialProvider.OPENAI.value]), - (CloudPlan.PROFESSIONAL, [HostedTrialProvider.OPENAI.value, HostedTrialProvider.X.value]), - ], -) -def test_get_trial_models_filters_providers_by_workspace_plan( - monkeypatch: pytest.MonkeyPatch, - config_overrides: Callable[..., None], - plan: CloudPlan, - expected: list[str], -) -> None: - values: dict[str, object] = {} - for provider in HostedTrialProvider: - values[f"HOSTED_{provider.config_key}_PAID_ENABLED"] = False - values[f"HOSTED_{provider.config_key}_TRIAL_ENABLED"] = False - values.update( - DEPLOYMENT_EDITION=DeploymentEdition.CLOUD, - HOSTED_OPENAI_TRIAL_ENABLED=True, - HOSTED_XAI_PAID_ENABLED=True, - ) - config_overrides(**values) - get_workspace_plan = Mock(return_value=plan) - monkeypatch.setattr(feature_service_module.FeatureService, "get_workspace_plan", get_workspace_plan) - - result = FeatureService.get_trial_models("tenant_1") - - assert result == expected - get_workspace_plan.assert_called_once_with("tenant_1") diff --git a/api/tests/unit_tests/services/test_feature_service_vector_space.py b/api/tests/unit_tests/services/test_feature_service_vector_space.py deleted file mode 100644 index 4eada6740a6..00000000000 --- a/api/tests/unit_tests/services/test_feature_service_vector_space.py +++ /dev/null @@ -1,49 +0,0 @@ -from collections.abc import Callable -from typing import cast -from unittest.mock import patch - -from enums import CloudPlan, DeploymentEdition -from services.billing_service import BillingInfo -from services.entities.feature_entities import LimitationModel -from services.feature_service import FeatureService - - -def test_get_features_exclude_vector_space_sets_vector_space_to_none(config_overrides: Callable[..., None]) -> None: - tenant_id = "tenant-id" - billing_info = { - "subscription": {"plan": CloudPlan.PROFESSIONAL, "interval": "monthly", "education": False}, - "members": {"size": 1, "limit": 10}, - "apps": {"size": 2, "limit": 20}, - "documents_upload_quota": {"size": 3, "limit": 100}, - "annotation_quota_limit": {"size": 4, "limit": 50}, - "can_replace_logo": True, - "model_load_balancing_enabled": True, - "knowledge_pipeline_publish_enabled": True, - } - - config_overrides( - DEPLOYMENT_EDITION=DeploymentEdition.CLOUD, - CAN_REPLACE_LOGO=False, - MODEL_LB_ENABLED=False, - EDUCATION_ENABLED=False, - ) - with ( - patch("services.feature_service.BillingService.get_info", return_value=billing_info) as get_info, - patch("services.feature_service.BillingService.get_quota_info", return_value={}), - ): - features = FeatureService.get_features(tenant_id, exclude_vector_space=True) - - assert features.vector_space is None - get_info.assert_called_once_with(tenant_id, exclude_vector_space=True) - - -def test_full_features_keep_treating_unknown_vector_usage_as_zero() -> None: - vector_space = LimitationModel() - - FeatureService._fulfill_vector_space_from_billing_info( - vector_space, - cast(BillingInfo, {"vector_space": {"size": 0.0, "limit": 50, "usage_unknown": True}}), - ) - - assert vector_space.size == 0 - assert vector_space.limit == 50 diff --git a/api/tests/unit_tests/services/test_file_service.py b/api/tests/unit_tests/services/test_file_service.py index bd35f39506a..85e3d4215d1 100644 --- a/api/tests/unit_tests/services/test_file_service.py +++ b/api/tests/unit_tests/services/test_file_service.py @@ -2,16 +2,21 @@ import base64 import hashlib import os from collections.abc import Callable, Iterator +from dataclasses import dataclass from datetime import UTC, datetime from unittest.mock import patch +import httpx import pytest -from sqlalchemy import Engine +from sqlalchemy import Connection, Engine, event, select from sqlalchemy.orm import Session, sessionmaker from werkzeug.exceptions import NotFound +from core.file import remote_fetcher from enums import DeploymentEdition +from extensions.ext_storage import storage from extensions.storage.storage_type import StorageType +from graphon.file import helpers as file_helpers from models.base import TypeBase from models.enums import CreatorUserRole from models.model import Account, EndUser, UploadFile @@ -21,508 +26,736 @@ from services.errors.file import ( FileTooLargeError, UnsupportedFileTypeError, ) -from services.file_service import FileService -from tests.unit_tests.model_factories import make_account +from services.file_service import FileService, FileUploadActor, FileUploadResult +from services.remote_file_service import RemoteFileService +from tests.helpers.unit.model_factories import make_account def _account() -> Account: return make_account(account_id="user_id", name="Test Account", email="test@example.com") -class TestFileService: - @pytest.fixture - def sqlite_session_maker(self, sqlite_engine: Engine) -> sessionmaker[Session]: - TypeBase.metadata.create_all(sqlite_engine, tables=[TypeBase.metadata.tables[UploadFile.__tablename__]]) - return sessionmaker(bind=sqlite_engine, expire_on_commit=False) +class TestFileServiceCases: + class TestFileService: + @pytest.fixture + def sqlite_session_maker(self, sqlite_engine: Engine) -> sessionmaker[Session]: + TypeBase.metadata.create_all(sqlite_engine, tables=[TypeBase.metadata.tables[UploadFile.__tablename__]]) + return sessionmaker(bind=sqlite_engine, expire_on_commit=False) - @pytest.fixture - def db_session(self, sqlite_session_maker: sessionmaker[Session]) -> Iterator[Session]: - with sqlite_session_maker() as session: - yield session + @pytest.fixture + def db_session(self, sqlite_session_maker: sessionmaker[Session]) -> Iterator[Session]: + with sqlite_session_maker() as session: + yield session + @pytest.fixture + def file_service(self, sqlite_session_maker: sessionmaker[Session]) -> FileService: + return FileService(session_factory=sqlite_session_maker) + + @staticmethod + def _persist_upload_file( + session: Session, + *, + file_id: str = "file_id", + tenant_id: str = "tenant_id", + extension: str = "txt", + mime_type: str = "text/plain", + key: str = "key", + ) -> UploadFile: + upload_file = UploadFile( + tenant_id=tenant_id, + storage_type=StorageType.LOCAL, + key=key, + name=f"test.{extension}", + size=10, + extension=extension, + mime_type=mime_type, + created_by_role=CreatorUserRole.ACCOUNT, + created_by="user_id", + created_at=datetime(2024, 1, 1, tzinfo=UTC), + used=False, + ) + upload_file.id = file_id + session.add(upload_file) + session.commit() + return upload_file + + def test_init_with_engine(self, sqlite_engine: Engine): + service = FileService(session_factory=sqlite_engine) + assert isinstance(service._session_maker, sessionmaker) + + def test_init_with_sessionmaker(self, sqlite_session_maker: sessionmaker[Session]): + service = FileService(session_factory=sqlite_session_maker) + assert service._session_maker == sqlite_session_maker + + def test_init_invalid_factory(self): + with pytest.raises(AssertionError, match="must be a sessionmaker or an Engine."): + FileService(session_factory="invalid") + + @patch("services.file_service.storage") + @patch("services.file_service.naive_utc_now") + @patch("services.file_service.extract_tenant_id") + @patch("services.file_service.file_helpers.get_signed_file_url") + def test_upload_file_success( + self, mock_get_url, mock_tenant_id, mock_now, mock_storage, file_service: FileService, db_session: Session + ): + # Setup + mock_tenant_id.return_value = "tenant_id" + mock_now.return_value = datetime(2024, 1, 1, tzinfo=UTC) + mock_get_url.return_value = "http://signed-url" + + user = Account(name="Test Account", email="test@example.com") + user.id = "user_id" + content = b"file content" + filename = "test.jpg" + mimetype = "image/jpeg" + + # Execute + result = file_service.upload_file(filename=filename, content=content, mimetype=mimetype, user=user) + + # Assert + assert isinstance(result, UploadFile) + assert result.name == filename + assert result.tenant_id == "tenant_id" + assert result.size == len(content) + assert result.extension == "jpg" + assert result.mime_type == mimetype + assert result.created_by_role == CreatorUserRole.ACCOUNT + assert result.created_by == "user_id" + assert result.hash == hashlib.sha3_256(content).hexdigest() + assert result.source_url == "http://signed-url" + + mock_storage.save.assert_called_once() + persisted = db_session.get(UploadFile, result.id) + assert persisted is not None + assert persisted.hash == result.hash + + @pytest.mark.parametrize("text", ["ASCII text", "包含多字节 UTF-8 文本 🚀"]) + def test_upload_text_uses_utf8_byte_length(self, text: str, file_service: FileService): + with patch("services.file_service.storage") as mock_storage: + result = file_service.upload_text( + text=text, + text_name="test.txt", + user_id="user_id", + tenant_id="tenant_id", + ) + + expected_content = text.encode("utf-8") + assert result.size == len(expected_content) + mock_storage.save.assert_called_once_with(result.key, expected_content) + + def test_upload_file_uses_explicit_resource_tenant(self, file_service: FileService): + user = Account(name="Test Account", email="test@example.com") + user.id = "user-id" + + with ( + patch("services.file_service.storage") as mock_storage, + patch("services.file_service.extract_tenant_id") as mock_extract_tenant_id, + patch("services.file_service.file_helpers.get_signed_file_url"), + ): + result = file_service.upload_file( + filename="test.txt", + content=b"test", + mimetype="text/plain", + user=user, + tenant_id="resource-tenant-id", + ) + + assert result.tenant_id == "resource-tenant-id" + assert mock_storage.save.call_args.args[0].startswith("upload_files/resource-tenant-id/") + mock_extract_tenant_id.assert_not_called() + + def test_upload_file_invalid_characters(self, file_service): + with pytest.raises(ValueError, match="Filename contains invalid characters"): + file_service.upload_file( + filename="invalid/file.txt", content=b"", mimetype="text/plain", user=_account() + ) + + def test_upload_file_long_filename(self, file_service: FileService, db_session: Session): + # Setup + long_name = "a" * 210 + ".txt" + user = Account(name="Test Account", email="test@example.com") + user.id = "user_id" + + with ( + patch("services.file_service.storage"), + patch("services.file_service.extract_tenant_id") as mock_tenant, + patch("services.file_service.file_helpers.get_signed_file_url"), + ): + mock_tenant.return_value = "tenant" + result = file_service.upload_file(filename=long_name, content=b"test", mimetype="text/plain", user=user) + assert len(result.name) <= 205 # 200 + . + extension + assert result.name.endswith(".txt") + assert db_session.get(UploadFile, result.id) is not None + + def test_upload_file_blocked_extension(self, file_service, config_overrides: Callable[..., None]): + config_overrides(inner_UPLOAD_FILE_EXTENSION_BLACKLIST="exe") + with pytest.raises(BlockedFileExtensionError): + file_service.upload_file( + filename="test.exe", content=b"", mimetype="application/octet-stream", user=_account() + ) + + def test_upload_file_unsupported_type_for_datasets(self, file_service): + with pytest.raises(UnsupportedFileTypeError): + file_service.upload_file( + filename="test.jpg", content=b"", mimetype="image/jpeg", user=_account(), source="datasets" + ) + + def test_upload_file_too_large(self, file_service, config_overrides: Callable[..., None]): + # 16MB file for an image with 15MB limit + content = b"a" * (16 * 1024 * 1024) + config_overrides(UPLOAD_IMAGE_FILE_SIZE_LIMIT=15) + with pytest.raises(FileTooLargeError): + file_service.upload_file(filename="test.jpg", content=content, mimetype="image/jpeg", user=_account()) + + def test_upload_file_end_user(self, file_service: FileService, db_session: Session): + user = EndUser( + id="end_user_id", + ) + + with ( + patch("services.file_service.storage"), + patch("services.file_service.extract_tenant_id") as mock_tenant, + patch("services.file_service.file_helpers.get_signed_file_url"), + ): + mock_tenant.return_value = "tenant" + result = file_service.upload_file( + filename="test.txt", content=b"test", mimetype="text/plain", user=user + ) + assert result.created_by_role == CreatorUserRole.END_USER + assert db_session.get(UploadFile, result.id) is not None + + def test_is_file_size_within_limit(self, config_overrides: Callable[..., None]): + config_overrides( + UPLOAD_IMAGE_FILE_SIZE_LIMIT=10, + UPLOAD_VIDEO_FILE_SIZE_LIMIT=20, + UPLOAD_AUDIO_FILE_SIZE_LIMIT=30, + UPLOAD_FILE_SIZE_LIMIT=5, + ) + # Image + assert FileService.is_file_size_within_limit(extension="jpg", file_size=10 * 1024 * 1024) is True + assert FileService.is_file_size_within_limit(extension="png", file_size=11 * 1024 * 1024) is False + + # Video + assert FileService.is_file_size_within_limit(extension="mp4", file_size=20 * 1024 * 1024) is True + assert FileService.is_file_size_within_limit(extension="avi", file_size=21 * 1024 * 1024) is False + + # Audio + assert FileService.is_file_size_within_limit(extension="mp3", file_size=30 * 1024 * 1024) is True + assert FileService.is_file_size_within_limit(extension="wav", file_size=31 * 1024 * 1024) is False + + # Default + assert FileService.is_file_size_within_limit(extension="txt", file_size=5 * 1024 * 1024) is True + assert FileService.is_file_size_within_limit(extension="pdf", file_size=6 * 1024 * 1024) is False + assert ( + FileService.is_file_size_within_limit(extension="txt", file_size=0, default_file_size_limit=0) is True + ) + assert ( + FileService.is_file_size_within_limit(extension="txt", file_size=1, default_file_size_limit=0) is False + ) + assert ( + FileService.is_file_size_within_limit( + extension="pdf", + file_size=6 * 1024 * 1024, + default_file_size_limit=7, + ) + is True + ) + assert ( + FileService.is_file_size_within_limit( + extension="pdf", + file_size=8 * 1024 * 1024, + default_file_size_limit=7, + ) + is False + ) + + # Media-specific limits are not affected by the knowledge document override. + assert ( + FileService.is_file_size_within_limit( + extension="jpg", + file_size=11 * 1024 * 1024, + default_file_size_limit=100, + ) + is False + ) + + def test_file_size_limit(self, config_overrides: Callable[..., None]): + config_overrides( + UPLOAD_IMAGE_FILE_SIZE_LIMIT=10, + UPLOAD_VIDEO_FILE_SIZE_LIMIT=20, + UPLOAD_AUDIO_FILE_SIZE_LIMIT=30, + UPLOAD_FILE_SIZE_LIMIT=5, + ) + + assert FileService.file_size_limit(extension="jpg") == 10 * 1024 * 1024 + assert FileService.file_size_limit(extension="mp4") == 20 * 1024 * 1024 + assert FileService.file_size_limit(extension="mp3") == 30 * 1024 * 1024 + assert FileService.file_size_limit(extension="txt") == 5 * 1024 * 1024 + assert FileService.file_size_limit(extension="txt", default_file_size_limit=None) == 5 * 1024 * 1024 + assert FileService.file_size_limit(extension="txt", default_file_size_limit=0) == 0 + assert FileService.file_size_limit(extension="jpg", default_file_size_limit=0) == 10 * 1024 * 1024 + assert FileService.file_size_limit(extension="txt", default_file_size_limit=7) == 7 * 1024 * 1024 + + def test_get_file_base64_success(self, file_service: FileService, db_session: Session): + self._persist_upload_file(db_session, key="test_key") + + with patch("services.file_service.storage") as mock_storage: + mock_storage.load_once.return_value = b"test content" + + # Execute + result = file_service.get_file_base64("file_id") + + # Assert + assert result == base64.b64encode(b"test content").decode() + mock_storage.load_once.assert_called_once_with("test_key") + + def test_get_file_base64_not_found(self, file_service: FileService): + with pytest.raises(NotFound, match="File not found"): + file_service.get_file_base64("non_existent") + + def test_get_file_presigned_url_success( + self, file_service: FileService, db_session: Session, config_overrides: Callable[..., None] + ): + config_overrides(FILES_ACCESS_TIMEOUT=300) + self._persist_upload_file( + db_session, + extension="png", + mime_type="image/png", + key="upload_files/tenant_id/icon.png", + ) + + with ( + patch("services.file_service.storage") as mock_storage, + ): + mock_storage.generate_presigned_url.return_value = "https://s3.example.com/icon.png?signature=test" + + result = file_service.get_file_presigned_url(file_id="file_id", tenant_id="tenant_id") + + assert result == "https://s3.example.com/icon.png?signature=test" + mock_storage.generate_presigned_url.assert_called_once_with( + "upload_files/tenant_id/icon.png", + expires_in=300, + content_type="image/png", + ) + + def test_get_file_presigned_url_not_found(self, file_service: FileService): + with pytest.raises(NotFound, match="File not found"): + file_service.get_file_presigned_url(file_id="file_id", tenant_id="tenant_id") + + def test_get_icon_url_uses_direct_storage_url_for_cloud_s3( + self, file_service: FileService, config_overrides: Callable[..., None] + ): + config_overrides(DEPLOYMENT_EDITION=DeploymentEdition.CLOUD, STORAGE_TYPE=StorageType.S3) + with patch.object(file_service, "get_file_presigned_url", return_value="direct-url") as get_presigned_url: + result = file_service.get_icon_url("file_id", "tenant_id") + + assert result == "direct-url" + get_presigned_url.assert_called_once_with(file_id="file_id", tenant_id="tenant_id") + + def test_get_icon_url_maps_missing_cloud_file_to_service_error( + self, file_service: FileService, config_overrides: Callable[..., None] + ) -> None: + config_overrides(DEPLOYMENT_EDITION=DeploymentEdition.CLOUD, STORAGE_TYPE=StorageType.S3) + with ( + patch.object(file_service, "get_file_presigned_url", side_effect=NotFound("File not found")), + pytest.raises(FileNotExistsError, match="File reference not found"), + ): + file_service.get_icon_url("file_id", "tenant_id") + + @pytest.mark.parametrize( + ("deployment_edition", "storage_type"), + [ + (DeploymentEdition.COMMUNITY, StorageType.S3), + (DeploymentEdition.CLOUD, StorageType.LOCAL), + ], + ) + def test_get_icon_url_uses_preview_url_outside_cloud_s3( + self, + file_service: FileService, + db_session: Session, + deployment_edition: DeploymentEdition, + storage_type: StorageType, + config_overrides: Callable[..., None], + ): + self._persist_upload_file(db_session) + config_overrides(DEPLOYMENT_EDITION=deployment_edition, STORAGE_TYPE=storage_type) + with patch("services.file_service.file_helpers.get_signed_file_url", return_value="preview-url") as get_url: + result = file_service.get_icon_url("file_id", "tenant_id") + + assert result == "preview-url" + get_url.assert_called_once_with(upload_file_id="file_id") + + @pytest.mark.parametrize( + ("deployment_edition", "storage_type"), + [ + (DeploymentEdition.COMMUNITY, StorageType.S3), + (DeploymentEdition.CLOUD, StorageType.LOCAL), + ], + ) + def test_get_icon_url_rejects_missing_file_outside_cloud_s3( + self, + file_service: FileService, + deployment_edition: DeploymentEdition, + storage_type: StorageType, + config_overrides: Callable[..., None], + ) -> None: + config_overrides(DEPLOYMENT_EDITION=deployment_edition, STORAGE_TYPE=storage_type) + + with pytest.raises(FileNotExistsError, match="File reference not found"): + file_service.get_icon_url("file_id", "tenant_id") + + def test_get_icon_url_rejects_cross_tenant_file_outside_cloud_s3( + self, + file_service: FileService, + db_session: Session, + config_overrides: Callable[..., None], + ) -> None: + self._persist_upload_file(db_session, tenant_id="other_tenant_id") + config_overrides(DEPLOYMENT_EDITION=DeploymentEdition.COMMUNITY, STORAGE_TYPE=StorageType.LOCAL) + + with pytest.raises(FileNotExistsError, match="File reference not found"): + file_service.get_icon_url("file_id", "tenant_id") + + def test_upload_text_success(self, file_service: FileService, db_session: Session): + # Setup + text = "sample text" + text_name = "test.txt" + user_id = "user_id" + tenant_id = "tenant_id" + + with patch("services.file_service.storage") as mock_storage: + # Execute + result = file_service.upload_text(text, text_name, user_id, tenant_id) + + # Assert + assert result.name == text_name + assert result.size == len(text) + assert result.tenant_id == tenant_id + assert result.created_by == user_id + assert result.used is True + assert result.extension == "txt" + mock_storage.save.assert_called_once() + assert db_session.get(UploadFile, result.id) is not None + + def test_upload_text_long_name(self, file_service: FileService, db_session: Session): + long_name = "a" * 210 + with patch("services.file_service.storage"): + result = file_service.upload_text("text", long_name, "user", "tenant") + assert len(result.name) == 200 + assert db_session.get(UploadFile, result.id) is not None + + def test_get_file_preview_success(self, file_service: FileService, db_session: Session): + self._persist_upload_file(db_session, extension="pdf", mime_type="application/pdf") + + with patch("services.file_service.ExtractProcessor.load_from_upload_file") as mock_extract: + mock_extract.return_value = "Extracted text content" + + # Execute + result = file_service.get_file_preview("file_id", "tenant_id") + + # Assert + assert result == "Extracted text content" + + def test_get_file_preview_not_found(self, file_service: FileService): + with pytest.raises(NotFound, match="File not found"): + file_service.get_file_preview("non_existent", "tenant_id") + + def test_get_file_preview_unsupported_type(self, file_service: FileService, db_session: Session): + self._persist_upload_file(db_session, extension="exe", mime_type="application/octet-stream") + with pytest.raises(UnsupportedFileTypeError): + file_service.get_file_preview("file_id", "tenant_id") + + def test_get_file_content_success(self, file_service: FileService, db_session: Session): + self._persist_upload_file(db_session) + + with patch("services.file_service.storage") as mock_storage: + mock_storage.load.return_value = b"hello world" + result = file_service.get_file_content("file_id") + assert result == "hello world" + + def test_get_file_content_not_found(self, file_service: FileService): + with pytest.raises(NotFound, match="File not found"): + file_service.get_file_content("file_id") + + def test_delete_file_success(self, file_service: FileService, db_session: Session): + self._persist_upload_file(db_session) + + with patch("services.file_service.storage") as mock_storage: + file_service.delete_file("file_id") + mock_storage.delete.assert_called_once_with("key") + db_session.expire_all() + assert db_session.get(UploadFile, "file_id") is None + + def test_delete_file_not_found(self, file_service: FileService): + file_service.delete_file("file_id") + # Should return without doing anything + + def test_get_upload_files_by_ids_empty(self, db_session: Session): + result = FileService.get_upload_files_by_ids("tenant_id", [], session=db_session) + assert result == {} + + def test_get_upload_file_by_id_scopes_to_tenant(self, db_session: Session) -> None: + upload_file = self._persist_upload_file(db_session) + + assert FileService.get_upload_file_by_id("tenant_id", "file_id", session=db_session) == upload_file + assert FileService.get_upload_file_by_id("other_tenant_id", "file_id", session=db_session) is None + + def test_get_upload_files_by_ids(self, db_session: Session): + upload_file = self._persist_upload_file(db_session, file_id="550e8400-e29b-41d4-a716-446655440000") + self._persist_upload_file( + db_session, + file_id="550e8400-e29b-41d4-a716-446655440001", + tenant_id="other-tenant", + ) + + result = FileService.get_upload_files_by_ids( + "tenant_id", + ["550e8400-e29b-41d4-a716-446655440000", "550e8400-e29b-41d4-a716-446655440001"], + session=db_session, + ) + assert result["550e8400-e29b-41d4-a716-446655440000"] == upload_file + + def test_sanitize_zip_entry_name(self): + assert FileService._sanitize_zip_entry_name("path/to/file.txt") == "file.txt" + assert FileService._sanitize_zip_entry_name("../../../etc/passwd") == "passwd" + assert FileService._sanitize_zip_entry_name(" ") == "file" + assert FileService._sanitize_zip_entry_name("a\\b") == "a_b" + + def test_dedupe_zip_entry_name(self): + used = {"a.txt"} + assert FileService._dedupe_zip_entry_name("b.txt", used) == "b.txt" + assert FileService._dedupe_zip_entry_name("a.txt", used) == "a (1).txt" + used.add("a (1).txt") + assert FileService._dedupe_zip_entry_name("a.txt", used) == "a (2).txt" + + def test_build_upload_files_zip_tempfile(self, db_session: Session): + upload_file = self._persist_upload_file(db_session) + + with ( + patch("services.file_service.storage") as mock_storage, + patch("services.file_service.os.remove") as mock_remove, + ): + mock_storage.load.return_value = [b"chunk1", b"chunk2"] + + with FileService.build_upload_files_zip_tempfile(upload_files=[upload_file]) as tmp_path: + assert os.path.exists(tmp_path) + + mock_remove.assert_called_once() + + +_TENANT_ID = "11111111-1111-1111-1111-111111111111" +_ACTOR_ID = "22222222-2222-2222-2222-222222222222" +_REMOTE_URL = "https://example.com/report.txt" + + +@dataclass +class _UploadHarness: + files: FileService + stored: dict[str, bytes] + active_transactions: set[Connection] + io_events: list[str] + database_events: list[str] + + +class TestFileUploadActorCases: + @staticmethod @pytest.fixture - def file_service(self, sqlite_session_maker: sessionmaker[Session]) -> FileService: - return FileService(session_factory=sqlite_session_maker) + def upload_harness( + sqlite_engine: Engine, + sqlite_session_factory: sessionmaker[Session], + monkeypatch: pytest.MonkeyPatch, + ) -> Iterator[_UploadHarness]: + harness = _UploadHarness( + files=FileService(session_factory=sqlite_session_factory), + stored={}, + active_transactions=set(), + io_events=[], + database_events=[], + ) + + def begin(connection: Connection) -> None: + harness.active_transactions.add(connection) + harness.database_events.append("begin") + + def end(connection: Connection) -> None: + harness.active_transactions.remove(connection) + harness.database_events.append("end") + + def save(key: str, content: bytes) -> None: + assert not harness.active_transactions + harness.io_events.append("storage") + harness.stored[key] = content + + def signed_url(upload_file_id: str) -> str: + assert not harness.active_transactions + harness.io_events.append("sign") + return f"https://files.example.com/{upload_file_id}?sign=test" + + event.listen(sqlite_engine, "begin", begin) + event.listen(sqlite_engine, "commit", end) + event.listen(sqlite_engine, "rollback", end) + monkeypatch.setattr(storage, "save", save) + monkeypatch.setattr(file_helpers, "get_signed_file_url", signed_url) + try: + yield harness + finally: + event.remove(sqlite_engine, "begin", begin) + event.remove(sqlite_engine, "commit", end) + event.remove(sqlite_engine, "rollback", end) @staticmethod - def _persist_upload_file( - session: Session, - *, - file_id: str = "file_id", - tenant_id: str = "tenant_id", - extension: str = "txt", - mime_type: str = "text/plain", - key: str = "key", - ) -> UploadFile: - upload_file = UploadFile( - tenant_id=tenant_id, - storage_type=StorageType.LOCAL, - key=key, - name=f"test.{extension}", - size=10, - extension=extension, - mime_type=mime_type, - created_by_role=CreatorUserRole.ACCOUNT, - created_by="user_id", - created_at=datetime(2024, 1, 1, tzinfo=UTC), - used=False, + @pytest.mark.parametrize("role", [CreatorUserRole.ACCOUNT, CreatorUserRole.END_USER]) + @pytest.mark.parametrize("source_url", ["", _REMOTE_URL]) + def test_actor_upload_preserves_owner_creator_and_source_url_without_extra_queries( + upload_harness: _UploadHarness, + sqlite_session_factory: sessionmaker[Session], + role: CreatorUserRole, + source_url: str, + ) -> None: + actor = FileUploadActor(id=_ACTOR_ID, creator_role=role) + content = "文件内容".encode() + + result = upload_harness.files.upload_file_for_actor( + actor=actor, + resource_tenant_id=_TENANT_ID, + filename="report.TXT", + content=content, + mimetype="text/plain", + source_url=source_url, ) - upload_file.id = file_id - session.add(upload_file) - session.commit() - return upload_file - def test_init_with_engine(self, sqlite_engine: Engine): - service = FileService(session_factory=sqlite_engine) - assert isinstance(service._session_maker, sessionmaker) - - def test_init_with_sessionmaker(self, sqlite_session_maker: sessionmaker[Session]): - service = FileService(session_factory=sqlite_session_maker) - assert service._session_maker == sqlite_session_maker - - def test_init_invalid_factory(self): - with pytest.raises(AssertionError, match="must be a sessionmaker or an Engine."): - FileService(session_factory="invalid") - - @patch("services.file_service.storage") - @patch("services.file_service.naive_utc_now") - @patch("services.file_service.extract_tenant_id") - @patch("services.file_service.file_helpers.get_signed_file_url") - def test_upload_file_success( - self, mock_get_url, mock_tenant_id, mock_now, mock_storage, file_service: FileService, db_session: Session - ): - # Setup - mock_tenant_id.return_value = "tenant_id" - mock_now.return_value = datetime(2024, 1, 1, tzinfo=UTC) - mock_get_url.return_value = "http://signed-url" - - user = Account(name="Test Account", email="test@example.com") - user.id = "user_id" - content = b"file content" - filename = "test.jpg" - mimetype = "image/jpeg" - - # Execute - result = file_service.upload_file(filename=filename, content=content, mimetype=mimetype, user=user) - - # Assert - assert isinstance(result, UploadFile) - assert result.name == filename - assert result.tenant_id == "tenant_id" + assert isinstance(result, FileUploadResult) + assert result.name == "report.TXT" assert result.size == len(content) - assert result.extension == "jpg" - assert result.mime_type == mimetype - assert result.created_by_role == CreatorUserRole.ACCOUNT - assert result.created_by == "user_id" - assert result.hash == hashlib.sha3_256(content).hexdigest() - assert result.source_url == "http://signed-url" + assert result.extension == "txt" + assert result.mime_type == "text/plain" + assert result.created_by == actor.id + assert isinstance(result.created_at, datetime) + assert result.tenant_id == _TENANT_ID + assert result.source_url == (source_url or f"https://files.example.com/{result.id}?sign=test") + assert upload_harness.database_events == ["begin", "end"] + assert not upload_harness.active_transactions + assert upload_harness.io_events == (["storage"] if source_url else ["storage", "sign"]) - mock_storage.save.assert_called_once() - persisted = db_session.get(UploadFile, result.id) - assert persisted is not None - assert persisted.hash == result.hash + with sqlite_session_factory() as session: + persisted = session.get(UploadFile, result.id) + assert persisted is not None + assert persisted.created_by_role == role + assert persisted.created_by == actor.id + assert persisted.tenant_id == _TENANT_ID + assert persisted.source_url == source_url + assert persisted.hash == hashlib.sha3_256(content).hexdigest() + assert persisted.key.startswith(f"upload_files/{_TENANT_ID}/") + assert upload_harness.stored == {persisted.key: content} - @pytest.mark.parametrize("text", ["ASCII text", "包含多字节 UTF-8 文本 🚀"]) - def test_upload_text_uses_utf8_byte_length(self, text: str, file_service: FileService): - with patch("services.file_service.storage") as mock_storage: - result = file_service.upload_text( - text=text, - text_name="test.txt", - user_id="user_id", - tenant_id="tenant_id", + @staticmethod + @pytest.mark.parametrize("head_status", [httpx.codes.OK, httpx.codes.METHOD_NOT_ALLOWED]) + def test_remote_actor_upload_fetches_before_transaction_and_retains_original_url( + upload_harness: _UploadHarness, + sqlite_session_factory: sessionmaker[Session], + monkeypatch: pytest.MonkeyPatch, + head_status: int, + ) -> None: + actor = FileUploadActor(id=_ACTOR_ID, creator_role=CreatorUserRole.ACCOUNT) + content = b"remote file content" + requests: list[tuple[str, dict[str, object]]] = [] + + def make_request(method: str, *, url: str, **kwargs: object) -> httpx.Response: + assert not upload_harness.active_transactions + assert url == _REMOTE_URL + requests.append((method, kwargs)) + upload_harness.io_events.append(method) + return httpx.Response( + head_status if method == "HEAD" else httpx.codes.OK, + content=content if method == "GET" else b"", + headers={"Content-Type": "text/plain", "Content-Length": str(len(content))}, + request=httpx.Request(method, url), ) - expected_content = text.encode("utf-8") - assert result.size == len(expected_content) - mock_storage.save.assert_called_once_with(result.key, expected_content) + monkeypatch.setattr(remote_fetcher, "make_request", make_request) - def test_upload_file_uses_explicit_resource_tenant(self, file_service: FileService): - user = Account(name="Test Account", email="test@example.com") - user.id = "user-id" + result = RemoteFileService(files=upload_harness.files).upload_from_url( + url=_REMOTE_URL, user=actor, tenant_id=_TENANT_ID + ) - with ( - patch("services.file_service.storage") as mock_storage, - patch("services.file_service.extract_tenant_id") as mock_extract_tenant_id, - patch("services.file_service.file_helpers.get_signed_file_url"), - ): - result = file_service.upload_file( - filename="test.txt", - content=b"test", + assert result.name == "report.txt" + assert result.size == len(content) + assert result.extension == "txt" + assert result.mime_type == "text/plain" + assert result.created_by == actor.id + assert result.url == f"https://files.example.com/{result.id}?sign=test" + assert requests == [ + ("HEAD", {}), + ("GET", {} if head_status == httpx.codes.OK else {"timeout": 3, "follow_redirects": True}), + ] + assert upload_harness.io_events == ["HEAD", "GET", "storage", "sign"] + assert upload_harness.database_events == ["begin", "end"] + assert not upload_harness.active_transactions + + with sqlite_session_factory() as session: + persisted = session.scalars(select(UploadFile)).one() + assert persisted.id == result.id + assert persisted.created_by == actor.id + assert persisted.created_by_role == CreatorUserRole.ACCOUNT + assert persisted.tenant_id == _TENANT_ID + assert persisted.source_url == _REMOTE_URL + assert persisted.key.startswith(f"upload_files/{_TENANT_ID}/") + assert upload_harness.stored == {persisted.key: content} + + @staticmethod + def test_same_actor_can_upload_to_separately_admitted_tenants( + upload_harness: _UploadHarness, + sqlite_session_factory: sessionmaker[Session], + ) -> None: + actor = FileUploadActor(id=_ACTOR_ID, creator_role=CreatorUserRole.ACCOUNT) + tenant_ids = (_TENANT_ID, "33333333-3333-3333-3333-333333333333") + + results = [ + upload_harness.files.upload_file_for_actor( + actor=actor, + resource_tenant_id=tenant_id, + filename="report.txt", + content=tenant_id.encode(), mimetype="text/plain", - user=user, - tenant_id="resource-tenant-id", ) + for tenant_id in tenant_ids + ] - assert result.tenant_id == "resource-tenant-id" - assert mock_storage.save.call_args.args[0].startswith("upload_files/resource-tenant-id/") - mock_extract_tenant_id.assert_not_called() + assert [result.tenant_id for result in results] == list(tenant_ids) + assert len({result.id for result in results}) == 2 + with sqlite_session_factory() as session: + for result, tenant_id in zip(results, tenant_ids, strict=True): + persisted = session.get(UploadFile, result.id) + assert persisted is not None + assert persisted.tenant_id == tenant_id + assert persisted.created_by == actor.id + assert persisted.created_by_role == actor.creator_role + assert persisted.key.startswith(f"upload_files/{tenant_id}/") + assert upload_harness.stored[persisted.key] == tenant_id.encode() - def test_upload_file_invalid_characters(self, file_service): - with pytest.raises(ValueError, match="Filename contains invalid characters"): - file_service.upload_file(filename="invalid/file.txt", content=b"", mimetype="text/plain", user=_account()) - - def test_upload_file_long_filename(self, file_service: FileService, db_session: Session): - # Setup - long_name = "a" * 210 + ".txt" - user = Account(name="Test Account", email="test@example.com") - user.id = "user_id" - - with ( - patch("services.file_service.storage"), - patch("services.file_service.extract_tenant_id") as mock_tenant, - patch("services.file_service.file_helpers.get_signed_file_url"), - ): - mock_tenant.return_value = "tenant" - result = file_service.upload_file(filename=long_name, content=b"test", mimetype="text/plain", user=user) - assert len(result.name) <= 205 # 200 + . + extension - assert result.name.endswith(".txt") - assert db_session.get(UploadFile, result.id) is not None - - def test_upload_file_blocked_extension(self, file_service, config_overrides: Callable[..., None]): - config_overrides(inner_UPLOAD_FILE_EXTENSION_BLACKLIST="exe") - with pytest.raises(BlockedFileExtensionError): - file_service.upload_file( - filename="test.exe", content=b"", mimetype="application/octet-stream", user=_account() - ) - - def test_upload_file_unsupported_type_for_datasets(self, file_service): - with pytest.raises(UnsupportedFileTypeError): - file_service.upload_file( - filename="test.jpg", content=b"", mimetype="image/jpeg", user=_account(), source="datasets" - ) - - def test_upload_file_too_large(self, file_service, config_overrides: Callable[..., None]): - # 16MB file for an image with 15MB limit - content = b"a" * (16 * 1024 * 1024) - config_overrides(UPLOAD_IMAGE_FILE_SIZE_LIMIT=15) - with pytest.raises(FileTooLargeError): - file_service.upload_file(filename="test.jpg", content=content, mimetype="image/jpeg", user=_account()) - - def test_upload_file_end_user(self, file_service: FileService, db_session: Session): - user = EndUser( - id="end_user_id", - ) - - with ( - patch("services.file_service.storage"), - patch("services.file_service.extract_tenant_id") as mock_tenant, - patch("services.file_service.file_helpers.get_signed_file_url"), - ): - mock_tenant.return_value = "tenant" - result = file_service.upload_file(filename="test.txt", content=b"test", mimetype="text/plain", user=user) - assert result.created_by_role == CreatorUserRole.END_USER - assert db_session.get(UploadFile, result.id) is not None - - def test_is_file_size_within_limit(self, config_overrides: Callable[..., None]): - config_overrides( - UPLOAD_IMAGE_FILE_SIZE_LIMIT=10, - UPLOAD_VIDEO_FILE_SIZE_LIMIT=20, - UPLOAD_AUDIO_FILE_SIZE_LIMIT=30, - UPLOAD_FILE_SIZE_LIMIT=5, - ) - # Image - assert FileService.is_file_size_within_limit(extension="jpg", file_size=10 * 1024 * 1024) is True - assert FileService.is_file_size_within_limit(extension="png", file_size=11 * 1024 * 1024) is False - - # Video - assert FileService.is_file_size_within_limit(extension="mp4", file_size=20 * 1024 * 1024) is True - assert FileService.is_file_size_within_limit(extension="avi", file_size=21 * 1024 * 1024) is False - - # Audio - assert FileService.is_file_size_within_limit(extension="mp3", file_size=30 * 1024 * 1024) is True - assert FileService.is_file_size_within_limit(extension="wav", file_size=31 * 1024 * 1024) is False - - # Default - assert FileService.is_file_size_within_limit(extension="txt", file_size=5 * 1024 * 1024) is True - assert FileService.is_file_size_within_limit(extension="pdf", file_size=6 * 1024 * 1024) is False - assert FileService.is_file_size_within_limit(extension="txt", file_size=0, default_file_size_limit=0) is True - assert FileService.is_file_size_within_limit(extension="txt", file_size=1, default_file_size_limit=0) is False - assert ( - FileService.is_file_size_within_limit( - extension="pdf", - file_size=6 * 1024 * 1024, - default_file_size_limit=7, - ) - is True - ) - assert ( - FileService.is_file_size_within_limit( - extension="pdf", - file_size=8 * 1024 * 1024, - default_file_size_limit=7, - ) - is False - ) - - # Media-specific limits are not affected by the knowledge document override. - assert ( - FileService.is_file_size_within_limit( - extension="jpg", - file_size=11 * 1024 * 1024, - default_file_size_limit=100, - ) - is False - ) - - def test_file_size_limit(self, config_overrides: Callable[..., None]): - config_overrides( - UPLOAD_IMAGE_FILE_SIZE_LIMIT=10, - UPLOAD_VIDEO_FILE_SIZE_LIMIT=20, - UPLOAD_AUDIO_FILE_SIZE_LIMIT=30, - UPLOAD_FILE_SIZE_LIMIT=5, - ) - - assert FileService.file_size_limit(extension="jpg") == 10 * 1024 * 1024 - assert FileService.file_size_limit(extension="mp4") == 20 * 1024 * 1024 - assert FileService.file_size_limit(extension="mp3") == 30 * 1024 * 1024 - assert FileService.file_size_limit(extension="txt") == 5 * 1024 * 1024 - assert FileService.file_size_limit(extension="txt", default_file_size_limit=None) == 5 * 1024 * 1024 - assert FileService.file_size_limit(extension="txt", default_file_size_limit=0) == 0 - assert FileService.file_size_limit(extension="jpg", default_file_size_limit=0) == 10 * 1024 * 1024 - assert FileService.file_size_limit(extension="txt", default_file_size_limit=7) == 7 * 1024 * 1024 - - def test_get_file_base64_success(self, file_service: FileService, db_session: Session): - self._persist_upload_file(db_session, key="test_key") - - with patch("services.file_service.storage") as mock_storage: - mock_storage.load_once.return_value = b"test content" - - # Execute - result = file_service.get_file_base64("file_id") - - # Assert - assert result == base64.b64encode(b"test content").decode() - mock_storage.load_once.assert_called_once_with("test_key") - - def test_get_file_base64_not_found(self, file_service: FileService): - with pytest.raises(NotFound, match="File not found"): - file_service.get_file_base64("non_existent") - - def test_get_file_presigned_url_success( - self, file_service: FileService, db_session: Session, config_overrides: Callable[..., None] - ): - config_overrides(FILES_ACCESS_TIMEOUT=300) - self._persist_upload_file( - db_session, - extension="png", - mime_type="image/png", - key="upload_files/tenant_id/icon.png", - ) - - with ( - patch("services.file_service.storage") as mock_storage, - ): - mock_storage.generate_presigned_url.return_value = "https://s3.example.com/icon.png?signature=test" - - result = file_service.get_file_presigned_url(file_id="file_id", tenant_id="tenant_id") - - assert result == "https://s3.example.com/icon.png?signature=test" - mock_storage.generate_presigned_url.assert_called_once_with( - "upload_files/tenant_id/icon.png", - expires_in=300, - content_type="image/png", - ) - - def test_get_file_presigned_url_not_found(self, file_service: FileService): - with pytest.raises(NotFound, match="File not found"): - file_service.get_file_presigned_url(file_id="file_id", tenant_id="tenant_id") - - def test_get_icon_url_uses_direct_storage_url_for_cloud_s3( - self, file_service: FileService, config_overrides: Callable[..., None] - ): - config_overrides(DEPLOYMENT_EDITION=DeploymentEdition.CLOUD, STORAGE_TYPE=StorageType.S3) - with patch.object(file_service, "get_file_presigned_url", return_value="direct-url") as get_presigned_url: - result = file_service.get_icon_url("file_id", "tenant_id") - - assert result == "direct-url" - get_presigned_url.assert_called_once_with(file_id="file_id", tenant_id="tenant_id") - - def test_get_icon_url_maps_missing_cloud_file_to_service_error( - self, file_service: FileService, config_overrides: Callable[..., None] + @staticmethod + @pytest.mark.parametrize("remote", [False, True], ids=["local", "remote"]) + def test_actor_upload_without_target_tenant_fails_before_external_io( + upload_harness: _UploadHarness, + monkeypatch: pytest.MonkeyPatch, + remote: bool, ) -> None: - config_overrides(DEPLOYMENT_EDITION=DeploymentEdition.CLOUD, STORAGE_TYPE=StorageType.S3) - with ( - patch.object(file_service, "get_file_presigned_url", side_effect=NotFound("File not found")), - pytest.raises(FileNotExistsError, match="File reference not found"), - ): - file_service.get_icon_url("file_id", "tenant_id") + actor = FileUploadActor(id=_ACTOR_ID, creator_role=CreatorUserRole.ACCOUNT) - @pytest.mark.parametrize( - ("deployment_edition", "storage_type"), - [ - (DeploymentEdition.COMMUNITY, StorageType.S3), - (DeploymentEdition.CLOUD, StorageType.LOCAL), - ], - ) - def test_get_icon_url_uses_preview_url_outside_cloud_s3( - self, - file_service: FileService, - db_session: Session, - deployment_edition: DeploymentEdition, - storage_type: StorageType, - config_overrides: Callable[..., None], - ): - self._persist_upload_file(db_session) - config_overrides(DEPLOYMENT_EDITION=deployment_edition, STORAGE_TYPE=storage_type) - with patch("services.file_service.file_helpers.get_signed_file_url", return_value="preview-url") as get_url: - result = file_service.get_icon_url("file_id", "tenant_id") + def unexpected_request(*_args: object, **_kwargs: object) -> httpx.Response: + pytest.fail("Missing target tenant must be rejected before fetching the remote file") - assert result == "preview-url" - get_url.assert_called_once_with(upload_file_id="file_id") + monkeypatch.setattr(remote_fetcher, "make_request", unexpected_request) - @pytest.mark.parametrize( - ("deployment_edition", "storage_type"), - [ - (DeploymentEdition.COMMUNITY, StorageType.S3), - (DeploymentEdition.CLOUD, StorageType.LOCAL), - ], - ) - def test_get_icon_url_rejects_missing_file_outside_cloud_s3( - self, - file_service: FileService, - deployment_edition: DeploymentEdition, - storage_type: StorageType, - config_overrides: Callable[..., None], - ) -> None: - config_overrides(DEPLOYMENT_EDITION=deployment_edition, STORAGE_TYPE=storage_type) + if remote: + with pytest.raises(TypeError, match="tenant_id"): + RemoteFileService(files=upload_harness.files).upload_from_url(url=_REMOTE_URL, user=actor) + else: + with pytest.raises(TypeError, match="tenant_id"): + upload_harness.files.upload_file( + filename="report.txt", content=b"file content", mimetype="text/plain", user=actor + ) - with pytest.raises(FileNotExistsError, match="File reference not found"): - file_service.get_icon_url("file_id", "tenant_id") - - def test_get_icon_url_rejects_cross_tenant_file_outside_cloud_s3( - self, - file_service: FileService, - db_session: Session, - config_overrides: Callable[..., None], - ) -> None: - self._persist_upload_file(db_session, tenant_id="other_tenant_id") - config_overrides(DEPLOYMENT_EDITION=DeploymentEdition.COMMUNITY, STORAGE_TYPE=StorageType.LOCAL) - - with pytest.raises(FileNotExistsError, match="File reference not found"): - file_service.get_icon_url("file_id", "tenant_id") - - def test_upload_text_success(self, file_service: FileService, db_session: Session): - # Setup - text = "sample text" - text_name = "test.txt" - user_id = "user_id" - tenant_id = "tenant_id" - - with patch("services.file_service.storage") as mock_storage: - # Execute - result = file_service.upload_text(text, text_name, user_id, tenant_id) - - # Assert - assert result.name == text_name - assert result.size == len(text) - assert result.tenant_id == tenant_id - assert result.created_by == user_id - assert result.used is True - assert result.extension == "txt" - mock_storage.save.assert_called_once() - assert db_session.get(UploadFile, result.id) is not None - - def test_upload_text_long_name(self, file_service: FileService, db_session: Session): - long_name = "a" * 210 - with patch("services.file_service.storage"): - result = file_service.upload_text("text", long_name, "user", "tenant") - assert len(result.name) == 200 - assert db_session.get(UploadFile, result.id) is not None - - def test_get_file_preview_success(self, file_service: FileService, db_session: Session): - self._persist_upload_file(db_session, extension="pdf", mime_type="application/pdf") - - with patch("services.file_service.ExtractProcessor.load_from_upload_file") as mock_extract: - mock_extract.return_value = "Extracted text content" - - # Execute - result = file_service.get_file_preview("file_id", "tenant_id") - - # Assert - assert result == "Extracted text content" - - def test_get_file_preview_not_found(self, file_service: FileService): - with pytest.raises(NotFound, match="File not found"): - file_service.get_file_preview("non_existent", "tenant_id") - - def test_get_file_preview_unsupported_type(self, file_service: FileService, db_session: Session): - self._persist_upload_file(db_session, extension="exe", mime_type="application/octet-stream") - with pytest.raises(UnsupportedFileTypeError): - file_service.get_file_preview("file_id", "tenant_id") - - def test_get_file_content_success(self, file_service: FileService, db_session: Session): - self._persist_upload_file(db_session) - - with patch("services.file_service.storage") as mock_storage: - mock_storage.load.return_value = b"hello world" - result = file_service.get_file_content("file_id") - assert result == "hello world" - - def test_get_file_content_not_found(self, file_service: FileService): - with pytest.raises(NotFound, match="File not found"): - file_service.get_file_content("file_id") - - def test_delete_file_success(self, file_service: FileService, db_session: Session): - self._persist_upload_file(db_session) - - with patch("services.file_service.storage") as mock_storage: - file_service.delete_file("file_id") - mock_storage.delete.assert_called_once_with("key") - db_session.expire_all() - assert db_session.get(UploadFile, "file_id") is None - - def test_delete_file_not_found(self, file_service: FileService): - file_service.delete_file("file_id") - # Should return without doing anything - - def test_get_upload_files_by_ids_empty(self, db_session: Session): - result = FileService.get_upload_files_by_ids("tenant_id", [], session=db_session) - assert result == {} - - def test_get_upload_file_by_id_scopes_to_tenant(self, db_session: Session) -> None: - upload_file = self._persist_upload_file(db_session) - - assert FileService.get_upload_file_by_id("tenant_id", "file_id", session=db_session) == upload_file - assert FileService.get_upload_file_by_id("other_tenant_id", "file_id", session=db_session) is None - - def test_get_upload_files_by_ids(self, db_session: Session): - upload_file = self._persist_upload_file(db_session, file_id="550e8400-e29b-41d4-a716-446655440000") - self._persist_upload_file( - db_session, - file_id="550e8400-e29b-41d4-a716-446655440001", - tenant_id="other-tenant", - ) - - result = FileService.get_upload_files_by_ids( - "tenant_id", - ["550e8400-e29b-41d4-a716-446655440000", "550e8400-e29b-41d4-a716-446655440001"], - session=db_session, - ) - assert result["550e8400-e29b-41d4-a716-446655440000"] == upload_file - - def test_sanitize_zip_entry_name(self): - assert FileService._sanitize_zip_entry_name("path/to/file.txt") == "file.txt" - assert FileService._sanitize_zip_entry_name("../../../etc/passwd") == "passwd" - assert FileService._sanitize_zip_entry_name(" ") == "file" - assert FileService._sanitize_zip_entry_name("a\\b") == "a_b" - - def test_dedupe_zip_entry_name(self): - used = {"a.txt"} - assert FileService._dedupe_zip_entry_name("b.txt", used) == "b.txt" - assert FileService._dedupe_zip_entry_name("a.txt", used) == "a (1).txt" - used.add("a (1).txt") - assert FileService._dedupe_zip_entry_name("a.txt", used) == "a (2).txt" - - def test_build_upload_files_zip_tempfile(self, db_session: Session): - upload_file = self._persist_upload_file(db_session) - - with ( - patch("services.file_service.storage") as mock_storage, - patch("services.file_service.os.remove") as mock_remove, - ): - mock_storage.load.return_value = [b"chunk1", b"chunk2"] - - with FileService.build_upload_files_zip_tempfile(upload_files=[upload_file]) as tmp_path: - assert os.path.exists(tmp_path) - - mock_remove.assert_called_once() + assert upload_harness.io_events == [] + assert upload_harness.database_events == [] + assert upload_harness.stored == {} diff --git a/api/tests/unit_tests/services/test_file_upload_actor.py b/api/tests/unit_tests/services/test_file_upload_actor.py deleted file mode 100644 index 9efe37cf4d9..00000000000 --- a/api/tests/unit_tests/services/test_file_upload_actor.py +++ /dev/null @@ -1,232 +0,0 @@ -import hashlib -from collections.abc import Iterator -from dataclasses import dataclass -from datetime import datetime - -import httpx -import pytest -from sqlalchemy import Connection, Engine, event, select -from sqlalchemy.orm import Session, sessionmaker - -from core.file import remote_fetcher -from extensions.ext_storage import storage -from graphon.file import helpers as file_helpers -from models.enums import CreatorUserRole -from models.model import UploadFile -from services.file_service import FileService, FileUploadActor, FileUploadResult -from services.remote_file_service import RemoteFileService - -_TENANT_ID = "11111111-1111-1111-1111-111111111111" -_ACTOR_ID = "22222222-2222-2222-2222-222222222222" -_REMOTE_URL = "https://example.com/report.txt" - - -@dataclass -class _UploadHarness: - files: FileService - stored: dict[str, bytes] - active_transactions: set[Connection] - io_events: list[str] - database_events: list[str] - - -@pytest.fixture -def upload_harness( - sqlite_engine: Engine, - sqlite_session_factory: sessionmaker[Session], - monkeypatch: pytest.MonkeyPatch, -) -> Iterator[_UploadHarness]: - harness = _UploadHarness( - files=FileService(session_factory=sqlite_session_factory), - stored={}, - active_transactions=set(), - io_events=[], - database_events=[], - ) - - def begin(connection: Connection) -> None: - harness.active_transactions.add(connection) - harness.database_events.append("begin") - - def end(connection: Connection) -> None: - harness.active_transactions.remove(connection) - harness.database_events.append("end") - - def save(key: str, content: bytes) -> None: - assert not harness.active_transactions - harness.io_events.append("storage") - harness.stored[key] = content - - def signed_url(upload_file_id: str) -> str: - assert not harness.active_transactions - harness.io_events.append("sign") - return f"https://files.example.com/{upload_file_id}?sign=test" - - event.listen(sqlite_engine, "begin", begin) - event.listen(sqlite_engine, "commit", end) - event.listen(sqlite_engine, "rollback", end) - monkeypatch.setattr(storage, "save", save) - monkeypatch.setattr(file_helpers, "get_signed_file_url", signed_url) - try: - yield harness - finally: - event.remove(sqlite_engine, "begin", begin) - event.remove(sqlite_engine, "commit", end) - event.remove(sqlite_engine, "rollback", end) - - -@pytest.mark.parametrize("role", [CreatorUserRole.ACCOUNT, CreatorUserRole.END_USER]) -@pytest.mark.parametrize("source_url", ["", _REMOTE_URL]) -def test_actor_upload_preserves_owner_creator_and_source_url_without_extra_queries( - upload_harness: _UploadHarness, - sqlite_session_factory: sessionmaker[Session], - role: CreatorUserRole, - source_url: str, -) -> None: - actor = FileUploadActor(id=_ACTOR_ID, creator_role=role) - content = "文件内容".encode() - - result = upload_harness.files.upload_file_for_actor( - actor=actor, - resource_tenant_id=_TENANT_ID, - filename="report.TXT", - content=content, - mimetype="text/plain", - source_url=source_url, - ) - - assert isinstance(result, FileUploadResult) - assert result.name == "report.TXT" - assert result.size == len(content) - assert result.extension == "txt" - assert result.mime_type == "text/plain" - assert result.created_by == actor.id - assert isinstance(result.created_at, datetime) - assert result.tenant_id == _TENANT_ID - assert result.source_url == (source_url or f"https://files.example.com/{result.id}?sign=test") - assert upload_harness.database_events == ["begin", "end"] - assert not upload_harness.active_transactions - assert upload_harness.io_events == (["storage"] if source_url else ["storage", "sign"]) - - with sqlite_session_factory() as session: - persisted = session.get(UploadFile, result.id) - assert persisted is not None - assert persisted.created_by_role == role - assert persisted.created_by == actor.id - assert persisted.tenant_id == _TENANT_ID - assert persisted.source_url == source_url - assert persisted.hash == hashlib.sha3_256(content).hexdigest() - assert persisted.key.startswith(f"upload_files/{_TENANT_ID}/") - assert upload_harness.stored == {persisted.key: content} - - -@pytest.mark.parametrize("head_status", [httpx.codes.OK, httpx.codes.METHOD_NOT_ALLOWED]) -def test_remote_actor_upload_fetches_before_transaction_and_retains_original_url( - upload_harness: _UploadHarness, - sqlite_session_factory: sessionmaker[Session], - monkeypatch: pytest.MonkeyPatch, - head_status: int, -) -> None: - actor = FileUploadActor(id=_ACTOR_ID, creator_role=CreatorUserRole.ACCOUNT) - content = b"remote file content" - requests: list[tuple[str, dict[str, object]]] = [] - - def make_request(method: str, *, url: str, **kwargs: object) -> httpx.Response: - assert not upload_harness.active_transactions - assert url == _REMOTE_URL - requests.append((method, kwargs)) - upload_harness.io_events.append(method) - return httpx.Response( - head_status if method == "HEAD" else httpx.codes.OK, - content=content if method == "GET" else b"", - headers={"Content-Type": "text/plain", "Content-Length": str(len(content))}, - request=httpx.Request(method, url), - ) - - monkeypatch.setattr(remote_fetcher, "make_request", make_request) - - result = RemoteFileService(files=upload_harness.files).upload_from_url( - url=_REMOTE_URL, user=actor, tenant_id=_TENANT_ID - ) - - assert result.name == "report.txt" - assert result.size == len(content) - assert result.extension == "txt" - assert result.mime_type == "text/plain" - assert result.created_by == actor.id - assert result.url == f"https://files.example.com/{result.id}?sign=test" - assert requests == [ - ("HEAD", {}), - ("GET", {} if head_status == httpx.codes.OK else {"timeout": 3, "follow_redirects": True}), - ] - assert upload_harness.io_events == ["HEAD", "GET", "storage", "sign"] - assert upload_harness.database_events == ["begin", "end"] - assert not upload_harness.active_transactions - - with sqlite_session_factory() as session: - persisted = session.scalars(select(UploadFile)).one() - assert persisted.id == result.id - assert persisted.created_by == actor.id - assert persisted.created_by_role == CreatorUserRole.ACCOUNT - assert persisted.tenant_id == _TENANT_ID - assert persisted.source_url == _REMOTE_URL - assert persisted.key.startswith(f"upload_files/{_TENANT_ID}/") - assert upload_harness.stored == {persisted.key: content} - - -def test_same_actor_can_upload_to_separately_admitted_tenants( - upload_harness: _UploadHarness, - sqlite_session_factory: sessionmaker[Session], -) -> None: - actor = FileUploadActor(id=_ACTOR_ID, creator_role=CreatorUserRole.ACCOUNT) - tenant_ids = (_TENANT_ID, "33333333-3333-3333-3333-333333333333") - - results = [ - upload_harness.files.upload_file_for_actor( - actor=actor, - resource_tenant_id=tenant_id, - filename="report.txt", - content=tenant_id.encode(), - mimetype="text/plain", - ) - for tenant_id in tenant_ids - ] - - assert [result.tenant_id for result in results] == list(tenant_ids) - assert len({result.id for result in results}) == 2 - with sqlite_session_factory() as session: - for result, tenant_id in zip(results, tenant_ids, strict=True): - persisted = session.get(UploadFile, result.id) - assert persisted is not None - assert persisted.tenant_id == tenant_id - assert persisted.created_by == actor.id - assert persisted.created_by_role == actor.creator_role - assert persisted.key.startswith(f"upload_files/{tenant_id}/") - assert upload_harness.stored[persisted.key] == tenant_id.encode() - - -@pytest.mark.parametrize("remote", [False, True], ids=["local", "remote"]) -def test_actor_upload_without_target_tenant_fails_before_external_io( - upload_harness: _UploadHarness, - monkeypatch: pytest.MonkeyPatch, - remote: bool, -) -> None: - actor = FileUploadActor(id=_ACTOR_ID, creator_role=CreatorUserRole.ACCOUNT) - - def unexpected_request(*_args: object, **_kwargs: object) -> httpx.Response: - pytest.fail("Missing target tenant must be rejected before fetching the remote file") - - monkeypatch.setattr(remote_fetcher, "make_request", unexpected_request) - - if remote: - with pytest.raises(TypeError, match="tenant_id"): - RemoteFileService(files=upload_harness.files).upload_from_url(url=_REMOTE_URL, user=actor) - else: - with pytest.raises(TypeError, match="tenant_id"): - upload_harness.files.upload_file( - filename="report.txt", content=b"file content", mimetype="text/plain", user=actor - ) - - assert upload_harness.io_events == [] - assert upload_harness.database_events == [] - assert upload_harness.stored == {} diff --git a/api/tests/unit_tests/services/test_human_input_service.py b/api/tests/unit_tests/services/test_human_input_service.py index a8aab964af5..1489c106d01 100644 --- a/api/tests/unit_tests/services/test_human_input_service.py +++ b/api/tests/unit_tests/services/test_human_input_service.py @@ -44,7 +44,7 @@ from services.human_input_service import ( HumanInputService, InvalidFormDataError, ) -from tests.unit_tests.config_override import apply_config_overrides +from tests.helpers.unit.config_override import apply_config_overrides def _make_app(mode: AppMode) -> App: diff --git a/api/tests/unit_tests/services/test_knowledge_fs_app_execution.py b/api/tests/unit_tests/services/test_knowledge_fs_app_execution.py deleted file mode 100644 index 47fdd8bb54b..00000000000 --- a/api/tests/unit_tests/services/test_knowledge_fs_app_execution.py +++ /dev/null @@ -1,247 +0,0 @@ -from __future__ import annotations - -from datetime import UTC, datetime -from typing import override - -import pytest -from pydantic import JsonValue -from sqlalchemy.orm import Session, sessionmaker - -from core.app.entities.app_invoke_entities import DifyRunContext, InvokeFrom, UserFrom -from enums import DeploymentEdition -from models.knowledge_fs import ( - AppKnowledgeFSSpaceJoin, - KnowledgeFSAppSpaceJoinStatus, - KnowledgeFSAppSpaceJoinType, - KnowledgeFSAuthorizationRevision, - KnowledgeFSCapabilityIssuanceAudit, - KnowledgeFSCapabilityIssuanceReservation, - KnowledgeFSControlSpace, - KnowledgeFSControlSpaceState, - KnowledgeFSExternalAccessPolicy, - KnowledgeFSLifecycleOutbox, -) -from services.knowledge_fs.app_admission_service import KnowledgeFSAppAdmissionError, KnowledgeFSAppAdmissionService -from services.knowledge_fs.app_execution_capability import ( - KnowledgeFSAppExecutionCapabilityService, - KnowledgeResourceRef, -) -from services.knowledge_fs.capability_broker import KnowledgeFSIssuedProductCapability -from services.knowledge_fs.product_dto import KnowledgeFSResearchTaskCreatePayload -from services.knowledge_fs.product_remote import KnowledgeFSRemoteJSONRequest, UnavailableKnowledgeFSProductRemote -from tests.unit_tests.config_override import apply_config_overrides - -pytestmark = pytest.mark.parametrize( - "sqlite_session", - [ - ( - KnowledgeFSControlSpace, - KnowledgeFSExternalAccessPolicy, - KnowledgeFSAuthorizationRevision, - KnowledgeFSCapabilityIssuanceAudit, - KnowledgeFSCapabilityIssuanceReservation, - KnowledgeFSLifecycleOutbox, - AppKnowledgeFSSpaceJoin, - ) - ], - indirect=True, -) - - -@pytest.fixture(autouse=True) -def _enable_knowledge_fs(monkeypatch: pytest.MonkeyPatch) -> None: - apply_config_overrides(monkeypatch, DEPLOYMENT_EDITION=DeploymentEdition.CLOUD, KNOWLEDGE_FS_ENABLED=True) - - -class Broker: - def __init__(self) -> None: - self.calls: list[dict[str, object]] = [] - - def issue_app(self, **kwargs: object) -> KnowledgeFSIssuedProductCapability: - self.calls.append(kwargs) - return KnowledgeFSIssuedProductCapability( - token="capability", - expires_at=datetime(2030, 1, 1, tzinfo=UTC), - operation_id="createResearchTask", - knowledge_space_id="space-1", - knowledge_space_revision=3, - trace_id="trace-issued", - ) - - -class Remote(UnavailableKnowledgeFSProductRemote): - def __init__(self) -> None: - self.calls: list[KnowledgeFSRemoteJSONRequest] = [] - - @override - def execute_json(self, request: KnowledgeFSRemoteJSONRequest) -> dict[str, JsonValue]: - self.calls.append(request) - return { - "id": "task-1", - "knowledgeSpaceId": "space-1", - "query": "question", - "cost": {}, - "stage": "queued", - "metadata": {}, - "createdAt": 1.0, - "updatedAt": 1.0, - } - - -def _run_context(*, app_id: str) -> DifyRunContext: - return DifyRunContext( - tenant_id="tenant-1", - app_id=app_id, - user_id="user-1", - user_from=UserFrom.END_USER, - invoke_from=InvokeFrom.WEB_APP, - trace_session_id="trace-run", - ) - - -def _seed( - session: Session, - *, - join_status: KnowledgeFSAppSpaceJoinStatus = KnowledgeFSAppSpaceJoinStatus.ACTIVE, - caller_kind: KnowledgeFSAppSpaceJoinType = KnowledgeFSAppSpaceJoinType.AGENT, - channel_enabled: bool = True, -) -> KnowledgeFSControlSpace: - space = KnowledgeFSControlSpace( - tenant_id="tenant-1", - owner_account_id="owner-1", - provisioning_key="provision-1", - knowledge_space_id="space-1", - knowledge_space_revision=3, - state=KnowledgeFSControlSpaceState.ACTIVE, - ) - session.add_all( - [ - space, - KnowledgeFSExternalAccessPolicy( - tenant_id="tenant-1", - control_space_id=space.id, - agent_enabled=channel_enabled if caller_kind is KnowledgeFSAppSpaceJoinType.AGENT else False, - workflow_enabled=channel_enabled if caller_kind is KnowledgeFSAppSpaceJoinType.WORKFLOW else False, - ), - KnowledgeFSAuthorizationRevision( - tenant_id="tenant-1", - control_space_id=space.id, - ), - AppKnowledgeFSSpaceJoin( - tenant_id="tenant-1", - control_space_id=space.id, - app_id="app-1", - join_type=caller_kind, - status=join_status, - ), - ] - ) - session.commit() - return space - - -@pytest.mark.parametrize( - ("join_status", "caller_kind", "channel_enabled", "context_app_id"), - [ - ( - KnowledgeFSAppSpaceJoinStatus.REVOKED, - KnowledgeFSAppSpaceJoinType.AGENT, - True, - "app-1", - ), - ( - KnowledgeFSAppSpaceJoinStatus.ACTIVE, - KnowledgeFSAppSpaceJoinType.AGENT, - False, - "app-1", - ), - ( - KnowledgeFSAppSpaceJoinStatus.ACTIVE, - KnowledgeFSAppSpaceJoinType.WORKFLOW, - False, - "app-1", - ), - ( - KnowledgeFSAppSpaceJoinStatus.ACTIVE, - KnowledgeFSAppSpaceJoinType.AGENT, - True, - "app-2", - ), - ], - ids=("binding-disabled", "agent-channel-off", "workflow-channel-off", "cross-app-space"), -) -def test_execution_fails_closed_before_capability_or_remote_io( - sqlite_session: Session, - join_status: KnowledgeFSAppSpaceJoinStatus, - caller_kind: KnowledgeFSAppSpaceJoinType, - channel_enabled: bool, - context_app_id: str, -) -> None: - space = _seed( - sqlite_session, - join_status=join_status, - caller_kind=caller_kind, - channel_enabled=channel_enabled, - ) - broker = Broker() - remote = Remote() - service = KnowledgeFSAppExecutionCapabilityService( # type: ignore[arg-type] - admission=KnowledgeFSAppAdmissionService(sessionmaker(bind=sqlite_session.get_bind(), expire_on_commit=False)), - broker=broker, - remote=remote, - ) - - with pytest.raises(KnowledgeFSAppAdmissionError): - service.create_research_task( - run_context=_run_context(app_id=context_app_id), - caller_kind=caller_kind, - resource=KnowledgeResourceRef(kind="knowledge_fs", control_space_id=space.id), - payload=KnowledgeFSResearchTaskCreatePayload(query="question"), - ) - - assert broker.calls == [] - assert remote.calls == [] - - -@pytest.mark.parametrize( - "caller_kind", - [KnowledgeFSAppSpaceJoinType.AGENT, KnowledgeFSAppSpaceJoinType.WORKFLOW], -) -def test_execution_issues_capability_then_performs_bounded_remote_io( - sqlite_session: Session, - caller_kind: KnowledgeFSAppSpaceJoinType, -) -> None: - space = _seed(sqlite_session, caller_kind=caller_kind) - broker = Broker() - remote = Remote() - service = KnowledgeFSAppExecutionCapabilityService( # type: ignore[arg-type] - admission=KnowledgeFSAppAdmissionService(sessionmaker(bind=sqlite_session.get_bind(), expire_on_commit=False)), - broker=broker, - remote=remote, - ) - - response = service.create_research_task( - run_context=_run_context(app_id="app-1"), - caller_kind=caller_kind, - resource=KnowledgeResourceRef(kind="knowledge_fs", control_space_id=space.id), - payload=KnowledgeFSResearchTaskCreatePayload(query="question"), - ) - - assert response.id == "task-1" - assert len(broker.calls) == 1 - assert remote.calls == [ - KnowledgeFSRemoteJSONRequest( - operation_id="createResearchTask", - method="POST", - path="/research-tasks", - namespace_id="tenant-1", - knowledge_space_id="space-1", - capability_token="capability", - trace_id="trace-issued", - payload={ - "query": "question", - "metadata": {}, - "knowledgeSpaceId": "space-1", - }, - ) - ] diff --git a/api/tests/unit_tests/services/test_knowledge_fs_app_execution_capability.py b/api/tests/unit_tests/services/test_knowledge_fs_app_execution_capability.py deleted file mode 100644 index 89b7b915015..00000000000 --- a/api/tests/unit_tests/services/test_knowledge_fs_app_execution_capability.py +++ /dev/null @@ -1,652 +0,0 @@ -from __future__ import annotations - -import base64 -import json -from datetime import UTC, datetime -from typing import override - -import pytest -from pydantic import JsonValue, ValidationError - -from core.app.entities.app_invoke_entities import DifyRunContext, InvokeFrom, UserFrom -from enums import DeploymentEdition -from models.knowledge_fs import KnowledgeFSAppSpaceJoinType -from services.knowledge_fs import app_execution_capability -from services.knowledge_fs.app_admission_service import KnowledgeFSAppPrincipalProfile -from services.knowledge_fs.app_execution_capability import ( - KnowledgeFSAppExecutionCapabilityService, - KnowledgeResourceRef, -) -from services.knowledge_fs.capability_broker import KnowledgeFSIssuedProductCapability -from services.knowledge_fs.product_dto import ( - KnowledgeFSResearchTaskCreatePayload, - KnowledgeFSRetrievalTestPayload, - KnowledgeFSWorkflowFailedRetrievalCapturePayload, -) -from services.knowledge_fs.product_remote import ( - KNOWLEDGE_FS_QUERY_IMAGE_GRANTS_HEADER, - KnowledgeFSOperationUnavailableError, - KnowledgeFSRemoteJSONRequest, - UnavailableKnowledgeFSProductRemote, -) -from tests.unit_tests.config_override import apply_config_overrides - - -@pytest.fixture(autouse=True) -def _enable_knowledge_fs(monkeypatch: pytest.MonkeyPatch) -> None: - apply_config_overrides(monkeypatch, DEPLOYMENT_EDITION=DeploymentEdition.CLOUD, KNOWLEDGE_FS_ENABLED=True) - - -class Admission: - def __init__(self) -> None: - self.calls: list[dict[str, object]] = [] - self.profile = KnowledgeFSAppPrincipalProfile( - tenant_id="tenant-1", - control_space_id="control-1", - app_id="app-1", - join_id="join-1", - caller_kind=KnowledgeFSAppSpaceJoinType.WORKFLOW, - action="read", - knowledge_space_id="space-1", - knowledge_space_revision=2, - membership_epoch=1, - space_acl_epoch=1, - external_access_epoch=1, - content_policy_revision=1, - ) - - def admit(self, **kwargs: object) -> KnowledgeFSAppPrincipalProfile: - self.calls.append(kwargs) - return self.profile - - -class Broker: - def __init__(self) -> None: - self.calls: list[dict[str, object]] = [] - - def issue_app(self, **kwargs: object) -> KnowledgeFSIssuedProductCapability: - self.calls.append(kwargs) - return KnowledgeFSIssuedProductCapability( - token="token", - expires_at=datetime(2030, 1, 1, tzinfo=UTC), - operation_id=str(kwargs["operation_id"]), - knowledge_space_id="space-1", - knowledge_space_revision=2, - trace_id="trace-1", - ) - - -class Remote(UnavailableKnowledgeFSProductRemote): - def __init__(self) -> None: - self.calls: list[KnowledgeFSRemoteJSONRequest] = [] - - @override - def execute_json(self, request: KnowledgeFSRemoteJSONRequest) -> dict[str, JsonValue]: - self.calls.append(request) - if request.operation_id == "captureAgentKnowledgeInvestigation": - return {"investigationId": "investigation-1", "failedQueries": []} - if request.operation_id == "captureWorkflowFailedRetrieval": - return { - "failedQueryId": "019fac9f-bfb0-75ee-9af5-252ebafbac1c", - "verdict": "retrieval-miss", - "badCaseId": "019fac9f-bfb0-75ee-9af5-252ebafbac1d", - } - if request.operation_id == "listMetadataFields": - return { - "items": [ - { - "count": 3, - "createdAt": "2026-01-01T00:00:00Z", - "id": "field-1", - "name": "department", - "rowVersion": 1, - "type": "string", - "updatedAt": "2026-01-01T00:00:00Z", - } - ], - "nextCursor": "cursor-2", - } - if request.operation_id == "retrieveEvidence": - return { - "items": [ - { - "citation": { - "artifactHash": "a" * 64, - "documentAssetId": "document-1", - "documentVersion": 1, - "sectionPath": ["Camera"], - }, - "nodeId": "node-1", - "projectionIds": ["projection-1"], - "score": 0.8, - "sources": ["dense"], - "text": "Camera evidence", - } - ], - "metrics": {"denseCandidates": 1, "totalMs": 10}, - "mode": "fast", - "traceId": "trace-1", - } - return { - "id": "research-1", - "knowledgeSpaceId": "space-1", - "query": "Compare the evidence", - "cost": {}, - "stage": "queued", - "metadata": {}, - "createdAt": 1.0, - "updatedAt": 1.0, - } - - -class FailingRemote(Remote): - @override - def execute_json(self, request: KnowledgeFSRemoteJSONRequest) -> dict[str, JsonValue]: - self.calls.append(request) - raise RuntimeError("remote failed") - - -def test_app_execution_capability_always_admits_binding_before_broker_issuance() -> None: - admission = Admission() - broker = Broker() - service = KnowledgeFSAppExecutionCapabilityService( # type: ignore[arg-type] - admission=admission, - broker=broker, - remote=Remote(), - ) - - issued = service.issue( - tenant_id="tenant-1", - app_id="app-1", - control_space_id="control-1", - caller_kind=KnowledgeFSAppSpaceJoinType.WORKFLOW, - operation_id="createResearchTask", - resource_id="research-1", - trace_id="trace-1", - ) - - assert issued.token == "token" - assert admission.calls == [ - { - "tenant_id": "tenant-1", - "app_id": "app-1", - "control_space_id": "control-1", - "caller_kind": KnowledgeFSAppSpaceJoinType.WORKFLOW, - "operation_id": "createResearchTask", - } - ] - assert broker.calls == [ - { - "profile": admission.profile, - "operation_id": "createResearchTask", - "resource_id": "research-1", - "trace_id": "trace-1", - } - ] - - -def test_disabled_feature_blocks_app_issuance_and_investigation_capture_before_io( - monkeypatch: pytest.MonkeyPatch, -) -> None: - admission = Admission() - broker = Broker() - remote = Remote() - service = KnowledgeFSAppExecutionCapabilityService( # type: ignore[arg-type] - admission=admission, - broker=broker, - remote=remote, - ) - previously_issued = KnowledgeFSIssuedProductCapability( - token="token", - expires_at=datetime(2030, 1, 1, tzinfo=UTC), - operation_id="captureAgentKnowledgeInvestigation", - knowledge_space_id="space-1", - knowledge_space_revision=2, - trace_id="trace-1", - ) - apply_config_overrides(monkeypatch, KNOWLEDGE_FS_ENABLED=False) - - with pytest.raises(KnowledgeFSOperationUnavailableError, match="disabled"): - service.issue( - tenant_id="tenant-1", - app_id="app-1", - control_space_id="control-1", - caller_kind=KnowledgeFSAppSpaceJoinType.AGENT, - operation_id="retrieveEvidence", - ) - with pytest.raises(KnowledgeFSOperationUnavailableError, match="disabled"): - service.create_research_task( - run_context=DifyRunContext( - tenant_id="tenant-1", - app_id="app-1", - user_id="user-1", - user_from=UserFrom.END_USER, - invoke_from=InvokeFrom.WEB_APP, - ), - caller_kind=KnowledgeFSAppSpaceJoinType.AGENT, - resource=KnowledgeResourceRef(kind="knowledge_fs", control_space_id="control-1"), - payload=KnowledgeFSResearchTaskCreatePayload(query="Compare the evidence"), - ) - with pytest.raises(KnowledgeFSOperationUnavailableError, match="disabled"): - service.run_retrieval( - run_context=DifyRunContext( - tenant_id="tenant-1", - app_id="app-1", - user_id="user-1", - user_from=UserFrom.END_USER, - invoke_from=InvokeFrom.WEB_APP, - ), - caller_kind=KnowledgeFSAppSpaceJoinType.WORKFLOW, - resource=KnowledgeResourceRef(kind="knowledge_fs", control_space_id="control-1"), - payload=KnowledgeFSRetrievalTestPayload.model_validate( - { - "includeText": True, - "mode": "fast", - "query": "camera", - "queryId": "10000000-0000-4000-8000-000000000010", - } - ), - ) - with pytest.raises(KnowledgeFSOperationUnavailableError, match="disabled"): - service.capture_agent_investigation(issued=previously_issued, tenant_id="tenant-1", payload={}) - - assert admission.calls == [] - assert broker.calls == [] - assert remote.calls == [] - - -def test_typed_knowledge_resource_ref_rejects_dataset_and_extra_authority_fields() -> None: - with pytest.raises(ValidationError): - KnowledgeResourceRef.model_validate( - { - "kind": "dataset", - "control_space_id": "control-1", - } - ) - - with pytest.raises(ValidationError): - KnowledgeResourceRef.model_validate( - { - "kind": "knowledge_fs", - "control_space_id": "control-1", - "dataset_ids": ["control-1"], - } - ) - - with pytest.raises(ValidationError): - KnowledgeResourceRef.model_validate( - { - "kind": "knowledge_fs", - "control_space_id": " ", - } - ) - - -def test_create_research_task_fails_closed_when_product_operation_is_unavailable( - monkeypatch: pytest.MonkeyPatch, -) -> None: - admission = Admission() - broker = Broker() - remote = Remote() - service = KnowledgeFSAppExecutionCapabilityService( # type: ignore[arg-type] - admission=admission, - broker=broker, - remote=remote, - ) - monkeypatch.setattr(app_execution_capability, "is_product_operation_ready", lambda _operation_id: False) - - with pytest.raises(KnowledgeFSOperationUnavailableError, match="creation is unavailable"): - service.create_research_task( - run_context=DifyRunContext( - tenant_id="tenant-1", - app_id="app-1", - user_id="user-1", - user_from=UserFrom.END_USER, - invoke_from=InvokeFrom.WEB_APP, - ), - caller_kind=KnowledgeFSAppSpaceJoinType.AGENT, - resource=KnowledgeResourceRef(kind="knowledge_fs", control_space_id="control-1"), - payload=KnowledgeFSResearchTaskCreatePayload(query="Compare the evidence"), - ) - - assert admission.calls == [] - assert broker.calls == [] - assert remote.calls == [] - - -def test_create_research_task_uses_only_dify_run_context_identity_then_calls_kfs() -> None: - admission = Admission() - broker = Broker() - remote = Remote() - service = KnowledgeFSAppExecutionCapabilityService( # type: ignore[arg-type] - admission=admission, - broker=broker, - remote=remote, - ) - run_context = DifyRunContext( - tenant_id="tenant-from-run-context", - app_id="app-from-run-context", - user_id="user-1", - user_from=UserFrom.END_USER, - invoke_from=InvokeFrom.WEB_APP, - trace_session_id="trace-from-run-context", - ) - - result = service.create_research_task( - run_context=run_context, - caller_kind=KnowledgeFSAppSpaceJoinType.AGENT, - resource=KnowledgeResourceRef(kind="knowledge_fs", control_space_id="control-1"), - payload=KnowledgeFSResearchTaskCreatePayload(query="Compare the evidence"), - ) - - assert result.id == "research-1" - assert admission.calls == [ - { - "tenant_id": "tenant-from-run-context", - "app_id": "app-from-run-context", - "control_space_id": "control-1", - "caller_kind": KnowledgeFSAppSpaceJoinType.AGENT, - "operation_id": "createResearchTask", - } - ] - assert broker.calls == [ - { - "profile": admission.profile, - "operation_id": "createResearchTask", - "resource_id": None, - "trace_id": "trace-from-run-context", - } - ] - assert remote.calls == [ - KnowledgeFSRemoteJSONRequest( - operation_id="createResearchTask", - method="POST", - path="/research-tasks", - namespace_id="tenant-from-run-context", - knowledge_space_id="space-1", - capability_token="token", - trace_id="trace-1", - payload={ - "query": "Compare the evidence", - "metadata": {}, - "knowledgeSpaceId": "space-1", - }, - ) - ] - - -def test_create_research_task_propagates_remote_io_failure() -> None: - service = KnowledgeFSAppExecutionCapabilityService( # type: ignore[arg-type] - admission=Admission(), - broker=Broker(), - remote=FailingRemote(), - ) - - with pytest.raises(RuntimeError, match="remote failed"): - service.create_research_task( - run_context=DifyRunContext( - tenant_id="tenant-1", - app_id="app-1", - user_id="user-1", - user_from=UserFrom.END_USER, - invoke_from=InvokeFrom.WEB_APP, - ), - caller_kind=KnowledgeFSAppSpaceJoinType.AGENT, - resource=KnowledgeResourceRef(kind="knowledge_fs", control_space_id="control-1"), - payload=KnowledgeFSResearchTaskCreatePayload(query="Compare the evidence"), - ) - - -def test_run_retrieval_issues_read_capability_and_calls_bounded_product_operation() -> None: - admission = Admission() - broker = Broker() - remote = Remote() - service = KnowledgeFSAppExecutionCapabilityService( # type: ignore[arg-type] - admission=admission, - broker=broker, - remote=remote, - ) - run_context = DifyRunContext( - tenant_id="tenant-1", - app_id="app-1", - user_id="user-1", - user_from=UserFrom.ACCOUNT, - invoke_from=InvokeFrom.DEBUGGER, - trace_session_id="trace-from-workflow", - ) - - result = service.run_retrieval( - run_context=run_context, - caller_kind=KnowledgeFSAppSpaceJoinType.WORKFLOW, - resource=KnowledgeResourceRef(kind="knowledge_fs", control_space_id="control-1"), - payload=KnowledgeFSRetrievalTestPayload.model_validate( - { - "includeText": True, - "mode": "fast", - "query": "camera", - "queryId": "10000000-0000-4000-8000-000000000010", - "queryImages": [ - { - "accessGrant": "short-lived-grant", - "uploadFileId": "00000000-0000-4000-8000-000000000001", - } - ], - } - ), - ) - - assert result.items[0].text == "Camera evidence" - assert admission.calls[0]["operation_id"] == "retrieveEvidence" - assert broker.calls[0]["operation_id"] == "retrieveEvidence" - assert len(remote.calls) == 1 - request = remote.calls[0] - assert request.payload == { - "includeText": True, - "mode": "fast", - "query": "camera", - "queryId": "10000000-0000-4000-8000-000000000010", - "queryImages": [{"uploadFileId": "00000000-0000-4000-8000-000000000001"}], - } - assert len(request.headers) == 1 - assert request.headers[0][0] == KNOWLEDGE_FS_QUERY_IMAGE_GRANTS_HEADER - encoded_grants = request.headers[0][1] - padded_grants = encoded_grants + "=" * (-len(encoded_grants) % 4) - assert json.loads(base64.urlsafe_b64decode(padded_grants)) == { - "g": ["short-lived-grant"], - "v": 1, - } - - -def test_list_metadata_fields_issues_workflow_capability_and_calls_catalog_operation() -> None: - admission = Admission() - broker = Broker() - remote = Remote() - service = KnowledgeFSAppExecutionCapabilityService( # type: ignore[arg-type] - admission=admission, - broker=broker, - remote=remote, - ) - run_context = DifyRunContext( - tenant_id="tenant-1", - app_id="app-1", - user_id="user-1", - user_from=UserFrom.ACCOUNT, - invoke_from=InvokeFrom.DEBUGGER, - trace_session_id="trace-from-workflow", - ) - - first_page = service.list_metadata_fields( - run_context=run_context, - caller_kind=KnowledgeFSAppSpaceJoinType.WORKFLOW, - resource=KnowledgeResourceRef(kind="knowledge_fs", control_space_id="control-1"), - ) - second_page = service.list_metadata_fields( - run_context=run_context, - caller_kind=KnowledgeFSAppSpaceJoinType.WORKFLOW, - resource=KnowledgeResourceRef(kind="knowledge_fs", control_space_id="control-1"), - cursor=first_page.next_cursor, - limit=50, - ) - - assert [item.name for item in first_page.data] == ["department"] - assert first_page.data[0].type == "string" - assert first_page.next_cursor == "cursor-2" - assert second_page.next_cursor == "cursor-2" - assert [call["operation_id"] for call in admission.calls] == ["listMetadataFields", "listMetadataFields"] - assert [call["caller_kind"] for call in admission.calls] == [ - KnowledgeFSAppSpaceJoinType.WORKFLOW, - KnowledgeFSAppSpaceJoinType.WORKFLOW, - ] - assert [call["operation_id"] for call in broker.calls] == ["listMetadataFields", "listMetadataFields"] - assert len(remote.calls) == 2 - first_request, second_request = remote.calls - assert first_request.operation_id == "listMetadataFields" - assert first_request.method == "GET" - assert first_request.path == "/knowledge-spaces/space-1/metadata-fields" - assert first_request.namespace_id == "tenant-1" - assert first_request.knowledge_space_id == "space-1" - assert first_request.capability_token == "token" - assert first_request.payload is None - assert first_request.query == (("limit", "100"),) - assert second_request.query == (("cursor", "cursor-2"), ("limit", "50")) - - -def test_list_metadata_fields_fails_closed_when_product_operation_is_unavailable( - monkeypatch: pytest.MonkeyPatch, -) -> None: - admission = Admission() - broker = Broker() - remote = Remote() - service = KnowledgeFSAppExecutionCapabilityService( # type: ignore[arg-type] - admission=admission, - broker=broker, - remote=remote, - ) - monkeypatch.setattr(app_execution_capability, "is_product_operation_ready", lambda _operation_id: False) - - with pytest.raises(KnowledgeFSOperationUnavailableError): - service.list_metadata_fields( - run_context=DifyRunContext( - tenant_id="tenant-1", - app_id="app-1", - user_id="user-1", - user_from=UserFrom.ACCOUNT, - invoke_from=InvokeFrom.DEBUGGER, - ), - caller_kind=KnowledgeFSAppSpaceJoinType.WORKFLOW, - resource=KnowledgeResourceRef(kind="knowledge_fs", control_space_id="control-1"), - ) - - assert admission.calls == [] - assert broker.calls == [] - assert remote.calls == [] - - -def test_capture_workflow_failed_retrieval_uses_fresh_transport_trace_and_business_event_id() -> None: - admission = Admission() - broker = Broker() - remote = Remote() - service = KnowledgeFSAppExecutionCapabilityService( # type: ignore[arg-type] - admission=admission, - broker=broker, - remote=remote, - ) - - result = service.capture_workflow_failed_retrieval( - tenant_id="tenant-1", - app_id="app-1", - resource=KnowledgeResourceRef(kind="knowledge_fs", control_space_id="control-1"), - payload=KnowledgeFSWorkflowFailedRetrievalCapturePayload( - event_id="019fac9f-bfb0-75ee-9af5-252ebafbac1e", - query=" missing answer ", - mode="fast", - retrieval_trace_id=" trace-retrieval-1 ", - ), - ) - - assert str(result.failed_query_id) == "019fac9f-bfb0-75ee-9af5-252ebafbac1c" - assert result.verdict == "retrieval-miss" - assert str(result.bad_case_id) == "019fac9f-bfb0-75ee-9af5-252ebafbac1d" - assert admission.calls == [ - { - "tenant_id": "tenant-1", - "app_id": "app-1", - "control_space_id": "control-1", - "caller_kind": KnowledgeFSAppSpaceJoinType.WORKFLOW, - "operation_id": "captureWorkflowFailedRetrieval", - } - ] - assert broker.calls == [ - { - "profile": admission.profile, - "operation_id": "captureWorkflowFailedRetrieval", - "resource_id": None, - "trace_id": None, - } - ] - assert remote.calls == [ - KnowledgeFSRemoteJSONRequest( - operation_id="captureWorkflowFailedRetrieval", - method="POST", - path="/knowledge-spaces/space-1/failed-queries/workflow-retrieval-misses", - namespace_id="tenant-1", - knowledge_space_id="space-1", - capability_token="token", - trace_id="trace-1", - payload={ - "eventId": "019fac9f-bfb0-75ee-9af5-252ebafbac1e", - "query": "missing answer", - "mode": "fast", - "retrievalTraceId": "trace-retrieval-1", - }, - ) - ] - - -def test_agent_investigation_uses_authorized_space_and_preserves_report_identity() -> None: - broker = Broker() - remote = Remote() - service = KnowledgeFSAppExecutionCapabilityService(admission=Admission(), broker=broker, remote=remote) - issued = broker.issue_app(operation_id="captureAgentKnowledgeInvestigation") - payload: dict[str, JsonValue] = {"investigationId": "investigation-1", "query": "refund", "attempts": []} - - result = service.capture_agent_investigation(issued=issued, tenant_id="tenant-1", payload=payload) - - assert result == {"investigationId": "investigation-1", "failedQueries": []} - assert remote.calls == [ - KnowledgeFSRemoteJSONRequest( - operation_id="captureAgentKnowledgeInvestigation", - method="POST", - path="/knowledge-spaces/space-1/agent-investigations", - namespace_id="tenant-1", - knowledge_space_id="space-1", - capability_token="token", - trace_id="trace-1", - payload=payload, - ) - ] - - -@pytest.mark.parametrize("unavailable", ["wrong-capability", "unready-operation", "missing-path"]) -def test_agent_investigation_refuses_unavailable_or_wrong_capability_before_io( - monkeypatch: pytest.MonkeyPatch, unavailable: str -) -> None: - broker = Broker() - remote = Remote() - service = KnowledgeFSAppExecutionCapabilityService(admission=Admission(), broker=broker, remote=remote) - issued = broker.issue_app(operation_id="captureAgentKnowledgeInvestigation") - if unavailable == "wrong-capability": - issued = issued._replace(operation_id="retrieveEvidence") - elif unavailable == "unready-operation": - monkeypatch.setattr(app_execution_capability, "is_product_operation_ready", lambda _operation_id: False) - else: - operation_id = "captureAgentKnowledgeInvestigation" - operations = dict(app_execution_capability.KNOWLEDGE_FS_PRODUCT_OPERATIONS) - operations[operation_id] = operations[operation_id]._replace(kfs_path=None) - monkeypatch.setattr(app_execution_capability, "KNOWLEDGE_FS_PRODUCT_OPERATIONS", operations) - monkeypatch.setattr(app_execution_capability, "is_product_operation_ready", lambda _operation_id: True) - - with pytest.raises(KnowledgeFSOperationUnavailableError): - service.capture_agent_investigation(issued=issued, tenant_id="tenant-1", payload={}) - - assert remote.calls == [] diff --git a/api/tests/unit_tests/services/test_knowledge_fs_capability.py b/api/tests/unit_tests/services/test_knowledge_fs_capability.py index db7f1764f44..6c1e575b7dd 100644 --- a/api/tests/unit_tests/services/test_knowledge_fs_capability.py +++ b/api/tests/unit_tests/services/test_knowledge_fs_capability.py @@ -25,8 +25,8 @@ from services.knowledge_fs_capability import ( RotatingCapabilityKeyRing, create_configured_knowledge_fs_capability_issuer, ) -from tests.unit_tests.config_override import apply_config_overrides -from tests.unit_tests.services.knowledge_fs_services_b_fakes import ( +from tests.helpers.unit.config_override import apply_config_overrides +from tests.helpers.unit.services.knowledge_fs_services_b_fakes import ( RecordingCapabilityAuditor, RecordingOperationalMetrics, ) diff --git a/api/tests/unit_tests/services/test_knowledge_fs_download_facade.py b/api/tests/unit_tests/services/test_knowledge_fs_download_facade.py deleted file mode 100644 index bd77833095a..00000000000 --- a/api/tests/unit_tests/services/test_knowledge_fs_download_facade.py +++ /dev/null @@ -1,238 +0,0 @@ -from __future__ import annotations - -from datetime import UTC, datetime -from typing import Literal, override - -import pytest - -from services.knowledge_fs.data_facade import KnowledgeFSDataFacade -from services.knowledge_fs.product_dto import ( - KnowledgeFSDocumentResponse, - KnowledgeFSDocumentRevisionListResponse, - KnowledgeFSDocumentRevisionResponse, - KnowledgeFSLogicalDocumentResponse, -) -from services.knowledge_fs.product_remote import KnowledgeFSProductResourceNotFoundError - - -def logical_document( - *, - active: KnowledgeFSDocumentRevisionResponse | None, - status: Literal["deleting", "failed", "pending", "ready"] = "ready", -) -> KnowledgeFSLogicalDocumentResponse: - return KnowledgeFSLogicalDocumentResponse( - active=active, - active_revision=active.revision if active else None, - created_at=datetime.now(UTC), - enabled=True, - id="logical-1", - knowledge_space_id="space-1", - row_version=1, - status=status, - title="Document", - updated_at=datetime.now(UTC), - user_metadata={}, - ) - - -def revision( - *, state: Literal["active", "candidate", "failed", "superseded"] = "active" -) -> KnowledgeFSDocumentRevisionResponse: - return KnowledgeFSDocumentRevisionResponse( - activated_at=datetime.now(UTC) if state == "active" else None, - content_hash="a" * 64, - created_at=datetime.now(UTC), - document_asset_id="asset-1", - document_asset_version=2, - document_id="logical-1", - knowledge_space_id="space-1", - mime_type="text/markdown", - revision=3, - size_bytes=4, - state=state, - ) - - -def asset() -> KnowledgeFSDocumentResponse: - return KnowledgeFSDocumentResponse( - id="asset-1", - knowledge_space_id="space-1", - filename="page.md", - metadata={}, - mime_type="text/markdown", - object_key="tenant/spaces/space-1/documents/page.md", - parser_status="parsed", - sha256="a" * 64, - size_bytes=4, - version=2, - created_at=datetime.now(UTC), - ) - - -class FakeDownloadFacade(KnowledgeFSDataFacade): - """Serves canned facade reads and records the calls download preparation makes.""" - - def __init__( - self, - *, - logical: KnowledgeFSLogicalDocumentResponse, - revisions: KnowledgeFSDocumentRevisionListResponse | None = None, - document: KnowledgeFSDocumentResponse | None = None, - ) -> None: - self._logical = logical - self._revisions = revisions - self._document = document - self.get_logical_document_calls: list[dict[str, object]] = [] - self.list_document_revisions_calls: list[dict[str, object]] = [] - self.get_document_calls: list[dict[str, object]] = [] - - @override - def get_logical_document( - self, *, tenant_id: str, account_id: str, control_space_id: str, document_id: str - ) -> KnowledgeFSLogicalDocumentResponse: - self.get_logical_document_calls.append( - { - "tenant_id": tenant_id, - "account_id": account_id, - "control_space_id": control_space_id, - "document_id": document_id, - } - ) - return self._logical - - @override - def list_document_revisions( - self, - *, - tenant_id: str, - account_id: str, - control_space_id: str, - document_id: str, - cursor: str | None = None, - limit: int | None = None, - ) -> KnowledgeFSDocumentRevisionListResponse: - self.list_document_revisions_calls.append( - { - "tenant_id": tenant_id, - "account_id": account_id, - "control_space_id": control_space_id, - "document_id": document_id, - "cursor": cursor, - "limit": limit, - } - ) - assert self._revisions is not None - return self._revisions - - @override - def get_document( - self, *, tenant_id: str, account_id: str, control_space_id: str, document_id: str - ) -> KnowledgeFSDocumentResponse: - self.get_document_calls.append( - { - "tenant_id": tenant_id, - "account_id": account_id, - "control_space_id": control_space_id, - "document_id": document_id, - } - ) - assert self._document is not None - return self._document - - -def test_prepare_logical_document_download_resolves_active_asset() -> None: - facade = FakeDownloadFacade(logical=logical_document(active=revision()), document=asset()) - - result = facade.prepare_logical_document_download( - tenant_id="tenant-1", - account_id="account-1", - control_space_id="control-1", - document_id="logical-1", - ) - - assert result.document_id == "logical-1" - assert result.object_key == "tenant/spaces/space-1/documents/page.md" - assert facade.list_document_revisions_calls == [] - assert facade.get_document_calls == [ - { - "tenant_id": "tenant-1", - "account_id": "account-1", - "control_space_id": "control-1", - "document_id": "asset-1", - } - ] - - -def test_prepare_logical_document_download_resolves_latest_failed_revision_without_active_revision() -> None: - facade = FakeDownloadFacade( - logical=logical_document(active=None, status="failed"), - revisions=KnowledgeFSDocumentRevisionListResponse(data=[revision(state="failed")]), - document=asset(), - ) - - result = facade.prepare_logical_document_download( - tenant_id="tenant-1", - account_id="account-1", - control_space_id="control-1", - document_id="logical-1", - ) - - assert result.document_id == "logical-1" - assert result.object_key == "tenant/spaces/space-1/documents/page.md" - assert facade.list_document_revisions_calls == [ - { - "tenant_id": "tenant-1", - "account_id": "account-1", - "control_space_id": "control-1", - "document_id": "logical-1", - "cursor": None, - "limit": 1, - } - ] - - -def test_prepare_logical_document_download_rejects_document_without_visible_revision() -> None: - facade = FakeDownloadFacade( - logical=logical_document(active=None, status="failed"), - revisions=KnowledgeFSDocumentRevisionListResponse(data=[]), - ) - - with pytest.raises(KnowledgeFSProductResourceNotFoundError): - facade.prepare_logical_document_download( - tenant_id="tenant-1", - account_id="account-1", - control_space_id="control-1", - document_id="logical-1", - ) - - -def test_prepare_logical_document_download_rejects_latest_candidate_revision() -> None: - facade = FakeDownloadFacade( - logical=logical_document(active=None, status="failed"), - revisions=KnowledgeFSDocumentRevisionListResponse(data=[revision(state="candidate")]), - ) - - with pytest.raises(KnowledgeFSProductResourceNotFoundError): - facade.prepare_logical_document_download( - tenant_id="tenant-1", - account_id="account-1", - control_space_id="control-1", - document_id="logical-1", - ) - - assert facade.get_document_calls == [] - - -def test_prepare_logical_document_download_rejects_pending_document_without_active_revision() -> None: - facade = FakeDownloadFacade(logical=logical_document(active=None, status="pending")) - - with pytest.raises(KnowledgeFSProductResourceNotFoundError): - facade.prepare_logical_document_download( - tenant_id="tenant-1", - account_id="account-1", - control_space_id="control-1", - document_id="logical-1", - ) - - assert facade.list_document_revisions_calls == [] - assert facade.get_document_calls == [] diff --git a/api/tests/unit_tests/services/test_knowledge_fs_initial_members.py b/api/tests/unit_tests/services/test_knowledge_fs_initial_members.py deleted file mode 100644 index 53aabe0c123..00000000000 --- a/api/tests/unit_tests/services/test_knowledge_fs_initial_members.py +++ /dev/null @@ -1,259 +0,0 @@ -from __future__ import annotations - -from typing import cast - -import pytest -from pydantic import ValidationError -from sqlalchemy import delete, func, select -from sqlalchemy.orm import Session, sessionmaker - -from models import TenantAccountJoin -from models.knowledge_fs import ( - KnowledgeFSAuthorizationRevision, - KnowledgeFSControlSpace, - KnowledgeFSControlSpacePermission, - KnowledgeFSControlSpacePermissionRole, - KnowledgeFSControlSpacePermissionStatus, - KnowledgeFSControlSpaceVisibility, - KnowledgeFSExternalAccessPolicy, - KnowledgeFSLifecycleOperation, - KnowledgeFSLifecycleOutbox, - KnowledgeFSProvisionCommandPayload, -) -from repositories.sqlalchemy_knowledge_fs_lifecycle_outbox_repository import ( - SQLAlchemyKnowledgeFSLifecycleOutboxRepository, -) -from services.knowledge_fs.control_space_commands import ( - KnowledgeFSControlSpaceCommandService, - KnowledgeFSControlSpaceIntentConflictError, - KnowledgeFSProvisionIntent, -) -from services.knowledge_fs.product_dto import KnowledgeFSSpaceCreatePayload -from services.knowledge_fs.workspace_members import KnowledgeFSControlPlaneInvariantError - - -def _intent() -> KnowledgeFSProvisionIntent: - return KnowledgeFSProvisionIntent( - tenant_id="tenant-1", - owner_account_id="owner-1", - provisioning_key="provision-1", - operation_id="operation-1", - idempotency_key="create-1", - name="Selected members", - slug="selected-members", - icon=None, - description=None, - model_intent=None, - profile_intent=None, - visibility=KnowledgeFSControlSpaceVisibility.PARTIAL_MEMBERS, - member_account_ids=("member-2", "member-1"), - ) - - -def _provision_payload(outbox: KnowledgeFSLifecycleOutbox) -> KnowledgeFSProvisionCommandPayload: - assert outbox.operation is KnowledgeFSLifecycleOperation.PROVISION - return cast(KnowledgeFSProvisionCommandPayload, outbox.command_payload) - - -def _seed_members(session: Session) -> None: - session.add_all( - [ - TenantAccountJoin(tenant_id="tenant-1", account_id="member-1"), - TenantAccountJoin(tenant_id="tenant-1", account_id="member-2"), - TenantAccountJoin(tenant_id="other-tenant", account_id="outsider"), - ] - ) - session.commit() - - -def _assert_no_creation(session: Session) -> None: - for model in ( - KnowledgeFSControlSpace, - KnowledgeFSControlSpacePermission, - KnowledgeFSAuthorizationRevision, - KnowledgeFSExternalAccessPolicy, - KnowledgeFSLifecycleOutbox, - ): - assert session.scalar(select(func.count()).select_from(model)) == 0 - - -def test_partial_members_create_dto_accepts_viewers_and_preserves_default_creation() -> None: - payload = KnowledgeFSSpaceCreatePayload.model_validate( - { - "name": "Selected members", - "slug": "selected-members", - "visibility": "partial_members", - "members": [{"account_id": "member-1", "role": "viewer"}], - } - ) - assert payload.members[0].account_id == "member-1" - assert payload.members[0].role == "viewer" - assert KnowledgeFSSpaceCreatePayload(name="Private", slug="private").members == [] - assert KnowledgeFSSpaceCreatePayload(name="Team", slug="team", visibility="all_team_members").members == [] - - -@pytest.mark.parametrize( - "updates", - [ - {"visibility": "partial_members"}, - {"visibility": "partial_members", "members": []}, - {"visibility": "partial_members", "members": None}, - {"members": [{"account_id": "member-1", "role": "viewer"}]}, - {"visibility": "all_team_members", "members": [{"account_id": "member-1", "role": "viewer"}]}, - {"visibility": "partial_members", "members": [{"account_id": "member-1", "role": "owner"}]}, - {"visibility": "partial_members", "members": [{"account_id": "member-1", "role": "editor"}]}, - {"visibility": "partial_members", "members": [{"account_id": "member-1"}]}, - {"visibility": "partial_members", "members": [{"account_id": "", "role": "viewer"}]}, - { - "visibility": "partial_members", - "members": [{"account_id": "member-1", "role": "viewer"}] * 2, - }, - { - "visibility": "partial_members", - "members": [{"account_id": f"member-{i}", "role": "viewer"} for i in range(1_001)], - }, - ], -) -def test_partial_members_create_dto_rejects_invalid_or_privileged_members(updates: dict[str, object]) -> None: - with pytest.raises(ValidationError): - KnowledgeFSSpaceCreatePayload.model_validate({"name": "Space", "slug": "space", **updates}) - - -def test_initial_members_and_visibility_are_committed_with_the_space_and_outbox( - sqlite_session: Session, sqlite_session_factory: sessionmaker[Session] -) -> None: - _seed_members(sqlite_session) - service = KnowledgeFSControlSpaceCommandService(sqlite_session_factory) - - result = service.create_provision_intent(_intent()) - replay = service.create_provision_intent(_intent()._replace(member_account_ids=("member-1", "member-2"))) - - assert result.control_space.visibility is KnowledgeFSControlSpaceVisibility.PARTIAL_MEMBERS - assert replay.control_space.id == result.control_space.id - assert replay.outbox.id == result.outbox.id - assert result.outbox.expected_control_space_version == result.control_space.resource_version == 0 - command_payload = _provision_payload(result.outbox) - assert command_payload["initial_visibility"] == "partial_members" - assert command_payload["initial_member_account_ids"] == ["member-1", "member-2"] - permissions = list(sqlite_session.scalars(select(KnowledgeFSControlSpacePermission))) - assert {permission.account_id: permission.role for permission in permissions} == { - "owner-1": KnowledgeFSControlSpacePermissionRole.OWNER, - "member-1": KnowledgeFSControlSpacePermissionRole.VIEWER, - "member-2": KnowledgeFSControlSpacePermissionRole.VIEWER, - } - assert all(permission.granted_by_account_id == "owner-1" for permission in permissions) - assert sqlite_session.scalar(select(func.count()).select_from(KnowledgeFSLifecycleOutbox)) == 1 - assert sqlite_session.scalar(select(func.count()).select_from(KnowledgeFSAuthorizationRevision)) == 1 - assert sqlite_session.scalar(select(func.count()).select_from(KnowledgeFSExternalAccessPolicy)) == 1 - - -@pytest.mark.parametrize("member_ids", [(), ("owner-1",), ("member-1", "member-1"), ("outsider",), ("missing",)]) -def test_initial_member_validation_precedes_any_creation( - sqlite_session: Session, sqlite_session_factory: sessionmaker[Session], member_ids: tuple[str, ...] -) -> None: - _seed_members(sqlite_session) - service = KnowledgeFSControlSpaceCommandService(sqlite_session_factory) - - with pytest.raises(KnowledgeFSControlPlaneInvariantError): - service.create_provision_intent(_intent()._replace(member_account_ids=member_ids)) - - _assert_no_creation(sqlite_session) - - -def test_initial_members_roll_back_if_the_outbox_write_fails( - sqlite_session: Session, - sqlite_session_factory: sessionmaker[Session], - monkeypatch: pytest.MonkeyPatch, -) -> None: - _seed_members(sqlite_session) - original_add = SQLAlchemyKnowledgeFSLifecycleOutboxRepository.add - - def fail_after_flush( - repository: SQLAlchemyKnowledgeFSLifecycleOutboxRepository, command: KnowledgeFSLifecycleOutbox - ) -> KnowledgeFSLifecycleOutbox: - original_add(repository, command) - raise RuntimeError("outbox write interrupted") - - monkeypatch.setattr(SQLAlchemyKnowledgeFSLifecycleOutboxRepository, "add", fail_after_flush) - - with pytest.raises(RuntimeError, match="outbox write interrupted"): - KnowledgeFSControlSpaceCommandService(sqlite_session_factory).create_provision_intent(_intent()) - - _assert_no_creation(sqlite_session) - - -def test_creation_replay_does_not_restore_revoked_members_or_initial_visibility( - sqlite_session: Session, sqlite_session_factory: sessionmaker[Session] -) -> None: - _seed_members(sqlite_session) - service = KnowledgeFSControlSpaceCommandService(sqlite_session_factory) - result = service.create_provision_intent(_intent()) - space = sqlite_session.get(KnowledgeFSControlSpace, result.control_space.id) - assert space is not None - space.visibility = KnowledgeFSControlSpaceVisibility.ONLY_ME - member = sqlite_session.scalar( - select(KnowledgeFSControlSpacePermission).where(KnowledgeFSControlSpacePermission.account_id == "member-1") - ) - assert member is not None - member.status = KnowledgeFSControlSpacePermissionStatus.REVOKED - sqlite_session.execute(delete(TenantAccountJoin).where(TenantAccountJoin.account_id == "member-1")) - sqlite_session.commit() - - replay = service.create_provision_intent(_intent()) - - assert replay.control_space.visibility is KnowledgeFSControlSpaceVisibility.ONLY_ME - assert replay.outbox.id == result.outbox.id - sqlite_session.refresh(member) - assert member.status is KnowledgeFSControlSpacePermissionStatus.REVOKED - - -@pytest.mark.parametrize( - ("visibility", "member_account_ids"), - [ - (KnowledgeFSControlSpaceVisibility.PARTIAL_MEMBERS, ("member-1",)), - (KnowledgeFSControlSpaceVisibility.ALL_TEAM_MEMBERS, ()), - (KnowledgeFSControlSpaceVisibility.ONLY_ME, ()), - ], -) -def test_creation_replay_cannot_change_initial_access( - sqlite_session: Session, - sqlite_session_factory: sessionmaker[Session], - visibility: KnowledgeFSControlSpaceVisibility, - member_account_ids: tuple[str, ...], -) -> None: - _seed_members(sqlite_session) - service = KnowledgeFSControlSpaceCommandService(sqlite_session_factory) - service.create_provision_intent(_intent()) - - with pytest.raises(KnowledgeFSControlSpaceIntentConflictError, match="different initial access"): - service.create_provision_intent( - _intent()._replace(visibility=visibility, member_account_ids=member_account_ids) - ) - - assert sqlite_session.scalar(select(func.count()).select_from(KnowledgeFSControlSpacePermission)) == 3 - assert sqlite_session.scalar(select(func.count()).select_from(KnowledgeFSLifecycleOutbox)) == 1 - - -def test_legacy_creation_replay_keeps_current_access_and_rejects_added_members( - sqlite_session: Session, sqlite_session_factory: sessionmaker[Session] -) -> None: - _seed_members(sqlite_session) - service = KnowledgeFSControlSpaceCommandService(sqlite_session_factory) - private_intent = _intent()._replace(visibility=KnowledgeFSControlSpaceVisibility.ONLY_ME, member_account_ids=()) - result = service.create_provision_intent(private_intent) - command = sqlite_session.get(KnowledgeFSLifecycleOutbox, result.outbox.id) - assert command is not None - legacy_payload = _provision_payload(command).copy() - del legacy_payload["initial_visibility"] - del legacy_payload["initial_member_account_ids"] - command.command_payload = legacy_payload - sqlite_session.commit() - - replay = service.create_provision_intent( - private_intent._replace(visibility=KnowledgeFSControlSpaceVisibility.ALL_TEAM_MEMBERS) - ) - - assert replay.control_space.visibility is KnowledgeFSControlSpaceVisibility.ONLY_ME - assert replay.outbox.id == result.outbox.id - with pytest.raises(KnowledgeFSControlSpaceIntentConflictError): - service.create_provision_intent(_intent()) diff --git a/api/tests/unit_tests/services/test_knowledge_fs_initial_source_reliability.py b/api/tests/unit_tests/services/test_knowledge_fs_initial_source_reliability.py deleted file mode 100644 index 32cffba51be..00000000000 --- a/api/tests/unit_tests/services/test_knowledge_fs_initial_source_reliability.py +++ /dev/null @@ -1,77 +0,0 @@ -import pytest -from pydantic import ValidationError - -from services.knowledge_fs.product_dto import ( - KnowledgeFSCrawlImportPayload, - KnowledgeFSInitialWebsiteSourcePayload, - knowledge_fs_initial_preview_configuration_fingerprint, -) - - -def _payload(*, limit: int = 2, urls: list[str] | None = None) -> dict[str, object]: - selected = urls or ["HTTPS://Docs.Dify.AI:443/a/#section", "https://docs.dify.ai/b/"] - return { - "kind": "website_crawl", - "name": "Dify docs", - "provider": "firecrawl", - "root_url": "https://docs.dify.ai", - "crawl_options": {"include_subpages": True, "limit": limit}, - "selection": [{"source_url": url} for url in selected], - } - - -def test_initial_website_selection_populates_stable_canonical_urls() -> None: - payload = KnowledgeFSInitialWebsiteSourcePayload.model_validate(_payload()) - - assert [item.canonical_url for item in payload.selection] == [ - "https://docs.dify.ai/a", - "https://docs.dify.ai/b", - ] - - -def test_initial_website_selection_preserves_canonical_duplicates() -> None: - payload = KnowledgeFSInitialWebsiteSourcePayload.model_validate( - _payload(urls=["https://docs.dify.ai/a/", "https://DOCS.dify.ai:443/a#section"]) - ) - - assert [item.source_url for item in payload.selection] == [ - "https://docs.dify.ai/a/", - "https://DOCS.dify.ai:443/a#section", - ] - assert [item.canonical_url for item in payload.selection] == [ - "https://docs.dify.ai/a", - "https://docs.dify.ai/a", - ] - - -def test_crawl_import_preserves_duplicate_source_urls() -> None: - payload = KnowledgeFSCrawlImportPayload.model_validate( - {"sourceUrls": [" https://docs.dify.ai/a ", "https://docs.dify.ai/a"]} - ) - - assert payload.source_urls == ["https://docs.dify.ai/a", "https://docs.dify.ai/a"] - - -def test_initial_website_selection_requires_crawl_limit_to_cover_selection() -> None: - with pytest.raises(ValidationError, match="crawl limit must cover every selected URL"): - KnowledgeFSInitialWebsiteSourcePayload.model_validate(_payload(limit=1)) - - -def test_initial_website_source_rejects_stale_preview_configuration() -> None: - payload = _payload() - payload["previewConfigurationFingerprint"] = "0" * 64 - - with pytest.raises(ValidationError, match="preview configuration no longer matches"): - KnowledgeFSInitialWebsiteSourcePayload.model_validate(payload) - - -def test_initial_website_source_accepts_matching_preview_configuration() -> None: - payload = KnowledgeFSInitialWebsiteSourcePayload.model_validate(_payload()) - raw = _payload() - raw["previewConfigurationFingerprint"] = knowledge_fs_initial_preview_configuration_fingerprint(payload) - - validated = KnowledgeFSInitialWebsiteSourcePayload.model_validate(raw) - - assert validated.preview_configuration_fingerprint == knowledge_fs_initial_preview_configuration_fingerprint( - validated - ) diff --git a/api/tests/unit_tests/services/test_knowledge_fs_retrieval_v2_contracts.py b/api/tests/unit_tests/services/test_knowledge_fs_retrieval_v2_contracts.py deleted file mode 100644 index 66253ac9066..00000000000 --- a/api/tests/unit_tests/services/test_knowledge_fs_retrieval_v2_contracts.py +++ /dev/null @@ -1,205 +0,0 @@ -from __future__ import annotations - -import math -from datetime import UTC, datetime - -import pytest -from pydantic import ValidationError -from sqlalchemy.orm import Session - -from graphon.model_runtime.entities.rerank_entities import RerankDocument, RerankResult -from models.knowledge_fs import KnowledgeFSControlSpace, KnowledgeFSControlSpaceState -from services.knowledge_fs.knowledge_retrieval_v2 import knowledge_retrieval_v2_node as node_module -from services.knowledge_fs.knowledge_retrieval_v2.entities import ( - KNOWLEDGE_RETRIEVAL_V2_NODE_TYPE, - KnowledgeRetrievalV2NodeData, -) -from services.knowledge_fs.knowledge_retrieval_v2.exc import KnowledgeFSRetrievalContractError -from services.knowledge_fs.knowledge_retrieval_v2.knowledge_retrieval_v2_node import KnowledgeRetrievalV2Node -from services.knowledge_fs.knowledge_retrieval_v2.validation import ( - collect_control_space_ids, - missing_control_space_ids, - validate_control_space_references, -) -from services.knowledge_fs.retrieval_contracts import KnowledgeFSRetrievalTestItemResponse - - -def _node_data(**overrides: object) -> dict[str, object]: - return { - "type": KNOWLEDGE_RETRIEVAL_V2_NODE_TYPE, - "title": "Agent Knowledge Retrieval", - "control_space_ids": ["space-a"], - "query_variable_selector": ["start", "query"], - **overrides, - } - - -def _graph(*nodes: object) -> dict[str, object]: - return {"nodes": list(nodes)} - - -def _item( - *, text: str | None, section_path: list[str], node_id: str = "node-a" -) -> KnowledgeFSRetrievalTestItemResponse: - return KnowledgeFSRetrievalTestItemResponse.model_validate( - { - "citation": { - "artifactHash": "a" * 64, - "documentAssetId": "document-a", - "documentTitle": "a.pdf", - "documentVersion": 1, - "knowledgeSpaceName": "Space A", - "logicalDocumentId": "logical-a", - "logicalDocumentRevision": 1, - "pageNumber": 1, - "sectionPath": section_path, - }, - "nodeId": node_id, - "projectionIds": ["projection-a"], - "score": 0.5, - "sources": ["dense"], - "text": text, - } - ) - - -def _candidates(count: int) -> list[node_module._WorkflowRerankCandidate]: - return [ - node_module._WorkflowRerankCandidate( - control_space_id="space-a", - item=_item(text=f"text {index}", section_path=[], node_id=f"node-{index}"), - item_index=index, - space_index=0, - ) - for index in range(count) - ] - - -def _active_space(space_id: str, *, tenant_id: str = "tenant-1") -> KnowledgeFSControlSpace: - space = KnowledgeFSControlSpace( - tenant_id=tenant_id, - owner_account_id="account-1", - provisioning_key=f"provision-{space_id}", - knowledge_space_id=f"knowledge-{space_id}", - state=KnowledgeFSControlSpaceState.ACTIVE, - ) - space.id = space_id - return space - - -def test_collect_control_space_ids_ignores_malformed_graphs_and_foreign_nodes() -> None: - assert collect_control_space_ids({"nodes": "not-a-list"}) == () - assert collect_control_space_ids( - _graph( - "not-a-node", - {"data": "not-a-mapping"}, - {"data": {"type": "knowledge-retrieval"}}, - {"data": _node_data(control_space_ids=["space-b", "space-a"])}, - ) - ) == ("space-b", "space-a") - - -def test_reference_validation_skips_the_database_without_v2_nodes() -> None: - # A graph without KnowledgeFS nodes must not need a session at all. - assert missing_control_space_ids(session=Session(), tenant_id="tenant-1", graph=_graph()) == () - validate_control_space_references(session=Session(), tenant_id="tenant-1", graph=_graph()) - - -def test_reference_validation_rejects_spaces_outside_the_tenant(sqlite_session: Session) -> None: - sqlite_session.add_all([_active_space("space-a"), _active_space("space-b", tenant_id="tenant-2")]) - sqlite_session.flush() - graph = _graph({"data": _node_data(control_space_ids=["space-a", "space-b"])}) - - with pytest.raises(ValueError, match="missing or inactive in this workspace: space-b$"): - validate_control_space_references(session=sqlite_session, tenant_id="tenant-1", graph=graph) - - validate_control_space_references( - session=sqlite_session, - tenant_id="tenant-1", - graph=_graph({"data": _node_data(control_space_ids=["space-a"])}), - ) - - -@pytest.mark.parametrize( - ("overrides", "message"), - [ - ({"control_space_ids": ["space-a", " "]}, "IDs must be non-empty"), - ({"query_variable_selector": ["start"]}, "query variable selector must contain at least two parts"), - ({"query_variable_selector": ["start", " "]}, "query variable selector must be non-empty"), - ({"query_attachment_selector": ["start"]}, "query attachment selector must contain at least two parts"), - ({"query_attachment_selector": ["start", " "]}, "query attachment selector must be non-empty"), - ], -) -def test_node_data_rejects_malformed_identifiers_and_selectors(overrides: dict[str, object], message: str) -> None: - with pytest.raises(ValidationError, match=message): - KnowledgeRetrievalV2NodeData.model_validate(_node_data(**overrides)) - - -def test_node_data_treats_empty_selectors_as_unbound() -> None: - node_data = KnowledgeRetrievalV2NodeData.model_validate( - _node_data(query_variable_selector=[], query_attachment_selector=["start", "image"]) - ) - - assert node_data.query_variable_selector is None - assert node_data.query_attachment_selector == ["start", "image"] - assert ( - KnowledgeRetrievalV2NodeData.model_validate(_node_data(query_attachment_selector=[])).query_attachment_selector - is None - ) - - -def test_variable_mapping_omits_unbound_query_attachment() -> None: - node_data = KnowledgeRetrievalV2NodeData.model_validate(_node_data()) - - assert KnowledgeRetrievalV2Node._extract_variable_selector_to_variable_mapping( - graph_config={}, - node_id="retrieval", - node_data=node_data, - ) == {"retrieval.query": ["start", "query"]} - - -@pytest.mark.parametrize( - ("value", "expected"), - [ - (None, None), - ("tag", "tag"), - (1.5, 1.5), - (3, 3), - (True, "True"), - (datetime(2026, 1, 2, tzinfo=UTC), "2026-01-02 00:00:00+00:00"), - ], -) -def test_manual_metadata_values_are_normalized_to_contract_scalars(value: object, expected: object) -> None: - assert node_module._normalize_metadata_filter_scalar(value) == expected - - -def test_rerank_text_falls_back_to_section_path_then_node_id() -> None: - assert KnowledgeRetrievalV2Node._rerank_text(_item(text=" body ", section_path=["A"])) == "body" - assert KnowledgeRetrievalV2Node._rerank_text(_item(text=" ", section_path=[" Camera ", " ", "Lens"])) == ( - "Camera / Lens" - ) - assert KnowledgeRetrievalV2Node._rerank_text(_item(text=None, section_path=[], node_id="node-z")) == "node-z" - - -def test_rerank_result_validation_keeps_model_order_and_rejects_bad_scores() -> None: - candidates = _candidates(2) - - ranked = KnowledgeRetrievalV2Node._validate_rerank_result( - result=RerankResult( - model="rerank", - docs=[RerankDocument(index=1, text="text 1", score=0.9), RerankDocument(index=0, text="text 0", score=0.1)], - ), - candidates=candidates, - ) - assert ranked == [(0.9, candidates[1]), (0.1, candidates[0])] - - with pytest.raises(KnowledgeFSRetrievalContractError, match="non-finite score"): - KnowledgeRetrievalV2Node._validate_rerank_result( - result=RerankResult(model="rerank", docs=[RerankDocument(index=0, text="text 0", score=math.nan)]), - candidates=candidates, - ) - with pytest.raises(KnowledgeFSRetrievalContractError, match="invalid document indices"): - KnowledgeRetrievalV2Node._validate_rerank_result( - result=RerankResult(model="rerank", docs=[RerankDocument(index=-1, text="text", score=0.5)]), - candidates=candidates, - ) diff --git a/api/tests/unit_tests/services/test_knowledge_fs_vector_backends.py b/api/tests/unit_tests/services/test_knowledge_fs_vector_backends.py deleted file mode 100644 index e4ff1a681fb..00000000000 --- a/api/tests/unit_tests/services/test_knowledge_fs_vector_backends.py +++ /dev/null @@ -1,685 +0,0 @@ -from collections.abc import Callable, Mapping -from dataclasses import dataclass, field -from typing import TYPE_CHECKING, cast -from unittest.mock import patch - -import pytest -from elastic_transport import ApiResponseMeta, HttpHeaders, NodeConfig -from elasticsearch import BadRequestError, NotFoundError -from weaviate.classes.config import ConsistencyLevel, DataType, VectorDistances -from weaviate.classes.data import DataObject -from weaviate.classes.query import Filter - -from services.knowledge_fs.vector_store import ( - VectorRequest, - VectorScope, - VectorStoreUnavailableError, - configured_vector_client, - execute_vector_request, -) -from services.knowledge_fs.vector_store_elasticsearch import ElasticsearchVectorStore -from services.knowledge_fs.vector_store_weaviate import WeaviateVectorStore - -if TYPE_CHECKING: - from elasticsearch import Elasticsearch - from weaviate import WeaviateClient - -A, B, GENERATION = (f"00000000-0000-4000-8000-{n:012d}" for n in (1, 2, 3)) -SCOPE: dict[str, object] = { - "tenant_id": A, - "knowledge_space_id": B, - "vector_space_id": "test", - "kind": "dense", - "dimension": 2, -} -POINT: dict[str, object] = {"id": A, "generation_id": GENERATION, "content_hash": "a" * 64, "vector": [3, 4]} - -type FieldMappings = dict[str, dict[str, object]] -type IndexMapping = dict[str, dict[str, FieldMappings]] -type ElasticsearchBody = dict[str, object] -type WeaviatePoint = DataObject[dict[str, str], None] - - -def request(operation: str, **kwargs: object) -> VectorRequest: - return VectorRequest.model_validate({"operation": operation, "scope": SCOPE, **kwargs}) - - -def fail_with_configured_store(expected: type[ElasticsearchVectorStore] | type[WeaviateVectorStore]) -> None: - with configured_vector_client(A) as store: - assert isinstance(store, expected) - raise RuntimeError("failure") - - -def es_error(status: int, name: str) -> NotFoundError | BadRequestError: - meta = ApiResponseMeta(status, "1.1", HttpHeaders(), 0, NodeConfig("http", "localhost", 9200)) - return (NotFoundError if status == 404 else BadRequestError)(name, meta, {"error": name}) - - -def es_mapping() -> IndexMapping: - return { - "mappings": { - "properties": { - "vector": {"type": "dense_vector", "dims": 2, "index": True, "similarity": "cosine"}, - "generation_id": {"type": "keyword"}, - "content_hash": {"type": "keyword"}, - } - } - } - - -@dataclass -class FakeElasticsearchIndices: - """Serves stored mappings per index; a missing index raises Elasticsearch's 404 error.""" - - mappings: dict[str, IndexMapping] = field(default_factory=dict) - queued_get_mapping_results: list[dict[str, IndexMapping] | Exception] = field(default_factory=list) - create_failure: Exception | None = None - get_mapping_calls: list[str] = field(default_factory=list) - created_indexes: list[str] = field(default_factory=list) - created_mappings: list[dict[str, object]] = field(default_factory=list) - - def get_mapping(self, *, index: str) -> dict[str, IndexMapping]: - self.get_mapping_calls.append(index) - if self.queued_get_mapping_results: - result = self.queued_get_mapping_results.pop(0) - if isinstance(result, Exception): - raise result - return result - if index not in self.mappings: - raise es_error(404, "missing") - return {index: self.mappings[index]} - - def create(self, *, index: str, mappings: dict[str, object]) -> None: - self.created_indexes.append(index) - self.created_mappings.append(mappings) - if self.create_failure is not None: - raise self.create_failure - - -@dataclass(frozen=True) -class BulkCall: - operations: list[dict[str, object]] - refresh: str - - def target_indexes(self, action: str) -> set[object]: - """Return the `_index` of every bulk action line of the given kind.""" - indexes: set[object] = set() - for operation in self.operations: - target = operation.get(action) - if isinstance(target, dict): - indexes.add(target.get("_index")) - return indexes - - -@dataclass -class FakeElasticsearch: - indices: FakeElasticsearchIndices = field(default_factory=FakeElasticsearchIndices) - bulk_result: ElasticsearchBody = field(default_factory=lambda: {"items": [{"index": {"status": 201}}]}) - mget_result: ElasticsearchBody = field(default_factory=lambda: {"docs": list[object]()}) - search_result: ElasticsearchBody = field(default_factory=lambda: {"hits": {"hits": list[object]()}}) - bulk_calls: list[BulkCall] = field(default_factory=list) - mget_calls: list[dict[str, object]] = field(default_factory=list) - search_calls: list[dict[str, object]] = field(default_factory=list) - close_count: int = 0 - - def bulk(self, *, operations: list[dict[str, object]], refresh: str) -> ElasticsearchBody: - self.bulk_calls.append(BulkCall(operations=operations, refresh=refresh)) - return self.bulk_result - - def mget(self, *, index: str, ids: list[str]) -> ElasticsearchBody: - self.mget_calls.append({"index": index, "ids": ids}) - return self.mget_result - - def search(self, **kwargs: object) -> ElasticsearchBody: - self.search_calls.append(kwargs) - return self.search_result - - def close(self) -> None: - self.close_count += 1 - - -type EsHarness = tuple[ElasticsearchVectorStore, FakeElasticsearch] - - -@pytest.fixture -def es() -> EsHarness: - scope = VectorScope.model_validate(SCOPE) - client = FakeElasticsearch() - client.indices.mappings = {name: es_mapping() for name in (scope.collection_name, scope.legacy_collection_name)} - return ElasticsearchVectorStore(cast("Elasticsearch", client)), client - - -def test_elasticsearch_retry_write_readback_filter_and_cosine(es: EsHarness) -> None: - store, client = es - for _ in range(2): - execute_vector_request(store, request("upsert", points=[POINT])) - write = client.bulk_calls[-1] - assert write.operations[0] == {"index": {"_index": VectorScope.model_validate(SCOPE).collection_name, "_id": A}} - assert write.operations[1] == {key: val for key, val in POINT.items() if key != "id"} - assert write.refresh == "wait_for" - client.mget_result = {"docs": [{"_id": A, "found": True, "_source": write.operations[1]}]} - assert execute_vector_request(store, request("get", ids=[A]))["points"] == [POINT] - client.search_result = {"hits": {"hits": [{"_id": A, "_score": 0.8}]}} - matches = execute_vector_request(store, request("search", ids=[A], query_vector=[1, 0]))["matches"] - assert matches[0]["score"] == pytest.approx(0.6) - search = client.search_calls[-1] - knn = search["knn"] - assert isinstance(knn, dict) - assert knn["filter"] == {"ids": {"values": [A]}} - assert search["allow_partial_search_results"] is False - - -@pytest.mark.parametrize("operation", ["get", "search", "delete"]) -@pytest.mark.parametrize("names", ["legacy", "current", "both"]) -def test_elasticsearch_names_preserve_existing_vectors_and_cleanup(es: EsHarness, operation: str, names: str) -> None: - store, client = es - scope = VectorScope.model_validate(SCOPE) - available = { - name: mapping - for name, mapping in client.indices.mappings.items() - if names == "both" or name == (scope.legacy_collection_name if names == "legacy" else scope.collection_name) - } - client.indices.mappings = available - client.mget_result = { - "docs": [{"_id": A, "found": operation != "delete", "_source": {k: v for k, v in POINT.items() if k != "id"}}] - } - client.search_result = {"hits": {"hits": [{"_id": A, "_score": 1.0}]}} - client.bulk_result = {"items": [{"delete": {"status": 200}}]} - result = execute_vector_request( - store, request(operation, ids=[A], **({"query_vector": [1, 0], "limit": 1} if operation == "search" else {})) - ) - assert result == { - "points": [POINT] if operation == "get" else [], - "matches": [{"id": A, "score": 1.0}] if operation == "search" else [], - } - if operation == "delete": - assert set().union(*(call.target_indexes("delete") for call in client.bulk_calls)) == set(available) - else: - calls = client.mget_calls if operation == "get" else client.search_calls - assert {call["index"] for call in calls} == set(available) - assert client.indices.created_indexes == [] - - -@pytest.mark.parametrize("race", [False, True]) -def test_elasticsearch_first_writer_or_compatible_creation_race(es: EsHarness, race: bool) -> None: - store, client = es - mapping = dict(client.indices.mappings) - client.indices.queued_get_mapping_results = [es_error(404, "index_not_found_exception"), mapping] - if race: - client.indices.create_failure = es_error(400, "resource_already_exists_exception") - execute_vector_request(store, request("upsert", points=[POINT])) - created = client.indices.created_mappings[-1] - properties = created["properties"] - assert isinstance(properties, dict) - assert properties["vector"]["index_options"] == {"type": "hnsw"} - assert created["dynamic"] == "strict" - - -def test_elasticsearch_creation_error_does_not_hide_configuration_failure(es: EsHarness) -> None: - store, client = es - client.indices.mappings = {} - client.indices.create_failure = es_error(400, "mapper_parsing_exception") - with pytest.raises(BadRequestError): - execute_vector_request(store, request("upsert", points=[POINT])) - assert client.bulk_calls == [] - - -@pytest.mark.parametrize("operation", ["search", "get", "delete"]) -def test_elasticsearch_missing_index_is_explicit_for_search(es: EsHarness, operation: str) -> None: - store, client = es - client.indices.mappings = {} - payload = request(operation, ids=[A], **({"query_vector": [1, 0]} if operation == "search" else {})) - if operation == "search": - with pytest.raises(VectorStoreUnavailableError, match="missing"): - execute_vector_request(store, payload) - else: - assert execute_vector_request(store, payload) == {"points": [], "matches": []} - assert client.indices.created_indexes == [] - - -@pytest.mark.parametrize("change", [{"dims": 3}, {"similarity": "l2_norm"}, {"index": False}, {"element_type": "byte"}]) -def test_elasticsearch_incompatible_index_is_never_overwritten(es: EsHarness, change: dict[str, object]) -> None: - store, client = es - mapping = next(iter(client.indices.mappings.values()))["mappings"] - mapping["properties"]["vector"].update(change) - with pytest.raises(VectorStoreUnavailableError, match="incompatible"): - execute_vector_request(store, request("upsert", points=[POINT])) - assert client.bulk_calls == [] - assert client.indices.created_indexes == [] - - -def test_elasticsearch_dimension_cap_fails_before_native_io(es: EsHarness) -> None: - store, client = es - scope = {**SCOPE, "dimension": 4097} - with pytest.raises(VectorStoreUnavailableError, match="4096"): - execute_vector_request(store, VectorRequest(operation="get", scope=VectorScope.model_validate(scope), ids=[A])) - assert client.indices.get_mapping_calls == [] - - -@pytest.mark.parametrize("result", [{"items": []}, {"items": [{"index": {"status": 429}}]}]) -def test_elasticsearch_partial_write_is_not_success(es: EsHarness, result: dict[str, object]) -> None: - store, client = es - client.bulk_result = result - with pytest.raises(VectorStoreUnavailableError, match="acknowledged"): - execute_vector_request(store, request("upsert", points=[POINT])) - - -@pytest.mark.parametrize("result", [{"docs": []}, {"docs": [{"error": "shard unavailable"}]}]) -def test_elasticsearch_partial_read_is_not_missing_data(es: EsHarness, result: dict[str, object]) -> None: - store, client = es - client.mget_result = result - with pytest.raises(VectorStoreUnavailableError, match="incomplete"): - execute_vector_request(store, request("get", ids=[A])) - - -@pytest.mark.parametrize("failure", [{"timed_out": True}, {"_shards": {"failed": 1}}]) -def test_elasticsearch_partial_search_is_not_success(es: EsHarness, failure: dict[str, object]) -> None: - store, client = es - client.search_result = failure - with pytest.raises(VectorStoreUnavailableError, match="incomplete"): - execute_vector_request(store, request("search", ids=[A], query_vector=[1, 0])) - - -@pytest.mark.parametrize("status", [200, 404]) -def test_elasticsearch_deletion_requires_visible_absence(es: EsHarness, status: int) -> None: - store, client = es - client.bulk_result = {"items": [{"delete": {"status": status}}]} - client.mget_result = {"docs": [{"_id": A, "found": False}]} - execute_vector_request(store, request("delete", ids=[A])) - client.mget_result = {"docs": [{"_id": A, "found": True, "_source": {k: v for k, v in POINT.items() if k != "id"}}]} - with pytest.raises(VectorStoreUnavailableError, match="not yet visible"): - execute_vector_request(store, request("delete", ids=[A])) - - -@pytest.mark.parametrize( - ("host", "expected"), - [ - ("elastic.internal", "http://elastic.internal:9200"), - ("https://host:9243", "https://host:9243"), - ("http://[::1]:9201", "http://[::1]:9201"), - ], -) -def test_elasticsearch_uses_dify_config_and_closes_on_failure( - config_overrides: Callable[..., None], host: str, expected: str -) -> None: - config_overrides( - VECTOR_STORE="elasticsearch", - ELASTICSEARCH_HOST=host, - ELASTICSEARCH_USE_CLOUD=False, - ELASTICSEARCH_CA_CERTS="/ca.pem", - ELASTICSEARCH_VERIFY_CERTS=True, - ) - with patch("elasticsearch.Elasticsearch") as native: - with pytest.raises(RuntimeError): - fail_with_configured_store(ElasticsearchVectorStore) - assert native.call_args.kwargs["hosts"] == [expected] - assert native.call_args.kwargs["request_timeout"] == 60 - assert native.call_args.kwargs["max_retries"] == 2 - if expected.startswith("https"): - assert native.call_args.kwargs["ca_certs"] == "/ca.pem" - native.return_value.close.assert_called_once() - - -def test_elasticsearch_cloud_key_and_missing_endpoints(config_overrides: Callable[..., None]) -> None: - config_overrides( - VECTOR_STORE="elasticsearch", - ELASTICSEARCH_USE_CLOUD=True, - ELASTICSEARCH_CLOUD_URL="https://cloud:443", - ELASTICSEARCH_API_KEY="secret", - ) - with patch("elasticsearch.Elasticsearch") as native: - with configured_vector_client(A): - assert native.call_args.kwargs["api_key"] == "secret" - assert "basic_auth" not in native.call_args.kwargs - config_overrides(ELASTICSEARCH_CLOUD_URL="") - with pytest.raises(VectorStoreUnavailableError, match="endpoint"): - ElasticsearchVectorStore.from_config() - config_overrides(ELASTICSEARCH_USE_CLOUD=False, ELASTICSEARCH_HOST="") - with pytest.raises(VectorStoreUnavailableError, match="endpoint"): - ElasticsearchVectorStore.from_config() - - -@dataclass -class WeaviateMetadata: - distance: float | None = None - - -@dataclass -class WeaviateObject: - uuid: str = A - properties: Mapping[str, object] = field(default_factory=dict) - vector: dict[str, list[float]] = field(default_factory=dict) - metadata: WeaviateMetadata = field(default_factory=WeaviateMetadata) - - -@dataclass -class WeaviateQueryResult: - objects: list[WeaviateObject] = field(default_factory=list) - - -@dataclass -class WeaviateInsertResult: - has_errors: bool - uuids: dict[int, str] - - -@dataclass -class WeaviateDeleteResult: - failed: int = 0 - - -@dataclass(frozen=True) -class WeaviateVectorIndexConfig: - distance_metric: VectorDistances - - -@dataclass(frozen=True) -class WeaviateMultiTenancyConfig: - enabled: bool - - -@dataclass(frozen=True) -class WeaviateProperty: - name: str - data_type: DataType - - -@dataclass -class WeaviateCollectionConfig: - description: str - vectorizer: str - vector_config: dict[str, object] | None - vector_index_config: WeaviateVectorIndexConfig | None - multi_tenancy_config: WeaviateMultiTenancyConfig - properties: list[WeaviateProperty] - - -@dataclass -class FakeWeaviateConfigAccessor: - current: WeaviateCollectionConfig = field( - default_factory=lambda: WeaviateCollectionConfig( - description="KnowledgeFS v1; dimension=2; metric=cosine", - vectorizer="none", - vector_config=None, - vector_index_config=WeaviateVectorIndexConfig(distance_metric=VectorDistances.COSINE), - multi_tenancy_config=WeaviateMultiTenancyConfig(enabled=False), - properties=[ - WeaviateProperty(name=name, data_type=DataType.TEXT) for name in ("generation_id", "content_hash") - ], - ) - ) - - def get(self) -> WeaviateCollectionConfig: - return self.current - - -@dataclass -class FakeWeaviateData: - insert_result: WeaviateInsertResult = field( - default_factory=lambda: WeaviateInsertResult(has_errors=False, uuids={0: A}) - ) - delete_result: WeaviateDeleteResult = field(default_factory=WeaviateDeleteResult) - insert_calls: list[list[WeaviatePoint]] = field(default_factory=list) - delete_calls: list[object] = field(default_factory=list) - - def insert_many(self, objects: list[WeaviatePoint]) -> WeaviateInsertResult: - self.insert_calls.append(objects) - return self.insert_result - - def delete_many(self, *, where: object) -> WeaviateDeleteResult: - self.delete_calls.append(where) - return self.delete_result - - -@dataclass -class FakeWeaviateQuery: - fetch_result: WeaviateQueryResult = field(default_factory=WeaviateQueryResult) - near_result: WeaviateQueryResult = field(default_factory=WeaviateQueryResult) - fetch_calls: list[dict[str, object]] = field(default_factory=list) - near_vector_calls: list[dict[str, object]] = field(default_factory=list) - - def fetch_objects(self, **kwargs: object) -> WeaviateQueryResult: - self.fetch_calls.append(kwargs) - return self.fetch_result - - def near_vector(self, **kwargs: object) -> WeaviateQueryResult: - self.near_vector_calls.append(kwargs) - return self.near_result - - -@dataclass -class FakeWeaviateCollection: - config: FakeWeaviateConfigAccessor = field(default_factory=FakeWeaviateConfigAccessor) - data: FakeWeaviateData = field(default_factory=FakeWeaviateData) - query: FakeWeaviateQuery = field(default_factory=FakeWeaviateQuery) - - -@dataclass -class FakeWeaviateCollectionHandle: - collection: FakeWeaviateCollection - consistency_levels: list[ConsistencyLevel] - - def with_consistency_level(self, level: ConsistencyLevel) -> FakeWeaviateCollection: - self.consistency_levels.append(level) - return self.collection - - -@dataclass -class FakeWeaviateCollections: - """Every collection exists unless `existing` or queued answers say otherwise.""" - - collection: FakeWeaviateCollection = field(default_factory=FakeWeaviateCollection) - existing: set[str] | None = None - queued_exists: list[bool] = field(default_factory=list) - create_failure: Exception | None = None - exists_calls: list[str] = field(default_factory=list) - create_calls: list[dict[str, object]] = field(default_factory=list) - use_calls: list[str] = field(default_factory=list) - consistency_levels: list[ConsistencyLevel] = field(default_factory=list) - - def exists(self, name: str) -> bool: - self.exists_calls.append(name) - if self.queued_exists: - return self.queued_exists.pop(0) - return self.existing is None or name in self.existing - - def create(self, **kwargs: object) -> None: - self.create_calls.append(kwargs) - if self.create_failure is not None: - raise self.create_failure - - def use(self, name: str) -> FakeWeaviateCollectionHandle: - self.use_calls.append(name) - return FakeWeaviateCollectionHandle(collection=self.collection, consistency_levels=self.consistency_levels) - - -@dataclass -class FakeWeaviateClient: - collections: FakeWeaviateCollections = field(default_factory=FakeWeaviateCollections) - close_count: int = 0 - - def close(self) -> None: - self.close_count += 1 - - -type WeaviateHarness = tuple[WeaviateVectorStore, FakeWeaviateClient, FakeWeaviateCollection] - - -@pytest.fixture -def weaviate() -> WeaviateHarness: - client = FakeWeaviateClient() - return WeaviateVectorStore(cast("WeaviateClient", client)), client, client.collections.collection - - -def test_weaviate_retry_readback_authorized_search_and_cosine(weaviate: WeaviateHarness) -> None: - store, client, collection = weaviate - for _ in range(2): - execute_vector_request(store, request("upsert", points=[POINT])) - name = VectorScope.model_validate(SCOPE).collection_name - assert client.collections.use_calls[-1] == name[0].upper() + name[1:] - assert store.collection_name(VectorScope.model_validate(SCOPE)) == client.collections.use_calls[-1] - point = collection.data.insert_calls[-1][0] - assert str(point.uuid) == A - assert point.properties == {"generation_id": GENERATION, "content_hash": "a" * 64} - obj = WeaviateObject( - uuid=A, properties=point.properties, vector={"default": [3, 4]}, metadata=WeaviateMetadata(distance=0.4) - ) - collection.query.fetch_result = WeaviateQueryResult(objects=[obj]) - assert execute_vector_request(store, request("get", ids=[A]))["points"] == [POINT] - collection.query.near_result = WeaviateQueryResult(objects=[obj]) - result = execute_vector_request(store, request("search", ids=[A], query_vector=[1, 0])) - assert result["matches"] == [{"id": A, "score": 0.6}] - assert collection.query.near_vector_calls[-1]["filters"] == Filter.by_id().contains_any([A]) - assert client.collections.consistency_levels[-1] == ConsistencyLevel.ALL - assert client.collections.use_calls[-1].startswith("Knowledgefs_v1_") - - -@pytest.mark.parametrize("operation", ["get", "search", "delete"]) -@pytest.mark.parametrize("names", ["legacy", "current", "both"]) -def test_weaviate_names_preserve_existing_vectors_and_cleanup( - weaviate: WeaviateHarness, operation: str, names: str -) -> None: - store, client, collection = weaviate - scope = VectorScope.model_validate(SCOPE) - available = { - name[0].upper() + name[1:] - for name in (scope.collection_name, scope.legacy_collection_name) - if names == "both" or name == (scope.legacy_collection_name if names == "legacy" else scope.collection_name) - } - client.collections.existing = available - obj = WeaviateObject( - uuid=A, - properties={k: v for k, v in POINT.items() if k not in {"id", "vector"}}, - vector={"default": [3, 4]}, - metadata=WeaviateMetadata(distance=0.0), - ) - collection.query.fetch_result = WeaviateQueryResult(objects=[] if operation == "delete" else [obj]) - collection.query.near_result = WeaviateQueryResult(objects=[obj]) - result = execute_vector_request( - store, request(operation, ids=[A], **({"query_vector": [1, 0], "limit": 1} if operation == "search" else {})) - ) - assert result == { - "points": [POINT] if operation == "get" else [], - "matches": [{"id": A, "score": 1.0}] if operation == "search" else [], - } - assert set(client.collections.use_calls) == available - if operation == "delete": - assert len(collection.data.delete_calls) == len(available) - assert client.collections.create_calls == [] - - -@pytest.mark.parametrize("race", [False, True]) -def test_weaviate_create_or_compatible_race(weaviate: WeaviateHarness, race: bool) -> None: - store, client, collection = weaviate - client.collections.queued_exists = [False, True] - if race: - client.collections.create_failure = RuntimeError("concurrent create") - execute_vector_request(store, request("upsert", points=[POINT])) - description = client.collections.create_calls[-1]["description"] - assert isinstance(description, str) - assert description.endswith("dimension=2; metric=cosine") - assert len(collection.data.insert_calls) == 1 - - -def test_weaviate_failed_creation_is_not_hidden(weaviate: WeaviateHarness) -> None: - store, client, collection = weaviate - client.collections.existing = set() - client.collections.create_failure = RuntimeError("creation failed") - with pytest.raises(RuntimeError): - execute_vector_request(store, request("upsert", points=[POINT])) - assert collection.data.insert_calls == [] - - -@pytest.mark.parametrize("operation", ["search", "get", "delete"]) -def test_weaviate_missing_collection_search_fails(weaviate: WeaviateHarness, operation: str) -> None: - store, client, collection = weaviate - client.collections.existing = set() - payload = request(operation, ids=[A], **({"query_vector": [1, 0]} if operation == "search" else {})) - if operation == "search": - with pytest.raises(VectorStoreUnavailableError, match="missing"): - execute_vector_request(store, payload) - else: - assert execute_vector_request(store, payload) == {"points": [], "matches": []} - assert collection.data.insert_calls == [] - - -@pytest.mark.parametrize( - ("field", "value"), - [ - ("description", "wrong dimensions"), - ("vectorizer", "text2vec-openai"), - ("vector_config", {"named": {}}), - ("vector_index_config", None), - ("properties", []), - ], -) -def test_weaviate_incompatible_schema_is_never_written(weaviate: WeaviateHarness, field: str, value: object) -> None: - store, client, collection = weaviate - setattr(collection.config.current, field, value) - with pytest.raises(VectorStoreUnavailableError, match="incompatible"): - execute_vector_request(store, request("upsert", points=[POINT])) - assert collection.data.insert_calls == [] - - -@pytest.mark.parametrize( - "result", [WeaviateInsertResult(has_errors=True, uuids={}), WeaviateInsertResult(has_errors=False, uuids={})] -) -def test_weaviate_partial_batch_is_not_success(weaviate: WeaviateHarness, result: WeaviateInsertResult) -> None: - store, client, collection = weaviate - collection.data.insert_result = result - with pytest.raises(VectorStoreUnavailableError, match="acknowledged"): - execute_vector_request(store, request("upsert", points=[POINT])) - - -def test_weaviate_search_missing_distance_is_not_success(weaviate: WeaviateHarness) -> None: - store, client, collection = weaviate - collection.query.near_result = WeaviateQueryResult( - objects=[WeaviateObject(metadata=WeaviateMetadata(distance=None))] - ) - with pytest.raises(VectorStoreUnavailableError, match="omitted distance"): - execute_vector_request(store, request("search", ids=[A], query_vector=[1, 0])) - - -def test_weaviate_delete_requires_ack_and_visible_absence(weaviate: WeaviateHarness) -> None: - store, client, collection = weaviate - for _ in range(2): - execute_vector_request(store, request("delete", ids=[A])) - collection.data.delete_result.failed = 1 - with pytest.raises(VectorStoreUnavailableError, match="not yet visible"): - execute_vector_request(store, request("delete", ids=[A])) - collection.data.delete_result.failed = 0 - collection.query.fetch_result.objects = [WeaviateObject(uuid=A)] - with pytest.raises(VectorStoreUnavailableError, match="not yet visible"): - execute_vector_request(store, request("delete", ids=[A])) - - -@pytest.mark.parametrize( - ("endpoint", "grpc", "secure", "port"), - [ - ("http://host:8080", None, False, 50051), - ("https://cloud", None, True, 443), - ("http://host", "grpcs://rpc:443", True, 443), - ("host", "rpc:50052", False, 50052), - ], -) -def test_weaviate_config_and_resource_cleanup( - config_overrides: Callable[..., None], endpoint: str, grpc: str | None, secure: bool, port: int -) -> None: - config_overrides( - VECTOR_STORE="weaviate", WEAVIATE_ENDPOINT=endpoint, WEAVIATE_GRPC_ENDPOINT=grpc, WEAVIATE_API_KEY="secret" - ) - with patch("weaviate.connect_to_custom") as connect: - with pytest.raises(RuntimeError): - fail_with_configured_store(WeaviateVectorStore) - assert connect.call_args.kwargs["grpc_secure"] is secure - assert connect.call_args.kwargs["grpc_port"] == port - connect.return_value.close.assert_called_once() - - -def test_weaviate_missing_and_invalid_endpoints(config_overrides: Callable[..., None]) -> None: - config_overrides(VECTOR_STORE="weaviate", WEAVIATE_ENDPOINT="") - with pytest.raises(VectorStoreUnavailableError, match="endpoint"): - WeaviateVectorStore.from_config() - config_overrides(WEAVIATE_ENDPOINT="http://host", WEAVIATE_GRPC_ENDPOINT="https://wrong") - with pytest.raises(VectorStoreUnavailableError, match="gRPC"): - WeaviateVectorStore.from_config() diff --git a/api/tests/unit_tests/services/test_knowledge_fs_vector_store.py b/api/tests/unit_tests/services/test_knowledge_fs_vector_store.py deleted file mode 100644 index 3895ce5e2f9..00000000000 --- a/api/tests/unit_tests/services/test_knowledge_fs_vector_store.py +++ /dev/null @@ -1,732 +0,0 @@ -from __future__ import annotations - -from collections.abc import Callable, Iterator -from dataclasses import dataclass, field -from typing import TYPE_CHECKING, cast -from unittest.mock import patch - -import pytest -from pydantic import ValidationError - -from models.dataset import TidbAuthBinding -from models.enums import TidbAuthBindingStatus -from services.knowledge_fs.vector_store import ( - TIDB_POINT_ID_FIELD, - VectorRequest, - VectorScope, - VectorStoreUnavailableError, - configured_vector_client, - execute_vector_request, -) - -if TYPE_CHECKING: - from qdrant_client import QdrantClient as QdrantClientType - from qdrant_client.http.models import CollectionInfo, FilterSelector, PointIdsList, PointStruct - - from services.knowledge_fs.vector_store import VectorResult - -QdrantClient = pytest.importorskip("qdrant_client").QdrantClient -models = pytest.importorskip("qdrant_client.http.models") -UnexpectedResponse = pytest.importorskip("qdrant_client.http.exceptions").UnexpectedResponse - -TENANT = "018f0d60-7a49-7cc2-9c1b-5b36f18f0001" -SPACE = "018f0d60-7a49-7cc2-9c1b-5b36f18f0002" -GENERATION = "018f0d60-7a49-7cc2-9c1b-5b36f18f0003" -A = "018f0d60-7a49-7cc2-9c1b-5b36f18f0004" -B = "018f0d60-7a49-7cc2-9c1b-5b36f18f0005" -SCOPE: dict[str, object] = { - "tenant_id": TENANT, - "knowledge_space_id": SPACE, - "vector_space_id": "embedding-v1", - "kind": "dense", - "dimension": 2, -} - - -@dataclass -class FakeQdrant: - """Qdrant client surface used by the bridge; unknown collections raise the remote 404 contract.""" - - collections: dict[str, CollectionInfo] = field(default_factory=dict) - search_result: list[object] = field(default_factory=list) - retrieve_result: list[object] = field(default_factory=list) - retrieve_hook: Callable[..., list[object]] | None = None - retrieve_failure: Exception | None = None - delete_failure: Exception | None = None - create_payload_index_hook: Callable[..., None] | None = None - get_collection_calls: list[str] = field(default_factory=list) - create_collection_calls: list[dict[str, object]] = field(default_factory=list) - create_payload_index_calls: list[dict[str, object]] = field(default_factory=list) - upsert_calls: list[dict[str, object]] = field(default_factory=list) - upserted_points: list[list[PointStruct]] = field(default_factory=list) - search_calls: list[dict[str, object]] = field(default_factory=list) - retrieve_calls: list[dict[str, object]] = field(default_factory=list) - delete_calls: list[dict[str, object]] = field(default_factory=list) - delete_selectors: list[PointIdsList | FilterSelector] = field(default_factory=list) - close_count: int = 0 - - def execute(self, payload: VectorRequest) -> VectorResult: - """Run the bridge against this fake through the Qdrant client contract.""" - return execute_vector_request(cast("QdrantClientType", self), payload) - - def get_collection(self, name: str) -> CollectionInfo: - self.get_collection_calls.append(name) - if name not in self.collections: - raise UnexpectedResponse(404, "missing", b"{}", {}) - return self.collections[name] - - def create_collection(self, **kwargs: object) -> None: - self.create_collection_calls.append(kwargs) - - def create_payload_index(self, **kwargs: object) -> None: - self.create_payload_index_calls.append(kwargs) - if self.create_payload_index_hook is not None: - self.create_payload_index_hook(**kwargs) - - def upsert(self, *, collection_name: str, wait: bool, points: list[PointStruct]) -> None: - self.upsert_calls.append({"collection_name": collection_name, "wait": wait, "points": points}) - self.upserted_points.append(points) - - def search(self, **kwargs: object) -> list[object]: - self.search_calls.append(kwargs) - return self.search_result - - def retrieve(self, **kwargs: object) -> list[object]: - self.retrieve_calls.append(kwargs) - if self.retrieve_failure is not None: - raise self.retrieve_failure - if self.retrieve_hook is not None: - return self.retrieve_hook(**kwargs) - return self.retrieve_result - - def delete(self, *, collection_name: str, points_selector: PointIdsList | FilterSelector, wait: bool) -> None: - self.delete_calls.append({"collection_name": collection_name, "points_selector": points_selector, "wait": wait}) - self.delete_selectors.append(points_selector) - if self.delete_failure is not None: - raise self.delete_failure - - def close(self) -> None: - self.close_count += 1 - - -def request(operation: str, **kwargs: object) -> VectorRequest: - return VectorRequest.model_validate({"operation": operation, "scope": SCOPE, **kwargs}) - - -@pytest.fixture(autouse=True) -def _vector_store_config(config_overrides: Callable[..., None]) -> None: - config_overrides(DEPLOYMENT_EDITION="COMMUNITY") - - -@pytest.fixture -def client() -> Iterator[QdrantClientType]: - client = QdrantClient(":memory:") - client.create_collection( - VectorScope.model_validate(SCOPE).collection_name, - vectors_config=models.VectorParams(size=2, distance=models.Distance.COSINE), - ) - native_get = client.get_collection - - def get_collection(name: str) -> CollectionInfo: - # Match the remote API's 404 contract while exercising native vector I/O. - if not client.collection_exists(name): - raise UnexpectedResponse(404, "missing", b"{}", {}) - return native_get(name) - - with patch.object(client, "get_collection", side_effect=get_collection): - yield client - client.close() - - -def test_native_vector_roundtrip_authorized_id_filter_and_idempotent_deletion(client: QdrantClientType) -> None: - points = [ - {"id": A, "generation_id": GENERATION, "content_hash": "a" * 64, "vector": [3, 4]}, - {"id": B, "generation_id": GENERATION, "content_hash": "b" * 64, "vector": [1, 0]}, - ] - execute_vector_request(client, request("upsert", points=points)) - execute_vector_request(client, request("upsert", points=points)) - actual = execute_vector_request(client, request("get", ids=[A, B]))["points"] - assert len(actual) == 2 - assert actual[0]["vector"] == pytest.approx([0.6, 0.8]) - assert actual[0]["content_hash"] == "a" * 64 - # The nearest point B is unauthorized. It must not crowd out the authorized A. - matches = execute_vector_request(client, request("search", ids=[A], query_vector=[1, 0], limit=1))["matches"] - assert [item["id"] for item in matches] == [A] - assert matches[0]["score"] == pytest.approx(0.6) - for _ in range(2): - execute_vector_request(client, request("delete", ids=[A])) - assert execute_vector_request(client, request("get", ids=[A]))["points"] == [] - assert len(execute_vector_request(client, request("get", ids=[B]))["points"]) == 1 - - -@pytest.mark.parametrize( - ("field", "value"), - [("tenant_id", B), ("knowledge_space_id", B), ("vector_space_id", "other"), ("kind", "visual"), ("dimension", 3)], -) -def test_each_scope_dimension_has_a_distinct_reserved_namespace(field: str, value: object) -> None: - original = VectorScope.model_validate(SCOPE) - other = VectorScope.model_validate({**SCOPE, field: value}) - assert other.collection_name != original.collection_name - assert other.legacy_collection_name != original.legacy_collection_name - assert original.collection_name.startswith("knowledgefs_v1_dense_") - assert SPACE not in original.collection_name - - -@pytest.mark.parametrize("kind", ["dense", "visual", "graph-entity", "graph-relation"]) -def test_typed_collection_names_fit_tidb_delete_endpoint_and_normalize_uuid_case(kind: str) -> None: - scope = VectorScope.model_validate({**SCOPE, "kind": kind}) - prefix = f"knowledgefs_v1_{kind.replace('-', '_')}_" - assert scope.collection_name.startswith(prefix) - assert len(scope.collection_name.removeprefix(prefix)) == min(32, 56 - len(prefix)) - assert len(scope.collection_name.removeprefix(prefix)) >= 26 - assert len(scope.collection_name) <= 56 - assert VectorScope.model_validate({**SCOPE, "kind": kind, "tenant_id": TENANT.upper()}) == scope - - -def test_legacy_namespace_keeps_pre_rename_hash() -> None: - # Pin the pre-rename format so compatibility cannot drift with the new one. - scope = VectorScope.model_validate(SCOPE) - assert scope.legacy_collection_name == "knowledgefs_v1_488d91e70fe33beedcda4cb5e47b954dd5af7ca3" - assert ( - scope.collection_name - == "knowledgefs_v1_dense_" + scope.legacy_collection_name.removeprefix("knowledgefs_v1_")[:32] - ) - - -def test_legacy_and_typed_points_remain_searchable_and_deletable_together(client: QdrantClientType) -> None: - scope = VectorScope.model_validate(SCOPE) - legacy = scope.legacy_collection_name - current = scope.collection_name - client.create_collection(legacy, vectors_config=models.VectorParams(size=2, distance=models.Distance.COSINE)) - points = [ - {"id": A, "generation_id": GENERATION, "content_hash": "a" * 64, "vector": [3, 4]}, - {"id": B, "generation_id": GENERATION, "content_hash": "b" * 64, "vector": [1, 0]}, - ] - client.upsert( - legacy, - points=[ - models.PointStruct( - id=point["id"], - vector=point["vector"], - payload={k: v for k, v in point.items() if k not in {"id", "vector"}}, - ) - for point in points - ], - ) - client.delete_collection(current) - assert len(execute_vector_request(client, request("get", ids=[A, B]))["points"]) == 2 - assert execute_vector_request(client, request("search", ids=[B], query_vector=[1, 0]))["matches"][0]["id"] == B - assert not client.collection_exists(current) - - # A new document in this scope must create the typed name, without hiding old A. - execute_vector_request(client, request("upsert", points=[points[1]])) - assert client.count(current).count == 1 - assert client.count(legacy).count == 2 - assert len(execute_vector_request(client, request("get", ids=[A, B]))["points"]) == 2 - matches = execute_vector_request(client, request("search", ids=[A, B], query_vector=[1, 0], limit=1))["matches"] - assert [match["id"] for match in matches] == [B] - # B exists in both collections, but cannot crowd out an authorized point A. - matches = execute_vector_request(client, request("search", ids=[A], query_vector=[1, 0], limit=1))["matches"] - assert [match["id"] for match in matches] == [A] - for _ in range(2): - execute_vector_request(client, request("delete", ids=[B])) - assert client.retrieve(current, [B]) == [] - assert client.retrieve(legacy, [B]) == [] - assert len(execute_vector_request(client, request("get", ids=[A]))["points"]) == 1 - - -@pytest.mark.parametrize("operation", ["get", "search", "delete"]) -def test_legacy_backend_failure_never_returns_partial_success(client: QdrantClientType, operation: str) -> None: - execute_vector_request( - client, - request("upsert", points=[{"id": A, "generation_id": GENERATION, "content_hash": "a" * 64, "vector": [1, 0]}]), - ) - native_get = client.get_collection - - def get_collection(name: str) -> CollectionInfo: - if name == VectorScope.model_validate(SCOPE).legacy_collection_name: - raise UnexpectedResponse(503, "unavailable", b"{}", {}) - return native_get(name) - - payload = request(operation, ids=[A], **({"query_vector": [1, 0]} if operation == "search" else {})) - with patch.object(client, "get_collection", side_effect=get_collection): - with pytest.raises(UnexpectedResponse): - execute_vector_request(client, payload) - - -@pytest.mark.parametrize( - "payload", - [ - {"operation": "search", "ids": [], "query_vector": [1, 0]}, - {"operation": "search", "ids": [A], "query_vector": [1]}, - {"operation": "search", "ids": [A], "query_vector": [0, 0]}, - {"operation": "search", "ids": [A], "query_vector": [float("nan"), 1]}, - {"operation": "search", "ids": [A, A], "query_vector": [1, 0]}, - {"operation": "search", "ids": [A], "query_vector": [1, 0], "collection_name": "dataset"}, - {"operation": "search", "ids": [A], "query_vector": [1, 0], "endpoint": "http://example.com"}, - {"operation": "get", "ids": ["not-a-uuid"]}, - {"operation": "delete", "ids": [A], "query_vector": [1, 0]}, - {"operation": "upsert", "points": []}, - ], -) -def test_invalid_operations_fail_before_backend_io(payload: dict[str, object]) -> None: - with pytest.raises(ValidationError): - VectorRequest.model_validate({"scope": SCOPE, **payload}) - - -def test_missing_collection_search_is_not_reported_as_empty_success() -> None: - client = FakeQdrant() - with pytest.raises(VectorStoreUnavailableError, match="missing"): - client.execute(request("search", ids=[A], query_vector=[1, 0])) - assert client.create_collection_calls == [] - assert client.execute(request("delete", ids=[A]))["points"] == [] - - -def test_collection_creation_race_accepts_only_compatible_existing_collection(client: QdrantClientType) -> None: - real_get = client.get_collection - with ( - patch.object( - client, - "get_collection", - side_effect=[ - UnexpectedResponse(404, "missing", b"{}", {}), - real_get(VectorScope.model_validate(SCOPE).collection_name), - ], - ), - patch.object(client, "create_collection", side_effect=RuntimeError("already exists")), - ): - execute_vector_request( - client, - request( - "upsert", points=[{"id": A, "generation_id": GENERATION, "content_hash": "a" * 64, "vector": [1, 0]}] - ), - ) - assert len(execute_vector_request(client, request("get", ids=[A]))["points"]) == 1 - - -def test_no_binding_never_allocates_a_paid_cluster(config_overrides: Callable[..., None]) -> None: - from services.tidb_binding_service import TidbBindingUnavailableError - - config_overrides(VECTOR_STORE="tidb_on_qdrant") - with patch("services.tidb_binding_service._load_binding", return_value=None): - with pytest.raises(TidbBindingUnavailableError, match="no active"): - with configured_vector_client(TENANT): - pytest.fail("unconfigured backend was opened") - - -def test_unsupported_backend_fails_explicitly(config_overrides: Callable[..., None]) -> None: - config_overrides(VECTOR_STORE="milvus") - with pytest.raises(VectorStoreUnavailableError, match="does not support"): - with configured_vector_client(TENANT): - pytest.fail("unsupported backend was opened") - - -@pytest.mark.parametrize("allow_create", [True, False]) -def test_tidb_client_uses_tenant_binding(config_overrides: Callable[..., None], allow_create: bool) -> None: - config_overrides(VECTOR_STORE="tidb_on_qdrant", TIDB_ON_QDRANT_URL="https://global.invalid") - binding = TidbAuthBinding( - tenant_id=TENANT, - cluster_id="cluster-1", - cluster_name="tenant-cluster", - active=True, - status=TidbAuthBindingStatus.ACTIVE, - account="tenant-account", - password="tenant-secret", - qdrant_endpoint="https://tenant.invalid", - ) - client = FakeQdrant() - - def open_client(**kwargs: object) -> FakeQdrant: - assert kwargs["url"] == "https://tenant.invalid" - assert kwargs["api_key"] == "tenant-account:tenant-secret" - return client - - with ( - patch("services.tidb_binding_service.resolve_tidb_auth_binding", return_value=binding) as resolve, - patch("qdrant_client.QdrantClient", side_effect=open_client), - ): - with configured_vector_client(TENANT, allow_create=allow_create) as configured: - assert configured is client - resolve.assert_called_once_with(TENANT, allow_create=allow_create) - assert client.close_count == 1 - - -def test_qdrant_configuration_closes_client_even_when_operation_fails(config_overrides: Callable[..., None]) -> None: - config_overrides(VECTOR_STORE="qdrant", QDRANT_URL="https://configured.invalid", QDRANT_API_KEY="private") - with patch("qdrant_client.QdrantClient") as factory: - with pytest.raises(RuntimeError): - with configured_vector_client(TENANT): - raise RuntimeError("operation failed") - assert factory.call_args.kwargs["url"] == "https://configured.invalid" - factory.return_value.close.assert_called_once() - config_overrides(QDRANT_URL="") - with pytest.raises(VectorStoreUnavailableError, match="endpoint"): - with configured_vector_client(TENANT): - pytest.fail("missing endpoint accepted") - - -def test_cloud_quota_is_enforced_before_any_vector_mutation() -> None: - from services.vector_space_admission_service import VectorSpaceAdmissionError - - client = FakeQdrant() - with patch( - "services.vector_space_admission_service.VectorSpaceAdmissionService.ensure_external_points_can_be_indexed", - side_effect=VectorSpaceAdmissionError("limit reached"), - ): - with pytest.raises(VectorSpaceAdmissionError): - client.execute( - request( - "upsert", - points=[{"id": A, "generation_id": GENERATION, "content_hash": "a" * 64, "vector": [1, 0]}], - ), - ) - assert client.get_collection_calls == [] - assert client.upsert_calls == [] - - -def test_cloud_quota_uses_existing_watermark_without_embedding_calls(config_overrides: Callable[..., None]) -> None: - from enums import CloudPlan, DeploymentEdition - from services.vector_space_admission_service import VectorSpaceAdmissionError, VectorSpaceAdmissionService - - config_overrides( - DEPLOYMENT_EDITION=DeploymentEdition.CLOUD, - VECTOR_STORE="tidb_on_qdrant", - TIDB_ON_QDRANT_ESTIMATED_STORAGE_LIMITS_MB="sandbox:60,professional:6144,team:24576", - ) - service = VectorSpaceAdmissionService() - with ( - patch.object(service, "_get_plan", return_value=CloudPlan.SANDBOX), - patch.object(service, "_get_usage_and_limit_mb", return_value=(0.0, 50)), - patch.object(service, "_get_embedding_dimension") as dimension, - patch.object(service, "_reserve_projected_usage", return_value=(0, 60 * 1024 * 1024)) as reserve, - ): - for _ in range(2): - service.ensure_external_points_can_be_indexed( - tenant_id=TENANT, batch_id="stable-batch", point_count=2, dimension=3 - ) - assert reserve.call_args.kwargs["document_id"] == "knowledgefs:stable-batch" - assert reserve.call_args.kwargs["document_estimate_bytes"] == 2 * (3 * 8 + 3584) - assert reserve.call_args.kwargs["estimate_limit_bytes"] == 60 * 1024 * 1024 - dimension.assert_not_called() - reserve.return_value = (0, 61 * 1024 * 1024) - with pytest.raises(VectorSpaceAdmissionError, match="exceeding the 50 MB limit"): - service.ensure_external_points_can_be_indexed( - tenant_id=TENANT, batch_id="large", point_count=2, dimension=3 - ) - with pytest.raises(VectorSpaceAdmissionError, match="Invalid"): - service.ensure_external_points_can_be_indexed(tenant_id=TENANT, batch_id="bad", point_count=0, dimension=3) - - -def test_unmetered_workspace_does_not_contact_billing_or_reserve_quota(config_overrides: Callable[..., None]) -> None: - from services.vector_space_admission_service import VectorSpaceAdmissionService - - config_overrides(DEPLOYMENT_EDITION="CLOUD", VECTOR_STORE="tidb_on_qdrant") - service = VectorSpaceAdmissionService() - with ( - patch.object(service, "_get_plan", return_value=None), - patch.object(service, "_get_usage_and_limit_mb") as billing, - patch.object(service, "_reserve_projected_usage") as reserve, - ): - service.ensure_external_points_can_be_indexed( - tenant_id=TENANT, batch_id="unmetered", point_count=2, dimension=3 - ) - billing.assert_not_called() - reserve.assert_not_called() - - -@pytest.mark.parametrize("limit", [50, -1]) -def test_external_points_with_unknown_usage_fail_before_reservation( - config_overrides: Callable[..., None], limit: int -) -> None: - from enums import CloudPlan - from services.vector_space_admission_service import ( - VectorSpaceAdmissionService, - VectorSpaceAdmissionUnavailableError, - ) - - config_overrides(DEPLOYMENT_EDITION="CLOUD", VECTOR_STORE="tidb_on_qdrant") - service = VectorSpaceAdmissionService() - with ( - patch.object(service, "_get_plan", return_value=CloudPlan.SANDBOX), - patch( - "services.vector_space_admission_service.BillingService.get_vector_space", - return_value={"size": 0, "limit": limit, "usage_unknown": True}, - ), - patch.object(service, "_reserve_projected_usage") as reserve, - ): - with pytest.raises(VectorSpaceAdmissionUnavailableError, match="Unable to verify vector storage usage"): - service.ensure_external_points_can_be_indexed( - tenant_id=TENANT, batch_id="unknown", point_count=2, dimension=3 - ) - reserve.assert_not_called() - - -@pytest.mark.parametrize(("edition", "vector_store"), [("COMMUNITY", "tidb_on_qdrant"), ("CLOUD", "qdrant")]) -def test_external_points_skip_billing_outside_cloud_tidb( - config_overrides: Callable[..., None], edition: str, vector_store: str -) -> None: - from services.vector_space_admission_service import VectorSpaceAdmissionService - - config_overrides(DEPLOYMENT_EDITION=edition, VECTOR_STORE=vector_store) - service = VectorSpaceAdmissionService() - with ( - patch("services.vector_space_admission_service.BillingService.get_vector_space") as billing, - patch.object(service, "_reserve_projected_usage") as reserve, - ): - service.ensure_external_points_can_be_indexed( - tenant_id=TENANT, batch_id="unmetered", point_count=2, dimension=3 - ) - billing.assert_not_called() - reserve.assert_not_called() - - -def test_incompatible_collection_is_never_recreated_or_written(client: QdrantClientType) -> None: - collection = VectorScope.model_validate(SCOPE).collection_name - client.delete_collection(collection) - client.create_collection(collection, vectors_config=models.VectorParams(size=3, distance=models.Distance.COSINE)) - with pytest.raises(VectorStoreUnavailableError, match="incompatible"): - execute_vector_request( - client, - request( - "upsert", points=[{"id": A, "generation_id": GENERATION, "content_hash": "a" * 64, "vector": [1, 0]}] - ), - ) - from qdrant_client.http.models import VectorParams - - vectors = client.get_collection(collection).config.params.vectors - assert isinstance(vectors, VectorParams) - assert vectors.size == 3 - - -def test_uncertain_native_delete_is_retryable(client: QdrantClientType) -> None: - execute_vector_request( - client, - request("upsert", points=[{"id": A, "generation_id": GENERATION, "content_hash": "a" * 64, "vector": [1, 0]}]), - ) - with patch.object(client, "delete"): - with pytest.raises(VectorStoreUnavailableError, match="not yet visible"): - execute_vector_request(client, request("delete", ids=[A])) - execute_vector_request(client, request("delete", ids=[A])) - assert execute_vector_request(client, request("get", ids=[A]))["points"] == [] - - -type TidbHarness = tuple[FakeQdrant, CollectionInfo] - - -@pytest.fixture -def tidb_client(client: QdrantClientType, config_overrides: Callable[..., None]) -> tuple[FakeQdrant, CollectionInfo]: - config_overrides(VECTOR_STORE="tidb_on_qdrant") - scope = VectorScope.model_validate(SCOPE) - info = client.get_collection(scope.collection_name) - return FakeQdrant(collections={scope.collection_name: info}), info - - -def test_tidb_new_collection_indexes_and_stores_point_ids(tidb_client: TidbHarness) -> None: - client, info = tidb_client - client.execute( - request("upsert", points=[{"id": A, "generation_id": GENERATION, "content_hash": "a" * 64, "vector": [1, 0]}]), - ) - assert client.create_payload_index_calls == [ - { - "collection_name": VectorScope.model_validate(SCOPE).collection_name, - "field_name": TIDB_POINT_ID_FIELD, - "field_schema": models.PayloadSchemaType.KEYWORD, - "wait": True, - } - ] - assert client.upserted_points[-1][0].payload == { - "generation_id": GENERATION, - "content_hash": "a" * 64, - TIDB_POINT_ID_FIELD: A, - } - - -@pytest.mark.parametrize("compatible", [False, True]) -def test_tidb_payload_index_creation_race_requires_compatible_index(tidb_client: TidbHarness, compatible: bool) -> None: - client, info = tidb_client - - def create_index(**_kwargs: object) -> None: - if compatible: - info.payload_schema[TIDB_POINT_ID_FIELD] = models.PayloadIndexInfo(data_type="keyword", points=0) - raise RuntimeError("index creation raced or failed") - - client.create_payload_index_hook = create_index - payload = request( - "upsert", points=[{"id": A, "generation_id": GENERATION, "content_hash": "a" * 64, "vector": [1, 0]}] - ) - if compatible: - client.execute(payload) - assert len(client.upsert_calls) == 1 - else: - with pytest.raises(RuntimeError, match="creation"): - client.execute(payload) - assert client.upsert_calls == [] - - -def test_tidb_nonempty_old_collection_is_not_marked_indexed(tidb_client: TidbHarness) -> None: - client, info = tidb_client - info.points_count = 1 - client.execute( - request("upsert", points=[{"id": A, "generation_id": GENERATION, "content_hash": "a" * 64, "vector": [1, 0]}]), - ) - assert client.create_payload_index_calls == [] - assert len(client.upsert_calls) == 1 - - -@pytest.mark.parametrize("operation", ["search", "upsert"]) -def test_tidb_wrong_id_index_type_fails_closed(tidb_client: TidbHarness, operation: str) -> None: - client, info = tidb_client - info.payload_schema[TIDB_POINT_ID_FIELD] = models.PayloadIndexInfo(data_type="integer", points=0) - payload = ( - request("search", ids=[A], query_vector=[1, 0]) - if operation == "search" - else request( - "upsert", points=[{"id": A, "generation_id": GENERATION, "content_hash": "a" * 64, "vector": [1, 0]}] - ) - ) - with pytest.raises(VectorStoreUnavailableError, match="ID index"): - client.execute(payload) - assert client.upsert_calls == [] - assert client.search_calls == [] - - -def test_tidb_indexed_search_uses_supported_filter_and_payload_selection(tidb_client: TidbHarness) -> None: - client, info = tidb_client - info.payload_schema[TIDB_POINT_ID_FIELD] = models.PayloadIndexInfo(data_type="keyword", points=1) - client.search_result = [models.ScoredPoint(id=A, version=0, score=0.6)] - result = client.execute(request("search", ids=[A], query_vector=[1, 0], limit=1)) - assert result["matches"] == [{"id": A, "score": 0.6}] - query = client.search_calls[-1] - assert query["with_payload"] is True - assert query["with_vectors"] is False - assert query["query_filter"] == models.Filter( - must=[models.FieldCondition(key=TIDB_POINT_ID_FIELD, match=models.MatchAny(any=[A]))] - ) - assert client.retrieve_calls == [] - client.search_result = [models.ScoredPoint(id=B, version=0, score=1.0)] - with pytest.raises(VectorStoreUnavailableError, match="unauthorized"): - client.execute(request("search", ids=[A], query_vector=[1, 0])) - - -def test_tidb_unindexed_legacy_search_reads_only_authorized_ids( - client: QdrantClientType, config_overrides: Callable[..., None] -) -> None: - scope = VectorScope.model_validate(SCOPE) - client.create_collection( - scope.legacy_collection_name, vectors_config=models.VectorParams(size=2, distance="Cosine") - ) - client.upsert( - scope.legacy_collection_name, - points=[ - models.PointStruct(id=A, vector=[3, 4]), - models.PointStruct(id=B, vector=[1, 0]), - ], - ) - config_overrides(VECTOR_STORE="tidb_on_qdrant") - with patch.object(client, "search", side_effect=AssertionError("Legacy search must not query unfiltered")): - result = execute_vector_request(client, request("search", ids=[A], query_vector=[1, 0], limit=1)) - assert [match["id"] for match in result["matches"]] == [A] - assert result["matches"][0]["score"] == pytest.approx(0.6) - result = execute_vector_request(client, request("search", ids=[A, B], query_vector=[1, 0], limit=1)) - assert result["matches"] == [{"id": B, "score": 1.0}] - - -def test_tidb_unindexed_search_bounds_reads_and_propagates_failures(tidb_client: TidbHarness) -> None: - client, info = tidb_client - info.points_count = 130 - ids = [f"00000000-0000-4000-8000-{n:012d}" for n in range(130)] - - def retrieve(**kwargs: object) -> list[object]: - assert kwargs["with_payload"] is False - assert kwargs["with_vectors"] is True - ids_batch = kwargs["ids"] - assert isinstance(ids_batch, list) - assert len(ids_batch) <= 64 - return [models.Record(id=point_id, vector=[1, 0]) for point_id in ids_batch] - - client.retrieve_hook = retrieve - result = client.execute(request("search", ids=ids, query_vector=[1, 0], limit=3)) - assert result["matches"] == [{"id": point_id, "score": 1.0} for point_id in ids[:3]] - assert len(client.retrieve_calls) == 3 - assert client.search_calls == [] - client.retrieve_failure = RuntimeError("read failed") - with pytest.raises(RuntimeError, match="read failed"): - client.execute(request("search", ids=ids, query_vector=[1, 0])) - - -@pytest.mark.parametrize("vector", [None, [1], [0, 0], [float("nan"), 0], [1e20, 0]]) -def test_tidb_unindexed_search_rejects_unscorable_vectors(tidb_client: TidbHarness, vector: list[float] | None) -> None: - client, _ = tidb_client - client.retrieve_result = [models.Record(id=A, vector=vector)] - with pytest.raises(VectorStoreUnavailableError, match="dimensions|scored"): - client.execute(request("search", ids=[A], query_vector=[1, 0])) - - -def test_tidb_unindexed_search_rejects_unrequested_points(tidb_client: TidbHarness) -> None: - client, _ = tidb_client - client.retrieve_result = [models.Record(id=B, vector=[1, 0])] - with pytest.raises(VectorStoreUnavailableError, match="unauthorized"): - client.execute(request("search", ids=[A], query_vector=[1, 0])) - - -def test_tidb_indexed_delete_uses_field_filter_and_verifies_absence(tidb_client: TidbHarness) -> None: - client, info = tidb_client - info.payload_schema[TIDB_POINT_ID_FIELD] = models.PayloadIndexInfo(data_type="keyword", points=2) - client.retrieve_result = list[object]() - client.execute(request("delete", ids=[A, B])) - assert client.delete_selectors[-1] == models.FilterSelector( - filter=models.Filter(must=[models.FieldCondition(key=TIDB_POINT_ID_FIELD, match=models.MatchAny(any=[A, B]))]) - ) - assert len(client.delete_calls) == 1 - client.retrieve_result = [models.Record(id=A)] - with pytest.raises(VectorStoreUnavailableError, match="not yet visible"): - client.execute(request("delete", ids=[A, B])) - - -def test_tidb_unindexed_delete_uses_single_ids_and_propagates_errors(tidb_client: TidbHarness) -> None: - client, _ = tidb_client - client.retrieve_result = list[object]() - client.execute(request("delete", ids=[A, B])) - assert client.create_payload_index_calls == [ - { - "collection_name": VectorScope.model_validate(SCOPE).collection_name, - "field_name": "content_hash", - "field_schema": models.PayloadSchemaType.KEYWORD, - "wait": True, - } - ] - from qdrant_client.http.models import PointIdsList as NativePointIdsList - - selectors = [selector for selector in client.delete_selectors if isinstance(selector, NativePointIdsList)] - assert len(selectors) == len(client.delete_selectors) - assert sorted(selector.points for selector in selectors) == [[A], [B]] - assert all(call["wait"] for call in client.delete_calls) - client.delete_failure = RuntimeError("delete failed") - with pytest.raises(RuntimeError, match="delete failed"): - client.execute(request("delete", ids=[A, B])) - - -@pytest.mark.parametrize("remaining", [[], [models.Record(id=A)]]) -def test_tidb_delete_not_found_still_requires_readback(tidb_client: TidbHarness, remaining: list[object]) -> None: - client, info = tidb_client - info.payload_schema[TIDB_POINT_ID_FIELD] = models.PayloadIndexInfo(data_type="keyword", points=1) - client.delete_failure = UnexpectedResponse(404, "missing", b"{}", {}) - client.retrieve_result = remaining - if remaining: - with pytest.raises(VectorStoreUnavailableError, match="not yet visible"): - client.execute(request("delete", ids=[A])) - else: - client.execute(request("delete", ids=[A])) - assert len(client.retrieve_calls) == 1 - - -@pytest.mark.parametrize("status", [404, 503]) -def test_tidb_delete_readback_only_accepts_missing_collection(tidb_client: TidbHarness, status: int) -> None: - client, info = tidb_client - info.payload_schema[TIDB_POINT_ID_FIELD] = models.PayloadIndexInfo(data_type="keyword", points=1) - client.retrieve_failure = UnexpectedResponse(status, "unavailable", b"{}", {}) - if status == 404: - client.execute(request("delete", ids=[A])) - else: - with pytest.raises(UnexpectedResponse): - client.execute(request("delete", ids=[A])) diff --git a/api/tests/unit_tests/services/test_knowledge_retrieval_inner_service.py b/api/tests/unit_tests/services/test_knowledge_retrieval_inner_service.py index 66a289d24fa..5e693750538 100644 --- a/api/tests/unit_tests/services/test_knowledge_retrieval_inner_service.py +++ b/api/tests/unit_tests/services/test_knowledge_retrieval_inner_service.py @@ -16,7 +16,7 @@ from services.errors.knowledge_retrieval import ( InnerKnowledgeRetrieveDatasetTenantMismatchError, ) from services.knowledge_retrieval_inner_service import InnerKnowledgeRetrievalService -from tests.unit_tests.model_factories import make_app, make_dataset +from tests.helpers.unit.model_factories import make_app, make_dataset TENANT_ID = "11111111-1111-1111-1111-111111111111" OTHER_TENANT_ID = "22222222-2222-2222-2222-222222222222" diff --git a/api/tests/unit_tests/services/test_legacy_workflow_task_recovery.py b/api/tests/unit_tests/services/test_legacy_workflow_task_recovery.py deleted file mode 100644 index e7943f0c8a0..00000000000 --- a/api/tests/unit_tests/services/test_legacy_workflow_task_recovery.py +++ /dev/null @@ -1,426 +0,0 @@ -"""Legacy pauses recover task IDs on demand without an offline backfill.""" - -import json -from collections.abc import Callable -from datetime import datetime -from uuid import UUID, uuid4 - -import pytest -from redis import Redis -from sqlalchemy import event, select -from sqlalchemy.engine import Connection, Engine -from sqlalchemy.orm import Session, sessionmaker - -import services.app_task_service as task_module -from core.app.app_config.entities import WorkflowUIBasedAppConfig -from core.app.entities.app_invoke_entities import AdvancedChatAppGenerateEntity, InvokeFrom, WorkflowAppGenerateEntity -from core.app.layers.pause_state_persist_layer import ( - WorkflowResumptionContext, - _AdvancedChatAppGenerateEntityWrapper, - _WorkflowGenerateEntityWrapper, -) -from graphon.enums import WorkflowExecutionStatus -from libs.broadcast_channel.redis.pubsub_channel import Topic -from models.enums import CreatorUserRole, WorkflowRunTriggeredFrom -from models.model import AppMode -from models.workflow import WorkflowPause, WorkflowRun, WorkflowType -from repositories.sqlalchemy_api_workflow_run_repository import DifyAPISQLAlchemyWorkflowRunRepository -from services.app_task_service import AppTaskControlService, AppTaskService -from tests.unit_tests.services.test_app_task_service import _StopRedis - - -@pytest.fixture -def pause_states(monkeypatch: pytest.MonkeyPatch) -> dict[str, bytes]: - states: dict[str, bytes] = {} - - def load(key: str) -> bytes: - try: - return states[key] - except KeyError as error: - raise FileNotFoundError(key) from error - - monkeypatch.setattr("repositories.sqlalchemy_api_workflow_run_repository.storage.load", load) - return states - - -@pytest.fixture -def recovery_redis() -> _StopRedis: - return _StopRedis(read_error=AssertionError("Recovery must use persisted ownership")) - - -@pytest.fixture -def recovery_service( - sqlite_engine: Engine, monkeypatch: pytest.MonkeyPatch, recovery_redis: _StopRedis -) -> AppTaskControlService: - monkeypatch.setattr(type(task_module.db), "engine", property(lambda _: sqlite_engine)) - monkeypatch.setattr("repositories.sqlalchemy_api_workflow_run_repository.storage.delete", lambda _: None) - monkeypatch.setattr(AppTaskService, "_persist_stopped_history", lambda _: None) - monkeypatch.setattr( - "core.workflow.nodes.agent_v2.session_store.WorkflowAgentWorkspaceStore.retire_workflow_run", - lambda *_args, **_kwargs: [], - ) - monkeypatch.setattr("tasks.collect_agent_resources_task.enqueue_agent_resource_collection", lambda **_kwargs: None) - topic = Topic(Redis(), "recovery-test") - monkeypatch.setattr(topic, "publish", lambda _payload: None) - monkeypatch.setattr(task_module.MessageBasedAppGenerator, "get_response_topic", lambda *_: topic) - return AppTaskControlService(redis_client=recovery_redis) - - -def _seed_pause( - session: Session, - states: dict[str, bytes], - app_mode: AppMode = AppMode.WORKFLOW, - *, - pause_id: str | None = None, -) -> tuple[WorkflowRun, WorkflowPause, WorkflowResumptionContext]: - run = WorkflowRun( - id=str(uuid4()), - tenant_id="tenant-1", - app_id="app-1", - workflow_id="workflow-1", - type=WorkflowType.WORKFLOW if app_mode == AppMode.WORKFLOW else WorkflowType.CHAT, - triggered_from=WorkflowRunTriggeredFrom.APP_RUN, - version="1", - status=WorkflowExecutionStatus.PAUSED, - created_by_role=CreatorUserRole.END_USER, - created_by="user-1", - created_at=datetime(2026, 10, 10), - elapsed_time=0, - total_tokens=0, - total_steps=1, - exceptions_count=0, - ) - pause = WorkflowPause(workflow_id=run.workflow_id, workflow_run_id=run.id, state_object_key=f"pause/{run.id}") - if pause_id is not None: - pause.id = pause_id - config = WorkflowUIBasedAppConfig( - tenant_id=run.tenant_id, app_id=run.app_id, app_mode=app_mode, workflow_id=run.workflow_id - ) - if app_mode == AppMode.WORKFLOW: - wrapper = _WorkflowGenerateEntityWrapper( - entity=WorkflowAppGenerateEntity( - task_id=str(uuid4()), - app_config=config, - inputs={}, - files=[], - user_id=run.created_by, - stream=True, - invoke_from=InvokeFrom.WEB_APP, - workflow_execution_id=run.id, - ) - ) - else: - wrapper = _AdvancedChatAppGenerateEntityWrapper( - entity=AdvancedChatAppGenerateEntity( - task_id=str(uuid4()), - app_config=config, - inputs={}, - files=[], - user_id=run.created_by, - stream=True, - invoke_from=InvokeFrom.WEB_APP, - workflow_run_id=run.id, - query="Approve this run", - conversation_id="conversation-1", - ) - ) - context = WorkflowResumptionContext(generate_entity=wrapper, serialized_graph_runtime_state="{}") - states[pause.state_object_key] = context.dumps().encode() - session.add_all([run, pause]) - session.commit() - return run, pause, context - - -def _stop(service: AppTaskControlService, task_id: str, app_mode: AppMode = AppMode.WORKFLOW) -> bool: - return service.stop_workflow_task( - tenant_id="tenant-1", - app_id="app-1", - task_id=task_id, - app_mode=app_mode, - owner=(CreatorUserRole.END_USER, "user-1"), - ) - - -@pytest.mark.parametrize("app_mode", [AppMode.WORKFLOW, AppMode.ADVANCED_CHAT]) -def test_stop_recovers_and_cancels_a_legacy_pause_without_prior_replay( - sqlite_session: Session, - pause_states: dict[str, bytes], - recovery_service: AppTaskControlService, - app_mode: AppMode, - recovery_redis: _StopRedis, -) -> None: - run, pause, context = _seed_pause(sqlite_session, pause_states, app_mode) - task_id = context.get_generate_entity().task_id - assert task_id != run.id - assert run.task_id is None - context.get_generate_entity().task_id = task_id.upper() - pause_states[pause.state_object_key] = context.dumps().encode() - - assert _stop(recovery_service, task_id.upper(), app_mode) - - sqlite_session.refresh(run) - assert run.task_id == task_id - assert run.status == WorkflowExecutionStatus.STOPPED - assert run.stop_requested_at is not None - assert sqlite_session.scalar(select(WorkflowPause.id).where(WorkflowPause.id == pause.id)) is None - sqlite_session.commit() - pause_states.clear() - assert _stop(recovery_service, task_id, app_mode) - assert recovery_redis.reads == [] - assert recovery_redis.operations == [] - - -@pytest.mark.parametrize("field", ["tenant_id", "app_id", "created_by", "created_by_role"]) -def test_recovery_excludes_unauthorized_pauses_before_storage_reads( - sqlite_session: Session, - pause_states: dict[str, bytes], - recovery_service: AppTaskControlService, - monkeypatch: pytest.MonkeyPatch, - field: str, -) -> None: - run, _, context = _seed_pause(sqlite_session, pause_states) - setattr(run, field, CreatorUserRole.ACCOUNT if field == "created_by_role" else "other-owner") - sqlite_session.commit() - - def reject_load(_key: str) -> bytes: - pytest.fail("Unauthorized pauses must not reach snapshot storage") - - monkeypatch.setattr("repositories.sqlalchemy_api_workflow_run_repository.storage.load", reject_load) - assert not _stop(recovery_service, context.get_generate_entity().task_id) - sqlite_session.refresh(run) - assert run.task_id is None - assert run.stop_requested_at is None - assert run.status == WorkflowExecutionStatus.PAUSED - - -@pytest.mark.parametrize("field", ["task_id", "run_id", "tenant_id", "app_id", "user_id", "app_config"]) -def test_recovery_rejects_snapshot_identity_mismatches_and_continues( - sqlite_session: Session, - pause_states: dict[str, bytes], - recovery_service: AppTaskControlService, - field: str, -) -> None: - invalid, pause, invalid_context = _seed_pause(sqlite_session, pause_states, pause_id=str(UUID(int=1))) - valid, _, valid_context = _seed_pause(sqlite_session, pause_states, pause_id=str(UUID(int=2))) - state = json.loads(pause_states[pause.state_object_key]) - entity = state["generate_entity"]["entity"] - if field in {"tenant_id", "app_id"}: - entity["app_config"][field] = "other-owner" - elif field == "app_config": - del entity["app_config"] - else: - entity["workflow_execution_id" if field == "run_id" else field] = "invalid-task-or-owner" - pause_states[pause.state_object_key] = json.dumps(state).encode() - - assert not _stop(recovery_service, invalid_context.get_generate_entity().task_id) - sqlite_session.refresh(invalid) - sqlite_session.refresh(valid) - assert invalid.task_id is None - assert invalid.status == WorkflowExecutionStatus.PAUSED - assert invalid.stop_requested_at is None - assert valid.task_id is None - sqlite_session.commit() - - assert _stop(recovery_service, valid_context.get_generate_entity().task_id) - sqlite_session.refresh(invalid) - sqlite_session.refresh(valid) - assert invalid.task_id is None - assert valid.task_id == valid_context.get_generate_entity().task_id - assert valid.status == WorkflowExecutionStatus.STOPPED - - -@pytest.mark.parametrize("missing", ["pause", "state", "corrupt_state"]) -def test_recovery_skips_unrecoverable_pause_and_finds_later_match( - sqlite_session: Session, - pause_states: dict[str, bytes], - recovery_service: AppTaskControlService, - missing: str, -) -> None: - invalid, pause, invalid_context = _seed_pause(sqlite_session, pause_states, pause_id=str(UUID(int=1))) - valid, _, valid_context = _seed_pause(sqlite_session, pause_states, pause_id=str(UUID(int=2))) - if missing == "pause": - sqlite_session.delete(pause) - sqlite_session.commit() - elif missing == "state": - del pause_states[pause.state_object_key] - else: - pause_states[pause.state_object_key] = b"invalid JSON" - - assert not _stop(recovery_service, invalid_context.get_generate_entity().task_id) - assert _stop(recovery_service, valid_context.get_generate_entity().task_id) - - sqlite_session.refresh(invalid) - sqlite_session.refresh(valid) - assert invalid.task_id is None - assert invalid.status == WorkflowExecutionStatus.PAUSED - assert invalid.stop_requested_at is None - assert valid.status == WorkflowExecutionStatus.STOPPED - - -@pytest.mark.parametrize("change", ["replaced_pause", "resumed", "already_bound"]) -def test_indexing_does_not_overwrite_a_concurrent_transition( - sqlite_session: Session, - sqlite_session_factory: sessionmaker[Session], - pause_states: dict[str, bytes], - change: str, -) -> None: - run, pause, context = _seed_pause(sqlite_session, pause_states) - repository = DifyAPISQLAlchemyWorkflowRunRepository(sqlite_session_factory) - saved_pause = repository.get_workflow_pause(run.id) - assert saved_pause is not None - task_id = str(uuid4()) if change == "already_bound" else None - if change == "replaced_pause": - sqlite_session.delete(pause) - sqlite_session.flush() - sqlite_session.add( - WorkflowPause(workflow_id=run.workflow_id, workflow_run_id=run.id, state_object_key="replacement") - ) - elif change == "resumed": - run.status = WorkflowExecutionStatus.RUNNING - pause.resumed_at = datetime(2026, 10, 10, 1) - else: - run.task_id = task_id - sqlite_session.commit() - - assert not AppTaskService.index_workflow_pause_task( - workflow_run=run, - pause_entity=saved_pause, - resumption_context=context, - workflow_run_repo=repository, - ) - - sqlite_session.refresh(run) - assert run.task_id == task_id - - -@pytest.mark.parametrize("resumed", [False, True]) -def test_stop_retries_durable_admission_after_another_request_binds_the_task( - sqlite_session: Session, - sqlite_session_factory: sessionmaker[Session], - pause_states: dict[str, bytes], - recovery_service: AppTaskControlService, - monkeypatch: pytest.MonkeyPatch, - resumed: bool, - recovery_redis: _StopRedis, -) -> None: - run, pause, context = _seed_pause(sqlite_session, pause_states) - task_id = context.get_generate_entity().task_id - - def load(key: str) -> bytes: - with sqlite_session_factory() as session, session.begin(): - concurrent_run = session.get(WorkflowRun, run.id) - assert concurrent_run is not None - concurrent_run.task_id = task_id - if resumed: - concurrent_run.status = WorkflowExecutionStatus.RUNNING - concurrent_pause = session.get(WorkflowPause, pause.id) - assert concurrent_pause is not None - concurrent_pause.resumed_at = datetime(2026, 10, 10, 1) - return pause_states[key] - - monkeypatch.setattr("repositories.sqlalchemy_api_workflow_run_repository.storage.load", load) - assert _stop(recovery_service, task_id) - sqlite_session.refresh(run) - assert run.task_id == task_id - assert run.stop_requested_at is not None - assert run.status == (WorkflowExecutionStatus.RUNNING if resumed else WorkflowExecutionStatus.STOPPED) - assert recovery_redis.reads == [] - assert recovery_redis.operations == (["legacy_flag", "graph_command"] if resumed else []) - - -@pytest.mark.parametrize("status", [WorkflowExecutionStatus.PAUSED, WorkflowExecutionStatus.RUNNING]) -def test_prebound_tasks_stop_without_reading_legacy_snapshots( - sqlite_session: Session, - pause_states: dict[str, bytes], - recovery_service: AppTaskControlService, - monkeypatch: pytest.MonkeyPatch, - status: WorkflowExecutionStatus, -) -> None: - run, _, context = _seed_pause(sqlite_session, pause_states) - run.task_id = context.get_generate_entity().task_id - run.status = status - sqlite_session.commit() - - def reject_load(_key: str) -> bytes: - pytest.fail("Already indexed tasks must not read snapshots") - - monkeypatch.setattr("repositories.sqlalchemy_api_workflow_run_repository.storage.load", reject_load) - assert _stop(recovery_service, context.get_generate_entity().task_id) - sqlite_session.refresh(run) - assert run.stop_requested_at is not None - - -@pytest.mark.parametrize("excluded", ["terminal", "pipeline", "resumed", "running"]) -def test_recovery_excludes_nonrecoverable_runs_before_storage_reads( - sqlite_session: Session, - pause_states: dict[str, bytes], - recovery_service: AppTaskControlService, - monkeypatch: pytest.MonkeyPatch, - excluded: str, -) -> None: - run, pause, context = _seed_pause(sqlite_session, pause_states) - if excluded == "terminal": - run.status = WorkflowExecutionStatus.SUCCEEDED - elif excluded == "pipeline": - run.type = WorkflowType.RAG_PIPELINE - elif excluded == "resumed": - pause.resumed_at = datetime(2026, 10, 10, 1) - else: - run.status = WorkflowExecutionStatus.RUNNING - sqlite_session.commit() - - def reject_load(_key: str) -> bytes: - pytest.fail("Ineligible runs must not read snapshots") - - monkeypatch.setattr("repositories.sqlalchemy_api_workflow_run_repository.storage.load", reject_load) - assert not _stop(recovery_service, context.get_generate_entity().task_id) - sqlite_session.refresh(run) - assert run.task_id is None - assert run.stop_requested_at is None - - -def test_recovery_pages_pauses_without_backfilling_unrelated_runs( - sqlite_session: Session, - pause_states: dict[str, bytes], - recovery_service: AppTaskControlService, -) -> None: - for index in range(1, 101): - _seed_pause(sqlite_session, pause_states, pause_id=str(UUID(int=index))) - run, _, context = _seed_pause(sqlite_session, pause_states, pause_id=str(UUID(int=101))) - task_id = context.get_generate_entity().task_id - - assert _stop(recovery_service, task_id) - - indexed = sqlite_session.execute( - select(WorkflowRun.id, WorkflowRun.task_id).where(WorkflowRun.task_id.is_not(None)) - ) - assert indexed.tuples().all() == [(run.id, task_id)] - - -def test_snapshot_storage_reads_do_not_hold_a_database_transaction( - sqlite_session: Session, - sqlite_engine: Engine, - pause_states: dict[str, bytes], - recovery_service: AppTaskControlService, - monkeypatch: pytest.MonkeyPatch, -) -> None: - _, _, context = _seed_pause(sqlite_session, pause_states) - active_transactions: set[Connection] = set() - - def load(key: str) -> bytes: - assert not active_transactions - return pause_states[key] - - monkeypatch.setattr("repositories.sqlalchemy_api_workflow_run_repository.storage.load", load) - listeners: tuple[tuple[str, Callable[[Connection], None]], ...] = ( - ("begin", active_transactions.add), - ("commit", active_transactions.discard), - ("rollback", active_transactions.discard), - ) - for name, listener in listeners: - event.listen(sqlite_engine, name, listener) - try: - assert _stop(recovery_service, context.get_generate_entity().task_id) - finally: - for name, listener in listeners: - event.remove(sqlite_engine, name, listener) diff --git a/api/tests/unit_tests/services/test_message_file_preview_service.py b/api/tests/unit_tests/services/test_message_file_preview_service.py index 75881ab89d6..3a37f7b28e8 100644 --- a/api/tests/unit_tests/services/test_message_file_preview_service.py +++ b/api/tests/unit_tests/services/test_message_file_preview_service.py @@ -17,7 +17,7 @@ from services.message_file_preview_service import ( MessageFilePreviewRecord, MessageFilePreviewService, ) -from tests.unit_tests.model_factories import make_message, make_upload_file +from tests.helpers.unit.model_factories import make_message, make_upload_file def _record() -> MessageFilePreviewRecord: diff --git a/api/tests/unit_tests/services/test_messages_clean_service.py b/api/tests/unit_tests/services/test_messages_clean_service.py deleted file mode 100644 index fc659e7b5cf..00000000000 --- a/api/tests/unit_tests/services/test_messages_clean_service.py +++ /dev/null @@ -1,685 +0,0 @@ -import datetime -from typing import Any -from unittest.mock import MagicMock, patch - -import pytest - -from enums import CloudPlan, DeploymentEdition -from services.retention.conversation.messages_clean_policy import ( - BillingDisabledPolicy, - BillingSandboxPolicy, - SimpleMessage, - create_message_clean_policy, -) -from services.retention.conversation.messages_clean_service import MessagesCleanService - - -def make_simple_message(msg_id: str, app_id: str) -> SimpleMessage: - """Helper to create a SimpleMessage with a fixed created_at timestamp.""" - return SimpleMessage(id=msg_id, app_id=app_id, created_at=datetime.datetime(2024, 1, 1)) - - -def make_plan_provider(tenant_plans: dict[str, Any]) -> MagicMock: - """Helper to create a mock plan_provider that returns the given tenant_plans.""" - provider = MagicMock() - provider.return_value = tenant_plans - return provider - - -class TestBillingSandboxPolicyFilterMessageIds: - """Unit tests for BillingSandboxPolicy.filter_message_ids method.""" - - # Fixed timestamp for deterministic tests - CURRENT_TIMESTAMP = 1000000 - GRACEFUL_PERIOD_DAYS = 8 - GRACEFUL_PERIOD_SECONDS = GRACEFUL_PERIOD_DAYS * 24 * 60 * 60 - - def test_missing_tenant_mapping_excluded(self): - """Test that messages with missing app-to-tenant mapping are excluded.""" - # Arrange - messages = [ - make_simple_message("msg1", "app1"), - make_simple_message("msg2", "app2"), - ] - app_to_tenant = {} # No mapping - tenant_plans = {"tenant1": {"plan": CloudPlan.SANDBOX, "expiration_date": -1}} - plan_provider = make_plan_provider(tenant_plans) - - policy = BillingSandboxPolicy( - plan_provider=plan_provider, - graceful_period_days=self.GRACEFUL_PERIOD_DAYS, - current_timestamp=self.CURRENT_TIMESTAMP, - ) - - # Act - result = policy.filter_message_ids(messages, app_to_tenant) - - # Assert - assert list(result) == [] - - def test_missing_tenant_plan_excluded(self): - """Test that messages with missing tenant plan are excluded (safe default).""" - # Arrange - messages = [ - make_simple_message("msg1", "app1"), - make_simple_message("msg2", "app2"), - ] - app_to_tenant = {"app1": "tenant1", "app2": "tenant2"} - tenant_plans = {} # No plans - plan_provider = make_plan_provider(tenant_plans) - - policy = BillingSandboxPolicy( - plan_provider=plan_provider, - graceful_period_days=self.GRACEFUL_PERIOD_DAYS, - current_timestamp=self.CURRENT_TIMESTAMP, - ) - - # Act - result = policy.filter_message_ids(messages, app_to_tenant) - - # Assert - assert list(result) == [] - - def test_non_sandbox_plan_excluded(self): - """Test that messages from non-sandbox plans (PROFESSIONAL/TEAM) are excluded.""" - # Arrange - messages = [ - make_simple_message("msg1", "app1"), - make_simple_message("msg2", "app2"), - make_simple_message("msg3", "app3"), - ] - app_to_tenant = {"app1": "tenant1", "app2": "tenant2", "app3": "tenant3"} - tenant_plans = { - "tenant1": {"plan": CloudPlan.PROFESSIONAL, "expiration_date": -1}, - "tenant2": {"plan": CloudPlan.TEAM, "expiration_date": -1}, - "tenant3": {"plan": CloudPlan.SANDBOX, "expiration_date": -1}, # Only this one - } - plan_provider = make_plan_provider(tenant_plans) - - policy = BillingSandboxPolicy( - plan_provider=plan_provider, - graceful_period_days=self.GRACEFUL_PERIOD_DAYS, - current_timestamp=self.CURRENT_TIMESTAMP, - ) - - # Act - result = policy.filter_message_ids(messages, app_to_tenant) - - # Assert - only msg3 (sandbox tenant) should be included - assert set(result) == {"msg3"} - - def test_whitelist_skip(self): - """Test that whitelisted tenants are excluded even if sandbox + expired.""" - # Arrange - messages = [ - make_simple_message("msg1", "app1"), # Whitelisted - excluded - make_simple_message("msg2", "app2"), # Not whitelisted - included - make_simple_message("msg3", "app3"), # Whitelisted - excluded - ] - app_to_tenant = {"app1": "tenant1", "app2": "tenant2", "app3": "tenant3"} - tenant_plans = { - "tenant1": {"plan": CloudPlan.SANDBOX, "expiration_date": -1}, - "tenant2": {"plan": CloudPlan.SANDBOX, "expiration_date": -1}, - "tenant3": {"plan": CloudPlan.SANDBOX, "expiration_date": -1}, - } - plan_provider = make_plan_provider(tenant_plans) - tenant_whitelist = ["tenant1", "tenant3"] - - policy = BillingSandboxPolicy( - plan_provider=plan_provider, - graceful_period_days=self.GRACEFUL_PERIOD_DAYS, - tenant_whitelist=tenant_whitelist, - current_timestamp=self.CURRENT_TIMESTAMP, - ) - - # Act - result = policy.filter_message_ids(messages, app_to_tenant) - - # Assert - only msg2 should be included - assert set(result) == {"msg2"} - - def test_no_previous_subscription_included(self): - """Test that messages with expiration_date=-1 (no previous subscription) are included.""" - # Arrange - messages = [ - make_simple_message("msg1", "app1"), - make_simple_message("msg2", "app2"), - ] - app_to_tenant = {"app1": "tenant1", "app2": "tenant2"} - tenant_plans = { - "tenant1": {"plan": CloudPlan.SANDBOX, "expiration_date": -1}, - "tenant2": {"plan": CloudPlan.SANDBOX, "expiration_date": -1}, - } - plan_provider = make_plan_provider(tenant_plans) - - policy = BillingSandboxPolicy( - plan_provider=plan_provider, - graceful_period_days=self.GRACEFUL_PERIOD_DAYS, - current_timestamp=self.CURRENT_TIMESTAMP, - ) - - # Act - result = policy.filter_message_ids(messages, app_to_tenant) - - # Assert - all messages should be included - assert set(result) == {"msg1", "msg2"} - - def test_within_grace_period_excluded(self): - """Test that messages within grace period are excluded.""" - # Arrange - now = self.CURRENT_TIMESTAMP - expired_1_day_ago = now - (1 * 24 * 60 * 60) - expired_5_days_ago = now - (5 * 24 * 60 * 60) - expired_7_days_ago = now - (7 * 24 * 60 * 60) - - messages = [ - make_simple_message("msg1", "app1"), - make_simple_message("msg2", "app2"), - make_simple_message("msg3", "app3"), - ] - app_to_tenant = {"app1": "tenant1", "app2": "tenant2", "app3": "tenant3"} - tenant_plans = { - "tenant1": {"plan": CloudPlan.SANDBOX, "expiration_date": expired_1_day_ago}, - "tenant2": {"plan": CloudPlan.SANDBOX, "expiration_date": expired_5_days_ago}, - "tenant3": {"plan": CloudPlan.SANDBOX, "expiration_date": expired_7_days_ago}, - } - plan_provider = make_plan_provider(tenant_plans) - - policy = BillingSandboxPolicy( - plan_provider=plan_provider, - graceful_period_days=self.GRACEFUL_PERIOD_DAYS, # 8 days - current_timestamp=now, - ) - - # Act - result = policy.filter_message_ids(messages, app_to_tenant) - - # Assert - all within 8-day grace period, none should be included - assert list(result) == [] - - def test_exactly_at_boundary_excluded(self): - """Test that messages exactly at grace period boundary are excluded (code uses >).""" - # Arrange - now = self.CURRENT_TIMESTAMP - expired_exactly_8_days_ago = now - self.GRACEFUL_PERIOD_SECONDS # Exactly at boundary - - messages = [make_simple_message("msg1", "app1")] - app_to_tenant = {"app1": "tenant1"} - tenant_plans = { - "tenant1": {"plan": CloudPlan.SANDBOX, "expiration_date": expired_exactly_8_days_ago}, - } - plan_provider = make_plan_provider(tenant_plans) - - policy = BillingSandboxPolicy( - plan_provider=plan_provider, - graceful_period_days=self.GRACEFUL_PERIOD_DAYS, - current_timestamp=now, - ) - - # Act - result = policy.filter_message_ids(messages, app_to_tenant) - - # Assert - exactly at boundary (==) should be excluded (code uses >) - assert list(result) == [] - - def test_beyond_grace_period_included(self): - """Test that messages beyond grace period are included.""" - # Arrange - now = self.CURRENT_TIMESTAMP - expired_9_days_ago = now - (9 * 24 * 60 * 60) # Just beyond 8-day grace - expired_30_days_ago = now - (30 * 24 * 60 * 60) # Well beyond - - messages = [ - make_simple_message("msg1", "app1"), - make_simple_message("msg2", "app2"), - ] - app_to_tenant = {"app1": "tenant1", "app2": "tenant2"} - tenant_plans = { - "tenant1": {"plan": CloudPlan.SANDBOX, "expiration_date": expired_9_days_ago}, - "tenant2": {"plan": CloudPlan.SANDBOX, "expiration_date": expired_30_days_ago}, - } - plan_provider = make_plan_provider(tenant_plans) - - policy = BillingSandboxPolicy( - plan_provider=plan_provider, - graceful_period_days=self.GRACEFUL_PERIOD_DAYS, - current_timestamp=now, - ) - - # Act - result = policy.filter_message_ids(messages, app_to_tenant) - - # Assert - both beyond grace period, should be included - assert set(result) == {"msg1", "msg2"} - - def test_empty_messages_returns_empty(self): - """Test that empty messages returns empty list.""" - # Arrange - messages: list[SimpleMessage] = [] - app_to_tenant = {"app1": "tenant1"} - plan_provider = make_plan_provider({"tenant1": {"plan": CloudPlan.SANDBOX, "expiration_date": -1}}) - - policy = BillingSandboxPolicy( - plan_provider=plan_provider, - graceful_period_days=self.GRACEFUL_PERIOD_DAYS, - current_timestamp=self.CURRENT_TIMESTAMP, - ) - - # Act - result = policy.filter_message_ids(messages, app_to_tenant) - - # Assert - assert list(result) == [] - - def test_plan_provider_called_with_correct_tenant_ids(self): - """Test that plan_provider is called with correct tenant_ids.""" - # Arrange - messages = [ - make_simple_message("msg1", "app1"), - make_simple_message("msg2", "app2"), - make_simple_message("msg3", "app3"), - ] - app_to_tenant = {"app1": "tenant1", "app2": "tenant2", "app3": "tenant1"} # tenant1 appears twice - plan_provider = make_plan_provider({}) - - policy = BillingSandboxPolicy( - plan_provider=plan_provider, - graceful_period_days=self.GRACEFUL_PERIOD_DAYS, - current_timestamp=self.CURRENT_TIMESTAMP, - ) - - # Act - policy.filter_message_ids(messages, app_to_tenant) - - # Assert - plan_provider should be called once with unique tenant_ids - plan_provider.assert_called_once() - called_tenant_ids = set(plan_provider.call_args[0][0]) - assert called_tenant_ids == {"tenant1", "tenant2"} - - def test_complex_mixed_scenario(self): - """Test complex scenario with mixed plans, expirations, whitelist, and missing mappings.""" - # Arrange - now = self.CURRENT_TIMESTAMP - sandbox_expired_old = now - (15 * 24 * 60 * 60) # Beyond grace - sandbox_expired_recent = now - (3 * 24 * 60 * 60) # Within grace - future_expiration = now + (30 * 24 * 60 * 60) - - messages = [ - make_simple_message("msg1", "app1"), # Sandbox, no subscription - included - make_simple_message("msg2", "app2"), # Sandbox, expired old - included - make_simple_message("msg3", "app3"), # Sandbox, within grace - excluded - make_simple_message("msg4", "app4"), # Team plan, active - excluded - make_simple_message("msg5", "app5"), # No tenant mapping - excluded - make_simple_message("msg6", "app6"), # No plan info - excluded - make_simple_message("msg7", "app7"), # Sandbox, expired old, whitelisted - excluded - ] - app_to_tenant = { - "app1": "tenant1", - "app2": "tenant2", - "app3": "tenant3", - "app4": "tenant4", - "app6": "tenant6", # Has mapping but no plan - "app7": "tenant7", - # app5 has no mapping - } - tenant_plans = { - "tenant1": {"plan": CloudPlan.SANDBOX, "expiration_date": -1}, - "tenant2": {"plan": CloudPlan.SANDBOX, "expiration_date": sandbox_expired_old}, - "tenant3": {"plan": CloudPlan.SANDBOX, "expiration_date": sandbox_expired_recent}, - "tenant4": {"plan": CloudPlan.TEAM, "expiration_date": future_expiration}, - "tenant7": {"plan": CloudPlan.SANDBOX, "expiration_date": sandbox_expired_old}, - # tenant6 has no plan - } - plan_provider = make_plan_provider(tenant_plans) - tenant_whitelist = ["tenant7"] - - policy = BillingSandboxPolicy( - plan_provider=plan_provider, - graceful_period_days=self.GRACEFUL_PERIOD_DAYS, - tenant_whitelist=tenant_whitelist, - current_timestamp=now, - ) - - # Act - result = policy.filter_message_ids(messages, app_to_tenant) - - # Assert - only msg1 and msg2 should be included - assert set(result) == {"msg1", "msg2"} - - -class TestBillingDisabledPolicyFilterMessageIds: - """Unit tests for BillingDisabledPolicy.filter_message_ids method.""" - - def test_returns_all_message_ids(self): - """Test that all message IDs are returned (order-preserving).""" - # Arrange - messages = [ - make_simple_message("msg1", "app1"), - make_simple_message("msg2", "app2"), - make_simple_message("msg3", "app3"), - ] - app_to_tenant = {"app1": "tenant1", "app2": "tenant2"} - - policy = BillingDisabledPolicy() - - # Act - result = policy.filter_message_ids(messages, app_to_tenant) - - # Assert - all message IDs returned in order - assert list(result) == ["msg1", "msg2", "msg3"] - - def test_ignores_app_to_tenant(self): - """Test that app_to_tenant mapping is ignored.""" - # Arrange - messages = [ - make_simple_message("msg1", "app1"), - make_simple_message("msg2", "app2"), - ] - app_to_tenant: dict[str, str] = {} # Empty - should be ignored - - policy = BillingDisabledPolicy() - - # Act - result = policy.filter_message_ids(messages, app_to_tenant) - - # Assert - all message IDs still returned - assert list(result) == ["msg1", "msg2"] - - def test_empty_messages_returns_empty(self): - """Test that empty messages returns empty list.""" - # Arrange - messages: list[SimpleMessage] = [] - app_to_tenant = {"app1": "tenant1"} - - policy = BillingDisabledPolicy() - - # Act - result = policy.filter_message_ids(messages, app_to_tenant) - - # Assert - assert list(result) == [] - - -class TestCreateMessageCleanPolicy: - """Unit tests for create_message_clean_policy factory function.""" - - def test_non_cloud_edition_returns_billing_disabled_policy(self, config_overrides): - """Test that the Community edition returns BillingDisabledPolicy.""" - # Arrange - config_overrides(DEPLOYMENT_EDITION=DeploymentEdition.COMMUNITY) - - # Act - policy = create_message_clean_policy(graceful_period_days=21) - - # Assert - assert isinstance(policy, BillingDisabledPolicy) - - @patch("services.retention.conversation.messages_clean_policy.BillingService", autospec=True) - def test_cloud_edition_policy_has_correct_internals(self, mock_billing_service, config_overrides): - """Test that BillingSandboxPolicy is created with correct internal values.""" - # Arrange - config_overrides(DEPLOYMENT_EDITION=DeploymentEdition.CLOUD) - whitelist = ["tenant1", "tenant2"] - mock_billing_service.get_expired_subscription_cleanup_whitelist.return_value = whitelist - mock_plan_provider = MagicMock() - mock_billing_service.get_plan_bulk_with_cache = mock_plan_provider - - # Act - policy = create_message_clean_policy(graceful_period_days=14, current_timestamp=1234567) - - # Assert - mock_billing_service.get_expired_subscription_cleanup_whitelist.assert_called_once() - assert isinstance(policy, BillingSandboxPolicy) - assert policy._graceful_period_days == 14 - assert list(policy._tenant_whitelist) == whitelist - assert policy._plan_provider == mock_plan_provider - assert policy._current_timestamp == 1234567 - - -class TestMessagesCleanServiceFromTimeRange: - """Unit tests for MessagesCleanService.from_time_range factory method.""" - - def test_start_from_end_before_raises_value_error(self): - """Test that start_from == end_before raises ValueError.""" - policy = BillingDisabledPolicy() - - # Arrange - same_time = datetime.datetime(2024, 1, 1, 12, 0, 0) - - # Act & Assert - with pytest.raises(ValueError, match="start_from .* must be less than end_before"): - MessagesCleanService.from_time_range( - policy=policy, - start_from=same_time, - end_before=same_time, - ) - - # Arrange - start_from = datetime.datetime(2024, 12, 31) - end_before = datetime.datetime(2024, 1, 1) - - # Act & Assert - with pytest.raises(ValueError, match="start_from .* must be less than end_before"): - MessagesCleanService.from_time_range( - policy=policy, - start_from=start_from, - end_before=end_before, - ) - - def test_batch_size_raises_value_error(self): - """Test that batch_size=0 raises ValueError.""" - # Arrange - start_from = datetime.datetime(2024, 1, 1) - end_before = datetime.datetime(2024, 2, 1) - policy = BillingDisabledPolicy() - - # Act & Assert - with pytest.raises(ValueError, match="batch_size .* must be greater than 0"): - MessagesCleanService.from_time_range( - policy=policy, - start_from=start_from, - end_before=end_before, - batch_size=0, - ) - - start_from = datetime.datetime(2024, 1, 1) - end_before = datetime.datetime(2024, 2, 1) - policy = BillingDisabledPolicy() - - # Act & Assert - with pytest.raises(ValueError, match="batch_size .* must be greater than 0"): - MessagesCleanService.from_time_range( - policy=policy, - start_from=start_from, - end_before=end_before, - batch_size=-100, - ) - - def test_valid_params_creates_instance(self): - """Test that valid parameters create a correctly configured instance.""" - # Arrange - start_from = datetime.datetime(2024, 1, 1, 0, 0, 0) - end_before = datetime.datetime(2024, 12, 31, 23, 59, 59) - policy = BillingDisabledPolicy() - batch_size = 500 - dry_run = True - - # Act - service = MessagesCleanService.from_time_range( - policy=policy, - start_from=start_from, - end_before=end_before, - batch_size=batch_size, - dry_run=dry_run, - ) - - # Assert - assert isinstance(service, MessagesCleanService) - assert service._policy is policy - assert service._start_from == start_from - assert service._end_before == end_before - assert service._batch_size == batch_size - assert service._dry_run == dry_run - - def test_default_params(self): - """Test that default parameters are applied correctly.""" - # Arrange - start_from = datetime.datetime(2024, 1, 1) - end_before = datetime.datetime(2024, 2, 1) - policy = BillingDisabledPolicy() - - # Act - service = MessagesCleanService.from_time_range( - policy=policy, - start_from=start_from, - end_before=end_before, - ) - - # Assert - assert service._batch_size == 1000 # default - assert service._dry_run is False # default - - def test_explicit_task_label(self): - start_from = datetime.datetime(2024, 1, 1) - end_before = datetime.datetime(2024, 1, 2) - policy = BillingDisabledPolicy() - - service = MessagesCleanService.from_time_range( - policy=policy, - start_from=start_from, - end_before=end_before, - task_label="60to30", - ) - - assert service._metrics._base_attributes["task_label"] == "60to30" - - -class TestMessagesCleanServiceFromDays: - """Unit tests for MessagesCleanService.from_days factory method.""" - - def test_days_raises_value_error(self): - """Test that days < 0 raises ValueError.""" - # Arrange - policy = BillingDisabledPolicy() - - # Act & Assert - with pytest.raises(ValueError, match="days .* must be greater than or equal to 0"): - MessagesCleanService.from_days(policy=policy, days=-1) - - # Act - with patch("services.retention.conversation.messages_clean_service.naive_utc_now") as mock_now: - fixed_now = datetime.datetime(2024, 6, 15, 14, 0, 0) - mock_now.return_value = fixed_now - service = MessagesCleanService.from_days(policy=policy, days=0) - - # Assert - assert service._end_before == fixed_now - - def test_batch_size_raises_value_error(self): - """Test that batch_size=0 raises ValueError.""" - # Arrange - policy = BillingDisabledPolicy() - - # Act & Assert - with pytest.raises(ValueError, match="batch_size .* must be greater than 0"): - MessagesCleanService.from_days(policy=policy, days=30, batch_size=0) - - # Act & Assert - with pytest.raises(ValueError, match="batch_size .* must be greater than 0"): - MessagesCleanService.from_days(policy=policy, days=30, batch_size=-500) - - def test_valid_params_creates_instance(self): - """Test that valid parameters create a correctly configured instance.""" - # Arrange - policy = BillingDisabledPolicy() - days = 90 - batch_size = 500 - dry_run = True - - # Act - with patch("services.retention.conversation.messages_clean_service.naive_utc_now") as mock_now: - fixed_now = datetime.datetime(2024, 6, 15, 10, 30, 0) - mock_now.return_value = fixed_now - service = MessagesCleanService.from_days( - policy=policy, - days=days, - batch_size=batch_size, - dry_run=dry_run, - ) - - # Assert - expected_end_before = fixed_now - datetime.timedelta(days=days) - assert isinstance(service, MessagesCleanService) - assert service._policy is policy - assert service._start_from is None - assert service._end_before == expected_end_before - assert service._batch_size == batch_size - assert service._dry_run == dry_run - - def test_default_params(self): - """Test that default parameters are applied correctly.""" - # Arrange - policy = BillingDisabledPolicy() - - # Act - with patch("services.retention.conversation.messages_clean_service.naive_utc_now") as mock_now: - fixed_now = datetime.datetime(2024, 6, 15, 10, 30, 0) - mock_now.return_value = fixed_now - service = MessagesCleanService.from_days(policy=policy) - - # Assert - expected_end_before = fixed_now - datetime.timedelta(days=30) # default days=30 - assert service._end_before == expected_end_before - assert service._batch_size == 1000 # default - assert service._dry_run is False # default - assert service._metrics._base_attributes["task_label"] == "custom" - - -class TestMessagesCleanServiceRun: - """Unit tests for MessagesCleanService.run instrumentation behavior.""" - - def test_run_records_completion_metrics_on_success(self): - # Arrange - service = MessagesCleanService( - policy=BillingDisabledPolicy(), - start_from=datetime.datetime(2024, 1, 1), - end_before=datetime.datetime(2024, 1, 2), - batch_size=100, - dry_run=False, - ) - expected_stats = { - "batches": 1, - "total_messages": 10, - "filtered_messages": 5, - "total_deleted": 5, - } - service._clean_messages_by_time_range = MagicMock(return_value=expected_stats) # type: ignore[method-assign] - completion_calls: list[dict[str, object]] = [] - service._metrics.record_completion = lambda **kwargs: completion_calls.append(kwargs) # type: ignore[method-assign] - - # Act - result = service.run() - - # Assert - assert result == expected_stats - assert len(completion_calls) == 1 - assert completion_calls[0]["status"] == "success" - - def test_run_records_completion_metrics_on_failure(self): - # Arrange - service = MessagesCleanService( - policy=BillingDisabledPolicy(), - start_from=datetime.datetime(2024, 1, 1), - end_before=datetime.datetime(2024, 1, 2), - batch_size=100, - dry_run=False, - ) - service._clean_messages_by_time_range = MagicMock(side_effect=RuntimeError("clean failed")) # type: ignore[method-assign] - completion_calls: list[dict[str, object]] = [] - service._metrics.record_completion = lambda **kwargs: completion_calls.append(kwargs) # type: ignore[method-assign] - - # Act & Assert - with pytest.raises(RuntimeError, match="clean failed"): - service.run() - assert len(completion_calls) == 1 - assert completion_calls[0]["status"] == "failed" diff --git a/api/tests/unit_tests/services/test_metadata_bug_complete.py b/api/tests/unit_tests/services/test_metadata_bug_complete.py deleted file mode 100644 index 7c5f978642d..00000000000 --- a/api/tests/unit_tests/services/test_metadata_bug_complete.py +++ /dev/null @@ -1,193 +0,0 @@ -from pathlib import Path -from typing import cast -from unittest.mock import Mock, create_autospec - -import pytest -from sqlalchemy.orm import Session, sessionmaker - -from models import Account, Tenant -from repositories.knowledge.metadata_repository import SQLAlchemyMetadataRepository -from services.knowledge.dataset_access import DatasetAccess -from services.knowledge.entities.knowledge_entities import MetadataArgs -from services.knowledge.metadata.application import MetadataService -from services.knowledge.resource_scope import DatasetRef - - -def _metadata_service(session: Session) -> MetadataService: - return MetadataService( - store=SQLAlchemyMetadataRepository( - session_factory=sessionmaker(bind=session.get_bind(), expire_on_commit=False) - ), - dataset_access=create_autospec(DatasetAccess, instance=True), - ) - - -def _make_account(account_id: str = "user-456", tenant_id: str = "tenant-123") -> Account: - account = Account(name="Test User", email=f"{account_id}@example.com") - account.id = account_id - tenant = Tenant(name="Test Tenant") - tenant.id = tenant_id - account._current_tenant = tenant - return account - - -class TestMetadataBugCompleteValidation: - """Complete test suite to verify the metadata nullable bug and its fix.""" - - def test_1_pydantic_layer_validation(self) -> None: - """Test Layer 1: Pydantic model validation correctly rejects None values.""" - # Pydantic should reject None values for required fields - with pytest.raises((ValueError, TypeError)): - MetadataArgs(type=None, name=None) # pyrefly: ignore[bad-argument-type] - - with pytest.raises((ValueError, TypeError)): - MetadataArgs(type="string", name=None) # pyrefly: ignore[bad-argument-type] - - with pytest.raises((ValueError, TypeError)): - MetadataArgs(type=None, name="test") # pyrefly: ignore[bad-argument-type] - - # Valid values should work - valid_args = MetadataArgs(type="string", name="test_name") - assert valid_args.type == "string" - assert valid_args.name == "test_name" - - @pytest.mark.parametrize("sqlite_session", [()], indirect=True) - def test_2_business_logic_layer_crashes_on_none(self, sqlite_session: Session) -> None: - """Test Layer 2: Business logic crashes when None values slip through.""" - # Create mock that bypasses Pydantic validation - mock_metadata_args = Mock() - mock_metadata_args.name = None - mock_metadata_args.type = "string" - - account = _make_account() - # Should crash with TypeError - with pytest.raises(TypeError, match="object of type 'NoneType' has no len"): - _metadata_service(sqlite_session).create_metadata( - DatasetRef("tenant-123", "dataset-123"), mock_metadata_args, actor_id=account.id - ) - - # Test update method as well - account = _make_account() - none_name = cast(str, None) - with pytest.raises(TypeError, match="object of type 'NoneType' has no len"): - _metadata_service(sqlite_session).update_metadata_name( - DatasetRef(Mock().tenant_id, Mock().id), "metadata-456", none_name, actor_id=account.id - ) - assert not sqlite_session.in_transaction() - - def test_3_database_constraints_verification(self) -> None: - """Test Layer 3: Verify database model has nullable=False constraints.""" - from sqlalchemy import inspect - - from models.dataset import DatasetMetadata - - # Get table info - mapper = inspect(DatasetMetadata) - - # Check that type and name columns are not nullable - type_column = mapper.columns["type"] - name_column = mapper.columns["name"] - - assert type_column.nullable is False, "type column should be nullable=False" - assert name_column.nullable is False, "name column should be nullable=False" - - def test_4_fixed_api_layer_rejects_null(self) -> None: - """Test Layer 4: Fixed API configuration properly rejects null values using Pydantic.""" - with pytest.raises((ValueError, TypeError)): - MetadataArgs.model_validate({"type": None, "name": None}) - - with pytest.raises((ValueError, TypeError)): - MetadataArgs.model_validate({"type": "string", "name": None}) - - with pytest.raises((ValueError, TypeError)): - MetadataArgs.model_validate({"type": None, "name": "test"}) - - def test_5_fixed_api_accepts_valid_values(self) -> None: - """Test that fixed API still accepts valid non-null values.""" - args = MetadataArgs.model_validate({"type": "string", "name": "valid_name"}) - assert args.type == "string" - assert args.name == "valid_name" - - @pytest.mark.parametrize("sqlite_session", [()], indirect=True) - def test_6_simulated_buggy_behavior(self, sqlite_session: Session) -> None: - """Test simulating the original buggy behavior by bypassing Pydantic validation.""" - mock_metadata_args = Mock() - mock_metadata_args.name = None - mock_metadata_args.type = None - - account = _make_account() - with pytest.raises(TypeError, match="object of type 'NoneType' has no len"): - _metadata_service(sqlite_session).create_metadata( - DatasetRef("tenant-123", "dataset-123"), mock_metadata_args, actor_id=account.id - ) - assert not sqlite_session.in_transaction() - - def test_7_end_to_end_validation_layers(self) -> None: - """Test all validation layers work together correctly.""" - # Layer 1: API should reject null at parameter level (with fix) - # Layer 2: Pydantic should reject null at model level - # Layer 3: Business logic expects non-null - # Layer 4: Database enforces non-null - - # Test that valid data flows through all layers - valid_data = {"type": "string", "name": "test_metadata"} - - # Should create valid Pydantic object - metadata_args = MetadataArgs.model_validate(valid_data) - assert metadata_args.type == "string" - assert metadata_args.name == "test_metadata" - - # Should not crash in business logic length check - assert len(metadata_args.name) <= 255 # This should not crash - assert len(metadata_args.type) > 0 # This should not crash - - def test_8_verify_specific_fix_locations(self) -> None: - """Verify that the specific locations mentioned in bug report are fixed.""" - # Read the actual files to verify fixes - import os - - # Console API create - console_create_file = "api/controllers/console/datasets/metadata.py" - if os.path.exists(console_create_file): - content = Path(console_create_file).read_text() - # Should contain nullable=False, not nullable=True - assert "nullable=True" not in content.split("class DatasetMetadataCreateApi")[1].split("class")[0] - - # Service API create - service_create_file = "api/controllers/service_api/dataset/metadata.py" - if os.path.exists(service_create_file): - content = Path(service_create_file).read_text() - # Should contain nullable=False, not nullable=True - create_api_section = content.split("class DatasetMetadataCreateServiceApi")[1].split("class")[0] - assert "nullable=True" not in create_api_section - - -class TestMetadataValidationSummary: - """Summary tests that demonstrate the complete validation architecture.""" - - def test_validation_layer_architecture(self) -> None: - """Document and test the 4-layer validation architecture.""" - # Layer 1: API Parameter Validation (Flask-RESTful reqparse) - # - Role: First line of defense, validates HTTP request parameters - # - Fixed: nullable=False ensures null values are rejected at API boundary - - # Layer 2: Pydantic Model Validation - # - Role: Validates data structure and types before business logic - # - Working: Required fields without Optional[] reject None values - - # Layer 3: Business Logic Validation - # - Role: Domain-specific validation (length checks, uniqueness, etc.) - # - Vulnerable: Direct len() calls crash on None values - - # Layer 4: Database Constraints - # - Role: Final data integrity enforcement - # - Working: nullable=False prevents None values in database - - # The bug was: Layer 1 allowed None, but Layers 2-4 expected non-None - # The fix: Make Layer 1 consistent with Layers 2-4 - - assert True # This test documents the architecture - - -if __name__ == "__main__": - pytest.main([__file__, "-v"]) diff --git a/api/tests/unit_tests/services/test_metadata_nullable_bug.py b/api/tests/unit_tests/services/test_metadata_nullable_bug.py deleted file mode 100644 index 241a32bcb52..00000000000 --- a/api/tests/unit_tests/services/test_metadata_nullable_bug.py +++ /dev/null @@ -1,88 +0,0 @@ -from typing import cast -from unittest.mock import Mock, create_autospec - -import pytest -from sqlalchemy.orm import Session, sessionmaker - -from models import Account, Tenant -from repositories.knowledge.metadata_repository import SQLAlchemyMetadataRepository -from services.knowledge.dataset_access import DatasetAccess -from services.knowledge.entities.knowledge_entities import MetadataArgs -from services.knowledge.metadata.application import MetadataService -from services.knowledge.resource_scope import DatasetRef - - -def _metadata_service(session: Session) -> MetadataService: - return MetadataService( - store=SQLAlchemyMetadataRepository( - session_factory=sessionmaker(bind=session.get_bind(), expire_on_commit=False) - ), - dataset_access=create_autospec(DatasetAccess, instance=True), - ) - - -pytestmark = [ - pytest.mark.usefixtures("sqlite_session"), - pytest.mark.parametrize("sqlite_session", [()], indirect=True), -] - - -def _make_account(account_id: str = "user-456", tenant_id: str = "tenant-123") -> Account: - account = Account(name="Test User", email=f"{account_id}@example.com") - account.id = account_id - tenant = Tenant(name="Test Tenant") - tenant.id = tenant_id - account._current_tenant = tenant - return account - - -class TestMetadataNullableBug: - """Test case to reproduce the metadata nullable validation bug.""" - - def test_metadata_args_with_none_values_should_fail(self) -> None: - """Test that MetadataArgs validation should reject None values.""" - # This test demonstrates the expected behavior - should fail validation - with pytest.raises((ValueError, TypeError)): - # This should fail because Pydantic expects non-None values - MetadataArgs(type=None, name=None) # pyrefly: ignore[bad-argument-type] - - def test_metadata_service_create_with_none_name_crashes(self, sqlite_session: Session) -> None: - """Test that MetadataService.create_metadata crashes when name is None.""" - # Mock the MetadataArgs to bypass Pydantic validation - mock_metadata_args = Mock() - mock_metadata_args.name = None # This will cause len() to crash - mock_metadata_args.type = "string" - - account = _make_account() - # This should crash with TypeError when calling len(None) - with pytest.raises(TypeError, match="object of type 'NoneType' has no len"): - _metadata_service(sqlite_session).create_metadata( - DatasetRef("tenant-123", "dataset-123"), mock_metadata_args, actor_id=account.id - ) - assert not sqlite_session.in_transaction() - - def test_metadata_service_update_with_none_name_crashes(self, sqlite_session: Session) -> None: - """Test that MetadataService.update_metadata_name crashes when name is None.""" - account = _make_account() - none_name = cast(str, None) - # This should crash with TypeError when calling len(None) - with pytest.raises(TypeError, match="object of type 'NoneType' has no len"): - _metadata_service(sqlite_session).update_metadata_name( - DatasetRef(Mock().tenant_id, Mock().id), "metadata-456", none_name, actor_id=account.id - ) - assert not sqlite_session.in_transaction() - - def test_api_layer_now_uses_pydantic_validation(self) -> None: - """Verify that API layer relies on Pydantic validation instead of reqparse.""" - invalid_payload: dict[str, object] = {"type": None, "name": None} - with pytest.raises((ValueError, TypeError)): - MetadataArgs.model_validate(invalid_payload) - - valid_payload = {"type": "string", "name": "valid"} - args = MetadataArgs.model_validate(valid_payload) - assert args.type == "string" - assert args.name == "valid" - - -if __name__ == "__main__": - pytest.main([__file__, "-v"]) diff --git a/api/tests/unit_tests/services/test_metadata_service_session_boundary.py b/api/tests/unit_tests/services/test_metadata_service_session_boundary.py deleted file mode 100644 index cf72155954b..00000000000 --- a/api/tests/unit_tests/services/test_metadata_service_session_boundary.py +++ /dev/null @@ -1,202 +0,0 @@ -"""Metadata operations persist through scoped, independently owned transactions.""" - -from collections.abc import Iterator -from datetime import datetime -from unittest.mock import create_autospec - -import pytest -from sqlalchemy import event, select -from sqlalchemy.orm import Session, sessionmaker - -from models.dataset import DatasetMetadataBinding, Document -from models.enums import DocumentCreatedFrom -from repositories.knowledge.metadata_repository import SQLAlchemyMetadataRepository -from services.errors.metadata import MetadataResourceNotFoundError -from services.knowledge.dataset_access import DatasetAccess -from services.knowledge.entities.knowledge_entities import ( - DocumentMetadataOperation, - MetadataArgs, - MetadataDetail, - MetadataOperationData, -) -from services.knowledge.metadata.application import MetadataService -from services.knowledge.resource_scope import DatasetRef -from tests.unit_tests.model_factories import make_account, make_dataset, make_document - -REF = DatasetRef("tenant-1", "dataset-1") -DOCUMENT_ID = "11111111-1111-1111-1111-111111111111" -FOREIGN_ID = "22222222-2222-2222-2222-222222222222" - - -@pytest.fixture -def metadata_service(sqlite_session_factory: sessionmaker[Session]) -> Iterator[MetadataService]: - with sqlite_session_factory.begin() as session: - account = make_account(account_id="actor", name="User", email="metadata@example.com") - session.add(account) - for ref, document_id in [(REF, DOCUMENT_ID), (DatasetRef("tenant-2", "dataset-2"), FOREIGN_ID)]: - session.add( - make_dataset(dataset_id=ref.dataset_id, tenant_id=ref.tenant_id, name="Dataset", created_by="actor") - ) - document = make_document( - document_id=document_id, - tenant_id=ref.tenant_id, - dataset_id=ref.dataset_id, - batch="batch", - created_from=DocumentCreatedFrom.API, - created_by="actor", - created_at=datetime(2026, 1, 1), - updated_at=datetime(2026, 1, 2), - ) - document.enabled = True - document.archived = False - document.indexing_status = "completed" - session.add(document) - return MetadataService( - store=SQLAlchemyMetadataRepository(session_factory=sqlite_session_factory), - dataset_access=create_autospec(DatasetAccess, instance=True), - ) - - -def test_create_commits_through_owned_session( - metadata_service: MetadataService, sqlite_session_factory: sessionmaker[Session] -) -> None: - result = metadata_service.create_metadata(REF, MetadataArgs(type="string", name="author"), actor_id="actor") - fields = metadata_service.get_dataset_metadatas(REF) - assert fields["doc_metadata"] == [{"id": result.id, "type": "string", "name": "author", "count": 0}] - with sqlite_session_factory() as session: - assert session.scalar(select(DatasetMetadataBinding)) is None - - -@pytest.mark.parametrize("name", ["x" * 256, "document_name", "uploader", "upload_date", "last_update_date", "source"]) -def test_invalid_name_is_rejected(metadata_service: MetadataService, name: str) -> None: - with pytest.raises( - ValueError, - match=( - "(?:Metadata name already exists in Built\\-in fields\\.|Metadata name cannot exceed 255 characters\\.)" - ), - ): - metadata_service.create_metadata(REF, MetadataArgs(type="string", name=name), actor_id="actor") - assert metadata_service.get_dataset_metadatas(REF)["doc_metadata"] == [] - - -def test_duplicate_name_is_rejected(metadata_service: MetadataService) -> None: - metadata_service.create_metadata(REF, MetadataArgs(type="string", name="author"), actor_id="actor") - with pytest.raises(ValueError, match="already exists"): - metadata_service.create_metadata(REF, MetadataArgs(type="string", name="author"), actor_id="actor") - - -def _operation(document_id: str, metadata_id: str, *, partial: bool = False) -> MetadataOperationData: - return MetadataOperationData( - operation_data=[ - DocumentMetadataOperation( - document_id=document_id, - metadata_list=[MetadataDetail(id=metadata_id, name="spoofed", value="Alice")], - partial_update=partial, - ) - ] - ) - - -@pytest.mark.parametrize("partial", [True, False]) -def test_update_uses_canonical_names_and_preserves_partial_values( - metadata_service: MetadataService, sqlite_session_factory: sessionmaker[Session], partial: bool -) -> None: - field = metadata_service.create_metadata(REF, MetadataArgs(type="string", name="author"), actor_id="actor") - with sqlite_session_factory.begin() as session: - document = session.get(Document, DOCUMENT_ID) - assert document is not None - document.doc_metadata = {"existing": "kept"} - metadata_service.update_documents_metadata( - REF, _operation(DOCUMENT_ID, field.id, partial=partial), actor_id="actor" - ) - metadata_service.update_documents_metadata( - REF, _operation(DOCUMENT_ID, field.id, partial=partial), actor_id="actor" - ) - with sqlite_session_factory() as session: - document = session.get(Document, DOCUMENT_ID) - assert document is not None - assert document.doc_metadata == ({"existing": "kept", "author": "Alice"} if partial else {"author": "Alice"}) - assert len(session.scalars(select(DatasetMetadataBinding)).all()) == 1 - - -@pytest.mark.parametrize("foreign", ["dataset", "document", "metadata"]) -def test_owner_chain_is_validated_before_any_batch_write( - metadata_service: MetadataService, sqlite_session_factory: sessionmaker[Session], foreign: str -) -> None: - field = metadata_service.create_metadata(REF, MetadataArgs(type="string", name="author"), actor_id="actor") - other = metadata_service.create_metadata( - DatasetRef("tenant-2", "dataset-2"), MetadataArgs(type="string", name="foreign"), actor_id="actor" - ) - args = _operation(DOCUMENT_ID, field.id) - args.operation_data.extend( - _operation( - FOREIGN_ID if foreign == "document" else DOCUMENT_ID, other.id if foreign == "metadata" else field.id - ).operation_data - ) - ref = DatasetRef("tenant-2", REF.dataset_id) if foreign == "dataset" else REF - with pytest.raises(MetadataResourceNotFoundError): - metadata_service.update_documents_metadata(ref, args, actor_id="actor") - with sqlite_session_factory() as session: - assert all(document.doc_metadata is None for document in session.scalars(select(Document))) - assert not session.scalars(select(DatasetMetadataBinding)).all() - - -def test_rename_delete_and_counts_share_scoped_bindings( - metadata_service: MetadataService, sqlite_session_factory: sessionmaker[Session] -) -> None: - field = metadata_service.create_metadata(REF, MetadataArgs(type="string", name="author"), actor_id="actor") - metadata_service.update_documents_metadata(REF, _operation(DOCUMENT_ID, field.id), actor_id="actor") - renamed = metadata_service.update_metadata_name(REF, field.id, "writer", actor_id="actor") - assert renamed.name == "writer" - assert metadata_service.get_dataset_metadatas(REF)["doc_metadata"] == [ - {"id": field.id, "name": "writer", "type": "string", "count": 1} - ] - with sqlite_session_factory() as session: - document = session.get(Document, DOCUMENT_ID) - assert document is not None - assert document.doc_metadata == {"writer": "Alice"} - metadata_service.delete_metadata(REF, field.id) - with sqlite_session_factory() as session: - document = session.get(Document, DOCUMENT_ID) - assert document is not None - assert document.doc_metadata == {} - assert not session.scalars(select(DatasetMetadataBinding)).all() - - -def test_builtin_fields_use_owned_document_uploader( - metadata_service: MetadataService, sqlite_session_factory: sessionmaker[Session] -) -> None: - metadata_service.enable_built_in_field(REF) - with sqlite_session_factory() as session: - document = session.get(Document, DOCUMENT_ID) - foreign = session.get(Document, FOREIGN_ID) - assert document is not None - assert foreign is not None - assert document.doc_metadata["uploader"] == "User" - assert foreign.doc_metadata is None - metadata_service.disable_built_in_field(REF) - with sqlite_session_factory() as session: - document = session.get(Document, DOCUMENT_ID) - assert document is not None - assert document.doc_metadata == {} - - -def test_commit_failure_rolls_back_document_and_bindings( - metadata_service: MetadataService, sqlite_session_factory: sessionmaker[Session] -) -> None: - field = metadata_service.create_metadata(REF, MetadataArgs(type="string", name="author"), actor_id="actor") - - def fail(_session: Session) -> None: - raise RuntimeError("commit failed") - - event.listen(sqlite_session_factory.class_, "before_commit", fail) - try: - with pytest.raises(RuntimeError, match="commit failed"): - metadata_service.update_documents_metadata(REF, _operation(DOCUMENT_ID, field.id), actor_id="actor") - finally: - event.remove(sqlite_session_factory.class_, "before_commit", fail) - with sqlite_session_factory() as session: - document = session.get(Document, DOCUMENT_ID) - assert document is not None - assert document.doc_metadata is None - assert not session.scalars(select(DatasetMetadataBinding)).all() diff --git a/api/tests/unit_tests/services/test_mlflow_tracing_config_gateway.py b/api/tests/unit_tests/services/test_mlflow_tracing_config_gateway.py deleted file mode 100644 index bd024bc1677..00000000000 --- a/api/tests/unit_tests/services/test_mlflow_tracing_config_gateway.py +++ /dev/null @@ -1,273 +0,0 @@ -"""MLflow config persistence uses real RSA/SQLite without constructing an exporter.""" - -import json -from dataclasses import dataclass -from uuid import uuid4 - -import pytest -from dify_trace_mlflow.oss_configuration import MLflowConfigurationError -from sqlalchemy.orm import Session, sessionmaker - -from core.ops.ops_trace_manager import OpsTraceManager -from machinery.context import RequestContext -from repositories.app.console_repository import ConsoleAppRepository -from repositories.app.tracing_config_repository import SQLAlchemyAppTracingConfigRepository -from services.app_tracing_config_gateway import OpsTraceManagerGateway -from services.app_tracing_config_service import ( - AppTracingConfigInvalidConfigurationError, - AppTracingConfigRecord, - AppTracingConfigService, - AppTracingConfigVerificationFailedError, -) -from tests.unit_tests.model_factories import make_app -from tests.unit_tests.services.tracing.test_databricks_configuration import ExperimentAPI, response -from tests.unit_tests.services.tracing.test_mlflow_configuration import ( - experiment_api as experiment_api, # noqa: PLC0414 - Re-export pytest fixtures only in this module. -) -from tests.unit_tests.services.tracing.test_mlflow_configuration import ( - isolated_sdk_state as isolated_sdk_state, # noqa: PLC0414 - Re-export pytest fixtures only in this module. -) - -pytest_plugins = ("tests.unit_tests.services.tracing.test_nullable_tracing_credentials",) - - -@dataclass -class ConfigurationCase: - api: ExperimentAPI - gateway: OpsTraceManagerGateway - service: AppTracingConfigService - store: SQLAlchemyAppTracingConfigRepository - reader: ConsoleAppRepository - context: RequestContext - app_id: str - - def read(self, provider: str = "mlflow") -> AppTracingConfigRecord | None: - return self.store.get( - workspace_id=self.context.active_workspace_id, app_id=self.app_id, tracing_provider=provider - ) - - -@pytest.fixture -def case( - workspace: str, sqlite_session_factory: sessionmaker[Session], experiment_api: ExperimentAPI -) -> ConfigurationCase: - app_id = str(uuid4()) - with sqlite_session_factory.begin() as session: - session.add( - make_app( - app_id=app_id, - tenant_id=workspace, - tracing=json.dumps( - {"enabled": True, "tracing_provider": "langsmith", "capture_content": True, "capture_agent": False} - ), - ) - ) - gateway = OpsTraceManagerGateway() - store = SQLAlchemyAppTracingConfigRepository(session_factory=sqlite_session_factory) - previous = OpsTraceManager.encrypt_tracing_config( - workspace, "langsmith", {"api_key": "synthetic-other-key", "project": "other", "endpoint": "https://other.test"} - ) - assert store.create(workspace_id=workspace, app_id=app_id, tracing_provider="langsmith", tracing_config=previous) - return ConfigurationCase( - experiment_api, - gateway, - AppTracingConfigService(configs=store, provider=gateway), - store, - ConsoleAppRepository(session_factory=sqlite_session_factory), - RequestContext("request", None, "account", workspace), - app_id, - ) - - -def test_post_masked_patch_and_get_preserve_ciphertext_and_exact_destination(case: ConfigurationCase) -> None: - original = case.read("langsmith") - submitted = { - "tracking_uri": " HTTP://FIRST.Example:5000/team/ ", - "experiment_id": "00012", - "username": "test-user", - "password": "synthetic-password", - } - untouched = dict(submitted) - case.service.create(case.context, case.app_id, "mlflow", submitted, activate=True) - before = case.read() - assert before is not None - assert before.tracing_config is not None - config = before.tracing_config - assert config["tracking_uri"] == "http://first.example:5000/team" - assert config["experiment_id"] == "12" - assert config["password"] != "synthetic-password" - assert OpsTraceManager.decrypt_tracing_config(case.context.active_workspace_id, "mlflow", config)["password"] == ( - "synthetic-password" - ) - assert submitted == untouched - shown = case.service.get(case.context, case.app_id, "mlflow") - assert shown is not None - assert shown.tracing_config is not None - masked = shown.tracing_config["password"] - assert isinstance(masked, str) - assert "*" in masked - assert masked != "synthetic-password" - assert len(case.api.calls) == 1 - case.service.update( - case.context, - case.app_id, - "mlflow", - { - "tracking_uri": "https://SECOND.test/new-prefix/", - "experiment_id": "00034", - "username": "test-user", - "password": masked, - }, - ) - after = case.read() - assert after is not None - assert after.tracing_config is not None - assert after.tracing_config["password"] == config["password"] - assert after.tracing_config["tracking_uri"] == "https://second.test/new-prefix" - assert after.tracing_config["experiment_id"] == "34" - assert after.tracing_config["project_url"] == "https://second.test/new-prefix/#/experiments/34/traces" - assert len(case.api.calls) == 2 - assert case.api.calls[-1][1] == "https://second.test/new-prefix/api/2.0/mlflow/experiments/get" - case.gateway.validate_saved_config( - workspace_id=case.context.active_workspace_id, tracing_provider="mlflow", tracing_config=after.tracing_config - ) - assert case.service.get(case.context, case.app_id, "mlflow") is not None - assert len(case.api.calls) == 2 - assert case.read("langsmith") == original - policy = case.reader.get_trace(case.context, case.app_id) - assert policy.enabled - assert policy.tracing_provider == "mlflow" - assert policy.capture_content is True - assert policy.capture_agent is False - - -@pytest.mark.parametrize("clear_with_null", [False, True]) -def test_basic_to_anonymous_and_back_does_not_reuse_credentials(case: ConfigurationCase, clear_with_null: bool) -> None: - destination = {"tracking_uri": "https://same.test/prefix", "experiment_id": "12"} - case.service.create( - case.context, case.app_id, "mlflow", {**destination, "username": "first", "password": "first-secret"} - ) - empty = None if clear_with_null else "" - case.service.update(case.context, case.app_id, "mlflow", {**destination, "username": empty, "password": empty}) - assert case.api.calls[-1][2]["headers"] == { - "Accept": "application/json", - "Accept-Encoding": "identity", - "Cookie": "", - } - stored = case.read() - assert stored is not None - assert stored.tracing_config is not None - assert ( - OpsTraceManager.decrypt_tracing_config(case.context.active_workspace_id, "mlflow", stored.tracing_config)[ - "password" - ] - == empty - ) - case.service.update( - case.context, case.app_id, "mlflow", {**destination, "username": "second", "password": "second-secret"} - ) - assert case.api.calls[-1][2]["headers"] != case.api.calls[0][2]["headers"] - - -def test_failed_patch_keeps_config_policy_and_other_provider( - case: ConfigurationCase, caplog: pytest.LogCaptureFixture -) -> None: - original = { - "tracking_uri": "https://first.test", - "experiment_id": "12", - "username": "user", - "password": "private-password", - } - case.service.create(case.context, case.app_id, "mlflow", original) - before, other, policy = case.read(), case.read("langsmith"), case.reader.get_trace(case.context, case.app_id) - case.api.override = lambda _method, url, _options: response(url, {"message": "private-vendor-response"}, 403) - with pytest.raises(AppTracingConfigVerificationFailedError) as caught: - case.service.update( - case.context, - case.app_id, - "mlflow", - {**original, "tracking_uri": "https://second.test", "password": "******"}, - activate=True, - ) - assert isinstance(caught.value.__cause__, MLflowConfigurationError) - assert caught.value.__cause__.code == "permission_denied" - assert case.read() == before - assert case.read("langsmith") == other - assert case.reader.get_trace(case.context, case.app_id) == policy - assert "permission_denied" in caplog.text - assert "private-vendor-response" not in caplog.text - assert "private-password" not in caplog.text - - -def test_failed_auth_pair_patch_does_not_erase_a_password(case: ConfigurationCase) -> None: - original = {"tracking_uri": "https://first.test", "experiment_id": "12", "username": "user", "password": "password"} - case.service.create(case.context, case.app_id, "mlflow", original) - before = case.read() - with pytest.raises(AppTracingConfigVerificationFailedError) as caught: - case.service.update(case.context, case.app_id, "mlflow", {**original, "password": None}) - assert isinstance(caught.value.__cause__, MLflowConfigurationError) - assert caught.value.__cause__.code == "incomplete_basic_auth" - assert len(case.api.calls) == 1 - assert case.read() == before - - -@pytest.mark.parametrize("saved_url", [None, "https://stale.test/wrong-experiment"]) -def test_historical_get_is_pure_even_with_partial_auth(case: ConfigurationCase, saved_url: str | None) -> None: - stored = OpsTraceManager.encrypt_tracing_config( - case.context.active_workspace_id, - "mlflow", - { - "tracking_uri": "https://server.test/prefix/", - "experiment_id": "00012", - "username": "historical-user", - "password": None, - }, - ) - if saved_url is not None: - stored["project_url"] = saved_url - shown = case.gateway.present_config( - workspace_id=case.context.active_workspace_id, tracing_provider="mlflow", tracing_config=stored - ) - assert shown["project_url"] == "https://server.test/prefix/#/experiments/12/traces" - assert shown["password"] is None - case.gateway.validate_saved_config( - workspace_id=case.context.active_workspace_id, tracing_provider="mlflow", tracing_config=stored - ) - assert case.api.calls == [] - - -def test_invalid_historical_url_has_no_fabricated_link_or_network(case: ConfigurationCase) -> None: - stored = OpsTraceManager.encrypt_tracing_config( - case.context.active_workspace_id, - "mlflow", - { - "tracking_uri": "https://user:secret@server.test?token=private", - "experiment_id": "12", - "username": None, - "password": None, - }, - ) - shown = case.gateway.present_config( - workspace_id=case.context.active_workspace_id, tracing_provider="mlflow", tracing_config=stored - ) - assert shown["project_url"] is None - assert case.api.calls == [] - - -def test_invalid_new_destination_never_reaches_storage_or_network(case: ConfigurationCase) -> None: - with pytest.raises(AppTracingConfigInvalidConfigurationError): - case.service.create( - case.context, - case.app_id, - "mlflow", - { - "tracking_uri": "https://server.test/prefix/../other", - "experiment_id": "12", - "username": "user", - "password": "secret", - }, - activate=True, - ) - assert case.read() is None - assert case.api.calls == [] - assert case.reader.get_trace(case.context, case.app_id).tracing_provider == "langsmith" diff --git a/api/tests/unit_tests/services/test_model_provider_service.py b/api/tests/unit_tests/services/test_model_provider_service.py deleted file mode 100644 index 980bb86b7a4..00000000000 --- a/api/tests/unit_tests/services/test_model_provider_service.py +++ /dev/null @@ -1,1369 +0,0 @@ -from types import SimpleNamespace -from typing import Any -from unittest.mock import MagicMock, create_autospec - -import pytest -from sqlalchemy import event -from sqlalchemy.engine import Engine -from sqlalchemy.exc import IntegrityError -from sqlalchemy.orm import Session - -from core.entities.model_entities import ModelStatus, ModelWithProviderEntity, SimpleModelProviderEntity -from core.entities.provider_configuration import ProviderConfigurations -from core.entities.provider_entities import CredentialConfiguration -from core.plugin.entities.plugin import PluginInstallationSource -from core.plugin.entities.plugin_daemon import PluginModelProviderBinding -from enums import DeploymentEdition -from graphon.model_runtime.entities.common_entities import I18nObject -from graphon.model_runtime.entities.model_entities import FetchFrom, ModelType, ParameterRule, ParameterType -from graphon.model_runtime.entities.provider_entities import ConfigurateMethod, ProviderEntity -from models.account import Account, TenantAccountRole -from models.provider import ( - Provider, - ProviderCredential, - ProviderModel, - ProviderType, - TenantDefaultModel, - TenantPreferredModelProvider, -) -from services.credentials.query import CredentialQuery, ModelCredentialRecord -from services.errors.app_model_config import ProviderNotFoundError -from services.model_provider import service as service_module -from services.model_provider.service import ModelProviderService, _ProviderSummaryState - - -def _create_service_with_mocked_manager() -> tuple[ModelProviderService, MagicMock]: - manager = MagicMock() - service = ModelProviderService() - service._get_provider_manager = MagicMock(return_value=manager) - return service, manager - - -def _build_available_llm(model: str) -> ModelWithProviderEntity: - provider = ProviderEntity( - provider="langgenius/openai/openai", - label=I18nObject(en_US="OpenAI"), - supported_model_types=[ModelType.LLM], - configurate_methods=[ConfigurateMethod.PREDEFINED_MODEL], - ) - return ModelWithProviderEntity( - model=model, - label=I18nObject(en_US=model), - model_type=ModelType.LLM, - fetch_from=FetchFrom.PREDEFINED_MODEL, - model_properties={}, - status=ModelStatus.ACTIVE, - provider=SimpleModelProviderEntity(provider), - ) - - -@pytest.mark.parametrize("has_active_models", [True, False]) -def test_initialize_default_llm_selects_first_active_model_once( - sqlite_session: Session, has_active_models: bool -) -> None: - sqlite_session.add_all( - [ - TenantDefaultModel( - tenant_id="other-tenant", model_type=ModelType.LLM, provider_name="other", model_name="other-model" - ), - TenantDefaultModel( - tenant_id="tenant-1", - model_type=ModelType.TEXT_EMBEDDING, - provider_name="embedding", - model_name="embedding-model", - ), - ] - ) - sqlite_session.commit() - service, manager = _create_service_with_mocked_manager() - models = [_build_available_llm("gpt-4o-mini"), _build_available_llm("gpt-4o")] if has_active_models else [] - configurations = manager.get_configurations.return_value - configurations.get_models.return_value = models - - service.initialize_default_llm("tenant-1") - service.initialize_default_llm("tenant-1") - - expected_defaults = { - ("other-tenant", ModelType.LLM, "other", "other-model"), - ("tenant-1", ModelType.TEXT_EMBEDDING, "embedding", "embedding-model"), - } - if has_active_models: - expected_defaults.add(("tenant-1", ModelType.LLM, "langgenius/openai/openai", "gpt-4o-mini")) - assert { - (record.tenant_id, record.model_type, record.provider_name, record.model_name) - for record in sqlite_session.query(TenantDefaultModel).all() - } == expected_defaults - assert configurations.get_models.call_count == (1 if has_active_models else 2) - configurations.get_models.assert_called_with(model_type=ModelType.LLM, only_active=True) - - -def test_initialize_default_llm_preserves_existing_selection(sqlite_session: Session) -> None: - saved_default = TenantDefaultModel( - tenant_id="tenant-1", model_type=ModelType.LLM, provider_name="chosen-provider", model_name="chosen-model" - ) - sqlite_session.add(saved_default) - sqlite_session.commit() - service, manager = _create_service_with_mocked_manager() - - service.initialize_default_llm("tenant-1") - - manager.get_configurations.assert_not_called() - sqlite_session.refresh(saved_default) - assert (saved_default.provider_name, saved_default.model_name) == ("chosen-provider", "chosen-model") - - -@pytest.mark.parametrize("selection_time", ["discovery", "insert"]) -def test_initialize_default_llm_preserves_concurrent_selection( - sqlite_session: Session, sqlite_engine: Engine, selection_time: str -) -> None: - service, manager = _create_service_with_mocked_manager() - selected = False - - def save_selection() -> None: - nonlocal selected - selected = True - sqlite_session.add( - TenantDefaultModel( - tenant_id="tenant-1", - model_type=ModelType.LLM, - provider_name="chosen-provider", - model_name="chosen-model", - ) - ) - sqlite_session.commit() - - def discover_models(*, model_type: ModelType, only_active: bool) -> list[ModelWithProviderEntity]: - assert model_type == ModelType.LLM - assert only_active - if selection_time == "discovery": - save_selection() - return [_build_available_llm("gpt-4o-mini")] - - manager.get_configurations.return_value.get_models.side_effect = discover_models - - def select_before_insert( - _connection: Any, _cursor: Any, statement: str, _parameters: Any, _context: Any, _executemany: bool - ) -> None: - if selection_time == "insert" and not selected and statement.startswith("INSERT INTO tenant_default_models"): - save_selection() - - event.listen(sqlite_engine, "before_cursor_execute", select_before_insert) - try: - service.initialize_default_llm("tenant-1") - finally: - event.remove(sqlite_engine, "before_cursor_execute", select_before_insert) - - defaults = sqlite_session.query(TenantDefaultModel).all() - assert selected - assert len(defaults) == 1 - assert (defaults[0].provider_name, defaults[0].model_name) == ("chosen-provider", "chosen-model") - - -def test_initialize_default_llm_propagates_database_failure(sqlite_session: Session, sqlite_engine: Engine) -> None: - service, manager = _create_service_with_mocked_manager() - manager.get_configurations.return_value.get_models.return_value = [_build_available_llm("gpt-4o-mini")] - with sqlite_engine.begin() as connection: - connection.exec_driver_sql( - "CREATE TRIGGER reject_default_model BEFORE INSERT ON tenant_default_models " - "BEGIN SELECT RAISE(ABORT, 'default model write rejected'); END" - ) - - with pytest.raises(IntegrityError, match="default model write rejected"): - service.initialize_default_llm("tenant-1") - - assert sqlite_session.query(TenantDefaultModel).count() == 0 - - -def test_available_credentials_use_injected_query_for_admin_viewer() -> None: - service, manager = _create_service_with_mocked_manager() - user = Account(name="Admin", email="admin@example.com") - user.id = "actor" - user.role = TenantAccountRole.ADMIN - query = create_autospec(CredentialQuery, instance=True, spec_set=True) - query.list_models.return_value = [ModelCredentialRecord(id="credential", name="Shared credential")] - - result = service.get_provider_available_credentials("tenant", "openai", user=user, credential_query=query) - - query.list_models.assert_called_once_with(workspace_id="tenant", provider="openai", actor_id="actor") - assert result == [CredentialConfiguration(credential_id="credential", credential_name="Shared credential")] - assert manager.mock_calls == [] - - -def _build_provider_configuration( - *, - provider_name: str = "openai", - supported_model_types: list[ModelType] | None = None, - custom_models: list[Any] | None = None, - custom_config_available: bool = True, -) -> SimpleNamespace: - if supported_model_types is None: - supported_model_types = [ModelType.LLM] - - return SimpleNamespace( - provider=SimpleNamespace( - provider=provider_name, - label=I18nObject(en_US=provider_name), - description=None, - icon_small=None, - icon_small_dark=None, - background=None, - help=None, - supported_model_types=supported_model_types, - configurate_methods=[], - provider_credential_schema=None, - model_credential_schema=None, - ), - preferred_provider_type=ProviderType.CUSTOM, - custom_configuration=SimpleNamespace( - provider=SimpleNamespace( - current_credential_id="cred-1", - current_credential_name="Credential 1", - available_credentials=[], - ), - models=custom_models, - can_added_models=[], - ), - system_configuration=SimpleNamespace(enabled=False, current_quota_type=None, quota_configurations=[]), - is_custom_configuration_available=lambda: custom_config_available, - ) - - -def _build_model_provider_binding( - source: PluginInstallationSource, - *, - installation_id: str = "installation-1", - plugin_id: str = "langgenius/openai", - plugin_unique_identifier: str = "langgenius/openai:1.0.0@checksum", - verified: bool = True, -) -> PluginModelProviderBinding: - return PluginModelProviderBinding( - provider="openai", - installation_id=installation_id, - plugin_id=plugin_id, - plugin_unique_identifier=plugin_unique_identifier, - runtime_type="remote" if source == PluginInstallationSource.Remote else "local", - source=source, - version="1.0.0", - verified=verified, - ) - - -class TestModelProviderServiceConfiguration: - def test__get_provider_configuration_should_return_configuration_when_provider_exists(self) -> None: - service, manager = _create_service_with_mocked_manager() - provider_configuration = SimpleNamespace(name="provider-config") - manager.get_configurations.return_value = {"openai": provider_configuration} - - result = service._get_provider_configuration(tenant_id="tenant-1", provider="openai") - - assert result is provider_configuration - - def test__get_provider_configuration_should_raise_error_when_provider_is_missing(self) -> None: - service, manager = _create_service_with_mocked_manager() - manager.get_configurations.return_value = {} - - with pytest.raises(ProviderNotFoundError, match="does not exist"): - service._get_provider_configuration(tenant_id="tenant-1", provider="missing") - - def test_get_provider_list_should_filter_by_model_type_and_build_no_configure_status(self) -> None: - service, manager = _create_service_with_mocked_manager() - allowed = _build_provider_configuration( - provider_name="openai", - supported_model_types=[ModelType.LLM], - custom_config_available=False, - ) - filtered = _build_provider_configuration( - provider_name="embedding", - supported_model_types=[ModelType.TEXT_EMBEDDING], - custom_config_available=True, - ) - manager.get_configurations.return_value = {"openai": allowed, "embedding": filtered} - - result = service.get_provider_list(tenant_id="tenant-1", model_type=ModelType.LLM) - - assert len(result) == 1 - assert result[0].provider == "openai" - assert result[0].custom_configuration.status.value == "no-configure" - - def test_get_provider_summary_list_uses_lightweight_state_and_plugin_bindings( - self, monkeypatch: pytest.MonkeyPatch - ) -> None: - service = ModelProviderService() - provider = SimpleNamespace( - provider="langgenius/openai/openai", - label=I18nObject(en_US="OpenAI"), - description=I18nObject(en_US="OpenAI models"), - icon_small=I18nObject(en_US="icon.svg"), - icon_small_dark=I18nObject(en_US="icon-dark.svg"), - supported_model_types=[ModelType.LLM], - configurate_methods=[ConfigurateMethod.PREDEFINED_MODEL], - ) - binding = SimpleNamespace( - provider="openai", - plugin_id="langgenius/openai", - installation_id="installation-1", - plugin_unique_identifier="langgenius/openai:1.2.3@checksum", - runtime_type="local", - source=PluginInstallationSource.Marketplace, - version="1.2.3", - verified=True, - ) - state = _ProviderSummaryState( - has_custom_provider=True, - available_credentials=[ - CredentialConfiguration( - credential_id="credential-1", - credential_name="Production", - ), - CredentialConfiguration( - credential_id="credential-2", - credential_name="Backup", - ), - ], - has_custom_models=True, - current_credential_id="credential-1", - current_credential_name="Production", - current_credential_usable=True, - preferred_provider_type=ProviderType.CUSTOM, - ) - call_order: list[str] = [] - manager_constructor = MagicMock(side_effect=AssertionError("summary must not construct ProviderManager")) - monkeypatch.setattr(service, "_get_provider_manager", manager_constructor) - monkeypatch.setattr( - service_module.PluginService, - "list_model_provider_bindings", - MagicMock(side_effect=lambda *_args, **_kwargs: call_order.append("bindings") or [binding]), - ) - monkeypatch.setattr( - service_module.PluginService, - "fetch_plugin_model_providers", - MagicMock(side_effect=lambda *_args, **_kwargs: call_order.append("providers") or [provider]), - ) - monkeypatch.setattr( - service, - "_load_provider_summary_states", - MagicMock(return_value={provider.provider: state}), - ) - monkeypatch.setattr( - service_module.ext_hosting_provider.hosting_configuration, - "provider_map", - {provider.provider: SimpleNamespace(enabled=True, quotas=[SimpleNamespace()])}, - ) - - providers, plugins = service.get_provider_summary_list(tenant_id="tenant-1") - - assert len(providers) == 1 - assert providers[0].provider == provider.provider - assert providers[0].plugin_id == "langgenius/openai" - assert providers[0].is_configured is True - assert providers[0].custom_configuration.available_credentials == [ - CredentialConfiguration( - credential_id="credential-1", - credential_name="Production", - ), - CredentialConfiguration( - credential_id="credential-2", - credential_name="Backup", - ), - ] - assert providers[0].custom_configuration.has_custom_models is True - assert providers[0].custom_configuration.current_credential_name == "Production" - assert providers[0].custom_configuration.current_credential_usable is True - assert providers[0].system_configuration.enabled is True - assert plugins["langgenius/openai"].installation_id == "installation-1" - assert plugins["langgenius/openai"].version == "1.2.3" - assert call_order == ["bindings", "providers"] - manager_constructor.assert_not_called() - - def test_get_provider_summary_list_enables_system_only_for_verified_hosted_non_package_bindings( - self, - monkeypatch: pytest.MonkeyPatch, - ) -> None: - marketplace_binding = _build_model_provider_binding(PluginInstallationSource.Marketplace) - package_binding = _build_model_provider_binding( - PluginInstallationSource.Package, - installation_id="installation-package", - plugin_id="langgenius/package", - plugin_unique_identifier="langgenius/package:1.0.0@checksum", - ) - unhosted_binding = _build_model_provider_binding( - PluginInstallationSource.Marketplace, - installation_id="installation-unhosted", - plugin_id="langgenius/unhosted", - plugin_unique_identifier="langgenius/unhosted:1.0.0@checksum", - ) - remote_binding = _build_model_provider_binding( - PluginInstallationSource.Remote, - installation_id="installation-remote", - plugin_id="langgenius/remote", - plugin_unique_identifier="langgenius/remote:1.0.0@checksum", - ) - unverified_binding = _build_model_provider_binding( - PluginInstallationSource.Marketplace, - installation_id="installation-unverified", - plugin_id="langgenius/unverified", - plugin_unique_identifier="langgenius/unverified:1.0.0@checksum", - verified=False, - ) - bindings = [ - marketplace_binding, - package_binding, - unhosted_binding, - remote_binding, - unverified_binding, - ] - provider_entities = [ - SimpleNamespace( - provider=f"{binding.plugin_id}/openai", - label=I18nObject(en_US=binding.plugin_id), - description=None, - icon_small=None, - icon_small_dark=None, - supported_model_types=[ModelType.LLM], - configurate_methods=[], - ) - for binding in bindings - ] - monkeypatch.setattr( - service_module.ext_hosting_provider.hosting_configuration, - "provider_map", - { - "langgenius/openai/openai": SimpleNamespace(enabled=True, quotas=[SimpleNamespace()]), - "langgenius/package/openai": SimpleNamespace(enabled=True, quotas=[SimpleNamespace()]), - "langgenius/remote/openai": SimpleNamespace(enabled=True, quotas=[SimpleNamespace()]), - "langgenius/unverified/openai": SimpleNamespace(enabled=True, quotas=[SimpleNamespace()]), - }, - ) - monkeypatch.setattr( - service_module.PluginService, - "list_model_provider_bindings", - MagicMock(return_value=bindings), - ) - monkeypatch.setattr( - service_module.PluginService, - "fetch_plugin_model_providers", - MagicMock(return_value=provider_entities), - ) - monkeypatch.setattr(ModelProviderService, "_load_provider_summary_states", MagicMock(return_value={})) - monkeypatch.setattr(service_module, "is_filtered", MagicMock(return_value=False)) - - providers, _ = ModelProviderService().get_provider_summary_list("tenant-1") - - assert {provider.provider: provider.system_configuration.enabled for provider in providers} == { - "langgenius/openai/openai": True, - "langgenius/package/openai": False, - "langgenius/unhosted/openai": False, - "langgenius/remote/openai": True, - "langgenius/unverified/openai": False, - } - - def test_model_provider_binding_without_verified_field_fails_closed(self) -> None: - binding = PluginModelProviderBinding.model_validate( - { - "provider": "openai", - "installation_id": "installation-1", - "plugin_id": "langgenius/openai", - "plugin_unique_identifier": "langgenius/openai:1.0.0@checksum", - "runtime_type": "local", - "source": "marketplace", - "version": "1.0.0", - } - ) - - assert binding.verified is False - - def test_get_provider_summary_list_returns_all_unique_provider_metadata( - self, monkeypatch: pytest.MonkeyPatch - ) -> None: - service = ModelProviderService() - llm_provider = SimpleNamespace( - provider="langgenius/openai/openai", - label=I18nObject(en_US="OpenAI"), - description=None, - icon_small=None, - icon_small_dark=None, - supported_model_types=[ModelType.LLM], - configurate_methods=[], - ) - embedding_provider = SimpleNamespace( - provider="langgenius/embedding/embedding", - label=I18nObject(en_US="Embedding"), - description=None, - icon_small=None, - icon_small_dark=None, - supported_model_types=[ModelType.TEXT_EMBEDDING], - configurate_methods=[], - ) - llm_binding = SimpleNamespace( - provider="openai", - plugin_id="langgenius/openai", - installation_id="installation-openai", - plugin_unique_identifier="langgenius/openai:1.0.0@checksum", - runtime_type="local", - source=service_module.PluginInstallationSource.Marketplace, - version="1.0.0", - verified=False, - ) - embedding_binding = SimpleNamespace( - provider="embedding", - plugin_id="langgenius/embedding", - installation_id="installation-embedding", - plugin_unique_identifier="langgenius/embedding:1.0.0@checksum", - runtime_type="local", - source=service_module.PluginInstallationSource.Marketplace, - version="1.0.0", - verified=False, - ) - monkeypatch.setattr( - service_module.PluginService, - "list_model_provider_bindings", - MagicMock(return_value=[llm_binding, embedding_binding]), - ) - monkeypatch.setattr( - service_module.PluginService, - "fetch_plugin_model_providers", - MagicMock(return_value=[llm_provider, llm_provider, embedding_provider]), - ) - monkeypatch.setattr(service, "_load_provider_summary_states", MagicMock(return_value={})) - monkeypatch.setattr(service_module, "is_filtered", MagicMock(return_value=False)) - - providers, plugins = service.get_provider_summary_list(tenant_id="tenant-1") - - assert [provider.provider for provider in providers] == [ - "langgenius/openai/openai", - "langgenius/embedding/embedding", - ] - assert providers[0].is_configured is False - assert providers[0].custom_configuration.status.value == "no-configure" - assert providers[0].custom_configuration.has_custom_models is False - assert providers[0].custom_configuration.available_credentials == [] - assert set(plugins) == {"langgenius/openai", "langgenius/embedding"} - - def test_preferred_provider_fallback_uses_custom_presence_not_configuration_status( - self, monkeypatch: pytest.MonkeyPatch - ) -> None: - from tests.unit_tests.config_override import apply_config_overrides - - apply_config_overrides(monkeypatch, DEPLOYMENT_EDITION=DeploymentEdition.COMMUNITY) - state = _ProviderSummaryState(has_custom_provider=True) - - preferred_provider_type = ModelProviderService._get_preferred_provider_type( - state, - custom_present=True, - system_enabled=True, - ) - - assert preferred_provider_type == ProviderType.CUSTOM - - def test_load_provider_summary_states_reads_only_lightweight_columns( - self, - sqlite_engine: Engine, - sqlite_session: Session, - ) -> None: - canonical_provider = "langgenius/openai/openai" - legacy_credential = ProviderCredential( - tenant_id="tenant-1", - provider_name="openai", - credential_name="Legacy", - encrypted_config="legacy-secret", - ) - legacy_credential.id = "credential-legacy" - current_credential = ProviderCredential( - tenant_id="tenant-1", - provider_name=canonical_provider, - credential_name="Production", - encrypted_config="production-secret", - ) - current_credential.id = "credential-current" - foreign_credential = ProviderCredential( - tenant_id="tenant-2", - provider_name=canonical_provider, - credential_name="Foreign", - encrypted_config="foreign-secret", - ) - foreign_credential.id = "credential-foreign" - sqlite_session.add_all( - [ - legacy_credential, - current_credential, - foreign_credential, - Provider( - tenant_id="tenant-1", - provider_name="openai", - provider_type=ProviderType.CUSTOM, - is_valid=True, - credential_id=legacy_credential.id, - ), - Provider( - tenant_id="tenant-1", - provider_name=canonical_provider, - provider_type=ProviderType.CUSTOM, - is_valid=True, - credential_id=current_credential.id, - ), - Provider( - tenant_id="tenant-2", - provider_name=canonical_provider, - provider_type=ProviderType.CUSTOM, - is_valid=True, - credential_id=foreign_credential.id, - ), - ProviderModel( - tenant_id="tenant-1", - provider_name="openai", - model_name="gpt-4o", - model_type=ModelType.LLM, - is_valid=True, - ), - ProviderModel( - tenant_id="tenant-2", - provider_name="langgenius/foreign/foreign", - model_name="foreign-model", - model_type=ModelType.LLM, - is_valid=True, - ), - TenantPreferredModelProvider( - tenant_id="tenant-1", - provider_name=canonical_provider, - preferred_provider_type=ProviderType.SYSTEM, - ), - TenantPreferredModelProvider( - tenant_id="tenant-2", - provider_name=canonical_provider, - preferred_provider_type=ProviderType.CUSTOM, - ), - ] - ) - sqlite_session.commit() - statements: list[str] = [] - - def capture_statement(_connection, _cursor, statement, _parameters, _context, _executemany) -> None: - statements.append(statement) - - event.listen(sqlite_engine, "before_cursor_execute", capture_statement) - try: - states = ModelProviderService._load_provider_summary_states("tenant-1") - finally: - event.remove(sqlite_engine, "before_cursor_execute", capture_statement) - - state = states[canonical_provider] - assert state.has_custom_provider is True - assert state.available_credentials == [ - CredentialConfiguration( - credential_id="credential-legacy", - credential_name="Legacy", - ), - CredentialConfiguration( - credential_id="credential-current", - credential_name="Production", - ), - ] - assert state.has_custom_models is True - assert state.current_credential_id == "credential-current" - assert state.current_credential_name == "Production" - assert state.current_credential_usable is True - assert state.preferred_provider_type == ProviderType.SYSTEM - - assert len(statements) == 4 - assert all("encrypted_config" not in statement for statement in statements) - assert "count(" not in statements[1].lower() - assert "provider_credentials.id" in statements[1] - assert "provider_credentials.credential_name" in statements[1] - assert "ORDER BY provider_credentials.created_at DESC, provider_credentials.id DESC" in statements[1] - assert "provider_model_credentials" in statements[2] - - def test_list_models_should_pass_filters_through(self, monkeypatch: pytest.MonkeyPatch) -> None: - service, manager = _create_service_with_mocked_manager() - filters: list[dict[str, object]] = [] - - def get_models(_self: ProviderConfigurations, **kwargs: object) -> list[ModelWithProviderEntity]: - filters.append(kwargs) - return [] - - monkeypatch.setattr(ProviderConfigurations, "get_models", get_models) - manager.get_configurations.return_value = ProviderConfigurations(tenant_id="tenant-1") - - assert service.list_models("tenant-1", provider=None, model_type=ModelType.LLM) == [] - assert filters == [{"provider": None, "model_type": ModelType.LLM}] - - def test_get_models_by_provider_should_wrap_model_entities_with_tenant_context(self) -> None: - service, manager = _create_service_with_mocked_manager() - - class _Model: - def __init__(self, model_name: str) -> None: - self.model_name = model_name - - def model_dump(self) -> dict[str, Any]: - return { - "model": self.model_name, - "label": {"en_US": self.model_name}, - "model_type": ModelType.LLM, - "features": [], - "fetch_from": FetchFrom.PREDEFINED_MODEL, - "model_properties": {}, - "deprecated": False, - "status": ModelStatus.ACTIVE, - "load_balancing_enabled": False, - "has_invalid_load_balancing_configs": False, - "provider": { - "provider": "openai", - "label": {"en_US": "OpenAI"}, - "icon_small": None, - "icon_small_dark": None, - "supported_model_types": [ModelType.LLM], - }, - } - - provider_configurations = SimpleNamespace( - get_models=MagicMock(return_value=[_Model("gpt-4o"), _Model("gpt-4o-mini")]) - ) - manager.get_configurations.return_value = provider_configurations - - result = service.get_models_by_provider(tenant_id="tenant-1", provider="openai") - - assert len(result) == 2 - assert result[0].model == "gpt-4o" - assert result[1].provider.provider == "openai" - provider_configurations.get_models.assert_called_once_with(provider="openai") - - -class TestModelProviderServiceDelegation: - @pytest.mark.parametrize( - ("method_name", "method_kwargs", "provider_method_name", "provider_call_kwargs", "provider_return"), - [ - ( - "get_provider_credential", - {"tenant_id": "tenant-1", "provider": "openai", "credential_id": "cred-1"}, - "get_provider_credential", - {"credential_id": "cred-1"}, - {"token": "abc"}, - ), - ( - "validate_provider_credentials", - {"tenant_id": "tenant-1", "provider": "openai", "credentials": {"token": "abc"}}, - "validate_provider_credentials", - ({"token": "abc"},), - None, - ), - ( - "create_provider_credential", - { - "tenant_id": "tenant-1", - "provider": "openai", - "credentials": {"token": "abc"}, - "credential_name": "A", - }, - "create_provider_credential", - ({"token": "abc"}, "A"), - "new-credential-id", - ), - ( - "update_provider_credential", - { - "tenant_id": "tenant-1", - "provider": "openai", - "credentials": {"token": "abc"}, - "credential_id": "cred-1", - "credential_name": "B", - }, - "update_provider_credential", - {"credential_id": "cred-1", "credentials": {"token": "abc"}, "credential_name": "B"}, - None, - ), - ( - "remove_provider_credential", - {"tenant_id": "tenant-1", "provider": "openai", "credential_id": "cred-1"}, - "delete_provider_credential", - {"credential_id": "cred-1"}, - None, - ), - ( - "switch_active_provider_credential", - {"tenant_id": "tenant-1", "provider": "openai", "credential_id": "cred-1"}, - "switch_active_provider_credential", - {"credential_id": "cred-1"}, - None, - ), - ], - ) - def test_provider_credential_methods_should_delegate_to_provider_configuration( - self, - method_name: str, - method_kwargs: dict[str, Any], - provider_method_name: str, - provider_call_kwargs: Any, - provider_return: Any, - monkeypatch: pytest.MonkeyPatch, - ) -> None: - service = ModelProviderService() - provider_configuration = MagicMock() - getattr(provider_configuration, provider_method_name).return_value = provider_return - get_provider_config_mock = MagicMock(return_value=provider_configuration) - monkeypatch.setattr(service, "_get_provider_configuration", get_provider_config_mock) - - result = getattr(service, method_name)(**method_kwargs) - - get_provider_config_mock.assert_called_once_with("tenant-1", "openai") - provider_method = getattr(provider_configuration, provider_method_name) - match provider_call_kwargs: - case tuple(): - provider_method.assert_called_once_with(*provider_call_kwargs) - case dict(): - provider_method.assert_called_once_with(**provider_call_kwargs) - case _: - provider_method.assert_called_once_with(provider_call_kwargs) - if method_name == "get_provider_credential": - assert result == {"token": "abc"} - if method_name == "create_provider_credential": - assert result == "new-credential-id" - - @pytest.mark.parametrize( - ("method_name", "method_kwargs", "provider_method_name", "expected_kwargs", "provider_return"), - [ - ( - "get_model_credential", - { - "tenant_id": "tenant-1", - "provider": "openai", - "model_type": ModelType.LLM, - "model": "gpt-4o", - "credential_id": "cred-1", - }, - "get_custom_model_credential", - {"model_type": ModelType.LLM, "model": "gpt-4o", "credential_id": "cred-1"}, - {"api_key": "x"}, - ), - ( - "validate_model_credentials", - { - "tenant_id": "tenant-1", - "provider": "openai", - "model_type": ModelType.LLM, - "model": "gpt-4o", - "credentials": {"api_key": "x"}, - }, - "validate_custom_model_credentials", - {"model_type": ModelType.LLM, "model": "gpt-4o", "credentials": {"api_key": "x"}}, - None, - ), - ( - "create_model_credential", - { - "tenant_id": "tenant-1", - "provider": "openai", - "model_type": ModelType.LLM, - "model": "gpt-4o", - "credentials": {"api_key": "x"}, - "credential_name": "cred-a", - }, - "create_custom_model_credential", - { - "model_type": ModelType.LLM, - "model": "gpt-4o", - "credentials": {"api_key": "x"}, - "credential_name": "cred-a", - }, - "new-credential-id", - ), - ( - "update_model_credential", - { - "tenant_id": "tenant-1", - "provider": "openai", - "model_type": ModelType.LLM, - "model": "gpt-4o", - "credentials": {"api_key": "x"}, - "credential_id": "cred-1", - "credential_name": "cred-b", - }, - "update_custom_model_credential", - { - "model_type": ModelType.LLM, - "model": "gpt-4o", - "credentials": {"api_key": "x"}, - "credential_id": "cred-1", - "credential_name": "cred-b", - }, - None, - ), - ( - "remove_model_credential", - { - "tenant_id": "tenant-1", - "provider": "openai", - "model_type": ModelType.LLM, - "model": "gpt-4o", - "credential_id": "cred-1", - }, - "delete_custom_model_credential", - {"model_type": ModelType.LLM, "model": "gpt-4o", "credential_id": "cred-1"}, - None, - ), - ( - "switch_active_custom_model_credential", - { - "tenant_id": "tenant-1", - "provider": "openai", - "model_type": ModelType.LLM, - "model": "gpt-4o", - "credential_id": "cred-1", - }, - "switch_custom_model_credential", - {"model_type": ModelType.LLM, "model": "gpt-4o", "credential_id": "cred-1"}, - None, - ), - ( - "add_model_credential_to_model_list", - { - "tenant_id": "tenant-1", - "provider": "openai", - "model_type": ModelType.LLM, - "model": "gpt-4o", - "credential_id": "cred-1", - }, - "add_model_credential_to_model", - {"model_type": ModelType.LLM, "model": "gpt-4o", "credential_id": "cred-1"}, - None, - ), - ( - "remove_model", - { - "tenant_id": "tenant-1", - "provider": "openai", - "model_type": ModelType.LLM, - "model": "gpt-4o", - }, - "delete_custom_model", - {"model_type": ModelType.LLM, "model": "gpt-4o"}, - None, - ), - ], - ) - def test_custom_model_methods_should_convert_model_type_and_delegate( - self, - method_name: str, - method_kwargs: dict[str, Any], - provider_method_name: str, - expected_kwargs: dict[str, Any], - provider_return: Any, - monkeypatch: pytest.MonkeyPatch, - ) -> None: - service = ModelProviderService() - provider_configuration = MagicMock() - getattr(provider_configuration, provider_method_name).return_value = provider_return - get_provider_config_mock = MagicMock(return_value=provider_configuration) - monkeypatch.setattr(service, "_get_provider_configuration", get_provider_config_mock) - - result = getattr(service, method_name)(**method_kwargs) - - get_provider_config_mock.assert_called_once_with("tenant-1", "openai") - getattr(provider_configuration, provider_method_name).assert_called_once_with(**expected_kwargs) - if method_name == "get_model_credential": - assert result == {"api_key": "x"} - if method_name == "create_model_credential": - assert result == "new-credential-id" - - @pytest.mark.parametrize( - ("method_name", "method_kwargs", "provider_method_name", "expected_kwargs"), - [ - ( - "get_model_credential", - { - "tenant_id": "tenant-1", - "provider": "openai", - "model_type": "llm", - "model": "gpt-4o", - "credential_id": "cred-1", - }, - "get_custom_model_credential", - {"model_type": ModelType.LLM, "model": "gpt-4o", "credential_id": "cred-1"}, - ), - ( - "create_model_credential", - { - "tenant_id": "tenant-1", - "provider": "openai", - "model_type": "llm", - "model": "gpt-4o", - "credentials": {"api_key": "x"}, - "credential_name": "cred-a", - }, - "create_custom_model_credential", - { - "model_type": ModelType.LLM, - "model": "gpt-4o", - "credentials": {"api_key": "x"}, - "credential_name": "cred-a", - }, - ), - ( - "remove_model", - { - "tenant_id": "tenant-1", - "provider": "openai", - "model_type": "llm", - "model": "gpt-4o", - }, - "delete_custom_model", - {"model_type": ModelType.LLM, "model": "gpt-4o"}, - ), - ], - ) - def test_custom_model_methods_use_model_type_constructor_directly( - self, - method_name: str, - method_kwargs: dict[str, Any], - provider_method_name: str, - expected_kwargs: dict[str, Any], - monkeypatch: pytest.MonkeyPatch, - ) -> None: - service = ModelProviderService() - provider_configuration = MagicMock() - get_provider_config_mock = MagicMock(return_value=provider_configuration) - monkeypatch.setattr(service, "_get_provider_configuration", get_provider_config_mock) - - getattr(service, method_name)(**method_kwargs) - - get_provider_config_mock.assert_called_once_with("tenant-1", "openai") - getattr(provider_configuration, provider_method_name).assert_called_once_with(**expected_kwargs) - - -class TestModelProviderServiceListingsAndDefaults: - def test_get_models_by_model_type_should_group_active_non_deprecated_models(self) -> None: - service, manager = _create_service_with_mocked_manager() - openai_provider = SimpleNamespace( - provider="openai", - label=I18nObject(en_US="OpenAI"), - icon_small=None, - icon_small_dark=None, - ) - anthropic_provider = SimpleNamespace( - provider="anthropic", - label=I18nObject(en_US="Anthropic"), - icon_small=None, - icon_small_dark=None, - ) - models = [ - SimpleNamespace( - provider=openai_provider, - model="gpt-4o", - label=I18nObject(en_US="GPT-4o"), - model_type=ModelType.LLM, - features=[], - fetch_from=FetchFrom.PREDEFINED_MODEL, - model_properties={}, - status=ModelStatus.ACTIVE, - load_balancing_enabled=False, - deprecated=False, - ), - SimpleNamespace( - provider=openai_provider, - model="old-openai", - label=I18nObject(en_US="Old OpenAI"), - model_type=ModelType.LLM, - features=[], - fetch_from=FetchFrom.PREDEFINED_MODEL, - model_properties={}, - status=ModelStatus.ACTIVE, - load_balancing_enabled=False, - deprecated=True, - ), - SimpleNamespace( - provider=anthropic_provider, - model="old-anthropic", - label=I18nObject(en_US="Old Anthropic"), - model_type=ModelType.LLM, - features=[], - fetch_from=FetchFrom.PREDEFINED_MODEL, - model_properties={}, - status=ModelStatus.ACTIVE, - load_balancing_enabled=False, - deprecated=True, - ), - ] - provider_configurations = SimpleNamespace(get_models=MagicMock(return_value=models)) - manager.get_configurations.return_value = provider_configurations - - result = service.get_models_by_model_type(tenant_id="tenant-1", model_type=ModelType.LLM) - - provider_configurations.get_models.assert_called_once_with(model_type=ModelType.LLM, only_active=True) - assert len(result) == 1 - assert result[0].provider == "openai" - assert len(result[0].models) == 1 - assert result[0].models[0].model == "gpt-4o" - - @pytest.mark.parametrize( - ("credentials", "schema", "expected_count"), - [ - (None, None, 0), - ({"api_key": "x"}, None, 0), - ( - {"api_key": "x"}, - SimpleNamespace( - parameter_rules=[ - ParameterRule( - name="temperature", - label=I18nObject(en_US="Temperature"), - type=ParameterType.FLOAT, - ) - ] - ), - 1, - ), - ], - ) - def test_get_model_parameter_rules_should_handle_missing_credentials_and_schema( - self, - credentials: dict[str, Any] | None, - schema: Any, - expected_count: int, - monkeypatch: pytest.MonkeyPatch, - ) -> None: - service = ModelProviderService() - provider_configuration = MagicMock() - provider_configuration.get_current_credentials.return_value = credentials - provider_configuration.get_model_schema.return_value = schema - monkeypatch.setattr(service, "_get_provider_configuration", MagicMock(return_value=provider_configuration)) - - result = service.get_model_parameter_rules(tenant_id="tenant-1", provider="openai", model="gpt-4o") - - assert len(result) == expected_count - provider_configuration.get_current_credentials.assert_called_once_with( - model_type=ModelType.LLM, - model="gpt-4o", - ) - if credentials: - provider_configuration.get_model_schema.assert_called_once_with( - model_type=ModelType.LLM, - model="gpt-4o", - credentials=credentials, - ) - else: - provider_configuration.get_model_schema.assert_not_called() - - def test_get_default_model_selection_uses_saved_workspace_choice(self, sqlite_session: Session) -> None: - sqlite_session.add( - TenantDefaultModel( - tenant_id="tenant-1", - model_type=ModelType.LLM, - provider_name="langgenius/openai/openai", - model_name="gpt-4o", - ) - ) - sqlite_session.commit() - service, manager = _create_service_with_mocked_manager() - - result = service.get_default_model_selection("tenant-1", ModelType.LLM, session=sqlite_session) - - assert result == ("langgenius/openai/openai", "gpt-4o") - manager.get_configurations.assert_not_called() - - def test_get_default_model_selection_uses_first_active_model_without_writing(self, sqlite_session: Session) -> None: - service, manager = _create_service_with_mocked_manager() - configurations = manager.get_configurations.return_value - configurations.get_models.return_value = [ - SimpleNamespace(model="gpt-4o", provider=SimpleNamespace(provider="langgenius/openai/openai")) - ] - - result = service.get_default_model_selection("tenant-1", ModelType.LLM, session=sqlite_session) - - assert result == ("langgenius/openai/openai", "gpt-4o") - configurations.get_models.assert_called_once_with(model_type=ModelType.LLM, only_active=True) - assert sqlite_session.query(TenantDefaultModel).filter_by(tenant_id="tenant-1").first() is None - - def test_get_default_model_selection_returns_none_when_provider_discovery_fails( - self, sqlite_session: Session - ) -> None: - service, manager = _create_service_with_mocked_manager() - manager.get_configurations.side_effect = RuntimeError("provider unavailable") - - result = service.get_default_model_selection("tenant-1", ModelType.LLM, session=sqlite_session) - - assert result is None - - def test_get_default_model_selection_returns_none_when_no_active_models(self, sqlite_session: Session) -> None: - service, manager = _create_service_with_mocked_manager() - manager.get_configurations.return_value.get_models.return_value = [] - - result = service.get_default_model_selection("tenant-1", ModelType.LLM, session=sqlite_session) - - assert result is None - manager.get_configurations.return_value.get_models.assert_called_once_with( - model_type=ModelType.LLM, only_active=True - ) - - def test_get_default_model_of_model_type_should_return_response_when_manager_returns_model(self) -> None: - service, manager = _create_service_with_mocked_manager() - manager.get_default_model.return_value = SimpleNamespace( - model="gpt-4o", - model_type=ModelType.LLM, - provider=SimpleNamespace( - provider="openai", - label=I18nObject(en_US="OpenAI"), - icon_small=None, - supported_model_types=[ModelType.LLM], - ), - ) - - result = service.get_default_model_of_model_type(tenant_id="tenant-1", model_type=ModelType.LLM) - - assert result is not None - assert result.model == "gpt-4o" - assert result.provider.provider == "openai" - manager.get_default_model.assert_called_once_with(tenant_id="tenant-1", model_type=ModelType.LLM) - - def test_get_default_model_of_model_type_should_return_none_when_manager_returns_none(self) -> None: - service, manager = _create_service_with_mocked_manager() - manager.get_default_model.return_value = None - - result = service.get_default_model_of_model_type(tenant_id="tenant-1", model_type=ModelType.LLM) - - assert result is None - - def test_get_default_model_of_model_type_should_return_none_when_no_configuration_is_saved(self) -> None: - service, manager = _create_service_with_mocked_manager() - manager.get_default_model.side_effect = RuntimeError("boom") - - result = service.get_default_model_of_model_type(tenant_id="tenant-1", model_type=ModelType.LLM) - - assert result is None - - @pytest.mark.parametrize( - "error", - [ValueError("Invalid provider: langgenius/openai/openai"), RuntimeError("Plugin daemon unavailable")], - ) - def test_get_default_model_of_model_type_should_preserve_saved_configuration_on_provider_failure( - self, sqlite_session: Session, error: Exception - ) -> None: - saved_model = TenantDefaultModel( - tenant_id="tenant-1", - model_type=ModelType.LLM, - provider_name="langgenius/openai/openai", - model_name="gpt-4o", - ) - sqlite_session.add(saved_model) - sqlite_session.commit() - service, manager = _create_service_with_mocked_manager() - manager.get_default_model.side_effect = error - - result = service.get_default_model_of_model_type(tenant_id="tenant-1", model_type=ModelType.LLM) - - assert result is not None - assert result.model == "gpt-4o" - assert result.model_type == ModelType.LLM - assert result.provider.provider == "langgenius/openai/openai" - assert result.provider.tenant_id == "tenant-1" - assert result.provider.label.en_us == "langgenius/openai/openai" - assert result.provider.label.zh_hans == "langgenius/openai/openai" - assert result.provider.icon_small is None - assert result.provider.supported_model_types == [] - assert result.model_dump(mode="json")["model_type"] == "llm" - sqlite_session.refresh(saved_model) - assert saved_model.model_name == "gpt-4o" - assert saved_model.provider_name == "langgenius/openai/openai" - - @pytest.mark.parametrize( - ("tenant_id", "model_type"), - [("other-tenant", ModelType.LLM), ("tenant-1", ModelType.TEXT_EMBEDDING)], - ) - def test_get_default_model_of_model_type_should_not_fall_back_to_another_tenant_or_model_type( - self, sqlite_session: Session, tenant_id: str, model_type: ModelType - ) -> None: - sqlite_session.add( - TenantDefaultModel( - tenant_id=tenant_id, - model_type=model_type, - provider_name="langgenius/openai/openai", - model_name="configured-model", - ) - ) - sqlite_session.commit() - service, manager = _create_service_with_mocked_manager() - manager.get_default_model.side_effect = ValueError("Invalid provider: langgenius/openai/openai") - - result = service.get_default_model_of_model_type(tenant_id="tenant-1", model_type=ModelType.LLM) - - assert result is None - - def test_update_default_model_of_model_type_should_delegate_to_provider_manager(self) -> None: - service, manager = _create_service_with_mocked_manager() - - service.update_default_model_of_model_type( - tenant_id="tenant-1", - model_type=ModelType.LLM, - provider="openai", - model="gpt-4o", - ) - - manager.update_default_model_record.assert_called_once_with( - tenant_id="tenant-1", - model_type=ModelType.LLM, - provider="openai", - model="gpt-4o", - ) - - def test_get_model_provider_icon_should_fetch_icon_bytes_from_factory( - self, - monkeypatch: pytest.MonkeyPatch, - ) -> None: - service = ModelProviderService() - factory_instance = MagicMock() - factory_instance.get_provider_icon.return_value = (b"icon-bytes", "image/png") - factory_constructor = MagicMock(return_value=factory_instance) - monkeypatch.setattr(service_module, "create_plugin_model_provider_factory", factory_constructor) - - result = service.get_model_provider_icon( - tenant_id="tenant-1", - provider="openai", - icon_type="icon_small", - lang="en_US", - ) - - factory_constructor.assert_called_once_with(tenant_id="tenant-1") - factory_instance.get_provider_icon.assert_called_once_with("openai", "icon_small", "en_US") - assert result == (b"icon-bytes", "image/png") - - def test_switch_preferred_provider_should_convert_enum_and_delegate( - self, - monkeypatch: pytest.MonkeyPatch, - ) -> None: - service = ModelProviderService() - provider_configuration = MagicMock() - monkeypatch.setattr(service, "_get_provider_configuration", MagicMock(return_value=provider_configuration)) - - service.switch_preferred_provider( - tenant_id="tenant-1", - provider="openai", - preferred_provider_type=ProviderType.SYSTEM.value, - ) - - provider_configuration.switch_preferred_provider_type.assert_called_once_with(ProviderType.SYSTEM) - - @pytest.mark.parametrize( - ("method_name", "provider_method_name"), - [ - ("enable_model", "enable_model"), - ("disable_model", "disable_model"), - ], - ) - def test_model_enablement_methods_should_convert_model_type_and_delegate( - self, - method_name: str, - provider_method_name: str, - monkeypatch: pytest.MonkeyPatch, - ) -> None: - service = ModelProviderService() - provider_configuration = MagicMock() - monkeypatch.setattr(service, "_get_provider_configuration", MagicMock(return_value=provider_configuration)) - - getattr(service, method_name)( - tenant_id="tenant-1", - provider="openai", - model="gpt-4o", - model_type=ModelType.LLM, - ) - - getattr(provider_configuration, provider_method_name).assert_called_once_with( - model="gpt-4o", - model_type=ModelType.LLM, - ) diff --git a/api/tests/unit_tests/services/test_model_provider_service_sanitization.py b/api/tests/unit_tests/services/test_model_provider_service_sanitization.py deleted file mode 100644 index f0474364faa..00000000000 --- a/api/tests/unit_tests/services/test_model_provider_service_sanitization.py +++ /dev/null @@ -1,87 +0,0 @@ -import types - -import pytest - -from core.entities.provider_entities import CredentialConfiguration, CustomModelConfiguration -from graphon.model_runtime.entities.common_entities import I18nObject -from graphon.model_runtime.entities.model_entities import ModelType -from graphon.model_runtime.entities.provider_entities import ConfigurateMethod -from models.provider import ProviderType -from services.model_provider.service import ModelProviderService - - -class _FakeConfigurations: - def __init__(self, provider_configuration: types.SimpleNamespace) -> None: - self._provider_configuration = provider_configuration - - def values(self) -> list[types.SimpleNamespace]: - return [self._provider_configuration] - - -@pytest.fixture -def service_with_fake_configurations(monkeypatch: pytest.MonkeyPatch) -> ModelProviderService: - # Build a fake provider schema with minimal fields used by ProviderResponse - fake_provider = types.SimpleNamespace( - provider="langgenius/openai_api_compatible/openai_api_compatible", - label=I18nObject(en_US="OpenAI API Compatible", zh_Hans="OpenAI API Compatible"), - description=None, - icon_small=None, - icon_small_dark=None, - background=None, - help=None, - supported_model_types=[ModelType.LLM], - configurate_methods=[ConfigurateMethod.CUSTOMIZABLE_MODEL], - provider_credential_schema=None, - model_credential_schema=None, - ) - - # Include decrypted credentials to simulate the leak source - custom_model = CustomModelConfiguration( - model="gpt-4o-mini", - model_type=ModelType.LLM, - credentials={"api_key": "sk-plain-text", "endpoint": "https://example.com"}, - current_credential_id="cred-1", - current_credential_name="API KEY 1", - available_model_credentials=[], - unadded_to_model_list=False, - ) - - fake_custom_provider = types.SimpleNamespace( - current_credential_id="cred-1", - current_credential_name="API KEY 1", - available_credentials=[CredentialConfiguration(credential_id="cred-1", credential_name="API KEY 1")], - ) - - fake_custom_configuration = types.SimpleNamespace( - provider=fake_custom_provider, models=[custom_model], can_added_models=[] - ) - - fake_system_configuration = types.SimpleNamespace(enabled=False, current_quota_type=None, quota_configurations=[]) - - fake_provider_configuration = types.SimpleNamespace( - provider=fake_provider, - preferred_provider_type=ProviderType.CUSTOM, - custom_configuration=fake_custom_configuration, - system_configuration=fake_system_configuration, - is_custom_configuration_available=lambda: True, - ) - - class _FakeProviderManager: - def get_configurations(self, _tenant_id: str) -> _FakeConfigurations: - return _FakeConfigurations(fake_provider_configuration) - - svc = ModelProviderService() - monkeypatch.setattr(svc, "_get_provider_manager", lambda _tenant_id: _FakeProviderManager()) - return svc - - -def test_get_provider_list_strips_credentials(service_with_fake_configurations: ModelProviderService) -> None: - providers = service_with_fake_configurations.get_provider_list(tenant_id="tenant-1", model_type=None) - - assert len(providers) == 1 - custom_models = providers[0].custom_configuration.custom_models - - assert custom_models is not None - assert len(custom_models) == 1 - # The sanitizer should drop credentials in list response - assert custom_models[0].credentials is None diff --git a/api/tests/unit_tests/services/test_oauth_device_application_service.py b/api/tests/unit_tests/services/test_oauth_device_application_service.py index 3f6e4991186..004b2cfbcb2 100644 --- a/api/tests/unit_tests/services/test_oauth_device_application_service.py +++ b/api/tests/unit_tests/services/test_oauth_device_application_service.py @@ -1,14 +1,24 @@ +from __future__ import annotations + import logging from dataclasses import dataclass, field from datetime import UTC, datetime -from unittest.mock import patch +from unittest.mock import create_autospec, patch import pytest +from constants.oauth_bearer import TokenType +from libs.oauth_bearer import sha256_hex from machinery.context import RequestContext from services.entities.account_entities import AccountSnapshot -from services.oauth_device_application_service import OAuthDeviceApplicationService +from services.oauth_device_adapters import OAuthDeviceTokenIssuanceGateway +from services.oauth_device_application_service import ( + OAuthDeviceApplicationService, + OAuthDeviceTokenPersistence, + OAuthDeviceTokenTTLPolicy, +) from services.oauth_device_contracts import ( + ACCOUNT_ISSUER_SENTINEL, AccessDeniedError, AlreadyResolvedError, ApprovalInProgressError, @@ -1008,3 +1018,71 @@ def test_approval_context_requires_verified_grant() -> None: assert result.user_code == harness.sso.grant.user_code assert result.csrf_token == harness.sso.grant.csrf_token assert result.expires_at == harness.sso.grant.expires_at + + +"""Minting normalises the issuer by the subject's declared account binding.""" + + +@pytest.mark.parametrize("token_type", [TokenType.OAUTH_ACCOUNT, TokenType.OAUTH_EXTERNAL_SSO], ids=lambda t: t.value) +def test_the_issuer_follows_the_subject_binding(token_type: TokenType) -> None: + tokens = create_autospec(OAuthDeviceTokenPersistence, instance=True) + tokens.rotate_token.return_value = OAuthDeviceTokenRotation( + token_id="token-1", replaced_token_id=None, replaced_token_hash=None + ) + ttl_policy = create_autospec(OAuthDeviceTokenTTLPolicy, instance=True) + ttl_policy.ttl_days.return_value = 1 + gateway = OAuthDeviceTokenIssuanceGateway(tokens=tokens, ttl_policy=ttl_policy) + + if token_type.subject.bound_to_account: + account = AccountSnapshot( + id="account-1", + name="Caller", + email="who@example.com", + avatar=None, + is_password_set=True, + interface_language=None, + interface_theme=None, + timezone=None, + last_login_at=None, + last_login_ip=None, + status="active", + initialized_at=None, + created_at=datetime.now(UTC), + ) + issued = gateway.issue_account_token( + account=account, workspace_id="workspace-1", client_id="client", device_label="device" + ) + expected_issuer = ACCOUNT_ISSUER_SENTINEL + expected_account = account.id + else: + issued = gateway.issue_external_token( + subject_email="who@example.com", + subject_issuer="https://idp.example.com", + client_id="client", + device_label="device", + ) + expected_issuer = "https://idp.example.com" + expected_account = None + + tokens.rotate_token.assert_called_once() + written = tokens.rotate_token.call_args.args[0] + assert written.subject_issuer == expected_issuer + assert written.account_id == expected_account + assert written.prefix == token_type.prefix + assert TokenType.for_token(issued.token) is token_type + assert written.token_hash == sha256_hex(issued.token) + + +@pytest.mark.parametrize("issuer", ["", " "]) +def test_external_issuer_is_required_before_persistence(issuer: str) -> None: + tokens = create_autospec(OAuthDeviceTokenPersistence, instance=True) + gateway = OAuthDeviceTokenIssuanceGateway( + tokens=tokens, ttl_policy=create_autospec(OAuthDeviceTokenTTLPolicy, instance=True) + ) + + with pytest.raises(ValueError, match="subject_issuer"): + gateway.issue_external_token( + subject_email="who@example.com", subject_issuer=issuer, client_id="client", device_label="device" + ) + + tokens.rotate_token.assert_not_called() diff --git a/api/tests/unit_tests/services/test_oauth_device_flow_mint.py b/api/tests/unit_tests/services/test_oauth_device_flow_mint.py deleted file mode 100644 index db27de6b91d..00000000000 --- a/api/tests/unit_tests/services/test_oauth_device_flow_mint.py +++ /dev/null @@ -1,80 +0,0 @@ -"""Minting normalises the issuer by the subject's declared account binding.""" - -from __future__ import annotations - -from datetime import UTC, datetime -from unittest.mock import create_autospec - -import pytest - -from constants.oauth_bearer import TokenType -from libs.oauth_bearer import sha256_hex -from services.entities.account_entities import AccountSnapshot -from services.oauth_device_adapters import OAuthDeviceTokenIssuanceGateway -from services.oauth_device_application_service import OAuthDeviceTokenPersistence, OAuthDeviceTokenTTLPolicy -from services.oauth_device_contracts import ACCOUNT_ISSUER_SENTINEL, OAuthDeviceTokenRotation - - -@pytest.mark.parametrize("token_type", [TokenType.OAUTH_ACCOUNT, TokenType.OAUTH_EXTERNAL_SSO], ids=lambda t: t.value) -def test_the_issuer_follows_the_subject_binding(token_type: TokenType) -> None: - tokens = create_autospec(OAuthDeviceTokenPersistence, instance=True) - tokens.rotate_token.return_value = OAuthDeviceTokenRotation( - token_id="token-1", replaced_token_id=None, replaced_token_hash=None - ) - ttl_policy = create_autospec(OAuthDeviceTokenTTLPolicy, instance=True) - ttl_policy.ttl_days.return_value = 1 - gateway = OAuthDeviceTokenIssuanceGateway(tokens=tokens, ttl_policy=ttl_policy) - - if token_type.subject.bound_to_account: - account = AccountSnapshot( - id="account-1", - name="Caller", - email="who@example.com", - avatar=None, - is_password_set=True, - interface_language=None, - interface_theme=None, - timezone=None, - last_login_at=None, - last_login_ip=None, - status="active", - initialized_at=None, - created_at=datetime.now(UTC), - ) - issued = gateway.issue_account_token( - account=account, workspace_id="workspace-1", client_id="client", device_label="device" - ) - expected_issuer = ACCOUNT_ISSUER_SENTINEL - expected_account = account.id - else: - issued = gateway.issue_external_token( - subject_email="who@example.com", - subject_issuer="https://idp.example.com", - client_id="client", - device_label="device", - ) - expected_issuer = "https://idp.example.com" - expected_account = None - - tokens.rotate_token.assert_called_once() - written = tokens.rotate_token.call_args.args[0] - assert written.subject_issuer == expected_issuer - assert written.account_id == expected_account - assert written.prefix == token_type.prefix - assert TokenType.for_token(issued.token) is token_type - assert written.token_hash == sha256_hex(issued.token) - - -@pytest.mark.parametrize("issuer", ["", " "]) -def test_external_issuer_is_required_before_persistence(issuer: str) -> None: - tokens = create_autospec(OAuthDeviceTokenPersistence, instance=True) - gateway = OAuthDeviceTokenIssuanceGateway( - tokens=tokens, ttl_policy=create_autospec(OAuthDeviceTokenTTLPolicy, instance=True) - ) - - with pytest.raises(ValueError, match="subject_issuer"): - gateway.issue_external_token( - subject_email="who@example.com", subject_issuer=issuer, client_id="client", device_label="device" - ) - - tokens.rotate_token.assert_not_called() diff --git a/api/tests/unit_tests/services/test_opik_tracing_config_gateway.py b/api/tests/unit_tests/services/test_opik_tracing_config_gateway.py deleted file mode 100644 index e72e22706b1..00000000000 --- a/api/tests/unit_tests/services/test_opik_tracing_config_gateway.py +++ /dev/null @@ -1,211 +0,0 @@ -"""Opik configuration transactions use real RSA and never initialize an SDK streamer.""" - -import json -from collections.abc import Iterator -from dataclasses import dataclass -from typing import Never -from uuid import uuid4 - -import httpx -import opik -import pytest -from dify_trace_opik.configuration import OpikConfigurationError -from dify_trace_opik.opik_trace import OpikDataTrace -from opik.config import OpikConfig as SDKOpikConfig -from sqlalchemy.orm import Session, sessionmaker - -from core.helper import ssrf_proxy -from core.ops.ops_trace_manager import OpsTraceManager -from machinery.context import RequestContext -from models.model import App -from repositories.app.console_repository import ConsoleAppRepository -from repositories.app.tracing_config_repository import SQLAlchemyAppTracingConfigRepository -from services.app_tracing_config_gateway import OpsTraceManagerGateway -from services.app_tracing_config_service import AppTracingConfigService, AppTracingConfigVerificationFailedError -from tests.unit_tests.model_factories import make_app -from tests.unit_tests.services.tracing.test_opik_configuration import AccessAPI - -pytest_plugins = ("tests.unit_tests.services.tracing.test_nullable_tracing_credentials",) - - -@dataclass -class ConfigurationCase: - api: AccessAPI - gateway: OpsTraceManagerGateway - service: AppTracingConfigService - store: SQLAlchemyAppTracingConfigRepository - reader: ConsoleAppRepository - context: RequestContext - app_id: str - sessions: sessionmaker[Session] - - -@pytest.fixture -def case( - workspace: str, sqlite_session_factory: sessionmaker[Session], monkeypatch: pytest.MonkeyPatch -) -> Iterator[ConfigurationCase]: - app_id = str(uuid4()) - with sqlite_session_factory.begin() as session: - session.add(make_app(app_id=app_id, tenant_id=workspace)) - api = AccessAPI() - gateway = OpsTraceManagerGateway() - store = SQLAlchemyAppTracingConfigRepository(session_factory=sqlite_session_factory) - forbidden_calls: list[str] = [] - - def forbidden(*_args: object, **_kwargs: object) -> Never: - forbidden_calls.append("sdk_constructor") - raise AssertionError("Configuration must not initialize an SDK client or ambient configuration") - - monkeypatch.setattr(ssrf_proxy, "make_request", api.request) - monkeypatch.setattr(OpikDataTrace, "__init__", forbidden) - monkeypatch.setattr(opik.Opik, "__init__", forbidden) - monkeypatch.setattr(SDKOpikConfig, "__init__", forbidden) - yield ConfigurationCase( - api, - gateway, - AppTracingConfigService(configs=store, provider=gateway), - store, - ConsoleAppRepository(session_factory=sqlite_session_factory), - RequestContext("request", None, "account", workspace), - app_id, - sqlite_session_factory, - ) - assert forbidden_calls == [] - - -def test_post_get_patch_preserve_real_ciphertext_and_verified_destination(case: ConfigurationCase) -> None: - secret = "synthetic-private-opik-key" - case.service.create( - case.context, - case.app_id, - "opik", - {"api_key": secret, "workspace": "", "project": " first ", "url": "https://first.test/prefix/api/"}, - activate=True, - ) - before = case.store.get(workspace_id=case.context.active_workspace_id, app_id=case.app_id, tracing_provider="opik") - assert before is not None - assert before.tracing_config is not None - config = before.tracing_config - assert config["api_key"] != secret - assert config["workspace"] == "resolved-team" - assert config["project"] == "first" - assert config["project_url"] == "https://first.test/prefix/resolved-team/redirect/projects?name=first" - assert OpsTraceManager.decrypt_tracing_config(case.context.active_workspace_id, "opik", config)["api_key"] == secret - assert case.reader.get_trace(case.context, case.app_id).tracing_provider == "opik" - - before_get = len(case.api.calls) - presented = case.service.get(case.context, case.app_id, "opik") - assert presented is not None - assert presented.tracing_config is not None - masked = presented.tracing_config["api_key"] - assert masked != secret - assert "*" in masked - assert len(case.api.calls) == before_get - case.api.default_workspace = "new-team" - case.service.update( - case.context, - case.app_id, - "opik", - {"api_key": masked, "workspace": "default", "project": "A & B", "url": "https://second.test/api/"}, - ) - after = case.store.get(workspace_id=case.context.active_workspace_id, app_id=case.app_id, tracing_provider="opik") - assert after is not None - assert after.tracing_config is not None - assert after.tracing_config["api_key"] == config["api_key"] - assert after.tracing_config["workspace"] == "new-team" - assert after.tracing_config["project"] == "A & B" - assert after.tracing_config["url"] == "https://second.test/api/" - assert after.tracing_config["project_url"] == "https://second.test/new-team/redirect/projects?name=A+%26+B" - assert case.api.calls[-1][2]["headers"] == { - "Authorization": secret, - "Comet-Workspace": "default", - "Accept": "application/json", - } - before_get = len(case.api.calls) - presented = case.service.get(case.context, case.app_id, "opik") - assert presented is not None - assert presented.tracing_config is not None - assert presented.tracing_config["project_url"] == after.tracing_config["project_url"] - case.gateway.validate_saved_config( - workspace_id=case.context.active_workspace_id, tracing_provider="opik", tracing_config=after.tracing_config - ) - assert len(case.api.calls) == before_get - - -@pytest.mark.parametrize("saved_url", [None, "https://wrong-global.test/old-project"]) -@pytest.mark.parametrize("workspace_name", ["team", "default", None]) -def test_historical_get_uses_only_configuration_host_without_network( - case: ConfigurationCase, saved_url: str | None, workspace_name: str | None -) -> None: - stored = OpsTraceManager.encrypt_tracing_config( - case.context.active_workspace_id, - "opik", - { - "api_key": "synthetic-key", - "workspace": workspace_name, - "project": "A & B", - "url": "https://private.test/opik/api/", - }, - ) - if saved_url: - stored["project_url"] = saved_url - result = case.gateway.present_config( - workspace_id=case.context.active_workspace_id, tracing_provider="opik", tracing_config=stored - ) - expected = ( - "https://private.test/opik/team/redirect/projects?name=A+%26+B" - if workspace_name == "team" - else "https://private.test/opik/" - ) - assert result["project_url"] == expected - assert result["api_key"] != "synthetic-key" - assert case.api.calls == [] - - -@pytest.mark.parametrize("status", [401, 403]) -@pytest.mark.parametrize("existing", [False, True]) -def test_bad_key_or_workspace_does_not_save_or_change_active_provider( - case: ConfigurationCase, status: int, existing: bool, caplog: pytest.LogCaptureFixture -) -> None: - original = {"api_key": "synthetic-private-key", "workspace": "team", "project": "original"} - if existing: - case.service.create(case.context, case.app_id, "opik", original) - with case.sessions.begin() as session: - app = session.get(App, case.app_id) - assert app is not None - app.tracing = json.dumps({"enabled": True, "tracing_provider": "langsmith", "capture_content": False}) - before_policy = case.reader.get_trace(case.context, case.app_id) - before = case.store.get(workspace_id=case.context.active_workspace_id, app_id=case.app_id, tracing_provider="opik") - case.api.override = lambda method, url, _options: httpx.Response( - status, text="private-vendor-details", request=httpx.Request(method, url) - ) - operation = case.service.update if existing else case.service.create - with pytest.raises(AppTracingConfigVerificationFailedError) as error: - operation(case.context, case.app_id, "opik", {**original, "workspace": "inaccessible"}, activate=True) - assert isinstance(error.value.__cause__, OpikConfigurationError) - assert error.value.__cause__.code == ("authentication_failed" if status == 401 else "permission_denied") - after = case.store.get(workspace_id=case.context.active_workspace_id, app_id=case.app_id, tracing_provider="opik") - assert after == before - assert case.reader.get_trace(case.context, case.app_id) == before_policy - assert "private-vendor-details" not in caplog.text - assert "synthetic-private-key" not in caplog.text - - -def test_clear_optional_api_key_keeps_keyless_self_host_valid( - case: ConfigurationCase, monkeypatch: pytest.MonkeyPatch -) -> None: - target = {"url": "http://local.test:5173/api/", "workspace": "team", "project": "sample"} - case.service.create(case.context, case.app_id, "opik", {**target, "api_key": "synthetic-key"}) - monkeypatch.setenv("OPIK_API_KEY", "unrelated-key") - case.api.calls.clear() - case.service.update(case.context, case.app_id, "opik", {**target, "api_key": None}) - stored = case.store.get(workspace_id=case.context.active_workspace_id, app_id=case.app_id, tracing_provider="opik") - assert stored is not None - assert stored.tracing_config is not None - assert stored.tracing_config["api_key"] is None - presented = case.service.get(case.context, case.app_id, "opik") - assert presented is not None - assert presented.tracing_config is not None - assert presented.tracing_config["api_key"] is None - assert case.api.calls[0][2]["headers"] == {"Comet-Workspace": "team", "Accept": "application/json"} - assert len(case.api.calls) == 1 diff --git a/api/tests/unit_tests/services/test_package_icons.py b/api/tests/unit_tests/services/test_package_icons.py deleted file mode 100644 index c5ea988eafe..00000000000 --- a/api/tests/unit_tests/services/test_package_icons.py +++ /dev/null @@ -1,207 +0,0 @@ -"""Packaged image icons travel with Apps and receive destination-owned file ids.""" - -import hashlib -import io -import zipfile -from collections.abc import Callable -from types import SimpleNamespace -from unittest.mock import MagicMock, Mock -from uuid import uuid4 - -import pytest -import yaml - -from services.agent.errors import ( - InvalidRosterAgentPackageError, - RosterAgentPackageExportFailedError, - RosterAgentPackageTooLargeError, -) -from services.agent.package_resource_exporter import AgentPackageResourceExporter, _FileSource, _SkillSource -from services.agent.package_resource_importer import AgentPackageResourceImporter -from services.agent.roster_package_entities import PackageIcon, PreparedPackageArchive, RosterAgentPackageMember -from services.app_package_service import AppPackageService - - -def _icon_archive(*, payload: bytes = b"image", digest: str | None = None, reference: str = "i_000001") -> io.BytesIO: - dsl = yaml.safe_dump({"kind": "app", "app": {"mode": "chat", "icon_type": "image", "icon": reference}}) - with AppPackageService().export(dsl=dsl.replace(reference, "legacy-id"), name="Example") as exported: - with zipfile.ZipFile(exported.archive) as package: - manifest = yaml.safe_load(package.read("manifest.yaml")) - app_bytes = dsl.encode() - manifest["apps"][0].update(size=len(app_bytes), sha256=hashlib.sha256(app_bytes).hexdigest()) - manifest["icons"] = [ - { - "id": "i_000001", - "path": "i_000001.png", - "size": len(payload), - "sha256": digest or hashlib.sha256(payload).hexdigest(), - } - ] - output = io.BytesIO() - with zipfile.ZipFile(output, "w") as package: - package.writestr("manifest.yaml", yaml.safe_dump(manifest)) - package.writestr("app.yaml", app_bytes) - package.writestr("i_000001.png", payload) - output.seek(0) - return output - - -def test_app_icon_export_deduplicates_uploads_and_restores_references(monkeypatch: pytest.MonkeyPatch) -> None: - source_id = str(uuid4()) - payload = b"image payload" - backend = Mock() - backend.load_stream.side_effect = lambda _: iter([payload]) - exporter = AgentPackageResourceExporter(storage_backend=backend) - upload_lookup = Mock(return_value={source_id: SimpleNamespace(key="source/icon.png", extension="png")}) - monkeypatch.setattr(exporter, "_upload_files", upload_lookup) - metadata = {"mode": "chat", "icon_type": "image", "icon": source_id} - duplicate = dict(metadata) - exporter.collect_icon(session=Mock(), tenant_id="source-tenant", metadata=metadata) - exporter.collect_icon(session=Mock(), tenant_id="source-tenant", metadata=duplicate) - assert metadata["icon"] == duplicate["icon"] == "i_000001" - assert upload_lookup.call_args.kwargs["tenant_id"] == "source-tenant" - assert upload_lookup.call_count == 1 - dsl = yaml.safe_dump({"kind": "app", "app": metadata}) - with AppPackageService().export(dsl=dsl, name="Example", resources=exporter) as exported: - with zipfile.ZipFile(exported.archive) as package: - assert package.read("i_000001.png") == payload - exported.archive.seek(0) - prepared = AppPackageService().read_package(exported.archive) - assert prepared is not None - with prepared: - restore = Mock(return_value={"i_000001": "destination-id"}) - monkeypatch.setattr(AgentPackageResourceImporter, "materialize_icons", restore) - data = yaml.safe_load(prepared.dsl) - prepared.materialize_icons(data=data, tenant_id="destination", account_id="account") - assert data["app"]["icon"] == "destination-id" - assert restore.call_args.kwargs["tenant_id"] == "destination" - - -def test_app_package_restores_independent_site_image_icon(monkeypatch: pytest.MonkeyPatch) -> None: - app_icon_id, site_icon_id = str(uuid4()), str(uuid4()) - backend = Mock() - backend.load_stream.side_effect = lambda key: iter([key.encode()]) - exporter = AgentPackageResourceExporter(storage_backend=backend) - monkeypatch.setattr( - exporter, - "_upload_files", - Mock( - side_effect=lambda *, file_ids, **_kwargs: { - file_id: SimpleNamespace(key=f"source/{file_id}.png", extension="png") for file_id in file_ids - } - ), - ) - app = {"mode": "chat", "icon_type": "image", "icon": app_icon_id} - site = {"title": "Site", "icon_type": "image", "icon": site_icon_id} - exporter.collect_icon(session=Mock(), tenant_id="source", metadata=app) - exporter.collect_icon(session=Mock(), tenant_id="source", metadata=site) - with AppPackageService().export( - dsl=yaml.safe_dump({"kind": "app", "app": app, "site": site}), name="App", resources=exporter - ) as exported: - prepared = AppPackageService().read_package(exported.archive) - assert prepared is not None - with prepared: - assert len(prepared.icons) == 2 - monkeypatch.setattr( - AgentPackageResourceImporter, - "materialize_icons", - Mock(return_value={app["icon"]: "new-app-icon", site["icon"]: "new-site-icon"}), - ) - data = yaml.safe_load(prepared.dsl) - prepared.materialize_icons(data=data, tenant_id="destination", account_id="account") - assert data["app"]["icon"] == "new-app-icon" - assert data["site"]["icon"] == "new-site-icon" - - -def test_icon_export_rejects_unavailable_tenant_upload(monkeypatch: pytest.MonkeyPatch) -> None: - exporter = AgentPackageResourceExporter(storage_backend=Mock()) - monkeypatch.setattr(exporter, "_upload_files", Mock(return_value={})) - with pytest.raises(RosterAgentPackageExportFailedError, match="unavailable"): - exporter.collect_icon(session=Mock(), tenant_id="tenant", metadata={"icon_type": "image", "icon": str(uuid4())}) - - -def test_icon_checksum_is_required() -> None: - with pytest.raises(InvalidRosterAgentPackageError, match="integrity"): - AppPackageService().read_package(_icon_archive(digest="0" * 64)) - - -@pytest.mark.parametrize("reference", ["i_000002", "unrelated-upload"]) -def test_icon_must_resolve_to_metadata(reference: str) -> None: - with pytest.raises(InvalidRosterAgentPackageError, match="references"): - AppPackageService().read_package(_icon_archive(reference=reference)) - - -def test_icon_image_limit_is_enforced(config_overrides: Callable[..., None]) -> None: - config_overrides(UPLOAD_IMAGE_FILE_SIZE_LIMIT=0) - with pytest.raises(RosterAgentPackageTooLargeError, match="image size"): - AppPackageService().read_package(_icon_archive()) - - -def test_materialized_icon_is_owned_by_destination(monkeypatch: pytest.MonkeyPatch) -> None: - payload = b"image payload" - digest = hashlib.sha256(payload).hexdigest() - archive = io.BytesIO() - with zipfile.ZipFile(archive, "w") as package: - package.writestr("i_000001.png", payload) - prepared = PreparedPackageArchive( - archive=archive, members={"i_000001.png": RosterAgentPackageMember(size=len(payload), sha256=digest)} - ) - icon = PackageIcon(id="i_000001", path="i_000001.png", size=len(payload), sha256=digest) - session = MagicMock() - monkeypatch.setattr( - "services.agent.package_resource_importer.session_factory.create_session", - MagicMock(return_value=session), - ) - backend = Mock() - with prepared: - mapping = AgentPackageResourceImporter(storage_backend=backend).materialize_icons( - archive=prepared, icons=[icon], tenant_id="destination", account_id="importer" - ) - row = session.__enter__.return_value.add_all.call_args.args[0][0] - assert row.id == mapping["i_000001"] - assert row.tenant_id == "destination" - assert row.created_by == "importer" - assert row.mime_type == "image/png" - assert row.key.startswith("upload_files/destination/") - backend.save.assert_called_once_with(row.key, payload) - - -def test_roster_export_embeds_agent_icon(monkeypatch: pytest.MonkeyPatch) -> None: - from models.agent_config_entities import AgentSoulConfig - from models.model import App, AppMode - from services.agent.dsl_entities import AgentPackage, AgentPackageMetadata, make_agent_app_dsl - from services.agent.roster_package_exporter import RosterAgentPackageExporter - from services.agent.roster_package_reader import RosterAgentPackageReader - from services.entities.site_dsl import SiteDsl - - backend = Mock() - backend.load_stream.side_effect = lambda _: iter([b"icon"]) - resources = AgentPackageResourceExporter(storage_backend=backend) - resources.sources["agent_1"] = (list[_SkillSource](), list[_FileSource]()) - metadata = {"name": "Agent", "icon_type": "image", "icon": "source-id"} - monkeypatch.setattr( - resources, - "_upload_files", - Mock( - return_value={ - "source-id": SimpleNamespace(key="icon.png", extension="png"), - } - ), - ) - resources.collect_icon(session=Mock(), tenant_id="source", metadata=metadata) - app = make_agent_app_dsl( - App(name="Agent", mode=AppMode.AGENT), - package_ref="agent_1", - packages={ - "agent_1": AgentPackage(metadata=AgentPackageMetadata.model_validate(metadata), soul=AgentSoulConfig()) - }, - dependencies=[], - ) - app.site = SiteDsl(title="Agent Site", icon_type="image", icon="i_000001") - with RosterAgentPackageExporter(storage_backend=backend)._build_archive(app=app, resources=resources) as exported: - with RosterAgentPackageReader().read(exported.archive) as prepared: - assert prepared.apps["app.yaml"].package.metadata.icon == "i_000001" - assert prepared.apps["app.yaml"].site is not None - assert prepared.apps["app.yaml"].site.icon == "i_000001" - assert prepared.manifest.icons[0].path == "i_000001.png" - assert RosterAgentPackageReader().read_member_bytes(prepared, "i_000001.png", max_bytes=4) == b"icon" diff --git a/api/tests/unit_tests/services/test_paused_workflow_stop.py b/api/tests/unit_tests/services/test_paused_workflow_stop.py deleted file mode 100644 index 14f9378e827..00000000000 --- a/api/tests/unit_tests/services/test_paused_workflow_stop.py +++ /dev/null @@ -1,296 +0,0 @@ -"""Paused stops persist before publishing terminal events or retiring resources.""" - -import json -from collections.abc import Callable -from datetime import datetime - -import pytest -from redis import Redis -from sqlalchemy.orm import Session, sessionmaker - -import services.app_task_service as task_module -from core.repositories.sqlalchemy_workflow_node_execution_repository import SQLAlchemyWorkflowNodeExecutionRepository -from core.workflow.nodes.agent_v2.session_store import WorkflowAgentWorkspaceStore -from extensions.logstore.aliyun_logstore import AliyunLogStore -from graphon.entities import WorkflowNodeExecution -from graphon.enums import WorkflowExecutionStatus, WorkflowNodeExecutionStatus, WorkflowType -from libs.broadcast_channel.redis.pubsub_channel import Topic -from models import EndUser -from models.enums import CreatorUserRole, WorkflowRunTriggeredFrom -from models.model import AppMode -from models.workflow import WorkflowNodeExecutionModel, WorkflowNodeExecutionTriggeredFrom, WorkflowPause, WorkflowRun -from repositories.sqlalchemy_api_workflow_run_repository import DifyAPISQLAlchemyWorkflowRunRepository -from services.app_task_service import AppTaskService -from tests.unit_tests.services.test_app_task_service import _StopRedis - - -@pytest.mark.parametrize("app_mode", [AppMode.WORKFLOW, AppMode.ADVANCED_CHAT]) -@pytest.mark.parametrize("backend", ["sql", "celery", "logstore"]) -@pytest.mark.parametrize("node_status", [WorkflowNodeExecutionStatus.PAUSED, WorkflowNodeExecutionStatus.RETRY]) -def test_paused_stop_publishes_terminal_event_after_commit_and_retires_run_resources( - sqlite_session: Session, - sqlite_session_factory: sessionmaker[Session], - monkeypatch: pytest.MonkeyPatch, - app_mode: AppMode, - backend: str, - node_status: WorkflowNodeExecutionStatus, - config_overrides: Callable[..., None], -) -> None: - config_overrides( - CORE_WORKFLOW_EXECUTION_REPOSITORY={ - "sql": "core.repositories.sqlalchemy_workflow_execution_repository.SQLAlchemyWorkflowExecutionRepository", - "celery": "core.repositories.celery_workflow_execution_repository.CeleryWorkflowExecutionRepository", - "logstore": ( - "extensions.logstore.repositories.logstore_workflow_execution_repository." - "LogstoreWorkflowExecutionRepository" - ), - }[backend], - CORE_WORKFLOW_NODE_EXECUTION_REPOSITORY={ - "sql": ( - "core.repositories.sqlalchemy_workflow_node_execution_repository." - "SQLAlchemyWorkflowNodeExecutionRepository" - ), - "celery": ( - "core.repositories.celery_workflow_node_execution_repository.CeleryWorkflowNodeExecutionRepository" - ), - "logstore": ( - "extensions.logstore.repositories.logstore_workflow_node_execution_repository." - "LogstoreWorkflowNodeExecutionRepository" - ), - }[backend], - LOGSTORE_DUAL_WRITE_ENABLED=False, - ) - monkeypatch.setattr( - "repositories.sqlalchemy_api_workflow_run_repository.naive_utc_now", lambda: datetime(2026, 10, 10, 0, 1) - ) - logstore_records: dict[str, list[dict[str, str]]] = {} - monkeypatch.setattr(AliyunLogStore, "_instance", None) - monkeypatch.setattr(AliyunLogStore, "_initialized", False) - monkeypatch.setenv("ALIYUN_SLS_ENDPOINT", "localhost") - monkeypatch.setattr(AliyunLogStore, "_check_endpoint_connectivity", staticmethod(lambda *_: None)) - monkeypatch.setattr( - AliyunLogStore, - "put_log", - lambda _self, name, record: logstore_records.setdefault(name, []).append(dict(record)), - ) - queued_node_saves: list[object] = [] - queued_saves: list[object] = [] - monkeypatch.setattr( - "core.repositories.celery_workflow_node_execution_repository.save_workflow_node_executions_task.delay", - lambda *args, **kwargs: queued_node_saves.append((args, kwargs)), - ) - monkeypatch.setattr( - "core.repositories.celery_workflow_execution_repository.save_workflow_execution_task.delay", - lambda *args, **kwargs: queued_saves.append((args, kwargs)), - ) - run = WorkflowRun( - id="run-1", - tenant_id="tenant-1", - app_id="app-1", - workflow_id="workflow-1", - type=WorkflowType.WORKFLOW, - triggered_from=WorkflowRunTriggeredFrom.APP_RUN, - version="1", - status=WorkflowExecutionStatus.PAUSED, - created_by_role=CreatorUserRole.END_USER, - created_by="user-1", - created_at=datetime(2026, 10, 10), - elapsed_time=1, - total_tokens=0, - total_steps=2, - exceptions_count=0, - ) - run.task_id = "11111111-1111-1111-1111-111111111111" - pause = WorkflowPause(workflow_id=run.workflow_id, workflow_run_id=run.id, state_object_key="state") - sqlite_session.add_all([run, pause]) - sqlite_session.commit() - for node_app, origin in ( - ("app-1", WorkflowNodeExecutionTriggeredFrom.WORKFLOW_RUN), - ("source-app", WorkflowNodeExecutionTriggeredFrom.WORKFLOW_TOOL), - ): - SQLAlchemyWorkflowNodeExecutionRepository( - session_factory=sqlite_session_factory, - tenant_id="tenant-1", - user=EndUser(id="user-1"), - app_id=node_app, - triggered_from=origin, - ).save( - WorkflowNodeExecution( - id=f"node-{node_app}", - node_execution_id=f"node-{node_app}", - workflow_id=f"workflow-{node_app}", - workflow_execution_id=run.id, - index=1, - node_id="review", - node_type="human-input", - title="Review", - status=node_status, - process_data={"__dify_retry_history": [{"retry_index": 1, "error": "retry failed"}]}, - created_at=run.created_at, - ) - ) - monkeypatch.setattr(type(task_module.db), "engine", property(lambda _: sqlite_session_factory.kw["bind"])) - redis = _StopRedis() - monkeypatch.setattr(task_module, "redis_client", redis) - monkeypatch.setattr("repositories.sqlalchemy_api_workflow_run_repository.storage.delete", lambda _: None) - calls: list[str] = [] - - def retire(*, tenant_id: str, app_id: str, workflow_run_id: str) -> list[str]: - assert (tenant_id, app_id, workflow_run_id) == ("tenant-1", "app-1", "run-1") - with sqlite_session_factory() as session: - persisted = session.get(WorkflowRun, workflow_run_id) - assert persisted is not None - assert persisted.status == WorkflowExecutionStatus.STOPPED - assert persisted.elapsed_time == 1 - node = session.get(WorkflowNodeExecutionModel, "node-app-1") - assert node is not None - assert node.status == WorkflowNodeExecutionStatus.FAILED - assert node.finished_at == datetime(2026, 10, 10, 0, 1) - assert node.process_data_dict == {"__dify_retry_history": [{"retry_index": 1, "error": "retry failed"}]} - calls.append("retire") - return ["workspace-1"] - - monkeypatch.setattr(WorkflowAgentWorkspaceStore, "retire_workflow_run", staticmethod(retire)) - collections: list[dict[str, object]] = [] - - def enqueue(**kwargs: object) -> None: - calls.append("collect") - collections.append(kwargs) - - monkeypatch.setattr("tasks.collect_agent_resources_task.enqueue_agent_resource_collection", enqueue) - publications: list[bytes] = [] - - def publish(_channel: str, payload: bytes) -> None: - calls.append("publish") - publications.append(payload) - - topic_client = Redis() - monkeypatch.setattr(topic_client, "publish", publish) - topic = Topic(topic_client, "workflow-events") - monkeypatch.setattr(task_module.MessageBasedAppGenerator, "get_response_topic", lambda *_: topic) - - assert run.task_id is not None - AppTaskService.stop_workflow_task( - tenant_id="tenant-1", - app_id="app-1", - task_id=run.task_id, - app_mode=app_mode, - owner=(CreatorUserRole.END_USER, "user-1"), - ) - - assert calls == ["retire", "collect", "publish"] - assert collections == [{"tenant_id": "tenant-1", "workspace_ids": ["workspace-1"]}] - assert len(publications) == 1 - event = json.loads(publications[0]) - assert event["event"] == "workflow_finished" - assert event["data"]["status"] == "stopped" - assert event["data"]["elapsed_time"] == 1 - assert event["workflow_run_id"] == "run-1" - assert event["task_id"] == run.task_id - assert redis.reads == [] - assert redis.operations == [] - assert queued_saves == [] - if backend == "logstore": - assert len(logstore_records[AliyunLogStore.workflow_execution_logstore]) == 1 - history = logstore_records[AliyunLogStore.workflow_execution_logstore][0] - assert history["status"] == "stopped" - assert float(history["elapsed_time"]) == 1 - assert (history["tenant_id"], history["app_id"], history["id"]) == ("tenant-1", "app-1", "run-1") - node_history = logstore_records[AliyunLogStore.workflow_node_execution_logstore] - assert all(float(record["elapsed_time"]) == 60 for record in node_history) - assert all( - json.loads(record["process_data"])["__dify_retry_history"][0]["error"] == "retry failed" - for record in node_history - ) - assert {(record["app_id"], record["triggered_from"], record["status"]) for record in node_history} == { - ("app-1", "workflow-run", "failed"), - ("source-app", "workflow-tool", "failed"), - } - - -@pytest.mark.parametrize("task_id_input", ["11111111-aaaa-1111-1111-111111111111", "11111111AAAA11111111111111111111"]) -def test_stop_after_resume_claim_uses_durable_owner_before_queue_is_recreated( - sqlite_session: Session, - sqlite_session_factory: sessionmaker[Session], - monkeypatch: pytest.MonkeyPatch, - task_id_input: str, -) -> None: - run = WorkflowRun( - id="run-1", - tenant_id="tenant-1", - app_id="app-1", - workflow_id="workflow-1", - type=WorkflowType.WORKFLOW, - triggered_from=WorkflowRunTriggeredFrom.APP_RUN, - version="1", - status=WorkflowExecutionStatus.PAUSED, - created_by_role=CreatorUserRole.END_USER, - created_by="user-1", - ) - run.task_id = "11111111-aaaa-1111-1111-111111111111" - pause = WorkflowPause(workflow_id=run.workflow_id, workflow_run_id=run.id, state_object_key="state") - sqlite_session.add_all([run, pause]) - sqlite_session.commit() - repository = DifyAPISQLAlchemyWorkflowRunRepository(sqlite_session_factory) - entity = repository.get_workflow_pause(run.id) - assert entity is not None - repository.resume_workflow_pause(run.id, entity) - redis = _StopRedis() - monkeypatch.setattr(type(task_module.db), "engine", property(lambda _: sqlite_session_factory.kw["bind"])) - monkeypatch.setattr(task_module, "redis_client", redis) - AppTaskService.stop_workflow_task( - tenant_id="tenant-1", - app_id="app-1", - task_id=task_id_input, - app_mode=AppMode.ADVANCED_CHAT, - owner=(CreatorUserRole.END_USER, "user-1"), - ) - assert redis.reads == [] - assert redis.values[f"generate_task_stopped:{run.task_id}"] == b"1" - assert json.loads(redis.commands[f"workflow:{run.task_id}:commands"][0])["command_type"] == "abort" - - -@pytest.mark.parametrize( - ("tenant_id", "app_id", "owner"), - [ - ("other-tenant", "app-1", None), - ("tenant-1", "other-app", None), - ("tenant-1", "app-1", (CreatorUserRole.END_USER, "other-user")), - ("tenant-1", "app-1", (CreatorUserRole.ACCOUNT, "user-1")), - ], -) -def test_running_task_admission_rejects_foreign_tenant_app_user_and_role( - sqlite_session: Session, - sqlite_session_factory: sessionmaker[Session], - monkeypatch: pytest.MonkeyPatch, - tenant_id: str, - app_id: str, - owner: tuple[CreatorUserRole, str] | None, -) -> None: - task_id = "22222222-2222-2222-2222-222222222222" - run = WorkflowRun( - id="run-1", - tenant_id="tenant-1", - app_id="app-1", - workflow_id="workflow-1", - task_id=task_id, - type=WorkflowType.WORKFLOW, - triggered_from=WorkflowRunTriggeredFrom.APP_RUN, - version="1", - status=WorkflowExecutionStatus.RUNNING, - created_by_role=CreatorUserRole.END_USER, - created_by="user-1", - ) - sqlite_session.add(run) - sqlite_session.commit() - monkeypatch.setattr(type(task_module.db), "engine", property(lambda _: sqlite_session_factory.kw["bind"])) - redis = _StopRedis(values={f"generate_task_belong:{task_id}": b"end-user-user-1"}) - monkeypatch.setattr(task_module, "redis_client", redis) - AppTaskService.stop_workflow_task( - tenant_id=tenant_id, app_id=app_id, task_id=task_id, app_mode=AppMode.WORKFLOW, owner=owner - ) - sqlite_session.expire_all() - assert run.status == WorkflowExecutionStatus.RUNNING - assert run.stop_requested_at is None - assert redis.reads == [] - assert redis.operations == [] diff --git a/api/tests/unit_tests/services/test_phoenix_tracing_config_gateway.py b/api/tests/unit_tests/services/test_phoenix_tracing_config_gateway.py deleted file mode 100644 index 3f07363c34f..00000000000 --- a/api/tests/unit_tests/services/test_phoenix_tracing_config_gateway.py +++ /dev/null @@ -1,285 +0,0 @@ -"""Persist verified Phoenix configuration while preserving owner policy and masked keys.""" - -import json -from collections import deque -from dataclasses import dataclass -from uuid import uuid4 - -import httpx -import pytest -from dify_trace_arize_phoenix.configuration import PhoenixConfigurationError -from sqlalchemy.orm import Session, sessionmaker - -from core.ops.ops_trace_manager import OpsTraceManager -from machinery.context import RequestContext -from repositories.app.console_repository import ConsoleAppRepository -from repositories.app.tracing_config_repository import SQLAlchemyAppTracingConfigRepository -from services.app_tracing_config_gateway import OpsTraceManagerGateway -from services.app_tracing_config_service import ( - AppTracingConfigInvalidConfigurationError, - AppTracingConfigRecord, - AppTracingConfigService, - AppTracingConfigVerificationFailedError, -) -from tests.unit_tests.model_factories import make_app -from tests.unit_tests.services.tracing.test_phoenix_configuration import ( - api_calls as api_calls, # noqa: PLC0414 - Re-export pytest fixtures only in this module. -) -from tests.unit_tests.services.tracing.test_phoenix_configuration import ( - api_responses as api_responses, # noqa: PLC0414 - Re-export pytest fixtures only in this module. -) -from tests.unit_tests.services.tracing.test_phoenix_configuration import ( - no_legacy_or_environment_auth as no_legacy_or_environment_auth, # noqa: PLC0414 - Re-export pytest fixtures. -) -from tests.unit_tests.services.tracing.test_phoenix_configuration import response - -pytest_plugins = ("tests.unit_tests.services.tracing.test_nullable_tracing_credentials",) - - -@dataclass -class PhoenixCase: - gateway: OpsTraceManagerGateway - service: AppTracingConfigService - store: SQLAlchemyAppTracingConfigRepository - reader: ConsoleAppRepository - context: RequestContext - app_id: str - - def read(self, provider: str = "phoenix") -> AppTracingConfigRecord | None: - return self.store.get( - workspace_id=self.context.active_workspace_id, app_id=self.app_id, tracing_provider=provider - ) - - -@pytest.fixture -def case(workspace: str, sqlite_session_factory: sessionmaker[Session]) -> PhoenixCase: - app_id = str(uuid4()) - with sqlite_session_factory.begin() as session: - session.add( - make_app( - app_id=app_id, - tenant_id=workspace, - tracing=json.dumps( - { - "enabled": True, - "tracing_provider": "langsmith", - "capture_content": True, - "capture_agent": False, - } - ), - ) - ) - gateway = OpsTraceManagerGateway() - store = SQLAlchemyAppTracingConfigRepository(session_factory=sqlite_session_factory) - old = OpsTraceManager.encrypt_tracing_config( - workspace, "langsmith", {"api_key": "synthetic-other-key", "project": "other", "endpoint": "https://other.test"} - ) - assert store.create(workspace_id=workspace, app_id=app_id, tracing_provider="langsmith", tracing_config=old) - return PhoenixCase( - gateway, - AppTracingConfigService(configs=store, provider=gateway), - store, - ConsoleAppRepository(session_factory=sqlite_session_factory), - RequestContext("request", None, "account", workspace), - app_id, - ) - - -def test_post_masked_patch_get_keep_normalized_destination_and_ciphertext( - case: PhoenixCase, - api_calls: list[tuple[str, dict[str, object]]], -) -> None: - other = case.read("langsmith") - submitted = { - "endpoint": " HTTP://FIRST.Example:6006/reverse/ ", - "project": " first project ", - "api_key": "synthetic-secret", - } - untouched = dict(submitted) - case.service.create(case.context, case.app_id, "phoenix", submitted, activate=True) - before = case.read() - assert before is not None - assert before.tracing_config is not None - stored = before.tracing_config - assert stored["endpoint"] == "http://first.example:6006/reverse" - assert stored["project"] == "first project" - assert stored["api_key"] != "synthetic-secret" - assert submitted == untouched - assert OpsTraceManager.decrypt_tracing_config(case.context.active_workspace_id, "phoenix", stored)["api_key"] == ( - "synthetic-secret" - ) - shown = case.service.get(case.context, case.app_id, "phoenix") - assert shown is not None - assert shown.tracing_config is not None - masked = shown.tracing_config["api_key"] - assert isinstance(masked, str) - assert "*" in masked - assert masked != "synthetic-secret" - assert shown.tracing_config["project_url"] == "http://first.example:6006/reverse/redirects/projects/first%20project" - assert len(api_calls) == 1 - case.service.update( - case.context, - case.app_id, - "phoenix", - { - "endpoint": " https://SECOND.test/s/space/ ", - "project": " second & 项目 ", - "api_key": masked, - }, - ) - after = case.read() - assert after is not None - assert after.tracing_config is not None - assert after.tracing_config["api_key"] == stored["api_key"] - assert after.tracing_config["endpoint"] == "https://second.test/s/space" - assert after.tracing_config["project"] == "second & 项目" - assert after.tracing_config["project_url"] == ( - "https://second.test/s/space/redirects/projects/second%20%26%20%E9%A1%B9%E7%9B%AE" - ) - assert api_calls[-1][0] == "https://second.test/s/space/v1/projects" - headers = api_calls[-1][1]["headers"] - assert isinstance(headers, dict) - assert headers["Authorization"] == "Bearer synthetic-secret" - case.gateway.validate_saved_config( - workspace_id=case.context.active_workspace_id, tracing_provider="phoenix", tracing_config=after.tracing_config - ) - assert case.service.get(case.context, case.app_id, "phoenix") is not None - assert len(api_calls) == 2 - assert case.read("langsmith") == other - policy = case.reader.get_trace(case.context, case.app_id) - assert policy.enabled - assert policy.tracing_provider == "phoenix" - assert policy.capture_content is True - assert policy.capture_agent is False - - -@pytest.mark.parametrize("empty", [None, ""]) -def test_bearer_to_keyless_and_back_is_explicit( - case: PhoenixCase, - api_calls: list[tuple[str, dict[str, object]]], - empty: str | None, -) -> None: - destination = {"endpoint": "https://host.test", "project": "project"} - case.service.create(case.context, case.app_id, "phoenix", {**destination, "api_key": "first-key"}) - case.service.update(case.context, case.app_id, "phoenix", {**destination, "api_key": empty}) - assert api_calls[-1][1]["headers"] == {"Accept": "application/json", "Accept-Encoding": "identity", "Cookie": ""} - stored = case.read() - assert stored is not None - assert stored.tracing_config is not None - assert ( - OpsTraceManager.decrypt_tracing_config(case.context.active_workspace_id, "phoenix", stored.tracing_config)[ - "api_key" - ] - == empty - ) - case.service.update(case.context, case.app_id, "phoenix", {**destination, "api_key": "second-key"}) - assert api_calls[-1][1]["headers"] != api_calls[0][1]["headers"] - assert len(api_calls) == 3 - - -@pytest.mark.parametrize("failure", ["denied", "wrong-path-html"]) -def test_failed_patch_preserves_policy_config_and_other_provider( - case: PhoenixCase, - api_calls: list[tuple[str, dict[str, object]]], - api_responses: deque[httpx.Response | Exception], - caplog: pytest.LogCaptureFixture, - failure: str, -) -> None: - original = {"endpoint": "https://host.test", "project": "project", "api_key": "private-secret"} - case.service.create(case.context, case.app_id, "phoenix", original) - before, other, policy = case.read(), case.read("langsmith"), case.reader.get_trace(case.context, case.app_id) - if failure == "denied": - api_responses.append(response({"secret": "private-vendor-response"}, 403)) - else: - api_responses.append( - httpx.Response( - 200, - text="private-vendor-response", - request=httpx.Request("GET", "https://host.test/not-a-collector/v1/projects"), - ) - ) - with pytest.raises(AppTracingConfigVerificationFailedError) as caught: - case.service.update( - case.context, - case.app_id, - "phoenix", - { - **original, - "endpoint": "https://host.test/not-a-collector", - "api_key": "******", - }, - activate=True, - ) - assert isinstance(caught.value.__cause__, PhoenixConfigurationError) - assert caught.value.__cause__.code == ("permission_denied" if failure == "denied" else "malformed_response") - assert case.read() == before - assert case.read("langsmith") == other - assert case.reader.get_trace(case.context, case.app_id) == policy - assert "private-vendor-response" not in caplog.text - assert "private-secret" not in caplog.text - assert len(api_calls) == 2 - - -@pytest.mark.parametrize("saved_url", [None, "https://stale.test/wrong-project"]) -def test_historical_get_calculates_url_without_verification_or_sdk( - case: PhoenixCase, - api_calls: list[tuple[str, dict[str, object]]], - saved_url: str | None, -) -> None: - stored = OpsTraceManager.encrypt_tracing_config( - case.context.active_workspace_id, - "phoenix", - { - "endpoint": "https://host.test/s/space/", - "project": "historical project", - "api_key": None, - }, - ) - if saved_url: - stored["project_url"] = saved_url - shown = case.gateway.present_config( - workspace_id=case.context.active_workspace_id, tracing_provider="phoenix", tracing_config=stored - ) - assert shown["project_url"] == "https://host.test/s/space/redirects/projects/historical%20project" - assert shown["api_key"] is None - case.gateway.validate_saved_config( - workspace_id=case.context.active_workspace_id, tracing_provider="phoenix", tracing_config=stored - ) - assert api_calls == [] - - -@pytest.mark.parametrize("endpoint", ["https://user:private@host.test?token=private", "https://host.test/v1/traces"]) -def test_unsafe_historical_destination_is_readable_without_link_or_io( - case: PhoenixCase, - api_calls: list[tuple[str, dict[str, object]]], - endpoint: str, -) -> None: - stored = OpsTraceManager.encrypt_tracing_config( - case.context.active_workspace_id, - "phoenix", - { - "endpoint": endpoint, - "project": "old-project", - "api_key": None, - }, - ) - shown = case.gateway.present_config( - workspace_id=case.context.active_workspace_id, tracing_provider="phoenix", tracing_config=stored - ) - assert shown["endpoint"] == endpoint - assert shown["project_url"] is None - assert api_calls == [] - - -@pytest.mark.parametrize("values", [{"endpoint": "https://host.test/v1/traces"}, {"project": "bad\nproject"}]) -def test_invalid_new_destination_has_no_io_or_storage_effect( - case: PhoenixCase, - api_calls: list[tuple[str, dict[str, object]]], - values: dict[str, str], -) -> None: - policy = case.reader.get_trace(case.context, case.app_id) - with pytest.raises(AppTracingConfigInvalidConfigurationError): - case.service.create(case.context, case.app_id, "phoenix", {"api_key": "synthetic-key", **values}, activate=True) - assert case.read() is None - assert case.reader.get_trace(case.context, case.app_id) == policy - assert api_calls == [] diff --git a/api/tests/unit_tests/services/test_plugin_file_upload_service.py b/api/tests/unit_tests/services/test_plugin_file_upload_service.py index b4d20230318..16b91caf4ed 100644 --- a/api/tests/unit_tests/services/test_plugin_file_upload_service.py +++ b/api/tests/unit_tests/services/test_plugin_file_upload_service.py @@ -21,7 +21,7 @@ from services.plugin_file_upload_service import ( PluginFileUploadService, PluginUploadUserFrom, ) -from tests.unit_tests.model_factories import make_account, make_end_user, make_tenant +from tests.helpers.unit.model_factories import make_account, make_end_user, make_tenant @pytest.fixture diff --git a/api/tests/unit_tests/services/test_rag_pipeline_task_proxy.py b/api/tests/unit_tests/services/test_rag_pipeline_task_proxy.py deleted file mode 100644 index eff006c8b09..00000000000 --- a/api/tests/unit_tests/services/test_rag_pipeline_task_proxy.py +++ /dev/null @@ -1,504 +0,0 @@ -import json -import logging -from datetime import datetime -from unittest.mock import MagicMock, Mock, call, patch - -import pytest - -from core.app.entities.rag_pipeline_invoke_entities import RagPipelineInvokeEntity -from core.rag.pipeline.queue import TenantIsolatedTaskQueue -from enums import CloudPlan, DeploymentEdition -from extensions.storage.storage_type import StorageType -from models.enums import CreatorUserRole -from models.model import UploadFile -from services.rag_pipeline.rag_pipeline_task_proxy import RagPipelineTaskProxy -from tests.unit_tests.config_override import config_overrides_context - - -class RagPipelineTaskProxyTestDataFactory: - """Factory class for creating test data and mock objects for RagPipelineTaskProxy tests.""" - - @staticmethod - def create_mock_features(plan: CloudPlan = CloudPlan.SANDBOX) -> Mock: - """Create mock features with billing configuration.""" - features = Mock() - features.billing = Mock() - features.billing.subscription = Mock() - features.billing.subscription.plan = plan - return features - - @staticmethod - def create_rag_pipeline_invoke_entity( - pipeline_id: str = "pipeline-123", - user_id: str = "user-456", - tenant_id: str = "tenant-789", - workflow_id: str = "workflow-101", - streaming: bool = True, - workflow_execution_id: str | None = None, - workflow_thread_pool_id: str | None = None, - ) -> RagPipelineInvokeEntity: - """Create RagPipelineInvokeEntity instance for testing.""" - return RagPipelineInvokeEntity( - pipeline_id=pipeline_id, - application_generate_entity={"key": "value"}, - user_id=user_id, - tenant_id=tenant_id, - workflow_id=workflow_id, - streaming=streaming, - workflow_execution_id=workflow_execution_id, - workflow_thread_pool_id=workflow_thread_pool_id, - ) - - @staticmethod - def create_rag_pipeline_task_proxy( - dataset_tenant_id: str = "tenant-123", - user_id: str = "user-456", - rag_pipeline_invoke_entities: list[RagPipelineInvokeEntity] | None = None, - ) -> RagPipelineTaskProxy: - """Create RagPipelineTaskProxy instance for testing.""" - if rag_pipeline_invoke_entities is None: - rag_pipeline_invoke_entities = [RagPipelineTaskProxyTestDataFactory.create_rag_pipeline_invoke_entity()] - return RagPipelineTaskProxy(dataset_tenant_id, user_id, rag_pipeline_invoke_entities) - - @staticmethod - def create_upload_file(file_id: str = "file-123") -> UploadFile: - """Create a transient upload file returned by the mocked storage service.""" - upload_file = UploadFile( - tenant_id="tenant-123", - storage_type=StorageType.LOCAL, - key="rag-pipeline.json", - name="rag-pipeline.json", - size=1, - extension="json", - mime_type="application/json", - created_by_role=CreatorUserRole.ACCOUNT, - created_by="user-456", - created_at=datetime(2025, 1, 1), - used=True, - ) - upload_file.id = file_id - return upload_file - - -class TestRagPipelineTaskProxy: - """Test cases for RagPipelineTaskProxy class.""" - - def test_initialization(self): - """Test RagPipelineTaskProxy initialization.""" - # Arrange - dataset_tenant_id = "tenant-123" - user_id = "user-456" - rag_pipeline_invoke_entities = [RagPipelineTaskProxyTestDataFactory.create_rag_pipeline_invoke_entity()] - - # Act - proxy = RagPipelineTaskProxy(dataset_tenant_id, user_id, rag_pipeline_invoke_entities) - - # Assert - assert proxy._dataset_tenant_id == dataset_tenant_id - assert proxy._user_id == user_id - assert proxy._rag_pipeline_invoke_entities == rag_pipeline_invoke_entities - assert isinstance(proxy._tenant_isolated_task_queue, TenantIsolatedTaskQueue) - assert proxy._tenant_isolated_task_queue._tenant_id == dataset_tenant_id - assert proxy._tenant_isolated_task_queue._unique_key == "pipeline" - - def test_initialization_with_empty_entities(self): - """Test initialization with empty rag_pipeline_invoke_entities.""" - # Arrange - dataset_tenant_id = "tenant-123" - user_id = "user-456" - rag_pipeline_invoke_entities = [] - - # Act - proxy = RagPipelineTaskProxy(dataset_tenant_id, user_id, rag_pipeline_invoke_entities) - - # Assert - assert proxy._dataset_tenant_id == dataset_tenant_id - assert proxy._user_id == user_id - assert proxy._rag_pipeline_invoke_entities == [] - - def test_initialization_with_multiple_entities(self): - """Test initialization with multiple rag_pipeline_invoke_entities.""" - # Arrange - dataset_tenant_id = "tenant-123" - user_id = "user-456" - rag_pipeline_invoke_entities = [ - RagPipelineTaskProxyTestDataFactory.create_rag_pipeline_invoke_entity(pipeline_id="pipeline-1"), - RagPipelineTaskProxyTestDataFactory.create_rag_pipeline_invoke_entity(pipeline_id="pipeline-2"), - RagPipelineTaskProxyTestDataFactory.create_rag_pipeline_invoke_entity(pipeline_id="pipeline-3"), - ] - - # Act - proxy = RagPipelineTaskProxy(dataset_tenant_id, user_id, rag_pipeline_invoke_entities) - - # Assert - assert len(proxy._rag_pipeline_invoke_entities) == 3 - assert proxy._rag_pipeline_invoke_entities[0].pipeline_id == "pipeline-1" - assert proxy._rag_pipeline_invoke_entities[1].pipeline_id == "pipeline-2" - assert proxy._rag_pipeline_invoke_entities[2].pipeline_id == "pipeline-3" - - @patch("services.rag_pipeline.rag_pipeline_task_proxy.FeatureService") - def test_features_property(self, mock_feature_service: MagicMock): - """Test cached_property features.""" - # Arrange - mock_features = RagPipelineTaskProxyTestDataFactory.create_mock_features() - mock_feature_service.get_features.return_value = mock_features - proxy = RagPipelineTaskProxyTestDataFactory.create_rag_pipeline_task_proxy() - - # Act - features1 = proxy.features - features2 = proxy.features # Second call should use cached property - - # Assert - assert features1 == mock_features - assert features2 == mock_features - assert features1 is features2 # Should be the same instance due to caching - mock_feature_service.get_features.assert_called_once_with("tenant-123", exclude_vector_space=True) - - @patch("services.rag_pipeline.rag_pipeline_task_proxy.FileService") - @patch("services.rag_pipeline.rag_pipeline_task_proxy.db") - def test_upload_invoke_entities(self, mock_db: MagicMock, mock_file_service_class: MagicMock): - """Test _upload_invoke_entities method.""" - # Arrange - proxy = RagPipelineTaskProxyTestDataFactory.create_rag_pipeline_task_proxy() - mock_file_service = Mock() - mock_file_service_class.return_value = mock_file_service - upload_file = RagPipelineTaskProxyTestDataFactory.create_upload_file("file-123") - mock_file_service.upload_text.return_value = upload_file - - # Act - result = proxy._upload_invoke_entities() - - # Assert - assert result == "file-123" - mock_file_service_class.assert_called_once_with(mock_db.engine) - - # Verify upload_text was called with correct parameters - mock_file_service.upload_text.assert_called_once() - call_args = mock_file_service.upload_text.call_args - json_text, name, user_id, tenant_id = call_args[0] - - assert name == "rag_pipeline_invoke_entities.json" - assert user_id == "user-456" - assert tenant_id == "tenant-123" - - # Verify JSON content - parsed_json = json.loads(json_text) - assert len(parsed_json) == 1 - assert parsed_json[0]["pipeline_id"] == "pipeline-123" - - @patch("services.rag_pipeline.rag_pipeline_task_proxy.FileService") - @patch("services.rag_pipeline.rag_pipeline_task_proxy.db") - def test_upload_invoke_entities_with_multiple_entities( - self, mock_db: MagicMock, mock_file_service_class: MagicMock - ): - """Test _upload_invoke_entities method with multiple entities.""" - # Arrange - entities = [ - RagPipelineTaskProxyTestDataFactory.create_rag_pipeline_invoke_entity(pipeline_id="pipeline-1"), - RagPipelineTaskProxyTestDataFactory.create_rag_pipeline_invoke_entity(pipeline_id="pipeline-2"), - ] - proxy = RagPipelineTaskProxy("tenant-123", "user-456", entities) - mock_file_service = Mock() - mock_file_service_class.return_value = mock_file_service - upload_file = RagPipelineTaskProxyTestDataFactory.create_upload_file("file-456") - mock_file_service.upload_text.return_value = upload_file - - # Act - result = proxy._upload_invoke_entities() - - # Assert - assert result == "file-456" - - # Verify JSON content contains both entities - call_args = mock_file_service.upload_text.call_args - json_text = call_args[0][0] - parsed_json = json.loads(json_text) - assert len(parsed_json) == 2 - assert parsed_json[0]["pipeline_id"] == "pipeline-1" - assert parsed_json[1]["pipeline_id"] == "pipeline-2" - - def test_send_to_direct_queue(self, tenant_queue_commands: MagicMock): - """Test _send_to_direct_queue method.""" - # Arrange - proxy = RagPipelineTaskProxyTestDataFactory.create_rag_pipeline_task_proxy() - upload_file_id = "file-123" - mock_task = Mock() - - # Act - proxy._send_to_direct_queue(upload_file_id, mock_task) - - # If sent to direct queue, tenant_isolated_task_queue should not be called - assert all(entry.args[0] != "LPUSH" for entry in tenant_queue_commands.call_args_list) - - # Celery should be called directly - mock_task.delay.assert_called_once_with( - rag_pipeline_invoke_entities_file_id=upload_file_id, tenant_id="tenant-123" - ) - - def test_send_to_tenant_queue_with_existing_task_key(self, tenant_queue_commands: MagicMock): - """Test _send_to_tenant_queue when task key exists.""" - # Arrange - proxy = RagPipelineTaskProxyTestDataFactory.create_rag_pipeline_task_proxy() - tenant_queue_commands.return_value = "task_key" - upload_file_id = "file-123" - mock_task = Mock() - - # Act - proxy._send_to_tenant_queue(upload_file_id, mock_task) - - # If task key exists, should push tasks to the queue - tenant_queue_commands.assert_any_call("LPUSH", "tenant_self_pipeline_task_queue:tenant-123", upload_file_id) - # Celery should not be called directly - mock_task.delay.assert_not_called() - - def test_send_to_tenant_queue_without_task_key(self, tenant_queue_commands: MagicMock): - """Test _send_to_tenant_queue when no task key exists.""" - # Arrange - proxy = RagPipelineTaskProxyTestDataFactory.create_rag_pipeline_task_proxy() - tenant_queue_commands.return_value = None - upload_file_id = "file-123" - mock_task = Mock() - - # Act - proxy._send_to_tenant_queue(upload_file_id, mock_task) - - # If no task key, should set task waiting time key first - tenant_queue_commands.assert_has_calls( - [ - call("GET", "tenant_pipeline_task:tenant-123", keys=["tenant_pipeline_task:tenant-123"]), - call("SETEX", "tenant_pipeline_task:tenant-123", 3600, 1), - ] - ) - mock_task.delay.assert_called_once_with( - rag_pipeline_invoke_entities_file_id=upload_file_id, tenant_id="tenant-123" - ) - - # The first task should be sent to celery directly, so push tasks should not be called - assert all(entry.args[0] != "LPUSH" for entry in tenant_queue_commands.call_args_list) - - def test_send_to_default_tenant_queue(self): - """Test _send_to_default_tenant_queue method.""" - # Arrange - proxy = RagPipelineTaskProxyTestDataFactory.create_rag_pipeline_task_proxy() - proxy._send_to_tenant_queue = Mock() - upload_file_id = "file-123" - - # Act - proxy._send_to_default_tenant_queue(upload_file_id) - - # Assert - proxy._send_to_tenant_queue.assert_called_once() - sent_file_id, task = proxy._send_to_tenant_queue.call_args.args - assert sent_file_id == upload_file_id - assert task.task == "tasks.rag_pipeline.rag_pipeline_run_task.rag_pipeline_run_task" - assert task.options["queue"] == "pipeline" - - def test_send_to_priority_tenant_queue(self): - """Test _send_to_priority_tenant_queue method.""" - # Arrange - proxy = RagPipelineTaskProxyTestDataFactory.create_rag_pipeline_task_proxy() - proxy._send_to_tenant_queue = Mock() - upload_file_id = "file-123" - - # Act - proxy._send_to_priority_tenant_queue(upload_file_id) - - # Assert - proxy._send_to_tenant_queue.assert_called_once() - sent_file_id, task = proxy._send_to_tenant_queue.call_args.args - assert sent_file_id == upload_file_id - assert task.task == "tasks.rag_pipeline.priority_rag_pipeline_run_task.priority_rag_pipeline_run_task" - assert task.options["queue"] == "priority_pipeline" - - def test_send_to_priority_direct_queue(self): - """Test _send_to_priority_direct_queue method.""" - # Arrange - proxy = RagPipelineTaskProxyTestDataFactory.create_rag_pipeline_task_proxy() - proxy._send_to_direct_queue = Mock() - upload_file_id = "file-123" - - # Act - proxy._send_to_priority_direct_queue(upload_file_id) - - # Assert - proxy._send_to_direct_queue.assert_called_once() - sent_file_id, task = proxy._send_to_direct_queue.call_args.args - assert sent_file_id == upload_file_id - assert task.task == "tasks.rag_pipeline.priority_rag_pipeline_run_task.priority_rag_pipeline_run_task" - assert task.options["queue"] == "priority_pipeline" - - @config_overrides_context(DEPLOYMENT_EDITION=DeploymentEdition.CLOUD) - @patch("services.rag_pipeline.rag_pipeline_task_proxy.FeatureService") - @patch("services.rag_pipeline.rag_pipeline_task_proxy.FileService") - @patch("services.rag_pipeline.rag_pipeline_task_proxy.db") - def test_dispatch_with_cloud_sandbox_plan( - self, mock_db: MagicMock, mock_file_service_class: MagicMock, mock_feature_service: MagicMock - ): - """Test _dispatch method in Cloud with Sandbox plan.""" - # Arrange - mock_features = RagPipelineTaskProxyTestDataFactory.create_mock_features(plan=CloudPlan.SANDBOX) - mock_feature_service.get_features.return_value = mock_features - proxy = RagPipelineTaskProxyTestDataFactory.create_rag_pipeline_task_proxy() - proxy._send_to_default_tenant_queue = Mock() - - mock_file_service = Mock() - mock_file_service_class.return_value = mock_file_service - mock_upload_file = RagPipelineTaskProxyTestDataFactory.create_upload_file("file-123") - mock_file_service.upload_text.return_value = mock_upload_file - - # Act - proxy._dispatch() - - # If billing is enabled with sandbox plan, should send to default tenant queue - proxy._send_to_default_tenant_queue.assert_called_once_with("file-123") - - @config_overrides_context(DEPLOYMENT_EDITION=DeploymentEdition.CLOUD) - @patch("services.rag_pipeline.rag_pipeline_task_proxy.FeatureService") - @patch("services.rag_pipeline.rag_pipeline_task_proxy.FileService") - @patch("services.rag_pipeline.rag_pipeline_task_proxy.db") - def test_dispatch_with_cloud_paid_plan(self, mock_db, mock_file_service_class, mock_feature_service): - """Test _dispatch method in Cloud with a paid plan.""" - # Arrange - mock_features = RagPipelineTaskProxyTestDataFactory.create_mock_features(plan=CloudPlan.TEAM) - mock_feature_service.get_features.return_value = mock_features - proxy = RagPipelineTaskProxyTestDataFactory.create_rag_pipeline_task_proxy() - proxy._send_to_priority_tenant_queue = Mock() - - mock_file_service = Mock() - mock_file_service_class.return_value = mock_file_service - mock_upload_file = RagPipelineTaskProxyTestDataFactory.create_upload_file("file-123") - mock_file_service.upload_text.return_value = mock_upload_file - - # Act - proxy._dispatch() - - # If billing is enabled with non-sandbox plan, should send to priority tenant queue - proxy._send_to_priority_tenant_queue.assert_called_once_with("file-123") - - @config_overrides_context(DEPLOYMENT_EDITION=DeploymentEdition.COMMUNITY) - @patch("services.rag_pipeline.rag_pipeline_task_proxy.FeatureService") - @patch("services.rag_pipeline.rag_pipeline_task_proxy.FileService") - @patch("services.rag_pipeline.rag_pipeline_task_proxy.db") - def test_dispatch_outside_cloud( - self, mock_db: MagicMock, mock_file_service_class: MagicMock, mock_feature_service: MagicMock - ): - """Test _dispatch method outside Cloud.""" - # Arrange - mock_features = RagPipelineTaskProxyTestDataFactory.create_mock_features() - mock_feature_service.get_features.return_value = mock_features - proxy = RagPipelineTaskProxyTestDataFactory.create_rag_pipeline_task_proxy() - proxy._send_to_priority_direct_queue = Mock() - - mock_file_service = Mock() - mock_file_service_class.return_value = mock_file_service - mock_upload_file = RagPipelineTaskProxyTestDataFactory.create_upload_file("file-123") - mock_file_service.upload_text.return_value = mock_upload_file - - # Act - proxy._dispatch() - - # If billing is disabled, for example: self-hosted or enterprise, should send to priority direct queue - proxy._send_to_priority_direct_queue.assert_called_once_with("file-123") - - @patch("services.rag_pipeline.rag_pipeline_task_proxy.FileService") - @patch("services.rag_pipeline.rag_pipeline_task_proxy.db") - def test_dispatch_with_empty_upload_file_id(self, mock_db: MagicMock, mock_file_service_class: MagicMock): - """Test _dispatch method when upload_file_id is empty.""" - # Arrange - proxy = RagPipelineTaskProxyTestDataFactory.create_rag_pipeline_task_proxy() - - mock_file_service = Mock() - mock_file_service_class.return_value = mock_file_service - upload_file = RagPipelineTaskProxyTestDataFactory.create_upload_file("") - mock_file_service.upload_text.return_value = upload_file - - # Act & Assert - with pytest.raises(ValueError, match="upload_file_id is empty"): - proxy._dispatch() - - @config_overrides_context(DEPLOYMENT_EDITION=DeploymentEdition.CLOUD) - @patch("services.rag_pipeline.rag_pipeline_task_proxy.FeatureService") - @patch("services.rag_pipeline.rag_pipeline_task_proxy.FileService") - @patch("services.rag_pipeline.rag_pipeline_task_proxy.db") - def test_dispatch_edge_case_empty_plan( - self, mock_db: MagicMock, mock_file_service_class: MagicMock, mock_feature_service: MagicMock - ): - """Test _dispatch method with empty plan string.""" - # Arrange - mock_features = RagPipelineTaskProxyTestDataFactory.create_mock_features(plan="") - mock_feature_service.get_features.return_value = mock_features - proxy = RagPipelineTaskProxyTestDataFactory.create_rag_pipeline_task_proxy() - proxy._send_to_priority_tenant_queue = Mock() - - mock_file_service = Mock() - mock_file_service_class.return_value = mock_file_service - mock_upload_file = RagPipelineTaskProxyTestDataFactory.create_upload_file("file-123") - mock_file_service.upload_text.return_value = mock_upload_file - - # Act - proxy._dispatch() - - # Assert - proxy._send_to_priority_tenant_queue.assert_called_once_with("file-123") - - @config_overrides_context(DEPLOYMENT_EDITION=DeploymentEdition.CLOUD) - @patch("services.rag_pipeline.rag_pipeline_task_proxy.FeatureService") - @patch("services.rag_pipeline.rag_pipeline_task_proxy.FileService") - @patch("services.rag_pipeline.rag_pipeline_task_proxy.db") - def test_dispatch_edge_case_none_plan( - self, mock_db: MagicMock, mock_file_service_class: MagicMock, mock_feature_service: MagicMock - ): - """Test _dispatch method with None plan.""" - # Arrange - mock_features = RagPipelineTaskProxyTestDataFactory.create_mock_features(plan=None) - mock_feature_service.get_features.return_value = mock_features - proxy = RagPipelineTaskProxyTestDataFactory.create_rag_pipeline_task_proxy() - proxy._send_to_priority_tenant_queue = Mock() - - mock_file_service = Mock() - mock_file_service_class.return_value = mock_file_service - mock_upload_file = RagPipelineTaskProxyTestDataFactory.create_upload_file("file-123") - mock_file_service.upload_text.return_value = mock_upload_file - - # Act - proxy._dispatch() - - # Assert - proxy._send_to_priority_tenant_queue.assert_called_once_with("file-123") - - @config_overrides_context(DEPLOYMENT_EDITION=DeploymentEdition.CLOUD) - @patch("services.rag_pipeline.rag_pipeline_task_proxy.FeatureService") - @patch("services.rag_pipeline.rag_pipeline_task_proxy.FileService") - @patch("services.rag_pipeline.rag_pipeline_task_proxy.db") - def test_delay_method( - self, mock_db: MagicMock, mock_file_service_class: MagicMock, mock_feature_service: MagicMock - ): - """Test delay method integration.""" - # Arrange - mock_features = RagPipelineTaskProxyTestDataFactory.create_mock_features(plan=CloudPlan.SANDBOX) - mock_feature_service.get_features.return_value = mock_features - proxy = RagPipelineTaskProxyTestDataFactory.create_rag_pipeline_task_proxy() - proxy._dispatch = Mock() - - mock_file_service = Mock() - mock_file_service_class.return_value = mock_file_service - mock_upload_file = RagPipelineTaskProxyTestDataFactory.create_upload_file("file-123") - mock_file_service.upload_text.return_value = mock_upload_file - - # Act - proxy.delay() - - # Assert - proxy._dispatch.assert_called_once() - - def test_delay_method_with_empty_entities(self, caplog: pytest.LogCaptureFixture): - """Test delay method with empty rag_pipeline_invoke_entities.""" - # Arrange - proxy = RagPipelineTaskProxy("tenant-123", "user-456", []) - - # Act - with caplog.at_level(logging.WARNING, logger="services.rag_pipeline.rag_pipeline_task_proxy"): - proxy.delay() - - # Assert - assert "Received empty rag pipeline invoke entities, no tasks delivered: tenant-123 user-456" in caplog.text diff --git a/api/tests/unit_tests/services/test_rbac_agent_access_service.py b/api/tests/unit_tests/services/test_rbac_agent_access_service.py index e9da8e7aa7a..7da1dad4997 100644 --- a/api/tests/unit_tests/services/test_rbac_agent_access_service.py +++ b/api/tests/unit_tests/services/test_rbac_agent_access_service.py @@ -5,7 +5,7 @@ from collections.abc import Callable import pytest from services import rbac_agent_access_service -from tests.unit_tests.rbac_fakes import Initialization, RBACDomain +from tests.helpers.unit.rbac_fakes import Initialization, RBACDomain def test_initialize_agent_rbac_access_seeds_scope_members_and_creator_policy( diff --git a/api/tests/unit_tests/services/test_recommended_app_catalog_gateway.py b/api/tests/unit_tests/services/test_recommended_app_catalog_gateway.py index 55308f0825b..e6ff933c328 100644 --- a/api/tests/unit_tests/services/test_recommended_app_catalog_gateway.py +++ b/api/tests/unit_tests/services/test_recommended_app_catalog_gateway.py @@ -22,7 +22,7 @@ from services.recommended_app_query_service import ( RecommendedAppInfoRecord, RecommendedAppRecord, ) -from tests.unit_tests.config_override import apply_config_overrides +from tests.helpers.unit.config_override import apply_config_overrides def _page_payload(*app_ids: str, learn_dify_ids: frozenset[str] = frozenset()) -> dict[str, object]: diff --git a/api/tests/unit_tests/services/test_remote_file_service.py b/api/tests/unit_tests/services/test_remote_file_service.py index 3273d53ccfe..0b380b3936b 100644 --- a/api/tests/unit_tests/services/test_remote_file_service.py +++ b/api/tests/unit_tests/services/test_remote_file_service.py @@ -26,7 +26,7 @@ from services.remote_file_service import ( RemoteFileUploadResult, RemoteFileUrlBlockedError, ) -from tests.unit_tests.model_factories import make_tenant +from tests.helpers.unit.model_factories import make_tenant REMOTE_URL = "https://example.com/files/report.pdf" diff --git a/api/tests/unit_tests/services/test_skill_management_service.py b/api/tests/unit_tests/services/test_skill_management_service.py index ce323a0ec01..452266a3644 100644 --- a/api/tests/unit_tests/services/test_skill_management_service.py +++ b/api/tests/unit_tests/services/test_skill_management_service.py @@ -65,7 +65,7 @@ from services.skill_management_service import ( validate_skill_description, validate_skill_name, ) -from tests.unit_tests.config_override import apply_config_overrides +from tests.helpers.unit.config_override import apply_config_overrides TENANT = "11111111-1111-1111-1111-111111111111" AGENT = "22222222-2222-2222-2222-222222222222" diff --git a/api/tests/unit_tests/services/test_snippet_dsl_service.py b/api/tests/unit_tests/services/test_snippet_dsl_service.py index 2471268fcdd..3d203cf08e9 100644 --- a/api/tests/unit_tests/services/test_snippet_dsl_service.py +++ b/api/tests/unit_tests/services/test_snippet_dsl_service.py @@ -25,7 +25,7 @@ from services.snippet_dsl_service import ( _check_version_compatibility, ) from services.snippet_service import SnippetService -from tests.unit_tests.model_factories import make_account, make_tenant, make_workflow +from tests.helpers.unit.model_factories import make_account, make_tenant, make_workflow SQLITE_MODELS = (CustomizedSnippet,) pytestmark = [ diff --git a/api/tests/unit_tests/services/test_snippet_generate_service.py b/api/tests/unit_tests/services/test_snippet_generate_service.py index 44517ec3de0..de105149333 100644 --- a/api/tests/unit_tests/services/test_snippet_generate_service.py +++ b/api/tests/unit_tests/services/test_snippet_generate_service.py @@ -20,7 +20,7 @@ from models.workflow import Workflow, WorkflowKind, WorkflowNodeExecutionModel from services.snippet_generate_service import SnippetGenerateService from services.snippet_service import SnippetService from services.workflow_service import WorkflowService -from tests.unit_tests.model_factories import make_account, make_workflow +from tests.helpers.unit.model_factories import make_account, make_workflow from tests.workflow_test_utils import build_test_graph_init_params diff --git a/api/tests/unit_tests/services/test_snippet_service.py b/api/tests/unit_tests/services/test_snippet_service.py index 6414a31f9b0..22c403f5aab 100644 --- a/api/tests/unit_tests/services/test_snippet_service.py +++ b/api/tests/unit_tests/services/test_snippet_service.py @@ -37,7 +37,7 @@ from models.workflow import ( from services.errors.app import IsDraftWorkflowError, WorkflowHashNotEqualError, WorkflowNotFoundError from services.errors.workflow_service import DraftWorkflowDeletionError, WorkflowInUseError from services.snippet_service import SnippetService -from tests.unit_tests.model_factories import make_account +from tests.helpers.unit.model_factories import make_account def _create_workflow(*, workflow_id: str, version: str, graph: dict, features: dict) -> Workflow: @@ -945,7 +945,7 @@ def test_delete_draft_variable_files_removes_storage_objects( def test_delete_archived_workflow_run_files_removes_prefixed_objects(monkeypatch: pytest.MonkeyPatch) -> None: - from tests.unit_tests.config_override import apply_config_overrides + from tests.helpers.unit.config_override import apply_config_overrides snippet = _snippet() archive_storage = SimpleNamespace( diff --git a/api/tests/unit_tests/services/test_step_by_step_tour_service.py b/api/tests/unit_tests/services/test_step_by_step_tour_service.py index 44d6a8f37fc..a5b11b487b7 100644 --- a/api/tests/unit_tests/services/test_step_by_step_tour_service.py +++ b/api/tests/unit_tests/services/test_step_by_step_tour_service.py @@ -12,7 +12,7 @@ from repositories.account.repository import SQLAlchemyAccountRepository from services.entities.account_entities import AccountSnapshot from services.entities.onboarding_entities import StepByStepTourPatch, StepByStepTourResult, StepByStepTourState from services.step_by_step_tour_service import StepByStepTourService -from tests.unit_tests.model_factories import make_account, make_account_snapshot +from tests.helpers.unit.model_factories import make_account, make_account_snapshot def _context(*, workspace_id: str = "workspace-1") -> RequestContext: diff --git a/api/tests/unit_tests/services/test_system_feature_service.py b/api/tests/unit_tests/services/test_system_feature_service.py new file mode 100644 index 00000000000..2ddd731db6e --- /dev/null +++ b/api/tests/unit_tests/services/test_system_feature_service.py @@ -0,0 +1,667 @@ +import logging +from collections.abc import Callable +from unittest.mock import MagicMock + +import pytest +from pydantic import ValidationError + +from enums import DeploymentEdition +from services import system_feature_service as feature_service_module +from services.entities.feature_entities import ( + LicenseModel, + LicenseStatus, + PluginInstallationScope, + SSOProtocol, + SystemFeatureModel, +) +from services.system_feature_service import SystemFeatureService + +"""Tests for the public SystemFeatureService app DSL contract.""" + + +def test_get_system_features_excludes_app_dsl_version() -> None: + result = SystemFeatureService.get_public_system_features().model_dump() + + assert "app_dsl_version" not in result + + +@pytest.mark.parametrize("enabled", [False, True]) +def test_get_system_features_reads_enable_change_email( + config_overrides: Callable[..., None], + enabled: bool, +) -> None: + config_overrides(ENABLE_CHANGE_EMAIL=enabled) + + result = SystemFeatureService.get_public_system_features() + + assert result.enable_change_email is enabled + + +@pytest.mark.parametrize( + ("deployment_edition", "configured", "expected"), + [ + (DeploymentEdition.COMMUNITY, False, False), + (DeploymentEdition.COMMUNITY, True, True), + (DeploymentEdition.ENTERPRISE, True, False), + ], +) +def test_change_email_policy( + config_overrides: Callable[..., None], + deployment_edition: DeploymentEdition, + configured: bool, + expected: bool, +) -> None: + config_overrides(DEPLOYMENT_EDITION=deployment_edition, ENABLE_CHANGE_EMAIL=configured) + + assert SystemFeatureService.is_change_email_enabled() is expected + + +"""Tests for SystemFeatureService deployment-edition behavior.""" + + +def test_system_feature_model_requires_deployment_edition() -> None: + with pytest.raises(ValidationError): + SystemFeatureModel.model_validate({"agent_tracing_api_collection_enabled": False}) + + +@pytest.mark.parametrize("edition", list(DeploymentEdition)) +@pytest.mark.parametrize("enabled", [False, True]) +def test_public_agent_tracing_gate_uses_configured_value_in_every_edition( + monkeypatch: pytest.MonkeyPatch, + config_overrides: Callable[..., None], + edition: DeploymentEdition, + enabled: bool, +) -> None: + config_overrides(DEPLOYMENT_EDITION=edition, OPS_AGENT_TRACE_ENABLED=enabled) + monkeypatch.setattr( + "services.system_feature_service.EnterpriseService.get_info", + staticmethod(lambda: {}), + ) + + result = SystemFeatureService.get_public_system_features() + + assert result.agent_tracing_api_collection_enabled is enabled + assert result.model_dump(mode="json")["agent_tracing_api_collection_enabled"] is enabled + + +@pytest.mark.parametrize( + "edition", + [ + DeploymentEdition.COMMUNITY, + DeploymentEdition.ENTERPRISE, + DeploymentEdition.CLOUD, + ], +) +def test_get_system_features_uses_configured_deployment_edition( + monkeypatch: pytest.MonkeyPatch, + config_overrides: Callable[..., None], + edition: DeploymentEdition, +) -> None: + fulfill_from_enterprise = MagicMock() + config_overrides(DEPLOYMENT_EDITION=edition) + monkeypatch.setattr( + "services.system_feature_service.SystemFeatureService._fulfill_params_from_enterprise", + fulfill_from_enterprise, + ) + + result = SystemFeatureService.get_public_system_features() + + assert result.deployment_edition is edition + assert result.model_dump(mode="json")["deployment_edition"] == edition.value + webapp_auth_enabled = edition is DeploymentEdition.ENTERPRISE + assert SystemFeatureService.is_webapp_auth_enabled() is webapp_auth_enabled + assert result.webapp_auth.enabled is webapp_auth_enabled + if edition is DeploymentEdition.ENTERPRISE: + fulfill_from_enterprise.assert_called_once_with(result) + else: + fulfill_from_enterprise.assert_not_called() + + +@pytest.mark.parametrize( + ("edition", "feature_enabled", "expected"), + [ + (DeploymentEdition.CLOUD, True, True), + (DeploymentEdition.CLOUD, False, False), + (DeploymentEdition.COMMUNITY, True, False), + (DeploymentEdition.ENTERPRISE, True, False), + ], +) +def test_trial_app_policy_is_cloud_only( + config_overrides: Callable[..., None], + edition: DeploymentEdition, + feature_enabled: bool, + expected: bool, +) -> None: + config_overrides(DEPLOYMENT_EDITION=edition, ENABLE_TRIAL_APP=feature_enabled) + + assert SystemFeatureService.is_trial_app_enabled() is expected + + +"""Tests for the SystemFeatureService app-deployment policy.""" + + +@pytest.mark.parametrize( + ("enterprise_info", "initial", "expected"), + [ + # Enterprise reports the feature on -> mirrored through. + ({"EnableAppDeploy": True}, False, True), + # Enterprise may turn it off; the read runs after the hardcoded default + # and overrides it (forward-compat with a future entitlement gate). + ({"EnableAppDeploy": False}, True, False), + # Old enterprise without the key -> the existing value is left untouched. + ({}, True, True), + ], + ids=["enabled", "override_off", "missing_keeps_default"], +) +def test_fulfill_params_from_enterprise_enable_app_deploy( + monkeypatch: pytest.MonkeyPatch, + enterprise_info: dict[str, object], + initial: bool, + expected: bool, +) -> None: + monkeypatch.setattr( + feature_service_module.EnterpriseService, + "get_info", + staticmethod(lambda: enterprise_info), + ) + + features = SystemFeatureModel( + agent_tracing_api_collection_enabled=False, deployment_edition=DeploymentEdition.COMMUNITY + ) + features.enable_app_deploy = initial + + SystemFeatureService._fulfill_params_from_enterprise(features) + + assert features.enable_app_deploy is expected + + +"""Tests for the SystemFeatureService explore-banner policy.""" + + +@pytest.mark.parametrize( + ("edition", "configured", "expected"), + [ + (DeploymentEdition.CLOUD, True, True), + (DeploymentEdition.CLOUD, False, False), + (DeploymentEdition.COMMUNITY, True, False), + (DeploymentEdition.ENTERPRISE, True, False), + ], +) +def test_get_system_features_enables_explore_banner_only_for_cloud( + monkeypatch: pytest.MonkeyPatch, + config_overrides: Callable[..., None], + edition: DeploymentEdition, + configured: bool, + expected: bool, +) -> None: + config_overrides(DEPLOYMENT_EDITION=edition, ENABLE_EXPLORE_BANNER=configured) + monkeypatch.setattr(SystemFeatureService, "_fulfill_params_from_enterprise", lambda *_: None) + + result = SystemFeatureService.get_public_system_features() + + assert SystemFeatureService.is_explore_banner_enabled() is expected + assert result.enable_explore_banner is expected + + +def test_workspace_creation_uses_environment_policy( + monkeypatch: pytest.MonkeyPatch, config_overrides: Callable[..., None] +) -> None: + config_overrides(DEPLOYMENT_EDITION=DeploymentEdition.COMMUNITY, ALLOW_CREATE_WORKSPACE=True) + monkeypatch.setattr( + "services.system_feature_service.EnterpriseService.get_info", + lambda: (_ for _ in ()).throw(AssertionError("enterprise API should not be called")), + ) + + assert SystemFeatureService.is_workspace_creation_allowed() is True + + +def test_workspace_creation_uses_enterprise_policy( + monkeypatch: pytest.MonkeyPatch, config_overrides: Callable[..., None] +) -> None: + config_overrides(DEPLOYMENT_EDITION=DeploymentEdition.ENTERPRISE) + monkeypatch.setattr( + "services.system_feature_service.EnterpriseService.get_info", + lambda: {"IsAllowCreateWorkspace": False}, + ) + + assert SystemFeatureService.is_workspace_creation_allowed() is False + + +def test_workspace_creation_keeps_environment_policy_when_enterprise_value_is_missing( + monkeypatch: pytest.MonkeyPatch, + config_overrides: Callable[..., None], +) -> None: + config_overrides(DEPLOYMENT_EDITION=DeploymentEdition.ENTERPRISE, ALLOW_CREATE_WORKSPACE=True) + monkeypatch.setattr("services.system_feature_service.EnterpriseService.get_info", lambda: {}) + + assert SystemFeatureService.is_workspace_creation_allowed() is True + + +def test_plugin_manager_is_enabled_only_for_enterprise(config_overrides: Callable[..., None]) -> None: + config_overrides(DEPLOYMENT_EDITION=DeploymentEdition.ENTERPRISE) + assert SystemFeatureService.is_plugin_manager_enabled() is True + + config_overrides(DEPLOYMENT_EDITION=DeploymentEdition.COMMUNITY) + assert SystemFeatureService.is_plugin_manager_enabled() is False + + +def test_webapp_auth_enabled_does_not_query_enterprise_info( + monkeypatch: pytest.MonkeyPatch, config_overrides: Callable[..., None] +) -> None: + config_overrides(DEPLOYMENT_EDITION=DeploymentEdition.ENTERPRISE) + monkeypatch.setattr( + "services.system_feature_service.EnterpriseService.get_info", + lambda: (_ for _ in ()).throw(AssertionError("enterprise info should not be queried")), + ) + + assert SystemFeatureService.is_webapp_auth_enabled() is True + + +def test_registration_policy_uses_enterprise_override( + monkeypatch: pytest.MonkeyPatch, config_overrides: Callable[..., None] +) -> None: + config_overrides(DEPLOYMENT_EDITION=DeploymentEdition.ENTERPRISE, ALLOW_REGISTER=True) + monkeypatch.setattr( + "services.system_feature_service.EnterpriseService.get_info", + lambda: {"IsAllowRegister": False}, + ) + + assert SystemFeatureService.is_registration_allowed() is False + + +def test_password_login_policy_uses_enterprise_override( + monkeypatch: pytest.MonkeyPatch, config_overrides: Callable[..., None] +) -> None: + config_overrides(DEPLOYMENT_EDITION=DeploymentEdition.ENTERPRISE, ENABLE_EMAIL_PASSWORD_LOGIN=True) + monkeypatch.setattr( + "services.system_feature_service.EnterpriseService.get_info", + lambda: {"EnableEmailPasswordLogin": False}, + ) + + assert SystemFeatureService.is_email_password_login_enabled() is False + + +def test_branding_reads_enterprise_configuration( + monkeypatch: pytest.MonkeyPatch, config_overrides: Callable[..., None] +) -> None: + config_overrides(DEPLOYMENT_EDITION=DeploymentEdition.ENTERPRISE) + monkeypatch.setattr( + "services.system_feature_service.EnterpriseService.get_info", + lambda: { + "Branding": { + "applicationTitle": "Enterprise Dify", + "loginPageLogo": "login-logo", + "workspaceLogo": "workspace-logo", + "favicon": "favicon", + } + }, + ) + + branding = SystemFeatureService.get_branding() + + assert branding.enabled is True + assert branding.application_title == "Enterprise Dify" + assert branding.login_page_logo == "login-logo" + assert branding.workspace_logo == "workspace-logo" + assert branding.favicon == "favicon" + + +def test_license_status_ignores_unrelated_malformed_quota_fields( + monkeypatch: pytest.MonkeyPatch, config_overrides: Callable[..., None] +) -> None: + config_overrides(DEPLOYMENT_EDITION=DeploymentEdition.ENTERPRISE) + monkeypatch.setattr( + "services.system_feature_service.EnterpriseService.get_info", + lambda: { + "License": { + "status": "active", + "workspaces": {"enabled": True, "limit": 3, "used": {"unexpected": "shape"}}, + "licensedSeats": "unexpected-shape", + } + }, + ) + + assert SystemFeatureService.get_license_status() == LicenseStatus.ACTIVE + + +"""Tests for SystemFeatureService learn-app and tour policies.""" + + +def test_system_feature_model_defaults_enable_learn_app() -> None: + system_features = SystemFeatureModel( + agent_tracing_api_collection_enabled=False, deployment_edition=DeploymentEdition.COMMUNITY + ) + + assert system_features.enable_learn_app is True + assert system_features.enable_step_by_step_tour is False + + +@pytest.mark.parametrize("enabled", [True, False]) +def test_get_system_features_reads_enable_learn_app(config_overrides: Callable[..., None], enabled: bool) -> None: + config_overrides(ENABLE_LEARN_APP=enabled) + + result = SystemFeatureService.get_public_system_features() + + assert result.enable_learn_app is enabled + + +@pytest.mark.parametrize("enabled", [True, False]) +def test_get_system_features_reads_enable_step_by_step_tour( + config_overrides: Callable[..., None], enabled: bool +) -> None: + config_overrides(ENABLE_STEP_BY_STEP_TOUR=enabled) + + result = SystemFeatureService.get_public_system_features() + + assert result.enable_step_by_step_tour is enabled + + +"""Tests for the SystemFeatureService license-expiry policy.""" + + +_SYSTEM_FEATURE_SERVICE_LICENSE_EXPIRY_NOTICE_ENTERPRISE_INFO = { + "License": {"status": LicenseStatus.EXPIRING, "expiredAt": "2026-12-31"} +} + + +def test_license_model_defaults_license_expiry_notice_disabled() -> None: + """Without a license there is no expiry to announce, so the notice is off unless enabled explicitly.""" + assert LicenseModel().license_expiry_notice_enabled is False + + +@pytest.mark.parametrize("enabled", [True, False]) +def test_get_license_non_enterprise_ignores_expiry_notice_config( + config_overrides: Callable[..., None], enabled: bool +) -> None: + """Non-enterprise deployments have no license, so the env toggle never turns the notice on.""" + config_overrides( + ENABLE_LICENSE_EXPIRY_NOTICE=enabled, + DEPLOYMENT_EDITION=DeploymentEdition.COMMUNITY, + ) + + result = SystemFeatureService.get_license() + + assert result.license_expiry_notice_enabled is False + + +@pytest.mark.parametrize("enabled", [True, False]) +def test_get_license_enterprise_reads_license_expiry_notice_enabled( + monkeypatch: pytest.MonkeyPatch, config_overrides: Callable[..., None], enabled: bool +) -> None: + """The enterprise-sourced license carries the env-resolved notice flag alongside its real status.""" + config_overrides( + ENABLE_LICENSE_EXPIRY_NOTICE=enabled, + DEPLOYMENT_EDITION=DeploymentEdition.ENTERPRISE, + ) + monkeypatch.setattr( + feature_service_module.EnterpriseService, + "get_info", + staticmethod(lambda: _SYSTEM_FEATURE_SERVICE_LICENSE_EXPIRY_NOTICE_ENTERPRISE_INFO), + ) + + result = SystemFeatureService.get_license() + + assert result.status == LicenseStatus.EXPIRING + assert result.expired_at == "2026-12-31" + assert result.license_expiry_notice_enabled is enabled + + +"""Tests for SystemFeatureService licensed-seat parsing.""" + + +_SYSTEM_FEATURE_SERVICE_LICENSED_SEATS_ENTERPRISE_INFO = { + "License": {"licensedSeats": {"enabled": True, "limit": 3, "used": 1}} +} + + +def test_get_license_parses_licensed_seats( + monkeypatch: pytest.MonkeyPatch, config_overrides: Callable[..., None] +) -> None: + """The authenticated license accessor copies the licensed-seat quota out of the enterprise payload.""" + config_overrides(DEPLOYMENT_EDITION=DeploymentEdition.ENTERPRISE) + monkeypatch.setattr( + feature_service_module.EnterpriseService, + "get_info", + staticmethod(lambda: _SYSTEM_FEATURE_SERVICE_LICENSED_SEATS_ENTERPRISE_INFO), + ) + + license_model = SystemFeatureService.get_license() + + assert isinstance(license_model, LicenseModel) + assert license_model.seats.enabled is True + assert license_model.seats.limit == 3 + assert license_model.seats.size == 1 + + +def test_get_license_non_enterprise_is_unconstrained(config_overrides: Callable[..., None]) -> None: + """Non-enterprise deployments have no license; seat allocation is unconstrained.""" + config_overrides(DEPLOYMENT_EDITION=DeploymentEdition.COMMUNITY) + + license_model = SystemFeatureService.get_license() + + assert license_model.status == LicenseStatus.NONE + assert license_model.seats.enabled is False + assert license_model.seats.is_available() is True + + +"""Tests for SystemFeatureService plugin-installation policy.""" + + +def test_get_plugin_installation_permission_defaults_to_all_for_non_enterprise( + config_overrides: Callable[..., None], +) -> None: + config_overrides(DEPLOYMENT_EDITION=DeploymentEdition.COMMUNITY) + + permission = SystemFeatureService.get_plugin_installation_permission() + + assert permission.plugin_installation_scope is PluginInstallationScope.ALL + assert permission.restrict_to_marketplace_only is False + + +def test_get_plugin_installation_permission_parses_enterprise_policy( + monkeypatch: pytest.MonkeyPatch, + config_overrides: Callable[..., None], +) -> None: + config_overrides(DEPLOYMENT_EDITION=DeploymentEdition.ENTERPRISE) + monkeypatch.setattr( + feature_service_module.EnterpriseService, + "get_info", + staticmethod( + lambda: { + "PluginInstallationPermission": { + "pluginInstallationScope": "official_only", + "restrictToMarketplaceOnly": True, + } + } + ), + ) + + permission = SystemFeatureService.get_plugin_installation_permission() + + assert permission.plugin_installation_scope is PluginInstallationScope.OFFICIAL_ONLY + assert permission.restrict_to_marketplace_only is True + + +@pytest.mark.parametrize( + "invalid_permission", + [ + { + "pluginInstallationScope": "unknown-scope", + "restrictToMarketplaceOnly": False, + }, + { + "pluginInstallationScope": "all", + "restrictToMarketplaceOnly": "false", + }, + ], + ids=["unknown_scope", "non_boolean_marketplace_restriction"], +) +def test_invalid_enterprise_policy_denies_all_plugin_installations( + caplog: pytest.LogCaptureFixture, + invalid_permission: dict[str, object], +) -> None: + with caplog.at_level(logging.ERROR, logger="services.system_feature_service"): + permission = SystemFeatureService._resolve_plugin_installation_permission( + {"PluginInstallationPermission": invalid_permission} + ) + + assert permission.plugin_installation_scope is PluginInstallationScope.NONE + assert permission.restrict_to_marketplace_only is True + assert "denying all plugin installations" in caplog.text + + +def test_system_features_exposes_only_validated_plugin_installation_policy( + monkeypatch: pytest.MonkeyPatch, +) -> None: + monkeypatch.setattr( + feature_service_module.EnterpriseService, + "get_info", + staticmethod( + lambda: { + "PluginInstallationPermission": { + "pluginInstallationScope": "unknown-scope", + "restrictToMarketplaceOnly": False, + } + } + ), + ) + features = SystemFeatureModel( + agent_tracing_api_collection_enabled=False, deployment_edition=DeploymentEdition.ENTERPRISE + ) + + SystemFeatureService._fulfill_params_from_enterprise(features) + + assert features.plugin_installation_permission.plugin_installation_scope is PluginInstallationScope.NONE + assert features.plugin_installation_permission.restrict_to_marketplace_only is True + + +"""Tests for SystemFeatureService SSO protocol parsing.""" + + +def test_system_features_exposes_valid_enterprise_sso_protocols( + monkeypatch: pytest.MonkeyPatch, +) -> None: + monkeypatch.setattr( + feature_service_module.EnterpriseService, + "get_info", + staticmethod( + lambda: { + "SSOEnforcedForSigninProtocol": "saml", + "WebAppAuth": {}, + "SSOEnforcedForWebProtocol": "oidc", + } + ), + ) + features = SystemFeatureModel( + agent_tracing_api_collection_enabled=False, deployment_edition=DeploymentEdition.ENTERPRISE + ) + + SystemFeatureService._fulfill_params_from_enterprise(features) + + assert features.sso_enforced_for_signin_protocol is SSOProtocol.SAML + assert features.webapp_auth.sso_config.protocol is SSOProtocol.OIDC + + +@pytest.mark.parametrize("empty_protocol", [None, "", " "]) +def test_system_features_normalizes_empty_enterprise_sso_protocols_to_none( + monkeypatch: pytest.MonkeyPatch, + empty_protocol: object, +) -> None: + monkeypatch.setattr( + feature_service_module.EnterpriseService, + "get_info", + staticmethod( + lambda: { + "SSOEnforcedForSigninProtocol": empty_protocol, + "WebAppAuth": {}, + "SSOEnforcedForWebProtocol": empty_protocol, + } + ), + ) + features = SystemFeatureModel( + agent_tracing_api_collection_enabled=False, deployment_edition=DeploymentEdition.ENTERPRISE + ) + + SystemFeatureService._fulfill_params_from_enterprise(features) + + assert features.sso_enforced_for_signin_protocol is None + assert features.webapp_auth.sso_config.protocol is None + + +@pytest.mark.parametrize("invalid_protocol", ["unknown", 42]) +def test_system_features_rejects_invalid_enterprise_sso_protocols( + caplog: pytest.LogCaptureFixture, + monkeypatch: pytest.MonkeyPatch, + invalid_protocol: object, +) -> None: + monkeypatch.setattr( + feature_service_module.EnterpriseService, + "get_info", + staticmethod( + lambda: { + "SSOEnforcedForSigninProtocol": invalid_protocol, + "WebAppAuth": {}, + "SSOEnforcedForWebProtocol": invalid_protocol, + } + ), + ) + features = SystemFeatureModel( + agent_tracing_api_collection_enabled=False, deployment_edition=DeploymentEdition.ENTERPRISE + ) + + with caplog.at_level(logging.ERROR, logger="services.system_feature_service"): + SystemFeatureService._fulfill_params_from_enterprise(features) + + assert features.sso_enforced_for_signin_protocol is None + assert features.webapp_auth.sso_config.protocol is None + assert caplog.text.count("Invalid Enterprise SSO protocol") == 2 + + +def test_system_features_defaults_sso_protocols_to_none() -> None: + features = SystemFeatureModel( + agent_tracing_api_collection_enabled=False, deployment_edition=DeploymentEdition.COMMUNITY + ) + + assert features.sso_enforced_for_signin_protocol is None + assert features.webapp_auth.sso_config.protocol is None + + +"""Tests for the public SystemFeatureService hosted-model contract.""" + + +def test_get_system_features_excludes_trial_models() -> None: + result = SystemFeatureService.get_public_system_features().model_dump() + + assert "trial_models" not in result + + +"""Tests for the SystemFeatureService WebApp public-access policy.""" + + +@pytest.mark.parametrize( + ("env_value", "expected"), + [ + (False, False), + (True, True), + ], + ids=["disabled_by_env", "enabled_by_env"], +) +def test_fulfill_system_params_from_env_sets_allow_public_access( + config_overrides: Callable[..., None], + env_value: bool, + expected: bool, +) -> None: + config_overrides(WEBAPP_PUBLIC_ACCESS_ENABLED=env_value) + + system_features = SystemFeatureModel( + agent_tracing_api_collection_enabled=False, deployment_edition=DeploymentEdition.COMMUNITY + ) + SystemFeatureService._fulfill_system_params_from_env(system_features) + + assert system_features.webapp_auth.allow_public_access is expected + + +def test_get_system_features_defaults_allow_public_access_to_true() -> None: + system_features = SystemFeatureService.get_public_system_features() + + assert system_features.webapp_auth.allow_public_access is True diff --git a/api/tests/unit_tests/services/test_system_feature_service_app_dsl_version.py b/api/tests/unit_tests/services/test_system_feature_service_app_dsl_version.py deleted file mode 100644 index 0fca8efa601..00000000000 --- a/api/tests/unit_tests/services/test_system_feature_service_app_dsl_version.py +++ /dev/null @@ -1,9 +0,0 @@ -"""Tests for the public SystemFeatureService app DSL contract.""" - -from services.system_feature_service import SystemFeatureService - - -def test_get_system_features_excludes_app_dsl_version() -> None: - result = SystemFeatureService.get_public_system_features().model_dump() - - assert "app_dsl_version" not in result diff --git a/api/tests/unit_tests/services/test_system_feature_service_change_email.py b/api/tests/unit_tests/services/test_system_feature_service_change_email.py deleted file mode 100644 index d4cc1666c47..00000000000 --- a/api/tests/unit_tests/services/test_system_feature_service_change_email.py +++ /dev/null @@ -1,37 +0,0 @@ -from collections.abc import Callable - -import pytest - -from enums import DeploymentEdition -from services.system_feature_service import SystemFeatureService - - -@pytest.mark.parametrize("enabled", [False, True]) -def test_get_system_features_reads_enable_change_email( - config_overrides: Callable[..., None], - enabled: bool, -) -> None: - config_overrides(ENABLE_CHANGE_EMAIL=enabled) - - result = SystemFeatureService.get_public_system_features() - - assert result.enable_change_email is enabled - - -@pytest.mark.parametrize( - ("deployment_edition", "configured", "expected"), - [ - (DeploymentEdition.COMMUNITY, False, False), - (DeploymentEdition.COMMUNITY, True, True), - (DeploymentEdition.ENTERPRISE, True, False), - ], -) -def test_change_email_policy( - config_overrides: Callable[..., None], - deployment_edition: DeploymentEdition, - configured: bool, - expected: bool, -) -> None: - config_overrides(DEPLOYMENT_EDITION=deployment_edition, ENABLE_CHANGE_EMAIL=configured) - - assert SystemFeatureService.is_change_email_enabled() is expected diff --git a/api/tests/unit_tests/services/test_system_feature_service_deployment_edition.py b/api/tests/unit_tests/services/test_system_feature_service_deployment_edition.py deleted file mode 100644 index c8ee83d8668..00000000000 --- a/api/tests/unit_tests/services/test_system_feature_service_deployment_edition.py +++ /dev/null @@ -1,89 +0,0 @@ -"""Tests for SystemFeatureService deployment-edition behavior.""" - -from collections.abc import Callable -from unittest.mock import MagicMock - -import pytest -from pydantic import ValidationError - -from enums import DeploymentEdition -from services.entities.feature_entities import SystemFeatureModel -from services.system_feature_service import SystemFeatureService - - -def test_system_feature_model_requires_deployment_edition() -> None: - with pytest.raises(ValidationError): - SystemFeatureModel.model_validate({"agent_tracing_api_collection_enabled": False}) - - -@pytest.mark.parametrize("edition", list(DeploymentEdition)) -@pytest.mark.parametrize("enabled", [False, True]) -def test_public_agent_tracing_gate_uses_configured_value_in_every_edition( - monkeypatch: pytest.MonkeyPatch, - config_overrides: Callable[..., None], - edition: DeploymentEdition, - enabled: bool, -) -> None: - config_overrides(DEPLOYMENT_EDITION=edition, OPS_AGENT_TRACE_ENABLED=enabled) - monkeypatch.setattr( - "services.system_feature_service.EnterpriseService.get_info", - staticmethod(lambda: {}), - ) - - result = SystemFeatureService.get_public_system_features() - - assert result.agent_tracing_api_collection_enabled is enabled - assert result.model_dump(mode="json")["agent_tracing_api_collection_enabled"] is enabled - - -@pytest.mark.parametrize( - "edition", - [ - DeploymentEdition.COMMUNITY, - DeploymentEdition.ENTERPRISE, - DeploymentEdition.CLOUD, - ], -) -def test_get_system_features_uses_configured_deployment_edition( - monkeypatch: pytest.MonkeyPatch, - config_overrides: Callable[..., None], - edition: DeploymentEdition, -) -> None: - fulfill_from_enterprise = MagicMock() - config_overrides(DEPLOYMENT_EDITION=edition) - monkeypatch.setattr( - "services.system_feature_service.SystemFeatureService._fulfill_params_from_enterprise", - fulfill_from_enterprise, - ) - - result = SystemFeatureService.get_public_system_features() - - assert result.deployment_edition is edition - assert result.model_dump(mode="json")["deployment_edition"] == edition.value - webapp_auth_enabled = edition is DeploymentEdition.ENTERPRISE - assert SystemFeatureService.is_webapp_auth_enabled() is webapp_auth_enabled - assert result.webapp_auth.enabled is webapp_auth_enabled - if edition is DeploymentEdition.ENTERPRISE: - fulfill_from_enterprise.assert_called_once_with(result) - else: - fulfill_from_enterprise.assert_not_called() - - -@pytest.mark.parametrize( - ("edition", "feature_enabled", "expected"), - [ - (DeploymentEdition.CLOUD, True, True), - (DeploymentEdition.CLOUD, False, False), - (DeploymentEdition.COMMUNITY, True, False), - (DeploymentEdition.ENTERPRISE, True, False), - ], -) -def test_trial_app_policy_is_cloud_only( - config_overrides: Callable[..., None], - edition: DeploymentEdition, - feature_enabled: bool, - expected: bool, -) -> None: - config_overrides(DEPLOYMENT_EDITION=edition, ENABLE_TRIAL_APP=feature_enabled) - - assert SystemFeatureService.is_trial_app_enabled() is expected diff --git a/api/tests/unit_tests/services/test_system_feature_service_enable_app_deploy.py b/api/tests/unit_tests/services/test_system_feature_service_enable_app_deploy.py deleted file mode 100644 index 8fbf8cc674e..00000000000 --- a/api/tests/unit_tests/services/test_system_feature_service_enable_app_deploy.py +++ /dev/null @@ -1,43 +0,0 @@ -"""Tests for the SystemFeatureService app-deployment policy.""" - -import pytest - -from enums import DeploymentEdition -from services import system_feature_service as feature_service_module -from services.entities.feature_entities import SystemFeatureModel -from services.system_feature_service import SystemFeatureService - - -@pytest.mark.parametrize( - ("enterprise_info", "initial", "expected"), - [ - # Enterprise reports the feature on -> mirrored through. - ({"EnableAppDeploy": True}, False, True), - # Enterprise may turn it off; the read runs after the hardcoded default - # and overrides it (forward-compat with a future entitlement gate). - ({"EnableAppDeploy": False}, True, False), - # Old enterprise without the key -> the existing value is left untouched. - ({}, True, True), - ], - ids=["enabled", "override_off", "missing_keeps_default"], -) -def test_fulfill_params_from_enterprise_enable_app_deploy( - monkeypatch: pytest.MonkeyPatch, - enterprise_info: dict[str, object], - initial: bool, - expected: bool, -) -> None: - monkeypatch.setattr( - feature_service_module.EnterpriseService, - "get_info", - staticmethod(lambda: enterprise_info), - ) - - features = SystemFeatureModel( - agent_tracing_api_collection_enabled=False, deployment_edition=DeploymentEdition.COMMUNITY - ) - features.enable_app_deploy = initial - - SystemFeatureService._fulfill_params_from_enterprise(features) - - assert features.enable_app_deploy is expected diff --git a/api/tests/unit_tests/services/test_system_feature_service_explore_banner.py b/api/tests/unit_tests/services/test_system_feature_service_explore_banner.py deleted file mode 100644 index 7b6fb8ba2c1..00000000000 --- a/api/tests/unit_tests/services/test_system_feature_service_explore_banner.py +++ /dev/null @@ -1,33 +0,0 @@ -"""Tests for the SystemFeatureService explore-banner policy.""" - -from collections.abc import Callable - -import pytest - -from enums import DeploymentEdition -from services.system_feature_service import SystemFeatureService - - -@pytest.mark.parametrize( - ("edition", "configured", "expected"), - [ - (DeploymentEdition.CLOUD, True, True), - (DeploymentEdition.CLOUD, False, False), - (DeploymentEdition.COMMUNITY, True, False), - (DeploymentEdition.ENTERPRISE, True, False), - ], -) -def test_get_system_features_enables_explore_banner_only_for_cloud( - monkeypatch: pytest.MonkeyPatch, - config_overrides: Callable[..., None], - edition: DeploymentEdition, - configured: bool, - expected: bool, -) -> None: - config_overrides(DEPLOYMENT_EDITION=edition, ENABLE_EXPLORE_BANNER=configured) - monkeypatch.setattr(SystemFeatureService, "_fulfill_params_from_enterprise", lambda *_: None) - - result = SystemFeatureService.get_public_system_features() - - assert SystemFeatureService.is_explore_banner_enabled() is expected - assert result.enable_explore_banner is expected diff --git a/api/tests/unit_tests/services/test_system_feature_service_internal_policies.py b/api/tests/unit_tests/services/test_system_feature_service_internal_policies.py deleted file mode 100644 index 4b0eebe7f0c..00000000000 --- a/api/tests/unit_tests/services/test_system_feature_service_internal_policies.py +++ /dev/null @@ -1,128 +0,0 @@ -from collections.abc import Callable - -import pytest - -from enums import DeploymentEdition -from services.entities.feature_entities import LicenseStatus -from services.system_feature_service import SystemFeatureService - - -def test_workspace_creation_uses_environment_policy( - monkeypatch: pytest.MonkeyPatch, config_overrides: Callable[..., None] -) -> None: - config_overrides(DEPLOYMENT_EDITION=DeploymentEdition.COMMUNITY, ALLOW_CREATE_WORKSPACE=True) - monkeypatch.setattr( - "services.system_feature_service.EnterpriseService.get_info", - lambda: (_ for _ in ()).throw(AssertionError("enterprise API should not be called")), - ) - - assert SystemFeatureService.is_workspace_creation_allowed() is True - - -def test_workspace_creation_uses_enterprise_policy( - monkeypatch: pytest.MonkeyPatch, config_overrides: Callable[..., None] -) -> None: - config_overrides(DEPLOYMENT_EDITION=DeploymentEdition.ENTERPRISE) - monkeypatch.setattr( - "services.system_feature_service.EnterpriseService.get_info", - lambda: {"IsAllowCreateWorkspace": False}, - ) - - assert SystemFeatureService.is_workspace_creation_allowed() is False - - -def test_workspace_creation_keeps_environment_policy_when_enterprise_value_is_missing( - monkeypatch: pytest.MonkeyPatch, - config_overrides: Callable[..., None], -) -> None: - config_overrides(DEPLOYMENT_EDITION=DeploymentEdition.ENTERPRISE, ALLOW_CREATE_WORKSPACE=True) - monkeypatch.setattr("services.system_feature_service.EnterpriseService.get_info", lambda: {}) - - assert SystemFeatureService.is_workspace_creation_allowed() is True - - -def test_plugin_manager_is_enabled_only_for_enterprise(config_overrides: Callable[..., None]) -> None: - config_overrides(DEPLOYMENT_EDITION=DeploymentEdition.ENTERPRISE) - assert SystemFeatureService.is_plugin_manager_enabled() is True - - config_overrides(DEPLOYMENT_EDITION=DeploymentEdition.COMMUNITY) - assert SystemFeatureService.is_plugin_manager_enabled() is False - - -def test_webapp_auth_enabled_does_not_query_enterprise_info( - monkeypatch: pytest.MonkeyPatch, config_overrides: Callable[..., None] -) -> None: - config_overrides(DEPLOYMENT_EDITION=DeploymentEdition.ENTERPRISE) - monkeypatch.setattr( - "services.system_feature_service.EnterpriseService.get_info", - lambda: (_ for _ in ()).throw(AssertionError("enterprise info should not be queried")), - ) - - assert SystemFeatureService.is_webapp_auth_enabled() is True - - -def test_registration_policy_uses_enterprise_override( - monkeypatch: pytest.MonkeyPatch, config_overrides: Callable[..., None] -) -> None: - config_overrides(DEPLOYMENT_EDITION=DeploymentEdition.ENTERPRISE, ALLOW_REGISTER=True) - monkeypatch.setattr( - "services.system_feature_service.EnterpriseService.get_info", - lambda: {"IsAllowRegister": False}, - ) - - assert SystemFeatureService.is_registration_allowed() is False - - -def test_password_login_policy_uses_enterprise_override( - monkeypatch: pytest.MonkeyPatch, config_overrides: Callable[..., None] -) -> None: - config_overrides(DEPLOYMENT_EDITION=DeploymentEdition.ENTERPRISE, ENABLE_EMAIL_PASSWORD_LOGIN=True) - monkeypatch.setattr( - "services.system_feature_service.EnterpriseService.get_info", - lambda: {"EnableEmailPasswordLogin": False}, - ) - - assert SystemFeatureService.is_email_password_login_enabled() is False - - -def test_branding_reads_enterprise_configuration( - monkeypatch: pytest.MonkeyPatch, config_overrides: Callable[..., None] -) -> None: - config_overrides(DEPLOYMENT_EDITION=DeploymentEdition.ENTERPRISE) - monkeypatch.setattr( - "services.system_feature_service.EnterpriseService.get_info", - lambda: { - "Branding": { - "applicationTitle": "Enterprise Dify", - "loginPageLogo": "login-logo", - "workspaceLogo": "workspace-logo", - "favicon": "favicon", - } - }, - ) - - branding = SystemFeatureService.get_branding() - - assert branding.enabled is True - assert branding.application_title == "Enterprise Dify" - assert branding.login_page_logo == "login-logo" - assert branding.workspace_logo == "workspace-logo" - assert branding.favicon == "favicon" - - -def test_license_status_ignores_unrelated_malformed_quota_fields( - monkeypatch: pytest.MonkeyPatch, config_overrides: Callable[..., None] -) -> None: - config_overrides(DEPLOYMENT_EDITION=DeploymentEdition.ENTERPRISE) - monkeypatch.setattr( - "services.system_feature_service.EnterpriseService.get_info", - lambda: { - "License": { - "status": "active", - "workspaces": {"enabled": True, "limit": 3, "used": {"unexpected": "shape"}}, - "licensedSeats": "unexpected-shape", - } - }, - ) - - assert SystemFeatureService.get_license_status() == LicenseStatus.ACTIVE diff --git a/api/tests/unit_tests/services/test_system_feature_service_learn_app.py b/api/tests/unit_tests/services/test_system_feature_service_learn_app.py deleted file mode 100644 index 370db54a8bf..00000000000 --- a/api/tests/unit_tests/services/test_system_feature_service_learn_app.py +++ /dev/null @@ -1,38 +0,0 @@ -"""Tests for SystemFeatureService learn-app and tour policies.""" - -from collections.abc import Callable - -import pytest - -from enums import DeploymentEdition -from services.entities.feature_entities import SystemFeatureModel -from services.system_feature_service import SystemFeatureService - - -def test_system_feature_model_defaults_enable_learn_app() -> None: - system_features = SystemFeatureModel( - agent_tracing_api_collection_enabled=False, deployment_edition=DeploymentEdition.COMMUNITY - ) - - assert system_features.enable_learn_app is True - assert system_features.enable_step_by_step_tour is False - - -@pytest.mark.parametrize("enabled", [True, False]) -def test_get_system_features_reads_enable_learn_app(config_overrides: Callable[..., None], enabled: bool) -> None: - config_overrides(ENABLE_LEARN_APP=enabled) - - result = SystemFeatureService.get_public_system_features() - - assert result.enable_learn_app is enabled - - -@pytest.mark.parametrize("enabled", [True, False]) -def test_get_system_features_reads_enable_step_by_step_tour( - config_overrides: Callable[..., None], enabled: bool -) -> None: - config_overrides(ENABLE_STEP_BY_STEP_TOUR=enabled) - - result = SystemFeatureService.get_public_system_features() - - assert result.enable_step_by_step_tour is enabled diff --git a/api/tests/unit_tests/services/test_system_feature_service_license_expiry_notice.py b/api/tests/unit_tests/services/test_system_feature_service_license_expiry_notice.py deleted file mode 100644 index 0ca89ae1933..00000000000 --- a/api/tests/unit_tests/services/test_system_feature_service_license_expiry_notice.py +++ /dev/null @@ -1,54 +0,0 @@ -"""Tests for the SystemFeatureService license-expiry policy.""" - -from collections.abc import Callable - -import pytest - -from enums import DeploymentEdition -from services import system_feature_service as feature_service_module -from services.entities.feature_entities import LicenseModel, LicenseStatus -from services.system_feature_service import SystemFeatureService - -_ENTERPRISE_INFO = {"License": {"status": LicenseStatus.EXPIRING, "expiredAt": "2026-12-31"}} - - -def test_license_model_defaults_license_expiry_notice_disabled() -> None: - """Without a license there is no expiry to announce, so the notice is off unless enabled explicitly.""" - assert LicenseModel().license_expiry_notice_enabled is False - - -@pytest.mark.parametrize("enabled", [True, False]) -def test_get_license_non_enterprise_ignores_expiry_notice_config( - config_overrides: Callable[..., None], enabled: bool -) -> None: - """Non-enterprise deployments have no license, so the env toggle never turns the notice on.""" - config_overrides( - ENABLE_LICENSE_EXPIRY_NOTICE=enabled, - DEPLOYMENT_EDITION=DeploymentEdition.COMMUNITY, - ) - - result = SystemFeatureService.get_license() - - assert result.license_expiry_notice_enabled is False - - -@pytest.mark.parametrize("enabled", [True, False]) -def test_get_license_enterprise_reads_license_expiry_notice_enabled( - monkeypatch: pytest.MonkeyPatch, config_overrides: Callable[..., None], enabled: bool -) -> None: - """The enterprise-sourced license carries the env-resolved notice flag alongside its real status.""" - config_overrides( - ENABLE_LICENSE_EXPIRY_NOTICE=enabled, - DEPLOYMENT_EDITION=DeploymentEdition.ENTERPRISE, - ) - monkeypatch.setattr( - feature_service_module.EnterpriseService, - "get_info", - staticmethod(lambda: _ENTERPRISE_INFO), - ) - - result = SystemFeatureService.get_license() - - assert result.status == LicenseStatus.EXPIRING - assert result.expired_at == "2026-12-31" - assert result.license_expiry_notice_enabled is enabled diff --git a/api/tests/unit_tests/services/test_system_feature_service_licensed_seats.py b/api/tests/unit_tests/services/test_system_feature_service_licensed_seats.py deleted file mode 100644 index a3f07e2e2ef..00000000000 --- a/api/tests/unit_tests/services/test_system_feature_service_licensed_seats.py +++ /dev/null @@ -1,42 +0,0 @@ -"""Tests for SystemFeatureService licensed-seat parsing.""" - -from collections.abc import Callable - -import pytest - -from enums import DeploymentEdition -from services import system_feature_service as feature_service_module -from services.entities.feature_entities import LicenseModel, LicenseStatus -from services.system_feature_service import SystemFeatureService - -_ENTERPRISE_INFO = {"License": {"licensedSeats": {"enabled": True, "limit": 3, "used": 1}}} - - -def test_get_license_parses_licensed_seats( - monkeypatch: pytest.MonkeyPatch, config_overrides: Callable[..., None] -) -> None: - """The authenticated license accessor copies the licensed-seat quota out of the enterprise payload.""" - config_overrides(DEPLOYMENT_EDITION=DeploymentEdition.ENTERPRISE) - monkeypatch.setattr( - feature_service_module.EnterpriseService, - "get_info", - staticmethod(lambda: _ENTERPRISE_INFO), - ) - - license_model = SystemFeatureService.get_license() - - assert isinstance(license_model, LicenseModel) - assert license_model.seats.enabled is True - assert license_model.seats.limit == 3 - assert license_model.seats.size == 1 - - -def test_get_license_non_enterprise_is_unconstrained(config_overrides: Callable[..., None]) -> None: - """Non-enterprise deployments have no license; seat allocation is unconstrained.""" - config_overrides(DEPLOYMENT_EDITION=DeploymentEdition.COMMUNITY) - - license_model = SystemFeatureService.get_license() - - assert license_model.status == LicenseStatus.NONE - assert license_model.seats.enabled is False - assert license_model.seats.is_available() is True diff --git a/api/tests/unit_tests/services/test_system_feature_service_plugin_installation_permission.py b/api/tests/unit_tests/services/test_system_feature_service_plugin_installation_permission.py deleted file mode 100644 index 9ed70688d3e..00000000000 --- a/api/tests/unit_tests/services/test_system_feature_service_plugin_installation_permission.py +++ /dev/null @@ -1,99 +0,0 @@ -"""Tests for SystemFeatureService plugin-installation policy.""" - -import logging -from collections.abc import Callable - -import pytest - -from enums import DeploymentEdition -from services import system_feature_service as feature_service_module -from services.entities.feature_entities import PluginInstallationScope, SystemFeatureModel -from services.system_feature_service import SystemFeatureService - - -def test_get_plugin_installation_permission_defaults_to_all_for_non_enterprise( - config_overrides: Callable[..., None], -) -> None: - config_overrides(DEPLOYMENT_EDITION=DeploymentEdition.COMMUNITY) - - permission = SystemFeatureService.get_plugin_installation_permission() - - assert permission.plugin_installation_scope is PluginInstallationScope.ALL - assert permission.restrict_to_marketplace_only is False - - -def test_get_plugin_installation_permission_parses_enterprise_policy( - monkeypatch: pytest.MonkeyPatch, - config_overrides: Callable[..., None], -) -> None: - config_overrides(DEPLOYMENT_EDITION=DeploymentEdition.ENTERPRISE) - monkeypatch.setattr( - feature_service_module.EnterpriseService, - "get_info", - staticmethod( - lambda: { - "PluginInstallationPermission": { - "pluginInstallationScope": "official_only", - "restrictToMarketplaceOnly": True, - } - } - ), - ) - - permission = SystemFeatureService.get_plugin_installation_permission() - - assert permission.plugin_installation_scope is PluginInstallationScope.OFFICIAL_ONLY - assert permission.restrict_to_marketplace_only is True - - -@pytest.mark.parametrize( - "invalid_permission", - [ - { - "pluginInstallationScope": "unknown-scope", - "restrictToMarketplaceOnly": False, - }, - { - "pluginInstallationScope": "all", - "restrictToMarketplaceOnly": "false", - }, - ], - ids=["unknown_scope", "non_boolean_marketplace_restriction"], -) -def test_invalid_enterprise_policy_denies_all_plugin_installations( - caplog: pytest.LogCaptureFixture, - invalid_permission: dict[str, object], -) -> None: - with caplog.at_level(logging.ERROR, logger="services.system_feature_service"): - permission = SystemFeatureService._resolve_plugin_installation_permission( - {"PluginInstallationPermission": invalid_permission} - ) - - assert permission.plugin_installation_scope is PluginInstallationScope.NONE - assert permission.restrict_to_marketplace_only is True - assert "denying all plugin installations" in caplog.text - - -def test_system_features_exposes_only_validated_plugin_installation_policy( - monkeypatch: pytest.MonkeyPatch, -) -> None: - monkeypatch.setattr( - feature_service_module.EnterpriseService, - "get_info", - staticmethod( - lambda: { - "PluginInstallationPermission": { - "pluginInstallationScope": "unknown-scope", - "restrictToMarketplaceOnly": False, - } - } - ), - ) - features = SystemFeatureModel( - agent_tracing_api_collection_enabled=False, deployment_edition=DeploymentEdition.ENTERPRISE - ) - - SystemFeatureService._fulfill_params_from_enterprise(features) - - assert features.plugin_installation_permission.plugin_installation_scope is PluginInstallationScope.NONE - assert features.plugin_installation_permission.restrict_to_marketplace_only is True diff --git a/api/tests/unit_tests/services/test_system_feature_service_sso_protocol.py b/api/tests/unit_tests/services/test_system_feature_service_sso_protocol.py deleted file mode 100644 index 48e42cfe8b3..00000000000 --- a/api/tests/unit_tests/services/test_system_feature_service_sso_protocol.py +++ /dev/null @@ -1,98 +0,0 @@ -"""Tests for SystemFeatureService SSO protocol parsing.""" - -import logging - -import pytest - -from enums import DeploymentEdition -from services import system_feature_service as feature_service_module -from services.entities.feature_entities import SSOProtocol, SystemFeatureModel -from services.system_feature_service import SystemFeatureService - - -def test_system_features_exposes_valid_enterprise_sso_protocols( - monkeypatch: pytest.MonkeyPatch, -) -> None: - monkeypatch.setattr( - feature_service_module.EnterpriseService, - "get_info", - staticmethod( - lambda: { - "SSOEnforcedForSigninProtocol": "saml", - "WebAppAuth": {}, - "SSOEnforcedForWebProtocol": "oidc", - } - ), - ) - features = SystemFeatureModel( - agent_tracing_api_collection_enabled=False, deployment_edition=DeploymentEdition.ENTERPRISE - ) - - SystemFeatureService._fulfill_params_from_enterprise(features) - - assert features.sso_enforced_for_signin_protocol is SSOProtocol.SAML - assert features.webapp_auth.sso_config.protocol is SSOProtocol.OIDC - - -@pytest.mark.parametrize("empty_protocol", [None, "", " "]) -def test_system_features_normalizes_empty_enterprise_sso_protocols_to_none( - monkeypatch: pytest.MonkeyPatch, - empty_protocol: object, -) -> None: - monkeypatch.setattr( - feature_service_module.EnterpriseService, - "get_info", - staticmethod( - lambda: { - "SSOEnforcedForSigninProtocol": empty_protocol, - "WebAppAuth": {}, - "SSOEnforcedForWebProtocol": empty_protocol, - } - ), - ) - features = SystemFeatureModel( - agent_tracing_api_collection_enabled=False, deployment_edition=DeploymentEdition.ENTERPRISE - ) - - SystemFeatureService._fulfill_params_from_enterprise(features) - - assert features.sso_enforced_for_signin_protocol is None - assert features.webapp_auth.sso_config.protocol is None - - -@pytest.mark.parametrize("invalid_protocol", ["unknown", 42]) -def test_system_features_rejects_invalid_enterprise_sso_protocols( - caplog: pytest.LogCaptureFixture, - monkeypatch: pytest.MonkeyPatch, - invalid_protocol: object, -) -> None: - monkeypatch.setattr( - feature_service_module.EnterpriseService, - "get_info", - staticmethod( - lambda: { - "SSOEnforcedForSigninProtocol": invalid_protocol, - "WebAppAuth": {}, - "SSOEnforcedForWebProtocol": invalid_protocol, - } - ), - ) - features = SystemFeatureModel( - agent_tracing_api_collection_enabled=False, deployment_edition=DeploymentEdition.ENTERPRISE - ) - - with caplog.at_level(logging.ERROR, logger="services.system_feature_service"): - SystemFeatureService._fulfill_params_from_enterprise(features) - - assert features.sso_enforced_for_signin_protocol is None - assert features.webapp_auth.sso_config.protocol is None - assert caplog.text.count("Invalid Enterprise SSO protocol") == 2 - - -def test_system_features_defaults_sso_protocols_to_none() -> None: - features = SystemFeatureModel( - agent_tracing_api_collection_enabled=False, deployment_edition=DeploymentEdition.COMMUNITY - ) - - assert features.sso_enforced_for_signin_protocol is None - assert features.webapp_auth.sso_config.protocol is None diff --git a/api/tests/unit_tests/services/test_system_feature_service_trial_models.py b/api/tests/unit_tests/services/test_system_feature_service_trial_models.py deleted file mode 100644 index 8760f79da7b..00000000000 --- a/api/tests/unit_tests/services/test_system_feature_service_trial_models.py +++ /dev/null @@ -1,9 +0,0 @@ -"""Tests for the public SystemFeatureService hosted-model contract.""" - -from services.system_feature_service import SystemFeatureService - - -def test_get_system_features_excludes_trial_models() -> None: - result = SystemFeatureService.get_public_system_features().model_dump() - - assert "trial_models" not in result diff --git a/api/tests/unit_tests/services/test_system_feature_service_webapp_public_access.py b/api/tests/unit_tests/services/test_system_feature_service_webapp_public_access.py deleted file mode 100644 index 6764dbf39d8..00000000000 --- a/api/tests/unit_tests/services/test_system_feature_service_webapp_public_access.py +++ /dev/null @@ -1,38 +0,0 @@ -"""Tests for the SystemFeatureService WebApp public-access policy.""" - -from collections.abc import Callable - -import pytest - -from enums import DeploymentEdition -from services.entities.feature_entities import SystemFeatureModel -from services.system_feature_service import SystemFeatureService - - -@pytest.mark.parametrize( - ("env_value", "expected"), - [ - (False, False), - (True, True), - ], - ids=["disabled_by_env", "enabled_by_env"], -) -def test_fulfill_system_params_from_env_sets_allow_public_access( - config_overrides: Callable[..., None], - env_value: bool, - expected: bool, -) -> None: - config_overrides(WEBAPP_PUBLIC_ACCESS_ENABLED=env_value) - - system_features = SystemFeatureModel( - agent_tracing_api_collection_enabled=False, deployment_edition=DeploymentEdition.COMMUNITY - ) - SystemFeatureService._fulfill_system_params_from_env(system_features) - - assert system_features.webapp_auth.allow_public_access is expected - - -def test_get_system_features_defaults_allow_public_access_to_true() -> None: - system_features = SystemFeatureService.get_public_system_features() - - assert system_features.webapp_auth.allow_public_access is True diff --git a/api/tests/unit_tests/services/test_tencent_tracing_config_gateway.py b/api/tests/unit_tests/services/test_tencent_tracing_config_gateway.py deleted file mode 100644 index ff5ab51464e..00000000000 --- a/api/tests/unit_tests/services/test_tencent_tracing_config_gateway.py +++ /dev/null @@ -1,324 +0,0 @@ -"""Tencent configuration persistence uses real RSA and bounded protocol checks.""" - -import json -import os -from dataclasses import dataclass -from typing import Never -from uuid import uuid4 - -import grpc # type: ignore[import-untyped] -import pytest -from dify_trace_tencent.configuration import TencentConfigurationError -from sqlalchemy.orm import Session, sessionmaker - -from core.ops.ops_trace_manager import OpsTraceManager -from machinery.context import RequestContext -from repositories.app.console_repository import ConsoleAppRepository -from repositories.app.tracing_config_repository import SQLAlchemyAppTracingConfigRepository -from services.app_tracing_config_gateway import OpsTraceManagerGateway -from services.app_tracing_config_service import ( - AppTracingConfigInvalidConfigurationError, - AppTracingConfigRecord, - AppTracingConfigService, - AppTracingConfigVerificationFailedError, -) -from tests.unit_tests.model_factories import make_app -from tests.unit_tests.services.tracing.test_tencent_configuration import CollectorBoundary -from tests.unit_tests.services.tracing.test_tencent_configuration import ( - _forbid_provider_io as _forbid_provider_io, # noqa: PLC0414 - Re-export pytest fixtures only in this module. -) -from tests.unit_tests.services.tracing.test_tencent_configuration import ( - collector as collector, # noqa: PLC0414 - Re-export pytest fixtures only in this module. -) - -pytest_plugins = ("tests.unit_tests.services.tracing.test_nullable_tracing_credentials",) - - -@dataclass -class ConfigurationCase: - collector: CollectorBoundary - gateway: OpsTraceManagerGateway - service: AppTracingConfigService - store: SQLAlchemyAppTracingConfigRepository - reader: ConsoleAppRepository - context: RequestContext - app_id: str - - def read(self, provider: str = "tencent") -> AppTracingConfigRecord | None: - return self.store.get( - workspace_id=self.context.active_workspace_id, app_id=self.app_id, tracing_provider=provider - ) - - -@pytest.fixture -def case( - workspace: str, sqlite_session_factory: sessionmaker[Session], collector: CollectorBoundary -) -> ConfigurationCase: - app_id = str(uuid4()) - with sqlite_session_factory.begin() as session: - session.add( - make_app( - app_id=app_id, - tenant_id=workspace, - tracing=json.dumps( - { - "enabled": True, - "tracing_provider": "langsmith", - "capture_content": True, - "capture_agent": False, - } - ), - ) - ) - gateway = OpsTraceManagerGateway() - store = SQLAlchemyAppTracingConfigRepository(session_factory=sqlite_session_factory) - old_config = OpsTraceManager.encrypt_tracing_config( - workspace, - "langsmith", - {"api_key": "synthetic-existing-key", "project": "existing-project", "endpoint": "https://existing.example"}, - ) - assert store.create(workspace_id=workspace, app_id=app_id, tracing_provider="langsmith", tracing_config=old_config) - return ConfigurationCase( - collector, - gateway, - AppTracingConfigService(configs=store, provider=gateway), - store, - ConsoleAppRepository(session_factory=sqlite_session_factory), - RequestContext("request", None, "account", workspace), - app_id, - ) - - -def test_post_get_masked_patch_persist_normalization_and_preserve_real_ciphertext(case: ConfigurationCase) -> None: - secret = "synthetic-private-tencent-token" - old_provider = case.read("langsmith") - submitted = {"token": secret, "endpoint": " HTTPS://FIRST.Example:4320/ ", "service_name": " first-app "} - original_submission = dict(submitted) - - case.service.create(case.context, case.app_id, "tencent", submitted, activate=True) - - before = case.read() - assert before is not None - assert before.tracing_config is not None - config = before.tracing_config - assert config["token"] != secret - assert config["endpoint"] == "https://first.example:4320" - assert config["service_name"] == "first-app" - assert config["project_url"] == "https://console.cloud.tencent.com/apm" - assert ( - OpsTraceManager.decrypt_tracing_config(case.context.active_workspace_id, "tencent", config)["token"] == secret - ) - assert submitted == original_submission - assert len(case.collector.channels) == 1 - assert case.collector.channels[0].calls[0].metadata == (("authorization", f"Bearer {secret}"),) - policy = case.reader.get_trace(case.context, case.app_id) - assert policy.enabled is True - assert policy.tracing_provider == "tencent" - assert policy.capture_content is True - assert policy.capture_agent is False - assert case.read("langsmith") == old_provider - - presented = case.service.get(case.context, case.app_id, "tencent") - assert presented is not None - assert presented.tracing_config is not None - masked = presented.tracing_config["token"] - assert isinstance(masked, str) - assert masked != secret - assert "*" in masked - assert presented.tracing_config["project_url"] == "https://console.cloud.tencent.com/apm" - assert len(case.collector.channels) == 1 - - replacement = {"token": masked, "endpoint": " http://SECOND.Example:4319/ ", "service_name": " second-app "} - original_replacement = dict(replacement) - case.service.update(case.context, case.app_id, "tencent", replacement, activate=True) - - after = case.read() - assert after is not None - assert after.tracing_config is not None - assert after.tracing_config["token"] == config["token"] - assert after.tracing_config["endpoint"] == "http://second.example:4319" - assert after.tracing_config["service_name"] == "second-app" - assert after.tracing_config["project_url"] == "https://console.cloud.tencent.com/apm" - assert replacement == original_replacement - assert len(case.collector.channels) == 2 - channel = case.collector.channels[-1] - assert channel.target == "second.example:4319" - assert channel.secure is False - assert channel.calls[0].metadata == (("authorization", f"Bearer {secret}"),) - assert all(channel.close_count == 1 for channel in case.collector.channels) - case.gateway.validate_saved_config( - workspace_id=case.context.active_workspace_id, tracing_provider="tencent", tracing_config=after.tracing_config - ) - assert case.service.get(case.context, case.app_id, "tencent") is not None - assert len(case.collector.channels) == 2 - assert case.read("langsmith") == old_provider - - -@pytest.mark.parametrize("saved_url", [None, "https://wrong.example/previous-project"]) -def test_historical_name_remains_readable_and_saved_validation_does_not_connect( - case: ConfigurationCase, saved_url: str | None -) -> None: - stored = OpsTraceManager.encrypt_tracing_config( - case.context.active_workspace_id, - "tencent", - {"token": "synthetic-historical-token", "endpoint": "https://legacy.example:4320", "service_name": "dify_app"}, - ) - if saved_url is not None: - stored["project_url"] = saved_url - assert case.store.create( - workspace_id=case.context.active_workspace_id, - app_id=case.app_id, - tracing_provider="tencent", - tracing_config=stored, - ) - before = case.read() - before_policy = case.reader.get_trace(case.context, case.app_id) - - presented = case.service.get(case.context, case.app_id, "tencent") - case.gateway.validate_saved_config( - workspace_id=case.context.active_workspace_id, tracing_provider="tencent", tracing_config=stored - ) - - assert presented is not None - assert presented.tracing_config is not None - assert presented.tracing_config["service_name"] == "dify_app" - assert presented.tracing_config["project_url"] == "https://console.cloud.tencent.com/apm" - assert presented.tracing_config["token"] != "synthetic-historical-token" - assert "*" in presented.tracing_config["token"] - assert case.collector.channels == [] - assert case.read() == before - assert case.reader.get_trace(case.context, case.app_id) == before_policy - - -@pytest.mark.parametrize("existing", [False, True]) -@pytest.mark.parametrize( - ("field_name", "invalid_value", "error_code"), - [ - ("service_name", "dify_app", "invalid_service_name"), - ("service_name", "", "invalid_service_name"), - ("endpoint", "https://collector.example:4320/v1/traces", "invalid_endpoint"), - ("endpoint", "https://collector.example:4320?token=private-endpoint-detail", "invalid_endpoint"), - ("endpoint", "https://", "invalid_endpoint"), - ], -) -def test_invalid_create_or_patch_rejects_before_rpc_encryption_or_persistence( - case: ConfigurationCase, - monkeypatch: pytest.MonkeyPatch, - caplog: pytest.LogCaptureFixture, - existing: bool, - field_name: str, - invalid_value: str, - error_code: str, -) -> None: - original = { - "token": "synthetic-private-token", - "endpoint": "https://original.example:4320", - "service_name": "original-app", - } - if existing: - case.service.create(case.context, case.app_id, "tencent", original) - case.collector.channels.clear() - before = case.read() - old_provider = case.read("langsmith") - before_policy = case.reader.get_trace(case.context, case.app_id) - forbidden_attempts: list[str] = [] - - def forbidden(*_args: object, **_kwargs: object) -> Never: - forbidden_attempts.append("encryption_or_write") - raise AssertionError("Invalid new input must fail before encryption or persistence") - - monkeypatch.setattr(OpsTraceManager, "encrypt_tracing_config", forbidden) - monkeypatch.setattr(case.store, "create", forbidden) - monkeypatch.setattr(case.store, "update", forbidden) - operation = case.service.update if existing else case.service.create - with pytest.raises(AppTracingConfigInvalidConfigurationError) as caught: - operation(case.context, case.app_id, "tencent", {**original, field_name: invalid_value}, activate=True) - - assert isinstance(caught.value.__cause__, TencentConfigurationError) - assert caught.value.__cause__.code == error_code - assert forbidden_attempts == [] - assert case.collector.channels == [] - assert case.read() == before - assert case.read("langsmith") == old_provider - assert case.reader.get_trace(case.context, case.app_id) == before_policy - assert "synthetic-private-token" not in caplog.text - assert "private-endpoint-detail" not in caplog.text - - -@pytest.mark.parametrize("existing", [False, True]) -@pytest.mark.parametrize("endpoint", ["https://collector.example:4320", "http://localhost:4317"]) -def test_rpc_failure_never_saves_or_activates_even_for_localhost( - case: ConfigurationCase, caplog: pytest.LogCaptureFixture, existing: bool, endpoint: str -) -> None: - secret = "synthetic-private-token" - original = {"token": secret, "endpoint": "https://original.example:4320", "service_name": "original-app"} - if existing: - case.service.create(case.context, case.app_id, "tencent", original) - case.collector.channels.clear() - before = case.read() - old_provider = case.read("langsmith") - before_policy = case.reader.get_trace(case.context, case.app_id) - case.collector.failure = grpc.RpcError("private-provider-details synthetic-private-token") - operation = case.service.update if existing else case.service.create - - with pytest.raises(AppTracingConfigVerificationFailedError) as caught: - operation( - case.context, - case.app_id, - "tencent", - {"token": "******" if existing else secret, "endpoint": endpoint, "service_name": "changed-app"}, - activate=True, - ) - - assert isinstance(caught.value.__cause__, TencentConfigurationError) - assert caught.value.__cause__.code == "connection_failed" - assert len(case.collector.channels) == 1 - channel = case.collector.channels[0] - assert channel.close_count == 1 - assert len(channel.calls) == 1 - assert channel.calls[0].metadata == (("authorization", f"Bearer {secret}"),) - assert case.read() == before - assert case.read("langsmith") == old_provider - assert case.reader.get_trace(case.context, case.app_id) == before_policy - assert "connection_failed" in caplog.text - assert "private-provider-details" not in caplog.text - assert secret not in caplog.text - assert secret not in str(caught.value) - assert secret not in str(caught.value.__cause__) - - -def test_explicit_configuration_ignores_and_preserves_otel_environment( - case: ConfigurationCase, monkeypatch: pytest.MonkeyPatch -) -> None: - environment = { - "OTEL_EXPORTER_OTLP_ENDPOINT": "http://ambient.example:4317", - "OTEL_EXPORTER_OTLP_TRACES_ENDPOINT": "http://ambient-traces.example:4317", - "OTEL_EXPORTER_OTLP_HEADERS": "authorization=Bearer unrelated-token", - "OTEL_EXPORTER_OTLP_TRACES_HEADERS": "authorization=Bearer unrelated-traces-token", - "OTEL_EXPORTER_OTLP_PROTOCOL": "http/protobuf", - "OTEL_SERVICE_NAME": "unrelated-service", - "OTEL_RESOURCE_ATTRIBUTES": "token=unrelated-resource-token,service.name=ambient-app", - } - for key, value in environment.items(): - monkeypatch.setenv(key, value) - - case.service.create( - case.context, - case.app_id, - "tencent", - {"token": "explicit-token", "endpoint": "https://explicit.example:4320", "service_name": "explicit-app"}, - ) - - assert len(case.collector.channels) == 1 - channel = case.collector.channels[0] - assert channel.target == "explicit.example:4320" - assert channel.secure is True - assert channel.calls[0].metadata == (("authorization", "Bearer explicit-token"),) - assert channel.calls[0].request.SerializeToString() == b"" - assert channel.close_count == 1 - presented = case.service.get(case.context, case.app_id, "tencent") - assert presented is not None - assert presented.tracing_config is not None - assert presented.tracing_config["service_name"] == "explicit-app" - assert len(case.collector.channels) == 1 - assert {key: os.environ.get(key) for key in environment} == environment diff --git a/api/tests/unit_tests/services/test_tidb_binding_service.py b/api/tests/unit_tests/services/test_tidb_binding_service.py index 66ce70c9081..45d67c3c134 100644 --- a/api/tests/unit_tests/services/test_tidb_binding_service.py +++ b/api/tests/unit_tests/services/test_tidb_binding_service.py @@ -20,7 +20,7 @@ from models.account import Tenant from models.dataset import Dataset, TidbAuthBinding from models.enums import TidbAuthBindingStatus from services import tidb_binding_service as service -from tests.unit_tests.knowledge_fs_core_fakes import RecordingCallable +from tests.helpers.unit.knowledge_fs_core_fakes import RecordingCallable @dataclass(frozen=True) diff --git a/api/tests/unit_tests/services/test_upload_file_delivery_service.py b/api/tests/unit_tests/services/test_upload_file_delivery_service.py index f8fde4a0bc8..d791e6a569a 100644 --- a/api/tests/unit_tests/services/test_upload_file_delivery_service.py +++ b/api/tests/unit_tests/services/test_upload_file_delivery_service.py @@ -16,7 +16,7 @@ from services.upload_file_delivery_service import ( UploadFileDeliveryRecord, UploadFileDeliveryService, ) -from tests.unit_tests.model_factories import make_tenant, make_upload_file +from tests.helpers.unit.model_factories import make_tenant, make_upload_file def _record(*, extension: str = "png") -> UploadFileDeliveryRecord: diff --git a/api/tests/unit_tests/services/test_variable_truncator.py b/api/tests/unit_tests/services/test_variable_truncator.py index 7cb31751b82..84d9054ecff 100644 --- a/api/tests/unit_tests/services/test_variable_truncator.py +++ b/api/tests/unit_tests/services/test_variable_truncator.py @@ -1,14 +1,3 @@ -""" -Comprehensive unit tests for VariableTruncator class based on current implementation. - -This test suite covers all functionality of the current VariableTruncator including: -- JSON size calculation for different data types -- String, array, and object truncation logic -- Segment-based truncation interface -- Helper methods for budget-based truncation -- Edge cases and error handling -""" - import functools import json import uuid @@ -18,6 +7,7 @@ from uuid import uuid4 import pytest from graphon.file import File, FileTransferMethod, FileType +from graphon.nodes.variable_assigner.common.helpers import UpdatedVariable from graphon.variables.segments import ( ArrayFileSegment, ArrayNumberSegment, @@ -29,6 +19,8 @@ from graphon.variables.segments import ( ObjectSegment, StringSegment, ) +from graphon.variables.types import SegmentType +from services import variable_truncator as truncator_module from services.variable_truncator import ( DummyVariableTruncator, MaxDepthExceededError, @@ -36,755 +28,915 @@ from services.variable_truncator import ( UnknownTypeError, VariableTruncator, ) +from tests.helpers.unit.config_override import apply_config_overrides +""" +Comprehensive unit tests for VariableTruncator class based on current implementation. -@pytest.fixture -def file() -> File: - return File( - file_id=str(uuid4()), # Generate new UUID for File.id - tenant_id=str(uuid.uuid4()), - file_type=FileType.DOCUMENT, - transfer_method=FileTransferMethod.LOCAL_FILE, - related_id=str(uuid.uuid4()), - filename="test_file.txt", - extension=".txt", - mime_type="text/plain", - size=1024, - storage_key="initial_key", - ) +This test suite covers all functionality of the current VariableTruncator including: +- JSON size calculation for different data types +- String, array, and object truncation logic +- Segment-based truncation interface +- Helper methods for budget-based truncation +- Edge cases and error handling +""" _compact_json_dumps = functools.partial(json.dumps, separators=(",", ":")) -class TestCalculateJsonSize: - """Test calculate_json_size method with different data types.""" - - @pytest.fixture - def truncator(self): - return VariableTruncator() - - def test_string_size_calculation(self): - """Test JSON size calculation for strings.""" - # Simple ASCII string - assert VariableTruncator.calculate_json_size("hello") == 7 # "hello" + 2 quotes - - # Empty string - assert VariableTruncator.calculate_json_size("") == 2 # Just quotes - - # Unicode string - assert VariableTruncator.calculate_json_size("你好") == 4 - - def test_number_size_calculation(self, truncator): - """Test JSON size calculation for numbers.""" - assert truncator.calculate_json_size(123) == 3 - assert truncator.calculate_json_size(12.34) == 5 - assert truncator.calculate_json_size(-456) == 4 - assert truncator.calculate_json_size(0) == 1 - - def test_boolean_size_calculation(self, truncator): - """Test JSON size calculation for booleans.""" - assert truncator.calculate_json_size(True) == 4 # "true" - assert truncator.calculate_json_size(False) == 5 # "false" - - def test_null_size_calculation(self, truncator): - """Test JSON size calculation for None/null.""" - assert truncator.calculate_json_size(None) == 4 # "null" - - def test_array_size_calculation(self, truncator): - """Test JSON size calculation for arrays.""" - # Empty array - assert truncator.calculate_json_size([]) == 2 # "[]" - - # Simple array - simple_array = [1, 2, 3] - # [1,2,3] = 1 + 1 + 1 + 1 + 1 + 2 = 7 (numbers + commas + brackets) - assert truncator.calculate_json_size(simple_array) == 7 - - # Array with strings - string_array = ["a", "b"] - # ["a","b"] = 3 + 3 + 1 + 2 = 9 (quoted strings + comma + brackets) - assert truncator.calculate_json_size(string_array) == 9 - - def test_object_size_calculation(self, truncator): - """Test JSON size calculation for objects.""" - # Empty object - assert truncator.calculate_json_size({}) == 2 # "{}" - - # Simple object - simple_obj = {"a": 1} - # {"a":1} = 3 + 1 + 1 + 2 = 7 (key + colon + value + brackets) - assert truncator.calculate_json_size(simple_obj) == 7 - - # Multiple keys - multi_obj = {"a": 1, "b": 2} - # {"a":1,"b":2} = 3 + 1 + 1 + 1 + 3 + 1 + 1 + 2 = 13 - assert truncator.calculate_json_size(multi_obj) == 13 - - def test_nested_structure_size_calculation(self, truncator): - """Test JSON size calculation for nested structures.""" - nested = {"items": [1, 2, {"nested": "value"}]} - size = truncator.calculate_json_size(nested) - assert size > 0 # Should calculate without error - - # Verify it matches actual JSON length roughly - - actual_json = _compact_json_dumps(nested) - # Should be close but not exact due to UTF-8 encoding considerations - assert abs(size - len(actual_json.encode())) <= 5 - - def test_calculate_json_size_max_depth_exceeded(self, truncator): - """Test that calculate_json_size handles deep nesting gracefully.""" - # Create deeply nested structure - nested: dict[str, Any] = {"level": 0} - current = nested - for i in range(105): # Create deep nesting - current["next"] = {"level": i + 1} - current = current["next"] - - # Should either raise an error or handle gracefully - with pytest.raises(MaxDepthExceededError): - truncator.calculate_json_size(nested) - - def test_calculate_json_size_unknown_type(self, truncator): - """Test that calculate_json_size raises error for unknown types.""" - - class CustomType: - pass - - with pytest.raises(UnknownTypeError): - truncator.calculate_json_size(CustomType()) - - -class TestStringTruncation: - LENGTH_LIMIT = 10 - """Test string truncation functionality.""" - - @pytest.fixture - def small_truncator(self): - return VariableTruncator(string_length_limit=10) - - def test_short_string_no_truncation(self, small_truncator): - """Test that short strings are not truncated.""" - short_str = "hello" - result = small_truncator._truncate_string(short_str, self.LENGTH_LIMIT) - assert result.value == short_str - assert result.truncated is False - assert result.value_size == VariableTruncator.calculate_json_size(short_str) - - def test_long_string_truncation(self, small_truncator: VariableTruncator): - """Test that long strings are truncated with ellipsis.""" - long_str = "this is a very long string that exceeds the limit" - result = small_truncator._truncate_string(long_str, self.LENGTH_LIMIT) - - assert result.truncated is True - assert result.value == long_str[:5] + "..." - assert result.value_size == 10 # 10 chars + "..." - - def test_exact_limit_string(self, small_truncator: VariableTruncator): - """Test string exactly at limit.""" - exact_str = "1234567890" # Exactly 10 chars - result = small_truncator._truncate_string(exact_str, self.LENGTH_LIMIT) - assert result.value == "12345..." - assert result.truncated is True - assert result.value_size == 10 - - -class TestArrayTruncation: - """Test array truncation functionality.""" - - @pytest.fixture - def small_truncator(self): - return VariableTruncator(array_element_limit=3, max_size_bytes=100) - - def test_small_array_no_truncation(self, small_truncator: VariableTruncator): - """Test that small arrays are not truncated.""" - small_array: list[object] = [1, 2] - result = small_truncator._truncate_array(small_array, 1000) - assert result.value == small_array - assert result.truncated is False - - def test_array_element_limit_truncation(self, small_truncator: VariableTruncator): - """Test that arrays over element limit are truncated.""" - large_array: list[object] = [1, 2, 3, 4, 5, 6] # Exceeds limit of 3 - result = small_truncator._truncate_array(large_array, 1000) - - assert result.truncated is True - assert result.value == [1, 2, 3] - - def test_array_size_budget_truncation(self, small_truncator: VariableTruncator): - """Test array truncation due to size budget constraints.""" - # Create array with strings that will exceed size budget - large_strings: list[object] = ["very long string " * 5, "another long string " * 5] - result = small_truncator._truncate_array(large_strings, 50) - - assert result.truncated is True - # Should have truncated the strings within the array - for item in result.value: - assert isinstance(item, str) - assert VariableTruncator.calculate_json_size(result.value) <= 50 - - def test_array_with_nested_objects(self, small_truncator): - """Test array truncation with nested objects.""" - nested_array = [ - {"name": "item1", "data": "some data"}, - {"name": "item2", "data": "more data"}, - {"name": "item3", "data": "even more data"}, - ] - result = small_truncator._truncate_array(nested_array, 30) - - assert isinstance(result.value, list) - assert len(result.value) <= 3 - for item in result.value: - assert isinstance(item, dict) - - -class TestObjectTruncation: - """Test object truncation functionality.""" - - @pytest.fixture - def small_truncator(self): - return VariableTruncator(max_size_bytes=100) - - def test_small_object_no_truncation(self, small_truncator): - """Test that small objects are not truncated.""" - small_obj = {"a": 1, "b": 2} - result = small_truncator._truncate_object(small_obj, 1000) - assert result.value == small_obj - assert result.truncated is False - - def test_empty_object_no_truncation(self, small_truncator): - """Test that empty objects are not truncated.""" - empty_obj = {} - result = small_truncator._truncate_object(empty_obj, 100) - assert result.value == empty_obj - assert result.truncated is False - - def test_object_value_truncation(self, small_truncator): - """Test object truncation where values are truncated to fit budget.""" - obj_with_long_values = { - "key1": "very long string " * 10, - "key2": "another long string " * 10, - "key3": "third long string " * 10, - } - result = small_truncator._truncate_object(obj_with_long_values, 80) - - assert result.truncated is True - assert isinstance(result.value, dict) - - assert set(result.value.keys()).issubset(obj_with_long_values.keys()) - - # Values should be truncated if they exist - for key, value in result.value.items(): - assert isinstance(value, str) - original_value = obj_with_long_values[key] - # Value should be same or smaller - assert len(value) <= len(original_value) - - def test_object_key_dropping(self, small_truncator): - """Test object truncation where keys are dropped due to size constraints.""" - large_obj = {f"key{i:02d}": f"value{i}" for i in range(20)} - result = small_truncator._truncate_object(large_obj, 50) - - assert result.truncated is True - assert len(result.value) < len(large_obj) - - # Should maintain sorted key order - result_keys = list(result.value.keys()) - assert result_keys == sorted(result_keys) - - def test_object_with_nested_structures(self, small_truncator): - """Test object truncation with nested arrays and objects.""" - nested_obj = {"simple": "value", "array": [1, 2, 3, 4, 5], "nested": {"inner": "data", "more": ["a", "b", "c"]}} - result = small_truncator._truncate_object(nested_obj, 60) - - assert isinstance(result.value, dict) - - -class TestSegmentBasedTruncation: - """Test the main truncate method that works with Segments.""" - - @pytest.fixture - def truncator(self): - return VariableTruncator() - - @pytest.fixture - def small_truncator(self): - return VariableTruncator(string_length_limit=20, array_element_limit=3, max_size_bytes=200) - - def test_integer_segment_no_truncation(self, truncator): - """Test that integer segments are never truncated.""" - segment = IntegerSegment(value=12345) - result = truncator.truncate(segment) - - assert isinstance(result, TruncationResult) - assert result.truncated is False - assert result.result == segment - - def test_boolean_as_integer_segment(self, truncator): - """Test boolean values in IntegerSegment are converted to int.""" - segment = IntegerSegment(value=True) - result = truncator.truncate(segment) - - assert isinstance(result, TruncationResult) - assert result.truncated is False - assert isinstance(result.result, IntegerSegment) - assert result.result.value == 1 # True converted to 1 - - def test_float_segment_no_truncation(self, truncator): - """Test that float segments are never truncated.""" - segment = FloatSegment(value=123.456) - result = truncator.truncate(segment) - - assert isinstance(result, TruncationResult) - assert result.truncated is False - assert result.result == segment - - def test_none_segment_no_truncation(self, truncator): - """Test that None segments are never truncated.""" - segment = NoneSegment() - result = truncator.truncate(segment) - - assert isinstance(result, TruncationResult) - assert result.truncated is False - assert result.result == segment - - def test_file_segment_no_truncation(self, truncator, file): - """Test that file segments are never truncated.""" - file_segment = FileSegment(value=file) - result = truncator.truncate(file_segment) - assert result.result == file_segment - assert result.truncated is False - - def test_array_file_segment_no_truncation(self, truncator, file): - """Test that array file segments are never truncated.""" - - array_file_segment = ArrayFileSegment(value=[file] * 20) - result = truncator.truncate(array_file_segment) - assert result.result == array_file_segment - assert result.truncated is False - - def test_string_segment_small_no_truncation(self, truncator): - """Test small string segments are not truncated.""" - segment = StringSegment(value="hello world") - result = truncator.truncate(segment) - - assert isinstance(result, TruncationResult) - assert result.truncated is False - assert result.result == segment - - def test_string_segment_large_truncation(self, small_truncator): - """Test large string segments are truncated.""" - long_text = "this is a very long string that will definitely exceed the limit" - segment = StringSegment(value=long_text) - result = small_truncator.truncate(segment) - - assert isinstance(result, TruncationResult) - assert result.truncated is True - assert isinstance(result.result, StringSegment) - assert len(result.result.value) < len(long_text) - assert result.result.value.endswith("...") - - def test_array_segment_small_no_truncation(self, truncator): - """Test small array segments are not truncated.""" - from factories.variable_factory import build_segment - - segment = build_segment([1, 2, 3]) - result = truncator.truncate(segment) - - assert isinstance(result, TruncationResult) - assert result.truncated is False - assert result.result == segment - - def test_array_segment_large_truncation(self, small_truncator): - """Test large array segments are truncated.""" - from factories.variable_factory import build_segment - - large_array = list(range(10)) # Exceeds element limit of 3 - segment = build_segment(large_array) - result = small_truncator.truncate(segment) - - assert isinstance(result, TruncationResult) - assert result.truncated is True - assert isinstance(result.result, ArraySegment) - assert len(result.result.value) <= 3 - - def test_object_segment_small_no_truncation(self, truncator): - """Test small object segments are not truncated.""" - segment = ObjectSegment(value={"key": "value"}) - result = truncator.truncate(segment) - - assert isinstance(result, TruncationResult) - assert result.truncated is False - assert result.result == segment - - def test_object_segment_large_truncation(self, small_truncator): - """Test large object segments are truncated.""" - large_obj = {f"key{i}": f"very long value {i}" * 5 for i in range(5)} - segment = ObjectSegment(value=large_obj) - result = small_truncator.truncate(segment) - - assert isinstance(result, TruncationResult) - assert result.truncated is True - assert isinstance(result.result, ObjectSegment) - # Object should be smaller or equal than original - original_size = small_truncator.calculate_json_size(large_obj) - result_size = small_truncator.calculate_json_size(result.result.value) - assert result_size <= original_size - - def test_final_size_fallback_to_json_string(self, small_truncator): - """Test final fallback when truncated result still exceeds size limit.""" - # Create data that will still be large after initial truncation - large_nested_data = {"data": ["very long string " * 5] * 5, "more": {"nested": "content " * 20}} - segment = ObjectSegment(value=large_nested_data) - - # Use very small limit to force JSON string fallback - tiny_truncator = VariableTruncator(max_size_bytes=50) - result = tiny_truncator.truncate(segment) - - assert isinstance(result, TruncationResult) - assert result.truncated is True - assert isinstance(result.result, StringSegment) - # Should be JSON string with possible truncation - assert len(result.result.value) <= 53 # 50 + "..." = 53 - - def test_final_size_fallback_string_truncation(self, small_truncator): - """Test final fallback for string that still exceeds limit.""" - # Create very long string that exceeds string length limit - very_long_string = "x" * 6000 # Exceeds default string_length_limit of 5000 - segment = StringSegment(value=very_long_string) - - # Use small limit to test string fallback path - tiny_truncator = VariableTruncator(string_length_limit=100, max_size_bytes=50) - result = tiny_truncator.truncate(segment) - - assert isinstance(result, TruncationResult) - assert result.truncated is True - assert isinstance(result.result, StringSegment) - # Should be truncated due to string limit or final size limit - assert len(result.result.value) <= 1000 # Much smaller than original - - -class TestEdgeCases: - """Test edge cases and error conditions.""" - - def test_empty_inputs(self): - """Test truncator with empty inputs.""" - truncator = VariableTruncator() - - # Empty string - result = truncator.truncate(StringSegment(value="")) - assert not result.truncated - assert result.result.value == "" - - # Empty array - from factories.variable_factory import build_segment - - result = truncator.truncate(build_segment([])) - assert not result.truncated - assert result.result.value == [] - - # Empty object - result = truncator.truncate(ObjectSegment(value={})) - assert not result.truncated - assert result.result.value == {} - - def test_zero_and_negative_limits(self): - """Test truncator behavior with zero or very small limits.""" - # Zero string limit - with pytest.raises(ValueError, match="string_length_limit should be greater than 3\\."): - VariableTruncator(string_length_limit=3) - - with pytest.raises(ValueError, match="array_element_limit should be greater than 0\\."): - VariableTruncator(array_element_limit=0) - - with pytest.raises(ValueError, match="max_size_bytes should be greater than 0\\."): - VariableTruncator(max_size_bytes=0) - - def test_unicode_and_special_characters(self): - """Test truncator with unicode and special characters.""" - truncator = VariableTruncator(string_length_limit=10) - - # Unicode characters - unicode_text = "你好世界你好世界你好世界" # Multi-byte UTF-8 characters - result = truncator.truncate(StringSegment(value=unicode_text)) - assert result.truncated is True - - # Special JSON characters - special_chars = '{"key": "value with \\"quotes\\" and \\n newlines"}' - result = truncator.truncate(StringSegment(value=special_chars)) - assert isinstance(result.result, StringSegment) - - -class TestTruncateJsonPrimitives: - """Test _truncate_json_primitives method with different data types.""" - - @pytest.fixture - def truncator(self): - return VariableTruncator() - - def test_truncate_json_primitives_file_type(self, truncator, file): - """Test that File objects are handled correctly in _truncate_json_primitives.""" - # Test File object is returned as-is without truncation - result = truncator._truncate_json_primitives(file, 1000) - - assert result.value == file - assert result.truncated is False - # Size should be calculated correctly - expected_size = VariableTruncator.calculate_json_size(file) - assert result.value_size == expected_size - - def test_truncate_json_primitives_file_type_small_budget(self, truncator, file): - """Test that File objects are returned as-is even with small budget.""" - # Even with a small size budget, File objects should not be truncated - result = truncator._truncate_json_primitives(file, 10) - - assert result.value == file - assert result.truncated is False - - def test_truncate_json_primitives_file_type_in_array(self, truncator, file): - """Test File objects in arrays are handled correctly.""" - array_with_files = [file, file] - result = truncator._truncate_json_primitives(array_with_files, 1000) - - assert isinstance(result.value, list) - assert len(result.value) == 2 - assert result.value[0] == file - assert result.value[1] == file - assert result.truncated is False - - def test_truncate_json_primitives_file_type_in_object(self, truncator, file): - """Test File objects in objects are handled correctly.""" - obj_with_files = {"file1": file, "file2": file} - result = truncator._truncate_json_primitives(obj_with_files, 1000) - - assert isinstance(result.value, dict) - assert len(result.value) == 2 - assert result.value["file1"] == file - assert result.value["file2"] == file - assert result.truncated is False - - -class TestIntegrationScenarios: - """Test realistic integration scenarios.""" - - def test_workflow_output_scenario(self): - """Test truncation of typical workflow output data.""" - truncator = VariableTruncator() - - workflow_data = { - "result": "success", - "data": { - "users": [ - {"id": 1, "name": "Alice", "email": "alice@example.com"}, - {"id": 2, "name": "Bob", "email": "bob@example.com"}, - ] - * 3, # Multiply to make it larger - "metadata": { - "count": 6, - "processing_time": "1.23s", - "details": "x" * 200, # Long string but not too long - }, - }, - } - - segment = ObjectSegment(value=workflow_data) - result = truncator.truncate(segment) - - assert isinstance(result, TruncationResult) - assert isinstance(result.result, (ObjectSegment, StringSegment)) - # Should handle complex nested structure appropriately - - def test_large_text_processing_scenario(self): - """Test truncation of large text data.""" - truncator = VariableTruncator(string_length_limit=100) - - large_text = "This is a very long text document. " * 20 # Make it larger than limit - - segment = StringSegment(value=large_text) - result = truncator.truncate(segment) - - assert isinstance(result, TruncationResult) - assert result.truncated is True - assert isinstance(result.result, StringSegment) - assert len(result.result.value) <= 103 # 100 + "..." - assert result.result.value.endswith("...") - - def test_mixed_data_types_scenario(self): - """Test truncation with mixed data types in complex structure.""" - truncator = VariableTruncator(string_length_limit=30, array_element_limit=3, max_size_bytes=300) - - mixed_data = { - "strings": ["short", "medium length", "very long string " * 3], - "numbers": [1, 2.5, 999999], - "booleans": [True, False, True], - "nested": { - "more_strings": ["nested string " * 2], - "more_numbers": list(range(5)), - "deep": {"level": 3, "content": "deep content " * 3}, - }, - "nulls": [None, None], - } - - segment = ObjectSegment(value=mixed_data) - result = truncator.truncate(segment) - - assert isinstance(result, TruncationResult) - assert result.truncated is True - assert isinstance(result.result, ObjectSegment) - # Verify the result is smaller or equal than original - original_size = truncator.calculate_json_size(mixed_data) - result_size = truncator.calculate_json_size(result.result.value) - assert result_size <= original_size - - def test_file_and_array_file_variable_mapping(self, file): - truncator = VariableTruncator(string_length_limit=30, array_element_limit=3, max_size_bytes=300) - - mapping = {"array_file": [file]} - truncated_mapping, truncated = truncator.truncate_variable_mapping(mapping) - assert truncated is False - assert truncated_mapping == mapping - - -def test_dummy_variable_truncator_methods(): - """Test DummyVariableTruncator methods work correctly.""" - truncator = DummyVariableTruncator() - - # Test truncate_variable_mapping - test_data: dict[str, Any] = { - "key1": "value1", - "key2": ["item1", "item2"], - "large_array": list(range(2000)), - } - result, is_truncated = truncator.truncate_variable_mapping(test_data) - - assert result == test_data - assert not is_truncated - - # Test truncate method - segment = StringSegment(value="test string") - result = truncator.truncate(segment) - assert isinstance(result, TruncationResult) - assert result.result == segment - assert result.truncated is False - - segment = ArrayNumberSegment(value=list(range(2000))) - result = truncator.truncate(segment) - assert isinstance(result, TruncationResult) - assert result.result == segment - assert result.truncated is False - - -# --------------------------------------------------------------------------- -# Regression tests for langgenius/dify#39218. -# -# Before the fix, ``_truncate_array`` had a "Dirty fix" branch that -# unconditionally appended every ``File`` element to ``truncated_value`` -# *before* the count cap and the byte-budget check, and *before* -# ``used_size`` was ever incremented. That made ``list[File]`` arrays: -# 1. uncapped by ``array_element_limit``, -# 2. uncounted against ``max_size_bytes``, and -# 3. always reported ``truncated=False``. -# The fix routes ``File`` through ``_truncate_json_primitives``'s dedicated -# ``File`` branch, which returns the file as-is with its real serialized -# size, while preserving the count cap and the byte budget. -# --------------------------------------------------------------------------- - - -class TestFileArrayTruncationRegression39218: - """``list[File]`` must respect ``array_element_limit`` and the byte budget.""" - - @pytest.fixture - def truncator(self) -> VariableTruncator: - return VariableTruncator( - array_element_limit=3, - max_size_bytes=1000, - string_length_limit=50, - ) - +class TestVariableTruncatorCases: @staticmethod - def _make_file(name: str = "f") -> File: + @pytest.fixture + def file() -> File: return File( - id=name, - type=FileType.DOCUMENT, - transfer_method=FileTransferMethod.REMOTE_URL, - remote_url=f"https://example.com/{name}.txt", - filename=f"{name}.txt", + file_id=str(uuid4()), # Generate new UUID for File.id + tenant_id=str(uuid.uuid4()), + file_type=FileType.DOCUMENT, + transfer_method=FileTransferMethod.LOCAL_FILE, + related_id=str(uuid.uuid4()), + filename="test_file.txt", extension=".txt", mime_type="text/plain", size=1024, + storage_key="initial_key", ) - def test_file_array_respects_element_count_cap(self, truncator: VariableTruncator) -> None: - # Use a target_size larger than ``count * file_size`` so the byte - # budget never binds — only the count cap should fire. - # Each File serializes to ~237 bytes; 3 files = ~713 bytes. - files = [self._make_file(f"f{i}") for i in range(500)] + class TestCalculateJsonSize: + """Test calculate_json_size method with different data types.""" - result = truncator._truncate_array(files, target_size=10_000_000) + @pytest.fixture + def truncator(self): + return VariableTruncator() - # Before the fix, all 500 File entries survived (``len(value)==500``, - # ``truncated==False``). After the fix, the array is capped at - # ``array_element_limit=3`` and ``truncated`` flips to True. - assert len(result.value) == 3 - assert result.truncated is True + def test_string_size_calculation(self): + """Test JSON size calculation for strings.""" + # Simple ASCII string + assert VariableTruncator.calculate_json_size("hello") == 7 # "hello" + 2 quotes - def test_file_array_reports_real_used_size(self, truncator: VariableTruncator) -> None: - # Large budget so the count cap fires before the byte budget does. - files = [self._make_file(f"f{i}") for i in range(500)] + # Empty string + assert VariableTruncator.calculate_json_size("") == 2 # Just quotes - result = truncator._truncate_array(files, target_size=10_000_000) + # Unicode string + assert VariableTruncator.calculate_json_size("你好") == 4 - # Before the fix, ``used_size`` for a File array was the empty-array - # baseline of 2 bytes (``[]``), regardless of how many File entries - # actually returned. After the fix, ``used_size`` reflects the real - # serialized size of the returned ``File`` payload. - assert result.value_size > 100 - assert result.truncated is True + def test_number_size_calculation(self, truncator): + """Test JSON size calculation for numbers.""" + assert truncator.calculate_json_size(123) == 3 + assert truncator.calculate_json_size(12.34) == 5 + assert truncator.calculate_json_size(-456) == 4 + assert truncator.calculate_json_size(0) == 1 - def test_file_array_respects_byte_budget(self, truncator: VariableTruncator) -> None: - # Use a small ``target_size`` so the byte budget is the binding - # constraint. Each File serializes to ~237 bytes, so even one File - # blows the 200-byte budget. - files = [self._make_file(f"f{i}") for i in range(50)] + def test_boolean_size_calculation(self, truncator): + """Test JSON size calculation for booleans.""" + assert truncator.calculate_json_size(True) == 4 # "true" + assert truncator.calculate_json_size(False) == 5 # "false" - result = truncator._truncate_array(files, target_size=200) + def test_null_size_calculation(self, truncator): + """Test JSON size calculation for None/null.""" + assert truncator.calculate_json_size(None) == 4 # "null" - # Before the fix, all 50 File entries survived and ``used_size`` - # reported ``2`` (the empty-array baseline). After the fix, the - # loop sees the File payload: ``value_size`` reflects the real - # serialized size, and the loop stops after the first File because - # adding the next one would exceed ``target_size``. - assert len(result.value) == 1 - assert result.value_size > 100 # the File's real serialized size - assert result.value_size <= 250 # in the ballpark of the budget + def test_array_size_calculation(self, truncator): + """Test JSON size calculation for arrays.""" + # Empty array + assert truncator.calculate_json_size([]) == 2 # "[]" - def test_mixed_array_counts_files_toward_cap(self, truncator: VariableTruncator) -> None: - mixed: list[object] = [ - self._make_file("f0"), - "a", - self._make_file("f1"), - "b", - self._make_file("f2"), - "c", - self._make_file("f3"), - "d", - ] + # Simple array + simple_array = [1, 2, 3] + # [1,2,3] = 1 + 1 + 1 + 1 + 1 + 2 = 7 (numbers + commas + brackets) + assert truncator.calculate_json_size(simple_array) == 7 - result = truncator._truncate_array(mixed, target_size=10_000_000) + # Array with strings + string_array = ["a", "b"] + # ["a","b"] = 3 + 3 + 1 + 2 = 9 (quoted strings + comma + brackets) + assert truncator.calculate_json_size(string_array) == 9 - # 8 items, cap of 3 → exactly 3 items. Files and primitives are - # counted together toward the cap. - assert len(result.value) == 3 - assert result.truncated is True + def test_object_size_calculation(self, truncator): + """Test JSON size calculation for objects.""" + # Empty object + assert truncator.calculate_json_size({}) == 2 # "{}" - def test_single_file_in_array_is_preserved(self, truncator: VariableTruncator) -> None: - result = truncator._truncate_array([self._make_file("only")], target_size=10_000_000) + # Simple object + simple_obj = {"a": 1} + # {"a":1} = 3 + 1 + 1 + 2 = 7 (key + colon + value + brackets) + assert truncator.calculate_json_size(simple_obj) == 7 - # The File itself is not truncated — the dedicated ``File`` branch - # in ``_truncate_json_primitives`` returns the file untouched. Only - # the array-shape accounting changes. - assert len(result.value) == 1 - assert isinstance(result.value[0], File) - assert result.value[0].id == "only" + # Multiple keys + multi_obj = {"a": 1, "b": 2} + # {"a":1,"b":2} = 3 + 1 + 1 + 1 + 3 + 1 + 1 + 2 = 13 + assert truncator.calculate_json_size(multi_obj) == 13 + + def test_nested_structure_size_calculation(self, truncator): + """Test JSON size calculation for nested structures.""" + nested = {"items": [1, 2, {"nested": "value"}]} + size = truncator.calculate_json_size(nested) + assert size > 0 # Should calculate without error + + # Verify it matches actual JSON length roughly + + actual_json = _compact_json_dumps(nested) + # Should be close but not exact due to UTF-8 encoding considerations + assert abs(size - len(actual_json.encode())) <= 5 + + def test_calculate_json_size_max_depth_exceeded(self, truncator): + """Test that calculate_json_size handles deep nesting gracefully.""" + # Create deeply nested structure + nested: dict[str, Any] = {"level": 0} + current = nested + for i in range(105): # Create deep nesting + current["next"] = {"level": i + 1} + current = current["next"] + + # Should either raise an error or handle gracefully + with pytest.raises(MaxDepthExceededError): + truncator.calculate_json_size(nested) + + def test_calculate_json_size_unknown_type(self, truncator): + """Test that calculate_json_size raises error for unknown types.""" + + class CustomType: + pass + + with pytest.raises(UnknownTypeError): + truncator.calculate_json_size(CustomType()) + + class TestStringTruncation: + LENGTH_LIMIT = 10 + """Test string truncation functionality.""" + + @pytest.fixture + def small_truncator(self): + return VariableTruncator(string_length_limit=10) + + def test_short_string_no_truncation(self, small_truncator): + """Test that short strings are not truncated.""" + short_str = "hello" + result = small_truncator._truncate_string(short_str, self.LENGTH_LIMIT) + assert result.value == short_str + assert result.truncated is False + assert result.value_size == VariableTruncator.calculate_json_size(short_str) + + def test_long_string_truncation(self, small_truncator: VariableTruncator): + """Test that long strings are truncated with ellipsis.""" + long_str = "this is a very long string that exceeds the limit" + result = small_truncator._truncate_string(long_str, self.LENGTH_LIMIT) + + assert result.truncated is True + assert result.value == long_str[:5] + "..." + assert result.value_size == 10 # 10 chars + "..." + + def test_exact_limit_string(self, small_truncator: VariableTruncator): + """Test string exactly at limit.""" + exact_str = "1234567890" # Exactly 10 chars + result = small_truncator._truncate_string(exact_str, self.LENGTH_LIMIT) + assert result.value == "12345..." + assert result.truncated is True + assert result.value_size == 10 + + class TestArrayTruncation: + """Test array truncation functionality.""" + + @pytest.fixture + def small_truncator(self): + return VariableTruncator(array_element_limit=3, max_size_bytes=100) + + def test_small_array_no_truncation(self, small_truncator: VariableTruncator): + """Test that small arrays are not truncated.""" + small_array: list[object] = [1, 2] + result = small_truncator._truncate_array(small_array, 1000) + assert result.value == small_array + assert result.truncated is False + + def test_array_element_limit_truncation(self, small_truncator: VariableTruncator): + """Test that arrays over element limit are truncated.""" + large_array: list[object] = [1, 2, 3, 4, 5, 6] # Exceeds limit of 3 + result = small_truncator._truncate_array(large_array, 1000) + + assert result.truncated is True + assert result.value == [1, 2, 3] + + def test_array_size_budget_truncation(self, small_truncator: VariableTruncator): + """Test array truncation due to size budget constraints.""" + # Create array with strings that will exceed size budget + large_strings: list[object] = ["very long string " * 5, "another long string " * 5] + result = small_truncator._truncate_array(large_strings, 50) + + assert result.truncated is True + # Should have truncated the strings within the array + for item in result.value: + assert isinstance(item, str) + assert VariableTruncator.calculate_json_size(result.value) <= 50 + + def test_array_with_nested_objects(self, small_truncator): + """Test array truncation with nested objects.""" + nested_array = [ + {"name": "item1", "data": "some data"}, + {"name": "item2", "data": "more data"}, + {"name": "item3", "data": "even more data"}, + ] + result = small_truncator._truncate_array(nested_array, 30) + + assert isinstance(result.value, list) + assert len(result.value) <= 3 + for item in result.value: + assert isinstance(item, dict) + + class TestObjectTruncation: + """Test object truncation functionality.""" + + @pytest.fixture + def small_truncator(self): + return VariableTruncator(max_size_bytes=100) + + def test_small_object_no_truncation(self, small_truncator): + """Test that small objects are not truncated.""" + small_obj = {"a": 1, "b": 2} + result = small_truncator._truncate_object(small_obj, 1000) + assert result.value == small_obj + assert result.truncated is False + + def test_empty_object_no_truncation(self, small_truncator): + """Test that empty objects are not truncated.""" + empty_obj = {} + result = small_truncator._truncate_object(empty_obj, 100) + assert result.value == empty_obj + assert result.truncated is False + + def test_object_value_truncation(self, small_truncator): + """Test object truncation where values are truncated to fit budget.""" + obj_with_long_values = { + "key1": "very long string " * 10, + "key2": "another long string " * 10, + "key3": "third long string " * 10, + } + result = small_truncator._truncate_object(obj_with_long_values, 80) + + assert result.truncated is True + assert isinstance(result.value, dict) + + assert set(result.value.keys()).issubset(obj_with_long_values.keys()) + + # Values should be truncated if they exist + for key, value in result.value.items(): + assert isinstance(value, str) + original_value = obj_with_long_values[key] + # Value should be same or smaller + assert len(value) <= len(original_value) + + def test_object_key_dropping(self, small_truncator): + """Test object truncation where keys are dropped due to size constraints.""" + large_obj = {f"key{i:02d}": f"value{i}" for i in range(20)} + result = small_truncator._truncate_object(large_obj, 50) + + assert result.truncated is True + assert len(result.value) < len(large_obj) + + # Should maintain sorted key order + result_keys = list(result.value.keys()) + assert result_keys == sorted(result_keys) + + def test_object_with_nested_structures(self, small_truncator): + """Test object truncation with nested arrays and objects.""" + nested_obj = { + "simple": "value", + "array": [1, 2, 3, 4, 5], + "nested": {"inner": "data", "more": ["a", "b", "c"]}, + } + result = small_truncator._truncate_object(nested_obj, 60) + + assert isinstance(result.value, dict) + + class TestSegmentBasedTruncation: + """Test the main truncate method that works with Segments.""" + + @pytest.fixture + def truncator(self): + return VariableTruncator() + + @pytest.fixture + def small_truncator(self): + return VariableTruncator(string_length_limit=20, array_element_limit=3, max_size_bytes=200) + + def test_integer_segment_no_truncation(self, truncator): + """Test that integer segments are never truncated.""" + segment = IntegerSegment(value=12345) + result = truncator.truncate(segment) + + assert isinstance(result, TruncationResult) + assert result.truncated is False + assert result.result == segment + + def test_boolean_as_integer_segment(self, truncator): + """Test boolean values in IntegerSegment are converted to int.""" + segment = IntegerSegment(value=True) + result = truncator.truncate(segment) + + assert isinstance(result, TruncationResult) + assert result.truncated is False + assert isinstance(result.result, IntegerSegment) + assert result.result.value == 1 # True converted to 1 + + def test_float_segment_no_truncation(self, truncator): + """Test that float segments are never truncated.""" + segment = FloatSegment(value=123.456) + result = truncator.truncate(segment) + + assert isinstance(result, TruncationResult) + assert result.truncated is False + assert result.result == segment + + def test_none_segment_no_truncation(self, truncator): + """Test that None segments are never truncated.""" + segment = NoneSegment() + result = truncator.truncate(segment) + + assert isinstance(result, TruncationResult) + assert result.truncated is False + assert result.result == segment + + def test_file_segment_no_truncation(self, truncator, file): + """Test that file segments are never truncated.""" + file_segment = FileSegment(value=file) + result = truncator.truncate(file_segment) + assert result.result == file_segment + assert result.truncated is False + + def test_array_file_segment_no_truncation(self, truncator, file): + """Test that array file segments are never truncated.""" + + array_file_segment = ArrayFileSegment(value=[file] * 20) + result = truncator.truncate(array_file_segment) + assert result.result == array_file_segment + assert result.truncated is False + + def test_string_segment_small_no_truncation(self, truncator): + """Test small string segments are not truncated.""" + segment = StringSegment(value="hello world") + result = truncator.truncate(segment) + + assert isinstance(result, TruncationResult) + assert result.truncated is False + assert result.result == segment + + def test_string_segment_large_truncation(self, small_truncator): + """Test large string segments are truncated.""" + long_text = "this is a very long string that will definitely exceed the limit" + segment = StringSegment(value=long_text) + result = small_truncator.truncate(segment) + + assert isinstance(result, TruncationResult) + assert result.truncated is True + assert isinstance(result.result, StringSegment) + assert len(result.result.value) < len(long_text) + assert result.result.value.endswith("...") + + def test_array_segment_small_no_truncation(self, truncator): + """Test small array segments are not truncated.""" + from factories.variable_factory import build_segment + + segment = build_segment([1, 2, 3]) + result = truncator.truncate(segment) + + assert isinstance(result, TruncationResult) + assert result.truncated is False + assert result.result == segment + + def test_array_segment_large_truncation(self, small_truncator): + """Test large array segments are truncated.""" + from factories.variable_factory import build_segment + + large_array = list(range(10)) # Exceeds element limit of 3 + segment = build_segment(large_array) + result = small_truncator.truncate(segment) + + assert isinstance(result, TruncationResult) + assert result.truncated is True + assert isinstance(result.result, ArraySegment) + assert len(result.result.value) <= 3 + + def test_object_segment_small_no_truncation(self, truncator): + """Test small object segments are not truncated.""" + segment = ObjectSegment(value={"key": "value"}) + result = truncator.truncate(segment) + + assert isinstance(result, TruncationResult) + assert result.truncated is False + assert result.result == segment + + def test_object_segment_large_truncation(self, small_truncator): + """Test large object segments are truncated.""" + large_obj = {f"key{i}": f"very long value {i}" * 5 for i in range(5)} + segment = ObjectSegment(value=large_obj) + result = small_truncator.truncate(segment) + + assert isinstance(result, TruncationResult) + assert result.truncated is True + assert isinstance(result.result, ObjectSegment) + # Object should be smaller or equal than original + original_size = small_truncator.calculate_json_size(large_obj) + result_size = small_truncator.calculate_json_size(result.result.value) + assert result_size <= original_size + + def test_final_size_fallback_to_json_string(self, small_truncator): + """Test final fallback when truncated result still exceeds size limit.""" + # Create data that will still be large after initial truncation + large_nested_data = {"data": ["very long string " * 5] * 5, "more": {"nested": "content " * 20}} + segment = ObjectSegment(value=large_nested_data) + + # Use very small limit to force JSON string fallback + tiny_truncator = VariableTruncator(max_size_bytes=50) + result = tiny_truncator.truncate(segment) + + assert isinstance(result, TruncationResult) + assert result.truncated is True + assert isinstance(result.result, StringSegment) + # Should be JSON string with possible truncation + assert len(result.result.value) <= 53 # 50 + "..." = 53 + + def test_final_size_fallback_string_truncation(self, small_truncator): + """Test final fallback for string that still exceeds limit.""" + # Create very long string that exceeds string length limit + very_long_string = "x" * 6000 # Exceeds default string_length_limit of 5000 + segment = StringSegment(value=very_long_string) + + # Use small limit to test string fallback path + tiny_truncator = VariableTruncator(string_length_limit=100, max_size_bytes=50) + result = tiny_truncator.truncate(segment) + + assert isinstance(result, TruncationResult) + assert result.truncated is True + assert isinstance(result.result, StringSegment) + # Should be truncated due to string limit or final size limit + assert len(result.result.value) <= 1000 # Much smaller than original + + class TestEdgeCases: + """Test edge cases and error conditions.""" + + def test_empty_inputs(self): + """Test truncator with empty inputs.""" + truncator = VariableTruncator() + + # Empty string + result = truncator.truncate(StringSegment(value="")) + assert not result.truncated + assert result.result.value == "" + + # Empty array + from factories.variable_factory import build_segment + + result = truncator.truncate(build_segment([])) + assert not result.truncated + assert result.result.value == [] + + # Empty object + result = truncator.truncate(ObjectSegment(value={})) + assert not result.truncated + assert result.result.value == {} + + def test_zero_and_negative_limits(self): + """Test truncator behavior with zero or very small limits.""" + # Zero string limit + with pytest.raises(ValueError, match="string_length_limit should be greater than 3\\."): + VariableTruncator(string_length_limit=3) + + with pytest.raises(ValueError, match="array_element_limit should be greater than 0\\."): + VariableTruncator(array_element_limit=0) + + with pytest.raises(ValueError, match="max_size_bytes should be greater than 0\\."): + VariableTruncator(max_size_bytes=0) + + def test_unicode_and_special_characters(self): + """Test truncator with unicode and special characters.""" + truncator = VariableTruncator(string_length_limit=10) + + # Unicode characters + unicode_text = "你好世界你好世界你好世界" # Multi-byte UTF-8 characters + result = truncator.truncate(StringSegment(value=unicode_text)) + assert result.truncated is True + + # Special JSON characters + special_chars = '{"key": "value with \\"quotes\\" and \\n newlines"}' + result = truncator.truncate(StringSegment(value=special_chars)) + assert isinstance(result.result, StringSegment) + + class TestTruncateJsonPrimitives: + """Test _truncate_json_primitives method with different data types.""" + + @pytest.fixture + def truncator(self): + return VariableTruncator() + + def test_truncate_json_primitives_file_type(self, truncator, file): + """Test that File objects are handled correctly in _truncate_json_primitives.""" + # Test File object is returned as-is without truncation + result = truncator._truncate_json_primitives(file, 1000) + + assert result.value == file + assert result.truncated is False + # Size should be calculated correctly + expected_size = VariableTruncator.calculate_json_size(file) + assert result.value_size == expected_size + + def test_truncate_json_primitives_file_type_small_budget(self, truncator, file): + """Test that File objects are returned as-is even with small budget.""" + # Even with a small size budget, File objects should not be truncated + result = truncator._truncate_json_primitives(file, 10) + + assert result.value == file + assert result.truncated is False + + def test_truncate_json_primitives_file_type_in_array(self, truncator, file): + """Test File objects in arrays are handled correctly.""" + array_with_files = [file, file] + result = truncator._truncate_json_primitives(array_with_files, 1000) + + assert isinstance(result.value, list) + assert len(result.value) == 2 + assert result.value[0] == file + assert result.value[1] == file + assert result.truncated is False + + def test_truncate_json_primitives_file_type_in_object(self, truncator, file): + """Test File objects in objects are handled correctly.""" + obj_with_files = {"file1": file, "file2": file} + result = truncator._truncate_json_primitives(obj_with_files, 1000) + + assert isinstance(result.value, dict) + assert len(result.value) == 2 + assert result.value["file1"] == file + assert result.value["file2"] == file + assert result.truncated is False + + class TestIntegrationScenarios: + """Test realistic integration scenarios.""" + + def test_workflow_output_scenario(self): + """Test truncation of typical workflow output data.""" + truncator = VariableTruncator() + + workflow_data = { + "result": "success", + "data": { + "users": [ + {"id": 1, "name": "Alice", "email": "alice@example.com"}, + {"id": 2, "name": "Bob", "email": "bob@example.com"}, + ] + * 3, # Multiply to make it larger + "metadata": { + "count": 6, + "processing_time": "1.23s", + "details": "x" * 200, # Long string but not too long + }, + }, + } + + segment = ObjectSegment(value=workflow_data) + result = truncator.truncate(segment) + + assert isinstance(result, TruncationResult) + assert isinstance(result.result, (ObjectSegment, StringSegment)) + # Should handle complex nested structure appropriately + + def test_large_text_processing_scenario(self): + """Test truncation of large text data.""" + truncator = VariableTruncator(string_length_limit=100) + + large_text = "This is a very long text document. " * 20 # Make it larger than limit + + segment = StringSegment(value=large_text) + result = truncator.truncate(segment) + + assert isinstance(result, TruncationResult) + assert result.truncated is True + assert isinstance(result.result, StringSegment) + assert len(result.result.value) <= 103 # 100 + "..." + assert result.result.value.endswith("...") + + def test_mixed_data_types_scenario(self): + """Test truncation with mixed data types in complex structure.""" + truncator = VariableTruncator(string_length_limit=30, array_element_limit=3, max_size_bytes=300) + + mixed_data = { + "strings": ["short", "medium length", "very long string " * 3], + "numbers": [1, 2.5, 999999], + "booleans": [True, False, True], + "nested": { + "more_strings": ["nested string " * 2], + "more_numbers": list(range(5)), + "deep": {"level": 3, "content": "deep content " * 3}, + }, + "nulls": [None, None], + } + + segment = ObjectSegment(value=mixed_data) + result = truncator.truncate(segment) + + assert isinstance(result, TruncationResult) + assert result.truncated is True + assert isinstance(result.result, ObjectSegment) + # Verify the result is smaller or equal than original + original_size = truncator.calculate_json_size(mixed_data) + result_size = truncator.calculate_json_size(result.result.value) + assert result_size <= original_size + + def test_file_and_array_file_variable_mapping(self, file): + truncator = VariableTruncator(string_length_limit=30, array_element_limit=3, max_size_bytes=300) + + mapping = {"array_file": [file]} + truncated_mapping, truncated = truncator.truncate_variable_mapping(mapping) + assert truncated is False + assert truncated_mapping == mapping + + @staticmethod + def test_dummy_variable_truncator_methods(): + """Test DummyVariableTruncator methods work correctly.""" + truncator = DummyVariableTruncator() + + # Test truncate_variable_mapping + test_data: dict[str, Any] = { + "key1": "value1", + "key2": ["item1", "item2"], + "large_array": list(range(2000)), + } + result, is_truncated = truncator.truncate_variable_mapping(test_data) + + assert result == test_data + assert not is_truncated + + # Test truncate method + segment = StringSegment(value="test string") + result = truncator.truncate(segment) + assert isinstance(result, TruncationResult) + assert result.result == segment assert result.truncated is False + + segment = ArrayNumberSegment(value=list(range(2000))) + result = truncator.truncate(segment) + assert isinstance(result, TruncationResult) + assert result.result == segment + assert result.truncated is False + + # --------------------------------------------------------------------------- + # Regression tests for langgenius/dify#39218. + # + # Before the fix, ``_truncate_array`` had a "Dirty fix" branch that + # unconditionally appended every ``File`` element to ``truncated_value`` + # *before* the count cap and the byte-budget check, and *before* + # ``used_size`` was ever incremented. That made ``list[File]`` arrays: + # 1. uncapped by ``array_element_limit``, + # 2. uncounted against ``max_size_bytes``, and + # 3. always reported ``truncated=False``. + # The fix routes ``File`` through ``_truncate_json_primitives``'s dedicated + # ``File`` branch, which returns the file as-is with its real serialized + # size, while preserving the count cap and the byte budget. + # --------------------------------------------------------------------------- + + class TestFileArrayTruncationRegression39218: + """``list[File]`` must respect ``array_element_limit`` and the byte budget.""" + + @pytest.fixture + def truncator(self) -> VariableTruncator: + return VariableTruncator( + array_element_limit=3, + max_size_bytes=1000, + string_length_limit=50, + ) + + @staticmethod + def _make_file(name: str = "f") -> File: + return File( + id=name, + type=FileType.DOCUMENT, + transfer_method=FileTransferMethod.REMOTE_URL, + remote_url=f"https://example.com/{name}.txt", + filename=f"{name}.txt", + extension=".txt", + mime_type="text/plain", + size=1024, + ) + + def test_file_array_respects_element_count_cap(self, truncator: VariableTruncator) -> None: + # Use a target_size larger than ``count * file_size`` so the byte + # budget never binds — only the count cap should fire. + # Each File serializes to ~237 bytes; 3 files = ~713 bytes. + files = [self._make_file(f"f{i}") for i in range(500)] + + result = truncator._truncate_array(files, target_size=10_000_000) + + # Before the fix, all 500 File entries survived (``len(value)==500``, + # ``truncated==False``). After the fix, the array is capped at + # ``array_element_limit=3`` and ``truncated`` flips to True. + assert len(result.value) == 3 + assert result.truncated is True + + def test_file_array_reports_real_used_size(self, truncator: VariableTruncator) -> None: + # Large budget so the count cap fires before the byte budget does. + files = [self._make_file(f"f{i}") for i in range(500)] + + result = truncator._truncate_array(files, target_size=10_000_000) + + # Before the fix, ``used_size`` for a File array was the empty-array + # baseline of 2 bytes (``[]``), regardless of how many File entries + # actually returned. After the fix, ``used_size`` reflects the real + # serialized size of the returned ``File`` payload. + assert result.value_size > 100 + assert result.truncated is True + + def test_file_array_respects_byte_budget(self, truncator: VariableTruncator) -> None: + # Use a small ``target_size`` so the byte budget is the binding + # constraint. Each File serializes to ~237 bytes, so even one File + # blows the 200-byte budget. + files = [self._make_file(f"f{i}") for i in range(50)] + + result = truncator._truncate_array(files, target_size=200) + + # Before the fix, all 50 File entries survived and ``used_size`` + # reported ``2`` (the empty-array baseline). After the fix, the + # loop sees the File payload: ``value_size`` reflects the real + # serialized size, and the loop stops after the first File because + # adding the next one would exceed ``target_size``. + assert len(result.value) == 1 + assert result.value_size > 100 # the File's real serialized size + assert result.value_size <= 250 # in the ballpark of the budget + + def test_mixed_array_counts_files_toward_cap(self, truncator: VariableTruncator) -> None: + mixed: list[object] = [ + self._make_file("f0"), + "a", + self._make_file("f1"), + "b", + self._make_file("f2"), + "c", + self._make_file("f3"), + "d", + ] + + result = truncator._truncate_array(mixed, target_size=10_000_000) + + # 8 items, cap of 3 → exactly 3 items. Files and primitives are + # counted together toward the cap. + assert len(result.value) == 3 + assert result.truncated is True + + def test_single_file_in_array_is_preserved(self, truncator: VariableTruncator) -> None: + result = truncator._truncate_array([self._make_file("only")], target_size=10_000_000) + + # The File itself is not truncated — the dedicated ``File`` branch + # in ``_truncate_json_primitives`` returns the file untouched. Only + # the array-shape accounting changes. + assert len(result.value) == 1 + assert isinstance(result.value[0], File) + assert result.value[0].id == "only" + assert result.truncated is False + + +class TestVariableTruncatorAdditionalCases: + class TestVariableTruncatorAdditionalBehavior: + def test_default_should_use_dify_config_limits(self, monkeypatch: pytest.MonkeyPatch) -> None: + apply_config_overrides( + monkeypatch, + WORKFLOW_VARIABLE_TRUNCATION_MAX_SIZE=111, + WORKFLOW_VARIABLE_TRUNCATION_ARRAY_LENGTH=7, + WORKFLOW_VARIABLE_TRUNCATION_STRING_LENGTH=33, + ) + + truncator = VariableTruncator.default() + + assert truncator._max_size_bytes == 111 + assert truncator._array_element_limit == 7 + assert truncator._string_length_limit == 33 + + def test_truncate_variable_mapping_should_mark_over_budget_keys_with_ellipsis(self) -> None: + truncator = VariableTruncator(max_size_bytes=5) + mapping = {"very_long_key": "value"} + + result, truncated = truncator.truncate_variable_mapping(mapping) + + assert result == {"very_long_key": "..."} + assert truncated is True + + def test_truncate_variable_mapping_should_handle_segment_values(self) -> None: + truncator = VariableTruncator(max_size_bytes=100) + mapping = {"seg": StringSegment(value="hello")} + + result, truncated = truncator.truncate_variable_mapping(mapping) + + assert isinstance(result["seg"], StringSegment) + assert result["seg"].value == "hello" + assert truncated is False + + @pytest.mark.parametrize( + ("value", "expected"), + [ + (None, False), + (True, False), + (1, False), + (1.5, False), + ("x", True), + ({"k": "v"}, True), + ], + ) + def test_json_value_needs_truncation_should_match_expected_rules( + self, + value: object, + expected: bool, + ) -> None: + result = VariableTruncator._json_value_needs_truncation(value) + assert result is expected + + def test_truncate_should_use_string_fallback_when_truncated_value_size_exceeds_limit( + self, + monkeypatch: pytest.MonkeyPatch, + ) -> None: + truncator = VariableTruncator(max_size_bytes=10) + forced_result = truncator_module._PartResult( + value=StringSegment(value="this is too long"), + value_size=100, + truncated=True, + ) + monkeypatch.setattr(truncator, "_truncate_segment", lambda *_args, **_kwargs: forced_result) + + result = truncator.truncate(StringSegment(value="input")) + + assert result.truncated is True + assert isinstance(result.result, StringSegment) + assert not result.result.value.startswith('"') + + def test_truncate_segment_should_raise_assertion_for_unexpected_truncatable_segment( + self, + monkeypatch: pytest.MonkeyPatch, + ) -> None: + truncator = VariableTruncator() + monkeypatch.setattr(VariableTruncator, "_segment_need_truncation", lambda _segment: True) + + with pytest.raises(AssertionError): + truncator._truncate_segment(IntegerSegment(value=1), 10) + + def test_calculate_json_size_should_unwrap_segment_values(self) -> None: + segment = StringSegment(value="abc") + + size = VariableTruncator.calculate_json_size(segment) + + assert size == VariableTruncator.calculate_json_size("abc") + + def test_calculate_json_size_should_handle_updated_variable_instances(self) -> None: + updated = UpdatedVariable(name="n", selector=["node", "var"], value_type=SegmentType.STRING, new_value="v") + + size = VariableTruncator.calculate_json_size(updated) + + assert size > 0 + + def test_maybe_qa_structure_should_validate_shape(self) -> None: + assert VariableTruncator._maybe_qa_structure({"qa_chunks": []}) is True + assert VariableTruncator._maybe_qa_structure({"qa_chunks": "not-list"}) is False + assert VariableTruncator._maybe_qa_structure({}) is False + + def test_maybe_parent_child_structure_should_validate_shape(self) -> None: + assert ( + VariableTruncator._maybe_parent_child_structure({"parent_mode": "full", "parent_child_chunks": []}) + is True + ) + assert ( + VariableTruncator._maybe_parent_child_structure({"parent_mode": 1, "parent_child_chunks": []}) is False + ) + assert ( + VariableTruncator._maybe_parent_child_structure({"parent_mode": "full", "parent_child_chunks": "bad"}) + is False + ) + + def test_truncate_object_should_truncate_segment_values_inside_object(self) -> None: + truncator = VariableTruncator(string_length_limit=8, max_size_bytes=30) + mapping = {"s": StringSegment(value="long-content")} + + result = truncator._truncate_object(mapping, 20) + + assert result.truncated is True + assert isinstance(result.value["s"], StringSegment) + + def test_truncate_json_primitives_should_handle_updated_variable_input(self) -> None: + truncator = VariableTruncator(max_size_bytes=100) + updated = UpdatedVariable(name="n", selector=["node", "var"], value_type=SegmentType.STRING, new_value="v") + + result = truncator._truncate_json_primitives(updated, 100) + + assert isinstance(result.value, dict) + + def test_truncate_json_primitives_should_raise_assertion_for_unsupported_value_type(self) -> None: + truncator = VariableTruncator() + + with pytest.raises(AssertionError): + truncator._truncate_json_primitives(object(), 100) # type: ignore[arg-type] + + def test_truncate_should_apply_json_string_fallback_for_large_non_string_segment( + self, + monkeypatch: pytest.MonkeyPatch, + ) -> None: + truncator = VariableTruncator(max_size_bytes=10) + forced_segment = ObjectSegment(value={"k": "v"}) + forced_result = truncator_module._PartResult(value=forced_segment, value_size=100, truncated=True) + monkeypatch.setattr(truncator, "_truncate_segment", lambda *_args, **_kwargs: forced_result) + + result = truncator.truncate(ObjectSegment(value={"a": "b"})) + + assert result.truncated is True + assert isinstance(result.result, StringSegment) diff --git a/api/tests/unit_tests/services/test_variable_truncator_additional.py b/api/tests/unit_tests/services/test_variable_truncator_additional.py deleted file mode 100644 index 3678e985673..00000000000 --- a/api/tests/unit_tests/services/test_variable_truncator_additional.py +++ /dev/null @@ -1,156 +0,0 @@ -import pytest - -from graphon.nodes.variable_assigner.common.helpers import UpdatedVariable -from graphon.variables.segments import IntegerSegment, ObjectSegment, StringSegment -from graphon.variables.types import SegmentType -from services import variable_truncator as truncator_module -from services.variable_truncator import VariableTruncator -from tests.unit_tests.config_override import apply_config_overrides - - -class TestVariableTruncatorAdditionalBehavior: - def test_default_should_use_dify_config_limits(self, monkeypatch: pytest.MonkeyPatch) -> None: - apply_config_overrides( - monkeypatch, - WORKFLOW_VARIABLE_TRUNCATION_MAX_SIZE=111, - WORKFLOW_VARIABLE_TRUNCATION_ARRAY_LENGTH=7, - WORKFLOW_VARIABLE_TRUNCATION_STRING_LENGTH=33, - ) - - truncator = VariableTruncator.default() - - assert truncator._max_size_bytes == 111 - assert truncator._array_element_limit == 7 - assert truncator._string_length_limit == 33 - - def test_truncate_variable_mapping_should_mark_over_budget_keys_with_ellipsis(self) -> None: - truncator = VariableTruncator(max_size_bytes=5) - mapping = {"very_long_key": "value"} - - result, truncated = truncator.truncate_variable_mapping(mapping) - - assert result == {"very_long_key": "..."} - assert truncated is True - - def test_truncate_variable_mapping_should_handle_segment_values(self) -> None: - truncator = VariableTruncator(max_size_bytes=100) - mapping = {"seg": StringSegment(value="hello")} - - result, truncated = truncator.truncate_variable_mapping(mapping) - - assert isinstance(result["seg"], StringSegment) - assert result["seg"].value == "hello" - assert truncated is False - - @pytest.mark.parametrize( - ("value", "expected"), - [ - (None, False), - (True, False), - (1, False), - (1.5, False), - ("x", True), - ({"k": "v"}, True), - ], - ) - def test_json_value_needs_truncation_should_match_expected_rules( - self, - value: object, - expected: bool, - ) -> None: - result = VariableTruncator._json_value_needs_truncation(value) - assert result is expected - - def test_truncate_should_use_string_fallback_when_truncated_value_size_exceeds_limit( - self, - monkeypatch: pytest.MonkeyPatch, - ) -> None: - truncator = VariableTruncator(max_size_bytes=10) - forced_result = truncator_module._PartResult( - value=StringSegment(value="this is too long"), - value_size=100, - truncated=True, - ) - monkeypatch.setattr(truncator, "_truncate_segment", lambda *_args, **_kwargs: forced_result) - - result = truncator.truncate(StringSegment(value="input")) - - assert result.truncated is True - assert isinstance(result.result, StringSegment) - assert not result.result.value.startswith('"') - - def test_truncate_segment_should_raise_assertion_for_unexpected_truncatable_segment( - self, - monkeypatch: pytest.MonkeyPatch, - ) -> None: - truncator = VariableTruncator() - monkeypatch.setattr(VariableTruncator, "_segment_need_truncation", lambda _segment: True) - - with pytest.raises(AssertionError): - truncator._truncate_segment(IntegerSegment(value=1), 10) - - def test_calculate_json_size_should_unwrap_segment_values(self) -> None: - segment = StringSegment(value="abc") - - size = VariableTruncator.calculate_json_size(segment) - - assert size == VariableTruncator.calculate_json_size("abc") - - def test_calculate_json_size_should_handle_updated_variable_instances(self) -> None: - updated = UpdatedVariable(name="n", selector=["node", "var"], value_type=SegmentType.STRING, new_value="v") - - size = VariableTruncator.calculate_json_size(updated) - - assert size > 0 - - def test_maybe_qa_structure_should_validate_shape(self) -> None: - assert VariableTruncator._maybe_qa_structure({"qa_chunks": []}) is True - assert VariableTruncator._maybe_qa_structure({"qa_chunks": "not-list"}) is False - assert VariableTruncator._maybe_qa_structure({}) is False - - def test_maybe_parent_child_structure_should_validate_shape(self) -> None: - assert ( - VariableTruncator._maybe_parent_child_structure({"parent_mode": "full", "parent_child_chunks": []}) is True - ) - assert VariableTruncator._maybe_parent_child_structure({"parent_mode": 1, "parent_child_chunks": []}) is False - assert ( - VariableTruncator._maybe_parent_child_structure({"parent_mode": "full", "parent_child_chunks": "bad"}) - is False - ) - - def test_truncate_object_should_truncate_segment_values_inside_object(self) -> None: - truncator = VariableTruncator(string_length_limit=8, max_size_bytes=30) - mapping = {"s": StringSegment(value="long-content")} - - result = truncator._truncate_object(mapping, 20) - - assert result.truncated is True - assert isinstance(result.value["s"], StringSegment) - - def test_truncate_json_primitives_should_handle_updated_variable_input(self) -> None: - truncator = VariableTruncator(max_size_bytes=100) - updated = UpdatedVariable(name="n", selector=["node", "var"], value_type=SegmentType.STRING, new_value="v") - - result = truncator._truncate_json_primitives(updated, 100) - - assert isinstance(result.value, dict) - - def test_truncate_json_primitives_should_raise_assertion_for_unsupported_value_type(self) -> None: - truncator = VariableTruncator() - - with pytest.raises(AssertionError): - truncator._truncate_json_primitives(object(), 100) # type: ignore[arg-type] - - def test_truncate_should_apply_json_string_fallback_for_large_non_string_segment( - self, - monkeypatch: pytest.MonkeyPatch, - ) -> None: - truncator = VariableTruncator(max_size_bytes=10) - forced_segment = ObjectSegment(value={"k": "v"}) - forced_result = truncator_module._PartResult(value=forced_segment, value_size=100, truncated=True) - monkeypatch.setattr(truncator, "_truncate_segment", lambda *_args, **_kwargs: forced_result) - - result = truncator.truncate(ObjectSegment(value={"a": "b"})) - - assert result.truncated is True - assert isinstance(result.result, StringSegment) diff --git a/api/tests/unit_tests/services/test_vector_space_admission_service.py b/api/tests/unit_tests/services/test_vector_space_admission_service.py index e4283461151..903db9dff84 100644 --- a/api/tests/unit_tests/services/test_vector_space_admission_service.py +++ b/api/tests/unit_tests/services/test_vector_space_admission_service.py @@ -27,7 +27,7 @@ from services.vector_space_admission_service import ( get_vector_space_admission_error_fields, parse_vector_space_estimate_limits, ) -from tests.unit_tests.model_factories import make_dataset +from tests.helpers.unit.model_factories import make_dataset _MEBIBYTE = 1024 * 1024 _ESTIMATE_LIMITS = "sandbox:60,professional:6400,team:25600" diff --git a/api/tests/unit_tests/services/test_weave_tracing_config_gateway.py b/api/tests/unit_tests/services/test_weave_tracing_config_gateway.py deleted file mode 100644 index 918d30bc8cc..00000000000 --- a/api/tests/unit_tests/services/test_weave_tracing_config_gateway.py +++ /dev/null @@ -1,167 +0,0 @@ -"""Config reads and writes must not construct Weave's stateful legacy exporter.""" - -from collections.abc import Callable, Iterator, Sequence -from dataclasses import dataclass, field -from typing import Never -from uuid import uuid4 - -import pytest -import wandb -import weave -from dify_trace_weave.configuration import WeaveConfigurationError -from dify_trace_weave.weave_trace import WeaveDataTrace -from sqlalchemy.orm import Session, sessionmaker - -from core.helper import ssrf_proxy -from core.ops import ops_trace_manager -from core.ops.ops_trace_manager import OpsTraceManager -from machinery.context import RequestContext -from repositories.app.tracing_config_repository import SQLAlchemyAppTracingConfigRepository -from services.app_tracing_config_gateway import OpsTraceManagerGateway -from services.app_tracing_config_service import AppTracingConfigService, AppTracingConfigVerificationFailedError -from tests.unit_tests.model_factories import make_app -from tests.unit_tests.services.tracing.test_weave_configuration import ReadAPI - - -@dataclass -class GatewayEnvironment: - gateway: OpsTraceManagerGateway = field(default_factory=OpsTraceManagerGateway) - api: ReadAPI = field(default_factory=ReadAPI) - forbidden_calls: list[str] = field(default_factory=list) - - -@pytest.fixture -def environment(monkeypatch: pytest.MonkeyPatch) -> Iterator[GatewayEnvironment]: - environment = GatewayEnvironment() - - def encrypt(_tenant: str, value: str) -> str: - return "test-cipher:" + value - - def decrypt(_tenant: str, values: Sequence[str]) -> list[str]: - assert all(value.startswith("test-cipher:") for value in values) - return [value.removeprefix("test-cipher:") for value in values] - - def mask(_value: str) -> str: - return "******" - - def forbidden(name: str) -> Callable[..., Never]: - def block(*_args: object, **_kwargs: object) -> Never: - environment.forbidden_calls.append(name) - raise AssertionError("Configuration must not initialize a stateful exporter") - - return block - - monkeypatch.setattr(ops_trace_manager, "encrypt_token", encrypt) - monkeypatch.setattr(ops_trace_manager, "batch_decrypt_token", decrypt) - monkeypatch.setattr(ops_trace_manager, "obfuscated_token", mask) - monkeypatch.setattr(ssrf_proxy, "post", environment.api.post) - monkeypatch.setattr(WeaveDataTrace, "__init__", forbidden("legacy_constructor")) - monkeypatch.setattr(wandb, "login", forbidden("wandb_login")) - monkeypatch.setattr(weave, "init", forbidden("weave_init")) - yield environment - assert environment.forbidden_calls == [] - - -def test_create_persists_resolved_default_entity_and_canonical_project(environment: GatewayEnvironment) -> None: - submitted = {"api_key": "synthetic-key", "project": " demo "} - stored = environment.gateway.prepare_new_config( - workspace_id=str(uuid4()), tracing_provider="weave", tracing_config=submitted - ) - assert stored["entity"] == "default-team" - assert stored["project"] == "demo" - assert stored["project_url"] == "https://wandb.ai/default-team/demo" - assert stored["api_key"] == "test-cipher:synthetic-key" - assert submitted == {"api_key": "synthetic-key", "project": " demo "} - assert len(environment.api.calls) == 3 - - -def test_patch_persists_verified_destination_with_the_existing_ciphertext( - environment: GatewayEnvironment, sqlite_session_factory: sessionmaker[Session] -) -> None: - workspace, app_id = str(uuid4()), str(uuid4()) - with sqlite_session_factory.begin() as session: - session.add(make_app(app_id=app_id, tenant_id=workspace)) - store = SQLAlchemyAppTracingConfigRepository(session_factory=sqlite_session_factory) - service = AppTracingConfigService(configs=store, provider=environment.gateway) - context = RequestContext("request", None, "account", workspace) - service.create(context, app_id, "weave", {"api_key": "synthetic-key", "entity": "old", "project": "original"}) - before = store.get(workspace_id=workspace, app_id=app_id, tracing_provider="weave") - assert before is not None - assert before.tracing_config is not None - - submitted = {"api_key": "******", "project": "new-team/changed"} - service.update(context, app_id, "weave", submitted) - after = store.get(workspace_id=workspace, app_id=app_id, tracing_provider="weave") - assert after is not None - assert after.tracing_config is not None - assert after.tracing_config["entity"] == "new-team" - assert after.tracing_config["project"] == "changed" - assert after.tracing_config["project_url"] == "https://wandb.ai/new-team/changed" - assert after.tracing_config["api_key"] == before.tracing_config["api_key"] == "test-cipher:synthetic-key" - assert before.tracing_config["project_url"] == "https://wandb.ai/old/original" - assert submitted == {"api_key": "******", "project": "new-team/changed"} - assert len(environment.api.calls) == 4 - - presented = service.get(context, app_id, "weave") - assert presented is not None - assert presented.tracing_config is not None - assert presented.tracing_config["api_key"] == "******" - assert presented.tracing_config["project_url"] == after.tracing_config["project_url"] - environment.gateway.validate_saved_config( - workspace_id=workspace, tracing_provider="weave", tracing_config=after.tracing_config - ) - assert len(environment.api.calls) == 4 - - -@pytest.mark.parametrize( - ("entity", "project", "host", "expected"), - [ - ("team", "demo", None, "https://wandb.ai/team/demo"), - (None, "team/demo", None, "https://wandb.ai/team/demo"), - (None, "demo", None, "https://wandb.ai"), - (None, "demo", "http://private.test/wandb", "http://private.test/wandb"), - ("team", "a & b", "https://private.test/wandb", "https://private.test/wandb/team/a%20%26%20b"), - ], -) -def test_get_computes_project_or_console_link_without_vendor_requests( - environment: GatewayEnvironment, entity: str | None, project: str, host: str | None, expected: str -) -> None: - workspace = str(uuid4()) - encrypted = OpsTraceManager.encrypt_tracing_config( - workspace, "weave", {"api_key": "synthetic-key", "entity": entity, "project": project, "host": host} - ) - encrypted["project_url"] = "https://stale.test/previous-project" - result = environment.gateway.present_config( - workspace_id=workspace, tracing_provider="weave", tracing_config=encrypted - ) - assert result["project_url"] == expected - assert result["api_key"] == "******" - assert environment.api.calls == [] - - -def test_inaccessible_project_is_a_sanitized_domain_error_and_is_not_saved( - environment: GatewayEnvironment, sqlite_session_factory: sessionmaker[Session], caplog: pytest.LogCaptureFixture -) -> None: - workspace, app_id = str(uuid4()), str(uuid4()) - with sqlite_session_factory.begin() as session: - session.add(make_app(app_id=app_id, tenant_id=workspace)) - environment.api.project_exists = False - store = SQLAlchemyAppTracingConfigRepository(session_factory=sqlite_session_factory) - service = AppTracingConfigService(configs=store, provider=environment.gateway) - with pytest.raises(AppTracingConfigVerificationFailedError) as caught: - service.create( - RequestContext("request", None, "account", workspace), - app_id, - "weave", - { - "api_key": "synthetic-private-key", - "entity": "team", - "project": "missing", - }, - ) - assert isinstance(caught.value.__cause__, WeaveConfigurationError) - assert caught.value.__cause__.code == "project_not_found_or_inaccessible" - assert store.get(workspace_id=workspace, app_id=app_id, tracing_provider="weave") is None - assert "synthetic-private-key" not in caplog.text - assert "project_not_found_or_inaccessible" in caplog.text - assert len(environment.api.calls) == 1 diff --git a/api/tests/unit_tests/services/test_web_app_runtime_query_service.py b/api/tests/unit_tests/services/test_web_app_runtime_query_service.py index 57a36ac6635..5fab2c575ae 100644 --- a/api/tests/unit_tests/services/test_web_app_runtime_query_service.py +++ b/api/tests/unit_tests/services/test_web_app_runtime_query_service.py @@ -16,7 +16,7 @@ from services.web_app_runtime_query_service import ( WebAppRuntimeRecord, WebAppRuntimeUnavailableError, ) -from tests.unit_tests.model_factories import make_upload_file +from tests.helpers.unit.model_factories import make_upload_file _FILES_URL = "https://files.example.com" diff --git a/api/tests/unit_tests/services/test_webhook_service.py b/api/tests/unit_tests/services/test_webhook_service.py deleted file mode 100644 index 09954f977b0..00000000000 --- a/api/tests/unit_tests/services/test_webhook_service.py +++ /dev/null @@ -1,965 +0,0 @@ -import logging -from io import BytesIO -from types import SimpleNamespace -from unittest.mock import MagicMock, patch - -import pytest -from flask import Flask -from sqlalchemy.engine import Engine -from sqlalchemy.orm import Session -from werkzeug.datastructures import FileStorage - -from models.enums import AppTriggerStatus, AppTriggerType, EndUserType -from models.model import App, AppMode, EndUser -from models.tools import ToolFile -from models.trigger import AppTrigger, WorkflowWebhookTrigger -from models.workflow import Workflow -from services.errors.app import QuotaExceededError -from services.trigger import webhook_service as webhook_service_module -from services.trigger.webhook_service import WebhookService -from tests.unit_tests.model_factories import make_app, make_end_user, make_workflow - - -class _EndUserServiceStub: - def __init__(self, result: EndUser | Exception) -> None: - self._result = result - self.calls: list[tuple[EndUserType, str, str, str | None]] = [] - - def get_or_create_end_user_by_type( - self, - type: EndUserType, - tenant_id: str, - app_id: str, - user_id: str | None = None, - ) -> EndUser: - self.calls.append((type, tenant_id, app_id, user_id)) - if isinstance(self._result, Exception): - raise self._result - return self._result - - -def _webhook_trigger( - *, - webhook_id: str = "webhook-123", - tenant_id: str = "tenant-123", - app_id: str = "app-123", - node_id: str = "node-123", - created_by: str = "account-123", -) -> WorkflowWebhookTrigger: - return WorkflowWebhookTrigger( - webhook_id=webhook_id, - tenant_id=tenant_id, - app_id=app_id, - node_id=node_id, - created_by=created_by, - ) - - -def _tool_file() -> ToolFile: - tool_file = ToolFile( - user_id="user-123", - tenant_id="test_tenant", - conversation_id=None, - file_key="webhook/test.txt", - mimetype="text/plain", - name="test.txt", - size=7, - ) - tool_file.id = "test_file_id" - return tool_file - - -def _workflow( - *, - workflow_id: str = "workflow-123", - tenant_id: str = "tenant-123", - app_id: str = "app-123", - version: str = Workflow.VERSION_DRAFT, -) -> Workflow: - return make_workflow( - workflow_id=workflow_id, - tenant_id=tenant_id, - app_id=app_id, - version=version, - created_by="account-123", - ) - - -def _app(*, tenant_id: str = "tenant-123", app_id: str = "app-123") -> App: - return make_app( - app_id=app_id, - tenant_id=tenant_id, - name="Webhook App", - mode=AppMode.WORKFLOW, - icon_type=None, - max_active_requests=0, - ) - - -def _app_trigger(*, status: AppTriggerStatus = AppTriggerStatus.ENABLED) -> AppTrigger: - return AppTrigger( - tenant_id="tenant-123", - app_id="app-123", - node_id="node-123", - trigger_type=AppTriggerType.TRIGGER_WEBHOOK, - title="Webhook", - status=status, - ) - - -def _end_user() -> EndUser: - return make_end_user( - end_user_id="end-user-123", - tenant_id="tenant-123", - app_id="app-123", - end_user_type=EndUserType.TRIGGER, - session_id="webhook-session", - ) - - -class TestWebhookServiceLookup: - def test_debug_lookup_scopes_draft_workflow_to_trigger_owner( - self, - sqlite_session: Session, - sqlite_engine: Engine, - monkeypatch: pytest.MonkeyPatch, - ) -> None: - target_workflow = _workflow() - tenant_decoy = _workflow( - workflow_id="workflow-decoy", - tenant_id="tenant-other", - app_id="app-123", - ) - sqlite_session.add_all([_webhook_trigger(), target_workflow, tenant_decoy]) - sqlite_session.commit() - monkeypatch.setattr(webhook_service_module, "db", SimpleNamespace(engine=sqlite_engine)) - node_config = {"id": "node-123", "data": {}} - - with patch.object(Workflow, "get_node_config_by_id", autospec=True, return_value=node_config) as get_node: - webhook_trigger, workflow, result_node_config = WebhookService.get_webhook_trigger_and_workflow( - "webhook-123", is_debug=True - ) - - assert webhook_trigger.tenant_id == "tenant-123" - assert workflow.id == target_workflow.id - assert result_node_config is node_config - get_node.assert_called_once_with(workflow, "node-123") - - def test_published_lookup_uses_persisted_owner_chain( - self, - sqlite_session: Session, - sqlite_engine: Engine, - monkeypatch: pytest.MonkeyPatch, - ) -> None: - sqlite_session.add_all([_webhook_trigger(), _app_trigger(), _app()]) - sqlite_session.commit() - monkeypatch.setattr(webhook_service_module, "db", SimpleNamespace(engine=sqlite_engine)) - published_workflow = _workflow(version="published") - node_config = {"id": "node-123", "data": {}} - - with ( - patch.object(webhook_service_module, "WorkflowService") as workflow_service_class, - patch.object(Workflow, "get_node_config_by_id", autospec=True, return_value=node_config), - ): - get_published_workflow = workflow_service_class.return_value.get_published_workflow - get_published_workflow.return_value = published_workflow - webhook_trigger, workflow, result_node_config = WebhookService.get_webhook_trigger_and_workflow( - "webhook-123" - ) - - assert webhook_trigger.app_id == "app-123" - assert workflow is published_workflow - assert result_node_config is node_config - persisted_app = get_published_workflow.call_args.args[0] - assert isinstance(persisted_app, App) - assert persisted_app.id == "app-123" - assert persisted_app.tenant_id == "tenant-123" - - def test_published_lookup_rejects_cross_tenant_app_decoy( - self, - sqlite_session: Session, - sqlite_engine: Engine, - monkeypatch: pytest.MonkeyPatch, - ) -> None: - sqlite_session.add_all( - [ - _webhook_trigger(), - _app_trigger(), - _app(tenant_id="tenant-other"), - ] - ) - sqlite_session.commit() - monkeypatch.setattr(webhook_service_module, "db", SimpleNamespace(engine=sqlite_engine)) - - with pytest.raises(ValueError, match="App not found"): - WebhookService.get_webhook_trigger_and_workflow("webhook-123") - - @pytest.mark.parametrize( - "status", - [AppTriggerStatus.DISABLED, AppTriggerStatus.UNAUTHORIZED], - ) - def test_published_lookup_rejects_inactive_trigger( - self, - status: AppTriggerStatus, - sqlite_session: Session, - sqlite_engine: Engine, - monkeypatch: pytest.MonkeyPatch, - ) -> None: - sqlite_session.add_all([_webhook_trigger(), _app_trigger(status=status)]) - sqlite_session.commit() - monkeypatch.setattr(webhook_service_module, "db", SimpleNamespace(engine=sqlite_engine)) - - with pytest.raises(ValueError, match="disabled"): - WebhookService.get_webhook_trigger_and_workflow("webhook-123") - - def test_published_lookup_reports_rate_limited_trigger( - self, - sqlite_session: Session, - sqlite_engine: Engine, - monkeypatch: pytest.MonkeyPatch, - ) -> None: - sqlite_session.add_all([_webhook_trigger(), _app_trigger(status=AppTriggerStatus.RATE_LIMITED)]) - sqlite_session.commit() - monkeypatch.setattr(webhook_service_module, "db", SimpleNamespace(engine=sqlite_engine)) - - with pytest.raises(QuotaExceededError): - WebhookService.get_webhook_trigger_and_workflow("webhook-123") - - -class TestWebhookServiceUnit: - """Webhook business-logic tests with isolated sessions where the service owns their lifecycle.""" - - def test_trigger_workflow_execution_propagates_quota_error_without_error_log( - self, - caplog: pytest.LogCaptureFixture, - monkeypatch: pytest.MonkeyPatch, - sqlite_engine: Engine, - ) -> None: - """Quota failures refund the charge and close the real service-owned session.""" - monkeypatch.setattr(webhook_service_module, "db", SimpleNamespace(engine=sqlite_engine)) - webhook_trigger = _webhook_trigger() - workflow = _workflow() - quota_charge = MagicMock() - quota_error = QuotaExceededError(feature="workflow", tenant_id="tenant-123", required=1) - - caplog.set_level(logging.INFO) - with ( - patch("services.trigger.webhook_service.QuotaService.reserve", return_value=quota_charge), - patch( - "services.trigger.webhook_service.AsyncWorkflowService.trigger_workflow_async", - side_effect=quota_error, - ) as mock_trigger_workflow_async, - ): - with pytest.raises(QuotaExceededError) as exc_info: - WebhookService.trigger_workflow_execution( - webhook_trigger, - {"body": {}, "headers": {}, "query_params": {}, "files": {}, "method": "POST"}, - workflow, - end_users=_EndUserServiceStub(_end_user()), - ) - - assert exc_info.value is quota_error - quota_charge.refund.assert_called_once_with() - session = mock_trigger_workflow_async.call_args.kwargs["session"] - assert isinstance(session, Session) - assert session.in_transaction() is False - assert len(caplog.records) == 1 - assert caplog.records[0].levelno == logging.INFO - assert caplog.records[0].message == ( - "Tenant tenant-123 quota exceeded for feature workflow, skipping webhook trigger webhook-123" - ) - - def test_trigger_workflow_execution_success( - self, - monkeypatch: pytest.MonkeyPatch, - sqlite_engine: Engine, - ) -> None: - monkeypatch.setattr(webhook_service_module, "db", SimpleNamespace(engine=sqlite_engine)) - webhook_trigger = _webhook_trigger() - workflow = _workflow() - end_user = _end_user() - quota_charge = MagicMock() - webhook_data = { - "method": "POST", - "headers": {"Authorization": "Bearer token"}, - "query_params": {"version": "1"}, - "body": {"message": "hello"}, - "files": {}, - } - - with ( - patch.object(webhook_service_module.QuotaService, "reserve", return_value=quota_charge), - patch.object( - webhook_service_module.AsyncWorkflowService, - "trigger_workflow_async", - ) as mock_trigger, - ): - WebhookService.trigger_workflow_execution( - webhook_trigger, - webhook_data, - workflow, - end_users=_EndUserServiceStub(end_user), - ) - - call_session = mock_trigger.call_args.kwargs["session"] - assert call_session.get_bind() is sqlite_engine - quota_charge.commit.assert_called_once_with() - quota_charge.refund.assert_not_called() - - def test_trigger_workflow_execution_end_user_service_failure(self) -> None: - webhook_trigger = _webhook_trigger() - workflow = _workflow() - webhook_data = {"method": "POST", "headers": {}, "query_params": {}, "body": {}, "files": {}} - - with pytest.raises(ValueError, match="Failed to create end user"): - WebhookService.trigger_workflow_execution( - webhook_trigger, - webhook_data, - workflow, - end_users=_EndUserServiceStub(ValueError("Failed to create end user")), - ) - - def test_extract_webhook_data_json(self): - """Test webhook data extraction from JSON request.""" - app = Flask(__name__) - - with app.test_request_context( - "/webhook", - method="POST", - headers={"Content-Type": "application/json", "Authorization": "Bearer token"}, - query_string="version=1&format=json", - json={"message": "hello", "count": 42}, - ): - webhook_trigger = _webhook_trigger() - webhook_data = WebhookService.extract_webhook_data(webhook_trigger) - - assert webhook_data["method"] == "POST" - assert webhook_data["headers"]["Authorization"] == "Bearer token" - # Query params are now extracted as raw strings - assert webhook_data["query_params"]["version"] == "1" - assert webhook_data["query_params"]["format"] == "json" - assert webhook_data["body"]["message"] == "hello" - assert webhook_data["body"]["count"] == 42 - assert webhook_data["files"] == {} - - def test_extract_webhook_data_query_params_remain_strings(self): - """Query parameters should be extracted as raw strings without automatic conversion.""" - app = Flask(__name__) - - with app.test_request_context( - "/webhook", - method="GET", - headers={"Content-Type": "application/json"}, - query_string="count=42&threshold=3.14&enabled=true¬e=text", - ): - webhook_trigger = _webhook_trigger() - webhook_data = WebhookService.extract_webhook_data(webhook_trigger) - - # After refactoring, raw extraction keeps query params as strings - assert webhook_data["query_params"]["count"] == "42" - assert webhook_data["query_params"]["threshold"] == "3.14" - assert webhook_data["query_params"]["enabled"] == "true" - assert webhook_data["query_params"]["note"] == "text" - - def test_extract_webhook_data_form_urlencoded(self): - """Test webhook data extraction from form URL encoded request.""" - app = Flask(__name__) - - with app.test_request_context( - "/webhook", - method="POST", - headers={"Content-Type": "application/x-www-form-urlencoded"}, - data={"username": "test", "password": "secret"}, - ): - webhook_trigger = _webhook_trigger() - webhook_data = WebhookService.extract_webhook_data(webhook_trigger) - - assert webhook_data["method"] == "POST" - assert webhook_data["body"]["username"] == "test" - assert webhook_data["body"]["password"] == "secret" - - def test_extract_webhook_data_multipart_with_files(self): - """Test webhook data extraction from multipart form with files.""" - app = Flask(__name__) - - # Create a mock file - file_content = b"test file content" - file_storage = FileStorage(stream=BytesIO(file_content), filename="test.txt", content_type="text/plain") - - with app.test_request_context( - "/webhook", - method="POST", - headers={"Content-Type": "multipart/form-data"}, - data={"message": "test", "file": file_storage}, - ): - webhook_trigger = _webhook_trigger() - webhook_trigger.tenant_id = "test_tenant" - - with patch.object(WebhookService, "_process_file_uploads", autospec=True) as mock_process_files: - mock_process_files.return_value = {"file": "mocked_file_obj"} - - webhook_data = WebhookService.extract_webhook_data(webhook_trigger) - - assert webhook_data["method"] == "POST" - assert webhook_data["body"]["message"] == "test" - assert webhook_data["files"]["file"] == "mocked_file_obj" - mock_process_files.assert_called_once() - - def test_extract_webhook_data_raw_text(self): - """Test webhook data extraction from raw text request.""" - app = Flask(__name__) - - with app.test_request_context( - "/webhook", method="POST", headers={"Content-Type": "text/plain"}, data="raw text content" - ): - webhook_trigger = _webhook_trigger() - webhook_data = WebhookService.extract_webhook_data(webhook_trigger) - - assert webhook_data["method"] == "POST" - assert webhook_data["body"]["raw"] == "raw text content" - - def test_extract_octet_stream_body_uses_detected_mime(self): - """Octet-stream uploads should rely on detected MIME type.""" - app = Flask(__name__) - binary_content = b"plain text data" - - with app.test_request_context( - "/webhook", method="POST", headers={"Content-Type": "application/octet-stream"}, data=binary_content - ): - webhook_trigger = _webhook_trigger() - mock_file = MagicMock() - mock_file.to_dict.return_value = {"file": "data"} - - with ( - patch.object( - WebhookService, "_detect_binary_mimetype", return_value="text/plain", autospec=True - ) as mock_detect, - patch.object(WebhookService, "_create_file_from_binary", autospec=True) as mock_create, - ): - mock_create.return_value = mock_file - body, files = WebhookService._extract_octet_stream_body(webhook_trigger) - - assert body["raw"] == {"file": "data"} - assert files == {} - mock_detect.assert_called_once_with(binary_content) - mock_create.assert_called_once() - args = mock_create.call_args[0] - assert args[0] == binary_content - assert args[1] == "text/plain" - assert args[2] is webhook_trigger - - def test_detect_binary_mimetype_uses_magic(self, monkeypatch: pytest.MonkeyPatch): - """python-magic output should be used when available.""" - fake_magic = MagicMock() - fake_magic.from_buffer.return_value = "image/png" - monkeypatch.setattr("services.trigger.webhook_service.magic", fake_magic) - - result = WebhookService._detect_binary_mimetype(b"binary data") - - assert result == "image/png" - fake_magic.from_buffer.assert_called_once() - - def test_detect_binary_mimetype_fallback_without_magic(self, monkeypatch: pytest.MonkeyPatch): - """Fallback MIME type should be used when python-magic is unavailable.""" - monkeypatch.setattr("services.trigger.webhook_service.magic", None) - - result = WebhookService._detect_binary_mimetype(b"binary data") - - assert result == "application/octet-stream" - - def test_detect_binary_mimetype_handles_magic_exception( - self, monkeypatch: pytest.MonkeyPatch, caplog: pytest.LogCaptureFixture - ): - """Fallback MIME type should be used when python-magic raises an exception.""" - try: - import magic as real_magic - except ImportError: - pytest.skip("python-magic is not installed") - - fake_magic = MagicMock() - fake_magic.from_buffer.side_effect = real_magic.MagicException("magic error") - monkeypatch.setattr("services.trigger.webhook_service.magic", fake_magic) - caplog.set_level(logging.DEBUG, logger="services.trigger.webhook_service") - - result = WebhookService._detect_binary_mimetype(b"binary data") - - assert result == "application/octet-stream" - assert "python-magic detection failed for octet-stream payload" in caplog.messages - - def test_extract_webhook_data_invalid_json(self): - """Test webhook data extraction with invalid JSON.""" - app = Flask(__name__) - - with app.test_request_context( - "/webhook", method="POST", headers={"Content-Type": "application/json"}, data="invalid json" - ): - webhook_trigger = _webhook_trigger() - with pytest.raises(ValueError, match="Invalid JSON body"): - WebhookService.extract_webhook_data(webhook_trigger) - - def test_generate_webhook_response_default(self): - """Test webhook response generation with default values.""" - node_config = {"data": {}} - - response_data, status_code = WebhookService.generate_webhook_response(node_config) - - assert status_code == 200 - assert response_data["status"] == "success" - assert "Webhook processed successfully" in response_data["message"] - - def test_generate_webhook_response_custom_json(self): - """Test webhook response generation with custom JSON response.""" - node_config = {"data": {"status_code": 201, "response_body": '{"result": "created", "id": 123}'}} - - response_data, status_code = WebhookService.generate_webhook_response(node_config) - - assert status_code == 201 - assert response_data["result"] == "created" - assert response_data["id"] == 123 - - def test_generate_webhook_response_custom_text(self): - """Test webhook response generation with custom text response.""" - node_config = {"data": {"status_code": 202, "response_body": "Request accepted for processing"}} - - response_data, status_code = WebhookService.generate_webhook_response(node_config) - - assert status_code == 202 - assert response_data["message"] == "Request accepted for processing" - - def test_generate_webhook_response_invalid_json(self): - """Test webhook response generation with invalid JSON response.""" - node_config = {"data": {"status_code": 400, "response_body": '{"invalid": json}'}} - - response_data, status_code = WebhookService.generate_webhook_response(node_config) - - assert status_code == 400 - assert response_data["message"] == '{"invalid": json}' - - def test_generate_webhook_response_empty_response_body(self): - """Test webhook response generation with empty response body.""" - node_config = {"data": {"status_code": 204, "response_body": ""}} - - response_data, status_code = WebhookService.generate_webhook_response(node_config) - - assert status_code == 204 - assert response_data["status"] == "success" - assert "Webhook processed successfully" in response_data["message"] - - def test_generate_webhook_response_array_json(self): - """Test webhook response generation with JSON array response.""" - node_config = {"data": {"status_code": 200, "response_body": '[{"id": 1}, {"id": 2}]'}} - - response_data, status_code = WebhookService.generate_webhook_response(node_config) - - assert status_code == 200 - assert isinstance(response_data, list) - assert len(response_data) == 2 - assert response_data[0]["id"] == 1 - assert response_data[1]["id"] == 2 - - @patch("services.trigger.webhook_service.ToolFileManager", autospec=True) - @patch("services.trigger.webhook_service.file_factory", autospec=True) - def test_process_file_uploads_success(self, mock_file_factory, mock_tool_file_manager): - """Test successful file upload processing.""" - # Mock ToolFileManager - mock_tool_file_instance = mock_tool_file_manager.return_value # Mock file creation - mock_tool_file_instance.create_file_by_raw.return_value = _tool_file() - - # Mock file factory - mock_file_obj = MagicMock() - mock_file_factory.build_from_mapping.return_value = mock_file_obj - - # Create mock files - files = { - "file1": MagicMock(filename="test1.txt", content_type="text/plain"), - "file2": MagicMock(filename="test2.jpg", content_type="image/jpeg"), - } - - # Mock file reads - files["file1"].stream.read.return_value = b"content1" - files["file2"].stream.read.return_value = b"content2" - - webhook_trigger = _webhook_trigger() - webhook_trigger.tenant_id = "test_tenant" - - result = WebhookService._process_file_uploads(files, webhook_trigger) - - assert len(result) == 2 - assert "file1" in result - assert "file2" in result - - # Verify file processing was called for each file - assert mock_tool_file_manager.call_count == 2 - assert mock_file_factory.build_from_mapping.call_count == 2 - - @patch("services.trigger.webhook_service.ToolFileManager", autospec=True) - @patch("services.trigger.webhook_service.file_factory", autospec=True) - def test_process_file_uploads_with_errors(self, mock_file_factory, mock_tool_file_manager): - """Test file upload processing with errors.""" - # Mock ToolFileManager - mock_tool_file_instance = mock_tool_file_manager.return_value # Mock file creation - mock_tool_file_instance.create_file_by_raw.return_value = _tool_file() - - # Mock file factory - mock_file_obj = MagicMock() - mock_file_factory.build_from_mapping.return_value = mock_file_obj - - # Create mock files, one will fail - files = { - "good_file": MagicMock(filename="test.txt", content_type="text/plain"), - "bad_file": MagicMock(filename="test.bad", content_type="text/plain"), - } - - files["good_file"].stream.read.return_value = b"content" - files["bad_file"].stream.read.side_effect = Exception("Read error") - - webhook_trigger = _webhook_trigger() - webhook_trigger.tenant_id = "test_tenant" - - result = WebhookService._process_file_uploads(files, webhook_trigger) - - # Should process the good file and skip the bad one - assert len(result) == 1 - assert "good_file" in result - assert "bad_file" not in result - - def test_process_file_uploads_empty_filename(self): - """Test file upload processing with empty filename.""" - files = { - "no_filename": MagicMock(filename="", content_type="text/plain"), - "none_filename": MagicMock(filename=None, content_type="text/plain"), - } - - webhook_trigger = _webhook_trigger() - webhook_trigger.tenant_id = "test_tenant" - - result = WebhookService._process_file_uploads(files, webhook_trigger) - - # Should skip files without filenames - assert len(result) == 0 - - def test_validate_json_value_string(self): - """Test JSON value validation for string type.""" - # Valid string - result = WebhookService._validate_json_value("name", "hello", "string") - assert result == "hello" - - # Invalid string (number) - should raise ValueError - with pytest.raises(ValueError, match="Expected string, got int"): - WebhookService._validate_json_value("name", 123, "string") - - def test_validate_json_value_number(self): - """Test JSON value validation for number type.""" - # Valid integer - result = WebhookService._validate_json_value("count", 42, "number") - assert result == 42 - - # Valid float - result = WebhookService._validate_json_value("price", 19.99, "number") - assert result == 19.99 - - # Invalid number (string) - should raise ValueError - with pytest.raises(ValueError, match="Expected number, got str"): - WebhookService._validate_json_value("count", "42", "number") - - def test_validate_json_value_bool(self): - """Test JSON value validation for boolean type.""" - # Valid boolean - result = WebhookService._validate_json_value("enabled", True, "boolean") - assert result is True - - result = WebhookService._validate_json_value("enabled", False, "boolean") - assert result is False - - # Invalid boolean (string) - should raise ValueError - with pytest.raises(ValueError, match="Expected boolean, got str"): - WebhookService._validate_json_value("enabled", "true", "boolean") - - def test_validate_json_value_object(self): - """Test JSON value validation for object type.""" - # Valid object - result = WebhookService._validate_json_value("user", {"name": "John", "age": 30}, "object") - assert result == {"name": "John", "age": 30} - - # Invalid object (string) - should raise ValueError - with pytest.raises(ValueError, match="Expected object, got str"): - WebhookService._validate_json_value("user", "not_an_object", "object") - - def test_validate_json_value_array_string(self): - """Test JSON value validation for array[string] type.""" - # Valid array of strings - result = WebhookService._validate_json_value("tags", ["tag1", "tag2", "tag3"], "array[string]") - assert result == ["tag1", "tag2", "tag3"] - - # Invalid - not an array - with pytest.raises(ValueError, match="Expected array of strings, got str"): - WebhookService._validate_json_value("tags", "not_an_array", "array[string]") - - # Invalid - array with non-strings - with pytest.raises(ValueError, match="Expected array of strings, got list"): - WebhookService._validate_json_value("tags", ["tag1", 123, "tag3"], "array[string]") - - def test_validate_json_value_array_number(self): - """Test JSON value validation for array[number] type.""" - # Valid array of numbers - result = WebhookService._validate_json_value("scores", [1, 2.5, 3, 4.7], "array[number]") - assert result == [1, 2.5, 3, 4.7] - - # Invalid - array with non-numbers - with pytest.raises(ValueError, match="Expected array of numbers, got list"): - WebhookService._validate_json_value("scores", [1, "2", 3], "array[number]") - - def test_validate_json_value_array_bool(self): - """Test JSON value validation for array[boolean] type.""" - # Valid array of booleans - result = WebhookService._validate_json_value("flags", [True, False, True], "array[boolean]") - assert result == [True, False, True] - - # Invalid - array with non-booleans - with pytest.raises(ValueError, match="Expected array of booleans, got list"): - WebhookService._validate_json_value("flags", [True, "false", True], "array[boolean]") - - def test_validate_json_value_array_object(self): - """Test JSON value validation for array[object] type.""" - # Valid array of objects - result = WebhookService._validate_json_value("users", [{"name": "John"}, {"name": "Jane"}], "array[object]") - assert result == [{"name": "John"}, {"name": "Jane"}] - - # Invalid - array with non-objects - with pytest.raises(ValueError, match="Expected array of objects, got list"): - WebhookService._validate_json_value("users", [{"name": "John"}, "not_object"], "array[object]") - - def test_convert_form_value_string(self): - """Test form value conversion for string type.""" - result = WebhookService._convert_form_value("test", "hello", "string") - assert result == "hello" - - def test_convert_form_value_number(self): - """Test form value conversion for number type.""" - # Integer - result = WebhookService._convert_form_value("count", "42", "number") - assert result == 42 - - # Float - result = WebhookService._convert_form_value("price", "19.99", "number") - assert result == 19.99 - - # Invalid number - with pytest.raises(ValueError, match="Cannot convert 'not_a_number' to number"): - WebhookService._convert_form_value("count", "not_a_number", "number") - - def test_convert_form_value_boolean(self): - """Test form value conversion for boolean type.""" - # True values - assert WebhookService._convert_form_value("flag", "true", "boolean") is True - assert WebhookService._convert_form_value("flag", "1", "boolean") is True - assert WebhookService._convert_form_value("flag", "yes", "boolean") is True - - # False values - assert WebhookService._convert_form_value("flag", "false", "boolean") is False - assert WebhookService._convert_form_value("flag", "0", "boolean") is False - assert WebhookService._convert_form_value("flag", "no", "boolean") is False - - # Invalid boolean - with pytest.raises(ValueError, match="Cannot convert 'maybe' to boolean"): - WebhookService._convert_form_value("flag", "maybe", "boolean") - - def test_extract_and_validate_webhook_data_success(self): - """Test successful unified data extraction and validation.""" - app = Flask(__name__) - - with app.test_request_context( - "/webhook", - method="POST", - headers={"Content-Type": "application/json"}, - query_string="count=42&enabled=true", - json={"message": "hello", "age": 25}, - ): - webhook_trigger = _webhook_trigger() - node_config = { - "data": { - "method": "post", - "content_type": "application/json", - "params": [ - {"name": "count", "type": "number", "required": True}, - {"name": "enabled", "type": "boolean", "required": True}, - ], - "body": [ - {"name": "message", "type": "string", "required": True}, - {"name": "age", "type": "number", "required": True}, - ], - } - } - - result = WebhookService.extract_and_validate_webhook_data(webhook_trigger, node_config) - - # Check that types are correctly converted - assert result["query_params"]["count"] == 42 # Converted to int - assert result["query_params"]["enabled"] is True # Converted to bool - assert result["body"]["message"] == "hello" # Already string - assert result["body"]["age"] == 25 # Already number - - def test_extract_and_validate_webhook_data_invalid_json_error(self): - """Invalid JSON should bubble up as a ValueError with details.""" - app = Flask(__name__) - - with app.test_request_context( - "/webhook", - method="POST", - headers={"Content-Type": "application/json"}, - data='{"invalid": }', - ): - webhook_trigger = _webhook_trigger() - node_config = { - "data": { - "method": "post", - "content_type": "application/json", - } - } - - with pytest.raises(ValueError, match="Invalid JSON body"): - WebhookService.extract_and_validate_webhook_data(webhook_trigger, node_config) - - def test_extract_and_validate_webhook_data_validation_error(self): - """Test unified data extraction with validation error.""" - app = Flask(__name__) - - with app.test_request_context( - "/webhook", - method="GET", # Wrong method - headers={"Content-Type": "application/json"}, - ): - webhook_trigger = _webhook_trigger() - node_config = { - "data": { - "method": "post", # Expects POST - "content_type": "application/json", - } - } - - with pytest.raises(ValueError, match="HTTP method mismatch"): - WebhookService.extract_and_validate_webhook_data(webhook_trigger, node_config) - - def test_extract_and_validate_webhook_request_missing_required_header(self) -> None: - app = Flask(__name__) - with app.test_request_context( - "/webhook", - method="POST", - headers={"Content-Type": "application/json"}, - ): - node_config = { - "data": { - "method": "post", - "content_type": "application/json", - "headers": [{"name": "Authorization", "required": True}], - } - } - - with pytest.raises(ValueError, match="Required header missing: Authorization"): - WebhookService.extract_and_validate_webhook_data(_webhook_trigger(), node_config) - - def test_extract_and_validate_webhook_request_case_insensitive_headers(self) -> None: - app = Flask(__name__) - with app.test_request_context( - "/webhook", - method="POST", - headers={"Content-Type": "application/json", "authorization": "Bearer token"}, - json={"message": "hello"}, - ): - node_config = { - "data": { - "method": "post", - "content_type": "application/json", - "headers": [{"name": "Authorization", "required": True}], - "body": [{"name": "message", "type": "string", "required": True}], - } - } - - result = WebhookService.extract_and_validate_webhook_data(_webhook_trigger(), node_config) - - assert result["headers"].get("Authorization") == "Bearer token" - - def test_extract_and_validate_webhook_request_missing_required_param(self) -> None: - app = Flask(__name__) - with app.test_request_context( - "/webhook", - method="POST", - headers={"Content-Type": "application/json"}, - json={"message": "hello"}, - ): - node_config = { - "data": { - "method": "post", - "content_type": "application/json", - "params": [{"name": "version", "required": True}], - "body": [{"name": "message", "type": "string", "required": True}], - } - } - - with pytest.raises(ValueError, match="Required parameter missing: version"): - WebhookService.extract_and_validate_webhook_data(_webhook_trigger(), node_config) - - def test_extract_and_validate_webhook_request_missing_required_body_param(self) -> None: - app = Flask(__name__) - with app.test_request_context( - "/webhook", - method="POST", - headers={"Content-Type": "application/json"}, - json={}, - ): - node_config = { - "data": { - "method": "post", - "content_type": "application/json", - "body": [{"name": "message", "type": "string", "required": True}], - } - } - - with pytest.raises(ValueError, match="Required body parameter missing: message"): - WebhookService.extract_and_validate_webhook_data(_webhook_trigger(), node_config) - - def test_extract_and_validate_webhook_request_missing_required_file(self) -> None: - app = Flask(__name__) - with app.test_request_context( - "/webhook", - method="POST", - data={"note": "test"}, - content_type="multipart/form-data", - ): - node_config = { - "data": { - "method": "post", - "content_type": "multipart/form-data", - "body": [{"name": "file", "type": "file", "required": True}], - } - } - - result = WebhookService.extract_and_validate_webhook_data(_webhook_trigger(), node_config) - - assert result["files"] == {} - - def test_debug_mode_parameter_handling(self): - """Test that the debug mode parameter is properly handled in _prepare_webhook_execution.""" - from controllers.trigger.webhook import _prepare_webhook_execution - - # Mock the WebhookService methods - with ( - patch.object(WebhookService, "get_webhook_trigger_and_workflow", autospec=True) as mock_get_trigger, - patch.object(WebhookService, "extract_and_validate_webhook_data", autospec=True) as mock_extract, - ): - webhook_trigger = _webhook_trigger() - workflow = _workflow() - mock_config = {"data": {"test": "config"}} - mock_data = {"test": "data"} - - mock_get_trigger.return_value = (webhook_trigger, workflow, mock_config) - mock_extract.return_value = mock_data - - result = _prepare_webhook_execution("test_webhook", is_debug=False) - assert result == (webhook_trigger, workflow, mock_config, mock_data, None) - - # Reset mock - mock_get_trigger.reset_mock() - - result = _prepare_webhook_execution("test_webhook", is_debug=True) - assert result == (webhook_trigger, workflow, mock_config, mock_data, None) diff --git a/api/tests/unit_tests/services/test_webhook_service_additional.py b/api/tests/unit_tests/services/test_webhook_service_additional.py deleted file mode 100644 index c4197b575fa..00000000000 --- a/api/tests/unit_tests/services/test_webhook_service_additional.py +++ /dev/null @@ -1,286 +0,0 @@ -import logging -from types import SimpleNamespace -from typing import Any -from unittest.mock import MagicMock - -import pytest -from flask import Flask -from werkzeug.exceptions import RequestEntityTooLarge - -from core.workflow.nodes.trigger_webhook.entities import ( - ContentType, - WebhookBodyParameter, - WebhookData, - WebhookParameter, -) -from graphon.variables.types import SegmentType -from models.trigger import WorkflowWebhookTrigger -from services.trigger import webhook_service as service_module -from services.trigger.webhook_service import WebhookService - - -@pytest.fixture -def flask_app() -> Flask: - return Flask(__name__) - - -def _workflow_trigger(**kwargs: Any) -> WorkflowWebhookTrigger: - values = { - "webhook_id": "webhook-123", - "tenant_id": "tenant-1", - "app_id": "app-1", - "node_id": "node-1", - "created_by": "user-1", - } - values.update(kwargs) - return WorkflowWebhookTrigger(**values) - - -class TestWebhookServiceExtractionFallbacks: - def test_extract_webhook_data_should_use_text_fallback_for_unknown_content_type( - self, - flask_app: Flask, - monkeypatch: pytest.MonkeyPatch, - caplog: pytest.LogCaptureFixture, - ) -> None: - webhook_trigger = _workflow_trigger() - - with caplog.at_level(logging.WARNING, logger="services.trigger.webhook_service"): - with flask_app.test_request_context( - "/webhook", - method="POST", - headers={"Content-Type": "application/vnd.custom"}, - data="plain content", - ): - result = WebhookService.extract_webhook_data(webhook_trigger) - - assert result["body"] == {"raw": "plain content"} - assert any(r.levelno >= logging.WARNING for r in caplog.records) - - def test_extract_webhook_data_should_raise_for_request_too_large( - self, - flask_app: Flask, - monkeypatch: pytest.MonkeyPatch, - ) -> None: - from tests.unit_tests.config_override import apply_config_overrides - - apply_config_overrides(monkeypatch, WEBHOOK_REQUEST_BODY_MAX_SIZE=1) - - with flask_app.test_request_context("/webhook", method="POST", data="ab"): - with pytest.raises(RequestEntityTooLarge): - WebhookService.extract_webhook_data(_workflow_trigger()) - - def test_extract_octet_stream_body_should_return_none_when_empty_payload(self, flask_app: Flask) -> None: - webhook_trigger = _workflow_trigger() - - with flask_app.test_request_context("/webhook", method="POST", data=b""): - body, files = WebhookService._extract_octet_stream_body(webhook_trigger) - - assert body == {"raw": None} - assert files == {} - - def test_extract_octet_stream_body_should_return_none_when_processing_raises( - self, - flask_app: Flask, - monkeypatch: pytest.MonkeyPatch, - ) -> None: - webhook_trigger = _workflow_trigger() - monkeypatch.setattr( - WebhookService, "_detect_binary_mimetype", MagicMock(return_value="application/octet-stream") - ) - monkeypatch.setattr(WebhookService, "_create_file_from_binary", MagicMock(side_effect=RuntimeError("boom"))) - - with flask_app.test_request_context("/webhook", method="POST", data=b"abc"): - body, files = WebhookService._extract_octet_stream_body(webhook_trigger) - - assert body == {"raw": None} - assert files == {} - - def test_extract_text_body_should_return_empty_string_when_request_read_fails( - self, - flask_app: Flask, - monkeypatch: pytest.MonkeyPatch, - ) -> None: - monkeypatch.setattr("flask.wrappers.Request.get_data", MagicMock(side_effect=RuntimeError("read error"))) - - with flask_app.test_request_context("/webhook", method="POST", data="abc"): - body, files = WebhookService._extract_text_body() - - assert body == {"raw": ""} - assert files == {} - - def test_detect_binary_mimetype_should_fallback_when_magic_raises( - self, - monkeypatch: pytest.MonkeyPatch, - ) -> None: - fake_magic = MagicMock() - fake_magic.from_buffer.side_effect = RuntimeError("magic failed") - monkeypatch.setattr(service_module, "magic", fake_magic) - - result = WebhookService._detect_binary_mimetype(b"binary") - - assert result == "application/octet-stream" - - def test_process_file_uploads_should_use_octet_stream_fallback_when_mimetype_unknown( - self, - monkeypatch: pytest.MonkeyPatch, - ) -> None: - webhook_trigger = _workflow_trigger(created_by="user-1", tenant_id="tenant-1") - file_obj = MagicMock() - file_obj.to_dict.return_value = {"id": "f-1"} - monkeypatch.setattr(WebhookService, "_create_file_from_binary", MagicMock(return_value=file_obj)) - monkeypatch.setattr(service_module.mimetypes, "guess_type", MagicMock(return_value=(None, None))) - - uploaded = MagicMock() - uploaded.filename = "file.unknown" - uploaded.content_type = None - uploaded.read.return_value = b"content" - - result = WebhookService._process_file_uploads({"f": uploaded}, webhook_trigger) - - assert result == {"f": {"id": "f-1"}} - - def test_create_file_from_binary_should_call_tool_file_manager_and_file_factory( - self, - monkeypatch: pytest.MonkeyPatch, - ) -> None: - webhook_trigger = _workflow_trigger(created_by="user-1", tenant_id="tenant-1") - manager = MagicMock() - manager.create_file_by_raw.return_value = SimpleNamespace(id="tool-file-1") - monkeypatch.setattr(service_module, "ToolFileManager", MagicMock(return_value=manager)) - expected_file = MagicMock() - monkeypatch.setattr(service_module.file_factory, "build_from_mapping", MagicMock(return_value=expected_file)) - - result = WebhookService._create_file_from_binary(b"abc", "text/plain", webhook_trigger) - - assert result is expected_file - manager.create_file_by_raw.assert_called_once() - - -class TestWebhookServiceValidationAndConversion: - @pytest.mark.parametrize( - ("raw_value", "param_type", "expected"), - [ - ("42", SegmentType.NUMBER, 42), - ("3.14", SegmentType.NUMBER, float("3.14")), - ("yes", SegmentType.BOOLEAN, True), - ("no", SegmentType.BOOLEAN, False), - ], - ) - def test_convert_form_value_should_convert_supported_types( - self, - raw_value: str, - param_type: str, - expected: Any, - ) -> None: - result = WebhookService._convert_form_value("param", raw_value, param_type) - assert result == expected - - def test_convert_form_value_should_raise_for_unsupported_type(self) -> None: - with pytest.raises(ValueError, match="Unsupported type"): - WebhookService._convert_form_value("p", "x", SegmentType.FILE) - - def test_validate_json_value_should_return_original_for_unmapped_supported_segment_type( - self, - monkeypatch: pytest.MonkeyPatch, - caplog: pytest.LogCaptureFixture, - ) -> None: - with caplog.at_level(logging.WARNING, logger="services.trigger.webhook_service"): - result = WebhookService._validate_json_value("param", {"x": 1}, "unsupported-type") - - assert result == {"x": 1} - assert any(r.levelno >= logging.WARNING for r in caplog.records) - - def test_validate_and_convert_value_should_wrap_conversion_errors(self) -> None: - with pytest.raises(ValueError, match="validation failed"): - WebhookService._validate_and_convert_value("param", "bad", SegmentType.NUMBER, is_form_data=True) - - def test_process_parameters_should_raise_when_required_parameter_missing(self) -> None: - raw_params = {"optional": "x"} - config = [WebhookParameter(name="required_param", type=SegmentType.STRING, required=True)] - - with pytest.raises(ValueError, match="Required parameter missing"): - WebhookService._process_parameters(raw_params, config, is_form_data=True) - - def test_process_parameters_should_include_unconfigured_parameters(self) -> None: - raw_params = {"known": "1", "unknown": "x"} - config = [WebhookParameter(name="known", type=SegmentType.NUMBER, required=False)] - - result = WebhookService._process_parameters(raw_params, config, is_form_data=True) - - assert result == {"known": 1, "unknown": "x"} - - def test_process_body_parameters_should_raise_when_required_text_raw_is_missing(self) -> None: - with pytest.raises(ValueError, match="Required body content missing"): - WebhookService._process_body_parameters( - raw_body={"raw": ""}, - body_configs=[WebhookBodyParameter(name="raw", required=True)], - content_type=ContentType.TEXT, - ) - - def test_process_body_parameters_should_skip_file_config_for_multipart_form_data(self) -> None: - raw_body = {"message": "hello", "extra": "x"} - body_configs = [ - WebhookBodyParameter(name="upload", type=SegmentType.FILE, required=True), - WebhookBodyParameter(name="message", type=SegmentType.STRING, required=True), - ] - - result = WebhookService._process_body_parameters(raw_body, body_configs, ContentType.FORM_DATA) - - assert result == {"message": "hello", "extra": "x"} - - def test_validate_required_headers_should_accept_sanitized_header_names(self) -> None: - headers = {"x_api_key": "123"} - configs = [WebhookParameter(name="x-api-key", required=True)] - - WebhookService._validate_required_headers(headers, configs) - - def test_validate_required_headers_should_raise_when_required_header_missing(self) -> None: - headers = {"x-other": "123"} - configs = [WebhookParameter(name="x-api-key", required=True)] - - with pytest.raises(ValueError, match="Required header missing"): - WebhookService._validate_required_headers(headers, configs) - - def test_validate_http_metadata_should_return_content_type_mismatch_error(self) -> None: - webhook_data = {"method": "POST", "headers": {"Content-Type": "application/json"}} - node_data = WebhookData(method="post", content_type=ContentType.TEXT) - - result = WebhookService._validate_http_metadata(webhook_data, node_data) - - assert result["valid"] is False - assert "Content-type mismatch" in result["error"] - - def test_extract_content_type_should_fallback_to_lowercase_header_key(self) -> None: - headers = {"content-type": "application/json; charset=utf-8"} - assert WebhookService._extract_content_type(headers) == "application/json" - - def test_build_workflow_inputs_should_include_expected_keys(self) -> None: - webhook_data = {"headers": {"h": "v"}, "query_params": {"q": 1}, "body": {"b": 2}} - - result = WebhookService.build_workflow_inputs(webhook_data) - - assert result["webhook_data"] == webhook_data - assert result["webhook_headers"] == {"h": "v"} - assert result["webhook_query_params"] == {"q": 1} - assert result["webhook_body"] == {"b": 2} - - -class TestWebhookServiceUtilities: - def test_generate_webhook_response_should_fallback_when_response_body_is_not_json(self) -> None: - node_config = {"data": {"status_code": 200, "response_body": "{bad-json"}} - - body, status = WebhookService.generate_webhook_response(node_config) - - assert status == 200 - assert "message" in body - - def test_generate_webhook_id_should_return_24_character_identifier(self) -> None: - webhook_id = WebhookService.generate_webhook_id() - - assert isinstance(webhook_id, str) - assert len(webhook_id) == 24 - - def test_sanitize_key_should_return_original_value_for_non_string_input(self) -> None: - result = WebhookService._sanitize_key(123) # type: ignore[arg-type] - assert result == 123 diff --git a/api/tests/unit_tests/services/test_webhook_service_relationships.py b/api/tests/unit_tests/services/test_webhook_service_relationships.py deleted file mode 100644 index 55e16db554a..00000000000 --- a/api/tests/unit_tests/services/test_webhook_service_relationships.py +++ /dev/null @@ -1,96 +0,0 @@ -from types import SimpleNamespace -from typing import cast -from unittest.mock import MagicMock - -import pytest -from sqlalchemy import Engine, Table, select -from sqlalchemy.orm import Session, sessionmaker - -from core.trigger.constants import TRIGGER_WEBHOOK_NODE_TYPE -from models.base import TypeBase -from models.model import App -from models.trigger import WorkflowWebhookTrigger -from models.workflow import Workflow -from services.trigger import webhook_service -from services.trigger.webhook_service import WebhookService - - -def test_draft_sync_preserves_webhook_id_when_deleted_node_is_restored( - monkeypatch: pytest.MonkeyPatch, sqlite_engine: Engine -) -> None: - webhook_trigger_table = cast(Table, WorkflowWebhookTrigger.__table__) - TypeBase.metadata.create_all(sqlite_engine, tables=[webhook_trigger_table]) - session_maker = sessionmaker(bind=sqlite_engine, expire_on_commit=False) - app = cast( - App, - SimpleNamespace(id="00000000-0000-0000-0000-000000000001", tenant_id="tenant-1", created_by="user-1"), - ) - assert app.created_by is not None - webhook_trigger = WorkflowWebhookTrigger( - app_id=app.id, - tenant_id=app.tenant_id, - node_id="webhook-node", - webhook_id="stable-webhook-id", - created_by=app.created_by, - ) - with session_maker.begin() as session: - session.add(webhook_trigger) - original_record_id = webhook_trigger.id - - redis = MagicMock() - redis.get.return_value = None - redis.lock.return_value.acquire.return_value = True - monkeypatch.setattr(webhook_service, "db", SimpleNamespace(engine=sqlite_engine)) - monkeypatch.setattr(webhook_service, "redis_client", redis) - - deleted_draft = cast(Workflow, SimpleNamespace(walk_nodes=lambda _node_type: iter(()))) - WebhookService.sync_webhook_relationships(app, deleted_draft, remove_stale=False) - - restored_draft = cast( - Workflow, - SimpleNamespace(walk_nodes=lambda _node_type: iter([("webhook-node", {"type": TRIGGER_WEBHOOK_NODE_TYPE})])), - ) - WebhookService.sync_webhook_relationships(app, restored_draft, remove_stale=False) - - with Session(sqlite_engine) as session: - records = session.scalars(select(WorkflowWebhookTrigger)).all() - - assert len(records) == 1 - assert records[0].id == original_record_id - assert records[0].webhook_id == "stable-webhook-id" - - -def test_published_sync_removes_stale_webhook_relationships( - monkeypatch: pytest.MonkeyPatch, sqlite_engine: Engine -) -> None: - webhook_trigger_table = cast(Table, WorkflowWebhookTrigger.__table__) - TypeBase.metadata.create_all(sqlite_engine, tables=[webhook_trigger_table]) - session_maker = sessionmaker(bind=sqlite_engine, expire_on_commit=False) - app = cast( - App, - SimpleNamespace(id="00000000-0000-0000-0000-000000000001", tenant_id="tenant-1", created_by="user-1"), - ) - assert app.created_by is not None - webhook_trigger = WorkflowWebhookTrigger( - app_id=app.id, - tenant_id=app.tenant_id, - node_id="deleted-webhook-node", - webhook_id="deleted-webhook-id", - created_by=app.created_by, - ) - with session_maker.begin() as session: - session.add(webhook_trigger) - - redis = MagicMock() - redis.get.return_value = None - redis.lock.return_value.acquire.return_value = True - monkeypatch.setattr(webhook_service, "db", SimpleNamespace(engine=sqlite_engine)) - monkeypatch.setattr(webhook_service, "redis_client", redis) - - published_workflow = cast(Workflow, SimpleNamespace(walk_nodes=lambda _node_type: iter(()))) - WebhookService.sync_webhook_relationships(app, published_workflow, remove_stale=True) - - with session_maker() as session: - records = session.scalars(select(WorkflowWebhookTrigger)).all() - assert records == [] - redis.delete.assert_called_once_with(f"{WebhookService.__WEBHOOK_NODE_CACHE_KEY__}:{app.id}:deleted-webhook-node") diff --git a/api/tests/unit_tests/services/test_workflow_agent_preview_identity.py b/api/tests/unit_tests/services/test_workflow_agent_preview_identity.py deleted file mode 100644 index dd4c5c56375..00000000000 --- a/api/tests/unit_tests/services/test_workflow_agent_preview_identity.py +++ /dev/null @@ -1,195 +0,0 @@ -"""Agent preview callers are committed before execution and retain their identity.""" - -import json -from collections.abc import Generator, Iterator -from unittest.mock import patch -from uuid import uuid4 - -import pytest -from flask import Flask -from sqlalchemy import Engine, func, select -from sqlalchemy.orm import sessionmaker - -from core.repositories.sqlalchemy_workflow_node_execution_repository import SQLAlchemyWorkflowNodeExecutionRepository -from core.workflow.nodes.agent_v2 import DifyAgentNode -from core.workflow.nodes.agent_v2.session_store import WorkflowAgentWorkspaceStore -from core.workflow.system_variables import SystemVariableKey, get_system_text -from extensions.ext_database import db -from graphon.engine_events import NodeRunFailedEvent, NodeRunStartedEvent, NodeRunSucceededEvent -from graphon.enums import WorkflowNodeExecutionStatus -from graphon.node_events import NodeRunResult -from libs.datetime_utils import naive_utc_now -from models.workflow import WorkflowNodeExecutionModel, WorkflowNodeExecutionTriggeredFrom -from services.workflow_service import WorkflowService -from tests.unit_tests.config_override import apply_config_overrides -from tests.unit_tests.services.test_workflow_service import TestWorkflowAssociatedDataFactory as Factory - - -class PreviewCancelled(BaseException): - """Cancellation can bypass Exception handlers, like a stopped greenlet.""" - - -@pytest.fixture -def preview_flask_app(sqlite_engine: Engine, monkeypatch: pytest.MonkeyPatch) -> Iterator[Flask]: - apply_config_overrides(monkeypatch, AGENT_BACKEND_BASE_URL="https://agent.test") - app = Flask(__name__) - app.config["SQLALCHEMY_DATABASE_URI"] = sqlite_engine.url.render_as_string(hide_password=False) - db.init_app(app) - with app.app_context(): - yield app - - -@pytest.mark.parametrize("outcome", ["success", "failure", "exception", "cancel"]) -@pytest.mark.usefixtures("preview_flask_app") -def test_agent_preview_commits_runtime_identity_and_updates_same_row(sqlite_engine: Engine, outcome: str) -> None: - tenant_id, app_id, workflow_id, account_id = (str(uuid4()) for _ in range(4)) - node_id = "agent-preview" - app = Factory.create_app(app_id=app_id, tenant_id=tenant_id) - account = Factory.create_account(account_id=account_id) - workflow = Factory.create_workflow( - workflow_id=workflow_id, - tenant_id=tenant_id, - app_id=app_id, - graph={ - "nodes": [ - { - "id": node_id, - "data": {"type": "agent", "version": "2", "agent_node_kind": "dify_agent", "title": "Agent"}, - } - ], - "edges": [], - }, - ) - sessions = sessionmaker(sqlite_engine, expire_on_commit=False) - with sessions.begin() as session: - session.add_all([app, account, workflow]) - repository = SQLAlchemyWorkflowNodeExecutionRepository( - sessions, tenant_id, account, app_id, WorkflowNodeExecutionTriggeredFrom.SINGLE_STEP - ) - observed_created_at = [] - execution_ids: list[str] = [] - - def events(node: DifyAgentNode) -> Generator[NodeRunStartedEvent | NodeRunSucceededEvent | NodeRunFailedEvent]: - execution_id = node.execution_id - execution_ids.append(execution_id) - started_at = naive_utc_now() - # Node.run now emits its typed start before entering the node body. The - # service observer must commit that event before resuming this generator. - yield NodeRunStartedEvent( - id=execution_id, - node_id=node_id, - node_type="agent", - node_version="2", - node_title="Agent", - start_at=started_at, - ) - # A separate transaction must see the committed caller before any - # Agent workspace or tracing request can be dispatched. - variable_pool = node.runtime_state.variable_pool - workflow_run_id = get_system_text(variable_pool, SystemVariableKey.WORKFLOW_EXECUTION_ID) - assert workflow_run_id is None - with sessions.begin() as session: - caller = WorkflowAgentWorkspaceStore._load_execution_by_identity( - session=session, - tenant_id=tenant_id, - app_id=app_id, - workflow_id=workflow_id, - workflow_run_id=workflow_run_id, - node_id=node_id, - node_execution_id=execution_id, - ) - assert caller.id == caller.node_execution_id == execution_id - assert caller.status == WorkflowNodeExecutionStatus.RUNNING - assert caller.finished_at is None - observed_created_at.append(caller.created_at) - # Session materialization writes these independently; final persistence - # must not discard either piece of caller ownership. - caller.agent_workspace_binding_id = "workspace-binding" - caller.process_data = json.dumps({"workflow_agent_binding_id": "workflow-binding"}) - if outcome == "exception": - raise RuntimeError("private runtime error details") - if outcome == "cancel": - raise PreviewCancelled() - status = WorkflowNodeExecutionStatus.SUCCEEDED if outcome == "success" else WorkflowNodeExecutionStatus.FAILED - result = NodeRunResult( - status=status, outputs={"answer": "done"}, error="" if outcome == "success" else "failed" - ) - if outcome == "success": - yield NodeRunSucceededEvent( - id=execution_id, node_id=node_id, node_type="agent", node_run_result=result, start_at=started_at - ) - else: - yield NodeRunFailedEvent( - id=execution_id, - node_id=node_id, - node_type="agent", - node_run_result=result, - start_at=started_at, - error="failed", - ) - - service = WorkflowService(sessions) - with ( - patch("services.workflow_service.WorkflowDraftVariableService.prefill_conversation_variable_default_values"), - patch.object(DifyAgentNode, "run", events), - patch( - "services.workflow_service.DifyCoreRepositoryFactory.create_workflow_node_execution_repository", - return_value=repository, - ), - patch("services.workflow_service.DraftVariableSaver.save"), - patch("services.workflow_service.enqueue_draft_node_execution_trace"), - ): - if outcome in {"exception", "cancel"}: - with pytest.raises(RuntimeError if outcome == "exception" else PreviewCancelled): - service.run_draft_workflow_node(app, workflow, node_id, {}, account) - else: - response = service.run_draft_workflow_node(app, workflow, node_id, {}, account) - assert response.id == response.node_execution_id == execution_ids[0] - - assert len(execution_ids) == 1 - execution_id = execution_ids[0] - with sessions() as session: - assert session.scalar(select(func.count()).select_from(WorkflowNodeExecutionModel)) == 1 - row = session.get(WorkflowNodeExecutionModel, execution_id) - assert row is not None - assert row.id == row.node_execution_id == execution_id - assert row.workflow_run_id is None - assert row.workflow_id == workflow_id - assert row.tenant_id == tenant_id - assert row.app_id == app_id - assert row.created_at == observed_created_at[0] - assert row.finished_at is not None - assert row.elapsed_time >= 0 - assert row.agent_workspace_binding_id == "workspace-binding" - assert row.process_data_dict is not None - assert row.process_data_dict["workflow_agent_binding_id"] == "workflow-binding" - assert row.status == ( - WorkflowNodeExecutionStatus.SUCCEEDED if outcome == "success" else WorkflowNodeExecutionStatus.FAILED - ) - if outcome in {"exception", "cancel"}: - assert row.error == "Agent preview execution interrupted" - - -@pytest.mark.usefixtures("preview_flask_app") -def test_agent_preview_rejects_mismatched_workflow_before_creating_caller(sqlite_engine: Engine) -> None: - app = Factory.create_app(app_id="app", tenant_id="tenant") - account = Factory.create_account() - workflow = Factory.create_workflow( - app_id="another-app", - tenant_id="another-tenant", - graph={ - "nodes": [{"id": "agent", "data": {"type": "agent", "version": "2", "agent_node_kind": "dify_agent"}}], - "edges": [], - }, - ) - sessions = sessionmaker(sqlite_engine, expire_on_commit=False) - with ( - patch( - "services.workflow_service.WorkflowDraftVariableService.prefill_conversation_variable_default_values" - ) as prefill, - pytest.raises(ValueError, match="does not belong"), - ): - WorkflowService(sessions).run_draft_workflow_node(app, workflow, "agent", {}, account) - prefill.assert_not_called() - with sessions() as session: - assert session.scalar(select(func.count()).select_from(WorkflowNodeExecutionModel)) == 0 diff --git a/api/tests/unit_tests/services/test_workflow_collaboration_service.py b/api/tests/unit_tests/services/test_workflow_collaboration_service.py index 8f607a6b683..99fde0bca4c 100644 --- a/api/tests/unit_tests/services/test_workflow_collaboration_service.py +++ b/api/tests/unit_tests/services/test_workflow_collaboration_service.py @@ -15,9 +15,9 @@ from models.base import TypeBase from models.model import App, AppMode from repositories.workflow_collaboration_repository import WorkflowCollaborationRepository from services.workflow_collaboration_service import SYNC_REQUEST_TIMEOUT_SECONDS, WorkflowCollaborationService -from tests.unit_tests.config_override import config_overrides_context -from tests.unit_tests.model_factories import make_app -from tests.unit_tests.services.workflow_collaboration_fixtures import RedisState, seed_session, seed_sessions +from tests.helpers.unit.config_override import config_overrides_context +from tests.helpers.unit.model_factories import make_app +from tests.helpers.unit.services.workflow_collaboration_fixtures import RedisState, seed_session, seed_sessions ServiceFixture = tuple[WorkflowCollaborationService, WorkflowCollaborationRepository, Mock] diff --git a/api/tests/unit_tests/services/test_workflow_draft_variable_service.py b/api/tests/unit_tests/services/test_workflow_draft_variable_service.py new file mode 100644 index 00000000000..1a318171f63 --- /dev/null +++ b/api/tests/unit_tests/services/test_workflow_draft_variable_service.py @@ -0,0 +1,1879 @@ +import dataclasses +import json +import secrets +import unittest +import uuid +from datetime import datetime +from pathlib import Path +from typing import override +from unittest.mock import Mock, patch + +import pytest +from sqlalchemy import Engine, delete, func, select +from sqlalchemy.orm import Session + +from core.workflow.file_reference import build_file_reference +from core.workflow.system_variables import SystemVariableKey +from core.workflow.variable_prefixes import ( + CONVERSATION_VARIABLE_NODE_ID, + ENVIRONMENT_VARIABLE_NODE_ID, + SYSTEM_VARIABLE_NODE_ID, +) +from extensions.ext_storage import storage +from extensions.storage.opendal_storage import OpenDALStorage +from extensions.storage.storage_type import StorageType +from factories.variable_factory import build_segment +from graphon.enums import BuiltinNodeTypes as _workflow_draft_variable_service_BuiltinNodeTypes +from graphon.file import File, FileTransferMethod, FileType +from graphon.nodes import BuiltinNodeTypes as _workflow_draft_variable_service_sqlite_BuiltinNodeTypes +from graphon.variables.segments import StringSegment +from graphon.variables.types import SegmentType +from graphon.variables.variables import StringVariable +from libs import datetime_utils +from libs.datetime_utils import naive_utc_now +from libs.uuid_utils import uuidv7 +from models import Account as _workflow_draft_variable_service_sqlite_Account +from models import Tenant, TenantAccountJoin +from models.account import Account as _workflow_draft_variable_service_Account +from models.account import TenantAccountRole +from models.enums import CreatorUserRole, DraftVariableType +from models.model import UploadFile +from models.workflow import ( + Workflow, + WorkflowDraftVariable, + WorkflowDraftVariableFile, + WorkflowNodeExecutionModel, + is_system_variable_editable, +) +from services.variable_truncator import TruncationResult +from services.workflow_draft_variable_service import ( + DraftVariableSaver, + DraftVarLoader, + VariableResetError, + WorkflowDraftVariableService, + _model_to_insertion_dict, +) + +"""Simplified unit tests for DraftVarLoader focusing on core functionality.""" + + +def _persist_offloaded_variable( + sqlite_session: Session, + *, + node_id: str, + name: str, +) -> WorkflowDraftVariable: + upload_file = UploadFile( + tenant_id="test-tenant-id", + storage_type=StorageType.LOCAL, + key=f"storage/key/{name}.txt", + name=f"{name}.txt", + size=10, + extension=".txt", + mime_type="text/plain", + created_by_role=CreatorUserRole.ACCOUNT, + created_by="test-user-id", + created_at=datetime(2025, 1, 1), + used=True, + ) + variable_file = WorkflowDraftVariableFile( + tenant_id="test-tenant-id", + app_id="test-app-id", + user_id="test-user-id", + upload_file_id=upload_file.id, + size=10, + length=None, + value_type=SegmentType.STRING, + ) + draft_variable = WorkflowDraftVariable.new_node_variable( + app_id="test-app-id", + user_id="test-user-id", + node_id=node_id, + name=name, + value=StringSegment(value="truncated"), + node_execution_id=f"execution-{node_id}", + file_id=variable_file.id, + ) + sqlite_session.add_all([upload_file, variable_file, draft_variable]) + return draft_variable + + +class TestDraftVarLoaderSimpleCases: + class TestDraftVarLoaderSimple: + """Simplified unit tests for DraftVarLoader core methods.""" + + @pytest.fixture + def draft_var_loader(self, sqlite_engine: Engine): + """Create DraftVarLoader instance for testing.""" + return DraftVarLoader( + engine=sqlite_engine, + app_id="test-app-id", + tenant_id="test-tenant-id", + user_id="test-user-id", + fallback_variables=[], + ) + + def test_load_offloaded_variable_object_type_unit(self, draft_var_loader): + """Test _load_offloaded_variable with object type - isolated unit test.""" + # Create mock objects + upload_file = UploadFile( + tenant_id="tenant-id", + storage_type="opendal", + key="storage/key/test.json", + name="test.txt", + size=0, + extension="txt", + mime_type="text/plain", + created_by_role="account", + created_by="account-id", + created_at=datetime.now(), + used=False, + ) + + variable_file = WorkflowDraftVariableFile( + tenant_id="tenant-id", + app_id="app-id", + user_id="user-id", + upload_file_id="upload-file-id", + size=0, + length=0, + value_type=SegmentType.OBJECT, + ) + variable_file.upload_file = upload_file + + draft_var = WorkflowDraftVariable( + id="draft-var-id", + node_id="test-node-id", + name="test_object", + description="test description", + selector=json.dumps(["test-node-id", "test_object"]), + variable_file=variable_file, + ) + + test_object = {"key1": "value1", "key2": 42} + test_json_content = json.dumps(test_object, ensure_ascii=False, separators=(",", ":")) + + with patch("services.workflow_draft_variable_service.storage") as mock_storage: + mock_storage.load.return_value = test_json_content.encode() + + # Execute the method + selector_tuple, variable = draft_var_loader._load_offloaded_variable(draft_var) + + # Verify results + assert selector_tuple == ("test-node-id", "test_object") + assert variable.id == "draft-var-id" + assert variable.name == "test_object" + assert variable.description == "test description" + assert variable.value == test_object + + # Verify method calls + mock_storage.load.assert_called_once_with("storage/key/test.json") + + def test_load_offloaded_variable_missing_variable_file_unit(self, draft_var_loader): + """Test that assertion error is raised when variable_file is None.""" + draft_var = WorkflowDraftVariable( + variable_file=None, + ) + + with pytest.raises(AssertionError): + draft_var_loader._load_offloaded_variable(draft_var) + + def test_load_offloaded_variable_missing_upload_file_unit(self, draft_var_loader): + """Test that assertion error is raised when upload_file is None.""" + variable_file = WorkflowDraftVariableFile( + tenant_id="tenant-id", + app_id="app-id", + user_id="user-id", + upload_file_id="upload-file-id", + size=0, + length=0, + value_type="file", + ) + variable_file.upload_file = None + + draft_var = WorkflowDraftVariable( + variable_file=variable_file, + ) + + with pytest.raises(AssertionError): + draft_var_loader._load_offloaded_variable(draft_var) + + def test_load_variables_empty_selectors_unit(self, draft_var_loader): + """Test load_variables returns empty list for empty selectors.""" + result = draft_var_loader.load_variables([]) + assert result == [] + + def test_selector_to_tuple_unit(self, draft_var_loader): + """Test _selector_to_tuple method.""" + selector = ["node_id", "var_name", "extra_field"] + result = draft_var_loader._selector_to_tuple(selector) + assert result == ("node_id", "var_name") + + def test_load_offloaded_variable_array_type_unit(self, draft_var_loader): + """Test _load_offloaded_variable with array type - isolated unit test.""" + # Create mock objects + upload_file = UploadFile( + tenant_id="tenant-id", + storage_type="opendal", + key="storage/key/test_array.json", + name="test.txt", + size=0, + extension="txt", + mime_type="text/plain", + created_by_role="account", + created_by="account-id", + created_at=datetime.now(), + used=False, + ) + + variable_file = WorkflowDraftVariableFile( + tenant_id="tenant-id", + app_id="app-id", + user_id="user-id", + upload_file_id="upload-file-id", + size=0, + length=0, + value_type=SegmentType.ARRAY_ANY, + ) + variable_file.upload_file = upload_file + + draft_var = WorkflowDraftVariable( + id="draft-var-id", + node_id="test-node-id", + name="test_array", + description="test array description", + selector=json.dumps(["test-node-id", "test_array"]), + variable_file=variable_file, + ) + + test_array = ["item1", 2, True] + test_json_content = json.dumps(test_array) + + with patch("services.workflow_draft_variable_service.storage") as mock_storage: + mock_storage.load.return_value = test_json_content.encode() + # Execute the method + selector_tuple, variable = draft_var_loader._load_offloaded_variable(draft_var) + + # Verify results + assert selector_tuple == ("test-node-id", "test_array") + assert variable.id == "draft-var-id" + assert variable.name == "test_array" + assert variable.description == "test array description" + + # Verify method calls + mock_storage.load.assert_called_once_with("storage/key/test_array.json") + + def test_load_offloaded_variable_file_type_rebuilds_storage_backed_payload(self, draft_var_loader): + upload_file = UploadFile( + tenant_id="tenant-id", + storage_type="opendal", + key="storage/key/test_file.json", + name="test.txt", + size=0, + extension="txt", + mime_type="text/plain", + created_by_role="account", + created_by="account-id", + created_at=datetime.now(), + used=False, + ) + + variable_file = WorkflowDraftVariableFile( + tenant_id="tenant-id", + app_id="app-id", + user_id="user-id", + upload_file_id="upload-file-id", + size=0, + length=0, + value_type=SegmentType.FILE, + ) + variable_file.upload_file = upload_file + + draft_var = WorkflowDraftVariable( + id="draft-var-id", + app_id="app-1", + node_id="test-node-id", + name="test_file", + description="test file description", + ) + draft_var._set_selector(["test-node-id", "test_file"]) + draft_var.variable_file = variable_file + + persisted_file = File( + file_id="file-1", + file_type=FileType.DOCUMENT, + transfer_method=FileTransferMethod.LOCAL_FILE, + reference=build_file_reference(record_id="upload-1", storage_key="legacy-storage-key"), + filename="test.txt", + extension=".txt", + mime_type="text/plain", + size=12, + ) + rebuilt_file = File( + file_id="file-1", + file_type=FileType.DOCUMENT, + transfer_method=FileTransferMethod.LOCAL_FILE, + reference=build_file_reference(record_id="upload-1"), + filename="test.txt", + extension=".txt", + mime_type="text/plain", + size=12, + storage_key="canonical-storage-key", + ) + + raw_file = { + **persisted_file.model_dump(mode="json"), + "tenant_id": "legacy-tenant", + } + + with ( + patch("services.workflow_draft_variable_service.storage") as mock_storage, + patch("models.workflow._resolve_workflow_app_tenant_id", return_value="tenant-1"), + patch("models.workflow.build_file_from_stored_mapping", return_value=rebuilt_file) as rebuild_file, + ): + mock_storage.load.return_value = json.dumps(raw_file).encode() + + selector_tuple, variable = draft_var_loader._load_offloaded_variable(draft_var) + + assert selector_tuple == ("test-node-id", "test_file") + assert variable.id == "draft-var-id" + assert variable.name == "test_file" + assert variable.description == "test file description" + assert variable.value == rebuilt_file + rebuild_file.assert_called_once_with(file_mapping=raw_file, tenant_id="tenant-1") + + @pytest.mark.parametrize( + "sqlite_session", + [(WorkflowDraftVariable, WorkflowDraftVariableFile, UploadFile)], + indirect=True, + ) + def test_load_variables_with_offloaded_variables_unit( + self, + draft_var_loader: DraftVarLoader, + sqlite_session: Session, + ): + """Test load_variables method with mix of regular and offloaded variables.""" + selectors = [["node1", "regular_var"], ["node2", "offloaded_var"]] + regular_draft_var = WorkflowDraftVariable.new_node_variable( + app_id="test-app-id", + user_id="test-user-id", + node_id="node1", + name="regular_var", + value=StringSegment(value="regular_value"), + node_execution_id="execution-node1", + ) + regular_draft_var.description = "regular description" + offloaded_draft_var = _persist_offloaded_variable( + sqlite_session, + node_id="node2", + name="offloaded_var", + ) + distractor = WorkflowDraftVariable.new_node_variable( + app_id="test-app-id", + user_id="another-user", + node_id="node1", + name="regular_var", + value=StringSegment(value="wrong user"), + node_execution_id="execution-distractor", + ) + sqlite_session.add_all([regular_draft_var, distractor]) + sqlite_session.commit() + + offloaded_variable = Mock() + offloaded_variable.id = offloaded_draft_var.id + offloaded_variable.selector = ["node2", "offloaded_var"] + + with ( + patch("services.workflow_draft_variable_service.StorageKeyLoader"), + patch.object( + draft_var_loader, + "_load_offloaded_variable", + return_value=(("node2", "offloaded_var"), offloaded_variable), + ) as load_offloaded, + patch("services.workflow_draft_variable_service.ThreadPoolExecutor") as executor_cls, + ): + executor = executor_cls.return_value.__enter__.return_value + executor.map.side_effect = lambda function, values: [function(value) for value in values] + + result = draft_var_loader.load_variables(selectors) + + assert {variable.id for variable in result} == {regular_draft_var.id, offloaded_draft_var.id} + load_offloaded.assert_called_once() + loaded_offloaded = load_offloaded.call_args.args[0] + assert isinstance(loaded_offloaded, WorkflowDraftVariable) + assert loaded_offloaded.id == offloaded_draft_var.id + assert loaded_offloaded.variable_file is not None + assert loaded_offloaded.variable_file.upload_file is not None + assert loaded_offloaded.variable_file.upload_file.key == "storage/key/offloaded_var.txt" + + @pytest.mark.parametrize( + "sqlite_session", + [(WorkflowDraftVariable, WorkflowDraftVariableFile, UploadFile)], + indirect=True, + ) + def test_load_variables_all_offloaded_variables_unit( + self, + draft_var_loader: DraftVarLoader, + sqlite_session: Session, + ): + """Test load_variables method with only offloaded variables.""" + selectors = [["node1", "offloaded_var1"], ["node2", "offloaded_var2"]] + offloaded_var1 = _persist_offloaded_variable( + sqlite_session, + node_id="node1", + name="offloaded_var1", + ) + offloaded_var2 = _persist_offloaded_variable( + sqlite_session, + node_id="node2", + name="offloaded_var2", + ) + sqlite_session.commit() + + with ( + patch("services.workflow_draft_variable_service.StorageKeyLoader"), + patch("services.workflow_draft_variable_service.ThreadPoolExecutor") as executor_cls, + ): + executor = executor_cls.return_value.__enter__.return_value + executor.map.return_value = [ + (("node1", "offloaded_var1"), Mock()), + (("node2", "offloaded_var2"), Mock()), + ] + + result = draft_var_loader.load_variables(selectors) + + assert len(result) == 2 + executor_cls.assert_called_once_with(max_workers=10) + executor.map.assert_called_once() + loaded_draft_vars = executor.map.call_args.args[1] + assert {variable.id for variable in loaded_draft_vars} == {offloaded_var1.id, offloaded_var2.id} + assert all(variable.variable_file.upload_file is not None for variable in loaded_draft_vars) + + +SQLITE_MODELS = (Workflow, WorkflowDraftVariable, WorkflowDraftVariableFile, WorkflowNodeExecutionModel) + + +class TestWorkflowDraftVariableServiceCases: + pytestmark = [ + pytest.mark.usefixtures("sqlite_session"), + pytest.mark.parametrize("sqlite_session", [SQLITE_MODELS], indirect=True), + ] + + class TestDraftVariableSaver: + def _get_test_app_id(self): + suffix = secrets.token_hex(6) + return f"test_app_id_{suffix}" + + def test__should_variable_be_visible(self, sqlite_session: Session): + mock_user = _workflow_draft_variable_service_Account(name="test", email="test@example.com") + mock_user.id = str(uuid.uuid4()) + test_app_id = self._get_test_app_id() + saver = DraftVariableSaver( + session=sqlite_session, + tenant_id="test-tenant-id", + app_id=test_app_id, + node_id="test_node_id", + node_type=_workflow_draft_variable_service_BuiltinNodeTypes.START, + node_execution_id="test_execution_id", + user=mock_user, + ) + assert ( + saver._should_variable_be_visible( + "123_456", _workflow_draft_variable_service_BuiltinNodeTypes.IF_ELSE, "output" + ) + is False + ) + assert ( + saver._should_variable_be_visible( + "123", _workflow_draft_variable_service_BuiltinNodeTypes.START, "output" + ) + is True + ) + + def test__normalize_variable_for_start_node(self, sqlite_session: Session): + @dataclasses.dataclass(frozen=True) + class TestCase: + name: str + input_node_id: str + input_name: str + expected_node_id: str + expected_name: str + + _node_id = "1747228642872" + cases = [ + TestCase( + name="name with `sys.` prefix should return the system node_id", + input_node_id=_node_id, + input_name="sys.workflow_id", + expected_node_id=SYSTEM_VARIABLE_NODE_ID, + expected_name="workflow_id", + ), + TestCase( + name="name without `sys.` prefix should return the original input node_id", + input_node_id=_node_id, + input_name="start_input", + expected_node_id=_node_id, + expected_name="start_input", + ), + TestCase( + name="name with `env.` prefix should return the environment node_id", + input_node_id=_node_id, + input_name="env.API_KEY", + expected_node_id=ENVIRONMENT_VARIABLE_NODE_ID, + expected_name="API_KEY", + ), + TestCase( + name="name with `conversation.` prefix should return the conversation node_id", + input_node_id=_node_id, + input_name="conversation.session_id", + expected_node_id=CONVERSATION_VARIABLE_NODE_ID, + expected_name="session_id", + ), + TestCase( + name="dummy_variable should return the original input node_id", + input_node_id=_node_id, + input_name="__dummy__", + expected_node_id=_node_id, + expected_name="__dummy__", + ), + ] + + mock_user = _workflow_draft_variable_service_Account(name="Test Account", email="test@example.com") + mock_user.id = str(uuid.uuid4()) + test_app_id = self._get_test_app_id() + saver = DraftVariableSaver( + session=sqlite_session, + tenant_id="test-tenant-id", + app_id=test_app_id, + node_id=_node_id, + node_type=_workflow_draft_variable_service_BuiltinNodeTypes.START, + node_execution_id="test_execution_id", + user=mock_user, + ) + for idx, c in enumerate(cases, 1): + fail_msg = f"Test case {c.name} failed, index={idx}" + node_id, name = saver._normalize_variable_for_start_node(c.input_name) + assert node_id == c.expected_node_id, fail_msg + assert name == c.expected_name, fail_msg + + def test_build_variables_from_start_mapping_rebuilds_system_files(self, sqlite_session: Session): + mock_user = _workflow_draft_variable_service_Account(name="Test Account", email="test@example.com") + mock_user.id = str(uuid.uuid4()) + saver = DraftVariableSaver( + session=sqlite_session, + tenant_id="tenant-1", + app_id=self._get_test_app_id(), + node_id="start", + node_type=_workflow_draft_variable_service_BuiltinNodeTypes.START, + node_execution_id="exec-1", + user=mock_user, + ) + rebuilt_file = File( + file_id="file-1", + file_type=FileType.DOCUMENT, + transfer_method=FileTransferMethod.LOCAL_FILE, + reference="upload-1", + filename="test.txt", + extension=".txt", + mime_type="text/plain", + size=12, + storage_key="canonical-storage-key", + ) + raw_file = { + **rebuilt_file.model_dump(mode="json"), + "tenant_id": "legacy-tenant", + } + + with patch( + "services.workflow_draft_variable_service.build_file_from_stored_mapping", + return_value=rebuilt_file, + ) as rebuild_file: + draft_vars = saver._build_variables_from_start_mapping({"sys.files": [raw_file]}) + + sys_var = draft_vars[0] + assert sys_var.get_value().value[0] == rebuilt_file + rebuild_file.assert_called_once_with(file_mapping=raw_file, tenant_id="tenant-1") + + @pytest.fixture + def draft_saver(self, sqlite_session: Session): + """Create DraftVariableSaver instance with user context.""" + # Create a mock user + mock_user = _workflow_draft_variable_service_Account(name="Test Account", email="test@example.com") + mock_user.id = "test-user-id" + + return DraftVariableSaver( + session=sqlite_session, + tenant_id="test-tenant-id", + app_id="test-app-id", + node_id="test-node-id", + node_type=_workflow_draft_variable_service_BuiltinNodeTypes.LLM, + node_execution_id="test-execution-id", + user=mock_user, + ) + + def test_draft_saver_with_small_variables(self, draft_saver: DraftVariableSaver): + with patch( + "services.workflow_draft_variable_service.DraftVariableSaver._try_offload_large_variable", autospec=True + ) as _mock_try_offload: + _mock_try_offload.return_value = None + mock_segment = StringSegment(value="small value") + draft_var = draft_saver._create_draft_variable(name="small_var", value=mock_segment, visible=True) + + # Should not have large variable metadata + assert draft_var.file_id is None + _mock_try_offload.return_value = None + + def test_draft_saver_with_large_variables(self, draft_saver: DraftVariableSaver): + with patch( + "services.workflow_draft_variable_service.DraftVariableSaver._try_offload_large_variable", autospec=True + ) as _mock_try_offload: + mock_segment = StringSegment(value="small value") + mock_draft_var_file = WorkflowDraftVariableFile( + tenant_id=str(uuidv7()), + app_id=str(uuidv7()), + user_id=str(uuidv7()), + size=1024, + length=10, + value_type=SegmentType.ARRAY_STRING, + upload_file_id=str(uuidv7()), + ) + mock_draft_var_file.id = str(uuidv7()) + + _mock_try_offload.return_value = mock_segment, mock_draft_var_file + draft_var = draft_saver._create_draft_variable(name="small_var", value=mock_segment, visible=True) + + # Should not have large variable metadata + assert draft_var.file_id == mock_draft_var_file.id + + def test_try_offload_large_variable_uses_resource_tenant(self, sqlite_session: Session): + mock_user = _workflow_draft_variable_service_Account(name="Test Account", email="test@example.com") + mock_user.id = "test-user-id" + saver = DraftVariableSaver( + session=sqlite_session, + tenant_id="app-tenant-id", + app_id="test-app-id", + node_id="test-node-id", + node_type=_workflow_draft_variable_service_BuiltinNodeTypes.LLM, + node_execution_id="test-execution-id", + user=mock_user, + ) + upload_file = UploadFile( + tenant_id="app-tenant-id", + storage_type=StorageType.LOCAL, + key="workflow/draft-variable.txt", + name="draft-variable.txt", + size=11, + extension="txt", + mime_type="text/plain", + created_by_role=CreatorUserRole.ACCOUNT, + created_by=mock_user.id, + created_at=naive_utc_now(), + used=True, + ) + sqlite_session.add(upload_file) + sqlite_session.commit() + truncation_result = TruncationResult(result=StringSegment(value="..."), truncated=True) + + with ( + patch( + "services.workflow_draft_variable_service.VariableTruncator.truncate", + return_value=truncation_result, + ), + patch("services.workflow_draft_variable_service.FileService") as file_service_class, + ): + file_service_class.return_value.upload_file.return_value = upload_file + result = saver._try_offload_large_variable("large_var", StringSegment(value="large value")) + + assert result is not None + _, variable_file = result + assert file_service_class.return_value.upload_file.call_args.kwargs["tenant_id"] == "app-tenant-id" + assert variable_file.tenant_id == "app-tenant-id" + sqlite_session.expire_all() + stored_variable_file = sqlite_session.get(WorkflowDraftVariableFile, variable_file.id) + assert stored_variable_file is not None + assert stored_variable_file.upload_file_id == upload_file.id + + @patch("services.workflow_draft_variable_service._batch_upsert_draft_variable", autospec=True) + def test_save_method_integration(self, mock_batch_upsert, draft_saver): + """Test complete save workflow.""" + outputs = {"result": {"data": "test_output"}, "metadata": {"type": "llm_response"}} + + draft_saver.save(outputs=outputs) + + # Should batch upsert draft variables + mock_batch_upsert.assert_called_once() + draft_vars = mock_batch_upsert.call_args[0][1] + assert len(draft_vars) == 2 + + @patch("services.workflow_draft_variable_service._batch_upsert_draft_variable", autospec=True) + def test_start_node_save_persists_sys_timestamp_and_workflow_run_id( + self, mock_batch_upsert, sqlite_session: Session + ): + """Start node should persist common `sys.*` variables, not only `sys.files`.""" + mock_user = _workflow_draft_variable_service_Account(name="Test Account", email="test@example.com") + mock_user.id = "test-user-id" + mock_user.tenant_id = "test-tenant-id" + + saver = DraftVariableSaver( + session=sqlite_session, + tenant_id="test-tenant-id", + app_id="test-app-id", + node_id="start-node-id", + node_type=_workflow_draft_variable_service_BuiltinNodeTypes.START, + node_execution_id="exec-id", + user=mock_user, + ) + + outputs = { + f"{SYSTEM_VARIABLE_NODE_ID}.{SystemVariableKey.TIMESTAMP}": 1700000000, + f"{SYSTEM_VARIABLE_NODE_ID}.{SystemVariableKey.WORKFLOW_EXECUTION_ID}": "run-id-123", + } + + saver.save(outputs=outputs) + + mock_batch_upsert.assert_called_once() + draft_vars = mock_batch_upsert.call_args[0][1] + + # plus one dummy output because there are no non-sys Start inputs + assert len(draft_vars) == 3 + + sys_vars = [v for v in draft_vars if v.node_id == SYSTEM_VARIABLE_NODE_ID] + assert {v.name for v in sys_vars} == { + str(SystemVariableKey.TIMESTAMP), + str(SystemVariableKey.WORKFLOW_EXECUTION_ID), + } + + @patch("services.workflow_draft_variable_service._batch_upsert_draft_variable", autospec=True) + def test_start_node_save_normalizes_reserved_prefix_outputs(self, mock_batch_upsert, sqlite_session: Session): + mock_user = _workflow_draft_variable_service_Account(name="Test Account", email="test@example.com") + mock_user.id = "test-user-id" + mock_user.tenant_id = "test-tenant-id" + + saver = DraftVariableSaver( + session=sqlite_session, + tenant_id="test-tenant-id", + app_id="test-app-id", + node_id="start-node-id", + node_type=_workflow_draft_variable_service_BuiltinNodeTypes.START, + node_execution_id="exec-id", + user=mock_user, + ) + + saver.save( + outputs={ + "env.API_KEY": "secret", + "conversation.session_id": "conversation-1", + "sys.workflow_run_id": "run-id-123", + } + ) + + mock_batch_upsert.assert_called_once() + draft_vars = mock_batch_upsert.call_args[0][1] + + assert len(draft_vars) == 3 + + env_var = next(v for v in draft_vars if v.node_id == ENVIRONMENT_VARIABLE_NODE_ID) + assert env_var.name == "API_KEY" + assert env_var.editable is False + + conversation_var = next(v for v in draft_vars if v.node_id == CONVERSATION_VARIABLE_NODE_ID) + assert conversation_var.name == "session_id" + assert conversation_var.node_execution_id is None + + sys_var = next(v for v in draft_vars if v.node_id == SYSTEM_VARIABLE_NODE_ID) + assert sys_var.name == str(SystemVariableKey.WORKFLOW_EXECUTION_ID) + + class TestWorkflowDraftVariableService: + def _get_test_app_id(self): + suffix = secrets.token_hex(6) + return f"test_app_id_{suffix}" + + def _create_test_workflow(self, app_id: str) -> Workflow: + """Create a real Workflow instance for testing""" + return Workflow.new( + tenant_id="test_tenant_id", + app_id=app_id, + type="workflow", + version="draft", + graph='{"nodes": [], "edges": []}', + features="{}", + created_by="test_user_id", + environment_variables=[], + conversation_variables=[], + rag_pipeline_variables=[], + ) + + def test_list_variables_without_values_excludes_node_ids(self, sqlite_session: Session): + service = WorkflowDraftVariableService(sqlite_session) + variable = WorkflowDraftVariable.new_node_variable( + app_id="app-1", + node_id="node-1", + name="output", + value=StringSegment(value="value"), + node_execution_id="execution-1", + ) + variable.user_id = "user-1" + excluded_system = WorkflowDraftVariable.new_sys_variable( + app_id="app-1", name="query", value=StringSegment(value="hidden"), node_execution_id="execution-1" + ) + excluded_system.user_id = "user-1" + other_user = WorkflowDraftVariable.new_node_variable( + app_id="app-1", + node_id="node-2", + name="output", + value=StringSegment(value="other"), + node_execution_id="execution-2", + ) + other_user.user_id = "user-2" + sqlite_session.add_all([variable, excluded_system, other_user]) + sqlite_session.commit() + + result = service.list_variables_without_values( + app_id="app-1", + page=1, + limit=20, + user_id="user-1", + exclude_node_ids={SYSTEM_VARIABLE_NODE_ID, CONVERSATION_VARIABLE_NODE_ID}, + ) + + assert result.total == 1 + assert [item.id for item in result.variables] == [variable.id] + + def test_reset_conversation_variable(self, sqlite_session: Session): + """Test resetting a conversation variable""" + service = WorkflowDraftVariableService(sqlite_session) + + test_app_id = self._get_test_app_id() + workflow = self._create_test_workflow(test_app_id) + + # Create real conversation variable + test_value = StringSegment(value="test_value") + variable = WorkflowDraftVariable.new_conversation_variable( + app_id=test_app_id, name="test_var", value=test_value, description="Test conversation variable" + ) + + # Mock the _reset_conv_var method + expected_result = WorkflowDraftVariable.new_conversation_variable( + app_id=test_app_id, + name="test_var", + value=StringSegment(value="reset_value"), + ) + with patch.object( + service, "_reset_conv_var", return_value=expected_result, autospec=True + ) as mock_reset_conv: + result = service.reset_variable(workflow, variable) + + mock_reset_conv.assert_called_once_with(workflow, variable) + assert result == expected_result + + def test_reset_node_variable_with_no_execution_id(self, sqlite_session: Session): + """Test resetting a node variable with no execution ID - should delete variable""" + service = WorkflowDraftVariableService(sqlite_session) + + test_app_id = self._get_test_app_id() + workflow = self._create_test_workflow(test_app_id) + + # Create real node variable with no execution ID + test_value = StringSegment(value="test_value") + variable = WorkflowDraftVariable.new_node_variable( + app_id=test_app_id, + node_id="test_node_id", + name="test_var", + value=test_value, + node_execution_id="exec-id", # Set initially + ) + # Manually set to None to simulate the test condition + variable.node_execution_id = None + sqlite_session.add(variable) + sqlite_session.commit() + + result = service._reset_node_var_or_sys_var(workflow, variable) + + # Should delete the variable and return None + assert sqlite_session.get(WorkflowDraftVariable, variable.id) is None + assert result is None + + def test_reset_node_variable_with_missing_execution_record(self, sqlite_session: Session): + """Test resetting a node variable when execution record doesn't exist""" + service = WorkflowDraftVariableService(sqlite_session) + + test_app_id = self._get_test_app_id() + workflow = self._create_test_workflow(test_app_id) + + # Create real node variable with execution ID + test_value = StringSegment(value="test_value") + variable = WorkflowDraftVariable.new_node_variable( + app_id=test_app_id, + node_id="test_node_id", + name="test_var", + value=test_value, + node_execution_id="exec-id", + ) + sqlite_session.add(variable) + sqlite_session.commit() + + result = service._reset_node_var_or_sys_var(workflow, variable) + + assert sqlite_session.get(WorkflowDraftVariable, variable.id) is None + assert result is None + + def test_reset_node_variable_with_valid_execution_record(self, sqlite_session: Session): + """Reset a node variable from its execution output and flush the restored value.""" + service = WorkflowDraftVariableService(sqlite_session) + + execution = WorkflowNodeExecutionModel(outputs=json.dumps({"test_var": "output_value"})) + + # Mock the repository to return the execution record + service._api_node_execution_repo = Mock() + service._api_node_execution_repo.get_execution_by_id.return_value = execution + + test_app_id = self._get_test_app_id() + workflow = self._create_test_workflow(test_app_id) + + # Create real node variable with execution ID + test_value = StringSegment(value="original_value") + variable = WorkflowDraftVariable.new_node_variable( + app_id=test_app_id, + node_id="test_node_id", + name="test_var", + value=test_value, + node_execution_id="exec-id", + ) + sqlite_session.add(variable) + sqlite_session.commit() + + # Mock workflow methods + mock_node_config = {"type": "test_node"} + with ( + patch.object(sqlite_session, "flush", wraps=sqlite_session.flush) as flush, + patch.object(workflow, "get_node_config_by_id", return_value=mock_node_config, autospec=True), + patch.object( + workflow, + "get_node_type_from_node_config", + return_value=_workflow_draft_variable_service_BuiltinNodeTypes.LLM, + autospec=True, + ), + ): + result = service._reset_node_var_or_sys_var(workflow, variable) + + # Verify last_edited_at was reset + assert variable.last_edited_at is None + flush.assert_called() + # Should return the updated variable + assert result == variable + + def test_reset_non_editable_system_variable_raises_error(self, sqlite_session: Session): + """Test that resetting a non-editable system variable raises an error""" + service = WorkflowDraftVariableService(sqlite_session) + + test_app_id = self._get_test_app_id() + workflow = self._create_test_workflow(test_app_id) + + # Create a non-editable system variable (workflow_id is not editable) + test_value = StringSegment(value="test_workflow_id") + variable = WorkflowDraftVariable.new_sys_variable( + app_id=test_app_id, + name="workflow_id", # This is not in _EDITABLE_SYSTEM_VARIABLE + value=test_value, + node_execution_id="exec-id", + editable=False, # Non-editable system variable + ) + + with pytest.raises(VariableResetError) as exc_info: + service.reset_variable(workflow, variable) + assert "cannot reset system variable" in str(exc_info.value) + assert f"variable_id={variable.id}" in str(exc_info.value) + + def test_reset_editable_system_variable_succeeds(self, sqlite_session: Session): + """Test that resetting an editable system variable succeeds""" + service = WorkflowDraftVariableService(sqlite_session) + + test_app_id = self._get_test_app_id() + workflow = self._create_test_workflow(test_app_id) + + # Create an editable system variable (files is editable) + test_value = StringSegment(value="[]") + variable = WorkflowDraftVariable.new_sys_variable( + app_id=test_app_id, + name="files", # This is in _EDITABLE_SYSTEM_VARIABLE + value=test_value, + node_execution_id="exec-id", + editable=True, # Editable system variable + ) + sqlite_session.add(variable) + sqlite_session.commit() + + execution = WorkflowNodeExecutionModel(outputs=json.dumps({"sys.files": "[]"})) + + # Mock the repository to return the execution record + service._api_node_execution_repo = Mock() + service._api_node_execution_repo.get_execution_by_id.return_value = execution + + with patch.object(sqlite_session, "flush", wraps=sqlite_session.flush) as flush: + result = service._reset_node_var_or_sys_var(workflow, variable) + + # Should succeed and return the variable + assert result == variable + assert variable.last_edited_at is None + flush.assert_called() + + def test_reset_query_system_variable_succeeds(self, sqlite_session: Session): + """Test that resetting query system variable (another editable one) succeeds""" + service = WorkflowDraftVariableService(sqlite_session) + + test_app_id = self._get_test_app_id() + workflow = self._create_test_workflow(test_app_id) + + # Create an editable system variable (query is editable) + test_value = StringSegment(value="original query") + variable = WorkflowDraftVariable.new_sys_variable( + app_id=test_app_id, + name="query", # This is in _EDITABLE_SYSTEM_VARIABLE + value=test_value, + node_execution_id="exec-id", + editable=True, # Editable system variable + ) + sqlite_session.add(variable) + sqlite_session.commit() + + execution = WorkflowNodeExecutionModel(outputs=json.dumps({"sys.query": "reset query"})) + + # Mock the repository to return the execution record + service._api_node_execution_repo = Mock() + service._api_node_execution_repo.get_execution_by_id.return_value = execution + + with patch.object(sqlite_session, "flush", wraps=sqlite_session.flush) as flush: + result = service._reset_node_var_or_sys_var(workflow, variable) + + # Should succeed and return the variable + assert result == variable + assert variable.last_edited_at is None + flush.assert_called() + + def test_system_variable_editability_check(self): + """Test the system variable editability function directly""" + # Test editable system variables + assert is_system_variable_editable("files") is True + assert is_system_variable_editable("query") is True + + # Test non-editable system variables + assert is_system_variable_editable("workflow_id") is False + assert is_system_variable_editable("conversation_id") is False + assert is_system_variable_editable("user_id") is False + + def test_workflow_draft_variable_factory_methods(self): + """Test that factory methods create proper instances""" + test_app_id = self._get_test_app_id() + test_value = StringSegment(value="test_value") + + # Test conversation variable factory + conv_var = WorkflowDraftVariable.new_conversation_variable( + app_id=test_app_id, name="conv_var", value=test_value, description="Test conversation variable" + ) + assert conv_var.get_variable_type() == DraftVariableType.CONVERSATION + assert conv_var.editable is True + assert conv_var.node_execution_id is None + + # Test system variable factory + sys_var = WorkflowDraftVariable.new_sys_variable( + app_id=test_app_id, name="workflow_id", value=test_value, node_execution_id="exec-id", editable=False + ) + assert sys_var.get_variable_type() == DraftVariableType.SYS + assert sys_var.editable is False + assert sys_var.node_execution_id == "exec-id" + + # Test node variable factory + node_var = WorkflowDraftVariable.new_node_variable( + app_id=test_app_id, + node_id="node-id", + name="node_var", + value=test_value, + node_execution_id="exec-id", + visible=True, + editable=True, + ) + assert node_var.get_variable_type() == DraftVariableType.NODE + assert node_var.visible is True + assert node_var.editable is True + assert node_var.node_execution_id == "exec-id" + + class TestModelToInsertionDict: + """Reproduce two production errors in _model_to_insertion_dict / _new().""" + + def test_visible_and_is_default_value_always_present(self): + """Problem 1: _new() did not set visible/is_default_value, causing + inconsistent dict keys across rows in multi-row INSERT and missing + is_default_value in the insertion dict entirely. + """ + conv_var = WorkflowDraftVariable.new_conversation_variable( + app_id="app-1", + name="counter", + value=StringSegment(value="0"), + ) + # _new() should explicitly set these fields so they are not None + assert conv_var.visible is not None + assert conv_var.is_default_value is not None + + d = _model_to_insertion_dict(conv_var) + # visible must appear in every row's dict + assert "visible" in d + # is_default_value must always be present + assert "is_default_value" in d + + def test_description_passthrough(self): + """_model_to_insertion_dict passes description as-is; + length validation is enforced earlier in build_conversation_variable_from_mapping. + """ + desc = "a" * 200 + conv_var = WorkflowDraftVariable.new_conversation_variable( + app_id="app-1", + name="counter", + value=StringSegment(value="0"), + description=desc, + ) + d = _model_to_insertion_dict(conv_var) + assert d["description"] == desc + + def test_is_default_value_omitted_when_none(self): + conv_var = WorkflowDraftVariable.new_conversation_variable( + app_id="app-1", + name="counter", + value=StringSegment(value="0"), + ) + conv_var.is_default_value = None + + d = _model_to_insertion_dict(conv_var) + + assert "is_default_value" not in d + + +"""SQLite-backed workflow draft-variable service coverage.""" + + +def _persist_account(*, engine: Engine, tenant_id: str) -> _workflow_draft_variable_service_sqlite_Account: + tenant = Tenant(name="draft-variable-tenant") + tenant.id = tenant_id + account = _workflow_draft_variable_service_sqlite_Account( + email=f"draft-{uuid.uuid4()}@example.com", name="Draft Variable User" + ) + membership = TenantAccountJoin(tenant_id=tenant.id, account_id=account.id, role=TenantAccountRole.OWNER) + with Session(bind=engine, expire_on_commit=False) as session: + session.add_all([tenant, account, membership]) + session.commit() + return account + + +class TestWorkflowDraftVariableServiceSqliteCases: + @staticmethod + @pytest.fixture(autouse=True) + def _bind_storage( + monkeypatch: pytest.MonkeyPatch, + tmp_path: Path, + ) -> None: + monkeypatch.setattr( + storage, + "storage_runner", + OpenDALStorage(scheme="fs", root=str(tmp_path / "storage")), + raising=False, + ) + + class TestWorkflowDraftVariableService(unittest.TestCase): + _test_app_id: str + _session: Session + _test_user_id: str + _node1_id = "test_node_1" + _node2_id = "test_node_2" + _node_exec_id = str(uuid.uuid4()) + + @pytest.fixture(autouse=True) + def _bind_session(self, sqlite_session: Session) -> None: + self._session = sqlite_session + + @override + def setUp(self) -> None: + self._test_app_id = str(uuid.uuid4()) + self._test_user_id = str(uuid.uuid4()) + sys_var = WorkflowDraftVariable.new_sys_variable( + app_id=self._test_app_id, + user_id=self._test_user_id, + name="sys_var", + value=build_segment("sys_value"), + node_execution_id=self._node_exec_id, + ) + conv_var = WorkflowDraftVariable.new_conversation_variable( + app_id=self._test_app_id, + user_id=self._test_user_id, + name="conv_var", + value=build_segment("conv_value"), + ) + node2_vars = [ + WorkflowDraftVariable.new_node_variable( + app_id=self._test_app_id, + user_id=self._test_user_id, + node_id=self._node2_id, + name="int_var", + value=build_segment(1), + visible=False, + node_execution_id=self._node_exec_id, + ), + WorkflowDraftVariable.new_node_variable( + app_id=self._test_app_id, + user_id=self._test_user_id, + node_id=self._node2_id, + name="str_var", + value=build_segment("str_value"), + visible=True, + node_execution_id=self._node_exec_id, + ), + ] + node1_var = WorkflowDraftVariable.new_node_variable( + app_id=self._test_app_id, + user_id=self._test_user_id, + node_id=self._node1_id, + name="str_var", + value=build_segment("str_value"), + visible=True, + node_execution_id=self._node_exec_id, + ) + _variables = list(node2_vars) + _variables.extend( + [ + node1_var, + sys_var, + conv_var, + ] + ) + + self._session.add_all(_variables) + self._session.flush() + self._variable_ids = [v.id for v in _variables] + self._node1_str_var_id = node1_var.id + self._sys_var_id = sys_var.id + self._conv_var_id = conv_var.id + self._node2_var_ids = [v.id for v in node2_vars] + + def _get_test_srv(self) -> WorkflowDraftVariableService: + return WorkflowDraftVariableService(session=self._session) + + @override + def tearDown(self) -> None: + self._session.rollback() + + def test_list_variables(self) -> None: + srv = self._get_test_srv() + var_list = srv.list_variables_without_values(self._test_app_id, page=1, limit=2, user_id=self._test_user_id) + assert var_list.total == 5 + assert len(var_list.variables) == 2 + page1_var_ids = {v.id for v in var_list.variables} + assert page1_var_ids.issubset(self._variable_ids) + + var_list_2 = srv.list_variables_without_values( + self._test_app_id, page=2, limit=2, user_id=self._test_user_id + ) + assert var_list_2.total is None + assert len(var_list_2.variables) == 2 + page2_var_ids = {v.id for v in var_list_2.variables} + assert page2_var_ids.isdisjoint(page1_var_ids) + assert page2_var_ids.issubset(self._variable_ids) + + def test_get_node_variable(self) -> None: + srv = self._get_test_srv() + node_var = srv.get_node_variable(self._test_app_id, self._node1_id, "str_var", user_id=self._test_user_id) + assert node_var is not None + assert node_var.id == self._node1_str_var_id + assert node_var.name == "str_var" + assert node_var.get_value() == build_segment("str_value") + + def test_get_system_variable(self) -> None: + srv = self._get_test_srv() + sys_var = srv.get_system_variable(self._test_app_id, "sys_var", user_id=self._test_user_id) + assert sys_var is not None + assert sys_var.id == self._sys_var_id + assert sys_var.name == "sys_var" + assert sys_var.get_value() == build_segment("sys_value") + + def test_get_conversation_variable(self) -> None: + srv = self._get_test_srv() + conv_var = srv.get_conversation_variable(self._test_app_id, "conv_var", user_id=self._test_user_id) + assert conv_var is not None + assert conv_var.id == self._conv_var_id + assert conv_var.name == "conv_var" + assert conv_var.get_value() == build_segment("conv_value") + + def test_delete_node_variables(self) -> None: + srv = self._get_test_srv() + srv.delete_node_variables(self._test_app_id, self._node2_id, user_id=self._test_user_id) + node2_var_count = self._session.scalar( + select(func.count()) + .select_from(WorkflowDraftVariable) + .where( + WorkflowDraftVariable.app_id == self._test_app_id, + WorkflowDraftVariable.node_id == self._node2_id, + WorkflowDraftVariable.user_id == self._test_user_id, + ) + ) + assert node2_var_count == 0 + + def test_delete_variable(self) -> None: + srv = self._get_test_srv() + node_1_var = self._session.scalars( + select(WorkflowDraftVariable).where(WorkflowDraftVariable.id == self._node1_str_var_id) + ).one() + srv.delete_variable(node_1_var) + exists = bool( + self._session.scalars( + select(WorkflowDraftVariable).where(WorkflowDraftVariable.id == self._node1_str_var_id) + ).first() + ) + assert exists is False + + def test__list_node_variables(self) -> None: + srv = self._get_test_srv() + node_vars = srv._list_node_variables(self._test_app_id, self._node2_id, user_id=self._test_user_id) + assert len(node_vars.variables) == 2 + assert {v.id for v in node_vars.variables} == set(self._node2_var_ids) + + def test_get_draft_variables_by_selectors(self) -> None: + srv = self._get_test_srv() + selectors = [ + [self._node1_id, "str_var"], + [self._node2_id, "str_var"], + [self._node2_id, "int_var"], + ] + variables = srv.get_draft_variables_by_selectors(self._test_app_id, selectors, user_id=self._test_user_id) + assert len(variables) == 3 + assert {v.id for v in variables} == {self._node1_str_var_id} | set(self._node2_var_ids) + + class TestDraftVariableLoader(unittest.TestCase): + _test_app_id: str + _test_tenant_id: str + _engine: Engine + + _node1_id = "test_loader_node_1" + _node_exec_id = str(uuid.uuid4()) + + @pytest.fixture(autouse=True) + def _bind_engine(self, sqlite_engine: Engine) -> None: + self._engine = sqlite_engine + + @override + def setUp(self) -> None: + self._test_app_id = str(uuid.uuid4()) + self._test_tenant_id = str(uuid.uuid4()) + self._test_user_id = str(uuid.uuid4()) + sys_var = WorkflowDraftVariable.new_sys_variable( + app_id=self._test_app_id, + user_id=self._test_user_id, + name="sys_var", + value=build_segment("sys_value"), + node_execution_id=self._node_exec_id, + ) + conv_var = WorkflowDraftVariable.new_conversation_variable( + app_id=self._test_app_id, + user_id=self._test_user_id, + name="conv_var", + value=build_segment("conv_value"), + ) + node_var = WorkflowDraftVariable.new_node_variable( + app_id=self._test_app_id, + user_id=self._test_user_id, + node_id=self._node1_id, + name="str_var", + value=build_segment("str_value"), + visible=True, + node_execution_id=self._node_exec_id, + ) + _variables = [ + node_var, + sys_var, + conv_var, + ] + + with Session(bind=self._engine, expire_on_commit=False) as session: + session.add_all(_variables) + session.flush() + session.commit() + self._variable_ids = [v.id for v in _variables] + self._node_var_id = node_var.id + self._sys_var_id = sys_var.id + self._conv_var_id = conv_var.id + + @override + def tearDown(self) -> None: + with Session(bind=self._engine, expire_on_commit=False) as session: + session.execute(delete(WorkflowDraftVariable).where(WorkflowDraftVariable.app_id == self._test_app_id)) + session.commit() + + def test_variable_loader_with_empty_selector(self) -> None: + var_loader = DraftVarLoader( + engine=self._engine, + app_id=self._test_app_id, + tenant_id=self._test_tenant_id, + user_id=self._test_user_id, + ) + variables = var_loader.load_variables([]) + assert len(variables) == 0 + + def test_variable_loader_with_non_empty_selector(self) -> None: + var_loader = DraftVarLoader( + engine=self._engine, + app_id=self._test_app_id, + tenant_id=self._test_tenant_id, + user_id=self._test_user_id, + ) + variables = var_loader.load_variables( + [ + [SYSTEM_VARIABLE_NODE_ID, "sys_var"], + [CONVERSATION_VARIABLE_NODE_ID, "conv_var"], + [self._node1_id, "str_var"], + ] + ) + assert len(variables) == 3 + conv_var = next(v for v in variables if v.selector[0] == CONVERSATION_VARIABLE_NODE_ID) + assert conv_var.id == self._conv_var_id + sys_var = next(v for v in variables if v.selector[0] == SYSTEM_VARIABLE_NODE_ID) + assert sys_var.id == self._sys_var_id + node1_var = next(v for v in variables if v.selector[0] == self._node1_id) + assert node1_var.id == self._node_var_id + + def test_load_offloaded_variable_string_type_integration(self) -> None: + """Test _load_offloaded_variable with string type using DraftVariableSaver for data creation.""" + account = _persist_account(engine=self._engine, tenant_id=self._test_tenant_id) + + # Create a large string that will be offloaded + test_content = "x" * 15000 # Create a string larger than LARGE_VARIABLE_THRESHOLD (10KB) + large_string_segment = StringSegment(value=test_content) + + node_execution_id = str(uuid.uuid4()) + + try: + with Session(bind=self._engine, expire_on_commit=False) as session: + # Use DraftVariableSaver to create offloaded variable (this mimics production) + saver = DraftVariableSaver( + session=session, + tenant_id=self._test_tenant_id, + app_id=self._test_app_id, + node_id="test_offload_node", + node_type=_workflow_draft_variable_service_sqlite_BuiltinNodeTypes.LLM, # Use a real node type + node_execution_id=node_execution_id, + user=account, + ) + + # Save the variable - this will trigger offloading due to large size + saver.save(outputs={"offloaded_string_var": large_string_segment}) + session.commit() + + # Now test loading using DraftVarLoader + var_loader = DraftVarLoader( + engine=self._engine, + app_id=self._test_app_id, + tenant_id=self._test_tenant_id, + user_id=account.id, + ) + + # Load the variable using the standard workflow + variables = var_loader.load_variables([["test_offload_node", "offloaded_string_var"]]) + + # Verify results + assert len(variables) == 1 + loaded_variable = variables[0] + assert loaded_variable.name == "offloaded_string_var" + assert loaded_variable.selector == ["test_offload_node", "offloaded_string_var"] + assert loaded_variable.value == test_content + + finally: + # Clean up - delete all draft variables for this app + with Session(bind=self._engine) as session: + service = WorkflowDraftVariableService(session) + service.delete_app_workflow_variables(self._test_app_id) + session.commit() + + def test_load_offloaded_variable_object_type_integration(self) -> None: + """Test _load_offloaded_variable with object type using real storage and service.""" + + # Create a test object + test_object = {"key1": "value1", "key2": 42, "nested": {"inner": "data"}} + test_json = json.dumps(test_object, ensure_ascii=False, separators=(",", ":")) + content_bytes = test_json.encode() + + # Create an upload file record + upload_file = UploadFile( + tenant_id=self._test_tenant_id, + storage_type=StorageType.LOCAL, + key=f"test_offload_{uuid.uuid4()}.json", + name="test_offload.json", + size=len(content_bytes), + extension="json", + mime_type="application/json", + created_by_role=CreatorUserRole.ACCOUNT, + created_by=str(uuid.uuid4()), + created_at=datetime_utils.naive_utc_now(), + used=True, + used_by=str(uuid.uuid4()), + used_at=datetime_utils.naive_utc_now(), + ) + + # Store the content in storage + storage.save(upload_file.key, content_bytes) + + # Create a variable file record + variable_file = WorkflowDraftVariableFile( + upload_file_id=upload_file.id, + value_type=SegmentType.OBJECT, + tenant_id=self._test_tenant_id, + app_id=self._test_app_id, + user_id=str(uuid.uuid4()), + size=len(content_bytes), + length=len(test_object), + created_at=datetime_utils.naive_utc_now(), + ) + + try: + with Session(bind=self._engine, expire_on_commit=False) as session: + # Add upload file and variable file first to get their IDs + session.add_all([upload_file, variable_file]) + session.flush() # This generates the IDs + + # Now create the offloaded draft variable with the correct file_id + offloaded_var = WorkflowDraftVariable.new_node_variable( + app_id=self._test_app_id, + user_id=self._test_user_id, + node_id="test_offload_node", + name="offloaded_object_var", + value=build_segment({"truncated": True}), + visible=True, + node_execution_id=str(uuid.uuid4()), + ) + offloaded_var.file_id = variable_file.id + + session.add(offloaded_var) + session.flush() + session.commit() + + # Use the service method that properly preloads relationships + service = WorkflowDraftVariableService(session) + draft_vars = service.get_draft_variables_by_selectors( + self._test_app_id, + [["test_offload_node", "offloaded_object_var"]], + user_id=self._test_user_id, + ) + + assert len(draft_vars) == 1 + loaded_var = draft_vars[0] + assert loaded_var.is_truncated() + + # Create DraftVarLoader and test loading + var_loader = DraftVarLoader( + engine=self._engine, + app_id=self._test_app_id, + tenant_id=self._test_tenant_id, + user_id=self._test_user_id, + ) + + # Test the _load_offloaded_variable method + selector_tuple, variable = var_loader._load_offloaded_variable(loaded_var) + + # Verify the results + assert selector_tuple == ("test_offload_node", "offloaded_object_var") + assert variable.id == loaded_var.id + assert variable.name == "offloaded_object_var" + assert variable.value == test_object + + finally: + # Clean up + with Session(bind=self._engine) as session: + # Query and delete by ID to ensure they're tracked in this session + session.execute(delete(WorkflowDraftVariable).where(WorkflowDraftVariable.id == offloaded_var.id)) + session.execute( + delete(WorkflowDraftVariableFile).where(WorkflowDraftVariableFile.id == variable_file.id) + ) + session.execute(delete(UploadFile).where(UploadFile.id == upload_file.id)) + session.commit() + storage.delete(upload_file.key) + + def test_load_variables_with_offloaded_variables_integration(self) -> None: + """Test load_variables method with mix of regular and offloaded variables using real storage.""" + # Create a regular variable (already exists from setUp) + # Create offloaded variable content + test_content = "This is offloaded content for integration test" + content_bytes = test_content.encode() + + # Create upload file record + upload_file = UploadFile( + tenant_id=self._test_tenant_id, + storage_type=StorageType.LOCAL, + key=f"test_integration_{uuid.uuid4()}.txt", + name="test_integration.txt", + size=len(content_bytes), + extension="txt", + mime_type="text/plain", + created_by_role=CreatorUserRole.ACCOUNT, + created_by=str(uuid.uuid4()), + created_at=datetime_utils.naive_utc_now(), + used=True, + used_by=str(uuid.uuid4()), + used_at=datetime_utils.naive_utc_now(), + ) + + # Store the content + storage.save(upload_file.key, content_bytes) + + # Create variable file + variable_file = WorkflowDraftVariableFile( + upload_file_id=upload_file.id, + value_type=SegmentType.STRING, + tenant_id=self._test_tenant_id, + app_id=self._test_app_id, + user_id=str(uuid.uuid4()), + size=len(content_bytes), + length=None, + created_at=datetime_utils.naive_utc_now(), + ) + + try: + with Session(bind=self._engine, expire_on_commit=False) as session: + # Add upload file and variable file first to get their IDs + session.add_all([upload_file, variable_file]) + session.flush() # This generates the IDs + + # Now create the offloaded draft variable with the correct file_id + offloaded_var = WorkflowDraftVariable.new_node_variable( + app_id=self._test_app_id, + user_id=self._test_user_id, + node_id="test_integration_node", + name="offloaded_integration_var", + value=build_segment("truncated"), + visible=True, + node_execution_id=str(uuid.uuid4()), + ) + offloaded_var.file_id = variable_file.id + + session.add(offloaded_var) + session.flush() + session.commit() + + # Test load_variables with both regular and offloaded variables + # This method should handle the relationship preloading internally + var_loader = DraftVarLoader( + engine=self._engine, + app_id=self._test_app_id, + tenant_id=self._test_tenant_id, + user_id=self._test_user_id, + ) + + variables = var_loader.load_variables( + [ + [SYSTEM_VARIABLE_NODE_ID, "sys_var"], # Regular variable from setUp + ["test_integration_node", "offloaded_integration_var"], # Offloaded variable + ] + ) + + # Verify results + assert len(variables) == 2 + + # Find regular variable + regular_var = next(v for v in variables if v.selector[0] == SYSTEM_VARIABLE_NODE_ID) + assert regular_var.id == self._sys_var_id + assert regular_var.value == "sys_value" + + # Find offloaded variable + offloaded_loaded_var = next(v for v in variables if v.selector[0] == "test_integration_node") + assert offloaded_loaded_var.id == offloaded_var.id + assert offloaded_loaded_var.value == test_content + + finally: + # Clean up + with Session(bind=self._engine) as session: + # Query and delete by ID to ensure they're tracked in this session + session.execute(delete(WorkflowDraftVariable).where(WorkflowDraftVariable.id == offloaded_var.id)) + session.execute( + delete(WorkflowDraftVariableFile).where(WorkflowDraftVariableFile.id == variable_file.id) + ) + session.execute(delete(UploadFile).where(UploadFile.id == upload_file.id)) + session.commit() + storage.delete(upload_file.key) + + class TestWorkflowDraftVariableServiceResetVariable(unittest.TestCase): + """Integration tests for reset_variable functionality using real database""" + + _test_app_id: str + _test_tenant_id: str + _test_workflow_id: str + _session: Session + _engine: Engine + _node_id = "test_reset_node" + _node_exec_id: str + _workflow_node_exec_id: str + + @pytest.fixture(autouse=True) + def _bind_sqlite(self, sqlite_session: Session, sqlite_engine: Engine) -> None: + self._session = sqlite_session + self._engine = sqlite_engine + + @override + def setUp(self) -> None: + self._test_app_id = str(uuid.uuid4()) + self._test_tenant_id = str(uuid.uuid4()) + self._test_workflow_id = str(uuid.uuid4()) + self._test_user_id = str(uuid.uuid4()) + self._node_exec_id = str(uuid.uuid4()) + self._workflow_node_exec_id = str(uuid.uuid4()) + + # Create a workflow node execution record with outputs + # Note: The WorkflowNodeExecutionModel.id should match the node_execution_id in WorkflowDraftVariable + self._workflow_node_execution = WorkflowNodeExecutionModel( + id=self._node_exec_id, # This should match the node_execution_id in the variable + tenant_id=self._test_tenant_id, + app_id=self._test_app_id, + workflow_id=self._test_workflow_id, + triggered_from="workflow-run", + workflow_run_id=str(uuid.uuid4()), + index=1, + node_execution_id=str(uuid.uuid4()), + node_id=self._node_id, + node_type=_workflow_draft_variable_service_sqlite_BuiltinNodeTypes.LLM, + title="Test Node", + inputs='{"input": "test input"}', + process_data='{"test_var": "process_value", "other_var": "other_process"}', + outputs='{"test_var": "output_value", "other_var": "other_output"}', + status="succeeded", + elapsed_time=1.5, + created_by_role=CreatorUserRole.ACCOUNT, + created_by=str(uuid.uuid4()), + ) + + # Create conversation variables for the workflow + self._conv_variables = [ + StringVariable( + id=str(uuid.uuid4()), + name="conv_var_1", + description="Test conversation variable 1", + value="default_value_1", + ), + StringVariable( + id=str(uuid.uuid4()), + name="conv_var_2", + description="Test conversation variable 2", + value="default_value_2", + ), + ] + + # Create test variables + self._node_var_with_exec = WorkflowDraftVariable.new_node_variable( + app_id=self._test_app_id, + user_id=self._test_user_id, + node_id=self._node_id, + name="test_var", + value=build_segment("old_value"), + node_execution_id=self._node_exec_id, + ) + self._node_var_with_exec.last_edited_at = datetime_utils.naive_utc_now() + + self._node_var_without_exec = WorkflowDraftVariable.new_node_variable( + app_id=self._test_app_id, + user_id=self._test_user_id, + node_id=self._node_id, + name="no_exec_var", + value=build_segment("some_value"), + node_execution_id="temp_exec_id", + ) + # Manually set node_execution_id to None after creation + self._node_var_without_exec.node_execution_id = None + + self._node_var_missing_exec = WorkflowDraftVariable.new_node_variable( + app_id=self._test_app_id, + user_id=self._test_user_id, + node_id=self._node_id, + name="missing_exec_var", + value=build_segment("some_value"), + node_execution_id=str(uuid.uuid4()), # Use a valid UUID that doesn't exist in database + ) + + self._conv_var = WorkflowDraftVariable.new_conversation_variable( + app_id=self._test_app_id, + user_id=self._test_user_id, + name="conv_var_1", + value=build_segment("old_conv_value"), + ) + self._conv_var.last_edited_at = datetime_utils.naive_utc_now() + + with Session(self._engine, expire_on_commit=False) as persistent_session, persistent_session.begin(): + persistent_session.add( + self._workflow_node_execution, + ) + + # Add all to database + self._session.add_all( + [ + self._node_var_with_exec, + self._node_var_without_exec, + self._node_var_missing_exec, + self._conv_var, + ] + ) + self._session.flush() + + # Store IDs for assertions + self._node_var_with_exec_id = self._node_var_with_exec.id + self._node_var_without_exec_id = self._node_var_without_exec.id + self._node_var_missing_exec_id = self._node_var_missing_exec.id + self._conv_var_id = self._conv_var.id + + @override + def tearDown(self) -> None: + self._session.rollback() + with Session(self._engine) as session, session.begin(): + stmt = delete(WorkflowNodeExecutionModel).where( + WorkflowNodeExecutionModel.id == self._workflow_node_execution.id + ) + session.execute(stmt) + + def _get_test_srv(self) -> WorkflowDraftVariableService: + return WorkflowDraftVariableService(session=self._session) + + def _create_mock_workflow(self) -> Workflow: + """Create a real workflow with conversation variables and graph""" + conversation_vars = self._conv_variables + + # Create a simple graph with the test node + graph: dict[str, object] = { + "nodes": [{"id": "test_reset_node", "type": "llm", "title": "Test Node", "data": {"type": "llm"}}], + "edges": [], + } + + workflow = Workflow.new( + tenant_id=str(uuid.uuid4()), + app_id=self._test_app_id, + type="workflow", + version="1.0", + graph=json.dumps(graph), + features="{}", + created_by=str(uuid.uuid4()), + environment_variables=[], + conversation_variables=conversation_vars, + rag_pipeline_variables=[], + ) + return workflow + + def test_reset_system_variable_raises_error(self) -> None: + """Test that resetting a system variable raises an error""" + srv = self._get_test_srv() + mock_workflow = self._create_mock_workflow() + + # Create a system variable + sys_var = WorkflowDraftVariable.new_sys_variable( + app_id=self._test_app_id, + user_id=self._test_user_id, + name="sys_var", + value=build_segment("sys_value"), + node_execution_id=self._node_exec_id, + ) + self._session.add(sys_var) + self._session.flush() + + # Attempt to reset the system variable + with pytest.raises(VariableResetError) as exc_info: + srv.reset_variable(mock_workflow, sys_var) + + assert "cannot reset system variable" in str(exc_info.value) + assert sys_var.id in str(exc_info.value) diff --git a/api/tests/unit_tests/services/workflow/test_workflow_event_snapshot_service.py b/api/tests/unit_tests/services/test_workflow_event_snapshot_service.py similarity index 100% rename from api/tests/unit_tests/services/workflow/test_workflow_event_snapshot_service.py rename to api/tests/unit_tests/services/test_workflow_event_snapshot_service.py diff --git a/api/tests/unit_tests/services/test_workflow_generator_service.py b/api/tests/unit_tests/services/test_workflow_generator_service.py index b0dc2620e7e..14ef9a2e6ff 100644 --- a/api/tests/unit_tests/services/test_workflow_generator_service.py +++ b/api/tests/unit_tests/services/test_workflow_generator_service.py @@ -37,7 +37,7 @@ from graphon.model_runtime.entities.message_entities import AssistantPromptMessa from graphon.model_runtime.entities.model_entities import ModelType from graphon.model_runtime.model_providers.base.large_language_model import LargeLanguageModel from services.workflow_generator_service import WorkflowGeneratorService -from tests.unit_tests.core.model_fixtures import make_model_config +from tests.helpers.unit.core.model_fixtures import make_model_config TENANT_ID = "t-1" PROVIDER = "langgenius/openai/openai" diff --git a/api/tests/unit_tests/services/test_workflow_persistence.py b/api/tests/unit_tests/services/test_workflow_persistence.py new file mode 100644 index 00000000000..cf508c3f14e --- /dev/null +++ b/api/tests/unit_tests/services/test_workflow_persistence.py @@ -0,0 +1,1409 @@ +from __future__ import annotations + +from dataclasses import replace +from datetime import UTC, datetime, timedelta +from types import SimpleNamespace +from typing import Any +from unittest.mock import MagicMock + +import pytest +from flask import Flask +from sqlalchemy import event as sqlalchemy_event +from sqlalchemy import select +from sqlalchemy.orm import Session, sessionmaker + +from core.app.entities.app_invoke_entities import AdvancedChatAppGenerateEntity, WorkflowAppGenerateEntity +from core.app.workflow.retry_history import RETRY_HISTORY_PROCESS_DATA_KEY +from core.ops.base_trace_instance import bind_application_tracing_hooks +from core.ops.ops_trace_manager import TraceTask, TraceTaskName +from core.repositories.sqlalchemy_workflow_execution_repository import SQLAlchemyWorkflowExecutionRepository +from core.repositories.sqlalchemy_workflow_node_execution_repository import SQLAlchemyWorkflowNodeExecutionRepository +from core.tools.workflow_as_tool.repository import WorkflowToolSource +from core.workflow.node_execution_process_data import ( + WORKFLOW_TOOL_INVOCATION_ID_KEY, + WORKFLOW_TOOL_ROOT_APP_ID_KEY, +) +from core.workflow.system_variables import SystemVariableKey, build_system_variables +from enums import WorkflowKind +from graphon.engine.event.node_failure import NodeFailureHandler +from graphon.engine.event.processor import NodeEventProcessor +from graphon.engine.event.stream import EventStream +from graphon.engine.frame import ExecutionFrame, FrameRegistry +from graphon.engine.scheduler import Scheduler +from graphon.engine.worker import NodeEventTask +from graphon.engine_events import ( + EngineEvent, + GraphRunAbortedEvent, + GraphRunFailedEvent, + GraphRunPartialSucceededEvent, + GraphRunPausedEvent, + GraphRunStartedEvent, + GraphRunSucceededEvent, + NodeEvent, + NodeRunExceptionEvent, + NodeRunFailedEvent, + NodeRunPauseRequestedEvent, + NodeRunRetryEvent, + NodeRunStartedEvent, + NodeRunSucceededEvent, +) +from graphon.entities import WorkflowNodeExecution, WorkflowStartReason +from graphon.entities.pause_reason import SchedulingPause +from graphon.enums import ( + BuiltinNodeTypes, + WorkflowExecutionStatus, + WorkflowNodeExecutionMetadataKey, + WorkflowNodeExecutionStatus, + WorkflowType, +) +from graphon.graph import Graph +from graphon.model_runtime.entities.llm_entities import LLMUsage +from graphon.node_events import NodeRunResult +from graphon.nodes.start.entities import StartNodeData +from graphon.nodes.start.start_node import StartNode +from graphon.runtime import ReadOnlyRuntimeStateWrapper, RuntimeState, VariablePool +from graphon.runtime.execution import ROOT_FRAME_ID +from models import Account, WorkflowRun +from models.enums import WorkflowRunTriggeredFrom +from models.workflow import WorkflowNodeExecutionModel, WorkflowNodeExecutionTriggeredFrom +from repositories.entities.workflow_execution import WorkflowExecution +from repositories.sqlalchemy_api_workflow_node_execution_repository import ( + DifyAPISQLAlchemyWorkflowNodeExecutionRepository, +) +from services import workflow_persistence as persistence_mod +from services.tracing.runtime_bridge import RuntimeApplicationTracingHooks +from services.workflow_persistence import PersistenceWorkflowInfo, WorkflowPersistenceLayer +from tests.helpers.unit.config_override import apply_config_overrides +from tests.workflow_test_utils import build_test_graph_init_params + +"""Verify the workflow persistence layer fans Inspector deltas to redis pub/sub. + +The hook lives in ``services/workflow_persistence.py``: +every ``_handle_node_*`` and the terminal ``_handle_graph_run_*`` handlers +call into ``services.workflow.inspector_events.publish_node_changed`` / +``publish_workflow_completed`` after the DB write succeeds. Those calls are +the only thing the Inspector SSE stream listens to, so any future refactor of +the persistence layer must keep them in place. + +We don't reconstruct a full workflow engine here — the handlers are tested +in isolation by patching just the moving parts they touch +(``_workflow_execution`` + ``_node_execution_cache``) and asserting against +the publisher module's call sites. This keeps the test compact and tied to +the contract, not the implementation. +""" + + +# ────────────────────────────────────────────────────────────────────────────── +# Graph-level publish hooks +# ────────────────────────────────────────────────────────────────────────────── + + +def _graph_event(**kwargs: Any) -> MagicMock: + return MagicMock(**kwargs) + + +# ────────────────────────────────────────────────────────────────────────────── +# Node-level publish hooks +# ────────────────────────────────────────────────────────────────────────────── + + +def _node_started_event(node_id: str = "agent-1", exec_id: str = "exec-1") -> NodeRunStartedEvent: + return NodeRunStartedEvent( + id=exec_id, + node_id=node_id, + node_type="agent", + node_title="Greeter", + predecessor_node_id=None, + container_id="", + node_version="1", + start_at=datetime(2026, 5, 26, 0, 0, 0), + ) + + +def _seed_node_execution(layer: WorkflowPersistenceLayer, exec_id: str, node_id: str) -> None: + """Inject a domain execution into the cache so the success / fail / etc + handlers (which look it up by id) can run without going through started.""" + layer._node_execution_cache[exec_id] = MagicMock( + id=exec_id, node_id=node_id, status=MagicMock(value="running"), outputs={}, error=None + ) + + +class TestPersistenceInspectorPublishCases: + @staticmethod + @pytest.fixture + def layer() -> WorkflowPersistenceLayer: + """Use real layer initialization with repository and runtime test doubles.""" + instance = WorkflowPersistenceLayer( + application_generate_entity=MagicMock(), + workflow_info=MagicMock(workflow_id="wf-1"), + workflow_execution_repository=MagicMock(), + workflow_node_execution_repository=MagicMock(), + ) + # Use a SimpleNamespace-like spec so Pydantic-validated callsites (e.g. + # ``WorkflowNodeExecution.new`` requires real strings) get the right types. + workflow_execution = MagicMock() + workflow_execution.id_ = "run-1" + workflow_execution.workflow_id = "wf-1" + workflow_execution.status = MagicMock(value="succeeded") + workflow_execution.outputs = {} + workflow_execution.error_message = None + workflow_execution.exceptions_count = 0 + workflow_execution.finished_at = None + instance.set_node_run_indices({"exec-1": 1}) + instance._workflow_execution = workflow_execution + # `runtime_state` is a layer-base property; stub it. + instance._runtime_state = MagicMock(total_tokens=0, node_run_steps=0, outputs={}, exceptions_count=0) + return instance + + @staticmethod + @pytest.fixture + def capture_publishes(monkeypatch: pytest.MonkeyPatch) -> dict[str, list]: + """Replace the two publishers with capture lists so each test can assert + on the exact arguments.""" + calls: dict[str, list] = {"node": [], "workflow": []} + + def fake_node(*, workflow_run_id: str, node_id: str, status: str) -> None: + calls["node"].append({"workflow_run_id": workflow_run_id, "node_id": node_id, "status": status}) + + def fake_workflow(*, workflow_run_id: str, status: str) -> None: + calls["workflow"].append({"workflow_run_id": workflow_run_id, "status": status}) + + monkeypatch.setattr(persistence_mod, "_inspector_publish_node_changed", fake_node) + monkeypatch.setattr(persistence_mod, "_inspector_publish_workflow_completed", fake_workflow) + return calls + + @staticmethod + def test_graph_run_succeeded_publishes_workflow_completed(layer, capture_publishes): + layer._workflow_execution.status = MagicMock(value="succeeded") + layer._handle_graph_run_succeeded(_graph_event(outputs={"text": "hi"})) + assert capture_publishes["workflow"] == [{"workflow_run_id": "run-1", "status": "succeeded"}] + assert capture_publishes["node"] == [] + + @staticmethod + def test_graph_run_partial_succeeded_publishes_workflow_completed(layer, capture_publishes): + layer._workflow_execution.status = MagicMock(value="partial-succeeded") + layer._handle_graph_run_partial_succeeded(_graph_event(outputs={}, exceptions_count=1)) + assert capture_publishes["workflow"] == [{"workflow_run_id": "run-1", "status": "partial-succeeded"}] + + @staticmethod + def test_graph_run_failed_publishes_workflow_completed(layer, capture_publishes): + layer._workflow_execution.status = MagicMock(value="failed") + layer._handle_graph_run_failed(_graph_event(error="boom", exceptions_count=0)) + assert capture_publishes["workflow"] == [{"workflow_run_id": "run-1", "status": "failed"}] + + @staticmethod + def test_graph_run_aborted_publishes_workflow_completed(layer, capture_publishes): + layer._workflow_execution.status = MagicMock(value="stopped") + layer._handle_graph_run_aborted(_graph_event(reason="user stop")) + assert capture_publishes["workflow"] == [{"workflow_run_id": "run-1", "status": "stopped"}] + + @staticmethod + def test_graph_run_paused_does_not_publish_completion(layer, capture_publishes): + """Pause is not a terminal state — the Inspector keeps waiting for either + resume or a real terminal event.""" + layer.on_event(GraphRunPausedEvent(outputs={})) + assert capture_publishes["workflow"] == [] + assert capture_publishes["node"] == [] + + @staticmethod + def test_node_started_publishes_running(layer, capture_publishes): + layer._handle_node_started(_node_started_event()) + assert capture_publishes["node"] == [{"workflow_run_id": "run-1", "node_id": "agent-1", "status": "running"}] + + @staticmethod + def test_node_retry_publishes_retry(layer, capture_publishes, monkeypatch: pytest.MonkeyPatch): + _seed_node_execution(layer, exec_id="exec-1", node_id="agent-1") + monkeypatch.setattr(layer, "_append_retry_history", lambda *args: None) + event = MagicMock(id="exec-1", error="rate limit") + layer._handle_node_retry(event) + assert capture_publishes["node"] == [{"workflow_run_id": "run-1", "node_id": "agent-1", "status": "retry"}] + + @staticmethod + def test_node_succeeded_publishes_succeeded(layer, capture_publishes, monkeypatch: pytest.MonkeyPatch): + _seed_node_execution(layer, exec_id="exec-1", node_id="agent-1") + # Stub the inner _update_node_execution so we don't have to construct a + # full NodeRunResult — we only want to confirm the publish happens after. + monkeypatch.setattr(layer, "_update_node_execution", lambda *a, **kw: None) + event = MagicMock(id="exec-1", node_run_result=MagicMock(), finished_at=datetime.now()) + layer._handle_node_succeeded(event) + assert capture_publishes["node"] == [{"workflow_run_id": "run-1", "node_id": "agent-1", "status": "succeeded"}] + + @staticmethod + def test_node_failed_publishes_failed(layer, capture_publishes, monkeypatch: pytest.MonkeyPatch): + _seed_node_execution(layer, exec_id="exec-1", node_id="agent-1") + monkeypatch.setattr(layer, "_update_node_execution", lambda *a, **kw: None) + event = MagicMock(id="exec-1", node_run_result=MagicMock(), error="bad", finished_at=datetime.now()) + layer._handle_node_failed(event) + assert capture_publishes["node"] == [{"workflow_run_id": "run-1", "node_id": "agent-1", "status": "failed"}] + + @staticmethod + def test_node_exception_publishes_exception(layer, capture_publishes, monkeypatch: pytest.MonkeyPatch): + _seed_node_execution(layer, exec_id="exec-1", node_id="agent-1") + monkeypatch.setattr(layer, "_update_node_execution", lambda *a, **kw: None) + event = MagicMock(id="exec-1", node_run_result=MagicMock(), error="oom", finished_at=datetime.now()) + layer._handle_node_exception(event) + assert capture_publishes["node"] == [{"workflow_run_id": "run-1", "node_id": "agent-1", "status": "exception"}] + + @staticmethod + def test_node_pause_requested_does_not_publish(layer, capture_publishes, monkeypatch: pytest.MonkeyPatch): + """Node pause is not an Inspector-visible state — no publish.""" + _seed_node_execution(layer, exec_id="exec-1", node_id="agent-1") + monkeypatch.setattr(layer, "_update_node_execution", lambda *a, **kw: None) + event = MagicMock(id="exec-1", node_run_result=MagicMock()) + layer._handle_node_pause_requested(event) + assert capture_publishes["node"] == [] + + +class _RepoRecorder: + def __init__(self) -> None: + self.saved: list[object] = [] + self.synchronously_saved: list[object] = [] + self.saved_exec_data: list[object] = [] + + def save(self, entity): + self.saved.append(entity) + + def save_many(self, entities): + self.saved.extend(entities) + + def save_synchronously(self, entity): + self.synchronously_saved.append(entity) + + def save_execution_data(self, entity): + self.saved_exec_data.append(entity) + + +def _naive_utc_now() -> datetime: + return datetime.now(UTC).replace(tzinfo=None) + + +def _make_layer( + system_variables: list | None = None, + *, + extras: dict | None = None, + graph_data: dict | None = None, + trace_manager: object | None = None, + workflow_execution_repo=None, + workflow_node_execution_repo=None, + node_run_indices=None, +): + system_variables = system_variables or build_system_variables( + workflow_execution_id="run-id", + conversation_id="conv-id", + ) + runtime_state = RuntimeState( + workflow_id="test-workflow", + variable_pool=VariablePool.from_bootstrap(system_variables=system_variables), + start_at=0.0, + ) + read_only_state = ReadOnlyRuntimeStateWrapper(runtime_state) + + application_generate_entity = WorkflowAppGenerateEntity.model_construct( + task_id="task", + app_config=SimpleNamespace(app_id="app", tenant_id="tenant"), + inputs={"foo": "bar"}, + files=[], + user_id="user", + stream=False, + invoke_from=None, + trace_manager=None, + workflow_execution_id="run-id", + extras=extras or {}, + call_depth=0, + ) + + workflow_info = PersistenceWorkflowInfo( + workflow_id="workflow-id", + workflow_type=WorkflowType.WORKFLOW, + version="1", + graph_data=graph_data if graph_data is not None else {"nodes": [], "edges": []}, + ) + + workflow_execution_repo = workflow_execution_repo if workflow_execution_repo is not None else _RepoRecorder() + workflow_node_execution_repo = ( + workflow_node_execution_repo if workflow_node_execution_repo is not None else _RepoRecorder() + ) + + layer = WorkflowPersistenceLayer( + application_generate_entity=application_generate_entity, + workflow_info=workflow_info, + workflow_execution_repository=workflow_execution_repo, + workflow_node_execution_repository=workflow_node_execution_repo, + trace_manager=trace_manager, + ) + layer.initialize(read_only_state, command_channel=None) + layer.set_node_run_indices(node_run_indices or {}) + + return layer, workflow_execution_repo, workflow_node_execution_repo, runtime_state + + +def _make_sql_layer(session_factory, *, node_run_indices=None): + user = Account(name="Test", email="test@example.com") + user.id = "user" + layer, _, _, _ = _make_layer( + node_run_indices=node_run_indices, + workflow_execution_repo=SQLAlchemyWorkflowExecutionRepository( + session_factory=session_factory, + tenant_id="tenant", + user=user, + app_id="caller-app", + triggered_from=WorkflowRunTriggeredFrom.APP_RUN, + ), + workflow_node_execution_repo=SQLAlchemyWorkflowNodeExecutionRepository( + session_factory=session_factory, + tenant_id="tenant", + user=user, + app_id="caller-app", + triggered_from=WorkflowNodeExecutionTriggeredFrom.WORKFLOW_RUN, + ), + ) + layer._application_generate_entity.app_config.app_id = "caller-app" + return layer + + +def _tool_source(): + return WorkflowToolSource( + app_id="source-app", + workflow_id="source-workflow", + graph_config={"nodes": [{"id": "source-loop", "data": {"type": "loop"}}], "edges": []}, + features_dict={}, + environment_variables=[], + workflow_kind=WorkflowKind.STANDARD, + ) + + +class TestPersistenceLayerCases: + @staticmethod + @pytest.fixture(autouse=True) + def _application_tracing_disabled(monkeypatch: pytest.MonkeyPatch) -> None: + apply_config_overrides(monkeypatch, OPS_AGENT_TRACE_ENABLED=False) + + @staticmethod + def test_tool_trace_identity_survives_start_pause_and_resume(): + layer, _, node_repo, _ = _make_layer(node_run_indices={"tool-exec": 1}) + layer.on_event(GraphRunStartedEvent()) + started_at = _naive_utc_now() + layer.on_event( + NodeRunStartedEvent( + id="tool-exec", + node_id="tool", + node_type=BuiltinNodeTypes.TOOL, + node_title="Nested tool", + provider_type="workflow", + provider_id="provider", + start_at=started_at, + node_run_result=NodeRunResult(process_data={WORKFLOW_TOOL_INVOCATION_ID_KEY: "invocation"}), + ) + ) + for event in ( + NodeRunPauseRequestedEvent( + id="tool-exec", node_id="tool", node_type=BuiltinNodeTypes.TOOL, reason=SchedulingPause(message="pause") + ), + NodeRunSucceededEvent(id="tool-exec", node_id="tool", node_type=BuiltinNodeTypes.TOOL, start_at=started_at), + ): + layer.on_event(event) + execution = node_repo.saved_exec_data[-1] + assert execution.process_data[WORKFLOW_TOOL_INVOCATION_ID_KEY] == "invocation" + assert execution.metadata[WorkflowNodeExecutionMetadataKey.TOOL_INFO] == { + "provider_type": "workflow", + "provider_id": "provider", + } + + @staticmethod + def test_workflow_tool_agent_nodes_persist_in_source_app_under_only_root_run(sqlite_session_factory): + layer = _make_sql_layer(sqlite_session_factory, node_run_indices={"caller-exec": 1, "source-agent-exec": 1}) + listen = layer.create_workflow_tool_event_listener(_tool_source(), "workflow-id", "caller-exec") + layer.on_graph_start() + layer.on_event(GraphRunStartedEvent()) + started_at = _naive_utc_now() + layer.on_event( + NodeRunStartedEvent( + id="caller-exec", + node_id="caller", + node_type=BuiltinNodeTypes.TOOL, + node_title="Workflow Tool", + start_at=started_at, + ) + ) + listen( + NodeRunStartedEvent( + id="source-agent-exec", + node_id="agent", + node_type=BuiltinNodeTypes.AGENT, + node_version="2", + node_title="Agent", + start_at=started_at, + container_id="source-loop", + node_run_result=NodeRunResult( + process_data={WORKFLOW_TOOL_ROOT_APP_ID_KEY: "untrusted-app"}, metadata={"loop_id": "source-loop"} + ), + ) + ) + + with sqlite_session_factory() as session: + rows = {row.id: row for row in session.scalars(select(WorkflowNodeExecutionModel))} + assert rows["caller-exec"].app_id == "caller-app" + source_row = rows["source-agent-exec"] + assert (source_row.tenant_id, source_row.app_id, source_row.workflow_id, source_row.workflow_run_id) == ( + "tenant", + "source-app", + "source-workflow", + "run-id", + ) + assert source_row.status == "running" + assert source_row.process_data_dict[WORKFLOW_TOOL_ROOT_APP_ID_KEY] == "caller-app" + assert source_row.execution_metadata_dict["loop_id"] == "source-loop" + assert session.scalars(select(WorkflowRun)).one().workflow_id == "workflow-id" + + listen( + NodeRunSucceededEvent( + id="source-agent-exec", + node_id="agent", + node_type=BuiltinNodeTypes.AGENT, + node_title="Agent", + start_at=started_at, + node_run_result=NodeRunResult(status=WorkflowNodeExecutionStatus.SUCCEEDED, outputs={"answer": "ok"}), + ) + ) + layer.on_event(GraphRunSucceededEvent(outputs={"result": "caller-result"})) + + with sqlite_session_factory() as session: + run = session.scalars(select(WorkflowRun)).one() + assert (run.app_id, run.workflow_id, run.status, run.outputs_dict) == ( + "caller-app", + "workflow-id", + "succeeded", + {"result": "caller-result"}, + ) + row = session.get(WorkflowNodeExecutionModel, "source-agent-exec") + assert row.status == "succeeded" + assert row.outputs_dict == {"answer": "ok"} + assert row.process_data_dict[WORKFLOW_TOOL_ROOT_APP_ID_KEY] == "caller-app" + + @staticmethod + def test_nested_tool_callers_keep_distinct_invocations_in_shared_source_layers( + sqlite_session_factory: sessionmaker[Session], + ) -> None: + ids = ["a-call-1", "a-call-2", "b-call-1", "b-call-2", "c-child"] + layer = _make_sql_layer(sqlite_session_factory, node_run_indices=dict.fromkeys(ids, 1)) + source = _tool_source() + nested = replace(source, app_id="nested-app", workflow_id="nested-workflow") + first_call = layer.create_workflow_tool_event_listener(source, "workflow-id", "a-call-1") + second_call = layer.create_workflow_tool_event_listener(source, "workflow-id", "a-call-2") + nested_call = layer.create_workflow_tool_event_listener(nested, source.workflow_id, "b-call-1") + layer.on_event(GraphRunStartedEvent()) + for execution_id, listener in zip( + ids, [layer.on_event, layer.on_event, first_call, second_call, nested_call], strict=True + ): + listener( + NodeRunStartedEvent( + id=execution_id, + node_id="same-canvas-node", + node_type=BuiltinNodeTypes.TOOL, + node_title="Tool", + start_at=_naive_utc_now(), + ) + ) + repository = DifyAPISQLAlchemyWorkflowNodeExecutionRepository(sqlite_session_factory) + for caller, expected, source_app, source_workflow in ( + ("a-call-1", "b-call-1", "source-app", "source-workflow"), + ("a-call-2", "b-call-2", "source-app", "source-workflow"), + ("b-call-1", "c-child", "nested-app", "nested-workflow"), + ): + children = repository.get_workflow_tool_executions("tenant", "run-id", caller) + assert [child.id for child in children] == [expected] + child = children[0] + assert (child.app_id, child.workflow_id, child.workflow_run_id) == (source_app, source_workflow, "run-id") + assert child.triggered_from_workflow_id == ("source-workflow" if caller == "b-call-1" else "workflow-id") + assert child.triggered_from_node_execution_id == caller + + @staticmethod + def test_workflow_tool_resume_updates_existing_running_container(sqlite_session_factory): + layer = _make_sql_layer(sqlite_session_factory, node_run_indices={"loop-exec": 1}) + listen = layer.create_workflow_tool_event_listener(_tool_source(), "workflow-id", "caller-exec") + layer.on_event(GraphRunStartedEvent()) + started_at = _naive_utc_now() + listen( + NodeRunStartedEvent( + id="loop-exec", + node_id="source-loop", + node_type=BuiltinNodeTypes.LOOP, + node_title="Approval Loop", + start_at=started_at, + ) + ) + layer.on_event(GraphRunPausedEvent(outputs={}, reasons=[])) + + resumed = _make_sql_layer(sqlite_session_factory, node_run_indices={"loop-exec": 1}) + resume_listener = resumed.create_workflow_tool_event_listener( + _tool_source(), + "workflow-id", + "caller-exec", + node_executions=resumed._workflow_node_execution_repository.for_workflow_tool( + "source-app" + ).get_by_workflow_execution("run-id"), + ) + resumed.on_graph_start() + resumed.on_event(GraphRunStartedEvent(reason=WorkflowStartReason.RESUMPTION)) + resume_listener( + NodeRunSucceededEvent( + id="loop-exec", + node_id="source-loop", + node_type=BuiltinNodeTypes.LOOP, + node_title="Approval Loop", + start_at=started_at, + node_run_result=NodeRunResult(status=WorkflowNodeExecutionStatus.SUCCEEDED, outputs={"approved": True}), + ) + ) + resumed.on_event(GraphRunSucceededEvent(outputs={"answer": "approved"})) + + with sqlite_session_factory() as session: + node = session.scalars(select(WorkflowNodeExecutionModel)).one() + assert (node.app_id, node.workflow_id, node.workflow_run_id, node.status) == ( + "source-app", + "source-workflow", + "run-id", + "succeeded", + ) + assert node.outputs_dict == {"approved": True} + run = session.scalars(select(WorkflowRun)).one() + assert (run.app_id, run.workflow_id, run.status) == ("caller-app", "workflow-id", "succeeded") + + @staticmethod + @pytest.mark.parametrize("source_workflow_id", ["older-workflow", "workflow-id"]) + def test_root_resume_isolates_same_app_workflow_tool_origin(sqlite_session_factory, source_workflow_id): + source = replace(_tool_source(), app_id="caller-app", workflow_id=source_workflow_id) + layer = _make_sql_layer(sqlite_session_factory, node_run_indices={"caller-exec": 1, "older-exec": 1}) + listen = layer.create_workflow_tool_event_listener(source, "workflow-id", "caller-exec") + layer.on_event(GraphRunStartedEvent()) + started_at = _naive_utc_now() + layer.on_event( + NodeRunStartedEvent( + id="caller-exec", + node_id="caller", + node_type=BuiltinNodeTypes.TOOL, + node_title="Caller", + start_at=started_at, + ) + ) + listen( + NodeRunStartedEvent( + id="older-exec", + node_id="older-node", + node_type=BuiltinNodeTypes.LLM, + node_title="Older version", + start_at=started_at, + ) + ) + layer.on_event(GraphRunPausedEvent()) + + resumed = _make_sql_layer(sqlite_session_factory, node_run_indices={"caller-exec": 1, "older-exec": 1}) + resumed.set_node_execution_history( + resumed._workflow_node_execution_repository.get_by_workflow_execution("run-id") + ) + resume_listener = resumed.create_workflow_tool_event_listener( + source, + "workflow-id", + "caller-exec", + node_executions=resumed._workflow_node_execution_repository.for_workflow_tool( + source.app_id + ).get_by_workflow_execution("run-id"), + ) + resumed.on_event(GraphRunStartedEvent(reason=WorkflowStartReason.RESUMPTION)) + resume_listener( + NodeRunSucceededEvent( + id="older-exec", + node_id="older-node", + node_type=BuiltinNodeTypes.LLM, + start_at=started_at, + node_run_result=NodeRunResult(status=WorkflowNodeExecutionStatus.SUCCEEDED), + ) + ) + resumed.on_event(GraphRunFailedEvent(error="caller failed later")) + + with sqlite_session_factory() as session: + rows = {row.id: row for row in session.scalars(select(WorkflowNodeExecutionModel))} + assert rows["caller-exec"].status == "failed" + assert rows["older-exec"].status == "succeeded" + assert rows["older-exec"].workflow_id == source_workflow_id + + @staticmethod + @pytest.mark.parametrize("source", [False, True]) + @pytest.mark.parametrize("outcome", ["succeeded", "failed", "aborted", "paused"]) + def test_retry_starts_preserve_original_execution_and_attempt_history( + sqlite_session_factory: sessionmaker[Session], + monkeypatch: pytest.MonkeyPatch, + source: bool, + outcome: str, + ) -> None: + layer = _make_sql_layer(sqlite_session_factory, node_run_indices={"retry-exec": 1}) + listen = ( + layer.create_workflow_tool_event_listener(_tool_source(), "workflow-id", "caller-exec") + if source + else layer.on_event + ) + layer.on_event(GraphRunStartedEvent()) + workflow_id = "source-workflow" if source else "workflow-id" + state = RuntimeState(workflow_id=workflow_id, variable_pool=VariablePool(), start_at=0) + frames = FrameRegistry() + node = StartNode( + node_id="child", + data=StartNodeData(title="Child", variables=[]), + init_params=build_test_graph_init_params(workflow_id=workflow_id), + runtime_state=state, + ) + graph = Graph.new().add_root(node).build() + frames.register( + ExecutionFrame( + frame_id=ROOT_FRAME_ID, + graph=graph, + state=state, + scheduler=Scheduler(graph, state, ROOT_FRAME_ID), + failure_handler=NodeFailureHandler(graph, state.graph_execution), + ) + ) + + event_stream = EventStream(layers=[]) + collect = event_stream.collect + + def collect_source_event(event: EngineEvent) -> None: + assert isinstance(event, NodeEvent) + listen(event) + collect(event) + + monkeypatch.setattr(event_stream, "collect", collect_source_event) + processor = NodeEventProcessor(state.graph_execution, event_stream, frames, {}) + + def dispatch(event: NodeEvent) -> None: + processor.dispatch(NodeEventTask(frame_id=ROOT_FRAME_ID, event=event)) + + started_at = _naive_utc_now() + started = NodeRunStartedEvent( + id="retry-exec", + node_id="child", + node_type=BuiltinNodeTypes.LLM, + node_title="Child", + start_at=started_at, + ) + dispatch(started) + for attempt in (1, 2): + dispatch( + NodeRunRetryEvent( + id="retry-exec", + node_id="child", + node_type=BuiltinNodeTypes.LLM, + node_title="Child", + start_at=started_at, + retry_index=attempt, + error=f"attempt-{attempt}", + node_run_result=NodeRunResult(outputs={"attempt": attempt}), + ) + ) + dispatch(started) + if outcome == "succeeded": + dispatch( + NodeRunSucceededEvent( + id="retry-exec", + node_id="child", + node_type=BuiltinNodeTypes.LLM, + start_at=started_at, + node_run_result=NodeRunResult(status=WorkflowNodeExecutionStatus.SUCCEEDED, outputs={"ok": True}), + ) + ) + elif outcome == "failed": + layer.on_event(GraphRunFailedEvent(error="terminal failure")) + elif outcome == "aborted": + layer.on_event(GraphRunAbortedEvent(reason="terminal failure")) + else: + commits = [] + + def record_commit(_session: Session) -> None: + commits.append(True) + + sqlalchemy_event.listen(Session, "after_commit", record_commit) + try: + layer.persist_unfinished_nodes_synchronously() + finally: + sqlalchemy_event.remove(Session, "after_commit", record_commit) + assert len(commits) == 1 + + with sqlite_session_factory() as session: + row = session.scalars(select(WorkflowNodeExecutionModel)).one() + assert row.index == 1 + assert row.created_at == started_at + assert row.workflow_id == workflow_id + assert row.app_id == ("source-app" if source else "caller-app") + if outcome == "succeeded": + assert row.status == "succeeded" + assert row.outputs_dict == {"ok": True} + elif outcome == "paused": + assert row.status == "retry" + assert row.finished_at is None + else: + assert row.status == "failed" + assert row.error == "terminal failure" + assert row.finished_at is not None + history = row.process_data_dict[RETRY_HISTORY_PROCESS_DATA_KEY] + assert [(attempt["retry_index"], attempt["error"]) for attempt in history] == [ + (1, "attempt-1"), + (2, "attempt-2"), + ] + + @staticmethod + @pytest.mark.parametrize( + "event", [GraphRunFailedEvent(error="failed", exceptions_count=1), GraphRunAbortedEvent(reason="aborted")] + ) + def test_root_terminal_failure_finishes_running_workflow_tool_nodes(sqlite_session_factory, event): + layer = _make_sql_layer(sqlite_session_factory, node_run_indices={"source-exec": 1}) + listen = layer.create_workflow_tool_event_listener(_tool_source(), "workflow-id", "caller-exec") + layer.on_event(GraphRunStartedEvent()) + listen( + NodeRunStartedEvent( + id="source-exec", + node_id="child", + node_type=BuiltinNodeTypes.LLM, + node_title="Child", + start_at=_naive_utc_now(), + ) + ) + + layer.on_event(event) + + with sqlite_session_factory() as session: + node = session.scalars(select(WorkflowNodeExecutionModel)).one() + assert node.status == "failed" + assert node.error == (event.error if isinstance(event, GraphRunFailedEvent) else event.reason) + assert node.finished_at is not None + assert session.scalars(select(WorkflowRun)).one().workflow_id == "workflow-id" + + @staticmethod + def test_terminal_failure_batches_only_unfinished_nodes_per_source( + sqlite_session_factory: sessionmaker[Session], + ) -> None: + ids = ( + "root-a", + "root-b", + "root-retry", + "root-paused", + "source-a", + "source-b", + "source-retry", + "source-paused", + "source-done", + ) + layer = _make_sql_layer(sqlite_session_factory, node_run_indices=dict.fromkeys(ids, 1)) + listen = layer.create_workflow_tool_event_listener(_tool_source(), "workflow-id", "root-a") + layer.on_event(GraphRunStartedEvent()) + started_at = _naive_utc_now() + for execution_id in ids: + listener = listen if execution_id.startswith("source") else layer.on_event + listener( + NodeRunStartedEvent( + id=execution_id, + node_id=execution_id, + node_type=BuiltinNodeTypes.CODE, + node_title="Code", + start_at=started_at, + ) + ) + if execution_id.endswith("paused"): + listener( + NodeRunPauseRequestedEvent( + id=execution_id, + node_id=execution_id, + node_type=BuiltinNodeTypes.CODE, + reason=SchedulingPause(message="pause"), + ) + ) + elif execution_id.endswith("retry"): + listener( + NodeRunRetryEvent( + id=execution_id, + node_id=execution_id, + node_type=BuiltinNodeTypes.CODE, + node_title="Code", + start_at=started_at, + retry_index=1, + error="retry failed", + ) + ) + listen( + NodeRunSucceededEvent( + id="source-done", + node_id="source-done", + node_type=BuiltinNodeTypes.CODE, + start_at=started_at, + node_run_result=NodeRunResult(status=WorkflowNodeExecutionStatus.SUCCEEDED, outputs={"answer": "kept"}), + ) + ) + commits = [] + + def record_commit(_session: Session) -> None: + commits.append(True) + + sqlalchemy_event.listen(Session, "after_commit", record_commit) + try: + layer.on_event(GraphRunFailedEvent(error="terminal failure", exceptions_count=0)) + finally: + sqlalchemy_event.remove(Session, "after_commit", record_commit) + + with sqlite_session_factory() as session: + rows = {row.id: row for row in session.scalars(select(WorkflowNodeExecutionModel))} + assert {key: row.status for key, row in rows.items()} == { + "root-a": "failed", + "root-b": "failed", + "root-retry": "failed", + "root-paused": "paused", + "source-a": "failed", + "source-b": "failed", + "source-retry": "failed", + "source-paused": "paused", + "source-done": "succeeded", + } + assert rows["source-done"].outputs_dict == {"answer": "kept"} + for execution_id in ("root-a", "root-b", "root-retry", "source-a", "source-b", "source-retry"): + assert rows[execution_id].error == "terminal failure" + assert rows[execution_id].finished_at is not None + for execution_id in ("root-retry", "source-retry"): + assert ( + rows[execution_id].process_data_dict[RETRY_HISTORY_PROCESS_DATA_KEY][0]["error"] == "retry failed" + ) + assert rows["source-a"].triggered_from_node_execution_id == "root-a" + assert rows["source-a"].app_id == "source-app" + # One commit per source owner, plus the workflow summary. + assert len(commits) == 3 + + class TestWorkflowPersistenceLayer: + def test_start_freezes_tracing_only_after_workflow_save(self, monkeypatch: pytest.MonkeyPatch) -> None: + apply_config_overrides(monkeypatch, OPS_AGENT_TRACE_ENABLED=True) + layer, workflow_repo, _, _ = _make_layer(extras={"trace_session_id": "external-session"}) + calls: list[tuple[str, str, str, str, object]] = [] + + def freeze( + entity: AdvancedChatAppGenerateEntity | WorkflowAppGenerateEntity, + *, + workflow_id: str, + workflow_run_id: str, + conversation_id: object, + graph_data: object, + ) -> None: + assert workflow_repo.synchronously_saved + assert graph_data == {"nodes": [], "edges": []} + saved_execution = workflow_repo.synchronously_saved[-1] + assert isinstance(saved_execution, WorkflowExecution) + assert saved_execution.id_ == workflow_run_id + calls.append( + ( + entity.app_config.app_id, + entity.app_config.tenant_id, + workflow_id, + workflow_run_id, + conversation_id, + ) + ) + + monkeypatch.setattr(RuntimeApplicationTracingHooks, "freeze_workflow_run", staticmethod(freeze)) + hooks = RuntimeApplicationTracingHooks() + app = Flask(__name__) + bind_application_tracing_hooks(app, hooks) + + with app.app_context(): + layer._handle_graph_run_started() + assert calls == [("app", "tenant", "workflow-id", "run-id", "conv-id")] + + def test_get_execution_id_requires_system_variable(self): + layer, _, _, _ = _make_layer(build_system_variables()) + + with pytest.raises(ValueError, match="workflow_execution_id must be provided"): + layer._get_execution_id() + + def test_prepare_workflow_inputs_excludes_conversation_id(self, monkeypatch: pytest.MonkeyPatch): + layer, _, _, _ = _make_layer() + + monkeypatch.setattr( + "core.workflow.workflow_entry.WorkflowEntry.handle_special_values", + lambda inputs: inputs, + ) + + inputs = layer._prepare_workflow_inputs() + + assert "sys.conversation_id" not in inputs + assert inputs[f"sys.{SystemVariableKey.WORKFLOW_EXECUTION_ID.value}"] == "run-id" + + @pytest.mark.parametrize( + "status", + [ + WorkflowNodeExecutionStatus.RUNNING, + WorkflowNodeExecutionStatus.RETRY, + WorkflowNodeExecutionStatus.PAUSED, + WorkflowNodeExecutionStatus.SUCCEEDED, + ], + ) + def test_fail_running_node_executions_marks_only_live_attempts_failed(self, status): + layer, _, node_repo, _ = _make_layer() + + execution = WorkflowNodeExecution( + id="exec-id", + workflow_id="workflow-id", + workflow_execution_id="run-id", + index=1, + node_id="node", + node_type=BuiltinNodeTypes.START, + title="Start", + status=status, + created_at=_naive_utc_now(), + ) + layer._node_execution_cache[execution.id] = execution + + layer._fail_running_node_executions(error_message="boom") + + if status in (WorkflowNodeExecutionStatus.RUNNING, WorkflowNodeExecutionStatus.RETRY): + assert execution.status == WorkflowNodeExecutionStatus.FAILED + assert execution.error == "boom" + assert execution.finished_at is not None + assert node_repo.saved == [execution] + else: + assert execution.status == status + assert execution.finished_at is None + assert node_repo.saved == [] + + def test_resumption_restores_container_execution_before_terminal_event(self): + layer, _, node_repo, _ = _make_layer() + started_at = _naive_utc_now() + execution = WorkflowNodeExecution( + id="loop-exec", + workflow_id="workflow-id", + workflow_execution_id="run-id", + index=4, + node_id="loop", + node_type=BuiltinNodeTypes.LOOP, + title="Loop", + status=WorkflowNodeExecutionStatus.RUNNING, + created_at=started_at, + ) + layer.set_node_execution_history([execution]) + + layer.on_event(GraphRunStartedEvent(reason=WorkflowStartReason.RESUMPTION)) + layer.on_event( + NodeRunSucceededEvent( + id=execution.id, + node_id=execution.node_id, + node_type=execution.node_type, + start_at=started_at, + finished_at=started_at + timedelta(seconds=2), + node_run_result=NodeRunResult(status=WorkflowNodeExecutionStatus.SUCCEEDED), + ) + ) + + assert execution.status == WorkflowNodeExecutionStatus.SUCCEEDED + assert execution.elapsed_time == 2.0 + assert execution.index == 4 + + def test_handle_graph_run_succeeded_updates_execution(self): + layer, exec_repo, _, runtime_state = _make_layer() + layer.on_event(GraphRunStartedEvent()) + usage = LLMUsage.empty_usage() + usage.total_tokens = 3 + runtime_state.add_llm_usage(usage) + for _ in range(2): + runtime_state.increment_node_run_steps() + runtime_state.set_output("out", "v") + + layer.on_event(GraphRunSucceededEvent(outputs={"ok": True})) + + saved = exec_repo.synchronously_saved[-1] + assert saved.status == WorkflowExecutionStatus.SUCCEEDED + assert saved.total_tokens == 3 + assert saved.total_steps == 2 + + def test_handle_graph_run_partial_succeeded_updates_execution(self): + layer, exec_repo, _, runtime_state = _make_layer() + layer.on_event(GraphRunStartedEvent()) + usage = LLMUsage.empty_usage() + usage.total_tokens = 5 + runtime_state.add_llm_usage(usage) + for _ in range(4): + runtime_state.increment_node_run_steps() + runtime_state._graph_execution = SimpleNamespace(exceptions_count=2) + + layer.on_event(GraphRunPartialSucceededEvent(outputs={"ok": True}, exceptions_count=2)) + + saved = exec_repo.synchronously_saved[-1] + assert saved.status == WorkflowExecutionStatus.PARTIAL_SUCCEEDED + assert saved.exceptions_count == 2 + assert saved.total_tokens == 5 + + def test_handle_graph_run_failed_marks_nodes_and_enqueues_trace(self): + trace_tasks: list[object] = [] + trace_manager = SimpleNamespace(user_id="user", add_trace_task=lambda task: trace_tasks.append(task)) + layer, exec_repo, node_repo, _ = _make_layer( + extras={"external_trace_id": "trace"}, trace_manager=trace_manager + ) + layer.on_event(GraphRunStartedEvent()) + + running = WorkflowNodeExecution( + id="node-exec", + workflow_id="workflow-id", + workflow_execution_id="run-id", + index=1, + node_id="node", + node_type=BuiltinNodeTypes.START, + title="Start", + created_at=_naive_utc_now(), + ) + layer._node_execution_cache[running.id] = running + + layer.on_event(GraphRunFailedEvent(error="boom", exceptions_count=1)) + + assert node_repo.saved + assert exec_repo.synchronously_saved[-1].status == WorkflowExecutionStatus.FAILED + assert trace_tasks + + def test_handle_graph_run_succeeded_enqueues_parent_trace_context(self, monkeypatch: pytest.MonkeyPatch): + trace_tasks: list[TraceTask] = [] + trace_manager = SimpleNamespace(user_id="user", add_trace_task=lambda task: trace_tasks.append(task)) + layer, _, _, _ = _make_layer( + extras={ + "external_trace_id": "trace", + "trace_session_id": "session-1", + "parent_trace_context": { + "parent_workflow_run_id": "outer-workflow-run-1", + "parent_node_execution_id": "outer-node-execution-1", + }, + }, + trace_manager=trace_manager, + ) + layer.on_event(GraphRunStartedEvent()) + + captured: dict[str, object] = {} + + def fake_workflow_trace( + self: TraceTask, + *, + workflow_run_id: str | None, + conversation_id: str | None, + user_id: str | None, + total_tokens_override: int | None = None, + ): + captured["trace_type"] = self.trace_type + captured["external_trace_id"] = self.kwargs.get("external_trace_id") + captured["trace_session_id"] = self.kwargs.get("trace_session_id") + captured["parent_trace_context"] = self.kwargs.get("parent_trace_context") + captured["workflow_run_id"] = workflow_run_id + return {"ok": True} + + monkeypatch.setattr(TraceTask, "workflow_trace", fake_workflow_trace) + + layer.on_event(GraphRunSucceededEvent(outputs={"ok": True})) + + assert trace_tasks + trace_task = trace_tasks[0] + assert trace_task.trace_type == TraceTaskName.WORKFLOW_TRACE + assert trace_task.kwargs["external_trace_id"] == "trace" + assert trace_task.kwargs["trace_session_id"] == "session-1" + assert trace_task.kwargs["parent_trace_context"] == { + "parent_workflow_run_id": "outer-workflow-run-1", + "parent_node_execution_id": "outer-node-execution-1", + } + + trace_task.execute() + + assert captured["trace_type"] == TraceTaskName.WORKFLOW_TRACE + assert captured["external_trace_id"] == "trace" + assert captured["trace_session_id"] == "session-1" + assert captured["parent_trace_context"] == { + "parent_workflow_run_id": "outer-workflow-run-1", + "parent_node_execution_id": "outer-node-execution-1", + } + + def test_handle_graph_run_aborted_sets_status(self): + layer, exec_repo, _, _ = _make_layer() + layer.on_event(GraphRunStartedEvent()) + + layer.on_event(GraphRunAbortedEvent(reason=None, outputs={})) + + saved = exec_repo.synchronously_saved[-1] + assert saved.status == WorkflowExecutionStatus.STOPPED + assert saved.error_message + + def test_handle_node_started_and_retry(self): + layer, _, node_repo, _ = _make_layer(node_run_indices={"exec": 1}) + layer.on_event(GraphRunStartedEvent()) + + start_event = NodeRunStartedEvent( + id="exec", + node_id="node", + node_type=BuiltinNodeTypes.START, + node_title="Start", + start_at=_naive_utc_now(), + predecessor_node_id="prev", + ) + layer.on_event(start_event) + + assert node_repo.saved + assert "exec" in layer._node_execution_cache + assert node_repo.saved[-1].created_at == start_event.start_at + + retry_event = NodeRunRetryEvent( + id="exec", + node_id="node", + node_type=BuiltinNodeTypes.START, + node_title="Start", + start_at=_naive_utc_now(), + error="retry", + retry_index=1, + ) + layer.on_event(retry_event) + assert node_repo.saved_exec_data + + def test_agent_v2_caller_row_is_saved_synchronously_before_node_run(self): + layer, _, node_repo, _ = _make_layer(node_run_indices={"agent-exec": 1}) + layer.on_event(GraphRunStartedEvent()) + + layer.on_event( + NodeRunStartedEvent( + id="agent-exec", + node_id="agent-node", + node_type=BuiltinNodeTypes.AGENT, + node_version="2", + node_title="Agent", + start_at=_naive_utc_now(), + ) + ) + + assert [execution.id for execution in node_repo.synchronously_saved] == ["agent-exec"] + assert node_repo.saved == [] + + def test_retry_history_is_preserved_after_node_succeeds(self): + layer, _, node_repo, _ = _make_layer(node_run_indices={"exec": 1}) + layer.on_event(GraphRunStartedEvent()) + started_at = _naive_utc_now() + layer.on_event( + NodeRunStartedEvent( + id="exec", + node_id="node", + node_type=BuiltinNodeTypes.HTTP_REQUEST, + node_title="HTTP", + start_at=started_at, + ) + ) + + for retry_index in (1, 2): + layer.on_event( + NodeRunRetryEvent( + id="exec", + node_id="node", + node_type=BuiltinNodeTypes.HTTP_REQUEST, + node_title="HTTP", + start_at=started_at, + error=f"attempt {retry_index} failed", + retry_index=retry_index, + node_run_result=NodeRunResult( + inputs={"attempt": retry_index}, + process_data={"request": f"attempt-{retry_index}"}, + outputs={"status_code": 500, "body": f"failure-{retry_index}"}, + metadata={WorkflowNodeExecutionMetadataKey.ITERATION_ID: "iteration-1"}, + ), + ) + ) + + layer.on_event( + NodeRunSucceededEvent( + id="exec", + node_id="node", + node_type=BuiltinNodeTypes.HTTP_REQUEST, + start_at=started_at, + node_run_result=NodeRunResult( + inputs={"attempt": 3}, + process_data={"request": "successful-attempt"}, + outputs={"status_code": 200, "body": "ok"}, + metadata={}, + ), + ) + ) + + saved_execution = node_repo.saved_exec_data[-1] + assert saved_execution.status == WorkflowNodeExecutionStatus.SUCCEEDED + assert saved_execution.process_data["request"] == "successful-attempt" + retry_history = saved_execution.process_data["__dify_retry_history"] + assert [attempt["retry_index"] for attempt in retry_history] == [1, 2] + assert retry_history[0]["inputs"] == {"attempt": 1} + assert retry_history[0]["process_data"] == {"request": "attempt-1"} + assert retry_history[0]["outputs"] == {"status_code": 500, "body": "failure-1"} + assert retry_history[0]["error"] == "attempt 1 failed" + assert retry_history[0]["execution_metadata"] == {"iteration_id": "iteration-1"} + + @pytest.mark.parametrize( + ("event_type", "expected_status"), + [ + (NodeRunFailedEvent, WorkflowNodeExecutionStatus.FAILED), + (NodeRunExceptionEvent, WorkflowNodeExecutionStatus.EXCEPTION), + (GraphRunFailedEvent, WorkflowNodeExecutionStatus.FAILED), + (GraphRunAbortedEvent, WorkflowNodeExecutionStatus.FAILED), + ], + ) + def test_retry_history_is_preserved_after_terminal_error(self, event_type, expected_status): + layer, _, node_repo, _ = _make_layer(node_run_indices={"exec": 1}) + layer.on_event(GraphRunStartedEvent()) + started_at = _naive_utc_now() + layer.on_event( + NodeRunStartedEvent( + id="exec", + node_id="node", + node_type=BuiltinNodeTypes.LLM, + node_title="LLM", + start_at=started_at, + ) + ) + layer.on_event( + NodeRunRetryEvent( + id="exec", + node_id="node", + node_type=BuiltinNodeTypes.LLM, + node_title="LLM", + start_at=started_at, + error="retry failed", + retry_index=1, + node_run_result=NodeRunResult(outputs={"attempt": 1}), + ) + ) + + if event_type is GraphRunFailedEvent: + terminal_event = event_type(error="terminal failure") + elif event_type is GraphRunAbortedEvent: + terminal_event = event_type(reason="terminal failure") + else: + terminal_event = event_type( + id="exec", + node_id="node", + node_type=BuiltinNodeTypes.LLM, + start_at=started_at, + error="terminal failure", + node_run_result=NodeRunResult(process_data={"terminal": True}), + ) + node_repo.saved.clear() + layer.on_event(terminal_event) + + saved_execution = node_repo.saved[-1] + assert saved_execution.status == expected_status + assert saved_execution.error == "terminal failure" + assert saved_execution.finished_at is not None + if event_type in (NodeRunFailedEvent, NodeRunExceptionEvent): + assert saved_execution.process_data["terminal"] is True + assert saved_execution.process_data["__dify_retry_history"][0]["retry_index"] == 1 + + def test_pause_flush_persists_retry_history_synchronously(self): + layer, _, node_repo, _ = _make_layer(node_run_indices={"exec": 1}) + layer.on_event(GraphRunStartedEvent()) + started_at = _naive_utc_now() + layer.on_event( + NodeRunStartedEvent( + id="exec", node_id="node", node_type=BuiltinNodeTypes.LLM, node_title="LLM", start_at=started_at + ) + ) + layer.on_event( + NodeRunRetryEvent( + id="exec", + node_id="node", + node_type=BuiltinNodeTypes.LLM, + node_title="LLM", + start_at=started_at, + error="retry failed", + retry_index=1, + ) + ) + assert node_repo.synchronously_saved == [] + + layer.persist_unfinished_nodes_synchronously() + + saved_execution = node_repo.synchronously_saved[-1] + assert saved_execution.status == WorkflowNodeExecutionStatus.RETRY + assert saved_execution.finished_at is None + assert saved_execution.process_data["__dify_retry_history"][0]["error"] == "retry failed" + + def test_handle_node_pause_requested_skips_outputs(self): + layer, _, _, _ = _make_layer(node_run_indices={"exec": 1}) + layer.on_event(GraphRunStartedEvent()) + start_event = NodeRunStartedEvent( + id="exec", + node_id="node", + node_type=BuiltinNodeTypes.LLM, + node_title="LLM", + start_at=_naive_utc_now(), + ) + layer.on_event(start_event) + + domain_execution = layer._node_execution_cache["exec"] + domain_execution.inputs = {"old": True} + + result = NodeRunResult( + inputs={"new": True}, + outputs={"out": 1}, + process_data={ + "p": 1, + "workflow_agent_binding_id": "workflow-binding-1", + "workflow_tool_invocation_id": "tool-call-1", + }, + metadata={}, + ) + pause_event = NodeRunPauseRequestedEvent( + id="exec", + node_id="node", + node_type=BuiltinNodeTypes.LLM, + reason=SchedulingPause(message="pause"), + node_run_result=result, + ) + layer.on_event(pause_event) + + assert domain_execution.status == WorkflowNodeExecutionStatus.PAUSED + assert domain_execution.inputs == {"old": True} + assert domain_execution.process_data == { + "workflow_agent_binding_id": "workflow-binding-1", + "workflow_tool_invocation_id": "tool-call-1", + } + + def test_handle_node_retry_preserves_workflow_agent_binding_identity(self): + layer, _, _, _ = _make_layer(node_run_indices={"exec": 1}) + layer.on_event(GraphRunStartedEvent()) + started_at = _naive_utc_now() + layer.on_event( + NodeRunStartedEvent( + id="exec", + node_id="node", + node_type=BuiltinNodeTypes.AGENT, + node_title="Agent", + start_at=started_at, + ) + ) + + layer.on_event( + NodeRunRetryEvent( + id="exec", + node_id="node", + node_type=BuiltinNodeTypes.AGENT, + node_title="Agent", + start_at=started_at, + error="retry", + retry_index=1, + node_run_result=NodeRunResult( + process_data={ + "workflow_agent_binding_id": "workflow-binding-1", + "workflow_tool_invocation_id": "tool-call-1", + }, + ), + ) + ) + + assert layer._node_execution_cache["exec"].process_data["workflow_agent_binding_id"] == "workflow-binding-1" + assert layer._node_execution_cache["exec"].process_data["workflow_tool_invocation_id"] == "tool-call-1" + + def test_get_node_execution_raises_for_missing(self): + layer, _, _, _ = _make_layer() + with pytest.raises(ValueError, match="Node execution not found"): + layer._get_node_execution("missing") + + def test_get_workflow_execution_raises_when_uninitialized(self): + layer, _, _, _ = _make_layer() + + with pytest.raises(ValueError, match="workflow execution not initialized"): + layer._get_workflow_execution() diff --git a/api/tests/unit_tests/services/workflow/test_workflow_restore.py b/api/tests/unit_tests/services/test_workflow_restore.py similarity index 100% rename from api/tests/unit_tests/services/workflow/test_workflow_restore.py rename to api/tests/unit_tests/services/test_workflow_restore.py diff --git a/api/tests/unit_tests/services/test_workflow_run_agg.py b/api/tests/unit_tests/services/test_workflow_run_agg.py new file mode 100644 index 00000000000..1d395a6e67c --- /dev/null +++ b/api/tests/unit_tests/services/test_workflow_run_agg.py @@ -0,0 +1,1047 @@ +import json +from datetime import UTC, datetime, timedelta +from pathlib import Path +from uuid import uuid4 + +import pytest +from sqlalchemy import Engine, select +from sqlalchemy.orm import Session + +from clients.agent_backend import FakeAgentBackendRunClient, FakeAgentBackendScenario +from core.app.apps.pipeline.pipeline_runner import PipelineRunner +from core.app.apps.workflow.app_config_manager import WorkflowAppConfig +from core.app.apps.workflow.app_queue_manager import WorkflowAppQueueManager +from core.app.apps.workflow.app_runner import WorkflowAppRunner +from core.app.entities.app_invoke_entities import InvokeFrom, WorkflowAppGenerateEntity +from core.app.entities.queue_entities import ( + QueueNodeStartedEvent, + QueueWorkflowFailedEvent, + QueueWorkflowPausedEvent, + QueueWorkflowSucceededEvent, +) +from core.app.layers.pause_state_persist_layer import PauseStateLayerConfig, WorkflowResumptionContext +from core.app.layers.trigger_post_layer import TriggerPostLayer +from core.repositories.sqlalchemy_workflow_execution_repository import SQLAlchemyWorkflowExecutionRepository +from core.repositories.sqlalchemy_workflow_node_execution_repository import SQLAlchemyWorkflowNodeExecutionRepository +from core.workflow.nodes.human_input.entities import HumanInputNodeData, ParagraphInputConfig, UserActionConfig +from extensions.ext_storage import storage +from extensions.storage.opendal_storage import OpenDALStorage +from graphon.engine_events import ( + GraphRunAbortedEvent, + GraphRunFailedEvent, + GraphRunPartialSucceededEvent, + GraphRunPausedEvent, + GraphRunStartedEvent, + GraphRunSucceededEvent, +) +from graphon.enums import WorkflowExecutionStatus +from graphon.nodes.start.entities import StartNodeData +from graphon.runtime import RuntimeState +from graphon.variable_loader import DUMMY_VARIABLE_LOADER +from models import Account +from models.agent import AgentWorkspace, AgentWorkspaceBinding, WorkflowAgentBindingType +from models.dataset import Document +from models.enums import AppTriggerType, CreatorUserRole, WorkflowRunTriggeredFrom, WorkflowTriggerStatus +from models.human_input import HumanInputForm +from models.model import App, AppMode +from models.tools import WorkflowToolProvider +from models.trigger import WorkflowTriggerLog +from models.workflow import ( + Workflow, + WorkflowNodeExecutionModel, + WorkflowNodeExecutionTriggeredFrom, + WorkflowPause, + WorkflowRun, +) +from services.agent.workspace_service import AgentWorkspaceService +from services.workflow.entities import WorkflowScheduleCFSPlanEntity +from services.workflow_run_agg import WorkflowRunAgg +from tasks.workflow_cfs_scheduler.cfs_scheduler import AsyncWorkflowCFSPlanEntity +from tasks.workflow_cfs_scheduler.entities import AsyncWorkflowQueue +from tests.unit_tests.core.app.apps.pipeline.test_pipeline_runner import ( + _build_app_generate_entity, + _build_runner, + _document, + _persist_scope, + _workflow, +) +from tests.unit_tests.core.workflow.nodes.agent_v2.test_binding_resolver import _agent, _binding, _snapshot +from tests.unit_tests.core.workflow.nodes.agent_v2.test_session_store import _install_backend_client + +"""Pipeline tool lifecycle through the real Engine and SQLAlchemy repositories.""" + + +def _pipeline_tool_runner( + sqlite_engine: Engine, invoke_from: InvokeFrom, source_node_data: dict[str, object] +) -> tuple[PipelineRunner, SQLAlchemyWorkflowExecutionRepository, SQLAlchemyWorkflowNodeExecutionRepository]: + source_graph = { + "nodes": [ + {"id": "source-start", "data": {"type": "start", "title": "Start", "variables": list[object]()}}, + {"id": "source-node", "data": source_node_data}, + {"id": "source-end", "data": {"type": "end", "title": "End", "outputs": list[object]()}}, + ], + "edges": [ + {"source": "source-start", "target": "source-node"}, + {"source": "source-node", "target": "source-end", "sourceHandle": "source"}, + ], + } + with Session(sqlite_engine, expire_on_commit=False) as session: + session.add( + App( + id="source-app", + tenant_id="tenant", + name="Source", + mode=AppMode.WORKFLOW, + enable_site=False, + enable_api=False, + ) + ) + session.add( + Workflow( + id="source-workflow", + tenant_id="tenant", + app_id="source-app", + type="workflow", + version="v1", + graph=json.dumps(source_graph), + features="{}", + created_by="user", + ) + ) + provider = WorkflowToolProvider( + tenant_id="tenant", + app_id="source-app", + user_id="user", + version="v1", + name="source", + label="Source", + description="Source", + icon="{}", + ) + session.add(provider) + session.flush() + graph = { + "nodes": [ + {"id": "start", "data": {"type": "start", "title": "Start", "variables": list[object]()}}, + { + "id": "tool", + "data": { + "type": "tool", + "title": "Tool", + "provider_id": provider.id, + "provider_type": "workflow", + "provider_name": "source", + "tool_name": "source", + "tool_label": "Source", + "tool_configurations": dict[str, object](), + "tool_parameters": dict[str, object](), + }, + }, + ], + "edges": [{"source": "start", "target": "tool"}], + } + _, _, workflow = _persist_scope(session, workflow=_workflow(graph=graph), documents=(_document(),)) + entity = _build_app_generate_entity().model_copy( + update={"invoke_from": invoke_from, "document_id": None if invoke_from == InvokeFrom.DEBUGGER else "doc"} + ) + runner = _build_runner(entity, sqlite_engine, workflow) + user = Account(name="User", email="pipeline@example.com") + user.id = entity.user_id + runs = SQLAlchemyWorkflowExecutionRepository( + sqlite_engine, + "tenant", + user, + "pipe", + WorkflowRunTriggeredFrom.RAG_PIPELINE_DEBUGGING + if invoke_from == InvokeFrom.DEBUGGER + else WorkflowRunTriggeredFrom.RAG_PIPELINE_RUN, + ) + nodes = SQLAlchemyWorkflowNodeExecutionRepository( + sqlite_engine, "tenant", user, "pipe", WorkflowNodeExecutionTriggeredFrom.RAG_PIPELINE_RUN + ) + return runner, runs, nodes + + +class TestPipelineWorkflowToolsCases: + @staticmethod + @pytest.mark.parametrize("invoke_from", [InvokeFrom.DEBUGGER, InvokeFrom.PUBLISHED_PIPELINE]) + def test_pipeline_rejects_human_input_tool_before_creating_run_or_form( + sqlite_engine: Engine, invoke_from: InvokeFrom + ) -> None: + human_input = HumanInputNodeData( + title="Approval", user_actions=[UserActionConfig(id="approve", title="Approve")] + ) + runner, runs, nodes = _pipeline_tool_runner(sqlite_engine, invoke_from, human_input.model_dump(mode="json")) + + with pytest.raises(ValueError, match="Human Input nodes are not supported"): + WorkflowRunAgg.run(runner, None, runs, nodes) + + with Session(sqlite_engine) as session: + assert session.scalar(select(WorkflowRun)) is None + assert session.scalar(select(WorkflowNodeExecutionModel)) is None + assert session.scalar(select(HumanInputForm)) is None + document = session.get(Document, "doc") + assert document is not None + if invoke_from == InvokeFrom.PUBLISHED_PIPELINE: + assert document.indexing_status == "error" + assert document.error == "Human Input nodes are not supported in this execution context." + else: + assert document.indexing_status == "waiting" + assert document.error is None + + @staticmethod + @pytest.mark.parametrize("invoke_from", [InvokeFrom.DEBUGGER, InvokeFrom.PUBLISHED_PIPELINE]) + @pytest.mark.parametrize( + ("scenario", "status"), + [(FakeAgentBackendScenario.SUCCESS, "succeeded"), (FakeAgentBackendScenario.FAILED, "failed")], + ) + def test_pipeline_retires_nested_agent_workspace_when_run_ends( + sqlite_engine: Engine, + monkeypatch: pytest.MonkeyPatch, + invoke_from: InvokeFrom, + scenario: FakeAgentBackendScenario, + status: str, + ) -> None: + runner, runs, nodes = _pipeline_tool_runner( + sqlite_engine, + invoke_from, + {"type": "agent", "version": "2", "agent_node_kind": "dify_agent", "title": "Agent"}, + ) + with Session(sqlite_engine) as session: + agent = _agent(tenant_id="tenant") + session.add(agent) + session.flush() + snapshot = _snapshot(tenant_id="tenant", agent_id=agent.id) + session.add(snapshot) + session.flush() + session.add( + _binding( + ids={ + "tenant_id": "tenant", + "app_id": "source-app", + "workflow_id": "source-workflow", + "node_id": "source-node", + }, + agent_id=agent.id, + snapshot_id=snapshot.id, + binding_type=WorkflowAgentBindingType.INLINE_AGENT, + ) + ) + session.commit() + backend_client = _install_backend_client(monkeypatch) + monkeypatch.setattr( + "clients.agent_backend.factory.create_agent_backend_run_client", + lambda **_kwargs: FakeAgentBackendRunClient(scenario=scenario), + ) + monkeypatch.setattr( + "core.workflow.nodes.agent_v2.runtime_request_builder.resolve_model_context_window", lambda **_kwargs: None + ) + collections: list[dict[str, object]] = [] + + def collect_task(**kwargs: object) -> None: + collections.append(kwargs) + + monkeypatch.setattr("tasks.collect_agent_resources_task.collect_agent_resources.delay", collect_task) + + WorkflowRunAgg.run(runner, None, runs, nodes) + list(runner._queue_manager.listen()) + + with Session(sqlite_engine) as session: + run = session.get(WorkflowRun, runner.application_generate_entity.workflow_execution_id) + assert run is not None + assert run.status == status + agent_execution = session.scalars( + select(WorkflowNodeExecutionModel).where(WorkflowNodeExecutionModel.node_id == "source-node") + ).one() + assert agent_execution.status == status + workspace = session.scalars(select(AgentWorkspace)).one() + binding = session.scalars(select(AgentWorkspaceBinding)).one() + assert (workspace.app_id, workspace.owner_id, workspace.status) == ("source-app", run.id, "retired") + assert binding.status == "retired" + workspace_id, binding_id = workspace.id, binding.id + assert collections == [ + { + "tenant_id": "tenant", + "workspace_ids": [workspace_id], + "binding_ids": [], + "home_snapshot_ids": [], + "purge_agent_ids": [], + } + ] + AgentWorkspaceService.collect_retired_workspace(tenant_id="tenant", workspace_id=workspace_id) + backend_client.destroy_execution_binding_sync.assert_called_once() + with Session(sqlite_engine) as session: + assert session.get(AgentWorkspace, workspace_id) is None + assert session.get(AgentWorkspaceBinding, binding_id) is None + + +"""Workflow orchestration against real Engine, SQLAlchemy, and filesystem storage.""" + + +def make_workflow_runner( + sqlite_engine: Engine, *, human_input: bool = False +) -> tuple[WorkflowAppRunner, SQLAlchemyWorkflowExecutionRepository, SQLAlchemyWorkflowNodeExecutionRepository]: + """Build the same runner/repositories/queue used by a workflow worker.""" + tenant_id, app_id, workflow_id, user_id, run_id = (str(uuid4()) for _ in range(5)) + graph: dict[str, list[dict[str, object]]] = { + "nodes": [{"id": "start", "data": {"type": "start", **StartNodeData(title="Start").model_dump(mode="json")}}], + "edges": [], + } + if human_input: + data = HumanInputNodeData( + title="Approval", + form_content="Decision: {{#$output.answer#}}", + inputs=[ParagraphInputConfig(output_variable_name="answer")], + user_actions=[UserActionConfig(id="approve", title="Approve")], + ) + graph["nodes"].append({"id": "approval", "data": {"type": "human-input", **data.model_dump(mode="json")}}) + graph["edges"].append({"id": "start-approval", "source": "start", "target": "approval"}) + graph["nodes"].append( + { + "id": "end", + "data": { + "type": "end", + "title": "End", + "outputs": [ + {"variable": "answer", "value_selector": ["approval", "answer"]}, + ], + }, + } + ) + graph["edges"].append({"id": "approval-end", "source": "approval", "target": "end", "sourceHandle": "approve"}) + workflow = Workflow( + id=workflow_id, + tenant_id=tenant_id, + app_id=app_id, + type="workflow", + version="draft", + graph=json.dumps(graph), + features="{}", + created_by=user_id, + ) + entity = WorkflowAppGenerateEntity( + task_id=str(uuid4()), + app_config=WorkflowAppConfig( + tenant_id=tenant_id, app_id=app_id, workflow_id=workflow_id, app_mode=AppMode.WORKFLOW + ), + inputs={}, + files=[], + user_id=user_id, + stream=True, + invoke_from=InvokeFrom.DEBUGGER, + workflow_execution_id=run_id, + ) + user = Account(name="Workflow Tester", email=f"{user_id}@example.com") + user.id = user_id + queue = WorkflowAppQueueManager(entity.task_id, user_id, entity.invoke_from, AppMode.WORKFLOW) + workflow_repository = SQLAlchemyWorkflowExecutionRepository( + sqlite_engine, + tenant_id=tenant_id, + user=user, + app_id=app_id, + triggered_from=WorkflowRunTriggeredFrom.DEBUGGING, + ) + node_repository = SQLAlchemyWorkflowNodeExecutionRepository( + sqlite_engine, + tenant_id=tenant_id, + user=user, + app_id=app_id, + triggered_from=WorkflowNodeExecutionTriggeredFrom.WORKFLOW_RUN, + ) + runner = WorkflowAppRunner( + application_generate_entity=entity, + queue_manager=queue, + variable_loader=DUMMY_VARIABLE_LOADER, + workflow=workflow, + system_user_id=user_id, + execution_driver=WorkflowRunAgg.run, + ) + return runner, workflow_repository, node_repository + + +class TestWorkflowRunAggCases: + @staticmethod + def test_driver_publishes_persisted_start_index_and_success(sqlite_engine: Engine) -> None: + runner, workflow_repository, node_repository = make_workflow_runner(sqlite_engine) + WorkflowRunAgg.run(runner, None, workflow_repository, node_repository) + events = [message.event for message in runner._queue_manager.listen()] + starts = [event for event in events if isinstance(event, QueueNodeStartedEvent)] + assert len(starts) == 1 + assert isinstance(events[-1], QueueWorkflowSucceededEvent) + with Session(sqlite_engine) as session: + execution = session.scalar(select(WorkflowNodeExecutionModel)) + assert execution is not None + assert execution.index == starts[0].node_run_index == 1 + run = session.get(WorkflowRun, runner.application_generate_entity.workflow_execution_id) + assert run is not None + assert run.status == "succeeded" + + @staticmethod + @pytest.mark.parametrize("snapshot", [None, {}, {"tools": []}]) + def test_invalid_saved_tool_snapshot_is_rejected(sqlite_engine: Engine, snapshot: object) -> None: + from pydantic import ValidationError + + runner, workflow_repository, node_repository = make_workflow_runner(sqlite_engine) + runner.application_generate_entity.extras["workflow_tool_snapshot"] = snapshot + with pytest.raises(ValidationError): + WorkflowRunAgg.run(runner, None, workflow_repository, node_repository) + with Session(sqlite_engine) as session: + assert session.scalar(select(WorkflowRun)) is None + + @staticmethod + @pytest.fixture + def _local_storage(tmp_path: Path, monkeypatch: pytest.MonkeyPatch) -> None: + monkeypatch.setattr(storage, "storage_runner", OpenDALStorage(scheme="fs", root=str(tmp_path)), raising=False) + from core.app.apps import workflow_app_runner + + monkeypatch.setattr(workflow_app_runner.dispatch_human_input_email_task, "apply_async", lambda **_kwargs: None) + + @staticmethod + @pytest.mark.usefixtures("_local_storage") + @pytest.mark.parametrize("persist_snapshot", [False, True]) + def test_pause_publication_persists_summary_and_optional_snapshot( + sqlite_engine: Engine, persist_snapshot: bool, caplog: pytest.LogCaptureFixture + ) -> None: + runner, workflow_repository, node_repository = make_workflow_runner(sqlite_engine, human_input=True) + config = ( + PauseStateLayerConfig(sqlite_engine, runner.application_generate_entity.user_id) + if persist_snapshot + else None + ) + WorkflowRunAgg.run(runner, config, workflow_repository, node_repository) + assert "Failed to enqueue human input email task" not in caplog.text + events = [message.event for message in runner._queue_manager.listen()] + assert isinstance(events[-1], QueueWorkflowPausedEvent), events + with Session(sqlite_engine) as session: + run = session.get(WorkflowRun, runner.application_generate_entity.workflow_execution_id) + assert run is not None + pause = session.scalar(select(WorkflowPause).where(WorkflowPause.workflow_run_id == run.id)) + assert run.status == "paused" + assert run.finished_at is None + assert run.total_steps == 2 + if not persist_snapshot: + assert pause is None + return + assert pause is not None + snapshot = WorkflowResumptionContext.loads(storage.load(pause.state_object_key).decode()) + restored = RuntimeState.from_snapshot(snapshot.serialized_graph_runtime_state) + assert restored.graph_execution.paused + assert snapshot.get_response_stream_filter().dumps() == snapshot.serialized_response_stream_filter_state + assert snapshot.get_generate_entity().model_dump( + mode="json" + ) == runner.application_generate_entity.model_dump(mode="json") + + @staticmethod + @pytest.mark.usefixtures("_local_storage") + def test_pause_database_failure_publishes_failure_instead_of_resume_ready( + sqlite_engine: Engine, monkeypatch: pytest.MonkeyPatch + ) -> None: + def fail_pause(*_args: object, **_kwargs: object) -> None: + raise RuntimeError("pause-write-failed") + + monkeypatch.setattr( + "repositories.sqlalchemy_api_workflow_run_repository.DifyAPISQLAlchemyWorkflowRunRepository.create_workflow_pause", + fail_pause, + ) + retirements: list[dict[str, object]] = [] + collections: list[dict[str, object]] = [] + + def retire(_store: object, **kwargs: object) -> list[str]: + retirements.append(kwargs) + return ["workspace-1"] + + monkeypatch.setattr( + "core.workflow.nodes.agent_v2.workspace_retirement_layer.WorkflowAgentWorkspaceStore.retire_workflow_run", + retire, + ) + monkeypatch.setattr( + "core.workflow.nodes.agent_v2.workspace_retirement_layer.enqueue_agent_resource_collection", + lambda **kwargs: collections.append(kwargs), + ) + runner, workflow_repository, node_repository = make_workflow_runner(sqlite_engine, human_input=True) + trigger_log = WorkflowTriggerLog( + tenant_id=runner.application_generate_entity.app_config.tenant_id, + app_id=runner.application_generate_entity.app_config.app_id, + workflow_id=runner._workflow.id, + workflow_run_id=None, + root_node_id=None, + trigger_metadata="{}", + trigger_type=AppTriggerType.TRIGGER_SCHEDULE, + trigger_data="{}", + inputs="{}", + outputs=None, + status=WorkflowTriggerStatus.RUNNING, + error=None, + queue_name="workflow", + celery_task_id=None, + created_by_role=CreatorUserRole.ACCOUNT, + created_by=runner.application_generate_entity.user_id, + elapsed_time=5.0, + ) + with Session(sqlite_engine, expire_on_commit=False) as session: + session.add(trigger_log) + session.commit() + start_time = datetime.now(UTC) - timedelta(seconds=10) + plan = AsyncWorkflowCFSPlanEntity( + schedule_strategy=WorkflowScheduleCFSPlanEntity.Strategy.Nop, + queue=AsyncWorkflowQueue.PROFESSIONAL_QUEUE, + ) + runner._graph_engine_layers = (TriggerPostLayer(plan, start_time, trigger_log.id),) + WorkflowRunAgg.run( + runner, + PauseStateLayerConfig(sqlite_engine, runner.application_generate_entity.user_id), + workflow_repository, + node_repository, + ) + events = [message.event for message in runner._queue_manager.listen()] + assert isinstance(events[-1], QueueWorkflowFailedEvent), events + assert "pause-write-failed" in events[-1].error + assert not any(isinstance(event, QueueWorkflowPausedEvent) for event in events) + assert retirements == [ + { + "tenant_id": runner.application_generate_entity.app_config.tenant_id, + "app_id": runner.application_generate_entity.app_config.app_id, + "workflow_run_id": runner.application_generate_entity.workflow_execution_id, + } + ] + assert collections == [ + { + "tenant_id": runner.application_generate_entity.app_config.tenant_id, + "workspace_ids": ["workspace-1"], + } + ] + with Session(sqlite_engine) as session: + persisted_trigger = session.get(WorkflowTriggerLog, trigger_log.id) + assert persisted_trigger is not None + assert persisted_trigger.status == WorkflowTriggerStatus.FAILED + assert persisted_trigger.finished_at is not None + assert persisted_trigger.elapsed_time == pytest.approx( + 5.0 + (persisted_trigger.finished_at.replace(tzinfo=UTC) - start_time).total_seconds() + ) + run = session.get(WorkflowRun, runner.application_generate_entity.workflow_execution_id) + assert run is not None + assert run.status == "failed" + assert session.scalar(select(WorkflowPause)) is None + + @staticmethod + def test_late_form_submission_is_refreshed_before_the_real_queue_closes(sqlite_session: Session) -> None: + from core.app.apps.workflow_app_runner import WorkflowBasedAppRunner + from core.app.entities.queue_entities import QueueHumanInputFormFilledEvent + from core.repositories.human_input_repository import HumanInputFormSubmissionRepository + from core.workflow.nodes.human_input.enums import HumanInputFormStatus + from graphon.engine_events import NodeRunStartedEvent + from tests.unit_tests.core.app.apps.test_workflow_app_runner import ( + _make_paused_workflow, + _resume_events, + _save_form, + ) + + first = _save_form(sqlite_session, status=HumanInputFormStatus.SUBMITTED) + late = _save_form(sqlite_session, status=HumanInputFormStatus.WAITING) + queue = WorkflowAppQueueManager("form-refresh-task", "user", InvokeFrom.DEBUGGER, AppMode.WORKFLOW) + runner = WorkflowBasedAppRunner(queue_manager=queue, app_id="app") + entry = _make_paused_workflow(runner, [first, late]) + for event in _resume_events(runner, entry): + runner.handle_event(entry, event) + if isinstance(event, NodeRunStartedEvent): + HumanInputFormSubmissionRepository().mark_submitted( + form_id=late.id, + recipient_id=None, + selected_action_id="approve", + form_data={"answer": "late answer"}, + submission_user_id=None, + submission_end_user_id=None, + ) + events = [message.event for message in queue.listen()] + completions = [event for event in events if isinstance(event, QueueHumanInputFormFilledEvent)] + assert [event.form_id for event in completions] == [first.id, late.id] + assert completions[-1].rendered_content == "Decision: late answer" + assert isinstance(events[-1], QueueWorkflowSucceededEvent) + + @staticmethod + @pytest.mark.usefixtures("_local_storage") + def test_durable_pause_resumes_with_one_form_completion_and_continuing_indexes(sqlite_engine: Engine) -> None: + from core.app.entities.queue_entities import QueueHumanInputFormFilledEvent, QueueWorkflowStartedEvent + from core.repositories.human_input_repository import HumanInputFormSubmissionRepository + from models.human_input import HumanInputForm + + original, workflow_repository, node_repository = make_workflow_runner(sqlite_engine, human_input=True) + config = PauseStateLayerConfig(sqlite_engine, original.application_generate_entity.user_id) + WorkflowRunAgg.run(original, config, workflow_repository, node_repository) + before = [message.event for message in original._queue_manager.listen()] + assert isinstance(before[-1], QueueWorkflowPausedEvent) + with Session(sqlite_engine) as session: + pause = session.scalar(select(WorkflowPause)) + form = session.scalar(select(HumanInputForm)) + assert pause is not None + assert form is not None + form_id = form.id + snapshot = WorkflowResumptionContext.loads(storage.load(pause.state_object_key).decode()) + HumanInputFormSubmissionRepository().mark_submitted( + form_id=form_id, + recipient_id=None, + selected_action_id="approve", + form_data={"answer": "approved result"}, + submission_user_id=None, + submission_end_user_id=None, + ) + entity = original.application_generate_entity.model_copy(update={"task_id": str(uuid4())}) + queue = WorkflowAppQueueManager(entity.task_id, entity.user_id, entity.invoke_from, AppMode.WORKFLOW) + resumed = WorkflowAppRunner( + application_generate_entity=entity, + queue_manager=queue, + variable_loader=DUMMY_VARIABLE_LOADER, + workflow=original._workflow, + system_user_id=entity.user_id, + graph_runtime_state=RuntimeState.from_snapshot(snapshot.serialized_graph_runtime_state), + response_stream_filter=snapshot.get_response_stream_filter(), + execution_driver=WorkflowRunAgg.run, + ) + WorkflowRunAgg.run(resumed, config, workflow_repository, node_repository) + after = [message.event for message in queue.listen()] + assert isinstance(after[0], QueueWorkflowStartedEvent) + assert len(after[0].node_execution_snapshots) == 2 + completions = [event for event in after if isinstance(event, QueueHumanInputFormFilledEvent)] + assert [event.form_id for event in completions] == [form_id] + starts = [event for event in after if isinstance(event, QueueNodeStartedEvent)] + assert [(event.node_id, event.node_run_index) for event in starts] == [("approval", 2), ("end", 3)] + assert isinstance(after[-1], QueueWorkflowSucceededEvent) + assert after[-1].outputs == {"answer": "approved result"} + with Session(sqlite_engine) as session: + executions = session.scalars( + select(WorkflowNodeExecutionModel).order_by(WorkflowNodeExecutionModel.index) + ).all() + assert [(row.node_id, row.index) for row in executions] == [("start", 1), ("approval", 2), ("end", 3)] + + @staticmethod + @pytest.mark.usefixtures("_local_storage") + def test_pause_before_first_tool_call_preserves_its_published_definition(sqlite_engine: Engine) -> None: + from core.app.entities.queue_entities import QueueHumanInputFormFilledEvent + from core.repositories.human_input_repository import HumanInputFormSubmissionRepository + from graphon.nodes.tool.entities import ToolNodeData + from models.human_input import HumanInputForm + from models.model import App + from models.tools import WorkflowToolProvider + + original, workflow_repository, node_repository = make_workflow_runner(sqlite_engine, human_input=True) + original_entity = original.application_generate_entity + tenant_id = original_entity.app_config.tenant_id + source_app_id = str(uuid4()) + source_workflow_ids = [str(uuid4()), str(uuid4())] + with Session(sqlite_engine) as session, session.begin(): + session.add( + App( + id=source_app_id, + tenant_id=tenant_id, + name="Source", + mode=AppMode.WORKFLOW, + enable_site=False, + enable_api=False, + ) + ) + for workflow_id, version, marker in zip( + source_workflow_ids, ("old", "new"), ("original", "republished"), strict=True + ): + session.add( + Workflow( + id=workflow_id, + app_id=source_app_id, + tenant_id=tenant_id, + created_by=original_entity.user_id, + type="workflow", + version=version, + features="{}", + graph=json.dumps( + { + "nodes": [ + { + "id": "source-start", + "data": { + "type": "start", + "title": "Start", + "variables": [ + { + "variable": "marker", + "type": "text-input", + "label": "Marker", + "default": marker, + "required": False, + } + ], + }, + }, + { + "id": "source-end", + "data": { + "type": "end", + "title": "End", + "outputs": [ + {"variable": "marker", "value_selector": ["source-start", "marker"]} + ], + }, + }, + ], + "edges": [{"source": "source-start", "target": "source-end"}], + } + ), + ) + ) + provider = WorkflowToolProvider( + tenant_id=tenant_id, + app_id=source_app_id, + version="old", + name="source", + label="Original source", + icon="{}", + user_id=original_entity.user_id, + description="Source", + parameter_configuration='[{"name":"marker","description":"Marker","form":"llm"}]', + ) + session.add(provider) + session.flush() + provider_id = provider.id + graph = dict(original._workflow.graph_dict) + graph["nodes"].append( + { + "id": "tool", + "data": { + "type": "tool", + **ToolNodeData( + title="Tool", + provider_id=provider_id, + provider_type="workflow", + provider_name="source", + tool_name="source", + tool_label="Source", + tool_configurations={}, + tool_parameters={}, + ).model_dump(mode="json"), + }, + } + ) + graph["nodes"][2]["data"]["outputs"] = [{"variable": "result", "value_selector": ["tool", "json"]}] + graph["edges"][1]["target"] = "tool" + graph["edges"].append({"source": "tool", "target": "end"}) + original._workflow.graph = json.dumps(graph) + config = PauseStateLayerConfig(sqlite_engine, original_entity.user_id) + WorkflowRunAgg.run(original, config, workflow_repository, node_repository) + before = [message.event for message in original._queue_manager.listen()] + assert isinstance(before[-1], QueueWorkflowPausedEvent) + assert not any(isinstance(event, QueueNodeStartedEvent) and event.node_id == "tool" for event in before) + with Session(sqlite_engine) as session, session.begin(): + pause = session.scalar(select(WorkflowPause)) + form = session.scalar(select(HumanInputForm)) + provider = session.get(WorkflowToolProvider, provider_id) + assert pause is not None + assert form is not None + assert provider is not None + form_id = form.id + snapshot = WorkflowResumptionContext.loads(storage.load(pause.state_object_key).decode()) + provider.version = "new" + provider.label = "Republished source" + HumanInputFormSubmissionRepository().mark_submitted( + form_id=form_id, + recipient_id=None, + selected_action_id="approve", + form_data={"answer": "approved"}, + submission_user_id=None, + submission_end_user_id=None, + ) + entity = snapshot.get_generate_entity().model_copy(update={"task_id": str(uuid4())}) + assert isinstance(entity, WorkflowAppGenerateEntity) + queue = WorkflowAppQueueManager(entity.task_id, entity.user_id, entity.invoke_from, AppMode.WORKFLOW) + resumed = WorkflowAppRunner( + application_generate_entity=entity, + queue_manager=queue, + variable_loader=DUMMY_VARIABLE_LOADER, + workflow=original._workflow, + system_user_id=entity.user_id, + graph_runtime_state=RuntimeState.from_snapshot(snapshot.serialized_graph_runtime_state), + response_stream_filter=snapshot.get_response_stream_filter(), + execution_driver=WorkflowRunAgg.run, + ) + WorkflowRunAgg.run(resumed, config, workflow_repository, node_repository) + after = [message.event for message in queue.listen()] + assert [event.form_id for event in after if isinstance(event, QueueHumanInputFormFilledEvent)] == [form_id] + assert isinstance(after[-1], QueueWorkflowSucceededEvent), after + assert after[-1].outputs == {"result": [{"marker": "original"}]} + assert entity.extras == original_entity.extras + with Session(sqlite_engine) as session: + child_workflows = session.scalars( + select(WorkflowNodeExecutionModel.workflow_id).where(WorkflowNodeExecutionModel.app_id == source_app_id) + ).all() + assert child_workflows + assert set(child_workflows) == {source_workflow_ids[0]} + + # A new execution resolves the newly published graph and parameter defaults. + graph["nodes"] = [node for node in graph["nodes"] if node["id"] != "approval"] + graph["edges"] = [{"source": "start", "target": "tool"}, {"source": "tool", "target": "end"}] + original._workflow.graph = json.dumps(graph) + fresh_entity = original_entity.model_copy( + update={"task_id": str(uuid4()), "workflow_execution_id": str(uuid4()), "extras": {}} + ) + fresh_queue = WorkflowAppQueueManager( + fresh_entity.task_id, fresh_entity.user_id, fresh_entity.invoke_from, AppMode.WORKFLOW + ) + fresh = WorkflowAppRunner( + application_generate_entity=fresh_entity, + queue_manager=fresh_queue, + variable_loader=DUMMY_VARIABLE_LOADER, + workflow=original._workflow, + system_user_id=fresh_entity.user_id, + execution_driver=WorkflowRunAgg.run, + ) + WorkflowRunAgg.run(fresh, None, workflow_repository, node_repository) + fresh_events = [message.event for message in fresh_queue.listen()] + assert isinstance(fresh_events[-1], QueueWorkflowSucceededEvent), fresh_events + assert fresh_events[-1].outputs == {"result": [{"marker": "republished"}]} + + @staticmethod + def test_active_execution_snapshot_failure_is_not_published_as_a_pause(sqlite_engine: Engine) -> None: + from graphon.engine_events import GraphRunFailedEvent, GraphRunPausedEvent + + runner, workflow_repository, node_repository = make_workflow_runner(sqlite_engine, human_input=True) + prepared = runner.prepare() + assert prepared is not None + aggregate = WorkflowRunAgg( + prepared, + runner, + PauseStateLayerConfig(sqlite_engine, runner.application_generate_entity.user_id), + workflow_execution_repository=workflow_repository, + workflow_node_execution_repository=node_repository, + ) + # A participant that has not shut down must prevent a resume-ready snapshot. + with prepared.entry.graph_engine.runtime_state.graph_execution.track_execution(): + events = list(aggregate.iter_events()) + assert isinstance(events[-1], GraphRunFailedEvent) + assert "during active execution" in events[-1].error + assert not any(isinstance(event, GraphRunPausedEvent) for event in events) + with Session(sqlite_engine) as session: + run = session.get(WorkflowRun, runner.application_generate_entity.workflow_execution_id) + assert run is not None + assert run.status == "failed" + assert session.scalar(select(WorkflowPause)) is None + + @staticmethod + @pytest.mark.usefixtures("_local_storage") + @pytest.mark.parametrize( + ("terminal_event", "status"), + [ + (GraphRunSucceededEvent(), WorkflowExecutionStatus.SUCCEEDED), + (GraphRunPartialSucceededEvent(exceptions_count=1), WorkflowExecutionStatus.PARTIAL_SUCCEEDED), + (GraphRunFailedEvent(error="failed"), WorkflowExecutionStatus.FAILED), + (GraphRunAbortedEvent(reason="stopped"), WorkflowExecutionStatus.STOPPED), + ], + ) + def test_elapsed_time_accumulates_across_pauses_and_matches_delayed_stream_responses( + monkeypatch: pytest.MonkeyPatch, + sqlite_engine: Engine, + terminal_event: GraphRunSucceededEvent + | GraphRunPartialSucceededEvent + | GraphRunFailedEvent + | GraphRunAbortedEvent, + status: WorkflowExecutionStatus, + ) -> None: + from core.app.apps.common.workflow_response_converter import WorkflowResponseConverter + from core.app.entities.task_entities import WorkflowPauseStreamResponse + from core.repositories.human_input_repository import HumanInputFormSubmissionRepository + from core.workflow.nodes.human_input.enums import HumanInputFormStatus + from core.workflow.system_variables import build_system_variables + from graphon.entities import WorkflowStartReason + from models.human_input import HumanInputForm + + clock = [100.0] + monkeypatch.setattr("repositories.entities.workflow_execution.perf_counter", lambda: clock[0]) + original, workflow_repository, node_repository = make_workflow_runner(sqlite_engine, human_input=True) + graph = original._workflow.graph_dict + graph["nodes"].append({"id": "second-approval", "data": graph["nodes"][1]["data"]}) + graph["edges"][1]["target"] = "second-approval" + graph["edges"].append( + {"id": "second-end", "source": "second-approval", "target": "end", "sourceHandle": "approve"} + ) + original._workflow.graph = json.dumps(graph) + runner = original + run_id = runner.application_generate_entity.workflow_execution_id + config = PauseStateLayerConfig(sqlite_engine, runner.application_generate_entity.user_id) + account = Account(name="Test", email="test@example.com") + account.id = runner.application_generate_entity.user_id + total = 0.0 + + for index, duration in enumerate([2.0, 3.0, 4.0]): + prepared = runner.prepare() + assert prepared is not None + aggregate = WorkflowRunAgg( + prepared, + runner, + config, + workflow_execution_repository=workflow_repository, + workflow_node_execution_repository=node_repository, + ) + is_paused = index < 2 + if is_paused: + events = [] + for event in aggregate.iter_events(): + if isinstance(event, GraphRunStartedEvent): + with Session(sqlite_engine) as session: + started = session.get(WorkflowRun, run_id) + assert started is not None + assert started.elapsed_time == total + clock[0] += duration + events.append(event) + assert isinstance(events[-1], GraphRunPausedEvent), events + else: + aggregate._persistence_layer.on_event(GraphRunStartedEvent()) + clock[0] += duration + aggregate._persistence_layer.on_event(terminal_event) + total += duration + with Session(sqlite_engine) as session: + execution = session.get(WorkflowRun, run_id) + assert execution is not None + assert execution.elapsed_time == total + assert execution.status == (WorkflowExecutionStatus.PAUSED if is_paused else status) + assert (execution.finished_at is None) == is_paused + + # Queue processing and pause waits must not extend a settled segment. + clock[0] += 1000 + entity = runner.application_generate_entity + runtime = prepared.entry.graph_engine.runtime_state + converter = WorkflowResponseConverter( + application_generate_entity=entity, + user=account, + system_variables=build_system_variables(workflow_execution_id=run_id), + ) + converter.workflow_start_to_stream_response( + task_id=entity.task_id, + workflow_run_id=run_id, + workflow_id=original._workflow.id, + reason=WorkflowStartReason.INITIAL if index == 0 else WorkflowStartReason.RESUMPTION, + ) + if is_paused: + responses = converter.workflow_pause_to_stream_response( + event=QueueWorkflowPausedEvent(), task_id=entity.task_id, graph_runtime_state=runtime + ) + pause_response = responses[-1] + assert isinstance(pause_response, WorkflowPauseStreamResponse) + assert pause_response.data.elapsed_time == total + with Session(sqlite_engine) as session: + pause = session.scalar(select(WorkflowPause).where(WorkflowPause.workflow_run_id == run_id)) + form = session.scalar( + select(HumanInputForm).where( + HumanInputForm.workflow_run_id == run_id, + HumanInputForm.status == HumanInputFormStatus.WAITING, + ) + ) + assert pause is not None + assert form is not None + form_id = form.id + snapshot = WorkflowResumptionContext.loads(storage.load(pause.state_object_key).decode()) + entity = snapshot.get_generate_entity() + assert isinstance(entity, WorkflowAppGenerateEntity) + assert entity.workflow_execution_timing.elapsed_time == snapshot.elapsed_time == total + HumanInputFormSubmissionRepository().mark_submitted( + form_id=form_id, + recipient_id=None, + selected_action_id="approve", + form_data={"answer": "approved result"}, + submission_user_id=None, + submission_end_user_id=None, + ) + entity.task_id = str(uuid4()) + runner = WorkflowAppRunner( + application_generate_entity=entity, + queue_manager=WorkflowAppQueueManager( + entity.task_id, entity.user_id, entity.invoke_from, AppMode.WORKFLOW + ), + variable_loader=DUMMY_VARIABLE_LOADER, + workflow=original._workflow, + system_user_id=entity.user_id, + graph_runtime_state=RuntimeState.from_snapshot(snapshot.serialized_graph_runtime_state), + response_stream_filter=snapshot.get_response_stream_filter(), + execution_driver=WorkflowRunAgg.run, + ) + else: + response = converter.workflow_finish_to_stream_response( + task_id=entity.task_id, + workflow_id=original._workflow.id, + status=status, + graph_runtime_state=runtime, + ) + assert response.data.elapsed_time == 9.0 + assert entity.workflow_execution_timing.finish() == 9.0 + + @staticmethod + @pytest.mark.usefixtures("_local_storage") + @pytest.mark.parametrize("stop_timing", ["before_summary", "during_upload"]) + def test_stop_wins_while_engine_is_persisting_a_pause( + sqlite_engine: Engine, monkeypatch: pytest.MonkeyPatch, stop_timing: str + ) -> None: + from sqlalchemy.orm import sessionmaker + + from core.app.entities.queue_entities import QueueStopEvent + from core.workflow.nodes.human_input.enums import HumanInputFormStatus + from graphon.enums import WorkflowExecutionStatus + from models.human_input import HumanInputForm + from repositories.sqlalchemy_api_workflow_run_repository import DifyAPISQLAlchemyWorkflowRunRepository + + runner, workflow_repository, node_repository = make_workflow_runner(sqlite_engine, human_input=True) + entity = runner.application_generate_entity + repository = DifyAPISQLAlchemyWorkflowRunRepository(sessionmaker(sqlite_engine)) + stop_calls = [] + + def request_stop() -> None: + admitted = repository.stop_paused_workflow_task( + tenant_id=entity.app_config.tenant_id, + app_id=entity.app_config.app_id, + task_id=entity.task_id, + owner=None, + ) + assert admitted is not None + stop_calls.append(admitted.status) + + if stop_timing == "during_upload": + original_save = storage.save + + def save(key: str, data: bytes) -> None: + if key.startswith("workflow-state-"): + request_stop() + return original_save(key, data) + + monkeypatch.setattr(storage, "save", save) + else: + original_persist = WorkflowRunAgg._persist_pause + + def persist(aggregate: WorkflowRunAgg, event: GraphRunPausedEvent) -> None: + request_stop() + return original_persist(aggregate, event) + + monkeypatch.setattr(WorkflowRunAgg, "_persist_pause", persist) + + WorkflowRunAgg.run( + runner, PauseStateLayerConfig(sqlite_engine, entity.user_id), workflow_repository, node_repository + ) + events = [message.event for message in runner._queue_manager.listen()] + assert stop_calls + assert isinstance(events[-1], QueueStopEvent), events + assert not any(isinstance(event, QueueWorkflowPausedEvent) for event in events) + with Session(sqlite_engine) as session: + run = session.get(WorkflowRun, entity.workflow_execution_id) + assert run is not None + assert run.status == WorkflowExecutionStatus.STOPPED + assert run.stop_requested_at is not None + assert run.task_id == entity.task_id + assert session.scalar(select(WorkflowPause).where(WorkflowPause.workflow_run_id == run.id)) is None + forms = session.scalars(select(HumanInputForm).where(HumanInputForm.workflow_run_id == run.id)).all() + assert forms + assert all(form.status == HumanInputFormStatus.EXPIRED for form in forms) + + @staticmethod + def test_legacy_caller_rejects_human_input_before_creating_run_or_form(sqlite_engine: Engine) -> None: + from models.human_input import HumanInputForm + + runner, workflow_repository, node_repository = make_workflow_runner(sqlite_engine, human_input=True) + runner.application_generate_entity.allow_human_input = False + with pytest.raises(ValueError, match="Human Input nodes are not supported"): + WorkflowRunAgg.run(runner, None, workflow_repository, node_repository) + with Session(sqlite_engine) as session: + assert session.scalar(select(WorkflowRun)) is None + assert session.scalar(select(HumanInputForm)) is None diff --git a/api/tests/unit_tests/services/workflow/test_workflow_run_index.py b/api/tests/unit_tests/services/test_workflow_run_index.py similarity index 100% rename from api/tests/unit_tests/services/workflow/test_workflow_run_index.py rename to api/tests/unit_tests/services/test_workflow_run_index.py diff --git a/api/tests/unit_tests/services/test_workflow_run_service.py b/api/tests/unit_tests/services/test_workflow_run_service.py index db7fe74c131..b065b2692ca 100644 --- a/api/tests/unit_tests/services/test_workflow_run_service.py +++ b/api/tests/unit_tests/services/test_workflow_run_service.py @@ -1,5 +1,4 @@ -"""Unit tests for the Console workflow-run application service.""" - +from datetime import datetime from types import SimpleNamespace from unittest.mock import ANY, MagicMock, create_autospec @@ -9,22 +8,25 @@ from core.workflow.node_execution_process_data import ( WORKFLOW_TOOL_INVOCATION_ID_KEY, WORKFLOW_TOOL_ROOT_APP_ID_KEY, ) +from core.workflow.nodes.human_input.pause_reason import HumanInputRequired +from graphon.entities.pause_reason import SchedulingPause from graphon.enums import WorkflowExecutionStatus from machinery.context import RequestContext from models import WorkflowRun, WorkflowRunTriggeredFrom, WorkflowType from models.enums import CreatorUserRole -from repositories.sqlalchemy_api_workflow_run_repository import WorkflowRunMessageRef +from repositories.sqlalchemy_api_workflow_run_repository import WorkflowRunMessageRef, WorkflowRunPauseRecord from services import workflow_run_service as service_module from services.workflow_node_execution_trace_service import WorkflowNodeExecutionTrace -from services.workflow_run_service import WorkflowRunService +from services.workflow_run_service import ( + WorkflowRunPauseDetails, + WorkflowRunPausedNode, + WorkflowRunService, +) + +"""Unit tests for the Console workflow-run application service.""" -@pytest.fixture -def service_dependencies() -> tuple[MagicMock, MagicMock]: - return MagicMock(), MagicMock() - - -def _service(dependencies: tuple[MagicMock, MagicMock]) -> WorkflowRunService: +def _workflow_run_service_service(dependencies: tuple[MagicMock, MagicMock]) -> WorkflowRunService: node_executions, workflow_runs = dependencies return WorkflowRunService( workflow_runs=workflow_runs, @@ -32,7 +34,7 @@ def _service(dependencies: tuple[MagicMock, MagicMock]) -> WorkflowRunService: ) -def _request_context(*, workspace_id: str = "tenant-1") -> RequestContext: +def _workflow_run_service_request_context(*, workspace_id: str = "tenant-1") -> RequestContext: return RequestContext( request_id="request-1", trace_id="trace-1", @@ -60,195 +62,309 @@ def _workflow_run( ) -class TestWorkflowRunServiceQueries: - def test_workflow_tool_details_require_the_root_app_run_and_hide_internal_correlation( - self, service_dependencies: tuple[MagicMock, MagicMock], monkeypatch: pytest.MonkeyPatch - ) -> None: - node_executions, workflow_runs = service_dependencies - service = _service(service_dependencies) - workflow_runs.get_workflow_run_by_id.return_value = None - context = _request_context() - assert ( - service.get_workflow_tool_node_executions( +class TestWorkflowRunServiceCases: + @staticmethod + @pytest.fixture + def service_dependencies() -> tuple[MagicMock, MagicMock]: + return MagicMock(), MagicMock() + + class TestWorkflowRunServiceQueries: + def test_workflow_tool_details_require_the_root_app_run_and_hide_internal_correlation( + self, service_dependencies: tuple[MagicMock, MagicMock], monkeypatch: pytest.MonkeyPatch + ) -> None: + node_executions, workflow_runs = service_dependencies + service = _workflow_run_service_service(service_dependencies) + workflow_runs.get_workflow_run_by_id.return_value = None + context = _workflow_run_service_request_context() + assert ( + service.get_workflow_tool_node_executions( + context, app_id="app-1", run_id="run-1", node_execution_id="parent-1" + ) + is None + ) + node_executions.get_workflow_tool_executions.assert_not_called() + + workflow_runs.get_workflow_run_by_id.return_value = _workflow_run() + trace = WorkflowNodeExecutionTrace( + id="execution-1", + tenant_id="tenant-1", + app_id="source-app", + workflow_id="source-workflow", + triggered_from="workflow-tool", + created_by="account-1", + process_data={ + "answer": "approved", + WORKFLOW_TOOL_INVOCATION_ID_KEY: "invocation-1", + WORKFLOW_TOOL_ROOT_APP_ID_KEY: "app-1", + }, + ) + assemble = create_autospec(service_module.assemble_workflow_node_execution_traces, return_value=[trace]) + monkeypatch.setattr(service_module, "assemble_workflow_node_execution_traces", assemble) + result = service.get_workflow_tool_node_executions( context, app_id="app-1", run_id="run-1", node_execution_id="parent-1" ) - is None - ) - node_executions.get_workflow_tool_executions.assert_not_called() + workflow_runs.get_workflow_run_by_id.assert_called_with( + tenant_id="tenant-1", app_id="app-1", run_id="run-1" + ) + node_executions.get_workflow_tool_executions.assert_called_once_with( + tenant_id="tenant-1", workflow_run_id="run-1", parent_node_execution_id="parent-1" + ) + assemble.assert_called_once_with( + node_executions.get_workflow_tool_executions.return_value, node_executions, session=ANY + ) + assert result == [trace] + assert trace.process_data == {"answer": "approved"} - workflow_runs.get_workflow_run_by_id.return_value = _workflow_run() - trace = WorkflowNodeExecutionTrace( - id="execution-1", - tenant_id="tenant-1", - app_id="source-app", - workflow_id="source-workflow", - triggered_from="workflow-tool", - created_by="account-1", - process_data={ - "answer": "approved", - WORKFLOW_TOOL_INVOCATION_ID_KEY: "invocation-1", - WORKFLOW_TOOL_ROOT_APP_ID_KEY: "app-1", - }, - ) - assemble = create_autospec(service_module.assemble_workflow_node_execution_traces, return_value=[trace]) - monkeypatch.setattr(service_module, "assemble_workflow_node_execution_traces", assemble) - result = service.get_workflow_tool_node_executions( - context, app_id="app-1", run_id="run-1", node_execution_id="parent-1" - ) - workflow_runs.get_workflow_run_by_id.assert_called_with(tenant_id="tenant-1", app_id="app-1", run_id="run-1") - node_executions.get_workflow_tool_executions.assert_called_once_with( - tenant_id="tenant-1", workflow_run_id="run-1", parent_node_execution_id="parent-1" - ) - assemble.assert_called_once_with( - node_executions.get_workflow_tool_executions.return_value, node_executions, session=ANY - ) - assert result == [trace] - assert trace.process_data == {"answer": "approved"} + def test_get_paginate_workflow_runs_should_forward_filters_and_parse_limit( + self, + service_dependencies: tuple[MagicMock, MagicMock], + ) -> None: + _, workflow_runs = service_dependencies + service = _workflow_run_service_service(service_dependencies) + expected = MagicMock(name="pagination") + workflow_runs.get_paginated_workflow_runs.return_value = expected + args = {"limit": "7", "last_id": "last-1", "status": "succeeded"} - def test_get_paginate_workflow_runs_should_forward_filters_and_parse_limit( - self, - service_dependencies: tuple[MagicMock, MagicMock], - ) -> None: - _, workflow_runs = service_dependencies - service = _service(service_dependencies) - expected = MagicMock(name="pagination") - workflow_runs.get_paginated_workflow_runs.return_value = expected - args = {"limit": "7", "last_id": "last-1", "status": "succeeded"} + result = service.get_paginate_workflow_runs( + _workflow_run_service_request_context(workspace_id="tenant-1"), + app_id="app-1", + args=args, + triggered_from=WorkflowRunTriggeredFrom.APP_RUN, + ) - result = service.get_paginate_workflow_runs( - _request_context(workspace_id="tenant-1"), - app_id="app-1", - args=args, - triggered_from=WorkflowRunTriggeredFrom.APP_RUN, + assert result is expected + workflow_runs.get_paginated_workflow_runs.assert_called_once_with( + tenant_id="tenant-1", + app_id="app-1", + triggered_from=WorkflowRunTriggeredFrom.APP_RUN, + limit=7, + last_id="last-1", + status="succeeded", + ) + + def test_get_paginate_advanced_chat_workflow_runs_should_attach_message_fields_when_message_exists( + self, + service_dependencies: tuple[MagicMock, MagicMock], + monkeypatch: pytest.MonkeyPatch, + ) -> None: + _, workflow_runs = service_dependencies + service = _workflow_run_service_service(service_dependencies) + run_with_message = _workflow_run(status=WorkflowExecutionStatus.RUNNING) + run_without_message = _workflow_run(run_id="run-2") + pagination = SimpleNamespace(data=[run_with_message, run_without_message]) + monkeypatch.setattr(service, "get_paginate_workflow_runs", MagicMock(return_value=pagination)) + workflow_runs.get_message_refs.return_value = { + "run-1": WorkflowRunMessageRef(message_id="msg-1", conversation_id="conv-1") + } + + result = service.get_paginate_advanced_chat_workflow_runs( + _workflow_run_service_request_context(), + app_id="app-1", + args={"limit": "2"}, + ) + + assert result is pagination + assert len(result.data) == 2 + assert result.data[0].message_id == "msg-1" + assert result.data[0].conversation_id == "conv-1" + assert result.data[0].status == "running" + assert not hasattr(result.data[1], "message_id") + assert result.data[1].id == "run-2" + workflow_runs.get_message_refs.assert_called_once_with( + app_id="app-1", + workflow_run_ids=["run-1", "run-2"], + ) + + def test_get_workflow_run_should_delegate_to_repository_by_tenant_and_app( + self, + service_dependencies: tuple[MagicMock, MagicMock], + ) -> None: + _, workflow_runs = service_dependencies + service = _workflow_run_service_service(service_dependencies) + expected = _workflow_run() + workflow_runs.get_workflow_run_by_id.return_value = expected + + result = service.get_workflow_run(_workflow_run_service_request_context(), app_id="app-1", run_id="run-1") + + assert result is expected + workflow_runs.get_workflow_run_by_id.assert_called_once_with( + tenant_id="tenant-1", + app_id="app-1", + run_id="run-1", + ) + + def test_get_workflow_runs_count_should_forward_optional_filters( + self, + service_dependencies: tuple[MagicMock, MagicMock], + ) -> None: + _, workflow_runs = service_dependencies + service = _workflow_run_service_service(service_dependencies) + expected = {"total": 3, "succeeded": 2} + workflow_runs.get_workflow_runs_count.return_value = expected + + result = service.get_workflow_runs_count( + _workflow_run_service_request_context(), + app_id="app-1", + status="succeeded", + time_range="7d", + triggered_from=WorkflowRunTriggeredFrom.APP_RUN, + ) + + assert result == expected + workflow_runs.get_workflow_runs_count.assert_called_once_with( + tenant_id="tenant-1", + app_id="app-1", + triggered_from=WorkflowRunTriggeredFrom.APP_RUN, + status="succeeded", + time_range="7d", + ) + + def test_get_workflow_run_node_executions_should_return_empty_list_when_run_not_found( + self, + service_dependencies: tuple[MagicMock, MagicMock], + monkeypatch: pytest.MonkeyPatch, + ) -> None: + service = _workflow_run_service_service(service_dependencies) + monkeypatch.setattr(service, "get_workflow_run", MagicMock(return_value=None)) + + result = service.get_workflow_run_node_executions( + _workflow_run_service_request_context(), + app_id="app-1", + run_id="run-1", + ) + + assert result == [] + + def test_get_workflow_run_node_executions_should_use_request_workspace( + self, + service_dependencies: tuple[MagicMock, MagicMock], + monkeypatch: pytest.MonkeyPatch, + ) -> None: + node_executions, _ = service_dependencies + service = _workflow_run_service_service(service_dependencies) + monkeypatch.setattr(service, "get_workflow_run", MagicMock(return_value=_workflow_run())) + expected_executions = [SimpleNamespace(id="exec-1"), SimpleNamespace(id="exec-2")] + expected_traces = [SimpleNamespace(id="exec-1:retry:1"), SimpleNamespace(id="exec-1")] + node_executions.get_executions_by_workflow_run.return_value = expected_executions + mock_assemble = MagicMock(return_value=expected_traces) + monkeypatch.setattr(service_module, "assemble_workflow_node_execution_traces", mock_assemble) + + result = service.get_workflow_run_node_executions( + _workflow_run_service_request_context(workspace_id="tenant-context"), + app_id="app-1", + run_id="run-1", + ) + + assert result == expected_traces + node_executions.get_executions_by_workflow_run.assert_called_once_with( + tenant_id="tenant-context", + app_id="app-1", + workflow_run_id="run-1", + ) + mock_assemble.assert_called_once_with(expected_executions, node_executions, session=ANY) + + +"""Tests for Console workflow pause details.""" + + +def _workflow_run_service_pause_service(workflow_runs: MagicMock) -> WorkflowRunService: + return WorkflowRunService( + workflow_runs=workflow_runs, + node_executions=MagicMock(), + ) + + +def _workflow_run_service_pause_request_context(*, workspace_id: str = "tenant-1") -> RequestContext: + return RequestContext( + request_id="request-1", + trace_id="trace-1", + account_id="account-1", + active_workspace_id=workspace_id, + ) + + +class TestWorkflowRunServicePauseCases: + @staticmethod + @pytest.fixture + def workflow_runs() -> MagicMock: + return MagicMock() + + @staticmethod + def test_get_pause_details_returns_none_when_run_is_not_found(workflow_runs: MagicMock) -> None: + workflow_runs.get_pause_record.return_value = None + + result = _workflow_run_service_pause_service(workflow_runs).get_pause_details( + _workflow_run_service_pause_request_context(), workflow_run_id="run-1" ) - assert result is expected - workflow_runs.get_paginated_workflow_runs.assert_called_once_with( - tenant_id="tenant-1", - app_id="app-1", - triggered_from=WorkflowRunTriggeredFrom.APP_RUN, - limit=7, - last_id="last-1", - status="succeeded", - ) - - def test_get_paginate_advanced_chat_workflow_runs_should_attach_message_fields_when_message_exists( - self, - service_dependencies: tuple[MagicMock, MagicMock], - monkeypatch: pytest.MonkeyPatch, - ) -> None: - _, workflow_runs = service_dependencies - service = _service(service_dependencies) - run_with_message = _workflow_run(status=WorkflowExecutionStatus.RUNNING) - run_without_message = _workflow_run(run_id="run-2") - pagination = SimpleNamespace(data=[run_with_message, run_without_message]) - monkeypatch.setattr(service, "get_paginate_workflow_runs", MagicMock(return_value=pagination)) - workflow_runs.get_message_refs.return_value = { - "run-1": WorkflowRunMessageRef(message_id="msg-1", conversation_id="conv-1") - } - - result = service.get_paginate_advanced_chat_workflow_runs( - _request_context(), - app_id="app-1", - args={"limit": "2"}, - ) - - assert result is pagination - assert len(result.data) == 2 - assert result.data[0].message_id == "msg-1" - assert result.data[0].conversation_id == "conv-1" - assert result.data[0].status == "running" - assert not hasattr(result.data[1], "message_id") - assert result.data[1].id == "run-2" - workflow_runs.get_message_refs.assert_called_once_with( - app_id="app-1", - workflow_run_ids=["run-1", "run-2"], - ) - - def test_get_workflow_run_should_delegate_to_repository_by_tenant_and_app( - self, - service_dependencies: tuple[MagicMock, MagicMock], - ) -> None: - _, workflow_runs = service_dependencies - service = _service(service_dependencies) - expected = _workflow_run() - workflow_runs.get_workflow_run_by_id.return_value = expected - - result = service.get_workflow_run(_request_context(), app_id="app-1", run_id="run-1") - - assert result is expected - workflow_runs.get_workflow_run_by_id.assert_called_once_with( - tenant_id="tenant-1", - app_id="app-1", - run_id="run-1", - ) - - def test_get_workflow_runs_count_should_forward_optional_filters( - self, - service_dependencies: tuple[MagicMock, MagicMock], - ) -> None: - _, workflow_runs = service_dependencies - service = _service(service_dependencies) - expected = {"total": 3, "succeeded": 2} - workflow_runs.get_workflow_runs_count.return_value = expected - - result = service.get_workflow_runs_count( - _request_context(), - app_id="app-1", - status="succeeded", - time_range="7d", - triggered_from=WorkflowRunTriggeredFrom.APP_RUN, - ) - - assert result == expected - workflow_runs.get_workflow_runs_count.assert_called_once_with( - tenant_id="tenant-1", - app_id="app-1", - triggered_from=WorkflowRunTriggeredFrom.APP_RUN, - status="succeeded", - time_range="7d", - ) - - def test_get_workflow_run_node_executions_should_return_empty_list_when_run_not_found( - self, - service_dependencies: tuple[MagicMock, MagicMock], - monkeypatch: pytest.MonkeyPatch, - ) -> None: - service = _service(service_dependencies) - monkeypatch.setattr(service, "get_workflow_run", MagicMock(return_value=None)) - - result = service.get_workflow_run_node_executions( - _request_context(), - app_id="app-1", - run_id="run-1", - ) - - assert result == [] - - def test_get_workflow_run_node_executions_should_use_request_workspace( - self, - service_dependencies: tuple[MagicMock, MagicMock], - monkeypatch: pytest.MonkeyPatch, - ) -> None: - node_executions, _ = service_dependencies - service = _service(service_dependencies) - monkeypatch.setattr(service, "get_workflow_run", MagicMock(return_value=_workflow_run())) - expected_executions = [SimpleNamespace(id="exec-1"), SimpleNamespace(id="exec-2")] - expected_traces = [SimpleNamespace(id="exec-1:retry:1"), SimpleNamespace(id="exec-1")] - node_executions.get_executions_by_workflow_run.return_value = expected_executions - mock_assemble = MagicMock(return_value=expected_traces) - monkeypatch.setattr(service_module, "assemble_workflow_node_execution_traces", mock_assemble) - - result = service.get_workflow_run_node_executions( - _request_context(workspace_id="tenant-context"), - app_id="app-1", - run_id="run-1", - ) - - assert result == expected_traces - node_executions.get_executions_by_workflow_run.assert_called_once_with( - tenant_id="tenant-context", - app_id="app-1", + assert result is None + workflow_runs.get_pause_record.assert_called_once_with( + workspace_id="tenant-1", workflow_run_id="run-1", ) - mock_assemble.assert_called_once_with(expected_executions, node_executions, session=ANY) + + @staticmethod + def test_get_pause_details_returns_empty_details_for_non_paused_run(workflow_runs: MagicMock) -> None: + workflow_runs.get_pause_record.return_value = WorkflowRunPauseRecord( + status=WorkflowExecutionStatus.SUCCEEDED, + paused_at=None, + reasons=(), + form_tokens={}, + ) + + result = _workflow_run_service_pause_service(workflow_runs).get_pause_details( + _workflow_run_service_pause_request_context(), workflow_run_id="run-1" + ) + + assert result == WorkflowRunPauseDetails(paused_at=None, paused_nodes=()) + + @staticmethod + def test_get_pause_details_maps_human_input_and_token(workflow_runs: MagicMock) -> None: + reason = HumanInputRequired( + form_id="form-1", + form_content="Approve?", + node_id="node-1", + node_title="Approval", + ) + paused_at = datetime(2026, 1, 2, 3, 4, 5) + workflow_runs.get_pause_record.return_value = WorkflowRunPauseRecord( + status=WorkflowExecutionStatus.PAUSED, + paused_at=paused_at, + reasons=(reason,), + form_tokens={"form-1": "form-token"}, + ) + + result = _workflow_run_service_pause_service(workflow_runs).get_pause_details( + _workflow_run_service_pause_request_context(workspace_id="tenant-context"), + workflow_run_id="run-1", + ) + + assert result == WorkflowRunPauseDetails( + paused_at=paused_at, + paused_nodes=( + WorkflowRunPausedNode( + node_id="node-1", + node_title="Approval", + form_id="form-1", + form_token="form-token", + ), + ), + ) + workflow_runs.get_pause_record.assert_called_once_with( + workspace_id="tenant-context", + workflow_run_id="run-1", + ) + + @staticmethod + def test_get_pause_details_rejects_unsupported_pause_reason(workflow_runs: MagicMock) -> None: + workflow_runs.get_pause_record.return_value = WorkflowRunPauseRecord( + status=WorkflowExecutionStatus.PAUSED, + paused_at=None, + reasons=(SchedulingPause(message="Waiting for external input"),), + form_tokens={}, + ) + + with pytest.raises(NotImplementedError, match="Pause details do not support SchedulingPause"): + _workflow_run_service_pause_service(workflow_runs).get_pause_details( + _workflow_run_service_pause_request_context(), workflow_run_id="run-1" + ) diff --git a/api/tests/unit_tests/services/test_workflow_run_service_pause.py b/api/tests/unit_tests/services/test_workflow_run_service_pause.py deleted file mode 100644 index 3552cbd7c9d..00000000000 --- a/api/tests/unit_tests/services/test_workflow_run_service_pause.py +++ /dev/null @@ -1,112 +0,0 @@ -"""Tests for Console workflow pause details.""" - -from datetime import datetime -from unittest.mock import MagicMock - -import pytest - -from core.workflow.nodes.human_input.pause_reason import HumanInputRequired -from graphon.entities.pause_reason import SchedulingPause -from graphon.enums import WorkflowExecutionStatus -from machinery.context import RequestContext -from repositories.sqlalchemy_api_workflow_run_repository import WorkflowRunPauseRecord -from services.workflow_run_service import ( - WorkflowRunPauseDetails, - WorkflowRunPausedNode, - WorkflowRunService, -) - - -@pytest.fixture -def workflow_runs() -> MagicMock: - return MagicMock() - - -def _service(workflow_runs: MagicMock) -> WorkflowRunService: - return WorkflowRunService( - workflow_runs=workflow_runs, - node_executions=MagicMock(), - ) - - -def _request_context(*, workspace_id: str = "tenant-1") -> RequestContext: - return RequestContext( - request_id="request-1", - trace_id="trace-1", - account_id="account-1", - active_workspace_id=workspace_id, - ) - - -def test_get_pause_details_returns_none_when_run_is_not_found(workflow_runs: MagicMock) -> None: - workflow_runs.get_pause_record.return_value = None - - result = _service(workflow_runs).get_pause_details(_request_context(), workflow_run_id="run-1") - - assert result is None - workflow_runs.get_pause_record.assert_called_once_with( - workspace_id="tenant-1", - workflow_run_id="run-1", - ) - - -def test_get_pause_details_returns_empty_details_for_non_paused_run(workflow_runs: MagicMock) -> None: - workflow_runs.get_pause_record.return_value = WorkflowRunPauseRecord( - status=WorkflowExecutionStatus.SUCCEEDED, - paused_at=None, - reasons=(), - form_tokens={}, - ) - - result = _service(workflow_runs).get_pause_details(_request_context(), workflow_run_id="run-1") - - assert result == WorkflowRunPauseDetails(paused_at=None, paused_nodes=()) - - -def test_get_pause_details_maps_human_input_and_token(workflow_runs: MagicMock) -> None: - reason = HumanInputRequired( - form_id="form-1", - form_content="Approve?", - node_id="node-1", - node_title="Approval", - ) - paused_at = datetime(2026, 1, 2, 3, 4, 5) - workflow_runs.get_pause_record.return_value = WorkflowRunPauseRecord( - status=WorkflowExecutionStatus.PAUSED, - paused_at=paused_at, - reasons=(reason,), - form_tokens={"form-1": "form-token"}, - ) - - result = _service(workflow_runs).get_pause_details( - _request_context(workspace_id="tenant-context"), - workflow_run_id="run-1", - ) - - assert result == WorkflowRunPauseDetails( - paused_at=paused_at, - paused_nodes=( - WorkflowRunPausedNode( - node_id="node-1", - node_title="Approval", - form_id="form-1", - form_token="form-token", - ), - ), - ) - workflow_runs.get_pause_record.assert_called_once_with( - workspace_id="tenant-context", - workflow_run_id="run-1", - ) - - -def test_get_pause_details_rejects_unsupported_pause_reason(workflow_runs: MagicMock) -> None: - workflow_runs.get_pause_record.return_value = WorkflowRunPauseRecord( - status=WorkflowExecutionStatus.PAUSED, - paused_at=None, - reasons=(SchedulingPause(message="Waiting for external input"),), - form_tokens={}, - ) - - with pytest.raises(NotImplementedError, match="Pause details do not support SchedulingPause"): - _service(workflow_runs).get_pause_details(_request_context(), workflow_run_id="run-1") diff --git a/api/tests/unit_tests/services/test_workflow_service.py b/api/tests/unit_tests/services/test_workflow_service.py index facd0b020a7..48e6bd04fb8 100644 --- a/api/tests/unit_tests/services/test_workflow_service.py +++ b/api/tests/unit_tests/services/test_workflow_service.py @@ -1,34 +1,35 @@ -""" -Unit tests for WorkflowService. - -This test suite covers: -- Workflow creation from template -- Workflow validation (graph and features structure) -- Draft/publish transitions -- Version management -- Execution triggering -""" - import json import uuid -from collections.abc import Callable, Generator +from collections.abc import Callable, Generator, Iterator from datetime import datetime, timedelta from types import SimpleNamespace from typing import Any, cast from unittest.mock import ANY, MagicMock, patch, sentinel +from uuid import uuid4 import pytest +from flask import Flask from pytest_mock import MockerFixture -from sqlalchemy import event, select +from sqlalchemy import Engine as _workflow_agent_preview_identity_Engine +from sqlalchemy import event, func, select from sqlalchemy.dialects import postgresql -from sqlalchemy.engine import Engine +from sqlalchemy.engine import Engine as _workflow_service_Engine from sqlalchemy.orm import Session, sessionmaker from core.repositories.celery_workflow_node_execution_repository import CeleryWorkflowNodeExecutionRepository from core.repositories.sqlalchemy_workflow_node_execution_repository import SQLAlchemyWorkflowNodeExecutionRepository +from core.workflow.human_input_adapter import ( + EmailDeliveryConfig, + EmailDeliveryMethod, + EmailRecipients, + ExternalRecipient, + MemberRecipient, +) from core.workflow.llm_environment_variable import LLMEnvironmentVariable from core.workflow.node_factory import DifyNodeFactory +from core.workflow.nodes.agent_v2 import DifyAgentNode from core.workflow.nodes.agent_v2.session_store import WorkflowAgentWorkspaceStore +from core.workflow.nodes.human_input.entities import HumanInputNodeData from core.workflow.system_variables import SystemVariableKey, get_system_text from enums import DeploymentEdition from extensions.ext_database import db @@ -50,7 +51,8 @@ from graphon.runtime import RuntimeState, VariablePool from graphon.variables import StringVariable from graphon.variables.input_entities import VariableEntityType from libs.datetime_utils import naive_utc_now -from models.account import Account +from models import Account as _workflow_human_input_delivery_Account +from models.account import Account as _workflow_service_Account from models.agent import ( Agent, AgentConfigSnapshot, @@ -67,6 +69,7 @@ from models.model import App, AppMode from models.tools import BuiltinToolProvider, WorkflowToolProvider from models.workflow import Workflow, WorkflowNodeExecutionModel, WorkflowNodeExecutionTriggeredFrom, WorkflowType from repositories.factory import DifyAPIRepositoryFactory +from services import workflow_service as workflow_service_module from services.agent.retirement_service import WorkflowAgentRetirementService from services.errors.app import IsDraftWorkflowError, TriggerNodeLimitExceededError, WorkflowHashNotEqualError from services.errors.workflow_service import DraftWorkflowDeletionError, WorkflowInUseError @@ -78,8 +81,21 @@ from services.workflow_service import ( _rebuild_single_file, _setup_variable_pool, ) +from tests.helpers.unit.config_override import apply_config_overrides +from tests.helpers.unit.model_factories import make_account, make_app from tests.workflow_test_utils import build_test_graph_init_params +""" +Unit tests for WorkflowService. + +This test suite covers: +- Workflow creation from template +- Workflow validation (graph and features structure) +- Draft/publish transitions +- Version management +- Execution triggering +""" + class TestWorkflowAssociatedDataFactory: """ @@ -157,9 +173,9 @@ class TestWorkflowAssociatedDataFactory: return workflow @staticmethod - def create_account(account_id: str = "user-123", **kwargs: Any) -> Account: + def create_account(account_id: str = "user-123", **kwargs: Any) -> _workflow_service_Account: """Create an Account model with specified attributes.""" - account = Account(name="Test User", email=f"{account_id}@example.com") + account = _workflow_service_Account(name="Test User", email=f"{account_id}@example.com") account.id = account_id for key, value in kwargs.items(): setattr(account, key, value) @@ -228,984 +244,958 @@ class TestWorkflowAssociatedDataFactory: return {"nodes": nodes, "edges": edges} -@pytest.mark.usefixtures("sqlite_session") -class TestWorkflowService: - @pytest.fixture(autouse=True) - def _community_edition(self, config_overrides: Callable[..., None]) -> None: - config_overrides(DEPLOYMENT_EDITION=DeploymentEdition.COMMUNITY) +def _env(variable_id: str, name: str) -> StringVariable: + return StringVariable(id=variable_id, name=name, value="v", selector=["env", name]) - """ - Comprehensive unit tests for WorkflowService methods. - This test suite covers: - - Workflow creation from template - - Workflow validation (graph and features) - - Draft/publish transitions - - Version management - - Workflow deletion and error handling - """ +class TestWorkflowServiceCases: + @pytest.mark.usefixtures("sqlite_session") + class TestWorkflowService: + @pytest.fixture(autouse=True) + def _community_edition(self, config_overrides: Callable[..., None]) -> None: + config_overrides(DEPLOYMENT_EDITION=DeploymentEdition.COMMUNITY) - @pytest.fixture - def workflow_service(self, sqlite_engine: Engine) -> WorkflowService: - """Create a WorkflowService whose repositories use the test SQLite engine.""" - return WorkflowService(sessionmaker(bind=sqlite_engine, expire_on_commit=False)) - - def test_get_tenant_app_maintainers_scopes_requested_apps( - self, workflow_service: WorkflowService, sqlite_session: Session - ) -> None: - sqlite_session.add_all( - [ - TestWorkflowAssociatedDataFactory.create_app( - app_id="app-1", tenant_id="tenant-1", maintainer="owner-1" - ), - TestWorkflowAssociatedDataFactory.create_app(app_id="app-2", tenant_id="tenant-1"), - TestWorkflowAssociatedDataFactory.create_app( - app_id="app-3", tenant_id="tenant-2", maintainer="owner-2" - ), - ] - ) - sqlite_session.commit() - - assert workflow_service.get_tenant_app_maintainers( - ["app-1", "app-2", "app-3", "missing"], "tenant-1", session=sqlite_session - ) == {"app-1": "owner-1", "app-2": None} - - # ==================== Workflow Existence Tests ==================== - # These tests verify the service can check if a draft workflow exists - - def test_is_workflow_exist_returns_true(self, workflow_service: WorkflowService, sqlite_session: Session): """ - Test is_workflow_exist returns True when draft workflow exists. + Comprehensive unit tests for WorkflowService methods. - Verifies that the service correctly identifies when an app has a draft workflow. - This is used to determine whether to create or update a workflow. + This test suite covers: + - Workflow creation from template + - Workflow validation (graph and features) + - Draft/publish transitions + - Version management + - Workflow deletion and error handling """ - app = TestWorkflowAssociatedDataFactory.create_app() - sqlite_session.add(TestWorkflowAssociatedDataFactory.create_workflow()) - sqlite_session.commit() - result = workflow_service.is_workflow_exist(app, session=sqlite_session) + @pytest.fixture + def workflow_service(self, sqlite_engine: _workflow_service_Engine) -> WorkflowService: + """Create a WorkflowService whose repositories use the test SQLite engine.""" + return WorkflowService(sessionmaker(bind=sqlite_engine, expire_on_commit=False)) - assert result is True + def test_get_tenant_app_maintainers_scopes_requested_apps( + self, workflow_service: WorkflowService, sqlite_session: Session + ) -> None: + sqlite_session.add_all( + [ + TestWorkflowAssociatedDataFactory.create_app( + app_id="app-1", tenant_id="tenant-1", maintainer="owner-1" + ), + TestWorkflowAssociatedDataFactory.create_app(app_id="app-2", tenant_id="tenant-1"), + TestWorkflowAssociatedDataFactory.create_app( + app_id="app-3", tenant_id="tenant-2", maintainer="owner-2" + ), + ] + ) + sqlite_session.commit() - def test_is_workflow_exist_returns_false(self, workflow_service: WorkflowService, sqlite_session: Session): - """Test is_workflow_exist returns False when no draft workflow exists.""" - app = TestWorkflowAssociatedDataFactory.create_app() + assert workflow_service.get_tenant_app_maintainers( + ["app-1", "app-2", "app-3", "missing"], "tenant-1", session=sqlite_session + ) == {"app-1": "owner-1", "app-2": None} - result = workflow_service.is_workflow_exist(app, session=sqlite_session) + # ==================== Workflow Existence Tests ==================== + # These tests verify the service can check if a draft workflow exists - assert result is False + def test_is_workflow_exist_returns_true(self, workflow_service: WorkflowService, sqlite_session: Session): + """ + Test is_workflow_exist returns True when draft workflow exists. - # ==================== Get Draft Workflow Tests ==================== - # These tests verify retrieval of draft workflows (version="draft") + Verifies that the service correctly identifies when an app has a draft workflow. + This is used to determine whether to create or update a workflow. + """ + app = TestWorkflowAssociatedDataFactory.create_app() + sqlite_session.add(TestWorkflowAssociatedDataFactory.create_workflow()) + sqlite_session.commit() - def test_get_draft_workflow_success(self, workflow_service: WorkflowService, sqlite_session: Session): - """ - Test get_draft_workflow returns draft workflow successfully. + result = workflow_service.is_workflow_exist(app, session=sqlite_session) - Draft workflows are the working copy that users edit before publishing. - Each app can have only one draft workflow at a time. - """ - app = TestWorkflowAssociatedDataFactory.create_app() - workflow = TestWorkflowAssociatedDataFactory.create_workflow() - sqlite_session.add(workflow) - sqlite_session.commit() + assert result is True - result = workflow_service.get_draft_workflow(app, session=sqlite_session) + def test_is_workflow_exist_returns_false(self, workflow_service: WorkflowService, sqlite_session: Session): + """Test is_workflow_exist returns False when no draft workflow exists.""" + app = TestWorkflowAssociatedDataFactory.create_app() - assert result is workflow + result = workflow_service.is_workflow_exist(app, session=sqlite_session) - def test_get_draft_workflow_uses_provided_session(self, workflow_service: WorkflowService, sqlite_session: Session): - """Test get_draft_workflow can reuse an injected SQLAlchemy session.""" - app = TestWorkflowAssociatedDataFactory.create_app() - workflow = TestWorkflowAssociatedDataFactory.create_workflow() - sqlite_session.add(workflow) - sqlite_session.commit() + assert result is False - result = workflow_service.get_draft_workflow(app, session=sqlite_session) + # ==================== Get Draft Workflow Tests ==================== + # These tests verify retrieval of draft workflows (version="draft") - assert result is workflow + def test_get_draft_workflow_success(self, workflow_service: WorkflowService, sqlite_session: Session): + """ + Test get_draft_workflow returns draft workflow successfully. - def test_get_draft_workflow_returns_none(self, workflow_service: WorkflowService, sqlite_session: Session): - """Test get_draft_workflow returns None when no draft exists.""" - app = TestWorkflowAssociatedDataFactory.create_app() + Draft workflows are the working copy that users edit before publishing. + Each app can have only one draft workflow at a time. + """ + app = TestWorkflowAssociatedDataFactory.create_app() + workflow = TestWorkflowAssociatedDataFactory.create_workflow() + sqlite_session.add(workflow) + sqlite_session.commit() - result = workflow_service.get_draft_workflow(app, session=sqlite_session) + result = workflow_service.get_draft_workflow(app, session=sqlite_session) - assert result is None + assert result is workflow - # ==================== Get Published Workflow Tests ==================== - # These tests verify retrieval of published workflows (versioned snapshots) - - def test_get_published_workflow_by_id_reuses_provided_session(self, sqlite_session: Session): - """Published lookup sees uncommitted rows through the caller's session.""" - app = TestWorkflowAssociatedDataFactory.create_app() - workflow_id = "workflow-123" - workflow = TestWorkflowAssociatedDataFactory.create_workflow(workflow_id=workflow_id, version="v1") - sqlite_session.add(workflow) - sqlite_session.flush() - - result = WorkflowService.get_published_workflow_by_id(app, workflow_id, session=sqlite_session) - - assert result is workflow - - def test_get_published_workflow_by_id_success(self, workflow_service: WorkflowService, sqlite_session: Session): - """Test get_published_workflow_by_id returns published workflow.""" - app = TestWorkflowAssociatedDataFactory.create_app() - workflow_id = "workflow-123" - workflow = TestWorkflowAssociatedDataFactory.create_workflow(workflow_id=workflow_id, version="v1") - sqlite_session.add(workflow) - sqlite_session.commit() - - result = workflow_service.get_published_workflow_by_id(app, workflow_id, session=sqlite_session) - - assert result is workflow - - def test_get_published_workflow_by_id_can_lock_restore_source( - self, workflow_service: WorkflowService, sqlite_session: Session, mocker: MockerFixture - ): - app = TestWorkflowAssociatedDataFactory.create_app() - workflow = TestWorkflowAssociatedDataFactory.create_workflow(version="v1") - session = sqlite_session - session.add(workflow) - session.commit() - lookup = mocker.spy(session, "scalar") - - result = workflow_service.get_published_workflow_by_id( - app, - workflow.id, - session=session, - for_update=True, - ) - - stmt = lookup.call_args.args[0] - sql = str(stmt.compile(dialect=postgresql.dialect())) - assert result is workflow - assert "FOR UPDATE" in sql - - def test_get_published_workflow_by_id_raises_error_for_draft( - self, workflow_service: WorkflowService, sqlite_session: Session - ): - """ - Test get_published_workflow_by_id raises error when workflow is draft. - - This prevents using draft workflows in production contexts where only - published, stable versions should be used (e.g., API execution). - """ - app = TestWorkflowAssociatedDataFactory.create_app() - workflow_id = "workflow-123" - workflow = TestWorkflowAssociatedDataFactory.create_workflow( - workflow_id=workflow_id, version=Workflow.VERSION_DRAFT - ) - sqlite_session.add(workflow) - sqlite_session.commit() - - with pytest.raises(IsDraftWorkflowError): - workflow_service.get_published_workflow_by_id(app, workflow_id, session=sqlite_session) - - def test_get_published_workflow_by_id_returns_none( - self, workflow_service: WorkflowService, sqlite_session: Session - ): - """Test get_published_workflow_by_id returns None when workflow not found.""" - app = TestWorkflowAssociatedDataFactory.create_app() - workflow_id = "nonexistent-workflow" - - result = workflow_service.get_published_workflow_by_id(app, workflow_id, session=sqlite_session) - - assert result is None - - @pytest.mark.parametrize( - ("tenant_id", "app_id"), - [("other-tenant", "app-123"), ("tenant-456", "other-app")], - ) - def test_get_published_workflow_by_id_rejects_foreign_workflow( - self, - tenant_id: str, - app_id: str, - workflow_service: WorkflowService, - sqlite_session: Session, - ): - app = TestWorkflowAssociatedDataFactory.create_app() - workflow = TestWorkflowAssociatedDataFactory.create_workflow( - workflow_id="workflow-123", - tenant_id=tenant_id, - app_id=app_id, - version="v1", - ) - sqlite_session.add(workflow) - sqlite_session.commit() - - result = workflow_service.get_published_workflow_by_id(app, workflow.id, session=sqlite_session) - - assert result is None - - def test_get_published_workflow_success(self, workflow_service: WorkflowService, sqlite_session: Session): - """Test get_published_workflow returns published workflow.""" - workflow_id = "workflow-123" - app = TestWorkflowAssociatedDataFactory.create_app(workflow_id=workflow_id) - workflow = TestWorkflowAssociatedDataFactory.create_workflow(workflow_id=workflow_id, version="v1") - sqlite_session.add(workflow) - sqlite_session.commit() - - result = workflow_service.get_published_workflow(app, session=sqlite_session) - - assert result is workflow - - def test_get_published_workflow_returns_none_when_no_workflow_id( - self, workflow_service: WorkflowService, sqlite_session: Session - ): - """Test get_published_workflow returns None when app has no workflow_id.""" - app = TestWorkflowAssociatedDataFactory.create_app(workflow_id=None) - - result = workflow_service.get_published_workflow(app, session=sqlite_session) - - assert result is None - - # ==================== Sync Draft Workflow Tests ==================== - # These tests verify creating and updating draft workflows with validation - - def test_sync_draft_workflow_creates_new_draft(self, workflow_service: WorkflowService, sqlite_session: Session): - """ - Test sync_draft_workflow creates new draft workflow when none exists. - - When a user first creates a workflow app, this creates the initial draft. - The draft is validated before creation to ensure graph and features are valid. - """ - app = TestWorkflowAssociatedDataFactory.create_app() - account = TestWorkflowAssociatedDataFactory.create_account() - graph = TestWorkflowAssociatedDataFactory.create_valid_workflow_graph() - features = {"file_upload": {"enabled": False}} - - with ( - patch("services.workflow_service.app_draft_workflow_was_synced"), - patch( - "services.agent.workflow_publish_service.WorkflowAgentPublishService.sync_agent_bindings_for_draft", - return_value={"retired-agent"}, - ), - patch("services.workflow_service.WorkflowAgentRetirementService.retire_unowned") as retire_unowned, + def test_get_draft_workflow_uses_provided_session( + self, workflow_service: WorkflowService, sqlite_session: Session ): - result = workflow_service.sync_draft_workflow( - app_model=app, - graph=graph, - features=features, - unique_hash=None, - account=account, - environment_variables=[], - conversation_variables=[], - session=sqlite_session, + """Test get_draft_workflow can reuse an injected SQLAlchemy session.""" + app = TestWorkflowAssociatedDataFactory.create_app() + workflow = TestWorkflowAssociatedDataFactory.create_workflow() + sqlite_session.add(workflow) + sqlite_session.commit() + + result = workflow_service.get_draft_workflow(app, session=sqlite_session) + + assert result is workflow + + def test_get_draft_workflow_returns_none(self, workflow_service: WorkflowService, sqlite_session: Session): + """Test get_draft_workflow returns None when no draft exists.""" + app = TestWorkflowAssociatedDataFactory.create_app() + + result = workflow_service.get_draft_workflow(app, session=sqlite_session) + + assert result is None + + # ==================== Get Published Workflow Tests ==================== + # These tests verify retrieval of published workflows (versioned snapshots) + + def test_get_published_workflow_by_id_reuses_provided_session(self, sqlite_session: Session): + """Published lookup sees uncommitted rows through the caller's session.""" + app = TestWorkflowAssociatedDataFactory.create_app() + workflow_id = "workflow-123" + workflow = TestWorkflowAssociatedDataFactory.create_workflow(workflow_id=workflow_id, version="v1") + sqlite_session.add(workflow) + sqlite_session.flush() + + result = WorkflowService.get_published_workflow_by_id(app, workflow_id, session=sqlite_session) + + assert result is workflow + + def test_get_published_workflow_by_id_success(self, workflow_service: WorkflowService, sqlite_session: Session): + """Test get_published_workflow_by_id returns published workflow.""" + app = TestWorkflowAssociatedDataFactory.create_app() + workflow_id = "workflow-123" + workflow = TestWorkflowAssociatedDataFactory.create_workflow(workflow_id=workflow_id, version="v1") + sqlite_session.add(workflow) + sqlite_session.commit() + + result = workflow_service.get_published_workflow_by_id(app, workflow_id, session=sqlite_session) + + assert result is workflow + + def test_get_published_workflow_by_id_can_lock_restore_source( + self, workflow_service: WorkflowService, sqlite_session: Session, mocker: MockerFixture + ): + app = TestWorkflowAssociatedDataFactory.create_app() + workflow = TestWorkflowAssociatedDataFactory.create_workflow(version="v1") + session = sqlite_session + session.add(workflow) + session.commit() + lookup = mocker.spy(session, "scalar") + + result = workflow_service.get_published_workflow_by_id( + app, + workflow.id, + session=session, + for_update=True, ) - persisted_workflow = sqlite_session.scalar(select(Workflow).where(Workflow.id == result.id)) - assert persisted_workflow is result - assert result.graph_dict == graph - assert result.features_dict == features - retire_unowned.assert_called_once_with( - tenant_id=app.tenant_id, - agent_ids={"retired-agent"}, - account_id=account.id, + stmt = lookup.call_args.args[0] + sql = str(stmt.compile(dialect=postgresql.dialect())) + assert result is workflow + assert "FOR UPDATE" in sql + + def test_get_published_workflow_by_id_raises_error_for_draft( + self, workflow_service: WorkflowService, sqlite_session: Session + ): + """ + Test get_published_workflow_by_id raises error when workflow is draft. + + This prevents using draft workflows in production contexts where only + published, stable versions should be used (e.g., API execution). + """ + app = TestWorkflowAssociatedDataFactory.create_app() + workflow_id = "workflow-123" + workflow = TestWorkflowAssociatedDataFactory.create_workflow( + workflow_id=workflow_id, version=Workflow.VERSION_DRAFT + ) + sqlite_session.add(workflow) + sqlite_session.commit() + + with pytest.raises(IsDraftWorkflowError): + workflow_service.get_published_workflow_by_id(app, workflow_id, session=sqlite_session) + + def test_get_published_workflow_by_id_returns_none( + self, workflow_service: WorkflowService, sqlite_session: Session + ): + """Test get_published_workflow_by_id returns None when workflow not found.""" + app = TestWorkflowAssociatedDataFactory.create_app() + workflow_id = "nonexistent-workflow" + + result = workflow_service.get_published_workflow_by_id(app, workflow_id, session=sqlite_session) + + assert result is None + + @pytest.mark.parametrize( + ("tenant_id", "app_id"), + [("other-tenant", "app-123"), ("tenant-456", "other-app")], ) + def test_get_published_workflow_by_id_rejects_foreign_workflow( + self, + tenant_id: str, + app_id: str, + workflow_service: WorkflowService, + sqlite_session: Session, + ): + app = TestWorkflowAssociatedDataFactory.create_app() + workflow = TestWorkflowAssociatedDataFactory.create_workflow( + workflow_id="workflow-123", + tenant_id=tenant_id, + app_id=app_id, + version="v1", + ) + sqlite_session.add(workflow) + sqlite_session.commit() - def test_sync_draft_workflow_updates_existing_draft( - self, workflow_service: WorkflowService, sqlite_session: Session - ): - """ - Test sync_draft_workflow updates existing draft workflow. + result = workflow_service.get_published_workflow_by_id(app, workflow.id, session=sqlite_session) - When users edit their workflow, this updates the existing draft. - The unique_hash is used for optimistic locking to prevent conflicts. - """ - app = TestWorkflowAssociatedDataFactory.create_app() - account = TestWorkflowAssociatedDataFactory.create_account() - graph = TestWorkflowAssociatedDataFactory.create_valid_workflow_graph() - features = {"file_upload": {"enabled": False}} - workflow = TestWorkflowAssociatedDataFactory.create_workflow() - sqlite_session.add(workflow) - sqlite_session.commit() - unique_hash = workflow.unique_hash + assert result is None - with patch("services.workflow_service.app_draft_workflow_was_synced"): - result = workflow_service.sync_draft_workflow( - app_model=app, - graph=graph, - features=features, - unique_hash=unique_hash, - account=account, - environment_variables=[], - conversation_variables=[], - session=sqlite_session, + def test_get_published_workflow_success(self, workflow_service: WorkflowService, sqlite_session: Session): + """Test get_published_workflow returns published workflow.""" + workflow_id = "workflow-123" + app = TestWorkflowAssociatedDataFactory.create_app(workflow_id=workflow_id) + workflow = TestWorkflowAssociatedDataFactory.create_workflow(workflow_id=workflow_id, version="v1") + sqlite_session.add(workflow) + sqlite_session.commit() + + result = workflow_service.get_published_workflow(app, session=sqlite_session) + + assert result is workflow + + def test_get_published_workflow_returns_none_when_no_workflow_id( + self, workflow_service: WorkflowService, sqlite_session: Session + ): + """Test get_published_workflow returns None when app has no workflow_id.""" + app = TestWorkflowAssociatedDataFactory.create_app(workflow_id=None) + + result = workflow_service.get_published_workflow(app, session=sqlite_session) + + assert result is None + + # ==================== Sync Draft Workflow Tests ==================== + # These tests verify creating and updating draft workflows with validation + + def test_sync_draft_workflow_creates_new_draft( + self, workflow_service: WorkflowService, sqlite_session: Session + ): + """ + Test sync_draft_workflow creates new draft workflow when none exists. + + When a user first creates a workflow app, this creates the initial draft. + The draft is validated before creation to ensure graph and features are valid. + """ + app = TestWorkflowAssociatedDataFactory.create_app() + account = TestWorkflowAssociatedDataFactory.create_account() + graph = TestWorkflowAssociatedDataFactory.create_valid_workflow_graph() + features = {"file_upload": {"enabled": False}} + + with ( + patch("services.workflow_service.app_draft_workflow_was_synced"), + patch( + "services.agent.workflow_publish_service.WorkflowAgentPublishService.sync_agent_bindings_for_draft", + return_value={"retired-agent"}, + ), + patch("services.workflow_service.WorkflowAgentRetirementService.retire_unowned") as retire_unowned, + ): + result = workflow_service.sync_draft_workflow( + app_model=app, + graph=graph, + features=features, + unique_hash=None, + account=account, + environment_variables=[], + conversation_variables=[], + session=sqlite_session, + ) + + persisted_workflow = sqlite_session.scalar(select(Workflow).where(Workflow.id == result.id)) + assert persisted_workflow is result + assert result.graph_dict == graph + assert result.features_dict == features + retire_unowned.assert_called_once_with( + tenant_id=app.tenant_id, + agent_ids={"retired-agent"}, + account_id=account.id, ) - sqlite_session.refresh(workflow) - assert result is workflow - assert workflow.graph_dict == graph - assert workflow.features_dict == features - assert workflow.updated_by == account.id + def test_sync_draft_workflow_updates_existing_draft( + self, workflow_service: WorkflowService, sqlite_session: Session + ): + """ + Test sync_draft_workflow updates existing draft workflow. - def test_sync_draft_workflow_collaborative_save_preserves_environment_variables_and_locks_row( - self, workflow_service: WorkflowService, sqlite_session: Session - ) -> None: - """A collaborative graph-only save locks the draft and keeps server environment values.""" - app = TestWorkflowAssociatedDataFactory.create_app() - account = TestWorkflowAssociatedDataFactory.create_account() - workflow = TestWorkflowAssociatedDataFactory.create_workflow() - remote_variable = StringVariable( - id="env-shared", - name="shared", - value="remote-value", - selector=["env", "shared"], - ) - stale_variable = remote_variable.model_copy(update={"value": "stale-client-value"}) - workflow.environment_variables = [remote_variable] - sqlite_session.add(workflow) - sqlite_session.commit() - unique_hash = workflow.unique_hash - statements = [] - commits = [] + When users edit their workflow, this updates the existing draft. + The unique_hash is used for optimistic locking to prevent conflicts. + """ + app = TestWorkflowAssociatedDataFactory.create_app() + account = TestWorkflowAssociatedDataFactory.create_account() + graph = TestWorkflowAssociatedDataFactory.create_valid_workflow_graph() + features = {"file_upload": {"enabled": False}} + workflow = TestWorkflowAssociatedDataFactory.create_workflow() + sqlite_session.add(workflow) + sqlite_session.commit() + unique_hash = workflow.unique_hash - @event.listens_for(sqlite_session, "do_orm_execute") - def capture_statement(execute_state): - statements.append(execute_state.statement) + with patch("services.workflow_service.app_draft_workflow_was_synced"): + result = workflow_service.sync_draft_workflow( + app_model=app, + graph=graph, + features=features, + unique_hash=unique_hash, + account=account, + environment_variables=[], + conversation_variables=[], + session=sqlite_session, + ) - @event.listens_for(sqlite_session, "before_commit") - def capture_commit(_session): - commits.append(True) + sqlite_session.refresh(workflow) + assert result is workflow + assert workflow.graph_dict == graph + assert workflow.features_dict == features + assert workflow.updated_by == account.id - result = workflow_service.sync_draft_workflow( - app_model=app, - graph=TestWorkflowAssociatedDataFactory.create_valid_workflow_graph(), - features={}, - unique_hash=unique_hash, - account=account, - environment_variables=[stale_variable], - conversation_variables=[], - session=sqlite_session, - preserve_environment_variables=True, - commit=False, - sync_agent_bindings=False, - ) - - assert "FOR UPDATE" in str(statements[0].compile(dialect=postgresql.dialect())) - assert result.environment_variables == [remote_variable] - assert commits == [] - - def test_sync_draft_workflow_merges_environment_patch_with_graph_update( - self, workflow_service: WorkflowService, sqlite_session: Session - ) -> None: - """Graph changes and a per-ID environment patch commit without replacing untouched aliases.""" - app = TestWorkflowAssociatedDataFactory.create_app() - account = TestWorkflowAssociatedDataFactory.create_account() - workflow = TestWorkflowAssociatedDataFactory.create_workflow() - remote_variable = StringVariable(id="env-a", name="a", value="remote-a", selector=["env", "a"]) - replaced_variable = StringVariable(id="env-b", name="b", value="old-b", selector=["env", "b"]) - workflow.environment_variables = [remote_variable, replaced_variable] - sqlite_session.add(workflow) - sqlite_session.commit() - unique_hash = workflow.unique_hash - next_graph = TestWorkflowAssociatedDataFactory.create_valid_workflow_graph() - next_graph["viewport"] = {"x": 10, "y": 20, "zoom": 1} - - with patch("services.workflow_service.app_draft_workflow_was_synced"): - result = workflow_service.sync_draft_workflow( - app_model=app, - graph=next_graph, - features={}, - unique_hash=unique_hash, - account=account, - environment_variables=[ - remote_variable.model_copy(update={"value": "stale-a"}), - replaced_variable, - ], - conversation_variables=[], - session=sqlite_session, - environment_variable_upserts=[replaced_variable.model_copy(update={"value": "new-b"})], - deleted_environment_variable_ids=[], - preserve_environment_variables=True, + def test_sync_draft_workflow_collaborative_save_preserves_environment_variables_and_locks_row( + self, workflow_service: WorkflowService, sqlite_session: Session + ) -> None: + """A collaborative graph-only save locks the draft and keeps server environment values.""" + app = TestWorkflowAssociatedDataFactory.create_app() + account = TestWorkflowAssociatedDataFactory.create_account() + workflow = TestWorkflowAssociatedDataFactory.create_workflow() + remote_variable = StringVariable( + id="env-shared", + name="shared", + value="remote-value", + selector=["env", "shared"], ) + stale_variable = remote_variable.model_copy(update={"value": "stale-client-value"}) + workflow.environment_variables = [remote_variable] + sqlite_session.add(workflow) + sqlite_session.commit() + unique_hash = workflow.unique_hash + statements = [] + commits = [] - assert result.graph_dict == next_graph - assert [(variable.id, variable.value) for variable in result.environment_variables] == [ - ("env-a", "remote-a"), - ("env-b", "new-b"), - ] + @event.listens_for(sqlite_session, "do_orm_execute") + def capture_statement(execute_state): + statements.append(execute_state.statement) - def test_sync_draft_workflow_noncollaborative_save_replaces_environment_variables( - self, workflow_service: WorkflowService, sqlite_session: Session - ) -> None: - """Legacy full-draft saves retain their full environment replacement contract.""" - app = TestWorkflowAssociatedDataFactory.create_app() - account = TestWorkflowAssociatedDataFactory.create_account() - workflow = TestWorkflowAssociatedDataFactory.create_workflow() - workflow.environment_variables = [ - StringVariable(id="env-old", name="old", value="old", selector=["env", "old"]) - ] - sqlite_session.add(workflow) - sqlite_session.commit() - replacement = StringVariable(id="env-new", name="new", value="new", selector=["env", "new"]) + @event.listens_for(sqlite_session, "before_commit") + def capture_commit(_session): + commits.append(True) - with patch("services.workflow_service.app_draft_workflow_was_synced"): result = workflow_service.sync_draft_workflow( app_model=app, graph=TestWorkflowAssociatedDataFactory.create_valid_workflow_graph(), features={}, - unique_hash=workflow.unique_hash, - account=account, - environment_variables=[replacement], - conversation_variables=[], - session=sqlite_session, - ) - - assert result.environment_variables == [replacement] - - def test_sync_draft_workflow_graph_only_preserves_independently_updated_draft_fields( - self, workflow_service: WorkflowService, sqlite_session: Session - ): - app = TestWorkflowAssociatedDataFactory.create_app() - account = TestWorkflowAssociatedDataFactory.create_account() - existing_features = {"file_upload": {"enabled": True}} - existing_environment_variables = [ - StringVariable(id="env-new", name="region", value="new", selector=["env", "region"]) - ] - existing_conversation_variables = [ - StringVariable( - id="conversation-new", - name="topic", - value="new", - selector=["conversation", "topic"], - ) - ] - workflow = TestWorkflowAssociatedDataFactory.create_workflow( - features=existing_features, - environment_variables=existing_environment_variables, - conversation_variables=existing_conversation_variables, - ) - sqlite_session.add(workflow) - sqlite_session.commit() - unique_hash = workflow.unique_hash - updated_graph = TestWorkflowAssociatedDataFactory.create_valid_workflow_graph() - updated_graph["nodes"][0]["data"]["title"] = "Updated graph" - - with patch("services.workflow_service.app_draft_workflow_was_synced"): - result = workflow_service.sync_draft_workflow( - app_model=app, - graph=updated_graph, - features={"file_upload": {"enabled": False}}, unique_hash=unique_hash, account=account, - environment_variables=[ - StringVariable(id="env-old", name="region", value="old", selector=["env", "region"]) - ], - conversation_variables=[ - StringVariable( - id="conversation-old", - name="topic", - value="old", - selector=["conversation", "topic"], - ) - ], - session=sqlite_session, - graph_only=True, - ) - - sqlite_session.refresh(workflow) - assert result is workflow - assert workflow.graph_dict == updated_graph - assert workflow.features_dict == existing_features - assert workflow.environment_variables == existing_environment_variables - assert workflow.conversation_variables == existing_conversation_variables - - def test_sync_draft_workflow_graph_only_creates_complete_initial_draft( - self, workflow_service: WorkflowService, sqlite_session: Session - ): - app = TestWorkflowAssociatedDataFactory.create_app() - account = TestWorkflowAssociatedDataFactory.create_account() - graph = TestWorkflowAssociatedDataFactory.create_valid_workflow_graph() - features = {"file_upload": {"enabled": True}} - environment_variables = [StringVariable(id="env-id", name="region", value="new", selector=["env", "region"])] - conversation_variables = [ - StringVariable( - id="conversation-id", - name="topic", - value="new", - selector=["conversation", "topic"], - ) - ] - - with patch("services.workflow_service.app_draft_workflow_was_synced"): - workflow = workflow_service.sync_draft_workflow( - app_model=app, - graph=graph, - features=features, - unique_hash=None, - account=account, - environment_variables=environment_variables, - conversation_variables=conversation_variables, - session=sqlite_session, - graph_only=True, - ) - - assert workflow.graph_dict == graph - assert workflow.features_dict == features - assert workflow.environment_variables == environment_variables - assert workflow.conversation_variables == conversation_variables - - def test_sync_draft_workflow_raises_hash_not_equal_error( - self, workflow_service: WorkflowService, sqlite_session: Session - ): - """ - Test sync_draft_workflow raises error when hash doesn't match. - - This implements optimistic locking: if the workflow was modified by another - user/session since it was loaded, the hash won't match and the update fails. - This prevents overwriting concurrent changes. - """ - app = TestWorkflowAssociatedDataFactory.create_app() - account = TestWorkflowAssociatedDataFactory.create_account() - graph = TestWorkflowAssociatedDataFactory.create_valid_workflow_graph() - features = {} - - sqlite_session.add(TestWorkflowAssociatedDataFactory.create_workflow()) - sqlite_session.commit() - - with pytest.raises(WorkflowHashNotEqualError): - workflow_service.sync_draft_workflow( - app_model=app, - graph=graph, - features=features, - unique_hash="new-hash", - account=account, - environment_variables=[], + environment_variables=[stale_variable], conversation_variables=[], session=sqlite_session, + preserve_environment_variables=True, + commit=False, + sync_agent_bindings=False, ) - def test_restore_published_workflow_to_draft_keeps_source_features_unmodified( - self, workflow_service: WorkflowService, sqlite_session: Session - ): - app = TestWorkflowAssociatedDataFactory.create_app() - account = TestWorkflowAssociatedDataFactory.create_account() - legacy_features = { - "file_upload": { - "image": { - "enabled": True, - "number_limits": 6, - "transfer_methods": ["remote_url", "local_file"], - } - }, - "opening_statement": "", - "retriever_resource": {"enabled": True}, - "sensitive_word_avoidance": {"enabled": False}, - "speech_to_text": {"enabled": False}, - "suggested_questions": [], - "suggested_questions_after_answer": {"enabled": False}, - "text_to_speech": {"enabled": False, "language": "", "voice": ""}, - } - source_workflow = Workflow( - id="published-workflow-id", - tenant_id=app.tenant_id, - app_id=app.id, - type=WorkflowType.WORKFLOW, - version="2026-03-19T00:00:00", - graph=json.dumps(TestWorkflowAssociatedDataFactory.create_valid_workflow_graph()), - features=json.dumps(legacy_features), - created_by=account.id, - environment_variables=[], - conversation_variables=[], - rag_pipeline_variables=[], - ) - draft_workflow = Workflow( - id="draft-workflow-id", - tenant_id=app.tenant_id, - app_id=app.id, - type=WorkflowType.WORKFLOW, - version=Workflow.VERSION_DRAFT, - graph=json.dumps({"nodes": [], "edges": []}), - features=json.dumps({}), - created_by=account.id, - environment_variables=[], - conversation_variables=[], - rag_pipeline_variables=[], - ) - sqlite_session.add_all([source_workflow, draft_workflow]) - sqlite_session.commit() + assert "FOR UPDATE" in str(statements[0].compile(dialect=postgresql.dialect())) + assert result.environment_variables == [remote_variable] + assert commits == [] - with ( - patch("services.workflow_service.app_draft_workflow_was_synced"), - patch.object( - workflow_service, - "get_published_workflow_by_id", - wraps=workflow_service.get_published_workflow_by_id, - ) as get_published_workflow_by_id, - patch( - "services.agent.workflow_publish_service.WorkflowAgentPublishService.restore_agent_node_bindings_to_draft", - return_value={"retired-agent"}, - ), - patch("services.workflow_service.WorkflowAgentRetirementService.retire_unowned") as retire_unowned, + def test_sync_draft_workflow_merges_environment_patch_with_graph_update( + self, workflow_service: WorkflowService, sqlite_session: Session + ) -> None: + """Graph changes and a per-ID environment patch commit without replacing untouched aliases.""" + app = TestWorkflowAssociatedDataFactory.create_app() + account = TestWorkflowAssociatedDataFactory.create_account() + workflow = TestWorkflowAssociatedDataFactory.create_workflow() + remote_variable = StringVariable(id="env-a", name="a", value="remote-a", selector=["env", "a"]) + replaced_variable = StringVariable(id="env-b", name="b", value="old-b", selector=["env", "b"]) + workflow.environment_variables = [remote_variable, replaced_variable] + sqlite_session.add(workflow) + sqlite_session.commit() + unique_hash = workflow.unique_hash + next_graph = TestWorkflowAssociatedDataFactory.create_valid_workflow_graph() + next_graph["viewport"] = {"x": 10, "y": 20, "zoom": 1} + + with patch("services.workflow_service.app_draft_workflow_was_synced"): + result = workflow_service.sync_draft_workflow( + app_model=app, + graph=next_graph, + features={}, + unique_hash=unique_hash, + account=account, + environment_variables=[ + remote_variable.model_copy(update={"value": "stale-a"}), + replaced_variable, + ], + conversation_variables=[], + session=sqlite_session, + environment_variable_upserts=[replaced_variable.model_copy(update={"value": "new-b"})], + deleted_environment_variable_ids=[], + preserve_environment_variables=True, + ) + + assert result.graph_dict == next_graph + assert [(variable.id, variable.value) for variable in result.environment_variables] == [ + ("env-a", "remote-a"), + ("env-b", "new-b"), + ] + + def test_sync_draft_workflow_noncollaborative_save_replaces_environment_variables( + self, workflow_service: WorkflowService, sqlite_session: Session + ) -> None: + """Legacy full-draft saves retain their full environment replacement contract.""" + app = TestWorkflowAssociatedDataFactory.create_app() + account = TestWorkflowAssociatedDataFactory.create_account() + workflow = TestWorkflowAssociatedDataFactory.create_workflow() + workflow.environment_variables = [ + StringVariable(id="env-old", name="old", value="old", selector=["env", "old"]) + ] + sqlite_session.add(workflow) + sqlite_session.commit() + replacement = StringVariable(id="env-new", name="new", value="new", selector=["env", "new"]) + + with patch("services.workflow_service.app_draft_workflow_was_synced"): + result = workflow_service.sync_draft_workflow( + app_model=app, + graph=TestWorkflowAssociatedDataFactory.create_valid_workflow_graph(), + features={}, + unique_hash=workflow.unique_hash, + account=account, + environment_variables=[replacement], + conversation_variables=[], + session=sqlite_session, + ) + + assert result.environment_variables == [replacement] + + def test_sync_draft_workflow_graph_only_preserves_independently_updated_draft_fields( + self, workflow_service: WorkflowService, sqlite_session: Session ): - result = workflow_service.restore_published_workflow_to_draft( + app = TestWorkflowAssociatedDataFactory.create_app() + account = TestWorkflowAssociatedDataFactory.create_account() + existing_features = {"file_upload": {"enabled": True}} + existing_environment_variables = [ + StringVariable(id="env-new", name="region", value="new", selector=["env", "region"]) + ] + existing_conversation_variables = [ + StringVariable( + id="conversation-new", + name="topic", + value="new", + selector=["conversation", "topic"], + ) + ] + workflow = TestWorkflowAssociatedDataFactory.create_workflow( + features=existing_features, + environment_variables=existing_environment_variables, + conversation_variables=existing_conversation_variables, + ) + sqlite_session.add(workflow) + sqlite_session.commit() + unique_hash = workflow.unique_hash + updated_graph = TestWorkflowAssociatedDataFactory.create_valid_workflow_graph() + updated_graph["nodes"][0]["data"]["title"] = "Updated graph" + + with patch("services.workflow_service.app_draft_workflow_was_synced"): + result = workflow_service.sync_draft_workflow( + app_model=app, + graph=updated_graph, + features={"file_upload": {"enabled": False}}, + unique_hash=unique_hash, + account=account, + environment_variables=[ + StringVariable(id="env-old", name="region", value="old", selector=["env", "region"]) + ], + conversation_variables=[ + StringVariable( + id="conversation-old", + name="topic", + value="old", + selector=["conversation", "topic"], + ) + ], + session=sqlite_session, + graph_only=True, + ) + + sqlite_session.refresh(workflow) + assert result is workflow + assert workflow.graph_dict == updated_graph + assert workflow.features_dict == existing_features + assert workflow.environment_variables == existing_environment_variables + assert workflow.conversation_variables == existing_conversation_variables + + def test_sync_draft_workflow_graph_only_creates_complete_initial_draft( + self, workflow_service: WorkflowService, sqlite_session: Session + ): + app = TestWorkflowAssociatedDataFactory.create_app() + account = TestWorkflowAssociatedDataFactory.create_account() + graph = TestWorkflowAssociatedDataFactory.create_valid_workflow_graph() + features = {"file_upload": {"enabled": True}} + environment_variables = [ + StringVariable(id="env-id", name="region", value="new", selector=["env", "region"]) + ] + conversation_variables = [ + StringVariable( + id="conversation-id", + name="topic", + value="new", + selector=["conversation", "topic"], + ) + ] + + with patch("services.workflow_service.app_draft_workflow_was_synced"): + workflow = workflow_service.sync_draft_workflow( + app_model=app, + graph=graph, + features=features, + unique_hash=None, + account=account, + environment_variables=environment_variables, + conversation_variables=conversation_variables, + session=sqlite_session, + graph_only=True, + ) + + assert workflow.graph_dict == graph + assert workflow.features_dict == features + assert workflow.environment_variables == environment_variables + assert workflow.conversation_variables == conversation_variables + + def test_sync_draft_workflow_raises_hash_not_equal_error( + self, workflow_service: WorkflowService, sqlite_session: Session + ): + """ + Test sync_draft_workflow raises error when hash doesn't match. + + This implements optimistic locking: if the workflow was modified by another + user/session since it was loaded, the hash won't match and the update fails. + This prevents overwriting concurrent changes. + """ + app = TestWorkflowAssociatedDataFactory.create_app() + account = TestWorkflowAssociatedDataFactory.create_account() + graph = TestWorkflowAssociatedDataFactory.create_valid_workflow_graph() + features = {} + + sqlite_session.add(TestWorkflowAssociatedDataFactory.create_workflow()) + sqlite_session.commit() + + with pytest.raises(WorkflowHashNotEqualError): + workflow_service.sync_draft_workflow( + app_model=app, + graph=graph, + features=features, + unique_hash="new-hash", + account=account, + environment_variables=[], + conversation_variables=[], + session=sqlite_session, + ) + + def test_restore_published_workflow_to_draft_keeps_source_features_unmodified( + self, workflow_service: WorkflowService, sqlite_session: Session + ): + app = TestWorkflowAssociatedDataFactory.create_app() + account = TestWorkflowAssociatedDataFactory.create_account() + legacy_features = { + "file_upload": { + "image": { + "enabled": True, + "number_limits": 6, + "transfer_methods": ["remote_url", "local_file"], + } + }, + "opening_statement": "", + "retriever_resource": {"enabled": True}, + "sensitive_word_avoidance": {"enabled": False}, + "speech_to_text": {"enabled": False}, + "suggested_questions": [], + "suggested_questions_after_answer": {"enabled": False}, + "text_to_speech": {"enabled": False, "language": "", "voice": ""}, + } + source_workflow = Workflow( + id="published-workflow-id", + tenant_id=app.tenant_id, + app_id=app.id, + type=WorkflowType.WORKFLOW, + version="2026-03-19T00:00:00", + graph=json.dumps(TestWorkflowAssociatedDataFactory.create_valid_workflow_graph()), + features=json.dumps(legacy_features), + created_by=account.id, + environment_variables=[], + conversation_variables=[], + rag_pipeline_variables=[], + ) + draft_workflow = Workflow( + id="draft-workflow-id", + tenant_id=app.tenant_id, + app_id=app.id, + type=WorkflowType.WORKFLOW, + version=Workflow.VERSION_DRAFT, + graph=json.dumps({"nodes": [], "edges": []}), + features=json.dumps({}), + created_by=account.id, + environment_variables=[], + conversation_variables=[], + rag_pipeline_variables=[], + ) + sqlite_session.add_all([source_workflow, draft_workflow]) + sqlite_session.commit() + + with ( + patch("services.workflow_service.app_draft_workflow_was_synced"), + patch.object( + workflow_service, + "get_published_workflow_by_id", + wraps=workflow_service.get_published_workflow_by_id, + ) as get_published_workflow_by_id, + patch( + "services.agent.workflow_publish_service.WorkflowAgentPublishService.restore_agent_node_bindings_to_draft", + return_value={"retired-agent"}, + ), + patch("services.workflow_service.WorkflowAgentRetirementService.retire_unowned") as retire_unowned, + ): + result = workflow_service.restore_published_workflow_to_draft( + app_model=app, + workflow_id=source_workflow.id, + account=account, + session=sqlite_session, + ) + + assert result is draft_workflow + assert source_workflow.serialized_features == json.dumps(legacy_features) + assert draft_workflow.serialized_features == json.dumps(legacy_features) + sqlite_session.refresh(draft_workflow) + assert draft_workflow.serialized_features == json.dumps(legacy_features) + get_published_workflow_by_id.assert_called_once_with( app_model=app, workflow_id=source_workflow.id, - account=account, session=sqlite_session, + for_update=True, + ) + retire_unowned.assert_called_once_with( + tenant_id=app.tenant_id, + agent_ids={"retired-agent"}, + account_id=account.id, ) - assert result is draft_workflow - assert source_workflow.serialized_features == json.dumps(legacy_features) - assert draft_workflow.serialized_features == json.dumps(legacy_features) - sqlite_session.refresh(draft_workflow) - assert draft_workflow.serialized_features == json.dumps(legacy_features) - get_published_workflow_by_id.assert_called_once_with( - app_model=app, - workflow_id=source_workflow.id, - session=sqlite_session, - for_update=True, - ) - retire_unowned.assert_called_once_with( - tenant_id=app.tenant_id, - agent_ids={"retired-agent"}, - account_id=account.id, - ) - - def test_restore_historical_inline_agent_after_current_pointer_moves_uses_real_clone( - self, - workflow_service: WorkflowService, - sqlite_session: Session, - ) -> None: - app = TestWorkflowAssociatedDataFactory.create_app(workflow_id=None) - account = TestWorkflowAssociatedDataFactory.create_account() - graph = { - "nodes": [ - { - "id": "agent-node", - "data": { - "type": "agent", - "version": "2", - "agent_node_kind": "dify_agent", - }, - } - ], - "edges": [], - } - historical = TestWorkflowAssociatedDataFactory.create_workflow( - workflow_id="historical-workflow", - version="historical-version", - graph=graph, - ) - current = TestWorkflowAssociatedDataFactory.create_workflow( - workflow_id="current-workflow", - version="current-version", - graph=graph, - ) - draft = TestWorkflowAssociatedDataFactory.create_workflow( - workflow_id="draft-workflow", - version=Workflow.VERSION_DRAFT, - ) - source_agent = Agent( - id="historical-agent", - tenant_id=app.tenant_id, - name="Historical inline Agent", - description="", - role="", - agent_kind=AgentKind.DIFY_AGENT, - scope=AgentScope.WORKFLOW_ONLY, - source=AgentSource.WORKFLOW, - app_id=app.id, - workflow_id=historical.id, - workflow_node_id="agent-node", - active_config_snapshot_id="historical-snapshot", - active_config_has_model=False, - active_config_is_published=True, - status=AgentStatus.ACTIVE, - created_by=account.id, - updated_by=account.id, - ) - source_snapshot = AgentConfigSnapshot( - id="historical-snapshot", - tenant_id=app.tenant_id, - agent_id=source_agent.id, - version=1, - config_snapshot=AgentSoulConfig(config_note="historical soul"), - created_by=account.id, - ) - historical_binding = WorkflowAgentNodeBinding( - id="historical-binding", - tenant_id=app.tenant_id, - app_id=app.id, - workflow_id=historical.id, - workflow_version=historical.version, - node_id="agent-node", - binding_type=WorkflowAgentBindingType.INLINE_AGENT, - agent_id=source_agent.id, - current_snapshot_id=source_snapshot.id, - node_job_config={}, - created_by=account.id, - ) - sqlite_session.add_all([app, historical, current, draft, source_agent, source_snapshot, historical_binding]) - app.workflow_id = historical.id - sqlite_session.commit() - - app.workflow_id = current.id - sqlite_session.commit() - - WorkflowAgentRetirementService.retire_unowned( - tenant_id=app.tenant_id, - agent_ids=[source_agent.id], - account_id=account.id, - ) - sqlite_session.expire_all() - retained_agent = sqlite_session.get(Agent, source_agent.id) - assert retained_agent is not None - assert retained_agent.status is AgentStatus.ACTIVE - - with patch("services.workflow_service.app_draft_workflow_was_synced"): - restored_draft = workflow_service.restore_published_workflow_to_draft( - app_model=app, + def test_restore_historical_inline_agent_after_current_pointer_moves_uses_real_clone( + self, + workflow_service: WorkflowService, + sqlite_session: Session, + ) -> None: + app = TestWorkflowAssociatedDataFactory.create_app(workflow_id=None) + account = TestWorkflowAssociatedDataFactory.create_account() + graph = { + "nodes": [ + { + "id": "agent-node", + "data": { + "type": "agent", + "version": "2", + "agent_node_kind": "dify_agent", + }, + } + ], + "edges": [], + } + historical = TestWorkflowAssociatedDataFactory.create_workflow( + workflow_id="historical-workflow", + version="historical-version", + graph=graph, + ) + current = TestWorkflowAssociatedDataFactory.create_workflow( + workflow_id="current-workflow", + version="current-version", + graph=graph, + ) + draft = TestWorkflowAssociatedDataFactory.create_workflow( + workflow_id="draft-workflow", + version=Workflow.VERSION_DRAFT, + ) + source_agent = Agent( + id="historical-agent", + tenant_id=app.tenant_id, + name="Historical inline Agent", + description="", + role="", + agent_kind=AgentKind.DIFY_AGENT, + scope=AgentScope.WORKFLOW_ONLY, + source=AgentSource.WORKFLOW, + app_id=app.id, workflow_id=historical.id, - account=account, - session=sqlite_session, + workflow_node_id="agent-node", + active_config_snapshot_id="historical-snapshot", + active_config_has_model=False, + active_config_is_published=True, + status=AgentStatus.ACTIVE, + created_by=account.id, + updated_by=account.id, ) - - restored_binding = sqlite_session.scalar( - select(WorkflowAgentNodeBinding).where( - WorkflowAgentNodeBinding.workflow_id == draft.id, - WorkflowAgentNodeBinding.workflow_version == Workflow.VERSION_DRAFT, - WorkflowAgentNodeBinding.node_id == "agent-node", + source_snapshot = AgentConfigSnapshot( + id="historical-snapshot", + tenant_id=app.tenant_id, + agent_id=source_agent.id, + version=1, + config_snapshot=AgentSoulConfig(config_note="historical soul"), + created_by=account.id, ) - ) - assert restored_draft is draft - assert app.workflow_id == current.id - assert sqlite_session.get(Workflow, historical.id) is historical - assert sqlite_session.get(WorkflowAgentNodeBinding, historical_binding.id) is historical_binding - assert restored_binding is not None - assert restored_binding.agent_id not in (None, source_agent.id) - assert restored_binding.current_snapshot_id not in (None, source_snapshot.id) - restored_agent = sqlite_session.get(Agent, restored_binding.agent_id) - restored_snapshot = sqlite_session.get(AgentConfigSnapshot, restored_binding.current_snapshot_id) - assert restored_agent is not None - assert restored_agent.workflow_id == draft.id - assert restored_agent.workflow_node_id == "agent-node" - assert restored_snapshot is not None - assert restored_snapshot.config_snapshot_dict == source_snapshot.config_snapshot_dict + historical_binding = WorkflowAgentNodeBinding( + id="historical-binding", + tenant_id=app.tenant_id, + app_id=app.id, + workflow_id=historical.id, + workflow_version=historical.version, + node_id="agent-node", + binding_type=WorkflowAgentBindingType.INLINE_AGENT, + agent_id=source_agent.id, + current_snapshot_id=source_snapshot.id, + node_job_config={}, + created_by=account.id, + ) + sqlite_session.add_all([app, historical, current, draft, source_agent, source_snapshot, historical_binding]) + app.workflow_id = historical.id + sqlite_session.commit() - # ==================== Workflow Validation Tests ==================== - # These tests verify graph structure and feature configuration validation + app.workflow_id = current.id + sqlite_session.commit() - def test_validate_graph_structure_empty_graph(self, workflow_service: WorkflowService): - """Test validate_graph_structure accepts empty graph.""" - graph = {"nodes": []} + WorkflowAgentRetirementService.retire_unowned( + tenant_id=app.tenant_id, + agent_ids=[source_agent.id], + account_id=account.id, + ) + sqlite_session.expire_all() + retained_agent = sqlite_session.get(Agent, source_agent.id) + assert retained_agent is not None + assert retained_agent.status is AgentStatus.ACTIVE - # Should not raise any exception - workflow_service.validate_graph_structure(graph) + with patch("services.workflow_service.app_draft_workflow_was_synced"): + restored_draft = workflow_service.restore_published_workflow_to_draft( + app_model=app, + workflow_id=historical.id, + account=account, + session=sqlite_session, + ) - def test_validate_graph_structure_valid_graph(self, workflow_service: WorkflowService): - """Test validate_graph_structure accepts valid graph.""" - graph = TestWorkflowAssociatedDataFactory.create_valid_workflow_graph() + restored_binding = sqlite_session.scalar( + select(WorkflowAgentNodeBinding).where( + WorkflowAgentNodeBinding.workflow_id == draft.id, + WorkflowAgentNodeBinding.workflow_version == Workflow.VERSION_DRAFT, + WorkflowAgentNodeBinding.node_id == "agent-node", + ) + ) + assert restored_draft is draft + assert app.workflow_id == current.id + assert sqlite_session.get(Workflow, historical.id) is historical + assert sqlite_session.get(WorkflowAgentNodeBinding, historical_binding.id) is historical_binding + assert restored_binding is not None + assert restored_binding.agent_id not in (None, source_agent.id) + assert restored_binding.current_snapshot_id not in (None, source_snapshot.id) + restored_agent = sqlite_session.get(Agent, restored_binding.agent_id) + restored_snapshot = sqlite_session.get(AgentConfigSnapshot, restored_binding.current_snapshot_id) + assert restored_agent is not None + assert restored_agent.workflow_id == draft.id + assert restored_agent.workflow_node_id == "agent-node" + assert restored_snapshot is not None + assert restored_snapshot.config_snapshot_dict == source_snapshot.config_snapshot_dict - # Should not raise any exception - workflow_service.validate_graph_structure(graph) + # ==================== Workflow Validation Tests ==================== + # These tests verify graph structure and feature configuration validation - def test_validate_graph_structure_start_and_trigger_coexist_raises_error(self, workflow_service: WorkflowService): - """ - Test validate_graph_structure raises error when start and trigger nodes coexist. + def test_validate_graph_structure_empty_graph(self, workflow_service: WorkflowService): + """Test validate_graph_structure accepts empty graph.""" + graph = {"nodes": []} - Workflows can be either: - - User-initiated (with START node): User provides input to start execution - - Event-driven (with trigger nodes): External events trigger execution - - These two patterns cannot be mixed in a single workflow. - """ - # Create a graph with both start and trigger nodes - # Use actual trigger node types: trigger-webhook, trigger-schedule, trigger-plugin - graph = { - "nodes": [ - { - "id": "start", - "data": { - "type": "start", - "title": "START", - }, - }, - { - "id": "trigger-1", - "data": { - "type": "trigger-webhook", - "title": "Webhook Trigger", - }, - }, - ], - "edges": [], - } - - with pytest.raises(ValueError, match="Start node and trigger nodes cannot coexist"): + # Should not raise any exception workflow_service.validate_graph_structure(graph) - def test_validate_features_structure_workflow_mode(self, workflow_service: WorkflowService): - """ - Test validate_features_structure for workflow mode. + def test_validate_graph_structure_valid_graph(self, workflow_service: WorkflowService): + """Test validate_graph_structure accepts valid graph.""" + graph = TestWorkflowAssociatedDataFactory.create_valid_workflow_graph() - Different app modes have different feature configurations. - This ensures the features match the expected schema for workflow apps. - """ - app = TestWorkflowAssociatedDataFactory.create_app(mode=AppMode.WORKFLOW) - features = {"file_upload": {"enabled": False}} + # Should not raise any exception + workflow_service.validate_graph_structure(graph) - with patch("services.workflow_service.WorkflowAppConfigManager.config_validate") as mock_validate: - workflow_service.validate_features_structure(app, features) - mock_validate.assert_called_once_with( - tenant_id=app.tenant_id, config=features, only_structure_validate=True - ) + def test_validate_graph_structure_start_and_trigger_coexist_raises_error( + self, workflow_service: WorkflowService + ): + """ + Test validate_graph_structure raises error when start and trigger nodes coexist. - def test_validate_features_structure_advanced_chat_mode(self, workflow_service: WorkflowService): - """Test validate_features_structure for advanced chat mode.""" - app = TestWorkflowAssociatedDataFactory.create_app(mode=AppMode.ADVANCED_CHAT) - features = {"opening_statement": "Hello"} + Workflows can be either: + - User-initiated (with START node): User provides input to start execution + - Event-driven (with trigger nodes): External events trigger execution - with patch("services.workflow_service.AdvancedChatAppConfigManager.config_validate") as mock_validate: - workflow_service.validate_features_structure(app, features) - mock_validate.assert_called_once_with( - tenant_id=app.tenant_id, config=features, only_structure_validate=True - ) + These two patterns cannot be mixed in a single workflow. + """ + # Create a graph with both start and trigger nodes + # Use actual trigger node types: trigger-webhook, trigger-schedule, trigger-plugin + graph = { + "nodes": [ + { + "id": "start", + "data": { + "type": "start", + "title": "START", + }, + }, + { + "id": "trigger-1", + "data": { + "type": "trigger-webhook", + "title": "Webhook Trigger", + }, + }, + ], + "edges": [], + } - def test_validate_features_structure_invalid_mode_raises_error(self, workflow_service: WorkflowService): - """Test validate_features_structure raises error for invalid mode.""" - app = TestWorkflowAssociatedDataFactory.create_app(mode=AppMode.COMPLETION) - features = {} + with pytest.raises(ValueError, match="Start node and trigger nodes cannot coexist"): + workflow_service.validate_graph_structure(graph) - with pytest.raises(ValueError, match="Invalid app mode"): - workflow_service.validate_features_structure(app, features) + def test_validate_features_structure_workflow_mode(self, workflow_service: WorkflowService): + """ + Test validate_features_structure for workflow mode. - # ==================== Draft Workflow Variable Update Tests ==================== - # These tests verify updating draft workflow environment/conversation variables + Different app modes have different feature configurations. + This ensures the features match the expected schema for workflow apps. + """ + app = TestWorkflowAssociatedDataFactory.create_app(mode=AppMode.WORKFLOW) + features = {"file_upload": {"enabled": False}} - def test_update_draft_workflow_environment_variables_updates_workflow( - self, workflow_service: WorkflowService, sqlite_session: Session - ): - """Test update_draft_workflow_environment_variables updates draft fields.""" - app = TestWorkflowAssociatedDataFactory.create_app() - account = TestWorkflowAssociatedDataFactory.create_account() - workflow = TestWorkflowAssociatedDataFactory.create_workflow() - sqlite_session.add(workflow) - sqlite_session.commit() - variables = [StringVariable(id="env-id", name="region", value="us-east-1", selector=["env", "region"])] + with patch("services.workflow_service.WorkflowAppConfigManager.config_validate") as mock_validate: + workflow_service.validate_features_structure(app, features) + mock_validate.assert_called_once_with( + tenant_id=app.tenant_id, config=features, only_structure_validate=True + ) - workflow_service.update_draft_workflow_environment_variables( - app_model=app, - environment_variables=variables, - account=account, - session=sqlite_session, - ) + def test_validate_features_structure_advanced_chat_mode(self, workflow_service: WorkflowService): + """Test validate_features_structure for advanced chat mode.""" + app = TestWorkflowAssociatedDataFactory.create_app(mode=AppMode.ADVANCED_CHAT) + features = {"opening_statement": "Hello"} - assert workflow.environment_variables == variables - assert workflow.updated_by == account.id - sqlite_session.expire_all() - persisted_workflow = sqlite_session.get(Workflow, workflow.id) - assert persisted_workflow is not None - assert persisted_workflow.environment_variables == variables + with patch("services.workflow_service.AdvancedChatAppConfigManager.config_validate") as mock_validate: + workflow_service.validate_features_structure(app, features) + mock_validate.assert_called_once_with( + tenant_id=app.tenant_id, config=features, only_structure_validate=True + ) - def test_update_draft_workflow_environment_variables_raises_when_missing( - self, workflow_service: WorkflowService, sqlite_session: Session - ): - """Test update_draft_workflow_environment_variables raises when draft missing.""" - app = TestWorkflowAssociatedDataFactory.create_app() - account = TestWorkflowAssociatedDataFactory.create_account() + def test_validate_features_structure_invalid_mode_raises_error(self, workflow_service: WorkflowService): + """Test validate_features_structure raises error for invalid mode.""" + app = TestWorkflowAssociatedDataFactory.create_app(mode=AppMode.COMPLETION) + features = {} + + with pytest.raises(ValueError, match="Invalid app mode"): + workflow_service.validate_features_structure(app, features) + + # ==================== Draft Workflow Variable Update Tests ==================== + # These tests verify updating draft workflow environment/conversation variables + + def test_update_draft_workflow_environment_variables_updates_workflow( + self, workflow_service: WorkflowService, sqlite_session: Session + ): + """Test update_draft_workflow_environment_variables updates draft fields.""" + app = TestWorkflowAssociatedDataFactory.create_app() + account = TestWorkflowAssociatedDataFactory.create_account() + workflow = TestWorkflowAssociatedDataFactory.create_workflow() + sqlite_session.add(workflow) + sqlite_session.commit() + variables = [StringVariable(id="env-id", name="region", value="us-east-1", selector=["env", "region"])] - with pytest.raises(ValueError, match="No draft workflow found."): workflow_service.update_draft_workflow_environment_variables( app_model=app, - environment_variables=[], + environment_variables=variables, account=account, session=sqlite_session, ) - def test_patch_draft_workflow_environment_variables_merges_by_id( - self, workflow_service: WorkflowService, sqlite_session: Session - ) -> None: - """A patch preserves untouched variables while applying ordered upserts and deletions.""" - app = TestWorkflowAssociatedDataFactory.create_app() - account = TestWorkflowAssociatedDataFactory.create_account() - workflow = TestWorkflowAssociatedDataFactory.create_workflow() - workflow.environment_variables = [ - StringVariable(id="env-a", name="a", value="remote-a", selector=["env", "a"]), - StringVariable(id="env-b", name="b", value="old-b", selector=["env", "b"]), - StringVariable(id="env-c", name="c", value="remote-c", selector=["env", "c"]), - ] - sqlite_session.add(workflow) - sqlite_session.commit() + assert workflow.environment_variables == variables + assert workflow.updated_by == account.id + sqlite_session.expire_all() + persisted_workflow = sqlite_session.get(Workflow, workflow.id) + assert persisted_workflow is not None + assert persisted_workflow.environment_variables == variables - workflow_service.patch_draft_workflow_environment_variables( - app_model=app, - environment_variables=[ - StringVariable(id="env-b", name="b", value="new-b", selector=["env", "b"]), - StringVariable(id="env-d", name="d", value="new-d", selector=["env", "d"]), - ], - deleted_environment_variable_ids=["env-a"], - account=account, - session=sqlite_session, - ) + def test_update_draft_workflow_environment_variables_raises_when_missing( + self, workflow_service: WorkflowService, sqlite_session: Session + ): + """Test update_draft_workflow_environment_variables raises when draft missing.""" + app = TestWorkflowAssociatedDataFactory.create_app() + account = TestWorkflowAssociatedDataFactory.create_account() - sqlite_session.expire_all() - persisted_workflow = sqlite_session.get(Workflow, workflow.id) - assert persisted_workflow is not None - assert [(variable.id, variable.value) for variable in persisted_workflow.environment_variables] == [ - ("env-b", "new-b"), - ("env-c", "remote-c"), - ("env-d", "new-d"), - ] - assert persisted_workflow.updated_by == account.id + with pytest.raises(ValueError, match="No draft workflow found."): + workflow_service.update_draft_workflow_environment_variables( + app_model=app, + environment_variables=[], + account=account, + session=sqlite_session, + ) - def test_patch_draft_workflow_environment_variables_locks_draft_row( - self, workflow_service: WorkflowService, sqlite_session: Session - ) -> None: - """The merge reads the draft with a row lock before applying a partial update.""" - app = TestWorkflowAssociatedDataFactory.create_app() - account = TestWorkflowAssociatedDataFactory.create_account() - workflow = TestWorkflowAssociatedDataFactory.create_workflow() - sqlite_session.add(workflow) - sqlite_session.commit() - statements = [] + def test_patch_draft_workflow_environment_variables_merges_by_id( + self, workflow_service: WorkflowService, sqlite_session: Session + ) -> None: + """A patch preserves untouched variables while applying ordered upserts and deletions.""" + app = TestWorkflowAssociatedDataFactory.create_app() + account = TestWorkflowAssociatedDataFactory.create_account() + workflow = TestWorkflowAssociatedDataFactory.create_workflow() + workflow.environment_variables = [ + StringVariable(id="env-a", name="a", value="remote-a", selector=["env", "a"]), + StringVariable(id="env-b", name="b", value="old-b", selector=["env", "b"]), + StringVariable(id="env-c", name="c", value="remote-c", selector=["env", "c"]), + ] + sqlite_session.add(workflow) + sqlite_session.commit() - @event.listens_for(sqlite_session, "do_orm_execute") - def capture_statement(execute_state): - statements.append(execute_state.statement) - - workflow_service.patch_draft_workflow_environment_variables( - app_model=app, - environment_variables=[], - deleted_environment_variable_ids=[], - account=account, - session=sqlite_session, - ) - - compiled_statement = str(statements[0].compile(dialect=postgresql.dialect())) - assert "FOR UPDATE" in compiled_statement - assert sqlite_session.get(Workflow, workflow.id) is workflow - - def test_patch_draft_workflow_environment_variables_rejects_conflicting_ids( - self, workflow_service: WorkflowService, sqlite_session: Session - ) -> None: - """A patch cannot both upsert and delete the same variable ID.""" - app = TestWorkflowAssociatedDataFactory.create_app() - account = TestWorkflowAssociatedDataFactory.create_account() - variable = StringVariable(id="env-a", name="a", value="a", selector=["env", "a"]) - sqlite_session.add(TestWorkflowAssociatedDataFactory.create_workflow()) - sqlite_session.commit() - - with pytest.raises(ValueError, match="cannot be upserted and deleted"): workflow_service.patch_draft_workflow_environment_variables( app_model=app, - environment_variables=[variable], + environment_variables=[ + StringVariable(id="env-b", name="b", value="new-b", selector=["env", "b"]), + StringVariable(id="env-d", name="d", value="new-d", selector=["env", "d"]), + ], deleted_environment_variable_ids=["env-a"], account=account, session=sqlite_session, ) - def test_patch_draft_workflow_environment_variables_raises_when_missing( - self, workflow_service: WorkflowService, sqlite_session: Session - ) -> None: - """A patch fails when the app has no draft workflow to lock.""" - app = TestWorkflowAssociatedDataFactory.create_app() - account = TestWorkflowAssociatedDataFactory.create_account() + sqlite_session.expire_all() + persisted_workflow = sqlite_session.get(Workflow, workflow.id) + assert persisted_workflow is not None + assert [(variable.id, variable.value) for variable in persisted_workflow.environment_variables] == [ + ("env-b", "new-b"), + ("env-c", "remote-c"), + ("env-d", "new-d"), + ] + assert persisted_workflow.updated_by == account.id + + def test_patch_draft_workflow_environment_variables_locks_draft_row( + self, workflow_service: WorkflowService, sqlite_session: Session + ) -> None: + """The merge reads the draft with a row lock before applying a partial update.""" + app = TestWorkflowAssociatedDataFactory.create_app() + account = TestWorkflowAssociatedDataFactory.create_account() + workflow = TestWorkflowAssociatedDataFactory.create_workflow() + sqlite_session.add(workflow) + sqlite_session.commit() + statements = [] + + @event.listens_for(sqlite_session, "do_orm_execute") + def capture_statement(execute_state): + statements.append(execute_state.statement) - with pytest.raises(ValueError, match="No draft workflow found."): workflow_service.patch_draft_workflow_environment_variables( app_model=app, environment_variables=[], @@ -1214,2692 +1204,3169 @@ class TestWorkflowService: session=sqlite_session, ) - def test_update_draft_workflow_conversation_variables_updates_workflow( - self, workflow_service: WorkflowService, sqlite_session: Session - ): - """Test update_draft_workflow_conversation_variables updates draft fields.""" - app = TestWorkflowAssociatedDataFactory.create_app() - account = TestWorkflowAssociatedDataFactory.create_account() - workflow = TestWorkflowAssociatedDataFactory.create_workflow() - sqlite_session.add(workflow) - sqlite_session.commit() - variables = [ - StringVariable(id="conversation-id", name="topic", value="sqlite", selector=["conversation", "topic"]) - ] + compiled_statement = str(statements[0].compile(dialect=postgresql.dialect())) + assert "FOR UPDATE" in compiled_statement + assert sqlite_session.get(Workflow, workflow.id) is workflow - workflow_service.update_draft_workflow_conversation_variables( - app_model=app, - conversation_variables=variables, - account=account, - session=sqlite_session, - ) + def test_patch_draft_workflow_environment_variables_rejects_conflicting_ids( + self, workflow_service: WorkflowService, sqlite_session: Session + ) -> None: + """A patch cannot both upsert and delete the same variable ID.""" + app = TestWorkflowAssociatedDataFactory.create_app() + account = TestWorkflowAssociatedDataFactory.create_account() + variable = StringVariable(id="env-a", name="a", value="a", selector=["env", "a"]) + sqlite_session.add(TestWorkflowAssociatedDataFactory.create_workflow()) + sqlite_session.commit() - assert workflow.conversation_variables == variables - assert workflow.updated_by == account.id - sqlite_session.expire_all() - persisted_workflow = sqlite_session.get(Workflow, workflow.id) - assert persisted_workflow is not None - assert persisted_workflow.conversation_variables == variables + with pytest.raises(ValueError, match="cannot be upserted and deleted"): + workflow_service.patch_draft_workflow_environment_variables( + app_model=app, + environment_variables=[variable], + deleted_environment_variable_ids=["env-a"], + account=account, + session=sqlite_session, + ) - def test_update_draft_workflow_conversation_variables_raises_when_missing( - self, workflow_service: WorkflowService, sqlite_session: Session - ): - """Test update_draft_workflow_conversation_variables raises when draft missing.""" - app = TestWorkflowAssociatedDataFactory.create_app() - account = TestWorkflowAssociatedDataFactory.create_account() + def test_patch_draft_workflow_environment_variables_raises_when_missing( + self, workflow_service: WorkflowService, sqlite_session: Session + ) -> None: + """A patch fails when the app has no draft workflow to lock.""" + app = TestWorkflowAssociatedDataFactory.create_app() + account = TestWorkflowAssociatedDataFactory.create_account() + + with pytest.raises(ValueError, match="No draft workflow found."): + workflow_service.patch_draft_workflow_environment_variables( + app_model=app, + environment_variables=[], + deleted_environment_variable_ids=[], + account=account, + session=sqlite_session, + ) + + def test_update_draft_workflow_conversation_variables_updates_workflow( + self, workflow_service: WorkflowService, sqlite_session: Session + ): + """Test update_draft_workflow_conversation_variables updates draft fields.""" + app = TestWorkflowAssociatedDataFactory.create_app() + account = TestWorkflowAssociatedDataFactory.create_account() + workflow = TestWorkflowAssociatedDataFactory.create_workflow() + sqlite_session.add(workflow) + sqlite_session.commit() + variables = [ + StringVariable(id="conversation-id", name="topic", value="sqlite", selector=["conversation", "topic"]) + ] - with pytest.raises(ValueError, match="No draft workflow found."): workflow_service.update_draft_workflow_conversation_variables( app_model=app, - conversation_variables=[], + conversation_variables=variables, account=account, session=sqlite_session, ) - # ==================== Publish Workflow Tests ==================== - # These tests verify creating published versions from draft workflows + assert workflow.conversation_variables == variables + assert workflow.updated_by == account.id + sqlite_session.expire_all() + persisted_workflow = sqlite_session.get(Workflow, workflow.id) + assert persisted_workflow is not None + assert persisted_workflow.conversation_variables == variables - def test_publish_workflow_success(self, workflow_service: WorkflowService, sqlite_session: Session): - """ - Test publish_workflow creates new published version. - - Publishing creates a timestamped snapshot of the draft workflow. - This allows users to: - - Roll back to previous versions - - Use stable versions in production - - Continue editing draft without affecting published version - """ - app = TestWorkflowAssociatedDataFactory.create_app() - account = TestWorkflowAssociatedDataFactory.create_account() - graph = TestWorkflowAssociatedDataFactory.create_valid_workflow_graph() - - draft = TestWorkflowAssociatedDataFactory.create_workflow(version=Workflow.VERSION_DRAFT, graph=graph) - sqlite_session.add(draft) - sqlite_session.commit() - - with patch("services.workflow_service.app_published_workflow_was_updated"): - result = workflow_service.publish_workflow( - session=sqlite_session, - app_model=app, - account=account, - marked_name="Version 1", - marked_comment="Initial release", - ) - - sqlite_session.flush() - assert result in sqlite_session - assert result.version != Workflow.VERSION_DRAFT - assert result.marked_name == "Version 1" - assert result.marked_comment == "Initial release" - - def test_publish_workflow_numbers_versions_from_one( - self, workflow_service: WorkflowService, sqlite_session: Session - ): - """ - Test publish_workflow assigns an app-scoped version number starting at #1. - - The number is what users see when a version carries no name, so it has to be - stable and monotonic per app rather than derived from list position. - """ - app = TestWorkflowAssociatedDataFactory.create_app() - account = TestWorkflowAssociatedDataFactory.create_account() - graph = TestWorkflowAssociatedDataFactory.create_valid_workflow_graph() - - draft = TestWorkflowAssociatedDataFactory.create_workflow(version=Workflow.VERSION_DRAFT, graph=graph) - sqlite_session.add(draft) - sqlite_session.commit() - - with ( - patch("services.workflow_service.app_published_workflow_was_updated"), + def test_update_draft_workflow_conversation_variables_raises_when_missing( + self, workflow_service: WorkflowService, sqlite_session: Session ): - first = workflow_service.publish_workflow(session=sqlite_session, app_model=app, account=account) - second = workflow_service.publish_workflow(session=sqlite_session, app_model=app, account=account) + """Test update_draft_workflow_conversation_variables raises when draft missing.""" + app = TestWorkflowAssociatedDataFactory.create_app() + account = TestWorkflowAssociatedDataFactory.create_account() - assert first.version_number == 1 - assert second.version_number == 2 + with pytest.raises(ValueError, match="No draft workflow found."): + workflow_service.update_draft_workflow_conversation_variables( + app_model=app, + conversation_variables=[], + account=account, + session=sqlite_session, + ) - def test_publish_workflow_does_not_reuse_a_deleted_version_number( - self, workflow_service: WorkflowService, sqlite_session: Session - ): - """ - Test version numbers are never handed out twice, even after a version is deleted. + # ==================== Publish Workflow Tests ==================== + # These tests verify creating published versions from draft workflows - Deployment records and audit logs refer to versions by number, so reusing one - would make two different workflows share an identity. - """ - app = TestWorkflowAssociatedDataFactory.create_app() - account = TestWorkflowAssociatedDataFactory.create_account() - graph = TestWorkflowAssociatedDataFactory.create_valid_workflow_graph() + def test_publish_workflow_success(self, workflow_service: WorkflowService, sqlite_session: Session): + """ + Test publish_workflow creates new published version. - draft = TestWorkflowAssociatedDataFactory.create_workflow(version=Workflow.VERSION_DRAFT, graph=graph) - sqlite_session.add(draft) - sqlite_session.commit() + Publishing creates a timestamped snapshot of the draft workflow. + This allows users to: + - Roll back to previous versions + - Use stable versions in production + - Continue editing draft without affecting published version + """ + app = TestWorkflowAssociatedDataFactory.create_app() + account = TestWorkflowAssociatedDataFactory.create_account() + graph = TestWorkflowAssociatedDataFactory.create_valid_workflow_graph() - with ( - patch("services.workflow_service.app_published_workflow_was_updated"), + draft = TestWorkflowAssociatedDataFactory.create_workflow(version=Workflow.VERSION_DRAFT, graph=graph) + sqlite_session.add(draft) + sqlite_session.commit() + + with patch("services.workflow_service.app_published_workflow_was_updated"): + result = workflow_service.publish_workflow( + session=sqlite_session, + app_model=app, + account=account, + marked_name="Version 1", + marked_comment="Initial release", + ) + + sqlite_session.flush() + assert result in sqlite_session + assert result.version != Workflow.VERSION_DRAFT + assert result.marked_name == "Version 1" + assert result.marked_comment == "Initial release" + + def test_publish_workflow_numbers_versions_from_one( + self, workflow_service: WorkflowService, sqlite_session: Session ): - published = workflow_service.publish_workflow(session=sqlite_session, app_model=app, account=account) - sqlite_session.flush() - sqlite_session.delete(published) - sqlite_session.flush() + """ + Test publish_workflow assigns an app-scoped version number starting at #1. - republished = workflow_service.publish_workflow(session=sqlite_session, app_model=app, account=account) + The number is what users see when a version carries no name, so it has to be + stable and monotonic per app rather than derived from list position. + """ + app = TestWorkflowAssociatedDataFactory.create_app() + account = TestWorkflowAssociatedDataFactory.create_account() + graph = TestWorkflowAssociatedDataFactory.create_valid_workflow_graph() - assert republished.version_number == 2 - - def test_publish_workflow_numbers_each_app_independently( - self, workflow_service: WorkflowService, sqlite_session: Session - ): - """ - Test the counter is scoped per app rather than global. - - Every app starts its own sequence at #1; a busy neighbour must not advance it. - """ - account = TestWorkflowAssociatedDataFactory.create_account() - graph = TestWorkflowAssociatedDataFactory.create_valid_workflow_graph() - - published: list[Workflow] = [] - for app_id in ("app-first", "app-second"): - app = TestWorkflowAssociatedDataFactory.create_app(app_id=app_id) - draft = TestWorkflowAssociatedDataFactory.create_workflow( - workflow_id=f"draft-{app_id}", - app_id=app_id, - version=Workflow.VERSION_DRAFT, - graph=graph, - ) + draft = TestWorkflowAssociatedDataFactory.create_workflow(version=Workflow.VERSION_DRAFT, graph=graph) sqlite_session.add(draft) sqlite_session.commit() with ( patch("services.workflow_service.app_published_workflow_was_updated"), ): - workflow = workflow_service.publish_workflow(session=sqlite_session, app_model=app, account=account) - published.append(workflow) + first = workflow_service.publish_workflow(session=sqlite_session, app_model=app, account=account) + second = workflow_service.publish_workflow(session=sqlite_session, app_model=app, account=account) - assert [workflow.version_number for workflow in published] == [1, 1] + assert first.version_number == 1 + assert second.version_number == 2 - def test_publish_workflow_no_draft_raises_error(self, workflow_service: WorkflowService, sqlite_session: Session): - """ - Test publish_workflow raises error when no draft exists. - - Cannot publish if there's no draft to publish from. - Users must create and save a draft before publishing. - """ - app = TestWorkflowAssociatedDataFactory.create_app() - account = TestWorkflowAssociatedDataFactory.create_account() - - with pytest.raises(ValueError, match="No valid workflow found"): - workflow_service.publish_workflow(session=sqlite_session, app_model=app, account=account) - - @pytest.mark.parametrize( - ("environment_variables", "error"), - [ - ([], "was not found or is not an LLM variable"), - ( - [ - LLMEnvironmentVariable( - name="shared_model", - value={"provider": "provider", "name": "model", "mode": "completion"}, - ) - ], - "uses mode 'completion'.*uses mode 'chat'", - ), - ], - ) - def test_publish_workflow_rejects_invalid_llm_environment_reference_without_credential_validation( - self, - workflow_service: WorkflowService, - sqlite_session: Session, - environment_variables: list[LLMEnvironmentVariable], - error: str, - ) -> None: - app = TestWorkflowAssociatedDataFactory.create_app() - account = TestWorkflowAssociatedDataFactory.create_account() - graph = TestWorkflowAssociatedDataFactory.create_valid_workflow_graph() - llm_node = next(node for node in graph["nodes"] if node["data"]["type"] == BuiltinNodeTypes.LLM) - llm_node["data"]["model"] = { - "provider": "provider", - "name": "model", - "mode": "chat", - "completion_params": {}, - } - llm_node["data"]["model_selector"] = ["env", "shared_model"] - draft = TestWorkflowAssociatedDataFactory.create_workflow(graph=graph) - draft.environment_variables = environment_variables - sqlite_session.add(draft) - sqlite_session.commit() - - with ( - patch( - "services.system_feature_service.SystemFeatureService.is_plugin_manager_enabled", - return_value=False, - ), - pytest.raises(ValueError, match=error), + def test_publish_workflow_does_not_reuse_a_deleted_version_number( + self, workflow_service: WorkflowService, sqlite_session: Session ): - workflow_service.publish_workflow(session=sqlite_session, app_model=app, account=account) + """ + Test version numbers are never handed out twice, even after a version is deleted. - def test_publish_workflow_trigger_limit_exceeded( - self, - workflow_service: WorkflowService, - sqlite_session: Session, - config_overrides: Callable[..., None], - ): - config_overrides(DEPLOYMENT_EDITION=DeploymentEdition.CLOUD) - """ - Test publish_workflow raises error when trigger node limit exceeded in SANDBOX plan. + Deployment records and audit logs refer to versions by number, so reusing one + would make two different workflows share an identity. + """ + app = TestWorkflowAssociatedDataFactory.create_app() + account = TestWorkflowAssociatedDataFactory.create_account() + graph = TestWorkflowAssociatedDataFactory.create_valid_workflow_graph() - Free/sandbox tier users have limits on the number of trigger nodes. - This prevents resource abuse while allowing users to test the feature. - The limit is enforced at publish time, not during draft editing. - """ - app = TestWorkflowAssociatedDataFactory.create_app() - account = TestWorkflowAssociatedDataFactory.create_account() + draft = TestWorkflowAssociatedDataFactory.create_workflow(version=Workflow.VERSION_DRAFT, graph=graph) + sqlite_session.add(draft) + sqlite_session.commit() - # Create graph with 3 trigger nodes (exceeds SANDBOX limit of 2) - # Trigger nodes enable event-driven automation which consumes resources - graph = { - "nodes": [ - {"id": "trigger-1", "data": {"type": "trigger-webhook"}}, - {"id": "trigger-2", "data": {"type": "trigger-schedule"}}, - {"id": "trigger-3", "data": {"type": "trigger-plugin"}}, - ], - "edges": [], - } - draft = TestWorkflowAssociatedDataFactory.create_workflow(version=Workflow.VERSION_DRAFT, graph=graph) - sqlite_session.add(draft) - sqlite_session.commit() + with ( + patch("services.workflow_service.app_published_workflow_was_updated"), + ): + published = workflow_service.publish_workflow(session=sqlite_session, app_model=app, account=account) + sqlite_session.flush() + sqlite_session.delete(published) + sqlite_session.flush() - with ( - patch("services.workflow_service.BillingService") as mock_billing_service, + republished = workflow_service.publish_workflow(session=sqlite_session, app_model=app, account=account) + + assert republished.version_number == 2 + + def test_publish_workflow_numbers_each_app_independently( + self, workflow_service: WorkflowService, sqlite_session: Session ): - mock_billing_service.get_info.return_value = {"subscription": {"plan": "sandbox"}} + """ + Test the counter is scoped per app rather than global. - with pytest.raises(TriggerNodeLimitExceededError): + Every app starts its own sequence at #1; a busy neighbour must not advance it. + """ + account = TestWorkflowAssociatedDataFactory.create_account() + graph = TestWorkflowAssociatedDataFactory.create_valid_workflow_graph() + + published: list[Workflow] = [] + for app_id in ("app-first", "app-second"): + app = TestWorkflowAssociatedDataFactory.create_app(app_id=app_id) + draft = TestWorkflowAssociatedDataFactory.create_workflow( + workflow_id=f"draft-{app_id}", + app_id=app_id, + version=Workflow.VERSION_DRAFT, + graph=graph, + ) + sqlite_session.add(draft) + sqlite_session.commit() + + with ( + patch("services.workflow_service.app_published_workflow_was_updated"), + ): + workflow = workflow_service.publish_workflow(session=sqlite_session, app_model=app, account=account) + published.append(workflow) + + assert [workflow.version_number for workflow in published] == [1, 1] + + def test_publish_workflow_no_draft_raises_error( + self, workflow_service: WorkflowService, sqlite_session: Session + ): + """ + Test publish_workflow raises error when no draft exists. + + Cannot publish if there's no draft to publish from. + Users must create and save a draft before publishing. + """ + app = TestWorkflowAssociatedDataFactory.create_app() + account = TestWorkflowAssociatedDataFactory.create_account() + + with pytest.raises(ValueError, match="No valid workflow found"): workflow_service.publish_workflow(session=sqlite_session, app_model=app, account=account) - # ==================== Version Management Tests ==================== - # These tests verify listing and managing published workflow versions - - def test_get_all_published_workflow_with_pagination( - self, workflow_service: WorkflowService, sqlite_session: Session - ): - """ - Test get_all_published_workflow returns paginated results. - - Apps can have many published versions over time. - Pagination prevents loading all versions at once, improving performance. - """ - app = TestWorkflowAssociatedDataFactory.create_app(workflow_id="workflow-123") - - sqlite_session.add_all( + @pytest.mark.parametrize( + ("environment_variables", "error"), [ - TestWorkflowAssociatedDataFactory.create_workflow(workflow_id=f"workflow-{i}", version=f"v{i}") - for i in range(5) - ] + ([], "was not found or is not an LLM variable"), + ( + [ + LLMEnvironmentVariable( + name="shared_model", + value={"provider": "provider", "name": "model", "mode": "completion"}, + ) + ], + "uses mode 'completion'.*uses mode 'chat'", + ), + ], ) - sqlite_session.commit() + def test_publish_workflow_rejects_invalid_llm_environment_reference_without_credential_validation( + self, + workflow_service: WorkflowService, + sqlite_session: Session, + environment_variables: list[LLMEnvironmentVariable], + error: str, + ) -> None: + app = TestWorkflowAssociatedDataFactory.create_app() + account = TestWorkflowAssociatedDataFactory.create_account() + graph = TestWorkflowAssociatedDataFactory.create_valid_workflow_graph() + llm_node = next(node for node in graph["nodes"] if node["data"]["type"] == BuiltinNodeTypes.LLM) + llm_node["data"]["model"] = { + "provider": "provider", + "name": "model", + "mode": "chat", + "completion_params": {}, + } + llm_node["data"]["model_selector"] = ["env", "shared_model"] + draft = TestWorkflowAssociatedDataFactory.create_workflow(graph=graph) + draft.environment_variables = environment_variables + sqlite_session.add(draft) + sqlite_session.commit() - workflows, has_more = workflow_service.get_all_published_workflow( - session=sqlite_session, app_model=app, page=1, limit=10, user_id=None - ) + with ( + patch( + "services.system_feature_service.SystemFeatureService.is_plugin_manager_enabled", + return_value=False, + ), + pytest.raises(ValueError, match=error), + ): + workflow_service.publish_workflow(session=sqlite_session, app_model=app, account=account) - assert len(workflows) == 5 - assert has_more is False + def test_publish_workflow_trigger_limit_exceeded( + self, + workflow_service: WorkflowService, + sqlite_session: Session, + config_overrides: Callable[..., None], + ): + config_overrides(DEPLOYMENT_EDITION=DeploymentEdition.CLOUD) + """ + Test publish_workflow raises error when trigger node limit exceeded in SANDBOX plan. - def test_get_all_published_workflow_lists_the_draft_first( - self, workflow_service: WorkflowService, sqlite_session: Session - ): - """ - Test the draft heads the version list no matter how old it is. + Free/sandbox tier users have limits on the number of trigger nodes. + This prevents resource abuse while allowing users to test the feature. + The limit is enforced at publish time, not during draft editing. + """ + app = TestWorkflowAssociatedDataFactory.create_app() + account = TestWorkflowAssociatedDataFactory.create_account() - A draft is created together with its app and its `created_at` is never refreshed, - so ordering purely by publish time would put it last — off the first page entirely - once the app has accumulated enough published versions. - """ - app = TestWorkflowAssociatedDataFactory.create_app(workflow_id="workflow-3") - app_created_at = datetime(2026, 1, 1) + # Create graph with 3 trigger nodes (exceeds SANDBOX limit of 2) + # Trigger nodes enable event-driven automation which consumes resources + graph = { + "nodes": [ + {"id": "trigger-1", "data": {"type": "trigger-webhook"}}, + {"id": "trigger-2", "data": {"type": "trigger-schedule"}}, + {"id": "trigger-3", "data": {"type": "trigger-plugin"}}, + ], + "edges": [], + } + draft = TestWorkflowAssociatedDataFactory.create_workflow(version=Workflow.VERSION_DRAFT, graph=graph) + sqlite_session.add(draft) + sqlite_session.commit() - sqlite_session.add( - TestWorkflowAssociatedDataFactory.create_workflow( - workflow_id="workflow-draft", - version=Workflow.VERSION_DRAFT, - created_at=app_created_at, + with ( + patch("services.workflow_service.BillingService") as mock_billing_service, + ): + mock_billing_service.get_info.return_value = {"subscription": {"plan": "sandbox"}} + + with pytest.raises(TriggerNodeLimitExceededError): + workflow_service.publish_workflow(session=sqlite_session, app_model=app, account=account) + + # ==================== Version Management Tests ==================== + # These tests verify listing and managing published workflow versions + + def test_get_all_published_workflow_with_pagination( + self, workflow_service: WorkflowService, sqlite_session: Session + ): + """ + Test get_all_published_workflow returns paginated results. + + Apps can have many published versions over time. + Pagination prevents loading all versions at once, improving performance. + """ + app = TestWorkflowAssociatedDataFactory.create_app(workflow_id="workflow-123") + + sqlite_session.add_all( + [ + TestWorkflowAssociatedDataFactory.create_workflow(workflow_id=f"workflow-{i}", version=f"v{i}") + for i in range(5) + ] ) - ) - sqlite_session.add_all( - [ + sqlite_session.commit() + + workflows, has_more = workflow_service.get_all_published_workflow( + session=sqlite_session, app_model=app, page=1, limit=10, user_id=None + ) + + assert len(workflows) == 5 + assert has_more is False + + def test_get_all_published_workflow_lists_the_draft_first( + self, workflow_service: WorkflowService, sqlite_session: Session + ): + """ + Test the draft heads the version list no matter how old it is. + + A draft is created together with its app and its `created_at` is never refreshed, + so ordering purely by publish time would put it last — off the first page entirely + once the app has accumulated enough published versions. + """ + app = TestWorkflowAssociatedDataFactory.create_app(workflow_id="workflow-3") + app_created_at = datetime(2026, 1, 1) + + sqlite_session.add( TestWorkflowAssociatedDataFactory.create_workflow( - workflow_id=f"workflow-{i}", - version=f"2026-02-0{i} 00:00:00", - created_at=app_created_at + timedelta(days=i), + workflow_id="workflow-draft", + version=Workflow.VERSION_DRAFT, + created_at=app_created_at, ) - for i in range(1, 4) - ] - ) - sqlite_session.commit() - - workflows, _ = workflow_service.get_all_published_workflow( - session=sqlite_session, app_model=app, page=1, limit=2, user_id=None - ) - - assert [workflow.id for workflow in workflows] == ["workflow-draft", "workflow-3"] - - def test_get_all_published_workflow_excludes_draft_when_requested( - self, workflow_service: WorkflowService, sqlite_session: Session - ): - """`include_draft=False` filters the draft out at query level so callers that - only understand published versions (e.g. the openapi version-list op, where the - draft breaks restore and ordering) get an exact page. - """ - app = TestWorkflowAssociatedDataFactory.create_app(workflow_id="workflow-3") - app_created_at = datetime(2026, 1, 1) - - sqlite_session.add( - TestWorkflowAssociatedDataFactory.create_workflow( - workflow_id="workflow-draft", - version=Workflow.VERSION_DRAFT, - created_at=app_created_at, ) - ) - sqlite_session.add( - TestWorkflowAssociatedDataFactory.create_workflow( - workflow_id="workflow-3", - version="2026-02-01 00:00:00", - created_at=app_created_at + timedelta(days=1), + sqlite_session.add_all( + [ + TestWorkflowAssociatedDataFactory.create_workflow( + workflow_id=f"workflow-{i}", + version=f"2026-02-0{i} 00:00:00", + created_at=app_created_at + timedelta(days=i), + ) + for i in range(1, 4) + ] ) - ) - sqlite_session.commit() + sqlite_session.commit() - workflows, has_more = workflow_service.get_all_published_workflow( - session=sqlite_session, app_model=app, page=1, limit=10, user_id=None, include_draft=False - ) + workflows, _ = workflow_service.get_all_published_workflow( + session=sqlite_session, app_model=app, page=1, limit=2, user_id=None + ) - assert [workflow.id for workflow in workflows] == ["workflow-3"] - assert has_more is False + assert [workflow.id for workflow in workflows] == ["workflow-draft", "workflow-3"] - def test_get_all_published_workflow_has_more(self, workflow_service: WorkflowService, sqlite_session: Session): - """ - Test get_all_published_workflow indicates has_more when results exceed limit. + def test_get_all_published_workflow_excludes_draft_when_requested( + self, workflow_service: WorkflowService, sqlite_session: Session + ): + """`include_draft=False` filters the draft out at query level so callers that + only understand published versions (e.g. the openapi version-list op, where the + draft breaks restore and ordering) get an exact page. + """ + app = TestWorkflowAssociatedDataFactory.create_app(workflow_id="workflow-3") + app_created_at = datetime(2026, 1, 1) - The has_more flag tells the UI whether to show a "Load More" button. - This is determined by fetching limit+1 records and checking if we got that many. - """ - app = TestWorkflowAssociatedDataFactory.create_app(workflow_id="workflow-123") + sqlite_session.add( + TestWorkflowAssociatedDataFactory.create_workflow( + workflow_id="workflow-draft", + version=Workflow.VERSION_DRAFT, + created_at=app_created_at, + ) + ) + sqlite_session.add( + TestWorkflowAssociatedDataFactory.create_workflow( + workflow_id="workflow-3", + version="2026-02-01 00:00:00", + created_at=app_created_at + timedelta(days=1), + ) + ) + sqlite_session.commit() - sqlite_session.add_all( - [ - TestWorkflowAssociatedDataFactory.create_workflow(workflow_id=f"workflow-{i}", version=f"v{i}") - for i in range(11) - ] - ) - sqlite_session.commit() + workflows, has_more = workflow_service.get_all_published_workflow( + session=sqlite_session, app_model=app, page=1, limit=10, user_id=None, include_draft=False + ) - workflows, has_more = workflow_service.get_all_published_workflow( - session=sqlite_session, app_model=app, page=1, limit=10, user_id=None - ) + assert [workflow.id for workflow in workflows] == ["workflow-3"] + assert has_more is False - assert len(workflows) == 10 - assert has_more is True + def test_get_all_published_workflow_has_more(self, workflow_service: WorkflowService, sqlite_session: Session): + """ + Test get_all_published_workflow indicates has_more when results exceed limit. - def test_get_all_published_workflow_no_workflow_id( - self, workflow_service: WorkflowService, sqlite_session: Session - ): - """Test get_all_published_workflow returns empty when app has no workflow_id.""" - app = TestWorkflowAssociatedDataFactory.create_app(workflow_id=None) + The has_more flag tells the UI whether to show a "Load More" button. + This is determined by fetching limit+1 records and checking if we got that many. + """ + app = TestWorkflowAssociatedDataFactory.create_app(workflow_id="workflow-123") - workflows, has_more = workflow_service.get_all_published_workflow( - session=sqlite_session, app_model=app, page=1, limit=10, user_id=None - ) + sqlite_session.add_all( + [ + TestWorkflowAssociatedDataFactory.create_workflow(workflow_id=f"workflow-{i}", version=f"v{i}") + for i in range(11) + ] + ) + sqlite_session.commit() - assert workflows == [] - assert has_more is False + workflows, has_more = workflow_service.get_all_published_workflow( + session=sqlite_session, app_model=app, page=1, limit=10, user_id=None + ) - # ==================== Update Workflow Tests ==================== - # These tests verify updating workflow metadata (name, comments, etc.) + assert len(workflows) == 10 + assert has_more is True - def test_update_workflow_success(self, workflow_service: WorkflowService, sqlite_session: Session): - """ - Test update_workflow updates workflow attributes. + def test_get_all_published_workflow_no_workflow_id( + self, workflow_service: WorkflowService, sqlite_session: Session + ): + """Test get_all_published_workflow returns empty when app has no workflow_id.""" + app = TestWorkflowAssociatedDataFactory.create_app(workflow_id=None) - Allows updating metadata like marked_name and marked_comment - without creating a new version. Only specific fields are allowed - to prevent accidental modification of workflow logic. - """ - workflow_id = "workflow-123" - tenant_id = "tenant-456" - app_id = "app-789" - workflow_ref = WorkflowRef(tenant_id=tenant_id, owner_id=app_id, workflow_id=workflow_id) - account_id = "user-123" - workflow = TestWorkflowAssociatedDataFactory.create_workflow( - workflow_id=workflow_id, tenant_id=tenant_id, app_id=app_id - ) - sqlite_session.add(workflow) - sqlite_session.commit() + workflows, has_more = workflow_service.get_all_published_workflow( + session=sqlite_session, app_model=app, page=1, limit=10, user_id=None + ) - result = workflow_service.update_workflow( - session=sqlite_session, - account_id=account_id, - data={"marked_name": "Updated Name", "marked_comment": "Updated Comment"}, - workflow_ref=workflow_ref, - ) + assert workflows == [] + assert has_more is False - sqlite_session.flush() - assert result is workflow - assert workflow.marked_name == "Updated Name" - assert workflow.marked_comment == "Updated Comment" - assert workflow.updated_by == account_id + # ==================== Update Workflow Tests ==================== + # These tests verify updating workflow metadata (name, comments, etc.) - def test_update_workflow_not_found(self, workflow_service: WorkflowService, sqlite_session: Session): - """Test update_workflow returns None when workflow not found.""" - result = workflow_service.update_workflow( - session=sqlite_session, - account_id="user-123", - data={"marked_name": "Test"}, - workflow_ref=WorkflowRef(tenant_id="tenant-456", owner_id="app-789", workflow_id="nonexistent"), - ) + def test_update_workflow_success(self, workflow_service: WorkflowService, sqlite_session: Session): + """ + Test update_workflow updates workflow attributes. - assert result is None + Allows updating metadata like marked_name and marked_comment + without creating a new version. Only specific fields are allowed + to prevent accidental modification of workflow logic. + """ + workflow_id = "workflow-123" + tenant_id = "tenant-456" + app_id = "app-789" + workflow_ref = WorkflowRef(tenant_id=tenant_id, owner_id=app_id, workflow_id=workflow_id) + account_id = "user-123" + workflow = TestWorkflowAssociatedDataFactory.create_workflow( + workflow_id=workflow_id, tenant_id=tenant_id, app_id=app_id + ) + sqlite_session.add(workflow) + sqlite_session.commit() - def test_update_workflow_with_ref_scopes_lookup_to_app( - self, workflow_service: WorkflowService, sqlite_session: Session - ): - """Test update_workflow includes the trusted app owner in the lookup.""" - workflow_id = "workflow-123" - tenant_id = "tenant-456" - app_id = "app-789" - account_id = "user-123" - workflow_ref = WorkflowRef(tenant_id=tenant_id, owner_id="other-app", workflow_id=workflow_id) - workflow = TestWorkflowAssociatedDataFactory.create_workflow( - workflow_id=workflow_id, tenant_id=tenant_id, app_id=app_id - ) - sqlite_session.add(workflow) - sqlite_session.commit() + result = workflow_service.update_workflow( + session=sqlite_session, + account_id=account_id, + data={"marked_name": "Updated Name", "marked_comment": "Updated Comment"}, + workflow_ref=workflow_ref, + ) - result = workflow_service.update_workflow( - session=sqlite_session, - account_id=account_id, - data={"marked_name": "Updated Name"}, - workflow_ref=workflow_ref, - ) + sqlite_session.flush() + assert result is workflow + assert workflow.marked_name == "Updated Name" + assert workflow.marked_comment == "Updated Comment" + assert workflow.updated_by == account_id - assert result is None - assert workflow.marked_name == "" + def test_update_workflow_not_found(self, workflow_service: WorkflowService, sqlite_session: Session): + """Test update_workflow returns None when workflow not found.""" + result = workflow_service.update_workflow( + session=sqlite_session, + account_id="user-123", + data={"marked_name": "Test"}, + workflow_ref=WorkflowRef(tenant_id="tenant-456", owner_id="app-789", workflow_id="nonexistent"), + ) - # ==================== Delete Workflow Tests ==================== - # These tests verify workflow deletion with safety checks + assert result is None - def test_delete_workflow_success(self, workflow_service: WorkflowService, sqlite_session: Session): - """ - Test delete_workflow successfully deletes a published workflow. + def test_update_workflow_with_ref_scopes_lookup_to_app( + self, workflow_service: WorkflowService, sqlite_session: Session + ): + """Test update_workflow includes the trusted app owner in the lookup.""" + workflow_id = "workflow-123" + tenant_id = "tenant-456" + app_id = "app-789" + account_id = "user-123" + workflow_ref = WorkflowRef(tenant_id=tenant_id, owner_id="other-app", workflow_id=workflow_id) + workflow = TestWorkflowAssociatedDataFactory.create_workflow( + workflow_id=workflow_id, tenant_id=tenant_id, app_id=app_id + ) + sqlite_session.add(workflow) + sqlite_session.commit() - Users can delete old published versions they no longer need. - This helps manage storage and keeps the version list clean. - """ - workflow_id = "workflow-123" - tenant_id = "tenant-456" - app_id = "app-789" - workflow_ref = WorkflowRef(tenant_id=tenant_id, owner_id=app_id, workflow_id=workflow_id) - workflow = TestWorkflowAssociatedDataFactory.create_workflow( - workflow_id=workflow_id, tenant_id=tenant_id, app_id=app_id, version="v1" - ) - inline_binding = WorkflowAgentNodeBinding( - id="inline-binding", - tenant_id=tenant_id, - app_id=app_id, - workflow_id=workflow_id, - workflow_version="v1", - node_id="inline-node", - binding_type=WorkflowAgentBindingType.INLINE_AGENT, - agent_id="inline-agent", - current_snapshot_id="snapshot-1", - node_job_config={}, - ) - roster_binding = WorkflowAgentNodeBinding( - id="roster-binding", - tenant_id=tenant_id, - app_id=app_id, - workflow_id=workflow_id, - workflow_version="v1", - node_id="roster-node", - binding_type=WorkflowAgentBindingType.ROSTER_AGENT, - agent_id="roster-agent", - current_snapshot_id="snapshot-2", - node_job_config={}, - ) - non_target_bindings = [ - WorkflowAgentNodeBinding( - id=f"non-target-{key}", - tenant_id="other-tenant" if key == "tenant" else tenant_id, - app_id="other-app" if key == "app" else app_id, - workflow_id="other-workflow" if key == "workflow" else workflow_id, - workflow_version="other-version" if key == "version" else workflow.version, - node_id=f"{key}-node", + result = workflow_service.update_workflow( + session=sqlite_session, + account_id=account_id, + data={"marked_name": "Updated Name"}, + workflow_ref=workflow_ref, + ) + + assert result is None + assert workflow.marked_name == "" + + # ==================== Delete Workflow Tests ==================== + # These tests verify workflow deletion with safety checks + + def test_delete_workflow_success(self, workflow_service: WorkflowService, sqlite_session: Session): + """ + Test delete_workflow successfully deletes a published workflow. + + Users can delete old published versions they no longer need. + This helps manage storage and keeps the version list clean. + """ + workflow_id = "workflow-123" + tenant_id = "tenant-456" + app_id = "app-789" + workflow_ref = WorkflowRef(tenant_id=tenant_id, owner_id=app_id, workflow_id=workflow_id) + workflow = TestWorkflowAssociatedDataFactory.create_workflow( + workflow_id=workflow_id, tenant_id=tenant_id, app_id=app_id, version="v1" + ) + inline_binding = WorkflowAgentNodeBinding( + id="inline-binding", + tenant_id=tenant_id, + app_id=app_id, + workflow_id=workflow_id, + workflow_version="v1", + node_id="inline-node", binding_type=WorkflowAgentBindingType.INLINE_AGENT, - agent_id=f"{key}-inline-agent", - current_snapshot_id=f"{key}-snapshot", + agent_id="inline-agent", + current_snapshot_id="snapshot-1", node_job_config={}, ) - for key in ("tenant", "app", "workflow", "version") - ] - sqlite_session.add_all([workflow, inline_binding, roster_binding, *non_target_bindings]) - sqlite_session.commit() + roster_binding = WorkflowAgentNodeBinding( + id="roster-binding", + tenant_id=tenant_id, + app_id=app_id, + workflow_id=workflow_id, + workflow_version="v1", + node_id="roster-node", + binding_type=WorkflowAgentBindingType.ROSTER_AGENT, + agent_id="roster-agent", + current_snapshot_id="snapshot-2", + node_job_config={}, + ) + non_target_bindings = [ + WorkflowAgentNodeBinding( + id=f"non-target-{key}", + tenant_id="other-tenant" if key == "tenant" else tenant_id, + app_id="other-app" if key == "app" else app_id, + workflow_id="other-workflow" if key == "workflow" else workflow_id, + workflow_version="other-version" if key == "version" else workflow.version, + node_id=f"{key}-node", + binding_type=WorkflowAgentBindingType.INLINE_AGENT, + agent_id=f"{key}-inline-agent", + current_snapshot_id=f"{key}-snapshot", + node_job_config={}, + ) + for key in ("tenant", "app", "workflow", "version") + ] + sqlite_session.add_all([workflow, inline_binding, roster_binding, *non_target_bindings]) + sqlite_session.commit() - result = workflow_service.delete_workflow(session=sqlite_session, workflow_ref=workflow_ref) - sqlite_session.flush() + result = workflow_service.delete_workflow(session=sqlite_session, workflow_ref=workflow_ref) + sqlite_session.flush() - assert result == ["inline-agent"] - assert sqlite_session.get(Workflow, workflow_id) is None - assert sqlite_session.get(WorkflowAgentNodeBinding, inline_binding.id) is None - assert sqlite_session.get(WorkflowAgentNodeBinding, roster_binding.id) is None - for binding in non_target_bindings: - assert sqlite_session.get(WorkflowAgentNodeBinding, binding.id) is binding + assert result == ["inline-agent"] + assert sqlite_session.get(Workflow, workflow_id) is None + assert sqlite_session.get(WorkflowAgentNodeBinding, inline_binding.id) is None + assert sqlite_session.get(WorkflowAgentNodeBinding, roster_binding.id) is None + for binding in non_target_bindings: + assert sqlite_session.get(WorkflowAgentNodeBinding, binding.id) is binding - def test_delete_workflow_locks_source_until_caller_commits( - self, workflow_service: WorkflowService, sqlite_session: Session, mocker: MockerFixture - ): - workflow = TestWorkflowAssociatedDataFactory.create_workflow(version="v1") - workflow_ref = WorkflowRef( - tenant_id=workflow.tenant_id, - owner_id=workflow.app_id, - workflow_id=workflow.id, - ) - session = sqlite_session - session.add(workflow) - session.commit() - lookup = mocker.spy(session, "scalar") + def test_delete_workflow_locks_source_until_caller_commits( + self, workflow_service: WorkflowService, sqlite_session: Session, mocker: MockerFixture + ): + workflow = TestWorkflowAssociatedDataFactory.create_workflow(version="v1") + workflow_ref = WorkflowRef( + tenant_id=workflow.tenant_id, + owner_id=workflow.app_id, + workflow_id=workflow.id, + ) + session = sqlite_session + session.add(workflow) + session.commit() + lookup = mocker.spy(session, "scalar") - result = workflow_service.delete_workflow(session=session, workflow_ref=workflow_ref) + result = workflow_service.delete_workflow(session=session, workflow_ref=workflow_ref) - stmt = lookup.call_args_list[0].args[0] - sql = str(stmt.compile(dialect=postgresql.dialect())) - assert result == [] - assert "FOR UPDATE" in sql - assert session.in_transaction() - assert workflow in session.deleted - session.commit() - assert session.get(Workflow, workflow.id) is None + stmt = lookup.call_args_list[0].args[0] + sql = str(stmt.compile(dialect=postgresql.dialect())) + assert result == [] + assert "FOR UPDATE" in sql + assert session.in_transaction() + assert workflow in session.deleted + session.commit() + assert session.get(Workflow, workflow.id) is None - def test_delete_workflow_with_ref_scopes_lookup_to_app( - self, workflow_service: WorkflowService, sqlite_session: Session - ): - """Test delete_workflow includes the trusted app owner in the lookup.""" - workflow_id = "workflow-123" - tenant_id = "tenant-456" - app_id = "app-789" - workflow_ref = WorkflowRef(tenant_id=tenant_id, owner_id="other-app", workflow_id=workflow_id) - workflow = TestWorkflowAssociatedDataFactory.create_workflow( - workflow_id=workflow_id, tenant_id=tenant_id, app_id=app_id, version="v1" - ) - sqlite_session.add(workflow) - sqlite_session.commit() + def test_delete_workflow_with_ref_scopes_lookup_to_app( + self, workflow_service: WorkflowService, sqlite_session: Session + ): + """Test delete_workflow includes the trusted app owner in the lookup.""" + workflow_id = "workflow-123" + tenant_id = "tenant-456" + app_id = "app-789" + workflow_ref = WorkflowRef(tenant_id=tenant_id, owner_id="other-app", workflow_id=workflow_id) + workflow = TestWorkflowAssociatedDataFactory.create_workflow( + workflow_id=workflow_id, tenant_id=tenant_id, app_id=app_id, version="v1" + ) + sqlite_session.add(workflow) + sqlite_session.commit() - with pytest.raises(ValueError, match="not found"): - workflow_service.delete_workflow(session=sqlite_session, workflow_ref=workflow_ref) + with pytest.raises(ValueError, match="not found"): + workflow_service.delete_workflow(session=sqlite_session, workflow_ref=workflow_ref) - assert sqlite_session.get(Workflow, workflow_id) is workflow + assert sqlite_session.get(Workflow, workflow_id) is workflow - def test_delete_workflow_draft_raises_error(self, workflow_service: WorkflowService, sqlite_session: Session): - """ - Test delete_workflow raises error when trying to delete draft. + def test_delete_workflow_draft_raises_error(self, workflow_service: WorkflowService, sqlite_session: Session): + """ + Test delete_workflow raises error when trying to delete draft. - Draft workflows cannot be deleted - they're the working copy. - Users can only delete published versions to clean up old snapshots. - """ - workflow_id = "workflow-123" - tenant_id = "tenant-456" - workflow_ref = WorkflowRef(tenant_id=tenant_id, owner_id="app-789", workflow_id=workflow_id) - workflow = TestWorkflowAssociatedDataFactory.create_workflow( - workflow_id=workflow_id, - tenant_id=tenant_id, - app_id=workflow_ref.owner_id, - version=Workflow.VERSION_DRAFT, - ) - sqlite_session.add(workflow) - sqlite_session.commit() + Draft workflows cannot be deleted - they're the working copy. + Users can only delete published versions to clean up old snapshots. + """ + workflow_id = "workflow-123" + tenant_id = "tenant-456" + workflow_ref = WorkflowRef(tenant_id=tenant_id, owner_id="app-789", workflow_id=workflow_id) + workflow = TestWorkflowAssociatedDataFactory.create_workflow( + workflow_id=workflow_id, + tenant_id=tenant_id, + app_id=workflow_ref.owner_id, + version=Workflow.VERSION_DRAFT, + ) + sqlite_session.add(workflow) + sqlite_session.commit() - with pytest.raises(DraftWorkflowDeletionError, match="Cannot delete draft workflow"): - workflow_service.delete_workflow(session=sqlite_session, workflow_ref=workflow_ref) + with pytest.raises(DraftWorkflowDeletionError, match="Cannot delete draft workflow"): + workflow_service.delete_workflow(session=sqlite_session, workflow_ref=workflow_ref) - def test_delete_workflow_in_use_by_app_raises_error( - self, workflow_service: WorkflowService, sqlite_session: Session - ): - """ - Test delete_workflow raises error when workflow is in use by app. + def test_delete_workflow_in_use_by_app_raises_error( + self, workflow_service: WorkflowService, sqlite_session: Session + ): + """ + Test delete_workflow raises error when workflow is in use by app. - Cannot delete a workflow version that's currently published/active. - This would break the app for users. Must publish a different version first. - """ - workflow_id = "workflow-123" - tenant_id = "tenant-456" - workflow_ref = WorkflowRef(tenant_id=tenant_id, owner_id="app-789", workflow_id=workflow_id) - workflow = TestWorkflowAssociatedDataFactory.create_workflow( - workflow_id=workflow_id, tenant_id=tenant_id, app_id=workflow_ref.owner_id, version="v1" - ) - app = App( - id="active-app", - tenant_id=tenant_id, - name="Active App", - description="", - mode=AppMode.WORKFLOW, - workflow_id=workflow_id, - enable_site=True, - enable_api=True, - max_active_requests=0, - ) - sqlite_session.add_all([workflow, app]) - sqlite_session.commit() + Cannot delete a workflow version that's currently published/active. + This would break the app for users. Must publish a different version first. + """ + workflow_id = "workflow-123" + tenant_id = "tenant-456" + workflow_ref = WorkflowRef(tenant_id=tenant_id, owner_id="app-789", workflow_id=workflow_id) + workflow = TestWorkflowAssociatedDataFactory.create_workflow( + workflow_id=workflow_id, tenant_id=tenant_id, app_id=workflow_ref.owner_id, version="v1" + ) + app = App( + id="active-app", + tenant_id=tenant_id, + name="Active App", + description="", + mode=AppMode.WORKFLOW, + workflow_id=workflow_id, + enable_site=True, + enable_api=True, + max_active_requests=0, + ) + sqlite_session.add_all([workflow, app]) + sqlite_session.commit() - with pytest.raises(WorkflowInUseError, match="currently in use by app"): - workflow_service.delete_workflow(session=sqlite_session, workflow_ref=workflow_ref) + with pytest.raises(WorkflowInUseError, match="currently in use by app"): + workflow_service.delete_workflow(session=sqlite_session, workflow_ref=workflow_ref) - def test_delete_workflow_published_as_tool_raises_error( - self, workflow_service: WorkflowService, sqlite_session: Session - ): - """ - Test delete_workflow raises error when workflow is published as tool. + def test_delete_workflow_published_as_tool_raises_error( + self, workflow_service: WorkflowService, sqlite_session: Session + ): + """ + Test delete_workflow raises error when workflow is published as tool. - Workflows can be published as reusable tools for other workflows. - Cannot delete a version that's being used as a tool, as this would - break other workflows that depend on it. - """ - workflow_id = "workflow-123" - tenant_id = "tenant-456" - workflow_ref = WorkflowRef(tenant_id=tenant_id, owner_id="app-789", workflow_id=workflow_id) - workflow = TestWorkflowAssociatedDataFactory.create_workflow( - workflow_id=workflow_id, tenant_id=tenant_id, app_id=workflow_ref.owner_id, version="v1" - ) - tool_provider = WorkflowToolProvider( - name="workflow-tool", - label="Workflow Tool", - icon="icon.svg", - app_id=workflow.app_id, - version=workflow.version, - user_id="user-123", - tenant_id=workflow.tenant_id, - description="Test provider", - parameter_configuration="[]", - ) - sqlite_session.add_all([workflow, tool_provider]) - sqlite_session.commit() + Workflows can be published as reusable tools for other workflows. + Cannot delete a version that's being used as a tool, as this would + break other workflows that depend on it. + """ + workflow_id = "workflow-123" + tenant_id = "tenant-456" + workflow_ref = WorkflowRef(tenant_id=tenant_id, owner_id="app-789", workflow_id=workflow_id) + workflow = TestWorkflowAssociatedDataFactory.create_workflow( + workflow_id=workflow_id, tenant_id=tenant_id, app_id=workflow_ref.owner_id, version="v1" + ) + tool_provider = WorkflowToolProvider( + name="workflow-tool", + label="Workflow Tool", + icon="icon.svg", + app_id=workflow.app_id, + version=workflow.version, + user_id="user-123", + tenant_id=workflow.tenant_id, + description="Test provider", + parameter_configuration="[]", + ) + sqlite_session.add_all([workflow, tool_provider]) + sqlite_session.commit() - with pytest.raises(WorkflowInUseError, match="published as a tool"): - workflow_service.delete_workflow(session=sqlite_session, workflow_ref=workflow_ref) + with pytest.raises(WorkflowInUseError, match="published as a tool"): + workflow_service.delete_workflow(session=sqlite_session, workflow_ref=workflow_ref) - def test_delete_workflow_not_found_raises_error(self, workflow_service: WorkflowService, sqlite_session: Session): - """Test delete_workflow raises error when workflow not found.""" - workflow_id = "nonexistent" - tenant_id = "tenant-456" - workflow_ref = WorkflowRef(tenant_id=tenant_id, owner_id="app-789", workflow_id=workflow_id) + def test_delete_workflow_not_found_raises_error( + self, workflow_service: WorkflowService, sqlite_session: Session + ): + """Test delete_workflow raises error when workflow not found.""" + workflow_id = "nonexistent" + tenant_id = "tenant-456" + workflow_ref = WorkflowRef(tenant_id=tenant_id, owner_id="app-789", workflow_id=workflow_id) - with pytest.raises(ValueError, match="not found"): - workflow_service.delete_workflow(session=sqlite_session, workflow_ref=workflow_ref) + with pytest.raises(ValueError, match="not found"): + workflow_service.delete_workflow(session=sqlite_session, workflow_ref=workflow_ref) - # ==================== Get Default Block Config Tests ==================== - # These tests verify retrieval of default node configurations + # ==================== Get Default Block Config Tests ==================== + # These tests verify retrieval of default node configurations - def test_get_default_block_configs(self, workflow_service: WorkflowService): - """ - Test get_default_block_configs returns list of default configs. + def test_get_default_block_configs(self, workflow_service: WorkflowService): + """ + Test get_default_block_configs returns list of default configs. - Returns default configurations for all available node types. - Used by the UI to populate the node palette and provide sensible defaults - when users add new nodes to their workflow. - """ - with patch("services.workflow_service.get_node_type_classes_mapping") as mock_mapping: - # Mock node class with default config - mock_node_class = MagicMock() - mock_node_class.get_default_config.return_value = {"type": "llm", "config": {}} + Returns default configurations for all available node types. + Used by the UI to populate the node palette and provide sensible defaults + when users add new nodes to their workflow. + """ + with patch("services.workflow_service.get_node_type_classes_mapping") as mock_mapping: + # Mock node class with default config + mock_node_class = MagicMock() + mock_node_class.get_default_config.return_value = {"type": "llm", "config": {}} - mock_mapping.return_value = {BuiltinNodeTypes.LLM: {"latest": mock_node_class}} + mock_mapping.return_value = {BuiltinNodeTypes.LLM: {"latest": mock_node_class}} + + with patch("services.workflow_service.LATEST_VERSION", "latest"): + result = workflow_service.get_default_block_configs() + + assert len(result) > 0 + + def test_get_default_block_configs_http_request_injects_default_config(self, workflow_service: WorkflowService): + injected_config = HttpRequestNodeConfig( + max_connect_timeout=15, + max_read_timeout=25, + max_write_timeout=35, + max_binary_size=4096, + max_text_size=2048, + ssl_verify=True, + ssrf_default_max_retries=6, + ) + + with ( + patch("services.workflow_service.get_node_type_classes_mapping") as mock_mapping, + patch("services.workflow_service.LATEST_VERSION", "latest"), + patch( + "services.workflow_service.build_http_request_config", + return_value=injected_config, + ) as mock_build_config, + ): + mock_http_node_class = MagicMock() + mock_http_node_class.get_default_config.return_value = {"type": "http-request", "config": {}} + mock_llm_node_class = MagicMock() + mock_llm_node_class.get_default_config.return_value = {"type": "llm", "config": {}} + mock_mapping.return_value = { + BuiltinNodeTypes.HTTP_REQUEST: {"latest": mock_http_node_class}, + BuiltinNodeTypes.LLM: {"latest": mock_llm_node_class}, + } - with patch("services.workflow_service.LATEST_VERSION", "latest"): result = workflow_service.get_default_block_configs() - assert len(result) > 0 + assert result == [ + {"type": "http-request", "config": {}}, + {"type": "llm", "config": {}}, + ] + mock_build_config.assert_called_once() + passed_http_filters = mock_http_node_class.get_default_config.call_args.kwargs["filters"] + assert passed_http_filters[HTTP_REQUEST_CONFIG_FILTER_KEY] is injected_config + mock_llm_node_class.get_default_config.assert_called_once_with(filters=None) - def test_get_default_block_configs_http_request_injects_default_config(self, workflow_service: WorkflowService): - injected_config = HttpRequestNodeConfig( - max_connect_timeout=15, - max_read_timeout=25, - max_write_timeout=35, - max_binary_size=4096, - max_text_size=2048, - ssl_verify=True, - ssrf_default_max_retries=6, - ) + def test_get_default_block_config_for_node_type(self, workflow_service: WorkflowService): + """ + Test get_default_block_config returns config for specific node type. - with ( - patch("services.workflow_service.get_node_type_classes_mapping") as mock_mapping, - patch("services.workflow_service.LATEST_VERSION", "latest"), - patch( - "services.workflow_service.build_http_request_config", - return_value=injected_config, - ) as mock_build_config, - ): - mock_http_node_class = MagicMock() - mock_http_node_class.get_default_config.return_value = {"type": "http-request", "config": {}} - mock_llm_node_class = MagicMock() - mock_llm_node_class.get_default_config.return_value = {"type": "llm", "config": {}} - mock_mapping.return_value = { - BuiltinNodeTypes.HTTP_REQUEST: {"latest": mock_http_node_class}, - BuiltinNodeTypes.LLM: {"latest": mock_llm_node_class}, - } + Returns the default configuration for a specific node type (e.g., LLM, HTTP). + This includes default values for all required and optional parameters. + """ + with ( + patch("services.workflow_service.get_node_type_classes_mapping") as mock_mapping, + patch("services.workflow_service.LATEST_VERSION", "latest"), + ): + # Mock node class with default config + mock_node_class = MagicMock() + mock_config = {"type": "llm", "config": {"provider": "openai"}} + mock_node_class.get_default_config.return_value = mock_config - result = workflow_service.get_default_block_configs() + # Create a mock mapping that includes BuiltinNodeTypes.LLM + mock_mapping.return_value = {BuiltinNodeTypes.LLM: {"latest": mock_node_class}} - assert result == [ - {"type": "http-request", "config": {}}, - {"type": "llm", "config": {}}, - ] - mock_build_config.assert_called_once() - passed_http_filters = mock_http_node_class.get_default_config.call_args.kwargs["filters"] - assert passed_http_filters[HTTP_REQUEST_CONFIG_FILTER_KEY] is injected_config - mock_llm_node_class.get_default_config.assert_called_once_with(filters=None) + result = workflow_service.get_default_block_config(BuiltinNodeTypes.LLM) - def test_get_default_block_config_for_node_type(self, workflow_service: WorkflowService): - """ - Test get_default_block_config returns config for specific node type. + assert result == mock_config + mock_node_class.get_default_config.assert_called_once() - Returns the default configuration for a specific node type (e.g., LLM, HTTP). - This includes default values for all required and optional parameters. - """ - with ( - patch("services.workflow_service.get_node_type_classes_mapping") as mock_mapping, - patch("services.workflow_service.LATEST_VERSION", "latest"), - ): - # Mock node class with default config - mock_node_class = MagicMock() - mock_config = {"type": "llm", "config": {"provider": "openai"}} - mock_node_class.get_default_config.return_value = mock_config + def test_get_default_block_config_invalid_node_type(self, workflow_service: WorkflowService): + """Test get_default_block_config returns empty dict for invalid node type.""" + with patch("services.workflow_service.get_node_type_classes_mapping") as mock_mapping: + mock_mapping.return_value = {} - # Create a mock mapping that includes BuiltinNodeTypes.LLM - mock_mapping.return_value = {BuiltinNodeTypes.LLM: {"latest": mock_node_class}} + # Use a valid NodeType but one that's not in the mapping + result = workflow_service.get_default_block_config(BuiltinNodeTypes.LLM) - result = workflow_service.get_default_block_config(BuiltinNodeTypes.LLM) + assert result == {} - assert result == mock_config - mock_node_class.get_default_config.assert_called_once() - - def test_get_default_block_config_invalid_node_type(self, workflow_service: WorkflowService): - """Test get_default_block_config returns empty dict for invalid node type.""" - with patch("services.workflow_service.get_node_type_classes_mapping") as mock_mapping: - mock_mapping.return_value = {} - - # Use a valid NodeType but one that's not in the mapping - result = workflow_service.get_default_block_config(BuiltinNodeTypes.LLM) - - assert result == {} - - def test_get_default_block_config_http_request_injects_default_config(self, workflow_service: WorkflowService): - injected_config = HttpRequestNodeConfig( - max_connect_timeout=11, - max_read_timeout=22, - max_write_timeout=33, - max_binary_size=4096, - max_text_size=2048, - ssl_verify=False, - ssrf_default_max_retries=7, - ) - - with ( - patch("services.workflow_service.get_node_type_classes_mapping") as mock_mapping, - patch("services.workflow_service.LATEST_VERSION", "latest"), - patch( - "services.workflow_service.build_http_request_config", - return_value=injected_config, - ) as mock_build_config, - ): - mock_node_class = MagicMock() - expected = {"type": "http-request", "config": {}} - mock_node_class.get_default_config.return_value = expected - mock_mapping.return_value = {BuiltinNodeTypes.HTTP_REQUEST: {"latest": mock_node_class}} - - result = workflow_service.get_default_block_config(BuiltinNodeTypes.HTTP_REQUEST) - - assert result == expected - mock_build_config.assert_called_once() - passed_filters = mock_node_class.get_default_config.call_args.kwargs["filters"] - assert passed_filters[HTTP_REQUEST_CONFIG_FILTER_KEY] is injected_config - - def test_get_default_block_config_http_request_uses_passed_config(self, workflow_service: WorkflowService): - provided_config = HttpRequestNodeConfig( - max_connect_timeout=13, - max_read_timeout=23, - max_write_timeout=34, - max_binary_size=8192, - max_text_size=4096, - ssl_verify=True, - ssrf_default_max_retries=2, - ) - - with ( - patch("services.workflow_service.get_node_type_classes_mapping") as mock_mapping, - patch("services.workflow_service.LATEST_VERSION", "latest"), - patch("services.workflow_service.build_http_request_config") as mock_build_config, - ): - mock_node_class = MagicMock() - expected = {"type": "http-request", "config": {}} - mock_node_class.get_default_config.return_value = expected - mock_mapping.return_value = {BuiltinNodeTypes.HTTP_REQUEST: {"latest": mock_node_class}} - - result = workflow_service.get_default_block_config( - BuiltinNodeTypes.HTTP_REQUEST, - filters={HTTP_REQUEST_CONFIG_FILTER_KEY: provided_config}, + def test_get_default_block_config_http_request_injects_default_config(self, workflow_service: WorkflowService): + injected_config = HttpRequestNodeConfig( + max_connect_timeout=11, + max_read_timeout=22, + max_write_timeout=33, + max_binary_size=4096, + max_text_size=2048, + ssl_verify=False, + ssrf_default_max_retries=7, ) - assert result == expected - mock_build_config.assert_not_called() - passed_filters = mock_node_class.get_default_config.call_args.kwargs["filters"] - assert passed_filters[HTTP_REQUEST_CONFIG_FILTER_KEY] is provided_config + with ( + patch("services.workflow_service.get_node_type_classes_mapping") as mock_mapping, + patch("services.workflow_service.LATEST_VERSION", "latest"), + patch( + "services.workflow_service.build_http_request_config", + return_value=injected_config, + ) as mock_build_config, + ): + mock_node_class = MagicMock() + expected = {"type": "http-request", "config": {}} + mock_node_class.get_default_config.return_value = expected + mock_mapping.return_value = {BuiltinNodeTypes.HTTP_REQUEST: {"latest": mock_node_class}} - def test_get_default_block_config_http_request_malformed_config_raises_type_error( - self, workflow_service: WorkflowService - ): - with ( - patch( - "services.workflow_service.get_node_type_classes_mapping", - return_value={BuiltinNodeTypes.HTTP_REQUEST: {"latest": HttpRequestNode}}, - ), - patch("services.workflow_service.LATEST_VERSION", "latest"), - ): - with pytest.raises(TypeError, match="http_request_config must be an HttpRequestNodeConfig instance"): - workflow_service.get_default_block_config( + result = workflow_service.get_default_block_config(BuiltinNodeTypes.HTTP_REQUEST) + + assert result == expected + mock_build_config.assert_called_once() + passed_filters = mock_node_class.get_default_config.call_args.kwargs["filters"] + assert passed_filters[HTTP_REQUEST_CONFIG_FILTER_KEY] is injected_config + + def test_get_default_block_config_http_request_uses_passed_config(self, workflow_service: WorkflowService): + provided_config = HttpRequestNodeConfig( + max_connect_timeout=13, + max_read_timeout=23, + max_write_timeout=34, + max_binary_size=8192, + max_text_size=4096, + ssl_verify=True, + ssrf_default_max_retries=2, + ) + + with ( + patch("services.workflow_service.get_node_type_classes_mapping") as mock_mapping, + patch("services.workflow_service.LATEST_VERSION", "latest"), + patch("services.workflow_service.build_http_request_config") as mock_build_config, + ): + mock_node_class = MagicMock() + expected = {"type": "http-request", "config": {}} + mock_node_class.get_default_config.return_value = expected + mock_mapping.return_value = {BuiltinNodeTypes.HTTP_REQUEST: {"latest": mock_node_class}} + + result = workflow_service.get_default_block_config( BuiltinNodeTypes.HTTP_REQUEST, - filters={HTTP_REQUEST_CONFIG_FILTER_KEY: "invalid"}, + filters={HTTP_REQUEST_CONFIG_FILTER_KEY: provided_config}, ) - # ==================== Workflow Conversion Tests ==================== - # These tests verify converting basic apps to workflow apps + assert result == expected + mock_build_config.assert_not_called() + passed_filters = mock_node_class.get_default_config.call_args.kwargs["filters"] + assert passed_filters[HTTP_REQUEST_CONFIG_FILTER_KEY] is provided_config - def test_convert_to_workflow_from_chat_app(self, workflow_service: WorkflowService, sqlite_session: Session): - """ - Test convert_to_workflow converts chat app to workflow. - - Allows users to migrate from simple chat apps to advanced workflow apps. - The conversion creates equivalent workflow nodes from the chat configuration, - giving users more control and customization options. - """ - app = TestWorkflowAssociatedDataFactory.create_app(mode=AppMode.CHAT) - account = TestWorkflowAssociatedDataFactory.create_account() - args = { - "name": "Converted Workflow", - "icon_type": "emoji", - "icon": "🤖", - "icon_background": "#FFEAD5", - } - - with patch("services.workflow_service.WorkflowConverter") as mock_converter_type: - mock_converter = mock_converter_type.return_value - mock_new_app = TestWorkflowAssociatedDataFactory.create_app(mode=AppMode.WORKFLOW) - mock_converter.convert_to_workflow.return_value = mock_new_app - - result = workflow_service.convert_to_workflow(app, account, args, session=sqlite_session) - - assert result == mock_new_app - mock_converter.convert_to_workflow.assert_called_once() - - def test_convert_to_workflow_from_completion_app(self, workflow_service: WorkflowService, sqlite_session: Session): - """ - Test convert_to_workflow converts completion app to workflow. - - Similar to chat conversion, but for completion-style apps. - Completion apps are simpler (single prompt-response), so the - conversion creates a basic workflow with fewer nodes. - """ - app = TestWorkflowAssociatedDataFactory.create_app(mode=AppMode.COMPLETION) - account = TestWorkflowAssociatedDataFactory.create_account() - args = {"name": "Converted Workflow"} - - with patch("services.workflow_service.WorkflowConverter") as mock_converter_type: - mock_converter = mock_converter_type.return_value - mock_new_app = TestWorkflowAssociatedDataFactory.create_app(mode=AppMode.WORKFLOW) - mock_converter.convert_to_workflow.return_value = mock_new_app - - result = workflow_service.convert_to_workflow(app, account, args, session=sqlite_session) - - assert result == mock_new_app - - def test_convert_to_workflow_invalid_mode_raises_error( - self, workflow_service: WorkflowService, sqlite_session: Session - ): - """ - Test convert_to_workflow raises error for invalid app mode. - - Only chat and completion apps can be converted to workflows. - Apps that are already workflows or have other modes cannot be converted. - """ - app = TestWorkflowAssociatedDataFactory.create_app(mode=AppMode.WORKFLOW) - account = TestWorkflowAssociatedDataFactory.create_account() - args = {} - - with pytest.raises(ValueError, match="not supported convert to workflow"): - workflow_service.convert_to_workflow(app, account, args, session=sqlite_session) - - -# =========================================================================== -# TestWorkflowServiceCredentialValidation -# Tests for _validate_workflow_credentials and related private helpers -# =========================================================================== - - -@pytest.mark.usefixtures("sqlite_session") -class TestWorkflowServiceCredentialValidation: - """ - Tests for the private credential-validation helpers on WorkflowService. - - These helpers gate `publish_workflow` when `PluginManager` is enabled. - Each test focuses on a distinct branch inside `_validate_workflow_credentials`, - `_validate_llm_model_config`, `_check_default_tool_credential`, and the - load-balancing path. - """ - - @pytest.fixture - def service(self, sqlite_engine: Engine) -> WorkflowService: - return WorkflowService(sessionmaker(bind=sqlite_engine, expire_on_commit=False)) - - @staticmethod - def _make_workflow(nodes: list[dict]) -> Workflow: - return TestWorkflowAssociatedDataFactory.create_workflow( - tenant_id="tenant-1", - app_id="app-1", - graph={"nodes": nodes, "edges": []}, - ) - - # --- _validate_workflow_credentials: tool node (with credential_id) --- - - def test_validate_workflow_credentials_should_check_tool_credential_when_credential_id_present( - self, service: WorkflowService, sqlite_session: Session - ) -> None: - # Arrange - nodes = [ - { - "id": "tool-node", - "data": { - "type": "tool", - "provider_id": "my-provider", - "credential_id": "cred-123", - }, - } - ] - workflow = self._make_workflow(nodes) - - # Act + Assert - with patch("core.helper.credential_utils.check_credential_policy_compliance") as mock_check: - # Should not raise; mock allows the call - service._validate_workflow_credentials(workflow, session=sqlite_session) - mock_check.assert_called_once() - - def test_validate_workflow_credentials_should_check_default_credential_when_no_credential_id( - self, service: WorkflowService, sqlite_session: Session - ) -> None: - # Arrange - nodes = [ - { - "id": "tool-node", - "data": { - "type": "tool", - "provider_id": "my-provider", - # No credential_id — should fall back to default - }, - } - ] - workflow = self._make_workflow(nodes) - - # Act - with patch.object(service, "_check_default_tool_credential") as mock_default: - service._validate_workflow_credentials(workflow, session=sqlite_session) - - # Assert - mock_default.assert_called_once_with("tenant-1", "my-provider", session=sqlite_session) - - def test_validate_workflow_credentials_should_skip_tool_node_without_provider( - self, service: WorkflowService, sqlite_session: Session - ) -> None: - """Tool nodes without a provider_id should be silently skipped.""" - # Arrange - nodes = [{"id": "tool-node", "data": {"type": "tool"}}] - workflow = self._make_workflow(nodes) - - # Act + Assert (no error raised) - with patch.object(service, "_check_default_tool_credential") as mock_default: - service._validate_workflow_credentials(workflow, session=sqlite_session) - mock_default.assert_not_called() - - def test_validate_workflow_credentials_should_validate_llm_node_with_model_config( - self, service: WorkflowService, sqlite_session: Session - ) -> None: - # Arrange - nodes = [ - { - "id": "llm-node", - "data": { - "type": "llm", - "model": {"provider": "openai", "name": "gpt-4"}, - }, - } - ] - workflow = self._make_workflow(nodes) - - # Act - with ( - patch.object(service, "_validate_llm_model_config") as mock_llm, - patch.object(service, "_validate_load_balancing_credentials"), + def test_get_default_block_config_http_request_malformed_config_raises_type_error( + self, workflow_service: WorkflowService ): - service._validate_workflow_credentials(workflow, session=sqlite_session) + with ( + patch( + "services.workflow_service.get_node_type_classes_mapping", + return_value={BuiltinNodeTypes.HTTP_REQUEST: {"latest": HttpRequestNode}}, + ), + patch("services.workflow_service.LATEST_VERSION", "latest"), + ): + with pytest.raises(TypeError, match="http_request_config must be an HttpRequestNodeConfig instance"): + workflow_service.get_default_block_config( + BuiltinNodeTypes.HTTP_REQUEST, + filters={HTTP_REQUEST_CONFIG_FILTER_KEY: "invalid"}, + ) - # Assert - mock_llm.assert_called_once_with("tenant-1", "openai", "gpt-4") + # ==================== Workflow Conversion Tests ==================== + # These tests verify converting basic apps to workflow apps - def test_validate_workflow_credentials_should_use_llm_environment_variable_model( - self, service: WorkflowService, sqlite_session: Session - ) -> None: - workflow = self._make_workflow( - [ + def test_convert_to_workflow_from_chat_app(self, workflow_service: WorkflowService, sqlite_session: Session): + """ + Test convert_to_workflow converts chat app to workflow. + + Allows users to migrate from simple chat apps to advanced workflow apps. + The conversion creates equivalent workflow nodes from the chat configuration, + giving users more control and customization options. + """ + app = TestWorkflowAssociatedDataFactory.create_app(mode=AppMode.CHAT) + account = TestWorkflowAssociatedDataFactory.create_account() + args = { + "name": "Converted Workflow", + "icon_type": "emoji", + "icon": "🤖", + "icon_background": "#FFEAD5", + } + + with patch("services.workflow_service.WorkflowConverter") as mock_converter_type: + mock_converter = mock_converter_type.return_value + mock_new_app = TestWorkflowAssociatedDataFactory.create_app(mode=AppMode.WORKFLOW) + mock_converter.convert_to_workflow.return_value = mock_new_app + + result = workflow_service.convert_to_workflow(app, account, args, session=sqlite_session) + + assert result == mock_new_app + mock_converter.convert_to_workflow.assert_called_once() + + def test_convert_to_workflow_from_completion_app( + self, workflow_service: WorkflowService, sqlite_session: Session + ): + """ + Test convert_to_workflow converts completion app to workflow. + + Similar to chat conversion, but for completion-style apps. + Completion apps are simpler (single prompt-response), so the + conversion creates a basic workflow with fewer nodes. + """ + app = TestWorkflowAssociatedDataFactory.create_app(mode=AppMode.COMPLETION) + account = TestWorkflowAssociatedDataFactory.create_account() + args = {"name": "Converted Workflow"} + + with patch("services.workflow_service.WorkflowConverter") as mock_converter_type: + mock_converter = mock_converter_type.return_value + mock_new_app = TestWorkflowAssociatedDataFactory.create_app(mode=AppMode.WORKFLOW) + mock_converter.convert_to_workflow.return_value = mock_new_app + + result = workflow_service.convert_to_workflow(app, account, args, session=sqlite_session) + + assert result == mock_new_app + + def test_convert_to_workflow_invalid_mode_raises_error( + self, workflow_service: WorkflowService, sqlite_session: Session + ): + """ + Test convert_to_workflow raises error for invalid app mode. + + Only chat and completion apps can be converted to workflows. + Apps that are already workflows or have other modes cannot be converted. + """ + app = TestWorkflowAssociatedDataFactory.create_app(mode=AppMode.WORKFLOW) + account = TestWorkflowAssociatedDataFactory.create_account() + args = {} + + with pytest.raises(ValueError, match="not supported convert to workflow"): + workflow_service.convert_to_workflow(app, account, args, session=sqlite_session) + + # =========================================================================== + # TestWorkflowServiceCredentialValidation + # Tests for _validate_workflow_credentials and related private helpers + # =========================================================================== + + @pytest.mark.usefixtures("sqlite_session") + class TestWorkflowServiceCredentialValidation: + """ + Tests for the private credential-validation helpers on WorkflowService. + + These helpers gate `publish_workflow` when `PluginManager` is enabled. + Each test focuses on a distinct branch inside `_validate_workflow_credentials`, + `_validate_llm_model_config`, `_check_default_tool_credential`, and the + load-balancing path. + """ + + @pytest.fixture + def service(self, sqlite_engine: _workflow_service_Engine) -> WorkflowService: + return WorkflowService(sessionmaker(bind=sqlite_engine, expire_on_commit=False)) + + @staticmethod + def _make_workflow(nodes: list[dict]) -> Workflow: + return TestWorkflowAssociatedDataFactory.create_workflow( + tenant_id="tenant-1", + app_id="app-1", + graph={"nodes": nodes, "edges": []}, + ) + + # --- _validate_workflow_credentials: tool node (with credential_id) --- + + def test_validate_workflow_credentials_should_check_tool_credential_when_credential_id_present( + self, service: WorkflowService, sqlite_session: Session + ) -> None: + # Arrange + nodes = [ { - "id": "llm-node", + "id": "tool-node", "data": { - "type": "llm", - "model": { - "provider": "old-provider", - "name": "old-model", - "mode": "chat", - "completion_params": {"temperature": 0.2}, - }, - "model_selector": ["env", "shared_model"], + "type": "tool", + "provider_id": "my-provider", + "credential_id": "cred-123", }, } ] - ) - workflow.environment_variables = [ - LLMEnvironmentVariable( - name="shared_model", - value={ - "provider": "new-provider", - "name": "new-model", - "mode": "chat", - "completion_params": {"temperature": 0.8}, - }, - ) - ] + workflow = self._make_workflow(nodes) - with ( - patch.object(service, "_validate_llm_model_config") as validate_model, - patch.object(service, "_validate_load_balancing_credentials") as validate_load_balancing, - ): - service._validate_workflow_credentials(workflow, session=sqlite_session) + # Act + Assert + with patch("core.helper.credential_utils.check_credential_policy_compliance") as mock_check: + # Should not raise; mock allows the call + service._validate_workflow_credentials(workflow, session=sqlite_session) + mock_check.assert_called_once() - validate_model.assert_called_once_with("tenant-1", "new-provider", "new-model") - validated_node_data = validate_load_balancing.call_args.args[1] - assert validated_node_data["model"] == { - "provider": "new-provider", - "name": "new-model", - "mode": "chat", - "completion_params": {"temperature": 0.8}, - } - - def test_validate_workflow_credentials_should_reject_llm_environment_variable_mode_mismatch( - self, service: WorkflowService, sqlite_session: Session - ) -> None: - workflow = self._make_workflow( - [ + def test_validate_workflow_credentials_should_check_default_credential_when_no_credential_id( + self, service: WorkflowService, sqlite_session: Session + ) -> None: + # Arrange + nodes = [ { - "id": "llm-node", + "id": "tool-node", "data": { - "type": "llm", - "model": {"provider": "provider", "name": "model", "mode": "chat"}, - "model_selector": ["env", "shared_model"], + "type": "tool", + "provider_id": "my-provider", + # No credential_id — should fall back to default }, } ] - ) - workflow.environment_variables = [ - LLMEnvironmentVariable( - name="shared_model", - value={"provider": "provider", "name": "model", "mode": "completion"}, - ) - ] + workflow = self._make_workflow(nodes) - with pytest.raises(ValueError, match="uses mode 'completion'.*uses mode 'chat'"): - service._validate_workflow_credentials(workflow, session=sqlite_session) - - def test_validate_workflow_credentials_should_raise_for_llm_node_missing_model( - self, service: WorkflowService, sqlite_session: Session - ) -> None: - """LLM nodes without provider AND name should raise ValueError.""" - # Arrange - nodes = [ - { - "id": "llm-node", - "data": {"type": "llm", "model": {"provider": "openai"}}, # name missing - } - ] - workflow = self._make_workflow(nodes) - - # Act + Assert - with pytest.raises(ValueError, match="Missing provider or model configuration"): - service._validate_workflow_credentials(workflow, session=sqlite_session) - - def test_validate_workflow_credentials_should_wrap_unexpected_exception_in_value_error( - self, service: WorkflowService, sqlite_session: Session - ) -> None: - """Non-ValueError exceptions from validation must be re-raised as ValueError.""" - # Arrange - nodes = [ - { - "id": "llm-node", - "data": { - "type": "llm", - "model": {"provider": "openai", "name": "gpt-4"}, - }, - } - ] - workflow = self._make_workflow(nodes) - - # Act + Assert - with patch.object(service, "_validate_llm_model_config", side_effect=RuntimeError("boom")): - with pytest.raises(ValueError, match="boom"): + # Act + with patch.object(service, "_check_default_tool_credential") as mock_default: service._validate_workflow_credentials(workflow, session=sqlite_session) - def test_validate_workflow_credentials_should_validate_agent_node_model( - self, service: WorkflowService, sqlite_session: Session - ) -> None: - # Arrange - nodes = [ - { - "id": "agent-node", - "data": { - "type": "agent", - "agent_parameters": { - "model": {"value": {"provider": "openai", "model": "gpt-4"}}, - "tools": {"value": []}, + # Assert + mock_default.assert_called_once_with("tenant-1", "my-provider", session=sqlite_session) + + def test_validate_workflow_credentials_should_skip_tool_node_without_provider( + self, service: WorkflowService, sqlite_session: Session + ) -> None: + """Tool nodes without a provider_id should be silently skipped.""" + # Arrange + nodes = [{"id": "tool-node", "data": {"type": "tool"}}] + workflow = self._make_workflow(nodes) + + # Act + Assert (no error raised) + with patch.object(service, "_check_default_tool_credential") as mock_default: + service._validate_workflow_credentials(workflow, session=sqlite_session) + mock_default.assert_not_called() + + def test_validate_workflow_credentials_should_validate_llm_node_with_model_config( + self, service: WorkflowService, sqlite_session: Session + ) -> None: + # Arrange + nodes = [ + { + "id": "llm-node", + "data": { + "type": "llm", + "model": {"provider": "openai", "name": "gpt-4"}, }, - }, - } - ] - workflow = self._make_workflow(nodes) + } + ] + workflow = self._make_workflow(nodes) - # Act - with ( - patch.object(service, "_validate_llm_model_config") as mock_llm, - patch.object(service, "_validate_load_balancing_credentials"), - ): - service._validate_workflow_credentials(workflow, session=sqlite_session) - - # Assert - mock_llm.assert_called_once_with("tenant-1", "openai", "gpt-4") - - def test_validate_workflow_credentials_should_validate_agent_tools( - self, service: WorkflowService, sqlite_session: Session - ) -> None: - """Each agent tool with a provider should be checked for credential compliance.""" - # Arrange - nodes = [ - { - "id": "agent-node", - "data": { - "type": "agent", - "agent_parameters": { - "model": {"value": {}}, # no model config - "tools": { - "value": [ - {"provider_name": "provider-a", "credential_id": "cred-a"}, - {"provider_name": "provider-b"}, # uses default - ] - }, - }, - }, - } - ] - workflow = self._make_workflow(nodes) - - # Act - with ( - patch("core.helper.credential_utils.check_credential_policy_compliance") as mock_check, - patch.object(service, "_check_default_tool_credential") as mock_default, - ): - service._validate_workflow_credentials(workflow, session=sqlite_session) - - # Assert - mock_check.assert_called_once() # provider-a has credential_id - mock_default.assert_called_once_with("tenant-1", "provider-b", session=sqlite_session) - - # --- _validate_llm_model_config --- - - def test_validate_llm_model_config_should_raise_value_error_on_failure(self, service: WorkflowService) -> None: - """If ModelManager raises any exception it must be wrapped into ValueError.""" - # Arrange - assembly = MagicMock() - assembly.model_manager.get_model_instance.side_effect = RuntimeError("no key") - - with patch("services.workflow_service.create_plugin_model_assembly", return_value=assembly): - # Act + Assert - with pytest.raises(ValueError, match="Failed to validate LLM model configuration"): - service._validate_llm_model_config("tenant-1", "openai", "gpt-4") - - def test_validate_llm_model_config_success(self, service: WorkflowService) -> None: - """Test success path with ProviderManager and Model entities.""" - mock_model = MagicMock() - mock_model.model = "gpt-4" - mock_model.provider.provider = "openai" - - mock_configs = MagicMock() - mock_configs.get_models.return_value = [mock_model] - assembly = MagicMock() - assembly.provider_manager.get_configurations.return_value = mock_configs - - with patch("services.workflow_service.create_plugin_model_assembly", return_value=assembly): # Act - service._validate_llm_model_config("tenant-1", "openai", "gpt-4") + with ( + patch.object(service, "_validate_llm_model_config") as mock_llm, + patch.object(service, "_validate_load_balancing_credentials"), + ): + service._validate_workflow_credentials(workflow, session=sqlite_session) # Assert - mock_model.raise_for_status.assert_called_once() - assembly.model_manager.get_model_instance.assert_called_once_with( - tenant_id="tenant-1", - provider="openai", - model_type=ModelType.LLM, - model="gpt-4", + mock_llm.assert_called_once_with("tenant-1", "openai", "gpt-4") + + def test_validate_workflow_credentials_should_use_llm_environment_variable_model( + self, service: WorkflowService, sqlite_session: Session + ) -> None: + workflow = self._make_workflow( + [ + { + "id": "llm-node", + "data": { + "type": "llm", + "model": { + "provider": "old-provider", + "name": "old-model", + "mode": "chat", + "completion_params": {"temperature": 0.2}, + }, + "model_selector": ["env", "shared_model"], + }, + } + ] ) + workflow.environment_variables = [ + LLMEnvironmentVariable( + name="shared_model", + value={ + "provider": "new-provider", + "name": "new-model", + "mode": "chat", + "completion_params": {"temperature": 0.8}, + }, + ) + ] - def test_validate_llm_model_config_model_not_found(self, service: WorkflowService) -> None: - """Test ValueError when model is not found in provider configurations.""" - mock_configs = MagicMock() - mock_configs.get_models.return_value = [] # No models - assembly = MagicMock() - assembly.provider_manager.get_configurations.return_value = mock_configs + with ( + patch.object(service, "_validate_llm_model_config") as validate_model, + patch.object(service, "_validate_load_balancing_credentials") as validate_load_balancing, + ): + service._validate_workflow_credentials(workflow, session=sqlite_session) + + validate_model.assert_called_once_with("tenant-1", "new-provider", "new-model") + validated_node_data = validate_load_balancing.call_args.args[1] + assert validated_node_data["model"] == { + "provider": "new-provider", + "name": "new-model", + "mode": "chat", + "completion_params": {"temperature": 0.8}, + } + + def test_validate_workflow_credentials_should_reject_llm_environment_variable_mode_mismatch( + self, service: WorkflowService, sqlite_session: Session + ) -> None: + workflow = self._make_workflow( + [ + { + "id": "llm-node", + "data": { + "type": "llm", + "model": {"provider": "provider", "name": "model", "mode": "chat"}, + "model_selector": ["env", "shared_model"], + }, + } + ] + ) + workflow.environment_variables = [ + LLMEnvironmentVariable( + name="shared_model", + value={"provider": "provider", "name": "model", "mode": "completion"}, + ) + ] + + with pytest.raises(ValueError, match="uses mode 'completion'.*uses mode 'chat'"): + service._validate_workflow_credentials(workflow, session=sqlite_session) + + def test_validate_workflow_credentials_should_raise_for_llm_node_missing_model( + self, service: WorkflowService, sqlite_session: Session + ) -> None: + """LLM nodes without provider AND name should raise ValueError.""" + # Arrange + nodes = [ + { + "id": "llm-node", + "data": {"type": "llm", "model": {"provider": "openai"}}, # name missing + } + ] + workflow = self._make_workflow(nodes) - with patch("services.workflow_service.create_plugin_model_assembly", return_value=assembly): # Act + Assert - with pytest.raises(ValueError, match="Model gpt-4 not found for provider openai"): + with pytest.raises(ValueError, match="Missing provider or model configuration"): + service._validate_workflow_credentials(workflow, session=sqlite_session) + + def test_validate_workflow_credentials_should_wrap_unexpected_exception_in_value_error( + self, service: WorkflowService, sqlite_session: Session + ) -> None: + """Non-ValueError exceptions from validation must be re-raised as ValueError.""" + # Arrange + nodes = [ + { + "id": "llm-node", + "data": { + "type": "llm", + "model": {"provider": "openai", "name": "gpt-4"}, + }, + } + ] + workflow = self._make_workflow(nodes) + + # Act + Assert + with patch.object(service, "_validate_llm_model_config", side_effect=RuntimeError("boom")): + with pytest.raises(ValueError, match="boom"): + service._validate_workflow_credentials(workflow, session=sqlite_session) + + def test_validate_workflow_credentials_should_validate_agent_node_model( + self, service: WorkflowService, sqlite_session: Session + ) -> None: + # Arrange + nodes = [ + { + "id": "agent-node", + "data": { + "type": "agent", + "agent_parameters": { + "model": {"value": {"provider": "openai", "model": "gpt-4"}}, + "tools": {"value": []}, + }, + }, + } + ] + workflow = self._make_workflow(nodes) + + # Act + with ( + patch.object(service, "_validate_llm_model_config") as mock_llm, + patch.object(service, "_validate_load_balancing_credentials"), + ): + service._validate_workflow_credentials(workflow, session=sqlite_session) + + # Assert + mock_llm.assert_called_once_with("tenant-1", "openai", "gpt-4") + + def test_validate_workflow_credentials_should_validate_agent_tools( + self, service: WorkflowService, sqlite_session: Session + ) -> None: + """Each agent tool with a provider should be checked for credential compliance.""" + # Arrange + nodes = [ + { + "id": "agent-node", + "data": { + "type": "agent", + "agent_parameters": { + "model": {"value": {}}, # no model config + "tools": { + "value": [ + {"provider_name": "provider-a", "credential_id": "cred-a"}, + {"provider_name": "provider-b"}, # uses default + ] + }, + }, + }, + } + ] + workflow = self._make_workflow(nodes) + + # Act + with ( + patch("core.helper.credential_utils.check_credential_policy_compliance") as mock_check, + patch.object(service, "_check_default_tool_credential") as mock_default, + ): + service._validate_workflow_credentials(workflow, session=sqlite_session) + + # Assert + mock_check.assert_called_once() # provider-a has credential_id + mock_default.assert_called_once_with("tenant-1", "provider-b", session=sqlite_session) + + # --- _validate_llm_model_config --- + + def test_validate_llm_model_config_should_raise_value_error_on_failure(self, service: WorkflowService) -> None: + """If ModelManager raises any exception it must be wrapped into ValueError.""" + # Arrange + assembly = MagicMock() + assembly.model_manager.get_model_instance.side_effect = RuntimeError("no key") + + with patch("services.workflow_service.create_plugin_model_assembly", return_value=assembly): + # Act + Assert + with pytest.raises(ValueError, match="Failed to validate LLM model configuration"): + service._validate_llm_model_config("tenant-1", "openai", "gpt-4") + + def test_validate_llm_model_config_success(self, service: WorkflowService) -> None: + """Test success path with ProviderManager and Model entities.""" + mock_model = MagicMock() + mock_model.model = "gpt-4" + mock_model.provider.provider = "openai" + + mock_configs = MagicMock() + mock_configs.get_models.return_value = [mock_model] + assembly = MagicMock() + assembly.provider_manager.get_configurations.return_value = mock_configs + + with patch("services.workflow_service.create_plugin_model_assembly", return_value=assembly): + # Act service._validate_llm_model_config("tenant-1", "openai", "gpt-4") - # --- _check_default_tool_credential --- + # Assert + mock_model.raise_for_status.assert_called_once() + assembly.model_manager.get_model_instance.assert_called_once_with( + tenant_id="tenant-1", + provider="openai", + model_type=ModelType.LLM, + model="gpt-4", + ) - def test_check_default_tool_credential_should_silently_pass_when_no_provider_found( - self, service: WorkflowService, sqlite_session: Session - ) -> None: - """Missing BuiltinToolProvider → plugin requires no credentials → no error.""" - service._check_default_tool_credential("tenant-1", "some-provider", session=sqlite_session) - - def test_check_default_tool_credential_should_raise_when_compliance_fails( - self, service: WorkflowService, sqlite_session: Session - ) -> None: - provider = BuiltinToolProvider( - name="API key", - tenant_id="tenant-1", - user_id="user-1", - provider="some-provider", - ) - sqlite_session.add(provider) - sqlite_session.commit() - - with patch("core.helper.credential_utils.check_credential_policy_compliance", side_effect=Exception("denied")): - with pytest.raises(ValueError, match="Failed to validate default credential"): - service._check_default_tool_credential("tenant-1", "some-provider", session=sqlite_session) - - # --- _is_load_balancing_enabled --- - - def test_is_load_balancing_enabled_should_return_false_when_provider_not_found( - self, service: WorkflowService - ) -> None: - with patch("core.provider_manager.ProviderManager.get_configurations") as mock_get_configs: + def test_validate_llm_model_config_model_not_found(self, service: WorkflowService) -> None: + """Test ValueError when model is not found in provider configurations.""" mock_configs = MagicMock() - mock_configs.get.return_value = None # provider not found - mock_get_configs.return_value = mock_configs + mock_configs.get_models.return_value = [] # No models + assembly = MagicMock() + assembly.provider_manager.get_configurations.return_value = mock_configs - # Act - result = service._is_load_balancing_enabled("tenant-1", "openai", "gpt-4") + with patch("services.workflow_service.create_plugin_model_assembly", return_value=assembly): + # Act + Assert + with pytest.raises(ValueError, match="Model gpt-4 not found for provider openai"): + service._validate_llm_model_config("tenant-1", "openai", "gpt-4") - # Assert - assert result is False + # --- _check_default_tool_credential --- - def test_is_load_balancing_enabled_should_return_true_when_setting_enabled(self, service: WorkflowService) -> None: - # Arrange - with patch("core.provider_manager.ProviderManager.get_configurations") as mock_get_configs: - mock_provider_config = MagicMock() - mock_provider_model_setting = MagicMock() - mock_provider_model_setting.load_balancing_enabled = True - mock_provider_config.get_provider_model_setting.return_value = mock_provider_model_setting + def test_check_default_tool_credential_should_silently_pass_when_no_provider_found( + self, service: WorkflowService, sqlite_session: Session + ) -> None: + """Missing BuiltinToolProvider → plugin requires no credentials → no error.""" + service._check_default_tool_credential("tenant-1", "some-provider", session=sqlite_session) - mock_configs = MagicMock() - mock_configs.get.return_value = mock_provider_config - mock_get_configs.return_value = mock_configs + def test_check_default_tool_credential_should_raise_when_compliance_fails( + self, service: WorkflowService, sqlite_session: Session + ) -> None: + provider = BuiltinToolProvider( + name="API key", + tenant_id="tenant-1", + user_id="user-1", + provider="some-provider", + ) + sqlite_session.add(provider) + sqlite_session.commit() - # Act - result = service._is_load_balancing_enabled("tenant-1", "openai", "gpt-4") + with patch( + "core.helper.credential_utils.check_credential_policy_compliance", side_effect=Exception("denied") + ): + with pytest.raises(ValueError, match="Failed to validate default credential"): + service._check_default_tool_credential("tenant-1", "some-provider", session=sqlite_session) - # Assert - assert result is True + # --- _is_load_balancing_enabled --- - def test_is_load_balancing_enabled_should_return_false_on_exception(self, service: WorkflowService) -> None: - """Any exception should be swallowed and return False.""" - # Arrange - with patch("core.provider_manager.ProviderManager.get_configurations", side_effect=RuntimeError("db down")): - # Act - result = service._is_load_balancing_enabled("tenant-1", "openai", "gpt-4") + def test_is_load_balancing_enabled_should_return_false_when_provider_not_found( + self, service: WorkflowService + ) -> None: + with patch("core.provider_manager.ProviderManager.get_configurations") as mock_get_configs: + mock_configs = MagicMock() + mock_configs.get.return_value = None # provider not found + mock_get_configs.return_value = mock_configs - # Assert - assert result is False + # Act + result = service._is_load_balancing_enabled("tenant-1", "openai", "gpt-4") - # --- _get_load_balancing_configs --- + # Assert + assert result is False - def test_get_load_balancing_configs_should_return_empty_list_on_exception( - self, service: WorkflowService, sqlite_session: Session - ) -> None: - """Any exception during LB config retrieval should return an empty list.""" - # Arrange - with patch( - "services.model_load_balancing_service.ModelLoadBalancingService.get_load_balancing_configs", - side_effect=RuntimeError("fail"), - ): - # Act - result = service._get_load_balancing_configs("tenant-1", "openai", "gpt-4", session=sqlite_session) + def test_is_load_balancing_enabled_should_return_true_when_setting_enabled( + self, service: WorkflowService + ) -> None: + # Arrange + with patch("core.provider_manager.ProviderManager.get_configurations") as mock_get_configs: + mock_provider_config = MagicMock() + mock_provider_model_setting = MagicMock() + mock_provider_model_setting.load_balancing_enabled = True + mock_provider_config.get_provider_model_setting.return_value = mock_provider_model_setting - # Assert - assert result == [] + mock_configs = MagicMock() + mock_configs.get.return_value = mock_provider_config + mock_get_configs.return_value = mock_configs - def test_get_load_balancing_configs_should_merge_predefined_and_custom( - self, service: WorkflowService, sqlite_session: Session - ) -> None: - # Arrange - predefined = [{"credential_id": "cred-a"}, {"credential_id": None}] - custom = [{"credential_id": "cred-b"}] - with patch( - "services.model_load_balancing_service.ModelLoadBalancingService.get_load_balancing_configs", - side_effect=[ - (None, predefined), # first call: predefined-model - (None, custom), # second call: custom-model - ], - ): - # Act - result = service._get_load_balancing_configs("tenant-1", "openai", "gpt-4", session=sqlite_session) + # Act + result = service._is_load_balancing_enabled("tenant-1", "openai", "gpt-4") - # Assert — only entries with a credential_id should be returned - assert len(result) == 2 - assert all(c["credential_id"] for c in result) + # Assert + assert result is True - # --- _validate_load_balancing_credentials --- + def test_is_load_balancing_enabled_should_return_false_on_exception(self, service: WorkflowService) -> None: + """Any exception should be swallowed and return False.""" + # Arrange + with patch("core.provider_manager.ProviderManager.get_configurations", side_effect=RuntimeError("db down")): + # Act + result = service._is_load_balancing_enabled("tenant-1", "openai", "gpt-4") - def test_validate_load_balancing_credentials_should_skip_when_no_model_config( - self, service: WorkflowService, sqlite_session: Session - ) -> None: - """Missing provider or model in node_data should be a no-op.""" - # Arrange - node_data: dict[str, Any] = {} # no model key + # Assert + assert result is False - # Act + Assert (no error expected) - service._validate_load_balancing_credentials("tenant-1", node_data, "node-1", session=sqlite_session) + # --- _get_load_balancing_configs --- - def test_validate_load_balancing_credentials_should_skip_when_lb_not_enabled( - self, service: WorkflowService, sqlite_session: Session - ) -> None: - # Arrange - node_data = {"model": {"provider": "openai", "name": "gpt-4"}} + def test_get_load_balancing_configs_should_return_empty_list_on_exception( + self, service: WorkflowService, sqlite_session: Session + ) -> None: + """Any exception during LB config retrieval should return an empty list.""" + # Arrange + with patch( + "services.model_load_balancing_service.ModelLoadBalancingService.get_load_balancing_configs", + side_effect=RuntimeError("fail"), + ): + # Act + result = service._get_load_balancing_configs("tenant-1", "openai", "gpt-4", session=sqlite_session) - # Act + Assert (no error expected) - with patch.object(service, "_is_load_balancing_enabled", return_value=False): + # Assert + assert result == [] + + def test_get_load_balancing_configs_should_merge_predefined_and_custom( + self, service: WorkflowService, sqlite_session: Session + ) -> None: + # Arrange + predefined = [{"credential_id": "cred-a"}, {"credential_id": None}] + custom = [{"credential_id": "cred-b"}] + with patch( + "services.model_load_balancing_service.ModelLoadBalancingService.get_load_balancing_configs", + side_effect=[ + (None, predefined), # first call: predefined-model + (None, custom), # second call: custom-model + ], + ): + # Act + result = service._get_load_balancing_configs("tenant-1", "openai", "gpt-4", session=sqlite_session) + + # Assert — only entries with a credential_id should be returned + assert len(result) == 2 + assert all(c["credential_id"] for c in result) + + # --- _validate_load_balancing_credentials --- + + def test_validate_load_balancing_credentials_should_skip_when_no_model_config( + self, service: WorkflowService, sqlite_session: Session + ) -> None: + """Missing provider or model in node_data should be a no-op.""" + # Arrange + node_data: dict[str, Any] = {} # no model key + + # Act + Assert (no error expected) service._validate_load_balancing_credentials("tenant-1", node_data, "node-1", session=sqlite_session) - def test_validate_load_balancing_credentials_should_raise_when_compliance_fails( - self, service: WorkflowService, sqlite_session: Session - ) -> None: - # Arrange - node_data = {"model": {"provider": "openai", "name": "gpt-4"}} - lb_configs = [{"credential_id": "cred-lb-1"}] + def test_validate_load_balancing_credentials_should_skip_when_lb_not_enabled( + self, service: WorkflowService, sqlite_session: Session + ) -> None: + # Arrange + node_data = {"model": {"provider": "openai", "name": "gpt-4"}} - # Act + Assert - with ( - patch.object(service, "_is_load_balancing_enabled", return_value=True), - patch.object(service, "_get_load_balancing_configs", return_value=lb_configs), - patch( - "core.helper.credential_utils.check_credential_policy_compliance", - side_effect=Exception("policy violation"), - ), - ): - with pytest.raises(ValueError, match="Invalid load balancing credentials"): + # Act + Assert (no error expected) + with patch.object(service, "_is_load_balancing_enabled", return_value=False): service._validate_load_balancing_credentials("tenant-1", node_data, "node-1", session=sqlite_session) - -# =========================================================================== -# TestWorkflowServiceExecutionHelpers -# Tests for _apply_error_strategy, _populate_execution_result, _execute_node_safely -# =========================================================================== - - -class TestWorkflowServiceExecutionHelpers: - """ - Tests for the private execution-result handling methods: - _apply_error_strategy, _populate_execution_result, _execute_node_safely. - """ - - @pytest.fixture - def service(self, sqlite_engine: Engine) -> WorkflowService: - return WorkflowService(sessionmaker(bind=sqlite_engine, expire_on_commit=False)) - - # --- _apply_error_strategy --- - - def test_apply_error_strategy_should_return_exception_status_noderunresult(self, service: WorkflowService) -> None: - # Arrange - node = MagicMock() - node.error_strategy = ErrorStrategy.FAIL_BRANCH - node.default_value_dict = {} - original = NodeRunResult( - status=WorkflowNodeExecutionStatus.FAILED, - error="something went wrong", - error_type="SomeError", - inputs={"x": 1}, - outputs={}, - ) - - # Act - result = service._apply_error_strategy(node, original) - - # Assert - assert result.status == WorkflowNodeExecutionStatus.EXCEPTION - assert result.error == "something went wrong" - assert result.metadata[WorkflowNodeExecutionMetadataKey.ERROR_STRATEGY] == ErrorStrategy.FAIL_BRANCH - - def test_apply_error_strategy_should_include_default_values_for_default_value_strategy( - self, service: WorkflowService - ) -> None: - # Arrange - node = MagicMock() - node.error_strategy = ErrorStrategy.DEFAULT_VALUE - node.default_value_dict = {"output_key": "fallback"} - original = NodeRunResult( - status=WorkflowNodeExecutionStatus.FAILED, - error="err", - ) - - # Act - result = service._apply_error_strategy(node, original) - - # Assert - assert result.outputs.get("output_key") == "fallback" - assert result.status == WorkflowNodeExecutionStatus.EXCEPTION - - # --- _populate_execution_result --- - - def test_populate_execution_result_should_set_succeeded_fields_when_run_succeeded( - self, service: WorkflowService - ) -> None: - # Arrange - node_execution = MagicMock(error=None) - node_run_result = NodeRunResult( - status=WorkflowNodeExecutionStatus.SUCCEEDED, - inputs={"q": "hello"}, - process_data={"steps": 3}, - outputs={"answer": "hi"}, - metadata={WorkflowNodeExecutionMetadataKey.TOTAL_TOKENS: 10}, - ) - - # Act - with patch("services.workflow_service.WorkflowEntry.handle_special_values", side_effect=lambda x: x): - service._populate_execution_result(node_execution, node_run_result, True, None) - - # Assert - assert node_execution.status == WorkflowNodeExecutionStatus.SUCCEEDED - assert node_execution.outputs == {"answer": "hi"} - assert node_execution.error is None # SUCCEEDED status doesn't set error - - def test_populate_execution_result_should_set_failed_status_and_error_when_not_succeeded( - self, service: WorkflowService - ) -> None: - # Arrange - node_execution = MagicMock(error=None) - - # Act - service._populate_execution_result(node_execution, None, False, "catastrophic failure") - - # Assert - assert node_execution.status == WorkflowNodeExecutionStatus.FAILED - assert node_execution.error == "catastrophic failure" - - def test_populate_execution_result_should_set_error_field_for_exception_status( - self, service: WorkflowService - ) -> None: - """A succeeded=True result with EXCEPTION status should still populate the error field.""" - # Arrange - node_execution = MagicMock() - node_run_result = NodeRunResult( - status=WorkflowNodeExecutionStatus.EXCEPTION, - error="constraint violated", - ) - - # Act - with patch("services.workflow_service.WorkflowEntry.handle_special_values", side_effect=lambda x: x): - service._populate_execution_result(node_execution, node_run_result, True, None) - - # Assert - assert node_execution.status == WorkflowNodeExecutionStatus.EXCEPTION - assert node_execution.error == "constraint violated" - - # --- _execute_node_safely --- - - def test_execute_node_safely_should_return_succeeded_result_on_happy_path(self, service: WorkflowService) -> None: - # Arrange - node = MagicMock() - node.error_strategy = None - node_run_result = NodeRunResult(status=WorkflowNodeExecutionStatus.SUCCEEDED) - succeeded_event = NodeRunSucceededEvent( - id=str(uuid.uuid4()), - node_id="node-1", - node_type=BuiltinNodeTypes.LLM, - node_run_result=node_run_result, - start_at=naive_utc_now(), - ) - - def invoke_fn(): - def _gen(): - yield succeeded_event - - return node, _gen() - - # Act - out_node, out_result, run_succeeded, error = service._execute_node_safely(invoke_fn) - - # Assert - assert out_node is node - assert run_succeeded is True - assert error is None - - def test_execute_node_safely_should_return_failed_result_on_failed_event(self, service: WorkflowService) -> None: - # Arrange - node = MagicMock() - node.error_strategy = None - node_run_result = NodeRunResult( - status=WorkflowNodeExecutionStatus.FAILED, - error="node exploded", - ) - failed_event = NodeRunFailedEvent( - id=str(uuid.uuid4()), - node_id="node-1", - node_type=BuiltinNodeTypes.LLM, - node_run_result=node_run_result, - start_at=naive_utc_now(), - error="node exploded", - ) - - def invoke_fn(): - def _gen(): - yield failed_event - - return node, _gen() - - # Act - _, _, run_succeeded, error = service._execute_node_safely(invoke_fn) - - # Assert - assert run_succeeded is False - assert error == "node exploded" - - def test_execute_node_safely_should_handle_workflow_node_run_failed_error(self, service: WorkflowService) -> None: - # Arrange - node = MagicMock() - exc = WorkflowNodeRunFailedError(node, "runtime failure") - - def invoke_fn(): - raise exc - - # Act - out_node, out_result, run_succeeded, error = service._execute_node_safely(invoke_fn) - - # Assert - assert out_node is node - assert out_result is None - assert run_succeeded is False - assert error == "runtime failure" - - def test_execute_node_safely_should_raise_when_no_result_event(self, service: WorkflowService) -> None: - """If the generator produces no NodeRunSucceededEvent/NodeRunFailedEvent, ValueError is expected.""" - # Arrange - node = MagicMock() - node.error_strategy = None - - def invoke_fn(): - def _gen(): - yield from [] - - return node, _gen() - - # Act + Assert - with pytest.raises(ValueError, match="no result returned"): - service._execute_node_safely(invoke_fn) - - # --- _apply_error_strategy with FAIL_BRANCH strategy --- - - def test_execute_node_safely_should_apply_error_strategy_on_failed_status(self, service: WorkflowService) -> None: - # Arrange - node = MagicMock() - node.error_strategy = ErrorStrategy.FAIL_BRANCH - node.default_value_dict = {} - - original_result = NodeRunResult( - status=WorkflowNodeExecutionStatus.FAILED, - error="oops", - error_type="ValueError", - inputs={}, - ) - failed_event = NodeRunFailedEvent( - id=str(uuid.uuid4()), - node_id="node-1", - node_type=BuiltinNodeTypes.LLM, - node_run_result=original_result, - start_at=naive_utc_now(), - error="oops", - ) - - def invoke_fn(): - def _gen(): - yield failed_event - - return node, _gen() - - # Act - _, result, run_succeeded, _ = service._execute_node_safely(invoke_fn) - - # Assert — after applying error strategy status becomes EXCEPTION - assert result is not None - assert result.status == WorkflowNodeExecutionStatus.EXCEPTION - # run_succeeded should be True because EXCEPTION is in the succeeded set - assert run_succeeded is True - - -# =========================================================================== -# TestWorkflowServiceGetNodeLastRun -# Tests for get_node_last_run delegation to repository -# =========================================================================== - - -class TestWorkflowServiceGetNodeLastRun: - @pytest.fixture - def service(self, sqlite_engine: Engine) -> WorkflowService: - return WorkflowService(sessionmaker(bind=sqlite_engine, expire_on_commit=False)) - - def test_get_node_last_run_should_delegate_to_repository(self, service: WorkflowService) -> None: - # Arrange - app = TestWorkflowAssociatedDataFactory.create_app(app_id="app-1", tenant_id="tenant-1") - workflow = TestWorkflowAssociatedDataFactory.create_workflow(workflow_id="wf-1") - expected = MagicMock() - - service._node_execution_service_repo = MagicMock() - service._node_execution_service_repo.get_node_last_execution.return_value = expected - - # Act - result = service.get_node_last_run(app, workflow, "node-42") - - # Assert - assert result is expected - service._node_execution_service_repo.get_node_last_execution.assert_called_once_with( - tenant_id="tenant-1", - app_id="app-1", - workflow_id="wf-1", - node_id="node-42", - ) - - def test_get_node_last_run_should_return_none_when_repository_returns_none(self, service: WorkflowService) -> None: - # Arrange - app = TestWorkflowAssociatedDataFactory.create_app(app_id="a", tenant_id="t") - workflow = TestWorkflowAssociatedDataFactory.create_workflow(workflow_id="w") - service._node_execution_service_repo = MagicMock() - service._node_execution_service_repo.get_node_last_execution.return_value = None - - # Act - result = service.get_node_last_run(app, workflow, "node-x") - - # Assert - assert result is None - - -# =========================================================================== -# TestWorkflowServiceModuleLevelHelpers -# Tests for module-level helper functions exported from workflow_service -# =========================================================================== - - -class TestSetupVariablePool: - """ - Tests for the module-level `_setup_variable_pool` function. - This helper initialises the VariablePool used for single-step workflow execution. - """ - - def _make_workflow(self, workflow_type: str = WorkflowType.WORKFLOW) -> Workflow: - return TestWorkflowAssociatedDataFactory.create_workflow( - workflow_id="wf-1", - app_id="app-1", - workflow_type=WorkflowType(workflow_type), - ) - - def test_setup_variable_pool_should_use_full_system_variables_for_start_node( - self, - ) -> None: - # Arrange - workflow = self._make_workflow() - - # Act - with ( - patch("services.workflow_service.VariablePool") as mock_pool, - patch("services.workflow_service.build_system_variables") as mock_build_system_variables, - patch("services.workflow_service.build_bootstrap_variables") as mock_build_bootstrap_variables, - patch("services.workflow_service.add_variables_to_pool") as mock_add_variables_to_pool, - patch("services.workflow_service.add_node_inputs_to_pool") as mock_add_node_inputs_to_pool, - ): - _setup_variable_pool( - query="hello", - files=[], - user_id="u-1", - user_inputs={"k": "v"}, - workflow=workflow, + def test_validate_load_balancing_credentials_should_raise_when_compliance_fails( + self, service: WorkflowService, sqlite_session: Session + ) -> None: + # Arrange + node_data = {"model": {"provider": "openai", "name": "gpt-4"}} + lb_configs = [{"credential_id": "cred-lb-1"}] + + # Act + Assert + with ( + patch.object(service, "_is_load_balancing_enabled", return_value=True), + patch.object(service, "_get_load_balancing_configs", return_value=lb_configs), + patch( + "core.helper.credential_utils.check_credential_policy_compliance", + side_effect=Exception("policy violation"), + ), + ): + with pytest.raises(ValueError, match="Invalid load balancing credentials"): + service._validate_load_balancing_credentials( + "tenant-1", node_data, "node-1", session=sqlite_session + ) + + # =========================================================================== + # TestWorkflowServiceExecutionHelpers + # Tests for _apply_error_strategy, _populate_execution_result, _execute_node_safely + # =========================================================================== + + class TestWorkflowServiceExecutionHelpers: + """ + Tests for the private execution-result handling methods: + _apply_error_strategy, _populate_execution_result, _execute_node_safely. + """ + + @pytest.fixture + def service(self, sqlite_engine: _workflow_service_Engine) -> WorkflowService: + return WorkflowService(sessionmaker(bind=sqlite_engine, expire_on_commit=False)) + + # --- _apply_error_strategy --- + + def test_apply_error_strategy_should_return_exception_status_noderunresult( + self, service: WorkflowService + ) -> None: + # Arrange + node = MagicMock() + node.error_strategy = ErrorStrategy.FAIL_BRANCH + node.default_value_dict = {} + original = NodeRunResult( + status=WorkflowNodeExecutionStatus.FAILED, + error="something went wrong", + error_type="SomeError", + inputs={"x": 1}, + outputs={}, + ) + + # Act + result = service._apply_error_strategy(node, original) + + # Assert + assert result.status == WorkflowNodeExecutionStatus.EXCEPTION + assert result.error == "something went wrong" + assert result.metadata[WorkflowNodeExecutionMetadataKey.ERROR_STRATEGY] == ErrorStrategy.FAIL_BRANCH + + def test_apply_error_strategy_should_include_default_values_for_default_value_strategy( + self, service: WorkflowService + ) -> None: + # Arrange + node = MagicMock() + node.error_strategy = ErrorStrategy.DEFAULT_VALUE + node.default_value_dict = {"output_key": "fallback"} + original = NodeRunResult( + status=WorkflowNodeExecutionStatus.FAILED, + error="err", + ) + + # Act + result = service._apply_error_strategy(node, original) + + # Assert + assert result.outputs.get("output_key") == "fallback" + assert result.status == WorkflowNodeExecutionStatus.EXCEPTION + + # --- _populate_execution_result --- + + def test_populate_execution_result_should_set_succeeded_fields_when_run_succeeded( + self, service: WorkflowService + ) -> None: + # Arrange + node_execution = MagicMock(error=None) + node_run_result = NodeRunResult( + status=WorkflowNodeExecutionStatus.SUCCEEDED, + inputs={"q": "hello"}, + process_data={"steps": 3}, + outputs={"answer": "hi"}, + metadata={WorkflowNodeExecutionMetadataKey.TOTAL_TOKENS: 10}, + ) + + # Act + with patch("services.workflow_service.WorkflowEntry.handle_special_values", side_effect=lambda x: x): + service._populate_execution_result(node_execution, node_run_result, True, None) + + # Assert + assert node_execution.status == WorkflowNodeExecutionStatus.SUCCEEDED + assert node_execution.outputs == {"answer": "hi"} + assert node_execution.error is None # SUCCEEDED status doesn't set error + + def test_populate_execution_result_should_set_failed_status_and_error_when_not_succeeded( + self, service: WorkflowService + ) -> None: + # Arrange + node_execution = MagicMock(error=None) + + # Act + service._populate_execution_result(node_execution, None, False, "catastrophic failure") + + # Assert + assert node_execution.status == WorkflowNodeExecutionStatus.FAILED + assert node_execution.error == "catastrophic failure" + + def test_populate_execution_result_should_set_error_field_for_exception_status( + self, service: WorkflowService + ) -> None: + """A succeeded=True result with EXCEPTION status should still populate the error field.""" + # Arrange + node_execution = MagicMock() + node_run_result = NodeRunResult( + status=WorkflowNodeExecutionStatus.EXCEPTION, + error="constraint violated", + ) + + # Act + with patch("services.workflow_service.WorkflowEntry.handle_special_values", side_effect=lambda x: x): + service._populate_execution_result(node_execution, node_run_result, True, None) + + # Assert + assert node_execution.status == WorkflowNodeExecutionStatus.EXCEPTION + assert node_execution.error == "constraint violated" + + # --- _execute_node_safely --- + + def test_execute_node_safely_should_return_succeeded_result_on_happy_path( + self, service: WorkflowService + ) -> None: + # Arrange + node = MagicMock() + node.error_strategy = None + node_run_result = NodeRunResult(status=WorkflowNodeExecutionStatus.SUCCEEDED) + succeeded_event = NodeRunSucceededEvent( + id=str(uuid.uuid4()), + node_id="node-1", + node_type=BuiltinNodeTypes.LLM, + node_run_result=node_run_result, + start_at=naive_utc_now(), + ) + + def invoke_fn(): + def _gen(): + yield succeeded_event + + return node, _gen() + + # Act + out_node, out_result, run_succeeded, error = service._execute_node_safely(invoke_fn) + + # Assert + assert out_node is node + assert run_succeeded is True + assert error is None + + def test_execute_node_safely_should_return_failed_result_on_failed_event( + self, service: WorkflowService + ) -> None: + # Arrange + node = MagicMock() + node.error_strategy = None + node_run_result = NodeRunResult( + status=WorkflowNodeExecutionStatus.FAILED, + error="node exploded", + ) + failed_event = NodeRunFailedEvent( + id=str(uuid.uuid4()), + node_id="node-1", + node_type=BuiltinNodeTypes.LLM, + node_run_result=node_run_result, + start_at=naive_utc_now(), + error="node exploded", + ) + + def invoke_fn(): + def _gen(): + yield failed_event + + return node, _gen() + + # Act + _, _, run_succeeded, error = service._execute_node_safely(invoke_fn) + + # Assert + assert run_succeeded is False + assert error == "node exploded" + + def test_execute_node_safely_should_handle_workflow_node_run_failed_error( + self, service: WorkflowService + ) -> None: + # Arrange + node = MagicMock() + exc = WorkflowNodeRunFailedError(node, "runtime failure") + + def invoke_fn(): + raise exc + + # Act + out_node, out_result, run_succeeded, error = service._execute_node_safely(invoke_fn) + + # Assert + assert out_node is node + assert out_result is None + assert run_succeeded is False + assert error == "runtime failure" + + def test_execute_node_safely_should_raise_when_no_result_event(self, service: WorkflowService) -> None: + """If the generator produces no NodeRunSucceededEvent/NodeRunFailedEvent, ValueError is expected.""" + # Arrange + node = MagicMock() + node.error_strategy = None + + def invoke_fn(): + def _gen(): + yield from [] + + return node, _gen() + + # Act + Assert + with pytest.raises(ValueError, match="no result returned"): + service._execute_node_safely(invoke_fn) + + # --- _apply_error_strategy with FAIL_BRANCH strategy --- + + def test_execute_node_safely_should_apply_error_strategy_on_failed_status( + self, service: WorkflowService + ) -> None: + # Arrange + node = MagicMock() + node.error_strategy = ErrorStrategy.FAIL_BRANCH + node.default_value_dict = {} + + original_result = NodeRunResult( + status=WorkflowNodeExecutionStatus.FAILED, + error="oops", + error_type="ValueError", + inputs={}, + ) + failed_event = NodeRunFailedEvent( + id=str(uuid.uuid4()), + node_id="node-1", + node_type=BuiltinNodeTypes.LLM, + node_run_result=original_result, + start_at=naive_utc_now(), + error="oops", + ) + + def invoke_fn(): + def _gen(): + yield failed_event + + return node, _gen() + + # Act + _, result, run_succeeded, _ = service._execute_node_safely(invoke_fn) + + # Assert — after applying error strategy status becomes EXCEPTION + assert result is not None + assert result.status == WorkflowNodeExecutionStatus.EXCEPTION + # run_succeeded should be True because EXCEPTION is in the succeeded set + assert run_succeeded is True + + # =========================================================================== + # TestWorkflowServiceGetNodeLastRun + # Tests for get_node_last_run delegation to repository + # =========================================================================== + + class TestWorkflowServiceGetNodeLastRun: + @pytest.fixture + def service(self, sqlite_engine: _workflow_service_Engine) -> WorkflowService: + return WorkflowService(sessionmaker(bind=sqlite_engine, expire_on_commit=False)) + + def test_get_node_last_run_should_delegate_to_repository(self, service: WorkflowService) -> None: + # Arrange + app = TestWorkflowAssociatedDataFactory.create_app(app_id="app-1", tenant_id="tenant-1") + workflow = TestWorkflowAssociatedDataFactory.create_workflow(workflow_id="wf-1") + expected = MagicMock() + + service._node_execution_service_repo = MagicMock() + service._node_execution_service_repo.get_node_last_execution.return_value = expected + + # Act + result = service.get_node_last_run(app, workflow, "node-42") + + # Assert + assert result is expected + service._node_execution_service_repo.get_node_last_execution.assert_called_once_with( + tenant_id="tenant-1", + app_id="app-1", + workflow_id="wf-1", + node_id="node-42", + ) + + def test_get_node_last_run_should_return_none_when_repository_returns_none( + self, service: WorkflowService + ) -> None: + # Arrange + app = TestWorkflowAssociatedDataFactory.create_app(app_id="a", tenant_id="t") + workflow = TestWorkflowAssociatedDataFactory.create_workflow(workflow_id="w") + service._node_execution_service_repo = MagicMock() + service._node_execution_service_repo.get_node_last_execution.return_value = None + + # Act + result = service.get_node_last_run(app, workflow, "node-x") + + # Assert + assert result is None + + # =========================================================================== + # TestWorkflowServiceModuleLevelHelpers + # Tests for module-level helper functions exported from workflow_service + # =========================================================================== + + class TestSetupVariablePool: + """ + Tests for the module-level `_setup_variable_pool` function. + This helper initialises the VariablePool used for single-step workflow execution. + """ + + def _make_workflow(self, workflow_type: str = WorkflowType.WORKFLOW) -> Workflow: + return TestWorkflowAssociatedDataFactory.create_workflow( + workflow_id="wf-1", + app_id="app-1", + workflow_type=WorkflowType(workflow_type), + ) + + def test_setup_variable_pool_should_use_full_system_variables_for_start_node( + self, + ) -> None: + # Arrange + workflow = self._make_workflow() + + # Act + with ( + patch("services.workflow_service.VariablePool") as mock_pool, + patch("services.workflow_service.build_system_variables") as mock_build_system_variables, + patch("services.workflow_service.build_bootstrap_variables") as mock_build_bootstrap_variables, + patch("services.workflow_service.add_variables_to_pool") as mock_add_variables_to_pool, + patch("services.workflow_service.add_node_inputs_to_pool") as mock_add_node_inputs_to_pool, + ): + _setup_variable_pool( + query="hello", + files=[], + user_id="u-1", + user_inputs={"k": "v"}, + workflow=workflow, + node_id="start-node", + node_type=BuiltinNodeTypes.START, + conversation_id="conv-1", + conversation_variables=[], + ) + + # Assert — start nodes should build bootstrap variables and attach node inputs. + mock_pool.assert_called_once_with() + mock_build_system_variables.assert_called_once() + mock_add_variables_to_pool.assert_called_once_with( + mock_pool.return_value, + mock_build_bootstrap_variables.return_value, + ) + mock_add_node_inputs_to_pool.assert_called_once_with( + mock_pool.return_value, node_id="start-node", - node_type=BuiltinNodeTypes.START, - conversation_id="conv-1", - conversation_variables=[], + inputs={"k": "v"}, ) - # Assert — start nodes should build bootstrap variables and attach node inputs. - mock_pool.assert_called_once_with() - mock_build_system_variables.assert_called_once() - mock_add_variables_to_pool.assert_called_once_with( - mock_pool.return_value, - mock_build_bootstrap_variables.return_value, - ) - mock_add_node_inputs_to_pool.assert_called_once_with( - mock_pool.return_value, - node_id="start-node", - inputs={"k": "v"}, - ) + def test_setup_variable_pool_should_use_default_system_variables_for_non_start_node( + self, + ) -> None: + # Arrange + workflow = self._make_workflow() - def test_setup_variable_pool_should_use_default_system_variables_for_non_start_node( - self, - ) -> None: - # Arrange - workflow = self._make_workflow() + # Act + with ( + patch("services.workflow_service.VariablePool") as mock_pool, + patch("services.workflow_service.default_system_variables") as mock_default_system_variables, + patch("services.workflow_service.build_bootstrap_variables") as mock_build_bootstrap_variables, + patch("services.workflow_service.add_variables_to_pool") as mock_add_variables_to_pool, + patch("services.workflow_service.add_node_inputs_to_pool") as mock_add_node_inputs_to_pool, + ): + _setup_variable_pool( + query="", + files=[], + user_id="u-1", + user_inputs={}, + workflow=workflow, + node_id="llm-node", + node_type=BuiltinNodeTypes.LLM, # not a start/trigger node + conversation_id="conv-1", + conversation_variables=[], + ) - # Act - with ( - patch("services.workflow_service.VariablePool") as mock_pool, - patch("services.workflow_service.default_system_variables") as mock_default_system_variables, - patch("services.workflow_service.build_bootstrap_variables") as mock_build_bootstrap_variables, - patch("services.workflow_service.add_variables_to_pool") as mock_add_variables_to_pool, - patch("services.workflow_service.add_node_inputs_to_pool") as mock_add_node_inputs_to_pool, - ): - _setup_variable_pool( - query="", - files=[], - user_id="u-1", - user_inputs={}, - workflow=workflow, - node_id="llm-node", - node_type=BuiltinNodeTypes.LLM, # not a start/trigger node - conversation_id="conv-1", - conversation_variables=[], + # Assert — default system variables should be used and node inputs should not be added. + mock_default_system_variables.assert_called_once() + mock_pool.assert_called_once_with() + mock_add_variables_to_pool.assert_called_once_with( + mock_pool.return_value, + mock_build_bootstrap_variables.return_value, ) + mock_add_node_inputs_to_pool.assert_not_called() - # Assert — default system variables should be used and node inputs should not be added. - mock_default_system_variables.assert_called_once() - mock_pool.assert_called_once_with() - mock_add_variables_to_pool.assert_called_once_with( - mock_pool.return_value, - mock_build_bootstrap_variables.return_value, - ) - mock_add_node_inputs_to_pool.assert_not_called() + def test_setup_variable_pool_should_set_chatflow_specifics_for_non_workflow_type( + self, + ) -> None: + """For ADVANCED_CHAT workflows on a START node, query/conversation_id/dialogue_count should be set.""" + from models.workflow import WorkflowType - def test_setup_variable_pool_should_set_chatflow_specifics_for_non_workflow_type( - self, - ) -> None: - """For ADVANCED_CHAT workflows on a START node, query/conversation_id/dialogue_count should be set.""" - from models.workflow import WorkflowType + # Arrange + workflow = self._make_workflow(workflow_type=WorkflowType.CHAT) - # Arrange - workflow = self._make_workflow(workflow_type=WorkflowType.CHAT) + # Act + with ( + patch("services.workflow_service.VariablePool") as mock_pool, + patch("services.workflow_service.build_system_variables") as mock_build_system_variables, + patch("services.workflow_service.build_bootstrap_variables"), + patch("services.workflow_service.add_variables_to_pool"), + patch("services.workflow_service.add_node_inputs_to_pool"), + ): + _setup_variable_pool( + query="what is AI?", + files=[], + user_id="u-1", + user_inputs={}, + workflow=workflow, + node_id="start-node", + node_type=BuiltinNodeTypes.START, + conversation_id="conv-abc", + conversation_variables=[], + ) - # Act - with ( - patch("services.workflow_service.VariablePool") as mock_pool, - patch("services.workflow_service.build_system_variables") as mock_build_system_variables, - patch("services.workflow_service.build_bootstrap_variables"), - patch("services.workflow_service.add_variables_to_pool"), - patch("services.workflow_service.add_node_inputs_to_pool"), - ): - _setup_variable_pool( - query="what is AI?", - files=[], - user_id="u-1", - user_inputs={}, - workflow=workflow, - node_id="start-node", - node_type=BuiltinNodeTypes.START, - conversation_id="conv-abc", - conversation_variables=[], - ) + # Assert — chatflow system variables should include query, conversation_id and dialogue_count. + mock_pool.assert_called_once_with() + system_variable_values = mock_build_system_variables.call_args.args[0] + assert system_variable_values["query"] == "what is AI?" + assert system_variable_values["conversation_id"] == "conv-abc" + assert system_variable_values["dialogue_count"] == 1 - # Assert — chatflow system variables should include query, conversation_id and dialogue_count. - mock_pool.assert_called_once_with() - system_variable_values = mock_build_system_variables.call_args.args[0] - assert system_variable_values["query"] == "what is AI?" - assert system_variable_values["conversation_id"] == "conv-abc" - assert system_variable_values["dialogue_count"] == 1 + class TestRebuildSingleFile: + """ + Tests for the module-level `_rebuild_single_file` function. + Ensures correct delegation to `build_from_mapping` / `build_from_mappings`. + """ + def test_rebuild_single_file_should_call_build_from_mapping_for_file_type( + self, + ) -> None: + # Arrange + tenant_id = "tenant-1" + value = {"url": "https://example.com/file.pdf", "type": "document"} + mock_file = MagicMock() -class TestRebuildSingleFile: - """ - Tests for the module-level `_rebuild_single_file` function. - Ensures correct delegation to `build_from_mapping` / `build_from_mappings`. - """ + # Act + with patch("services.workflow_service.build_from_mapping", return_value=mock_file) as mock_build: + result = _rebuild_single_file(tenant_id, value, VariableEntityType.FILE) - def test_rebuild_single_file_should_call_build_from_mapping_for_file_type( - self, - ) -> None: - # Arrange - tenant_id = "tenant-1" - value = {"url": "https://example.com/file.pdf", "type": "document"} - mock_file = MagicMock() + # Assert + assert result is mock_file + mock_build.assert_called_once_with(mapping=value, tenant_id=tenant_id, access_controller=ANY) - # Act - with patch("services.workflow_service.build_from_mapping", return_value=mock_file) as mock_build: - result = _rebuild_single_file(tenant_id, value, VariableEntityType.FILE) + def test_rebuild_single_file_should_raise_when_file_value_not_dict( + self, + ) -> None: + # Arrange + Act + Assert + with pytest.raises(ValueError, match="expected dict for file object"): + _rebuild_single_file("tenant-1", "not-a-dict", VariableEntityType.FILE) - # Assert - assert result is mock_file - mock_build.assert_called_once_with(mapping=value, tenant_id=tenant_id, access_controller=ANY) + def test_rebuild_single_file_should_call_build_from_mappings_for_file_list( + self, + ) -> None: + # Arrange + tenant_id = "tenant-1" + value = [{"url": "https://example.com/a.pdf"}, {"url": "https://example.com/b.pdf"}] + mock_files = [MagicMock(), MagicMock()] - def test_rebuild_single_file_should_raise_when_file_value_not_dict( - self, - ) -> None: - # Arrange + Act + Assert - with pytest.raises(ValueError, match="expected dict for file object"): - _rebuild_single_file("tenant-1", "not-a-dict", VariableEntityType.FILE) + # Act + with patch("services.workflow_service.build_from_mappings", return_value=mock_files) as mock_build: + result = _rebuild_single_file(tenant_id, value, VariableEntityType.FILE_LIST) - def test_rebuild_single_file_should_call_build_from_mappings_for_file_list( - self, - ) -> None: - # Arrange - tenant_id = "tenant-1" - value = [{"url": "https://example.com/a.pdf"}, {"url": "https://example.com/b.pdf"}] - mock_files = [MagicMock(), MagicMock()] + # Assert + assert result is mock_files + mock_build.assert_called_once_with(mappings=value, tenant_id=tenant_id, access_controller=ANY) - # Act - with patch("services.workflow_service.build_from_mappings", return_value=mock_files) as mock_build: - result = _rebuild_single_file(tenant_id, value, VariableEntityType.FILE_LIST) + def test_rebuild_single_file_should_raise_when_file_list_value_not_list( + self, + ) -> None: + # Arrange + Act + Assert + with pytest.raises(ValueError, match="expected list for file list object"): + _rebuild_single_file("tenant-1", "not-a-list", VariableEntityType.FILE_LIST) - # Assert - assert result is mock_files - mock_build.assert_called_once_with(mappings=value, tenant_id=tenant_id, access_controller=ANY) + def test_rebuild_single_file_should_return_empty_list_for_empty_file_list( + self, + ) -> None: + # Arrange + Act + result = _rebuild_single_file("tenant-1", [], VariableEntityType.FILE_LIST) - def test_rebuild_single_file_should_raise_when_file_list_value_not_list( - self, - ) -> None: - # Arrange + Act + Assert - with pytest.raises(ValueError, match="expected list for file list object"): - _rebuild_single_file("tenant-1", "not-a-list", VariableEntityType.FILE_LIST) + # Assert + assert result == [] - def test_rebuild_single_file_should_return_empty_list_for_empty_file_list( - self, - ) -> None: - # Arrange + Act - result = _rebuild_single_file("tenant-1", [], VariableEntityType.FILE_LIST) + def test_rebuild_single_file_should_raise_when_first_element_not_dict( + self, + ) -> None: + # Arrange + Act + Assert + with pytest.raises(ValueError, match="expected dict for first element"): + _rebuild_single_file("tenant-1", ["not-a-dict"], VariableEntityType.FILE_LIST) - # Assert - assert result == [] + class TestRebuildFileForUserInputsInStartNode: + """ + Tests for the module-level `_rebuild_file_for_user_inputs_in_start_node` function. + """ - def test_rebuild_single_file_should_raise_when_first_element_not_dict( - self, - ) -> None: - # Arrange + Act + Assert - with pytest.raises(ValueError, match="expected dict for first element"): - _rebuild_single_file("tenant-1", ["not-a-dict"], VariableEntityType.FILE_LIST) + def _make_start_node_data(self, variables: list) -> MagicMock: + start_data = MagicMock() + start_data.variables = variables + return start_data + def _make_variable(self, name: str, var_type: VariableEntityType) -> MagicMock: + var = MagicMock() + var.variable = name + var.type = var_type + return var -class TestRebuildFileForUserInputsInStartNode: - """ - Tests for the module-level `_rebuild_file_for_user_inputs_in_start_node` function. - """ + def test_rebuild_should_pass_through_non_file_variables( + self, + ) -> None: + # Arrange + text_var = self._make_variable("query", VariableEntityType.TEXT_INPUT) + start_data = self._make_start_node_data([text_var]) + user_inputs = {"query": "hello world"} - def _make_start_node_data(self, variables: list) -> MagicMock: - start_data = MagicMock() - start_data.variables = variables - return start_data - - def _make_variable(self, name: str, var_type: VariableEntityType) -> MagicMock: - var = MagicMock() - var.variable = name - var.type = var_type - return var - - def test_rebuild_should_pass_through_non_file_variables( - self, - ) -> None: - # Arrange - text_var = self._make_variable("query", VariableEntityType.TEXT_INPUT) - start_data = self._make_start_node_data([text_var]) - user_inputs = {"query": "hello world"} - - # Act - result = _rebuild_file_for_user_inputs_in_start_node( - tenant_id="tenant-1", - start_node_data=start_data, - user_inputs=user_inputs, - ) - - # Assert — non-file inputs are untouched - assert result["query"] == "hello world" - - def test_rebuild_should_rebuild_file_variable( - self, - ) -> None: - # Arrange - file_var = self._make_variable("attachment", VariableEntityType.FILE) - start_data = self._make_start_node_data([file_var]) - file_value = {"url": "https://example.com/file.pdf"} - user_inputs = {"attachment": file_value} - mock_file = MagicMock() - - # Act - with patch("services.workflow_service.build_from_mapping", return_value=mock_file): + # Act result = _rebuild_file_for_user_inputs_in_start_node( tenant_id="tenant-1", start_node_data=start_data, user_inputs=user_inputs, ) - # Assert — the dict value should be replaced by the rebuilt File object - assert result["attachment"] is mock_file + # Assert — non-file inputs are untouched + assert result["query"] == "hello world" - def test_rebuild_should_skip_variable_not_in_inputs( - self, - ) -> None: - # Arrange - file_var = self._make_variable("attachment", VariableEntityType.FILE) - start_data = self._make_start_node_data([file_var]) - user_inputs: dict[str, Any] = {} # attachment not provided - - # Act - result = _rebuild_file_for_user_inputs_in_start_node( - tenant_id="tenant-1", - start_node_data=start_data, - user_inputs=user_inputs, - ) - - # Assert — no key should be added for missing inputs - assert "attachment" not in result - - -class TestWorkflowServiceResolveDeliveryMethod: - """ - Tests for the static helper `_resolve_human_input_delivery_method`. - """ - - def _make_method(self, method_id) -> MagicMock: - m = MagicMock() - m.id = method_id - return m - - def test_resolve_delivery_method_should_return_method_when_id_matches(self) -> None: - # Arrange - method_a = self._make_method("method-1") - method_b = self._make_method("method-2") - - # Act - with patch("services.workflow_service.parse_human_input_delivery_methods", return_value=[method_a, method_b]): - result = WorkflowService._resolve_human_input_delivery_method( - node_data=MagicMock(), delivery_method_id="method-2" - ) - - # Assert - assert result is method_b - - def test_resolve_delivery_method_should_return_none_when_no_match(self) -> None: - # Arrange - method_a = self._make_method("method-1") - - # Act - with patch("services.workflow_service.parse_human_input_delivery_methods", return_value=[method_a]): - result = WorkflowService._resolve_human_input_delivery_method( - node_data=MagicMock(), delivery_method_id="does-not-exist" - ) - - # Assert - assert result is None - - def test_resolve_delivery_method_should_return_none_for_empty_methods(self) -> None: - # Act - with patch("services.workflow_service.parse_human_input_delivery_methods", return_value=[]): - result = WorkflowService._resolve_human_input_delivery_method( - node_data=MagicMock(), delivery_method_id="method-1" - ) - - # Assert - assert result is None - - -# =========================================================================== -# TestWorkflowServiceDraftExecution -# Tests for run_draft_workflow_node -# =========================================================================== - - -class TestWorkflowServiceDraftExecution: - @pytest.fixture - def service(self, sqlite_engine: Engine) -> WorkflowService: - return WorkflowService(sessionmaker(bind=sqlite_engine, expire_on_commit=False)) - - def test_run_draft_workflow_node_should_execute_start_node_successfully( - self, service: WorkflowService, sqlite_engine: Engine - ) -> None: - # Arrange - app = TestWorkflowAssociatedDataFactory.create_app(app_id="app-1", tenant_id="tenant-1") - account = TestWorkflowAssociatedDataFactory.create_account(account_id="user-1") - node_id = "start-node" - draft_workflow = TestWorkflowAssociatedDataFactory.create_workflow( - workflow_id="wf-1", - tenant_id="tenant-1", - app_id="app-1", - graph={ - "nodes": [ - { - "id": node_id, - "data": {"type": BuiltinNodeTypes.START, "title": "Start", "variables": []}, - } - ], - "edges": [], - }, - ) - - # Mocking complex dependencies - with ( - patch("services.workflow_service.db", SimpleNamespace(engine=sqlite_engine)), - patch("services.workflow_service.WorkflowDraftVariableService"), - patch("services.workflow_service.StartNodeData"), - patch( - "services.workflow_service._rebuild_file_for_user_inputs_in_start_node", - side_effect=lambda **kwargs: kwargs["user_inputs"], - ), - patch("services.workflow_service._setup_variable_pool"), - patch("services.workflow_service.DraftVarLoader"), - patch("services.workflow_service.WorkflowEntry.single_step_run") as mock_run, - patch("services.workflow_service.DifyCoreRepositoryFactory") as mock_repo_factory, - patch("services.workflow_service.DraftVariableSaver") as mock_saver_cls, - patch("services.workflow_service.storage"), - ): - mock_node = MagicMock() - mock_node.node_type = BuiltinNodeTypes.START - mock_node.title = "Start Node" - mock_run_result = NodeRunResult( - status=WorkflowNodeExecutionStatus.SUCCEEDED, inputs={}, outputs={"result": "ok"} - ) - mock_event = NodeRunSucceededEvent( - id=str(uuid.uuid4()), - node_id="start-node", - node_type=BuiltinNodeTypes.START, - node_run_result=mock_run_result, - start_at=naive_utc_now(), - ) - mock_run.return_value = (mock_node, [mock_event]) - - mock_repo = MagicMock() - mock_repo_factory.create_workflow_node_execution_repository.return_value = mock_repo - - service._node_execution_service_repo = MagicMock() - mock_execution_record = MagicMock() - mock_execution_record.node_type = "start" - mock_execution_record.node_id = "start-node" - mock_execution_record.load_full_outputs.return_value = {} - service._node_execution_service_repo.get_execution_by_id.return_value = mock_execution_record + def test_rebuild_should_rebuild_file_variable( + self, + ) -> None: + # Arrange + file_var = self._make_variable("attachment", VariableEntityType.FILE) + start_data = self._make_start_node_data([file_var]) + file_value = {"url": "https://example.com/file.pdf"} + user_inputs = {"attachment": file_value} + mock_file = MagicMock() # Act - result = service.run_draft_workflow_node( - app_model=app, - draft_workflow=draft_workflow, - account=account, - node_id=node_id, - user_inputs={"key": "val"}, - query="hi", - files=[], - ) + with patch("services.workflow_service.build_from_mapping", return_value=mock_file): + result = _rebuild_file_for_user_inputs_in_start_node( + tenant_id="tenant-1", + start_node_data=start_data, + user_inputs=user_inputs, + ) - # Assert - assert result is not None - mock_run.assert_called_once() - mock_repo.save_synchronously.assert_called_once() - mock_saver_cls.return_value.save.assert_called_once() + # Assert — the dict value should be replaced by the rebuilt File object + assert result["attachment"] is mock_file - def test_run_draft_workflow_node_should_execute_non_start_node_successfully( - self, service: WorkflowService, sqlite_engine: Engine - ) -> None: - # Arrange - app = TestWorkflowAssociatedDataFactory.create_app(app_id="app-1", tenant_id="tenant-1") - account = TestWorkflowAssociatedDataFactory.create_account(account_id="user-1") - node_id = "llm-node" - draft_workflow = TestWorkflowAssociatedDataFactory.create_workflow( - workflow_id="wf-1", - tenant_id="tenant-1", - app_id="app-1", - graph={ - "nodes": [ - { - "id": node_id, - "data": { - "type": BuiltinNodeTypes.LLM, - "title": "LLM", - "model": {"provider": "openai", "name": "gpt-4"}, - }, - } - ], - "edges": [], - }, - ) - - with ( - patch("services.workflow_service.db", SimpleNamespace(engine=sqlite_engine)), - patch("services.workflow_service.WorkflowDraftVariableService"), - patch("services.workflow_service.VariablePool") as mock_pool_cls, - patch("services.workflow_service.build_system_variables") as mock_build_system_variables, - patch("services.workflow_service.build_bootstrap_variables") as mock_build_bootstrap_variables, - patch("services.workflow_service.add_variables_to_pool") as mock_add_variables_to_pool, - patch("services.workflow_service.DraftVarLoader"), - patch("services.workflow_service.WorkflowEntry.single_step_run") as mock_run, - patch("services.workflow_service.DifyCoreRepositoryFactory"), - patch("services.workflow_service.DraftVariableSaver"), - patch("services.workflow_service.storage"), - ): - mock_node = MagicMock() - mock_node.node_type = BuiltinNodeTypes.LLM - mock_node.title = "LLM Node" - mock_run_result = NodeRunResult( - status=WorkflowNodeExecutionStatus.SUCCEEDED, inputs={}, outputs={"result": "ok"} - ) - mock_event = NodeRunSucceededEvent( - id=str(uuid.uuid4()), - node_id="llm-node", - node_type=BuiltinNodeTypes.LLM, - node_run_result=mock_run_result, - start_at=naive_utc_now(), - ) - mock_run.return_value = (mock_node, [mock_event]) - - service._node_execution_service_repo = MagicMock() - mock_execution_record = MagicMock() - mock_execution_record.node_type = "llm" - mock_execution_record.node_id = "llm-node" - mock_execution_record.load_full_outputs.return_value = {"answer": "hello"} - service._node_execution_service_repo.get_execution_by_id.return_value = mock_execution_record + def test_rebuild_should_skip_variable_not_in_inputs( + self, + ) -> None: + # Arrange + file_var = self._make_variable("attachment", VariableEntityType.FILE) + start_data = self._make_start_node_data([file_var]) + user_inputs: dict[str, Any] = {} # attachment not provided # Act - service.run_draft_workflow_node( - app_model=app, - draft_workflow=draft_workflow, - account=account, - node_id=node_id, - user_inputs={}, - query="", - files=None, + result = _rebuild_file_for_user_inputs_in_start_node( + tenant_id="tenant-1", + start_node_data=start_data, + user_inputs=user_inputs, ) + # Assert — no key should be added for missing inputs + assert "attachment" not in result + + class TestWorkflowServiceResolveDeliveryMethod: + """ + Tests for the static helper `_resolve_human_input_delivery_method`. + """ + + def _make_method(self, method_id) -> MagicMock: + m = MagicMock() + m.id = method_id + return m + + def test_resolve_delivery_method_should_return_method_when_id_matches(self) -> None: + # Arrange + method_a = self._make_method("method-1") + method_b = self._make_method("method-2") + + # Act + with patch( + "services.workflow_service.parse_human_input_delivery_methods", return_value=[method_a, method_b] + ): + result = WorkflowService._resolve_human_input_delivery_method( + node_data=MagicMock(), delivery_method_id="method-2" + ) + # Assert - # For non-start nodes, bootstrap variables should be loaded into an empty pool. - mock_pool_cls.assert_called_once_with() - mock_build_system_variables.assert_called_once_with() - mock_build_bootstrap_variables.assert_called_once_with( - system_variables=mock_build_system_variables.return_value, - environment_variables=draft_workflow.environment_variables, - ) - mock_add_variables_to_pool.assert_called_once_with( - mock_pool_cls.return_value, mock_build_bootstrap_variables.return_value - ) + assert result is method_b - @pytest.mark.parametrize( - "repository_type", - [ - SQLAlchemyWorkflowNodeExecutionRepository, - CeleryWorkflowNodeExecutionRepository, - LogstoreWorkflowNodeExecutionRepository, - ], - ) - @pytest.mark.parametrize("status", [WorkflowNodeExecutionStatus.SUCCEEDED, WorkflowNodeExecutionStatus.FAILED]) - def test_single_step_persists_caller_before_body_and_keeps_workspace_on_completion( - self, - service: WorkflowService, - sqlite_engine: Engine, - config_overrides: Callable[..., None], - repository_type: type[SQLAlchemyWorkflowNodeExecutionRepository] - | type[CeleryWorkflowNodeExecutionRepository] - | type[LogstoreWorkflowNodeExecutionRepository], - status: WorkflowNodeExecutionStatus, - ) -> None: - app = TestWorkflowAssociatedDataFactory.create_app() - account = TestWorkflowAssociatedDataFactory.create_account() - node_id = "agent-node" - workflow = TestWorkflowAssociatedDataFactory.create_workflow( - graph={ - "nodes": [{"id": node_id, "data": {"type": "agent", "version": "2", "title": "Agent"}}], - "edges": [], - } - ) - execution_id = str(uuid.uuid4()) - workspace_binding_id = str(uuid.uuid4()) - workflow_binding_id = str(uuid.uuid4()) - started_at = naive_utc_now() - config_overrides(AGENT_BACKEND_USE_FAKE=True) - node = DifyNodeFactory( - init_params=build_test_graph_init_params( - workflow_id=workflow.id, - graph_config=workflow.graph_dict, - tenant_id=app.tenant_id, - app_id=app.id, - user_id=account.id, - ), - runtime_state=RuntimeState(workflow_id=workflow.id, variable_pool=VariablePool(), start_at=0), - ).create_node( - { - "id": node_id, - "data": {"type": "agent", "version": "2", "title": "Agent", "agent_node_kind": "dify_agent"}, - } - ) - config_overrides(LOGSTORE_DUAL_WRITE_ENABLED=False) - logstore_records: list[dict[str, str]] = [] - logstore = MagicMock() - logstore.put_log.side_effect = lambda _name, record: logstore_records.append(dict(record)) - logstore.execute_sql.side_effect = lambda **_kwargs: logstore_records[-1:] - with ( - patch( - "extensions.logstore.repositories.logstore_workflow_node_execution_repository.AliyunLogStore", - return_value=logstore, - ), - patch( - "extensions.logstore.repositories.logstore_api_workflow_node_execution_repository.AliyunLogStore", - return_value=logstore, - ), - ): - repository = repository_type( - session_factory=sqlite_engine, - tenant_id=app.tenant_id, - user=account, - app_id=app.id, - triggered_from=WorkflowNodeExecutionTriggeredFrom.SINGLE_STEP, - ) - if repository_type is LogstoreWorkflowNodeExecutionRepository: - config_overrides( - API_WORKFLOW_NODE_EXECUTION_REPOSITORY=( - "extensions.logstore.repositories.logstore_api_workflow_node_execution_repository." - "LogstoreAPIWorkflowNodeExecutionRepository" - ) - ) - service._node_execution_service_repo = ( - DifyAPIRepositoryFactory.create_api_workflow_node_execution_repository( - session_maker=sessionmaker(sqlite_engine) - ) + def test_resolve_delivery_method_should_return_none_when_no_match(self) -> None: + # Arrange + method_a = self._make_method("method-1") + + # Act + with patch("services.workflow_service.parse_human_input_delivery_methods", return_value=[method_a]): + result = WorkflowService._resolve_human_input_delivery_method( + node_data=MagicMock(), delivery_method_id="does-not-exist" ) - def node_events(variable_pool: VariablePool) -> Generator[NodeEvent]: - yield NodeRunStartedEvent( - id=execution_id, - node_id=node_id, - node_type=BuiltinNodeTypes.AGENT, - node_version="2", - node_title="Agent", - start_at=started_at, + # Assert + assert result is None + + def test_resolve_delivery_method_should_return_none_for_empty_methods(self) -> None: + # Act + with patch("services.workflow_service.parse_human_input_delivery_methods", return_value=[]): + result = WorkflowService._resolve_human_input_delivery_method( + node_data=MagicMock(), delivery_method_id="method-1" + ) + + # Assert + assert result is None + + # =========================================================================== + # TestWorkflowServiceDraftExecution + # Tests for run_draft_workflow_node + # =========================================================================== + + class TestWorkflowServiceDraftExecution: + @pytest.fixture + def service(self, sqlite_engine: _workflow_service_Engine) -> WorkflowService: + return WorkflowService(sessionmaker(bind=sqlite_engine, expire_on_commit=False)) + + def test_run_draft_workflow_node_should_execute_start_node_successfully( + self, service: WorkflowService, sqlite_engine: _workflow_service_Engine + ) -> None: + # Arrange + app = TestWorkflowAssociatedDataFactory.create_app(app_id="app-1", tenant_id="tenant-1") + account = TestWorkflowAssociatedDataFactory.create_account(account_id="user-1") + node_id = "start-node" + draft_workflow = TestWorkflowAssociatedDataFactory.create_workflow( + workflow_id="wf-1", + tenant_id="tenant-1", + app_id="app-1", + graph={ + "nodes": [ + { + "id": node_id, + "data": {"type": BuiltinNodeTypes.START, "title": "Start", "variables": []}, + } + ], + "edges": [], + }, ) - with Session(sqlite_engine) as session: - execution = WorkflowAgentWorkspaceStore._load_execution_by_identity( - session=session, + + # Mocking complex dependencies + with ( + patch("services.workflow_service.db", SimpleNamespace(engine=sqlite_engine)), + patch("services.workflow_service.WorkflowDraftVariableService"), + patch("services.workflow_service.StartNodeData"), + patch( + "services.workflow_service._rebuild_file_for_user_inputs_in_start_node", + side_effect=lambda **kwargs: kwargs["user_inputs"], + ), + patch("services.workflow_service._setup_variable_pool"), + patch("services.workflow_service.DraftVarLoader"), + patch("services.workflow_service.WorkflowEntry.single_step_run") as mock_run, + patch("services.workflow_service.DifyCoreRepositoryFactory") as mock_repo_factory, + patch("services.workflow_service.DraftVariableSaver") as mock_saver_cls, + patch("services.workflow_service.storage"), + ): + mock_node = MagicMock() + mock_node.node_type = BuiltinNodeTypes.START + mock_node.title = "Start Node" + mock_run_result = NodeRunResult( + status=WorkflowNodeExecutionStatus.SUCCEEDED, inputs={}, outputs={"result": "ok"} + ) + mock_event = NodeRunSucceededEvent( + id=str(uuid.uuid4()), + node_id="start-node", + node_type=BuiltinNodeTypes.START, + node_run_result=mock_run_result, + start_at=naive_utc_now(), + ) + mock_run.return_value = (mock_node, [mock_event]) + + mock_repo = MagicMock() + mock_repo_factory.create_workflow_node_execution_repository.return_value = mock_repo + + service._node_execution_service_repo = MagicMock() + mock_execution_record = MagicMock() + mock_execution_record.node_type = "start" + mock_execution_record.node_id = "start-node" + mock_execution_record.load_full_outputs.return_value = {} + service._node_execution_service_repo.get_execution_by_id.return_value = mock_execution_record + + # Act + result = service.run_draft_workflow_node( + app_model=app, + draft_workflow=draft_workflow, + account=account, + node_id=node_id, + user_inputs={"key": "val"}, + query="hi", + files=[], + ) + + # Assert + assert result is not None + mock_run.assert_called_once() + mock_repo.save_synchronously.assert_called_once() + mock_saver_cls.return_value.save.assert_called_once() + + def test_run_draft_workflow_node_should_execute_non_start_node_successfully( + self, service: WorkflowService, sqlite_engine: _workflow_service_Engine + ) -> None: + # Arrange + app = TestWorkflowAssociatedDataFactory.create_app(app_id="app-1", tenant_id="tenant-1") + account = TestWorkflowAssociatedDataFactory.create_account(account_id="user-1") + node_id = "llm-node" + draft_workflow = TestWorkflowAssociatedDataFactory.create_workflow( + workflow_id="wf-1", + tenant_id="tenant-1", + app_id="app-1", + graph={ + "nodes": [ + { + "id": node_id, + "data": { + "type": BuiltinNodeTypes.LLM, + "title": "LLM", + "model": {"provider": "openai", "name": "gpt-4"}, + }, + } + ], + "edges": [], + }, + ) + + with ( + patch("services.workflow_service.db", SimpleNamespace(engine=sqlite_engine)), + patch("services.workflow_service.WorkflowDraftVariableService"), + patch("services.workflow_service.VariablePool") as mock_pool_cls, + patch("services.workflow_service.build_system_variables") as mock_build_system_variables, + patch("services.workflow_service.build_bootstrap_variables") as mock_build_bootstrap_variables, + patch("services.workflow_service.add_variables_to_pool") as mock_add_variables_to_pool, + patch("services.workflow_service.DraftVarLoader"), + patch("services.workflow_service.WorkflowEntry.single_step_run") as mock_run, + patch("services.workflow_service.DifyCoreRepositoryFactory"), + patch("services.workflow_service.DraftVariableSaver"), + patch("services.workflow_service.storage"), + ): + mock_node = MagicMock() + mock_node.node_type = BuiltinNodeTypes.LLM + mock_node.title = "LLM Node" + mock_run_result = NodeRunResult( + status=WorkflowNodeExecutionStatus.SUCCEEDED, inputs={}, outputs={"result": "ok"} + ) + mock_event = NodeRunSucceededEvent( + id=str(uuid.uuid4()), + node_id="llm-node", + node_type=BuiltinNodeTypes.LLM, + node_run_result=mock_run_result, + start_at=naive_utc_now(), + ) + mock_run.return_value = (mock_node, [mock_event]) + + service._node_execution_service_repo = MagicMock() + mock_execution_record = MagicMock() + mock_execution_record.node_type = "llm" + mock_execution_record.node_id = "llm-node" + mock_execution_record.load_full_outputs.return_value = {"answer": "hello"} + service._node_execution_service_repo.get_execution_by_id.return_value = mock_execution_record + + # Act + service.run_draft_workflow_node( + app_model=app, + draft_workflow=draft_workflow, + account=account, + node_id=node_id, + user_inputs={}, + query="", + files=None, + ) + + # Assert + # For non-start nodes, bootstrap variables should be loaded into an empty pool. + mock_pool_cls.assert_called_once_with() + mock_build_system_variables.assert_called_once_with() + mock_build_bootstrap_variables.assert_called_once_with( + system_variables=mock_build_system_variables.return_value, + environment_variables=draft_workflow.environment_variables, + ) + mock_add_variables_to_pool.assert_called_once_with( + mock_pool_cls.return_value, mock_build_bootstrap_variables.return_value + ) + + @pytest.mark.parametrize( + "repository_type", + [ + SQLAlchemyWorkflowNodeExecutionRepository, + CeleryWorkflowNodeExecutionRepository, + LogstoreWorkflowNodeExecutionRepository, + ], + ) + @pytest.mark.parametrize("status", [WorkflowNodeExecutionStatus.SUCCEEDED, WorkflowNodeExecutionStatus.FAILED]) + def test_single_step_persists_caller_before_body_and_keeps_workspace_on_completion( + self, + service: WorkflowService, + sqlite_engine: _workflow_service_Engine, + config_overrides: Callable[..., None], + repository_type: type[SQLAlchemyWorkflowNodeExecutionRepository] + | type[CeleryWorkflowNodeExecutionRepository] + | type[LogstoreWorkflowNodeExecutionRepository], + status: WorkflowNodeExecutionStatus, + ) -> None: + app = TestWorkflowAssociatedDataFactory.create_app() + account = TestWorkflowAssociatedDataFactory.create_account() + node_id = "agent-node" + workflow = TestWorkflowAssociatedDataFactory.create_workflow( + graph={ + "nodes": [{"id": node_id, "data": {"type": "agent", "version": "2", "title": "Agent"}}], + "edges": [], + } + ) + execution_id = str(uuid.uuid4()) + workspace_binding_id = str(uuid.uuid4()) + workflow_binding_id = str(uuid.uuid4()) + started_at = naive_utc_now() + config_overrides(AGENT_BACKEND_USE_FAKE=True) + node = DifyNodeFactory( + init_params=build_test_graph_init_params( + workflow_id=workflow.id, + graph_config=workflow.graph_dict, tenant_id=app.tenant_id, app_id=app.id, - workflow_id=workflow.id, - workflow_run_id=get_system_text(variable_pool, SystemVariableKey.WORKFLOW_EXECUTION_ID), - node_id=node_id, - node_execution_id=execution_id, + user_id=account.id, + ), + runtime_state=RuntimeState(workflow_id=workflow.id, variable_pool=VariablePool(), start_at=0), + ).create_node( + { + "id": node_id, + "data": {"type": "agent", "version": "2", "title": "Agent", "agent_node_kind": "dify_agent"}, + } + ) + config_overrides(LOGSTORE_DUAL_WRITE_ENABLED=False) + logstore_records: list[dict[str, str]] = [] + logstore = MagicMock() + logstore.put_log.side_effect = lambda _name, record: logstore_records.append(dict(record)) + logstore.execute_sql.side_effect = lambda **_kwargs: logstore_records[-1:] + with ( + patch( + "extensions.logstore.repositories.logstore_workflow_node_execution_repository.AliyunLogStore", + return_value=logstore, + ), + patch( + "extensions.logstore.repositories.logstore_api_workflow_node_execution_repository.AliyunLogStore", + return_value=logstore, + ), + ): + repository = repository_type( + session_factory=sqlite_engine, + tenant_id=app.tenant_id, + user=account, + app_id=app.id, + triggered_from=WorkflowNodeExecutionTriggeredFrom.SINGLE_STEP, ) - assert execution.status == WorkflowNodeExecutionStatus.RUNNING - assert execution.node_execution_id == execution_id - assert execution.triggered_from == WorkflowNodeExecutionTriggeredFrom.SINGLE_STEP - assert execution.finished_at is None - execution.agent_workspace_binding_id = workspace_binding_id - execution.process_data = json.dumps({"workflow_agent_binding_id": workflow_binding_id}) - session.commit() + if repository_type is LogstoreWorkflowNodeExecutionRepository: + config_overrides( + API_WORKFLOW_NODE_EXECUTION_REPOSITORY=( + "extensions.logstore.repositories.logstore_api_workflow_node_execution_repository." + "LogstoreAPIWorkflowNodeExecutionRepository" + ) + ) + service._node_execution_service_repo = ( + DifyAPIRepositoryFactory.create_api_workflow_node_execution_repository( + session_maker=sessionmaker(sqlite_engine) + ) + ) - result = NodeRunResult( - status=status, - outputs={"text": "done"} if status == WorkflowNodeExecutionStatus.SUCCEEDED else {}, - process_data={"step": "completed"}, - error="runtime failure" if status == WorkflowNodeExecutionStatus.FAILED else "", - ) - event_kwargs = { - "id": execution_id, - "node_id": node_id, - "node_type": BuiltinNodeTypes.AGENT, - "node_run_result": result, - "start_at": started_at, - } - if status == WorkflowNodeExecutionStatus.FAILED: - yield NodeRunFailedEvent(**event_kwargs, error="runtime failure") + def node_events(variable_pool: VariablePool) -> Generator[NodeEvent]: + yield NodeRunStartedEvent( + id=execution_id, + node_id=node_id, + node_type=BuiltinNodeTypes.AGENT, + node_version="2", + node_title="Agent", + start_at=started_at, + ) + with Session(sqlite_engine) as session: + execution = WorkflowAgentWorkspaceStore._load_execution_by_identity( + session=session, + tenant_id=app.tenant_id, + app_id=app.id, + workflow_id=workflow.id, + workflow_run_id=get_system_text(variable_pool, SystemVariableKey.WORKFLOW_EXECUTION_ID), + node_id=node_id, + node_execution_id=execution_id, + ) + assert execution.status == WorkflowNodeExecutionStatus.RUNNING + assert execution.node_execution_id == execution_id + assert execution.triggered_from == WorkflowNodeExecutionTriggeredFrom.SINGLE_STEP + assert execution.finished_at is None + execution.agent_workspace_binding_id = workspace_binding_id + execution.process_data = json.dumps({"workflow_agent_binding_id": workflow_binding_id}) + session.commit() + + result = NodeRunResult( + status=status, + outputs={"text": "done"} if status == WorkflowNodeExecutionStatus.SUCCEEDED else {}, + process_data={"step": "completed"}, + error="runtime failure" if status == WorkflowNodeExecutionStatus.FAILED else "", + ) + event_kwargs = { + "id": execution_id, + "node_id": node_id, + "node_type": BuiltinNodeTypes.AGENT, + "node_run_result": result, + "start_at": started_at, + } + if status == WorkflowNodeExecutionStatus.FAILED: + yield NodeRunFailedEvent(**event_kwargs, error="runtime failure") + else: + yield NodeRunSucceededEvent(**event_kwargs) + + with ( + patch("services.workflow_service.db", SimpleNamespace(engine=sqlite_engine)), + patch("services.workflow_service.WorkflowDraftVariableService"), + patch("services.workflow_service.DraftVarLoader"), + patch( + "services.workflow_service.WorkflowEntry.single_step_run", + side_effect=lambda **kwargs: (node, node_events(kwargs["variable_pool"])), + ), + patch( + "services.workflow_service.DifyCoreRepositoryFactory.create_workflow_node_execution_repository", + return_value=repository, + ), + patch("services.workflow_service.DraftVariableSaver"), + patch("services.workflow_service.enqueue_draft_node_execution_trace"), + patch( + "core.repositories.celery_workflow_node_execution_repository.save_workflow_node_execution_task.delay" + ) as async_save, + ): + result = service.run_draft_workflow_node( + app_model=app, draft_workflow=workflow, node_id=node_id, user_inputs={}, account=account + ) + + assert result.id == execution_id + assert result.node_execution_id == execution_id + assert result.status == status + assert result.created_at == started_at + assert result.finished_at is not None + if status == WorkflowNodeExecutionStatus.SUCCEEDED: + assert result.outputs_dict == {"text": "done"} + assert result.process_data_dict is not None + assert result.process_data_dict["step"] == "completed" else: - yield NodeRunSucceededEvent(**event_kwargs) + assert result.error == "runtime failure" + with Session(sqlite_engine) as session: + assert list(session.scalars(select(WorkflowNodeExecutionModel.id))) == [execution_id] + persisted = session.get(WorkflowNodeExecutionModel, execution_id) + assert persisted is not None + assert persisted.status == status + assert persisted.agent_workspace_binding_id == workspace_binding_id + assert persisted.process_data_dict is not None + assert persisted.process_data_dict["workflow_agent_binding_id"] == workflow_binding_id + if repository_type is CeleryWorkflowNodeExecutionRepository: + async_save.assert_called_once() + assert async_save.call_args.kwargs["execution_data"]["status"] == status + else: + async_save.assert_not_called() + if repository_type is LogstoreWorkflowNodeExecutionRepository: + assert len(logstore_records) == 1 + assert logstore_records[0]["status"] == status - with ( - patch("services.workflow_service.db", SimpleNamespace(engine=sqlite_engine)), - patch("services.workflow_service.WorkflowDraftVariableService"), - patch("services.workflow_service.DraftVarLoader"), - patch( - "services.workflow_service.WorkflowEntry.single_step_run", - side_effect=lambda **kwargs: (node, node_events(kwargs["variable_pool"])), - ), - patch( - "services.workflow_service.DifyCoreRepositoryFactory.create_workflow_node_execution_repository", - return_value=repository, - ), - patch("services.workflow_service.DraftVariableSaver"), - patch("services.workflow_service.enqueue_draft_node_execution_trace"), - patch( - "core.repositories.celery_workflow_node_execution_repository.save_workflow_node_execution_task.delay" - ) as async_save, - ): - result = service.run_draft_workflow_node( - app_model=app, draft_workflow=workflow, node_id=node_id, user_inputs={}, account=account + # =========================================================================== + # TestWorkflowServiceHumanInputOperations + # Tests for Human Input related methods + # =========================================================================== + + class TestWorkflowServiceHumanInputOperations: + @pytest.fixture + def service(self, sqlite_engine: _workflow_service_Engine) -> WorkflowService: + return WorkflowService(sessionmaker(bind=sqlite_engine, expire_on_commit=False)) + + @staticmethod + def _create_human_input_workflow() -> Workflow: + return TestWorkflowAssociatedDataFactory.create_workflow( + workflow_id="wf-1", + tenant_id="tenant-1", + app_id="app-1", + graph={ + "nodes": [ + { + "id": "node-1", + "data": {"type": BuiltinNodeTypes.HUMAN_INPUT, "title": "Human Input"}, + } + ], + "edges": [], + }, ) - assert result.id == execution_id - assert result.node_execution_id == execution_id - assert result.status == status - assert result.created_at == started_at - assert result.finished_at is not None - if status == WorkflowNodeExecutionStatus.SUCCEEDED: - assert result.outputs_dict == {"text": "done"} - assert result.process_data_dict is not None - assert result.process_data_dict["step"] == "completed" - else: - assert result.error == "runtime failure" - with Session(sqlite_engine) as session: - assert list(session.scalars(select(WorkflowNodeExecutionModel.id))) == [execution_id] - persisted = session.get(WorkflowNodeExecutionModel, execution_id) - assert persisted is not None - assert persisted.status == status - assert persisted.agent_workspace_binding_id == workspace_binding_id - assert persisted.process_data_dict is not None - assert persisted.process_data_dict["workflow_agent_binding_id"] == workflow_binding_id - if repository_type is CeleryWorkflowNodeExecutionRepository: - async_save.assert_called_once() - assert async_save.call_args.kwargs["execution_data"]["status"] == status - else: - async_save.assert_not_called() - if repository_type is LogstoreWorkflowNodeExecutionRepository: - assert len(logstore_records) == 1 - assert logstore_records[0]["status"] == status + def test_get_human_input_form_preview_should_raise_if_workflow_not_init( + self, service: WorkflowService, sqlite_session: Session + ) -> None: + service.get_draft_workflow = MagicMock(return_value=None) + with pytest.raises(ValueError, match="Workflow not initialized"): + service.get_human_input_form_preview( + app_model=TestWorkflowAssociatedDataFactory.create_app(), + account=TestWorkflowAssociatedDataFactory.create_account(), + node_id="node-1", + session=sqlite_session, + ) - -# =========================================================================== -# TestWorkflowServiceHumanInputOperations -# Tests for Human Input related methods -# =========================================================================== - - -class TestWorkflowServiceHumanInputOperations: - @pytest.fixture - def service(self, sqlite_engine: Engine) -> WorkflowService: - return WorkflowService(sessionmaker(bind=sqlite_engine, expire_on_commit=False)) - - @staticmethod - def _create_human_input_workflow() -> Workflow: - return TestWorkflowAssociatedDataFactory.create_workflow( - workflow_id="wf-1", - tenant_id="tenant-1", - app_id="app-1", - graph={ - "nodes": [ - { - "id": "node-1", - "data": {"type": BuiltinNodeTypes.HUMAN_INPUT, "title": "Human Input"}, - } - ], - "edges": [], - }, - ) - - def test_get_human_input_form_preview_should_raise_if_workflow_not_init( - self, service: WorkflowService, sqlite_session: Session - ) -> None: - service.get_draft_workflow = MagicMock(return_value=None) - with pytest.raises(ValueError, match="Workflow not initialized"): - service.get_human_input_form_preview( - app_model=TestWorkflowAssociatedDataFactory.create_app(), - account=TestWorkflowAssociatedDataFactory.create_account(), - node_id="node-1", - session=sqlite_session, + def test_get_human_input_form_preview_should_raise_if_wrong_node_type( + self, service: WorkflowService, sqlite_session: Session + ) -> None: + draft = TestWorkflowAssociatedDataFactory.create_workflow( + graph={"nodes": [{"id": "node-1", "data": {"type": "llm", "title": "LLM"}}], "edges": []} ) + service.get_draft_workflow = MagicMock(return_value=draft) + with pytest.raises(ValueError, match="Node type must be human-input"): + service.get_human_input_form_preview( + app_model=TestWorkflowAssociatedDataFactory.create_app(), + account=TestWorkflowAssociatedDataFactory.create_account(), + node_id="node-1", + session=sqlite_session, + ) - def test_get_human_input_form_preview_should_raise_if_wrong_node_type( - self, service: WorkflowService, sqlite_session: Session - ) -> None: - draft = TestWorkflowAssociatedDataFactory.create_workflow( - graph={"nodes": [{"id": "node-1", "data": {"type": "llm", "title": "LLM"}}], "edges": []} - ) - service.get_draft_workflow = MagicMock(return_value=draft) - with pytest.raises(ValueError, match="Node type must be human-input"): - service.get_human_input_form_preview( - app_model=TestWorkflowAssociatedDataFactory.create_app(), - account=TestWorkflowAssociatedDataFactory.create_account(), - node_id="node-1", - session=sqlite_session, - ) + def test_get_human_input_form_preview_success(self, service: WorkflowService, sqlite_session: Session) -> None: + app_model = TestWorkflowAssociatedDataFactory.create_app(app_id="app-1", tenant_id="tenant-1") + account = TestWorkflowAssociatedDataFactory.create_account(account_id="user-1") + draft = self._create_human_input_workflow() + service.get_draft_workflow = MagicMock(return_value=draft) - def test_get_human_input_form_preview_success(self, service: WorkflowService, sqlite_session: Session) -> None: - app_model = TestWorkflowAssociatedDataFactory.create_app(app_id="app-1", tenant_id="tenant-1") - account = TestWorkflowAssociatedDataFactory.create_account(account_id="user-1") - draft = self._create_human_input_workflow() - service.get_draft_workflow = MagicMock(return_value=draft) + mock_node = MagicMock() + mock_node.render_form_content_before_submission.return_value = "rendered" + mock_node.resolve_default_values.return_value = {"def": 1} + mock_node.title = "Form Title" + mock_node.node_data = MagicMock() - mock_node = MagicMock() - mock_node.render_form_content_before_submission.return_value = "rendered" - mock_node.resolve_default_values.return_value = {"def": 1} - mock_node.title = "Form Title" - mock_node.node_data = MagicMock() + with ( + patch.object(service, "_build_human_input_variable_pool"), + patch("services.workflow_service.HumanInputNode", return_value=mock_node), + patch("services.workflow_service.HumanInputRequired") as mock_required_cls, + ): + service.get_human_input_form_preview( + app_model=app_model, account=account, node_id="node-1", session=sqlite_session + ) + mock_node.render_form_content_before_submission.assert_called_once() + mock_required_cls.return_value.model_dump.assert_called_once() - with ( - patch.object(service, "_build_human_input_variable_pool"), - patch("services.workflow_service.HumanInputNode", return_value=mock_node), - patch("services.workflow_service.HumanInputRequired") as mock_required_cls, - ): - service.get_human_input_form_preview( - app_model=app_model, account=account, node_id="node-1", session=sqlite_session - ) - mock_node.render_form_content_before_submission.assert_called_once() - mock_required_cls.return_value.model_dump.assert_called_once() + def test_submit_human_input_form_preview_success( + self, service: WorkflowService, sqlite_session: Session + ) -> None: + app_model = TestWorkflowAssociatedDataFactory.create_app(app_id="app-1", tenant_id="tenant-1") + account = TestWorkflowAssociatedDataFactory.create_account(account_id="user-1") + draft = self._create_human_input_workflow() + service.get_draft_workflow = MagicMock(return_value=draft) - def test_submit_human_input_form_preview_success(self, service: WorkflowService, sqlite_session: Session) -> None: - app_model = TestWorkflowAssociatedDataFactory.create_app(app_id="app-1", tenant_id="tenant-1") - account = TestWorkflowAssociatedDataFactory.create_account(account_id="user-1") - draft = self._create_human_input_workflow() - service.get_draft_workflow = MagicMock(return_value=draft) + mock_node = MagicMock() + mock_node.node_data = MagicMock() + mock_node.node_data.user_actions = [ + SimpleNamespace(id="submit", title="card_visa_enterprise_001"), + ] + mock_node.node_data.outputs_field_names.return_value = ["field1"] + mock_node.node_data.inputs = [] + mock_node.render_form_content_before_submission.return_value = "Ticket: {{#$output.field1#}}" + mock_node.render_form_content_with_outputs.return_value = "Ticket: val1" - mock_node = MagicMock() - mock_node.node_data = MagicMock() - mock_node.node_data.user_actions = [ - SimpleNamespace(id="submit", title="card_visa_enterprise_001"), - ] - mock_node.node_data.outputs_field_names.return_value = ["field1"] - mock_node.node_data.inputs = [] - mock_node.render_form_content_before_submission.return_value = "Ticket: {{#$output.field1#}}" - mock_node.render_form_content_with_outputs.return_value = "Ticket: val1" + with ( + patch("services.workflow_service.db"), + patch.object(service, "_build_human_input_variable_pool"), + patch("services.workflow_service.HumanInputNode", return_value=mock_node), + patch( + "services.workflow_service.HumanInputService.validate_and_normalize_submission", + return_value={"field1": "val1"}, + ) as mock_validate, + patch("services.workflow_service.DraftVariableSaver") as mock_saver_cls, + ): + result = service.submit_human_input_form_preview( + app_model=app_model, + account=account, + node_id="node-1", + form_inputs={"field1": "val1"}, + action="submit", + session=sqlite_session, + ) + assert result["__action_id"] == "submit" + mock_validate.assert_called_once() + assert result["__action_value"] == "card_visa_enterprise_001" + assert result["__rendered_content"] == "Ticket: val1" + mock_saver_cls.return_value.save.assert_called_once() - with ( - patch("services.workflow_service.db"), - patch.object(service, "_build_human_input_variable_pool"), - patch("services.workflow_service.HumanInputNode", return_value=mock_node), - patch( - "services.workflow_service.HumanInputService.validate_and_normalize_submission", - return_value={"field1": "val1"}, - ) as mock_validate, - patch("services.workflow_service.DraftVariableSaver") as mock_saver_cls, - ): - result = service.submit_human_input_form_preview( - app_model=app_model, - account=account, - node_id="node-1", - form_inputs={"field1": "val1"}, - action="submit", - session=sqlite_session, - ) - assert result["__action_id"] == "submit" - mock_validate.assert_called_once() - assert result["__action_value"] == "card_visa_enterprise_001" - assert result["__rendered_content"] == "Ticket: val1" - mock_saver_cls.return_value.save.assert_called_once() + def test_test_human_input_delivery_success(self, service: WorkflowService, sqlite_session: Session) -> None: + draft = self._create_human_input_workflow() + service.get_draft_workflow = MagicMock(return_value=draft) - def test_test_human_input_delivery_success(self, service: WorkflowService, sqlite_session: Session) -> None: - draft = self._create_human_input_workflow() - service.get_draft_workflow = MagicMock(return_value=draft) - - with ( - patch("services.workflow_service.HumanInputNodeData.model_validate"), - patch.object(service, "_resolve_human_input_delivery_method") as mock_resolve, - patch("services.workflow_service.apply_dify_debug_email_recipient"), - patch.object(service, "_build_human_input_variable_pool"), - patch.object(service, "_build_human_input_node_for_debugging"), - patch.object(service, "_create_human_input_delivery_test_form", return_value=("form-1", [])), - patch("services.workflow_service.HumanInputDeliveryTestService") as mock_test_srv, - ): - mock_resolve.return_value = MagicMock() - service.test_human_input_delivery( - app_model=TestWorkflowAssociatedDataFactory.create_app(), - account=TestWorkflowAssociatedDataFactory.create_account(), - node_id="node-1", - delivery_method_id="method-1", - session=sqlite_session, - ) - mock_test_srv.return_value.send_test.assert_called_once() - - def test_test_human_input_delivery_failure_cases(self, service: WorkflowService, sqlite_session: Session) -> None: - draft = self._create_human_input_workflow() - service.get_draft_workflow = MagicMock(return_value=draft) - - with ( - patch("services.workflow_service.HumanInputNodeData.model_validate"), - patch.object(service, "_resolve_human_input_delivery_method", return_value=None), - ): - with pytest.raises(ValueError, match="Delivery method not found"): + with ( + patch("services.workflow_service.HumanInputNodeData.model_validate"), + patch.object(service, "_resolve_human_input_delivery_method") as mock_resolve, + patch("services.workflow_service.apply_dify_debug_email_recipient"), + patch.object(service, "_build_human_input_variable_pool"), + patch.object(service, "_build_human_input_node_for_debugging"), + patch.object(service, "_create_human_input_delivery_test_form", return_value=("form-1", [])), + patch("services.workflow_service.HumanInputDeliveryTestService") as mock_test_srv, + ): + mock_resolve.return_value = MagicMock() service.test_human_input_delivery( app_model=TestWorkflowAssociatedDataFactory.create_app(), account=TestWorkflowAssociatedDataFactory.create_account(), node_id="node-1", - delivery_method_id="none", + delivery_method_id="method-1", session=sqlite_session, ) + mock_test_srv.return_value.send_test.assert_called_once() - def test_load_email_recipients_parsing_failure(self, service: WorkflowService, sqlite_session: Session) -> None: - """Malformed persisted recipient payloads are skipped instead of aborting delivery tests.""" - recipient = HumanInputFormRecipient( - form_id="form-1", - delivery_id="delivery-1", - recipient_payload="invalid json", - recipient_type=RecipientType.EMAIL_MEMBER, - access_token="recipient-token", - ) - sqlite_session.add(recipient) - sqlite_session.commit() + def test_test_human_input_delivery_failure_cases( + self, service: WorkflowService, sqlite_session: Session + ) -> None: + draft = self._create_human_input_workflow() + service.get_draft_workflow = MagicMock(return_value=draft) - with patch("services.workflow_service.db") as mock_db: - mock_db.engine = sqlite_session.get_bind() - result = WorkflowService._load_email_recipients("form-1") + with ( + patch("services.workflow_service.HumanInputNodeData.model_validate"), + patch.object(service, "_resolve_human_input_delivery_method", return_value=None), + ): + with pytest.raises(ValueError, match="Delivery method not found"): + service.test_human_input_delivery( + app_model=TestWorkflowAssociatedDataFactory.create_app(), + account=TestWorkflowAssociatedDataFactory.create_account(), + node_id="node-1", + delivery_method_id="none", + session=sqlite_session, + ) - assert result == [] + def test_load_email_recipients_parsing_failure(self, service: WorkflowService, sqlite_session: Session) -> None: + """Malformed persisted recipient payloads are skipped instead of aborting delivery tests.""" + recipient = HumanInputFormRecipient( + form_id="form-1", + delivery_id="delivery-1", + recipient_payload="invalid json", + recipient_type=RecipientType.EMAIL_MEMBER, + access_token="recipient-token", + ) + sqlite_session.add(recipient) + sqlite_session.commit() - def test_form_preview_resolves_nested_inputs_without_loading_email_dependencies( - self, service: WorkflowService, sqlite_engine: Engine, sqlite_session: Session, monkeypatch: pytest.MonkeyPatch - ) -> None: - workflow = TestWorkflowAssociatedDataFactory.create_workflow( - graph={ + with patch("services.workflow_service.db") as mock_db: + mock_db.engine = sqlite_session.get_bind() + result = WorkflowService._load_email_recipients("form-1") + + assert result == [] + + def test_form_preview_resolves_nested_inputs_without_loading_email_dependencies( + self, + service: WorkflowService, + sqlite_engine: _workflow_service_Engine, + sqlite_session: Session, + monkeypatch: pytest.MonkeyPatch, + ) -> None: + workflow = TestWorkflowAssociatedDataFactory.create_workflow( + graph={ + "nodes": [ + { + "id": "human", + "data": { + "type": "human-input", + "title": "Human", + "form_content": "Hello {{#source.object.name#}}", + "inputs": [ + { + "type": "paragraph", + "output_variable_name": "comment", + "default": {"type": "variable", "selector": ["source", "object", "default"]}, + } + ], + "delivery_methods": [ + { + "type": "email", + "enabled": True, + "config": { + "recipients": {"items": []}, + "subject": "Review", + "body": "{{#email.missing#}} {{#url#}}", + }, + } + ], + }, + } + ], + "edges": [], + } + ) + monkeypatch.setattr(service, "get_draft_workflow", lambda **_: workflow) + monkeypatch.setattr(type(db), "engine", property(lambda _: sqlite_engine)) + result = service.get_human_input_form_preview( + app_model=TestWorkflowAssociatedDataFactory.create_app(), + account=TestWorkflowAssociatedDataFactory.create_account(), + node_id="human", + inputs={"#source.object.name#": "Alice", "#source.object.default#": "Approved"}, + session=sqlite_session, + ) + + assert result["form_content"] == "Hello Alice" + assert result["resolved_default_values"] == {"comment": "Approved"} + + # =========================================================================== + # TestWorkflowServiceFreeNodeExecution + # Tests for run_free_workflow_node and handle_single_step_result + # =========================================================================== + + class TestWorkflowServiceFreeNodeExecution: + @pytest.fixture + def service(self, sqlite_engine: _workflow_service_Engine) -> WorkflowService: + return WorkflowService(sessionmaker(bind=sqlite_engine, expire_on_commit=False)) + + def test_run_free_workflow_node_success(self, service: WorkflowService) -> None: + node_execution = MagicMock() + with ( + patch.object(service, "_handle_single_step_result", return_value=node_execution), + patch("services.workflow_service.WorkflowEntry.run_free_node"), + ): + result = service.run_free_workflow_node({}, "tenant-1", "user-1", "node-1", {}) + assert result == node_execution + + def test_validate_graph_structure_coexist_error(self, service: WorkflowService) -> None: + graph = { + "nodes": [ + {"data": {"type": "start"}}, + {"data": {"type": "trigger-webhook"}}, # is_trigger_node=True + ] + } + with pytest.raises(ValueError, match="Start node and trigger nodes cannot coexist"): + service.validate_graph_structure(graph) + + def test_validate_features_structure_success(self, service: WorkflowService) -> None: + app = TestWorkflowAssociatedDataFactory.create_app(mode=AppMode.WORKFLOW) + features = {} + with patch("services.workflow_service.WorkflowAppConfigManager.config_validate") as mock_val: + service.validate_features_structure(app, features) + mock_val.assert_called_once() + + def test_validate_features_structure_invalid_mode(self, service: WorkflowService) -> None: + app = TestWorkflowAssociatedDataFactory.create_app() + app.mode = "invalid" + with pytest.raises(ValueError, match="Invalid app mode"): + service.validate_features_structure(app, {}) + + def test_validate_human_input_node_data_error(self, service: WorkflowService) -> None: + with patch("services.workflow_service.HumanInputNodeData.model_validate", side_effect=Exception("error")): + with pytest.raises(ValueError, match="Invalid HumanInput node data"): + service._validate_human_input_node_data({}) + + def test_rebuild_single_file_unreachable(self) -> None: + # Test line 1523 (unreachable) + with pytest.raises(Exception, match="unreachable"): + _rebuild_single_file("tenant-1", {}, cast(Any, "invalid_type")) + + def test_build_human_input_node_for_debugging(self, service: WorkflowService) -> None: + """Cover _build_human_input_node_for_debugging.""" + workflow = TestWorkflowAssociatedDataFactory.create_workflow() + account = TestWorkflowAssociatedDataFactory.create_account() + node_config = {"id": "n-1", "data": {"type": BuiltinNodeTypes.HUMAN_INPUT, "title": "Human Input"}} + variable_pool = MagicMock() + node_data = MagicMock() + node_data.title = "Human Input" + + with ( + patch( + "services.workflow_service.adapt_human_input_node_data_for_graph", + return_value=sentinel.adapted_node_data, + ) as mock_adapt_node_data, + patch("services.workflow_service.HumanInputNodeData.model_validate", return_value=node_data), + ): + node = service._build_human_input_node_for_debugging( + workflow=workflow, account=account, node_config=node_config, variable_pool=variable_pool + ) + mock_adapt_node_data.assert_called_once_with(node_config["data"]) + assert node.node_id == "n-1" + assert node.title == "Human Input" + assert node.node_data is node_data + assert node.variable_pool is variable_pool + + @staticmethod + @pytest.mark.parametrize( + "upsert", + [_env("new-id", "TAKEN"), _env("other-id", "TAKEN")], + ids=["new id reuses a name", "rename onto another variable's name"], + ) + def test_environment_patch_rejects_a_name_another_id_uses(upsert: StringVariable) -> None: + current = [_env("taken-id", "TAKEN"), _env("other-id", "OTHER")] + with pytest.raises(ValueError, match="Environment variable names must be unique"): + _merge_environment_variable_patch(current, [upsert], []) + + +"""Tests for human-input delivery with a persisted draft workflow.""" + + +TENANT_ID = "11111111-1111-1111-1111-111111111111" +APP_ID = "22222222-2222-2222-2222-222222222222" +ACCOUNT_ID = "33333333-3333-3333-3333-333333333333" + + +def _make_service(sqlite_session: Session) -> WorkflowService: + return WorkflowService( + session_maker=sessionmaker(bind=sqlite_session.get_bind(), expire_on_commit=False), + ) + + +def _app() -> App: + return make_app( + app_id=APP_ID, + tenant_id=TENANT_ID, + mode=AppMode.WORKFLOW, + icon_type=None, + enable_site=False, + enable_api=False, + ) + + +def _account() -> _workflow_human_input_delivery_Account: + return make_account(account_id=ACCOUNT_ID, name="Test User", email="test@example.com") + + +def _build_node_config(delivery_methods: list[EmailDeliveryMethod]) -> dict[str, object]: + return { + "id": "node-1", + "data": HumanInputNodeData( + title="Human Input", + type=BuiltinNodeTypes.HUMAN_INPUT, + delivery_methods=delivery_methods, + form_content="Test content", + inputs=[], + user_actions=[], + ), + } + + +def _persist_workflow(sqlite_session: Session, delivery_methods: list[EmailDeliveryMethod]) -> None: + node_config = _build_node_config(delivery_methods) + node_data = node_config["data"] + assert isinstance(node_data, HumanInputNodeData) + workflow = Workflow.new( + tenant_id=TENANT_ID, + app_id=APP_ID, + type=WorkflowType.WORKFLOW, + version=Workflow.VERSION_DRAFT, + graph=json.dumps( + { "nodes": [ { - "id": "human", - "data": { - "type": "human-input", - "title": "Human", - "form_content": "Hello {{#source.object.name#}}", - "inputs": [ - { - "type": "paragraph", - "output_variable_name": "comment", - "default": {"type": "variable", "selector": ["source", "object", "default"]}, - } - ], - "delivery_methods": [ - { - "type": "email", - "enabled": True, - "config": { - "recipients": {"items": []}, - "subject": "Review", - "body": "{{#email.missing#}} {{#url#}}", - }, - } - ], - }, + "id": node_config["id"], + "data": node_data.model_dump(mode="json"), } ], "edges": [], } + ), + features="{}", + created_by=ACCOUNT_ID, + environment_variables=[], + conversation_variables=[], + rag_pipeline_variables=[], + ) + sqlite_session.add(workflow) + sqlite_session.commit() + sqlite_session.expunge_all() + + +def _make_email_method(enabled: bool = True, debug_mode: bool = False) -> EmailDeliveryMethod: + return EmailDeliveryMethod( + id=uuid.uuid4(), + enabled=enabled, + config=EmailDeliveryConfig( + recipients=EmailRecipients( + include_bound_group=False, + items=[ExternalRecipient(email="tester@example.com")], + ), + subject="Test subject", + body="Test body", + debug_mode=debug_mode, + ), + ) + + +class TestWorkflowHumanInputDeliveryCases: + @staticmethod + @pytest.mark.parametrize("sqlite_session", [(Workflow,)], indirect=True) + def test_human_input_delivery_requires_draft_workflow(sqlite_session: Session): + service = _make_service(sqlite_session) + app_model = _app() + account = _account() + + with pytest.raises(ValueError, match="Workflow not initialized"): + service.test_human_input_delivery( + app_model=app_model, + account=account, + node_id="node-1", + delivery_method_id="delivery-1", + session=sqlite_session, + ) + assert sqlite_session.in_transaction() + + @staticmethod + @pytest.mark.parametrize("sqlite_session", [(Workflow,)], indirect=True) + def test_human_input_delivery_allows_disabled_method( + monkeypatch: pytest.MonkeyPatch, + sqlite_session: Session, + ): + service = _make_service(sqlite_session) + delivery_method = _make_email_method(enabled=False) + _persist_workflow(sqlite_session, [delivery_method]) + service._build_human_input_variable_pool = MagicMock(return_value=MagicMock()) # type: ignore[attr-defined] + node_stub = MagicMock() + node_stub.render_form_content_before_submission.return_value = "rendered" + node_stub.resolve_default_values.return_value = {} + service._build_human_input_node_for_debugging = MagicMock(return_value=node_stub) # type: ignore[attr-defined] + service._create_human_input_delivery_test_form = MagicMock( # type: ignore[attr-defined] + return_value=("form-1", []) ) - monkeypatch.setattr(service, "get_draft_workflow", lambda **_: workflow) - monkeypatch.setattr(type(db), "engine", property(lambda _: sqlite_engine)) - result = service.get_human_input_form_preview( - app_model=TestWorkflowAssociatedDataFactory.create_app(), - account=TestWorkflowAssociatedDataFactory.create_account(), - node_id="human", - inputs={"#source.object.name#": "Alice", "#source.object.default#": "Approved"}, + + test_service_instance = MagicMock() + monkeypatch.setattr( + workflow_service_module, + "HumanInputDeliveryTestService", + MagicMock(return_value=test_service_instance), + ) + + app_model = _app() + account = _account() + + service.test_human_input_delivery( + app_model=app_model, + account=account, + node_id="node-1", + delivery_method_id=str(delivery_method.id), session=sqlite_session, ) - assert result["form_content"] == "Hello Alice" - assert result["resolved_default_values"] == {"comment": "Approved"} + test_service_instance.send_test.assert_called_once() + assert sqlite_session.in_transaction() + + @staticmethod + @pytest.mark.parametrize("sqlite_session", [(Workflow,)], indirect=True) + def test_human_input_delivery_dispatches_to_test_service( + monkeypatch: pytest.MonkeyPatch, + sqlite_session: Session, + ): + service = _make_service(sqlite_session) + delivery_method = _make_email_method(enabled=True) + _persist_workflow(sqlite_session, [delivery_method]) + service._build_human_input_variable_pool = MagicMock(return_value=MagicMock()) # type: ignore[attr-defined] + node_stub = MagicMock() + node_stub.render_form_content_before_submission.return_value = "rendered" + node_stub.resolve_default_values.return_value = {} + service._build_human_input_node_for_debugging = MagicMock(return_value=node_stub) # type: ignore[attr-defined] + service._create_human_input_delivery_test_form = MagicMock( # type: ignore[attr-defined] + return_value=("form-1", []) + ) + + test_service_instance = MagicMock() + monkeypatch.setattr( + workflow_service_module, + "HumanInputDeliveryTestService", + MagicMock(return_value=test_service_instance), + ) + + app_model = _app() + account = _account() + + service.test_human_input_delivery( + app_model=app_model, + account=account, + node_id="node-1", + delivery_method_id=str(delivery_method.id), + inputs={"#node-1.output#": "value"}, + session=sqlite_session, + ) + + pool_args = service._build_human_input_variable_pool.call_args.kwargs + assert pool_args["manual_inputs"] == {"#node-1.output#": "value"} + test_service_instance.send_test.assert_called_once() + assert sqlite_session.in_transaction() + + @staticmethod + @pytest.mark.parametrize("sqlite_session", [(Workflow,)], indirect=True) + def test_human_input_delivery_debug_mode_overrides_recipients( + monkeypatch: pytest.MonkeyPatch, + sqlite_session: Session, + ): + service = _make_service(sqlite_session) + delivery_method = _make_email_method(enabled=True, debug_mode=True) + _persist_workflow(sqlite_session, [delivery_method]) + service._build_human_input_variable_pool = MagicMock(return_value=MagicMock()) # type: ignore[attr-defined] + node_stub = MagicMock() + node_stub.render_form_content_before_submission.return_value = "rendered" + node_stub.resolve_default_values.return_value = {} + service._build_human_input_node_for_debugging = MagicMock(return_value=node_stub) # type: ignore[attr-defined] + service._create_human_input_delivery_test_form = MagicMock( # type: ignore[attr-defined] + return_value=("form-1", []) + ) + + test_service_instance = MagicMock() + monkeypatch.setattr( + workflow_service_module, + "HumanInputDeliveryTestService", + MagicMock(return_value=test_service_instance), + ) + + app_model = _app() + account = _account() + + service.test_human_input_delivery( + app_model=app_model, + account=account, + node_id="node-1", + delivery_method_id=str(delivery_method.id), + session=sqlite_session, + ) + + test_service_instance.send_test.assert_called_once() + sent_method = test_service_instance.send_test.call_args.kwargs["method"] + assert isinstance(sent_method, EmailDeliveryMethod) + assert sent_method.config.debug_mode is True + assert sent_method.config.recipients.include_bound_group is False + assert len(sent_method.config.recipients.items) == 1 + recipient = sent_method.config.recipients.items[0] + assert isinstance(recipient, MemberRecipient) + assert recipient.reference_id == account.id + assert sqlite_session.in_transaction() -# =========================================================================== -# TestWorkflowServiceFreeNodeExecution -# Tests for run_free_workflow_node and handle_single_step_result -# =========================================================================== +"""Agent preview callers are committed before execution and retain their identity.""" -class TestWorkflowServiceFreeNodeExecution: +Factory = TestWorkflowAssociatedDataFactory + + +class PreviewCancelled(BaseException): + """Cancellation can bypass Exception handlers, like a stopped greenlet.""" + + +class TestWorkflowAgentPreviewIdentityCases: + @staticmethod @pytest.fixture - def service(self, sqlite_engine: Engine) -> WorkflowService: - return WorkflowService(sessionmaker(bind=sqlite_engine, expire_on_commit=False)) + def preview_flask_app( + sqlite_engine: _workflow_agent_preview_identity_Engine, monkeypatch: pytest.MonkeyPatch + ) -> Iterator[Flask]: + apply_config_overrides(monkeypatch, AGENT_BACKEND_BASE_URL="https://agent.test") + app = Flask(__name__) + app.config["SQLALCHEMY_DATABASE_URI"] = sqlite_engine.url.render_as_string(hide_password=False) + db.init_app(app) + with app.app_context(): + yield app - def test_run_free_workflow_node_success(self, service: WorkflowService) -> None: - node_execution = MagicMock() - with ( - patch.object(service, "_handle_single_step_result", return_value=node_execution), - patch("services.workflow_service.WorkflowEntry.run_free_node"), - ): - result = service.run_free_workflow_node({}, "tenant-1", "user-1", "node-1", {}) - assert result == node_execution + @staticmethod + @pytest.mark.parametrize("outcome", ["success", "failure", "exception", "cancel"]) + @pytest.mark.usefixtures("preview_flask_app") + def test_agent_preview_commits_runtime_identity_and_updates_same_row( + sqlite_engine: _workflow_agent_preview_identity_Engine, outcome: str + ) -> None: + tenant_id, app_id, workflow_id, account_id = (str(uuid4()) for _ in range(4)) + node_id = "agent-preview" + app = Factory.create_app(app_id=app_id, tenant_id=tenant_id) + account = Factory.create_account(account_id=account_id) + workflow = Factory.create_workflow( + workflow_id=workflow_id, + tenant_id=tenant_id, + app_id=app_id, + graph={ + "nodes": [ + { + "id": node_id, + "data": {"type": "agent", "version": "2", "agent_node_kind": "dify_agent", "title": "Agent"}, + } + ], + "edges": [], + }, + ) + sessions = sessionmaker(sqlite_engine, expire_on_commit=False) + with sessions.begin() as session: + session.add_all([app, account, workflow]) + repository = SQLAlchemyWorkflowNodeExecutionRepository( + sessions, tenant_id, account, app_id, WorkflowNodeExecutionTriggeredFrom.SINGLE_STEP + ) + observed_created_at = [] + execution_ids: list[str] = [] - def test_validate_graph_structure_coexist_error(self, service: WorkflowService) -> None: - graph = { - "nodes": [ - {"data": {"type": "start"}}, - {"data": {"type": "trigger-webhook"}}, # is_trigger_node=True - ] - } - with pytest.raises(ValueError, match="Start node and trigger nodes cannot coexist"): - service.validate_graph_structure(graph) - - def test_validate_features_structure_success(self, service: WorkflowService) -> None: - app = TestWorkflowAssociatedDataFactory.create_app(mode=AppMode.WORKFLOW) - features = {} - with patch("services.workflow_service.WorkflowAppConfigManager.config_validate") as mock_val: - service.validate_features_structure(app, features) - mock_val.assert_called_once() - - def test_validate_features_structure_invalid_mode(self, service: WorkflowService) -> None: - app = TestWorkflowAssociatedDataFactory.create_app() - app.mode = "invalid" - with pytest.raises(ValueError, match="Invalid app mode"): - service.validate_features_structure(app, {}) - - def test_validate_human_input_node_data_error(self, service: WorkflowService) -> None: - with patch("services.workflow_service.HumanInputNodeData.model_validate", side_effect=Exception("error")): - with pytest.raises(ValueError, match="Invalid HumanInput node data"): - service._validate_human_input_node_data({}) - - def test_rebuild_single_file_unreachable(self) -> None: - # Test line 1523 (unreachable) - with pytest.raises(Exception, match="unreachable"): - _rebuild_single_file("tenant-1", {}, cast(Any, "invalid_type")) - - def test_build_human_input_node_for_debugging(self, service: WorkflowService) -> None: - """Cover _build_human_input_node_for_debugging.""" - workflow = TestWorkflowAssociatedDataFactory.create_workflow() - account = TestWorkflowAssociatedDataFactory.create_account() - node_config = {"id": "n-1", "data": {"type": BuiltinNodeTypes.HUMAN_INPUT, "title": "Human Input"}} - variable_pool = MagicMock() - node_data = MagicMock() - node_data.title = "Human Input" + def events(node: DifyAgentNode) -> Generator[NodeRunStartedEvent | NodeRunSucceededEvent | NodeRunFailedEvent]: + execution_id = node.execution_id + execution_ids.append(execution_id) + started_at = naive_utc_now() + # Node.run now emits its typed start before entering the node body. The + # service observer must commit that event before resuming this generator. + yield NodeRunStartedEvent( + id=execution_id, + node_id=node_id, + node_type="agent", + node_version="2", + node_title="Agent", + start_at=started_at, + ) + # A separate transaction must see the committed caller before any + # Agent workspace or tracing request can be dispatched. + variable_pool = node.runtime_state.variable_pool + workflow_run_id = get_system_text(variable_pool, SystemVariableKey.WORKFLOW_EXECUTION_ID) + assert workflow_run_id is None + with sessions.begin() as session: + caller = WorkflowAgentWorkspaceStore._load_execution_by_identity( + session=session, + tenant_id=tenant_id, + app_id=app_id, + workflow_id=workflow_id, + workflow_run_id=workflow_run_id, + node_id=node_id, + node_execution_id=execution_id, + ) + assert caller.id == caller.node_execution_id == execution_id + assert caller.status == WorkflowNodeExecutionStatus.RUNNING + assert caller.finished_at is None + observed_created_at.append(caller.created_at) + # Session materialization writes these independently; final persistence + # must not discard either piece of caller ownership. + caller.agent_workspace_binding_id = "workspace-binding" + caller.process_data = json.dumps({"workflow_agent_binding_id": "workflow-binding"}) + if outcome == "exception": + raise RuntimeError("private runtime error details") + if outcome == "cancel": + raise PreviewCancelled() + status = ( + WorkflowNodeExecutionStatus.SUCCEEDED if outcome == "success" else WorkflowNodeExecutionStatus.FAILED + ) + result = NodeRunResult( + status=status, outputs={"answer": "done"}, error="" if outcome == "success" else "failed" + ) + if outcome == "success": + yield NodeRunSucceededEvent( + id=execution_id, node_id=node_id, node_type="agent", node_run_result=result, start_at=started_at + ) + else: + yield NodeRunFailedEvent( + id=execution_id, + node_id=node_id, + node_type="agent", + node_run_result=result, + start_at=started_at, + error="failed", + ) + service = WorkflowService(sessions) with ( patch( - "services.workflow_service.adapt_human_input_node_data_for_graph", - return_value=sentinel.adapted_node_data, - ) as mock_adapt_node_data, - patch("services.workflow_service.HumanInputNodeData.model_validate", return_value=node_data), + "services.workflow_service.WorkflowDraftVariableService.prefill_conversation_variable_default_values" + ), + patch.object(DifyAgentNode, "run", events), + patch( + "services.workflow_service.DifyCoreRepositoryFactory.create_workflow_node_execution_repository", + return_value=repository, + ), + patch("services.workflow_service.DraftVariableSaver.save"), + patch("services.workflow_service.enqueue_draft_node_execution_trace"), ): - node = service._build_human_input_node_for_debugging( - workflow=workflow, account=account, node_config=node_config, variable_pool=variable_pool + if outcome in {"exception", "cancel"}: + with pytest.raises(RuntimeError if outcome == "exception" else PreviewCancelled): + service.run_draft_workflow_node(app, workflow, node_id, {}, account) + else: + response = service.run_draft_workflow_node(app, workflow, node_id, {}, account) + assert response.id == response.node_execution_id == execution_ids[0] + + assert len(execution_ids) == 1 + execution_id = execution_ids[0] + with sessions() as session: + assert session.scalar(select(func.count()).select_from(WorkflowNodeExecutionModel)) == 1 + row = session.get(WorkflowNodeExecutionModel, execution_id) + assert row is not None + assert row.id == row.node_execution_id == execution_id + assert row.workflow_run_id is None + assert row.workflow_id == workflow_id + assert row.tenant_id == tenant_id + assert row.app_id == app_id + assert row.created_at == observed_created_at[0] + assert row.finished_at is not None + assert row.elapsed_time >= 0 + assert row.agent_workspace_binding_id == "workspace-binding" + assert row.process_data_dict is not None + assert row.process_data_dict["workflow_agent_binding_id"] == "workflow-binding" + assert row.status == ( + WorkflowNodeExecutionStatus.SUCCEEDED if outcome == "success" else WorkflowNodeExecutionStatus.FAILED ) - mock_adapt_node_data.assert_called_once_with(node_config["data"]) - assert node.node_id == "n-1" - assert node.title == "Human Input" - assert node.node_data is node_data - assert node.variable_pool is variable_pool + if outcome in {"exception", "cancel"}: + assert row.error == "Agent preview execution interrupted" - -def _env(variable_id: str, name: str) -> StringVariable: - return StringVariable(id=variable_id, name=name, value="v", selector=["env", name]) - - -@pytest.mark.parametrize( - "upsert", - [_env("new-id", "TAKEN"), _env("other-id", "TAKEN")], - ids=["new id reuses a name", "rename onto another variable's name"], -) -def test_environment_patch_rejects_a_name_another_id_uses(upsert: StringVariable) -> None: - current = [_env("taken-id", "TAKEN"), _env("other-id", "OTHER")] - with pytest.raises(ValueError, match="Environment variable names must be unique"): - _merge_environment_variable_patch(current, [upsert], []) + @staticmethod + @pytest.mark.usefixtures("preview_flask_app") + def test_agent_preview_rejects_mismatched_workflow_before_creating_caller( + sqlite_engine: _workflow_agent_preview_identity_Engine, + ) -> None: + app = Factory.create_app(app_id="app", tenant_id="tenant") + account = Factory.create_account() + workflow = Factory.create_workflow( + app_id="another-app", + tenant_id="another-tenant", + graph={ + "nodes": [{"id": "agent", "data": {"type": "agent", "version": "2", "agent_node_kind": "dify_agent"}}], + "edges": [], + }, + ) + sessions = sessionmaker(sqlite_engine, expire_on_commit=False) + with ( + patch( + "services.workflow_service.WorkflowDraftVariableService.prefill_conversation_variable_default_values" + ) as prefill, + pytest.raises(ValueError, match="does not belong"), + ): + WorkflowService(sessions).run_draft_workflow_node(app, workflow, "agent", {}, account) + prefill.assert_not_called() + with sessions() as session: + assert session.scalar(select(func.count()).select_from(WorkflowNodeExecutionModel)) == 0 diff --git a/api/tests/unit_tests/services/tools/test_builtin_tools_manage_service.py b/api/tests/unit_tests/services/tools/test_builtin_tools_manage_service.py index 4345014bf6e..dc8887ffe56 100644 --- a/api/tests/unit_tests/services/tools/test_builtin_tools_manage_service.py +++ b/api/tests/unit_tests/services/tools/test_builtin_tools_manage_service.py @@ -480,7 +480,7 @@ class TestGetOauthClientSchema: monkeypatch.setattr(BuiltinToolManageService, "is_oauth_custom_client_enabled", MagicMock(return_value=True)) monkeypatch.setattr(BuiltinToolManageService, "is_oauth_system_client_exists", MagicMock(return_value=False)) monkeypatch.setattr(BuiltinToolManageService, "get_custom_oauth_client_params", MagicMock(return_value={})) - from tests.unit_tests.config_override import apply_config_overrides + from tests.helpers.unit.config_override import apply_config_overrides apply_config_overrides(monkeypatch, CONSOLE_API_URL="https://api.example.com") diff --git a/api/tests/unit_tests/services/tools/test_mcp_tools_manage_service_lookup.py b/api/tests/unit_tests/services/tools/test_mcp_tools_manage_service.py similarity index 100% rename from api/tests/unit_tests/services/tools/test_mcp_tools_manage_service_lookup.py rename to api/tests/unit_tests/services/tools/test_mcp_tools_manage_service.py diff --git a/api/tests/unit_tests/services/tools/test_mcp_tools_transform.py b/api/tests/unit_tests/services/tools/test_mcp_tools_transform.py deleted file mode 100644 index 533577f8955..00000000000 --- a/api/tests/unit_tests/services/tools/test_mcp_tools_transform.py +++ /dev/null @@ -1,455 +0,0 @@ -"""Test cases for MCP tool transformation functionality.""" - -from unittest.mock import Mock - -import pytest -from pytest_mock import MockerFixture - -from core.mcp.types import Tool as MCPTool -from core.tools.entities.api_entities import ToolApiEntity, ToolProviderApiEntity -from core.tools.entities.common_entities import I18nObject -from core.tools.entities.tool_entities import ToolParameter, ToolProviderType -from models.account import Account -from models.tools import MCPToolProvider -from services.tools.tools_transform_service import ToolTransformService - - -@pytest.fixture -def mock_user(): - """Provides a real mapped account returned by the provider lookup.""" - return Account(name="Test User", email="user@example.com") - - -@pytest.fixture -def mock_provider(mock_user, mocker: MockerFixture): - """Provides a mock MCPToolProvider with a loaded user.""" - provider = MCPToolProvider( - name="Test Provider", - server_identifier="test-provider", - server_url="https://example.com", - server_url_hash="hash", - icon="icon", - tenant_id="tenant-id", - user_id="user-id", - ) - mocker.patch.object(provider, "load_user", return_value=mock_user) - return provider - - -@pytest.fixture -def mock_provider_no_user(mocker: MockerFixture): - """Provides a mock MCPToolProvider with no user.""" - provider = MCPToolProvider( - name="Test Provider", - server_identifier="test-provider", - server_url="https://example.com", - server_url_hash="hash", - icon="icon", - tenant_id="tenant-id", - user_id="user-id", - ) - mocker.patch.object(provider, "load_user", return_value=None) - return provider - - -@pytest.fixture -def mock_provider_full(mock_user, mocker: MockerFixture): - """Provides a fully configured mock MCPToolProvider for detailed tests.""" - provider = MCPToolProvider( - name="Test MCP Provider", - server_identifier="server-identifier-456", - server_url="https://example.com", - server_url_hash="hash", - icon="icon", - tenant_id="tenant-id", - user_id="user-id", - authed=True, - timeout=30, - sse_read_timeout=300, - ) - provider.id = "provider-id-123" - provider.provider_icon = "icon.png" - provider.masked_server_url = "https://*****.com/mcp" - provider.masked_headers = {"Authorization": "Bearer *****"} - provider.decrypted_headers = {"Authorization": "Bearer secret-token"} - - # Mock timestamp - mock_updated_at = Mock() - mock_updated_at.timestamp.return_value = 1234567890 - provider.updated_at = mock_updated_at - - mocker.patch.object(provider, "load_user", return_value=mock_user) - return provider - - -@pytest.fixture -def sample_mcp_tools(): - """Provides sample MCP tools for testing.""" - return { - "simple": MCPTool( - name="simple_tool", description="A simple test tool", inputSchema={"type": "object", "properties": {}} - ), - "none_desc": MCPTool(name="tool_none_desc", description=None, inputSchema={"type": "object", "properties": {}}), - "complex": MCPTool( - name="complex_tool", - description="A tool with complex parameters", - inputSchema={ - "type": "object", - "properties": { - "text": {"type": "string", "description": "Input text"}, - "count": {"type": "integer", "description": "Number of items", "minimum": 1, "maximum": 100}, - "options": {"type": "array", "items": {"type": "string"}, "description": "List of options"}, - }, - "required": ["text"], - }, - ), - } - - -class TestMCPToolTransform: - """Test cases for MCP tool transformation methods.""" - - def test_mcp_tool_to_user_tool_with_none_description(self, mock_provider): - """Test that mcp_tool_to_user_tool handles None description correctly.""" - # Create MCP tools with None description - tools = [ - MCPTool( - name="tool1", - description=None, # This is the case that caused the error - inputSchema={"type": "object", "properties": {}}, - ), - MCPTool( - name="tool2", - description=None, - inputSchema={ - "type": "object", - "properties": {"param1": {"type": "string", "description": "A parameter"}}, - }, - ), - ] - - # Call the method - result = ToolTransformService.mcp_tool_to_user_tool(mock_provider, tools) - - # Verify the result - assert len(result) == 2 - assert all(isinstance(tool, ToolApiEntity) for tool in result) - - # Check first tool - assert result[0].name == "tool1" - assert result[0].author == "Test User" - assert isinstance(result[0].label, I18nObject) - assert result[0].label.en_US == "tool1" - assert isinstance(result[0].description, I18nObject) - assert result[0].description.en_US == "" # Should be empty string, not None - assert result[0].description.zh_Hans == "" - - # Check second tool - assert result[1].name == "tool2" - assert result[1].description.en_US == "" - assert result[1].description.zh_Hans == "" - - def test_mcp_tool_to_user_tool_with_description(self, mock_provider): - """Test that mcp_tool_to_user_tool handles normal description correctly.""" - # Create MCP tools with description - tools = [ - MCPTool( - name="tool_with_desc", - description="This is a test tool that does something useful", - inputSchema={"type": "object", "properties": {}}, - ) - ] - - # Call the method - result = ToolTransformService.mcp_tool_to_user_tool(mock_provider, tools) - - # Verify the result - assert len(result) == 1 - assert isinstance(result[0], ToolApiEntity) - assert result[0].name == "tool_with_desc" - assert result[0].description.en_US == "This is a test tool that does something useful" - assert result[0].description.zh_Hans == "This is a test tool that does something useful" - - def test_mcp_tool_to_user_tool_with_no_user(self, mock_provider_no_user): - """Test that mcp_tool_to_user_tool handles None user correctly.""" - # Create MCP tool - tools = [MCPTool(name="tool1", description="Test tool", inputSchema={"type": "object", "properties": {}})] - - # Call the method - result = ToolTransformService.mcp_tool_to_user_tool(mock_provider_no_user, tools) - - # Verify the result - assert len(result) == 1 - assert result[0].author == "Anonymous" - - def test_mcp_tool_to_user_tool_with_complex_schema(self, mock_provider, sample_mcp_tools): - """Test that mcp_tool_to_user_tool correctly converts complex input schemas.""" - # Use complex tool from fixtures - tools = [sample_mcp_tools["complex"]] - - # Call the method - result = ToolTransformService.mcp_tool_to_user_tool(mock_provider, tools) - - # Verify the result - assert len(result) == 1 - assert result[0].name == "complex_tool" - assert result[0].parameters is not None - # The actual parameter conversion is handled by convert_mcp_schema_to_parameter - # which should be tested separately - - def test_convert_mcp_schema_to_parameter_preserves_anyof_object_type(self): - """Nullable object schemas should keep the object parameter type.""" - schema = { - "type": "object", - "properties": { - "retrieval_model": { - "anyOf": [{"type": "object"}, {"type": "null"}], - "description": "检索模型配置", - } - }, - } - - result = ToolTransformService.convert_mcp_schema_to_parameter(schema) - - assert len(result) == 1 - assert result[0].name == "retrieval_model" - assert result[0].type == ToolParameter.ToolParameterType.OBJECT - assert result[0].input_schema == schema["properties"]["retrieval_model"] - - def test_convert_mcp_schema_to_parameter_preserves_oneof_object_type(self): - """Nullable oneOf object schemas should keep the object parameter type.""" - schema = { - "type": "object", - "properties": { - "retrieval_model": { - "oneOf": [{"type": "object"}, {"type": "null"}], - "description": "检索模型配置", - } - }, - } - - result = ToolTransformService.convert_mcp_schema_to_parameter(schema) - - assert len(result) == 1 - assert result[0].name == "retrieval_model" - assert result[0].type == ToolParameter.ToolParameterType.OBJECT - assert result[0].input_schema == schema["properties"]["retrieval_model"] - - def test_convert_mcp_schema_to_parameter_handles_null_type(self): - """Schemas with only a null type should fall back to string.""" - schema = { - "type": "object", - "properties": { - "null_prop_str": {"type": "null"}, - "null_prop_list": {"type": ["null"]}, - }, - } - - result = ToolTransformService.convert_mcp_schema_to_parameter(schema) - - assert len(result) == 2 - param_map = {parameter.name: parameter for parameter in result} - assert "null_prop_str" in param_map - assert param_map["null_prop_str"].type == ToolParameter.ToolParameterType.STRING - assert "null_prop_list" in param_map - assert param_map["null_prop_list"].type == ToolParameter.ToolParameterType.STRING - - def test_convert_mcp_schema_to_parameter_preserves_allof_object_type_with_multiple_object_items(self): - """Property-level allOf with multiple object items should still resolve to object.""" - schema = { - "type": "object", - "properties": { - "config": { - "allOf": [ - { - "type": "object", - "properties": { - "enabled": {"type": "boolean"}, - }, - "required": ["enabled"], - }, - { - "type": "object", - "properties": { - "priority": {"type": "integer", "minimum": 1, "maximum": 10}, - }, - "required": ["priority"], - }, - ], - "description": "Config must match all schemas (allOf)", - } - }, - } - - result = ToolTransformService.convert_mcp_schema_to_parameter(schema) - - assert len(result) == 1 - assert result[0].name == "config" - assert result[0].type == ToolParameter.ToolParameterType.OBJECT - assert result[0].input_schema == schema["properties"]["config"] - - def test_convert_mcp_schema_to_parameter_preserves_allof_object_type(self): - """Composed property schemas should keep the object parameter type.""" - schema = { - "type": "object", - "properties": { - "retrieval_model": { - "allOf": [ - {"type": "object"}, - {"properties": {"top_k": {"type": "integer"}}}, - ], - "description": "检索模型配置", - } - }, - } - - result = ToolTransformService.convert_mcp_schema_to_parameter(schema) - - assert len(result) == 1 - assert result[0].name == "retrieval_model" - assert result[0].type == ToolParameter.ToolParameterType.OBJECT - assert result[0].input_schema == schema["properties"]["retrieval_model"] - - def test_convert_mcp_schema_to_parameter_limits_recursive_schema_depth(self): - """Self-referential composed schemas should stop resolving after the configured max depth.""" - recursive_property: dict[str, object] = {"description": "Recursive schema"} - recursive_property["anyOf"] = [recursive_property] - schema = { - "type": "object", - "properties": { - "recursive_config": recursive_property, - }, - } - - result = ToolTransformService.convert_mcp_schema_to_parameter(schema) - - assert len(result) == 1 - assert result[0].name == "recursive_config" - assert result[0].type == ToolParameter.ToolParameterType.STRING - assert result[0].input_schema is None - - def test_mcp_provider_to_user_provider_minimal_response(self, mock_provider_full, mocker: MockerFixture): - """Test mcp_provider_to_user_provider with a response that omits optional fields.""" - # Set tools data with null description - mock_provider_full.tools = '[{"name": "tool1", "description": null, "inputSchema": {}}]' - - # Mock the to_entity and to_api_response methods - mock_entity = Mock() - mock_entity.to_api_response.return_value = { - "name": "Test MCP Provider", - "type": ToolProviderType.MCP, - "is_team_authorization": True, - "server_url": "https://*****.com/mcp", - "provider_icon": "icon.png", - "masked_headers": {"Authorization": "Bearer *****"}, - "updated_at": 1234567890, - "labels": [], - "author": "Test User", - "description": I18nObject(en_US="Test MCP Provider Description", zh_Hans="Test MCP Provider Description"), - "icon": "icon.png", - "label": I18nObject(en_US="Test MCP Provider", zh_Hans="Test MCP Provider"), - "masked_credentials": {}, - } - mocker.patch.object(mock_provider_full, "to_entity", return_value=mock_entity) - - result = ToolTransformService.mcp_provider_to_user_provider(mock_provider_full) - - # Verify the result - assert isinstance(result, ToolProviderApiEntity) - assert result.id == "provider-id-123" - assert result.name == "Test MCP Provider" - assert result.type == ToolProviderType.MCP - assert result.is_team_authorization is True - assert result.server_url == "https://*****.com/mcp" - assert len(result.tools) == 1 - assert result.tools[0].description.en_US == "" # Should handle None description - - def test_mcp_provider_to_user_provider_full_response(self, mock_provider_full, mocker: MockerFixture): - """Test mcp_provider_to_user_provider with configuration and sensitive fields.""" - # Set tools data with description - mock_provider_full.tools = '[{"name": "tool1", "description": "Tool description", "inputSchema": {}}]' - - # Mock the to_entity and to_api_response methods - mock_entity = Mock() - mock_entity.to_api_response.return_value = { - "name": "Test MCP Provider", - "type": ToolProviderType.MCP, - "is_team_authorization": True, - "server_url": "https://*****.com/mcp", - "provider_icon": "icon.png", - "masked_headers": {"Authorization": "Bearer *****"}, - "updated_at": 1234567890, - "labels": [], - "configuration": {"timeout": "30", "sse_read_timeout": "300"}, - "original_headers": {"Authorization": "Bearer secret-token"}, - "author": "Test User", - "description": I18nObject(en_US="Test MCP Provider Description", zh_Hans="Test MCP Provider Description"), - "icon": "icon.png", - "label": I18nObject(en_US="Test MCP Provider", zh_Hans="Test MCP Provider"), - "masked_credentials": {}, - } - mocker.patch.object(mock_provider_full, "to_entity", return_value=mock_entity) - - result = ToolTransformService.mcp_provider_to_user_provider(mock_provider_full) - - # Verify the result - assert isinstance(result, ToolProviderApiEntity) - assert result.id == "provider-id-123" - assert result.server_identifier == "server-identifier-456" - assert result.configuration is not None - assert result.configuration.timeout == 30 - assert result.configuration.sse_read_timeout == 300 - assert result.original_headers == {"Authorization": "Bearer secret-token"} - assert len(result.tools) == 1 - assert result.tools[0].description.en_US == "Tool description" - - def test_mcp_tool_to_user_tool_falls_back_to_name_when_title_is_null(self, mock_provider): - """Tools where the server returned ``title: null`` should fall back to the tool name. - - Regression test for langgenius/dify#42453: Exa's MCP server returns - ``title: null`` for some tools, breaking tool conversion when the - server-provided title is required. - """ - tools = [ - MCPTool( - name="search_web", - title=None, - description="Search the web", - inputSchema={"type": "object", "properties": {}}, - ), - MCPTool( - name="get_content", - description="Fetch page content", - inputSchema={"type": "object", "properties": {}}, - ), - ] - - result = ToolTransformService.mcp_tool_to_user_tool(mock_provider, tools) - - assert len(result) == 2 - # title is explicitly null -> must fall back to ``name`` - assert result[0].name == "search_web" - assert result[0].label.en_US == "search_web" - assert result[0].label.zh_Hans == "search_web" - # title key is absent entirely -> must also fall back to ``name`` - assert result[1].name == "get_content" - assert result[1].label.en_US == "get_content" - assert result[1].label.zh_Hans == "get_content" - - def test_mcp_tool_to_user_tool_uses_title_when_provided(self, mock_provider): - """When the server provides a non-null title, it should be used for the label.""" - tools = [ - MCPTool( - name="search_web", - title="Search the Web", - description="Search the web", - inputSchema={"type": "object", "properties": {}}, - ) - ] - - result = ToolTransformService.mcp_tool_to_user_tool(mock_provider, tools) - - assert len(result) == 1 - assert result[0].name == "search_web" - assert result[0].label.en_US == "Search the Web" - assert result[0].label.zh_Hans == "Search the Web" diff --git a/api/tests/unit_tests/services/tools/test_tools_transform_service.py b/api/tests/unit_tests/services/tools/test_tools_transform_service.py index 2bc491a4452..e1501f1ec7f 100644 --- a/api/tests/unit_tests/services/tools/test_tools_transform_service.py +++ b/api/tests/unit_tests/services/tools/test_tools_transform_service.py @@ -1,7 +1,9 @@ -from unittest.mock import ANY, MagicMock, patch +from unittest.mock import ANY, MagicMock, Mock, patch import pytest +from pytest_mock import MockerFixture +from core.mcp.types import Tool as MCPTool from core.tools.entities.api_entities import ToolApiEntity, ToolProviderApiEntity from core.tools.entities.common_entities import I18nObject from core.tools.entities.tool_entities import ( @@ -11,9 +13,462 @@ from core.tools.entities.tool_entities import ( ToolProviderIdentity, ToolProviderType, ) +from models.account import Account +from models.tools import MCPToolProvider from services.tools.tools_transform_service import ToolTransformService from tests.tool_fixtures import make_runtime_tool +"""Test cases for MCP tool transformation functionality.""" + + +class TestMcpToolsTransformCases: + @staticmethod + @pytest.fixture + def mock_user(): + """Provides a real mapped account returned by the provider lookup.""" + return Account(name="Test User", email="user@example.com") + + @staticmethod + @pytest.fixture + def mock_provider(mock_user, mocker: MockerFixture): + """Provides a mock MCPToolProvider with a loaded user.""" + provider = MCPToolProvider( + name="Test Provider", + server_identifier="test-provider", + server_url="https://example.com", + server_url_hash="hash", + icon="icon", + tenant_id="tenant-id", + user_id="user-id", + ) + mocker.patch.object(provider, "load_user", return_value=mock_user) + return provider + + @staticmethod + @pytest.fixture + def mock_provider_no_user(mocker: MockerFixture): + """Provides a mock MCPToolProvider with no user.""" + provider = MCPToolProvider( + name="Test Provider", + server_identifier="test-provider", + server_url="https://example.com", + server_url_hash="hash", + icon="icon", + tenant_id="tenant-id", + user_id="user-id", + ) + mocker.patch.object(provider, "load_user", return_value=None) + return provider + + @staticmethod + @pytest.fixture + def mock_provider_full(mock_user, mocker: MockerFixture): + """Provides a fully configured mock MCPToolProvider for detailed tests.""" + provider = MCPToolProvider( + name="Test MCP Provider", + server_identifier="server-identifier-456", + server_url="https://example.com", + server_url_hash="hash", + icon="icon", + tenant_id="tenant-id", + user_id="user-id", + authed=True, + timeout=30, + sse_read_timeout=300, + ) + provider.id = "provider-id-123" + provider.provider_icon = "icon.png" + provider.masked_server_url = "https://*****.com/mcp" + provider.masked_headers = {"Authorization": "Bearer *****"} + provider.decrypted_headers = {"Authorization": "Bearer secret-token"} + + # Mock timestamp + mock_updated_at = Mock() + mock_updated_at.timestamp.return_value = 1234567890 + provider.updated_at = mock_updated_at + + mocker.patch.object(provider, "load_user", return_value=mock_user) + return provider + + @staticmethod + @pytest.fixture + def sample_mcp_tools(): + """Provides sample MCP tools for testing.""" + return { + "simple": MCPTool( + name="simple_tool", description="A simple test tool", inputSchema={"type": "object", "properties": {}} + ), + "none_desc": MCPTool( + name="tool_none_desc", description=None, inputSchema={"type": "object", "properties": {}} + ), + "complex": MCPTool( + name="complex_tool", + description="A tool with complex parameters", + inputSchema={ + "type": "object", + "properties": { + "text": {"type": "string", "description": "Input text"}, + "count": {"type": "integer", "description": "Number of items", "minimum": 1, "maximum": 100}, + "options": {"type": "array", "items": {"type": "string"}, "description": "List of options"}, + }, + "required": ["text"], + }, + ), + } + + class TestMCPToolTransform: + """Test cases for MCP tool transformation methods.""" + + def test_mcp_tool_to_user_tool_with_none_description(self, mock_provider): + """Test that mcp_tool_to_user_tool handles None description correctly.""" + # Create MCP tools with None description + tools = [ + MCPTool( + name="tool1", + description=None, # This is the case that caused the error + inputSchema={"type": "object", "properties": {}}, + ), + MCPTool( + name="tool2", + description=None, + inputSchema={ + "type": "object", + "properties": {"param1": {"type": "string", "description": "A parameter"}}, + }, + ), + ] + + # Call the method + result = ToolTransformService.mcp_tool_to_user_tool(mock_provider, tools) + + # Verify the result + assert len(result) == 2 + assert all(isinstance(tool, ToolApiEntity) for tool in result) + + # Check first tool + assert result[0].name == "tool1" + assert result[0].author == "Test User" + assert isinstance(result[0].label, I18nObject) + assert result[0].label.en_US == "tool1" + assert isinstance(result[0].description, I18nObject) + assert result[0].description.en_US == "" # Should be empty string, not None + assert result[0].description.zh_Hans == "" + + # Check second tool + assert result[1].name == "tool2" + assert result[1].description.en_US == "" + assert result[1].description.zh_Hans == "" + + def test_mcp_tool_to_user_tool_with_description(self, mock_provider): + """Test that mcp_tool_to_user_tool handles normal description correctly.""" + # Create MCP tools with description + tools = [ + MCPTool( + name="tool_with_desc", + description="This is a test tool that does something useful", + inputSchema={"type": "object", "properties": {}}, + ) + ] + + # Call the method + result = ToolTransformService.mcp_tool_to_user_tool(mock_provider, tools) + + # Verify the result + assert len(result) == 1 + assert isinstance(result[0], ToolApiEntity) + assert result[0].name == "tool_with_desc" + assert result[0].description.en_US == "This is a test tool that does something useful" + assert result[0].description.zh_Hans == "This is a test tool that does something useful" + + def test_mcp_tool_to_user_tool_with_no_user(self, mock_provider_no_user): + """Test that mcp_tool_to_user_tool handles None user correctly.""" + # Create MCP tool + tools = [MCPTool(name="tool1", description="Test tool", inputSchema={"type": "object", "properties": {}})] + + # Call the method + result = ToolTransformService.mcp_tool_to_user_tool(mock_provider_no_user, tools) + + # Verify the result + assert len(result) == 1 + assert result[0].author == "Anonymous" + + def test_mcp_tool_to_user_tool_with_complex_schema(self, mock_provider, sample_mcp_tools): + """Test that mcp_tool_to_user_tool correctly converts complex input schemas.""" + # Use complex tool from fixtures + tools = [sample_mcp_tools["complex"]] + + # Call the method + result = ToolTransformService.mcp_tool_to_user_tool(mock_provider, tools) + + # Verify the result + assert len(result) == 1 + assert result[0].name == "complex_tool" + assert result[0].parameters is not None + # The actual parameter conversion is handled by convert_mcp_schema_to_parameter + # which should be tested separately + + def test_convert_mcp_schema_to_parameter_preserves_anyof_object_type(self): + """Nullable object schemas should keep the object parameter type.""" + schema = { + "type": "object", + "properties": { + "retrieval_model": { + "anyOf": [{"type": "object"}, {"type": "null"}], + "description": "检索模型配置", + } + }, + } + + result = ToolTransformService.convert_mcp_schema_to_parameter(schema) + + assert len(result) == 1 + assert result[0].name == "retrieval_model" + assert result[0].type == ToolParameter.ToolParameterType.OBJECT + assert result[0].input_schema == schema["properties"]["retrieval_model"] + + def test_convert_mcp_schema_to_parameter_preserves_oneof_object_type(self): + """Nullable oneOf object schemas should keep the object parameter type.""" + schema = { + "type": "object", + "properties": { + "retrieval_model": { + "oneOf": [{"type": "object"}, {"type": "null"}], + "description": "检索模型配置", + } + }, + } + + result = ToolTransformService.convert_mcp_schema_to_parameter(schema) + + assert len(result) == 1 + assert result[0].name == "retrieval_model" + assert result[0].type == ToolParameter.ToolParameterType.OBJECT + assert result[0].input_schema == schema["properties"]["retrieval_model"] + + def test_convert_mcp_schema_to_parameter_handles_null_type(self): + """Schemas with only a null type should fall back to string.""" + schema = { + "type": "object", + "properties": { + "null_prop_str": {"type": "null"}, + "null_prop_list": {"type": ["null"]}, + }, + } + + result = ToolTransformService.convert_mcp_schema_to_parameter(schema) + + assert len(result) == 2 + param_map = {parameter.name: parameter for parameter in result} + assert "null_prop_str" in param_map + assert param_map["null_prop_str"].type == ToolParameter.ToolParameterType.STRING + assert "null_prop_list" in param_map + assert param_map["null_prop_list"].type == ToolParameter.ToolParameterType.STRING + + def test_convert_mcp_schema_to_parameter_preserves_allof_object_type_with_multiple_object_items(self): + """Property-level allOf with multiple object items should still resolve to object.""" + schema = { + "type": "object", + "properties": { + "config": { + "allOf": [ + { + "type": "object", + "properties": { + "enabled": {"type": "boolean"}, + }, + "required": ["enabled"], + }, + { + "type": "object", + "properties": { + "priority": {"type": "integer", "minimum": 1, "maximum": 10}, + }, + "required": ["priority"], + }, + ], + "description": "Config must match all schemas (allOf)", + } + }, + } + + result = ToolTransformService.convert_mcp_schema_to_parameter(schema) + + assert len(result) == 1 + assert result[0].name == "config" + assert result[0].type == ToolParameter.ToolParameterType.OBJECT + assert result[0].input_schema == schema["properties"]["config"] + + def test_convert_mcp_schema_to_parameter_preserves_allof_object_type(self): + """Composed property schemas should keep the object parameter type.""" + schema = { + "type": "object", + "properties": { + "retrieval_model": { + "allOf": [ + {"type": "object"}, + {"properties": {"top_k": {"type": "integer"}}}, + ], + "description": "检索模型配置", + } + }, + } + + result = ToolTransformService.convert_mcp_schema_to_parameter(schema) + + assert len(result) == 1 + assert result[0].name == "retrieval_model" + assert result[0].type == ToolParameter.ToolParameterType.OBJECT + assert result[0].input_schema == schema["properties"]["retrieval_model"] + + def test_convert_mcp_schema_to_parameter_limits_recursive_schema_depth(self): + """Self-referential composed schemas should stop resolving after the configured max depth.""" + recursive_property: dict[str, object] = {"description": "Recursive schema"} + recursive_property["anyOf"] = [recursive_property] + schema = { + "type": "object", + "properties": { + "recursive_config": recursive_property, + }, + } + + result = ToolTransformService.convert_mcp_schema_to_parameter(schema) + + assert len(result) == 1 + assert result[0].name == "recursive_config" + assert result[0].type == ToolParameter.ToolParameterType.STRING + assert result[0].input_schema is None + + def test_mcp_provider_to_user_provider_minimal_response(self, mock_provider_full, mocker: MockerFixture): + """Test mcp_provider_to_user_provider with a response that omits optional fields.""" + # Set tools data with null description + mock_provider_full.tools = '[{"name": "tool1", "description": null, "inputSchema": {}}]' + + # Mock the to_entity and to_api_response methods + mock_entity = Mock() + mock_entity.to_api_response.return_value = { + "name": "Test MCP Provider", + "type": ToolProviderType.MCP, + "is_team_authorization": True, + "server_url": "https://*****.com/mcp", + "provider_icon": "icon.png", + "masked_headers": {"Authorization": "Bearer *****"}, + "updated_at": 1234567890, + "labels": [], + "author": "Test User", + "description": I18nObject( + en_US="Test MCP Provider Description", zh_Hans="Test MCP Provider Description" + ), + "icon": "icon.png", + "label": I18nObject(en_US="Test MCP Provider", zh_Hans="Test MCP Provider"), + "masked_credentials": {}, + } + mocker.patch.object(mock_provider_full, "to_entity", return_value=mock_entity) + + result = ToolTransformService.mcp_provider_to_user_provider(mock_provider_full) + + # Verify the result + assert isinstance(result, ToolProviderApiEntity) + assert result.id == "provider-id-123" + assert result.name == "Test MCP Provider" + assert result.type == ToolProviderType.MCP + assert result.is_team_authorization is True + assert result.server_url == "https://*****.com/mcp" + assert len(result.tools) == 1 + assert result.tools[0].description.en_US == "" # Should handle None description + + def test_mcp_provider_to_user_provider_full_response(self, mock_provider_full, mocker: MockerFixture): + """Test mcp_provider_to_user_provider with configuration and sensitive fields.""" + # Set tools data with description + mock_provider_full.tools = '[{"name": "tool1", "description": "Tool description", "inputSchema": {}}]' + + # Mock the to_entity and to_api_response methods + mock_entity = Mock() + mock_entity.to_api_response.return_value = { + "name": "Test MCP Provider", + "type": ToolProviderType.MCP, + "is_team_authorization": True, + "server_url": "https://*****.com/mcp", + "provider_icon": "icon.png", + "masked_headers": {"Authorization": "Bearer *****"}, + "updated_at": 1234567890, + "labels": [], + "configuration": {"timeout": "30", "sse_read_timeout": "300"}, + "original_headers": {"Authorization": "Bearer secret-token"}, + "author": "Test User", + "description": I18nObject( + en_US="Test MCP Provider Description", zh_Hans="Test MCP Provider Description" + ), + "icon": "icon.png", + "label": I18nObject(en_US="Test MCP Provider", zh_Hans="Test MCP Provider"), + "masked_credentials": {}, + } + mocker.patch.object(mock_provider_full, "to_entity", return_value=mock_entity) + + result = ToolTransformService.mcp_provider_to_user_provider(mock_provider_full) + + # Verify the result + assert isinstance(result, ToolProviderApiEntity) + assert result.id == "provider-id-123" + assert result.server_identifier == "server-identifier-456" + assert result.configuration is not None + assert result.configuration.timeout == 30 + assert result.configuration.sse_read_timeout == 300 + assert result.original_headers == {"Authorization": "Bearer secret-token"} + assert len(result.tools) == 1 + assert result.tools[0].description.en_US == "Tool description" + + def test_mcp_tool_to_user_tool_falls_back_to_name_when_title_is_null(self, mock_provider): + """Tools where the server returned ``title: null`` should fall back to the tool name. + + Regression test for langgenius/dify#42453: Exa's MCP server returns + ``title: null`` for some tools, breaking tool conversion when the + server-provided title is required. + """ + tools = [ + MCPTool( + name="search_web", + title=None, + description="Search the web", + inputSchema={"type": "object", "properties": {}}, + ), + MCPTool( + name="get_content", + description="Fetch page content", + inputSchema={"type": "object", "properties": {}}, + ), + ] + + result = ToolTransformService.mcp_tool_to_user_tool(mock_provider, tools) + + assert len(result) == 2 + # title is explicitly null -> must fall back to ``name`` + assert result[0].name == "search_web" + assert result[0].label.en_US == "search_web" + assert result[0].label.zh_Hans == "search_web" + # title key is absent entirely -> must also fall back to ``name`` + assert result[1].name == "get_content" + assert result[1].label.en_US == "get_content" + assert result[1].label.zh_Hans == "get_content" + + def test_mcp_tool_to_user_tool_uses_title_when_provided(self, mock_provider): + """When the server provides a non-null title, it should be used for the label.""" + tools = [ + MCPTool( + name="search_web", + title="Search the Web", + description="Search the web", + inputSchema={"type": "object", "properties": {}}, + ) + ] + + result = ToolTransformService.mcp_tool_to_user_tool(mock_provider, tools) + + assert len(result) == 1 + assert result[0].name == "search_web" + assert result[0].label.en_US == "Search the Web" + assert result[0].label.zh_Hans == "Search the Web" + + MODULE = "services.tools.tools_transform_service" @@ -31,490 +486,486 @@ def _parameter( ) -class TestToolTransformService: - """Test cases for ToolTransformService.convert_tool_entity_to_api_entity method""" +class TestToolsTransformServiceCases: + class TestToolTransformService: + """Test cases for ToolTransformService.convert_tool_entity_to_api_entity method""" - def test_convert_tool_with_parameter_override(self): - """Test that runtime parameters correctly override base parameters""" - base_param1 = _parameter("param1", "Base Param 1") - base_param2 = _parameter("param2", "Base Param 2") + def test_convert_tool_with_parameter_override(self): + """Test that runtime parameters correctly override base parameters""" + base_param1 = _parameter("param1", "Base Param 1") + base_param2 = _parameter("param2", "Base Param 2") - runtime_param1 = _parameter("param1", "Runtime Param 1") + runtime_param1 = _parameter("param1", "Runtime Param 1") - tool = make_runtime_tool(base_params=[base_param1, base_param2], runtime_params=[runtime_param1]) + tool = make_runtime_tool(base_params=[base_param1, base_param2], runtime_params=[runtime_param1]) - # Call the method - result = ToolTransformService.convert_tool_entity_to_api_entity(tool, "test_tenant", None) + # Call the method + result = ToolTransformService.convert_tool_entity_to_api_entity(tool, "test_tenant", None) - # Verify the result - assert isinstance(result, ToolApiEntity) - assert result.author == "test_author" - assert result.name == "test_tool" - assert result.parameters is not None - assert len(result.parameters) == 2 + # Verify the result + assert isinstance(result, ToolApiEntity) + assert result.author == "test_author" + assert result.name == "test_tool" + assert result.parameters is not None + assert len(result.parameters) == 2 - # Find the overridden parameter - overridden_param = next((p for p in result.parameters if p.name == "param1"), None) - assert overridden_param is not None - assert overridden_param.label.en_US == "Runtime Param 1" # Should be runtime version + # Find the overridden parameter + overridden_param = next((p for p in result.parameters if p.name == "param1"), None) + assert overridden_param is not None + assert overridden_param.label.en_US == "Runtime Param 1" # Should be runtime version - # Find the non-overridden parameter - original_param = next((p for p in result.parameters if p.name == "param2"), None) - assert original_param is not None - assert original_param.label.en_US == "Base Param 2" # Should be base version + # Find the non-overridden parameter + original_param = next((p for p in result.parameters if p.name == "param2"), None) + assert original_param is not None + assert original_param.label.en_US == "Base Param 2" # Should be base version - def test_convert_tool_with_additional_runtime_parameters(self): - """Test that additional runtime parameters are added to the final list""" - base_param1 = _parameter("param1", "Base Param 1") + def test_convert_tool_with_additional_runtime_parameters(self): + """Test that additional runtime parameters are added to the final list""" + base_param1 = _parameter("param1", "Base Param 1") - runtime_param1 = _parameter("param1", "Runtime Param 1") - runtime_param2 = _parameter("runtime_only", "Runtime Only Param") + runtime_param1 = _parameter("param1", "Runtime Param 1") + runtime_param2 = _parameter("runtime_only", "Runtime Only Param") - tool = make_runtime_tool(base_params=[base_param1], runtime_params=[runtime_param1, runtime_param2]) + tool = make_runtime_tool(base_params=[base_param1], runtime_params=[runtime_param1, runtime_param2]) - # Call the method - result = ToolTransformService.convert_tool_entity_to_api_entity(tool, "test_tenant", None) + # Call the method + result = ToolTransformService.convert_tool_entity_to_api_entity(tool, "test_tenant", None) - # Verify the result - assert isinstance(result, ToolApiEntity) - assert result.parameters is not None - assert len(result.parameters) == 2 + # Verify the result + assert isinstance(result, ToolApiEntity) + assert result.parameters is not None + assert len(result.parameters) == 2 - # Check that both parameters are present - param_names = [p.name for p in result.parameters] - assert "param1" in param_names - assert "runtime_only" in param_names + # Check that both parameters are present + param_names = [p.name for p in result.parameters] + assert "param1" in param_names + assert "runtime_only" in param_names - # Verify the overridden parameter has runtime version - overridden_param = next((p for p in result.parameters if p.name == "param1"), None) - assert overridden_param is not None - assert overridden_param.label.en_US == "Runtime Param 1" + # Verify the overridden parameter has runtime version + overridden_param = next((p for p in result.parameters if p.name == "param1"), None) + assert overridden_param is not None + assert overridden_param.label.en_US == "Runtime Param 1" - # Verify the new runtime parameter is included - new_param = next((p for p in result.parameters if p.name == "runtime_only"), None) - assert new_param is not None - assert new_param.label.en_US == "Runtime Only Param" + # Verify the new runtime parameter is included + new_param = next((p for p in result.parameters if p.name == "runtime_only"), None) + assert new_param is not None + assert new_param.label.en_US == "Runtime Only Param" - def test_convert_tool_with_non_form_runtime_parameters(self): - """Test that non-FORM runtime parameters are not added as new parameters""" - base_param1 = _parameter("param1", "Base Param 1") + def test_convert_tool_with_non_form_runtime_parameters(self): + """Test that non-FORM runtime parameters are not added as new parameters""" + base_param1 = _parameter("param1", "Base Param 1") - runtime_param1 = _parameter("param1", "Runtime Param 1") - runtime_param2 = _parameter("llm_param", "LLM Param", ToolParameter.ToolParameterForm.LLM) + runtime_param1 = _parameter("param1", "Runtime Param 1") + runtime_param2 = _parameter("llm_param", "LLM Param", ToolParameter.ToolParameterForm.LLM) - tool = make_runtime_tool(base_params=[base_param1], runtime_params=[runtime_param1, runtime_param2]) + tool = make_runtime_tool(base_params=[base_param1], runtime_params=[runtime_param1, runtime_param2]) - # Call the method - result = ToolTransformService.convert_tool_entity_to_api_entity(tool, "test_tenant", None) + # Call the method + result = ToolTransformService.convert_tool_entity_to_api_entity(tool, "test_tenant", None) - # Verify the result - assert isinstance(result, ToolApiEntity) - assert result.parameters is not None - assert len(result.parameters) == 1 # Only the FORM parameter should be present + # Verify the result + assert isinstance(result, ToolApiEntity) + assert result.parameters is not None + assert len(result.parameters) == 1 # Only the FORM parameter should be present - # Check that only the FORM parameter is present - param_names = [p.name for p in result.parameters] - assert "param1" in param_names - assert "llm_param" not in param_names + # Check that only the FORM parameter is present + param_names = [p.name for p in result.parameters] + assert "param1" in param_names + assert "llm_param" not in param_names - def test_convert_tool_with_empty_parameters(self): - """Test conversion with empty base and runtime parameters""" - tool = make_runtime_tool(base_params=[], runtime_params=[]) + def test_convert_tool_with_empty_parameters(self): + """Test conversion with empty base and runtime parameters""" + tool = make_runtime_tool(base_params=[], runtime_params=[]) - # Call the method - result = ToolTransformService.convert_tool_entity_to_api_entity(tool, "test_tenant", None) + # Call the method + result = ToolTransformService.convert_tool_entity_to_api_entity(tool, "test_tenant", None) - # Verify the result - assert isinstance(result, ToolApiEntity) - assert result.parameters is not None - assert len(result.parameters) == 0 + # Verify the result + assert isinstance(result, ToolApiEntity) + assert result.parameters is not None + assert len(result.parameters) == 0 - def test_convert_tool_with_none_parameters(self): - """Test conversion when base parameters is None""" - tool = make_runtime_tool(base_params=None, runtime_params=[]) + def test_convert_tool_with_none_parameters(self): + """Test conversion when base parameters is None""" + tool = make_runtime_tool(base_params=None, runtime_params=[]) - # Call the method - result = ToolTransformService.convert_tool_entity_to_api_entity(tool, "test_tenant", None) + # Call the method + result = ToolTransformService.convert_tool_entity_to_api_entity(tool, "test_tenant", None) - # Verify the result - assert isinstance(result, ToolApiEntity) - assert result.parameters is not None - assert len(result.parameters) == 0 + # Verify the result + assert isinstance(result, ToolApiEntity) + assert result.parameters is not None + assert len(result.parameters) == 0 - def test_convert_tool_parameter_order_preserved(self): - """Test that parameter order is preserved correctly""" - base_param1 = _parameter("param1", "Base Param 1") - base_param2 = _parameter("param2", "Base Param 2") - base_param3 = _parameter("param3", "Base Param 3") + def test_convert_tool_parameter_order_preserved(self): + """Test that parameter order is preserved correctly""" + base_param1 = _parameter("param1", "Base Param 1") + base_param2 = _parameter("param2", "Base Param 2") + base_param3 = _parameter("param3", "Base Param 3") - # Create runtime parameter that overrides middle parameter - runtime_param2 = _parameter("param2", "Runtime Param 2") + # Create runtime parameter that overrides middle parameter + runtime_param2 = _parameter("param2", "Runtime Param 2") - # Create new runtime parameter - runtime_param4 = _parameter("param4", "Runtime Param 4") + # Create new runtime parameter + runtime_param4 = _parameter("param4", "Runtime Param 4") - tool = make_runtime_tool( - base_params=[base_param1, base_param2, base_param3], runtime_params=[runtime_param2, runtime_param4] - ) + tool = make_runtime_tool( + base_params=[base_param1, base_param2, base_param3], runtime_params=[runtime_param2, runtime_param4] + ) - # Call the method - result = ToolTransformService.convert_tool_entity_to_api_entity(tool, "test_tenant", None) + # Call the method + result = ToolTransformService.convert_tool_entity_to_api_entity(tool, "test_tenant", None) - # Verify the result - assert isinstance(result, ToolApiEntity) - assert result.parameters is not None - assert len(result.parameters) == 4 + # Verify the result + assert isinstance(result, ToolApiEntity) + assert result.parameters is not None + assert len(result.parameters) == 4 - # Check that order is maintained: base parameters first, then new runtime parameters - param_names = [p.name for p in result.parameters] - assert param_names == ["param1", "param2", "param3", "param4"] + # Check that order is maintained: base parameters first, then new runtime parameters + param_names = [p.name for p in result.parameters] + assert param_names == ["param1", "param2", "param3", "param4"] - # Verify that param2 was overridden with runtime version - param2 = result.parameters[1] - assert param2.name == "param2" - assert param2.label.en_US == "Runtime Param 2" + # Verify that param2 was overridden with runtime version + param2 = result.parameters[1] + assert param2.name == "param2" + assert param2.label.en_US == "Runtime Param 2" + class TestWorkflowProviderToUserProvider: + """Test cases for ToolTransformService.workflow_provider_to_user_provider method""" -class TestWorkflowProviderToUserProvider: - """Test cases for ToolTransformService.workflow_provider_to_user_provider method""" + def test_workflow_provider_to_user_provider_with_workflow_app_id(self): + """Test that workflow_provider_to_user_provider correctly sets workflow_app_id.""" + from core.tools.workflow_as_tool.provider import WorkflowToolProviderController - def test_workflow_provider_to_user_provider_with_workflow_app_id(self): - """Test that workflow_provider_to_user_provider correctly sets workflow_app_id.""" - from core.tools.workflow_as_tool.provider import WorkflowToolProviderController - - workflow_app_id = "app_123" - provider_id = "provider_123" - controller = WorkflowToolProviderController( - provider_id=provider_id, - entity=ToolProviderEntity( - # Retain legacy emoji mappings accepted by the API transformation. - identity=ToolProviderIdentity.model_construct( - author="test_author", - name="test_workflow_tool", - description=I18nObject(en_US="Test description"), - icon={"type": "emoji", "content": "🔧"}, - icon_dark=None, - label=I18nObject(en_US="Test Workflow Tool"), + workflow_app_id = "app_123" + provider_id = "provider_123" + controller = WorkflowToolProviderController( + provider_id=provider_id, + entity=ToolProviderEntity( + # Retain legacy emoji mappings accepted by the API transformation. + identity=ToolProviderIdentity.model_construct( + author="test_author", + name="test_workflow_tool", + description=I18nObject(en_US="Test description"), + icon={"type": "emoji", "content": "🔧"}, + icon_dark=None, + label=I18nObject(en_US="Test Workflow Tool"), + ), ), - ), - ) + ) - # Call the method - result = ToolTransformService.workflow_provider_to_user_provider( - provider_controller=controller, - labels=["label1", "label2"], - workflow_app_id=workflow_app_id, - ) + # Call the method + result = ToolTransformService.workflow_provider_to_user_provider( + provider_controller=controller, + labels=["label1", "label2"], + workflow_app_id=workflow_app_id, + ) - # Verify the result - assert isinstance(result, ToolProviderApiEntity) - assert result.id == provider_id - assert result.author == "test_author" - assert result.name == "test_workflow_tool" - assert result.type == ToolProviderType.WORKFLOW - assert result.workflow_app_id == workflow_app_id - assert result.labels == ["label1", "label2"] - assert result.is_team_authorization is True - assert result.plugin_id is None - assert result.plugin_unique_identifier is None - assert result.tools == [] + # Verify the result + assert isinstance(result, ToolProviderApiEntity) + assert result.id == provider_id + assert result.author == "test_author" + assert result.name == "test_workflow_tool" + assert result.type == ToolProviderType.WORKFLOW + assert result.workflow_app_id == workflow_app_id + assert result.labels == ["label1", "label2"] + assert result.is_team_authorization is True + assert result.plugin_id is None + assert result.plugin_unique_identifier is None + assert result.tools == [] - def test_workflow_provider_to_user_provider_without_workflow_app_id(self): - """Test that workflow_provider_to_user_provider works when workflow_app_id is not provided.""" - from core.tools.workflow_as_tool.provider import WorkflowToolProviderController + def test_workflow_provider_to_user_provider_without_workflow_app_id(self): + """Test that workflow_provider_to_user_provider works when workflow_app_id is not provided.""" + from core.tools.workflow_as_tool.provider import WorkflowToolProviderController - provider_id = "provider_123" - controller = WorkflowToolProviderController( - provider_id=provider_id, - entity=ToolProviderEntity( - # Retain legacy emoji mappings accepted by the API transformation. - identity=ToolProviderIdentity.model_construct( - author="test_author", - name="test_workflow_tool", - description=I18nObject(en_US="Test description"), - icon={"type": "emoji", "content": "🔧"}, - icon_dark=None, - label=I18nObject(en_US="Test Workflow Tool"), + provider_id = "provider_123" + controller = WorkflowToolProviderController( + provider_id=provider_id, + entity=ToolProviderEntity( + # Retain legacy emoji mappings accepted by the API transformation. + identity=ToolProviderIdentity.model_construct( + author="test_author", + name="test_workflow_tool", + description=I18nObject(en_US="Test description"), + icon={"type": "emoji", "content": "🔧"}, + icon_dark=None, + label=I18nObject(en_US="Test Workflow Tool"), + ), ), - ), - ) + ) - # Call the method without workflow_app_id - result = ToolTransformService.workflow_provider_to_user_provider( - provider_controller=controller, - labels=["label1"], - ) + # Call the method without workflow_app_id + result = ToolTransformService.workflow_provider_to_user_provider( + provider_controller=controller, + labels=["label1"], + ) - # Verify the result - assert isinstance(result, ToolProviderApiEntity) - assert result.id == provider_id - assert result.workflow_app_id is None - assert result.labels == ["label1"] + # Verify the result + assert isinstance(result, ToolProviderApiEntity) + assert result.id == provider_id + assert result.workflow_app_id is None + assert result.labels == ["label1"] - def test_workflow_provider_to_user_provider_workflow_app_id_none(self): - """Test that workflow_provider_to_user_provider handles None workflow_app_id explicitly.""" - from core.tools.workflow_as_tool.provider import WorkflowToolProviderController + def test_workflow_provider_to_user_provider_workflow_app_id_none(self): + """Test that workflow_provider_to_user_provider handles None workflow_app_id explicitly.""" + from core.tools.workflow_as_tool.provider import WorkflowToolProviderController - provider_id = "provider_123" - controller = WorkflowToolProviderController( - provider_id=provider_id, - entity=ToolProviderEntity( - # Retain legacy emoji mappings accepted by the API transformation. - identity=ToolProviderIdentity.model_construct( - author="test_author", - name="test_workflow_tool", - description=I18nObject(en_US="Test description"), - icon={"type": "emoji", "content": "🔧"}, - icon_dark=None, - label=I18nObject(en_US="Test Workflow Tool"), + provider_id = "provider_123" + controller = WorkflowToolProviderController( + provider_id=provider_id, + entity=ToolProviderEntity( + # Retain legacy emoji mappings accepted by the API transformation. + identity=ToolProviderIdentity.model_construct( + author="test_author", + name="test_workflow_tool", + description=I18nObject(en_US="Test description"), + icon={"type": "emoji", "content": "🔧"}, + icon_dark=None, + label=I18nObject(en_US="Test Workflow Tool"), + ), ), - ), - ) + ) - # Call the method with explicit None values - result = ToolTransformService.workflow_provider_to_user_provider( - provider_controller=controller, - labels=None, - workflow_app_id=None, - ) + # Call the method with explicit None values + result = ToolTransformService.workflow_provider_to_user_provider( + provider_controller=controller, + labels=None, + workflow_app_id=None, + ) - # Verify the result - assert isinstance(result, ToolProviderApiEntity) - assert result.id == provider_id - assert result.workflow_app_id is None - assert result.labels == [] + # Verify the result + assert isinstance(result, ToolProviderApiEntity) + assert result.id == provider_id + assert result.workflow_app_id is None + assert result.labels == [] - def test_workflow_provider_to_user_provider_preserves_other_fields(self): - """Test that workflow_provider_to_user_provider preserves all other entity fields.""" - from core.tools.workflow_as_tool.provider import WorkflowToolProviderController + def test_workflow_provider_to_user_provider_preserves_other_fields(self): + """Test that workflow_provider_to_user_provider preserves all other entity fields.""" + from core.tools.workflow_as_tool.provider import WorkflowToolProviderController - workflow_app_id = "app_456" - provider_id = "provider_456" - controller = WorkflowToolProviderController( - provider_id=provider_id, - entity=ToolProviderEntity( - # Retain legacy emoji mappings accepted by the API transformation. - identity=ToolProviderIdentity.model_construct( - author="another_author", - name="another_workflow_tool", - description=I18nObject(en_US="Another description", zh_Hans="Another description"), - icon={"type": "emoji", "content": "⚙️"}, - icon_dark={"type": "emoji", "content": "🔧"}, - label=I18nObject(en_US="Another Workflow Tool", zh_Hans="Another Workflow Tool"), + workflow_app_id = "app_456" + provider_id = "provider_456" + controller = WorkflowToolProviderController( + provider_id=provider_id, + entity=ToolProviderEntity( + # Retain legacy emoji mappings accepted by the API transformation. + identity=ToolProviderIdentity.model_construct( + author="another_author", + name="another_workflow_tool", + description=I18nObject(en_US="Another description", zh_Hans="Another description"), + icon={"type": "emoji", "content": "⚙️"}, + icon_dark={"type": "emoji", "content": "🔧"}, + label=I18nObject(en_US="Another Workflow Tool", zh_Hans="Another Workflow Tool"), + ), ), - ), + ) + + # Call the method + result = ToolTransformService.workflow_provider_to_user_provider( + provider_controller=controller, + labels=["automation", "workflow"], + workflow_app_id=workflow_app_id, + ) + + # Verify all fields are preserved correctly + assert isinstance(result, ToolProviderApiEntity) + assert result.id == provider_id + assert result.author == "another_author" + assert result.name == "another_workflow_tool" + assert result.description.en_US == "Another description" + assert result.description.zh_Hans == "Another description" + assert result.icon == {"type": "emoji", "content": "⚙️"} + assert result.icon_dark == {"type": "emoji", "content": "🔧"} + assert result.label.en_US == "Another Workflow Tool" + assert result.label.zh_Hans == "Another Workflow Tool" + assert result.type == ToolProviderType.WORKFLOW + assert result.workflow_app_id == workflow_app_id + assert result.labels == ["automation", "workflow"] + assert result.masked_credentials == {} + assert result.is_team_authorization is True + assert result.allow_delete is True + assert result.plugin_id is None + assert result.plugin_unique_identifier is None + assert result.tools == [] + + class TestGetToolProviderIconUrl: + def test_builtin_provider_returns_console_url(self, config_overrides): + config_overrides(CONSOLE_API_URL="https://app.dify.ai") + url = ToolTransformService.get_tool_provider_icon_url("builtin", "google", "icon.png") + + assert "/builtin/google/icon" in url + assert url.startswith("https://app.dify.ai/console/api/workspaces/current/tool-provider") + + def test_builtin_provider_with_no_console_url(self, config_overrides): + config_overrides(CONSOLE_API_URL=None) + url = ToolTransformService.get_tool_provider_icon_url("builtin", "slack", "icon.png") + + assert "/builtin/slack/icon" in url + + def test_api_provider_parses_json_icon(self): + icon_json = '{"background": "#fff", "content": "A"}' + result = ToolTransformService.get_tool_provider_icon_url("api", "my-api", icon_json) + assert result == {"background": "#fff", "content": "A"} + + def test_api_provider_returns_dict_icon_directly(self): + icon = {"background": "#000", "content": "B"} + result = ToolTransformService.get_tool_provider_icon_url("api", "my-api", icon) + assert result == icon + + def test_api_provider_returns_fallback_on_invalid_json(self): + result = ToolTransformService.get_tool_provider_icon_url("api", "my-api", "not-json") + assert result == {"background": "#252525", "content": "\ud83d\ude01"} + + def test_workflow_provider_behaves_like_api(self): + icon = {"background": "#123", "content": "W"} + assert ToolTransformService.get_tool_provider_icon_url("workflow", "wf", icon) == icon + + def test_mcp_returns_icon_as_is(self): + assert ToolTransformService.get_tool_provider_icon_url("mcp", "srv", "icon-value") == "icon-value" + + def test_unknown_type_returns_empty(self): + assert ToolTransformService.get_tool_provider_icon_url("unknown", "x", "i") == "" + + class TestRepackProvider: + def test_repacks_dict_provider_icon(self): + provider = {"type": "builtin", "name": "google", "icon": "old"} + with patch.object(ToolTransformService, "get_tool_provider_icon_url", return_value="/new-url") as mock_fn: + ToolTransformService.repack_provider("t1", provider) + + assert provider["icon"] == "/new-url" + mock_fn.assert_called_once_with(provider_type="builtin", provider_name="google", icon="old") + + def test_repacks_tool_provider_api_entity_without_plugin(self): + entity = ToolProviderApiEntity( + id="slack", + author="test_author", + name="slack", + description=I18nObject(en_US="Slack"), + label=I18nObject(en_US="Slack"), + type=ToolProviderType.BUILT_IN, + plugin_id=None, + icon="icon.svg", + icon_dark="dark.svg", + ) + + with patch.object(ToolTransformService, "get_tool_provider_icon_url", return_value="/url"): + ToolTransformService.repack_provider("t1", entity) + + assert entity.icon == "/url" + assert entity.icon_dark == "/url" + + class TestConvertMcpSchemaToParameter: + def test_simple_object_schema(self): + schema = { + "type": "object", + "properties": { + "query": {"type": "string", "description": "Search query"}, + "count": {"type": "integer", "description": "Result count"}, + }, + "required": ["query"], + } + + params = ToolTransformService.convert_mcp_schema_to_parameter(schema) + + assert len(params) == 2 + query_param = next(p for p in params if p.name == "query") + count_param = next(p for p in params if p.name == "count") + assert query_param.required is True + assert count_param.required is False + assert count_param.type.value == "number" + + def test_float_maps_to_number(self): + schema = {"type": "object", "properties": {"rate": {"type": "float"}}, "required": []} + assert ToolTransformService.convert_mcp_schema_to_parameter(schema)[0].type.value == "number" + + def test_array_type_attaches_input_schema(self): + prop = {"type": "array", "description": "Items", "items": {"type": "string"}} + schema = {"type": "object", "properties": {"items": prop}, "required": []} + param = ToolTransformService.convert_mcp_schema_to_parameter(schema)[0] + assert param.input_schema is not None + + def test_non_object_schema_returns_empty(self): + assert ToolTransformService.convert_mcp_schema_to_parameter({"type": "string"}) == [] + + def test_missing_properties_returns_empty(self): + assert ToolTransformService.convert_mcp_schema_to_parameter({"type": "object"}) == [] + + def test_list_type_uses_first_element(self): + schema = {"type": "object", "properties": {"f": {"type": ["string", "null"]}}, "required": []} + assert ToolTransformService.convert_mcp_schema_to_parameter(schema)[0].type.value == "string" + + @pytest.mark.parametrize( + ("prop", "expected_type"), + [ + ({"anyOf": [{"type": "string"}, {"type": "array", "items": {"type": "string"}}]}, "string"), + ({"anyOf": [{"type": "array", "items": {"type": "string"}}, {"type": "string"}]}, "array"), + ({"oneOf": [{"type": "string"}, {"type": "array", "items": {"type": "string"}}]}, "string"), + ({"oneOf": [{"type": "array", "items": {"type": "string"}}, {"type": "string"}]}, "array"), + ({"type": ["string", "array"], "items": {"type": "string"}}, "string"), + ({"type": ["array", "string"], "items": {"type": "string"}}, "array"), + ({"type": ["string", "array", "null"]}, "string"), + ], ) + def test_mixed_type_union_attaches_input_schema(self, prop, expected_type): + schema = {"type": "object", "properties": {"domain": prop}, "required": []} + param = ToolTransformService.convert_mcp_schema_to_parameter(schema)[0] + assert param.type.value == expected_type + assert param.input_schema == prop - # Call the method - result = ToolTransformService.workflow_provider_to_user_provider( - provider_controller=controller, - labels=["automation", "workflow"], - workflow_app_id=workflow_app_id, + @pytest.mark.parametrize( + "prop", + [ + {"type": ["string", "null"]}, + {"anyOf": [{"type": "string"}, {"type": "null"}]}, + {"oneOf": [{"type": "null"}, {"type": "string"}]}, + {"oneOf": [{"type": "integer"}, {"type": "number"}]}, + ], ) + def test_single_type_union_has_no_input_schema(self, prop): + schema = {"type": "object", "properties": {"f": prop}, "required": []} + assert ToolTransformService.convert_mcp_schema_to_parameter(schema)[0].input_schema is None - # Verify all fields are preserved correctly - assert isinstance(result, ToolProviderApiEntity) - assert result.id == provider_id - assert result.author == "another_author" - assert result.name == "another_workflow_tool" - assert result.description.en_US == "Another description" - assert result.description.zh_Hans == "Another description" - assert result.icon == {"type": "emoji", "content": "⚙️"} - assert result.icon_dark == {"type": "emoji", "content": "🔧"} - assert result.label.en_US == "Another Workflow Tool" - assert result.label.zh_Hans == "Another Workflow Tool" - assert result.type == ToolProviderType.WORKFLOW - assert result.workflow_app_id == workflow_app_id - assert result.labels == ["automation", "workflow"] - assert result.masked_credentials == {} - assert result.is_team_authorization is True - assert result.allow_delete is True - assert result.plugin_id is None - assert result.plugin_unique_identifier is None - assert result.tools == [] + def test_missing_description_defaults_empty(self): + schema = {"type": "object", "properties": {"f": {"type": "string"}}, "required": []} + assert ToolTransformService.convert_mcp_schema_to_parameter(schema)[0].llm_description == "" + class TestApiProviderToController: + def test_api_key_header_auth(self): + db_provider = MagicMock() + db_provider.credentials = {"auth_type": "api_key_header"} + with patch(f"{MODULE}.db.session"), patch(f"{MODULE}.ApiToolProviderController") as ctrl_cls: + ctrl_cls.from_db.return_value = MagicMock() + ToolTransformService.api_provider_to_controller(db_provider) + ctrl_cls.from_db.assert_called_once_with( + db_provider=db_provider, auth_type=ApiProviderAuthType.API_KEY_HEADER, session=ANY + ) -class TestGetToolProviderIconUrl: - def test_builtin_provider_returns_console_url(self, config_overrides): - config_overrides(CONSOLE_API_URL="https://app.dify.ai") - url = ToolTransformService.get_tool_provider_icon_url("builtin", "google", "icon.png") + def test_api_key_query_auth(self): + db_provider = MagicMock() + db_provider.credentials = {"auth_type": "api_key_query"} + with patch(f"{MODULE}.db.session"), patch(f"{MODULE}.ApiToolProviderController") as ctrl_cls: + ctrl_cls.from_db.return_value = MagicMock() + ToolTransformService.api_provider_to_controller(db_provider) + ctrl_cls.from_db.assert_called_once_with( + db_provider=db_provider, auth_type=ApiProviderAuthType.API_KEY_QUERY, session=ANY + ) - assert "/builtin/google/icon" in url - assert url.startswith("https://app.dify.ai/console/api/workspaces/current/tool-provider") + def test_legacy_api_key_maps_to_header(self): + db_provider = MagicMock() + db_provider.credentials = {"auth_type": "api_key"} + with patch(f"{MODULE}.db.session"), patch(f"{MODULE}.ApiToolProviderController") as ctrl_cls: + ctrl_cls.from_db.return_value = MagicMock() + ToolTransformService.api_provider_to_controller(db_provider) + ctrl_cls.from_db.assert_called_once_with( + db_provider=db_provider, auth_type=ApiProviderAuthType.API_KEY_HEADER, session=ANY + ) - def test_builtin_provider_with_no_console_url(self, config_overrides): - config_overrides(CONSOLE_API_URL=None) - url = ToolTransformService.get_tool_provider_icon_url("builtin", "slack", "icon.png") - - assert "/builtin/slack/icon" in url - - def test_api_provider_parses_json_icon(self): - icon_json = '{"background": "#fff", "content": "A"}' - result = ToolTransformService.get_tool_provider_icon_url("api", "my-api", icon_json) - assert result == {"background": "#fff", "content": "A"} - - def test_api_provider_returns_dict_icon_directly(self): - icon = {"background": "#000", "content": "B"} - result = ToolTransformService.get_tool_provider_icon_url("api", "my-api", icon) - assert result == icon - - def test_api_provider_returns_fallback_on_invalid_json(self): - result = ToolTransformService.get_tool_provider_icon_url("api", "my-api", "not-json") - assert result == {"background": "#252525", "content": "\ud83d\ude01"} - - def test_workflow_provider_behaves_like_api(self): - icon = {"background": "#123", "content": "W"} - assert ToolTransformService.get_tool_provider_icon_url("workflow", "wf", icon) == icon - - def test_mcp_returns_icon_as_is(self): - assert ToolTransformService.get_tool_provider_icon_url("mcp", "srv", "icon-value") == "icon-value" - - def test_unknown_type_returns_empty(self): - assert ToolTransformService.get_tool_provider_icon_url("unknown", "x", "i") == "" - - -class TestRepackProvider: - def test_repacks_dict_provider_icon(self): - provider = {"type": "builtin", "name": "google", "icon": "old"} - with patch.object(ToolTransformService, "get_tool_provider_icon_url", return_value="/new-url") as mock_fn: - ToolTransformService.repack_provider("t1", provider) - - assert provider["icon"] == "/new-url" - mock_fn.assert_called_once_with(provider_type="builtin", provider_name="google", icon="old") - - def test_repacks_tool_provider_api_entity_without_plugin(self): - entity = ToolProviderApiEntity( - id="slack", - author="test_author", - name="slack", - description=I18nObject(en_US="Slack"), - label=I18nObject(en_US="Slack"), - type=ToolProviderType.BUILT_IN, - plugin_id=None, - icon="icon.svg", - icon_dark="dark.svg", - ) - - with patch.object(ToolTransformService, "get_tool_provider_icon_url", return_value="/url"): - ToolTransformService.repack_provider("t1", entity) - - assert entity.icon == "/url" - assert entity.icon_dark == "/url" - - -class TestConvertMcpSchemaToParameter: - def test_simple_object_schema(self): - schema = { - "type": "object", - "properties": { - "query": {"type": "string", "description": "Search query"}, - "count": {"type": "integer", "description": "Result count"}, - }, - "required": ["query"], - } - - params = ToolTransformService.convert_mcp_schema_to_parameter(schema) - - assert len(params) == 2 - query_param = next(p for p in params if p.name == "query") - count_param = next(p for p in params if p.name == "count") - assert query_param.required is True - assert count_param.required is False - assert count_param.type.value == "number" - - def test_float_maps_to_number(self): - schema = {"type": "object", "properties": {"rate": {"type": "float"}}, "required": []} - assert ToolTransformService.convert_mcp_schema_to_parameter(schema)[0].type.value == "number" - - def test_array_type_attaches_input_schema(self): - prop = {"type": "array", "description": "Items", "items": {"type": "string"}} - schema = {"type": "object", "properties": {"items": prop}, "required": []} - param = ToolTransformService.convert_mcp_schema_to_parameter(schema)[0] - assert param.input_schema is not None - - def test_non_object_schema_returns_empty(self): - assert ToolTransformService.convert_mcp_schema_to_parameter({"type": "string"}) == [] - - def test_missing_properties_returns_empty(self): - assert ToolTransformService.convert_mcp_schema_to_parameter({"type": "object"}) == [] - - def test_list_type_uses_first_element(self): - schema = {"type": "object", "properties": {"f": {"type": ["string", "null"]}}, "required": []} - assert ToolTransformService.convert_mcp_schema_to_parameter(schema)[0].type.value == "string" - - @pytest.mark.parametrize( - ("prop", "expected_type"), - [ - ({"anyOf": [{"type": "string"}, {"type": "array", "items": {"type": "string"}}]}, "string"), - ({"anyOf": [{"type": "array", "items": {"type": "string"}}, {"type": "string"}]}, "array"), - ({"oneOf": [{"type": "string"}, {"type": "array", "items": {"type": "string"}}]}, "string"), - ({"oneOf": [{"type": "array", "items": {"type": "string"}}, {"type": "string"}]}, "array"), - ({"type": ["string", "array"], "items": {"type": "string"}}, "string"), - ({"type": ["array", "string"], "items": {"type": "string"}}, "array"), - ({"type": ["string", "array", "null"]}, "string"), - ], - ) - def test_mixed_type_union_attaches_input_schema(self, prop, expected_type): - schema = {"type": "object", "properties": {"domain": prop}, "required": []} - param = ToolTransformService.convert_mcp_schema_to_parameter(schema)[0] - assert param.type.value == expected_type - assert param.input_schema == prop - - @pytest.mark.parametrize( - "prop", - [ - {"type": ["string", "null"]}, - {"anyOf": [{"type": "string"}, {"type": "null"}]}, - {"oneOf": [{"type": "null"}, {"type": "string"}]}, - {"oneOf": [{"type": "integer"}, {"type": "number"}]}, - ], - ) - def test_single_type_union_has_no_input_schema(self, prop): - schema = {"type": "object", "properties": {"f": prop}, "required": []} - assert ToolTransformService.convert_mcp_schema_to_parameter(schema)[0].input_schema is None - - def test_missing_description_defaults_empty(self): - schema = {"type": "object", "properties": {"f": {"type": "string"}}, "required": []} - assert ToolTransformService.convert_mcp_schema_to_parameter(schema)[0].llm_description == "" - - -class TestApiProviderToController: - def test_api_key_header_auth(self): - db_provider = MagicMock() - db_provider.credentials = {"auth_type": "api_key_header"} - with patch(f"{MODULE}.db.session"), patch(f"{MODULE}.ApiToolProviderController") as ctrl_cls: - ctrl_cls.from_db.return_value = MagicMock() - ToolTransformService.api_provider_to_controller(db_provider) - ctrl_cls.from_db.assert_called_once_with( - db_provider=db_provider, auth_type=ApiProviderAuthType.API_KEY_HEADER, session=ANY - ) - - def test_api_key_query_auth(self): - db_provider = MagicMock() - db_provider.credentials = {"auth_type": "api_key_query"} - with patch(f"{MODULE}.db.session"), patch(f"{MODULE}.ApiToolProviderController") as ctrl_cls: - ctrl_cls.from_db.return_value = MagicMock() - ToolTransformService.api_provider_to_controller(db_provider) - ctrl_cls.from_db.assert_called_once_with( - db_provider=db_provider, auth_type=ApiProviderAuthType.API_KEY_QUERY, session=ANY - ) - - def test_legacy_api_key_maps_to_header(self): - db_provider = MagicMock() - db_provider.credentials = {"auth_type": "api_key"} - with patch(f"{MODULE}.db.session"), patch(f"{MODULE}.ApiToolProviderController") as ctrl_cls: - ctrl_cls.from_db.return_value = MagicMock() - ToolTransformService.api_provider_to_controller(db_provider) - ctrl_cls.from_db.assert_called_once_with( - db_provider=db_provider, auth_type=ApiProviderAuthType.API_KEY_HEADER, session=ANY - ) - - def test_unknown_auth_defaults_to_none(self): - db_provider = MagicMock() - db_provider.credentials = {"auth_type": "something_else"} - with patch(f"{MODULE}.db.session"), patch(f"{MODULE}.ApiToolProviderController") as ctrl_cls: - ctrl_cls.from_db.return_value = MagicMock() - ToolTransformService.api_provider_to_controller(db_provider) - ctrl_cls.from_db.assert_called_once_with( - db_provider=db_provider, auth_type=ApiProviderAuthType.NONE, session=ANY - ) + def test_unknown_auth_defaults_to_none(self): + db_provider = MagicMock() + db_provider.credentials = {"auth_type": "something_else"} + with patch(f"{MODULE}.db.session"), patch(f"{MODULE}.ApiToolProviderController") as ctrl_cls: + ctrl_cls.from_db.return_value = MagicMock() + ToolTransformService.api_provider_to_controller(db_provider) + ctrl_cls.from_db.assert_called_once_with( + db_provider=db_provider, auth_type=ApiProviderAuthType.NONE, session=ANY + ) diff --git a/api/tests/unit_tests/services/tracing/test_application_trace_collection.py b/api/tests/unit_tests/services/tracing/test_application_trace_collection.py deleted file mode 100644 index 84f421cc9f2..00000000000 --- a/api/tests/unit_tests/services/tracing/test_application_trace_collection.py +++ /dev/null @@ -1,519 +0,0 @@ -"""Exercise collection transactions and delivery state against real SQLite rows.""" - -import json -from dataclasses import dataclass -from datetime import UTC, datetime, timedelta -from decimal import Decimal -from uuid import NAMESPACE_URL, UUID, uuid4, uuid5 - -import pytest -from dify_agent.protocol.tracing import AgentTraceBatch, AgentTraceEvent, AgentTraceSpan -from sqlalchemy import func, select -from sqlalchemy.orm import Session, sessionmaker - -from core.ops.entities.trace_entity import MessageTraceInfo -from core.ops.exceptions import RetryableTraceDispatchError -from core.ops.unified_trace.entities import CanonicalSpan, CanonicalSpanKind, CanonicalSpanStatus, CanonicalTrace -from libs.datetime_utils import naive_utc_now -from models.application_trace import ApplicationTraceBatch, ApplicationTraceInvocation, ApplicationTraceRun -from models.enums import ConversationFromSource -from models.model import App, AppMode, Message, TraceAppConfig -from repositories.application_trace_repository import ApplicationTraceRepository, TraceCollectionError -from services.tracing import service as service_module -from services.tracing.provider_dispatch import TraceExportResult, TraceExportUncertainError -from services.tracing.service import ApplicationTracingService -from tests.unit_tests.config_override import apply_config_overrides -from tests.unit_tests.model_factories import make_app, make_message - - -def identifier(value: str) -> str: - return str(uuid5(NAMESPACE_URL, value)) - - -@dataclass -class Collection: - repository: ApplicationTraceRepository - sessions: sessionmaker[Session] - run: ApplicationTraceRun - - def invocation(self, name: str = "first") -> ApplicationTraceInvocation: - return self.repository.invocation( - run_id=self.run.id, agent_id=identifier("agent"), node_execution_id=None, attempt=0, call_key=name - ) - - -def seed( - sessions: sessionmaker[Session], - name: str, - *, - capture_content: bool = False, - enabled: bool = True, - capture_agent: bool = True, -) -> Collection: - app_id, tenant_id, message_id = identifier(f"app-{name}"), identifier("tenant"), identifier(f"message-{name}") - with sessions.begin() as session: - session.add( - make_app( - app_id=app_id, - tenant_id=tenant_id, - mode=AppMode.AGENT, - tracing=json.dumps( - { - "enabled": enabled, - "tracing_provider": "langfuse", - "capture_content": capture_content, - "capture_agent": capture_agent, - } - ), - ) - ) - session.add( - make_message( - message_id=message_id, - app_id=app_id, - conversation_id=identifier(f"conversation-{name}"), - inputs={}, - query="hello", - message={}, - answer="done", - message_unit_price=Decimal(0), - answer_unit_price=Decimal(0), - currency="USD", - from_source=ConversationFromSource.CONSOLE, - ) - ) - session.add( - TraceAppConfig( - app_id=app_id, - tracing_provider="langfuse", - tracing_config={ - "host": "https://langfuse.test", - "secret_key": "encrypted-original", - "project_url": "https://langfuse.test/project/original", - }, - ) - ) - repository = ApplicationTraceRepository(sessions) - run = repository.start( - tenant_id=tenant_id, app_id=app_id, source_kind="message", source_id=message_id, session_id="session-original" - ) - assert run is not None - return Collection(repository, sessions, run) - - -@pytest.fixture -def collection(sqlite_session_factory: sessionmaker[Session]) -> Collection: - return seed(sqlite_session_factory, "first") - - -def batch_for( - run: ApplicationTraceRun, - invocation: ApplicationTraceInvocation, - *, - final: bool = True, - batch_id: UUID | None = None, - span_id: str = "a" * 16, -) -> AgentTraceBatch: - return AgentTraceBatch( - run_id=UUID(run.id), - invocation_id=UUID(invocation.id), - batch_id=batch_id or uuid4(), - final=final, - spans=[ - AgentTraceSpan( - trace_id=run.trace_id, - span_id=span_id, - parent_span_id=invocation.span_id, - name="Agent execution", - start_time_ns=1_000_000_000, - end_time_ns=2_000_000_000, - status="ok", - attributes={"gen_ai.operation.name": "invoke_agent"}, - ) - ], - ) - - -def ingest(collection: Collection, batch: AgentTraceBatch) -> tuple[int, str]: - return collection.repository.ingest(batch, max_spans=100, max_bytes=1_000_000) - - -def terminal(collection: Collection) -> None: - run = collection.run - root = CanonicalSpan( - id=run.root_span_id, - parent_id=None, - name="Agent turn", - kind=CanonicalSpanKind.CHAIN, - start_time=datetime.fromtimestamp(0, UTC), - end_time=datetime.fromtimestamp(3, UTC), - status=CanonicalSpanStatus.OK, - ) - trace = CanonicalTrace(trace_id=run.trace_id, session_id=run.session_id, root_span_id=root.id, spans=(root,)) - collection.repository.terminal(run.id, trace.model_dump(mode="json"), root.id) - - -def claim(collection: Collection) -> tuple[ApplicationTraceRun, str] | None: - return collection.repository.claim(collection.run.id, collection_timeout=60, lease_seconds=30) - - -def test_batch_duplicate_is_idempotent_after_final_and_acceptance(collection: Collection) -> None: - invocation = collection.invocation() - batch = batch_for(collection.run, invocation) - assert ingest(collection, batch) == (1, collection.run.id) - assert ingest(collection, batch) == (1, collection.run.id) - terminal(collection) - leased = claim(collection) - assert leased is not None - collection.repository.complete( - collection.run.id, leased[1], state="accepted", result=TraceExportResult("provider-id") - ) - assert ingest(collection, batch) == (1, collection.run.id) - stored = collection.repository.get(collection.run.id) - assert stored is not None - assert stored.span_count == 1 - with collection.sessions() as session: - assert session.scalar(select(func.count()).select_from(ApplicationTraceBatch)) == 1 - - -def test_conflicting_batch_id_rolls_back_without_replacing_payload(collection: Collection) -> None: - batch = batch_for(collection.run, collection.invocation(), final=False) - ingest(collection, batch) - changed = batch.model_copy(update={"final": True}) - with pytest.raises(TraceCollectionError, match="tracing_batch_conflict"): - ingest(collection, changed) - invocations, batches = collection.repository.collected(collection.run.id) - assert invocations[0].finalized is False - assert batches[0].final is False - - -def test_run_cannot_ingest_another_runs_invocation(collection: Collection) -> None: - other = seed(collection.sessions, "other") - foreign = other.invocation() - forged = batch_for(collection.run, foreign) - with pytest.raises(TraceCollectionError, match="tracing_invocation_not_found"): - ingest(collection, forged) - with collection.sessions() as session: - assert session.scalar(select(func.count()).select_from(ApplicationTraceBatch)) == 0 - - -def test_payload_limits_reject_whole_transaction(collection: Collection) -> None: - batch = batch_for(collection.run, collection.invocation()) - with pytest.raises(TraceCollectionError, match="tracing_payload_limit"): - collection.repository.ingest(batch, max_spans=100, max_bytes=1) - invocations, batches = collection.repository.collected(collection.run.id) - assert not batches - assert invocations[0].finalized is False - stored = collection.repository.get(collection.run.id) - assert stored is not None - assert stored.span_count == 0 - - -def test_wrong_trace_id_is_rejected_before_persistence(collection: Collection) -> None: - batch = batch_for(collection.run, collection.invocation()) - altered = batch.model_copy(update={"spans": [batch.spans[0].model_copy(update={"trace_id": "f" * 32})]}) - with pytest.raises(TraceCollectionError, match="tracing_trace_id_mismatch"): - ingest(collection, altered) - _, batches = collection.repository.collected(collection.run.id) - assert not batches - - -def test_content_off_cleans_the_durable_batch_not_only_the_export(collection: Collection) -> None: - original = batch_for(collection.run, collection.invocation(), final=False) - observation = original.spans[0].model_copy( - update={ - "status": "error", - "status_message": "private error", - "attributes": { - "gen_ai.operation.name": "chat", - "gen_ai.request.model": "model-one", - "gen_ai.input.messages": "private prompt", - "gen_ai.tool.call.arguments": "private arguments", - "dify.payload": "private metadata", - "authorization": "private credential", - }, - "events": [ - AgentTraceEvent( - name="exception", timestamp_ns=1_500_000_000, attributes={"exception.message": "private event"} - ) - ], - } - ) - batch = original.model_copy(update={"spans": [observation]}) - ingest(collection, batch) - _, batches = collection.repository.collected(collection.run.id) - assert "private" not in json.dumps(batches[0].spans) - assert batches[0].spans[0]["events"] == [] - assert batches[0].spans[0]["status"] == "error" - assert batches[0].spans[0]["attributes"]["gen_ai.request.model"] == "model-one" - # A changed original payload is a conflict even if content-off sanitization - # would produce exactly the same persisted span. - changed = batch.model_copy( - update={"spans": [observation.model_copy(update={"status_message": "different private error"})]} - ) - with pytest.raises(TraceCollectionError, match="tracing_batch_conflict"): - ingest(collection, changed) - - -def test_invocation_key_replay_must_not_change_agent_binding(collection: Collection) -> None: - original = collection.invocation() - repeated = collection.invocation() - assert original.id == repeated.id - with pytest.raises(TraceCollectionError, match="tracing_invocation_conflict"): - collection.repository.invocation( - run_id=collection.run.id, - agent_id=identifier("different-agent"), - node_execution_id=identifier("different-node"), - attempt=2, - call_key="first", - ) - - -def test_claim_waits_for_all_invocations_final_and_then_only_one_lease(collection: Collection) -> None: - first, second = collection.invocation("first"), collection.invocation("second") - terminal(collection) - ingest(collection, batch_for(collection.run, first)) - assert claim(collection) is None - ingest(collection, batch_for(collection.run, second, span_id="b" * 16)) - leased = claim(collection) - assert leased is not None - assert claim(collection) is None - assert leased[0].attempts == 1 - - -def test_expired_delivery_lease_accepts_confirmed_late_ack_for_same_attempt(collection: Collection) -> None: - ingest(collection, batch_for(collection.run, collection.invocation())) - terminal(collection) - leased = claim(collection) - assert leased is not None - with collection.sessions.begin() as session: - stored = session.get(ApplicationTraceRun, collection.run.id) - assert stored is not None - stored.export_started_at = naive_utc_now() - timedelta(seconds=60) - stored.export_heartbeat_at = stored.export_started_at - assert claim(collection) is None - current = collection.repository.get(collection.run.id) - assert current is not None - assert current.state == "uncertain" - collection.repository.complete(collection.run.id, leased[1], state="accepted", result=TraceExportResult("late")) - stored = collection.repository.get(collection.run.id) - assert stored is not None - assert stored.state == "accepted" - assert stored.error_code is None - assert stored.accepted_at is not None - assert stored.provider_trace_id == "late" - assert claim(collection) is None - - -def test_incomplete_final_times_out_without_export(collection: Collection) -> None: - collection.invocation() - terminal(collection) - with collection.sessions.begin() as session: - stored = session.get(ApplicationTraceRun, collection.run.id) - assert stored is not None - stored.terminal_at = naive_utc_now() - timedelta(seconds=120) - assert claim(collection) is None - stored = collection.repository.get(collection.run.id) - assert stored is not None - assert stored.state == "failed" - assert stored.error_code == "agent_trace_collection_incomplete" - - -def test_running_trace_freezes_credentials_content_policy_and_session(collection: Collection) -> None: - original = collection.run - with collection.sessions.begin() as session: - app = session.get(App, original.app_id) - assert app is not None - app.tracing = json.dumps({"enabled": False, "tracing_provider": "opik", "capture_content": True}) - config = session.scalar(select(TraceAppConfig).where(TraceAppConfig.app_id == original.app_id)) - assert config is not None - config.tracing_config = {"secret_key": "encrypted-new"} - repeated = collection.repository.start( - tenant_id=original.tenant_id, - app_id=original.app_id, - source_kind="message", - source_id=original.source_id, - session_id="changed", - ) - assert repeated is not None - assert repeated.provider == "langfuse" - assert repeated.provider_config["secret_key"] == "encrypted-original" - assert repeated.capture_content is False - assert repeated.config_revision == original.config_revision - assert repeated.session_id == "session-original" - - -@pytest.mark.parametrize("outcome", ["accepted", "uncertain", "retrying"]) -def test_delivery_checkpoint_controls_replays( - collection: Collection, monkeypatch: pytest.MonkeyPatch, outcome: str -) -> None: - ingest(collection, batch_for(collection.run, collection.invocation())) - terminal(collection) - exports: list[tuple[str, dict[str, object], CanonicalTrace]] = [] - decryptions: list[dict[str, object]] = [] - - def decrypt(_tenant: str, _provider: str, config: dict[str, object]) -> dict[str, object]: - decryptions.append(config) - return {"secret_key": "plaintext-only-in-worker"} - - def export(provider: str, config: dict[str, object], trace: CanonicalTrace) -> TraceExportResult: - exports.append((provider, config, trace)) - if outcome == "uncertain": - raise TraceExportUncertainError("ack lost") - if outcome == "retrying": - raise RetryableTraceDispatchError("busy") - return TraceExportResult("vendor-trace", "https://trace.test/example") - - monkeypatch.setattr(service_module.OpsTraceManager, "decrypt_tracing_config", decrypt) - monkeypatch.setattr(service_module, "export_trace", export) - service = ApplicationTracingService(collection.repository) - service.deliver(collection.run.id) - service.deliver(collection.run.id) - assert len(exports) == 1 - assert decryptions[0]["secret_key"] == "encrypted-original" - stored = collection.repository.get(collection.run.id) - assert stored is not None - assert stored.state == outcome - assert stored.provider_config["secret_key"] == "encrypted-original" - assert "plaintext-only-in-worker" not in json.dumps(stored.outer_trace) - assert exports[0][1]["host"] == "https://langfuse.test" - - -def test_source_owner_checks_app_and_tenant(collection: Collection) -> None: - run = collection.run - assert ( - collection.repository.start( - tenant_id=identifier("foreign-tenant"), - app_id=run.app_id, - source_kind="message", - source_id=run.source_id, - session_id="x", - ) - is None - ) - other = seed(collection.sessions, "second-app") - assert ( - collection.repository.start( - tenant_id=run.tenant_id, - app_id=other.run.app_id, - source_kind="message", - source_id=run.source_id, - session_id="x", - ) - is None - ) - - -def test_disabled_source_decision_stays_disabled_after_config_enable( - sqlite_session_factory: sessionmaker[Session], monkeypatch: pytest.MonkeyPatch -) -> None: - apply_config_overrides(monkeypatch, OPS_AGENT_TRACE_ENABLED=True) - case = seed(sqlite_session_factory, "initially-disabled", enabled=False) - service = ApplicationTracingService(case.repository) - assert case.run.state == "not_collected" - with case.sessions.begin() as session: - app = session.get(App, case.run.app_id) - assert app is not None - app.tracing = json.dumps({"enabled": True, "tracing_provider": "langfuse", "capture_content": True}) - message = session.get(Message, case.run.source_id) - assert message is not None - assert ( - service.start( - tenant_id=case.run.tenant_id, - app_id=case.run.app_id, - source_kind="message", - source_id=case.run.source_id, - session_id="later-session", - ) - is None - ) - info = MessageTraceInfo( - metadata={}, - message_data=message, - message_id=message.id, - conversation_model="agent", - conversation_mode="agent", - message_tokens=1, - answer_tokens=1, - total_tokens=2, - ) - assert service.capture_terminal(info) is True - stored = case.repository.get(case.run.id) - assert stored is not None - assert stored.state == "not_collected" - assert stored.provider_config == {} - assert stored.outer_trace is None - - -def test_capture_agent_disabled_preserves_outer_only_content_policy( - sqlite_session_factory: sessionmaker[Session], monkeypatch: pytest.MonkeyPatch -) -> None: - apply_config_overrides(monkeypatch, OPS_AGENT_TRACE_ENABLED=True) - case = seed(sqlite_session_factory, "outer-only", capture_agent=False) - service = ApplicationTracingService(case.repository) - scheduled: list[str] = [] - monkeypatch.setattr(service, "schedule", scheduled.append) - assert ( - service.prepare_invocation( - tenant_id=case.run.tenant_id, - app_id=case.run.app_id, - agent_id=identifier("agent"), - source_kind="message", - source_id=case.run.source_id, - session_id="session", - node_execution_id=None, - attempt=0, - call_key="call", - ) - is None - ) - with case.sessions() as session: - message = session.get(Message, case.run.source_id) - assert message is not None - info = MessageTraceInfo( - metadata={}, - message_data=message, - message_id=message.id, - conversation_model="agent", - conversation_mode="agent", - message_tokens=1, - answer_tokens=1, - total_tokens=2, - inputs="private input", - outputs="private output", - ) - assert service.capture_terminal(info) is True - assert scheduled == [case.run.id] - stored = case.repository.get(case.run.id) - assert stored is not None - assert stored.outer_trace is not None - assert "private" not in json.dumps(stored.outer_trace) - assert len(stored.outer_trace["spans"]) == 1 - invocations, _ = case.repository.collected(case.run.id) - assert invocations == [] - - -def test_deleting_selected_configuration_keeps_historical_destination(collection: Collection) -> None: - from machinery.context import RequestContext - from repositories.app.tracing_config_repository import SQLAlchemyAppTracingConfigRepository - from services.app_tracing_config_gateway import OpsTraceManagerGateway - from services.app_tracing_config_service import AppTracingConfigService - - original = collection.run - service = AppTracingConfigService( - configs=SQLAlchemyAppTracingConfigRepository(session_factory=collection.sessions), - provider=OpsTraceManagerGateway(), - ) - service.delete(RequestContext("request", None, "account", original.tenant_id), original.app_id, "langfuse") - stored = collection.repository.get(original.id) - assert stored is not None - assert stored.provider_config == original.provider_config - assert stored.provider_config["project_url"] == "https://langfuse.test/project/original" - assert stored.config_revision == original.config_revision - with collection.sessions() as session: - assert session.scalar(select(TraceAppConfig).where(TraceAppConfig.app_id == original.app_id)) is None - app = session.get(App, original.app_id) - assert app is not None - assert app.tracing is not None - assert json.loads(app.tracing)["enabled"] is False - assert json.loads(app.tracing)["tracing_provider"] is None diff --git a/api/tests/unit_tests/services/tracing/test_application_trace_handoff.py b/api/tests/unit_tests/services/tracing/test_application_trace_handoff.py deleted file mode 100644 index 55ed619edfa..00000000000 --- a/api/tests/unit_tests/services/tracing/test_application_trace_handoff.py +++ /dev/null @@ -1,468 +0,0 @@ -"""SQLite ownership regressions at the actual legacy-task handoff boundary.""" - -import json -from dataclasses import dataclass -from datetime import datetime, timedelta -from uuid import UUID, uuid4 - -import httpx -import pytest -from dify_agent.protocol.tracing import AgentTraceBatch, AgentTraceSpan -from flask import Flask -from opentelemetry.proto.collector.trace.v1.trace_service_pb2 import ExportTraceServiceRequest -from sqlalchemy import delete, func, select -from sqlalchemy.exc import OperationalError -from sqlalchemy.orm import Session, sessionmaker - -from core.app.apps.agent_app.runtime_request_builder import AgentAppRuntimeBuildContext -from core.app.entities.app_invoke_entities import DifyRunContext, InvokeFrom, UserFrom -from core.ops.base_trace_instance import bind_application_tracing_hooks -from core.ops.entities.trace_entity import ( - BaseTraceInfo, - GenerateNameTraceInfo, - MessageTraceInfo, - ModerationTraceInfo, - WorkflowTraceInfo, -) -from core.ops.ops_trace_manager import OpsTraceManager -from extensions.ext_database import db -from graphon.enums import WorkflowExecutionStatus, WorkflowNodeExecutionStatus -from models.agent_config_entities import AgentSoulConfig -from models.application_trace import ApplicationTraceRun -from models.enums import ConversationFromSource, CreatorUserRole, WorkflowRunTriggeredFrom -from models.model import App, AppMode, Message -from models.workflow import WorkflowNodeExecutionModel, WorkflowNodeExecutionTriggeredFrom, WorkflowRun, WorkflowType -from services.tracing import runtime_bridge -from services.tracing.provider_dispatch import export_trace -from services.tracing.service import ApplicationTracingService -from tasks import ops_trace_task -from tests.unit_tests.config_override import apply_config_overrides -from tests.unit_tests.model_factories import make_account, make_conversation, make_workflow -from tests.unit_tests.services.tracing.test_application_trace_collection import Collection, identifier, seed, terminal - -AGENT_GRAPH = { - "nodes": [{"id": "agent", "data": {"type": "agent", "version": "2", "agent_node_kind": "dify_agent"}}], - "edges": list[object](), -} - - -@dataclass -class Case: - collection: Collection - service: ApplicationTracingService - scheduled: list[str] - - -def setup( - sessions: sessionmaker[Session], monkeypatch: pytest.MonkeyPatch, name: str, mode: AppMode, *, snapshot: bool -) -> Case: - apply_config_overrides(monkeypatch, OPS_AGENT_TRACE_ENABLED=True) - collection = seed(sessions, name) - with sessions.begin() as session: - app = session.get(App, collection.run.app_id) - message = session.get(Message, collection.run.source_id) - assert app is not None - assert message is not None - app.mode = mode - session.add( - make_conversation( - conversation_id=message.conversation_id, - app_id=app.id, - mode=mode, - inputs={}, - from_source=ConversationFromSource.CONSOLE, - ) - ) - if not snapshot: - session.execute(delete(ApplicationTraceRun).where(ApplicationTraceRun.id == collection.run.id)) - service = ApplicationTracingService(collection.repository) - scheduled: list[str] = [] - monkeypatch.setattr(service, "schedule", scheduled.append) - monkeypatch.setattr(runtime_bridge, "_service", lambda: service) - return Case(collection, service, scheduled) - - -def message_info(case: Case) -> MessageTraceInfo: - with case.collection.sessions() as session: - message = session.get(Message, case.collection.run.source_id) - assert message is not None - return MessageTraceInfo( - metadata={"app_id": message.app_id}, - message_id=message.id, - message_data=message.to_dict(session=session), - conversation_model="agent", - conversation_mode="agent", - message_tokens=0, - answer_tokens=0, - total_tokens=0, - inputs="private input", - outputs="private output", - ) - - -def dispatch(case: Case, info: BaseTraceInfo, monkeypatch: pytest.MonkeyPatch) -> list[BaseTraceInfo]: - calls: list[BaseTraceInfo] = [] - failures: list[str] = [] - - class Legacy: - def trace(self, restored: BaseTraceInfo) -> None: - calls.append(restored) - - payload = json.dumps({"trace_info_type": type(info).__name__, "trace_info": info.model_dump(mode="json")}) - monkeypatch.setattr(ops_trace_task.storage, "load", lambda _path: payload) - monkeypatch.setattr(ops_trace_task.storage, "delete", lambda _path: None) - monkeypatch.setattr(ops_trace_task.redis_client, "incr", failures.append) - monkeypatch.setattr(OpsTraceManager, "get_ops_trace_instance", lambda _app: Legacy()) - monkeypatch.setattr("extensions.ext_enterprise_telemetry.is_enabled", lambda: False) - app = Flask(__name__) - bind_application_tracing_hooks(app, runtime_bridge.RuntimeApplicationTracingHooks()) - with app.app_context(): - ops_trace_task.process_trace_tasks.run({"app_id": case.collection.run.app_id, "file_id": "fixture"}) - assert failures == [] - return calls - - -def workflow( - case: Case, name: str, *, has_agent: bool, snapshot: bool -) -> tuple[WorkflowRun, ApplicationTraceRun | None]: - graph = AGENT_GRAPH if has_agent else {"nodes": [{"data": {"type": "llm"}}], "edges": list[object]()} - now = datetime(2026, 10, 15) - with case.collection.sessions.begin() as session: - definition = make_workflow( - workflow_id=identifier(f"workflow-{name}"), - tenant_id=case.collection.run.tenant_id, - app_id=case.collection.run.app_id, - graph=graph, - ) - session.add(definition) - app = session.get(App, case.collection.run.app_id) - assert app is not None - app.workflow_id = definition.id - row = WorkflowRun( - id=identifier(f"workflow-run-{name}"), - app_id=app.id, - tenant_id=app.tenant_id, - workflow_id=definition.id, - type=WorkflowType.WORKFLOW, - version="1", - graph=json.dumps(graph), - inputs="{}", - outputs="{}", - status=WorkflowExecutionStatus.SUCCEEDED, - triggered_from=WorkflowRunTriggeredFrom.APP_RUN, - created_by_role=CreatorUserRole.ACCOUNT, - created_by=identifier("owner"), - created_at=now, - finished_at=now + timedelta(seconds=2), - elapsed_time=2, - total_tokens=0, - total_steps=0, - ) - session.add(row) - trace_run = ( - case.collection.repository.start( - tenant_id=row.tenant_id, - app_id=row.app_id, - source_kind="workflow", - source_id=row.id, - session_id="workflow-session", - ) - if snapshot - else None - ) - return row, trace_run - - -def workflow_info(row: WorkflowRun, parent: WorkflowRun | None = None) -> WorkflowTraceInfo: - metadata: dict[str, object] = {"app_id": row.app_id} - if parent is not None: - metadata["parent_trace_context"] = { - "parent_workflow_run_id": parent.id, - "parent_node_execution_id": identifier("parent-tool-node"), - } - return WorkflowTraceInfo( - workflow_data=row, - workflow_id=row.workflow_id, - tenant_id=row.tenant_id, - workflow_run_id=row.id, - workflow_run_elapsed_time=2, - workflow_run_status="succeeded", - workflow_run_inputs={}, - workflow_run_outputs={}, - workflow_run_version="1", - total_tokens=0, - file_list=[], - query="", - metadata=metadata, - start_time=row.created_at, - end_time=row.finished_at, - ) - - -def test_chatflow_message_event_cannot_export_again_via_legacy( - sqlite_session_factory: sessionmaker[Session], monkeypatch: pytest.MonkeyPatch -) -> None: - case = setup(sqlite_session_factory, monkeypatch, "chatflow", AppMode.ADVANCED_CHAT, snapshot=False) - row, trace_run = workflow(case, "chatflow", has_agent=True, snapshot=True) - assert trace_run is not None - terminal(Collection(case.collection.repository, case.collection.sessions, trace_run)) - before = case.collection.repository.get(trace_run.id) - assert before is not None - with case.collection.sessions.begin() as session: - message = session.get(Message, case.collection.run.source_id) - assert message is not None - message.workflow_run_id = row.id - assert dispatch(case, message_info(case), monkeypatch) == [] - after = case.collection.repository.get(trace_run.id) - assert after is not None - assert after.outer_trace == before.outer_trace - assert case.scheduled == [] - with case.collection.sessions() as session: - assert session.scalar(select(func.count()).select_from(ApplicationTraceRun)) == 1 - - -@pytest.mark.parametrize("mode", [AppMode.AGENT, AppMode.ADVANCED_CHAT]) -def test_snapshot_failure_is_persisted_and_never_falls_back( - sqlite_session_factory: sessionmaker[Session], monkeypatch: pytest.MonkeyPatch, mode: AppMode -) -> None: - case = setup(sqlite_session_factory, monkeypatch, "failure", mode, snapshot=False) - - def fail_start(**_kwargs: object) -> None: - raise OperationalError("snapshot", {}, Exception("temporary database unavailable")) - - monkeypatch.setattr(case.collection.repository, "start", fail_start) - context = AgentAppRuntimeBuildContext( - dify_context=DifyRunContext( - tenant_id=case.collection.run.tenant_id, - app_id=case.collection.run.app_id, - user_id=identifier("user"), - user_from=UserFrom.ACCOUNT, - invoke_from=InvokeFrom.DEBUGGER, - ), - agent_id=identifier("agent"), - idempotency_key=case.collection.run.source_id, - conversation_id="conversation", - agent_config_snapshot_id=identifier("snapshot"), - agent_soul=AgentSoulConfig(), - user_query="query", - binding_id=identifier("binding"), - backend_binding_ref="backend", - ) - assert runtime_bridge.prepare_agent_app(context) is None - assert dispatch(case, message_info(case), monkeypatch) == [] - stored = case.collection.repository.source_run( - tenant_id=case.collection.run.tenant_id, - app_id=case.collection.run.app_id, - source_kind="message", - source_id=case.collection.run.source_id, - ) - assert stored is not None - assert stored.state == "failed" - assert stored.error_code == "tracing_initialization_failed" - assert stored.provider_config == {} - - -@pytest.mark.parametrize("mode", [AppMode.AGENT, AppMode.ADVANCED_CHAT, AppMode.CHAT]) -@pytest.mark.parametrize("kind", ["title", "moderation", "aliased_moderation"]) -def test_auxiliary_events_resolve_owner_without_bypassing_content_policy( - sqlite_session_factory: sessionmaker[Session], monkeypatch: pytest.MonkeyPatch, mode: AppMode, kind: str -) -> None: - case = setup(sqlite_session_factory, monkeypatch, "aux", mode, snapshot=False) - with case.collection.sessions() as session: - message = session.get(Message, case.collection.run.source_id) - assert message is not None - if kind == "title": - info: BaseTraceInfo = GenerateNameTraceInfo( - tenant_id=case.collection.run.tenant_id, - conversation_id=message.conversation_id, - metadata={"conversation_id": message.conversation_id}, - inputs="private conversation", - outputs="private generated title", - ) - else: - info = ModerationTraceInfo( - message_id=identifier("workflow-app-log") if kind == "aliased_moderation" else message.id, - message_data=message.to_dict(session=session), - metadata={"message_id": message.id}, - flagged=False, - action="direct_output", - preset_response="", - query="private prompt", - inputs="private prompt", - outputs="private output", - ) - legacy = dispatch(case, info, monkeypatch) - assert len(legacy) == (1 if mode is AppMode.CHAT else 0) - assert case.scheduled == [] - - -def test_plain_workflow_preserves_legacy_parent_context( - sqlite_session_factory: sessionmaker[Session], monkeypatch: pytest.MonkeyPatch -) -> None: - parent_case = setup(sqlite_session_factory, monkeypatch, "plain-parent", AppMode.WORKFLOW, snapshot=False) - parent, _ = workflow(parent_case, "plain-parent", has_agent=False, snapshot=False) - child_case = setup(sqlite_session_factory, monkeypatch, "plain-child", AppMode.WORKFLOW, snapshot=False) - child, _ = workflow(child_case, "plain-child", has_agent=False, snapshot=False) - info = workflow_info(child, parent) - original = dict(info.metadata) - assert child_case.service.capture_terminal(info) is False - assert info.metadata == original - with child_case.collection.sessions() as session: - assert session.scalar(select(func.count()).select_from(ApplicationTraceRun)) == 0 - - -def test_plain_chatflow_auxiliary_events_keep_the_legacy_route( - sqlite_session_factory: sessionmaker[Session], monkeypatch: pytest.MonkeyPatch -) -> None: - case = setup(sqlite_session_factory, monkeypatch, "plain-chatflow", AppMode.ADVANCED_CHAT, snapshot=False) - row, _ = workflow(case, "plain-chatflow", has_agent=False, snapshot=False) - with case.collection.sessions.begin() as session: - message = session.get(Message, case.collection.run.source_id) - assert message is not None - message.workflow_run_id = row.id - info = GenerateNameTraceInfo( - tenant_id=row.tenant_id, - conversation_id=message.conversation_id, - metadata={"conversation_id": message.conversation_id}, - inputs="legacy input", - outputs="legacy title", - ) - assert len(dispatch(case, info, monkeypatch)) == 1 - - -def test_legacy_child_of_managed_parent_uses_explicit_link_not_unavailable_context( - sqlite_session_factory: sessionmaker[Session], monkeypatch: pytest.MonkeyPatch -) -> None: - parent_case = setup(sqlite_session_factory, monkeypatch, "managed-parent", AppMode.WORKFLOW, snapshot=False) - parent, parent_trace = workflow(parent_case, "managed-parent", has_agent=True, snapshot=True) - assert parent_trace is not None - child_case = setup(sqlite_session_factory, monkeypatch, "legacy-child", AppMode.WORKFLOW, snapshot=False) - child, _ = workflow(child_case, "legacy-child", has_agent=False, snapshot=False) - info = workflow_info(child, parent) - assert child_case.service.capture_terminal(info) is False - assert info.resolved_parent_context == (None, None) - assert info.metadata["linked_parent_workflow_run_id"] == parent.id - assert info.metadata["linked_parent_node_execution_id"] == identifier("parent-tool-node") - assert info.metadata["linked_parent_application_trace_id"] == parent_trace.id - assert info.metadata["linked_parent_trace_id"] == parent_trace.trace_id - - -def test_mixed_workflow_exports_normal_llm_and_agent_calls_without_double_usage( - sqlite_session_factory: sessionmaker[Session], monkeypatch: pytest.MonkeyPatch -) -> None: - case = setup(sqlite_session_factory, monkeypatch, "mixed", AppMode.WORKFLOW, snapshot=False) - row, managed = workflow(case, "mixed", has_agent=True, snapshot=True) - assert managed is not None - agent_node_id = identifier("agent-execution") - owner = identifier("mixed-owner") - with case.collection.sessions.begin() as session: - app = session.get(App, row.app_id) - assert app is not None - app.created_by = owner - session.add(make_account(account_id=owner)) - for index, kind in enumerate(("llm", "agent")): - node_id = identifier("normal-execution") if kind == "llm" else agent_node_id - usage: dict[str, int] = ( - {"prompt_tokens": 10, "completion_tokens": 5, "total_tokens": 15} if kind == "llm" else {} - ) - session.add( - WorkflowNodeExecutionModel( - id=node_id, - node_execution_id=node_id, - tenant_id=row.tenant_id, - app_id=row.app_id, - workflow_id=row.workflow_id, - workflow_run_id=row.id, - triggered_from=WorkflowNodeExecutionTriggeredFrom.WORKFLOW_RUN, - index=index, - node_id=kind, - node_type=kind, - title=kind, - inputs="{}", - outputs="{}", - process_data=json.dumps( - { - "model_mode": "chat", - "model_name": "ordinary-model", - "model_provider": "openai", - "usage": usage, - } - ) - if kind == "llm" - else "{}", - status=WorkflowNodeExecutionStatus.SUCCEEDED, - elapsed_time=1, - execution_metadata=json.dumps( - {"total_tokens": 15 if kind == "llm" else 10, "total_price": "0.02", "currency": "USD"} - ), - created_at=row.created_at, - finished_at=row.finished_at, - created_by=owner, - created_by_role=CreatorUserRole.ACCOUNT, - ) - ) - invocation = case.collection.repository.invocation( - run_id=managed.id, agent_id=identifier("agent"), node_execution_id=agent_node_id, attempt=0, call_key="agent:0" - ) - spans = [ - AgentTraceSpan( - trace_id=managed.trace_id, - span_id="a" * 16, - parent_span_id=invocation.span_id, - name="agent run", - start_time_ns=1_000_000_000, - end_time_ns=3_000_000_000, - status="ok", - attributes={"gen_ai.operation.name": "invoke_agent"}, - ) - ] - spans.extend( - AgentTraceSpan( - trace_id=managed.trace_id, - span_id=key * 16, - parent_span_id="a" * 16, - name="model call", - start_time_ns=1_000_000_000, - end_time_ns=2_000_000_000, - status="ok", - attributes={ - "gen_ai.operation.name": "chat", - "gen_ai.request.model": "agent-model", - "gen_ai.usage.input_tokens": 3, - "gen_ai.usage.output_tokens": 2, - "dify.llm.total_price": "0.01", - "dify.llm.currency": "USD", - }, - ) - for key in ("b", "c") - ) - case.collection.repository.ingest( - AgentTraceBatch( - run_id=UUID(managed.id), invocation_id=UUID(invocation.id), batch_id=uuid4(), final=True, spans=spans - ), - max_spans=100, - max_bytes=100000, - ) - monkeypatch.setattr(type(db), "engine", property(lambda _db: sqlite_session_factory.kw["bind"])) - assert case.service.capture_terminal(workflow_info(row)) is True - stored = case.collection.repository.get(managed.id) - assert stored is not None - trace = case.service.assemble(stored) - payloads: list[ExportTraceServiceRequest] = [] - - def send(method: str, url: str, **kwargs: object) -> httpx.Response: - body = kwargs["content"] - assert isinstance(body, bytes) - payloads.append(ExportTraceServiceRequest.FromString(body)) - return httpx.Response(200, content=b"", request=httpx.Request(method, url)) - - monkeypatch.setattr("core.helper.ssrf_proxy.make_request", send) - export_trace("langfuse", {"host": "https://langfuse.test", "public_key": "pk", "secret_key": "sk"}, trace) - observations = payloads[0].resource_spans[0].scope_spans[0].spans - attributes = [{item.key: item.value.string_value for item in span.attributes} for span in observations] - generations = [item for item in attributes if item["langfuse.observation.type"] == "generation"] - assert len(generations) == 3 - assert sum(json.loads(item["langfuse.observation.usage_details"])["total"] for item in generations) == 25 - assert sum(json.loads(item["langfuse.observation.cost_details"])["total"] for item in generations) == pytest.approx( - 0.04 - ) diff --git a/api/tests/unit_tests/services/tracing/test_application_trace_logstore.py b/api/tests/unit_tests/services/tracing/test_application_trace_logstore.py deleted file mode 100644 index 22bc2132e8f..00000000000 --- a/api/tests/unit_tests/services/tracing/test_application_trace_logstore.py +++ /dev/null @@ -1,422 +0,0 @@ -"""Configured LogStore readers work without SQL mirrors and retain source ownership.""" - -import json -from collections.abc import Iterator -from dataclasses import dataclass -from datetime import timedelta -from typing import Literal - -import pytest -from aliyun.log.logclient import LogClient -from sqlalchemy import Connection, Engine, event, select -from sqlalchemy.orm import Session, sessionmaker - -from core.ops.entities.trace_entity import ModerationTraceInfo -from extensions.logstore.aliyun_logstore import AliyunLogStore -from extensions.logstore.repositories.logstore_api_workflow_run_repository import _dict_to_workflow_run -from libs.datetime_utils import naive_utc_now -from models.application_trace import ApplicationTraceRun -from models.model import App, AppMode -from models.workflow import WorkflowNodeExecutionModel, WorkflowRun -from repositories.application_trace_repository import TraceCollectionError -from services.tracing.provider_dispatch import TraceExportResult -from services.tracing.recovery import _workflow_task -from services.tracing.service import ApplicationTracingService -from tests.unit_tests.config_override import apply_config_overrides -from tests.unit_tests.model_factories import make_workflow -from tests.unit_tests.services.tracing.test_application_trace_collection import ( - Collection, - batch_for, - claim, - identifier, - ingest, - seed, - terminal, -) -from tests.unit_tests.services.tracing.test_application_trace_handoff import AGENT_GRAPH, workflow_info - - -@dataclass -class StoredSource: - case: Collection - kind: Literal["workflow", "node"] - row: dict[str, object] - requests: list[str] - - def start(self) -> ApplicationTraceRun: - run = self.case.repository.start( - tenant_id=self.case.run.tenant_id, - app_id=self.case.run.app_id, - source_kind=self.kind, - source_id=str(self.row["id"]), - session_id="logstore-session", - ) - assert run is not None - return run - - def service(self) -> ApplicationTracingService: - return ApplicationTracingService(self.case.repository) - - def status(self, node_execution_id: str | None = None) -> dict[str, object]: - return self.service().status( - tenant_id=self.case.run.tenant_id, - app_id=self.case.run.app_id, - source_kind=self.kind, - source_id=str(self.row["id"]), - node_execution_id=node_execution_id, - ) - - -@pytest.fixture(params=["workflow", "node"]) -def stored( - request: pytest.FixtureRequest, - sqlite_engine: Engine, - sqlite_session_factory: sessionmaker[Session], - monkeypatch: pytest.MonkeyPatch, -) -> Iterator[StoredSource]: - kind = request.param - assert kind in {"workflow", "node"} - case = seed(sqlite_session_factory, "logstore") - with case.sessions.begin() as session: - app = session.get(App, case.run.app_id) - assert app is not None - app.mode = AppMode.WORKFLOW - now = naive_utc_now() - row: dict[str, object] = { - "id": identifier("logstore-source"), - "tenant_id": case.run.tenant_id, - "app_id": case.run.app_id, - "workflow_id": identifier("definition"), - "created_by": identifier("owner"), - "created_at": (now - timedelta(seconds=30)).isoformat(), - "finished_at": (now - timedelta(seconds=10)).isoformat(), - "status": "succeeded", - "elapsed_time": 20, - "inputs": json.dumps({"query": "private input"}), - "outputs": json.dumps({"result": "private output"}), - "node_id": "agent-node", - "node_type": "agent", - "title": "Agent preview", - "graph": json.dumps(AGENT_GRAPH), - } - apply_config_overrides( - monkeypatch, - LOGSTORE_DUAL_WRITE_ENABLED=False, - API_WORKFLOW_RUN_REPOSITORY=( - "extensions.logstore.repositories.logstore_api_workflow_run_repository.LogstoreAPIWorkflowRunRepository" - ), - API_WORKFLOW_NODE_EXECUTION_REPOSITORY=( - "extensions.logstore.repositories.logstore_api_workflow_node_execution_repository." - "LogstoreAPIWorkflowNodeExecutionRepository" - ), - ) - monkeypatch.setenv("LOGSTORE_DUAL_READ_ENABLED", "false") - monkeypatch.setenv("LOGSTORE_PG_MODE_ENABLED", "false") - monkeypatch.setenv("ALIYUN_SLS_ACCESS_KEY_ID", "test-access-key") - monkeypatch.setenv("ALIYUN_SLS_ACCESS_KEY_SECRET", "test-secret") - monkeypatch.setenv("ALIYUN_SLS_ENDPOINT", "https://sls.example.test") - monkeypatch.setenv("ALIYUN_SLS_REGION", "test-region") - monkeypatch.setattr(AliyunLogStore, "_instance", None) - monkeypatch.setattr(AliyunLogStore, "_initialized", False) - monkeypatch.setattr(AliyunLogStore, "_check_endpoint_connectivity", staticmethod(lambda _endpoint, _timeout: None)) - - def reject_http_transport(*_args: object, **_kwargs: object) -> None: - pytest.fail("Configured LogStore test escaped its isolated transport boundary") - - monkeypatch.setattr(LogClient, "_getHttpResponse", reject_http_transport) - transactions: set[Connection] = set() - - def begun(connection: Connection) -> None: - transactions.add(connection) - - def ended(connection: Connection) -> None: - transactions.discard(connection) - - requests: list[str] = [] - - def get_logs(_client: AliyunLogStore, *, query: str, logstore: str, **_kwargs: object) -> list[dict[str, object]]: - # Exercise real configured repository construction, scoped query generation, - # and model decoding; replace only the remote transport boundary. - assert not transactions, "LogStore I/O must not hold a SQL transaction" - assert f'tenant_id:"{case.run.tenant_id}"' in query.replace(": ", ":") - assert logstore == ( - AliyunLogStore.workflow_execution_logstore - if kind == "workflow" - else AliyunLogStore.workflow_node_execution_logstore - ) - if kind == "workflow": - assert f'app_id:"{case.run.app_id}"' in query - requests.append(query) - return [dict(row)] if row else [] - - monkeypatch.setattr(AliyunLogStore, "get_logs", get_logs) - - def execute_sql( - client: AliyunLogStore, *, sql: str, query: str, logstore: str, **_kwargs: object - ) -> list[dict[str, object]]: - # The node reader now uses SQL in both SDK and PG modes. Preserve the - # real repository's query/decoding path while isolating this transport. - assert kind == "node" - assert f'id:"{identifier("logstore-source")}"' in query.replace(": ", ":") - assert f"tenant_id = '{case.run.tenant_id}'" in sql - assert f"id = '{identifier('logstore-source')}'" in sql - assert f'FROM "{logstore}"' in sql - assert "ORDER BY log_version DESC LIMIT 1" in sql - return get_logs(client, query=query, logstore=logstore) - - monkeypatch.setattr(AliyunLogStore, "execute_sql", execute_sql) - event.listen(sqlite_engine, "begin", begun) - event.listen(sqlite_engine, "commit", ended) - event.listen(sqlite_engine, "rollback", ended) - try: - yield StoredSource(case, kind, row, requests) - finally: - event.remove(sqlite_engine, "begin", begun) - event.remove(sqlite_engine, "commit", ended) - event.remove(sqlite_engine, "rollback", ended) - - -def test_status_and_retry_use_logstore_without_sql_source_rows( - stored: StoredSource, monkeypatch: pytest.MonkeyPatch -) -> None: - run = stored.start() - with stored.case.sessions() as session: - assert session.scalar(select(WorkflowRun)) is None - assert session.scalar(select(WorkflowNodeExecutionModel)) is None - with stored.case.sessions.begin() as session: - current = session.get(ApplicationTraceRun, run.id) - assert current is not None - current.state = "failed" - current.error_code = "provider_export_rejected" - current.outer_trace = {"persisted": True} - assert stored.status()["status"] == "failed" - scheduled: list[str] = [] - service = stored.service() - monkeypatch.setattr(service, "schedule", scheduled.append) - stored.requests.clear() - result = service.retry( - tenant_id=run.tenant_id, - app_id=run.app_id, - source_kind=stored.kind, - source_id=run.source_id, - acknowledge_possible_duplicates=False, - ) - assert result["status"] == "ready" - assert scheduled == [run.id] - assert len(stored.requests) == 1 - - -@pytest.mark.parametrize("field", ["id", "tenant_id", "app_id"]) -def test_backend_result_must_match_complete_source_owner(stored: StoredSource, field: str) -> None: - run = stored.start() - stored.row[field] = identifier("foreign") - with pytest.raises(TraceCollectionError, match="tracing_source_not_found") as caught: - stored.service().status( - tenant_id=run.tenant_id, app_id=run.app_id, source_kind=stored.kind, source_id=run.source_id - ) - assert caught.value.status == 404 - assert stored.case.repository.get(run.id) is not None - - -def test_foreign_app_owner_is_rejected_before_logstore_request(stored: StoredSource) -> None: - stored.start() - stored.requests.clear() - with stored.case.sessions.begin() as session: - app = session.get(App, stored.case.run.app_id) - assert app is not None - app.tenant_id = identifier("other-tenant") - with pytest.raises(TraceCollectionError, match="tracing_source_not_found"): - stored.status() - assert stored.requests == [] - - -def test_initialization_failure_uses_configured_source_reader(stored: StoredSource) -> None: - run = stored.case.repository.record_initialization_failure( - tenant_id=stored.case.run.tenant_id, - app_id=stored.case.run.app_id, - source_kind=stored.kind, - source_id=str(stored.row["id"]), - ) - assert run is not None - assert run.state == "failed" - assert run.error_code == "tracing_initialization_failed" - - -def test_existing_snapshot_reentry_does_not_consult_logstore(stored: StoredSource) -> None: - original = stored.start() - stored.requests.clear() - stored.row.clear() - current = stored.case.repository.start( - tenant_id=original.tenant_id, - app_id=original.app_id, - source_kind=stored.kind, - source_id=original.source_id, - session_id="must-not-replace", - ) - assert current is not None - assert current.id == original.id - assert current.session_id == original.session_id - assert stored.requests == [] - - -@pytest.mark.parametrize("stored", ["workflow"], indirect=True) -def test_workflow_node_deep_link_is_scoped_to_durable_invocation(stored: StoredSource) -> None: - run = stored.start() - case = Collection(stored.case.repository, stored.case.sessions, run) - node_id = identifier("logstore-agent-node") - invocation = case.repository.invocation( - run_id=run.id, agent_id=identifier("agent"), node_execution_id=node_id, attempt=0, call_key="node" - ) - ingest(case, batch_for(run, invocation)) - terminal(case) - lease = claim(case) - assert lease is not None - case.repository.complete( - run.id, - lease[1], - state="accepted", - result=TraceExportResult("provider-id", span_urls={invocation.span_id: "https://trace.test/span"}), - ) - assert stored.status(node_id)["span_url"] == "https://trace.test/span" - with pytest.raises(TraceCollectionError, match="tracing_node_not_found"): - stored.status(identifier("another-run-node")) - - -def test_preview_settlement_and_workflow_recovery_use_configured_source(stored: StoredSource) -> None: - run = stored.start() - if stored.kind == "workflow": - task = _workflow_task(stored.case.repository, run, naive_utc_now()) - assert task is not None - assert task.workflow_run_id == run.source_id - assert task.kwargs["trace_session_id"] == run.session_id - stored.row["status"] = "running" - assert _workflow_task(stored.case.repository, run, naive_utc_now()) is None - else: - stored.service().settle_preview(run) - current = stored.case.repository.get(run.id) - assert current is not None - assert current.outer_trace is not None - root = current.outer_trace["spans"][0] - assert root["status"] == "ok" - assert root["inputs"] is None - assert root["outputs"] is None - assert root["metadata"]["node_execution_id"] == run.source_id - - -@pytest.mark.parametrize("stored", ["workflow"], indirect=True) -def test_missing_snapshot_classifies_logstore_agent_workflow_and_persists_failure(stored: StoredSource) -> None: - info = workflow_info(_dict_to_workflow_run(stored.row)) - assert stored.service().capture_terminal(info) is True - result = stored.status() - assert result["status"] == "failed" - assert result["error_code"] == "tracing_initialization_failed" - - -@pytest.mark.parametrize("stored", ["workflow"], indirect=True) -@pytest.mark.parametrize("has_agent", [True, False]) -def test_unindexed_workflow_uses_owned_definition_for_safe_legacy_decision( - stored: StoredSource, has_agent: bool -) -> None: - info = workflow_info(_dict_to_workflow_run(stored.row)) - with stored.case.sessions.begin() as session: - app = session.get(App, stored.case.run.app_id) - assert app is not None - definition = make_workflow( - workflow_id=str(stored.row["workflow_id"]), - tenant_id=app.tenant_id, - app_id=app.id, - graph=AGENT_GRAPH if has_agent else {"nodes": [{"data": {"type": "llm"}}]}, - ) - app.workflow_id = definition.id - session.add(definition) - stored.row.clear() - assert stored.service().capture_terminal(info) is has_agent - - -@pytest.mark.parametrize("stored", ["workflow"], indirect=True) -@pytest.mark.parametrize("snapshot", [True, False]) -def test_auxiliary_content_is_suppressed_for_logstore_workflow(stored: StoredSource, snapshot: bool) -> None: - if snapshot: - stored.start() - info = ModerationTraceInfo( - message_id=identifier("workflow-log"), - metadata={ - "workflow_run_id": stored.row["id"], - **({} if snapshot else {"app_id": stored.case.run.app_id}), - }, - flagged=False, - action="direct_output", - preset_response="", - query="private query", - inputs="private input", - outputs="private output", - ) - assert stored.service().capture_terminal(info) is True - if snapshot: - # Durable ownership is sufficient to suppress auxiliary payload even when - # the remote run has not been indexed yet. - stored.row.clear() - assert stored.service().capture_terminal(info) is True - else: - stored.row["graph"] = json.dumps({"nodes": [{"data": {"type": "llm"}}]}) - assert stored.service().capture_terminal(info) is False - - -@pytest.mark.parametrize("stored", ["workflow"], indirect=True) -@pytest.mark.parametrize("mismatch", [None, "tenant_id", "app_id"]) -def test_linked_parent_requires_logstore_run_and_durable_owner_to_agree( - stored: StoredSource, mismatch: str | None -) -> None: - parent = stored.start() - info = workflow_info(_dict_to_workflow_run(stored.row), _dict_to_workflow_run(stored.row)) - original = dict(info.metadata) - if mismatch: - stored.row[mismatch] = identifier("foreign") - stored.service()._link_managed_parent(info, parent.tenant_id) - if mismatch: - assert info.metadata == original - else: - assert info.resolved_parent_context == (None, None) - assert info.metadata["linked_parent_trace_id"] == parent.trace_id - assert info.metadata["linked_parent_workflow_run_id"] == parent.source_id - - -@pytest.mark.parametrize("stored", ["workflow"], indirect=True) -def test_auxiliary_durable_owner_rejects_app_reassigned_to_other_tenant(stored: StoredSource) -> None: - run = stored.start() - with stored.case.sessions.begin() as session: - app = session.get(App, run.app_id) - assert app is not None - app.tenant_id = identifier("foreign") - info = ModerationTraceInfo( - metadata={"workflow_run_id": run.source_id}, query="", flagged=False, action="direct_output", preset_response="" - ) - assert stored.service().capture_terminal(info) is False - - -def test_sql_workflow_source_still_uses_scoped_configured_reader( - sqlite_session_factory: sessionmaker[Session], -) -> None: - case = seed(sqlite_session_factory, "sql-reader") - row = _dict_to_workflow_run( - { - "id": identifier("sql-run"), - "tenant_id": case.run.tenant_id, - "app_id": case.run.app_id, - "workflow_id": identifier("definition"), - "created_by": identifier("owner"), - } - ) - with case.sessions.begin() as session: - session.add(row) - source = case.repository._source( - tenant_id=case.run.tenant_id, app_id=case.run.app_id, kind="workflow", source_id=row.id - ) - assert isinstance(source, WorkflowRun) - assert source.id == row.id - assert ( - case.repository._source( - tenant_id=identifier("wrong-tenant"), app_id=case.run.app_id, kind="workflow", source_id=row.id - ) - is None - ) diff --git a/api/tests/unit_tests/services/tracing/test_application_trace_node_alias.py b/api/tests/unit_tests/services/tracing/test_application_trace_node_alias.py deleted file mode 100644 index 59f9116f177..00000000000 --- a/api/tests/unit_tests/services/tracing/test_application_trace_node_alias.py +++ /dev/null @@ -1,386 +0,0 @@ -"""Node logs resolve the caller's trace using SQL or configured LogStore ownership.""" - -import json -from collections.abc import Iterator -from dataclasses import dataclass - -import pytest -from aliyun.log.logclient import LogClient -from sqlalchemy import Connection, Engine, event, select -from sqlalchemy.orm import Session, sessionmaker - -from extensions.logstore.aliyun_logstore import AliyunLogStore -from extensions.logstore.repositories.logstore_api_workflow_node_execution_repository import ( - _dict_to_workflow_node_execution_model, -) -from extensions.logstore.repositories.logstore_api_workflow_run_repository import _dict_to_workflow_run -from libs.datetime_utils import naive_utc_now -from models.application_trace import ApplicationTraceRun -from models.model import App, AppMode -from models.workflow import WorkflowNodeExecutionModel, WorkflowRun -from repositories.application_trace_repository import TraceCollectionError -from services.tracing.provider_dispatch import TraceExportResult -from services.tracing.service import ApplicationTracingService -from tests.unit_tests.config_override import apply_config_overrides -from tests.unit_tests.services.tracing.test_application_trace_collection import ( - Collection, - batch_for, - claim, - identifier, - ingest, - seed, - terminal, -) - - -@dataclass -class NodeAlias: - case: Collection - node: dict[str, object] - workflow: dict[str, object] - requests: list[str] - logstore: bool - - def service(self) -> ApplicationTracingService: - return ApplicationTracingService(self.case.repository) - - def status(self, node_execution_id: str | None = None) -> dict[str, object]: - return self.service().status( - tenant_id=self.case.run.tenant_id, - app_id=self.case.run.app_id, - source_kind="node", - source_id=str(self.node["id"]), - node_execution_id=node_execution_id, - ) - - def accept(self, *, spans: bool = True) -> tuple[str, str]: - span_urls: dict[str, str] = {} - for attempt in (0, 1): - invocation = self.case.repository.invocation( - run_id=self.case.run.id, - agent_id=identifier("referenced-agent"), - node_execution_id=str(self.node["id"]), - attempt=attempt, - call_key=f"attempt-{attempt}", - ) - ingest(self.case, batch_for(self.case.run, invocation, span_id=str(attempt + 1) * 16)) - span_urls[invocation.span_id] = f"https://provider.test/span/{attempt}" - terminal(self.case) - lease = claim(self.case) - assert lease is not None - self.case.repository.complete( - self.case.run.id, - lease[1], - state="accepted", - result=TraceExportResult( - "provider-trace", "https://provider.test/trace", span_urls=span_urls if spans else {} - ), - ) - return "https://provider.test/trace", "https://provider.test/span/1" - - -@pytest.fixture(params=["sql", "logstore"]) -def alias( - request: pytest.FixtureRequest, - sqlite_engine: Engine, - sqlite_session_factory: sessionmaker[Session], - monkeypatch: pytest.MonkeyPatch, -) -> Iterator[NodeAlias]: - case = seed(sqlite_session_factory, "node-alias") - with case.sessions.begin() as session: - app = session.get(App, case.run.app_id) - assert app is not None - app.mode = AppMode.WORKFLOW - common: dict[str, object] = { - "tenant_id": case.run.tenant_id, - "app_id": case.run.app_id, - "workflow_id": identifier("definition"), - "created_by": identifier("actor"), - "created_at": naive_utc_now().isoformat(), - "status": "succeeded", - } - workflow = {**common, "id": identifier("workflow-source"), "graph": json.dumps({"nodes": []})} - node = { - **common, - "id": identifier("node-source"), - "workflow_run_id": workflow["id"], - "node_id": "reference-agent", - "node_type": "agent", - "title": "Reference Agent", - } - requests: list[str] = [] - transactions: set[Connection] = set() - - def begun(connection: Connection) -> None: - transactions.add(connection) - - def ended(connection: Connection) -> None: - transactions.discard(connection) - - if request.param == "sql": - with case.sessions.begin() as session: - session.add(_dict_to_workflow_run(workflow)) - session.add(_dict_to_workflow_node_execution_model(node)) - else: - apply_config_overrides( - monkeypatch, - LOGSTORE_DUAL_WRITE_ENABLED=False, - API_WORKFLOW_RUN_REPOSITORY=( - "extensions.logstore.repositories.logstore_api_workflow_run_repository.LogstoreAPIWorkflowRunRepository" - ), - API_WORKFLOW_NODE_EXECUTION_REPOSITORY=( - "extensions.logstore.repositories.logstore_api_workflow_node_execution_repository." - "LogstoreAPIWorkflowNodeExecutionRepository" - ), - ) - for key, value in { - "LOGSTORE_DUAL_READ_ENABLED": "false", - "LOGSTORE_PG_MODE_ENABLED": "false", - "ALIYUN_SLS_ACCESS_KEY_ID": "test-access-key", - "ALIYUN_SLS_ACCESS_KEY_SECRET": "test-secret", - "ALIYUN_SLS_ENDPOINT": "https://sls.example.test", - "ALIYUN_SLS_REGION": "test-region", - }.items(): - monkeypatch.setenv(key, value) - monkeypatch.setattr(AliyunLogStore, "_instance", None) - monkeypatch.setattr(AliyunLogStore, "_initialized", False) - monkeypatch.setattr( - AliyunLogStore, "_check_endpoint_connectivity", staticmethod(lambda _endpoint, _timeout: None) - ) - - def reject_http_transport(*_args: object, **_kwargs: object) -> None: - pytest.fail("Configured LogStore test escaped its isolated transport boundary") - - monkeypatch.setattr(LogClient, "_getHttpResponse", reject_http_transport) - - def get_logs( - _client: AliyunLogStore, *, query: str, logstore: str, **_kwargs: object - ) -> list[dict[str, object]]: - assert not transactions, "Configured source reads must not hold a SQL transaction" - assert f'tenant_id:"{case.run.tenant_id}"' in query.replace(": ", ":") - assert logstore in { - AliyunLogStore.workflow_execution_logstore, - AliyunLogStore.workflow_node_execution_logstore, - } - is_workflow = logstore == AliyunLogStore.workflow_execution_logstore - if is_workflow: - assert f'app_id:"{case.run.app_id}"' in query - requests.append("workflow" if is_workflow else "node") - row = workflow if is_workflow else node - return [dict(row)] if row else [] - - monkeypatch.setattr(AliyunLogStore, "get_logs", get_logs) - - def execute_sql( - client: AliyunLogStore, *, sql: str, query: str, logstore: str, **_kwargs: object - ) -> list[dict[str, object]]: - # Node lookups use the SQL transport in SDK mode too; the real - # service still validates the returned node-to-workflow owner chain. - assert logstore == AliyunLogStore.workflow_node_execution_logstore - assert f'id:"{identifier("node-source")}"' in query.replace(": ", ":") - assert f"tenant_id = '{case.run.tenant_id}'" in sql - assert f"id = '{identifier('node-source')}'" in sql - assert f'FROM "{logstore}"' in sql - assert "ORDER BY log_version DESC LIMIT 1" in sql - return get_logs(client, query=query, logstore=logstore) - - monkeypatch.setattr(AliyunLogStore, "execute_sql", execute_sql) - with case.sessions() as session: - assert session.scalar(select(WorkflowRun)) is None - assert session.scalar(select(WorkflowNodeExecutionModel)) is None - event.listen(sqlite_engine, "begin", begun) - event.listen(sqlite_engine, "commit", ended) - event.listen(sqlite_engine, "rollback", ended) - try: - run = case.repository.start( - tenant_id=case.run.tenant_id, - app_id=case.run.app_id, - source_kind="workflow", - source_id=str(workflow["id"]), - session_id="caller-session", - ) - assert run is not None - requests.clear() - yield NodeAlias( - Collection(case.repository, case.sessions, run), node, workflow, requests, request.param == "logstore" - ) - finally: - event.remove(sqlite_engine, "begin", begun) - event.remove(sqlite_engine, "commit", ended) - event.remove(sqlite_engine, "rollback", ended) - - -@pytest.mark.parametrize("spans", [True, False]) -def test_alias_returns_parent_trace_and_latest_owned_invocation(alias: NodeAlias, spans: bool) -> None: - trace_url, span_url = alias.accept(spans=spans) - status = alias.status() - assert status["status"] == "accepted" - assert status["source_kind"] == "node" - assert status["source_id"] == alias.node["id"] - assert status["trace_id"] == "provider-trace" - assert status["trace_url"] == trace_url - assert status["span_url"] == (span_url if spans else None) - assert status["config_revision"] == alias.case.run.config_revision - if alias.logstore: - assert alias.requests == ["node", "workflow"] - - -@pytest.mark.parametrize("capture_agent", [True, False]) -@pytest.mark.parametrize("accepted", [True, False]) -def test_alias_without_internal_invocation_returns_parent_status( - alias: NodeAlias, capture_agent: bool, accepted: bool -) -> None: - with alias.case.sessions.begin() as session: - run = session.get(ApplicationTraceRun, alias.case.run.id) - assert run is not None - run.capture_agent = capture_agent - if accepted: - terminal(alias.case) - lease = claim(alias.case) - assert lease is not None - alias.case.repository.complete( - alias.case.run.id, - lease[1], - state="accepted", - result=TraceExportResult("outer-only", "https://provider.test/outer"), - ) - status = alias.status() - assert status["status"] == ("accepted" if accepted else "collecting") - assert status["span_url"] is None - assert status["capture_agent"] is capture_agent - assert status["trace_url"] == ("https://provider.test/outer" if accepted else None) - - -@pytest.mark.parametrize("state", ["collecting", "not_collected"]) -def test_standalone_node_snapshot_wins_over_parent(alias: NodeAlias, state: str) -> None: - alias.accept() - node_run = alias.case.repository.start( - tenant_id=alias.case.run.tenant_id, - app_id=alias.case.run.app_id, - source_kind="node", - source_id=str(alias.node["id"]), - session_id="preview-session", - ) - assert node_run is not None - with alias.case.sessions.begin() as session: - run = session.get(ApplicationTraceRun, node_run.id) - assert run is not None - run.state = state - alias.requests.clear() - status = alias.status() - assert status["status"] == state - assert status.get("trace_url") is None - if alias.logstore: - assert alias.requests == ["node"] - - -def test_explicit_other_node_filter_cannot_select_sibling_span( - alias: NodeAlias, monkeypatch: pytest.MonkeyPatch -) -> None: - sibling = alias.case.repository.invocation( - run_id=alias.case.run.id, - agent_id=identifier("referenced-agent"), - node_execution_id=identifier("sibling-node"), - attempt=0, - call_key="sibling", - ) - ingest(alias.case, batch_for(alias.case.run, sibling, span_id="f" * 16)) - alias.accept() - with pytest.raises(TraceCollectionError, match="tracing_node_not_found") as error: - alias.status(identifier("sibling-node")) - assert error.value.status == 404 - service = alias.service() - scheduled: list[str] = [] - monkeypatch.setattr(service, "schedule", scheduled.append) - with pytest.raises(TraceCollectionError, match="tracing_node_not_found"): - service.retry( - tenant_id=alias.case.run.tenant_id, - app_id=alias.case.run.app_id, - source_kind="node", - source_id=str(alias.node["id"]), - node_execution_id=identifier("sibling-node"), - acknowledge_possible_duplicates=True, - ) - assert not scheduled - - -@pytest.mark.parametrize("state", ["failed", "uncertain", "accepted"]) -def test_alias_retry_uses_parent_guards_and_resolves_sources_once( - alias: NodeAlias, state: str, monkeypatch: pytest.MonkeyPatch -) -> None: - alias.accept() - with alias.case.sessions.begin() as session: - run = session.get(ApplicationTraceRun, alias.case.run.id) - assert run is not None - run.state = state - service = alias.service() - scheduled: list[str] = [] - monkeypatch.setattr(service, "schedule", scheduled.append) - - def retry(acknowledge: bool) -> dict[str, object]: - alias.requests.clear() - return service.retry( - tenant_id=alias.case.run.tenant_id, - app_id=alias.case.run.app_id, - source_kind="node", - source_id=str(alias.node["id"]), - acknowledge_possible_duplicates=acknowledge, - ) - - if state == "accepted": - with pytest.raises(TraceCollectionError, match="tracing_already_accepted") as error: - retry(True) - assert error.value.status == 409 - assert not scheduled - else: - if state == "uncertain": - with pytest.raises(TraceCollectionError, match="tracing_duplicate_acknowledgement_required"): - retry(False) - assert not scheduled - assert retry(state == "uncertain")["status"] == "ready" - assert scheduled == [alias.case.run.id] - if alias.logstore: - assert alias.requests == ["node", "workflow"] - - -@pytest.mark.parametrize("source", ["node", "workflow"]) -@pytest.mark.parametrize("field", ["tenant_id", "app_id", "workflow_id"]) -def test_alias_rejects_broken_owner_chain(alias: NodeAlias, source: str, field: str) -> None: - alias.accept() - row = alias.node if source == "node" else alias.workflow - if alias.logstore: - row[field] = identifier("foreign") - else: - with alias.case.sessions.begin() as session: - model = session.get(WorkflowNodeExecutionModel if source == "node" else WorkflowRun, str(row["id"])) - assert model is not None - # Keep field assignments explicit so ownership failures are visible. - if field == "tenant_id": - model.tenant_id = identifier("foreign") - elif field == "app_id": - model.app_id = identifier("foreign") - else: - model.workflow_id = identifier("foreign") - with pytest.raises(TraceCollectionError, match="tracing_source_not_found") as error: - alias.status() - assert error.value.status == 404 - - -def test_alias_missing_parent_business_source_is_not_a_durable_owner_bypass(alias: NodeAlias) -> None: - if alias.logstore: - alias.workflow.clear() - else: - with alias.case.sessions.begin() as session: - parent = session.get(WorkflowRun, str(alias.workflow["id"])) - assert parent is not None - session.delete(parent) - with pytest.raises(TraceCollectionError, match="tracing_source_not_found"): - alias.status() - - -def test_owned_alias_without_parent_trace_is_not_collected(alias: NodeAlias) -> None: - with alias.case.sessions.begin() as session: - run = session.get(ApplicationTraceRun, alias.case.run.id) - assert run is not None - session.delete(run) - assert alias.status() == {"status": "not_collected", "source_kind": "node", "source_id": alias.node["id"]} diff --git a/api/tests/unit_tests/services/tracing/test_application_trace_repository_boundaries.py b/api/tests/unit_tests/services/tracing/test_application_trace_repository_boundaries.py deleted file mode 100644 index 7122829140b..00000000000 --- a/api/tests/unit_tests/services/tracing/test_application_trace_repository_boundaries.py +++ /dev/null @@ -1,209 +0,0 @@ -"""Persistence boundaries that must fail closed or remain immutable across retries.""" - -from datetime import timedelta -from uuid import UUID - -import pytest -from sqlalchemy import delete, func, select -from sqlalchemy.orm import Session, sessionmaker - -from libs.datetime_utils import naive_utc_now -from models.application_trace import ApplicationTraceBatch, ApplicationTraceInvocation, ApplicationTraceRun -from models.model import App -from repositories.application_trace_repository import TraceCollectionError -from tests.unit_tests.services.tracing.test_application_trace_collection import ( - batch_for, - claim, - identifier, - ingest, - seed, - terminal, -) - - -@pytest.mark.parametrize("policy", ["invalid-json", "null", "[]"]) -def test_malformed_policy_freezes_disabled_decision(sqlite_session_factory: sessionmaker[Session], policy: str) -> None: - case = seed(sqlite_session_factory, "malformed") - with case.sessions.begin() as session: - session.execute(delete(ApplicationTraceRun).where(ApplicationTraceRun.id == case.run.id)) - app = session.get(App, case.run.app_id) - assert app is not None - app.tracing = policy - run = case.repository.start( - tenant_id=case.run.tenant_id, - app_id=case.run.app_id, - source_kind="message", - source_id=case.run.source_id, - session_id="", - ) - assert run is not None - assert run.state == "not_collected" - assert run.provider_config == {} - assert run.session_id == case.run.source_id - assert not run.capture_content - assert not run.capture_agent - - -def test_reserved_preview_requires_persisted_owner_workflow(sqlite_session_factory: sessionmaker[Session]) -> None: - case = seed(sqlite_session_factory, "missing-preview-owner") - assert ( - case.repository.start( - tenant_id=case.run.tenant_id, - app_id=case.run.app_id, - source_kind="node", - source_id=identifier("preview"), - session_id="preview", - workflow_id=identifier("missing-workflow"), - ) - is None - ) - - -def test_initialization_failure_cannot_replace_snapshot_or_create_foreign_row( - sqlite_session_factory: sessionmaker[Session], -) -> None: - case = seed(sqlite_session_factory, "failure-snapshot") - same = case.repository.record_initialization_failure( - tenant_id=case.run.tenant_id, - app_id=case.run.app_id, - source_kind="message", - source_id=case.run.source_id, - ) - assert same is not None - assert same.id == case.run.id - assert same.state == "collecting" - assert same.provider_config == case.run.provider_config - assert ( - case.repository.record_initialization_failure( - tenant_id=identifier("foreign-tenant"), - app_id=case.run.app_id, - source_kind="message", - source_id=case.run.source_id, - ) - is None - ) - - -@pytest.mark.parametrize("target", ["missing_run", "revoked_owner", "closed_invocation"]) -def test_ingestion_rejects_missing_or_closed_owner_without_writing( - sqlite_session_factory: sessionmaker[Session], target: str -) -> None: - case = seed(sqlite_session_factory, "ingestion-owner") - invocation = case.invocation() - batch = batch_for(case.run, invocation) - if target == "missing_run": - batch = batch.model_copy(update={"run_id": UUID(identifier("missing-run"))}) - error = "tracing_run_not_found" - elif target == "revoked_owner": - with case.sessions.begin() as session: - app = session.get(App, case.run.app_id) - assert app is not None - app.tenant_id = identifier("another-tenant") - error = "tracing_owner_not_found" - else: - with case.sessions.begin() as session: - stored = session.get(ApplicationTraceInvocation, invocation.id) - assert stored is not None - stored.finalized = True - error = "tracing_collection_closed" - with pytest.raises(TraceCollectionError, match=error): - ingest(case, batch) - with case.sessions() as session: - assert session.scalar(select(func.count()).select_from(ApplicationTraceBatch)) == 0 - stored_run = case.repository.get(case.run.id) - assert stored_run is not None - assert stored_run.span_count == 0 - - -def test_first_terminal_is_immutable_and_revoked_owner_is_never_exported( - sqlite_session_factory: sessionmaker[Session], -) -> None: - case = seed(sqlite_session_factory, "terminal-owner") - terminal(case) - before = case.repository.get(case.run.id) - assert before is not None - case.repository.terminal(case.run.id, {"replacement": True}, "f" * 16) - after = case.repository.get(case.run.id) - assert after is not None - assert after.outer_trace == before.outer_trace - assert after.root_span_id == before.root_span_id - with case.sessions.begin() as session: - app = session.get(App, case.run.app_id) - assert app is not None - app.tenant_id = identifier("moved-owner") - assert claim(case) is None - failed = case.repository.get(case.run.id) - assert failed is not None - assert failed.error_code == "tracing_owner_not_found" - assert failed.state == "failed" - case.repository.terminal(case.run.id, {"replacement": True}, "f" * 16) - frozen = case.repository.get(case.run.id) - assert frozen is not None - assert frozen.outer_trace == before.outer_trace - - -@pytest.mark.parametrize( - ("state", "error"), - [ - ("not_collected", "tracing_run_not_found"), - ("accepted", "tracing_already_accepted"), - ("collecting", "tracing_retry_not_allowed"), - ], -) -def test_retry_rejects_ineligible_states(sqlite_session_factory: sessionmaker[Session], state: str, error: str) -> None: - case = seed(sqlite_session_factory, "retry-state") - terminal(case) - with case.sessions.begin() as session: - row = session.get(ApplicationTraceRun, case.run.id) - assert row is not None - row.state = state - with pytest.raises(TraceCollectionError, match=error): - case.repository.retry(run_id=case.run.id, acknowledge_possible_duplicates=True) - row = case.repository.get(case.run.id) - assert row is not None - assert row.state == state - - -def test_retry_unknown_run_and_invocation_reject_before_creation(sqlite_session_factory: sessionmaker[Session]) -> None: - case = seed(sqlite_session_factory, "retry-unknown") - with pytest.raises(TraceCollectionError, match="tracing_run_not_found"): - case.repository.retry(run_id=identifier("missing"), acknowledge_possible_duplicates=True) - with pytest.raises(TraceCollectionError, match="tracing_run_not_collecting"): - case.repository.invocation( - run_id=identifier("missing"), - agent_id=identifier("agent"), - node_execution_id=None, - attempt=0, - call_key="x", - ) - assert case.repository.claim(identifier("missing"), collection_timeout=1, lease_seconds=1) is None - - -def test_collection_without_terminal_is_not_exportable(sqlite_session_factory: sessionmaker[Session]) -> None: - case = seed(sqlite_session_factory, "no-terminal") - ingest(case, batch_for(case.run, case.invocation())) - assert claim(case) is None - row = case.repository.get(case.run.id) - assert row is not None - assert row.attempts == 0 - assert row.state == "collecting" - - -def test_retention_removes_all_payload_rows_and_preserves_newer_runs( - sqlite_session_factory: sessionmaker[Session], -) -> None: - expired = seed(sqlite_session_factory, "expired") - fresh = seed(sqlite_session_factory, "fresh") - for case in (expired, fresh): - ingest(case, batch_for(case.run, case.invocation())) - with expired.sessions.begin() as session: - row = session.get(ApplicationTraceRun, expired.run.id) - assert row is not None - row.created_at = naive_utc_now() - timedelta(days=8) - assert expired.repository.purge(retention_days=7) == 1 - assert expired.repository.get(expired.run.id) is None - assert expired.repository.collected(expired.run.id) == ([], []) - assert fresh.repository.get(fresh.run.id) is not None - invocations, batches = fresh.repository.collected(fresh.run.id) - assert len(invocations) == len(batches) == 1 - assert expired.repository.purge(retention_days=7) == 0 diff --git a/api/tests/unit_tests/services/tracing/test_application_trace_service_boundaries.py b/api/tests/unit_tests/services/tracing/test_application_trace_service_boundaries.py deleted file mode 100644 index 52baaeaf778..00000000000 --- a/api/tests/unit_tests/services/tracing/test_application_trace_service_boundaries.py +++ /dev/null @@ -1,425 +0,0 @@ -"""Exercise public orchestration, privacy, and delivery decisions against durable rows.""" - -from datetime import datetime -from uuid import UUID - -import pytest -from sqlalchemy import select -from sqlalchemy.orm import Session, sessionmaker - -from core.ops.exceptions import RetryableTraceDispatchError -from graphon.enums import WorkflowNodeExecutionStatus -from models.application_trace import ApplicationTraceRun -from models.enums import CreatorUserRole -from models.model import Message, TraceAppConfig -from models.workflow import WorkflowNodeExecutionModel, WorkflowNodeExecutionTriggeredFrom -from repositories.application_trace_repository import TraceCollectionError -from services.tracing import service as service_module -from services.tracing.provider_dispatch import TraceExportResult -from services.tracing.service import ApplicationTracingService -from tests.unit_tests.config_override import apply_config_overrides -from tests.unit_tests.model_factories import make_message, make_workflow -from tests.unit_tests.services.tracing.test_application_trace_collection import ( - Collection, - batch_for, - claim, - identifier, - ingest, - seed, - terminal, -) - - -def source(case: Collection) -> dict[str, str]: - return { - "tenant_id": case.run.tenant_id, - "app_id": case.run.app_id, - "source_kind": "message", - "source_id": case.run.source_id, - } - - -def test_prepare_invocation_returns_only_opaque_context_and_respects_capture_gate( - sqlite_session_factory: sessionmaker[Session], monkeypatch: pytest.MonkeyPatch -) -> None: - apply_config_overrides(monkeypatch, OPS_AGENT_TRACE_ENABLED=True) - case = seed(sqlite_session_factory, "prepare", capture_content=True) - service = ApplicationTracingService(case.repository) - context = service.prepare_invocation( - tenant_id=case.run.tenant_id, - app_id=case.run.app_id, - agent_id=identifier("agent"), - source_kind="message", - source_id=case.run.source_id, - session_id="ignored-new-session", - node_execution_id=None, - attempt=0, - call_key="call", - ) - assert context is not None - assert context.run_id == UUID(case.run.id) - assert context.capture_content is True - assert context.session_id == "session-original" - assert "encrypted-original" not in context.model_dump_json() - apply_config_overrides(monkeypatch, OPS_AGENT_TRACE_ENABLED=False) - assert ( - service.start( - tenant_id=case.run.tenant_id, - app_id=case.run.app_id, - source_kind="message", - source_id=case.run.source_id, - session_id="new", - ) - is None - ) - - -@pytest.mark.parametrize("missing", ["tenant", "app", "message"]) -def test_status_denies_foreign_or_absent_source_without_destination_details( - sqlite_session_factory: sessionmaker[Session], missing: str -) -> None: - case = seed(sqlite_session_factory, "scope") - identity = source(case) - identity[{"tenant": "tenant_id", "app": "app_id", "message": "source_id"}[missing]] = identifier("foreign") - with pytest.raises(TraceCollectionError) as caught: - ApplicationTracingService(case.repository).status(**identity) - assert caught.value.code == "tracing_source_not_found" - assert caught.value.status == 404 - assert "encrypted-original" not in str(caught.value) - - -def test_status_uses_frozen_destination_and_selects_latest_node_attempt( - sqlite_session_factory: sessionmaker[Session], - monkeypatch: pytest.MonkeyPatch, -) -> None: - case = seed(sqlite_session_factory, "status") - with case.sessions.begin() as session: - config = session.scalar(select(TraceAppConfig).where(TraceAppConfig.app_id == case.run.app_id)) - assert config is not None - config.tracing_config = {"project_url": "https://new.example/project/new", "secret_key": "encrypted-new"} - - def forbidden(*_args: object, **_kwargs: object) -> None: - pytest.fail("status must not decrypt credentials or resolve current provider metadata") - - monkeypatch.setattr(service_module.OpsTraceManager, "decrypt_tracing_config", forbidden) - monkeypatch.setattr(service_module.OpsTraceManager, "get_trace_config_project_url", forbidden) - node_id = identifier("selected-node") - invocation = case.repository.invocation( - run_id=case.run.id, agent_id=identifier("agent"), node_execution_id=node_id, attempt=1, call_key="node:1" - ) - ingest(case, batch_for(case.run, invocation)) - terminal(case) - lease = claim(case) - assert lease is not None - case.repository.complete( - case.run.id, - lease[1], - state="accepted", - result=TraceExportResult( - "provider-trace", - trace_url="https://frozen.example/trace", - span_urls={invocation.span_id: "https://frozen.example/span"}, - ), - ) - service = ApplicationTracingService(case.repository) - result = service.status(**source(case), node_execution_id=node_id) - assert result["status"] == "accepted" - assert result["trace_id"] == "provider-trace" - assert result["span_url"] == "https://frozen.example/span" - assert result["project_url"] == "https://langfuse.test/project/original" - assert result["config_revision"] == case.run.config_revision - assert result["queryable"] is None - assert result["capture_content"] is False - assert "provider_config" not in result - assert "encrypted-original" not in str(result) - with pytest.raises(TraceCollectionError, match="tracing_node_not_found"): - service.status(**source(case), node_execution_id=identifier("another-node")) - - -def test_new_run_freezes_updated_project_without_changing_previous_run( - sqlite_session_factory: sessionmaker[Session], -) -> None: - case = seed(sqlite_session_factory, "updated-project") - new_id = identifier("new-message") - with case.sessions.begin() as session: - config = session.scalar(select(TraceAppConfig).where(TraceAppConfig.app_id == case.run.app_id)) - assert config is not None - config.tracing_config = {"project_url": "https://new.example/project/new", "secret_key": "encrypted-new"} - message = session.get(Message, case.run.source_id) - assert message is not None - session.add( - make_message( - message_id=new_id, - app_id=case.run.app_id, - conversation_id=message.conversation_id, - inputs={}, - query="new run", - message={}, - answer="done", - message_unit_price=message.message_unit_price, - answer_unit_price=message.answer_unit_price, - currency=message.currency, - from_source=message.from_source, - ) - ) - new_run = case.repository.start( - tenant_id=case.run.tenant_id, - app_id=case.run.app_id, - source_kind="message", - source_id=new_id, - session_id="new-session", - ) - assert new_run is not None - service = ApplicationTracingService(case.repository) - assert service.status(**source(case))["project_url"] == "https://langfuse.test/project/original" - assert service.status(**{**source(case), "source_id": new_id})["project_url"] == "https://new.example/project/new" - - -@pytest.mark.parametrize( - "project_url", - [ - None, - "", - 1, - "javascript:alert(1)", - "//example.com/project", - "https://user:secret@example.com/project", - "https://@example.com", - "https://example.com\\@other.test", - "https://example.com:invalid/project", - "https://[invalid/project", - "https://example.com/\nproject", - ], -) -def test_status_hides_unavailable_or_unsafe_frozen_project_urls( - sqlite_session_factory: sessionmaker[Session], project_url: object -) -> None: - case = seed(sqlite_session_factory, "invalid-project") - with case.sessions.begin() as session: - run = session.get(ApplicationTraceRun, case.run.id) - assert run is not None - run.provider_config = { - "secret_key": "encrypted-original", - **({"project_url": project_url} if project_url is not None else {}), - } - result = ApplicationTracingService(case.repository).status(**source(case)) - assert result["project_url"] is None - assert "encrypted-original" not in str(result) - - -@pytest.mark.parametrize("provider", ["aliyun", "tencent"]) -def test_status_does_not_present_platform_console_as_project( - sqlite_session_factory: sessionmaker[Session], provider: str -) -> None: - case = seed(sqlite_session_factory, "platform-console") - with case.sessions.begin() as session: - run = session.get(ApplicationTraceRun, case.run.id) - assert run is not None - run.provider = provider - run.provider_config = {"project_url": "https://console.example.com"} - assert ApplicationTracingService(case.repository).status(**source(case))["project_url"] is None - - -@pytest.mark.parametrize("disabled", [True, False]) -def test_uncollected_status_and_missing_run_retry_do_not_construct_provider( - sqlite_session_factory: sessionmaker[Session], disabled: bool -) -> None: - case = seed(sqlite_session_factory, "missing", enabled=not disabled) - if not disabled: - with case.sessions.begin() as session: - stored = session.get(ApplicationTraceRun, case.run.id) - assert stored is not None - session.delete(stored) - service = ApplicationTracingService(case.repository) - assert service.status(**source(case)) == { - "status": "not_collected", - "source_kind": "message", - "source_id": case.run.source_id, - } - if not disabled: - with pytest.raises(TraceCollectionError, match="tracing_run_not_found"): - service.retry(**source(case), acknowledge_possible_duplicates=True) - - -def test_retry_acknowledgement_guards_schedule_and_keeps_frozen_snapshot( - sqlite_session_factory: sessionmaker[Session], monkeypatch: pytest.MonkeyPatch -) -> None: - case = seed(sqlite_session_factory, "retry") - ingest(case, batch_for(case.run, case.invocation())) - terminal(case) - lease = claim(case) - assert lease is not None - case.repository.complete(case.run.id, lease[1], state="uncertain", error_code="provider_acceptance_unknown") - scheduled: list[str] = [] - service = ApplicationTracingService(case.repository) - monkeypatch.setattr(service, "schedule", scheduled.append) - with pytest.raises(TraceCollectionError, match="tracing_duplicate_acknowledgement_required"): - service.retry(**source(case), acknowledge_possible_duplicates=False) - assert scheduled == [] - result = service.retry(**source(case), acknowledge_possible_duplicates=True) - assert scheduled == [case.run.id] - assert result["status"] == "ready" - assert result["config_revision"] == case.run.config_revision - current = case.repository.get(case.run.id) - assert current is not None - assert current.provider_config == case.run.provider_config - - -@pytest.mark.parametrize( - ("failure", "code"), - [ - (RetryableTraceDispatchError("private-transport-error"), "provider_retry_exhausted"), - (ValueError("private-credential-error"), "provider_export_rejected"), - ], -) -def test_delivery_exhaustion_and_rejection_never_log_vendor_details( - sqlite_session_factory: sessionmaker[Session], - monkeypatch: pytest.MonkeyPatch, - caplog: pytest.LogCaptureFixture, - failure: Exception, - code: str, -) -> None: - apply_config_overrides(monkeypatch, OPS_AGENT_TRACE_MAX_ATTEMPTS=1) - case = seed(sqlite_session_factory, "delivery") - ingest(case, batch_for(case.run, case.invocation())) - terminal(case) - - def export(*_args: object) -> TraceExportResult: - raise failure - - monkeypatch.setattr(service_module.OpsTraceManager, "decrypt_tracing_config", lambda *_args: {}) - monkeypatch.setattr(service_module, "export_trace", export) - service = ApplicationTracingService(case.repository) - service.deliver(case.run.id) - current = case.repository.get(case.run.id) - assert current is not None - assert current.state == "failed" - assert current.error_code == code - assert "private-" not in caplog.text - service.deliver(identifier("missing-run")) - - -@pytest.mark.parametrize("capture_content", [False, True]) -@pytest.mark.parametrize( - "status", - [WorkflowNodeExecutionStatus.RUNNING, WorkflowNodeExecutionStatus.SUCCEEDED, WorkflowNodeExecutionStatus.FAILED], -) -def test_preview_terminal_uses_persisted_node_state_and_frozen_privacy( - sqlite_session_factory: sessionmaker[Session], status: WorkflowNodeExecutionStatus, capture_content: bool -) -> None: - case = seed(sqlite_session_factory, "preview", capture_content=capture_content) - node_id = identifier("preview-node") - workflow_id = identifier("preview-definition") - with case.sessions.begin() as session: - session.add(make_workflow(workflow_id=workflow_id, tenant_id=case.run.tenant_id, app_id=case.run.app_id)) - session.add( - WorkflowNodeExecutionModel( - id=node_id, - node_execution_id=node_id, - tenant_id=case.run.tenant_id, - app_id=case.run.app_id, - workflow_id=workflow_id, - triggered_from=WorkflowNodeExecutionTriggeredFrom.SINGLE_STEP, - index=0, - node_id="agent", - node_type="agent", - title="Agent preview", - inputs='{"secret":"private-input"}', - outputs='{"answer":"private-output"}', - process_data="{}", - status=status, - elapsed_time=2, - error="private-model-error" if status == WorkflowNodeExecutionStatus.FAILED else None, - created_at=datetime(2026, 1, 1), - created_by=identifier("account"), - created_by_role=CreatorUserRole.ACCOUNT, - ) - ) - run = case.repository.start( - tenant_id=case.run.tenant_id, - app_id=case.run.app_id, - source_kind="node", - source_id=node_id, - workflow_id=workflow_id, - session_id="preview-session", - ) - assert run is not None - ApplicationTracingService(case.repository).settle_preview(run) - current = case.repository.get(run.id) - assert current is not None - if status == WorkflowNodeExecutionStatus.RUNNING: - assert current.outer_trace is None - else: - assert current.outer_trace is not None - trace = service_module.CanonicalTrace.model_validate(current.outer_trace) - root = trace.spans[0] - assert root.status.value == ("error" if status == WorkflowNodeExecutionStatus.FAILED else "ok") - assert ("private-" in trace.model_dump_json()) is capture_content - assert root.metadata["node_execution_id"] == node_id - if not capture_content and status == WorkflowNodeExecutionStatus.FAILED: - assert root.error == "Node execution failed" - - -@pytest.mark.parametrize("owner_state", ["no_creator", "missing_account", "wrong_app"]) -def test_missing_owner_identity_cannot_authorize_trace_expansion( - sqlite_session_factory: sessionmaker[Session], owner_state: str -) -> None: - from models.model import App - - case = seed(sqlite_session_factory, "owner") - with case.sessions.begin() as session: - app = session.get(App, case.run.app_id) - assert app is not None - app.created_by = None if owner_state == "no_creator" else identifier("missing-owner") - app_id = identifier("another-app") if owner_state == "wrong_app" else case.run.app_id - with pytest.raises(TraceCollectionError, match="tracing_owner_not_found"): - ApplicationTracingService(case.repository)._service_account(case.run, app_id) - - -@pytest.mark.parametrize( - ("invalid", "code"), - [ - ("empty", "tracing_invocation_empty"), - ("foreign_parent", "tracing_invocation_parent_invalid"), - ("no_invocation_root", "tracing_invocation_root_missing"), - ("conflicting_span", "tracing_span_conflict"), - ], -) -def test_invalid_runtime_tree_is_rejected_before_provider_credentials_or_io( - sqlite_session_factory: sessionmaker[Session], monkeypatch: pytest.MonkeyPatch, invalid: str, code: str -) -> None: - case = seed(sqlite_session_factory, "invalid-runtime") - batch = batch_for(case.run, case.invocation()) - if invalid == "empty": - batch.spans = [] - elif invalid == "foreign_parent": - batch.spans[0] = batch.spans[0].model_copy(update={"parent_span_id": "f" * 16}) - elif invalid == "no_invocation_root": - batch.spans[0] = batch.spans[0].model_copy(update={"parent_span_id": batch.spans[0].span_id}) - else: - ingest(case, batch.model_copy(update={"final": False})) - from uuid import uuid4 - - batch = batch.model_copy( - update={"batch_id": uuid4(), "spans": [batch.spans[0].model_copy(update={"name": "conflicting-version"})]} - ) - ingest(case, batch) - terminal(case) - - def forbidden(*_args: object, **_kwargs: object) -> None: - pytest.fail("invalid trees must fail before credentials or provider requests") - - monkeypatch.setattr(service_module.OpsTraceManager, "decrypt_tracing_config", forbidden) - monkeypatch.setattr(service_module, "export_trace", forbidden) - service = ApplicationTracingService(case.repository) - service.deliver(case.run.id) - current = case.repository.get(case.run.id) - assert current is not None - assert current.state == "failed" - assert current.error_code == code - - -def test_assemble_requires_a_persisted_business_terminal(sqlite_session_factory: sessionmaker[Session]) -> None: - case = seed(sqlite_session_factory, "no-terminal") - with pytest.raises(TraceCollectionError, match="tracing_terminal_missing"): - ApplicationTracingService(case.repository).assemble(case.run) diff --git a/api/tests/unit_tests/services/tracing/test_arize_configuration.py b/api/tests/unit_tests/services/tracing/test_arize_configuration.py deleted file mode 100644 index 7e6a1e587d1..00000000000 --- a/api/tests/unit_tests/services/tracing/test_arize_configuration.py +++ /dev/null @@ -1,260 +0,0 @@ -"""Arize metadata reads are optional, bounded and pinned to the exact destination.""" - -from collections import deque -from collections.abc import Iterator -from typing import Never - -import httpx -import pytest -from dify_trace_arize_phoenix import arize_configuration as configuration -from dify_trace_arize_phoenix.arize_configuration import ( - DESTINATION_KEY, - bind_destination, - project_url, - resolve_project_destination, -) - -from core.helper import ssrf_proxy - -SPACE = "U3BhY2U6dGVzdA==" -ORG = "T3JnOnRlc3Q=" -PROJECT = "UHJvamVjdDovKz0=" -VALUES = {"endpoint": "https://otlp.arize.com", "space_id": SPACE, "project": "exact", "api_key": "synthetic-key"} -DESTINATION = { - "endpoint": VALUES["endpoint"], - "space_id": SPACE, - "project": "exact", - "organization_id": ORG, - "project_id": PROJECT, -} -URL = "https://app.arize.com/organizations/T3JnOnRlc3Q%3D/spaces/U3BhY2U6dGVzdA%3D%3D/projects/UHJvamVjdDovKz0%3D" - - -def response(body: object, status: int = 200) -> httpx.Response: - return httpx.Response(status, json=body, request=httpx.Request("GET", "https://api.arize.com/v2/projects")) - - -def space() -> httpx.Response: - return response({"id": SPACE, "organization_id": ORG}) - - -def page(*projects: dict[str, str], more: bool = False, cursor: str | None = None) -> httpx.Response: - return response({"projects": list(projects), "pagination": {"has_more": more, "next_cursor": cursor}}) - - -def project(name: str = "exact", project_id: str = PROJECT) -> dict[str, str]: - return {"id": project_id, "name": name, "space_id": SPACE} - - -@pytest.fixture -def responses() -> deque[httpx.Response | Exception]: - return deque() - - -@pytest.fixture -def calls( - monkeypatch: pytest.MonkeyPatch, responses: deque[httpx.Response | Exception] -) -> list[tuple[str, dict[str, object]]]: - calls: list[tuple[str, dict[str, object]]] = [] - - def get(url: str, **options: object) -> httpx.Response: - calls.append((url, options)) - assert responses, "Unplanned metadata request" - item = responses.popleft() - if isinstance(item, Exception): - raise item - return item - - monkeypatch.setattr(ssrf_proxy, "get", get) - return calls - - -def test_exact_scoped_project_and_opaque_ids_use_official_route( - calls: list[tuple[str, dict[str, object]]], responses: deque[httpx.Response | Exception] -) -> None: - received = [space(), page(project("exact-substring"), more=True, cursor="next"), page(project())] - responses.extend(received) - assert resolve_project_destination(VALUES) == DESTINATION - assert ( - project_url( - {**VALUES, DESTINATION_KEY: bind_destination(VALUES, DESTINATION), "project_url": "https://untrusted.test"} - ) - == URL - ) - assert len(calls) == 3 - assert calls[0][0] == "https://api.arize.com/v2/spaces/U3BhY2U6dGVzdA%3D%3D" - assert calls[1][1]["params"] == {"space_id": SPACE, "name": "exact", "limit": 100} - assert calls[2][1]["params"] == {"space_id": SPACE, "name": "exact", "limit": 100, "cursor": "next"} - for url, options in calls: - assert url.startswith("https://api.arize.com/v2/") - assert options["headers"] == { - "Authorization": "Bearer synthetic-key", - "Accept": "application/json", - "Accept-Encoding": "identity", - "Cookie": "", - } - assert options["auth"] is None - assert options["follow_redirects"] is False - assert options["max_retries"] == 0 - assert options["stream_response"] is True - timeout = options["timeout"] - assert isinstance(timeout, float) - assert 0 < timeout <= 5 - assert all(item.is_closed for item in received) - - -@pytest.mark.parametrize("status", [301, 401, 403, 404, 429, 500]) -def test_unreadable_metadata_has_no_link_and_closes_response( - calls: list[tuple[str, dict[str, object]]], responses: deque[httpx.Response | Exception], status: int -) -> None: - denied = response({"error": "private-response"}, status) - responses.append(denied) - assert resolve_project_destination(VALUES) is None - assert len(calls) == 1 - assert denied.is_closed - - -@pytest.mark.parametrize( - "body", [{"id": SPACE}, {"id": "other", "organization_id": ORG}, {"id": SPACE, "organization_id": "../bad"}] -) -def test_missing_or_mismatched_space_metadata_is_not_a_project( - calls: list[tuple[str, dict[str, object]]], responses: deque[httpx.Response | Exception], body: dict[str, str] -) -> None: - responses.append(response(body)) - assert resolve_project_destination(VALUES) is None - assert len(calls) == 1 - - -@pytest.mark.parametrize( - "body", - [ - {"projects": [project("EXACT")], "pagination": {"has_more": False}}, - {"projects": [project(), project(project_id="other")], "pagination": {"has_more": False}}, - {"projects": [{**project(), "space_id": "other"}], "pagination": {"has_more": False}}, - {"projects": [{**project(), "id": "../bad"}], "pagination": {"has_more": False}}, - {"projects": [project()] * 101, "pagination": {"has_more": False}}, - {"projects": [project()], "pagination": {"has_more": True}}, - {"projects": [project()], "pagination": {"has_more": "false"}}, - {"projects": [project()]}, - {"data": [project()], "pagination": {"has_more": False}}, - ], -) -def test_incomplete_ambiguous_or_non_exact_projects_do_not_fabricate_links( - calls: list[tuple[str, dict[str, object]]], responses: deque[httpx.Response | Exception], body: object -) -> None: - responses.extend([space(), response(body)]) - assert resolve_project_destination(VALUES) is None - assert len(calls) == 2 - - -@pytest.mark.parametrize("repeat", [False, True]) -def test_pagination_is_bounded_and_never_accepts_a_partial_match( - calls: list[tuple[str, dict[str, object]]], responses: deque[httpx.Response | Exception], repeat: bool -) -> None: - responses.append(space()) - responses.extend(page(project(), more=True, cursor="same" if repeat else str(index)) for index in range(3)) - assert resolve_project_destination(VALUES) is None - assert len(calls) == (3 if repeat else 4) - - -@pytest.mark.parametrize( - "item", - [ - httpx.ReadTimeout("synthetic-timeout"), - httpx.Response(200, content=b"not-json"), - httpx.Response(200, json=[]), - httpx.Response(200, content=b" " * (1024 * 1024 + 1)), - httpx.Response(200, json={}, headers={"content-encoding": "unsupported"}), - ], -) -def test_transport_and_bounded_decode_failures_remain_optional( - calls: list[tuple[str, dict[str, object]]], - responses: deque[httpx.Response | Exception], - item: httpx.Response | Exception, -) -> None: - responses.append(item) - assert resolve_project_destination(VALUES) is None - assert len(calls) == 1 - if isinstance(item, httpx.Response): - assert item.is_closed - - -def test_all_pages_share_one_deadline(monkeypatch: pytest.MonkeyPatch) -> None: - clock = [100.0] - calls: list[float] = [] - received = [space(), page(project(), more=True, cursor="next")] - - def get(_url: str, **options: object) -> httpx.Response: - timeout = options["timeout"] - assert isinstance(timeout, float) - calls.append(timeout) - clock[0] += 2.5 - return received[len(calls) - 1] - - monkeypatch.setattr(configuration.time, "monotonic", lambda: clock[0]) - monkeypatch.setattr(ssrf_proxy, "get", get) - assert resolve_project_destination(VALUES) is None - assert calls == [5.0, 2.5] - assert all(item.is_closed for item in received) - - -def test_slow_response_chunks_stop_at_the_shared_deadline(monkeypatch: pytest.MonkeyPatch) -> None: - clock = [100.0] - chunks: list[int] = [] - received = httpx.Response(200) - - def stream() -> Iterator[bytes]: - for index in range(10): - clock[0] += 2.0 - chunks.append(index) - yield b" " - - monkeypatch.setattr(configuration.time, "monotonic", lambda: clock[0]) - monkeypatch.setattr(received, "iter_bytes", stream) - monkeypatch.setattr(ssrf_proxy, "get", lambda _url, **_options: received) - assert resolve_project_destination(VALUES) is None - assert chunks == [0, 1, 2] - assert received.is_closed - - -@pytest.mark.parametrize( - "change", - [ - {"endpoint": "https://custom.test"}, - {"endpoint": "https://otlp.eu-west-1a.arize.com"}, - {"endpoint": "http://otlp.arize.com"}, - {"endpoint": "https://user:secret@otlp.arize.com"}, - {"endpoint": "https://otlp.arize.com:444"}, - {"space_id": "../bad"}, - {"api_key": None}, - {"api_key": ""}, - {"api_key": "******"}, - {"api_key": "key\nheader"}, - ], -) -def test_unsupported_destination_or_credentials_do_not_send_to_an_assumed_host( - calls: list[tuple[str, dict[str, object]]], change: dict[str, str | None] -) -> None: - assert resolve_project_destination({**VALUES, **change}) is None - assert calls == [] - - -@pytest.mark.parametrize( - "change", - [ - {"endpoint": "https://otlp.arize.com:443"}, - {"space_id": "other"}, - {"project": "later"}, - {DESTINATION_KEY: {**DESTINATION, "organization_id": "bad?query"}}, - {DESTINATION_KEY: {**DESTINATION, "project_id": "../bad"}}, - {DESTINATION_KEY: None}, - ], -) -def test_get_formatter_rejects_stale_or_forged_metadata_without_io( - monkeypatch: pytest.MonkeyPatch, change: dict[str, object] -) -> None: - def forbidden(*_args: object, **_kwargs: object) -> Never: - raise AssertionError("GET must not call the vendor") - - monkeypatch.setattr(ssrf_proxy, "get", forbidden) - assert project_url({**VALUES, DESTINATION_KEY: bind_destination(VALUES, DESTINATION), **change}) is None diff --git a/api/tests/unit_tests/services/tracing/test_arize_configuration_ack.py b/api/tests/unit_tests/services/tracing/test_arize_configuration_ack.py deleted file mode 100644 index 1bf6b77dcfc..00000000000 --- a/api/tests/unit_tests/services/tracing/test_arize_configuration_ack.py +++ /dev/null @@ -1,315 +0,0 @@ -"""A real gRPC collector must acknowledge Arize credentials before config persistence.""" - -from collections.abc import Iterator -from concurrent.futures import ThreadPoolExecutor -from threading import Event -from typing import Never - -import grpc # type: ignore[import-untyped] -import httpx -import pytest -from dify_trace_arize_phoenix import arize_configuration as configuration -from dify_trace_arize_phoenix.arize_configuration import ArizeConfigurationError, verify_configuration -from dify_trace_arize_phoenix.arize_phoenix_trace import ArizePhoenixDataTrace -from dify_trace_arize_phoenix.config import ArizeConfig -from opentelemetry.context import _SUPPRESS_INSTRUMENTATION_KEY, get_value -from opentelemetry.proto.collector.trace.v1.trace_service_pb2 import ( - ExportTraceServiceRequest, - ExportTraceServiceResponse, -) - -from core.helper import ssrf_proxy -from core.ops.ops_trace_manager import OpsTraceManager -from services.app_tracing_config_gateway import OpsTraceManagerGateway -from services.app_tracing_config_service import AppTracingConfigVerificationFailedError - -pytest_plugins = ("tests.unit_tests.services.tracing.test_nullable_tracing_credentials",) - - -@pytest.fixture -def collector() -> Iterator[tuple[str, list[tuple[bytes, dict[str, str]]], dict[str, object]]]: - calls: list[tuple[bytes, dict[str, str]]] = [] - state: dict[str, object] = {"status": grpc.StatusCode.OK} - release = Event() - state["release"] = release - - def export(request: ExportTraceServiceRequest, context: grpc.ServicerContext) -> ExportTraceServiceResponse: - headers = { - key: value.decode() if isinstance(value, bytes) else value for key, value in context.invocation_metadata() - } - calls.append((request.SerializeToString(), headers)) - if headers.get("api_key") != "synthetic-good" or headers.get("space_id") != "synthetic-space": - context.abort(grpc.StatusCode.PERMISSION_DENIED, "private-vendor-details") - code = state["status"] - if isinstance(code, grpc.StatusCode) and code is not grpc.StatusCode.OK: - context.abort(code, "private-vendor-details") - if state.get("delay"): - release.wait(2) - result = ExportTraceServiceResponse() - if state.get("partial"): - result.partial_success.rejected_spans = 1 - if state.get("message"): - result.partial_success.error_message = "private-vendor-details" - return result - - def dispatch(*args: object, **kwargs: object) -> ExportTraceServiceResponse: - assert kwargs == {} - request, context = args - assert isinstance(request, ExportTraceServiceRequest) - assert isinstance(context, grpc.ServicerContext) - return export(request, context) - - handler = grpc.unary_unary_rpc_method_handler( - dispatch, - request_deserializer=ExportTraceServiceRequest.FromString, - response_serializer=ExportTraceServiceResponse.SerializeToString, - ) - with ThreadPoolExecutor(max_workers=2) as pool: - server = grpc.server(pool) - server.add_generic_rpc_handlers( - ( - grpc.method_handlers_generic_handler( - "opentelemetry.proto.collector.trace.v1.TraceService", {"Export": handler} - ), - ) - ) - port = server.add_insecure_port("127.0.0.1:0") - assert port > 0 - server.start() - try: - yield f"http://127.0.0.1:{port}", calls, state - finally: - release.set() - server.stop(0).wait(2) - - -def test_rejected_real_collector_ack_cannot_be_saved( - workspace: str, - collector: tuple[str, list[tuple[bytes, dict[str, str]]], dict[str, object]], -) -> None: - endpoint, calls, _ = collector - with pytest.raises(AppTracingConfigVerificationFailedError): - OpsTraceManagerGateway().prepare_new_config( - workspace_id=workspace, - tracing_provider="arize", - tracing_config={ - "endpoint": endpoint, - "project": "synthetic-project", - "space_id": "synthetic-space", - "api_key": "synthetic-invalid", - }, - ) - assert len(calls) == 1 - - -def test_synchronous_ack_is_resource_only_and_has_no_sdk_or_environment_side_effects( - collector: tuple[str, list[tuple[bytes, dict[str, str]]], dict[str, object]], - monkeypatch: pytest.MonkeyPatch, -) -> None: - endpoint, calls, _ = collector - attempted: list[bool] = [] - - def forbidden(*_args: object, **_kwargs: object) -> Never: - attempted.append(True) - raise AssertionError("Arize verification must not construct the SDK tracer") - - monkeypatch.setattr(ArizePhoenixDataTrace, "__init__", forbidden) - monkeypatch.setenv("OTEL_EXPORTER_OTLP_HEADERS", "api_key=ambient-invalid,space_id=ambient-other") - verify_configuration( - ArizeConfig( - endpoint=endpoint, api_key="synthetic-good", space_id="synthetic-space", project="synthetic-project" - ) - ) - assert attempted == [] - assert len(calls) == 1 - payload = ExportTraceServiceRequest.FromString(calls[0][0]) - assert len(payload.resource_spans) == 1 - assert {a.key: a.value.string_value for a in payload.resource_spans[0].resource.attributes} == { - "model_id": "synthetic-project", - "openinference.project.name": "synthetic-project", - } - assert len(payload.resource_spans[0].scope_spans) == 0 - assert calls[0][1]["api_key"] == "synthetic-good" - assert calls[0][1]["space_id"] == "synthetic-space" - assert calls[0][1]["authorization"] == "Bearer synthetic-good" - assert "ambient" not in str(calls) - - -@pytest.mark.parametrize( - ("code", "expected"), - [ - (grpc.StatusCode.UNAUTHENTICATED, "authentication_failed"), - (grpc.StatusCode.PERMISSION_DENIED, "permission_denied"), - (grpc.StatusCode.UNAVAILABLE, "connection_failed"), - (grpc.StatusCode.RESOURCE_EXHAUSTED, "connection_failed"), - ], -) -def test_real_rpc_rejection_uses_only_safe_fixed_diagnostics( - collector: tuple[str, list[tuple[bytes, dict[str, str]]], dict[str, object]], - code: grpc.StatusCode, - expected: str, - caplog: pytest.LogCaptureFixture, -) -> None: - endpoint, calls, state = collector - state["status"] = code - with pytest.raises(ArizeConfigurationError) as caught: - verify_configuration(ArizeConfig(endpoint=endpoint, api_key="synthetic-good", space_id="synthetic-space")) - assert caught.value.code == expected - assert caught.value.__suppress_context__ - assert "synthetic-good" not in str(caught.value) - assert "private-vendor-details" not in str(caught.value) - assert "private-vendor-details" not in caplog.text - assert len(calls) == 1 - - -@pytest.mark.parametrize("partial", ["partial", "message"]) -def test_partial_or_error_ack_is_not_success( - collector: tuple[str, list[tuple[bytes, dict[str, str]]], dict[str, object]], - partial: str, -) -> None: - endpoint, calls, state = collector - state[partial] = True - with pytest.raises(ArizeConfigurationError, match="unexpected_response"): - verify_configuration(ArizeConfig(endpoint=endpoint, api_key="synthetic-good", space_id="synthetic-space")) - assert len(calls) == 1 - - -def test_deadline_is_finite_and_does_not_retry( - collector: tuple[str, list[tuple[bytes, dict[str, str]]], dict[str, object]], - monkeypatch: pytest.MonkeyPatch, -) -> None: - endpoint, calls, state = collector - state["delay"] = True - monkeypatch.setattr(configuration, "_VERIFY_TIMEOUT_SECONDS", 0.05) - with pytest.raises(ArizeConfigurationError, match="connection_failed"): - verify_configuration(ArizeConfig(endpoint=endpoint, api_key="synthetic-good", space_id="synthetic-space")) - assert len(calls) == 1 - - -@pytest.mark.parametrize("reject", [False, True]) -def test_tls_channel_is_explicit_suppressed_and_closed_on_all_outcomes( - collector: tuple[str, list[tuple[bytes, dict[str, str]]], dict[str, object]], - monkeypatch: pytest.MonkeyPatch, - reject: bool, -) -> None: - endpoint, _, state = collector - if reject: - state["status"] = grpc.StatusCode.PERMISSION_DENIED - original = grpc.insecure_channel - options_seen: list[tuple[str, object, object]] = [] - closed: list[bool] = [] - - def secure(target: str, credentials: grpc.ChannelCredentials, options: object) -> grpc.Channel: - assert isinstance(credentials, grpc.ChannelCredentials) - options_seen.append((target, options, get_value(_SUPPRESS_INSTRUMENTATION_KEY))) - channel = original(endpoint.removeprefix("http://")) - close = channel.close - - def close_recording() -> None: - closed.append(True) - close() - - monkeypatch.setattr(channel, "close", close_recording) - return channel - - monkeypatch.setattr(grpc, "secure_channel", secure) - before = get_value(_SUPPRESS_INSTRUMENTATION_KEY) - config = ArizeConfig(api_key="synthetic-good", space_id="synthetic-space") - if reject: - with pytest.raises(ArizeConfigurationError): - verify_configuration(config) - else: - verify_configuration(config) - assert options_seen == [ - ("otlp.arize.com:443", (("grpc.enable_retries", 0), ("grpc.max_receive_message_length", 1024 * 1024)), True) - ] - assert closed == [True] - assert get_value(_SUPPRESS_INSTRUMENTATION_KEY) == before - - -@pytest.mark.parametrize( - ("field", "value"), - [ - ("api_key", None), - ("api_key", ""), - ("api_key", "******"), - ("api_key", "key\nheader"), - ("space_id", None), - ("space_id", " "), - ("space_id", "******"), - ("endpoint", "https://user:secret@otlp.arize.com"), - ("endpoint", "https://otlp.arize.com:0"), - ], -) -def test_invalid_auth_values_never_construct_a_channel( - monkeypatch: pytest.MonkeyPatch, - field: str, - value: str | None, -) -> None: - attempted: list[bool] = [] - - def forbidden(*_args: object, **_kwargs: object) -> Never: - attempted.append(True) - raise AssertionError("Invalid configuration must be rejected before transport construction") - - monkeypatch.setattr(grpc, "secure_channel", forbidden) - monkeypatch.setattr(grpc, "insecure_channel", forbidden) - config = ArizeConfig.model_validate({"api_key": "synthetic-good", "space_id": "synthetic-space", field: value}) - with pytest.raises(ArizeConfigurationError, match="invalid_configuration"): - verify_configuration(config) - assert attempted == [] - - -@pytest.mark.parametrize("accepted", [False, True]) -def test_legacy_api_check_also_requires_the_real_ack( - collector: tuple[str, list[tuple[bytes, dict[str, str]]], dict[str, object]], - accepted: bool, -) -> None: - endpoint, calls, _ = collector - instance = ArizePhoenixDataTrace( - ArizeConfig( - endpoint=endpoint, api_key="synthetic-good" if accepted else "synthetic-invalid", space_id="synthetic-space" - ) - ) - try: - if accepted: - assert instance.api_check() - else: - with pytest.raises(ArizeConfigurationError, match="permission_denied"): - instance.api_check() - finally: - instance.processor.shutdown() - assert len(calls) == 1 - assert len(ExportTraceServiceRequest.FromString(calls[0][0]).resource_spans[0].scope_spans) == 0 - - -def test_real_ingest_ack_without_rest_read_permission_can_still_save( - workspace: str, - collector: tuple[str, list[tuple[bytes, dict[str, str]]], dict[str, object]], - monkeypatch: pytest.MonkeyPatch, -) -> None: - endpoint, rpc_calls, _ = collector - original = grpc.insecure_channel - reads: list[str] = [] - - def secure(_target: str, _credentials: grpc.ChannelCredentials, **_options: object) -> grpc.Channel: - return original(endpoint.removeprefix("http://")) - - def get(url: str, **_options: object) -> httpx.Response: - reads.append(url) - return httpx.Response(403, json={"message": "private-vendor-details"}) - - monkeypatch.setattr(grpc, "secure_channel", secure) - monkeypatch.setattr(ssrf_proxy, "get", get) - saved = OpsTraceManagerGateway().prepare_new_config( - workspace_id=workspace, - tracing_provider="arize", - tracing_config={"api_key": "synthetic-good", "space_id": "synthetic-space", "project": "not-created"}, - ) - assert len(rpc_calls) == len(reads) == 1 - assert reads[0] == "https://api.arize.com/v2/spaces/synthetic-space" - assert "project_url" not in saved - assert "arize_project_destination" not in saved - assert OpsTraceManager.decrypt_tracing_config(workspace, "arize", saved)["api_key"] == "synthetic-good" - payload = ExportTraceServiceRequest.FromString(rpc_calls[0][0]) - assert len(payload.resource_spans[0].scope_spans) == 0 diff --git a/api/tests/unit_tests/services/tracing/test_export_lease.py b/api/tests/unit_tests/services/tracing/test_export_lease.py deleted file mode 100644 index 7d742207d6b..00000000000 --- a/api/tests/unit_tests/services/tracing/test_export_lease.py +++ /dev/null @@ -1,223 +0,0 @@ -"""Real SQLite state transitions and background renewal for delivery leases. - -PostgreSQL row-lock scheduling and the deployed Celery pool still need dev -validation; these tests establish token and acknowledgement semantics. -""" - -import threading -from datetime import timedelta - -import pytest -from sqlalchemy.orm import Session, sessionmaker - -from libs.datetime_utils import naive_utc_now -from models.application_trace import ApplicationTraceRun -from repositories import application_trace_repository as repository_module -from repositories.application_trace_repository import TraceCollectionError -from services.tracing import service as service_module -from services.tracing.provider_dispatch import TraceExportResult, TraceExportUncertainError -from services.tracing.service import ApplicationTracingService -from tests.unit_tests.config_override import apply_config_overrides -from tests.unit_tests.services.tracing.test_application_trace_collection import ( - Collection, - batch_for, - claim, - ingest, - seed, - terminal, -) - - -@pytest.fixture -def case(sqlite_session_factory: sessionmaker[Session]) -> Collection: - result = seed(sqlite_session_factory, "lease-test") - ingest(result, batch_for(result.run, result.invocation())) - terminal(result) - return result - - -def _stored(case: Collection) -> ApplicationTraceRun: - row = case.repository.get(case.run.id) - assert row is not None - return row - - -def _expire(case: Collection) -> str: - leased = claim(case) - assert leased is not None - with case.sessions.begin() as session: - row = session.get(ApplicationTraceRun, case.run.id) - assert row is not None - row.export_started_at = naive_utc_now() - timedelta(seconds=60) - row.export_heartbeat_at = row.export_started_at - assert claim(case) is None - assert _stored(case).state == "uncertain" - return leased[1] - - -def test_renewal_keeps_long_export_alive_then_missing_heartbeat_expires( - case: Collection, monkeypatch: pytest.MonkeyPatch -) -> None: - clock = [naive_utc_now()] - monkeypatch.setattr(repository_module, "naive_utc_now", lambda: clock[0]) - leased = claim(case) - assert leased is not None - token = leased[1] - clock[0] += timedelta(seconds=20) - assert case.repository.renew_export_lease(case.run.id, token) - clock[0] += timedelta(seconds=20) - assert claim(case) is None - current = _stored(case) - assert current.state == "exporting" - assert current.attempts == 1 - assert current.export_started_at is not None - assert current.export_started_at + timedelta(seconds=30) < clock[0] - clock[0] += timedelta(seconds=11) - assert claim(case) is None - expired = _stored(case) - assert expired.state == "uncertain" - assert expired.error_code == "export_lease_expired" - assert expired.export_token == token - assert not case.repository.renew_export_lease(case.run.id, token) - - -def test_old_rows_without_heartbeat_use_started_at(case: Collection) -> None: - leased = claim(case) - assert leased is not None - with case.sessions.begin() as session: - row = session.get(ApplicationTraceRun, case.run.id) - assert row is not None - row.export_started_at = naive_utc_now() - timedelta(seconds=60) - row.export_heartbeat_at = None - assert claim(case) is None - current = _stored(case) - assert current.state == "uncertain" - assert current.export_token == leased[1] - - -def test_wrong_token_cannot_renew_or_complete(case: Collection) -> None: - leased = claim(case) - assert leased is not None - before = _stored(case) - assert not case.repository.renew_export_lease(case.run.id, "other-token") - case.repository.complete(case.run.id, "other-token", state="accepted", result=TraceExportResult("wrong")) - current = _stored(case) - assert current.export_token == leased[1] - assert current.export_heartbeat_at == before.export_heartbeat_at - assert current.state == "exporting" - assert current.provider_trace_id is None - - -def test_late_success_wins_before_manual_retry_and_is_idempotent(case: Collection) -> None: - token = _expire(case) - result = TraceExportResult("confirmed", "https://example.com/trace") - case.repository.complete(case.run.id, token, state="accepted", result=result) - current = _stored(case) - assert current.state == "accepted" - assert current.error_code is None - assert current.export_token is None - assert current.accepted_at is not None - assert current.trace_url == result.trace_url - case.repository.complete(case.run.id, token, state="accepted", result=TraceExportResult("duplicate-ack")) - assert _stored(case).provider_trace_id == "confirmed" - with pytest.raises(TraceCollectionError, match="tracing_already_accepted"): - case.repository.retry(run_id=case.run.id, acknowledge_possible_duplicates=True) - - -def test_manual_retry_revokes_old_heartbeat_and_ack_before_and_after_new_claim(case: Collection) -> None: - old_token = _expire(case) - case.repository.retry(run_id=case.run.id, acknowledge_possible_duplicates=True) - assert not case.repository.renew_export_lease(case.run.id, old_token) - case.repository.complete(case.run.id, old_token, state="accepted", result=TraceExportResult("old")) - assert _stored(case).state == "ready" - second = claim(case) - assert second is not None - assert second[1] != old_token - case.repository.complete(case.run.id, old_token, state="accepted", result=TraceExportResult("old")) - current = _stored(case) - assert current.state == "exporting" - assert current.export_token == second[1] - assert current.provider_trace_id is None - case.repository.complete(case.run.id, second[1], state="accepted", result=TraceExportResult("new")) - assert _stored(case).provider_trace_id == "new" - - -@pytest.mark.parametrize("late_state", ["failed", "retrying", "uncertain"]) -def test_late_non_success_cannot_overwrite_lease_expiry(case: Collection, late_state: str) -> None: - token = _expire(case) - case.repository.complete(case.run.id, token, state=late_state, error_code="late-failure") - current = _stored(case) - assert current.state == "uncertain" - assert current.error_code == "export_lease_expired" - - -def test_other_uncertain_reason_cannot_be_promoted_by_late_ack(case: Collection) -> None: - token = _expire(case) - with case.sessions.begin() as session: - row = session.get(ApplicationTraceRun, case.run.id) - assert row is not None - row.error_code = "provider_acceptance_unknown" - case.repository.complete(case.run.id, token, state="accepted", result=TraceExportResult("late")) - assert _stored(case).state == "uncertain" - - -@pytest.mark.parametrize("transient_failure", [False, True]) -@pytest.mark.parametrize("outcome", ["accepted", "uncertain"]) -def test_service_heartbeats_during_export_and_stops_on_every_outcome( - case: Collection, - monkeypatch: pytest.MonkeyPatch, - caplog: pytest.LogCaptureFixture, - transient_failure: bool, - outcome: str, -) -> None: - apply_config_overrides(monkeypatch, OPS_AGENT_TRACE_EXPORT_LEASE_SECONDS=1) - clock = [naive_utc_now()] - monkeypatch.setattr(repository_module, "naive_utc_now", lambda: clock[0]) - monkeypatch.setattr(service_module.OpsTraceManager, "decrypt_tracing_config", lambda *_args: {}) - renewed = threading.Event() - calls = [] - original = case.repository.renew_export_lease - - def renew(run_id: str, token: str) -> bool: - calls.append(threading.get_ident()) - if transient_failure and len(calls) == 1: - raise RuntimeError("private-provider-value") - result = original(run_id, token) - renewed.set() - return result - - monkeypatch.setattr(case.repository, "renew_export_lease", renew) - - def export(*_args: object) -> TraceExportResult: - # The original start time is now beyond the lease; only a real background - # transaction can prevent a second claimant from marking it uncertain. - clock[0] += timedelta(seconds=3) - assert renewed.wait(3), "background heartbeat did not commit" - assert case.repository.claim(case.run.id, collection_timeout=60, lease_seconds=1) is None - assert _stored(case).state == "exporting" - if outcome == "uncertain": - raise TraceExportUncertainError("lost acknowledgement") - return TraceExportResult("ack") - - monkeypatch.setattr(service_module, "export_trace", export) - ApplicationTracingService(case.repository).deliver(case.run.id) - assert _stored(case).state == outcome - assert calls - assert all(thread_id != threading.get_ident() for thread_id in calls) - assert not any(thread.name == f"application-trace-lease-{case.run.id}" for thread in threading.enumerate()) - assert "private-provider-value" not in caplog.text - - -def test_heartbeat_stops_even_when_completion_database_write_raises( - case: Collection, monkeypatch: pytest.MonkeyPatch -) -> None: - monkeypatch.setattr(service_module.OpsTraceManager, "decrypt_tracing_config", lambda *_args: {}) - monkeypatch.setattr(service_module, "export_trace", lambda *_args: TraceExportResult("ack")) - - def fail_complete(*_args: object, **_kwargs: object) -> None: - raise RuntimeError("database unavailable") - - monkeypatch.setattr(case.repository, "complete", fail_complete) - with pytest.raises(RuntimeError, match="database unavailable"): - ApplicationTracingService(case.repository).deliver(case.run.id) - assert not any(thread.name == f"application-trace-lease-{case.run.id}" for thread in threading.enumerate()) diff --git a/api/tests/unit_tests/services/tracing/test_invalid_trace_delivery.py b/api/tests/unit_tests/services/tracing/test_invalid_trace_delivery.py deleted file mode 100644 index fbfd62dd790..00000000000 --- a/api/tests/unit_tests/services/tracing/test_invalid_trace_delivery.py +++ /dev/null @@ -1,55 +0,0 @@ -"""Malformed durable traces must never be sent or offered for provider retry.""" - -from uuid import uuid4 - -import pytest -from dify_agent.protocol.tracing import AgentTraceSpan -from sqlalchemy.orm import Session, sessionmaker - -from models.application_trace import ApplicationTraceRun -from repositories.application_trace_repository import TraceCollectionError -from services.tracing import service as service_module -from services.tracing.service import ApplicationTracingService -from tests.unit_tests.services.tracing.test_application_trace_collection import batch_for, ingest, seed, terminal - - -@pytest.mark.parametrize("invalid", ["cycle", "invalid_envelope"]) -def test_invalid_tree_is_permanent_before_decryption_or_provider_io( - sqlite_session_factory: sessionmaker[Session], monkeypatch: pytest.MonkeyPatch, invalid: str -) -> None: - case = seed(sqlite_session_factory, str(uuid4())) - batch = batch_for(case.run, case.invocation()) - if invalid == "cycle": - for span_id, parent_id in [("b" * 16, "c" * 16), ("c" * 16, "b" * 16)]: - batch.spans.append( - AgentTraceSpan( - trace_id=case.run.trace_id, - span_id=span_id, - parent_span_id=parent_id, - name="cyclic tool", - start_time_ns=1_000_000_000, - end_time_ns=2_000_000_000, - status="ok", - ) - ) - ingest(case, batch) - terminal(case) - if invalid == "invalid_envelope": - with case.sessions.begin() as session: - row = session.get(ApplicationTraceRun, case.run.id) - assert row is not None - row.outer_trace = {"corrupt": True} - - def forbidden(*_args: object, **_kwargs: object) -> None: - pytest.fail("invalid traces must fail before credentials or provider I/O") - - monkeypatch.setattr(service_module.OpsTraceManager, "decrypt_tracing_config", forbidden) - monkeypatch.setattr(service_module, "export_trace", forbidden) - ApplicationTracingService(case.repository).deliver(case.run.id) - current = case.repository.get(case.run.id) - assert current is not None - assert current.state == "failed" - assert current.error_code == "tracing_payload_invalid" - assert current.attempts == 1 - with pytest.raises(TraceCollectionError, match="tracing_retry_not_allowed"): - case.repository.retry(run_id=case.run.id, acknowledge_possible_duplicates=True) diff --git a/api/tests/unit_tests/services/tracing/test_legacy_aliyun_transport.py b/api/tests/unit_tests/services/tracing/test_legacy_aliyun_transport.py deleted file mode 100644 index 811cf01feb4..00000000000 --- a/api/tests/unit_tests/services/tracing/test_legacy_aliyun_transport.py +++ /dev/null @@ -1,344 +0,0 @@ -"""Persisted legacy sources through real OTel protobuf, replacing only HTTP I/O.""" - -import json -import os -from collections.abc import Mapping -from concurrent.futures import ThreadPoolExecutor -from datetime import UTC, datetime, timedelta -from decimal import Decimal -from threading import Barrier, Thread -from typing import TypedDict, Unpack -from uuid import UUID - -import httpx -import pytest -from dify_trace_aliyun.aliyun_trace import AliyunDataTrace -from dify_trace_aliyun.config import AliyunConfig -from dify_trace_aliyun.data_exporter.traceclient import TraceClient, convert_to_span_id -from dify_trace_aliyun.entities.aliyun_trace_entity import SpanData -from dify_trace_aliyun.entities.semconv import GEN_AI_USAGE_TOTAL_TOKENS -from opentelemetry.exporter.otlp.proto.http.trace_exporter import OTLPSpanExporter -from opentelemetry.proto.collector.trace.v1.trace_service_pb2 import ( - ExportTraceServiceRequest, - ExportTraceServiceResponse, -) -from sqlalchemy.orm import Session - -from core.helper import ssrf_proxy -from core.ops.entities.trace_entity import MessageTraceInfo, ModerationTraceInfo, WorkflowTraceInfo -from core.ops.exceptions import RetryableTraceDispatchError -from graphon.enums import BuiltinNodeTypes, WorkflowExecutionStatus, WorkflowNodeExecutionStatus -from models.account import TenantAccountJoin, TenantAccountRole -from models.enums import ConversationFromSource, CreatorUserRole, MessageStatus -from models.model import AppMode -from models.workflow import ( - WorkflowNodeExecutionModel, - WorkflowNodeExecutionTriggeredFrom, - WorkflowRun, - WorkflowRunTriggeredFrom, - WorkflowType, -) -from services.tracing.provider_dispatch import TraceExportRejectedError, TraceExportUncertainError -from tests.unit_tests.model_factories import make_account, make_app, make_message, make_tenant - -START = datetime(2026, 10, 19, tzinfo=UTC) -RUN_ID = "a0000000-0000-4000-8000-000000000001" -MESSAGE_ID = "a0000000-0000-4000-8000-000000000002" -NODE_IDS = [f"b0000000-0000-4000-8000-00000000000{i}" for i in range(3)] - - -class RequestOptions(TypedDict, total=False): - content: bytes - headers: Mapping[str, str] - max_retries: int - follow_redirects: bool - timeout: float - - -def config(**overrides: object) -> AliyunConfig: - return AliyunConfig.model_validate( - { - "app_name": "owned-service", - "license_key": "owned-key", - "endpoint": "https://collector.test/prefix", - **overrides, - } - ) - - -def decoded(options: RequestOptions) -> ExportTraceServiceRequest: - payload = ExportTraceServiceRequest() - payload.ParseFromString(options["content"]) - return payload - - -def test_constructor_and_presentation_have_no_network_sdk_queue_or_environment_mutation( - monkeypatch: pytest.MonkeyPatch, -) -> None: - def forbidden(*_args: object, **_kwargs: object) -> None: - pytest.fail("Constructor and presentation must not create I/O, exporters or worker threads") - - monkeypatch.setattr(ssrf_proxy, "make_request", forbidden) - monkeypatch.setattr(OTLPSpanExporter, "__init__", forbidden) - monkeypatch.setattr(Thread, "start", forbidden) - monkeypatch.setenv("OTEL_EXPORTER_OTLP_ENDPOINT", "https://ambient.test") - monkeypatch.setenv("OTEL_EXPORTER_OTLP_HEADERS", "Authorization=ambient") - before = dict(os.environ) - original = config() - adapter = AliyunDataTrace(original) - original.license_key = "mutated-after-construction" - assert adapter.trace_client.endpoint == "https://collector.test/prefix/adapt_owned-key/api/otlp/traces" - assert adapter.get_project_url() == "https://arms.console.aliyun.com/#/llm" - assert not hasattr(adapter.trace_client, "exporter") - assert not hasattr(adapter.trace_client, "queue") - assert not hasattr(adapter.trace_client, "worker_thread") - assert dict(os.environ) == before - - -@pytest.mark.usefixtures("agent_runtime_app") -@pytest.mark.parametrize("failure", [None, 400, 401, 403, 429, 503, "timeout", "partial", "malformed", "redirect"]) -def test_persisted_workflow_uses_one_acknowledged_protobuf_batch_and_preserves_business_state( - sqlite_session: Session, monkeypatch: pytest.MonkeyPatch, failure: int | str | None -) -> None: - tenant = make_tenant(tenant_id="tenant") - account = make_account(account_id="owner", tenant=tenant) - app = make_app(app_id="app", tenant_id=tenant.id, mode=AppMode.WORKFLOW, created_by=account.id) - membership = TenantAccountJoin(tenant_id=tenant.id, account_id=account.id, role=TenantAccountRole.OWNER) - run = WorkflowRun( - id=RUN_ID, - tenant_id=tenant.id, - app_id=app.id, - workflow_id="workflow", - type=WorkflowType.WORKFLOW, - triggered_from=WorkflowRunTriggeredFrom.DEBUGGING, - version="v1", - inputs=json.dumps({"query": "hello"}), - outputs=json.dumps({"answer": "ok"}), - status=WorkflowExecutionStatus.SUCCEEDED, - total_tokens=32, - created_by=account.id, - created_by_role=CreatorUserRole.ACCOUNT, - created_at=START, - finished_at=START + timedelta(seconds=1), - elapsed_time=1, - ) - sqlite_session.add_all([tenant, account, app, membership, run]) - nodes = [ - WorkflowNodeExecutionModel( - id=NODE_IDS[i], - tenant_id=tenant.id, - app_id=app.id, - workflow_id="workflow", - workflow_run_id=RUN_ID, - triggered_from=WorkflowNodeExecutionTriggeredFrom.WORKFLOW_RUN, - index=i, - node_id=f"node-{i}", - title=title, - node_type=kind, - status=WorkflowNodeExecutionStatus.SUCCEEDED, - created_by=account.id, - created_by_role=CreatorUserRole.ACCOUNT, - inputs=json.dumps({"synthetic": True}), - outputs=json.dumps( - {"text": "ok", "usage": {"prompt_tokens": 21, "completion_tokens": 11, "total_tokens": 32}} - ), - created_at=START + timedelta(milliseconds=i), - finished_at=START + timedelta(milliseconds=i + 100), - elapsed_time=0.1, - process_data=json.dumps( - {"prompts": [{"role": "user", "text": "hello"}], "model_provider": "test", "model_name": "test-model"} - ) - if i == 1 - else None, - ) - for i, (title, kind) in enumerate( - [("Start", BuiltinNodeTypes.START), ("LLM", BuiltinNodeTypes.LLM), ("End", BuiltinNodeTypes.END)] - ) - ] - sqlite_session.add_all(nodes) - sqlite_session.commit() - info = WorkflowTraceInfo( - query="hello", - workflow_id="workflow", - tenant_id=tenant.id, - workflow_run_id=RUN_ID, - workflow_run_elapsed_time=1.0, - workflow_run_status="succeeded", - workflow_run_inputs={"query": "hello"}, - workflow_run_outputs={"answer": "ok"}, - workflow_run_version="v1", - total_tokens=32, - file_list=[], - metadata={"app_id": app.id, "conversation_id": "owned-conversation"}, - start_time=START, - end_time=START + timedelta(seconds=1), - ) - sent: list[tuple[str, RequestOptions, ExportTraceServiceRequest]] = [] - - def send(method: str, url: str, **options: Unpack[RequestOptions]) -> httpx.Response: - assert method == "POST" - assert url == "https://collector.test/prefix/adapt_owned-key/api/otlp/traces" - sent.append((url, options, decoded(options))) - if len(sent) == 1: - if failure == "timeout": - raise httpx.ReadTimeout("synthetic lost ACK", request=httpx.Request("POST", url)) - if isinstance(failure, int): - return httpx.Response(failure, content=b"DO_NOT_EXPOSE_BODY") - if failure == "malformed": - return httpx.Response(200, content=b"not protobuf") - if failure == "partial": - response = ExportTraceServiceResponse() - response.partial_success.rejected_spans = 1 - response.partial_success.error_message = "DO_NOT_EXPOSE_BODY" - return httpx.Response(200, content=response.SerializeToString()) - if failure == "redirect": - return httpx.Response(307, headers={"Location": "https://unrelated.test"}) - return httpx.Response(200, content=ExportTraceServiceResponse().SerializeToString()) - - monkeypatch.setattr(ssrf_proxy, "make_request", send) - adapter = AliyunDataTrace(config()) - if failure is None: - adapter.trace(info) - else: - expected = ( - RetryableTraceDispatchError - if failure in {429, 503} - else TraceExportUncertainError - if failure in {"timeout", "partial", "malformed"} - else TraceExportRejectedError - ) - with pytest.raises(expected) as error: - adapter.trace(info) - assert "DO_NOT_EXPOSE_BODY" not in str(error.value) - if isinstance(error.value, RetryableTraceDispatchError): - assert error.value.acceptance_unknown is (expected is TraceExportUncertainError) - assert len(sent) == 1 # No SDK retry, per-span request, batching thread or silent queue success. - _, options, payload = sent[0] - assert options["max_retries"] == 0 - assert options["follow_redirects"] is False - assert options["headers"] == {"Content-Type": "application/x-protobuf", "Accept": "application/x-protobuf"} - resource = payload.resource_spans[0] - attributes = {item.key: item.value for item in resource.resource.attributes} - assert attributes["service.name"].string_value == "owned-service" - spans = list(resource.scope_spans[0].spans) - assert len(spans) == 4 - assert [s.name for s in spans] == ["workflow", "Start", "LLM", "End"] - root_id = convert_to_span_id(RUN_ID, "workflow").to_bytes(8) - assert [s.span_id for s in spans] == [root_id, *[convert_to_span_id(i, "node").to_bytes(8) for i in NODE_IDS]] - assert all(s.trace_id == UUID(RUN_ID).bytes for s in spans) - assert [s.parent_span_id for s in spans] == [b"", root_id, root_id, root_id] - assert spans[0].start_time_unix_nano == int(START.timestamp() * 1_000_000_000) - assert spans[0].end_time_unix_nano == int((START + timedelta(seconds=1)).timestamp() * 1_000_000_000) - model_attributes = {item.key: item.value for item in spans[2].attributes} - assert model_attributes[GEN_AI_USAGE_TOTAL_TOKENS].string_value == "32" - assert all(s.status.code == 1 for s in spans) - assert adapter.trace_client._dispatch_spans.get() is None - assert info.workflow_run_status == "succeeded" - sqlite_session.refresh(run) - assert run.status == WorkflowExecutionStatus.SUCCEEDED - assert run.total_tokens == 32 - assert all(node.status == WorkflowNodeExecutionStatus.SUCCEEDED for node in nodes) - # A fresh business operation does not inherit any old batch after success or error. - adapter.trace(ModerationTraceInfo(flagged=False, action="allow", preset_response="", query="", metadata={})) - assert len(sent) == 1 - adapter.trace(info.model_copy(update={"workflow_run_id": "a0000000-0000-4000-8000-000000000003"})) - assert len(sent) == 2 - recovered_spans = list(sent[1][2].resource_spans[0].scope_spans[0].spans) - assert [s.name for s in recovered_spans] == ["workflow"] - - -@pytest.mark.parametrize("failed", [False, True]) -def test_message_mapping_keeps_original_ids_raw_usage_error_and_business_state( - sqlite_session: Session, monkeypatch: pytest.MonkeyPatch, failed: bool -) -> None: - message = make_message( - message_id=MESSAGE_ID, - app_id="app", - conversation_id="conversation", - query="hello", - answer="ok", - status=MessageStatus.ERROR if failed else MessageStatus.NORMAL, - inputs={}, - message={}, - message_unit_price=Decimal(0), - answer_unit_price=Decimal(0), - currency="USD", - from_source=ConversationFromSource.CONSOLE, - from_account_id="owner", - message_tokens=21, - answer_tokens=11, - ) - sqlite_session.add(message) - sqlite_session.commit() - info = MessageTraceInfo( - conversation_model="chat", - conversation_mode="chat", - message_id=MESSAGE_ID, - message_data=message, - inputs={"prompt": "hello"}, - outputs="ok", - message_tokens=21, - answer_tokens=11, - total_tokens=32, - error="synthetic model failure" if failed else None, - metadata={ - "app_id": "app", - "conversation_id": "conversation", - "ls_model_name": "test-model", - "ls_provider": "test", - }, - start_time=START, - end_time=START + timedelta(seconds=1), - ) - sent: list[ExportTraceServiceRequest] = [] - - def send(_method: str, _url: str, **options: Unpack[RequestOptions]) -> httpx.Response: - sent.append(decoded(options)) - return httpx.Response(200, content=b"") - - monkeypatch.setattr(ssrf_proxy, "make_request", send) - adapter = AliyunDataTrace(config()) - adapter.trace(info) - assert len(sent) == 1 - spans = list(sent[0].resource_spans[0].scope_spans[0].spans) - assert [s.name for s in spans] == ["message", "llm"] - root_id = convert_to_span_id(MESSAGE_ID, "message").to_bytes(8) - assert [s.parent_span_id for s in spans] == [b"", root_id] - assert [s.span_id for s in spans] == [root_id, convert_to_span_id(MESSAGE_ID, "llm").to_bytes(8)] - assert all(s.trace_id == UUID(MESSAGE_ID).bytes and s.status.code == (2 if failed else 1) for s in spans) - assert {a.key: a.value for a in spans[1].attributes}[GEN_AI_USAGE_TOTAL_TOKENS].string_value == "32" - assert message.status == (MessageStatus.ERROR if failed else MessageStatus.NORMAL) - - -def test_one_cached_client_keeps_interleaved_batches_and_failure_cleanup_isolated( - monkeypatch: pytest.MonkeyPatch, -) -> None: - client = TraceClient(service_name="owned-service", endpoint="https://collector.test/adapt_key/api/otlp/traces") - barrier = Barrier(2) - observed: list[list[str]] = [] - - def send(_method: str, _url: str, **options: Unpack[RequestOptions]) -> httpx.Response: - payload = decoded(options) - observed.append([s.name for rs in payload.resource_spans for scope in rs.scope_spans for s in scope.spans]) - return httpx.Response(200, content=b"") - - monkeypatch.setattr(ssrf_proxy, "make_request", send) - - def run(label: str) -> None: - try: - with client.batch(): - client.add_span(SpanData(name=label, trace_id=1, span_id=2, start_time=1, end_time=2)) - barrier.wait(timeout=5) - if label == "failed": - raise ValueError("mapper failed before submission") - client.add_span(SpanData(name=label + "-child", trace_id=1, span_id=3, start_time=1, end_time=2)) - except ValueError: - assert label == "failed" - assert client._dispatch_spans.get() is None - - with ThreadPoolExecutor(max_workers=2) as pool: - list(pool.map(run, ["accepted", "failed"])) - assert observed == [["accepted", "accepted-child"]] - with client.batch(): - client.add_span(SpanData(name="next", trace_id=1, span_id=4, start_time=1, end_time=2)) - assert observed == [["accepted", "accepted-child"], ["next"]] diff --git a/api/tests/unit_tests/services/tracing/test_legacy_weave_transport.py b/api/tests/unit_tests/services/tracing/test_legacy_weave_transport.py deleted file mode 100644 index 6ab4749fd89..00000000000 --- a/api/tests/unit_tests/services/tracing/test_legacy_weave_transport.py +++ /dev/null @@ -1,521 +0,0 @@ -"""Real Weave SDK serialization with only the outbound HTTP boundary replaced.""" - -import atexit -import json -import os -from collections.abc import Mapping -from concurrent.futures import ThreadPoolExecutor -from datetime import UTC, datetime, timedelta -from pathlib import Path -from threading import Barrier -from typing import TypedDict, Unpack - -import httpx -import pytest -import wandb -import weave -from dify_trace_weave.config import WeaveConfig -from dify_trace_weave.entities.weave_trace_entity import WeaveTraceModel -from dify_trace_weave.weave_trace import WeaveDataTrace -from sqlalchemy.orm import Session -from weave.trace_server.trace_server_interface import CallEndReq, CallStartReq, TraceStatus - -from core.helper import ssrf_proxy -from core.ops.entities.trace_entity import GenerateNameTraceInfo, WorkflowTraceInfo -from core.ops.exceptions import RetryableTraceDispatchError -from graphon.enums import BuiltinNodeTypes, WorkflowNodeExecutionStatus -from models.account import TenantAccountJoin, TenantAccountRole -from models.enums import CreatorUserRole -from models.model import AppMode -from models.workflow import WorkflowNodeExecutionModel, WorkflowNodeExecutionTriggeredFrom -from services.tracing.provider_dispatch import TraceExportRejectedError, TraceExportUncertainError -from tests.unit_tests.model_factories import make_account, make_app, make_tenant - -START = datetime(2026, 10, 19, tzinfo=UTC) - - -class RequestOptions(TypedDict, total=False): - content: bytes - json: dict[str, object] - auth: tuple[str, str] - headers: Mapping[str, str] - max_retries: int - follow_redirects: bool - timeout: float - - -def config(**overrides: object) -> WeaveConfig: - return WeaveConfig.model_validate({"api_key": "fixture-key", "entity": "team", "project": "project", **overrides}) - - -def call(call_id: str = "call-1", **overrides: object) -> WeaveTraceModel: - return WeaveTraceModel.model_validate( - { - "id": call_id, - "op": "tool", - "inputs": {"synthetic": True}, - "outputs": {"result": "ok"}, - "attributes": {"trace_id": call_id, "start_time": START, "end_time": START + timedelta(seconds=1)}, - **overrides, - } - ) - - -def response(url: str, body: object, *, status: int = 200) -> httpx.Response: - return httpx.Response(status, request=httpx.Request("POST", url), json=body) - - -@pytest.fixture -def wire( - monkeypatch: pytest.MonkeyPatch, failure: tuple[str, str, int | str] | None -) -> list[tuple[str, RequestOptions, dict[str, object]]]: - sent: list[tuple[str, RequestOptions, dict[str, object]]] = [] - fault_pending = True - - def send(method: str, url: str, **options: Unpack[RequestOptions]) -> httpx.Response: - nonlocal fault_pending - assert method == "POST" - assert url in {"https://trace.wandb.ai/call/start", "https://trace.wandb.ai/call/end"} - body = json.loads(options["content"]) - sent.append((url, options, body)) - stage = "start" if "start" in body else "end" - record = body[stage] - if fault_pending and failure is not None and (stage, record["id"]) == failure[:2]: - fault_pending = False - outcome = failure[2] - if outcome == "timeout": - raise httpx.ReadTimeout("synthetic lost acknowledgment", request=httpx.Request("POST", url)) - if outcome == "malformed": - return httpx.Response(200, request=httpx.Request("POST", url), content=b"invalid-json") - assert isinstance(outcome, int) - return response(url, {"error": "DO_NOT_EXPOSE_BODY"}, status=outcome) - if url.endswith("/call/start"): - req = CallStartReq.model_validate(body) - return response(url, {"id": req.start.id, "trace_id": req.start.trace_id}) - CallEndReq.model_validate(body) - return response(url, {}) - - monkeypatch.setattr(ssrf_proxy, "make_request", send) - return sent - - -def test_constructor_and_presentation_have_no_login_network_files_or_background_state( - monkeypatch: pytest.MonkeyPatch, -) -> None: - def forbidden(*_args: object, **_kwargs: object) -> None: - pytest.fail("Constructor must not initialize global credentials, files, I/O or background work") - - monkeypatch.setattr(Path, "touch", forbidden) - monkeypatch.setattr(wandb, "login", forbidden) - monkeypatch.setattr(weave, "init", forbidden) - monkeypatch.setattr("wandb.sdk.lib.wbauth.wbnetrc.write_netrc_auth", forbidden) - monkeypatch.setattr("weave.trace_server_bindings.remote_http_trace_server.AsyncBatchProcessor", forbidden) - monkeypatch.setattr("weave.trace_server_bindings.remote_http_trace_server.CallBatchProcessor", forbidden) - monkeypatch.setattr(ssrf_proxy, "make_request", forbidden) - monkeypatch.setattr(atexit, "register", forbidden) - monkeypatch.setenv("WANDB_API_KEY", "unrelated-key") - monkeypatch.setenv("WANDB_BASE_URL", "https://unrelated.test") - monkeypatch.setenv("WEAVE_TRACE_SERVER", "https://unrelated-traces.test") - before = dict(os.environ) - original = config(entity=None) - adapter = WeaveDataTrace(original) - original.api_key = "changed-after-construction" - assert adapter.weave_api_key == "fixture-key" - assert adapter.get_project_url() == "https://wandb.ai" - assert adapter.server.trace_server_url == "https://trace.wandb.ai" - assert adapter.server._auth == ("api", "fixture-key") - assert not adapter.server.should_batch - assert adapter.server.call_processor is None - assert adapter.server.feedback_processor is None - assert dict(os.environ) == before - - -@pytest.mark.usefixtures("agent_runtime_app") -@pytest.mark.parametrize( - "failure", - [None] - + [ - (stage, call_id, status) - for stage, call_id in [("start", "legacy-run"), ("end", "legacy-run"), ("start", "node-0"), ("end", "node-0")] - for status in [429, 503] - ] - + [ - ("start", "legacy-run", "timeout"), - ("start", "legacy-run", "malformed"), - ("end", "node-0", "malformed"), - ("end", "legacy-run", "timeout"), - ] - + [("start", "legacy-run", 401), ("end", "legacy-run", 401)], -) -def test_workflow_emits_four_real_sdk_calls_with_original_ids_payload_and_status( - sqlite_session: Session, - wire: list[tuple[str, RequestOptions, dict[str, object]]], - failure: tuple[str, str, int | str] | None, -) -> None: - # The normal Flask/SQLite fixtures bind the actual service-account lookup and - # production repository to persisted rows. Only outbound HTTP is substituted. - tenant = make_tenant(tenant_id="tenant") - account = make_account(account_id="owner", tenant=tenant) - app = make_app(app_id="app", tenant_id=tenant.id, mode=AppMode.WORKFLOW, created_by=account.id) - membership = TenantAccountJoin(tenant_id=tenant.id, account_id=account.id, role=TenantAccountRole.OWNER) - sqlite_session.add_all([tenant, account, app, membership]) - nodes = [ - WorkflowNodeExecutionModel( - id=f"node-{i}", - tenant_id=tenant.id, - app_id=app.id, - workflow_id="workflow", - workflow_run_id="legacy-run", - triggered_from=WorkflowNodeExecutionTriggeredFrom.WORKFLOW_RUN, - index=i, - node_id=f"node-{i}", - title=title, - node_type=kind, - status=WorkflowNodeExecutionStatus.SUCCEEDED, - created_by=account.id, - created_by_role=CreatorUserRole.ACCOUNT, - inputs=json.dumps({"synthetic": True}), - outputs=json.dumps({"answer": "ok"}), - created_at=START + timedelta(milliseconds=i), - finished_at=START + timedelta(milliseconds=i + 100), - elapsed_time=0.1, - process_data=json.dumps( - { - "prompts": [{"role": "user", "text": "hello"}], - "model_mode": "chat", - "model_provider": "test", - "model_name": "test-model", - } - ) - if i == 1 - else None, - execution_metadata=json.dumps( - {"total_tokens": 32, "total_price": "0.0000098", "currency": "USD"} if i == 1 else {} - ), - ) - for i, (title, kind) in enumerate( - [("Start", BuiltinNodeTypes.START), ("LLM", BuiltinNodeTypes.LLM), ("End", BuiltinNodeTypes.END)] - ) - ] - sqlite_session.add_all(nodes) - sqlite_session.commit() - adapter = WeaveDataTrace(config()) - info = WorkflowTraceInfo( - query="hello", - workflow_id="workflow", - tenant_id="tenant", - workflow_run_id="legacy-run", - workflow_run_elapsed_time=1.0, - workflow_run_status="succeeded", - workflow_run_inputs={"query": "hello"}, - workflow_run_outputs={"answer": "ok"}, - workflow_run_version="v1", - total_tokens=32, - file_list=[], - metadata={"app_id": "app", "conversation_id": None}, - trace_id="external-correlation", - start_time=START, - end_time=START + timedelta(seconds=1), - ) - - if failure is not None: - stage, call_id, outcome = failure - first_start = (stage, call_id) == ("start", "legacy-run") - expected = ( - TraceExportRejectedError - if outcome == 401 - else RetryableTraceDispatchError - if first_start and outcome in {429, 503} - else TraceExportUncertainError - ) - with pytest.raises(expected) as error: - adapter.trace(info) - if isinstance(error.value, RetryableTraceDispatchError): - assert error.value.acceptance_unknown is (expected is TraceExportUncertainError) - assert len(wire) == (1 if first_start else 8 if call_id == "legacy-run" else 2 if stage == "start" else 3) - assert "DO_NOT_EXPOSE" not in str(error.value) - assert adapter.calls == {} - assert info.workflow_run_status == "succeeded" - assert all(node.status == WorkflowNodeExecutionStatus.SUCCEEDED for node in nodes) - before_recovery = len(wire) - adapter.trace( - GenerateNameTraceInfo( - tenant_id=tenant.id, - inputs="new business operation", - outputs="new title", - metadata={}, - start_time=START, - end_time=START + timedelta(seconds=1), - ) - ) - assert len(wire) == before_recovery + 2 - assert adapter.calls == {} - return - - adapter.trace(info) - - starts = [CallStartReq.model_validate(body).start for _, _, body in wire if "start" in body] - ends = [CallEndReq.model_validate(body).end for _, _, body in wire if "end" in body] - assert len(starts) == len(ends) == 4 - assert [s.id for s in starts] == ["legacy-run", "node-0", "node-1", "node-2"] - assert [s.parent_id for s in starts] == [None, "legacy-run", "legacy-run", "legacy-run"] - assert all(s.trace_id == "external-correlation" and s.project_id == "team/project" for s in starts) - # Legacy has no native thread/turn contract: keep the existing raw metadata. - assert all(s.thread_id is None and s.turn_id is None for s in starts) - assert starts[0].attributes["conversation_id"] is None - model = starts[2] - assert model.attributes["total_tokens"] == 32 - assert model.attributes["total_price"] == "0.0000098" - assert model.inputs["messages"][0]["content"] == "hello" - model_end = next(e for e in ends if e.id == "node-1") - assert isinstance(model_end.output, dict) - assert model_end.output["usage_metadata"]["total_tokens"] == 32 - assert all(e.exception is None and e.summary["status_counts"][TraceStatus.SUCCESS] == 1 for e in ends) - assert all("usage" not in e.summary for e in ends) # Do not invent native aggregation for the old mapper. - assert all(options["auth"] == ("api", "fixture-key") and options["max_retries"] == 0 for _, options, _ in wire) - assert adapter.calls == {} - - -def test_interleaved_instances_resolve_only_their_own_credentials_and_targets(monkeypatch: pytest.MonkeyPatch) -> None: - barrier = Barrier(2) - observed: list[tuple[str, tuple[str, str], str]] = [] - monkeypatch.setenv("WANDB_API_KEY", "ambient-key") - monkeypatch.setenv("WANDB_ENTITY", "ambient-entity") - monkeypatch.setenv("WANDB_PROJECT", "ambient-project") - before = dict(os.environ) - - def send(_method: str, url: str, **options: Unpack[RequestOptions]) -> httpx.Response: - label = ( - "a" - if url - in { - "https://control-a.test/prefix/graphql", - "https://trace-a.test/prefix/call/start", - "https://trace-a.test/prefix/call/end", - } - else "b" - ) - assert url in { - f"https://control-{label}.test/prefix/graphql", - f"https://trace-{label}.test/prefix/call/start", - f"https://trace-{label}.test/prefix/call/end", - } - assert options["auth"] == ("api", f"key-{label}") - if url.endswith("/graphql"): - return response(url, {"data": {"viewer": {"entity": f"team-{label}"}}}) - body = json.loads(options["content"]) - record = body.get("start", body.get("end")) - assert record["project_id"] == f"team-{label}/project-{label}" - observed.append((url, options["auth"], record["project_id"])) - if "start" in body: - barrier.wait(timeout=5) - return response(url, {"id": record["id"], "trace_id": record["trace_id"]}) - return response(url, {}) - - monkeypatch.setattr(ssrf_proxy, "make_request", send) - - def run(label: str) -> None: - adapter = WeaveDataTrace( - config( - api_key=f"key-{label}", - entity=None, - project=f"project-{label}", - host=f"https://control-{label}.test/prefix", - endpoint=f"https://trace-{label}.test/prefix/", - ) - ) - data = call(f"call-{label}") - adapter.start_call(data) - adapter.finish_call(data) - - with ThreadPoolExecutor(max_workers=2) as pool: - list(pool.map(run, ["a", "b"])) - assert len(observed) == 4 - assert dict(os.environ) == before - - -def test_concurrent_dispatches_on_one_cached_instance_keep_prefix_state_isolated( - monkeypatch: pytest.MonkeyPatch, -) -> None: - barrier = Barrier(2) - - def send(_method: str, url: str, **options: Unpack[RequestOptions]) -> httpx.Response: - body = json.loads(options["content"]) - if "start" in body: - start = CallStartReq.model_validate(body).start - if start.inputs["label"] == "unaccepted": - barrier.wait(timeout=5) - return response(url, {}, status=503) - return response(url, {"id": start.id, "trace_id": start.trace_id}) - CallEndReq.model_validate(body) - barrier.wait(timeout=5) - return response(url, {}, status=503) - - monkeypatch.setattr(ssrf_proxy, "make_request", send) - adapter = WeaveDataTrace(config()) - - def run(label: str) -> RetryableTraceDispatchError: - info = GenerateNameTraceInfo( - tenant_id="tenant", - inputs={"label": label}, - outputs={"label": label}, - metadata={}, - start_time=START, - end_time=START + timedelta(seconds=1), - ) - with pytest.raises(RetryableTraceDispatchError) as error: - adapter.trace(info) - return error.value - - with ThreadPoolExecutor(max_workers=2) as pool: - accepted, unaccepted = list(pool.map(run, ["accepted", "unaccepted"])) - assert isinstance(accepted, TraceExportUncertainError) - assert accepted.acceptance_unknown - assert not unaccepted.acceptance_unknown - assert adapter.calls == {} - - -@pytest.mark.parametrize( - ("stage", "failure"), - [ - (stage, failure) - for stage in ("start", "end") - for failure in ("rejected", "retryable", "malformed", "mismatched_id") - if not (stage == "end" and failure == "mismatched_id") - ], -) -def test_synchronous_failures_propagate_and_next_call_recovers( - monkeypatch: pytest.MonkeyPatch, stage: str, failure: str -) -> None: - broken = True - requests = [] - - def send(_method: str, url: str, **options: Unpack[RequestOptions]) -> httpx.Response: - body = json.loads(options["content"]) - requests.append(url) - if broken and url.endswith("/call/" + stage): - if failure == "rejected": - return response(url, {"error": "DO_NOT_EXPOSE_KEY_OR_BODY"}, status=401) - if failure == "retryable": - return response(url, {}, status=503) - if failure == "malformed": - return httpx.Response(200, request=httpx.Request("POST", url), content=b"invalid-json") - return response(url, {"id": "wrong-id", "trace_id": "wrong-trace"}) - if "start" in body: - return response(url, {"id": body["start"]["id"], "trace_id": body["start"]["trace_id"]}) - return response(url, {}) - - monkeypatch.setattr(ssrf_proxy, "make_request", send) - adapter = WeaveDataTrace(config()) - data = call() - expected = { - "rejected": TraceExportRejectedError, - "retryable": RetryableTraceDispatchError, - "malformed": TraceExportUncertainError, - "mismatched_id": TraceExportUncertainError, - }[failure] - if stage == "end": - adapter.start_call(data) - with pytest.raises(expected) as error: - adapter.start_call(data) if stage == "start" else adapter.finish_call(data) - assert "DO_NOT_EXPOSE" not in str(error.value) - assert len(requests) == (1 if stage == "start" else 2) - assert adapter.calls == {} - broken = False - recovered = call("next-call", exception="synthetic tool error") - adapter.start_call(recovered) - adapter.finish_call(recovered) - assert adapter.calls == {} - - -def test_legacy_api_check_uses_read_only_project_and_stats_requests(monkeypatch: pytest.MonkeyPatch) -> None: - sent = [] - - def post(url: str, **options: Unpack[RequestOptions]) -> httpx.Response: - assert options["auth"] == ("api", "fixture-key") - assert options["follow_redirects"] is False - sent.append(url) - if url == "https://api.wandb.ai/graphql": - return response( - url, - { - "data": { - "viewer": {"id": "viewer"}, - "project": {"id": "project-id", "name": "project", "entityName": "team", "readOnly": False}, - } - }, - ) - assert url == "https://trace.wandb.ai/calls/query_stats" - assert options["json"]["project_id"] == "team/project" - return response(url, {"count": 0}) - - monkeypatch.setattr(ssrf_proxy, "post", post) - adapter = WeaveDataTrace(config()) - assert sent == [] - assert adapter.api_check() - assert sent == ["https://api.wandb.ai/graphql", "https://trace.wandb.ai/calls/query_stats"] - assert adapter.calls == {} - - -@pytest.mark.parametrize( - ("sent_id", "ack_id"), - [ - ("1234567812345678123456781234abcd", "12345678-1234-5678-1234-56781234abcd"), - ("12345678-1234-5678-1234-56781234abcd", "1234567812345678123456781234ABCD"), - ("12345678-1234-5678-1234-56781234ABCD", "1234567812345678123456781234abcd"), - ], -) -def test_uuid_equivalent_start_ack_preserves_sent_ids( - monkeypatch: pytest.MonkeyPatch, sent_id: str, ack_id: str -) -> None: - sent: list[dict[str, object]] = [] - - def send(_method: str, url: str, **options: Unpack[RequestOptions]) -> httpx.Response: - body = json.loads(options["content"]) - sent.append(body) - if "start" in body: - start = CallStartReq.model_validate(body).start - assert start.id == sent_id - assert start.trace_id == sent_id - return response(url, {"id": ack_id, "trace_id": ack_id}) - end = CallEndReq.model_validate(body).end - assert end.id == sent_id - return response(url, {}) - - monkeypatch.setattr(ssrf_proxy, "make_request", send) - adapter = WeaveDataTrace(config()) - data = call(sent_id) - adapter.start_call(data) - assert adapter.calls[sent_id]["trace_id"] == sent_id - adapter.finish_call(data) - assert len(sent) == 2 - assert adapter.calls == {} - - -@pytest.mark.parametrize("field", ["id", "trace_id"]) -@pytest.mark.parametrize( - ("sent_id", "wrong_ack"), - [ - ("1234567812345678123456781234abcd", "12345678-1234-5678-1234-56781234abce"), - ("legacy-correlation", "LEGACY-CORRELATION"), - ("1234567812345678123456781234abcd", "not-a-uuid"), - ], -) -def test_start_ack_rejects_different_uuid_or_opaque_identifiers( - monkeypatch: pytest.MonkeyPatch, field: str, sent_id: str, wrong_ack: str -) -> None: - sent = [] - - def send(_method: str, url: str, **options: Unpack[RequestOptions]) -> httpx.Response: - start = CallStartReq.model_validate_json(options["content"]).start - sent.append(start) - ack = {"id": start.id, "trace_id": start.trace_id} - ack[field] = wrong_ack - return response(url, ack) - - monkeypatch.setattr(ssrf_proxy, "make_request", send) - adapter = WeaveDataTrace(config()) - with pytest.raises(TraceExportUncertainError, match="did not acknowledge the submitted identifiers"): - adapter.start_call(call(sent_id)) - assert len(sent) == 1 - assert adapter.calls == {} diff --git a/api/tests/unit_tests/services/tracing/test_native_weave_mlflow.py b/api/tests/unit_tests/services/tracing/test_native_weave_mlflow.py deleted file mode 100644 index 14a7b9964ad..00000000000 --- a/api/tests/unit_tests/services/tracing/test_native_weave_mlflow.py +++ /dev/null @@ -1,516 +0,0 @@ -"""Real SDK schema/serialization boundaries with HTTP transport substituted. - -These tests do not claim acceptance by a live third-party account. -""" - -import base64 -import json -import os -from collections.abc import Mapping -from copy import deepcopy -from datetime import UTC, datetime, timedelta -from typing import TypedDict, Unpack -from urllib.parse import parse_qs, urlsplit - -import httpx -import pytest -from mlflow.entities import Span, TraceData -from mlflow.tracing.constant import SpanAttributeKey -from mlflow.tracing.otel.translation import translate_span_when_storing -from mlflow.types.chat import ChatTool -from opentelemetry.proto.collector.trace.v1.trace_service_pb2 import ( - ExportTraceServiceRequest, - ExportTraceServiceResponse, -) -from weave.trace_server.trace_server_interface import StartedCallSchemaForInsert - -from core.helper import ssrf_proxy -from core.ops.unified_trace.entities import CanonicalSpan, CanonicalSpanKind, CanonicalSpanStatus, CanonicalTrace -from services.tracing.normalization import remap_outer_trace -from services.tracing.provider_dispatch import TraceExportUncertainError, export_trace - - -class TransportOptions(TypedDict, total=False): - content: bytes | str - headers: Mapping[str, str] - auth: tuple[str, str] - max_retries: int - - -@pytest.fixture -def trace() -> CanonicalTrace: - start = datetime(2026, 1, 1, tzinfo=UTC) - root = CanonicalSpan( - id="0000000000000001", - parent_id=None, - name="Agent", - kind=CanonicalSpanKind.AGENT, - start_time=start, - end_time=start + timedelta(seconds=2), - inputs={"query": "hello"}, - outputs="answer", - status=CanonicalSpanStatus.OK, - metadata={"total_price": 999, "total_tokens": 999}, - ) - child = CanonicalSpan( - id="0000000000000002", - parent_id=root.id, - name="Model", - kind=CanonicalSpanKind.LLM, - start_time=start, - end_time=start + timedelta(seconds=1), - inputs={"prompt": "hello"}, - outputs="answer", - status=CanonicalSpanStatus.OK, - metadata={ - "model_name": "test-model", - "prompt_tokens": 2, - "completion_tokens": 3, - "total_tokens": 5, - "total_price": "0.01", - "currency": "USD", - }, - ) - return CanonicalTrace( - trace_id="1234567890abcdef1234567890abcdef", session_id="session-1", root_span_id=root.id, spans=(root, child) - ) - - -def _response(url: str, *, body: object = None, content: bytes | None = None, status: int = 200) -> httpx.Response: - request = httpx.Request("POST", url) - if content is not None: - return httpx.Response(status, request=request, content=content) - return httpx.Response(status, request=request, json=body or {}) - - -@pytest.fixture -def weave_started_calls(monkeypatch: pytest.MonkeyPatch) -> list[StartedCallSchemaForInsert]: - starts: list[StartedCallSchemaForInsert] = [] - - def send(_method: str, url: str, **kwargs: Unpack[TransportOptions]) -> httpx.Response: - body = json.loads(kwargs["content"]) - if start := body.get("start"): - started = StartedCallSchemaForInsert.model_validate(start) - starts.append(started) - return _response(url, body={"id": started.id, "trace_id": started.trace_id}) - return _response(url) - - monkeypatch.setattr(ssrf_proxy, "make_request", send) - return starts - - -@pytest.mark.parametrize("host", ["https://api.wandb.ai", "https://API.WANDB.AI:443"]) -def test_weave_sdk_exports_tree_session_usage_without_global_credentials( - monkeypatch: pytest.MonkeyPatch, trace: CanonicalTrace, host: str -) -> None: - calls = [] - before = dict(os.environ) - - def send(_method: str, url: str, **kwargs: Unpack[TransportOptions]) -> httpx.Response: - body = json.loads(kwargs["content"]) - calls.append((_method, url, kwargs, body)) - start = body.get("start") - return _response(url, body={"id": start["id"], "trace_id": start["trace_id"]} if start else {}) - - monkeypatch.setattr(ssrf_proxy, "make_request", send) - result = export_trace("weave", {"api_key": "tenant-key", "entity": "team", "project": "demo", "host": host}, trace) - assert len(calls) == 4 - root, child = calls[0][3]["start"], calls[1][3]["start"] - assert child["parent_id"] == root["id"] - assert child["thread_id"] == root["thread_id"] == "session-1" - assert child["turn_id"] == root["turn_id"] == root["id"] - assert calls[2][3]["end"]["summary"]["usage"]["test-model"]["total_tokens"] == 5 - assert calls[3][3]["end"]["summary"] == {} - assert result.trace_url == f"https://wandb.ai/team/demo/r/call/{root['id']}" - assert result.span_urls[trace.spans[1].id].endswith(child["id"]) - assert all(call[2]["auth"] == ("api", "tenant-key") and call[2]["max_retries"] == 0 for call in calls) - assert dict(os.environ) == before - - -@pytest.mark.parametrize("root_kind", [CanonicalSpanKind.CHAIN, CanonicalSpanKind.AGENT]) -def test_weave_nested_agent_calls_form_one_outer_turn( - trace: CanonicalTrace, weave_started_calls: list[StartedCallSchemaForInsert], root_kind: CanonicalSpanKind -) -> None: - root = trace.spans[0].model_copy(update={"kind": root_kind}) - agent = trace.spans[1].model_copy(update={"kind": CanonicalSpanKind.AGENT}) - model = trace.spans[1].model_copy(update={"id": "0000000000000003", "parent_id": agent.id}) - nested = trace.model_copy(update={"spans": (root, agent, model)}) - - export_trace("weave", {"api_key": "tenant-key", "entity": "team", "project": "demo"}, nested) - - outer, inner_agent, inner_model = weave_started_calls - assert inner_agent.parent_id == outer.id - assert inner_model.parent_id == inner_agent.id - assert all(call.thread_id == trace.session_id and call.turn_id == outer.id for call in weave_started_calls) - # Native Threads queries count only call.id == call.turn_id, including for - # Workflow and single-step preview outer roots with nested Agent details. - assert [call.id for call in weave_started_calls if call.id == call.turn_id] == [outer.id] - - -def test_weave_same_session_has_distinct_stable_turns( - trace: CanonicalTrace, weave_started_calls: list[StartedCallSchemaForInsert] -) -> None: - config = {"api_key": "tenant-key", "entity": "team", "project": "demo"} - first = export_trace("weave", config, trace) - repeated = export_trace("weave", config, trace) - next_trace = trace.model_copy(update={"trace_id": "abcdef1234567890abcdef1234567890"}) - second = export_trace("weave", config, next_trace) - - first_calls, repeated_calls, second_calls = ( - weave_started_calls[:2], - weave_started_calls[2:4], - weave_started_calls[4:], - ) - assert repeated_calls == first_calls - assert repeated == first - assert second.trace_url != first.trace_url - assert all(call.thread_id == trace.session_id for call in weave_started_calls) - assert first_calls[0].turn_id == first_calls[0].id - assert second_calls[0].turn_id == second_calls[0].id - assert second_calls[0].turn_id != first_calls[0].turn_id - assert all(call.turn_id == first_calls[0].id for call in first_calls) - assert all(call.turn_id == second_calls[0].id for call in second_calls) - - -def test_weave_without_session_does_not_fabricate_a_thread_or_turn( - trace: CanonicalTrace, weave_started_calls: list[StartedCallSchemaForInsert] -) -> None: - trace = trace.model_copy(update={"session_id": ""}) - - export_trace("weave", {"api_key": "tenant-key", "entity": "team", "project": "demo"}, trace) - - assert len(weave_started_calls) == len(trace.spans) - assert all(call.thread_id is None and call.turn_id is None for call in weave_started_calls) - - -def test_weave_resolves_optional_entity_using_only_supplied_credentials( - monkeypatch: pytest.MonkeyPatch, trace: CanonicalTrace -) -> None: - def send(_method: str, url: str, **kwargs: Unpack[TransportOptions]) -> httpx.Response: - assert kwargs["auth"] == ("api", "tenant-key") - if url.endswith("/graphql"): - return _response(url, body={"data": {"viewer": {"entity": "default-team"}}}) - body = json.loads(kwargs["content"]) - if start := body.get("start"): - assert start["project_id"] == "default-team/demo" - return _response(url, body={"id": start["id"], "trace_id": start["trace_id"]}) - return _response(url) - - monkeypatch.setattr(ssrf_proxy, "make_request", send) - result = export_trace("weave", {"api_key": "tenant-key", "project": "demo"}, trace) - assert result.trace_url is not None - assert "/default-team/demo/" in result.trace_url - - -def test_weave_partial_tree_is_uncertain_and_not_retried( - monkeypatch: pytest.MonkeyPatch, trace: CanonicalTrace -) -> None: - calls = [] - - def send(_method: str, url: str, **kwargs: Unpack[TransportOptions]) -> httpx.Response: - calls.append(url) - if len(calls) > 1: - return _response(url, status=503) - start = json.loads(kwargs["content"])["start"] - return _response(url, body={"id": start["id"], "trace_id": start["trace_id"]}) - - monkeypatch.setattr(ssrf_proxy, "make_request", send) - with pytest.raises(TraceExportUncertainError, match="partially accepted"): - export_trace("weave", {"api_key": "secret", "entity": "team", "project": "demo"}, trace) - assert len(calls) == 2 - - -@pytest.mark.parametrize("tracking_uri", ["https://mlflow.example", "https://mlflow.example/proxy/mlflow"]) -def test_mlflow_sdk_span_roundtrip_and_credential_isolation( - monkeypatch: pytest.MonkeyPatch, trace: CanonicalTrace, tracking_uri: str -) -> None: - calls = [] - before = dict(os.environ) - - def send(_method: str, url: str, **kwargs: Unpack[TransportOptions]) -> httpx.Response: - calls.append((url, kwargs)) - return _response(url, content=b"") - - monkeypatch.setattr(ssrf_proxy, "make_request", send) - for username in ("tenant-a", "tenant-b"): - result = export_trace( - "mlflow", - { - "tracking_uri": tracking_uri, - "experiment_id": "42", - "username": username, - "password": "password", - }, - trace, - ) - assert result.trace_id == "tr-" + trace.trace_id - assert result.trace_url is not None - fragment_path, _, query = urlsplit(result.trace_url).fragment.partition("?") - assert result.trace_url.startswith(tracking_uri + "/#/") - assert fragment_path == "/experiments/42/traces" - assert parse_qs(query) == { - "traceId": [result.trace_id], - "selectedTraceId": [result.trace_id], - "selectedEvaluationId": [result.trace_id], - } - for index, (url, options) in enumerate(calls): - assert url == tracking_uri + "/v1/traces" - assert options["headers"]["x-mlflow-experiment-id"] == "42" - assert ( - options["headers"]["Authorization"] - == "Basic " + base64.b64encode(f"tenant-{'a' if index == 0 else 'b'}:password".encode()).decode() - ) - content = options["content"] - assert isinstance(content, bytes) - payload = ExportTraceServiceRequest.FromString(content) - spans = [Span.from_otel_proto(span) for span in payload.resource_spans[0].scope_spans[0].spans] - assert spans[1].parent_id == spans[0].span_id - assert spans[1].get_attribute(SpanAttributeKey.CHAT_USAGE) == { - "input_tokens": 2, - "output_tokens": 3, - "total_tokens": 5, - } - assert spans[1].get_attribute(SpanAttributeKey.LLM_COST) == {"total_cost": 0.01} - assert spans[0].get_attribute(SpanAttributeKey.CHAT_USAGE) is None - assert spans[0].get_attribute(SpanAttributeKey.SESSION_ID) == "session-1" - assert dict(os.environ) == before - - -def _tool_definitions() -> list[dict[str, object]]: - return [ - { - "type": "function", - "name": "calculate", - "description": "TOOL_CONTENT_CANARY", - "parameters": { - "type": "object", - "properties": {"value": {"type": "integer"}}, - "required": ["value"], - "additionalProperties": False, - "$comment": "Preserve schema extension fields", - }, - "strict": True, - }, - {"type": "function", "function": {"name": "already_nested", "description": "Nested definition"}}, - ] - - -@pytest.mark.parametrize("serialized", [False, True]) -def test_mlflow_tool_definitions_have_native_chat_schema_without_losing_source( - monkeypatch: pytest.MonkeyPatch, trace: CanonicalTrace, serialized: bool -) -> None: - definitions = _tool_definitions() - source = json.dumps(definitions) if serialized else definitions - trace.spans[1].metadata["gen_ai.tool.definitions"] = source - before = deepcopy(trace.model_dump()) - native: list[Span] = [] - - def send(_method: str, url: str, **options: Unpack[TransportOptions]) -> httpx.Response: - content = options["content"] - assert isinstance(content, bytes) - payload = ExportTraceServiceRequest.FromString(content) - native.extend(Span.from_otel_proto(item) for item in payload.resource_spans[0].scope_spans[0].spans) - return _response(url, content=b"") - - monkeypatch.setattr(ssrf_proxy, "make_request", send) - export_trace("mlflow", {"tracking_uri": "https://mlflow.example", "experiment_id": "42"}, trace) - tools = native[1].get_attribute(SpanAttributeKey.CHAT_TOOLS) - assert tools == [ - {"type": "function", "function": {key: value for key, value in definitions[0].items() if key != "type"}}, - definitions[1], - ] - assert isinstance(tools, list) - for tool in tools: - ChatTool.model_validate(tool, strict=True) - assert native[1].get_attribute("gen_ai.tool.definitions") == source - stored = translate_span_when_storing(native[1]) - assert json.loads(stored["attributes"][SpanAttributeKey.CHAT_TOOLS]) == tools - assert trace.model_dump() == before - - -def test_mlflow_redacted_trace_does_not_recreate_tool_definition_content( - monkeypatch: pytest.MonkeyPatch, trace: CanonicalTrace -) -> None: - trace.spans[1].metadata["gen_ai.tool.definitions"] = _tool_definitions() - trace.spans[1].metadata[SpanAttributeKey.CHAT_TOOLS] = [ - {"type": "function", "function": {"name": "TOOL_CONTENT_CANARY"}} - ] - redacted = remap_outer_trace(trace, trace_id=trace.trace_id, capture_content=False, agent_message=False) - native: list[Span] = [] - - def send(_method: str, url: str, **options: Unpack[TransportOptions]) -> httpx.Response: - content = options["content"] - assert isinstance(content, bytes) - assert b"TOOL_CONTENT_CANARY" not in content - payload = ExportTraceServiceRequest.FromString(content) - native.extend(Span.from_otel_proto(item) for item in payload.resource_spans[0].scope_spans[0].spans) - return _response(url, content=b"") - - monkeypatch.setattr(ssrf_proxy, "make_request", send) - export_trace("mlflow", {"tracking_uri": "https://mlflow.example", "experiment_id": "42"}, redacted) - for span in native: - assert span.get_attribute(SpanAttributeKey.CHAT_TOOLS) is None - assert span.get_attribute("gen_ai.tool.definitions") is None - assert span.inputs is None - assert span.outputs is None - - -@pytest.mark.parametrize("definitions", [[None, 12, {"type": "other"}, {"type": "function"}], "invalid-json"]) -def test_mlflow_unsupported_tool_shapes_remain_in_metadata_without_fake_functions( - trace: CanonicalTrace, definitions: object -) -> None: - from dify_trace_mlflow.agent_trace import _spans - - trace.spans[1].metadata["gen_ai.tool.definitions"] = definitions - before = deepcopy(trace.model_dump()) - result = _spans(trace, "42", "tr-" + trace.trace_id) - assert result[1].get_attribute("gen_ai.tool.definitions") == definitions - assert result[1].get_attribute(SpanAttributeKey.CHAT_TOOLS) == ([] if isinstance(definitions, list) else None) - assert trace.model_dump() == before - - -def test_mlflow_existing_native_chat_tools_take_precedence(trace: CanonicalTrace) -> None: - from dify_trace_mlflow.agent_trace import _spans - - native_tools = [{"type": "function", "function": {"name": "already_native"}}] - trace.spans[1].metadata[SpanAttributeKey.CHAT_TOOLS] = native_tools - trace.spans[1].metadata["gen_ai.tool.definitions"] = _tool_definitions() - result = _spans(trace, "42", "tr-" + trace.trace_id) - assert result[1].get_attribute(SpanAttributeKey.CHAT_TOOLS) == native_tools - - -def test_mlflow_preserves_json_schema_beyond_the_sdk_chat_tool_model(trace: CanonicalTrace) -> None: - from dify_trace_mlflow.agent_trace import _spans - - parameters = { - "$ref": "#/$defs/Arguments", - "$defs": { - "Arguments": { - "type": "object", - "properties": {"choice": {"enum": [None, {"mode": "auto"}]}}, - "additionalProperties": {"type": "string"}, - } - }, - } - source = [{"type": "function", "name": "schema_extension", "parameters": parameters}] - trace.spans[1].metadata["gen_ai.tool.definitions"] = source - result = _spans(trace, "42", "tr-" + trace.trace_id) - expected = [{"type": "function", "function": {"name": "schema_extension", "parameters": parameters}}] - assert result[1].get_attribute(SpanAttributeKey.CHAT_TOOLS) == expected - assert result[1].get_attribute("gen_ai.tool.definitions") == source - stored = translate_span_when_storing(result[1]) - assert json.loads(stored["attributes"][SpanAttributeKey.CHAT_TOOLS]) == expected - - -def test_mlflow_partial_otlp_ack_is_uncertain(monkeypatch: pytest.MonkeyPatch, trace: CanonicalTrace) -> None: - ack = ExportTraceServiceResponse() - ack.partial_success.rejected_spans = 1 - monkeypatch.setattr( - ssrf_proxy, "make_request", lambda _method, url, **_kwargs: _response(url, content=ack.SerializeToString()) - ) - with pytest.raises(TraceExportUncertainError, match="partially accepted"): - export_trace("mlflow", {"tracking_uri": "https://mlflow.example", "experiment_id": "42"}, trace) - - -@pytest.mark.parametrize( - "credential_type", ["AWS_PRESIGNED_URL", "GCP_SIGNED_URL", "AZURE_SAS_URI", "AZURE_ADLS_GEN2_SAS_URI"] -) -@pytest.mark.parametrize("oauth", [False, True]) -def test_databricks_sdk_v3_metadata_and_native_artifact( - monkeypatch: pytest.MonkeyPatch, trace: CanonicalTrace, credential_type: str, oauth: bool -) -> None: - calls: list[tuple[str, str, TransportOptions]] = [] - uploaded: list[TraceData] = [] - before = dict(os.environ) - server_trace_id = "tr-abcdef0123456789abcdef0123456789" - definitions = _tool_definitions() - trace.spans[1].metadata["gen_ai.tool.definitions"] = definitions - - def send(_method: str, url: str, **kwargs: Unpack[TransportOptions]) -> httpx.Response: - calls.append((_method, url, kwargs)) - if url.endswith("/oidc/v1/token"): - assert kwargs["auth"] == ("client", "client-secret") - return _response(url, body={"access_token": "oauth-token"}) - if url.endswith("/api/3.0/mlflow/traces"): - assert kwargs["headers"]["Authorization"] == ("Bearer oauth-token" if oauth else "Bearer pat-token") - body = json.loads(kwargs["content"]) - assert body["trace"]["trace_info"]["trace_metadata"]["mlflow.trace.session"] == "session-1" - body["trace"]["trace_info"]["trace_id"] = server_trace_id - return _response(url, body=body) - if url.endswith("/credentials-for-data-upload"): - assert f"/traces/{server_trace_id}/" in url - return _response( - url, - body={ - "credential_info": { - "signed_uri": "https://storage.example/trace?sig=fake", - "type": credential_type, - "headers": [], - } - }, - ) - assert "Authorization" not in kwargs["headers"] - if kwargs.get("content"): - uploaded.append(TraceData.from_dict(json.loads(kwargs["content"]))) - return _response(url) - - monkeypatch.setattr(ssrf_proxy, "make_request", send) - config = {"host": "https://databricks.example", "experiment_id": "42"} - config.update( - {"client_id": "client", "client_secret": "client-secret"} if oauth else {"personal_access_token": "pat-token"} - ) - result = export_trace("databricks", config, trace) - assert result.trace_id == server_trace_id - assert result.trace_url == ( - f"https://databricks.example/ml/experiments/42/traces?selectedEvaluationId={server_trace_id}" - ) - data = uploaded[0] - assert len(data.spans) == 2 - assert all(span.trace_id == server_trace_id for span in data.spans) - assert data.spans[1].parent_id == data.spans[0].span_id - usage = data.spans[1].get_attribute(SpanAttributeKey.CHAT_USAGE) - assert usage is not None - assert usage["total_tokens"] == 5 - assert data.spans[1].get_attribute(SpanAttributeKey.LLM_COST) == {"total_cost": 0.01} - tools = data.spans[1].get_attribute(SpanAttributeKey.CHAT_TOOLS) - assert isinstance(tools, list) - assert len(tools) == 2 - for tool in tools: - ChatTool.model_validate(tool, strict=True) - assert data.spans[1].get_attribute("gen_ai.tool.definitions") == definitions - assert dict(os.environ) == before - assert len(calls) == 3 + int(oauth) + int(credential_type == "AZURE_ADLS_GEN2_SAS_URI") - - -def test_databricks_partial_metadata_delivery_is_uncertain( - monkeypatch: pytest.MonkeyPatch, trace: CanonicalTrace -) -> None: - def send(_method: str, url: str, **kwargs: Unpack[TransportOptions]) -> httpx.Response: - if url.endswith("/api/3.0/mlflow/traces"): - return _response(url, body=json.loads(kwargs["content"])) - return _response(url, status=503) - - monkeypatch.setattr(ssrf_proxy, "make_request", send) - with pytest.raises(TraceExportUncertainError, match="metadata"): - export_trace( - "databricks", - {"host": "https://databricks.example", "experiment_id": "42", "personal_access_token": "secret"}, - trace, - ) - - -@pytest.mark.parametrize("provider", ["weave", "databricks"]) -def test_success_status_with_unreadable_ack_is_uncertain( - monkeypatch: pytest.MonkeyPatch, trace: CanonicalTrace, provider: str -) -> None: - monkeypatch.setattr(ssrf_proxy, "make_request", lambda _method, url, **_kwargs: _response(url, content=b"not-json")) - config = { - "api_key": "secret", - "entity": "team", - "project": "demo", - "host": "https://example.com", - "experiment_id": "42", - "personal_access_token": "secret", - } - with pytest.raises(TraceExportUncertainError, match="unreadable"): - export_trace(provider, config, trace) diff --git a/api/tests/unit_tests/services/tracing/test_application_trace_normalization.py b/api/tests/unit_tests/services/tracing/test_normalization.py similarity index 100% rename from api/tests/unit_tests/services/tracing/test_application_trace_normalization.py rename to api/tests/unit_tests/services/tracing/test_normalization.py diff --git a/api/tests/unit_tests/services/tracing/test_nullable_tracing_credentials.py b/api/tests/unit_tests/services/tracing/test_nullable_tracing_credentials.py deleted file mode 100644 index fa8d4215b31..00000000000 --- a/api/tests/unit_tests/services/tracing/test_nullable_tracing_credentials.py +++ /dev/null @@ -1,92 +0,0 @@ -"""Nullable provider credentials never reach RSA/base64 as non-string tokens.""" - -from uuid import uuid4 - -import pytest -from Crypto.PublicKey import RSA -from pydantic import ValidationError -from sqlalchemy.orm import Session - -import models.engine as model_engine -from core.ops.ops_trace_manager import OpsTraceManager -from extensions.ext_key_provider import key_provider_manager -from libs import gmpy2_pkcs10aep_cipher, rsa -from libs.key_providers.rsa_key_provider import RSAKeyProvider -from tests.unit_tests.model_factories import make_tenant - - -@pytest.fixture -def workspace(sqlite_session: Session, monkeypatch: pytest.MonkeyPatch) -> str: - key = RSA.generate(2048) - tenant_id = str(uuid4()) - sqlite_session.add(make_tenant(tenant_id=tenant_id, encrypt_public_key=key.publickey().export_key().decode())) - sqlite_session.commit() - monkeypatch.setattr(model_engine.db, "session", sqlite_session) - monkeypatch.setattr(key_provider_manager, "_provider", RSAKeyProvider()) - monkeypatch.setattr(rsa, "get_decrypt_decoding", lambda _tenant: (key, gmpy2_pkcs10aep_cipher.new(key))) - return tenant_id - - -def test_databricks_switches_auth_with_real_encryption_and_explicit_null(workspace: str) -> None: - pat = OpsTraceManager.encrypt_tracing_config( - workspace, - "databricks", - { - "host": "https://workspace.test", - "experiment_id": "12", - "personal_access_token": "synthetic-pat", - "client_id": None, - "client_secret": None, - }, - ) - assert pat["personal_access_token"] != "synthetic-pat" - decrypted_pat = OpsTraceManager.decrypt_tracing_config(workspace, "databricks", pat) - assert decrypted_pat["personal_access_token"] == "synthetic-pat" - assert decrypted_pat["client_secret"] is None - - oauth = OpsTraceManager.encrypt_tracing_config( - workspace, - "databricks", - { - "host": "https://workspace.test", - "experiment_id": "12", - "personal_access_token": None, - "client_id": "synthetic-client", - "client_secret": "synthetic-secret", - }, - pat, - ) - assert oauth["personal_access_token"] is None - assert oauth["client_secret"] != "synthetic-secret" - decrypted_oauth = OpsTraceManager.decrypt_tracing_config(workspace, "databricks", oauth) - assert decrypted_oauth["personal_access_token"] is None - assert decrypted_oauth["client_secret"] == "synthetic-secret" - - -def test_other_optional_secret_model_can_be_cleared_without_losing_validation(workspace: str) -> None: - password = OpsTraceManager.encrypt_tracing_config( - workspace, - "mlflow", - { - "tracking_uri": "https://mlflow.test", - "experiment_id": "0", - "username": "user", - "password": "synthetic-password", - }, - ) - assert OpsTraceManager.decrypt_tracing_config(workspace, "mlflow", password)["password"] == "synthetic-password" - cleared = OpsTraceManager.encrypt_tracing_config(workspace, "mlflow", {**password, "password": None}, password) - assert cleared["password"] is None - assert OpsTraceManager.decrypt_tracing_config(workspace, "mlflow", cleared)["password"] is None - - -def test_non_nullable_secret_is_still_rejected_by_provider_model(workspace: str) -> None: - with pytest.raises(ValidationError, match="secret_key"): - OpsTraceManager.encrypt_tracing_config( - workspace, "langfuse", {"public_key": "synthetic-public", "secret_key": None} - ) - encrypted = OpsTraceManager.encrypt_tracing_config( - workspace, "langfuse", {"public_key": "synthetic-public", "secret_key": "synthetic-secret"} - ) - with pytest.raises(ValidationError, match="secret_key"): - OpsTraceManager.decrypt_tracing_config(workspace, "langfuse", {**encrypted, "secret_key": None}) diff --git a/api/tests/unit_tests/services/tracing/test_aliyun_otlp_ack.py b/api/tests/unit_tests/services/tracing/test_otlp_export.py similarity index 100% rename from api/tests/unit_tests/services/tracing/test_aliyun_otlp_ack.py rename to api/tests/unit_tests/services/tracing/test_otlp_export.py diff --git a/api/tests/unit_tests/services/tracing/test_phoenix_export_auth.py b/api/tests/unit_tests/services/tracing/test_phoenix_export_auth.py deleted file mode 100644 index ed5f8c1d3ac..00000000000 --- a/api/tests/unit_tests/services/tracing/test_phoenix_export_auth.py +++ /dev/null @@ -1,142 +0,0 @@ -"""Exercise anonymous Phoenix auth at the HTTP/1.1 and legacy SDK boundaries.""" - -from collections.abc import Iterator -from datetime import UTC, datetime, timedelta - -import httpcore -import httpx -import pytest -import requests -from dify_trace_arize_phoenix.arize_phoenix_trace import setup_tracer -from dify_trace_arize_phoenix.config import PhoenixConfig -from opentelemetry.proto.collector.trace.v1.trace_service_pb2 import ExportTraceServiceRequest - -from core.helper import ssrf_proxy -from core.ops.unified_trace.entities import CanonicalSpan, CanonicalSpanKind, CanonicalSpanStatus, CanonicalTrace -from services.tracing.provider_dispatch import export_trace - - -@pytest.fixture -def trace() -> CanonicalTrace: - start = datetime(2026, 10, 20, tzinfo=UTC) - return CanonicalTrace( - trace_id="1" * 32, - session_id="phoenix-auth-session", - root_span_id="2" * 16, - spans=( - CanonicalSpan( - id="2" * 16, - parent_id=None, - name="Phoenix auth regression", - kind=CanonicalSpanKind.AGENT, - start_time=start, - end_time=start + timedelta(seconds=1), - status=CanonicalSpanStatus.OK, - ), - ), - ) - - -@pytest.fixture -def http11_wire(monkeypatch: pytest.MonkeyPatch) -> Iterator[list[bytes]]: - """Keep real HTTPX/httpcore/h11 serialization, replacing only socket I/O.""" - chunks: list[bytes] = [] - - def write(_stream: httpcore.MockStream, buffer: bytes, timeout: float | None = None) -> None: - del timeout - chunks.append(buffer) - - monkeypatch.setattr(httpcore.MockStream, "write", write) - transport = httpx.HTTPTransport() - transport.close() - transport._pool = httpcore.ConnectionPool( - network_backend=httpcore.MockBackend([b"HTTP/1.1 200 OK\r\nContent-Length: 0\r\nConnection: close\r\n\r\n"]) - ) - with httpx.Client(transport=transport, trust_env=False) as client: - monkeypatch.setattr(ssrf_proxy, "_get_ssrf_client", lambda _verify: client) - yield chunks - - -def test_managed_exports_serialize_optional_auth_without_cross_destination_credentials( - monkeypatch: pytest.MonkeyPatch, trace: CanonicalTrace, http11_wire: list[bytes] -) -> None: - monkeypatch.setenv("OTEL_EXPORTER_OTLP_HEADERS", "authorization=Bearer%20ambient-general,api_key=ambient-general") - monkeypatch.setenv( - "OTEL_EXPORTER_OTLP_TRACES_HEADERS", "authorization=Bearer%20ambient-traces,api_key=ambient-traces" - ) - - for index, key in enumerate(("first-synthetic-key", "", None, "second-synthetic-key", None, "")): - previous_chunks = len(http11_wire) - result = export_trace( - "phoenix", {"endpoint": "http://phoenix.test/proxy", "api_key": key, "project": f"project-{index}"}, trace - ) - assert result.status == "accepted" - assert result.trace_id == trace.trace_id - head, separator, body = b"".join(http11_wire[previous_chunks:]).partition(b"\r\n\r\n") - assert separator - lines = head.split(b"\r\n") - assert lines[0] == b"POST /proxy/v1/traces HTTP/1.1" - header_values: list[tuple[bytes, bytes]] = [] - for line in lines[1:]: - name, value = line.split(b": ", 1) - header_values.append((name, value)) - headers = httpx.Headers(header_values) - assert headers["Content-Type"] == "application/x-protobuf" - if key: - assert headers["api_key"] == key - assert headers["authorization"] == f"Bearer {key}" - else: - assert "api_key" not in headers - assert "authorization" not in headers - payload = ExportTraceServiceRequest.FromString(body) - resource = payload.resource_spans[0] - assert {item.key: item.value.string_value for item in resource.resource.attributes}[ - "openinference.project.name" - ] == f"project-{index}" - assert resource.scope_spans[0].spans[0].trace_id.hex() == trace.trace_id - - -def test_legacy_sdk_prepares_anonymous_requests_without_ambient_or_previous_auth( - monkeypatch: pytest.MonkeyPatch, -) -> None: - monkeypatch.setenv("OTEL_EXPORTER_OTLP_HEADERS", "authorization=Bearer%20ambient-general,api_key=ambient-general") - monkeypatch.setenv( - "OTEL_EXPORTER_OTLP_TRACES_HEADERS", "authorization=Bearer%20ambient-traces,api_key=ambient-traces" - ) - monkeypatch.delenv("OTEL_EXPORTER_OTLP_COMPRESSION", raising=False) - monkeypatch.delenv("OTEL_EXPORTER_OTLP_TRACES_COMPRESSION", raising=False) - prepared: list[requests.PreparedRequest] = [] - - def send(_session: requests.Session, request: requests.PreparedRequest, **_kwargs: object) -> requests.Response: - prepared.append(request) - response = requests.Response() - response.status_code = 200 - response._content = b"" - response.request = request - return response - - monkeypatch.setattr(requests.Session, "send", send) - for index, key in enumerate(("first-synthetic-key", "", None, "second-synthetic-key", None, "")): - tracer, processor = setup_tracer( - PhoenixConfig(endpoint="http://phoenix.test/proxy", api_key=key, project=f"project-{index}") - ) - try: - tracer.start_span("Legacy Phoenix auth regression").end() - finally: - processor.shutdown() - assert len(prepared) == index + 1 - request = prepared[-1] - assert request.method == "POST" - assert request.url == "http://phoenix.test/proxy/v1/traces" - assert request.headers["Content-Type"] == "application/x-protobuf" - if key: - assert request.headers["api_key"] == key - assert request.headers["authorization"] == f"Bearer {key}" - else: - assert "api_key" not in request.headers - assert "authorization" not in request.headers - assert isinstance(request.body, bytes) - payload = ExportTraceServiceRequest.FromString(request.body) - assert {item.key: item.value.string_value for item in payload.resource_spans[0].resource.attributes}[ - "openinference.project.name" - ] == f"project-{index}" diff --git a/api/tests/unit_tests/services/tracing/test_provider_dispatch.py b/api/tests/unit_tests/services/tracing/test_provider_dispatch.py index 32ffd59b8e7..8cf3a86cb92 100644 --- a/api/tests/unit_tests/services/tracing/test_provider_dispatch.py +++ b/api/tests/unit_tests/services/tracing/test_provider_dispatch.py @@ -1,28 +1,2133 @@ -"""Exercise real vendor serializers and OTLP stubs at their network boundary.""" - +import atexit +import base64 import json -from collections.abc import Callable, Mapping, Sequence +import os +import threading +from collections.abc import Callable, Iterator, Mapping, Sequence +from concurrent.futures import ThreadPoolExecutor +from copy import deepcopy +from dataclasses import dataclass from datetime import UTC, datetime, timedelta -from typing import Never, NotRequired, TypedDict, cast +from decimal import Decimal +from pathlib import Path +from threading import Barrier, Thread +from typing import Never, NotRequired, TypedDict, Unpack, cast +from urllib.parse import parse_qs, urlsplit from uuid import UUID +import httpcore import httpx import pytest +import requests +import wandb +import weave +from aliyun.log.logclient import LogClient +from dify_trace_aliyun.aliyun_trace import AliyunDataTrace +from dify_trace_aliyun.config import AliyunConfig +from dify_trace_aliyun.data_exporter.traceclient import TraceClient, convert_to_span_id +from dify_trace_aliyun.entities.aliyun_trace_entity import SpanData +from dify_trace_aliyun.entities.semconv import GEN_AI_USAGE_TOTAL_TOKENS +from dify_trace_arize_phoenix.arize_phoenix_trace import setup_tracer +from dify_trace_arize_phoenix.config import PhoenixConfig +from dify_trace_weave.config import WeaveConfig +from dify_trace_weave.entities.weave_trace_entity import WeaveTraceModel +from dify_trace_weave.weave_trace import WeaveDataTrace +from mlflow.entities import Span as _native_weave_mlflow_Span +from mlflow.entities import TraceData +from mlflow.tracing.constant import SpanAttributeKey +from mlflow.tracing.otel.translation import translate_span_when_storing +from mlflow.types.chat import ChatTool +from opentelemetry.exporter.otlp.proto.http.trace_exporter import OTLPSpanExporter from opentelemetry.proto.collector.trace.v1.trace_service_pb2 import ( ExportTraceServiceRequest, ExportTraceServiceResponse, ) -from opentelemetry.proto.trace.v1.trace_pb2 import Span +from opentelemetry.proto.trace.v1.trace_pb2 import Span as _provider_dispatch_Span +from sqlalchemy import Connection, Engine, event, select +from sqlalchemy.orm import Session, sessionmaker +from weave.trace_server.trace_server_interface import CallEndReq, CallStartReq, StartedCallSchemaForInsert, TraceStatus +from core.helper import ssrf_proxy +from core.ops.entities.trace_entity import ( + GenerateNameTraceInfo, + MessageTraceInfo, + ModerationTraceInfo, + WorkflowTraceInfo, +) from core.ops.exceptions import RetryableTraceDispatchError from core.ops.unified_trace.entities import CanonicalSpan, CanonicalSpanKind, CanonicalSpanStatus, CanonicalTrace +from extensions.logstore.aliyun_logstore import AliyunLogStore +from extensions.logstore.repositories.logstore_api_workflow_node_execution_repository import ( + _dict_to_workflow_node_execution_model, +) +from extensions.logstore.repositories.logstore_api_workflow_run_repository import _dict_to_workflow_run +from graphon.enums import BuiltinNodeTypes, WorkflowExecutionStatus, WorkflowNodeExecutionStatus +from libs.datetime_utils import naive_utc_now +from models.account import TenantAccountJoin, TenantAccountRole +from models.application_trace import ApplicationTraceRun +from models.enums import ConversationFromSource, CreatorUserRole, MessageStatus +from models.model import App, AppMode +from models.workflow import ( + WorkflowNodeExecutionModel, + WorkflowNodeExecutionTriggeredFrom, + WorkflowRun, + WorkflowRunTriggeredFrom, + WorkflowType, +) +from repositories import application_trace_repository as repository_module +from repositories.application_trace_repository import TraceCollectionError +from services.tracing import service as service_module from services.tracing.normalization import remap_outer_trace from services.tracing.provider_dispatch import ( TraceExportRejectedError, + TraceExportResult, TraceExportUncertainError, export_trace, provider_capabilities, ) +from services.tracing.service import ApplicationTracingService +from tests.helpers.unit.config_override import apply_config_overrides +from tests.helpers.unit.model_factories import make_account, make_app, make_message, make_tenant +from tests.unit_tests.services.tracing.test_service import _application_trace_collection_batch_for as batch_for +from tests.unit_tests.services.tracing.test_service import _application_trace_collection_claim as claim +from tests.unit_tests.services.tracing.test_service import _application_trace_collection_identifier as identifier +from tests.unit_tests.services.tracing.test_service import _application_trace_collection_ingest as ingest +from tests.unit_tests.services.tracing.test_service import _application_trace_collection_seed as seed +from tests.unit_tests.services.tracing.test_service import _application_trace_collection_terminal as terminal +from tests.unit_tests.services.tracing.test_service import _ApplicationTraceCollectionCollection as Collection + +"""Node logs resolve the caller's trace using SQL or configured LogStore ownership.""" + + +@dataclass +class NodeAlias: + case: Collection + node: dict[str, object] + workflow: dict[str, object] + requests: list[str] + logstore: bool + + def service(self) -> ApplicationTracingService: + return ApplicationTracingService(self.case.repository) + + def status(self, node_execution_id: str | None = None) -> dict[str, object]: + return self.service().status( + tenant_id=self.case.run.tenant_id, + app_id=self.case.run.app_id, + source_kind="node", + source_id=str(self.node["id"]), + node_execution_id=node_execution_id, + ) + + def accept(self, *, spans: bool = True) -> tuple[str, str]: + span_urls: dict[str, str] = {} + for attempt in (0, 1): + invocation = self.case.repository.invocation( + run_id=self.case.run.id, + agent_id=identifier("referenced-agent"), + node_execution_id=str(self.node["id"]), + attempt=attempt, + call_key=f"attempt-{attempt}", + ) + ingest(self.case, batch_for(self.case.run, invocation, span_id=str(attempt + 1) * 16)) + span_urls[invocation.span_id] = f"https://provider.test/span/{attempt}" + terminal(self.case) + lease = claim(self.case) + assert lease is not None + self.case.repository.complete( + self.case.run.id, + lease[1], + state="accepted", + result=TraceExportResult( + "provider-trace", "https://provider.test/trace", span_urls=span_urls if spans else {} + ), + ) + return "https://provider.test/trace", "https://provider.test/span/1" + + +class TestApplicationTraceNodeAliasCases: + @staticmethod + @pytest.fixture(params=["sql", "logstore"]) + def alias( + request: pytest.FixtureRequest, + sqlite_engine: Engine, + sqlite_session_factory: sessionmaker[Session], + monkeypatch: pytest.MonkeyPatch, + ) -> Iterator[NodeAlias]: + case = seed(sqlite_session_factory, "node-alias") + with case.sessions.begin() as session: + app = session.get(App, case.run.app_id) + assert app is not None + app.mode = AppMode.WORKFLOW + common: dict[str, object] = { + "tenant_id": case.run.tenant_id, + "app_id": case.run.app_id, + "workflow_id": identifier("definition"), + "created_by": identifier("actor"), + "created_at": naive_utc_now().isoformat(), + "status": "succeeded", + } + workflow = {**common, "id": identifier("workflow-source"), "graph": json.dumps({"nodes": []})} + node = { + **common, + "id": identifier("node-source"), + "workflow_run_id": workflow["id"], + "node_id": "reference-agent", + "node_type": "agent", + "title": "Reference Agent", + } + requests: list[str] = [] + transactions: set[Connection] = set() + + def begun(connection: Connection) -> None: + transactions.add(connection) + + def ended(connection: Connection) -> None: + transactions.discard(connection) + + if request.param == "sql": + with case.sessions.begin() as session: + session.add(_dict_to_workflow_run(workflow)) + session.add(_dict_to_workflow_node_execution_model(node)) + else: + apply_config_overrides( + monkeypatch, + LOGSTORE_DUAL_WRITE_ENABLED=False, + API_WORKFLOW_RUN_REPOSITORY=( + "extensions.logstore.repositories.logstore_api_workflow_run_repository.LogstoreAPIWorkflowRunRepository" + ), + API_WORKFLOW_NODE_EXECUTION_REPOSITORY=( + "extensions.logstore.repositories.logstore_api_workflow_node_execution_repository." + "LogstoreAPIWorkflowNodeExecutionRepository" + ), + ) + for key, value in { + "LOGSTORE_DUAL_READ_ENABLED": "false", + "LOGSTORE_PG_MODE_ENABLED": "false", + "ALIYUN_SLS_ACCESS_KEY_ID": "test-access-key", + "ALIYUN_SLS_ACCESS_KEY_SECRET": "test-secret", + "ALIYUN_SLS_ENDPOINT": "https://sls.example.test", + "ALIYUN_SLS_REGION": "test-region", + }.items(): + monkeypatch.setenv(key, value) + monkeypatch.setattr(AliyunLogStore, "_instance", None) + monkeypatch.setattr(AliyunLogStore, "_initialized", False) + monkeypatch.setattr( + AliyunLogStore, "_check_endpoint_connectivity", staticmethod(lambda _endpoint, _timeout: None) + ) + + def reject_http_transport(*_args: object, **_kwargs: object) -> None: + pytest.fail("Configured LogStore test escaped its isolated transport boundary") + + monkeypatch.setattr(LogClient, "_getHttpResponse", reject_http_transport) + + def get_logs( + _client: AliyunLogStore, *, query: str, logstore: str, **_kwargs: object + ) -> list[dict[str, object]]: + assert not transactions, "Configured source reads must not hold a SQL transaction" + assert f'tenant_id:"{case.run.tenant_id}"' in query.replace(": ", ":") + assert logstore in { + AliyunLogStore.workflow_execution_logstore, + AliyunLogStore.workflow_node_execution_logstore, + } + is_workflow = logstore == AliyunLogStore.workflow_execution_logstore + if is_workflow: + assert f'app_id:"{case.run.app_id}"' in query + requests.append("workflow" if is_workflow else "node") + row = workflow if is_workflow else node + return [dict(row)] if row else [] + + monkeypatch.setattr(AliyunLogStore, "get_logs", get_logs) + + def execute_sql( + client: AliyunLogStore, *, sql: str, query: str, logstore: str, **_kwargs: object + ) -> list[dict[str, object]]: + # Node lookups use the SQL transport in SDK mode too; the real + # service still validates the returned node-to-workflow owner chain. + assert logstore == AliyunLogStore.workflow_node_execution_logstore + assert f'id:"{identifier("node-source")}"' in query.replace(": ", ":") + assert f"tenant_id = '{case.run.tenant_id}'" in sql + assert f"id = '{identifier('node-source')}'" in sql + assert f'FROM "{logstore}"' in sql + assert "ORDER BY log_version DESC LIMIT 1" in sql + return get_logs(client, query=query, logstore=logstore) + + monkeypatch.setattr(AliyunLogStore, "execute_sql", execute_sql) + with case.sessions() as session: + assert session.scalar(select(WorkflowRun)) is None + assert session.scalar(select(WorkflowNodeExecutionModel)) is None + event.listen(sqlite_engine, "begin", begun) + event.listen(sqlite_engine, "commit", ended) + event.listen(sqlite_engine, "rollback", ended) + try: + run = case.repository.start( + tenant_id=case.run.tenant_id, + app_id=case.run.app_id, + source_kind="workflow", + source_id=str(workflow["id"]), + session_id="caller-session", + ) + assert run is not None + requests.clear() + yield NodeAlias( + Collection(case.repository, case.sessions, run), node, workflow, requests, request.param == "logstore" + ) + finally: + event.remove(sqlite_engine, "begin", begun) + event.remove(sqlite_engine, "commit", ended) + event.remove(sqlite_engine, "rollback", ended) + + @staticmethod + @pytest.mark.parametrize("spans", [True, False]) + def test_alias_returns_parent_trace_and_latest_owned_invocation(alias: NodeAlias, spans: bool) -> None: + trace_url, span_url = alias.accept(spans=spans) + status = alias.status() + assert status["status"] == "accepted" + assert status["source_kind"] == "node" + assert status["source_id"] == alias.node["id"] + assert status["trace_id"] == "provider-trace" + assert status["trace_url"] == trace_url + assert status["span_url"] == (span_url if spans else None) + assert status["config_revision"] == alias.case.run.config_revision + if alias.logstore: + assert alias.requests == ["node", "workflow"] + + @staticmethod + @pytest.mark.parametrize("capture_agent", [True, False]) + @pytest.mark.parametrize("accepted", [True, False]) + def test_alias_without_internal_invocation_returns_parent_status( + alias: NodeAlias, capture_agent: bool, accepted: bool + ) -> None: + with alias.case.sessions.begin() as session: + run = session.get(ApplicationTraceRun, alias.case.run.id) + assert run is not None + run.capture_agent = capture_agent + if accepted: + terminal(alias.case) + lease = claim(alias.case) + assert lease is not None + alias.case.repository.complete( + alias.case.run.id, + lease[1], + state="accepted", + result=TraceExportResult("outer-only", "https://provider.test/outer"), + ) + status = alias.status() + assert status["status"] == ("accepted" if accepted else "collecting") + assert status["span_url"] is None + assert status["capture_agent"] is capture_agent + assert status["trace_url"] == ("https://provider.test/outer" if accepted else None) + + @staticmethod + @pytest.mark.parametrize("state", ["collecting", "not_collected"]) + def test_standalone_node_snapshot_wins_over_parent(alias: NodeAlias, state: str) -> None: + alias.accept() + node_run = alias.case.repository.start( + tenant_id=alias.case.run.tenant_id, + app_id=alias.case.run.app_id, + source_kind="node", + source_id=str(alias.node["id"]), + session_id="preview-session", + ) + assert node_run is not None + with alias.case.sessions.begin() as session: + run = session.get(ApplicationTraceRun, node_run.id) + assert run is not None + run.state = state + alias.requests.clear() + status = alias.status() + assert status["status"] == state + assert status.get("trace_url") is None + if alias.logstore: + assert alias.requests == ["node"] + + @staticmethod + def test_explicit_other_node_filter_cannot_select_sibling_span( + alias: NodeAlias, monkeypatch: pytest.MonkeyPatch + ) -> None: + sibling = alias.case.repository.invocation( + run_id=alias.case.run.id, + agent_id=identifier("referenced-agent"), + node_execution_id=identifier("sibling-node"), + attempt=0, + call_key="sibling", + ) + ingest(alias.case, batch_for(alias.case.run, sibling, span_id="f" * 16)) + alias.accept() + with pytest.raises(TraceCollectionError, match="tracing_node_not_found") as error: + alias.status(identifier("sibling-node")) + assert error.value.status == 404 + service = alias.service() + scheduled: list[str] = [] + monkeypatch.setattr(service, "schedule", scheduled.append) + with pytest.raises(TraceCollectionError, match="tracing_node_not_found"): + service.retry( + tenant_id=alias.case.run.tenant_id, + app_id=alias.case.run.app_id, + source_kind="node", + source_id=str(alias.node["id"]), + node_execution_id=identifier("sibling-node"), + acknowledge_possible_duplicates=True, + ) + assert not scheduled + + @staticmethod + @pytest.mark.parametrize("state", ["failed", "uncertain", "accepted"]) + def test_alias_retry_uses_parent_guards_and_resolves_sources_once( + alias: NodeAlias, state: str, monkeypatch: pytest.MonkeyPatch + ) -> None: + alias.accept() + with alias.case.sessions.begin() as session: + run = session.get(ApplicationTraceRun, alias.case.run.id) + assert run is not None + run.state = state + service = alias.service() + scheduled: list[str] = [] + monkeypatch.setattr(service, "schedule", scheduled.append) + + def retry(acknowledge: bool) -> dict[str, object]: + alias.requests.clear() + return service.retry( + tenant_id=alias.case.run.tenant_id, + app_id=alias.case.run.app_id, + source_kind="node", + source_id=str(alias.node["id"]), + acknowledge_possible_duplicates=acknowledge, + ) + + if state == "accepted": + with pytest.raises(TraceCollectionError, match="tracing_already_accepted") as error: + retry(True) + assert error.value.status == 409 + assert not scheduled + else: + if state == "uncertain": + with pytest.raises(TraceCollectionError, match="tracing_duplicate_acknowledgement_required"): + retry(False) + assert not scheduled + assert retry(state == "uncertain")["status"] == "ready" + assert scheduled == [alias.case.run.id] + if alias.logstore: + assert alias.requests == ["node", "workflow"] + + @staticmethod + @pytest.mark.parametrize("source", ["node", "workflow"]) + @pytest.mark.parametrize("field", ["tenant_id", "app_id", "workflow_id"]) + def test_alias_rejects_broken_owner_chain(alias: NodeAlias, source: str, field: str) -> None: + alias.accept() + row = alias.node if source == "node" else alias.workflow + if alias.logstore: + row[field] = identifier("foreign") + else: + with alias.case.sessions.begin() as session: + model = session.get(WorkflowNodeExecutionModel if source == "node" else WorkflowRun, str(row["id"])) + assert model is not None + # Keep field assignments explicit so ownership failures are visible. + if field == "tenant_id": + model.tenant_id = identifier("foreign") + elif field == "app_id": + model.app_id = identifier("foreign") + else: + model.workflow_id = identifier("foreign") + with pytest.raises(TraceCollectionError, match="tracing_source_not_found") as error: + alias.status() + assert error.value.status == 404 + + @staticmethod + def test_alias_missing_parent_business_source_is_not_a_durable_owner_bypass(alias: NodeAlias) -> None: + if alias.logstore: + alias.workflow.clear() + else: + with alias.case.sessions.begin() as session: + parent = session.get(WorkflowRun, str(alias.workflow["id"])) + assert parent is not None + session.delete(parent) + with pytest.raises(TraceCollectionError, match="tracing_source_not_found"): + alias.status() + + @staticmethod + def test_owned_alias_without_parent_trace_is_not_collected(alias: NodeAlias) -> None: + with alias.case.sessions.begin() as session: + run = session.get(ApplicationTraceRun, alias.case.run.id) + assert run is not None + session.delete(run) + assert alias.status() == {"status": "not_collected", "source_kind": "node", "source_id": alias.node["id"]} + + +"""Real SQLite state transitions and background renewal for delivery leases. + +PostgreSQL row-lock scheduling and the deployed Celery pool still need dev +validation; these tests establish token and acknowledgement semantics. +""" + + +def _stored(case: Collection) -> ApplicationTraceRun: + row = case.repository.get(case.run.id) + assert row is not None + return row + + +def _expire(case: Collection) -> str: + leased = claim(case) + assert leased is not None + with case.sessions.begin() as session: + row = session.get(ApplicationTraceRun, case.run.id) + assert row is not None + row.export_started_at = naive_utc_now() - timedelta(seconds=60) + row.export_heartbeat_at = row.export_started_at + assert claim(case) is None + assert _stored(case).state == "uncertain" + return leased[1] + + +class TestExportLeaseCases: + @staticmethod + @pytest.fixture + def case(sqlite_session_factory: sessionmaker[Session]) -> Collection: + result = seed(sqlite_session_factory, "lease-test") + ingest(result, batch_for(result.run, result.invocation())) + terminal(result) + return result + + @staticmethod + def test_renewal_keeps_long_export_alive_then_missing_heartbeat_expires( + case: Collection, monkeypatch: pytest.MonkeyPatch + ) -> None: + clock = [naive_utc_now()] + monkeypatch.setattr(repository_module, "naive_utc_now", lambda: clock[0]) + leased = claim(case) + assert leased is not None + token = leased[1] + clock[0] += timedelta(seconds=20) + assert case.repository.renew_export_lease(case.run.id, token) + clock[0] += timedelta(seconds=20) + assert claim(case) is None + current = _stored(case) + assert current.state == "exporting" + assert current.attempts == 1 + assert current.export_started_at is not None + assert current.export_started_at + timedelta(seconds=30) < clock[0] + clock[0] += timedelta(seconds=11) + assert claim(case) is None + expired = _stored(case) + assert expired.state == "uncertain" + assert expired.error_code == "export_lease_expired" + assert expired.export_token == token + assert not case.repository.renew_export_lease(case.run.id, token) + + @staticmethod + def test_old_rows_without_heartbeat_use_started_at(case: Collection) -> None: + leased = claim(case) + assert leased is not None + with case.sessions.begin() as session: + row = session.get(ApplicationTraceRun, case.run.id) + assert row is not None + row.export_started_at = naive_utc_now() - timedelta(seconds=60) + row.export_heartbeat_at = None + assert claim(case) is None + current = _stored(case) + assert current.state == "uncertain" + assert current.export_token == leased[1] + + @staticmethod + def test_wrong_token_cannot_renew_or_complete(case: Collection) -> None: + leased = claim(case) + assert leased is not None + before = _stored(case) + assert not case.repository.renew_export_lease(case.run.id, "other-token") + case.repository.complete(case.run.id, "other-token", state="accepted", result=TraceExportResult("wrong")) + current = _stored(case) + assert current.export_token == leased[1] + assert current.export_heartbeat_at == before.export_heartbeat_at + assert current.state == "exporting" + assert current.provider_trace_id is None + + @staticmethod + def test_late_success_wins_before_manual_retry_and_is_idempotent(case: Collection) -> None: + token = _expire(case) + result = TraceExportResult("confirmed", "https://example.com/trace") + case.repository.complete(case.run.id, token, state="accepted", result=result) + current = _stored(case) + assert current.state == "accepted" + assert current.error_code is None + assert current.export_token is None + assert current.accepted_at is not None + assert current.trace_url == result.trace_url + case.repository.complete(case.run.id, token, state="accepted", result=TraceExportResult("duplicate-ack")) + assert _stored(case).provider_trace_id == "confirmed" + with pytest.raises(TraceCollectionError, match="tracing_already_accepted"): + case.repository.retry(run_id=case.run.id, acknowledge_possible_duplicates=True) + + @staticmethod + def test_manual_retry_revokes_old_heartbeat_and_ack_before_and_after_new_claim(case: Collection) -> None: + old_token = _expire(case) + case.repository.retry(run_id=case.run.id, acknowledge_possible_duplicates=True) + assert not case.repository.renew_export_lease(case.run.id, old_token) + case.repository.complete(case.run.id, old_token, state="accepted", result=TraceExportResult("old")) + assert _stored(case).state == "ready" + second = claim(case) + assert second is not None + assert second[1] != old_token + case.repository.complete(case.run.id, old_token, state="accepted", result=TraceExportResult("old")) + current = _stored(case) + assert current.state == "exporting" + assert current.export_token == second[1] + assert current.provider_trace_id is None + case.repository.complete(case.run.id, second[1], state="accepted", result=TraceExportResult("new")) + assert _stored(case).provider_trace_id == "new" + + @staticmethod + @pytest.mark.parametrize("late_state", ["failed", "retrying", "uncertain"]) + def test_late_non_success_cannot_overwrite_lease_expiry(case: Collection, late_state: str) -> None: + token = _expire(case) + case.repository.complete(case.run.id, token, state=late_state, error_code="late-failure") + current = _stored(case) + assert current.state == "uncertain" + assert current.error_code == "export_lease_expired" + + @staticmethod + def test_other_uncertain_reason_cannot_be_promoted_by_late_ack(case: Collection) -> None: + token = _expire(case) + with case.sessions.begin() as session: + row = session.get(ApplicationTraceRun, case.run.id) + assert row is not None + row.error_code = "provider_acceptance_unknown" + case.repository.complete(case.run.id, token, state="accepted", result=TraceExportResult("late")) + assert _stored(case).state == "uncertain" + + @staticmethod + @pytest.mark.parametrize("transient_failure", [False, True]) + @pytest.mark.parametrize("outcome", ["accepted", "uncertain"]) + def test_service_heartbeats_during_export_and_stops_on_every_outcome( + case: Collection, + monkeypatch: pytest.MonkeyPatch, + caplog: pytest.LogCaptureFixture, + transient_failure: bool, + outcome: str, + ) -> None: + apply_config_overrides(monkeypatch, OPS_AGENT_TRACE_EXPORT_LEASE_SECONDS=1) + clock = [naive_utc_now()] + monkeypatch.setattr(repository_module, "naive_utc_now", lambda: clock[0]) + monkeypatch.setattr(service_module.OpsTraceManager, "decrypt_tracing_config", lambda *_args: {}) + renewed = threading.Event() + calls = [] + original = case.repository.renew_export_lease + + def renew(run_id: str, token: str) -> bool: + calls.append(threading.get_ident()) + if transient_failure and len(calls) == 1: + raise RuntimeError("private-provider-value") + result = original(run_id, token) + renewed.set() + return result + + monkeypatch.setattr(case.repository, "renew_export_lease", renew) + + def export(*_args: object) -> TraceExportResult: + # The original start time is now beyond the lease; only a real background + # transaction can prevent a second claimant from marking it uncertain. + clock[0] += timedelta(seconds=3) + assert renewed.wait(3), "background heartbeat did not commit" + assert case.repository.claim(case.run.id, collection_timeout=60, lease_seconds=1) is None + assert _stored(case).state == "exporting" + if outcome == "uncertain": + raise TraceExportUncertainError("lost acknowledgement") + return TraceExportResult("ack") + + monkeypatch.setattr(service_module, "export_trace", export) + ApplicationTracingService(case.repository).deliver(case.run.id) + assert _stored(case).state == outcome + assert calls + assert all(thread_id != threading.get_ident() for thread_id in calls) + assert not any(thread.name == f"application-trace-lease-{case.run.id}" for thread in threading.enumerate()) + assert "private-provider-value" not in caplog.text + + @staticmethod + def test_heartbeat_stops_even_when_completion_database_write_raises( + case: Collection, monkeypatch: pytest.MonkeyPatch + ) -> None: + monkeypatch.setattr(service_module.OpsTraceManager, "decrypt_tracing_config", lambda *_args: {}) + monkeypatch.setattr(service_module, "export_trace", lambda *_args: TraceExportResult("ack")) + + def fail_complete(*_args: object, **_kwargs: object) -> None: + raise RuntimeError("database unavailable") + + monkeypatch.setattr(case.repository, "complete", fail_complete) + with pytest.raises(RuntimeError, match="database unavailable"): + ApplicationTracingService(case.repository).deliver(case.run.id) + assert not any(thread.name == f"application-trace-lease-{case.run.id}" for thread in threading.enumerate()) + + +"""Persisted legacy sources through real OTel protobuf, replacing only HTTP I/O.""" + + +_LEGACY_ALIYUN_TRANSPORT_START = datetime(2026, 10, 19, tzinfo=UTC) +RUN_ID = "a0000000-0000-4000-8000-000000000001" +MESSAGE_ID = "a0000000-0000-4000-8000-000000000002" +NODE_IDS = [f"b0000000-0000-4000-8000-00000000000{i}" for i in range(3)] + + +class _LegacyAliyunTransportRequestoptions(TypedDict, total=False): + content: bytes + headers: Mapping[str, str] + max_retries: int + follow_redirects: bool + timeout: float + + +def _legacy_aliyun_transport_config(**overrides: object) -> AliyunConfig: + return AliyunConfig.model_validate( + { + "app_name": "owned-service", + "license_key": "owned-key", + "endpoint": "https://collector.test/prefix", + **overrides, + } + ) + + +def decoded(options: _LegacyAliyunTransportRequestoptions) -> ExportTraceServiceRequest: + payload = ExportTraceServiceRequest() + payload.ParseFromString(options["content"]) + return payload + + +class TestLegacyAliyunTransportCases: + @staticmethod + def test_constructor_and_presentation_have_no_network_sdk_queue_or_environment_mutation( + monkeypatch: pytest.MonkeyPatch, + ) -> None: + def forbidden(*_args: object, **_kwargs: object) -> None: + pytest.fail("Constructor and presentation must not create I/O, exporters or worker threads") + + monkeypatch.setattr(ssrf_proxy, "make_request", forbidden) + monkeypatch.setattr(OTLPSpanExporter, "__init__", forbidden) + monkeypatch.setattr(Thread, "start", forbidden) + monkeypatch.setenv("OTEL_EXPORTER_OTLP_ENDPOINT", "https://ambient.test") + monkeypatch.setenv("OTEL_EXPORTER_OTLP_HEADERS", "Authorization=ambient") + before = dict(os.environ) + original = _legacy_aliyun_transport_config() + adapter = AliyunDataTrace(original) + original.license_key = "mutated-after-construction" + assert adapter.trace_client.endpoint == "https://collector.test/prefix/adapt_owned-key/api/otlp/traces" + assert adapter.get_project_url() == "https://arms.console.aliyun.com/#/llm" + assert not hasattr(adapter.trace_client, "exporter") + assert not hasattr(adapter.trace_client, "queue") + assert not hasattr(adapter.trace_client, "worker_thread") + assert dict(os.environ) == before + + @staticmethod + @pytest.mark.usefixtures("agent_runtime_app") + @pytest.mark.parametrize("failure", [None, 400, 401, 403, 429, 503, "timeout", "partial", "malformed", "redirect"]) + def test_persisted_workflow_uses_one_acknowledged_protobuf_batch_and_preserves_business_state( + sqlite_session: Session, monkeypatch: pytest.MonkeyPatch, failure: int | str | None + ) -> None: + tenant = make_tenant(tenant_id="tenant") + account = make_account(account_id="owner", tenant=tenant) + app = make_app(app_id="app", tenant_id=tenant.id, mode=AppMode.WORKFLOW, created_by=account.id) + membership = TenantAccountJoin(tenant_id=tenant.id, account_id=account.id, role=TenantAccountRole.OWNER) + run = WorkflowRun( + id=RUN_ID, + tenant_id=tenant.id, + app_id=app.id, + workflow_id="workflow", + type=WorkflowType.WORKFLOW, + triggered_from=WorkflowRunTriggeredFrom.DEBUGGING, + version="v1", + inputs=json.dumps({"query": "hello"}), + outputs=json.dumps({"answer": "ok"}), + status=WorkflowExecutionStatus.SUCCEEDED, + total_tokens=32, + created_by=account.id, + created_by_role=CreatorUserRole.ACCOUNT, + created_at=_LEGACY_ALIYUN_TRANSPORT_START, + finished_at=_LEGACY_ALIYUN_TRANSPORT_START + timedelta(seconds=1), + elapsed_time=1, + ) + sqlite_session.add_all([tenant, account, app, membership, run]) + nodes = [ + WorkflowNodeExecutionModel( + id=NODE_IDS[i], + tenant_id=tenant.id, + app_id=app.id, + workflow_id="workflow", + workflow_run_id=RUN_ID, + triggered_from=WorkflowNodeExecutionTriggeredFrom.WORKFLOW_RUN, + index=i, + node_id=f"node-{i}", + title=title, + node_type=kind, + status=WorkflowNodeExecutionStatus.SUCCEEDED, + created_by=account.id, + created_by_role=CreatorUserRole.ACCOUNT, + inputs=json.dumps({"synthetic": True}), + outputs=json.dumps( + {"text": "ok", "usage": {"prompt_tokens": 21, "completion_tokens": 11, "total_tokens": 32}} + ), + created_at=_LEGACY_ALIYUN_TRANSPORT_START + timedelta(milliseconds=i), + finished_at=_LEGACY_ALIYUN_TRANSPORT_START + timedelta(milliseconds=i + 100), + elapsed_time=0.1, + process_data=json.dumps( + { + "prompts": [{"role": "user", "text": "hello"}], + "model_provider": "test", + "model_name": "test-model", + } + ) + if i == 1 + else None, + ) + for i, (title, kind) in enumerate( + [("Start", BuiltinNodeTypes.START), ("LLM", BuiltinNodeTypes.LLM), ("End", BuiltinNodeTypes.END)] + ) + ] + sqlite_session.add_all(nodes) + sqlite_session.commit() + info = WorkflowTraceInfo( + query="hello", + workflow_id="workflow", + tenant_id=tenant.id, + workflow_run_id=RUN_ID, + workflow_run_elapsed_time=1.0, + workflow_run_status="succeeded", + workflow_run_inputs={"query": "hello"}, + workflow_run_outputs={"answer": "ok"}, + workflow_run_version="v1", + total_tokens=32, + file_list=[], + metadata={"app_id": app.id, "conversation_id": "owned-conversation"}, + start_time=_LEGACY_ALIYUN_TRANSPORT_START, + end_time=_LEGACY_ALIYUN_TRANSPORT_START + timedelta(seconds=1), + ) + sent: list[tuple[str, _LegacyAliyunTransportRequestoptions, ExportTraceServiceRequest]] = [] + + def send(method: str, url: str, **options: Unpack[_LegacyAliyunTransportRequestoptions]) -> httpx.Response: + assert method == "POST" + assert url == "https://collector.test/prefix/adapt_owned-key/api/otlp/traces" + sent.append((url, options, decoded(options))) + if len(sent) == 1: + if failure == "timeout": + raise httpx.ReadTimeout("synthetic lost ACK", request=httpx.Request("POST", url)) + if isinstance(failure, int): + return httpx.Response(failure, content=b"DO_NOT_EXPOSE_BODY") + if failure == "malformed": + return httpx.Response(200, content=b"not protobuf") + if failure == "partial": + response = ExportTraceServiceResponse() + response.partial_success.rejected_spans = 1 + response.partial_success.error_message = "DO_NOT_EXPOSE_BODY" + return httpx.Response(200, content=response.SerializeToString()) + if failure == "redirect": + return httpx.Response(307, headers={"Location": "https://unrelated.test"}) + return httpx.Response(200, content=ExportTraceServiceResponse().SerializeToString()) + + monkeypatch.setattr(ssrf_proxy, "make_request", send) + adapter = AliyunDataTrace(_legacy_aliyun_transport_config()) + if failure is None: + adapter.trace(info) + else: + expected = ( + RetryableTraceDispatchError + if failure in {429, 503} + else TraceExportUncertainError + if failure in {"timeout", "partial", "malformed"} + else TraceExportRejectedError + ) + with pytest.raises(expected) as error: + adapter.trace(info) + assert "DO_NOT_EXPOSE_BODY" not in str(error.value) + if isinstance(error.value, RetryableTraceDispatchError): + assert error.value.acceptance_unknown is (expected is TraceExportUncertainError) + assert len(sent) == 1 # No SDK retry, per-span request, batching thread or silent queue success. + _, options, payload = sent[0] + assert options["max_retries"] == 0 + assert options["follow_redirects"] is False + assert options["headers"] == {"Content-Type": "application/x-protobuf", "Accept": "application/x-protobuf"} + resource = payload.resource_spans[0] + attributes = {item.key: item.value for item in resource.resource.attributes} + assert attributes["service.name"].string_value == "owned-service" + spans = list(resource.scope_spans[0].spans) + assert len(spans) == 4 + assert [s.name for s in spans] == ["workflow", "Start", "LLM", "End"] + root_id = convert_to_span_id(RUN_ID, "workflow").to_bytes(8) + assert [s.span_id for s in spans] == [root_id, *[convert_to_span_id(i, "node").to_bytes(8) for i in NODE_IDS]] + assert all(s.trace_id == UUID(RUN_ID).bytes for s in spans) + assert [s.parent_span_id for s in spans] == [b"", root_id, root_id, root_id] + assert spans[0].start_time_unix_nano == int(_LEGACY_ALIYUN_TRANSPORT_START.timestamp() * 1_000_000_000) + assert spans[0].end_time_unix_nano == int( + (_LEGACY_ALIYUN_TRANSPORT_START + timedelta(seconds=1)).timestamp() * 1_000_000_000 + ) + model_attributes = {item.key: item.value for item in spans[2].attributes} + assert model_attributes[GEN_AI_USAGE_TOTAL_TOKENS].string_value == "32" + assert all(s.status.code == 1 for s in spans) + assert adapter.trace_client._dispatch_spans.get() is None + assert info.workflow_run_status == "succeeded" + sqlite_session.refresh(run) + assert run.status == WorkflowExecutionStatus.SUCCEEDED + assert run.total_tokens == 32 + assert all(node.status == WorkflowNodeExecutionStatus.SUCCEEDED for node in nodes) + # A fresh business operation does not inherit any old batch after success or error. + adapter.trace(ModerationTraceInfo(flagged=False, action="allow", preset_response="", query="", metadata={})) + assert len(sent) == 1 + adapter.trace(info.model_copy(update={"workflow_run_id": "a0000000-0000-4000-8000-000000000003"})) + assert len(sent) == 2 + recovered_spans = list(sent[1][2].resource_spans[0].scope_spans[0].spans) + assert [s.name for s in recovered_spans] == ["workflow"] + + @staticmethod + @pytest.mark.parametrize("failed", [False, True]) + def test_message_mapping_keeps_original_ids_raw_usage_error_and_business_state( + sqlite_session: Session, monkeypatch: pytest.MonkeyPatch, failed: bool + ) -> None: + message = make_message( + message_id=MESSAGE_ID, + app_id="app", + conversation_id="conversation", + query="hello", + answer="ok", + status=MessageStatus.ERROR if failed else MessageStatus.NORMAL, + inputs={}, + message={}, + message_unit_price=Decimal(0), + answer_unit_price=Decimal(0), + currency="USD", + from_source=ConversationFromSource.CONSOLE, + from_account_id="owner", + message_tokens=21, + answer_tokens=11, + ) + sqlite_session.add(message) + sqlite_session.commit() + info = MessageTraceInfo( + conversation_model="chat", + conversation_mode="chat", + message_id=MESSAGE_ID, + message_data=message, + inputs={"prompt": "hello"}, + outputs="ok", + message_tokens=21, + answer_tokens=11, + total_tokens=32, + error="synthetic model failure" if failed else None, + metadata={ + "app_id": "app", + "conversation_id": "conversation", + "ls_model_name": "test-model", + "ls_provider": "test", + }, + start_time=_LEGACY_ALIYUN_TRANSPORT_START, + end_time=_LEGACY_ALIYUN_TRANSPORT_START + timedelta(seconds=1), + ) + sent: list[ExportTraceServiceRequest] = [] + + def send(_method: str, _url: str, **options: Unpack[_LegacyAliyunTransportRequestoptions]) -> httpx.Response: + sent.append(decoded(options)) + return httpx.Response(200, content=b"") + + monkeypatch.setattr(ssrf_proxy, "make_request", send) + adapter = AliyunDataTrace(_legacy_aliyun_transport_config()) + adapter.trace(info) + assert len(sent) == 1 + spans = list(sent[0].resource_spans[0].scope_spans[0].spans) + assert [s.name for s in spans] == ["message", "llm"] + root_id = convert_to_span_id(MESSAGE_ID, "message").to_bytes(8) + assert [s.parent_span_id for s in spans] == [b"", root_id] + assert [s.span_id for s in spans] == [root_id, convert_to_span_id(MESSAGE_ID, "llm").to_bytes(8)] + assert all(s.trace_id == UUID(MESSAGE_ID).bytes and s.status.code == (2 if failed else 1) for s in spans) + assert {a.key: a.value for a in spans[1].attributes}[GEN_AI_USAGE_TOTAL_TOKENS].string_value == "32" + assert message.status == (MessageStatus.ERROR if failed else MessageStatus.NORMAL) + + @staticmethod + def test_one_cached_client_keeps_interleaved_batches_and_failure_cleanup_isolated( + monkeypatch: pytest.MonkeyPatch, + ) -> None: + client = TraceClient(service_name="owned-service", endpoint="https://collector.test/adapt_key/api/otlp/traces") + barrier = Barrier(2) + observed: list[list[str]] = [] + + def send(_method: str, _url: str, **options: Unpack[_LegacyAliyunTransportRequestoptions]) -> httpx.Response: + payload = decoded(options) + observed.append([s.name for rs in payload.resource_spans for scope in rs.scope_spans for s in scope.spans]) + return httpx.Response(200, content=b"") + + monkeypatch.setattr(ssrf_proxy, "make_request", send) + + def run(label: str) -> None: + try: + with client.batch(): + client.add_span(SpanData(name=label, trace_id=1, span_id=2, start_time=1, end_time=2)) + barrier.wait(timeout=5) + if label == "failed": + raise ValueError("mapper failed before submission") + client.add_span(SpanData(name=label + "-child", trace_id=1, span_id=3, start_time=1, end_time=2)) + except ValueError: + assert label == "failed" + assert client._dispatch_spans.get() is None + + with ThreadPoolExecutor(max_workers=2) as pool: + list(pool.map(run, ["accepted", "failed"])) + assert observed == [["accepted", "accepted-child"]] + with client.batch(): + client.add_span(SpanData(name="next", trace_id=1, span_id=4, start_time=1, end_time=2)) + assert observed == [["accepted", "accepted-child"], ["next"]] + + +"""Real Weave SDK serialization with only the outbound HTTP boundary replaced.""" + + +_LEGACY_WEAVE_TRANSPORT_START = datetime(2026, 10, 19, tzinfo=UTC) + + +class _LegacyWeaveTransportRequestoptions(TypedDict, total=False): + content: bytes + json: dict[str, object] + auth: tuple[str, str] + headers: Mapping[str, str] + max_retries: int + follow_redirects: bool + timeout: float + + +def _legacy_weave_transport_config(**overrides: object) -> WeaveConfig: + return WeaveConfig.model_validate({"api_key": "fixture-key", "entity": "team", "project": "project", **overrides}) + + +def call(call_id: str = "call-1", **overrides: object) -> WeaveTraceModel: + return WeaveTraceModel.model_validate( + { + "id": call_id, + "op": "tool", + "inputs": {"synthetic": True}, + "outputs": {"result": "ok"}, + "attributes": { + "trace_id": call_id, + "start_time": _LEGACY_WEAVE_TRANSPORT_START, + "end_time": _LEGACY_WEAVE_TRANSPORT_START + timedelta(seconds=1), + }, + **overrides, + } + ) + + +def response(url: str, body: object, *, status: int = 200) -> httpx.Response: + return httpx.Response(status, request=httpx.Request("POST", url), json=body) + + +class TestLegacyWeaveTransportCases: + @staticmethod + @pytest.fixture + def wire( + monkeypatch: pytest.MonkeyPatch, failure: tuple[str, str, int | str] | None + ) -> list[tuple[str, _LegacyWeaveTransportRequestoptions, dict[str, object]]]: + sent: list[tuple[str, _LegacyWeaveTransportRequestoptions, dict[str, object]]] = [] + fault_pending = True + + def send(method: str, url: str, **options: Unpack[_LegacyWeaveTransportRequestoptions]) -> httpx.Response: + nonlocal fault_pending + assert method == "POST" + assert url in {"https://trace.wandb.ai/call/start", "https://trace.wandb.ai/call/end"} + body = json.loads(options["content"]) + sent.append((url, options, body)) + stage = "start" if "start" in body else "end" + record = body[stage] + if fault_pending and failure is not None and (stage, record["id"]) == failure[:2]: + fault_pending = False + outcome = failure[2] + if outcome == "timeout": + raise httpx.ReadTimeout("synthetic lost acknowledgment", request=httpx.Request("POST", url)) + if outcome == "malformed": + return httpx.Response(200, request=httpx.Request("POST", url), content=b"invalid-json") + assert isinstance(outcome, int) + return response(url, {"error": "DO_NOT_EXPOSE_BODY"}, status=outcome) + if url.endswith("/call/start"): + req = CallStartReq.model_validate(body) + return response(url, {"id": req.start.id, "trace_id": req.start.trace_id}) + CallEndReq.model_validate(body) + return response(url, {}) + + monkeypatch.setattr(ssrf_proxy, "make_request", send) + return sent + + @staticmethod + def test_constructor_and_presentation_have_no_login_network_files_or_background_state( + monkeypatch: pytest.MonkeyPatch, + ) -> None: + def forbidden(*_args: object, **_kwargs: object) -> None: + pytest.fail("Constructor must not initialize global credentials, files, I/O or background work") + + monkeypatch.setattr(Path, "touch", forbidden) + monkeypatch.setattr(wandb, "login", forbidden) + monkeypatch.setattr(weave, "init", forbidden) + monkeypatch.setattr("wandb.sdk.lib.wbauth.wbnetrc.write_netrc_auth", forbidden) + monkeypatch.setattr("weave.trace_server_bindings.remote_http_trace_server.AsyncBatchProcessor", forbidden) + monkeypatch.setattr("weave.trace_server_bindings.remote_http_trace_server.CallBatchProcessor", forbidden) + monkeypatch.setattr(ssrf_proxy, "make_request", forbidden) + monkeypatch.setattr(atexit, "register", forbidden) + monkeypatch.setenv("WANDB_API_KEY", "unrelated-key") + monkeypatch.setenv("WANDB_BASE_URL", "https://unrelated.test") + monkeypatch.setenv("WEAVE_TRACE_SERVER", "https://unrelated-traces.test") + before = dict(os.environ) + original = _legacy_weave_transport_config(entity=None) + adapter = WeaveDataTrace(original) + original.api_key = "changed-after-construction" + assert adapter.weave_api_key == "fixture-key" + assert adapter.get_project_url() == "https://wandb.ai" + assert adapter.server.trace_server_url == "https://trace.wandb.ai" + assert adapter.server._auth == ("api", "fixture-key") + assert not adapter.server.should_batch + assert adapter.server.call_processor is None + assert adapter.server.feedback_processor is None + assert dict(os.environ) == before + + @staticmethod + @pytest.mark.usefixtures("agent_runtime_app") + @pytest.mark.parametrize( + "failure", + [None] + + [ + (stage, call_id, status) + for stage, call_id in [ + ("start", "legacy-run"), + ("end", "legacy-run"), + ("start", "node-0"), + ("end", "node-0"), + ] + for status in [429, 503] + ] + + [ + ("start", "legacy-run", "timeout"), + ("start", "legacy-run", "malformed"), + ("end", "node-0", "malformed"), + ("end", "legacy-run", "timeout"), + ] + + [("start", "legacy-run", 401), ("end", "legacy-run", 401)], + ) + def test_workflow_emits_four_real_sdk_calls_with_original_ids_payload_and_status( + sqlite_session: Session, + wire: list[tuple[str, _LegacyWeaveTransportRequestoptions, dict[str, object]]], + failure: tuple[str, str, int | str] | None, + ) -> None: + # The normal Flask/SQLite fixtures bind the actual service-account lookup and + # production repository to persisted rows. Only outbound HTTP is substituted. + tenant = make_tenant(tenant_id="tenant") + account = make_account(account_id="owner", tenant=tenant) + app = make_app(app_id="app", tenant_id=tenant.id, mode=AppMode.WORKFLOW, created_by=account.id) + membership = TenantAccountJoin(tenant_id=tenant.id, account_id=account.id, role=TenantAccountRole.OWNER) + sqlite_session.add_all([tenant, account, app, membership]) + nodes = [ + WorkflowNodeExecutionModel( + id=f"node-{i}", + tenant_id=tenant.id, + app_id=app.id, + workflow_id="workflow", + workflow_run_id="legacy-run", + triggered_from=WorkflowNodeExecutionTriggeredFrom.WORKFLOW_RUN, + index=i, + node_id=f"node-{i}", + title=title, + node_type=kind, + status=WorkflowNodeExecutionStatus.SUCCEEDED, + created_by=account.id, + created_by_role=CreatorUserRole.ACCOUNT, + inputs=json.dumps({"synthetic": True}), + outputs=json.dumps({"answer": "ok"}), + created_at=_LEGACY_WEAVE_TRANSPORT_START + timedelta(milliseconds=i), + finished_at=_LEGACY_WEAVE_TRANSPORT_START + timedelta(milliseconds=i + 100), + elapsed_time=0.1, + process_data=json.dumps( + { + "prompts": [{"role": "user", "text": "hello"}], + "model_mode": "chat", + "model_provider": "test", + "model_name": "test-model", + } + ) + if i == 1 + else None, + execution_metadata=json.dumps( + {"total_tokens": 32, "total_price": "0.0000098", "currency": "USD"} if i == 1 else {} + ), + ) + for i, (title, kind) in enumerate( + [("Start", BuiltinNodeTypes.START), ("LLM", BuiltinNodeTypes.LLM), ("End", BuiltinNodeTypes.END)] + ) + ] + sqlite_session.add_all(nodes) + sqlite_session.commit() + adapter = WeaveDataTrace(_legacy_weave_transport_config()) + info = WorkflowTraceInfo( + query="hello", + workflow_id="workflow", + tenant_id="tenant", + workflow_run_id="legacy-run", + workflow_run_elapsed_time=1.0, + workflow_run_status="succeeded", + workflow_run_inputs={"query": "hello"}, + workflow_run_outputs={"answer": "ok"}, + workflow_run_version="v1", + total_tokens=32, + file_list=[], + metadata={"app_id": "app", "conversation_id": None}, + trace_id="external-correlation", + start_time=_LEGACY_WEAVE_TRANSPORT_START, + end_time=_LEGACY_WEAVE_TRANSPORT_START + timedelta(seconds=1), + ) + + if failure is not None: + stage, call_id, outcome = failure + first_start = (stage, call_id) == ("start", "legacy-run") + expected = ( + TraceExportRejectedError + if outcome == 401 + else RetryableTraceDispatchError + if first_start and outcome in {429, 503} + else TraceExportUncertainError + ) + with pytest.raises(expected) as error: + adapter.trace(info) + if isinstance(error.value, RetryableTraceDispatchError): + assert error.value.acceptance_unknown is (expected is TraceExportUncertainError) + assert len(wire) == (1 if first_start else 8 if call_id == "legacy-run" else 2 if stage == "start" else 3) + assert "DO_NOT_EXPOSE" not in str(error.value) + assert adapter.calls == {} + assert info.workflow_run_status == "succeeded" + assert all(node.status == WorkflowNodeExecutionStatus.SUCCEEDED for node in nodes) + before_recovery = len(wire) + adapter.trace( + GenerateNameTraceInfo( + tenant_id=tenant.id, + inputs="new business operation", + outputs="new title", + metadata={}, + start_time=_LEGACY_WEAVE_TRANSPORT_START, + end_time=_LEGACY_WEAVE_TRANSPORT_START + timedelta(seconds=1), + ) + ) + assert len(wire) == before_recovery + 2 + assert adapter.calls == {} + return + + adapter.trace(info) + + starts = [CallStartReq.model_validate(body).start for _, _, body in wire if "start" in body] + ends = [CallEndReq.model_validate(body).end for _, _, body in wire if "end" in body] + assert len(starts) == len(ends) == 4 + assert [s.id for s in starts] == ["legacy-run", "node-0", "node-1", "node-2"] + assert [s.parent_id for s in starts] == [None, "legacy-run", "legacy-run", "legacy-run"] + assert all(s.trace_id == "external-correlation" and s.project_id == "team/project" for s in starts) + # Legacy has no native thread/turn contract: keep the existing raw metadata. + assert all(s.thread_id is None and s.turn_id is None for s in starts) + assert starts[0].attributes["conversation_id"] is None + model = starts[2] + assert model.attributes["total_tokens"] == 32 + assert model.attributes["total_price"] == "0.0000098" + assert model.inputs["messages"][0]["content"] == "hello" + model_end = next(e for e in ends if e.id == "node-1") + assert isinstance(model_end.output, dict) + assert model_end.output["usage_metadata"]["total_tokens"] == 32 + assert all(e.exception is None and e.summary["status_counts"][TraceStatus.SUCCESS] == 1 for e in ends) + assert all("usage" not in e.summary for e in ends) # Do not invent native aggregation for the old mapper. + assert all(options["auth"] == ("api", "fixture-key") and options["max_retries"] == 0 for _, options, _ in wire) + assert adapter.calls == {} + + @staticmethod + def test_interleaved_instances_resolve_only_their_own_credentials_and_targets( + monkeypatch: pytest.MonkeyPatch, + ) -> None: + barrier = Barrier(2) + observed: list[tuple[str, tuple[str, str], str]] = [] + monkeypatch.setenv("WANDB_API_KEY", "ambient-key") + monkeypatch.setenv("WANDB_ENTITY", "ambient-entity") + monkeypatch.setenv("WANDB_PROJECT", "ambient-project") + before = dict(os.environ) + + def send(_method: str, url: str, **options: Unpack[_LegacyWeaveTransportRequestoptions]) -> httpx.Response: + label = ( + "a" + if url + in { + "https://control-a.test/prefix/graphql", + "https://trace-a.test/prefix/call/start", + "https://trace-a.test/prefix/call/end", + } + else "b" + ) + assert url in { + f"https://control-{label}.test/prefix/graphql", + f"https://trace-{label}.test/prefix/call/start", + f"https://trace-{label}.test/prefix/call/end", + } + assert options["auth"] == ("api", f"key-{label}") + if url.endswith("/graphql"): + return response(url, {"data": {"viewer": {"entity": f"team-{label}"}}}) + body = json.loads(options["content"]) + record = body.get("start", body.get("end")) + assert record["project_id"] == f"team-{label}/project-{label}" + observed.append((url, options["auth"], record["project_id"])) + if "start" in body: + barrier.wait(timeout=5) + return response(url, {"id": record["id"], "trace_id": record["trace_id"]}) + return response(url, {}) + + monkeypatch.setattr(ssrf_proxy, "make_request", send) + + def run(label: str) -> None: + adapter = WeaveDataTrace( + _legacy_weave_transport_config( + api_key=f"key-{label}", + entity=None, + project=f"project-{label}", + host=f"https://control-{label}.test/prefix", + endpoint=f"https://trace-{label}.test/prefix/", + ) + ) + data = call(f"call-{label}") + adapter.start_call(data) + adapter.finish_call(data) + + with ThreadPoolExecutor(max_workers=2) as pool: + list(pool.map(run, ["a", "b"])) + assert len(observed) == 4 + assert dict(os.environ) == before + + @staticmethod + def test_concurrent_dispatches_on_one_cached_instance_keep_prefix_state_isolated( + monkeypatch: pytest.MonkeyPatch, + ) -> None: + barrier = Barrier(2) + + def send(_method: str, url: str, **options: Unpack[_LegacyWeaveTransportRequestoptions]) -> httpx.Response: + body = json.loads(options["content"]) + if "start" in body: + start = CallStartReq.model_validate(body).start + if start.inputs["label"] == "unaccepted": + barrier.wait(timeout=5) + return response(url, {}, status=503) + return response(url, {"id": start.id, "trace_id": start.trace_id}) + CallEndReq.model_validate(body) + barrier.wait(timeout=5) + return response(url, {}, status=503) + + monkeypatch.setattr(ssrf_proxy, "make_request", send) + adapter = WeaveDataTrace(_legacy_weave_transport_config()) + + def run(label: str) -> RetryableTraceDispatchError: + info = GenerateNameTraceInfo( + tenant_id="tenant", + inputs={"label": label}, + outputs={"label": label}, + metadata={}, + start_time=_LEGACY_WEAVE_TRANSPORT_START, + end_time=_LEGACY_WEAVE_TRANSPORT_START + timedelta(seconds=1), + ) + with pytest.raises(RetryableTraceDispatchError) as error: + adapter.trace(info) + return error.value + + with ThreadPoolExecutor(max_workers=2) as pool: + accepted, unaccepted = list(pool.map(run, ["accepted", "unaccepted"])) + assert isinstance(accepted, TraceExportUncertainError) + assert accepted.acceptance_unknown + assert not unaccepted.acceptance_unknown + assert adapter.calls == {} + + @staticmethod + @pytest.mark.parametrize( + ("stage", "failure"), + [ + (stage, failure) + for stage in ("start", "end") + for failure in ("rejected", "retryable", "malformed", "mismatched_id") + if not (stage == "end" and failure == "mismatched_id") + ], + ) + def test_synchronous_failures_propagate_and_next_call_recovers( + monkeypatch: pytest.MonkeyPatch, stage: str, failure: str + ) -> None: + broken = True + requests = [] + + def send(_method: str, url: str, **options: Unpack[_LegacyWeaveTransportRequestoptions]) -> httpx.Response: + body = json.loads(options["content"]) + requests.append(url) + if broken and url.endswith("/call/" + stage): + if failure == "rejected": + return response(url, {"error": "DO_NOT_EXPOSE_KEY_OR_BODY"}, status=401) + if failure == "retryable": + return response(url, {}, status=503) + if failure == "malformed": + return httpx.Response(200, request=httpx.Request("POST", url), content=b"invalid-json") + return response(url, {"id": "wrong-id", "trace_id": "wrong-trace"}) + if "start" in body: + return response(url, {"id": body["start"]["id"], "trace_id": body["start"]["trace_id"]}) + return response(url, {}) + + monkeypatch.setattr(ssrf_proxy, "make_request", send) + adapter = WeaveDataTrace(_legacy_weave_transport_config()) + data = call() + expected = { + "rejected": TraceExportRejectedError, + "retryable": RetryableTraceDispatchError, + "malformed": TraceExportUncertainError, + "mismatched_id": TraceExportUncertainError, + }[failure] + if stage == "end": + adapter.start_call(data) + with pytest.raises(expected) as error: + adapter.start_call(data) if stage == "start" else adapter.finish_call(data) + assert "DO_NOT_EXPOSE" not in str(error.value) + assert len(requests) == (1 if stage == "start" else 2) + assert adapter.calls == {} + broken = False + recovered = call("next-call", exception="synthetic tool error") + adapter.start_call(recovered) + adapter.finish_call(recovered) + assert adapter.calls == {} + + @staticmethod + def test_legacy_api_check_uses_read_only_project_and_stats_requests(monkeypatch: pytest.MonkeyPatch) -> None: + sent = [] + + def post(url: str, **options: Unpack[_LegacyWeaveTransportRequestoptions]) -> httpx.Response: + assert options["auth"] == ("api", "fixture-key") + assert options["follow_redirects"] is False + sent.append(url) + if url == "https://api.wandb.ai/graphql": + return response( + url, + { + "data": { + "viewer": {"id": "viewer"}, + "project": {"id": "project-id", "name": "project", "entityName": "team", "readOnly": False}, + } + }, + ) + assert url == "https://trace.wandb.ai/calls/query_stats" + assert options["json"]["project_id"] == "team/project" + return response(url, {"count": 0}) + + monkeypatch.setattr(ssrf_proxy, "post", post) + adapter = WeaveDataTrace(_legacy_weave_transport_config()) + assert sent == [] + assert adapter.api_check() + assert sent == ["https://api.wandb.ai/graphql", "https://trace.wandb.ai/calls/query_stats"] + assert adapter.calls == {} + + @staticmethod + @pytest.mark.parametrize( + ("sent_id", "ack_id"), + [ + ("1234567812345678123456781234abcd", "12345678-1234-5678-1234-56781234abcd"), + ("12345678-1234-5678-1234-56781234abcd", "1234567812345678123456781234ABCD"), + ("12345678-1234-5678-1234-56781234ABCD", "1234567812345678123456781234abcd"), + ], + ) + def test_uuid_equivalent_start_ack_preserves_sent_ids( + monkeypatch: pytest.MonkeyPatch, sent_id: str, ack_id: str + ) -> None: + sent: list[dict[str, object]] = [] + + def send(_method: str, url: str, **options: Unpack[_LegacyWeaveTransportRequestoptions]) -> httpx.Response: + body = json.loads(options["content"]) + sent.append(body) + if "start" in body: + start = CallStartReq.model_validate(body).start + assert start.id == sent_id + assert start.trace_id == sent_id + return response(url, {"id": ack_id, "trace_id": ack_id}) + end = CallEndReq.model_validate(body).end + assert end.id == sent_id + return response(url, {}) + + monkeypatch.setattr(ssrf_proxy, "make_request", send) + adapter = WeaveDataTrace(_legacy_weave_transport_config()) + data = call(sent_id) + adapter.start_call(data) + assert adapter.calls[sent_id]["trace_id"] == sent_id + adapter.finish_call(data) + assert len(sent) == 2 + assert adapter.calls == {} + + @staticmethod + @pytest.mark.parametrize("field", ["id", "trace_id"]) + @pytest.mark.parametrize( + ("sent_id", "wrong_ack"), + [ + ("1234567812345678123456781234abcd", "12345678-1234-5678-1234-56781234abce"), + ("legacy-correlation", "LEGACY-CORRELATION"), + ("1234567812345678123456781234abcd", "not-a-uuid"), + ], + ) + def test_start_ack_rejects_different_uuid_or_opaque_identifiers( + monkeypatch: pytest.MonkeyPatch, field: str, sent_id: str, wrong_ack: str + ) -> None: + sent = [] + + def send(_method: str, url: str, **options: Unpack[_LegacyWeaveTransportRequestoptions]) -> httpx.Response: + start = CallStartReq.model_validate_json(options["content"]).start + sent.append(start) + ack = {"id": start.id, "trace_id": start.trace_id} + ack[field] = wrong_ack + return response(url, ack) + + monkeypatch.setattr(ssrf_proxy, "make_request", send) + adapter = WeaveDataTrace(_legacy_weave_transport_config()) + with pytest.raises(TraceExportUncertainError, match="did not acknowledge the submitted identifiers"): + adapter.start_call(call(sent_id)) + assert len(sent) == 1 + assert adapter.calls == {} + + +"""Real SDK schema/serialization boundaries with HTTP transport substituted. + +These tests do not claim acceptance by a live third-party account. +""" + + +class TransportOptions(TypedDict, total=False): + content: bytes | str + headers: Mapping[str, str] + auth: tuple[str, str] + max_retries: int + + +def _response(url: str, *, body: object = None, content: bytes | None = None, status: int = 200) -> httpx.Response: + request = httpx.Request("POST", url) + if content is not None: + return httpx.Response(status, request=request, content=content) + return httpx.Response(status, request=request, json=body or {}) + + +def _tool_definitions() -> list[dict[str, object]]: + return [ + { + "type": "function", + "name": "calculate", + "description": "TOOL_CONTENT_CANARY", + "parameters": { + "type": "object", + "properties": {"value": {"type": "integer"}}, + "required": ["value"], + "additionalProperties": False, + "$comment": "Preserve schema extension fields", + }, + "strict": True, + }, + {"type": "function", "function": {"name": "already_nested", "description": "Nested definition"}}, + ] + + +class TestNativeWeaveMlflowCases: + @staticmethod + @pytest.fixture + def trace() -> CanonicalTrace: + start = datetime(2026, 1, 1, tzinfo=UTC) + root = CanonicalSpan( + id="0000000000000001", + parent_id=None, + name="Agent", + kind=CanonicalSpanKind.AGENT, + start_time=start, + end_time=start + timedelta(seconds=2), + inputs={"query": "hello"}, + outputs="answer", + status=CanonicalSpanStatus.OK, + metadata={"total_price": 999, "total_tokens": 999}, + ) + child = CanonicalSpan( + id="0000000000000002", + parent_id=root.id, + name="Model", + kind=CanonicalSpanKind.LLM, + start_time=start, + end_time=start + timedelta(seconds=1), + inputs={"prompt": "hello"}, + outputs="answer", + status=CanonicalSpanStatus.OK, + metadata={ + "model_name": "test-model", + "prompt_tokens": 2, + "completion_tokens": 3, + "total_tokens": 5, + "total_price": "0.01", + "currency": "USD", + }, + ) + return CanonicalTrace( + trace_id="1234567890abcdef1234567890abcdef", + session_id="session-1", + root_span_id=root.id, + spans=(root, child), + ) + + @staticmethod + @pytest.fixture + def weave_started_calls(monkeypatch: pytest.MonkeyPatch) -> list[StartedCallSchemaForInsert]: + starts: list[StartedCallSchemaForInsert] = [] + + def send(_method: str, url: str, **kwargs: Unpack[TransportOptions]) -> httpx.Response: + body = json.loads(kwargs["content"]) + if start := body.get("start"): + started = StartedCallSchemaForInsert.model_validate(start) + starts.append(started) + return _response(url, body={"id": started.id, "trace_id": started.trace_id}) + return _response(url) + + monkeypatch.setattr(ssrf_proxy, "make_request", send) + return starts + + @staticmethod + @pytest.mark.parametrize("host", ["https://api.wandb.ai", "https://API.WANDB.AI:443"]) + def test_weave_sdk_exports_tree_session_usage_without_global_credentials( + monkeypatch: pytest.MonkeyPatch, trace: CanonicalTrace, host: str + ) -> None: + calls = [] + before = dict(os.environ) + + def send(_method: str, url: str, **kwargs: Unpack[TransportOptions]) -> httpx.Response: + body = json.loads(kwargs["content"]) + calls.append((_method, url, kwargs, body)) + start = body.get("start") + return _response(url, body={"id": start["id"], "trace_id": start["trace_id"]} if start else {}) + + monkeypatch.setattr(ssrf_proxy, "make_request", send) + result = export_trace( + "weave", {"api_key": "tenant-key", "entity": "team", "project": "demo", "host": host}, trace + ) + assert len(calls) == 4 + root, child = calls[0][3]["start"], calls[1][3]["start"] + assert child["parent_id"] == root["id"] + assert child["thread_id"] == root["thread_id"] == "session-1" + assert child["turn_id"] == root["turn_id"] == root["id"] + assert calls[2][3]["end"]["summary"]["usage"]["test-model"]["total_tokens"] == 5 + assert calls[3][3]["end"]["summary"] == {} + assert result.trace_url == f"https://wandb.ai/team/demo/r/call/{root['id']}" + assert result.span_urls[trace.spans[1].id].endswith(child["id"]) + assert all(call[2]["auth"] == ("api", "tenant-key") and call[2]["max_retries"] == 0 for call in calls) + assert dict(os.environ) == before + + @staticmethod + @pytest.mark.parametrize("root_kind", [CanonicalSpanKind.CHAIN, CanonicalSpanKind.AGENT]) + def test_weave_nested_agent_calls_form_one_outer_turn( + trace: CanonicalTrace, weave_started_calls: list[StartedCallSchemaForInsert], root_kind: CanonicalSpanKind + ) -> None: + root = trace.spans[0].model_copy(update={"kind": root_kind}) + agent = trace.spans[1].model_copy(update={"kind": CanonicalSpanKind.AGENT}) + model = trace.spans[1].model_copy(update={"id": "0000000000000003", "parent_id": agent.id}) + nested = trace.model_copy(update={"spans": (root, agent, model)}) + + export_trace("weave", {"api_key": "tenant-key", "entity": "team", "project": "demo"}, nested) + + outer, inner_agent, inner_model = weave_started_calls + assert inner_agent.parent_id == outer.id + assert inner_model.parent_id == inner_agent.id + assert all(call.thread_id == trace.session_id and call.turn_id == outer.id for call in weave_started_calls) + # Native Threads queries count only call.id == call.turn_id, including for + # Workflow and single-step preview outer roots with nested Agent details. + assert [call.id for call in weave_started_calls if call.id == call.turn_id] == [outer.id] + + @staticmethod + def test_weave_same_session_has_distinct_stable_turns( + trace: CanonicalTrace, weave_started_calls: list[StartedCallSchemaForInsert] + ) -> None: + config = {"api_key": "tenant-key", "entity": "team", "project": "demo"} + first = export_trace("weave", config, trace) + repeated = export_trace("weave", config, trace) + next_trace = trace.model_copy(update={"trace_id": "abcdef1234567890abcdef1234567890"}) + second = export_trace("weave", config, next_trace) + + first_calls, repeated_calls, second_calls = ( + weave_started_calls[:2], + weave_started_calls[2:4], + weave_started_calls[4:], + ) + assert repeated_calls == first_calls + assert repeated == first + assert second.trace_url != first.trace_url + assert all(call.thread_id == trace.session_id for call in weave_started_calls) + assert first_calls[0].turn_id == first_calls[0].id + assert second_calls[0].turn_id == second_calls[0].id + assert second_calls[0].turn_id != first_calls[0].turn_id + assert all(call.turn_id == first_calls[0].id for call in first_calls) + assert all(call.turn_id == second_calls[0].id for call in second_calls) + + @staticmethod + def test_weave_without_session_does_not_fabricate_a_thread_or_turn( + trace: CanonicalTrace, weave_started_calls: list[StartedCallSchemaForInsert] + ) -> None: + trace = trace.model_copy(update={"session_id": ""}) + + export_trace("weave", {"api_key": "tenant-key", "entity": "team", "project": "demo"}, trace) + + assert len(weave_started_calls) == len(trace.spans) + assert all(call.thread_id is None and call.turn_id is None for call in weave_started_calls) + + @staticmethod + def test_weave_resolves_optional_entity_using_only_supplied_credentials( + monkeypatch: pytest.MonkeyPatch, trace: CanonicalTrace + ) -> None: + def send(_method: str, url: str, **kwargs: Unpack[TransportOptions]) -> httpx.Response: + assert kwargs["auth"] == ("api", "tenant-key") + if url.endswith("/graphql"): + return _response(url, body={"data": {"viewer": {"entity": "default-team"}}}) + body = json.loads(kwargs["content"]) + if start := body.get("start"): + assert start["project_id"] == "default-team/demo" + return _response(url, body={"id": start["id"], "trace_id": start["trace_id"]}) + return _response(url) + + monkeypatch.setattr(ssrf_proxy, "make_request", send) + result = export_trace("weave", {"api_key": "tenant-key", "project": "demo"}, trace) + assert result.trace_url is not None + assert "/default-team/demo/" in result.trace_url + + @staticmethod + def test_weave_partial_tree_is_uncertain_and_not_retried( + monkeypatch: pytest.MonkeyPatch, trace: CanonicalTrace + ) -> None: + calls = [] + + def send(_method: str, url: str, **kwargs: Unpack[TransportOptions]) -> httpx.Response: + calls.append(url) + if len(calls) > 1: + return _response(url, status=503) + start = json.loads(kwargs["content"])["start"] + return _response(url, body={"id": start["id"], "trace_id": start["trace_id"]}) + + monkeypatch.setattr(ssrf_proxy, "make_request", send) + with pytest.raises(TraceExportUncertainError, match="partially accepted"): + export_trace("weave", {"api_key": "secret", "entity": "team", "project": "demo"}, trace) + assert len(calls) == 2 + + @staticmethod + @pytest.mark.parametrize("tracking_uri", ["https://mlflow.example", "https://mlflow.example/proxy/mlflow"]) + def test_mlflow_sdk_span_roundtrip_and_credential_isolation( + monkeypatch: pytest.MonkeyPatch, trace: CanonicalTrace, tracking_uri: str + ) -> None: + calls = [] + before = dict(os.environ) + + def send(_method: str, url: str, **kwargs: Unpack[TransportOptions]) -> httpx.Response: + calls.append((url, kwargs)) + return _response(url, content=b"") + + monkeypatch.setattr(ssrf_proxy, "make_request", send) + for username in ("tenant-a", "tenant-b"): + result = export_trace( + "mlflow", + { + "tracking_uri": tracking_uri, + "experiment_id": "42", + "username": username, + "password": "password", + }, + trace, + ) + assert result.trace_id == "tr-" + trace.trace_id + assert result.trace_url is not None + fragment_path, _, query = urlsplit(result.trace_url).fragment.partition("?") + assert result.trace_url.startswith(tracking_uri + "/#/") + assert fragment_path == "/experiments/42/traces" + assert parse_qs(query) == { + "traceId": [result.trace_id], + "selectedTraceId": [result.trace_id], + "selectedEvaluationId": [result.trace_id], + } + for index, (url, options) in enumerate(calls): + assert url == tracking_uri + "/v1/traces" + assert options["headers"]["x-mlflow-experiment-id"] == "42" + assert ( + options["headers"]["Authorization"] + == "Basic " + base64.b64encode(f"tenant-{'a' if index == 0 else 'b'}:password".encode()).decode() + ) + content = options["content"] + assert isinstance(content, bytes) + payload = ExportTraceServiceRequest.FromString(content) + spans = [ + _native_weave_mlflow_Span.from_otel_proto(span) + for span in payload.resource_spans[0].scope_spans[0].spans + ] + assert spans[1].parent_id == spans[0].span_id + assert spans[1].get_attribute(SpanAttributeKey.CHAT_USAGE) == { + "input_tokens": 2, + "output_tokens": 3, + "total_tokens": 5, + } + assert spans[1].get_attribute(SpanAttributeKey.LLM_COST) == {"total_cost": 0.01} + assert spans[0].get_attribute(SpanAttributeKey.CHAT_USAGE) is None + assert spans[0].get_attribute(SpanAttributeKey.SESSION_ID) == "session-1" + assert dict(os.environ) == before + + @staticmethod + @pytest.mark.parametrize("serialized", [False, True]) + def test_mlflow_tool_definitions_have_native_chat_schema_without_losing_source( + monkeypatch: pytest.MonkeyPatch, trace: CanonicalTrace, serialized: bool + ) -> None: + definitions = _tool_definitions() + source = json.dumps(definitions) if serialized else definitions + trace.spans[1].metadata["gen_ai.tool.definitions"] = source + before = deepcopy(trace.model_dump()) + native: list[_native_weave_mlflow_Span] = [] + + def send(_method: str, url: str, **options: Unpack[TransportOptions]) -> httpx.Response: + content = options["content"] + assert isinstance(content, bytes) + payload = ExportTraceServiceRequest.FromString(content) + native.extend( + _native_weave_mlflow_Span.from_otel_proto(item) + for item in payload.resource_spans[0].scope_spans[0].spans + ) + return _response(url, content=b"") + + monkeypatch.setattr(ssrf_proxy, "make_request", send) + export_trace("mlflow", {"tracking_uri": "https://mlflow.example", "experiment_id": "42"}, trace) + tools = native[1].get_attribute(SpanAttributeKey.CHAT_TOOLS) + assert tools == [ + {"type": "function", "function": {key: value for key, value in definitions[0].items() if key != "type"}}, + definitions[1], + ] + assert isinstance(tools, list) + for tool in tools: + ChatTool.model_validate(tool, strict=True) + assert native[1].get_attribute("gen_ai.tool.definitions") == source + stored = translate_span_when_storing(native[1]) + assert json.loads(stored["attributes"][SpanAttributeKey.CHAT_TOOLS]) == tools + assert trace.model_dump() == before + + @staticmethod + def test_mlflow_redacted_trace_does_not_recreate_tool_definition_content( + monkeypatch: pytest.MonkeyPatch, trace: CanonicalTrace + ) -> None: + trace.spans[1].metadata["gen_ai.tool.definitions"] = _tool_definitions() + trace.spans[1].metadata[SpanAttributeKey.CHAT_TOOLS] = [ + {"type": "function", "function": {"name": "TOOL_CONTENT_CANARY"}} + ] + redacted = remap_outer_trace(trace, trace_id=trace.trace_id, capture_content=False, agent_message=False) + native: list[_native_weave_mlflow_Span] = [] + + def send(_method: str, url: str, **options: Unpack[TransportOptions]) -> httpx.Response: + content = options["content"] + assert isinstance(content, bytes) + assert b"TOOL_CONTENT_CANARY" not in content + payload = ExportTraceServiceRequest.FromString(content) + native.extend( + _native_weave_mlflow_Span.from_otel_proto(item) + for item in payload.resource_spans[0].scope_spans[0].spans + ) + return _response(url, content=b"") + + monkeypatch.setattr(ssrf_proxy, "make_request", send) + export_trace("mlflow", {"tracking_uri": "https://mlflow.example", "experiment_id": "42"}, redacted) + for span in native: + assert span.get_attribute(SpanAttributeKey.CHAT_TOOLS) is None + assert span.get_attribute("gen_ai.tool.definitions") is None + assert span.inputs is None + assert span.outputs is None + + @staticmethod + @pytest.mark.parametrize("definitions", [[None, 12, {"type": "other"}, {"type": "function"}], "invalid-json"]) + def test_mlflow_unsupported_tool_shapes_remain_in_metadata_without_fake_functions( + trace: CanonicalTrace, definitions: object + ) -> None: + from dify_trace_mlflow.agent_trace import _spans + + trace.spans[1].metadata["gen_ai.tool.definitions"] = definitions + before = deepcopy(trace.model_dump()) + result = _spans(trace, "42", "tr-" + trace.trace_id) + assert result[1].get_attribute("gen_ai.tool.definitions") == definitions + assert result[1].get_attribute(SpanAttributeKey.CHAT_TOOLS) == ([] if isinstance(definitions, list) else None) + assert trace.model_dump() == before + + @staticmethod + def test_mlflow_existing_native_chat_tools_take_precedence(trace: CanonicalTrace) -> None: + from dify_trace_mlflow.agent_trace import _spans + + native_tools = [{"type": "function", "function": {"name": "already_native"}}] + trace.spans[1].metadata[SpanAttributeKey.CHAT_TOOLS] = native_tools + trace.spans[1].metadata["gen_ai.tool.definitions"] = _tool_definitions() + result = _spans(trace, "42", "tr-" + trace.trace_id) + assert result[1].get_attribute(SpanAttributeKey.CHAT_TOOLS) == native_tools + + @staticmethod + def test_mlflow_preserves_json_schema_beyond_the_sdk_chat_tool_model(trace: CanonicalTrace) -> None: + from dify_trace_mlflow.agent_trace import _spans + + parameters = { + "$ref": "#/$defs/Arguments", + "$defs": { + "Arguments": { + "type": "object", + "properties": {"choice": {"enum": [None, {"mode": "auto"}]}}, + "additionalProperties": {"type": "string"}, + } + }, + } + source = [{"type": "function", "name": "schema_extension", "parameters": parameters}] + trace.spans[1].metadata["gen_ai.tool.definitions"] = source + result = _spans(trace, "42", "tr-" + trace.trace_id) + expected = [{"type": "function", "function": {"name": "schema_extension", "parameters": parameters}}] + assert result[1].get_attribute(SpanAttributeKey.CHAT_TOOLS) == expected + assert result[1].get_attribute("gen_ai.tool.definitions") == source + stored = translate_span_when_storing(result[1]) + assert json.loads(stored["attributes"][SpanAttributeKey.CHAT_TOOLS]) == expected + + @staticmethod + def test_mlflow_partial_otlp_ack_is_uncertain(monkeypatch: pytest.MonkeyPatch, trace: CanonicalTrace) -> None: + ack = ExportTraceServiceResponse() + ack.partial_success.rejected_spans = 1 + monkeypatch.setattr( + ssrf_proxy, "make_request", lambda _method, url, **_kwargs: _response(url, content=ack.SerializeToString()) + ) + with pytest.raises(TraceExportUncertainError, match="partially accepted"): + export_trace("mlflow", {"tracking_uri": "https://mlflow.example", "experiment_id": "42"}, trace) + + @staticmethod + @pytest.mark.parametrize( + "credential_type", ["AWS_PRESIGNED_URL", "GCP_SIGNED_URL", "AZURE_SAS_URI", "AZURE_ADLS_GEN2_SAS_URI"] + ) + @pytest.mark.parametrize("oauth", [False, True]) + def test_databricks_sdk_v3_metadata_and_native_artifact( + monkeypatch: pytest.MonkeyPatch, trace: CanonicalTrace, credential_type: str, oauth: bool + ) -> None: + calls: list[tuple[str, str, TransportOptions]] = [] + uploaded: list[TraceData] = [] + before = dict(os.environ) + server_trace_id = "tr-abcdef0123456789abcdef0123456789" + definitions = _tool_definitions() + trace.spans[1].metadata["gen_ai.tool.definitions"] = definitions + + def send(_method: str, url: str, **kwargs: Unpack[TransportOptions]) -> httpx.Response: + calls.append((_method, url, kwargs)) + if url.endswith("/oidc/v1/token"): + assert kwargs["auth"] == ("client", "client-secret") + return _response(url, body={"access_token": "oauth-token"}) + if url.endswith("/api/3.0/mlflow/traces"): + assert kwargs["headers"]["Authorization"] == ("Bearer oauth-token" if oauth else "Bearer pat-token") + body = json.loads(kwargs["content"]) + assert body["trace"]["trace_info"]["trace_metadata"]["mlflow.trace.session"] == "session-1" + body["trace"]["trace_info"]["trace_id"] = server_trace_id + return _response(url, body=body) + if url.endswith("/credentials-for-data-upload"): + assert f"/traces/{server_trace_id}/" in url + return _response( + url, + body={ + "credential_info": { + "signed_uri": "https://storage.example/trace?sig=fake", + "type": credential_type, + "headers": [], + } + }, + ) + assert "Authorization" not in kwargs["headers"] + if kwargs.get("content"): + uploaded.append(TraceData.from_dict(json.loads(kwargs["content"]))) + return _response(url) + + monkeypatch.setattr(ssrf_proxy, "make_request", send) + config = {"host": "https://databricks.example", "experiment_id": "42"} + config.update( + {"client_id": "client", "client_secret": "client-secret"} + if oauth + else {"personal_access_token": "pat-token"} + ) + result = export_trace("databricks", config, trace) + assert result.trace_id == server_trace_id + assert result.trace_url == ( + f"https://databricks.example/ml/experiments/42/traces?selectedEvaluationId={server_trace_id}" + ) + data = uploaded[0] + assert len(data.spans) == 2 + assert all(span.trace_id == server_trace_id for span in data.spans) + assert data.spans[1].parent_id == data.spans[0].span_id + usage = data.spans[1].get_attribute(SpanAttributeKey.CHAT_USAGE) + assert usage is not None + assert usage["total_tokens"] == 5 + assert data.spans[1].get_attribute(SpanAttributeKey.LLM_COST) == {"total_cost": 0.01} + tools = data.spans[1].get_attribute(SpanAttributeKey.CHAT_TOOLS) + assert isinstance(tools, list) + assert len(tools) == 2 + for tool in tools: + ChatTool.model_validate(tool, strict=True) + assert data.spans[1].get_attribute("gen_ai.tool.definitions") == definitions + assert dict(os.environ) == before + assert len(calls) == 3 + int(oauth) + int(credential_type == "AZURE_ADLS_GEN2_SAS_URI") + + @staticmethod + def test_databricks_partial_metadata_delivery_is_uncertain( + monkeypatch: pytest.MonkeyPatch, trace: CanonicalTrace + ) -> None: + def send(_method: str, url: str, **kwargs: Unpack[TransportOptions]) -> httpx.Response: + if url.endswith("/api/3.0/mlflow/traces"): + return _response(url, body=json.loads(kwargs["content"])) + return _response(url, status=503) + + monkeypatch.setattr(ssrf_proxy, "make_request", send) + with pytest.raises(TraceExportUncertainError, match="metadata"): + export_trace( + "databricks", + {"host": "https://databricks.example", "experiment_id": "42", "personal_access_token": "secret"}, + trace, + ) + + @staticmethod + @pytest.mark.parametrize("provider", ["weave", "databricks"]) + def test_success_status_with_unreadable_ack_is_uncertain( + monkeypatch: pytest.MonkeyPatch, trace: CanonicalTrace, provider: str + ) -> None: + monkeypatch.setattr( + ssrf_proxy, "make_request", lambda _method, url, **_kwargs: _response(url, content=b"not-json") + ) + config = { + "api_key": "secret", + "entity": "team", + "project": "demo", + "host": "https://example.com", + "experiment_id": "42", + "personal_access_token": "secret", + } + with pytest.raises(TraceExportUncertainError, match="unreadable"): + export_trace(provider, config, trace) + + +"""Exercise anonymous Phoenix auth at the HTTP/1.1 and legacy SDK boundaries.""" + + +class TestPhoenixExportAuthCases: + @staticmethod + @pytest.fixture + def trace() -> CanonicalTrace: + start = datetime(2026, 10, 20, tzinfo=UTC) + return CanonicalTrace( + trace_id="1" * 32, + session_id="phoenix-auth-session", + root_span_id="2" * 16, + spans=( + CanonicalSpan( + id="2" * 16, + parent_id=None, + name="Phoenix auth regression", + kind=CanonicalSpanKind.AGENT, + start_time=start, + end_time=start + timedelta(seconds=1), + status=CanonicalSpanStatus.OK, + ), + ), + ) + + @staticmethod + @pytest.fixture + def http11_wire(monkeypatch: pytest.MonkeyPatch) -> Iterator[list[bytes]]: + """Keep real HTTPX/httpcore/h11 serialization, replacing only socket I/O.""" + chunks: list[bytes] = [] + + def write(_stream: httpcore.MockStream, buffer: bytes, timeout: float | None = None) -> None: + del timeout + chunks.append(buffer) + + monkeypatch.setattr(httpcore.MockStream, "write", write) + transport = httpx.HTTPTransport() + transport.close() + transport._pool = httpcore.ConnectionPool( + network_backend=httpcore.MockBackend([b"HTTP/1.1 200 OK\r\nContent-Length: 0\r\nConnection: close\r\n\r\n"]) + ) + with httpx.Client(transport=transport, trust_env=False) as client: + monkeypatch.setattr(ssrf_proxy, "_get_ssrf_client", lambda _verify: client) + yield chunks + + @staticmethod + def test_managed_exports_serialize_optional_auth_without_cross_destination_credentials( + monkeypatch: pytest.MonkeyPatch, trace: CanonicalTrace, http11_wire: list[bytes] + ) -> None: + monkeypatch.setenv( + "OTEL_EXPORTER_OTLP_HEADERS", "authorization=Bearer%20ambient-general,api_key=ambient-general" + ) + monkeypatch.setenv( + "OTEL_EXPORTER_OTLP_TRACES_HEADERS", "authorization=Bearer%20ambient-traces,api_key=ambient-traces" + ) + + for index, key in enumerate(("first-synthetic-key", "", None, "second-synthetic-key", None, "")): + previous_chunks = len(http11_wire) + result = export_trace( + "phoenix", + {"endpoint": "http://phoenix.test/proxy", "api_key": key, "project": f"project-{index}"}, + trace, + ) + assert result.status == "accepted" + assert result.trace_id == trace.trace_id + head, separator, body = b"".join(http11_wire[previous_chunks:]).partition(b"\r\n\r\n") + assert separator + lines = head.split(b"\r\n") + assert lines[0] == b"POST /proxy/v1/traces HTTP/1.1" + header_values: list[tuple[bytes, bytes]] = [] + for line in lines[1:]: + name, value = line.split(b": ", 1) + header_values.append((name, value)) + headers = httpx.Headers(header_values) + assert headers["Content-Type"] == "application/x-protobuf" + if key: + assert headers["api_key"] == key + assert headers["authorization"] == f"Bearer {key}" + else: + assert "api_key" not in headers + assert "authorization" not in headers + payload = ExportTraceServiceRequest.FromString(body) + resource = payload.resource_spans[0] + assert {item.key: item.value.string_value for item in resource.resource.attributes}[ + "openinference.project.name" + ] == f"project-{index}" + assert resource.scope_spans[0].spans[0].trace_id.hex() == trace.trace_id + + @staticmethod + def test_legacy_sdk_prepares_anonymous_requests_without_ambient_or_previous_auth( + monkeypatch: pytest.MonkeyPatch, + ) -> None: + monkeypatch.setenv( + "OTEL_EXPORTER_OTLP_HEADERS", "authorization=Bearer%20ambient-general,api_key=ambient-general" + ) + monkeypatch.setenv( + "OTEL_EXPORTER_OTLP_TRACES_HEADERS", "authorization=Bearer%20ambient-traces,api_key=ambient-traces" + ) + monkeypatch.delenv("OTEL_EXPORTER_OTLP_COMPRESSION", raising=False) + monkeypatch.delenv("OTEL_EXPORTER_OTLP_TRACES_COMPRESSION", raising=False) + prepared: list[requests.PreparedRequest] = [] + + def send(_session: requests.Session, request: requests.PreparedRequest, **_kwargs: object) -> requests.Response: + prepared.append(request) + response = requests.Response() + response.status_code = 200 + response._content = b"" + response.request = request + return response + + monkeypatch.setattr(requests.Session, "send", send) + for index, key in enumerate(("first-synthetic-key", "", None, "second-synthetic-key", None, "")): + tracer, processor = setup_tracer( + PhoenixConfig(endpoint="http://phoenix.test/proxy", api_key=key, project=f"project-{index}") + ) + try: + tracer.start_span("Legacy Phoenix auth regression").end() + finally: + processor.shutdown() + assert len(prepared) == index + 1 + request = prepared[-1] + assert request.method == "POST" + assert request.url == "http://phoenix.test/proxy/v1/traces" + assert request.headers["Content-Type"] == "application/x-protobuf" + if key: + assert request.headers["api_key"] == key + assert request.headers["authorization"] == f"Bearer {key}" + else: + assert "api_key" not in request.headers + assert "authorization" not in request.headers + assert isinstance(request.body, bytes) + payload = ExportTraceServiceRequest.FromString(request.body) + assert {item.key: item.value.string_value for item in payload.resource_spans[0].resource.attributes}[ + "openinference.project.name" + ] == f"project-{index}" + + +"""Exercise real vendor serializers and OTLP stubs at their network boundary.""" class HTTPArguments(TypedDict): @@ -35,83 +2140,12 @@ class HTTPArguments(TypedDict): type HTTPCall = tuple[str, str, HTTPArguments] -@pytest.fixture -def trace() -> CanonicalTrace: - start = datetime(2026, 10, 15, 1, tzinfo=UTC) - return CanonicalTrace( - trace_id="1" * 32, - session_id="conversation-1", - root_span_id="1" * 16, - spans=( - CanonicalSpan( - id="1" * 16, - parent_id=None, - name="Agent turn", - kind=CanonicalSpanKind.AGENT, - start_time=start, - end_time=start + timedelta(seconds=3), - status=CanonicalSpanStatus.OK, - inputs={"question": "hi"}, - outputs={"answer": "done"}, - metadata={"app_id": "app-one", "tenant_id": "tenant-one", "total_tokens": 12, "total_price": "0.03"}, - ), - CanonicalSpan( - id="2" * 16, - parent_id="1" * 16, - name="model call", - kind=CanonicalSpanKind.LLM, - start_time=start, - end_time=start + timedelta(seconds=1), - status=CanonicalSpanStatus.OK, - inputs=[{"role": "user", "content": "hi"}], - outputs="done", - metadata={ - "model": "gpt-4o-mini", - "model_provider": "openai", - "prompt_tokens": 8, - "completion_tokens": 4, - "total_tokens": 12, - "total_price": "0.03", - "currency": "USD", - "ttft": 0.2, - }, - ), - CanonicalSpan( - id="3" * 16, - parent_id="1" * 16, - name="search", - kind=CanonicalSpanKind.TOOL, - start_time=start + timedelta(seconds=1), - end_time=start + timedelta(seconds=2), - status=CanonicalSpanStatus.ERROR, - inputs={"q": "x"}, - outputs=None, - error="tool failed", - ), - ), - ) - - -@pytest.fixture -def http_calls(monkeypatch: pytest.MonkeyPatch) -> list[HTTPCall]: - from core.helper import ssrf_proxy - - calls: list[HTTPCall] = [] - - def send(method: str, url: str, **kwargs: object) -> httpx.Response: - calls.append((method, url, cast(HTTPArguments, kwargs))) - return httpx.Response(200, content=b"", request=httpx.Request(method, url)) - - monkeypatch.setattr(ssrf_proxy, "make_request", send) - return calls - - -def spans_from_call(call: HTTPCall) -> list[Span]: +def spans_from_call(call: HTTPCall) -> list[_provider_dispatch_Span]: request = ExportTraceServiceRequest.FromString(call[2]["content"]) return list(request.resource_spans[0].scope_spans[0].spans) -def attributes(span: Span) -> dict[str, object]: +def attributes(span: _provider_dispatch_Span) -> dict[str, object]: result: dict[str, object] = {} for item in span.attributes: kind = item.value.WhichOneof("value") @@ -128,352 +2162,6 @@ def attributes(span: Span) -> dict[str, object]: return result -def test_langfuse_v4_tree_and_native_usage(trace: CanonicalTrace, http_calls: list[HTTPCall]) -> None: - config = { - "host": "https://us.cloud.langfuse.com", - "public_key": "pk-test", - "secret_key": "sk-test", - "project_url": "https://us.cloud.langfuse.com/project/example", - } - result = export_trace("langfuse", config, trace) - assert result.status == "accepted" - assert result.trace_id == trace.trace_id - assert len(http_calls) == 1 - _, url, args = http_calls[0] - assert url.endswith("/api/public/otel/v1/traces") - assert args["headers"]["x-langfuse-ingestion-version"] == "4" - assert args["max_retries"] == 0 - assert args["follow_redirects"] is False - spans = spans_from_call(http_calls[0]) - assert [span.span_id.hex() for span in spans] == [s.id for s in trace.spans] - assert spans[1].parent_span_id == spans[0].span_id - assert spans[2].status.code == 2 - root, llm = attributes(spans[0]), attributes(spans[1]) - assert root["langfuse.observation.type"] == "agent" - assert "langfuse.observation.usage_details" not in root - usage_details, cost_details = llm["langfuse.observation.usage_details"], llm["langfuse.observation.cost_details"] - assert isinstance(usage_details, str) - assert isinstance(cost_details, str) - assert json.loads(usage_details) == {"input": 8, "output": 4, "total": 12} - assert json.loads(cost_details) == {"total": 0.03} - assert all(attributes(span)["langfuse.session.id"] == trace.session_id for span in spans) - assert result.span_urls[trace.spans[1].id].endswith(trace.spans[1].id) - - -@pytest.mark.parametrize( - ("provider", "config", "suffix", "kind_key", "kind"), - [ - ( - "phoenix", - {"endpoint": "https://phoenix.test/proxy", "api_key": "key", "project": "one"}, - "/proxy/v1/traces", - "openinference.span.kind", - "AGENT", - ), - ( - "aliyun", - {"endpoint": "https://region.log.aliyuncs.com", "license_key": "license", "app_name": "one"}, - "/adapt_license/api/v1/traces", - "gen_ai.span.kind", - "AGENT", - ), - ], -) -def test_other_http_providers_have_native_semantics( - trace: CanonicalTrace, - http_calls: list[HTTPCall], - provider: str, - config: dict[str, str], - suffix: str, - kind_key: str, - kind: str, -) -> None: - result = export_trace(provider, config, trace) - assert result.status == "accepted" - assert http_calls[0][1].endswith(suffix) - spans = spans_from_call(http_calls[0]) - assert attributes(spans[0])[kind_key] == kind - token_key = "llm.token_count.prompt" if provider == "phoenix" else "gen_ai.usage.input_tokens" - assert attributes(spans[1])[token_key] == 8 - assert token_key not in attributes(spans[0]) - - -@pytest.mark.parametrize( - ("endpoint_url", "suffix"), - [ - ("https://log.aliyuncs.com", "api/v1/traces"), - ("https://cn-hangzhou.log.aliyuncs.com", "api/v1/traces"), - ("https://project.cn-heyuan.log.aliyuncs.com/apm/trace/opentelemetry", "api/v1/traces"), - ("https://project.cn-heyuan-intranet.log.aliyuncs.com/apm/trace/opentelemetry", "api/v1/traces"), - ("https://cn-shanghai-share.log.aliyuncs.com", "api/v1/traces"), - ("https://cn-beijing.dualstack.log.aliyuncs.com", "api/v1/traces"), - ("https://cn-hangzhou.log.aliyuncs.com.:443", "api/v1/traces"), - ("https://tracing-analysis-dc-hz.aliyuncs.com", "api/otlp/traces"), - ("https://collector.example.test:4318/custom-prefix", "api/otlp/traces"), - ("https://evil-log.aliyuncs.com.evil.test", "api/otlp/traces"), - ("https://log.aliyuncs.com.evil.test", "api/otlp/traces"), - ("https://evil-log.aliyuncs.com", "api/otlp/traces"), - ("https://collector.example.test/log.aliyuncs.com", "api/otlp/traces"), - ], -) -def test_aliyun_endpoint_family_uses_dns_labels_and_preserves_custom_destinations( - trace: CanonicalTrace, http_calls: list[HTTPCall], endpoint_url: str, suffix: str -) -> None: - result = export_trace( - "aliyun", {"endpoint": endpoint_url, "license_key": "synthetic-license", "app_name": "test-app"}, trace - ) - assert result.status == "accepted" - assert len(http_calls) == 1 - assert http_calls[0][1] == f"{endpoint_url}/adapt_synthetic-license/{suffix}" - spans = spans_from_call(http_calls[0]) - assert [span.span_id.hex() for span in spans] == [span.id for span in trace.spans] - assert spans[1].parent_span_id == spans[0].span_id - - -def test_aliyun_raw_license_is_encoded_once_at_the_real_otlp_boundary( - trace: CanonicalTrace, http_calls: list[HTTPCall] -) -> None: - result = export_trace( - "aliyun", - {"endpoint": "https://COLLECTOR.test:8443/prefix/", "license_key": "raw%2F@a+b/c", "app_name": "test_app"}, - trace, - ) - assert result.status == "accepted" - assert len(http_calls) == 1 - assert http_calls[0][1] == "https://collector.test:8443/prefix/adapt_raw%252F@a%2Bb%2Fc/api/otlp/traces" - assert [span.span_id.hex() for span in spans_from_call(http_calls[0])] == [span.id for span in trace.spans] - - -def test_aliyun_complete_token_url_is_rejected_without_any_ingestion( - trace: CanonicalTrace, http_calls: list[HTTPCall] -) -> None: - with pytest.raises(TraceExportRejectedError, match="Aliyun tracing configuration is invalid"): - export_trace( - "aliyun", - {"endpoint": "https://collector.test/adapt_private/api/otlp/traces", "license_key": "private"}, - trace, - ) - assert http_calls == [] - - -@pytest.mark.parametrize( - ("provider", "config", "header"), - [ - ( - "arize", - {"endpoint": "https://otlp.arize.com", "api_key": "key", "space_id": "space", "project": "one"}, - ("space_id", "space"), - ), - ( - "tencent", - {"endpoint": "https://apm.test:4317", "token": "token", "service_name": "one"}, - ("authorization", "Bearer token"), - ), - ], -) -def test_grpc_real_stub_serialization( - trace: CanonicalTrace, - monkeypatch: pytest.MonkeyPatch, - provider: str, - config: dict[str, str], - header: tuple[str, str], -) -> None: - calls: list[tuple[ExportTraceServiceRequest, Sequence[tuple[str, str]]]] = [] - - class Channel: - def unary_unary( - self, - method: str, - request_serializer: Callable[[ExportTraceServiceRequest], bytes], - response_deserializer: Callable[[bytes], ExportTraceServiceResponse], - **_kwargs: object, - ) -> Callable[..., ExportTraceServiceResponse]: - assert method == "/opentelemetry.proto.collector.trace.v1.TraceService/Export" - - def invoke( - payload: ExportTraceServiceRequest, metadata: Sequence[tuple[str, str]], timeout: float - ) -> ExportTraceServiceResponse: - assert timeout == 30 - calls.append((ExportTraceServiceRequest.FromString(request_serializer(payload)), metadata)) - return response_deserializer(ExportTraceServiceResponse().SerializeToString()) - - return invoke - - def close(self) -> None: - pass - - monkeypatch.setattr("grpc.secure_channel", lambda *_args, **_kwargs: Channel()) - result = export_trace(provider, config, trace) - assert result.status == "accepted" - assert header in calls[0][1] - spans = calls[0][0].resource_spans[0].scope_spans[0].spans - assert spans[1].parent_span_id == spans[0].span_id - - -def test_langsmith_real_sdk_uses_isolated_sync_runs(trace: CanonicalTrace, http_calls: list[HTTPCall]) -> None: - result = export_trace( - "langsmith", {"api_key": "ls-test", "endpoint": "https://langsmith.test", "project": "one"}, trace - ) - writes = [json.loads(call[2]["content"]) for call in http_calls if call[0] == "POST"] - assert len(writes) == 3 - root, llm, tool = writes - assert root["run_type"] == "chain" - assert llm["run_type"] == "llm" - assert llm["parent_run_id"] == root["id"] - assert llm["dotted_order"].startswith(root["dotted_order"]) - assert llm["input_tokens"] == 8 - assert llm["output_tokens"] == 4 - assert llm["total_cost"] == 0.03 - assert "input_tokens" not in root - assert all(write["extra"]["metadata"]["session_id"] == trace.session_id for write in writes) - assert tool["error"] == "tool failed" - assert UUID(result.trace_id) - - -def test_opik_real_sdk_preserves_tree_thread_and_native_usage( - trace: CanonicalTrace, http_calls: list[HTTPCall] -) -> None: - result = export_trace( - "opik", - {"api_key": "opik-test", "url": "https://opik.test/api", "workspace": "org-one", "project": "one"}, - trace, - ) - assert len(http_calls) == 2 - root = json.loads(http_calls[0][2]["content"]) - batch = json.loads(http_calls[1][2]["content"]) - spans = batch["spans"] - assert UUID(result.trace_id).version == 7 - assert root["thread_id"] == trace.session_id - assert spans[1]["parent_span_id"] == spans[0]["id"] - assert spans[1]["usage"]["prompt_tokens"] == 8 - assert spans[1]["total_estimated_cost"] == 0.03 - assert not spans[0].get("usage") - assert spans[2]["error_info"]["message"] == "tool failed" - assert http_calls[0][2]["headers"]["comet-workspace"] == "org-one" - - -@pytest.mark.parametrize( - ("value", "wire_value"), - [ - (None, None), - ("323", "323"), - ({"result": 323}, {"result": 323}), - ([{"result": 323}], [{"result": 323}]), - ([], []), - (323, "323"), - (0, "0"), - (-5, "-5"), - (3.5, "3.5"), - (True, "true"), - (False, "false"), - ([1, 2], "[1, 2]"), - (["one", "two"], '["one", "two"]'), - ([{"result": 1}, "text"], '[{"result": 1}, "text"]'), - ([[1], [2]], "[[1], [2]]"), - ], -) -def test_opik_real_sdk_preserves_json_content_across_trace_and_span_boundaries( - trace: CanonicalTrace, http_calls: list[HTTPCall], value: object, wire_value: object -) -> None: - content_trace = trace.model_copy( - update={"spans": tuple(span.model_copy(update={"inputs": value, "outputs": value}) for span in trace.spans)} - ) - result = export_trace( - "opik", - {"api_key": "synthetic-key", "url": "https://opik.test/api", "workspace": "team", "project": "one"}, - content_trace, - ) - assert result.status == "accepted" - assert len(http_calls) == 2 - root = json.loads(http_calls[0][2]["content"]) - spans = json.loads(http_calls[1][2]["content"])["spans"] - for payload in [root, *spans]: - assert payload.get("input") == wire_value - assert payload.get("output") == wire_value - assert spans[1]["parent_span_id"] == spans[0]["id"] - assert spans[2]["parent_span_id"] == spans[0]["id"] - assert root["thread_id"] == trace.session_id - assert spans[1]["usage"]["total_tokens"] == 12 - assert not spans[0].get("usage") - assert not spans[2].get("usage") - - -@pytest.mark.parametrize("capture_content", [True, False]) -@pytest.mark.parametrize("root_error", [True, False]) -def test_opik_trace_error_matches_sanitized_root_and_not_handled_tool_failure( - trace: CanonicalTrace, http_calls: list[HTTPCall], capture_content: bool, root_error: bool -) -> None: - root = trace.spans[0].model_copy( - update={ - "status": CanonicalSpanStatus.ERROR if root_error else CanonicalSpanStatus.OK, - "error": "private-model-error" if root_error else None, - "inputs": {"query": "private-error-input"}, - } - ) - source = trace.model_copy(update={"spans": (root, *trace.spans[1:])}) - canonical = remap_outer_trace(source, trace_id=trace.trace_id, capture_content=capture_content, agent_message=True) - export_trace( - "opik", - {"api_key": "synthetic-key", "url": "https://opik.test/api", "workspace": "team", "project": "one"}, - canonical, - ) - root_payload = json.loads(http_calls[0][2]["content"]) - spans = json.loads(http_calls[1][2]["content"])["spans"] - assert spans[2]["error_info"] is not None - if root_error: - expected = { - "exception_type": "DifyExecutionError", - "message": "private-model-error" if capture_content else "Application execution failed", - "traceback": "", - } - assert root_payload["error_info"] == spans[0]["error_info"] == expected - else: - assert root_payload.get("error_info") is None - assert spans[0].get("error_info") is None - if not capture_content: - assert all(b"private-model-error" not in call[2]["content"] for call in http_calls) - assert all(b"private-error-input" not in call[2]["content"] for call in http_calls) - - -@pytest.mark.parametrize( - ("status", "error"), - [(401, TraceExportRejectedError), (429, RetryableTraceDispatchError), (503, RetryableTraceDispatchError)], -) -def test_explicit_http_rejections_are_classified_without_response_body( - trace: CanonicalTrace, monkeypatch: pytest.MonkeyPatch, status: int, error: type[Exception] -) -> None: - from core.helper import ssrf_proxy - - monkeypatch.setattr( - ssrf_proxy, - "make_request", - lambda method, url, **_kwargs: httpx.Response(status, text="secret-token", request=httpx.Request(method, url)), - ) - with pytest.raises(error) as caught: - export_trace( - "langfuse", {"host": "https://langfuse.test", "public_key": "pk", "secret_key": "secret-token"}, trace - ) - assert "secret-token" not in str(caught.value) - - -def test_partial_otlp_acceptance_requires_reconciliation( - trace: CanonicalTrace, monkeypatch: pytest.MonkeyPatch -) -> None: - from core.helper import ssrf_proxy - - response = ExportTraceServiceResponse() - response.partial_success.rejected_spans = 1 - monkeypatch.setattr( - ssrf_proxy, - "make_request", - lambda method, url, **_kwargs: httpx.Response( - 200, content=response.SerializeToString(), request=httpx.Request(method, url) - ), - ) - with pytest.raises(TraceExportUncertainError): - export_trace("langfuse", {"host": "https://langfuse.test", "public_key": "pk", "secret_key": "sk"}, trace) - - def _langfuse_queue_job_ack() -> dict[str, object]: # Synthetic values only. Real acknowledgements contain auth context and must # never be logged or persisted as acceptance-test evidence. @@ -504,212 +2192,650 @@ def _replace_ack_value(ack: dict[str, object], path: tuple[str, ...], value: obj ack[path[-1]] = value -def test_langfuse_v4_queue_ack_confirms_ingestion_acceptance_only( - trace: CanonicalTrace, - monkeypatch: pytest.MonkeyPatch, - caplog: pytest.LogCaptureFixture, - capsys: pytest.CaptureFixture[str], -) -> None: - from core.helper import ssrf_proxy - - monkeypatch.setattr( - ssrf_proxy, - "make_request", - lambda method, url, **_kwargs: httpx.Response( - 200, json=_langfuse_queue_job_ack(), request=httpx.Request(method, url) - ), - ) - result = export_trace("langfuse", {"host": "https://langfuse.test", "public_key": "pk", "secret_key": "sk"}, trace) - assert result.status == "accepted" - assert result.trace_id == trace.trace_id - assert result.trace_url is None - captured = capsys.readouterr() - evidence = caplog.text + captured.out + captured.err + repr(result) - assert "never-log-public-key" not in evidence - assert "never-log-project" not in evidence - assert "authCheck" not in evidence - - -@pytest.mark.parametrize( - "scenario", - [ - "wrong_name", - "wrong_nested_name", - "invalid_id", - "invalid_timestamp", - "invalid_auth", - "missing_scope", - "missing_file", - "partial_success", - "error", - "failed_job", - "unknown", - "null", - "array", - ], -) -def test_langfuse_ambiguous_json_ack_remains_uncertain( - trace: CanonicalTrace, - monkeypatch: pytest.MonkeyPatch, - caplog: pytest.LogCaptureFixture, - capsys: pytest.CaptureFixture[str], - scenario: str, -) -> None: - from core.helper import ssrf_proxy - - ack = _langfuse_queue_job_ack() - body: object = ack - if scenario == "wrong_name": - ack["name"] = "different-job" - elif scenario == "wrong_nested_name": - _replace_ack_value(ack, ("data", "name"), "different-job") - elif scenario == "invalid_id": - _replace_ack_value(ack, ("data", "id"), "not-a-uuid") - elif scenario == "invalid_timestamp": - _replace_ack_value(ack, ("data", "timestamp"), "2026-10-15T01:00:00") - elif scenario == "invalid_auth": - _replace_ack_value(ack, ("data", "payload", "authCheck", "validKey"), False) - elif scenario == "missing_scope": - _replace_ack_value(ack, ("data", "payload", "authCheck", "scope"), dict[str, object]()) - elif scenario == "missing_file": - _replace_ack_value(ack, ("data", "payload", "data", "fileKey"), "") - elif scenario == "partial_success": - ack["partialSuccess"] = {"rejectedSpans": "1"} - elif scenario == "error": - ack["error"] = "never-log-public-key" - elif scenario == "failed_job": - ack["failedReason"] = "never-log-public-key" - elif scenario == "unknown": - body = {"accepted": True, "publicKey": "never-log-public-key"} - elif scenario == "null": - body = None - elif scenario == "array": - body = list[object]() - monkeypatch.setattr( - ssrf_proxy, - "make_request", - lambda method, url, **_kwargs: httpx.Response( - 200, - content=json.dumps(body).encode(), - headers={"Content-Type": "application/json"}, - request=httpx.Request(method, url), - ), - ) - with pytest.raises(TraceExportUncertainError) as caught: - export_trace("langfuse", {"host": "https://langfuse.test", "public_key": "pk", "secret_key": "sk"}, trace) - captured = capsys.readouterr() - evidence = caplog.text + captured.out + captured.err + str(caught.value) - assert "never-log-public-key" not in evidence - assert "never-log-project" not in evidence - assert "authCheck" not in evidence - - -@pytest.mark.parametrize( - ("provider", "version", "status"), [("phoenix", "4", 200), ("langfuse", "3", 200), ("langfuse", "4", 202)] -) -def test_langfuse_queue_ack_is_not_a_generic_otlp_success( - monkeypatch: pytest.MonkeyPatch, provider: str, version: str, status: int -) -> None: - from core.helper import ssrf_proxy - from services.tracing.otlp_export import export_http - - monkeypatch.setattr( - ssrf_proxy, - "make_request", - lambda method, url, **_kwargs: httpx.Response( - status, json=_langfuse_queue_job_ack(), request=httpx.Request(method, url) - ), - ) - with pytest.raises(TraceExportUncertainError): - export_http( - provider, - "https://collector.test/v1/traces", - {"x-langfuse-ingestion-version": version}, - ExportTraceServiceRequest(), +class TestProviderDispatchCases: + @staticmethod + @pytest.fixture + def trace() -> CanonicalTrace: + start = datetime(2026, 10, 15, 1, tzinfo=UTC) + return CanonicalTrace( + trace_id="1" * 32, + session_id="conversation-1", + root_span_id="1" * 16, + spans=( + CanonicalSpan( + id="1" * 16, + parent_id=None, + name="Agent turn", + kind=CanonicalSpanKind.AGENT, + start_time=start, + end_time=start + timedelta(seconds=3), + status=CanonicalSpanStatus.OK, + inputs={"question": "hi"}, + outputs={"answer": "done"}, + metadata={ + "app_id": "app-one", + "tenant_id": "tenant-one", + "total_tokens": 12, + "total_price": "0.03", + }, + ), + CanonicalSpan( + id="2" * 16, + parent_id="1" * 16, + name="model call", + kind=CanonicalSpanKind.LLM, + start_time=start, + end_time=start + timedelta(seconds=1), + status=CanonicalSpanStatus.OK, + inputs=[{"role": "user", "content": "hi"}], + outputs="done", + metadata={ + "model": "gpt-4o-mini", + "model_provider": "openai", + "prompt_tokens": 8, + "completion_tokens": 4, + "total_tokens": 12, + "total_price": "0.03", + "currency": "USD", + "ttft": 0.2, + }, + ), + CanonicalSpan( + id="3" * 16, + parent_id="1" * 16, + name="search", + kind=CanonicalSpanKind.TOOL, + start_time=start + timedelta(seconds=1), + end_time=start + timedelta(seconds=2), + status=CanonicalSpanStatus.ERROR, + inputs={"q": "x"}, + outputs=None, + error="tool failed", + ), + ), ) + @staticmethod + @pytest.fixture + def http_calls(monkeypatch: pytest.MonkeyPatch) -> list[HTTPCall]: + from core.helper import ssrf_proxy -@pytest.mark.parametrize("body", [{}, {"partialSuccess": {}}, {"partialSuccess": {"rejectedSpans": "1"}}]) -def test_langfuse_standard_otlp_json_ack_remains_supported( - trace: CanonicalTrace, monkeypatch: pytest.MonkeyPatch, body: dict[str, dict[str, str]] -) -> None: - from core.helper import ssrf_proxy + calls: list[HTTPCall] = [] - monkeypatch.setattr( - ssrf_proxy, - "make_request", - lambda method, url, **_kwargs: httpx.Response(200, json=body, request=httpx.Request(method, url)), + def send(method: str, url: str, **kwargs: object) -> httpx.Response: + calls.append((method, url, cast(HTTPArguments, kwargs))) + return httpx.Response(200, content=b"", request=httpx.Request(method, url)) + + monkeypatch.setattr(ssrf_proxy, "make_request", send) + return calls + + @staticmethod + def test_langfuse_v4_tree_and_native_usage(trace: CanonicalTrace, http_calls: list[HTTPCall]) -> None: + config = { + "host": "https://us.cloud.langfuse.com", + "public_key": "pk-test", + "secret_key": "sk-test", + "project_url": "https://us.cloud.langfuse.com/project/example", + } + result = export_trace("langfuse", config, trace) + assert result.status == "accepted" + assert result.trace_id == trace.trace_id + assert len(http_calls) == 1 + _, url, args = http_calls[0] + assert url.endswith("/api/public/otel/v1/traces") + assert args["headers"]["x-langfuse-ingestion-version"] == "4" + assert args["max_retries"] == 0 + assert args["follow_redirects"] is False + spans = spans_from_call(http_calls[0]) + assert [span.span_id.hex() for span in spans] == [s.id for s in trace.spans] + assert spans[1].parent_span_id == spans[0].span_id + assert spans[2].status.code == 2 + root, llm = attributes(spans[0]), attributes(spans[1]) + assert root["langfuse.observation.type"] == "agent" + assert "langfuse.observation.usage_details" not in root + usage_details, cost_details = ( + llm["langfuse.observation.usage_details"], + llm["langfuse.observation.cost_details"], + ) + assert isinstance(usage_details, str) + assert isinstance(cost_details, str) + assert json.loads(usage_details) == {"input": 8, "output": 4, "total": 12} + assert json.loads(cost_details) == {"total": 0.03} + assert all(attributes(span)["langfuse.session.id"] == trace.session_id for span in spans) + assert result.span_urls[trace.spans[1].id].endswith(trace.spans[1].id) + + @staticmethod + @pytest.mark.parametrize( + ("provider", "config", "suffix", "kind_key", "kind"), + [ + ( + "phoenix", + {"endpoint": "https://phoenix.test/proxy", "api_key": "key", "project": "one"}, + "/proxy/v1/traces", + "openinference.span.kind", + "AGENT", + ), + ( + "aliyun", + {"endpoint": "https://region.log.aliyuncs.com", "license_key": "license", "app_name": "one"}, + "/adapt_license/api/v1/traces", + "gen_ai.span.kind", + "AGENT", + ), + ], ) - config = {"host": "https://langfuse.test", "public_key": "pk", "secret_key": "sk"} - if body.get("partialSuccess", {}).get("rejectedSpans"): - with pytest.raises(TraceExportUncertainError): - export_trace("langfuse", config, trace) - else: - assert export_trace("langfuse", config, trace).status == "accepted" + def test_other_http_providers_have_native_semantics( + trace: CanonicalTrace, + http_calls: list[HTTPCall], + provider: str, + config: dict[str, str], + suffix: str, + kind_key: str, + kind: str, + ) -> None: + result = export_trace(provider, config, trace) + assert result.status == "accepted" + assert http_calls[0][1].endswith(suffix) + spans = spans_from_call(http_calls[0]) + assert attributes(spans[0])[kind_key] == kind + token_key = "llm.token_count.prompt" if provider == "phoenix" else "gen_ai.usage.input_tokens" + assert attributes(spans[1])[token_key] == 8 + assert token_key not in attributes(spans[0]) + @staticmethod + @pytest.mark.parametrize( + ("endpoint_url", "suffix"), + [ + ("https://log.aliyuncs.com", "api/v1/traces"), + ("https://cn-hangzhou.log.aliyuncs.com", "api/v1/traces"), + ("https://project.cn-heyuan.log.aliyuncs.com/apm/trace/opentelemetry", "api/v1/traces"), + ("https://project.cn-heyuan-intranet.log.aliyuncs.com/apm/trace/opentelemetry", "api/v1/traces"), + ("https://cn-shanghai-share.log.aliyuncs.com", "api/v1/traces"), + ("https://cn-beijing.dualstack.log.aliyuncs.com", "api/v1/traces"), + ("https://cn-hangzhou.log.aliyuncs.com.:443", "api/v1/traces"), + ("https://tracing-analysis-dc-hz.aliyuncs.com", "api/otlp/traces"), + ("https://collector.example.test:4318/custom-prefix", "api/otlp/traces"), + ("https://evil-log.aliyuncs.com.evil.test", "api/otlp/traces"), + ("https://log.aliyuncs.com.evil.test", "api/otlp/traces"), + ("https://evil-log.aliyuncs.com", "api/otlp/traces"), + ("https://collector.example.test/log.aliyuncs.com", "api/otlp/traces"), + ], + ) + def test_aliyun_endpoint_family_uses_dns_labels_and_preserves_custom_destinations( + trace: CanonicalTrace, http_calls: list[HTTPCall], endpoint_url: str, suffix: str + ) -> None: + result = export_trace( + "aliyun", {"endpoint": endpoint_url, "license_key": "synthetic-license", "app_name": "test-app"}, trace + ) + assert result.status == "accepted" + assert len(http_calls) == 1 + assert http_calls[0][1] == f"{endpoint_url}/adapt_synthetic-license/{suffix}" + spans = spans_from_call(http_calls[0]) + assert [span.span_id.hex() for span in spans] == [span.id for span in trace.spans] + assert spans[1].parent_span_id == spans[0].span_id -def test_lost_http_ack_is_uncertain_not_rejected(trace: CanonicalTrace, monkeypatch: pytest.MonkeyPatch) -> None: - from core.helper import ssrf_proxy + @staticmethod + def test_aliyun_raw_license_is_encoded_once_at_the_real_otlp_boundary( + trace: CanonicalTrace, http_calls: list[HTTPCall] + ) -> None: + result = export_trace( + "aliyun", + {"endpoint": "https://COLLECTOR.test:8443/prefix/", "license_key": "raw%2F@a+b/c", "app_name": "test_app"}, + trace, + ) + assert result.status == "accepted" + assert len(http_calls) == 1 + assert http_calls[0][1] == "https://collector.test:8443/prefix/adapt_raw%252F@a%2Bb%2Fc/api/otlp/traces" + assert [span.span_id.hex() for span in spans_from_call(http_calls[0])] == [span.id for span in trace.spans] - def timeout(*_args: object, **_kwargs: object) -> Never: - raise httpx.ReadTimeout("response lost") + @staticmethod + def test_aliyun_complete_token_url_is_rejected_without_any_ingestion( + trace: CanonicalTrace, http_calls: list[HTTPCall] + ) -> None: + with pytest.raises(TraceExportRejectedError, match="Aliyun tracing configuration is invalid"): + export_trace( + "aliyun", + {"endpoint": "https://collector.test/adapt_private/api/otlp/traces", "license_key": "private"}, + trace, + ) + assert http_calls == [] - monkeypatch.setattr(ssrf_proxy, "make_request", timeout) - with pytest.raises(TraceExportUncertainError): - export_trace("langfuse", {"host": "https://langfuse.test", "public_key": "pk", "secret_key": "sk"}, trace) + @staticmethod + @pytest.mark.parametrize( + ("provider", "config", "header"), + [ + ( + "arize", + {"endpoint": "https://otlp.arize.com", "api_key": "key", "space_id": "space", "project": "one"}, + ("space_id", "space"), + ), + ( + "tencent", + {"endpoint": "https://apm.test:4317", "token": "token", "service_name": "one"}, + ("authorization", "Bearer token"), + ), + ], + ) + def test_grpc_real_stub_serialization( + trace: CanonicalTrace, + monkeypatch: pytest.MonkeyPatch, + provider: str, + config: dict[str, str], + header: tuple[str, str], + ) -> None: + calls: list[tuple[ExportTraceServiceRequest, Sequence[tuple[str, str]]]] = [] + class Channel: + def unary_unary( + self, + method: str, + request_serializer: Callable[[ExportTraceServiceRequest], bytes], + response_deserializer: Callable[[bytes], ExportTraceServiceResponse], + **_kwargs: object, + ) -> Callable[..., ExportTraceServiceResponse]: + assert method == "/opentelemetry.proto.collector.trace.v1.TraceService/Export" -@pytest.mark.parametrize("provider", ["langsmith", "opik"]) -@pytest.mark.parametrize("accepted_prefix", [False, True]) -def test_native_sdk_preserves_first_failure_or_partial_acceptance( - trace: CanonicalTrace, monkeypatch: pytest.MonkeyPatch, provider: str, accepted_prefix: bool -) -> None: - from core.helper import ssrf_proxy + def invoke( + payload: ExportTraceServiceRequest, metadata: Sequence[tuple[str, str]], timeout: float + ) -> ExportTraceServiceResponse: + assert timeout == 30 + calls.append((ExportTraceServiceRequest.FromString(request_serializer(payload)), metadata)) + return response_deserializer(ExportTraceServiceResponse().SerializeToString()) - calls: list[str] = [] + return invoke - def send(method: str, url: str, **_kwargs: object) -> httpx.Response: - calls.append(url) - status = 200 if accepted_prefix and len(calls) == 1 else 429 - return httpx.Response(status, content=b"", request=httpx.Request(method, url)) + def close(self) -> None: + pass - monkeypatch.setattr(ssrf_proxy, "make_request", send) - config = {"api_key": "test-key", "project": "one", "endpoint": "https://smith.test", "url": "https://opik.test/api"} - expected = TraceExportUncertainError if accepted_prefix else RetryableTraceDispatchError - with pytest.raises(expected) as caught: - export_trace(provider, config, trace) - assert isinstance(caught.value, TraceExportUncertainError) is accepted_prefix - assert len(calls) == (2 if accepted_prefix else 1) + monkeypatch.setattr("grpc.secure_channel", lambda *_args, **_kwargs: Channel()) + result = export_trace(provider, config, trace) + assert result.status == "accepted" + assert header in calls[0][1] + spans = calls[0][0].resource_spans[0].scope_spans[0].spans + assert spans[1].parent_span_id == spans[0].span_id + @staticmethod + def test_langsmith_real_sdk_uses_isolated_sync_runs(trace: CanonicalTrace, http_calls: list[HTTPCall]) -> None: + result = export_trace( + "langsmith", {"api_key": "ls-test", "endpoint": "https://langsmith.test", "project": "one"}, trace + ) + writes = [json.loads(call[2]["content"]) for call in http_calls if call[0] == "POST"] + assert len(writes) == 3 + root, llm, tool = writes + assert root["run_type"] == "chain" + assert llm["run_type"] == "llm" + assert llm["parent_run_id"] == root["id"] + assert llm["dotted_order"].startswith(root["dotted_order"]) + assert llm["input_tokens"] == 8 + assert llm["output_tokens"] == 4 + assert llm["total_cost"] == 0.03 + assert "input_tokens" not in root + assert all(write["extra"]["metadata"]["session_id"] == trace.session_id for write in writes) + assert tool["error"] == "tool failed" + assert UUID(result.trace_id) -def test_http_exports_keep_destination_credentials_local(trace: CanonicalTrace, http_calls: list[HTTPCall]) -> None: - for project in ("a", "b"): + @staticmethod + def test_opik_real_sdk_preserves_tree_thread_and_native_usage( + trace: CanonicalTrace, http_calls: list[HTTPCall] + ) -> None: + result = export_trace( + "opik", + {"api_key": "opik-test", "url": "https://opik.test/api", "workspace": "org-one", "project": "one"}, + trace, + ) + assert len(http_calls) == 2 + root = json.loads(http_calls[0][2]["content"]) + batch = json.loads(http_calls[1][2]["content"]) + spans = batch["spans"] + assert UUID(result.trace_id).version == 7 + assert root["thread_id"] == trace.session_id + assert spans[1]["parent_span_id"] == spans[0]["id"] + assert spans[1]["usage"]["prompt_tokens"] == 8 + assert spans[1]["total_estimated_cost"] == 0.03 + assert not spans[0].get("usage") + assert spans[2]["error_info"]["message"] == "tool failed" + assert http_calls[0][2]["headers"]["comet-workspace"] == "org-one" + + @staticmethod + @pytest.mark.parametrize( + ("value", "wire_value"), + [ + (None, None), + ("323", "323"), + ({"result": 323}, {"result": 323}), + ([{"result": 323}], [{"result": 323}]), + ([], []), + (323, "323"), + (0, "0"), + (-5, "-5"), + (3.5, "3.5"), + (True, "true"), + (False, "false"), + ([1, 2], "[1, 2]"), + (["one", "two"], '["one", "two"]'), + ([{"result": 1}, "text"], '[{"result": 1}, "text"]'), + ([[1], [2]], "[[1], [2]]"), + ], + ) + def test_opik_real_sdk_preserves_json_content_across_trace_and_span_boundaries( + trace: CanonicalTrace, http_calls: list[HTTPCall], value: object, wire_value: object + ) -> None: + content_trace = trace.model_copy( + update={"spans": tuple(span.model_copy(update={"inputs": value, "outputs": value}) for span in trace.spans)} + ) + result = export_trace( + "opik", + {"api_key": "synthetic-key", "url": "https://opik.test/api", "workspace": "team", "project": "one"}, + content_trace, + ) + assert result.status == "accepted" + assert len(http_calls) == 2 + root = json.loads(http_calls[0][2]["content"]) + spans = json.loads(http_calls[1][2]["content"])["spans"] + for payload in [root, *spans]: + assert payload.get("input") == wire_value + assert payload.get("output") == wire_value + assert spans[1]["parent_span_id"] == spans[0]["id"] + assert spans[2]["parent_span_id"] == spans[0]["id"] + assert root["thread_id"] == trace.session_id + assert spans[1]["usage"]["total_tokens"] == 12 + assert not spans[0].get("usage") + assert not spans[2].get("usage") + + @staticmethod + @pytest.mark.parametrize("capture_content", [True, False]) + @pytest.mark.parametrize("root_error", [True, False]) + def test_opik_trace_error_matches_sanitized_root_and_not_handled_tool_failure( + trace: CanonicalTrace, http_calls: list[HTTPCall], capture_content: bool, root_error: bool + ) -> None: + root = trace.spans[0].model_copy( + update={ + "status": CanonicalSpanStatus.ERROR if root_error else CanonicalSpanStatus.OK, + "error": "private-model-error" if root_error else None, + "inputs": {"query": "private-error-input"}, + } + ) + source = trace.model_copy(update={"spans": (root, *trace.spans[1:])}) + canonical = remap_outer_trace( + source, trace_id=trace.trace_id, capture_content=capture_content, agent_message=True + ) export_trace( - "langfuse", {"host": f"https://{project}.test", "public_key": project, "secret_key": project}, trace + "opik", + {"api_key": "synthetic-key", "url": "https://opik.test/api", "workspace": "team", "project": "one"}, + canonical, ) - assert http_calls[0][1].startswith("https://a.test/") - assert http_calls[1][1].startswith("https://b.test/") - assert http_calls[0][2]["headers"]["Authorization"] != http_calls[1][2]["headers"]["Authorization"] - assert http_calls[0][2]["content"] == http_calls[1][2]["content"] + root_payload = json.loads(http_calls[0][2]["content"]) + spans = json.loads(http_calls[1][2]["content"])["spans"] + assert spans[2]["error_info"] is not None + if root_error: + expected = { + "exception_type": "DifyExecutionError", + "message": "private-model-error" if capture_content else "Application execution failed", + "traceback": "", + } + assert root_payload["error_info"] == spans[0]["error_info"] == expected + else: + assert root_payload.get("error_info") is None + assert spans[0].get("error_info") is None + if not capture_content: + assert all(b"private-model-error" not in call[2]["content"] for call in http_calls) + assert all(b"private-error-input" not in call[2]["content"] for call in http_calls) + @staticmethod + @pytest.mark.parametrize( + ("status", "error"), + [(401, TraceExportRejectedError), (429, RetryableTraceDispatchError), (503, RetryableTraceDispatchError)], + ) + def test_explicit_http_rejections_are_classified_without_response_body( + trace: CanonicalTrace, monkeypatch: pytest.MonkeyPatch, status: int, error: type[Exception] + ) -> None: + from core.helper import ssrf_proxy -@pytest.mark.parametrize( - "provider", - ["langfuse", "langsmith", "phoenix", "arize", "opik", "weave", "mlflow", "databricks", "aliyun", "tencent"], -) -def test_capabilities_state_implemented_contract_not_platform_validation(provider: str) -> None: - result = provider_capabilities(provider) - assert result["agent_spans"] is True - assert result["export_ack"] == "accepted" - assert "platform_validated" not in result + monkeypatch.setattr( + ssrf_proxy, + "make_request", + lambda method, url, **_kwargs: httpx.Response( + status, text="secret-token", request=httpx.Request(method, url) + ), + ) + with pytest.raises(error) as caught: + export_trace( + "langfuse", {"host": "https://langfuse.test", "public_key": "pk", "secret_key": "secret-token"}, trace + ) + assert "secret-token" not in str(caught.value) + @staticmethod + def test_partial_otlp_acceptance_requires_reconciliation( + trace: CanonicalTrace, monkeypatch: pytest.MonkeyPatch + ) -> None: + from core.helper import ssrf_proxy -def test_rejects_unfinished_and_unknown_provider_before_sending( - trace: CanonicalTrace, http_calls: list[HTTPCall] -) -> None: - unfinished = trace.model_copy(update={"spans": (trace.spans[0].model_copy(update={"end_time": None}),)}) - with pytest.raises(TraceExportRejectedError): - export_trace("langfuse", {}, unfinished) - with pytest.raises(TraceExportRejectedError): - export_trace("unknown", {}, trace) - assert not http_calls + response = ExportTraceServiceResponse() + response.partial_success.rejected_spans = 1 + monkeypatch.setattr( + ssrf_proxy, + "make_request", + lambda method, url, **_kwargs: httpx.Response( + 200, content=response.SerializeToString(), request=httpx.Request(method, url) + ), + ) + with pytest.raises(TraceExportUncertainError): + export_trace("langfuse", {"host": "https://langfuse.test", "public_key": "pk", "secret_key": "sk"}, trace) + + @staticmethod + def test_langfuse_v4_queue_ack_confirms_ingestion_acceptance_only( + trace: CanonicalTrace, + monkeypatch: pytest.MonkeyPatch, + caplog: pytest.LogCaptureFixture, + capsys: pytest.CaptureFixture[str], + ) -> None: + from core.helper import ssrf_proxy + + monkeypatch.setattr( + ssrf_proxy, + "make_request", + lambda method, url, **_kwargs: httpx.Response( + 200, json=_langfuse_queue_job_ack(), request=httpx.Request(method, url) + ), + ) + result = export_trace( + "langfuse", {"host": "https://langfuse.test", "public_key": "pk", "secret_key": "sk"}, trace + ) + assert result.status == "accepted" + assert result.trace_id == trace.trace_id + assert result.trace_url is None + captured = capsys.readouterr() + evidence = caplog.text + captured.out + captured.err + repr(result) + assert "never-log-public-key" not in evidence + assert "never-log-project" not in evidence + assert "authCheck" not in evidence + + @staticmethod + @pytest.mark.parametrize( + "scenario", + [ + "wrong_name", + "wrong_nested_name", + "invalid_id", + "invalid_timestamp", + "invalid_auth", + "missing_scope", + "missing_file", + "partial_success", + "error", + "failed_job", + "unknown", + "null", + "array", + ], + ) + def test_langfuse_ambiguous_json_ack_remains_uncertain( + trace: CanonicalTrace, + monkeypatch: pytest.MonkeyPatch, + caplog: pytest.LogCaptureFixture, + capsys: pytest.CaptureFixture[str], + scenario: str, + ) -> None: + from core.helper import ssrf_proxy + + ack = _langfuse_queue_job_ack() + body: object = ack + if scenario == "wrong_name": + ack["name"] = "different-job" + elif scenario == "wrong_nested_name": + _replace_ack_value(ack, ("data", "name"), "different-job") + elif scenario == "invalid_id": + _replace_ack_value(ack, ("data", "id"), "not-a-uuid") + elif scenario == "invalid_timestamp": + _replace_ack_value(ack, ("data", "timestamp"), "2026-10-15T01:00:00") + elif scenario == "invalid_auth": + _replace_ack_value(ack, ("data", "payload", "authCheck", "validKey"), False) + elif scenario == "missing_scope": + _replace_ack_value(ack, ("data", "payload", "authCheck", "scope"), dict[str, object]()) + elif scenario == "missing_file": + _replace_ack_value(ack, ("data", "payload", "data", "fileKey"), "") + elif scenario == "partial_success": + ack["partialSuccess"] = {"rejectedSpans": "1"} + elif scenario == "error": + ack["error"] = "never-log-public-key" + elif scenario == "failed_job": + ack["failedReason"] = "never-log-public-key" + elif scenario == "unknown": + body = {"accepted": True, "publicKey": "never-log-public-key"} + elif scenario == "null": + body = None + elif scenario == "array": + body = list[object]() + monkeypatch.setattr( + ssrf_proxy, + "make_request", + lambda method, url, **_kwargs: httpx.Response( + 200, + content=json.dumps(body).encode(), + headers={"Content-Type": "application/json"}, + request=httpx.Request(method, url), + ), + ) + with pytest.raises(TraceExportUncertainError) as caught: + export_trace("langfuse", {"host": "https://langfuse.test", "public_key": "pk", "secret_key": "sk"}, trace) + captured = capsys.readouterr() + evidence = caplog.text + captured.out + captured.err + str(caught.value) + assert "never-log-public-key" not in evidence + assert "never-log-project" not in evidence + assert "authCheck" not in evidence + + @staticmethod + @pytest.mark.parametrize( + ("provider", "version", "status"), [("phoenix", "4", 200), ("langfuse", "3", 200), ("langfuse", "4", 202)] + ) + def test_langfuse_queue_ack_is_not_a_generic_otlp_success( + monkeypatch: pytest.MonkeyPatch, provider: str, version: str, status: int + ) -> None: + from core.helper import ssrf_proxy + from services.tracing.otlp_export import export_http + + monkeypatch.setattr( + ssrf_proxy, + "make_request", + lambda method, url, **_kwargs: httpx.Response( + status, json=_langfuse_queue_job_ack(), request=httpx.Request(method, url) + ), + ) + with pytest.raises(TraceExportUncertainError): + export_http( + provider, + "https://collector.test/v1/traces", + {"x-langfuse-ingestion-version": version}, + ExportTraceServiceRequest(), + ) + + @staticmethod + @pytest.mark.parametrize("body", [{}, {"partialSuccess": {}}, {"partialSuccess": {"rejectedSpans": "1"}}]) + def test_langfuse_standard_otlp_json_ack_remains_supported( + trace: CanonicalTrace, monkeypatch: pytest.MonkeyPatch, body: dict[str, dict[str, str]] + ) -> None: + from core.helper import ssrf_proxy + + monkeypatch.setattr( + ssrf_proxy, + "make_request", + lambda method, url, **_kwargs: httpx.Response(200, json=body, request=httpx.Request(method, url)), + ) + config = {"host": "https://langfuse.test", "public_key": "pk", "secret_key": "sk"} + if body.get("partialSuccess", {}).get("rejectedSpans"): + with pytest.raises(TraceExportUncertainError): + export_trace("langfuse", config, trace) + else: + assert export_trace("langfuse", config, trace).status == "accepted" + + @staticmethod + def test_lost_http_ack_is_uncertain_not_rejected(trace: CanonicalTrace, monkeypatch: pytest.MonkeyPatch) -> None: + from core.helper import ssrf_proxy + + def timeout(*_args: object, **_kwargs: object) -> Never: + raise httpx.ReadTimeout("response lost") + + monkeypatch.setattr(ssrf_proxy, "make_request", timeout) + with pytest.raises(TraceExportUncertainError): + export_trace("langfuse", {"host": "https://langfuse.test", "public_key": "pk", "secret_key": "sk"}, trace) + + @staticmethod + @pytest.mark.parametrize("provider", ["langsmith", "opik"]) + @pytest.mark.parametrize("accepted_prefix", [False, True]) + def test_native_sdk_preserves_first_failure_or_partial_acceptance( + trace: CanonicalTrace, monkeypatch: pytest.MonkeyPatch, provider: str, accepted_prefix: bool + ) -> None: + from core.helper import ssrf_proxy + + calls: list[str] = [] + + def send(method: str, url: str, **_kwargs: object) -> httpx.Response: + calls.append(url) + status = 200 if accepted_prefix and len(calls) == 1 else 429 + return httpx.Response(status, content=b"", request=httpx.Request(method, url)) + + monkeypatch.setattr(ssrf_proxy, "make_request", send) + config = { + "api_key": "test-key", + "project": "one", + "endpoint": "https://smith.test", + "url": "https://opik.test/api", + } + expected = TraceExportUncertainError if accepted_prefix else RetryableTraceDispatchError + with pytest.raises(expected) as caught: + export_trace(provider, config, trace) + assert isinstance(caught.value, TraceExportUncertainError) is accepted_prefix + assert len(calls) == (2 if accepted_prefix else 1) + + @staticmethod + def test_http_exports_keep_destination_credentials_local(trace: CanonicalTrace, http_calls: list[HTTPCall]) -> None: + for project in ("a", "b"): + export_trace( + "langfuse", {"host": f"https://{project}.test", "public_key": project, "secret_key": project}, trace + ) + assert http_calls[0][1].startswith("https://a.test/") + assert http_calls[1][1].startswith("https://b.test/") + assert http_calls[0][2]["headers"]["Authorization"] != http_calls[1][2]["headers"]["Authorization"] + assert http_calls[0][2]["content"] == http_calls[1][2]["content"] + + @staticmethod + @pytest.mark.parametrize( + "provider", + ["langfuse", "langsmith", "phoenix", "arize", "opik", "weave", "mlflow", "databricks", "aliyun", "tencent"], + ) + def test_capabilities_state_implemented_contract_not_platform_validation(provider: str) -> None: + result = provider_capabilities(provider) + assert result["agent_spans"] is True + assert result["export_ack"] == "accepted" + assert "platform_validated" not in result + + @staticmethod + def test_rejects_unfinished_and_unknown_provider_before_sending( + trace: CanonicalTrace, http_calls: list[HTTPCall] + ) -> None: + unfinished = trace.model_copy(update={"spans": (trace.spans[0].model_copy(update={"end_time": None}),)}) + with pytest.raises(TraceExportRejectedError): + export_trace("langfuse", {}, unfinished) + with pytest.raises(TraceExportRejectedError): + export_trace("unknown", {}, trace) + assert not http_calls diff --git a/api/tests/unit_tests/services/tracing/test_application_trace_recovery.py b/api/tests/unit_tests/services/tracing/test_recovery.py similarity index 98% rename from api/tests/unit_tests/services/tracing/test_application_trace_recovery.py rename to api/tests/unit_tests/services/tracing/test_recovery.py index 471366713c5..91cda1b9d1d 100644 --- a/api/tests/unit_tests/services/tracing/test_application_trace_recovery.py +++ b/api/tests/unit_tests/services/tracing/test_recovery.py @@ -26,14 +26,12 @@ from services.tracing import service as service_module from services.tracing.normalization import span_id from services.tracing.provider_dispatch import TraceExportResult, TraceExportUncertainError from services.tracing.service import ApplicationTracingService -from tests.unit_tests.model_factories import make_account, make_conversation -from tests.unit_tests.services.tracing.test_application_trace_collection import ( - Collection, - batch_for, - identifier, - ingest, - seed, -) +from tests.helpers.unit.model_factories import make_account, make_conversation +from tests.unit_tests.services.tracing.test_service import _application_trace_collection_batch_for as batch_for +from tests.unit_tests.services.tracing.test_service import _application_trace_collection_identifier as identifier +from tests.unit_tests.services.tracing.test_service import _application_trace_collection_ingest as ingest +from tests.unit_tests.services.tracing.test_service import _application_trace_collection_seed as seed +from tests.unit_tests.services.tracing.test_service import _ApplicationTraceCollectionCollection as Collection @dataclass @@ -130,12 +128,12 @@ def save_agent_stop( from core.ops import ops_trace_manager from core.ops.base_trace_instance import bind_application_tracing_hooks from services.tracing import runtime_bridge - from tests.unit_tests.config_override import apply_config_overrides + from tests.helpers.unit.config_override import apply_config_overrides + from tests.helpers.unit.core.model_fixtures import make_model_config from tests.unit_tests.core.app.task_pipeline.test_easy_ui_based_generate_task_pipeline import ( _QueueManager, queue_message, ) - from tests.unit_tests.core.model_fixtures import make_model_config apply_config_overrides(monkeypatch, OPS_AGENT_TRACE_ENABLED=True) monkeypatch.setattr(runtime_bridge, "_service", lambda: case.service) @@ -447,7 +445,7 @@ def test_stopped_runner_hands_off_only_an_observed_runtime_terminal( from core.app.entities.app_invoke_entities import DifyRunContext, InvokeFrom, UserFrom from core.ops.base_trace_instance import bind_application_tracing_hooks from services.tracing import runtime_bridge - from tests.unit_tests.config_override import apply_config_overrides + from tests.helpers.unit.config_override import apply_config_overrides from tests.unit_tests.core.app.apps.agent_app.test_app_runner import ( _CancelAndWaitFailingClient, _ExplodingSessionStore, @@ -565,7 +563,7 @@ def test_stopped_runner_hands_off_only_an_observed_runtime_terminal( def test_missing_runtime_stop_terminal_converges_without_inventing_success( case: RecoveryCase, monkeypatch: pytest.MonkeyPatch, incomplete: bool ) -> None: - from tests.unit_tests.config_override import apply_config_overrides + from tests.helpers.unit.config_override import apply_config_overrides apply_config_overrides(monkeypatch, OPS_AGENT_TRACE_COLLECTION_TIMEOUT=60) save_agent_stop(case, monkeypatch, "partial answer") @@ -697,7 +695,7 @@ def test_runtime_cleanup_cancellation_does_not_reclassify_a_business_stream_fail from core.app.entities.app_invoke_entities import DifyRunContext, InvokeFrom, UserFrom from core.ops.base_trace_instance import bind_application_tracing_hooks from services.tracing import runtime_bridge - from tests.unit_tests.config_override import apply_config_overrides + from tests.helpers.unit.config_override import apply_config_overrides from tests.unit_tests.core.app.apps.agent_app.test_app_runner import ( _FakeQueueManager, _FakeSessionStore, @@ -800,7 +798,7 @@ def test_runtime_cleanup_cancellation_does_not_reclassify_a_business_stream_fail def test_missing_stop_worker_without_runtime_spans_also_reaches_a_delivery_failure( case: RecoveryCase, monkeypatch: pytest.MonkeyPatch, capture_agent: bool ) -> None: - from tests.unit_tests.config_override import apply_config_overrides + from tests.helpers.unit.config_override import apply_config_overrides collection = seed(case.collection.sessions, f"stop-without-spans-{capture_agent}", capture_agent=capture_agent) attach_conversation(collection) diff --git a/api/tests/unit_tests/services/tracing/test_runtime_bridge.py b/api/tests/unit_tests/services/tracing/test_runtime_bridge.py index bb2ae7ca369..b262125f42c 100644 --- a/api/tests/unit_tests/services/tracing/test_runtime_bridge.py +++ b/api/tests/unit_tests/services/tracing/test_runtime_bridge.py @@ -1,20 +1,39 @@ -"""Verify execution hook ownership, policy gating, and fail-open business behavior.""" - -from dataclasses import dataclass, field +import json +from copy import deepcopy +from dataclasses import dataclass, field, replace +from datetime import UTC, datetime, timedelta from typing import cast -from uuid import uuid4 +from uuid import UUID, uuid4 +import httpx import pytest from dify_agent.layers.execution_context import ApplicationTracingContext +from dify_agent.protocol.tracing import AgentTraceBatch, AgentTraceSpan +from flask import Flask +from opentelemetry.proto.collector.trace.v1.trace_service_pb2 import ExportTraceServiceRequest +from sqlalchemy import delete, func, select from sqlalchemy.exc import OperationalError from sqlalchemy.orm import Session, sessionmaker from core.app.app_config.entities import WorkflowUIBasedAppConfig from core.app.apps.agent_app.runtime_request_builder import AgentAppRuntimeBuildContext from core.app.entities.app_invoke_entities import DifyRunContext, InvokeFrom, UserFrom, WorkflowAppGenerateEntity -from core.ops.entities.trace_entity import BaseTraceInfo +from core.ops.base_trace_instance import bind_application_tracing_hooks +from core.ops.entities.trace_entity import ( + BaseTraceInfo, + GenerateNameTraceInfo, + MessageTraceInfo, + ModerationTraceInfo, + WorkflowTraceInfo, +) from core.ops.exceptions import RetryableTraceDispatchError +from core.ops.ops_trace_manager import OpsTraceManager +from core.ops.unified_trace.trace_builder import CanonicalTraceBuilder +from core.repositories.sqlalchemy_workflow_node_execution_repository import SQLAlchemyWorkflowNodeExecutionRepository from core.workflow.nodes.agent_v2.runtime_request_builder import WorkflowAgentRuntimeBuildContext +from extensions.ext_database import db +from graphon.enums import WorkflowExecutionStatus, WorkflowNodeExecutionStatus +from graphon.enums import WorkflowType as _workflow_tool_tracing_boundary_WorkflowType from graphon.runtime import VariablePool from models.agent import ( Agent, @@ -25,11 +44,474 @@ from models.agent import ( WorkflowAgentNodeBinding, ) from models.agent_config_entities import AgentSoulConfig, WorkflowNodeJobConfig -from models.model import AppMode +from models.application_trace import ApplicationTraceRun +from models.enums import ConversationFromSource, CreatorUserRole +from models.enums import WorkflowRunTriggeredFrom as _application_trace_handoff_WorkflowRunTriggeredFrom +from models.model import App, AppMode, Message, TraceAppConfig +from models.workflow import ( + WorkflowNodeExecutionModel, + WorkflowNodeExecutionTriggeredFrom, + WorkflowRun, +) +from models.workflow import ( + WorkflowRunTriggeredFrom as _workflow_tool_tracing_boundary_WorkflowRunTriggeredFrom, +) +from models.workflow import WorkflowType as _application_trace_handoff_WorkflowType from repositories.application_trace_repository import ApplicationTraceRepository from services.tracing import runtime_bridge +from services.tracing.normalization import remap_outer_trace +from services.tracing.provider_dispatch import export_trace from services.tracing.service import ApplicationTracingService -from tests.unit_tests.config_override import apply_config_overrides +from tasks import ops_trace_task +from tests.helpers.unit.config_override import apply_config_overrides +from tests.helpers.unit.model_factories import make_account, make_app, make_conversation, make_workflow +from tests.unit_tests.services.tracing.test_service import _application_trace_collection_batch_for as batch_for +from tests.unit_tests.services.tracing.test_service import _application_trace_collection_identifier as identifier +from tests.unit_tests.services.tracing.test_service import _application_trace_collection_ingest as ingest +from tests.unit_tests.services.tracing.test_service import ( + _application_trace_collection_seed as _application_trace_handoff_seed, +) +from tests.unit_tests.services.tracing.test_service import _application_trace_collection_terminal as terminal +from tests.unit_tests.services.tracing.test_service import _ApplicationTraceCollectionCollection as Collection + +"""SQLite ownership regressions at the actual legacy-task handoff boundary.""" + + +AGENT_GRAPH = { + "nodes": [{"id": "agent", "data": {"type": "agent", "version": "2", "agent_node_kind": "dify_agent"}}], + "edges": list[object](), +} + + +@dataclass +class Case: + collection: Collection + service: ApplicationTracingService + scheduled: list[str] + + +def setup( + sessions: sessionmaker[Session], monkeypatch: pytest.MonkeyPatch, name: str, mode: AppMode, *, snapshot: bool +) -> Case: + apply_config_overrides(monkeypatch, OPS_AGENT_TRACE_ENABLED=True) + collection = _application_trace_handoff_seed(sessions, name) + with sessions.begin() as session: + app = session.get(App, collection.run.app_id) + message = session.get(Message, collection.run.source_id) + assert app is not None + assert message is not None + app.mode = mode + session.add( + make_conversation( + conversation_id=message.conversation_id, + app_id=app.id, + mode=mode, + inputs={}, + from_source=ConversationFromSource.CONSOLE, + ) + ) + if not snapshot: + session.execute(delete(ApplicationTraceRun).where(ApplicationTraceRun.id == collection.run.id)) + service = ApplicationTracingService(collection.repository) + scheduled: list[str] = [] + monkeypatch.setattr(service, "schedule", scheduled.append) + monkeypatch.setattr(runtime_bridge, "_service", lambda: service) + return Case(collection, service, scheduled) + + +def message_info(case: Case) -> MessageTraceInfo: + with case.collection.sessions() as session: + message = session.get(Message, case.collection.run.source_id) + assert message is not None + return MessageTraceInfo( + metadata={"app_id": message.app_id}, + message_id=message.id, + message_data=message.to_dict(session=session), + conversation_model="agent", + conversation_mode="agent", + message_tokens=0, + answer_tokens=0, + total_tokens=0, + inputs="private input", + outputs="private output", + ) + + +def dispatch(case: Case, info: BaseTraceInfo, monkeypatch: pytest.MonkeyPatch) -> list[BaseTraceInfo]: + calls: list[BaseTraceInfo] = [] + failures: list[str] = [] + + class Legacy: + def trace(self, restored: BaseTraceInfo) -> None: + calls.append(restored) + + payload = json.dumps({"trace_info_type": type(info).__name__, "trace_info": info.model_dump(mode="json")}) + monkeypatch.setattr(ops_trace_task.storage, "load", lambda _path: payload) + monkeypatch.setattr(ops_trace_task.storage, "delete", lambda _path: None) + monkeypatch.setattr(ops_trace_task.redis_client, "incr", failures.append) + monkeypatch.setattr(OpsTraceManager, "get_ops_trace_instance", lambda _app: Legacy()) + monkeypatch.setattr("extensions.ext_enterprise_telemetry.is_enabled", lambda: False) + app = Flask(__name__) + bind_application_tracing_hooks(app, runtime_bridge.RuntimeApplicationTracingHooks()) + with app.app_context(): + ops_trace_task.process_trace_tasks.run({"app_id": case.collection.run.app_id, "file_id": "fixture"}) + assert failures == [] + return calls + + +def workflow( + case: Case, name: str, *, has_agent: bool, snapshot: bool +) -> tuple[WorkflowRun, ApplicationTraceRun | None]: + graph = AGENT_GRAPH if has_agent else {"nodes": [{"data": {"type": "llm"}}], "edges": list[object]()} + now = datetime(2026, 10, 15) + with case.collection.sessions.begin() as session: + definition = make_workflow( + workflow_id=identifier(f"workflow-{name}"), + tenant_id=case.collection.run.tenant_id, + app_id=case.collection.run.app_id, + graph=graph, + ) + session.add(definition) + app = session.get(App, case.collection.run.app_id) + assert app is not None + app.workflow_id = definition.id + row = WorkflowRun( + id=identifier(f"workflow-run-{name}"), + app_id=app.id, + tenant_id=app.tenant_id, + workflow_id=definition.id, + type=_application_trace_handoff_WorkflowType.WORKFLOW, + version="1", + graph=json.dumps(graph), + inputs="{}", + outputs="{}", + status=WorkflowExecutionStatus.SUCCEEDED, + triggered_from=_application_trace_handoff_WorkflowRunTriggeredFrom.APP_RUN, + created_by_role=CreatorUserRole.ACCOUNT, + created_by=identifier("owner"), + created_at=now, + finished_at=now + timedelta(seconds=2), + elapsed_time=2, + total_tokens=0, + total_steps=0, + ) + session.add(row) + trace_run = ( + case.collection.repository.start( + tenant_id=row.tenant_id, + app_id=row.app_id, + source_kind="workflow", + source_id=row.id, + session_id="workflow-session", + ) + if snapshot + else None + ) + return row, trace_run + + +def workflow_info(row: WorkflowRun, parent: WorkflowRun | None = None) -> WorkflowTraceInfo: + metadata: dict[str, object] = {"app_id": row.app_id} + if parent is not None: + metadata["parent_trace_context"] = { + "parent_workflow_run_id": parent.id, + "parent_node_execution_id": identifier("parent-tool-node"), + } + return WorkflowTraceInfo( + workflow_data=row, + workflow_id=row.workflow_id, + tenant_id=row.tenant_id, + workflow_run_id=row.id, + workflow_run_elapsed_time=2, + workflow_run_status="succeeded", + workflow_run_inputs={}, + workflow_run_outputs={}, + workflow_run_version="1", + total_tokens=0, + file_list=[], + query="", + metadata=metadata, + start_time=row.created_at, + end_time=row.finished_at, + ) + + +class TestApplicationTraceHandoffCases: + @staticmethod + def test_chatflow_message_event_cannot_export_again_via_legacy( + sqlite_session_factory: sessionmaker[Session], monkeypatch: pytest.MonkeyPatch + ) -> None: + case = setup(sqlite_session_factory, monkeypatch, "chatflow", AppMode.ADVANCED_CHAT, snapshot=False) + row, trace_run = workflow(case, "chatflow", has_agent=True, snapshot=True) + assert trace_run is not None + terminal(Collection(case.collection.repository, case.collection.sessions, trace_run)) + before = case.collection.repository.get(trace_run.id) + assert before is not None + with case.collection.sessions.begin() as session: + message = session.get(Message, case.collection.run.source_id) + assert message is not None + message.workflow_run_id = row.id + assert dispatch(case, message_info(case), monkeypatch) == [] + after = case.collection.repository.get(trace_run.id) + assert after is not None + assert after.outer_trace == before.outer_trace + assert case.scheduled == [] + with case.collection.sessions() as session: + assert session.scalar(select(func.count()).select_from(ApplicationTraceRun)) == 1 + + @staticmethod + @pytest.mark.parametrize("mode", [AppMode.AGENT, AppMode.ADVANCED_CHAT]) + def test_snapshot_failure_is_persisted_and_never_falls_back( + sqlite_session_factory: sessionmaker[Session], monkeypatch: pytest.MonkeyPatch, mode: AppMode + ) -> None: + case = setup(sqlite_session_factory, monkeypatch, "failure", mode, snapshot=False) + + def fail_start(**_kwargs: object) -> None: + raise OperationalError("snapshot", {}, Exception("temporary database unavailable")) + + monkeypatch.setattr(case.collection.repository, "start", fail_start) + context = AgentAppRuntimeBuildContext( + dify_context=DifyRunContext( + tenant_id=case.collection.run.tenant_id, + app_id=case.collection.run.app_id, + user_id=identifier("user"), + user_from=UserFrom.ACCOUNT, + invoke_from=InvokeFrom.DEBUGGER, + ), + agent_id=identifier("agent"), + idempotency_key=case.collection.run.source_id, + conversation_id="conversation", + agent_config_snapshot_id=identifier("snapshot"), + agent_soul=AgentSoulConfig(), + user_query="query", + binding_id=identifier("binding"), + backend_binding_ref="backend", + ) + assert runtime_bridge.prepare_agent_app(context) is None + assert dispatch(case, message_info(case), monkeypatch) == [] + stored = case.collection.repository.source_run( + tenant_id=case.collection.run.tenant_id, + app_id=case.collection.run.app_id, + source_kind="message", + source_id=case.collection.run.source_id, + ) + assert stored is not None + assert stored.state == "failed" + assert stored.error_code == "tracing_initialization_failed" + assert stored.provider_config == {} + + @staticmethod + @pytest.mark.parametrize("mode", [AppMode.AGENT, AppMode.ADVANCED_CHAT, AppMode.CHAT]) + @pytest.mark.parametrize("kind", ["title", "moderation", "aliased_moderation"]) + def test_auxiliary_events_resolve_owner_without_bypassing_content_policy( + sqlite_session_factory: sessionmaker[Session], monkeypatch: pytest.MonkeyPatch, mode: AppMode, kind: str + ) -> None: + case = setup(sqlite_session_factory, monkeypatch, "aux", mode, snapshot=False) + with case.collection.sessions() as session: + message = session.get(Message, case.collection.run.source_id) + assert message is not None + if kind == "title": + info: BaseTraceInfo = GenerateNameTraceInfo( + tenant_id=case.collection.run.tenant_id, + conversation_id=message.conversation_id, + metadata={"conversation_id": message.conversation_id}, + inputs="private conversation", + outputs="private generated title", + ) + else: + info = ModerationTraceInfo( + message_id=identifier("workflow-app-log") if kind == "aliased_moderation" else message.id, + message_data=message.to_dict(session=session), + metadata={"message_id": message.id}, + flagged=False, + action="direct_output", + preset_response="", + query="private prompt", + inputs="private prompt", + outputs="private output", + ) + legacy = dispatch(case, info, monkeypatch) + assert len(legacy) == (1 if mode is AppMode.CHAT else 0) + assert case.scheduled == [] + + @staticmethod + def test_plain_workflow_preserves_legacy_parent_context( + sqlite_session_factory: sessionmaker[Session], monkeypatch: pytest.MonkeyPatch + ) -> None: + parent_case = setup(sqlite_session_factory, monkeypatch, "plain-parent", AppMode.WORKFLOW, snapshot=False) + parent, _ = workflow(parent_case, "plain-parent", has_agent=False, snapshot=False) + child_case = setup(sqlite_session_factory, monkeypatch, "plain-child", AppMode.WORKFLOW, snapshot=False) + child, _ = workflow(child_case, "plain-child", has_agent=False, snapshot=False) + info = workflow_info(child, parent) + original = dict(info.metadata) + assert child_case.service.capture_terminal(info) is False + assert info.metadata == original + with child_case.collection.sessions() as session: + assert session.scalar(select(func.count()).select_from(ApplicationTraceRun)) == 0 + + @staticmethod + def test_plain_chatflow_auxiliary_events_keep_the_legacy_route( + sqlite_session_factory: sessionmaker[Session], monkeypatch: pytest.MonkeyPatch + ) -> None: + case = setup(sqlite_session_factory, monkeypatch, "plain-chatflow", AppMode.ADVANCED_CHAT, snapshot=False) + row, _ = workflow(case, "plain-chatflow", has_agent=False, snapshot=False) + with case.collection.sessions.begin() as session: + message = session.get(Message, case.collection.run.source_id) + assert message is not None + message.workflow_run_id = row.id + info = GenerateNameTraceInfo( + tenant_id=row.tenant_id, + conversation_id=message.conversation_id, + metadata={"conversation_id": message.conversation_id}, + inputs="legacy input", + outputs="legacy title", + ) + assert len(dispatch(case, info, monkeypatch)) == 1 + + @staticmethod + def test_legacy_child_of_managed_parent_uses_explicit_link_not_unavailable_context( + sqlite_session_factory: sessionmaker[Session], monkeypatch: pytest.MonkeyPatch + ) -> None: + parent_case = setup(sqlite_session_factory, monkeypatch, "managed-parent", AppMode.WORKFLOW, snapshot=False) + parent, parent_trace = workflow(parent_case, "managed-parent", has_agent=True, snapshot=True) + assert parent_trace is not None + child_case = setup(sqlite_session_factory, monkeypatch, "legacy-child", AppMode.WORKFLOW, snapshot=False) + child, _ = workflow(child_case, "legacy-child", has_agent=False, snapshot=False) + info = workflow_info(child, parent) + assert child_case.service.capture_terminal(info) is False + assert info.resolved_parent_context == (None, None) + assert info.metadata["linked_parent_workflow_run_id"] == parent.id + assert info.metadata["linked_parent_node_execution_id"] == identifier("parent-tool-node") + assert info.metadata["linked_parent_application_trace_id"] == parent_trace.id + assert info.metadata["linked_parent_trace_id"] == parent_trace.trace_id + + @staticmethod + def test_mixed_workflow_exports_normal_llm_and_agent_calls_without_double_usage( + sqlite_session_factory: sessionmaker[Session], monkeypatch: pytest.MonkeyPatch + ) -> None: + case = setup(sqlite_session_factory, monkeypatch, "mixed", AppMode.WORKFLOW, snapshot=False) + row, managed = workflow(case, "mixed", has_agent=True, snapshot=True) + assert managed is not None + agent_node_id = identifier("agent-execution") + owner = identifier("mixed-owner") + with case.collection.sessions.begin() as session: + app = session.get(App, row.app_id) + assert app is not None + app.created_by = owner + session.add(make_account(account_id=owner)) + for index, kind in enumerate(("llm", "agent")): + node_id = identifier("normal-execution") if kind == "llm" else agent_node_id + usage: dict[str, int] = ( + {"prompt_tokens": 10, "completion_tokens": 5, "total_tokens": 15} if kind == "llm" else {} + ) + session.add( + WorkflowNodeExecutionModel( + id=node_id, + node_execution_id=node_id, + tenant_id=row.tenant_id, + app_id=row.app_id, + workflow_id=row.workflow_id, + workflow_run_id=row.id, + triggered_from=WorkflowNodeExecutionTriggeredFrom.WORKFLOW_RUN, + index=index, + node_id=kind, + node_type=kind, + title=kind, + inputs="{}", + outputs="{}", + process_data=json.dumps( + { + "model_mode": "chat", + "model_name": "ordinary-model", + "model_provider": "openai", + "usage": usage, + } + ) + if kind == "llm" + else "{}", + status=WorkflowNodeExecutionStatus.SUCCEEDED, + elapsed_time=1, + execution_metadata=json.dumps( + {"total_tokens": 15 if kind == "llm" else 10, "total_price": "0.02", "currency": "USD"} + ), + created_at=row.created_at, + finished_at=row.finished_at, + created_by=owner, + created_by_role=CreatorUserRole.ACCOUNT, + ) + ) + invocation = case.collection.repository.invocation( + run_id=managed.id, + agent_id=identifier("agent"), + node_execution_id=agent_node_id, + attempt=0, + call_key="agent:0", + ) + spans = [ + AgentTraceSpan( + trace_id=managed.trace_id, + span_id="a" * 16, + parent_span_id=invocation.span_id, + name="agent run", + start_time_ns=1_000_000_000, + end_time_ns=3_000_000_000, + status="ok", + attributes={"gen_ai.operation.name": "invoke_agent"}, + ) + ] + spans.extend( + AgentTraceSpan( + trace_id=managed.trace_id, + span_id=key * 16, + parent_span_id="a" * 16, + name="model call", + start_time_ns=1_000_000_000, + end_time_ns=2_000_000_000, + status="ok", + attributes={ + "gen_ai.operation.name": "chat", + "gen_ai.request.model": "agent-model", + "gen_ai.usage.input_tokens": 3, + "gen_ai.usage.output_tokens": 2, + "dify.llm.total_price": "0.01", + "dify.llm.currency": "USD", + }, + ) + for key in ("b", "c") + ) + case.collection.repository.ingest( + AgentTraceBatch( + run_id=UUID(managed.id), invocation_id=UUID(invocation.id), batch_id=uuid4(), final=True, spans=spans + ), + max_spans=100, + max_bytes=100000, + ) + monkeypatch.setattr(type(db), "engine", property(lambda _db: sqlite_session_factory.kw["bind"])) + assert case.service.capture_terminal(workflow_info(row)) is True + stored = case.collection.repository.get(managed.id) + assert stored is not None + trace = case.service.assemble(stored) + payloads: list[ExportTraceServiceRequest] = [] + + def send(method: str, url: str, **kwargs: object) -> httpx.Response: + body = kwargs["content"] + assert isinstance(body, bytes) + payloads.append(ExportTraceServiceRequest.FromString(body)) + return httpx.Response(200, content=b"", request=httpx.Request(method, url)) + + monkeypatch.setattr("core.helper.ssrf_proxy.make_request", send) + export_trace("langfuse", {"host": "https://langfuse.test", "public_key": "pk", "secret_key": "sk"}, trace) + observations = payloads[0].resource_spans[0].scope_spans[0].spans + attributes = [{item.key: item.value.string_value for item in span.attributes} for span in observations] + generations = [item for item in attributes if item["langfuse.observation.type"] == "generation"] + assert len(generations) == 3 + assert sum(json.loads(item["langfuse.observation.usage_details"])["total"] for item in generations) == 25 + assert sum( + json.loads(item["langfuse.observation.cost_details"])["total"] for item in generations + ) == pytest.approx(0.04) + + +"""Verify execution hook ownership, policy gating, and fail-open business behavior.""" @dataclass @@ -49,41 +531,6 @@ class BridgeCase: ) -@pytest.fixture -def service(monkeypatch: pytest.MonkeyPatch, sqlite_session_factory: sessionmaker[Session]) -> BridgeCase: - apply_config_overrides(monkeypatch, OPS_AGENT_TRACE_ENABLED=True) - target = ApplicationTracingService(ApplicationTraceRepository(sqlite_session_factory)) - case = BridgeCase(target) - - def prepare(**values: object) -> ApplicationTracingContext: - if case.failure is not None: - raise case.failure - case.calls.append(("prepare", values)) - return case.context - - def start(**values: object) -> None: - if case.failure is not None: - raise case.failure - case.calls.append(("start", values)) - - def terminal(_info: BaseTraceInfo) -> bool: - if case.failure is not None: - raise case.failure - return True - - def agent_terminal(**values: object) -> None: - if case.failure is not None: - raise case.failure - case.calls.append(("terminal", values)) - - monkeypatch.setattr(target, "prepare_invocation", prepare) - monkeypatch.setattr(target, "start", start) - monkeypatch.setattr(target, "capture_terminal", terminal) - monkeypatch.setattr(target, "capture_agent_app_terminal", agent_terminal) - monkeypatch.setattr(runtime_bridge, "_service", lambda: target) - return case - - def run_context(app_id: str) -> DifyRunContext: return DifyRunContext( tenant_id="tenant", @@ -108,7 +555,7 @@ def agent_context() -> AgentAppRuntimeBuildContext: ) -def workflow_context(preview: bool) -> WorkflowAgentRuntimeBuildContext: +def _runtime_bridge_workflow_context(preview: bool) -> WorkflowAgentRuntimeBuildContext: return WorkflowAgentRuntimeBuildContext( dify_context=run_context("workflow-app"), agent=Agent( @@ -143,7 +590,7 @@ def workflow_context(preview: bool) -> WorkflowAgentRuntimeBuildContext: ) -def workflow_entity(extras: dict[str, object] | None = None) -> WorkflowAppGenerateEntity: +def _runtime_bridge_workflow_entity(extras: dict[str, object] | None = None) -> WorkflowAppGenerateEntity: return WorkflowAppGenerateEntity.model_construct( app_config=WorkflowUIBasedAppConfig( app_id="app", tenant_id="tenant", app_mode=AppMode.WORKFLOW, workflow_id="workflow" @@ -152,216 +599,521 @@ def workflow_entity(extras: dict[str, object] | None = None) -> WorkflowAppGener ) -def test_disabled_gate_never_constructs_database_service_or_reads_context(monkeypatch: pytest.MonkeyPatch) -> None: - apply_config_overrides(monkeypatch, OPS_AGENT_TRACE_ENABLED=False) - calls = [] +class TestRuntimeBridgeCases: + @staticmethod + @pytest.fixture + def service(monkeypatch: pytest.MonkeyPatch, sqlite_session_factory: sessionmaker[Session]) -> BridgeCase: + apply_config_overrides(monkeypatch, OPS_AGENT_TRACE_ENABLED=True) + target = ApplicationTracingService(ApplicationTraceRepository(sqlite_session_factory)) + case = BridgeCase(target) - def forbidden() -> ApplicationTracingService: - calls.append(True) - raise AssertionError("disabled tracing must not construct a service") + def prepare(**values: object) -> ApplicationTracingContext: + if case.failure is not None: + raise case.failure + case.calls.append(("prepare", values)) + return case.context - monkeypatch.setattr(runtime_bridge, "_service", forbidden) - assert runtime_bridge.prepare_agent_app(cast(AgentAppRuntimeBuildContext, object())) is None - assert runtime_bridge.prepare_workflow_agent(cast(WorkflowAgentRuntimeBuildContext, object()), None) is None - runtime_bridge.freeze_workflow_run( - cast(WorkflowAppGenerateEntity, object()), - workflow_id="workflow", - workflow_run_id="run", - conversation_id=None, - graph_data={}, - ) - assert runtime_bridge.capture_terminal(cast(BaseTraceInfo, object())) is False - assert runtime_bridge.queue_enabled_at_start("app") is False - assert calls == [] + def start(**values: object) -> None: + if case.failure is not None: + raise case.failure + case.calls.append(("start", values)) + def terminal(_info: BaseTraceInfo) -> bool: + if case.failure is not None: + raise case.failure + return True -def test_agent_app_uses_message_identity_and_stable_conversation(service: BridgeCase) -> None: - assert runtime_bridge.prepare_agent_app(agent_context()) is service.context - assert service.calls == [ - ( - "prepare", - { - "tenant_id": "tenant", - "app_id": "agent-app", - "agent_id": "agent", - "source_kind": "message", - "source_id": "message-id", - "session_id": "conversation", - "node_execution_id": None, - "attempt": 0, - "call_key": "message-id", - }, + def agent_terminal(**values: object) -> None: + if case.failure is not None: + raise case.failure + case.calls.append(("terminal", values)) + + monkeypatch.setattr(target, "prepare_invocation", prepare) + monkeypatch.setattr(target, "start", start) + monkeypatch.setattr(target, "capture_terminal", terminal) + monkeypatch.setattr(target, "capture_agent_app_terminal", agent_terminal) + monkeypatch.setattr(runtime_bridge, "_service", lambda: target) + return case + + @staticmethod + def test_disabled_gate_never_constructs_database_service_or_reads_context(monkeypatch: pytest.MonkeyPatch) -> None: + apply_config_overrides(monkeypatch, OPS_AGENT_TRACE_ENABLED=False) + calls = [] + + def forbidden() -> ApplicationTracingService: + calls.append(True) + raise AssertionError("disabled tracing must not construct a service") + + monkeypatch.setattr(runtime_bridge, "_service", forbidden) + assert runtime_bridge.prepare_agent_app(cast(AgentAppRuntimeBuildContext, object())) is None + assert runtime_bridge.prepare_workflow_agent(cast(WorkflowAgentRuntimeBuildContext, object()), None) is None + runtime_bridge.freeze_workflow_run( + cast(WorkflowAppGenerateEntity, object()), + workflow_id="workflow", + workflow_run_id="run", + conversation_id=None, + graph_data={}, ) - ] + assert runtime_bridge.capture_terminal(cast(BaseTraceInfo, object())) is False + assert runtime_bridge.queue_enabled_at_start("app") is False + assert calls == [] + + @staticmethod + def test_agent_app_uses_message_identity_and_stable_conversation(service: BridgeCase) -> None: + assert runtime_bridge.prepare_agent_app(agent_context()) is service.context + assert service.calls == [ + ( + "prepare", + { + "tenant_id": "tenant", + "app_id": "agent-app", + "agent_id": "agent", + "source_kind": "message", + "source_id": "message-id", + "session_id": "conversation", + "node_execution_id": None, + "attempt": 0, + "call_key": "message-id", + }, + ) + ] + + @staticmethod + @pytest.mark.parametrize("preview", [False, True]) + def test_workflow_invocation_uses_owner_app_and_execution_attempt(service: BridgeCase, preview: bool) -> None: + assert runtime_bridge.prepare_workflow_agent(_runtime_bridge_workflow_context(preview), None) is service.context + values = service.calls[0][1] + assert values["app_id"] == "workflow-app" + assert values["agent_id"] == "referenced-agent" + assert values["source_kind"] == ("node" if preview else "workflow") + assert values["source_id"] == ("node-execution" if preview else "run") + assert values["session_id"] == values["source_id"] + assert values["call_key"] == "node-execution:2" + assert values["workflow_id"] == "workflow" + + @staticmethod + def test_workflow_start_prefers_explicit_trace_session(service: BridgeCase) -> None: + entity = _runtime_bridge_workflow_entity({"trace_session_id": "external-session"}) + runtime_bridge.freeze_workflow_run( + entity, + workflow_id="workflow", + workflow_run_id="run", + conversation_id="conversation", + graph_data={"nodes": [{"data": {"type": "agent", "version": "2", "agent_node_kind": "dify_agent"}}]}, + ) + assert service.calls == [ + ( + "start", + { + "tenant_id": "tenant", + "app_id": "app", + "source_kind": "workflow", + "source_id": "run", + "session_id": "external-session", + "workflow_id": "workflow", + }, + ) + ] + + @staticmethod + def test_business_hook_failure_logs_type_without_sensitive_exception_message( + service: BridgeCase, caplog: pytest.LogCaptureFixture + ) -> None: + service.failure = ValueError("password=private-secret") + assert runtime_bridge.prepare_agent_app(agent_context()) is None + assert "ValueError" in caplog.text + assert "private-secret" not in caplog.text + + @staticmethod + def test_worker_database_failure_retries_without_legacy_fallback(service: BridgeCase) -> None: + service.failure = OperationalError("statement", {}, Exception("sensitive database error")) + with pytest.raises(RetryableTraceDispatchError, match="persistence temporarily unavailable"): + runtime_bridge.capture_terminal(cast(BaseTraceInfo, object())) + + @staticmethod + def test_ordinary_workflow_does_not_create_managed_snapshot(service: BridgeCase) -> None: + runtime_bridge.freeze_workflow_run( + cast(WorkflowAppGenerateEntity, object()), + workflow_id="workflow", + workflow_run_id="run", + conversation_id=None, + graph_data={"nodes": [{"data": {"type": "llm"}}, {"data": {"type": "agent", "version": "1"}}]}, + ) + assert service.calls == [] + + @staticmethod + @pytest.mark.parametrize( + "graph", ["not-json", "null", "[]", {}, {"nodes": {}}, {"nodes": [None, {"data": {"type": "llm"}}]}] + ) + def test_invalid_or_ordinary_graph_cannot_enable_agent_collection(graph: object) -> None: + assert runtime_bridge.workflow_has_agent_v2(graph) is False + + @staticmethod + def test_graph_string_with_agent_is_recognized() -> None: + assert ( + runtime_bridge.workflow_has_agent_v2( + '{"nodes":[{"data":{"type":"agent","version":"2","agent_node_kind":"dify_agent"}}]}' + ) + is True + ) + + @staticmethod + @pytest.mark.parametrize("outcome", ["succeeded", "failed", "cancelled"]) + def test_worker_terminal_handoff_preserves_outcome_without_sse(service: BridgeCase, outcome: str) -> None: + from typing import Literal, cast + + runtime_bridge.capture_agent_app_terminal( + tenant_id="tenant", + app_id="agent-app", + message_id="message", + output="answer", + outcome=cast(Literal["succeeded", "failed", "cancelled"], outcome), + ) + assert service.calls == [ + ( + "terminal", + { + "tenant_id": "tenant", + "app_id": "agent-app", + "message_id": "message", + "outcome": outcome, + "output": "answer", + }, + ) + ] + + @staticmethod + def test_terminal_handoff_gate_and_fail_open_log_redact_sensitive_details( + service: BridgeCase, monkeypatch: pytest.MonkeyPatch, caplog: pytest.LogCaptureFixture + ) -> None: + from repositories.application_trace_repository import TraceCollectionError + + apply_config_overrides(monkeypatch, OPS_AGENT_TRACE_ENABLED=False) + runtime_bridge.capture_agent_app_terminal( + tenant_id="tenant", app_id="app", message_id="message", outcome="failed", output=None + ) + assert service.calls == [] + apply_config_overrides(monkeypatch, OPS_AGENT_TRACE_ENABLED=True) + service.failure = TraceCollectionError("safe_error_code") + runtime_bridge.capture_agent_app_terminal( + tenant_id="tenant", app_id="app", message_id="message", outcome="failed", output="private-content" + ) + assert "safe_error_code" in caplog.text + assert "private-content" not in caplog.text + + @staticmethod + def test_workflow_fail_open_logs_error_kind_without_content( + service: BridgeCase, caplog: pytest.LogCaptureFixture + ) -> None: + service.failure = ValueError("private-prompt") + assert runtime_bridge.prepare_workflow_agent(_runtime_bridge_workflow_context(False), "conversation") is None + runtime_bridge.freeze_workflow_run( + _runtime_bridge_workflow_entity(), + workflow_id="workflow", + workflow_run_id="run", + conversation_id="conversation", + graph_data={"nodes": [{"data": {"type": "agent", "version": "2", "agent_node_kind": "dify_agent"}}]}, + ) + assert "workflow_agent" in caplog.text + assert "workflow_start" in caplog.text + assert "ValueError" in caplog.text + assert "private-prompt" not in caplog.text + + @staticmethod + @pytest.mark.parametrize( + ("policy", "expected"), + [ + (None, False), + ('{"enabled":true}', True), + ('{"enabled":"true"}', False), + ("[]", False), + ("invalid-private", False), + ], + ) + def test_queue_policy_reads_real_app_rows_without_constructing_vendor_sdk( + sqlite_session_factory: sessionmaker[Session], + monkeypatch: pytest.MonkeyPatch, + caplog: pytest.LogCaptureFixture, + policy: str | None, + expected: bool, + ) -> None: + from extensions.ext_database import db + from tests.helpers.unit.model_factories import make_app + + apply_config_overrides(monkeypatch, OPS_AGENT_TRACE_ENABLED=True) + with sqlite_session_factory.begin() as session: + session.add(make_app(app_id="queue-app", tenant_id="tenant", tracing=policy)) + with sqlite_session_factory() as session: + monkeypatch.setattr(db, "session", session) + assert runtime_bridge.queue_enabled_at_start("queue-app") is expected + assert runtime_bridge.queue_enabled_at_start("missing") is False + assert runtime_bridge.queue_enabled_at_start(None) is False + assert "invalid-private" not in caplog.text -@pytest.mark.parametrize("preview", [False, True]) -def test_workflow_invocation_uses_owner_app_and_execution_attempt(service: BridgeCase, preview: bool) -> None: - assert runtime_bridge.prepare_workflow_agent(workflow_context(preview), None) is service.context - values = service.calls[0][1] - assert values["app_id"] == "workflow-app" - assert values["agent_id"] == "referenced-agent" - assert values["source_kind"] == ("node" if preview else "workflow") - assert values["source_id"] == ("node-execution" if preview else "run") - assert values["session_id"] == values["source_id"] - assert values["call_key"] == "node-execution:2" - assert values["workflow_id"] == "workflow" +"""Hidden Workflow Tool Agents must not invent a trace owned by the source app.""" -def test_workflow_start_prefers_explicit_trace_session(service: BridgeCase) -> None: - entity = workflow_entity({"trace_session_id": "external-session"}) +_workflow_tool_tracing_boundary_workflow_context = _runtime_bridge_workflow_context +_workflow_tool_tracing_boundary_workflow_entity = _runtime_bridge_workflow_entity + + +def _workflow_tool_tracing_boundary_seed( + sessions: sessionmaker[Session], + monkeypatch: pytest.MonkeyPatch, + *, + root_has_agent: bool, + source_enabled: bool = True, +) -> tuple[ApplicationTracingService, WorkflowAgentRuntimeBuildContext, dict[str, str]]: + ids = { + name: identifier("tool-boundary-" + name) + for name in [ + "tenant", + "root", + "source", + "root_workflow", + "source_workflow", + "run", + "root_agent", + "root_tool", + "child", + "actor", + ] + } + root_graph: dict[str, list[dict[str, object]]] = { + "nodes": [{"id": "tool", "data": {"type": "tool", "title": "Root tool"}}], + "edges": [], + } + if root_has_agent: + root_graph["nodes"].append( + { + "id": "root-agent", + "data": {"type": "agent", "version": "2", "agent_node_kind": "dify_agent", "title": "Root Agent"}, + } + ) + now = datetime(2026, 10, 21, tzinfo=UTC) + with sessions.begin() as session: + for app, provider, enabled, content in [ + ("root", "langfuse", True, False), + ("source", "langsmith", source_enabled, True), + ]: + session.add( + make_app( + app_id=ids[app], + tenant_id=ids["tenant"], + mode=AppMode.WORKFLOW, + tracing=json.dumps({"enabled": enabled, "tracing_provider": provider, "capture_content": content}), + ) + ) + session.add( + make_workflow( + workflow_id=ids[app + "_workflow"], + app_id=ids[app], + tenant_id=ids["tenant"], + graph=root_graph if app == "root" else {"nodes": [], "edges": []}, + ) + ) + session.add( + TraceAppConfig( + app_id=ids[app], + tracing_provider=provider, + tracing_config={"project_url": "https://tracing.test/" + app}, + ) + ) + session.add( + WorkflowRun( + id=ids["run"], + tenant_id=ids["tenant"], + app_id=ids["root"], + workflow_id=ids["root_workflow"], + type=_workflow_tool_tracing_boundary_WorkflowType.WORKFLOW, + triggered_from=_workflow_tool_tracing_boundary_WorkflowRunTriggeredFrom.APP_RUN, + version="1", + graph=json.dumps(root_graph), + inputs="{}", + status=WorkflowExecutionStatus.SUCCEEDED, + created_by=ids["actor"], + created_by_role=CreatorUserRole.ACCOUNT, + created_at=now, + finished_at=now + timedelta(seconds=3), + ) + ) + for key, node_type in [("root_agent", "agent"), ("root_tool", "tool"), ("child", "agent")]: + child = key == "child" + app = "source" if child else "root" + session.add( + WorkflowNodeExecutionModel( + id=ids[key], + tenant_id=ids["tenant"], + app_id=ids[app], + workflow_id=ids[app + "_workflow"], + workflow_run_id=ids["run"], + triggered_from=WorkflowNodeExecutionTriggeredFrom.WORKFLOW_TOOL + if child + else WorkflowNodeExecutionTriggeredFrom.WORKFLOW_RUN, + triggered_from_workflow_id=ids["root_workflow"] if child else None, + triggered_from_node_execution_id=ids["root_tool"] if child else None, + node_execution_id=ids[key], + index=1, + node_id=key, + node_type=node_type, + title=key, + status=WorkflowNodeExecutionStatus.SUCCEEDED, + created_by=ids["actor"], + created_by_role=CreatorUserRole.ACCOUNT, + created_at=now, + finished_at=now + timedelta(seconds=2), + inputs="{}", + outputs="{}", + process_data=json.dumps( + {"workflow_tool_root_app_id": ids["root"], "workflow_tool_invocation_id": "tool-call"} + ) + if child + else "{}", + ) + ) + service = ApplicationTracingService(ApplicationTraceRepository(sessions)) + apply_config_overrides(monkeypatch, OPS_AGENT_TRACE_ENABLED=True) + monkeypatch.setattr(runtime_bridge, "_service", lambda: service) + entity = _workflow_tool_tracing_boundary_workflow_entity() + entity.app_config = entity.app_config.model_copy( + update={"app_id": ids["root"], "tenant_id": ids["tenant"], "workflow_id": ids["root_workflow"]} + ) runtime_bridge.freeze_workflow_run( entity, - workflow_id="workflow", - workflow_run_id="run", - conversation_id="conversation", - graph_data={"nodes": [{"data": {"type": "agent", "version": "2", "agent_node_kind": "dify_agent"}}]}, - ) - assert service.calls == [ - ( - "start", - { - "tenant_id": "tenant", - "app_id": "app", - "source_kind": "workflow", - "source_id": "run", - "session_id": "external-session", - "workflow_id": "workflow", - }, - ) - ] - - -def test_business_hook_failure_logs_type_without_sensitive_exception_message( - service: BridgeCase, caplog: pytest.LogCaptureFixture -) -> None: - service.failure = ValueError("password=private-secret") - assert runtime_bridge.prepare_agent_app(agent_context()) is None - assert "ValueError" in caplog.text - assert "private-secret" not in caplog.text - - -def test_worker_database_failure_retries_without_legacy_fallback(service: BridgeCase) -> None: - service.failure = OperationalError("statement", {}, Exception("sensitive database error")) - with pytest.raises(RetryableTraceDispatchError, match="persistence temporarily unavailable"): - runtime_bridge.capture_terminal(cast(BaseTraceInfo, object())) - - -def test_ordinary_workflow_does_not_create_managed_snapshot(service: BridgeCase) -> None: - runtime_bridge.freeze_workflow_run( - cast(WorkflowAppGenerateEntity, object()), - workflow_id="workflow", - workflow_run_id="run", + workflow_id=ids["root_workflow"], + workflow_run_id=ids["run"], conversation_id=None, - graph_data={"nodes": [{"data": {"type": "llm"}}, {"data": {"type": "agent", "version": "1"}}]}, + graph_data=root_graph, ) - assert service.calls == [] + original = _workflow_tool_tracing_boundary_workflow_context(False) + context = replace( + original, + dify_context=original.dify_context.model_copy( + update={"tenant_id": ids["tenant"], "app_id": ids["source"], "workflow_tool_invocation_id": "tool-call"} + ), + workflow_id=ids["source_workflow"], + workflow_run_id=ids["run"], + node_execution_id=ids["child"], + ) + return service, context, ids -@pytest.mark.parametrize( - "graph", ["not-json", "null", "[]", {}, {"nodes": {}}, {"nodes": [None, {"data": {"type": "llm"}}]}] -) -def test_invalid_or_ordinary_graph_cannot_enable_agent_collection(graph: object) -> None: - assert runtime_bridge.workflow_has_agent_v2(graph) is False - - -def test_graph_string_with_agent_is_recognized() -> None: - assert ( - runtime_bridge.workflow_has_agent_v2( - '{"nodes":[{"data":{"type":"agent","version":"2","agent_node_kind":"dify_agent"}}]}' +class TestWorkflowToolTracingBoundaryCases: + @staticmethod + @pytest.mark.parametrize("root_has_agent", [False, True]) + @pytest.mark.parametrize("source_enabled", [False, True]) + def test_hidden_child_never_reserves_source_policy_or_changes_root_freeze( + sqlite_session_factory: sessionmaker[Session], + monkeypatch: pytest.MonkeyPatch, + root_has_agent: bool, + source_enabled: bool, + ) -> None: + service, context, ids = _workflow_tool_tracing_boundary_seed( + sqlite_session_factory, monkeypatch, root_has_agent=root_has_agent, source_enabled=source_enabled ) - is True - ) + with sqlite_session_factory() as session: + before = deepcopy([dict(row) for row in session.execute(select(ApplicationTraceRun.__table__)).mappings()]) + service_lookups: list[bool] = [] + def lookup() -> ApplicationTracingService: + service_lookups.append(True) + return service -@pytest.mark.parametrize("outcome", ["succeeded", "failed", "cancelled"]) -def test_worker_terminal_handoff_preserves_outcome_without_sse(service: BridgeCase, outcome: str) -> None: - from typing import Literal, cast - - runtime_bridge.capture_agent_app_terminal( - tenant_id="tenant", - app_id="agent-app", - message_id="message", - output="answer", - outcome=cast(Literal["succeeded", "failed", "cancelled"], outcome), - ) - assert service.calls == [ - ( - "terminal", - { - "tenant_id": "tenant", - "app_id": "agent-app", - "message_id": "message", - "outcome": outcome, - "output": "answer", - }, + monkeypatch.setattr(runtime_bridge, "_service", lookup) + # Marked children never consult policy or ownership, even for malformed + # empty markers or a foreign tenant. They cannot create an orphan B run. + for tenant in (ids["tenant"], identifier("foreign-tool-tenant")): + for marker in ("tool-call", ""): + child = replace( + context, + dify_context=context.dify_context.model_copy( + update={ + "tenant_id": tenant, + "workflow_tool_invocation_id": marker, + } + ), + ) + assert runtime_bridge.prepare_workflow_agent(child, None) is None + assert service_lookups == [] + with sqlite_session_factory() as session: + observed = [ + (r.app_id, r.provider, r.capture_content, r.source_id) + for r in session.scalars(select(ApplicationTraceRun)) + ] + assert observed == ([(ids["root"], "langfuse", False, ids["run"])] if root_has_agent else []) + after = [dict(row) for row in session.execute(select(ApplicationTraceRun.__table__)).mappings()] + assert after == before + assert ( + service.repository.source_run( + tenant_id=ids["tenant"], app_id=ids["source"], source_kind="workflow", source_id=ids["run"] + ) + is None ) - ] - -def test_terminal_handoff_gate_and_fail_open_log_redact_sensitive_details( - service: BridgeCase, monkeypatch: pytest.MonkeyPatch, caplog: pytest.LogCaptureFixture -) -> None: - from repositories.application_trace_repository import TraceCollectionError - - apply_config_overrides(monkeypatch, OPS_AGENT_TRACE_ENABLED=False) - runtime_bridge.capture_agent_app_terminal( - tenant_id="tenant", app_id="app", message_id="message", outcome="failed", output=None - ) - assert service.calls == [] - apply_config_overrides(monkeypatch, OPS_AGENT_TRACE_ENABLED=True) - service.failure = TraceCollectionError("safe_error_code") - runtime_bridge.capture_agent_app_terminal( - tenant_id="tenant", app_id="app", message_id="message", outcome="failed", output="private-content" - ) - assert "safe_error_code" in caplog.text - assert "private-content" not in caplog.text - - -def test_workflow_fail_open_logs_error_kind_without_content( - service: BridgeCase, caplog: pytest.LogCaptureFixture -) -> None: - service.failure = ValueError("private-prompt") - assert runtime_bridge.prepare_workflow_agent(workflow_context(False), "conversation") is None - runtime_bridge.freeze_workflow_run( - workflow_entity(), - workflow_id="workflow", - workflow_run_id="run", - conversation_id="conversation", - graph_data={"nodes": [{"data": {"type": "agent", "version": "2", "agent_node_kind": "dify_agent"}}]}, - ) - assert "workflow_agent" in caplog.text - assert "workflow_start" in caplog.text - assert "ValueError" in caplog.text - assert "private-prompt" not in caplog.text - - -@pytest.mark.parametrize( - ("policy", "expected"), - [ - (None, False), - ('{"enabled":true}', True), - ('{"enabled":"true"}', False), - ("[]", False), - ("invalid-private", False), - ], -) -def test_queue_policy_reads_real_app_rows_without_constructing_vendor_sdk( - sqlite_session_factory: sessionmaker[Session], - monkeypatch: pytest.MonkeyPatch, - caplog: pytest.LogCaptureFixture, - policy: str | None, - expected: bool, -) -> None: - from extensions.ext_database import db - from tests.unit_tests.model_factories import make_app - - apply_config_overrides(monkeypatch, OPS_AGENT_TRACE_ENABLED=True) - with sqlite_session_factory.begin() as session: - session.add(make_app(app_id="queue-app", tenant_id="tenant", tracing=policy)) - with sqlite_session_factory() as session: - monkeypatch.setattr(db, "session", session) - assert runtime_bridge.queue_enabled_at_start("queue-app") is expected - assert runtime_bridge.queue_enabled_at_start("missing") is False - assert runtime_bridge.queue_enabled_at_start(None) is False - assert "invalid-private" not in caplog.text + @staticmethod + def test_hidden_child_skip_keeps_normal_root_agent_trace_assemblable( + sqlite_session_factory: sessionmaker[Session], + monkeypatch: pytest.MonkeyPatch, + ) -> None: + service, child, ids = _workflow_tool_tracing_boundary_seed( + sqlite_session_factory, monkeypatch, root_has_agent=True + ) + assert runtime_bridge.prepare_workflow_agent(child, None) is None + root_context = replace( + child, + dify_context=child.dify_context.model_copy( + update={"app_id": ids["root"], "workflow_tool_invocation_id": None} + ), + workflow_id=ids["root_workflow"], + node_execution_id=ids["root_agent"], + ) + runtime = runtime_bridge.prepare_workflow_agent(root_context, None) + assert runtime is not None + run = service.repository.get(str(runtime.run_id)) + assert run is not None + assert run.app_id == ids["root"] + frozen = deepcopy(run.provider_config) + with sqlite_session_factory.begin() as session: + owner = session.get(App, ids["root"]) + assert owner is not None + owner.tracing = json.dumps({"enabled": True, "tracing_provider": "langsmith", "capture_content": True}) + workflow = session.get(WorkflowRun, ids["run"]) + assert workflow is not None + actor = make_account(account_id=ids["actor"]) + nodes = SQLAlchemyWorkflowNodeExecutionRepository( + session_factory=sqlite_session_factory, + tenant_id=ids["tenant"], + user=actor, + app_id=ids["root"], + triggered_from=WorkflowNodeExecutionTriggeredFrom.WORKFLOW_RUN, + ) + info = WorkflowTraceInfo( + workflow_data=workflow, + workflow_id=ids["root_workflow"], + tenant_id=ids["tenant"], + workflow_run_id=ids["run"], + workflow_run_elapsed_time=3, + workflow_run_status="succeeded", + workflow_run_inputs={}, + workflow_run_outputs={}, + workflow_run_version="1", + total_tokens=0, + file_list=[], + query="", + metadata={"app_id": ids["root"]}, + ) + outer = CanonicalTraceBuilder(lambda _info: nodes.get_by_workflow_execution(ids["run"])).build(info) + assert outer is not None + outer = remap_outer_trace(outer, trace_id=run.trace_id, capture_content=False, agent_message=False) + service.repository.terminal(run.id, outer.model_dump(mode="json"), outer.root_span_id) + invocation = service.repository.collected(run.id)[0][0] + ingest(Collection(service.repository, sqlite_session_factory, run), batch_for(run, invocation)) + current = service.repository.get(run.id) + assert current is not None + trace = service.assemble(current) + span_ids = {span.id for span in trace.spans} + assert all(span.parent_id is None or span.parent_id in span_ids for span in trace.spans) + assert current.provider_config == frozen + assert current.provider == "langfuse" + assert current.capture_content is False diff --git a/api/tests/unit_tests/services/tracing/test_service.py b/api/tests/unit_tests/services/tracing/test_service.py new file mode 100644 index 00000000000..d33f8da2223 --- /dev/null +++ b/api/tests/unit_tests/services/tracing/test_service.py @@ -0,0 +1,1052 @@ +import json +from dataclasses import dataclass +from datetime import UTC, datetime, timedelta +from decimal import Decimal +from uuid import NAMESPACE_URL, UUID, uuid4, uuid5 + +import pytest +from dify_agent.protocol.tracing import AgentTraceBatch, AgentTraceEvent, AgentTraceSpan +from sqlalchemy import func, select +from sqlalchemy.orm import Session, sessionmaker + +from core.ops.entities.trace_entity import MessageTraceInfo +from core.ops.exceptions import RetryableTraceDispatchError +from core.ops.unified_trace.entities import CanonicalSpan, CanonicalSpanKind, CanonicalSpanStatus, CanonicalTrace +from graphon.enums import WorkflowNodeExecutionStatus +from libs.datetime_utils import naive_utc_now +from models.application_trace import ApplicationTraceBatch, ApplicationTraceInvocation, ApplicationTraceRun +from models.enums import ConversationFromSource, CreatorUserRole +from models.model import App, AppMode, Message, TraceAppConfig +from models.workflow import WorkflowNodeExecutionModel, WorkflowNodeExecutionTriggeredFrom +from repositories.application_trace_repository import ApplicationTraceRepository, TraceCollectionError +from services.tracing import service as service_module +from services.tracing.provider_dispatch import TraceExportResult, TraceExportUncertainError +from services.tracing.service import ApplicationTracingService +from tests.helpers.unit.config_override import apply_config_overrides +from tests.helpers.unit.model_factories import make_app, make_message, make_workflow + +"""Exercise collection transactions and delivery state against real SQLite rows.""" + + +def _application_trace_collection_identifier(value: str) -> str: + return str(uuid5(NAMESPACE_URL, value)) + + +@dataclass +class _ApplicationTraceCollectionCollection: + repository: ApplicationTraceRepository + sessions: sessionmaker[Session] + run: ApplicationTraceRun + + def invocation(self, name: str = "first") -> ApplicationTraceInvocation: + return self.repository.invocation( + run_id=self.run.id, + agent_id=_application_trace_collection_identifier("agent"), + node_execution_id=None, + attempt=0, + call_key=name, + ) + + +def _application_trace_collection_seed( + sessions: sessionmaker[Session], + name: str, + *, + capture_content: bool = False, + enabled: bool = True, + capture_agent: bool = True, +) -> _ApplicationTraceCollectionCollection: + app_id, tenant_id, message_id = ( + _application_trace_collection_identifier(f"app-{name}"), + _application_trace_collection_identifier("tenant"), + _application_trace_collection_identifier(f"message-{name}"), + ) + with sessions.begin() as session: + session.add( + make_app( + app_id=app_id, + tenant_id=tenant_id, + mode=AppMode.AGENT, + tracing=json.dumps( + { + "enabled": enabled, + "tracing_provider": "langfuse", + "capture_content": capture_content, + "capture_agent": capture_agent, + } + ), + ) + ) + session.add( + make_message( + message_id=message_id, + app_id=app_id, + conversation_id=_application_trace_collection_identifier(f"conversation-{name}"), + inputs={}, + query="hello", + message={}, + answer="done", + message_unit_price=Decimal(0), + answer_unit_price=Decimal(0), + currency="USD", + from_source=ConversationFromSource.CONSOLE, + ) + ) + session.add( + TraceAppConfig( + app_id=app_id, + tracing_provider="langfuse", + tracing_config={ + "host": "https://langfuse.test", + "secret_key": "encrypted-original", + "project_url": "https://langfuse.test/project/original", + }, + ) + ) + repository = ApplicationTraceRepository(sessions) + run = repository.start( + tenant_id=tenant_id, app_id=app_id, source_kind="message", source_id=message_id, session_id="session-original" + ) + assert run is not None + return _ApplicationTraceCollectionCollection(repository, sessions, run) + + +def _application_trace_collection_batch_for( + run: ApplicationTraceRun, + invocation: ApplicationTraceInvocation, + *, + final: bool = True, + batch_id: UUID | None = None, + span_id: str = "a" * 16, +) -> AgentTraceBatch: + return AgentTraceBatch( + run_id=UUID(run.id), + invocation_id=UUID(invocation.id), + batch_id=batch_id or uuid4(), + final=final, + spans=[ + AgentTraceSpan( + trace_id=run.trace_id, + span_id=span_id, + parent_span_id=invocation.span_id, + name="Agent execution", + start_time_ns=1_000_000_000, + end_time_ns=2_000_000_000, + status="ok", + attributes={"gen_ai.operation.name": "invoke_agent"}, + ) + ], + ) + + +def _application_trace_collection_ingest( + collection: _ApplicationTraceCollectionCollection, batch: AgentTraceBatch +) -> tuple[int, str]: + return collection.repository.ingest(batch, max_spans=100, max_bytes=1_000_000) + + +def _application_trace_collection_terminal(collection: _ApplicationTraceCollectionCollection) -> None: + run = collection.run + root = CanonicalSpan( + id=run.root_span_id, + parent_id=None, + name="Agent turn", + kind=CanonicalSpanKind.CHAIN, + start_time=datetime.fromtimestamp(0, UTC), + end_time=datetime.fromtimestamp(3, UTC), + status=CanonicalSpanStatus.OK, + ) + trace = CanonicalTrace(trace_id=run.trace_id, session_id=run.session_id, root_span_id=root.id, spans=(root,)) + collection.repository.terminal(run.id, trace.model_dump(mode="json"), root.id) + + +def _application_trace_collection_claim( + collection: _ApplicationTraceCollectionCollection, +) -> tuple[ApplicationTraceRun, str] | None: + return collection.repository.claim(collection.run.id, collection_timeout=60, lease_seconds=30) + + +class TestApplicationTraceCollectionCases: + @staticmethod + @pytest.fixture + def collection(sqlite_session_factory: sessionmaker[Session]) -> _ApplicationTraceCollectionCollection: + return _application_trace_collection_seed(sqlite_session_factory, "first") + + @staticmethod + def test_batch_duplicate_is_idempotent_after_final_and_acceptance( + collection: _ApplicationTraceCollectionCollection, + ) -> None: + invocation = collection.invocation() + batch = _application_trace_collection_batch_for(collection.run, invocation) + assert _application_trace_collection_ingest(collection, batch) == (1, collection.run.id) + assert _application_trace_collection_ingest(collection, batch) == (1, collection.run.id) + _application_trace_collection_terminal(collection) + leased = _application_trace_collection_claim(collection) + assert leased is not None + collection.repository.complete( + collection.run.id, leased[1], state="accepted", result=TraceExportResult("provider-id") + ) + assert _application_trace_collection_ingest(collection, batch) == (1, collection.run.id) + stored = collection.repository.get(collection.run.id) + assert stored is not None + assert stored.span_count == 1 + with collection.sessions() as session: + assert session.scalar(select(func.count()).select_from(ApplicationTraceBatch)) == 1 + + @staticmethod + def test_conflicting_batch_id_rolls_back_without_replacing_payload( + collection: _ApplicationTraceCollectionCollection, + ) -> None: + batch = _application_trace_collection_batch_for(collection.run, collection.invocation(), final=False) + _application_trace_collection_ingest(collection, batch) + changed = batch.model_copy(update={"final": True}) + with pytest.raises(TraceCollectionError, match="tracing_batch_conflict"): + _application_trace_collection_ingest(collection, changed) + invocations, batches = collection.repository.collected(collection.run.id) + assert invocations[0].finalized is False + assert batches[0].final is False + + @staticmethod + def test_run_cannot_ingest_another_runs_invocation(collection: _ApplicationTraceCollectionCollection) -> None: + other = _application_trace_collection_seed(collection.sessions, "other") + foreign = other.invocation() + forged = _application_trace_collection_batch_for(collection.run, foreign) + with pytest.raises(TraceCollectionError, match="tracing_invocation_not_found"): + _application_trace_collection_ingest(collection, forged) + with collection.sessions() as session: + assert session.scalar(select(func.count()).select_from(ApplicationTraceBatch)) == 0 + + @staticmethod + def test_payload_limits_reject_whole_transaction(collection: _ApplicationTraceCollectionCollection) -> None: + batch = _application_trace_collection_batch_for(collection.run, collection.invocation()) + with pytest.raises(TraceCollectionError, match="tracing_payload_limit"): + collection.repository.ingest(batch, max_spans=100, max_bytes=1) + invocations, batches = collection.repository.collected(collection.run.id) + assert not batches + assert invocations[0].finalized is False + stored = collection.repository.get(collection.run.id) + assert stored is not None + assert stored.span_count == 0 + + @staticmethod + def test_wrong_trace_id_is_rejected_before_persistence(collection: _ApplicationTraceCollectionCollection) -> None: + batch = _application_trace_collection_batch_for(collection.run, collection.invocation()) + altered = batch.model_copy(update={"spans": [batch.spans[0].model_copy(update={"trace_id": "f" * 32})]}) + with pytest.raises(TraceCollectionError, match="tracing_trace_id_mismatch"): + _application_trace_collection_ingest(collection, altered) + _, batches = collection.repository.collected(collection.run.id) + assert not batches + + @staticmethod + def test_content_off_cleans_the_durable_batch_not_only_the_export( + collection: _ApplicationTraceCollectionCollection, + ) -> None: + original = _application_trace_collection_batch_for(collection.run, collection.invocation(), final=False) + observation = original.spans[0].model_copy( + update={ + "status": "error", + "status_message": "private error", + "attributes": { + "gen_ai.operation.name": "chat", + "gen_ai.request.model": "model-one", + "gen_ai.input.messages": "private prompt", + "gen_ai.tool.call.arguments": "private arguments", + "dify.payload": "private metadata", + "authorization": "private credential", + }, + "events": [ + AgentTraceEvent( + name="exception", timestamp_ns=1_500_000_000, attributes={"exception.message": "private event"} + ) + ], + } + ) + batch = original.model_copy(update={"spans": [observation]}) + _application_trace_collection_ingest(collection, batch) + _, batches = collection.repository.collected(collection.run.id) + assert "private" not in json.dumps(batches[0].spans) + assert batches[0].spans[0]["events"] == [] + assert batches[0].spans[0]["status"] == "error" + assert batches[0].spans[0]["attributes"]["gen_ai.request.model"] == "model-one" + # A changed original payload is a conflict even if content-off sanitization + # would produce exactly the same persisted span. + changed = batch.model_copy( + update={"spans": [observation.model_copy(update={"status_message": "different private error"})]} + ) + with pytest.raises(TraceCollectionError, match="tracing_batch_conflict"): + _application_trace_collection_ingest(collection, changed) + + @staticmethod + def test_invocation_key_replay_must_not_change_agent_binding( + collection: _ApplicationTraceCollectionCollection, + ) -> None: + original = collection.invocation() + repeated = collection.invocation() + assert original.id == repeated.id + with pytest.raises(TraceCollectionError, match="tracing_invocation_conflict"): + collection.repository.invocation( + run_id=collection.run.id, + agent_id=_application_trace_collection_identifier("different-agent"), + node_execution_id=_application_trace_collection_identifier("different-node"), + attempt=2, + call_key="first", + ) + + @staticmethod + def test_claim_waits_for_all_invocations_final_and_then_only_one_lease( + collection: _ApplicationTraceCollectionCollection, + ) -> None: + first, second = collection.invocation("first"), collection.invocation("second") + _application_trace_collection_terminal(collection) + _application_trace_collection_ingest(collection, _application_trace_collection_batch_for(collection.run, first)) + assert _application_trace_collection_claim(collection) is None + _application_trace_collection_ingest( + collection, _application_trace_collection_batch_for(collection.run, second, span_id="b" * 16) + ) + leased = _application_trace_collection_claim(collection) + assert leased is not None + assert _application_trace_collection_claim(collection) is None + assert leased[0].attempts == 1 + + @staticmethod + def test_expired_delivery_lease_accepts_confirmed_late_ack_for_same_attempt( + collection: _ApplicationTraceCollectionCollection, + ) -> None: + _application_trace_collection_ingest( + collection, _application_trace_collection_batch_for(collection.run, collection.invocation()) + ) + _application_trace_collection_terminal(collection) + leased = _application_trace_collection_claim(collection) + assert leased is not None + with collection.sessions.begin() as session: + stored = session.get(ApplicationTraceRun, collection.run.id) + assert stored is not None + stored.export_started_at = naive_utc_now() - timedelta(seconds=60) + stored.export_heartbeat_at = stored.export_started_at + assert _application_trace_collection_claim(collection) is None + current = collection.repository.get(collection.run.id) + assert current is not None + assert current.state == "uncertain" + collection.repository.complete(collection.run.id, leased[1], state="accepted", result=TraceExportResult("late")) + stored = collection.repository.get(collection.run.id) + assert stored is not None + assert stored.state == "accepted" + assert stored.error_code is None + assert stored.accepted_at is not None + assert stored.provider_trace_id == "late" + assert _application_trace_collection_claim(collection) is None + + @staticmethod + def test_incomplete_final_times_out_without_export(collection: _ApplicationTraceCollectionCollection) -> None: + collection.invocation() + _application_trace_collection_terminal(collection) + with collection.sessions.begin() as session: + stored = session.get(ApplicationTraceRun, collection.run.id) + assert stored is not None + stored.terminal_at = naive_utc_now() - timedelta(seconds=120) + assert _application_trace_collection_claim(collection) is None + stored = collection.repository.get(collection.run.id) + assert stored is not None + assert stored.state == "failed" + assert stored.error_code == "agent_trace_collection_incomplete" + + @staticmethod + def test_running_trace_freezes_credentials_content_policy_and_session( + collection: _ApplicationTraceCollectionCollection, + ) -> None: + original = collection.run + with collection.sessions.begin() as session: + app = session.get(App, original.app_id) + assert app is not None + app.tracing = json.dumps({"enabled": False, "tracing_provider": "opik", "capture_content": True}) + config = session.scalar(select(TraceAppConfig).where(TraceAppConfig.app_id == original.app_id)) + assert config is not None + config.tracing_config = {"secret_key": "encrypted-new"} + repeated = collection.repository.start( + tenant_id=original.tenant_id, + app_id=original.app_id, + source_kind="message", + source_id=original.source_id, + session_id="changed", + ) + assert repeated is not None + assert repeated.provider == "langfuse" + assert repeated.provider_config["secret_key"] == "encrypted-original" + assert repeated.capture_content is False + assert repeated.config_revision == original.config_revision + assert repeated.session_id == "session-original" + + @staticmethod + @pytest.mark.parametrize("outcome", ["accepted", "uncertain", "retrying"]) + def test_delivery_checkpoint_controls_replays( + collection: _ApplicationTraceCollectionCollection, monkeypatch: pytest.MonkeyPatch, outcome: str + ) -> None: + _application_trace_collection_ingest( + collection, _application_trace_collection_batch_for(collection.run, collection.invocation()) + ) + _application_trace_collection_terminal(collection) + exports: list[tuple[str, dict[str, object], CanonicalTrace]] = [] + decryptions: list[dict[str, object]] = [] + + def decrypt(_tenant: str, _provider: str, config: dict[str, object]) -> dict[str, object]: + decryptions.append(config) + return {"secret_key": "plaintext-only-in-worker"} + + def export(provider: str, config: dict[str, object], trace: CanonicalTrace) -> TraceExportResult: + exports.append((provider, config, trace)) + if outcome == "uncertain": + raise TraceExportUncertainError("ack lost") + if outcome == "retrying": + raise RetryableTraceDispatchError("busy") + return TraceExportResult("vendor-trace", "https://trace.test/example") + + monkeypatch.setattr(service_module.OpsTraceManager, "decrypt_tracing_config", decrypt) + monkeypatch.setattr(service_module, "export_trace", export) + service = ApplicationTracingService(collection.repository) + service.deliver(collection.run.id) + service.deliver(collection.run.id) + assert len(exports) == 1 + assert decryptions[0]["secret_key"] == "encrypted-original" + stored = collection.repository.get(collection.run.id) + assert stored is not None + assert stored.state == outcome + assert stored.provider_config["secret_key"] == "encrypted-original" + assert "plaintext-only-in-worker" not in json.dumps(stored.outer_trace) + assert exports[0][1]["host"] == "https://langfuse.test" + + @staticmethod + def test_source_owner_checks_app_and_tenant(collection: _ApplicationTraceCollectionCollection) -> None: + run = collection.run + assert ( + collection.repository.start( + tenant_id=_application_trace_collection_identifier("foreign-tenant"), + app_id=run.app_id, + source_kind="message", + source_id=run.source_id, + session_id="x", + ) + is None + ) + other = _application_trace_collection_seed(collection.sessions, "second-app") + assert ( + collection.repository.start( + tenant_id=run.tenant_id, + app_id=other.run.app_id, + source_kind="message", + source_id=run.source_id, + session_id="x", + ) + is None + ) + + @staticmethod + def test_disabled_source_decision_stays_disabled_after_config_enable( + sqlite_session_factory: sessionmaker[Session], monkeypatch: pytest.MonkeyPatch + ) -> None: + apply_config_overrides(monkeypatch, OPS_AGENT_TRACE_ENABLED=True) + case = _application_trace_collection_seed(sqlite_session_factory, "initially-disabled", enabled=False) + service = ApplicationTracingService(case.repository) + assert case.run.state == "not_collected" + with case.sessions.begin() as session: + app = session.get(App, case.run.app_id) + assert app is not None + app.tracing = json.dumps({"enabled": True, "tracing_provider": "langfuse", "capture_content": True}) + message = session.get(Message, case.run.source_id) + assert message is not None + assert ( + service.start( + tenant_id=case.run.tenant_id, + app_id=case.run.app_id, + source_kind="message", + source_id=case.run.source_id, + session_id="later-session", + ) + is None + ) + info = MessageTraceInfo( + metadata={}, + message_data=message, + message_id=message.id, + conversation_model="agent", + conversation_mode="agent", + message_tokens=1, + answer_tokens=1, + total_tokens=2, + ) + assert service.capture_terminal(info) is True + stored = case.repository.get(case.run.id) + assert stored is not None + assert stored.state == "not_collected" + assert stored.provider_config == {} + assert stored.outer_trace is None + + @staticmethod + def test_capture_agent_disabled_preserves_outer_only_content_policy( + sqlite_session_factory: sessionmaker[Session], monkeypatch: pytest.MonkeyPatch + ) -> None: + apply_config_overrides(monkeypatch, OPS_AGENT_TRACE_ENABLED=True) + case = _application_trace_collection_seed(sqlite_session_factory, "outer-only", capture_agent=False) + service = ApplicationTracingService(case.repository) + scheduled: list[str] = [] + monkeypatch.setattr(service, "schedule", scheduled.append) + assert ( + service.prepare_invocation( + tenant_id=case.run.tenant_id, + app_id=case.run.app_id, + agent_id=_application_trace_collection_identifier("agent"), + source_kind="message", + source_id=case.run.source_id, + session_id="session", + node_execution_id=None, + attempt=0, + call_key="call", + ) + is None + ) + with case.sessions() as session: + message = session.get(Message, case.run.source_id) + assert message is not None + info = MessageTraceInfo( + metadata={}, + message_data=message, + message_id=message.id, + conversation_model="agent", + conversation_mode="agent", + message_tokens=1, + answer_tokens=1, + total_tokens=2, + inputs="private input", + outputs="private output", + ) + assert service.capture_terminal(info) is True + assert scheduled == [case.run.id] + stored = case.repository.get(case.run.id) + assert stored is not None + assert stored.outer_trace is not None + assert "private" not in json.dumps(stored.outer_trace) + assert len(stored.outer_trace["spans"]) == 1 + invocations, _ = case.repository.collected(case.run.id) + assert invocations == [] + + @staticmethod + def test_deleting_selected_configuration_keeps_historical_destination( + collection: _ApplicationTraceCollectionCollection, + ) -> None: + from machinery.context import RequestContext + from repositories.app.tracing_config_repository import SQLAlchemyAppTracingConfigRepository + from services.app_tracing_config_gateway import OpsTraceManagerGateway + from services.app_tracing_config_service import AppTracingConfigService + + original = collection.run + service = AppTracingConfigService( + configs=SQLAlchemyAppTracingConfigRepository(session_factory=collection.sessions), + provider=OpsTraceManagerGateway(), + ) + service.delete(RequestContext("request", None, "account", original.tenant_id), original.app_id, "langfuse") + stored = collection.repository.get(original.id) + assert stored is not None + assert stored.provider_config == original.provider_config + assert stored.provider_config["project_url"] == "https://langfuse.test/project/original" + assert stored.config_revision == original.config_revision + with collection.sessions() as session: + assert session.scalar(select(TraceAppConfig).where(TraceAppConfig.app_id == original.app_id)) is None + app = session.get(App, original.app_id) + assert app is not None + assert app.tracing is not None + assert json.loads(app.tracing)["enabled"] is False + assert json.loads(app.tracing)["tracing_provider"] is None + + +"""Exercise public orchestration, privacy, and delivery decisions against durable rows.""" + + +_ApplicationTraceServiceBoundariesCollection = _ApplicationTraceCollectionCollection +_application_trace_service_boundaries_batch_for = _application_trace_collection_batch_for +_application_trace_service_boundaries_claim = _application_trace_collection_claim +_application_trace_service_boundaries_identifier = _application_trace_collection_identifier +_application_trace_service_boundaries_ingest = _application_trace_collection_ingest +_application_trace_service_boundaries_seed = _application_trace_collection_seed +_application_trace_service_boundaries_terminal = _application_trace_collection_terminal + + +def source(case: _ApplicationTraceServiceBoundariesCollection) -> dict[str, str]: + return { + "tenant_id": case.run.tenant_id, + "app_id": case.run.app_id, + "source_kind": "message", + "source_id": case.run.source_id, + } + + +class TestApplicationTraceServiceBoundariesCases: + @staticmethod + def test_prepare_invocation_returns_only_opaque_context_and_respects_capture_gate( + sqlite_session_factory: sessionmaker[Session], monkeypatch: pytest.MonkeyPatch + ) -> None: + apply_config_overrides(monkeypatch, OPS_AGENT_TRACE_ENABLED=True) + case = _application_trace_service_boundaries_seed(sqlite_session_factory, "prepare", capture_content=True) + service = ApplicationTracingService(case.repository) + context = service.prepare_invocation( + tenant_id=case.run.tenant_id, + app_id=case.run.app_id, + agent_id=_application_trace_service_boundaries_identifier("agent"), + source_kind="message", + source_id=case.run.source_id, + session_id="ignored-new-session", + node_execution_id=None, + attempt=0, + call_key="call", + ) + assert context is not None + assert context.run_id == UUID(case.run.id) + assert context.capture_content is True + assert context.session_id == "session-original" + assert "encrypted-original" not in context.model_dump_json() + apply_config_overrides(monkeypatch, OPS_AGENT_TRACE_ENABLED=False) + assert ( + service.start( + tenant_id=case.run.tenant_id, + app_id=case.run.app_id, + source_kind="message", + source_id=case.run.source_id, + session_id="new", + ) + is None + ) + + @staticmethod + @pytest.mark.parametrize("missing", ["tenant", "app", "message"]) + def test_status_denies_foreign_or_absent_source_without_destination_details( + sqlite_session_factory: sessionmaker[Session], missing: str + ) -> None: + case = _application_trace_service_boundaries_seed(sqlite_session_factory, "scope") + identity = source(case) + identity[{"tenant": "tenant_id", "app": "app_id", "message": "source_id"}[missing]] = ( + _application_trace_service_boundaries_identifier("foreign") + ) + with pytest.raises(TraceCollectionError) as caught: + ApplicationTracingService(case.repository).status(**identity) + assert caught.value.code == "tracing_source_not_found" + assert caught.value.status == 404 + assert "encrypted-original" not in str(caught.value) + + @staticmethod + def test_status_uses_frozen_destination_and_selects_latest_node_attempt( + sqlite_session_factory: sessionmaker[Session], + monkeypatch: pytest.MonkeyPatch, + ) -> None: + case = _application_trace_service_boundaries_seed(sqlite_session_factory, "status") + with case.sessions.begin() as session: + config = session.scalar(select(TraceAppConfig).where(TraceAppConfig.app_id == case.run.app_id)) + assert config is not None + config.tracing_config = {"project_url": "https://new.example/project/new", "secret_key": "encrypted-new"} + + def forbidden(*_args: object, **_kwargs: object) -> None: + pytest.fail("status must not decrypt credentials or resolve current provider metadata") + + monkeypatch.setattr(service_module.OpsTraceManager, "decrypt_tracing_config", forbidden) + monkeypatch.setattr(service_module.OpsTraceManager, "get_trace_config_project_url", forbidden) + node_id = _application_trace_service_boundaries_identifier("selected-node") + invocation = case.repository.invocation( + run_id=case.run.id, + agent_id=_application_trace_service_boundaries_identifier("agent"), + node_execution_id=node_id, + attempt=1, + call_key="node:1", + ) + _application_trace_service_boundaries_ingest( + case, _application_trace_service_boundaries_batch_for(case.run, invocation) + ) + _application_trace_service_boundaries_terminal(case) + lease = _application_trace_service_boundaries_claim(case) + assert lease is not None + case.repository.complete( + case.run.id, + lease[1], + state="accepted", + result=TraceExportResult( + "provider-trace", + trace_url="https://frozen.example/trace", + span_urls={invocation.span_id: "https://frozen.example/span"}, + ), + ) + service = ApplicationTracingService(case.repository) + result = service.status(**source(case), node_execution_id=node_id) + assert result["status"] == "accepted" + assert result["trace_id"] == "provider-trace" + assert result["span_url"] == "https://frozen.example/span" + assert result["project_url"] == "https://langfuse.test/project/original" + assert result["config_revision"] == case.run.config_revision + assert result["queryable"] is None + assert result["capture_content"] is False + assert "provider_config" not in result + assert "encrypted-original" not in str(result) + with pytest.raises(TraceCollectionError, match="tracing_node_not_found"): + service.status( + **source(case), node_execution_id=_application_trace_service_boundaries_identifier("another-node") + ) + + @staticmethod + def test_new_run_freezes_updated_project_without_changing_previous_run( + sqlite_session_factory: sessionmaker[Session], + ) -> None: + case = _application_trace_service_boundaries_seed(sqlite_session_factory, "updated-project") + new_id = _application_trace_service_boundaries_identifier("new-message") + with case.sessions.begin() as session: + config = session.scalar(select(TraceAppConfig).where(TraceAppConfig.app_id == case.run.app_id)) + assert config is not None + config.tracing_config = {"project_url": "https://new.example/project/new", "secret_key": "encrypted-new"} + message = session.get(Message, case.run.source_id) + assert message is not None + session.add( + make_message( + message_id=new_id, + app_id=case.run.app_id, + conversation_id=message.conversation_id, + inputs={}, + query="new run", + message={}, + answer="done", + message_unit_price=message.message_unit_price, + answer_unit_price=message.answer_unit_price, + currency=message.currency, + from_source=message.from_source, + ) + ) + new_run = case.repository.start( + tenant_id=case.run.tenant_id, + app_id=case.run.app_id, + source_kind="message", + source_id=new_id, + session_id="new-session", + ) + assert new_run is not None + service = ApplicationTracingService(case.repository) + assert service.status(**source(case))["project_url"] == "https://langfuse.test/project/original" + assert ( + service.status(**{**source(case), "source_id": new_id})["project_url"] == "https://new.example/project/new" + ) + + @staticmethod + @pytest.mark.parametrize( + "project_url", + [ + None, + "", + 1, + "javascript:alert(1)", + "//example.com/project", + "https://user:secret@example.com/project", + "https://@example.com", + "https://example.com\\@other.test", + "https://example.com:invalid/project", + "https://[invalid/project", + "https://example.com/\nproject", + ], + ) + def test_status_hides_unavailable_or_unsafe_frozen_project_urls( + sqlite_session_factory: sessionmaker[Session], project_url: object + ) -> None: + case = _application_trace_service_boundaries_seed(sqlite_session_factory, "invalid-project") + with case.sessions.begin() as session: + run = session.get(ApplicationTraceRun, case.run.id) + assert run is not None + run.provider_config = { + "secret_key": "encrypted-original", + **({"project_url": project_url} if project_url is not None else {}), + } + result = ApplicationTracingService(case.repository).status(**source(case)) + assert result["project_url"] is None + assert "encrypted-original" not in str(result) + + @staticmethod + @pytest.mark.parametrize("provider", ["aliyun", "tencent"]) + def test_status_does_not_present_platform_console_as_project( + sqlite_session_factory: sessionmaker[Session], provider: str + ) -> None: + case = _application_trace_service_boundaries_seed(sqlite_session_factory, "platform-console") + with case.sessions.begin() as session: + run = session.get(ApplicationTraceRun, case.run.id) + assert run is not None + run.provider = provider + run.provider_config = {"project_url": "https://console.example.com"} + assert ApplicationTracingService(case.repository).status(**source(case))["project_url"] is None + + @staticmethod + @pytest.mark.parametrize("disabled", [True, False]) + def test_uncollected_status_and_missing_run_retry_do_not_construct_provider( + sqlite_session_factory: sessionmaker[Session], disabled: bool + ) -> None: + case = _application_trace_service_boundaries_seed(sqlite_session_factory, "missing", enabled=not disabled) + if not disabled: + with case.sessions.begin() as session: + stored = session.get(ApplicationTraceRun, case.run.id) + assert stored is not None + session.delete(stored) + service = ApplicationTracingService(case.repository) + assert service.status(**source(case)) == { + "status": "not_collected", + "source_kind": "message", + "source_id": case.run.source_id, + } + if not disabled: + with pytest.raises(TraceCollectionError, match="tracing_run_not_found"): + service.retry(**source(case), acknowledge_possible_duplicates=True) + + @staticmethod + def test_retry_acknowledgement_guards_schedule_and_keeps_frozen_snapshot( + sqlite_session_factory: sessionmaker[Session], monkeypatch: pytest.MonkeyPatch + ) -> None: + case = _application_trace_service_boundaries_seed(sqlite_session_factory, "retry") + _application_trace_service_boundaries_ingest( + case, _application_trace_service_boundaries_batch_for(case.run, case.invocation()) + ) + _application_trace_service_boundaries_terminal(case) + lease = _application_trace_service_boundaries_claim(case) + assert lease is not None + case.repository.complete(case.run.id, lease[1], state="uncertain", error_code="provider_acceptance_unknown") + scheduled: list[str] = [] + service = ApplicationTracingService(case.repository) + monkeypatch.setattr(service, "schedule", scheduled.append) + with pytest.raises(TraceCollectionError, match="tracing_duplicate_acknowledgement_required"): + service.retry(**source(case), acknowledge_possible_duplicates=False) + assert scheduled == [] + result = service.retry(**source(case), acknowledge_possible_duplicates=True) + assert scheduled == [case.run.id] + assert result["status"] == "ready" + assert result["config_revision"] == case.run.config_revision + current = case.repository.get(case.run.id) + assert current is not None + assert current.provider_config == case.run.provider_config + + @staticmethod + @pytest.mark.parametrize( + ("failure", "code"), + [ + (RetryableTraceDispatchError("private-transport-error"), "provider_retry_exhausted"), + (ValueError("private-credential-error"), "provider_export_rejected"), + ], + ) + def test_delivery_exhaustion_and_rejection_never_log_vendor_details( + sqlite_session_factory: sessionmaker[Session], + monkeypatch: pytest.MonkeyPatch, + caplog: pytest.LogCaptureFixture, + failure: Exception, + code: str, + ) -> None: + apply_config_overrides(monkeypatch, OPS_AGENT_TRACE_MAX_ATTEMPTS=1) + case = _application_trace_service_boundaries_seed(sqlite_session_factory, "delivery") + _application_trace_service_boundaries_ingest( + case, _application_trace_service_boundaries_batch_for(case.run, case.invocation()) + ) + _application_trace_service_boundaries_terminal(case) + + def export(*_args: object) -> TraceExportResult: + raise failure + + monkeypatch.setattr(service_module.OpsTraceManager, "decrypt_tracing_config", lambda *_args: {}) + monkeypatch.setattr(service_module, "export_trace", export) + service = ApplicationTracingService(case.repository) + service.deliver(case.run.id) + current = case.repository.get(case.run.id) + assert current is not None + assert current.state == "failed" + assert current.error_code == code + assert "private-" not in caplog.text + service.deliver(_application_trace_service_boundaries_identifier("missing-run")) + + @staticmethod + @pytest.mark.parametrize("capture_content", [False, True]) + @pytest.mark.parametrize( + "status", + [ + WorkflowNodeExecutionStatus.RUNNING, + WorkflowNodeExecutionStatus.SUCCEEDED, + WorkflowNodeExecutionStatus.FAILED, + ], + ) + def test_preview_terminal_uses_persisted_node_state_and_frozen_privacy( + sqlite_session_factory: sessionmaker[Session], status: WorkflowNodeExecutionStatus, capture_content: bool + ) -> None: + case = _application_trace_service_boundaries_seed( + sqlite_session_factory, "preview", capture_content=capture_content + ) + node_id = _application_trace_service_boundaries_identifier("preview-node") + workflow_id = _application_trace_service_boundaries_identifier("preview-definition") + with case.sessions.begin() as session: + session.add(make_workflow(workflow_id=workflow_id, tenant_id=case.run.tenant_id, app_id=case.run.app_id)) + session.add( + WorkflowNodeExecutionModel( + id=node_id, + node_execution_id=node_id, + tenant_id=case.run.tenant_id, + app_id=case.run.app_id, + workflow_id=workflow_id, + triggered_from=WorkflowNodeExecutionTriggeredFrom.SINGLE_STEP, + index=0, + node_id="agent", + node_type="agent", + title="Agent preview", + inputs='{"secret":"private-input"}', + outputs='{"answer":"private-output"}', + process_data="{}", + status=status, + elapsed_time=2, + error="private-model-error" if status == WorkflowNodeExecutionStatus.FAILED else None, + created_at=datetime(2026, 1, 1), + created_by=_application_trace_service_boundaries_identifier("account"), + created_by_role=CreatorUserRole.ACCOUNT, + ) + ) + run = case.repository.start( + tenant_id=case.run.tenant_id, + app_id=case.run.app_id, + source_kind="node", + source_id=node_id, + workflow_id=workflow_id, + session_id="preview-session", + ) + assert run is not None + ApplicationTracingService(case.repository).settle_preview(run) + current = case.repository.get(run.id) + assert current is not None + if status == WorkflowNodeExecutionStatus.RUNNING: + assert current.outer_trace is None + else: + assert current.outer_trace is not None + trace = service_module.CanonicalTrace.model_validate(current.outer_trace) + root = trace.spans[0] + assert root.status.value == ("error" if status == WorkflowNodeExecutionStatus.FAILED else "ok") + assert ("private-" in trace.model_dump_json()) is capture_content + assert root.metadata["node_execution_id"] == node_id + if not capture_content and status == WorkflowNodeExecutionStatus.FAILED: + assert root.error == "Node execution failed" + + @staticmethod + @pytest.mark.parametrize("owner_state", ["no_creator", "missing_account", "wrong_app"]) + def test_missing_owner_identity_cannot_authorize_trace_expansion( + sqlite_session_factory: sessionmaker[Session], owner_state: str + ) -> None: + from models.model import App + + case = _application_trace_service_boundaries_seed(sqlite_session_factory, "owner") + with case.sessions.begin() as session: + app = session.get(App, case.run.app_id) + assert app is not None + app.created_by = ( + None + if owner_state == "no_creator" + else _application_trace_service_boundaries_identifier("missing-owner") + ) + app_id = ( + _application_trace_service_boundaries_identifier("another-app") + if owner_state == "wrong_app" + else case.run.app_id + ) + with pytest.raises(TraceCollectionError, match="tracing_owner_not_found"): + ApplicationTracingService(case.repository)._service_account(case.run, app_id) + + @staticmethod + @pytest.mark.parametrize( + ("invalid", "code"), + [ + ("empty", "tracing_invocation_empty"), + ("foreign_parent", "tracing_invocation_parent_invalid"), + ("no_invocation_root", "tracing_invocation_root_missing"), + ("conflicting_span", "tracing_span_conflict"), + ], + ) + def test_invalid_runtime_tree_is_rejected_before_provider_credentials_or_io( + sqlite_session_factory: sessionmaker[Session], monkeypatch: pytest.MonkeyPatch, invalid: str, code: str + ) -> None: + case = _application_trace_service_boundaries_seed(sqlite_session_factory, "invalid-runtime") + batch = _application_trace_service_boundaries_batch_for(case.run, case.invocation()) + if invalid == "empty": + batch.spans = [] + elif invalid == "foreign_parent": + batch.spans[0] = batch.spans[0].model_copy(update={"parent_span_id": "f" * 16}) + elif invalid == "no_invocation_root": + batch.spans[0] = batch.spans[0].model_copy(update={"parent_span_id": batch.spans[0].span_id}) + else: + _application_trace_service_boundaries_ingest(case, batch.model_copy(update={"final": False})) + from uuid import uuid4 + + batch = batch.model_copy( + update={ + "batch_id": uuid4(), + "spans": [batch.spans[0].model_copy(update={"name": "conflicting-version"})], + } + ) + _application_trace_service_boundaries_ingest(case, batch) + _application_trace_service_boundaries_terminal(case) + + def forbidden(*_args: object, **_kwargs: object) -> None: + pytest.fail("invalid trees must fail before credentials or provider requests") + + monkeypatch.setattr(service_module.OpsTraceManager, "decrypt_tracing_config", forbidden) + monkeypatch.setattr(service_module, "export_trace", forbidden) + service = ApplicationTracingService(case.repository) + service.deliver(case.run.id) + current = case.repository.get(case.run.id) + assert current is not None + assert current.state == "failed" + assert current.error_code == code + + @staticmethod + def test_assemble_requires_a_persisted_business_terminal(sqlite_session_factory: sessionmaker[Session]) -> None: + case = _application_trace_service_boundaries_seed(sqlite_session_factory, "no-terminal") + with pytest.raises(TraceCollectionError, match="tracing_terminal_missing"): + ApplicationTracingService(case.repository).assemble(case.run) + + +"""Malformed durable traces must never be sent or offered for provider retry.""" + + +_invalid_trace_delivery_batch_for = _application_trace_collection_batch_for +_invalid_trace_delivery_ingest = _application_trace_collection_ingest +_invalid_trace_delivery_seed = _application_trace_collection_seed +_invalid_trace_delivery_terminal = _application_trace_collection_terminal + + +class TestInvalidTraceDeliveryCases: + @staticmethod + @pytest.mark.parametrize("invalid", ["cycle", "invalid_envelope"]) + def test_invalid_tree_is_permanent_before_decryption_or_provider_io( + sqlite_session_factory: sessionmaker[Session], monkeypatch: pytest.MonkeyPatch, invalid: str + ) -> None: + case = _invalid_trace_delivery_seed(sqlite_session_factory, str(uuid4())) + batch = _invalid_trace_delivery_batch_for(case.run, case.invocation()) + if invalid == "cycle": + for span_id, parent_id in [("b" * 16, "c" * 16), ("c" * 16, "b" * 16)]: + batch.spans.append( + AgentTraceSpan( + trace_id=case.run.trace_id, + span_id=span_id, + parent_span_id=parent_id, + name="cyclic tool", + start_time_ns=1_000_000_000, + end_time_ns=2_000_000_000, + status="ok", + ) + ) + _invalid_trace_delivery_ingest(case, batch) + _invalid_trace_delivery_terminal(case) + if invalid == "invalid_envelope": + with case.sessions.begin() as session: + row = session.get(ApplicationTraceRun, case.run.id) + assert row is not None + row.outer_trace = {"corrupt": True} + + def forbidden(*_args: object, **_kwargs: object) -> None: + pytest.fail("invalid traces must fail before credentials or provider I/O") + + monkeypatch.setattr(service_module.OpsTraceManager, "decrypt_tracing_config", forbidden) + monkeypatch.setattr(service_module, "export_trace", forbidden) + ApplicationTracingService(case.repository).deliver(case.run.id) + current = case.repository.get(case.run.id) + assert current is not None + assert current.state == "failed" + assert current.error_code == "tracing_payload_invalid" + assert current.attempts == 1 + with pytest.raises(TraceCollectionError, match="tracing_retry_not_allowed"): + case.repository.retry(run_id=case.run.id, acknowledge_possible_duplicates=True) diff --git a/api/tests/unit_tests/services/tracing/test_workflow_tool_tracing_boundary.py b/api/tests/unit_tests/services/tracing/test_workflow_tool_tracing_boundary.py deleted file mode 100644 index 4c0a881ba12..00000000000 --- a/api/tests/unit_tests/services/tracing/test_workflow_tool_tracing_boundary.py +++ /dev/null @@ -1,292 +0,0 @@ -"""Hidden Workflow Tool Agents must not invent a trace owned by the source app.""" - -import json -from copy import deepcopy -from dataclasses import replace -from datetime import UTC, datetime, timedelta - -import pytest -from sqlalchemy import select -from sqlalchemy.orm import Session, sessionmaker - -from core.ops.entities.trace_entity import WorkflowTraceInfo -from core.ops.unified_trace.trace_builder import CanonicalTraceBuilder -from core.repositories.sqlalchemy_workflow_node_execution_repository import SQLAlchemyWorkflowNodeExecutionRepository -from core.workflow.nodes.agent_v2.runtime_request_builder import WorkflowAgentRuntimeBuildContext -from graphon.enums import WorkflowExecutionStatus, WorkflowNodeExecutionStatus, WorkflowType -from models.application_trace import ApplicationTraceRun -from models.enums import CreatorUserRole -from models.model import App, AppMode, TraceAppConfig -from models.workflow import ( - WorkflowNodeExecutionModel, - WorkflowNodeExecutionTriggeredFrom, - WorkflowRun, - WorkflowRunTriggeredFrom, -) -from repositories.application_trace_repository import ApplicationTraceRepository -from services.tracing import runtime_bridge -from services.tracing.normalization import remap_outer_trace -from services.tracing.service import ApplicationTracingService -from tests.unit_tests.config_override import apply_config_overrides -from tests.unit_tests.model_factories import make_account, make_app, make_workflow -from tests.unit_tests.services.tracing.test_application_trace_collection import ( - Collection, - batch_for, - identifier, - ingest, -) -from tests.unit_tests.services.tracing.test_runtime_bridge import workflow_context, workflow_entity - - -def seed( - sessions: sessionmaker[Session], - monkeypatch: pytest.MonkeyPatch, - *, - root_has_agent: bool, - source_enabled: bool = True, -) -> tuple[ApplicationTracingService, WorkflowAgentRuntimeBuildContext, dict[str, str]]: - ids = { - name: identifier("tool-boundary-" + name) - for name in [ - "tenant", - "root", - "source", - "root_workflow", - "source_workflow", - "run", - "root_agent", - "root_tool", - "child", - "actor", - ] - } - root_graph: dict[str, list[dict[str, object]]] = { - "nodes": [{"id": "tool", "data": {"type": "tool", "title": "Root tool"}}], - "edges": [], - } - if root_has_agent: - root_graph["nodes"].append( - { - "id": "root-agent", - "data": {"type": "agent", "version": "2", "agent_node_kind": "dify_agent", "title": "Root Agent"}, - } - ) - now = datetime(2026, 10, 21, tzinfo=UTC) - with sessions.begin() as session: - for app, provider, enabled, content in [ - ("root", "langfuse", True, False), - ("source", "langsmith", source_enabled, True), - ]: - session.add( - make_app( - app_id=ids[app], - tenant_id=ids["tenant"], - mode=AppMode.WORKFLOW, - tracing=json.dumps({"enabled": enabled, "tracing_provider": provider, "capture_content": content}), - ) - ) - session.add( - make_workflow( - workflow_id=ids[app + "_workflow"], - app_id=ids[app], - tenant_id=ids["tenant"], - graph=root_graph if app == "root" else {"nodes": [], "edges": []}, - ) - ) - session.add( - TraceAppConfig( - app_id=ids[app], - tracing_provider=provider, - tracing_config={"project_url": "https://tracing.test/" + app}, - ) - ) - session.add( - WorkflowRun( - id=ids["run"], - tenant_id=ids["tenant"], - app_id=ids["root"], - workflow_id=ids["root_workflow"], - type=WorkflowType.WORKFLOW, - triggered_from=WorkflowRunTriggeredFrom.APP_RUN, - version="1", - graph=json.dumps(root_graph), - inputs="{}", - status=WorkflowExecutionStatus.SUCCEEDED, - created_by=ids["actor"], - created_by_role=CreatorUserRole.ACCOUNT, - created_at=now, - finished_at=now + timedelta(seconds=3), - ) - ) - for key, node_type in [("root_agent", "agent"), ("root_tool", "tool"), ("child", "agent")]: - child = key == "child" - app = "source" if child else "root" - session.add( - WorkflowNodeExecutionModel( - id=ids[key], - tenant_id=ids["tenant"], - app_id=ids[app], - workflow_id=ids[app + "_workflow"], - workflow_run_id=ids["run"], - triggered_from=WorkflowNodeExecutionTriggeredFrom.WORKFLOW_TOOL - if child - else WorkflowNodeExecutionTriggeredFrom.WORKFLOW_RUN, - triggered_from_workflow_id=ids["root_workflow"] if child else None, - triggered_from_node_execution_id=ids["root_tool"] if child else None, - node_execution_id=ids[key], - index=1, - node_id=key, - node_type=node_type, - title=key, - status=WorkflowNodeExecutionStatus.SUCCEEDED, - created_by=ids["actor"], - created_by_role=CreatorUserRole.ACCOUNT, - created_at=now, - finished_at=now + timedelta(seconds=2), - inputs="{}", - outputs="{}", - process_data=json.dumps( - {"workflow_tool_root_app_id": ids["root"], "workflow_tool_invocation_id": "tool-call"} - ) - if child - else "{}", - ) - ) - service = ApplicationTracingService(ApplicationTraceRepository(sessions)) - apply_config_overrides(monkeypatch, OPS_AGENT_TRACE_ENABLED=True) - monkeypatch.setattr(runtime_bridge, "_service", lambda: service) - entity = workflow_entity() - entity.app_config = entity.app_config.model_copy( - update={"app_id": ids["root"], "tenant_id": ids["tenant"], "workflow_id": ids["root_workflow"]} - ) - runtime_bridge.freeze_workflow_run( - entity, - workflow_id=ids["root_workflow"], - workflow_run_id=ids["run"], - conversation_id=None, - graph_data=root_graph, - ) - original = workflow_context(False) - context = replace( - original, - dify_context=original.dify_context.model_copy( - update={"tenant_id": ids["tenant"], "app_id": ids["source"], "workflow_tool_invocation_id": "tool-call"} - ), - workflow_id=ids["source_workflow"], - workflow_run_id=ids["run"], - node_execution_id=ids["child"], - ) - return service, context, ids - - -@pytest.mark.parametrize("root_has_agent", [False, True]) -@pytest.mark.parametrize("source_enabled", [False, True]) -def test_hidden_child_never_reserves_source_policy_or_changes_root_freeze( - sqlite_session_factory: sessionmaker[Session], - monkeypatch: pytest.MonkeyPatch, - root_has_agent: bool, - source_enabled: bool, -) -> None: - service, context, ids = seed( - sqlite_session_factory, monkeypatch, root_has_agent=root_has_agent, source_enabled=source_enabled - ) - with sqlite_session_factory() as session: - before = deepcopy([dict(row) for row in session.execute(select(ApplicationTraceRun.__table__)).mappings()]) - service_lookups: list[bool] = [] - - def lookup() -> ApplicationTracingService: - service_lookups.append(True) - return service - - monkeypatch.setattr(runtime_bridge, "_service", lookup) - # Marked children never consult policy or ownership, even for malformed - # empty markers or a foreign tenant. They cannot create an orphan B run. - for tenant in (ids["tenant"], identifier("foreign-tool-tenant")): - for marker in ("tool-call", ""): - child = replace( - context, - dify_context=context.dify_context.model_copy( - update={ - "tenant_id": tenant, - "workflow_tool_invocation_id": marker, - } - ), - ) - assert runtime_bridge.prepare_workflow_agent(child, None) is None - assert service_lookups == [] - with sqlite_session_factory() as session: - observed = [ - (r.app_id, r.provider, r.capture_content, r.source_id) for r in session.scalars(select(ApplicationTraceRun)) - ] - assert observed == ([(ids["root"], "langfuse", False, ids["run"])] if root_has_agent else []) - after = [dict(row) for row in session.execute(select(ApplicationTraceRun.__table__)).mappings()] - assert after == before - assert ( - service.repository.source_run( - tenant_id=ids["tenant"], app_id=ids["source"], source_kind="workflow", source_id=ids["run"] - ) - is None - ) - - -def test_hidden_child_skip_keeps_normal_root_agent_trace_assemblable( - sqlite_session_factory: sessionmaker[Session], - monkeypatch: pytest.MonkeyPatch, -) -> None: - service, child, ids = seed(sqlite_session_factory, monkeypatch, root_has_agent=True) - assert runtime_bridge.prepare_workflow_agent(child, None) is None - root_context = replace( - child, - dify_context=child.dify_context.model_copy(update={"app_id": ids["root"], "workflow_tool_invocation_id": None}), - workflow_id=ids["root_workflow"], - node_execution_id=ids["root_agent"], - ) - runtime = runtime_bridge.prepare_workflow_agent(root_context, None) - assert runtime is not None - run = service.repository.get(str(runtime.run_id)) - assert run is not None - assert run.app_id == ids["root"] - frozen = deepcopy(run.provider_config) - with sqlite_session_factory.begin() as session: - owner = session.get(App, ids["root"]) - assert owner is not None - owner.tracing = json.dumps({"enabled": True, "tracing_provider": "langsmith", "capture_content": True}) - workflow = session.get(WorkflowRun, ids["run"]) - assert workflow is not None - actor = make_account(account_id=ids["actor"]) - nodes = SQLAlchemyWorkflowNodeExecutionRepository( - session_factory=sqlite_session_factory, - tenant_id=ids["tenant"], - user=actor, - app_id=ids["root"], - triggered_from=WorkflowNodeExecutionTriggeredFrom.WORKFLOW_RUN, - ) - info = WorkflowTraceInfo( - workflow_data=workflow, - workflow_id=ids["root_workflow"], - tenant_id=ids["tenant"], - workflow_run_id=ids["run"], - workflow_run_elapsed_time=3, - workflow_run_status="succeeded", - workflow_run_inputs={}, - workflow_run_outputs={}, - workflow_run_version="1", - total_tokens=0, - file_list=[], - query="", - metadata={"app_id": ids["root"]}, - ) - outer = CanonicalTraceBuilder(lambda _info: nodes.get_by_workflow_execution(ids["run"])).build(info) - assert outer is not None - outer = remap_outer_trace(outer, trace_id=run.trace_id, capture_content=False, agent_message=False) - service.repository.terminal(run.id, outer.model_dump(mode="json"), outer.root_span_id) - invocation = service.repository.collected(run.id)[0][0] - ingest(Collection(service.repository, sqlite_session_factory, run), batch_for(run, invocation)) - current = service.repository.get(run.id) - assert current is not None - trace = service.assemble(current) - span_ids = {span.id for span in trace.spans} - assert all(span.parent_id is None or span.parent_id in span_ids for span in trace.spans) - assert current.provider_config == frozen - assert current.provider == "langfuse" - assert current.capture_content is False diff --git a/api/tests/unit_tests/services/test_schedule_service.py b/api/tests/unit_tests/services/trigger/test_schedule_service.py similarity index 99% rename from api/tests/unit_tests/services/test_schedule_service.py rename to api/tests/unit_tests/services/trigger/test_schedule_service.py index 18f68560579..a68123ed6e5 100644 --- a/api/tests/unit_tests/services/test_schedule_service.py +++ b/api/tests/unit_tests/services/trigger/test_schedule_service.py @@ -12,7 +12,7 @@ from core.workflow.nodes.trigger_schedule.exc import ScheduleConfigError from libs.schedule_utils import calculate_next_run_at, convert_12h_to_24h from models.workflow import Workflow from services.trigger.schedule_service import ScheduleService -from tests.unit_tests.model_factories import make_workflow +from tests.helpers.unit.model_factories import make_workflow class TestScheduleService(unittest.TestCase): diff --git a/api/tests/unit_tests/services/test_trigger_provider_service.py b/api/tests/unit_tests/services/trigger/test_trigger_provider_service.py similarity index 100% rename from api/tests/unit_tests/services/test_trigger_provider_service.py rename to api/tests/unit_tests/services/trigger/test_trigger_provider_service.py diff --git a/api/tests/unit_tests/services/test_trigger_subscription_builder_service.py b/api/tests/unit_tests/services/trigger/test_trigger_subscription_builder_service.py similarity index 100% rename from api/tests/unit_tests/services/test_trigger_subscription_builder_service.py rename to api/tests/unit_tests/services/trigger/test_trigger_subscription_builder_service.py diff --git a/api/tests/unit_tests/services/trigger/test_webhook_service.py b/api/tests/unit_tests/services/trigger/test_webhook_service.py new file mode 100644 index 00000000000..9aa4af9dcad --- /dev/null +++ b/api/tests/unit_tests/services/trigger/test_webhook_service.py @@ -0,0 +1,1335 @@ +import logging +from io import BytesIO +from types import SimpleNamespace +from typing import Any, cast +from unittest.mock import MagicMock, patch + +import pytest +from flask import Flask +from sqlalchemy import Engine as _webhook_service_relationships_Engine +from sqlalchemy import Table, select +from sqlalchemy.engine import Engine as _webhook_service_Engine +from sqlalchemy.orm import Session, sessionmaker +from werkzeug.datastructures import FileStorage +from werkzeug.exceptions import RequestEntityTooLarge + +from core.trigger.constants import TRIGGER_WEBHOOK_NODE_TYPE +from core.workflow.nodes.trigger_webhook.entities import ( + ContentType, + WebhookBodyParameter, + WebhookData, + WebhookParameter, +) +from graphon.variables.types import SegmentType +from models.base import TypeBase +from models.enums import AppTriggerStatus, AppTriggerType, EndUserType +from models.model import App, AppMode, EndUser +from models.tools import ToolFile +from models.trigger import AppTrigger, WorkflowWebhookTrigger +from models.workflow import Workflow +from services.errors.app import QuotaExceededError +from services.trigger import webhook_service +from services.trigger import webhook_service as service_module +from services.trigger import webhook_service as webhook_service_module +from services.trigger.webhook_service import WebhookService +from tests.helpers.unit.model_factories import make_app, make_end_user, make_workflow + + +class _EndUserServiceStub: + def __init__(self, result: EndUser | Exception) -> None: + self._result = result + self.calls: list[tuple[EndUserType, str, str, str | None]] = [] + + def get_or_create_end_user_by_type( + self, + type: EndUserType, + tenant_id: str, + app_id: str, + user_id: str | None = None, + ) -> EndUser: + self.calls.append((type, tenant_id, app_id, user_id)) + if isinstance(self._result, Exception): + raise self._result + return self._result + + +def _webhook_trigger( + *, + webhook_id: str = "webhook-123", + tenant_id: str = "tenant-123", + app_id: str = "app-123", + node_id: str = "node-123", + created_by: str = "account-123", +) -> WorkflowWebhookTrigger: + return WorkflowWebhookTrigger( + webhook_id=webhook_id, + tenant_id=tenant_id, + app_id=app_id, + node_id=node_id, + created_by=created_by, + ) + + +def _tool_file() -> ToolFile: + tool_file = ToolFile( + user_id="user-123", + tenant_id="test_tenant", + conversation_id=None, + file_key="webhook/test.txt", + mimetype="text/plain", + name="test.txt", + size=7, + ) + tool_file.id = "test_file_id" + return tool_file + + +def _workflow( + *, + workflow_id: str = "workflow-123", + tenant_id: str = "tenant-123", + app_id: str = "app-123", + version: str = Workflow.VERSION_DRAFT, +) -> Workflow: + return make_workflow( + workflow_id=workflow_id, + tenant_id=tenant_id, + app_id=app_id, + version=version, + created_by="account-123", + ) + + +def _app(*, tenant_id: str = "tenant-123", app_id: str = "app-123") -> App: + return make_app( + app_id=app_id, + tenant_id=tenant_id, + name="Webhook App", + mode=AppMode.WORKFLOW, + icon_type=None, + max_active_requests=0, + ) + + +def _app_trigger(*, status: AppTriggerStatus = AppTriggerStatus.ENABLED) -> AppTrigger: + return AppTrigger( + tenant_id="tenant-123", + app_id="app-123", + node_id="node-123", + trigger_type=AppTriggerType.TRIGGER_WEBHOOK, + title="Webhook", + status=status, + ) + + +def _end_user() -> EndUser: + return make_end_user( + end_user_id="end-user-123", + tenant_id="tenant-123", + app_id="app-123", + end_user_type=EndUserType.TRIGGER, + session_id="webhook-session", + ) + + +class TestWebhookServiceCases: + class TestWebhookServiceLookup: + def test_debug_lookup_scopes_draft_workflow_to_trigger_owner( + self, + sqlite_session: Session, + sqlite_engine: _webhook_service_Engine, + monkeypatch: pytest.MonkeyPatch, + ) -> None: + target_workflow = _workflow() + tenant_decoy = _workflow( + workflow_id="workflow-decoy", + tenant_id="tenant-other", + app_id="app-123", + ) + sqlite_session.add_all([_webhook_trigger(), target_workflow, tenant_decoy]) + sqlite_session.commit() + monkeypatch.setattr(webhook_service_module, "db", SimpleNamespace(engine=sqlite_engine)) + node_config = {"id": "node-123", "data": {}} + + with patch.object(Workflow, "get_node_config_by_id", autospec=True, return_value=node_config) as get_node: + webhook_trigger, workflow, result_node_config = WebhookService.get_webhook_trigger_and_workflow( + "webhook-123", is_debug=True + ) + + assert webhook_trigger.tenant_id == "tenant-123" + assert workflow.id == target_workflow.id + assert result_node_config is node_config + get_node.assert_called_once_with(workflow, "node-123") + + def test_published_lookup_uses_persisted_owner_chain( + self, + sqlite_session: Session, + sqlite_engine: _webhook_service_Engine, + monkeypatch: pytest.MonkeyPatch, + ) -> None: + sqlite_session.add_all([_webhook_trigger(), _app_trigger(), _app()]) + sqlite_session.commit() + monkeypatch.setattr(webhook_service_module, "db", SimpleNamespace(engine=sqlite_engine)) + published_workflow = _workflow(version="published") + node_config = {"id": "node-123", "data": {}} + + with ( + patch.object(webhook_service_module, "WorkflowService") as workflow_service_class, + patch.object(Workflow, "get_node_config_by_id", autospec=True, return_value=node_config), + ): + get_published_workflow = workflow_service_class.return_value.get_published_workflow + get_published_workflow.return_value = published_workflow + webhook_trigger, workflow, result_node_config = WebhookService.get_webhook_trigger_and_workflow( + "webhook-123" + ) + + assert webhook_trigger.app_id == "app-123" + assert workflow is published_workflow + assert result_node_config is node_config + persisted_app = get_published_workflow.call_args.args[0] + assert isinstance(persisted_app, App) + assert persisted_app.id == "app-123" + assert persisted_app.tenant_id == "tenant-123" + + def test_published_lookup_rejects_cross_tenant_app_decoy( + self, + sqlite_session: Session, + sqlite_engine: _webhook_service_Engine, + monkeypatch: pytest.MonkeyPatch, + ) -> None: + sqlite_session.add_all( + [ + _webhook_trigger(), + _app_trigger(), + _app(tenant_id="tenant-other"), + ] + ) + sqlite_session.commit() + monkeypatch.setattr(webhook_service_module, "db", SimpleNamespace(engine=sqlite_engine)) + + with pytest.raises(ValueError, match="App not found"): + WebhookService.get_webhook_trigger_and_workflow("webhook-123") + + @pytest.mark.parametrize( + "status", + [AppTriggerStatus.DISABLED, AppTriggerStatus.UNAUTHORIZED], + ) + def test_published_lookup_rejects_inactive_trigger( + self, + status: AppTriggerStatus, + sqlite_session: Session, + sqlite_engine: _webhook_service_Engine, + monkeypatch: pytest.MonkeyPatch, + ) -> None: + sqlite_session.add_all([_webhook_trigger(), _app_trigger(status=status)]) + sqlite_session.commit() + monkeypatch.setattr(webhook_service_module, "db", SimpleNamespace(engine=sqlite_engine)) + + with pytest.raises(ValueError, match="disabled"): + WebhookService.get_webhook_trigger_and_workflow("webhook-123") + + def test_published_lookup_reports_rate_limited_trigger( + self, + sqlite_session: Session, + sqlite_engine: _webhook_service_Engine, + monkeypatch: pytest.MonkeyPatch, + ) -> None: + sqlite_session.add_all([_webhook_trigger(), _app_trigger(status=AppTriggerStatus.RATE_LIMITED)]) + sqlite_session.commit() + monkeypatch.setattr(webhook_service_module, "db", SimpleNamespace(engine=sqlite_engine)) + + with pytest.raises(QuotaExceededError): + WebhookService.get_webhook_trigger_and_workflow("webhook-123") + + class TestWebhookServiceUnit: + """Webhook business-logic tests with isolated sessions where the service owns their lifecycle.""" + + def test_trigger_workflow_execution_propagates_quota_error_without_error_log( + self, + caplog: pytest.LogCaptureFixture, + monkeypatch: pytest.MonkeyPatch, + sqlite_engine: _webhook_service_Engine, + ) -> None: + """Quota failures refund the charge and close the real service-owned session.""" + monkeypatch.setattr(webhook_service_module, "db", SimpleNamespace(engine=sqlite_engine)) + webhook_trigger = _webhook_trigger() + workflow = _workflow() + quota_charge = MagicMock() + quota_error = QuotaExceededError(feature="workflow", tenant_id="tenant-123", required=1) + + caplog.set_level(logging.INFO) + with ( + patch("services.trigger.webhook_service.QuotaService.reserve", return_value=quota_charge), + patch( + "services.trigger.webhook_service.AsyncWorkflowService.trigger_workflow_async", + side_effect=quota_error, + ) as mock_trigger_workflow_async, + ): + with pytest.raises(QuotaExceededError) as exc_info: + WebhookService.trigger_workflow_execution( + webhook_trigger, + {"body": {}, "headers": {}, "query_params": {}, "files": {}, "method": "POST"}, + workflow, + end_users=_EndUserServiceStub(_end_user()), + ) + + assert exc_info.value is quota_error + quota_charge.refund.assert_called_once_with() + session = mock_trigger_workflow_async.call_args.kwargs["session"] + assert isinstance(session, Session) + assert session.in_transaction() is False + assert len(caplog.records) == 1 + assert caplog.records[0].levelno == logging.INFO + assert caplog.records[0].message == ( + "Tenant tenant-123 quota exceeded for feature workflow, skipping webhook trigger webhook-123" + ) + + def test_trigger_workflow_execution_success( + self, + monkeypatch: pytest.MonkeyPatch, + sqlite_engine: _webhook_service_Engine, + ) -> None: + monkeypatch.setattr(webhook_service_module, "db", SimpleNamespace(engine=sqlite_engine)) + webhook_trigger = _webhook_trigger() + workflow = _workflow() + end_user = _end_user() + quota_charge = MagicMock() + webhook_data = { + "method": "POST", + "headers": {"Authorization": "Bearer token"}, + "query_params": {"version": "1"}, + "body": {"message": "hello"}, + "files": {}, + } + + with ( + patch.object(webhook_service_module.QuotaService, "reserve", return_value=quota_charge), + patch.object( + webhook_service_module.AsyncWorkflowService, + "trigger_workflow_async", + ) as mock_trigger, + ): + WebhookService.trigger_workflow_execution( + webhook_trigger, + webhook_data, + workflow, + end_users=_EndUserServiceStub(end_user), + ) + + call_session = mock_trigger.call_args.kwargs["session"] + assert call_session.get_bind() is sqlite_engine + quota_charge.commit.assert_called_once_with() + quota_charge.refund.assert_not_called() + + def test_trigger_workflow_execution_end_user_service_failure(self) -> None: + webhook_trigger = _webhook_trigger() + workflow = _workflow() + webhook_data = {"method": "POST", "headers": {}, "query_params": {}, "body": {}, "files": {}} + + with pytest.raises(ValueError, match="Failed to create end user"): + WebhookService.trigger_workflow_execution( + webhook_trigger, + webhook_data, + workflow, + end_users=_EndUserServiceStub(ValueError("Failed to create end user")), + ) + + def test_extract_webhook_data_json(self): + """Test webhook data extraction from JSON request.""" + app = Flask(__name__) + + with app.test_request_context( + "/webhook", + method="POST", + headers={"Content-Type": "application/json", "Authorization": "Bearer token"}, + query_string="version=1&format=json", + json={"message": "hello", "count": 42}, + ): + webhook_trigger = _webhook_trigger() + webhook_data = WebhookService.extract_webhook_data(webhook_trigger) + + assert webhook_data["method"] == "POST" + assert webhook_data["headers"]["Authorization"] == "Bearer token" + # Query params are now extracted as raw strings + assert webhook_data["query_params"]["version"] == "1" + assert webhook_data["query_params"]["format"] == "json" + assert webhook_data["body"]["message"] == "hello" + assert webhook_data["body"]["count"] == 42 + assert webhook_data["files"] == {} + + def test_extract_webhook_data_query_params_remain_strings(self): + """Query parameters should be extracted as raw strings without automatic conversion.""" + app = Flask(__name__) + + with app.test_request_context( + "/webhook", + method="GET", + headers={"Content-Type": "application/json"}, + query_string="count=42&threshold=3.14&enabled=true¬e=text", + ): + webhook_trigger = _webhook_trigger() + webhook_data = WebhookService.extract_webhook_data(webhook_trigger) + + # After refactoring, raw extraction keeps query params as strings + assert webhook_data["query_params"]["count"] == "42" + assert webhook_data["query_params"]["threshold"] == "3.14" + assert webhook_data["query_params"]["enabled"] == "true" + assert webhook_data["query_params"]["note"] == "text" + + def test_extract_webhook_data_form_urlencoded(self): + """Test webhook data extraction from form URL encoded request.""" + app = Flask(__name__) + + with app.test_request_context( + "/webhook", + method="POST", + headers={"Content-Type": "application/x-www-form-urlencoded"}, + data={"username": "test", "password": "secret"}, + ): + webhook_trigger = _webhook_trigger() + webhook_data = WebhookService.extract_webhook_data(webhook_trigger) + + assert webhook_data["method"] == "POST" + assert webhook_data["body"]["username"] == "test" + assert webhook_data["body"]["password"] == "secret" + + def test_extract_webhook_data_multipart_with_files(self): + """Test webhook data extraction from multipart form with files.""" + app = Flask(__name__) + + # Create a mock file + file_content = b"test file content" + file_storage = FileStorage(stream=BytesIO(file_content), filename="test.txt", content_type="text/plain") + + with app.test_request_context( + "/webhook", + method="POST", + headers={"Content-Type": "multipart/form-data"}, + data={"message": "test", "file": file_storage}, + ): + webhook_trigger = _webhook_trigger() + webhook_trigger.tenant_id = "test_tenant" + + with patch.object(WebhookService, "_process_file_uploads", autospec=True) as mock_process_files: + mock_process_files.return_value = {"file": "mocked_file_obj"} + + webhook_data = WebhookService.extract_webhook_data(webhook_trigger) + + assert webhook_data["method"] == "POST" + assert webhook_data["body"]["message"] == "test" + assert webhook_data["files"]["file"] == "mocked_file_obj" + mock_process_files.assert_called_once() + + def test_extract_webhook_data_raw_text(self): + """Test webhook data extraction from raw text request.""" + app = Flask(__name__) + + with app.test_request_context( + "/webhook", method="POST", headers={"Content-Type": "text/plain"}, data="raw text content" + ): + webhook_trigger = _webhook_trigger() + webhook_data = WebhookService.extract_webhook_data(webhook_trigger) + + assert webhook_data["method"] == "POST" + assert webhook_data["body"]["raw"] == "raw text content" + + def test_extract_octet_stream_body_uses_detected_mime(self): + """Octet-stream uploads should rely on detected MIME type.""" + app = Flask(__name__) + binary_content = b"plain text data" + + with app.test_request_context( + "/webhook", method="POST", headers={"Content-Type": "application/octet-stream"}, data=binary_content + ): + webhook_trigger = _webhook_trigger() + mock_file = MagicMock() + mock_file.to_dict.return_value = {"file": "data"} + + with ( + patch.object( + WebhookService, "_detect_binary_mimetype", return_value="text/plain", autospec=True + ) as mock_detect, + patch.object(WebhookService, "_create_file_from_binary", autospec=True) as mock_create, + ): + mock_create.return_value = mock_file + body, files = WebhookService._extract_octet_stream_body(webhook_trigger) + + assert body["raw"] == {"file": "data"} + assert files == {} + mock_detect.assert_called_once_with(binary_content) + mock_create.assert_called_once() + args = mock_create.call_args[0] + assert args[0] == binary_content + assert args[1] == "text/plain" + assert args[2] is webhook_trigger + + def test_detect_binary_mimetype_uses_magic(self, monkeypatch: pytest.MonkeyPatch): + """python-magic output should be used when available.""" + fake_magic = MagicMock() + fake_magic.from_buffer.return_value = "image/png" + monkeypatch.setattr("services.trigger.webhook_service.magic", fake_magic) + + result = WebhookService._detect_binary_mimetype(b"binary data") + + assert result == "image/png" + fake_magic.from_buffer.assert_called_once() + + def test_detect_binary_mimetype_fallback_without_magic(self, monkeypatch: pytest.MonkeyPatch): + """Fallback MIME type should be used when python-magic is unavailable.""" + monkeypatch.setattr("services.trigger.webhook_service.magic", None) + + result = WebhookService._detect_binary_mimetype(b"binary data") + + assert result == "application/octet-stream" + + def test_detect_binary_mimetype_handles_magic_exception( + self, monkeypatch: pytest.MonkeyPatch, caplog: pytest.LogCaptureFixture + ): + """Fallback MIME type should be used when python-magic raises an exception.""" + try: + import magic as real_magic + except ImportError: + pytest.skip("python-magic is not installed") + + fake_magic = MagicMock() + fake_magic.from_buffer.side_effect = real_magic.MagicException("magic error") + monkeypatch.setattr("services.trigger.webhook_service.magic", fake_magic) + caplog.set_level(logging.DEBUG, logger="services.trigger.webhook_service") + + result = WebhookService._detect_binary_mimetype(b"binary data") + + assert result == "application/octet-stream" + assert "python-magic detection failed for octet-stream payload" in caplog.messages + + def test_extract_webhook_data_invalid_json(self): + """Test webhook data extraction with invalid JSON.""" + app = Flask(__name__) + + with app.test_request_context( + "/webhook", method="POST", headers={"Content-Type": "application/json"}, data="invalid json" + ): + webhook_trigger = _webhook_trigger() + with pytest.raises(ValueError, match="Invalid JSON body"): + WebhookService.extract_webhook_data(webhook_trigger) + + def test_generate_webhook_response_default(self): + """Test webhook response generation with default values.""" + node_config = {"data": {}} + + response_data, status_code = WebhookService.generate_webhook_response(node_config) + + assert status_code == 200 + assert response_data["status"] == "success" + assert "Webhook processed successfully" in response_data["message"] + + def test_generate_webhook_response_custom_json(self): + """Test webhook response generation with custom JSON response.""" + node_config = {"data": {"status_code": 201, "response_body": '{"result": "created", "id": 123}'}} + + response_data, status_code = WebhookService.generate_webhook_response(node_config) + + assert status_code == 201 + assert response_data["result"] == "created" + assert response_data["id"] == 123 + + def test_generate_webhook_response_custom_text(self): + """Test webhook response generation with custom text response.""" + node_config = {"data": {"status_code": 202, "response_body": "Request accepted for processing"}} + + response_data, status_code = WebhookService.generate_webhook_response(node_config) + + assert status_code == 202 + assert response_data["message"] == "Request accepted for processing" + + def test_generate_webhook_response_invalid_json(self): + """Test webhook response generation with invalid JSON response.""" + node_config = {"data": {"status_code": 400, "response_body": '{"invalid": json}'}} + + response_data, status_code = WebhookService.generate_webhook_response(node_config) + + assert status_code == 400 + assert response_data["message"] == '{"invalid": json}' + + def test_generate_webhook_response_empty_response_body(self): + """Test webhook response generation with empty response body.""" + node_config = {"data": {"status_code": 204, "response_body": ""}} + + response_data, status_code = WebhookService.generate_webhook_response(node_config) + + assert status_code == 204 + assert response_data["status"] == "success" + assert "Webhook processed successfully" in response_data["message"] + + def test_generate_webhook_response_array_json(self): + """Test webhook response generation with JSON array response.""" + node_config = {"data": {"status_code": 200, "response_body": '[{"id": 1}, {"id": 2}]'}} + + response_data, status_code = WebhookService.generate_webhook_response(node_config) + + assert status_code == 200 + assert isinstance(response_data, list) + assert len(response_data) == 2 + assert response_data[0]["id"] == 1 + assert response_data[1]["id"] == 2 + + @patch("services.trigger.webhook_service.ToolFileManager", autospec=True) + @patch("services.trigger.webhook_service.file_factory", autospec=True) + def test_process_file_uploads_success(self, mock_file_factory, mock_tool_file_manager): + """Test successful file upload processing.""" + # Mock ToolFileManager + mock_tool_file_instance = mock_tool_file_manager.return_value # Mock file creation + mock_tool_file_instance.create_file_by_raw.return_value = _tool_file() + + # Mock file factory + mock_file_obj = MagicMock() + mock_file_factory.build_from_mapping.return_value = mock_file_obj + + # Create mock files + files = { + "file1": MagicMock(filename="test1.txt", content_type="text/plain"), + "file2": MagicMock(filename="test2.jpg", content_type="image/jpeg"), + } + + # Mock file reads + files["file1"].stream.read.return_value = b"content1" + files["file2"].stream.read.return_value = b"content2" + + webhook_trigger = _webhook_trigger() + webhook_trigger.tenant_id = "test_tenant" + + result = WebhookService._process_file_uploads(files, webhook_trigger) + + assert len(result) == 2 + assert "file1" in result + assert "file2" in result + + # Verify file processing was called for each file + assert mock_tool_file_manager.call_count == 2 + assert mock_file_factory.build_from_mapping.call_count == 2 + + @patch("services.trigger.webhook_service.ToolFileManager", autospec=True) + @patch("services.trigger.webhook_service.file_factory", autospec=True) + def test_process_file_uploads_with_errors(self, mock_file_factory, mock_tool_file_manager): + """Test file upload processing with errors.""" + # Mock ToolFileManager + mock_tool_file_instance = mock_tool_file_manager.return_value # Mock file creation + mock_tool_file_instance.create_file_by_raw.return_value = _tool_file() + + # Mock file factory + mock_file_obj = MagicMock() + mock_file_factory.build_from_mapping.return_value = mock_file_obj + + # Create mock files, one will fail + files = { + "good_file": MagicMock(filename="test.txt", content_type="text/plain"), + "bad_file": MagicMock(filename="test.bad", content_type="text/plain"), + } + + files["good_file"].stream.read.return_value = b"content" + files["bad_file"].stream.read.side_effect = Exception("Read error") + + webhook_trigger = _webhook_trigger() + webhook_trigger.tenant_id = "test_tenant" + + result = WebhookService._process_file_uploads(files, webhook_trigger) + + # Should process the good file and skip the bad one + assert len(result) == 1 + assert "good_file" in result + assert "bad_file" not in result + + def test_process_file_uploads_empty_filename(self): + """Test file upload processing with empty filename.""" + files = { + "no_filename": MagicMock(filename="", content_type="text/plain"), + "none_filename": MagicMock(filename=None, content_type="text/plain"), + } + + webhook_trigger = _webhook_trigger() + webhook_trigger.tenant_id = "test_tenant" + + result = WebhookService._process_file_uploads(files, webhook_trigger) + + # Should skip files without filenames + assert len(result) == 0 + + def test_validate_json_value_string(self): + """Test JSON value validation for string type.""" + # Valid string + result = WebhookService._validate_json_value("name", "hello", "string") + assert result == "hello" + + # Invalid string (number) - should raise ValueError + with pytest.raises(ValueError, match="Expected string, got int"): + WebhookService._validate_json_value("name", 123, "string") + + def test_validate_json_value_number(self): + """Test JSON value validation for number type.""" + # Valid integer + result = WebhookService._validate_json_value("count", 42, "number") + assert result == 42 + + # Valid float + result = WebhookService._validate_json_value("price", 19.99, "number") + assert result == 19.99 + + # Invalid number (string) - should raise ValueError + with pytest.raises(ValueError, match="Expected number, got str"): + WebhookService._validate_json_value("count", "42", "number") + + def test_validate_json_value_bool(self): + """Test JSON value validation for boolean type.""" + # Valid boolean + result = WebhookService._validate_json_value("enabled", True, "boolean") + assert result is True + + result = WebhookService._validate_json_value("enabled", False, "boolean") + assert result is False + + # Invalid boolean (string) - should raise ValueError + with pytest.raises(ValueError, match="Expected boolean, got str"): + WebhookService._validate_json_value("enabled", "true", "boolean") + + def test_validate_json_value_object(self): + """Test JSON value validation for object type.""" + # Valid object + result = WebhookService._validate_json_value("user", {"name": "John", "age": 30}, "object") + assert result == {"name": "John", "age": 30} + + # Invalid object (string) - should raise ValueError + with pytest.raises(ValueError, match="Expected object, got str"): + WebhookService._validate_json_value("user", "not_an_object", "object") + + def test_validate_json_value_array_string(self): + """Test JSON value validation for array[string] type.""" + # Valid array of strings + result = WebhookService._validate_json_value("tags", ["tag1", "tag2", "tag3"], "array[string]") + assert result == ["tag1", "tag2", "tag3"] + + # Invalid - not an array + with pytest.raises(ValueError, match="Expected array of strings, got str"): + WebhookService._validate_json_value("tags", "not_an_array", "array[string]") + + # Invalid - array with non-strings + with pytest.raises(ValueError, match="Expected array of strings, got list"): + WebhookService._validate_json_value("tags", ["tag1", 123, "tag3"], "array[string]") + + def test_validate_json_value_array_number(self): + """Test JSON value validation for array[number] type.""" + # Valid array of numbers + result = WebhookService._validate_json_value("scores", [1, 2.5, 3, 4.7], "array[number]") + assert result == [1, 2.5, 3, 4.7] + + # Invalid - array with non-numbers + with pytest.raises(ValueError, match="Expected array of numbers, got list"): + WebhookService._validate_json_value("scores", [1, "2", 3], "array[number]") + + def test_validate_json_value_array_bool(self): + """Test JSON value validation for array[boolean] type.""" + # Valid array of booleans + result = WebhookService._validate_json_value("flags", [True, False, True], "array[boolean]") + assert result == [True, False, True] + + # Invalid - array with non-booleans + with pytest.raises(ValueError, match="Expected array of booleans, got list"): + WebhookService._validate_json_value("flags", [True, "false", True], "array[boolean]") + + def test_validate_json_value_array_object(self): + """Test JSON value validation for array[object] type.""" + # Valid array of objects + result = WebhookService._validate_json_value("users", [{"name": "John"}, {"name": "Jane"}], "array[object]") + assert result == [{"name": "John"}, {"name": "Jane"}] + + # Invalid - array with non-objects + with pytest.raises(ValueError, match="Expected array of objects, got list"): + WebhookService._validate_json_value("users", [{"name": "John"}, "not_object"], "array[object]") + + def test_convert_form_value_string(self): + """Test form value conversion for string type.""" + result = WebhookService._convert_form_value("test", "hello", "string") + assert result == "hello" + + def test_convert_form_value_number(self): + """Test form value conversion for number type.""" + # Integer + result = WebhookService._convert_form_value("count", "42", "number") + assert result == 42 + + # Float + result = WebhookService._convert_form_value("price", "19.99", "number") + assert result == 19.99 + + # Invalid number + with pytest.raises(ValueError, match="Cannot convert 'not_a_number' to number"): + WebhookService._convert_form_value("count", "not_a_number", "number") + + def test_convert_form_value_boolean(self): + """Test form value conversion for boolean type.""" + # True values + assert WebhookService._convert_form_value("flag", "true", "boolean") is True + assert WebhookService._convert_form_value("flag", "1", "boolean") is True + assert WebhookService._convert_form_value("flag", "yes", "boolean") is True + + # False values + assert WebhookService._convert_form_value("flag", "false", "boolean") is False + assert WebhookService._convert_form_value("flag", "0", "boolean") is False + assert WebhookService._convert_form_value("flag", "no", "boolean") is False + + # Invalid boolean + with pytest.raises(ValueError, match="Cannot convert 'maybe' to boolean"): + WebhookService._convert_form_value("flag", "maybe", "boolean") + + def test_extract_and_validate_webhook_data_success(self): + """Test successful unified data extraction and validation.""" + app = Flask(__name__) + + with app.test_request_context( + "/webhook", + method="POST", + headers={"Content-Type": "application/json"}, + query_string="count=42&enabled=true", + json={"message": "hello", "age": 25}, + ): + webhook_trigger = _webhook_trigger() + node_config = { + "data": { + "method": "post", + "content_type": "application/json", + "params": [ + {"name": "count", "type": "number", "required": True}, + {"name": "enabled", "type": "boolean", "required": True}, + ], + "body": [ + {"name": "message", "type": "string", "required": True}, + {"name": "age", "type": "number", "required": True}, + ], + } + } + + result = WebhookService.extract_and_validate_webhook_data(webhook_trigger, node_config) + + # Check that types are correctly converted + assert result["query_params"]["count"] == 42 # Converted to int + assert result["query_params"]["enabled"] is True # Converted to bool + assert result["body"]["message"] == "hello" # Already string + assert result["body"]["age"] == 25 # Already number + + def test_extract_and_validate_webhook_data_invalid_json_error(self): + """Invalid JSON should bubble up as a ValueError with details.""" + app = Flask(__name__) + + with app.test_request_context( + "/webhook", + method="POST", + headers={"Content-Type": "application/json"}, + data='{"invalid": }', + ): + webhook_trigger = _webhook_trigger() + node_config = { + "data": { + "method": "post", + "content_type": "application/json", + } + } + + with pytest.raises(ValueError, match="Invalid JSON body"): + WebhookService.extract_and_validate_webhook_data(webhook_trigger, node_config) + + def test_extract_and_validate_webhook_data_validation_error(self): + """Test unified data extraction with validation error.""" + app = Flask(__name__) + + with app.test_request_context( + "/webhook", + method="GET", # Wrong method + headers={"Content-Type": "application/json"}, + ): + webhook_trigger = _webhook_trigger() + node_config = { + "data": { + "method": "post", # Expects POST + "content_type": "application/json", + } + } + + with pytest.raises(ValueError, match="HTTP method mismatch"): + WebhookService.extract_and_validate_webhook_data(webhook_trigger, node_config) + + def test_extract_and_validate_webhook_request_missing_required_header(self) -> None: + app = Flask(__name__) + with app.test_request_context( + "/webhook", + method="POST", + headers={"Content-Type": "application/json"}, + ): + node_config = { + "data": { + "method": "post", + "content_type": "application/json", + "headers": [{"name": "Authorization", "required": True}], + } + } + + with pytest.raises(ValueError, match="Required header missing: Authorization"): + WebhookService.extract_and_validate_webhook_data(_webhook_trigger(), node_config) + + def test_extract_and_validate_webhook_request_case_insensitive_headers(self) -> None: + app = Flask(__name__) + with app.test_request_context( + "/webhook", + method="POST", + headers={"Content-Type": "application/json", "authorization": "Bearer token"}, + json={"message": "hello"}, + ): + node_config = { + "data": { + "method": "post", + "content_type": "application/json", + "headers": [{"name": "Authorization", "required": True}], + "body": [{"name": "message", "type": "string", "required": True}], + } + } + + result = WebhookService.extract_and_validate_webhook_data(_webhook_trigger(), node_config) + + assert result["headers"].get("Authorization") == "Bearer token" + + def test_extract_and_validate_webhook_request_missing_required_param(self) -> None: + app = Flask(__name__) + with app.test_request_context( + "/webhook", + method="POST", + headers={"Content-Type": "application/json"}, + json={"message": "hello"}, + ): + node_config = { + "data": { + "method": "post", + "content_type": "application/json", + "params": [{"name": "version", "required": True}], + "body": [{"name": "message", "type": "string", "required": True}], + } + } + + with pytest.raises(ValueError, match="Required parameter missing: version"): + WebhookService.extract_and_validate_webhook_data(_webhook_trigger(), node_config) + + def test_extract_and_validate_webhook_request_missing_required_body_param(self) -> None: + app = Flask(__name__) + with app.test_request_context( + "/webhook", + method="POST", + headers={"Content-Type": "application/json"}, + json={}, + ): + node_config = { + "data": { + "method": "post", + "content_type": "application/json", + "body": [{"name": "message", "type": "string", "required": True}], + } + } + + with pytest.raises(ValueError, match="Required body parameter missing: message"): + WebhookService.extract_and_validate_webhook_data(_webhook_trigger(), node_config) + + def test_extract_and_validate_webhook_request_missing_required_file(self) -> None: + app = Flask(__name__) + with app.test_request_context( + "/webhook", + method="POST", + data={"note": "test"}, + content_type="multipart/form-data", + ): + node_config = { + "data": { + "method": "post", + "content_type": "multipart/form-data", + "body": [{"name": "file", "type": "file", "required": True}], + } + } + + result = WebhookService.extract_and_validate_webhook_data(_webhook_trigger(), node_config) + + assert result["files"] == {} + + def test_debug_mode_parameter_handling(self): + """Test that the debug mode parameter is properly handled in _prepare_webhook_execution.""" + from controllers.trigger.webhook import _prepare_webhook_execution + + # Mock the WebhookService methods + with ( + patch.object(WebhookService, "get_webhook_trigger_and_workflow", autospec=True) as mock_get_trigger, + patch.object(WebhookService, "extract_and_validate_webhook_data", autospec=True) as mock_extract, + ): + webhook_trigger = _webhook_trigger() + workflow = _workflow() + mock_config = {"data": {"test": "config"}} + mock_data = {"test": "data"} + + mock_get_trigger.return_value = (webhook_trigger, workflow, mock_config) + mock_extract.return_value = mock_data + + result = _prepare_webhook_execution("test_webhook", is_debug=False) + assert result == (webhook_trigger, workflow, mock_config, mock_data, None) + + # Reset mock + mock_get_trigger.reset_mock() + + result = _prepare_webhook_execution("test_webhook", is_debug=True) + assert result == (webhook_trigger, workflow, mock_config, mock_data, None) + + +def _workflow_trigger(**kwargs: Any) -> WorkflowWebhookTrigger: + values = { + "webhook_id": "webhook-123", + "tenant_id": "tenant-1", + "app_id": "app-1", + "node_id": "node-1", + "created_by": "user-1", + } + values.update(kwargs) + return WorkflowWebhookTrigger(**values) + + +class TestWebhookServiceAdditionalCases: + @staticmethod + @pytest.fixture + def flask_app() -> Flask: + return Flask(__name__) + + class TestWebhookServiceExtractionFallbacks: + def test_extract_webhook_data_should_use_text_fallback_for_unknown_content_type( + self, + flask_app: Flask, + monkeypatch: pytest.MonkeyPatch, + caplog: pytest.LogCaptureFixture, + ) -> None: + webhook_trigger = _workflow_trigger() + + with caplog.at_level(logging.WARNING, logger="services.trigger.webhook_service"): + with flask_app.test_request_context( + "/webhook", + method="POST", + headers={"Content-Type": "application/vnd.custom"}, + data="plain content", + ): + result = WebhookService.extract_webhook_data(webhook_trigger) + + assert result["body"] == {"raw": "plain content"} + assert any(r.levelno >= logging.WARNING for r in caplog.records) + + def test_extract_webhook_data_should_raise_for_request_too_large( + self, + flask_app: Flask, + monkeypatch: pytest.MonkeyPatch, + ) -> None: + from tests.helpers.unit.config_override import apply_config_overrides + + apply_config_overrides(monkeypatch, WEBHOOK_REQUEST_BODY_MAX_SIZE=1) + + with flask_app.test_request_context("/webhook", method="POST", data="ab"): + with pytest.raises(RequestEntityTooLarge): + WebhookService.extract_webhook_data(_workflow_trigger()) + + def test_extract_octet_stream_body_should_return_none_when_empty_payload(self, flask_app: Flask) -> None: + webhook_trigger = _workflow_trigger() + + with flask_app.test_request_context("/webhook", method="POST", data=b""): + body, files = WebhookService._extract_octet_stream_body(webhook_trigger) + + assert body == {"raw": None} + assert files == {} + + def test_extract_octet_stream_body_should_return_none_when_processing_raises( + self, + flask_app: Flask, + monkeypatch: pytest.MonkeyPatch, + ) -> None: + webhook_trigger = _workflow_trigger() + monkeypatch.setattr( + WebhookService, "_detect_binary_mimetype", MagicMock(return_value="application/octet-stream") + ) + monkeypatch.setattr(WebhookService, "_create_file_from_binary", MagicMock(side_effect=RuntimeError("boom"))) + + with flask_app.test_request_context("/webhook", method="POST", data=b"abc"): + body, files = WebhookService._extract_octet_stream_body(webhook_trigger) + + assert body == {"raw": None} + assert files == {} + + def test_extract_text_body_should_return_empty_string_when_request_read_fails( + self, + flask_app: Flask, + monkeypatch: pytest.MonkeyPatch, + ) -> None: + monkeypatch.setattr("flask.wrappers.Request.get_data", MagicMock(side_effect=RuntimeError("read error"))) + + with flask_app.test_request_context("/webhook", method="POST", data="abc"): + body, files = WebhookService._extract_text_body() + + assert body == {"raw": ""} + assert files == {} + + def test_detect_binary_mimetype_should_fallback_when_magic_raises( + self, + monkeypatch: pytest.MonkeyPatch, + ) -> None: + fake_magic = MagicMock() + fake_magic.from_buffer.side_effect = RuntimeError("magic failed") + monkeypatch.setattr(service_module, "magic", fake_magic) + + result = WebhookService._detect_binary_mimetype(b"binary") + + assert result == "application/octet-stream" + + def test_process_file_uploads_should_use_octet_stream_fallback_when_mimetype_unknown( + self, + monkeypatch: pytest.MonkeyPatch, + ) -> None: + webhook_trigger = _workflow_trigger(created_by="user-1", tenant_id="tenant-1") + file_obj = MagicMock() + file_obj.to_dict.return_value = {"id": "f-1"} + monkeypatch.setattr(WebhookService, "_create_file_from_binary", MagicMock(return_value=file_obj)) + monkeypatch.setattr(service_module.mimetypes, "guess_type", MagicMock(return_value=(None, None))) + + uploaded = MagicMock() + uploaded.filename = "file.unknown" + uploaded.content_type = None + uploaded.read.return_value = b"content" + + result = WebhookService._process_file_uploads({"f": uploaded}, webhook_trigger) + + assert result == {"f": {"id": "f-1"}} + + def test_create_file_from_binary_should_call_tool_file_manager_and_file_factory( + self, + monkeypatch: pytest.MonkeyPatch, + ) -> None: + webhook_trigger = _workflow_trigger(created_by="user-1", tenant_id="tenant-1") + manager = MagicMock() + manager.create_file_by_raw.return_value = SimpleNamespace(id="tool-file-1") + monkeypatch.setattr(service_module, "ToolFileManager", MagicMock(return_value=manager)) + expected_file = MagicMock() + monkeypatch.setattr( + service_module.file_factory, "build_from_mapping", MagicMock(return_value=expected_file) + ) + + result = WebhookService._create_file_from_binary(b"abc", "text/plain", webhook_trigger) + + assert result is expected_file + manager.create_file_by_raw.assert_called_once() + + class TestWebhookServiceValidationAndConversion: + @pytest.mark.parametrize( + ("raw_value", "param_type", "expected"), + [ + ("42", SegmentType.NUMBER, 42), + ("3.14", SegmentType.NUMBER, float("3.14")), + ("yes", SegmentType.BOOLEAN, True), + ("no", SegmentType.BOOLEAN, False), + ], + ) + def test_convert_form_value_should_convert_supported_types( + self, + raw_value: str, + param_type: str, + expected: Any, + ) -> None: + result = WebhookService._convert_form_value("param", raw_value, param_type) + assert result == expected + + def test_convert_form_value_should_raise_for_unsupported_type(self) -> None: + with pytest.raises(ValueError, match="Unsupported type"): + WebhookService._convert_form_value("p", "x", SegmentType.FILE) + + def test_validate_json_value_should_return_original_for_unmapped_supported_segment_type( + self, + monkeypatch: pytest.MonkeyPatch, + caplog: pytest.LogCaptureFixture, + ) -> None: + with caplog.at_level(logging.WARNING, logger="services.trigger.webhook_service"): + result = WebhookService._validate_json_value("param", {"x": 1}, "unsupported-type") + + assert result == {"x": 1} + assert any(r.levelno >= logging.WARNING for r in caplog.records) + + def test_validate_and_convert_value_should_wrap_conversion_errors(self) -> None: + with pytest.raises(ValueError, match="validation failed"): + WebhookService._validate_and_convert_value("param", "bad", SegmentType.NUMBER, is_form_data=True) + + def test_process_parameters_should_raise_when_required_parameter_missing(self) -> None: + raw_params = {"optional": "x"} + config = [WebhookParameter(name="required_param", type=SegmentType.STRING, required=True)] + + with pytest.raises(ValueError, match="Required parameter missing"): + WebhookService._process_parameters(raw_params, config, is_form_data=True) + + def test_process_parameters_should_include_unconfigured_parameters(self) -> None: + raw_params = {"known": "1", "unknown": "x"} + config = [WebhookParameter(name="known", type=SegmentType.NUMBER, required=False)] + + result = WebhookService._process_parameters(raw_params, config, is_form_data=True) + + assert result == {"known": 1, "unknown": "x"} + + def test_process_body_parameters_should_raise_when_required_text_raw_is_missing(self) -> None: + with pytest.raises(ValueError, match="Required body content missing"): + WebhookService._process_body_parameters( + raw_body={"raw": ""}, + body_configs=[WebhookBodyParameter(name="raw", required=True)], + content_type=ContentType.TEXT, + ) + + def test_process_body_parameters_should_skip_file_config_for_multipart_form_data(self) -> None: + raw_body = {"message": "hello", "extra": "x"} + body_configs = [ + WebhookBodyParameter(name="upload", type=SegmentType.FILE, required=True), + WebhookBodyParameter(name="message", type=SegmentType.STRING, required=True), + ] + + result = WebhookService._process_body_parameters(raw_body, body_configs, ContentType.FORM_DATA) + + assert result == {"message": "hello", "extra": "x"} + + def test_validate_required_headers_should_accept_sanitized_header_names(self) -> None: + headers = {"x_api_key": "123"} + configs = [WebhookParameter(name="x-api-key", required=True)] + + WebhookService._validate_required_headers(headers, configs) + + def test_validate_required_headers_should_raise_when_required_header_missing(self) -> None: + headers = {"x-other": "123"} + configs = [WebhookParameter(name="x-api-key", required=True)] + + with pytest.raises(ValueError, match="Required header missing"): + WebhookService._validate_required_headers(headers, configs) + + def test_validate_http_metadata_should_return_content_type_mismatch_error(self) -> None: + webhook_data = {"method": "POST", "headers": {"Content-Type": "application/json"}} + node_data = WebhookData(method="post", content_type=ContentType.TEXT) + + result = WebhookService._validate_http_metadata(webhook_data, node_data) + + assert result["valid"] is False + assert "Content-type mismatch" in result["error"] + + def test_extract_content_type_should_fallback_to_lowercase_header_key(self) -> None: + headers = {"content-type": "application/json; charset=utf-8"} + assert WebhookService._extract_content_type(headers) == "application/json" + + def test_build_workflow_inputs_should_include_expected_keys(self) -> None: + webhook_data = {"headers": {"h": "v"}, "query_params": {"q": 1}, "body": {"b": 2}} + + result = WebhookService.build_workflow_inputs(webhook_data) + + assert result["webhook_data"] == webhook_data + assert result["webhook_headers"] == {"h": "v"} + assert result["webhook_query_params"] == {"q": 1} + assert result["webhook_body"] == {"b": 2} + + class TestWebhookServiceUtilities: + def test_generate_webhook_response_should_fallback_when_response_body_is_not_json(self) -> None: + node_config = {"data": {"status_code": 200, "response_body": "{bad-json"}} + + body, status = WebhookService.generate_webhook_response(node_config) + + assert status == 200 + assert "message" in body + + def test_generate_webhook_id_should_return_24_character_identifier(self) -> None: + webhook_id = WebhookService.generate_webhook_id() + + assert isinstance(webhook_id, str) + assert len(webhook_id) == 24 + + def test_sanitize_key_should_return_original_value_for_non_string_input(self) -> None: + result = WebhookService._sanitize_key(123) # type: ignore[arg-type] + assert result == 123 + + +class TestWebhookServiceRelationshipsCases: + @staticmethod + def test_draft_sync_preserves_webhook_id_when_deleted_node_is_restored( + monkeypatch: pytest.MonkeyPatch, sqlite_engine: _webhook_service_relationships_Engine + ) -> None: + webhook_trigger_table = cast(Table, WorkflowWebhookTrigger.__table__) + TypeBase.metadata.create_all(sqlite_engine, tables=[webhook_trigger_table]) + session_maker = sessionmaker(bind=sqlite_engine, expire_on_commit=False) + app = cast( + App, + SimpleNamespace(id="00000000-0000-0000-0000-000000000001", tenant_id="tenant-1", created_by="user-1"), + ) + assert app.created_by is not None + webhook_trigger = WorkflowWebhookTrigger( + app_id=app.id, + tenant_id=app.tenant_id, + node_id="webhook-node", + webhook_id="stable-webhook-id", + created_by=app.created_by, + ) + with session_maker.begin() as session: + session.add(webhook_trigger) + original_record_id = webhook_trigger.id + + redis = MagicMock() + redis.get.return_value = None + redis.lock.return_value.acquire.return_value = True + monkeypatch.setattr(webhook_service, "db", SimpleNamespace(engine=sqlite_engine)) + monkeypatch.setattr(webhook_service, "redis_client", redis) + + deleted_draft = cast(Workflow, SimpleNamespace(walk_nodes=lambda _node_type: iter(()))) + WebhookService.sync_webhook_relationships(app, deleted_draft, remove_stale=False) + + restored_draft = cast( + Workflow, + SimpleNamespace( + walk_nodes=lambda _node_type: iter([("webhook-node", {"type": TRIGGER_WEBHOOK_NODE_TYPE})]) + ), + ) + WebhookService.sync_webhook_relationships(app, restored_draft, remove_stale=False) + + with Session(sqlite_engine) as session: + records = session.scalars(select(WorkflowWebhookTrigger)).all() + + assert len(records) == 1 + assert records[0].id == original_record_id + assert records[0].webhook_id == "stable-webhook-id" + + @staticmethod + def test_published_sync_removes_stale_webhook_relationships( + monkeypatch: pytest.MonkeyPatch, sqlite_engine: _webhook_service_relationships_Engine + ) -> None: + webhook_trigger_table = cast(Table, WorkflowWebhookTrigger.__table__) + TypeBase.metadata.create_all(sqlite_engine, tables=[webhook_trigger_table]) + session_maker = sessionmaker(bind=sqlite_engine, expire_on_commit=False) + app = cast( + App, + SimpleNamespace(id="00000000-0000-0000-0000-000000000001", tenant_id="tenant-1", created_by="user-1"), + ) + assert app.created_by is not None + webhook_trigger = WorkflowWebhookTrigger( + app_id=app.id, + tenant_id=app.tenant_id, + node_id="deleted-webhook-node", + webhook_id="deleted-webhook-id", + created_by=app.created_by, + ) + with session_maker.begin() as session: + session.add(webhook_trigger) + + redis = MagicMock() + redis.get.return_value = None + redis.lock.return_value.acquire.return_value = True + monkeypatch.setattr(webhook_service, "db", SimpleNamespace(engine=sqlite_engine)) + monkeypatch.setattr(webhook_service, "redis_client", redis) + + published_workflow = cast(Workflow, SimpleNamespace(walk_nodes=lambda _node_type: iter(()))) + WebhookService.sync_webhook_relationships(app, published_workflow, remove_stale=True) + + with session_maker() as session: + records = session.scalars(select(WorkflowWebhookTrigger)).all() + assert records == [] + redis.delete.assert_called_once_with( + f"{WebhookService.__WEBHOOK_NODE_CACHE_KEY__}:{app.id}:deleted-webhook-node" + ) diff --git a/api/tests/unit_tests/services/workflow/test_draft_var_loader_simple.py b/api/tests/unit_tests/services/workflow/test_draft_var_loader_simple.py deleted file mode 100644 index 2fe76f7becc..00000000000 --- a/api/tests/unit_tests/services/workflow/test_draft_var_loader_simple.py +++ /dev/null @@ -1,411 +0,0 @@ -"""Simplified unit tests for DraftVarLoader focusing on core functionality.""" - -import json -from datetime import datetime -from unittest.mock import Mock, patch - -import pytest -from sqlalchemy import Engine -from sqlalchemy.orm import Session - -from core.workflow.file_reference import build_file_reference -from extensions.storage.storage_type import StorageType -from graphon.file import File, FileTransferMethod, FileType -from graphon.variables.segments import StringSegment -from graphon.variables.types import SegmentType -from models.enums import CreatorUserRole -from models.model import UploadFile -from models.workflow import WorkflowDraftVariable, WorkflowDraftVariableFile -from services.workflow_draft_variable_service import DraftVarLoader - - -def _persist_offloaded_variable( - sqlite_session: Session, - *, - node_id: str, - name: str, -) -> WorkflowDraftVariable: - upload_file = UploadFile( - tenant_id="test-tenant-id", - storage_type=StorageType.LOCAL, - key=f"storage/key/{name}.txt", - name=f"{name}.txt", - size=10, - extension=".txt", - mime_type="text/plain", - created_by_role=CreatorUserRole.ACCOUNT, - created_by="test-user-id", - created_at=datetime(2025, 1, 1), - used=True, - ) - variable_file = WorkflowDraftVariableFile( - tenant_id="test-tenant-id", - app_id="test-app-id", - user_id="test-user-id", - upload_file_id=upload_file.id, - size=10, - length=None, - value_type=SegmentType.STRING, - ) - draft_variable = WorkflowDraftVariable.new_node_variable( - app_id="test-app-id", - user_id="test-user-id", - node_id=node_id, - name=name, - value=StringSegment(value="truncated"), - node_execution_id=f"execution-{node_id}", - file_id=variable_file.id, - ) - sqlite_session.add_all([upload_file, variable_file, draft_variable]) - return draft_variable - - -class TestDraftVarLoaderSimple: - """Simplified unit tests for DraftVarLoader core methods.""" - - @pytest.fixture - def draft_var_loader(self, sqlite_engine: Engine): - """Create DraftVarLoader instance for testing.""" - return DraftVarLoader( - engine=sqlite_engine, - app_id="test-app-id", - tenant_id="test-tenant-id", - user_id="test-user-id", - fallback_variables=[], - ) - - def test_load_offloaded_variable_object_type_unit(self, draft_var_loader): - """Test _load_offloaded_variable with object type - isolated unit test.""" - # Create mock objects - upload_file = UploadFile( - tenant_id="tenant-id", - storage_type="opendal", - key="storage/key/test.json", - name="test.txt", - size=0, - extension="txt", - mime_type="text/plain", - created_by_role="account", - created_by="account-id", - created_at=datetime.now(), - used=False, - ) - - variable_file = WorkflowDraftVariableFile( - tenant_id="tenant-id", - app_id="app-id", - user_id="user-id", - upload_file_id="upload-file-id", - size=0, - length=0, - value_type=SegmentType.OBJECT, - ) - variable_file.upload_file = upload_file - - draft_var = WorkflowDraftVariable( - id="draft-var-id", - node_id="test-node-id", - name="test_object", - description="test description", - selector=json.dumps(["test-node-id", "test_object"]), - variable_file=variable_file, - ) - - test_object = {"key1": "value1", "key2": 42} - test_json_content = json.dumps(test_object, ensure_ascii=False, separators=(",", ":")) - - with patch("services.workflow_draft_variable_service.storage") as mock_storage: - mock_storage.load.return_value = test_json_content.encode() - - # Execute the method - selector_tuple, variable = draft_var_loader._load_offloaded_variable(draft_var) - - # Verify results - assert selector_tuple == ("test-node-id", "test_object") - assert variable.id == "draft-var-id" - assert variable.name == "test_object" - assert variable.description == "test description" - assert variable.value == test_object - - # Verify method calls - mock_storage.load.assert_called_once_with("storage/key/test.json") - - def test_load_offloaded_variable_missing_variable_file_unit(self, draft_var_loader): - """Test that assertion error is raised when variable_file is None.""" - draft_var = WorkflowDraftVariable( - variable_file=None, - ) - - with pytest.raises(AssertionError): - draft_var_loader._load_offloaded_variable(draft_var) - - def test_load_offloaded_variable_missing_upload_file_unit(self, draft_var_loader): - """Test that assertion error is raised when upload_file is None.""" - variable_file = WorkflowDraftVariableFile( - tenant_id="tenant-id", - app_id="app-id", - user_id="user-id", - upload_file_id="upload-file-id", - size=0, - length=0, - value_type="file", - ) - variable_file.upload_file = None - - draft_var = WorkflowDraftVariable( - variable_file=variable_file, - ) - - with pytest.raises(AssertionError): - draft_var_loader._load_offloaded_variable(draft_var) - - def test_load_variables_empty_selectors_unit(self, draft_var_loader): - """Test load_variables returns empty list for empty selectors.""" - result = draft_var_loader.load_variables([]) - assert result == [] - - def test_selector_to_tuple_unit(self, draft_var_loader): - """Test _selector_to_tuple method.""" - selector = ["node_id", "var_name", "extra_field"] - result = draft_var_loader._selector_to_tuple(selector) - assert result == ("node_id", "var_name") - - def test_load_offloaded_variable_array_type_unit(self, draft_var_loader): - """Test _load_offloaded_variable with array type - isolated unit test.""" - # Create mock objects - upload_file = UploadFile( - tenant_id="tenant-id", - storage_type="opendal", - key="storage/key/test_array.json", - name="test.txt", - size=0, - extension="txt", - mime_type="text/plain", - created_by_role="account", - created_by="account-id", - created_at=datetime.now(), - used=False, - ) - - variable_file = WorkflowDraftVariableFile( - tenant_id="tenant-id", - app_id="app-id", - user_id="user-id", - upload_file_id="upload-file-id", - size=0, - length=0, - value_type=SegmentType.ARRAY_ANY, - ) - variable_file.upload_file = upload_file - - draft_var = WorkflowDraftVariable( - id="draft-var-id", - node_id="test-node-id", - name="test_array", - description="test array description", - selector=json.dumps(["test-node-id", "test_array"]), - variable_file=variable_file, - ) - - test_array = ["item1", 2, True] - test_json_content = json.dumps(test_array) - - with patch("services.workflow_draft_variable_service.storage") as mock_storage: - mock_storage.load.return_value = test_json_content.encode() - # Execute the method - selector_tuple, variable = draft_var_loader._load_offloaded_variable(draft_var) - - # Verify results - assert selector_tuple == ("test-node-id", "test_array") - assert variable.id == "draft-var-id" - assert variable.name == "test_array" - assert variable.description == "test array description" - - # Verify method calls - mock_storage.load.assert_called_once_with("storage/key/test_array.json") - - def test_load_offloaded_variable_file_type_rebuilds_storage_backed_payload(self, draft_var_loader): - upload_file = UploadFile( - tenant_id="tenant-id", - storage_type="opendal", - key="storage/key/test_file.json", - name="test.txt", - size=0, - extension="txt", - mime_type="text/plain", - created_by_role="account", - created_by="account-id", - created_at=datetime.now(), - used=False, - ) - - variable_file = WorkflowDraftVariableFile( - tenant_id="tenant-id", - app_id="app-id", - user_id="user-id", - upload_file_id="upload-file-id", - size=0, - length=0, - value_type=SegmentType.FILE, - ) - variable_file.upload_file = upload_file - - draft_var = WorkflowDraftVariable( - id="draft-var-id", - app_id="app-1", - node_id="test-node-id", - name="test_file", - description="test file description", - ) - draft_var._set_selector(["test-node-id", "test_file"]) - draft_var.variable_file = variable_file - - persisted_file = File( - file_id="file-1", - file_type=FileType.DOCUMENT, - transfer_method=FileTransferMethod.LOCAL_FILE, - reference=build_file_reference(record_id="upload-1", storage_key="legacy-storage-key"), - filename="test.txt", - extension=".txt", - mime_type="text/plain", - size=12, - ) - rebuilt_file = File( - file_id="file-1", - file_type=FileType.DOCUMENT, - transfer_method=FileTransferMethod.LOCAL_FILE, - reference=build_file_reference(record_id="upload-1"), - filename="test.txt", - extension=".txt", - mime_type="text/plain", - size=12, - storage_key="canonical-storage-key", - ) - - raw_file = { - **persisted_file.model_dump(mode="json"), - "tenant_id": "legacy-tenant", - } - - with ( - patch("services.workflow_draft_variable_service.storage") as mock_storage, - patch("models.workflow._resolve_workflow_app_tenant_id", return_value="tenant-1"), - patch("models.workflow.build_file_from_stored_mapping", return_value=rebuilt_file) as rebuild_file, - ): - mock_storage.load.return_value = json.dumps(raw_file).encode() - - selector_tuple, variable = draft_var_loader._load_offloaded_variable(draft_var) - - assert selector_tuple == ("test-node-id", "test_file") - assert variable.id == "draft-var-id" - assert variable.name == "test_file" - assert variable.description == "test file description" - assert variable.value == rebuilt_file - rebuild_file.assert_called_once_with(file_mapping=raw_file, tenant_id="tenant-1") - - @pytest.mark.parametrize( - "sqlite_session", - [(WorkflowDraftVariable, WorkflowDraftVariableFile, UploadFile)], - indirect=True, - ) - def test_load_variables_with_offloaded_variables_unit( - self, - draft_var_loader: DraftVarLoader, - sqlite_session: Session, - ): - """Test load_variables method with mix of regular and offloaded variables.""" - selectors = [["node1", "regular_var"], ["node2", "offloaded_var"]] - regular_draft_var = WorkflowDraftVariable.new_node_variable( - app_id="test-app-id", - user_id="test-user-id", - node_id="node1", - name="regular_var", - value=StringSegment(value="regular_value"), - node_execution_id="execution-node1", - ) - regular_draft_var.description = "regular description" - offloaded_draft_var = _persist_offloaded_variable( - sqlite_session, - node_id="node2", - name="offloaded_var", - ) - distractor = WorkflowDraftVariable.new_node_variable( - app_id="test-app-id", - user_id="another-user", - node_id="node1", - name="regular_var", - value=StringSegment(value="wrong user"), - node_execution_id="execution-distractor", - ) - sqlite_session.add_all([regular_draft_var, distractor]) - sqlite_session.commit() - - offloaded_variable = Mock() - offloaded_variable.id = offloaded_draft_var.id - offloaded_variable.selector = ["node2", "offloaded_var"] - - with ( - patch("services.workflow_draft_variable_service.StorageKeyLoader"), - patch.object( - draft_var_loader, - "_load_offloaded_variable", - return_value=(("node2", "offloaded_var"), offloaded_variable), - ) as load_offloaded, - patch("services.workflow_draft_variable_service.ThreadPoolExecutor") as executor_cls, - ): - executor = executor_cls.return_value.__enter__.return_value - executor.map.side_effect = lambda function, values: [function(value) for value in values] - - result = draft_var_loader.load_variables(selectors) - - assert {variable.id for variable in result} == {regular_draft_var.id, offloaded_draft_var.id} - load_offloaded.assert_called_once() - loaded_offloaded = load_offloaded.call_args.args[0] - assert isinstance(loaded_offloaded, WorkflowDraftVariable) - assert loaded_offloaded.id == offloaded_draft_var.id - assert loaded_offloaded.variable_file is not None - assert loaded_offloaded.variable_file.upload_file is not None - assert loaded_offloaded.variable_file.upload_file.key == "storage/key/offloaded_var.txt" - - @pytest.mark.parametrize( - "sqlite_session", - [(WorkflowDraftVariable, WorkflowDraftVariableFile, UploadFile)], - indirect=True, - ) - def test_load_variables_all_offloaded_variables_unit( - self, - draft_var_loader: DraftVarLoader, - sqlite_session: Session, - ): - """Test load_variables method with only offloaded variables.""" - selectors = [["node1", "offloaded_var1"], ["node2", "offloaded_var2"]] - offloaded_var1 = _persist_offloaded_variable( - sqlite_session, - node_id="node1", - name="offloaded_var1", - ) - offloaded_var2 = _persist_offloaded_variable( - sqlite_session, - node_id="node2", - name="offloaded_var2", - ) - sqlite_session.commit() - - with ( - patch("services.workflow_draft_variable_service.StorageKeyLoader"), - patch("services.workflow_draft_variable_service.ThreadPoolExecutor") as executor_cls, - ): - executor = executor_cls.return_value.__enter__.return_value - executor.map.return_value = [ - (("node1", "offloaded_var1"), Mock()), - (("node2", "offloaded_var2"), Mock()), - ] - - result = draft_var_loader.load_variables(selectors) - - assert len(result) == 2 - executor_cls.assert_called_once_with(max_workers=10) - executor.map.assert_called_once() - loaded_draft_vars = executor.map.call_args.args[1] - assert {variable.id for variable in loaded_draft_vars} == {offloaded_var1.id, offloaded_var2.id} - assert all(variable.variable_file.upload_file is not None for variable in loaded_draft_vars) diff --git a/api/tests/unit_tests/services/workflow/test_pipeline_workflow_tools.py b/api/tests/unit_tests/services/workflow/test_pipeline_workflow_tools.py deleted file mode 100644 index 2696f03daf6..00000000000 --- a/api/tests/unit_tests/services/workflow/test_pipeline_workflow_tools.py +++ /dev/null @@ -1,231 +0,0 @@ -"""Pipeline tool lifecycle through the real Engine and SQLAlchemy repositories.""" - -import json - -import pytest -from sqlalchemy import Engine, select -from sqlalchemy.orm import Session - -from clients.agent_backend import FakeAgentBackendRunClient, FakeAgentBackendScenario -from core.app.apps.pipeline.pipeline_runner import PipelineRunner -from core.app.entities.app_invoke_entities import InvokeFrom -from core.repositories.sqlalchemy_workflow_execution_repository import SQLAlchemyWorkflowExecutionRepository -from core.repositories.sqlalchemy_workflow_node_execution_repository import SQLAlchemyWorkflowNodeExecutionRepository -from core.workflow.nodes.human_input.entities import HumanInputNodeData, UserActionConfig -from models import Account -from models.agent import AgentWorkspace, AgentWorkspaceBinding, WorkflowAgentBindingType -from models.dataset import Document -from models.enums import WorkflowRunTriggeredFrom -from models.human_input import HumanInputForm -from models.model import App, AppMode -from models.tools import WorkflowToolProvider -from models.workflow import Workflow, WorkflowNodeExecutionModel, WorkflowNodeExecutionTriggeredFrom, WorkflowRun -from services.agent.workspace_service import AgentWorkspaceService -from services.workflow_run_agg import WorkflowRunAgg -from tests.unit_tests.core.app.apps.pipeline.test_pipeline_runner import ( - _build_app_generate_entity, - _build_runner, - _document, - _persist_scope, - _workflow, -) -from tests.unit_tests.core.workflow.nodes.agent_v2.test_binding_resolver import _agent, _binding, _snapshot -from tests.unit_tests.core.workflow.nodes.agent_v2.test_session_store import _install_backend_client - - -def _pipeline_tool_runner( - sqlite_engine: Engine, invoke_from: InvokeFrom, source_node_data: dict[str, object] -) -> tuple[PipelineRunner, SQLAlchemyWorkflowExecutionRepository, SQLAlchemyWorkflowNodeExecutionRepository]: - source_graph = { - "nodes": [ - {"id": "source-start", "data": {"type": "start", "title": "Start", "variables": list[object]()}}, - {"id": "source-node", "data": source_node_data}, - {"id": "source-end", "data": {"type": "end", "title": "End", "outputs": list[object]()}}, - ], - "edges": [ - {"source": "source-start", "target": "source-node"}, - {"source": "source-node", "target": "source-end", "sourceHandle": "source"}, - ], - } - with Session(sqlite_engine, expire_on_commit=False) as session: - session.add( - App( - id="source-app", - tenant_id="tenant", - name="Source", - mode=AppMode.WORKFLOW, - enable_site=False, - enable_api=False, - ) - ) - session.add( - Workflow( - id="source-workflow", - tenant_id="tenant", - app_id="source-app", - type="workflow", - version="v1", - graph=json.dumps(source_graph), - features="{}", - created_by="user", - ) - ) - provider = WorkflowToolProvider( - tenant_id="tenant", - app_id="source-app", - user_id="user", - version="v1", - name="source", - label="Source", - description="Source", - icon="{}", - ) - session.add(provider) - session.flush() - graph = { - "nodes": [ - {"id": "start", "data": {"type": "start", "title": "Start", "variables": list[object]()}}, - { - "id": "tool", - "data": { - "type": "tool", - "title": "Tool", - "provider_id": provider.id, - "provider_type": "workflow", - "provider_name": "source", - "tool_name": "source", - "tool_label": "Source", - "tool_configurations": dict[str, object](), - "tool_parameters": dict[str, object](), - }, - }, - ], - "edges": [{"source": "start", "target": "tool"}], - } - _, _, workflow = _persist_scope(session, workflow=_workflow(graph=graph), documents=(_document(),)) - entity = _build_app_generate_entity().model_copy( - update={"invoke_from": invoke_from, "document_id": None if invoke_from == InvokeFrom.DEBUGGER else "doc"} - ) - runner = _build_runner(entity, sqlite_engine, workflow) - user = Account(name="User", email="pipeline@example.com") - user.id = entity.user_id - runs = SQLAlchemyWorkflowExecutionRepository( - sqlite_engine, - "tenant", - user, - "pipe", - WorkflowRunTriggeredFrom.RAG_PIPELINE_DEBUGGING - if invoke_from == InvokeFrom.DEBUGGER - else WorkflowRunTriggeredFrom.RAG_PIPELINE_RUN, - ) - nodes = SQLAlchemyWorkflowNodeExecutionRepository( - sqlite_engine, "tenant", user, "pipe", WorkflowNodeExecutionTriggeredFrom.RAG_PIPELINE_RUN - ) - return runner, runs, nodes - - -@pytest.mark.parametrize("invoke_from", [InvokeFrom.DEBUGGER, InvokeFrom.PUBLISHED_PIPELINE]) -def test_pipeline_rejects_human_input_tool_before_creating_run_or_form( - sqlite_engine: Engine, invoke_from: InvokeFrom -) -> None: - human_input = HumanInputNodeData(title="Approval", user_actions=[UserActionConfig(id="approve", title="Approve")]) - runner, runs, nodes = _pipeline_tool_runner(sqlite_engine, invoke_from, human_input.model_dump(mode="json")) - - with pytest.raises(ValueError, match="Human Input nodes are not supported"): - WorkflowRunAgg.run(runner, None, runs, nodes) - - with Session(sqlite_engine) as session: - assert session.scalar(select(WorkflowRun)) is None - assert session.scalar(select(WorkflowNodeExecutionModel)) is None - assert session.scalar(select(HumanInputForm)) is None - document = session.get(Document, "doc") - assert document is not None - if invoke_from == InvokeFrom.PUBLISHED_PIPELINE: - assert document.indexing_status == "error" - assert document.error == "Human Input nodes are not supported in this execution context." - else: - assert document.indexing_status == "waiting" - assert document.error is None - - -@pytest.mark.parametrize("invoke_from", [InvokeFrom.DEBUGGER, InvokeFrom.PUBLISHED_PIPELINE]) -@pytest.mark.parametrize( - ("scenario", "status"), - [(FakeAgentBackendScenario.SUCCESS, "succeeded"), (FakeAgentBackendScenario.FAILED, "failed")], -) -def test_pipeline_retires_nested_agent_workspace_when_run_ends( - sqlite_engine: Engine, - monkeypatch: pytest.MonkeyPatch, - invoke_from: InvokeFrom, - scenario: FakeAgentBackendScenario, - status: str, -) -> None: - runner, runs, nodes = _pipeline_tool_runner( - sqlite_engine, invoke_from, {"type": "agent", "version": "2", "agent_node_kind": "dify_agent", "title": "Agent"} - ) - with Session(sqlite_engine) as session: - agent = _agent(tenant_id="tenant") - session.add(agent) - session.flush() - snapshot = _snapshot(tenant_id="tenant", agent_id=agent.id) - session.add(snapshot) - session.flush() - session.add( - _binding( - ids={ - "tenant_id": "tenant", - "app_id": "source-app", - "workflow_id": "source-workflow", - "node_id": "source-node", - }, - agent_id=agent.id, - snapshot_id=snapshot.id, - binding_type=WorkflowAgentBindingType.INLINE_AGENT, - ) - ) - session.commit() - backend_client = _install_backend_client(monkeypatch) - monkeypatch.setattr( - "clients.agent_backend.factory.create_agent_backend_run_client", - lambda **_kwargs: FakeAgentBackendRunClient(scenario=scenario), - ) - monkeypatch.setattr( - "core.workflow.nodes.agent_v2.runtime_request_builder.resolve_model_context_window", lambda **_kwargs: None - ) - collections: list[dict[str, object]] = [] - - def collect_task(**kwargs: object) -> None: - collections.append(kwargs) - - monkeypatch.setattr("tasks.collect_agent_resources_task.collect_agent_resources.delay", collect_task) - - WorkflowRunAgg.run(runner, None, runs, nodes) - list(runner._queue_manager.listen()) - - with Session(sqlite_engine) as session: - run = session.get(WorkflowRun, runner.application_generate_entity.workflow_execution_id) - assert run is not None - assert run.status == status - agent_execution = session.scalars( - select(WorkflowNodeExecutionModel).where(WorkflowNodeExecutionModel.node_id == "source-node") - ).one() - assert agent_execution.status == status - workspace = session.scalars(select(AgentWorkspace)).one() - binding = session.scalars(select(AgentWorkspaceBinding)).one() - assert (workspace.app_id, workspace.owner_id, workspace.status) == ("source-app", run.id, "retired") - assert binding.status == "retired" - workspace_id, binding_id = workspace.id, binding.id - assert collections == [ - { - "tenant_id": "tenant", - "workspace_ids": [workspace_id], - "binding_ids": [], - "home_snapshot_ids": [], - "purge_agent_ids": [], - } - ] - AgentWorkspaceService.collect_retired_workspace(tenant_id="tenant", workspace_id=workspace_id) - backend_client.destroy_execution_binding_sync.assert_called_once() - with Session(sqlite_engine) as session: - assert session.get(AgentWorkspace, workspace_id) is None - assert session.get(AgentWorkspaceBinding, binding_id) is None diff --git a/api/tests/unit_tests/services/workflow/test_workflow_converter_additional.py b/api/tests/unit_tests/services/workflow/test_workflow_converter.py similarity index 100% rename from api/tests/unit_tests/services/workflow/test_workflow_converter_additional.py rename to api/tests/unit_tests/services/workflow/test_workflow_converter.py diff --git a/api/tests/unit_tests/services/workflow/test_workflow_draft_variable_service.py b/api/tests/unit_tests/services/workflow/test_workflow_draft_variable_service.py deleted file mode 100644 index f6ee73514b8..00000000000 --- a/api/tests/unit_tests/services/workflow/test_workflow_draft_variable_service.py +++ /dev/null @@ -1,709 +0,0 @@ -import dataclasses -import json -import secrets -import uuid -from unittest.mock import Mock, patch - -import pytest -from sqlalchemy.orm import Session - -from core.workflow.system_variables import SystemVariableKey -from core.workflow.variable_prefixes import ( - CONVERSATION_VARIABLE_NODE_ID, - ENVIRONMENT_VARIABLE_NODE_ID, - SYSTEM_VARIABLE_NODE_ID, -) -from extensions.storage.storage_type import StorageType -from graphon.enums import BuiltinNodeTypes -from graphon.file import File, FileTransferMethod, FileType -from graphon.variables.segments import StringSegment -from graphon.variables.types import SegmentType -from libs.datetime_utils import naive_utc_now -from libs.uuid_utils import uuidv7 -from models.account import Account -from models.enums import CreatorUserRole, DraftVariableType -from models.model import UploadFile -from models.workflow import ( - Workflow, - WorkflowDraftVariable, - WorkflowDraftVariableFile, - WorkflowNodeExecutionModel, - is_system_variable_editable, -) -from services.variable_truncator import TruncationResult -from services.workflow_draft_variable_service import ( - DraftVariableSaver, - VariableResetError, - WorkflowDraftVariableService, - _model_to_insertion_dict, -) - -SQLITE_MODELS = (Workflow, WorkflowDraftVariable, WorkflowDraftVariableFile, WorkflowNodeExecutionModel) -pytestmark = [ - pytest.mark.usefixtures("sqlite_session"), - pytest.mark.parametrize("sqlite_session", [SQLITE_MODELS], indirect=True), -] - - -class TestDraftVariableSaver: - def _get_test_app_id(self): - suffix = secrets.token_hex(6) - return f"test_app_id_{suffix}" - - def test__should_variable_be_visible(self, sqlite_session: Session): - mock_user = Account(name="test", email="test@example.com") - mock_user.id = str(uuid.uuid4()) - test_app_id = self._get_test_app_id() - saver = DraftVariableSaver( - session=sqlite_session, - tenant_id="test-tenant-id", - app_id=test_app_id, - node_id="test_node_id", - node_type=BuiltinNodeTypes.START, - node_execution_id="test_execution_id", - user=mock_user, - ) - assert saver._should_variable_be_visible("123_456", BuiltinNodeTypes.IF_ELSE, "output") is False - assert saver._should_variable_be_visible("123", BuiltinNodeTypes.START, "output") is True - - def test__normalize_variable_for_start_node(self, sqlite_session: Session): - @dataclasses.dataclass(frozen=True) - class TestCase: - name: str - input_node_id: str - input_name: str - expected_node_id: str - expected_name: str - - _node_id = "1747228642872" - cases = [ - TestCase( - name="name with `sys.` prefix should return the system node_id", - input_node_id=_node_id, - input_name="sys.workflow_id", - expected_node_id=SYSTEM_VARIABLE_NODE_ID, - expected_name="workflow_id", - ), - TestCase( - name="name without `sys.` prefix should return the original input node_id", - input_node_id=_node_id, - input_name="start_input", - expected_node_id=_node_id, - expected_name="start_input", - ), - TestCase( - name="name with `env.` prefix should return the environment node_id", - input_node_id=_node_id, - input_name="env.API_KEY", - expected_node_id=ENVIRONMENT_VARIABLE_NODE_ID, - expected_name="API_KEY", - ), - TestCase( - name="name with `conversation.` prefix should return the conversation node_id", - input_node_id=_node_id, - input_name="conversation.session_id", - expected_node_id=CONVERSATION_VARIABLE_NODE_ID, - expected_name="session_id", - ), - TestCase( - name="dummy_variable should return the original input node_id", - input_node_id=_node_id, - input_name="__dummy__", - expected_node_id=_node_id, - expected_name="__dummy__", - ), - ] - - mock_user = Account(name="Test Account", email="test@example.com") - mock_user.id = str(uuid.uuid4()) - test_app_id = self._get_test_app_id() - saver = DraftVariableSaver( - session=sqlite_session, - tenant_id="test-tenant-id", - app_id=test_app_id, - node_id=_node_id, - node_type=BuiltinNodeTypes.START, - node_execution_id="test_execution_id", - user=mock_user, - ) - for idx, c in enumerate(cases, 1): - fail_msg = f"Test case {c.name} failed, index={idx}" - node_id, name = saver._normalize_variable_for_start_node(c.input_name) - assert node_id == c.expected_node_id, fail_msg - assert name == c.expected_name, fail_msg - - def test_build_variables_from_start_mapping_rebuilds_system_files(self, sqlite_session: Session): - mock_user = Account(name="Test Account", email="test@example.com") - mock_user.id = str(uuid.uuid4()) - saver = DraftVariableSaver( - session=sqlite_session, - tenant_id="tenant-1", - app_id=self._get_test_app_id(), - node_id="start", - node_type=BuiltinNodeTypes.START, - node_execution_id="exec-1", - user=mock_user, - ) - rebuilt_file = File( - file_id="file-1", - file_type=FileType.DOCUMENT, - transfer_method=FileTransferMethod.LOCAL_FILE, - reference="upload-1", - filename="test.txt", - extension=".txt", - mime_type="text/plain", - size=12, - storage_key="canonical-storage-key", - ) - raw_file = { - **rebuilt_file.model_dump(mode="json"), - "tenant_id": "legacy-tenant", - } - - with patch( - "services.workflow_draft_variable_service.build_file_from_stored_mapping", - return_value=rebuilt_file, - ) as rebuild_file: - draft_vars = saver._build_variables_from_start_mapping({"sys.files": [raw_file]}) - - sys_var = draft_vars[0] - assert sys_var.get_value().value[0] == rebuilt_file - rebuild_file.assert_called_once_with(file_mapping=raw_file, tenant_id="tenant-1") - - @pytest.fixture - def draft_saver(self, sqlite_session: Session): - """Create DraftVariableSaver instance with user context.""" - # Create a mock user - mock_user = Account(name="Test Account", email="test@example.com") - mock_user.id = "test-user-id" - - return DraftVariableSaver( - session=sqlite_session, - tenant_id="test-tenant-id", - app_id="test-app-id", - node_id="test-node-id", - node_type=BuiltinNodeTypes.LLM, - node_execution_id="test-execution-id", - user=mock_user, - ) - - def test_draft_saver_with_small_variables(self, draft_saver: DraftVariableSaver): - with patch( - "services.workflow_draft_variable_service.DraftVariableSaver._try_offload_large_variable", autospec=True - ) as _mock_try_offload: - _mock_try_offload.return_value = None - mock_segment = StringSegment(value="small value") - draft_var = draft_saver._create_draft_variable(name="small_var", value=mock_segment, visible=True) - - # Should not have large variable metadata - assert draft_var.file_id is None - _mock_try_offload.return_value = None - - def test_draft_saver_with_large_variables(self, draft_saver: DraftVariableSaver): - with patch( - "services.workflow_draft_variable_service.DraftVariableSaver._try_offload_large_variable", autospec=True - ) as _mock_try_offload: - mock_segment = StringSegment(value="small value") - mock_draft_var_file = WorkflowDraftVariableFile( - tenant_id=str(uuidv7()), - app_id=str(uuidv7()), - user_id=str(uuidv7()), - size=1024, - length=10, - value_type=SegmentType.ARRAY_STRING, - upload_file_id=str(uuidv7()), - ) - mock_draft_var_file.id = str(uuidv7()) - - _mock_try_offload.return_value = mock_segment, mock_draft_var_file - draft_var = draft_saver._create_draft_variable(name="small_var", value=mock_segment, visible=True) - - # Should not have large variable metadata - assert draft_var.file_id == mock_draft_var_file.id - - def test_try_offload_large_variable_uses_resource_tenant(self, sqlite_session: Session): - mock_user = Account(name="Test Account", email="test@example.com") - mock_user.id = "test-user-id" - saver = DraftVariableSaver( - session=sqlite_session, - tenant_id="app-tenant-id", - app_id="test-app-id", - node_id="test-node-id", - node_type=BuiltinNodeTypes.LLM, - node_execution_id="test-execution-id", - user=mock_user, - ) - upload_file = UploadFile( - tenant_id="app-tenant-id", - storage_type=StorageType.LOCAL, - key="workflow/draft-variable.txt", - name="draft-variable.txt", - size=11, - extension="txt", - mime_type="text/plain", - created_by_role=CreatorUserRole.ACCOUNT, - created_by=mock_user.id, - created_at=naive_utc_now(), - used=True, - ) - sqlite_session.add(upload_file) - sqlite_session.commit() - truncation_result = TruncationResult(result=StringSegment(value="..."), truncated=True) - - with ( - patch( - "services.workflow_draft_variable_service.VariableTruncator.truncate", return_value=truncation_result - ), - patch("services.workflow_draft_variable_service.FileService") as file_service_class, - ): - file_service_class.return_value.upload_file.return_value = upload_file - result = saver._try_offload_large_variable("large_var", StringSegment(value="large value")) - - assert result is not None - _, variable_file = result - assert file_service_class.return_value.upload_file.call_args.kwargs["tenant_id"] == "app-tenant-id" - assert variable_file.tenant_id == "app-tenant-id" - sqlite_session.expire_all() - stored_variable_file = sqlite_session.get(WorkflowDraftVariableFile, variable_file.id) - assert stored_variable_file is not None - assert stored_variable_file.upload_file_id == upload_file.id - - @patch("services.workflow_draft_variable_service._batch_upsert_draft_variable", autospec=True) - def test_save_method_integration(self, mock_batch_upsert, draft_saver): - """Test complete save workflow.""" - outputs = {"result": {"data": "test_output"}, "metadata": {"type": "llm_response"}} - - draft_saver.save(outputs=outputs) - - # Should batch upsert draft variables - mock_batch_upsert.assert_called_once() - draft_vars = mock_batch_upsert.call_args[0][1] - assert len(draft_vars) == 2 - - @patch("services.workflow_draft_variable_service._batch_upsert_draft_variable", autospec=True) - def test_start_node_save_persists_sys_timestamp_and_workflow_run_id( - self, mock_batch_upsert, sqlite_session: Session - ): - """Start node should persist common `sys.*` variables, not only `sys.files`.""" - mock_user = Account(name="Test Account", email="test@example.com") - mock_user.id = "test-user-id" - mock_user.tenant_id = "test-tenant-id" - - saver = DraftVariableSaver( - session=sqlite_session, - tenant_id="test-tenant-id", - app_id="test-app-id", - node_id="start-node-id", - node_type=BuiltinNodeTypes.START, - node_execution_id="exec-id", - user=mock_user, - ) - - outputs = { - f"{SYSTEM_VARIABLE_NODE_ID}.{SystemVariableKey.TIMESTAMP}": 1700000000, - f"{SYSTEM_VARIABLE_NODE_ID}.{SystemVariableKey.WORKFLOW_EXECUTION_ID}": "run-id-123", - } - - saver.save(outputs=outputs) - - mock_batch_upsert.assert_called_once() - draft_vars = mock_batch_upsert.call_args[0][1] - - # plus one dummy output because there are no non-sys Start inputs - assert len(draft_vars) == 3 - - sys_vars = [v for v in draft_vars if v.node_id == SYSTEM_VARIABLE_NODE_ID] - assert {v.name for v in sys_vars} == { - str(SystemVariableKey.TIMESTAMP), - str(SystemVariableKey.WORKFLOW_EXECUTION_ID), - } - - @patch("services.workflow_draft_variable_service._batch_upsert_draft_variable", autospec=True) - def test_start_node_save_normalizes_reserved_prefix_outputs(self, mock_batch_upsert, sqlite_session: Session): - mock_user = Account(name="Test Account", email="test@example.com") - mock_user.id = "test-user-id" - mock_user.tenant_id = "test-tenant-id" - - saver = DraftVariableSaver( - session=sqlite_session, - tenant_id="test-tenant-id", - app_id="test-app-id", - node_id="start-node-id", - node_type=BuiltinNodeTypes.START, - node_execution_id="exec-id", - user=mock_user, - ) - - saver.save( - outputs={ - "env.API_KEY": "secret", - "conversation.session_id": "conversation-1", - "sys.workflow_run_id": "run-id-123", - } - ) - - mock_batch_upsert.assert_called_once() - draft_vars = mock_batch_upsert.call_args[0][1] - - assert len(draft_vars) == 3 - - env_var = next(v for v in draft_vars if v.node_id == ENVIRONMENT_VARIABLE_NODE_ID) - assert env_var.name == "API_KEY" - assert env_var.editable is False - - conversation_var = next(v for v in draft_vars if v.node_id == CONVERSATION_VARIABLE_NODE_ID) - assert conversation_var.name == "session_id" - assert conversation_var.node_execution_id is None - - sys_var = next(v for v in draft_vars if v.node_id == SYSTEM_VARIABLE_NODE_ID) - assert sys_var.name == str(SystemVariableKey.WORKFLOW_EXECUTION_ID) - - -class TestWorkflowDraftVariableService: - def _get_test_app_id(self): - suffix = secrets.token_hex(6) - return f"test_app_id_{suffix}" - - def _create_test_workflow(self, app_id: str) -> Workflow: - """Create a real Workflow instance for testing""" - return Workflow.new( - tenant_id="test_tenant_id", - app_id=app_id, - type="workflow", - version="draft", - graph='{"nodes": [], "edges": []}', - features="{}", - created_by="test_user_id", - environment_variables=[], - conversation_variables=[], - rag_pipeline_variables=[], - ) - - def test_list_variables_without_values_excludes_node_ids(self, sqlite_session: Session): - service = WorkflowDraftVariableService(sqlite_session) - variable = WorkflowDraftVariable.new_node_variable( - app_id="app-1", - node_id="node-1", - name="output", - value=StringSegment(value="value"), - node_execution_id="execution-1", - ) - variable.user_id = "user-1" - excluded_system = WorkflowDraftVariable.new_sys_variable( - app_id="app-1", name="query", value=StringSegment(value="hidden"), node_execution_id="execution-1" - ) - excluded_system.user_id = "user-1" - other_user = WorkflowDraftVariable.new_node_variable( - app_id="app-1", - node_id="node-2", - name="output", - value=StringSegment(value="other"), - node_execution_id="execution-2", - ) - other_user.user_id = "user-2" - sqlite_session.add_all([variable, excluded_system, other_user]) - sqlite_session.commit() - - result = service.list_variables_without_values( - app_id="app-1", - page=1, - limit=20, - user_id="user-1", - exclude_node_ids={SYSTEM_VARIABLE_NODE_ID, CONVERSATION_VARIABLE_NODE_ID}, - ) - - assert result.total == 1 - assert [item.id for item in result.variables] == [variable.id] - - def test_reset_conversation_variable(self, sqlite_session: Session): - """Test resetting a conversation variable""" - service = WorkflowDraftVariableService(sqlite_session) - - test_app_id = self._get_test_app_id() - workflow = self._create_test_workflow(test_app_id) - - # Create real conversation variable - test_value = StringSegment(value="test_value") - variable = WorkflowDraftVariable.new_conversation_variable( - app_id=test_app_id, name="test_var", value=test_value, description="Test conversation variable" - ) - - # Mock the _reset_conv_var method - expected_result = WorkflowDraftVariable.new_conversation_variable( - app_id=test_app_id, - name="test_var", - value=StringSegment(value="reset_value"), - ) - with patch.object(service, "_reset_conv_var", return_value=expected_result, autospec=True) as mock_reset_conv: - result = service.reset_variable(workflow, variable) - - mock_reset_conv.assert_called_once_with(workflow, variable) - assert result == expected_result - - def test_reset_node_variable_with_no_execution_id(self, sqlite_session: Session): - """Test resetting a node variable with no execution ID - should delete variable""" - service = WorkflowDraftVariableService(sqlite_session) - - test_app_id = self._get_test_app_id() - workflow = self._create_test_workflow(test_app_id) - - # Create real node variable with no execution ID - test_value = StringSegment(value="test_value") - variable = WorkflowDraftVariable.new_node_variable( - app_id=test_app_id, - node_id="test_node_id", - name="test_var", - value=test_value, - node_execution_id="exec-id", # Set initially - ) - # Manually set to None to simulate the test condition - variable.node_execution_id = None - sqlite_session.add(variable) - sqlite_session.commit() - - result = service._reset_node_var_or_sys_var(workflow, variable) - - # Should delete the variable and return None - assert sqlite_session.get(WorkflowDraftVariable, variable.id) is None - assert result is None - - def test_reset_node_variable_with_missing_execution_record(self, sqlite_session: Session): - """Test resetting a node variable when execution record doesn't exist""" - service = WorkflowDraftVariableService(sqlite_session) - - test_app_id = self._get_test_app_id() - workflow = self._create_test_workflow(test_app_id) - - # Create real node variable with execution ID - test_value = StringSegment(value="test_value") - variable = WorkflowDraftVariable.new_node_variable( - app_id=test_app_id, node_id="test_node_id", name="test_var", value=test_value, node_execution_id="exec-id" - ) - sqlite_session.add(variable) - sqlite_session.commit() - - result = service._reset_node_var_or_sys_var(workflow, variable) - - assert sqlite_session.get(WorkflowDraftVariable, variable.id) is None - assert result is None - - def test_reset_node_variable_with_valid_execution_record(self, sqlite_session: Session): - """Reset a node variable from its execution output and flush the restored value.""" - service = WorkflowDraftVariableService(sqlite_session) - - execution = WorkflowNodeExecutionModel(outputs=json.dumps({"test_var": "output_value"})) - - # Mock the repository to return the execution record - service._api_node_execution_repo = Mock() - service._api_node_execution_repo.get_execution_by_id.return_value = execution - - test_app_id = self._get_test_app_id() - workflow = self._create_test_workflow(test_app_id) - - # Create real node variable with execution ID - test_value = StringSegment(value="original_value") - variable = WorkflowDraftVariable.new_node_variable( - app_id=test_app_id, node_id="test_node_id", name="test_var", value=test_value, node_execution_id="exec-id" - ) - sqlite_session.add(variable) - sqlite_session.commit() - - # Mock workflow methods - mock_node_config = {"type": "test_node"} - with ( - patch.object(sqlite_session, "flush", wraps=sqlite_session.flush) as flush, - patch.object(workflow, "get_node_config_by_id", return_value=mock_node_config, autospec=True), - patch.object(workflow, "get_node_type_from_node_config", return_value=BuiltinNodeTypes.LLM, autospec=True), - ): - result = service._reset_node_var_or_sys_var(workflow, variable) - - # Verify last_edited_at was reset - assert variable.last_edited_at is None - flush.assert_called() - # Should return the updated variable - assert result == variable - - def test_reset_non_editable_system_variable_raises_error(self, sqlite_session: Session): - """Test that resetting a non-editable system variable raises an error""" - service = WorkflowDraftVariableService(sqlite_session) - - test_app_id = self._get_test_app_id() - workflow = self._create_test_workflow(test_app_id) - - # Create a non-editable system variable (workflow_id is not editable) - test_value = StringSegment(value="test_workflow_id") - variable = WorkflowDraftVariable.new_sys_variable( - app_id=test_app_id, - name="workflow_id", # This is not in _EDITABLE_SYSTEM_VARIABLE - value=test_value, - node_execution_id="exec-id", - editable=False, # Non-editable system variable - ) - - with pytest.raises(VariableResetError) as exc_info: - service.reset_variable(workflow, variable) - assert "cannot reset system variable" in str(exc_info.value) - assert f"variable_id={variable.id}" in str(exc_info.value) - - def test_reset_editable_system_variable_succeeds(self, sqlite_session: Session): - """Test that resetting an editable system variable succeeds""" - service = WorkflowDraftVariableService(sqlite_session) - - test_app_id = self._get_test_app_id() - workflow = self._create_test_workflow(test_app_id) - - # Create an editable system variable (files is editable) - test_value = StringSegment(value="[]") - variable = WorkflowDraftVariable.new_sys_variable( - app_id=test_app_id, - name="files", # This is in _EDITABLE_SYSTEM_VARIABLE - value=test_value, - node_execution_id="exec-id", - editable=True, # Editable system variable - ) - sqlite_session.add(variable) - sqlite_session.commit() - - execution = WorkflowNodeExecutionModel(outputs=json.dumps({"sys.files": "[]"})) - - # Mock the repository to return the execution record - service._api_node_execution_repo = Mock() - service._api_node_execution_repo.get_execution_by_id.return_value = execution - - with patch.object(sqlite_session, "flush", wraps=sqlite_session.flush) as flush: - result = service._reset_node_var_or_sys_var(workflow, variable) - - # Should succeed and return the variable - assert result == variable - assert variable.last_edited_at is None - flush.assert_called() - - def test_reset_query_system_variable_succeeds(self, sqlite_session: Session): - """Test that resetting query system variable (another editable one) succeeds""" - service = WorkflowDraftVariableService(sqlite_session) - - test_app_id = self._get_test_app_id() - workflow = self._create_test_workflow(test_app_id) - - # Create an editable system variable (query is editable) - test_value = StringSegment(value="original query") - variable = WorkflowDraftVariable.new_sys_variable( - app_id=test_app_id, - name="query", # This is in _EDITABLE_SYSTEM_VARIABLE - value=test_value, - node_execution_id="exec-id", - editable=True, # Editable system variable - ) - sqlite_session.add(variable) - sqlite_session.commit() - - execution = WorkflowNodeExecutionModel(outputs=json.dumps({"sys.query": "reset query"})) - - # Mock the repository to return the execution record - service._api_node_execution_repo = Mock() - service._api_node_execution_repo.get_execution_by_id.return_value = execution - - with patch.object(sqlite_session, "flush", wraps=sqlite_session.flush) as flush: - result = service._reset_node_var_or_sys_var(workflow, variable) - - # Should succeed and return the variable - assert result == variable - assert variable.last_edited_at is None - flush.assert_called() - - def test_system_variable_editability_check(self): - """Test the system variable editability function directly""" - # Test editable system variables - assert is_system_variable_editable("files") is True - assert is_system_variable_editable("query") is True - - # Test non-editable system variables - assert is_system_variable_editable("workflow_id") is False - assert is_system_variable_editable("conversation_id") is False - assert is_system_variable_editable("user_id") is False - - def test_workflow_draft_variable_factory_methods(self): - """Test that factory methods create proper instances""" - test_app_id = self._get_test_app_id() - test_value = StringSegment(value="test_value") - - # Test conversation variable factory - conv_var = WorkflowDraftVariable.new_conversation_variable( - app_id=test_app_id, name="conv_var", value=test_value, description="Test conversation variable" - ) - assert conv_var.get_variable_type() == DraftVariableType.CONVERSATION - assert conv_var.editable is True - assert conv_var.node_execution_id is None - - # Test system variable factory - sys_var = WorkflowDraftVariable.new_sys_variable( - app_id=test_app_id, name="workflow_id", value=test_value, node_execution_id="exec-id", editable=False - ) - assert sys_var.get_variable_type() == DraftVariableType.SYS - assert sys_var.editable is False - assert sys_var.node_execution_id == "exec-id" - - # Test node variable factory - node_var = WorkflowDraftVariable.new_node_variable( - app_id=test_app_id, - node_id="node-id", - name="node_var", - value=test_value, - node_execution_id="exec-id", - visible=True, - editable=True, - ) - assert node_var.get_variable_type() == DraftVariableType.NODE - assert node_var.visible is True - assert node_var.editable is True - assert node_var.node_execution_id == "exec-id" - - -class TestModelToInsertionDict: - """Reproduce two production errors in _model_to_insertion_dict / _new().""" - - def test_visible_and_is_default_value_always_present(self): - """Problem 1: _new() did not set visible/is_default_value, causing - inconsistent dict keys across rows in multi-row INSERT and missing - is_default_value in the insertion dict entirely. - """ - conv_var = WorkflowDraftVariable.new_conversation_variable( - app_id="app-1", - name="counter", - value=StringSegment(value="0"), - ) - # _new() should explicitly set these fields so they are not None - assert conv_var.visible is not None - assert conv_var.is_default_value is not None - - d = _model_to_insertion_dict(conv_var) - # visible must appear in every row's dict - assert "visible" in d - # is_default_value must always be present - assert "is_default_value" in d - - def test_description_passthrough(self): - """_model_to_insertion_dict passes description as-is; - length validation is enforced earlier in build_conversation_variable_from_mapping. - """ - desc = "a" * 200 - conv_var = WorkflowDraftVariable.new_conversation_variable( - app_id="app-1", - name="counter", - value=StringSegment(value="0"), - description=desc, - ) - d = _model_to_insertion_dict(conv_var) - assert d["description"] == desc - - def test_is_default_value_omitted_when_none(self): - conv_var = WorkflowDraftVariable.new_conversation_variable( - app_id="app-1", - name="counter", - value=StringSegment(value="0"), - ) - conv_var.is_default_value = None - - d = _model_to_insertion_dict(conv_var) - - assert "is_default_value" not in d diff --git a/api/tests/unit_tests/services/workflow/test_workflow_draft_variable_service_sqlite.py b/api/tests/unit_tests/services/workflow/test_workflow_draft_variable_service_sqlite.py deleted file mode 100644 index 2f66d12e6e2..00000000000 --- a/api/tests/unit_tests/services/workflow/test_workflow_draft_variable_service_sqlite.py +++ /dev/null @@ -1,759 +0,0 @@ -"""SQLite-backed workflow draft-variable service coverage.""" - -import json -import unittest -import uuid -from pathlib import Path -from typing import override - -import pytest -from sqlalchemy import Engine, delete, func, select -from sqlalchemy.orm import Session - -from core.workflow.variable_prefixes import CONVERSATION_VARIABLE_NODE_ID, SYSTEM_VARIABLE_NODE_ID -from extensions.ext_storage import storage -from extensions.storage.opendal_storage import OpenDALStorage -from extensions.storage.storage_type import StorageType -from factories.variable_factory import build_segment -from graphon.nodes import BuiltinNodeTypes -from graphon.variables.segments import StringSegment -from graphon.variables.types import SegmentType -from graphon.variables.variables import StringVariable -from libs import datetime_utils -from models import Account, Tenant, TenantAccountJoin -from models.account import TenantAccountRole -from models.enums import CreatorUserRole -from models.model import UploadFile -from models.workflow import Workflow, WorkflowDraftVariable, WorkflowDraftVariableFile, WorkflowNodeExecutionModel -from services.workflow_draft_variable_service import ( - DraftVariableSaver, - DraftVarLoader, - VariableResetError, - WorkflowDraftVariableService, -) - - -@pytest.fixture(autouse=True) -def _bind_storage( - monkeypatch: pytest.MonkeyPatch, - tmp_path: Path, -) -> None: - monkeypatch.setattr( - storage, - "storage_runner", - OpenDALStorage(scheme="fs", root=str(tmp_path / "storage")), - raising=False, - ) - - -def _persist_account(*, engine: Engine, tenant_id: str) -> Account: - tenant = Tenant(name="draft-variable-tenant") - tenant.id = tenant_id - account = Account(email=f"draft-{uuid.uuid4()}@example.com", name="Draft Variable User") - membership = TenantAccountJoin(tenant_id=tenant.id, account_id=account.id, role=TenantAccountRole.OWNER) - with Session(bind=engine, expire_on_commit=False) as session: - session.add_all([tenant, account, membership]) - session.commit() - return account - - -class TestWorkflowDraftVariableService(unittest.TestCase): - _test_app_id: str - _session: Session - _test_user_id: str - _node1_id = "test_node_1" - _node2_id = "test_node_2" - _node_exec_id = str(uuid.uuid4()) - - @pytest.fixture(autouse=True) - def _bind_session(self, sqlite_session: Session) -> None: - self._session = sqlite_session - - @override - def setUp(self) -> None: - self._test_app_id = str(uuid.uuid4()) - self._test_user_id = str(uuid.uuid4()) - sys_var = WorkflowDraftVariable.new_sys_variable( - app_id=self._test_app_id, - user_id=self._test_user_id, - name="sys_var", - value=build_segment("sys_value"), - node_execution_id=self._node_exec_id, - ) - conv_var = WorkflowDraftVariable.new_conversation_variable( - app_id=self._test_app_id, - user_id=self._test_user_id, - name="conv_var", - value=build_segment("conv_value"), - ) - node2_vars = [ - WorkflowDraftVariable.new_node_variable( - app_id=self._test_app_id, - user_id=self._test_user_id, - node_id=self._node2_id, - name="int_var", - value=build_segment(1), - visible=False, - node_execution_id=self._node_exec_id, - ), - WorkflowDraftVariable.new_node_variable( - app_id=self._test_app_id, - user_id=self._test_user_id, - node_id=self._node2_id, - name="str_var", - value=build_segment("str_value"), - visible=True, - node_execution_id=self._node_exec_id, - ), - ] - node1_var = WorkflowDraftVariable.new_node_variable( - app_id=self._test_app_id, - user_id=self._test_user_id, - node_id=self._node1_id, - name="str_var", - value=build_segment("str_value"), - visible=True, - node_execution_id=self._node_exec_id, - ) - _variables = list(node2_vars) - _variables.extend( - [ - node1_var, - sys_var, - conv_var, - ] - ) - - self._session.add_all(_variables) - self._session.flush() - self._variable_ids = [v.id for v in _variables] - self._node1_str_var_id = node1_var.id - self._sys_var_id = sys_var.id - self._conv_var_id = conv_var.id - self._node2_var_ids = [v.id for v in node2_vars] - - def _get_test_srv(self) -> WorkflowDraftVariableService: - return WorkflowDraftVariableService(session=self._session) - - @override - def tearDown(self) -> None: - self._session.rollback() - - def test_list_variables(self) -> None: - srv = self._get_test_srv() - var_list = srv.list_variables_without_values(self._test_app_id, page=1, limit=2, user_id=self._test_user_id) - assert var_list.total == 5 - assert len(var_list.variables) == 2 - page1_var_ids = {v.id for v in var_list.variables} - assert page1_var_ids.issubset(self._variable_ids) - - var_list_2 = srv.list_variables_without_values(self._test_app_id, page=2, limit=2, user_id=self._test_user_id) - assert var_list_2.total is None - assert len(var_list_2.variables) == 2 - page2_var_ids = {v.id for v in var_list_2.variables} - assert page2_var_ids.isdisjoint(page1_var_ids) - assert page2_var_ids.issubset(self._variable_ids) - - def test_get_node_variable(self) -> None: - srv = self._get_test_srv() - node_var = srv.get_node_variable(self._test_app_id, self._node1_id, "str_var", user_id=self._test_user_id) - assert node_var is not None - assert node_var.id == self._node1_str_var_id - assert node_var.name == "str_var" - assert node_var.get_value() == build_segment("str_value") - - def test_get_system_variable(self) -> None: - srv = self._get_test_srv() - sys_var = srv.get_system_variable(self._test_app_id, "sys_var", user_id=self._test_user_id) - assert sys_var is not None - assert sys_var.id == self._sys_var_id - assert sys_var.name == "sys_var" - assert sys_var.get_value() == build_segment("sys_value") - - def test_get_conversation_variable(self) -> None: - srv = self._get_test_srv() - conv_var = srv.get_conversation_variable(self._test_app_id, "conv_var", user_id=self._test_user_id) - assert conv_var is not None - assert conv_var.id == self._conv_var_id - assert conv_var.name == "conv_var" - assert conv_var.get_value() == build_segment("conv_value") - - def test_delete_node_variables(self) -> None: - srv = self._get_test_srv() - srv.delete_node_variables(self._test_app_id, self._node2_id, user_id=self._test_user_id) - node2_var_count = self._session.scalar( - select(func.count()) - .select_from(WorkflowDraftVariable) - .where( - WorkflowDraftVariable.app_id == self._test_app_id, - WorkflowDraftVariable.node_id == self._node2_id, - WorkflowDraftVariable.user_id == self._test_user_id, - ) - ) - assert node2_var_count == 0 - - def test_delete_variable(self) -> None: - srv = self._get_test_srv() - node_1_var = self._session.scalars( - select(WorkflowDraftVariable).where(WorkflowDraftVariable.id == self._node1_str_var_id) - ).one() - srv.delete_variable(node_1_var) - exists = bool( - self._session.scalars( - select(WorkflowDraftVariable).where(WorkflowDraftVariable.id == self._node1_str_var_id) - ).first() - ) - assert exists is False - - def test__list_node_variables(self) -> None: - srv = self._get_test_srv() - node_vars = srv._list_node_variables(self._test_app_id, self._node2_id, user_id=self._test_user_id) - assert len(node_vars.variables) == 2 - assert {v.id for v in node_vars.variables} == set(self._node2_var_ids) - - def test_get_draft_variables_by_selectors(self) -> None: - srv = self._get_test_srv() - selectors = [ - [self._node1_id, "str_var"], - [self._node2_id, "str_var"], - [self._node2_id, "int_var"], - ] - variables = srv.get_draft_variables_by_selectors(self._test_app_id, selectors, user_id=self._test_user_id) - assert len(variables) == 3 - assert {v.id for v in variables} == {self._node1_str_var_id} | set(self._node2_var_ids) - - -class TestDraftVariableLoader(unittest.TestCase): - _test_app_id: str - _test_tenant_id: str - _engine: Engine - - _node1_id = "test_loader_node_1" - _node_exec_id = str(uuid.uuid4()) - - @pytest.fixture(autouse=True) - def _bind_engine(self, sqlite_engine: Engine) -> None: - self._engine = sqlite_engine - - @override - def setUp(self) -> None: - self._test_app_id = str(uuid.uuid4()) - self._test_tenant_id = str(uuid.uuid4()) - self._test_user_id = str(uuid.uuid4()) - sys_var = WorkflowDraftVariable.new_sys_variable( - app_id=self._test_app_id, - user_id=self._test_user_id, - name="sys_var", - value=build_segment("sys_value"), - node_execution_id=self._node_exec_id, - ) - conv_var = WorkflowDraftVariable.new_conversation_variable( - app_id=self._test_app_id, - user_id=self._test_user_id, - name="conv_var", - value=build_segment("conv_value"), - ) - node_var = WorkflowDraftVariable.new_node_variable( - app_id=self._test_app_id, - user_id=self._test_user_id, - node_id=self._node1_id, - name="str_var", - value=build_segment("str_value"), - visible=True, - node_execution_id=self._node_exec_id, - ) - _variables = [ - node_var, - sys_var, - conv_var, - ] - - with Session(bind=self._engine, expire_on_commit=False) as session: - session.add_all(_variables) - session.flush() - session.commit() - self._variable_ids = [v.id for v in _variables] - self._node_var_id = node_var.id - self._sys_var_id = sys_var.id - self._conv_var_id = conv_var.id - - @override - def tearDown(self) -> None: - with Session(bind=self._engine, expire_on_commit=False) as session: - session.execute(delete(WorkflowDraftVariable).where(WorkflowDraftVariable.app_id == self._test_app_id)) - session.commit() - - def test_variable_loader_with_empty_selector(self) -> None: - var_loader = DraftVarLoader( - engine=self._engine, - app_id=self._test_app_id, - tenant_id=self._test_tenant_id, - user_id=self._test_user_id, - ) - variables = var_loader.load_variables([]) - assert len(variables) == 0 - - def test_variable_loader_with_non_empty_selector(self) -> None: - var_loader = DraftVarLoader( - engine=self._engine, - app_id=self._test_app_id, - tenant_id=self._test_tenant_id, - user_id=self._test_user_id, - ) - variables = var_loader.load_variables( - [ - [SYSTEM_VARIABLE_NODE_ID, "sys_var"], - [CONVERSATION_VARIABLE_NODE_ID, "conv_var"], - [self._node1_id, "str_var"], - ] - ) - assert len(variables) == 3 - conv_var = next(v for v in variables if v.selector[0] == CONVERSATION_VARIABLE_NODE_ID) - assert conv_var.id == self._conv_var_id - sys_var = next(v for v in variables if v.selector[0] == SYSTEM_VARIABLE_NODE_ID) - assert sys_var.id == self._sys_var_id - node1_var = next(v for v in variables if v.selector[0] == self._node1_id) - assert node1_var.id == self._node_var_id - - def test_load_offloaded_variable_string_type_integration(self) -> None: - """Test _load_offloaded_variable with string type using DraftVariableSaver for data creation.""" - account = _persist_account(engine=self._engine, tenant_id=self._test_tenant_id) - - # Create a large string that will be offloaded - test_content = "x" * 15000 # Create a string larger than LARGE_VARIABLE_THRESHOLD (10KB) - large_string_segment = StringSegment(value=test_content) - - node_execution_id = str(uuid.uuid4()) - - try: - with Session(bind=self._engine, expire_on_commit=False) as session: - # Use DraftVariableSaver to create offloaded variable (this mimics production) - saver = DraftVariableSaver( - session=session, - tenant_id=self._test_tenant_id, - app_id=self._test_app_id, - node_id="test_offload_node", - node_type=BuiltinNodeTypes.LLM, # Use a real node type - node_execution_id=node_execution_id, - user=account, - ) - - # Save the variable - this will trigger offloading due to large size - saver.save(outputs={"offloaded_string_var": large_string_segment}) - session.commit() - - # Now test loading using DraftVarLoader - var_loader = DraftVarLoader( - engine=self._engine, - app_id=self._test_app_id, - tenant_id=self._test_tenant_id, - user_id=account.id, - ) - - # Load the variable using the standard workflow - variables = var_loader.load_variables([["test_offload_node", "offloaded_string_var"]]) - - # Verify results - assert len(variables) == 1 - loaded_variable = variables[0] - assert loaded_variable.name == "offloaded_string_var" - assert loaded_variable.selector == ["test_offload_node", "offloaded_string_var"] - assert loaded_variable.value == test_content - - finally: - # Clean up - delete all draft variables for this app - with Session(bind=self._engine) as session: - service = WorkflowDraftVariableService(session) - service.delete_app_workflow_variables(self._test_app_id) - session.commit() - - def test_load_offloaded_variable_object_type_integration(self) -> None: - """Test _load_offloaded_variable with object type using real storage and service.""" - - # Create a test object - test_object = {"key1": "value1", "key2": 42, "nested": {"inner": "data"}} - test_json = json.dumps(test_object, ensure_ascii=False, separators=(",", ":")) - content_bytes = test_json.encode() - - # Create an upload file record - upload_file = UploadFile( - tenant_id=self._test_tenant_id, - storage_type=StorageType.LOCAL, - key=f"test_offload_{uuid.uuid4()}.json", - name="test_offload.json", - size=len(content_bytes), - extension="json", - mime_type="application/json", - created_by_role=CreatorUserRole.ACCOUNT, - created_by=str(uuid.uuid4()), - created_at=datetime_utils.naive_utc_now(), - used=True, - used_by=str(uuid.uuid4()), - used_at=datetime_utils.naive_utc_now(), - ) - - # Store the content in storage - storage.save(upload_file.key, content_bytes) - - # Create a variable file record - variable_file = WorkflowDraftVariableFile( - upload_file_id=upload_file.id, - value_type=SegmentType.OBJECT, - tenant_id=self._test_tenant_id, - app_id=self._test_app_id, - user_id=str(uuid.uuid4()), - size=len(content_bytes), - length=len(test_object), - created_at=datetime_utils.naive_utc_now(), - ) - - try: - with Session(bind=self._engine, expire_on_commit=False) as session: - # Add upload file and variable file first to get their IDs - session.add_all([upload_file, variable_file]) - session.flush() # This generates the IDs - - # Now create the offloaded draft variable with the correct file_id - offloaded_var = WorkflowDraftVariable.new_node_variable( - app_id=self._test_app_id, - user_id=self._test_user_id, - node_id="test_offload_node", - name="offloaded_object_var", - value=build_segment({"truncated": True}), - visible=True, - node_execution_id=str(uuid.uuid4()), - ) - offloaded_var.file_id = variable_file.id - - session.add(offloaded_var) - session.flush() - session.commit() - - # Use the service method that properly preloads relationships - service = WorkflowDraftVariableService(session) - draft_vars = service.get_draft_variables_by_selectors( - self._test_app_id, - [["test_offload_node", "offloaded_object_var"]], - user_id=self._test_user_id, - ) - - assert len(draft_vars) == 1 - loaded_var = draft_vars[0] - assert loaded_var.is_truncated() - - # Create DraftVarLoader and test loading - var_loader = DraftVarLoader( - engine=self._engine, - app_id=self._test_app_id, - tenant_id=self._test_tenant_id, - user_id=self._test_user_id, - ) - - # Test the _load_offloaded_variable method - selector_tuple, variable = var_loader._load_offloaded_variable(loaded_var) - - # Verify the results - assert selector_tuple == ("test_offload_node", "offloaded_object_var") - assert variable.id == loaded_var.id - assert variable.name == "offloaded_object_var" - assert variable.value == test_object - - finally: - # Clean up - with Session(bind=self._engine) as session: - # Query and delete by ID to ensure they're tracked in this session - session.execute(delete(WorkflowDraftVariable).where(WorkflowDraftVariable.id == offloaded_var.id)) - session.execute( - delete(WorkflowDraftVariableFile).where(WorkflowDraftVariableFile.id == variable_file.id) - ) - session.execute(delete(UploadFile).where(UploadFile.id == upload_file.id)) - session.commit() - storage.delete(upload_file.key) - - def test_load_variables_with_offloaded_variables_integration(self) -> None: - """Test load_variables method with mix of regular and offloaded variables using real storage.""" - # Create a regular variable (already exists from setUp) - # Create offloaded variable content - test_content = "This is offloaded content for integration test" - content_bytes = test_content.encode() - - # Create upload file record - upload_file = UploadFile( - tenant_id=self._test_tenant_id, - storage_type=StorageType.LOCAL, - key=f"test_integration_{uuid.uuid4()}.txt", - name="test_integration.txt", - size=len(content_bytes), - extension="txt", - mime_type="text/plain", - created_by_role=CreatorUserRole.ACCOUNT, - created_by=str(uuid.uuid4()), - created_at=datetime_utils.naive_utc_now(), - used=True, - used_by=str(uuid.uuid4()), - used_at=datetime_utils.naive_utc_now(), - ) - - # Store the content - storage.save(upload_file.key, content_bytes) - - # Create variable file - variable_file = WorkflowDraftVariableFile( - upload_file_id=upload_file.id, - value_type=SegmentType.STRING, - tenant_id=self._test_tenant_id, - app_id=self._test_app_id, - user_id=str(uuid.uuid4()), - size=len(content_bytes), - length=None, - created_at=datetime_utils.naive_utc_now(), - ) - - try: - with Session(bind=self._engine, expire_on_commit=False) as session: - # Add upload file and variable file first to get their IDs - session.add_all([upload_file, variable_file]) - session.flush() # This generates the IDs - - # Now create the offloaded draft variable with the correct file_id - offloaded_var = WorkflowDraftVariable.new_node_variable( - app_id=self._test_app_id, - user_id=self._test_user_id, - node_id="test_integration_node", - name="offloaded_integration_var", - value=build_segment("truncated"), - visible=True, - node_execution_id=str(uuid.uuid4()), - ) - offloaded_var.file_id = variable_file.id - - session.add(offloaded_var) - session.flush() - session.commit() - - # Test load_variables with both regular and offloaded variables - # This method should handle the relationship preloading internally - var_loader = DraftVarLoader( - engine=self._engine, - app_id=self._test_app_id, - tenant_id=self._test_tenant_id, - user_id=self._test_user_id, - ) - - variables = var_loader.load_variables( - [ - [SYSTEM_VARIABLE_NODE_ID, "sys_var"], # Regular variable from setUp - ["test_integration_node", "offloaded_integration_var"], # Offloaded variable - ] - ) - - # Verify results - assert len(variables) == 2 - - # Find regular variable - regular_var = next(v for v in variables if v.selector[0] == SYSTEM_VARIABLE_NODE_ID) - assert regular_var.id == self._sys_var_id - assert regular_var.value == "sys_value" - - # Find offloaded variable - offloaded_loaded_var = next(v for v in variables if v.selector[0] == "test_integration_node") - assert offloaded_loaded_var.id == offloaded_var.id - assert offloaded_loaded_var.value == test_content - - finally: - # Clean up - with Session(bind=self._engine) as session: - # Query and delete by ID to ensure they're tracked in this session - session.execute(delete(WorkflowDraftVariable).where(WorkflowDraftVariable.id == offloaded_var.id)) - session.execute( - delete(WorkflowDraftVariableFile).where(WorkflowDraftVariableFile.id == variable_file.id) - ) - session.execute(delete(UploadFile).where(UploadFile.id == upload_file.id)) - session.commit() - storage.delete(upload_file.key) - - -class TestWorkflowDraftVariableServiceResetVariable(unittest.TestCase): - """Integration tests for reset_variable functionality using real database""" - - _test_app_id: str - _test_tenant_id: str - _test_workflow_id: str - _session: Session - _engine: Engine - _node_id = "test_reset_node" - _node_exec_id: str - _workflow_node_exec_id: str - - @pytest.fixture(autouse=True) - def _bind_sqlite(self, sqlite_session: Session, sqlite_engine: Engine) -> None: - self._session = sqlite_session - self._engine = sqlite_engine - - @override - def setUp(self) -> None: - self._test_app_id = str(uuid.uuid4()) - self._test_tenant_id = str(uuid.uuid4()) - self._test_workflow_id = str(uuid.uuid4()) - self._test_user_id = str(uuid.uuid4()) - self._node_exec_id = str(uuid.uuid4()) - self._workflow_node_exec_id = str(uuid.uuid4()) - - # Create a workflow node execution record with outputs - # Note: The WorkflowNodeExecutionModel.id should match the node_execution_id in WorkflowDraftVariable - self._workflow_node_execution = WorkflowNodeExecutionModel( - id=self._node_exec_id, # This should match the node_execution_id in the variable - tenant_id=self._test_tenant_id, - app_id=self._test_app_id, - workflow_id=self._test_workflow_id, - triggered_from="workflow-run", - workflow_run_id=str(uuid.uuid4()), - index=1, - node_execution_id=str(uuid.uuid4()), - node_id=self._node_id, - node_type=BuiltinNodeTypes.LLM, - title="Test Node", - inputs='{"input": "test input"}', - process_data='{"test_var": "process_value", "other_var": "other_process"}', - outputs='{"test_var": "output_value", "other_var": "other_output"}', - status="succeeded", - elapsed_time=1.5, - created_by_role=CreatorUserRole.ACCOUNT, - created_by=str(uuid.uuid4()), - ) - - # Create conversation variables for the workflow - self._conv_variables = [ - StringVariable( - id=str(uuid.uuid4()), - name="conv_var_1", - description="Test conversation variable 1", - value="default_value_1", - ), - StringVariable( - id=str(uuid.uuid4()), - name="conv_var_2", - description="Test conversation variable 2", - value="default_value_2", - ), - ] - - # Create test variables - self._node_var_with_exec = WorkflowDraftVariable.new_node_variable( - app_id=self._test_app_id, - user_id=self._test_user_id, - node_id=self._node_id, - name="test_var", - value=build_segment("old_value"), - node_execution_id=self._node_exec_id, - ) - self._node_var_with_exec.last_edited_at = datetime_utils.naive_utc_now() - - self._node_var_without_exec = WorkflowDraftVariable.new_node_variable( - app_id=self._test_app_id, - user_id=self._test_user_id, - node_id=self._node_id, - name="no_exec_var", - value=build_segment("some_value"), - node_execution_id="temp_exec_id", - ) - # Manually set node_execution_id to None after creation - self._node_var_without_exec.node_execution_id = None - - self._node_var_missing_exec = WorkflowDraftVariable.new_node_variable( - app_id=self._test_app_id, - user_id=self._test_user_id, - node_id=self._node_id, - name="missing_exec_var", - value=build_segment("some_value"), - node_execution_id=str(uuid.uuid4()), # Use a valid UUID that doesn't exist in database - ) - - self._conv_var = WorkflowDraftVariable.new_conversation_variable( - app_id=self._test_app_id, - user_id=self._test_user_id, - name="conv_var_1", - value=build_segment("old_conv_value"), - ) - self._conv_var.last_edited_at = datetime_utils.naive_utc_now() - - with Session(self._engine, expire_on_commit=False) as persistent_session, persistent_session.begin(): - persistent_session.add( - self._workflow_node_execution, - ) - - # Add all to database - self._session.add_all( - [ - self._node_var_with_exec, - self._node_var_without_exec, - self._node_var_missing_exec, - self._conv_var, - ] - ) - self._session.flush() - - # Store IDs for assertions - self._node_var_with_exec_id = self._node_var_with_exec.id - self._node_var_without_exec_id = self._node_var_without_exec.id - self._node_var_missing_exec_id = self._node_var_missing_exec.id - self._conv_var_id = self._conv_var.id - - @override - def tearDown(self) -> None: - self._session.rollback() - with Session(self._engine) as session, session.begin(): - stmt = delete(WorkflowNodeExecutionModel).where( - WorkflowNodeExecutionModel.id == self._workflow_node_execution.id - ) - session.execute(stmt) - - def _get_test_srv(self) -> WorkflowDraftVariableService: - return WorkflowDraftVariableService(session=self._session) - - def _create_mock_workflow(self) -> Workflow: - """Create a real workflow with conversation variables and graph""" - conversation_vars = self._conv_variables - - # Create a simple graph with the test node - graph: dict[str, object] = { - "nodes": [{"id": "test_reset_node", "type": "llm", "title": "Test Node", "data": {"type": "llm"}}], - "edges": [], - } - - workflow = Workflow.new( - tenant_id=str(uuid.uuid4()), - app_id=self._test_app_id, - type="workflow", - version="1.0", - graph=json.dumps(graph), - features="{}", - created_by=str(uuid.uuid4()), - environment_variables=[], - conversation_variables=conversation_vars, - rag_pipeline_variables=[], - ) - return workflow - - def test_reset_system_variable_raises_error(self) -> None: - """Test that resetting a system variable raises an error""" - srv = self._get_test_srv() - mock_workflow = self._create_mock_workflow() - - # Create a system variable - sys_var = WorkflowDraftVariable.new_sys_variable( - app_id=self._test_app_id, - user_id=self._test_user_id, - name="sys_var", - value=build_segment("sys_value"), - node_execution_id=self._node_exec_id, - ) - self._session.add(sys_var) - self._session.flush() - - # Attempt to reset the system variable - with pytest.raises(VariableResetError) as exc_info: - srv.reset_variable(mock_workflow, sys_var) - - assert "cannot reset system variable" in str(exc_info.value) - assert sys_var.id in str(exc_info.value) diff --git a/api/tests/unit_tests/services/workflow/test_workflow_human_input_delivery.py b/api/tests/unit_tests/services/workflow/test_workflow_human_input_delivery.py deleted file mode 100644 index 1f75c576b78..00000000000 --- a/api/tests/unit_tests/services/workflow/test_workflow_human_input_delivery.py +++ /dev/null @@ -1,255 +0,0 @@ -"""Tests for human-input delivery with a persisted draft workflow.""" - -import json -import uuid -from unittest.mock import MagicMock - -import pytest -from sqlalchemy.orm import Session, sessionmaker - -from core.workflow.human_input_adapter import ( - EmailDeliveryConfig, - EmailDeliveryMethod, - EmailRecipients, - ExternalRecipient, - MemberRecipient, -) -from core.workflow.nodes.human_input.entities import HumanInputNodeData -from graphon.enums import BuiltinNodeTypes -from models import Account -from models.model import App, AppMode -from models.workflow import Workflow, WorkflowType -from services import workflow_service as workflow_service_module -from services.workflow_service import WorkflowService -from tests.unit_tests.model_factories import make_account, make_app - -TENANT_ID = "11111111-1111-1111-1111-111111111111" -APP_ID = "22222222-2222-2222-2222-222222222222" -ACCOUNT_ID = "33333333-3333-3333-3333-333333333333" - - -def _make_service(sqlite_session: Session) -> WorkflowService: - return WorkflowService( - session_maker=sessionmaker(bind=sqlite_session.get_bind(), expire_on_commit=False), - ) - - -def _app() -> App: - return make_app( - app_id=APP_ID, - tenant_id=TENANT_ID, - mode=AppMode.WORKFLOW, - icon_type=None, - enable_site=False, - enable_api=False, - ) - - -def _account() -> Account: - return make_account(account_id=ACCOUNT_ID, name="Test User", email="test@example.com") - - -def _build_node_config(delivery_methods: list[EmailDeliveryMethod]) -> dict[str, object]: - return { - "id": "node-1", - "data": HumanInputNodeData( - title="Human Input", - type=BuiltinNodeTypes.HUMAN_INPUT, - delivery_methods=delivery_methods, - form_content="Test content", - inputs=[], - user_actions=[], - ), - } - - -def _persist_workflow(sqlite_session: Session, delivery_methods: list[EmailDeliveryMethod]) -> None: - node_config = _build_node_config(delivery_methods) - node_data = node_config["data"] - assert isinstance(node_data, HumanInputNodeData) - workflow = Workflow.new( - tenant_id=TENANT_ID, - app_id=APP_ID, - type=WorkflowType.WORKFLOW, - version=Workflow.VERSION_DRAFT, - graph=json.dumps( - { - "nodes": [ - { - "id": node_config["id"], - "data": node_data.model_dump(mode="json"), - } - ], - "edges": [], - } - ), - features="{}", - created_by=ACCOUNT_ID, - environment_variables=[], - conversation_variables=[], - rag_pipeline_variables=[], - ) - sqlite_session.add(workflow) - sqlite_session.commit() - sqlite_session.expunge_all() - - -def _make_email_method(enabled: bool = True, debug_mode: bool = False) -> EmailDeliveryMethod: - return EmailDeliveryMethod( - id=uuid.uuid4(), - enabled=enabled, - config=EmailDeliveryConfig( - recipients=EmailRecipients( - include_bound_group=False, - items=[ExternalRecipient(email="tester@example.com")], - ), - subject="Test subject", - body="Test body", - debug_mode=debug_mode, - ), - ) - - -@pytest.mark.parametrize("sqlite_session", [(Workflow,)], indirect=True) -def test_human_input_delivery_requires_draft_workflow(sqlite_session: Session): - service = _make_service(sqlite_session) - app_model = _app() - account = _account() - - with pytest.raises(ValueError, match="Workflow not initialized"): - service.test_human_input_delivery( - app_model=app_model, - account=account, - node_id="node-1", - delivery_method_id="delivery-1", - session=sqlite_session, - ) - assert sqlite_session.in_transaction() - - -@pytest.mark.parametrize("sqlite_session", [(Workflow,)], indirect=True) -def test_human_input_delivery_allows_disabled_method( - monkeypatch: pytest.MonkeyPatch, - sqlite_session: Session, -): - service = _make_service(sqlite_session) - delivery_method = _make_email_method(enabled=False) - _persist_workflow(sqlite_session, [delivery_method]) - service._build_human_input_variable_pool = MagicMock(return_value=MagicMock()) # type: ignore[attr-defined] - node_stub = MagicMock() - node_stub.render_form_content_before_submission.return_value = "rendered" - node_stub.resolve_default_values.return_value = {} - service._build_human_input_node_for_debugging = MagicMock(return_value=node_stub) # type: ignore[attr-defined] - service._create_human_input_delivery_test_form = MagicMock( # type: ignore[attr-defined] - return_value=("form-1", []) - ) - - test_service_instance = MagicMock() - monkeypatch.setattr( - workflow_service_module, - "HumanInputDeliveryTestService", - MagicMock(return_value=test_service_instance), - ) - - app_model = _app() - account = _account() - - service.test_human_input_delivery( - app_model=app_model, - account=account, - node_id="node-1", - delivery_method_id=str(delivery_method.id), - session=sqlite_session, - ) - - test_service_instance.send_test.assert_called_once() - assert sqlite_session.in_transaction() - - -@pytest.mark.parametrize("sqlite_session", [(Workflow,)], indirect=True) -def test_human_input_delivery_dispatches_to_test_service( - monkeypatch: pytest.MonkeyPatch, - sqlite_session: Session, -): - service = _make_service(sqlite_session) - delivery_method = _make_email_method(enabled=True) - _persist_workflow(sqlite_session, [delivery_method]) - service._build_human_input_variable_pool = MagicMock(return_value=MagicMock()) # type: ignore[attr-defined] - node_stub = MagicMock() - node_stub.render_form_content_before_submission.return_value = "rendered" - node_stub.resolve_default_values.return_value = {} - service._build_human_input_node_for_debugging = MagicMock(return_value=node_stub) # type: ignore[attr-defined] - service._create_human_input_delivery_test_form = MagicMock( # type: ignore[attr-defined] - return_value=("form-1", []) - ) - - test_service_instance = MagicMock() - monkeypatch.setattr( - workflow_service_module, - "HumanInputDeliveryTestService", - MagicMock(return_value=test_service_instance), - ) - - app_model = _app() - account = _account() - - service.test_human_input_delivery( - app_model=app_model, - account=account, - node_id="node-1", - delivery_method_id=str(delivery_method.id), - inputs={"#node-1.output#": "value"}, - session=sqlite_session, - ) - - pool_args = service._build_human_input_variable_pool.call_args.kwargs - assert pool_args["manual_inputs"] == {"#node-1.output#": "value"} - test_service_instance.send_test.assert_called_once() - assert sqlite_session.in_transaction() - - -@pytest.mark.parametrize("sqlite_session", [(Workflow,)], indirect=True) -def test_human_input_delivery_debug_mode_overrides_recipients( - monkeypatch: pytest.MonkeyPatch, - sqlite_session: Session, -): - service = _make_service(sqlite_session) - delivery_method = _make_email_method(enabled=True, debug_mode=True) - _persist_workflow(sqlite_session, [delivery_method]) - service._build_human_input_variable_pool = MagicMock(return_value=MagicMock()) # type: ignore[attr-defined] - node_stub = MagicMock() - node_stub.render_form_content_before_submission.return_value = "rendered" - node_stub.resolve_default_values.return_value = {} - service._build_human_input_node_for_debugging = MagicMock(return_value=node_stub) # type: ignore[attr-defined] - service._create_human_input_delivery_test_form = MagicMock( # type: ignore[attr-defined] - return_value=("form-1", []) - ) - - test_service_instance = MagicMock() - monkeypatch.setattr( - workflow_service_module, - "HumanInputDeliveryTestService", - MagicMock(return_value=test_service_instance), - ) - - app_model = _app() - account = _account() - - service.test_human_input_delivery( - app_model=app_model, - account=account, - node_id="node-1", - delivery_method_id=str(delivery_method.id), - session=sqlite_session, - ) - - test_service_instance.send_test.assert_called_once() - sent_method = test_service_instance.send_test.call_args.kwargs["method"] - assert isinstance(sent_method, EmailDeliveryMethod) - assert sent_method.config.debug_mode is True - assert sent_method.config.recipients.include_bound_group is False - assert len(sent_method.config.recipients.items) == 1 - recipient = sent_method.config.recipients.items[0] - assert isinstance(recipient, MemberRecipient) - assert recipient.reference_id == account.id - assert sqlite_session.in_transaction() diff --git a/api/tests/unit_tests/services/workflow/test_workflow_run_agg.py b/api/tests/unit_tests/services/workflow/test_workflow_run_agg.py deleted file mode 100644 index 64c12fbf217..00000000000 --- a/api/tests/unit_tests/services/workflow/test_workflow_run_agg.py +++ /dev/null @@ -1,811 +0,0 @@ -"""Workflow orchestration against real Engine, SQLAlchemy, and filesystem storage.""" - -import json -from datetime import UTC, datetime, timedelta -from pathlib import Path -from uuid import uuid4 - -import pytest -from sqlalchemy import Engine, select -from sqlalchemy.orm import Session - -from core.app.apps.workflow.app_config_manager import WorkflowAppConfig -from core.app.apps.workflow.app_queue_manager import WorkflowAppQueueManager -from core.app.apps.workflow.app_runner import WorkflowAppRunner -from core.app.entities.app_invoke_entities import InvokeFrom, WorkflowAppGenerateEntity -from core.app.entities.queue_entities import ( - QueueNodeStartedEvent, - QueueWorkflowFailedEvent, - QueueWorkflowPausedEvent, - QueueWorkflowSucceededEvent, -) -from core.app.layers.pause_state_persist_layer import PauseStateLayerConfig, WorkflowResumptionContext -from core.app.layers.trigger_post_layer import TriggerPostLayer -from core.repositories.sqlalchemy_workflow_execution_repository import SQLAlchemyWorkflowExecutionRepository -from core.repositories.sqlalchemy_workflow_node_execution_repository import SQLAlchemyWorkflowNodeExecutionRepository -from core.workflow.nodes.human_input.entities import HumanInputNodeData, ParagraphInputConfig, UserActionConfig -from extensions.ext_storage import storage -from extensions.storage.opendal_storage import OpenDALStorage -from graphon.engine_events import ( - GraphRunAbortedEvent, - GraphRunFailedEvent, - GraphRunPartialSucceededEvent, - GraphRunPausedEvent, - GraphRunStartedEvent, - GraphRunSucceededEvent, -) -from graphon.enums import WorkflowExecutionStatus -from graphon.nodes.start.entities import StartNodeData -from graphon.runtime import RuntimeState -from graphon.variable_loader import DUMMY_VARIABLE_LOADER -from models import Account -from models.enums import AppTriggerType, CreatorUserRole, WorkflowRunTriggeredFrom, WorkflowTriggerStatus -from models.model import AppMode -from models.trigger import WorkflowTriggerLog -from models.workflow import ( - Workflow, - WorkflowNodeExecutionModel, - WorkflowNodeExecutionTriggeredFrom, - WorkflowPause, - WorkflowRun, -) -from services.workflow.entities import WorkflowScheduleCFSPlanEntity -from services.workflow_run_agg import WorkflowRunAgg -from tasks.workflow_cfs_scheduler.cfs_scheduler import AsyncWorkflowCFSPlanEntity -from tasks.workflow_cfs_scheduler.entities import AsyncWorkflowQueue - - -def make_workflow_runner( - sqlite_engine: Engine, *, human_input: bool = False -) -> tuple[WorkflowAppRunner, SQLAlchemyWorkflowExecutionRepository, SQLAlchemyWorkflowNodeExecutionRepository]: - """Build the same runner/repositories/queue used by a workflow worker.""" - tenant_id, app_id, workflow_id, user_id, run_id = (str(uuid4()) for _ in range(5)) - graph: dict[str, list[dict[str, object]]] = { - "nodes": [{"id": "start", "data": {"type": "start", **StartNodeData(title="Start").model_dump(mode="json")}}], - "edges": [], - } - if human_input: - data = HumanInputNodeData( - title="Approval", - form_content="Decision: {{#$output.answer#}}", - inputs=[ParagraphInputConfig(output_variable_name="answer")], - user_actions=[UserActionConfig(id="approve", title="Approve")], - ) - graph["nodes"].append({"id": "approval", "data": {"type": "human-input", **data.model_dump(mode="json")}}) - graph["edges"].append({"id": "start-approval", "source": "start", "target": "approval"}) - graph["nodes"].append( - { - "id": "end", - "data": { - "type": "end", - "title": "End", - "outputs": [ - {"variable": "answer", "value_selector": ["approval", "answer"]}, - ], - }, - } - ) - graph["edges"].append({"id": "approval-end", "source": "approval", "target": "end", "sourceHandle": "approve"}) - workflow = Workflow( - id=workflow_id, - tenant_id=tenant_id, - app_id=app_id, - type="workflow", - version="draft", - graph=json.dumps(graph), - features="{}", - created_by=user_id, - ) - entity = WorkflowAppGenerateEntity( - task_id=str(uuid4()), - app_config=WorkflowAppConfig( - tenant_id=tenant_id, app_id=app_id, workflow_id=workflow_id, app_mode=AppMode.WORKFLOW - ), - inputs={}, - files=[], - user_id=user_id, - stream=True, - invoke_from=InvokeFrom.DEBUGGER, - workflow_execution_id=run_id, - ) - user = Account(name="Workflow Tester", email=f"{user_id}@example.com") - user.id = user_id - queue = WorkflowAppQueueManager(entity.task_id, user_id, entity.invoke_from, AppMode.WORKFLOW) - workflow_repository = SQLAlchemyWorkflowExecutionRepository( - sqlite_engine, - tenant_id=tenant_id, - user=user, - app_id=app_id, - triggered_from=WorkflowRunTriggeredFrom.DEBUGGING, - ) - node_repository = SQLAlchemyWorkflowNodeExecutionRepository( - sqlite_engine, - tenant_id=tenant_id, - user=user, - app_id=app_id, - triggered_from=WorkflowNodeExecutionTriggeredFrom.WORKFLOW_RUN, - ) - runner = WorkflowAppRunner( - application_generate_entity=entity, - queue_manager=queue, - variable_loader=DUMMY_VARIABLE_LOADER, - workflow=workflow, - system_user_id=user_id, - execution_driver=WorkflowRunAgg.run, - ) - return runner, workflow_repository, node_repository - - -def test_driver_publishes_persisted_start_index_and_success(sqlite_engine: Engine) -> None: - runner, workflow_repository, node_repository = make_workflow_runner(sqlite_engine) - WorkflowRunAgg.run(runner, None, workflow_repository, node_repository) - events = [message.event for message in runner._queue_manager.listen()] - starts = [event for event in events if isinstance(event, QueueNodeStartedEvent)] - assert len(starts) == 1 - assert isinstance(events[-1], QueueWorkflowSucceededEvent) - with Session(sqlite_engine) as session: - execution = session.scalar(select(WorkflowNodeExecutionModel)) - assert execution is not None - assert execution.index == starts[0].node_run_index == 1 - run = session.get(WorkflowRun, runner.application_generate_entity.workflow_execution_id) - assert run is not None - assert run.status == "succeeded" - - -@pytest.mark.parametrize("snapshot", [None, {}, {"tools": []}]) -def test_invalid_saved_tool_snapshot_is_rejected(sqlite_engine: Engine, snapshot: object) -> None: - from pydantic import ValidationError - - runner, workflow_repository, node_repository = make_workflow_runner(sqlite_engine) - runner.application_generate_entity.extras["workflow_tool_snapshot"] = snapshot - with pytest.raises(ValidationError): - WorkflowRunAgg.run(runner, None, workflow_repository, node_repository) - with Session(sqlite_engine) as session: - assert session.scalar(select(WorkflowRun)) is None - - -@pytest.fixture -def _local_storage(tmp_path: Path, monkeypatch: pytest.MonkeyPatch) -> None: - monkeypatch.setattr(storage, "storage_runner", OpenDALStorage(scheme="fs", root=str(tmp_path)), raising=False) - from core.app.apps import workflow_app_runner - - monkeypatch.setattr(workflow_app_runner.dispatch_human_input_email_task, "apply_async", lambda **_kwargs: None) - - -@pytest.mark.usefixtures("_local_storage") -@pytest.mark.parametrize("persist_snapshot", [False, True]) -def test_pause_publication_persists_summary_and_optional_snapshot( - sqlite_engine: Engine, persist_snapshot: bool, caplog: pytest.LogCaptureFixture -) -> None: - runner, workflow_repository, node_repository = make_workflow_runner(sqlite_engine, human_input=True) - config = ( - PauseStateLayerConfig(sqlite_engine, runner.application_generate_entity.user_id) if persist_snapshot else None - ) - WorkflowRunAgg.run(runner, config, workflow_repository, node_repository) - assert "Failed to enqueue human input email task" not in caplog.text - events = [message.event for message in runner._queue_manager.listen()] - assert isinstance(events[-1], QueueWorkflowPausedEvent), events - with Session(sqlite_engine) as session: - run = session.get(WorkflowRun, runner.application_generate_entity.workflow_execution_id) - assert run is not None - pause = session.scalar(select(WorkflowPause).where(WorkflowPause.workflow_run_id == run.id)) - assert run.status == "paused" - assert run.finished_at is None - assert run.total_steps == 2 - if not persist_snapshot: - assert pause is None - return - assert pause is not None - snapshot = WorkflowResumptionContext.loads(storage.load(pause.state_object_key).decode()) - restored = RuntimeState.from_snapshot(snapshot.serialized_graph_runtime_state) - assert restored.graph_execution.paused - assert snapshot.get_response_stream_filter().dumps() == snapshot.serialized_response_stream_filter_state - assert snapshot.get_generate_entity().model_dump(mode="json") == runner.application_generate_entity.model_dump( - mode="json" - ) - - -@pytest.mark.usefixtures("_local_storage") -def test_pause_database_failure_publishes_failure_instead_of_resume_ready( - sqlite_engine: Engine, monkeypatch: pytest.MonkeyPatch -) -> None: - def fail_pause(*_args: object, **_kwargs: object) -> None: - raise RuntimeError("pause-write-failed") - - monkeypatch.setattr( - "repositories.sqlalchemy_api_workflow_run_repository.DifyAPISQLAlchemyWorkflowRunRepository.create_workflow_pause", - fail_pause, - ) - retirements: list[dict[str, object]] = [] - collections: list[dict[str, object]] = [] - - def retire(_store: object, **kwargs: object) -> list[str]: - retirements.append(kwargs) - return ["workspace-1"] - - monkeypatch.setattr( - "core.workflow.nodes.agent_v2.workspace_retirement_layer.WorkflowAgentWorkspaceStore.retire_workflow_run", - retire, - ) - monkeypatch.setattr( - "core.workflow.nodes.agent_v2.workspace_retirement_layer.enqueue_agent_resource_collection", - lambda **kwargs: collections.append(kwargs), - ) - runner, workflow_repository, node_repository = make_workflow_runner(sqlite_engine, human_input=True) - trigger_log = WorkflowTriggerLog( - tenant_id=runner.application_generate_entity.app_config.tenant_id, - app_id=runner.application_generate_entity.app_config.app_id, - workflow_id=runner._workflow.id, - workflow_run_id=None, - root_node_id=None, - trigger_metadata="{}", - trigger_type=AppTriggerType.TRIGGER_SCHEDULE, - trigger_data="{}", - inputs="{}", - outputs=None, - status=WorkflowTriggerStatus.RUNNING, - error=None, - queue_name="workflow", - celery_task_id=None, - created_by_role=CreatorUserRole.ACCOUNT, - created_by=runner.application_generate_entity.user_id, - elapsed_time=5.0, - ) - with Session(sqlite_engine, expire_on_commit=False) as session: - session.add(trigger_log) - session.commit() - start_time = datetime.now(UTC) - timedelta(seconds=10) - plan = AsyncWorkflowCFSPlanEntity( - schedule_strategy=WorkflowScheduleCFSPlanEntity.Strategy.Nop, - queue=AsyncWorkflowQueue.PROFESSIONAL_QUEUE, - ) - runner._graph_engine_layers = (TriggerPostLayer(plan, start_time, trigger_log.id),) - WorkflowRunAgg.run( - runner, - PauseStateLayerConfig(sqlite_engine, runner.application_generate_entity.user_id), - workflow_repository, - node_repository, - ) - events = [message.event for message in runner._queue_manager.listen()] - assert isinstance(events[-1], QueueWorkflowFailedEvent), events - assert "pause-write-failed" in events[-1].error - assert not any(isinstance(event, QueueWorkflowPausedEvent) for event in events) - assert retirements == [ - { - "tenant_id": runner.application_generate_entity.app_config.tenant_id, - "app_id": runner.application_generate_entity.app_config.app_id, - "workflow_run_id": runner.application_generate_entity.workflow_execution_id, - } - ] - assert collections == [ - { - "tenant_id": runner.application_generate_entity.app_config.tenant_id, - "workspace_ids": ["workspace-1"], - } - ] - with Session(sqlite_engine) as session: - persisted_trigger = session.get(WorkflowTriggerLog, trigger_log.id) - assert persisted_trigger is not None - assert persisted_trigger.status == WorkflowTriggerStatus.FAILED - assert persisted_trigger.finished_at is not None - assert persisted_trigger.elapsed_time == pytest.approx( - 5.0 + (persisted_trigger.finished_at.replace(tzinfo=UTC) - start_time).total_seconds() - ) - run = session.get(WorkflowRun, runner.application_generate_entity.workflow_execution_id) - assert run is not None - assert run.status == "failed" - assert session.scalar(select(WorkflowPause)) is None - - -def test_late_form_submission_is_refreshed_before_the_real_queue_closes(sqlite_session: Session) -> None: - from core.app.apps.workflow_app_runner import WorkflowBasedAppRunner - from core.app.entities.queue_entities import QueueHumanInputFormFilledEvent - from core.repositories.human_input_repository import HumanInputFormSubmissionRepository - from core.workflow.nodes.human_input.enums import HumanInputFormStatus - from graphon.engine_events import NodeRunStartedEvent - from tests.unit_tests.core.app.apps.test_workflow_human_input_completion import ( - _make_paused_workflow, - _resume_events, - _save_form, - ) - - first = _save_form(sqlite_session, status=HumanInputFormStatus.SUBMITTED) - late = _save_form(sqlite_session, status=HumanInputFormStatus.WAITING) - queue = WorkflowAppQueueManager("form-refresh-task", "user", InvokeFrom.DEBUGGER, AppMode.WORKFLOW) - runner = WorkflowBasedAppRunner(queue_manager=queue, app_id="app") - entry = _make_paused_workflow(runner, [first, late]) - for event in _resume_events(runner, entry): - runner.handle_event(entry, event) - if isinstance(event, NodeRunStartedEvent): - HumanInputFormSubmissionRepository().mark_submitted( - form_id=late.id, - recipient_id=None, - selected_action_id="approve", - form_data={"answer": "late answer"}, - submission_user_id=None, - submission_end_user_id=None, - ) - events = [message.event for message in queue.listen()] - completions = [event for event in events if isinstance(event, QueueHumanInputFormFilledEvent)] - assert [event.form_id for event in completions] == [first.id, late.id] - assert completions[-1].rendered_content == "Decision: late answer" - assert isinstance(events[-1], QueueWorkflowSucceededEvent) - - -@pytest.mark.usefixtures("_local_storage") -def test_durable_pause_resumes_with_one_form_completion_and_continuing_indexes(sqlite_engine: Engine) -> None: - from core.app.entities.queue_entities import QueueHumanInputFormFilledEvent, QueueWorkflowStartedEvent - from core.repositories.human_input_repository import HumanInputFormSubmissionRepository - from models.human_input import HumanInputForm - - original, workflow_repository, node_repository = make_workflow_runner(sqlite_engine, human_input=True) - config = PauseStateLayerConfig(sqlite_engine, original.application_generate_entity.user_id) - WorkflowRunAgg.run(original, config, workflow_repository, node_repository) - before = [message.event for message in original._queue_manager.listen()] - assert isinstance(before[-1], QueueWorkflowPausedEvent) - with Session(sqlite_engine) as session: - pause = session.scalar(select(WorkflowPause)) - form = session.scalar(select(HumanInputForm)) - assert pause is not None - assert form is not None - form_id = form.id - snapshot = WorkflowResumptionContext.loads(storage.load(pause.state_object_key).decode()) - HumanInputFormSubmissionRepository().mark_submitted( - form_id=form_id, - recipient_id=None, - selected_action_id="approve", - form_data={"answer": "approved result"}, - submission_user_id=None, - submission_end_user_id=None, - ) - entity = original.application_generate_entity.model_copy(update={"task_id": str(uuid4())}) - queue = WorkflowAppQueueManager(entity.task_id, entity.user_id, entity.invoke_from, AppMode.WORKFLOW) - resumed = WorkflowAppRunner( - application_generate_entity=entity, - queue_manager=queue, - variable_loader=DUMMY_VARIABLE_LOADER, - workflow=original._workflow, - system_user_id=entity.user_id, - graph_runtime_state=RuntimeState.from_snapshot(snapshot.serialized_graph_runtime_state), - response_stream_filter=snapshot.get_response_stream_filter(), - execution_driver=WorkflowRunAgg.run, - ) - WorkflowRunAgg.run(resumed, config, workflow_repository, node_repository) - after = [message.event for message in queue.listen()] - assert isinstance(after[0], QueueWorkflowStartedEvent) - assert len(after[0].node_execution_snapshots) == 2 - completions = [event for event in after if isinstance(event, QueueHumanInputFormFilledEvent)] - assert [event.form_id for event in completions] == [form_id] - starts = [event for event in after if isinstance(event, QueueNodeStartedEvent)] - assert [(event.node_id, event.node_run_index) for event in starts] == [("approval", 2), ("end", 3)] - assert isinstance(after[-1], QueueWorkflowSucceededEvent) - assert after[-1].outputs == {"answer": "approved result"} - with Session(sqlite_engine) as session: - executions = session.scalars( - select(WorkflowNodeExecutionModel).order_by(WorkflowNodeExecutionModel.index) - ).all() - assert [(row.node_id, row.index) for row in executions] == [("start", 1), ("approval", 2), ("end", 3)] - - -@pytest.mark.usefixtures("_local_storage") -def test_pause_before_first_tool_call_preserves_its_published_definition(sqlite_engine: Engine) -> None: - from core.app.entities.queue_entities import QueueHumanInputFormFilledEvent - from core.repositories.human_input_repository import HumanInputFormSubmissionRepository - from graphon.nodes.tool.entities import ToolNodeData - from models.human_input import HumanInputForm - from models.model import App - from models.tools import WorkflowToolProvider - - original, workflow_repository, node_repository = make_workflow_runner(sqlite_engine, human_input=True) - original_entity = original.application_generate_entity - tenant_id = original_entity.app_config.tenant_id - source_app_id = str(uuid4()) - source_workflow_ids = [str(uuid4()), str(uuid4())] - with Session(sqlite_engine) as session, session.begin(): - session.add( - App( - id=source_app_id, - tenant_id=tenant_id, - name="Source", - mode=AppMode.WORKFLOW, - enable_site=False, - enable_api=False, - ) - ) - for workflow_id, version, marker in zip( - source_workflow_ids, ("old", "new"), ("original", "republished"), strict=True - ): - session.add( - Workflow( - id=workflow_id, - app_id=source_app_id, - tenant_id=tenant_id, - created_by=original_entity.user_id, - type="workflow", - version=version, - features="{}", - graph=json.dumps( - { - "nodes": [ - { - "id": "source-start", - "data": { - "type": "start", - "title": "Start", - "variables": [ - { - "variable": "marker", - "type": "text-input", - "label": "Marker", - "default": marker, - "required": False, - } - ], - }, - }, - { - "id": "source-end", - "data": { - "type": "end", - "title": "End", - "outputs": [ - {"variable": "marker", "value_selector": ["source-start", "marker"]} - ], - }, - }, - ], - "edges": [{"source": "source-start", "target": "source-end"}], - } - ), - ) - ) - provider = WorkflowToolProvider( - tenant_id=tenant_id, - app_id=source_app_id, - version="old", - name="source", - label="Original source", - icon="{}", - user_id=original_entity.user_id, - description="Source", - parameter_configuration='[{"name":"marker","description":"Marker","form":"llm"}]', - ) - session.add(provider) - session.flush() - provider_id = provider.id - graph = dict(original._workflow.graph_dict) - graph["nodes"].append( - { - "id": "tool", - "data": { - "type": "tool", - **ToolNodeData( - title="Tool", - provider_id=provider_id, - provider_type="workflow", - provider_name="source", - tool_name="source", - tool_label="Source", - tool_configurations={}, - tool_parameters={}, - ).model_dump(mode="json"), - }, - } - ) - graph["nodes"][2]["data"]["outputs"] = [{"variable": "result", "value_selector": ["tool", "json"]}] - graph["edges"][1]["target"] = "tool" - graph["edges"].append({"source": "tool", "target": "end"}) - original._workflow.graph = json.dumps(graph) - config = PauseStateLayerConfig(sqlite_engine, original_entity.user_id) - WorkflowRunAgg.run(original, config, workflow_repository, node_repository) - before = [message.event for message in original._queue_manager.listen()] - assert isinstance(before[-1], QueueWorkflowPausedEvent) - assert not any(isinstance(event, QueueNodeStartedEvent) and event.node_id == "tool" for event in before) - with Session(sqlite_engine) as session, session.begin(): - pause = session.scalar(select(WorkflowPause)) - form = session.scalar(select(HumanInputForm)) - provider = session.get(WorkflowToolProvider, provider_id) - assert pause is not None - assert form is not None - assert provider is not None - form_id = form.id - snapshot = WorkflowResumptionContext.loads(storage.load(pause.state_object_key).decode()) - provider.version = "new" - provider.label = "Republished source" - HumanInputFormSubmissionRepository().mark_submitted( - form_id=form_id, - recipient_id=None, - selected_action_id="approve", - form_data={"answer": "approved"}, - submission_user_id=None, - submission_end_user_id=None, - ) - entity = snapshot.get_generate_entity().model_copy(update={"task_id": str(uuid4())}) - assert isinstance(entity, WorkflowAppGenerateEntity) - queue = WorkflowAppQueueManager(entity.task_id, entity.user_id, entity.invoke_from, AppMode.WORKFLOW) - resumed = WorkflowAppRunner( - application_generate_entity=entity, - queue_manager=queue, - variable_loader=DUMMY_VARIABLE_LOADER, - workflow=original._workflow, - system_user_id=entity.user_id, - graph_runtime_state=RuntimeState.from_snapshot(snapshot.serialized_graph_runtime_state), - response_stream_filter=snapshot.get_response_stream_filter(), - execution_driver=WorkflowRunAgg.run, - ) - WorkflowRunAgg.run(resumed, config, workflow_repository, node_repository) - after = [message.event for message in queue.listen()] - assert [event.form_id for event in after if isinstance(event, QueueHumanInputFormFilledEvent)] == [form_id] - assert isinstance(after[-1], QueueWorkflowSucceededEvent), after - assert after[-1].outputs == {"result": [{"marker": "original"}]} - assert entity.extras == original_entity.extras - with Session(sqlite_engine) as session: - child_workflows = session.scalars( - select(WorkflowNodeExecutionModel.workflow_id).where(WorkflowNodeExecutionModel.app_id == source_app_id) - ).all() - assert child_workflows - assert set(child_workflows) == {source_workflow_ids[0]} - - # A new execution resolves the newly published graph and parameter defaults. - graph["nodes"] = [node for node in graph["nodes"] if node["id"] != "approval"] - graph["edges"] = [{"source": "start", "target": "tool"}, {"source": "tool", "target": "end"}] - original._workflow.graph = json.dumps(graph) - fresh_entity = original_entity.model_copy( - update={"task_id": str(uuid4()), "workflow_execution_id": str(uuid4()), "extras": {}} - ) - fresh_queue = WorkflowAppQueueManager( - fresh_entity.task_id, fresh_entity.user_id, fresh_entity.invoke_from, AppMode.WORKFLOW - ) - fresh = WorkflowAppRunner( - application_generate_entity=fresh_entity, - queue_manager=fresh_queue, - variable_loader=DUMMY_VARIABLE_LOADER, - workflow=original._workflow, - system_user_id=fresh_entity.user_id, - execution_driver=WorkflowRunAgg.run, - ) - WorkflowRunAgg.run(fresh, None, workflow_repository, node_repository) - fresh_events = [message.event for message in fresh_queue.listen()] - assert isinstance(fresh_events[-1], QueueWorkflowSucceededEvent), fresh_events - assert fresh_events[-1].outputs == {"result": [{"marker": "republished"}]} - - -def test_active_execution_snapshot_failure_is_not_published_as_a_pause(sqlite_engine: Engine) -> None: - from graphon.engine_events import GraphRunFailedEvent, GraphRunPausedEvent - - runner, workflow_repository, node_repository = make_workflow_runner(sqlite_engine, human_input=True) - prepared = runner.prepare() - assert prepared is not None - aggregate = WorkflowRunAgg( - prepared, - runner, - PauseStateLayerConfig(sqlite_engine, runner.application_generate_entity.user_id), - workflow_execution_repository=workflow_repository, - workflow_node_execution_repository=node_repository, - ) - # A participant that has not shut down must prevent a resume-ready snapshot. - with prepared.entry.graph_engine.runtime_state.graph_execution.track_execution(): - events = list(aggregate.iter_events()) - assert isinstance(events[-1], GraphRunFailedEvent) - assert "during active execution" in events[-1].error - assert not any(isinstance(event, GraphRunPausedEvent) for event in events) - with Session(sqlite_engine) as session: - run = session.get(WorkflowRun, runner.application_generate_entity.workflow_execution_id) - assert run is not None - assert run.status == "failed" - assert session.scalar(select(WorkflowPause)) is None - - -@pytest.mark.usefixtures("_local_storage") -@pytest.mark.parametrize( - ("terminal_event", "status"), - [ - (GraphRunSucceededEvent(), WorkflowExecutionStatus.SUCCEEDED), - (GraphRunPartialSucceededEvent(exceptions_count=1), WorkflowExecutionStatus.PARTIAL_SUCCEEDED), - (GraphRunFailedEvent(error="failed"), WorkflowExecutionStatus.FAILED), - (GraphRunAbortedEvent(reason="stopped"), WorkflowExecutionStatus.STOPPED), - ], -) -def test_elapsed_time_accumulates_across_pauses_and_matches_delayed_stream_responses( - monkeypatch: pytest.MonkeyPatch, - sqlite_engine: Engine, - terminal_event: GraphRunSucceededEvent | GraphRunPartialSucceededEvent | GraphRunFailedEvent | GraphRunAbortedEvent, - status: WorkflowExecutionStatus, -) -> None: - from core.app.apps.common.workflow_response_converter import WorkflowResponseConverter - from core.app.entities.task_entities import WorkflowPauseStreamResponse - from core.repositories.human_input_repository import HumanInputFormSubmissionRepository - from core.workflow.nodes.human_input.enums import HumanInputFormStatus - from core.workflow.system_variables import build_system_variables - from graphon.entities import WorkflowStartReason - from models.human_input import HumanInputForm - - clock = [100.0] - monkeypatch.setattr("repositories.entities.workflow_execution.perf_counter", lambda: clock[0]) - original, workflow_repository, node_repository = make_workflow_runner(sqlite_engine, human_input=True) - graph = original._workflow.graph_dict - graph["nodes"].append({"id": "second-approval", "data": graph["nodes"][1]["data"]}) - graph["edges"][1]["target"] = "second-approval" - graph["edges"].append({"id": "second-end", "source": "second-approval", "target": "end", "sourceHandle": "approve"}) - original._workflow.graph = json.dumps(graph) - runner = original - run_id = runner.application_generate_entity.workflow_execution_id - config = PauseStateLayerConfig(sqlite_engine, runner.application_generate_entity.user_id) - account = Account(name="Test", email="test@example.com") - account.id = runner.application_generate_entity.user_id - total = 0.0 - - for index, duration in enumerate([2.0, 3.0, 4.0]): - prepared = runner.prepare() - assert prepared is not None - aggregate = WorkflowRunAgg( - prepared, - runner, - config, - workflow_execution_repository=workflow_repository, - workflow_node_execution_repository=node_repository, - ) - is_paused = index < 2 - if is_paused: - events = [] - for event in aggregate.iter_events(): - if isinstance(event, GraphRunStartedEvent): - with Session(sqlite_engine) as session: - started = session.get(WorkflowRun, run_id) - assert started is not None - assert started.elapsed_time == total - clock[0] += duration - events.append(event) - assert isinstance(events[-1], GraphRunPausedEvent), events - else: - aggregate._persistence_layer.on_event(GraphRunStartedEvent()) - clock[0] += duration - aggregate._persistence_layer.on_event(terminal_event) - total += duration - with Session(sqlite_engine) as session: - execution = session.get(WorkflowRun, run_id) - assert execution is not None - assert execution.elapsed_time == total - assert execution.status == (WorkflowExecutionStatus.PAUSED if is_paused else status) - assert (execution.finished_at is None) == is_paused - - # Queue processing and pause waits must not extend a settled segment. - clock[0] += 1000 - entity = runner.application_generate_entity - runtime = prepared.entry.graph_engine.runtime_state - converter = WorkflowResponseConverter( - application_generate_entity=entity, - user=account, - system_variables=build_system_variables(workflow_execution_id=run_id), - ) - converter.workflow_start_to_stream_response( - task_id=entity.task_id, - workflow_run_id=run_id, - workflow_id=original._workflow.id, - reason=WorkflowStartReason.INITIAL if index == 0 else WorkflowStartReason.RESUMPTION, - ) - if is_paused: - responses = converter.workflow_pause_to_stream_response( - event=QueueWorkflowPausedEvent(), task_id=entity.task_id, graph_runtime_state=runtime - ) - pause_response = responses[-1] - assert isinstance(pause_response, WorkflowPauseStreamResponse) - assert pause_response.data.elapsed_time == total - with Session(sqlite_engine) as session: - pause = session.scalar(select(WorkflowPause).where(WorkflowPause.workflow_run_id == run_id)) - form = session.scalar( - select(HumanInputForm).where( - HumanInputForm.workflow_run_id == run_id, HumanInputForm.status == HumanInputFormStatus.WAITING - ) - ) - assert pause is not None - assert form is not None - form_id = form.id - snapshot = WorkflowResumptionContext.loads(storage.load(pause.state_object_key).decode()) - entity = snapshot.get_generate_entity() - assert isinstance(entity, WorkflowAppGenerateEntity) - assert entity.workflow_execution_timing.elapsed_time == snapshot.elapsed_time == total - HumanInputFormSubmissionRepository().mark_submitted( - form_id=form_id, - recipient_id=None, - selected_action_id="approve", - form_data={"answer": "approved result"}, - submission_user_id=None, - submission_end_user_id=None, - ) - entity.task_id = str(uuid4()) - runner = WorkflowAppRunner( - application_generate_entity=entity, - queue_manager=WorkflowAppQueueManager( - entity.task_id, entity.user_id, entity.invoke_from, AppMode.WORKFLOW - ), - variable_loader=DUMMY_VARIABLE_LOADER, - workflow=original._workflow, - system_user_id=entity.user_id, - graph_runtime_state=RuntimeState.from_snapshot(snapshot.serialized_graph_runtime_state), - response_stream_filter=snapshot.get_response_stream_filter(), - execution_driver=WorkflowRunAgg.run, - ) - else: - response = converter.workflow_finish_to_stream_response( - task_id=entity.task_id, workflow_id=original._workflow.id, status=status, graph_runtime_state=runtime - ) - assert response.data.elapsed_time == 9.0 - assert entity.workflow_execution_timing.finish() == 9.0 - - -@pytest.mark.usefixtures("_local_storage") -@pytest.mark.parametrize("stop_timing", ["before_summary", "during_upload"]) -def test_stop_wins_while_engine_is_persisting_a_pause( - sqlite_engine: Engine, monkeypatch: pytest.MonkeyPatch, stop_timing: str -) -> None: - from sqlalchemy.orm import sessionmaker - - from core.app.entities.queue_entities import QueueStopEvent - from core.workflow.nodes.human_input.enums import HumanInputFormStatus - from graphon.enums import WorkflowExecutionStatus - from models.human_input import HumanInputForm - from repositories.sqlalchemy_api_workflow_run_repository import DifyAPISQLAlchemyWorkflowRunRepository - - runner, workflow_repository, node_repository = make_workflow_runner(sqlite_engine, human_input=True) - entity = runner.application_generate_entity - repository = DifyAPISQLAlchemyWorkflowRunRepository(sessionmaker(sqlite_engine)) - stop_calls = [] - - def request_stop() -> None: - admitted = repository.stop_paused_workflow_task( - tenant_id=entity.app_config.tenant_id, - app_id=entity.app_config.app_id, - task_id=entity.task_id, - owner=None, - ) - assert admitted is not None - stop_calls.append(admitted.status) - - if stop_timing == "during_upload": - original_save = storage.save - - def save(key: str, data: bytes) -> None: - if key.startswith("workflow-state-"): - request_stop() - return original_save(key, data) - - monkeypatch.setattr(storage, "save", save) - else: - original_persist = WorkflowRunAgg._persist_pause - - def persist(aggregate: WorkflowRunAgg, event: GraphRunPausedEvent) -> None: - request_stop() - return original_persist(aggregate, event) - - monkeypatch.setattr(WorkflowRunAgg, "_persist_pause", persist) - - WorkflowRunAgg.run( - runner, PauseStateLayerConfig(sqlite_engine, entity.user_id), workflow_repository, node_repository - ) - events = [message.event for message in runner._queue_manager.listen()] - assert stop_calls - assert isinstance(events[-1], QueueStopEvent), events - assert not any(isinstance(event, QueueWorkflowPausedEvent) for event in events) - with Session(sqlite_engine) as session: - run = session.get(WorkflowRun, entity.workflow_execution_id) - assert run is not None - assert run.status == WorkflowExecutionStatus.STOPPED - assert run.stop_requested_at is not None - assert run.task_id == entity.task_id - assert session.scalar(select(WorkflowPause).where(WorkflowPause.workflow_run_id == run.id)) is None - forms = session.scalars(select(HumanInputForm).where(HumanInputForm.workflow_run_id == run.id)).all() - assert forms - assert all(form.status == HumanInputFormStatus.EXPIRED for form in forms) - - -def test_legacy_caller_rejects_human_input_before_creating_run_or_form(sqlite_engine: Engine) -> None: - from models.human_input import HumanInputForm - - runner, workflow_repository, node_repository = make_workflow_runner(sqlite_engine, human_input=True) - runner.application_generate_entity.allow_human_input = False - with pytest.raises(ValueError, match="Human Input nodes are not supported"): - WorkflowRunAgg.run(runner, None, workflow_repository, node_repository) - with Session(sqlite_engine) as session: - assert session.scalar(select(WorkflowRun)) is None - assert session.scalar(select(HumanInputForm)) is None diff --git a/api/tests/unit_tests/services/workspace/test_gateways.py b/api/tests/unit_tests/services/workspace/test_gateways.py index 3e0a5e87b69..08e1792686f 100644 --- a/api/tests/unit_tests/services/workspace/test_gateways.py +++ b/api/tests/unit_tests/services/workspace/test_gateways.py @@ -1,226 +1,465 @@ -"""Workspace policies and notifications with only external clients replaced.""" - from collections.abc import Callable from unittest.mock import Mock, create_autospec import pytest -from enums import DeploymentEdition +from enums import CloudPlan, DeploymentEdition from extensions.ext_redis import RedisClientWrapper +from machinery.context import RequestContext +from models.account import Account from services.account.adapters import RedisInvitationTokenStore +from services.credit_pool_service import CreditPoolBalance +from services.enterprise.base import EnterpriseRequest from services.enterprise.enterprise_service import WorkspacePermission from services.entities.feature_entities import FeatureModel, LicenseLimitationModel, LicenseModel, LimitationModel from services.errors.base import NoPermissionError from services.errors.workspace import WorkspaceInvitationQuotaError, WorkspaceMemberLicenseQuotaError +from services.file_service import FileService from services.workspace import gateways -from tests.unit_tests.account_domain import AccountDomain +from services.workspace.contracts import WorkspaceMemberRole, WorkspaceMemberRoleSubject +from services.workspace.gateways import ( + DeploymentWorkspaceFeatureGateway, + DeploymentWorkspaceMemberRoleResolver, + DeploymentWorkspacePlanGateway, + WorkspaceFileGateway, +) +from tests.helpers.unit.account_domain import AccountDomain +from tests.helpers.unit.rbac_fakes import RBACTransport + +"""Workspace policies and notifications with only external clients replaced.""" -@pytest.fixture -def invitations() -> gateways.WorkspaceInvitationGateway: - redis = create_autospec(RedisClientWrapper, instance=True) - return gateways.WorkspaceInvitationGateway(tokens=RedisInvitationTokenStore(redis=redis), redis=redis) +class TestGatewaysCases: + @staticmethod + @pytest.fixture + def invitations() -> gateways.WorkspaceInvitationGateway: + redis = create_autospec(RedisClientWrapper, instance=True) + return gateways.WorkspaceInvitationGateway(tokens=RedisInvitationTokenStore(redis=redis), redis=redis) + @staticmethod + @pytest.mark.parametrize("allowed", [True, False]) + def test_invitation_policy_exposes_domain_errors( + invitations: gateways.WorkspaceInvitationGateway, + monkeypatch: pytest.MonkeyPatch, + config_overrides: Callable[..., None], + allowed: bool, + ) -> None: + config_overrides(DEPLOYMENT_EDITION=DeploymentEdition.ENTERPRISE) + requested: list[str] = [] -@pytest.mark.parametrize("allowed", [True, False]) -def test_invitation_policy_exposes_domain_errors( - invitations: gateways.WorkspaceInvitationGateway, - monkeypatch: pytest.MonkeyPatch, - config_overrides: Callable[..., None], - allowed: bool, -) -> None: - config_overrides(DEPLOYMENT_EDITION=DeploymentEdition.ENTERPRISE) - requested: list[str] = [] + def get_permission(workspace_id: str) -> WorkspacePermission: + requested.append(workspace_id) + return WorkspacePermission(workspaceId=workspace_id, allowMemberInvite=allowed) - def get_permission(workspace_id: str) -> WorkspacePermission: - requested.append(workspace_id) - return WorkspacePermission(workspaceId=workspace_id, allowMemberInvite=allowed) - - monkeypatch.setattr(gateways.EnterpriseService.WorkspacePermissionService, "get_permission", get_permission) - if allowed: - invitations.ensure_allowed("workspace") - else: - with pytest.raises(NoPermissionError, match="Workspace policy prohibits member invitations"): + monkeypatch.setattr(gateways.EnterpriseService.WorkspacePermissionService, "get_permission", get_permission) + if allowed: invitations.ensure_allowed("workspace") - assert requested == ["workspace"] + else: + with pytest.raises(NoPermissionError, match="Workspace policy prohibits member invitations"): + invitations.ensure_allowed("workspace") + assert requested == ["workspace"] + @staticmethod + @pytest.mark.parametrize(("limit", "allowed"), [(0, True), (-1, True), (4, True), (5, True), (3, False)]) + def test_cloud_invitation_member_limit( + invitations: gateways.WorkspaceInvitationGateway, + monkeypatch: pytest.MonkeyPatch, + config_overrides: Callable[..., None], + limit: int, + allowed: bool, + ) -> None: + config_overrides(DEPLOYMENT_EDITION=DeploymentEdition.CLOUD) + # Use the repository's current count, not a potentially stale billing size. + features = Mock(return_value=FeatureModel(members=LimitationModel(size=99, limit=limit))) + license_query = create_autospec(gateways.SystemFeatureService.get_license, instance=True) + monkeypatch.setattr(gateways.FeatureService, "get_features", features) + monkeypatch.setattr(gateways.SystemFeatureService, "get_license", license_query) -@pytest.mark.parametrize(("limit", "allowed"), [(0, True), (-1, True), (4, True), (5, True), (3, False)]) -def test_cloud_invitation_member_limit( - invitations: gateways.WorkspaceInvitationGateway, - monkeypatch: pytest.MonkeyPatch, - config_overrides: Callable[..., None], - limit: int, - allowed: bool, -) -> None: - config_overrides(DEPLOYMENT_EDITION=DeploymentEdition.CLOUD) - # Use the repository's current count, not a potentially stale billing size. - features = Mock(return_value=FeatureModel(members=LimitationModel(size=99, limit=limit))) - license_query = create_autospec(gateways.SystemFeatureService.get_license, instance=True) - monkeypatch.setattr(gateways.FeatureService, "get_features", features) - monkeypatch.setattr(gateways.SystemFeatureService, "get_license", license_query) - - assert invitations.requires_capacity_check - if allowed: - invitations.check_capacity("workspace", current_members=2, new_members=2, new_accounts=1) - else: - with pytest.raises(WorkspaceInvitationQuotaError) as error: + assert invitations.requires_capacity_check + if allowed: invitations.check_capacity("workspace", current_members=2, new_members=2, new_accounts=1) - assert error.value.seats is False - features.assert_called_once_with(tenant_id="workspace", exclude_vector_space=True) - license_query.assert_not_called() + else: + with pytest.raises(WorkspaceInvitationQuotaError) as error: + invitations.check_capacity("workspace", current_members=2, new_members=2, new_accounts=1) + assert error.value.seats is False + features.assert_called_once_with(tenant_id="workspace", exclude_vector_space=True) + license_query.assert_not_called() - -@pytest.mark.parametrize( - ("member_quota", "seat_quota", "new_accounts", "failure"), - [ - ( - LicenseLimitationModel(enabled=True, size=2, limit=4), - LicenseLimitationModel(enabled=True, size=9, limit=10), - 1, - None, - ), - ( - LicenseLimitationModel(enabled=True, size=3, limit=4), - LicenseLimitationModel(enabled=True, size=9, limit=10), - 1, - "members", - ), - (LicenseLimitationModel(), LicenseLimitationModel(enabled=True, size=10, limit=10), 1, "seats"), - (LicenseLimitationModel(), LicenseLimitationModel(enabled=True, size=9, limit=10), 2, "seats"), - (LicenseLimitationModel(), LicenseLimitationModel(enabled=True, size=8, limit=10), 2, None), - (LicenseLimitationModel(), LicenseLimitationModel(enabled=True, size=10, limit=10), 0, None), - ( - LicenseLimitationModel(enabled=True, size=99, limit=0), - LicenseLimitationModel(enabled=True, size=99, limit=0), - 2, - None, - ), - ( - LicenseLimitationModel(enabled=False, size=99, limit=1), - LicenseLimitationModel(enabled=False, size=99, limit=1), - 2, - None, - ), - ], - ids=[ - "exact-capacity", - "workspace-full", - "seats-full", - "not-enough-seats-for-batch", - "exact-seats-for-batch", - "existing-accounts", - "unlimited", - "disabled", - ], -) -def test_enterprise_invitation_member_and_seat_limits( - invitations: gateways.WorkspaceInvitationGateway, - monkeypatch: pytest.MonkeyPatch, - config_overrides: Callable[..., None], - member_quota: LicenseLimitationModel, - seat_quota: LicenseLimitationModel, - new_accounts: int, - failure: str | None, -) -> None: - config_overrides(DEPLOYMENT_EDITION=DeploymentEdition.ENTERPRISE) - monkeypatch.setattr( - gateways.FeatureService, "get_features", Mock(return_value=FeatureModel(workspace_members=member_quota)) + @staticmethod + @pytest.mark.parametrize( + ("member_quota", "seat_quota", "new_accounts", "failure"), + [ + ( + LicenseLimitationModel(enabled=True, size=2, limit=4), + LicenseLimitationModel(enabled=True, size=9, limit=10), + 1, + None, + ), + ( + LicenseLimitationModel(enabled=True, size=3, limit=4), + LicenseLimitationModel(enabled=True, size=9, limit=10), + 1, + "members", + ), + (LicenseLimitationModel(), LicenseLimitationModel(enabled=True, size=10, limit=10), 1, "seats"), + (LicenseLimitationModel(), LicenseLimitationModel(enabled=True, size=9, limit=10), 2, "seats"), + (LicenseLimitationModel(), LicenseLimitationModel(enabled=True, size=8, limit=10), 2, None), + (LicenseLimitationModel(), LicenseLimitationModel(enabled=True, size=10, limit=10), 0, None), + ( + LicenseLimitationModel(enabled=True, size=99, limit=0), + LicenseLimitationModel(enabled=True, size=99, limit=0), + 2, + None, + ), + ( + LicenseLimitationModel(enabled=False, size=99, limit=1), + LicenseLimitationModel(enabled=False, size=99, limit=1), + 2, + None, + ), + ], + ids=[ + "exact-capacity", + "workspace-full", + "seats-full", + "not-enough-seats-for-batch", + "exact-seats-for-batch", + "existing-accounts", + "unlimited", + "disabled", + ], ) - license_query = Mock(return_value=LicenseModel(seats=seat_quota)) - monkeypatch.setattr(gateways.SystemFeatureService, "get_license", license_query) + def test_enterprise_invitation_member_and_seat_limits( + invitations: gateways.WorkspaceInvitationGateway, + monkeypatch: pytest.MonkeyPatch, + config_overrides: Callable[..., None], + member_quota: LicenseLimitationModel, + seat_quota: LicenseLimitationModel, + new_accounts: int, + failure: str | None, + ) -> None: + config_overrides(DEPLOYMENT_EDITION=DeploymentEdition.ENTERPRISE) + monkeypatch.setattr( + gateways.FeatureService, "get_features", Mock(return_value=FeatureModel(workspace_members=member_quota)) + ) + license_query = Mock(return_value=LicenseModel(seats=seat_quota)) + monkeypatch.setattr(gateways.SystemFeatureService, "get_license", license_query) - assert invitations.requires_capacity_check - if failure is None: - invitations.check_capacity("workspace", current_members=2, new_members=2, new_accounts=new_accounts) - else: - with pytest.raises(WorkspaceInvitationQuotaError) as error: + assert invitations.requires_capacity_check + if failure is None: invitations.check_capacity("workspace", current_members=2, new_members=2, new_accounts=new_accounts) - assert error.value.seats is (failure == "seats") - assert license_query.call_count == int(new_accounts > 0 and failure != "members") + else: + with pytest.raises(WorkspaceInvitationQuotaError) as error: + invitations.check_capacity("workspace", current_members=2, new_members=2, new_accounts=new_accounts) + assert error.value.seats is (failure == "seats") + assert license_query.call_count == int(new_accounts > 0 and failure != "members") + @staticmethod + def test_community_invitation_does_not_query_quotas( + invitations: gateways.WorkspaceInvitationGateway, + monkeypatch: pytest.MonkeyPatch, + config_overrides: Callable[..., None], + ) -> None: + config_overrides(DEPLOYMENT_EDITION=DeploymentEdition.COMMUNITY) + features = create_autospec(gateways.FeatureService.get_features, instance=True) + license_query = create_autospec(gateways.SystemFeatureService.get_license, instance=True) + monkeypatch.setattr(gateways.FeatureService, "get_features", features) + monkeypatch.setattr(gateways.SystemFeatureService, "get_license", license_query) -def test_community_invitation_does_not_query_quotas( - invitations: gateways.WorkspaceInvitationGateway, - monkeypatch: pytest.MonkeyPatch, - config_overrides: Callable[..., None], -) -> None: - config_overrides(DEPLOYMENT_EDITION=DeploymentEdition.COMMUNITY) - features = create_autospec(gateways.FeatureService.get_features, instance=True) - license_query = create_autospec(gateways.SystemFeatureService.get_license, instance=True) - monkeypatch.setattr(gateways.FeatureService, "get_features", features) - monkeypatch.setattr(gateways.SystemFeatureService, "get_license", license_query) + assert not invitations.requires_capacity_check + invitations.check_capacity("workspace", current_members=2, new_members=2, new_accounts=1) + features.assert_not_called() + license_query.assert_not_called() - assert not invitations.requires_capacity_check - invitations.check_capacity("workspace", current_members=2, new_members=2, new_accounts=1) - features.assert_not_called() - license_query.assert_not_called() - - -@pytest.mark.parametrize( - ("edition", "plan_limit", "license_quota", "error"), - [ - (DeploymentEdition.CLOUD, 2, LicenseLimitationModel(), WorkspaceInvitationQuotaError), - (DeploymentEdition.CLOUD, 3, LicenseLimitationModel(), None), - (DeploymentEdition.CLOUD, 0, LicenseLimitationModel(), None), - (DeploymentEdition.CLOUD, -1, LicenseLimitationModel(), None), - ( - DeploymentEdition.ENTERPRISE, - 2, - LicenseLimitationModel(enabled=True, size=2, limit=2), - WorkspaceMemberLicenseQuotaError, - ), - (DeploymentEdition.ENTERPRISE, 2, LicenseLimitationModel(enabled=True, size=2, limit=3), None), - (DeploymentEdition.ENTERPRISE, 2, LicenseLimitationModel(enabled=True, size=2, limit=0), None), - (DeploymentEdition.COMMUNITY, 2, LicenseLimitationModel(enabled=False, size=2, limit=2), None), - ( - DeploymentEdition.CLOUD, - 3, - LicenseLimitationModel(enabled=True, size=2, limit=2), - WorkspaceMemberLicenseQuotaError, - ), - ], -) -def test_single_invitation_quota_preserves_plan_and_license_limits( - invitations: gateways.WorkspaceInvitationGateway, - monkeypatch: pytest.MonkeyPatch, - config_overrides: Callable[..., None], - edition: DeploymentEdition, - plan_limit: int, - license_quota: LicenseLimitationModel, - error: type[WorkspaceInvitationQuotaError] | None, -) -> None: - config_overrides(DEPLOYMENT_EDITION=edition) - features = FeatureModel(members=LimitationModel(size=2, limit=plan_limit), workspace_members=license_quota) - requested: list[str] = [] - - def get_features(tenant_id: str, *, exclude_vector_space: bool = False) -> FeatureModel: - requested.append(tenant_id) - assert exclude_vector_space - return features - - monkeypatch.setattr(gateways.FeatureService, "get_features", get_features) - - if error is None: - invitations.check_invitation_quota("workspace") - else: - with pytest.raises(error): - invitations.check_invitation_quota("workspace") - assert requested == ["workspace"] - - -def test_owner_transfer_notifies_both_owners(account_domain: AccountDomain, monkeypatch: pytest.MonkeyPatch) -> None: - old_owner = account_domain.accounts.create_account("old@example.com", "Old", "en-US") - new_owner = account_domain.accounts.create_account("new@example.com", "New", "en-US") - old_notification = create_autospec(gateways.send_old_owner_transfer_notify_email_task.delay, instance=True) - new_notification = create_autospec(gateways.send_new_owner_transfer_notify_email_task.delay, instance=True) - monkeypatch.setattr(gateways.send_old_owner_transfer_notify_email_task, "delay", old_notification) - monkeypatch.setattr(gateways.send_new_owner_transfer_notify_email_task, "delay", new_notification) - - gateway = gateways.WorkspaceOwnerTransferGateway(redis=create_autospec(RedisClientWrapper, instance=True)) - gateway.notify(old_owner=old_owner, new_owner=new_owner, workspace_name="Workspace") - - new_notification.assert_called_once_with(language="en-US", to=new_owner.email, workspace="Workspace") - old_notification.assert_called_once_with( - language="en-US", to=old_owner.email, workspace="Workspace", new_owner_email=new_owner.email + @staticmethod + @pytest.mark.parametrize( + ("edition", "plan_limit", "license_quota", "error"), + [ + (DeploymentEdition.CLOUD, 2, LicenseLimitationModel(), WorkspaceInvitationQuotaError), + (DeploymentEdition.CLOUD, 3, LicenseLimitationModel(), None), + (DeploymentEdition.CLOUD, 0, LicenseLimitationModel(), None), + (DeploymentEdition.CLOUD, -1, LicenseLimitationModel(), None), + ( + DeploymentEdition.ENTERPRISE, + 2, + LicenseLimitationModel(enabled=True, size=2, limit=2), + WorkspaceMemberLicenseQuotaError, + ), + (DeploymentEdition.ENTERPRISE, 2, LicenseLimitationModel(enabled=True, size=2, limit=3), None), + (DeploymentEdition.ENTERPRISE, 2, LicenseLimitationModel(enabled=True, size=2, limit=0), None), + (DeploymentEdition.COMMUNITY, 2, LicenseLimitationModel(enabled=False, size=2, limit=2), None), + ( + DeploymentEdition.CLOUD, + 3, + LicenseLimitationModel(enabled=True, size=2, limit=2), + WorkspaceMemberLicenseQuotaError, + ), + ], ) + def test_single_invitation_quota_preserves_plan_and_license_limits( + invitations: gateways.WorkspaceInvitationGateway, + monkeypatch: pytest.MonkeyPatch, + config_overrides: Callable[..., None], + edition: DeploymentEdition, + plan_limit: int, + license_quota: LicenseLimitationModel, + error: type[WorkspaceInvitationQuotaError] | None, + ) -> None: + config_overrides(DEPLOYMENT_EDITION=edition) + features = FeatureModel(members=LimitationModel(size=2, limit=plan_limit), workspace_members=license_quota) + requested: list[str] = [] + + def get_features(tenant_id: str, *, exclude_vector_space: bool = False) -> FeatureModel: + requested.append(tenant_id) + assert exclude_vector_space + return features + + monkeypatch.setattr(gateways.FeatureService, "get_features", get_features) + + if error is None: + invitations.check_invitation_quota("workspace") + else: + with pytest.raises(error): + invitations.check_invitation_quota("workspace") + assert requested == ["workspace"] + + @staticmethod + def test_owner_transfer_notifies_both_owners( + account_domain: AccountDomain, monkeypatch: pytest.MonkeyPatch + ) -> None: + old_owner = account_domain.accounts.create_account("old@example.com", "Old", "en-US") + new_owner = account_domain.accounts.create_account("new@example.com", "New", "en-US") + old_notification = create_autospec(gateways.send_old_owner_transfer_notify_email_task.delay, instance=True) + new_notification = create_autospec(gateways.send_new_owner_transfer_notify_email_task.delay, instance=True) + monkeypatch.setattr(gateways.send_old_owner_transfer_notify_email_task, "delay", old_notification) + monkeypatch.setattr(gateways.send_new_owner_transfer_notify_email_task, "delay", new_notification) + + gateway = gateways.WorkspaceOwnerTransferGateway(redis=create_autospec(RedisClientWrapper, instance=True)) + gateway.notify(old_owner=old_owner, new_owner=new_owner, workspace_name="Workspace") + + new_notification.assert_called_once_with(language="en-US", to=new_owner.email, workspace="Workspace") + old_notification.assert_called_once_with( + language="en-US", to=old_owner.email, workspace="Workspace", new_owner_email=new_owner.email + ) + + +"""Deployment billing adapter and paid-to-trial fallback regressions.""" + + +class TestWorkspaceCreditPoolCases: + @staticmethod + @pytest.mark.parametrize( + ("plan", "paid_quota", "paid_used", "trial_quota", "expected_type", "remaining"), + [ + (CloudPlan.TEAM, 500, 120, 100, "paid", 380), + (CloudPlan.TEAM, -1, 999, 100, "paid", -1), + (CloudPlan.TEAM, 500, 500, 100, "trial", 60), + (CloudPlan.TEAM, None, 0, 100, "trial", 60), + (CloudPlan.SANDBOX, 500, 0, 100, "trial", 60), + (CloudPlan.SANDBOX, None, 0, None, None, None), + ], + ) + def test_effective_credit_pool( + monkeypatch: pytest.MonkeyPatch, + config_overrides: Callable[..., None], + plan: CloudPlan | None, + paid_quota: int | None, + paid_used: int, + trial_quota: int | None, + expected_type: str | None, + remaining: int | None, + ) -> None: + config_overrides(DEPLOYMENT_EDITION=DeploymentEdition.CLOUD) + monkeypatch.setattr( + gateways.BillingService, + "get_info", + lambda *_a, **_k: {"subscription": {"plan": plan}, "next_credit_reset_date": 1775001600}, + ) + paid = CreditPoolBalance("tenant", "paid", paid_quota, paid_used) if paid_quota is not None else None + trial = CreditPoolBalance("tenant", "trial", trial_quota, 40) if trial_quota is not None else None + calls = [] + + def get_pool(*, tenant_id: str, pool_type: str) -> CreditPoolBalance | None: + calls.append((tenant_id, pool_type)) + return paid if pool_type == "paid" else trial + + monkeypatch.setattr(gateways.CreditPoolService, "get_pool", get_pool) + result = DeploymentWorkspaceFeatureGateway().get_effective_credit_pool("tenant") + assert result.pool_type == expected_type + assert result.remaining_credits == remaining + assert result.next_credit_reset_date == 1775001600 + if plan == CloudPlan.SANDBOX: + assert calls == [("tenant", "trial")] + elif expected_type == "paid": + assert calls == [("tenant", "paid")] + else: + assert calls == [("tenant", "paid"), ("tenant", "trial")] + + @staticmethod + @pytest.mark.parametrize( + ("quota", "used", "exhausted_at", "expected"), + [(100, 100, 123, 123), (100, 10, 123, None), (-1, 100, 123, None), (100, 100, 0, None)], + ) + def test_feature_info_uses_same_pool_rules( + monkeypatch: pytest.MonkeyPatch, + config_overrides: Callable[..., None], + quota: int, + used: int, + exhausted_at: int | None, + expected: int | None, + ) -> None: + config_overrides(DEPLOYMENT_EDITION=DeploymentEdition.CLOUD) + features = FeatureModel(can_replace_logo=True) + monkeypatch.setattr(gateways.FeatureService, "get_features", lambda *_a, **_k: features) + monkeypatch.setattr( + gateways.CreditPoolService, + "get_pool", + lambda **_k: CreditPoolBalance("tenant", "trial", quota, used, exhausted_at), + ) + result = DeploymentWorkspaceFeatureGateway().get_features("tenant") + assert result.can_replace_logo + assert result.credits.exhausted_at == expected + + @staticmethod + @pytest.mark.parametrize("edition", [DeploymentEdition.COMMUNITY, DeploymentEdition.ENTERPRISE]) + def test_non_cloud_summary_skips_external_io( + monkeypatch: pytest.MonkeyPatch, config_overrides: Callable[..., None], edition: DeploymentEdition + ) -> None: + config_overrides(DEPLOYMENT_EDITION=edition) + billing = Mock(side_effect=AssertionError("billing must not be queried")) + monkeypatch.setattr(gateways.BillingService, "get_info", billing) + result = DeploymentWorkspaceFeatureGateway().get_effective_credit_pool("tenant") + assert result.plan is None + assert result.remaining_credits is None + billing.assert_not_called() + + @staticmethod + @pytest.mark.parametrize( + ("edition", "bulk", "expected", "fallback_ids"), + [ + (DeploymentEdition.CLOUD, {"w1": {"plan": "team"}}, {"w1": "team", "w2": "professional"}, ["w2"]), + (DeploymentEdition.CLOUD, {}, {"w1": "professional", "w2": "professional"}, ["w1", "w2"]), + (DeploymentEdition.COMMUNITY, {}, {"w1": "professional", "w2": "professional"}, ["w1", "w2"]), + (DeploymentEdition.ENTERPRISE, {}, {"w1": "sandbox", "w2": "sandbox"}, []), + ], + ) + def test_plan_gateway_bulk_and_fallback( + monkeypatch: pytest.MonkeyPatch, + config_overrides: Callable[..., None], + edition: DeploymentEdition, + bulk: dict[str, dict[str, str]], + expected: int | None, + fallback_ids: list[str], + ) -> None: + config_overrides(DEPLOYMENT_EDITION=edition) + bulk_query = Mock(return_value=bulk) + features = FeatureModel() + features.billing.subscription.plan = CloudPlan.PROFESSIONAL + feature_query = Mock(return_value=features) + monkeypatch.setattr(gateways.BillingService, "get_plan_bulk", bulk_query) + monkeypatch.setattr(gateways.FeatureService, "get_features", feature_query) + assert DeploymentWorkspacePlanGateway().resolve_many(["w1", "w2"]) == expected + assert [call.args[0] for call in feature_query.call_args_list] == fallback_ids + assert bulk_query.call_count == int(edition == DeploymentEdition.CLOUD) + + +"""File adapter keeps admitted identity explicit across the FileService boundary.""" + + +class TestWorkspaceFileGatewayCases: + @staticmethod + def test_upload_uses_admitted_account_and_explicit_workspace() -> None: + files = create_autospec(FileService, instance=True) + files.upload_file.return_value = Mock(id="file-id") + gateway = WorkspaceFileGateway(files=files) + context = RequestContext("request", None, "account", "workspace") + + assert gateway.upload(context, filename="logo.png", content=b"image", mimetype="image/png") == "file-id" + + kwargs = files.upload_file.call_args.kwargs + assert isinstance(kwargs["user"], Account) + assert kwargs["user"].id == "account" + assert kwargs["tenant_id"] == "workspace" + assert kwargs["content"] == b"image" + + +def make_subject(account_id: str, *, legacy_role: str = "normal") -> WorkspaceMemberRoleSubject: + return WorkspaceMemberRoleSubject(account_id=account_id, legacy_role=legacy_role) + + +class RoleResolutionError(Exception): + pass + + +class TestWorkspaceMemberRoleResolverCases: + @staticmethod + @pytest.fixture + def transport(monkeypatch: pytest.MonkeyPatch) -> RBACTransport: + transport = RBACTransport() + monkeypatch.setattr(EnterpriseRequest, "send_inner_rbac_request", transport.send) + return transport + + @staticmethod + def test_builtin_mode_projects_join_roles_without_enterprise_call( + config_overrides: Callable[..., None], transport: RBACTransport + ) -> None: + config_overrides(RBAC_ENABLED=False) + + result = DeploymentWorkspaceMemberRoleResolver().resolve_many( + "workspace-1", "actor-1", [make_subject("owner", legacy_role="owner"), make_subject("member")] + ) + + assert result == { + "owner": (WorkspaceMemberRole(id="owner", name="owner"),), + "member": (WorkspaceMemberRole(id="normal", name="normal"),), + } + assert transport.requests == [] + + @staticmethod + @pytest.mark.parametrize("language", [None, "ja"]) + def test_rbac_mode_maps_batch_response_without_builtin_fallback( + config_overrides: Callable[..., None], transport: RBACTransport, language: str | None + ) -> None: + config_overrides(RBAC_ENABLED=True) + transport.response = { + "owner": [ + {"id": "workspace.owner", "name": "Owner", "type": "builtin"}, + {"id": "workspace.editor", "name": "Editor", "type": "builtin"}, + ] + } + + result = DeploymentWorkspaceMemberRoleResolver().resolve_many( + "workspace-1", + "actor-1", + [make_subject("owner", legacy_role="owner"), make_subject("omitted", legacy_role="admin")], + language=language, + ) + + assert result == { + "owner": ( + WorkspaceMemberRole(id="workspace.owner", name="Owner"), + WorkspaceMemberRole(id="workspace.editor", name="Editor"), + ) + } + sent = transport.only_request + assert (sent.method, sent.endpoint) == ("POST", "/rbac/members/rbac-roles/batch") + assert (sent.tenant_id, sent.account_id) == ("workspace-1", "actor-1") + assert sent.json == {"member_ids": ["owner", "omitted"]} + assert sent.params == ({"language": language} if language else None) + + @staticmethod + def test_rbac_failure_propagates(config_overrides: Callable[..., None], transport: RBACTransport) -> None: + config_overrides(RBAC_ENABLED=True) + transport.failure = RoleResolutionError("enterprise unavailable") + + with pytest.raises(RoleResolutionError, match="enterprise unavailable"): + DeploymentWorkspaceMemberRoleResolver().resolve_many("workspace-1", "actor-1", [make_subject("member-1")]) + + @staticmethod + def test_empty_member_list_skips_enterprise_call( + config_overrides: Callable[..., None], transport: RBACTransport + ) -> None: + config_overrides(RBAC_ENABLED=True) + + assert DeploymentWorkspaceMemberRoleResolver().resolve_many("workspace-1", "actor-1", []) == {} + assert transport.requests == [] diff --git a/api/tests/unit_tests/services/workspace/test_member_service.py b/api/tests/unit_tests/services/workspace/test_member_service.py index d66e7e04af9..8528b98604d 100644 --- a/api/tests/unit_tests/services/workspace/test_member_service.py +++ b/api/tests/unit_tests/services/workspace/test_member_service.py @@ -1,8 +1,7 @@ -"""Membership and invitation behavior across the application/persistence boundary.""" - -from collections.abc import Callable +from collections.abc import Callable, Mapping, Sequence from contextlib import nullcontext from dataclasses import dataclass, field +from datetime import datetime from unittest.mock import Mock, create_autospec import pytest @@ -10,9 +9,11 @@ from sqlalchemy import event, select from sqlalchemy.orm import Session, sessionmaker from enums import DeploymentEdition +from enums.account import TenantAccountRole as _workspace_member_query_service_TenantAccountRole from extensions.ext_redis import RedisClientWrapper from machinery.context import RequestContext -from models.account import Account, AccountStatus, Tenant, TenantAccountJoin, TenantAccountRole +from models.account import Account, AccountStatus, Tenant, TenantAccountJoin +from models.account import TenantAccountRole as _member_service_TenantAccountRole from models.dataset import Dataset from models.model import App, AppMode from services.account.adapters import RedisInvitationTokenStore @@ -37,7 +38,13 @@ from services.errors.workspace import ( from services.feature_service import FeatureService from services.rbac.contracts import MemberRolesResponse, Paginated, RBACRole from services.system_feature_service import SystemFeatureService -from services.workspace.contracts import OwnerTransferToken +from services.workspace.contracts import ( + OwnerTransferToken, + WorkspaceMemberRecord, + WorkspaceMemberRole, + WorkspaceMemberRoleSubject, + WorkspaceMemberSummary, +) from services.workspace.gateways import ( DeploymentWorkspaceMemberAccessGateway, WorkspaceInvitationGateway, @@ -45,483 +52,13 @@ from services.workspace.gateways import ( from services.workspace.member_service import ( OwnerTransferGateway, WorkspaceInvitationService, + WorkspaceMemberQueryService, WorkspaceMemberService, WorkspaceOwnerTransferService, ) -from tests.unit_tests.account_domain import AccountDomain +from tests.helpers.unit.account_domain import AccountDomain - -@pytest.fixture -def members(sqlite_session: Session) -> tuple[str, str, str]: - owner = Account(name="Owner", email="owner@example.com") - member = Account(name="Member", email="member@example.com") - tenant = Tenant(name="Workspace") - sqlite_session.add_all( - [ - owner, - member, - tenant, - TenantAccountJoin(account_id=owner.id, tenant_id=tenant.id, role=TenantAccountRole.OWNER, current=True), - TenantAccountJoin(account_id=member.id, tenant_id=tenant.id, role=TenantAccountRole.NORMAL), - ] - ) - sqlite_session.commit() - return tenant.id, owner.id, member.id - - -@pytest.mark.parametrize("existing_member", [True, False]) -def test_join_rejects_a_second_owner_without_changing_membership( - account_domain: AccountDomain, members: tuple[str, str, str], existing_member: bool -) -> None: - workspace_id, owner_id, member_id = members - if existing_member: - account = account_domain.accounts.get_account_by_id(member_id) - assert account is not None - else: - account = account_domain.accounts.create_account("new-owner@example.com", "New owner", "en-US") - - with pytest.raises(WorkspaceAlreadyHasOwnerError, match="Tenant already has an owner"): - account_domain.members.join_member( - workspace_id=workspace_id, - account_id=account.id, - email=account.email, - role=TenantAccountRole.OWNER, - operator_account_id=owner_id, - ) - - assert account_domain.members.get_role(workspace_id, owner_id) == TenantAccountRole.OWNER - assert account_domain.members.get_role(workspace_id, account.id) == ( - TenantAccountRole.NORMAL if existing_member else None - ) - account_domain.access.membership_changed.assert_not_called() - - -@pytest.mark.parametrize( - ("status", "other_workspace", "deleted"), - [("pending", False, True), ("pending", True, False), ("active", False, False)], -) -def test_removal_deletes_only_orphaned_pending_accounts( - account_domain: AccountDomain, - sqlite_session: Session, - members: tuple[str, str, str], - status: str, - other_workspace: bool, - deleted: bool, -) -> None: - workspace_id, owner_id, member_id = members - member = sqlite_session.get(Account, member_id) - assert member is not None - member.status = AccountStatus(status) - if other_workspace: - tenant = Tenant(name="Other") - sqlite_session.add_all( - [tenant, TenantAccountJoin(account_id=member_id, tenant_id=tenant.id, role=TenantAccountRole.NORMAL)] - ) - sqlite_session.commit() - sqlite_session.expunge_all() - account_domain.members.remove(workspace_id, member_id, owner_id) - assert (sqlite_session.get(Account, member_id) is None) == deleted - assert account_domain.members.get_role(workspace_id, member_id) is None - account_domain.access.member_removed.assert_called_once() - assert account_domain.access.member_removed.call_args.args[1].account_deleted == deleted - - -@pytest.mark.parametrize("action", ["add", "remove", "update"]) -def test_self_operations_and_non_administrators_are_denied( - account_domain: AccountDomain, members: tuple[str, str, str], action: str -) -> None: - workspace_id, owner_id, member_id = members - with pytest.raises(CannotOperateSelfError): - account_domain.members.check_permission(workspace_id, owner_id, owner_id, action) - with pytest.raises(NoPermissionError): - account_domain.members.check_permission(workspace_id, member_id, owner_id, action) - - -def test_owner_transfer_demotes_previous_owner(account_domain: AccountDomain, members: tuple[str, str, str]) -> None: - workspace_id, owner_id, member_id = members - account_domain.members.update_role(workspace_id, member_id, "owner", owner_id) - assert account_domain.members.get_role(workspace_id, owner_id) == TenantAccountRole.NORMAL - assert account_domain.members.get_role(workspace_id, member_id) == TenantAccountRole.OWNER - - -@pytest.mark.parametrize("new_role", ["normal", "owner"]) -def test_admin_cannot_change_owner_or_promote_owner( - account_domain: AccountDomain, members: tuple[str, str, str], sqlite_session: Session, new_role: str -) -> None: - workspace_id, owner_id, member_id = members - membership = sqlite_session.scalar(select(TenantAccountJoin).where(TenantAccountJoin.account_id == member_id)) - assert membership is not None - membership.role = TenantAccountRole.ADMIN - sqlite_session.commit() - with pytest.raises(NoPermissionError): - account_domain.members.update_role(workspace_id, owner_id, new_role, member_id) - - -def test_unchanged_role_does_not_sync(account_domain: AccountDomain, members: tuple[str, str, str]) -> None: - workspace_id, owner_id, member_id = members - with pytest.raises(RoleAlreadyAssignedError): - account_domain.members.update_role(workspace_id, member_id, "normal", owner_id) - account_domain.access.change_role.assert_not_called() - - -@pytest.mark.parametrize("role", ["invalid", "dataset_operator"]) -def test_role_policy_prevents_persistence_and_rbac_changes( - account_domain: AccountDomain, - members: tuple[str, str, str], - config_overrides: Callable[..., None], - role: str, -) -> None: - config_overrides(DATASET_OPERATOR_ENABLED=False) - account_domain.access.ensure_role_enabled.side_effect = DeploymentWorkspaceMemberAccessGateway().ensure_role_enabled - workspace_id, owner_id, member_id = members - - with pytest.raises(InvalidWorkspaceMemberRoleError): - account_domain.members.update_role(workspace_id, member_id, role, owner_id) - - assert account_domain.members.get_role(workspace_id, member_id) == TenantAccountRole.NORMAL - account_domain.access.change_role.assert_not_called() - - -@pytest.mark.parametrize("operation", ["remove", "update_role"]) -@pytest.mark.parametrize("account_exists", [True, False]) -def test_member_operations_distinguish_missing_account_from_missing_membership( - account_domain: AccountDomain, members: tuple[str, str, str], operation: str, account_exists: bool -) -> None: - workspace_id, owner_id, _ = members - account_id = "00000000-0000-0000-0000-000000000000" - if account_exists: - account_id = account_domain.accounts.create_account("outsider@example.com", "Outsider", "en-US").id - - error = MemberNotInTenantError if account_exists else AccountNotFoundError - if operation == "remove": - with pytest.raises(error): - account_domain.members.remove(workspace_id, account_id, owner_id) - else: - with pytest.raises(error): - account_domain.members.update_role(workspace_id, account_id, "editor", owner_id) - - account_domain.access.member_removed.assert_not_called() - account_domain.access.change_role.assert_not_called() - - -@pytest.mark.parametrize( - ("permission", "owner", "allowed"), [(True, False, True), (False, False, False), (True, True, False)] -) -def test_rbac_member_removal_uses_permissions_and_owner_protection( - account_domain: AccountDomain, members: tuple[str, str, str], permission: bool, owner: bool, allowed: bool -) -> None: - workspace_id, owner_id, member_id = members - account_domain.access.rbac_enabled = True - account_domain.access.permission_keys.return_value = {"workspace.member.manage"} if permission else set() - account_domain.access.is_owner.return_value = owner - account_domain.access.owner_id.return_value = owner_id - if allowed: - account_domain.members.remove(workspace_id, member_id, owner_id) - assert account_domain.members.get_role(workspace_id, member_id) is None - else: - with pytest.raises(NoPermissionError): - account_domain.members.remove(workspace_id, member_id, owner_id) - - -def test_new_invitation_normalizes_email_and_sets_current_workspace( - account_domain: AccountDomain, members: tuple[str, str, str] -) -> None: - workspace_id, owner_id, _ = members - result = account_domain.invitations.invite( - RequestContext("request", None, owner_id, workspace_id), "New@Example.com", "en-US", "editor" - ) - assert result.token == "invitation-token" - assert result.email == "new@example.com" - assert result.workspace_id == workspace_id - assert result.role == "editor" - invitee = account_domain.repository.find_by_email("new@example.com") - assert invitee is not None - assert result.account_id == invitee.id - assert invitee.status == "pending" - assert account_domain.members.get_role(workspace_id, invitee.id) == TenantAccountRole.EDITOR - membership = account_domain.workspaces.find_membership(invitee.id, workspace_id) - assert membership is not None - assert membership.current - invitation = account_domain.delivery.create.call_args.args[0] - assert invitation.requires_setup is True - assert invitation.email == "new@example.com" - - -def test_active_invitee_must_accept_before_joining( - account_domain: AccountDomain, members: tuple[str, str, str] -) -> None: - workspace_id, owner_id, _ = members - account = account_domain.accounts.create_account("new@example.com", "New", "en-US") - account_domain.invitations.invite(RequestContext("request", None, owner_id, workspace_id), account.email, None) - assert account_domain.members.get_role(workspace_id, account.id) is None - assert account_domain.delivery.create.call_args.args[0].requires_setup is False - - -def test_pending_member_can_be_reinvited( - account_domain: AccountDomain, members: tuple[str, str, str], sqlite_session: Session -) -> None: - workspace_id, owner_id, member_id = members - account = sqlite_session.get(Account, member_id) - assert account is not None - account.status = AccountStatus.PENDING - sqlite_session.commit() - account_domain.invitations.invite(RequestContext("request", None, owner_id, workspace_id), account.email, None) - assert account_domain.delivery.create.call_args.args[0].requires_setup is True - - -def test_active_member_is_reported_without_sending_invitation( - account_domain: AccountDomain, members: tuple[str, str, str] -) -> None: - workspace_id, owner_id, _ = members - with pytest.raises(AccountAlreadyInTenantError): - account_domain.invitations.invite( - RequestContext("request", None, owner_id, workspace_id), "member@example.com", None - ) - account_domain.delivery.send.assert_not_called() - - -def test_removal_transfers_maintainers_only_in_the_target_workspace( - account_domain: AccountDomain, members: tuple[str, str, str], sqlite_session: Session -) -> None: - workspace_id, owner_id, member_id = members - other = Tenant(name="Other workspace") - sqlite_session.add(other) - apps = [ - App( - tenant_id=tenant_id, - name="App", - mode=AppMode.CHAT, - maintainer=member_id, - enable_site=True, - enable_api=True, - max_active_requests=0, - ) - for tenant_id in (workspace_id, other.id) - ] - datasets = [ - Dataset(tenant_id=tenant_id, name="Dataset", created_by=member_id, maintainer=member_id) - for tenant_id in (workspace_id, other.id) - ] - sqlite_session.add_all([*apps, *datasets]) - sqlite_session.commit() - - account_domain.members.remove(workspace_id, member_id, owner_id) - - sqlite_session.expire_all() - assert [app.maintainer for app in apps] == [owner_id, member_id] - assert [dataset.maintainer for dataset in datasets] == [owner_id, member_id] - - -def test_community_invitation_continues_after_a_historical_duplicate_email( - account_domain: AccountDomain, members: tuple[str, str, str], sqlite_session: Session -) -> None: - workspace_id, owner_id, _ = members - sqlite_session.add_all( - [ - Account(name="First", email="duplicate@example.com"), - Account(name="Second", email="duplicate@example.com"), - ] - ) - sqlite_session.commit() - account_domain.delivery.acquire.return_value = nullcontext() - - results = account_domain.invitations.invite_many( - RequestContext("request", None, owner_id, workspace_id), - emails=["duplicate@example.com", "valid@example.com"], - language=None, - role="normal", - ) - - assert [result.status for result in results] == ["failed", "success"] - account_domain.delivery.check_capacity.assert_not_called() - account_domain.delivery.send.assert_called_once() - assert account_domain.repository.find_by_email("valid@example.com") is not None - - -def test_bulk_invitation_counts_new_accounts_and_members_before_writing( - account_domain: AccountDomain, - members: tuple[str, str, str], - monkeypatch: pytest.MonkeyPatch, - config_overrides: Callable[..., None], -) -> None: - config_overrides(DEPLOYMENT_EDITION=DeploymentEdition.ENTERPRISE) - monkeypatch.setattr(FeatureService, "get_features", Mock(return_value=FeatureModel())) - monkeypatch.setattr( - SystemFeatureService, - "get_license", - Mock(return_value=LicenseModel(seats=LicenseLimitationModel(enabled=True, size=3, limit=3))), - ) - redis = create_autospec(RedisClientWrapper, instance=True) - gateway = WorkspaceInvitationGateway(tokens=RedisInvitationTokenStore(redis=redis), redis=redis) - workspace_id, owner_id, _ = members - account_domain.accounts.create_account("outside@example.com", "Outside", "en-US") - account_domain.delivery.acquire.return_value = nullcontext() - account_domain.delivery.check_capacity.side_effect = gateway.check_capacity - account_domain.delivery.requires_capacity_check = True - - with pytest.raises(WorkspaceInvitationQuotaError) as error: - account_domain.invitations.invite_many( - RequestContext("request", None, owner_id, workspace_id), - emails=["new@example.com", "outside@example.com", "member@example.com"], - language=None, - role="normal", - ) - - assert error.value.seats is True - account_domain.delivery.check_capacity.assert_called_once_with( - workspace_id, current_members=2, new_members=2, new_accounts=1 - ) - assert account_domain.repository.find_by_email("new@example.com") is None - account_domain.delivery.create.assert_not_called() - - -def test_bulk_invitation_retains_partial_results_under_workspace_lock( - account_domain: AccountDomain, members: tuple[str, str, str] -) -> None: - workspace_id, owner_id, _ = members - account_domain.delivery.acquire.return_value = nullcontext() - account_domain.delivery.send.side_effect = [None, RuntimeError("Mail unavailable")] - - results = account_domain.invitations.invite_many( - RequestContext("request", None, owner_id, workspace_id), - emails=["new@example.com", "member@example.com", "failed@example.com"], - language=None, - role="normal", - ) - - assert [(result.email, result.status) for result in results] == [ - ("new@example.com", "success"), - ("member@example.com", "already_member"), - ("failed@example.com", "failed"), - ] - assert results[-1].message == "Mail unavailable" - account_domain.delivery.acquire.assert_called_once_with(workspace_id) - - -@pytest.fixture -def owner_transfer(account_domain: AccountDomain) -> tuple[WorkspaceOwnerTransferService, Mock]: - challenges = create_autospec(OwnerTransferGateway, instance=True) - challenges.is_ip_limited.return_value = False - challenges.is_verification_limited.return_value = False - challenges.read_token.return_value = OwnerTransferToken("owner@example.com", "123456") - challenges.issue_token.return_value = "promoted-token" - return WorkspaceOwnerTransferService( - accounts=account_domain.repository, - workspaces=account_domain.workspaces, - members=account_domain.members, - challenges=challenges, - ), challenges - - -@pytest.mark.parametrize( - ("challenge", "code", "error"), - [ - (None, "123456", InvalidOwnerTransferTokenError), - (OwnerTransferToken("other@example.com", "123456"), "123456", InvalidOwnerTransferEmailError), - (OwnerTransferToken("owner@example.com", "123456"), "000000", InvalidOwnerTransferCodeError), - ], -) -def test_owner_verification_failure_preserves_challenge( - owner_transfer: tuple[WorkspaceOwnerTransferService, Mock], - members: tuple[str, str, str], - challenge: OwnerTransferToken | None, - code: str, - error: type[Exception], -) -> None: - service, challenges = owner_transfer - workspace_id, owner_id, _ = members - challenges.read_token.return_value = challenge - with pytest.raises(error): - service.verify_code(RequestContext("request", None, owner_id, workspace_id), token="challenge", code=code) - challenges.revoke_token.assert_not_called() - challenges.issue_token.assert_not_called() - if error is InvalidOwnerTransferCodeError: - challenges.record_verification_failure.assert_called_once_with("owner@example.com") - - -def test_owner_verification_promotes_token_and_clears_failure_count( - owner_transfer: tuple[WorkspaceOwnerTransferService, Mock], - members: tuple[str, str, str], -) -> None: - service, challenges = owner_transfer - workspace_id, owner_id, _ = members - result = service.verify_code( - RequestContext("request", None, owner_id, workspace_id), token="challenge", code="123456" - ) - assert result == ("owner@example.com", "promoted-token") - challenges.revoke_token.assert_called_once_with("challenge") - challenges.reset_verification_failures.assert_called_once_with("owner@example.com") - - -def test_owner_transfer_commits_roles_before_notification( - owner_transfer: tuple[WorkspaceOwnerTransferService, Mock], - members: tuple[str, str, str], - account_domain: AccountDomain, -) -> None: - service, challenges = owner_transfer - workspace_id, owner_id, member_id = members - - def notify(*, old_owner: AccountSnapshot, new_owner: AccountSnapshot, workspace_name: str) -> None: - assert (old_owner.id, new_owner.id, workspace_name) == (owner_id, member_id, "Workspace") - assert account_domain.members.get_role(workspace_id, owner_id) == TenantAccountRole.NORMAL - assert account_domain.members.get_role(workspace_id, member_id) == TenantAccountRole.OWNER - - challenges.notify.side_effect = notify - service.transfer(RequestContext("request", None, owner_id, workspace_id), member_id=member_id, token="challenge") - challenges.notify.assert_called_once() - - -def test_owner_transfer_rejects_accounts_outside_workspace( - owner_transfer: tuple[WorkspaceOwnerTransferService, Mock], - members: tuple[str, str, str], - account_domain: AccountDomain, -) -> None: - service, challenges = owner_transfer - workspace_id, owner_id, _ = members - outside = account_domain.accounts.create_account("outside@example.com", "Outside", "en-US") - with pytest.raises(MemberNotInTenantError): - service.transfer( - RequestContext("request", None, owner_id, workspace_id), member_id=outside.id, token="challenge" - ) - assert account_domain.members.get_role(workspace_id, owner_id) == TenantAccountRole.OWNER - challenges.notify.assert_not_called() - - -@pytest.mark.parametrize("operation", ["send_code", "transfer"]) -def test_owner_transfer_reports_missing_workspace( - owner_transfer: tuple[WorkspaceOwnerTransferService, Mock], - members: tuple[str, str, str], - account_domain: AccountDomain, - monkeypatch: pytest.MonkeyPatch, - operation: str, -) -> None: - service, challenges = owner_transfer - workspace_id, owner_id, member_id = members - monkeypatch.setattr(account_domain.workspaces, "get", lambda _workspace_id: None) - context = RequestContext("request", None, owner_id, workspace_id) - if operation == "send_code": - with pytest.raises(WorkspaceNotFoundError): - service.send_code(context, ip_address="127.0.0.1", language=None) - else: - with pytest.raises(WorkspaceNotFoundError): - service.transfer(context, member_id=member_id, token="challenge") - challenges.send_confirmation.assert_not_called() - challenges.notify.assert_not_called() - assert account_domain.members.get_role(workspace_id, owner_id) == TenantAccountRole.OWNER - - -def test_rbac_transfer_failure_leaves_local_owner_unchanged( - account_domain: AccountDomain, - members: tuple[str, str, str], -) -> None: - workspace_id, owner_id, member_id = members - account_domain.access.rbac_enabled = True - account_domain.access.is_owner.side_effect = lambda _workspace_id, _actor_id, member_id: member_id == owner_id - account_domain.access.transfer_owner.side_effect = RuntimeError("RBAC unavailable") - with pytest.raises(RuntimeError, match="RBAC unavailable"): - account_domain.members.update_role(workspace_id, member_id, "owner", owner_id) - assert account_domain.members.get_role(workspace_id, owner_id) == TenantAccountRole.OWNER - assert account_domain.members.get_role(workspace_id, member_id) == TenantAccountRole.NORMAL +"""Membership and invitation behavior across the application/persistence boundary.""" @dataclass(frozen=True) @@ -551,407 +88,1112 @@ class RecordedRoleAssignments: return MemberRolesResponse(account_id=member_account_id) -@pytest.fixture -def rbac_access( - monkeypatch: pytest.MonkeyPatch, - config_overrides: Callable[..., None], - members: tuple[str, str, str], -) -> tuple[DeploymentWorkspaceMemberAccessGateway, RecordedRoleAssignments]: - config_overrides(RBAC_ENABLED=True, DATASET_OPERATOR_ENABLED=True) - gateway = DeploymentWorkspaceMemberAccessGateway() - roles = [ - RBACRole(id=f"role-{tag}", name=tag, type="", category="global_system_default", role_tag=tag, is_builtin=True) - for tag in [*(role.value for role in TenantAccountRole), "no_access"] - ] - _, owner_id, _ = members - roles_by_tag = {role.role_tag: role for role in roles} +class TestMemberServiceCases: + @staticmethod + @pytest.fixture + def members(sqlite_session: Session) -> tuple[str, str, str]: + owner = Account(name="Owner", email="owner@example.com") + member = Account(name="Member", email="member@example.com") + tenant = Tenant(name="Workspace") + sqlite_session.add_all( + [ + owner, + member, + tenant, + TenantAccountJoin( + account_id=owner.id, tenant_id=tenant.id, role=_member_service_TenantAccountRole.OWNER, current=True + ), + TenantAccountJoin( + account_id=member.id, tenant_id=tenant.id, role=_member_service_TenantAccountRole.NORMAL + ), + ] + ) + sqlite_session.commit() + return tenant.id, owner.id, member.id - def get_roles(tenant_id: str, account_id: str | None, member_account_id: str) -> MemberRolesResponse: - del tenant_id, account_id - tag = "owner" if member_account_id == owner_id else "normal" - return MemberRolesResponse(account_id=member_account_id, roles=[roles_by_tag[tag]]) + @staticmethod + @pytest.mark.parametrize("existing_member", [True, False]) + def test_join_rejects_a_second_owner_without_changing_membership( + account_domain: AccountDomain, members: tuple[str, str, str], existing_member: bool + ) -> None: + workspace_id, owner_id, member_id = members + if existing_member: + account = account_domain.accounts.get_account_by_id(member_id) + assert account is not None + else: + account = account_domain.accounts.create_account("new-owner@example.com", "New owner", "en-US") - monkeypatch.setattr(RBACService.Roles, "list", lambda *_args, **_kwargs: Paginated[RBACRole](data=roles)) - monkeypatch.setattr(RBACService.MemberRoles, "get", get_roles) - monkeypatch.setattr(gateway, "permission_keys", lambda _workspace_id, _actor_id: {"workspace.role.manage"}) - replace = RecordedRoleAssignments() - monkeypatch.setattr(RBACService.MemberRoles, "replace", replace) - return gateway, replace + with pytest.raises(WorkspaceAlreadyHasOwnerError, match="Tenant already has an owner"): + account_domain.members.join_member( + workspace_id=workspace_id, + account_id=account.id, + email=account.email, + role=_member_service_TenantAccountRole.OWNER, + operator_account_id=owner_id, + ) + assert account_domain.members.get_role(workspace_id, owner_id) == _member_service_TenantAccountRole.OWNER + assert account_domain.members.get_role(workspace_id, account.id) == ( + _member_service_TenantAccountRole.NORMAL if existing_member else None + ) + account_domain.access.membership_changed.assert_not_called() -@pytest.mark.parametrize( - ("status", "already_joined"), - [ - (None, False), - (AccountStatus.PENDING, False), - (AccountStatus.PENDING, True), - (AccountStatus.ACTIVE, False), - (AccountStatus.ACTIVE, True), - ], - ids=["new-account", "pending-account", "pending-member", "active-account", "active-member"], -) -def test_rbac_invitation_assigns_requested_role( - account_domain: AccountDomain, - members: tuple[str, str, str], - sqlite_session: Session, - rbac_access: tuple[DeploymentWorkspaceMemberAccessGateway, RecordedRoleAssignments], - monkeypatch: pytest.MonkeyPatch, - status: AccountStatus | None, - already_joined: bool, -) -> None: - from tasks.initialize_created_app_rbac_access_task import sync_joined_workspace_member_rbac_access_task - - workspace_id, owner_id, _ = members - email = "invited@example.com" - if status is not None: - account = Account(name="Invited", email=email, status=status) - sqlite_session.add(account) - if already_joined: - sqlite_session.add( - TenantAccountJoin(tenant_id=workspace_id, account_id=account.id, role=TenantAccountRole.NORMAL) + @staticmethod + @pytest.mark.parametrize( + ("status", "other_workspace", "deleted"), + [("pending", False, True), ("pending", True, False), ("active", False, False)], + ) + def test_removal_deletes_only_orphaned_pending_accounts( + account_domain: AccountDomain, + sqlite_session: Session, + members: tuple[str, str, str], + status: str, + other_workspace: bool, + deleted: bool, + ) -> None: + workspace_id, owner_id, member_id = members + member = sqlite_session.get(Account, member_id) + assert member is not None + member.status = AccountStatus(status) + if other_workspace: + tenant = Tenant(name="Other") + sqlite_session.add_all( + [ + tenant, + TenantAccountJoin( + account_id=member_id, tenant_id=tenant.id, role=_member_service_TenantAccountRole.NORMAL + ), + ] ) sqlite_session.commit() - access, replace = rbac_access - monkeypatch.setattr(access, "permission_keys", lambda _workspace, _actor: {"workspace.member.manage"}) - synced_members: list[tuple[str, str, str | None]] = [] + sqlite_session.expunge_all() + account_domain.members.remove(workspace_id, member_id, owner_id) + assert (sqlite_session.get(Account, member_id) is None) == deleted + assert account_domain.members.get_role(workspace_id, member_id) is None + account_domain.access.member_removed.assert_called_once() + assert account_domain.access.member_removed.call_args.args[1].account_deleted == deleted - def sync(tenant_id: str, member_account_id: str, operator_account_id: str | None = None) -> None: - synced_members.append((tenant_id, member_account_id, operator_account_id)) + @staticmethod + @pytest.mark.parametrize("action", ["add", "remove", "update"]) + def test_self_operations_and_non_administrators_are_denied( + account_domain: AccountDomain, members: tuple[str, str, str], action: str + ) -> None: + workspace_id, owner_id, member_id = members + with pytest.raises(CannotOperateSelfError): + account_domain.members.check_permission(workspace_id, owner_id, owner_id, action) + with pytest.raises(NoPermissionError): + account_domain.members.check_permission(workspace_id, member_id, owner_id, action) - monkeypatch.setattr(sync_joined_workspace_member_rbac_access_task, "delay", sync) - member_service = WorkspaceMemberService( - workspaces=account_domain.workspaces, accounts=account_domain.repository, access=access - ) - invitations = WorkspaceInvitationService( - accounts=account_domain.accounts, - workspaces=account_domain.workspaces, - members=member_service, - invitations=account_domain.delivery, - ) - account_domain.delivery.acquire.return_value = nullcontext() + @staticmethod + def test_owner_transfer_demotes_previous_owner( + account_domain: AccountDomain, members: tuple[str, str, str] + ) -> None: + workspace_id, owner_id, member_id = members + account_domain.members.update_role(workspace_id, member_id, "owner", owner_id) + assert account_domain.members.get_role(workspace_id, owner_id) == _member_service_TenantAccountRole.NORMAL + assert account_domain.members.get_role(workspace_id, member_id) == _member_service_TenantAccountRole.OWNER - results = invitations.invite_many( - RequestContext("request", None, owner_id, workspace_id), - emails=[email], - language="en-US", - role="custom-role-id", - ) + @staticmethod + @pytest.mark.parametrize("new_role", ["normal", "owner"]) + def test_admin_cannot_change_owner_or_promote_owner( + account_domain: AccountDomain, members: tuple[str, str, str], sqlite_session: Session, new_role: str + ) -> None: + workspace_id, owner_id, member_id = members + membership = sqlite_session.scalar(select(TenantAccountJoin).where(TenantAccountJoin.account_id == member_id)) + assert membership is not None + membership.role = _member_service_TenantAccountRole.ADMIN + sqlite_session.commit() + with pytest.raises(NoPermissionError): + account_domain.members.update_role(workspace_id, owner_id, new_role, member_id) - invited = account_domain.repository.find_by_email(email) - assert invited is not None - assert member_service.get_role(workspace_id, invited.id) == TenantAccountRole.NORMAL - if status == AccountStatus.ACTIVE and already_joined: - assert replace.calls == [] - else: - assignments = [item for item in replace.calls if item.tenant_id == workspace_id] - assert len(assignments) == 1 - assert assignments[0].account_id == owner_id - assert assignments[0].member_account_id == invited.id - assert assignments[0].role_ids == ["custom-role-id"] - if status is None: - # Registration also provisions the invitee's own workspace. - assert len(replace.calls) == 2 - personal_workspace = replace.calls[0] - assert personal_workspace.tenant_id != workspace_id - assert personal_workspace.member_account_id == invited.id - assert personal_workspace.role_ids == ["role-owner"] - else: - assert len(replace.calls) == 1 - if status == AccountStatus.ACTIVE: - assert results[0].status == "already_member" - account_domain.delivery.create.assert_not_called() - account_domain.delivery.send.assert_not_called() - else: - assert results[0].status == "success" - assert invited.status == "pending" - invitation = account_domain.delivery.create.call_args.args[0] - assert invitation.account_id == invited.id - assert invitation.role == "custom-role-id" - assert invitation.requires_setup is True - account_domain.delivery.send.assert_called_once_with( - language="en-US", email=email, token="invitation-token", inviter_name="Owner", workspace_name="Workspace" + @staticmethod + def test_unchanged_role_does_not_sync(account_domain: AccountDomain, members: tuple[str, str, str]) -> None: + workspace_id, owner_id, member_id = members + with pytest.raises(RoleAlreadyAssignedError): + account_domain.members.update_role(workspace_id, member_id, "normal", owner_id) + account_domain.access.change_role.assert_not_called() + + @staticmethod + @pytest.mark.parametrize("role", ["invalid", "dataset_operator"]) + def test_role_policy_prevents_persistence_and_rbac_changes( + account_domain: AccountDomain, + members: tuple[str, str, str], + config_overrides: Callable[..., None], + role: str, + ) -> None: + config_overrides(DATASET_OPERATOR_ENABLED=False) + account_domain.access.ensure_role_enabled.side_effect = ( + DeploymentWorkspaceMemberAccessGateway().ensure_role_enabled ) - assert synced_members == ( - [(workspace_id, invited.id, owner_id)] if status == AccountStatus.ACTIVE and not already_joined else [] + workspace_id, owner_id, member_id = members + + with pytest.raises(InvalidWorkspaceMemberRoleError): + account_domain.members.update_role(workspace_id, member_id, role, owner_id) + + assert account_domain.members.get_role(workspace_id, member_id) == _member_service_TenantAccountRole.NORMAL + account_domain.access.change_role.assert_not_called() + + @staticmethod + @pytest.mark.parametrize("operation", ["remove", "update_role"]) + @pytest.mark.parametrize("account_exists", [True, False]) + def test_member_operations_distinguish_missing_account_from_missing_membership( + account_domain: AccountDomain, members: tuple[str, str, str], operation: str, account_exists: bool + ) -> None: + workspace_id, owner_id, _ = members + account_id = "00000000-0000-0000-0000-000000000000" + if account_exists: + account_id = account_domain.accounts.create_account("outsider@example.com", "Outsider", "en-US").id + + error = MemberNotInTenantError if account_exists else AccountNotFoundError + if operation == "remove": + with pytest.raises(error): + account_domain.members.remove(workspace_id, account_id, owner_id) + else: + with pytest.raises(error): + account_domain.members.update_role(workspace_id, account_id, "editor", owner_id) + + account_domain.access.member_removed.assert_not_called() + account_domain.access.change_role.assert_not_called() + + @staticmethod + @pytest.mark.parametrize( + ("permission", "owner", "allowed"), [(True, False, True), (False, False, False), (True, True, False)] ) + def test_rbac_member_removal_uses_permissions_and_owner_protection( + account_domain: AccountDomain, members: tuple[str, str, str], permission: bool, owner: bool, allowed: bool + ) -> None: + workspace_id, owner_id, member_id = members + account_domain.access.rbac_enabled = True + account_domain.access.permission_keys.return_value = {"workspace.member.manage"} if permission else set() + account_domain.access.is_owner.return_value = owner + account_domain.access.owner_id.return_value = owner_id + if allowed: + account_domain.members.remove(workspace_id, member_id, owner_id) + assert account_domain.members.get_role(workspace_id, member_id) is None + else: + with pytest.raises(NoPermissionError): + account_domain.members.remove(workspace_id, member_id, owner_id) + @staticmethod + def test_new_invitation_normalizes_email_and_sets_current_workspace( + account_domain: AccountDomain, members: tuple[str, str, str] + ) -> None: + workspace_id, owner_id, _ = members + result = account_domain.invitations.invite( + RequestContext("request", None, owner_id, workspace_id), "New@Example.com", "en-US", "editor" + ) + assert result.token == "invitation-token" + assert result.email == "new@example.com" + assert result.workspace_id == workspace_id + assert result.role == "editor" + invitee = account_domain.repository.find_by_email("new@example.com") + assert invitee is not None + assert result.account_id == invitee.id + assert invitee.status == "pending" + assert account_domain.members.get_role(workspace_id, invitee.id) == _member_service_TenantAccountRole.EDITOR + membership = account_domain.workspaces.find_membership(invitee.id, workspace_id) + assert membership is not None + assert membership.current + invitation = account_domain.delivery.create.call_args.args[0] + assert invitation.requires_setup is True + assert invitation.email == "new@example.com" -@pytest.mark.parametrize( - "operator_role", - [TenantAccountRole.NORMAL, TenantAccountRole.EDITOR, TenantAccountRole.ADMIN, TenantAccountRole.DATASET_OPERATOR], -) -@pytest.mark.parametrize( - "new_role", - [TenantAccountRole.NORMAL, TenantAccountRole.EDITOR, TenantAccountRole.ADMIN, TenantAccountRole.DATASET_OPERATOR], -) -def test_rbac_role_managers_cannot_demote_owner( - account_domain: AccountDomain, - members: tuple[str, str, str], - sqlite_session: Session, - rbac_access: tuple[DeploymentWorkspaceMemberAccessGateway, RecordedRoleAssignments], - operator_role: TenantAccountRole, - new_role: TenantAccountRole, -) -> None: - workspace_id, owner_id, operator_id = members - membership = sqlite_session.scalar(select(TenantAccountJoin).where(TenantAccountJoin.account_id == operator_id)) - assert membership is not None - membership.role = operator_role - sqlite_session.commit() - gateway, replace = rbac_access - service = WorkspaceMemberService( - workspaces=account_domain.workspaces, accounts=account_domain.repository, access=gateway - ) + @staticmethod + def test_active_invitee_must_accept_before_joining( + account_domain: AccountDomain, members: tuple[str, str, str] + ) -> None: + workspace_id, owner_id, _ = members + account = account_domain.accounts.create_account("new@example.com", "New", "en-US") + account_domain.invitations.invite(RequestContext("request", None, owner_id, workspace_id), account.email, None) + assert account_domain.members.get_role(workspace_id, account.id) is None + assert account_domain.delivery.create.call_args.args[0].requires_setup is False - with pytest.raises(NoPermissionError): - service.update_role(workspace_id, owner_id, new_role.value, operator_id) + @staticmethod + def test_pending_member_can_be_reinvited( + account_domain: AccountDomain, members: tuple[str, str, str], sqlite_session: Session + ) -> None: + workspace_id, owner_id, member_id = members + account = sqlite_session.get(Account, member_id) + assert account is not None + account.status = AccountStatus.PENDING + sqlite_session.commit() + account_domain.invitations.invite(RequestContext("request", None, owner_id, workspace_id), account.email, None) + assert account_domain.delivery.create.call_args.args[0].requires_setup is True - assert replace.calls == [] - assert service.get_role(workspace_id, owner_id) == TenantAccountRole.OWNER - assert service.get_role(workspace_id, operator_id) == operator_role + @staticmethod + def test_active_member_is_reported_without_sending_invitation( + account_domain: AccountDomain, members: tuple[str, str, str] + ) -> None: + workspace_id, owner_id, _ = members + with pytest.raises(AccountAlreadyInTenantError): + account_domain.invitations.invite( + RequestContext("request", None, owner_id, workspace_id), "member@example.com", None + ) + account_domain.delivery.send.assert_not_called() + @staticmethod + def test_removal_transfers_maintainers_only_in_the_target_workspace( + account_domain: AccountDomain, members: tuple[str, str, str], sqlite_session: Session + ) -> None: + workspace_id, owner_id, member_id = members + other = Tenant(name="Other workspace") + sqlite_session.add(other) + apps = [ + App( + tenant_id=tenant_id, + name="App", + mode=AppMode.CHAT, + maintainer=member_id, + enable_site=True, + enable_api=True, + max_active_requests=0, + ) + for tenant_id in (workspace_id, other.id) + ] + datasets = [ + Dataset(tenant_id=tenant_id, name="Dataset", created_by=member_id, maintainer=member_id) + for tenant_id in (workspace_id, other.id) + ] + sqlite_session.add_all([*apps, *datasets]) + sqlite_session.commit() -@pytest.mark.parametrize( - "operator_role", - [TenantAccountRole.NORMAL, TenantAccountRole.EDITOR, TenantAccountRole.ADMIN, TenantAccountRole.DATASET_OPERATOR], -) -def test_rbac_role_managers_cannot_transfer_ownership( - account_domain: AccountDomain, - members: tuple[str, str, str], - sqlite_session: Session, - operator_role: TenantAccountRole, -) -> None: - workspace_id, owner_id, operator_id = members - membership = sqlite_session.scalar(select(TenantAccountJoin).where(TenantAccountJoin.account_id == operator_id)) - assert membership is not None - membership.role = operator_role - target = Account(name="Target", email="target@example.com") - sqlite_session.add_all( - [target, TenantAccountJoin(tenant_id=workspace_id, account_id=target.id, role=TenantAccountRole.NORMAL)] - ) - sqlite_session.commit() - account_domain.access.rbac_enabled = True - account_domain.access.permission_keys.return_value = {"workspace.role.manage"} + account_domain.members.remove(workspace_id, member_id, owner_id) - with pytest.raises(NoPermissionError): - account_domain.members.update_role(workspace_id, target.id, "owner", operator_id) + sqlite_session.expire_all() + assert [app.maintainer for app in apps] == [owner_id, member_id] + assert [dataset.maintainer for dataset in datasets] == [owner_id, member_id] - account_domain.access.transfer_owner.assert_not_called() - account_domain.access.change_role.assert_not_called() - assert account_domain.members.get_role(workspace_id, owner_id) == TenantAccountRole.OWNER - assert account_domain.members.get_role(workspace_id, target.id) == TenantAccountRole.NORMAL + @staticmethod + def test_community_invitation_continues_after_a_historical_duplicate_email( + account_domain: AccountDomain, members: tuple[str, str, str], sqlite_session: Session + ) -> None: + workspace_id, owner_id, _ = members + sqlite_session.add_all( + [ + Account(name="First", email="duplicate@example.com"), + Account(name="Second", email="duplicate@example.com"), + ] + ) + sqlite_session.commit() + account_domain.delivery.acquire.return_value = nullcontext() + results = account_domain.invitations.invite_many( + RequestContext("request", None, owner_id, workspace_id), + emails=["duplicate@example.com", "valid@example.com"], + language=None, + role="normal", + ) -def test_rbac_normal_role_manager_can_update_another_non_owner( - account_domain: AccountDomain, - members: tuple[str, str, str], - sqlite_session: Session, - rbac_access: tuple[DeploymentWorkspaceMemberAccessGateway, RecordedRoleAssignments], -) -> None: - workspace_id, owner_id, operator_id = members - target = Account(name="Target", email="target@example.com") - sqlite_session.add_all( - [target, TenantAccountJoin(tenant_id=workspace_id, account_id=target.id, role=TenantAccountRole.NORMAL)] - ) - sqlite_session.commit() - gateway, replace = rbac_access - service = WorkspaceMemberService( - workspaces=account_domain.workspaces, accounts=account_domain.repository, access=gateway - ) + assert [result.status for result in results] == ["failed", "success"] + account_domain.delivery.check_capacity.assert_not_called() + account_domain.delivery.send.assert_called_once() + assert account_domain.repository.find_by_email("valid@example.com") is not None - service.update_role(workspace_id, target.id, "editor", operator_id) + @staticmethod + def test_bulk_invitation_counts_new_accounts_and_members_before_writing( + account_domain: AccountDomain, + members: tuple[str, str, str], + monkeypatch: pytest.MonkeyPatch, + config_overrides: Callable[..., None], + ) -> None: + config_overrides(DEPLOYMENT_EDITION=DeploymentEdition.ENTERPRISE) + monkeypatch.setattr(FeatureService, "get_features", Mock(return_value=FeatureModel())) + monkeypatch.setattr( + SystemFeatureService, + "get_license", + Mock(return_value=LicenseModel(seats=LicenseLimitationModel(enabled=True, size=3, limit=3))), + ) + redis = create_autospec(RedisClientWrapper, instance=True) + gateway = WorkspaceInvitationGateway(tokens=RedisInvitationTokenStore(redis=redis), redis=redis) + workspace_id, owner_id, _ = members + account_domain.accounts.create_account("outside@example.com", "Outside", "en-US") + account_domain.delivery.acquire.return_value = nullcontext() + account_domain.delivery.check_capacity.side_effect = gateway.check_capacity + account_domain.delivery.requires_capacity_check = True - assert len(replace.calls) == 1 - assert replace.calls[0].member_account_id == target.id - assert replace.calls[0].role_ids == ["role-editor"] - assert service.get_role(workspace_id, target.id) == TenantAccountRole.EDITOR - assert service.get_role(workspace_id, owner_id) == TenantAccountRole.OWNER + with pytest.raises(WorkspaceInvitationQuotaError) as error: + account_domain.invitations.invite_many( + RequestContext("request", None, owner_id, workspace_id), + emails=["new@example.com", "outside@example.com", "member@example.com"], + language=None, + role="normal", + ) + assert error.value.seats is True + account_domain.delivery.check_capacity.assert_called_once_with( + workspace_id, current_members=2, new_members=2, new_accounts=1 + ) + assert account_domain.repository.find_by_email("new@example.com") is None + account_domain.delivery.create.assert_not_called() -@pytest.mark.parametrize( - "role", - [TenantAccountRole.ADMIN, TenantAccountRole.EDITOR, TenantAccountRole.NORMAL, TenantAccountRole.DATASET_OPERATOR], -) -@pytest.mark.parametrize("remote_fails", [False, True]) -def test_rbac_role_update_assigns_requested_role_and_synchronizes_membership( - account_domain: AccountDomain, - members: tuple[str, str, str], - sqlite_session: Session, - rbac_access: tuple[DeploymentWorkspaceMemberAccessGateway, RecordedRoleAssignments], - role: TenantAccountRole, - remote_fails: bool, -) -> None: - workspace_id, owner_id, member_id = members - previous = TenantAccountRole.EDITOR if role == TenantAccountRole.NORMAL else TenantAccountRole.NORMAL - membership = sqlite_session.scalar(select(TenantAccountJoin).where(TenantAccountJoin.account_id == member_id)) - assert membership is not None - membership.role = previous - sqlite_session.commit() - gateway, replace = rbac_access - service = WorkspaceMemberService( - workspaces=account_domain.workspaces, accounts=account_domain.repository, access=gateway - ) - if remote_fails: - replace.failure = RuntimeError("RBAC unavailable") - with pytest.raises(RuntimeError, match="RBAC unavailable"): - service.update_role(workspace_id, member_id, role.value, owner_id) - else: - service.update_role(workspace_id, member_id, role.value, owner_id) + @staticmethod + def test_bulk_invitation_retains_partial_results_under_workspace_lock( + account_domain: AccountDomain, members: tuple[str, str, str] + ) -> None: + workspace_id, owner_id, _ = members + account_domain.delivery.acquire.return_value = nullcontext() + account_domain.delivery.send.side_effect = [None, RuntimeError("Mail unavailable")] - assert len(replace.calls) == 1 - assert replace.calls[0].role_ids == [f"role-{role.value}"] - assert replace.calls[0].member_account_id == member_id - assert service.get_role(workspace_id, member_id) == (previous if remote_fails else role) - assert service.get_role(workspace_id, owner_id) == TenantAccountRole.OWNER + results = account_domain.invitations.invite_many( + RequestContext("request", None, owner_id, workspace_id), + emails=["new@example.com", "member@example.com", "failed@example.com"], + language=None, + role="normal", + ) + assert [(result.email, result.status) for result in results] == [ + ("new@example.com", "success"), + ("member@example.com", "already_member"), + ("failed@example.com", "failed"), + ] + assert results[-1].message == "Mail unavailable" + account_domain.delivery.acquire.assert_called_once_with(workspace_id) -def test_rbac_owner_transfer_preserves_other_roles_and_updates_both_memberships( - account_domain: AccountDomain, - members: tuple[str, str, str], - rbac_access: tuple[DeploymentWorkspaceMemberAccessGateway, RecordedRoleAssignments], - monkeypatch: pytest.MonkeyPatch, -) -> None: - workspace_id, owner_id, member_id = members - gateway, replace = rbac_access - monkeypatch.setattr(gateway, "owner_id", lambda _workspace_id, _actor_id: owner_id) - current_roles = MemberRolesResponse( - account_id=owner_id, - roles=[ - RBACRole( - id="role-owner", - name="Owner", - type="", - category="global_system_default", - role_tag="owner", - is_builtin=True, - ), - RBACRole(id="custom", name="Custom", type=""), + @staticmethod + @pytest.fixture + def owner_transfer(account_domain: AccountDomain) -> tuple[WorkspaceOwnerTransferService, Mock]: + challenges = create_autospec(OwnerTransferGateway, instance=True) + challenges.is_ip_limited.return_value = False + challenges.is_verification_limited.return_value = False + challenges.read_token.return_value = OwnerTransferToken("owner@example.com", "123456") + challenges.issue_token.return_value = "promoted-token" + return WorkspaceOwnerTransferService( + accounts=account_domain.repository, + workspaces=account_domain.workspaces, + members=account_domain.members, + challenges=challenges, + ), challenges + + @staticmethod + @pytest.mark.parametrize( + ("challenge", "code", "error"), + [ + (None, "123456", InvalidOwnerTransferTokenError), + (OwnerTransferToken("other@example.com", "123456"), "123456", InvalidOwnerTransferEmailError), + (OwnerTransferToken("owner@example.com", "123456"), "000000", InvalidOwnerTransferCodeError), ], ) + def test_owner_verification_failure_preserves_challenge( + owner_transfer: tuple[WorkspaceOwnerTransferService, Mock], + members: tuple[str, str, str], + challenge: OwnerTransferToken | None, + code: str, + error: type[Exception], + ) -> None: + service, challenges = owner_transfer + workspace_id, owner_id, _ = members + challenges.read_token.return_value = challenge + with pytest.raises(error): + service.verify_code(RequestContext("request", None, owner_id, workspace_id), token="challenge", code=code) + challenges.revoke_token.assert_not_called() + challenges.issue_token.assert_not_called() + if error is InvalidOwnerTransferCodeError: + challenges.record_verification_failure.assert_called_once_with("owner@example.com") - def get_roles(tenant_id: str, account_id: str | None, member_account_id: str) -> MemberRolesResponse: - del tenant_id, account_id - return current_roles if member_account_id == owner_id else MemberRolesResponse(account_id=member_account_id) + @staticmethod + def test_owner_verification_promotes_token_and_clears_failure_count( + owner_transfer: tuple[WorkspaceOwnerTransferService, Mock], + members: tuple[str, str, str], + ) -> None: + service, challenges = owner_transfer + workspace_id, owner_id, _ = members + result = service.verify_code( + RequestContext("request", None, owner_id, workspace_id), token="challenge", code="123456" + ) + assert result == ("owner@example.com", "promoted-token") + challenges.revoke_token.assert_called_once_with("challenge") + challenges.reset_verification_failures.assert_called_once_with("owner@example.com") - monkeypatch.setattr(RBACService.MemberRoles, "get", get_roles) - service = WorkspaceMemberService( - workspaces=account_domain.workspaces, accounts=account_domain.repository, access=gateway - ) + @staticmethod + def test_owner_transfer_commits_roles_before_notification( + owner_transfer: tuple[WorkspaceOwnerTransferService, Mock], + members: tuple[str, str, str], + account_domain: AccountDomain, + ) -> None: + service, challenges = owner_transfer + workspace_id, owner_id, member_id = members - service.update_role(workspace_id, member_id, "owner", owner_id) + def notify(*, old_owner: AccountSnapshot, new_owner: AccountSnapshot, workspace_name: str) -> None: + assert (old_owner.id, new_owner.id, workspace_name) == (owner_id, member_id, "Workspace") + assert account_domain.members.get_role(workspace_id, owner_id) == _member_service_TenantAccountRole.NORMAL + assert account_domain.members.get_role(workspace_id, member_id) == _member_service_TenantAccountRole.OWNER - assert [(call.member_account_id, call.role_ids) for call in replace.calls] == [ - (owner_id, ["custom"]), - (member_id, ["role-owner"]), - ] - assert service.get_role(workspace_id, owner_id) == TenantAccountRole.NORMAL - assert service.get_role(workspace_id, member_id) == TenantAccountRole.OWNER + challenges.notify.side_effect = notify + service.transfer( + RequestContext("request", None, owner_id, workspace_id), member_id=member_id, token="challenge" + ) + challenges.notify.assert_called_once() + @staticmethod + def test_owner_transfer_rejects_accounts_outside_workspace( + owner_transfer: tuple[WorkspaceOwnerTransferService, Mock], + members: tuple[str, str, str], + account_domain: AccountDomain, + ) -> None: + service, challenges = owner_transfer + workspace_id, owner_id, _ = members + outside = account_domain.accounts.create_account("outside@example.com", "Outside", "en-US") + with pytest.raises(MemberNotInTenantError): + service.transfer( + RequestContext("request", None, owner_id, workspace_id), member_id=outside.id, token="challenge" + ) + assert account_domain.members.get_role(workspace_id, owner_id) == _member_service_TenantAccountRole.OWNER + challenges.notify.assert_not_called() -@pytest.mark.parametrize("operation", ["demote_new_owner", "transfer_to_another_member"]) -def test_rbac_owner_is_protected_after_local_transfer_commit_fails( - account_domain: AccountDomain, - members: tuple[str, str, str], - sqlite_session: Session, - sqlite_session_factory: sessionmaker[Session], - rbac_access: tuple[DeploymentWorkspaceMemberAccessGateway, RecordedRoleAssignments], - monkeypatch: pytest.MonkeyPatch, - operation: str, -) -> None: - workspace_id, old_owner_id, new_owner_id = members - third_member = Account(name="Third", email="third@example.com") - sqlite_session.add_all( + @staticmethod + @pytest.mark.parametrize("operation", ["send_code", "transfer"]) + def test_owner_transfer_reports_missing_workspace( + owner_transfer: tuple[WorkspaceOwnerTransferService, Mock], + members: tuple[str, str, str], + account_domain: AccountDomain, + monkeypatch: pytest.MonkeyPatch, + operation: str, + ) -> None: + service, challenges = owner_transfer + workspace_id, owner_id, member_id = members + monkeypatch.setattr(account_domain.workspaces, "get", lambda _workspace_id: None) + context = RequestContext("request", None, owner_id, workspace_id) + if operation == "send_code": + with pytest.raises(WorkspaceNotFoundError): + service.send_code(context, ip_address="127.0.0.1", language=None) + else: + with pytest.raises(WorkspaceNotFoundError): + service.transfer(context, member_id=member_id, token="challenge") + challenges.send_confirmation.assert_not_called() + challenges.notify.assert_not_called() + assert account_domain.members.get_role(workspace_id, owner_id) == _member_service_TenantAccountRole.OWNER + + @staticmethod + def test_rbac_transfer_failure_leaves_local_owner_unchanged( + account_domain: AccountDomain, + members: tuple[str, str, str], + ) -> None: + workspace_id, owner_id, member_id = members + account_domain.access.rbac_enabled = True + account_domain.access.is_owner.side_effect = lambda _workspace_id, _actor_id, member_id: member_id == owner_id + account_domain.access.transfer_owner.side_effect = RuntimeError("RBAC unavailable") + with pytest.raises(RuntimeError, match="RBAC unavailable"): + account_domain.members.update_role(workspace_id, member_id, "owner", owner_id) + assert account_domain.members.get_role(workspace_id, owner_id) == _member_service_TenantAccountRole.OWNER + assert account_domain.members.get_role(workspace_id, member_id) == _member_service_TenantAccountRole.NORMAL + + @staticmethod + @pytest.fixture + def rbac_access( + monkeypatch: pytest.MonkeyPatch, + config_overrides: Callable[..., None], + members: tuple[str, str, str], + ) -> tuple[DeploymentWorkspaceMemberAccessGateway, RecordedRoleAssignments]: + config_overrides(RBAC_ENABLED=True, DATASET_OPERATOR_ENABLED=True) + gateway = DeploymentWorkspaceMemberAccessGateway() + roles = [ + RBACRole( + id=f"role-{tag}", name=tag, type="", category="global_system_default", role_tag=tag, is_builtin=True + ) + for tag in [*(role.value for role in _member_service_TenantAccountRole), "no_access"] + ] + _, owner_id, _ = members + roles_by_tag = {role.role_tag: role for role in roles} + + def get_roles(tenant_id: str, account_id: str | None, member_account_id: str) -> MemberRolesResponse: + del tenant_id, account_id + tag = "owner" if member_account_id == owner_id else "normal" + return MemberRolesResponse(account_id=member_account_id, roles=[roles_by_tag[tag]]) + + monkeypatch.setattr(RBACService.Roles, "list", lambda *_args, **_kwargs: Paginated[RBACRole](data=roles)) + monkeypatch.setattr(RBACService.MemberRoles, "get", get_roles) + monkeypatch.setattr(gateway, "permission_keys", lambda _workspace_id, _actor_id: {"workspace.role.manage"}) + replace = RecordedRoleAssignments() + monkeypatch.setattr(RBACService.MemberRoles, "replace", replace) + return gateway, replace + + @staticmethod + @pytest.mark.parametrize( + ("status", "already_joined"), [ - third_member, - TenantAccountJoin(tenant_id=workspace_id, account_id=third_member.id, role=TenantAccountRole.NORMAL), + (None, False), + (AccountStatus.PENDING, False), + (AccountStatus.PENDING, True), + (AccountStatus.ACTIVE, False), + (AccountStatus.ACTIVE, True), + ], + ids=["new-account", "pending-account", "pending-member", "active-account", "active-member"], + ) + def test_rbac_invitation_assigns_requested_role( + account_domain: AccountDomain, + members: tuple[str, str, str], + sqlite_session: Session, + rbac_access: tuple[DeploymentWorkspaceMemberAccessGateway, RecordedRoleAssignments], + monkeypatch: pytest.MonkeyPatch, + status: AccountStatus | None, + already_joined: bool, + ) -> None: + from tasks.initialize_created_app_rbac_access_task import sync_joined_workspace_member_rbac_access_task + + workspace_id, owner_id, _ = members + email = "invited@example.com" + if status is not None: + account = Account(name="Invited", email=email, status=status) + sqlite_session.add(account) + if already_joined: + sqlite_session.add( + TenantAccountJoin( + tenant_id=workspace_id, account_id=account.id, role=_member_service_TenantAccountRole.NORMAL + ) + ) + sqlite_session.commit() + access, replace = rbac_access + monkeypatch.setattr(access, "permission_keys", lambda _workspace, _actor: {"workspace.member.manage"}) + synced_members: list[tuple[str, str, str | None]] = [] + + def sync(tenant_id: str, member_account_id: str, operator_account_id: str | None = None) -> None: + synced_members.append((tenant_id, member_account_id, operator_account_id)) + + monkeypatch.setattr(sync_joined_workspace_member_rbac_access_task, "delay", sync) + member_service = WorkspaceMemberService( + workspaces=account_domain.workspaces, accounts=account_domain.repository, access=access + ) + invitations = WorkspaceInvitationService( + accounts=account_domain.accounts, + workspaces=account_domain.workspaces, + members=member_service, + invitations=account_domain.delivery, + ) + account_domain.delivery.acquire.return_value = nullcontext() + + results = invitations.invite_many( + RequestContext("request", None, owner_id, workspace_id), + emails=[email], + language="en-US", + role="custom-role-id", + ) + + invited = account_domain.repository.find_by_email(email) + assert invited is not None + assert member_service.get_role(workspace_id, invited.id) == _member_service_TenantAccountRole.NORMAL + if status == AccountStatus.ACTIVE and already_joined: + assert replace.calls == [] + else: + assignments = [item for item in replace.calls if item.tenant_id == workspace_id] + assert len(assignments) == 1 + assert assignments[0].account_id == owner_id + assert assignments[0].member_account_id == invited.id + assert assignments[0].role_ids == ["custom-role-id"] + if status is None: + # Registration also provisions the invitee's own workspace. + assert len(replace.calls) == 2 + personal_workspace = replace.calls[0] + assert personal_workspace.tenant_id != workspace_id + assert personal_workspace.member_account_id == invited.id + assert personal_workspace.role_ids == ["role-owner"] + else: + assert len(replace.calls) == 1 + if status == AccountStatus.ACTIVE: + assert results[0].status == "already_member" + account_domain.delivery.create.assert_not_called() + account_domain.delivery.send.assert_not_called() + else: + assert results[0].status == "success" + assert invited.status == "pending" + invitation = account_domain.delivery.create.call_args.args[0] + assert invitation.account_id == invited.id + assert invitation.role == "custom-role-id" + assert invitation.requires_setup is True + account_domain.delivery.send.assert_called_once_with( + language="en-US", + email=email, + token="invitation-token", + inviter_name="Owner", + workspace_name="Workspace", + ) + assert synced_members == ( + [(workspace_id, invited.id, owner_id)] if status == AccountStatus.ACTIVE and not already_joined else [] + ) + + @staticmethod + @pytest.mark.parametrize( + "operator_role", + [ + _member_service_TenantAccountRole.NORMAL, + _member_service_TenantAccountRole.EDITOR, + _member_service_TenantAccountRole.ADMIN, + _member_service_TenantAccountRole.DATASET_OPERATOR, + ], + ) + @pytest.mark.parametrize( + "new_role", + [ + _member_service_TenantAccountRole.NORMAL, + _member_service_TenantAccountRole.EDITOR, + _member_service_TenantAccountRole.ADMIN, + _member_service_TenantAccountRole.DATASET_OPERATOR, + ], + ) + def test_rbac_role_managers_cannot_demote_owner( + account_domain: AccountDomain, + members: tuple[str, str, str], + sqlite_session: Session, + rbac_access: tuple[DeploymentWorkspaceMemberAccessGateway, RecordedRoleAssignments], + operator_role: _member_service_TenantAccountRole, + new_role: _member_service_TenantAccountRole, + ) -> None: + workspace_id, owner_id, operator_id = members + membership = sqlite_session.scalar(select(TenantAccountJoin).where(TenantAccountJoin.account_id == operator_id)) + assert membership is not None + membership.role = operator_role + sqlite_session.commit() + gateway, replace = rbac_access + service = WorkspaceMemberService( + workspaces=account_domain.workspaces, accounts=account_domain.repository, access=gateway + ) + + with pytest.raises(NoPermissionError): + service.update_role(workspace_id, owner_id, new_role.value, operator_id) + + assert replace.calls == [] + assert service.get_role(workspace_id, owner_id) == _member_service_TenantAccountRole.OWNER + assert service.get_role(workspace_id, operator_id) == operator_role + + @staticmethod + @pytest.mark.parametrize( + "operator_role", + [ + _member_service_TenantAccountRole.NORMAL, + _member_service_TenantAccountRole.EDITOR, + _member_service_TenantAccountRole.ADMIN, + _member_service_TenantAccountRole.DATASET_OPERATOR, + ], + ) + def test_rbac_role_managers_cannot_transfer_ownership( + account_domain: AccountDomain, + members: tuple[str, str, str], + sqlite_session: Session, + operator_role: _member_service_TenantAccountRole, + ) -> None: + workspace_id, owner_id, operator_id = members + membership = sqlite_session.scalar(select(TenantAccountJoin).where(TenantAccountJoin.account_id == operator_id)) + assert membership is not None + membership.role = operator_role + target = Account(name="Target", email="target@example.com") + sqlite_session.add_all( + [ + target, + TenantAccountJoin( + tenant_id=workspace_id, account_id=target.id, role=_member_service_TenantAccountRole.NORMAL + ), + ] + ) + sqlite_session.commit() + account_domain.access.rbac_enabled = True + account_domain.access.permission_keys.return_value = {"workspace.role.manage"} + + with pytest.raises(NoPermissionError): + account_domain.members.update_role(workspace_id, target.id, "owner", operator_id) + + account_domain.access.transfer_owner.assert_not_called() + account_domain.access.change_role.assert_not_called() + assert account_domain.members.get_role(workspace_id, owner_id) == _member_service_TenantAccountRole.OWNER + assert account_domain.members.get_role(workspace_id, target.id) == _member_service_TenantAccountRole.NORMAL + + @staticmethod + def test_rbac_normal_role_manager_can_update_another_non_owner( + account_domain: AccountDomain, + members: tuple[str, str, str], + sqlite_session: Session, + rbac_access: tuple[DeploymentWorkspaceMemberAccessGateway, RecordedRoleAssignments], + ) -> None: + workspace_id, owner_id, operator_id = members + target = Account(name="Target", email="target@example.com") + sqlite_session.add_all( + [ + target, + TenantAccountJoin( + tenant_id=workspace_id, account_id=target.id, role=_member_service_TenantAccountRole.NORMAL + ), + ] + ) + sqlite_session.commit() + gateway, replace = rbac_access + service = WorkspaceMemberService( + workspaces=account_domain.workspaces, accounts=account_domain.repository, access=gateway + ) + + service.update_role(workspace_id, target.id, "editor", operator_id) + + assert len(replace.calls) == 1 + assert replace.calls[0].member_account_id == target.id + assert replace.calls[0].role_ids == ["role-editor"] + assert service.get_role(workspace_id, target.id) == _member_service_TenantAccountRole.EDITOR + assert service.get_role(workspace_id, owner_id) == _member_service_TenantAccountRole.OWNER + + @staticmethod + @pytest.mark.parametrize( + "role", + [ + _member_service_TenantAccountRole.ADMIN, + _member_service_TenantAccountRole.EDITOR, + _member_service_TenantAccountRole.NORMAL, + _member_service_TenantAccountRole.DATASET_OPERATOR, + ], + ) + @pytest.mark.parametrize("remote_fails", [False, True]) + def test_rbac_role_update_assigns_requested_role_and_synchronizes_membership( + account_domain: AccountDomain, + members: tuple[str, str, str], + sqlite_session: Session, + rbac_access: tuple[DeploymentWorkspaceMemberAccessGateway, RecordedRoleAssignments], + role: _member_service_TenantAccountRole, + remote_fails: bool, + ) -> None: + workspace_id, owner_id, member_id = members + previous = ( + _member_service_TenantAccountRole.EDITOR + if role == _member_service_TenantAccountRole.NORMAL + else _member_service_TenantAccountRole.NORMAL + ) + membership = sqlite_session.scalar(select(TenantAccountJoin).where(TenantAccountJoin.account_id == member_id)) + assert membership is not None + membership.role = previous + sqlite_session.commit() + gateway, replace = rbac_access + service = WorkspaceMemberService( + workspaces=account_domain.workspaces, accounts=account_domain.repository, access=gateway + ) + if remote_fails: + replace.failure = RuntimeError("RBAC unavailable") + with pytest.raises(RuntimeError, match="RBAC unavailable"): + service.update_role(workspace_id, member_id, role.value, owner_id) + else: + service.update_role(workspace_id, member_id, role.value, owner_id) + + assert len(replace.calls) == 1 + assert replace.calls[0].role_ids == [f"role-{role.value}"] + assert replace.calls[0].member_account_id == member_id + assert service.get_role(workspace_id, member_id) == (previous if remote_fails else role) + assert service.get_role(workspace_id, owner_id) == _member_service_TenantAccountRole.OWNER + + @staticmethod + def test_rbac_owner_transfer_preserves_other_roles_and_updates_both_memberships( + account_domain: AccountDomain, + members: tuple[str, str, str], + rbac_access: tuple[DeploymentWorkspaceMemberAccessGateway, RecordedRoleAssignments], + monkeypatch: pytest.MonkeyPatch, + ) -> None: + workspace_id, owner_id, member_id = members + gateway, replace = rbac_access + monkeypatch.setattr(gateway, "owner_id", lambda _workspace_id, _actor_id: owner_id) + current_roles = MemberRolesResponse( + account_id=owner_id, + roles=[ + RBACRole( + id="role-owner", + name="Owner", + type="", + category="global_system_default", + role_tag="owner", + is_builtin=True, + ), + RBACRole(id="custom", name="Custom", type=""), + ], + ) + + def get_roles(tenant_id: str, account_id: str | None, member_account_id: str) -> MemberRolesResponse: + del tenant_id, account_id + return current_roles if member_account_id == owner_id else MemberRolesResponse(account_id=member_account_id) + + monkeypatch.setattr(RBACService.MemberRoles, "get", get_roles) + service = WorkspaceMemberService( + workspaces=account_domain.workspaces, accounts=account_domain.repository, access=gateway + ) + + service.update_role(workspace_id, member_id, "owner", owner_id) + + assert [(call.member_account_id, call.role_ids) for call in replace.calls] == [ + (owner_id, ["custom"]), + (member_id, ["role-owner"]), ] - ) - sqlite_session.commit() - gateway, replace = rbac_access - owner_role = RBACRole( - id="role-owner", name="Owner", type="", category="global_system_default", role_tag="owner", is_builtin=True - ) - custom_role = RBACRole(id="custom-manager", name="Role manager", type="") - remote_roles = {old_owner_id: [owner_role, custom_role], new_owner_id: [], third_member.id: []} - roles_by_id = {role.id: role for role in [owner_role, custom_role]} + assert service.get_role(workspace_id, owner_id) == _member_service_TenantAccountRole.NORMAL + assert service.get_role(workspace_id, member_id) == _member_service_TenantAccountRole.OWNER - def get_roles(tenant_id: str, account_id: str | None, member_account_id: str) -> MemberRolesResponse: - del tenant_id, account_id - return MemberRolesResponse(account_id=member_account_id, roles=remote_roles[member_account_id]) + @staticmethod + @pytest.mark.parametrize("operation", ["demote_new_owner", "transfer_to_another_member"]) + def test_rbac_owner_is_protected_after_local_transfer_commit_fails( + account_domain: AccountDomain, + members: tuple[str, str, str], + sqlite_session: Session, + sqlite_session_factory: sessionmaker[Session], + rbac_access: tuple[DeploymentWorkspaceMemberAccessGateway, RecordedRoleAssignments], + monkeypatch: pytest.MonkeyPatch, + operation: str, + ) -> None: + workspace_id, old_owner_id, new_owner_id = members + third_member = Account(name="Third", email="third@example.com") + sqlite_session.add_all( + [ + third_member, + TenantAccountJoin( + tenant_id=workspace_id, account_id=third_member.id, role=_member_service_TenantAccountRole.NORMAL + ), + ] + ) + sqlite_session.commit() + gateway, replace = rbac_access + owner_role = RBACRole( + id="role-owner", name="Owner", type="", category="global_system_default", role_tag="owner", is_builtin=True + ) + custom_role = RBACRole(id="custom-manager", name="Role manager", type="") + remote_roles = {old_owner_id: [owner_role, custom_role], new_owner_id: [], third_member.id: []} + roles_by_id = {role.id: role for role in [owner_role, custom_role]} - def replace_roles(*, tenant_id: str, account_id: str, member_account_id: str, role_ids: list[str]) -> None: - del tenant_id, account_id - remote_roles[member_account_id] = [ - roles_by_id.get(role_id, RBACRole(id=role_id, name=role_id, type="")) for role_id in role_ids + def get_roles(tenant_id: str, account_id: str | None, member_account_id: str) -> MemberRolesResponse: + del tenant_id, account_id + return MemberRolesResponse(account_id=member_account_id, roles=remote_roles[member_account_id]) + + def replace_roles(*, tenant_id: str, account_id: str, member_account_id: str, role_ids: list[str]) -> None: + del tenant_id, account_id + remote_roles[member_account_id] = [ + roles_by_id.get(role_id, RBACRole(id=role_id, name=role_id, type="")) for role_id in role_ids + ] + + def fail_commit(session: Session) -> None: + session.flush() + raise RuntimeError("Membership commit failed") + + monkeypatch.setattr(RBACService.MemberRoles, "get", get_roles) + monkeypatch.setattr( + gateway, + "owner_id", + lambda _workspace_id, _actor_id: next( + member_id for member_id, roles in remote_roles.items() if owner_role in roles + ), + ) + replace.on_apply = replace_roles + service = WorkspaceMemberService( + workspaces=account_domain.workspaces, accounts=account_domain.repository, access=gateway + ) + + event.listen(sqlite_session_factory, "before_commit", fail_commit) + try: + with pytest.raises(RuntimeError, match="Membership commit failed"): + service.update_role(workspace_id, new_owner_id, "owner", old_owner_id) + finally: + event.remove(sqlite_session_factory, "before_commit", fail_commit) + + assert remote_roles[old_owner_id] == [custom_role] + assert remote_roles[new_owner_id] == [owner_role] + assert service.get_role(workspace_id, old_owner_id) == _member_service_TenantAccountRole.OWNER + assert service.get_role(workspace_id, new_owner_id) == _member_service_TenantAccountRole.NORMAL + replace.calls.clear() + + target_id, role = (new_owner_id, "editor") if operation == "demote_new_owner" else (third_member.id, "owner") + with pytest.raises(NoPermissionError): + service.update_role(workspace_id, target_id, role, old_owner_id) + + assert replace.calls == [] + assert remote_roles[old_owner_id] == [custom_role] + assert remote_roles[new_owner_id] == [owner_role] + assert remote_roles[third_member.id] == [] + assert service.get_role(workspace_id, old_owner_id) == _member_service_TenantAccountRole.OWNER + assert service.get_role(workspace_id, new_owner_id) == _member_service_TenantAccountRole.NORMAL + assert service.get_role(workspace_id, third_member.id) == _member_service_TenantAccountRole.NORMAL + + @staticmethod + @pytest.mark.parametrize("new_role", ["editor", "owner"]) + def test_rbac_owner_lookup_failure_prevents_role_changes( + account_domain: AccountDomain, + members: tuple[str, str, str], + rbac_access: tuple[DeploymentWorkspaceMemberAccessGateway, RecordedRoleAssignments], + monkeypatch: pytest.MonkeyPatch, + new_role: str, + ) -> None: + workspace_id, owner_id, member_id = members + gateway, replace = rbac_access + + def get_roles(tenant_id: str, account_id: str | None, member_account_id: str) -> MemberRolesResponse: + del tenant_id, account_id + if new_role == "editor" or member_account_id == owner_id: + raise RuntimeError("RBAC unavailable") + return MemberRolesResponse(account_id=member_account_id) + + monkeypatch.setattr(RBACService.MemberRoles, "get", get_roles) + service = WorkspaceMemberService( + workspaces=account_domain.workspaces, accounts=account_domain.repository, access=gateway + ) + + with pytest.raises(RuntimeError, match="RBAC unavailable"): + service.update_role(workspace_id, member_id, new_role, owner_id) + + assert replace.calls == [] + assert service.get_role(workspace_id, owner_id) == _member_service_TenantAccountRole.OWNER + assert service.get_role(workspace_id, member_id) == _member_service_TenantAccountRole.NORMAL + + +def make_context() -> RequestContext: + return RequestContext( + request_id="request-1", + trace_id="trace-1", + account_id="actor-1", + active_workspace_id="workspace-1", + ) + + +def make_member( + member_id: str, + *, + status: str = "active", + legacy_role: str = "normal", +) -> WorkspaceMemberRecord: + created_at = datetime(2026, 1, 1) + return WorkspaceMemberRecord( + id=member_id, + name=f"Member {member_id}", + email=f"{member_id}@example.com", + avatar=None, + last_login_at=None, + last_active_at=created_at, + created_at=created_at, + status=status, + legacy_role=legacy_role, + ) + + +class RecordingMemberQuery: + def __init__(self, records: Sequence[WorkspaceMemberRecord]) -> None: + self.records = tuple(records) + self.workspace_ids: list[str] = [] + self.requested_roles: list[_workspace_member_query_service_TenantAccountRole | None] = [] + + def list_for_workspace( + self, workspace_id: str, *, role: _workspace_member_query_service_TenantAccountRole | None = None + ) -> Sequence[WorkspaceMemberRecord]: + self.workspace_ids.append(workspace_id) + self.requested_roles.append(role) + return self.records + + +class RecordingRoleResolver: + def __init__(self, roles: Mapping[str, Sequence[WorkspaceMemberRole]]) -> None: + self.roles = roles + self.calls: list[tuple[str, str, tuple[WorkspaceMemberRoleSubject, ...]]] = [] + self.languages: list[str | None] = [] + + def resolve_many( + self, + workspace_id: str, + actor_account_id: str, + subjects: Sequence[WorkspaceMemberRoleSubject], + *, + language: str | None = None, + ) -> Mapping[str, Sequence[WorkspaceMemberRole]]: + self.calls.append((workspace_id, actor_account_id, tuple(subjects))) + self.languages.append(language) + return self.roles + + +class FailingRoleResolver: + def resolve_many( + self, + workspace_id: str, + actor_account_id: str, + subjects: Sequence[WorkspaceMemberRoleSubject], + *, + language: str | None = None, + ) -> Mapping[str, Sequence[WorkspaceMemberRole]]: + del workspace_id, actor_account_id, subjects, language + raise RoleResolutionError + + +class RoleResolutionError(Exception): + pass + + +class TestWorkspaceMemberQueryServiceCases: + @staticmethod + @pytest.mark.parametrize("language", [None, "ja"]) + def test_list_current_projects_members_and_merges_roles_by_account_id(language: str | None) -> None: + active = make_member("active", legacy_role="owner") + pending = make_member("pending", status="pending") + members = RecordingMemberQuery([active, pending]) + roles = RecordingRoleResolver( + { + active.id: [ + WorkspaceMemberRole(id="workspace.owner", name="Owner"), + WorkspaceMemberRole(id="workspace.editor", name="Editor"), + ] + } + ) + service = WorkspaceMemberQueryService(members=members, roles=roles) + + result = service.list_current(make_context(), language=language) + + by_id = {member.id: member for member in result} + assert set(by_id) == {"active", "pending"} + assert by_id["active"] == WorkspaceMemberSummary( + id=active.id, + name=active.name, + email=active.email, + avatar=active.avatar, + last_login_at=active.last_login_at, + last_active_at=active.last_active_at, + created_at=active.created_at, + role=active.legacy_role, + roles=( + WorkspaceMemberRole(id="workspace.owner", name="Owner"), + WorkspaceMemberRole(id="workspace.editor", name="Editor"), + ), + status=active.status, + ) + assert by_id["pending"].status == "pending" + assert by_id["pending"].roles == () + assert members.workspace_ids == ["workspace-1"] + assert roles.languages == [language] + assert roles.calls == [ + ( + "workspace-1", + "actor-1", + ( + WorkspaceMemberRoleSubject(account_id=active.id, legacy_role=active.legacy_role), + WorkspaceMemberRoleSubject(account_id=pending.id, legacy_role=pending.legacy_role), + ), + ) ] - def fail_commit(session: Session) -> None: - session.flush() - raise RuntimeError("Membership commit failed") + @staticmethod + def test_list_current_propagates_role_resolution_failure() -> None: + members = RecordingMemberQuery([make_member("member-1")]) + service = WorkspaceMemberQueryService(members=members, roles=FailingRoleResolver()) - monkeypatch.setattr(RBACService.MemberRoles, "get", get_roles) - monkeypatch.setattr( - gateway, - "owner_id", - lambda _workspace_id, _actor_id: next( - member_id for member_id, roles in remote_roles.items() if owner_role in roles - ), - ) - replace.on_apply = replace_roles - service = WorkspaceMemberService( - workspaces=account_domain.workspaces, accounts=account_domain.repository, access=gateway + with pytest.raises(RoleResolutionError): + service.list_current(make_context()) + + @staticmethod + @pytest.mark.parametrize("operators_only", [True, False]) + def test_member_listing_passes_role_filter_to_repository(operators_only: bool) -> None: + records = [make_member("operator", legacy_role="dataset_operator")] + members = RecordingMemberQuery(records) + roles = RecordingRoleResolver({}) + service = WorkspaceMemberQueryService(members=members, roles=roles) + + assert service.list_members("workspace-1", dataset_operators_only=operators_only) == tuple(records) + assert members.workspace_ids == ["workspace-1"] + assert members.requested_roles == [ + _workspace_member_query_service_TenantAccountRole.DATASET_OPERATOR if operators_only else None + ] + assert roles.calls == [] + + @staticmethod + @pytest.mark.parametrize( + ("count", "page", "limit", "expected_ids"), + [(5, 2, 2, ["2", "3"]), (5, 3, 2, ["4"]), (5, 4, 2, []), (0, 1, 20, [])], ) + def test_member_page_preserves_total_and_order(count: int, page: int, limit: int, expected_ids: list[str]) -> None: + members = RecordingMemberQuery([make_member(str(index)) for index in range(count)]) + roles = RecordingRoleResolver({}) + service = WorkspaceMemberQueryService(members=members, roles=roles) - event.listen(sqlite_session_factory, "before_commit", fail_commit) - try: - with pytest.raises(RuntimeError, match="Membership commit failed"): - service.update_role(workspace_id, new_owner_id, "owner", old_owner_id) - finally: - event.remove(sqlite_session_factory, "before_commit", fail_commit) + result = service.list_page(make_context(), page=page, limit=limit) - assert remote_roles[old_owner_id] == [custom_role] - assert remote_roles[new_owner_id] == [owner_role] - assert service.get_role(workspace_id, old_owner_id) == TenantAccountRole.OWNER - assert service.get_role(workspace_id, new_owner_id) == TenantAccountRole.NORMAL - replace.calls.clear() - - target_id, role = (new_owner_id, "editor") if operation == "demote_new_owner" else (third_member.id, "owner") - with pytest.raises(NoPermissionError): - service.update_role(workspace_id, target_id, role, old_owner_id) - - assert replace.calls == [] - assert remote_roles[old_owner_id] == [custom_role] - assert remote_roles[new_owner_id] == [owner_role] - assert remote_roles[third_member.id] == [] - assert service.get_role(workspace_id, old_owner_id) == TenantAccountRole.OWNER - assert service.get_role(workspace_id, new_owner_id) == TenantAccountRole.NORMAL - assert service.get_role(workspace_id, third_member.id) == TenantAccountRole.NORMAL - - -@pytest.mark.parametrize("new_role", ["editor", "owner"]) -def test_rbac_owner_lookup_failure_prevents_role_changes( - account_domain: AccountDomain, - members: tuple[str, str, str], - rbac_access: tuple[DeploymentWorkspaceMemberAccessGateway, RecordedRoleAssignments], - monkeypatch: pytest.MonkeyPatch, - new_role: str, -) -> None: - workspace_id, owner_id, member_id = members - gateway, replace = rbac_access - - def get_roles(tenant_id: str, account_id: str | None, member_account_id: str) -> MemberRolesResponse: - del tenant_id, account_id - if new_role == "editor" or member_account_id == owner_id: - raise RuntimeError("RBAC unavailable") - return MemberRolesResponse(account_id=member_account_id) - - monkeypatch.setattr(RBACService.MemberRoles, "get", get_roles) - service = WorkspaceMemberService( - workspaces=account_domain.workspaces, accounts=account_domain.repository, access=gateway - ) - - with pytest.raises(RuntimeError, match="RBAC unavailable"): - service.update_role(workspace_id, member_id, new_role, owner_id) - - assert replace.calls == [] - assert service.get_role(workspace_id, owner_id) == TenantAccountRole.OWNER - assert service.get_role(workspace_id, member_id) == TenantAccountRole.NORMAL + assert result.total == count + assert [member.id for member in result.members] == expected_ids + assert members.workspace_ids == ["workspace-1"] + assert roles.calls == [] diff --git a/api/tests/unit_tests/services/workspace/test_workspace_provisioning.py b/api/tests/unit_tests/services/workspace/test_provisioning_service.py similarity index 99% rename from api/tests/unit_tests/services/workspace/test_workspace_provisioning.py rename to api/tests/unit_tests/services/workspace/test_provisioning_service.py index 2c010d5f277..5436e97c4f8 100644 --- a/api/tests/unit_tests/services/workspace/test_workspace_provisioning.py +++ b/api/tests/unit_tests/services/workspace/test_provisioning_service.py @@ -37,7 +37,7 @@ from services.workspace.provisioning_service import ( WorkspaceProvisioningService, WorkspaceProvisioningStore, ) -from tests.unit_tests.account_domain import AccountDomain +from tests.helpers.unit.account_domain import AccountDomain OWNER = AccountSnapshot( "account", diff --git a/api/tests/unit_tests/services/workspace/test_workspace_service.py b/api/tests/unit_tests/services/workspace/test_service.py similarity index 100% rename from api/tests/unit_tests/services/workspace/test_workspace_service.py rename to api/tests/unit_tests/services/workspace/test_service.py diff --git a/api/tests/unit_tests/services/workspace/test_workspace_credit_pool.py b/api/tests/unit_tests/services/workspace/test_workspace_credit_pool.py deleted file mode 100644 index 73c5d44eb31..00000000000 --- a/api/tests/unit_tests/services/workspace/test_workspace_credit_pool.py +++ /dev/null @@ -1,127 +0,0 @@ -"""Deployment billing adapter and paid-to-trial fallback regressions.""" - -from collections.abc import Callable -from unittest.mock import Mock - -import pytest - -from enums import CloudPlan, DeploymentEdition -from services.credit_pool_service import CreditPoolBalance -from services.entities.feature_entities import FeatureModel -from services.workspace import gateways -from services.workspace.gateways import DeploymentWorkspaceFeatureGateway, DeploymentWorkspacePlanGateway - - -@pytest.mark.parametrize( - ("plan", "paid_quota", "paid_used", "trial_quota", "expected_type", "remaining"), - [ - (CloudPlan.TEAM, 500, 120, 100, "paid", 380), - (CloudPlan.TEAM, -1, 999, 100, "paid", -1), - (CloudPlan.TEAM, 500, 500, 100, "trial", 60), - (CloudPlan.TEAM, None, 0, 100, "trial", 60), - (CloudPlan.SANDBOX, 500, 0, 100, "trial", 60), - (CloudPlan.SANDBOX, None, 0, None, None, None), - ], -) -def test_effective_credit_pool( - monkeypatch: pytest.MonkeyPatch, - config_overrides: Callable[..., None], - plan: CloudPlan | None, - paid_quota: int | None, - paid_used: int, - trial_quota: int | None, - expected_type: str | None, - remaining: int | None, -) -> None: - config_overrides(DEPLOYMENT_EDITION=DeploymentEdition.CLOUD) - monkeypatch.setattr( - gateways.BillingService, - "get_info", - lambda *_a, **_k: {"subscription": {"plan": plan}, "next_credit_reset_date": 1775001600}, - ) - paid = CreditPoolBalance("tenant", "paid", paid_quota, paid_used) if paid_quota is not None else None - trial = CreditPoolBalance("tenant", "trial", trial_quota, 40) if trial_quota is not None else None - calls = [] - - def get_pool(*, tenant_id: str, pool_type: str) -> CreditPoolBalance | None: - calls.append((tenant_id, pool_type)) - return paid if pool_type == "paid" else trial - - monkeypatch.setattr(gateways.CreditPoolService, "get_pool", get_pool) - result = DeploymentWorkspaceFeatureGateway().get_effective_credit_pool("tenant") - assert result.pool_type == expected_type - assert result.remaining_credits == remaining - assert result.next_credit_reset_date == 1775001600 - if plan == CloudPlan.SANDBOX: - assert calls == [("tenant", "trial")] - elif expected_type == "paid": - assert calls == [("tenant", "paid")] - else: - assert calls == [("tenant", "paid"), ("tenant", "trial")] - - -@pytest.mark.parametrize( - ("quota", "used", "exhausted_at", "expected"), - [(100, 100, 123, 123), (100, 10, 123, None), (-1, 100, 123, None), (100, 100, 0, None)], -) -def test_feature_info_uses_same_pool_rules( - monkeypatch: pytest.MonkeyPatch, - config_overrides: Callable[..., None], - quota: int, - used: int, - exhausted_at: int | None, - expected: int | None, -) -> None: - config_overrides(DEPLOYMENT_EDITION=DeploymentEdition.CLOUD) - features = FeatureModel(can_replace_logo=True) - monkeypatch.setattr(gateways.FeatureService, "get_features", lambda *_a, **_k: features) - monkeypatch.setattr( - gateways.CreditPoolService, - "get_pool", - lambda **_k: CreditPoolBalance("tenant", "trial", quota, used, exhausted_at), - ) - result = DeploymentWorkspaceFeatureGateway().get_features("tenant") - assert result.can_replace_logo - assert result.credits.exhausted_at == expected - - -@pytest.mark.parametrize("edition", [DeploymentEdition.COMMUNITY, DeploymentEdition.ENTERPRISE]) -def test_non_cloud_summary_skips_external_io( - monkeypatch: pytest.MonkeyPatch, config_overrides: Callable[..., None], edition: DeploymentEdition -) -> None: - config_overrides(DEPLOYMENT_EDITION=edition) - billing = Mock(side_effect=AssertionError("billing must not be queried")) - monkeypatch.setattr(gateways.BillingService, "get_info", billing) - result = DeploymentWorkspaceFeatureGateway().get_effective_credit_pool("tenant") - assert result.plan is None - assert result.remaining_credits is None - billing.assert_not_called() - - -@pytest.mark.parametrize( - ("edition", "bulk", "expected", "fallback_ids"), - [ - (DeploymentEdition.CLOUD, {"w1": {"plan": "team"}}, {"w1": "team", "w2": "professional"}, ["w2"]), - (DeploymentEdition.CLOUD, {}, {"w1": "professional", "w2": "professional"}, ["w1", "w2"]), - (DeploymentEdition.COMMUNITY, {}, {"w1": "professional", "w2": "professional"}, ["w1", "w2"]), - (DeploymentEdition.ENTERPRISE, {}, {"w1": "sandbox", "w2": "sandbox"}, []), - ], -) -def test_plan_gateway_bulk_and_fallback( - monkeypatch: pytest.MonkeyPatch, - config_overrides: Callable[..., None], - edition: DeploymentEdition, - bulk: dict[str, dict[str, str]], - expected: int | None, - fallback_ids: list[str], -) -> None: - config_overrides(DEPLOYMENT_EDITION=edition) - bulk_query = Mock(return_value=bulk) - features = FeatureModel() - features.billing.subscription.plan = CloudPlan.PROFESSIONAL - feature_query = Mock(return_value=features) - monkeypatch.setattr(gateways.BillingService, "get_plan_bulk", bulk_query) - monkeypatch.setattr(gateways.FeatureService, "get_features", feature_query) - assert DeploymentWorkspacePlanGateway().resolve_many(["w1", "w2"]) == expected - assert [call.args[0] for call in feature_query.call_args_list] == fallback_ids - assert bulk_query.call_count == int(edition == DeploymentEdition.CLOUD) diff --git a/api/tests/unit_tests/services/workspace/test_workspace_file_gateway.py b/api/tests/unit_tests/services/workspace/test_workspace_file_gateway.py deleted file mode 100644 index dec7b089fb8..00000000000 --- a/api/tests/unit_tests/services/workspace/test_workspace_file_gateway.py +++ /dev/null @@ -1,23 +0,0 @@ -"""File adapter keeps admitted identity explicit across the FileService boundary.""" - -from unittest.mock import Mock, create_autospec - -from machinery.context import RequestContext -from models.account import Account -from services.file_service import FileService -from services.workspace.gateways import WorkspaceFileGateway - - -def test_upload_uses_admitted_account_and_explicit_workspace() -> None: - files = create_autospec(FileService, instance=True) - files.upload_file.return_value = Mock(id="file-id") - gateway = WorkspaceFileGateway(files=files) - context = RequestContext("request", None, "account", "workspace") - - assert gateway.upload(context, filename="logo.png", content=b"image", mimetype="image/png") == "file-id" - - kwargs = files.upload_file.call_args.kwargs - assert isinstance(kwargs["user"], Account) - assert kwargs["user"].id == "account" - assert kwargs["tenant_id"] == "workspace" - assert kwargs["content"] == b"image" diff --git a/api/tests/unit_tests/services/workspace/test_workspace_member_query_service.py b/api/tests/unit_tests/services/workspace/test_workspace_member_query_service.py deleted file mode 100644 index d142750f4d4..00000000000 --- a/api/tests/unit_tests/services/workspace/test_workspace_member_query_service.py +++ /dev/null @@ -1,181 +0,0 @@ -from collections.abc import Mapping, Sequence -from datetime import datetime - -import pytest - -from enums.account import TenantAccountRole -from machinery.context import RequestContext -from services.workspace.contracts import ( - WorkspaceMemberRecord, - WorkspaceMemberRole, - WorkspaceMemberRoleSubject, - WorkspaceMemberSummary, -) -from services.workspace.member_service import WorkspaceMemberQueryService - - -def make_context() -> RequestContext: - return RequestContext( - request_id="request-1", - trace_id="trace-1", - account_id="actor-1", - active_workspace_id="workspace-1", - ) - - -def make_member( - member_id: str, - *, - status: str = "active", - legacy_role: str = "normal", -) -> WorkspaceMemberRecord: - created_at = datetime(2026, 1, 1) - return WorkspaceMemberRecord( - id=member_id, - name=f"Member {member_id}", - email=f"{member_id}@example.com", - avatar=None, - last_login_at=None, - last_active_at=created_at, - created_at=created_at, - status=status, - legacy_role=legacy_role, - ) - - -class RecordingMemberQuery: - def __init__(self, records: Sequence[WorkspaceMemberRecord]) -> None: - self.records = tuple(records) - self.workspace_ids: list[str] = [] - self.requested_roles: list[TenantAccountRole | None] = [] - - def list_for_workspace( - self, workspace_id: str, *, role: TenantAccountRole | None = None - ) -> Sequence[WorkspaceMemberRecord]: - self.workspace_ids.append(workspace_id) - self.requested_roles.append(role) - return self.records - - -class RecordingRoleResolver: - def __init__(self, roles: Mapping[str, Sequence[WorkspaceMemberRole]]) -> None: - self.roles = roles - self.calls: list[tuple[str, str, tuple[WorkspaceMemberRoleSubject, ...]]] = [] - self.languages: list[str | None] = [] - - def resolve_many( - self, - workspace_id: str, - actor_account_id: str, - subjects: Sequence[WorkspaceMemberRoleSubject], - *, - language: str | None = None, - ) -> Mapping[str, Sequence[WorkspaceMemberRole]]: - self.calls.append((workspace_id, actor_account_id, tuple(subjects))) - self.languages.append(language) - return self.roles - - -class FailingRoleResolver: - def resolve_many( - self, - workspace_id: str, - actor_account_id: str, - subjects: Sequence[WorkspaceMemberRoleSubject], - *, - language: str | None = None, - ) -> Mapping[str, Sequence[WorkspaceMemberRole]]: - del workspace_id, actor_account_id, subjects, language - raise RoleResolutionError - - -class RoleResolutionError(Exception): - pass - - -@pytest.mark.parametrize("language", [None, "ja"]) -def test_list_current_projects_members_and_merges_roles_by_account_id(language: str | None) -> None: - active = make_member("active", legacy_role="owner") - pending = make_member("pending", status="pending") - members = RecordingMemberQuery([active, pending]) - roles = RecordingRoleResolver( - { - active.id: [ - WorkspaceMemberRole(id="workspace.owner", name="Owner"), - WorkspaceMemberRole(id="workspace.editor", name="Editor"), - ] - } - ) - service = WorkspaceMemberQueryService(members=members, roles=roles) - - result = service.list_current(make_context(), language=language) - - by_id = {member.id: member for member in result} - assert set(by_id) == {"active", "pending"} - assert by_id["active"] == WorkspaceMemberSummary( - id=active.id, - name=active.name, - email=active.email, - avatar=active.avatar, - last_login_at=active.last_login_at, - last_active_at=active.last_active_at, - created_at=active.created_at, - role=active.legacy_role, - roles=( - WorkspaceMemberRole(id="workspace.owner", name="Owner"), - WorkspaceMemberRole(id="workspace.editor", name="Editor"), - ), - status=active.status, - ) - assert by_id["pending"].status == "pending" - assert by_id["pending"].roles == () - assert members.workspace_ids == ["workspace-1"] - assert roles.languages == [language] - assert roles.calls == [ - ( - "workspace-1", - "actor-1", - ( - WorkspaceMemberRoleSubject(account_id=active.id, legacy_role=active.legacy_role), - WorkspaceMemberRoleSubject(account_id=pending.id, legacy_role=pending.legacy_role), - ), - ) - ] - - -def test_list_current_propagates_role_resolution_failure() -> None: - members = RecordingMemberQuery([make_member("member-1")]) - service = WorkspaceMemberQueryService(members=members, roles=FailingRoleResolver()) - - with pytest.raises(RoleResolutionError): - service.list_current(make_context()) - - -@pytest.mark.parametrize("operators_only", [True, False]) -def test_member_listing_passes_role_filter_to_repository(operators_only: bool) -> None: - records = [make_member("operator", legacy_role="dataset_operator")] - members = RecordingMemberQuery(records) - roles = RecordingRoleResolver({}) - service = WorkspaceMemberQueryService(members=members, roles=roles) - - assert service.list_members("workspace-1", dataset_operators_only=operators_only) == tuple(records) - assert members.workspace_ids == ["workspace-1"] - assert members.requested_roles == [TenantAccountRole.DATASET_OPERATOR if operators_only else None] - assert roles.calls == [] - - -@pytest.mark.parametrize( - ("count", "page", "limit", "expected_ids"), - [(5, 2, 2, ["2", "3"]), (5, 3, 2, ["4"]), (5, 4, 2, []), (0, 1, 20, [])], -) -def test_member_page_preserves_total_and_order(count: int, page: int, limit: int, expected_ids: list[str]) -> None: - members = RecordingMemberQuery([make_member(str(index)) for index in range(count)]) - roles = RecordingRoleResolver({}) - service = WorkspaceMemberQueryService(members=members, roles=roles) - - result = service.list_page(make_context(), page=page, limit=limit) - - assert result.total == count - assert [member.id for member in result.members] == expected_ids - assert members.workspace_ids == ["workspace-1"] - assert roles.calls == [] diff --git a/api/tests/unit_tests/services/workspace/test_workspace_member_role_resolver.py b/api/tests/unit_tests/services/workspace/test_workspace_member_role_resolver.py deleted file mode 100644 index 17e7e896ca8..00000000000 --- a/api/tests/unit_tests/services/workspace/test_workspace_member_role_resolver.py +++ /dev/null @@ -1,88 +0,0 @@ -from collections.abc import Callable - -import pytest - -from services.enterprise.base import EnterpriseRequest -from services.workspace.contracts import WorkspaceMemberRole, WorkspaceMemberRoleSubject -from services.workspace.gateways import DeploymentWorkspaceMemberRoleResolver -from tests.unit_tests.rbac_fakes import RBACTransport - - -def make_subject(account_id: str, *, legacy_role: str = "normal") -> WorkspaceMemberRoleSubject: - return WorkspaceMemberRoleSubject(account_id=account_id, legacy_role=legacy_role) - - -@pytest.fixture -def transport(monkeypatch: pytest.MonkeyPatch) -> RBACTransport: - transport = RBACTransport() - monkeypatch.setattr(EnterpriseRequest, "send_inner_rbac_request", transport.send) - return transport - - -def test_builtin_mode_projects_join_roles_without_enterprise_call( - config_overrides: Callable[..., None], transport: RBACTransport -) -> None: - config_overrides(RBAC_ENABLED=False) - - result = DeploymentWorkspaceMemberRoleResolver().resolve_many( - "workspace-1", "actor-1", [make_subject("owner", legacy_role="owner"), make_subject("member")] - ) - - assert result == { - "owner": (WorkspaceMemberRole(id="owner", name="owner"),), - "member": (WorkspaceMemberRole(id="normal", name="normal"),), - } - assert transport.requests == [] - - -@pytest.mark.parametrize("language", [None, "ja"]) -def test_rbac_mode_maps_batch_response_without_builtin_fallback( - config_overrides: Callable[..., None], transport: RBACTransport, language: str | None -) -> None: - config_overrides(RBAC_ENABLED=True) - transport.response = { - "owner": [ - {"id": "workspace.owner", "name": "Owner", "type": "builtin"}, - {"id": "workspace.editor", "name": "Editor", "type": "builtin"}, - ] - } - - result = DeploymentWorkspaceMemberRoleResolver().resolve_many( - "workspace-1", - "actor-1", - [make_subject("owner", legacy_role="owner"), make_subject("omitted", legacy_role="admin")], - language=language, - ) - - assert result == { - "owner": ( - WorkspaceMemberRole(id="workspace.owner", name="Owner"), - WorkspaceMemberRole(id="workspace.editor", name="Editor"), - ) - } - sent = transport.only_request - assert (sent.method, sent.endpoint) == ("POST", "/rbac/members/rbac-roles/batch") - assert (sent.tenant_id, sent.account_id) == ("workspace-1", "actor-1") - assert sent.json == {"member_ids": ["owner", "omitted"]} - assert sent.params == ({"language": language} if language else None) - - -def test_rbac_failure_propagates(config_overrides: Callable[..., None], transport: RBACTransport) -> None: - config_overrides(RBAC_ENABLED=True) - transport.failure = RoleResolutionError("enterprise unavailable") - - with pytest.raises(RoleResolutionError, match="enterprise unavailable"): - DeploymentWorkspaceMemberRoleResolver().resolve_many("workspace-1", "actor-1", [make_subject("member-1")]) - - -def test_empty_member_list_skips_enterprise_call( - config_overrides: Callable[..., None], transport: RBACTransport -) -> None: - config_overrides(RBAC_ENABLED=True) - - assert DeploymentWorkspaceMemberRoleResolver().resolve_many("workspace-1", "actor-1", []) == {} - assert transport.requests == [] - - -class RoleResolutionError(Exception): - pass diff --git a/api/tests/unit_tests/tasks/test_workflow_execute_task.py b/api/tests/unit_tests/tasks/app_generate/test_workflow_execute_task.py similarity index 100% rename from api/tests/unit_tests/tasks/test_workflow_execute_task.py rename to api/tests/unit_tests/tasks/app_generate/test_workflow_execute_task.py diff --git a/api/tests/unit_tests/tasks/test_application_trace_task.py b/api/tests/unit_tests/tasks/test_application_trace_task.py index f9713f5be4a..0768f6732b2 100644 --- a/api/tests/unit_tests/tasks/test_application_trace_task.py +++ b/api/tests/unit_tests/tasks/test_application_trace_task.py @@ -12,8 +12,11 @@ from services.tracing import service as service_module from services.tracing.provider_dispatch import TraceExportResult from services.tracing.service import ApplicationTracingService from tasks.application_trace_task import export_application_trace, reconcile_application_traces -from tests.unit_tests.config_override import apply_config_overrides -from tests.unit_tests.services.tracing.test_application_trace_collection import batch_for, ingest, seed, terminal +from tests.helpers.unit.config_override import apply_config_overrides +from tests.unit_tests.services.tracing.test_service import _application_trace_collection_batch_for as batch_for +from tests.unit_tests.services.tracing.test_service import _application_trace_collection_ingest as ingest +from tests.unit_tests.services.tracing.test_service import _application_trace_collection_seed as seed +from tests.unit_tests.services.tracing.test_service import _application_trace_collection_terminal as terminal def test_worker_task_exports_once_and_replayed_task_observes_checkpoint( diff --git a/api/tests/unit_tests/tasks/test_clean_document_task.py b/api/tests/unit_tests/tasks/test_clean_document_task.py index 9fbb3f9b91e..36aa6c59581 100644 --- a/api/tests/unit_tests/tasks/test_clean_document_task.py +++ b/api/tests/unit_tests/tasks/test_clean_document_task.py @@ -26,7 +26,7 @@ from models.enums import CreatorUserRole, DataSourceType from models.knowledge_fs import KnowledgeFSUpgradeFileLease, KnowledgeFSUpgradeJob from models.model import UploadFile from tasks.clean_document_task import clean_document_task -from tests.unit_tests.model_factories import make_document +from tests.helpers.unit.model_factories import make_document SQLITE_MODELS = ( Dataset, diff --git a/api/tests/unit_tests/tasks/test_delete_conversation_task.py b/api/tests/unit_tests/tasks/test_delete_conversation_task.py index c21c3391183..a0a6e53454e 100644 --- a/api/tests/unit_tests/tasks/test_delete_conversation_task.py +++ b/api/tests/unit_tests/tasks/test_delete_conversation_task.py @@ -38,7 +38,7 @@ from models.enums import ( from models.human_input import HumanInputDelivery, HumanInputFormRecipient, RecipientType from models.tools import ToolConversationVariables, ToolFile from tasks.delete_conversation_task import _cleanup_conversation_related_data, sweep_deleted_conversations -from tests.unit_tests.model_factories import make_conversation, make_message +from tests.helpers.unit.model_factories import make_conversation, make_message TENANT_ID = "11111111-1111-1111-1111-111111111111" APP_ID = "22222222-2222-2222-2222-222222222222" diff --git a/api/tests/unit_tests/tasks/test_segment_index_cleanup_tasks.py b/api/tests/unit_tests/tasks/test_delete_segment_from_index_task.py similarity index 100% rename from api/tests/unit_tests/tasks/test_segment_index_cleanup_tasks.py rename to api/tests/unit_tests/tasks/test_delete_segment_from_index_task.py diff --git a/api/tests/unit_tests/tasks/test_dataset_indexing_task.py b/api/tests/unit_tests/tasks/test_document_indexing_task.py similarity index 99% rename from api/tests/unit_tests/tasks/test_dataset_indexing_task.py rename to api/tests/unit_tests/tasks/test_document_indexing_task.py index 51c61f45e9d..49edd1d96b5 100644 --- a/api/tests/unit_tests/tasks/test_dataset_indexing_task.py +++ b/api/tests/unit_tests/tasks/test_document_indexing_task.py @@ -29,7 +29,7 @@ from tasks.document_indexing_task import ( priority_document_indexing_task, ) from tasks.recover_document_indexing_task import recover_document_indexing_task -from tests.unit_tests.config_override import apply_config_overrides, config_overrides_context +from tests.helpers.unit.config_override import apply_config_overrides, config_overrides_context @pytest.mark.parametrize( diff --git a/api/tests/unit_tests/tasks/test_enable_segment_index_tasks.py b/api/tests/unit_tests/tasks/test_enable_segments_to_index_task.py similarity index 100% rename from api/tests/unit_tests/tasks/test_enable_segment_index_tasks.py rename to api/tests/unit_tests/tasks/test_enable_segments_to_index_task.py diff --git a/api/tests/unit_tests/tasks/test_summary_queue_isolation.py b/api/tests/unit_tests/tasks/test_generate_summary_index_task.py similarity index 100% rename from api/tests/unit_tests/tasks/test_summary_queue_isolation.py rename to api/tests/unit_tests/tasks/test_generate_summary_index_task.py diff --git a/api/tests/unit_tests/tasks/test_human_input_timeout_tasks.py b/api/tests/unit_tests/tasks/test_human_input_timeout_tasks.py index 1b69892571b..d24c0f24c56 100644 --- a/api/tests/unit_tests/tasks/test_human_input_timeout_tasks.py +++ b/api/tests/unit_tests/tasks/test_human_input_timeout_tasks.py @@ -14,7 +14,7 @@ from core.workflow.nodes.human_input.entities import FormDefinition from core.workflow.nodes.human_input.enums import HumanInputFormKind, HumanInputFormStatus from models.human_input import HumanInputForm from tasks import human_input_timeout_tasks as task_module -from tests.unit_tests.config_override import apply_config_overrides +from tests.helpers.unit.config_override import apply_config_overrides class _FakeService: diff --git a/api/tests/unit_tests/tasks/test_initialize_created_app_rbac_access_task.py b/api/tests/unit_tests/tasks/test_initialize_created_app_rbac_access_task.py index 42d0497e3b5..b7ac59dae43 100644 --- a/api/tests/unit_tests/tasks/test_initialize_created_app_rbac_access_task.py +++ b/api/tests/unit_tests/tasks/test_initialize_created_app_rbac_access_task.py @@ -17,7 +17,7 @@ from tasks.initialize_created_app_rbac_access_task import ( initialize_created_app_rbac_access_task, sync_joined_workspace_member_rbac_access_task, ) -from tests.unit_tests.rbac_fakes import RBACDomain +from tests.helpers.unit.rbac_fakes import RBACDomain class _TaskWithQueue(Protocol): diff --git a/api/tests/unit_tests/tasks/test_install_default_plugins_task.py b/api/tests/unit_tests/tasks/test_install_default_plugins_task.py index 2a24a374951..32ee68baeb1 100644 --- a/api/tests/unit_tests/tasks/test_install_default_plugins_task.py +++ b/api/tests/unit_tests/tasks/test_install_default_plugins_task.py @@ -7,7 +7,7 @@ import pytest from celery.exceptions import Retry from core.plugin.entities.plugin_daemon import PluginInstallTaskStartResponse -from tests.unit_tests.config_override import apply_config_overrides +from tests.helpers.unit.config_override import apply_config_overrides class _TaskWithQueue(Protocol): diff --git a/api/tests/unit_tests/tasks/test_knowledge_fs_agent_investigation_tasks.py b/api/tests/unit_tests/tasks/test_knowledge_fs_agent_investigation_tasks.py index 711b91220c5..579d9da459e 100644 --- a/api/tests/unit_tests/tasks/test_knowledge_fs_agent_investigation_tasks.py +++ b/api/tests/unit_tests/tasks/test_knowledge_fs_agent_investigation_tasks.py @@ -9,8 +9,8 @@ from dify_agent.protocol.knowledge_investigation import KnowledgeInvestigationPa from services.knowledge_fs.product_remote import KnowledgeFSProductRequestRejectedError from tasks import knowledge_fs_agent_investigation_tasks as module -from tests.unit_tests.tasks.knowledge_fs_task_fakes import RecordingRetry -from tests.unit_tests.tasks.task_options import task_options +from tests.helpers.unit.tasks.knowledge_fs_task_fakes import RecordingRetry +from tests.helpers.unit.tasks.task_options import task_options @dataclass diff --git a/api/tests/unit_tests/tasks/test_knowledge_fs_failed_retrieval_tasks.py b/api/tests/unit_tests/tasks/test_knowledge_fs_failed_retrieval_tasks.py index e471feff134..fe72007a5a4 100644 --- a/api/tests/unit_tests/tasks/test_knowledge_fs_failed_retrieval_tasks.py +++ b/api/tests/unit_tests/tasks/test_knowledge_fs_failed_retrieval_tasks.py @@ -16,8 +16,8 @@ from services.knowledge_fs.product_dto import ( from services.knowledge_fs.product_remote import KnowledgeFSProductRequestRejectedError from services.knowledge_fs.resource import KnowledgeResourceRef from tasks import knowledge_fs_failed_retrieval_tasks as task_module -from tests.unit_tests.tasks.knowledge_fs_task_fakes import RecordingRetry -from tests.unit_tests.tasks.task_options import task_options +from tests.helpers.unit.tasks.knowledge_fs_task_fakes import RecordingRetry +from tests.helpers.unit.tasks.task_options import task_options EVENT_ID = "019fac9f-bfb0-75ee-9af5-252ebafbac1e" diff --git a/api/tests/unit_tests/tasks/test_knowledge_fs_initial_source_preview_tasks.py b/api/tests/unit_tests/tasks/test_knowledge_fs_initial_source_preview_tasks.py index d98ac2903dd..274f791ea35 100644 --- a/api/tests/unit_tests/tasks/test_knowledge_fs_initial_source_preview_tasks.py +++ b/api/tests/unit_tests/tasks/test_knowledge_fs_initial_source_preview_tasks.py @@ -13,7 +13,7 @@ from services.knowledge_fs.product_dto import ( KnowledgeFSInitialSourcePreviewResponse, ) from tasks.knowledge_fs_initial_source_preview_tasks import run_knowledge_fs_initial_source_preview -from tests.unit_tests.tasks.task_options import task_options +from tests.helpers.unit.tasks.task_options import task_options _JOB_KEY = {"tenant_id": "tenant-1", "account_id": "account-1", "job_id": "job-1"} diff --git a/api/tests/unit_tests/tasks/test_knowledge_fs_lifecycle_tasks.py b/api/tests/unit_tests/tasks/test_knowledge_fs_lifecycle_tasks.py index ac63311604d..7723658d7ce 100644 --- a/api/tests/unit_tests/tasks/test_knowledge_fs_lifecycle_tasks.py +++ b/api/tests/unit_tests/tasks/test_knowledge_fs_lifecycle_tasks.py @@ -13,7 +13,7 @@ from tasks.knowledge_fs_lifecycle_tasks import ( cleanup_knowledge_fs_staged_uploads, run_knowledge_fs_lifecycle_worker, ) -from tests.unit_tests.config_override import apply_config_overrides +from tests.helpers.unit.config_override import apply_config_overrides @dataclass diff --git a/api/tests/unit_tests/tasks/test_knowledge_fs_source_import_tasks.py b/api/tests/unit_tests/tasks/test_knowledge_fs_source_import_tasks.py index 2a108586672..f234f5ae3ac 100644 --- a/api/tests/unit_tests/tasks/test_knowledge_fs_source_import_tasks.py +++ b/api/tests/unit_tests/tasks/test_knowledge_fs_source_import_tasks.py @@ -19,7 +19,7 @@ from tasks.knowledge_fs_source_import_tasks import ( KnowledgeFSSourceImportNotReadyError, finalize_source_import_once, ) -from tests.unit_tests.tasks.knowledge_fs_task_fakes import ( +from tests.helpers.unit.tasks.knowledge_fs_task_fakes import ( FacadeRuntime, RecordingKnowledgeFSFacade, RecordingRetry, diff --git a/api/tests/unit_tests/tasks/test_knowledge_fs_upgrade_tasks.py b/api/tests/unit_tests/tasks/test_knowledge_fs_upgrade_tasks.py index ba2df3ed108..00c06f85009 100644 --- a/api/tests/unit_tests/tasks/test_knowledge_fs_upgrade_tasks.py +++ b/api/tests/unit_tests/tasks/test_knowledge_fs_upgrade_tasks.py @@ -19,7 +19,7 @@ from tasks.knowledge_fs_upgrade_tasks import ( recover_stalled_knowledge_fs_upgrades, run_knowledge_fs_upgrade, ) -from tests.unit_tests.tasks.task_options import task_options +from tests.helpers.unit.tasks.task_options import task_options @dataclass diff --git a/api/tests/unit_tests/tasks/test_mail_send_task.py b/api/tests/unit_tests/tasks/test_mail_inner_task.py similarity index 100% rename from api/tests/unit_tests/tasks/test_mail_send_task.py rename to api/tests/unit_tests/tasks/test_mail_inner_task.py diff --git a/api/tests/unit_tests/tasks/test_ops_trace_task.py b/api/tests/unit_tests/tasks/test_ops_trace_task.py index 95960c69239..e82849abe71 100644 --- a/api/tests/unit_tests/tasks/test_ops_trace_task.py +++ b/api/tests/unit_tests/tasks/test_ops_trace_task.py @@ -19,7 +19,7 @@ from core.ops.exceptions import PendingTraceParentContextError, RetryableTraceDi from services.tracing.provider_dispatch import TraceExportUncertainError from services.tracing.runtime_bridge import RuntimeApplicationTracingHooks from tasks.ops_trace_task import process_trace_tasks -from tests.unit_tests.config_override import apply_config_overrides +from tests.helpers.unit.config_override import apply_config_overrides @pytest.fixture(autouse=True) diff --git a/api/tests/unit_tests/tasks/test_remove_app_and_related_data_task.py b/api/tests/unit_tests/tasks/test_remove_app_and_related_data_task.py index 693dcddf73c..eabaab82567 100644 --- a/api/tests/unit_tests/tasks/test_remove_app_and_related_data_task.py +++ b/api/tests/unit_tests/tasks/test_remove_app_and_related_data_task.py @@ -28,7 +28,7 @@ from tasks.remove_app_and_related_data_task import ( _delete_workflow_agent_node_bindings, delete_draft_variables_batch, ) -from tests.unit_tests.config_override import apply_config_overrides +from tests.helpers.unit.config_override import apply_config_overrides def test_delete_workflow_agent_node_bindings_is_scoped_to_tenant_and_app(sqlite_session: Session) -> None: diff --git a/api/tests/unit_tests/tasks/test_retry_document_indexing_task.py b/api/tests/unit_tests/tasks/test_retry_document_indexing_task.py index 2adf8cd37e5..e68ef04dc31 100644 --- a/api/tests/unit_tests/tasks/test_retry_document_indexing_task.py +++ b/api/tests/unit_tests/tasks/test_retry_document_indexing_task.py @@ -12,7 +12,7 @@ from models.account import TenantAccountRole from models.dataset import Dataset, Document from models.enums import DatasetRuntimeMode, DataSourceType, DocumentCreatedFrom, IndexingStatus from tasks.retry_document_indexing_task import retry_document_indexing_task -from tests.unit_tests.config_override import config_overrides_context +from tests.helpers.unit.config_override import config_overrides_context @config_overrides_context(DEPLOYMENT_EDITION=DeploymentEdition.COMMUNITY) diff --git a/api/tests/unit_tests/tasks/test_sync_website_document_indexing_task.py b/api/tests/unit_tests/tasks/test_sync_website_document_indexing_task.py index 0cfa3d62742..48b26055bb4 100644 --- a/api/tests/unit_tests/tasks/test_sync_website_document_indexing_task.py +++ b/api/tests/unit_tests/tasks/test_sync_website_document_indexing_task.py @@ -9,8 +9,8 @@ from enums import DeploymentEdition from models.dataset import Dataset, Document, DocumentSegment from models.enums import DataSourceType from tasks.sync_website_document_indexing_task import sync_website_document_indexing_task -from tests.unit_tests.config_override import config_overrides_context -from tests.unit_tests.model_factories import make_dataset, make_document +from tests.helpers.unit.config_override import config_overrides_context +from tests.helpers.unit.model_factories import make_dataset, make_document def _dataset(tenant_id: str) -> Dataset: diff --git a/api/tests/unit_tests/tasks/test_trigger_processing_tasks.py b/api/tests/unit_tests/tasks/test_trigger_processing_tasks.py index f8ba6ad57fe..a344b89a50b 100644 --- a/api/tests/unit_tests/tasks/test_trigger_processing_tasks.py +++ b/api/tests/unit_tests/tasks/test_trigger_processing_tasks.py @@ -12,7 +12,7 @@ from models.trigger import TriggerSubscription, WorkflowPluginTrigger from models.workflow import Workflow from services.errors.app import QuotaExceededError from tasks.trigger_processing_tasks import dispatch_triggered_workflow -from tests.unit_tests.model_factories import make_end_user, make_workflow +from tests.helpers.unit.model_factories import make_end_user, make_workflow def _workflow(*, app_id: str = "app-123") -> Workflow: diff --git a/api/tests/unit_tests/tasks/test_workflow_node_execution_identity.py b/api/tests/unit_tests/tasks/test_workflow_node_execution_identity.py deleted file mode 100644 index e1e8f8bc0cf..00000000000 --- a/api/tests/unit_tests/tasks/test_workflow_node_execution_identity.py +++ /dev/null @@ -1,28 +0,0 @@ -from datetime import UTC, datetime - -from graphon.entities import WorkflowNodeExecution -from models.workflow import WorkflowNodeExecutionModel -from tasks.workflow_node_execution_tasks import _update_node_execution_from_domain - - -def test_celery_update_preserves_workflow_agent_binding_identity() -> None: - stored = WorkflowNodeExecutionModel( - process_data='{"workflow_agent_binding_id": "workflow-binding-1"}', - ) - incoming = WorkflowNodeExecution( - id="execution-1", - workflow_id="workflow-1", - node_id="node-1", - node_type="agent", - title="Agent", - index=1, - created_at=datetime.now(UTC), - process_data={"retry": True}, - ) - - _update_node_execution_from_domain(stored, incoming) - - assert stored.process_data_dict == { - "retry": True, - "workflow_agent_binding_id": "workflow-binding-1", - } diff --git a/api/tests/unit_tests/tasks/test_workflow_node_execution_tasks.py b/api/tests/unit_tests/tasks/test_workflow_node_execution_tasks.py index 041261d45c0..a4c2324db81 100644 --- a/api/tests/unit_tests/tasks/test_workflow_node_execution_tasks.py +++ b/api/tests/unit_tests/tasks/test_workflow_node_execution_tasks.py @@ -1,5 +1,3 @@ -"""SQLite-backed tests for asynchronous workflow node-execution persistence.""" - import json from datetime import UTC, datetime, timedelta from decimal import Decimal @@ -25,6 +23,33 @@ from tasks.workflow_node_execution_tasks import ( save_workflow_node_executions_task, ) + +def test_celery_update_preserves_workflow_agent_binding_identity() -> None: + stored = WorkflowNodeExecutionModel( + process_data='{"workflow_agent_binding_id": "workflow-binding-1"}', + ) + incoming = WorkflowNodeExecution( + id="execution-1", + workflow_id="workflow-1", + node_id="node-1", + node_type="agent", + title="Agent", + index=1, + created_at=datetime.now(UTC), + process_data={"retry": True}, + ) + + _update_node_execution_from_domain(stored, incoming) + + assert stored.process_data_dict == { + "retry": True, + "workflow_agent_binding_id": "workflow-binding-1", + } + + +"""SQLite-backed tests for asynchronous workflow node-execution persistence.""" + + TENANT_ID = "00000000-0000-0000-0000-000000000010" APP_ID = "00000000-0000-0000-0000-000000000020" WORKFLOW_ID = "00000000-0000-0000-0000-000000000030" diff --git a/api/tests/unit_tests/test_app_factory.py b/api/tests/unit_tests/test_app_factory.py index 816f49dd045..3e8ba6a6806 100644 --- a/api/tests/unit_tests/test_app_factory.py +++ b/api/tests/unit_tests/test_app_factory.py @@ -16,7 +16,7 @@ from app_factory import create_flask_app_with_configs from enums import DeploymentEdition from libs.external_api import ExternalApi from services.entities.feature_entities import LicenseStatus -from tests.unit_tests.config_override import config_overrides_context +from tests.helpers.unit.config_override import config_overrides_context INVALID_STATUSES = [LicenseStatus.INACTIVE, LicenseStatus.EXPIRED, LicenseStatus.LOST] VALID_STATUSES = [LicenseStatus.ACTIVE, LicenseStatus.EXPIRING] diff --git a/api/tests/unit_tests/tools/test_api_tool.py b/api/tests/unit_tests/tools/test_api_tool.py deleted file mode 100644 index 948482269ae..00000000000 --- a/api/tests/unit_tests/tools/test_api_tool.py +++ /dev/null @@ -1,233 +0,0 @@ -import json -import operator -from collections.abc import Iterator -from unittest.mock import Mock, patch - -import httpx -import pytest -from sqlalchemy import Engine -from sqlalchemy.orm import Session - -from core.tools.__base.tool_runtime import ToolRuntime -from core.tools.custom_tool.tool import ApiTool -from core.tools.entities.common_entities import I18nObject -from core.tools.entities.tool_bundle import ApiToolBundle -from core.tools.entities.tool_entities import ( - ToolEntity, - ToolIdentity, - ToolInvokeMessage, -) - - -@pytest.fixture -def database_session(sqlite_engine: Engine) -> Iterator[Session]: - """Yield the live ORM session required by the shared tool invocation contract.""" - with Session(sqlite_engine) as session: - yield session - - -def _get_message_by_type[T](msgs: list[ToolInvokeMessage], msg_type: type[T]) -> ToolInvokeMessage | None: - return next((i for i in msgs if isinstance(i.message, msg_type)), None) - - -class TestApiToolInvoke: - """Test suite for ApiTool._invoke method to ensure JSON responses are properly serialized.""" - - def setup_method(self): - """Setup test fixtures.""" - # Create a mock tool entity - self.mock_tool_identity = ToolIdentity( - author="test", - name="test_api_tool", - label=I18nObject(en_US="Test API Tool", zh_Hans="测试API工具"), - provider="test_provider", - ) - self.mock_tool_entity = ToolEntity(identity=self.mock_tool_identity) - - # Create a mock API bundle - self.mock_api_bundle = ApiToolBundle( - server_url="https://api.example.com/test", - method="GET", - openapi={}, - operation_id="test_operation", - parameters=[], - author="test_author", - ) - - self.runtime = ToolRuntime(tenant_id="test_tenant", credentials={"auth_type": "none"}) - - # Create the ApiTool instance - self.api_tool = ApiTool( - entity=self.mock_tool_entity, - api_bundle=self.mock_api_bundle, - runtime=self.runtime, - provider_id="test_provider", - ) - - @patch("core.tools.custom_tool.tool.ssrf_proxy.get") - def test_invoke_with_json_response_creates_text_message_with_serialized_json( - self, mock_get: Mock, database_session: Session - ) -> None: - """Test that when upstream returns JSON, the output Text message contains JSON-serialized string.""" - # Use the response's real JSON parsing and content decoding. - json_response_data = { - "key": "value", - "number": 123, - "nested": {"inner": "data"}, - } - mock_get.return_value = httpx.Response(200, json=json_response_data) - - # Invoke the tool - result_generator = self.api_tool._invoke(session=database_session, user_id="test_user", tool_parameters={}) - - # Get the result from the generator - result = list(result_generator) - assert len(result) == 2 - - # Verify _invoke yields text message - text_message = _get_message_by_type(result, ToolInvokeMessage.TextMessage) - assert text_message is not None, "_invoke should yield a text message" - assert isinstance(text_message, ToolInvokeMessage) - assert text_message.type == ToolInvokeMessage.MessageType.TEXT - assert text_message.message is not None - # Verify the text contains the JSON-serialized string - # Check if message is a TextMessage - assert isinstance(text_message.message, ToolInvokeMessage.TextMessage) - # Verify it's a valid JSON string and equals to the mock response - parsed_back = json.loads(text_message.message.text) - assert parsed_back == json_response_data - - # Verify _invoke yields json message - json_message = _get_message_by_type(result, ToolInvokeMessage.JsonMessage) - assert json_message is not None, "_invoke should yield a JSON message" - assert isinstance(json_message, ToolInvokeMessage) - assert json_message.type == ToolInvokeMessage.MessageType.JSON - assert json_message.message is not None - - assert isinstance(json_message.message, ToolInvokeMessage.JsonMessage) - assert json_message.message.json_object == json_response_data - - @patch("core.tools.custom_tool.tool.ssrf_proxy.get") - @pytest.mark.parametrize( - "test_case", - [ - ( - "array", - [ - {"id": 1, "name": "Item 1", "active": True}, - {"id": 2, "name": "Item 2", "active": False}, - {"id": 3, "name": "项目 3", "active": True}, - ], - ), - ( - "string", - "string", - ), - ( - "number", - 123.456, - ), - ( - "boolean", - True, - ), - ( - "null", - None, - ), - ], - ids=operator.itemgetter(0), - ) - def test_invoke_with_non_dict_json_response_creates_text_message_with_serialized_json( - self, mock_get: Mock, test_case, database_session: Session - ) -> None: - """Test that when upstream returns a non-dict JSON, the output Text message contains JSON-serialized string.""" - # Explicit content also covers JSON null, which is distinct from no body. - _, json_value = test_case - mock_get.return_value = httpx.Response( - 200, content=json.dumps(json_value), headers={"content-type": "application/json"} - ) - - # Invoke the tool - result_generator = self.api_tool._invoke(session=database_session, user_id="test_user", tool_parameters={}) - - # Get the result from the generator - result = list(result_generator) - assert len(result) == 1 - - # Verify _invoke yields a text message - text_message = _get_message_by_type(result, ToolInvokeMessage.TextMessage) - assert text_message is not None, "_invoke should yield a text message containing the serialized JSON." - assert isinstance(text_message, ToolInvokeMessage) - assert text_message.type == ToolInvokeMessage.MessageType.TEXT - assert text_message.message is not None - # Verify the text contains the JSON-serialized string - # Check if message is a TextMessage - assert isinstance(text_message.message, ToolInvokeMessage.TextMessage) - # Verify it's a valid JSON string - parsed_back = json.loads(text_message.message.text) - assert parsed_back == json_value - - # Verify _invoke yields json message - json_message = _get_message_by_type(result, ToolInvokeMessage.JsonMessage) - assert json_message is None, "_invoke should not yield a JSON message for JSON array response" - - @patch("core.tools.custom_tool.tool.ssrf_proxy.get") - def test_invoke_with_text_response_creates_text_message_with_original_text( - self, mock_get: Mock, database_session: Session - ) -> None: - """Test that when upstream returns plain text, the output Text message contains the original text.""" - # Plain text exercises the response's real JSON decoding failure. - text_response_data = "This is a plain text response" - mock_get.return_value = httpx.Response(200, text=text_response_data) - - # Invoke the tool - result_generator = self.api_tool._invoke(session=database_session, user_id="test_user", tool_parameters={}) - - # Get the result from the generator - result = list(result_generator) - assert len(result) == 1 - - # Verify it's a text message with the original text - message = result[0] - assert isinstance(message, ToolInvokeMessage) - assert message.type == ToolInvokeMessage.MessageType.TEXT - assert message.message is not None - # Check if message is a TextMessage - assert isinstance(message.message, ToolInvokeMessage.TextMessage) - assert message.message.text == text_response_data - - @patch("core.tools.custom_tool.tool.ssrf_proxy.get") - def test_invoke_with_empty_response(self, mock_get: Mock, database_session: Session) -> None: - """Test that empty responses are handled correctly.""" - mock_get.return_value = httpx.Response(200, content=b"", headers={"content-type": "application/json"}) - - # Invoke the tool - result_generator = self.api_tool._invoke(session=database_session, user_id="test_user", tool_parameters={}) - - # Get the result from the generator - result = list(result_generator) - assert len(result) == 1 - - # Verify it's a text message with the empty response message - message = result[0] - assert isinstance(message, ToolInvokeMessage) - assert message.type == ToolInvokeMessage.MessageType.TEXT - assert message.message is not None - # Check if message is a TextMessage - assert isinstance(message.message, ToolInvokeMessage.TextMessage) - assert "Empty response from the tool" in message.message.text - - @patch("core.tools.custom_tool.tool.ssrf_proxy.get") - def test_invoke_with_error_response(self, mock_get: Mock, database_session: Session) -> None: - """Test that error responses are handled correctly.""" - mock_get.return_value = httpx.Response(404, text="Not Found") - - result_generator = self.api_tool._invoke(session=database_session, user_id="test_user", tool_parameters={}) - - # Invoke the tool and expect an error - with pytest.raises(Exception, match="Request failed with status code 404 and Not Found") as exc_info: - list(result_generator) # Consume the generator to trigger the error - - # Verify the error message - assert "Request failed with status code 404" in str(exc_info.value) diff --git a/api/tests/unit_tests/tools/test_mcp_tool.py b/api/tests/unit_tests/tools/test_mcp_tool.py deleted file mode 100644 index e841bd03c50..00000000000 --- a/api/tests/unit_tests/tools/test_mcp_tool.py +++ /dev/null @@ -1,376 +0,0 @@ -import base64 -from collections.abc import Iterator -from decimal import Decimal -from typing import Any -from unittest.mock import patch - -import pytest -from sqlalchemy.engine import Engine -from sqlalchemy.orm import Session - -from core.mcp.types import ( - AudioContent, - BlobResourceContents, - CallToolResult, - EmbeddedResource, - ImageContent, - TextContent, - TextResourceContents, -) -from core.tools.__base.tool_runtime import ToolRuntime -from core.tools.entities.common_entities import I18nObject -from core.tools.entities.tool_entities import ToolEntity, ToolIdentity, ToolInvokeMessage -from core.tools.mcp_tool.tool import MCPTool -from graphon.model_runtime.entities.llm_entities import LLMUsage - - -@pytest.fixture -def orm_session(sqlite_engine: Engine) -> Iterator[Session]: - """Use a real ORM session while MCP transport remains mocked.""" - with Session(sqlite_engine) as session: - yield session - - -def _make_mcp_tool(output_schema: dict[str, Any] | None = None) -> MCPTool: - identity = ToolIdentity( - author="test", - name="test_mcp_tool", - label=I18nObject(en_US="Test MCP Tool", zh_Hans="测试MCP工具"), - provider="test_provider", - ) - entity = ToolEntity(identity=identity, output_schema=output_schema or {}) - runtime = ToolRuntime(tenant_id="test_tenant") - return MCPTool( - entity=entity, - runtime=runtime, - tenant_id="test_tenant", - icon="", - server_url="https://server.invalid", - server_identifier="provider_1", - headers={}, - ) - - -class TestMCPToolInvoke: - @pytest.mark.parametrize( - ("content_factory", "mime_type"), - [ - ( - lambda b64, mt: ImageContent(type="image", data=b64, mimeType=mt), - "image/png", - ), - ( - lambda b64, mt: AudioContent(type="audio", data=b64, mimeType=mt), - "audio/mpeg", - ), - ], - ) - def test_invoke_image_or_audio_yields_blob(self, content_factory, mime_type, orm_session: Session) -> None: - tool = _make_mcp_tool() - raw = b"\x00\x01test-bytes\x02" - b64 = base64.b64encode(raw).decode() - content = content_factory(b64, mime_type) - result = CallToolResult(content=[content]) - - with patch.object(tool, "invoke_remote_mcp_tool", return_value=result): - messages = list(tool._invoke(session=orm_session, user_id="test_user", tool_parameters={})) - - assert len(messages) == 1 - msg = messages[0] - assert msg.type == ToolInvokeMessage.MessageType.BLOB - assert isinstance(msg.message, ToolInvokeMessage.BlobMessage) - assert msg.message.blob == raw - assert msg.meta == {"mime_type": mime_type} - - def test_invoke_embedded_text_resource_yields_text(self, orm_session: Session) -> None: - tool = _make_mcp_tool() - text_resource = TextResourceContents(uri="file://test.txt", mimeType="text/plain", text="hello world") - content = EmbeddedResource(type="resource", resource=text_resource) - result = CallToolResult(content=[content]) - - with patch.object(tool, "invoke_remote_mcp_tool", return_value=result): - messages = list(tool._invoke(session=orm_session, user_id="test_user", tool_parameters={})) - - assert len(messages) == 1 - msg = messages[0] - assert msg.type == ToolInvokeMessage.MessageType.TEXT - assert isinstance(msg.message, ToolInvokeMessage.TextMessage) - assert msg.message.text == "hello world" - - @pytest.mark.parametrize( - ("mime_type", "expected_mime"), - [("application/pdf", "application/pdf"), (None, "application/octet-stream")], - ) - def test_invoke_embedded_blob_resource_yields_blob(self, mime_type, expected_mime, orm_session: Session) -> None: - tool = _make_mcp_tool() - raw = b"binary-data" - b64 = base64.b64encode(raw).decode() - blob_resource = BlobResourceContents(uri="file://doc.bin", mimeType=mime_type, blob=b64) - content = EmbeddedResource(type="resource", resource=blob_resource) - result = CallToolResult(content=[content]) - - with patch.object(tool, "invoke_remote_mcp_tool", return_value=result): - messages = list(tool._invoke(session=orm_session, user_id="test_user", tool_parameters={})) - - assert len(messages) == 1 - msg = messages[0] - assert msg.type == ToolInvokeMessage.MessageType.BLOB - assert isinstance(msg.message, ToolInvokeMessage.BlobMessage) - assert msg.message.blob == raw - assert msg.meta == {"mime_type": expected_mime} - - def test_invoke_yields_variables_when_structured_content_and_schema(self, orm_session: Session) -> None: - tool = _make_mcp_tool(output_schema={"type": "object"}) - result = CallToolResult(content=[], structuredContent={"a": 1, "b": "x"}) - - with patch.object(tool, "invoke_remote_mcp_tool", return_value=result): - messages = list(tool._invoke(session=orm_session, user_id="test_user", tool_parameters={})) - - # Expect two variable messages corresponding to keys a and b - assert len(messages) == 2 - var_msgs = [m for m in messages if isinstance(m.message, ToolInvokeMessage.VariableMessage)] - assert {m.message.variable_name for m in var_msgs} == {"a", "b"} - # Validate values - values = {m.message.variable_name: m.message.variable_value for m in var_msgs} - assert values == {"a": 1, "b": "x"} - - def test_invoke_yields_json_when_structured_content_has_no_output_schema(self, orm_session: Session) -> None: - tool = _make_mcp_tool() - result = CallToolResult(content=[], structuredContent={"a": 1, "b": "x"}) - - with patch.object(tool, "invoke_remote_mcp_tool", return_value=result): - messages = list(tool._invoke(session=orm_session, user_id="test_user", tool_parameters={})) - - assert len(messages) == 1 - msg = messages[0] - assert msg.type == ToolInvokeMessage.MessageType.JSON - assert isinstance(msg.message, ToolInvokeMessage.JsonMessage) - assert msg.message.json_object == {"a": 1, "b": "x"} - - -class TestMCPToolUsageExtraction: - """Test usage metadata extraction from MCP tool results.""" - - def test_extract_usage_dict_from_direct_usage_field(self) -> None: - """Test extraction when usage is directly in meta.usage field.""" - meta = { - "usage": { - "prompt_tokens": 100, - "completion_tokens": 50, - "total_tokens": 150, - "total_price": "0.001", - "currency": "USD", - } - } - usage_dict = MCPTool._extract_usage_dict(meta) - assert usage_dict is not None - assert usage_dict["prompt_tokens"] == 100 - assert usage_dict["completion_tokens"] == 50 - assert usage_dict["total_tokens"] == 150 - assert usage_dict["total_price"] == "0.001" - assert usage_dict["currency"] == "USD" - - def test_extract_usage_dict_from_nested_metadata(self) -> None: - """Test extraction when usage is nested in meta.metadata.usage.""" - meta = { - "metadata": { - "usage": { - "prompt_tokens": 200, - "completion_tokens": 100, - "total_tokens": 300, - } - } - } - usage_dict = MCPTool._extract_usage_dict(meta) - assert usage_dict is not None - assert usage_dict["prompt_tokens"] == 200 - assert usage_dict["total_tokens"] == 300 - - def test_extract_usage_dict_from_flat_token_fields(self) -> None: - """Test extraction when token counts are directly in meta.""" - meta = { - "prompt_tokens": 150, - "completion_tokens": 75, - "total_tokens": 225, - "currency": "EUR", - } - usage_dict = MCPTool._extract_usage_dict(meta) - assert usage_dict is not None - assert usage_dict["prompt_tokens"] == 150 - assert usage_dict["completion_tokens"] == 75 - assert usage_dict["total_tokens"] == 225 - assert usage_dict["currency"] == "EUR" - - def test_extract_usage_dict_recursive(self) -> None: - """Test recursive search through nested structures.""" - meta = { - "custom": { - "nested": { - "usage": { - "total_tokens": 500, - "prompt_tokens": 300, - "completion_tokens": 200, - } - } - } - } - usage_dict = MCPTool._extract_usage_dict(meta) - assert usage_dict is not None - assert usage_dict["total_tokens"] == 500 - - def test_extract_usage_dict_from_list(self) -> None: - """Test extraction from nested list structures.""" - meta = { - "items": [ - {"usage": {"total_tokens": 100}}, - {"other": "data"}, - ] - } - usage_dict = MCPTool._extract_usage_dict(meta) - assert usage_dict is not None - assert usage_dict["total_tokens"] == 100 - - def test_extract_usage_dict_returns_none_when_missing(self) -> None: - """Test that None is returned when no usage data is present.""" - meta = {"other": "data", "custom": {"nested": {"value": 123}}} - usage_dict = MCPTool._extract_usage_dict(meta) - assert usage_dict is None - - def test_extract_usage_dict_empty_meta(self) -> None: - """Test with empty meta dict.""" - usage_dict = MCPTool._extract_usage_dict({}) - assert usage_dict is None - - def test_derive_usage_from_result_with_meta(self) -> None: - """Test _derive_usage_from_result with populated meta.""" - meta = { - "usage": { - "prompt_tokens": 100, - "completion_tokens": 50, - "total_tokens": 150, - "total_price": "0.0015", - "currency": "USD", - } - } - result = CallToolResult(content=[], _meta=meta) - usage = MCPTool._derive_usage_from_result(result) - - assert isinstance(usage, LLMUsage) - assert usage.prompt_tokens == 100 - assert usage.completion_tokens == 50 - assert usage.total_tokens == 150 - assert usage.total_price == Decimal("0.0015") - assert usage.currency == "USD" - - def test_derive_usage_from_result_without_meta(self) -> None: - """Test _derive_usage_from_result with no meta returns empty usage.""" - result = CallToolResult(content=[], meta=None) - usage = MCPTool._derive_usage_from_result(result) - - assert isinstance(usage, LLMUsage) - assert usage.total_tokens == 0 - assert usage.prompt_tokens == 0 - assert usage.completion_tokens == 0 - - def test_derive_usage_from_result_calculates_total_tokens(self) -> None: - """Test that total_tokens is calculated when missing.""" - meta = { - "usage": { - "prompt_tokens": 100, - "completion_tokens": 50, - # total_tokens is missing - } - } - result = CallToolResult(content=[], _meta=meta) - usage = MCPTool._derive_usage_from_result(result) - - assert usage.total_tokens == 150 # 100 + 50 - assert usage.prompt_tokens == 100 - assert usage.completion_tokens == 50 - - def test_invoke_sets_latest_usage_from_meta(self, orm_session: Session) -> None: - """Test that _invoke sets _latest_usage from result meta.""" - tool = _make_mcp_tool() - meta = { - "usage": { - "prompt_tokens": 200, - "completion_tokens": 100, - "total_tokens": 300, - "total_price": "0.003", - "currency": "USD", - } - } - result = CallToolResult(content=[TextContent(type="text", text="test")], _meta=meta) - - with patch.object(tool, "invoke_remote_mcp_tool", return_value=result): - list(tool._invoke(session=orm_session, user_id="test_user", tool_parameters={})) - - # Verify latest_usage was set correctly - assert tool.latest_usage.prompt_tokens == 200 - assert tool.latest_usage.completion_tokens == 100 - assert tool.latest_usage.total_tokens == 300 - assert tool.latest_usage.total_price == Decimal("0.003") - - def test_invoke_with_no_meta_returns_empty_usage(self, orm_session: Session) -> None: - """Test that _invoke returns empty usage when no meta is present.""" - tool = _make_mcp_tool() - result = CallToolResult(content=[TextContent(type="text", text="test")], _meta=None) - - with patch.object(tool, "invoke_remote_mcp_tool", return_value=result): - list(tool._invoke(session=orm_session, user_id="test_user", tool_parameters={})) - - # Verify latest_usage is empty - assert tool.latest_usage.total_tokens == 0 - assert tool.latest_usage.prompt_tokens == 0 - assert tool.latest_usage.completion_tokens == 0 - - def test_latest_usage_property_returns_llm_usage(self) -> None: - """Test that latest_usage property returns LLMUsage instance.""" - tool = _make_mcp_tool() - assert isinstance(tool.latest_usage, LLMUsage) - - def test_initial_usage_is_empty(self) -> None: - """Test that MCPTool is initialized with empty usage.""" - tool = _make_mcp_tool() - assert tool.latest_usage.total_tokens == 0 - assert tool.latest_usage.prompt_tokens == 0 - assert tool.latest_usage.completion_tokens == 0 - assert tool.latest_usage.total_price == Decimal(0) - - @pytest.mark.parametrize( - "meta_data", - [ - # Direct usage field - {"usage": {"prompt_tokens": 10, "completion_tokens": 5, "total_tokens": 15}}, - # Nested metadata - {"metadata": {"usage": {"total_tokens": 100}}}, - # Flat token fields - {"total_tokens": 50, "prompt_tokens": 30, "completion_tokens": 20}, - # With price info - { - "usage": { - "total_tokens": 150, - "total_price": "0.002", - "currency": "EUR", - } - }, - # Deep nested - {"level1": {"level2": {"usage": {"total_tokens": 200}}}}, - ], - ) - def test_various_meta_formats(self, meta_data) -> None: - """Test that various meta formats are correctly parsed.""" - result = CallToolResult(content=[], _meta=meta_data) - usage = MCPTool._derive_usage_from_result(result) - - assert isinstance(usage, LLMUsage) - # Should have at least some usage data - if meta_data.get("usage", {}).get("total_tokens") or meta_data.get("total_tokens"): - expected_total = ( - meta_data.get("usage", {}).get("total_tokens") - or meta_data.get("total_tokens") - or meta_data.get("metadata", {}).get("usage", {}).get("total_tokens") - or meta_data.get("level1", {}).get("level2", {}).get("usage", {}).get("total_tokens") - ) - if expected_total: - assert usage.total_tokens == expected_total diff --git a/api/tests/unit_tests/utils/http_parser/test_oauth_convert_request_to_raw_data.py b/api/tests/unit_tests/utils/http_parser/test_oauth_convert_request_to_raw_data.py deleted file mode 100644 index 1fa7c863835..00000000000 --- a/api/tests/unit_tests/utils/http_parser/test_oauth_convert_request_to_raw_data.py +++ /dev/null @@ -1,75 +0,0 @@ -import json - -from werkzeug import Request -from werkzeug.datastructures import Headers -from werkzeug.test import EnvironBuilder - -from core.plugin.impl.oauth import OAuthHandler - - -def test_oauth_convert_request_to_raw_data() -> None: - oauth_handler = OAuthHandler() - builder = EnvironBuilder( - method="GET", - path="/test", - headers=Headers({"Content-Type": "application/json"}), - ) - request = Request(builder.get_environ()) - raw_request_bytes = oauth_handler._convert_request_to_raw_data(request) - - assert b"GET /test? HTTP/1.1" in raw_request_bytes - assert b"Content-Type: application/json" in raw_request_bytes - assert b"\r\n\r\n" in raw_request_bytes - - -def test_oauth_convert_request_to_raw_data_with_query_params() -> None: - oauth_handler = OAuthHandler() - builder = EnvironBuilder( - method="GET", - path="/test", - query_string="code=abc123&state=xyz789", - headers=Headers({"Content-Type": "application/json"}), - ) - request = Request(builder.get_environ()) - raw_request_bytes = oauth_handler._convert_request_to_raw_data(request) - - assert b"GET /test?code=abc123&state=xyz789 HTTP/1.1" in raw_request_bytes - assert b"Content-Type: application/json" in raw_request_bytes - assert b"\r\n\r\n" in raw_request_bytes - - -def test_oauth_convert_request_to_raw_data_with_post_body() -> None: - oauth_handler = OAuthHandler() - builder = EnvironBuilder( - method="POST", - path="/test", - data="param1=value1¶m2=value2", - headers=Headers({"Content-Type": "application/x-www-form-urlencoded"}), - ) - request = Request(builder.get_environ()) - raw_request_bytes = oauth_handler._convert_request_to_raw_data(request) - - assert b"POST /test? HTTP/1.1" in raw_request_bytes - assert b"Content-Type: application/x-www-form-urlencoded" in raw_request_bytes - assert b"\r\n\r\n" in raw_request_bytes - assert b"param1=value1¶m2=value2" in raw_request_bytes - - -def test_oauth_convert_request_to_raw_data_with_json_body() -> None: - oauth_handler = OAuthHandler() - json_data = {"code": "abc123", "state": "xyz789", "grant_type": "authorization_code"} - builder = EnvironBuilder( - method="POST", - path="/test", - data=json.dumps(json_data), - headers=Headers({"Content-Type": "application/json"}), - ) - request = Request(builder.get_environ()) - raw_request_bytes = oauth_handler._convert_request_to_raw_data(request) - - assert b"POST /test? HTTP/1.1" in raw_request_bytes - assert b"Content-Type: application/json" in raw_request_bytes - assert b"\r\n\r\n" in raw_request_bytes - assert b'"code": "abc123"' in raw_request_bytes - assert b'"state": "xyz789"' in raw_request_bytes - assert b'"grant_type": "authorization_code"' in raw_request_bytes diff --git a/api/tests/unit_tests/utils/structured_output_parser/test_structured_output_parser.py b/api/tests/unit_tests/utils/structured_output_parser/test_structured_output_parser.py deleted file mode 100644 index 372a8a25ec8..00000000000 --- a/api/tests/unit_tests/utils/structured_output_parser/test_structured_output_parser.py +++ /dev/null @@ -1,497 +0,0 @@ -from decimal import Decimal -from typing import Literal, NotRequired -from unittest.mock import MagicMock, patch - -import pytest -from typing_extensions import TypedDict - -from core.llm_generator.output_parser.errors import OutputParserError -from core.llm_generator.output_parser.structured_output import invoke_llm_with_structured_output -from graphon.model_runtime.entities.common_entities import I18nObject -from graphon.model_runtime.entities.llm_entities import ( - LLMResult, - LLMResultChunk, - LLMResultChunkDelta, - LLMResultChunkWithStructuredOutput, - LLMResultWithStructuredOutput, - LLMUsage, -) -from graphon.model_runtime.entities.message_entities import ( - AssistantPromptMessage, - SystemPromptMessage, - TextPromptMessageContent, - UserPromptMessage, -) -from graphon.model_runtime.entities.model_entities import AIModelEntity, ModelType, ParameterRule, ParameterType - - -class _StructuredOutputCase(TypedDict): - name: str - provider: str - model_name: str - support_structure_output: bool - stream: bool - json_schema: dict[str, object] - expected_llm_response: LLMResult | list[LLMResultChunk] | list[MagicMock] - expected_result_type: type[LLMResultWithStructuredOutput] | Literal["generator"] | None - should_raise: bool - parameter_rules: NotRequired[list[ParameterRule]] - expected_response_format: NotRequired[str] - expected_error: NotRequired[type[Exception]] - - -def _response_format_rule(*options: str) -> ParameterRule: - return ParameterRule( - name="response_format", - label=I18nObject(en_US="Response Format"), - type=ParameterType.STRING, - options=list(options), - ) - - -def create_mock_usage(prompt_tokens: int = 10, completion_tokens: int = 5) -> LLMUsage: - """Create a mock LLMUsage with all required fields""" - return LLMUsage( - prompt_tokens=prompt_tokens, - prompt_unit_price=Decimal("0.001"), - prompt_price_unit=Decimal(1), - prompt_price=Decimal(str(prompt_tokens)) * Decimal("0.001"), - completion_tokens=completion_tokens, - completion_unit_price=Decimal("0.002"), - completion_price_unit=Decimal(1), - completion_price=Decimal(str(completion_tokens)) * Decimal("0.002"), - total_tokens=prompt_tokens + completion_tokens, - total_price=Decimal(str(prompt_tokens)) * Decimal("0.001") + Decimal(str(completion_tokens)) * Decimal("0.002"), - currency="USD", - latency=1.5, - ) - - -def get_model_entity(provider: str, model_name: str, support_structure_output: bool = False) -> AIModelEntity: - """Create a mock AIModelEntity for testing""" - model_schema = MagicMock() - model_schema.model = model_name - model_schema.provider = provider - model_schema.model_type = ModelType.LLM - model_schema.model_provider = provider - model_schema.model_name = model_name - model_schema.support_structure_output = support_structure_output - model_schema.parameter_rules = list[object]() - - return model_schema - - -def get_model_instance() -> MagicMock: - """Create a mock ModelInstance for testing""" - mock_instance = MagicMock() - mock_instance.provider = "openai" - mock_instance.credentials = dict[str, object]() - return mock_instance - - -def test_structured_output_parser() -> None: - """Test cases for invoke_llm_with_structured_output function""" - - testcases: list[_StructuredOutputCase] = [ - # Test case 1: Model with native structured output support, non-streaming - { - "name": "native_structured_output_non_streaming", - "provider": "openai", - "model_name": "gpt-4o", - "support_structure_output": True, - "stream": False, - "json_schema": {"type": "object", "properties": {"name": {"type": "string"}}}, - "expected_llm_response": LLMResult( - model="gpt-4o", - message=AssistantPromptMessage(content='{"name": "test"}'), - usage=create_mock_usage(prompt_tokens=10, completion_tokens=5), - ), - "expected_result_type": LLMResultWithStructuredOutput, - "should_raise": False, - }, - # Test case 2: Model with native structured output support, streaming - { - "name": "native_structured_output_streaming", - "provider": "openai", - "model_name": "gpt-4o", - "support_structure_output": True, - "stream": True, - "json_schema": {"type": "object", "properties": {"name": {"type": "string"}}}, - "expected_llm_response": [ - LLMResultChunk( - model="gpt-4o", - prompt_messages=[UserPromptMessage(content="test")], - system_fingerprint="test", - delta=LLMResultChunkDelta( - index=0, - message=AssistantPromptMessage(content='{"name":'), - usage=create_mock_usage(prompt_tokens=10, completion_tokens=2), - ), - ), - LLMResultChunk( - model="gpt-4o", - prompt_messages=[UserPromptMessage(content="test")], - system_fingerprint="test", - delta=LLMResultChunkDelta( - index=0, - message=AssistantPromptMessage(content=' "test"}'), - usage=create_mock_usage(prompt_tokens=10, completion_tokens=3), - ), - ), - ], - "expected_result_type": "generator", - "should_raise": False, - }, - # Test case 3: Model without native structured output support, non-streaming - { - "name": "prompt_based_structured_output_non_streaming", - "provider": "anthropic", - "model_name": "claude-3-sonnet", - "support_structure_output": False, - "stream": False, - "json_schema": {"type": "object", "properties": {"answer": {"type": "string"}}}, - "expected_llm_response": LLMResult( - model="claude-3-sonnet", - message=AssistantPromptMessage(content='{"answer": "test response"}'), - usage=create_mock_usage(prompt_tokens=15, completion_tokens=8), - ), - "expected_result_type": LLMResultWithStructuredOutput, - "should_raise": False, - }, - # Test case 4: Model without native structured output support, streaming - { - "name": "prompt_based_structured_output_streaming", - "provider": "anthropic", - "model_name": "claude-3-sonnet", - "support_structure_output": False, - "stream": True, - "json_schema": {"type": "object", "properties": {"answer": {"type": "string"}}}, - "expected_llm_response": [ - LLMResultChunk( - model="claude-3-sonnet", - prompt_messages=[UserPromptMessage(content="test")], - system_fingerprint="test", - delta=LLMResultChunkDelta( - index=0, - message=AssistantPromptMessage(content='{"answer": "test'), - usage=create_mock_usage(prompt_tokens=15, completion_tokens=3), - ), - ), - LLMResultChunk( - model="claude-3-sonnet", - prompt_messages=[UserPromptMessage(content="test")], - system_fingerprint="test", - delta=LLMResultChunkDelta( - index=0, - message=AssistantPromptMessage(content=' response"}'), - usage=create_mock_usage(prompt_tokens=15, completion_tokens=5), - ), - ), - ], - "expected_result_type": "generator", - "should_raise": False, - }, - # Test case 5: Streaming with list content - { - "name": "streaming_with_list_content", - "provider": "openai", - "model_name": "gpt-4o", - "support_structure_output": True, - "stream": True, - "json_schema": {"type": "object", "properties": {"data": {"type": "string"}}}, - "expected_llm_response": [ - LLMResultChunk( - model="gpt-4o", - prompt_messages=[UserPromptMessage(content="test")], - system_fingerprint="test", - delta=LLMResultChunkDelta( - index=0, - message=AssistantPromptMessage( - content=[ - TextPromptMessageContent(data='{"data":'), - ] - ), - usage=create_mock_usage(prompt_tokens=10, completion_tokens=2), - ), - ), - LLMResultChunk( - model="gpt-4o", - prompt_messages=[UserPromptMessage(content="test")], - system_fingerprint="test", - delta=LLMResultChunkDelta( - index=0, - message=AssistantPromptMessage( - content=[ - TextPromptMessageContent(data=' "value"}'), - ] - ), - usage=create_mock_usage(prompt_tokens=10, completion_tokens=3), - ), - ), - ], - "expected_result_type": "generator", - "should_raise": False, - }, - # Test case 6: Error case - non-string LLM response content (non-streaming) - { - "name": "error_non_string_content_non_streaming", - "provider": "openai", - "model_name": "gpt-4o", - "support_structure_output": True, - "stream": False, - "json_schema": {"type": "object", "properties": {"name": {"type": "string"}}}, - "expected_llm_response": LLMResult( - model="gpt-4o", - message=AssistantPromptMessage(content=None), # Non-string content - usage=create_mock_usage(prompt_tokens=10, completion_tokens=5), - ), - "expected_result_type": None, - "should_raise": True, - "expected_error": OutputParserError, - }, - # Test case 7: JSON repair scenario - { - "name": "json_repair_scenario", - "provider": "openai", - "model_name": "gpt-4o", - "support_structure_output": True, - "stream": False, - "json_schema": {"type": "object", "properties": {"name": {"type": "string"}}}, - "expected_llm_response": LLMResult( - model="gpt-4o", - message=AssistantPromptMessage(content='{"name": "test"'), # Invalid JSON - missing closing brace - usage=create_mock_usage(prompt_tokens=10, completion_tokens=5), - ), - "expected_result_type": LLMResultWithStructuredOutput, - "should_raise": False, - }, - # Test case 8: Model with parameter rules for response format - { - "name": "model_with_parameter_rules", - "provider": "openai", - "model_name": "gpt-4o", - "support_structure_output": True, - "stream": False, - "json_schema": {"type": "object", "properties": {"result": {"type": "string"}}}, - "parameter_rules": [_response_format_rule("json_schema")], - "expected_response_format": "json_schema", - "expected_llm_response": LLMResult( - model="gpt-4o", - message=AssistantPromptMessage(content='{"result": "success"}'), - usage=create_mock_usage(prompt_tokens=10, completion_tokens=5), - ), - "expected_result_type": LLMResultWithStructuredOutput, - "should_raise": False, - }, - # Test case 9: Model without native support but with JSON response format rules - { - "name": "non_native_with_json_rules", - "provider": "anthropic", - "model_name": "claude-3-sonnet", - "support_structure_output": False, - "stream": False, - "json_schema": {"type": "object", "properties": {"output": {"type": "string"}}}, - "parameter_rules": [_response_format_rule("JSON")], - "expected_response_format": "JSON", - "expected_llm_response": LLMResult( - model="claude-3-sonnet", - message=AssistantPromptMessage(content='{"output": "result"}'), - usage=create_mock_usage(prompt_tokens=15, completion_tokens=8), - ), - "expected_result_type": LLMResultWithStructuredOutput, - "should_raise": False, - }, - ] - - for case in testcases: - # Setup model entity - model_schema = get_model_entity(case["provider"], case["model_name"], case["support_structure_output"]) - - # Add parameter rules if specified - if "parameter_rules" in case: - model_schema.parameter_rules = case["parameter_rules"] - - # Setup model instance - model_instance = get_model_instance() - model_instance.invoke_llm.return_value = case["expected_llm_response"] - - # Setup prompt messages - prompt_messages = [ - SystemPromptMessage(content="You are a helpful assistant."), - UserPromptMessage(content="Generate a response according to the schema."), - ] - - if case["should_raise"]: - # Test error cases - expected_error = case.get("expected_error") - assert expected_error is not None - with pytest.raises(expected_error): # noqa: PT012 - if case["stream"]: - result_generator = invoke_llm_with_structured_output( - provider=case["provider"], - model_schema=model_schema, - model_instance=model_instance, - prompt_messages=prompt_messages, - json_schema=case["json_schema"], - stream=case["stream"], - ) - # Consume the generator to trigger the error - list(result_generator) - else: - invoke_llm_with_structured_output( - provider=case["provider"], - model_schema=model_schema, - model_instance=model_instance, - prompt_messages=prompt_messages, - json_schema=case["json_schema"], - stream=case["stream"], - ) - else: - # Test successful cases - with patch( - "core.llm_generator.output_parser.structured_output.json_repair.loads", autospec=True - ) as mock_json_repair: - # Configure json_repair mock for cases that need it - if case["name"] == "json_repair_scenario": - mock_json_repair.return_value = {"name": "test"} - - result = invoke_llm_with_structured_output( - provider=case["provider"], - model_schema=model_schema, - model_instance=model_instance, - prompt_messages=prompt_messages, - json_schema=case["json_schema"], - stream=case["stream"], - model_parameters={"temperature": 0.7, "max_tokens": 100}, - ) - - if case["expected_result_type"] == "generator": - # Test streaming results - assert hasattr(result, "__iter__") - chunks = list(result) - assert len(chunks) > 0 - - # Verify all chunks are LLMResultChunkWithStructuredOutput - for chunk in chunks[:-1]: # All except last - assert isinstance(chunk, LLMResultChunkWithStructuredOutput) - assert chunk.model == case["model_name"] - - # Last chunk should have structured output - last_chunk = chunks[-1] - assert isinstance(last_chunk, LLMResultChunkWithStructuredOutput) - assert last_chunk.structured_output is not None - assert isinstance(last_chunk.structured_output, dict) - else: - # Test non-streaming results - expected_type = case["expected_result_type"] - assert isinstance(expected_type, type) - assert isinstance(result, expected_type) - assert result.model == case["model_name"] - assert result.structured_output is not None - assert isinstance(result.structured_output, dict) - - # Verify model_instance.invoke_llm was called with correct parameters - model_instance.invoke_llm.assert_called_once() - call_args = model_instance.invoke_llm.call_args - - assert call_args.kwargs["stream"] == case["stream"] - assert "user" not in call_args.kwargs - assert "temperature" in call_args.kwargs["model_parameters"] - if "expected_response_format" in case: - assert call_args.kwargs["model_parameters"]["response_format"] == case["expected_response_format"] - assert "max_tokens" in call_args.kwargs["model_parameters"] - - -def test_parse_structured_output_edge_cases() -> None: - """Test edge cases for structured output parsing""" - - # Test case with list that contains dict (reasoning model scenario) - testcase_list_with_dict = { - "name": "list_with_dict_parsing", - "provider": "deepseek", - "model_name": "deepseek-r1", - "support_structure_output": False, - "stream": False, - "json_schema": {"type": "object", "properties": {"thought": {"type": "string"}}}, - "expected_llm_response": LLMResult( - model="deepseek-r1", - message=AssistantPromptMessage(content='[{"thought": "reasoning process"}, "other content"]'), - usage=create_mock_usage(prompt_tokens=10, completion_tokens=5), - ), - "expected_result_type": LLMResultWithStructuredOutput, - "should_raise": False, - } - - # Setup for list parsing test - model_schema = get_model_entity( - testcase_list_with_dict["provider"], - testcase_list_with_dict["model_name"], - testcase_list_with_dict["support_structure_output"], - ) - - model_instance = get_model_instance() - model_instance.invoke_llm.return_value = testcase_list_with_dict["expected_llm_response"] - - prompt_messages = [UserPromptMessage(content="Test reasoning")] - - with patch( - "core.llm_generator.output_parser.structured_output.json_repair.loads", autospec=True - ) as mock_json_repair: - # Mock json_repair to return a list with dict - mock_json_repair.return_value = [{"thought": "reasoning process"}, "other content"] - - result = invoke_llm_with_structured_output( - provider=testcase_list_with_dict["provider"], - model_schema=model_schema, - model_instance=model_instance, - prompt_messages=prompt_messages, - json_schema=testcase_list_with_dict["json_schema"], - stream=testcase_list_with_dict["stream"], - ) - - assert isinstance(result, LLMResultWithStructuredOutput) - assert result.structured_output == {"thought": "reasoning process"} - - -def test_model_specific_schema_preparation() -> None: - """Test schema preparation for different model types""" - - # Test Gemini model - gemini_case = { - "provider": "google", - "model_name": "gemini-pro", - "support_structure_output": True, - "stream": False, - "json_schema": {"type": "object", "properties": {"result": {"type": "boolean"}}, "additionalProperties": False}, - } - - model_schema = get_model_entity( - gemini_case["provider"], gemini_case["model_name"], gemini_case["support_structure_output"] - ) - - model_instance = get_model_instance() - model_instance.invoke_llm.return_value = LLMResult( - model="gemini-pro", - message=AssistantPromptMessage(content='{"result": "true"}'), - usage=create_mock_usage(prompt_tokens=10, completion_tokens=5), - ) - - prompt_messages = [UserPromptMessage(content="Test")] - - result = invoke_llm_with_structured_output( - provider=gemini_case["provider"], - model_schema=model_schema, - model_instance=model_instance, - prompt_messages=prompt_messages, - json_schema=gemini_case["json_schema"], - stream=gemini_case["stream"], - ) - - assert isinstance(result, LLMResultWithStructuredOutput) - - # Verify model_instance.invoke_llm was called and check the schema preparation - model_instance.invoke_llm.assert_called_once() - call_args = model_instance.invoke_llm.call_args - - # For Gemini, the schema should not have additionalProperties and boolean should be converted to string - assert "json_schema" in call_args.kwargs["model_parameters"] diff --git a/dev/pyrefly-check-local b/dev/pyrefly-check-local index 01a15db79be..34274b3cbcf 100755 --- a/dev/pyrefly-check-local +++ b/dev/pyrefly-check-local @@ -8,7 +8,7 @@ cd "$REPO_ROOT" EXCLUDES_FILE="api/pyrefly-local-excludes.txt" UNIT_TESTS_CONFIG="tests/unit_tests/pyrefly.toml" -TEST_CONTAINERS_DIR="tests/test_containers_integration_tests" +TEST_CONTAINERS_DIR="tests/e2e_tests/containers" TEST_CONTAINERS_CONFIG="$TEST_CONTAINERS_DIR/pyrefly.toml" target_paths=() diff --git a/e2e/scripts/common.ts b/e2e/scripts/common.ts index d27bbc06b3c..f041d28d8e5 100644 --- a/e2e/scripts/common.ts +++ b/e2e/scripts/common.ts @@ -42,7 +42,13 @@ export const middlewareEnvFile = path.join(dockerDir, 'middleware.env') export const middlewareEnvExampleFile = path.join(dockerDir, 'envs', 'middleware.env.example') export const webEnvLocalFile = path.join(webDir, '.env.local') export const webEnvExampleFile = path.join(webDir, '.env.example') -export const apiEnvExampleFile = path.join(apiDir, 'tests', 'integration_tests', '.env.example') +export const apiEnvExampleFile = path.join( + apiDir, + 'tests', + 'e2e_tests', + 'middleware', + '.env.example', +) export const e2eWebEnvOverrides = { NEXT_PUBLIC_API_PREFIX: 'http://127.0.0.1:5001/console/api', NEXT_PUBLIC_ENABLE_AGENT_V2: 'true', diff --git a/scripts/ast_grep_guard.py b/scripts/ast_grep_guard.py index 9f67bfd6d08..85415cdbf63 100644 --- a/scripts/ast_grep_guard.py +++ b/scripts/ast_grep_guard.py @@ -4,18 +4,20 @@ from __future__ import annotations import argparse +import ast import json import re import shutil import subprocess import sys import tempfile +import textwrap +from collections import Counter from collections.abc import Callable from dataclasses import dataclass from pathlib import Path from typing import Any - REPO_ROOT = Path(__file__).resolve().parents[1] SCAN_ROOT = Path.cwd() RULES_ROOT = REPO_ROOT / "scripts" / "ast_grep_rules" @@ -106,11 +108,19 @@ def git_output(*args: str, allow_missing: bool = False) -> str: def collect_diff_text(args: argparse.Namespace) -> str: if args.staged: - return git_output("diff", "--cached", "--unified=0", "--diff-filter=AM", "--no-ext-diff") + return git_output( + "diff", + "--cached", + "--unified=0", + "--no-renames", + "--diff-filter=AMD", + "--no-ext-diff", + ) return git_output( "diff", "--unified=0", - "--diff-filter=AM", + "--no-renames", + "--diff-filter=AMD", "--no-ext-diff", f"{args.base_rev}..HEAD", ) @@ -160,12 +170,12 @@ def load_file_versions(path: str, args: argparse.Namespace) -> tuple[str, str]: if args.staged: return ( git_output("show", f"HEAD:{path}", allow_missing=True), - git_output("show", f":{path}"), + git_output("show", f":{path}", allow_missing=True), ) return ( git_output("show", f"{args.base_rev}:{path}", allow_missing=True), - git_output("show", f"HEAD:{path}"), + git_output("show", f"HEAD:{path}", allow_missing=True), ) @@ -254,29 +264,50 @@ def has_reasoned_guard_ignore(source_line: str, rule_id: str) -> bool: return bool(reason.strip()) -def collect_hunk_violations( - path: str, +def match_key(match: Match) -> str: + """Ignore formatting when comparing moved Python syntax, retaining literal values.""" + return ast.dump(ast.parse(textwrap.dedent(match.text))) + + +def filter_moved_matches( + added: list[tuple[str, Match]], + removed: Counter[str], +) -> list[tuple[str, Match]]: + unmatched = [] + for path, match in added: + key = match_key(match) + if removed[key]: + removed[key] -= 1 + else: + unmatched.append((path, match)) + return unmatched + + +def collect_hunk_changes( old_matches: list[Match], new_matches: list[Match], hunk: Hunk, *, is_reportable_match: MatchPredicate, - violation_message: str, -) -> list[Violation]: +) -> tuple[Counter[str], list[Match]]: old_in_hunk = [ match for match in old_matches if hunk.contains_old(match.line_number) and is_reportable_match(match) ] new_in_hunk = [ match for match in new_matches if hunk.contains_new(match.line_number) and is_reportable_match(match) ] - surplus = len(new_in_hunk) - len(old_in_hunk) - if surplus <= 0: - return [] - - return [ - Violation(path=path, line_number=match.line_number, message=violation_message) - for match in new_in_hunk[-surplus:] - ] + existing = Counter(match_key(match) for match in old_in_hunk) + added = [] + for match in new_in_hunk: + key = match_key(match) + if existing[key]: + existing[key] -= 1 + else: + added.append(match) + # Preserve the existing allowance for replacing legacy calls within a hunk. + remaining = list(existing.elements()) + replacements = min(len(remaining), len(added)) + return Counter(remaining[replacements:]), added[replacements:] def find_net_new_violations( @@ -288,7 +319,8 @@ def find_net_new_violations( is_reportable_match: MatchPredicate, violation_message: str, ) -> list[Violation]: - violations: list[Violation] = [] + removed: Counter[str] = Counter() + added: list[tuple[str, Match]] = [] for path, hunks in changed_hunks.items(): if not is_scanned_path(path): continue @@ -296,16 +328,17 @@ def find_net_new_violations( old_matches = run_ast_grep(old_source, rule=rule) new_matches = run_ast_grep(new_source, rule=rule) for hunk in hunks: - violations.extend( - collect_hunk_violations( - path, - old_matches, - new_matches, - hunk, - is_reportable_match=is_reportable_match, - violation_message=violation_message, - ) + removed_in_hunk, added_in_hunk = collect_hunk_changes( + old_matches, + new_matches, + hunk, + is_reportable_match=is_reportable_match, ) + removed.update(removed_in_hunk) + added.extend((path, match) for match in added_in_hunk) + violations = [ + Violation(path, match.line_number, violation_message) for path, match in filter_moved_matches(added, removed) + ] return sorted(violations, key=lambda item: (item.path, item.line_number)) diff --git a/scripts/ast_grep_rules/no_new_mock.yml b/scripts/ast_grep_rules/no_new_mock.yml index 1b77d10acfe..ff0b62255d1 100644 --- a/scripts/ast_grep_rules/no_new_mock.yml +++ b/scripts/ast_grep_rules/no_new_mock.yml @@ -44,14 +44,18 @@ rule: field: module_name stopBy: end regex: '^types$' - - inside: - kind: module - stopBy: end - has: - stopBy: end - all: - - pattern: $CALLEE = $ORIGINAL - - has: - field: right - stopBy: end - regex: '(^|\.)(Mock|MagicMock|SimpleNamespace)$' + - all: + - has: + field: function + kind: identifier + - inside: + kind: module + stopBy: end + has: + stopBy: end + all: + - pattern: $CALLEE = $ORIGINAL + - has: + field: right + stopBy: end + regex: '(^|\.)(Mock|MagicMock|SimpleNamespace)$' diff --git a/scripts/check_no_new_test_doubles.py b/scripts/check_no_new_test_doubles.py index da7ed8f27cf..d5923838f1e 100644 --- a/scripts/check_no_new_test_doubles.py +++ b/scripts/check_no_new_test_doubles.py @@ -12,7 +12,8 @@ is a syntactic guard, not general Python inheritance or execution analysis. Protocol and Pydantic model implementations are not targeted. Both rules retain the provided Git baseline. Identical mock calls are counted -per file: an additional copy or a different constructor call is rejected. +per file, with removed occurrences available to match code moved between files: +an additional copy or a different constructor call is rejected. Git-detected renames retain the source file's baseline, including paths with spaces. Subclass matches are compared by class name within each changed file, not only by the line of the class header: @@ -25,16 +26,21 @@ For a justified subclass exception only, add from __future__ import annotations import argparse +import ast import sys +import textwrap from collections import Counter +from dataclasses import replace from pathlib import Path from ast_grep_guard import ( Match, Violation, + filter_moved_matches, git_output, has_reasoned_guard_ignore, is_python_source_path, + match_key, parse_args, print_violations, rule_path, @@ -48,7 +54,13 @@ def changed_file_versions(args: argparse.Namespace) -> dict[str, tuple[str, str] """Read Git blobs, preserving rename baselines and NUL-delimited filenames.""" revisions = ["--cached"] if args.staged else [f"{args.base_rev}..HEAD"] status_text = git_output( - "diff", "--name-status", "-z", "--find-renames", "--diff-filter=AMR", "--no-ext-diff", *revisions + "diff", + "--name-status", + "-z", + "--find-renames", + "--diff-filter=AMRD", + "--no-ext-diff", + *revisions, ) if not status_text: return {} @@ -65,27 +77,56 @@ def changed_file_versions(args: argparse.Namespace) -> dict[str, tuple[str, str] if not is_python_source_path(path): continue old_source = "" if status == "A" else git_output("show", f"{old_revision}:{old_path}") - versions[path] = (old_source, git_output("show", f"{new_revision}:{path}")) + new_source = "" if status == "D" else git_output("show", f"{new_revision}:{path}") + versions[path] = (old_source, new_source) return versions +def mock_matches(source: str) -> list[Match]: + """Keep import renames from turning existing constructors into new calls.""" + bindings: dict[str, str] = {} + for node in ast.parse(source).body: + if isinstance(node, ast.Import): + for alias in node.names: + bindings[alias.asname or alias.name.split(".")[0]] = ( + alias.name if alias.asname else alias.name.split(".")[0] + ) + elif isinstance(node, ast.ImportFrom): + for alias in node.names: + bindings[alias.asname or alias.name] = alias.name + matches = [] + for match in run_ast_grep(source, rule=rule_path("no_new_mock.yml")): + tree = ast.parse(textwrap.dedent(match.text)) + for node in ast.walk(tree): + if isinstance(node, ast.Name) and node.id in bindings: + node.id = bindings[node.id] + matches.append(replace(match, text=ast.unparse(tree))) + return matches + + def find_new_mock_calls(changed: dict[str, tuple[str, str]]) -> list[Violation]: """Keep identical baseline calls, but catch alias changes outside their hunks.""" - violations: list[Violation] = [] + removed: Counter[str] = Counter() + added: list[tuple[str, Match]] = [] for path, (old_source, new_source) in changed.items(): - old_calls = Counter(match.text for match in run_ast_grep(old_source, rule=rule_path("no_new_mock.yml"))) - for match in run_ast_grep(new_source, rule=rule_path("no_new_mock.yml")): - if old_calls[match.text]: - old_calls[match.text] -= 1 + if old_source == new_source: + continue + old_calls = Counter(match_key(match) for match in mock_matches(old_source)) + for match in mock_matches(new_source): + key = match_key(match) + if old_calls[key]: + old_calls[key] -= 1 continue - violations.append( - Violation( - path=path, - line_number=match.line_number, - message="no-new-mock: instantiate a real class instead of Mock, MagicMock or SimpleNamespace", - ) - ) - return violations + added.append((path, match)) + removed.update(old_calls) + return [ + Violation( + path, + match.line_number, + "no-new-mock: instantiate a real class instead of Mock, MagicMock or SimpleNamespace", + ) + for path, match in filter_moved_matches(added, removed) + ] def is_reportable_subclass(match: Match) -> bool: @@ -93,28 +134,35 @@ def is_reportable_subclass(match: Match) -> bool: def find_new_stub_subclasses(changed: dict[str, tuple[str, str]]) -> list[Violation]: - violations: list[Violation] = [] + removed: Counter[str] = Counter() + added: list[tuple[str, Match]] = [] for path, (old_source, new_source) in changed.items(): if not Path(path).is_relative_to("api/tests"): continue old_matches = run_ast_grep(old_source, rule=rule_path("no_new_stub_subclass.yml")) new_matches = run_ast_grep(new_source, rule=rule_path("no_new_stub_subclass.yml")) - old_names = Counter(match.meta_variables["NAME"] for match in old_matches if is_reportable_subclass(match)) + old_names: dict[str, list[Match]] = {} + for match in old_matches: + if is_reportable_subclass(match): + old_names.setdefault(match.meta_variables["NAME"], []).append(match) for match in new_matches: if not is_reportable_subclass(match): continue name = match.meta_variables["NAME"] - if old_names[name]: - old_names[name] -= 1 + if old_names.get(name): + old_names[name].pop() continue - violations.append( - Violation( - path=path, - line_number=match.line_number, - message=f"{STUB_RULE_ID}: {name} bypasses a runtime constructor; instantiate the real class", - ) - ) - return violations + added.append((path, match)) + removed.update(match_key(match) for matches in old_names.values() for match in matches) + return [ + Violation( + path, + match.line_number, + f"{STUB_RULE_ID}: {match.meta_variables['NAME']} bypasses a runtime constructor; " + "instantiate the real class", + ) + for path, match in filter_moved_matches(added, removed) + ] def main() -> int: